diff --git a/packages/mcp/src/client.ts b/packages/mcp/src/client.ts index 036f251b..fa92a345 100644 --- a/packages/mcp/src/client.ts +++ b/packages/mcp/src/client.ts @@ -1,6 +1,6 @@ import { env } from './env.js'; -import { listReposResponseSchema, searchResponseSchema, fileSourceResponseSchema, listCommitsResponseSchema } from './schemas.js'; -import { FileSourceRequest, ListReposQueryParams, SearchRequest, ListCommitsQueryParamsSchema } from './types.js'; +import { listReposResponseSchema, searchResponseSchema, fileSourceResponseSchema, listCommitsResponseSchema, askCodebaseResponseSchema } from './schemas.js'; +import { AskCodebaseRequest, AskCodebaseResponse, FileSourceRequest, ListReposQueryParams, SearchRequest, ListCommitsQueryParamsSchema } from './types.js'; import { isServiceError, ServiceErrorException } from './utils.js'; import { z } from 'zod'; @@ -103,3 +103,23 @@ export const listCommits = async (queryParams: ListCommitsQueryParamsSchema) => const totalCount = parseInt(response.headers.get('X-Total-Count') ?? '0', 10); return { commits, totalCount }; } + +/** + * Asks a natural language question about the codebase using the Sourcebot AI agent. + * This is a blocking call that runs the full agent loop and returns when complete. + * + * @param request - The question and optional repo filters + * @returns The agent's answer, chat URL, sources, and metadata + */ +export const askCodebase = async (request: AskCodebaseRequest): Promise => { + const response = await fetch(`${env.SOURCEBOT_HOST}/api/chat/blocking`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + ...(env.SOURCEBOT_API_KEY ? { 'X-Sourcebot-Api-Key': env.SOURCEBOT_API_KEY } : {}) + }, + body: JSON.stringify(request), + }); + + return parseResponse(response, askCodebaseResponseSchema); +} diff --git a/packages/mcp/src/index.ts b/packages/mcp/src/index.ts index e306d867..80ad6bc0 100644 --- a/packages/mcp/src/index.ts +++ b/packages/mcp/src/index.ts @@ -6,10 +6,10 @@ import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js' import _dedent from "dedent"; import escapeStringRegexp from 'escape-string-regexp'; import { z } from 'zod'; -import { getFileSource, listCommits, listRepos, search } from './client.js'; +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 }); @@ -239,7 +239,53 @@ server.tool( } ); +server.tool( + "ask_codebase", + dedent` + Ask a natural language question about the codebase. This tool uses an AI agent to autonomously search code, read files, and find symbol references/definitions to answer your question. + + The agent will: + - Analyze your question and determine what context it needs + - Search the codebase using multiple strategies (code search, symbol lookup, file reading) + - Synthesize findings into a comprehensive answer with code references + + Returns a detailed answer in markdown format with code references, plus a link to view the full research session (including all tool calls and reasoning) in the Sourcebot web UI. + + 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."), + 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 + 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 { + content: [{ + type: "text", + text: formattedResponse, + }], + }; + } +); const runServer = async () => { const transport = new StdioServerTransport(); diff --git a/packages/mcp/src/schemas.ts b/packages/mcp/src/schemas.ts index fd89bdfb..a87aea37 100644 --- a/packages/mcp/src/schemas.ts +++ b/packages/mcp/src/schemas.ts @@ -274,3 +274,35 @@ export const listCommitsResponseSchema = z.array(z.object({ author_name: z.string(), 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({ + type: z.literal('file'), + repo: z.string(), + path: z.string(), + name: z.string(), + language: z.string(), + revision: z.string(), +}); + +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"), +}); diff --git a/packages/mcp/src/types.ts b/packages/mcp/src/types.ts index cd64cb08..8a721970 100644 --- a/packages/mcp/src/types.ts +++ b/packages/mcp/src/types.ts @@ -12,6 +12,8 @@ import { serviceErrorSchema, listCommitsQueryParamsSchema, listCommitsResponseSchema, + askCodebaseRequestSchema, + askCodebaseResponseSchema, } from "./schemas.js"; import { z } from "zod"; @@ -34,3 +36,6 @@ export type ServiceError = z.infer; export type ListCommitsQueryParamsSchema = z.infer; export type ListCommitsResponse = z.infer; + +export type AskCodebaseRequest = z.infer; +export type AskCodebaseResponse = z.infer; diff --git a/packages/web/src/app/api/(server)/chat/blocking/route.ts b/packages/web/src/app/api/(server)/chat/blocking/route.ts new file mode 100644 index 00000000..fd6a0a61 --- /dev/null +++ b/packages/web/src/app/api/(server)/chat/blocking/route.ts @@ -0,0 +1,257 @@ +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 { ErrorCode } from "@/lib/errorCodes"; +import { requestBodySchemaValidationError, ServiceError, serviceErrorResponse } from "@/lib/serviceError"; +import { isServiceError } from "@/lib/utils"; +import { getBaseUrl } from "@/lib/utils.server"; +import { withOptionalAuthV2 } from "@/withAuthV2"; +import { ChatVisibility, Prisma } from "@sourcebot/db"; +import { createLogger } from "@sourcebot/shared"; +import { randomUUID } from "crypto"; +import { StatusCodes } from "http-status-codes"; +import { headers } from "next/headers"; +import { NextResponse } from "next/server"; +import { z } from "zod"; + +const logger = createLogger('chat-blocking-api'); + +/** + * Request schema for the blocking chat 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"), + // 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; + }; +} + +/** + * POST /api/chat/blocking + * + * A blocking (non-streaming) chat endpoint designed for MCP and other integrations. + * Creates a chat session, runs the agent to completion, and returns the final answer. + * + * The chat session is persisted to the database, allowing users to view the full + * conversation (including tool calls and reasoning) in the web UI. + */ +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 response: BlockingChatResponse | ServiceError = await sew(() => + withOptionalAuthV2(async ({ org, user, prisma }) => { + // Get all configured language models + const configuredModels = await _getConfiguredLanguageModelsFull(); + if (configuredModels.length === 0) { + return { + statusCode: StatusCodes.BAD_REQUEST, + errorCode: ErrorCode.INVALID_REQUEST_BODY, + message: "No language models are configured. Please configure at least one language model.", + } 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; + } + + + 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: { + orgId: org.id, + createdById: user?.id, + visibility: ChatVisibility.PRIVATE, + messages: [] as unknown as Prisma.InputJsonValue, + }, + }); + + const traceId = randomUUID(); + + // Run the agent to completion + logger.info(`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 + 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, + }, + }; + + // Persist the messages to the chat + await updateChatMessages({ + chatId: chat.id, + messages: [userMessage, assistantMessage], + }); + + logger.info(`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; + }) + ); + + if (isServiceError(response)) { + return serviceErrorResponse(response); + } + + console.log(response); + + return NextResponse.json(response); +} diff --git a/packages/web/src/app/api/(server)/chat/route.ts b/packages/web/src/app/api/(server)/chat/route.ts index 2874c48f..37fc3180 100644 --- a/packages/web/src/app/api/(server)/chat/route.ts +++ b/packages/web/src/app/api/(server)/chat/route.ts @@ -4,7 +4,7 @@ import { createAgentStream } from "@/features/chat/agent"; import { additionalChatRequestParamsSchema, LanguageModelInfo, SBChatMessage, SearchScope } from "@/features/chat/types"; import { getAnswerPartFromAssistantMessage, getLanguageModelKey } from "@/features/chat/utils"; import { ErrorCode } from "@/lib/errorCodes"; -import { notFound, requestBodySchemaValidationError, serviceErrorResponse } from "@/lib/serviceError"; +import { notFound, requestBodySchemaValidationError, ServiceError, serviceErrorResponse } from "@/lib/serviceError"; import { isServiceError } from "@/lib/utils"; import { withOptionalAuthV2 } from "@/withAuthV2"; import { LanguageModelV2 as AISDKLanguageModelV2 } from "@ai-sdk/provider"; @@ -61,11 +61,11 @@ export async function POST(req: Request) { } if (chat.isReadonly) { - return serviceErrorResponse({ + return { statusCode: StatusCodes.BAD_REQUEST, errorCode: ErrorCode.INVALID_REQUEST_BODY, message: "Chat is readonly and cannot be edited.", - }); + } satisfies ServiceError; } // From the language model ID, attempt to find the @@ -75,11 +75,11 @@ export async function POST(req: Request) { .find((model) => getLanguageModelKey(model) === getLanguageModelKey(languageModel)); if (!languageModelConfig) { - return serviceErrorResponse({ + return { statusCode: StatusCodes.BAD_REQUEST, errorCode: ErrorCode.INVALID_REQUEST_BODY, message: `Language model ${languageModel.model} is not configured.`, - }); + } satisfies ServiceError; } const { model, providerOptions } = await _getAISDKLanguageModelAndOptions(languageModelConfig); diff --git a/packages/web/src/features/chat/agent.ts b/packages/web/src/features/chat/agent.ts index efeb12f8..f482a54c 100644 --- a/packages/web/src/features/chat/agent.ts +++ b/packages/web/src/features/chat/agent.ts @@ -4,7 +4,7 @@ import { getFileSource } from "@/features/search/fileSourceApi"; import { isServiceError } from "@/lib/utils"; import { ProviderOptions } from "@ai-sdk/provider-utils"; import { createLogger } from "@sourcebot/shared"; -import { LanguageModel, ModelMessage, StopCondition, streamText } from "ai"; +import { generateText, LanguageModel, ModelMessage, StopCondition, streamText } from "ai"; import { ANSWER_TAG, FILE_REFERENCE_PREFIX, toolNames } from "./constants"; import { createCodeSearchTool, findSymbolDefinitionsTool, findSymbolReferencesTool, readFilesTool, searchReposTool, listAllReposTool } from "./tools"; import { FileSource, Source } from "./types"; @@ -266,4 +266,165 @@ const resolveFileSource = async ({ path, repo, revision }: FileSource) => { language: fileSource.language, revision, } +} + +// ============================================================================ +// Blocking Agent Execution (for MCP and other non-streaming use cases) +// ============================================================================ + +interface BlockingAgentOptions { + model: LanguageModel; + providerOptions?: ProviderOptions; + searchScopeRepoNames: string[]; + inputMessages: ModelMessage[]; + inputSources: Source[]; + traceId: string; +} + +export interface BlockingAgentResult { + text: string; + sources: Source[]; + usage: { + inputTokens: number; + outputTokens: number; + totalTokens: number; + }; + responseTimeMs: number; +} + +/** + * Runs the chat agent in blocking mode, waiting for the complete response. + * This is used by the MCP server and other integrations that don't support streaming. + */ +export const runAgentBlocking = async ({ + model, + providerOptions, + inputMessages, + inputSources, + searchScopeRepoNames, + traceId, +}: BlockingAgentOptions): Promise => { + const startTime = Date.now(); + const collectedSources: Source[] = []; + + const onWriteSource = (source: Source) => { + // Deduplicate sources by checking if we already have this file + const exists = collectedSources.some( + (s) => s.type === source.type && + s.type === 'file' && source.type === 'file' && + s.repo === source.repo && + s.path === source.path + ); + if (!exists) { + collectedSources.push(source); + } + }; + + const baseSystemPrompt = createBaseSystemPrompt({ searchScopeRepoNames }); + + // Resolve any input file sources for the first step + let systemPromptWithSources = baseSystemPrompt; + if (inputSources.length > 0) { + const fileSources = inputSources.filter((source) => source.type === 'file'); + const resolvedFileSources = ( + await Promise.all(fileSources.map(resolveFileSource)) + ).filter((source) => source !== undefined); + + if (resolvedFileSources.length > 0) { + const fileSourcesSystemPrompt = await createFileSourcesSystemPrompt({ + files: resolvedFileSources + }); + systemPromptWithSources = `${baseSystemPrompt}\n\n${fileSourcesSystemPrompt}`; + } + } + + const result = await generateText({ + model, + providerOptions, + system: systemPromptWithSources, + messages: inputMessages, + tools: { + [toolNames.searchCode]: createCodeSearchTool(searchScopeRepoNames), + [toolNames.readFiles]: readFilesTool, + [toolNames.findSymbolReferences]: findSymbolReferencesTool, + [toolNames.findSymbolDefinitions]: findSymbolDefinitionsTool, + [toolNames.searchRepos]: searchReposTool, + [toolNames.listAllRepos]: listAllReposTool, + }, + temperature: env.SOURCEBOT_CHAT_MODEL_TEMPERATURE, + stopWhen: [ + stepCountIsGTE(env.SOURCEBOT_CHAT_MAX_STEP_COUNT), + ], + toolChoice: "auto", + onStepFinish: ({ toolResults }) => { + // Extract sources from tool results (same logic as streaming version) + toolResults.forEach(({ toolName, output, dynamic }) => { + // We don't care about dynamic tool results here. + if (dynamic) { + return; + } + + if (isServiceError(output)) { + return; + } + + if (toolName === toolNames.readFiles) { + (output as { path: string; repository: string; language: string; revision: string }[]).forEach((file) => { + onWriteSource({ + type: 'file', + language: file.language, + repo: file.repository, + path: file.path, + revision: file.revision, + name: file.path.split('/').pop() ?? file.path, + }); + }); + } + else if (toolName === toolNames.searchCode) { + const searchOutput = output as { files: { language: string; repository: string; fileName: string; revision: string }[] }; + searchOutput.files.forEach((file) => { + onWriteSource({ + type: 'file', + language: file.language, + repo: file.repository, + path: file.fileName, + revision: file.revision, + name: file.fileName.split('/').pop() ?? file.fileName, + }); + }); + } + else if (toolName === toolNames.findSymbolDefinitions || toolName === toolNames.findSymbolReferences) { + (output as { language: string; repository: string; fileName: string; revision: string }[]).forEach((file) => { + onWriteSource({ + type: 'file', + language: file.language, + repo: file.repository, + path: file.fileName, + revision: file.revision, + name: file.fileName.split('/').pop() ?? file.fileName, + }); + }); + } + }); + }, + experimental_telemetry: { + isEnabled: clientEnv.NEXT_PUBLIC_SOURCEBOT_CLOUD_ENVIRONMENT !== undefined, + metadata: { + langfuseTraceId: traceId, + }, + }, + }); + + const responseTimeMs = Date.now() - startTime; + + return { + text: result.text, + sources: collectedSources, + usage: { + inputTokens: result.totalUsage.inputTokens ?? 0, + outputTokens: result.totalUsage.outputTokens ?? 0, + totalTokens: result.totalUsage.totalTokens ?? 0, + }, + responseTimeMs, + }; } \ No newline at end of file