Merge pull request #180 from tickernelz/fix-kiro-token-count

refactor(claude-kiro): improve token counting accuracy using API contextUsagePercentage

src/claude/claude-kiro.js

@@ -10,6 +10,7 @@ import { getProviderModels } from '../provider-models.js';
 import { countTokens } from '@anthropic-ai/tokenizer';
 import { configureAxiosProxy } from '../proxy-utils.js';
 import { isRetryableNetworkError } from '../common.js';
+import { CLAUDE_DEFAULT_MAX_TOKENS } from '../converters/utils.js';
 
 const KIRO_CONSTANTS = {
     REFRESH_URL: 'https://prod.{{region}}.auth.desktop.kiro.dev/refreshToken',
@@ -1102,7 +1103,6 @@ async initializeAuth(forceRefresh = false) {
 
     _processApiResponse(response) {
         const rawResponseText = Buffer.isBuffer(response.data) ? response.data.toString('utf8') : String(response.data);
-        //console.log(`[Kiro] Raw response length: ${rawResponseText.length}`);
         if (rawResponseText.includes("[Called")) {
             console.log("[Kiro] Raw response contains [Called marker.");
         }
@@ -1154,13 +1154,22 @@ async initializeAuth(forceRefresh = false) {
         const finalModel = MODEL_MAPPING[model] ? model : this.modelName;
         console.log(`[Kiro] Calling generateContent with model: ${finalModel}`);
 
-        // Estimate input tokens before making the API call
-        const inputTokens = this.estimateInputTokens(requestBody);
-        
         const response = await this.callApi('', finalModel, requestBody);
 
         try {
             const { responseText, toolCalls } = this._processApiResponse(response);
+
+            let inputTokens = 0;
+            const rawResponseText = Buffer.isBuffer(response.data)
+                ? response.data.toString('utf8')
+                : String(response.data);
+
+            const contextUsageMatch = rawResponseText.match(/"contextUsagePercentage":\s*([\d.]+)/);
+            if (contextUsageMatch) {
+                const percentage = parseFloat(contextUsageMatch[1]);
+                inputTokens = this.calculateInputTokensFromPercentage(percentage);
+            }
+
             return this.buildClaudeResponse(responseText, false, 'assistant', model, toolCalls, inputTokens);
         } catch (error) {
             console.error('[Kiro] Error in generateContent:', error);
@@ -1192,9 +1201,10 @@ async initializeAuth(forceRefresh = false) {
             const followupStart = remaining.indexOf('{"followupPrompt":', searchStart);
             const inputStart = remaining.indexOf('{"input":', searchStart);
             const stopStart = remaining.indexOf('{"stop":', searchStart);
+            const contextUsageStart = remaining.indexOf('{"contextUsagePercentage":', searchStart);
 
             // 找到最早出现的有效 JSON 模式
-            const candidates = [contentStart, nameStart, followupStart, inputStart, stopStart].filter(pos => pos >= 0);
+            const candidates = [contentStart, nameStart, followupStart, inputStart, stopStart, contextUsageStart].filter(pos => pos >= 0);
             if (candidates.length === 0) break;
 
             const jsonStart = Math.min(...candidates);
@@ -1284,6 +1294,15 @@ async initializeAuth(forceRefresh = false) {
                         }
                     });
                 }
+                // 处理 context usage percentage 事件
+                else if (parsed.contextUsagePercentage !== undefined) {
+                    events.push({
+                        type: 'contextUsage',
+                        data: {
+                            percentage: parsed.contextUsagePercentage
+                        }
+                    });
+                }
             } catch (e) {
                 // JSON 解析失败，跳过这个位置继续搜索
             }
@@ -1330,7 +1349,7 @@ async initializeAuth(forceRefresh = false) {
 
             stream = response.data;
             let buffer = '';
-            let lastContentEvent = null;  // 用于检测连续重复的 content 事件
+            let lastContentEvent = null;
 
             for await (const chunk of stream) {
                 buffer += chunk.toString();
@@ -1355,6 +1374,8 @@ async initializeAuth(forceRefresh = false) {
                         yield { type: 'toolUseInput', input: event.data.input };
                     } else if (event.type === 'toolUseStop') {
                         yield { type: 'toolUseStop', stop: event.data.stop };
+                    } else if (event.type === 'contextUsage') {
+                        yield { type: 'contextUsage', percentage: event.data.percentage };
                     }
                 }
             }
@@ -1438,11 +1459,25 @@ async initializeAuth(forceRefresh = false) {
         const finalModel = MODEL_MAPPING[model] ? model : this.modelName;
         console.log(`[Kiro] Calling generateContentStream with model: ${finalModel} (real streaming)`);
 
-        const inputTokens = this.estimateInputTokens(requestBody);
+        let inputTokens = 0;
+        let contextUsagePercentage = null;
         const messageId = `${uuidv4()}`;
 
+        let messageStartSent = false;
+        const bufferedEvents = [];
+
         try {
-            // 1. 先发送 message_start 事件
+            let totalContent = '';
+            let outputTokens = 0;
+            const toolCalls = [];
+            let currentToolCall = null;
+
+            for await (const event of this.streamApiReal('', finalModel, requestBody)) {
+                if (event.type === 'contextUsage' && event.percentage) {
+                    contextUsagePercentage = event.percentage;
+                    inputTokens = this.calculateInputTokensFromPercentage(contextUsagePercentage);
+
+                    if (!messageStartSent) {
                         yield {
                             type: "message_start",
                             message: {
@@ -1450,34 +1485,43 @@ async initializeAuth(forceRefresh = false) {
                                 type: "message",
                                 role: "assistant",
                                 model: model,
-                    usage: { input_tokens: inputTokens, output_tokens: 0 },
+                                usage: {
+                                    input_tokens: inputTokens,
+                                    output_tokens: 0,
+                                    cache_creation_input_tokens: 0,
+                                    cache_read_input_tokens: 0
+                                },
                                 content: []
                             }
                         };
 
-            // 2. 发送 content_block_start 事件
                         yield {
                             type: "content_block_start",
                             index: 0,
                             content_block: { type: "text", text: "" }
                         };
 
-            let totalContent = '';
+                        messageStartSent = true;
-            let outputTokens = 0;
-            const toolCalls = [];
-            let currentToolCall = null;  // 用于累积结构化工具调用
 
-            // 3. 流式接收并发送每个 content_block_delta
+                        for (const buffered of bufferedEvents) {
-            for await (const event of this.streamApiReal('', finalModel, requestBody)) {
+                            yield buffered;
-                if (event.type === 'content' && event.content) {
+                        }
+                        bufferedEvents.length = 0;
+                    }
+                } else if (event.type === 'content' && event.content) {
                     totalContent += event.content;
-                    // 不再每个 chunk 都计算 token，改为最后统一计算，避免阻塞事件循环
 
-                    yield {
+                    const contentEvent = {
                         type: "content_block_delta",
                         index: 0,
                         delta: { type: "text_delta", text: event.content }
                     };
+
+                    if (messageStartSent) {
+                        yield contentEvent;
+                    } else {
+                        bufferedEvents.push(contentEvent);
+                    }
                 } else if (event.type === 'toolUse') {
                     const tc = event.toolUse;
                     // 工具调用事件（包含 name 和 toolUseId）
@@ -1541,6 +1585,12 @@ async initializeAuth(forceRefresh = false) {
                 currentToolCall = null;
             }
 
+            // Fallback: 如果 contextUsagePercentage 没有收到，抛出错误
+            if (!messageStartSent) {
+                console.error('[Kiro Stream] contextUsagePercentage not received from API - cannot calculate accurate input tokens');
+                throw new Error('Failed to receive contextUsagePercentage from Kiro API. Input token calculation requires this data.');
+            }
+
             // 检查文本内容中的 bracket 格式工具调用
             const bracketToolCalls = parseBracketToolCalls(totalContent);
             if (bracketToolCalls && bracketToolCalls.length > 0) {
@@ -1595,8 +1645,16 @@ async initializeAuth(forceRefresh = false) {
 
             yield {
                 type: "message_delta",
-                delta: { stop_reason: toolCalls.length > 0 ? "tool_use" : "end_turn" },
+                delta: {
-                usage: { output_tokens: outputTokens }
+                    stop_reason: toolCalls.length > 0 ? "tool_use" : "end_turn",
+                    stop_sequence: null
+                },
+                usage: {
+                    input_tokens: inputTokens,
+                    output_tokens: outputTokens,
+                    cache_creation_input_tokens: 0,
+                    cache_read_input_tokens: 0
+                }
             };
 
             // 7. 发送 message_stop 事件
@@ -1623,9 +1681,27 @@ async initializeAuth(forceRefresh = false) {
     }
 
     /**
+     * Convert context usage percentage to actual input tokens
+     * @param {number} percentage - Context usage percentage (0-100)
+     * @returns {number} Actual input tokens
+     */
+    calculateInputTokensFromPercentage(percentage) {
+        if (!percentage || percentage <= 0) {
+            return 0;
+        }
+
+        const contextWindow = CLAUDE_DEFAULT_MAX_TOKENS;
+        const inputTokens = Math.round((percentage / 100) * contextWindow);
+
+        return inputTokens;
+    }
+
+    /**
+     * @deprecated Use contextUsagePercentage from API response instead
      * Calculate input tokens from request body using Claude's official tokenizer
      */
     estimateInputTokens(requestBody) {
+        console.warn('[Kiro] estimateInputTokens() is deprecated. Use contextUsagePercentage from API response instead.');
         let totalTokens = 0;
 
         // Count system prompt tokens
@@ -1824,7 +1900,9 @@ async initializeAuth(forceRefresh = false) {
                 stop_sequence: null,
                 usage: {
                     input_tokens: inputTokens,
-                    output_tokens: outputTokens
+                    output_tokens: outputTokens,
+                    cache_creation_input_tokens: 0,
+                    cache_read_input_tokens: 0
                 },
                 content: contentArray
             };