Add direct upload tool

Author: BrokenDuck

mcp-run-python/src/files.ts

@@ -3,7 +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'
+import { decodeBase64, encodeBase64 } from '@std/encoding/base64'
 
 /**
  * Returns the temporary directory in the local filesystem for file persistence.
@@ -18,6 +18,37 @@ export function createRootDir(): string {
  * @param rootDir Directory in the local file system to read/write to.
  */
 export function registerFileFunctions(server: McpServer, rootDir: string) {
+  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.nullable(z.string().describe('Text content of the file, if the type is "text".')),
+      blob: z.nullable(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
   server.registerTool(
     'upload_file_from_uri',

mcp-run-python/test_mcp_servers.py

@@ -39,6 +39,7 @@
 
 class McpTools(StrEnum):
     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'
@@ -169,7 +170,7 @@ async def test_list_tools(
             )
         else:
             # Check tools
-            assert len(tools.tools) == 4
+            assert len(tools.tools) == 5
             # sort tools by their name
             sorted_tools = sorted(tools.tools, key=lambda t: t.name)
 
@@ -345,8 +346,59 @@ def do_GET(self):
             httpd.server_close()
             t.join(timeout=2)
 
-    @pytest.mark.parametrize('uri_type', ['http', 'file'])
+    @pytest.mark.parametrize('file_type', ['text', 'bytes'])
     async def test_upload_files(
+        self,
+        mcp_session: ClientSession,
+        server_type: Literal['stdio', 'sse', 'streamable_http'],
+        mount: bool | str,
+        file_type: Literal['text', 'bytes'],
+        tmp_path: Path,
+    ) -> None:
+        if mount is False:
+            pytest.skip('No directory mounted.')
+        result = await mcp_session.initialize()
+
+        # Extract directory from response
+        storageDir = self.get_dir_from_instructions(result.instructions)
+        assert storageDir.is_dir()
+
+        match file_type:
+            case 'text':
+                filename = 'data.csv'
+                ctype = 'text/csv'
+                result = await mcp_session.call_tool(
+                    McpTools.UPLOAD_FILE, {'type': 'text', 'filename': filename, 'text': CSV_DATA, 'blob': None}
+                )
+
+            case 'bytes':
+                filename = 'image.png'
+                ctype = 'image/png'
+                result = await mcp_session.call_tool(
+                    McpTools.UPLOAD_FILE, {'type': 'bytes', 'filename': filename, 'blob': BASE_64_IMAGE, 'text': None}
+                )
+
+        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)
+
+        createdFile = storageDir / filename
+        assert createdFile.exists()
+        assert createdFile.is_file()
+
+        match file_type:
+            case 'text':
+                assert createdFile.read_text() == CSV_DATA
+            case 'bytes':
+                assert base64.b64encode(createdFile.read_bytes()).decode() == BASE_64_IMAGE
+
+    @pytest.mark.parametrize('uri_type', ['http', 'file'])
+    async def test_upload_files_with_uri(
         self,
         mcp_session: ClientSession,
         server_type: Literal['stdio', 'sse', 'streamable_http'],
@@ -400,7 +452,7 @@ async def test_download_files(
         server_type: Literal['stdio', 'sse', 'streamable_http'],
         mount: bool | str,
         content_type: Literal['bytes', 'text'],
-    ):
+    ) -> None:
         if mount is False:
             pytest.skip('No directory mounted.')
         result = await mcp_session.initialize()