mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +00:00
feat(durable): deferred tools and the tool_search extension
- ToolRegistration.deferred: not offered to the model until modelTools names it
- modelTools gains { remove?, add? }; control.addTools writes names to add, so
loaded tools follow the others in load order and a load only appends
- Agent.deferred lists deferred tools not loaded yet
- @earendil-works/pi-durable/tool-search: the coding agent's tool_search (BM25)
- the codemode section leaves deferred tools out
This commit is contained in:
parent
f11e50f72e
commit
77332f7a34
14 changed files with 448 additions and 40 deletions
|
|
@ -18,6 +18,7 @@
|
|||
- `ToolExecutionApi.executeTool(name, args, context, { key?, progress? })`: a tool calls another tool as a nested call. The nested call is its own `pi.tool` task owned by the caller, runs validation, the `ToolTask` hooks (whose `ToolHookCall` carries `parent`), output limits, and its replay policy, and returns a `NestedToolExecutionResult` (its `taskId`, `structuredOutput`, `details`, `diagnostics`, `usage`, `isError`, `durationMs`, but not the model's `output`) to the caller instead of the transcript. Each nested call is a member of `NestedCallDoc`, a task-scoped document family of the caller keyed by the call's key, holding its `taskId` and, once it ends, its `result` as returned, so it is gone from storage once the caller settles; its task's receipt stays small. Keys, by default the call's position, let a replay-safe caller's rerun reattach to its nested calls; a caller that is not replay-safe has its unfinished nested calls abandoned after a restart. Nested calls show in `pi.live.nestedTools`, with the arguments they run with and a `summary` (error, duration, usage) once done (`progress: false` keeps their running output out), and as tool events with `parentToolCallId` and `parentTaskId`; every tool event also carries its call's `taskId`. A caller aborts the nested calls it left running before it settles.
|
||||
- Structured tool output: a tool declares `structuredOutputSchema` and returns `structuredOutput`, which the Harness validates and gives to programs that call the tool; a tool without a schema gives them its bounded output, one text item as a string, one image as itself, and otherwise the content list, for errors too. `structuredOutput` is always the tool's own value: a schema tool's error result keeps the value it returned, and the text of one without becomes a `tool_error` diagnostic of the nested result, as a thrown error's message does. A nested call with invalid structured output gets an error result; a model-issued one only loses it, and the host gets a report. `api.retainedOutput()` reads the bounded output the model will see. `bash` and `powershell` return `{ output, truncated, fullOutputPath?, exitCode }`.
|
||||
- `ToolRegistration.callers` (`model`, `tools`, or both, the default) and `pi.agent` `modelTools`: who may call a tool, and which of its tools a conversation offers the model; `Agent.callable` lists the tools nested calls resolve among. `addTools` also extends `modelTools`.
|
||||
- Deferred tools: `ToolRegistration.deferred` keeps a tool from the model until `modelTools` names it, in its array or its new `add` (`{ remove?, add? }`); `control.addTools` writes names to `add`, and loaded tools follow the others in load order, so a load only appends declarations. `Agent.deferred` lists the deferred tools not loaded yet. `@earendil-works/pi-durable/tool-search` provides the coding agent's `tool_search` tool, which searches them with BM25 and loads the matches. The `codemode` section leaves deferred tools out.
|
||||
- `TaskOptions.abandonOnRestart`: a task its creator awaits only in memory is aborted, with its owned work, when a later Harness starts scheduling, before any of it runs again. Its mark carries `abortReason: "restart"`, which cascades; such a task waits for a missing definition instead of becoming `orphaned`. Nested calls and `api.createTask()` children of tools that are not replay-safe get it by default. No task below an owner with cancellation intent starts a run phase any more, even before the cascade marks it.
|
||||
- `TaskRuntime.abortOwned(id, context)`: abort a task this task owns and wait until it is terminal.
|
||||
- `TaskRuntime.ownedTasks(context)`: the live tasks a task owns directly, whichever invocation created them, as copies. A tool call settling finds the nested calls it left running with it.
|
||||
|
|
|
|||
|
|
@ -250,6 +250,8 @@ await root.configure({ modelTools: { remove: [bash] } }, context);
|
|||
|
||||
`conversation.agent()` resolves both lists: `tools`, offered to the model, and `callable`, which nested calls resolve among. `tools` still decides what is enabled at all; `modelTools` cannot add to it.
|
||||
|
||||
A `deferred: true` tool is not offered to the model until it is loaded, though other tools can call it any time. The `tool_search` tool from `@earendil-works/pi-durable/tool-search` (`createToolSearchExtension()`) lets the model search the deferred tools and loads the matches with `control.addTools`, which records them in `modelTools`; loaded tools stay offered on that branch, also after a compaction, and the `codemode` section leaves deferred tools out. Use it for large tool sets, such as an MCP server's.
|
||||
|
||||
### Tasks a restart abandons
|
||||
|
||||
A task that its creator awaits only in memory, and that the creator never resumes after a restart, can be created with `abandonOnRestart: true`. When a later Harness starts scheduling and finds it still live, it is aborted with `abortReason: "restart"`, together with everything it owns, before any of it runs again. Nested calls of tools that are not replay-safe get this automatically, and so do the child tasks such a tool creates with `api.createTask()`. A task abandoned this way whose definition is missing waits until its extension is installed, so its abort handler still cleans up. Its caller, and the conversation, wait with it; `harness.inspect()` shows it as blocked. Aborting it then orphans it, as any abort does.
|
||||
|
|
|
|||
|
|
@ -440,8 +440,12 @@ type AgentState = {
|
|||
extensions?: string[] | { add?: string[]; remove?: string[] };
|
||||
/** Filters the selected extensions' tools: the enabled tools. An array enables exactly these, in order. */
|
||||
tools?: string[] | { remove: string[] };
|
||||
/** Filters the enabled tools the model may call: the tools it is offered. An array offers exactly these, in order. */
|
||||
modelTools?: string[] | { remove: string[] };
|
||||
/**
|
||||
* Filters the enabled tools the model may call: the tools it is offered. Unset: all but deferred ones. An array:
|
||||
* exactly these, in order, deferred ones too. An object: all but deferred ones and those in `remove`, in registry
|
||||
* order, followed by the deferred tools in `add` that are not in `remove`, in `add` order.
|
||||
*/
|
||||
modelTools?: string[] | { remove?: string[]; add?: string[] };
|
||||
/** Rendered after every extension section, as the section `instructions`. */
|
||||
instructions?: string;
|
||||
/** Directory within the environment's file system, passed to `HarnessOptions.env`. */
|
||||
|
|
@ -460,7 +464,10 @@ type AgentChange = {
|
|||
| { readonly add?: readonly Extension[]; readonly remove?: readonly Extension[] }
|
||||
| null;
|
||||
readonly tools?: readonly ToolRegistration[] | { readonly remove: readonly ToolRegistration[] } | null;
|
||||
readonly modelTools?: readonly ToolRegistration[] | { readonly remove: readonly ToolRegistration[] } | null;
|
||||
readonly modelTools?:
|
||||
| readonly ToolRegistration[]
|
||||
| { readonly remove?: readonly ToolRegistration[]; readonly add?: readonly ToolRegistration[] }
|
||||
| null;
|
||||
readonly instructions?: string | null;
|
||||
readonly cwd?: string | null;
|
||||
};
|
||||
|
|
@ -477,6 +484,8 @@ type Agent<Tool extends ToolRegistration = ToolRegistration> = {
|
|||
readonly tools: readonly Tool[];
|
||||
/** Nested calls resolve among these: enabled, and `callers` includes `tools`. */
|
||||
readonly callable: readonly Tool[];
|
||||
/** Enabled deferred tools the model may call and is not offered yet, in registry order. */
|
||||
readonly deferred: readonly Tool[];
|
||||
/** Extension sections, then `instructions` when set. */
|
||||
readonly sections: readonly PromptSection<Tool>[];
|
||||
readonly instructions?: string;
|
||||
|
|
@ -3016,6 +3025,19 @@ narrows which of its model-callable tools the model is offered, for example
|
|||
only the code mode tool, while the others stay callable by tools. Neither can
|
||||
widen what `tools` enables.
|
||||
|
||||
A `deferred` tool is not offered to the model until it is loaded: `modelTools`
|
||||
must name it, in its array or its `add`. Other tools can call it any time. A
|
||||
tool loads deferred tools with `control.addTools`, which writes their names to
|
||||
`modelTools` (section 8.5); since `pi.agent` is rewindable and forks as of the
|
||||
fork's entry, loaded tools are per branch, and a baseline after a head marker
|
||||
declares them, built from the resolved tools. Loaded tools follow the other
|
||||
tools in load order, so a load appends declarations to the transcript and the
|
||||
cached prefix holds. `@earendil-works/pi-durable/tool-search` provides the
|
||||
coding agent's `tool_search` (`createToolSearchExtension()`): it ranks
|
||||
`Agent.deferred` with BM25 over the tools' names, descriptions, and parameter
|
||||
names and descriptions, and loads the matches. The `codemode` section does not
|
||||
list deferred tools; scripts find them with `searchTools()`.
|
||||
|
||||
Tools reach files and processes only through `api.env`, never through an
|
||||
environment captured when the tool was built. The tool task builds it for each
|
||||
call with `HarnessOptions.env` from the conversation's ID and `cwd` (section
|
||||
|
|
@ -4104,9 +4126,12 @@ and then commits once:
|
|||
|
||||
- It edits the conversation's stored `pi.agent` `tools` and `modelTools` for
|
||||
every `addTools` name: an array gets the name appended unless it holds it
|
||||
already, and `{ remove }` loses the name. An unset filter already lets every
|
||||
tool through, and nothing is written for it. The next preparation offers the
|
||||
tool when it resolves and the model may call it.
|
||||
already, and `{ remove }` loses the name. An unset `tools` already enables
|
||||
every tool, and nothing is written for it. An unset or object `modelTools`
|
||||
also gets the name appended to its `add`, which loads a deferred tool and has
|
||||
no effect for others. Names are written without consulting the registry, so a
|
||||
tool that is not registered now is offered once it is again. The next
|
||||
preparation offers the tool when it resolves and the model may call it.
|
||||
- When every result of the round requests `terminate`, or any requests
|
||||
`handoff`, it appends the handoff's `pi.reset` entry, if any, settles the
|
||||
run's inputs `done` with the tool-calling answer, removes `run` and `tools`,
|
||||
|
|
|
|||
|
|
@ -30,6 +30,12 @@
|
|||
"import": "./dist/codemode/index.js",
|
||||
"default": "./dist/codemode/index.js"
|
||||
},
|
||||
"./tool-search": {
|
||||
"source": "./src/tool-search/index.ts",
|
||||
"types": "./dist/tool-search/index.d.ts",
|
||||
"import": "./dist/tool-search/index.js",
|
||||
"default": "./dist/tool-search/index.js"
|
||||
},
|
||||
"./tools": {
|
||||
"source": "./src/tools/index.ts",
|
||||
"types": "./dist/tools/index.d.ts",
|
||||
|
|
|
|||
|
|
@ -58,6 +58,7 @@ import type {
|
|||
ToolExecutionResult,
|
||||
ToolRegistration,
|
||||
} from "../harness/types.ts";
|
||||
import { rankDocuments } from "../search.ts";
|
||||
import type { JsonObject } from "../types.ts";
|
||||
import {
|
||||
type CodemodeNamespace,
|
||||
|
|
@ -77,7 +78,6 @@ import {
|
|||
MODEL_CALLS,
|
||||
} from "./models.ts";
|
||||
import { DEFAULT_MAX_OUTPUT_TOKENS, formatResult, type OutputFile } from "./output.ts";
|
||||
import { rankDocuments } from "./search.ts";
|
||||
|
||||
/**
|
||||
* Version of everything in this module that turns stored results into what a script sees: result mapping, globals,
|
||||
|
|
@ -252,7 +252,10 @@ export function createCodemodeExtension(options: CodemodeToolOptions): Extension
|
|||
sections: [
|
||||
section("codemode", (input) =>
|
||||
codemodeSection(
|
||||
input.agent.callable.filter((tool) => !tool.name.startsWith(CODEMODE_INTERNAL_PREFIX)),
|
||||
// Deferred tools are found with `searchTools()`, as the coding agent leaves them out too.
|
||||
input.agent.callable.filter(
|
||||
(tool) => !tool.name.startsWith(CODEMODE_INTERNAL_PREFIX) && tool.deferred !== true,
|
||||
),
|
||||
options.namespaces ?? [],
|
||||
options.inlineBudget ?? DEFAULT_CODEMODE_INLINE_BUDGET,
|
||||
),
|
||||
|
|
|
|||
|
|
@ -82,23 +82,29 @@ export async function configure(tx: Tx, conversationId: ConversationId, change:
|
|||
|
||||
/**
|
||||
* `addTools` of a tool round, applied to `tools` and `modelTools`, so the next request offers the tools: an array gets
|
||||
* each name it lacks appended, `{ remove }` loses the names, and an unset filter already lets every tool through, so
|
||||
* nothing is written.
|
||||
* each name it lacks appended, and `{ remove }` loses the names. An unset `tools` already enables every tool, so nothing
|
||||
* is written for it. An unset or object `modelTools` also gets each name appended to its `add`, which loads deferred
|
||||
* tools and has no effect for others. By name, without the registry: a tool that is not registered now, such as one of
|
||||
* an MCP server that is reconnecting, is offered once it is again.
|
||||
*/
|
||||
export async function addTools(tx: Tx, conversationId: ConversationId, added: readonly string[]): Promise<void> {
|
||||
const state = await tx.doc(AgentDoc, conversationId);
|
||||
for (const field of ["tools", "modelTools"] as const) {
|
||||
const filter = state[field];
|
||||
if (filter === undefined) continue;
|
||||
if (Array.isArray(filter)) {
|
||||
for (const name of added) if (!filter.includes(name)) filter.push(name);
|
||||
} else {
|
||||
const remove = (filter as { remove: string[] }).remove;
|
||||
if (remove.some((name) => added.includes(name))) {
|
||||
state[field] = { remove: remove.filter((name) => !added.includes(name)) };
|
||||
}
|
||||
}
|
||||
const tools = state.tools;
|
||||
if (Array.isArray(tools)) {
|
||||
for (const name of added) if (!tools.includes(name)) tools.push(name);
|
||||
} else if (tools?.remove.some((name) => added.includes(name))) {
|
||||
state.tools = { remove: tools.remove.filter((name) => !added.includes(name)) };
|
||||
}
|
||||
const modelTools = state.modelTools;
|
||||
if (Array.isArray(modelTools)) {
|
||||
for (const name of added) if (!modelTools.includes(name)) modelTools.push(name);
|
||||
return;
|
||||
}
|
||||
// Assigned whole; a document write of an equal value records no change.
|
||||
const remove = (modelTools?.remove ?? []).filter((name) => !added.includes(name));
|
||||
const add = [...(modelTools?.add ?? [])];
|
||||
for (const name of added) if (!add.includes(name)) add.push(name);
|
||||
state.modelTools = { ...(remove.length === 0 ? {} : { remove }), add };
|
||||
}
|
||||
|
||||
function applyChange(state: Draft<AgentState>, change: AgentChange): void {
|
||||
|
|
@ -121,17 +127,23 @@ function applyChange(state: Draft<AgentState>, change: AgentChange): void {
|
|||
...(extensions.remove === undefined ? {} : { remove: names(extensions.remove) }),
|
||||
},
|
||||
);
|
||||
for (const field of ["tools", "modelTools"] as const) {
|
||||
const filter = change[field];
|
||||
set(
|
||||
field,
|
||||
filter === undefined || filter === null
|
||||
? filter
|
||||
: isList(filter)
|
||||
? names(filter)
|
||||
: { remove: names(filter.remove) },
|
||||
);
|
||||
}
|
||||
const tools = change.tools;
|
||||
set(
|
||||
"tools",
|
||||
tools === undefined || tools === null ? tools : isList(tools) ? names(tools) : { remove: names(tools.remove) },
|
||||
);
|
||||
const modelTools = change.modelTools;
|
||||
set(
|
||||
"modelTools",
|
||||
modelTools === undefined || modelTools === null
|
||||
? modelTools
|
||||
: isList(modelTools)
|
||||
? names(modelTools)
|
||||
: {
|
||||
...(modelTools.remove === undefined ? {} : { remove: names(modelTools.remove) }),
|
||||
...(modelTools.add === undefined ? {} : { add: names(modelTools.add) }),
|
||||
},
|
||||
);
|
||||
set("instructions", change.instructions);
|
||||
set("cwd", change.cwd);
|
||||
}
|
||||
|
|
@ -208,7 +220,10 @@ export function resolveAgent<Tool extends ToolRegistration>(
|
|||
const enabled = filterTools([...composed.values()], state?.tools);
|
||||
const callableBy = (caller: "model" | "tools") => (tool: Tool) =>
|
||||
tool.callers === undefined || tool.callers.includes(caller);
|
||||
const tools = filterTools(enabled.filter(callableBy("model")), state?.modelTools);
|
||||
const modelCallable = enabled.filter(callableBy("model"));
|
||||
const tools = offeredTools(modelCallable, state?.modelTools);
|
||||
const offered = new Set(tools);
|
||||
const deferred = modelCallable.filter((tool) => tool.deferred === true && !offered.has(tool));
|
||||
const callable = enabled.filter(callableBy("tools"));
|
||||
|
||||
const instructions = state?.instructions;
|
||||
|
|
@ -221,6 +236,7 @@ export function resolveAgent<Tool extends ToolRegistration>(
|
|||
extensions,
|
||||
tools,
|
||||
callable,
|
||||
deferred,
|
||||
sections: agentSections,
|
||||
...(instructions === undefined ? {} : { instructions }),
|
||||
...(state?.cwd === undefined ? {} : { cwd: state.cwd }),
|
||||
|
|
@ -228,6 +244,20 @@ export function resolveAgent<Tool extends ToolRegistration>(
|
|||
return agent;
|
||||
}
|
||||
|
||||
/** The tools offered to the model, as `AgentState.modelTools` describes. */
|
||||
function offeredTools<Tool extends ToolRegistration>(tools: Tool[], filter: AgentState["modelTools"]): Tool[] {
|
||||
if (Array.isArray(filter)) return filterTools(tools, filter);
|
||||
const removed = new Set(filter?.remove ?? []);
|
||||
const kept = tools.filter((tool) => tool.deferred !== true && !removed.has(tool.name));
|
||||
// Loaded deferred tools follow in load order, so loading one appends to the declared tools.
|
||||
const byName = new Map(tools.map((tool) => [tool.name, tool]));
|
||||
const loaded = [...new Set(filter?.add ?? [])].flatMap((name) => {
|
||||
const tool = byName.get(name);
|
||||
return tool?.deferred === true && !removed.has(name) ? [tool] : [];
|
||||
});
|
||||
return [...kept, ...loaded];
|
||||
}
|
||||
|
||||
/** `tools` through a stored filter: an array selects exactly its names, in its order; `{ remove }` drops names. */
|
||||
function filterTools<Tool extends ToolRegistration>(tools: Tool[], filter: AgentState["tools"]): Tool[] {
|
||||
if (filter === undefined) return tools;
|
||||
|
|
|
|||
|
|
@ -315,6 +315,11 @@ export type ToolRegistration<
|
|||
* model is never offered them.
|
||||
*/
|
||||
readonly callers?: readonly ("model" | "tools")[];
|
||||
/**
|
||||
* Not offered to the model until loaded: the model gets the tool only when `modelTools` names it, in its array or its
|
||||
* `add`, which `control.addTools` (from a tool such as `tool_search`) writes. Other tools can call it any time.
|
||||
*/
|
||||
readonly deferred?: true;
|
||||
readonly outputLimits?: {
|
||||
readonly maxBytes?: number;
|
||||
readonly maxLines?: number;
|
||||
|
|
@ -405,10 +410,15 @@ export type AgentState = {
|
|||
/** Filters the selected extensions' tools: the enabled tools. An array enables exactly these, in order. */
|
||||
tools?: string[] | { remove: string[] };
|
||||
/**
|
||||
* Filters the enabled tools the model may call (`callers` includes `model`): the tools offered to it. An array offers
|
||||
* exactly these, in order. Never widens: other tools can still call the ones it leaves out.
|
||||
* Filters the enabled tools the model may call (`callers` includes `model`): the tools offered to it. Never widens:
|
||||
* other tools can still call the ones it leaves out.
|
||||
* - Unset: every such tool except deferred ones.
|
||||
* - An array: exactly these, in order, deferred ones too.
|
||||
* - An object: every such tool except deferred ones and those in `remove`, in registry order, followed by the
|
||||
* deferred tools in `add` that are not in `remove`, in `add` order. Names of tools that are not deferred or not
|
||||
* registered have no effect in `add`.
|
||||
*/
|
||||
modelTools?: string[] | { remove: string[] };
|
||||
modelTools?: string[] | { remove?: string[]; add?: string[] };
|
||||
/** Rendered after every extension section, as the section `instructions`. */
|
||||
instructions?: string;
|
||||
/** Directory within the environment's file system, passed to `HarnessOptions.env`. */
|
||||
|
|
@ -424,7 +434,11 @@ export type AgentChange = {
|
|||
| { readonly add?: readonly Extension[]; readonly remove?: readonly Extension[] }
|
||||
| null;
|
||||
readonly tools?: readonly ToolRegistration[] | { readonly remove: readonly ToolRegistration[] } | null;
|
||||
readonly modelTools?: readonly ToolRegistration[] | { readonly remove: readonly ToolRegistration[] } | null;
|
||||
/** Replaces the stored filter as a whole; tools loaded with `control.addTools` are in its `add`. */
|
||||
readonly modelTools?:
|
||||
| readonly ToolRegistration[]
|
||||
| { readonly remove?: readonly ToolRegistration[]; readonly add?: readonly ToolRegistration[] }
|
||||
| null;
|
||||
readonly instructions?: string | null;
|
||||
readonly cwd?: string | null;
|
||||
};
|
||||
|
|
@ -438,6 +452,8 @@ export type Agent<Tool extends ToolRegistration = ToolRegistration> = {
|
|||
readonly tools: readonly Tool[];
|
||||
/** The tools nested calls (`executeTool()`) resolve among, in order: enabled and callable by tools. */
|
||||
readonly callable: readonly Tool[];
|
||||
/** Enabled deferred tools the model may call and is not offered yet, in registry order: what `tool_search` loads. */
|
||||
readonly deferred: readonly Tool[];
|
||||
/** Extension sections, then `instructions` when set. */
|
||||
readonly sections: readonly PromptSection<Tool>[];
|
||||
readonly instructions?: string;
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
/**
|
||||
* BM25 tool search for code mode's `searchTools()`, a copy of the coding agent's tool search ranking: durable code mode
|
||||
* keeps its own copy of what it needs from the coding agent.
|
||||
* BM25 tool search for code mode's `searchTools()` and the `tool_search` tool, a copy of the coding agent's tool search
|
||||
* ranking: durable keeps its own copy of what it needs from the coding agent.
|
||||
*/
|
||||
|
||||
export type SearchDocument = { readonly name: string; readonly text: string };
|
||||
7
packages/durable/src/tool-search/index.ts
Normal file
7
packages/durable/src/tool-search/index.ts
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
export {
|
||||
createToolSearchExtension,
|
||||
createToolSearchTool,
|
||||
DEFAULT_TOOL_SEARCH_LIMIT,
|
||||
TOOL_SEARCH_TOOL_NAME,
|
||||
type ToolSearchToolDetails,
|
||||
} from "./tool-search.ts";
|
||||
92
packages/durable/src/tool-search/tool-search.ts
Normal file
92
packages/durable/src/tool-search/tool-search.ts
Normal file
|
|
@ -0,0 +1,92 @@
|
|||
/**
|
||||
* The `tool_search` tool, a copy of the coding agent's: the model searches the deferred tools it is not offered yet,
|
||||
* and the matches are loaded with `control.addTools`, so the next request offers them. Loading writes the
|
||||
* conversation's `modelTools` (its `add`), so loaded tools stay offered on that branch, also after a compaction.
|
||||
*/
|
||||
import { Type } from "@earendil-works/pi-ai";
|
||||
import { defineExtension, defineTool } from "../harness/define.ts";
|
||||
import type { ToolRegistration } from "../harness/types.ts";
|
||||
import { rankDocuments } from "../search.ts";
|
||||
|
||||
export const TOOL_SEARCH_TOOL_NAME = "tool_search";
|
||||
export const DEFAULT_TOOL_SEARCH_LIMIT = 8;
|
||||
|
||||
/** Fixed, so it does not change while deferred tools are registered, for example as MCP servers connect. */
|
||||
const TOOL_SEARCH_DESCRIPTION = `# Tool discovery\n\nSearches over deferred tool metadata with BM25 and exposes matching tools for the next model call.\n\nSome of the tools, such as tools of MCP servers, may not have been provided to you upfront, and you should use this tool (\`${TOOL_SEARCH_TOOL_NAME}\`) to search for the required tools. For MCP tool discovery, always use \`${TOOL_SEARCH_TOOL_NAME}\`.`;
|
||||
|
||||
const toolSearchSchema = Type.Object({
|
||||
query: Type.String({ description: "Search query for deferred tools." }),
|
||||
limit: Type.Optional(
|
||||
Type.Number({ description: `Maximum number of tools to return. Defaults to ${DEFAULT_TOOL_SEARCH_LIMIT}.` }),
|
||||
),
|
||||
});
|
||||
|
||||
/** `tool_search` result details: the tools it loaded. */
|
||||
export type ToolSearchToolDetails = { loaded: string[] };
|
||||
|
||||
function isObject(value: unknown): value is Record<string, unknown> {
|
||||
return typeof value === "object" && value !== null && !Array.isArray(value);
|
||||
}
|
||||
|
||||
/** Schema descriptions and property names, recursively. */
|
||||
function schemaText(schema: unknown, parts: string[]): void {
|
||||
if (!isObject(schema)) return;
|
||||
if (typeof schema.description === "string") parts.push(schema.description);
|
||||
if (isObject(schema.properties)) {
|
||||
for (const [name, property] of Object.entries(schema.properties)) {
|
||||
parts.push(name);
|
||||
schemaText(property, parts);
|
||||
}
|
||||
}
|
||||
schemaText(schema.items, parts);
|
||||
for (const key of ["anyOf", "oneOf", "allOf"]) {
|
||||
const variants = schema[key];
|
||||
if (Array.isArray(variants)) for (const variant of variants) schemaText(variant, parts);
|
||||
}
|
||||
}
|
||||
|
||||
/** Search text of a tool: its name, the name with `_` as spaces, its description, and its parameters' text. */
|
||||
function searchText(tool: ToolRegistration): string {
|
||||
const parts = [tool.name, tool.name.replaceAll("_", " "), tool.description];
|
||||
schemaText(tool.parameters, parts);
|
||||
return parts.filter((part) => part.trim() !== "").join(" ");
|
||||
}
|
||||
|
||||
export function createToolSearchTool() {
|
||||
return defineTool({
|
||||
name: TOOL_SEARCH_TOOL_NAME,
|
||||
description: TOOL_SEARCH_DESCRIPTION,
|
||||
parameters: toolSearchSchema,
|
||||
// Searching changes what the model is offered; scripts have their own `searchTools()`.
|
||||
callers: ["model"],
|
||||
replay: "safe",
|
||||
async execute({ query, limit }, api, context) {
|
||||
if (query.trim() === "") throw new Error("query must not be empty");
|
||||
const max = limit ?? DEFAULT_TOOL_SEARCH_LIMIT;
|
||||
if (!Number.isInteger(max) || max <= 0) throw new Error("limit must be a positive integer");
|
||||
const candidates = (await api.agent(context)).deferred;
|
||||
const documents = candidates.map((tool) => ({ name: tool.name, text: searchText(tool) }));
|
||||
const loaded = rankDocuments(query, documents, max);
|
||||
const text =
|
||||
loaded.length === 0
|
||||
? "No matching tools found."
|
||||
: `Loaded ${loaded.length} tool${loaded.length === 1 ? "" : "s"}. They are available from your next call:\n${loaded
|
||||
.map((name) => {
|
||||
const description = candidates.find((tool) => tool.name === name)?.description ?? "";
|
||||
return `- ${name}: ${description.trim().split(/\r?\n/)[0]}`;
|
||||
})
|
||||
.join("\n")}`;
|
||||
const details: ToolSearchToolDetails = { loaded };
|
||||
return {
|
||||
output: [{ type: "text", text }],
|
||||
details,
|
||||
...(loaded.length === 0 ? {} : { control: { addTools: loaded } }),
|
||||
};
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
/** The `tool_search` extension: install it and select it where the model should find deferred tools. */
|
||||
export function createToolSearchExtension() {
|
||||
return defineExtension({ name: "tool-search", tools: [createToolSearchTool()] });
|
||||
}
|
||||
|
|
@ -73,6 +73,7 @@ describe("code mode description and section", () => {
|
|||
tools: [
|
||||
tool("read", "Read a file."),
|
||||
tool("mcp__github__list_issues", "List issues.", { callers: ["tools"] }),
|
||||
tool("mcp__github__get_issue", "Get an issue.", { deferred: true }),
|
||||
],
|
||||
});
|
||||
setup.faux.setResponses([fauxAssistantMessage([fauxText("hi")]) as AssistantMessage]);
|
||||
|
|
@ -82,6 +83,8 @@ describe("code mode description and section", () => {
|
|||
const sections = (system.model![0] as { sections?: Record<string, string> }).sections!;
|
||||
expect(sections.codemode).toContain("### `read`");
|
||||
expect(sections.codemode).toContain("### `mcp__github__list_issues`");
|
||||
// Deferred tools are found with searchTools(), not listed.
|
||||
expect(sections.codemode).not.toContain("mcp__github__get_issue");
|
||||
// The code mode tool itself is not callable from scripts.
|
||||
expect(sections.codemode).not.toContain("### `codemode`");
|
||||
// The tool declaration carries the script grammar; pi-ai applies it only for models that support it.
|
||||
|
|
|
|||
|
|
@ -1015,6 +1015,25 @@ describe("durable code mode output and tool search", () => {
|
|||
expect(message.content.some((item) => item.type === "image")).toBe(true);
|
||||
});
|
||||
|
||||
it("calls a deferred tool the model is not offered, and finds it in ALL_TOOLS", async () => {
|
||||
const chat = codemodeChat();
|
||||
addTool(
|
||||
chat.setup.registry,
|
||||
defineTool({
|
||||
name: "later",
|
||||
description: "A deferred tool",
|
||||
parameters: Type.Object({}),
|
||||
deferred: true,
|
||||
execute: async () => ({ output: [{ type: "text", text: "later ran" }] }),
|
||||
}),
|
||||
);
|
||||
const { text } = await runScript(
|
||||
`return [await tools.later({}), ALL_TOOLS.some((tool) => tool.name === "later")];`,
|
||||
chat,
|
||||
);
|
||||
expect(text).toContain('["later ran",true]');
|
||||
});
|
||||
|
||||
it("reports a failed script with its error and the calls already made", async () => {
|
||||
const { text, message } = await runScript(`await tools.echo({ text: "a" }); throw new Error("boom");`);
|
||||
expect(message.isError).toBe(true);
|
||||
|
|
|
|||
|
|
@ -62,7 +62,7 @@ function section(key: string, render: PromptSection<ToolRegistration>["render"],
|
|||
|
||||
const input: PromptInput<ToolRegistration> = {
|
||||
conversationId: 1 as never,
|
||||
agent: { thinkingLevel: "off", extensions: [], tools: [], callable: [], sections: [] },
|
||||
agent: { thinkingLevel: "off", extensions: [], tools: [], callable: [], deferred: [], sections: [] },
|
||||
env: undefined,
|
||||
shown: {},
|
||||
read: { snapshot: async () => undefined, snapshotAsOf: async () => undefined },
|
||||
|
|
|
|||
204
packages/durable/test/tool-search.test.ts
Normal file
204
packages/durable/test/tool-search.test.ts
Normal file
|
|
@ -0,0 +1,204 @@
|
|||
import { type AssistantMessage, fauxAssistantMessage, fauxText, fauxToolCall, Type } from "@earendil-works/pi-ai";
|
||||
import {
|
||||
AgentDoc,
|
||||
type Conversation,
|
||||
defineExtension,
|
||||
defineTool,
|
||||
type EntryRecord,
|
||||
type JsonObject,
|
||||
MemoryStorage,
|
||||
type ToolRegistration,
|
||||
} from "@earendil-works/pi-durable";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { createToolSearchExtension } from "../src/tool-search/index.ts";
|
||||
import { allEntries, type ChatSetup, chatSetup, openChat } from "./chat-support.ts";
|
||||
import { addTool } from "./harness-support.ts";
|
||||
import { context } from "./session-support.ts";
|
||||
|
||||
function tool(name: string, description: string, extra: Partial<ToolRegistration> = {}): ToolRegistration {
|
||||
return defineTool({
|
||||
name,
|
||||
description,
|
||||
parameters: Type.Object({ number: Type.Optional(Type.Number({ description: "Issue number" })) }),
|
||||
execute: async () => ({ output: [{ type: "text", text: `${name} ran` }] }),
|
||||
...extra,
|
||||
});
|
||||
}
|
||||
|
||||
const DONE = fauxAssistantMessage([fauxText("done")]);
|
||||
|
||||
function call(name: string, args: JsonObject = {}, id = "c1"): AssistantMessage {
|
||||
return fauxAssistantMessage([fauxToolCall(name, args, { id })], { stopReason: "toolUse" });
|
||||
}
|
||||
|
||||
/** `alpha`, the deferred `github_get_issue` and `github_list_repos`, then `omega`, and `tool_search` last. */
|
||||
function setup(): ChatSetup {
|
||||
const chat = chatSetup();
|
||||
addTool(chat.registry, tool("alpha", "Alpha tool"));
|
||||
addTool(chat.registry, tool("github_get_issue", "Get a GitHub issue by number", { deferred: true }));
|
||||
addTool(chat.registry, tool("github_list_repos", "List the repositories of a user", { deferred: true }));
|
||||
addTool(chat.registry, tool("omega", "Omega tool"));
|
||||
chat.registry.install(createToolSearchExtension());
|
||||
return chat;
|
||||
}
|
||||
|
||||
async function offered(root: Conversation): Promise<string[]> {
|
||||
return (await root.agent(context)).tools.map((each) => each.name);
|
||||
}
|
||||
|
||||
/** The model messages of the `pi.system` entries, oldest first. */
|
||||
function systemMessages(entries: readonly EntryRecord[]) {
|
||||
return entries
|
||||
.filter((entry) => entry.kind === "pi.system")
|
||||
.map((entry) => entry.model![0] as { toolsAdded?: { name: string }[]; toolsRemoved?: { name: string }[] });
|
||||
}
|
||||
|
||||
function toolResultTexts(entries: readonly EntryRecord[]): string[] {
|
||||
return entries
|
||||
.filter((entry) => entry.kind === "pi.tool-result")
|
||||
.map((entry) => (entry.model![0] as { content: { text?: string }[] }).content.map((item) => item.text).join(""));
|
||||
}
|
||||
|
||||
describe("deferred tools and tool_search", () => {
|
||||
it("offers a deferred tool only once tool_search loads it, by appending its declaration", async () => {
|
||||
const chat = setup();
|
||||
chat.faux.setResponses([
|
||||
call("tool_search", { query: "github issue", limit: 1 }),
|
||||
call("github_get_issue", { number: 42 }, "c2"),
|
||||
DONE,
|
||||
]);
|
||||
const { harness, root } = await openChat(new MemoryStorage(), chat);
|
||||
expect(await offered(root)).toEqual(["alpha", "omega", "tool_search"]);
|
||||
expect((await root.agent(context)).deferred.map((each) => each.name)).toEqual([
|
||||
"github_get_issue",
|
||||
"github_list_repos",
|
||||
]);
|
||||
// Deferred tools stay callable by other tools, such as code mode scripts.
|
||||
expect((await root.agent(context)).callable.map((each) => each.name)).toContain("github_get_issue");
|
||||
await (await root.submit({ type: "input", content: "go" }, context)).wait(context);
|
||||
const entries = await allEntries(root);
|
||||
const [search, issue] = toolResultTexts(entries);
|
||||
expect(search).toBe(
|
||||
"Loaded 1 tool. They are available from your next call:\n- github_get_issue: Get a GitHub issue by number",
|
||||
);
|
||||
expect(issue).toBe("github_get_issue ran");
|
||||
expect((await harness.snapshot(AgentDoc, root.id, context))?.modelTools).toEqual({ add: ["github_get_issue"] });
|
||||
expect(await offered(root)).toEqual(["alpha", "omega", "tool_search", "github_get_issue"]);
|
||||
// A later search sees only the tools not loaded yet.
|
||||
expect((await root.agent(context)).deferred.map((each) => each.name)).toEqual(["github_list_repos"]);
|
||||
// The load appended the declaration: no tool was removed and declared again, so the cached prefix holds.
|
||||
const [baseline, load] = systemMessages(entries);
|
||||
expect(baseline!.toolsAdded!.map((each) => each.name)).toEqual(["alpha", "omega", "tool_search"]);
|
||||
expect(load).toEqual(
|
||||
expect.objectContaining({ toolsAdded: [expect.objectContaining({ name: "github_get_issue" })] }),
|
||||
);
|
||||
expect(load!.toolsRemoved).toBeUndefined();
|
||||
await harness.close(context);
|
||||
});
|
||||
|
||||
it("appends each loaded tool in load order, not registry order", async () => {
|
||||
const chat = setup();
|
||||
chat.faux.setResponses([
|
||||
call("tool_search", { query: "repositories", limit: 1 }),
|
||||
call("tool_search", { query: "issue", limit: 1 }, "c2"),
|
||||
DONE,
|
||||
]);
|
||||
const { harness, root } = await openChat(new MemoryStorage(), chat);
|
||||
await (await root.submit({ type: "input", content: "go" }, context)).wait(context);
|
||||
expect(await offered(root)).toEqual(["alpha", "omega", "tool_search", "github_list_repos", "github_get_issue"]);
|
||||
const loads = systemMessages(await allEntries(root)).slice(1);
|
||||
expect(loads.map((message) => message.toolsAdded?.map((each) => each.name))).toEqual([
|
||||
["github_list_repos"],
|
||||
["github_get_issue"],
|
||||
]);
|
||||
expect(loads.every((message) => message.toolsRemoved === undefined)).toBe(true);
|
||||
await harness.close(context);
|
||||
});
|
||||
|
||||
it("records loads by name: unknown and non-deferred names change nothing, remove wins, and a load lifts it", async () => {
|
||||
const chat = setup();
|
||||
const grow = defineTool({
|
||||
name: "grow",
|
||||
description: "grow",
|
||||
parameters: Type.Object({}),
|
||||
execute: async () => ({ output: [], control: { addTools: ["github_get_issue", "alpha", "ghost"] } }),
|
||||
});
|
||||
addTool(chat.registry, grow);
|
||||
chat.faux.setResponses([call("grow"), DONE, call("grow"), DONE]);
|
||||
const { harness, root } = await openChat(new MemoryStorage(), chat);
|
||||
await (await root.submit({ type: "input", content: "go" }, context)).wait(context);
|
||||
expect((await harness.snapshot(AgentDoc, root.id, context))?.modelTools).toEqual({
|
||||
add: ["github_get_issue", "alpha", "ghost"],
|
||||
});
|
||||
expect(await offered(root)).toEqual(["alpha", "omega", "tool_search", "grow", "github_get_issue"]);
|
||||
// A name in both remove and add is not offered.
|
||||
const [getIssue] = (await root.agent(context)).callable.filter((each) => each.name === "github_get_issue");
|
||||
await root.configure({ modelTools: { remove: [getIssue!], add: [getIssue!] } }, context);
|
||||
expect(await offered(root)).not.toContain("github_get_issue");
|
||||
// Loading it again drops it from remove.
|
||||
await (await root.submit({ type: "input", content: "again" }, context)).wait(context);
|
||||
expect((await harness.snapshot(AgentDoc, root.id, context))?.modelTools).toEqual({
|
||||
add: ["github_get_issue", "alpha", "ghost"],
|
||||
});
|
||||
expect(await offered(root)).toContain("github_get_issue");
|
||||
await harness.close(context);
|
||||
});
|
||||
|
||||
it("keeps a loaded tool when its extension is gone for a while, and appends to an explicit array", async () => {
|
||||
const chat = setup();
|
||||
const issue = tool("github_get_issue", "Get a GitHub issue by number", { deferred: true });
|
||||
chat.faux.setResponses([call("tool_search", { query: "github issue", limit: 1 }), DONE]);
|
||||
const { harness, root } = await openChat(new MemoryStorage(), chat);
|
||||
await (await root.submit({ type: "input", content: "go" }, context)).wait(context);
|
||||
// The server goes away: its extension is installed again without the tool.
|
||||
chat.registry.install(defineExtension({ name: "tool:github_get_issue", tools: [] }));
|
||||
expect(await offered(root)).not.toContain("github_get_issue");
|
||||
chat.registry.install(defineExtension({ name: "tool:github_get_issue", tools: [issue] }));
|
||||
expect(await offered(root)).toContain("github_get_issue");
|
||||
// With an explicit array, a load appends to it as before.
|
||||
const alpha = (await root.agent(context)).tools.find((each) => each.name === "alpha")!;
|
||||
const search = (await root.agent(context)).tools.find((each) => each.name === "tool_search")!;
|
||||
await root.configure({ modelTools: [alpha, search] }, context);
|
||||
chat.faux.setResponses([call("tool_search", { query: "repositories", limit: 1 }), DONE]);
|
||||
await (await root.submit({ type: "input", content: "more" }, context)).wait(context);
|
||||
expect((await harness.snapshot(AgentDoc, root.id, context))?.modelTools).toEqual([
|
||||
"alpha",
|
||||
"tool_search",
|
||||
"github_list_repos",
|
||||
]);
|
||||
await harness.close(context);
|
||||
});
|
||||
|
||||
it("keeps loaded tools per branch and in the baseline after a reset", async () => {
|
||||
const chat = setup();
|
||||
chat.faux.setResponses([call("tool_search", { query: "github issue", limit: 1 }), DONE]);
|
||||
const { harness, root } = await openChat(new MemoryStorage(), chat);
|
||||
await (await root.submit({ type: "input", content: "go" }, context)).wait(context);
|
||||
const entries = await allEntries(root);
|
||||
const before = await root.fork(entries[0]!.id, { ownership: { kind: "ownerless" } }, context);
|
||||
const after = await root.fork(entries.at(-1)!.id, { ownership: { kind: "ownerless" } }, context);
|
||||
expect(await offered(before)).not.toContain("github_get_issue");
|
||||
expect(await offered(after)).toContain("github_get_issue");
|
||||
await root.reset(undefined, context);
|
||||
chat.faux.setResponses([DONE]);
|
||||
await (await root.submit({ type: "input", content: "fresh" }, context)).wait(context);
|
||||
const baseline = systemMessages(await allEntries(root)).at(-1)!;
|
||||
expect(baseline.toolsAdded!.map((each) => each.name)).toEqual([
|
||||
"alpha",
|
||||
"omega",
|
||||
"tool_search",
|
||||
"github_get_issue",
|
||||
]);
|
||||
await harness.close(context);
|
||||
});
|
||||
|
||||
it("finds nothing for an unrelated query and loads nothing", async () => {
|
||||
const chat = setup();
|
||||
chat.faux.setResponses([call("tool_search", { query: "weather forecast" }), DONE]);
|
||||
const { harness, root } = await openChat(new MemoryStorage(), chat);
|
||||
await (await root.submit({ type: "input", content: "go" }, context)).wait(context);
|
||||
expect(toolResultTexts(await allEntries(root))).toEqual(["No matching tools found."]);
|
||||
expect((await harness.snapshot(AgentDoc, root.id, context))?.modelTools).toBeUndefined();
|
||||
await harness.close(context);
|
||||
});
|
||||
});
|
||||
Loading…
Add table
Add a link
Reference in a new issue