Merge branch 'main' into dallin/limit-glob-output

Author: Patrick-Erichsen

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/config/types.ts

@@ -719,7 +719,7 @@ declare global {
   
     getPinnedFiles(): Promise<string[]>;
   
-    getSearchResults(query: string): Promise<string>;
+    getSearchResults(query: string, maxResults?: number): Promise<string>;
   
     subprocess(command: string, cwd?: string): Promise<[string, string]>;
   

core/context/providers/SearchContextProvider.ts

@@ -6,6 +6,7 @@ import {
 import { formatGrepSearchResults } from "../../util/grepSearch.js";
 import { BaseContextProvider } from "../index.js";
 
+const DEFAULT_MAX_SEARCH_CONTEXT_RESULTS = 200;
 class SearchContextProvider extends BaseContextProvider {
   static description: ContextProviderDescription = {
     title: "search",
@@ -19,8 +20,12 @@ class SearchContextProvider extends BaseContextProvider {
     query: string,
     extras: ContextProviderExtras,
   ): Promise<ContextItem[]> {
-    const results = await extras.ide.getSearchResults(query);
-    const formatted = formatGrepSearchResults(results);
+    const results = await extras.ide.getSearchResults(
+      query,
+      this.options?.maxResults ?? DEFAULT_MAX_SEARCH_CONTEXT_RESULTS,
+    );
+    // Note, search context provider will not truncate result chars, but will limit number of results
+    const { formatted } = formatGrepSearchResults(results);
     return [
       {
         description: "Search results",

core/core.ts

@@ -42,12 +42,13 @@ import {
   IdeSettings,
   ModelDescription,
   RangeInFile,
+  ToolCall,
   type ContextItem,
   type ContextItemId,
   type IDE,
 } from ".";
 
-import { ConfigYaml } from "@continuedev/config-yaml";
+import { BLOCK_TYPES, ConfigYaml } from "@continuedev/config-yaml";
 import { getDiffFn, GitDiffCache } from "./autocomplete/snippets/gitDiffCache";
 import { isLocalDefinitionFile } from "./config/loadLocalAssistants";
 import {
@@ -58,6 +59,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 +69,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 => {
@@ -147,6 +148,18 @@ export class Core {
 
     MCPManagerSingleton.getInstance().onConnectionsRefreshed = () => {
       void this.configHandler.reloadConfig();
+
+      // Refresh @mention dropdown submenu items for MCP providers
+      const mcpManager = MCPManagerSingleton.getInstance();
+      const mcpProviderNames = Array.from(mcpManager.connections.keys()).map(
+        (mcpId) => `mcp-${mcpId}`,
+      );
+
+      if (mcpProviderNames.length > 0) {
+        this.messenger.send("refreshSubmenuItems", {
+          providers: mcpProviderNames,
+        });
+      }
     };
 
     this.codeBaseIndexer = new CodebaseIndexer(
@@ -512,6 +525,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 +543,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));
@@ -738,44 +756,9 @@ export class Core {
       return { url };
     });
 
-    on("tools/call", async ({ data: { toolCall } }) => {
-      const { config } = await this.configHandler.loadConfig();
-      if (!config) {
-        throw new Error("Config not loaded");
-      }
-
-      const tool = config.tools.find(
-        (t) => t.function.name === toolCall.function.name,
-      );
-
-      if (!tool) {
-        throw new Error(`Tool ${toolCall.function.name} not found`);
-      }
-
-      if (!config.selectedModelByRole.chat) {
-        throw new Error("No chat model selected");
-      }
-
-      // Define a callback for streaming output updates
-      const onPartialOutput = (params: {
-        toolCallId: string;
-        contextItems: ContextItem[];
-      }) => {
-        this.messenger.send("toolCallPartialOutput", params);
-      };
-
-      return await callTool(tool, toolCall, {
-        config,
-        ide: this.ide,
-        llm: config.selectedModelByRole.chat,
-        fetch: (url, init) =>
-          fetchwithRequestOptions(url, init, config.requestOptions),
-        tool,
-        toolCallId: toolCall.id,
-        onPartialOutput,
-        codeBaseIndexer: this.codeBaseIndexer,
-      });
-    });
+    on("tools/call", async ({ data: { toolCall } }) =>
+      this.handleToolCall(toolCall),
+    );
 
     on("isItemTooBig", async ({ data: { item } }) => {
       return this.isItemTooBig(item);
@@ -800,6 +783,45 @@ export class Core {
     });
   }
 
+  private async handleToolCall(toolCall: ToolCall) {
+    const { config } = await this.configHandler.loadConfig();
+    if (!config) {
+      throw new Error("Config not loaded");
+    }
+
+    const tool = config.tools.find(
+      (t) => t.function.name === toolCall.function.name,
+    );
+
+    if (!tool) {
+      throw new Error(`Tool ${toolCall.function.name} not found`);
+    }
+
+    if (!config.selectedModelByRole.chat) {
+      throw new Error("No chat model selected");
+    }
+
+    // Define a callback for streaming output updates
+    const onPartialOutput = (params: {
+      toolCallId: string;
+      contextItems: ContextItem[];
+    }) => {
+      this.messenger.send("toolCallPartialOutput", params);
+    };
+
+    return await callTool(tool, toolCall, {
+      config,
+      ide: this.ide,
+      llm: config.selectedModelByRole.chat,
+      fetch: (url, init) =>
+        fetchwithRequestOptions(url, init, config.requestOptions),
+      tool,
+      toolCallId: toolCall.id,
+      onPartialOutput,
+      codeBaseIndexer: this.codeBaseIndexer,
+    });
+  }
+
   private async isItemTooBig(item: ContextItemWithId) {
     const { config } = await this.configHandler.loadConfig();
     if (!config) {
@@ -887,7 +909,12 @@ export class Core {
           uri.endsWith(".prompt") ||
           uri.endsWith(SYSTEM_PROMPT_DOT_FILE) ||
           (uri.includes(".continue") && uri.endsWith(".yaml")) ||
-          uri.endsWith(RULES_MARKDOWN_FILENAME)
+          uri.endsWith(RULES_MARKDOWN_FILENAME) ||
+          BLOCK_TYPES.some(
+            (blockType) =>
+              uri.includes(`.continue/${blockType}`) ||
+              uri.includes(`.continue\\${blockType}`),
+          )
         ) {
           await this.configHandler.reloadConfig();
         } else if (

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