Kiro Executor Stability and API Compatibility Improvements #18
Conversation
## English Description ### Request Format Fixes - Fix conversationState field order (chatTriggerType must be first) - Add conditional profileArn inclusion based on auth method - builder-id auth (AWS SSO) doesn't require profileArn - social auth (Google OAuth) requires profileArn ### Stream Processing Enhancements - Add headersLen boundary validation to prevent slice out of bounds - Handle incomplete tool use at EOF by flushing pending data - Separate message_delta and message_stop events for proper streaming - Add error logging for JSON unmarshal failures ### JSON Repair Improvements - Add escapeNewlinesInStrings() to handle control characters in JSON strings - Remove incorrect unquotedKeyPattern that broke valid JSON content - Fix handling of streaming fragments with embedded newlines/tabs ### Debug Info Filtering (Optional) - Add filterHeliosDebugInfo() to remove [HELIOS_CHK] blocks - Pattern matches internal state tracking from Kiro/Amazon Q - Currently disabled pending further testing ### Usage Tracking - Add usage information extraction in message_delta response - Include prompt_tokens, completion_tokens, total_tokens in OpenAI format --- ## 中文描述 ### 请求格式修复 - 修复 conversationState 字段顺序（chatTriggerType 必须在第一位） - 根据认证方式条件性包含 profileArn - builder-id 认证（AWS SSO）不需要 profileArn - social 认证（Google OAuth）需要 profileArn ### 流处理增强 - 添加 headersLen 边界验证，防止切片越界 - 在 EOF 时处理未完成的工具调用，刷新待处理数据 - 分离 message_delta 和 message_stop 事件以实现正确的流式传输 - 添加 JSON 反序列化失败的错误日志 ### JSON 修复改进 - 添加 escapeNewlinesInStrings() 处理 JSON 字符串中的控制字符 - 移除错误的 unquotedKeyPattern，该模式会破坏有效的 JSON 内容 - 修复包含嵌入换行符/制表符的流式片段处理 ### 调试信息过滤（可选） - 添加 filterHeliosDebugInfo() 移除 [HELIOS_CHK] 块 - 模式匹配来自 Kiro/Amazon Q 的内部状态跟踪信息 - 目前已禁用，等待进一步测试 ### 使用量跟踪 - 在 message_delta 响应中添加 usage 信息提取 - 以 OpenAI 格式包含 prompt_tokens、completion_tokens、total_tokens
## Changes Overview This commit includes multiple improvements to the Kiro executor for better stability, API compatibility, and code quality. ## Detailed Changes ### 1. Output Token Calculation Improvement (lines 317-330) - Replace simple len(content)/4 estimation with tiktoken-based calculation - Add fallback to character count estimation if tiktoken fails - Improves token counting accuracy for usage tracking ### 2. Stream Handler Panic Recovery (lines 528-533) - Add defer/recover block in streamToChannel goroutine - Prevents single request crashes from affecting the entire service ### 3. Struct Field Reordering (lines 670-673) - Reorder kiroToolResult struct fields: Content, Status, ToolUseID - Ensures consistency with API expectations ### 4. Message Merging Function (lines 778-780, 2356-2483) - Add mergeAdjacentMessages() to combine consecutive messages with same role - Add helper functions: mergeMessageContent(), blockToMap(), createMergedMessage() - Required by Kiro API which doesn't allow adjacent messages from same role ### 5. Empty Content Handling (lines 791-800) - Add default content for empty history messages - User messages with tool results: "Tool results provided." - User messages without tool results: "Continue" ### 6. Assistant Last Message Handling (lines 811-830) - Detect when last message is from assistant - Create synthetic "Continue" user message to satisfy Kiro API requirements - Kiro API requires currentMessage to be userInputMessage type ### 7. Duplicate Content Event Detection (lines 1650-1660) - Track lastContentEvent to detect duplicate streaming events - Skip redundant events to prevent duplicate content in responses - Based on AIClient-2-API implementation for Kiro ### 8. Streaming Token Calculation Enhancement (lines 1785-1817) - Add accumulatedContent buffer for streaming token calculation - Use tiktoken for accurate output token counting during streaming - Add fallback to character count estimation with proper logging ### 9. JSON Repair Enhancement (lines 2665-2818) - Implement conservative JSON repair strategy - First try to parse JSON directly - if valid, return unchanged - Add bracket balancing detection and repair - Only repair when necessary to avoid corrupting valid JSON - Validate repaired JSON before returning ### 10. HELIOS_CHK Filtering Removal (lines 2500-2504, 3004-3039) - Remove filterHeliosDebugInfo function - Remove heliosDebugPattern regex - HELIOS_CHK fields now handled by client-side processing ### 11. Comment Translation - Translate Chinese comments to English for code consistency - Affected areas: token calculation, buffer handling, message processing
Summary of ChangesHello @Ravens2121, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the stability, API compatibility, and overall code quality of the Kiro executor. It addresses various critical aspects, from refining request formats and bolstering stream processing to improving JSON repair mechanisms and ensuring accurate token calculations. The changes aim to create a more robust and compliant interaction with the Kiro API, providing a smoother and more reliable user experience. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a comprehensive set of improvements to the Kiro executor, enhancing its stability, API compatibility, and overall robustness. The changes are well-structured and address several important issues, including request formatting, stream processing reliability, JSON repair, and message handling. The addition of panic recovery, boundary checks, and more accurate token counting significantly improves the executor's quality.
My review focuses on identifying opportunities to reduce code duplication and improve maintainability. I've suggested extracting some repeated logic into helper functions. These changes should make the code cleaner and easier to manage in the future. Overall, this is a very strong contribution that makes the Kiro executor much more reliable.
| // Determine if profileArn should be included based on auth method | ||
| // profileArn is only needed for social auth (Google OAuth), not for builder-id (AWS SSO) | ||
| effectiveProfileArn := profileArn | ||
| if auth != nil && auth.Metadata != nil { | ||
| if authMethod, ok := auth.Metadata["auth_method"].(string); ok && authMethod == "builder-id" { | ||
| effectiveProfileArn = "" // Don't include profileArn for builder-id auth | ||
| } | ||
| } |
There was a problem hiding this comment.
This logic to determine effectiveProfileArn is duplicated in ExecuteStream at lines 401-408. To improve maintainability and reduce code duplication, consider extracting this into a helper function.
For example:
func getEffectiveProfileArn(auth *cliproxyauth.Auth, profileArn string) string {
if auth != nil && auth.Metadata != nil {
if authMethod, ok := auth.Metadata["auth_method"].(string); ok && authMethod == "builder-id" {
return "" // Don't include profileArn for builder-id auth
}
}
return profileArn
}You could then call this function in both Execute and ExecuteStream.
| // Deduplicate currentToolResults before adding to context | ||
| // Kiro API does not accept duplicate toolUseIds | ||
| if len(currentToolResults) > 0 { | ||
| seenIDs := make(map[string]bool) | ||
| uniqueToolResults := make([]kiroToolResult, 0, len(currentToolResults)) | ||
| for _, tr := range currentToolResults { | ||
| if !seenIDs[tr.ToolUseID] { | ||
| seenIDs[tr.ToolUseID] = true | ||
| uniqueToolResults = append(uniqueToolResults, tr) | ||
| } else { | ||
| log.Debugf("kiro: skipping duplicate toolResult in currentMessage: %s", tr.ToolUseID) | ||
| } | ||
| } | ||
| currentToolResults = uniqueToolResults | ||
| } |
There was a problem hiding this comment.
This block deduplicates currentToolResults. However, currentToolResults is assigned from toolResults which is returned by buildUserMessageStruct. The buildUserMessageStruct function already performs deduplication of tool results (as seen in the changes at lines 965-970).
This makes the deduplication here redundant. Please consider removing this block to simplify the code and avoid unnecessary processing.
| // Send tool_use content block | ||
| blockStart := e.buildClaudeContentBlockStartEvent(contentBlockIndex, "tool_use", currentToolUse.toolUseID, currentToolUse.name) | ||
| sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, blockStart, &translatorParam) | ||
| for _, chunk := range sseData { | ||
| if chunk != "" { | ||
| out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")} | ||
| } | ||
| } | ||
|
|
||
| // Send tool input as delta | ||
| inputBytes, _ := json.Marshal(finalInput) | ||
| inputDelta := e.buildClaudeInputJsonDeltaEvent(string(inputBytes), contentBlockIndex) | ||
| sseData = sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, inputDelta, &translatorParam) | ||
| for _, chunk := range sseData { | ||
| if chunk != "" { | ||
| out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")} | ||
| } | ||
| } | ||
|
|
||
| // Close block | ||
| blockStop := e.buildClaudeContentBlockStopEvent(contentBlockIndex) | ||
| sseData = sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, blockStop, &translatorParam) | ||
| for _, chunk := range sseData { | ||
| if chunk != "" { | ||
| out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")} | ||
| } | ||
| } | ||
|
|
There was a problem hiding this comment.
The pattern of translating an event to SSE and sending it to the output channel is repeated three times in this block for blockStart, inputDelta, and blockStop. To improve readability and reduce duplication, you could extract this into a helper function.
For example:
func sendSSEChunk(ctx context.Context, out chan<- cliproxyexecutor.StreamChunk, targetFormat sdktranslator.Format, model string, originalReq, claudeBody, eventPayload []byte, translatorParam *any) {
sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, eventPayload, translatorParam)
for _, chunk := range sseData {
if chunk != "" {
out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
}
}
}This would make the flushing logic for incomplete tool use more concise.
| if totalUsage.OutputTokens == 0 && accumulatedContent.Len() > 0 { | ||
| // Try to use tiktoken for accurate counting | ||
| if enc, err := tokenizerForModel(model); err == nil { | ||
| if tokenCount, countErr := enc.Count(accumulatedContent.String()); countErr == nil { | ||
| totalUsage.OutputTokens = int64(tokenCount) | ||
| log.Debugf("kiro: streamToChannel calculated output tokens using tiktoken: %d", totalUsage.OutputTokens) | ||
| } else { | ||
| // Fallback on count error: estimate from character count | ||
| totalUsage.OutputTokens = int64(accumulatedContent.Len() / 4) | ||
| if totalUsage.OutputTokens == 0 { | ||
| totalUsage.OutputTokens = 1 | ||
| } | ||
| log.Debugf("kiro: streamToChannel tiktoken count failed, estimated from chars: %d", totalUsage.OutputTokens) | ||
| } | ||
| } else { | ||
| // Fallback: estimate from character count (roughly 4 chars per token) | ||
| totalUsage.OutputTokens = int64(accumulatedContent.Len() / 4) | ||
| if totalUsage.OutputTokens == 0 { | ||
| totalUsage.OutputTokens = 1 | ||
| } | ||
| log.Debugf("kiro: streamToChannel estimated output tokens from chars: %d (content len: %d)", totalUsage.OutputTokens, accumulatedContent.Len()) | ||
| } |
There was a problem hiding this comment.
The fallback logic to estimate tokens from character count is duplicated within this if block (lines 1797-1802 and 1805-1810).
To avoid this repetition, you could extract the estimation into a helper function.
For example:
func estimateTokensFromChars(length int) int64 {
tokens := int64(length / 4)
if tokens == 0 && length > 0 {
return 1
}
return tokens
}Then you could simplify the fallback logic to:
// ...
} else {
// Fallback on count error: estimate from character count
totalUsage.OutputTokens = estimateTokensFromChars(accumulatedContent.Len())
log.Debugf("kiro: streamToChannel tiktoken count failed, estimated from chars: %d", totalUsage.OutputTokens)
}
} else {
// Fallback: estimate from character count (roughly 4 chars per token)
totalUsage.OutputTokens = estimateTokensFromChars(accumulatedContent.Len())
log.Debugf("kiro: streamToChannel estimated output tokens from chars: %d (content len: %d)", totalUsage.OutputTokens, accumulatedContent.Len())
}
2 participants
Pull Request: Kiro Executor Stability and API Compatibility Improvements
Summary
This PR includes multiple improvements to the Kiro executor for better stability, API compatibility, and code quality. The changes span two commits that address request format issues, stream processing reliability, JSON repair robustness, token calculation accuracy, and message handling enhancements.
Commits Included
Commit 1:
cd4e84a- feat(kiro): enhance request format, stream handling, and usage trackingCommit 2:
6133bac- feat(kiro): enhance Kiro executor stability and compatibilityDetailed Changes
🔧 Request Format Fixes (Commit 1)
conversationStatefield order -chatTriggerTypemust be first to comply with Kiro API requirementsprofileArninclusion based on authentication methodbuilder-idauth (AWS SSO): Does not requireprofileArnsocialauth (Google OAuth): RequiresprofileArn🌊 Stream Processing Enhancements (Commits 1 & 2)
headersLenboundary validation to prevent slice out of bounds panicsmessage_deltaandmessage_stopevents for proper streaming protocol compliancestreamToChannelgoroutine to prevent single request crashes from affecting the entire service🔨 JSON Repair Improvements (Commits 1 & 2)
escapeNewlinesInStrings()to handle control characters (newlines, tabs) in JSON stringsunquotedKeyPatternthat was breaking valid JSON content📊 Token Calculation Improvements (Commit 2)
len(content)/4estimation with tiktoken-based calculationaccumulatedContentbuffer for streaming token calculation with real-time counting📝 Message Processing Enhancements (Commit 2)
mergeAdjacentMessages()to combine consecutive messages with same rolemergeMessageContent(),blockToMap(),createMergedMessage()currentMessageto beuserInputMessagetype🔄 Duplicate Content Detection (Commit 2)
lastContentEventto detect duplicate streaming events📈 Usage Tracking (Commit 1)
message_deltaresponseprompt_tokens,completion_tokens,total_tokensin OpenAI format🏗️ Struct Field Reordering (Commit 2)
Content,Status,ToolUseID🧹 Code Quality Improvements (Commit 2)
filterHeliosDebugInfofunctionheliosDebugPatternregex🔍 Debug Info Filtering - Optional (Commit 1)
filterHeliosDebugInfo()to remove[HELIOS_CHK]blocksFiles Changed
internal/runtime/executor/kiro_executor.gointernal/translator/kiro/openai/chat-completions/kiro_openai_response.goTotal Changes: 2 files, ~740 insertions, ~100 deletions
Key Line References (Final State)
Testing Checklist
Breaking Changes
None. All changes are backward compatible.
Related Issues
N/A
中文说明
提交概览
本次 Pull Request 包含两个提交：
cd4e84a): 增强请求格式、流处理和使用量跟踪6133bac): 增强 Kiro executor 稳定性和兼容性请求格式修复 (Commit 1)
conversationState字段顺序（chatTriggerType必须在第一位）profileArnbuilder-id认证（AWS SSO）不需要profileArnsocial认证（Google OAuth）需要profileArn流处理增强 (Commits 1 & 2)
headersLen边界验证，防止切片越界message_delta和message_stop事件以实现正确的流式传输JSON 修复改进 (Commits 1 & 2)
escapeNewlinesInStrings()处理 JSON 字符串中的控制字符unquotedKeyPattern，该模式会破坏有效的 JSON 内容Token 计算改进 (Commit 2)
len(content)/4估算消息处理增强 (Commit 2)
mergeAdjacentMessages()合并相同角色的连续消息mergeMessageContent(),blockToMap(),createMergedMessage()currentMessage必须是userInputMessage类型重复内容检测 (Commit 2)
lastContentEvent跟踪以检测重复的流式事件使用量跟踪 (Commit 1)
message_delta响应中添加 usage 信息提取prompt_tokens、completion_tokens、total_tokens结构体字段重排序 (Commit 2)
kiroToolResult结构体字段：Content,Status,ToolUseID代码质量改进 (Commit 2)
filterHeliosDebugInfo函数和heliosDebugPattern正则表达式