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',
@@ -33,9 +34,9 @@ const KIRO_MODELS = getProviderModels('claude-kiro-oauth');
 
 // 完整的模型映射表
 const FULL_MODEL_MAPPING = {
-    "claude-opus-4-5":"claude-opus-4.5",
+    "claude-opus-4-5": "claude-opus-4.5",
-    "claude-opus-4-5-20251101":"claude-opus-4.5",
+    "claude-opus-4-5-20251101": "claude-opus-4.5",
-    "claude-haiku-4-5":"claude-haiku-4.5",
+    "claude-haiku-4-5": "claude-haiku-4.5",
     "claude-sonnet-4-5": "CLAUDE_SONNET_4_5_20250929_V1_0",
     "claude-sonnet-4-5-20250929": "CLAUDE_SONNET_4_5_20250929_V1_0",
     "claude-sonnet-4-20250514": "CLAUDE_SONNET_4_20250514_V1_0",
@@ -370,197 +371,197 @@ export class KiroApiService {
         this.isInitialized = true;
     }
 
-async initializeAuth(forceRefresh = false) {
+    async initializeAuth(forceRefresh = false) {
-    if (this.accessToken && !forceRefresh) {
+        if (this.accessToken && !forceRefresh) {
-        console.debug('[Kiro Auth] Access token already available and not forced refresh.');
+            console.debug('[Kiro Auth] Access token already available and not forced refresh.');
-        return;
+            return;
-    }
-
-    // Helper to load credentials from a file
-    const loadCredentialsFromFile = async (filePath) => {
-        try {
-            const fileContent = await fs.readFile(filePath, 'utf8');
-            return JSON.parse(fileContent);
-        } catch (error) {
-            if (error.code === 'ENOENT') {
-                console.debug(`[Kiro Auth] Credential file not found: ${filePath}`);
-            } else if (error instanceof SyntaxError) {
-                console.warn(`[Kiro Auth] Failed to parse JSON from ${filePath}: ${error.message}`);
-            } else {
-                console.warn(`[Kiro Auth] Failed to read credential file ${filePath}: ${error.message}`);
-            }
-            return null;
         }
-    };
 
-    // Helper to save credentials to a file
+        // Helper to load credentials from a file
-    const saveCredentialsToFile = async (filePath, newData) => {
+        const loadCredentialsFromFile = async (filePath) => {
-        try {
-            let existingData = {};
             try {
                 const fileContent = await fs.readFile(filePath, 'utf8');
-                existingData = JSON.parse(fileContent);
+                return JSON.parse(fileContent);
-            } catch (readError) {
+            } catch (error) {
-                if (readError.code === 'ENOENT') {
+                if (error.code === 'ENOENT') {
-                    console.debug(`[Kiro Auth] Token file not found, creating new one: ${filePath}`);
+                    console.debug(`[Kiro Auth] Credential file not found: ${filePath}`);
+                } else if (error instanceof SyntaxError) {
+                    console.warn(`[Kiro Auth] Failed to parse JSON from ${filePath}: ${error.message}`);
                 } else {
-                    console.warn(`[Kiro Auth] Could not read existing token file ${filePath}: ${readError.message}`);
+                    console.warn(`[Kiro Auth] Failed to read credential file ${filePath}: ${error.message}`);
                 }
+                return null;
             }
-            const mergedData = { ...existingData, ...newData };
+        };
-            await fs.writeFile(filePath, JSON.stringify(mergedData, null, 2), 'utf8');
-            console.info(`[Kiro Auth] Updated token file: ${filePath}`);
-        } catch (error) {
-            console.error(`[Kiro Auth] Failed to write token to file ${filePath}: ${error.message}`);
-        }
-    };
 
-    try {
+        // Helper to save credentials to a file
-        let mergedCredentials = {};
+        const saveCredentialsToFile = async (filePath, newData) => {
-
+            try {
-        // Priority 1: Load from Base64 credentials if available
+                let existingData = {};
-        if (this.base64Creds) {
+                try {
-            Object.assign(mergedCredentials, this.base64Creds);
+                    const fileContent = await fs.readFile(filePath, 'utf8');
-            console.info('[Kiro Auth] Successfully loaded credentials from Base64 (constructor).');
+                    existingData = JSON.parse(fileContent);
-            // Clear base64Creds after use to prevent re-processing
+                } catch (readError) {
-            this.base64Creds = null;
+                    if (readError.code === 'ENOENT') {
-        }
+                        console.debug(`[Kiro Auth] Token file not found, creating new one: ${filePath}`);
-
+                    } else {
-        // Priority 2 & 3 合并: 从指定文件路径或目录加载凭证
+                        console.warn(`[Kiro Auth] Could not read existing token file ${filePath}: ${readError.message}`);
-        // 读取指定的 credPath 文件以及目录下的其他 JSON 文件(排除当前文件)
-        const targetFilePath = this.credsFilePath || path.join(this.credPath, KIRO_AUTH_TOKEN_FILE);
-        const dirPath = path.dirname(targetFilePath);
-        const targetFileName = path.basename(targetFilePath);
-        
-        console.debug(`[Kiro Auth] Attempting to load credentials from directory: ${dirPath}`);
-        
-        try {
-            // 首先尝试读取目标文件
-            const targetCredentials = await loadCredentialsFromFile(targetFilePath);
-            if (targetCredentials) {
-                Object.assign(mergedCredentials, targetCredentials);
-                console.info(`[Kiro Auth] Successfully loaded OAuth credentials from ${targetFilePath}`);
-            }
-            
-            // 然后读取目录下的其他 JSON 文件(排除目标文件本身)
-            const files = await fs.readdir(dirPath);
-            for (const file of files) {
-                if (file.endsWith('.json') && file !== targetFileName) {
-                    const filePath = path.join(dirPath, file);
-                    const credentials = await loadCredentialsFromFile(filePath);
-                    if (credentials) {
-                        // 保留已有的 expiresAt,避免被覆盖
-                        credentials.expiresAt = mergedCredentials.expiresAt;
-                        Object.assign(mergedCredentials, credentials);
-                        console.debug(`[Kiro Auth] Loaded Client credentials from ${file}`);
                     }
                 }
+                const mergedData = { ...existingData, ...newData };
+                await fs.writeFile(filePath, JSON.stringify(mergedData, null, 2), 'utf8');
+                console.info(`[Kiro Auth] Updated token file: ${filePath}`);
+            } catch (error) {
+                console.error(`[Kiro Auth] Failed to write token to file ${filePath}: ${error.message}`);
             }
-        } catch (error) {
+        };
-            console.warn(`[Kiro Auth] Error loading credentials from directory ${dirPath}: ${error.message}`);
-        }
 
-        // console.log('[Kiro Auth] Merged credentials:', mergedCredentials);
-        // Apply loaded credentials, prioritizing existing values if they are not null/undefined
-        this.accessToken = this.accessToken || mergedCredentials.accessToken;
-        this.refreshToken = this.refreshToken || mergedCredentials.refreshToken;
-        this.clientId = this.clientId || mergedCredentials.clientId;
-        this.clientSecret = this.clientSecret || mergedCredentials.clientSecret;
-        this.authMethod = this.authMethod || mergedCredentials.authMethod;
-        this.expiresAt = this.expiresAt || mergedCredentials.expiresAt;
-        this.profileArn = this.profileArn || mergedCredentials.profileArn;
-        this.region = this.region || mergedCredentials.region;
-
-        // Ensure region is set before using it in URLs
-        if (!this.region) {
-            console.warn('[Kiro Auth] Region not found in credentials. Using default region us-east-1 for URLs.');
-            this.region = 'us-east-1'; // Set default region
-        }
-
-        this.refreshUrl = (this.config.KIRO_REFRESH_URL || KIRO_CONSTANTS.REFRESH_URL).replace("{{region}}", this.region);
-        this.refreshIDCUrl = (this.config.KIRO_REFRESH_IDC_URL || KIRO_CONSTANTS.REFRESH_IDC_URL).replace("{{region}}", this.region);
-        this.baseUrl = (this.config.KIRO_BASE_URL || KIRO_CONSTANTS.BASE_URL).replace("{{region}}", this.region);
-        this.amazonQUrl = (KIRO_CONSTANTS.AMAZON_Q_URL).replace("{{region}}", this.region);
-    } catch (error) {
-        console.warn(`[Kiro Auth] Error during credential loading: ${error.message}`);
-    }
-
-    // Refresh token if forced or if access token is missing but refresh token is available
-    if (forceRefresh || (!this.accessToken && this.refreshToken)) {
-        if (!this.refreshToken) {
-            throw new Error('No refresh token available to refresh access token.');
-        }
         try {
-            const requestBody = {
+            let mergedCredentials = {};
-                refreshToken: this.refreshToken,
-            };
 
-            let refreshUrl = this.refreshUrl;
+            // Priority 1: Load from Base64 credentials if available
-            if (this.authMethod !== KIRO_CONSTANTS.AUTH_METHOD_SOCIAL) {
+            if (this.base64Creds) {
-                refreshUrl = this.refreshIDCUrl;
+                Object.assign(mergedCredentials, this.base64Creds);
-                requestBody.clientId = this.clientId;
+                console.info('[Kiro Auth] Successfully loaded credentials from Base64 (constructor).');
-                requestBody.clientSecret = this.clientSecret;
+                // Clear base64Creds after use to prevent re-processing
-                requestBody.grantType = 'refresh_token';
+                this.base64Creds = null;
             }
 
-            let response = null;
+            // Priority 2 & 3 合并: 从指定文件路径或目录加载凭证
-            if (this.authMethod === KIRO_CONSTANTS.AUTH_METHOD_SOCIAL) {
+            // 读取指定的 credPath 文件以及目录下的其他 JSON 文件(排除当前文件)
-                response = await this.axiosSocialRefreshInstance.post(refreshUrl, requestBody);
+            const targetFilePath = this.credsFilePath || path.join(this.credPath, KIRO_AUTH_TOKEN_FILE);
-                console.log('[Kiro Auth] Token refresh social response: ok');
+            const dirPath = path.dirname(targetFilePath);
-            }else{
+            const targetFileName = path.basename(targetFilePath);
-                response = await this.axiosInstance.post(refreshUrl, requestBody);
-                console.log('[Kiro Auth] Token refresh idc response: ok');
-            }
 
-            if (response.data && response.data.accessToken) {
+            console.debug(`[Kiro Auth] Attempting to load credentials from directory: ${dirPath}`);
-                this.accessToken = response.data.accessToken;
-                this.refreshToken = response.data.refreshToken;
-                this.profileArn = response.data.profileArn;
-                const expiresIn = response.data.expiresIn;
-                const expiresAt = new Date(Date.now() + expiresIn * 1000).toISOString();
-                this.expiresAt = expiresAt;
-                console.info('[Kiro Auth] Access token refreshed successfully');
 
-                // Update the token file - use specified path if configured, otherwise use default
+            try {
-                const tokenFilePath = this.credsFilePath || path.join(this.credPath, KIRO_AUTH_TOKEN_FILE);
+                // 首先尝试读取目标文件
-                const updatedTokenData = {
+                const targetCredentials = await loadCredentialsFromFile(targetFilePath);
-                    accessToken: this.accessToken,
+                if (targetCredentials) {
-                    refreshToken: this.refreshToken,
+                    Object.assign(mergedCredentials, targetCredentials);
-                    expiresAt: expiresAt,
+                    console.info(`[Kiro Auth] Successfully loaded OAuth credentials from ${targetFilePath}`);
-                };
-                if(this.profileArn){
-                    updatedTokenData.profileArn = this.profileArn;
                 }
-                await saveCredentialsToFile(tokenFilePath, updatedTokenData);
-            } else {
-                throw new Error('Invalid refresh response: Missing accessToken');
-            }
-        } catch (error) {
-            console.error('[Kiro Auth] Token refresh failed:', error.message);
-            throw new Error(`Token refresh failed: ${error.message}`);
-        }
-    }
 
-    if (!this.accessToken) {
+                // 然后读取目录下的其他 JSON 文件(排除目标文件本身)
-        throw new Error('No access token available after initialization and refresh attempts.');
+                const files = await fs.readdir(dirPath);
-    }
+                for (const file of files) {
+                    if (file.endsWith('.json') && file !== targetFileName) {
+                        const filePath = path.join(dirPath, file);
+                        const credentials = await loadCredentialsFromFile(filePath);
+                        if (credentials) {
+                            // 保留已有的 expiresAt,避免被覆盖
+                            credentials.expiresAt = mergedCredentials.expiresAt;
+                            Object.assign(mergedCredentials, credentials);
+                            console.debug(`[Kiro Auth] Loaded Client credentials from ${file}`);
+                        }
+                    }
+                }
+            } catch (error) {
+                console.warn(`[Kiro Auth] Error loading credentials from directory ${dirPath}: ${error.message}`);
+            }
+
+            // console.log('[Kiro Auth] Merged credentials:', mergedCredentials);
+            // Apply loaded credentials, prioritizing existing values if they are not null/undefined
+            this.accessToken = this.accessToken || mergedCredentials.accessToken;
+            this.refreshToken = this.refreshToken || mergedCredentials.refreshToken;
+            this.clientId = this.clientId || mergedCredentials.clientId;
+            this.clientSecret = this.clientSecret || mergedCredentials.clientSecret;
+            this.authMethod = this.authMethod || mergedCredentials.authMethod;
+            this.expiresAt = this.expiresAt || mergedCredentials.expiresAt;
+            this.profileArn = this.profileArn || mergedCredentials.profileArn;
+            this.region = this.region || mergedCredentials.region;
+
+            // Ensure region is set before using it in URLs
+            if (!this.region) {
+                console.warn('[Kiro Auth] Region not found in credentials. Using default region us-east-1 for URLs.');
+                this.region = 'us-east-1'; // Set default region
+            }
+
+            this.refreshUrl = (this.config.KIRO_REFRESH_URL || KIRO_CONSTANTS.REFRESH_URL).replace("{{region}}", this.region);
+            this.refreshIDCUrl = (this.config.KIRO_REFRESH_IDC_URL || KIRO_CONSTANTS.REFRESH_IDC_URL).replace("{{region}}", this.region);
+            this.baseUrl = (this.config.KIRO_BASE_URL || KIRO_CONSTANTS.BASE_URL).replace("{{region}}", this.region);
+            this.amazonQUrl = (KIRO_CONSTANTS.AMAZON_Q_URL).replace("{{region}}", this.region);
+        } catch (error) {
+            console.warn(`[Kiro Auth] Error during credential loading: ${error.message}`);
+        }
+
+        // Refresh token if forced or if access token is missing but refresh token is available
+        if (forceRefresh || (!this.accessToken && this.refreshToken)) {
+            if (!this.refreshToken) {
+                throw new Error('No refresh token available to refresh access token.');
+            }
+            try {
+                const requestBody = {
+                    refreshToken: this.refreshToken,
+                };
+
+                let refreshUrl = this.refreshUrl;
+                if (this.authMethod !== KIRO_CONSTANTS.AUTH_METHOD_SOCIAL) {
+                    refreshUrl = this.refreshIDCUrl;
+                    requestBody.clientId = this.clientId;
+                    requestBody.clientSecret = this.clientSecret;
+                    requestBody.grantType = 'refresh_token';
+                }
+
+                let response = null;
+                if (this.authMethod === KIRO_CONSTANTS.AUTH_METHOD_SOCIAL) {
+                    response = await this.axiosSocialRefreshInstance.post(refreshUrl, requestBody);
+                    console.log('[Kiro Auth] Token refresh social response: ok');
+                } else {
+                    response = await this.axiosInstance.post(refreshUrl, requestBody);
+                    console.log('[Kiro Auth] Token refresh idc response: ok');
+                }
+
+                if (response.data && response.data.accessToken) {
+                    this.accessToken = response.data.accessToken;
+                    this.refreshToken = response.data.refreshToken;
+                    this.profileArn = response.data.profileArn;
+                    const expiresIn = response.data.expiresIn;
+                    const expiresAt = new Date(Date.now() + expiresIn * 1000).toISOString();
+                    this.expiresAt = expiresAt;
+                    console.info('[Kiro Auth] Access token refreshed successfully');
+
+                    // Update the token file - use specified path if configured, otherwise use default
+                    const tokenFilePath = this.credsFilePath || path.join(this.credPath, KIRO_AUTH_TOKEN_FILE);
+                    const updatedTokenData = {
+                        accessToken: this.accessToken,
+                        refreshToken: this.refreshToken,
+                        expiresAt: expiresAt,
+                    };
+                    if (this.profileArn) {
+                        updatedTokenData.profileArn = this.profileArn;
+                    }
+                    await saveCredentialsToFile(tokenFilePath, updatedTokenData);
+                } else {
+                    throw new Error('Invalid refresh response: Missing accessToken');
+                }
+            } catch (error) {
+                console.error('[Kiro Auth] Token refresh failed:', error.message);
+                throw new Error(`Token refresh failed: ${error.message}`);
+            }
+        }
+
+        if (!this.accessToken) {
+            throw new Error('No access token available after initialization and refresh attempts.');
+        }
     }
 
     /**
      * Extract text content from OpenAI message format
      */
     getContentText(message) {
-        if(message==null){
+        if (message == null) {
             return "";
         }
-        if (Array.isArray(message) ) {
+        if (Array.isArray(message)) {
             return message
                 .filter(part => part.type === 'text' && part.text)
                 .map(part => part.text)
                 .join('');
         } else if (typeof message.content === 'string') {
             return message.content;
-        } else if (Array.isArray(message.content) ) {
+        } else if (Array.isArray(message.content)) {
             return message.content
                 .filter(part => part.type === 'text' && part.text)
                 .map(part => part.text)
@@ -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,46 +1459,69 @@ 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 事件
-            yield {
-                type: "message_start",
-                message: {
-                    id: messageId,
-                    type: "message",
-                    role: "assistant",
-                    model: model,
-                    usage: { input_tokens: inputTokens, output_tokens: 0 },
-                    content: []
-                }
-            };
-
-            // 2. 发送 content_block_start 事件
-            yield {
-                type: "content_block_start",
-                index: 0,
-                content_block: { type: "text", text: "" }
-            };
-
             let totalContent = '';
             let outputTokens = 0;
             const toolCalls = [];
-            let currentToolCall = null;  // 用于累积结构化工具调用
+            let currentToolCall = null;
 
-            // 3. 流式接收并发送每个 content_block_delta
             for await (const event of this.streamApiReal('', finalModel, requestBody)) {
-                if (event.type === 'content' && event.content) {
+                if (event.type === 'contextUsage' && event.percentage) {
-                    totalContent += event.content;
+                    contextUsagePercentage = event.percentage;
-                    // 不再每个 chunk 都计算 token，改为最后统一计算，避免阻塞事件循环
+                    inputTokens = this.calculateInputTokensFromPercentage(contextUsagePercentage);
 
-                    yield {
+                    if (!messageStartSent) {
+                        yield {
+                            type: "message_start",
+                            message: {
+                                id: messageId,
+                                type: "message",
+                                role: "assistant",
+                                model: model,
+                                usage: {
+                                    input_tokens: inputTokens,
+                                    output_tokens: 0,
+                                    cache_creation_input_tokens: 0,
+                                    cache_read_input_tokens: 0
+                                },
+                                content: []
+                            }
+                        };
+
+                        yield {
+                            type: "content_block_start",
+                            index: 0,
+                            content_block: { type: "text", text: "" }
+                        };
+
+                        messageStartSent = true;
+
+                        for (const buffered of bufferedEvents) {
+                            yield buffered;
+                        }
+                        bufferedEvents.length = 0;
+                    }
+                } else if (event.type === 'content' && event.content) {
+                    totalContent += event.content;
+
+                    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）
@@ -1508,7 +1552,7 @@ async initializeAuth(forceRefresh = false) {
                         if (tc.stop) {
                             try {
                                 currentToolCall.input = JSON.parse(currentToolCall.input);
-                            } catch (e) {}
+                            } catch (e) { }
                             toolCalls.push(currentToolCall);
                             currentToolCall = null;
                         }
@@ -1536,11 +1580,17 @@ async initializeAuth(forceRefresh = false) {
             if (currentToolCall) {
                 try {
                     currentToolCall.input = JSON.parse(currentToolCall.input);
-                } catch (e) {}
+                } catch (e) { }
                 toolCalls.push(currentToolCall);
                 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
             };
@@ -1950,7 +2028,7 @@ async initializeAuth(forceRefresh = false) {
             origin: KIRO_CONSTANTS.ORIGIN_AI_EDITOR,
             resourceType: resourceType
         });
-         if (this.authMethod === KIRO_CONSTANTS.AUTH_METHOD_SOCIAL && this.profileArn) {
+        if (this.authMethod === KIRO_CONSTANTS.AUTH_METHOD_SOCIAL && this.profileArn) {
             params.append('profileArn', this.profileArn);
         }
         const fullUrl = `${usageLimitsUrl}?${params.toString()}`;