Call Anthropic token count endpoint via our web api every step. Use in context pruner

Author: jahooma

.agents/context-pruner.ts

@@ -241,8 +241,7 @@ const definition: AgentDefinition = {
     }
 
     // Initial check - if already under limit, return early (skip all pruning)
-    const initialTokens = countMessagesTokens(currentMessages)
-    if (initialTokens < maxMessageTokens) {
+    if (agentState.contextTokenCount < maxMessageTokens) {
       yield {
         toolName: 'set_messages',
         input: { messages: currentMessages },

.agents/types/agent-definition.ts

@@ -291,6 +291,12 @@ export interface AgentState {
     string,
     { description: string | undefined; inputSchema: {} }
   >
+
+  /**
+   * The token count from the Anthropic API.
+   * This is updated on every agent step via the /api/v1/token-count endpoint.
+   */
+  contextTokenCount: number
 }
 
 /**

.env.example

@@ -2,6 +2,7 @@
 CLAUDE_CODE_KEY=dummy_claude_code_key
 OPEN_ROUTER_API_KEY=dummy_openrouter_key
 OPENAI_API_KEY=dummy_openai_key
+ANTHROPIC_API_KEY=dummy_anthropic_key
 
 # Database & Server
 DATABASE_URL=postgresql://manicode_user_local:secretpassword_local@localhost:5432/manicode_db_local

common/src/constants/analytics-events.ts

@@ -122,6 +122,12 @@ export enum AnalyticsEvent {
   // Web - Ads API
   ADS_API_AUTH_ERROR = 'api.ads_auth_error',
 
+  // Web - Token Count API
+  TOKEN_COUNT_REQUEST = 'api.token_count_request',
+  TOKEN_COUNT_AUTH_ERROR = 'api.token_count_auth_error',
+  TOKEN_COUNT_VALIDATION_ERROR = 'api.token_count_validation_error',
+  TOKEN_COUNT_ERROR = 'api.token_count_error',
+
   // Common
   FLUSH_FAILED = 'common.flush_failed',
 }

common/src/templates/initial-agents-dir/types/agent-definition.ts

@@ -291,6 +291,12 @@ export interface AgentState {
     string,
     { description: string | undefined; inputSchema: {} }
   >
+
+  /**
+   * The token count from the Anthropic API.
+   * This is updated on every agent step via the /api/v1/token-count endpoint.
+   */
+  contextTokenCount: number
 }
 
 /**

common/src/types/session-state.ts

@@ -44,6 +44,11 @@ export type AgentState = {
     string,
     { description: string | undefined; inputSchema: {} }
   >
+  /**
+   * The accurate token count from the Anthropic API.
+   * This is updated on every agent step via the /api/v1/token-count endpoint.
+   */
+  contextTokenCount: number
 }
 
 export const AgentOutputSchema = z.discriminatedUnion('type', [
@@ -127,6 +132,7 @@ export function getInitialAgentState(): AgentState {
     parentId: undefined,
     systemPrompt: '',
     toolDefinitions: {},
+    contextTokenCount: 0,
   }
 }
 export function getInitialSessionState(

packages/agent-runtime/src/__tests__/cost-aggregation.test.ts

@@ -124,6 +124,7 @@ describe('Cost Aggregation System', () => {
         directCreditsUsed: 50,
         systemPrompt: 'Test system prompt',
         toolDefinitions: {},
+        contextTokenCount: 0,
       }
 
       // Mock executeAgent to return results with different credit costs
@@ -136,7 +137,10 @@ describe('Cost Aggregation System', () => {
             stepsRemaining: 10,
             creditsUsed: 75, // First subagent uses 75 credits
           },
-          output: { type: 'lastMessage', value: [assistantMessage('Sub-agent 1 response')] },
+          output: {
+            type: 'lastMessage',
+            value: [assistantMessage('Sub-agent 1 response')],
+          },
         })
         .mockResolvedValueOnce({
           agentState: {
@@ -146,7 +150,10 @@ describe('Cost Aggregation System', () => {
             stepsRemaining: 10,
             creditsUsed: 100, // Second subagent uses 100 credits
           },
-          output: { type: 'lastMessage', value: [assistantMessage('Sub-agent 2 response')] },
+          output: {
+            type: 'lastMessage',
+            value: [assistantMessage('Sub-agent 2 response')],
+          },
         })
 
       const mockToolCall = {
@@ -200,7 +207,10 @@ describe('Cost Aggregation System', () => {
             stepsRemaining: 10,
             creditsUsed: 50, // Successful agent
           },
-          output: { type: 'lastMessage', value: [assistantMessage('Successful response')] },
+          output: {
+            type: 'lastMessage',
+            value: [assistantMessage('Successful response')],
+          },
         })
         .mockRejectedValueOnce(
           (() => {
@@ -347,7 +357,10 @@ describe('Cost Aggregation System', () => {
             stepsRemaining: 10,
             creditsUsed: subAgent1Cost,
           } as AgentState,
-          output: { type: 'lastMessage', value: [assistantMessage('Sub-agent 1 response')] },
+          output: {
+            type: 'lastMessage',
+            value: [assistantMessage('Sub-agent 1 response')],
+          },
         })
         .mockResolvedValueOnce({
           agentState: {
@@ -358,7 +371,10 @@ describe('Cost Aggregation System', () => {
             stepsRemaining: 10,
             creditsUsed: subAgent2Cost,
           } as AgentState,
-          output: { type: 'lastMessage', value: [assistantMessage('Sub-agent 2 response')] },
+          output: {
+            type: 'lastMessage',
+            value: [assistantMessage('Sub-agent 2 response')],
+          },
         })
 
       const mockToolCall = {

packages/agent-runtime/src/__tests__/main-prompt.test.ts

@@ -99,6 +99,12 @@ describe('mainPrompt', () => {
       onResponseChunk: () => {},
       localAgentTemplates: mockLocalAgentTemplates,
       signal: new AbortController().signal,
+      // Mock fetch to return a token count response
+      fetch: async () =>
+        ({
+          ok: true,
+          text: async () => JSON.stringify({ inputTokens: 1000 }),
+        }) as Response,
     }
 
     // Mock analytics and tracing
@@ -451,56 +457,4 @@ describe('mainPrompt', () => {
 
     expect(output.type).toBeDefined() // Output should exist even for empty response
   })
-
-  it('should unescape ampersands in run_terminal_command tool calls', async () => {
-    const sessionState = getInitialSessionState(mockFileContext)
-    const userPromptText = 'Run the backend tests'
-    const expectedCommand = 'cd backend && bun test'
-
-    mockAgentStream([
-      createToolCallChunk('run_terminal_command', {
-        command: expectedCommand,
-        process_type: 'SYNC',
-      }),
-      createToolCallChunk('end_turn', {}),
-    ])
-
-    // Get reference to the spy so we can check if it was called
-    const requestToolCallSpy = mainPromptBaseParams.requestToolCall
-
-    const action = {
-      type: 'prompt' as const,
-      prompt: userPromptText,
-      sessionState,
-      fingerprintId: 'test',
-      costMode: 'max' as const,
-      promptId: 'test',
-      toolResults: [],
-    }
-
-    await mainPrompt({
-      ...mainPromptBaseParams,
-      repoId: undefined,
-      repoUrl: undefined,
-      action,
-      userId: TEST_USER_ID,
-      clientSessionId: 'test-session',
-      onResponseChunk: () => {},
-      localAgentTemplates: mockLocalAgentTemplates,
-    })
-
-    // Assert that requestToolCall was called exactly once
-    expect(requestToolCallSpy).toHaveBeenCalledTimes(1)
-
-    // Verify the run_terminal_command call was made with the correct arguments
-    expect(requestToolCallSpy).toHaveBeenCalledWith({
-      userInputId: expect.any(String), // userInputId
-      toolName: 'run_terminal_command',
-      input: expect.objectContaining({
-        command: expectedCommand,
-        process_type: 'SYNC',
-        mode: 'assistant',
-      }),
-    })
-  })
 })

packages/agent-runtime/src/__tests__/prompt-caching-subagents.test.ts

@@ -44,7 +44,12 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
   let capturedMessages: Message[] = []
   let loopAgentStepsBaseParams: ParamsExcluding<
     typeof loopAgentSteps,
-    'agentState' | 'userInputId' | 'prompt' | 'agentType' | 'parentSystemPrompt'
+    | 'agentState'
+    | 'userInputId'
+    | 'prompt'
+    | 'agentType'
+    | 'parentSystemPrompt'
+    | 'agentTemplate'
   >
 
   beforeAll(() => {

packages/agent-runtime/src/__tests__/read-docs-tool.test.ts

@@ -30,7 +30,7 @@ import type { ParamsExcluding } from '@codebuff/common/types/function-params'
 let agentRuntimeImpl: AgentRuntimeDeps & AgentRuntimeScopedDeps
 let runAgentStepBaseParams: ParamsExcluding<
   typeof runAgentStep,
-  'fileContext' | 'localAgentTemplates' | 'agentState' | 'prompt'
+  'fileContext' | 'localAgentTemplates' | 'agentState' | 'prompt' | 'agentTemplate'
 >
 
 import type { StreamChunk } from '@codebuff/common/types/contracts/llm'
@@ -126,6 +126,7 @@ describe('read_docs tool with researcher agent (via web API facade)', () => {
       ...runAgentStepBaseParams,
       fileContext: mockFileContextWithAgents,
       localAgentTemplates: agentTemplates,
+      agentTemplate: agentTemplates['researcher'],
       agentState,
       prompt: 'Get React documentation',
     })
@@ -173,6 +174,7 @@ describe('read_docs tool with researcher agent (via web API facade)', () => {
       ...runAgentStepBaseParams,
       fileContext: mockFileContextWithAgents,
       localAgentTemplates: agentTemplates,
+      agentTemplate: agentTemplates['researcher'],
       agentState,
       prompt: 'Get React hooks documentation',
     })
@@ -212,6 +214,7 @@ describe('read_docs tool with researcher agent (via web API facade)', () => {
       ...runAgentStepBaseParams,
       fileContext: mockFileContextWithAgents,
       localAgentTemplates: agentTemplates,
+      agentTemplate: agentTemplates['researcher'],
       agentState,
       prompt: 'Get documentation for NonExistentLibrary',
     })
@@ -251,6 +254,7 @@ describe('read_docs tool with researcher agent (via web API facade)', () => {
       ...runAgentStepBaseParams,
       fileContext: mockFileContextWithAgents,
       localAgentTemplates: agentTemplates,
+      agentTemplate: agentTemplates['researcher'],
       agentState,
       prompt: 'Get React documentation',
     })
@@ -289,6 +293,7 @@ describe('read_docs tool with researcher agent (via web API facade)', () => {
       ...runAgentStepBaseParams,
       fileContext: mockFileContextWithAgents,
       localAgentTemplates: agentTemplates,
+      agentTemplate: agentTemplates['researcher'],
       agentState,
       prompt: 'Get React server components documentation',
     })
@@ -329,6 +334,7 @@ describe('read_docs tool with researcher agent (via web API facade)', () => {
       ...runAgentStepBaseParams,
       fileContext: mockFileContextWithAgents,
       localAgentTemplates: agentTemplates,
+      agentTemplate: agentTemplates['researcher'],
       agentState,
       prompt: 'Get React documentation',
     })
@@ -374,6 +380,7 @@ describe('read_docs tool with researcher agent (via web API facade)', () => {
       ...runAgentStepBaseParams,
       fileContext: mockFileContextWithAgents,
       localAgentTemplates: agentTemplates,
+      agentTemplate: agentTemplates['researcher'],
       agentState,
       prompt: 'Get React documentation',
     })