plumb repo selection to grep tool

Author: brendan-kellam

docs/docs/features/mcp-server.mdx

@@ -304,22 +304,19 @@ Pass the key as an `Authorization: Bearer <key>` header when connecting to the M
 
 ## Available Tools
 
-### `search_code`
+### `grep`
 
-Searches for code that matches the provided search query as a substring by default, or as a regular expression if `useRegex` is true.
+Searches for code matching a regular expression pattern across repositories, similar to `grep`/`ripgrep`. Always case-sensitive. Results are grouped by file and include line numbers.
 
 Parameters:
 | Name | Required | Description |
-|:----------------------|:---------|:---------------------------------------------------------------------------------------------------------------------|
-| `query` | yes | The search pattern to match against code contents. Do not escape quotes in your query. |
-| `useRegex` | no | Whether to use regular expression matching. When false, substring matching is used (default: false). |
-| `filterByRepos` | no | Scope the search to specific repositories. |
-| `filterByLanguages` | no | Scope the search to specific languages. |
-| `filterByFilepaths` | no | Scope the search to specific filepaths. |
-| `caseSensitive` | no | Whether the search should be case sensitive (default: false). |
-| `includeCodeSnippets` | no | Whether to include code snippets in the response (default: false). |
+|:----------|:---------|:--------------------------------------------------------------------------------------------------------------|
+| `pattern` | yes | The regex pattern to search for in file contents. |
+| `path` | no | Directory path to scope the search to. Defaults to the repository root. |
+| `include` | no | File glob pattern to include in the search (e.g. `*.ts`, `*.{ts,tsx}`). |
+| `repo` | no | Repository name to search in. If not provided, searches all repositories. Use the full name including host (e.g. `github.com/org/repo`). |
 | `ref` | no | Commit SHA, branch or tag name to search on. If not provided, defaults to the default branch. |
-| `maxTokens` | no | The maximum number of tokens to return (default: 10000). |
+| `limit` | no | Maximum number of matching files to return (default: 100). |
 
 ### `list_repos`
 
@@ -336,18 +333,20 @@ Parameters:
 
 ### `read_file`
 
-Reads the source code for a given file.
+Reads the source code for a given file, with optional line range control for large files.
 
 Parameters:
 | Name | Required | Description |
-|:-------|:---------|:-------------------------------------------------------------------------------------------------------|
+|:---------|:---------|:-------------------------------------------------------------------------------------------------------|
 | `repo` | yes | The repository name. |
 | `path` | yes | The path to the file. |
 | `ref` | no | Commit SHA, branch or tag name to fetch the source code for. If not provided, uses the default branch. |
+| `offset` | no | Line number to start reading from (1-indexed). Omit to start from the beginning. |
+| `limit` | no | Maximum number of lines to read (max: 500). Omit to read up to 500 lines. |
 
 ### `list_tree`
 
-Lists files and directories from a repository path. Can be used as a directory listing tool (`depth: 1`) or a repo-tree tool (`depth > 1`).
+Lists files and directories from a repository path. Directories are shown before files at each level.
 
 Parameters:
 | Name | Required | Description |

packages/web/src/features/chat/agent.ts

@@ -21,7 +21,7 @@ import { readFileDefinition } from "@/features/tools/readFile";
 import { grepDefinition } from "@/features/tools/grep";
 import { Source } from "./types";
 import { addLineNumbers, fileReferenceToString } from "./utils";
-import { tools } from "./tools";
+import { createTools } from "./tools";
 import { listTreeDefinition } from "../tools";
 
 const dedent = _dedent.withOptions({ alignValues: true });
