-
-
Notifications
You must be signed in to change notification settings - Fork 129
fix: handle Write tool truncation when content exceeds API limits #150
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
fix: handle Write tool truncation when content exceeds API limits #150
Conversation
When the Kiro/AWS CodeWhisperer API receives a Write tool request with content that exceeds transmission limits, it truncates the tool input. This can result in: - Empty input buffer (no input transmitted at all) - Missing 'content' field in the parsed JSON - Incomplete JSON that fails to parse This fix detects these truncation scenarios and converts them to Bash tool calls that echo an error message. This allows Claude Code to execute the Bash command, see the error output, and the agent can then retry with smaller chunks. Changes: - kiro_claude_tools.go: Detect three truncation scenarios in ProcessToolUseEvent: 1. Empty input buffer (no input transmitted) 2. JSON parse failure with file_path but no content field 3. Successfully parsed JSON missing content field When detected, emit a special '__truncated_write__' marker tool use - kiro_executor.go: Handle '__truncated_write__' markers in streamToChannel: 1. Extract file_path from the marker for context 2. Create a Bash tool_use that echoes an error message 3. Include retry guidance (700-line chunks recommended) 4. Set hasToolUses=true to ensure stop_reason='tool_use' for agent continuation This ensures the agent continues and can retry with smaller file chunks instead of failing silently or showing errors to the user.
Summary of ChangesHello @PancakeZik, 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 provides a robust workaround for handling Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. 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
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Code Review
This pull request introduces a clever workaround to handle API limits for the Write tool by detecting content truncation and converting it into a Bash tool call that echoes an error message. This allows the agent to receive feedback and retry the operation. The implementation correctly identifies various truncation scenarios in kiro_claude_tools.go and handles the special __truncated_write__ marker in kiro_executor.go.
My review focuses on improving maintainability by reducing code duplication, enhancing the robustness of string parsing, and addressing a potential error handling issue. Overall, the changes are well-thought-out and effectively address the problem described.
| bashInput := map[string]interface{}{ | ||
| "command": errorMsg, | ||
| } | ||
| inputJSON, _ := json.Marshal(bashInput) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The error returned by json.Marshal is being ignored. While it's unlikely to fail for this specific map structure, it's best practice to handle potential errors. If marshalling were to fail, inputJSON would be nil, resulting in an empty command being sent to the Bash tool. This would prevent the agent from receiving the intended error message.
Consider handling the error, for instance by logging it and skipping the tool use for this event.
inputJSON, err := json.Marshal(bashInput)
if err != nil {
log.Errorf("kiro: failed to marshal bash input for truncated write error: %v", err)
continue
}| if tu.Name == "__truncated_write__" { | ||
| filePath := "" | ||
| if fp, ok := tu.Input["file_path"].(string); ok && fp != "" { | ||
| filePath = fp | ||
| } | ||
|
|
||
| // Create a Bash tool that echoes the error message | ||
| // This will be executed by Claude Code and the agent will see the result | ||
| var errorMsg string | ||
| if filePath != "" { | ||
| errorMsg = fmt.Sprintf("echo '[WRITE TOOL ERROR] The file content for \"%s\" is too large to be transmitted by the upstream API. You MUST retry by writing the file in smaller chunks: First use Write to create the file with the first 700 lines, then use multiple Edit operations to append the remaining content in chunks of ~700 lines each.'", filePath) | ||
| } else { | ||
| errorMsg = "echo '[WRITE TOOL ERROR] The file content is too large to be transmitted by the upstream API. The Write tool input was truncated. You MUST retry by writing the file in smaller chunks: First use Write to create the file with the first 700 lines, then use multiple Edit operations to append the remaining content in chunks of ~700 lines each.'" | ||
| } | ||
|
|
||
| log.Warnf("kiro: converting truncated write to Bash echo for file: %s", filePath) | ||
|
|
||
| hasToolUses = true | ||
|
|
||
| // Close text block if open | ||
| if isTextBlockOpen && contentBlockIndex >= 0 { | ||
| blockStop := kiroclaude.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")} | ||
| } | ||
| } | ||
| isTextBlockOpen = false | ||
| } | ||
|
|
||
| contentBlockIndex++ | ||
|
|
||
| // Emit as Bash tool_use | ||
| blockStart := kiroclaude.BuildClaudeContentBlockStartEvent(contentBlockIndex, "tool_use", tu.ToolUseID, "Bash") | ||
| 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")} | ||
| } | ||
| } | ||
|
|
||
| // Emit the Bash command as input | ||
| bashInput := map[string]interface{}{ | ||
| "command": errorMsg, | ||
| } | ||
| inputJSON, _ := json.Marshal(bashInput) | ||
| inputDelta := kiroclaude.BuildClaudeInputJsonDeltaEvent(string(inputJSON), 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")} | ||
| } | ||
| } | ||
|
|
||
| blockStop := kiroclaude.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")} | ||
| } | ||
| } | ||
|
|
||
| continue // Skip the normal tool_use emission |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The logic within this block to handle __truncated_write__ markers is a bit repetitive. There are a couple of opportunities to refactor for improved readability and maintainability:
- Error Message Construction: The
errorMsgstrings for cases with and without afilePathare very similar. You can reduce duplication by defining a common guidance constant. - SSE Event Sending: The pattern for translating and sending SSE events is repeated multiple times. This could be extracted into a local helper function.
Here's an example of how you could refactor this block:
// At the top of streamToChannel or as a package constant
const writeToolTruncationGuidance = "You MUST retry by writing the file in smaller chunks: First use Write to create the file with the first 700 lines, then use multiple Edit operations to append the remaining content in chunks of ~700 lines each."
// Inside the loop
if tu.Name == "__truncated_write__" {
// ...
// Helper for sending SSE events
sendSSE := func(event []byte) {
sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, event, &translatorParam)
for _, chunk := range sseData {
if chunk != "" {
out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
}
}
}
// Refactored error message
var errorDetail string
if filePath != "" {
errorDetail = fmt.Sprintf("The file content for \"%s\" is too large to be transmitted by the upstream API.", filePath)
} else {
errorDetail = "The file content is too large to be transmitted by the upstream API. The Write tool input was truncated."
}
errorMsg := fmt.Sprintf("echo '[WRITE TOOL ERROR] %s %s'", errorDetail, writeToolTruncationGuidance)
log.Warnf("kiro: converting truncated write to Bash echo for file: %s", filePath)
hasToolUses = true
// Close text block if open
if isTextBlockOpen && contentBlockIndex >= 0 {
sendSSE(kiroclaude.BuildClaudeContentBlockStopEvent(contentBlockIndex))
isTextBlockOpen = false
}
contentBlockIndex++
// Emit as Bash tool_use
sendSSE(kiroclaude.BuildClaudeContentBlockStartEvent(contentBlockIndex, "tool_use", tu.ToolUseID, "Bash"))
// Emit the Bash command as input
// ... (rest of the logic using sendSSE)
continue
}| filePath := "" | ||
| if idx := strings.Index(fullInput, "file_path"); idx >= 0 { | ||
| // Try to extract the file path value | ||
| rest := fullInput[idx:] | ||
| if colonIdx := strings.Index(rest, ":"); colonIdx >= 0 { | ||
| rest = strings.TrimSpace(rest[colonIdx+1:]) | ||
| if len(rest) > 0 && rest[0] == '"' { | ||
| rest = rest[1:] | ||
| if endQuote := strings.Index(rest, "\""); endQuote >= 0 { | ||
| filePath = rest[:endQuote] | ||
| } | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The manual string manipulation to extract filePath from the malformed JSON is a bit brittle and could fail if the spacing or structure of the truncated input changes slightly.
Using a regular expression would provide a more robust way to extract the file_path value. You could define a package-level regex for performance.
Example:
// At package level
var filePathRegex = regexp.MustCompile(`"file_path"\s*:\s*"([^"]+)"`)
// In function
filePath := ""
if matches := filePathRegex.FindStringSubmatch(fullInput); len(matches) > 1 {
filePath = matches[1]
}This would make the extraction more resilient to variations in whitespace and surrounding text.
| // Check for Write tool with empty or missing input - this happens when Kiro API | ||
| // completely skips sending input for large file writes | ||
| if currentToolUse.Name == "Write" && len(strings.TrimSpace(fullInput)) == 0 { | ||
| log.Warnf("kiro: Write tool received no input from upstream API. The file content may be too large to transmit.") | ||
| // Return nil to skip this tool use - it will be handled as a truncation error | ||
| // The caller should emit a text block explaining the error instead | ||
| if processedIDs != nil { | ||
| processedIDs[currentToolUse.ToolUseID] = true | ||
| } | ||
| log.Infof("kiro: skipping Write tool use %s due to empty input (content too large)", currentToolUse.ToolUseID) | ||
| // Return a special marker tool use that indicates truncation | ||
| toolUse := KiroToolUse{ | ||
| ToolUseID: currentToolUse.ToolUseID, | ||
| Name: "__truncated_write__", // Special marker name | ||
| Input: map[string]interface{}{ | ||
| "error": "Write tool input was not transmitted by upstream API. The file content is too large.", | ||
| }, | ||
| } | ||
| toolUses = append(toolUses, toolUse) | ||
| return toolUses, nil | ||
| } | ||
|
|
||
| // Repair and parse the accumulated JSON | ||
| repairedJSON := RepairJSON(fullInput) | ||
| var finalInput map[string]interface{} | ||
| if err := json.Unmarshal([]byte(repairedJSON), &finalInput); err != nil { | ||
| log.Warnf("kiro: failed to parse accumulated tool input: %v, raw: %s", err, fullInput) | ||
| finalInput = make(map[string]interface{}) | ||
|
|
||
| // Check if this is a Write tool with truncated input (missing content field) | ||
| // This happens when the Kiro API truncates large tool inputs | ||
| if currentToolUse.Name == "Write" && strings.Contains(fullInput, "file_path") && !strings.Contains(fullInput, "content") { | ||
| log.Warnf("kiro: Write tool input was truncated by upstream API (content field missing). The file content may be too large.") | ||
| // Extract file_path if possible for error context | ||
| filePath := "" | ||
| if idx := strings.Index(fullInput, "file_path"); idx >= 0 { | ||
| // Try to extract the file path value | ||
| rest := fullInput[idx:] | ||
| if colonIdx := strings.Index(rest, ":"); colonIdx >= 0 { | ||
| rest = strings.TrimSpace(rest[colonIdx+1:]) | ||
| if len(rest) > 0 && rest[0] == '"' { | ||
| rest = rest[1:] | ||
| if endQuote := strings.Index(rest, "\""); endQuote >= 0 { | ||
| filePath = rest[:endQuote] | ||
| } | ||
| } | ||
| } | ||
| } | ||
| if processedIDs != nil { | ||
| processedIDs[currentToolUse.ToolUseID] = true | ||
| } | ||
| // Return a special marker tool use that indicates truncation | ||
| toolUse := KiroToolUse{ | ||
| ToolUseID: currentToolUse.ToolUseID, | ||
| Name: "__truncated_write__", // Special marker name | ||
| Input: map[string]interface{}{ | ||
| "error": "Write tool content was truncated by upstream API. The file content is too large.", | ||
| "file_path": filePath, | ||
| }, | ||
| } | ||
| toolUses = append(toolUses, toolUse) | ||
| return toolUses, nil | ||
| } | ||
| } | ||
|
|
||
| // Additional check: Write tool parsed successfully but missing content field | ||
| if currentToolUse.Name == "Write" { | ||
| if _, hasContent := finalInput["content"]; !hasContent { | ||
| if filePath, hasPath := finalInput["file_path"]; hasPath { | ||
| log.Warnf("kiro: Write tool input missing 'content' field, likely truncated by upstream API") | ||
| if processedIDs != nil { | ||
| processedIDs[currentToolUse.ToolUseID] = true | ||
| } | ||
| // Return a special marker tool use that indicates truncation | ||
| toolUse := KiroToolUse{ | ||
| ToolUseID: currentToolUse.ToolUseID, | ||
| Name: "__truncated_write__", // Special marker name | ||
| Input: map[string]interface{}{ | ||
| "error": "Write tool content field was missing. The file content is too large.", | ||
| "file_path": filePath, | ||
| }, | ||
| } | ||
| toolUses = append(toolUses, toolUse) | ||
| return toolUses, nil | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There's some code duplication across the three scenarios where a __truncated_write__ marker is created. Each block handles marking the tool ID as processed and then constructs and appends the special KiroToolUse marker.
To improve maintainability, you could extract this logic into a helper function. This function would take the necessary context (like currentToolUse, processedIDs, error message, and file path) and return the []KiroToolUse and *ToolUseState to be returned by ProcessToolUseEvent.
Here's a conceptual example of such a helper:
func handleTruncatedWrite(currentToolUse *ToolUseState, processedIDs map[string]bool, errorMsg string, filePath interface{}) ([]KiroToolUse, *ToolUseState) {
if processedIDs != nil {
processedIDs[currentToolUse.ToolUseID] = true
}
input := map[string]interface{}{"error": errorMsg}
if fp, ok := filePath.(string); ok && fp != "" {
input["file_path"] = fp
}
toolUse := KiroToolUse{
ToolUseID: currentToolUse.ToolUseID,
Name: "__truncated_write__",
Input: input,
}
return []KiroToolUse{toolUse}, nil
}Each of the three detection blocks could then call this helper, simplifying the main function's logic.
Problem
When the Kiro/AWS CodeWhisperer API receives a Write tool request with content that exceeds transmission limits, it truncates the tool input. This results in:
contentfield in the parsed JSONThis caused Claude Code to fail with validation errors and multiple sequential failed Write tool calls when attempting to write large files.
Workaround
This is a workaround solution that detects truncation scenarios and converts them to Bash tool calls that echo an error message. This allows:
Instead of silently failing or showing errors to the user, the agent now receives clear guidance and can automatically retry with smaller file chunks.
Changes
internal/translator/kiro/claude/kiro_claude_tools.goDetects three truncation scenarios in
ProcessToolUseEvent:When detected, emits a special
__truncated_write__marker tool use.internal/runtime/executor/kiro_executor.goHandles
__truncated_write__markers instreamToChannel:hasToolUses=trueto ensurestop_reason='tool_use'for agent continuationTesting