refactor: optimized mcp and tool elicitation

yamanahlawat · yamanahlawat · commit 04124c28dcb9 · 2025-07-26T00:11:57.000+05:30
diff --git a/mcp-run-python/src/main.ts b/mcp-run-python/src/main.ts
@@ -12,7 +12,7 @@ import { type LoggingLevel, SetLevelRequestSchema } from '@modelcontextprotocol/
 import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
 import { z } from 'zod'
 import { Buffer } from 'node:buffer'
-import { asXml, runCode, runCodeWithToolInjection, type ToolInjectionConfig } from './runCode.ts'
+import { asXml, runCode, type RunError, type RunSuccess, type ToolInjectionConfig } from './runCode.ts'
 
 const VERSION = '0.0.14'
 
@@ -50,6 +50,85 @@ options:
   }
 }
 
+/*
+ * Helper function to create a logging handler
+ */
+function createLogHandler(
+  setLogLevel: LoggingLevel,
+  logPromises: Promise<void>[],
+  server: McpServer,
+) {
+  return (level: LoggingLevel, data: string) => {
+    if (LogLevels.indexOf(level) >= LogLevels.indexOf(setLogLevel)) {
+      logPromises.push(server.server.sendLoggingMessage({ level, data }))
+    }
+  }
+}
+
+/*
+ * Helper function to build unified response
+ */
+async function buildResponse(
+  result: RunSuccess | RunError,
+  logPromises: Promise<void>[],
+) {
+  await Promise.all(logPromises)
+  return {
+    content: [{ type: 'text' as const, text: asXml(result) }],
+  }
+}
+
+/*
+ * Create elicitation callback for tool execution
+ */
+function createElicitationCallback(
+  server: McpServer,
+  logPromises: Promise<void>[],
+) {
+  // deno-lint-ignore no-explicit-any
+  return async (elicitationRequest: any) => {
+    // Convert Python dict to JavaScript object if needed
+    let jsRequest
+    if (elicitationRequest && typeof elicitationRequest === 'object' && elicitationRequest.toJs) {
+      jsRequest = elicitationRequest.toJs()
+    } else if (elicitationRequest && typeof elicitationRequest === 'object') {
+      // Handle Python dict-like objects
+      jsRequest = {
+        message: elicitationRequest.message || elicitationRequest.get?.('message'),
+        requestedSchema: elicitationRequest.requestedSchema || elicitationRequest.get?.('requestedSchema'),
+      }
+    } else {
+      jsRequest = elicitationRequest
+    }
+
+    try {
+      const elicitationResult = await server.server.request(
+        {
+          method: 'elicitation/create',
+          params: {
+            message: jsRequest.message,
+            requestedSchema: jsRequest.requestedSchema,
+          },
+        },
+        z.object({
+          action: z.enum(['accept', 'decline', 'cancel']),
+          content: z.optional(z.record(z.string(), z.unknown())),
+        }),
+      )
+
+      return elicitationResult
+    } catch (error) {
+      logPromises.push(
+        server.server.sendLoggingMessage({
+          level: 'error',
+          data: `Elicitation error: ${error}`,
+        }),
+      )
+      throw error
+    }
+  }
+}
+
 /*
  * Create an MCP server with the `run_python_code` tool registered.
  */
