From b5f425ad152491a28ed0f8c20461ba2170215b16 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Sun, 29 Mar 2026 21:06:12 +0200 Subject: [PATCH] feat(agent,coding-agent): add prepareArguments hook for pre-validation argument preparation Add AgentTool.prepareArguments and ToolDefinition.prepareArguments hook that runs before schema validation in the agent loop. This lets tools silently accept legacy argument shapes from resumed old sessions without polluting the public schema. The built-in edit tool uses this to fold legacy top-level oldText/newText into edits[] when resuming sessions that predate the edits-only schema. - AgentTool/ToolDefinition: typed prepareArguments returning Static - agent-loop: prepareToolCallArguments() runs before validateToolArguments() - edit tool: prepareEditArguments folds legacy fields, validateEditInput is strict - Documented in extensions.md with edit-tool example --- packages/agent/CHANGELOG.md | 4 + packages/agent/src/agent-loop.ts | 17 +++- packages/agent/src/types.ts | 5 +- packages/agent/test/agent-loop.test.ts | 80 +++++++++++++++++ packages/coding-agent/CHANGELOG.md | 5 ++ packages/coding-agent/docs/extensions.md | 61 +++++++++++++ .../coding-agent/src/core/extensions/types.ts | 3 + packages/coding-agent/src/core/tools/edit.ts | 28 ++++-- .../src/core/tools/tool-definition-wrapper.ts | 2 + .../test/edit-tool-legacy-input.test.ts | 90 +++++++++++++++++++ 10 files changed, 287 insertions(+), 8 deletions(-) create mode 100644 packages/coding-agent/test/edit-tool-legacy-input.test.ts diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index 0d9e9de1..75dbf623 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- Added `AgentTool.prepareArguments` hook to prepare raw tool call arguments before schema validation, enabling compatibility shims for resumed sessions with outdated tool schemas + ## [0.63.2] - 2026-03-29 ### Added diff --git a/packages/agent/src/agent-loop.ts b/packages/agent/src/agent-loop.ts index 1b04e809..586a1b7d 100644 --- a/packages/agent/src/agent-loop.ts +++ b/packages/agent/src/agent-loop.ts @@ -455,6 +455,20 @@ type ExecutedToolCallOutcome = { isError: boolean; }; +function prepareToolCallArguments(tool: AgentTool, toolCall: AgentToolCall): AgentToolCall { + if (!tool.prepareArguments) { + return toolCall; + } + const preparedArguments = tool.prepareArguments(toolCall.arguments); + if (preparedArguments === toolCall.arguments) { + return toolCall; + } + return { + ...toolCall, + arguments: preparedArguments as Record, + }; +} + async function prepareToolCall( currentContext: AgentContext, assistantMessage: AssistantMessage, @@ -472,7 +486,8 @@ async function prepareToolCall( } try { - const validatedArgs = validateToolArguments(tool, toolCall); + const preparedToolCall = prepareToolCallArguments(tool, toolCall); + const validatedArgs = validateToolArguments(tool, preparedToolCall); if (config.beforeToolCall) { const beforeResult = await config.beforeToolCall( { diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index 2e7b53a3..c391cc01 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -269,10 +269,13 @@ export interface AgentToolResult { // Callback for streaming tool execution updates export type AgentToolUpdateCallback = (partialResult: AgentToolResult) => void; -// AgentTool extends Tool but adds the execute function +// AgentTool extends Tool but adds argument preparation and execution hooks export interface AgentTool extends Tool { // A human-readable label for the tool to be displayed in UI label: string; + // Optional compatibility shim to prepare raw tool call arguments before schema validation. + // Must return an object conforming to TParameters. + prepareArguments?: (args: unknown) => Static; execute: ( toolCallId: string, params: Static, diff --git a/packages/agent/test/agent-loop.test.ts b/packages/agent/test/agent-loop.test.ts index 1e201e7d..00efc688 100644 --- a/packages/agent/test/agent-loop.test.ts +++ b/packages/agent/test/agent-loop.test.ts @@ -369,6 +369,86 @@ describe("agentLoop with AgentMessage", () => { expect(executed).toEqual([123]); }); + it("should prepare tool arguments for validation", async () => { + const replaceSchema = Type.Object({ oldText: Type.String(), newText: Type.String() }); + const toolSchema = Type.Object({ edits: Type.Array(replaceSchema) }); + const executed: Array> = []; + const tool: AgentTool = { + name: "edit", + label: "Edit", + description: "Edit tool", + parameters: toolSchema, + prepareArguments(args) { + if (!args || typeof args !== "object") { + return args as { edits: { oldText: string; newText: string }[] }; + } + const input = args as { + edits?: Array<{ oldText: string; newText: string }>; + oldText?: string; + newText?: string; + }; + if (typeof input.oldText !== "string" || typeof input.newText !== "string") { + return args as { edits: { oldText: string; newText: string }[] }; + } + return { + edits: [...(input.edits ?? []), { oldText: input.oldText, newText: input.newText }], + }; + }, + async execute(_toolCallId, params) { + executed.push(params.edits); + return { + content: [{ type: "text", text: `edited ${params.edits.length}` }], + details: { count: params.edits.length }, + }; + }, + }; + + const context: AgentContext = { + systemPrompt: "", + messages: [], + tools: [tool], + }; + + const userPrompt: AgentMessage = createUserMessage("edit something"); + const config: AgentLoopConfig = { + model: createModel(), + convertToLlm: identityConverter, + }; + + let callIndex = 0; + const streamFn = () => { + const stream = new MockAssistantStream(); + queueMicrotask(() => { + if (callIndex === 0) { + const message = createAssistantMessage( + [ + { + type: "toolCall", + id: "tool-1", + name: "edit", + arguments: { oldText: "before", newText: "after" }, + }, + ], + "toolUse", + ); + stream.push({ type: "done", reason: "toolUse", message }); + } else { + const message = createAssistantMessage([{ type: "text", text: "done" }]); + stream.push({ type: "done", reason: "stop", message }); + } + callIndex++; + }); + return stream; + }; + + const stream = agentLoop([userPrompt], context, config, undefined, streamFn); + for await (const _event of stream) { + // consume + } + + expect(executed).toEqual([[{ oldText: "before", newText: "after" }]]); + }); + it("should execute tool calls in parallel and emit tool results in source order", async () => { const toolSchema = Type.Object({ value: Type.String() }); let firstResolved = false; diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 84e6064f..3253a2d4 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,11 @@ ## [Unreleased] +### Added + +- Added `ToolDefinition.prepareArguments` hook to prepare raw tool call arguments before schema validation, enabling compatibility shims for resumed sessions with outdated tool schemas +- Built-in `edit` tool now uses `prepareArguments` to silently fold legacy top-level `oldText`/`newText` into `edits[]` when resuming old sessions + ### Fixed - Fixed extension-queued user messages to refresh the interactive pending-message list while preserving `input` event source semantics for `pi.sendUserMessage()` ([#2674](https://github.com/badlogic/pi-mono/pull/2674) by [@mrexodia](https://github.com/mrexodia)) diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index a450b211..a4f4f92b 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -1008,6 +1008,12 @@ pi.registerTool({ action: StringEnum(["list", "add"] as const), text: Type.Optional(Type.String()), }), + prepareArguments(args) { + // Optional compatibility shim. Runs before schema validation. + // Return the current schema shape, for example to fold legacy fields + // into the modern parameter object. + return args; + }, async execute(toolCallId, params, signal, onUpdate, ctx) { // Stream progress @@ -1472,6 +1478,14 @@ pi.registerTool({ action: StringEnum(["list", "add"] as const), // Use StringEnum for Google compatibility text: Type.Optional(Type.String()), }), + prepareArguments(args) { + if (!args || typeof args !== "object") return args; + const input = args as { action?: string; oldAction?: string }; + if (typeof input.oldAction === "string" && input.action === undefined) { + return { ...input, action: input.oldAction }; + } + return args; + }, async execute(toolCallId, params, signal, onUpdate, ctx) { // Check for cancellation @@ -1515,6 +1529,53 @@ async execute(toolCallId, params) { **Important:** Use `StringEnum` from `@mariozechner/pi-ai` for string enums. `Type.Union`/`Type.Literal` doesn't work with Google's API. +**Argument preparation:** `prepareArguments(args)` is optional. If defined, it runs before schema validation and before `execute()`. Use it to mimic an older accepted input shape when pi resumes an older session whose stored tool call arguments no longer match the current schema. Return the object you want validated against `parameters`. Keep the public schema strict. Do not add deprecated compatibility fields to `parameters` just to keep old resumed sessions working. + +Example: an older session may contain an `edit` tool call with top-level `oldText` and `newText`, while the current schema only accepts `edits: [{ oldText, newText }]`. + +```typescript +pi.registerTool({ + name: "edit", + label: "Edit", + description: "Edit a single file using exact text replacement", + parameters: Type.Object({ + path: Type.String(), + edits: Type.Array( + Type.Object({ + oldText: Type.String(), + newText: Type.String(), + }), + ), + }), + prepareArguments(args) { + if (!args || typeof args !== "object") return args; + + const input = args as { + path?: string; + edits?: Array<{ oldText: string; newText: string }>; + oldText?: unknown; + newText?: unknown; + }; + + if (typeof input.oldText !== "string" || typeof input.newText !== "string") { + return args; + } + + return { + ...input, + edits: [...(input.edits ?? []), { oldText: input.oldText, newText: input.newText }], + }; + }, + async execute(toolCallId, params, signal, onUpdate, ctx) { + // params now matches the current schema + return { + content: [{ type: "text", text: `Applying ${params.edits.length} edit block(s)` }], + details: {}, + }; + }, +}); +``` + ### Overriding Built-in Tools Extensions can override built-in tools (`read`, `bash`, `edit`, `write`, `grep`, `find`, `ls`) by registering a tool with the same name. Interactive mode displays a warning when this happens. diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index 4bd5adfe..6ede5f95 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -377,6 +377,9 @@ export interface ToolDefinition Static; + /** Execute the tool. */ execute( toolCallId: string, diff --git a/packages/coding-agent/src/core/tools/edit.ts b/packages/coding-agent/src/core/tools/edit.ts index ad821e7a..bb521aa9 100644 --- a/packages/coding-agent/src/core/tools/edit.ts +++ b/packages/coding-agent/src/core/tools/edit.ts @@ -44,6 +44,10 @@ const editSchema = Type.Object( ); export type EditToolInput = Static; +type LegacyEditToolInput = EditToolInput & { + oldText?: unknown; + newText?: unknown; +}; export interface EditToolDetails { /** Unified diff of the changes made */ @@ -76,13 +80,24 @@ export interface EditToolOptions { operations?: EditOperations; } -function validateEditInput(input: EditToolInput): { path: string; edits: Edit[] } { - if (!Array.isArray(input.edits)) { - throw new Error( - "Edit tool input is invalid. edits must be an array of replacements in the form { oldText: string, newText: string }.", - ); +function prepareEditArguments(input: unknown): EditToolInput { + if (!input || typeof input !== "object") { + return input as EditToolInput; } - if (input.edits.length === 0) { + + const args = input as LegacyEditToolInput; + if (typeof args.oldText !== "string" || typeof args.newText !== "string") { + return input as EditToolInput; + } + + const edits = Array.isArray(args.edits) ? [...args.edits] : []; + edits.push({ oldText: args.oldText, newText: args.newText }); + const { oldText: _oldText, newText: _newText, ...rest } = args; + return { ...rest, edits } as EditToolInput; +} + +function validateEditInput(input: EditToolInput): { path: string; edits: Edit[] } { + if (!Array.isArray(input.edits) || input.edits.length === 0) { throw new Error("Edit tool input is invalid. edits must contain at least one replacement."); } return { path: input.path, edits: input.edits }; @@ -154,6 +169,7 @@ export function createEditToolDefinition( "Keep edits[].oldText as small as possible while still being unique in the file. Do not pad with large unchanged regions.", ], parameters: editSchema, + prepareArguments: prepareEditArguments, async execute(_toolCallId, input: EditToolInput, signal?: AbortSignal, _onUpdate?, _ctx?) { const { path, edits } = validateEditInput(input); const absolutePath = resolveToCwd(path, cwd); diff --git a/packages/coding-agent/src/core/tools/tool-definition-wrapper.ts b/packages/coding-agent/src/core/tools/tool-definition-wrapper.ts index 913fbac1..6aac4a5e 100644 --- a/packages/coding-agent/src/core/tools/tool-definition-wrapper.ts +++ b/packages/coding-agent/src/core/tools/tool-definition-wrapper.ts @@ -11,6 +11,7 @@ export function wrapToolDefinition( label: definition.label, description: definition.description, parameters: definition.parameters, + prepareArguments: definition.prepareArguments, execute: (toolCallId, params, signal, onUpdate) => definition.execute(toolCallId, params, signal, onUpdate, ctxFactory?.() as ExtensionContext), }; @@ -36,6 +37,7 @@ export function createToolDefinitionFromAgentTool(tool: AgentTool): ToolDef label: tool.label, description: tool.description, parameters: tool.parameters as any, + prepareArguments: tool.prepareArguments, execute: async (toolCallId, params, signal, onUpdate) => tool.execute(toolCallId, params, signal, onUpdate), }; } diff --git a/packages/coding-agent/test/edit-tool-legacy-input.test.ts b/packages/coding-agent/test/edit-tool-legacy-input.test.ts new file mode 100644 index 00000000..d78acc65 --- /dev/null +++ b/packages/coding-agent/test/edit-tool-legacy-input.test.ts @@ -0,0 +1,90 @@ +import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, describe, expect, it } from "vitest"; +import type { ExtensionContext } from "../src/core/extensions/types.js"; +import { createEditToolDefinition } from "../src/core/tools/edit.js"; + +const tempDirs: string[] = []; + +async function createTempDir(): Promise { + const dir = await mkdtemp(join(tmpdir(), "pi-edit-legacy-input-")); + tempDirs.push(dir); + return dir; +} + +afterEach(async () => { + await Promise.all(tempDirs.splice(0, tempDirs.length).map((dir) => rm(dir, { recursive: true, force: true }))); +}); + +describe("edit tool prepareArguments", () => { + it("keeps legacy fields out of the public schema", () => { + const definition = createEditToolDefinition(process.cwd()); + expect(definition.parameters.properties).not.toHaveProperty("oldText"); + expect(definition.parameters.properties).not.toHaveProperty("newText"); + }); + + it("folds top-level oldText/newText into edits", () => { + const definition = createEditToolDefinition(process.cwd()); + const prepared = definition.prepareArguments!({ + path: "file.txt", + oldText: "before", + newText: "after", + }); + expect(prepared).toEqual({ + path: "file.txt", + edits: [{ oldText: "before", newText: "after" }], + }); + }); + + it("appends legacy replacement to existing edits", () => { + const definition = createEditToolDefinition(process.cwd()); + const prepared = definition.prepareArguments!({ + path: "file.txt", + edits: [{ oldText: "a", newText: "b" }], + oldText: "c", + newText: "d", + }); + expect(prepared).toEqual({ + path: "file.txt", + edits: [ + { oldText: "a", newText: "b" }, + { oldText: "c", newText: "d" }, + ], + }); + }); + + it("passes through valid input unchanged", () => { + const definition = createEditToolDefinition(process.cwd()); + const input = { + path: "file.txt", + edits: [{ oldText: "a", newText: "b" }], + }; + const prepared = definition.prepareArguments!(input); + expect(prepared).toBe(input); + }); + + it("passes through non-object input unchanged", () => { + const definition = createEditToolDefinition(process.cwd()); + expect(definition.prepareArguments!(null)).toBe(null); + expect(definition.prepareArguments!(undefined)).toBe(undefined); + expect(definition.prepareArguments!("garbage")).toBe("garbage"); + }); + + it("prepared args execute correctly", async () => { + const dir = await createTempDir(); + const filePath = join(dir, "legacy.txt"); + await writeFile(filePath, "before\n", "utf8"); + + const definition = createEditToolDefinition(dir); + const prepared = definition.prepareArguments!({ + path: "legacy.txt", + oldText: "before", + newText: "after", + }); + + const result = await definition.execute("tool-1", prepared, undefined, undefined, {} as ExtensionContext); + expect(result.content).toEqual([{ type: "text", text: "Successfully replaced 1 block(s) in legacy.txt." }]); + expect(await readFile(filePath, "utf8")).toBe("after\n"); + }); +});