diff --git a/packages/durable/CHANGELOG.md b/packages/durable/CHANGELOG.md index 49068ddf9..9fe3d7f9e 100644 --- a/packages/durable/CHANGELOG.md +++ b/packages/durable/CHANGELOG.md @@ -13,18 +13,19 @@ ### Added - `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. The result is kept, as returned, in `NestedResultDoc`, a task-scoped document family of the caller, 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. +- `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. 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. - `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; nothing loads it unless imported. A model without image input gets a diagnostic saying it sees a placeholder. ### Fixed - A throwing `onReport` no longer replaces the error being reported or becomes an unhandled rejection, a throwing `now()` is reported once and `Date.now` used, and a watch listener's error is reported as well as ending the watch with `listener_error`. - A task whose fault write failed no longer runs again in the same process. -- The scheduler no longer visits every live task on each pass: abort cascades start only from tasks with cancellation intent, owned live work and idle scopes are walks down an in-memory ownership index, and reservation, finalization, and failFast checks keep their own candidate sets. A chain of 500 owned tasks that took 12 s to settle takes 50 ms, 20,000 take 1.5 s; a tool making 8000 parallel nested calls finishes in half the time (`npm run bench:scheduler-fanout`). +- The scheduler no longer visits every live task on each pass: abort cascades start only from tasks with cancellation intent, owned live work and idle scopes are walks down an in-memory ownership index, and reservation, finalization, and failFast checks keep their own candidate sets. A chain of 500 owned tasks that took 12 s to settle takes 50 ms, 20,000 take 1.5 s; a tool making 8000 parallel nested calls finishes in 2.7 s instead of 18.8 s, with one document per nested call instead of one growing index (`npm run bench:scheduler-fanout`). ## [1.1.0] - 2026-10-07 diff --git a/packages/durable/README.md b/packages/durable/README.md index 9bb8629d0..24ec0cbbf 100644 --- a/packages/durable/README.md +++ b/packages/durable/README.md @@ -227,10 +227,10 @@ const testsPass = defineTool({ - The caller gets the nested call's `taskId`, `structuredOutput`, `details`, `diagnostics`, `usage`, and `isError`, not the `output` meant for the model. A blocked, invalid, failed, or aborted nested call returns an `isError` result. `executeTool()` itself rejects once `execute()` has returned, when the caller is aborted, or when the Harness closes; cancelling its `context` stops only the wait. - Hooks see a nested call's `call.parent`, so a guard can treat calls from scripts differently. -- While the caller runs, its nested calls show in `docs["pi.live"].nestedTools`, with the arguments they run with and a `summary` (error, duration, usage) once done, and as tool events with `parentToolCallId` and `parentTaskId`. `progress: false` keeps a nested call's running output out of `pi.live`. Before the caller settles, nested calls it left running are aborted. The caller's result message does not list its nested calls; their task records keep them. +- While the caller runs, its nested calls show in `docs["pi.live"].nestedTools`, with the arguments they run with and a `summary` (error, duration, usage) once done, and as tool events with `parentToolCallId` and `parentTaskId`. `progress: false` keeps a nested call's running output out of `pi.live`. Before the caller settles, nested calls it left running are aborted, also those an earlier attempt made. The caller's result message does not list its nested calls; their task records keep them. - Each nested call has a key, its position among the caller's nested calls unless the caller passes `{ key }`. A replay-safe caller that reruns after a crash and calls in the same order, or passes the same keys, gets the nested calls it already made back, finished or still running, instead of starting them again. A nested call the crash interrupted before its result was committed follows its own `replay` policy, so a replay-safe tool with external effects still needs its own idempotency. A caller that is not replay-safe never reruns, so after a restart its unfinished nested calls are abandoned, with everything they own, before they run again. - A nested call's `usage` is counted under its own tool; passing it on in the caller's result counts it twice. -- Nested results are not kept: they live in the caller's task documents (`NestedResultDoc`, one per nested call) only until the caller settles. The nested call's task record keeps only a small receipt. A stored result is exactly what `executeTool()` returns, without the model's `output`, so an image or a `bash` tail is stored once. +- Nested results are not kept: they live in the caller's task documents (`NestedCallDoc`, one per nested call, keyed by its key, holding `{ taskId, result? }`) only until the caller settles. The nested call's task record keeps only a small receipt. A stored result is exactly what `executeTool()` returns, without the model's `output`, so an image or a `bash` tail is stored once. - Nothing bounds how deeply calls nest; a tool with the default `callers` may even call itself. Restrict `callers` where that matters (below). ### Who may call a tool diff --git a/packages/durable/docs/spec.md b/packages/durable/docs/spec.md index 3f6a1652c..725c81f98 100644 --- a/packages/durable/docs/spec.md +++ b/packages/durable/docs/spec.md @@ -1727,6 +1727,8 @@ interface TaskRuntime extends DocumentObserver, Docum waitForTask(id: TaskId, context: Context): Promise>; /** `abortTask()` of a task this task owns, resolving once it is terminal; rejects for a task another owns. */ abortOwned(id: TaskId, context: Context): Promise; + /** Live tasks this task owns directly, in ID order, whichever invocation created them; not those in its conversations. */ + ownedTasks(context: Context): Promise[]>; /** Outcomes of terminal tasks, in order; rejects when one is not terminal. Used after a wait (section 5.5). */ outcomes(ids: readonly TaskId[], context: Context): Promise[]>; /** Committed entry visible from the task's conversation. */ @@ -2878,9 +2880,10 @@ The nested call's key names it within its caller; the nested call ID is `/`. By default the key is the call's position among the invocation's nested calls, `1`, `2`, ... An explicit `key` must be non-empty, without `/`, not `__proto__`, and not a positive integer, so call IDs never -collide. One commit creates the nested task, its slot, and an index entry from -the key to the task in `pi.tool.nested`, a task-scoped document of the caller, -which retires with it. A caller that is not replay-safe creates its nested +collide. One commit creates the nested task, its slot, and its member of +`NestedCallDoc` (`pi.tool.nested-call`), a task-scoped document family of the +caller keyed by the call's key, holding `{ taskId }`; it retires with the +caller. One document per call, so admitting a call rewrites no other call's. A caller that is not replay-safe creates its nested calls, and the child tasks it creates through `api.createTask()` unless the options say otherwise, with `abandonOnRestart` (section 5.4): it never resumes after a restart, so neither do they. With `progress: false`, a nested call @@ -2892,15 +2895,14 @@ never get the partial output of an interrupted or aborted call. Nothing bounds h deeply calls nest; `callers` keeps a tool from being called by tools at all. The nested call's result does not go into its receipt, which stays -`{ kind: "nested" }`. Its terminal commit writes the result to -`NestedResultDoc` (`pi.tool.nested-result`), a task-scoped document family of -the caller keyed by the nested task ID, one member per result, and sets its +`{ kind: "nested" }`. Its terminal commit writes the result as `result` into +its `NestedCallDoc` member and sets its slot's `summary` (`isError`, `durationMs`, `usage`, bounded error text). The stored result is exactly what `executeTool()` returns: it keeps no `output`, which only the model reads, so each result, an image or a `bash` tail, is stored once. The caller is never terminal before its owned nested calls, so the write always finds its -documents open; a nested call missing from the caller's index throws. The +documents open; a member that names another task throws. The caller reads the result after the wait, and every nested result is gone from storage once its caller is terminal. A nested call the scheduler faulted or orphaned has no stored result; the caller gets a synthesized `isError` one. @@ -2920,8 +2922,9 @@ settles (section 8.6); a caller that passes it on in its own result counts it again. Before a caller settles, by any path, it lets nested call admissions its -invocation started commit, then aborts every nested call it left running with -`runtime.abortOwned()`, all at once, and waits for them, so they settle into +invocation started commit, then aborts every live nested call it owns, found +with `runtime.ownedTasks()`, so also one an earlier attempt made and a rerun +never made again, with `runtime.abortOwned()`, all at once, and waits for them, so they settle into their slots before the caller's settlement removes them. `executeTool()` throws once `execute()` has returned, and rejects when the caller is aborted or the Harness closes; cancelling its `context` stops only the wait. An abort that @@ -3662,7 +3665,7 @@ with the arguments the caller passed, runs like a tool slot, gets the arguments it executes with in its intent commit when they differ (a model-issued call's are in its assistant entry), and in its terminal commit becomes `done`, with its progress removed like a tool slot's and its `summary` set, enough for a status line; its -result is in the caller's `NestedResultDoc`, never in `pi.live` (section 7.3). When any call settles, its terminal commit +result is in the caller's `NestedCallDoc`, never in `pi.live` (section 7.3). When any call settles, its terminal commit removes the nested slots below it, transitively by `parentTaskId`; its left-running nested calls were aborted first. `endRun` and the generation's `tools` phase remove `nestedTools` with `tools`. @@ -3860,7 +3863,7 @@ because the terminal record keeps it; the call is read from the assistant entry. A nested call (section 7.3) carries its call, since no entry holds it, and settles differently: its result commit records the result's `usage` (section 8.6), writes its result for programs to the caller's -`NestedResultDoc` (section 7.3), marks its nested slot `done` with its +`NestedCallDoc` (section 7.3), marks its nested slot `done` with its summary, appends no entry, and ends with `{ kind: "nested" }`; `control` does not apply. Below, `{ entryId }` is a model-issued call's result. @@ -4469,11 +4472,11 @@ Events derive from committed changes: result entry appended in the same commit, as for the unstarted calls of an aborted round, directly before its `message_start`, or without an entry. Nested slots produce the same events with `parentToolCallId`; a nested call - ends with its `result`, taken from the `NestedResultDoc` value its terminal + ends with its `result`, taken from the `NestedCallDoc` value its terminal commit published, or without one after a fault or orphan or when its slot disappears unfinished. Snapshots show done nested slots with their summary, not their result; a late subscriber reads a live caller's results from its - `NestedResultDoc`. + `NestedCallDoc`, keyed by call key (the slot's `callId` after `/`). - `inbox_update`, `agent_changed`, `usage_changed`: the document changed; a retired one reads as its initial value, as in a snapshot. Registry installs and settings changes make no commit and emit nothing; a UI showing resolved diff --git a/packages/durable/src/harness/events.ts b/packages/durable/src/harness/events.ts index 080ed6254..fad248a8a 100644 --- a/packages/durable/src/harness/events.ts +++ b/packages/durable/src/harness/events.ts @@ -16,7 +16,7 @@ import type { import { AgentDoc } from "./agent.ts"; import type { InboxItem, InboxState } from "./inbox.ts"; import type { CompactionStatus, LiveState, NestedToolSlot, ToolSlot } from "./live.ts"; -import { NestedResultDoc } from "./tool.ts"; +import { NestedCallDoc, type NestedCallState } from "./tool.ts"; import type { AgentState, CompactionReason, Harness, NestedToolExecutionResult, ToolDiagnostic } from "./types.ts"; import { UsageDoc, type UsageState } from "./usage.ts"; import { scanAll } from "./util.ts"; @@ -230,16 +230,18 @@ function translate( const tasks = new Map(); const submissions: SubmissionRecord[] = []; // Nested results this commit stored, by nested task ID; read now, since the caller's documents may retire next. - const nestedResults = new Map(); + const nestedResults = new Map(); for (const change of publication.changes) { if ( change.type === "document" && - change.record.kind === NestedResultDoc.definition.kind && - change.record.key !== undefined && + change.record.kind === NestedCallDoc.definition.kind && change.value !== null && change.conversationId === conversationId ) { - nestedResults.set(change.record.key, change.value.result as unknown as NestedToolExecutionResult); + const state = change.value as NestedCallState; + if (state.taskId !== undefined && state.result !== undefined) { + nestedResults.set(state.taskId, state.result as unknown as NestedToolExecutionResult); + } } if (change.type === "entry" && change.value.conversationId === conversationId) entries.push(change.value); if (change.type === "task" && change.value.conversationId === conversationId) @@ -324,7 +326,7 @@ function translate( // Nested calls that end in this commit, the same way, children before the calls that made them: the lists hold // parents first, so walk them backwards. const endNested = (slot: NestedToolSlot): void => { - const result = nestedResults.get(String(slot.taskId)); + const result = nestedResults.get(slot.taskId); events.push({ type: "tool_execution_end", ...callOf(slot), ...(result === undefined ? {} : { result }) }); }; const nestedNow = new Map(nested.map((slot) => [slot.taskId, slot])); diff --git a/packages/durable/src/harness/live.ts b/packages/durable/src/harness/live.ts index db5c5d73b..a7a5d6d32 100644 --- a/packages/durable/src/harness/live.ts +++ b/packages/durable/src/harness/live.ts @@ -47,7 +47,7 @@ export type NestedToolSlot = SlotProgress & { * once it starts executing. A model-issued call's arguments are in its assistant entry instead. */ arguments: JsonObject; - /** Once done: how the call ended, for a status line. The result itself is in the caller's `NestedResultDoc`. */ + /** Once done: how the call ended, for a status line. The result itself is in the caller's `NestedCallDoc`. */ summary?: NestedToolSummary; }; diff --git a/packages/durable/src/harness/scheduler.ts b/packages/durable/src/harness/scheduler.ts index 13ce79b75..e3f659334 100644 --- a/packages/durable/src/harness/scheduler.ts +++ b/packages/durable/src/harness/scheduler.ts @@ -1583,6 +1583,24 @@ export class TaskScheduler { this.#read(invocation, () => this.waitForTask(id, withAbortSignal(invocation.controller.signal, context)), )) as ErasedRuntime["waitForTask"], + ownedTasks: (context) => + this.#read(invocation, () => + this.#session.readOnLine(async () => { + // The step that ends the invocation may have run on the line before this. + if (invocation.ended) throw endedError(invocation); + context.abortSignal?.throwIfAborted(); + // Every live task is listed under its parent node, its owner task when it has one. + const owned: AnyTaskRecord[] = []; + for (const node of this.#below.get(invocation.taskId) ?? []) { + const record = typeof node === "string" ? undefined : this.#live.get(node); + if (record?.owner === invocation.taskId) owned.push(record); + } + // Copies, as from storage: the scheduler reads its own. + return owned + .sort((a, b) => a.id - b.id) + .map((record) => copyJson(record as unknown as JsonValue) as unknown as AnyTaskRecord); + }), + ), abortOwned: (id, context) => this.#read(invocation, async () => { const bound = withAbortSignal(invocation.controller.signal, context); diff --git a/packages/durable/src/harness/tool.ts b/packages/durable/src/harness/tool.ts index 5c75761e6..b301d0470 100644 --- a/packages/durable/src/harness/tool.ts +++ b/packages/durable/src/harness/tool.ts @@ -4,7 +4,7 @@ import { overlap } from "@earendil-works/chord/delta"; import type { ImageContent, TextContent, ToolCall, ToolResultMessage } from "@earendil-works/pi-ai"; import { validateToolArguments } from "@earendil-works/pi-ai/utils/validation"; import { Compile } from "typebox/compile"; -import { defineDoc, defineDocFamily } from "../documents.ts"; +import { defineDocFamily } from "../documents.ts"; import { AssistantEntry, ToolResultEntry } from "../entries.ts"; import { errorMessage } from "../errors.ts"; import { defineTask } from "../tasks.ts"; @@ -68,31 +68,37 @@ export type ToolTaskCheckpoint = /** * A model-issued call's result is its transcript entry; a nested call's result is in its caller's - * `NestedResultDoc`, so the receipt stays small and the result retires with the caller. A version 1 task that was + * `NestedCallDoc`, so the receipt stays small and the result retires with the caller. A version 1 task that was * already holding its outcome as `completing` when the Harness upgraded finishes with a version 1 result, * `{ entryId, control? }` without `kind`; treat a result without `kind` as a model-issued call's. */ export type ToolTaskResult = { kind: "model"; entryId: EntryId; control?: ToolControl } | { kind: "nested" }; -/** Index of a call's nested calls by key, so a rerun of the call finds the nested calls it already made. */ -const NestedCallsDoc = defineDoc<{ calls: Record> }>({ - kind: "pi.tool.nested", - version: 1, - scope: "task", - initial: () => ({ calls: {} }), -}); +/** + * Memo of a tool task that made a nested call, in any attempt, committed before the first admission of each attempt: + * only then can `pi.live` hold slots below it, which settling removes. A crash between the two only costs a scan. + */ +const MADE_NESTED_CALLS = "pi.tool.madeNestedCalls"; + +/** One nested call of a running tool: its task, and once it ended, its result. */ +export type NestedCallState = { + /** Absent only inside the commit that creates the call; a creation that fails leaves no member. */ + taskId?: TaskId; + /** Exactly what `executeTool()` returns; absent while running, and for a call the scheduler faulted or orphaned. */ + result?: JsonRepresentation; +}; /** - * The result of one nested call, exactly as `executeTool()` returns it, a member of its caller's family keyed by the - * nested task ID. Written in the nested call's terminal commit; one document per result, so no write rewrites other - * results. Retires with the caller. + * The nested calls of a tool call, a task-scoped family of the calling task keyed by the nested call's key: so a rerun + * finds the calls it already made, and the caller reads each result. Created with the nested task; the nested call's + * terminal commit adds its result. One document per call, so no commit rewrites the others. Retires with the caller. */ -export const NestedResultDoc = defineDocFamily<{ result: JsonRepresentation }, JsonObject>({ - kind: "pi.tool.nested-result", +export const NestedCallDoc = defineDocFamily({ + kind: "pi.tool.nested-call", version: 1, scope: "task", family: true, - initial: (seed) => seed as { result: JsonRepresentation }, + initial: () => ({}), }); type Runtime = TaskRuntime; @@ -102,7 +108,7 @@ type Content = (TextContent | ImageContent)[]; * Built-in tool task: resolves the called tool among its phase agent's tools, validates, runs `beforeTool`, records intent, * executes, runs `afterTool`, and settles the result, all in one `call` handler so nothing separates resolution from * settlement. `execute` is reached only by recovery and applies the replay rule. A model-issued call settles by - * appending its result entry; a nested call by writing its result to its caller's `NestedResultDoc`. + * appending its result entry; a nested call by writing its result to its caller's `NestedCallDoc`. */ export const ToolTask = defineTask({ name: "pi.tool", @@ -277,6 +283,8 @@ async function run( const resumes = tool.replay === "safe"; /** Default keys: the order of this invocation's calls, so a rerun that calls in the same order reattaches. */ let sequence = 0; + /** This attempt's commit of the `MADE_NESTED_CALLS` memo. */ + let madeNestedMemo: Promise | undefined; let ended = false; const assertLive = (): void => { if (ended) throw new Error(`Tool call ${call.id} has settled`); @@ -353,22 +361,30 @@ async function run( assertLive(); const key = options.key ?? String(++sequence); if (options.key !== undefined) checkKey(options.key); - const admission = startNestedCall( - runtime, - call.id, - name, - toolArgs, - { key, progress: options.progress, abandonOnRestart: !resumes }, - callContext, + // The memo first, once per attempt; a failed one is tried again by the next call. + madeNestedMemo ??= runtime.memo(MADE_NESTED_CALLS, true, callContext).catch((error: unknown) => { + madeNestedMemo = undefined; + throw error; + }); + // Registered now, so cleanup waits for it. + const admission = madeNestedMemo.then(() => + startNestedCall( + runtime, + call.id, + name, + toolArgs, + { key, progress: options.progress, abandonOnRestart: !resumes }, + callContext, + ), ); admissions.push(admission); const id = await admission; const settled = await runtime.waitForTask(id, callContext); - const stored = await runtime.snapshot(NestedResultDoc, runtime.taskId, String(id), callContext); + const stored = (await runtime.snapshot(NestedCallDoc, runtime.taskId, key, callContext))?.result; // A copy the tool may change; a nested call the scheduler faulted or orphaned stored none. return stored === undefined ? fallbackResult(id, name, settled.state.outcome) - : (copyJson(stored.result as unknown as JsonValue) as unknown as NestedToolExecutionResult); + : (copyJson(stored as unknown as JsonValue) as unknown as NestedToolExecutionResult); }, }; @@ -546,7 +562,7 @@ function structuredOutputError(tool: ToolRegistration, result: ToolExecutionResu /** * Commit the tool's terminal state and mark its slot done. A model-issued call appends its result entry and ends with - * the entry ID; a nested call writes its result to the caller's `NestedResultDoc`, records its usage, and ends with a + * the entry ID; a nested call writes its result to the caller's `NestedCallDoc`, records its usage, and ends with a * small receipt. Nested calls this call left running are aborted first, so they report into their slots before the * slots below the call leave `pi.live`. `build` receives the slot so interruption and abort can report the durable * partial output. @@ -561,8 +577,12 @@ async function settle( durationMs?: number, ): Promise { // A crash in between leaves the call in `execute`: a safe rerun reattaches to the calls, an interruption ends here. - const index = await runtime.snapshot(NestedCallsDoc, runtime.taskId, context); - const nestedIds = Object.values(index?.calls ?? {}).sort((a, b) => a - b); + // Every live nested call the task owns, also those an earlier attempt made and this one never called again. + const nestedIds = (await runtime.ownedTasks(context)) + .filter((record) => record.kind === ToolTask.definition.name && (record.input as ToolTaskInput).kind === "nested") + .map((record) => record.id); + // Read before the settling commit: no admission starts after `ended`, and those underway were awaited. + const madeNested = (await runtime.memo(MADE_NESTED_CALLS, context)) === true; // Mark them all before waiting for any: one may run until a sibling is cancelled. await Promise.all(nestedIds.map((id) => runtime.abortOwned(id, context))); await runtime.commit(async (tx) => { @@ -574,12 +594,11 @@ async function settle( if (input.kind === "nested") { const nested = nestedResult(runtime.taskId, result, durationMs); if (nested.usage !== undefined) await recordUsage(tx, conversationId, "tools", call.name, nested.usage); - // The caller is never terminal before its owned nested calls; a missing entry is a bug, not a race. - if ((await tx.doc(NestedCallsDoc, input.parent)).calls[input.key] !== runtime.taskId) { - throw new Error(`Nested call ${call.id} is not in its caller's index`); - } - const stored = { result: nested as unknown as JsonRepresentation }; - await tx.doc(NestedResultDoc, input.parent, String(runtime.taskId), stored); + // The caller is never terminal before its owned nested calls, so its document is open; another task in it is a + // bug, not a race. + const state = await tx.doc(NestedCallDoc, input.parent, input.key, null); + if (state.taskId !== runtime.taskId) throw new Error(`Nested call ${call.id} is not in its caller's index`); + state.result = nested as unknown as Draft>; if (slot !== undefined && "parentTaskId" in slot) { finishSlot(slot); slot.summary = summaryOf(nested, result.output ?? []); @@ -599,8 +618,9 @@ async function settle( : { control: copyJson(result.control as JsonValue, { omitUndefinedProperties: true }) as ToolControl }; settled = { kind: "model", entryId: entry.id, ...control }; } - // A call that made no nested calls has no slots below it; skipping the scan keeps many leaves linear. - if (nestedIds.length > 0) removeNestedSlots(live, runtime.taskId); + // A call that made no nested calls, in any attempt, has no slots below it; skipping the scan keeps many leaves + // linear. Committed before the read: no nested call is admitted after `ended`. + if (madeNested) removeNestedSlots(live, runtime.taskId); if (ending.status === "aborted") return { status: "terminal", outcome: { status: "aborted", result: settled } }; if (ending.status === "failed") { const error = { message: ending.message }; @@ -635,7 +655,7 @@ function nestedResult( /** * Create nested call `key` of the call `parentCallId` in one commit: its tool task, owned by the calling task, its slot, - * and its index entry. A key already in the index returns that call's task when it names the same tool and arguments. + * and its `NestedCallDoc`. A key already made returns that call's task when it names the same tool and arguments. */ async function startNestedCall( runtime: Runtime, @@ -654,8 +674,9 @@ async function startNestedCall( }; let id: TaskId | undefined; await runtime.commit(async (tx) => { - const index = await tx.doc(NestedCallsDoc, runtime.taskId); - const existing = Object.hasOwn(index.calls, key) ? index.calls[key] : undefined; + // Opened, or created empty for a key not made yet. + const state = await tx.doc(NestedCallDoc, runtime.taskId, key, null); + const existing = state.taskId; if (existing !== undefined) { const input = (await tx.task(existing))?.input as ToolTaskInput | undefined; const same = @@ -679,7 +700,7 @@ async function startNestedCall( ownership: { kind: "task", taskId: runtime.taskId }, abandonOnRestart: options.abandonOnRestart, }); - index.calls[key] = created; + state.taskId = created; const live = await tx.doc(LiveDoc, runtime.conversationId); live.nestedTools ??= []; const slot: NestedToolSlot = { @@ -846,7 +867,7 @@ function boundContent( /** * Explicit nested call keys are path segments of the call ID, so IDs of different nested calls never collide, and are - * never plain positive integers, which default keys use. `__proto__` would not land in the index as an own key. + * never plain positive integers, which default keys use. `__proto__` stays reserved, as it was when keys indexed a JSON object. */ function checkKey(key: string): void { if (key === "" || key.includes("/") || key === "__proto__" || /^[1-9][0-9]*$/.test(key)) { diff --git a/packages/durable/src/index.ts b/packages/durable/src/index.ts index 4a9e205ab..aab6967a1 100644 --- a/packages/durable/src/index.ts +++ b/packages/durable/src/index.ts @@ -51,7 +51,8 @@ export { ProviderDoc, type ProviderState } from "./harness/provider.ts"; export { createRegistry } from "./harness/registry.ts"; export type { TaskGraph, TaskGraphNode, TaskGraphState, TaskGraphWatch } from "./harness/task-graph.ts"; export { - NestedResultDoc, + NestedCallDoc, + type NestedCallState, ToolTask, type ToolTaskCheckpoint, type ToolTaskInput, diff --git a/packages/durable/src/types.ts b/packages/durable/src/types.ts index d7aa8d513..c54981ba3 100644 --- a/packages/durable/src/types.ts +++ b/packages/durable/src/types.ts @@ -213,6 +213,11 @@ export interface TaskRuntime extends DocumentObserver * another owns; resolves at once for a terminal one. */ abortOwned(id: TaskId, context: Context): Promise; + /** + * Live tasks this task owns directly, in ID order: also those an earlier invocation created, before a restart. Tasks + * in conversations it owns are not included. Copies of the committed records. + */ + ownedTasks(context: Context): Promise[]>; /** Outcomes of terminal tasks, in order; rejects when one is missing or not terminal. Used after a wait. */ outcomes(ids: readonly TaskId[], context: Context): Promise[]>; /** diff --git a/packages/durable/test/harness-nested-tools-restart.test.ts b/packages/durable/test/harness-nested-tools-restart.test.ts index 3212f1ae9..342e103c0 100644 --- a/packages/durable/test/harness-nested-tools-restart.test.ts +++ b/packages/durable/test/harness-nested-tools-restart.test.ts @@ -462,6 +462,104 @@ describe("nested tool calls across a restart", () => { await opened.harness.close(context); }); + it("clears the slots of nested calls an earlier attempt finished when the rerun makes none", async () => { + const path = await sqlitePath(); + const setup = chatSetup(); + const finished = deferred(); + let attempts = 0; + addTool( + setup.registry, + tool("quick", async () => ({ output: [{ type: "text", text: "ok" }] })), + ); + addTool( + setup.registry, + tool( + "caller", + async (_args, api, callContext) => { + attempts++; + if (attempts === 1) { + await api.executeTool("quick", {}, callContext); + finished.resolve(); + return aborted(callContext.abortSignal!); + } + return { output: [{ type: "text", text: "done" }] }; + }, + { replay: "safe" }, + ), + ); + setup.faux.setResponses([call("caller"), DONE]); + let opened = await open(path, setup); + const id = (await opened.root.submit({ type: "input", content: "go" }, context)).id; + await finished.promise; + // The finished nested call's slot is still listed, below its running caller. + expect((await opened.harness.snapshot(LiveDoc, opened.root.id, context))?.nestedTools).toHaveLength(1); + await opened.harness.close(context); + + opened = await open(path, setup); + let slotsAfterCaller: number | undefined; + opened.harness.subscribeCommits((publication) => { + for (const change of publication.changes) { + if (change.type !== "task" || change.value.kind !== ToolTask.definition.name) continue; + if ((change.value.input as ToolTaskInput).kind !== "model" || change.value.state.status !== "terminal") + continue; + const live = publication.changes.find((each) => each.type === "document" && each.record.kind === "pi.live"); + const value = live?.type === "document" ? (live.value as { nestedTools?: unknown[] } | null) : undefined; + slotsAfterCaller = value?.nestedTools?.length ?? 0; + } + }); + opened.harness.resume(); + expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done"); + expect(attempts).toBe(2); + // The caller's terminal commit removed the slot, though the rerun made no nested call. + expect(slotsAfterCaller).toBe(0); + await opened.harness.close(context); + }); + + it("aborts a nested call an earlier attempt of a replay-safe caller left running and the rerun never made again", async () => { + const path = await sqlitePath(); + const setup = chatSetup(); + const running = deferred(); + let attempts = 0; + addTool( + setup.registry, + tool( + "slow", + async (_args, _api, callContext) => { + running.resolve(); + return aborted(callContext.abortSignal!); + }, + { replay: "safe" }, + ), + ); + addTool( + setup.registry, + tool( + "caller", + async (_args, api, callContext) => { + attempts++; + // Only the first attempt calls `slow`; the rerun after the restart returns at once. + if (attempts === 1) await api.executeTool("slow", {}, callContext, { key: "slow" }); + return { output: [{ type: "text", text: "done" }] }; + }, + { replay: "safe" }, + ), + ); + setup.faux.setResponses([call("caller"), DONE]); + let opened = await open(path, setup); + const id = (await opened.root.submit({ type: "input", content: "go" }, context)).id; + await running.promise; + await opened.harness.close(context); + + opened = await open(path, setup); + opened.harness.resume(); + expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done"); + expect(attempts).toBe(2); + // The rerun's settlement found the leftover call among the tasks it owns and aborted it. + const [leftover] = await nestedTasks(opened.harness); + expect(leftover?.state).toMatchObject({ status: "terminal", outcome: { status: "aborted" } }); + await opened.harness.close(context); + }); + it("reattaches a replay-safe caller by default and explicit keys, and rejects a rerun that calls differently", async () => { for (const order of ["same", "swapped"] as const) { const path = await sqlitePath(); diff --git a/packages/durable/test/harness-nested-tools.test.ts b/packages/durable/test/harness-nested-tools.test.ts index 83cab2ae9..59c05560b 100644 --- a/packages/durable/test/harness-nested-tools.test.ts +++ b/packages/durable/test/harness-nested-tools.test.ts @@ -20,7 +20,7 @@ import { type Harness, LiveDoc, MemoryStorage, - NestedResultDoc, + NestedCallDoc, type NestedToolExecutionResult, type TaskId, type ToolHookCall, @@ -96,6 +96,79 @@ afterEach(async () => { }); describe("nested tool calls", () => { + it("stores one document per nested call, keyed by its key, written only by that call", async () => { + const setup = chatSetup(); + const echo = echoTool(); + addTool(setup.registry, echo.registration); + addTool( + setup.registry, + tool("batch", async (_args, api, callContext) => { + await Promise.all([ + api.executeTool("echo", { text: "a" }, callContext), + api.executeTool("echo", { text: "b" }, callContext, { key: "named" }), + ]); + return {}; + }), + ); + setup.faux.setResponses([call("batch"), DONE]); + const storage = new MemoryStorage(); + // Per document written, the keys of the nested calls it holds. + const written: { kind: string; key: string | undefined }[] = []; + const commit = storage.commit.bind(storage); + storage.commit = (writes, commitContext) => { + for (const write of writes) { + if (write.type === "document.create") written.push({ kind: write.record.kind, key: write.record.key }); + } + return commit(writes, commitContext); + }; + const { harness, root } = await openChat(storage, setup); + await (await root.submit({ type: "input", content: "go" }, context)).wait(context); + const nested = written.filter((write) => write.kind.startsWith("pi.tool.")); + expect(nested).toEqual([ + { kind: NestedCallDoc.definition.kind, key: "1" }, + { kind: NestedCallDoc.definition.kind, key: "named" }, + ]); + await harness.close(context); + }); + + it("rejects a nested call of an aborted caller and leaves no document for it", async () => { + const setup = chatSetup(); + const echo = echoTool(); + addTool(setup.registry, echo.registration); + const started = deferred(); + const marked = deferred(); + let rejected: unknown; + let caller: TaskId | undefined; + addTool( + setup.registry, + tool("late", async (_args, api) => { + caller = api.taskId; + started.resolve(); + await marked.promise; + // Its context is aborted too, so pass one that is not: the admission itself must refuse. + try { + await api.executeTool("echo", { text: "a" }, context); + } catch (error) { + rejected = error; + } + return {}; + }), + ); + setup.faux.setResponses([call("late"), DONE]); + const { harness, root } = await openChat(new MemoryStorage(), setup); + const submission = await root.submit({ type: "input", content: "go" }, context); + await started.promise; + await harness.commit(async (tx) => { + const record = (await tx.task(caller!))!; + (tx as unknown as { setTask(value: unknown): void }).setTask({ ...record, abortRequested: true }); + }, context); + marked.resolve(); + await submission.wait(context); + expect(String(rejected)).toMatch(/abort mark|aborted|has settled/); + expect(await harness.snapshot(NestedCallDoc, caller!, "1", context)).toBeUndefined(); + await harness.close(context); + }); + it("runs a nested call as its own tool task and returns its result to the caller, not the transcript", async () => { const setup = chatSetup(); const echo = echoTool(); @@ -135,7 +208,7 @@ describe("nested tool calls", () => { status: "terminal", outcome: { status: "completed", result: { kind: "nested" } }, }); - expect(await harness.snapshot(NestedResultDoc, parent!.id, String(child!.id), context)).toBeUndefined(); + expect(await harness.snapshot(NestedCallDoc, parent!.id, "1", context)).toBeUndefined(); expect(await harness.snapshot(LiveDoc, root.id, context)).toEqual({}); await harness.close(context); }); diff --git a/packages/durable/test/harness-scheduler-indexes.test.ts b/packages/durable/test/harness-scheduler-indexes.test.ts index 6db765f9d..853684f06 100644 --- a/packages/durable/test/harness-scheduler-indexes.test.ts +++ b/packages/durable/test/harness-scheduler-indexes.test.ts @@ -3,7 +3,7 @@ import { mkdtemp, rm } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; -import { createSession, defineTask, type Harness, type TaskId } from "@earendil-works/pi-durable"; +import { createSession, defineTask, type Harness, type TaskId, type TaskRuntime } from "@earendil-works/pi-durable"; import { afterEach, describe, expect, it } from "vitest"; import { openNodeSqliteStorage } from "../src/storage/sqlite/node.ts"; import { context } from "./session-support.ts"; @@ -35,9 +35,36 @@ const Hold = defineTask<{ name: string }, { phase: "hold" }, null>({ runtime.commit(() => ({ status: "terminal", outcome: { status: "aborted" } }), ctx), }); +/** What `ownedTasks()` returned to each run of a `Probe`, by its name, and the runtime of its last run. */ +const owned = new Map(); +const runtimes = new Map, "ownedTasks">>(); + +/** Records its owned tasks once its `look` gate opens, then holds until its own gate opens. */ +const Probe = defineTask<{ name: string }, { phase: "look" }, null>({ + name: "test.index.probe", + version: 1, + initial: () => ({ phase: "look" }), + phases: { + look: async (task, runtime, ctx) => { + runtimes.set(task.input.name, runtime); + await Promise.race([gate(`look.${task.input.name}`).promise, aborted(runtime.signal)]); + owned.set( + task.input.name, + (await runtime.ownedTasks(ctx)).map((record) => record.id), + ); + await Promise.race([gate(task.input.name).promise, aborted(runtime.signal)]); + await runtime.commit(() => ({ status: "terminal", outcome: { status: "completed", result: null } }), ctx); + }, + }, + abort: async (_task, runtime, ctx) => + runtime.commit(() => ({ status: "terminal", outcome: { status: "aborted" } }), ctx), +}); + const directories: string[] = []; afterEach(async () => { gates.clear(); + owned.clear(); + runtimes.clear(); for (const directory of directories.splice(0)) await rm(directory, { recursive: true, force: true }); }); @@ -171,3 +198,59 @@ describe("scheduler indexes", () => { await opened.harness.close(context); }); }); + +describe("TaskRuntime.ownedTasks", () => { + it("lists the live tasks a task owns directly, in ID order, also after a reopen", async () => { + const path = await sqlitePath(); + let opened = await openTasks(await openNodeSqliteStorage(path), [Hold, Probe]); + const root = await opened.harness.root(context); + opened.harness.resume(); + // The probe owns: two held tasks, one that ends, a grandchild below one of them, and a conversation with a task. + const ids = await root.commit(async (tx) => { + const probe = await tx.createTask(Probe, { name: "probe" }, { ownership: { kind: "conversation" } }); + const first = await tx.createTask(Hold, { name: "first" }, { ownership: { kind: "task", taskId: probe } }); + const ends = await tx.createTask(Hold, { name: "ends" }, { ownership: { kind: "task", taskId: probe } }); + const second = await tx.createTask(Hold, { name: "second" }, { ownership: { kind: "task", taskId: probe } }); + await tx.createTask(Hold, { name: "grandchild" }, { ownership: { kind: "task", taskId: first } }); + const child = await tx.createConversation({ ownership: { kind: "task", taskId: probe } }); + await tx.createTask( + Hold, + { name: "inside" }, + { ownership: { kind: "conversation" }, conversationId: child.id }, + ); + return { probe, first, ends, second }; + }, context); + gate("ends").resolve(); + await opened.harness.waitForTask(ids.ends, context); + gate("look.probe").resolve(); + await eventually(() => owned.has("probe")); + expect(owned.get("probe")).toEqual([ids.first, ids.second]); + await opened.harness.close(context); + + // After a reopen the rerun sees the same tasks, which an earlier invocation created. + gates.clear(); + owned.clear(); + opened = await openTasks(await openNodeSqliteStorage(path), [Hold, Probe]); + opened.harness.resume(); + gate("look.probe").resolve(); + await eventually(() => owned.has("probe")); + expect(owned.get("probe")).toEqual([ids.first, ids.second]); + await opened.harness.close(context); + }); + + it("rejects once its invocation has ended", async () => { + const opened = await openTasks(await openNodeSqliteStorage(await sqlitePath()), [Hold, Probe]); + const root = await opened.harness.root(context); + opened.harness.resume(); + const probe = await root.commit( + (tx) => tx.createTask(Probe, { name: "probe" }, { ownership: { kind: "conversation" } }), + context, + ); + await eventually(() => runtimes.has("probe")); + const runtime = runtimes.get("probe")!; + // Aborting joins the run invocation, which has ended by then. + await opened.harness.abortTask(probe, context); + await expect(runtime.ownedTasks(context)).rejects.toThrow("invocation has ended"); + await opened.harness.close(context); + }); +}); diff --git a/packages/durable/test/harness-structured-output.test.ts b/packages/durable/test/harness-structured-output.test.ts index 566d21556..a12535995 100644 --- a/packages/durable/test/harness-structured-output.test.ts +++ b/packages/durable/test/harness-structured-output.test.ts @@ -13,7 +13,8 @@ import { import { defineTool, MemoryStorage, - NestedResultDoc, + NestedCallDoc, + type NestedCallState, type NestedToolExecutionResult, type ToolExecutionResult, type ToolRegistration, @@ -152,17 +153,16 @@ describe("structured output", () => { ); setup.faux.setResponses([call("probe"), DONE]); const storage = new MemoryStorage(); - const stored: JsonValue[] = []; - const commit = storage.commit.bind(storage); - storage.commit = (writes, commitContext) => { - for (const write of writes) { - if (write.type === "document.create" && write.record.kind === NestedResultDoc.definition.kind) { - stored.push(write.content.value.result!); - } - } - return commit(writes, commitContext); - }; const { harness, root } = await openChat(storage, setup); + // The result each nested call's document got, as committed. + const stored: JsonValue[] = []; + harness.subscribeCommits((publication) => { + for (const change of publication.changes) { + if (change.type !== "document" || change.record.kind !== NestedCallDoc.definition.kind) continue; + const result = (change.value as NestedCallState | null)?.result; + if (result !== undefined && change.ops.length > 0) stored.push(result as unknown as JsonValue); + } + }); const stream = await watchEvents(harness, root.id, context); const ends: NestedToolExecutionResult[] = []; stream.start(async (batch) => { diff --git a/packages/durable/test/nested-images-bench.ts b/packages/durable/test/nested-images-bench.ts index 4030064eb..4cb854520 100644 --- a/packages/durable/test/nested-images-bench.ts +++ b/packages/durable/test/nested-images-bench.ts @@ -2,7 +2,7 @@ * Nested image storage benchmark. One model-issued `script` call, as a code mode call would, loops `count` times: a * nested `read_image` call returns a unique base64 image of `bytes` characters, then a nested `classify` call reads it. * The `inline` variant passes the image data in `classify`'s arguments; the `reference` variant passes the `taskId` of - * the `read_image` result, whose value is the image, and `classify` resolves it from its caller's `NestedResultDoc`. + * the `read_image` result, whose value is the image, and `classify` resolves it from its caller's `NestedCallDoc`. * Reports commits, wall time, logical bytes written (JSON of the storage writes) by write type, the stored size before * `script` returns, after the run settles, after close, and after reopening, the heap (peak during the run, sampled * every 5 ms and at each commit; after a GC once the run settles, after `waitForIdle()`, and after reopening; all over @@ -25,9 +25,10 @@ import { defineTool, Harness, type JsonObject, - NestedResultDoc, + NestedCallDoc, type Storage, type StorageWrite, + type TaskId, ToolResultEntry, type ToolTaskInput, } from "../src/index.ts"; @@ -122,16 +123,14 @@ function registryFor(scenario: Scenario, probe: Probe) { for (const image of args.images) { let data = image.data; if (image.source !== undefined) { - // The caller is the nested call's `parent`; its `NestedResultDoc` holds the referenced result. + // The caller is the nested call's `parent`; its `NestedCallDoc` holds the referenced result. const input = (await api.getTask(api.taskId, callContext))?.input as ToolTaskInput | undefined; if (input?.kind !== "nested") throw new Error("classify resolves references only as a nested call"); - const stored = await api.snapshot( - NestedResultDoc, - input.parent, - String(image.source.taskId), - callContext, - ); - const block = stored?.result.structuredOutput as ImageContent | undefined; + // Results are keyed by call key, which the referenced call's input carries. + const source = (await api.getTask(image.source.taskId as TaskId, callContext))?.input as ToolTaskInput; + if (source.kind !== "nested") throw new Error("A reference names a nested call"); + const stored = await api.snapshot(NestedCallDoc, input.parent, source.key, callContext); + const block = stored?.result?.structuredOutput as ImageContent | undefined; data = block?.type === "image" ? block.data : undefined; } if (data?.length !== scenario.bytes) throw new Error(`Image has ${data?.length} characters`);