diff --git a/packages/computer/src/runtime/runtime.ts b/packages/computer/src/runtime/runtime.ts index 9086b2c..3a7a61f 100644 --- a/packages/computer/src/runtime/runtime.ts +++ b/packages/computer/src/runtime/runtime.ts @@ -34,6 +34,14 @@ export class WorkspaceRuntime { this.#options = options; } + // Whether the named backend accepts a structured `input` value and + // returns a structured result. Consumers such as the exec tool ask + // this to know whether a backend is callable without the caller + // having to declare it a second time. + isCallable(id: string): boolean { + return this.#options.callableBackendIds.has(id); + } + exec(source: string): Promise>; exec( source: string, @@ -49,7 +57,7 @@ export class WorkspaceRuntime { ): Promise> { if (options.id !== undefined) assertExecutionId(options.id); const backend = this.#backend(options.backend); - if (options.input !== undefined && !this.#options.callableBackendIds.has(backend)) { + if (options.input !== undefined && !this.isCallable(backend)) { throw new Error(notCallableMessage(backend)); } const runtime = await this.#options.backendHandle(backend); diff --git a/packages/computer/src/tools/ai.test.ts b/packages/computer/src/tools/ai.test.ts index bb60d70..b2636f9 100644 --- a/packages/computer/src/tools/ai.test.ts +++ b/packages/computer/src/tools/ai.test.ts @@ -17,7 +17,53 @@ async function executeTool(tool: unknown, input: unknown): Promise { const execute = (tool as { execute?: (input: unknown, options: typeof toolOptions) => unknown }) .execute; if (!execute) throw new Error("tool has no execute function"); - return await execute(input, toolOptions); + const output = await execute(input, toolOptions); + if (output && typeof output === "object" && Symbol.asyncIterator in output) { + let last: unknown; + for await (const chunk of output as AsyncIterable) last = chunk; + return last; + } + return output; +} + +async function collectTool(tool: unknown, input: unknown): Promise { + const execute = (tool as { execute?: (input: unknown, options: typeof toolOptions) => unknown }) + .execute; + if (!execute) throw new Error("tool has no execute function"); + const output = await execute(input, toolOptions); + if (!output || typeof output !== "object" || !(Symbol.asyncIterator in output)) { + return [output]; + } + const chunks: unknown[] = []; + for await (const chunk of output as AsyncIterable) chunks.push(chunk); + return chunks; +} + +type ExecStreamEvent = + | { name: "stdout"; value: string } + | { name: "stderr"; value: string } + | { name: "exit"; code: number; result?: unknown }; + +function streamingHandle(events: ExecStreamEvent[]) { + return { + async *[Symbol.asyncIterator]() { + for (const event of events) yield event; + }, + result: async () => { + throw new Error("result() must not be called on a streamed handle"); + }, + }; +} + +// A clock that advances 200ms per read, past the 100ms coalescing +// floor, so every chunk produces its own running snapshot. Streaming +// tests that assert per-chunk output inject this to defeat coalescing. +function steppingClock(stepMs = 200): () => number { + let t = 0; + return () => { + t += stepMs; + return t; + }; } function toolDescription(tool: unknown): string { @@ -30,6 +76,50 @@ function makeWorkspace(): Workspace { return new Workspace({ storage: new SQLiteTestStorage(), now: () => 1_700_000_000_000 }); } +// An in-process command backend that streams a fixed event sequence. +// Registered on a real Workspace so the exec tool runs against the +// genuine WorkspaceRuntime handle rather than a hand-shaped fake: +// this pins the ExecWorkspaceLike binding and the assumption that the +// real handle is async-iterable. +function streamingCommandBackend(events: import("@cloudflare/computer-rpc").ExecEvent[]): { + id: string; + type: string; + connect(): Promise<{ + rpc: import("@cloudflare/computer-rpc").WorkspaceRPC; + sync: "none"; + close(): Promise; + }>; +} { + const shell: import("@cloudflare/computer-rpc").ShellRPC = { + async exec(input) { + const id = input.id ?? "cmd-1"; + return { + id, + events: new ReadableStream({ + start(controller) { + for (const event of events) controller.enqueue({ ...event, id }); + controller.close(); + }, + }), + }; + }, + getExec: () => Promise.reject(new Error("not used")), + killExec: () => Promise.resolve(), + disposeExec: () => Promise.resolve(), + }; + const noopSync = new Proxy( + {}, + { get: () => () => Promise.reject(new Error("sync: none")) }, + ) as import("@cloudflare/computer-rpc").SyncRPC; + return { + id: "shell", + type: "fake-command", + async connect() { + return { rpc: { sync: noopSync, shell }, sync: "none", close: async () => {} }; + }, + }; +} + function bytes(text: string): Uint8Array { return new TextEncoder().encode(text); } @@ -481,6 +571,476 @@ describe("createAITools exec tool", () => { }); }); +describe("createAITools callable exec", () => { + it("forwards env and input to the runtime and returns the result value", async () => { + const calls: Array<{ + command: string; + env: Record | undefined; + input: unknown; + backend: string | undefined; + }> = []; + const workspace = { + runtime: { + async exec( + command: string, + options: { + cwd?: string; + encoding: "utf8"; + backend?: string; + env?: Record; + input?: unknown; + }, + ) { + calls.push({ + command, + env: options.env, + input: options.input, + backend: options.backend, + }); + return { + result: async () => ({ + exitCode: 0, + stdout: "ran", + stderr: "", + value: { doubled: 84 }, + }), + }; + }, + isCallable: (id: string) => id === "js", + }, + }; + const tools = createAITools({ + workspace, + shell: { + defaultBackend: "js", + backends: { + js: { description: "JavaScript module runtime" }, + }, + }, + }); + + await expect( + executeTool(tools.exec, { + command: "export default (input) => ({ doubled: input.value * 2 })", + env: { API_KEY: "secret" }, + input: { value: 42 }, + }), + ).resolves.toEqual({ + command: "export default (input) => ({ doubled: input.value * 2 })", + cwd: null, + backend: "js", + exitCode: 0, + stdout: "ran", + stderr: "", + result: { doubled: 84 }, + }); + expect(calls).toEqual([ + { + command: "export default (input) => ({ doubled: input.value * 2 })", + env: { API_KEY: "secret" }, + input: { value: 42 }, + backend: "js", + }, + ]); + }); + + it("omits the result field when the backend returns no value", async () => { + const workspace = { + runtime: { + async exec() { + return { + result: async () => ({ exitCode: 0, stdout: "ok", stderr: "" }), + }; + }, + isCallable: (id: string) => id === "js", + }, + }; + const tools = createAITools({ + workspace, + shell: { + defaultBackend: "js", + backends: { js: { description: "JavaScript module runtime" } }, + }, + }); + + const output = (await executeTool(tools.exec, { command: "noop" })) as Record; + expect(output).not.toHaveProperty("result"); + expect(output).toMatchObject({ backend: "js", exitCode: 0, stdout: "ok" }); + }); + + it("errors quickly without calling the backend when input targets a non-callable backend", async () => { + let called = false; + const workspace = { + runtime: { + async exec() { + called = true; + return { result: async () => ({ exitCode: 0, stdout: "", stderr: "" }) }; + }, + isCallable: (id: string) => id === "js", + }, + }; + const tools = createAITools({ + workspace, + shell: { + defaultBackend: "shell", + backends: { + shell: { description: "fast shell" }, + js: { description: "JavaScript module runtime" }, + }, + }, + }); + + await expect( + executeTool(tools.exec, { command: "echo hi", input: { value: 1 }, backend: "shell" }), + ).resolves.toEqual({ + command: "echo hi", + cwd: null, + backend: "shell", + error: 'Backend "shell" is not callable; it does not accept structured input.', + }); + expect(called).toBe(false); + }); + + it("allows env on non-callable backends", async () => { + const calls: Array<{ env: Record | undefined; input: unknown }> = []; + const workspace = { + runtime: { + async exec(_command: string, options: { env?: Record; input?: unknown }) { + calls.push({ env: options.env, input: options.input }); + return { result: async () => ({ exitCode: 0, stdout: "", stderr: "" }) }; + }, + }, + }; + const tools = createAITools({ + workspace, + shell: { + defaultBackend: "shell", + backends: { shell: { description: "fast shell" } }, + }, + }); + + await expect( + executeTool(tools.exec, { command: "env", env: { FOO: "bar" } }), + ).resolves.toMatchObject({ backend: "shell", exitCode: 0 }); + expect(calls).toEqual([{ env: { FOO: "bar" }, input: undefined }]); + }); + + it("describes callable backends in the tool description", () => { + const workspace = { + runtime: { + async exec() { + throw new Error("not used"); + }, + isCallable: (id: string) => id === "js", + }, + }; + const tools = createAITools({ + workspace, + shell: { + defaultBackend: "js", + backends: { js: { description: "JavaScript module runtime" } }, + }, + }); + + expect(toolDescription(tools.exec)).toContain("callable"); + }); +}); + +describe("createAITools exec against a real Workspace", () => { + it("streams a real runtime handle end to end", async () => { + const workspace = new Workspace({ + storage: new SQLiteTestStorage(), + backends: [ + streamingCommandBackend([ + { id: "cmd-1", seq: 1, name: "stdout", value: new TextEncoder().encode("hello\n") }, + { id: "cmd-1", seq: 2, name: "exit", code: 0 }, + ]) as never, + ], + }); + const tools = createAITools({ + workspace, + shell: { + defaultBackend: "shell", + backends: { shell: { description: "fast shell" } }, + }, + }); + + const chunks = await collectTool(tools.exec, { command: "echo hello" }); + expect(chunks.at(-1)).toEqual({ + command: "echo hello", + cwd: null, + backend: "shell", + exitCode: 0, + stdout: "hello\n", + stderr: "", + }); + await workspace.close(); + }); +}); + +describe("createAITools exec streaming", () => { + it("streams stdout and stderr chunks and yields a final aggregate", async () => { + const workspace = { + runtime: { + async exec() { + return streamingHandle([ + { name: "stdout", value: "one\n" }, + { name: "stderr", value: "warn\n" }, + { name: "stdout", value: "two\n" }, + { name: "exit", code: 0 }, + ]); + }, + }, + }; + const tools = createAITools({ + workspace, + shell: { + defaultBackend: "shell", + backends: { shell: { description: "fast shell" } }, + now: steppingClock(), + }, + }); + + await expect(collectTool(tools.exec, { command: "run" })).resolves.toEqual([ + { command: "run", cwd: null, backend: "shell", exitCode: null, stdout: "one\n", stderr: "" }, + { + command: "run", + cwd: null, + backend: "shell", + exitCode: null, + stdout: "one\n", + stderr: "warn\n", + }, + { + command: "run", + cwd: null, + backend: "shell", + exitCode: null, + stdout: "one\ntwo\n", + stderr: "warn\n", + }, + { + command: "run", + cwd: null, + backend: "shell", + exitCode: 0, + stdout: "one\ntwo\n", + stderr: "warn\n", + }, + ]); + }); + + it("streams a callable backend's result folded onto the exit event", async () => { + const workspace = { + runtime: { + async exec() { + return streamingHandle([ + { name: "stdout", value: "working\n" }, + { name: "exit", code: 0, result: { ok: true } }, + ]); + }, + isCallable: (id: string) => id === "js", + }, + }; + const tools = createAITools({ + workspace, + shell: { + defaultBackend: "js", + backends: { js: { description: "JavaScript module runtime" } }, + }, + }); + + const chunks = await collectTool(tools.exec, { + command: "export default () => ({ ok: true })", + input: {}, + }); + expect(chunks).toHaveLength(2); + expect(chunks.at(-1)).toEqual({ + command: "export default () => ({ ok: true })", + cwd: null, + backend: "js", + exitCode: 0, + stdout: "working\n", + stderr: "", + result: { ok: true }, + }); + }); + + it("truncates streamed output on UTF-8 byte boundaries", async () => { + const workspace = { + runtime: { + async exec() { + return streamingHandle([ + { name: "stdout", value: "a\u{1f642}b" }, + { name: "exit", code: 0 }, + ]); + }, + }, + }; + const tools = createAITools({ + workspace, + shell: { + defaultBackend: "shell", + backends: { shell: { description: "fast shell" } }, + maxBytes: 5, + }, + }); + + const chunks = await collectTool(tools.exec, { command: "echo emoji" }); + expect(chunks.at(-1)).toMatchObject({ + exitCode: 0, + stdout: "a\u{1f642}\n\n[truncated, 1 more bytes]", + }); + }); + + it("yields a structured error when the stream fails mid-run", async () => { + const workspace = { + runtime: { + async exec() { + return { + async *[Symbol.asyncIterator]() { + yield { name: "stdout", value: "partial\n" } as ExecStreamEvent; + throw new Error("stream broke"); + }, + result: async () => { + throw new Error("result() must not be called on a streamed handle"); + }, + }; + }, + }, + }; + const tools = createAITools({ + workspace, + shell: { + defaultBackend: "shell", + backends: { shell: { description: "fast shell" } }, + }, + }); + + const chunks = await collectTool(tools.exec, { command: "run" }); + expect(chunks.at(-1)).toEqual({ + command: "run", + cwd: null, + backend: "shell", + error: "stream broke", + }); + }); + + it("coalesces running snapshots to at most one per interval", async () => { + const workspace = { + runtime: { + async exec() { + return streamingHandle([ + { name: "stdout", value: "a" }, + { name: "stdout", value: "b" }, + { name: "stdout", value: "c" }, + { name: "exit", code: 0 }, + ]); + }, + }, + }; + // A frozen clock never advances past the coalescing floor. The + // first chunk still yields a running snapshot for responsiveness; + // the later chunks coalesce into the terminal snapshot. + const tools = createAITools({ + workspace, + shell: { + defaultBackend: "shell", + backends: { shell: { description: "fast shell" } }, + now: () => 1000, + }, + }); + + const chunks = await collectTool(tools.exec, { command: "run" }); + expect(chunks).toEqual([ + { command: "run", cwd: null, backend: "shell", exitCode: null, stdout: "a", stderr: "" }, + { command: "run", cwd: null, backend: "shell", exitCode: 0, stdout: "abc", stderr: "" }, + ]); + }); + + it("caps streamed output in memory at streamMaxBytes", async () => { + const workspace = { + runtime: { + async exec() { + return streamingHandle([ + { name: "stdout", value: "a".repeat(10) }, + { name: "stdout", value: "b".repeat(10) }, + { name: "exit", code: 0 }, + ]); + }, + }, + }; + // Hold at most 12 bytes; show at most 8. The marker counts every + // byte seen (20), not just the 12 retained. + const tools = createAITools({ + workspace, + shell: { + defaultBackend: "shell", + backends: { shell: { description: "fast shell" } }, + maxBytes: 8, + streamMaxBytes: 12, + }, + }); + + const chunks = await collectTool(tools.exec, { command: "run" }); + expect(chunks.at(-1)).toMatchObject({ + exitCode: 0, + stdout: "aaaaaaaa\n\n[truncated, 12 more bytes]", + }); + }); + + it("kills the backend execution when the turn aborts", async () => { + let killed = 0; + const controller = new AbortController(); + const workspace = { + runtime: { + async exec() { + return { + async *[Symbol.asyncIterator]() { + yield { name: "stdout", value: "working\n" } as ExecStreamEvent; + controller.abort(); + // The abort listener calls kill(); yield once more so + // the iteration observes the signal before the stream + // ends on its own. + yield { name: "exit", code: 130 } as ExecStreamEvent; + }, + result: async () => { + throw new Error("result() must not be called on a streamed handle"); + }, + kill: async () => { + killed += 1; + }, + }; + }, + }, + }; + const tools = createAITools({ + workspace, + shell: { + defaultBackend: "shell", + backends: { shell: { description: "fast shell" } }, + }, + }); + + const execute = ( + tools.exec as { + execute: ( + input: unknown, + options: { toolCallId: string; messages: []; abortSignal: AbortSignal }, + ) => AsyncIterable; + } + ).execute; + const output = execute( + { command: "sleep" }, + { toolCallId: "t", messages: [], abortSignal: controller.signal }, + ); + for await (const _chunk of output) { + // Drain to completion. + } + expect(killed).toBe(1); + }); +}); + describe("createAITools publish tool", () => { it("adds publish by default when assets are configured", async () => { const calls: Array<{ path: string; expiresAfter: number; prefix?: string }> = []; diff --git a/packages/computer/src/tools/exec.ts b/packages/computer/src/tools/exec.ts index efbb93c..5a58fb3 100644 --- a/packages/computer/src/tools/exec.ts +++ b/packages/computer/src/tools/exec.ts @@ -1,18 +1,69 @@ import { type Tool, tool } from "ai"; import { z } from "zod"; +import { notCallableMessage } from "../runtime/runtime.js"; +import type { WorkspaceRuntimeValue } from "../runtime/types.js"; + +// A finite JSON value: what a callable backend accepts as `input` and +// returns as `result`. Declared as a concrete recursive schema rather +// than z.unknown() so the tool validates `input` at its own boundary +// and the generated JSON Schema describes a real shape (z.unknown() +// serializes to an empty schema that some providers reject under +// strict function calling). +const jsonValueSchema: z.ZodType = z.lazy(() => + z.union([ + z.string(), + z.number(), + z.boolean(), + z.null(), + z.array(jsonValueSchema), + z.record(z.string(), jsonValueSchema), + ]), +); + +// One event drained from a running execution. stdout / stderr carry +// output chunks as they arrive; exit carries the process exit code and, +// for a callable backend, the structured return value on `result`. The +// value settles at the same instant as the exit code, so it rides the +// same terminal event rather than a separate one. +export type ExecStreamEvent = + | { name: "stdout"; value: string } + | { name: "stderr"; value: string } + | { name: "exit"; code: number; result?: unknown }; + +// A detached execution handle. The tool streams stdout / stderr +// chunks by iterating the handle when it is async-iterable, and +// falls back to draining result() when it is not. +export interface ExecRuntimeHandle extends Partial> { + result(): Promise<{ + exitCode: number; + stdout: string; + stderr: string; + value?: unknown; + }>; + // Signal the running execution. The tool calls it when the model + // turn aborts, so the backend stops rather than running on after + // the tool stops iterating. + kill?(): Promise; +} + export interface ExecWorkspaceLike { runtime: { exec( command: string, - options: { cwd?: string; encoding: "utf8"; backend?: string }, - ): Promise<{ - result(): Promise<{ - exitCode: number; - stdout: string; - stderr: string; - }>; - }>; + options: { + cwd?: string; + encoding: "utf8"; + backend?: string; + env?: Record; + input?: WorkspaceRuntimeValue; + }, + ): Promise; + // Whether a backend accepts a structured `input` value and returns + // a structured result. The tool asks this to know which backends + // are callable; the runtime derives it from each backend's + // `callable` flag. Omit when no backend is callable. + isCallable?(id: string): boolean; }; } @@ -24,15 +75,54 @@ export interface ExecToolOptions { workspace: ExecWorkspaceLike; backends: Record; defaultBackend: string; + // Per-snapshot display cap for each of stdout and stderr, in bytes. + // Output past it is shown as a truncation marker. Defaults to 64 KiB. maxBytes?: number; + // In-memory cap per stream while streaming, in bytes. Output past it + // is counted toward the truncation marker but not retained, so a long + // run does not grow the buffer without bound. Defaults to 512 KiB. + streamMaxBytes?: number; + // Clock backing the running-snapshot coalescing floor. Defaults to + // Date.now; injectable so tests can drive the interval deterministically. + now?: () => number; } const DEFAULT_MAX_BYTES = 64 * 1024; +const DEFAULT_STREAM_MAX_BYTES = 512 * 1024; +// Minimum wall-clock gap between running snapshots. A chatty command +// yields at most one snapshot per interval instead of one per chunk; +// the terminal snapshot always fires regardless. +const STREAM_COALESCE_MS = 100; + +// Progressive snapshot emitted while a command streams (exitCode +// null until the run ends), and the terminal snapshot once the exit +// code lands. `result` appears only when a callable backend returned +// a value; `error` replaces the run fields when the exec fails. +export type ExecToolOutput = + | { + command: string; + cwd: string | null; + backend: string; + exitCode: number | null; + stdout: string; + stderr: string; + result?: unknown; + } + | { command: string; cwd: string | null; backend: string; error: string }; -export function createExecTool( - options: ExecToolOptions, -): Tool<{ command: string; cwd?: string; backend?: string }> { +export function createExecTool(options: ExecToolOptions): Tool< + { + command: string; + cwd?: string; + backend?: string; + env?: Record; + input?: WorkspaceRuntimeValue; + }, + ExecToolOutput +> { const maxBytes = options.maxBytes ?? DEFAULT_MAX_BYTES; + const streamMaxBytes = options.streamMaxBytes ?? DEFAULT_STREAM_MAX_BYTES; + const now = options.now ?? Date.now; const backendIds = Object.keys(options.backends); if (backendIds.length === 0) { throw new Error("createExecTool: pass at least one backend in `backends`"); @@ -43,9 +133,21 @@ export function createExecTool( ); } + const isCallable = options.workspace.runtime.isCallable?.bind(options.workspace.runtime); + const callableBackendIds = new Set(backendIds.filter((id) => isCallable?.(id) === true)); const backendGuidance = backendIds - .map((id) => `- ${JSON.stringify(id)}: ${options.backends[id].description}`) + .map((id) => { + const suffix = callableBackendIds.has(id) ? " (callable)" : ""; + return `- ${JSON.stringify(id)}${suffix}: ${options.backends[id].description}`; + }) .join("\n"); + const callableGuidance = + callableBackendIds.size > 0 + ? [ + "", + `Callable backends (${[...callableBackendIds].map((id) => JSON.stringify(id)).join(", ")}) run \`command\` as module source rather than a shell command. Pass \`input\` to hand the module a structured value, and read the module's returned value back from the \`result\` field. Other backends reject \`input\`.`, + ].join("\n") + : ""; const description = [ "Run a shell command in the workspace. The workspace exposes multiple backends, each with different capabilities.", "Pick the cheapest backend that can run the command; fall back to a heavier one only when the lighter backend's command set doesn't cover what you need.", @@ -55,6 +157,7 @@ export function createExecTool( "", `Default backend: ${JSON.stringify(options.defaultBackend)}. Try this first for any command you're not sure about; if it fails with a "command not found" or a similar capability error, retry on a backend whose description covers the missing tool.`, "Use for builds, test runs, typechecks, formatters, and git plumbing. Prefer the dedicated read, write, and edit tools for file operations. Long output is truncated to keep tool replies small.", + callableGuidance, ].join("\n"); const backendSchema = z @@ -71,41 +174,189 @@ export function createExecTool( return tool({ description, inputSchema: z.object({ - command: z.string().describe("Shell command, e.g. 'npm test' or 'git diff HEAD'."), + command: z + .string() + .describe( + "Shell command, e.g. 'npm test' or 'git diff HEAD'. For a callable backend this is the module source to run.", + ), cwd: z.string().optional().describe("Working directory. Defaults to the workspace root."), backend: backendSchema, + env: z + .record(z.string(), z.string()) + .optional() + .describe( + "Environment variables for this run only. Values override the backend's base environment without affecting later runs.", + ), + input: jsonValueSchema + .optional() + .describe( + "Structured value handed to a callable backend's module. Only callable backends accept it; other backends reject it.", + ), }), - execute: async ({ command, cwd, backend }) => { + execute: async function* ({ command, cwd, backend, env, input }, { abortSignal }) { const selectedBackend = backend ?? options.defaultBackend; + const base = { command, cwd: cwd ?? null, backend: selectedBackend }; + if (input !== undefined && !callableBackendIds.has(selectedBackend)) { + yield { ...base, error: notCallableMessage(selectedBackend) }; + return; + } + let handle: ExecRuntimeHandle; try { - const handle = await options.workspace.runtime.exec(command, { + handle = await options.workspace.runtime.exec(command, { cwd, encoding: "utf8", backend: selectedBackend, + env, + input, }); - const result = await handle.result(); - return { - command, - cwd: cwd ?? null, - backend: selectedBackend, - exitCode: result.exitCode, - stdout: truncate(result.stdout, maxBytes), - stderr: truncate(result.stderr, maxBytes), - }; } catch (err) { - return { - command, - cwd: cwd ?? null, - backend: selectedBackend, - error: err instanceof Error ? err.message : String(err), - }; + yield { ...base, error: errorMessage(err) }; + return; + } + + // Aborting the model turn kills the backend execution so it does + // not run on unobserved after the tool stops iterating. The run + // then emits its terminal event and the stream closes normally. + const onAbort = () => void handle.kill?.().catch(() => undefined); + if (abortSignal?.aborted) onAbort(); + else abortSignal?.addEventListener("abort", onAbort, { once: true }); + try { + yield* runExecution(); + } finally { + abortSignal?.removeEventListener("abort", onAbort); + } + + // Produce the run's snapshots. Streams the raw events when the + // handle is iterable; otherwise drains the aggregate result. + async function* runExecution(): AsyncGenerator { + // Stream stdout / stderr chunks as they arrive when the handle + // is iterable. Each chunk yields a fresh snapshot with the + // running output so the model sees progress before the run + // ends; the exit event settles the terminal snapshot. + if (typeof handle[Symbol.asyncIterator] === "function") { + const stdout = new StreamBuffer(streamMaxBytes); + const stderr = new StreamBuffer(streamMaxBytes); + let exitCode: number | null = null; + let value: unknown; + let hasValue = false; + // Coalesce running snapshots to at most one per interval. A + // chatty command would otherwise yield a full-buffer snapshot + // per chunk; the terminal snapshot below always fires. + let lastSnapshot = 0; + try { + for await (const event of handle as AsyncIterable) { + if (event.name === "stdout") stdout.push(event.value); + else if (event.name === "stderr") stderr.push(event.value); + else { + exitCode = event.code; + if ("result" in event) { + value = event.result; + hasValue = true; + } + continue; + } + const at = now(); + if (at - lastSnapshot < STREAM_COALESCE_MS) continue; + lastSnapshot = at; + yield { + ...base, + exitCode: null, + stdout: stdout.render(maxBytes), + stderr: stderr.render(maxBytes), + }; + } + } catch (err) { + yield { ...base, error: errorMessage(err) }; + return; + } + yield { + ...base, + exitCode, + stdout: stdout.render(maxBytes), + stderr: stderr.render(maxBytes), + ...(hasValue ? { result: value } : {}), + }; + return; + } + + // Non-streaming handle: drain the aggregate result. + try { + const result = await handle.result(); + yield { + ...base, + exitCode: result.exitCode, + stdout: truncate(result.stdout, maxBytes), + stderr: truncate(result.stderr, maxBytes), + ...(result.value === undefined ? {} : { result: result.value }), + }; + } catch (err) { + yield { ...base, error: errorMessage(err) }; + } } }, }); } +function errorMessage(err: unknown): string { + return err instanceof Error ? err.message : String(err); +} + const encoder = new TextEncoder(); +// A bounded, incrementally-counted accumulator for one output stream. +// It keeps at most `cap` bytes of head text in memory while tracking +// the total bytes seen, so a long run neither grows without bound nor +// re-encodes the whole buffer on every snapshot. `render` returns the +// display-capped view with a marker for the bytes not shown. +class StreamBuffer { + #head = ""; + #headBytes = 0; + #totalBytes = 0; + readonly #cap: number; + + constructor(cap: number) { + this.#cap = cap; + } + + push(chunk: string): void { + const chunkBytes = encoder.encode(chunk).byteLength; + this.#totalBytes += chunkBytes; + if (this.#headBytes >= this.#cap) return; + if (this.#headBytes + chunkBytes <= this.#cap) { + this.#head += chunk; + this.#headBytes += chunkBytes; + return; + } + // The chunk crosses the cap: keep the largest whole-character + // prefix that fits, then stop growing the head. + let used = this.#headBytes; + let end = 0; + for (const char of chunk) { + const charBytes = encoder.encode(char).byteLength; + if (used + charBytes > this.#cap) break; + used += charBytes; + end += char.length; + } + this.#head += chunk.slice(0, end); + this.#headBytes = used; + } + + render(maxBytes: number): string { + if (this.#totalBytes <= maxBytes && this.#totalBytes === this.#headBytes) { + return this.#head; + } + let used = 0; + let end = 0; + for (const char of this.#head) { + const charBytes = encoder.encode(char).byteLength; + if (used + charBytes > maxBytes) break; + used += charBytes; + end += char.length; + } + return `${this.#head.slice(0, end)}\n\n[truncated, ${this.#totalBytes - used} more bytes]`; + } +} + function truncate(value: string, maxBytes: number): string { if (!value) return value; const totalBytes = encoder.encode(value).byteLength; diff --git a/packages/computer/src/tools/index.ts b/packages/computer/src/tools/index.ts index 45dc297..b2686a5 100644 --- a/packages/computer/src/tools/index.ts +++ b/packages/computer/src/tools/index.ts @@ -1,5 +1,12 @@ export { type CreateAIToolsOptions, createAITools } from "./ai.js"; -export { createExecTool, type ExecBackendDescription, type ExecToolOptions } from "./exec.js"; +export { + createExecTool, + type ExecBackendDescription, + type ExecRuntimeHandle, + type ExecStreamEvent, + type ExecToolOptions, + type ExecToolOutput, +} from "./exec.js"; export { createEditTool, type EditToolOptions } from "./fs/edit.js"; export { createListTool, type ListToolOptions } from "./fs/list.js"; export { createReadTool, type ReadToolOptions } from "./fs/read.js";