move cache control into convertCbToModelMessages

Author: charleslien

backend/src/llm-apis/vercel-ai-sdk/ai-sdk.ts

@@ -75,6 +75,7 @@ export const promptAiSdkStream = async function* (
     agentId?: string
     maxRetries?: number
     onCostCalculated?: (credits: number) => Promise<void>
+    includeCacheControl?: boolean
   } & Omit<Parameters<typeof streamText>[0], 'model'>,
 ) {
   if (
@@ -233,6 +234,7 @@ export const promptAiSdk = async function (
     chargeUser?: boolean
     agentId?: string
     onCostCalculated?: (credits: number) => Promise<void>
+    includeCacheControl?: boolean
   } & Omit<Parameters<typeof generateText>[0], 'model'>,
 ): Promise<string> {
   if (
@@ -306,6 +308,7 @@ export const promptAiSdkStructured = async function <T>(options: {
   chargeUser?: boolean
   agentId?: string
   onCostCalculated?: (credits: number) => Promise<void>
+  includeCacheControl?: boolean
 }): Promise<T> {
   if (
     !checkLiveUserInput(

backend/src/prompt-agent-stream.ts

@@ -14,6 +14,7 @@ export const getAgentStreamFromTemplate = (params: {
   userId: string | undefined
   onCostCalculated?: (credits: number) => Promise<void>
   agentId?: string
+  includeCacheControl?: boolean
 
   template: AgentTemplate
 }) => {
@@ -24,6 +25,7 @@ export const getAgentStreamFromTemplate = (params: {
     userId,
     onCostCalculated,
     agentId,
+    includeCacheControl,
     template,
   } = params
 
@@ -44,6 +46,7 @@ export const getAgentStreamFromTemplate = (params: {
       userId,
       maxOutputTokens: 32_000,
       onCostCalculated,
+      includeCacheControl,
       agentId,
     }
 

backend/src/run-agent-step.ts

@@ -269,6 +269,7 @@ export const runAgentStep = async (
         )
       }
     },
+    includeCacheControl: supportsCacheControl(agentTemplate.model),
   })
 
   const iterationNum = agentState.messageHistory.length
@@ -287,7 +288,6 @@ export const runAgentStep = async (
   const agentMessages = getMessagesSubset(
     agentState.messageHistory,
     systemTokens,
-    supportsCacheControl(agentTemplate.model),
   )
 
   logger.debug(

backend/src/util/messages.ts

@@ -1,7 +1,6 @@
 import { AssertionError } from 'assert'
 
 import { buildArray } from '@codebuff/common/util/array'
-import { withCacheControl } from '@codebuff/common/util/messages'
 import { closeXml } from '@codebuff/common/util/xml'
 
 import { logger } from './logger'
@@ -250,7 +249,6 @@ export function trimMessagesToFitTokenLimit(
 export function getMessagesSubset(
   messages: CodebuffMessage[],
   otherTokens: number,
-  includeCacheControl: boolean = true,
 ): CodebuffMessage[] {
   const messagesSubset = trimMessagesToFitTokenLimit(messages, otherTokens)
 
@@ -271,23 +269,8 @@ export function getMessagesSubset(
       },
       'No last message found in messagesSubset!',
     )
-    return messagesSubset
   }
 
-  if (!includeCacheControl) {
-    return messagesSubset
-  }
-
-  // add cache control to specific messages
-  for (const ttl of ['agentStep', 'userPrompt'] as const) {
-    const index = messagesSubset.findIndex((m) => m.timeToLive === ttl)
-    if (index <= 0) {
-      continue
-    }
-    messagesSubset[index - 1] = withCacheControl(messagesSubset[index - 1])
-  }
-  messagesSubset[messagesSubset.length - 1] = withCacheControl(lastMessage)
-
   return messagesSubset
 }
 

common/src/types/messages/codebuff-message.ts

@@ -10,44 +10,60 @@ import {
 } from './content-part'
 import { providerMetadataSchema } from './provider-metadata'
 
-export const systemCodebuffMessageSchema = z.object({
-  role: z.literal('system'),
-  content: z.string(),
+const auxiliaryDataSchema = z.object({
+  providerOptions: providerMetadataSchema.optional(),
+  timeToLive: z
+    .union([z.literal('agentStep'), z.literal('userPrompt')])
+    .optional(),
+  keepDuringTruncation: z.boolean().optional(),
 })
+
+export const systemCodebuffMessageSchema = z
+  .object({
+    role: z.literal('system'),
+    content: z.string(),
+  })
+  .and(auxiliaryDataSchema)
 export type SystemCodebuffMessage = z.infer<typeof systemCodebuffMessageSchema>
 
-export const userCodebuffMessageSchema = z.object({
-  role: z.literal('user'),
-  content: z.union([
-    z.string(),
-    z.union([textPartSchema, imagePartSchema, filePartSchema]).array(),
-  ]),
-})
+export const userCodebuffMessageSchema = z
+  .object({
+    role: z.literal('user'),
+    content: z.union([
+      z.string(),
+      z.union([textPartSchema, imagePartSchema, filePartSchema]).array(),
+    ]),
+  })
+  .and(auxiliaryDataSchema)
 export type UserCodebuffMessage = z.infer<typeof userCodebuffMessageSchema>
 
-export const assistantCodebuffMessageSchema = z.object({
-  role: z.literal('assistant'),
-  content: z.union([
-    z.string(),
-    z
-      .union([
-        textPartSchema,
-        filePartSchema,
-        reasoningPartSchema,
-        toolCallPartSchema,
-        toolResultPartSchema,
-      ])
-      .array(),
-  ]),
-})
+export const assistantCodebuffMessageSchema = z
+  .object({
+    role: z.literal('assistant'),
+    content: z.union([
+      z.string(),
+      z
+        .union([
+          textPartSchema,
+          filePartSchema,
+          reasoningPartSchema,
+          toolCallPartSchema,
+          toolResultPartSchema,
+        ])
+        .array(),
+    ]),
+  })
+  .and(auxiliaryDataSchema)
 export type AssistantCodebuffMessage = z.infer<
   typeof assistantCodebuffMessageSchema
 >
 
-export const toolCodebuffMessageSchema = z.object({
-  role: z.literal('tool'),
-  content: toolResultPartSchema.array(),
-})
+export const toolCodebuffMessageSchema = z
+  .object({
+    role: z.literal('tool'),
+    content: toolResultPartSchema.array(),
+  })
+  .and(auxiliaryDataSchema)
 export type ToolCodebuffMessage = z.infer<typeof toolCodebuffMessageSchema>
 
 export const codebuffMessageSchema = z