Add upload tests

Author: BrokenDuck

mcp-run-python/deno.json

@@ -15,6 +15,7 @@
   "imports": {
     "@modelcontextprotocol/sdk": "npm:@modelcontextprotocol/sdk@^1.17.4",
     "@std/cli": "jsr:@std/cli@^1.0.15",
+    "@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.0.8",

mcp-run-python/deno.lock

@@ -3,6 +3,7 @@ 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 { encodeBase64 } from '@std/encoding/base64'
 
 /**
  * Returns the temporary directory in the local filesystem for file persistence.
@@ -11,56 +12,6 @@ export function createRootDir(): string {
   return Deno.makeTempDirSync({ prefix: 'mcp_run_python' })
 }
 
-/**
- * Streams a file from an http(s) or file:// URI into targetPath.
- * Returns when bytes are fully written.
- */
-async function uploadFromUri(
-  uri: string,
-  absPath: string,
-): Promise<void> {
-  await Deno.mkdir(
-    path.dirname(absPath),
-    { recursive: true },
-  )
-
-  if (uri.startsWith('http://') || uri.startsWith('https://')) {
-    const res = await fetch(uri)
-    if (!res.ok || !res.body) {
-      throw new Error(`Fetch failed: ${res.status} ${res.statusText}`)
-    }
-    const file = await Deno.open(absPath, {
-      write: true,
-      create: true,
-      truncate: true,
-    })
-    try {
-      await res.body.pipeTo(file.writable)
-    } finally {
-      file.close()
-    }
-    return
-  }
-
-  if (uri.startsWith('file://')) {
-    const src = await Deno.open(path.fromFileUrl(uri), { read: true })
-    const dest = await Deno.open(absPath, {
-      write: true,
-      create: true,
-      truncate: true,
-    })
-    try {
-      await src.readable.pipeTo(dest.writable)
-    } finally {
-      src.close()
-      dest.close()
-    }
-    return
-  }
-
-  throw new Error(`Unsupported URI scheme: ${uri}`)
-}
-
 /**
  * Register file related functions to the MCP server.
  * @param server The MCP Server
@@ -82,9 +33,18 @@ export function registerFileFunctions(server: McpServer, rootDir: string) {
     },
     async ({ uri, filename }: { uri: string; filename: string }) => {
       const absPath = path.join(rootDir, filename)
-      await uploadFromUri(uri, absPath)
+      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: 'text', text: 'Upload successful.' }],
+        content: [{
+          type: 'resource_link',
+          uri: `file:///${filename}`,
+          name: filename,
+          mimeType: contentType(path.extname(absPath)),
+        }],
       }
     },
   )
@@ -112,15 +72,15 @@ export function registerFileFunctions(server: McpServer, rootDir: string) {
     },
     async (uri, { filename }) => {
       const absPath = path.join(rootDir, ...(Array.isArray(filename) ? filename : [filename]))
-      const mime = contentType(absPath)
+      const mime = contentType(path.extname(absPath))
       const fileBytes = await Deno.readFile(absPath)
 
       // Check if it's text-based
-      if (mime && /^text\/|\/json$|\/csv$|\/javascript$|\/xml$/.test(mime)) {
+      if (mime && /^(text\/|.*\/json$|.*\/csv$|.*\/javascript$|.*\/xml$)/.test(mime.split(';')[0])) {
         const text = new TextDecoder().decode(fileBytes)
         return { contents: [{ uri: uri.href, name: filename, mimeType: mime, text: text }] }
       } else {
-        const base64 = btoa(String.fromCharCode(...fileBytes))
+        const base64 = encodeBase64(String.fromCharCode(...fileBytes))
         return { contents: [{ uri: uri.href, name: filename, mimeType: mime, blob: base64 }] }
       }
     },

mcp-run-python/src/files.ts

@@ -18,12 +18,13 @@ const VERSION = '0.0.13'
 export async function main() {
   // Parse global flags once, then branch on subcommand
   const flags = parseArgs(Deno.args, {
-    string: ['port'],
-    default: { port: '3001', mount: false },
+    string: ['port', 'mount'],
+    default: { port: '3001' },
   })
   const mode = (flags._[0] as string | undefined) ?? ''
   const port = parseInt(flags.port as string)
-  const mount = flags.mount as string | boolean
+  const rawMount = flags.mount as string | undefined
+  const mount: string | boolean = rawMount === undefined ? false : rawMount === '' ? true : rawMount
 
   if (mode === 'stdio') {
     await runStdio(mount)
@@ -32,7 +33,7 @@ export async function main() {
   } else if (mode === 'sse') {
     runSse(port, mount)
   } else if (mode === 'warmup') {
-    await warmup(mount)
+    await warmup()
   } else {
     console.error(
       `\
@@ -58,8 +59,11 @@ function createServer(rootDir: string | null, mount: string | boolean): McpServe
       version: VERSION,
     },
     {
-      instructions: 'Call the "run_python_code" tool with the Python code to run.',
+      instructions: 'Call the "run_python_code" tool with the Python code to run.' +
+        (rootDir != null ? ` Persistent storage is mounted at: "${rootDir}".` : ''),
       capabilities: {
+        resources: {},
+        tools: {},
         logging: {},
       },
     },
@@ -128,6 +132,36 @@ function httpSetJsonResponse(
   res.end()
 }
 
+function addDirCleanupCallback(server: http.Server | StdioServerTransport, dir: string) {
+  let cleaned = false
+  const cleanup = () => {
+    if (cleaned) return
+    cleaned = true
+    try {
+      Deno.removeSync(dir, { recursive: true })
+    } catch {
+      // ignore
+    }
+  }
+  if (server instanceof http.Server) {
+    server.on('close', cleanup)
+  } else {
+    server.onclose = cleanup
+  }
+  const handleSig = () => {
+    try {
+      server.close(() => {})
+    } catch {
+      // ignore
+    }
+    cleanup()
+    Deno.exit()
+  }
+  Deno.addSignalListener('SIGINT', handleSig)
+  Deno.addSignalListener('SIGTERM', handleSig)
+  addEventListener('unload', cleanup)
+}
+
 /*
  * Run the MCP server using the Streamable HTTP transport
  */
@@ -215,16 +249,14 @@ function runStreamableHttp(port: number, mount: string | boolean) {
     }
   })
 