@@ -203,7 +203,7 @@ const createAgentStream = async ({
         providerOptions,
         messages: inputMessages,
         system: systemPrompt,
-        tools,
+        tools: createTools({ source: 'sourcebot-ask-agent', selectedRepos }),
         temperature: env.SOURCEBOT_CHAT_MODEL_TEMPERATURE,
         stopWhen: [
             stepCountIsGTE(env.SOURCEBOT_CHAT_MAX_STEP_COUNT),

packages/web/src/features/chat/tools.ts

@@ -8,18 +8,19 @@ import {
     findSymbolDefinitionsDefinition,
     listTreeDefinition,
 } from "@/features/tools";
+import { ToolContext } from "@/features/tools/types";
 import { ToolUIPart } from "ai";
 import { SBChatMessageToolTypes } from "./types";
 
-export const tools = {
-    [readFileDefinition.name]: toVercelAITool(readFileDefinition),
-    [listCommitsDefinition.name]: toVercelAITool(listCommitsDefinition),
-    [listReposDefinition.name]: toVercelAITool(listReposDefinition),
-    [grepDefinition.name]: toVercelAITool(grepDefinition),
-    [findSymbolReferencesDefinition.name]: toVercelAITool(findSymbolReferencesDefinition),
-    [findSymbolDefinitionsDefinition.name]: toVercelAITool(findSymbolDefinitionsDefinition),
-    [listTreeDefinition.name]: toVercelAITool(listTreeDefinition),
-} as const;
+export const createTools = (context: ToolContext) => ({
+    [readFileDefinition.name]: toVercelAITool(readFileDefinition, context),
+    [listCommitsDefinition.name]: toVercelAITool(listCommitsDefinition, context),
+    [listReposDefinition.name]: toVercelAITool(listReposDefinition, context),
+    [grepDefinition.name]: toVercelAITool(grepDefinition, context),
+    [findSymbolReferencesDefinition.name]: toVercelAITool(findSymbolReferencesDefinition, context),
+    [findSymbolDefinitionsDefinition.name]: toVercelAITool(findSymbolDefinitionsDefinition, context),
+    [listTreeDefinition.name]: toVercelAITool(listTreeDefinition, context),
+});
 
 export type ReadFileToolUIPart = ToolUIPart<{ read_file: SBChatMessageToolTypes['read_file'] }>;
 export type ListCommitsToolUIPart = ToolUIPart<{ list_commits: SBChatMessageToolTypes['list_commits'] }>;

packages/web/src/features/mcp/server.ts

@@ -18,6 +18,7 @@ import {
     readFileDefinition,
     registerMcpTool,
     grepDefinition,
+    ToolContext,
 } from '../tools';
 
 const dedent = _dedent.withOptions({ alignValues: true });
@@ -28,13 +29,17 @@ export function createMcpServer(): McpServer {
         version: SOURCEBOT_VERSION,
     });
 
-    registerMcpTool(server, grepDefinition);
-    registerMcpTool(server, listCommitsDefinition);
-    registerMcpTool(server, listReposDefinition);
-    registerMcpTool(server, readFileDefinition);
-    registerMcpTool(server, listTreeDefinition);
-    registerMcpTool(server, findSymbolDefinitionsDefinition);
-    registerMcpTool(server, findSymbolReferencesDefinition);
+    const toolContext: ToolContext = {
+        source: 'sourcebot-mcp-server',
+    }
+
+    registerMcpTool(server, grepDefinition, toolContext);
+    registerMcpTool(server, listCommitsDefinition, toolContext);
+    registerMcpTool(server, listReposDefinition, toolContext);
+    registerMcpTool(server, readFileDefinition, toolContext);
+    registerMcpTool(server, listTreeDefinition, toolContext);
+    registerMcpTool(server, findSymbolDefinitionsDefinition, toolContext);
+    registerMcpTool(server, findSymbolReferencesDefinition, toolContext);
 
     server.registerTool(
         "list_language_models",

packages/web/src/features/tools/adapters.ts

@@ -1,15 +1,16 @@
 import { tool } from "ai";
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { z } from "zod";
-import { ToolDefinition } from "./types";
+import { ToolContext, ToolDefinition } from "./types";
 
 export function toVercelAITool<TName extends string, TShape extends z.ZodRawShape, TMetadata>(
     def: ToolDefinition<TName, TShape, TMetadata>,
+    context: ToolContext,
 ) {
     return tool({
         description: def.description,
         inputSchema: def.inputSchema,
-        execute: (input) => def.execute(input, { source: 'sourcebot-ask-agent' }),
+        execute: (input) => def.execute(input, context),
         toModelOutput: ({ output }) => ({
             type: "content",
             value: [{ type: "text", text: output.output }],
@@ -20,6 +21,7 @@ export function toVercelAITool<TName extends string, TShape extends z.ZodRawShap
 export function registerMcpTool<TName extends string, TShape extends z.ZodRawShape, TMetadata>(
     server: McpServer,
     def: ToolDefinition<TName, TShape, TMetadata>,
+    context: ToolContext,
 ) {
     // Widening .shape to z.ZodRawShape (its base constraint) gives TypeScript a
     // concrete InputArgs so it can fully resolve BaseToolCallback's conditional
@@ -37,7 +39,7 @@ export function registerMcpTool<TName extends string, TShape extends z.ZodRawSha
         async (input) => {
             try {
                 const parsed = def.inputSchema.parse(input);
-                const result = await def.execute(parsed, { source: 'mcp' });
+                const result = await def.execute(parsed, context);
                 return { content: [{ type: "text" as const, text: result.output }] };
             } catch (error) {
                 const message = error instanceof Error ? error.message : String(error);

packages/web/src/features/tools/grep.ts

@@ -84,6 +84,8 @@ export const grepDefinition: ToolDefinition<'grep', typeof grepShape, GrepMetada
 
         if (repo) {
             query += ` repo:${escapeStringRegexp(repo)}`;
+        } else if (context.selectedRepos && context.selectedRepos.length > 0) {
+            query += ` reposet:${context.selectedRepos.join(',')}`;
         }
 
         if (ref) {

packages/web/src/features/tools/index.ts

@@ -5,4 +5,5 @@ export * from './grep';
 export * from './findSymbolReferences';
 export * from './findSymbolDefinitions';
 export * from './listTree';
-export * from './adapters';
+export * from './adapters';
+export * from './types';

packages/web/src/features/tools/types.ts

@@ -2,6 +2,7 @@ import { z } from "zod";
 
 export interface ToolContext {
     source?: string;
+    selectedRepos?: string[];
 }
 
 export interface ToolDefinition<