sscoble/pi-mono
1
0
Fork 0
mirror of https://github.com/badlogic/pi-mono.git synced 2026-11-05 23:37:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

feat(durable): durable code mode (@earendil-works/pi-durable/codemode)

Browse source 
A codemode tool whose scripts call the conversation's tools as nested calls and
survive a restart. The script runs deterministically (pi-codemode start()); each
result is journaled before it is delivered. A rerun starts the script again from
the stored start data, delivers the journaled results in order (nested calls
reattach by key c<id>), checks every recorded call count, holds calls the
journal lacks until replay ends and checks them against earlier attempts' keys,
then continues live. A call launches only once the count of the drain that made
it is committed, so a rerun can prove which calls are new. Divergence, a changed
module version or sandbox fingerprint, and a passed deadline end the call with a
stored error; a decided outcome is returned without a VM. store()/load() keep
values per conversation, rewindable and forked as of the fork's entry.

pi-codemode: CodemodeRun.finished, set when the outcome is decided.
This commit is contained in:
Mario Zechner 2026-10-11 21:31:38 +00:00
parent 8a9c051f80
commit 7b597d1481
12 changed files with 1618 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

1
package-lock.json generated
View file

@ -5295,6 +5295,7 @@
"dependencies": {
"@earendil-works/chord": "^1.1.0",
"@earendil-works/pi-ai": "^1.1.0",
"@earendil-works/pi-codemode": "^1.1.0",
"@silvia-odwyer/photon": "0.3.3",
"diff": "8.0.4",
"typebox": "1.3.27"

2
packages/codemode/CHANGELOG.md
View file

@ -7,7 +7,7 @@
- Added `RemoteCodemodeSandbox`, which runs each script's VM in another isolate while tools run in the caller. `CodemodeSandboxDurableObject` from the new `@earendil-works/pi-codemode/cloudflare` entry is its remote end, deployed as its own Cloudflare Worker. `examples/cloudflare` shows the setup
- Added `InlineCodemodeSandbox`, which runs the VM on the calling thread for runtimes without worker threads. Its `interruptBudget` limits scripts that compute without awaiting a tool
- Added the `@earendil-works/pi-codemode/portable` entry, which has no Node imports, and `CodemodeSandboxBase`, the API all three sandboxes share
- Added `sandbox.start()`, a deterministic run whose calls the host answers itself (`onCall`, `run.deliver()`), so a host can replay a script after a restart: a virtual clock from `startedAt`, randomness from `seed`, UTC, no `WeakRef`/`FinalizationRegistry`, results applied in the order the host numbers them, the VM's fingerprint (`onStarted`), and the calls issued after each result (`onIdle`). Works with every sandbox; `execute()` is unchanged
- Added `sandbox.start()`, a deterministic run whose calls the host answers itself (`onCall`, `run.deliver()`), so a host can replay a script after a restart: a virtual clock from `startedAt`, randomness from `seed`, UTC, no `WeakRef`/`FinalizationRegistry`, results applied in the order the host numbers them, the VM's fingerprint (`onStarted`), the calls issued after each result (`onIdle`), and whether the run has its outcome (`run.finished`). Works with every sandbox; `execute()` is unchanged
### Changed

11
packages/codemode/src/runtime/execution.ts
View file

@ -200,7 +200,8 @@ class Execution {
private readonly output: CodemodeOutputItem[] = [];
private readonly calls: CodemodeCall[] = [];
private readonly pending = new Map<number, PendingCall>();
private finished = false;
/** Whether the execution has its outcome. */
finished = false;
constructor(options: ExecutionOptions) {
this.promise = new Promise<CodemodeResult>((resolve) => {
@ -555,7 +556,13 @@ export class CodemodeSandboxBase {
});
this.running.add(execution);
const result = execution.promise.finally(() => this.running.delete(execution));
return { result, deliver: (seq, now, id, outcome) => execution.deliver(seq, now, id, outcome) };
return {
result,
deliver: (seq, now, id, outcome) => execution.deliver(seq, now, id, outcome),
get finished() {
return execution.finished;
},
};
}
/** Aborts in-flight executions (they resolve with `kind: "aborted"`) and rejects new ones. */

5
packages/codemode/src/types.ts
View file

@ -214,6 +214,11 @@ export interface CodemodeRunOptions {
/** A running {@link CodemodeSandboxBase.start} execution. */
export interface CodemodeRun {
readonly result: Promise<CodemodeResult>;
/**
* Whether the run has its outcome: the script ended, timed out, or failed, or the run was aborted. Set when that is
* decided, before `result` settles (which waits for the VM to stop); calls can no longer be answered.
*/
readonly finished: boolean;
/**
* Answer call `id` as result number `seq`. Results are numbered 1, 2, ... without gaps, and the VM applies them
* strictly in that order, whatever order they arrive in: after a gap it waits for the missing number. While it

10
packages/codemode/test/deterministic.test.ts
View file

@ -340,6 +340,16 @@ for (const [backend, create] of Object.entries(BACKENDS)) {
expect(() => recorded.run.deliver(1, STARTED_AT, 1, { ok: true, value: null })).not.toThrow();
});
it("is finished once it has its outcome", async () => {
const sandbox = create();
const recorded = start(sandbox, "await tools.wait(1); return 2");
await waitFor(() => recorded.calls.length === 1);
expect(recorded.run.finished).toBe(false);
recorded.run.deliver(1, STARTED_AT, 1, { ok: true, value: null });
await recorded.run.result;
expect(recorded.run.finished).toBe(true);
});
it("rejects a call that is not waiting and a result number used twice", async () => {
const sandbox = create();
const recorded = start(sandbox, "await Promise.all([tools.wait(1), tools.wait(2)])");

1
packages/durable/CHANGELOG.md
View file

@ -21,6 +21,7 @@
- `TaskOptions.abandonOnRestart`: a task its creator awaits only in memory is aborted, with its owned work, when a later Harness starts scheduling, before any of it runs again. Its mark carries `abortReason: "restart"`, which cascades; such a task waits for a missing definition instead of becoming `orphaned`. Nested calls and `api.createTask()` children of tools that are not replay-safe get it by default. No task below an owner with cancellation intent starts a run phase any more, even before the cascade marks it.
- `TaskRuntime.abortOwned(id, context)`: abort a task this task owns and wait until it is terminal.
- `TaskRuntime.ownedTasks(context)`: the live tasks a task owns directly, whichever invocation created them, as copies. A tool call settling finds the nested calls it left running with it.
- `@earendil-works/pi-durable/codemode`: a `codemode` tool (`createCodemodeExtension({ sandbox })`) whose scripts call the conversation's tools as nested calls and survive a restart: the script is replayed deterministically from a journal of the results it got, written before each delivery, reattaching to its nested calls, and then continues. Divergence, a changed sandbox fingerprint or module version, and a passed deadline end the call with an error. Scripts keep `store()`/`load()` values per conversation, forked as of the fork's entry.
- `read` returns images as one image block instead of an `unsupported_image` error: PNG, JPEG, GIF, and WebP within the model's image limits (by default 2000x2000 pixels and 4.5 MB of base64) as they are, or, with `createCodingTools({ images })`, oriented, converted, and shrunk to fit. `@earendil-works/pi-durable/images` provides a Photon (WebAssembly) image processor, with `/images/node` and `/images/cloudflare` loaders; `/images/node` prepares images in a worker thread, so a large photo never blocks the event loop, and the thread exits after 30 s without images; nothing loads it unless imported. A model without image input gets a diagnostic saying it sees a placeholder.
- `pi.live.compactionFailure` keeps the newest compaction failure (task, reason, error) until a compaction produces a summary; the agent events' snapshot carries it as `compactionFailure`, and `compaction_end` carries `error` for a compaction that failed. A failed background or manual compaction was invisible once its status was removed.

16
packages/durable/README.md
View file

@ -14,6 +14,7 @@ Built on [`@earendil-works/pi-ai`](../ai/README.md) for model access and `@earen
- [Persist and Resume](#persist-and-resume)
- [Extensions](#extensions)
- [Tools](#tools)
- [Code Mode](#code-mode)
- [System Prompt](#system-prompt)
- [Per-Conversation Agent](#per-conversation-agent)
- [Settings](#settings)
@ -253,6 +254,21 @@ await root.configure({ modelTools: { remove: [bash] } }, context);
A task that its creator awaits only in memory, and that the creator never resumes after a restart, can be created with `abandonOnRestart: true`. When a later Harness starts scheduling and finds it still live, it is aborted with `abortReason: "restart"`, together with everything it owns, before any of it runs again. Nested calls of tools that are not replay-safe get this automatically, and so do the child tasks such a tool creates with `api.createTask()`. A task abandoned this way whose definition is missing waits until its extension is installed, so its abort handler still cleans up. Its caller, and the conversation, wait with it; `harness.inspect()` shows it as blocked. Aborting it then orphans it, as any abort does.
## Code Mode
`@earendil-works/pi-durable/codemode` adds a `codemode` tool: the model writes JavaScript that calls the conversation's tools as `await tools.<name>(args)`, in a QuickJS sandbox (`@earendil-works/pi-codemode`). Each call is a nested tool call, and only the script's output reaches the model. The script survives a restart: it resumes where it was, without running its finished calls again.
```typescript
import { CodemodeSandbox } from "@earendil-works/pi-codemode";
import { createCodemodeExtension } from "@earendil-works/pi-durable/codemode";
registry.install(createCodemodeExtension({ sandbox: new CodemodeSandbox({ memoryLimitBytes: 256 * 1024 * 1024 }) }));
```
On Cloudflare, pass a `RemoteCodemodeSandbox` instead. A failed nested call rejects in the script with its error text; its `structuredOutput`, if any, is on the error as `error.structuredOutput`. `store(key, value)` and `load(key)` keep JSON values per conversation; a fork sees them as they were at its entry.
How it survives a restart: the sandbox runs the script deterministically (a virtual clock, seeded randomness, UTC), and the tool writes each result the script gets to a journal before handing it over. After a restart, the tool runs the script again and hands it the journaled results in the same order, so it makes the same calls with the same arguments; these reattach to their nested calls by key, and then the script continues. A rerun ends the call with an error instead when the script takes another path than before, when the sandbox's fingerprint or the code mode version changed, or when its `timeout_ms` passed meanwhile. A nested call running at the restart follows its own replay policy.
## System Prompt
The system prompt is built from the selected extensions' sections, rendered in order before each request. A section sees the resolved agent, the environment built for the request, and committed documents:

61
packages/durable/docs/spec.md
View file

@ -3551,6 +3551,67 @@ Harness instances must never own the same Session concurrently: open the new one
after the old one's `close()` resolved, when none of its invocation code runs
any more (section 2.2).
### 7.6 Code mode
`@earendil-works/pi-durable/codemode` provides the `codemode` tool (`createCodemodeExtension({ sandbox })`):
`callers: ["model"]`, `replay: "safe"`. Its argument is a script, run with
pi-codemode's deterministic `sandbox.start()`. The script calls the agent's
`callable` tools as `tools.<name>(args)`; each call is a nested call with key
`c<id>`, the VM's own number for the call. A success resolves with the nested
result's `structuredOutput`; every error rejects, with that `structuredOutput`
as message when it is a non-empty string (the joined texts of a content list),
else the error diagnostics' messages, and the value as `error.structuredOutput`.
Three documents, two of them scoped to the code mode call's task:
- `pi.codemode.run`: `start`, written once before the VM starts (module
version, `startedAt`, seed, the VM's fingerprint once reported, the deadline
from `timeout_ms`, the tools and globals the script sees, the store
snapshot), and `result`, the outcome, written once.
- `pi.codemode.journal`: `entries`, one per result delivered to the VM, in
delivery order: the call's ID, the clock value, `issuedAfter` (the calls the VM
had issued after processing it, from its acknowledgment), and for a global its
name, arguments, and outcome; and `firstIssued`, the calls the VM had issued
after the script's first run.
- `pi.codemode.store`: conversation-scoped, `history: "rewindable"`,
`fork: "asOf"`; a successful script's `store()` writes are applied per key in
the transaction that writes its `result`.
Live, a finished call's result is appended to the journal and then delivered as
result number `entry index + 1`, with the clock at the entry's time; one journal
commit is in flight at a time and later results join the next. Acknowledged
counts are recorded with the next journal commit, or at once when the VM has
processed everything delivered. A call is launched only once the count of the
drain that made it is committed, so every key an attempt admitted is at most its
last recorded count.
A rerun (`execute()` after recovery) returns a stored `result` without a VM.
Otherwise it ends the call with a stored error result, admitting no new nested
call, when the module version differs or the deadline passed; it starts the VM
from `start` and does the same when the fingerprint the VM reports differs, before
delivering or reattaching anything. It then delivers the journal in order, each
entry once the VM has issued its call; a tool's result comes from its nested
call by key, a global's from the journal after comparing name and arguments.
Calls the journal does not have are held until the last journaled result is
acknowledged. Divergence ends the call with a stored error result, then stops the
VM with an error the script cannot catch: a count that differs from a recorded
`issuedAfter`, a VM that waits without having issued the next entry's call, a
key or global called with another name or arguments, or a script that ends while
processing an entry whose `issuedAfter` was recorded (or during its first run
with a non-empty journal), or before issuing a call that a later entry answers.
A live call whose key an earlier attempt admitted with another tool or other
arguments diverges the same way. Every recorded count, `firstIssued`
included, must equal the replay's. On a rerun, live calls are held until the
acknowledgment of the last journaled result, or of the first run when the
journal is empty; by then the counts have matched, so the script has made every
call an earlier attempt could have admitted. Every held call's key is checked
against its stored nested call, whatever the call's target, and only when none
differs are they launched; later calls are new. Nothing is admitted or delivered once the run
has its outcome. A script that ends before using every entry
is otherwise fine: the live run ended there too. A count the replay acknowledges
but the journal lacks is recorded with the first live commit. Nested calls admitted before the restart may
resume on their own whatever the rerun decides; the call's cleanup aborts them.
## 8. Built-in tasks
The initial implementation provides:

7
packages/durable/package.json
View file

@ -24,6 +24,12 @@
"import": "./dist/env/node.js",
"default": "./dist/env/node.js"
},
"./codemode": {
"source": "./src/codemode/index.ts",
"types": "./dist/codemode/index.d.ts",
"import": "./dist/codemode/index.js",
"default": "./dist/codemode/index.js"
},
"./tools": {
"source": "./src/tools/index.ts",
"types": "./dist/tools/index.d.ts",
@ -132,6 +138,7 @@
"dependencies": {
"@earendil-works/chord": "^1.1.0",
"@earendil-works/pi-ai": "^1.1.0",
"@earendil-works/pi-codemode": "^1.1.0",
"@silvia-odwyer/photon": "0.3.3",
"diff": "8.0.4",
"typebox": "1.3.27"

722
packages/durable/src/codemode/codemode.ts Normal file
View file

@ -0,0 +1,722 @@
/**
* Durable code mode: a tool that runs model-written JavaScript in a QuickJS sandbox, whose script survives a restart.
*
* Problem: a script is in-memory VM state. A restart while
*
* const results = await Promise.all(paths.map((path) => tools.read({ path })));
*
* waits for the last of 200 reads would lose the script and every read it made.
*
* Solution: deterministic replay (`CodemodeSandboxBase.start()` in pi-codemode). The VM only runs when the host hands
* it a result; given the same start data, clock, seed, and results in the same order, it makes the same calls with
* the same arguments. The tool is replay-safe: after a restart, `execute()` runs again, starts the script again from
* the stored start data, and hands it the results it got before, in the order it got them, from a journal written
* before each delivery. Nested tool calls reattach to their tasks by key (`c<VM call id>`), so they never run twice
* unless their own replay policy reruns them; their results come from `NestedCallDoc`. Then the script continues live.
*
* Documents, scoped to the code mode call's task and gone when it settles:
* - `pi.codemode.run`: the start data, written once before the VM starts, and the outcome, written once when the call
* decides it (with the store writes, in one transaction), so a rerun after that returns it without a VM.
* - `pi.codemode.journal`: one entry per delivered result, in delivery order: the VM call ID, the clock value, the
* calls the VM had issued after processing it (once acknowledged), and a global's outcome, which nothing else stores.
*
* Each result's acknowledged call count is recorded with the next journal commit, or at once when the VM has processed
* everything delivered: a script waiting on a long call is where a restart usually finds it. A sequential script
* therefore costs one journal commit and one count commit per result, on top of its nested calls' own; that is the
* price of checking the most common crash point.
*
* A call is launched only once the count of calls issued in the drain that made it is committed. So every key an earlier
* attempt admitted is at most its last recorded count; a replay checks every recorded count, so when it releases its
* calls it has made every call an earlier attempt could have admitted, and later calls are new.
*
* A rerun refuses to replay (no new nested calls; nested calls admitted earlier may still resume on their own and are
* aborted by the call's cleanup) when this module's version or the VM's fingerprint changed, when the deadline passed,
* or when the script diverges: it makes another call, or calls in another order, than the journal says. Each ends the
* call with an error result.
*/
import { type Context, copyJson, type JsonValue } from "@earendil-works/chord";
import type { ImageContent, TextContent } from "@earendil-works/pi-ai";
import {
type CodemodeHostCall,
type CodemodeOutcome,
type CodemodeResult,
type CodemodeRun,
type CodemodeSandboxBase,
toCodemodeIdentifier,
} from "@earendil-works/pi-codemode/portable";
import { parseCodemodeSource } from "@earendil-works/pi-codemode/source";
import { type Static, Type } from "typebox";
import { defineDoc } from "../documents.ts";
import { errorMessage } from "../errors.ts";
import { defineExtension, defineTool } from "../harness/define.ts";
import { NestedCallDoc, type ToolTaskInput } from "../harness/tool.ts";
import type {
Extension,
NestedToolExecutionResult,
ToolExecutionApi,
ToolExecutionResult,
ToolRegistration,
} from "../harness/types.ts";
import type { JsonObject } from "../types.ts";
/**
* Version of everything in this module that turns stored results into what a script sees: result mapping, globals,
* start data. Bump it when that changes, so a call started before an upgrade is not replayed by different code.
*/
export const CODEMODE_MODULE_VERSION = 1;
export const CODEMODE_TOOL_NAME = "codemode";
/** One tool as the script sees it: `tools.<jsName>`. */
type DeclaredTool = { name: string; jsName: string; description: string };
/** What a call starts with; a rerun starts the VM from exactly this. */
type Start = {
version: number;
startedAt: number;
seed: number;
/** The VM's fingerprint, recorded by the first run before any call is admitted. */
fingerprint?: string;
/** `startedAt + timeout_ms` when the script set one. */
deadline?: number;
tools: DeclaredTool[];
globals: { name: string; spread: boolean }[];
store: Record<string, JsonValue>;
};
type RunState = {
start?: Start;
/** The outcome, decided once. */
result?: JsonObject;
};
export const CodemodeRunDoc = defineDoc<RunState>({
kind: "pi.codemode.run",
version: 1,
scope: "task",
initial: () => ({}),
});
type JournalEntry = {
/** The VM's call ID. Its result is the entry's position + 1 (`seq`). */
id: number;
/** The VM clock while it processes this result. */
at: number;
/** Calls the VM had issued after processing this result; absent until acknowledged. */
issuedAfter?: number;
/** A global's call and outcome; a tool's result is its nested call's. */
global?: { name: string; args?: JsonValue; outcome: JsonObject };
};
type JournalState = {
entries: JournalEntry[];
/** Calls the VM had issued after the script's first run, before any result; absent until acknowledged. */
firstIssued?: number;
};
export const CodemodeJournalDoc = defineDoc<JournalState>({
kind: "pi.codemode.journal",
version: 1,
scope: "task",
initial: () => ({ entries: [] }),
});
/**
* What scripts keep with `store()` and read with `load()`, per conversation. Rewindable and copied as of the fork
* point, so a fork sees the values as they were when its entry was committed.
*/
export const CodemodeStoreDoc = defineDoc<{ values: Record<string, JsonValue> }>({
kind: "pi.codemode.store",
version: 1,
scope: "conversation",
history: "rewindable",
fork: "asOf",
initial: () => ({ values: {} }),
});
/** A function scripts call as a top-level global; its result is journaled. */
export type CodemodeGlobal = {
readonly name: string;
readonly spread?: boolean;
execute(args: unknown, api: ToolExecutionApi, context: Context): Promise<JsonValue | undefined>;
};
const codemodeSchema = Type.Object({ code: Type.String({ description: "Raw JavaScript source." }) });
export type CodemodeToolInput = Static<typeof codemodeSchema>;
export type CodemodeToolOptions = {
/** Runs the scripts: `CodemodeSandbox` on Node, `RemoteCodemodeSandbox` on Cloudflare. */
readonly sandbox: CodemodeSandboxBase;
/** @internal Host functions scripts call as globals. */
readonly globals?: readonly CodemodeGlobal[];
};
/** The code mode tool. Scripts call the conversation's callable tools as `tools.<name>(args)`. */
export function createCodemodeTool(options: CodemodeToolOptions): ToolRegistration<typeof codemodeSchema> {
return defineTool({
name: CODEMODE_TOOL_NAME,
description:
"Run JavaScript that calls other tools. The input is raw JavaScript, run as an async function body in a QuickJS sandbox: top-level `await` and `return` work. No Node, file system, network, or timers. `await tools.<name>({ ...args })` resolves to the tool's result and rejects with an Error on failure. `text(value)` and `console.log()` add output.",
parameters: codemodeSchema,
callers: ["model"],
replay: "safe",
execute: (args, api, context) => new CodemodeCall(options, api, context).execute(args.code),
});
}
/** The `codemode` extension: the code mode tool. */
export function createCodemodeExtension(options: CodemodeToolOptions): Extension {
return defineExtension({ name: "codemode", tools: [createCodemodeTool(options)] });
}
/** What a script gets for a nested call: success resolves with its value; every error rejects (design §4.5). */
function outcomeOf(name: string, result: NestedToolExecutionResult): CodemodeOutcome {
const value = result.structuredOutput;
if (!result.isError) return { ok: true, value };
const message =
typeof value === "string" && value !== ""
? value
: Array.isArray(value)
? value.flatMap((item) => (isText(item) ? [item.text] : [])).join("\n")
: "";
const diagnostics = result.diagnostics
.filter((diagnostic) => diagnostic.severity === "error")
.map((diagnostic) => diagnostic.message)
.join("\n");
return {
ok: false,
message: message || diagnostics || `Tool ${name} failed`,
...(value === undefined ? {} : { properties: { structuredOutput: value } }),
};
}
function isText(item: unknown): item is TextContent {
return typeof item === "object" && item !== null && (item as TextContent).type === "text";
}
/** The tool result for a script's result. */
function toolResult(result: CodemodeResult): ToolExecutionResult {
const output: (TextContent | ImageContent)[] = [];
for (const item of result.output) output.push(item.type === "image" ? item : { type: "text", text: item.text });
if (result.ok && result.value !== undefined) {
output.push({
type: "text",
text: typeof result.value === "string" ? result.value : JSON.stringify(result.value),
});
}
if (!result.ok) {
const { error } = result;
output.push({ type: "text", text: `Script error (${error.kind}): ${error.stack ?? error.message}` });
}
return { output, isError: !result.ok };
}
function errorResult(message: string): ToolExecutionResult {
return { output: [{ type: "text", text: message }], isError: true };
}
/** One execution of the code mode tool: a first run, or a rerun after a restart that replays the journal. */
class CodemodeCall {
private readonly options: CodemodeToolOptions;
private readonly api: ToolExecutionApi;
private readonly context: Context;
private readonly controller = new AbortController();
private run!: CodemodeRun;
private start!: Start;
/** The journal as this attempt found it; replay delivers exactly these. */
private replayed: JournalEntry[] = [];
/** Entries in the journal, replayed and live; the next delivery is `length + 1`. */
private journalLength = 0;
/** Replayed entries delivered so far. */
private nextReplayed = 0;
private replaying = false;
/** Whether an earlier attempt ran this call: its nested call keys may already exist. */
private rerun = false;
/** Whether live calls may launch; before, they are held (see `release()`). */
private released = false;
/** The count an earlier attempt recorded after the script's first run. */
private firstIssued: number | undefined;
/** The highest issued count known to be committed: calls up to it may launch. */
private durableIssued = 0;
/** Live calls waiting for their issued count to be committed. */
private waiting: CodemodeHostCall[] = [];
/** Journal index of each replayed call ID. */
private replayIndex = new Map<number, number>();
/** Calls the VM issued so far; its call IDs are 1, 2, ... */
private issued = 0;
/** Replay: results of journaled calls, ready to deliver in journal order. */
private readonly ready = new Map<number, CodemodeOutcome>();
/** Replay: calls the journal does not have, admitted once replay is done. */
private held: CodemodeHostCall[] = [];
/** Live: finished calls waiting for their journal commit. */
private queue: { id: number; at: number; outcome: CodemodeOutcome; global?: JournalEntry["global"] }[] = [];
/** Acknowledged counts to record, by result number (0: the script's first run). */
private readonly acks = new Map<number, number>();
private committing = false;
/** Set once the call decided an outcome other than the script's own (divergence, mismatch) or failed. */
private ended: ToolExecutionResult | undefined;
private failure: unknown;
/** Resolves once the VM's fingerprint is checked or recorded; nothing is delivered or admitted before. */
private gate!: Promise<void>;
private openGate!: () => void;
constructor(options: CodemodeToolOptions, api: ToolExecutionApi, context: Context) {
this.options = options;
this.api = api;
this.context = context;
this.gate = new Promise((resolve) => {
this.openGate = resolve;
});
}
async execute(source: string): Promise<ToolExecutionResult> {
const { code, options: sourceOptions } = parseCodemodeSource(source);
const { api, context } = this;
const agent = await api.agent(context);
const tools: DeclaredTool[] = agent.callable.map((tool) => ({
name: tool.name,
jsName: toCodemodeIdentifier(tool.name),
description: tool.description,
}));
const globals = (this.options.globals ?? []).map((global) => ({
name: global.name,
spread: global.spread === true,
}));
const now = Date.now();
const seed = crypto.getRandomValues(new Uint32Array(1))[0]!;
type Found =
| { done: ToolExecutionResult }
| { start: Start; journal: JournalEntry[]; firstIssued: number | undefined; rerun: boolean };
const found = await api.commit(async (tx): Promise<Found> => {
const run = await tx.doc(CodemodeRunDoc, api.taskId);
if (run.result !== undefined) return { done: copyJson(run.result as JsonValue) as ToolExecutionResult };
const rerun = run.start !== undefined;
if (run.start === undefined) {
const store = await tx.doc(CodemodeStoreDoc, api.conversationId);
const timeout = sourceOptions.timeoutMs;
run.start = {
version: CODEMODE_MODULE_VERSION,
startedAt: now,
seed,
...(timeout === undefined ? {} : { deadline: now + timeout }),
tools,
globals,
store: copyJson(store.values as JsonValue) as Record<string, JsonValue>,
};
}
const journal = await tx.doc(CodemodeJournalDoc, api.taskId);
return {
start: copyJson(run.start as unknown as JsonValue) as unknown as Start,
journal: copyJson(journal.entries as unknown as JsonValue) as unknown as JournalEntry[],
firstIssued: journal.firstIssued,
rerun,
};
}, context);
if ("done" in found) return found.done;
this.start = found.start;
this.replayed = found.journal;
this.replayIndex = new Map(found.journal.map((entry, index) => [entry.id, index]));
this.journalLength = found.journal.length;
this.replaying = found.journal.length > 0;
this.rerun = found.rerun;
this.firstIssued = found.firstIssued;
// A loop, not spread arguments: a journal can hold more entries than a call takes arguments.
this.durableIssued = found.journal.reduce(
(max, entry) => Math.max(max, entry.issuedAfter ?? 0),
found.firstIssued ?? 0,
);
if (this.start.version !== CODEMODE_MODULE_VERSION) {
return this.decide(
errorResult(
`The script cannot resume: it started under code mode version ${this.start.version}, this is ${CODEMODE_MODULE_VERSION}.`,
),
);
}
const remaining = this.start.deadline === undefined ? undefined : this.start.deadline - Date.now();
if (remaining !== undefined && remaining <= 0) {
return this.decide(errorResult(`Script timed out: its deadline passed while it was interrupted.`));
}
const onAbort = () => this.controller.abort(context.abortSignal?.reason);
context.abortSignal?.addEventListener("abort", onAbort, { once: true });
if (context.abortSignal?.aborted) onAbort();
try {
this.run = this.options.sandbox.start(code, {
startedAt: this.start.startedAt,
seed: this.start.seed,
tools: this.start.tools,
globals: this.start.globals,
store: this.start.store,
signal: this.controller.signal,
...(remaining === undefined ? {} : { timeoutMs: remaining }),
onStarted: (fingerprint) => void this.checkFingerprint(fingerprint),
onCall: (call) => this.onCall(call),
onIdle: (seq, issued) => this.onIdle(seq, issued),
});
const result = await this.run.result;
if (this.failure !== undefined) throw this.failure;
if (this.ended !== undefined) return this.ended;
context.abortSignal?.throwIfAborted();
// A replay that ends where the live run went on is a divergence (design §4.3).
const during = result.during;
// Within the replayed part, the first run included: a rerun without journal entries replays that much too.
if (this.rerun && during !== undefined && during.seq <= this.replayed.length) {
// The live run went on: it acknowledged that result, or it made a call that later results answer, which the
// replay never made.
const continued =
(during.seq === 0
? this.replayed.length > 0 || this.firstIssued !== undefined
: this.replayed[during.seq - 1]?.issuedAfter !== undefined) ||
this.replayed.slice(during.seq).some((entry) => entry.id > during.issued);
if (continued) {
return this.decide(
errorResult(
`The script took a different path when resumed after a restart: it ended after result ${during.seq}, where it went on before.`,
),
);
}
}
return this.complete(result);
} finally {
context.abortSignal?.removeEventListener("abort", onAbort);
}
}
/** Store an outcome the call decided, and end the VM if it runs; returns the outcome. */
private async decide(result: ToolExecutionResult): Promise<ToolExecutionResult> {
await this.api.commit(async (tx) => {
const run = await tx.doc(CodemodeRunDoc, this.api.taskId);
run.result ??= copyJson(result as JsonValue, { omitUndefinedProperties: true }) as JsonObject;
}, this.context);
return result;
}
/** End the call with `message`: store it as the outcome, then stop the VM with an error the script cannot catch. */
private diverge(message: string): void {
if (this.ended !== undefined || this.failure !== undefined) return;
const result = errorResult(`The script took a different path when resumed after a restart: ${message}`);
this.ended = result;
this.decide(result).then(
() => this.controller.abort(new Error(message)),
(error: unknown) => this.fail(error),
);
}
/** Whether nothing more may happen: the run has its outcome, or the call decided or failed. */
private get stopped(): boolean {
return this.ended !== undefined || this.failure !== undefined || this.run?.finished === true;
}
private fail(error: unknown): void {
if (this.failure !== undefined) return;
this.failure = error;
this.controller.abort(error);
}
/** The script's own outcome, stored with its store writes in one transaction. */
private async complete(result: CodemodeResult): Promise<ToolExecutionResult> {
const tool = toolResult(result);
const writes = result.ok ? result.storeWrites : undefined;
return this.api.commit(async (tx) => {
const run = await tx.doc(CodemodeRunDoc, this.api.taskId);
if (run.result !== undefined) return run.result as ToolExecutionResult;
run.result = copyJson(tool as JsonValue, { omitUndefinedProperties: true }) as JsonObject;
if (writes !== undefined && (Object.keys(writes.set).length > 0 || writes.delete.length > 0)) {
const store = await tx.doc(CodemodeStoreDoc, this.api.conversationId);
for (const key of writes.delete) delete store.values[key];
for (const [key, value] of Object.entries(writes.set)) store.values[key] = value as JsonValue;
}
return tool;
}, this.context);
}
private async checkFingerprint(fingerprint: string): Promise<void> {
const recorded = this.start.fingerprint;
if (recorded !== undefined) {
if (recorded !== fingerprint) {
this.ended = errorResult(
`The script cannot resume: the sandbox changed since it started (${recorded} -> ${fingerprint}).`,
);
await this.decide(this.ended).catch((error: unknown) => this.fail(error));
this.controller.abort(new Error("Sandbox fingerprint changed"));
return;
}
} else {
try {
await this.api.commit(async (tx) => {
const run = await tx.doc(CodemodeRunDoc, this.api.taskId);
if (run.start !== undefined) run.start.fingerprint = fingerprint;
}, this.context);
} catch (error) {
this.fail(error);
return;
}
}
this.openGate();
// Nothing was delivered before the gate: start now. A rerun without a journal releases its calls at the
// script's first acknowledgment instead.
if (this.replaying) this.deliverReplayed();
else if (!this.rerun) void this.release();
}
private onCall(call: CodemodeHostCall): void {
if (this.stopped) return;
this.issued = Math.max(this.issued, call.id);
const index = this.replaying ? (this.replayIndex.get(call.id) ?? -1) : -1;
if (index === -1) {
if (this.released) {
this.waiting.push(call);
this.launchCovered();
} else this.held.push(call);
return;
}
const entry = this.replayed[index]!;
if (entry.global !== undefined) {
if (call.target !== "global" || call.name !== entry.global.name || !sameJson(call.args, entry.global.args)) {
this.diverge(`call ${call.id} is ${describe(call)}, but was a call to ${entry.global.name}`);
return;
}
this.ready.set(call.id, entry.global.outcome as CodemodeOutcome);
void this.gate.then(() => this.deliverReplayed());
return;
}
if (call.target !== "tool") {
this.diverge(`call ${call.id} is ${describe(call)}, but was a tool call`);
return;
}
void this.gate
.then(() => this.callTool(call))
.then(
(outcome) => {
this.ready.set(call.id, outcome);
this.deliverReplayed();
},
(error: unknown) => {
if (this.context.abortSignal?.aborted) this.fail(error);
// The key names another tool or other arguments.
else this.diverge(`call ${call.id} is ${describe(call)}: ${errorMessage(error)}`);
},
);
}
/** Replay: deliver journaled results in journal order, each once the VM has issued its call. */
private deliverReplayed(): void {
while (this.nextReplayed < this.replayed.length && !this.stopped) {
const entry = this.replayed[this.nextReplayed]!;
const outcome = this.ready.get(entry.id);
if (entry.id > this.issued || outcome === undefined) return;
this.ready.delete(entry.id);
this.nextReplayed++;
this.run.deliver(this.nextReplayed, entry.at, entry.id, outcome);
}
}
private onIdle(seq: number, issued: number): void {
if (this.stopped) return;
// A count an earlier attempt recorded must match; a missing one is recorded now.
if (this.rerun && seq <= this.replayed.length) {
const recorded = seq === 0 ? this.firstIssued : this.replayed[seq - 1]!.issuedAfter;
if (recorded !== undefined && recorded !== issued) {
this.diverge(
seq === 0
? `its first run made ${issued} calls, where it had made ${recorded}`
: `after result ${seq} it had made ${issued} calls, where it had made ${recorded}`,
);
return;
}
}
this.acks.set(seq, issued);
if (!this.replaying) {
// A rerun without a journal: every call an earlier attempt could have admitted is made by now.
if (seq === 0 && this.rerun) void this.release();
// Committed at once when the script has processed everything delivered: calls it made wait for it, and a script
// waiting on a long call is where a restart usually finds it. Otherwise with the next commit; the VM catches up
// without waiting on anything.
if (seq === this.journalLength) void this.pump();
return;
}
// Caught up: the VM processed everything delivered and waits.
if (seq !== this.nextReplayed) return;
const next = this.replayed[seq];
if (next !== undefined) {
if (next.id > issued) this.diverge(`it waits without making call ${next.id}, which it made before`);
return;
}
// The last replayed result is acknowledged: continue live, admitting the calls the journal did not have once their
// count is committed.
this.replaying = false;
void this.pump();
void this.release();
}
/**
* Let the held calls launch, and later live calls, each once its issued count is committed. On a rerun, an earlier
* attempt may have admitted some of them under their keys, as other calls: when the replay releases (idle after the
* last journaled result, or after its first run without a journal) its counts matched every recorded one, so it has
* made every call an earlier attempt could have admitted. All held calls are checked first, whatever their target,
* and a divergence launches none. Calls made after this are new.
*/
private async release(): Promise<void> {
await this.gate;
if (this.stopped || this.released) return;
const calls = this.held.splice(0);
if (this.rerun) {
for (const call of calls) {
let other: string | undefined;
try {
other = await this.earlier(call);
} catch (error) {
this.fail(error);
return;
}
if (this.stopped) return;
if (other !== undefined) {
this.diverge(`call ${call.id} is ${describe(call)}, but was ${other}`);
return;
}
}
}
this.released = true;
this.waiting.push(...calls, ...this.held.splice(0));
this.launchCovered();
}
/** Launch the waiting calls whose issued count is committed. */
private launchCovered(): void {
if (!this.released || this.stopped) return;
const covered = this.waiting.filter((call) => call.id <= this.durableIssued);
if (covered.length === 0) return;
this.waiting = this.waiting.filter((call) => call.id > this.durableIssued);
for (const call of covered) this.launch(call);
}
/** The call an earlier attempt made under this call's key, when it is another one: another target, tool, or arguments. */
private async earlier(call: CodemodeHostCall): Promise<string | undefined> {
const { api, context } = this;
const taskId = (await api.snapshot(NestedCallDoc, api.taskId, `c${call.id}`, context))?.taskId;
if (taskId === undefined) return undefined;
const input = (await api.getTask(taskId, context))?.input as ToolTaskInput | undefined;
if (input?.kind !== "nested") return undefined;
const made = input.call;
if (call.target === "tool" && made.name === call.name && sameJson(made.arguments, call.args ?? {}))
return undefined;
return `tools.${made.name}(${JSON.stringify(made.arguments)})`;
}
/** Run one live call and journal its result. */
private launch(call: CodemodeHostCall): void {
if (this.stopped) return;
const global = call.target === "global" ? this.options.globals?.find((g) => g.name === call.name) : undefined;
const running: Promise<{ outcome: CodemodeOutcome; global?: JournalEntry["global"] }> =
call.target === "tool"
? this.callTool(call).then((outcome) => ({ outcome }))
: this.callGlobal(call, global).then((outcome) => ({
outcome,
global: {
name: call.name,
...(call.args === undefined ? {} : { args: call.args as JsonValue }),
outcome: outcome as JsonObject,
},
}));
running.then(
(done) => {
this.queue.push({ id: call.id, at: Date.now(), ...done });
void this.pump();
},
(error: unknown) => {
if (this.context.abortSignal?.aborted) this.fail(error);
// A key an earlier attempt admitted with another tool or other arguments: this attempt took another path.
else this.diverge(`call ${call.id} is ${describe(call)}: ${errorMessage(error)}`);
},
);
}
private async callTool(call: CodemodeHostCall): Promise<CodemodeOutcome> {
// Passed as the script gave them; validation reports arguments that are not an object.
const args = (call.args ?? {}) as JsonObject;
const result = await this.api.executeTool(call.name, args, this.context, { key: `c${call.id}` });
return outcomeOf(call.name, result);
}
private async callGlobal(call: CodemodeHostCall, global: CodemodeGlobal | undefined): Promise<CodemodeOutcome> {
if (global === undefined) return { ok: false, message: `Unknown global ${call.name}` };
try {
// `undefined` stays undefined: the journal omits the value, and a replay resolves with undefined again.
return { ok: true, value: await global.execute(call.args, this.api, this.context) };
} catch (error) {
if (this.context.abortSignal?.aborted) throw error;
return { ok: false, message: errorMessage(error) };
}
}
/**
* Journal finished calls and acknowledged counts in one commit, then deliver the calls and launch the calls the
* counts cover; one commit in flight at a time, so what arrives meanwhile goes into the next.
*/
private async pump(): Promise<void> {
if (this.committing || this.stopped) return;
if (this.queue.length === 0 && this.acks.size === 0) return;
this.committing = true;
const batch = this.queue.splice(0);
const acks = [...this.acks];
this.acks.clear();
try {
await this.api.commit(async (tx) => {
const journal = await tx.doc(CodemodeJournalDoc, this.api.taskId);
for (const [seq, issued] of acks) {
if (seq === 0) journal.firstIssued ??= issued;
else {
const entry = journal.entries[seq - 1];
if (entry !== undefined && entry.issuedAfter === undefined) entry.issuedAfter = issued;
}
}
for (const item of batch) {
const entry = {
id: item.id,
at: item.at,
...(item.global === undefined ? {} : { global: item.global }),
};
journal.entries.push(
copyJson(entry as unknown as JsonValue, { omitUndefinedProperties: true }) as unknown as JournalEntry,
);
}
}, this.context);
} catch (error) {
this.fail(error);
return;
}
for (const [, issued] of acks) this.durableIssued = Math.max(this.durableIssued, issued);
for (const item of batch) {
this.journalLength++;
if (!this.run.finished) this.run.deliver(this.journalLength, item.at, item.id, item.outcome);
}
this.launchCovered();
this.committing = false;
void this.pump();
}
}
function describe(call: CodemodeHostCall): string {
return `${call.target === "tool" ? "tools." : ""}${call.name}(${JSON.stringify(call.args) ?? ""})`;
}
/** Structural JSON equality; key order does not matter. */
function sameJson(a: unknown, b: unknown): boolean {
if (a === b) return true;
if (typeof a !== "object" || typeof b !== "object" || a === null || b === null) return false;
if (Array.isArray(a) || Array.isArray(b)) {
return (
Array.isArray(a) && Array.isArray(b) && a.length === b.length && a.every((item, i) => sameJson(item, b[i]))
);
}
const keys = Object.keys(a);
return (
keys.length === Object.keys(b).length &&
keys.every(
(key) =>
Object.hasOwn(b, key) && sameJson((a as Record<string, unknown>)[key], (b as Record<string, unknown>)[key]),
)
);
}

10
packages/durable/src/codemode/index.ts Normal file
View file

@ -0,0 +1,10 @@
export {
CODEMODE_MODULE_VERSION,
CODEMODE_TOOL_NAME,
type CodemodeGlobal,
CodemodeStoreDoc,
type CodemodeToolInput,
type CodemodeToolOptions,
createCodemodeExtension,
createCodemodeTool,
} from "./codemode.ts";

775
packages/durable/test/codemode.test.ts Normal file
View file

@ -0,0 +1,775 @@
import { mkdtemp, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { type AssistantMessage, fauxAssistantMessage, fauxText, fauxToolCall, Type } from "@earendil-works/pi-ai";
import { type CodemodeRun, CodemodeSandbox } from "@earendil-works/pi-codemode";
import {
defineTool,
type Harness,
type TaskId,
type ToolRegistration,
ToolResultEntry,
ToolTask,
type ToolTaskInput,
} from "@earendil-works/pi-durable";
import { afterEach, describe, expect, it } from "vitest";
import {
type CodemodeGlobal,
CodemodeJournalDoc,
CodemodeRunDoc,
createCodemodeTool,
} from "../src/codemode/codemode.ts";
import { openNodeSqliteStorage } from "../src/storage/sqlite/node.ts";
import { allEntries, type ChatSetup, chatSetup, openChat } from "./chat-support.ts";
import { addHooks, addTool } from "./harness-support.ts";
import { ControlledStorage, context } from "./session-support.ts";
import { aborted, type Deferred, deferred, eventually } from "./task-support.ts";
const Parameters = Type.Object({ text: Type.Optional(Type.String()), id: Type.Optional(Type.Number()) });
function tool(
name: string,
execute: ToolRegistration<typeof Parameters>["execute"],
replay: "safe" | "unsafe" = "unsafe",
) {
return defineTool({ name, description: name, parameters: Parameters, execute, replay });
}
const DONE = fauxAssistantMessage([fauxText("done")]);
function script(code: string): AssistantMessage {
return fauxAssistantMessage([fauxToolCall("codemode", { code }, { id: "s1" })], { stopReason: "toolUse" });
}
const directories = new Set<string>();
afterEach(async () => {
for (const directory of directories) await rm(directory, { recursive: true, force: true });
directories.clear();
});
async function sqlitePath(): Promise<string> {
const directory = await mkdtemp(join(tmpdir(), "pi-durable-codemode-"));
directories.add(directory);
return join(directory, "session.sqlite");
}
/**
* A chat with the code mode tool and test tools:
* - `echo` returns `echo <text>` and counts its runs by text;
* - `hold` (replay-safe) waits for `holds.gate` while `holds.block` is set, then returns `held <id>`;
* - `slow` returns `slow <id>` once `slows.release(id)` was called.
*/
function codemodeChat(globals: CodemodeGlobal[] = []) {
const setup = chatSetup();
const sandbox = new CodemodeSandbox({ memoryLimitBytes: 64 * 1024 * 1024 });
const vm: { starts: number; onStart?: (run: CodemodeRun) => void } = { starts: 0 };
const start = sandbox.start.bind(sandbox);
sandbox.start = (code, options) => {
vm.starts++;
const run = start(code, options);
vm.onStart?.(run);
return run;
};
addTool(setup.registry, createCodemodeTool({ sandbox, globals }));
const echo = new Map<string, number>();
addTool(
setup.registry,
tool("echo", async (args) => {
echo.set(args.text ?? "", (echo.get(args.text ?? "") ?? 0) + 1);
return { output: [{ type: "text", text: `echo ${args.text ?? ""}` }] };
}),
);
const holds = { block: true, reached: deferred(), runs: 0 };
addTool(
setup.registry,
tool(
"hold",
async (args, _api, callContext) => {
holds.runs++;
if (holds.block) {
holds.reached.resolve();
await aborted(callContext.abortSignal!);
}
return { output: [{ type: "text", text: `held ${args.id ?? ""}` }] };
},
"safe",
),
);
const pauses = { block: true, reached: deferred() };
addTool(
setup.registry,
tool(
"pause",
async (_args, _api, callContext) => {
if (pauses.block) {
pauses.reached.resolve();
await aborted(callContext.abortSignal!);
}
return { output: [{ type: "text", text: "paused" }] };
},
"safe",
),
);
const gates = new Map<number, Deferred>();
const finished: number[] = [];
const slows = {
finished,
gate: (id: number) => {
const gate = gates.get(id) ?? deferred();
gates.set(id, gate);
return gate;
},
release: (id: number) => slows.gate(id).resolve(),
};
addTool(
setup.registry,
tool("slow", async (args) => {
await slows.gate(args.id ?? 0).promise;
finished.push(args.id ?? 0);
return { output: [{ type: "text", text: `slow ${args.id}` }] };
}),
);
return { setup, sandbox, vm, echo, holds, pauses, slows };
}
async function scriptTask(harness: Harness): Promise<TaskId> {
const tasks = (await harness.commit((tx) => tx.scanTasks({ kind: ToolTask.definition.name }, 100), context)).items;
return tasks.find((task) => (task.input as ToolTaskInput).kind === "model")!.id;
}
async function nestedCount(harness: Harness): Promise<number> {
const tasks = (await harness.commit((tx) => tx.scanTasks({ kind: ToolTask.definition.name }, 100), context)).items;
return tasks.filter((task) => (task.input as ToolTaskInput).kind === "nested").length;
}
/** The code mode call's result text. */
async function resultText(setup: ChatSetup, root: Awaited<ReturnType<typeof openChat>>["root"]) {
void setup;
const entry = (await allEntries(root)).find((candidate) => ToolResultEntry.is(candidate))!;
const message = entry.model![0]!;
return {
text:
message.role === "toolResult"
? message.content.map((item) => (item.type === "text" ? item.text : "")).join("\n")
: "",
isError: message.role === "toolResult" && message.isError,
};
}
/**
* Run `code` until `holds.reached`, close the Harness (a restart), let `between` change the stored state, then reopen
* with `hold` no longer blocking and wait for the answer.
*/
async function restart(
chat: ReturnType<typeof codemodeChat>,
code: string,
between: (harness: Harness, task: TaskId) => Promise<void> = async () => {},
beforeClose: (harness: Harness, task: TaskId) => Promise<void> = async () => {},
) {
const path = await sqlitePath();
chat.setup.faux.setResponses([script(code), DONE]);
let opened = await openChat(await openNodeSqliteStorage(path), chat.setup);
const id = (await opened.root.submit({ type: "input", content: "go" }, context)).id;
await chat.holds.reached.promise;
await beforeClose(opened.harness, await scriptTask(opened.harness));
await opened.harness.close(context);
chat.holds.block = false;
opened = await openChat(await openNodeSqliteStorage(path), chat.setup);
await between(opened.harness, await scriptTask(opened.harness));
opened.harness.resume();
expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done");
return { ...opened, result: await resultText(chat.setup, opened.root) };
}
describe("durable code mode", () => {
it("runs a script that calls tools as nested calls, and journals each result before delivering it", async () => {
const chat = codemodeChat();
chat.holds.block = false;
chat.setup.faux.setResponses([
script(`const a = await tools.echo({ text: "a" }); const b = await tools.echo({ text: "b" }); return [a, b];`),
DONE,
]);
const { harness, root } = await openChat(await openNodeSqliteStorage(await sqlitePath()), chat.setup);
await (await root.submit({ type: "input", content: "go" }, context)).wait(context);
expect(await resultText(chat.setup, root)).toEqual({ text: '["echo a","echo b"]', isError: false });
expect(await nestedCount(harness)).toBe(2);
await harness.close(context);
});
it("resumes a script after a restart without running its finished calls again", async () => {
const chat = codemodeChat();
const { harness, result } = await restart(
chat,
`const a = await tools.echo({ text: "a" });
const held = await tools.hold({ id: 1 });
const b = await tools.echo({ text: "b" });
return [a, held, b, Date.now() > 0];`,
);
expect(result).toEqual({ text: '["echo a","held 1","echo b",true]', isError: false });
expect(Object.fromEntries(chat.echo)).toEqual({ a: 1, b: 1 });
// The VM ran twice; hold, replay-safe, reran once after the restart.
expect(chat.vm.starts).toBe(2);
expect(chat.holds.runs).toBe(2);
await harness.close(context);
});
it("replays results in the order they arrived, not the order the calls were made", async () => {
const chat = codemodeChat();
const code = `const order = [];
await Promise.all([tools.slow({ id: 1 }).then(() => order.push(1)), tools.slow({ id: 2 }).then(() => order.push(2))]);
await tools.echo({ text: order.join(",") });
await tools.hold({});
return order;`;
// 2 finishes first, then 1.
chat.slows.release(2);
void eventually(() => chat.slows.finished.includes(2)).then(() => chat.slows.release(1));
const { harness, result } = await restart(chat, code);
expect(result).toEqual({ text: "[2,1]", isError: false });
expect(Object.fromEntries(chat.echo)).toEqual({ "2,1": 1 });
await harness.close(context);
});
it("journals globals with their results and hands a replay the recorded ones", async () => {
let counter = 0;
const globals: CodemodeGlobal[] = [{ name: "next", execute: async () => ++counter }];
const chat = codemodeChat(globals);
const { harness, result } = await restart(
chat,
`const n = await next(); await tools.hold({}); return [n, await next()];`,
);
// The first next() ran once and was replayed from the journal; the second ran after the restart.
expect(result).toEqual({ text: "[1,2]", isError: false });
expect(counter).toBe(2);
await harness.close(context);
});
it("keeps the script's clock and randomness across a restart", async () => {
const chat = codemodeChat();
const { harness, result } = await restart(
chat,
`const before = [Date.now(), Math.random()];
await tools.hold({});
await tools.echo({ text: String(before[0]) + "/" + String(before[1]) });
return "ok";`,
);
expect(result.text).toBe("ok");
// One echo: the same arguments, so the rerun admitted no second call.
expect([...chat.echo.values()]).toEqual([1]);
await harness.close(context);
});
it("ends the call when a resumed script made a different number of calls than before", async () => {
const chat = codemodeChat();
const { harness, result } = await restart(
chat,
`await tools.echo({ text: "a" }); await tools.hold({});`,
async (opened, task) => {
await opened.commit(async (tx) => {
const journal = await tx.doc(CodemodeJournalDoc, task);
journal.entries[0]!.issuedAfter = 7;
}, context);
},
);
expect(result.isError).toBe(true);
expect(result.text).toContain("after result 1 it had made 2 calls, where it had made 7");
await harness.close(context);
});
it("ends the call at once when a resumed script waits without making a call it made before", async () => {
const chat = codemodeChat();
const { harness, result } = await restart(
chat,
`await tools.echo({ text: "a" }); await tools.hold({});`,
async (opened, task) => {
await opened.commit(async (tx) => {
const journal = await tx.doc(CodemodeJournalDoc, task);
journal.entries[0]!.id = 99;
}, context);
},
);
expect(result.isError).toBe(true);
expect(result.text).toContain("it waits without making call 99");
// hold, admitted before the restart, resumed on its own; no call was admitted by the replay.
expect(await nestedCount(harness)).toBe(2);
await harness.close(context);
});
it("ends the call when a resumed script calls a global with other arguments", async () => {
const chat = codemodeChat([{ name: "next", execute: async () => 1 }]);
const { harness, result } = await restart(chat, `await next(); await tools.hold({});`, async (opened, task) => {
await opened.commit(async (tx) => {
const journal = await tx.doc(CodemodeJournalDoc, task);
journal.entries[0]!.global!.args = "other";
}, context);
});
expect(result.isError).toBe(true);
expect(result.text).toContain("call 1 is next");
await harness.close(context);
});
it("admits no call of a diverging replay, even one the script never made before", async () => {
const chat = codemodeChat();
// With k set, the replay makes three extra calls after result 1, which the live run did not make; the count in
// result 1's acknowledgment shows the divergence only after they were issued.
const { harness, result } = await restart(
chat,
`await tools.echo({ text: "a" });
const extra = load("k") ?? 0;
for (let i = 0; i < extra; i++) tools.echo({ text: "extra" + i });
await tools.hold({});`,
async (opened, task) => {
await opened.commit(async (tx) => {
const run = await tx.doc(CodemodeRunDoc, task);
run.start!.store = { k: 3 };
}, context);
},
async (opened, task) => {
await eventually(
async () => (await opened.snapshot(CodemodeJournalDoc, task, context))?.entries[0]?.issuedAfter === 2,
);
},
);
expect(result.isError).toBe(true);
expect(result.text).toContain("after result 1 it had made 5 calls, where it had made 2");
expect(Object.fromEntries(chat.echo)).toEqual({ a: 1 });
expect(await nestedCount(harness)).toBe(2);
await harness.close(context);
});
it("ends the call when a resumed script ends where it went on before", async () => {
const chat = codemodeChat();
const { harness, result } = await restart(
chat,
`await tools.echo({ text: "a" }); if (load("k")) return "early"; await tools.hold({});`,
async (opened, task) => {
await opened.commit(async (tx) => {
const run = await tx.doc(CodemodeRunDoc, task);
run.start!.store = { k: true };
}, context);
},
// The live run acknowledged result 1 before the restart: it went on after it.
async (opened, task) => {
await eventually(
async () => (await opened.snapshot(CodemodeJournalDoc, task, context))?.entries[0]?.issuedAfter === 2,
);
},
);
expect(result.isError).toBe(true);
expect(result.text).toContain("it ended after result 1, where it went on before");
await harness.close(context);
});
it("ends the call when a resumed script ends before calls the live run made, even without a recorded count", async () => {
const chat = codemodeChat();
const { harness, result } = await restart(
chat,
`await tools.echo({ text: "a" }); if (load("k")) return "early"; await tools.echo({ text: "b" }); await tools.hold({});`,
async (opened, task) => {
await opened.commit(async (tx) => {
const run = await tx.doc(CodemodeRunDoc, task);
run.start!.store = { k: true };
// As if the live run crashed before recording result 1's count.
const journal = await tx.doc(CodemodeJournalDoc, task);
delete journal.entries[0]!.issuedAfter;
}, context);
},
);
expect(result.isError).toBe(true);
expect(result.text).toContain("it ended after result 1, where it went on before");
await harness.close(context);
});
it("ends the call when a resumed script makes a call an earlier attempt admitted with other arguments", async () => {
const chat = codemodeChat();
const { harness, result } = await restart(
chat,
`await tools.echo({ text: load("k") ?? "a" }); await tools.hold({});`,
async (opened, task) => {
await opened.commit(async (tx) => {
const run = await tx.doc(CodemodeRunDoc, task);
run.start!.store = { k: "other" };
// As if the restart came before the first journal commit: the rerun runs live from the start.
const journal = await tx.doc(CodemodeJournalDoc, task);
journal.entries.length = 0;
}, context);
},
);
expect(result.isError).toBe(true);
expect(result.text).toContain("different path");
expect(result.text).toContain("call 1 is tools.echo");
expect(Object.fromEntries(chat.echo)).toEqual({ a: 1 });
await harness.close(context);
});
it("records a count the live run never recorded as soon as a replay acknowledges it", async () => {
const chat = codemodeChat();
const path = await sqlitePath();
chat.setup.faux.setResponses([script(`await tools.echo({ text: "a" }); await tools.hold({});`), DONE]);
let opened = await openChat(await openNodeSqliteStorage(path), chat.setup);
const id = (await opened.root.submit({ type: "input", content: "go" }, context)).id;
await chat.holds.reached.promise;
await opened.harness.close(context);
// As if result 1's count was never recorded; hold blocks again after the restart.
chat.holds.reached = deferred();
opened = await openChat(await openNodeSqliteStorage(path), chat.setup);
const task = await scriptTask(opened.harness);
await opened.harness.commit(async (tx) => {
delete (await tx.doc(CodemodeJournalDoc, task)).entries[0]!.issuedAfter;
}, context);
opened.harness.resume();
// The replay acknowledges result 1 and goes live; the script waits on hold, and the count is already recorded.
await chat.holds.reached.promise;
await eventually(
async () => (await opened.harness.snapshot(CodemodeJournalDoc, task, context))?.entries[0]?.issuedAfter === 2,
);
await opened.harness.close(context);
chat.holds.block = false;
opened = await openChat(await openNodeSqliteStorage(path), chat.setup);
opened.harness.resume();
expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done");
await opened.harness.close(context);
});
it("aborts a call still running when the script ends", async () => {
const chat = codemodeChat();
chat.setup.faux.setResponses([script(`tools.hold({}); await tools.echo({ text: "a" }); return 1;`), DONE]);
const { harness, root } = await openChat(await openNodeSqliteStorage(await sqlitePath()), chat.setup);
await (await root.submit({ type: "input", content: "go" }, context)).wait(context);
expect(await resultText(chat.setup, root)).toEqual({ text: "1", isError: false });
const tasks = (await harness.commit((tx) => tx.scanTasks({ kind: ToolTask.definition.name }, 100), context))
.items;
const hold = tasks.find(
(task) =>
(task.input as ToolTaskInput).kind === "nested" &&
(task.input as { call: { name: string } }).call.name === "hold",
)!;
expect(chat.holds.runs).toBe(1);
expect(hold.state).toMatchObject({ status: "terminal", outcome: { status: "aborted" } });
await harness.close(context);
});
it("admits no call once the script has ended, even one it made before ending", async () => {
const chat = codemodeChat();
chat.holds.block = false;
const storage = new ControlledStorage();
// Hold the fingerprint commit, and so the admission of the script's call, until the script has ended.
chat.vm.onStart = (run) => {
const held = storage.holdCommits();
void run.result.then(() => held.release());
};
chat.setup.faux.setResponses([script(`tools.echo({ text: "late" }); return 1;`), DONE]);
const { harness, root } = await openChat(storage, chat.setup);
await (await root.submit({ type: "input", content: "go" }, context)).wait(context);
expect(await resultText(chat.setup, root)).toEqual({ text: "1", isError: false });
expect(chat.echo.size).toBe(0);
expect(await nestedCount(harness)).toBe(0);
await harness.close(context);
});
it("checks a rerun's calls against an earlier attempt's before admitting any later call", async () => {
const chat = codemodeChat();
// With k set, call 1 has other arguments and call 2 is another tool than before; call 3 would be new.
const { harness, result } = await restart(
chat,
`const k = load("k");
tools.echo({ text: k ?? "a" });
if (k) tools.echo({ text: "new" });
await tools.hold({});`,
async (opened, task) => {
await opened.commit(async (tx) => {
const run = await tx.doc(CodemodeRunDoc, task);
run.start!.store = { k: "x" };
// As if the restart came before the first journal commit.
(await tx.doc(CodemodeJournalDoc, task)).entries.length = 0;
}, context);
},
);
expect(result.isError).toBe(true);
// The first run already made more calls than before.
expect(result.text).toContain("its first run made 3 calls, where it had made 2");
expect(await nestedCount(harness)).toBe(2);
expect(Object.fromEntries(chat.echo)).toEqual({ a: 1 });
await harness.close(context);
});
it("ends the call when a rerun makes a global call where an earlier attempt made a tool call", async () => {
const chat = codemodeChat([{ name: "pick", execute: async () => 1 }]);
const { harness, result } = await restart(
chat,
`if (load("k")) { await pick(); await tools.echo({ text: "new" }); } else { await tools.hold({}); }`,
async (opened, task) => {
await opened.commit(async (tx) => {
const run = await tx.doc(CodemodeRunDoc, task);
run.start!.store = { k: true };
}, context);
},
);
expect(result.isError).toBe(true);
expect(result.text).toContain("call 1 is pick(), but was tools.hold({})");
expect(chat.echo.size).toBe(0);
await harness.close(context);
});
it("ends the call when a rerun's first run makes fewer calls than an earlier attempt admitted", async () => {
// Live, the first run admits hold 1 and hold 2. With k set, it makes only hold 1, and after its result a global and a
// new tool call: the global's ID is hold 2's key, which no check after the release would see.
const chat = codemodeChat([{ name: "pick", execute: async () => 1 }]);
const { harness, result } = await restart(
chat,
`const k = load("k");
const first = tools.hold({ id: 1 });
if (!k) tools.hold({ id: 2 });
await first;
await pick();
await tools.echo({ text: "new" });`,
async (opened, task) => {
await opened.commit(async (tx) => {
const run = await tx.doc(CodemodeRunDoc, task);
run.start!.store = { k: true };
}, context);
},
async (opened, task) => {
await eventually(async () => (await opened.snapshot(CodemodeJournalDoc, task, context))?.firstIssued === 2);
},
);
expect(result.isError).toBe(true);
expect(result.text).toContain("its first run made 1 calls, where it had made 2");
expect(chat.echo.size).toBe(0);
await harness.close(context);
});
it("ends the call when a rerun without journal entries ends in its first run, where it went on before", async () => {
const chat = codemodeChat();
const { harness, result } = await restart(
chat,
`if (load("k")) { store("x", 1); return "early"; } await tools.hold({});`,
async (opened, task) => {
await opened.commit(async (tx) => {
const run = await tx.doc(CodemodeRunDoc, task);
run.start!.store = { k: true };
}, context);
},
async (opened, task) => {
await eventually(async () => (await opened.snapshot(CodemodeJournalDoc, task, context))?.firstIssued === 1);
},
);
expect(result.isError).toBe(true);
expect(result.text).toContain("it ended after result 0, where it went on before");
await harness.close(context);
});
it("checks every call a rerun holds before launching any of them", async () => {
// Live, call 1 is a global that never finishes and call 2 is hold; with k set, call 1 is a new tool call and call 2
// holds with other arguments. Checking call 1 alone would launch it before call 2's conflict shows.
const chat = codemodeChat([
{
name: "forever",
execute: async (_args, _api, callContext) => {
await aborted(callContext.abortSignal!);
return null;
},
},
]);
const { harness, result } = await restart(
chat,
`const k = load("k");
if (k) tools.echo({ text: "new" }); else forever();
await tools.hold({ id: k ? 2 : 1 });`,
async (opened, task) => {
await opened.commit(async (tx) => {
const run = await tx.doc(CodemodeRunDoc, task);
run.start!.store = { k: true };
}, context);
},
);
expect(result.isError).toBe(true);
expect(result.text).toContain('call 2 is tools.hold({"id":2}), but was tools.hold({"id":1})');
expect(chat.echo.size).toBe(0);
expect(await nestedCount(harness)).toBe(1);
await harness.close(context);
});
it("keeps a global's undefined result undefined, live and replayed", async () => {
const chat = codemodeChat([{ name: "nothing", execute: async () => undefined }]);
const { harness, result } = await restart(
chat,
`const before = await nothing(); await tools.hold({}); return [before === undefined, (await nothing()) === undefined];`,
);
expect(result).toEqual({ text: "[true,true]", isError: false });
await harness.close(context);
});
it("rejects a failed tool call with its error text and structured output", async () => {
const chat = codemodeChat();
addTool(
chat.setup.registry,
tool("fails", async () => ({ output: [{ type: "text", text: "permission denied" }], isError: true })),
);
chat.holds.block = false;
chat.setup.faux.setResponses([
script(`try { await tools.fails({}); } catch (e) { return [e.message, e.structuredOutput]; }`),
DONE,
]);
const { harness, root } = await openChat(await openNodeSqliteStorage(await sqlitePath()), chat.setup);
await (await root.submit({ type: "input", content: "go" }, context)).wait(context);
expect(await resultText(chat.setup, root)).toEqual({
text: '["permission denied","permission denied"]',
isError: false,
});
await harness.close(context);
});
it("delivers a result only after its journal commit", async () => {
const storage = new ControlledStorage();
const seen: string[] = [];
const chat = codemodeChat([
{
name: "first",
execute: async () => {
seen.push("first");
// The next commit is this result's journal commit: it never lands.
storage.failNextCommit(new Error("lost"));
return 1;
},
},
{
name: "second",
execute: async () => {
seen.push("second");
return 2;
},
},
]);
chat.holds.block = false;
chat.setup.faux.setResponses([script(`await first(); await second();`), DONE]);
const { harness, root } = await openChat(storage, chat.setup);
await root.submit({ type: "input", content: "go" }, context);
await eventually(() => seen.includes("first"));
await new Promise((resolve) => setTimeout(resolve, 100));
expect(seen).toEqual(["first"]);
await harness.close(context).catch(() => {});
});
it("refuses to resume when the sandbox's fingerprint or the code mode version changed", async () => {
for (const field of ["fingerprint", "version"] as const) {
const chat = codemodeChat();
const { harness, result } = await restart(
chat,
`await tools.hold({}); await tools.echo({ text: "late" });`,
async (opened, task) => {
await opened.commit(async (tx) => {
const run = await tx.doc(CodemodeRunDoc, task);
if (field === "fingerprint") run.start!.fingerprint = "another sandbox";
else run.start!.version = 0;
}, context);
},
);
expect(result.isError).toBe(true);
expect(result.text).toContain("cannot resume");
expect(chat.echo.size).toBe(0);
await harness.close(context);
}
});
it("times a script out across a restart, without admitting new calls", async () => {
const chat = codemodeChat();
const { harness, result } = await restart(
chat,
`// @options: {"timeout_ms": 300}
await tools.hold({}); await tools.echo({ text: "late" });`,
() => new Promise((resolve) => setTimeout(resolve, 400)),
);
expect(result.isError).toBe(true);
expect(result.text).toContain("deadline passed");
expect(chat.echo.size).toBe(0);
expect(chat.vm.starts).toBe(1);
await harness.close(context);
});
it("returns a decided outcome after a restart without running the script again", async () => {
const chat = codemodeChat();
const path = await sqlitePath();
chat.holds.block = false;
chat.setup.faux.setResponses([script(`return "first"`), DONE]);
// Stop after execute() returned, so after the outcome is stored, and before the result commits.
const stopped = { block: true, reached: deferred() };
addHooks(chat.setup.registry, ToolTask, {
afterTool: async (call, _result, _api, hookContext) => {
if (call.name !== "codemode" || !stopped.block) return undefined;
stopped.reached.resolve();
await aborted(hookContext.abortSignal!);
return undefined;
},
});
let opened = await openChat(await openNodeSqliteStorage(path), chat.setup);
const id = (await opened.root.submit({ type: "input", content: "go" }, context)).id;
await stopped.reached.promise;
const task = await scriptTask(opened.harness);
expect((await opened.harness.snapshot(CodemodeRunDoc, task, context))?.result).toBeDefined();
await opened.harness.close(context);
stopped.block = false;
opened = await openChat(await openNodeSqliteStorage(path), chat.setup);
opened.harness.resume();
expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done");
expect(await resultText(chat.setup, opened.root)).toEqual({ text: "first", isError: false });
expect(chat.vm.starts).toBe(1);
await opened.harness.close(context);
});
it("keeps store() values per conversation, as of a fork's entry", async () => {
const chat = codemodeChat();
chat.holds.block = false;
chat.setup.faux.setResponses([
script(`store("k", 1); return load("k")`),
DONE,
script(`return [load("k") ?? null]`),
DONE,
]);
const { harness, root } = await openChat(await openNodeSqliteStorage(await sqlitePath()), chat.setup);
await (await root.submit({ type: "input", content: "first" }, context)).wait(context);
const firstUser = (await allEntries(root)).find((entry) => entry.kind === "pi.user")!;
await (await root.submit({ type: "input", content: "second" }, context)).wait(context);
const results = (await allEntries(root)).filter((entry) => ToolResultEntry.is(entry));
const texts = results.map((entry) => (entry.model![0] as { content: { text: string }[] }).content[0]!.text);
expect(texts).toEqual(["1", "[1]"]);
// A fork at the first input is from before the first script stored anything.
chat.setup.faux.setResponses([script(`return [load("k") ?? null]`), DONE]);
const fork = await root.fork(firstUser.id, { ownership: { kind: "ownerless" } }, context);
await (await fork.submit({ type: "input", content: "fork" }, context)).wait(context);
const forkResult = (await allEntries(fork)).filter((entry) => ToolResultEntry.is(entry)).at(-1)!;
expect((forkResult.model![0] as { content: { text: string }[] }).content[0]!.text).toBe("[null]");
await harness.close(context);
});
it("journals results that finish together in one commit", async () => {
const release = deferred();
let waiting = 0;
const chat = codemodeChat([
{
name: "wait",
execute: async () => {
if (++waiting === 50) release.resolve();
await release.promise;
return null;
},
},
]);
chat.holds.block = false;
chat.setup.faux.setResponses([
script(`return (await Promise.all(Array.from({ length: 50 }, () => wait()))).length`),
DONE,
]);
const { harness, root } = await openChat(await openNodeSqliteStorage(await sqlitePath()), chat.setup);
let journalCommits = 0;
harness.subscribeCommits((publication) => {
for (const change of publication.changes) {
if (change.type === "document" && change.record.kind === "pi.codemode.journal") journalCommits++;
}
});
await (await root.submit({ type: "input", content: "go" }, context)).wait(context);
expect(await resultText(chat.setup, root)).toEqual({ text: "50", isError: false });
// The journal's creation and a few groups, not one commit per result (51).
expect(journalCommits).toBeLessThan(10);
await harness.close(context);
});
});