diff --git a/build/build.py b/build/build.py index 55632c7..c8b3189 100644 --- a/build/build.py +++ b/build/build.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -"""Build script to inline ./prepare_env.py into mcp_run_python/deno/prepareEnvCode.ts""" +"""Build script to inline ./prepare_env.py and tool_injection.py into mcp_run_python/deno/prepareEnvCode.ts""" from pathlib import Path @@ -7,17 +7,23 @@ root_dir = this_dir.parent # Define source and destination paths -src = this_dir / 'prepare_env.py' -dst = root_dir / 'mcp_run_python' / 'deno' / 'src' / 'prepareEnvCode.ts' +src_dir = root_dir / 'mcp_run_python' / 'deno' / 'src' +prepare_env_src = this_dir / 'prepare_env.py' +tool_injection_src = src_dir / 'tool_injection.py' +dst = src_dir / 'prepareEnvCode.ts' -python_code = src.read_text() +prepare_env_code = prepare_env_src.read_text() +tool_injection_code = tool_injection_src.read_text() # Escape backslashes for JavaScript string literal -python_code = python_code.replace('\\', '\\\\') +prepare_env_code = prepare_env_code.replace('\\', '\\\\') +tool_injection_code = tool_injection_code.replace('\\', '\\\\') # Create the JavaScript/TypeScript code ts_code = f"""// DO NOT EDIT THIS FILE DIRECTLY, INSTEAD RUN "make build" -export const preparePythonCode = `{python_code}` +export const preparePythonCode = `{prepare_env_code}` + +export const toolInjectionCode = `{tool_injection_code}` """ dst.write_text(ts_code) diff --git a/mcp_run_python/deno/src/main.ts b/mcp_run_python/deno/src/main.ts index d00d549..06360d5 100644 --- a/mcp_run_python/deno/src/main.ts +++ b/mcp_run_python/deno/src/main.ts @@ -17,6 +17,126 @@ import { Buffer } from 'node:buffer' const VERSION = '0.0.13' +// Tool injection support +const EMPTY_TOOLS_RESULT = { toolNames: [], toolSchemas: {} } + +const ELICIT_RESPONSE_SCHEMA = z.object({ + action: z.enum(['accept', 'decline', 'cancel']), + content: z.optional(z.record(z.string(), z.unknown())), +}) + +// Schema for tool discovery responses +const TOOL_DISCOVERY_RESPONSE_SCHEMA = { + type: 'object', + properties: { + action: { + type: 'string', + enum: ['accept', 'decline', 'cancel'], + description: 'Action to take with the tool discovery request', + }, + content: { + type: 'object', + properties: { + data: { + type: 'string', + description: 'JSON string containing tool_names and tool_schemas', + }, + }, + description: 'Tool discovery result data', + }, + }, + required: ['action'], +} + +// Schema for tool execution responses +const TOOL_EXECUTION_RESPONSE_SCHEMA = { + type: 'object', + properties: { + action: { + type: 'string', + enum: ['accept', 'decline', 'cancel'], + description: 'Action to take with the tool execution request', + }, + content: { + type: 'object', + properties: { + result: { + type: 'string', + description: 'JSON string containing tool execution result', + }, + error: { + type: 'string', + description: 'Error message if execution failed', + }, + retry: { + type: 'string', + description: 'JSON string containing retry information if tool needs retry', + }, + }, + description: 'Tool execution result or error information', + }, + }, + required: ['action'], +} + +function makeElicitationRequest( + server: McpServer, + message: string, + requestedSchema: Record, +): Promise> { + return server.server.request( + { + method: 'elicitation/create', + params: { message, requestedSchema }, + }, + ELICIT_RESPONSE_SCHEMA, + ) +} + +async function discoverAvailableTools( + server: McpServer, +): Promise<{ toolNames: string[]; toolSchemas: Record> }> { + try { + const result = await makeElicitationRequest( + server, + JSON.stringify({ action: 'discover_tools' }), + TOOL_DISCOVERY_RESPONSE_SCHEMA, + ) + + if (result.action === 'accept' && result.content?.data) { + const data = JSON.parse(result.content.data as string) + return { + toolNames: data.tool_names || [], + toolSchemas: data.tool_schemas || {}, + } + } + } catch (error) { + console.warn('Tool discovery failed:', error) + } + + return EMPTY_TOOLS_RESULT +} + +function createToolExecutionCallback( + server: McpServer, +): (message: string) => Promise> { + return async (message: string) => { + try { + return await makeElicitationRequest( + server, + message, + TOOL_EXECUTION_RESPONSE_SCHEMA, + ) + } catch (error) { + console.error('Tool execution failed:', error) + return { + action: 'decline', + content: { error: `Tool execution failed: ${error}` }, + } + } + } +} + export async function main() { const { args } = Deno const flags = parseArgs(Deno.args, { @@ -86,6 +206,38 @@ The code will be executed with Python 3.13. return {} }) + server.registerTool( + 'discover_available_tools', + { + title: 'Discover available tools', + description: + 'Discover what tools are available for injection into Python code execution environment. Call this before writing Python code to know which tools you can use.', + inputSchema: {}, + }, + async () => { + // Discover available tools via elicitation + const { toolNames: availableTools, toolSchemas } = await discoverAvailableTools(server) + + return { + content: [ + { + type: 'text', + text: JSON.stringify( + { + available_tools: availableTools, + tool_schemas: toolSchemas, + usage_note: + "These tools will be available as Python functions when you run Python code. Tool names with hyphens become underscores (e.g., 'web-search' -> 'web_search').", + }, + null, + 2, + ), + }, + ], + } + }, + ) + server.registerTool( 'run_python_code', { @@ -100,6 +252,20 @@ The code will be executed with Python 3.13. }, async ({ python_code, global_variables }: { python_code: string; global_variables: Record }) => { const logPromises: Promise[] = [] + + // Discover available tools via elicitation + const { toolNames: availableTools, toolSchemas } = await discoverAvailableTools(server) + + // Create elicitation callback for tool execution if tools are available + let toolInjectionOptions + if (availableTools.length > 0) { + toolInjectionOptions = { + elicitationCallback: createToolExecutionCallback(server), + availableTools: availableTools, + toolSchemas: toolSchemas, + } + } + const result = await runCode.run( deps, (level, data) => { @@ -110,6 +276,7 @@ The code will be executed with Python 3.13. { name: 'main.py', content: python_code }, global_variables, returnMode !== 'xml', + toolInjectionOptions, ) await Promise.all(logPromises) return { diff --git a/mcp_run_python/deno/src/runCode.ts b/mcp_run_python/deno/src/runCode.ts index 1be8681..4dcd38a 100644 --- a/mcp_run_python/deno/src/runCode.ts +++ b/mcp_run_python/deno/src/runCode.ts @@ -1,6 +1,6 @@ // deno-lint-ignore-file no-explicit-any import { loadPyodide, type PyodideInterface } from 'pyodide' -import { preparePythonCode } from './prepareEnvCode.ts' +import { preparePythonCode, toolInjectionCode } from './prepareEnvCode.ts' import type { LoggingLevel } from '@modelcontextprotocol/sdk/types.js' export interface CodeFile { @@ -8,6 +8,17 @@ export interface CodeFile { content: string } +interface ElicitationResponse { + action: 'accept' | 'decline' | 'cancel' + content?: Record +} + +export interface ToolInjectionOptions { + elicitationCallback?: (message: string) => Promise + availableTools?: string[] + toolSchemas?: Record> +} + interface PrepResult { pyodide: PyodideInterface preparePyEnv: PreparePyEnv @@ -27,6 +38,7 @@ export class RunCode { file?: CodeFile, globals?: Record, alwaysReturnJson: boolean = false, + toolInjectionOptions?: ToolInjectionOptions, ): Promise { let pyodide: PyodideInterface let sys: any @@ -56,8 +68,50 @@ export class RunCode { } } else if (file) { try { + // Create globals for python execution, starting with user's global_variables + const globalVars = { ...(globals || {}), __name__: '__main__' } + const globalsProxy = pyodide.toPy(globalVars) + + // Setup tool injection if elicitation callback is provided + if (toolInjectionOptions?.elicitationCallback && toolInjectionOptions?.availableTools?.length) { + const dirPath = '/tmp/mcp_run_python' + const pathlib = pyodide.pyimport('pathlib') + const toolModuleName = '_tool_injection' + pathlib + .Path(`${dirPath}/${toolModuleName}.py`) + .write_text(toolInjectionCode) + + const toolInjectionModule = pyodide.pyimport(toolModuleName) + + // Create Javascript callback wrapper that handles promises + const jsElicitationCallback = async (message: string): Promise => { + try { + const result = await toolInjectionOptions.elicitationCallback!(message) + return result + } catch (error) { + log('error', `Elicitation callback error: ${error}`) + + return { + action: 'decline', + content: { error: `Elicitation failed: ${error}` }, + } + } + } + + // Convert to Python and inject tools + const pyCallback = pyodide.toPy(jsElicitationCallback) + const pyTools = pyodide.toPy(toolInjectionOptions.availableTools) + const pyToolSchemas = pyodide.toPy(toolInjectionOptions.toolSchemas || {}) + toolInjectionModule.inject_tool_functions(globalsProxy, pyTools, pyCallback, pyToolSchemas) + + log( + 'info', + `Tool injection enabled for: ${toolInjectionOptions.availableTools.join(', ')}`, + ) + } + const rawValue = await pyodide.runPythonAsync(file.content, { - globals: pyodide.toPy({ ...(globals || {}), __name__: '__main__' }), + globals: globalsProxy, filename: file.name, }) return { diff --git a/mcp_run_python/deno/src/tool_injection.py b/mcp_run_python/deno/src/tool_injection.py new file mode 100644 index 0000000..29b5b5d --- /dev/null +++ b/mcp_run_python/deno/src/tool_injection.py @@ -0,0 +1,120 @@ +import json +import uuid +from collections.abc import Callable +from typing import Any, cast + +from pyodide.webloop import run_sync # type: ignore[import-untyped] + + +def _extract_tool_result(elicit_result: Any, tool_name: str) -> Any: + # Convert JsProxy to Python dict if needed + if hasattr(elicit_result, 'to_py'): + elicit_result = elicit_result.to_py() + + # elicit_result is now a dict with 'action' and 'content' fields + action = elicit_result['action'] + content = elicit_result['content'] + + if action == 'accept': + # Content is always a dict with 'result' key containing JSON string + return json.loads(content['result']) + elif action == 'decline': + error_msg = content.get('error', 'Unknown error') + raise RuntimeError(f"Tool execution failed for '{tool_name}': {error_msg}") + else: + raise RuntimeError(f"MCP protocol error for tool '{tool_name}': unexpected action '{action}'") + + +def _handle_elicitation_result(result: Any, tool_name: str) -> Any: + if hasattr(result, 'then'): + try: + resolved_result = cast(dict[str, Any], run_sync(result)) + return _extract_tool_result(resolved_result, tool_name) + except Exception as e: + raise RuntimeError(f"Error resolving async result for tool '{tool_name}': {e}") from e + else: + return _extract_tool_result(result, tool_name) + + +def _create_tool_call_part( + tool_name: str, args: tuple[Any, ...], kwargs: dict[str, Any], tool_schema: dict[str, Any] +) -> dict[str, Any]: + # Schema-aware parameter mapping - tool_schema is required + if 'properties' not in tool_schema: + raise ValueError(f"Tool '{tool_name}' schema missing 'properties' field") + + properties = tool_schema['properties'] + + # Handle different argument patterns + if not args: + tool_args = kwargs.copy() + elif len(args) == 1 and properties: + # Single positional arg - map to first parameter or merge if dict + first_param_name = next(iter(properties)) + first_arg = args[0] + + if isinstance(first_arg, dict): + tool_args: dict[str, Any] = {**first_arg, **kwargs} + else: + tool_args = {first_param_name: first_arg, **kwargs} + else: + # Multiple positional args - map to parameters in order + param_names = list(properties.keys()) + if len(args) > len(param_names): + raise ValueError( + f"Tool '{tool_name}' received {len(args)} positional args but only has {len(param_names)} parameters" + ) + + tool_args = dict(zip(param_names, args)) + tool_args.update(kwargs) + + return { + 'tool_name': tool_name, + 'tool_call_id': str(uuid.uuid4()), + 'args': tool_args, + } + + +def _create_tool_function( + tool_name: str, elicitation_callback: Callable[[str], Any], tool_schema: dict[str, Any] +) -> Callable[..., Any]: + def tool_function(*args: Any, **kwargs: Any) -> Any: + # Create tool call with schema-aware parameter mapping + tool_call_data = _create_tool_call_part(tool_name, args, kwargs, tool_schema) + + elicitation_message = json.dumps(tool_call_data) + + try: + result = elicitation_callback(elicitation_message) + return _handle_elicitation_result(result, tool_name) + except Exception as e: + raise RuntimeError(f"Error calling tool '{tool_name}': {e}") from e + + return tool_function + + +def inject_tool_functions( + globals_dict: dict[str, Any], + available_tools: list[str], + elicitation_callback: Callable[[str], Any] | None = None, + tool_schemas: dict[str, dict[str, Any]] | None = None, +) -> None: + if not available_tools or elicitation_callback is None: + return + + if not tool_schemas: + raise ValueError('tool_schemas is required for tool injection') + + for tool_name in available_tools: + python_name = tool_name.replace('-', '_') + if python_name in globals_dict: + continue + + tool_schema = tool_schemas.get(tool_name) + + if not tool_schema: + raise ValueError(f"Schema missing for tool '{tool_name}' - cannot inject without schema") + + globals_dict[python_name] = _create_tool_function( + tool_name=tool_name, elicitation_callback=elicitation_callback, tool_schema=tool_schema + ) diff --git a/mcp_run_python/ext/__init__.py b/mcp_run_python/ext/__init__.py new file mode 100644 index 0000000..df06296 --- /dev/null +++ b/mcp_run_python/ext/__init__.py @@ -0,0 +1,8 @@ +"""Extensions for mcp-run-python. + +This module provides extensions for integrating mcp-run-python with various AI agent frameworks. +""" + +from .pydantic_ai import create_tool_elicitation_callback + +__all__ = ['create_tool_elicitation_callback'] diff --git a/mcp_run_python/ext/pydantic_ai.py b/mcp_run_python/ext/pydantic_ai.py new file mode 100644 index 0000000..57053b2 --- /dev/null +++ b/mcp_run_python/ext/pydantic_ai.py @@ -0,0 +1,130 @@ +"""PydanticAI integration for mcp-run-python tool injection. + +This module provides functionality to inject MCP tools from PydanticAI toolsets +into Python code execution environments via elicitation callbacks. +""" + +import json +from typing import Any + +from mcp.client.session import ClientSession, ElicitationFnT +from mcp.shared.context import RequestContext +from mcp.types import ElicitRequestParams, ElicitResult +from pydantic import ValidationError +from pydantic.type_adapter import TypeAdapter +from pydantic_ai._run_context import AgentDepsT, RunContext +from pydantic_ai._tool_manager import ToolManager +from pydantic_ai.exceptions import ModelRetry, ToolRetryError +from pydantic_ai.messages import ToolCallPart +from pydantic_ai.models.test import TestModel +from pydantic_ai.toolsets.abstract import AbstractToolset +from pydantic_ai.usage import RunUsage + +__all__ = ['create_tool_elicitation_callback'] + +# Protocol constants for tool discovery +TOOL_DISCOVERY_ACTION = 'discover_tools' + + +def create_tool_elicitation_callback(toolset: AbstractToolset[Any], deps: AgentDepsT | None = None) -> ElicitationFnT: + """Create an elicitation callback for tool injection into Python execution. + + This function creates a callback that handles tool discovery and execution + requests from mcp-run-python + + Args: + toolset: The PydanticAI toolset containing tools to make available + deps: Optional dependencies for the tool execution context + + Returns: + An elicitation callback function that handles tool discovery and execution + + Example: + ```python + from mcp_run_python.ext import create_tool_elicitation_callback + from pydantic_ai import Agent + from pydantic_ai.mcp import MCPServerStdio + from pydantic_ai.toolsets import CombinedToolset + + # Create combined toolset with your tools + my_toolset = CombinedToolset([search_toolset, email_toolset]) + + # Create elicitation callback + callback = create_tool_elicitation_callback(toolset=my_toolset) + + # Use with mcp-run-python + agent = Agent( + toolsets=[ + MCPServerStdio( + command='mcp-run-python', + elicitation_callback=callback + ), + elicitation_callback=my_toolset, # Also available directly to agent + ] + ) + ``` + """ + tool_call_adapter = TypeAdapter(ToolCallPart) + + # Create shared run context for both tool discovery and execution + shared_run_context = RunContext[Any]( + deps=deps, + model=TestModel(), + usage=RunUsage(), + ) + + async def elicitation_callback( + context: RequestContext[ClientSession, Any, Any], # pyright: ignore[reportUnusedParameter] + params: ElicitRequestParams, + ) -> ElicitResult: + """Handle elicitation requests for tool discovery and execution.""" + try: + # Parse the elicitation message + try: + data = json.loads(params.message) + except json.JSONDecodeError as e: + return ElicitResult(action='decline', content={'error': f'Invalid JSON: {e}'}) + + # Handle tool discovery requests + if data.get('action') == TOOL_DISCOVERY_ACTION: + try: + # Use shared context for tool discovery + available_tools = await toolset.get_tools(shared_run_context) + tool_names = list(available_tools.keys()) + tool_schemas = { + tool_name: toolset_tool.tool_def.parameters_json_schema + for tool_name, toolset_tool in available_tools.items() + } + discovery_result = {'tool_names': tool_names, 'tool_schemas': tool_schemas} + return ElicitResult(action='accept', content={'data': json.dumps(discovery_result)}) + except Exception as e: + return ElicitResult(action='decline', content={'error': f'Tool discovery failed: {e}'}) + + # Handle regular tool execution requests + try: + tool_call = tool_call_adapter.validate_python(data) + except ValidationError as e: + return ElicitResult(action='decline', content={'error': f'Invalid tool call: {e}'}) + + # Execute the tool + tool_manager = await ToolManager(toolset=toolset).for_run_step(ctx=shared_run_context) + result = await tool_manager.handle_call(call=tool_call) + + # Return result as JSON + return ElicitResult(action='accept', content={'result': json.dumps(result)}) + + except ToolRetryError as e: + # Handle tool retry requests with structured information + retry_info = { + 'error': 'Tool retry needed', + 'tool_name': e.tool_retry.tool_name, + 'message': e.tool_retry.content, + 'tool_call_id': e.tool_retry.tool_call_id, + } + return ElicitResult(action='decline', content={'retry': json.dumps(retry_info)}) + except ModelRetry as e: + return ElicitResult(action='decline', content={'error': f'Model retry failed: {e}'}) + except Exception as e: + return ElicitResult(action='decline', content={'error': f'Unexpected error: {e}'}) + + return elicitation_callback diff --git a/tests/test_mcp_servers.py b/tests/test_mcp_servers.py index 889b903..6485402 100644 --- a/tests/test_mcp_servers.py +++ b/tests/test_mcp_servers.py @@ -76,12 +76,24 @@ async def test_list_tools(run_mcp_session: Callable[[list[str]], AbstractAsyncCo async with run_mcp_session([]) as mcp_session: await mcp_session.initialize() tools = await mcp_session.list_tools() - assert len(tools.tools) == 1 - tool = tools.tools[0] - assert tool.name == 'run_python_code' - assert tool.description - assert tool.description.startswith('Tool to execute Python code and return stdout, stderr, and return value.') - assert tool.inputSchema == snapshot( + assert len(tools.tools) == 2 + + # Find tools by name since order may vary + tool_by_name = {t.name: t for t in tools.tools} + + # Check discover_available_tools tool + discover_tool = tool_by_name['discover_available_tools'] + assert discover_tool.description + assert 'Discover what tools are available' in discover_tool.description + assert discover_tool.inputSchema['properties'] == {} + + # Check run_python_code tool + run_tool = tool_by_name['run_python_code'] + assert run_tool.description + assert run_tool.description.startswith( + 'Tool to execute Python code and return stdout, stderr, and return value.' + ) + assert run_tool.inputSchema == snapshot( { 'type': 'object', 'properties': { @@ -226,3 +238,31 @@ def logging_callback(level: str, message: str) -> None: \ """ ) + + +async def test_discover_available_tools( + run_mcp_session: Callable[[list[str]], AbstractAsyncContextManager[ClientSession]], +) -> None: + """Test the discover_available_tools functionality.""" + async with run_mcp_session([]) as mcp_session: + await mcp_session.initialize() + result = await mcp_session.call_tool('discover_available_tools', {}) + + assert len(result.content) == 1 + content = result.content[0] + assert isinstance(content, types.TextContent) + + # Parse the response to verify structure + import json + + response_data = json.loads(content.text) + + assert 'available_tools' in response_data + assert 'tool_schemas' in response_data + assert 'usage_note' in response_data + assert isinstance(response_data['available_tools'], list) + assert isinstance(response_data['tool_schemas'], dict) + + # Should initially return empty tools since no elicitation callback is set up + assert response_data['available_tools'] == [] + assert response_data['tool_schemas'] == {}