packages/coding-agent/src/core/tools/pi/edit.ts

// Adapted from Pi c20cb09772bf4e2590a316cb54514cef76df4293.
// Copyright (c) 2025 Mario Zechner. MIT license: licenses/third-party-notices.txt.

import { constants } from "node:fs";
import { access as fsAccess, writeFile as fsWriteFile } from "node:fs/promises";
import { type Static, Type } from "typebox";
import type { ToolContext, ToolDefinition } from "../../types.ts";
import {
  applyEditsToNormalizedContent,
  type Edit,
  generateDiffString,
  generateUnifiedPatch,
  normalizeToLF,
} from "./edit-diff.ts";
import { withFileMutationQueue } from "./file-mutation-queue.ts";
import type { FileReadState } from "./file-read-state.ts";
import { resolveToCwd } from "./path-utils.ts";
import { readRegularFile } from "./utils/read-file.ts";
import { splitBom } from "./utils/text.ts";

const MAX_EDIT_DETAIL_BYTES = 64 * 1024;
const OMITTED_EDIT_DIFF = "[Diff omitted: exceeds the inline display or computation limit.]";

const replaceEditSchema = Type.Object(
  {
    oldText: Type.String({
      description:
        "Exact text for one targeted replacement. It must be unique in the original file and must not overlap with any other edits[].oldText in the same call.",
    }),
    newText: Type.String({ description: "Replacement text for this targeted edit." }),
  },
  {},
);

const editSchema = Type.Object(
  {
    path: Type.String({ description: "Path to the file to edit (relative or absolute)" }),
    edits: Type.Array(replaceEditSchema, {
      description:
        "One or more targeted replacements. Each edit is matched against the original file, not incrementally. Do not include overlapping or nested edits. If two changes touch the same block or nearby lines, merge them into one edit instead.",
    }),
  },
  {},
);

export const editToolSystemPromptContribution = {
  snippet:
    "Make precise file edits with exact text replacement, including multiple disjoint edits in one call",
  guidelines: [
    "Use edit for precise changes (edits[].oldText must match exactly)",
    "When changing multiple separate locations in one file, use one edit call with multiple entries in edits[] instead of multiple edit calls",
    "Each edits[].oldText is matched against the original file, not after earlier edits are applied. Do not emit overlapping or nested edits. Merge nearby changes into one edit.",
    "Keep edits[].oldText as small as possible while still being unique in the file. Do not pad with large unchanged regions.",
  ],
} as const;

export type EditToolInput = Static<typeof editSchema>;

type SingleEditInput = { oldText: string; newText: string };

function isSingleEditInput(value: unknown): value is SingleEditInput {
  if (!value || typeof value !== "object" || Array.isArray(value)) {
    return false;
  }

  const edit = value as Record<string, unknown>;
  return typeof edit.oldText === "string" && typeof edit.newText === "string";
}

export interface EditToolDetails {
  /** Display-oriented diff of the changes made */
  diff: string;
  /** Standard unified patch of the changes made */
  patch: string;
  /** Line number of the first change in the new file (for editor navigation) */
  firstChangedLine?: number;
}

/**
 * Pluggable operations for the edit tool.
 * Override these to delegate file editing to remote systems (for example SSH).
 */
export interface EditOperations {
  /** Read file contents as a Buffer */
  readFile: (absolutePath: string) => Promise<Buffer>;
  /** Write content to a file */
  writeFile: (absolutePath: string, content: string) => Promise<void>;
  /** Check if file is readable and writable (throw if not) */
  access: (absolutePath: string) => Promise<void>;
}

const defaultEditOperations: EditOperations = {
  readFile: readRegularFile,
  writeFile: (path, content) => fsWriteFile(path, content, "utf-8"),
  access: (path) => fsAccess(path, constants.R_OK | constants.W_OK),
};

export interface EditToolOptions {
  /** Shared knowledge of file contents in the current model context. */
  fileReadState?: FileReadState;
  /** Custom operations for file editing. Default: local filesystem */
  operations?: EditOperations;
}

function prepareEditArguments(input: unknown): EditToolInput {
  if (!input || typeof input !== "object") {
    return input as EditToolInput;
  }

  const args = input as Record<string, unknown>;

  // Some models (Opus 4.6, GLM-5.1) send edits as a JSON string instead of an array.
  // Others send a single edit object instead of a one-element edits array.
  if (typeof args.edits === "string") {
    try {
      const parsed = JSON.parse(args.edits);
      if (Array.isArray(parsed)) {
        args.edits = parsed;
      } else if (isSingleEditInput(parsed)) {
        args.edits = [parsed];
      }
    } catch {}
  } else if (isSingleEditInput(args.edits)) {
    args.edits = [args.edits];
  }

  return args as EditToolInput;
}

