fix: sanitize apply content before file edits

core/edit/lazy/applyCodeBlock.ts

@@ -1,5 +1,6 @@
 import { DiffLine, ILLM } from "../..";
 import { generateLines } from "../../diff/util";
+import { stripReasoningFromApplyContent } from "../../util";
 import { supportedLanguages } from "../../util/treeSitter";
 import { getUriFileExtension } from "../../util/uri";
 import { deterministicApplyLazyEdit } from "./deterministic";
@@ -21,6 +22,8 @@ export async function applyCodeBlock(
   isInstantApply: boolean;
   diffLinesGenerator: AsyncGenerator<DiffLine>;
 }> {
+  newLazyFile = stripReasoningFromApplyContent(newLazyFile);
+
   if (canUseInstantApply(filename)) {
     const diffLines = await deterministicApplyLazyEdit({
       oldFile,

core/edit/lazy/applyCodeBlock.vitest.ts

@@ -0,0 +1,83 @@
+import { beforeEach, describe, expect, test, vi } from "vitest";
+
+const { deterministicApplyLazyEdit } = vi.hoisted(() => ({
+  deterministicApplyLazyEdit: vi.fn(),
+}));
+
+vi.mock("../../util/treeSitter", () => ({
+  supportedLanguages: {
+    ts: {},
+  },
+}));
+
+vi.mock("./deterministic", () => ({
+  deterministicApplyLazyEdit,
+}));
+
+import { applyCodeBlock } from "./applyCodeBlock";
+
+describe("applyCodeBlock", () => {
+  beforeEach(() => {
+    deterministicApplyLazyEdit.mockReset();
+    deterministicApplyLazyEdit.mockResolvedValue([
+      { type: "old", line: "export const answer = 41;" },
+      { type: "new", line: "export const answer = 42;" },
+      { type: "same", line: "" },
+    ]);
+  });
+
+  test("strips think blocks before deterministic apply", async () => {
+    const oldFile = "export const answer = 41;\n";
+    const newLazyFile =
+      "<think>Need to update the constant first</think>\nexport const answer = 42;\n";
+
+    const { isInstantApply, diffLinesGenerator } = await applyCodeBlock(
+      oldFile,
+      newLazyFile,
+      "answer.ts",
+      {} as any,
+      new AbortController(),
+    );
+
+    expect(isInstantApply).toBe(true);
+    expect(deterministicApplyLazyEdit).toHaveBeenCalledWith({
+      oldFile,
+      newLazyFile: "\nexport const answer = 42;\n",
+      filename: "answer.ts",
+      onlyFullFileRewrite: true,
+    });
+
+    const diffLines = [];
+    for await (const diffLine of diffLinesGenerator) {
+      diffLines.push(diffLine);
+    }
+    expect(diffLines).toEqual([
+      { type: "old", line: "export const answer = 41;" },
+      { type: "new", line: "export const answer = 42;" },
+      { type: "same", line: "" },
+    ]);
+  });
+
+  test("extracts the final Harmony channel before deterministic apply", async () => {
+    const oldFile = "export const answer = 41;\n";
+    const newLazyFile =
+      "<|start|>assistant<|channel|>analysis<|message|>Thinking through the edit<|end|>" +
+      "<|start|>assistant<|channel|>final<|message|>export const answer = 42;\n<|end|>";
+
+    const { isInstantApply } = await applyCodeBlock(
+      oldFile,
+      newLazyFile,
+      "answer.ts",
+      {} as any,
+      new AbortController(),
+    );
+
+    expect(isInstantApply).toBe(true);
+    expect(deterministicApplyLazyEdit).toHaveBeenCalledWith({
+      oldFile,
+      newLazyFile: "export const answer = 42;\n",
+      filename: "answer.ts",
+      onlyFullFileRewrite: true,
+    });
+  });
+});

core/util/index.ts

@@ -206,6 +206,30 @@ export function removeCodeBlocksAndTrim(text: string): string {
   return processedText.trim();
 }
 
+const THINK_BLOCK_REGEX = /<think>[\s\S]*?<\/think>/gi;
+const HARMONY_TOKEN_REGEX =
+  /<\|(?:start|end|channel|message|constrain|call|return)\|>/g;
+const HARMONY_FINAL_CHANNEL_REGEX =
+  /(?:<\|start\|>[^\n<|]+)?\s*<\|channel\|>\s*final\s*<\|message\|>([\s\S]*?)(?=(?:<\|end\|>|<\|start\|>|$))/gi;
+
+/**
+ * Strips reasoning blocks and Harmony protocol wrappers from model output
+ * before it is treated as file content by apply/edit flows.
+ */
+export function stripReasoningFromApplyContent(content: string): string {
+  let processedText = content.replace(THINK_BLOCK_REGEX, "");
+
+  const finalChannelMatches = Array.from(
+    processedText.matchAll(HARMONY_FINAL_CHANNEL_REGEX),
+  );
+  if (finalChannelMatches.length > 0) {
+    processedText =
+      finalChannelMatches[finalChannelMatches.length - 1][1] ?? "";
+  }
+
+  return processedText.replace(HARMONY_TOKEN_REGEX, "");
+}
+
 export function splitCamelCaseAndNonAlphaNumeric(value: string) {
   return value
     .split(/(?<=[a-z0-9])(?=[A-Z])|[^a-zA-Z0-9]/)

core/util/stripReasoningFromApplyContent.vitest.ts

@@ -0,0 +1,28 @@
+import { describe, expect, it } from "vitest";
+import { stripReasoningFromApplyContent } from "./";
+
+describe("stripReasoningFromApplyContent", () => {
+  it("should remove think blocks before applying file content", () => {
+    const text =
+      "<think>Need to plan the edit first</think>\nexport const answer = 42;\n";
+
+    expect(stripReasoningFromApplyContent(text)).toBe(
+      "\nexport const answer = 42;\n",
+    );
+  });
+
+  it("should extract the final Harmony channel content", () => {
+    const text =
+      "<|start|>assistant<|channel|>analysis<|message|>Thinking about the edit<|end|><|start|>assistant<|channel|>final<|message|>export const answer = 42;\n<|end|>";
+
+    expect(stripReasoningFromApplyContent(text)).toBe(
+      "export const answer = 42;\n",
+    );
+  });
+
+  it("should leave ordinary file content unchanged", () => {
+    const text = "export function greet() {\n  return 'hello';\n}\n";
+
+    expect(stripReasoningFromApplyContent(text)).toBe(text);
+  });
+});

extensions/vscode/src/apply/ApplyManager.ts

@@ -9,7 +9,10 @@ import { generateLines } from "core/diff/util";
 import { ApplyAbortManager } from "core/edit/applyAbortManager";
 import { streamDiffLines } from "core/edit/streamDiffLines";
 import { pruneLinesFromBottom, pruneLinesFromTop } from "core/llm/countTokens";
-import { getMarkdownLanguageTagForFile } from "core/util";
+import {
+  getMarkdownLanguageTagForFile,
+  stripReasoningFromApplyContent,
+} from "core/util";
 import { VerticalDiffManager } from "../diff/vertical/manager";
 import { VsCodeIde } from "../VsCodeIde";
 import { VsCodeWebviewProtocol } from "../webviewProtocol";
@@ -32,6 +35,8 @@ export class ApplyManager {
     toolCallId,
     isSearchAndReplace,
   }: ApplyToFilePayload) {
+    text = stripReasoningFromApplyContent(text);
+
     if (filepath) {
       await this.ensureFileOpen(filepath);
     }