Merge pull request #5907 from continuedev/dallin/cancelable-apply-2

Author: sestinj

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 renamed to 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);