function validateEditInput(input: EditToolInput): { path: string; edits: Edit[] } {
  if (!Array.isArray(input.edits) || input.edits.length === 0) {
    throw new Error("Edit tool input is invalid. edits must contain at least one replacement.");
  }
  return { path: input.path, edits: input.edits };
}

export function createEditToolDefinition(cwd: string, options?: EditToolOptions) {
  const ops = options?.operations ?? defaultEditOperations;
  return {
    name: "edit",
    label: "edit",
    description:
      "Edit a single file using exact text replacement. Every edits[].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 edit instead of emitting overlapping edits. Do not include large unchanged regions just to connect distant changes.",
    promptSnippet: editToolSystemPromptContribution.snippet,
    promptGuidelines: [...editToolSystemPromptContribution.guidelines],
    parameters: editSchema,
    prepareArguments: prepareEditArguments,
    async execute(
      _toolCallId,
      input: EditToolInput,
      signal?: AbortSignal,
      _onUpdate?,
      ctx?: Partial<ToolContext>,
    ) {
      const { path, edits } = validateEditInput(input);
      const absolutePath = resolveToCwd(path, ctx?.cwd || cwd);

      return withFileMutationQueue(absolutePath, async () => {
        // Do not reject from an abort event listener here: that would release the
        // mutation queue while an in-flight filesystem operation may still finish.
        // Check cancellation before writing, and keep the queue locked until the
        // write settles. A successful write must still be reported as successful.
        const throwIfAborted = (): void => {
          if (signal?.aborted) throw new Error("Operation aborted");
        };

        throwIfAborted();

        // Check if file exists.
        try {
          await ops.access(absolutePath);
        } catch (error: unknown) {
          throwIfAborted();
          const errorMessage =
            error instanceof Error && "code" in error ? `Error code: ${error.code}` : String(error);
          throw new Error(`Could not edit file: ${path}. ${errorMessage}.`);
        }
        throwIfAborted();

        // Read the file.
        const buffer = await ops.readFile(absolutePath);
        const rawContent = buffer.toString("utf-8");
        throwIfAborted();

        // Strip BOM before matching. The model will not include an invisible BOM in oldText.
        const { bom, text: content } = splitBom(rawContent);
        const normalizedContent = normalizeToLF(content);
        const { baseContent, newContent } = applyEditsToNormalizedContent(
          normalizedContent,
          edits,
          path,
          content,
        );
        throwIfAborted();

        const finalContent = bom + newContent;
        // The queue covers ZPI writers only. Reject changes observed during edit preparation
        // rather than overwrite them; this is not an atomic lock against external writers.
        const current = await ops.readFile(absolutePath);
        throwIfAborted();
        if (!current.equals(buffer))
          throw new Error(`File changed while preparing edits: ${path}. Read it again before retrying.`);
        await ops.writeFile(absolutePath, finalContent);
        options?.fileReadState?.recordEdit(absolutePath, rawContent, finalContent);

        const diffResult = generateDiffString(baseContent, newContent);
        const patch =
          diffResult.diff === undefined ? undefined : generateUnifiedPatch(path, baseContent, newContent);
        // Long changed lines appear in both display formats. Bound only these
        // previews, not the independent fileChange snapshots used for rewind.
        const diffTooLarge =
          diffResult.diff === undefined || Buffer.byteLength(diffResult.diff) > MAX_EDIT_DETAIL_BYTES;
        const patchTooLarge = patch === undefined || Buffer.byteLength(patch) > MAX_EDIT_DETAIL_BYTES;
        return {
          content: [
            {
              type: "text",
              text:
                `Successfully replaced ${edits.length} block(s) in ${path}. ` +
                "No read-back needed while the edited text remains in context; re-read if it is missing or may have changed.",
            },
          ],
          details: {
            diff: diffTooLarge || patchTooLarge ? OMITTED_EDIT_DIFF : (diffResult.diff ?? OMITTED_EDIT_DIFF),
            patch: patchTooLarge ? "" : (patch ?? ""),
            firstChangedLine: diffResult.firstChangedLine,
          },
        };
      });
    },
  } satisfies ToolDefinition<typeof editSchema, EditToolDetails | undefined>;
}