automatic tool call repair for pi coding agent

tool-input-repair.ts

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