Fix apply cancellation and make slash commands cancelable

core/commands/index.ts

@@ -25,6 +25,7 @@ export function slashFromCustomCommand(
       config,
       selectedCode,
       fetch,
+      abortController,
     }) {
       // Render prompt template
       let renderedPrompt: string;
@@ -63,7 +64,7 @@ export function slashFromCustomCommand(
 
       for await (const chunk of llm.streamChat(
         messages,
-        new AbortController().signal,
+        abortController.signal,
         completionOptions,
       )) {
         yield renderChatMessage(chunk);

core/commands/slash/commit.ts

@@ -4,7 +4,7 @@ import { renderChatMessage } from "../../util/messageContent.js";
 const CommitMessageCommand: SlashCommand = {
   name: "commit",
   description: "Generate a commit message for current changes",
-  run: async function* ({ ide, llm, params }) {
+  run: async function* ({ ide, llm, params, abortController }) {
     const includeUnstaged = params?.includeUnstaged ?? false;
     const diff = await ide.getDiff(includeUnstaged);
 
@@ -16,7 +16,7 @@ const CommitMessageCommand: SlashCommand = {
     const prompt = `${diff.join("\n")}\n\nGenerate a commit message for the above set of changes. First, give a single sentence, no more than 80 characters. Then, after 2 line breaks, give a list of no more than 5 short bullet points, each no more than 40 characters. Output nothing except for the commit message, and don't surround it in quotes.`;
     for await (const chunk of llm.streamChat(
       [{ role: "user", content: prompt }],
-      new AbortController().signal,
+      abortController.signal,
     )) {
       yield renderChatMessage(chunk);
     }

core/commands/slash/draftIssue.ts

@@ -24,7 +24,7 @@ Body:\n\n`;
 const DraftIssueCommand: SlashCommand = {
   name: "issue",
   description: "Draft a GitHub issue",
-  run: async function* ({ input, llm, history, params }) {
+  run: async function* ({ input, llm, history, params, abortController }) {
     if (params?.repositoryUrl === undefined) {
       yield "This command requires a repository URL to be set in the config file.";
       return;
@@ -46,7 +46,7 @@ const DraftIssueCommand: SlashCommand = {
 
     for await (const chunk of llm.streamChat(
       messages,
-      new AbortController().signal,
+      abortController.signal,
     )) {
       body += chunk.content;
       yield renderChatMessage(chunk);

core/commands/slash/mcp.ts

@@ -60,7 +60,7 @@ export function constructMcpSlashCommand(
 
       for await (const chunk of context.llm.streamChat(
         messages,
-        new AbortController().signal,
+        context.abortController.signal,
         context.completionOptions,
       )) {
         yield renderChatMessage(chunk);

core/commands/slash/onboard.ts

@@ -41,15 +41,15 @@ const MAX_EXPLORE_DEPTH = 2;
 const OnboardSlashCommand: SlashCommand = {
   name: "onboard",
   description: "Familiarize yourself with the codebase",
-  run: async function* ({ llm, ide }) {
+  run: async function* ({ llm, ide, abortController }) {
     const [workspaceDir] = await ide.getWorkspaceDirs();
 
     const context = await gatherProjectContext(workspaceDir, ide);
     const prompt = createOnboardingPrompt(context);
 
     for await (const chunk of llm.streamChat(
       [{ role: "user", content: prompt }],
-      new AbortController().signal,
+      abortController.signal,
     )) {
       yield renderChatMessage(chunk);
     }

core/commands/slash/review.ts

@@ -38,14 +38,14 @@ function getLastUserHistory(history: ChatMessage[]): string {
 const ReviewMessageCommand: SlashCommand = {
   name: "review",
   description: "Review code and give feedback",
-  run: async function* ({ llm, history }) {
+  run: async function* ({ llm, history, abortController }) {
     const reviewText = getLastUserHistory(history).replace("\\review", "");
 
     const content = `${prompt} \r\n ${reviewText}`;
 
     for await (const chunk of llm.streamChat(
       [{ role: "user", content: content }],
-      new AbortController().signal,
+      abortController.signal,
     )) {
       yield renderChatMessage(chunk);
     }

core/core.ts

@@ -58,6 +58,7 @@ import {
 import { createNewWorkspaceBlockFile } from "./config/workspace/workspaceBlocks";
 import { MCPManagerSingleton } from "./context/mcp/MCPManagerSingleton";
 import { setMdmLicenseKey } from "./control-plane/mdm/mdm";
+import { ApplyAbortManager } from "./edit/applyAbortManager";
 import { streamDiffLines } from "./edit/streamDiffLines";
 import { shouldIgnore } from "./indexing/shouldIgnore";
 import { walkDirCache } from "./indexing/walkDir";
@@ -67,7 +68,6 @@ import { llmStreamChat } from "./llm/streamChat";
 import type { FromCoreProtocol, ToCoreProtocol } from "./protocol";
 import { OnboardingModes } from "./protocol/core";
 import type { IMessenger, Message } from "./protocol/messenger";
-import { StreamAbortManager } from "./util/abortManager";
 import { getUriPathBasename } from "./util/uri";
 
 const hasRulesFiles = (uris: string[]): boolean => {
@@ -512,6 +512,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,
@@ -525,13 +530,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("onboarding/complete", this.handleCompleteOnboarding.bind(this));

core/util/abortManager.ts → 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 {

core/edit/lazy/applyCodeBlock.ts

@@ -16,6 +16,7 @@ export async function applyCodeBlock(
   newLazyFile: string,
   filename: string,
   llm: ILLM,
+  abortController: AbortController,
 ): Promise<{
   isInstantApply: boolean;
   diffLinesGenerator: AsyncGenerator<DiffLine>;
@@ -51,6 +52,12 @@ export async function applyCodeBlock(
 
   return {
     isInstantApply: false,
-    diffLinesGenerator: streamLazyApply(oldFile, filename, newLazyFile, llm),
+    diffLinesGenerator: streamLazyApply(
+      oldFile,
+      filename,
+      newLazyFile,
+      llm,
+      abortController,
+    ),
   };
 }

core/edit/lazy/replace.ts

@@ -59,6 +59,7 @@ export async function* getReplacementWithLlm(
   linesBefore: string[],
   linesAfter: string[],
   llm: ILLM,
+  abortController: AbortController,
 ): AsyncGenerator<string> {
   const userPrompt = dedent`
     ORIGINAL CODE:
@@ -88,7 +89,7 @@ export async function* getReplacementWithLlm(
       { role: "user", content: userPrompt },
       { role: "assistant", content: assistantPrompt },
     ],
-    new AbortController().signal,
+    abortController.signal,
   );
 
   let lines = streamLines(completion);

core/edit/lazy/streamLazyApply.ts

@@ -16,17 +16,15 @@ export async function* streamLazyApply(
   filename: string,
   newCode: string,
   llm: ILLM,
+  abortController: AbortController,
 ): AsyncGenerator<DiffLine> {
   const promptFactory = lazyApplyPromptForModel(llm.model, llm.providerName);
   if (!promptFactory) {
     throw new Error(`Lazy apply not supported for model ${llm.model}`);
   }
 
   const promptMessages = promptFactory(oldCode, filename, newCode);
-  const lazyCompletion = llm.streamChat(
-    promptMessages,
-    new AbortController().signal,
-  );
+  const lazyCompletion = llm.streamChat(promptMessages, abortController.signal);
 
   // Do find and replace over the lazy edit response
   async function* replacementFunction(
@@ -39,6 +37,7 @@ export async function* streamLazyApply(
       linesBefore,
       linesAfter,
       llm,
+      abortController,
     )) {
       yield line;
     }

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",
     {

core/index.d.ts

@@ -833,6 +833,7 @@ export interface ContinueSDK {
   config: ContinueConfig;
   fetch: FetchFunction;
   completionOptions?: LLMFullCompletionOptions;
+  abortController: AbortController;
 }
 
 export interface SlashCommand {
@@ -1280,7 +1281,7 @@ export interface HighlightedCodePayload {
 }
 
 export interface AcceptOrRejectDiffPayload {
-  filepath: string;
+  filepath?: string;
   streamId?: string;
 }
 

core/llm/streamChat.ts

@@ -85,6 +85,7 @@ export async function* llmStreamChat(
           config.requestOptions,
         ),
       completionOptions,
+      abortController,
     });
     let next = await gen.next();
     while (!next.done) {

core/promptFiles/v1/slashCommandFromPromptFile.ts

@@ -79,7 +79,7 @@ export function slashCommandFromPromptFileV1(
 
       for await (const chunk of context.llm.streamChat(
         messages,
-        new AbortController().signal,
+        context.abortController.signal,
         context.completionOptions,
       )) {
         yield renderChatMessage(chunk);

extensions/intellij/src/main/kotlin/com/github/continuedev/continueintellijextension/types.kt

@@ -239,7 +239,7 @@ data class StreamDiffLinesPayload(
 )
 
 data class AcceptOrRejectDiffPayload(
-    val filepath: String,
+    val filepath: String? = null,
     val streamId: String? = null
 )
 

extensions/vscode/src/apply/ApplyManager.ts

@@ -3,6 +3,7 @@ import { applyCodeBlock } from "core/edit/lazy/applyCodeBlock";
 import { getUriPathBasename } from "core/util/uri";
 import * as vscode from "vscode";
 
+import { ApplyAbortManager } from "core/edit/applyAbortManager";
 import { VerticalDiffManager } from "../diff/vertical/manager";
 import { VsCodeIde } from "../VsCodeIde";
 import { VsCodeWebviewProtocol } from "../webviewProtocol";
@@ -116,11 +117,16 @@ export class ApplyManager {
       return;
     }
 
+    const fileUri = editor.document.uri.toString();
+    const abortManager = ApplyAbortManager.getInstance();
+    const abortController = abortManager.get(fileUri);
+
     const { isInstantApply, diffLinesGenerator } = await applyCodeBlock(
       editor.document.getText(),
       text,
-      getUriPathBasename(editor.document.uri.toString()),
+      getUriPathBasename(fileUri),
       llm,
+      abortController,
     );
 
     if (isInstantApply) {