edit-flat.ts

/**
 * edit-flat - drop-in replacement for the built-in `edit` tool that takes
 * replacement fields at the top level instead of an `edits[]` array.
 *
 * Models emit top-level string arguments more reliably than nested JSON arrays
 * (no array bracket/quote escaping), so this variant exposes `oldText` and
 * `newText` directly. To make several replacements in one call, append a
 * number to both names: `oldText2`/`newText2`, `oldText3`/`newText3`, ...
 *
 * Behaviour, result shape, diff details, and TUI rendering are identical to
 * the built-in `edit` tool: this definition delegates execute/render to the
 * real edit tool definition and only reshapes the arguments.
 */

import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { createEditToolDefinition } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";

type Edit = { oldText: string; newText: string };
type EditArgs = { path: string; edits: Edit[] };

/**
 * Collect the numbered oldText/newText pairs from a raw argument object.
 * Index 1 uses the bare names (`oldText`/`newText`); index >= 2 uses the
 * suffixed names (`oldText2`/`newText2`, ...). Stops at the first missing
 * pair so the schema does not have to enumerate the variants.
 */
function collectEdits(args: Record<string, unknown>): Edit[] {
  const edits: Edit[] = [];
  for (let i = 1; ; i++) {
    const oldKey = i === 1 ? "oldText" : `oldText${i}`;
    const newKey = i === 1 ? "newText" : `newText${i}`;
    const oldVal = args[oldKey];
    const newVal = args[newKey];
    if (oldVal === undefined && newVal === undefined) break;
    if (typeof oldVal !== "string" || typeof newVal !== "string") {
      throw new Error(
        `Edit ${i}: both ${oldKey} and ${newKey} must be provided as strings.`,
      );
    }
    edits.push({ oldText: oldVal, newText: newVal });
  }
  return edits;
}

function toEditArgs(args: unknown): EditArgs | unknown {
  if (!args || typeof args !== "object") return args;
  const input = args as Record<string, unknown>;

  // Already in the canonical shape (e.g. a resumed older session that stored
  // an edits array, or a model that emitted one anyway).
  if (Array.isArray(input.edits) && input.edits.length > 0) {
    return { path: input.path, edits: input.edits } as EditArgs;
  }

  // Some models serialize the array as a JSON string.
  if (typeof input.edits === "string") {
    try {
      const parsed = JSON.parse(input.edits);
      if (Array.isArray(parsed) && parsed.length > 0) {
        return { path: input.path, edits: parsed } as EditArgs;
      }
    } catch {
      // fall through to flat-argument collection
    }
  }

  const flat = collectEdits(input);
  if (flat.length === 0) return args;
  return { path: input.path, edits: flat } as EditArgs;
}

const schema = Type.Object(
  {
    path: Type.String({
      description: "Path to the file to edit (relative or absolute).",
    }),
    oldText: Type.Optional(
      Type.String({
        description:
          "Exact text for the first replacement. Must be unique in the original file and must not overlap with any other replacement in this call.",
      }),
    ),
    newText: Type.Optional(
      Type.String({
        description: "Replacement text for the first replacement.",
      }),
    ),
  },
  {
    additionalProperties: true,
    description:
      "To make multiple replacements in one call, append a number to both " +
      "oldText and newText for each additional block: oldText2/newText2, " +
      "oldText3/newText3, and so on. Every oldText is matched against the " +
      "original file, not incrementally. Do not emit overlapping or nested " +
      "replacements; merge nearby changes into one block instead. Keep each " +
      "oldText as small as possible while still being unique in the file.",
  },
);

export default function (pi: ExtensionAPI) {
  pi.registerTool({
    name: "edit",
    label: "edit",
    description:
      "Edit a single file using exact text replacement. Pass the first " +
      "replacement as oldText/newText; for additional replacements in the " +
      "same call, append a number (oldText2/newText2, oldText3/newText3, " +
      "...). Each oldText must match a unique, non-overlapping region of the " +
      "original file. If two changes affect the same block or nearby lines, " +
      "merge them into one replacement instead of emitting overlapping ones. " +
      "Do not include large unchanged regions just to connect distant changes.",
    promptSnippet:
      "Make precise file edits with exact text replacement, including multiple disjoint edits in one call",
    promptGuidelines: [
      "Use edit for precise changes (oldText must match exactly).",
      "When changing multiple separate locations in one file, pass multiple numbered pairs (oldText/newText, oldText2/newText2, ...) in a single edit call instead of multiple edit calls.",
      "Each oldText is matched against the original file, not after earlier replacements are applied. Do not emit overlapping or nested replacements. Merge nearby changes into one block.",
      "Keep each oldText as small as possible while still being unique in the file. Do not pad with large unchanged regions.",
    ],
    parameters: schema,
    prepareArguments(args) {
      return toEditArgs(args) as any;
    },

    async execute(toolCallId, params, signal, onUpdate, ctx) {
      const def = createEditToolDefinition(ctx.cwd);
      return def.execute(toolCallId, params as EditArgs, signal, onUpdate, ctx);
    },

    renderCall(args, theme, context) {
      const def = createEditToolDefinition(context.cwd);
      const transformed = toEditArgs(args) as EditArgs;
      return def.renderCall!(transformed, theme, { ...context, args: transformed });
    },

    renderResult(result, options, theme, context) {
      const def = createEditToolDefinition(context.cwd);
      const transformed = toEditArgs(context.args) as EditArgs;
      return def.renderResult!(result as any, options, theme, {
        ...context,
        args: transformed,
      });
    },
  });
}