feat(durable): code mode listing, tool search, output, models.*, grammar; tool error rule

- structuredOutput is always the tool's own value; a schema tool's error text
  without one becomes a tool_error diagnostic on nested results
- code mode: a schema tool's error result with its value resolves (bash nonzero
  exit), schema-less tools resolve with their text, as in the coding agent
- codemode prompt section with tool declarations by namespace within an inline
  budget; searchTools/describeTool/describeNamespace from the start data
- result layout of the coding agent: header, text markers, console block,
  truncation within max_output_tokens and the tool's output limits, spill and
  image files in a recorded output directory (a rerun overwrites them)
- models.*: catalog globals; classify/generateImages as hidden internal tools
  (nested calls with usage), at most four at once
- Lark grammar of the script source as constrainedSampling
This commit is contained in:
Mario Zechner 2026-10-12 00:06:09 +00:00
parent 04566631c2
commit 8bb180bec1
15 changed files with 1766 additions and 103 deletions

View file

@ -16,12 +16,12 @@
- Input submissions take `entry: { kind, data }` (`InputEntry`): the input is placed as an entry of that kind carrying that data, with its content as the user message, instead of `pi.user`, and behaves as any other input; the kind must not start with `pi.`. Queued inputs carry it in their `pi.inbox` item ([#10407](https://github.com/earendil-works/pi/issues/10407)).
- `Session.closed` settles once the Session has closed, Storage included: `{ reason: "closed" }` after `close()`, or `{ reason: "failed", error }` after a storage error. `WatchEnd` gains `{ reason: "session_failed", error }`, and the README has an Errors section.
- `ToolExecutionApi.executeTool(name, args, context, { key?, progress? })`: a tool calls another tool as a nested call. The nested call is its own `pi.tool` task owned by the caller, runs validation, the `ToolTask` hooks (whose `ToolHookCall` carries `parent`), output limits, and its replay policy, and returns a `NestedToolExecutionResult` (its `taskId`, `structuredOutput`, `details`, `diagnostics`, `usage`, `isError`, `durationMs`, but not the model's `output`) to the caller instead of the transcript. Each nested call is a member of `NestedCallDoc`, a task-scoped document family of the caller keyed by the call's key, holding its `taskId` and, once it ends, its `result` as returned, so it is gone from storage once the caller settles; its task's receipt stays small. Keys, by default the call's position, let a replay-safe caller's rerun reattach to its nested calls; a caller that is not replay-safe has its unfinished nested calls abandoned after a restart. Nested calls show in `pi.live.nestedTools`, with the arguments they run with and a `summary` (error, duration, usage) once done (`progress: false` keeps their running output out), and as tool events with `parentToolCallId` and `parentTaskId`; every tool event also carries its call's `taskId`. A caller aborts the nested calls it left running before it settles.
- Structured tool output: a tool declares `structuredOutputSchema` and returns `structuredOutput`, which the Harness validates and gives to programs that call the tool; a tool without a schema gives them its bounded output, one text item as a string, one image as itself, and otherwise the content list, for errors too, and so does a schema tool's error result without `structuredOutput`, so programs get the text of a returned error; a thrown error's message is in its `tool_error` diagnostic. A nested call with invalid structured output gets an error result; a model-issued one only loses it, and the host gets a report. `api.retainedOutput()` reads the bounded output the model will see. `bash` and `powershell` return `{ output, truncated, fullOutputPath?, exitCode }`.
- Structured tool output: a tool declares `structuredOutputSchema` and returns `structuredOutput`, which the Harness validates and gives to programs that call the tool; a tool without a schema gives them its bounded output, one text item as a string, one image as itself, and otherwise the content list, for errors too. `structuredOutput` is always the tool's own value: a schema tool's error result keeps the value it returned, and the text of one without becomes a `tool_error` diagnostic of the nested result, as a thrown error's message does. A nested call with invalid structured output gets an error result; a model-issued one only loses it, and the host gets a report. `api.retainedOutput()` reads the bounded output the model will see. `bash` and `powershell` return `{ output, truncated, fullOutputPath?, exitCode }`.
- `ToolRegistration.callers` (`model`, `tools`, or both, the default) and `pi.agent` `modelTools`: who may call a tool, and which of its tools a conversation offers the model; `Agent.callable` lists the tools nested calls resolve among. `addTools` also extends `modelTools`.
- `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. The result's `details.calls` list the script's tool calls with their status; live progress is in `pi.live.nestedTools` (`npm run bench:codemode`).
- `@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. As in the coding agent, a `codemode` prompt section lists the tools scripts can call (`namespaces`, an inline token budget), scripts find tools with `searchTools()`, `describeTool()`, and `describeNamespace()`, and the result is laid out with a header, text item markers, a console block, truncation to `max_output_tokens` with a spill file, and saved images. Scripts get `models.*` (the catalog, `models.classify()`, `models.generateImages()`; the last two as nested calls of internal tools, at most four at once), and the tool declares the Lark grammar of its source for models that support grammar-constrained sampling. A failed call rejects in the script with its text, except an error result of a schema tool with its own value, which resolves: `tools.bash` gives `{ output, exitCode, ... }` for a nonzero exit. Scripts keep `store()`/`load()` values per conversation, forked as of the fork's entry. The result's `details.calls` list the script's tool calls with their status; live progress is in `pi.live.nestedTools` (`npm run bench:codemode`).
- `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.

View file

@ -180,7 +180,7 @@ registry.install(defineExtension({ name: "count", tools: [count] }));
Each call runs as its own durable task. Its intent is committed before `execute()` runs. If the process dies mid-call, the tool reruns on reopen only when it is declared `replay: "safe"`; otherwise the model gets an `interrupted` error result with the output committed so far. Throwing from `execute()` gives the model an error result. A result can also return `usage`, which is added to the conversation's [usage](#usage-and-cost). It can also return `control: { terminate: true }`: when every result of the round asks for it, the run ends without another model request.
A result has three channels: `output` for the model, `structuredOutput` for programs that call the tool (see [Calling tools from tools](#calling-tools-from-tools)), and `details` for UIs. A tool that declares `structuredOutputSchema` returns a matching `structuredOutput`, which the Harness validates. Programs calling a tool without one get its bounded output: one text item as a string, one image as its `ImageContent`, nothing as `""`, and anything else as the content list, for errors too. They get the same for a schema tool's error result without `structuredOutput`, so they still see the text of a returned error; a thrown error's message is in its `tool_error` diagnostic. The transcript stores only `output` and `details`:
A result has three channels: `output` for the model, `structuredOutput` for programs that call the tool (see [Calling tools from tools](#calling-tools-from-tools)), and `details` for UIs. A tool that declares `structuredOutputSchema` returns a matching `structuredOutput`, which the Harness validates. Programs calling a tool without one get its bounded output: one text item as a string, one image as its `ImageContent`, nothing as `""`, and anything else as the content list, for errors too. `structuredOutput` is always the tool's own value: a schema tool's error result keeps the value it returned, and the text of one without is in a `tool_error` diagnostic, as a thrown error's message is. The transcript stores only `output` and `details`:
```typescript
const lines = defineTool({
@ -196,7 +196,7 @@ const lines = defineTool({
});
```
A missing, undeclared, or invalid `structuredOutput` turns a nested call's result into an error result with an `invalid_structured_output` diagnostic; for a call the model made, which never sees it, it is dropped and reported through `onReport`. An error result may omit it, and programs then get its bounded output; error results the Harness writes itself (blocked, invalid, interrupted, aborted) carry none, and their `diagnostics` say what went wrong. An `afterTool` hook that redacts a schema tool's `output` must redact its `structuredOutput` too. The built-in `bash` returns `{ output, truncated, fullOutputPath?, exitCode }`, where `output` is the retained tail the model sees, read with `api.retainedOutput()`, and answers a nonzero exit with an error result that still carries it.
A missing, undeclared, or invalid `structuredOutput` turns a nested call's result into an error result with an `invalid_structured_output` diagnostic; for a call the model made, which never sees it, it is dropped and reported through `onReport`. An error result may omit it, and programs then find its text in a `tool_error` diagnostic; error results the Harness writes itself (blocked, invalid, interrupted, aborted) carry none, and their `diagnostics` say what went wrong. An `afterTool` hook that redacts a schema tool's `output` must redact its `structuredOutput` too. The built-in `bash` returns `{ output, truncated, fullOutputPath?, exitCode }`, where `output` is the retained tail the model sees, read with `api.retainedOutput()`, and answers a nonzero exit with an error result that still carries it.
A later extension's tool with the same name replaces an earlier one where both are selected, and `wrapTool()` decorates whichever tool won:
@ -265,7 +265,7 @@ import { createCodemodeExtension } from "@earendil-works/pi-durable/codemode";
registry.install(createCodemodeExtension({ sandbox: new CodemodeSandbox({ memoryLimitBytes: 256 * 1024 * 1024 }) }));
```
On Cloudflare, pass a `RemoteCodemodeSandbox` instead. Closing the sandbox while a script runs (for example when shutting down) is not an outcome: the script resumes with the next process's sandbox. 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.
The extension also adds a `codemode` prompt section listing the tools scripts can call, as TypeScript declarations within a token budget; scripts find the others with `searchTools()`, `describeTool()`, and `describeNamespace()`. Group tools by name prefix with `namespaces: [{ name, description, prefix }]`, for example an MCP server's `mcp__github__`. The result reads like the coding agent's: a header, the script's output, and on failure the calls already made; text past `max_output_tokens` is cut, with the full text in a temp file. Scripts also get `models.*`: the model catalog, `models.classify()`, and `models.generateImages()`, which run as nested calls of internal tools (`models: false` turns them off). On Cloudflare, pass a `RemoteCodemodeSandbox` instead. Closing the sandbox while a script runs (for example when shutting down) is not an outcome: the script resumes with the next process's sandbox. A failed nested call rejects in the script with its error text, unless the tool declares a schema and returned its own value with the error: `tools.bash` resolves with `{ output, exitCode, ... }` also for a nonzero exit, so scripts check `exitCode`. `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.

View file

@ -2856,7 +2856,7 @@ interface ToolExecutionApi<TDetails extends JsonValue = JsonValue> extends Docum
type NestedToolExecutionResult = {
/** The nested call's tool task. */
readonly taskId: TaskId;
/** A schema tool's validated `structuredOutput`; else, and for an error result without one, its bounded output, collapsed (below). */
/** The tool's own value: a schema tool's validated `structuredOutput` (absent on an error result without one, whose text is then a `tool_error` diagnostic); a schema-less tool's bounded output, collapsed (below). */
readonly structuredOutput?: JsonValue;
readonly isError: boolean;
readonly details?: JsonValue;
@ -2984,19 +2984,20 @@ on an error result, where it may omit it; the Harness validates it after the
programs get its bounded output, the content the model sees without the
rendered diagnostics, for a success and for an error alike: one text item as
its string, one image as its `ImageContent`, no items as `""`, and anything
else as the content list. A schema tool's error result without
`structuredOutput` gives programs its bounded output the same way, so they get
the text of a returned error, for example `"permission denied: /etc"`. A
thrown error's message is in its `tool_error` diagnostic, not in the output,
so programs read it there, as they do for any tool. So a
hook that redacts a tool's output redacts what programs get of a schema-less
tool and of an error result without structured output. An
else as the content list. `structuredOutput` is always the tool's own value,
never text the Harness makes: a schema tool's error result keeps the value it
returned, and one without gives programs no `structuredOutput`; the text of its
output becomes a `tool_error` diagnostic of the nested result instead, unless an
error diagnostic already says what went wrong, so programs get it, for example
`"permission denied: /etc"`. A thrown error's message is in its `tool_error`
diagnostic too. So a hook that redacts a schema-less tool's output redacts
what programs get. An
error result the Harness writes itself (unavailable, invalid arguments,
blocked, interrupted, aborted, abandoned, or a scheduler fault) has no
`structuredOutput`; its diagnostics say what went wrong. A thrown error
becomes a `tool_error` diagnostic with its message. A nested call whose result breaks this gets an error result
with an `invalid_structured_output` diagnostic, without the broken value, and
with its bounded output as `structuredOutput`;
with an `invalid_structured_output` diagnostic, without the broken value: a
schema tool's then has no `structuredOutput`, a schema-less tool's its output;
a model-issued call, whose `structuredOutput` the model never sees, only loses
it, and the break is reported through `HarnessOptions.onReport`. `structuredOutput` has no size bound;
keeping it small is the tool's job, as `bash` keeps its retained tail and the
@ -3557,10 +3558,56 @@ any more (section 2.2).
`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`.
`c<id>`, the VM's own number for the call. As in the coding agent, a tool that
declares `structuredOutputSchema` resolves with its `structuredOutput`, also on an
error result that carries one (`bash` with a nonzero exit resolves with
`{ output, exitCode, ... }`, which the script checks), and a schema-less tool
resolves with its text output, images left out, as its declaration
(`Promise<string>`) says. Every other error rejects, with the schema-less tool's
text as message when it is not empty, else the error diagnostics' messages, and
the value, if any, as `error.structuredOutput`. Which tools declare a schema is
recorded in the start data, so a resumed script maps results as it began.
The tool's description is fixed: the script environment and its globals. The
extension's `codemode` prompt section lists the agent's `callable` tools as
TypeScript declarations, grouped by `namespaces` (a tool belongs to the one with
the longest matching name prefix), within an inline budget of 3000 estimated
tokens: every namespace places its cheapest tool before any places a second,
and what does not fit is said so and found with `searchTools()`. Scripts find
tools with `ALL_TOOLS`, `searchTools(query, { limit?, namespace? })` (BM25),
`describeTool(name)`, and `describeNamespace(name)`, answered from the start
data and journaled like other globals. The result's content is a
`Script completed`/`Script failed` header with the wall time since
`startedAt`, the text items (`==> text N/M <==` when there are several), the
`console` lines in one `<console_output>` block, the returned value, and on
failure the error with the calls already made. Text past `max_output_tokens`
(default 10000) keeps its start and end, the complete text written to a file
(`details.fullOutputPath`); images are saved to files, their paths next to them.
The Harness cuts a result past the tool's `outputLimits` (4 MiB, 100,000 lines)
without such a file, so the text is also cut when it exceeds 1 Mi characters or
the line limit, whatever `max_output_tokens` says. The files go to a temp
directory of the environment that the run creates once and records in its run
document (`outputDir`) before writing to it, under fixed names (`output.txt`,
`image-<n>.<ext>`): a rerun after a crash overwrites them instead of adding new
ones.
With `models` (default on), scripts get the model catalog
(`models.getModelsOfType`, `getAvailableOfType`, `getModelOfType`; journaled
globals, configured `headers` removed) and `models.classify()` and
`models.generateImages()`. Those two are paid effects, so each runs as an
internal tool of the extension, `pi.codemode.models.classify` and
`pi.codemode.models.generateImages` (`callers: ["tools"]`): the code mode tool
turns the script's call of the global into a nested call of the tool, which then
has its own task, live row, abort, usage accounting, and replay policy. Tools
named `pi.codemode.*` are recorded in the start data as `hidden`: never in the
section, `tools`, `ALL_TOOLS`, or tool search. They declare a schema, so like any
such tool they resolve with the model's result whatever its `stopReason`, which
the script checks; their rows show a failed one as an error. At most four such
calls of one script run at once, and a call waiting for a slot when the script
has ended is not started. A script that
received generated images and showed none gets a note in its result. The tool's
declaration carries the Lark grammar of the script source as
`constrainedSampling`, which pi-ai applies for models that support it.
The result's \`details\` list the script's tool calls in call order
(\`CodemodeToolDetails\`: id, name, cut arguments, status \`ok\`/\`error\`/

View file

@ -35,7 +35,7 @@
* 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 { TextContent } from "@earendil-works/pi-ai";
import {
type CodemodeHostCall,
type CodemodeOutcome,
@ -44,11 +44,12 @@ import {
type CodemodeSandboxBase,
toCodemodeIdentifier,
} from "@earendil-works/pi-codemode/portable";
import { parseCodemodeSource } from "@earendil-works/pi-codemode/source";
import { CODEMODE_SOURCE_GRAMMAR, parseCodemodeSource } from "@earendil-works/pi-codemode/source";
import { type Static, Type } from "typebox";
import { defineDoc } from "../documents.ts";
import type { ExecutionEnv } from "../env/index.ts";
import { errorMessage } from "../errors.ts";
import { defineExtension, defineTool } from "../harness/define.ts";
import { defineExtension, defineTool, section } from "../harness/define.ts";
import { NestedCallDoc, type ToolTaskInput } from "../harness/tool.ts";
import type {
Extension,
@ -58,17 +59,41 @@ import type {
ToolRegistration,
} from "../harness/types.ts";
import type { JsonObject } from "../types.ts";
import {
type CodemodeNamespace,
codemodeDescription,
codemodeSection,
DEFAULT_CODEMODE_INLINE_BUDGET,
namespaceOf,
sampleOf,
} from "./describe.ts";
import {
CATALOG_GLOBALS,
CODEMODE_INTERNAL_PREFIX,
catalogCall,
createModelTools,
GENERATE_IMAGES_TOOL,
MAX_CONCURRENT_MODEL_CALLS,
MODEL_CALLS,
} from "./models.ts";
import { DEFAULT_MAX_OUTPUT_TOKENS, formatResult, type OutputFile } from "./output.ts";
import { rankDocuments } from "./search.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_MODULE_VERSION = 2;
export const CODEMODE_TOOL_NAME = "codemode";
/** One tool as the script sees it: `tools.<jsName>`. */
type DeclaredTool = { name: string; jsName: string; description: string };
/**
* `structured`: the tool declares `structuredOutputSchema`, so an error result with a value resolves with it.
* `hidden`: an internal tool of this extension (`pi.codemode.*`), which scripts reach only through a global such as
* `models.classify()`; never in `tools`, `ALL_TOOLS`, or tool search.
*/
type DeclaredTool = { name: string; jsName: string; description: string; structured?: true; hidden?: true };
/** What a call starts with; a rerun starts the VM from exactly this. */
type Start = {
@ -80,12 +105,18 @@ type Start = {
/** `startedAt + timeout_ms` when the script set one. */
deadline?: number;
tools: DeclaredTool[];
/** The namespaces of the tools, as `describeNamespace()` and `searchTools()` see them. */
namespaces: { name: string; description?: string; instructions?: string; tools: string[] }[];
globals: { name: string; spread: boolean }[];
/** `max_output_tokens` of the script's options line, or the default. */
maxOutputTokens: number;
store: Record<string, JsonValue>;
};
type RunState = {
start?: Start;
/** Directory of the result's files (the full output, images), created once: a rerun writes the same files again. */
outputDir?: string;
/** The outcome, decided once. */
result?: JsonObject;
};
@ -136,7 +167,7 @@ export const CodemodeStoreDoc = defineDoc<{ values: Record<string, JsonValue> }>
/** One tool call a script made, as the code mode call's details show it. */
export type CodemodeNestedCall = {
/** The VM's number for the call; its nested call ID is \`<callId>/c<id>\`. */
/** The VM's number for the call; its nested call ID is `<callId>/c<id>`. */
id: number;
name: string;
/** Compact JSON of the arguments, cut for display. */
@ -149,10 +180,14 @@ export type CodemodeNestedCall = {
/**
* The code mode call's details, in its result: its tool calls in the order it made them, with their status. Calls to
* globals are not listed. While the script runs, UIs show its nested calls from \`pi.live.nestedTools\`, so the list is
* globals are not listed. While the script runs, UIs show its nested calls from `pi.live.nestedTools`, so the list is
* written once, with the result.
*/
export type CodemodeToolDetails = { calls: CodemodeNestedCall[] };
export type CodemodeToolDetails = {
calls: CodemodeNestedCall[];
/** Temp file with the complete text output, when the output was cut to its token budget. */
fullOutputPath?: string;
};
const ARGS_PREVIEW_CHARS = 200;
const ERROR_PREVIEW_CHARS = 500;
@ -175,38 +210,67 @@ export type CodemodeToolInput = Static<typeof codemodeSchema>;
export type CodemodeToolOptions = {
/** Runs the scripts: `CodemodeSandbox` on Node, `RemoteCodemodeSandbox` on Cloudflare. */
readonly sandbox: CodemodeSandboxBase;
/** Groups of tools by name prefix, such as MCP servers, for the `codemode` section and tool search. */
readonly namespaces?: readonly CodemodeNamespace[];
/** Estimated tokens the tool declarations in the `codemode` section may use. Default 3000. */
readonly inlineBudget?: number;
/** Give scripts `models.*`: the model catalog, classifiers, and image generation. Default true. */
readonly models?: boolean;
/** @internal Host functions scripts call as globals. */
readonly globals?: readonly CodemodeGlobal[];
};
/** Outputs bounded by the script's own token budget; the Harness's limits only stop a runaway. */
const OUTPUT_LIMITS = { maxBytes: 4 * 1024 * 1024, maxLines: 100_000 } as const;
/** Globals for finding tools, answered from the start data, so a replay sees what the first run saw. */
const DISCOVERY_GLOBALS = ["searchTools", "describeTool", "describeNamespace"] as const;
const DEFAULT_SEARCH_LIMIT = 8;
/** 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.",
description: codemodeDescription(options.models !== false),
parameters: codemodeSchema,
outputLimits: OUTPUT_LIMITS,
// Capable models write the script as raw text instead of a JSON-escaped string; pi-ai applies it only for them.
constrainedSampling: { type: "grammar", variants: { openai_lark: CODEMODE_SOURCE_GRAMMAR } },
callers: ["model"],
replay: "safe",
execute: (args, api, context) => new CodemodeCall(options, api, context).execute(args.code),
});
}
/** The `codemode` extension: the code mode tool. */
/**
* The `codemode` extension: the code mode tool, and the `codemode` prompt section listing the tools scripts can call.
*/
export function createCodemodeExtension(options: CodemodeToolOptions): Extension {
return defineExtension({ name: "codemode", tools: [createCodemodeTool(options)] });
return defineExtension({
name: "codemode",
tools: [createCodemodeTool(options), ...(options.models === false ? [] : createModelTools())],
sections: [
section("codemode", (input) =>
codemodeSection(
input.agent.callable.filter((tool) => !tool.name.startsWith(CODEMODE_INTERNAL_PREFIX)),
options.namespaces ?? [],
options.inlineBudget ?? DEFAULT_CODEMODE_INLINE_BUDGET,
),
),
],
});
}
/** 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")
: "";
/**
* What a script gets for a nested call, as in the coding agent. A tool with a schema resolves with its value, also on
* an error result that carries one (`bash`'s `{ output, exitCode }` for a nonzero exit); a schema-less tool
* resolves with its text output, images left out, as its declaration says. Any other error rejects, with its text: a
* schema-less tool's output, else its error diagnostics.
*/
function outcomeOf(name: string, structured: boolean, result: NestedToolExecutionResult): CodemodeOutcome {
const value = structured ? result.structuredOutput : textOf(result.structuredOutput);
if (!result.isError || (structured && value !== undefined)) return { ok: true, value };
const message = structured ? "" : (value as string);
const diagnostics = result.diagnostics
.filter((diagnostic) => diagnostic.severity === "error")
.map((diagnostic) => diagnostic.message)
@ -218,25 +282,15 @@ function outcomeOf(name: string, result: NestedToolExecutionResult): CodemodeOut
};
}
function isText(item: unknown): item is TextContent {
return typeof item === "object" && item !== null && (item as TextContent).type === "text";
/** The text of a schema-less tool's collapsed output: a string, one image, or an array of text and image blocks. */
function textOf(value: JsonValue | undefined): string {
if (typeof value === "string") return value;
if (!Array.isArray(value)) return "";
return value.flatMap((item) => (isText(item) ? [item.text] : [])).join("\n");
}
/** 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 isText(item: unknown): item is TextContent {
return typeof item === "object" && item !== null && (item as TextContent).type === "text";
}
function errorResult(message: string): ToolExecutionResult {
@ -284,7 +338,12 @@ class CodemodeCall {
/** Set once the call decided an outcome other than the script's own (divergence, mismatch) or failed. */
private ended: ToolExecutionResult | undefined;
private failure: unknown;
/** The script's tool calls for \`details\`, by VM call ID. */
/** `models.classify()`/`generateImages()` calls in flight, and those waiting for a slot. */
private modelCalls = 0;
private readonly modelWaiters: (() => void)[] = [];
/** Images `models.generateImages()` returned, to notice a script that never shows them. */
private generatedImages = 0;
/** The script's tool calls for `details`, by VM call ID. */
private readonly rows = new Map<number, CodemodeNestedCall>();
/** Resolves once the VM's fingerprint is checked or recorded; nothing is delivered or admitted before. */
private gate!: Promise<void>;
@ -306,12 +365,33 @@ class CodemodeCall {
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,
description: sampleOf(tool),
...(tool.structuredOutputSchema === undefined ? {} : { structured: true as const }),
...(tool.name.startsWith(CODEMODE_INTERNAL_PREFIX) ? { hidden: true as const } : {}),
}));
const callable = agent.callable.filter((tool) => !tool.name.startsWith(CODEMODE_INTERNAL_PREFIX));
const models = this.options.models === false ? [] : [...CATALOG_GLOBALS, ...Object.keys(MODEL_CALLS)];
const globals = [
...DISCOVERY_GLOBALS.map((name) => ({ name, spread: true })),
...models.map((name) => ({ name, spread: true })),
...(this.options.globals ?? []).map((global) => ({ name: global.name, spread: global.spread === true })),
];
const namespaces: Start["namespaces"] = [];
for (const tool of callable) {
const namespace = namespaceOf(tool.name, this.options.namespaces ?? []);
if (namespace === undefined) continue;
let entry = namespaces.find((candidate) => candidate.name === namespace.name);
if (entry === undefined) {
entry = {
name: namespace.name,
...(namespace.description === undefined ? {} : { description: namespace.description }),
...(namespace.instructions === undefined ? {} : { instructions: namespace.instructions }),
tools: [],
};
namespaces.push(entry);
}
entry.tools.push(tool.name);
}
const now = Date.now();
const seed = crypto.getRandomValues(new Uint32Array(1))[0]!;
type Found =
@ -330,7 +410,9 @@ class CodemodeCall {
seed,
...(timeout === undefined ? {} : { deadline: now + timeout }),
tools,
namespaces,
globals,
maxOutputTokens: sourceOptions.maxOutputTokens ?? DEFAULT_MAX_OUTPUT_TOKENS,
store: copyJson(store.values as JsonValue) as Record<string, JsonValue>,
};
}
@ -375,7 +457,7 @@ class CodemodeCall {
this.run = this.options.sandbox.start(code, {
startedAt: this.start.startedAt,
seed: this.start.seed,
tools: this.start.tools,
tools: this.start.tools.filter((tool) => tool.hidden !== true),
globals: this.start.globals,
store: this.start.store,
signal: this.controller.signal,
@ -458,7 +540,23 @@ class CodemodeCall {
/** The script's own outcome, stored with its store writes in one transaction. */
private async complete(result: CodemodeResult): Promise<ToolExecutionResult> {
const tool = { ...toolResult(result), details: this.details() as unknown as JsonValue };
const details = this.details();
const formatted = await formatResult(result, {
generatedImages: this.generatedImages,
wallMs: Date.now() - this.start.startedAt,
maxOutputTokens: this.start.maxOutputTokens,
calls: details.calls,
limits: OUTPUT_LIMITS,
write: this.outputFiles(),
});
const tool: ToolExecutionResult = {
output: formatted.content,
isError: !result.ok,
details: {
...details,
...(formatted.fullOutputPath === undefined ? {} : { fullOutputPath: formatted.fullOutputPath }),
} as unknown as JsonValue,
};
const writes = result.ok ? result.storeWrites : undefined;
return this.api.commit(async (tx) => {
const run = await tx.doc(CodemodeRunDoc, this.api.taskId);
@ -473,6 +571,39 @@ class CodemodeCall {
}, this.context);
}
/** Writes the result's files under fixed names into the run's output directory, so a rerun overwrites them. */
private outputFiles(): OutputFile {
const env = this.api.env;
let directory: Promise<string> | undefined;
return async (name, data) => {
if (env === undefined) return { error: "no execution environment" };
try {
directory ??= this.outputDir(env);
const path = `${await directory}/${name}`;
const written = await env.writeFile(path, data, this.context);
return written.ok ? { path } : { error: written.error.message };
} catch (error) {
return { error: errorMessage(error) };
}
};
}
/** The run's output directory: the recorded one, else a new temp directory, recorded before any file is written. */
private async outputDir(env: ExecutionEnv): Promise<string> {
const recorded = await this.api.commit(
async (tx) => (await tx.doc(CodemodeRunDoc, this.api.taskId)).outputDir,
this.context,
);
if (recorded !== undefined) return recorded;
const created = await env.createTempDir("pi-codemode-", this.context);
if (!created.ok) throw new Error(created.error.message);
return this.api.commit(async (tx) => {
const run = await tx.doc(CodemodeRunDoc, this.api.taskId);
run.outputDir ??= created.value;
return run.outputDir!;
}, this.context);
}
private async checkFingerprint(fingerprint: string): Promise<void> {
const recorded = this.start.fingerprint;
if (recorded !== undefined) {
@ -502,8 +633,9 @@ class CodemodeCall {
else if (!this.rerun) void this.release();
}
private onCall(call: CodemodeHostCall): void {
private onCall(made: CodemodeHostCall): void {
if (this.stopped) return;
const call = asToolCall(made);
if (call.target === "tool") this.started(call);
this.issued = Math.max(this.issued, call.id);
const index = this.replaying ? (this.replayIndex.get(call.id) ?? -1) : -1;
@ -536,6 +668,7 @@ class CodemodeCall {
this.deliverReplayed();
},
(error: unknown) => {
if (this.stopped) return;
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)}`);
@ -650,11 +783,10 @@ class CodemodeCall {
/** 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) => ({
: this.callGlobal(call).then((outcome) => ({
outcome,
global: {
name: call.name,
@ -668,6 +800,8 @@ class CodemodeCall {
void this.pump();
},
(error: unknown) => {
// Stopped meanwhile: the outcome is decided, or a waiting model call gave up.
if (this.stopped) return;
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)}`);
@ -678,9 +812,25 @@ class CodemodeCall {
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}` });
const outcome = outcomeOf(call.name, result);
this.finishRow(call.id, outcome, result.durationMs);
const model = call.name.startsWith(CODEMODE_INTERNAL_PREFIX);
if (model) {
await this.modelSlot();
// The call may have stopped while this one waited: a finished script admits no new paid call.
if (this.stopped) {
this.releaseModelSlot();
throw new Error("The script has ended");
}
}
let result: NestedToolExecutionResult;
try {
result = await this.api.executeTool(call.name, args, this.context, { key: `c${call.id}` });
} finally {
if (model) this.releaseModelSlot();
}
if (call.name === GENERATE_IMAGES_TOOL) this.generatedImages += countImages(result.structuredOutput);
const structured = this.start.tools.find((tool) => tool.name === call.name)?.structured === true;
const outcome = outcomeOf(call.name, structured, result);
this.finishRow(call.id, outcome, result);
return outcome;
}
@ -688,16 +838,22 @@ class CodemodeCall {
private started(call: CodemodeHostCall): void {
if (this.rows.has(call.id)) return;
const args = call.args === undefined ? "" : (JSON.stringify(call.args) ?? "");
this.rows.set(call.id, { id: call.id, name: call.name, args: cut(args, ARGS_PREVIEW_CHARS), status: "running" });
const name = call.name.startsWith(CODEMODE_INTERNAL_PREFIX)
? call.name.slice(CODEMODE_INTERNAL_PREFIX.length)
: call.name;
this.rows.set(call.id, { id: call.id, name, args: cut(args, ARGS_PREVIEW_CHARS), status: "running" });
}
/** A tool call ended: its row gets its status, duration, and error. */
private finishRow(id: number, outcome: CodemodeOutcome, durationMs: number | undefined): void {
/** The call ended: an error result is an error row, also when the script got its value (bash's nonzero exit). */
private finishRow(id: number, outcome: CodemodeOutcome, result: NestedToolExecutionResult): void {
const row = this.rows.get(id);
if (row === undefined) return;
row.status = outcome.ok ? "ok" : "error";
if (durationMs !== undefined) row.durationMs = durationMs;
if (!outcome.ok) row.error = cut(outcome.message, ERROR_PREVIEW_CHARS);
row.status = result.isError ? "error" : "ok";
if (result.durationMs !== undefined) row.durationMs = result.durationMs;
const errors = result.diagnostics.filter((diagnostic) => diagnostic.severity === "error");
const error = !outcome.ok ? outcome.message : errors.map((diagnostic) => diagnostic.message).join("\n");
if (result.isError && error !== "") row.error = cut(error, ERROR_PREVIEW_CHARS);
}
/** The rows in call order; calls still running when the script ended were cancelled. */
@ -708,17 +864,91 @@ class CodemodeCall {
return { calls };
}
private async callGlobal(call: CodemodeHostCall, global: CodemodeGlobal | undefined): Promise<CodemodeOutcome> {
if (global === undefined) return { ok: false, message: `Unknown global ${call.name}` };
private async callGlobal(call: CodemodeHostCall): Promise<CodemodeOutcome> {
const discovery = (DISCOVERY_GLOBALS as readonly string[]).includes(call.name);
const catalog = (CATALOG_GLOBALS as readonly string[]).includes(call.name);
const global =
discovery || catalog ? undefined : this.options.globals?.find((candidate) => candidate.name === call.name);
if (!discovery && !catalog && 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) };
const value = discovery
? this.discover(call)
: catalog
? await catalogCall(call.name, call.args, this.api.models, this.context)
: await global!.execute(call.args, this.api, this.context);
return { ok: true, value };
} catch (error) {
if (this.context.abortSignal?.aborted) throw error;
return { ok: false, message: errorMessage(error) };
}
}
private async modelSlot(): Promise<void> {
if (this.modelCalls >= MAX_CONCURRENT_MODEL_CALLS) {
await new Promise<void>((resolve) => this.modelWaiters.push(resolve));
}
this.modelCalls++;
}
private releaseModelSlot(): void {
this.modelCalls--;
this.modelWaiters.shift()?.();
}
/** `searchTools()`, `describeTool()`, `describeNamespace()`: answered from the start data. */
private discover(call: CodemodeHostCall): JsonValue | undefined {
const args = Array.isArray(call.args) ? call.args : [];
const { namespaces } = this.start;
const tools = this.start.tools.filter((tool) => tool.hidden !== true);
const namespaceOfTool = (name: string) => namespaces.find((namespace) => namespace.tools.includes(name));
const matchesNamespace = (namespace: Start["namespaces"][number], query: string) =>
namespace.name === query || toCodemodeIdentifier(namespace.name) === toCodemodeIdentifier(query);
if (call.name === "describeTool") {
const [name] = args;
if (typeof name !== "string") throw new Error("describeTool() expects a tool name");
return tools.find((tool) => tool.name === name || tool.jsName === name)?.description;
}
if (call.name === "describeNamespace") {
const [name] = args;
if (typeof name !== "string") throw new Error("describeNamespace() expects a namespace name");
const namespace = namespaces.find((candidate) => matchesNamespace(candidate, name));
if (namespace === undefined) return undefined;
return {
name: namespace.name,
...(namespace.description === undefined ? {} : { description: namespace.description }),
...(namespace.instructions === undefined ? {} : { instructions: namespace.instructions }),
tools: namespace.tools.map((tool) => toCodemodeIdentifier(tool)),
};
}
const [query, options] = args as [unknown, { limit?: unknown; namespace?: unknown } | undefined];
if (typeof query !== "string") throw new Error("searchTools() expects a query string");
const limit = options?.limit ?? DEFAULT_SEARCH_LIMIT;
if (typeof limit !== "number" || !Number.isInteger(limit) || limit <= 0) {
throw new Error("searchTools() limit must be a positive integer");
}
const namespace = options?.namespace;
if (namespace !== undefined && namespace !== null && typeof namespace !== "string") {
throw new Error("searchTools() namespace must be a string");
}
const documents = tools.flatMap((tool) => {
const own = namespaceOfTool(tool.name);
if (typeof namespace === "string" && (own === undefined || !matchesNamespace(own, namespace))) return [];
const text = [
tool.name,
tool.name.replaceAll("_", " "),
tool.description,
own?.name ?? "",
own?.description ?? "",
];
return [{ name: tool.name, text: text.join(" ") }];
});
return rankDocuments(query, documents, limit).map((name) => {
const tool = tools.find((candidate) => candidate.name === name)!;
return { name: tool.jsName, description: tool.description };
});
}
/**
* 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.
@ -766,6 +996,20 @@ class CodemodeCall {
}
}
/** A script's call of `models.classify()` or `models.generateImages()` as the call of its internal tool. */
function asToolCall(call: CodemodeHostCall): CodemodeHostCall {
const tool = call.target === "global" ? MODEL_CALLS[call.name] : undefined;
if (tool === undefined) return call;
const [model, context] = Array.isArray(call.args) ? call.args : [];
return { id: call.id, target: "tool", name: tool, args: { model: model ?? null, context: context ?? null } };
}
/** Image blocks in a `models.generateImages()` result. */
function countImages(value: unknown): number {
const output = typeof value === "object" && value !== null ? (value as { output?: unknown }).output : undefined;
return Array.isArray(output) ? output.filter((item) => (item as { type?: unknown })?.type === "image").length : 0;
}
function describe(call: CodemodeHostCall): string {
return `${call.target === "tool" ? "tools." : ""}${call.name}(${JSON.stringify(call.args) ?? ""})`;
}

View file

@ -0,0 +1,171 @@
/**
* What the model reads about code mode, copied from the coding agent's code mode: the tool's fixed description (the
* script environment and its globals), and the `codemode` prompt section, which lists the tools scripts can call as
* TypeScript declarations, grouped by namespace and limited to a token budget. Tools that do not fit are found with
* `searchTools()`. The section renders from the conversation's callable tools, so it changes with them; the tool's
* description never does.
*/
import type { CodemodeJsonSchema } from "@earendil-works/pi-codemode";
import {
MCP_TYPESCRIPT_PREAMBLE,
mcpStructuredContentSchema,
renderToolSample,
toCodemodeIdentifier,
} from "@earendil-works/pi-codemode/declarations";
import type { ToolRegistration } from "../harness/types.ts";
import { MODELS_DESCRIPTION } from "./models.ts";
/** A group of tools named by a common prefix, such as an MCP server's `mcp__github__`. */
export type CodemodeNamespace = {
readonly name: string;
readonly description?: string;
readonly instructions?: string;
/** Tool names starting with it belong to the namespace; the longest matching prefix wins. */
readonly prefix: string;
};
/** Default token budget of the tool declarations in the `codemode` section, as in the coding agent. */
export const DEFAULT_CODEMODE_INLINE_BUDGET = 3000;
/** Characters per token when estimating the cost of a tool section. */
const CHARS_PER_TOKEN = 4;
const TEXT_OUTPUT_SCHEMA: CodemodeJsonSchema = { type: "string" };
const DESCRIPTION_INTRO = `Run JavaScript that calls other tools. The input is raw JavaScript (not JSON, no code fence), 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 a string, or an object if the tool's declaration says so, and rejects with an Error on failure. A tool whose declaration returns an object resolves on failure too when it returned that object; check its fields. Calls still running when the script ends are cancelled.
- The script survives a restart: it resumes where it was, without running finished calls again.
- Optional first line: \`// @options: {"max_output_tokens": 10000, "timeout_ms": 60000}\``;
/** The code mode tool's description: the script environment and its globals, the same for every conversation. */
export function codemodeDescription(models: boolean): string {
const globals = [
"Globals:",
"- `text(value)`, `image(dataUrlOrImageBlock)`, `console.log(...)`, and top-level `return` add output; `exit()` ends the script. With several text items, each starts with a `==> text N/M <==` line, and `console` lines follow the other output in one `<console_output>` block. `image()` also saves the image to a temp file and the result names its path.",
"- `store(key, value)` and `load(key)` keep JSON values across codemode calls.",
"- `ALL_TOOLS`, `await searchTools(query, { limit?, namespace? })`, `await describeTool(name)`, `await describeNamespace(name)`: find tools the `codemode` section does not list.",
];
return [
DESCRIPTION_INTRO,
globals.join("\n"),
...(models ? [MODELS_DESCRIPTION] : []),
"The `codemode` section of the system prompt lists the tools scripts can call, by namespace, as far as space allows.",
].join("\n\n");
}
/** What a script sees of a tool: its declaration. Tools without an output schema resolve to their text output. */
export function declarationOf(tool: ToolRegistration): {
name: string;
description: string;
inputSchema: CodemodeJsonSchema;
outputSchema: CodemodeJsonSchema;
} {
return {
name: tool.name,
description: tool.description,
inputSchema: tool.parameters as CodemodeJsonSchema,
outputSchema: (tool.structuredOutputSchema as CodemodeJsonSchema | undefined) ?? TEXT_OUTPUT_SCHEMA,
};
}
/** `toolSample` of a tool: its description and TypeScript declaration. */
export function sampleOf(tool: ToolRegistration): string {
return renderToolSample(declarationOf(tool));
}
/** The namespace of a tool: the one with the longest prefix its name starts with. */
export function namespaceOf(name: string, namespaces: readonly CodemodeNamespace[]): CodemodeNamespace | undefined {
let found: CodemodeNamespace | undefined;
for (const namespace of namespaces) {
if (name.startsWith(namespace.prefix) && (found === undefined || namespace.prefix.length > found.prefix.length)) {
found = namespace;
}
}
return found;
}
/** `### \`id\` (\`raw name\`)` followed by the tool's description and declaration. */
function renderToolSection(tool: ToolRegistration): string {
const id = toCodemodeIdentifier(tool.name);
const heading = id === tool.name ? `### \`${id}\`` : `### \`${id}\` (\`${tool.name}\`)`;
return `${heading}\n${sampleOf(tool).trim()}`;
}
type CatalogEntry = { name: string; section: string; cost: number };
type CatalogGroup = { namespace: CodemodeNamespace | undefined; entries: CatalogEntry[] };
/**
* Pick the tool sections that fit the budget: in each round every group (tools without a namespace first, then
* namespaces by name) places its cheapest remaining tool; a group whose next tool does not fit drops out while the
* others continue. Every namespace is represented before any namespace is complete.
*/
function selectCatalog(groups: readonly CatalogGroup[], budget: number): Set<string> {
const queues = groups.map((group) => [...group.entries].sort((a, b) => a.cost - b.cost));
const shown = new Set<string>();
let remaining = budget;
let active = queues.filter((queue) => queue.length > 0);
while (active.length > 0) {
active = active.filter((queue) => {
const next = queue[0]!;
if (next.cost > remaining) return false;
remaining -= next.cost;
shown.add(next.name);
queue.shift();
return queue.length > 0;
});
}
return shown;
}
/**
* The `codemode` section: the tools scripts can call, grouped by namespace, with their declarations as far as the
* budget allows, and the shared MCP types when a listed tool needs them. `undefined` without callable tools.
*/
export function codemodeSection(
tools: readonly ToolRegistration[],
namespaces: readonly CodemodeNamespace[],
budget: number,
): string | undefined {
if (tools.length === 0) return undefined;
const groups = new Map<string, CatalogGroup>([["", { namespace: undefined, entries: [] }]]);
for (const tool of tools) {
const namespace = namespaceOf(tool.name, namespaces);
const key = namespace ? `ns:${namespace.name}` : "";
const group = groups.get(key) ?? { namespace, entries: [] };
groups.set(key, group);
const section = renderToolSection(tool);
group.entries.push({ name: tool.name, section, cost: Math.ceil(section.length / CHARS_PER_TOKEN) });
}
const ordered = [...groups.values()].sort((a, b) =>
a.namespace === undefined ? -1 : b.namespace === undefined ? 1 : a.namespace.name.localeCompare(b.namespace.name),
);
const shown = selectCatalog(ordered, budget);
const sections: string[] = [];
if (
tools.some(
(tool) => shown.has(tool.name) && mcpStructuredContentSchema(declarationOf(tool).outputSchema) !== undefined,
)
) {
sections.push(`Shared MCP Types:\n\`\`\`ts\n${MCP_TYPESCRIPT_PREAMBLE}\n\`\`\``);
}
const toolSections = ["Tools scripts can call (`await tools.<name>(args)`):"];
for (const { namespace, entries } of ordered) {
if (entries.length === 0) continue;
const visible = entries.filter((entry) => shown.has(entry.name));
if (namespace) {
const listing =
visible.length === entries.length
? ""
: visible.length === 0
? " (tools not listed)"
: " (some tools not listed)";
const description = namespace.description?.trim();
toolSections.push(`## ${namespace.name}${listing}${description ? `\n${description}` : ""}`);
} else if (visible.length < entries.length) {
toolSections.push("Some tools are not listed: find them with `searchTools()`.");
}
for (const entry of visible) toolSections.push(entry.section);
}
sections.push(toolSections.join("\n\n"));
return sections.join("\n\n");
}

View file

@ -10,3 +10,4 @@ export {
createCodemodeExtension,
createCodemodeTool,
} from "./codemode.ts";
export { type CodemodeNamespace, DEFAULT_CODEMODE_INLINE_BUDGET } from "./describe.ts";

View file

@ -0,0 +1,285 @@
/**
* `models.*` for code mode scripts, copied from the coding agent's code mode: the model catalog, and classifier and
* image model calls.
*
* The catalog functions are host globals, journaled like any global. `models.classify()` and `models.generateImages()`
* are paid effects, so each runs as an internal tool (`callers: ["tools"]`, a reserved name, hidden from scripts'
* tool listings): a nested durable call with its own task, live row, abort, usage accounting, and replay policy. The
* code mode tool turns a script's call of the global into a call of that tool.
*/
import type { Context, JsonValue } from "@earendil-works/chord";
import type {
AnyModel,
ClassifierContext,
ImagesContext,
Models,
ModelType,
ModelTypeMap,
Usage,
} from "@earendil-works/pi-ai";
import { Type } from "typebox";
import { defineTool } from "../harness/define.ts";
import type { ToolRegistration } from "../harness/types.ts";
/** Reserved names of code mode's internal tools; they are never listed to scripts or the model. */
export const CODEMODE_INTERNAL_PREFIX = "pi.codemode.";
export const CLASSIFY_TOOL = `${CODEMODE_INTERNAL_PREFIX}models.classify`;
export const GENERATE_IMAGES_TOOL = `${CODEMODE_INTERNAL_PREFIX}models.generateImages`;
/** The global a script calls, by internal tool. */
export const MODEL_CALLS: Readonly<Record<string, string>> = {
"models.classify": CLASSIFY_TOOL,
"models.generateImages": GENERATE_IMAGES_TOOL,
};
/** The catalog functions, host globals of the `models` namespace. */
export const CATALOG_GLOBALS = [
"models.getModelsOfType",
"models.getAvailableOfType",
"models.getModelOfType",
] as const;
/** `models.classify()` and `models.generateImages()` calls one script may have in flight; the rest wait. */
export const MAX_CONCURRENT_MODEL_CALLS = 4;
const MODEL_TYPES: ReadonlySet<string> = new Set<ModelType>(["chat", "image", "classifier"]);
export const MODELS_DESCRIPTION = [
"`models` (all functions return promises):",
'- `models.getModelsOfType(type, provider?)`, `models.getAvailableOfType(type, provider?)` (configured auth only), `models.getModelOfType(type, provider, id)`: the catalog; `type` is "chat", "image", or "classifier".',
'- `models.classify(model, { state: { ... }, images?: [{ type: "image", data, mimeType }], questions: { <id>: { type: "choice", instructions, criteria: { <label>: <meaning> } } | { type: "score", instructions, criteria: [<lowest level>, ..., <highest level>] } | { type: "bool", instructions, criteria: { true: <meaning>, false: <meaning> } } } })` resolves to `{ answers: { <id>: ... }, stopReason, errorMessage?, usage? }`.',
'- `models.generateImages(model, { input: [{ type: "text", text: <prompt> }, ...optional { type: "image", data, mimeType }] })` resolves to `{ output: [image blocks], stopReason, errorMessage?, usage? }`; show each image with `image(block)`.',
`At most ${MAX_CONCURRENT_MODEL_CALLS} classify/generateImages calls run at once; more wait.`,
].join("\n");
function isRecord(value: unknown): value is Record<string, unknown> {
return typeof value === "object" && value !== null && !Array.isArray(value);
}
/** A script value in an error message: `undefined`, `a string`, `an array`, or its keys (`{ prompt }`). */
function describeValue(value: unknown): string {
if (value === undefined || value === null) return String(value);
if (Array.isArray(value)) return value.length === 0 ? "an empty array" : "an array";
if (typeof value === "object") {
const keys = Object.keys(value);
if (keys.length === 0) return "{}";
return `{ ${keys.slice(0, 6).join(", ")}${keys.length > 6 ? ", ..." : ""} }`;
}
return typeof value === "string" ? "a string" : `a ${typeof value}`;
}
/** `an image`, `a classifier`. */
function withArticle(word: string): string {
return `${/^[aeiou]/.test(word) ? "an" : "a"} ${word}`;
}
function toModelType(value: unknown): ModelType {
if (typeof value === "string" && MODEL_TYPES.has(value)) return value as ModelType;
throw new Error(`Unknown model type ${JSON.stringify(value)}. Use "chat", "image", or "classifier".`);
}
function toProvider(value: unknown): string | undefined {
if (value === undefined || value === null) return undefined;
if (typeof value !== "string") throw new Error("provider must be a string");
return value;
}
/** Catalog entry for scripts. `headers` is dropped because configured headers can carry credentials. */
function toModelInfo(model: AnyModel): JsonValue {
const info: Record<string, unknown> = { ...model };
delete info.headers;
return JSON.parse(JSON.stringify(info)) as JsonValue;
}
/** The catalog functions: what a script's call of `name` with `args` (the argument list) resolves to. */
export async function catalogCall(
name: string,
args: unknown,
models: Models,
context: Context,
): Promise<JsonValue | undefined> {
const list = Array.isArray(args) ? args : [];
const [type, provider, id] = list;
if (name === "models.getModelsOfType")
return models.getModelsOfType(toModelType(type), toProvider(provider)).map(toModelInfo);
if (name === "models.getAvailableOfType") {
const available = await models.getAvailableOfType(toModelType(type), toProvider(provider), {
...(context.abortSignal === undefined ? {} : { signal: context.abortSignal }),
});
return available.map(toModelInfo);
}
if (typeof provider !== "string" || typeof id !== "string") {
throw new Error(
`models.getModelOfType(type, provider, id) expects three strings, got (${list.map(describeValue).join(", ")}). The provider and the id are separate arguments, for example models.getModelOfType("classifier", "typesafe", "jev-latest").`,
);
}
const model = models.getModelOfType(toModelType(type), provider, id);
return model === undefined ? undefined : toModelInfo(model);
}
const CLASSIFIER_CONTEXT_SHAPE =
'{ state: { ... }, images?: [{ type: "image", data: <base64>, mimeType }], questions: { <id>: { type: "choice", instructions, criteria: { <label>: <meaning> } } | { type: "score", instructions, criteria: [<lowest level>, ..., <highest level>] } | { type: "bool", instructions, criteria: { true: <meaning>, false: <meaning> } } } }';
/** Check a script's classifier context, so mistakes fail with the expected shape instead of a provider error. */
function checkClassifierContext(context: unknown): ClassifierContext {
const fail = (problem: string) =>
new Error(`models.classify() ${problem}. Expected context: ${CLASSIFIER_CONTEXT_SHAPE}.`);
if (!isRecord(context)) throw fail(`expects a context object as its second argument, got ${describeValue(context)}`);
if (!isRecord(context.state)) throw fail(`context.state must be an object, got ${describeValue(context.state)}`);
const { images } = context;
if (images !== undefined) {
if (!Array.isArray(images)) throw fail(`context.images must be an array, got ${describeValue(images)}`);
images.forEach((image: unknown, index) => {
if (
!isRecord(image) ||
image.type !== "image" ||
typeof image.data !== "string" ||
typeof image.mimeType !== "string"
) {
throw fail(`context.images[${index}] must be an image block, got ${describeValue(image)}`);
}
});
}
const { questions } = context;
if (!isRecord(questions) || Object.keys(questions).length === 0) {
throw fail(`context.questions must map question IDs to questions, got ${describeValue(questions)}`);
}
const isStrings = (values: unknown[]) => values.length > 0 && values.every((value) => typeof value === "string");
for (const [id, question] of Object.entries(questions)) {
const at = `context.questions.${id}`;
if (!isRecord(question)) throw fail(`${at} must be a question object, got ${describeValue(question)}`);
if (typeof question.instructions !== "string") throw fail(`${at}.instructions must be a string`);
const { criteria } = question;
if (question.type === "choice") {
if (!isRecord(criteria) || !isStrings(Object.values(criteria))) {
throw fail(`${at} is a "choice" question, so criteria must map each label to its meaning`);
}
} else if (question.type === "score") {
if (!Array.isArray(criteria) || !isStrings(criteria)) {
throw fail(`${at} is a "score" question, so criteria must list the levels as strings, lowest first`);
}
} else if (question.type === "bool") {
if (!isRecord(criteria) || typeof criteria.true !== "string" || typeof criteria.false !== "string") {
throw fail(`${at} is a "bool" question, so criteria must be { true: string, false: string }`);
}
} else {
throw fail(`${at}.type must be "choice", "score", or "bool", got ${JSON.stringify(question.type)}`);
}
}
return context as unknown as ClassifierContext;
}
/** Check a script's image context, so mistakes such as `{ prompt }` fail with the expected shape. */
function checkImagesContext(context: unknown): ImagesContext {
const fail = (problem: string) =>
new Error(
`models.generateImages() ${problem}. Expected context: { input: [{ type: "text", text: <prompt> }, ...optional { type: "image", data: <base64>, mimeType } references] }.`,
);
if (!isRecord(context)) throw fail(`expects a context object as its second argument, got ${describeValue(context)}`);
const { input } = context;
if (!Array.isArray(input) || input.length === 0) {
throw fail(`context.input must be a non-empty array of blocks, got ${describeValue(input)}`);
}
input.forEach((block: unknown, index) => {
if (isRecord(block) && block.type === "text" && typeof block.text === "string") return;
if (
isRecord(block) &&
block.type === "image" &&
typeof block.data === "string" &&
typeof block.mimeType === "string"
)
return;
throw fail(`context.input[${index}] must be a text or image block, got ${describeValue(block)}`);
});
return context as unknown as ImagesContext;
}
/**
* Resolve the script's model by provider and id only, so a script-supplied baseUrl or headers never receive the
* credentials, with errors that say what to do.
*/
function resolveModel<TType extends "classifier" | "image">(
name: string,
type: TType,
model: unknown,
models: Models,
): ModelTypeMap[TType] {
const listHint = `List the ${type} models you can use with models.getAvailableOfType("${type}").`;
if (!isRecord(model) || typeof model.provider !== "string" || typeof model.id !== "string") {
const undefinedHint =
model === undefined || model === null
? " models.getModelOfType() returns undefined for an unknown provider or id."
: "";
throw new Error(
`${name}() expects ${withArticle(type)} model as its first argument, got ${describeValue(model)}.${undefinedHint} ${listHint}`,
);
}
const { provider, id } = model;
const resolved = models.getModelOfType(type, provider, id);
if (resolved !== undefined) return resolved;
const actualType = [...MODEL_TYPES].find(
(other) => other !== type && models.getModelOfType(other as ModelType, provider, id) !== undefined,
);
throw new Error(
actualType
? `"${provider}/${id}" is ${withArticle(actualType)} model, not ${withArticle(type)} model. ${listHint}`
: `Unknown ${type} model "${provider}/${id}". ${listHint}`,
);
}
const ModelCallParameters = Type.Object({ model: Type.Unknown(), context: Type.Unknown() });
/** The result of a model call as the tool returns it: the whole result for the script, usage counted by the Harness. */
function modelResult(
label: string,
result: { stopReason: string; errorMessage?: string; usage?: Usage },
value: unknown,
) {
const ok = result.stopReason === "stop";
return {
output: [
{
type: "text" as const,
text: ok ? `${label} done` : `${label} ${result.stopReason}: ${result.errorMessage ?? ""}`,
},
],
structuredOutput: JSON.parse(JSON.stringify(value)) as JsonValue,
...(ok ? {} : { isError: true }),
...(result.usage === undefined ? {} : { usage: result.usage }),
};
}
/** The internal tools behind `models.classify()` and `models.generateImages()`. */
export function createModelTools(): ToolRegistration[] {
const classify = defineTool({
name: CLASSIFY_TOOL,
description: "models.classify() of code mode scripts.",
parameters: ModelCallParameters,
structuredOutputSchema: Type.Unknown(),
callers: ["tools"],
execute: async (args, api, context) => {
const model = resolveModel("models.classify", "classifier", args.model, api.models);
const checked = checkClassifierContext(args.context);
const result = await api.models.classify(model, checked, {
...(context.abortSignal === undefined ? {} : { signal: context.abortSignal }),
});
return modelResult(`classify ${model.provider}/${model.id}`, result, result);
},
});
const generateImages = defineTool({
name: GENERATE_IMAGES_TOOL,
description: "models.generateImages() of code mode scripts.",
parameters: ModelCallParameters,
structuredOutputSchema: Type.Unknown(),
callers: ["tools"],
execute: async (args, api, context) => {
const model = resolveModel("models.generateImages", "image", args.model, api.models);
const checked = checkImagesContext(args.context);
const result = await api.models.generateImages(model, checked, {
...(context.abortSignal === undefined ? {} : { signal: context.abortSignal }),
});
return modelResult(`generateImages ${model.provider}/${model.id}`, result, result);
},
});
return [classify as unknown as ToolRegistration, generateImages as unknown as ToolRegistration];
}

View file

@ -0,0 +1,190 @@
/**
* A script's result as the model reads it, copied from the coding agent's code mode: a "Script completed/failed"
* header with the wall time, the script's text items (`==> text N/M <==` when there are several), its `console` lines
* in one `<console_output>` block, the returned value, and on failure the error with the calls already made. Text past
* the token budget is cut to its start and end, with the complete text in a temp file; images are saved to temp files
* too, with their path next to them, because a script has no other way to hand their bytes on.
*/
import type { ImageContent, TextContent } from "@earendil-works/pi-ai";
import type { CodemodeOutputItem, CodemodeResult } from "@earendil-works/pi-codemode/portable";
type Item = TextContent | ImageContent;
/** Writes a file of the result under a fixed name; its path, or why it could not be written. */
export type OutputFile = (name: string, data: string | Uint8Array) => Promise<{ path: string } | { error: string }>;
/** The tool's output limits, which the Harness enforces by cutting without a spill file: results stay inside them. */
type Limits = { readonly maxBytes: number; readonly maxLines: number };
/** Lines of the result besides the script's text: the header, markers, notes. */
const LINE_MARGIN = 100;
/** Default token budget for script output. */
export const DEFAULT_MAX_OUTPUT_TOKENS = 10_000;
/** Characters per token when estimating. */
const CHARS_PER_TOKEN = 4;
/** File extensions of the image types `image()` accepts. */
const IMAGE_EXTENSIONS: Record<string, string> = {
"image/png": ".png",
"image/jpeg": ".jpg",
"image/gif": ".gif",
"image/webp": ".webp",
};
/** Like the script's `text()`: strings as is, other values as compact JSON. */
function valueText(value: unknown): string {
if (typeof value === "string") return value;
return JSON.stringify(value) ?? String(value);
}
/**
* Lay out the script's output so the model can tell items apart: with more than one text item, each starts with a
* `==> text N/M <==` line; `console` lines follow all other output in one `<console_output>` block.
*/
function formatItems(output: readonly CodemodeOutputItem[]): Item[] {
const total = output.filter((item) => item.type === "text" && !item.console).length;
const items: Item[] = [];
const consoleLines: string[] = [];
let index = 0;
for (const item of output) {
if (item.type === "image") items.push(item);
else if (item.console) consoleLines.push(item.text);
else {
index++;
items.push({ type: "text", text: total > 1 ? `==> text ${index}/${total} <==\n${item.text}` : item.text });
}
}
if (consoleLines.length > 0) {
items.push({ type: "text", text: `<console_output>\n${consoleLines.join("\n")}\n</console_output>` });
}
return items;
}
/** Join adjacent text items into one, each part starting on its own line. */
function joinAdjacentText(items: readonly Item[]): Item[] {
const joined: Item[] = [];
for (const item of items) {
const last = joined.at(-1);
if (item.type === "text" && last?.type === "text") {
const separator = last.text === "" || last.text.endsWith("\n") ? "" : "\n";
joined[joined.length - 1] = { type: "text", text: `${last.text}${separator}${item.text}` };
} else {
joined.push(item);
}
}
return joined;
}
/**
* Apply the token budget: when the combined text exceeds it, the text items become one item that keeps the start and
* end of the text, and images follow it. The complete text is written to a temp file.
*/
async function truncate(
items: readonly Item[],
maxTokens: number,
limits: Limits,
write: OutputFile,
): Promise<{ items: Item[]; fullOutputPath?: string }> {
const texts = items.filter((item): item is TextContent => item.type === "text").map((item) => item.text);
const combined = texts.join("\n");
// Within the limits too: a UTF-16 unit is at most 3 UTF-8 bytes, and the rest of the result is a few short lines.
const budget = Math.min(maxTokens * CHARS_PER_TOKEN, Math.floor(limits.maxBytes / 4));
const lineBudget = Math.floor((limits.maxLines - LINE_MARGIN) / 2);
const lines = combined.split("\n").length;
if (texts.length === 0 || (combined.length <= budget && lines <= 2 * lineBudget)) return { items: [...items] };
const headChars = Math.floor(budget / 2);
const tailChars = budget - headChars;
const head = combined.slice(0, headChars).split("\n").slice(0, lineBudget).join("\n");
const tail = tailChars > 0 ? combined.slice(-tailChars).split("\n").slice(-lineBudget).join("\n") : "";
const removed = combined.length - head.length - tail.length;
let text = `Warning: truncated output (original token count: ${Math.ceil(combined.length / CHARS_PER_TOKEN)})\nTotal output lines: ${combined.split("\n").length}\n\n${head}…${Math.ceil(removed / CHARS_PER_TOKEN)} tokens truncated…${tail}`;
const spilled = await write("output.txt", combined);
text +=
"path" in spilled
? `\n\n[Full output: ${spilled.path} (read with offset/limit)]`
: `\n\n[Could not save the full output: ${spilled.error}]`;
return {
items: [{ type: "text", text }, ...items.filter((item) => item.type === "image")],
...("path" in spilled ? { fullOutputPath: spilled.path } : {}),
};
}
/** Save each image to a file and put a text item with its path before it; an image shown twice is saved once. */
async function saveImages(items: readonly Item[], write: OutputFile): Promise<Item[]> {
const labels = new Map<string, Promise<string>>();
const label = async ({ data, mimeType }: ImageContent, index: number): Promise<string> => {
const bytes = Uint8Array.from(atob(data), (char) => char.charCodeAt(0));
const kind = `${mimeType}, ${bytes.length} bytes`;
// A failed write must not discard the result of a script whose calls already ran: it becomes part of the label.
const saved = await write(`image-${index}${IMAGE_EXTENSIONS[mimeType] ?? ".bin"}`, bytes);
return "path" in saved
? `[Image saved to ${saved.path} (${kind})]`
: `[Image (${kind}) could not be saved: ${saved.error}]`;
};
const result = await Promise.all(
items.map(async (item): Promise<Item[]> => {
if (item.type !== "image") return [item];
let pending = labels.get(item.data);
if (pending === undefined) {
pending = label(item, labels.size + 1);
labels.set(item.data, pending);
}
return [{ type: "text", text: await pending }, item];
}),
);
return result.flat();
}
function formatError(
result: Extract<CodemodeResult, { ok: false }>,
calls: readonly { name: string; status: string }[],
) {
const { error } = result;
const head =
error.kind === "script"
? (error.stack ?? `${error.name ?? "Error"}: ${error.message}`)
: error.kind === "timeout"
? `Script timed out: ${error.message}`
: error.kind === "aborted"
? `Script aborted: ${error.message}`
: `Script sandbox failed: ${error.message}`;
const summary =
calls.length === 0
? "No tool calls were made."
: `Tool calls made before the failure (they are not undone): ${calls.map((call) => `${call.name} (${call.status})`).join(", ")}`;
return `${head}\n\n${summary}`;
}
/** The tool result's content for a script's result, and the path of the complete text when it was cut. */
export async function formatResult(
result: CodemodeResult,
options: {
/** Images `models.generateImages()` returned; when the script showed none, the result says so. */
readonly generatedImages: number;
readonly wallMs: number;
readonly maxOutputTokens: number;
readonly calls: readonly { name: string; status: string }[];
readonly limits: Limits;
readonly write: OutputFile;
},
): Promise<{ content: Item[]; fullOutputPath?: string }> {
const output = [...result.output];
if (result.ok && result.value !== undefined) output.push({ type: "text", text: valueText(result.value) });
const items = formatItems(output);
if (!result.ok) items.push({ type: "text", text: `Script error:\n${formatError(result, options.calls)}` });
const generated = options.generatedImages;
if (generated > 0 && !items.some((item) => item.type === "image")) {
items.push({
type: "text",
text: `Note: models.generateImages() returned ${generated} image${generated === 1 ? "" : "s"} that the script did not show. Show each image block of result.output with image(block).`,
});
}
const truncated = await truncate(joinAdjacentText(items), options.maxOutputTokens, options.limits, options.write);
// After truncation, which joins the text items and moves images after them, so each path stays next to its image.
const content = joinAdjacentText(await saveImages(truncated.items, options.write));
const header = `${result.ok ? "Script completed" : "Script failed"}\nWall time ${(options.wallMs / 1000).toFixed(1)} seconds\nOutput:\n`;
return {
content: [{ type: "text", text: header }, ...content],
...(truncated.fullOutputPath === undefined ? {} : { fullOutputPath: truncated.fullOutputPath }),
};
}

View file

@ -0,0 +1,85 @@
/**
* BM25 tool search for code mode's `searchTools()`, a copy of the coding agent's tool search ranking: durable code mode
* keeps its own copy of what it needs from the coding agent.
*/
export type SearchDocument = { readonly name: string; readonly text: string };
const STOP_WORDS: ReadonlySet<string> = new Set([
"a",
"an",
"and",
"are",
"as",
"at",
"be",
"by",
"for",
"from",
"in",
"is",
"it",
"of",
"on",
"or",
"that",
"the",
"this",
"to",
"with",
]);
/** Naive singular form, so `issues` matches `issue` and `searches` matches `search`. */
function stem(term: string): string {
if (term.length > 4 && term.endsWith("ies")) return `${term.slice(0, -3)}y`;
if (term.length > 4 && /(ches|shes|sses|xes|zes)$/.test(term)) return term.slice(0, -2);
if (term.length > 3 && term.endsWith("s") && !term.endsWith("ss")) return term.slice(0, -1);
return term;
}
/** Lowercase terms, split at camelCase boundaries and non-alphanumerics, without stop words. */
export function tokenize(text: string): string[] {
return text
.replace(/([a-z0-9])([A-Z])/g, "$1 $2")
.replace(/([A-Z]+)([A-Z][a-z])/g, "$1 $2")
.toLowerCase()
.split(/[^a-z0-9]+/)
.filter((term) => term.length > 0 && !STOP_WORDS.has(term))
.map(stem);
}
/** Okapi BM25 (k1 1.2, b 0.75): the names of the best matches, best first; ties keep document order. */
export function rankDocuments(query: string, documents: readonly SearchDocument[], limit: number): string[] {
const k1 = 1.2;
const b = 0.75;
const queryTerms = [...new Set(tokenize(query))];
if (queryTerms.length === 0 || documents.length === 0 || limit <= 0) return [];
const termCounts = documents.map((document) => {
const counts = new Map<string, number>();
for (const term of tokenize(document.text)) counts.set(term, (counts.get(term) ?? 0) + 1);
return counts;
});
const lengths = termCounts.map((counts) => [...counts.values()].reduce((sum, count) => sum + count, 0));
const averageLength = lengths.reduce((sum, length) => sum + length, 0) / documents.length || 1;
const idf = new Map(
queryTerms.map((term) => {
const frequency = termCounts.filter((counts) => counts.has(term)).length;
return [term, Math.log(1 + (documents.length - frequency + 0.5) / (frequency + 0.5))] as const;
}),
);
const matches: { name: string; score: number }[] = [];
documents.forEach((document, index) => {
let score = 0;
for (const term of queryTerms) {
const count = termCounts[index]!.get(term);
if (!count) continue;
const norm = k1 * (1 - b + (b * lengths[index]!) / averageLength);
score += (idf.get(term) ?? 0) * ((count * (k1 + 1)) / (count + norm));
}
if (score > 0) matches.push({ name: document.name, score });
});
return matches
.sort((left, right) => right.score - left.score)
.slice(0, limit)
.map((match) => match.name);
}

View file

@ -528,13 +528,23 @@ async function finalResult(
runtime.report(new Error(broken));
}
}
// Without a schema, or for an error result without structured output, programs get the output itself, so they get
// the error's text; the model reads only `output`, so model-issued calls skip it.
const schemaValue =
tool.structuredOutputSchema !== undefined && !(final.isError === true && final.structuredOutput === undefined);
if (input.kind === "nested" && !schemaValue) {
// Programs get the tool's own value: without a schema, its output. The model reads only `output`, so model-issued
// calls skip both.
if (input.kind === "nested" && tool.structuredOutputSchema === undefined) {
final = { ...final, structuredOutput: outputValue(bounded.content) as JsonValue };
}
// A schema tool's error result without structured output has its text only in the output, which programs never
// get: it becomes an error diagnostic, as a thrown error's message does, unless one says what went wrong already.
if (
input.kind === "nested" &&
tool.structuredOutputSchema !== undefined &&
final.isError === true &&
final.structuredOutput === undefined &&
![...(final.diagnostics ?? []), ...harness].some((diagnostic) => diagnostic.severity === "error")
) {
const text = bounded.content.flatMap((item) => (item.type === "text" ? [item.text] : [])).join("\n");
if (text !== "") harness.push(toolDiagnostic("tool_error", text));
}
return { ...final, output: bounded.content, diagnostics: [...(final.diagnostics ?? []), ...harness] };
}

View file

@ -177,10 +177,11 @@ export type NestedToolExecutionResult = {
/** The nested call's tool task. */
readonly taskId: TaskId;
/**
* The value for programs. A tool with `structuredOutputSchema`: its validated `structuredOutput`. A tool without, and
* an error result without `structuredOutput`: its bounded output, as a string for one text item, an `ImageContent`
* for one image, `""` for none, and the content list otherwise. Absent on an error result the Harness wrote itself
* (unavailable, invalid, blocked, interrupted, aborted); `diagnostics` say what went wrong.
* The tool's value for programs. A tool with `structuredOutputSchema`: its validated `structuredOutput`, also on an
* error result that has one; absent otherwise, and an error's text is then in an error diagnostic. A tool without:
* its bounded output, success or error, as a string for one text item, an `ImageContent` for one image, `""` for
* none, and the content list otherwise. Absent on an error result the Harness wrote itself (unavailable, invalid,
* blocked, interrupted, aborted); `diagnostics` say what went wrong.
*/
readonly structuredOutput?: JsonValue;
readonly isError: boolean;

View file

@ -0,0 +1,96 @@
import { type AssistantMessage, fauxAssistantMessage, fauxText, Type } from "@earendil-works/pi-ai";
import { CodemodeSandbox } from "@earendil-works/pi-codemode";
import { CODEMODE_SOURCE_GRAMMAR } from "@earendil-works/pi-codemode/source";
import { defineTool, MemoryStorage, type ToolRegistration } from "@earendil-works/pi-durable";
import { describe, expect, it } from "vitest";
import { createCodemodeExtension } from "../src/codemode/codemode.ts";
import { type CodemodeNamespace, codemodeDescription, codemodeSection } from "../src/codemode/describe.ts";
import { allEntries, chatSetup, openChat } from "./chat-support.ts";
import { context } from "./session-support.ts";
function tool(name: string, description: string, extra: Partial<ToolRegistration> = {}): ToolRegistration {
return defineTool({
name,
description,
parameters: Type.Object({ path: Type.String({ description: "File path" }) }),
execute: async () => ({ output: [] }),
...extra,
});
}
const GITHUB: CodemodeNamespace = { name: "github", description: "GitHub MCP server", prefix: "mcp__github__" };
describe("code mode description and section", () => {
it("keeps the tool description the same for every conversation", () => {
const text = codemodeDescription(false);
expect(text).toContain("`await tools.<name>({ ...args })`");
expect(text).toContain("searchTools(query");
expect(text).toContain("The `codemode` section of the system prompt lists the tools scripts can call");
expect(text).not.toContain("`models`");
expect(codemodeDescription(true)).toContain("`models`");
});
it("lists callable tools as declarations, grouped by namespace, with namespace descriptions", () => {
const section = codemodeSection(
[tool("read", "Read a file."), tool("mcp__github__list_issues", "List issues.")],
[GITHUB],
3000,
)!;
expect(section).toContain("### `read`\nRead a file.");
expect(section).toContain("read(args: {");
expect(section).toContain("## github\nGitHub MCP server");
// Tools before namespaces; a namespace's tools after its heading.
expect(section.indexOf("### `read`")).toBeLessThan(section.indexOf("## github"));
expect(section.indexOf("## github")).toBeLessThan(section.indexOf("### `mcp__github__list_issues`"));
expect(section).toContain("Promise<string>");
});
it("lists what fits the budget, every namespace before any is complete, and says what is left out", () => {
const github = Array.from({ length: 30 }, (_, index) => tool(`mcp__github__tool${index}`, "x".repeat(200)));
const section = codemodeSection([tool("read", "Read a file."), ...github], [GITHUB], 400)!;
expect(section).toContain("### `read`");
expect(section).toContain("## github (some tools not listed)");
expect(section).not.toContain("tool29");
expect(codemodeSection([...github], [GITHUB], 10)!).toContain("## github (tools not listed)");
expect(codemodeSection([], [GITHUB], 3000)).toBeUndefined();
});
it("uses the longest matching namespace prefix", () => {
const section = codemodeSection(
[tool("mcp__github__x", "x")],
[{ name: "git", prefix: "mcp__git" }, GITHUB],
3000,
)!;
expect(section).toContain("## github");
expect(section).not.toContain("## git\n");
});
it("renders the section from the conversation's callable tools into the system prompt", async () => {
const setup = chatSetup();
setup.registry.install(createCodemodeExtension({ sandbox: new CodemodeSandbox(), namespaces: [GITHUB] }));
setup.registry.install({
name: "tools",
tools: [
tool("read", "Read a file."),
tool("mcp__github__list_issues", "List issues.", { callers: ["tools"] }),
],
});
setup.faux.setResponses([fauxAssistantMessage([fauxText("hi")]) as AssistantMessage]);
const { harness, root } = await openChat(new MemoryStorage(), setup);
await (await root.submit({ type: "input", content: "go" }, context)).wait(context);
const system = (await allEntries(root)).find((entry) => entry.kind === "pi.system")!;
const sections = (system.model![0] as { sections?: Record<string, string> }).sections!;
expect(sections.codemode).toContain("### `read`");
expect(sections.codemode).toContain("### `mcp__github__list_issues`");
// The code mode tool itself is not callable from scripts.
expect(sections.codemode).not.toContain("### `codemode`");
// The tool declaration carries the script grammar; pi-ai applies it only for models that support it.
const tools = (system.model![0] as { toolsAdded?: { name: string; constrainedSampling?: unknown }[] })
.toolsAdded!;
expect(tools.find((candidate) => candidate.name === "codemode")?.constrainedSampling).toEqual({
type: "grammar",
variants: { openai_lark: CODEMODE_SOURCE_GRAMMAR },
});
await harness.close(context);
});
});

View file

@ -0,0 +1,271 @@
import { mkdtemp, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import type { Context } from "@earendil-works/chord";
import {
type AssistantImages,
type ClassifierResult,
createProvider,
fauxAssistantMessage,
fauxText,
fauxToolCall,
type MutableModels,
} from "@earendil-works/pi-ai";
import { CodemodeSandbox } from "@earendil-works/pi-codemode";
import { MemoryStorage, ToolResultEntry, UsageDoc } from "@earendil-works/pi-durable";
import { describe, expect, it } from "vitest";
import { createCodemodeExtension } from "../src/codemode/codemode.ts";
import type { ExecutionEnv } from "../src/env/index.ts";
import { NodeExecutionEnv } from "../src/env/node.ts";
import { allEntries, chatSetup, openChat } from "./chat-support.ts";
import { context } from "./session-support.ts";
import { deferred } from "./task-support.ts";
const COST = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 };
const USAGE = {
input: 10,
output: 2,
cacheRead: 0,
cacheWrite: 0,
totalTokens: 12,
cost: { input: 0.001, output: 0.0002, cacheRead: 0, cacheWrite: 0, total: 0.0012 },
};
const PNG = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==";
/** A chat whose models also have a classifier and an image model of a test provider, and the code mode extension. */
function modelsChat() {
const setup = chatSetup();
const calls = { classify: 0, images: 0, inFlight: 0, maxInFlight: 0 };
// \`atFour\`: hold classifier calls until four run, then a little longer, so a fifth would show up.
const gate = { open: deferred(), atFour: false };
gate.open.resolve();
(setup.models as MutableModels).setProvider(
createProvider({
id: "test",
auth: { apiKey: { name: "Test", resolve: async () => ({ auth: {} }) } },
models: [
{
id: "judge",
name: "Judge",
api: "test-classifier",
provider: "test",
baseUrl: "https://example.invalid",
input: ["text"],
cost: COST,
headers: { authorization: "secret" },
type: "classifier",
contextWindow: 1000,
},
{
id: "painter",
name: "Painter",
api: "test-images",
provider: "test",
baseUrl: "https://example.invalid",
input: ["text"],
cost: COST,
type: "image",
output: ["image"],
},
],
classifiers: {
"test-classifier": {
classify: async (model, classifierContext): Promise<ClassifierResult> => {
calls.classify++;
if ((classifierContext.state as { fail?: boolean }).fail === true) {
return {
api: model.api,
provider: model.provider,
model: model.id,
answers: {},
stopReason: "error",
errorMessage: "quota exceeded",
timestamp: 0,
};
}
calls.inFlight++;
calls.maxInFlight = Math.max(calls.maxInFlight, calls.inFlight);
if (gate.atFour && calls.inFlight === 4) setTimeout(() => gate.open.resolve(), 50);
await gate.open.promise;
calls.inFlight--;
return {
api: model.api,
provider: model.provider,
model: model.id,
answers: { good: { value: true } } as unknown as ClassifierResult["answers"],
usage: USAGE,
stopReason: "stop",
timestamp: 0,
};
},
},
},
images: {
"test-images": {
generateImages: async (model): Promise<AssistantImages> => {
calls.images++;
return {
api: model.api,
provider: model.provider,
model: model.id,
output: [{ type: "image", data: PNG, mimeType: "image/png" }],
stopReason: "stop",
timestamp: 0,
};
},
},
},
}),
);
setup.registry.install(createCodemodeExtension({ sandbox: new CodemodeSandbox() }));
return { setup, calls, gate };
}
async function run(chat: ReturnType<typeof modelsChat>, code: string, env?: ExecutionEnv) {
chat.setup.faux.setResponses([
fauxAssistantMessage([fauxToolCall("codemode", { code }, { id: "s1" })], { stopReason: "toolUse" }),
fauxAssistantMessage([fauxText("done")]),
]);
const { harness, root } = await openChat(new MemoryStorage(), chat.setup, env === undefined ? {} : { env });
await (await root.submit({ type: "input", content: "go" }, context)).wait(context);
const entry = (await allEntries(root)).find((candidate) => ToolResultEntry.is(candidate))!;
const message = entry.model![0] as { content: { type: string; text?: string }[]; details?: unknown };
const text = message.content
.map((item) => item.text ?? `<${item.type}>`)
.join("")
.replace(/^Script (completed|failed)\nWall time [\d.]+ seconds\nOutput:\n/, "");
const usage = await harness.snapshot(UsageDoc, root.id, context);
await harness.close(context);
return { text, message, usage };
}
describe("code mode models.*", () => {
it("reads the model catalog without credentials in headers", async () => {
const chat = modelsChat();
const { text } = await run(
chat,
`const classifiers = await models.getModelsOfType("classifier", "test");
const one = await models.getModelOfType("classifier", "test", "judge");
const available = await models.getAvailableOfType("image");
return [classifiers.map((m) => m.id), one.id, "headers" in one, available.map((m) => m.id)];`,
);
expect(JSON.parse(text)).toEqual([["judge"], "judge", false, ["painter"]]);
});
it("classifies through an internal tool: a nested call with its usage, hidden from tools and listings", async () => {
const chat = modelsChat();
const { text, message, usage } = await run(
chat,
`const model = await models.getModelOfType("classifier", "test", "judge");
const result = await models.classify(model, {
state: { text: "hi" },
questions: { good: { type: "bool", instructions: "Is it good?", criteria: { true: "yes", false: "no" } } },
});
const found = (await searchTools("classify generate images models")).length;
const described = (await describeTool("pi.codemode.models.classify")) !== undefined;
return [result.stopReason, result.answers.good.value, Object.keys(tools).some((name) => name.includes("models")), ALL_TOOLS.some((tool) => tool.name.includes("models")), found, described];`,
);
expect(JSON.parse(text)).toEqual(["stop", true, false, false, 0, false]);
expect(chat.calls.classify).toBe(1);
expect((message.details as { calls: { name: string; status: string }[] }).calls).toMatchObject([
{ name: "models.classify", status: "ok" },
]);
// Counted once, under the internal tool, by the Harness.
expect(usage?.tools["pi.codemode.models.classify"]?.cost.total).toBeCloseTo(0.0012);
});
it("resolves a model call that failed with its result, as the coding agent does", async () => {
const chat = modelsChat();
const { text, message } = await run(
chat,
`const model = await models.getModelOfType("classifier", "test", "judge");
const result = await models.classify(model, {
state: { fail: true },
questions: { good: { type: "bool", instructions: "?", criteria: { true: "y", false: "n" } } },
});
return [result.stopReason, result.errorMessage];`,
);
expect(JSON.parse(text)).toEqual(["error", "quota exceeded"]);
// The row still shows the failure.
expect((message.details as { calls: { name: string; status: string }[] }).calls).toMatchObject([
{ name: "models.classify", status: "error" },
]);
});
it("rejects a bad classifier context with the expected shape, before any request", async () => {
const chat = modelsChat();
const { text } = await run(
chat,
`const model = await models.getModelOfType("classifier", "test", "judge");
try { await models.classify(model, { prompt: "x" }); } catch (error) { return error.message; }`,
);
expect(text).toContain("context.state must be an object");
expect(text).toContain("Expected context:");
expect(chat.calls.classify).toBe(0);
});
it("runs at most four model calls at once", async () => {
const chat = modelsChat();
chat.gate.open = deferred();
chat.gate.atFour = true;
const { text } = await run(
chat,
`const model = await models.getModelOfType("classifier", "test", "judge");
const context = { state: {}, questions: { good: { type: "bool", instructions: "?", criteria: { true: "y", false: "n" } } } };
return (await Promise.all(Array.from({ length: 8 }, () => models.classify(model, context)))).length;`,
);
expect(text).toBe("8");
expect(chat.calls.classify).toBe(8);
expect(chat.calls.maxInFlight).toBe(4);
});
it("does not start a model call that waited for a slot after the script ended", async () => {
const chat = modelsChat();
chat.gate.open = deferred();
const directory = await mkdtemp(join(tmpdir(), "pi-codemode-models-"));
// The result spills to a file after the script ended and before the tool returns. Then the running
// calls finish and free their slots, which the two waiting calls would take.
class SlowTempFiles extends NodeExecutionEnv {
override async createTempDir(prefix: string | undefined, ctx: Context) {
chat.gate.open.resolve();
await new Promise((resolve) => setTimeout(resolve, 100));
return super.createTempDir(prefix, ctx);
}
}
try {
const { text } = await run(
chat,
`const model = await models.getModelOfType("classifier", "test", "judge");
const context = { state: {}, questions: { good: { type: "bool", instructions: "?", criteria: { true: "y", false: "n" } } } };
for (let i = 0; i < 6; i++) models.classify(model, context).catch(() => {});
// A journaled call after them: by its result, the first four have been launched.
await models.getModelsOfType("image");
return "x".repeat(60000);`,
new SlowTempFiles({ cwd: directory }),
);
expect(text).toContain("[Full output: ");
expect(chat.calls.classify).toBe(4);
} finally {
await rm(directory, { recursive: true, force: true });
}
});
it("notes generated images the script did not show", async () => {
const chat = modelsChat();
const hidden = await run(
chat,
`const model = await models.getModelOfType("image", "test", "painter");
const result = await models.generateImages(model, { input: [{ type: "text", text: "a cat" }] });
return result.output.length;`,
);
expect(hidden.text).toContain("Note: models.generateImages() returned 1 image that the script did not show");
const shown = await run(
modelsChat(),
`const model = await models.getModelOfType("image", "test", "painter");
const result = await models.generateImages(model, { input: [{ type: "text", text: "a cat" }] });
image(result.output[0]);`,
);
expect(shown.text).not.toContain("Note:");
expect(shown.message.content.some((item) => item.type === "image")).toBe(true);
});
});

View file

@ -1,6 +1,7 @@
import { mkdtemp, rm } from "node:fs/promises";
import { mkdtemp, readdir, readFile, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { dirname, join } from "node:path";
import type { Context } from "@earendil-works/chord";
import { type AssistantMessage, fauxAssistantMessage, fauxText, fauxToolCall, Type } from "@earendil-works/pi-ai";
import { type CodemodeRun, CodemodeSandbox } from "@earendil-works/pi-codemode";
import {
@ -20,6 +21,8 @@ import {
type CodemodeToolDetails,
createCodemodeTool,
} from "../src/codemode/codemode.ts";
import type { CodemodeNamespace } from "../src/codemode/describe.ts";
import { NodeExecutionEnv } from "../src/env/node.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";
@ -60,7 +63,7 @@ async function sqlitePath(): Promise<string> {
* - `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[] = []) {
function codemodeChat(globals: CodemodeGlobal[] = [], namespaces: CodemodeNamespace[] = []) {
const setup = chatSetup();
const sandbox = new CodemodeSandbox({ memoryLimitBytes: 64 * 1024 * 1024 });
const vm: { starts: number; onStart?: (run: CodemodeRun) => void } = { starts: 0 };
@ -71,7 +74,7 @@ function codemodeChat(globals: CodemodeGlobal[] = []) {
vm.onStart?.(run);
return run;
};
addTool(setup.registry, createCodemodeTool({ sandbox, globals }));
addTool(setup.registry, createCodemodeTool({ sandbox, globals, namespaces }));
const echo = new Map<string, number>();
addTool(
setup.registry,
@ -149,9 +152,13 @@ async function resultText(setup: ChatSetup, root: Awaited<ReturnType<typeof open
const entry = (await allEntries(root)).find((candidate) => ToolResultEntry.is(candidate))!;
const message = entry.model![0]!;
return {
// Without the "Script completed" header and its wall time, which tests of it check themselves.
text:
message.role === "toolResult"
? message.content.map((item) => (item.type === "text" ? item.text : "")).join("\n")
? message.content
.map((item) => (item.type === "text" ? item.text : ""))
.join("")
.replace(/^Script (completed|failed)\nWall time [\d.]+ seconds\nOutput:\n/, "")
: "",
isError: message.role === "toolResult" && message.isError,
};
@ -676,6 +683,80 @@ describe("durable code mode", () => {
await harness.close(context);
});
it("resolves a schema tool's error result that carries its value, and rejects one without with its text", async () => {
const chat = codemodeChat();
const Exit = Type.Object({ output: Type.String(), exitCode: Type.Number() });
addTool(
chat.setup.registry,
defineTool({
name: "run",
description: "run",
parameters: Type.Object({}),
structuredOutputSchema: Exit,
// Like bash on a nonzero exit: an error result with its value.
execute: async () => ({
output: [{ type: "text", text: "no matches" }],
structuredOutput: { output: "no matches", exitCode: 1 },
isError: true,
}),
}),
);
addTool(
chat.setup.registry,
defineTool({
name: "denied",
description: "denied",
parameters: Type.Object({}),
structuredOutputSchema: Exit,
execute: async () => ({ output: [{ type: "text", text: "permission denied" }], isError: true }),
}),
);
chat.holds.block = false;
chat.setup.faux.setResponses([
script(
`const r = await tools.run({});
let e; try { await tools.denied({}); } catch (error) { e = [error.message, error.structuredOutput ?? null]; }
return [r.exitCode, r.output, e];`,
),
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,"no matches",["permission denied",null]]',
isError: false,
});
await harness.close(context);
});
it("resolves a schema-less tool with its text, as its declaration says, also when it returned images", async () => {
const chat = codemodeChat();
const PNG = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==";
addTool(
chat.setup.registry,
defineTool({
name: "picture",
description: "picture",
parameters: Type.Object({ caption: Type.Boolean() }),
execute: async (args) => ({
output: [
...(args.caption ? [{ type: "text" as const, text: "a cat" }] : []),
{ type: "image" as const, data: PNG, mimeType: "image/png" },
],
}),
}),
);
chat.holds.block = false;
chat.setup.faux.setResponses([
script(`return [await tools.picture({ caption: true }), await tools.picture({ caption: false })];`),
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: '["a cat",""]', isError: false });
await harness.close(context);
});
it("delivers a result only after its journal commit", async () => {
const storage = new ControlledStorage();
const seen: string[] = [];
@ -787,14 +868,14 @@ describe("durable code mode", () => {
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);
const texts = results.map((entry) => (entry.model![0] as { content: { text: string }[] }).content[1]!.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]");
expect((forkResult.model![0] as { content: { text: string }[] }).content[1]!.text).toBe("[null]");
await harness.close(context);
});
@ -830,3 +911,144 @@ describe("durable code mode", () => {
await harness.close(context);
});
});
describe("durable code mode output and tool search", () => {
/** Run `code` once in a chat with a Node environment; the result's content and details. */
async function runScript(code: string, chat = codemodeChat()) {
chat.holds.block = false;
const cwd = await mkdtemp(join(tmpdir(), "pi-codemode-env-"));
directories.add(cwd);
chat.setup.faux.setResponses([script(code), DONE]);
const { harness, root } = await openChat(await openNodeSqliteStorage(await sqlitePath()), chat.setup, {
env: new NodeExecutionEnv({ cwd }),
});
await (await root.submit({ type: "input", content: "go" }, context)).wait(context);
const entry = (await allEntries(root)).find((candidate) => ToolResultEntry.is(candidate))!;
const message = entry.model![0] as {
content: { type: string; text?: string }[];
isError: boolean;
details?: CodemodeToolDetails;
};
await harness.close(context);
return { text: message.content.map((item) => item.text ?? `<${item.type}>`).join(""), message };
}
it("lays out text items, console lines, and the returned value under a header", async () => {
const { text } = await runScript(`text("one"); console.log("log", 1); text({ two: 2 }); return "three";`);
expect(text).toMatch(/^Script completed\nWall time [\d.]+ seconds\nOutput:\n/);
expect(text).toContain('==> text 1/3 <==\none\n==> text 2/3 <==\n{"two":2}\n==> text 3/3 <==\nthree');
expect(text).toContain("<console_output>\nlog 1\n</console_output>");
});
it("cuts text past the token budget to its start and end, with the full text in a file", async () => {
const { text, message } = await runScript(
`// @options: {"max_output_tokens": 20}
text("BEGIN" + "x".repeat(1000) + "END");`,
);
expect(text).toContain("Warning: truncated output");
expect(text).toContain("BEGIN");
expect(text).toContain("END");
const path = message.details?.fullOutputPath;
expect(text).toContain(`[Full output: ${path}`);
expect(await readFile(path!, "utf8")).toBe(`BEGIN${"x".repeat(1000)}END`);
});
it("keeps a result inside the tool's output limits whatever max_output_tokens says, spilling the rest", async () => {
// 6 MiB of text: over the Harness's 4 MiB limit, which would cut it without a spill file.
const big = await runScript(
`// @options: {"max_output_tokens": 2000000}
return "x".repeat(6 * 1024 * 1024);`,
);
expect(big.text).toContain("Warning: truncated output");
expect(big.text.length).toBeLessThanOrEqual(1024 * 1024 + 1000);
expect((await readFile(big.message.details!.fullOutputPath!, "utf8")).length).toBe(6 * 1024 * 1024);
// 200,000 short lines: within the byte limit, over the 100,000 lines one.
const long = await runScript(
`// @options: {"max_output_tokens": 2000000}
return Array.from({ length: 200000 }, () => "a").join("\\n");`,
);
expect(long.text).toContain("Warning: truncated output");
expect(long.text.split("\n").length).toBeLessThan(100_000);
expect(long.message.details?.fullOutputPath).toBeDefined();
});
it("writes the result's files to the same paths when it reruns after a crash", async () => {
const chat = codemodeChat();
chat.holds.block = false;
const cwd = await mkdtemp(join(tmpdir(), "pi-codemode-env-"));
directories.add(cwd);
const written = deferred<string>();
// The process dies after the full output was written and before the result was committed.
class CrashAfterWrite extends NodeExecutionEnv {
override async writeFile(path: string, content: string | Uint8Array, callContext: Context) {
const result = await super.writeFile(path, content, callContext);
written.resolve(path);
await aborted(callContext.abortSignal!);
return result;
}
}
const storage = await sqlitePath();
chat.setup.faux.setResponses([script(`// @options: {"max_output_tokens": 20}\nreturn "x".repeat(1000);`), DONE]);
let opened = await openChat(await openNodeSqliteStorage(storage), chat.setup, {
env: new CrashAfterWrite({ cwd }),
});
const id = (await opened.root.submit({ type: "input", content: "go" }, context)).id;
const first = await written.promise;
await opened.harness.close(context);
opened = await openChat(await openNodeSqliteStorage(storage), chat.setup, { env: new NodeExecutionEnv({ cwd }) });
opened.harness.resume();
expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done");
const entry = (await allEntries(opened.root)).find((candidate) => ToolResultEntry.is(candidate))!;
const details = (entry.model![0] as { details?: CodemodeToolDetails }).details;
expect(details?.fullOutputPath).toBe(first);
directories.add(dirname(first));
expect(await readdir(dirname(first))).toEqual(["output.txt"]);
await opened.harness.close(context);
});
it("saves images to files and names them next to the image", async () => {
const png = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==";
const { text, message } = await runScript(`image("data:image/png;base64,${png}");`);
const saved = /\[Image saved to (\S+) \(image\/png, \d+ bytes\)\]/.exec(text);
expect(saved).not.toBeNull();
expect((await readFile(saved![1]!)).toString("base64")).toBe(png);
expect(message.content.some((item) => item.type === "image")).toBe(true);
});
it("reports a failed script with its error and the calls already made", async () => {
const { text, message } = await runScript(`await tools.echo({ text: "a" }); throw new Error("boom");`);
expect(message.isError).toBe(true);
expect(text).toMatch(/^Script failed\n/);
expect(text).toContain("Error: boom");
expect(text).toContain("Tool calls made before the failure (they are not undone): echo (ok)");
});
it("finds tools with searchTools(), describeTool(), and describeNamespace(), and lists them in ALL_TOOLS", async () => {
const chat = codemodeChat([], [{ name: "github", description: "GitHub MCP server", prefix: "mcp__github__" }]);
addTool(
chat.setup.registry,
tool("mcp__github__list_issues", async () => ({ output: [] })),
);
const { text } = await runScript(
`const found = await searchTools("issues");
const inNamespace = await searchTools("echo", { namespace: "github" });
return {
found: found.map((tool) => tool.name),
inNamespace,
described: (await describeTool("echo")).includes("echo(args:"),
missing: (await describeTool("nope")) ?? null,
namespace: await describeNamespace("github"),
all: ALL_TOOLS.map((tool) => tool.name).includes("mcp__github__list_issues"),
};`,
chat,
);
expect(JSON.parse(text.replace(/^Script completed\nWall time [\d.]+ seconds\nOutput:\n/, ""))).toEqual({
found: ["mcp__github__list_issues"],
inNamespace: [],
described: true,
missing: null,
namespace: { name: "github", description: "GitHub MCP server", tools: ["mcp__github__list_issues"] },
all: true,
});
});
});

View file

@ -208,8 +208,9 @@ describe("structured output", () => {
addTool(setup.registry, returning(name, result, schema ? { structuredOutputSchema: Count } : {}));
const received = await nested(setup, name);
expect(received.isError).toBe(true);
// The broken value is dropped; the error result carries the bounded output instead, here none.
expect(received.structuredOutput).toBe("");
// The broken value is dropped; a schema-less tool's error result still carries its output, here none.
if (schema) expect(received).not.toHaveProperty("structuredOutput");
else expect(received.structuredOutput).toBe("");
expect(codes(received)).toEqual(["invalid_structured_output"]);
expect(received.diagnostics[0]?.message).toContain(message);
@ -236,7 +237,7 @@ describe("structured output", () => {
expect(received).not.toHaveProperty("structuredOutput");
});
it("gives callers the output of a schema tool's error result without structured output", async () => {
it("gives callers a schema tool's error text as a diagnostic, never as structured output", async () => {
const setup = chatSetup();
addTool(
setup.registry,
@ -246,8 +247,26 @@ describe("structured output", () => {
{ structuredOutputSchema: Count },
),
);
// Its text reaches the caller: what a code mode script's rejection carries.
expect(await nested(setup, "fails")).toMatchObject({ isError: true, structuredOutput: "nope", diagnostics: [] });
// Its text reaches the caller as an error diagnostic: structuredOutput is only ever the tool's own value.
const received = await nested(setup, "fails");
expect(received).toMatchObject({
isError: true,
diagnostics: [{ severity: "error", code: "tool_error", message: "nope" }],
});
expect(received).not.toHaveProperty("structuredOutput");
// The model's call keeps the text in its output, without the extra diagnostic.
const modelSetup = chatSetup();
addTool(
modelSetup.registry,
returning(
"fails",
{ output: [{ type: "text", text: "nope" }], isError: true },
{ structuredOutputSchema: Count },
),
);
const stored = await direct(modelSetup, "fails");
expect(stored.content).toEqual([{ type: "text", text: "nope" }]);
// A thrown error's message is in its tool_error diagnostic, not in the output.
const thrown = chatSetup();
@ -263,11 +282,32 @@ describe("structured output", () => {
},
}),
);
expect(await nested(thrown, "throws")).toMatchObject({
const thrownResult = await nested(thrown, "throws");
// One diagnostic: the thrown message, not a second one from its (empty) output.
expect(thrownResult).toMatchObject({
isError: true,
structuredOutput: "",
diagnostics: [{ code: "tool_error", message: "permission denied" }],
});
expect(thrownResult.diagnostics).toHaveLength(1);
expect(thrownResult).not.toHaveProperty("structuredOutput");
// An error result whose own diagnostics say what went wrong gets no second one from its output.
const explained = chatSetup();
addTool(
explained.registry,
returning(
"fails",
{
output: [{ type: "text", text: "partial log" }],
isError: true,
diagnostics: [{ severity: "error", code: "quota", message: "quota exceeded" }],
},
{ structuredOutputSchema: Count },
),
);
expect((await nested(explained, "fails")).diagnostics).toEqual([
{ severity: "error", code: "quota", message: "quota exceeded" },
]);
// An error result that has structured output keeps it.
const kept = chatSetup();