fix apply cancellation and make slash commands cancelable

Author: RomneyDa

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/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:
@@ -83,10 +84,13 @@ export async function* getReplacementWithLlm(
     \`\`\`
   `;
 
-  const completion = await llm.streamChat([
-    { role: "user", content: userPrompt },
-    { role: "assistant", content: assistantPrompt },
-  ], new AbortController().signal);
+  const completion = await llm.streamChat(
+    [
+      { role: "user", content: userPrompt },
+      { role: "assistant", content: assistantPrompt },
+    ],
+    abortController.signal,
+  );
 
   let lines = streamLines(completion);
   lines = filterLeadingNewline(lines);

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/index.d.ts

@@ -836,6 +836,7 @@ export interface ContinueSDK {
   config: ContinueConfig;
   fetch: FetchFunction;
   completionOptions?: LLMFullCompletionOptions;
+  abortController: AbortController;
 }
 
 export interface SlashCommand {
@@ -1141,7 +1142,7 @@ export interface StreamableHTTPOptions {
   requestOptions?: RequestOptions;
 }
 
-export type TransportOptions = 
+export type TransportOptions =
   | StdioOptions
   | WebSocketOptions
   | SSEOptions