migrate readFile

Author: brendan-kellam

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

@@ -15,7 +15,9 @@ import {
 import { randomUUID } from "crypto";
 import _dedent from "dedent";
 import { ANSWER_TAG, FILE_REFERENCE_PREFIX, toolNames } from "./constants";
-import { createCodeSearchTool, findSymbolDefinitionsTool, findSymbolReferencesTool, listCommitsTool, listReposTool, readFileTool } from "./tools";
+import { findSymbolDefinitionsTool, findSymbolReferencesTool, listCommitsTool, listReposTool, searchCodeTool } from "./tools";
+import { toVercelAITool } from "@/features/tools/adapters";
+import { readFileDefinition } from "@/features/tools/readFile";
 import { Source } from "./types";
 import { addLineNumbers, fileReferenceToString } from "./utils";
 
@@ -199,8 +201,8 @@ const createAgentStream = async ({
         messages: inputMessages,
         system: systemPrompt,
         tools: {
-            [toolNames.searchCode]: createCodeSearchTool(selectedRepos),
-            [toolNames.readFile]: readFileTool,
+            [toolNames.searchCode]: searchCodeTool,
+            [toolNames.readFile]: toVercelAITool(readFileDefinition),
             [toolNames.findSymbolReferences]: findSymbolReferencesTool,
             [toolNames.findSymbolDefinitions]: findSymbolDefinitionsTool,
             [toolNames.listRepos]: listReposTool,
@@ -226,11 +228,11 @@ const createAgentStream = async ({
                 if (toolName === toolNames.readFile) {
                     onWriteSource({
                         type: 'file',
-                        language: output.language,
-                        repo: output.repository,
-                        path: output.path,
-                        revision: output.revision,
-                        name: output.path.split('/').pop() ?? output.path,
+                        language: output.metadata.language,
+                        repo: output.metadata.repository,
+                        path: output.metadata.path,
+                        revision: output.metadata.revision,
+                        name: output.metadata.path.split('/').pop() ?? output.metadata.path,
                     });
                 } else if (toolName === toolNames.searchCode) {
                     output.files.forEach((file) => {
@@ -310,6 +312,8 @@ const createPrompt = ({
         <selected_repositories>
         The user has explicitly selected the following repositories for analysis:
         ${repos.map(repo => `- ${repo}`).join('\n')}
+
+        When calling the searchCode tool, always pass these repositories as \`filterByRepos\` to scope results to the selected repositories.
         </selected_repositories>
     ` : ''}
 

packages/web/src/features/chat/components/chatThread/detailsCard.tsx

@@ -208,7 +208,16 @@ const DetailsCardComponent = ({
                                                         part={part}
                                                     />
                                                 )
+                                            case 'data-source':
+                                            case 'dynamic-tool':
+                                            case 'file':
+                                            case 'source-document':
+                                            case 'source-url':
+                                            case 'step-start':
+                                                return null;
                                             default:
+                                                // Guarantees this switch-case to be exhaustive
+                                                part satisfies never;
                                                 return null;
                                         }
                                     })}

packages/web/src/features/chat/components/chatThread/tools/readFileToolComponent.tsx

@@ -2,7 +2,7 @@
 
 import { CodeSnippet } from "@/app/components/codeSnippet";
 import { Separator } from "@/components/ui/separator";
-import { ReadFileToolUIPart } from "@/features/chat/tools";
+import { ReadFileToolUIPart } from "@/features/tools/registry";
 import { isServiceError } from "@/lib/utils";
 import { EyeIcon } from "lucide-react";
 import { useMemo, useState } from "react";
@@ -14,7 +14,7 @@ export const ReadFileToolComponent = ({ part }: { part: ReadFileToolUIPart }) =>
 
     const onCopy = () => {
         if (part.state !== 'output-available' || isServiceError(part.output)) return false;
-        navigator.clipboard.writeText(part.output.source);
+        navigator.clipboard.writeText(part.output.output);
         return true;
     };
 
@@ -30,7 +30,10 @@ export const ReadFileToolComponent = ({ part }: { part: ReadFileToolUIPart }) =>
                 if (isServiceError(part.output)) {
                     return 'Failed to read file';
                 }
-                return `Read ${part.output.path}`;
+                if (part.output.metadata.isTruncated || part.output.metadata.startLine > 1) {
+                    return `Read ${part.output.metadata.path} (lines ${part.output.metadata.startLine}–${part.output.metadata.endLine})`;
+                }
+                return `Read ${part.output.metadata.path}`;
         }
     }, [part]);
 
@@ -60,8 +63,8 @@ export const ReadFileToolComponent = ({ part }: { part: ReadFileToolUIPart }) =>
                             <span>Failed with the following error: <CodeSnippet className="text-sm text-destructive">{part.output.message}</CodeSnippet></span>
                         ) : (
                             <FileListItem
-                                path={part.output.path}
-                                repoName={part.output.repository}
+                                path={part.output.metadata.path}
+                                repoName={part.output.metadata.repository}
                             />
                         )}
                     </TreeList>

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

@@ -10,7 +10,6 @@
 
 export * from "./findSymbolReferences";
 export * from "./findSymbolDefinitions";
-export * from "./readFile";
 export * from "./searchCode";
 export * from "./listRepos";
 export * from "./listCommits";

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

@@ -10,7 +10,7 @@ import description from './searchCode.txt';
 
 const DEFAULT_SEARCH_LIMIT = 100;
 
-export const createCodeSearchTool = (selectedRepos: string[]) => tool({
+export const searchCodeTool = tool({
     description,
     inputSchema: z.object({
         query: z
@@ -64,10 +64,6 @@ export const createCodeSearchTool = (selectedRepos: string[]) => tool({
     }) => {
         logger.debug('searchCode', { query, useRegex, repos, languages, filepaths, caseSensitive, ref, limit });
 
-        if (selectedRepos.length > 0) {
-            query += ` reposet:${selectedRepos.join(',')}`;
-        }
-
         if (repos.length > 0) {
             query += ` (repo:${repos.map(id => escapeStringRegexp(id)).join(' or repo:')})`;
         }
@@ -115,7 +111,7 @@ export const createCodeSearchTool = (selectedRepos: string[]) => tool({
     },
 });
 
-export type SearchCodeTool = InferUITool<ReturnType<typeof createCodeSearchTool>>;
-export type SearchCodeToolInput = InferToolInput<ReturnType<typeof createCodeSearchTool>>;
-export type SearchCodeToolOutput = InferToolOutput<ReturnType<typeof createCodeSearchTool>>;
+export type SearchCodeTool = InferUITool<typeof searchCodeTool>;
+export type SearchCodeToolInput = InferToolInput<typeof searchCodeTool>;
+export type SearchCodeToolOutput = InferToolOutput<typeof searchCodeTool>;
 export type SearchCodeToolUIPart = ToolUIPart<{ [toolNames.searchCode]: SearchCodeTool }>;

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

@@ -3,8 +3,9 @@ import { BaseEditor, Descendant } from "slate";
 import { HistoryEditor } from "slate-history";
 import { ReactEditor, RenderElementProps } from "slate-react";
 import { z } from "zod";
-import { FindSymbolDefinitionsTool, FindSymbolReferencesTool, ReadFileTool, SearchCodeTool, ListReposTool, ListCommitsTool } from "./tools";
+import { FindSymbolDefinitionsTool, FindSymbolReferencesTool, SearchCodeTool, ListReposTool, ListCommitsTool } from "./tools";
 import { toolNames } from "./constants";
+import { ToolTypes } from "@/features/tools/registry";
 import { LanguageModel } from "@sourcebot/schemas/v3/index.type";
 
 const fileSourceSchema = z.object({
@@ -80,7 +81,7 @@ export type SBChatMessageMetadata = z.infer<typeof sbChatMessageMetadataSchema>;
 
 export type SBChatMessageToolTypes = {
     [toolNames.searchCode]: SearchCodeTool,
-    [toolNames.readFile]: ReadFileTool,
+    [toolNames.readFile]: ToolTypes['readFile'],
     [toolNames.findSymbolReferences]: FindSymbolReferencesTool,
     [toolNames.findSymbolDefinitions]: FindSymbolDefinitionsTool,
     [toolNames.listRepos]: ListReposTool,

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

@@ -0,0 +1,36 @@
+import { tool } from "ai";
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { z } from "zod";
+import { ToolDefinition } from "./types";
+
+export function toVercelAITool<TName extends string, TShape extends z.ZodRawShape, TMetadata>(
+    def: ToolDefinition<TName, TShape, TMetadata>,
+) {
+    return tool({
+        description: def.description,
+        inputSchema: def.inputSchema,
+        execute: def.execute,
+        toModelOutput: ({ output }) => ({
+            type: "content",
+            value: [{ type: "text", text: output.output }],
+        }),
+    });
+}
+
+export function registerMcpTool<TName extends string, TShape extends z.ZodRawShape, TMetadata>(
+    server: McpServer,
+    def: ToolDefinition<TName, TShape, TMetadata>,
+) {
+    // Widening .shape to z.ZodRawShape (its base constraint) gives TypeScript a
+    // concrete InputArgs so it can fully resolve BaseToolCallback's conditional
+    // type. def.inputSchema.parse() recovers the correctly typed value inside.
+    server.registerTool(
+        def.name,
+        { description: def.description, inputSchema: def.inputSchema.shape as z.ZodRawShape },
+        async (input) => {
+            const parsed = def.inputSchema.parse(input);
+            const result = await def.execute(parsed);
+            return { content: [{ type: "text" as const, text: result.output }] };
+        },
+    );
+}