feat: implement prompt caching for Anthropic

sdirix · sdirix · commit ca385edc1898 · 2025-06-05T15:08:08.000+02:00
This is a backport of #15731 for Theia 1.61.x Mark system prompts, tools and messages for caching in the Anthropic LanguageModel. For typical agent workflows like "@coder" this reduces token costs by a large factor and reduces the pressure on the rate limit. Currently caching is enabled by default for all Anthropic models. Also updates token tracking functionality to be aware of token caching and updates the Anthropic SDK.
diff --git a/package-lock.json b/package-lock.json
diff --git a/packages/ai-anthropic/package.json b/packages/ai-anthropic/package.json
@@ -3,7 +3,7 @@
   "version": "1.61.0",
   "description": "Theia - Anthropic Integration",
   "dependencies": {
-    "@anthropic-ai/sdk": "^0.39.0",
+    "@anthropic-ai/sdk": "^0.52.0",
     "@theia/ai-core": "1.61.0",
     "@theia/core": "1.61.0"
   },
diff --git a/packages/ai-anthropic/src/browser/anthropic-frontend-application-contribution.ts b/packages/ai-anthropic/src/browser/anthropic-frontend-application-contribution.ts
@@ -79,7 +79,8 @@ export class AnthropicFrontendApplicationContribution implements FrontendApplica
             id: id,
             model: modelId,
             apiKey: true,
-            enableStreaming: true
+            enableStreaming: true,
+            useCaching: true
         };
 
         if (maxTokens !== undefined) {
diff --git a/packages/ai-anthropic/src/common/anthropic-language-models-manager.ts b/packages/ai-anthropic/src/common/anthropic-language-models-manager.ts
@@ -32,6 +32,10 @@ export interface AnthropicModelDescription {
      * Indicate whether the streaming API shall be used.
      */
     enableStreaming: boolean;
+    /**
+     * Indicate whether the model supports prompt caching.
+     */
+    useCaching: boolean;
     /**
      * Maximum number of tokens to generate. Default is 4096.
      */
diff --git a/packages/ai-anthropic/src/node/anthropic-language-model.ts b/packages/ai-anthropic/src/node/anthropic-language-model.ts
@@ -28,7 +28,7 @@ import {
 } from '@theia/ai-core';
 import { CancellationToken, isArray } from '@theia/core';
 import { Anthropic } from '@anthropic-ai/sdk';
-import { Message, MessageParam } from '@anthropic-ai/sdk/resources';
+import type { Message, MessageParam } from '@anthropic-ai/sdk/resources';
 
 export const DEFAULT_MAX_TOKENS = 4096;
 
@@ -41,7 +41,7 @@ interface ToolCallback {
 
 const createMessageContent = (message: LanguageModelMessage): MessageParam['content'] => {
     if (LanguageModelMessage.isTextMessage(message)) {
-        return message.text;
+        return [{ type: 'text', text: message.text }];
     } else if (LanguageModelMessage.isThinkingMessage(message)) {
         return [{ signature: message.signature, thinking: message.thinking, type: 'thinking' }];
     } else if (LanguageModelMessage.isToolUseMessage(message)) {
@@ -52,17 +52,27 @@ const createMessageContent = (message: LanguageModelMessage): MessageParam['cont
     throw new Error(`Unknown message type:'${JSON.stringify(message)}'`);
 };
 
+type NonThinkingParam = Exclude<Anthropic.Messages.ContentBlockParam, Anthropic.Messages.ThinkingBlockParam | Anthropic.Messages.RedactedThinkingBlockParam>;
+function isNonThinkingParam(
+    content: Anthropic.Messages.ContentBlockParam
+): content is NonThinkingParam {
+    return content.type !== 'thinking' && content.type !== 'redacted_thinking';
+}
+
 /**
  * Transforms Theia language model messages to Anthropic API format
  * @param messages Array of LanguageModelRequestMessage to transform
  * @returns Object containing transformed messages and optional system message
  */
 function transformToAnthropicParams(
-    messages: readonly LanguageModelMessage[]
-): { messages: MessageParam[]; systemMessage?: string } {
+    messages: readonly LanguageModelMessage[],
+    addCacheControl: boolean = true
+): { messages: MessageParam[]; systemMessage?: Anthropic.Messages.TextBlockParam[] } {
     // Extract the system message (if any), as it is a separate parameter in the Anthropic API.
     const systemMessageObj = messages.find(message => message.actor === 'system');
-    const systemMessage = systemMessageObj && LanguageModelMessage.isTextMessage(systemMessageObj) && systemMessageObj.text || undefined;
+    const systemMessageText = systemMessageObj && LanguageModelMessage.isTextMessage(systemMessageObj) && systemMessageObj.text || undefined;
+    const systemMessage: Anthropic.Messages.TextBlockParam[] | undefined =
+        systemMessageText ? [{ type: 'text', text: systemMessageText, cache_control: addCacheControl ? { type: 'ephemeral' } : undefined }] : undefined;
 
     const convertedMessages = messages
         .filter(message => message.actor !== 'system')
@@ -77,6 +87,35 @@ function transformToAnthropicParams(
     };
 }
 
+/**
+ * If possible adds a cache control to the last message in the conversation.
+ * This is used to enable incremental caching of the conversation.
+ * @param messages The messages to process
+ * @returns A new messages array with the last message adapted to include cache control. If no cache control can be added, the original messages are returned.
+ * In any case, the original messages are not modified
+ */
+function addCacheControlToLastMessage(messages: Anthropic.Messages.MessageParam[]): Anthropic.Messages.MessageParam[] {
+    const clonedMessages = [...messages];
+    const latestMessage = clonedMessages.pop();
+    if (latestMessage) {
+        let content: NonThinkingParam | undefined = undefined;
+        if (typeof latestMessage.content === 'string') {
+            content = { type: 'text', text: latestMessage.content };
+        } else if (Array.isArray(latestMessage.content)) {
+            // we can't set cache control on thinking messages, so we only set it on the last non-thinking block
+            const filteredContent = latestMessage.content.filter(isNonThinkingParam);
+            if (filteredContent.length) {
+                content = filteredContent[filteredContent.length - 1];
+            }
+        }
+        if (content) {
+            const cachedContent: NonThinkingParam = { ...content, cache_control: { type: 'ephemeral' } };
+            return [...clonedMessages, { ...latestMessage, content: [cachedContent] }];
+        }
+    }
+    return messages;
+}
+
 export const AnthropicModelIdentifier = Symbol('AnthropicModelIdentifier');
 
 /**
@@ -102,6 +141,7 @@ export class AnthropicModel implements LanguageModel {
         public readonly id: string,
         public model: string,
         public enableStreaming: boolean,
+        public useCaching: boolean,
         public apiKey: () => string | undefined,
         public maxTokens: number = DEFAULT_MAX_TOKENS,
         protected readonly tokenUsageService?: TokenUsageService
@@ -153,11 +193,18 @@ export class AnthropicModel implements LanguageModel {
         toolMessages?: readonly Anthropic.Messages.MessageParam[]
     ): Promise<LanguageModelStreamResponse> {
         const settings = this.getSettings(request);
-        const { messages, systemMessage } = transformToAnthropicParams(request.messages);
+        const { messages, systemMessage } = transformToAnthropicParams(request.messages, this.useCaching);
+
+        let anthropicMessages = [...messages, ...(toolMessages ?? [])];
+
+        if (this.useCaching && anthropicMessages.length) {
+            anthropicMessages = addCacheControlToLastMessage(anthropicMessages);
+        }
+
         const tools = this.createTools(request);
         const params: Anthropic.MessageCreateParams = {
             max_tokens: this.maxTokens,
-            messages: [...messages, ...(toolMessages ?? [])],
+            messages: anthropicMessages,
             tools,
             tool_choice: tools ? { type: 'auto' } : undefined,
             model: this.model,
@@ -231,6 +278,8 @@ export class AnthropicModel implements LanguageModel {
                                 const tokenUsageParams: TokenUsageParams = {
                                     inputTokens: currentMessage.usage.input_tokens,
                                     outputTokens: currentMessage.usage.output_tokens,
+                                    cachedInputTokens: currentMessage.usage.cache_creation_input_tokens || undefined,
+                                    readCachedInputTokens: currentMessage.usage.cache_read_input_tokens || undefined,
                                     requestId: request.requestId
                                 };
                                 await that.tokenUsageService.recordTokenUsage(that.id, tokenUsageParams);
@@ -285,15 +334,21 @@ export class AnthropicModel implements LanguageModel {
         return { stream: asyncIterator };
     }
 
-    private createTools(request: LanguageModelRequest): Anthropic.Messages.Tool[] | undefined {
+    protected createTools(request: LanguageModelRequest): Anthropic.Messages.Tool[] | undefined {
         if (request.tools?.length === 0) {
             return undefined;
         }
-        return request.tools?.map(tool => ({
+        const tools = request.tools?.map(tool => ({
             name: tool.name,
             description: tool.description,
             input_schema: tool.parameters
         } as Anthropic.Messages.Tool));
+        if (this.useCaching) {
+            if (tools?.length) {
+                tools[tools.length - 1].cache_control = { type: 'ephemeral' };
+            }
+        }
+        return tools;
     }
 
     protected async handleNonStreamingRequest(
diff --git a/packages/ai-anthropic/src/node/anthropic-language-models-manager-impl.ts b/packages/ai-anthropic/src/node/anthropic-language-models-manager-impl.ts
@@ -66,6 +66,7 @@ export class AnthropicLanguageModelsManagerImpl implements AnthropicLanguageMode
                         modelDescription.id,
                         modelDescription.model,
                         modelDescription.enableStreaming,
+                        modelDescription.useCaching,
                         apiKeyProvider,
                         modelDescription.maxTokens,
                         this.tokenUsageService
diff --git a/packages/ai-core/src/browser/token-usage-frontend-service-impl.ts b/packages/ai-core/src/browser/token-usage-frontend-service-impl.ts
@@ -76,6 +76,8 @@ export class TokenUsageFrontendServiceImpl implements TokenUsageFrontendService
         const modelMap = new Map<string, {
             inputTokens: number;
             outputTokens: number;
+            cachedInputTokens: number;
+            readCachedInputTokens: number;
             lastUsed?: Date;
         }>();
 
@@ -87,6 +89,16 @@ export class TokenUsageFrontendServiceImpl implements TokenUsageFrontendService
                 existing.inputTokens += usage.inputTokens;
                 existing.outputTokens += usage.outputTokens;
 
+                // Add cached tokens if they exist
+                if (usage.cachedInputTokens !== undefined) {
+                    existing.cachedInputTokens += usage.cachedInputTokens;
+                }
+
+                // Add read cached tokens if they exist
+                if (usage.readCachedInputTokens !== undefined) {
+                    existing.readCachedInputTokens += usage.readCachedInputTokens;
+                }
+
                 // Update last used if this usage is more recent
                 if (!existing.lastUsed || (usage.timestamp && usage.timestamp > existing.lastUsed)) {
                     existing.lastUsed = usage.timestamp;
@@ -95,6 +107,8 @@ export class TokenUsageFrontendServiceImpl implements TokenUsageFrontendService
                 modelMap.set(usage.model, {
                     inputTokens: usage.inputTokens,
                     outputTokens: usage.outputTokens,
+                    cachedInputTokens: usage.cachedInputTokens || 0,
+                    readCachedInputTokens: usage.readCachedInputTokens || 0,
                     lastUsed: usage.timestamp
                 });
             }
@@ -104,12 +118,23 @@ export class TokenUsageFrontendServiceImpl implements TokenUsageFrontendService
         const result: ModelTokenUsageData[] = [];
 
         for (const [modelId, data] of modelMap.entries()) {
-            result.push({
+            const modelData: ModelTokenUsageData = {
                 modelId,
                 inputTokens: data.inputTokens,
                 outputTokens: data.outputTokens,
                 lastUsed: data.lastUsed
-            });
+            };
+
+            // Only include cache-related fields if they have non-zero values
+            if (data.cachedInputTokens > 0) {
+                modelData.cachedInputTokens = data.cachedInputTokens;
+            }
+
+            if (data.readCachedInputTokens > 0) {
+                modelData.readCachedInputTokens = data.readCachedInputTokens;
+            }
+
+            result.push(modelData);
         }
 
         return result;
diff --git a/packages/ai-core/src/browser/token-usage-frontend-service.ts b/packages/ai-core/src/browser/token-usage-frontend-service.ts
@@ -26,6 +26,10 @@ export interface ModelTokenUsageData {
     inputTokens: number;
     /** Number of output tokens used */
     outputTokens: number;
+    /** Number of input tokens written to cache */
+    cachedInputTokens?: number;
+    /** Number of input tokens read from cache */
+    readCachedInputTokens?: number;
     /** Date when the model was last used */
     lastUsed?: Date;
 }
diff --git a/packages/ai-core/src/common/token-usage-service.ts b/packages/ai-core/src/common/token-usage-service.ts
@@ -23,6 +23,10 @@ export interface TokenUsage {
     inputTokens: number;
     /** The output token count */
     outputTokens: number;
+    /** Input tokens written to cache */
+    cachedInputTokens?: number;
+    /** Input tokens read from cache */
+    readCachedInputTokens?: number;
     /** The model identifier */
     model: string;
     /** The timestamp of when the tokens were used */
@@ -36,6 +40,10 @@ export interface TokenUsageParams {
     inputTokens: number;
     /** The output token count */
     outputTokens: number;
+    /** Input tokens placed in cache */
+    cachedInputTokens?: number;
+    /** Input tokens read from cache */
+    readCachedInputTokens?: number;
     /** Request identifier */
     requestId: string;
 }
diff --git a/packages/ai-core/src/node/token-usage-service-impl.ts b/packages/ai-core/src/node/token-usage-service-impl.ts
@@ -41,6 +41,8 @@ export class TokenUsageServiceImpl implements TokenUsageService {
     async recordTokenUsage(model: string, params: TokenUsageParams): Promise<void> {
         const usage: TokenUsage = {
             inputTokens: params.inputTokens,
+            cachedInputTokens: params.cachedInputTokens,
+            readCachedInputTokens: params.readCachedInputTokens,
             outputTokens: params.outputTokens,
             model,
             timestamp: new Date(),
@@ -50,7 +52,23 @@ export class TokenUsageServiceImpl implements TokenUsageService {
         this.tokenUsages.push(usage);
         this.client?.notifyTokenUsage(usage);
 
-        console.log(`Input Tokens: ${params.inputTokens}; Output Tokens: ${params.outputTokens}; Model: ${model}${params.requestId ? `; RequestId: ${params.requestId}` : ''}`);
+        let logMessage = `Input Tokens: ${params.inputTokens};`;
+
+        if (params.cachedInputTokens) {
+            logMessage += ` Input Tokens written to cache: ${params.cachedInputTokens};`;
+        }
+
+        if (params.readCachedInputTokens) {
+            logMessage += ` Input Tokens read from cache: ${params.readCachedInputTokens};`;
+        }
+
+        logMessage += ` Output Tokens: ${params.outputTokens}; Model: ${model};`;
+
+        if (params.requestId) {
+            logMessage += `; RequestId: ${params.requestId}`;
+        }
+
+        console.debug(logMessage);
         // For now we just store in memory
         // In the future, this could be persisted to disk, a database, or sent to a service
         return Promise.resolve();
diff --git a/packages/ai-ide/src/browser/ai-configuration/token-usage-configuration-widget.tsx b/packages/ai-ide/src/browser/ai-configuration/token-usage-configuration-widget.tsx
