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

feat(durable): afterTool runs on every tool task result

Browse source
This commit is contained in:
Mario Zechner 2026-10-28 11:45:03 +00:00
parent 47dba56f92
commit 87eff524f9
7 changed files with 817 additions and 37 deletions
Download patch file Download diff file Expand all files Collapse all files

1
packages/durable/CHANGELOG.md
View file

@ -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

4
packages/durable/README.md
View file

@ -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] } })`.

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

@ -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

163
packages/durable/src/harness/tool.ts
View file

@ -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<ToolTaskInput, ToolTaskCheckpoint, ToolTaskRe
const tool = await resolveTool(runtime, task.input, call.name, context);
if (tool === undefined) {
const error = harnessError("tool_unavailable", `Tool ${call.name} is not available`);
return settle(runtime, task.input, call, COMPLETED, () => 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<ToolTaskInput, ToolTaskCheckpoint, ToolTaskRe
});
if (block !== undefined) {
const blocked = harnessError("blocked", `Tool call blocked: ${block}`);
return settle(runtime, task.input, call, COMPLETED, () => 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<ToolTaskInput, ToolTaskCheckpoint, ToolTaskRe
execute: async (task, runtime, context) => {
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<ToolTaskInput, ToolTaskCheckpoint, ToolTaskRe
const message = `Tool ${call.name} was interrupted and may have partially run`;
// `failed` records cancellation intent, so the call's owned conversations, left unsupervised, are aborted.
const ending = { status: "failed", message } as const;
await settle(runtime, task.input, call, ending, (slot) => 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<ToolTaskInput, ToolTaskCheckpoint, ToolTaskRe
: task.state.checkpoint.phase === "execute"
? ["interrupted", `Tool ${call.name} was interrupted by a restart and may have partially run`]
: ["abandoned", `Tool ${call.name} was not started: its caller ended with a restart`];
const ending = { status: "aborted" } as const;
await settle(runtime, task.input, call, ending, (slot) => 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<void> {
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<void> {
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<void> {
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<ToolRegistration | undefined> {
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<ToolHookCall> {
if (input.kind === "nested") {
@ -310,11 +386,7 @@ async function run(
args: JsonObject,
context: Context,
): Promise<void> {
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<ToolExecutionResult> {
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<object, ReturnType<typeof Compile>>();
/** 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<ToolSlot | NestedToolSlot> | undefined) => ToolExecutionResult,
result: ToolExecutionResult,
context: Context,
durationMs?: number,
): Promise<void> {
@ -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<ToolSlot | NestedToolSlot> | 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`. */

13
packages/durable/src/harness/types.ts
View file

@ -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,

626
packages/durable/test/harness-tool-hooks.test.ts Normal file
View file

@ -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<string>();
/** Storages of killed Harnesses, never closed by them: closed after each test. */
const killed = new Set<Storage>();
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<string> {
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<void> {
return new Promise((resolve) => setTimeout(resolve, 50));
}
const Parameters = Type.Object({ text: Type.Optional(Type.String()) });
function tool(
name: string,
execute: ToolRegistration<typeof Parameters>["execute"],
extra: Partial<ToolRegistration<typeof Parameters>> = {},
): ToolRegistration<typeof Parameters> {
return defineTool({ name, description: name, parameters: Parameters, execute, ...extra });
}
function calls(...list: readonly (readonly [string, Record<string, JsonValue>, 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<string, ToolResultMessage> {
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<TaskId> {
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<ReturnType<typeof openChat>>["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}|<harness>`)).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<ReturnType<typeof openChat>>["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<ToolRegistration<typeof Parameters>> = {}) {
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<undefined>(() => {});
},
});
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<undefined>(() => {});
},
});
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);
});
});

4
packages/durable/test/harness-tools.test.ts
View file

@ -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;