From 9435c19d47c1a12904c70d1a40d62bbe242a6e92 Mon Sep 17 00:00:00 2001
From: Ralph Khreish <35776126+Crunchyman-ralph@users.noreply.github.com>
Date: Thu, 1 May 2025 18:52:45 +0200
Subject: [PATCH] fix: env variables and wrapper

---
 mcp-server/src/tools/add-task.js | 19 ++------
 mcp-server/src/tools/utils.js    | 81 +++++++++++++++++++++++++++++++-
 2 files changed, 85 insertions(+), 15 deletions(-)

diff --git a/mcp-server/src/tools/add-task.js b/mcp-server/src/tools/add-task.js
index 7c726995..8546c24e 100644
--- a/mcp-server/src/tools/add-task.js
+++ b/mcp-server/src/tools/add-task.js
@@ -6,7 +6,7 @@
 import { z } from 'zod';
 import {
 	createErrorResponse,
-	getProjectRootFromSession,
+	withProjectContext,
 	handleApiResult
 } from './utils.js';
 import { addTaskDirect } from '../core/task-master-core.js';
@@ -63,26 +63,17 @@ export function registerAddTaskTool(server) {
 				.optional()
 				.describe('Whether to use research capabilities for task creation')
 		}),
-		execute: async (args, { log, session }) => {
+		execute: withProjectContext(async (args, { log, session }) => {
 			try {
 				log.info(`Starting add-task with args: ${JSON.stringify(args)}`);
 
-				// Get project root from args or session
-				const rootFolder =
-					args.projectRoot || getProjectRootFromSession(session, log);
-
-				// Ensure project root was determined
-				if (!rootFolder) {
-					return createErrorResponse(
-						'Could not determine project root. Please provide it explicitly or ensure your session contains valid root information.'
-					);
-				}
+				// Note: projectRoot is now guaranteed to be normalized by withProjectContext
 
 				// Resolve the path to tasks.json
 				let tasksJsonPath;
 				try {
 					tasksJsonPath = findTasksJsonPath(
-						{ projectRoot: rootFolder, file: args.file },
+						{ projectRoot: args.projectRoot, file: args.file },
 						log
 					);
 				} catch (error) {
@@ -117,6 +108,6 @@ export function registerAddTaskTool(server) {
 				log.error(`Error in add-task tool: ${error.message}`);
 				return createErrorResponse(error.message);
 			}
-		}
+		})
 	});
 }
diff --git a/mcp-server/src/tools/utils.js b/mcp-server/src/tools/utils.js
index 71b439f3..16b3922c 100644
--- a/mcp-server/src/tools/utils.js
+++ b/mcp-server/src/tools/utils.js
@@ -7,7 +7,7 @@ import { spawnSync } from 'child_process';
 import path from 'path';
 import fs from 'fs';
 import { contextManager } from '../core/context-manager.js'; // Import the singleton
-
+import { withSessionEnv } from '../core/utils/env-utils.js';
 // Import path utilities to ensure consistent path resolution
 import {
 	lastFoundProjectRoot,
@@ -474,6 +474,85 @@ function createLogWrapper(log) {
 	};
 }
 
+/**
+ * Resolves and normalizes a project root path
+ * @param {string} rawPath - The raw project root path that might be URI encoded
+ * @returns {string} Normalized absolute path
+ */
+function normalizeProjectRoot(rawPath) {
+	if (!rawPath) return null;
+
+	try {
+		// Handle URI encoded paths (e.g. /c%20/path/with%20spaces)
+		const decoded = decodeURIComponent(rawPath);
+
+		// Convert Windows-style paths if needed
+		const normalized = decoded.replace(/\\/g, '/');
+
+		// Remove any file:// prefix
+		const withoutProtocol = normalized.replace(/^file:\/\//, '');
+
+		// Ensure absolute path
+		return path.resolve(withoutProtocol);
+	} catch (error) {
+		return null;
+	}
+}
+
+/**
+ * Higher-order function that wraps an MCP tool's execute method to provide:
+ * 1. Normalized project root resolution
+ * 2. Environment variable setup from both .env and session
+ * 3. Error handling
+ *
+ * @param {Function} executeFn - The original execute function to wrap
+ * @returns {Function} Wrapped execute function with project context
+ */
+export function withProjectContext(executeFn) {
+	return async (args, context) => {
+		const { log, session } = context;
+
+		try {
+			// 1. Resolve and normalize project root
+			const rawRoot =
+				args.projectRoot || getProjectRootFromSession(session, log);
+			const projectRoot = normalizeProjectRoot(rawRoot);
+
+			if (!projectRoot) {
+				return createErrorResponse(
+					'Could not determine project root. Please provide it explicitly or ensure your session contains valid root information.'
+				);
+			}
+
+			// 2. Load .env from project root if it exists
+			const envPath = path.join(projectRoot, '.env');
+			let localEnv = {};
+
+			if (fs.existsSync(envPath)) {
+				try {
+					localEnv = require('dotenv').config({ path: envPath }).parsed || {};
+				} catch (error) {
+					log.warn(`Failed to load .env from ${envPath}: ${error.message}`);
+				}
+			}
+
+			// 3. Combine with session env and execute
+			const combinedEnv = {
+				...localEnv,
+				...(session?.env || {})
+			};
+
+			return await withSessionEnv(combinedEnv, async () => {
+				// Pass normalized projectRoot to the execute function
+				return await executeFn({ ...args, projectRoot }, context);
+			});
+		} catch (error) {
+			log.error(`Error in project context: ${error.message}`);
+			return createErrorResponse(error.message);
+		}
+	};
+}
+
 // Ensure all functions are exported
 export {
 	getProjectRoot,