add back selection

Author: brendan-kellam

packages/mcp/src/index.ts

@@ -8,8 +8,8 @@ import escapeStringRegexp from 'escape-string-regexp';
 import { z } from 'zod';
 import { askCodebase, getFileSource, listCommits, listRepos, search } from './client.js';
 import { env, numberSchema } from './env.js';
-import { fileSourceRequestSchema, listCommitsQueryParamsSchema, listReposQueryParamsSchema } from './schemas.js';
-import { FileSourceRequest, ListCommitsQueryParamsSchema, ListReposQueryParams, TextContent } from './types.js';
+import { askCodebaseRequestSchema, fileSourceRequestSchema, listCommitsQueryParamsSchema, listReposQueryParamsSchema } from './schemas.js';
+import { AskCodebaseRequest, FileSourceRequest, ListCommitsQueryParamsSchema, ListReposQueryParams, TextContent } from './types.js';
 
 const dedent = _dedent.withOptions({ alignValues: true });
 
@@ -253,26 +253,16 @@ server.tool(
 
     This is a blocking operation that may take 30-60+ seconds for complex questions as the agent researches the codebase.
     `,
-    {
-        question: z.string().describe("The question to ask about the codebase."),
-    },
-    async ({
-        question,
-    }) => {
-        const response = await askCodebase({
-            question,
-        });
+    askCodebaseRequestSchema.shape,
+    async (request: AskCodebaseRequest) => {
+        const response = await askCodebase(request);
 
         // Format the response with the answer and a link to the chat
         const formattedResponse = dedent`
         ${response.answer}
 
         ---
         **View full research session:** ${response.chatUrl}
-
-        **Sources referenced:** ${response.sources.length} files
-        **Response time:** ${(response.metadata.totalResponseTimeMs / 1000).toFixed(1)}s
-        **Model:** ${response.metadata.modelName}
         `;
 
         return {

packages/mcp/src/schemas.ts

@@ -275,7 +275,13 @@ export const listCommitsResponseSchema = z.array(z.object({
 }));
 
 export const askCodebaseRequestSchema = z.object({
-    question: z.string().describe("The question to ask about the codebase"),
+    query: z
+        .string()
+        .describe("The query to ask about the codebase."),
+    repos: z
+        .array(z.string())
+        .optional()
+        .describe("The repositories that are accessible to the agent during the chat. If not provided, all repositories are accessible."),
 });
 
 export const sourceSchema = z.object({
@@ -291,12 +297,4 @@ export const askCodebaseResponseSchema = z.object({
     answer: z.string().describe("The agent's final answer in markdown format"),
     chatId: z.string().describe("ID of the persisted chat session"),
     chatUrl: z.string().describe("URL to view the chat in the web UI"),
-    sources: z.array(sourceSchema).describe("Files the agent referenced during research"),
-    metadata: z.object({
-        totalTokens: z.number(),
-        inputTokens: z.number(),
-        outputTokens: z.number(),
-        totalResponseTimeMs: z.number(),
-        modelName: z.string(),
-    }).describe("Metadata about the response"),
 });

packages/web/src/app/api/(server)/chat/blocking/route.ts

@@ -1,9 +1,9 @@
 import { sew } from "@/actions";
 import { _getConfiguredLanguageModelsFull, _getAISDKLanguageModelAndOptions, updateChatMessages, generateAndUpdateChatNameFromMessage } from "@/features/chat/actions";
-import { SBChatMessage, Source } from "@/features/chat/types";
+import { SBChatMessage, SearchScope } from "@/features/chat/types";
 import { convertLLMOutputToPortableMarkdown, getAnswerPartFromAssistantMessage } from "@/features/chat/utils";
 import { ErrorCode } from "@/lib/errorCodes";
-import { requestBodySchemaValidationError, ServiceError, serviceErrorResponse } from "@/lib/serviceError";
+import { requestBodySchemaValidationError, ServiceError, ServiceErrorException, serviceErrorResponse } from "@/lib/serviceError";
 import { isServiceError } from "@/lib/utils";
 import { getBaseUrl } from "@/lib/utils.server";
 import { withOptionalAuthV2 } from "@/withAuthV2";
@@ -24,8 +24,13 @@ const logger = createLogger('chat-blocking-api');
  * This is a simpler interface designed for MCP and other programmatic integrations.
  */
 const blockingChatRequestSchema = z.object({
-    // The question to ask about the codebase
-    question: z.string().min(1, "Question is required"),
+    query: z
+        .string()
+        .describe("The query to ask about the codebase."),
+    repos: z
+        .array(z.string())
+        .optional()
+        .describe("The repositories that are accessible to the agent during the chat. If not provided, all repositories are accessible."),
 });
 
 /**
@@ -54,7 +59,7 @@ export async function POST(request: Request) {
         return serviceErrorResponse(requestBodySchemaValidationError(parsed.error));
     }
 
-    const { question } = parsed.data;
+    const { query, repos = [] } = parsed.data;
 
     const response: BlockingChatResponse | ServiceError = await sew(() =>
         withOptionalAuthV2(async ({ org, user, prisma }) => {
@@ -88,23 +93,47 @@ export async function POST(request: Request) {
             // Run the agent to completion
             logger.debug(`Starting blocking agent for chat ${chat.id}`, {
                 chatId: chat.id,
-                question: question.substring(0, 100),
+                query: query.substring(0, 100),
                 model: modelName,
             });
 
             // Create the initial user message
             const userMessage: SBChatMessage = {
                 id: randomUUID(),
                 role: 'user',
-                parts: [{ type: 'text', text: question }],
+                parts: [{ type: 'text', text: query }],
             };
 
+            const selectedSearchScopes: SearchScope[] = [];
+            for (const repo of repos) {
+                const repoDB = await prisma.repo.findFirst({
+                    where: {
+                        name: repo,
+                    },
+                });
+
+                if (!repoDB) {
+                    throw new ServiceErrorException({
+                        statusCode: StatusCodes.BAD_REQUEST,
+                        errorCode: ErrorCode.INVALID_REQUEST_BODY,
+                        message: `Repository '${repo}' not found.`,
+                    })
+                }
+
+                selectedSearchScopes.push({
+                    type: 'repo',
+                    value: repoDB.name,
+                    name: repoDB.displayName ?? repoDB.name.split('/').pop() ?? repoDB.name,
+                    codeHostType: repoDB.external_codeHostType,
+                })
+            }
+
             // We'll capture the final messages and usage from the stream
             let finalMessages: SBChatMessage[] = [];
 
             const stream = await createMessageStream({
                 messages: [userMessage],
-                selectedSearchScopes: [],
+                selectedSearchScopes,
                 model,
                 modelName,
                 modelProviderOptions: providerOptions,
@@ -122,7 +151,7 @@ export async function POST(request: Request) {
                 generateAndUpdateChatNameFromMessage({
                     chatId: chat.id,
                     languageModelId: languageModelConfig.model,
-                    message: question,
+                    message: query,
                 })
             ]);
 

packages/web/src/app/api/(server)/chat/route.ts

@@ -182,7 +182,7 @@ export const createMessageStream = async ({
 
             const startTime = new Date();
 
-            const expandedReposArrays = await Promise.all(selectedSearchScopes.map(async (scope) => {
+            const expandedRepos = (await Promise.all(selectedSearchScopes.map(async (scope) => {
                 if (scope.type === 'repo') {
                     return [scope.value];
                 }
@@ -204,15 +204,14 @@ export const createMessageStream = async ({
                 }
 
                 return [];
-            }));
-            const expandedRepos = expandedReposArrays.flat();
+            }))).flat()
 
             const researchStream = await createAgentStream({
                 model,
                 providerOptions: modelProviderOptions,
                 inputMessages: messageHistory,
                 inputSources: sources,
-                searchScopeRepoNames: expandedRepos,
+                selectedRepos: expandedRepos,
                 onWriteSource: (source) => {
                     writer.write({
                         type: 'data-source',

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

@@ -17,7 +17,7 @@ const logger = createLogger('chat-agent');
 interface AgentOptions {
     model: LanguageModel;
     providerOptions?: ProviderOptions;
-    searchScopeRepoNames: string[];
+    selectedRepos: string[];
     inputMessages: ModelMessage[];
     inputSources: Source[];
     onWriteSource: (source: Source) => void;
@@ -29,7 +29,7 @@ export const createAgentStream = async ({
     providerOptions,
     inputMessages,
     inputSources,
-    searchScopeRepoNames,
+    selectedRepos,
     onWriteSource,
     traceId,
 }: AgentOptions) => {
@@ -59,7 +59,7 @@ export const createAgentStream = async ({
     ).filter((source) => source !== undefined);
 
     const systemPrompt = createPrompt({
-        repos: searchScopeRepoNames,
+        repos: selectedRepos,
         files: resolvedFileSources,
     });
 
@@ -69,7 +69,7 @@ export const createAgentStream = async ({
         messages: inputMessages,
         system: systemPrompt,
         tools: {
-            [toolNames.searchCode]: createCodeSearchTool(searchScopeRepoNames),
+            [toolNames.searchCode]: createCodeSearchTool(selectedRepos),
             [toolNames.readFiles]: readFilesTool,
             [toolNames.findSymbolReferences]: findSymbolReferencesTool,
             [toolNames.findSymbolDefinitions]: findSymbolDefinitionsTool,
@@ -172,10 +172,12 @@ const createPrompt = ({
     During the research phase, use the tools available to you to gather comprehensive context before answering. Always explain why you're using each tool. Depending on the user's question, you may need to use multiple tools. If the question is vague, ask the user for more information.
     </research_phase_instructions>
 
-    <available_repositories>
-    The user has selected the following repositories for analysis:
-    ${repos.map(scope => `- ${scope}`).join('\n')}
-    </available_repositories>
+    ${repos.length > 0 ? dedent`
+        <selected_repositories>
+        The user has explicitly selected the following repositories for analysis:
+        ${repos.map(repo => `- ${repo}`).join('\n')}
+        </selected_repositories>
+    ` : ''}
 
     ${files ? dedent`
         <files>