change to reject diffs approach rather than cancel apply approach

RomneyDa · RomneyDa · commit a1dc1f8aa249 · 2025-05-29T16:04:35.000-07:00
diff --git a/core/core.ts b/core/core.ts
@@ -53,6 +53,7 @@ import {
 } from "./config/onboarding";
 import { createNewWorkspaceBlockFile } from "./config/workspace/workspaceBlocks";
 import { MCPManagerSingleton } from "./context/mcp/MCPManagerSingleton";
+import { ApplyAbortManager } from "./edit/applyAbortManager";
 import { streamDiffLines } from "./edit/streamDiffLines";
 import { shouldIgnore } from "./indexing/shouldIgnore";
 import { walkDirCache } from "./indexing/walkDir";
@@ -61,7 +62,6 @@ import { LLMLogger } from "./llm/logger";
 import { llmStreamChat } from "./llm/streamChat";
 import type { FromCoreProtocol, ToCoreProtocol } from "./protocol";
 import type { IMessenger, Message } from "./protocol/messenger";
-import { StreamAbortManager } from "./util/abortManager";
 
 export class Core {
   configHandler: ConfigHandler;
@@ -509,6 +509,11 @@ export class Core {
         throw new Error("No model selected");
       }
 
+      const abortManager = ApplyAbortManager.getInstance();
+      const abortController = abortManager.get(
+        data.fileUri ?? "current-file-stream",
+      ); // not super important since currently cancelling apply will cancel all streams it's one file at a time
+
       return streamDiffLines({
         highlighted: data.highlighted,
         prefix: data.prefix,
@@ -522,13 +527,13 @@ export class Core {
         language: data.language,
         onlyOneInsertion: false,
         overridePrompt: undefined,
-        abortControllerId: data.fileUri ?? "current-file-stream", // not super important since currently cancelling apply will cancel all streams it's one file at a time
+        abortController,
       });
     });
 
     on("cancelApply", async (msg) => {
-      const abortManager = StreamAbortManager.getInstance();
-      abortManager.clear();
+      const abortManager = ApplyAbortManager.getInstance();
+      abortManager.clear(); // for now abort all streams
     });
 
     on("completeOnboarding", this.handleCompleteOnboarding.bind(this));
