mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +00:00
feat(coding-agent): codemode and MCP
This adds codemode and MCP support to pi. Codemode (packages/codemode) runs model-written JavaScript in a QuickJS wasm VM inside a worker. From there the script calls pi's tools as async functions, uses a per-session store, and can read the model catalog and run classifiers. The codemode tool in coding-agent is compatible with Codex's exec tool, and codemode.mode works like Codex's tool modes (on or only). Large results are written to a JSON temp file. MCP (packages/mcp) is a standalone client for stdio and streamable HTTP servers, with OAuth. The core only gets general mechanisms; codemode, tool_search and MCP are built-in extensions that use them. Other extensions can replace these built-ins. Tools declare an exposure: direct, model-only, codemode, deferred or hidden. prepareLoadout() lets an orchestrating tool change which tool declarations the model sees. ctx.executeTool() runs nested calls through validation and the tool_call/tool_result hooks. Those calls emit events with parentToolCallId and are recorded as bounded nestedCalls on the parent result, which compaction and HTML export use. Tools can also return structuredContent with an outputSchema, and can report errors with isError instead of throwing. MCP servers come from mcp.json (global, or per project after the project is trusted) or from pi.registerMcpServer(). They are managed with /mcp and pi mcp list|login|logout, and their runtimes load on first use. Closes #10040
This commit is contained in:
parent
b485fa3128
commit
8562bcf66a
143 changed files with 17620 additions and 225 deletions
|
|
@ -511,6 +511,10 @@ Thrown errors are caught by the agent and reported to the LLM as tool errors wit
|
|||
|
||||
Return `terminate: true` from `execute()`, a blocked `beforeToolCall`, or `afterToolCall` to hint that the agent should stop after the current tool batch. This only takes effect when every finalized tool result in the batch is terminating. The hint is runtime-only; emitted `toolResult` transcript messages remain standard LLM tool results.
|
||||
|
||||
### MCP and Codemode
|
||||
|
||||
`@earendil-works/pi-mcp` connects to MCP servers and `@earendil-works/pi-codemode` runs model-written JavaScript that calls tools. [examples/mcp-codemode](examples/mcp-codemode) wraps both as `AgentTool`s: one tool per MCP tool, and a `codemode` tool whose scripts call the agent's tools through `runToolCall()`, so `beforeToolCall` and `afterToolCall` apply to those calls too.
|
||||
|
||||
## Proxy Usage
|
||||
|
||||
For browser apps that proxy through a backend:
|
||||
|
|
|
|||
50
packages/agent/examples/mcp-codemode/main.ts
Normal file
50
packages/agent/examples/mcp-codemode/main.ts
Normal file
|
|
@ -0,0 +1,50 @@
|
|||
/**
|
||||
* An Agent with the tools of an MCP server plus a codemode tool that calls them from JavaScript.
|
||||
*
|
||||
* ANTHROPIC_API_KEY=... node packages/agent/examples/mcp-codemode/main.ts \
|
||||
* npx -y @modelcontextprotocol/server-filesystem .
|
||||
*
|
||||
* The arguments are the command that starts a stdio MCP server.
|
||||
*/
|
||||
|
||||
import { Agent } from "@earendil-works/pi-agent-core";
|
||||
import { createModels } from "@earendil-works/pi-ai";
|
||||
import { anthropicProvider } from "@earendil-works/pi-ai/providers/anthropic";
|
||||
import { McpClient, StdioTransport } from "@earendil-works/pi-mcp";
|
||||
import { createCodemodeTool, createMcpTools, createNestedToolRunner } from "./tools.ts";
|
||||
|
||||
const [command, ...args] = process.argv.slice(2);
|
||||
if (!command) {
|
||||
console.error("Usage: main.ts <mcp server command> [args...]");
|
||||
process.exit(1);
|
||||
}
|
||||
|
||||
const models = createModels();
|
||||
models.setProvider(anthropicProvider());
|
||||
const model = models.getModel("anthropic", "claude-sonnet-4-6");
|
||||
if (!model) throw new Error("Model not found");
|
||||
|
||||
const client = new McpClient({ name: "pi-agent-example", version: "1.0.0" });
|
||||
await client.connect(new StdioTransport({ command, args }));
|
||||
|
||||
try {
|
||||
const agent = new Agent({
|
||||
initialState: { systemPrompt: "You are a helpful assistant.", model },
|
||||
streamFn: models.streamSimple.bind(models),
|
||||
});
|
||||
const mcpTools = await createMcpTools(client, "mcp");
|
||||
// The model can call the MCP tools directly or batch them in one codemode script.
|
||||
agent.state.tools = [...mcpTools, createCodemodeTool(mcpTools, createNestedToolRunner(agent))];
|
||||
|
||||
agent.subscribe((event) => {
|
||||
if (event.type === "message_update" && event.assistantMessageEvent.type === "text_delta") {
|
||||
process.stdout.write(event.assistantMessageEvent.delta);
|
||||
} else if (event.type === "tool_execution_start") {
|
||||
console.log(`\n[${event.toolName}]`);
|
||||
}
|
||||
});
|
||||
await agent.prompt("Use codemode to list what the MCP tools can do, then summarize it in two sentences.");
|
||||
console.log();
|
||||
} finally {
|
||||
await client.close();
|
||||
}
|
||||
172
packages/agent/examples/mcp-codemode/tools.ts
Normal file
172
packages/agent/examples/mcp-codemode/tools.ts
Normal file
|
|
@ -0,0 +1,172 @@
|
|||
/**
|
||||
* Adapters that give a pi-agent-core `Agent` MCP tools and a codemode tool, using
|
||||
* `@earendil-works/pi-mcp` and `@earendil-works/pi-codemode`. See main.ts for how they fit together.
|
||||
*/
|
||||
|
||||
import {
|
||||
type Agent,
|
||||
type AgentTool,
|
||||
type AgentToolCallOutcome,
|
||||
type AgentToolResult,
|
||||
runToolCall,
|
||||
} from "@earendil-works/pi-agent-core";
|
||||
import type { ImageContent, JsonObject, JsonValue, TextContent } from "@earendil-works/pi-ai";
|
||||
import {
|
||||
type CodemodeJsonSchema,
|
||||
CodemodeSandbox,
|
||||
type CodemodeTool,
|
||||
renderDeclarations,
|
||||
} from "@earendil-works/pi-codemode";
|
||||
import { type McpClient, toLlmContent } from "@earendil-works/pi-mcp";
|
||||
import { Type } from "typebox";
|
||||
|
||||
/** Provider tool names are limited to 64 characters of `[A-Za-z0-9_-]`. */
|
||||
function toToolName(prefix: string, name: string): string {
|
||||
return `${prefix}_${name}`.replace(/[^A-Za-z0-9_-]/g, "_").slice(0, 64);
|
||||
}
|
||||
|
||||
function textOf(result: AgentToolResult<unknown>): string {
|
||||
return result.content
|
||||
.filter((block): block is TextContent => block.type === "text")
|
||||
.map((block) => block.text)
|
||||
.join("\n");
|
||||
}
|
||||
|
||||
/**
|
||||
* One `AgentTool` per tool of a connected MCP server, named `<prefix>_<tool>`. Tools with an MCP
|
||||
* output schema declare it as `outputSchema` and return `structuredContent`, which codemode scripts
|
||||
* receive instead of the text.
|
||||
*/
|
||||
export async function createMcpTools(client: McpClient, prefix: string): Promise<AgentTool[]> {
|
||||
const tools = await client.listTools();
|
||||
return tools.map(
|
||||
(tool): AgentTool => ({
|
||||
name: toToolName(prefix, tool.name),
|
||||
label: tool.title ?? tool.name,
|
||||
description: tool.description ?? tool.title ?? `MCP tool ${tool.name}`,
|
||||
// Providers require an object schema, and some reject one without `properties`.
|
||||
parameters: Type.Unsafe({
|
||||
...tool.inputSchema,
|
||||
type: "object",
|
||||
properties: tool.inputSchema.properties ?? {},
|
||||
}),
|
||||
outputSchema: tool.outputSchema ? Type.Unsafe<Record<string, unknown>>(tool.outputSchema) : undefined,
|
||||
execute: async (_toolCallId, params, signal, onUpdate) => {
|
||||
const result = await client.callTool(tool.name, params as Record<string, unknown>, {
|
||||
signal,
|
||||
onProgress: (progress) => {
|
||||
const text = progress.message ?? `Progress ${progress.progress}`;
|
||||
onUpdate?.({ content: [{ type: "text", text }], details: undefined });
|
||||
},
|
||||
});
|
||||
const content = toLlmContent(result);
|
||||
if (result.isError && content.length === 0) {
|
||||
content.push({ type: "text", text: `MCP tool ${tool.name} failed` });
|
||||
}
|
||||
return {
|
||||
content,
|
||||
details: undefined,
|
||||
structuredContent: result.structuredContent as JsonValue | undefined,
|
||||
// MCP reports tool failures in the result instead of as a protocol error.
|
||||
isError: result.isError === true,
|
||||
};
|
||||
},
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Run a tool call on behalf of another tool through the agent's pipeline: argument validation and
|
||||
* the agent's `beforeToolCall`/`afterToolCall` hooks apply as for calls the model makes. Calling
|
||||
* `tool.execute()` directly also works when the agent has no hooks.
|
||||
*/
|
||||
export function createNestedToolRunner(
|
||||
agent: Agent,
|
||||
): (name: string, args: unknown, signal: AbortSignal) => Promise<AgentToolCallOutcome> {
|
||||
let nextId = 1;
|
||||
return (name, args, signal) => {
|
||||
const assistantMessage = agent.state.messages.findLast((message) => message.role === "assistant");
|
||||
if (assistantMessage?.role !== "assistant") throw new Error("Nested tool calls need an assistant message");
|
||||
const tools = agent.state.tools;
|
||||
return runToolCall(
|
||||
{ type: "toolCall", id: `nested-${nextId++}`, name, arguments: (args ?? {}) as JsonObject },
|
||||
{
|
||||
tools,
|
||||
assistantMessage,
|
||||
context: { messages: agent.state.messages, tools },
|
||||
beforeToolCall: agent.beforeToolCall,
|
||||
afterToolCall: agent.afterToolCall,
|
||||
signal,
|
||||
},
|
||||
);
|
||||
};
|
||||
}
|
||||
|
||||
export const CODEMODE_TOOL_NAME = "codemode";
|
||||
|
||||
const codemodeParameters = Type.Object({
|
||||
code: Type.String({ description: "JavaScript source. Top-level await and return work." }),
|
||||
});
|
||||
|
||||
/**
|
||||
* A tool that runs model-written JavaScript in the codemode sandbox, where `tools.<name>(args)`
|
||||
* calls `tools`. Nested results stay out of the model context; only the script's output and
|
||||
* return value come back.
|
||||
*/
|
||||
export function createCodemodeTool(
|
||||
tools: readonly AgentTool[],
|
||||
runTool: (name: string, args: unknown, signal: AbortSignal) => Promise<AgentToolCallOutcome>,
|
||||
): AgentTool<typeof codemodeParameters> {
|
||||
const callable = tools.filter((tool) => tool.name !== CODEMODE_TOOL_NAME);
|
||||
const sandboxTools: CodemodeTool[] = callable.map((tool) => ({
|
||||
name: tool.name,
|
||||
description: tool.description,
|
||||
// TypeBox schemas are JSON Schema.
|
||||
inputSchema: tool.parameters as CodemodeJsonSchema,
|
||||
// What the script receives: the structured result if the tool declares one, else its text.
|
||||
outputSchema: (tool.outputSchema as CodemodeJsonSchema | undefined) ?? { type: "string" },
|
||||
execute: async (args, { signal }) => {
|
||||
const outcome = await runTool(tool.name, args, signal);
|
||||
if (tool.outputSchema && outcome.result.structuredContent !== undefined) {
|
||||
return outcome.result.structuredContent;
|
||||
}
|
||||
const text = textOf(outcome.result);
|
||||
if (outcome.isError) throw new Error(text || `Tool ${tool.name} failed`);
|
||||
return text;
|
||||
},
|
||||
}));
|
||||
|
||||
return {
|
||||
name: CODEMODE_TOOL_NAME,
|
||||
label: "Codemode",
|
||||
description: [
|
||||
"Run JavaScript that calls other tools, to chain calls, run them in parallel, or filter large results.",
|
||||
"The code is the body of an async function. There is no Node, file system, network, or timers.",
|
||||
"Call tools as `await tools.<name>(args)`. Output with `text(value)` or `console.log()`, or `return` a value.",
|
||||
"Only the output and the return value reach you, not the results of the tool calls.",
|
||||
"",
|
||||
"```ts",
|
||||
renderDeclarations({ tools: sandboxTools }),
|
||||
"```",
|
||||
].join("\n"),
|
||||
parameters: codemodeParameters,
|
||||
execute: async (_toolCallId, params, signal) => {
|
||||
const sandbox = new CodemodeSandbox({ tools: sandboxTools, timeoutMs: 120_000 });
|
||||
try {
|
||||
const result = await sandbox.execute(params.code, { signal });
|
||||
// Items from text(), console.*, and image(), in order. They have pi-ai's content shapes.
|
||||
const content: (TextContent | ImageContent)[] = [...result.output];
|
||||
if (result.ok && result.value !== undefined) {
|
||||
const value = typeof result.value === "string" ? result.value : JSON.stringify(result.value);
|
||||
content.push({ type: "text", text: value });
|
||||
}
|
||||
if (!result.ok) {
|
||||
content.push({ type: "text", text: `Script failed: ${result.error.stack ?? result.error.message}` });
|
||||
}
|
||||
return { content, details: { calls: result.calls }, isError: !result.ok };
|
||||
} finally {
|
||||
await sandbox.close();
|
||||
}
|
||||
},
|
||||
};
|
||||
}
|
||||
|
|
@ -23,6 +23,7 @@ import type {
|
|||
AgentMessage,
|
||||
AgentTool,
|
||||
AgentToolCall,
|
||||
AgentToolCallOutcome,
|
||||
AgentToolResult,
|
||||
PrepareNextTurnContext,
|
||||
StreamFn,
|
||||
|
|
@ -554,7 +555,7 @@ async function executeToolCallsSequential(
|
|||
isError: preparation.isError,
|
||||
};
|
||||
} else {
|
||||
const executed = await executePreparedToolCall(preparation, signal, emit);
|
||||
const executed = await executePreparedToolCall(preparation, signal, emitToolExecutionUpdate(toolCall, emit));
|
||||
finalized = await finalizeExecutedToolCall(
|
||||
currentContext,
|
||||
assistantMessage,
|
||||
|
|
@ -625,7 +626,7 @@ async function executeToolCallsParallel(
|
|||
await emitToolExecutionEnd(finalized, emit);
|
||||
return finalized;
|
||||
}
|
||||
const executed = await executePreparedToolCall(preparation, signal, emit);
|
||||
const executed = await executePreparedToolCall(preparation, signal, emitToolExecutionUpdate(toolCall, emit));
|
||||
const finalized = await finalizeExecutedToolCall(
|
||||
currentContext,
|
||||
assistantMessage,
|
||||
|
|
@ -676,11 +677,12 @@ type ExecutedToolCallOutcome = {
|
|||
isError: boolean;
|
||||
};
|
||||
|
||||
type FinalizedToolCallOutcome = {
|
||||
toolCall: AgentToolCall;
|
||||
result: AgentToolResult<any>;
|
||||
isError: boolean;
|
||||
};
|
||||
type FinalizedToolCallOutcome = AgentToolCallOutcome;
|
||||
|
||||
/** The `beforeToolCall` and `afterToolCall` hooks of {@link AgentLoopConfig}. */
|
||||
export type ToolCallHooks = Pick<AgentLoopConfig, "beforeToolCall" | "afterToolCall">;
|
||||
|
||||
type ToolUpdateSink = (partialResult: AgentToolResult<any>) => Promise<void> | void;
|
||||
|
||||
type FinalizedToolCallEntry = FinalizedToolCallOutcome | (() => Promise<FinalizedToolCallOutcome>);
|
||||
|
||||
|
|
@ -706,10 +708,11 @@ async function prepareToolCall(
|
|||
currentContext: AgentContext,
|
||||
assistantMessage: AssistantMessage,
|
||||
toolCall: AgentToolCall,
|
||||
config: AgentLoopConfig,
|
||||
config: ToolCallHooks,
|
||||
signal: AbortSignal | undefined,
|
||||
tools: readonly AgentTool<any>[] = currentContext.tools ?? [],
|
||||
): Promise<PreparedToolCall | ImmediateToolCallOutcome> {
|
||||
const tool = currentContext.tools?.find((t) => t.name === toolCall.name);
|
||||
const tool = tools.find((t) => t.name === toolCall.name);
|
||||
if (!tool) {
|
||||
return {
|
||||
kind: "immediate",
|
||||
|
|
@ -772,10 +775,52 @@ async function prepareToolCall(
|
|||
}
|
||||
}
|
||||
|
||||
function emitToolExecutionUpdate(toolCall: AgentToolCall, emit: AgentEventSink): ToolUpdateSink {
|
||||
return (partialResult) =>
|
||||
emit({
|
||||
type: "tool_execution_update",
|
||||
toolCallId: toolCall.id,
|
||||
toolName: toolCall.name,
|
||||
args: toolCall.arguments,
|
||||
partialResult,
|
||||
});
|
||||
}
|
||||
|
||||
/** Options for {@link runToolCall}. */
|
||||
export interface RunToolCallOptions extends ToolCallHooks {
|
||||
/** Tools the call resolves against. */
|
||||
tools: readonly AgentTool<any>[];
|
||||
/** Passed to the hooks as the message that issued the call. */
|
||||
assistantMessage: AssistantMessage;
|
||||
/** Passed to the hooks as the current agent context. */
|
||||
context: AgentContext;
|
||||
signal?: AbortSignal;
|
||||
onUpdate?: ToolUpdateSink;
|
||||
}
|
||||
|
||||
/**
|
||||
* Run one tool call through the same steps as a model-issued call: argument preparation, schema
|
||||
* validation, `beforeToolCall`, execution, and `afterToolCall`. Emits no events and adds no
|
||||
* messages. Tools that call other tools use this so the hooks (for example permission checks)
|
||||
* apply to those calls too.
|
||||
*
|
||||
* Never rejects for tool failures: unknown tools, validation errors, blocked calls, and thrown
|
||||
* errors come back as `isError: true`.
|
||||
*/
|
||||
export async function runToolCall(toolCall: AgentToolCall, options: RunToolCallOptions): Promise<AgentToolCallOutcome> {
|
||||
const { assistantMessage, context, signal } = options;
|
||||
const preparation = await prepareToolCall(context, assistantMessage, toolCall, options, signal, options.tools);
|
||||
if (preparation.kind === "immediate") {
|
||||
return { toolCall, result: preparation.result, isError: preparation.isError };
|
||||
}
|
||||
const executed = await executePreparedToolCall(preparation, signal, options.onUpdate ?? (() => {}));
|
||||
return finalizeExecutedToolCall(context, assistantMessage, preparation, executed, options, signal);
|
||||
}
|
||||
|
||||
async function executePreparedToolCall(
|
||||
prepared: PreparedToolCall,
|
||||
signal: AbortSignal | undefined,
|
||||
emit: AgentEventSink,
|
||||
onUpdate: ToolUpdateSink,
|
||||
): Promise<ExecutedToolCallOutcome> {
|
||||
const updateEvents: Promise<void>[] = [];
|
||||
let acceptingUpdates = true;
|
||||
|
|
@ -787,22 +832,12 @@ async function executePreparedToolCall(
|
|||
signal,
|
||||
(partialResult) => {
|
||||
if (!acceptingUpdates) return;
|
||||
updateEvents.push(
|
||||
Promise.resolve(
|
||||
emit({
|
||||
type: "tool_execution_update",
|
||||
toolCallId: prepared.toolCall.id,
|
||||
toolName: prepared.toolCall.name,
|
||||
args: prepared.toolCall.arguments,
|
||||
partialResult,
|
||||
}),
|
||||
),
|
||||
);
|
||||
updateEvents.push(Promise.resolve(onUpdate(partialResult)));
|
||||
},
|
||||
);
|
||||
acceptingUpdates = false;
|
||||
await Promise.all(updateEvents);
|
||||
return { result, isError: false };
|
||||
return { result, isError: result.isError === true };
|
||||
} catch (error) {
|
||||
acceptingUpdates = false;
|
||||
await Promise.all(updateEvents);
|
||||
|
|
@ -820,7 +855,7 @@ async function finalizeExecutedToolCall(
|
|||
assistantMessage: AssistantMessage,
|
||||
prepared: PreparedToolCall,
|
||||
executed: ExecutedToolCallOutcome,
|
||||
config: AgentLoopConfig,
|
||||
config: ToolCallHooks,
|
||||
signal: AbortSignal | undefined,
|
||||
): Promise<FinalizedToolCallOutcome> {
|
||||
let result = executed.result;
|
||||
|
|
@ -840,6 +875,9 @@ async function finalizeExecutedToolCall(
|
|||
signal,
|
||||
);
|
||||
if (afterResult) {
|
||||
// Structured content not replaced along with the content may no longer match it.
|
||||
const structuredContent =
|
||||
afterResult.structuredContent ?? (afterResult.content ? undefined : result.structuredContent);
|
||||
result = {
|
||||
...result,
|
||||
content: afterResult.content ?? result.content,
|
||||
|
|
@ -847,6 +885,8 @@ async function finalizeExecutedToolCall(
|
|||
usage: afterResult.usage ?? result.usage,
|
||||
terminate: afterResult.terminate ?? result.terminate,
|
||||
};
|
||||
if (structuredContent === undefined) delete result.structuredContent;
|
||||
else result.structuredContent = structuredContent;
|
||||
isError = afterResult.isError ?? isError;
|
||||
}
|
||||
} catch (error) {
|
||||
|
|
|
|||
|
|
@ -82,13 +82,17 @@ export interface BeforeToolCallResult {
|
|||
* - `isError`: if provided, replaces the tool result error flag
|
||||
* - `usage`: if provided, replaces the tool result usage
|
||||
* - `terminate`: if provided, replaces the early-termination hint
|
||||
* - `structuredContent`: if provided, replaces the structured content. If `content` is provided
|
||||
* without it, the structured content is dropped, because it may no longer match the content.
|
||||
* Return it along with `content` to keep it.
|
||||
*
|
||||
* Omitted fields keep the original executed tool result values.
|
||||
* Other omitted fields keep the original executed tool result values.
|
||||
* There is no deep merge for `content`, `details`, or `usage`.
|
||||
*/
|
||||
export interface AfterToolCallResult {
|
||||
content?: (TextContent | ImageContent)[];
|
||||
details?: unknown;
|
||||
structuredContent?: JsonValue;
|
||||
isError?: boolean;
|
||||
/** Usage from the final tool execution itself, if available. Not used for main LLM context accounting. */
|
||||
usage?: Usage;
|
||||
|
|
@ -422,8 +426,18 @@ export interface AgentToolResult<T = JsonValue | undefined> {
|
|||
content: (TextContent | ImageContent)[];
|
||||
/** Arbitrary structured details for logs or UI rendering. */
|
||||
details: T;
|
||||
/**
|
||||
* Machine-readable result matching the tool's `outputSchema`, for programmatic callers. Not sent
|
||||
* to the model; `content` remains the model-facing result.
|
||||
*/
|
||||
structuredContent?: JsonValue;
|
||||
/** Usage from the final tool execution itself, if available. Not used for main LLM context accounting. */
|
||||
usage?: Usage;
|
||||
/**
|
||||
* Report a failure without throwing. The model sees `content` as an error result, like a thrown
|
||||
* error, but `details` and `structuredContent` are kept for the UI and programmatic callers.
|
||||
*/
|
||||
isError?: boolean;
|
||||
/**
|
||||
* Hint that the agent should stop after the current tool batch.
|
||||
* Early termination only happens when every finalized tool result in the batch sets this to true.
|
||||
|
|
@ -431,6 +445,13 @@ export interface AgentToolResult<T = JsonValue | undefined> {
|
|||
terminate?: boolean;
|
||||
}
|
||||
|
||||
/** Final outcome of a tool call after hooks ran. */
|
||||
export interface AgentToolCallOutcome {
|
||||
toolCall: AgentToolCall;
|
||||
result: AgentToolResult<any>;
|
||||
isError: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Callback used by tools to stream partial execution updates.
|
||||
*
|
||||
|
|
@ -448,7 +469,15 @@ export interface AgentTool<TParameters extends TSchema = TSchema, TDetails = any
|
|||
* Must return an object that matches `TParameters`.
|
||||
*/
|
||||
prepareArguments?: (args: unknown) => Static<TParameters>;
|
||||
/** Execute the tool call. Throw on failure instead of encoding errors in `content`. */
|
||||
/**
|
||||
* JSON Schema of `structuredContent` in successful results. Tools that declare it should always
|
||||
* set `structuredContent`.
|
||||
*/
|
||||
outputSchema?: TSchema;
|
||||
/**
|
||||
* Execute the tool call. Throw on failure, or return a result with `isError: true`; do not only
|
||||
* describe the failure in `content`.
|
||||
*/
|
||||
execute: (
|
||||
toolCallId: string,
|
||||
params: Static<TParameters>,
|
||||
|
|
|
|||
|
|
@ -8,9 +8,16 @@ import {
|
|||
} from "@earendil-works/pi-ai";
|
||||
import { Type } from "typebox";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { agentLoop, agentLoopContinue, runAgentLoop } from "../src/agent-loop.ts";
|
||||
import { agentLoop, agentLoopContinue, runAgentLoop, runToolCall } from "../src/agent-loop.ts";
|
||||
import { setDefaultStreamFn } from "../src/index.ts";
|
||||
import type { AgentContext, AgentEvent, AgentLoopConfig, AgentMessage, AgentTool } from "../src/types.ts";
|
||||
import type {
|
||||
AgentContext,
|
||||
AgentEvent,
|
||||
AgentLoopConfig,
|
||||
AgentMessage,
|
||||
AgentTool,
|
||||
AgentToolCall,
|
||||
} from "../src/types.ts";
|
||||
|
||||
// Mock stream for testing - mimics MockAssistantStream
|
||||
class MockAssistantStream extends EventStream<AssistantMessageEvent, AssistantMessage> {
|
||||
|
|
@ -2081,3 +2088,106 @@ describe("agentLoopContinue with AgentMessage", () => {
|
|||
expect(messages[0].role).toBe("assistant");
|
||||
});
|
||||
});
|
||||
|
||||
describe("runToolCall", () => {
|
||||
const echoSchema = Type.Object({ value: Type.String() });
|
||||
const echo: AgentTool<typeof echoSchema> = {
|
||||
name: "echo",
|
||||
label: "Echo",
|
||||
description: "Echo tool",
|
||||
parameters: echoSchema,
|
||||
outputSchema: Type.Object({ value: Type.String() }),
|
||||
async execute(_toolCallId, params, _signal, onUpdate) {
|
||||
onUpdate?.({ content: [{ type: "text", text: "partial" }], details: {} });
|
||||
return {
|
||||
content: [{ type: "text", text: params.value }],
|
||||
details: {},
|
||||
structuredContent: { value: params.value },
|
||||
};
|
||||
},
|
||||
};
|
||||
const failing: AgentTool = {
|
||||
name: "failing",
|
||||
label: "Failing",
|
||||
description: "Returns an error result",
|
||||
parameters: Type.Object({}),
|
||||
async execute() {
|
||||
return { content: [{ type: "text", text: "bad" }], details: { partial: true }, isError: true };
|
||||
},
|
||||
};
|
||||
const assistantMessage = createAssistantMessage([]);
|
||||
const call = (id: string, name: string, args: AgentToolCall["arguments"]): AgentToolCall => ({
|
||||
type: "toolCall",
|
||||
id,
|
||||
name,
|
||||
arguments: args,
|
||||
});
|
||||
|
||||
it("validates, runs the hooks, and reports failures as error outcomes", async () => {
|
||||
const hookCalls: string[] = [];
|
||||
const updates: unknown[] = [];
|
||||
const options = {
|
||||
tools: [echo, failing],
|
||||
assistantMessage,
|
||||
context: { messages: [] },
|
||||
beforeToolCall: async ({ toolCall, args }: { toolCall: { id: string }; args: unknown }) => {
|
||||
hookCalls.push(`before ${toolCall.id}`);
|
||||
if ((args as { value?: string }).value === "blocked") return { block: true, reason: "nope" };
|
||||
return undefined;
|
||||
},
|
||||
afterToolCall: async ({ toolCall }: { toolCall: { id: string } }) => {
|
||||
hookCalls.push(`after ${toolCall.id}`);
|
||||
return undefined;
|
||||
},
|
||||
onUpdate: (partial: unknown) => {
|
||||
updates.push(partial);
|
||||
},
|
||||
};
|
||||
|
||||
expect(await runToolCall(call("a", "echo", { value: "a" }), options)).toMatchObject({
|
||||
toolCall: { id: "a" },
|
||||
result: { structuredContent: { value: "a" } },
|
||||
isError: false,
|
||||
});
|
||||
expect(await runToolCall(call("b", "echo", { value: { nested: true } }), options)).toMatchObject({
|
||||
isError: true,
|
||||
});
|
||||
expect(await runToolCall(call("c", "echo", { value: "blocked" }), options)).toMatchObject({
|
||||
result: { content: [{ type: "text", text: "nope" }] },
|
||||
isError: true,
|
||||
});
|
||||
expect(await runToolCall(call("d", "missing", {}), options)).toMatchObject({
|
||||
result: { content: [{ type: "text", text: "Tool missing not found" }] },
|
||||
isError: true,
|
||||
});
|
||||
// Error results keep their details.
|
||||
expect(await runToolCall(call("e", "failing", {}), options)).toMatchObject({
|
||||
result: { details: { partial: true } },
|
||||
isError: true,
|
||||
});
|
||||
expect(updates).toEqual([{ content: [{ type: "text", text: "partial" }], details: {} }]);
|
||||
// Validation failures and unknown tools never reach the hooks; blocked calls skip afterToolCall.
|
||||
expect(hookCalls).toEqual(["before a", "after a", "before c", "before e", "after e"]);
|
||||
});
|
||||
|
||||
it("lets afterToolCall replace structured content and drops it when only content is replaced", async () => {
|
||||
const redacted = [{ type: "text" as const, text: "redacted" }];
|
||||
const results = [
|
||||
{ content: redacted },
|
||||
{ structuredContent: { value: "replaced" } },
|
||||
{ content: redacted, structuredContent: { value: "both" } },
|
||||
{ details: { note: "kept" } },
|
||||
];
|
||||
const seen: unknown[] = [];
|
||||
for (const afterResult of results) {
|
||||
const outcome = await runToolCall(call("x", "echo", { value: "original" }), {
|
||||
tools: [echo],
|
||||
assistantMessage,
|
||||
context: { messages: [] },
|
||||
afterToolCall: async () => afterResult,
|
||||
});
|
||||
seen.push(outcome.result.structuredContent);
|
||||
}
|
||||
expect(seen).toEqual([undefined, { value: "replaced" }, { value: "both" }, { value: "original" }]);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue