diff --git a/CHANGELOG.md b/CHANGELOG.md index fa46b346e..689718d36 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Changed +- [EE] Improved Ask Sourcebot prompt caching by splitting static and dynamic prompt sections and advancing cache breakpoints after every agent step instead of only after each message. [#1366](https://github.com/sourcebot-dev/sourcebot/pull/1366) + ### Added - Added per-step token cost tracking and estimated tool call token usage to Ask Sourcebot chat history. [#1353](https://github.com/sourcebot-dev/sourcebot/pull/1353) diff --git a/packages/shared/src/env.server.ts b/packages/shared/src/env.server.ts index ee294b1db..57c83f9f0 100644 --- a/packages/shared/src/env.server.ts +++ b/packages/shared/src/env.server.ts @@ -311,6 +311,14 @@ const options = { SOURCEBOT_CHAT_MAX_STEP_COUNT: numberSchema.default(100), SOURCEBOT_CHAT_PROMPT_CACHING_ENABLED: booleanSchema.default('true'), + /** TTL for the static block. The moving tail marker always uses the 5m default. */ + SOURCEBOT_CHAT_PROMPT_CACHE_STATIC_TTL: z.enum(['5m', '1h']).default('5m'), + /** + * Observability: when enabled, logs a warning on unexpected prompt-cache + * breaks (static-prefix signature changes, or zero cache reads on a + * continuation step). Does not affect request behavior. + */ + SOURCEBOT_CHAT_PROMPT_CACHE_BREAK_DETECTION_ENABLED: booleanSchema.default('false'), SOURCEBOT_MCP_TOOL_CALL_TIMEOUT_MS: numberSchema.int().positive().max(maxTimerDelayMs).default(60000), DEBUG_WRITE_CHAT_MESSAGES_TO_FILE: booleanSchema.default('false'), diff --git a/packages/web/src/app/api/(server)/ee/chat/route.ts b/packages/web/src/app/api/(server)/ee/chat/route.ts index fcca63437..0f20ee8e3 100644 --- a/packages/web/src/app/api/(server)/ee/chat/route.ts +++ b/packages/web/src/app/api/(server)/ee/chat/route.ts @@ -1,6 +1,7 @@ import { sew } from "@/middleware/sew"; import { getAskMcpAvailabilityAnalytics, getAskMcpTurnCompletedAnalytics } from "@/ee/features/chat/askMcpAnalytics.server"; import { createMessageStream } from "@/ee/features/chat/agent"; +import { getPromptCacheStrategy } from "@/ee/features/chat/promptCaching"; import { additionalChatRequestParamsSchema } from "@/features/chat/types"; import { getLanguageModelKey } from "@/features/chat/utils"; import { checkAskEntitlement, getConfiguredLanguageModels, isOwnerOfChat, updateChatMessages } from "@/features/chat/utils.server"; @@ -88,6 +89,13 @@ export const POST = apiHandler(async (req: NextRequest) => { const { model, providerOptions, temperature } = await getAISDKLanguageModelAndOptions(languageModelConfig); + // No-op for non-Anthropic providers / when caching is disabled, so + // it never perturbs other providers' requests. + const promptCacheStrategy = getPromptCacheStrategy( + languageModelConfig.provider, + env.SOURCEBOT_CHAT_PROMPT_CACHING_ENABLED === 'true', + ); + const expandedRepos = (await Promise.all(selectedSearchScopes.map(async (scope) => { if (scope.type === 'repo') return [scope.value]; if (scope.type === 'reposet') { @@ -131,6 +139,7 @@ export const POST = apiHandler(async (req: NextRequest) => { disabledMcpServerIds, model, modelName: languageModelConfig.displayName ?? languageModelConfig.model, + promptCacheStrategy, modelProviderOptions: providerOptions, modelTemperature: temperature, userId: user?.id, diff --git a/packages/web/src/ee/features/chat/agent.test.ts b/packages/web/src/ee/features/chat/agent.test.ts index 5ac87c551..a16c4041c 100644 --- a/packages/web/src/ee/features/chat/agent.test.ts +++ b/packages/web/src/ee/features/chat/agent.test.ts @@ -1,6 +1,8 @@ import { beforeEach, describe, expect, test, vi } from 'vitest'; import type { ModelMessage } from 'ai'; +import type { ProviderOptions } from '@ai-sdk/provider-utils'; import type { SBChatMessage, SBChatMessagePart } from '@/features/chat/types'; +import type { PromptCacheStrategy } from './promptCaching'; const mockLogger = vi.hoisted(() => ({ debug: vi.fn(), @@ -25,6 +27,8 @@ vi.mock('@sourcebot/shared', () => ({ SOURCEBOT_CHAT_MAX_STEP_COUNT: 8, SOURCEBOT_CHAT_MODEL_TEMPERATURE: 0, SOURCEBOT_TELEMETRY_PII_COLLECTION_ENABLED: 'false', + SOURCEBOT_CHAT_PROMPT_CACHE_STATIC_TTL: '5m', + SOURCEBOT_CHAT_PROMPT_CACHE_BREAK_DETECTION_ENABLED: 'false', }, getDBConnectionString: () => 'postgresql://sourcebot:sourcebot@db.example.com:5432/sourcebot', })); @@ -94,6 +98,11 @@ vi.mock('ai', async (importOriginal) => { }); const { createMessageStream } = await import('./agent'); +const { getPromptCacheStrategy } = await import('./promptCaching'); + +// Strategies reused across the prompt-caching tests below. +const anthropicStrategy = getPromptCacheStrategy('anthropic', true); +const noopStrategy = getPromptCacheStrategy('openai', true); const listReposInput = { sort: 'name', @@ -150,7 +159,29 @@ const createFakeStreamResult = () => ({ }), }); -const runCreateMessageStream = async (messages: SBChatMessage[]) => { +type FakePrepareStep = (opts: { + steps: Array<{ toolResults: Array<{ toolName: string; output: unknown }> }>; + stepNumber: number; + model: unknown; + messages: ModelMessage[]; +}) => + | { messages?: ModelMessage[]; activeTools?: string[] } + | Promise<{ messages?: ModelMessage[]; activeTools?: string[] }>; + +interface StreamTextArgs { + messages: ModelMessage[]; + system: Array<{ role: 'system'; content: string; providerOptions?: ProviderOptions }>; + tools: Record; + prepareStep?: FakePrepareStep; +} + +const runCreateMessageStream = async ( + messages: SBChatMessage[], + opts: { + promptCacheStrategy?: PromptCacheStrategy; + selectedRepos?: string[]; + } = {}, +): Promise => { const convertedLastTurn: ModelMessage = { role: 'assistant', content: 'converted-last-turn', @@ -161,10 +192,13 @@ const runCreateMessageStream = async (messages: SBChatMessage[]) => { const props = { chatId: 'chat-id', messages, - selectedRepos: [], + selectedRepos: opts.selectedRepos ?? [], prisma: {}, model: {}, modelName: 'test-model', + // Default to a no-op strategy so the approval-continuation tests below + // (which assert plain, unmarked messages) are unaffected by caching. + promptCacheStrategy: opts.promptCacheStrategy ?? noopStrategy, onFinish: vi.fn(), onError: () => 'error', } as unknown as Parameters[0]; @@ -188,7 +222,7 @@ const runCreateMessageStream = async (messages: SBChatMessage[]) => { throw new Error('Expected streamText to be called with messages.'); } - return streamTextArgs.messages as ModelMessage[]; + return streamTextArgs as StreamTextArgs; }; beforeEach(() => { @@ -216,7 +250,7 @@ describe('createMessageStream approval continuation', () => { approvalPart, ]); - const streamTextMessages = await runCreateMessageStream([ + const { messages: streamTextMessages } = await runCreateMessageStream([ createUserMessage(), assistantMessage, ]); @@ -249,7 +283,7 @@ describe('createMessageStream approval continuation', () => { dynamicApprovalRespondedPart, ]); - const streamTextMessages = await runCreateMessageStream([ + const { messages: streamTextMessages } = await runCreateMessageStream([ createUserMessage(), assistantMessage, ]); @@ -265,3 +299,151 @@ describe('createMessageStream approval continuation', () => { }); }); }); + +const EPHEMERAL = { type: 'ephemeral' }; + +describe('createMessageStream prompt caching', () => { + test('marks the static system block for the Anthropic family', async () => { + const { system, messages } = await runCreateMessageStream([createUserMessage()], { + promptCacheStrategy: anthropicStrategy, + }); + + // No repos / files / MCP tools → only the static system block. + expect(system).toHaveLength(1); + expect(system[0].providerOptions?.anthropic?.cacheControl).toEqual(EPHEMERAL); + + // The tail marker is applied per-step in prepareStep, not on the messages + // handed to streamText — those stay unmarked. + for (const message of messages) { + expect(message.providerOptions).toBeUndefined(); + } + }); + + test('moves the tail marker onto the last message of each step via prepareStep', async () => { + const { prepareStep } = await runCreateMessageStream([createUserMessage()], { + promptCacheStrategy: anthropicStrategy, + }); + expect(prepareStep).toBeTypeOf('function'); + + // Step 0: a single input message → marker lands on it. + const step0 = await prepareStep!({ + steps: [], + stepNumber: 0, + model: {}, + messages: [{ role: 'user', content: 'q' }], + }); + expect(step0.messages?.at(-1)?.providerOptions?.anthropic?.cacheControl).toEqual(EPHEMERAL); + + // Continuation step: the marker rides the NEW last message, and only it — + // earlier messages (including the prior tail) carry no marker. + const stepN = await prepareStep!({ + steps: [], + stepNumber: 1, + model: {}, + messages: [ + { role: 'user', content: 'q' }, + { role: 'assistant', content: 'searching' }, + { role: 'assistant', content: 'tool output' }, + ], + }); + const out = stepN.messages!; + expect(out[0].providerOptions?.anthropic?.cacheControl).toBeUndefined(); + expect(out[1].providerOptions?.anthropic?.cacheControl).toBeUndefined(); + expect(out.at(-1)?.providerOptions?.anthropic?.cacheControl).toEqual(EPHEMERAL); + }); + + test('prepareStep adds no tail marker for non-Anthropic providers', async () => { + // Force MCP so prepareStep exists even without a tail marker. + const { buildMcpToolRegistry } = await import('@/ee/features/chat/mcp/mcpToolRegistry'); + vi.mocked(buildMcpToolRegistry).mockReturnValueOnce([ + { name: 'mcp_linear__save_issue', description: 'Save an issue', serverName: 'linear' }, + ]); + + const { prepareStep } = await runCreateMessageStream([createUserMessage()], { + promptCacheStrategy: noopStrategy, + }); + expect(prepareStep).toBeTypeOf('function'); + + const result = await prepareStep!({ + steps: [], + stepNumber: 1, + model: {}, + messages: [ + { role: 'user', content: 'q' }, + { role: 'assistant', content: 'a' }, + ], + }); + + // activeTools still managed (MCP), but no message override / marker. + expect(result.messages).toBeUndefined(); + expect(result.activeTools).toContain('tool_request_activation'); + }); + + test('leaves the dynamic system block uncached', async () => { + const { system } = await runCreateMessageStream([createUserMessage()], { + promptCacheStrategy: anthropicStrategy, + selectedRepos: ['github.com/acme/repo'], + }); + + // Static checkpoint + dynamic (per-conversation) block. + expect(system).toHaveLength(2); + expect(system[0].providerOptions?.anthropic?.cacheControl).toEqual(EPHEMERAL); + expect(system[1].providerOptions).toBeUndefined(); + expect(system[1].content).toContain(''); + }); + + test('does not mark the tools block, so mid-run activeTools growth never busts it', async () => { + const { buildMcpToolRegistry } = await import('@/ee/features/chat/mcp/mcpToolRegistry'); + vi.mocked(buildMcpToolRegistry).mockReturnValueOnce([ + { name: 'mcp_linear__save_issue', description: 'Save an issue', serverName: 'linear' }, + ]); + + const { tools } = await runCreateMessageStream([createUserMessage()], { + promptCacheStrategy: anthropicStrategy, + }); + + // The static checkpoint sits on the system block (after the full tools + // section in render order), so no tool definition carries a breakpoint. + for (const tool of Object.values(tools)) { + expect(tool.providerOptions?.anthropic?.cacheControl).toBeUndefined(); + } + }); + + test.each([ + ['non-Anthropic provider', () => getPromptCacheStrategy('openai', true)], + ['caching disabled', () => getPromptCacheStrategy('anthropic', false)], + ])('emits no cache markers for %s (multi-provider regression guard)', async (_label, makeStrategy) => { + const { buildMcpToolRegistry } = await import('@/ee/features/chat/mcp/mcpToolRegistry'); + vi.mocked(buildMcpToolRegistry).mockReturnValueOnce([ + { name: 'mcp_linear__save_issue', description: 'Save an issue', serverName: 'linear' }, + ]); + + const { system, messages, tools } = await runCreateMessageStream([createUserMessage()], { + promptCacheStrategy: makeStrategy(), + selectedRepos: ['github.com/acme/repo'], + }); + + for (const block of system) { + expect(block.providerOptions).toBeUndefined(); + } + for (const message of messages) { + expect(message.providerOptions).toBeUndefined(); + } + for (const tool of Object.values(tools)) { + expect(tool.providerOptions).toBeUndefined(); + } + }); + + test('builds a byte-identical static prompt regardless of repos', async () => { + const first = await runCreateMessageStream([createUserMessage()], { + promptCacheStrategy: anthropicStrategy, + selectedRepos: ['github.com/acme/one'], + }); + const second = await runCreateMessageStream([createUserMessage()], { + promptCacheStrategy: anthropicStrategy, + selectedRepos: ['github.com/acme/two', 'github.com/acme/three'], + }); + + expect(first.system[0].content).toBe(second.system[0].content); + }); +}); diff --git a/packages/web/src/ee/features/chat/agent.ts b/packages/web/src/ee/features/chat/agent.ts index c415e80df..d2f3a4761 100644 --- a/packages/web/src/ee/features/chat/agent.ts +++ b/packages/web/src/ee/features/chat/agent.ts @@ -9,6 +9,7 @@ import { createLogger, env } from "@sourcebot/shared"; import { convertToModelMessages, createUIMessageStream, JSONValue, LanguageModel, ModelMessage, StopCondition, streamText, StreamTextResult, + SystemModelMessage, UIMessageStreamOnFinishCallback, UIMessageStreamOptions, UIMessageStreamWriter, @@ -26,6 +27,7 @@ import { createTools } from "./tools"; import { getConnectedMcpClients } from "@/ee/features/chat/mcp/mcpClientFactory"; import { getMcpTools, McpToolsResult } from "@/ee/features/chat/mcp/mcpToolSets"; import { buildMcpToolRegistry, McpToolRegistryEntry, searchMcpTools } from "@/ee/features/chat/mcp/mcpToolRegistry"; +import { PromptCacheStrategy, mergeProviderOptions, detectPromptCacheBreak, detectUnexpectedCacheMiss } from "./promptCaching"; import { hasEntitlement } from '@/lib/entitlements'; const dedent = _dedent.withOptions({ alignValues: true }); @@ -52,6 +54,7 @@ interface CreateMessageStreamResponseProps { disabledMcpServerIds?: string[]; model: AISDKLanguageModelV3; modelName: string; + promptCacheStrategy: PromptCacheStrategy; onFinish: UIMessageStreamOnFinishCallback; onError: (error: unknown) => string; modelProviderOptions?: Record>; @@ -70,6 +73,7 @@ export const createMessageStream = async ({ disabledMcpServerIds, model, modelName, + promptCacheStrategy, modelProviderOptions, modelTemperature, onFinish, @@ -152,6 +156,7 @@ export const createMessageStream = async ({ const researchStream = await createAgentStream({ model, + promptCacheStrategy, providerOptions: modelProviderOptions, temperature: modelTemperature, inputMessages: messageHistory, @@ -245,6 +250,20 @@ export const createMessageStream = async ({ stepTokenUsage[0].tools.unshift(...toolUsageByToolCallId.values()); } + // Observability only (default off): warn when a continuation step + // reports zero cache reads while the provider supports breakpoints — + // a likely byte-stability regression in the cached prefix. + if (env.SOURCEBOT_CHAT_PROMPT_CACHE_BREAK_DETECTION_ENABLED === 'true') { + steps.forEach((step, stepIndex) => { + detectUnexpectedCacheMiss({ + chatId, + stepIndex, + cacheReadTokens: step.usage.inputTokenDetails?.cacheReadTokens, + supportsBreakpoints: promptCacheStrategy.supportsBreakpoints, + }); + }); + } + writer.write({ type: 'message-metadata', messageMetadata: { @@ -278,6 +297,7 @@ export const createMessageStream = async ({ interface AgentOptions { model: LanguageModel; + promptCacheStrategy: PromptCacheStrategy; providerOptions?: ProviderOptions; temperature?: number; selectedRepos: string[]; @@ -296,6 +316,7 @@ interface AgentOptions { const createAgentStream = async ({ model, + promptCacheStrategy, providerOptions, temperature, inputMessages, @@ -311,6 +332,10 @@ const createAgentStream = async ({ userId, orgId, }: AgentOptions) => { + // Sort repos so the dynamic block is byte-stable + // across a chat's requests (prompt caching). + const sortedRepos = [...selectedRepos].sort((a, b) => a.localeCompare(b)); + // For every file source, resolve the source code so that we can include it in the system prompt. const fileSources = inputSources.filter((source) => source.type === 'file'); const resolvedFileSources = ( @@ -366,6 +391,8 @@ const createAgentStream = async ({ const mcpRegistry = buildMcpToolRegistry(mcpToolSetsObj.tools); const hasMcpTools = mcpRegistry.length > 0; + const staticTtl = env.SOURCEBOT_CHAT_PROMPT_CACHE_STATIC_TTL; + const toolRequestActivation = tool({ description: dedent` Activate an MCP tool by name so it becomes callable on your next step. @@ -390,69 +417,100 @@ const createAgentStream = async ({ }, }); - const systemPrompt = createPrompt({ - repos: selectedRepos, + const { staticPrompt, dynamicPrompt } = createPrompt({ + repos: sortedRepos, files: resolvedFileSources, mcpToolRegistry: mcpRegistry, }); - const builtinTools = createTools({ source: 'sourcebot-ask-agent', selectedRepos }); + const builtinTools = createTools({ source: 'sourcebot-ask-agent', selectedRepos: sortedRepos }); const builtinToolNames = Object.keys(builtinTools); const allTools: Record = { ...builtinTools, ...(hasMcpTools ? { tool_request_activation: toolRequestActivation, ...mcpToolSetsObj.tools } : {}), }; - // Anthropic prompt caching: mark the end of the prompt's static prefix — - // tool definitions, the system prompt (including any resolved file sources), - // and the conversation history — with an ephemeral (5m) cache breakpoint on - // the last input message. Anthropic caches everything up to and including - // this point, so the large prefix is written once (~1.25x) and read back at - // ~0.1x on every subsequent agent step and follow-up turn instead of being - // reprocessed in full. The `anthropic` provider-options namespace is ignored - // by non-Anthropic providers, so this is safe to apply unconditionally. + // Anthropic prompt caching uses two nested breakpoints over one cumulative + // prefix (render order: tools -> system -> messages): + // + // Static checkpoint: the static system block below caches tools + the + // static system instructions. This block is byte-identical across every + // chat and user, so it is a divergence-proof checkpoint a brand-new chat + // can read from instead of re-writing the large static prefix. + // Moving tail: a breakpoint on the last message of each step (applied in + // `prepareStep` below) caches tools + static + dynamic system + the full + // conversation so far. Because it advances to the new tail every step, + // the turn's growing delta (assistant tool calls and their outputs) is + // cached incrementally instead of reprocessed on each later step. + // + // The `anthropic` provider-options namespace is ignored by non-Anthropic + // providers, and a no-op strategy emits no markers at all, so this is safe + // for every provider. When the static prefix falls below the model's minimum + // cacheable size the marker is a harmless no-op. // // Caveat: when MCP tools are lazily activated mid-run via prepareStep, the - // tools section (which precedes everything else in the prefix) grows and - // invalidates the cache for that step; the cache re-warms on subsequent - // steps once the active tool set is stable. - const isPromptCachingEnabled = env.SOURCEBOT_CHAT_PROMPT_CACHING_ENABLED === 'true'; - const messagesWithCachedPrefix: ModelMessage[] = inputMessages.map((message, index) => { - if (!isPromptCachingEnabled || index !== inputMessages.length - 1) { - return message; - } + // tools section grows and invalidates both breakpoints for that step; the + // cache re-warms on subsequent steps once the active tool set is stable. + const staticMarker = promptCacheStrategy.cacheControl({ ttl: staticTtl }); + const systemMessages: SystemModelMessage[] = [ + { role: 'system', content: staticPrompt, providerOptions: staticMarker }, + ]; + if (dynamicPrompt) { + systemMessages.push({ role: 'system', content: dynamicPrompt }); + } - return { - ...message, - providerOptions: { - ...message.providerOptions, - anthropic: { - ...message.providerOptions?.anthropic, - cacheControl: { type: 'ephemeral' }, - }, - }, - }; - }); + // The moving-tail marker (see above), resolved once here. `prepareStep` + // merges it onto the last message's existing providerOptions so sibling + // namespaces (e.g. anthropic.thinking) survive; a no-op strategy leaves it + // undefined and the messages untouched. + const tailMarker = promptCacheStrategy.cacheControl(); + + if (env.SOURCEBOT_CHAT_PROMPT_CACHE_BREAK_DETECTION_ENABLED === 'true') { + detectPromptCacheBreak({ + chatId, + staticPrompt, + toolSignature: Object.entries(builtinTools) + .map(([name, builtinTool]) => `${name}:${builtinTool.description ?? ''}`) + .join('|'), + model: typeof model === 'string' ? model : model.modelId, + staticTtl, + }); + } try { const stream = streamText({ model, providerOptions, - messages: messagesWithCachedPrefix, - system: systemPrompt, + messages: inputMessages, + system: systemMessages, tools: allTools, activeTools: [ ...builtinToolNames, ...(hasMcpTools ? ['tool_request_activation'] : []), ], - prepareStep: hasMcpTools ? ({ steps }) => { + // `prepareStep` runs before every step (including the first). The SDK + // rebuilds the step's messages each time as the original input plus + // its own accumulated response messages. Re-applying the moving tail marker + // to the new last message each step is safe and does not accumulate. + prepareStep: (tailMarker || hasMcpTools) ? ({ steps, messages }) => { + const stepMessages = (tailMarker && messages.length > 0) + ? messages.map((message, index) => + index === messages.length - 1 + ? { ...message, providerOptions: mergeProviderOptions(message.providerOptions, tailMarker) } + : message) + : undefined; + + if (!hasMcpTools) { + return stepMessages ? { messages: stepMessages } : {}; + } + const activated = new Set(); for (const step of steps) { - for (const result of step.toolResults) { - if (!result || result.toolName !== 'tool_request_activation') { + for (const toolResult of step.toolResults) { + if (!toolResult || toolResult.toolName !== 'tool_request_activation') { continue; } - const output = result.output as { results?: Array<{ name: string }> }; + const output = toolResult.output as { results?: Array<{ name: string }> }; for (const { name } of output?.results ?? []) { if (name in mcpToolSetsObj.tools) { activated.add(name); @@ -460,7 +518,9 @@ const createAgentStream = async ({ } } } + return { + ...(stepMessages ? { messages: stepMessages } : {}), activeTools: [ ...builtinToolNames, 'tool_request_activation', @@ -544,8 +604,12 @@ const createPrompt = ({ }[], repos: string[], mcpToolRegistry: McpToolRegistryEntry[], -}) => { - return dedent` +}): { staticPrompt: string; dynamicPrompt: string } => { + // Static prefix: byte-identical across every chat and user. + // It interpolates only module-level constants. Keep it free of any + // per-conversation data — repos, files, and MCP tools live in the dynamic + // block below so their volatility never busts the shared static cache. + const staticPrompt = dedent` You are a powerful agentic AI code assistant built into Sourcebot, the world's best code-intelligence platform. Your job is to help developers understand and navigate their large codebases. @@ -567,39 +631,6 @@ 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. - ${repos.length > 0 ? dedent` - - The user has explicitly selected the following repositories for analysis: - ${repos.map(repo => `- ${repo}`).join('\n')} - - When calling tools that accept a \`repo\` parameter (e.g. \`read_file\`, \`list_commits\`, \`list_tree\`, \`get_diff\`, \`grep\`), use these repository names exactly as listed above, including the full host prefix (e.g. \`github.com/org/repo\`). - - When using \`grep\` to search across ALL selected repositories (e.g. "which repos have X?"), omit the \`repo\` parameter entirely — the tool will automatically search across all selected repositories in a single call. Do NOT call \`grep\` once per repository when a single broad search would suffice. - - ` : ''} - - ${(files && files.length > 0) ? dedent` - - The user has mentioned the following files, which are automatically included for analysis. - - ${files?.map(file => ` - ${addLineNumbers(file.source)} - `).join('\n\n')} - - `: ''} - - ${(mcpToolRegistry.length > 0) ? dedent` - - External MCP tools are available but must first be activated via \`tool_request_activation\`. - - **CRITICAL**: The list below is the complete and authoritative inventory of all tools available to you: - ${mcpToolRegistry.map(e => `- ${e.name}: ${e.description}`).join('\n')} - - **How to use tool_request_activation**: Pass the exact tool name from the list above as the \`tool_to_activate_name\` parameter. Do NOT pass natural language descriptions or sentences. If you need multiple tools, call \`tool_request_activation\` once per tool. - Example: to activate the comment tool, call \`tool_request_activation\` with tool_to_activate_name="mcp_linear__save_comment", NOT tool_to_activate_name="save a comment on an issue". - - ` : ''} - When you have sufficient context, output your answer as a structured markdown response. @@ -630,7 +661,52 @@ const createPrompt = ({ Authentication in Sourcebot is built on NextAuth.js with a session-based approach using JWT tokens and Prisma as the database adapter ${fileReferenceToString({ repo: 'github.com/sourcebot-dev/sourcebot', path: 'auth.ts', range: { startLine: 135, endLine: 140 } })}. The system supports multiple authentication providers and implements organization-based authorization with role-defined permissions. \`\`\` - ` + `; + + // Dynamic block: per-conversation context (selected repos, mentioned files, + // MCP tool registry). + const dynamicSections: string[] = []; + + if (repos.length > 0) { + dynamicSections.push(dedent` + + The user has explicitly selected the following repositories for analysis: + ${repos.map(repo => `- ${repo}`).join('\n')} + + When calling tools that accept a \`repo\` parameter (e.g. \`read_file\`, \`list_commits\`, \`list_tree\`, \`get_diff\`, \`grep\`), use these repository names exactly as listed above, including the full host prefix (e.g. \`github.com/org/repo\`). + + When using \`grep\` to search across ALL selected repositories (e.g. "which repos have X?"), omit the \`repo\` parameter entirely — the tool will automatically search across all selected repositories in a single call. Do NOT call \`grep\` once per repository when a single broad search would suffice. + + `); + } + + if (files && files.length > 0) { + dynamicSections.push(dedent` + + The user has mentioned the following files, which are automatically included for analysis. + + ${files.map(file => ` + ${addLineNumbers(file.source)} + `).join('\n\n')} + + `); + } + + if (mcpToolRegistry.length > 0) { + dynamicSections.push(dedent` + + External MCP tools are available but must first be activated via \`tool_request_activation\`. + + **CRITICAL**: The list below is the complete and authoritative inventory of all tools available to you: + ${mcpToolRegistry.map(e => `- ${e.name}: ${e.description}`).join('\n')} + + **How to use tool_request_activation**: Pass the exact tool name from the list above as the \`tool_to_activate_name\` parameter. Do NOT pass natural language descriptions or sentences. If you need multiple tools, call \`tool_request_activation\` once per tool. + Example: to activate the comment tool, call \`tool_request_activation\` with tool_to_activate_name="mcp_linear__save_comment", NOT tool_to_activate_name="save a comment on an issue". + + `); + } + + return { staticPrompt, dynamicPrompt: dynamicSections.join('\n\n') }; } // If the agent exceeds the step count, then we will stop. diff --git a/packages/web/src/ee/features/chat/mcp/mcpClientFactory.ts b/packages/web/src/ee/features/chat/mcp/mcpClientFactory.ts index 58b813101..6dfdfcbfa 100644 --- a/packages/web/src/ee/features/chat/mcp/mcpClientFactory.ts +++ b/packages/web/src/ee/features/chat/mcp/mcpClientFactory.ts @@ -36,6 +36,11 @@ export async function getConnectedMcpClients(prisma: PrismaClient, userId: strin clientInfo: { not: null }, }, }, + // Deterministic ordering is required for prompt caching: the MCP server + // order flows into both the serialized tool definitions and the + // system-prompt block, and an unordered query would let the + // byte layout drift between requests and bust the cache. + orderBy: { serverId: 'asc' }, select: { serverId: true, tokens: true, diff --git a/packages/web/src/ee/features/chat/mcp/mcpToolRegistry.ts b/packages/web/src/ee/features/chat/mcp/mcpToolRegistry.ts index 431710e9e..ead4f2be0 100644 --- a/packages/web/src/ee/features/chat/mcp/mcpToolRegistry.ts +++ b/packages/web/src/ee/features/chat/mcp/mcpToolRegistry.ts @@ -49,15 +49,19 @@ function expandTokens(tokens: string[]): string[] { } export function buildMcpToolRegistry(tools: McpToolRecord): McpToolRegistryEntry[] { - return Object.entries(tools).map(([name, tool]) => { - const match = name.match(/^mcp_(.+?)__/); - const serverName = match ? match[1] : ''; - return { - name, - description: tool.description ?? '', - serverName, - }; - }); + return Object.entries(tools) + // Sort by tool name so the system-prompt block is byte-stable + // across requests regardless of upstream iteration order (prompt caching). + .sort(([nameA], [nameB]) => nameA.localeCompare(nameB)) + .map(([name, tool]) => { + const match = name.match(/^mcp_(.+?)__/); + const serverName = match ? match[1] : ''; + return { + name, + description: tool.description ?? '', + serverName, + }; + }); } export function searchMcpTools( diff --git a/packages/web/src/ee/features/chat/promptCaching.test.ts b/packages/web/src/ee/features/chat/promptCaching.test.ts new file mode 100644 index 000000000..69dd96ef4 --- /dev/null +++ b/packages/web/src/ee/features/chat/promptCaching.test.ts @@ -0,0 +1,93 @@ +import { describe, expect, test, vi } from 'vitest'; + +vi.mock('@sourcebot/shared', () => ({ + createLogger: () => ({ + debug: vi.fn(), + error: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + }), +})); + +const { getPromptCacheStrategy, mergeProviderOptions } = await import('./promptCaching'); + +describe('getPromptCacheStrategy', () => { + test.each(['anthropic', 'google-vertex-anthropic'] as const)( + 'returns a 5m ephemeral marker (no ttl) for %s when enabled', + (provider) => { + const strategy = getPromptCacheStrategy(provider, true); + expect(strategy.supportsBreakpoints).toBe(true); + expect(strategy.cacheControl()).toEqual({ + anthropic: { cacheControl: { type: 'ephemeral' } }, + }); + }, + ); + + test('includes ttl only for the 1h bucket', () => { + const strategy = getPromptCacheStrategy('anthropic', true); + expect(strategy.cacheControl({ ttl: '1h' })).toEqual({ + anthropic: { cacheControl: { type: 'ephemeral', ttl: '1h' } }, + }); + // 5m is the default — ttl is omitted so the request matches Anthropic's + // default and stays byte-stable. + expect(strategy.cacheControl({ ttl: '5m' })).toEqual({ + anthropic: { cacheControl: { type: 'ephemeral' } }, + }); + }); + + test.each([ + 'amazon-bedrock', + 'azure', + 'deepseek', + 'google-generative-ai', + 'google-vertex', + 'mistral', + 'openai', + 'openai-compatible', + 'openrouter', + 'xai', + ] as const)('is a no-op for non-Anthropic provider %s', (provider) => { + const strategy = getPromptCacheStrategy(provider, true); + expect(strategy.supportsBreakpoints).toBe(false); + expect(strategy.cacheControl()).toBeUndefined(); + expect(strategy.cacheControl({ ttl: '1h' })).toBeUndefined(); + }); + + test('is a no-op when caching is disabled, even for Anthropic', () => { + const strategy = getPromptCacheStrategy('anthropic', false); + expect(strategy.supportsBreakpoints).toBe(false); + expect(strategy.cacheControl()).toBeUndefined(); + }); +}); + +describe('mergeProviderOptions', () => { + test('returns the existing options unchanged when there is no marker', () => { + const existing = { anthropic: { thinking: { type: 'enabled' } } }; + expect(mergeProviderOptions(existing, undefined)).toBe(existing); + }); + + test('returns the marker when there are no existing options', () => { + const marker = { anthropic: { cacheControl: { type: 'ephemeral' } } }; + expect(mergeProviderOptions(undefined, marker)).toBe(marker); + }); + + test('deep-merges the cacheControl while preserving sibling options in the same namespace', () => { + const existing = { anthropic: { thinking: { type: 'enabled', budgetTokens: 1024 } } }; + const marker = { anthropic: { cacheControl: { type: 'ephemeral' } } }; + expect(mergeProviderOptions(existing, marker)).toEqual({ + anthropic: { + thinking: { type: 'enabled', budgetTokens: 1024 }, + cacheControl: { type: 'ephemeral' }, + }, + }); + }); + + test('preserves sibling provider namespaces', () => { + const existing = { openai: { someOption: true } }; + const marker = { anthropic: { cacheControl: { type: 'ephemeral' } } }; + expect(mergeProviderOptions(existing, marker)).toEqual({ + openai: { someOption: true }, + anthropic: { cacheControl: { type: 'ephemeral' } }, + }); + }); +}); diff --git a/packages/web/src/ee/features/chat/promptCaching.ts b/packages/web/src/ee/features/chat/promptCaching.ts new file mode 100644 index 000000000..828f48d4b --- /dev/null +++ b/packages/web/src/ee/features/chat/promptCaching.ts @@ -0,0 +1,184 @@ +import { createHash } from "crypto"; +import { ProviderOptions } from "@ai-sdk/provider-utils"; +import { createLogger } from "@sourcebot/shared"; +import { LanguageModelProvider } from "@/features/chat/types"; + +// @note: Prompt-cache plumbing for the Ask agent — the single place that knows +// about provider-specific cache-control shapes. The agent loop stays +// provider-agnostic by asking the strategy for a `providerOptions` blob at each +// breakpoint and merging whatever it gets back (possibly `undefined`). + +const logger = createLogger('prompt-caching'); + +export type CacheTtl = '5m' | '1h'; + +export interface PromptCacheStrategy { + /** + * Whether the resolved provider supports explicit cache breakpoints. When + * false, every `cacheControl()` call is a no-op. + */ + readonly supportsBreakpoints: boolean; + /** + * Returns the `providerOptions` blob to merge onto a system block, tool + * definition, or message to mark a cache breakpoint — or `undefined` when + * caching is disabled or the provider does not support explicit breakpoints. + */ + cacheControl: (opts?: { ttl?: CacheTtl }) => ProviderOptions | undefined; +} + +const NOOP_STRATEGY: PromptCacheStrategy = { + supportsBreakpoints: false, + cacheControl: () => undefined, +}; + +// Providers whose `@ai-sdk/*` package consumes the `anthropic` providerOptions +// namespace and honors ephemeral `cacheControl` breakpoints. `google-vertex-anthropic` +// reuses the same namespace; if a given SDK version ignores it the marker is a +// harmless no-op. Bedrock uses a different shape (`cachePoint`) and is intentionally +// not covered here - adding it later is a single new branch. +const ANTHROPIC_FAMILY_PROVIDERS: ReadonlySet = new Set([ + 'anthropic', + 'google-vertex-anthropic', +]); + +const anthropicCacheControl = ({ ttl }: { ttl?: CacheTtl } = {}): ProviderOptions => ({ + anthropic: { + cacheControl: { + type: 'ephemeral', + // 5m is the Anthropic default; only emit `ttl` for the 1h bucket. + ...(ttl === '1h' ? { ttl: '1h' } : {}), + }, + }, +}); + +/** + * Resolves how (and whether) to emit prompt-cache breakpoints for a provider. + */ +export const getPromptCacheStrategy = ( + provider: LanguageModelProvider, + enabled: boolean, +): PromptCacheStrategy => { + if (!enabled || !ANTHROPIC_FAMILY_PROVIDERS.has(provider)) { + return NOOP_STRATEGY; + } + + return { + supportsBreakpoints: true, + cacheControl: (opts) => anthropicCacheControl(opts), + }; +}; + +/** + * Deep-merges a cache-control marker's `providerOptions` onto any existing + * `providerOptions`, preserving sibling provider namespaces and other options + * within the same namespace (e.g. `anthropic.thinking`). Returns the original + * value unchanged when there is no marker to apply. + */ +export const mergeProviderOptions = ( + existing: ProviderOptions | undefined, + marker: ProviderOptions | undefined, +): ProviderOptions | undefined => { + if (!marker) { + return existing; + } + if (!existing) { + return marker; + } + + const merged: ProviderOptions = { ...existing }; + for (const [namespace, options] of Object.entries(marker)) { + merged[namespace] = { + ...(existing[namespace] ?? {}), + ...options, + }; + } + return merged; +}; + +interface CacheBreakSnapshot { + signature: string; + requestCount: number; +} + +// Keyed by chatId, in-memory, observability-only. Entries are never removed on chat +// end, so the map is bounded by a FIFO cap (oldest insertion evicted first, below); +const MAX_CACHE_BREAK_SNAPSHOTS = 10_000; +const cacheBreakSnapshots = new Map(); + +const hashString = (input: string): string => + createHash('sha256').update(input).digest('hex').slice(0, 16); + +/** + * Records the cache-relevant prefix signature for a chat and logs a warning when + * it changes between requests in a way that would invalidate the static prompt + * cache. Observability only — never affects request behavior. + */ +export const detectPromptCacheBreak = ({ + chatId, + staticPrompt, + toolSignature, + model, + staticTtl, +}: { + chatId: string; + staticPrompt: string; + toolSignature: string; + model: string; + staticTtl: CacheTtl; +}): void => { + try { + const signature = hashString([ + model, + staticTtl, + hashString(staticPrompt), + hashString(toolSignature), + ].join('|')); + + const prev = cacheBreakSnapshots.get(chatId); + const requestCount = (prev?.requestCount ?? 0) + 1; + cacheBreakSnapshots.set(chatId, { signature, requestCount }); + + if (cacheBreakSnapshots.size > MAX_CACHE_BREAK_SNAPSHOTS) { + const oldestKey = cacheBreakSnapshots.keys().next().value; + if (oldestKey !== undefined) { + cacheBreakSnapshots.delete(oldestKey); + } + } + + if (prev && prev.signature !== signature) { + logger.warn( + `Prompt cache break detected for chat ${chatId} (request #${requestCount}): the ` + + `static system prompt, built-in tool definitions, model, or TTL changed between ` + + `requests, invalidating the static cache prefix.`, + ); + } + } catch (error) { + logger.debug(`detectPromptCacheBreak failed: ${error}`); + } +}; + +/** + * Logs a warning when a non-first agent step reports zero cache-read tokens while + * caching is supported — a likely byte-stability regression. Observability only. + */ +export const detectUnexpectedCacheMiss = ({ + chatId, + stepIndex, + cacheReadTokens, + supportsBreakpoints, +}: { + chatId: string; + stepIndex: number; + cacheReadTokens: number | undefined; + supportsBreakpoints: boolean; +}): void => { + if (!supportsBreakpoints || stepIndex === 0) { + return; + } + if ((cacheReadTokens ?? 0) === 0) { + logger.warn( + `Prompt cache miss for chat ${chatId} at step ${stepIndex}: cacheReadTokens=0 on a ` + + `continuation step where a cache hit was expected (likely a byte-stability regression).`, + ); + } +}; diff --git a/packages/web/src/ee/features/mcp/askCodebase.ts b/packages/web/src/ee/features/mcp/askCodebase.ts index 9eaaf3fc2..4b7cfb7b0 100644 --- a/packages/web/src/ee/features/mcp/askCodebase.ts +++ b/packages/web/src/ee/features/mcp/askCodebase.ts @@ -15,6 +15,7 @@ import { InferUIMessageChunk, UIDataTypes, UIMessage, UITools } from "ai"; import { captureEvent } from "@/lib/posthog"; import { createAudit } from "@/ee/features/audit/audit"; import { createMessageStream } from "@/ee/features/chat/agent"; +import { getPromptCacheStrategy } from "@/ee/features/chat/promptCaching"; const logger = createLogger('ask-codebase-api'); @@ -84,6 +85,12 @@ export const askCodebase = (params: AskCodebaseParams): Promise {