Re-implement token caching for Vercel AI SDK usage with Anthropic provider (fixes #58)

Author: bhouston

.changeset/implement-token-caching.md

@@ -0,0 +1,5 @@
+---
+"mycoder-agent": patch
+---
+
+Re-implemented token caching for Vercel AI SDK usage with Anthropic provider to reduce token consumption during repeated API calls.

packages/agent/src/core/toolAgent.ts

@@ -8,6 +8,7 @@ import {
   ToolResultPart,
   ToolSet,
   tool as makeTool,
+  Message
 } from 'ai';
 import chalk from 'chalk';
 
@@ -191,62 +192,39 @@ async function executeTools(
   };
 }
 
-/*
-// a function that takes a list of messages and returns a list of messages but with the last message having a cache_control of ephemeral
-function addCacheControlToTools<T>(messages: T[]): T[] {
-  return messages.map((m, i) => ({
-    ...m,
-    ...(i === messages.length - 1
-      ? { cache_control: { type: 'ephemeral' } }
-      : {}),
-  }));
-}
-
-function addCacheControlToContentBlocks(
-  content: ContentBlockParam[],
-): ContentBlockParam[] {
-  return content.map((c, i) => {
-    if (i === content.length - 1) {
-      if (
-        c.type === 'text' ||
-        c.type === 'document' ||
-        c.type === 'image' ||
-        c.type === 'tool_use' ||
-        c.type === 'tool_result' ||
-        c.type === 'thinking' ||
-        c.type === 'redacted_thinking'
-      ) {
-        return { ...c, cache_control: { type: 'ephemeral' } };
-      }
+/**
+ * Adds cache control to the messages for token caching with the Vercel AI SDK
+ * This marks the last message as ephemeral to prevent it from being cached
+ * which helps reduce token usage when making multiple API calls
+ */
+function addCacheControlToMessages(messages: CoreMessage[]): CoreMessage[] {
+  if (messages.length === 0) return messages;
+  
+  // Create a deep copy of the messages array to avoid mutating the original
+  const result = JSON.parse(JSON.stringify(messages)) as CoreMessage[];
+  
+  // Get the last message
+  const lastMessage = result[result.length - 1];
+  
+  // Add cache_control to the last message
+  if (lastMessage) {
+    // For the Vercel AI SDK, we need to add the cache_control property
+    // to indicate that this message should not be cached
+    (lastMessage as any).cache_control = { type: 'ephemeral' };
+    
+    // If the message has content array, mark each content item as ephemeral
+    if (Array.isArray(lastMessage.content)) {
+      lastMessage.content = lastMessage.content.map(item => {
+        return {
+          ...item,
+          cache_control: { type: 'ephemeral' }
+        };
+      });
     }
-    return c;
-  });
+  }
+  
+  return result;
 }
-function addCacheControlToMessages(
-  messages: Anthropic.Messages.MessageParam[],
-): Anthropic.Messages.MessageParam[] {
-  return messages.map((m, i) => {
-    if (typeof m.content === 'string') {
-      return {
-        ...m,
-        content: [
-          {
-            type: 'text',
-            text: m.content,
-            cache_control: { type: 'ephemeral' },
-          },
-        ] as ContentBlockParam[],
-      };
-    }
-    return {
-      ...m,
-      content:
-        i >= messages.length - 2
-          ? addCacheControlToContentBlocks(m.content)
-          : m.content,
-    };
-  });
-}*/
 
 export const toolAgent = async (
   initialPrompt: string,
@@ -293,10 +271,13 @@ export const toolAgent = async (
         parameters: tool.parameters,
       });
     });
+    // Apply cache control to messages for token caching
+    const messagesWithCacheControl = addCacheControlToMessages(messages);
+    
     const generateTextProps = {
       model: config.model,
       temperature: config.temperature,
-      messages,
+      messages: messagesWithCacheControl,
       system: systemPrompt,
       tools: toolSet,
     };