feat: forward count_tokens requests

bun.lock

@@ -25,7 +25,7 @@
         "lint-staged": "^16.1.2",
         "playwright": "^1.54.2",
         "prettier": "^3.4.2",
-        "typescript": "^5.8.3",
+        "typescript": "^5.9.2",
         "typescript-eslint": "^8.18.0",
       },
     },
@@ -916,7 +916,7 @@
 
     "type-is": ["[email protected]", "", { "dependencies": { "content-type": "^1.0.5", "media-typer": "^1.1.0", "mime-types": "^3.0.0" } }, "sha512-OZs6gsjF4vMp32qrCbiVSkrFmXtG/AZhY3t0iAMrMBiAZyV9oALtXO8hsrHbMXF9x6L3grlFuwW2oAz7cav+Gw=="],
 
-    "typescript": ["typescript@5.8.3", "", { "bin": { "tsc": "bin/tsc", "tsserver": "bin/tsserver" } }, "sha512-p1diW6TqL9L07nNxvRMM7hMMw4c5XOo/1ibL4aAIGmSAt9slTE1Xgw5KWuof2uTOvCg9BY7ZRi+GaF+7sfgPeQ=="],
+    "typescript": ["typescript@5.9.2", "", { "bin": { "tsc": "bin/tsc", "tsserver": "bin/tsserver" } }, "sha512-CWBzXQrc/qOkhidw1OzBTQuYRbfyxDXJMVJ1XNwUHGROVmuaeiEm3OslpZ1RV96d7SKKjZKrSJu3+t/xlw3R9A=="],
 
     "typescript-eslint": ["[email protected]", "", { "dependencies": { "@typescript-eslint/eslint-plugin": "8.35.1", "@typescript-eslint/parser": "8.35.1", "@typescript-eslint/utils": "8.35.1" }, "peerDependencies": { "eslint": "^8.57.0 || ^9.0.0", "typescript": ">=4.8.4 <5.9.0" } }, "sha512-xslJjFzhOmHYQzSB/QTeASAHbjmxOGEP6Coh93TXmUBFQoJ1VU35UHIDmG06Jd6taf3wqqC1ntBnCMeymy5Ovw=="],
 

package.json

@@ -100,7 +100,7 @@
     "lint-staged": "^16.1.2",
     "playwright": "^1.54.2",
     "prettier": "^3.4.2",