@@ -118,64 +197,18 @@ The tools are injected into the global namespace automatically - no discovery fu
 
       // Check if tools are provided
       if (tools.length > 0) {
-        // Create elicitation callback
-        // deno-lint-ignore no-explicit-any
-        const elicitationCallback = async (elicitationRequest: any) => {
-          // Convert Python dict to JavaScript object if needed
-          let jsRequest
-          if (elicitationRequest && typeof elicitationRequest === 'object' && elicitationRequest.toJs) {
-            jsRequest = elicitationRequest.toJs()
-          } else if (elicitationRequest && typeof elicitationRequest === 'object') {
-            // Handle Python dict-like objects
-            jsRequest = {
-              message: elicitationRequest.message || elicitationRequest.get?.('message'),
-              requestedSchema: elicitationRequest.requestedSchema || elicitationRequest.get?.('requestedSchema'),
-            }
-          } else {
-            jsRequest = elicitationRequest
-          }
-
-          try {
-            const elicitationResult = await server.server.request(
-              {
-                method: 'elicitation/create',
-                params: {
-                  message: jsRequest.message,
-                  requestedSchema: jsRequest.requestedSchema,
-                },
-              },
-              z.object({
-                action: z.enum(['accept', 'decline', 'cancel']),
-                content: z.optional(z.record(z.string(), z.unknown())),
-              }),
-            )
-
-            return elicitationResult
-          } catch (error) {
-            logPromises.push(
-              server.server.sendLoggingMessage({
-                level: 'error',
-                data: `Elicitation error: ${error}`,
-              }),
-            )
-            throw error
-          }
-        }
+        const elicitationCallback = createElicitationCallback(server, logPromises)
 
         // Use tool injection mode
-        const result = await runCodeWithToolInjection(
+        const result = await runCode(
           [
             {
               name: 'main.py',
               content: python_code,
               active: true,
             },
           ],
-          (level, data) => {
-            if (LogLevels.indexOf(level) >= LogLevels.indexOf(setLogLevel)) {
-              logPromises.push(server.server.sendLoggingMessage({ level, data }))
-            }
-          },
+          createLogHandler(setLogLevel, logPromises, server),
           {
             enableToolInjection: true,
             availableTools: tools,
@@ -184,11 +217,7 @@ The tools are injected into the global namespace automatically - no discovery fu
           } as ToolInjectionConfig,
         )
 
-        await Promise.all(logPromises)
-
-        return {
-          content: [{ type: 'text', text: asXml(result) }],
-        }
+        return await buildResponse(result, logPromises)
       } else {
         // Use basic mode without tool injection
         const result = await runCode(
@@ -199,16 +228,10 @@ The tools are injected into the global namespace automatically - no discovery fu
               active: true,
             },
           ],
-          (level, data) => {
-            if (LogLevels.indexOf(level) >= LogLevels.indexOf(setLogLevel)) {
-              logPromises.push(server.server.sendLoggingMessage({ level, data }))
-            }
-          },
+          createLogHandler(setLogLevel, logPromises, server),
+          undefined,
         )
-        await Promise.all(logPromises)
-        return {
-          content: [{ type: 'text', text: asXml(result) }],
-        }
+        return await buildResponse(result, logPromises)
       }
     },
   )