-  // Cleanup root dir on server close
+  // Cleanup root dir on server close and on process signals
   if (rootDir != null) {
-    server.on('close', () => {
-      Deno.removeSync(rootDir, { recursive: true })
-    })
+    addDirCleanupCallback(server, rootDir)
   }
 
   server.listen(port, () => {
     console.log(
-      `Running MCP Run Python version ${VERSION} with Streamable HTTP transport on port ${port}`,
+      `Running MCP Run Python version ${VERSION} with SSE transport on port ${port}.`,
     )
   })
 }
@@ -275,16 +307,14 @@ function runSse(port: number, mount: string | boolean) {
     }
   })
 
-  // Cleanup root dir on server close
+  // Cleanup root dir on server close and on process signals
   if (rootDir != null) {
-    server.on('close', () => {
-      Deno.removeSync(rootDir, { recursive: true })
-    })
+    addDirCleanupCallback(server, rootDir)
   }
 
   server.listen(port, () => {
     console.log(
-      `Running MCP Run Python version ${VERSION} with SSE transport on port ${port}`,
+      `Running MCP Run Python version ${VERSION} with SSE transport on port ${port}.`,
     )
   })
 }
@@ -297,11 +327,9 @@ async function runStdio(mount: string | boolean) {
   const mcpServer = createServer(rootDir, mount)
   const transport = new StdioServerTransport()
 
-  // Cleanup root dir on server close
+  // Cleanup root dir on server close and on process signals
   if (rootDir != null) {
-    transport.onclose = () => {
-      Deno.removeSync(rootDir, { recursive: true })
-    }
+    addDirCleanupCallback(transport, rootDir)
   }
 
   await mcpServer.connect(transport)
@@ -310,10 +338,9 @@ async function runStdio(mount: string | boolean) {
 /*
  * Run pyodide to download packages which can otherwise interrupt the server
  */
-async function warmup(mount?: string | boolean) {
+async function warmup() {
   console.error(
-    `Running warmup script for MCP Run Python version ${VERSION}...` +
-      (mount ? ` (mount: ${typeof mount === 'string' ? mount : 'enabled'})` : ''),
+    `Running warmup script for MCP Run Python version ${VERSION}...`,
   )
   const code = `
 import numpy