remove old CodebuffMessage

Author: charleslien

backend/src/__tests__/request-files-prompt.test.ts

@@ -12,7 +12,7 @@ import * as OriginalRequestFilesPromptModule from '../find-files/request-files-p
 import * as geminiWithFallbacksModule from '../llm-apis/gemini-with-fallbacks'
 
 import type { CostMode } from '@codebuff/common/constants'
-import type { CodebuffMessage } from '@codebuff/common/types/message'
+import type { CodebuffMessage } from '@codebuff/common/types/messages/codebuff-message'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
 import type { Mock } from 'bun:test'
 

backend/src/__tests__/spawn-agents-message-history.test.ts

@@ -1,3 +1,5 @@
+import { TEST_USER_ID } from '@codebuff/common/constants'
+import { getInitialSessionState } from '@codebuff/common/types/session-state'
 import {
   describe,
   expect,
@@ -7,16 +9,15 @@ import {
   mock,
   spyOn,
 } from 'bun:test'
-import { handleSpawnAgents } from '../tools/handlers/tool/spawn-agents'
-import { TEST_USER_ID } from '@codebuff/common/constants'
-import { getInitialSessionState } from '@codebuff/common/types/session-state'
+
 import { mockFileContext, MockWebSocket } from './test-utils'
-import * as loggerModule from '../util/logger'
 import * as runAgentStep from '../run-agent-step'
+import { handleSpawnAgents } from '../tools/handlers/tool/spawn-agents'
+import * as loggerModule from '../util/logger'
 
-import type { AgentTemplate } from '@codebuff/common/types/agent-template'
 import type { CodebuffToolCall } from '@codebuff/common/tools/list'
-import type { CodebuffMessage } from '@codebuff/common/types/message'
+import type { AgentTemplate } from '@codebuff/common/types/agent-template'
+import type { CodebuffMessage } from '@codebuff/common/types/messages/codebuff-message'
 import type { WebSocket } from 'ws'
 
 describe('Spawn Agents Message History', () => {

backend/src/admin/relabelRuns.ts

@@ -27,7 +27,7 @@ import type {
   GetRelevantFilesTrace,
   Relabel,
 } from '@codebuff/bigquery'
-import type { Message } from '@codebuff/common/types/message'
+import type { CodebuffMessage } from '@codebuff/common/types/messages/codebuff-message'
 import type { Request, Response } from 'express'
 
 // --- GET Handler Logic ---
@@ -172,7 +172,7 @@ export async function relabelForUserHandler(req: Request, res: Response) {
 
             output = await promptAiSdk({
               messages: transformMessages(
-                messages as Message[],
+                messages as CodebuffMessage[],
                 system as System,
               ),
               model: model,
@@ -397,7 +397,10 @@ export async function relabelWithClaudeWithFullFileContext(
   }
 
   const output = await promptAiSdk({
-    messages: transformMessages(trace.payload.messages as Message[], system),
+    messages: transformMessages(
+      trace.payload.messages as CodebuffMessage[],
+      system,
+    ),
     model: model as any, // Model type is string here for flexibility
     clientSessionId: 'relabel-trace-api',
     fingerprintId: 'relabel-trace-api',

backend/src/fast-rewrite.ts

@@ -11,7 +11,7 @@ import { promptRelaceAI } from './llm-apis/relace-api'
 import { promptAiSdk } from './llm-apis/vercel-ai-sdk/ai-sdk'
 import { logger } from './util/logger'
 
-import type { CodebuffMessage } from '@codebuff/common/types/message'
+import type { CodebuffMessage } from '@codebuff/common/types/messages/codebuff-message'
 
 export async function fastRewrite(
   initialContent: string,

backend/src/find-files/request-files-prompt.ts

@@ -29,7 +29,7 @@ import type {
   GetExpandedFileContextForTrainingTrace,
   GetRelevantFilesTrace,
 } from '@codebuff/bigquery'
-import type { CodebuffMessage } from '@codebuff/common/types/message'
+import type { CodebuffMessage } from '@codebuff/common/types/messages/codebuff-message'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
 
 const NUMBER_OF_EXAMPLE_FILES = 100

backend/src/get-file-reading-updates.ts

@@ -13,7 +13,7 @@ import {
 import { countTokensJson } from './util/token-counter'
 import { requestFiles } from './websockets/websocket-action'
 
-import type { CodebuffMessage } from '@codebuff/common/types/message'
+import type { CodebuffMessage } from '@codebuff/common/types/messages/codebuff-message'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
 import type { WebSocket } from 'ws'
 

backend/src/llm-apis/gemini-with-fallbacks.ts

@@ -8,7 +8,7 @@ import type {
   FinetunedVertexModel,
   GeminiModel,
 } from '@codebuff/common/constants'
-import type { CodebuffMessage } from '@codebuff/common/types/message'
+import type { CodebuffMessage } from '@codebuff/common/types/messages/codebuff-message'
 
 /**
  * Prompts a Gemini model with fallback logic.

backend/src/llm-apis/message-cost-tracker.ts

@@ -22,7 +22,7 @@ import { SWITCHBOARD } from '../websockets/server'
 import { sendAction } from '../websockets/websocket-action'
 
 import type { ClientState } from '../websockets/switchboard'
-import type { CodebuffMessage } from '@codebuff/common/types/message'
+import type { CodebuffMessage } from '@codebuff/common/types/messages/codebuff-message'
 
 export const PROFIT_MARGIN = 0.055
 

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

@@ -29,12 +29,12 @@ import type {
   Model,
   OpenAIModel,
 } from '@codebuff/common/constants'
-import type { CodebuffMessage, Message } from '@codebuff/common/types/message'
+import type { CodebuffMessage } from '@codebuff/common/types/messages/codebuff-message'
 import type {
   OpenRouterProviderOptions,
   OpenRouterUsageAccounting,
 } from '@openrouter/ai-sdk-provider'
-import type { AssistantModelMessage, UserModelMessage, LanguageModel } from 'ai'
+import type { LanguageModel } from 'ai'
 import type { z } from 'zod/v4'
 
 // TODO: We'll want to add all our models here!
@@ -370,7 +370,7 @@ export const promptAiSdkStructured = async function <T>(options: {
 // TODO: temporary - ideally we move to using CodebuffMessage[] directly
 // and don't need this transform!!
 export function transformMessages(
-  messages: (Message | CodebuffMessage)[],
+  messages: CodebuffMessage[],
   system?: System,
 ): CodebuffMessage[] {
   const codebuffMessages: CodebuffMessage[] = []
@@ -385,115 +385,14 @@ export function transformMessages(
     })
   }
 
-  for (const message of messages) {
-    if (message.role === 'system') {
-      if (typeof message.content === 'string') {
-        codebuffMessages.push({ role: 'system', content: message.content })
-        continue
-      } else {
-        throw new Error(
-          'Multiple part system message - unsupported (TODO: fix if we hit this.)',
-        )
-      }
-    }
-
-    if (message.role === 'user') {
-      if (typeof message.content === 'string') {
-        codebuffMessages.push({
-          ...message,
-          role: 'user',
-          content: message.content,
-        })
-        continue
-      } else {
-        const parts: UserModelMessage['content'] = []
-        const modelMessage: UserModelMessage = { role: 'user', content: parts }
-        for (const part of message.content) {
-          // Add ephemeral if present
-          if ('cache_control' in part) {
-            modelMessage.providerOptions = {
-              anthropic: { cacheControl: { type: 'ephemeral' } },
-              openrouter: { cacheControl: { type: 'ephemeral' } },
-            }
-          }
-          // Handle Message type image format
-          if (part.type === 'image' && 'source' in part) {
-            parts.push({
-              type: 'image' as const,
-              image: `data:${part.source.media_type};base64,${part.source.data}`,
-            })
-            continue
-          }
-          if (part.type === 'file') {
-            throw new Error('File messages not supported')
-          }
-          if (part.type === 'text') {
-            parts.push({
-              type: 'text' as const,
-              text: part.text,
-            })
-            continue
-          }
-          if (part.type === 'tool_use' || part.type === 'tool_result') {
-            // Skip tool parts in user messages - they should be in assistant/tool messages
-            continue
-          }
-        }
-        codebuffMessages.push(modelMessage)
-        continue
-      }
-    }
-
-    if (message.role === 'assistant') {
-      if (message.content === undefined || message.content === null) {
-        continue
-      }
-      if (typeof message.content === 'string') {
-        codebuffMessages.push({
-          ...message,
-          role: 'assistant',
-          content: message.content,
-        })
-        continue
-      } else {
-        let messageContent: AssistantModelMessage['content'] = []
-        const modelMessage: AssistantModelMessage = {
-          ...message,
-          role: 'assistant',
-          content: messageContent,
-        }
-        for (const part of message.content) {
-          // Add ephemeral if present
-          if ('cache_control' in part) {
-            modelMessage.providerOptions = {
-              anthropic: { cacheControl: { type: 'ephemeral' } },
-              openrouter: { cacheControl: { type: 'ephemeral' } },
-            }
-          }
-          if (part.type === 'text') {
-            messageContent.push({ type: 'text', text: part.text })
-          }
-          if (part.type === 'tool_use') {
-            messageContent.push({
-              type: 'tool-call',
-              toolCallId: part.id,
-              toolName: part.name,
-              input: part.input,
-            })
-          }
-        }
-        codebuffMessages.push(modelMessage)
-        continue
-      }
-    }
-
-    if (message.role === 'tool') {
-      codebuffMessages.push(message)
-      continue
-    }
-
-    throw new Error('Unknown message role received: ' + message)
-  }
-
-  return codebuffMessages
+  return buildArray<CodebuffMessage>([
+    system && {
+      role: 'system',
+      content:
+        typeof system === 'string'
+          ? system
+          : system.map((block) => block.text).join('\n\n'),
+    },
+    messages,
+  ])
 }

backend/src/process-file-block.ts

@@ -12,7 +12,7 @@ import { promptAiSdk } from './llm-apis/vercel-ai-sdk/ai-sdk'
 import { logger } from './util/logger'
 import { countTokens } from './util/token-counter'
 
-import type { CodebuffMessage } from '@codebuff/common/types/message'
+import type { CodebuffMessage } from '@codebuff/common/types/messages/codebuff-message'
 
 export async function processFileBlock(
   path: string,