diff --git a/.changeset/scoped-host-modules.md b/.changeset/scoped-host-modules.md new file mode 100644 index 00000000..7b723cd2 --- /dev/null +++ b/.changeset/scoped-host-modules.md @@ -0,0 +1,5 @@ +--- +"@cloudflare/computer": minor +--- + +Add `createScopedModule()` for host modules that hold state for a run, such as a browser session. A scope lasts from when a run's code starts in an isolate until it waits in `ctx.interrupt()` or ends. The backend starts one instance for each scope, on the scope's first call to the module, and closes it when the scope ends, with a `log()` on a run's end that appends to its stderr. Scoped modules are experimental. diff --git a/docs/17_isolate_javascript.md b/docs/17_isolate_javascript.md index 6c9cc94b..444aa6e3 100644 --- a/docs/17_isolate_javascript.md +++ b/docs/17_isolate_javascript.md @@ -276,7 +276,7 @@ Caller source can import four kinds of module, and all of them are fixed when th | Built in | Always installed | The isolate, backed by the Workspace | `node:fs`, `node:fs/promises` | | Node.js | `nodejs_compat` in `compatibilityFlags`, the default | The isolate, provided by the runtime | `node:path`, `node:crypto` | | Source | `modules: { name: "source" }` | The isolate | a bundled library | -| Host | `modules: { "codemode:name": { fn } }`, or a factory | The Durable Object | `codemode:git`, `codemode:container`, `codemode:tools`, your own | +| Host | `modules: { "codemode:name": { fn } }`, a factory, or `createScopedModule()` | The Durable Object | `codemode:git`, `codemode:container`, `codemode:tools`, your own | ```ts import { createArtifactsModule } from "@cloudflare/computer/modules/artifacts"; @@ -324,7 +324,7 @@ Modules code can import: - Node.js built-ins: `node:path`, `node:url`, ... Bare names such as `path` work too. ``` -A factory adds its own text through a `description` property, as the prebuilt modules do. The backend reads it each time it describes itself, not once when it is constructed, so a factory can define `description` as a getter and describe something that exists only later, as `codemode:tools` does with the agent's tools. An object of functions is listed by its export names; say more about it in the `exec` tool's backend description if the model needs it. +A factory adds its own text through a `description` property, as the prebuilt modules do. The backend reads it each time it describes itself, not once when it is constructed, so a factory can define `description` as a getter and describe something that exists only later, as `codemode:tools` does with the agent's tools. A scoped module takes `description` as an option. An object of functions, or a scoped module without a description, is listed by its export names; say more about it in the `exec` tool's backend description if the model needs it. ### Built-in filesystem @@ -429,6 +429,43 @@ The run's code never sees metadata: it isn't in `process.env` or the input, so c Specifiers and the export names of an object are checked at construction. A factory's export names are checked when the backend connects and the factory runs. An object must export at least one function. A factory may build none, as `codemode:tools` does for an agent with no tools: the backend then leaves that module out of the connection, so only code that imports it fails, with a message saying the module exports nothing, and every other run still starts. Every export name must be a JavaScript identifier name other than `default` or `then`. A reserved word such as `delete` is allowed, and caller code renames it on import: `import { delete as remove } from "codemode:files"`. Importing a name the module does not export fails when the module graph links, before any code runs. +#### Scoped modules + +A module that holds state for a run, such as a browser session it opens on the first call, can be made with `createScopedModule()`. The state lives for one scope of the run: from when the run's code starts in an isolate until it waits in `ctx.interrupt()` or ends. A run that resumes starts again from the top in a new scope. The backend makes one instance of the module for each scope, on the scope's first call to the module, and closes it when the scope ends: + +```ts +import { createScopedModule } from "@cloudflare/computer"; + +modules: { + "codemode:browser": createScopedModule({ + description: "A browser: `open(url)` and `text()`.", + exports: ["open", "text"], + async start(scope, host) { + const session = await browsers.connect({ conversation: scope.metadata.conversation }); + return { + functions: { + open: ([url]) => session.open(String(url)), + text: () => session.text(), + }, + async close(end) { + await session.close(); + if (end.reason !== "paused") end.log("Closed the browser.\n"); + }, + }; + }, + }), +} +``` + +`start` gets the scope's `{ id, metadata, signal }`, where `id` and `metadata` are the run's and stay the same in every scope, and the same services a factory gets. It returns the instance's `functions` with an optional `close`. The backend needs the export names before any run starts, so the module lists them in `exports`, which is checked at construction like an object's export names. If the functions `start` returns don't match `exports` exactly, the call fails and says so. Calls a scope makes while its instance starts wait for it, and if `start` throws, they fail. `scope.signal` aborts when the scope is cancelled or releases its isolate, and a call that was cancelled or timed out by the time `start` returns fails rather than running. A call answered from a replayable run's journal doesn't start an instance, so a scope whose calls to the module are all replayed never starts one. + +An instance closes only if it started, and then once, after the scope's host calls settle. `close` gets one of these reasons: + +- `"paused"`: the run is waiting in `ctx.interrupt()` and has released its isolate. Its code starts again from the top when it resumes, in a new scope with a new instance, so let go of everything. An answer that arrives during `close` resumes the run once `close` is done. +- `"completed"`, `"failed"`, or `"cancelled"`: the run ended. `close` runs before its exit is published, and only these reasons have `log(text)`, which appends text to the run's stderr ahead of its exit. Logged text is saved with the run's terminal state, so a run interrupted before it ends never resumes against output it didn't write. Text past what's left of the run's `maxStdioBytes` is cut, and ends with a line saying `Host module output truncated.` Text logged after `close` returns is dropped. + +A run killed while it waits has no instance open, so nothing closes. If the Durable Object restarts during a scope, the scope's instances go with it, unclosed. Every instance of a scope closes at once, with a `signal` and a `deadline`, and the backend waits for them up to `maxHostCallMs`, when `signal` aborts. A `close` that throws or runs late is reported with `console.error`, and the run goes on as it would have. Scoped modules are experimental and may change in any release. + ### `codemode:git` ```js diff --git a/packages/computer/README.md b/packages/computer/README.md index 85fc9f7e..0727b531 100644 --- a/packages/computer/README.md +++ b/packages/computer/README.md @@ -34,6 +34,7 @@ removed in a future release; each one names its replacement. | Resuming interrupted isolate JavaScript: `onInterrupt: "restart"` and `"replay"`, `workspace.runtime.recover()`, `unsafeToRepeat()`, and waiting for answers with `ctx.interrupt()`, `pending()` and `respond()` | `@cloudflare/computer`, `@cloudflare/computer/backends/worker-javascript` | Experimental | | `WorkerShellBackend` and the shell command groups | `@cloudflare/computer/backends/worker-shell`, `@cloudflare/computer/shell/*` | Experimental | | The `codemode:container` host module | `@cloudflare/computer/modules/container` | Experimental | +| Scoped host modules: `createScopedModule()` | `@cloudflare/computer`, `@cloudflare/computer/backends/worker-javascript` | Experimental | | `workspace.push()`, `workspace.pull()`, and the sync protocol between the durable object and `computerd` | `@cloudflare/computer` | Experimental | | The `computerd` daemon and its container image | `ghcr.io/cloudflare/computer-computerd-linux-x64` | Experimental | | The RPC wire protocol between the durable object and `computerd` (`@cloudflare/computer-rpc`) | Internal | Experimental | diff --git a/packages/computer/src/backends/worker-javascript/module-graph.ts b/packages/computer/src/backends/worker-javascript/module-graph.ts index 3a329c17..f05d8c48 100644 --- a/packages/computer/src/backends/worker-javascript/module-graph.ts +++ b/packages/computer/src/backends/worker-javascript/module-graph.ts @@ -1,6 +1,7 @@ import { parse } from "acorn"; import type { WorkspaceRuntimeCapability } from "../../runtime/capability"; +import { isScopedModule, type WorkspaceScopedModule } from "../../runtime/scoped-module"; import type { WorkspaceModule, WorkspaceModuleFactory, @@ -85,6 +86,7 @@ const RESERVED_EXPORT_NAMES = new Set(["default", "then"]); export interface ParsedModules { readonly source: Readonly>; readonly host: ReadonlyMap; + readonly scoped: ReadonlyMap; /** One markdown bullet per importable module, for a model. */ readonly description: string; } @@ -93,19 +95,22 @@ const FILESYSTEM_DESCRIPTION = "- `node:fs/promises` (also `node:fs`): the workspace's files. `readFile`, `writeFile`, `mkdir`, `rm`, `readdir`, `stat`, `lstat`, `readlink`, `symlink`, `chmod`, and `access`. Async only."; /** - * Parse the backend's `modules` option: split it into bundled source - * and host module factories, check every specifier and every object's - * export names, and describe each module for a model. A factory's - * export names are checked when the backend connects and it runs. + * Parse the backend's `modules` option: split it into bundled source, + * host module factories, and scoped host modules, check every + * specifier and the export names of every object and scoped module, + * and describe each module for a model. A factory's export names are + * checked when the backend connects and it runs. * * @param modules - The modules passed to the backend. - * @returns Source modules, host module factories, and their description. + * @returns Source modules, host module factories, scoped modules, and + * their description. * @throws When a specifier or an object's export names are not allowed. * The host configured the backend wrongly and no execution can use it. */ export function parseModules(modules: Readonly>): ParsedModules { const source: Record = Object.create(null); const host = new Map(); + const scoped = new Map(); const lines: Array string)> = [FILESYSTEM_DESCRIPTION]; for (const [specifier, module] of Object.entries(modules)) { const name = `\`${specifier}\``; @@ -137,6 +142,13 @@ export function parseModules(modules: Readonly>) `Host module ${JSON.stringify(specifier)} must use a simple codemode:* name, such as "codemode:git".`, ); } + if (isScopedModule(module)) { + assertExportNames(specifier, module.exports); + scoped.set(specifier, module); + const exports = module.exports.map((key) => `\`${key}\``).join(", "); + lines.push(`- ${name}: ${module.description ?? `exports ${exports}.`}`); + continue; + } if (typeof module === "function") { host.set(specifier, module); lines.push(() => `- ${name}: ${module.description ?? "a host module."}`); @@ -155,6 +167,7 @@ export function parseModules(modules: Readonly>) return { source, host, + scoped, get description() { return lines.map((line) => (typeof line === "string" ? line : line())).join("\n"); }, @@ -174,20 +187,29 @@ export function assertHostModuleExports( functions: WorkspaceModuleFunctions, ): void { const names = Object.keys(functions); + assertExportNames(specifier, names); + for (const name of names) { + if (typeof functions[name] !== "function") { + throw new Error( + `Host module ${JSON.stringify(specifier)} export ${JSON.stringify(name)} must be a function.`, + ); + } + } +} + +// Check a host module's export names. +function assertExportNames(specifier: string, names: readonly string[]): void { if (names.length === 0) { throw new Error(`Host module ${JSON.stringify(specifier)} must export a function.`); } + if (new Set(names).size !== names.length) + throw new Error(`Host module ${JSON.stringify(specifier)} exports a name twice.`); for (const name of names) { - if (!EXPORT_NAME.test(name) || RESERVED_EXPORT_NAMES.has(name)) { + if (typeof name !== "string" || !EXPORT_NAME.test(name) || RESERVED_EXPORT_NAMES.has(name)) { throw new Error( `Host module ${JSON.stringify(specifier)} export ${JSON.stringify(name)} must be a JavaScript identifier other than "default" or "then".`, ); } - if (typeof functions[name] !== "function") { - throw new Error( - `Host module ${JSON.stringify(specifier)} export ${JSON.stringify(name)} must be a function.`, - ); - } } } @@ -196,7 +218,8 @@ export interface BuildModuleGraphOptions { cwd: string; capability: WorkspaceRuntimeCapability; configuredModules: Readonly>; - hostModules: ReadonlyMap; + /** Each host module's export names, by specifier. */ + hostModules: ReadonlyMap; /** Host modules whose factory built no functions on this connection, so they are left out. */ emptyHostModules?: ReadonlySet; /** Whether the Dynamic Worker has the Node.js built-ins. */ @@ -365,9 +388,9 @@ export async function buildModuleGraph(options: BuildModuleGraphOptions) { modules["node:fs"] = { js: nodeFsModule() }; const toCapabilities = relativeSpecifier(storedName("module"), CAPABILITIES_MODULE); - for (const [specifier, functions] of options.hostModules) { + for (const [specifier, names] of options.hostModules) { modules[storedName(specifier)] = { - js: hostModule(toCapabilities, specifier, Object.keys(functions)), + js: hostModule(toCapabilities, specifier, names), }; } for (const [specifier, source] of Object.entries(options.configuredModules)) { diff --git a/packages/computer/src/backends/worker-javascript/worker-javascript.test.ts b/packages/computer/src/backends/worker-javascript/worker-javascript.test.ts index 499f045a..49be3474 100644 --- a/packages/computer/src/backends/worker-javascript/worker-javascript.test.ts +++ b/packages/computer/src/backends/worker-javascript/worker-javascript.test.ts @@ -2,6 +2,7 @@ import { Database, initializeSchema, WorkspaceFilesystem } from "@cloudflare/dof import { SQLiteTestStorage } from "@cloudflare/dofs/testing"; import { describe, expect, it, vi } from "vitest"; +import { createScopedModule } from "../../runtime/scoped-module"; import type { WorkspaceRuntimeLoader } from "../../runtime/types"; import { Workspace } from "../../workspace"; import { ExecutionJournal, fingerprint, loadResume, type ResumeStart, saveResume } from "./journal"; @@ -1048,6 +1049,26 @@ describe("WorkerJavaScriptBackend", () => { { "codemode:test": { run: "nope" } }, /must be a function/, ], + [ + "a scoped module with a non-identifier export", + { + "codemode:test": createScopedModule({ + exports: ["not-a-name"], + start: () => ({ functions: {} }), + }), + }, + /JavaScript identifier/, + ], + [ + "a scoped module that exports a name twice", + { + "codemode:test": createScopedModule({ + exports: ["run", "run"], + start: () => ({ functions: {} }), + }), + }, + /exports a name twice/, + ], ])("rejects %s at construction", (_label, modules, message) => { expect( () => @@ -1977,8 +1998,12 @@ describe("WorkerJavaScriptBackend", () => { new WorkerJavaScriptBackend({ loader: { load }, modules: { - "codemode:facade": `import { run } from "codemode:echo"; export const echo = { run };`, + "codemode:facade": `import { run } from "codemode:echo"; import { open } from "codemode:browser"; export const echo = { run, open };`, "codemode:echo": { run: async () => null }, + "codemode:browser": createScopedModule({ + exports: ["open"], + start: () => ({ functions: { open: () => null } }), + }), }, }), ], @@ -1992,7 +2017,7 @@ describe("WorkerJavaScriptBackend", () => { const modules = load.mock.calls[0]?.[0].modules ?? {}; expect(source(modules["__modules__/codemode:facade"])).toBe( - `import { run } from "./codemode:echo"; export const echo = { run };`, + `import { run } from "./codemode:echo"; import { open } from "./codemode:browser"; export const echo = { run, open };`, ); }); @@ -2443,3 +2468,314 @@ describe("WorkerJavaScriptBackend interrupted executions", () => { await workspace.close(); }); }); +describe("WorkerJavaScriptBackend scoped modules", () => { + interface Script { + // Host functions the run calls in turn, such as "host/codemode:test.run". + calls?: string[]; + // Where the run puts each call's response. + responses?: unknown[]; + stdout?: string; + // Hang after the calls, until the run is killed. + hang?: boolean; + } + + interface ScriptHost { + call(name: string, args: unknown[]): Promise; + attachOutput(readable: ReadableStream): Promise; + } + + // A stream of output frames, one per line, as the real runner writes. + function frames(...lines: string[]) { + const encoder = new TextEncoder(); + return new ReadableStream({ + start(controller) { + for (const line of lines) controller.enqueue(encoder.encode(`${line}\n`)); + controller.close(); + }, + }); + } + + // A loader whose runs make the script's host calls, write its stdout, + // and exit 0, the way the real runner frames output. + function scripted(script: Script) { + return { + load() { + return { + getEntrypoint() { + return { + async evaluate(_input: unknown, host: ScriptHost) { + for (const name of script.calls ?? []) { + const response = await host.call(name, []); + script.responses?.push(response); + } + if (script.hang === true) await new Promise(() => undefined); + const lines: string[] = []; + if (script.stdout !== undefined) { + lines.push(JSON.stringify({ name: "stdout", b64: btoa(script.stdout) })); + } + lines.push(JSON.stringify({ name: "exit", code: 0 })); + await host.attachOutput(frames(...lines)); + }, + }; + }, + }; + }, + }; + } + + function open( + script: Script, + modules: ConstructorParameters[0]["modules"], + options: Partial[0]> = {}, + ) { + return new Workspace({ + storage: new SQLiteTestStorage(), + backends: [new WorkerJavaScriptBackend({ loader: scripted(script), modules, ...options })], + }); + } + + async function start(workspace: Workspace, options: { onInterrupt?: "replay" } = {}) { + await workspace.fs.mkdir("/workspace", { recursive: true }); + return workspace.runtime.exec("export default 1", { + backend: "worker-javascript", + id: "run", + encoding: "utf8", + ...options, + }); + } + + function gate() { + let open!: () => void; + const opened = new Promise((resolve) => (open = resolve)); + return { open, opened }; + } + + // Let every task that isn't waiting on a gate run. + const idle = () => new Promise((resolve) => setTimeout(resolve, 0)); + + // A scoped module whose instance logs `text` when the run ends. + function logging(...texts: string[]) { + return createScopedModule({ + exports: ["run"], + start: () => ({ + functions: { run: () => null }, + close: (end) => { + if (end.reason !== "paused") for (const text of texts) end.log(text); + }, + }), + }); + } + + it("fails a call when the instance's functions don't match the module's exports", async () => { + const responses: unknown[] = []; + const workspace = open( + { calls: ["host/codemode:test.run"], responses }, + { + "codemode:test": createScopedModule({ + exports: ["run"], + start: () => ({ functions: { other: () => null } }), + }), + }, + ); + + await (await start(workspace)).result(); + + expect(responses).toEqual([ + expect.objectContaining({ + error: expect.objectContaining({ + message: + 'Host module "codemode:test" started with functions other, but declares exports run.', + }), + }), + ]); + await workspace.close(); + }); + + it("keeps what an instance logs whole when it fits the run's output limit", async () => { + const workspace = open( + { calls: ["host/codemode:test.run"], stdout: "123456" }, + { "codemode:test": logging("ok", "!\n") }, + { maxStdioBytes: 10 }, + ); + await expect((await start(workspace)).result()).resolves.toMatchObject({ + stdout: "123456", + stderr: "ok!\n", + }); + await workspace.close(); + }); + + it("cuts what an instance logs at the run's output limit, with a note", async () => { + const note = "Host module output truncated.\n"; + const workspace = open( + { calls: ["host/codemode:test.run"], stdout: "123456" }, + { "codemode:test": logging("ab\n", "cdefghijklmnopqrstuvwxyz".repeat(2), "more") }, + { maxStdioBytes: 6 + 3 + 6 + note.length }, + ); + + await expect((await start(workspace)).result()).resolves.toMatchObject({ + stdout: "123456", + stderr: `ab\ncdefg\n${note}`, + }); + await workspace.close(); + }); + + it("reports a close that throws or runs late, and the run goes on", async () => { + const reported = vi.spyOn(console, "error").mockImplementation(() => undefined); + const close = (fn: () => Promise) => + createScopedModule({ + exports: ["run"], + start: () => ({ functions: { run: () => null }, close: fn }), + }); + const workspace = open( + { + calls: ["host/codemode:broken.run", "host/codemode:slow.run", "host/codemode:fine.run"], + }, + { + "codemode:broken": close(async () => { + throw new Error("no browser"); + }), + "codemode:slow": close(() => new Promise(() => undefined)), + "codemode:fine": logging("fine\n"), + }, + { maxHostCallMs: 20 }, + ); + + try { + await expect((await start(workspace)).result()).resolves.toMatchObject({ + exitCode: 0, + stderr: "fine\n", + }); + expect(reported).toHaveBeenCalledWith( + expect.objectContaining({ + message: "Host module close failed", + module: "codemode:broken", + error: "Error: no browser", + }), + ); + expect(reported).toHaveBeenCalledWith( + expect.objectContaining({ message: "Host module close timed out", id: "run" }), + ); + } finally { + reported.mockRestore(); + await workspace.close(); + } + }); + + it("aborts a start's signal when the run is killed, so the kill doesn't wait out the start", async () => { + const starting = gate(); + const workspace = open( + { calls: ["host/codemode:test.run"] }, + { + "codemode:test": createScopedModule({ + exports: ["run"], + start: ({ signal }) => { + starting.open(); + return new Promise((_resolve, reject) => + signal.addEventListener("abort", () => reject(signal.reason)), + ); + }, + }), + }, + ); + const handle = await start(workspace); + await starting.opened; + + await workspace.runtime.killExec("run", { backend: "worker-javascript" }); + + await expect(handle.result()).resolves.toMatchObject({ status: "cancelled" }); + await workspace.close(); + }); + + it("closes the instance of a run killed while its interrupted call cleans up, after the cleanup", async () => { + const steps: string[] = []; + const cleaning = gate(); + const cleanup = gate(); + const workspace = open( + { calls: ["host/codemode:test.wait"], hang: true }, + { + "codemode:test": createScopedModule({ + exports: ["wait"], + start: () => ({ + functions: { + wait: async (_args, context) => { + try { + return await context.interrupt("approve?"); + } finally { + steps.push("cleanup starts"); + cleaning.open(); + await cleanup.opened; + steps.push("cleanup ends"); + } + }, + }, + close: (end) => { + steps.push(`close ${end.reason}`); + if (end.reason !== "paused") end.log("Closed.\n"); + }, + }), + }), + }, + ); + const handle = await start(workspace, { onInterrupt: "replay" }); + let exited = false; + const result = handle.result().then((settled) => { + exited = true; + return settled; + }); + await cleaning.opened; + + const killed = workspace.runtime.killExec("run", { backend: "worker-javascript" }); + await idle(); + expect(steps).toEqual(["cleanup starts"]); + expect(exited).toBe(false); + cleanup.open(); + + await killed; + await expect(result).resolves.toMatchObject({ + status: "cancelled", + exitCode: 130, + stderr: "Closed.\nExecution cancelled.\n", + }); + expect(steps).toEqual(["cleanup starts", "cleanup ends", "close cancelled"]); + await workspace.close(); + }); + + it("settles a run whose code finishes while a call it didn't await interrupts", async () => { + // The call prepares for `yields` turns of the microtask queue, so + // across the range its interrupt lands before, while, and after the + // run's completion is handled. + for (let yields = 0; yields <= 20; yields += 1) { + const interrupting = gate(); + const evaluate = async (_input: unknown, host: ScriptHost) => { + void host.call("host/codemode:test.wait", []).catch(() => undefined); + await host.attachOutput(frames('{"name":"exit","code":0}')); + }; + const workspace = new Workspace({ + storage: new SQLiteTestStorage(), + backends: [ + new WorkerJavaScriptBackend({ + loader: { load: () => ({ getEntrypoint: () => ({ evaluate }) }) }, + modules: { + "codemode:test": { + wait: async (_args, context) => { + for (let turn = 0; turn < yields; turn += 1) await Promise.resolve(); + interrupting.open(); + return context.interrupt("approve?"); + }, + }, + }, + }), + ], + }); + const handle = await start(workspace, { onInterrupt: "replay" }); + await interrupting.opened; + await idle(); + + await workspace.runtime.killExec("run", { backend: "worker-javascript" }); + await expect(handle.result()).resolves.toMatchObject({ + status: expect.stringMatching(/^(completed|cancelled)$/), + }); + await workspace.close(); + } + }); +}); diff --git a/packages/computer/src/backends/worker-javascript/worker-javascript.ts b/packages/computer/src/backends/worker-javascript/worker-javascript.ts index a2a79f8c..4a51f916 100644 --- a/packages/computer/src/backends/worker-javascript/worker-javascript.ts +++ b/packages/computer/src/backends/worker-javascript/worker-javascript.ts @@ -6,6 +6,11 @@ import { hostEgressFetch, type WorkspaceEgressPolicy, } from "../../runtime/egress"; +import type { + WorkspaceModuleScopeEnd, + WorkspaceScopedModuleInstance, + WorkspaceScopedModule, +} from "../../runtime/scoped-module"; import type { ModuleExecutionEnvelope, ModuleExecutionInput, @@ -16,6 +21,7 @@ import type { WorkspaceModuleBackendHost, WorkspaceModuleExecution, WorkspaceModuleFunctions, + WorkspaceModuleHost, WorkspaceRuntimeAccess, WorkspaceRuntimeEvent, WorkspaceRuntimeInterruptPolicy, @@ -250,6 +256,11 @@ interface JavaScriptEntrypoint { interface ActiveControl { cancel(): void; + // Settles once the isolate has stopped and been disposed, before the + // run's completion callbacks start. Never rejects. + readonly stopped: Promise; + // Settles once the completion callbacks, which settle the run, are + // done. readonly completion: Promise; } @@ -283,6 +294,14 @@ interface ExecutionRecord { // Set while a run's isolate is being released for an interrupt, and // once it is, until the run resumes. waiting?: "releasing" | "released"; + // Releasing the run's isolate, from when it starts letting go of the + // isolate's host calls until its scoped instances have closed. A + // run that ends meanwhile waits for it, so its end follows them. + releasing?: Promise; + // The scoped module instances the run's current scope started, by + // specifier. Whoever closes them takes them off the record first, so + // each closes once. + instances?: Map>; } export class WorkerJavaScriptBackend implements WorkspaceModuleBackend { @@ -420,6 +439,9 @@ class JavaScriptBackendHandle implements WorkspaceModuleBackendHandle { readonly #host: WorkspaceModuleBackendHost; readonly #hostModuleFunctions: ReadonlyMap; readonly #emptyHostModules: ReadonlySet; + // Every host module's export names, for the module graph. + readonly #hostModuleExports: ReadonlyMap; + readonly #moduleHost: WorkspaceModuleHost; #sourceModules: Readonly> | undefined; readonly #records = new Map(); readonly #pendingIds = new Set(); @@ -442,13 +464,15 @@ class JavaScriptBackendHandle implements WorkspaceModuleBackendHandle { this.#host = host; const functions = new Map(); const empty = new Set(); + // Shared by every factory and scoped module start, so frozen. + this.#moduleHost = Object.freeze({ + git: host.git, + artifacts: host.artifacts, + runtime: host.runtime, + ...(host.assets === undefined ? {} : { assets: host.assets }), + }); for (const [specifier, factory] of options.modules.host) { - const built = factory({ - git: host.git, - artifacts: host.artifacts, - runtime: host.runtime, - ...(host.assets === undefined ? {} : { assets: host.assets }), - }); + const built = factory(this.#moduleHost); // A factory can have nothing to offer on a connection, such as // codemode:tools for an agent with no tools. Leaving its module out means // only an import of it fails, not every run. An object of functions @@ -463,6 +487,12 @@ class JavaScriptBackendHandle implements WorkspaceModuleBackendHandle { } this.#hostModuleFunctions = functions; this.#emptyHostModules = empty; + const exports = new Map(); + for (const [specifier, built] of functions) exports.set(specifier, Object.keys(built)); + for (const [specifier, module] of options.modules.scoped) { + exports.set(specifier, module.exports); + } + this.#hostModuleExports = exports; host.db.run(` CREATE TABLE IF NOT EXISTS workspace_runtime_executions ( backend TEXT NOT NULL, @@ -599,15 +629,125 @@ class JavaScriptBackendHandle implements WorkspaceModuleBackendHandle { await bridge.quiesce(); // It may have settled meanwhile, for example by failing on its own. if (record.status !== "running" || record.finalization) return; + // From here the bridge is off the record, so a run that ends + // meanwhile can't drain it, and waits for this instead. + const releasing = this.#releaseIsolate(record, bridge); + record.releasing = releasing; + await releasing; + record.releasing = undefined; + if (record.status !== "running" || record.finalization) return; + record.waiting = "released"; + this.#markSettled(record.id); + this.#resumeIfAnswered(record); + } + + async #releaseIsolate(record: ExecutionRecord, bridge: WorkspaceRuntimeBridge) { record.bridge = undefined; record.control?.cancel(); - await record.control?.completion.catch(() => undefined); + // Only until the isolate stops: its completion callbacks may be + // settling this run, and that waits for this release. + await record.control?.stopped; await bridge.cancelAndDrain().catch(() => undefined); record.control = undefined; record.journal = undefined; - record.waiting = "released"; - this.#markSettled(record.id); - this.#resumeIfAnswered(record); + // The instances close before the run counts as released, so an + // answer that arrives meanwhile can't resume it under them. A run + // that ended while its host calls drained closes them as it ends. + if (record.status !== "running" || record.finalization) return; + await this.#closeInstances(record, "paused"); + } + + // Close the scoped module instances the run's scope started, all at + // once, and wait for them up to `maxHostCallMs`. A close that + // throws or runs late is reported and otherwise ignored. Returns what + // they logged, cut to what is left of the run's output limit. + async #closeInstances( + record: ExecutionRecord, + reason: WorkspaceModuleScopeEnd["reason"], + ): Promise { + const instances = record.instances; + record.instances = undefined; + if (instances === undefined || instances.size === 0) return []; + const output = closeOutput( + Math.max(0, this.#options.maxStdioBytes - stdioBytes(record.events)), + ); + let open = true; + const timeoutMs = this.#options.maxHostCallMs; + const controller = new AbortController(); + const deadline = Date.now() + timeoutMs; + const settled = Promise.all( + [...instances].map(async ([specifier, started]) => { + // Each instance's log closes when its close returns, even while + // another module's close is still running. + let active = true; + const log = (text: string) => { + if (open && active) output.write(String(text)); + }; + try { + // An instance whose start failed has nothing to close. + const instance = await started.catch(() => undefined); + await instance?.close?.( + Object.freeze( + reason === "paused" + ? { reason, signal: controller.signal, deadline } + : { reason, signal: controller.signal, deadline, log }, + ), + ); + } catch (error) { + console.error({ + message: "Host module close failed", + module: specifier, + reason, + id: record.id, + error: String(error), + }); + } finally { + active = false; + } + }), + ).then(() => false); + let timer: ReturnType | undefined; + const late = new Promise((resolve) => { + timer = setTimeout(() => resolve(true), timeoutMs); + }); + try { + if (await Promise.race([settled, late])) { + controller.abort(new Error(`Host module close passed its ${timeoutMs}ms deadline.`)); + console.error({ message: "Host module close timed out", reason, id: record.id, timeoutMs }); + } + } finally { + open = false; + clearTimeout(timer); + } + return output.written(); + } + + // Settle a run held in memory. Its scoped instances close once + // its host calls have settled, and what they log goes ahead of the + // run's last events, in the same transaction as its terminal state, + // so a restart can't replay the run against output it never wrote. + async #end( + record: ExecutionRecord, + status: "completed" | "failed" | "cancelled", + tail: TailEvent[], + ) { + // A run released for an interrupt ends once its released calls and + // closes are done. + if (record.releasing !== undefined) await record.releasing; + const logged = record.instances?.size ? await this.#closeInstances(record, status) : []; + const events: TailEvent[] = [ + ...logged.map((value) => ({ name: "stderr" as const, value })), + ...tail, + ]; + this.#finish( + record, + status, + events.map((event, index) => ({ + ...event, + id: record.id, + seq: record.events.length + 1 + index, + })), + ); } #notifyWait(record: ExecutionRecord, call: { seq: number; payload: unknown; since: number }) { @@ -808,7 +948,7 @@ class JavaScriptBackendHandle implements WorkspaceModuleBackendHandle { #preparedSourceModules(): Readonly> { this.#sourceModules ??= prepareSourceModules( this.#options.modules.source, - new Set(this.#options.modules.host.keys()), + new Set([...this.#options.modules.host.keys(), ...this.#options.modules.scoped.keys()]), ); return this.#sourceModules; } @@ -862,7 +1002,7 @@ class JavaScriptBackendHandle implements WorkspaceModuleBackendHandle { cwd: input.cwd ?? this.#options.root, capability, configuredModules: this.#preparedSourceModules(), - hostModules: this.#hostModuleFunctions, + hostModules: this.#hostModuleExports, emptyHostModules: this.#emptyHostModules, nodeModules: hasNodeModules(this.#options.compatibilityFlags), maxSourceBytes: this.#options.maxSourceBytes, @@ -954,6 +1094,27 @@ class JavaScriptBackendHandle implements WorkspaceModuleBackendHandle { } } + // This scope's scoped modules, each as the function that gets its + // instance. The first live call to a module starts it, + // and calls made meanwhile share that start. + #startInstances(record: ExecutionRecord, execution: WorkspaceModuleExecution) { + const instances = new Map>(); + record.instances = instances; + const starters = new Map Promise>(); + for (const [specifier, module] of this.#options.modules.scoped) { + starters.set(specifier, async (signal) => { + let started = instances.get(specifier); + if (started === undefined) { + const scope = Object.freeze({ ...execution, signal }); + started = Promise.resolve().then(() => module.start(scope, this.#moduleHost)); + instances.set(specifier, started); + } + return instanceFunctions(specifier, module, await started); + }); + } + return starters; + } + // Start a run's Dynamic Worker. A failure to start settles the run. async #launch( record: ExecutionRecord, @@ -970,8 +1131,10 @@ class JavaScriptBackendHandle implements WorkspaceModuleBackendHandle { ): Promise { this.#markRunning(record.id, spec.timeoutMs); try { + const execution = record.execution ?? this.#execution(record.id, undefined); const bridge = new WorkspaceRuntimeBridge(spec.capability, { hostModules: this.#hostModuleFunctions, + scopedModules: this.#startInstances(record, execution), maxPayloadBytes: this.#options.maxCapabilityBytes, maxCallDurationMs: this.#options.maxHostCallMs, maxConcurrentCalls: this.#options.maxConcurrentCapabilityCalls, @@ -980,7 +1143,7 @@ class JavaScriptBackendHandle implements WorkspaceModuleBackendHandle { maxTotalResponseBytes: this.#options.maxCapabilityResponseBytes, maxResultBytes: this.#options.maxResultBytes, onAttachOutput: (readable) => this.#pumpFrames(record, readable), - execution: record.execution ?? this.#execution(record.id, undefined), + execution, onWait: (call) => { this.#notifyWait(record, call); void this.#release(record); @@ -1123,27 +1286,12 @@ class JavaScriptBackendHandle implements WorkspaceModuleBackendHandle { try { await record.bridge?.cancelAndDrain(); } catch (error) { - this.#finish(record, "failed", [ - { - id: record.id, - seq: record.events.length + 1, - name: "stderr", - value: new TextEncoder().encode( - `${error instanceof Error ? error.message : String(error)}\n`, - ), - }, - { id: record.id, seq: record.events.length + 2, name: "exit", code: 1 }, - ]); + await this.#end(record, "failed", failure(error)); return; } - this.#finish(record, "cancelled", [ - { - id: record.id, - seq: record.events.length + 1, - name: "stderr", - value: new TextEncoder().encode(message), - }, - { id: record.id, seq: record.events.length + 2, name: "exit", code: 130 }, + await this.#end(record, "cancelled", [ + { name: "stderr", value: new TextEncoder().encode(message) }, + { name: "exit", code: 130 }, ]); })(); record.finalization = finalization; @@ -1305,31 +1453,15 @@ class JavaScriptBackendHandle implements WorkspaceModuleBackendHandle { try { await record.bridge?.cancelAndDrain(); } catch (error) { - this.#finish(record, "failed", [ - { - id: record.id, - seq: record.events.length + 1, - name: "stderr", - value: new TextEncoder().encode( - `${error instanceof Error ? error.message : String(error)}\n`, - ), - }, - { id: record.id, seq: record.events.length + 2, name: "exit", code: 1 }, - ]); + await this.#end(record, "failed", failure(error)); return; } if (errorMessage !== undefined) { - this.#finish(record, "failed", [ - { - id: record.id, - seq: record.events.length + 1, - name: "stderr", - value: new TextEncoder().encode( - `${truncateUtf8(errorMessage, Math.max(0, this.#options.maxStdioBytes - 1))}\n`, - ), - }, - { id: record.id, seq: record.events.length + 2, name: "exit", code: 1 }, - ]); + await this.#end( + record, + "failed", + failure(truncateUtf8(errorMessage, Math.max(0, this.#options.maxStdioBytes - 1))), + ); return; } // No exit frame means the runner never reported a terminal state: @@ -1337,29 +1469,15 @@ class JavaScriptBackendHandle implements WorkspaceModuleBackendHandle { // isolate crashed. Settle as a failure rather than a silent exit 0 // with no result. if (record.exitCode === undefined) { - this.#finish(record, "failed", [ - { - id: record.id, - seq: record.events.length + 1, - name: "stderr", - value: new TextEncoder().encode("Execution ended without reporting a result.\n"), - }, - { id: record.id, seq: record.events.length + 2, name: "exit", code: 1 }, - ]); + await this.#end(record, "failed", failure("Execution ended without reporting a result.")); return; } const exitCode = record.exitCode; - const exit: WorkspaceRuntimeEvent = + const exit: TailEvent = exitCode === 0 && record.hasResult - ? { - id: record.id, - seq: record.events.length + 1, - name: "exit", - code: exitCode, - result: record.result as WorkspaceRuntimeValue, - } - : { id: record.id, seq: record.events.length + 1, name: "exit", code: exitCode }; - this.#finish(record, exitCode === 0 ? "completed" : "failed", [exit]); + ? { name: "exit", code: exitCode, result: record.result as WorkspaceRuntimeValue } + : { name: "exit", code: exitCode }; + await this.#end(record, exitCode === 0 ? "completed" : "failed", [exit]); } #stream(record: ExecutionRecord, after?: number | "tail") { @@ -1668,10 +1786,15 @@ function startJavaScriptExecution(options: { options.timeoutMs, ); }); - const completion = Promise.race([run, timeout, cancellation]) + const race = Promise.race([run, timeout, cancellation]); + const stop = () => { + if (timer !== undefined) clearTimeout(timer); + dispose(); + }; + const stopped = race.then(stop, stop); + const completion = race .then(async () => { - if (timer !== undefined) clearTimeout(timer); - dispose(); + stop(); await options.onComplete(); }) .catch(async (error) => { @@ -1695,6 +1818,7 @@ function startJavaScriptExecution(options: { // settled the race; swallow that late rejection. run.catch(() => undefined); return { + stopped, completion, cancel() { if (cancelled) return; @@ -1948,6 +2072,85 @@ function assertLoaderGraph( } } +// A scoped module instance's functions, checked against the exports its +// module declared, which the module graph was built from. +function instanceFunctions( + specifier: string, + module: WorkspaceScopedModule, + instance: WorkspaceScopedModuleInstance, +): WorkspaceModuleFunctions { + const functions: unknown = instance?.functions; + const names = + typeof functions === "object" && functions !== null ? Object.keys(functions).sort() : []; + const exports = [...module.exports].sort(); + const matches = + names.length === exports.length && + names.every( + (name, index) => + name === exports[index] && + typeof (functions as Record)[name] === "function", + ); + if (!matches) { + throw new Error( + `Host module ${JSON.stringify(specifier)} started with functions ${names.join(", ") || "(none)"}, but declares exports ${exports.join(", ")}.`, + ); + } + // SAFETY: Checked above: an object of the declared functions. + return functions as WorkspaceModuleFunctions; +} + +// One of the last events of a run, before #end numbers it. +type TailEvent = + | { name: "stderr"; value: Uint8Array } + | { name: "exit"; code: number; result?: WorkspaceRuntimeValue }; + +// The last events of a run that failed with an error. +function failure(error: unknown): TailEvent[] { + const message = error instanceof Error ? error.message : String(error); + return [ + { name: "stderr", value: new TextEncoder().encode(`${message}\n`) }, + { name: "exit", code: 1 }, + ]; +} + +const CLOSE_OUTPUT_TRUNCATED = "Host module output truncated.\n"; + +// What a run's scoped instances log as they close, kept within +// `room` bytes. Text is held until they are done; if it doesn't all +// fit, it's cut and ends with a note saying so, and the whole is cut +// at `room`. +function closeOutput(room: number) { + const encoder = new TextEncoder(); + const chunks: Uint8Array[] = []; + let bytes = 0; + return { + write(text: string) { + if (bytes > room || text === "") return; + const chunk = encoder.encode(text); + chunks.push(chunk); + bytes += chunk.byteLength; + }, + written(): Uint8Array[] { + if (bytes <= room) return chunks; + const text = chunks.map((chunk) => new TextDecoder().decode(chunk)).join(""); + const note = encoder.encode(CLOSE_OUTPUT_TRUNCATED).byteLength; + let cut = truncateUtf8(text, Math.max(0, room - note - 1)); + if (cut !== "" && !cut.endsWith("\n")) cut += "\n"; + const kept = truncateUtf8(cut + CLOSE_OUTPUT_TRUNCATED, room); + return kept === "" ? [] : [encoder.encode(kept)]; + }, + }; +} + +// The bytes of stdout and stderr a run has written so far. +function stdioBytes(events: readonly WorkspaceRuntimeEvent[]): number { + let bytes = 0; + for (const event of events) { + if (event.name === "stdout" || event.name === "stderr") bytes += event.value.byteLength; + } + return bytes; +} + function truncateUtf8(value: string, maxBytes: number) { const bytes = new TextEncoder().encode(value); if (bytes.byteLength <= maxBytes) return value; diff --git a/packages/computer/src/index.ts b/packages/computer/src/index.ts index 9e475d71..49b3ef8e 100644 --- a/packages/computer/src/index.ts +++ b/packages/computer/src/index.ts @@ -119,6 +119,14 @@ export type { WorkspaceRuntimeValue, } from "./runtime/types"; export { unsafeToRepeat } from "./runtime/replay-safety"; +export { + createScopedModule, + type WorkspaceModuleScopeEnd, + type WorkspaceScopedModuleInstance, + type WorkspaceScopedModule, + type WorkspaceScopedModuleOptions, + type WorkspaceModuleScope, +} from "./runtime/scoped-module"; export { decodeRuntimeEvents, encodeRuntimeEvent } from "./runtime/wire"; export { type RawShellValue, type ShellValue, sh, shellQuote } from "./sh"; export type { ExecEncoding, ExecSyncResult, KillSignal } from "./shell"; diff --git a/packages/computer/src/runtime/bridge.ts b/packages/computer/src/runtime/bridge.ts index febcb2bc..9f25b1b0 100644 --- a/packages/computer/src/runtime/bridge.ts +++ b/packages/computer/src/runtime/bridge.ts @@ -13,6 +13,13 @@ import type { export class WorkspaceRuntimeBridge extends RpcTarget { readonly #capability: WorkspaceRuntimeCapability; readonly #hostModules: ReadonlyMap; + readonly #scopedModules: ReadonlyMap< + string, + (signal: AbortSignal) => Promise + >; + // Aborts when the scope is cancelled or drained, for scoped module + // starts, which the scope's calls share. + readonly #scope = new AbortController(); readonly #maxPayloadBytes: number; readonly #maxCallDurationMs: number; readonly #maxConcurrentCalls: number; @@ -48,6 +55,13 @@ export class WorkspaceRuntimeBridge extends RpcTarget { capability: WorkspaceRuntimeCapability, integrations: { hostModules?: ReadonlyMap; + // Scoped host modules, each as the function that gets this + // scope's instance of it, starting it on the first call with a + // signal that aborts when the scope is cancelled. + scopedModules?: ReadonlyMap< + string, + (signal: AbortSignal) => Promise + >; maxPayloadBytes?: number; maxCallDurationMs?: number; maxConcurrentCalls?: number; @@ -70,6 +84,7 @@ export class WorkspaceRuntimeBridge extends RpcTarget { super(); this.#capability = capability; this.#hostModules = integrations.hostModules ?? new Map(); + this.#scopedModules = integrations.scopedModules ?? new Map(); this.#maxPayloadBytes = integrations.maxPayloadBytes ?? 1024 * 1024; this.#maxCallDurationMs = integrations.maxCallDurationMs ?? 30_000; this.#maxConcurrentCalls = integrations.maxConcurrentCalls ?? 16; @@ -206,7 +221,7 @@ export class WorkspaceRuntimeBridge extends RpcTarget { let interrupted = false; const operation = respond(async () => { if (name.startsWith("host/")) { - return this.#callHostModule(name, args, { + return this.#callHostModule(name, args, step?.kind === "live" && step.rerun, { signal: abort.signal, deadline, callId, @@ -407,6 +422,7 @@ export class WorkspaceRuntimeBridge extends RpcTarget { async cancelAndDrain(): Promise { this.#cancelled = true; + this.#scope.abort(new Error("Workspace execution is being cancelled.")); for (const controller of this.#abortControllers) { controller.abort(new Error("Workspace execution is being cancelled.")); } @@ -420,25 +436,38 @@ export class WorkspaceRuntimeBridge extends RpcTarget { // identifiers and never contain a dot, so the last dot splits them // from a specifier such as `codemode:a.b`. Own-property checks keep // isolate code from reaching `toString` or other inherited members. - async #callHostModule(name: string, args: unknown[], context: WorkspaceModuleCallContext) { - const fn = this.#hostFunction(name); + async #callHostModule( + name: string, + args: unknown[], + rerun: boolean, + context: WorkspaceModuleCallContext, + ) { + const start = this.#scopedModules.get(moduleSpecifier(name) ?? ""); + const functions = + start === undefined ? this.#hostModuleOf(name) : await start(this.#scope.signal); + // The call may have been cancelled or timed out while the instance started. + context.signal.throwIfAborted(); + const fn = this.#hostFunction(name, functions); if (fn === undefined) { throw new Error(`Unknown Workspace host module call ${JSON.stringify(name)}.`); } + // #dispatch can't see a scoped module function until its instance starts. + if (start !== undefined && rerun && isUnsafeToRepeat(fn)) { + const message = `The execution was interrupted while ${name.slice("host/".length)} was running, which isn't safe to repeat.`; + this.#journal?.diverged(message); + throw new Error(message); + } // Called on its module, so a method that uses `this` still works. - const result = await fn.call(this.#hostModuleOf(name), hostArguments(args), context); + const result = await fn.call(functions, hostArguments(args), context); return hostValue(result ?? null); } #hostModuleOf(name: string) { - if (!name.startsWith("host/")) return undefined; - const target = name.slice("host/".length); - const dot = target.lastIndexOf("."); - return dot === -1 ? undefined : this.#hostModules.get(target.slice(0, dot)); + const specifier = moduleSpecifier(name); + return specifier === undefined ? undefined : this.#hostModules.get(specifier); } - #hostFunction(name: string) { - const functions = this.#hostModuleOf(name); + #hostFunction(name: string, functions = this.#hostModuleOf(name)) { const functionName = name.slice(name.lastIndexOf(".") + 1); const fn = functions !== undefined && Object.hasOwn(functions, functionName) @@ -531,6 +560,14 @@ async function readLimited(response: Response, maxBytes: number): Promise.`. +function moduleSpecifier(name: string): string | undefined { + if (!name.startsWith("host/")) return undefined; + const target = name.slice("host/".length); + const dot = target.lastIndexOf("."); + return dot === -1 ? undefined : target.slice(0, dot); +} + function hostArguments(args: unknown[]): WorkspaceRuntimeValue[] { return args.map((arg) => hostValue(arg ?? null)); } diff --git a/packages/computer/src/runtime/scoped-module.ts b/packages/computer/src/runtime/scoped-module.ts new file mode 100644 index 00000000..cc765852 --- /dev/null +++ b/packages/computer/src/runtime/scoped-module.ts @@ -0,0 +1,134 @@ +import type { + WorkspaceModuleExecution, + WorkspaceModuleFunctions, + WorkspaceModuleHost, +} from "./types"; + +declare const scopedModuleBrand: unique symbol; + +// Marks what `createScopedModule()` returns. A symbol, so it can't +// collide with an export name, and a registered one, so two copies of +// this package agree on it. +const SCOPED_MODULE = Symbol.for("@cloudflare/computer/scoped-module"); + +/** How a scope ended, as `WorkspaceScopedModuleInstance.close` is told. */ +export type WorkspaceModuleScopeEnd = + | { + /** + * The run is waiting in `ctx.interrupt()` and has released its + * isolate. Its code starts again from the top when it resumes, in + * a new scope with a new instance. + */ + readonly reason: "paused"; + /** Aborts at `deadline`, when the backend stops waiting for `close`. */ + readonly signal: AbortSignal; + /** Epoch milliseconds after which the backend stops waiting for `close`. */ + readonly deadline: number; + } + | { + /** How the run ended. */ + readonly reason: "completed" | "failed" | "cancelled"; + readonly signal: AbortSignal; + readonly deadline: number; + /** + * Append text to the run's stderr, before its exit. Text past what + * is left of the run's output limit is cut, and ends with a line + * saying so. Anything logged after `close` returns is dropped. + */ + log(text: string): void; + }; + +/** + * One scope of a run: from when its code starts in an isolate until it + * waits in `ctx.interrupt()` or ends. A run that resumes starts a new + * scope with the same `id`. + */ +export interface WorkspaceModuleScope extends WorkspaceModuleExecution { + /** Aborts when the scope is cancelled or its isolate is released. */ + readonly signal: AbortSignal; +} + +/** A scoped module's instance for one scope, from its `start`. */ +export interface WorkspaceScopedModuleInstance { + /** The functions code imports. Their names must match the module's `exports`. */ + readonly functions: WorkspaceModuleFunctions; + /** + * Let go of what the instance holds. Called once, after the scope's + * host calls settle: when the run pauses in `ctx.interrupt()`, before + * it can resume, or when it ends, before its exit is published. A + * close that throws or misses its deadline is reported with + * `console.error`, and the run goes on as it would have. + */ + close?(end: WorkspaceModuleScopeEnd): void | Promise; +} + +/** What `createScopedModule()` takes. */ +export interface WorkspaceScopedModuleOptions { + /** What the module does and how to call it, for a model. */ + readonly description?: string; + /** + * The names code can import. The backend needs them before any run + * starts, and checks that every instance's functions match them. + */ + readonly exports: readonly string[]; + /** + * Make the instance for one scope, on that scope's first call to the + * module. A call answered from the run's journal doesn't count. Calls + * made while it starts wait for it, and if it throws, they fail. + * `scope.signal` aborts when the scope is cancelled or released, and a + * call whose own signal has aborted by then fails. + */ + start( + scope: WorkspaceModuleScope, + host: WorkspaceModuleHost, + ): WorkspaceScopedModuleInstance | Promise; +} + +/** A host module with one instance per scope; see `createScopedModule()`. */ +export interface WorkspaceScopedModule extends WorkspaceScopedModuleOptions { + readonly [scopedModuleBrand]: true; +} + +/** + * A host module with an instance for each scope of a run, for state + * such as a browser session. A scope lasts from when the run's code + * starts in an isolate until it waits in `ctx.interrupt()` or ends. The + * backend calls `start` on the scope's first call to the module, shares + * the instance with the scope's other calls, and closes it when the + * scope ends. A run that resumes starts again from the top in a new + * scope, with a new instance, so nothing held in memory outlives a + * scope. An instance closes only if it started. Experimental: it may + * change in any release. + * + * ```ts + * "codemode:browser": createScopedModule({ + * exports: ["open"], + * async start(scope, host) { + * const session = await connect(scope.metadata.conversation); + * return { + * functions: { open: ([url]) => session.open(String(url)) }, + * async close(end) { + * await session.close(); + * if (end.reason !== "paused") end.log("Closed the browser.\n"); + * }, + * }; + * }, + * }), + * ``` + */ +export function createScopedModule(options: WorkspaceScopedModuleOptions): WorkspaceScopedModule { + // SAFETY: The symbol written here is the brand the type names. + return Object.freeze({ + ...options, + exports: Object.freeze([...options.exports]), + [SCOPED_MODULE]: true, + }) as unknown as WorkspaceScopedModule; +} + +export function isScopedModule(module: unknown): module is WorkspaceScopedModule { + return ( + typeof module === "object" && + module !== null && + (module as { [SCOPED_MODULE]?: unknown })[SCOPED_MODULE] === true + ); +} diff --git a/packages/computer/src/runtime/types.ts b/packages/computer/src/runtime/types.ts index 5a725ba7..e2302eae 100644 --- a/packages/computer/src/runtime/types.ts +++ b/packages/computer/src/runtime/types.ts @@ -2,6 +2,7 @@ import type { SkippedEntry } from "@cloudflare/dofs"; import type { ExecEncoding, ExecSyncResult, KillSignal } from "../shell"; import type { TruncatedOutput } from "./output-spool"; +import type { WorkspaceScopedModule } from "./scoped-module"; export type WorkspaceRuntimeAccess = "read" | "read-write"; @@ -132,11 +133,18 @@ export interface WorkspaceModuleFactory { * `import { forecast } from "codemode:weather"`. * - A factory is a host module that needs the Workspace's Git client, * Artifacts client, or runtime, such as `createGitModule()`. + * - `createScopedModule()` makes a host module with an instance for each + * scope of a run, for state such as a session the run opens and must + * close. A scope lasts until the run waits in `ctx.interrupt()` or ends. * * Host modules must use a `codemode:*` specifier. Source may use a simple one * too, for a library the host provides. */ -export type WorkspaceModule = string | WorkspaceModuleFunctions | WorkspaceModuleFactory; +export type WorkspaceModule = + | string + | WorkspaceModuleFunctions + | WorkspaceModuleFactory + | WorkspaceScopedModule; export type WorkspaceRuntimeValue = | null diff --git a/packages/computer/tests/script-runner-worker.ts b/packages/computer/tests/script-runner-worker.ts index 50bd0325..fb375118 100644 --- a/packages/computer/tests/script-runner-worker.ts +++ b/packages/computer/tests/script-runner-worker.ts @@ -12,7 +12,7 @@ import type { WorkspaceRuntimeValue, WorkspaceStub, } from "../src/index"; -import { unsafeToRepeat, Workspace } from "../src/index"; +import { createScopedModule, unsafeToRepeat, Workspace } from "../src/index"; import { createArtifactsModule } from "../src/modules/artifacts"; import { createContainerModule } from "../src/modules/container"; import { createGitModule } from "../src/modules/git"; @@ -77,6 +77,16 @@ const executions: Array<{ metadata: WorkspaceExecutionMetadata; }> = []; +// What codemode:test-session instances did, and each approve call, in +// order, such as "start 1" or "close paused 1". +const sessionEvents: Array<{ id: string; event: string }> = []; + +// How many codemode:test-session instances each execution started. +const sessionCounts = new Map(); + +// Paused closes held until the test lets them go, by execution id. +const heldCloses = new Map void>(); + function countCall(name: string): number { const count = (durableCalls.get(name) ?? 0) + 1; durableCalls.set(name, count); @@ -111,6 +121,34 @@ function javascriptBackend( "codemode:git": createGitModule(), "codemode:artifacts": createArtifactsModule(), "codemode:container": createContainerModule(), + // Records each instance's start and close. `use()` returns the + // instance's number within its run. A run's metadata can ask it + // to hold a paused close until /session-release, and to log + // `closeLog` when the run ends. + "codemode:test-session": createScopedModule({ + exports: ["use"], + async start({ id, metadata }) { + const instance = (sessionCounts.get(id) ?? 0) + 1; + sessionCounts.set(id, instance); + sessionEvents.push({ id, event: `start ${instance}` }); + // Slow enough that a run's concurrent first calls overlap. + await new Promise((resolve) => setTimeout(resolve, 10)); + return { + functions: { use: () => instance }, + async close(end) { + sessionEvents.push({ id, event: `close ${end.reason} ${instance}` }); + if (end.reason === "paused" && metadata.holdClose !== undefined) { + await new Promise((resolve) => heldCloses.set(id, resolve)); + heldCloses.delete(id); + } + if (end.reason !== "paused" && metadata.closeLog !== undefined) { + end.log(`${metadata.closeLog}\n`); + } + sessionEvents.push({ id, event: `closed ${instance}` }); + }, + }; + }, + }), "codemode:test-host": { async echo(args) { return { args: [...args] }; @@ -190,6 +228,7 @@ function javascriptBackend( async approve(args, ctx) { const action = String(args[0]); countCall(`ask:${action}`); + sessionEvents.push({ id: ctx.execution.id, event: "approve" }); const answer = await ctx.interrupt({ kind: "approval", action }); countCall(`answered:${action}`); return answer; @@ -391,6 +430,16 @@ export class HostDO extends DurableObject { return executions; } + sessionEvents(id: string) { + return sessionEvents.filter((event) => event.id === id).map((event) => event.event); + } + + releaseClose(id: string) { + const release = heldCloses.get(id); + release?.(); + return release !== undefined; + } + killRuntime(id: string) { return this.#workspace.runtime.killExec(id, { backend: "worker-javascript" }); } @@ -687,6 +736,14 @@ export default class extends WorkerEntrypoint { return Response.json(await stub.respond(body.id, body.callId, body.value)); } + if (url.pathname === "/session-release") { + return Response.json(await stub.releaseClose(url.searchParams.get("id") ?? "")); + } + + if (url.pathname === "/session-events") { + return Response.json(await stub.sessionEvents(url.searchParams.get("id") ?? "")); + } + if (url.pathname === "/durable-executions") { return Response.json(await stub.executions()); } diff --git a/packages/computer/tests/script-runner.test.ts b/packages/computer/tests/script-runner.test.ts index 83cc1e75..3c14b281 100644 --- a/packages/computer/tests/script-runner.test.ts +++ b/packages/computer/tests/script-runner.test.ts @@ -1505,6 +1505,88 @@ describe("interrupts", () => { }); }); + async function sessionEvents(id: string) { + const response = await SELF.fetch( + `https://example.test/session-events?id=${encodeURIComponent(id)}`, + ); + return (await response.json()) as string[]; + } + + // Two concurrent first calls to a scoped module, a wait, and one + // more call to it if the answer is ok. + const session = ` + import { approve } from "codemode:test-host"; + import { use } from "codemode:test-session"; + export default async function () { + const before = await Promise.all([use(), use()]); + const answer = await approve("deploy"); + const after = answer.ok ? await use() : "none"; + console.log(before.join(" "), after); + return answer; + } + `; + + it("starts a scoped module once per scope, and resumes only after it closes", async () => { + await start("session", session, { metadata: { holdClose: "1", closeLog: "closed" } }); + const call = await waitingCall("session"); + await vi.waitFor(async () => + expect(await sessionEvents("session")).toEqual(["start 1", "approve", "close paused 1"]), + ); + // Answered while the close holds the run. + await respond("session", call.callId, { ok: true }); + const released = await SELF.fetch("https://example.test/session-release?id=session"); + expect(await released.json()).toBe(true); + + // The calls the second scope replays don't start an instance; the + // live one after the answer does. + expect(await result("session")).toMatchObject({ + status: "completed", + stdout: "1 1 2\n", + stderr: "closed\n", + }); + expect(await sessionEvents("session")).toEqual([ + "start 1", + "approve", + "close paused 1", + "closed 1", + "approve", + "start 2", + "close completed 2", + "closed 2", + ]); + }); + + it("starts no instance for a scope whose calls to the module are all replayed", async () => { + await start("session-replayed", session); + const call = await waitingCall("session-replayed"); + await respond("session-replayed", call.callId, { ok: false }); + expect(await result("session-replayed")).toMatchObject({ + status: "completed", + stdout: "1 1 none\n", + }); + expect(await sessionEvents("session-replayed")).toEqual([ + "start 1", + "approve", + "close paused 1", + "closed 1", + "approve", + ]); + }); + + it("closes nothing when a run is killed while it waits", async () => { + await start("session-kill", session); + await waitingCall("session-kill"); + await vi.waitFor(async () => expect(await sessionEvents("session-kill")).toContain("closed 1")); + await SELF.fetch("https://example.test/runtime-kill?id=session-kill"); + expect(await result("session-kill")).toMatchObject({ status: "cancelled" }); + expect(await sessionEvents("session-kill")).toEqual([ + "start 1", + "approve", + "close paused 1", + "closed 1", + ]); + }); + it("keeps waiting across a restart", async () => { await start("wait-restart", approval, { metadata: { conversation: "c2" } }); expect((await waitingCall("wait-restart")).metadata).toEqual({ conversation: "c2" });