| 1 | import { DurableObject } from "cloudflare:workers"; |
| 2 | import { Container } from "@cloudflare/containers"; |
| 3 | import type { StopParams } from "@cloudflare/containers"; |
| 4 | import type { Env } from "../env.js"; |
| 5 | import { endCompute, startCompute } from "../lib/quotas.js"; |
| 6 | import { blockedBy, killSwitches } from "../lib/killswitch.js"; |
| 7 | |
| 8 | /** |
| 9 | * One container per workflow job. The runner agent inside fetches its job spec |
| 10 | * from the Diffinery API, runs the steps and reports back, then exits. |
| 11 | */ |
| 12 | export class RunnerContainer extends Container<Env> { |
| 13 | // Jobs make no inbound requests, so activity-based sleep must not stop them. |
| 14 | override sleepAfter = "6h"; |
| 15 | override enableInternet = true; |
| 16 | |
| 17 | async startJob(input: { repoId: string; jobId: string; token: string; apiUrl: string }): Promise<void> { |
| 18 | await this.ctx.storage.put("job", { repoId: input.repoId, jobId: input.jobId }); |
| 19 | await this.start({ |
| 20 | envVars: { |
| 21 | DIFFINERY_API_URL: input.apiUrl, |
| 22 | DIFFINERY_JOB_ID: input.jobId, |
| 23 | DIFFINERY_JOB_TOKEN: input.token, |
| 24 | }, |
| 25 | }); |
| 26 | } |
| 27 | |
| 28 | async stopJob(): Promise<void> { |
| 29 | // SIGINT cancels the job at once. The agent reads SIGTERM as a platform rollout and |
| 30 | // keeps the job running for up to 14 minutes, which a cancelled job must not do. |
| 31 | if (this.ctx.container?.running) await this.stop("SIGINT"); |
| 32 | } |
| 33 | |
| 34 | override async onStop(params: StopParams): Promise<void> { |
| 35 | const job = await this.ctx.storage.get<{ repoId: string; jobId: string }>("job"); |
| 36 | if (!job) return; |
| 37 | const repo = this.env.REPO.get(this.env.REPO.idFromName(job.repoId)); |
| 38 | await repo.containerStopped(job.jobId, params.exitCode ?? null); |
| 39 | await this.ctx.storage.delete("job"); |
| 40 | } |
| 41 | |
| 42 | override onError(error: unknown): void { |
| 43 | console.error("runner container error", error); |
| 44 | } |
| 45 | } |
| 46 | |
| 47 | /** |
| 48 | * Long-lived task sandbox for AI agents: the repo checked out with toolchains, |
| 49 | * and an HTTP API (exec, read/write files) on port 8080. |
| 50 | */ |
| 51 | export class SandboxContainer extends Container<Env> { |
| 52 | override defaultPort = 8080; |
| 53 | override sleepAfter = "30m"; |
| 54 | override enableInternet = true; |
| 55 | |
| 56 | async boot(envVars: Record<string, string>, secret: string, _name?: string, usage?: { id: string; accountId: string; repoId: string }): Promise<void> { |
| 57 | await this.ctx.storage.put("secret", secret); |
| 58 | // Its running time counts against the repository owner's hosted minutes (lib/quotas.ts). |
| 59 | if (usage) await this.ctx.storage.put("usage", usage); |
| 60 | // The first start on a host after a deploy pulls the runner image, which takes |
| 61 | // well over the library's 8 s / 20 s defaults. |
| 62 | await this.startAndWaitForPorts({ |
| 63 | ports: 8080, |
| 64 | startOptions: { envVars, entrypoint: ["node", "/opt/diffinery/agent.mjs", "--sandbox"] }, |
| 65 | cancellationOptions: { instanceGetTimeoutMS: 180_000, portReadyTimeoutMS: 120_000, waitInterval: 500 }, |
| 66 | }); |
| 67 | } |
| 68 | |
| 69 | /** Forward an API call (JSON text) to the sandbox's HTTP server with its secret attached. */ |
| 70 | async call(path: string, body: string): Promise<{ status: number; body: string; json: boolean }> { |
| 71 | const secret = await this.ctx.storage.get<string>("secret"); |
| 72 | // A sandbox that slept after 30 idle minutes starts again for this call: its time counts again. |
| 73 | const usage = await this.ctx.storage.get<{ id: string; accountId: string; repoId: string; part?: number }>("usage"); |
| 74 | if (usage && !this.ctx.container?.running) { |
| 75 | // It doesn't wake while a kill switch stops sandboxes (lib/killswitch.ts). |
| 76 | const block = await killSwitches(this.env, usage.repoId).then( |
| 77 | (k) => blockedBy(k, "sandbox"), |
| 78 | () => null, |
| 79 | ); |
| 80 | if (block) return { status: 403, body: JSON.stringify({ message: block }), json: true }; |
| 81 | const part = (usage.part ?? 0) + 1; |
| 82 | await startCompute(this.env, { id: `${usage.id}:${part}`, accountId: usage.accountId, repoId: usage.repoId, kind: "sandbox" }); |
| 83 | await this.ctx.storage.put("usage", { ...usage, part }); |
| 84 | } |
| 85 | const res = await this.containerFetch(`http://sandbox${path}`, { |
| 86 | method: "POST", |
| 87 | headers: { "content-type": "application/json", authorization: `Bearer ${secret}` }, |
| 88 | body: body || "{}", |
| 89 | }); |
| 90 | return { status: res.status, body: await res.text(), json: (res.headers.get("content-type") ?? "").includes("json") }; |
| 91 | } |
| 92 | |
| 93 | async shutdown(): Promise<void> { |
| 94 | if (this.ctx.container?.running) await this.stop(); |
| 95 | } |
| 96 | |
| 97 | /** Whether the container is up (calling into a stopped one would start it without its checkout). */ |
| 98 | async running(): Promise<boolean> { |
| 99 | return !!this.ctx.container?.running && !!(await this.ctx.storage.get<string>("secret")); |
| 100 | } |
| 101 | |
| 102 | /** |
| 103 | * A web terminal's socket (agentic/terminal-web.ts checked who is asking): only a WebSocket |
| 104 | * upgrade to /pty, with the sandbox's secret added. A sandbox that isn't running isn't started |
| 105 | * again (it would come up without its checkout), and nothing else of the sandbox is reachable here. |
| 106 | */ |
| 107 | override async fetch(request: Request): Promise<Response> { |
| 108 | const url = new URL(request.url); |
| 109 | if (url.pathname !== "/pty" || request.headers.get("upgrade")?.toLowerCase() !== "websocket") return new Response("not found", { status: 404 }); |
| 110 | const secret = await this.ctx.storage.get<string>("secret"); |
| 111 | if (!secret || !this.ctx.container?.running) return new Response("the terminal isn't running", { status: 410 }); |
| 112 | const headers = new Headers({ upgrade: "websocket", authorization: `Bearer ${secret}` }); |
| 113 | return this.containerFetch(new Request(`http://sandbox/pty${url.search}`, { headers }), 8080); |
| 114 | } |
| 115 | |
| 116 | override async onStop(): Promise<void> { |
| 117 | const usage = await this.ctx.storage.get<{ id: string; part?: number }>("usage"); |
| 118 | if (usage) await endCompute(this.env, usage.part ? `${usage.id}:${usage.part}` : usage.id); |
| 119 | } |
| 120 | } |
| 121 | |
| 122 | /** |
| 123 | * Local development and end-to-end tests only: `env.dev` binds SANDBOX to this |
| 124 | * instead of SandboxContainer, because containers don't run under `wrangler dev`. |
| 125 | * The same calls, served by one `agent.mjs --sandbox-pool` process at |
| 126 | * SANDBOX_POOL_URL that runs each sandbox's server. |
| 127 | */ |
| 128 | export class SandboxPoolDO extends DurableObject<Env> { |
| 129 | private poolAuth() { |
| 130 | return { authorization: `Bearer ${this.env.SANDBOX_POOL_SECRET ?? ""}` }; |
| 131 | } |
| 132 | |
| 133 | async boot(envVars: Record<string, string>, secret: string, name?: string): Promise<void> { |
| 134 | if (!this.env.SANDBOX_POOL_URL || !name) throw new Error("SANDBOX_POOL_URL is not set: start `agent.mjs --sandbox-pool` for sandboxes under wrangler dev"); |
| 135 | await this.ctx.storage.put({ secret, name }); |
| 136 | const r = await fetch(`${this.env.SANDBOX_POOL_URL}/boot`, { method: "POST", headers: this.poolAuth(), body: JSON.stringify({ id: name, secret, env: envVars }) }); |
| 137 | if (!r.ok) throw new Error(`sandbox pool boot failed: ${r.status} ${await r.text()}`); |
| 138 | } |
| 139 | |
| 140 | async call(path: string, body: string): Promise<{ status: number; body: string; json: boolean }> { |
| 141 | const [secret, name] = await Promise.all([this.ctx.storage.get<string>("secret"), this.ctx.storage.get<string>("name")]); |
| 142 | if (!name) return { status: 404, body: JSON.stringify({ error: "sandbox not found" }), json: true }; |
| 143 | const res = await fetch(`${this.env.SANDBOX_POOL_URL}/s/${name}${path}`, { method: "POST", headers: { "content-type": "application/json", authorization: `Bearer ${secret}` }, body: body || "{}" }); |
| 144 | return { status: res.status, body: await res.text(), json: (res.headers.get("content-type") ?? "").includes("json") }; |
| 145 | } |
| 146 | |
| 147 | async shutdown(): Promise<void> { |
| 148 | const name = await this.ctx.storage.get<string>("name"); |
| 149 | if (name) await fetch(`${this.env.SANDBOX_POOL_URL}/shutdown/${name}`, { method: "POST", headers: this.poolAuth() }).catch(() => {}); |
| 150 | } |
| 151 | |
| 152 | async running(): Promise<boolean> { |
| 153 | if (!(await this.ctx.storage.get<string>("name"))) return false; |
| 154 | return (await this.call("/pty/status", "{}").catch(() => ({ status: 0 }))).status === 200; |
| 155 | } |
| 156 | |
| 157 | /** As SandboxContainer.fetch: a terminal's WebSocket, through the pool. */ |
| 158 | override async fetch(request: Request): Promise<Response> { |
| 159 | const url = new URL(request.url); |
| 160 | if (url.pathname !== "/pty" || request.headers.get("upgrade")?.toLowerCase() !== "websocket") return new Response("not found", { status: 404 }); |
| 161 | const [secret, name] = await Promise.all([this.ctx.storage.get<string>("secret"), this.ctx.storage.get<string>("name")]); |
| 162 | if (!secret || !name) return new Response("the terminal isn't running", { status: 410 }); |
| 163 | const res = await fetch(`${this.env.SANDBOX_POOL_URL}/s/${name}/pty${url.search}`, { headers: { upgrade: "websocket", authorization: `Bearer ${secret}` } }); |
| 164 | if (!res.webSocket) return new Response(await res.text(), { status: res.status === 101 ? 502 : res.status }); |
| 165 | return pipeSocket(res.webSocket); |
| 166 | } |
| 167 | } |
| 168 | |
| 169 | /** Hands a WebSocket this Durable Object holds to the caller, relaying frames and closes both ways. */ |
| 170 | function pipeSocket(upstream: WebSocket): Response { |
| 171 | const [client, server] = Object.values(new WebSocketPair()) as [WebSocket, WebSocket]; |
| 172 | upstream.accept(); |
| 173 | server.accept(); |
| 174 | const code = (c: number) => (c === 1005 || c === 1006 ? 1000 : c); |
| 175 | // Binary frames can arrive as Blobs; reading one is async, so each direction keeps its order with a chain. |
| 176 | const relay = (from: WebSocket, to: WebSocket) => { |
| 177 | let chain = Promise.resolve(); |
| 178 | from.addEventListener("message", (e) => { |
| 179 | const data = e.data as string | ArrayBuffer | Blob; |
| 180 | chain = chain.then(async () => to.send(data instanceof Blob ? await data.arrayBuffer() : data)).catch(() => to.close(1011, "relay failed")); |
| 181 | }); |
| 182 | }; |
| 183 | relay(server, upstream); |
| 184 | relay(upstream, server); |
| 185 | server.addEventListener("close", (e) => upstream.close(code(e.code), e.reason)); |
| 186 | upstream.addEventListener("close", (e) => server.close(code(e.code), e.reason)); |
| 187 | server.addEventListener("error", () => upstream.close(1011, "client error")); |
| 188 | upstream.addEventListener("error", () => server.close(1011, "sandbox error")); |
| 189 | return new Response(null, { status: 101, webSocket: client }); |
| 190 | } |