Sandbox (#2)

Author: samuelcolvin

README.md

@@ -47,7 +47,7 @@ where:
 - `example` will run a minimal Python script using `numpy`, useful for checking that the package is working, for the code
   to run successfully, you'll need to install `numpy` using `uvx mcp-run-python --deps numpy example`
 
-## Usage in codes
+## Usage in codes as an MCP server
 
 ```bash
 pip install mcp-run-python
@@ -87,14 +87,40 @@ if __name__ == '__main__':
 As well as returning the args needed to run `mcp_run_python`, `deno_args_prepare` installs the dependencies
 so they can be used by the server.
 
+## Usage in code with `code_sandbox`
+
+`mcp-run-python` includes a helper function `code_sandbox` to allow you to easily run code in a sandbox.
+
+```py
+from mcp_run_python import code_sandbox
+
+code = """
+import numpy
+a = numpy.array([1, 2, 3])
+print(a)
+a
+"""
+
+async def main():
+    async with code_sandbox(dependencies=['numpy']) as sandbox:
+        result = await sandbox.eval(code)
+        print(result)
+
+
+if __name__ == '__main__':
+    import asyncio
+
+    asyncio.run(main())
+```
+
+Under the hood, `code_sandbox` runs an MCP server using `stdio`. You can run multiple code blocks with a single sandbox.
+
 ## Logging
 
 MCP Run Python supports emitting stdout and stderr from the python execution as [MCP logging messages](https://github.com/modelcontextprotocol/specification/blob/eb4abdf2bb91e0d5afd94510741eadd416982350/docs/specification/draft/server/utilities/logging.md?plain=1).
 
 For logs to be emitted you must set the logging level when connecting to the server. By default, the log level is set to the highest level, `emergency`.
 
-Currently, it's not possible to demonstrate this due to a bug in the Python MCP Client, see [modelcontextprotocol/python-sdk#201](https://github.com/modelcontextprotocol/python-sdk/issues/201#issuecomment-2727663121).
-
 ## Dependencies
 
 `mcp_run_python` uses a two step process to install dependencies while avoiding any risk that sandboxed code can
@@ -104,54 +130,3 @@ edit the filesystem.
 * `deno` is then run with read-only permissions to the `node_modules` directory to run untrusted code.
 
 Dependencies must be provided when initializing the server so they can be installed in the first step.
-
-Here's an example of manually running code with `mcp-run-python`:
-
-```python
-from mcp import ClientSession, StdioServerParameters
-from mcp.client.stdio import stdio_client
-
-from mcp_run_python import deno_args_prepare
-
-code = """
-import numpy
-a = numpy.array([1, 2, 3])
-print(a)
-a
-"""
-server_params = StdioServerParameters(
-    command='deno',
-    args=deno_args_prepare('stdio', deps=['numpy'])
-)
-
-
-async def main():
-    async with stdio_client(server_params) as (read, write):
-        async with ClientSession(read, write) as session:
-            await session.initialize()
-            tools = await session.list_tools()
-            print(len(tools.tools))
-            #> 1
-            print(repr(tools.tools[0].name))
-            #> 'run_python_code'
-            print(repr(tools.tools[0].inputSchema))
-            """
-            {'type': 'object', 'properties': {'python_code': {'type': 'string', 'description': 'Python code to run'}}, 'required': ['python_code'], 'additionalProperties': False, '$schema': 'http://json-schema.org/draft-07/schema#'}
-            """
-            result = await session.call_tool('run_python_code', {'python_code': code})
-            print(result.content[0].text)
-            """
-            <status>success</status>
-            <dependencies>["numpy"]</dependencies>
-            <output>
-            [1 2 3]
-            </output>
-            <return_value>
-            [
-              1,
-              2,
-              3
-            ]
-            </return_value>
-            """
-```

examples/direct.py

@@ -21,7 +21,9 @@ async def main():
             print(repr(tools.tools[0].name))
             print(repr(tools.tools[0].inputSchema))
             result = await session.call_tool('run_python_code', {'python_code': code})
-            print(result.content[0].text)
+            content_block = result.content[0]
+            assert content_block.type == 'text'
+            print(content_block.text)
 
 
 if __name__ == '__main__':

examples/pydantic_ai_ex.py

@@ -1,24 +1,8 @@
-# /// script
-# requires-python = ">=3.13"
-# dependencies = [
-#     "logfire",
-#     "mcp-run-python",
-#     "pydantic-ai-slim[mcp,anthropic]",
-# ]
-#
-# [tool.uv.sources]
-# mcp-run-python = { path = "." }
-# ///
-import logfire
 from pydantic_ai import Agent
 from pydantic_ai.mcp import MCPServerStdio
 
 from mcp_run_python import deno_args_prepare
 
-logfire.configure()
-logfire.instrument_mcp()
-logfire.instrument_pydantic_ai()
-
 server = MCPServerStdio('deno', args=deno_args_prepare('stdio'))
 agent_with_python = Agent('claude-3-5-haiku-latest', toolsets=[server])
 

examples/sandbox.py

@@ -0,0 +1,19 @@
+from mcp_run_python import code_sandbox
+
+
+async def print_wait(level: str, message: str):
+    await asyncio.sleep(1)
+    print(f'[{level}] {message}')
+
+
+async def main():
+    async with code_sandbox(dependencies=['numpy'], print_handler=print_wait) as sandbox:
+        for i in range(10):
+            result = await sandbox.eval(f'import numpy as np\na = np.array([1, 2, {i}])\nprint(a)\na')
+            print(result)
+
+
+if __name__ == '__main__':
+    import asyncio
+
+    asyncio.run(main())

mcp_run_python/__init__.py

@@ -2,7 +2,8 @@
 
 from importlib.metadata import version as _metadata_version
 
+from .code_sandbox import code_sandbox
 from .main import deno_args_prepare, deno_run_server
 
 __version__ = _metadata_version('mcp_run_python')
-__all__ = '__version__', 'deno_args_prepare', 'deno_run_server'
+__all__ = '__version__', 'deno_args_prepare', 'deno_run_server', 'code_sandbox'

mcp_run_python/_cli.py

@@ -30,6 +30,6 @@ def cli(args_list: Sequence[str] | None = None):  # pragma: no cover
         print(f'mcp-run-python {__version__}')
     elif args.mode:
         deps: list[str] = args.deps.split(',') if args.deps else []
-        deno_run_server(args.mode.replace('-', '_'), port=args.port, deps=deps, install_log_handler=print)
+        deno_run_server(args.mode.replace('-', '_'), port=args.port, deps=deps, prep_log_handler=print)
     else:
         parser.error('Mode is required')

mcp_run_python/code_sandbox.py

@@ -0,0 +1,76 @@
+import inspect
+import json
+from collections.abc import AsyncIterator, Awaitable, Callable
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import Literal, TypeAlias, TypedDict
+
+from mcp import ClientSession, StdioServerParameters, types as mcp_types
+from mcp.client.stdio import stdio_client
+
+from .main import deno_args_prepare
+
+JsonData: TypeAlias = 'str| bool | int | float | None | list[JsonData] | dict[str, JsonData]'
+
+
+class RunSuccess(TypedDict):
+    status: Literal['success']
+    output: list[str]
+    returnValueJson: JsonData
+
+
+class RunError(TypedDict):
+    status: Literal['install-error', 'run-error']
+    output: list[str]
+    error: str
+
+
+@dataclass
+class CodeSandbox:
+    _session: ClientSession
+
+    async def eval(self, code: str) -> RunSuccess | RunError:
+        result = await self._session.call_tool('run_python_code', {'python_code': code})
+        content_block = result.content[0]
+        if content_block.type == 'text':
+            return json.loads(content_block.text)
+        else:
+            raise ValueError(f'Unexpected content type: {content_block.type}')
+
+
+@asynccontextmanager
+async def code_sandbox(
+    *,
+    dependencies: list[str] | None = None,
+    print_handler: Callable[[mcp_types.LoggingLevel, str], None | Awaitable[None]] | None = None,
+    logging_level: mcp_types.LoggingLevel | None = None,
+    prep_log_handler: Callable[[str], None] | None = None,
+) -> AsyncIterator['CodeSandbox']:
+    """Run code in a secure sandbox.
+
+    Args:
+        dependencies: A list of dependencies to be installed.
+        print_handler: A callback function to handle print statements when code is running.
+        logging_level: The logging level to use for the print handler, defaults to `info` if `print_handler` is provided.
+        prep_log_handler: A callback function to run on log statements during initial install of dependencies.
+    """
+    args = deno_args_prepare('stdio', deps=dependencies, prep_log_handler=prep_log_handler, return_mode='json')
+    server_params = StdioServerParameters(command='deno', args=args)
+
+    logging_callback: Callable[[mcp_types.LoggingMessageNotificationParams], Awaitable[None]] | None = None
+
+    if print_handler:
+
+        async def logging_callback_(params: mcp_types.LoggingMessageNotificationParams) -> None:
+            if inspect.iscoroutinefunction(print_handler):
+                await print_handler(params.level, params.data)
+            else:
+                print_handler(params.level, params.data)
+
+        logging_callback = logging_callback_
+
+    async with stdio_client(server_params) as (read, write):
+        async with ClientSession(read, write, logging_callback=logging_callback) as session:
+            if print_handler:
+                await session.set_logging_level(logging_level or 'info')
+            yield CodeSandbox(session)

mcp_run_python/deno/main.ts

@@ -11,28 +11,25 @@ import { type LoggingLevel, SetLevelRequestSchema } from '@modelcontextprotocol/
 import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
 import { z } from 'zod'
 
-import { asXml, runCode } from './runCode.ts'
+import { asJson, asXml, runCode } from './runCode.ts'
 import { Buffer } from 'node:buffer'
 
 const VERSION = '0.0.13'
 
 export async function main() {
   const { args } = Deno
   const flags = parseArgs(Deno.args, {
-    string: ['deps'],
+    string: ['deps', 'return-mode', 'port'],
+    default: { port: '3001', 'return-mode': 'xml' },
   })
   const deps = flags.deps?.split(',') ?? []
   if (args.length >= 1) {
     if (args[0] === 'stdio') {
-      await runStdio(deps)
+      await runStdio(deps, flags['return-mode'])
       return
     } else if (args[0] === 'streamable_http') {
-      const flags = parseArgs(Deno.args, {
-        string: ['port'],
-        default: { port: '3001' },
-      })
       const port = parseInt(flags.port)
-      runStreamableHttp(port, deps)
+      runStreamableHttp(port, deps, flags['return-mode'])
       return
     } else if (args[0] === 'example') {
       await example(deps)
@@ -49,16 +46,17 @@ Invalid arguments: ${args.join(' ')}
 Usage: deno ... deno/main.ts [stdio|streamable_http|install_deps|noop]
 
 options:
---port <port>  Port to run the HTTP server on (default: 3001)
---deps <deps>  Comma separated list of dependencies to install`,
+--port <port>             Port to run the HTTP server on (default: 3001)
+--deps <deps>             Comma separated list of dependencies to install
+--return-mode <xml/json>  Return mode for output data (default: xml)`,
   )
   Deno.exit(1)
 }
 
 /*
  * Create an MCP server with the `run_python_code` tool registered.
  */
-function createServer(deps: string[]): McpServer {
+function createServer(deps: string[], returnMode: string): McpServer {
   const server = new McpServer(
     {
       name: 'MCP Run Python',
@@ -103,7 +101,7 @@ The code will be executed with Python 3.12.
       )
       await Promise.all(logPromises)
       return {
-        content: [{ type: 'text', text: asXml(result) }],
+        content: [{ type: 'text', text: returnMode === 'xml' ? asXml(result) : asJson(result) }],
       }
     },
   )
@@ -158,9 +156,9 @@ function httpSetJsonResponse(res: http.ServerResponse, status: number, text: str
 /*
  * Run the MCP server using the Streamable HTTP transport
  */
-function runStreamableHttp(port: number, deps: string[]) {
+function runStreamableHttp(port: number, deps: string[], returnMode: string) {
   // https://github.com/modelcontextprotocol/typescript-sdk?tab=readme-ov-file#with-session-management
-  const mcpServer = createServer(deps)
+  const mcpServer = createServer(deps, returnMode)
   const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {}
 
   const server = http.createServer(async (req, res) => {
@@ -245,8 +243,8 @@ function runStreamableHttp(port: number, deps: string[]) {
 /*
  * Run the MCP server using the Stdio transport.
  */
-async function runStdio(deps: string[]) {
-  const mcpServer = createServer(deps)
+async function runStdio(deps: string[], returnMode: string) {
+  const mcpServer = createServer(deps, returnMode)
   const transport = new StdioServerTransport()
   await mcpServer.connect(transport)
 }

mcp_run_python/deno/runCode.ts

@@ -1,4 +1,4 @@
-/* eslint @typescript-eslint/no-explicit-any: off */
+// deno-lint-ignore-file no-explicit-any
 import { loadPyodide } from 'pyodide'
 import { preparePythonCode } from './prepareEnvCode.ts'
 import type { LoggingLevel } from '@modelcontextprotocol/sdk/types.js'
@@ -15,7 +15,6 @@ export async function runCode(
 ): Promise<RunSuccess | RunError> {
   // remove once we can upgrade to pyodide 0.27.7 and console.log is no longer used.
   const realConsoleLog = console.log
-  // deno-lint-ignore no-explicit-any
   console.log = (...args: any[]) => log('debug', args.join(' '))
 
   const output: string[] = []
@@ -130,6 +129,17 @@ export function asXml(runResult: RunSuccess | RunError): string {
   return xml.join('\n')
 }
 
+export function asJson(runResult: RunSuccess | RunError): string {
+  const { status, output } = runResult
+  const json: Record<string, any> = { status, output }
+  if (runResult.status == 'success') {
+    json.return_value = JSON.parse(runResult.returnValueJson || 'null')
+  } else {
+    json.error = runResult.error
+  }
+  return JSON.stringify(json)
+}
+
 function escapeClosing(closingTag: string): (str: string) => string {
   const regex = new RegExp(`</?\\s*${closingTag}(?:.*?>)?`, 'gi')
   const onMatch = (match: string) => {
@@ -138,7 +148,6 @@ function escapeClosing(closingTag: string): (str: string) => string {
   return (str) => str.replace(regex, onMatch)
 }
 
-// deno-lint-ignore no-explicit-any
 function formatError(err: any): string {
   let errStr = err.toString()
   errStr = errStr.replace(/^PythonError: +/, '')
@@ -160,6 +169,5 @@ interface PrepareError {
 }
 interface PreparePyEnv {
   prepare_env: (files: CodeFile[]) => Promise<PrepareSuccess | PrepareError>
-  // deno-lint-ignore no-explicit-any
   dump_json: (value: any) => string | null
 }