@@ -439,7 +462,7 @@ print(f"Tool result: {result}")
 `
 
   try {
-    const toolResult = await runCodeWithToolInjection(
+    const toolResult = await runCode(
       [
         {
           name: 'tool_test.py',
diff --git a/mcp-run-python/src/runCode.ts b/mcp-run-python/src/runCode.ts
@@ -21,18 +21,6 @@ export interface ToolInjectionConfig {
 export async function runCode(
   files: CodeFile[],
   log: (level: LoggingLevel, data: string) => void,
-): Promise<RunSuccess | RunError> {
-  // Use enhanced version without tool injection for backwards compatibility
-  const result = await runCodeWithToolInjection(files, log, undefined)
-
-  // Convert to original format
-  return result
-}
-
-// Enhanced version that supports optional tool injection
-export async function runCodeWithToolInjection(
-  files: CodeFile[],
-  log: (level: LoggingLevel, data: string) => void,
   toolConfig?: ToolInjectionConfig,
 ): Promise<RunSuccess | RunError> {
   // remove once we can upgrade to pyodide 0.27.7 and console.log is no longer used.
@@ -71,7 +59,11 @@ export async function runCodeWithToolInjection(
   const dirPath = '/tmp/mcp_run_python'
   sys.path.append(dirPath)
   const pathlib = pyodide.pyimport('pathlib')
-  pathlib.Path(dirPath).mkdir()
+  try {
+    pathlib.Path(dirPath).mkdir()
+  } catch (_error) {
+    // Directory already exists, which is fine
+  }
   const moduleName = '_prepare_env'
 
   pathlib.Path(`${dirPath}/${moduleName}.py`).write_text(preparePythonCode)
@@ -162,15 +154,15 @@ function injectToolFunctions(
   log('info', `Tool injection complete. Available tools: ${config.availableTools.join(', ')}`)
 }
 
-interface RunSuccess {
+export interface RunSuccess {
   status: 'success'
   // we could record stdout and stderr separately, but I suspect simplicity is more important
   output: string[]
   dependencies: string[]
   returnValueJson: string | null
 }
 
-interface RunError {
+export interface RunError {
   status: 'install-error' | 'run-error'
   output: string[]
   dependencies?: string[]
diff --git a/mcp-run-python/src/tool_injection.py b/mcp-run-python/src/tool_injection.py
@@ -1,48 +1,46 @@
 """Tool injection for MCP elicitation support."""
 
 import json
+from functools import lru_cache
 from typing import Any, Callable, cast
 
+from pyodide.webloop import run_sync  # type: ignore[import-untyped]
+
 
 def _create_elicitation_request(tool_name: str, args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any]:
     """Create an elicitation request object for tool execution."""
-    tool_args: dict[str, Any] = {}
+    if args and len(args) == 1:
+        first_arg = args[0]
+        if isinstance(first_arg, str):
+            tool_args = {'query': first_arg, **kwargs}
+        elif isinstance(first_arg, dict):
+            tool_args = {**first_arg, **kwargs}  # type: ignore[arg-type]
+        else:
+            tool_args = kwargs.copy()
+    else:
+        tool_args = kwargs.copy()
 
-    # Handle positional arguments
-    if args:
-        if len(args) == 1 and isinstance(args[0], str):
-            # Single string argument - assume it's a query
-            tool_args['query'] = args[0]
-        elif len(args) == 1 and isinstance(args[0], dict):
-            # Single dict argument
-            tool_args.update(args[0])  # type: ignore[arg-type]
+    return {
+        'message': json.dumps({'tool_name': tool_name, 'arguments': tool_args}),
+        'requestedSchema': _get_tool_schema(tool_name),
+    }
 
-    # Add keyword arguments
-    tool_args.update(kwargs)
 
-    tool_execution_data: dict[str, Any] = {'tool_name': tool_name, 'arguments': tool_args}
+@lru_cache(maxsize=128)
+def _get_tool_schema(tool_name: str) -> dict[str, Any]:
+    """Get cached schema for a tool."""
     return {
-        'message': json.dumps(tool_execution_data),
-        'requestedSchema': {
-            'type': 'object',
-            'properties': {'result': {'type': 'string', 'description': f'Result of executing {tool_name} tool'}},
-            'required': ['result'],
-        },
+        'type': 'object',
+        'properties': {'result': {'type': 'string', 'description': f'Result of executing {tool_name} tool'}},
+        'required': ['result'],
     }
 
 
 def _handle_tool_callback_result(result: Any, tool_name: str) -> Any:
     """Handle the result from a tool callback, including promise resolution."""
-    # Handle PyodideFuture (JavaScript Promise)
     if hasattr(result, 'then'):
         try:
-            # Import at runtime to avoid dependency issues
-            from pyodide.webloop import run_sync  # type: ignore[import-untyped]
-
-            # Use cast to tell the type checker what we expect
             resolved_result = cast(dict[str, Any], run_sync(result))
-
-            # Extract result from elicitation response
             if isinstance(resolved_result, dict):
                 if resolved_result.get('action') == 'accept':
                     content = cast(dict[str, Any], resolved_result.get('content', {}))
