[codex] fix(tokens): fallback when provider lacks countTokens (#1624)

* fix(tokens): fallback when provider lacks countTokens

* test(tokens): isolate shim fallback coverage

Avoid process-wide api client mocks in token estimation tests by exercising the count-token dispatch helper directly. Add non-empty tool coverage for the rough fallback path so the local overhead remains covered.

src/services/tokenEstimation.test.ts

@@ -0,0 +1,80 @@
+import type { Anthropic } from '@anthropic-ai/sdk'
+import { expect, mock, test } from 'bun:test'
+import { jsonStringify } from '../utils/slowOperations.js'
+import { __test, roughTokenCountEstimation } from './tokenEstimation.js'
+
+function createTextTool(): Anthropic.Beta.Messages.BetaToolUnion {
+  return {
+    name: 'lookup_docs',
+    description: 'Look up project documentation.',
+    input_schema: {
+      type: 'object',
+      properties: {
+        query: { type: 'string' },
+      },
+      required: ['query'],
+    },
+  }
+}
+
+test('countMessagesTokensWithClient falls back when shim client lacks countTokens', async () => {
+  const content = 'hello from an openai-compatible provider'
+
+  const result = await __test.countMessagesTokensWithClient({
+    messagesClient: {},
+    model: 'gpt-4o',
+    messages: [{ role: 'user', content }],
+    tools: [],
+    filteredBetas: [],
+    containsThinking: false,
+  })
+
+  expect(result).toBe(roughTokenCountEstimation(content))
+})
+
+test('countMessagesTokensWithClient includes tool overhead in fallback estimates', async () => {
+  const content = 'count this request with tool definitions'
+  const tools = [createTextTool()]
+
+  const result = await __test.countMessagesTokensWithClient({
+    messagesClient: {},
+    model: 'gpt-4o',
+    messages: [{ role: 'user', content }],
+    tools,
+    filteredBetas: [],
+    containsThinking: false,
+  })
+
+  expect(result).toBe(
+    roughTokenCountEstimation(content) +
+      500 +
+      roughTokenCountEstimation(jsonStringify(tools)),
+  )
+})
+
+test('countMessagesTokensWithClient uses countTokens when the client supports it', async () => {
+  const countTokens = mock(async (_params: unknown) => ({ input_tokens: 42 }))
+  const messages: Anthropic.Beta.Messages.BetaMessageParam[] = [
+    { role: 'user', content: 'use exact count when available' },
+  ]
+
+  const result = await __test.countMessagesTokensWithClient({
+    messagesClient: {
+      countTokens:
+        countTokens as unknown as Anthropic['beta']['messages']['countTokens'],
+    },
+    model: 'gpt-4o',
+    messages,
+    tools: [],
+    filteredBetas: [],
+    containsThinking: false,
+  })
+
+  expect(countTokens).toHaveBeenCalledTimes(1)
+  expect(countTokens.mock.calls[0]?.[0]).toEqual({
+    model: 'gpt-4o',
+    messages,
+    tools: [],
+  })
+  expect(result).toBe(42)
+})

src/services/tokenEstimation.ts

@@ -31,6 +31,12 @@ import { withTokenCountVCR } from './vcr.js'
 // API constraint: max_tokens must be greater than thinking.budget_tokens
 const TOKEN_COUNT_THINKING_BUDGET = 1024
 const TOKEN_COUNT_MAX_TOKENS = 2048
+// Keep this local to avoid importing analyzeContext.ts, which already depends on tokenEstimation.
+const ROUGH_TOOL_TOKEN_COUNT_OVERHEAD = 500
+
+type CountTokensMessagesClient = {
+  countTokens?: Anthropic['beta']['messages']['countTokens']
+}
 
 /**
  * Check if messages contain thinking blocks
@@ -163,36 +169,27 @@ export async function countMessagesTokensWithAPI(
         model,
         source: 'count_tokens',
       })
+      const messagesClient = (
+        anthropic as {
+          beta?: {
+            messages?: CountTokensMessagesClient
+          }
+        }
+      ).beta?.messages
 
       const filteredBetas =
         getAPIProvider() === 'vertex'
           ? betas.filter(b => VERTEX_COUNT_TOKENS_ALLOWED_BETAS.has(b))
           : betas
 
-      const response = await anthropic.beta.messages.countTokens({
-        model: normalizeModelStringForAPI(model),
-        messages:
-          // When we pass tools and no messages, we need to pass a dummy message
-          // to get an accurate tool token count.
-          messages.length > 0 ? messages : [{ role: 'user', content: 'foo' }],
+      return countMessagesTokensWithClient({
+        messagesClient,
+        model,
+        messages,
         tools,
-        ...(filteredBetas.length > 0 && { betas: filteredBetas }),
-        // Enable thinking if messages contain thinking blocks
-        ...(containsThinking && {
-          thinking: {
-            type: 'enabled',
-            budget_tokens: TOKEN_COUNT_THINKING_BUDGET,
-          },
-        }),
+        filteredBetas,
+        containsThinking,
       })
-
-      if (typeof response.input_tokens !== 'number') {
-        // Vertex client throws
-        // Bedrock client succeeds with { Output: { __type: 'com.amazon.coral.service#UnknownOperationException' }, Version: '1.0' }
-        return null
-      }
-
-      return response.input_tokens
     } catch (error) {
       logError(error)
       return null
@@ -200,6 +197,79 @@ export async function countMessagesTokensWithAPI(
   })
 }
 
+async function countMessagesTokensWithClient({
+  messagesClient,
+  model,
+  messages,
+  tools,
+  filteredBetas,
+  containsThinking,
+}: {
+  messagesClient: CountTokensMessagesClient | undefined
+  model: string
+  messages: Anthropic.Beta.Messages.BetaMessageParam[]
+  tools: Anthropic.Beta.Messages.BetaToolUnion[]
+  filteredBetas: string[]
+  containsThinking: boolean
+}): Promise<number | null> {
+  if (typeof messagesClient?.countTokens !== 'function') {
+    return roughTokenCountEstimationForCountTokensFallback(messages, tools)
+  }
+
+  const response = await messagesClient.countTokens({
+    model: normalizeModelStringForAPI(model),
+    messages:
+      // When we pass tools and no messages, we need to pass a dummy message
+      // to get an accurate tool token count.
+      messages.length > 0 ? messages : [{ role: 'user', content: 'foo' }],
+    tools,
+    ...(filteredBetas.length > 0 && { betas: filteredBetas }),
+    // Enable thinking if messages contain thinking blocks
+    ...(containsThinking && {
+      thinking: {
+        type: 'enabled',
+        budget_tokens: TOKEN_COUNT_THINKING_BUDGET,
+      },
+    }),
+  })
+
+  if (typeof response.input_tokens !== 'number') {
+    // Vertex client throws
+    // Bedrock client succeeds with { Output: { __type: 'com.amazon.coral.service#UnknownOperationException' }, Version: '1.0' }
+    return null
+  }
+
+  return response.input_tokens
+}
+
+function roughTokenCountEstimationForCountTokensFallback(
+  messages: Anthropic.Beta.Messages.BetaMessageParam[],
+  tools: Anthropic.Beta.Messages.BetaToolUnion[],
+): number {
+  let totalTokens = 0
+
+  for (const message of messages) {
+    totalTokens += roughTokenCountEstimationForContent(
+      message.content as
+        | string
+        | Array<Anthropic.ContentBlock>
+        | Array<Anthropic.ContentBlockParam>
+        | undefined,
+    )
+  }
+
+  if (tools.length > 0) {
+    totalTokens +=
+      ROUGH_TOOL_TOKEN_COUNT_OVERHEAD +
+      roughTokenCountEstimation(jsonStringify(tools))
+  }
+
+  return totalTokens
+}
+
+// Test-only surface for fallback dispatch without process-wide module mocks.
+export const __test = { countMessagesTokensWithClient }
+
 export function roughTokenCountEstimation(
   content: string,
   bytesPerToken: number = 4,