Merge branch 'main' of https://github.com/badlogic/pi-mono

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

packages/agent/src/agent-loop.ts

@@ -455,6 +455,20 @@ type ExecutedToolCallOutcome = {
 	isError: boolean;
 };
 
+function prepareToolCallArguments(tool: AgentTool<any>, 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<string, any>,
+	};
+}
+
 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(
 				{

packages/agent/src/types.ts

@@ -269,10 +269,13 @@ export interface AgentToolResult<T> {
 // Callback for streaming tool execution updates
 export type AgentToolUpdateCallback<T = any> = (partialResult: AgentToolResult<T>) => void;
 
-// AgentTool extends Tool but adds the execute function
+// AgentTool extends Tool but adds argument preparation and execution hooks
 export interface AgentTool<TParameters extends TSchema = TSchema, TDetails = any> extends Tool<TParameters> {
 	// 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<TParameters>;
 	execute: (
 		toolCallId: string,
 		params: Static<TParameters>,

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<Array<{ oldText: string; newText: string }>> = [];
+		const tool: AgentTool<typeof toolSchema, { count: number }> = {
+			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;

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

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.

packages/coding-agent/src/core/extensions/types.ts

@@ -377,6 +377,9 @@ export interface ToolDefinition<TParams extends TSchema = TSchema, TDetails = un
 	/** Parameter schema (TypeBox) */
 	parameters: TParams;
 
+	/** Optional compatibility shim to prepare raw tool call arguments before schema validation. Must return an object conforming to TParams. */
+	prepareArguments?: (args: unknown) => Static<TParams>;
+
 	/** Execute the tool. */
 	execute(
 		toolCallId: string,

packages/coding-agent/src/core/tools/edit.ts

@@ -44,6 +44,10 @@ const editSchema = Type.Object(
 );
 
 export type EditToolInput = Static<typeof editSchema>;
+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);

packages/coding-agent/src/core/tools/tool-definition-wrapper.ts

@@ -11,6 +11,7 @@ export function wrapToolDefinition<TDetails = unknown>(
 		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<any>): 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),
 	};
 }

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<string> {
+	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");
+	});
+});