Merge branch 'main' into dallin/chat-symbols

.gitignore

@@ -142,7 +142,7 @@ codeql
 .tiktoken_cache
 
 
-# IntelliJ Plugin 
+# IntelliJ Plugin
 **/**/.gradle
 **/**/.qodana
 **/**/build
@@ -168,4 +168,6 @@ extensions/intellij/.idea/**
 **/.idea/shelf/
 
 extensions/intellij/bin
-extensions/.continue-debug/
+extensions/.continue-debug/
+
+*.vsix

core/autocomplete/CompletionProvider.ts

@@ -29,6 +29,7 @@ const autocompleteCache = AutocompleteLruCache.get();
 const ERRORS_TO_IGNORE = [
   // From Ollama
   "unexpected server status",
+  "operation was aborted"
 ];
 
 export type GetLspDefinitionsFunction = (

core/autocomplete/generation/CompletionStreamer.ts

@@ -25,10 +25,10 @@ export class CompletionStreamer {
     // Try to reuse pending requests if what the user typed matches start of completion
     const generator = this.generatorReuseManager.getGenerator(
       prefix,
-      () =>
+      (abortSignal: AbortSignal) =>
         llm.supportsFim()
-          ? llm.streamFim(prefix, suffix, completionOptions)
-          : llm.streamComplete(prompt, {
+          ? llm.streamFim(prefix, suffix, abortSignal, completionOptions)
+          : llm.streamComplete(prompt, abortSignal, {
               ...completionOptions,
               raw: true,
             }),

core/autocomplete/generation/GeneratorReuseManager.test.ts

@@ -4,7 +4,7 @@ import { GeneratorReuseManager } from "./GeneratorReuseManager";
 function createMockGenerator(
   data: string[],
   delay: number = 0,
-): () => AsyncGenerator<string> {
+): (abortSignal: AbortSignal) => AsyncGenerator<string> {
   const mockGenerator = async function* () {
     for (const chunk of data) {
       yield chunk;

core/autocomplete/generation/GeneratorReuseManager.ts

@@ -8,12 +8,13 @@ export class GeneratorReuseManager {
   constructor(private readonly onError: (err: any) => void) {}
 
   private _createListenableGenerator(
+    abortController: AbortController,
     gen: AsyncGenerator<string>,
     prefix: string,
   ) {
     this.currentGenerator?.cancel();
 
-    const listenableGen = new ListenableGenerator(gen, this.onError);
+    const listenableGen = new ListenableGenerator(gen, this.onError, abortController);
     listenableGen.listen((chunk) => (this.pendingCompletion += chunk ?? ""));
 
     this.pendingGeneratorPrefix = prefix;
@@ -35,13 +36,14 @@ export class GeneratorReuseManager {
 
   async *getGenerator(
     prefix: string,
-    newGenerator: () => AsyncGenerator<string>,
+    newGenerator: (abortSignal: AbortSignal) => AsyncGenerator<string>,
     multiline: boolean,
   ): AsyncGenerator<string> {
     // If we can't reuse, then create a new generator
     if (!this.shouldReuseExistingGenerator(prefix)) {
       // Create a wrapper over the current generator to fix the prompt
-      this._createListenableGenerator(newGenerator(), prefix);
+      const abortController = new AbortController();
+      this._createListenableGenerator(abortController, newGenerator(abortController.signal), prefix);
     }
 
     // Already typed characters are those that are new in the prefix from the old generator

core/autocomplete/generation/ListenableGenerator.ts

@@ -3,16 +3,20 @@ export class ListenableGenerator<T> {
   private _buffer: T[] = [];
   private _listeners: Set<(value: T) => void> = new Set();
   private _isEnded = false;
+  private _abortController: AbortController
 
   constructor(
     source: AsyncGenerator<T>,
     private readonly onError: (e: any) => void,
+    abortController: AbortController
   ) {
     this._source = source;
+    this._abortController = abortController;
     this._start();
   }
 
   public cancel() {
+    this._abortController.abort();
     this._isEnded = true;
   }
 

core/commands/index.ts

@@ -58,7 +58,7 @@ export function slashFromCustomCommand(
         }
       }
 
-      for await (const chunk of llm.streamChat(messages)) {
+      for await (const chunk of llm.streamChat(messages, new AbortController().signal)) {
         yield stripImages(chunk.content);
       }
     },

core/commands/slash/cmd.ts

@@ -24,7 +24,9 @@ const GenerateTerminalCommand: SlashCommand = {
 
 "${input}"
 
-Please write a shell command that will do what the user requested. Your output should consist of only the command itself, without any explanation or example output. Do not use any newlines. Only output the command that when inserted into the terminal will do precisely what was requested. Here is the command:`);
+Please write a shell command that will do what the user requested. Your output should consist of only the command itself, without any explanation or example output. Do not use any newlines. Only output the command that when inserted into the terminal will do precisely what was requested. Here is the command:`,
+        new AbortController().signal
+      );
 
     const lines = streamLines(gen);
     let cmd = "";

core/commands/slash/commit.ts

@@ -16,7 +16,7 @@ const CommitMessageCommand: SlashCommand = {
     const prompt = `${diff}\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)) {
       yield stripImages(chunk.content);
     }
   },

core/commands/slash/draftIssue.ts

@@ -30,7 +30,7 @@ const DraftIssueCommand: SlashCommand = {
       return;
     }
     let title = await llm.complete(
-      `Generate a title for the GitHub issue requested in this user input: '${input}'. Use no more than 20 words and output nothing other than the title. Do not surround it with quotes. The title is: `,
+      `Generate a title for the GitHub issue requested in this user input: '${input}'. Use no more than 20 words and output nothing other than the title. Do not surround it with quotes. The title is: `, new AbortController().signal,
       { maxTokens: 20 },
     );
 
@@ -43,7 +43,7 @@ const DraftIssueCommand: SlashCommand = {
       { role: "user", content: PROMPT(input, title) },
     ];
 
-    for await (const chunk of llm.streamChat(messages)) {
+    for await (const chunk of llm.streamChat(messages, new AbortController().signal)) {
       body += chunk.content;
       yield stripImages(chunk.content);
     }

core/commands/slash/edit.ts

@@ -478,7 +478,7 @@ const EditSlashCommand: SlashCommand = {
         messages = rendered;
       }
 
-      const completion = llm.streamComplete(rendered as string, {
+      const completion = llm.streamComplete(rendered as string,  new AbortController().signal, {
         maxTokens: Math.min(maxTokens, Math.floor(llm.contextLength / 2), 4096),
         raw: true,
       });
@@ -494,7 +494,7 @@ const EditSlashCommand: SlashCommand = {
       );
     } else {
       async function* gen() {
-        for await (const chunk of llm.streamChat(messages, {
+        for await (const chunk of llm.streamChat(messages,  new AbortController().signal,{
           temperature: 0.5, // TODO
           maxTokens: Math.min(
             maxTokens,
@@ -609,7 +609,7 @@ ${lines.join("\n")}
 
 Please briefly explain the changes made to the code above. Give no more than 2-3 sentences, and use markdown bullet points:`;
 
-      for await (const update of llm.streamComplete(prompt)) {
+      for await (const update of llm.streamComplete(prompt, new AbortController().signal)) {
         yield update;
       }
     }

core/commands/slash/multifileEdit.ts

@@ -39,7 +39,7 @@ const MultiFileEditSlashCommand: SlashCommand = {
 
     const content = createPrompt(filesToEditStr, additionalContextStr, input);
 
-    for await (const chunk of llm.streamChat([{ role: "user", content }])) {
+    for await (const chunk of llm.streamChat([{ role: "user", content }], new AbortController().signal)) {
       yield stripImages(chunk.content);
     }
   },

core/commands/slash/onboard.ts

@@ -47,7 +47,7 @@ const OnboardSlashCommand: SlashCommand = {
 
     for await (const chunk of llm.streamChat([
       { role: "user", content: prompt },
-    ])) {
+    ], new AbortController().signal)) {
       yield stripImages(chunk.content);
     }
   },
@@ -134,7 +134,7 @@ function createOnboardingPrompt(context: string): string {
     Your response should be structured, clear, and focused on giving the new developer both a detailed understanding of individual components and a high-level overview of the project as a whole.
 
     Here is an example of a valid response:
-    
+
     ## Important folders
 
     ### /folder1

core/commands/slash/review.ts

@@ -45,7 +45,7 @@ const ReviewMessageCommand: SlashCommand = {
 
     for await (const chunk of llm.streamChat([
       { role: "user", content: content },
-    ])) {
+    ], new AbortController().signal)) {
       yield stripImages(chunk.content);
     }
   },

core/config/promptFile.ts

@@ -125,7 +125,7 @@ export function slashCommandFromPromptFile(
         systemMessage,
       );
 
-      for await (const chunk of context.llm.streamChat(messages)) {
+      for await (const chunk of context.llm.streamChat(messages, new AbortController().signal)) {
         yield stripImages(chunk.content);
       }
 

core/context/rerankers/llm.ts

@@ -6,7 +6,7 @@ const RERANK_PROMPT = (
   documentId: string,
   document: string,
 ) => `You are an expert software developer responsible for helping detect whether the retrieved snippet of code is relevant to the query. For a given input, you need to output a single word: "Yes" or "No" indicating the retrieved snippet is relevant to the query.
-  
+
   Query: Where is the FastAPI server?
   Snippet:
   \`\`\`/Users/andrew/Desktop/server/main.py
@@ -17,7 +17,7 @@ const RERANK_PROMPT = (
       return {{"Hello": "World"}}
   \`\`\`
   Relevant: Yes
-  
+
   Query: Where in the documentation does it talk about the UI?
   Snippet:
   \`\`\`/Users/andrew/Projects/bubble_sort/src/lib.rs
@@ -32,13 +32,13 @@ const RERANK_PROMPT = (
   }}
   \`\`\`
   Relevant: No
-  
+
   Query: ${query}
   Snippet:
   \`\`\`${documentId}
   ${document}
   \`\`\`
-  Relevant: 
+  Relevant:
   `;
 
 export class LLMReranker implements Reranker {
@@ -49,6 +49,7 @@ export class LLMReranker implements Reranker {
   async scoreChunk(chunk: Chunk, query: string): Promise<number> {
     const completion = await this.llm.complete(
       RERANK_PROMPT(query, getBasename(chunk.filepath), chunk.content),
+      new AbortController().signal,
       {
         maxTokens: 1,
         model:

core/context/retrieval/repoMapRequest.ts

@@ -61,7 +61,7 @@ This is the question that you should select relevant files for: "${input}"`;
     const response = await llm.chat([
       { role: "user", content: prompt },
       { role: "assistant", content: "<reasoning>" },
-    ]);
+    ], new AbortController().signal);
     const content = stripImages(response.content);
     console.debug("Repo map retrieval response: ", content);
 

core/core.ts

@@ -393,6 +393,7 @@ export class Core {
       const model = await configHandler.llmFromTitle(msg.data.title);
       const gen = model.streamChat(
         msg.data.messages,
+        new AbortController().signal,
         msg.data.completionOptions,
       );
       let next = await gen.next();
@@ -434,6 +435,7 @@ export class Core {
       const model = await configHandler.llmFromTitle(msg.data.title);
       const gen = model.streamComplete(
         msg.data.prompt,
+        new AbortController().signal,
         msg.data.completionOptions,
       );
       let next = await gen.next();
@@ -466,6 +468,7 @@ export class Core {
       const model = await this.configHandler.llmFromTitle(msg.data.title);
       const completion = await model.complete(
         msg.data.prompt,
+        new AbortController().signal,
         msg.data.completionOptions,
       );
       return completion;

core/edit/lazy/replace.ts

@@ -86,7 +86,7 @@ export async function* getReplacementWithLlm(
   const completion = await llm.streamChat([
     { role: "user", content: userPrompt },
     { role: "assistant", content: assistantPrompt },
-  ]);
+  ], new AbortController().signal);
 
   let lines = streamLines(completion);
   lines = filterLeadingNewline(lines);

core/edit/lazy/streamLazyApply.ts

@@ -23,7 +23,7 @@ export async function* streamLazyApply(
   }
 
   const promptMessages = promptFactory(oldCode, filename, newCode);
-  const lazyCompletion = llm.streamChat(promptMessages);
+  const lazyCompletion = llm.streamChat(promptMessages, new AbortController().signal);
 
   // Do find and replace over the lazy edit response
   async function* replacementFunction(

core/edit/streamDiffLines.ts

@@ -96,8 +96,8 @@ export async function* streamDiffLines(
 
   const completion =
     typeof prompt === "string"
-      ? llm.streamComplete(prompt, { raw: true, prediction })
-      : llm.streamChat(prompt, {
+      ? llm.streamComplete(prompt,  new AbortController().signal, { raw: true, prediction })
+      : llm.streamChat(prompt,  new AbortController().signal, {
           prediction,
         });
 

core/index.d.ts

@@ -79,26 +79,34 @@ export interface ILLM extends LLMOptions {
   region?: string;
   projectId?: string;
 
-  complete(prompt: string, options?: LLMFullCompletionOptions): Promise<string>;
+  complete(
+    prompt: string,
+    signal: AbortSignal,
+    options?: LLMFullCompletionOptions
+  ): Promise<string>;
 
   streamComplete(
     prompt: string,
+    signal: AbortSignal,
     options?: LLMFullCompletionOptions,
   ): AsyncGenerator<string, PromptLog>;
 
   streamFim(
     prefix: string,
     suffix: string,
+    signal: AbortSignal,
     options?: LLMFullCompletionOptions,
   ): AsyncGenerator<string, PromptLog>;
 
   streamChat(
     messages: ChatMessage[],
+    signal: AbortSignal,
     options?: LLMFullCompletionOptions,
   ): AsyncGenerator<ChatMessage, PromptLog>;
 
   chat(
     messages: ChatMessage[],
+    signal: AbortSignal,
     options?: LLMFullCompletionOptions,
   ): Promise<ChatMessage>;
 
@@ -395,11 +403,13 @@ export interface CustomLLMWithOptionals {
   options: LLMOptions;
   streamCompletion?: (
     prompt: string,
+    signal: AbortSignal,
     options: CompletionOptions,
     fetch: (input: RequestInfo | URL, init?: RequestInit) => Promise<Response>,
   ) => AsyncGenerator<string>;
   streamChat?: (
     messages: ChatMessage[],
+    signal: AbortSignal,
     options: CompletionOptions,
     fetch: (input: RequestInfo | URL, init?: RequestInit) => Promise<Response>,
   ) => AsyncGenerator<string>;