fix(kiro): Handle tool results correctly in OpenAI format translation

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.
This commit is contained in:
TinyCoder
2025-12-24 09:58:34 +07:00
parent ee6fc4e8a1
commit 36a512fdf2
3 changed files with 436 additions and 4 deletions

View File

@@ -10,6 +10,7 @@ import (
// MergeAdjacentMessages merges adjacent messages with the same role.
// This reduces API call complexity and improves compatibility.
// Based on AIClient-2-API implementation.
// NOTE: Tool messages are NOT merged because each has a unique tool_call_id that must be preserved.
func MergeAdjacentMessages(messages []gjson.Result) []gjson.Result {
if len(messages) <= 1 {
return messages
@@ -26,6 +27,12 @@ func MergeAdjacentMessages(messages []gjson.Result) []gjson.Result {
currentRole := msg.Get("role").String()
lastRole := lastMsg.Get("role").String()
// Don't merge tool messages - each has a unique tool_call_id
if currentRole == "tool" || lastRole == "tool" {
merged = append(merged, msg)
continue
}
if currentRole == lastRole {
// Merge content from current message into last message
mergedContent := mergeMessageContent(lastMsg, msg)

View File

@@ -469,6 +469,11 @@ func processOpenAIMessages(messages gjson.Result, modelID, origin string) ([]Kir
}
}
// Track pending tool results that should be attached to the next user message
// This is critical for LiteLLM-translated requests where tool results appear
// as separate "tool" role messages between assistant and user messages
var pendingToolResults []KiroToolResult
for i, msg := range messagesArray {
role := msg.Get("role").String()
isLastMessage := i == len(messagesArray)-1
@@ -480,6 +485,10 @@ func processOpenAIMessages(messages gjson.Result, modelID, origin string) ([]Kir
case "user":
userMsg, toolResults := buildUserMessageFromOpenAI(msg, modelID, origin)
// Merge any pending tool results from preceding "tool" role messages
toolResults = append(pendingToolResults, toolResults...)
pendingToolResults = nil // Reset pending tool results
if isLastMessage {
currentUserMsg = &userMsg
currentToolResults = toolResults
@@ -505,6 +514,24 @@ func processOpenAIMessages(messages gjson.Result, modelID, origin string) ([]Kir
case "assistant":
assistantMsg := buildAssistantMessageFromOpenAI(msg)
// If there are pending tool results, we need to insert a synthetic user message
// before this assistant message to maintain proper conversation structure
if len(pendingToolResults) > 0 {
syntheticUserMsg := KiroUserInputMessage{
Content: "Tool results provided.",
ModelID: modelID,
Origin: origin,
UserInputMessageContext: &KiroUserInputMessageContext{
ToolResults: pendingToolResults,
},
}
history = append(history, KiroHistoryMessage{
UserInputMessage: &syntheticUserMsg,
})
pendingToolResults = nil
}
if isLastMessage {
history = append(history, KiroHistoryMessage{
AssistantResponseMessage: &assistantMsg,
@@ -524,7 +551,7 @@ func processOpenAIMessages(messages gjson.Result, modelID, origin string) ([]Kir
case "tool":
// Tool messages in OpenAI format provide results for tool_calls
// These are typically followed by user or assistant messages
// Process them and merge into the next user message's tool results
// Collect them as pending and attach to the next user message
toolCallID := msg.Get("tool_call_id").String()
content := msg.Get("content").String()
@@ -534,9 +561,21 @@ func processOpenAIMessages(messages gjson.Result, modelID, origin string) ([]Kir
Content: []KiroTextContent{{Text: content}},
Status: "success",
}
// Tool results should be included in the next user message
// For now, collect them and they'll be handled when we build the current message
currentToolResults = append(currentToolResults, toolResult)
// Collect pending tool results to attach to the next user message
pendingToolResults = append(pendingToolResults, toolResult)
}
}
}
// Handle case where tool results are at the end with no following user message
if len(pendingToolResults) > 0 {
currentToolResults = append(currentToolResults, pendingToolResults...)
// If there's no current user message, create a synthetic one for the tool results
if currentUserMsg == nil {
currentUserMsg = &KiroUserInputMessage{
Content: "Tool results provided.",
ModelID: modelID,
Origin: origin,
}
}
}

View File

@@ -0,0 +1,386 @@
package openai
import (
"encoding/json"
"testing"
)
// TestToolResultsAttachedToCurrentMessage verifies that tool results from "tool" role messages
// are properly attached to the current user message (the last message in the conversation).
// This is critical for LiteLLM-translated requests where tool results appear as separate messages.
func TestToolResultsAttachedToCurrentMessage(t *testing.T) {
// OpenAI format request simulating LiteLLM's translation from Anthropic format
// Sequence: user -> assistant (with tool_calls) -> tool (result) -> user
// The last user message should have the tool results attached
input := []byte(`{
"model": "kiro-claude-opus-4-5-agentic",
"messages": [
{"role": "user", "content": "Hello, can you read a file for me?"},
{
"role": "assistant",
"content": "I'll read that file for you.",
"tool_calls": [
{
"id": "call_abc123",
"type": "function",
"function": {
"name": "Read",
"arguments": "{\"file_path\": \"/tmp/test.txt\"}"
}
}
]
},
{
"role": "tool",
"tool_call_id": "call_abc123",
"content": "File contents: Hello World!"
},
{"role": "user", "content": "What did the file say?"}
]
}`)
result, _ := BuildKiroPayloadFromOpenAI(input, "kiro-model", "", "CLI", false, false, nil, nil)
var payload KiroPayload
if err := json.Unmarshal(result, &payload); err != nil {
t.Fatalf("Failed to unmarshal result: %v", err)
}
// The last user message becomes currentMessage
// History should have: user (first), assistant (with tool_calls)
t.Logf("History count: %d", len(payload.ConversationState.History))
if len(payload.ConversationState.History) != 2 {
t.Errorf("Expected 2 history entries (user + assistant), got %d", len(payload.ConversationState.History))
}
// Tool results should be attached to currentMessage (the last 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) != 1 {
t.Fatalf("Expected 1 tool result in currentMessage, got %d", len(ctx.ToolResults))
}
tr := ctx.ToolResults[0]
if tr.ToolUseID != "call_abc123" {
t.Errorf("Expected toolUseId 'call_abc123', got '%s'", tr.ToolUseID)
}
if len(tr.Content) == 0 || tr.Content[0].Text != "File contents: Hello World!" {
t.Errorf("Tool result content mismatch, got: %+v", tr.Content)
}
}
// TestToolResultsInHistoryUserMessage verifies that when there are multiple user messages
// after tool results, the tool results are attached to the correct user message in history.
func TestToolResultsInHistoryUserMessage(t *testing.T) {
// Sequence: user -> assistant (with tool_calls) -> tool (result) -> user -> assistant -> user
// The first user after tool should have tool results in history
input := []byte(`{
"model": "kiro-claude-opus-4-5-agentic",
"messages": [
{"role": "user", "content": "Hello"},
{
"role": "assistant",
"content": "I'll read the file.",
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {
"name": "Read",
"arguments": "{}"
}
}
]
},
{
"role": "tool",
"tool_call_id": "call_1",
"content": "File result"
},
{"role": "user", "content": "Thanks for the file"},
{"role": "assistant", "content": "You're welcome"},
{"role": "user", "content": "Bye"}
]
}`)
result, _ := BuildKiroPayloadFromOpenAI(input, "kiro-model", "", "CLI", false, false, nil, nil)
var payload KiroPayload
if err := json.Unmarshal(result, &payload); err != nil {
t.Fatalf("Failed to unmarshal result: %v", err)
}
// History should have: user, assistant, user (with tool results), assistant
// CurrentMessage should be: last user "Bye"
t.Logf("History count: %d", len(payload.ConversationState.History))
// Find the user message in history with tool results
foundToolResults := false
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 {
if len(h.UserInputMessage.UserInputMessageContext.ToolResults) > 0 {
foundToolResults = true
t.Logf(" Found %d tool results", len(h.UserInputMessage.UserInputMessageContext.ToolResults))
tr := h.UserInputMessage.UserInputMessageContext.ToolResults[0]
if tr.ToolUseID != "call_1" {
t.Errorf("Expected toolUseId 'call_1', got '%s'", tr.ToolUseID)
}
}
}
}
if h.AssistantResponseMessage != nil {
t.Logf("History[%d]: assistant message content=%q", i, h.AssistantResponseMessage.Content)
}
}
if !foundToolResults {
t.Error("Tool results were not attached to any user message in history")
}
}
// TestToolResultsWithMultipleToolCalls verifies handling of multiple tool calls
func TestToolResultsWithMultipleToolCalls(t *testing.T) {
input := []byte(`{
"model": "kiro-claude-opus-4-5-agentic",
"messages": [
{"role": "user", "content": "Read two files for me"},
{
"role": "assistant",
"content": "I'll read both files.",
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {
"name": "Read",
"arguments": "{\"file_path\": \"/tmp/file1.txt\"}"
}
},
{
"id": "call_2",
"type": "function",
"function": {
"name": "Read",
"arguments": "{\"file_path\": \"/tmp/file2.txt\"}"
}
}
]
},
{
"role": "tool",
"tool_call_id": "call_1",
"content": "Content of file 1"
},
{
"role": "tool",
"tool_call_id": "call_2",
"content": "Content of file 2"
},
{"role": "user", "content": "What do they say?"}
]
}`)
result, _ := BuildKiroPayloadFromOpenAI(input, "kiro-model", "", "CLI", false, false, nil, nil)
var payload KiroPayload
if err := json.Unmarshal(result, &payload); err != nil {
t.Fatalf("Failed to unmarshal result: %v", err)
}
t.Logf("History count: %d", len(payload.ConversationState.History))
t.Logf("CurrentMessage content: %q", payload.ConversationState.CurrentMessage.UserInputMessage.Content)
// 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)
}
}
// TestToolResultsAtEndOfConversation verifies tool results are handled when
// the conversation ends with tool results (no following user message)
func TestToolResultsAtEndOfConversation(t *testing.T) {
input := []byte(`{
"model": "kiro-claude-opus-4-5-agentic",
"messages": [
{"role": "user", "content": "Read a file"},
{
"role": "assistant",
"content": "Reading the file.",
"tool_calls": [
{
"id": "call_end",
"type": "function",
"function": {
"name": "Read",
"arguments": "{\"file_path\": \"/tmp/test.txt\"}"
}
}
]
},
{
"role": "tool",
"tool_call_id": "call_end",
"content": "File contents here"
}
]
}`)
result, _ := BuildKiroPayloadFromOpenAI(input, "kiro-model", "", "CLI", false, false, nil, nil)
var payload KiroPayload
if err := json.Unmarshal(result, &payload); err != nil {
t.Fatalf("Failed to unmarshal result: %v", err)
}
// When the last message is a tool result, a synthetic user message is created
// and tool results should be attached to it
ctx := payload.ConversationState.CurrentMessage.UserInputMessage.UserInputMessageContext
if ctx == nil || len(ctx.ToolResults) == 0 {
t.Error("Expected tool results to be attached to current message when conversation ends with tool result")
} else {
if ctx.ToolResults[0].ToolUseID != "call_end" {
t.Errorf("Expected toolUseId 'call_end', got '%s'", ctx.ToolResults[0].ToolUseID)
}
}
}
// TestToolResultsFollowedByAssistant verifies handling when tool results are followed
// by an assistant message (no intermediate user message).
// This is the pattern from LiteLLM translation of Anthropic format where:
// user message has ONLY tool_result blocks -> LiteLLM creates tool messages
// then the next message is assistant
func TestToolResultsFollowedByAssistant(t *testing.T) {
// Sequence: user -> assistant (with tool_calls) -> tool -> tool -> assistant -> user
// This simulates LiteLLM's translation of:
// user: "Read files"
// assistant: [tool_use, tool_use]
// user: [tool_result, tool_result] <- becomes multiple "tool" role messages
// assistant: "I've read them"
// user: "What did they say?"
input := []byte(`{
"model": "kiro-claude-opus-4-5-agentic",
"messages": [
{"role": "user", "content": "Read two files for me"},
{
"role": "assistant",
"content": "I'll read both files.",
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {
"name": "Read",
"arguments": "{\"file_path\": \"/tmp/a.txt\"}"
}
},
{
"id": "call_2",
"type": "function",
"function": {
"name": "Read",
"arguments": "{\"file_path\": \"/tmp/b.txt\"}"
}
}
]
},
{
"role": "tool",
"tool_call_id": "call_1",
"content": "Contents of file A"
},
{
"role": "tool",
"tool_call_id": "call_2",
"content": "Contents of file B"
},
{
"role": "assistant",
"content": "I've read both files."
},
{"role": "user", "content": "What did they say?"}
]
}`)
result, _ := BuildKiroPayloadFromOpenAI(input, "kiro-model", "", "CLI", false, false, nil, nil)
var payload KiroPayload
if err := json.Unmarshal(result, &payload); err != nil {
t.Fatalf("Failed to unmarshal result: %v", err)
}
t.Logf("History count: %d", len(payload.ConversationState.History))
// 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)
}
}
// TestAssistantEndsConversation verifies handling when assistant is the last message
func TestAssistantEndsConversation(t *testing.T) {
input := []byte(`{
"model": "kiro-claude-opus-4-5-agentic",
"messages": [
{"role": "user", "content": "Hello"},
{
"role": "assistant",
"content": "Hi there!"
}
]
}`)
result, _ := BuildKiroPayloadFromOpenAI(input, "kiro-model", "", "CLI", false, false, nil, nil)
var payload KiroPayload
if err := json.Unmarshal(result, &payload); err != nil {
t.Fatalf("Failed to unmarshal result: %v", err)
}
// When assistant is last, a "Continue" user message should be created
if payload.ConversationState.CurrentMessage.UserInputMessage.Content == "" {
t.Error("Expected a 'Continue' message to be created when assistant is last")
}
}