fix(kiro): Handle tool results correctly in OpenAI format translation #61
Conversation
Fix three issues in Kiro OpenAI translator that caused "Improperly formed request" errors when processing LiteLLM-translated requests with tool_use/tool_result: 1. Skip merging tool role messages in MergeAdjacentMessages() to preserve individual tool_call_id fields 2. Track pendingToolResults and attach to the next user message instead of only the last message. Create synthetic user message when conversation ends with tool results. 3. Insert synthetic user message with tool results before assistant messages to maintain proper alternating user/assistant structure. This fixes the case where LiteLLM translates Anthropic user messages containing only tool_result blocks into tool role messages followed by assistant. Adds unit tests covering all tool result handling scenarios.
Remove dead code that was never used: - toolCallIDToName map: built but never read from - seenToolCallIDs: declared but never populated, only suppressed with _
Summary of ChangesHello @tinyc0der, 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 resolves a critical bug in the Kiro OpenAI translator that caused API requests with 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
|
![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 effectively addresses a bug in handling tool results within the Kiro OpenAI translator. The core logic change, which involves introducing pendingToolResults to correctly associate tool results with subsequent messages, is well-implemented and handles various edge cases like trailing tool results or tool results followed by an assistant message. The dead code removal is a welcome cleanup. The addition of a comprehensive test suite with six new test cases significantly improves the robustness of the translator. My feedback focuses on making a few of the new tests even more specific and resilient to future changes.
| import ( | ||
| "encoding/json" | ||
| "testing" | ||
| ) |
There was a problem hiding this comment.
To support a more robust assertion in TestAssistantEndsConversation, the strings package should be imported.
| import ( | |
| "encoding/json" | |
| "testing" | |
| ) | |
| import ( | |
| "encoding/json" | |
| "strings" | |
| "testing" | |
| ) |
| // Check if there are any tool results anywhere | ||
| var totalToolResults int | ||
| for i, h := range payload.ConversationState.History { | ||
| if h.UserInputMessage != nil && h.UserInputMessage.UserInputMessageContext != nil { | ||
| count := len(h.UserInputMessage.UserInputMessageContext.ToolResults) | ||
| t.Logf("History[%d] user message has %d tool results", i, count) | ||
| totalToolResults += count | ||
| } | ||
| } | ||
|
|
||
| ctx := payload.ConversationState.CurrentMessage.UserInputMessage.UserInputMessageContext | ||
| if ctx != nil { | ||
| t.Logf("CurrentMessage has %d tool results", len(ctx.ToolResults)) | ||
| totalToolResults += len(ctx.ToolResults) | ||
| } else { | ||
| t.Logf("CurrentMessage has no UserInputMessageContext") | ||
| } | ||
|
|
||
| if totalToolResults != 2 { | ||
| t.Errorf("Expected 2 tool results total, got %d", totalToolResults) | ||
| } |
There was a problem hiding this comment.
The current assertion in this test checks for the total number of tool results but doesn't verify their precise location. A more specific assertion would ensure that both tool results are correctly attached to the currentMessage and not scattered throughout the history. This will make the test more robust against future refactoring.
// Check that the two tool results are attached to the current user message.
ctx := payload.ConversationState.CurrentMessage.UserInputMessage.UserInputMessageContext
if ctx == nil {
t.Fatal("Expected currentMessage to have UserInputMessageContext with tool results")
}
if len(ctx.ToolResults) != 2 {
t.Fatalf("Expected 2 tool results in currentMessage, got %d", len(ctx.ToolResults))
}
// Also ensure no tool results were mistakenly added to history messages.
for _, h := range payload.ConversationState.History {
if h.UserInputMessage != nil && h.UserInputMessage.UserInputMessageContext != nil && len(h.UserInputMessage.UserInputMessageContext.ToolResults) > 0 {
t.Errorf("Did not expect tool results in history, but found %d", len(h.UserInputMessage.UserInputMessageContext.ToolResults))
}
}| // Tool results should be attached to a synthetic user message or the history should be valid | ||
| var totalToolResults int | ||
| for i, h := range payload.ConversationState.History { | ||
| if h.UserInputMessage != nil { | ||
| t.Logf("History[%d]: user message content=%q", i, h.UserInputMessage.Content) | ||
| if h.UserInputMessage.UserInputMessageContext != nil { | ||
| count := len(h.UserInputMessage.UserInputMessageContext.ToolResults) | ||
| t.Logf(" Has %d tool results", count) | ||
| totalToolResults += count | ||
| } | ||
| } | ||
| if h.AssistantResponseMessage != nil { | ||
| t.Logf("History[%d]: assistant message content=%q", i, h.AssistantResponseMessage.Content) | ||
| } | ||
| } | ||
|
|
||
| ctx := payload.ConversationState.CurrentMessage.UserInputMessage.UserInputMessageContext | ||
| if ctx != nil { | ||
| t.Logf("CurrentMessage has %d tool results", len(ctx.ToolResults)) | ||
| totalToolResults += len(ctx.ToolResults) | ||
| } | ||
|
|
||
| if totalToolResults != 2 { | ||
| t.Errorf("Expected 2 tool results total, got %d", totalToolResults) | ||
| } |
There was a problem hiding this comment.
Similar to a previous comment, this test verifies the total count of tool results but doesn't check their specific location. For this test case, the tool results should be in a synthetic user message within the history. Making the assertion more specific will strengthen this test.
// Tool results should be attached to a synthetic user message in the history.
var foundSyntheticMessageWithTools bool
for _, h := range payload.ConversationState.History {
if h.UserInputMessage != nil && h.UserInputMessage.Content == "Tool results provided." {
if h.UserInputMessage.UserInputMessageContext != nil && len(h.UserInputMessage.UserInputMessageContext.ToolResults) == 2 {
foundSyntheticMessageWithTools = true
}
}
}
if !foundSyntheticMessageWithTools {
t.Error("Expected to find a synthetic user message with 2 tool results in history, but did not.")
}| if payload.ConversationState.CurrentMessage.UserInputMessage.Content == "" { | ||
| t.Error("Expected a 'Continue' message to be created when assistant is last") | ||
| } |
There was a problem hiding this comment.
The current assertion payload.ConversationState.CurrentMessage.UserInputMessage.Content == "" passes because the content is not empty, but it's not a very precise check. The final content will include a system prompt in addition to the "Continue" message. A more robust check would be to verify that the content contains "Continue".
if !strings.Contains(payload.ConversationState.CurrentMessage.UserInputMessage.Content, "Continue") {
t.Error("Expected a 'Continue' message to be created when assistant is last")
}
2 participants
Summary
Problem
Requests with
tool_use/tool_resultblocks in the conversation history were failing with "Improperly formed request" error from AWS Kiro API.Root Cause: The Kiro OpenAI translator didn't properly track tool results from
toolrole messages and attach them to the correct user message in history.Reproduction
Save this as
test-bug.json:{ "model": "kiro-claude-opus-4-5-agentic", "max_tokens": 1024, "messages": [ { "role": "user", "content": "What is 2+2?" }, { "role": "assistant", "content": "I'll calculate that for you.", "tool_calls": [ { "type": "function", "id": "calc_001", "function": { "name": "calculator", "arguments": "{\"expression\": \"2+2\"}" } } ] }, { "role": "tool", "content": "4", "tool_call_id": "calc_001" } ], "tools": [ { "type": "function", "function": { "name": "calculator", "description": "Simple calculator", "parameters": { "type": "object", "properties": { "expression": {"type": "string"} } } } } ] }Before fix (FAIL):
After fix (SUCCESS):
Fix
MergeAdjacentMessages()to preserve individualtool_call_idfieldspendingToolResultsand attach to the next user message instead of only the last messageTest plan
Files Modified
internal/translator/kiro/common/message_merge.go- Skip merging tool messagesinternal/translator/kiro/openai/kiro_openai_request.go- Fix tool result attachment + remove dead codeinternal/translator/kiro/openai/kiro_openai_request_test.go- Add unit tests