Added MCP elicitation support for tool injection in mcp-run-python

mcp-run-python/build.ts

@@ -5,12 +5,20 @@ if (!import.meta.dirname) {
   throw new Error('import.meta.dirname is not defined, unable to load prepare_env.py')
 }
 const src = path.join(import.meta.dirname, 'src/prepare_env.py')
+const toolInjectionSrc = path.join(import.meta.dirname, 'src/tool_injection.py')
 const dst = path.join(import.meta.dirname, 'src/prepareEnvCode.ts')
 
 let pythonCode = await Deno.readTextFile(src)
 pythonCode = pythonCode.replace(/\\/g, '\\\\')
+
+// Read tool injection code from separate Python file
+let toolInjectionCode = await Deno.readTextFile(toolInjectionSrc)
+toolInjectionCode = toolInjectionCode.replace(/\\/g, '\\\\')
+
 const jsCode = `\
 // DO NOT EDIT THIS FILE DIRECTLY, INSTEAD RUN "deno run build"
 export const preparePythonCode = \`${pythonCode}\`
+
+export const toolInjectionCode = \`${toolInjectionCode}\`
 `
 await Deno.writeTextFile(dst, jsCode)

mcp-run-python/pyproject.toml

@@ -10,7 +10,7 @@ dev = [
     "dirty-equals>=0.9.0",
     "httpx>=0.28.1",
     "inline-snapshot>=0.19.3",
-    "mcp>=1.4.1; python_version >= '3.10'",
+    "mcp>=1.12.2; python_version >= '3.10'",
     "micropip>=0.9.0; python_version >= '3.12'",
     "pytest>=8.3.3",
     "pytest-pretty>=1.2.0",

mcp-run-python/src/main.ts

@@ -1,8 +1,8 @@
 /// <reference types="npm:@types/[email protected]" />
 
 import './polyfill.ts'
-import http from 'node:http'
 import { randomUUID } from 'node:crypto'
+import http, { type IncomingMessage, type ServerResponse } from 'node:http'
 import { parseArgs } from '@std/cli/parse-args'
 import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js'
 import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'
@@ -11,14 +11,13 @@ import { isInitializeRequest } from '@modelcontextprotocol/sdk/types.js'
 import { type LoggingLevel, SetLevelRequestSchema } from '@modelcontextprotocol/sdk/types.js'
 import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
 import { z } from 'zod'
-
-import { asXml, runCode } from './runCode.ts'
 import { Buffer } from 'node:buffer'
+import { asXml, runCode, type RunError, type RunSuccess, type ToolInjectionConfig } from './runCode.ts'
 
-const VERSION = '0.0.13'
+const VERSION = '0.0.14'
 
 export async function main() {
-  const { args } = Deno
+  const args = Deno?.args || []
   if (args.length === 1 && args[0] === 'stdio') {
     await runStdio()
   } else if (args.length >= 1 && args[0] === 'streamable_http') {
@@ -29,7 +28,7 @@ export async function main() {
     const port = parseInt(flags.port)
     runStreamableHttp(port)
   } else if (args.length >= 1 && args[0] === 'sse') {
-    const flags = parseArgs(Deno.args, {
+    const flags = parseArgs(args, {
       string: ['port'],
       default: { port: '3001' },
     })
@@ -47,7 +46,86 @@ Usage: deno run -N -R=node_modules -W=node_modules --node-modules-dir=auto jsr:@
 options:
   --port <port>  Port to run the SSE server on (default: 3001)`,
     )
-    Deno.exit(1)
+    Deno?.exit(1)
+  }
+}
+
+/*
+ * 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
+    }
   }
 }
 
@@ -64,6 +142,7 @@ function createServer(): McpServer {
       instructions: 'Call the "run_python_code" tool with the Python code to run.',
       capabilities: {
         logging: {},
+        elicitation: {},
       },
     },
   )
@@ -81,6 +160,13 @@ with a comment of the form:
 # dependencies = ['pydantic']
 # ///
 print('python code here')
+
+TOOL INJECTION: When 'tools' parameter is provided, the specified tool functions become available directly in Python's global namespace. You can call them directly like any other function. For example, if 'web_search' is provided as a tool, you can call it directly:
+
+result = web_search("search query")
+print(result)
+
+The tools are injected into the global namespace automatically - no discovery functions needed.
 `
 
   let setLogLevel: LoggingLevel = 'emergency'
@@ -93,21 +179,59 @@ print('python code here')
   server.tool(
     'run_python_code',
     toolDescription,
-    { python_code: z.string().describe('Python code to run') },
-    async ({ python_code }: { python_code: string }) => {
+    {
+      python_code: z.string().describe('Python code to run'),
+      tools: z
+        .array(z.string())
+        .optional()
+        .describe('List of available tools for injection (enables tool injection when provided)'),
+    },
+    async ({
+      python_code,
+      tools = [],
+    }: {
+      python_code: string
+      tools?: string[]
+    }) => {
       const logPromises: Promise<void>[] = []
-      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 }))
-        }
-      })
-      await Promise.all(logPromises)
-      return {
-        content: [{ type: 'text', text: asXml(result) }],
+
+      // Check if tools are provided
+      if (tools.length > 0) {
+        const elicitationCallback = createElicitationCallback(server, logPromises)
+
+        // Use tool injection mode
+        const result = await runCode(
+          [
+            {
+              name: 'main.py',
+              content: python_code,
+              active: true,
+            },
+          ],
+          createLogHandler(setLogLevel, logPromises, server),
+          {
+            enableToolInjection: true,
+            availableTools: tools,
+            timeoutSeconds: 30,
+            elicitationCallback,
+          } as ToolInjectionConfig,
+        )
+
+        return await buildResponse(result, logPromises)
+      } else {
+        // Use basic mode without tool injection
+        const result = await runCode(
+          [
+            {
+              name: 'main.py',
+              content: python_code,
+              active: true,
+            },
+          ],
+          createLogHandler(setLogLevel, logPromises, server),
+          undefined,
+        )
+        return await buildResponse(result, logPromises)
       }
     },
   )
@@ -167,7 +291,7 @@ function runStreamableHttp(port: number) {
   const mcpServer = createServer()
   const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {}
 
-  const server = http.createServer(async (req, res) => {
+  const server = http.createServer(async (req: IncomingMessage, res: ServerResponse) => {
     const url = httpGetUrl(req)
     let pathMatch = false
     function match(method: string, path: string): boolean {
@@ -309,21 +433,66 @@ async function warmup() {
   console.error(
     `Running warmup script for MCP Run Python version ${VERSION}...`,
   )
+
   const code = `
 import numpy
 a = numpy.array([1, 2, 3])
 print('numpy array:', a)
 a
 `
-  const result = await runCode([{
-    name: 'warmup.py',
-    content: code,
-    active: true,
-  }], (level, data) =>
-    // use warn to avoid recursion since console.log is patched in runCode
-    console.error(`${level}: ${data}`))
-  console.log('Tool return value:')
+  const result = await runCode(
+    [
+      {
+        name: 'warmup.py',
+        content: code,
+        active: true,
+      },
+    ],
+    (level, data) => console.error(`${level}: ${data}`),
+  )
   console.log(asXml(result))
+
+  // Test tool injection functionality
+  console.error('Testing tool injection framework...')
+  const toolCode = `
+# Test tool injection - directly call an injected tool
+result = web_search("test query")
+print(f"Tool result: {result}")
+"tool_test_complete"
+`
+
+  try {
+    const toolResult = await runCode(
+      [
+        {
+          name: 'tool_test.py',
+          content: toolCode,
+          active: true,
+        },
+      ],
+      (level, data) => console.error(`${level}: ${data}`),
+      {
+        enableToolInjection: true,
+        availableTools: ['web_search', 'send_email'],
+        timeoutSeconds: 30,
+        // deno-lint-ignore no-explicit-any require-await
+        elicitationCallback: async (_elicitationRequest: any) => {
+          // Mock callback for warmup test
+          return {
+            action: 'accept',
+            content: {
+              result: '{"status": "mock success"}',
+            },
+          }
+        },
+      } as ToolInjectionConfig,
+    )
+    console.log('Tool injection result:')
+    console.log(asXml(toolResult))
+  } catch (error) {
+    console.error('Tool injection test failed:', error)
+  }
+
   console.log('\nwarmup successful 🎉')
 }