wip

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 { askCodebaseRequestSchema, fileSourceRequestSchema, listCommitsQueryParamsSchema, listReposQueryParamsSchema } from './schemas.js';
-import { AskCodebaseRequest, FileSourceRequest, ListCommitsQueryParamsSchema, ListReposQueryParams, TextContent } from './types.js';
+import { fileSourceRequestSchema, listCommitsQueryParamsSchema, listReposQueryParamsSchema } from './schemas.js';
+import { FileSourceRequest, ListCommitsQueryParamsSchema, ListReposQueryParams, TextContent } from './types.js';
 
 const dedent = _dedent.withOptions({ alignValues: true });
 
@@ -255,15 +255,12 @@ server.tool(
     `,
     {
         question: z.string().describe("The question to ask about the codebase."),
-        repo: z.string().describe("The repository to ask the question on."),
     },
     async ({
         question,
-        repo,
     }) => {
         const response = await askCodebase({
             question,
-            repos: [repo],
         });
 
         // Format the response with the answer and a link to the chat

packages/mcp/src/schemas.ts

@@ -274,13 +274,8 @@ export const listCommitsResponseSchema = z.array(z.object({
     author_email: z.string(),
 }));
 
-// ============================================================================
-// Ask Codebase (Blocking Chat API)
-// ============================================================================
-
 export const askCodebaseRequestSchema = z.object({
     question: z.string().describe("The question to ask about the codebase"),
-    repos: z.array(z.string()).optional().describe("Optional: filter to specific repositories by name"),
 });
 
 export const sourceSchema = z.object({

packages/web/package.json

@@ -134,6 +134,7 @@
     "codemirror-lang-spreadsheet": "^1.3.0",
     "codemirror-lang-zig": "^0.1.0",
     "date-fns": "^4.1.0",
+    "dedent": "^1.7.1",
     "embla-carousel-auto-scroll": "^8.3.0",
     "embla-carousel-react": "^8.3.0",
     "escape-string-regexp": "^5.0.0",

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

@@ -1,9 +1,7 @@
 import { sew } from "@/actions";
-import { _getConfiguredLanguageModelsFull, _getAISDKLanguageModelAndOptions, updateChatMessages } from "@/features/chat/actions";
-import { runAgentBlocking } from "@/features/chat/agent";
-import { ANSWER_TAG } from "@/features/chat/constants";
-import { LanguageModelInfo, SBChatMessage, Source } from "@/features/chat/types";
-import { convertLLMOutputToPortableMarkdown, getLanguageModelKey } from "@/features/chat/utils";
+import { _getConfiguredLanguageModelsFull, _getAISDKLanguageModelAndOptions, updateChatMessages, generateAndUpdateChatNameFromMessage } from "@/features/chat/actions";
+import { SBChatMessage, Source } from "@/features/chat/types";
+import { convertLLMOutputToPortableMarkdown, getAnswerPartFromAssistantMessage } from "@/features/chat/utils";
 import { ErrorCode } from "@/lib/errorCodes";
 import { requestBodySchemaValidationError, ServiceError, serviceErrorResponse } from "@/lib/serviceError";
 import { isServiceError } from "@/lib/utils";
@@ -16,6 +14,8 @@ import { StatusCodes } from "http-status-codes";
 import { headers } from "next/headers";
 import { NextResponse } from "next/server";
 import { z } from "zod";
+import { createMessageStream } from "../route";
+import { InferUIMessageChunk, UITools, UIDataTypes, UIMessage } from "ai";
 
 const logger = createLogger('chat-blocking-api');
 
@@ -26,36 +26,15 @@ const logger = createLogger('chat-blocking-api');
 const blockingChatRequestSchema = z.object({
     // The question to ask about the codebase
     question: z.string().min(1, "Question is required"),
-    // Optional: filter to specific repositories (by name)
-    repos: z.array(z.string()).optional(),
-    // Optional: specify a language model (defaults to first configured model)
-    languageModel: z.object({
-        provider: z.string(),
-        model: z.string(),
-        displayName: z.string().optional(),
-    }).optional(),
 });
 
 /**
  * Response schema for the blocking chat API.
  */
 interface BlockingChatResponse {
-    // The agent's final answer (markdown format)
     answer: string;
-    // ID of the persisted chat session
     chatId: string;
-    // URL to view the chat in the web UI
     chatUrl: string;
-    // Files the agent referenced during research
-    sources: Source[];
-    // Metadata about the response
-    metadata: {
-        totalTokens: number;
-        inputTokens: number;
-        outputTokens: number;
-        totalResponseTimeMs: number;
-        modelName: string;
-    };
 }
 
 /**
@@ -70,12 +49,12 @@ interface BlockingChatResponse {
 export async function POST(request: Request) {
     const requestBody = await request.json();
     const parsed = await blockingChatRequestSchema.safeParseAsync(requestBody);
-    
+
     if (!parsed.success) {
         return serviceErrorResponse(requestBodySchemaValidationError(parsed.error));
     }
 
-    const { question, repos, languageModel: requestedLanguageModel } = parsed.data;
+    const { question } = parsed.data;
 
     const response: BlockingChatResponse | ServiceError = await sew(() =>
         withOptionalAuthV2(async ({ org, user, prisma }) => {
@@ -89,64 +68,13 @@ export async function POST(request: Request) {
                 } satisfies ServiceError;
             }
 
-            // Select the language model to use
-            let languageModelConfig = configuredModels[0]; // Default to first configured model
-            
-            if (requestedLanguageModel) {
-                const requested = requestedLanguageModel as LanguageModelInfo;
-                const found = configuredModels.find(
-                    (model) => getLanguageModelKey(model) === getLanguageModelKey(requested)
-                );
-                if (!found) {
-                    return {
-                        statusCode: StatusCodes.BAD_REQUEST,
-                        errorCode: ErrorCode.INVALID_REQUEST_BODY,
-                        message: `Language model ${requested.model} is not configured.`,
-                    } satisfies ServiceError;
-                }
-                languageModelConfig = found;
-            }
+            // @todo: we should probably have a option of passing the language model
+            // into the request body. For now, just use the first configured model.
+            const languageModelConfig = configuredModels[0];
 
-            
             const { model, providerOptions } = await _getAISDKLanguageModelAndOptions(languageModelConfig);
             const modelName = languageModelConfig.displayName ?? languageModelConfig.model;
 
-            // Determine which repos to search
-            let searchScopeRepoNames: string[];
-            
-            if (repos && repos.length > 0) {
-                // Use the provided repos filter
-                // Validate that these repos exist and the user has access
-                const validRepos = await prisma.repo.findMany({
-                    where: {
-                        orgId: org.id,
-                        name: {
-                            in: repos,
-                        },
-                    },
-                    select: { name: true },
-                });
-                
-                searchScopeRepoNames = validRepos.map(r => r.name);
-                
-                if (searchScopeRepoNames.length === 0) {
-                    return {
-                        statusCode: StatusCodes.BAD_REQUEST,
-                        errorCode: ErrorCode.INVALID_REQUEST_BODY,
-                        message: "None of the specified repositories were found or accessible.",
-                    } satisfies ServiceError;
-                }
-            } else {
-                // Search all repos the user has access to
-                const allRepos = await prisma.repo.findMany({
-                    where: {
-                        orgId: org.id,
-                    },
-                    select: { name: true },
-                });
-                searchScopeRepoNames = allRepos.map(r => r.name);
-            }
-
             // Create a new chat session
             const chat = await prisma.chat.create({
                 data: {
@@ -157,92 +85,76 @@ export async function POST(request: Request) {
                 },
             });
 
-            const traceId = randomUUID();
-
             // Run the agent to completion
-            logger.info(`Starting blocking agent for chat ${chat.id}`, {
+            logger.debug(`Starting blocking agent for chat ${chat.id}`, {
                 chatId: chat.id,
                 question: question.substring(0, 100),
-                repoCount: searchScopeRepoNames.length,
                 model: modelName,
             });
 
-            const agentResult = await runAgentBlocking({
-                model,
-                providerOptions,
-                searchScopeRepoNames,
-                inputMessages: [{ role: 'user', content: question }],
-                inputSources: [],
-                traceId,
-            });
-
-            // Extract the answer (removing the answer tag if present)
-            let answer = agentResult.text;
-            if (answer.startsWith(ANSWER_TAG)) {
-                answer = answer.slice(ANSWER_TAG.length).trim();
-            }
-
-            // Convert to portable markdown (replaces @file: references with markdown links)
-            const portableAnswer = convertLLMOutputToPortableMarkdown(answer);
-
-            // Build the chat URL
-            const headersList = await headers();
-            const baseUrl = getBaseUrl(headersList);
-            const chatUrl = `${baseUrl}/${org.domain}/chat/${chat.id}`;
-
-            // Create the message history for persistence
+            // Create the initial user message
             const userMessage: SBChatMessage = {
                 id: randomUUID(),
                 role: 'user',
                 parts: [{ type: 'text', text: question }],
             };
 
-            const assistantMessage: SBChatMessage = {
-                id: randomUUID(),
-                role: 'assistant',
-                parts: [
-                    { type: 'text', text: agentResult.text },
-                    // Include sources as data parts
-                    ...agentResult.sources.map((source) => ({
-                        type: 'data-source' as const,
-                        data: source,
-                    })),
-                ],
-                metadata: {
-                    totalTokens: agentResult.usage.totalTokens,
-                    totalInputTokens: agentResult.usage.inputTokens,
-                    totalOutputTokens: agentResult.usage.outputTokens,
-                    totalResponseTimeMs: agentResult.responseTimeMs,
-                    modelName,
-                    traceId,
+            // We'll capture the final messages and usage from the stream
+            let finalMessages: SBChatMessage[] = [];
+
+            const stream = await createMessageStream({
+                messages: [userMessage],
+                selectedSearchScopes: [],
+                model,
+                modelName,
+                modelProviderOptions: providerOptions,
+                orgId: org.id,
+                prisma,
+                onFinish: async ({ messages }) => {
+                    finalMessages = messages;
                 },
-            };
+            })
+
+            await Promise.all([
+                // Consume the stream fully to trigger onFinish
+                blockStreamUntilFinish(stream),
+                // Generate and update the chat name
+                generateAndUpdateChatNameFromMessage({
+                    chatId: chat.id,
+                    languageModelId: languageModelConfig.model,
+                    message: question,
+                })
+            ]);
 
             // Persist the messages to the chat
             await updateChatMessages({
                 chatId: chat.id,
-                messages: [userMessage, assistantMessage],
+                messages: finalMessages,
             });
 
-            logger.info(`Completed blocking agent for chat ${chat.id}`, {
+            // Extract the answer text from the assistant message
+            const assistantMessage = finalMessages.find(m => m.role === 'assistant');
+            const answerPart = assistantMessage
+                ? getAnswerPartFromAssistantMessage(assistantMessage, false)
+                : undefined;
+            const answerText = answerPart?.text ?? '';
+
+            // Convert to portable markdown (replaces @file: references with markdown links)
+            const portableAnswer = convertLLMOutputToPortableMarkdown(answerText);
+
+            // Build the chat URL
+            const headersList = await headers();
+            const baseUrl = getBaseUrl(headersList);
+            const chatUrl = `${baseUrl}/${org.domain}/chat/${chat.id}`;
+
+            logger.debug(`Completed blocking agent for chat ${chat.id}`, {
                 chatId: chat.id,
-                responseTimeMs: agentResult.responseTimeMs,
-                totalTokens: agentResult.usage.totalTokens,
-                sourceCount: agentResult.sources.length,
             });
 
             return {
                 answer: portableAnswer,
                 chatId: chat.id,
                 chatUrl,
-                sources: agentResult.sources,
-                metadata: {
-                    totalTokens: agentResult.usage.totalTokens,
-                    inputTokens: agentResult.usage.inputTokens,
-                    outputTokens: agentResult.usage.outputTokens,
-                    totalResponseTimeMs: agentResult.responseTimeMs,
-                    modelName,
-                },
             } satisfies BlockingChatResponse;
         })
     );
@@ -251,7 +163,13 @@ export async function POST(request: Request) {
         return serviceErrorResponse(response);
     }
 
-    console.log(response);
-
     return NextResponse.json(response);
 }
+
+const blockStreamUntilFinish = async <T extends UIMessage<unknown, UIDataTypes, UITools>>(stream: ReadableStream<InferUIMessageChunk<T>>) => {
+    const reader = stream.getReader();
+    while (true as const) {
+        const { done } = await reader.read();
+        if (done) break;
+    }
+}