-    "typescript": "^5.8.3",
+    "typescript": "^5.9.2",
     "typescript-eslint": "^8.18.0"
   },
   "dependencies": {

services/proxy/src/app.ts

@@ -254,6 +254,11 @@ export async function createProxyApp(): Promise<
   app.post('/v1/messages', c => messageController.handle(c))
   app.options('/v1/messages', c => messageController.handleOptions(c))
 
+  // Token counting endpoint
+  const tokenCountController = container.getTokenCountController()
+  app.post('/v1/messages/count_tokens', c => tokenCountController.handle(c))
+  app.options('/v1/messages/count_tokens', c => tokenCountController.handleOptions(c))
+
   // Root endpoint
   app.get('/', c => {
     const endpoints: any = {

services/proxy/src/container.ts

@@ -1,5 +1,6 @@
 import { Pool } from 'pg'
 import { MessageController } from './controllers/MessageController.js'
+import { TokenCountController } from './controllers/TokenCountController.js'
 import { ProxyService } from './services/ProxyService.js'
 import { AuthenticationService } from './services/AuthenticationService.js'
 import { ClaudeApiClient } from './services/ClaudeApiClient.js'
@@ -28,6 +29,7 @@ class Container {
   private claudeApiClient?: ClaudeApiClient
   private proxyService?: ProxyService
   private messageController?: MessageController
+  private tokenCountController?: TokenCountController
   private mcpServer?: McpServer
   private promptRegistry?: PromptRegistryService
   private githubSyncService?: GitHubSyncService
@@ -115,6 +117,7 @@ class Container {
     )
 
     this.messageController = new MessageController(this.proxyService)
+    this.tokenCountController = new TokenCountController(this.proxyService)
 
     // Initialize MCP services if enabled
     if (config.mcp.enabled) {
@@ -202,6 +205,13 @@ class Container {
     return this.messageController
   }
 
+  getTokenCountController(): TokenCountController {
+    if (!this.tokenCountController) {
+      throw new Error('TokenCountController not initialized')
+    }
+    return this.tokenCountController
+  }
+
   getMcpHandler(): JsonRpcHandler | undefined {
     return this.jsonRpcHandler
   }
@@ -281,6 +291,10 @@ class LazyContainer {
     return this.ensureInstance().getMessageController()
   }
 
+  getTokenCountController(): TokenCountController {
+    return this.ensureInstance().getTokenCountController()
+  }
+
   getMcpHandler(): JsonRpcHandler | undefined {
     return this.ensureInstance().getMcpHandler()
   }

services/proxy/src/controllers/TokenCountController.ts

@@ -0,0 +1,70 @@
+import { Context } from 'hono'
+import { ProxyService } from '../services/ProxyService'
+import { RequestContext } from '../domain/value-objects/RequestContext'
+import { ValidationError, serializeError } from '@claude-nexus/shared'
+import { getRequestLogger } from '../middleware/logger'
+
+/**
+ * Controller for handling /v1/messages/count_tokens endpoint
+ * Simply forwards token counting requests to Claude API without validation
+ */
+export class TokenCountController {
+  constructor(private proxyService: ProxyService) {}
+
+  /**
+   * Handle POST /v1/messages/count_tokens
+   */
+  async handle(c: Context): Promise<Response> {
+    const logger = getRequestLogger(c)
+    const requestContext = RequestContext.fromHono(c)
+
+    try {
+      // Get the request body - no validation, let Anthropic handle it
+      const body = await c.req.json()
+
+      logger.debug('Processing token count request', {
+        model: body?.model,
+        messageCount: body?.messages?.length,
+        hasSystemField: !!body?.system,
+      })
+
+      // Forward the request to Claude API for token counting
+      const response = await this.proxyService.handleTokenCountRequest(body, requestContext)
+
+      return response
+    } catch (error) {
+      logger.error('Token count request failed', error instanceof Error ? error : undefined)
+
+      // Serialize error for response
+      const errorObj = error instanceof Error ? error : new Error(String(error))
+      const errorResponse = serializeError(errorObj)
+
+      // Determine status code
+      let statusCode = 500
+      if (error instanceof ValidationError) {
+        statusCode = 400
+      } else if ((error as any).statusCode) {
+        statusCode = (error as any).statusCode
+      } else if ((error as any).upstreamStatus) {
+        statusCode = (error as any).upstreamStatus
+      }
+
+      return c.json(errorResponse, statusCode as any)
+    }
+  }
+
+  /**
+   * Handle OPTIONS /v1/messages/count_tokens (CORS preflight)
+   */
+  async handleOptions(_c: Context): Promise<Response> {
+    return new Response('', {
+      status: 204,
+      headers: {
+        'Access-Control-Allow-Origin': '*',
+        'Access-Control-Allow-Methods': 'GET, POST, OPTIONS',
+        'Access-Control-Allow-Headers': 'Content-Type, Authorization, X-API-Key',
+        'Access-Control-Max-Age': '86400',
+      },
+    })
+  }
+}

services/proxy/src/services/ClaudeApiClient.ts

@@ -44,6 +44,123 @@ export class ClaudeApiClient {
     )
   }
 
+  /**
+   * Forward a token count request to Claude API
+   */
+  async forwardTokenCount(rawRequest: any, auth: AuthResult, requestId: string): Promise<Response> {
+    const url = `${this.config.baseUrl}/v1/messages/count_tokens`
+
+    // Build headers for token count request
+    const headers: Record<string, string> = {
+      'Content-Type': 'application/json',
+      'anthropic-version': '2023-06-01',
+      ...auth.headers,
+    }
+
+    // Make the request without retry logic as it's read-only
+    const controller = new AbortController()
+    const timeout = setTimeout(() => controller.abort(), 30000) // 30 second timeout for token counting
+
+    try {
+      const response = await fetch(url, {
+        method: 'POST',
+        headers,
+        body: JSON.stringify(rawRequest),
+        signal: controller.signal,
+      })
+
+      clearTimeout(timeout)
+
+      // Check for errors
+      if (!response.ok) {
+        const errorBody = await response.text()
+        let errorMessage = `Claude API token count error: ${response.status}`
+        let parsedError: any
+
+        try {
+          parsedError = JSON.parse(errorBody)
+          if (isClaudeError(parsedError)) {
+            errorMessage = `${parsedError.error.type}: ${parsedError.error.message}`
+          }
+        } catch {
+          // Use text error if not JSON
+          errorMessage = errorBody || errorMessage
+          parsedError = { error: { message: errorBody, type: 'api_error' } }
+        }
+
+        logger.error('Token count request failed', {
+          requestId,
+          metadata: {
+            status: response.status,
+            error: errorMessage,
+          },
+        })
+
+        // Return error response directly to client
+        return new Response(JSON.stringify(parsedError), {
+          status: response.status,
+          headers: {
+            'Content-Type': 'application/json',
+          },
+        })
+      }
+
+      // Return successful response directly
+      const responseBody = await response.text()
+      return new Response(responseBody, {
+        status: response.status,
+        headers: {
+          'Content-Type': 'application/json',
+        },
+      })
+    } catch (error) {
+      clearTimeout(timeout)
+
+      if (error instanceof Error && error.name === 'AbortError') {
+        logger.error('Token count request timeout', {
+          requestId,
+          metadata: {
+            timeout: 30000,
+          },
+        })
+        return new Response(
+          JSON.stringify({
+            error: {
+              type: 'timeout_error',
+              message: 'Token count request timed out',
+            },
+          }),
+          {
+            status: 504,
+            headers: {
+              'Content-Type': 'application/json',
+            },
+          }
+        )
+      }
+
+      logger.error('Token count request failed with exception', {
+        requestId,
+        error: getErrorMessage(error),
+      })
+
+      return new Response(
+        JSON.stringify({
+          error: {
+            type: 'internal_error',
+            message: 'Failed to process token count request',
+          },
+        }),
+        {
+          status: 500,
+          headers: {
+            'Content-Type': 'application/json',
+          },
+        }
+      )
+    }
+  }
+
   /**
    * Make the actual HTTP request
    */

services/proxy/src/services/ProxyService.ts

@@ -638,4 +638,79 @@ export class ProxyService {
 
     return response
   }
+
+  /**
+   * Handle a token count request
+   * Forwards the request to Claude API's count_tokens endpoint
+   */
+  async handleTokenCountRequest(rawRequest: any, context: RequestContext): Promise<Response> {
+    const log = {
+      debug: (message: string, metadata?: Record<string, any>) => {
+        logger.debug(message, { requestId: context.requestId, domain: context.host, metadata })
+      },
+      info: (message: string, metadata?: Record<string, any>) => {
+        logger.info(message, { requestId: context.requestId, domain: context.host, metadata })
+      },
+      error: (message: string, error?: Error, metadata?: Record<string, any>) => {
+        logger.error(message, {
+          requestId: context.requestId,
+          domain: context.host,
+          error: error
+            ? {
+                message: error.message,
+                stack: error.stack,
+                code: (error as any).code,
+              }
+            : undefined,
+          metadata,
+        })
+      },
+    }
+
+    try {
+      // Authenticate
+      let auth: AuthResult
+
+      // Passthrough mode when client auth disabled and Bearer token present
+      if (config.features.enableClientAuth === false && context.apiKey?.startsWith('Bearer ')) {
+        log.debug('Using passthrough authentication for token count (client auth disabled)', {
+          domain: context.host,
+          hasToken: true,
+        })
+
+        auth = {
+          type: 'oauth',
+          headers: {
+            Authorization: context.apiKey,
+            'anthropic-beta': 'oauth-2025-04-20',
+          },
+          key: context.apiKey.replace('Bearer ', ''),
+          betaHeader: 'oauth-2025-04-20',
+        }
+      } else {
+        // Existing domain-based routing
+        auth = context.host.toLowerCase().includes('personal')
+          ? await this.authService.authenticatePersonalDomain(context)
+          : await this.authService.authenticateNonPersonalDomain(context)
+
+        log.debug('Authentication successful for token count', {
+          authType: auth.type,
+          betaHeader: auth.betaHeader,
+        })
+      }
+
+      // Forward the request directly to the count_tokens endpoint
+      const response = await this.apiClient.forwardTokenCount(rawRequest, auth, context.requestId)
+
+      log.info('Token count request completed', {
+        model: rawRequest.model,
+        messageCount: rawRequest.messages?.length || 0,
+      })
+
+      return response
+    } catch (error) {
+      log.error('Token count request failed', error as Error)
+      throw error
+    }
+  }
 }