diff --git a/mcp_run_python/_cli.py b/mcp_run_python/_cli.py index 89258c5..5fed335 100644 --- a/mcp_run_python/_cli.py +++ b/mcp_run_python/_cli.py @@ -28,6 +28,7 @@ def cli_logic(args_list: Sequence[str] | None = None) -> int: ) parser.add_argument('--verbose', action='store_true', help='Enable verbose logging') parser.add_argument('--version', action='store_true', help='Show version and exit') + parser.add_argument('--mount-fs', action='store_true', help='Activate file persistence.') parser.add_argument( 'mode', choices=['stdio', 'streamable-http', 'example'], @@ -53,6 +54,7 @@ def cli_logic(args_list: Sequence[str] | None = None) -> int: http_port=args.port, dependencies=deps, deps_log_handler=deps_log_handler, + file_persistence=args.mount_fs, ) return return_code else: diff --git a/mcp_run_python/deno/deno.jsonc b/mcp_run_python/deno/deno.jsonc index 57866be..eeffc3a 100644 --- a/mcp_run_python/deno/deno.jsonc +++ b/mcp_run_python/deno/deno.jsonc @@ -16,8 +16,11 @@ }, "imports": { "@modelcontextprotocol/sdk": "npm:@modelcontextprotocol/sdk@^1.17.5", - "@std/cli": "jsr:@std/cli@^1.0.15", - "@std/path": "jsr:@std/path@^1.0.8", + "@std/cli": "jsr:@std/cli@^1.0.21", + "@std/encoding": "jsr:@std/encoding@^1.0.10", + "@std/fs": "jsr:@std/fs@^1.0.19", + "@std/media-types": "jsr:@std/media-types@^1.1.0", + "@std/path": "jsr:@std/path@^1.1.2", "pyodide": "npm:pyodide@0.28.2", "zod": "npm:zod@^3.24.4" }, diff --git a/mcp_run_python/deno/deno.lock b/mcp_run_python/deno/deno.lock index 4138faa..f9b937f 100644 --- a/mcp_run_python/deno/deno.lock +++ b/mcp_run_python/deno/deno.lock @@ -1,15 +1,45 @@ { "version": "5", "specifiers": { - "jsr:@std/cli@^1.0.15": "1.0.20", + "jsr:@std/cli@^1.0.21": "1.0.21", + "jsr:@std/encoding@^1.0.10": "1.0.10", + "jsr:@std/fs@^1.0.19": "1.0.19", + "jsr:@std/internal@^1.0.10": "1.0.10", + "jsr:@std/internal@^1.0.9": "1.0.10", + "jsr:@std/media-types@^1.1.0": "1.1.0", + "jsr:@std/path@^1.1.1": "1.1.2", + "jsr:@std/path@^1.1.2": "1.1.2", "npm:@modelcontextprotocol/sdk@^1.17.5": "1.17.5_express@5.1.0_zod@3.25.76", + "npm:@types/node@*": "24.2.0", "npm:@types/node@22.12.0": "22.12.0", "npm:pyodide@0.28.2": "0.28.2", "npm:zod@^3.24.4": "3.25.76" }, "jsr": { - "@std/cli@1.0.20": { - "integrity": "a8c384a2c98cec6ec6a2055c273a916e2772485eb784af0db004c5ab8ba52333" + "@std/cli@1.0.21": { + "integrity": "cd25b050bdf6282e321854e3822bee624f07aca7636a3a76d95f77a3a919ca2a" + }, + "@std/encoding@1.0.10": { + "integrity": "8783c6384a2d13abd5e9e87a7ae0520a30e9f56aeeaa3bdf910a3eaaf5c811a1" + }, + "@std/fs@1.0.19": { + "integrity": "051968c2b1eae4d2ea9f79a08a3845740ef6af10356aff43d3e2ef11ed09fb06", + "dependencies": [ + "jsr:@std/internal@^1.0.9", + "jsr:@std/path@^1.1.1" + ] + }, + "@std/internal@1.0.10": { + "integrity": "e3be62ce42cab0e177c27698e5d9800122f67b766a0bea6ca4867886cbde8cf7" + }, + "@std/media-types@1.1.0": { + "integrity": "c9d093f0c05c3512932b330e3cc1fe1d627b301db33a4c2c2185c02471d6eaa4" + }, + "@std/path@1.1.2": { + "integrity": "c0b13b97dfe06546d5e16bf3966b1cadf92e1cc83e56ba5476ad8b498d9e3038", + "dependencies": [ + "jsr:@std/internal@^1.0.10" + ] } }, "npm": { @@ -33,7 +63,13 @@ "@types/node@22.12.0": { "integrity": "sha512-Fll2FZ1riMjNmlmJOdAyY5pUbkftXslB5DgEzlIuNaiWhXd00FhWxVC/r4yV/4wBb9JfImTu+jiSvXTkJ7F/gA==", "dependencies": [ - "undici-types" + "undici-types@6.20.0" + ] + }, + "@types/node@24.2.0": { + "integrity": "sha512-3xyG3pMCq3oYCNg7/ZP+E1ooTaGB4cG8JWRsqqOYQdbWNY4zbaV0Ennrd7stjiJEFZCaybcIgpTjJWHRfBSIDw==", + "dependencies": [ + "undici-types@7.10.0" ] }, "accepts@2.0.0": { @@ -489,6 +525,9 @@ "undici-types@6.20.0": { "integrity": "sha512-Ny6QZ2Nju20vw1SRHe3d9jVu6gJ+4e3+MMpqu7pqE5HT6WsTSlce++GQmK5UXS8mzV8DSYHrQH+Xrf2jVcuKNg==" }, + "undici-types@7.10.0": { + "integrity": "sha512-t5Fy/nfn+14LuOc2KNYg75vZqClpAiqscVvMygNnlsHBFpSXdJaYtXMcdNLpl/Qvc3P2cB3s6lOV51nqsFq4ag==" + }, "unpipe@1.0.0": { "integrity": "sha512-pjy2bYhSsufwWlKwPc+l3cN7+wuJlK6uz0YdJEOlQDbl6jo/YlPi4mb8agUkVC8BF7V8NuzeyPNqRksA3hztKQ==" }, @@ -526,8 +565,11 @@ }, "workspace": { "dependencies": [ - "jsr:@std/cli@^1.0.15", - "jsr:@std/path@^1.0.8", + "jsr:@std/cli@^1.0.21", + "jsr:@std/encoding@^1.0.10", + "jsr:@std/fs@^1.0.19", + "jsr:@std/media-types@^1.1.0", + "jsr:@std/path@^1.1.2", "npm:@modelcontextprotocol/sdk@^1.17.5", "npm:pyodide@0.28.2", "npm:zod@^3.24.4" diff --git a/mcp_run_python/deno/src/files.ts b/mcp_run_python/deno/src/files.ts new file mode 100644 index 0000000..1294443 --- /dev/null +++ b/mcp_run_python/deno/src/files.ts @@ -0,0 +1,166 @@ +import * as path from '@std/path' +import { exists } from '@std/fs/exists' +import { contentType } from '@std/media-types' +import { type McpServer, ResourceTemplate } from '@modelcontextprotocol/sdk/server/mcp.js' +import z from 'zod' +import { decodeBase64, encodeBase64 } from '@std/encoding/base64' + +/** + * Register file related functions to the MCP server. + * @param server The MCP Server + * @param rootDir Directory in the local file system to read/write to. + */ +export function registerFileFunctions(server: McpServer, rootDir: string) { + // File upload + server.registerTool('upload_file', { + title: 'Upload file.', + description: 'Ingest a file from the given object. Returns a link to the resource that was created.', + inputSchema: { + type: z.union([z.literal('text'), z.literal('bytes')]), + filename: z.string().describe('Name of the file to write.'), + text: z.optional(z.string().describe('Text content of the file, if the type is "text".')), + blob: z.optional(z.string().describe('Base 64 encoded content of the file, if the type is "bytes".')), + }, + }, async ({ type, filename, text, blob }) => { + const absPath = path.join(rootDir, filename) + if (type === 'text') { + if (text == null) { + return { content: [{ type: 'text', text: "Type is 'text', but no text provided." }], isError: true } + } + await Deno.writeTextFile(absPath, text) + } else { + if (blob == null) { + return { content: [{ type: 'text', text: "Type is 'bytes', but no blob provided." }], isError: true } + } + await Deno.writeFile(absPath, decodeBase64(blob)) + } + return { + content: [{ + type: 'resource_link', + uri: `file:///${filename}`, + name: filename, + mimeType: contentType(path.extname(absPath)), + }], + } + }) + + // File Upload from URI + server.registerTool( + 'upload_file_from_uri', + { + title: 'Upload file from URI', + description: 'Ingest a file by URI and store it. Returns a canonical URL.', + inputSchema: { + uri: z.string().url().describe('file:// or https:// style URL'), + filename: z + .string() + .describe('The name of the file to write.'), + }, + }, + async ({ uri, filename }: { uri: string; filename: string }) => { + const absPath = path.join(rootDir, filename) + const fileResponse = await fetch(uri) + if (fileResponse.body) { + const file = await Deno.open(absPath, { write: true, create: true }) + await fileResponse.body.pipeTo(file.writable) + } + return { + content: [{ + type: 'resource_link', + uri: `file:///${filename}`, + name: filename, + mimeType: contentType(path.extname(absPath)), + }], + } + }, + ) + + // Register all the files in the local directory as resources + server.registerResource( + 'read-file', + new ResourceTemplate('file:///{filename}', { + list: async (_extra) => { + const resources = [] + for await (const dirEntry of Deno.readDir(rootDir)) { + if (!dirEntry.isFile) continue + resources.push({ + uri: `file:///${dirEntry.name}`, + name: dirEntry.name, + mimeType: contentType(path.extname(dirEntry.name)), + }) + } + return { resources: resources } + }, + }), + { + title: 'Read file.', + description: 'Read file from persistent storage', + }, + async (uri, { filename }) => { + if (filename == null) { + throw new Deno.errors.NotFound('File not found. No filename provided.') + } + const absPath = path.join(rootDir, ...(Array.isArray(filename) ? filename : [filename])) + const mime = contentType(path.extname(absPath)) || 'application/octet-stream' + const mimeType = mime.split(';')[0] || 'application/octet-stream' + const fileBytes = await Deno.readFile(absPath) + + // Check if it's text-based + if (/^(text\/|.*\/json$|.*\/csv$|.*\/javascript$|.*\/xml$)/.test(mimeType)) { + const text = new TextDecoder().decode(fileBytes) + return { contents: [{ uri: uri.href, mimeType: mime, text: text }] } + } else { + const base64 = encodeBase64(fileBytes) + return { contents: [{ uri: uri.href, mimeType: mime, blob: base64 }] } + } + }, + ) + + // This functions only checks if the file exits + // Download happens through the registered resource + server.registerTool('retrieve_file', { + title: 'Retrieve a file', + description: 'Retrieve a file from the persistent file store.', + inputSchema: { filename: z.string().describe('The name of the file to read.') }, + }, async ({ filename }) => { + const absPath = path.join(rootDir, filename) + if (await exists(absPath, { isFile: true })) { + return { + content: [{ + type: 'resource_link', + uri: `file:///${filename}`, + name: filename, + mimeType: contentType(path.extname(absPath)), + }], + } + } else { + return { + content: [{ 'type': 'text', 'text': `Failed to retrieve file ${filename}. File not found.` }], + isError: true, + } + } + }) + + // File deletion + server.registerTool('delete_file', { + title: 'Delete a file', + description: 'Delete a file from the persistent file store.', + inputSchema: { filename: z.string().describe('The name of the file to delete.') }, + }, async ({ filename }) => { + const absPath = path.join(rootDir, filename) + if (await exists(absPath, { isFile: true })) { + await Deno.remove(absPath) + return { + content: [{ + type: 'text', + text: `${filename} deleted successfully`, + }], + } + } else { + return { + content: [{ 'type': 'text', 'text': `Failed to delete file ${filename}. File not found.` }], + isError: true, + } + } + }) +} diff --git a/mcp_run_python/deno/src/main.ts b/mcp_run_python/deno/src/main.ts index d00d549..cbd27d9 100644 --- a/mcp_run_python/deno/src/main.ts +++ b/mcp_run_python/deno/src/main.ts @@ -12,6 +12,7 @@ import { type LoggingLevel, SetLevelRequestSchema } from '@modelcontextprotocol/ import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js' import { z } from 'zod' +import { registerFileFunctions } from './files.ts' import { asJson, asXml, RunCode } from './runCode.ts' import { Buffer } from 'node:buffer' @@ -20,17 +21,18 @@ const VERSION = '0.0.13' export async function main() { const { args } = Deno const flags = parseArgs(Deno.args, { + boolean: ['mount-fs'], string: ['deps', 'return-mode', 'port'], - default: { port: '3001', 'return-mode': 'xml' }, + default: { port: '3001', 'return-mode': 'xml', 'mount-fs': false }, }) const deps = flags.deps?.split(',') ?? [] if (args.length >= 1) { if (args[0] === 'stdio') { - await runStdio(deps, flags['return-mode']) + await runStdio(deps, flags['return-mode'], flags['mount-fs']) return } else if (args[0] === 'streamable_http') { const port = parseInt(flags.port) - runStreamableHttp(port, deps, flags['return-mode']) + runStreamableHttp(port, deps, flags['return-mode'], flags['mount-fs']) return } else if (args[0] === 'example') { await example(deps) @@ -49,7 +51,8 @@ Usage: deno ... deno/main.ts [stdio|streamable_http|install_deps|noop] options: --port Port to run the HTTP server on (default: 3001) --deps Comma separated list of dependencies to install ---return-mode Return mode for output data (default: xml)`, +--return-mode Return mode for output data (default: xml) +--mount-fs Activate file persistence in sandbox`, ) Deno.exit(1) } @@ -57,8 +60,14 @@ options: /* * Create an MCP server with the `run_python_code` tool registered. */ -function createServer(deps: string[], returnMode: string): McpServer { +function createServer( + deps: string[], + returnMode: string, + mountFS: boolean, + transport?: StdioServerTransport, +): McpServer { const runCode = new RunCode() + const server = new McpServer( { name: 'MCP Run Python', @@ -67,17 +76,41 @@ function createServer(deps: string[], returnMode: string): McpServer { { instructions: 'Call the "run_python_code" tool with the Python code to run.', capabilities: { + resources: {}, + tools: {}, logging: {}, }, }, ) + // Create storage directory + let extraDescription = '' + if (mountFS) { + const rootDir = Deno.makeTempDirSync({ prefix: 'mcp_run_python' }) + runCode.rootDir = rootDir + const signalHandler = () => { + Deno.removeSync(rootDir, { recursive: true }) + Deno.exit() + } + Deno.addSignalListener('SIGINT', signalHandler) + Deno.addSignalListener('SIGTERM', signalHandler) + extraDescription = ` + You can read and create persisted files at ~/storage/. + ` + registerFileFunctions(server, rootDir) + if (transport) { + transport.onclose = () => { + Deno.removeSync(rootDir, { recursive: true }) + } + } + } + const toolDescription = `Tool to execute Python code and return stdout, stderr, and return value. The code may be async, and the value on the last line will be returned as the return value. The code will be executed with Python 3.13. -` +` + extraDescription let setLogLevel: LoggingLevel = 'emergency' @@ -167,9 +200,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[], returnMode: string) { +function runStreamableHttp(port: number, deps: string[], returnMode: string, mountFS: boolean) { // https://github.com/modelcontextprotocol/typescript-sdk?tab=readme-ov-file#with-session-management - const mcpServer = createServer(deps, returnMode) + const mcpServer = createServer(deps, returnMode, mountFS) const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {} const server = http.createServer(async (req, res) => { @@ -252,9 +285,9 @@ function runStreamableHttp(port: number, deps: string[], returnMode: string) { /* * Run the MCP server using the Stdio transport. */ -async function runStdio(deps: string[], returnMode: string) { - const mcpServer = createServer(deps, returnMode) +async function runStdio(deps: string[], returnMode: string, mountFS: boolean) { const transport = new StdioServerTransport() + const mcpServer = createServer(deps, returnMode, mountFS, transport) await mcpServer.connect(transport) } diff --git a/mcp_run_python/deno/src/runCode.ts b/mcp_run_python/deno/src/runCode.ts index 1be8681..845d92c 100644 --- a/mcp_run_python/deno/src/runCode.ts +++ b/mcp_run_python/deno/src/runCode.ts @@ -21,6 +21,8 @@ export class RunCode { private preparePyEnv?: PreparePyEnv private prepPromise?: Promise + public rootDir?: string + async run( dependencies: string[], log: (level: LoggingLevel, data: string) => void, @@ -96,6 +98,19 @@ export class RunCode { }, }) + // Mount file system + if (this.rootDir) { + const mountDir = '/home/pyodide/storage/' + // Ensure emscriptem directory is created + pyodide.FS.mkdirTree(mountDir) + // Mount local directory + pyodide.FS.mount( + pyodide.FS.filesystems.NODEFS, + { root: this.rootDir }, + mountDir, + ) + } + // see https://github.com/pyodide/pyodide/discussions/5512 const origLoadPackage = pyodide.loadPackage pyodide.loadPackage = (pkgs, options) => diff --git a/mcp_run_python/main.py b/mcp_run_python/main.py index 34d177d..da4a49f 100644 --- a/mcp_run_python/main.py +++ b/mcp_run_python/main.py @@ -26,6 +26,7 @@ def run_mcp_server( return_mode: Literal['json', 'xml'] = 'xml', deps_log_handler: LogHandler | None = None, allow_networking: bool = True, + file_persistence: bool = False, ) -> int: """Install dependencies then run the mcp-run-python server. @@ -36,6 +37,7 @@ def run_mcp_server( return_mode: The mode to return tool results in. deps_log_handler: Optional function to receive logs emitted while installing dependencies. allow_networking: Whether to allow networking when running provided python code. + file_persistence: Files stored in the `~/storage/` directory will be persisted between MCP runs. """ with prepare_deno_env( mode, @@ -44,6 +46,7 @@ def run_mcp_server( return_mode=return_mode, deps_log_handler=deps_log_handler, allow_networking=allow_networking, + file_persistence=file_persistence, ) as env: if mode == 'streamable_http': logger.info('Running mcp-run-python via %s on port %d...', mode, http_port) @@ -74,6 +77,7 @@ def prepare_deno_env( return_mode: Literal['json', 'xml'] = 'xml', deps_log_handler: LogHandler | None = None, allow_networking: bool = True, + file_persistence: bool = False, ) -> Iterator[DenoEnv]: """Prepare the deno environment for running the mcp-run-python server with Deno. @@ -89,6 +93,7 @@ def prepare_deno_env( deps_log_handler: Optional function to receive logs emitted while installing dependencies. allow_networking: Whether the prepared DenoEnv should allow networking when running code. Note that we always allow networking during environment initialization to install dependencies. + file_persistence: Files stored in the `~/storage/` directory will be persisted between MCP runs. Returns: Yields the deno environment details. @@ -121,6 +126,7 @@ def prepare_deno_env( dependencies=dependencies, return_mode=return_mode, allow_networking=allow_networking, + file_persistence=file_persistence, ) yield DenoEnv(cwd, args) @@ -137,6 +143,7 @@ async def async_prepare_deno_env( return_mode: Literal['json', 'xml'] = 'xml', deps_log_handler: LogHandler | None = None, allow_networking: bool = True, + file_persistence: bool = False, ) -> AsyncIterator[DenoEnv]: """Async variant of `prepare_deno_env`.""" ct = await _asyncify( @@ -147,6 +154,7 @@ async def async_prepare_deno_env( return_mode=return_mode, deps_log_handler=deps_log_handler, allow_networking=allow_networking, + file_persistence=file_persistence, ) try: yield await _asyncify(ct.__enter__) @@ -176,12 +184,16 @@ def _deno_run_args( dependencies: list[str] | None = None, return_mode: Literal['json', 'xml'] = 'xml', allow_networking: bool = True, + file_persistence: bool = False, ) -> list[str]: args = ['run'] if allow_networking: args += ['--allow-net'] + if file_persistence: + args += ['--allow-read=./node_modules,/tmp', '--allow-write=/tmp'] + else: + args += ['--allow-read=./node_modules/'] args += [ - '--allow-read=./node_modules', '--node-modules-dir=auto', 'src/main.ts', mode, @@ -189,6 +201,8 @@ def _deno_run_args( ] if dependencies is not None: args.append(f'--deps={",".join(dependencies)}') + if file_persistence: + args += ['--mount-fs'] if http_port is not None: if mode == 'streamable_http': args.append(f'--port={http_port}') diff --git a/tests/test_mcp_servers.py b/tests/test_mcp_servers.py index 889b903..8b24fa4 100644 --- a/tests/test_mcp_servers.py +++ b/tests/test_mcp_servers.py @@ -3,9 +3,10 @@ import asyncio import re import subprocess -from collections.abc import AsyncIterator, Callable +from collections.abc import AsyncIterator from contextlib import AbstractAsyncContextManager, asynccontextmanager -from typing import TYPE_CHECKING +from enum import Enum +from typing import TYPE_CHECKING, Literal, Protocol import pytest from httpx import AsyncClient, HTTPError @@ -22,14 +23,28 @@ pytestmark = pytest.mark.anyio +class McpTools(str, Enum): + RUN_PYTHON_CODE = 'run_python_code' + UPLOAD_FILE = 'upload_file' + UPLOAD_FILE_FROM_URI = 'upload_file_from_uri' + RETRIEVE_FILE = 'retrieve_file' + DELETE_FILE = 'delete_file' + + +class SessionManagerFactory(Protocol): + def __call__( + self, deps: list[str], file_persistence: bool = False + ) -> AbstractAsyncContextManager[ClientSession]: ... + + @pytest.fixture(name='run_mcp_session', params=['stdio', 'streamable_http']) def fixture_run_mcp_session( request: pytest.FixtureRequest, -) -> Callable[[list[str]], AbstractAsyncContextManager[ClientSession]]: +) -> SessionManagerFactory: @asynccontextmanager - async def run_mcp(deps: list[str]) -> AsyncIterator[ClientSession]: + async def run_mcp(deps: list[str], file_persistence: bool = False) -> AsyncIterator[ClientSession]: if request.param == 'stdio': - async with async_prepare_deno_env('stdio', dependencies=deps) as env: + async with async_prepare_deno_env('stdio', dependencies=deps, file_persistence=file_persistence) as env: server_params = StdioServerParameters(command='deno', args=env.args, cwd=env.cwd) async with stdio_client(server_params) as (read, write): async with ClientSession(read, write) as session: @@ -37,7 +52,9 @@ async def run_mcp(deps: list[str]) -> AsyncIterator[ClientSession]: else: assert request.param == 'streamable_http', request.param port = 3101 - async with async_prepare_deno_env('streamable_http', http_port=port, dependencies=deps) as env: + async with async_prepare_deno_env( + 'streamable_http', http_port=port, dependencies=deps, file_persistence=file_persistence + ) as env: p = subprocess.Popen(['deno', *env.args], cwd=env.cwd) try: url = f'http://localhost:{port}/mcp' @@ -72,7 +89,7 @@ async def wait_for_server(url: str, timeout: float): raise TimeoutError(f'URL {url} did not become available within {timeout} seconds') -async def test_list_tools(run_mcp_session: Callable[[list[str]], AbstractAsyncContextManager[ClientSession]]) -> None: +async def test_list_tools(run_mcp_session: SessionManagerFactory) -> None: async with run_mcp_session([]) as mcp_session: await mcp_session.initialize() tools = await mcp_session.list_tools() @@ -178,19 +195,183 @@ async def test_list_tools(run_mcp_session: Callable[[list[str]], AbstractAsyncCo ], ) async def test_run_python_code( - run_mcp_session: Callable[[list[str]], AbstractAsyncContextManager[ClientSession]], + run_mcp_session: SessionManagerFactory, deps: list[str], code: list[str], expected_output: str, ) -> None: async with run_mcp_session(deps) as mcp_session: await mcp_session.initialize() + result = await mcp_session.call_tool(McpTools.RUN_PYTHON_CODE, {'python_code': '\n'.join(code)}) + assert len(result.content) == 1 + content = result.content[0] + assert isinstance(content, types.TextContent) + assert content.text == expected_output + + +CSV_DATA = """Name,Age,Department,Salary +Alice,25,Engineering,60000 +Bob,32,Marketing,52000 +Charlie,29,Engineering,70000 +Diana,45,HR,65000 +Ethan,35,Marketing,58000 +Fiona,28,Engineering,72000 +George,40,HR,64000 +Hannah,31,Engineering,68000 +Ian,38,Marketing,61000 +Julia,27,HR,59000 +""" + +BASE_64_IMAGE = 'iVBORw0KGgoAAAANSUhEUgAAAQAAAAEACAIAAADTED8xAAAEX0lEQVR4nOzdO8vX9R/HcS/56f8PWotGQkPBBUWESCQYNJR0GjIn6UBTgUMZTiGE4ZgRVKNkuDSEFtgBQqIiKunkEFdkWLmEBQUWiNUQYd2KNwTPx+MGvD7Tk/f2/S7O7tmyatKnJx8b3f/p6EOj+5euu2Z0/+Sxt0f3N++9fHR/+57/j+7vuPuT0f3Vo+vwHycA0gRAmgBIEwBpAiBNAKQJgDQBkCYA0gRAmgBIEwBpAiBNAKQJgDQBkCYA0gRAmgBIEwBpAiBNAKQJgDQBkCYA0gRAmgBIEwBpAiBNAKQtDr561+gDpzf9PLp/4eNzo/uXzv41uv/BM0+O7h9/bsPo/vqPdo3u7965GN13AUgTAGkCIE0ApAmANAGQJgDSBECaAEgTAGkCIE0ApAmANAGQJgDSBECaAEgTAGkCIE0ApAmANAGQJgDSBECaAEgTAGkCIE0ApAmANAGQJgDSlh5ce+XoA9+eODK6v3r7naP7b31zaHT/4p+3jO4f2/Tb6P7K41tH9zff+8LovgtAmgBIEwBpAiBNAKQJgDQBkCYA0gRAmgBIEwBpAiBNAKQJgDQBkCYA0gRAmgBIEwBpAiBNAKQJgDQBkCYA0gRAmgBIEwBpAiBNAKQJgDQBkLb09ZmLow8sb1ke3d92YXR+1dO7PhzdX7f2xtH9Q5fN/t/g2j9eHt3/cc350X0XgDQBkCYA0gRAmgBIEwBpAiBNAKQJgDQBkCYA0gRAmgBIEwBpAiBNAKQJgDQBkCYA0gRAmgBIEwBpAiBNAKQJgDQBkCYA0gRAmgBIEwBpAiBtcf3eW0cfePTE7Pf1D9yxMrq/4YrR+VWvnN84uv/lvs2j+2v3nx3dv3rT/0b3XQDSBECaAEgTAGkCIE0ApAmANAGQJgDSBECaAEgTAGkCIE0ApAmANAGQJgDSBECaAEgTAGkCIE0ApAmANAGQJgDSBECaAEgTAGkCIE0ApAmAtKWrzq0ffeD312f339h5ZnT/npsPj+7//cPDo/un739idP/Xg5+P7j/y/G2j+y4AaQIgTQCkCYA0AZAmANIEQJoASBMAaQIgTQCkCYA0AZAmANIEQJoASBMAaQIgTQCkCYA0AZAmANIEQJoASBMAaQIgTQCkCYA0AZAmANIEQNpi/5FfRh94753XRvcP7F0zuv/V7e+O7t906v3R/WdP/zO6f9/ixdH9G3Z/NrrvApAmANIEQJoASBMAaQIgTQCkCYA0AZAmANIEQJoASBMAaQIgTQCkCYA0AZAmANIEQJoASBMAaQIgTQCkCYA0AZAmANIEQJoASBMAaQIgTQCkLb25vDL6wLoHjo7ur7z03ej++u+fGt0/vm/2+/dfHF4e3d9xauPo/taN20b3XQDSBECaAEgTAGkCIE0ApAmANAGQJgDSBECaAEgTAGkCIE0ApAmANAGQJgDSBECaAEgTAGkCIE0ApAmANAGQJgDSBECaAEgTAGkCIE0ApAmAtH8DAAD//9drYGg9ROu9AAAAAElFTkSuQmCC' + + +@pytest.mark.parametrize( + 'deps,code,expected_output,data_type,expected_file', + [ + pytest.param( + ['pillow'], + [ + 'from PIL import Image, ImageFilter', + 'img = Image.open("storage/image.png")', + 'gray_img = img.convert("L")', + 'gray_img.save("storage/image-gray.png")', + 'print(f"Image size: {img.size}")', + ], + snapshot("""\ +success + +Image size: (256, 256) +\ +"""), + 'bytes', + snapshot( + 'iVBORw0KGgoAAAANSUhEUgAAAQAAAAEACAAAAAB5Gfe6AAAC6ElEQVR4nO3cMaqVZxhFYY9ea7EIiJ0DEEv7gEUghag4gEAGYCNoestg4wCcgE3ggmWIpULSJUJwBEEQQZEEnIG7sHggZ612f8Vi8Zb/OYdHZ77Mn2P/cez/jv33sT8f+52x/zP2s2P/31MALaApgBbQFEALaAqgBTQF0AKaAmgBTQG0gKYAWkBTAC2gKYAW0BRAC2gKoAU0BdACmgJoAU0BtIDmcH08ODf2G2P/MPZ3Y7869vdjPx370V9AAbSApgBaQFMALaApgBbQFEALaAqgBTQF0AKaAmgBTQG0gKYAWkBTAC2gKYAW0BRAC2gKoAU0BdACmsOV8eDe2K+N/eLY/xv7+n7g6djX9wVHfwEF0AKaAmgBTQG0gKYAWkBTAC2gKYAW0BRAC2gKoAU0BdACmgJoAU0BtICmAFpAUwAtoCmAFtAUQAtoDq/Hg/tjvzX2v8b+duyfvnL/OPajv4ACaAFNAbSApgBaQFMALaApgBbQFEALaAqgBTQF0AKaAmgBTQG0gKYAWkBTAC2gKYAW0BRAC2gKoAU0h4fjwfr9/W9jvzz2F2O/OfZnY7809qO/gAJoAU0BtICmAFpAUwAtoCmAFtAUQAtoCqAFNAXQApoCaAFNAbSApgBaQFMALaApgBbQFEALaAqgBTQn34wHp2P/aezfjf3nsb8c+4OxPxn70V9AAbSApgBaQFMALaApgBbQFEALaAqgBTQF0AKaAmgBTQG0gKYAWkBTAC2gKYAW0BRAC2gKoAU0BdACmpPz48H6//+7Y7859j/G/uvYb4/9+7Ef/QUUQAtoCqAFNAXQApoCaAFNAbSApgBaQFMALaApgBbQFEALaAqgBTQF0AKaAmgBTQG0gKYAWkBTAC2gOXk1Hnw79l/G/sPYH4/9zdj/HvuFsR/9BRRAC2gKoAU0BdACmgJoAU0BtICmAFpAUwAtoCmAFtAUQAtoCqAFNAXQApoCaAFNAbSApgBaQFMALaD5DBLPJU7x++9vAAAAAElFTkSuQmCC' + ), + id='image-transform', + ), + pytest.param( + ['pandas'], + [ + 'import pandas as pd', + 'df = pd.read_csv("storage/data.csv")', + 'df["Age_in_10_years"] = df["Age"] + 10', + 'df.to_csv("storage/data-processed.csv", index=False)', + 'print(df.describe())', + ], + snapshot("""\ +success + + Age Salary Age_in_10_years +count 10.000000 10.000000 10.000000 +mean 33.000000 62900.000000 43.000000 +std 6.394442 6100.091074 6.394442 +min 25.000000 52000.000000 35.000000 +25% 28.250000 59250.000000 38.250000 +50% 31.500000 62500.000000 41.500000 +75% 37.250000 67250.000000 47.250000 +max 45.000000 72000.000000 55.000000 +\ +"""), + 'text', + snapshot("""\ +Name,Age,Department,Salary,Age_in_10_years +Alice,25,Engineering,60000,35 +Bob,32,Marketing,52000,42 +Charlie,29,Engineering,70000,39 +Diana,45,HR,65000,55 +Ethan,35,Marketing,58000,45 +Fiona,28,Engineering,72000,38 +George,40,HR,64000,50 +Hannah,31,Engineering,68000,41 +Ian,38,Marketing,61000,48 +Julia,27,HR,59000,37 +"""), + id='dataframe-manipulation', + ), + ], +) +async def test_run_python_code_with_files( + run_mcp_session: SessionManagerFactory, + deps: list[str], + code: list[str], + expected_output: str, + data_type: Literal['bytes', 'text'], + expected_file: str, +) -> None: + async with run_mcp_session(deps, file_persistence=True) as mcp_session: + await mcp_session.initialize() + + match data_type: + case 'text': + filename = 'data.csv' + new_filename = 'data-processed.csv' + ctype = 'text/csv' + result = await mcp_session.call_tool( + McpTools.UPLOAD_FILE, {'type': 'text', 'filename': filename, 'text': CSV_DATA} + ) + + case 'bytes': + filename = 'image.png' + new_filename = 'image-gray.png' + ctype = 'image/png' + result = await mcp_session.call_tool( + McpTools.UPLOAD_FILE, {'type': 'bytes', 'filename': filename, 'blob': BASE_64_IMAGE} + ) + + assert result.isError is False + assert len(result.content) == 1 + content = result.content[0] + assert isinstance(content, types.ResourceLink) + assert str(content.uri) == f'file:///{filename}' + assert content.name == filename + assert content.mimeType is not None + assert content.mimeType.startswith(ctype) + + result = await mcp_session.list_resources() + assert len(result.resources) == 1 + content = result.resources[0] + assert str(content.uri) == f'file:///{filename}' + assert content.name == filename + assert content.mimeType is not None + assert content.mimeType.startswith(ctype) + result = await mcp_session.call_tool('run_python_code', {'python_code': '\n'.join(code)}) + assert result.isError is False assert len(result.content) == 1 content = result.content[0] assert isinstance(content, types.TextContent) assert content.text == expected_output + result = await mcp_session.list_resources() + assert len(result.resources) == 2 + assert {filename, new_filename} == set(resource.name for resource in result.resources) + + result = await mcp_session.call_tool(McpTools.RETRIEVE_FILE, {'filename': new_filename}) + assert result.isError is False + assert len(result.content) == 1 + content = result.content[0] + assert isinstance(content, types.ResourceLink) + assert str(content.uri) == f'file:///{new_filename}' + assert content.name == new_filename + assert content.mimeType is not None + assert content.mimeType.startswith(ctype) + + result = await mcp_session.read_resource(content.uri) + assert len(result.contents) == 1 + content = result.contents[0] + assert str(content.uri) == f'file:///{new_filename}' + assert content.mimeType is not None + assert content.mimeType.startswith(ctype) + + match data_type: + case 'text': + assert isinstance(content, types.TextResourceContents) + assert content.text == expected_file + + case 'bytes': + assert isinstance(content, types.BlobResourceContents) + assert content.blob == expected_file + async def test_install_run_python_code() -> None: logs: list[str] = []