| // 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} }], | |
| }, | |
| }; | |
| } | |
| }); | |
| } |