From d26b4ee1626e9d250612c0dfbc9ce1182afc3606 Mon Sep 17 00:00:00 2001 From: Zhafron Kautsar Date: Sun, 11 Jan 2026 05:06:01 -0500 Subject: [PATCH 1/2] feat(kiro): implement extended thinking support with streaming and token estimation Add comprehensive support for Claude's extended thinking feature in Kiro provider: - Add thinking block parsing and streaming with proper tag detection - Implement thinking prefix injection in system prompts with budget validation - Add bidirectional conversion between Kiro text format and Claude content blocks - Enhance token estimation with detailed breakdown for thinking, tools, and content types - Fix streaming to properly handle thinking blocks with start/stop events - Improve context usage percentage handling and input token calculation - Add helper functions for quote-aware tag detection to avoid false positives --- src/providers/claude/claude-kiro.js | 557 ++++++++++++++++++++++------ 1 file changed, 448 insertions(+), 109 deletions(-) diff --git a/src/providers/claude/claude-kiro.js b/src/providers/claude/claude-kiro.js index 8acae3e..9062d73 100644 --- a/src/providers/claude/claude-kiro.js +++ b/src/providers/claude/claude-kiro.js @@ -11,6 +11,15 @@ import { countTokens } from '@anthropic-ai/tokenizer'; import { configureAxiosProxy } from '../../utils/proxy-utils.js'; import { isRetryableNetworkError } from '../../utils/common.js'; +const KIRO_THINKING = { + MAX_BUDGET_TOKENS: 24576, + DEFAULT_BUDGET_TOKENS: 20000, + START_TAG: '', + END_TAG: '', + MODE_TAG: '', + MAX_LEN_TAG: '', +}; + const KIRO_CONSTANTS = { REFRESH_URL: 'https://prod.{{region}}.auth.desktop.kiro.dev/refreshToken', REFRESH_IDC_URL: 'https://oidc.{{region}}.amazonaws.com/token', @@ -87,7 +96,27 @@ function getSystemRuntimeInfo() { }; } -// Helper functions for tool calls and JSON parsing +function isQuoteCharAt(text, index) { + if (index < 0 || index >= text.length) return false; + const ch = text[index]; + return ch === '"' || ch === "'" || ch === '`'; +} + +function findRealTag(text, tag, startIndex = 0) { + let searchStart = Math.max(0, startIndex); + while (true) { + const pos = text.indexOf(tag, searchStart); + if (pos === -1) return -1; + + const hasQuoteBefore = isQuoteCharAt(text, pos - 1); + const hasQuoteAfter = isQuoteCharAt(text, pos + tag.length); + if (!hasQuoteBefore && !hasQuoteAfter) { + return pos; + } + + searchStart = pos + 1; + } +} /** * 通用的括号匹配函数 - 支持多种括号类型 @@ -554,26 +583,85 @@ async initializeAuth(forceRefresh = false) { if(message==null){ return ""; } - if (Array.isArray(message) ) { - return message - .filter(part => part.type === 'text' && part.text) - .map(part => part.text) - .join(''); + if (Array.isArray(message)) { + return message.map(part => { + if (typeof part === 'string') return part; + if (part && typeof part === 'object') { + if (part.type === 'text' && part.text) return part.text; + if (part.text) return part.text; + } + return ''; + }).join(''); } else if (typeof message.content === 'string') { return message.content; - } else if (Array.isArray(message.content) ) { - return message.content - .filter(part => part.type === 'text' && part.text) - .map(part => part.text) - .join(''); - } + } else if (Array.isArray(message.content)) { + return message.content.map(part => { + if (typeof part === 'string') return part; + if (part && typeof part === 'object') { + if (part.type === 'text' && part.text) return part.text; + if (part.text) return part.text; + } + return ''; + }).join(''); + } return String(message.content || message); } + _normalizeThinkingBudgetTokens(budgetTokens) { + let value = Number(budgetTokens); + if (!Number.isFinite(value) || value <= 0) { + value = KIRO_THINKING.DEFAULT_BUDGET_TOKENS; + } + value = Math.floor(value); + return Math.min(value, KIRO_THINKING.MAX_BUDGET_TOKENS); + } + + _generateThinkingPrefix(thinking) { + if (!thinking || thinking.type !== 'enabled') return null; + const budget = this._normalizeThinkingBudgetTokens(thinking.budget_tokens); + return `enabled${budget}`; + } + + _hasThinkingPrefix(text) { + if (!text) return false; + return text.includes(KIRO_THINKING.MODE_TAG) || text.includes(KIRO_THINKING.MAX_LEN_TAG); + } + + _toClaudeContentBlocksFromKiroText(content) { + const raw = content ?? ''; + if (!raw) return []; + + const startPos = findRealTag(raw, KIRO_THINKING.START_TAG); + if (startPos === -1) { + return [{ type: "text", text: raw }]; + } + + const before = raw.slice(0, startPos); + let rest = raw.slice(startPos + KIRO_THINKING.START_TAG.length); + + const endPosInRest = findRealTag(rest, KIRO_THINKING.END_TAG); + let thinking = ''; + let after = ''; + if (endPosInRest === -1) { + thinking = rest; + } else { + thinking = rest.slice(0, endPosInRest); + after = rest.slice(endPosInRest + KIRO_THINKING.END_TAG.length); + } + + if (after.startsWith('\n\n')) after = after.slice(2); + + const blocks = []; + if (before) blocks.push({ type: "text", text: before }); + blocks.push({ type: "thinking", thinking }); + if (after) blocks.push({ type: "text", text: after }); + return blocks; + } + /** * Build CodeWhisperer request from OpenAI messages */ - buildCodewhispererRequest(messages, model, tools = null, inSystemPrompt = null) { + buildCodewhispererRequest(messages, model, tools = null, inSystemPrompt = null, thinking = null) { const conversationId = uuidv4(); let systemPrompt = this.getContentText(inSystemPrompt); @@ -583,6 +671,15 @@ async initializeAuth(forceRefresh = false) { throw new Error('No user messages found'); } + const thinkingPrefix = this._generateThinkingPrefix(thinking); + if (thinkingPrefix) { + if (!systemPrompt) { + systemPrompt = thinkingPrefix; + } else if (!this._hasThinkingPrefix(systemPrompt)) { + systemPrompt = `${thinkingPrefix}\n${systemPrompt}`; + } + } + // 判断最后一条消息是否为 assistant,如果是则移除 const lastMessage = processedMessages[processedMessages.length - 1]; if (processedMessages.length > 0 && lastMessage.role === 'assistant') { @@ -818,11 +915,14 @@ async initializeAuth(forceRefresh = false) { content: '' }; let toolUses = []; + let thinkingText = ''; if (Array.isArray(message.content)) { for (const part of message.content) { if (part.type === 'text') { assistantResponseMessage.content += part.text; + } else if (part.type === 'thinking') { + thinkingText += (part.thinking ?? part.text ?? ''); } else if (part.type === 'tool_use') { toolUses.push({ input: part.input, @@ -835,7 +935,12 @@ async initializeAuth(forceRefresh = false) { assistantResponseMessage.content = this.getContentText(message); } - // 只添加非空字段 + if (thinkingText) { + assistantResponseMessage.content = assistantResponseMessage.content + ? `${KIRO_THINKING.START_TAG}${thinkingText}${KIRO_THINKING.END_TAG}\n\n${assistantResponseMessage.content}` + : `${KIRO_THINKING.START_TAG}${thinkingText}${KIRO_THINKING.END_TAG}`; + } + if (toolUses.length > 0) { assistantResponseMessage.toolUses = toolUses; } @@ -861,10 +966,13 @@ async initializeAuth(forceRefresh = false) { content: '', toolUses: [] }; + let thinkingText = ''; if (Array.isArray(currentMessage.content)) { for (const part of currentMessage.content) { if (part.type === 'text') { assistantResponseMessage.content += part.text; + } else if (part.type === 'thinking') { + thinkingText += (part.thinking ?? part.text ?? ''); } else if (part.type === 'tool_use') { assistantResponseMessage.toolUses.push({ input: part.input, @@ -876,6 +984,11 @@ async initializeAuth(forceRefresh = false) { } else { assistantResponseMessage.content = this.getContentText(currentMessage); } + if (thinkingText) { + assistantResponseMessage.content = assistantResponseMessage.content + ? `${KIRO_THINKING.START_TAG}${thinkingText}${KIRO_THINKING.END_TAG}\n\n${assistantResponseMessage.content}` + : `${KIRO_THINKING.START_TAG}${thinkingText}${KIRO_THINKING.END_TAG}`; + } if (assistantResponseMessage.toolUses.length === 0) { delete assistantResponseMessage.toolUses; } @@ -1100,9 +1213,9 @@ async initializeAuth(forceRefresh = false) { async callApi(method, model, body, isRetry = false, retryCount = 0) { if (!this.isInitialized) await this.initialize(); const maxRetries = this.config.REQUEST_MAX_RETRIES || 3; - const baseDelay = this.config.REQUEST_BASE_DELAY || 1000; // 1 second base delay + const baseDelay = this.config.REQUEST_BASE_DELAY || 1000; - const requestData = this.buildCodewhispererRequest(body.messages, model, body.tools, body.system); + const requestData = this.buildCodewhispererRequest(body.messages, model, body.tools, body.system, body.thinking); try { const token = this.accessToken; // Use the already initialized token @@ -1386,7 +1499,7 @@ async initializeAuth(forceRefresh = false) { const maxRetries = this.config.REQUEST_MAX_RETRIES || 3; const baseDelay = this.config.REQUEST_BASE_DELAY || 1000; - const requestData = this.buildCodewhispererRequest(body.messages, model, body.tools, body.system); + const requestData = this.buildCodewhispererRequest(body.messages, model, body.tools, body.system, body.thinking); const token = this.accessToken; const headers = { @@ -1505,8 +1618,7 @@ async initializeAuth(forceRefresh = false) { // 真正的流式传输实现 async * generateContentStream(model, requestBody) { if (!this.isInitialized) await this.initialize(); - - // 检查 token 是否即将过期,如果是则先刷新 + if (this.isExpiryDateNear()) { console.log('[Kiro] Token is near expiry, refreshing before generateContentStream request...'); await this.initializeAuth(true); @@ -1514,12 +1626,93 @@ 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()}`; - + + const thinkingRequested = requestBody?.thinking?.type === 'enabled'; + + const streamState = { + thinkingRequested, + buffer: '', + inThinking: false, + thinkingExtracted: false, + thinkingBlockIndex: null, + textBlockIndex: null, + nextBlockIndex: 0, + stoppedBlocks: new Set(), + }; + + const ensureBlockStart = (blockType) => { + if (blockType === 'thinking') { + if (streamState.thinkingBlockIndex != null) return []; + const idx = streamState.nextBlockIndex++; + streamState.thinkingBlockIndex = idx; + return [{ + type: "content_block_start", + index: idx, + content_block: { type: "thinking", thinking: "" } + }]; + } + if (blockType === 'text') { + if (streamState.textBlockIndex != null) return []; + const idx = streamState.nextBlockIndex++; + streamState.textBlockIndex = idx; + return [{ + type: "content_block_start", + index: idx, + content_block: { type: "text", text: "" } + }]; + } + return []; + }; + + const stopBlock = (index) => { + if (index == null) return []; + if (streamState.stoppedBlocks.has(index)) return []; + streamState.stoppedBlocks.add(index); + return [{ type: "content_block_stop", index }]; + }; + + const createTextDeltaEvents = (text) => { + if (!text) return []; + const events = []; + events.push(...ensureBlockStart('text')); + events.push({ + type: "content_block_delta", + index: streamState.textBlockIndex, + delta: { type: "text_delta", text } + }); + return events; + }; + + const createThinkingDeltaEvents = (thinking) => { + const events = []; + events.push(...ensureBlockStart('thinking')); + events.push({ + type: "content_block_delta", + index: streamState.thinkingBlockIndex, + delta: { type: "thinking_delta", thinking } + }); + return events; + }; + + function* pushEvents(events) { + for (const ev of events) { + yield ev; + } + } + try { - // 1. 先发送 message_start 事件 + let totalContent = ''; + let outputTokens = 0; + const toolCalls = []; + let currentToolCall = null; + + const estimatedInputTokens = this.estimateInputTokens(requestBody); + const tokenBreakdown = this._lastTokenBreakdown || {}; + yield { type: "message_start", message: { @@ -1527,66 +1720,120 @@ async initializeAuth(forceRefresh = false) { type: "message", role: "assistant", model: model, - usage: { input_tokens: inputTokens, output_tokens: 0 }, + usage: { + input_tokens: estimatedInputTokens, + 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 = ''; - let outputTokens = 0; - const toolCalls = []; - let currentToolCall = null; // 用于累积结构化工具调用 - let contextUsagePercentage = null; // 用于存储上下文使用百分比 - - // 3. 流式接收并发送每个 content_block_delta for await (const event of this.streamApiReal('', finalModel, requestBody)) { - if (event.type === 'content' && event.content) { - totalContent += event.content; - // 不再每个 chunk 都计算 token，改为最后统一计算，避免阻塞事件循环 + if (event.type === 'contextUsage' && event.percentage) { + contextUsagePercentage = event.percentage; + inputTokens = this.calculateInputTokensFromPercentage(contextUsagePercentage); - yield { - type: "content_block_delta", - index: 0, - delta: { type: "text_delta", text: event.content } - }; - } else if (event.type === 'contextUsage') { - // 捕获上下文使用百分比 - contextUsagePercentage = event.contextUsagePercentage; - console.log(`[Kiro] Received contextUsagePercentage: ${contextUsagePercentage}%`); + if (Math.abs(inputTokens - estimatedInputTokens) > estimatedInputTokens * 0.1) { + yield { + type: "message_delta", + delta: {}, + usage: { + input_tokens: inputTokens, + output_tokens: 0, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0 + } + }; + } + } else if (event.type === 'content' && event.content) { + totalContent += event.content; + + if (!thinkingRequested) { + yield* pushEvents(createTextDeltaEvents(event.content)); + continue; + } + + streamState.buffer += event.content; + const events = []; + + while (streamState.buffer.length > 0) { + if (!streamState.inThinking && !streamState.thinkingExtracted) { + const startPos = findRealTag(streamState.buffer, KIRO_THINKING.START_TAG); + if (startPos !== -1) { + const before = streamState.buffer.slice(0, startPos); + if (before) events.push(...createTextDeltaEvents(before)); + + streamState.buffer = streamState.buffer.slice(startPos + KIRO_THINKING.START_TAG.length); + streamState.inThinking = true; + continue; + } + + const safeLen = Math.max(0, streamState.buffer.length - KIRO_THINKING.START_TAG.length); + if (safeLen > 0) { + const safeText = streamState.buffer.slice(0, safeLen); + if (safeText) events.push(...createTextDeltaEvents(safeText)); + streamState.buffer = streamState.buffer.slice(safeLen); + } + break; + } + + if (streamState.inThinking) { + const endPos = findRealTag(streamState.buffer, KIRO_THINKING.END_TAG); + if (endPos !== -1) { + const thinkingPart = streamState.buffer.slice(0, endPos); + if (thinkingPart) events.push(...createThinkingDeltaEvents(thinkingPart)); + + streamState.buffer = streamState.buffer.slice(endPos + KIRO_THINKING.END_TAG.length); + streamState.inThinking = false; + streamState.thinkingExtracted = true; + + events.push(...createThinkingDeltaEvents("")); + events.push(...stopBlock(streamState.thinkingBlockIndex)); + + if (streamState.buffer.startsWith('\n\n')) { + streamState.buffer = streamState.buffer.slice(2); + } + continue; + } + + const safeLen = Math.max(0, streamState.buffer.length - KIRO_THINKING.END_TAG.length); + if (safeLen > 0) { + const safeThinking = streamState.buffer.slice(0, safeLen); + if (safeThinking) events.push(...createThinkingDeltaEvents(safeThinking)); + streamState.buffer = streamState.buffer.slice(safeLen); + } + break; + } + + if (streamState.thinkingExtracted) { + const rest = streamState.buffer; + streamState.buffer = ''; + if (rest) events.push(...createTextDeltaEvents(rest)); + break; + } + } + + yield* pushEvents(events); } else if (event.type === 'toolUse') { const tc = event.toolUse; - // 工具调用事件（包含 name 和 toolUseId） if (tc.name && tc.toolUseId) { - // 检查是否是同一个工具调用的续传（相同 toolUseId） if (currentToolCall && currentToolCall.toolUseId === tc.toolUseId) { - // 同一个工具调用，累积 input currentToolCall.input += tc.input || ''; } else { - // 不同的工具调用 - // 如果有未完成的工具调用，先保存它 if (currentToolCall) { try { currentToolCall.input = JSON.parse(currentToolCall.input); - } catch (e) { - // input 不是有效 JSON，保持原样 - } + } catch (e) {} toolCalls.push(currentToolCall); } - // 开始新的工具调用 currentToolCall = { toolUseId: tc.toolUseId, name: tc.name, input: tc.input || '' }; } - // 如果这个事件包含 stop，完成工具调用 if (tc.stop) { try { currentToolCall.input = JSON.parse(currentToolCall.input); @@ -1596,25 +1843,20 @@ async initializeAuth(forceRefresh = false) { } } } else if (event.type === 'toolUseInput') { - // 工具调用的 input 续传事件 if (currentToolCall) { currentToolCall.input += event.input || ''; } } else if (event.type === 'toolUseStop') { - // 工具调用结束事件 if (currentToolCall && event.stop) { try { currentToolCall.input = JSON.parse(currentToolCall.input); - } catch (e) { - // input 不是有效 JSON，保持原样 - } + } catch (e) {} toolCalls.push(currentToolCall); currentToolCall = null; } } } - // 处理未完成的工具调用（如果流提前结束） if (currentToolCall) { try { currentToolCall.input = JSON.parse(currentToolCall.input); @@ -1622,8 +1864,30 @@ async initializeAuth(forceRefresh = false) { toolCalls.push(currentToolCall); currentToolCall = null; } - - // 检查文本内容中的 bracket 格式工具调用 + + if (thinkingRequested && streamState.buffer) { + if (streamState.inThinking) { + console.warn('[Kiro] Incomplete thinking tag at stream end'); + yield* pushEvents(createThinkingDeltaEvents(streamState.buffer)); + streamState.buffer = ''; + yield* pushEvents(createThinkingDeltaEvents("")); + yield* pushEvents(stopBlock(streamState.thinkingBlockIndex)); + } else if (!streamState.thinkingExtracted) { + yield* pushEvents(createTextDeltaEvents(streamState.buffer)); + streamState.buffer = ''; + } else { + yield* pushEvents(createTextDeltaEvents(streamState.buffer)); + streamState.buffer = ''; + } + } + + yield* pushEvents(stopBlock(streamState.textBlockIndex)); + + if (contextUsagePercentage === null) { + console.warn('[Kiro Stream] contextUsagePercentage not received, using estimation'); + inputTokens = estimatedInputTokens; + } + const bracketToolCalls = parseBracketToolCalls(totalContent); if (bracketToolCalls && bracketToolCalls.length > 0) { for (const btc of bracketToolCalls) { @@ -1635,15 +1899,12 @@ async initializeAuth(forceRefresh = false) { } } - // 4. 发送 content_block_stop 事件 - yield { type: "content_block_stop", index: 0 }; - - // 5. 处理工具调用（如果有） if (toolCalls.length > 0) { + const baseIndex = streamState.nextBlockIndex; for (let i = 0; i < toolCalls.length; i++) { const tc = toolCalls[i]; - const blockIndex = i + 1; - + const blockIndex = baseIndex + i; + yield { type: "content_block_start", index: blockIndex, @@ -1668,32 +1929,29 @@ async initializeAuth(forceRefresh = false) { } } - // 6. 发送 message_delta 事件 - // 如果有 contextUsagePercentage，使用它来计算 token - // 总上下文 200k tokens，通过百分比计算总使用量，再减去输入 token 得到输出 token - let totalTokens = 0; - if (contextUsagePercentage !== null && contextUsagePercentage > 0) { - const totalContextTokens = KIRO_CONSTANTS.TOTAL_CONTEXT_TOKENS; - // totalUsedTokens 就是通过百分比计算出的总使用量，直接作为 total_tokens - totalTokens = Math.round(totalContextTokens * contextUsagePercentage / 100); - outputTokens = Math.max(0, totalTokens - inputTokens); - console.log(`[Kiro] Token calculation from contextUsagePercentage: total=${totalTokens}, input=${inputTokens}, output=${outputTokens}`); - } else { - // 回退到原来的计算方式 - outputTokens = this.countTextTokens(totalContent); - for (const tc of toolCalls) { - outputTokens += this.countTextTokens(JSON.stringify(tc.input || {})); - } - totalTokens = inputTokens + outputTokens; + const contentBlocksForCount = thinkingRequested + ? this._toClaudeContentBlocksFromKiroText(totalContent) + : [{ type: "text", text: totalContent }]; + const plainForCount = contentBlocksForCount + .map(b => (b.type === 'thinking' ? (b.thinking ?? '') : (b.text ?? ''))) + .join(''); + outputTokens = this.countTextTokens(plainForCount); + + for (const tc of toolCalls) { + outputTokens += this.countTextTokens(JSON.stringify(tc.input || {})); } - + yield { type: "message_delta", delta: { stop_reason: toolCalls.length > 0 ? "tool_use" : "end_turn" }, - usage: { input_tokens: inputTokens, output_tokens: outputTokens, total_tokens: totalTokens } + usage: { + input_tokens: inputTokens, + output_tokens: outputTokens, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0 + } }; - // 7. 发送 message_stop 事件 yield { type: "message_stop" }; } catch (error) { @@ -1702,9 +1960,6 @@ async initializeAuth(forceRefresh = false) { } } - /** - * Count tokens for a given text using Claude's official tokenizer - */ countTextTokens(text) { if (!text) return 0; try { @@ -1721,29 +1976,113 @@ async initializeAuth(forceRefresh = false) { */ estimateInputTokens(requestBody) { let totalTokens = 0; - - // Count system prompt tokens + const OVERHEAD_MULTIPLIERS = { + system: 1.0, + message: 1.0, + tools: 1.0, + thinking: 1.0, + tool_result: 1.0, + tool_use_input: 1.0, + image: 1500 + }; + + const breakdown = { + system: 0, + thinking: 0, + text: 0, + tool_result: 0, + tool_use_input: 0, + image: 0, + thinking_content: 0, + tools_def: 0 + }; + if (requestBody.system) { const systemText = this.getContentText(requestBody.system); - totalTokens += this.countTextTokens(systemText); + const systemTokens = this.countTextTokens(systemText); + const counted = Math.ceil(systemTokens * OVERHEAD_MULTIPLIERS.system); + breakdown.system = counted; + totalTokens += counted; } - - // Count all messages tokens + + if (requestBody.thinking?.type === 'enabled') { + const budget = this._normalizeThinkingBudgetTokens(requestBody.thinking.budget_tokens); + const prefixText = `enabled${budget}`; + const prefixTokens = this.countTextTokens(prefixText); + const counted = Math.ceil(prefixTokens * OVERHEAD_MULTIPLIERS.thinking); + breakdown.thinking = counted; + totalTokens += counted; + } + if (requestBody.messages && Array.isArray(requestBody.messages)) { for (const message of requestBody.messages) { - if (message.content) { - const contentText = this.getContentText(message); - totalTokens += this.countTextTokens(contentText); + if (!message.content) { + continue; + } + + if (Array.isArray(message.content)) { + for (const part of message.content) { + if (part.type === 'text' && part.text) { + const counted = Math.ceil(this.countTextTokens(part.text) * OVERHEAD_MULTIPLIERS.message); + breakdown.text += counted; + totalTokens += counted; + } + else if (part.type === 'tool_result') { + const toolResultText = this.getContentText(part.content); + const counted = Math.ceil(this.countTextTokens(toolResultText) * OVERHEAD_MULTIPLIERS.tool_result); + breakdown.tool_result += counted; + totalTokens += counted; + } + else if (part.type === 'tool_use' && part.input) { + const inputJson = JSON.stringify(part.input); + const counted = Math.ceil(this.countTextTokens(inputJson) * OVERHEAD_MULTIPLIERS.tool_use_input); + breakdown.tool_use_input += counted; + totalTokens += counted; + } + else if (part.type === 'image') { + breakdown.image += OVERHEAD_MULTIPLIERS.image; + totalTokens += OVERHEAD_MULTIPLIERS.image; + } + else if (part.type === 'thinking' && part.thinking) { + const counted = Math.ceil(this.countTextTokens(part.thinking) * OVERHEAD_MULTIPLIERS.message); + breakdown.thinking_content += counted; + totalTokens += counted; + } + } + } + else if (typeof message.content === 'string') { + const counted = Math.ceil(this.countTextTokens(message.content) * OVERHEAD_MULTIPLIERS.message); + breakdown.text += counted; + totalTokens += counted; } } } - - // Count tools definitions tokens if present + if (requestBody.tools && Array.isArray(requestBody.tools)) { - totalTokens += this.countTextTokens(JSON.stringify(requestBody.tools)); + for (const tool of requestBody.tools) { + const toolJson = JSON.stringify(tool); + const toolTokens = this.countTextTokens(toolJson); + const counted = Math.ceil(toolTokens * OVERHEAD_MULTIPLIERS.tools); + breakdown.tools_def += counted; + totalTokens += counted; + } } + + const hasTools = requestBody.tools && requestBody.tools.length > 0; + const toolsDefTokens = breakdown.tools_def || 0; + const isSmallToolsDef = toolsDefTokens > 0 && toolsDefTokens < 21000; + + const KIRO_BASE_OVERHEAD = 400; + const KIRO_PERCENTAGE_OVERHEAD = hasTools + ? (isSmallToolsDef ? 0.18 : 0.08) + : 0.25; + + const baseOverhead = KIRO_BASE_OVERHEAD; + const percentageOverhead = Math.ceil(totalTokens * KIRO_PERCENTAGE_OVERHEAD); + totalTokens += baseOverhead + percentageOverhead; - return totalTokens; + this._lastTokenBreakdown = breakdown; + return Math.ceil(totalTokens); } /** From 89083499bc2b44aa0fd0b0a4ef04c22461d6d7d2 Mon Sep 17 00:00:00 2001 From: Zhafron Kautsar Date: Sun, 11 Jan 2026 10:52:26 -0500 Subject: [PATCH 2/2] docs(kiro): restore deleted comments --- src/providers/claude/claude-kiro.js | 38 ++++++++++++++++++++++++++--- 1 file changed, 34 insertions(+), 4 deletions(-) diff --git a/src/providers/claude/claude-kiro.js b/src/providers/claude/claude-kiro.js index 9062d73..b80c60c 100644 --- a/src/providers/claude/claude-kiro.js +++ b/src/providers/claude/claude-kiro.js @@ -96,6 +96,8 @@ function getSystemRuntimeInfo() { }; } +// Helper functions for tool calls and JSON parsing + function isQuoteCharAt(text, index) { if (index < 0 || index >= text.length) return false; const ch = text[index]; @@ -941,6 +943,7 @@ async initializeAuth(forceRefresh = false) { : `${KIRO_THINKING.START_TAG}${thinkingText}${KIRO_THINKING.END_TAG}`; } + // 只添加非空字段 if (toolUses.length > 0) { assistantResponseMessage.toolUses = toolUses; } @@ -1618,7 +1621,8 @@ async initializeAuth(forceRefresh = false) { // 真正的流式传输实现 async * generateContentStream(model, requestBody) { if (!this.isInitialized) await this.initialize(); - + + // 检查 token 是否即将过期,如果是则先刷新 if (this.isExpiryDateNear()) { console.log('[Kiro] Token is near expiry, refreshing before generateContentStream request...'); await this.initializeAuth(true); @@ -1708,11 +1712,12 @@ async initializeAuth(forceRefresh = false) { let totalContent = ''; let outputTokens = 0; const toolCalls = []; - let currentToolCall = null; + let currentToolCall = null; // 用于累积结构化工具调用 const estimatedInputTokens = this.estimateInputTokens(requestBody); const tokenBreakdown = this._lastTokenBreakdown || {}; + // 1. 先发送 message_start 事件 yield { type: "message_start", message: { @@ -1730,8 +1735,10 @@ async initializeAuth(forceRefresh = false) { } }; + // 2. 流式接收并发送每个 content_block_delta for await (const event of this.streamApiReal('', finalModel, requestBody)) { if (event.type === 'contextUsage' && event.percentage) { + // 捕获上下文使用百分比 contextUsagePercentage = event.percentage; inputTokens = this.calculateInputTokensFromPercentage(contextUsagePercentage); @@ -1818,22 +1825,31 @@ async initializeAuth(forceRefresh = false) { yield* pushEvents(events); } else if (event.type === 'toolUse') { const tc = event.toolUse; + // 工具调用事件（包含 name 和 toolUseId） if (tc.name && tc.toolUseId) { + // 检查是否是同一个工具调用的续传（相同 toolUseId） if (currentToolCall && currentToolCall.toolUseId === tc.toolUseId) { + // 同一个工具调用，累积 input currentToolCall.input += tc.input || ''; } else { + // 不同的工具调用 + // 如果有未完成的工具调用，先保存它 if (currentToolCall) { try { currentToolCall.input = JSON.parse(currentToolCall.input); - } catch (e) {} + } catch (e) { + // input 不是有效 JSON，保持原样 + } toolCalls.push(currentToolCall); } + // 开始新的工具调用 currentToolCall = { toolUseId: tc.toolUseId, name: tc.name, input: tc.input || '' }; } + // 如果这个事件包含 stop，完成工具调用 if (tc.stop) { try { currentToolCall.input = JSON.parse(currentToolCall.input); @@ -1843,20 +1859,25 @@ async initializeAuth(forceRefresh = false) { } } } else if (event.type === 'toolUseInput') { + // 工具调用的 input 续传事件 if (currentToolCall) { currentToolCall.input += event.input || ''; } } else if (event.type === 'toolUseStop') { + // 工具调用结束事件 if (currentToolCall && event.stop) { try { currentToolCall.input = JSON.parse(currentToolCall.input); - } catch (e) {} + } catch (e) { + // input 不是有效 JSON，保持原样 + } toolCalls.push(currentToolCall); currentToolCall = null; } } } + // 处理未完成的工具调用（如果流提前结束） if (currentToolCall) { try { currentToolCall.input = JSON.parse(currentToolCall.input); @@ -1888,6 +1909,7 @@ async initializeAuth(forceRefresh = false) { inputTokens = estimatedInputTokens; } + // 检查文本内容中的 bracket 格式工具调用 const bracketToolCalls = parseBracketToolCalls(totalContent); if (bracketToolCalls && bracketToolCalls.length > 0) { for (const btc of bracketToolCalls) { @@ -1899,6 +1921,7 @@ async initializeAuth(forceRefresh = false) { } } + // 3. 处理工具调用（如果有） if (toolCalls.length > 0) { const baseIndex = streamState.nextBlockIndex; for (let i = 0; i < toolCalls.length; i++) { @@ -1941,6 +1964,7 @@ async initializeAuth(forceRefresh = false) { outputTokens += this.countTextTokens(JSON.stringify(tc.input || {})); } + // 4. 发送 message_delta 事件 yield { type: "message_delta", delta: { stop_reason: toolCalls.length > 0 ? "tool_use" : "end_turn" }, @@ -1952,6 +1976,7 @@ async initializeAuth(forceRefresh = false) { } }; + // 5. 发送 message_stop 事件 yield { type: "message_stop" }; } catch (error) { @@ -1976,6 +2001,8 @@ async initializeAuth(forceRefresh = false) { */ estimateInputTokens(requestBody) { let totalTokens = 0; + + // 定义各类内容的开销乘数 const OVERHEAD_MULTIPLIERS = { system: 1.0, message: 1.0, @@ -1997,6 +2024,7 @@ async initializeAuth(forceRefresh = false) { tools_def: 0 }; + // Count system prompt tokens if (requestBody.system) { const systemText = this.getContentText(requestBody.system); const systemTokens = this.countTextTokens(systemText); @@ -2014,6 +2042,7 @@ async initializeAuth(forceRefresh = false) { totalTokens += counted; } + // Count all messages tokens if (requestBody.messages && Array.isArray(requestBody.messages)) { for (const message of requestBody.messages) { if (!message.content) { @@ -2058,6 +2087,7 @@ async initializeAuth(forceRefresh = false) { } } + // Count tools definitions tokens if present if (requestBody.tools && Array.isArray(requestBody.tools)) { for (const tool of requestBody.tools) { const toolJson = JSON.stringify(tool);