Sync public snapshot from freebuff-private

Source: CodebuffAI/freebuff-private@3f892919d2b3ed8b094ac58b76014de866dcb53c

Author: github-actions[bot]

bun.lock

@@ -51,7 +51,11 @@ import {
   buildUserMessageContent,
   expireMessages,
 } from './util/messages'
-import { countTokensJson } from './util/token-counter'
+import {
+  countTokens,
+  countTokensJson,
+  countTokensMessages,
+} from './util/token-counter'
 
 import type { AgentTemplate } from '@codebuff/common/types/agent-template'
 import type { TrackEventFn } from '@codebuff/common/types/contracts/analytics'
@@ -99,7 +103,9 @@ export function toTokenCountInputSchema(
   if (inputSchema == null) return undefined
 
   let jsonSchema: Record<string, unknown>
-  if (typeof (inputSchema as { safeParse?: unknown }).safeParse === 'function') {
+  if (
+    typeof (inputSchema as { safeParse?: unknown }).safeParse === 'function'
+  ) {
     try {
       jsonSchema = z.toJSONSchema(inputSchema as z.ZodType, {
         io: 'input',
@@ -340,7 +346,9 @@ export const runAgentStep = async (
   }
 
   const iterationNum = agentState.messageHistory.length
-  const systemTokens = countTokensJson(system)
+  // system is a plain string; count it directly rather than JSON-stringifying
+  // it (which would add quotes and escape every newline).
+  const systemTokens = countTokens(system)
 
   let cacheDebugCorrelation:
     | ReturnType<typeof createCacheDebugSnapshot>
@@ -977,9 +985,12 @@ export async function loopAgentSteps(
           }),
       )
 
+      // Count structured message content (not JSON.stringify, which inflates the
+      // count and counts image base64 as text); system is a plain string; tool
+      // schemas stay JSON since that's roughly how the model sees them.
       const estimateContextTokensLocally = () =>
-        countTokensJson(messagesWithStepPrompt) +
-        countTokensJson(system) +
+        countTokensMessages(messagesWithStepPrompt) +
+        countTokens(system) +
         countTokensJson(toolsForTokenCount)
 
       // Free (freebuff) runs never call the token-count web API: the awaited
@@ -1014,11 +1025,7 @@ export async function loopAgentSteps(
             { error: tokenCountResult.error },
             'Failed to get token count from web API',
           )
-          const estimatedTokens =
-            countTokensJson(currentAgentState.messageHistory) +
-            countTokensJson(system) +
-            countTokensJson(toolDefinitions)
-          currentAgentState.contextTokenCount = estimatedTokens
+          currentAgentState.contextTokenCount = estimateContextTokensLocally()
         }
       }
 

packages/agent-runtime/src/run-agent-step.ts

@@ -377,7 +377,8 @@ describe('trimMessagesToFitTokenLimit', () => {
       const result = trimMessagesToFitTokenLimit({
         messages,
         systemTokens: 0,
-        maxTotalTokens: 1000,
+        // Below the messages' structured token total so truncation triggers.
+        maxTotalTokens: 300,
         logger,
       })
 
@@ -438,7 +439,8 @@ describe('trimMessagesToFitTokenLimit', () => {
       const result = trimMessagesToFitTokenLimit({
         messages,
         systemTokens: 0,
-        maxTotalTokens: 1000,
+        // Below the messages' structured token total so truncation triggers.
+        maxTotalTokens: 500,
         logger,
       })
 

packages/agent-runtime/src/util/__tests__/messages.test.ts

@@ -0,0 +1,120 @@
+import { describe, expect, test } from 'bun:test'
+
+import {
+  countTokens,
+  countTokensJson,
+  countTokensMessages,
+} from '../token-counter'
+
+import type { Message } from '@codebuff/common/types/messages/codebuff-message'
+
+describe('countTokensMessages', () => {
+  test('counts text content plus per-message overhead', () => {
+    const messages = [
+      { role: 'user', content: [{ type: 'text', text: 'hello world' }] },
+    ] as unknown as Message[]
+
+    const count = countTokensMessages(messages)
+    // At least the text tokens; overhead makes it strictly greater.
+    expect(count).toBeGreaterThan(countTokens('hello world'))
+  })
+
+  test('counts tool-call inputs and tool-result payloads', () => {
+    const messages = [
+      {
+        role: 'assistant',
+        content: [
+          { type: 'text', text: 'reading' },
+          {
+            type: 'tool-call',
+            toolCallId: '1',
+            toolName: 'read_files',
+            input: { paths: ['a.ts', 'b.ts'] },
+          },
+        ],
+      },
+      {
+        role: 'tool',
+        toolCallId: '1',
+        toolName: 'read_files',
+        content: [{ type: 'json', value: { files: ['contents here'] } }],
+      },
+    ] as unknown as Message[]
+
+    expect(countTokensMessages(messages)).toBeGreaterThan(0)
+  })
+
+  test('is cheaper than JSON.stringify counting for structured messages', () => {
+    const codeBlock = 'const x = "a \\"b\\" c";\n'.repeat(50)
+    const messages = [
+      {
+        role: 'tool',
+        toolCallId: '1',
+        toolName: 'read_files',
+        content: [
+          { type: 'json', value: { path: 'x.ts', content: codeBlock } },
+        ],
+      },
+      { role: 'user', content: [{ type: 'text', text: codeBlock }] },
+    ] as unknown as Message[]
+
+    // Structured counting removes the JSON envelope for text parts, so it must
+    // not exceed the whole-array JSON count.
+    expect(countTokensMessages(messages)).toBeLessThanOrEqual(
+      countTokensJson(messages),
+    )
+  })
+
+  test('does NOT count image/file base64 as text (uses a fixed estimate)', () => {
+    const bigB64 = 'A'.repeat(500_000)
+    const messages = [
+      {
+        role: 'user',
+        content: [
+          { type: 'text', text: 'screenshot:' },
+          { type: 'image', image: bigB64, mediaType: 'image/png' },
+        ],
+      },
+      {
+        role: 'tool',
+        toolCallId: '1',
+        toolName: 'browser',
+        content: [{ type: 'media', data: bigB64, mediaType: 'image/png' }],
+      },
+    ] as unknown as Message[]
+
+    const count = countTokensMessages(messages)
+    // Two images at a ~1600-token ceiling each, plus a little text/overhead —
+    // nowhere near the ~hundreds-of-thousands of tokens the base64 would add.
+    expect(count).toBeLessThan(10_000)
+  })
+
+  test('tolerates string content without iterating characters', () => {
+    const text = 'a plain string message content'
+    const messages = [{ role: 'user', content: text }] as unknown as Message[]
+
+    // Must count the string as one blob, not char-by-char via the default case.
+    expect(countTokensMessages(messages)).toBeGreaterThan(countTokens(text))
+    expect(countTokensMessages(messages)).toBeLessThan(countTokens(text) + 100)
+  })
+
+  test('does not throw on missing/non-array content', () => {
+    const messages = [
+      { role: 'assistant' },
+      { role: 'user', content: undefined },
+    ] as unknown as Message[]
+
+    expect(() => countTokensMessages(messages)).not.toThrow()
+  })
+
+  test('falls back to JSON for unknown part shapes so it never under-counts', () => {
+    const messages = [
+      {
+        role: 'user',
+        content: [{ type: 'mystery', payload: 'some meaningful content here' }],
+      },
+    ] as unknown as Message[]
+
+    expect(countTokensMessages(messages)).toBeGreaterThan(0)
+  })
+})

packages/agent-runtime/src/util/__tests__/token-counter.test.ts

@@ -7,7 +7,7 @@ import { closeXml } from '@codebuff/common/util/xml'
 import { cloneDeep, isEqual } from 'lodash'
 
 import { simplifyTerminalCommandResults } from './simplify-tool-results'
-import { countTokensJson } from './token-counter'
+import { countTokensMessages } from './token-counter'
 
 import type { System } from '../llm-api/claude'
 import type {
@@ -189,7 +189,7 @@ export function trimMessagesToFitTokenLimit(params: {
   const maxMessageTokens = maxTotalTokens - systemTokens
 
   // Check if we're already under the limit
-  const initialTokens = countTokensJson(messages)
+  const initialTokens = countTokensMessages(messages)
 
   if (initialTokens < maxMessageTokens) {
     return messages
@@ -231,7 +231,7 @@ export function trimMessagesToFitTokenLimit(params: {
   }
   shortenedMessages.reverse()
 
-  const requiredTokens = countTokensJson(
+  const requiredTokens = countTokensMessages(
     shortenedMessages.filter((m) => m.keepDuringTruncation),
   )
   let removedTokens = 0
@@ -245,13 +245,13 @@ export function trimMessagesToFitTokenLimit(params: {
       filteredMessages.push(message)
       continue
     }
-    removedTokens += countTokensJson(message)
+    removedTokens += countTokensMessages([message])
     if (
       filteredMessages.length === 0 ||
       filteredMessages[filteredMessages.length - 1] !== placeholder
     ) {
       filteredMessages.push(placeholder)
-      removedTokens -= countTokensJson(replacementMessage)
+      removedTokens -= countTokensMessages([replacementMessage])
     }
   }
 

packages/agent-runtime/src/util/messages.ts

@@ -1,8 +1,19 @@
 import { LRUCache } from '@codebuff/common/util/lru-cache'
 import { encode } from 'gpt-tokenizer/esm/model/gpt-4o'
 
+import type { Message } from '@codebuff/common/types/messages/codebuff-message'
+
 const ANTHROPIC_TOKEN_FUDGE_FACTOR = 1.35
 
+/** Flat per-image/file cost. Anthropic bills a large image at ~1600 tokens; we
+ *  use that ceiling instead of counting the base64 as text (which JSON.stringify
+ *  would do), where a single screenshot is hundreds of thousands of chars. */
+const IMAGE_TOKEN_ESTIMATE = 1600
+
+/** Per-message structural overhead (role marker, delimiters) Anthropic adds on
+ *  top of the raw content. */
+const PER_MESSAGE_TOKEN_OVERHEAD = 8
+
 const TOKEN_COUNT_CACHE = new LRUCache<string, number>(1000)
 
 export function countTokens(text: string): number {
@@ -27,8 +38,61 @@ export function countTokens(text: string): number {
   }
 }
 
-export function countTokensJson(text: string | object): number {
-  return countTokens(JSON.stringify(text))
+export function countTokensJson(value: unknown): number {
+  // JSON.stringify(undefined) returns undefined; fall back to '' so countTokens
+  // always gets a string.
+  return countTokens(JSON.stringify(value) ?? '')
+}
+
+/**
+ * Estimate tokens for a list of messages by counting the content the model
+ * actually tokenizes (text, tool inputs, tool results, a flat cost per image)
+ * plus a small per-message overhead — not the JSON envelope. Avoids the
+ * scaffolding inflation of `countTokensJson(messages)` and, crucially, counting
+ * image/file base64 character-for-character. `ANTHROPIC_TOKEN_FUDGE_FACTOR` still
+ * applies, so the estimate stays deliberately a touch above the true count.
+ */
+export function countTokensMessages(messages: Message[]): number {
+  let total = 0
+  for (const message of messages) {
+    total += PER_MESSAGE_TOKEN_OVERHEAD
+
+    // content is typed as an array, but tolerate string / missing content the
+    // same way the replaced JSON.stringify did (see getTextContent), so a stray
+    // shape can't crash or mis-count the estimate.
+    const content = (message as { content?: unknown }).content
+    if (typeof content === 'string') {
+      total += countTokens(content)
+      continue
+    }
+    if (!Array.isArray(content)) {
+      continue
+    }
+
+    for (const part of content as Array<Record<string, unknown>>) {
+      switch (part.type) {
+        case 'text':
+        case 'reasoning':
+          total += countTokens(part.text as string)
+          break
+        case 'tool-call':
+          total +=
+            countTokens(part.toolName as string) + countTokensJson(part.input)
+          break
+        case 'json': // tool result payload
+          total += countTokensJson(part.value)
+          break
+        case 'image':
+        case 'file':
+        case 'media':
+          total += IMAGE_TOKEN_ESTIMATE
+          break
+        default: // unknown shape: JSON fallback so we never under-count
+          total += countTokensJson(part)
+      }
+    }
+  }
+  return total
 }
 
 export function countTokensForFiles(