Replace the code interpreter

FROM node:24-trixie-slim

RUN apt-get update \
	&& apt-get install -y --no-install-recommends \
		ca-certificates git python3 \
		python3-ipython python3-matplotlib \
		python3-numpy python3-pandas \
	&& rm -rf /var/lib/apt/lists/*

COPY --from=docker.io/cloudflare/sandbox:1.0.0 /usr/local/bin/sandbox-shim /usr/local/bin/sandbox-shim

WORKDIR /workspace
COPY kernel.py /opt/kernel.py
CMD ["sleep", "infinity"]

import base64
import io
import json
import os
import signal
import sys
import traceback
from http.server import BaseHTTPRequestHandler
from socketserver import TCPServer

os.environ["MPLBACKEND"] = "Agg"

from IPython.core.interactiveshell import InteractiveShell
from IPython.utils.capture import capture_output

PORT = int(sys.argv[1])
FORMATS = {
    "text/plain": "text",
    "text/html": "html",
    "text/markdown": "markdown",
    "text/latex": "latex",
    "image/png": "png",
    "image/jpeg": "jpeg",
    "image/svg+xml": "svg",
    "application/json": "json",
    "application/javascript": "javascript",
}

shell = InteractiveShell.instance(colors="NoColor")
# Return values and errors in the response, not in stdout.
shell.displayhook.write_output_prompt = lambda: None
shell.displayhook.write_format_data = lambda *a, **k: None
shell.showtraceback = lambda *a, **k: None
shell.showsyntaxerror = lambda *a, **k: None
running = False


def interrupt(signum, frame):
    # Stop a running cell. An idle kernel ignores the signal.
    if running:
        raise KeyboardInterrupt


def formats(data):
    return {
        FORMATS[kind]: value
        for kind, value in data.items()
        if kind in FORMATS
    }


def figures():
    if "matplotlib.pyplot" not in sys.modules:
        return []
    import matplotlib.pyplot as plt

    images = []
    for number in plt.get_fignums():
        buffer = io.BytesIO()
        plt.figure(number).savefig(
            buffer, format="png", bbox_inches="tight"
        )
        png = base64.b64encode(buffer.getvalue()).decode()
        images.append({"png": png})
    plt.close("all")
    return images


def describe(error):
    if error is None:
        return None
    if isinstance(error, SyntaxError):
        lines = traceback.format_exception_only(error)
    else:
        # Skip the frame of the kernel itself.
        lines = traceback.format_exception(
            type(error), error, error.__traceback__.tb_next
        )
    return {
        "name": type(error).__name__,
        "message": str(error),
        "traceback": lines,
    }


def run(code):
    global running
    running = True
    try:
        with capture_output(display=True) as captured:
            cell = shell.run_cell(code, store_history=True)
    finally:
        running = False

    results = [formats(output.data) for output in captured.outputs]
    results += figures()
    if isinstance(cell.result, (dict, list)):
        results.append({"json": cell.result})
    elif cell.result is not None:
        data, _ = shell.display_formatter.format(cell.result)
        results.append(formats(data))

    stdout, stderr = captured.stdout, captured.stderr
    return {
        "logs": {
            "stdout": [stdout] if stdout else [],
            "stderr": [stderr] if stderr else [],
        },
        "results": results,
        "error": describe(
            cell.error_before_exec or cell.error_in_exec
        ),
        "executionCount": cell.execution_count,
    }


class Kernel(BaseHTTPRequestHandler):
    def do_GET(self):
        self.reply(b"ready")

    def do_POST(self):
        length = int(self.headers["Content-Length"])
        code = self.rfile.read(length).decode()
        result = json.dumps(run(code), default=repr)
        self.reply(result.encode())

    def reply(self, body):
        self.send_response(200)
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, *args):
        pass


signal.signal(signal.SIGINT, interrupt)
# HTTPServer looks up the container hostname, which is too
# long to resolve.
TCPServer.allow_reuse_address = True
TCPServer(("0.0.0.0", PORT), Kernel).serve_forever()

type Execution = {
	logs: { stdout: string[]; stderr: string[] };
	// One object for each value, keyed by format.
	results: Record<string, string | object>[];
	error: {
		name: string;
		message: string;
		traceback: string[];
	} | null;
	executionCount: number;
};

type CodeContext = {
	id: string;
	port: number;
	cwd: string;
	env?: Record<string, string>;
};

const DEFAULT_CONTEXT: CodeContext = {
	id: "default",
	port: 8800,
	cwd: "/workspace",
};

// Each kernel runs as a server in its own directory.
const kernelDir = (context: CodeContext) =>
	`${SERVERS}/kernel-${context.port}`;

export class MySandbox extends DurableObject<Env> {
	// ...

	// Ports of the kernels that answered in this container.
	private kernels = new Set<number>();

	private async startKernel(
		context: CodeContext,
	): Promise<Fetcher> {
		if (this.kernels.has(context.port)) {
			return this.container.getTcpPort(context.port);
		}

		await startServer(
			this.container,
			kernelDir(context),
			["python3", "/opt/kernel.py", String(context.port)],
			{ cwd: context.cwd, env: { ...ENV, ...context.env } },
		);
		const kernel = await waitForServer(
			this.container,
			kernelDir(context),
			context.port,
			30_000,
		);
		this.kernels.add(context.port);
		return kernel;
	}

	private async interruptKernel(context: CodeContext): Promise<void> {
		const kill = await this.container.exec([
			"sh",
			"-c",
			`${CURRENT}\ncurrent "$1" && kill -INT "$pid"`,
			"sh",
			kernelDir(context),
		]);
		await kill.exitCode;
	}
}

export class MySandbox extends DurableObject<Env> {
	// ...

	async runCode(
		code: string,
		options: { context?: string; timeout?: number } = {},
	): Promise<Execution> {
		await this.ensureRunning();
		const context = await this.getContext(
			options.context ?? DEFAULT_CONTEXT.id,
		);
		const kernel = await this.startKernel(context);

		try {
			const response = await kernel.fetch("http://kernel/", {
				method: "POST",
				body: code,
				signal: AbortSignal.timeout(options.timeout ?? 60_000),
			});
			return await response.json<Execution>();
		} catch (error) {
			if ((error as Error).name === "TimeoutError") {
				// Stop the cell, and keep the kernel variables.
				await this.interruptKernel(context);
			} else {
				// The kernel exited. The next call starts a new one.
				this.kernels.delete(context.port);
			}
			throw error;
		}
	}
}

export class MySandbox extends DurableObject<Env> {
	// ...

	private async getContext(id: string): Promise<CodeContext> {
		if (id === DEFAULT_CONTEXT.id) {
			return DEFAULT_CONTEXT;
		}

		const context = await this.ctx.storage.get<CodeContext>(
			`context:${id}`,
		);

		if (!context) {
			throw new Error(`No code context ${id}`);
		}

		return context;
	}

	async createCodeContext(
		options: { cwd?: string; env?: Record<string, string> } = {},
	): Promise<CodeContext> {
		// Take the lowest port that no stored context uses.
		const stored = await this.ctx.storage.list<CodeContext>({
			prefix: "context:",
		});
		const used = new Set([...stored.values()].map((c) => c.port));
		let port = DEFAULT_CONTEXT.port + 1;

		while (used.has(port)) {
			port++;
		}

		const context: CodeContext = {
			id: crypto.randomUUID(),
			port,
			cwd: options.cwd ?? "/workspace",
			env: options.env,
		};

		await this.ctx.storage.put(`context:${context.id}`, context);
		return context;
	}

	async listCodeContexts(): Promise<CodeContext[]> {
		const contexts = await this.ctx.storage.list<CodeContext>({
			prefix: "context:",
		});
		return [DEFAULT_CONTEXT, ...contexts.values()];
	}

	async deleteCodeContext(id: string): Promise<void> {
		const context = await this.getContext(id);

		if (this.container.running) {
			await stopServer(this.container, kernelDir(context));
			this.kernels.delete(context.port);
		}

		await this.ctx.storage.delete(`context:${id}`);
	}
}

const sandbox = getSandbox(env.Sandbox, "ada");
const context = await sandbox.createCodeContext({
	language: "python",
	envVars: { REGION: "east" },
});
const execution = await sandbox.runCode(
	"import os; os.environ['REGION']",
	{ context, timeout: 30_000 },
);

const sandbox = env.Sandbox.getByName("ada");
const context = await sandbox.createCodeContext({
	env: { REGION: "east" },
});
const execution = await sandbox.runCode(
	"import os; os.environ['REGION']",
	{ context: context.id, timeout: 30_000 },
);

const execution = await sandbox.runCode(
	"[2, 3, 5, 7].reduce((a, b) => a + b)",
	{ language: "javascript" },
);

import { runJavaScript } from "./sandbox";

const output = await runJavaScript(
	env,
	"return [2, 3, 5, 7].reduce((a, b) => a + b);",
);