diff --git a/.changeset/host-module-lifecycle.md b/.changeset/host-module-lifecycle.md new file mode 100644 index 000000000..f49875830 --- /dev/null +++ b/.changeset/host-module-lifecycle.md @@ -0,0 +1,5 @@ +--- +"@cloudflare/computer": minor +--- + +Host module calls get `executionId`, `log(text)` for writing to the run's standard error, and `defer(fn)` for cleanup when the run ends. diff --git a/docs/17_isolate_javascript.md b/docs/17_isolate_javascript.md index 0bf86598b..83243021c 100644 --- a/docs/17_isolate_javascript.md +++ b/docs/17_isolate_javascript.md @@ -291,10 +291,39 @@ Each function receives the arguments the isolate passed, as an array of JSON-com | Field | Meaning | | --- | --- | +| `executionId` | The run making the call. Every call from one run has the same id. | | `signal` | Aborts when the call passes its deadline or the execution is cancelled. | | `deadline` | Epoch milliseconds after which the isolate stops waiting. | | `access` | The backend's `"read"` or `"read-write"` access. Check it before any write. | | `resolvePath(path, { allowMissing })` | Confines a caller path to the backend root and rejects symlinks. | +| `log(text)` | Writes a line to the run's standard error. | +| `defer(fn)` | Runs `fn` once when the run ends. | + +A module that holds something for the length of a run, such as a connection, keys it by `executionId` and cleans up in `defer`: + +```ts +const sockets = new Map(); + +modules: { + "ws:remote": { + async send([message], context) { + let socket = sockets.get(context.executionId); + if (!socket) { + socket = await openSocket(); + sockets.set(context.executionId, socket); + context.defer(() => { + sockets.delete(context.executionId); + socket.close(); + }); + } + if (socket.reconnected) context.log("The remote end restarted; earlier state is gone."); + return socket.send(message); + }, + }, +} +``` + +Deferred functions run however the run ends, whether it completes, fails, times out or is cancelled. They run after the code has finished and its host calls have settled, and before the result is published, so the next run sees whatever they saved. They run one at a time, last registered first, and each gets the host call deadline. One that throws or misses its deadline is reported on standard error, and the rest still run. A deferred function can call `log`, which is how a module reports something it only knows at the end of the run. Log lines count toward their own `maxStdioBytes` limit, separate from the code's output. Past it they are dropped, as is any line written after the run settled. The arguments come from caller code, so parse them before use. A function may return a value or a promise. Arguments and results cross the isolate boundary as real values through Workers RPC, not as encoded text. The result must be JSON-compatible plain data, and the bridge checks it at runtime. As in JSON, an `undefined` result or array item becomes `null` and an `undefined` object field is left out, in arguments and results alike. Byte arrays work for `node:fs` calls but not for host modules. It fits within the same capability byte limits as every other host call. A function that ignores `signal` and never settles keeps the execution in its finalizing state. diff --git a/examples/rlm/worker/capability.test.ts b/examples/rlm/worker/capability.test.ts index d08381006..acc9ffcc1 100644 --- a/examples/rlm/worker/capability.test.ts +++ b/examples/rlm/worker/capability.test.ts @@ -36,10 +36,13 @@ function successfulResult(text = "ok") { function context(signal = new AbortController().signal): WorkspaceModuleCallContext { return { + executionId: "test-execution", signal, deadline: Date.now() + 1_000, access: "read", resolvePath: async (path) => path, + log: () => {}, + defer: () => {}, }; } diff --git a/packages/computer/src/backends/worker-javascript/worker-javascript.ts b/packages/computer/src/backends/worker-javascript/worker-javascript.ts index 3a8ee6e30..8edd2aa95 100644 --- a/packages/computer/src/backends/worker-javascript/worker-javascript.ts +++ b/packages/computer/src/backends/worker-javascript/worker-javascript.ts @@ -534,6 +534,17 @@ class JavaScriptBackendHandle implements WorkspaceModuleBackendHandle { maxTotalResponseBytes: this.#options.maxCapabilityResponseBytes, maxResultBytes: this.#options.maxResultBytes, onAttachOutput: (readable) => this.#pumpFrames(record, readable), + executionId: id, + maxLogBytes: this.#options.maxStdioBytes, + onLog: (text) => { + if (record.status !== "running") return; + this.#append(record, { + id: record.id, + seq: record.events.length + 1, + name: "stderr", + value: new TextEncoder().encode(text), + }); + }, }); record.bridge = bridge; record.control = startJavaScriptExecution({ @@ -681,6 +692,7 @@ class JavaScriptBackendHandle implements WorkspaceModuleBackendHandle { try { await record.bridge?.cancelAndDrain(); } catch (error) { + await record.bridge?.runDeferred(); this.#finish(record, "failed", [ { id: record.id, @@ -694,6 +706,7 @@ class JavaScriptBackendHandle implements WorkspaceModuleBackendHandle { ]); return; } + await record.bridge?.runDeferred(); this.#finish(record, "cancelled", [ { id: record.id, @@ -816,6 +829,7 @@ class JavaScriptBackendHandle implements WorkspaceModuleBackendHandle { try { await record.bridge?.cancelAndDrain(); } catch (error) { + await record.bridge?.runDeferred(); this.#finish(record, "failed", [ { id: record.id, @@ -829,6 +843,7 @@ class JavaScriptBackendHandle implements WorkspaceModuleBackendHandle { ]); return; } + await record.bridge?.runDeferred(); if (errorMessage !== undefined) { this.#finish(record, "failed", [ { diff --git a/packages/computer/src/modules/artifacts.test.ts b/packages/computer/src/modules/artifacts.test.ts index 90fd8ab30..5e0504b05 100644 --- a/packages/computer/src/modules/artifacts.test.ts +++ b/packages/computer/src/modules/artifacts.test.ts @@ -39,10 +39,13 @@ function build(artifacts: ReturnType["client"]) { function callContext(access: "read" | "read-write" = "read-write"): WorkspaceModuleCallContext { return { + executionId: "test-execution", signal: new AbortController().signal, deadline: Date.now() + 60_000, access, resolvePath: async (path) => path, + log: () => {}, + defer: () => {}, }; } diff --git a/packages/computer/src/modules/assets.test.ts b/packages/computer/src/modules/assets.test.ts index 170f33b60..6c050ef4e 100644 --- a/packages/computer/src/modules/assets.test.ts +++ b/packages/computer/src/modules/assets.test.ts @@ -37,10 +37,13 @@ function callContext( overrides: Partial = {}, ): WorkspaceModuleCallContext { return { + executionId: "test-execution", signal: new AbortController().signal, deadline: Date.now() + 60_000, access: "read-write", resolvePath: async (path) => (path.startsWith("/") ? path : `/workspace/${path}`), + log: () => {}, + defer: () => {}, ...overrides, }; } diff --git a/packages/computer/src/modules/container.test.ts b/packages/computer/src/modules/container.test.ts index 592418a8a..8962ef947 100644 --- a/packages/computer/src/modules/container.test.ts +++ b/packages/computer/src/modules/container.test.ts @@ -86,10 +86,13 @@ function callContext( overrides: Partial = {}, ): WorkspaceModuleCallContext { return { + executionId: "test-execution", signal: new AbortController().signal, deadline: Date.now() + 60_000, access: "read-write", resolvePath: async (path) => path, + log: () => {}, + defer: () => {}, ...overrides, }; } diff --git a/packages/computer/src/modules/tools.test.ts b/packages/computer/src/modules/tools.test.ts index f8b9709a9..4eb3cd53c 100644 --- a/packages/computer/src/modules/tools.test.ts +++ b/packages/computer/src/modules/tools.test.ts @@ -14,10 +14,13 @@ function callContext( overrides: Partial = {}, ): WorkspaceModuleCallContext { return { + executionId: "test-execution", signal: new AbortController().signal, deadline: Date.now() + 60_000, access: "read-write", resolvePath: async (path) => path, + log: () => {}, + defer: () => {}, ...overrides, }; } diff --git a/packages/computer/src/runtime/bridge.test.ts b/packages/computer/src/runtime/bridge.test.ts index 075e2d267..00349dd45 100644 --- a/packages/computer/src/runtime/bridge.test.ts +++ b/packages/computer/src/runtime/bridge.test.ts @@ -245,3 +245,49 @@ describe("WorkspaceRuntimeBridge assertResult", () => { ); }); }); + +describe("WorkspaceRuntimeBridge host module lifecycle", () => { + it("drops log lines past the byte limit", async () => { + const lines: string[] = []; + const target = new WorkspaceRuntimeBridge({} as WorkspaceRuntimeCapability, { + maxLogBytes: 8, + onLog: (text) => lines.push(text), + hostModules: new Map([ + [ + "ws:test", + { + run: async (_values, context) => { + context.log("abc"); + context.log("defghij"); + context.log("k"); + return null; + }, + }, + ], + ]), + }); + await expect(message(target.call("host/ws:test.run", []))).resolves.toBeUndefined(); + expect(lines).toEqual(["abc\n", "k\n"]); + }); + + it("refuses defer once the execution has ended", async () => { + let defer: ((fn: () => unknown) => void) | undefined; + const target = new WorkspaceRuntimeBridge({} as WorkspaceRuntimeCapability, { + executionId: "run-1", + hostModules: new Map([ + [ + "ws:test", + { + run: async (_values, context) => { + defer = context.defer; + return context.executionId; + }, + }, + ], + ]), + }); + await expect(target.call("host/ws:test.run", [])).resolves.toEqual({ result: "run-1" }); + await target.runDeferred(); + expect(() => defer?.(() => undefined)).toThrow("already ended"); + }); +}); diff --git a/packages/computer/src/runtime/bridge.ts b/packages/computer/src/runtime/bridge.ts index b1f18395b..de3261843 100644 --- a/packages/computer/src/runtime/bridge.ts +++ b/packages/computer/src/runtime/bridge.ts @@ -19,6 +19,12 @@ export class WorkspaceRuntimeBridge extends RpcTarget { readonly #maxTotalResponseBytes: number; readonly #maxResultBytes: number; readonly #onAttachOutput?: (readable: ReadableStream) => Promise; + readonly #executionId: string; + readonly #onLog?: (text: string) => void; + readonly #maxLogBytes: number; + readonly #deferred: Array<() => unknown> = []; + #logBytes = 0; + #ended = false; readonly #inFlight = new Set>(); readonly #abortControllers = new Set(); #cancelled = false; @@ -39,6 +45,9 @@ export class WorkspaceRuntimeBridge extends RpcTarget { maxTotalResponseBytes?: number; maxResultBytes?: number; onAttachOutput?: (readable: ReadableStream) => Promise; + executionId?: string; + onLog?: (text: string) => void; + maxLogBytes?: number; } = {}, ) { super(); @@ -52,6 +61,9 @@ export class WorkspaceRuntimeBridge extends RpcTarget { this.#maxTotalResponseBytes = integrations.maxTotalResponseBytes ?? 8 * 1024 * 1024; this.#maxResultBytes = integrations.maxResultBytes ?? 1024 * 1024; this.#onAttachOutput = integrations.onAttachOutput; + this.#executionId = integrations.executionId ?? ""; + this.#onLog = integrations.onLog; + this.#maxLogBytes = integrations.maxLogBytes ?? 1024 * 1024; } // Drain the runner's framed output stream. The isolate passes the @@ -114,11 +126,14 @@ export class WorkspaceRuntimeBridge extends RpcTarget { const operation = respond(async () => { if (name.startsWith("host/")) { return this.#callHostModule(name, args, { + executionId: this.#executionId, signal: abort.signal, deadline, access: this.#capability.access, resolvePath: (path, options) => this.#capability.resolveConfined(path, options?.allowMissing ?? false), + log: (text) => this.#log(text), + defer: (fn) => this.#defer(fn), }); } const operation = name.startsWith("fs.") ? name.slice(3) : name; @@ -215,6 +230,51 @@ export class WorkspaceRuntimeBridge extends RpcTarget { }); } + // Run what host modules deferred, last registered first, once the + // execution's calls have drained. A deferred function that throws or + // misses the host call deadline is reported on standard error and + // doesn't stop the rest. Runs once; later `defer` calls are refused. + async runDeferred(): Promise { + if (this.#ended) return; + while (this.#deferred.length > 0) { + const fn = this.#deferred.pop() as () => unknown; + let timer: ReturnType | undefined; + const timeout = new Promise((_, reject) => { + timer = setTimeout( + () => reject(new Error("did not settle before its deadline")), + this.#maxCallDurationMs, + ); + }); + try { + await Promise.race([Promise.resolve().then(fn), timeout]); + } catch (error) { + this.#log( + `A host module's deferred function failed: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + } finally { + clearTimeout(timer); + } + } + this.#ended = true; + } + + #defer(fn: () => unknown): void { + if (typeof fn !== "function") throw new TypeError("defer takes a function."); + if (this.#ended) throw new Error("The execution has already ended."); + this.#deferred.push(fn); + } + + #log(text: string): void { + if (this.#ended || this.#onLog === undefined) return; + const line = String(text).endsWith("\n") ? String(text) : `${String(text)}\n`; + const bytes = utf8ByteLength(line); + if (this.#logBytes + bytes > this.#maxLogBytes) return; + this.#logBytes += bytes; + this.#onLog(line); + } + async cancelAndDrain(): Promise { this.#cancelled = true; for (const controller of this.#abortControllers) { diff --git a/packages/computer/src/runtime/types.ts b/packages/computer/src/runtime/types.ts index c1e7d76ec..472a937e1 100644 --- a/packages/computer/src/runtime/types.ts +++ b/packages/computer/src/runtime/types.ts @@ -7,6 +7,11 @@ export type WorkspaceRuntimeAccess = "read" | "read-write"; /** Per-call context the backend passes to every host module function. */ export interface WorkspaceModuleCallContext { + /** + * The execution making the call. Every call from one run has the same + * id, so a module can keep state for the run, such as a connection. + */ + readonly executionId: string; /** Aborts when the call passes its deadline or the execution is cancelled. */ readonly signal: AbortSignal; /** Epoch milliseconds after which the caller stops waiting for this call. */ @@ -22,6 +27,23 @@ export interface WorkspaceModuleCallContext { * @returns The confined absolute path. */ resolvePath(path: string, options?: { readonly allowMissing?: boolean }): Promise; + /** + * Write a line to the execution's standard error, to tell the caller + * something the call's result can't carry. A newline is added when + * `text` doesn't end with one. Lines past `maxStdioBytes` in total are + * dropped, and a line written after the execution settled is ignored. + */ + log(text: string): void; + /** + * Run `fn` once when the execution ends, however it ends: after its + * code has finished and its host calls have settled, and before its + * result is published. Use it to close a connection or save state the + * run built up. Deferred functions run one at a time, last registered + * first, and can still call `log`. Each gets the host call deadline. + * One that throws or misses its deadline is reported on standard + * error, and the rest still run. + */ + defer(fn: () => unknown): void; } /** diff --git a/packages/computer/src/tools/ai-sdk/bindings.test.ts b/packages/computer/src/tools/ai-sdk/bindings.test.ts index c75a2ba4c..fd4936fbd 100644 --- a/packages/computer/src/tools/ai-sdk/bindings.test.ts +++ b/packages/computer/src/tools/ai-sdk/bindings.test.ts @@ -11,10 +11,13 @@ import { createAITools } from "./index.js"; function callContext(): WorkspaceModuleCallContext { return { + executionId: "test-execution", signal: new AbortController().signal, deadline: Date.now() + 60_000, access: "read-write", resolvePath: async (path) => path, + log: () => {}, + defer: () => {}, }; } diff --git a/packages/computer/src/tools/pi-ai/bindings.test.ts b/packages/computer/src/tools/pi-ai/bindings.test.ts index 6b09ec6ee..b015cd0fd 100644 --- a/packages/computer/src/tools/pi-ai/bindings.test.ts +++ b/packages/computer/src/tools/pi-ai/bindings.test.ts @@ -14,10 +14,13 @@ function callContext( access: WorkspaceModuleCallContext["access"] = "read-write", ): WorkspaceModuleCallContext { return { + executionId: "test-execution", signal: new AbortController().signal, deadline: Date.now() + 60_000, access, resolvePath: async (path) => path, + log: () => {}, + defer: () => {}, }; } diff --git a/packages/computer/tests/script-runner-worker.ts b/packages/computer/tests/script-runner-worker.ts index b2c0a1486..e475485e7 100644 --- a/packages/computer/tests/script-runner-worker.ts +++ b/packages/computer/tests/script-runner-worker.ts @@ -57,6 +57,11 @@ function fakeContainerBackend(): WorkspaceBackend { }; } +// Executions whose deferred functions have finished, so a later run can +// check that they ran before the earlier run's result was published. +const endedRuns = new Set(); +const trackedRuns = new Set(); + // The Worker Loader resolves imports differently under the legacy and // the new module registry, so the same backend runs on both. const NEW_REGISTRY_BACKEND = "worker-javascript-new-registry"; @@ -80,6 +85,37 @@ function javascriptBackend( "ws:git": createGitModule(), "ws:artifacts": createArtifactsModule(), "ws:container": createContainerModule(), + "ws:lifecycle": { + id(_args, context) { + return context.executionId; + }, + log([text], context) { + context.log(String(text)); + return null; + }, + track([label], context) { + trackedRuns.add(context.executionId); + context.defer(() => context.log(`${String(label)} 1`)); + context.defer(async () => { + await new Promise((resolve) => setTimeout(resolve, 20)); + endedRuns.add(context.executionId); + context.log(`${String(label)} 2`); + }); + return null; + }, + throwing(_args, context) { + context.defer(() => { + throw new Error("boom"); + }); + return null; + }, + ended([id]) { + return endedRuns.has(String(id)); + }, + tracked([id]) { + return trackedRuns.has(String(id)); + }, + }, "ws:test-host": { async echo(args) { return { args: [...args] }; diff --git a/packages/computer/tests/script-runner.test.ts b/packages/computer/tests/script-runner.test.ts index 7796e133b..4f718b0b8 100644 --- a/packages/computer/tests/script-runner.test.ts +++ b/packages/computer/tests/script-runner.test.ts @@ -340,6 +340,133 @@ describe("WorkspaceRuntime", () => { ]); }); + it("gives every host call in a run the run's execution id", async () => { + const response = await runtime({ + source: ` + import { id } from "ws:lifecycle"; + export default async () => [await id(), await id()]; + `, + cwd: "/workspace", + id: "lifecycle-id", + }); + const text = await response.text(); + expect(response.status, text).toBe(200); + expect(JSON.parse(text).result.value).toEqual(["lifecycle-id", "lifecycle-id"]); + }); + + it("writes a host module's log lines to standard error", async () => { + const response = await runtime({ + source: ` + import { log } from "ws:lifecycle"; + export default async () => { await log("from host"); return "done"; }; + `, + cwd: "/workspace", + }); + const text = await response.text(); + expect(response.status, text).toBe(200); + expect(JSON.parse(text).result).toMatchObject({ + status: "completed", + stderr: "from host\n", + value: "done", + }); + }); + + it("runs deferred functions last first, before the result is published", async () => { + const response = await runtime({ + source: ` + import { track } from "ws:lifecycle"; + export default async () => { await track("t"); return "done"; }; + `, + cwd: "/workspace", + id: "lifecycle-defer", + }); + const text = await response.text(); + expect(response.status, text).toBe(200); + expect(JSON.parse(text).result).toMatchObject({ + status: "completed", + stderr: "t 2\nt 1\n", + value: "done", + }); + + const check = await runtime({ + source: ` + import { ended } from "ws:lifecycle"; + export default () => ended("lifecycle-defer"); + `, + cwd: "/workspace", + }); + const checked = await check.text(); + expect(check.status, checked).toBe(200); + expect(JSON.parse(checked).result.value).toBe(true); + }); + + it("runs deferred functions when the run fails", async () => { + const response = await runtime({ + source: ` + import { track } from "ws:lifecycle"; + export default async () => { await track("f"); throw new Error("x"); }; + `, + cwd: "/workspace", + }); + const text = await response.text(); + expect(response.status, text).toBe(200); + const { result } = JSON.parse(text); + expect(result.status).toBe("failed"); + expect(result.stderr).toContain("f 2\nf 1\n"); + }); + + it("runs deferred functions when the run is cancelled", async () => { + const id = `lifecycle-cancel-${crypto.randomUUID()}`; + const start = await SELF.fetch("https://example.test/runtime-start", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + id, + source: ` + import { track } from "ws:lifecycle"; + export default async function () { + await track("c"); + await new Promise((resolve) => setTimeout(resolve, 10000)); + } + `, + }), + }); + expect(start.status).toBe(200); + const probe = (name: string) => + runtime({ + source: `import { ${name} } from "ws:lifecycle"; export default () => ${name}(${JSON.stringify(id)});`, + cwd: "/workspace", + }).then(async (response) => JSON.parse(await response.text()).result.value); + for (let attempt = 0; attempt < 50 && !(await probe("tracked")); attempt++) { + await new Promise((resolve) => setTimeout(resolve, 10)); + } + const killed = await SELF.fetch( + `https://example.test/runtime-kill?id=${encodeURIComponent(id)}`, + { method: "POST" }, + ); + expect(killed.status).toBe(204); + const result = await SELF.fetch( + `https://example.test/runtime-get?id=${encodeURIComponent(id)}`, + ); + expect(await result.json()).toMatchObject({ status: "cancelled", exitCode: 130 }); + expect(await probe("ended")).toBe(true); + }); + + it("reports a deferred function that throws and still settles the run", async () => { + const response = await runtime({ + source: ` + import { throwing } from "ws:lifecycle"; + export default async () => { await throwing(); return "done"; }; + `, + cwd: "/workspace", + }); + const text = await response.text(); + expect(response.status, text).toBe(200); + const { result } = JSON.parse(text); + expect(result).toMatchObject({ status: "completed", value: "done" }); + expect(result.stderr).toContain("deferred function failed: boom"); + }); + it("fails to link an import the host module does not export", async () => { const response = await runtime({ source: `