computer: Give host module calls an execution id, log and defer

.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.

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<string, Socket>();
+
+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.
 

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: () => {},
   };
 }
 

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", [
         {

packages/computer/src/modules/artifacts.test.ts

@@ -39,10 +39,13 @@ function build(artifacts: ReturnType<typeof fakeArtifacts>["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: () => {},
   };
 }
 

packages/computer/src/modules/assets.test.ts

@@ -37,10 +37,13 @@ function callContext(
   overrides: Partial<WorkspaceModuleCallContext> = {},
 ): 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,
   };
 }

packages/computer/src/modules/container.test.ts

@@ -86,10 +86,13 @@ function callContext(
   overrides: Partial<WorkspaceModuleCallContext> = {},
 ): WorkspaceModuleCallContext {
   return {
+    executionId: "test-execution",
     signal: new AbortController().signal,
     deadline: Date.now() + 60_000,
     access: "read-write",
     resolvePath: async (path) => path,
+    log: () => {},
+    defer: () => {},
     ...overrides,
   };
 }

packages/computer/src/modules/tools.test.ts

@@ -14,10 +14,13 @@ function callContext(
   overrides: Partial<WorkspaceModuleCallContext> = {},
 ): WorkspaceModuleCallContext {
   return {
+    executionId: "test-execution",
     signal: new AbortController().signal,
     deadline: Date.now() + 60_000,
     access: "read-write",
     resolvePath: async (path) => path,
+    log: () => {},
+    defer: () => {},
     ...overrides,
   };
 }

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");
+  });
+});

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<Uint8Array>) => Promise<void>;
+  readonly #executionId: string;
+  readonly #onLog?: (text: string) => void;
+  readonly #maxLogBytes: number;
+  readonly #deferred: Array<() => unknown> = [];
+  #logBytes = 0;
+  #ended = false;
   readonly #inFlight = new Set<Promise<MeasuredResponse>>();
   readonly #abortControllers = new Set<AbortController>();
   #cancelled = false;
@@ -39,6 +45,9 @@ export class WorkspaceRuntimeBridge extends RpcTarget {
       maxTotalResponseBytes?: number;
       maxResultBytes?: number;
       onAttachOutput?: (readable: ReadableStream<Uint8Array>) => Promise<void>;
+      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<void> {
+    if (this.#ended) return;
+    while (this.#deferred.length > 0) {
+      const fn = this.#deferred.pop() as () => unknown;
+      let timer: ReturnType<typeof setTimeout> | undefined;
+      const timeout = new Promise<never>((_, 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<void> {
     this.#cancelled = true;
     for (const controller of this.#abortControllers) {

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<string>;
+  /**
+   * 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;
 }
 
 /**

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: () => {},
   };
 }
 

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: () => {},
   };
 }