Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
18 kB · 546 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547import { constants, readFileSync as fsReadFileSync, statSync as fsStatSync } from "node:fs";import type { AgentTool } from "@earendil-works/pi-agent-core";import { Box, type Component, Container, Spacer, Text, wrapTextWithAnsi } from "@earendil-works/pi-tui";import { access as fsAccess, readFile as fsReadFile } from "fs/promises";import { type Static, Type } from "typebox";import { renderDiff } from "../../modes/interactive/components/diff.js";import { countChangedLines, formatFileChangeSummaryLine } from "../../modes/interactive/components/edit-summary.js";import { realpathIfPresentSync, writeFileAtomicSync } from "../../utils/atomic-file.js";import type { ToolDefinition } from "../extensions/types.js";import { applyEditsToNormalizedContent, computeEditsDiff, detectLineEnding, type Edit, type EditDiffError, type EditDiffResult, generateDiffString, normalizeToLF, restoreLineEndings, stripBom,} from "./edit-diff.js";import { withFileMutationQueue } from "./file-mutation-queue.js";import { resolveToCwd } from "./path-utils.js";import { invalidArgText, shortenPath, str } from "./render-utils.js";import { wrapToolDefinition } from "./tool-definition-wrapper.js";
type EditPreview = EditDiffResult | EditDiffError;
type EditRenderState = { callComponent?: EditCallRenderComponent;};
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." }), }, { additionalProperties: false },);
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.", }), }, { additionalProperties: false },);
export type EditToolInput = Static<typeof editSchema>;type LegacyEditToolInput = EditToolInput & { oldText?: unknown; newText?: unknown;};
export interface EditToolDetails { /** Unified diff of the changes made */ diff: 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, refusing to replace a changed snapshot. */ writeFile: (absolutePath: string, content: string, expectedContent?: Buffer) => Promise<void>; /** Check if file is readable and writable (throw if not) */ access: (absolutePath: string) => Promise<void>;}
function concurrentEditError(path: string): Error { return new Error(`File changed while editing ${path}; refusing to overwrite concurrent changes.`);}
const defaultEditOperations: EditOperations = { readFile: (path) => fsReadFile(path), writeFile: async (path, content, expectedContent) => { const destination = realpathIfPresentSync(path); const mode = fsStatSync(destination).mode & 0o7777; writeFileAtomicSync(destination, content, { mode, fsync: true, fsyncDir: true, beforeRename: expectedContent ? () => { if (!fsReadFileSync(destination).equals(expectedContent)) { throw concurrentEditError(path); } } : undefined, }); }, access: (path) => fsAccess(path, constants.R_OK | constants.W_OK),};
export interface EditToolOptions { /** 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 if (typeof args.edits === "string") { try { const parsed = JSON.parse(args.edits); if (Array.isArray(parsed)) args.edits = parsed; } catch { // Leave non-JSON input for schema validation to reject. } }
const legacy = args as LegacyEditToolInput; if (typeof legacy.oldText !== "string" || typeof legacy.newText !== "string") { return args as EditToolInput; }
const edits = Array.isArray(legacy.edits) ? [...legacy.edits] : []; edits.push({ oldText: legacy.oldText, newText: legacy.newText }); const { oldText: _oldText, newText: _newText, ...rest } = legacy; return { ...rest, edits } 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 };}
type RenderableEditArgs = { path?: string; file_path?: string; edits?: Edit[]; oldText?: string; newText?: string;};
type EditToolResultLike = { content: Array<{ type: string; text?: string; data?: string; mimeType?: string }>; details?: EditToolDetails;};
type EditCallRenderComponent = Box & { preview?: EditPreview; previewArgsKey?: string; previewPending?: boolean; settledError?: boolean;};
function createEditCallRenderComponent(): EditCallRenderComponent { return Object.assign(new Box(1, 1, (text: string) => text), { preview: undefined as EditPreview | undefined, previewArgsKey: undefined as string | undefined, previewPending: false, settledError: false, });}
function getEditCallRenderComponent(state: EditRenderState, lastComponent: unknown): EditCallRenderComponent { if (lastComponent instanceof Box) { const component = lastComponent as EditCallRenderComponent; state.callComponent = component; return component; } if (state.callComponent) { return state.callComponent; } const component = createEditCallRenderComponent(); state.callComponent = component; return component;}
function getRenderablePreviewInput(args: RenderableEditArgs | undefined): { path: string; edits: Edit[] } | null { if (!args) { return null; }
const path = typeof args.path === "string" ? args.path : typeof args.file_path === "string" ? args.file_path : null; if (!path) { return null; }
if ( Array.isArray(args.edits) && args.edits.length > 0 && args.edits.every((edit) => typeof edit?.oldText === "string" && typeof edit?.newText === "string") ) { return { path, edits: args.edits }; }
if (typeof args.oldText === "string" && typeof args.newText === "string") { return { path, edits: [{ oldText: args.oldText, newText: args.newText }] }; }
return null;}
function formatEditCall( args: RenderableEditArgs | undefined, theme: typeof import("../../modes/interactive/theme/theme.js").theme,): string { const invalidArg = invalidArgText(theme); const rawPath = str(args?.file_path ?? args?.path); const path = rawPath !== null ? shortenPath(rawPath) : null; const pathDisplay = path === null ? invalidArg : path ? theme.fg("dim", path) : theme.fg("toolOutput", "..."); return `${theme.fg("toolTitle", "edit")} ${pathDisplay}`;}
function formatEditResult( args: RenderableEditArgs | undefined, preview: EditPreview | undefined, result: EditToolResultLike, theme: typeof import("../../modes/interactive/theme/theme.js").theme, isError: boolean,): string | undefined { const rawPath = str(args?.file_path ?? args?.path); const previewDiff = preview && !("error" in preview) ? preview.diff : undefined; const previewError = preview && "error" in preview ? preview.error : undefined; if (isError) { const errorText = result.content .filter((c) => c.type === "text") .map((c) => c.text || "") .join("\n"); if (!errorText || errorText === previewError) { return undefined; } return theme.fg("error", errorText); }
const resultDiff = result.details?.diff; if (resultDiff && resultDiff !== previewDiff) { return renderDiff(resultDiff, { filePath: rawPath ?? undefined }); }
return undefined;}
function getEditHeaderBg( preview: EditPreview | undefined, settledError: boolean | undefined, theme: typeof import("../../modes/interactive/theme/theme.js").theme,): (text: string) => string { if (settledError || (preview && "error" in preview)) { return (text: string) => theme.bg("toolErrorBg", text); } if (preview) { return (text: string) => theme.bg("toolSuccessBg", text); } return (text: string) => theme.bg("toolPendingBg", text);}
// Width-aware file summary with optional diff rows at the containing box's content margin.class EditChangeSummaryComponent implements Component { constructor( private readonly rawPath: string, private readonly cwd: string, private readonly change: { added: number; removed: number }, private readonly diffLines: readonly string[] | undefined, ) {}
render(width: number): string[] { const safeWidth = Math.max(1, width); const lines = [formatFileChangeSummaryLine(this.rawPath, this.cwd, this.change, safeWidth)]; if (this.diffLines !== undefined) { for (const line of this.diffLines) { for (const row of wrapTextWithAnsi(line, safeWidth)) { lines.push(row); } } } return lines; }
invalidate(): void {}}
function buildEditCallComponent( component: EditCallRenderComponent, args: RenderableEditArgs | undefined, theme: typeof import("../../modes/interactive/theme/theme.js").theme, expanded: boolean, cwd: string,): EditCallRenderComponent { component.setBgFn(getEditHeaderBg(component.preview, component.settledError, theme)); component.clear(); component.addChild(new Text(formatEditCall(args, theme), 0, 0));
if (component.preview && "error" in component.preview) { component.addChild(new Spacer(1)); component.addChild(new Text(theme.fg("error", component.preview.error), 0, 0)); return component; } // A failed execution must not present the predicted diff as applied changes. if (!component.preview || component.settledError) { return component; }
// Keep the file summary visible while detail expansion reveals the diff. const rawPath = str(args?.file_path ?? args?.path); const change = countChangedLines(component.preview.diff); component.addChild(new Spacer(1)); component.addChild( new EditChangeSummaryComponent( rawPath ?? "...", cwd, change, expanded ? renderDiff(component.preview.diff).split("\n") : undefined, ), ); return component;}
function setEditPreview( component: EditCallRenderComponent, preview: EditPreview, argsKey: string | undefined,): boolean { const current = component.preview; const changed = current === undefined || ("error" in current && "error" in preview ? current.error !== preview.error : "error" in current !== "error" in preview) || (!("error" in current) && !("error" in preview) && (current.diff !== preview.diff || current.firstChangedLine !== preview.firstChangedLine)); component.preview = preview; component.previewArgsKey = argsKey; component.previewPending = false; return changed;}
export function createEditToolDefinition( cwd: string, options?: EditToolOptions,): ToolDefinition<typeof editSchema, EditToolDetails | undefined, EditRenderState> { const ops = options?.operations ?? defaultEditOperations; const definition: ToolDefinition<typeof editSchema, EditToolDetails | undefined, EditRenderState> = { 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: "Make precise file edits with exact text replacement, including multiple disjoint edits in one call", parameters: editSchema, renderShell: "self", prepareArguments: prepareEditArguments, async execute(_toolCallId, input: EditToolInput, signal?: AbortSignal, _onUpdate?, _ctx?) { const { path, edits } = validateEditInput(input); const absolutePath = resolveToCwd(path, cwd);
return withFileMutationQueue( absolutePath, () => new Promise<{ content: Array<{ type: "text"; text: string }>; details: EditToolDetails | undefined; }>((resolve, reject) => { if (signal?.aborted) { reject(new Error("Operation aborted")); return; }
let aborted = false;
const onAbort = () => { aborted = true; reject(new Error("Operation aborted")); };
if (signal) { signal.addEventListener("abort", onAbort, { once: true }); }
void (async () => { try { try { await ops.access(absolutePath); } catch (error: unknown) { const errorMessage = error instanceof Error && "code" in error ? `Error code: ${error.code}` : String(error); if (signal) { signal.removeEventListener("abort", onAbort); } reject(new Error(`Could not edit file: ${path}. ${errorMessage}.`)); return; }
if (aborted) { return; }
const buffer = await ops.readFile(absolutePath); const rawContent = buffer.toString("utf-8");
if (aborted) { return; }
// Strip BOM before matching. The model will not include an invisible BOM in oldText. const { bom, text: content } = stripBom(rawContent); const originalEnding = detectLineEnding(content); const normalizedContent = normalizeToLF(content); const { baseContent, newContent } = applyEditsToNormalizedContent( normalizedContent, edits, path, );
if (aborted) { return; }
const finalContent = bom + restoreLineEndings(newContent, originalEnding); // Re-read immediately before committing. The default writer repeats this // check in its pre-rename hook; this also protects custom writers. const latestBuffer = await ops.readFile(absolutePath); if (!latestBuffer.equals(buffer)) { throw concurrentEditError(path); } if (signal) { signal.removeEventListener("abort", onAbort); } if (aborted) { return; } // Atomic replacement is the commit point. Once it starts, report its // outcome instead of turning a successful mutation into a retryable abort. await ops.writeFile(absolutePath, finalContent, buffer);
const diffResult = generateDiffString(baseContent, newContent); resolve({ content: [ { type: "text", text: `Successfully replaced ${edits.length} block(s) in ${path}.`, }, ], details: { diff: diffResult.diff, firstChangedLine: diffResult.firstChangedLine }, }); } catch (error: unknown) { if (signal) { signal.removeEventListener("abort", onAbort); }
if (!aborted) { reject(error instanceof Error ? error : new Error(String(error))); } } })(); }), ); }, renderCall(args, theme, context) { const component = getEditCallRenderComponent(context.state, context.lastComponent); const previewInput = getRenderablePreviewInput(args as RenderableEditArgs | undefined); const argsKey = previewInput ? JSON.stringify({ path: previewInput.path, edits: previewInput.edits }) : undefined;
if (component.previewArgsKey !== argsKey) { component.preview = undefined; component.previewArgsKey = argsKey; component.previewPending = false; component.settledError = false; }
if (context.argsComplete && previewInput && !component.preview && !component.previewPending) { component.previewPending = true; const requestKey = argsKey; void computeEditsDiff(previewInput.path, previewInput.edits, context.cwd).then((preview) => { if (component.previewArgsKey === requestKey) { setEditPreview(component, preview, requestKey); context.invalidate(); } }); }
return buildEditCallComponent(component, args, theme, context.expanded, context.cwd); }, renderResult(result, _options, theme, context) { const callComponent = context.state.callComponent; const previewInput = getRenderablePreviewInput(context.args as RenderableEditArgs | undefined); const argsKey = previewInput ? JSON.stringify({ path: previewInput.path, edits: previewInput.edits }) : undefined; const typedResult = result as EditToolResultLike; const resultDiff = !context.isError ? typedResult.details?.diff : undefined; let changed = false; if (callComponent) { if (typeof resultDiff === "string") { changed = setEditPreview( callComponent, { diff: resultDiff, firstChangedLine: typedResult.details?.firstChangedLine }, argsKey, ) || changed; } if (callComponent.settledError !== context.isError) { callComponent.settledError = context.isError; changed = true; } if (changed) { buildEditCallComponent( callComponent, context.args as RenderableEditArgs | undefined, theme, context.expanded, context.cwd, ); } }
const output = formatEditResult(context.args, callComponent?.preview, typedResult, theme, context.isError); const component = (context.lastComponent as Container | undefined) ?? new Container(); component.clear(); if (!output) { return component; } component.addChild(new Spacer(1)); component.addChild(new Text(output, 1, 0)); return component; }, }; return Object.assign(definition, { replayBuiltInToolName: "edit" as const });}
export function createEditTool(cwd: string, options?: EditToolOptions): AgentTool<typeof editSchema> { return wrapToolDefinition(createEditToolDefinition(cwd, options));}