Update container read tool call (#54)

* Update read tool call

* forgot a test

* prettier

Author: courtney-sims

apps/sandbox-container/container/fileUtils.spec.ts

@@ -1,6 +1,21 @@
-import { describe, expect, it } from 'vitest'
+import mime from 'mime'
+import mock from 'mock-fs'
+import { afterEach, describe, expect, it, vi } from 'vitest'
 
-import { get_file_name_from_path } from './fileUtils'
+import { get_file_name_from_path, get_mime_type, list_files_in_directory } from './fileUtils'
+
+vi.mock('mime', () => {
+	return {
+		default: {
+			getType: vi.fn(),
+		},
+	}
+})
+
+afterEach(async () => {
+	mock.restore()
+	vi.restoreAllMocks()
+})
 
 describe('get_file_name_from_path', () => {
 	it('strips files/contents', async () => {
@@ -15,4 +30,50 @@ describe('get_file_name_from_path', () => {
 			const path = await get_file_name_from_path('/files/contents/birds/')
 			expect(path).toBe('/birds')
 		})
-})
+}),
+	describe('list_files_in_directory', () => {
+		it('lists the files in a directory', async () => {
+			mock({
+				testDir: {
+					cats: 'aurora, luna',
+					dogs: 'penny',
+				},
+			})
+			const listFiles = await list_files_in_directory('testDir')
+			expect(listFiles).toEqual(['file:///testDir/cats', 'file:///testDir/dogs'])
+		}),
+			it('throws an error if path is not a directory', async () => {
+				mock({
+					testDir: {
+						cats: 'aurora, luna',
+						dogs: 'penny',
+					},
+				})
+				await expect(async () => await list_files_in_directory('testDir/cats')).rejects.toThrow(
+					'Failed to read directory'
+				)
+			}),
+			it('treats empty strings as cwd', async () => {
+				mock({
+					testDir: {
+						cats: 'aurora, luna',
+						dogs: 'penny',
+					},
+				})
+
+				const listFiles = await list_files_in_directory('')
+				expect(listFiles).toEqual(['file:///../../../../../../testDir'])
+			})
+	}),
+	describe('get_mime_type', async () => {
+		it("provides the natural mime type when not 'inode/directory'", async () => {
+			vi.mocked(mime.getType).mockReturnValueOnce('theType')
+			const mimeType = await get_mime_type('someFile')
+			expect(mimeType).toEqual('theType')
+		})
+		it("overrides mime type for 'inode/directory'", async () => {
+			vi.mocked(mime.getType).mockReturnValueOnce('inode/directory')
+			const mimeType = await get_mime_type('someDirectory')
+			expect(mimeType).toEqual('text/directory')
+		})
+	})

apps/sandbox-container/container/fileUtils.ts

@@ -1,6 +1,39 @@
+import * as fs from 'node:fs/promises'
+import path from 'node:path'
+import mime from 'mime'
+
+// this is because there isn't a "real" directory mime type, so we're reusing the "text/directory" mime type
+// so claude doesn't give an error
+export const DIRECTORY_CONTENT_TYPE = 'text/directory'
+
 export async function get_file_name_from_path(path: string): Promise<string> {
 	path = path.replace('/files/contents', '')
 	path = path.endsWith('/') ? path.substring(0, path.length - 1) : path
 
 	return path
 }
+
+export async function list_files_in_directory(dirPath: string): Promise<string[]> {
+	const files: string[] = []
+	try {
+		const dir = await fs.readdir(path.join(process.cwd(), dirPath), {
+			withFileTypes: true,
+		})
+		for (const dirent of dir) {
+			const relPath = path.relative(process.cwd(), `${dirPath}/${dirent.name}`)
+			files.push(`file:///${relPath}`)
+		}
+	} catch (error) {
+		throw new Error('Failed to read directory')
+	}
+
+	return files
+}
+
+export async function get_mime_type(path: string): Promise<string | null> {
+	let mimeType = mime.getType(path)
+	if (mimeType && mimeType === 'inode/directory') {
+		mimeType = DIRECTORY_CONTENT_TYPE
+	}
+	return mimeType
+}

apps/sandbox-container/container/index.ts

@@ -7,8 +7,13 @@ import { Hono } from 'hono'
 import { streamText } from 'hono/streaming'
 import mime from 'mime'
 
-import { ExecParams, FilesWrite } from '../shared/schema'
-import { get_file_name_from_path } from './fileUtils'
+import { ExecParams, FilesWrite } from '../shared/schema.ts'
+import {
+	DIRECTORY_CONTENT_TYPE,
+	get_file_name_from_path,
+	get_mime_type,
+	list_files_in_directory,
+} from './fileUtils.ts'
 
 import type { FileList } from '../shared/schema.ts'
 
@@ -65,32 +70,18 @@ app.get('/files/ls', async (c) => {
 app.get('/files/contents/*', async (c) => {
 	const reqPath = await get_file_name_from_path(c.req.path)
 	try {
-		const mimeType = mime.getType(reqPath)
+		const mimeType = await get_mime_type(reqPath)
 		const headers = mimeType ? { 'Content-Type': mimeType } : undefined
 		const contents = await fs.readFile(path.join(process.cwd(), reqPath))
 		return c.newResponse(contents, 200, headers)
 	} catch (e: any) {
 		if (e.code) {
-			// handle directory
 			if (e.code === 'EISDIR') {
-				const files: string[] = []
-				const dir = await fs.readdir(path.join(process.cwd(), reqPath), {
-					withFileTypes: true,
-				})
-				for (const dirent of dir) {
-					const relPath = path.relative(process.cwd(), `${reqPath}/${dirent.name}`)
-					if (dirent.isDirectory()) {
-						files.push(`file:///${relPath}`)
-					} else {
-						const mimeType = mime.getType(dirent.name)
-						files.push(`file:///${relPath}`)
-					}
-				}
+				const files = await list_files_in_directory(reqPath)
 				return c.newResponse(files.join('\n'), 200, {
-					'Content-Type': 'inode/directory',
+					'Content-Type': DIRECTORY_CONTENT_TYPE,
 				})
 			}
-
 			if (e.code === 'ENOENT') {
 				return c.notFound()
 			}

apps/sandbox-container/package.json

@@ -40,6 +40,8 @@
 		"zod": "^3.24.2"
 	},
 	"devDependencies": {
+		"@types/mock-fs": "^4.13.4",
+		"mock-fs": "^5.5.0",
 		"@cloudflare/vitest-pool-workers": "0.8.14",
 		"ai": "^4.3.6",
 		"concurrently": "^9.1.2",

apps/sandbox-container/server/containerMcp.ts

@@ -1,13 +1,12 @@
 import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
 import { McpAgent } from 'agents/mcp'
-import { z } from 'zod'
 
 import { OPEN_CONTAINER_PORT } from '../shared/consts'
 import { ExecParams, FilePathParam, FilesWrite } from '../shared/schema'
 import { MAX_CONTAINERS, proxyFetch, startAndWaitForPort } from './containerHelpers'
 import { getContainerManager } from './containerManager'
 import { BASE_INSTRUCTIONS } from './prompts'
-import { fileToBase64 } from './utils'
+import { fileToBase64, stripProtocolFromFilePath } from './utils'
 
 import type { FileList } from '../shared/schema'
 import type { Env, Props } from '.'
@@ -77,7 +76,8 @@ export class ContainerMcpAgent extends McpAgent<Env, Props> {
 			'Delete file and its contents',
 			{ args: FilePathParam },
 			async ({ args }) => {
-				const deleted = await this.container_file_delete(args)
+				const path = await stripProtocolFromFilePath(args.path)
+				const deleted = await this.container_file_delete(path)
 				return {
 					content: [{ type: 'text', text: `File deleted: ${deleted}.` }],
 				}
@@ -88,6 +88,7 @@ export class ContainerMcpAgent extends McpAgent<Env, Props> {
 			'Write file contents',
 			{ args: FilesWrite },
 			async ({ args }) => {
+				args.path = await stripProtocolFromFilePath(args.path)
 				return {
 					content: [{ type: 'text', text: await this.container_files_write(args) }],
 				}
@@ -117,19 +118,13 @@ export class ContainerMcpAgent extends McpAgent<Env, Props> {
 		})
 		this.server.tool(
 			'container_file_read',
-			'Read a specific file',
-			{ path: z.string() },
-			async ({ path }) => {
-				// normalize
-				path = path.startsWith('file://') ? path.replace('file://', '') : path
-				let { blob, mimeType } = await this.container_files_read(path)
-
-				if (mimeType && (mimeType.startsWith('text') || mimeType === 'inode/directory')) {
-					// this is because there isn't a "real" directory mime type, so we're reusing the "text/directory" mime type
-					// so claude doesn't give an error
-					mimeType = mimeType === 'inode/directory' ? 'text/directory' : mimeType
+			'Read a specific file or list of files within a specific directory',
+			{ args: FilePathParam },
+			async ({ args }) => {
+				const path = await stripProtocolFromFilePath(args.path)
+				const { blob, mimeType } = await this.container_file_read(path)
 
-					// maybe "inode/directory" should list out multiple files in the contents list?
+				if (mimeType && mimeType.startsWith('text')) {
 					return {
 						content: [
 							{
@@ -253,7 +248,7 @@ export class ContainerMcpAgent extends McpAgent<Env, Props> {
 		)
 		return res.ok
 	}
-	async container_files_read(
+	async container_file_read(
 		filePath: string
 	): Promise<{ blob: Blob; mimeType: string | undefined }> {
 		const res = await proxyFetch(
@@ -267,15 +262,11 @@ export class ContainerMcpAgent extends McpAgent<Env, Props> {
 		}
 		return {
 			blob: await res.blob(),
-			mimeType: res.headers.get('content-type') ?? undefined,
+			mimeType: res.headers.get('Content-Type') ?? undefined,
 		}
 	}
 
 	async container_files_write(file: FilesWrite): Promise<string> {
-		if (file.path.startsWith('file://')) {
-			// normalize just in case the LLM sends the full resource URI
-			file.path = file.path.replace('file://', '')
-		}
 		const res = await proxyFetch(
 			this.env.ENVIRONMENT,
 			this.ctx.container,

apps/sandbox-container/server/utils.spec.ts

@@ -0,0 +1,14 @@
+import { describe, expect, it } from 'vitest'
+
+import { stripProtocolFromFilePath } from './utils'
+
+describe('get_file_name_from_path', () => {
+	it('strips file:// protocol from path', async () => {
+		const path = await stripProtocolFromFilePath('file:///files/contents/cats')
+		expect(path).toBe('/files/contents/cats')
+	}),
+		it('leaves protocol-less paths untouched', async () => {
+			const path = await stripProtocolFromFilePath('/files/contents/cats')
+			expect(path).toBe('/files/contents/cats')
+		})
+})

apps/sandbox-container/server/utils.ts

@@ -12,3 +12,8 @@ export async function fileToBase64(blob: Blob): Promise<string> {
 	// Apply base64 encoding
 	return btoa(binary)
 }
+
+// Used for file related tool calls in case the llm sends a full resource URI
+export async function stripProtocolFromFilePath(path: string): Promise<string> {
+	return path.startsWith('file://') ? path.replace('file://', '') : path
+}

apps/sandbox-container/shared/schema.ts

@@ -14,7 +14,9 @@ export const FilesWrite = z.object({
 })
 
 export type FilePathParam = z.infer<typeof FilePathParam>
-export const FilePathParam = z.string()
+export const FilePathParam = z.object({
+	path: z.string(),
+})
 
 export type FileList = z.infer<typeof FileList>
 export const FileList = z.object({