From 87eff524f96b69c919482ac46b26bd9333f26e22 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Wed, 28 Oct 2026 11:45:03 +0000 Subject: [PATCH] feat(durable): afterTool runs on every tool task result --- packages/durable/CHANGELOG.md | 1 + packages/durable/README.md | 4 +- packages/durable/docs/spec.md | 43 +- packages/durable/src/harness/tool.ts | 163 ++++- packages/durable/src/harness/types.ts | 13 +- .../durable/test/harness-tool-hooks.test.ts | 626 ++++++++++++++++++ packages/durable/test/harness-tools.test.ts | 4 +- 7 files changed, 817 insertions(+), 37 deletions(-) create mode 100644 packages/durable/test/harness-tool-hooks.test.ts diff --git a/packages/durable/CHANGELOG.md b/packages/durable/CHANGELOG.md index cf45dd551..441b8644c 100644 --- a/packages/durable/CHANGELOG.md +++ b/packages/durable/CHANGELOG.md @@ -92,6 +92,7 @@ - `FileChange` (`FileSystem.watch()` in `@earendil-works/pi-durable/env`) is tagged by `kind`: `{ kind: "paths", paths }`, `{ kind: "overflow" }`, or `{ kind: "failed", error }` (was `{ paths }`, `{ overflow: true }`, `{ error }`), and `FileWatcher.close(context)` is `stop(context)`. Before: `if ("overflow" in change) rescan(); else if ("paths" in change) reload(change.paths); await watcher.close(ctx);`. After: `if (change.kind === "overflow") rescan(); else if (change.kind === "paths") reload(change.paths); await watcher.stop(ctx);`. - `SessionEnd` (`Session.closed`, `Harness.closed`) is tagged by `kind`: `{ kind: "closed" }` or `{ kind: "failed", error }` (was `reason`). Before: `harness.closed.then((end) => { if (end.reason === "failed") reopen(end.error); })`. After: `harness.closed.then((end) => { if (end.kind === "failed") reopen(end.error); })`. - `ReportFrom`, the second argument of a task's `runtime.report(error, from)`, is tagged by `kind`: `{ kind: "extension", extension? }` or `{ kind: "host", callback }` (was `type`). Before: `runtime.report(error, { type: "host", callback: "env" })`. After: `runtime.report(error, { kind: "host", callback: "env" })`. +- `afterTool` runs on every result a tool task writes, not only on an executed call's: also on the error results the Harness writes itself, for a tool that does not resolve (`tool_unavailable`), invalid arguments (`invalid_arguments`, with `call.arguments` as the model sent them), a block (`blocked`), an interruption (`interrupted`), the restart limit (`restart_limit`), and an abort (`aborted`, `abandoned`); the diagnostic code tells them apart. Such a result is bounded with the tool's output limits after the hooks, never keeps a `structuredOutput`, and applies `control` only when the call completes; a nested call a hook turns into a success gets its bounded output as value unless the tool declares a schema. Once a tool task was cut off more than `settings.restartLimit` times its hooks no longer run, nor in the last attempt the scheduler gives its abort handler, and an abort's hooks delay `Conversation.abort()`. A non-replay-safe tool interrupted that often now ends with `restart_limit` too. Hooks that assumed the tool ran (cast `call.arguments`, read `structuredOutput`) check `result.diagnostics` first ([#11043](https://github.com/earendil-works/pi/issues/11043)) ### Added diff --git a/packages/durable/README.md b/packages/durable/README.md index 9b2882e22..996f484cf 100644 --- a/packages/durable/README.md +++ b/packages/durable/README.md @@ -292,7 +292,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 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 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, even when an `afterTool` hook returns one, and their `diagnostics` say what went wrong; one a hook turns into a success gives a schema-less tool's program its output. 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: @@ -630,7 +630,7 @@ const Guard = defineExtension({ ``` - **Generation:** `beforeRequest` (replace the messages of one request), `afterResponse`, `onYield` (continue the run with another user message), and `afterTools` (runs once a round's tools are done). -- **Tools:** `beforeTool` (block or rewrite arguments) and `afterTool` (replace the result). +- **Tools:** `beforeTool` (block or rewrite arguments) and `afterTool` (replace the result). `afterTool` sees every result a tool task writes, also the error results the Harness writes itself (unavailable, invalid, blocked, interrupted, restart limit, aborted), told apart by their diagnostic code; past `settings.restartLimit` it no longer runs. On an abort it runs inside the abort handler, so a slow hook delays `abort()`. To limit a hook to some conversations, select its extension only there, for example with `configure({ extensions: { add: [Guard] } })`. diff --git a/packages/durable/docs/spec.md b/packages/durable/docs/spec.md index 70a185a86..7fe639106 100644 --- a/packages/durable/docs/spec.md +++ b/packages/durable/docs/spec.md @@ -3217,7 +3217,10 @@ interface ToolHooks { api: HookApi, context: Context, ): HookResult<{ readonly arguments?: JsonObject; readonly block?: string }>; - /** After execution, before the result is committed; replaces the result. */ + /** + * Before the result is committed; replaces the result. Runs on every result the tool task writes, also the error + * results the Harness writes itself (section 8.4), except past the restart limit. + */ afterTool( call: ToolHookCall, result: ToolExecutionResult, @@ -3503,7 +3506,10 @@ what programs get. An error result the Harness writes itself (unavailable, invalid arguments, blocked, interrupted, aborted, abandoned, ended at the restart limit, or a scheduler fault) has no -`structuredOutput`; its diagnostics say what went wrong. A thrown error +`structuredOutput`; its diagnostics say what went wrong. When an `afterTool` +hook turns such a result into a success, programs get its bounded output for a +schema-less tool, or one that did not resolve, and no value for a schema tool; +the hook's own `structuredOutput` is dropped. 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: a schema tool's then has no `structuredOutput`, a schema-less tool's its output; @@ -4679,12 +4685,15 @@ Below, `{ entryId }` is a model-issued call's result. same handler executes the tool, runs the `afterTool` chain, and commits the result; it builds `api.env` with `HarnessOptions.env` right before executing (section 2.2). A tool that does not resolve, invalid arguments, or a block - commits the corresponding error result instead, without intent. + runs the `afterTool` chain on the corresponding error result and commits it + instead, without intent. - `execute` is reached only by recovery. When both the stored and the current policy, resolved as in `call`, are `safe`, it executes again with the stored arguments, without `beforeTool`, and settles like `call`. Otherwise it commits an `interrupted` error result from the slot's durable partial output, details, and - diagnostics, and ends `failed` with `{ entryId }`. + diagnostics, after the `afterTool` chain, and ends `failed` with `{ entryId }`. + Either way, a task cut off `settings.restartLimit` times ends with a + `restart_limit` result instead (section 5.4), also after the chain. - The result commit bounds the content, appends the diagnostics block (section 7.3), appends one `pi.tool-result` entry with `model: [{ role: "toolResult", toolCallId, toolName, content, details, isError, durationMs, timestamp }]` and @@ -4704,7 +4713,25 @@ Below, `{ entryId }` is a model-issued call's result. `completing` with the result entry already written (section 5.5); the generation resumes only when the tool task is terminal. - The abort handler commits an `aborted` error result from the slot's durable - partial output, details, and diagnostics and ends `aborted` with `{ entryId }`. + partial output, details, and diagnostics, after the `afterTool` chain, and ends + `aborted` with `{ entryId }`. The scheduler faults a task whose abort handler + was cut off `restartLimit` times, without a result (section 5.4), so the last + attempt it runs, at `restartLimit - 1` restarts, skips the chain. +- Every result the task writes passes the `afterTool` chain: the slot is read + before the result commit, since hooks run outside commits, and does not change + while the tool does not run. A result the Harness wrote itself is then bounded + with the tool's output limits (the defaults when it did not resolve), gets a + truncation diagnostic when that dropped content (a slot's truncation + diagnostic stays only while the hooks keep its output), and loses any + `structuredOutput`, except that a nested call the chain turns into a success + gets its bounded output as value, as a schema-less tool's success does + (unless the tool declares a schema, section 7.3); `control` applies only when + the task completes. Once the + task's restarts exceed `settings.restartLimit`, the chain no longer runs, + since it may be what keeps cutting the task off. Results that no tool task + writes never pass it: the generation's `aborted` results for calls a + sequential round never started (section 8.5), a nested call's fallback result + (section 7.3), and a missing result synthesized in context (section 2.1). ### 8.5 Tool rounds @@ -5916,6 +5943,12 @@ These are contracts, not invitations to add defensive machinery: background conversation-owned tasks awaited with `waitForTask()`; a non-background one would be aborted by any cascade that reaches its conversation. +- **Hooks in abort handlers:** a tool task's abort handler runs the `afterTool` + chain on its `aborted` result, and abort invocations cannot be aborted. A + slow hook delays the result, and with it `Conversation.abort()` and + `waitForIdle()`; a hanging one holds them until the host restarts, and the + handler's last attempt before the restart limit runs without hooks (section + 8.4). - **Hook memo names:** hooks share the asking task's memo namespace with the task and other hooks. Prefix memo names. - **Long transactions:** an async commit callback holds the Session mutation diff --git a/packages/durable/src/harness/tool.ts b/packages/durable/src/harness/tool.ts index 946294524..02179a817 100644 --- a/packages/durable/src/harness/tool.ts +++ b/packages/durable/src/harness/tool.ts @@ -42,6 +42,7 @@ import { type NestedToolSummary, PROGRESS_COMMIT, removeNestedSlots, + salvageable, type ToolSlot, toolSlot, } from "./live.ts"; @@ -141,12 +142,11 @@ export const ToolTask = defineTask error, context); + return reject(runtime, task.input, call, undefined, error, context); } const prepared = prepare(tool, call.arguments as JsonObject); const checked = "error" in prepared ? prepared : validate(tool, call, prepared.args); - if ("error" in checked) - return settle(runtime, task.input, call, COMPLETED, () => invalid(checked.error), context); + if ("error" in checked) return reject(runtime, task.input, call, tool, invalid(checked.error), context); let args = checked.args; let block: string | undefined; await runtime.hooks.each("beforeTool", async (hook) => { @@ -162,11 +162,10 @@ export const ToolTask = defineTask blocked, context); + return reject(runtime, task.input, call, tool, blocked, context); } const validated = validate(tool, call, args); - if ("error" in validated) - return settle(runtime, task.input, call, COMPLETED, () => invalid(validated.error), context); + if ("error" in validated) return reject(runtime, task.input, call, tool, invalid(validated.error), context); // Strict JSON: repair may set optional properties to undefined, which the slot and checkpoint cannot hold. const final = copyJson(validated.args, { omitUndefinedProperties: true }) as JsonObject; await runtime.commit(async (tx) => { @@ -185,9 +184,10 @@ export const ToolTask = defineTask { const { arguments: args, replay } = task.state.checkpoint; const call = await readCall(runtime, task.input, context); + // Also before an interruption result: its `afterTool` hooks may be what cut the task off. + if (cutOffTooOften(task, runtime)) return endCutOff(runtime, task, call, context); const tool = await resolveTool(runtime, task.input, call.name, context); if (replay === "safe" && tool?.replay === "safe") { - if (cutOffTooOften(task, runtime)) return endCutOff(runtime, task, call, context); // The rerun reports from scratch; clear what the interrupted attempt published. await runtime.commit(async (tx) => { const slot = toolSlot(await tx.doc(LiveDoc, runtime.conversationId), runtime.taskId); @@ -199,7 +199,8 @@ export const ToolTask = defineTask fromSlot(slot, "interrupted", message), context); + // Not cut off too often (checked above), so the hooks run. + await settleFromSlot(runtime, task.input, call, { tool }, ending, "interrupted", message, context); }, }, abort: async (task, runtime, context) => { @@ -211,8 +212,13 @@ export const ToolTask = defineTask fromSlot(slot, code, message), context); + // The scheduler runs an abort handler only while its restarts are below the limit, and faults the task without a + // result when the last attempt is cut off. That attempt skips the hooks, which may be what cuts it off. + const hooks = + (task.restarts ?? 0) < runtime.settings.restartLimit - 1 + ? { tool: await toolForLimits(runtime, task.input, call.name, context) } + : undefined; + await settleFromSlot(runtime, task.input, call, hooks, { status: "aborted" }, code, message, context); }, }); @@ -228,7 +234,68 @@ async function endCutOff( ): Promise { const message = `Tool ${call.name} was interrupted ${cutOffCount(task)}; it is not run again`; const ending = { status: "failed", message, reason: "restart_limit" } as const; - await settle(runtime, task.input, call, ending, (slot) => fromSlot(slot, "restart_limit", message), context); + // Past the limit the hooks may be what keeps cutting the task off (`salvageable`). + const hooks = salvageable(task, runtime) + ? { tool: await toolForLimits(runtime, task.input, call.name, context) } + : undefined; + await settleFromSlot(runtime, task.input, call, hooks, ending, "restart_limit", message, context); +} + +/** + * Settle a result the Harness decided in the `call` phase, without running the tool: the tool did not resolve, the + * arguments are invalid, or `beforeTool` blocked the call. `afterTool` applies as to an executed call's result. + */ +async function reject( + runtime: Runtime, + input: ToolTaskInput, + call: ToolHookCall, + tool: ToolRegistration | undefined, + result: ToolExecutionResult, + context: Context, +): Promise { + const settled = await afterHarnessResult(runtime, input, call, tool, result, undefined, context); + await settle(runtime, input, call, COMPLETED, settled, context); +} + +/** + * Settle an error result built from the call's slot (`fromSlot`): an interruption, the restart limit, or an abort. + * The tool is not running, so the slot read before the commit is still current when it commits. With `hooks`, the + * result passes `afterTool` and is bounded with the limits of `hooks.tool` (the defaults when it is `undefined`); + * without, it is settled as built, when the hooks may be what cuts the task off. + */ +async function settleFromSlot( + runtime: Runtime, + input: ToolTaskInput, + call: ToolHookCall, + hooks: { readonly tool: ToolRegistration | undefined } | undefined, + ending: Ending, + code: string, + message: string, + context: Context, +): Promise { + const live = await runtime.snapshot(LiveDoc, runtime.conversationId, context); + const { result, truncation } = fromSlot( + live === undefined ? undefined : toolSlot(live, runtime.taskId), + code, + message, + ); + const settled = + hooks !== undefined + ? await afterHarnessResult(runtime, input, call, hooks.tool, result, truncation, context) + : truncation === undefined + ? result + : { ...result, diagnostics: [...(result.diagnostics ?? []), truncation] }; + await settle(runtime, input, call, ending, settled, context); +} + +/** The tool whose output limits bound a hooks' result: the called tool, or none when it, or the agent, fails to resolve. */ +async function toolForLimits( + runtime: Runtime, + input: ToolTaskInput, + name: string, + context: Context, +): Promise { + return resolveTool(runtime, input, name, context).catch(() => undefined); } /** The called tool: among the tools the model is offered for a model-issued call, else among those tools may call. */ @@ -242,6 +309,15 @@ async function resolveTool( return (input.kind === "model" ? agent.tools : agent.callable).find((tool) => tool.name === name); } +/** The tool's output limits, the defaults filled in; all defaults when the tool did not resolve. */ +function outputLimits(tool: ToolRegistration | undefined): OutputLimits { + return { + maxBytes: tool?.outputLimits?.maxBytes ?? DEFAULT_MAX_BYTES, + maxLines: tool?.outputLimits?.maxLines ?? DEFAULT_MAX_LINES, + retain: tool?.outputLimits?.retain ?? "head", + }; +} + /** The call: from the input for a nested call, which hooks see with its parent, else from the assistant entry. */ async function readCall(runtime: Runtime, input: ToolTaskInput, context: Context): Promise { if (input.kind === "nested") { @@ -310,11 +386,7 @@ async function run( args: JsonObject, context: Context, ): Promise { - const limits: OutputLimits = { - maxBytes: tool.outputLimits?.maxBytes ?? DEFAULT_MAX_BYTES, - maxLines: tool.outputLimits?.maxLines ?? DEFAULT_MAX_LINES, - retain: tool.outputLimits?.retain ?? "head", - }; + const limits = outputLimits(tool); const reported: Reported = { output: new OutputBuffer(limits), limits, diagnostics: [], details: undefined }; const streams = input.kind === "model" || input.progress !== false; const progress = publishProgress(runtime, reported, streams, context); @@ -473,7 +545,7 @@ async function run( const pending = await progress.stop(); try { const settled = await finalResult(runtime, input, call, tool, result, reported, context); - await settle(runtime, input, call, ending, () => settled, context, durationMs); + await settle(runtime, input, call, ending, settled, context, durationMs); } catch (error) { for (const waiter of pending) waiter.reject(error); throw error; @@ -630,6 +702,41 @@ async function finalResult( return { ...final, output: bounded.content, diagnostics: [...(final.diagnostics ?? []), ...harness] }; } +/** + * A result the Harness wrote itself, without executing the tool, after the `afterTool` chain, normalized like an + * executed call's: the output bounded with the tool's limits (the defaults when it did not resolve), a truncation + * diagnostic when that dropped something, and the hooks' `structuredOutput` dropped, since the tool never produced + * one. A nested call that the hooks turned into a success gets the bounded output as its value, as a schema-less + * tool's success does, unless the tool declares a schema. `usage` is kept and recorded; `control` applies only when + * the task completes. `truncation` describes what the result's own output lost before (a slot's retained output); it + * is kept only while the hooks keep that output. + */ +async function afterHarnessResult( + runtime: Runtime, + input: ToolTaskInput, + call: ToolHookCall, + tool: ToolRegistration | undefined, + result: ToolExecutionResult, + truncation: ToolDiagnostic | undefined, + context: Context, +): Promise { + let final = result; + await runtime.hooks.each("afterTool", async (hook) => { + final = (await hook(call, final, runtime, context)) ?? final; + }); + const harness: ToolDiagnostic[] = []; + if (truncation !== undefined && final.output === result.output) harness.push(truncation); + const limits = outputLimits(tool); + const bounded = boundContent(final.output ?? [], limits); + if (bounded.droppedBytes > 0) harness.push(truncated(bounded, limits.retain)); + const { structuredOutput: _dropped, ...rest } = final; + const value = + input.kind === "nested" && final.isError !== true && tool?.structuredOutputSchema === undefined + ? { structuredOutput: outputValue(bounded.content) as JsonValue } + : {}; + return { ...rest, ...value, output: bounded.content, diagnostics: [...(rest.diagnostics ?? []), ...harness] }; +} + const structuredValidators = new WeakMap>(); /** Why a result's `structuredOutput` breaks the tool's contract, or `undefined` when it keeps it. */ @@ -659,15 +766,14 @@ 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 `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. + * slots below the call leave `pi.live`. `control` applies only when the task completes. */ async function settle( runtime: Runtime, input: ToolTaskInput, call: ToolHookCall, ending: Ending, - build: (slot: Readonly | undefined) => ToolExecutionResult, + result: ToolExecutionResult, context: Context, durationMs?: number, ): Promise { @@ -684,7 +790,6 @@ async function settle( const conversationId = runtime.conversationId; const live = await tx.doc(LiveDoc, conversationId); const slot = toolSlot(live, runtime.taskId); - const result = build(slot); let settled: ToolTaskResult; if (input.kind === "nested") { const nested = nestedResult(runtime.taskId, result, durationMs); @@ -854,22 +959,24 @@ type Ending = const COMPLETED: Ending = { status: "completed" }; -/** An error result from the slot's durable partial output, details, and diagnostics. */ +/** + * An error result from the slot's durable partial output, details, and diagnostics, ending with `code`; and the + * truncation diagnostic of that output, when its bounds dropped some, which the caller adds after `afterTool`. + */ function fromSlot( slot: Readonly | undefined, code: string, message: string, -): ToolExecutionResult { - const diagnostics = [...(slot?.diagnostics ?? [])]; +): { readonly result: ToolExecutionResult; readonly truncation?: ToolDiagnostic } { const droppedBytes = slot?.droppedBytes ?? 0; - if (droppedBytes > 0) diagnostics.push(truncated({ droppedBytes, droppedLines: slot?.droppedLines ?? 0 })); - diagnostics.push(toolDiagnostic(code, message)); - return { + const result: ToolExecutionResult = { output: slot?.output === undefined || slot.output === "" ? [] : [{ type: "text", text: slot.output }], isError: true, ...(slot?.details === undefined ? {} : { details: slot.details }), - diagnostics, + diagnostics: [...(slot?.diagnostics ?? []), toolDiagnostic(code, message)], }; + if (droppedBytes === 0) return { result }; + return { result, truncation: truncated({ droppedBytes, droppedLines: slot?.droppedLines ?? 0 }) }; } /** An error result the Harness writes itself: no content and one `error` diagnostic with `code`. */ diff --git a/packages/durable/src/harness/types.ts b/packages/durable/src/harness/types.ts index d787dd3cd..272af94f4 100644 --- a/packages/durable/src/harness/types.ts +++ b/packages/durable/src/harness/types.ts @@ -917,7 +917,18 @@ export interface ToolHooks { api: HookApi, context: Context, ): HookResult<{ readonly arguments?: JsonObject; readonly block?: string }>; - /** After execution, before the result is committed; replaces the result. */ + /** + * Before the result is committed; replaces the result. Runs on every result the tool task writes: an executed + * call's, and the error results the Harness writes itself, told apart by their diagnostic code: a tool that does not + * resolve (`tool_unavailable`), invalid arguments (`invalid_arguments`, with `call.arguments` as the model sent + * them), a block (`blocked`), and, rebuilt from the call's partial output, an interruption (`interrupted`), the + * restart limit (`restart_limit`), and an abort (`aborted`, `abandoned`). Once the task was cut off more than + * `settings.restartLimit` times the hooks no longer run, since they may be what cuts it off, nor in the last attempt + * the scheduler gives an abort handler (restarts at the limit minus one). A result's `structuredOutput` is dropped + * when the Harness wrote it (a nested call a hook turns into a success gets the bounded output as a schema-less tool's + * does), and its `control` applies only to a call that completes (not to an interrupted, cut-off, or aborted one). + * An abort's hooks run inside the abort handler, which cannot be aborted: a slow hook delays `Conversation.abort()`. + */ afterTool( call: ToolHookCall, result: ToolExecutionResult, diff --git a/packages/durable/test/harness-tool-hooks.test.ts b/packages/durable/test/harness-tool-hooks.test.ts new file mode 100644 index 000000000..aeecb2e1f --- /dev/null +++ b/packages/durable/test/harness-tool-hooks.test.ts @@ -0,0 +1,626 @@ +// `afterTool` runs on every result a tool task writes, also those the Harness decides without executing the tool +// (#11043). +import { mkdtemp, rm } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import type { JsonValue } from "@earendil-works/chord"; +import { + type AssistantMessage, + fauxAssistantMessage, + fauxText, + fauxToolCall, + type ToolResultMessage, + Type, +} from "@earendil-works/pi-ai"; +import { + defineTool, + type EntryRecord, + type Harness, + LiveDoc, + MemoryStorage, + type NestedToolExecutionResult, + ROOT_CONVERSATION_ID, + type Storage, + type TaskId, + type ToolExecutionResult, + type ToolHookCall, + type ToolRegistration, + ToolResultEntry, + ToolTask, +} from "@earendil-works/pi-durable"; +import { afterEach, describe, expect, it } from "vitest"; +import { openNodeSqliteStorage } from "../src/storage/sqlite/node.ts"; +import { allEntries, type ChatSetup, chatSetup, openChat, waitFor } from "./chat-support.ts"; +import { addHooks, addTool } from "./harness-support.ts"; +import { context } from "./session-support.ts"; +import { aborted, deferred } from "./task-support.ts"; + +const directories = new Set(); +/** Storages of killed Harnesses, never closed by them: closed after each test. */ +const killed = new Set(); + +afterEach(async () => { + for (const storage of killed) await storage.close(context).catch(() => {}); + killed.clear(); + for (const directory of directories) await rm(directory, { recursive: true, force: true }); + directories.clear(); +}); + +async function sqlitePath(): Promise { + const directory = await mkdtemp(join(tmpdir(), "pi-durable-tool-hooks-")); + directories.add(directory); + return join(directory, "session.sqlite"); +} + +/** Open a chat over the SQLite file at `path`; `kill()` simulates a crash: every later storage call hangs. */ +async function open(path: string, setup: ChatSetup) { + const inner = await openNodeSqliteStorage(path); + let dead = false; + const storage = new Proxy(inner, { + get(target, property) { + const value = Reflect.get(target, property, target); + if (typeof value !== "function") return value; + return (...args: unknown[]) => (dead ? new Promise(() => {}) : value.apply(target, args)); + }, + }); + const opened = await openChat(storage, setup); + opened.harness.resume(); + const kill = (): void => { + dead = true; + killed.add(inner); + }; + return { ...opened, kill }; +} + +/** Let deferred progress commits (`persistWithinMs: 0`) reach Storage before a crash. */ +function stored(): Promise { + return new Promise((resolve) => setTimeout(resolve, 50)); +} + +const Parameters = Type.Object({ text: Type.Optional(Type.String()) }); + +function tool( + name: string, + execute: ToolRegistration["execute"], + extra: Partial> = {}, +): ToolRegistration { + return defineTool({ name, description: name, parameters: Parameters, execute, ...extra }); +} + +function calls(...list: readonly (readonly [string, Record, string])[]): AssistantMessage { + return fauxAssistantMessage( + list.map(([name, args, id]) => fauxToolCall(name, args, { id })), + { stopReason: "toolUse" }, + ); +} + +const DONE = fauxAssistantMessage([fauxText("done")]); + +function results(entries: readonly EntryRecord[]): Map { + const messages = entries + .filter((entry) => ToolResultEntry.is(entry)) + .map((entry) => entry.model![0] as ToolResultMessage); + return new Map(messages.map((message) => [message.toolCallId, message])); +} + +function text(message: ToolResultMessage | undefined): string { + return (message?.content ?? []).map((item) => (item.type === "text" ? item.text : "")).join("|"); +} + +function codes(result: { readonly diagnostics?: readonly { readonly code?: string }[] }): (string | undefined)[] { + return (result.diagnostics ?? []).map((diagnostic) => diagnostic.code); +} + +async function diagnosticsOf(entries: readonly EntryRecord[], callId: string): Promise<(string | undefined)[]> { + const entry = entries.find( + (each) => ToolResultEntry.is(each) && (each.model![0] as ToolResultMessage).toolCallId === callId, + ); + return codes((entry?.data ?? {}) as { diagnostics?: { code?: string }[] }); +} + +type Seen = { readonly call: ToolHookCall; readonly result: ToolExecutionResult }; + +async function toolTaskId(harness: Harness, index = 0): Promise { + let id: TaskId | undefined; + await waitFor(async () => { + id = (await harness.snapshot(LiveDoc, ROOT_CONVERSATION_ID, context))?.tools?.[index]?.taskId; + return id !== undefined; + }); + return id!; +} + +describe("afterTool on results the call phase decides", () => { + // #11043 + it("sees and replaces unavailable, invalid, and blocked results, with the raw arguments", async () => { + const setup = chatSetup(); + let ran = 0; + const echo = tool("echo", async () => { + ran++; + return { output: [{ type: "text", text: "ran" }] }; + }); + addTool(setup.registry, echo); + addTool( + setup.registry, + tool("gone", async () => ({ output: [] })), + ); + const seen: Seen[] = []; + addHooks(setup.registry, ToolTask, { + beforeTool: (call) => { + if (call.id === "block") return { block: "not today" }; + if (call.id === "rewrite") return { arguments: { text: { not: "a string" } } }; + return undefined; + }, + afterTool: (call, result) => { + seen.push({ call, result }); + return { ...result, output: [{ type: "text", text: `hooked ${call.id}` }], details: { by: "hook" } }; + }, + }); + let root!: Awaited>["root"]; + setup.faux.setResponses([ + async () => { + // The tool is offered but no longer resolves when its task runs. + await root.configure({ tools: [echo] }, context); + return calls( + ["gone", {}, "unavailable"], + ["echo", { text: { not: "a string" } }, "invalid"], + ["echo", { text: "x" }, "rewrite"], + ["echo", {}, "block"], + ); + }, + DONE, + ]); + const opened = await openChat(new MemoryStorage(), setup); + root = opened.root; + const settled = await (await root.submit({ type: "input", content: "go" }, context)).wait(context); + expect(settled.status).toBe("done"); + expect(ran).toBe(0); + const byId = new Map(seen.map((each) => [each.call.id, each])); + expect([...byId.keys()].sort()).toEqual(["block", "invalid", "rewrite", "unavailable"]); + expect(codes(byId.get("unavailable")!.result)).toEqual(["tool_unavailable"]); + expect(codes(byId.get("invalid")!.result)).toEqual(["invalid_arguments"]); + expect(codes(byId.get("rewrite")!.result)).toEqual(["invalid_arguments"]); + expect(codes(byId.get("block")!.result)).toEqual(["blocked"]); + // Hooks see the call as the model made it, also when it is invalid or `beforeTool` rewrote it. + expect(byId.get("invalid")!.call.arguments).toEqual({ text: { not: "a string" } }); + expect(byId.get("rewrite")!.call.arguments).toEqual({ text: "x" }); + const entries = await allEntries(root); + const messages = results(entries); + for (const id of ["unavailable", "invalid", "rewrite", "block"]) { + const message = messages.get(id)!; + expect(message).toMatchObject({ isError: true, details: { by: "hook" } }); + expect(text(message).startsWith(`hooked ${id}|`)).toBe(true); + } + await opened.harness.close(context); + }); + + // #11043 + it("bounds a replacement output with the tool's limits, or the defaults for an unresolved tool", async () => { + const setup = chatSetup(); + addTool( + setup.registry, + tool("small", async () => ({ output: [] }), { outputLimits: { maxBytes: 10 } }), + ); + addHooks(setup.registry, ToolTask, { + beforeTool: () => ({ block: "no" }), + afterTool: (call, result) => ({ + ...result, + output: [{ type: "text", text: call.id === "c2" ? "x".repeat(100) : "y".repeat(60 * 1024) }], + }), + }); + let root!: Awaited>["root"]; + setup.faux.setResponses([ + async () => { + await root.configure({ tools: [] }, context); + return calls(["small", {}, "c1"]); + }, + DONE, + ]); + const opened = await openChat(new MemoryStorage(), setup); + root = opened.root; + await (await root.submit({ type: "input", content: "go" }, context)).wait(context); + await root.configure({ tools: null }, context); + setup.faux.setResponses([calls(["small", {}, "c2"]), DONE]); + await (await root.submit({ type: "input", content: "again" }, context)).wait(context); + const entries = await allEntries(root); + const messages = results(entries); + // Unresolved: the default 50 KB. + const unresolved = messages.get("c1")!.content[0]; + expect(unresolved?.type === "text" && unresolved.text.length).toBe(50 * 1024); + expect(await diagnosticsOf(entries, "c1")).toEqual(["tool_unavailable", "truncated"]); + const bounded = messages.get("c2")!.content[0]; + expect(bounded?.type === "text" && bounded.text).toBe("x".repeat(10)); + expect(await diagnosticsOf(entries, "c2")).toEqual(["blocked", "truncated"]); + await opened.harness.close(context); + }); + + // #11043 + it("lets a hook's terminate on a blocked call count toward the round", async () => { + const setup = chatSetup(); + addTool( + setup.registry, + tool("echo", async () => ({ output: [], control: { terminate: true } })), + ); + addHooks(setup.registry, ToolTask, { + beforeTool: (call) => (call.id === "c2" ? { block: "no" } : undefined), + afterTool: (_call, result) => ({ ...result, control: { terminate: true } }), + }); + setup.faux.setResponses([calls(["echo", {}, "c1"], ["echo", {}, "c2"]), DONE]); + const { harness, root } = await openChat(new MemoryStorage(), setup); + const settled = await (await root.submit({ type: "input", content: "go" }, context)).wait(context); + expect(settled.status).toBe("done"); + // The round terminated: no second request. + expect(setup.faux.state.callCount).toBe(1); + await harness.close(context); + }); + + // #11043 + it("runs on nested invalid and blocked calls; the caller gets diagnostics and details, not output", async () => { + const setup = chatSetup(); + addTool( + setup.registry, + tool("echo", async () => ({ output: [] })), + ); + const nested: NestedToolExecutionResult[] = []; + addTool( + setup.registry, + tool("batch", async (_args, api, callContext) => { + nested.push(await api.executeTool("echo", { text: { not: "a string" } }, callContext)); + nested.push(await api.executeTool("echo", { text: "secret" }, callContext)); + nested.push(await api.executeTool("missing", {}, callContext)); + return {}; + }), + ); + const seen: ToolHookCall[] = []; + addHooks(setup.registry, ToolTask, { + beforeTool: (call) => (call.arguments.text === "secret" ? { block: "no secrets" } : undefined), + afterTool: (call, result) => { + if (call.parent === undefined) return undefined; + seen.push(call); + return { + ...result, + output: [{ type: "text", text: "dropped" }], + structuredOutput: { dropped: true }, + details: { hooked: call.id }, + diagnostics: [...(result.diagnostics ?? []), { severity: "info", code: "hooked", message: "seen" }], + }; + }, + }); + setup.faux.setResponses([calls(["batch", {}, "c1"]), DONE]); + const { harness, root } = await openChat(new MemoryStorage(), setup); + await (await root.submit({ type: "input", content: "go" }, context)).wait(context); + expect(seen.map((call) => call.id)).toEqual(["c1/1", "c1/2", "c1/3"]); + expect(seen.every((call) => call.parent?.callId === "c1")).toBe(true); + expect(nested.map((result) => [result.isError, codes(result)])).toEqual([ + [true, ["invalid_arguments", "hooked"]], + [true, ["blocked", "hooked"]], + [true, ["tool_unavailable", "hooked"]], + ]); + expect(nested.map((result) => result.details)).toEqual([ + { hooked: "c1/1" }, + { hooked: "c1/2" }, + { hooked: "c1/3" }, + ]); + // A Harness error result has no structured output, even when a hook supplies one. + expect(nested[1]!.structuredOutput).toBeUndefined(); + expect(nested[2]!.structuredOutput).toBeUndefined(); + await harness.close(context); + }); + + // #11043 + it("gives a nested call a hook turns into a success its output as value, unless the tool has a schema", async () => { + const setup = chatSetup(); + addTool( + setup.registry, + tool("plain", async () => ({ output: [] })), + ); + addTool( + setup.registry, + tool("typed", async () => ({ output: [], structuredOutput: { ok: true } }), { + structuredOutputSchema: Type.Object({ ok: Type.Boolean() }), + }), + ); + const nested: NestedToolExecutionResult[] = []; + addTool( + setup.registry, + tool("batch", async (_args, api, callContext) => { + nested.push(await api.executeTool("plain", {}, callContext)); + nested.push(await api.executeTool("typed", {}, callContext)); + return {}; + }), + ); + addHooks(setup.registry, ToolTask, { + beforeTool: (call) => (call.parent === undefined ? undefined : { block: "no" }), + afterTool: (call) => + call.parent === undefined + ? undefined + : { output: [{ type: "text", text: `fine ${call.name}` }], structuredOutput: { ok: false } }, + }); + setup.faux.setResponses([calls(["batch", {}, "c1"]), DONE]); + const { harness, root } = await openChat(new MemoryStorage(), setup); + await (await root.submit({ type: "input", content: "go" }, context)).wait(context); + expect(nested.map((result) => [result.isError, result.structuredOutput])).toEqual([ + [false, "fine plain"], + [false, undefined], + ]); + await harness.close(context); + }); +}); + +describe("afterTool on results rebuilt from the slot", () => { + /** A tool that writes output, then blocks until its invocation is cancelled. */ + function blocking(name: string, extra: Partial> = {}) { + const started = deferred(); + const state = { runs: 0 }; + const registration = tool( + name, + async (_args, api, callContext) => { + state.runs++; + api.output(`${name} ${"o".repeat(40)}\n`); + await api.details({ run: state.runs }, callContext); + started.resolve(); + return aborted(callContext.abortSignal!); + }, + extra, + ); + return { registration, started, state }; + } + + // #11043 + it("runs on an interrupted call; a replaced output loses the slot's truncation diagnostic", async () => { + const path = await sqlitePath(); + const setup = chatSetup(); + setup.settings.progress = { outputIntervalMs: 0, persistWithinMs: 0 }; + const limits = { outputLimits: { maxBytes: 10, retain: "tail" } } as const; + const work = blocking("work", limits); + const keep = blocking("keep", limits); + addTool(setup.registry, work.registration); + addTool(setup.registry, keep.registration); + const seen: Seen[] = []; + addHooks(setup.registry, ToolTask, { + afterTool: (call, result) => { + seen.push({ call, result }); + return call.name === "work" ? { ...result, output: [{ type: "text", text: "redacted" }] } : undefined; + }, + }); + setup.faux.setResponses([calls(["work", {}, "c1"], ["keep", {}, "c2"]), DONE]); + let opened = await open(path, setup); + const id = (await opened.root.submit({ type: "input", content: "go" }, context)).id; + await work.started.promise; + await keep.started.promise; + await stored(); + opened.kill(); + + opened = await open(path, setup); + expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done"); + expect(seen.map((each) => [each.call.id, codes(each.result)]).sort()).toEqual([ + ["c1", ["interrupted"]], + ["c2", ["interrupted"]], + ]); + const entries = await allEntries(opened.root); + const messages = results(entries); + expect(text(messages.get("c1")).startsWith("redacted|")).toBe(true); + expect(messages.get("c1")).toMatchObject({ isError: true, details: { run: 1 } }); + expect(await diagnosticsOf(entries, "c1")).toEqual(["interrupted"]); + // The kept output keeps its truncation diagnostic, after the hooks' result. + expect(await diagnosticsOf(entries, "c2")).toEqual(["interrupted", "truncated"]); + await opened.harness.close(context); + }); + + // #11043 + it("runs on an abort result, but not on the aborted results of calls that never started", async () => { + const setup = chatSetup(); + setup.settings.toolExecution = "sequential"; + const first = blocking("first"); + addTool(setup.registry, first.registration); + addTool( + setup.registry, + tool("second", async () => ({ output: [] })), + ); + const seen: Seen[] = []; + addHooks(setup.registry, ToolTask, { + afterTool: (call, result) => { + seen.push({ call, result }); + return { ...result, details: { hooked: true } }; + }, + }); + setup.faux.setResponses([calls(["first", {}, "c1"], ["second", {}, "c2"]), DONE]); + const { harness, root } = await openChat(new MemoryStorage(), setup); + const submission = await root.submit({ type: "input", content: "go" }, context); + await first.started.promise; + await root.abort(context); + expect((await submission.wait(context)).status).toBe("unanswered"); + expect(seen.map((each) => [each.call.id, codes(each.result)])).toEqual([["c1", ["aborted"]]]); + const messages = results(await allEntries(root)); + expect(messages.get("c1")).toMatchObject({ isError: true, details: { hooked: true } }); + expect(messages.get("c2")?.details).toBeUndefined(); + await harness.close(context); + }); + + // #11043 + it("reports a throwing hook on an abort result and settles the result as built", async () => { + const setup = chatSetup(); + const work = blocking("work"); + addTool(setup.registry, work.registration); + addHooks(setup.registry, ToolTask, { + afterTool: () => { + throw new Error("hook failed"); + }, + }); + setup.faux.setResponses([calls(["work", {}, "c1"]), DONE]); + const { harness, root } = await openChat(new MemoryStorage(), setup); + const submission = await root.submit({ type: "input", content: "go" }, context); + await work.started.promise; + await root.abort(context); + expect((await submission.wait(context)).status).toBe("unanswered"); + const entries = await allEntries(root); + expect(text(results(entries).get("c1")).startsWith(`work ${"o".repeat(40)}`)).toBe(true); + expect(await diagnosticsOf(entries, "c1")).toEqual(["aborted"]); + expect(setup.reports.map((error) => (error as Error).message)).toContain("hook failed"); + await harness.close(context); + }); + + // #11043 + it("settles an abort without the hooks in the abort handler's last attempt", async () => { + const path = await sqlitePath(); + const setup = chatSetup(); + setup.settings.restartLimit = 2; + setup.settings.progress = { outputIntervalMs: 0, persistWithinMs: 0 }; + const work = blocking("work"); + addTool(setup.registry, work.registration); + const seen: Seen[] = []; + const hanging = deferred(); + addHooks(setup.registry, ToolTask, { + afterTool: async (call, result) => { + seen.push({ call, result }); + hanging.resolve(); + // Never returns: the host dies inside the hook. + return new Promise(() => {}); + }, + }); + setup.faux.setResponses([calls(["work", {}, "c1"]), DONE]); + let opened = await open(path, setup); + const id = (await opened.root.submit({ type: "input", content: "go" }, context)).id; + await work.started.promise; + const taskId = await toolTaskId(opened.harness); + await stored(); + void opened.root.abort(context).catch(() => {}); + await hanging.promise; + opened.kill(); + + // One restart, at the limit minus one: the attempt the scheduler would fault when cut off runs no hooks. + opened = await open(path, setup); + expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("unanswered"); + expect(seen.map((each) => codes(each.result))).toEqual([["aborted"]]); + const entries = await allEntries(opened.root); + expect(text(results(entries).get("c1")).startsWith(`work ${"o".repeat(40)}`)).toBe(true); + expect(await diagnosticsOf(entries, "c1")).toEqual(["aborted"]); + expect((await opened.harness.getTask(taskId, context))?.state).toMatchObject({ + status: "terminal", + outcome: { status: "aborted" }, + }); + await opened.harness.close(context); + }); + + // #11043 + it("sees the held progress of an aborted call, and the result stores it", async () => { + const path = await sqlitePath(); + const setup = chatSetup(); + // Progress commits are held in memory until the result commit stores them. + setup.settings.progress = { outputIntervalMs: 0, persistWithinMs: 60_000 }; + const work = blocking("work"); + addTool(setup.registry, work.registration); + const seen: Seen[] = []; + addHooks(setup.registry, ToolTask, { + afterTool: (call, result) => { + seen.push({ call, result }); + return undefined; + }, + }); + setup.faux.setResponses([calls(["work", {}, "c1"]), DONE]); + let opened = await open(path, setup); + const submission = await opened.root.submit({ type: "input", content: "go" }, context); + await work.started.promise; + await opened.root.abort(context); + expect((await submission.wait(context)).status).toBe("unanswered"); + expect(seen).toHaveLength(1); + expect(seen[0]!.result).toMatchObject({ details: { run: 1 } }); + expect(text({ content: seen[0]!.result.output } as ToolResultMessage)).toBe(`work ${"o".repeat(40)}\n`); + await opened.harness.close(context); + + opened = await open(path, setup); + const message = results(await allEntries(opened.root)).get("c1"); + expect(message).toMatchObject({ isError: true, details: { run: 1 } }); + expect(text(message).startsWith(`work ${"o".repeat(40)}`)).toBe(true); + await opened.harness.close(context); + }); + + // #11043 + it("runs on an abandoned nested call's result", async () => { + const path = await sqlitePath(); + const setup = chatSetup(); + const gated = deferred(); + let gate = true; + addTool( + setup.registry, + tool("gated", async () => ({ output: [] })), + ); + addTool( + setup.registry, + tool("batch", async (_args, api, callContext) => { + await api.executeTool("gated", {}, callContext); + return {}; + }), + ); + const seen: Seen[] = []; + addHooks(setup.registry, ToolTask, { + // Holds the nested call in its `call` phase, before intent, until the Harness closes. + beforeTool: async (call, _api, hookContext) => { + if (call.name !== "gated" || !gate) return undefined; + gated.resolve(); + return aborted(hookContext.abortSignal!); + }, + afterTool: (call, result) => { + if (call.parent === undefined) return undefined; + seen.push({ call, result }); + return { ...result, details: { hooked: true } }; + }, + }); + setup.faux.setResponses([calls(["batch", {}, "c1"]), DONE]); + let opened = await open(path, setup); + const id = (await opened.root.submit({ type: "input", content: "go" }, context)).id; + await gated.promise; + await opened.harness.close(context); + gate = false; + + opened = await open(path, setup); + expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done"); + expect(seen.map((each) => [each.call.id, codes(each.result)])).toEqual([["c1/1", ["abandoned"]]]); + await opened.harness.close(context); + }); + + // #11043 + it("runs on a restart_limit result at the limit, and not once the hooks themselves were cut off", async () => { + const path = await sqlitePath(); + const setup = chatSetup(); + setup.settings.restartLimit = 1; + setup.settings.progress = { outputIntervalMs: 0, persistWithinMs: 0 }; + const work = blocking("work"); + addTool(setup.registry, work.registration); + const seen: Seen[] = []; + const hanging = deferred(); + addHooks(setup.registry, ToolTask, { + afterTool: async (call, result) => { + seen.push({ call, result }); + hanging.resolve(); + // Never returns: the host dies inside the hook. + return new Promise(() => {}); + }, + }); + setup.faux.setResponses([calls(["work", {}, "c1"]), DONE]); + let opened = await open(path, setup); + const id = (await opened.root.submit({ type: "input", content: "go" }, context)).id; + await work.started.promise; + const taskId = await toolTaskId(opened.harness); + await stored(); + opened.kill(); + + // Cut off once, at the limit: the task ends itself, and the hooks run on its result. + opened = await open(path, setup); + await hanging.promise; + expect(seen.map((each) => codes(each.result))).toEqual([["restart_limit"]]); + opened.kill(); + + // Past the limit: ends without the hooks. + opened = await open(path, setup); + expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done"); + expect(seen).toHaveLength(1); + expect(work.state.runs).toBe(1); + const entries = await allEntries(opened.root); + expect(await diagnosticsOf(entries, "c1")).toEqual(["restart_limit"]); + expect((await opened.harness.getTask(taskId, context))?.state).toMatchObject({ + status: "terminal", + outcome: { status: "failed", error: { detail: { reason: "restart_limit" } } }, + }); + await opened.harness.close(context); + }); +}); diff --git a/packages/durable/test/harness-tools.test.ts b/packages/durable/test/harness-tools.test.ts index 29f6fd3f7..613975d52 100644 --- a/packages/durable/test/harness-tools.test.ts +++ b/packages/durable/test/harness-tools.test.ts @@ -1168,7 +1168,9 @@ describe("tool progress and lifetime", () => { }), ); addHooks(setup.registry, ToolTask, { - afterTool: async (_call, _result, _api, callContext) => { + afterTool: async (_call, result, _api, callContext) => { + // The abort handler's result runs the hooks too (#11043); only the executed call's waits. + if (result.diagnostics?.some((diagnostic) => diagnostic.code === "aborted")) return undefined; inAfterTool.resolve(); await aborted(callContext.abortSignal!); return undefined;