diff --git a/core/edit/lazy/applyCodeBlock.ts b/core/edit/lazy/applyCodeBlock.ts index b58b9d4ec25..45258dbc364 100644 --- a/core/edit/lazy/applyCodeBlock.ts +++ b/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; }> { + newLazyFile = stripReasoningFromApplyContent(newLazyFile); + if (canUseInstantApply(filename)) { const diffLines = await deterministicApplyLazyEdit({ oldFile, diff --git a/core/edit/lazy/applyCodeBlock.vitest.ts b/core/edit/lazy/applyCodeBlock.vitest.ts new file mode 100644 index 00000000000..ad7ad56bd9f --- /dev/null +++ b/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 = + "Need to update the constant first\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, + }); + }); +}); diff --git a/core/util/index.ts b/core/util/index.ts index 76b319bddf0..e11f43b4fa9 100644 --- a/core/util/index.ts +++ b/core/util/index.ts @@ -206,6 +206,30 @@ export function removeCodeBlocksAndTrim(text: string): string { return processedText.trim(); } +const THINK_BLOCK_REGEX = /[\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]/) diff --git a/core/util/stripReasoningFromApplyContent.vitest.ts b/core/util/stripReasoningFromApplyContent.vitest.ts new file mode 100644 index 00000000000..83bfc0f2437 --- /dev/null +++ b/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 = + "Need to plan the edit first\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); + }); +}); diff --git a/extensions/vscode/src/apply/ApplyManager.ts b/extensions/vscode/src/apply/ApplyManager.ts index de8344e1889..bf4cf672b38 100644 --- a/extensions/vscode/src/apply/ApplyManager.ts +++ b/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); }