diff --git a/core/edit/applyAbortManager.ts b/core/edit/applyAbortManager.ts
@@ -1,16 +1,16 @@
-export class StreamAbortManager {
-  private static instance: StreamAbortManager;
+export class ApplyAbortManager {
+  private static instance: ApplyAbortManager;
   private controllers: Map<string, AbortController>;
 
   private constructor() {
     this.controllers = new Map();
   }
 
-  public static getInstance(): StreamAbortManager {
-    if (!StreamAbortManager.instance) {
-      StreamAbortManager.instance = new StreamAbortManager();
+  public static getInstance(): ApplyAbortManager {
+    if (!ApplyAbortManager.instance) {
+      ApplyAbortManager.instance = new ApplyAbortManager();
     }
-    return StreamAbortManager.instance;
+    return ApplyAbortManager.instance;
   }
 
   public get(id: string): AbortController {
diff --git a/core/edit/recursiveStream.ts b/core/edit/recursiveStream.ts
@@ -68,6 +68,7 @@ export async function* recursiveStream(
     });
 
     for await (const chunk of generator) {
+      console.log(chunk);
       yield chunk;
       const rendered = renderChatMessage(chunk);
       buffer += rendered;
diff --git a/core/edit/streamDiffLines.ts b/core/edit/streamDiffLines.ts
@@ -20,7 +20,6 @@ import { streamDiff } from "../diff/streamDiff";
 import { streamLines } from "../diff/util";
 import { getSystemMessageWithRules } from "../llm/rules/getSystemMessageWithRules";
 import { gptEditPrompt } from "../llm/templates/edit";
-import { StreamAbortManager } from "../util/abortManager";
 import { findLast } from "../util/findLast";
 import { Telemetry } from "../util/posthog";
 import { recursiveStream } from "./recursiveStream";
@@ -64,7 +63,7 @@ export async function* streamDiffLines({
   highlighted,
   suffix,
   llm,
-  abortControllerId,
+  abortController,
   input,
   language,
   onlyOneInsertion,
@@ -75,15 +74,13 @@ export async function* streamDiffLines({
   highlighted: string;
   suffix: string;
   llm: ILLM;
-  abortControllerId: string;
+  abortController: AbortController;
   input: string;
   language: string | undefined;
   onlyOneInsertion: boolean;
   overridePrompt: ChatMessage[] | undefined;
   rulesToInclude: RuleWithSource[] | undefined;
 }): AsyncGenerator<DiffLine> {
-  const abortManager = StreamAbortManager.getInstance();
-  const abortController = abortManager.get(abortControllerId);
   void Telemetry.capture(
     "inlineEdit",
     {
diff --git a/core/index.d.ts b/core/index.d.ts
@@ -1263,7 +1263,7 @@ export interface HighlightedCodePayload {
 }
 
 export interface AcceptOrRejectDiffPayload {
-  filepath: string;
+  filepath?: string;
   streamId?: string;
 }
 
diff --git a/core/llm/utils/extractPathsFromCodeBlocks.ts b/core/llm/utils/extractPathsFromCodeBlocks.ts
@@ -2,8 +2,6 @@
  * Extracts file paths from markdown code blocks
  */
 export function extractPathsFromCodeBlocks(content: string): string[] {
-  console.log("CONTENT", content);
-
   const paths: string[] = [];
 
   // Match code block opening patterns:
diff --git a/extensions/intellij/src/main/kotlin/com/github/continuedev/continueintellijextension/types.kt b/extensions/intellij/src/main/kotlin/com/github/continuedev/continueintellijextension/types.kt
@@ -249,7 +249,7 @@ data class StreamDiffLinesPayload(
 )
 
 data class AcceptOrRejectDiffPayload(
-    val filepath: String,
+    val filepath: String? = null,
     val streamId: String? = null
 )
 
diff --git a/extensions/vscode/package-lock.json b/extensions/vscode/package-lock.json
diff --git a/extensions/vscode/src/apply/ApplyManager.ts b/extensions/vscode/src/apply/ApplyManager.ts
@@ -3,7 +3,7 @@ import { applyCodeBlock } from "core/edit/lazy/applyCodeBlock";
 import { getUriPathBasename } from "core/util/uri";
 import * as vscode from "vscode";
 
-import { StreamAbortManager } from "core/util/abortManager";
+import { ApplyAbortManager } from "core/edit/applyAbortManager";
 import { VerticalDiffManager } from "../diff/vertical/manager";
 import { VsCodeIde } from "../VsCodeIde";
 import { VsCodeWebviewProtocol } from "../webviewProtocol";
@@ -118,7 +118,7 @@ export class ApplyManager {
     }
 
     const fileUri = editor.document.uri.toString();
-    const abortManager = StreamAbortManager.getInstance();
+    const abortManager = ApplyAbortManager.getInstance();
     const abortController = abortManager.get(fileUri);
 
     const { isInstantApply, diffLinesGenerator } = await applyCodeBlock(
diff --git a/extensions/vscode/src/commands.ts b/extensions/vscode/src/commands.ts
@@ -35,6 +35,7 @@ import {
 } from "./autocomplete/statusBar";
 import { ContinueConsoleWebviewViewProvider } from "./ContinueConsoleWebviewViewProvider";
 import { ContinueGUIWebviewViewProvider } from "./ContinueGUIWebviewViewProvider";
+import { processDiff } from "./diff/processDiff";
 import { VerticalDiffManager } from "./diff/vertical/manager";
 import EditDecorationManager from "./quickEdit/EditDecorationManager";
 import { QuickEdit, QuickEditShowParams } from "./quickEdit/QuickEditQuickPick";
@@ -93,56 +94,6 @@ function hideGUI() {
   }
 }
 
-async function processDiff(
-  action: "accept" | "reject",
-  sidebar: ContinueGUIWebviewViewProvider,
-  ide: VsCodeIde,
-  core: Core,
-  verticalDiffManager: VerticalDiffManager,
-  newFileUri?: string,
-  streamId?: string,
-  toolCallId?: string,
-) {
-  captureCommandTelemetry(`${action}Diff`);
-
-  const currentFile = await ide.getCurrentFile();
-
-  let newOrCurrentUri = newFileUri;
-  if (!newOrCurrentUri) {
-    newOrCurrentUri = currentFile?.path;
-  }
-  if (!newOrCurrentUri) {
-    console.warn(
-      `No file provided or current file open while attempting to resolve diff`,
-    );
-    return;
-  }
-
-  await ide.openFile(newOrCurrentUri);
-
-  // Clear vertical diffs depending on action
-  verticalDiffManager.clearForfileUri(newOrCurrentUri, action === "accept");
-  if (action === "reject") {
-    core.invoke("cancelApply", undefined);
-  }
-
-  if (streamId) {
-    const fileContent = await ide.readFile(newOrCurrentUri);
-
-    await sidebar.webviewProtocol.request("updateApplyState", {
-      fileContent,
-      filepath: newOrCurrentUri,
-      streamId,
-      status: "closed",
-      numDiffs: 0,
-      toolCallId,
-    });
-  }
-
-  // Save the file
-  await ide.saveFile(newOrCurrentUri);
-}
-
 function waitForSidebarReady(
   sidebar: ContinueGUIWebviewViewProvider,
   timeout: number,
@@ -235,27 +186,31 @@ const getCommandsMap: (
   }
 
   return {
-    "continue.acceptDiff": async (newFileUri?: string, streamId?: string) =>
-      processDiff(
+    "continue.acceptDiff": async (newFileUri?: string, streamId?: string) => {
+      captureCommandTelemetry("acceptDiff");
+      void processDiff(
         "accept",
         sidebar,
         ide,
         core,
         verticalDiffManager,
         newFileUri,
         streamId,
-      ),
+      );
+    },
 
-    "continue.rejectDiff": async (newFileUri?: string, streamId?: string) =>
-      processDiff(
+    "continue.rejectDiff": async (newFileUri?: string, streamId?: string) => {
+      captureCommandTelemetry("rejectDiff");
+      void processDiff(
         "reject",
         sidebar,
         ide,
         core,
         verticalDiffManager,
         newFileUri,
         streamId,
-      ),
+      );
+    },
     "continue.acceptVerticalDiffBlock": (fileUri?: string, index?: number) => {
       captureCommandTelemetry("acceptVerticalDiffBlock");
       verticalDiffManager.acceptRejectVerticalDiffBlock(true, fileUri, index);
diff --git a/extensions/vscode/src/diff/processDiff.ts b/extensions/vscode/src/diff/processDiff.ts
@@ -0,0 +1,52 @@
+import { Core } from "core/core";
+import { ContinueGUIWebviewViewProvider } from "../ContinueGUIWebviewViewProvider";
+import { VsCodeIde } from "../VsCodeIde";
+import { VerticalDiffManager } from "./vertical/manager";
+
+export async function processDiff(
+  action: "accept" | "reject",
+  sidebar: ContinueGUIWebviewViewProvider,
+  ide: VsCodeIde,
+  core: Core,
+  verticalDiffManager: VerticalDiffManager,
+  newFileUri?: string,
+  streamId?: string,
+  toolCallId?: string,
+) {
+  let newOrCurrentUri = newFileUri;
+  if (!newOrCurrentUri) {
+    const currentFile = await ide.getCurrentFile();
+    newOrCurrentUri = currentFile?.path;
+  }
+  if (!newOrCurrentUri) {
+    console.warn(
+      `No file provided or current file open while attempting to resolve diff`,
+    );
+    return;
+  }
+
+  await ide.openFile(newOrCurrentUri);
+
+  // Clear vertical diffs depending on action
+  verticalDiffManager.clearForfileUri(newOrCurrentUri, action === "accept");
+  if (action === "reject") {
+    // this is so that IDE reject diff command can also cancel apply
+    core.invoke("cancelApply", undefined);
+  }
+
+  if (streamId) {
+    const fileContent = await ide.readFile(newOrCurrentUri);
+
+    await sidebar.webviewProtocol.request("updateApplyState", {
+      fileContent,
+      filepath: newOrCurrentUri,
+      streamId,
+      status: "closed",
+      numDiffs: 0,
+      toolCallId,
+    });
+  }
+
+  // Save the file
+  await ide.saveFile(newOrCurrentUri);
+}
diff --git a/extensions/vscode/src/diff/vertical/manager.ts b/extensions/vscode/src/diff/vertical/manager.ts
@@ -10,6 +10,7 @@ import EditDecorationManager from "../../quickEdit/EditDecorationManager";
 import { handleLLMError } from "../../util/errorHandling";
 import { VsCodeWebviewProtocol } from "../../webviewProtocol";
 
+import { ApplyAbortManager } from "core/edit/applyAbortManager";
 import { VerticalDiffHandler, VerticalDiffHandlerOptions } from "./handler";
 
 export interface VerticalDiffCodeLens {
@@ -437,6 +438,9 @@ export class VerticalDiffManager {
 
     this.editDecorationManager.clear();
 
+    const abortManager = ApplyAbortManager.getInstance();
+    const abortController = abortManager.get(fileUri);
+
     try {
       const streamedLines: string[] = [];
 
@@ -451,7 +455,7 @@ export class VerticalDiffManager {
           language: getMarkdownLanguageTagForFile(fileUri),
           onlyOneInsertion: !!onlyOneInsertion,
           overridePrompt,
-          abortControllerId: fileUri,
+          abortController,
         });
 
         for await (const line of stream) {
@@ -467,6 +471,10 @@ export class VerticalDiffManager {
       // enable a listener for user edits to file while diff is open
       this.enableDocumentChangeListener();
 
+      if (abortController.signal.aborted) {
+        vscode.commands.executeCommand("continue.rejectDiff");
+      }
+
       return `${prefix}${streamedLines.join("\n")}${suffix}`;
     } catch (e) {
       this.disableDocumentChangeListener();
diff --git a/gui/src/components/mainInput/Lump/LumpToolbar/IsApplyingToolbar.tsx b/gui/src/components/mainInput/Lump/LumpToolbar/IsApplyingToolbar.tsx
@@ -1,12 +1,14 @@
 import { useContext } from "react";
 import { IdeMessengerContext } from "../../../../context/IdeMessenger";
+import { useAppDispatch } from "../../../../redux/hooks";
+import { cancelStream } from "../../../../redux/thunks";
 import { getAltKeyLabel, getMetaKeyLabel } from "../../../../util";
 import { Container, StopButton } from "./components";
 import { GeneratingIndicator } from "./GeneratingIndicator";
 
 export const IsApplyingToolbar = () => {
   const ideMessenger = useContext(IdeMessengerContext);
-
+  const dispatch = useAppDispatch();
   const jetbrains = window.location.protocol === "jb-api:";
 
   return (
@@ -18,7 +20,9 @@ export const IsApplyingToolbar = () => {
           // Note that this will NOT stop generation but once apply is cancelled will show the Generating/cancel option
           // Apply is prioritized because it can be more catastrophic
           // Intentional to be more WYSIWYG for now
-          ideMessenger.post("cancelApply", undefined);
+          // Keyboard shortcut is handled in chat
+          dispatch(cancelStream());
+          ideMessenger.post("rejectDiff", {});
         }}
       >
         {/* JetBrains overrides cmd+backspace, so we have to use another shortcut */}
diff --git a/gui/src/pages/gui/Chat.tsx b/gui/src/pages/gui/Chat.tsx
diff --git a/gui/src/redux/thunks/streamResponseAfterToolCall.ts b/gui/src/redux/thunks/streamResponseAfterToolCall.ts

Original file line number	Diff line number	Diff line change
`@@ -1263,7 +1263,7 @@ export interface HighlightedCodePayload {`
`1263`	`1263`	`}`
`1264`	`1264`
`1265`	`1265`	`export interface AcceptOrRejectDiffPayload {`
`1266`		`- filepath: string;`
	`1266`	`+ filepath?: string;`
`1267`	`1267`	`streamId?: string;`
`1268`	`1268`	`}`
`1269`	`1269`
Original file line number	Diff line number	Diff line change
`@@ -249,7 +249,7 @@ data class StreamDiffLinesPayload(`
`249`	`249`	`)`
`250`	`250`
`251`	`251`	`data class AcceptOrRejectDiffPayload(`
`252`		`- val filepath: String,`
	`252`	`+ val filepath: String? = null,`
`253`	`253`	`val streamId: String? = null`
`254`	`254`	`)`
`255`	`255`