fix(diff): repair truncated Grok diffs with missing markers (#186) (#230)

* fix(diff): repair truncated Grok diffs by reinserting missing markers (#186)

Grok frequently truncates streamed diffs, leaving SEARCH blocks without the
======= separator and/or the >>>>>>> REPLACE closer, which makes applyDiff
fail with 'Expected ======= was not found'. repairTruncatedDiff() detects
incomplete blocks and reinserts the missing markers while preserving valid
blocks and escaped markers.

* test(diff): add fixture-based regression tests for truncated Grok diffs (#186)

Per review feedback: end-to-end applyDiff() regression guards using realistic
truncated-Grok fixtures (missing >>>>>>> REPLACE, missing ======= separator),
plus a well-formed multi-block diff that must pass through unchanged.

* refactor(diff): address CodeRabbit review on truncated-diff repair (#186)

- Use a local repairedDiff in applyDiff instead of reassigning the diffContent
  parameter, keeping the original input observable.
- When a block has a closer but no ======= separator, splice the separator in
  before the existing >>>>>>> REPLACE rather than synthesizing a second closer.
- Strip leading Grok header directives (:start_line:, :end_line:, -------) before
  the first-line-is-SEARCH heuristic so metadata isn't treated as content; the
  directives are preserved on the SEARCH section.

* fix: remove unused needsRepair variable in repairTruncatedDiff

Address review feedback from @edelauna (code review #3284681514):
needsRepair was assigned but never read, making it a dead store.
The variable served no functional purpose in the repair loop,
so it has been removed.

---------

Co-authored-by: Armando Vaquera <263793884+proyectoauraorg@users.noreply.github.com>

Author: proyectoauraorg

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")
+			}
+		})
+	})
 })

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 (!/(?<!\\)<<<<<<< SEARCH/.test(diffContent)) {
+			return diffContent
+		}
+
+		// Split into blocks based on SEARCH markers
+		const blocks = diffContent.split(/(?=(?<!\\)<<<<<<< SEARCH)/)
+
+		let repaired = ""
+
+		for (let i = 0; i < blocks.length; i++) {
+			const block = blocks[i]
+
+			if (block.trim() === "") {
+				continue
+			}
+
+			// Skip prefix blocks that don't contain a SEARCH marker
+			// (e.g., the filename line before the first <<<<<<< SEARCH)
+			if (!/(?<!\\)<<<<<<< SEARCH/.test(block)) {
+				repaired += block
+				continue
+			}
+
+			// Check if this block is complete (has both ======= and >>>>>>> REPLACE)
+			const hasSeparator = /(?<=\n)(?<!\\)=======\s*\n/.test(block)
+			const hasCloser = /(?<=\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<DiffResult> {
-		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)(?<!\\)<<<<<<< SEARCH>?\s*\n((?:\:start_line:\s*(\d+)\s*\n))?((?:\:end_line:\s*(\d+)\s*\n))?((?<!\\)-------\s*\n)?([\s\S]*?)(?:\n)?(?:(?<=\n)(?<!\\)=======\s*\n)([\s\S]*?)(?:\n)?(?:(?<=\n)(?<!\\)>>>>>>> REPLACE)(?=\n|$)/g,
 			),
 		]