automatic tool call repair for pi coding agent

A developer created an automatic tool call repair extension for the Pi coding agent that fixes invalid tool call arguments before execution. The extension validates tool call parameters against their schemas and attempts to repair common issues like incorrect paths or missing required fields, logging all repair attempts for debugging.

| // put in ~/.pi/agent/extensions/tool-input-repair.ts | | | // based off of https://x.com/MrAhmadAwais/status/2050956678502420612 | | | import { appendFile, mkdir } from "node:fs/promises"; | | | import { homedir } from "node:os"; | | | import { dirname, join } from "node:path"; | | | import { env } from "node:process"; | | | import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent"; | | | import { Compile } from "typebox/compile"; | | | import { Value } from "typebox/value"; | | | type Args = Record<string, unknown ; | | | type ToolLike = { name: string; parameters: unknown }; | | | type RepairResult = | | | | { status: "none" } | | | | { status: "repaired"; args: Args; paths: string } | | | | { status: "failed"; paths: string }; | | | const DEFAULT LOG FILE = join homedir , ".pi/agent/tool-input-repair.jsonl" ; | | | const LOG ENV = "PI TOOL INPUT REPAIR LOG"; | | | const MAX LOG TEXT = 2000; | | | const PATH FIELDS = new Set | | | "path", | | | "filePath", | | | "file path", | | | "absolutePath", | | | "absolute path", | | | "targetPath", | | | "target path", | | | ; | | | const validators = new WeakMap<object, ReturnType<typeof Compile ; | | | function isRecord value: unknown : value is Args { | | | return typeof value === "object" && value == null && Array.isArray value ; | | | } | | | function truncate value: string, max = MAX LOG TEXT : string { | | | return value.length max ? ${value.slice 0, max }… : value; | | | } | | | function summarizeValue value: unknown : unknown { | | | if value === null || value === undefined return value; | | | if typeof value === "string" return { type: "string", length: value.length, preview: truncate value, 160 }; | | | if typeof value === "number" || typeof value === "boolean" return value; | | | if Array.isArray value return { type: "array", length: value.length, items: value.slice 0, 3 .map summarizeValue }; | | | if isRecord value return { type: "object", keys: Object.keys value , values: summarizeArgs value }; | | | return { type: typeof value }; | | | } | | | function summarizeArgs args: Args : Args { | | | return Object.fromEntries Object.entries args .map k, v = k, summarizeValue v ; | | | } | | | function configuredLogFile : string | undefined { | | | const value = env LOG ENV ?.trim ; | | | if value || /^ 0|false|no|off $/i.test value return undefined; | | | if /^ 1|true|yes|on $/i.test value return DEFAULT LOG FILE; | | | return value; | | | } | | | async function logEvent ctx: ExtensionContext, event: Args : Promise<void { | | | const file = configuredLogFile ; | | | if file return; | | | const model = ctx.model; | | | try { | | | await mkdir dirname file , { recursive: true } ; | | | await appendFile file, ${JSON.stringify { timestamp: new Date .toISOString , cwd: ctx.cwd, model: model && { provider: model.provider, id: model.id, name: model.name }, ...event } }\n ; | | | } catch { | | | // Logging must never affect tool execution. | | | } | | | } | | | function validatorFor schema: unknown : ReturnType<typeof Compile | undefined { | | | if isRecord schema return undefined; | | | let validator = validators.get schema ; | | | if validator { | | | validator = Compile schema ; | | | validators.set schema, validator ; | | | } | | | return validator; | | | } | | | function validationPath error: any : string { | | | if error.keyword === "required" { | | | const required = error.params?.requiredProperties?. 0 ?? error.params?.missingProperty; | | | if required { | | | const base = String error.instancePath ?? "" .replace /^\//, "" .replace /\//g, "." ; | | | return base ? ${base}.${required} : required; | | | } | | | } | | | return String error.instancePath ?? "" .replace /^\//, "" .replace /\//g, "." || "root"; | | | } | | | function validate schema: unknown, args: Args, convert: boolean : { ok: true } | { ok: false; paths: string } { | | | try { | | | const $ = validatorFor schema ; | | | if $ return { ok: true }; | | | const checked = convert ? structuredClone args : args; | | | if convert Value.Convert schema, checked ; | | | if $.Check checked return { ok: true }; | | | return { ok: false, paths: ...$.Errors checked .map validationPath .filter p = p == "root" }; | | | } catch { | | | return { ok: true }; | | | } | | | } | | | function topField path: string : string { | | | return path.split "." 0 ; | | | } | | | function fieldSchema tool: ToolLike, field: string : unknown { | | | const props = isRecord tool.parameters ? tool.parameters.properties : undefined; | | | return isRecord props ? props field : undefined; | | | } | | | function isArraySchema schema: unknown, seen = new WeakSet<object : boolean { | | | if isRecord schema return false; | | | if seen.has schema return false; | | | seen.add schema ; | | | if schema.type === "array" return true; | | | const anyOf = Array.isArray schema.anyOf ? schema.anyOf : ; | | | const oneOf = Array.isArray schema.oneOf ? schema.oneOf : ; | | | return ...anyOf, ...oneOf .some item = isArraySchema item, seen ; | | | } | | | function fieldIsOptional tool: ToolLike, field: string : boolean { | | | const required = isRecord tool.parameters ? tool.parameters.required : undefined; | | | return Array.isArray required || required.includes field ; | | | } | | | function parseJsonArray value: string : unknown | undefined { | | | const trimmed = value.trim ; | | | if trimmed.startsWith " " return undefined; | | | try { | | | const parsed = JSON.parse trimmed ; | | | return Array.isArray parsed ? parsed : undefined; | | | } catch { | | | return undefined; | | | } | | | } | | | function repairField tool: ToolLike, args: Args, field: string : boolean { | | | const schema = fieldSchema tool, field ; | | | const value = args field ; | | | // Safe: only reaches this branch after strict validation rejected null . | | | if value === null && fieldIsOptional tool, field { | | | delete args field ; | | | return true; | | | } | | | if isArraySchema schema return false; | | | if typeof value === "string" { | | | args field = parseJsonArray value ?? value ; | | | return true; | | | } | | | if isRecord value { | | | args field = value ; | | | return true; | | | } | | | return false; | | | } | | | function unwrapAutoLink value: string : string { | | | const match = value.match /^\ ^\ + \ https?:\/\/ ^ + \ $/ ; | | | if match return value; | | | const , text, href, hrefWithoutScheme = match; | | | return text === href || text === hrefWithoutScheme ? text : value; | | | } | | | function repairPathLinks args: Args : string { | | | const repaired: string = ; | | | for const field of PATH FIELDS { | | | const value = args field ; | | | if typeof value == "string" continue; | | | const next = unwrapAutoLink value ; | | | if next === value continue; | | | args field = next; | | | repaired.push field ; | | | } | | | return repaired; | | | } | | | function repairArgs tool: ToolLike, args: Args : RepairResult { | | | const patched = structuredClone args ; | | | const pathFields = repairPathLinks patched ; | | | // Validate before shape repair; pre-normalizing can corrupt valid JSON-shaped tool input. | | | const strict = validate tool.parameters, patched, false ; | | | if strict.ok { | | | return pathFields.length 0 ? { status: "repaired", args: patched, paths: pathFields } : { status: "none" }; | | | } | | | const repaired = structuredClone patched ; | | | let changedShape = false; | | | for const path of strict.paths { | | | changedShape = repairField tool, repaired, topField path || changedShape; | | | } | | | if changedShape && validate tool.parameters, repaired, false .ok { | | | return { status: "repaired", args: repaired, paths: ...new Set ...pathFields, ...strict.paths }; | | | } | | | // Structural repairs failed; try again with TypeBox coercion e.g. "1" → 1 . | | | if validate tool.parameters, patched, true .ok { | | | return pathFields.length 0 ? { status: "repaired", args: patched, paths: pathFields } : { status: "none" }; | | | } | | | return { status: "failed", paths: strict.paths }; | | | } | | | function validationHint text: string : string | undefined { | | | if /^Validation failed for tool /.test text return undefined; | | | const issues = ...text.matchAll /^\s -\s+ ^: + :\s+ .+ $/gm .map | | | , path, message = ${path}: ${message} , | | | ; | | | if issues.length === 0 return "Tool arguments failed validation. Retry with corrected argument types."; | | | return Tool arguments failed validation: ${issues.join "; " }. Retry with corrected argument types. ; | | | } | | | export default function pi: ExtensionAPI { | | | pi.on "message end", async event, ctx = { | | | if event.message.role === "assistant" { | | | const tools = new Map pi.getAllTools .map tool = tool.name, tool ; | | | let changed = false; | | | const content = event.message.content.map part = { | | | if part?.type == "toolCall" || isRecord part.arguments return part; | | | const tool = tools.get part.name ; | | | if tool return part; | | | const result = repairArgs tool, part.arguments ; | | | if result.status === "repaired" { | | | changed = true; | | | void logEvent ctx, { | | | type: "tool input repaired", | | | toolName: part.name, | | | toolCallId: part.id, | | | paths: result.paths, | | | before: summarizeArgs part.arguments , | | | after: summarizeArgs result.args , | | | } ; | | | return { ...part, arguments: result.args }; | | | } | | | if result.status === "failed" { | | | void logEvent ctx, { | | | type: "tool input unrepaired", | | | toolName: part.name, | | | toolCallId: part.id, | | | paths: result.paths, | | | args: summarizeArgs part.arguments , | | | } ; | | | } | | | return part; | | | } ; | | | if changed { | | | if ctx.hasUI ctx.ui.notify "Repaired assistant tool-call arguments", "info" ; | | | return { message: { ...event.message, content } }; | | | } | | | } | | | if event.message.role === "toolResult" && event.message.isError { | | | const text = event.message.content | | | .filter part = part.type === "text" | | | .map part = part.text ?? "" | | | .join "\n" ; | | | const hint = validationHint text ; | | | void logEvent ctx, { | | | type: hint ? "tool validation error" : "tool result error", | | | toolName: event.message.toolName, | | | toolCallId: event.message.toolCallId, | | | error: truncate text , | | | } ; | | | if hint return; | | | return { | | | message: { | | | ...event.message, | | | content: { type: "text", text: ${hint}\n\n---\n${text} } , | | | }, | | | }; | | | } | | | } ; | | | } |