| 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 { egressStart, type Egress } from "../lib/egress.js"; |
| 6 | import { endCompute, startCompute } from "../lib/quotas.js"; |
| 7 | import { blockedBy, killSwitches } from "../lib/killswitch.js"; |
| 8 | |
| 9 | /** |
| 10 | * Hosted compute (jobs and sandboxes) under the instance's egress policy (lib/egress.ts), set before |
| 11 | * each start. A limited one is kept in storage, so the container starting again from this Durable |
| 12 | * Object with its defaults (a sandbox after its idle sleep) is limited too; the allowed hosts are |
| 13 | * kept by the library itself. |
| 14 | */ |
| 15 | abstract class HostedContainer extends Container<Env> { |
| 16 | override enableInternet = true; |
| 17 | |
| 18 | constructor(ctx: DurableObject["ctx"], env: Env) { |
| 19 | super(ctx, env); |
| 20 | const kept = ctx.storage.kv.get<Egress>("egress"); |
| 21 | if (kept) this.adopt(kept); |
| 22 | } |
| 23 | |
| 24 | /** Applies the policy before a start and returns the variables that tell the agent about it. */ |
| 25 | protected async useEgress(egress: Egress): Promise<Record<string, string>> { |
| 26 | const o = this.adopt(egress); |
| 27 | if (o.allowedHosts) await this.setAllowedHosts(o.allowedHosts); |
| 28 | if (egress !== "open") this.ctx.storage.kv.put("egress", egress); |
| 29 | return o.envVars; |
| 30 | } |
| 31 | |
| 32 | private adopt(egress: Egress) { |
| 33 | const o = egressStart(egress); |
| 34 | this.enableInternet = o.enableInternet; |
| 35 | this.interceptHttps = o.interceptHttps; |
| 36 | return o; |
| 37 | } |
| 38 | } |
| 39 | |
| 40 | /** The job a runner container runs; stopRequested once Diffinery asked it to stop (stopJob). */ |
| 41 | type RunnerJob = { repoId: string; jobId: string; container?: string; stopRequested?: boolean }; |
| 42 | |
| 43 | /** How long a job's container may take to come up, and how often it is checked meanwhile. */ |
| 44 | const START_WAIT_MS = 180_000; |
| 45 | const START_POLL_MS = 500; |
| 46 | |
| 47 | /** |
| 48 | * One container per workflow job. The runner agent inside fetches its job spec |
| 49 | * from the Diffinery API, runs the steps and reports back, then exits. |
| 50 | */ |
| 51 | export class RunnerContainer extends HostedContainer { |
| 52 | // Jobs make no inbound requests, so activity-based sleep must not stop them. |
| 53 | override sleepAfter = "6h"; |
| 54 | |
| 55 | async startJob(input: { repoId: string; jobId: string; container?: string; token: string; apiUrl: string; egress?: Egress; vcpu?: number }): Promise<void> { |
| 56 | await this.ctx.storage.put("job", { repoId: input.repoId, jobId: input.jobId, container: input.container }); |
| 57 | const egressEnv = await this.useEgress(input.egress ?? "open"); |
| 58 | await this.start( |
| 59 | { |
| 60 | envVars: { |
| 61 | DIFFINERY_API_URL: input.apiUrl, |
| 62 | DIFFINERY_JOB_ID: input.jobId, |
| 63 | DIFFINERY_JOB_TOKEN: input.token, |
| 64 | // The runner's share of CPU, which the agent measures the job's use against (runner/sizes.ts). |
| 65 | ...(input.vcpu ? { DIFFINERY_RUNNER_VCPU: String(input.vcpu) } : {}), |
| 66 | ...egressEnv, |
| 67 | }, |
| 68 | }, |
| 69 | // The library gives up after 8 s and reports "no container instance", but a container on a host |
| 70 | // that hasn't pulled the runner image yet takes longer than that to come up: such jobs went back |
| 71 | // in line for 30 to 90 s more (measured on 10-leg matrices). Wait for it instead. |
| 72 | // The runner listens on no port: the library's own fallback port (33) answers "not listening" once it is up. |
| 73 | // Each check can itself take up to 5 s, so the signal bounds the whole wait. |
| 74 | { portToCheck: 33, waitInterval: START_POLL_MS, retries: Math.ceil(START_WAIT_MS / START_POLL_MS), signal: AbortSignal.timeout(START_WAIT_MS + 60_000) }, |
| 75 | ); |
| 76 | } |
| 77 | |
| 78 | async stopJob(): Promise<void> { |
| 79 | // Noted first, so onStop can tell Diffinery's own stop (a cancel, a timeout) from the platform's. |
| 80 | const job = await this.ctx.storage.get<RunnerJob>("job"); |
| 81 | if (job) await this.ctx.storage.put("job", { ...job, stopRequested: true }); |
| 82 | // SIGINT cancels the job at once. The agent reads SIGTERM as a platform rollout and |
| 83 | // keeps the job running for up to 14 minutes, which a cancelled job must not do. |
| 84 | if (this.ctx.container?.running) await this.stop("SIGINT"); |
| 85 | } |
| 86 | |
| 87 | override async onStop(params: StopParams): Promise<void> { |
| 88 | const job = await this.ctx.storage.get<RunnerJob>("job"); |
| 89 | if (!job) return; |
| 90 | const repo = this.env.REPO.get(this.env.REPO.idFromName(job.repoId)); |
| 91 | await repo.containerStopped(job.jobId, params.exitCode ?? null, job.container, !!job.stopRequested); |
| 92 | await this.ctx.storage.delete("job"); |
| 93 | } |
| 94 | |
| 95 | override onError(error: unknown): void { |
| 96 | console.error("runner container error", error); |
| 97 | } |
| 98 | } |
| 99 | |
| 100 | /** |
| 101 | * The smaller hosted runner sizes (runner/sizes.ts): the same runner on smaller Cloudflare instance |
| 102 | * types, each its own container class because a container class has one instance type. |
| 103 | */ |
| 104 | export class RunnerContainerMedium extends RunnerContainer {} |
| 105 | export class RunnerContainerSmall extends RunnerContainer {} |
| 106 | export class RunnerContainerTiny extends RunnerContainer {} |
| 107 | |
| 108 | /** |
| 109 | * Long-lived task sandbox for AI agents: the repo checked out with toolchains, |
| 110 | * and an HTTP API (exec, read/write files) on port 8080. |
| 111 | */ |
| 112 | export class SandboxContainer extends HostedContainer { |
| 113 | override defaultPort = 8080; |
| 114 | override sleepAfter = "30m"; |
| 115 | |
| 116 | async boot(envVars: Record<string, string>, secret: string, _name?: string, usage?: { id: string; accountId: string; repoId: string }, egress: Egress = "open"): Promise<void> { |
| 117 | await this.ctx.storage.put("secret", secret); |
| 118 | // Its running time counts against the repository owner's hosted minutes (lib/quotas.ts). |
| 119 | if (usage) await this.ctx.storage.put("usage", usage); |
| 120 | envVars = { ...envVars, ...(await this.useEgress(egress)) }; |
| 121 | // The first start on a host after a deploy pulls the runner image, which takes |
| 122 | // well over the library's 8 s / 20 s defaults. |
| 123 | await this.startAndWaitForPorts({ |
| 124 | ports: 8080, |
| 125 | startOptions: { envVars, entrypoint: ["node", "/opt/diffinery/agent.mjs", "--sandbox"] }, |
| 126 | cancellationOptions: { instanceGetTimeoutMS: 180_000, portReadyTimeoutMS: 120_000, waitInterval: 500 }, |
| 127 | }); |
| 128 | } |
| 129 | |
| 130 | /** Forward an API call (JSON text) to the sandbox's HTTP server with its secret attached. */ |
| 131 | async call(path: string, body: string): Promise<{ status: number; body: string; json: boolean }> { |
| 132 | const secret = await this.ctx.storage.get<string>("secret"); |
| 133 | // A sandbox that slept after 30 idle minutes starts again for this call: its time counts again. |
| 134 | const usage = await this.ctx.storage.get<{ id: string; accountId: string; repoId: string; part?: number }>("usage"); |
| 135 | if (usage && !this.ctx.container?.running) { |
| 136 | // It doesn't wake while a kill switch stops sandboxes (lib/killswitch.ts). |
| 137 | const block = await killSwitches(this.env, usage.repoId).then( |
| 138 | (k) => blockedBy(k, "sandbox"), |
| 139 | () => null, |
| 140 | ); |
| 141 | if (block) return { status: 403, body: JSON.stringify({ message: block }), json: true }; |
| 142 | const part = (usage.part ?? 0) + 1; |
| 143 | await startCompute(this.env, { id: `${usage.id}:${part}`, accountId: usage.accountId, repoId: usage.repoId, kind: "sandbox" }); |
| 144 | await this.ctx.storage.put("usage", { ...usage, part }); |
| 145 | } |
| 146 | const res = await this.containerFetch(`http://sandbox${path}`, { |
| 147 | method: "POST", |
| 148 | headers: { "content-type": "application/json", authorization: `Bearer ${secret}` }, |
| 149 | body: body || "{}", |
| 150 | }); |
| 151 | return { status: res.status, body: await res.text(), json: (res.headers.get("content-type") ?? "").includes("json") }; |
| 152 | } |
| 153 | |
| 154 | async shutdown(): Promise<void> { |
| 155 | if (this.ctx.container?.running) await this.stop(); |
| 156 | } |
| 157 | |
| 158 | /** Whether the container is up (calling into a stopped one would start it without its checkout). */ |
| 159 | async running(): Promise<boolean> { |
| 160 | return !!this.ctx.container?.running && !!(await this.ctx.storage.get<string>("secret")); |
| 161 | } |
| 162 | |
| 163 | /** |
| 164 | * A web terminal's socket (agentic/terminal-web.ts checked who is asking): only a WebSocket |
| 165 | * upgrade to /pty, with the sandbox's secret added. A sandbox that isn't running isn't started |
| 166 | * again (it would come up without its checkout), and nothing else of the sandbox is reachable here. |
| 167 | */ |
| 168 | override async fetch(request: Request): Promise<Response> { |
| 169 | const url = new URL(request.url); |
| 170 | if (url.pathname !== "/pty" || request.headers.get("upgrade")?.toLowerCase() !== "websocket") return new Response("not found", { status: 404 }); |
| 171 | const secret = await this.ctx.storage.get<string>("secret"); |
| 172 | if (!secret || !this.ctx.container?.running) return new Response("the terminal isn't running", { status: 410 }); |
| 173 | const headers = new Headers({ upgrade: "websocket", authorization: `Bearer ${secret}` }); |
| 174 | return this.containerFetch(new Request(`http://sandbox/pty${url.search}`, { headers }), 8080); |
| 175 | } |
| 176 | |
| 177 | override async onStop(): Promise<void> { |
| 178 | const usage = await this.ctx.storage.get<{ id: string; part?: number }>("usage"); |
| 179 | if (usage) await endCompute(this.env, usage.part ? `${usage.id}:${usage.part}` : usage.id); |
| 180 | } |
| 181 | } |
| 182 | |
| 183 | /** |
| 184 | * Local development and end-to-end tests only: `env.dev` binds SANDBOX to this |
| 185 | * instead of SandboxContainer, because containers don't run under `wrangler dev`. |
| 186 | * The same calls, served by one `agent.mjs --sandbox-pool` process at |
| 187 | * SANDBOX_POOL_URL that runs each sandbox's server. |
| 188 | */ |
| 189 | export class SandboxPoolDO extends DurableObject<Env> { |
| 190 | private poolAuth() { |
| 191 | return { authorization: `Bearer ${this.env.SANDBOX_POOL_SECRET ?? ""}` }; |
| 192 | } |
| 193 | |
| 194 | async boot(envVars: Record<string, string>, secret: string, name?: string): Promise<void> { |
| 195 | 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"); |
| 196 | await this.ctx.storage.put({ secret, name }); |
| 197 | const r = await fetch(`${this.env.SANDBOX_POOL_URL}/boot`, { method: "POST", headers: this.poolAuth(), body: JSON.stringify({ id: name, secret, env: envVars }) }); |
| 198 | if (!r.ok) throw new Error(`sandbox pool boot failed: ${r.status} ${await r.text()}`); |
| 199 | } |
| 200 | |
| 201 | async call(path: string, body: string): Promise<{ status: number; body: string; json: boolean }> { |
| 202 | const [secret, name] = await Promise.all([this.ctx.storage.get<string>("secret"), this.ctx.storage.get<string>("name")]); |
| 203 | if (!name) return { status: 404, body: JSON.stringify({ error: "sandbox not found" }), json: true }; |
| 204 | const res = await fetch(`${this.env.SANDBOX_POOL_URL}/s/${name}${path}`, { method: "POST", headers: { "content-type": "application/json", authorization: `Bearer ${secret}` }, body: body || "{}" }); |
| 205 | return { status: res.status, body: await res.text(), json: (res.headers.get("content-type") ?? "").includes("json") }; |
| 206 | } |
| 207 | |
| 208 | async shutdown(): Promise<void> { |
| 209 | const name = await this.ctx.storage.get<string>("name"); |
| 210 | if (name) await fetch(`${this.env.SANDBOX_POOL_URL}/shutdown/${name}`, { method: "POST", headers: this.poolAuth() }).catch(() => {}); |
| 211 | } |
| 212 | |
| 213 | async running(): Promise<boolean> { |
| 214 | if (!(await this.ctx.storage.get<string>("name"))) return false; |
| 215 | return (await this.call("/pty/status", "{}").catch(() => ({ status: 0 }))).status === 200; |
| 216 | } |
| 217 | |
| 218 | /** As SandboxContainer.fetch: a terminal's WebSocket, through the pool. */ |
| 219 | override async fetch(request: Request): Promise<Response> { |
| 220 | const url = new URL(request.url); |
| 221 | if (url.pathname !== "/pty" || request.headers.get("upgrade")?.toLowerCase() !== "websocket") return new Response("not found", { status: 404 }); |
| 222 | const [secret, name] = await Promise.all([this.ctx.storage.get<string>("secret"), this.ctx.storage.get<string>("name")]); |
| 223 | if (!secret || !name) return new Response("the terminal isn't running", { status: 410 }); |
| 224 | const res = await fetch(`${this.env.SANDBOX_POOL_URL}/s/${name}/pty${url.search}`, { headers: { upgrade: "websocket", authorization: `Bearer ${secret}` } }); |
| 225 | if (!res.webSocket) return new Response(await res.text(), { status: res.status === 101 ? 502 : res.status }); |
| 226 | return pipeSocket(res.webSocket); |
| 227 | } |
| 228 | } |
| 229 | |
| 230 | /** Hands a WebSocket this Durable Object holds to the caller, relaying frames and closes both ways. */ |
| 231 | function pipeSocket(upstream: WebSocket): Response { |
| 232 | const [client, server] = Object.values(new WebSocketPair()) as [WebSocket, WebSocket]; |
| 233 | upstream.accept(); |
| 234 | server.accept(); |
| 235 | const code = (c: number) => (c === 1005 || c === 1006 ? 1000 : c); |
| 236 | // Binary frames can arrive as Blobs; reading one is async, so each direction keeps its order with a chain. |
| 237 | const relay = (from: WebSocket, to: WebSocket) => { |
| 238 | let chain = Promise.resolve(); |
| 239 | from.addEventListener("message", (e) => { |
| 240 | const data = e.data as string | ArrayBuffer | Blob; |
| 241 | chain = chain.then(async () => to.send(data instanceof Blob ? await data.arrayBuffer() : data)).catch(() => to.close(1011, "relay failed")); |
| 242 | }); |
| 243 | }; |
| 244 | relay(server, upstream); |
| 245 | relay(upstream, server); |
| 246 | server.addEventListener("close", (e) => upstream.close(code(e.code), e.reason)); |
| 247 | upstream.addEventListener("close", (e) => server.close(code(e.code), e.reason)); |
| 248 | server.addEventListener("error", () => upstream.close(1011, "client error")); |
| 249 | upstream.addEventListener("error", () => server.close(1011, "sandbox error")); |
| 250 | return new Response(null, { status: 101, webSocket: client }); |
| 251 | } |