Add user based container

Author: cmsparks

apps/sandbox-container/server/containerHelpers.ts

@@ -1,4 +1,4 @@
-export const MAX_CONTAINERS = 8
+export const MAX_CONTAINERS = 20
 export async function startAndWaitForPort(
 	environment: 'dev' | 'prod' | 'test',
 	container: Container | undefined,

apps/sandbox-container/server/containerManager.ts

@@ -28,8 +28,8 @@ export class ContainerManager extends DurableObject<Env> {
 			console.log(id, time, now, now.valueOf() - time.valueOf())
 
 			if (now.valueOf() - time.valueOf() > 10 * 60 * 1000) {
-				const doId = this.env.CONTAINER_MCP_AGENT.idFromString(id)
-				const stub = this.env.CONTAINER_MCP_AGENT.get(doId)
+				const doId = this.env.USER_CONTAINER.idFromString(id)
+				const stub = this.env.USER_CONTAINER.get(doId)
 				await stub.destroyContainer()
 				await this.killContainer(id)
 			}

apps/sandbox-container/server/containerMcp.ts

@@ -2,16 +2,12 @@ import { McpAgent } from 'agents/mcp'
 
 import { CloudflareMCPServer } from '@repo/mcp-common/src/server'
 
-import { OPEN_CONTAINER_PORT } from '../shared/consts'
 import { ExecParams, FilePathParam, FileWrite } from '../shared/schema'
-import { MAX_CONTAINERS, proxyFetch, startAndWaitForPort } from './containerHelpers'
-import { getContainerManager } from './containerManager'
 import { BASE_INSTRUCTIONS } from './prompts'
 import { fileToBase64, stripProtocolFromFilePath } from './utils'
 
-import type { FileList } from '../shared/schema'
 import type { Env } from './context'
-import type { Props } from '.'
+import type { Props, UserContainer } from '.'
 
 export class ContainerMcpAgent extends McpAgent<Env, {}, Props> {
 	_server: CloudflareMCPServer | undefined
@@ -27,6 +23,11 @@ export class ContainerMcpAgent extends McpAgent<Env, {}, Props> {
 		return this._server
 	}
 
+	get userContainer(): DurableObjectStub<UserContainer> {
+		const userContainer = this.env.USER_CONTAINER.idFromName(this.props.user.id)
+		return this.env.USER_CONTAINER.get(userContainer)
+	}
+
 	constructor(
 		public ctx: DurableObjectState,
 		public env: Env
@@ -35,21 +36,6 @@ export class ContainerMcpAgent extends McpAgent<Env, {}, Props> {
 		super(ctx, env)
 	}
 
-	async destroyContainer(): Promise<void> {
-		await this.ctx.container?.destroy()
-	}
-
-	async killContainer(): Promise<void> {
-		console.log('Reaping container')
-		const containerManager = getContainerManager(this.env)
-		const active = await containerManager.listActive()
-		if (this.ctx.id.toString() in active) {
-			console.log('killing container')
-			await this.destroyContainer()
-			await containerManager.killContainer(this.ctx.id.toString())
-		}
-	}
-
 	async init() {
 		this.props.user.id
 		this.server = new CloudflareMCPServer({
@@ -64,78 +50,63 @@ export class ContainerMcpAgent extends McpAgent<Env, {}, Props> {
 
 		this.server.tool(
 			'container_initialize',
-			'Start or reset the container',
-			{},
+			`Start or restart the container. 
+			Use this tool to initialize a container before running any python or node.js code that the user requests ro run.`,
 			// @ts-ignore
-			async ({}) => {
+			async () => {
+				const userInBlocklist = await this.env.USER_BLOCKLIST.get(this.props.user.id)
+				if (userInBlocklist) {
+					return {
+						content: [{ type: 'text', text: "Blocked from intializing container." }],
+					}
+				}
 				return {
-					content: [{ type: 'text', text: await this.container_initialize() }],
+					content: [{ type: 'text', text: await this.userContainer.container_initialize() }],
 				}
 			}
 		)
 
-		this.server.tool('container_ping', 'Ping the container for liveliness', {}, async ({}) => {
+		this.server.tool('container_ping', `Ping the container for liveliness. Use this tool to check if the container is running.`, {}, async ({}) => {
 			return {
-				content: [{ type: 'text', text: await this.container_ping() }],
+				content: [{ type: 'text', text: await this.userContainer.container_ping() }],
 			}
 		})
 		this.server.tool(
 			'container_exec',
-			'Run a command in a container and return the results from stdout',
+			'Run a command in a container and return the results from stdout. If necessary, set a timeout. To debug, stream back standard error.',
 			{ args: ExecParams },
 			async ({ args }) => {
 				return {
-					content: [{ type: 'text', text: await this.container_exec(args) }],
+					content: [{ type: 'text', text: await this.userContainer.container_exec(args) }],
 				}
 			}
 		)
 		this.server.tool(
 			'container_file_delete',
-			'Delete file and its contents',
+			'Delete file in the working directory',
 			{ args: FilePathParam },
 			async ({ args }) => {
 				const path = await stripProtocolFromFilePath(args.path)
-				const deleted = await this.container_file_delete(path)
+				const deleted = await this.userContainer.container_file_delete(path)
 				return {
 					content: [{ type: 'text', text: `File deleted: ${deleted}.` }],
 				}
 			}
 		)
 		this.server.tool(
 			'container_file_write',
-			'Create a new file with the provided contents, overwriting the file if it already exists',
+			'Create a new file with the provided contents in the working direcotry, overwriting the file if it already exists',
 			{ args: FileWrite },
 			async ({ args }) => {
 				args.path = await stripProtocolFromFilePath(args.path)
 				return {
-					content: [{ type: 'text', text: await this.container_file_write(args) }],
+					content: [{ type: 'text', text: await this.userContainer.container_file_write(args) }],
 				}
 			}
 		)
-		this.server.tool('container_files_list', 'List working directory file tree', {}, async ({}) => {
-			// This approach relies on resources, which aren't handled well by Claude right now. Until that's sorted, we can just use file read, since it lists all files in a directory if a directory is passed to it.
-			//const files = await this.container_ls()
-
-			// const resources: {
-			// 	type: 'resource'
-			// 	resource: { uri: string; text: string; mimeType: string }
-			// }[] = files.resources.map((r) => {
-			// 	return {
-			// 		type: 'resource',
-			// 		resource: {
-			// 			uri: r.uri,
-			// 			text: r.uri,
-			// 			mimeType: 'text/plain',
-			// 		},
-			// 	}
-			// })
-
-			// return {
-			// 	content: resources,
-			// }
-
+		this.server.tool('container_files_list', 'List working directory file tree. This just reads the contents of the current working directory', {}, async ({}) => {
 			// Begin workaround using container read rather than ls:
-			const { blob, mimeType } = await this.container_file_read('.')
+			const { blob, mimeType } = await this.userContainer.container_file_read('.')
 			return {
 				content: [
 					{
@@ -155,7 +126,7 @@ export class ContainerMcpAgent extends McpAgent<Env, {}, Props> {
 			{ args: FilePathParam },
 			async ({ args }) => {
 				const path = await stripProtocolFromFilePath(args.path)
-				const { blob, mimeType } = await this.container_file_read(path)
+				const { blob, mimeType } = await this.userContainer.container_file_read(path)
 
 				if (mimeType && mimeType.startsWith('text')) {
 					return {
@@ -176,6 +147,8 @@ export class ContainerMcpAgent extends McpAgent<Env, {}, Props> {
 							{
 								type: 'resource',
 								resource: {
+									// for some reason the RPC type for Blob is not exactly the same as the regular Blob type
+									// @ts-ignore
 									blob: await fileToBase64(blob),
 									uri: `file://${path}`,
 									mimeType: mimeType,
@@ -187,134 +160,4 @@ export class ContainerMcpAgent extends McpAgent<Env, {}, Props> {
 			}
 		)
 	}
-
-	async container_initialize(): Promise<string> {
-		// kill container
-		await this.killContainer()
-
-		// try to cleanup cleanup old containers
-		const containerManager = getContainerManager(this.env)
-
-		if ((await containerManager.listActive()).length >= MAX_CONTAINERS) {
-			await containerManager.tryKillOldContainers()
-			if ((await containerManager.listActive()).length >= MAX_CONTAINERS) {
-				throw new Error(
-					`Unable to reap enough containers. There are ${MAX_CONTAINERS} active container sandboxes, please wait`
-				)
-			}
-		}
-
-		// start container
-		let startedContainer = false
-		await this.ctx.blockConcurrencyWhile(async () => {
-			startedContainer = await startAndWaitForPort(
-				this.env.ENVIRONMENT,
-				this.ctx.container,
-				OPEN_CONTAINER_PORT
-			)
-		})
-		if (!startedContainer) {
-			throw new Error('Failed to start container')
-		}
-
-		// track and manage lifecycle
-		containerManager.trackContainer(this.ctx.id.toString())
-
-		return `Created new container`
-	}
-
-	async container_ping(): Promise<string> {
-		const res = await proxyFetch(
-			this.env.ENVIRONMENT,
-			this.ctx.container,
-			new Request(`http://host:${OPEN_CONTAINER_PORT}/ping`),
-			OPEN_CONTAINER_PORT
-		)
-		if (!res || !res.ok) {
-			throw new Error(`Request to container failed: ${await res.text()}`)
-		}
-		return await res.text()
-	}
-
-	async container_exec(params: ExecParams): Promise<string> {
-		const res = await proxyFetch(
-			this.env.ENVIRONMENT,
-			this.ctx.container,
-			new Request(`http://host:${OPEN_CONTAINER_PORT}/exec`, {
-				method: 'POST',
-				body: JSON.stringify(params),
-				headers: {
-					'content-type': 'application/json',
-				},
-			}),
-			OPEN_CONTAINER_PORT
-		)
-		if (!res || !res.ok) {
-			throw new Error(`Request to container failed: ${await res.text()}`)
-		}
-		const txt = await res.text()
-		return txt
-	}
-
-	async container_ls(): Promise<FileList> {
-		const res = await proxyFetch(
-			this.env.ENVIRONMENT,
-			this.ctx.container,
-			new Request(`http://host:${OPEN_CONTAINER_PORT}/files/ls`),
-			OPEN_CONTAINER_PORT
-		)
-		if (!res || !res.ok) {
-			throw new Error(`Request to container failed: ${await res.text()}`)
-		}
-		const json = (await res.json()) as FileList
-		return json
-	}
-
-	async container_file_delete(filePath: string): Promise<boolean> {
-		const res = await proxyFetch(
-			this.env.ENVIRONMENT,
-			this.ctx.container,
-			new Request(`http://host:${OPEN_CONTAINER_PORT}/files/contents/${filePath}`, {
-				method: 'DELETE',
-			}),
-			OPEN_CONTAINER_PORT
-		)
-		return res.ok
-	}
-	async container_file_read(
-		filePath: string
-	): Promise<{ blob: Blob; mimeType: string | undefined }> {
-		const res = await proxyFetch(
-			this.env.ENVIRONMENT,
-			this.ctx.container,
-			new Request(`http://host:${OPEN_CONTAINER_PORT}/files/contents/${filePath}`),
-			OPEN_CONTAINER_PORT
-		)
-		if (!res || !res.ok) {
-			throw new Error(`Request to container failed: ${await res.text()}`)
-		}
-		return {
-			blob: await res.blob(),
-			mimeType: res.headers.get('Content-Type') ?? undefined,
-		}
-	}
-
-	async container_file_write(file: FileWrite): Promise<string> {
-		const res = await proxyFetch(
-			this.env.ENVIRONMENT,
-			this.ctx.container,
-			new Request(`http://host:${OPEN_CONTAINER_PORT}/files/contents`, {
-				method: 'POST',
-				body: JSON.stringify(file),
-				headers: {
-					'content-type': 'application/json',
-				},
-			}),
-			OPEN_CONTAINER_PORT
-		)
-		if (!res || !res.ok) {
-			throw new Error(`Request to container failed: ${await res.text()}`)
-		}
-		return `Wrote file: ${file.path}`
-	}
 }

apps/sandbox-container/server/context.ts

@@ -1,4 +1,4 @@
-import type { ContainerManager, ContainerMcpAgent } from './index'
+import type { ContainerManager, ContainerMcpAgent, UserContainer } from './index'
 
 export interface Env {
 	OAUTH_KV: KVNamespace
@@ -10,6 +10,8 @@ export interface Env {
 	OPENAI_API_KEY: string
 	CONTAINER_MCP_AGENT: DurableObjectNamespace<ContainerMcpAgent>
 	CONTAINER_MANAGER: DurableObjectNamespace<ContainerManager>
+	USER_CONTAINER: DurableObjectNamespace<UserContainer>
+	USER_BLOCKLIST: KVNamespace
 	MCP_METRICS: AnalyticsEngineDataset
 	AI: Ai
 	DEV_DISABLE_OAUTH: string

apps/sandbox-container/server/index.ts

@@ -16,8 +16,9 @@ import { ContainerMcpAgent } from './containerMcp'
 import type { McpAgent } from 'agents/mcp'
 import type { AuthProps } from '@repo/mcp-common/src/cloudflare-oauth-handler'
 import type { Env } from './context'
+import { UserContainer } from './userContainer'
 
-export { ContainerManager, ContainerMcpAgent }
+export { ContainerManager, ContainerMcpAgent, UserContainer }
 
 const env = getEnv<Env>()