Improve token caching.

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)
 

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'),

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,

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<string, { providerOptions?: ProviderOptions }>;
+    prepareStep?: FakePrepareStep;
+}
+
+const runCreateMessageStream = async (
+    messages: SBChatMessage[],
+    opts: {
+        promptCacheStrategy?: PromptCacheStrategy;
+        selectedRepos?: string[];
+    } = {},
+): Promise<StreamTextArgs> => {
     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<typeof createMessageStream>[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('<selected_repositories>');
+    });
+
+    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);
+    });
+});