diff --git a/package-lock.json b/package-lock.json index 79b762c..2d5d2be 100644 --- a/package-lock.json +++ b/package-lock.json @@ -9,6 +9,8 @@ "version": "1.2.6", "license": "MIT", "dependencies": { + "@anthropic-ai/tokenizer": "^0.0.4", + "@lenml/tokenizer-gemini": "^3.7.2", "async-mutex": "^0.5.0", "better-sqlite3": "^12.5.0", "cors": "^2.8.5", @@ -42,6 +44,16 @@ "url": "https://github.com/sponsors/sindresorhus" } }, + "node_modules/@anthropic-ai/tokenizer": { + "version": "0.0.4", + "resolved": "https://registry.npmjs.org/@anthropic-ai/tokenizer/-/tokenizer-0.0.4.tgz", + "integrity": "sha512-EHRKbxlxlc8W4KCBEseByJ7YwyYCmgu9OyN59H9+IYIGPoKv8tXyQXinkeGDI+cI8Tiuz9wk2jZb/kK7AyvL7g==", + "license": "Apache-2.0", + "dependencies": { + "@types/node": "^18.11.18", + "tiktoken": "^1.0.10" + } + }, "node_modules/@babel/runtime": { "version": "7.28.4", "resolved": "https://registry.npmjs.org/@babel/runtime/-/runtime-7.28.4.tgz", @@ -91,6 +103,21 @@ "@jridgewell/sourcemap-codec": "^1.4.14" } }, + "node_modules/@lenml/tokenizer-gemini": { + "version": "3.7.2", + "resolved": "https://registry.npmjs.org/@lenml/tokenizer-gemini/-/tokenizer-gemini-3.7.2.tgz", + "integrity": "sha512-sdSfXqjGSZWRHtf4toMcjzpBm/tOPPAtUQ5arTx4neQ2nzHUtJQJyHkoiB9KRyEfvVjW6WtQU+WbvU9glsFT2g==", + "license": "Apache-2.0", + "dependencies": { + "@lenml/tokenizers": "^3.7.2" + } + }, + "node_modules/@lenml/tokenizers": { + "version": "3.7.2", + "resolved": "https://registry.npmjs.org/@lenml/tokenizers/-/tokenizers-3.7.2.tgz", + "integrity": "sha512-tuap9T7Q80Czor8NHzxjlLNvxEX8MgFINzsBTV+lq1v7G+78YR3ZvBhmLsPHtgqExB4Q4kCJH6dhXOYWSLdHLw==", + "license": "Apache-2.0" + }, "node_modules/@nodelib/fs.scandir": { "version": "2.1.5", "resolved": "https://registry.npmjs.org/@nodelib/fs.scandir/-/fs.scandir-2.1.5.tgz", @@ -142,6 +169,15 @@ "tailwindcss": ">=3.0.0 || >= 3.0.0-alpha.1 || >= 4.0.0-alpha.20 || >= 4.0.0-beta.1" } }, + "node_modules/@types/node": { + "version": "18.19.130", + "resolved": "https://registry.npmjs.org/@types/node/-/node-18.19.130.tgz", + "integrity": "sha512-GRaXQx6jGfL8sKfaIDD6OupbIHBr9jv7Jnaml9tB7l4v068PAOXqfcujMMo5PhbIs6ggR1XODELqahT2R8v0fg==", + "license": "MIT", + "dependencies": { + "undici-types": "~5.26.4" + } + }, "node_modules/accepts": { "version": "1.3.8", "resolved": "https://registry.npmjs.org/accepts/-/accepts-1.3.8.tgz", @@ -395,7 +431,6 @@ } ], "license": "MIT", - "peer": true, "dependencies": { "baseline-browser-mapping": "^2.9.0", "caniuse-lite": "^1.0.30001759", @@ -1416,7 +1451,6 @@ "integrity": "sha512-/imKNG4EbWNrVjoNC/1H5/9GFy+tqjGBHCaSsN+P2RnPqjsLmv6UD3Ej+Kj8nBWaRAwyk7kK5ZUc+OEatnTR3A==", "dev": true, "license": "MIT", - "peer": true, "bin": { "jiti": "bin/jiti.js" } @@ -1793,7 +1827,6 @@ } ], "license": "MIT", - "peer": true, "dependencies": { "nanoid": "^3.3.11", "picocolors": "^1.1.1", @@ -2611,7 +2644,6 @@ "integrity": "sha512-3ofp+LL8E+pK/JuPLPggVAIaEuhvIz4qNcf3nA1Xn2o/7fb7s/TYpHhwGDv1ZU3PkBluUVaF8PyCHcm48cKLWQ==", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "@alloc/quick-lru": "^5.2.0", "arg": "^5.0.2", @@ -2695,6 +2727,12 @@ "node": ">=0.8" } }, + "node_modules/tiktoken": { + "version": "1.0.22", + "resolved": "https://registry.npmjs.org/tiktoken/-/tiktoken-1.0.22.tgz", + "integrity": "sha512-PKvy1rVF1RibfF3JlXBSP0Jrcw2uq3yXdgcEXtKTYn3QJ/cBRBHDnrJ5jHky+MENZ6DIPwNUGWpkVx+7joCpNA==", + "license": "MIT" + }, "node_modules/tinyglobby": { "version": "0.2.15", "resolved": "https://registry.npmjs.org/tinyglobby/-/tinyglobby-0.2.15.tgz", @@ -2736,7 +2774,6 @@ "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==", "dev": true, "license": "MIT", - "peer": true, "engines": { "node": ">=12" }, @@ -2814,6 +2851,12 @@ "node": ">= 0.6" } }, + "node_modules/undici-types": { + "version": "5.26.5", + "resolved": "https://registry.npmjs.org/undici-types/-/undici-types-5.26.5.tgz", + "integrity": "sha512-JlCMO+ehdEIKqlFxk6IfVoAUVmgz7cU7zD/h9XZ0qzeosSHmUJVOzSQvvYSYWXkFXC+IfLKSIffhv0sVZup6pA==", + "license": "MIT" + }, "node_modules/unpipe": { "version": "1.0.0", "resolved": "https://registry.npmjs.org/unpipe/-/unpipe-1.0.0.tgz", diff --git a/package.json b/package.json index dec5b4f..563fa77 100644 --- a/package.json +++ b/package.json @@ -34,7 +34,8 @@ "test:crossmodel": "node tests/test-cross-model-thinking.cjs", "test:oauth": "node tests/test-oauth-no-browser.cjs", "test:emptyretry": "node tests/test-empty-response-retry.cjs", - "test:sanitizer": "node tests/test-schema-sanitizer.cjs" + "test:sanitizer": "node tests/test-schema-sanitizer.cjs", + "test:counttokens": "node tests/test-count-tokens.cjs" }, "keywords": [ "claude", @@ -57,6 +58,8 @@ "node": ">=18.0.0" }, "dependencies": { + "@anthropic-ai/tokenizer": "^0.0.4", + "@lenml/tokenizer-gemini": "^3.7.2", "async-mutex": "^0.5.0", "better-sqlite3": "^12.5.0", "cors": "^2.8.5", diff --git a/src/cloudcode/count-tokens.js b/src/cloudcode/count-tokens.js new file mode 100644 index 0000000..be12a70 --- /dev/null +++ b/src/cloudcode/count-tokens.js @@ -0,0 +1,302 @@ +/** + * Token Counter Implementation for antigravity-claude-proxy + * + * Implements Anthropic's /v1/messages/count_tokens endpoint + * Uses official tokenizers for each model family: + * - Claude: @anthropic-ai/tokenizer + * - Gemini: @lenml/tokenizer-gemini + * + * @see https://platform.claude.com/docs/en/api/messages-count-tokens + */ + +import { countTokens as claudeCountTokens } from '@anthropic-ai/tokenizer'; +import { fromPreTrained as loadGeminiTokenizer } from '@lenml/tokenizer-gemini'; +import { logger } from '../utils/logger.js'; +import { getModelFamily } from '../constants.js'; + +// Lazy-loaded Gemini tokenizer (138MB, loaded once on first use) +let geminiTokenizer = null; +let geminiTokenizerLoading = null; + +/** + * Get or initialize the Gemini tokenizer + * Uses singleton pattern with loading lock to prevent multiple loads + * + * @returns {Promise} Gemini tokenizer instance + */ +async function getGeminiTokenizer() { + if (geminiTokenizer) { + return geminiTokenizer; + } + + // Prevent multiple simultaneous loads + if (geminiTokenizerLoading) { + return geminiTokenizerLoading; + } + + geminiTokenizerLoading = (async () => { + try { + logger.debug('[TokenCounter] Loading Gemini tokenizer...'); + geminiTokenizer = await loadGeminiTokenizer(); + logger.debug('[TokenCounter] Gemini tokenizer loaded successfully'); + return geminiTokenizer; + } catch (error) { + logger.warn(`[TokenCounter] Failed to load Gemini tokenizer: ${error.message}`); + throw error; + } finally { + geminiTokenizerLoading = null; + } + })(); + + return geminiTokenizerLoading; +} + +/** + * Count tokens for text using Claude tokenizer + * + * @param {string} text - Text to tokenize + * @returns {number} Token count + */ +function countClaudeTokens(text) { + if (!text) return 0; + try { + return claudeCountTokens(text); + } catch (error) { + logger.debug(`[TokenCounter] Claude tokenizer error: ${error.message}`); + return Math.ceil(text.length / 4); + } +} + +/** + * Count tokens for text using Gemini tokenizer + * + * @param {Object} tokenizer - Gemini tokenizer instance + * @param {string} text - Text to tokenize + * @returns {number} Token count + */ +function countGeminiTokens(tokenizer, text) { + if (!text) return 0; + try { + const tokens = tokenizer.encode(text); + // Remove BOS token if present (token id 2) + return tokens[0] === 2 ? tokens.length - 1 : tokens.length; + } catch (error) { + logger.debug(`[TokenCounter] Gemini tokenizer error: ${error.message}`); + return Math.ceil(text.length / 4); + } +} + +/** + * Estimate tokens for text content using appropriate tokenizer + * + * @param {string} text - Text to tokenize + * @param {string} model - Model name to determine tokenizer + * @param {Object} geminiTok - Gemini tokenizer instance (optional) + * @returns {number} Token count + */ +function estimateTextTokens(text, model, geminiTok = null) { + if (!text) return 0; + + const family = getModelFamily(model); + + if (family === 'claude') { + return countClaudeTokens(text); + } else if (family === 'gemini' && geminiTok) { + return countGeminiTokens(geminiTok, text); + } + + // Fallback for unknown models: rough estimate + return Math.ceil(text.length / 4); +} + +/** + * Extract text from message content + * + * Note: This function only extracts text from 'text' type blocks. + * Image blocks (type: 'image') and document blocks (type: 'document') are not tokenized + * and will not contribute to the token count. This is intentional as binary content + * requires different handling and Anthropic's actual token counting for images uses + * a fixed estimate (~1600 tokens per image) that depends on image dimensions. + * + * @param {string|Array} content - Message content + * @returns {string} Concatenated text + */ +function extractText(content) { + if (typeof content === 'string') { + return content; + } + + if (Array.isArray(content)) { + return content + .filter(block => block.type === 'text') + .map(block => block.text) + .join('\n'); + } + + return ''; +} + +/** + * Count tokens locally using model-specific tokenizer + * + * @param {Object} request - Anthropic format request + * @param {Object} geminiTok - Gemini tokenizer instance (optional) + * @returns {number} Token count + */ +function countTokensLocally(request, geminiTok = null) { + const { messages = [], system, tools, model } = request; + let totalTokens = 0; + + // Count system prompt tokens + if (system) { + if (typeof system === 'string') { + totalTokens += estimateTextTokens(system, model, geminiTok); + } else if (Array.isArray(system)) { + for (const block of system) { + if (block.type === 'text') { + totalTokens += estimateTextTokens(block.text, model, geminiTok); + } + } + } + } + + // Count message tokens + for (const message of messages) { + // Add overhead for role and structure (~4 tokens per message) + totalTokens += 4; + totalTokens += estimateTextTokens(extractText(message.content), model, geminiTok); + + // Handle tool_use and tool_result blocks + if (Array.isArray(message.content)) { + for (const block of message.content) { + if (block.type === 'tool_use') { + totalTokens += estimateTextTokens(block.name, model, geminiTok); + totalTokens += estimateTextTokens(JSON.stringify(block.input), model, geminiTok); + } else if (block.type === 'tool_result') { + if (typeof block.content === 'string') { + totalTokens += estimateTextTokens(block.content, model, geminiTok); + } else if (Array.isArray(block.content)) { + totalTokens += estimateTextTokens(extractText(block.content), model, geminiTok); + } + } else if (block.type === 'thinking') { + totalTokens += estimateTextTokens(block.thinking, model, geminiTok); + } + } + } + } + + // Count tool definitions + if (tools && tools.length > 0) { + for (const tool of tools) { + totalTokens += estimateTextTokens(tool.name, model, geminiTok); + totalTokens += estimateTextTokens(tool.description || '', model, geminiTok); + totalTokens += estimateTextTokens(JSON.stringify(tool.input_schema || {}), model, geminiTok); + } + } + + return totalTokens; +} + +/** + * Count tokens in a message request + * Implements Anthropic's /v1/messages/count_tokens endpoint + * Uses local tokenization for all content types + * + * @param {Object} anthropicRequest - Anthropic format request with messages, model, system, tools + * @param {Object} accountManager - Account manager instance (unused, kept for API compatibility) + * @param {Object} options - Options (unused, kept for API compatibility) + * @returns {Promise} Response with input_tokens count + */ +export async function countTokens(anthropicRequest, accountManager = null, options = {}) { + try { + const family = getModelFamily(anthropicRequest.model); + let geminiTok = null; + + // Load Gemini tokenizer if needed + if (family === 'gemini') { + try { + geminiTok = await getGeminiTokenizer(); + } catch (error) { + logger.warn(`[TokenCounter] Gemini tokenizer unavailable, using fallback`); + } + } + + const inputTokens = countTokensLocally(anthropicRequest, geminiTok); + logger.debug(`[TokenCounter] Local count (${family}): ${inputTokens} tokens`); + + return { + input_tokens: inputTokens + }; + + } catch (error) { + logger.warn(`[TokenCounter] Error: ${error.message}, using character-based fallback`); + + // Ultimate fallback: character-based estimation + const { messages = [], system } = anthropicRequest; + let charCount = 0; + + if (system) { + charCount += typeof system === 'string' ? system.length : JSON.stringify(system).length; + } + + for (const message of messages) { + charCount += JSON.stringify(message.content).length; + } + + return { + input_tokens: Math.ceil(charCount / 4) + }; + } +} + +/** + * Express route handler for /v1/messages/count_tokens + * + * @param {Object} accountManager - Account manager instance + * @returns {Function} Express middleware + */ +export function createCountTokensHandler(accountManager) { + return async (req, res) => { + try { + const { messages, model, system, tools, tool_choice, thinking } = req.body; + + // Validate required fields + if (!messages || !Array.isArray(messages)) { + return res.status(400).json({ + type: 'error', + error: { + type: 'invalid_request_error', + message: 'messages is required and must be an array' + } + }); + } + + if (!model) { + return res.status(400).json({ + type: 'error', + error: { + type: 'invalid_request_error', + message: 'model is required' + } + }); + } + + const result = await countTokens( + { messages, model, system, tools, tool_choice, thinking }, + accountManager + ); + + res.json(result); + + } catch (error) { + logger.error(`[TokenCounter] Handler error: ${error.message}`); + res.status(500).json({ + type: 'error', + error: { + type: 'api_error', + message: error.message + } + }); + } + }; +} diff --git a/src/server.js b/src/server.js index be2c3af..86c92a2 100644 --- a/src/server.js +++ b/src/server.js @@ -9,6 +9,7 @@ import cors from 'cors'; import path from 'path'; import { fileURLToPath } from 'url'; import { sendMessage, sendMessageStream, listModels, getModelQuotas, getSubscriptionTier } from './cloudcode/index.js'; +import { createCountTokensHandler } from './cloudcode/count-tokens.js'; import { mountWebUI } from './webui/index.js'; import { config } from './config.js'; @@ -597,16 +598,19 @@ app.get('/v1/models', async (req, res) => { }); /** - * Count tokens endpoint (not supported) + * Count tokens endpoint - Anthropic Messages API compatible + * Uses local tokenization with official tokenizers (@anthropic-ai/tokenizer for Claude, @lenml/tokenizer-gemini for Gemini) */ -app.post('/v1/messages/count_tokens', (req, res) => { - res.status(501).json({ - type: 'error', - error: { - type: 'not_implemented', - message: 'Token counting is not implemented. Use /v1/messages with max_tokens or configure your client to skip token counting.' - } - }); +app.post('/v1/messages/count_tokens', async (req, res) => { + try { + // Ensure account manager is initialized for API-based counting + await ensureInitialized(); + } catch (error) { + // If initialization fails, handler will fall back to local estimation + logger.debug(`[TokenCounter] Account manager not initialized: ${error.message}`); + } + + return createCountTokensHandler(accountManager)(req, res); }); /** diff --git a/tests/run-all.cjs b/tests/run-all.cjs index 95015d3..8c1dc2f 100644 --- a/tests/run-all.cjs +++ b/tests/run-all.cjs @@ -18,7 +18,8 @@ const tests = [ { name: 'Cross-Model Thinking', file: 'test-cross-model-thinking.cjs' }, { name: 'OAuth No-Browser Mode', file: 'test-oauth-no-browser.cjs' }, { name: 'Empty Response Retry', file: 'test-empty-response-retry.cjs' }, - { name: 'Schema Sanitizer', file: 'test-schema-sanitizer.cjs' } + { name: 'Schema Sanitizer', file: 'test-schema-sanitizer.cjs' }, + { name: 'Count Tokens', file: 'test-count-tokens.cjs' } ]; async function runTest(test) { diff --git a/tests/test-count-tokens.cjs b/tests/test-count-tokens.cjs new file mode 100644 index 0000000..325a72d --- /dev/null +++ b/tests/test-count-tokens.cjs @@ -0,0 +1,492 @@ +/** + * Test Count Tokens - Tests for the /v1/messages/count_tokens endpoint + * + * Verifies token counting functionality: + * - Local estimation using official tokenizers (@anthropic-ai/tokenizer for Claude, @lenml/tokenizer-gemini for Gemini) + * - Request validation + * - Different content types (text, tools, system prompts) + */ +const http = require('http'); + +// Server configuration +const BASE_URL = 'localhost'; +const PORT = 8080; + +/** + * Make a request to the count_tokens endpoint + * @param {Object} body - Request body + * @returns {Promise} - Parsed JSON response with statusCode + */ +function countTokensRequest(body) { + return new Promise((resolve, reject) => { + const data = JSON.stringify(body); + const req = http.request({ + host: BASE_URL, + port: PORT, + path: '/v1/messages/count_tokens', + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'x-api-key': 'test', + 'anthropic-version': '2023-06-01', + 'Content-Length': Buffer.byteLength(data) + } + }, res => { + let fullData = ''; + res.on('data', chunk => fullData += chunk.toString()); + res.on('end', () => { + try { + const parsed = JSON.parse(fullData); + resolve({ ...parsed, statusCode: res.statusCode }); + } catch (e) { + reject(new Error(`Parse error: ${e.message}\nRaw: ${fullData.substring(0, 500)}`)); + } + }); + }); + req.on('error', reject); + req.write(data); + req.end(); + }); +} + +async function runTests() { + console.log('╔══════════════════════════════════════════════════════════════╗'); + console.log('║ COUNT TOKENS ENDPOINT TEST SUITE ║'); + console.log('╚══════════════════════════════════════════════════════════════╝\n'); + + let passed = 0; + let failed = 0; + + function test(name, fn) { + return fn() + .then(() => { + console.log(`✓ ${name}`); + passed++; + }) + .catch(e => { + console.log(`✗ ${name}`); + console.log(` Error: ${e.message}`); + failed++; + }); + } + + function assert(condition, message) { + if (!condition) throw new Error(message); + } + + function assertType(value, type, name) { + if (typeof value !== type) { + throw new Error(`${name} should be ${type}, got ${typeof value}`); + } + } + + function assertGreater(value, min, name) { + if (value <= min) { + throw new Error(`${name} should be greater than ${min}, got ${value}`); + } + } + + // Test 1: Simple text message + await test('Simple text message returns token count', async () => { + const response = await countTokensRequest({ + model: 'claude-sonnet-4-5', + messages: [ + { role: 'user', content: 'Hello, how are you?' } + ] + }); + + assert(response.statusCode === 200, `Expected 200, got ${response.statusCode}`); + assertType(response.input_tokens, 'number', 'input_tokens'); + assertGreater(response.input_tokens, 0, 'input_tokens'); + }); + + // Test 2: Multi-turn conversation + await test('Multi-turn conversation counts all messages', async () => { + const response = await countTokensRequest({ + model: 'claude-sonnet-4-5', + messages: [ + { role: 'user', content: 'What is the capital of France?' }, + { role: 'assistant', content: 'The capital of France is Paris.' }, + { role: 'user', content: 'And what about Germany?' } + ] + }); + + assert(response.statusCode === 200, `Expected 200, got ${response.statusCode}`); + assertType(response.input_tokens, 'number', 'input_tokens'); + // Multi-turn should have more tokens than single message + assertGreater(response.input_tokens, 10, 'input_tokens for multi-turn'); + }); + + // Test 3: System prompt + await test('System prompt tokens are counted', async () => { + const responseWithSystem = await countTokensRequest({ + model: 'claude-sonnet-4-5', + system: 'You are a helpful assistant that speaks like a pirate.', + messages: [ + { role: 'user', content: 'Hello' } + ] + }); + + const responseWithoutSystem = await countTokensRequest({ + model: 'claude-sonnet-4-5', + messages: [ + { role: 'user', content: 'Hello' } + ] + }); + + assert(responseWithSystem.statusCode === 200, `Expected 200, got ${responseWithSystem.statusCode}`); + // With system prompt should have more tokens + assertGreater(responseWithSystem.input_tokens, responseWithoutSystem.input_tokens, + 'tokens with system prompt'); + }); + + // Test 4: System prompt as array + await test('System prompt as array is counted', async () => { + const response = await countTokensRequest({ + model: 'claude-sonnet-4-5', + system: [ + { type: 'text', text: 'You are a helpful assistant.' }, + { type: 'text', text: 'Be concise and clear.' } + ], + messages: [ + { role: 'user', content: 'Hello' } + ] + }); + + assert(response.statusCode === 200, `Expected 200, got ${response.statusCode}`); + assertType(response.input_tokens, 'number', 'input_tokens'); + assertGreater(response.input_tokens, 5, 'input_tokens'); + }); + + // Test 5: With tools + await test('Tool definitions are counted', async () => { + const responseWithTools = await countTokensRequest({ + model: 'claude-sonnet-4-5', + messages: [ + { role: 'user', content: 'Get the weather in Tokyo' } + ], + tools: [ + { + name: 'get_weather', + description: 'Get the current weather for a location', + input_schema: { + type: 'object', + properties: { + location: { type: 'string', description: 'City name' } + }, + required: ['location'] + } + } + ] + }); + + const responseWithoutTools = await countTokensRequest({ + model: 'claude-sonnet-4-5', + messages: [ + { role: 'user', content: 'Get the weather in Tokyo' } + ] + }); + + assert(responseWithTools.statusCode === 200, `Expected 200, got ${responseWithTools.statusCode}`); + // With tools should have more tokens + assertGreater(responseWithTools.input_tokens, responseWithoutTools.input_tokens, + 'tokens with tools'); + }); + + // Test 6: Content as array with text blocks + await test('Content array with text blocks', async () => { + const response = await countTokensRequest({ + model: 'claude-sonnet-4-5', + messages: [ + { + role: 'user', + content: [ + { type: 'text', text: 'First part of the message.' }, + { type: 'text', text: 'Second part of the message.' } + ] + } + ] + }); + + assert(response.statusCode === 200, `Expected 200, got ${response.statusCode}`); + assertType(response.input_tokens, 'number', 'input_tokens'); + assertGreater(response.input_tokens, 5, 'input_tokens'); + }); + + // Test 7: Tool use and tool result blocks + await test('Tool use and tool result blocks are counted', async () => { + const response = await countTokensRequest({ + model: 'claude-sonnet-4-5', + messages: [ + { role: 'user', content: 'What is the weather in Paris?' }, + { + role: 'assistant', + content: [ + { + type: 'tool_use', + id: 'tool_123', + name: 'get_weather', + input: { location: 'Paris' } + } + ] + }, + { + role: 'user', + content: [ + { + type: 'tool_result', + tool_use_id: 'tool_123', + content: 'The weather in Paris is sunny with 22°C' + } + ] + } + ], + tools: [ + { + name: 'get_weather', + description: 'Get weather for a location', + input_schema: { + type: 'object', + properties: { + location: { type: 'string' } + } + } + } + ] + }); + + assert(response.statusCode === 200, `Expected 200, got ${response.statusCode}`); + assertType(response.input_tokens, 'number', 'input_tokens'); + assertGreater(response.input_tokens, 20, 'input_tokens for tool conversation'); + }); + + // Test 8: Thinking blocks + await test('Thinking blocks are counted', async () => { + const response = await countTokensRequest({ + model: 'claude-sonnet-4-5', + messages: [ + { role: 'user', content: 'Solve this problem step by step' }, + { + role: 'assistant', + content: [ + { + type: 'thinking', + thinking: 'Let me think about this problem carefully. First, I need to understand what is being asked...' + }, + { type: 'text', text: 'Here is my solution.' } + ] + }, + { role: 'user', content: 'Can you explain further?' } + ] + }); + + assert(response.statusCode === 200, `Expected 200, got ${response.statusCode}`); + assertType(response.input_tokens, 'number', 'input_tokens'); + assertGreater(response.input_tokens, 20, 'input_tokens with thinking'); + }); + + // Test 9: Long text + await test('Long text message', async () => { + const longText = 'This is a test message. '.repeat(100); + const response = await countTokensRequest({ + model: 'claude-sonnet-4-5', + messages: [ + { role: 'user', content: longText } + ] + }); + + assert(response.statusCode === 200, `Expected 200, got ${response.statusCode}`); + assertType(response.input_tokens, 'number', 'input_tokens'); + // Long text should have many tokens + assertGreater(response.input_tokens, 100, 'input_tokens for long text'); + }); + + // Test 10: Missing messages field (error case) + await test('Missing messages returns error', async () => { + const response = await countTokensRequest({ + model: 'claude-sonnet-4-5' + }); + + assert(response.statusCode === 400, `Expected 400, got ${response.statusCode}`); + assert(response.type === 'error', 'Should return error type'); + assert(response.error.type === 'invalid_request_error', + `Expected invalid_request_error, got ${response.error?.type}`); + }); + + // Test 11: Missing model field (error case) + await test('Missing model returns error', async () => { + const response = await countTokensRequest({ + messages: [ + { role: 'user', content: 'Hello' } + ] + }); + + assert(response.statusCode === 400, `Expected 400, got ${response.statusCode}`); + assert(response.type === 'error', 'Should return error type'); + assert(response.error.type === 'invalid_request_error', + `Expected invalid_request_error, got ${response.error?.type}`); + }); + + // Test 12: Invalid messages type (error case) + await test('Invalid messages type returns error', async () => { + const response = await countTokensRequest({ + model: 'claude-sonnet-4-5', + messages: 'not an array' + }); + + assert(response.statusCode === 400, `Expected 400, got ${response.statusCode}`); + assert(response.type === 'error', 'Should return error type'); + }); + + // Test 13: Empty messages array + await test('Empty messages array returns token count', async () => { + const response = await countTokensRequest({ + model: 'claude-sonnet-4-5', + messages: [] + }); + + assert(response.statusCode === 200, `Expected 200, got ${response.statusCode}`); + assertType(response.input_tokens, 'number', 'input_tokens'); + }); + + // Test 14: Multiple tools with complex schemas + await test('Multiple tools with complex schemas', async () => { + const response = await countTokensRequest({ + model: 'claude-sonnet-4-5', + messages: [ + { role: 'user', content: 'Help me with file operations' } + ], + tools: [ + { + name: 'read_file', + description: 'Read a file from the filesystem', + input_schema: { + type: 'object', + properties: { + path: { type: 'string', description: 'Path to the file' }, + encoding: { type: 'string', description: 'File encoding' } + }, + required: ['path'] + } + }, + { + name: 'write_file', + description: 'Write content to a file', + input_schema: { + type: 'object', + properties: { + path: { type: 'string', description: 'Path to the file' }, + content: { type: 'string', description: 'Content to write' }, + append: { type: 'boolean', description: 'Append mode' } + }, + required: ['path', 'content'] + } + }, + { + name: 'list_directory', + description: 'List files in a directory', + input_schema: { + type: 'object', + properties: { + path: { type: 'string', description: 'Directory path' }, + recursive: { type: 'boolean', description: 'List recursively' } + }, + required: ['path'] + } + } + ] + }); + + assert(response.statusCode === 200, `Expected 200, got ${response.statusCode}`); + assertType(response.input_tokens, 'number', 'input_tokens'); + // Multiple tools should have significant token count + assertGreater(response.input_tokens, 50, 'input_tokens for multiple tools'); + }); + + // Test 15: Tool result as array content + await test('Tool result with array content', async () => { + const response = await countTokensRequest({ + model: 'claude-sonnet-4-5', + messages: [ + { role: 'user', content: 'Search for files' }, + { + role: 'assistant', + content: [ + { type: 'tool_use', id: 'tool_456', name: 'search', input: { query: 'test' } } + ] + }, + { + role: 'user', + content: [ + { + type: 'tool_result', + tool_use_id: 'tool_456', + content: [ + { type: 'text', text: 'Found file1.txt' }, + { type: 'text', text: 'Found file2.txt' } + ] + } + ] + } + ] + }); + + assert(response.statusCode === 200, `Expected 200, got ${response.statusCode}`); + assertType(response.input_tokens, 'number', 'input_tokens'); + assertGreater(response.input_tokens, 10, 'input_tokens'); + }); + + // Test 16: Gemini model token counting + await test('Gemini model returns token count', async () => { + const response = await countTokensRequest({ + model: 'gemini-3-flash', + messages: [ + { role: 'user', content: 'Hello, how are you?' } + ] + }); + + assert(response.statusCode === 200, `Expected 200, got ${response.statusCode}`); + assertType(response.input_tokens, 'number', 'input_tokens'); + assertGreater(response.input_tokens, 0, 'input_tokens'); + }); + + // Test 17: Gemini model with system prompt and tools + await test('Gemini model with system prompt and tools', async () => { + const response = await countTokensRequest({ + model: 'gemini-3-flash', + system: 'You are a helpful assistant.', + messages: [ + { role: 'user', content: 'What is the weather in Tokyo?' } + ], + tools: [ + { + name: 'get_weather', + description: 'Get weather for a location', + input_schema: { + type: 'object', + properties: { + location: { type: 'string' } + } + } + } + ] + }); + + assert(response.statusCode === 200, `Expected 200, got ${response.statusCode}`); + assertType(response.input_tokens, 'number', 'input_tokens'); + assertGreater(response.input_tokens, 10, 'input_tokens for Gemini with tools'); + }); + + // Summary + console.log('\n' + '═'.repeat(60)); + console.log(`Tests completed: ${passed} passed, ${failed} failed`); + + if (failed > 0) { + process.exit(1); + } +} + +runTests().catch(err => { + console.error('Test suite failed:', err); + process.exit(1); +});