locode/src/tools/multiEdit.ts

import { createPatch } from "diff";
import { randomBytes } from "node:crypto";
import { readFile as fsReadFile, rename as fsRename, unlink as fsUnlink, writeFile as fsWriteFile } from "node:fs/promises";
import path from "node:path";
import { z } from "zod";
import { assertWithinWorkspace } from "../utils/path.js";
import { applyEdit, countOccurrences } from "./editFile.js";
import type { ToolDef } from "./types.js";

// A single edit within a multi_edit batch. Mirrors edit_file's args minus `path` (which is shared
// across the whole batch). Each edit is applied in array order to the result of the previous one.
const editSchema = z.object({
  old_string: z.string().describe("Exact text to replace. Must match the current file content exactly at this point in the batch."),
  new_string: z.string().describe("Replacement text."),
  replace_all: z.boolean().optional().describe("Replace every occurrence instead of requiring a unique match."),
});

const schema = z.object({
  path: z.string().describe("File path to edit, relative to the working directory or absolute."),
  edits: z.array(editSchema).min(1).describe("Ordered list of edits to apply to the same file, one after another."),
});

/** Applies a batch of edits to an in-memory string, validating each. Throws on the first edit that
 *  doesn't match uniquely (unless its replace_all is set) or doesn't match at all. Edits apply to the
 *  running result, so an earlier edit can change the text a later edit matches. */
function applyBatch(original: string, edits: { old_string: string; new_string: string; replace_all?: boolean }[], filePath: string): string {
  let current = original;
  edits.forEach((edit, i) => {
    const occurrences = countOccurrences(current, edit.old_string);
    if (occurrences === 0) {
      throw new Error(
        `Edit ${i + 1}: old_string not found in ${filePath}. Earlier edits may have shifted the text — re-read the file and adjust. Make sure it matches exactly, including whitespace.`,
      );
    }
    if (occurrences > 1 && !edit.replace_all) {
      throw new Error(
        `Edit ${i + 1}: old_string appears ${occurrences} times in ${filePath}. Provide more surrounding context to make it unique, or set replace_all: true.`,
      );
    }
    current = applyEdit(current, edit.old_string, edit.new_string, edit.replace_all);
  });
  return current;
}

export const multiEditTool: ToolDef<z.infer<typeof schema>> = {
  name: "multi_edit",
  description:
    "Apply several edits to the same file in one call, in order. Each edit is {old_string, new_string, replace_all?}. " +
    "Use this instead of repeated edit_file calls when you have multiple distinct changes to one file — it's one confirmation " +
    "and one atomic write. Each old_string must match uniquely at its point in the batch (unless replace_all is set). Read the file first.",
  schema,
  mutating: true,
  preview: async ({ path: filePath, edits }, ctx) => {
    const resolved = path.resolve(ctx.cwd, filePath);
    assertWithinWorkspace(resolved, ctx.cwd, filePath);
    let original: string;
    try {
      original = await fsReadFile(resolved, "utf-8");
    } catch {
      return `File ${resolved} does not exist.`;
    }
    try {
      const updated = applyBatch(original, edits, filePath);
      return createPatch(resolved, original, updated, "", "");
    } catch (err) {
      return `Warning: ${(err as Error).message} — this edit will fail.`;
    }
  },
  handler: async ({ path: filePath, edits }, ctx) => {
    const resolved = path.resolve(ctx.cwd, filePath);
    assertWithinWorkspace(resolved, ctx.cwd, filePath);
    const original = await fsReadFile(resolved, "utf-8");
    const updated = applyBatch(original, edits, filePath);
    // Atomic write via temp+rename (same rationale as edit_file): a crash mid-write can't leave the
    // user's source file half-overwritten — the live file stays intact until the rename swaps in the
    // full new content. Clean up the temp file if anything fails.
    const tmp = `${resolved}.locode-${randomBytes(4).toString("hex")}.tmp`;
    try {
      await fsWriteFile(tmp, updated, "utf-8");
      await fsRename(tmp, resolved);
    } catch (err) {
      await fsUnlink(tmp).catch(() => {});
      throw err;
    }
    if (ctx.setLastEdit) {
      ctx.setLastEdit({ path: filePath, previousContent: original });
    }
    return { path: resolved, applied: edits.length };
  },
};