Include tools in token count api. Only gemini/openai get 30% token increase for token counter

Author: jahooma

packages/agent-runtime/src/llm-api/codebuff-web-api.ts

@@ -230,13 +230,14 @@ export async function callTokenCountAPI(params: {
   messages: unknown[]
   system?: string
   model?: string
+  tools?: Array<{ name: string; description?: string; input_schema?: unknown }>
   fetch: typeof globalThis.fetch
   logger: Logger
   env: CodebuffWebApiEnv
   baseUrl?: string
   apiKey?: string
 }): Promise<{ inputTokens?: number; error?: string }> {
-  const { messages, system, model, fetch, logger, env } = params
+  const { messages, system, model, tools, fetch, logger, env } = params
   const baseUrl = params.baseUrl ?? env.clientEnv.NEXT_PUBLIC_CODEBUFF_APP_URL
   const apiKey = params.apiKey ?? env.ciEnv.CODEBUFF_API_KEY
 
@@ -248,6 +249,7 @@ export async function callTokenCountAPI(params: {
   const payload: Record<string, unknown> = { messages }
   if (system) payload.system = system
   if (model) payload.model = model
+  if (tools) payload.tools = tools
 
   try {
     const res = await withTimeout(

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

@@ -806,6 +806,18 @@ export async function loopAgentSteps(
     systemPrompt: system,
     toolDefinitions,
   }
+
+  // Convert tool definitions to Anthropic format for accurate token counting
+  // Tool definitions are stored as { [name]: { description, inputSchema } }
+  // Anthropic count_tokens API expects [{ name, description, input_schema }]
+  const toolsForTokenCount = Object.entries(toolDefinitions).map(
+    ([name, def]) => ({
+      name,
+      ...(def.description && { description: def.description }),
+      ...(def.inputSchema && { input_schema: def.inputSchema }),
+    }),
+  )
+
   let shouldEndTurn = false
   let hasRetriedOutputSchema = false
   let currentPrompt = prompt
@@ -845,6 +857,7 @@ export async function loopAgentSteps(
         messages: messagesWithStepPrompt,
         system,
         model: agentTemplate.model,
+        tools: toolsForTokenCount,
         fetch,
         logger,
         env: { clientEnv, ciEnv },

web/src/app/api/v1/token-count/_post.ts

@@ -3,6 +3,7 @@ import {
   isClaudeModel,
   toAnthropicModelId,
 } from '@codebuff/common/constants/claude-oauth'
+import { isOpenAIProviderModel } from '@codebuff/common/constants/chatgpt-oauth'
 import { getErrorObject } from '@codebuff/common/util/error'
 import { env } from '@codebuff/internal/env'
 import { NextResponse } from 'next/server'
@@ -22,6 +23,11 @@ const tokenCountRequestSchema = z.object({
   messages: z.array(z.any()),
   system: z.string().optional(),
   model: z.string().optional(),
+  tools: z.array(z.object({
+    name: z.string(),
+    description: z.string().optional(),
+    input_schema: z.any().optional(),
+  })).optional(),
 })
 
 type TokenCountRequest = z.infer<typeof tokenCountRequestSchema>
@@ -74,24 +80,27 @@ export async function postTokenCount(params: {
     return bodyResult.response
   }
 
-  const { messages, system, model } = bodyResult.data
+  const { messages, system, model, tools } = bodyResult.data
 
   try {
     const useOpenAI = model != null && false // isOpenAIProviderModel(model)
     const inputTokens = useOpenAI
       ? await countTokensViaOpenAI({ messages, system, model, fetch, logger })
       : await countTokensViaAnthropic({
-          messages,
-          system,
-          model,
-          fetch,
-          logger,
-        })
+        messages,
+        system,
+        model,
+        tools,
+        fetch,
+        logger,
+      })
 
     logger.info({
       userId,
       messageCount: messages.length,
       hasSystem: !!system,
+      hasTools: !!tools,
+      toolCount: tools?.length,
       model: model ?? DEFAULT_ANTHROPIC_MODEL,
       tokenCount: inputTokens,
       provider: useOpenAI ? 'openai' : 'anthropic',
@@ -285,10 +294,11 @@ async function countTokensViaAnthropic(params: {
   messages: TokenCountRequest['messages']
   system: string | undefined
   model: string | undefined
+  tools: TokenCountRequest['tools']
   fetch: typeof globalThis.fetch
   logger: Logger
 }): Promise<number> {
-  const { messages, system, model, fetch, logger } = params
+  const { messages, system, model, tools, fetch, logger } = params
 
   // Convert messages to Anthropic format
   const anthropicMessages = convertToAnthropicMessages(messages)
@@ -315,6 +325,7 @@ async function countTokensViaAnthropic(params: {
         model: anthropicModelId,
         messages: anthropicMessages,
         ...(system && { system }),
+        ...(tools && { tools }),
       }),
     },
   )
@@ -337,8 +348,12 @@ async function countTokensViaAnthropic(params: {
   const data = await response.json()
   const baseTokens = data.input_tokens
 
-  // Add 30% buffer for non-Anthropic models since tokenizers differ
-  if (isNonAnthropicModel) {
+  // Add 30% buffer for OpenAI and Gemini models since their tokenizers differ from Anthropic's
+  // Other non-Anthropic models (x-ai, qwen, deepseek, etc.) are routed through providers that
+  // use similar tokenization, so the buffer is not needed and was causing premature context pruning.
+  const isOpenAIModel = model ? isOpenAIProviderModel(model) : false
+  const isGeminiModel = model?.startsWith('google/') ?? false
+  if (isOpenAIModel || isGeminiModel) {
     return Math.ceil(baseTokens * (1 + NON_ANTHROPIC_TOKEN_BUFFER))
   }