diff --git a/src/core/diff/strategies/__tests__/multi-search-replace.spec.ts b/src/core/diff/strategies/__tests__/multi-search-replace.spec.ts index 3b6696e387..b58e331f8f 100644 --- a/src/core/diff/strategies/__tests__/multi-search-replace.spec.ts +++ b/src/core/diff/strategies/__tests__/multi-search-replace.spec.ts @@ -1204,4 +1204,174 @@ function sum(a, b) { expect(result.error).toContain(":start_line:5 <-- Invalid location") }) }) + + describe("repairTruncatedDiff", () => { + let strategy: MultiSearchReplaceDiffStrategy + + beforeEach(() => { + strategy = new MultiSearchReplaceDiffStrategy() + }) + + it("should not modify a complete diff", () => { + const diff = "<<<<<<< SEARCH\n" + "original content\n" + "=======\n" + "new content\n" + ">>>>>>> REPLACE" + const result = strategy["repairTruncatedDiff"](diff) + expect(result).toBe(diff) + }) + + it("should not modify a diff with multiple complete blocks", () => { + const diff = + "<<<<<<< SEARCH\n" + + "content1\n" + + "=======\n" + + "new1\n" + + ">>>>>>> REPLACE\n\n" + + "<<<<<<< SEARCH\n" + + "content2\n" + + "=======\n" + + "new2\n" + + ">>>>>>> REPLACE" + const result = strategy["repairTruncatedDiff"](diff) + expect(result).toBe(diff) + }) + + it("should repair diff missing ======= and >>>>>>> REPLACE", () => { + const diff = "<<<<<<< SEARCH\n" + "original content\n" + "new content" + const result = strategy["repairTruncatedDiff"](diff) + expect(result).toBe( + "<<<<<<< SEARCH\n" + "original content\n" + "=======\n" + "new content\n" + ">>>>>>> REPLACE", + ) + }) + + it("should repair diff missing only >>>>>>> REPLACE", () => { + const diff = "<<<<<<< SEARCH\n" + "original content\n" + "=======\n" + "new content" + const result = strategy["repairTruncatedDiff"](diff) + expect(result).toBe( + "<<<<<<< SEARCH\n" + "original content\n" + "=======\n" + "new content\n" + ">>>>>>> REPLACE", + ) + }) + + it("should repair first truncated block while preserving subsequent complete blocks", () => { + const diff = + "<<<<<<< SEARCH\n" + + "content1\n" + + "new1\n\n" + + "<<<<<<< SEARCH\n" + + "content2\n" + + "=======\n" + + "new2\n" + + ">>>>>>> REPLACE" + const result = strategy["repairTruncatedDiff"](diff) + expect(result).toBe( + "<<<<<<< SEARCH\n" + + "content1\n" + + "=======\n" + + "new1\n" + + ">>>>>>> REPLACE\n\n" + + "<<<<<<< SEARCH\n" + + "content2\n" + + "=======\n" + + "new2\n" + + ">>>>>>> REPLACE", + ) + }) + + it("should handle empty search content with missing replace marker", () => { + const diff = "<<<<<<< SEARCH\n" + "replacement text" + const result = strategy["repairTruncatedDiff"](diff) + expect(result).toBe("<<<<<<< SEARCH\n" + "=======\n" + "replacement text\n" + ">>>>>>> REPLACE") + }) + + it("should not add trailing newline if content already ends with one", () => { + const diff = "<<<<<<< SEARCH\n" + "original\n" + "=======\n" + "new content\n" + const result = strategy["repairTruncatedDiff"](diff) + expect(result).toBe("<<<<<<< SEARCH\n" + "original\n" + "=======\n" + "new content\n" + ">>>>>>> REPLACE") + }) + + it("inserts ======= before an existing closer instead of synthesizing a second one", () => { + // Has >>>>>>> REPLACE but no ======= separator. + const diff = "<<<<<<< SEARCH\n" + "old line\n" + ">>>>>>> REPLACE" + const result = strategy["repairTruncatedDiff"](diff) + expect(result).toBe("<<<<<<< SEARCH\n" + "old line\n" + "=======\n" + ">>>>>>> REPLACE") + // Exactly one closer, exactly one separator. + expect(result.match(/>>>>>>> REPLACE/g)).toHaveLength(1) + expect(result.match(/^=======$/gm)).toHaveLength(1) + }) + + it("preserves :start_line: / ------- directives instead of treating them as SEARCH content", () => { + const diff = "<<<<<<< SEARCH\n" + ":start_line:5\n" + "-------\n" + "old line\n" + "new line" + const result = strategy["repairTruncatedDiff"](diff) + expect(result).toBe( + "<<<<<<< SEARCH\n" + + ":start_line:5\n" + + "-------\n" + + "old line\n" + + "=======\n" + + "new line\n" + + ">>>>>>> REPLACE", + ) + }) + + it("treats a single content line after a directive header as the SEARCH target", () => { + const diff = "<<<<<<< SEARCH\n" + ":start_line:5\n" + "-------\n" + "old line" + const result = strategy["repairTruncatedDiff"](diff) + expect(result).toBe( + "<<<<<<< SEARCH\n" + + ":start_line:5\n" + + "-------\n" + + "old line\n" + + "=======\n" + + "\n" + + ">>>>>>> REPLACE", + ) + }) + }) + + // Regression guards for #186: Grok sometimes truncates the streamed diff and drops + // the closing markers, which previously surfaced as "Unable to apply diff - Expected + // '=======' was not found". These fixtures exercise the full applyDiff() path end-to-end. + describe("truncated Grok diff regression (#186)", () => { + const grokStrategy = new MultiSearchReplaceDiffStrategy(1.0, 5) + const originalContent = 'function greet() {\n\treturn "hello"\n}\n' + const expectedContent = 'function greet() {\n\treturn "hi there"\n}\n' + + it("applies a diff whose closing >>>>>>> REPLACE marker was truncated", async () => { + const diff = + "src/greet.ts\n" + "<<<<<<< SEARCH\n" + '\treturn "hello"\n' + "=======\n" + '\treturn "hi there"' + const result = await grokStrategy.applyDiff(originalContent, diff) + expect(result.success).toBe(true) + if (result.success) { + expect(result.content).toBe(expectedContent) + } + }) + + it("applies a diff truncated before the ======= separator", async () => { + const diff = "src/greet.ts\n" + "<<<<<<< SEARCH\n" + '\treturn "hello"\n' + '\treturn "hi there"' + const result = await grokStrategy.applyDiff(originalContent, diff) + expect(result.success).toBe(true) + if (result.success) { + expect(result.content).toBe(expectedContent) + } + }) + + it("leaves a well-formed multi-block diff unchanged", async () => { + const multiBlock = "export const a = 1\nexport const b = 2\n" + const diff = + "src/consts.ts\n" + + "<<<<<<< SEARCH\n" + + "export const a = 1\n" + + "=======\n" + + "export const a = 10\n" + + ">>>>>>> REPLACE\n" + + "<<<<<<< SEARCH\n" + + "export const b = 2\n" + + "=======\n" + + "export const b = 20\n" + + ">>>>>>> REPLACE" + const result = await grokStrategy.applyDiff(multiBlock, diff) + expect(result.success).toBe(true) + if (result.success) { + expect(result.content).toBe("export const a = 10\nexport const b = 20\n") + } + }) + }) }) diff --git a/src/core/diff/strategies/multi-search-replace.ts b/src/core/diff/strategies/multi-search-replace.ts index f43bbee0dc..85460da3a7 100644 --- a/src/core/diff/strategies/multi-search-replace.ts +++ b/src/core/diff/strategies/multi-search-replace.ts @@ -242,13 +242,117 @@ export class MultiSearchReplaceDiffStrategy implements DiffStrategy { } } + /** + * Repairs truncated diffs (common with Grok) by adding missing ======= and >>>>>>> REPLACE markers. + * When the model's output gets cut off mid-stream, the diff may end after SEARCH content + * without the separator or closing marker. This method detects that pattern and appends + * the missing markers so the diff can still be parsed and applied. + */ + private repairTruncatedDiff(diffContent: string): string { + // Only repair if the diff has at least one SEARCH marker + if (!/(?>>>>>> REPLACE) + const hasSeparator = /(?<=\n)(?>>>>>> REPLACE(?=\n|$)/.test(block) + + if (hasSeparator && hasCloser) { + // Block is complete — emit verbatim (keeps its own trailing separator) + repaired += block + continue + } + + // Block needs repair. Build a clean block ending at >>>>>>> REPLACE, then + // re-add an inter-block separator if more (non-empty) blocks follow, so the + // appended closer never gets glued to the next "<<<<<<< SEARCH". + const isLast = blocks.slice(i + 1).every((b) => b.trim() === "") + const separator = isLast ? "" : "\n\n" + + if (hasSeparator && !hasCloser) { + // Has ======= but missing >>>>>>> REPLACE — append closing marker + const body = block.replace(/\s+$/, "") + repaired += body + "\n>>>>>>> REPLACE" + separator + } else if (hasCloser && !hasSeparator) { + // Has >>>>>>> REPLACE but missing the ======= separator. Don't synthesize a + // second closer; splice the separator in right before the existing closer so + // everything above it becomes the SEARCH section. + const body = block.replace(/\s+$/, "") + repaired += body.replace(/(\n)(>>>>>>> REPLACE)(?=\n|$)/, "$1=======\n$2") + separator + } else { + // Missing both ======= and >>>>>>> REPLACE. + const searchMatch = block.match(/^<<<<<<< SEARCH\n?([\s\S]*)$/) + let content = (searchMatch?.[1] ?? "").replace(/\s+$/, "") + + // Peel off any leading Grok header directives (:start_line:, :end_line:, -------) + // so the "first line is SEARCH" heuristic sees real content, not metadata. The + // directives are preserved as a header on the SEARCH section. + let header = "" + const directiveLine = /^(?::start_line:\s*\d+|:end_line:\s*\d+|-------)\s*$/ + let nlIdx: number + while ((nlIdx = content.indexOf("\n")) !== -1 && directiveLine.test(content.slice(0, nlIdx))) { + header += content.slice(0, nlIdx + 1) + content = content.slice(nlIdx + 1) + } + + const firstNewlineIdx = content.indexOf("\n") + if (firstNewlineIdx !== -1) { + // First line is SEARCH content, rest is REPLACE content + const searchContent = content.substring(0, firstNewlineIdx) + const replaceContent = content.substring(firstNewlineIdx + 1) + repaired += + "<<<<<<< SEARCH\n" + + header + + searchContent + + "\n=======\n" + + replaceContent + + "\n>>>>>>> REPLACE" + + separator + } else if (header) { + // Only a directive header plus a single content line: that line is the SEARCH + // target (the user pinned it with start_line); the REPLACE section is empty. + repaired += "<<<<<<< SEARCH\n" + header + content + "\n=======\n\n>>>>>>> REPLACE" + separator + } else { + // Single line — treat as empty SEARCH with content as REPLACE + repaired += "<<<<<<< SEARCH\n=======\n" + content + "\n>>>>>>> REPLACE" + separator + } + } + } + + return repaired || diffContent + } + async applyDiff( originalContent: string, diffContent: string, _paramStartLine?: number, _paramEndLine?: number, ): Promise { - const validseq = this.validateMarkerSequencing(diffContent) + // Repair truncated diffs before validation (common with Grok and other models + // whose output gets cut off mid-stream, leaving missing ======= and >>>>>>> REPLACE markers) + const repairedDiff = this.repairTruncatedDiff(diffContent) + + const validseq = this.validateMarkerSequencing(repairedDiff) if (!validseq.success) { return { success: false, @@ -288,7 +392,7 @@ export class MultiSearchReplaceDiffStrategy implements DiffStrategy { */ let matches = [ - ...diffContent.matchAll( + ...repairedDiff.matchAll( /(?:^|\n)(??\s*\n((?:\:start_line:\s*(\d+)\s*\n))?((?:\:end_line:\s*(\d+)\s*\n))?((?>>>>>> REPLACE)(?=\n|$)/g, ), ]