Skip to content

Commit fb02f9b

Browse files
committed
fix: use correct slice index for Bedrock tool result coalescing
When system messages are present in the request body, promoteAnthropicSystemMessagesToParam filters them out, producing a shorter messages slice. The coalescing loop was using body.Messages[i+1] (the original unfiltered slice) instead of messages[i+1] (the filtered slice), causing index misalignment that could read wrong messages or panic with out-of-bounds access. Additionally, tool_result messages have Role user, so they entered the MessageRoleUser branch and went through convertUserMessage() which does NOT coalesce consecutive tool results. The coalescing logic in the default branch was unreachable for this role. Moved the hasToolResult check into the MessageRoleUser branch so consecutive tool_result messages are properly coalesced. Updated existing test to reflect the now-working coalescing behavior, and added a regression test with system messages to verify index alignment. Signed-off-by: liuhy <liuhongyu@apache.org>
1 parent 2e23b28 commit fb02f9b

3 files changed

Lines changed: 159 additions & 41 deletions

File tree

.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,3 +52,4 @@ inference-extension-conformance-test-report.yaml
5252
/aigw
5353
site/.cursorrules
5454
/.cursor
55+
.worktrees/

internal/translator/anthropic_awsbedrock.go

Lines changed: 16 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -78,36 +78,33 @@ func (a *anthropicToAWSBedrockTranslator) RequestBody(_ []byte, body *anthropics
7878
msg := &messages[i]
7979
switch msg.Role {
8080
case anthropicschema.MessageRoleUser:
81-
bedrockMsg, convErr := a.convertUserMessage(msg)
82-
if convErr != nil {
83-
return nil, nil, convErr
84-
}
85-
bedrockReq.Messages = append(bedrockReq.Messages, bedrockMsg)
86-
i++
87-
case anthropicschema.MessageRoleAssistant:
88-
bedrockMsg, convErr := a.convertAssistantMessage(msg)
89-
if convErr != nil {
90-
return nil, nil, convErr
91-
}
92-
bedrockReq.Messages = append(bedrockReq.Messages, bedrockMsg)
93-
i++
94-
default:
95-
// Check for tool_result content blocks (these come as "user" role in Anthropic
96-
// but we handle them specially).
9781
if hasToolResult(msg) {
9882
bedrockMsg := a.convertToolResultMessage(msg)
9983
// Coalesce consecutive tool result messages.
100-
for i+1 < msgLen && hasToolResult(&body.Messages[i+1]) {
101-
nextMsg := &body.Messages[i+1]
84+
for i+1 < msgLen && hasToolResult(&messages[i+1]) {
85+
nextMsg := &messages[i+1]
10286
nextBedrockMsg := a.convertToolResultMessage(nextMsg)
10387
bedrockMsg.Content = append(bedrockMsg.Content, nextBedrockMsg.Content...)
10488
i++
10589
}
10690
bedrockReq.Messages = append(bedrockReq.Messages, bedrockMsg)
10791
} else {
108-
return nil, nil, fmt.Errorf("%w: unexpected role: %s", internalapi.ErrInvalidRequestBody, msg.Role)
92+
bedrockMsg, convErr := a.convertUserMessage(msg)
93+
if convErr != nil {
94+
return nil, nil, convErr
95+
}
96+
bedrockReq.Messages = append(bedrockReq.Messages, bedrockMsg)
10997
}
11098
i++
99+
case anthropicschema.MessageRoleAssistant:
100+
bedrockMsg, convErr := a.convertAssistantMessage(msg)
101+
if convErr != nil {
102+
return nil, nil, convErr
103+
}
104+
bedrockReq.Messages = append(bedrockReq.Messages, bedrockMsg)
105+
i++
106+
default:
107+
return nil, nil, fmt.Errorf("%w: unexpected role: %s", internalapi.ErrInvalidRequestBody, msg.Role)
111108
}
112109
}
113110

internal/translator/anthropic_awsbedrock_test.go

Lines changed: 142 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1032,35 +1032,155 @@ func TestAnthropicToAWSBedrockTranslator_RequestBody_ToolResultMessages(t *testi
10321032
err = json.Unmarshal(body, &bedrockReq)
10331033
require.NoError(t, err)
10341034

1035-
// Should be: user, assistant, user (tool result 1), user (tool result 2).
1036-
// Tool result messages with role "user" go through convertUserMessage individually.
1037-
require.Len(t, bedrockReq.Messages, 4)
1038-
1039-
// First tool result message.
1040-
toolResultMsg1 := bedrockReq.Messages[2]
1041-
assert.Equal(t, awsbedrock.ConversationRoleUser, toolResultMsg1.Role)
1042-
require.Len(t, toolResultMsg1.Content, 1)
1043-
require.NotNil(t, toolResultMsg1.Content[0].ToolResult)
1044-
assert.Equal(t, "tu_abc", *toolResultMsg1.Content[0].ToolResult.ToolUseID)
1045-
require.Len(t, toolResultMsg1.Content[0].ToolResult.Content, 1)
1046-
assert.Equal(t, "72°F and sunny", *toolResultMsg1.Content[0].ToolResult.Content[0].Text)
1047-
1048-
// Second tool result message (error).
1049-
toolResultMsg2 := bedrockReq.Messages[3]
1050-
assert.Equal(t, awsbedrock.ConversationRoleUser, toolResultMsg2.Role)
1051-
require.Len(t, toolResultMsg2.Content, 1)
1052-
require.NotNil(t, toolResultMsg2.Content[0].ToolResult)
1053-
assert.Equal(t, "tu_def", *toolResultMsg2.Content[0].ToolResult.ToolUseID)
1054-
assert.Equal(t, "error", *toolResultMsg2.Content[0].ToolResult.Status)
1055-
require.Len(t, toolResultMsg2.Content[0].ToolResult.Content, 1)
1056-
assert.Equal(t, "Error: city not found", *toolResultMsg2.Content[0].ToolResult.Content[0].Text)
1035+
// Should be: user, assistant, user (coalesced tool results).
1036+
// Consecutive tool_result messages are coalesced into a single Bedrock message.
1037+
require.Len(t, bedrockReq.Messages, 3)
1038+
1039+
// Coalesced tool result message contains both tool results.
1040+
toolResultMsg := bedrockReq.Messages[2]
1041+
assert.Equal(t, awsbedrock.ConversationRoleUser, toolResultMsg.Role)
1042+
require.Len(t, toolResultMsg.Content, 2)
1043+
1044+
// First tool result.
1045+
require.NotNil(t, toolResultMsg.Content[0].ToolResult)
1046+
assert.Equal(t, "tu_abc", *toolResultMsg.Content[0].ToolResult.ToolUseID)
1047+
require.Len(t, toolResultMsg.Content[0].ToolResult.Content, 1)
1048+
assert.Equal(t, "72°F and sunny", *toolResultMsg.Content[0].ToolResult.Content[0].Text)
1049+
1050+
// Second tool result (error) coalesced into the same message.
1051+
require.NotNil(t, toolResultMsg.Content[1].ToolResult)
1052+
assert.Equal(t, "tu_def", *toolResultMsg.Content[1].ToolResult.ToolUseID)
1053+
assert.Equal(t, "error", *toolResultMsg.Content[1].ToolResult.Status)
1054+
require.Len(t, toolResultMsg.Content[1].ToolResult.Content, 1)
1055+
assert.Equal(t, "Error: city not found", *toolResultMsg.Content[1].ToolResult.Content[0].Text)
10571056

10581057
// Verify tool choice "any".
10591058
require.NotNil(t, bedrockReq.ToolConfig)
10601059
require.NotNil(t, bedrockReq.ToolConfig.ToolChoice)
10611060
assert.NotNil(t, bedrockReq.ToolConfig.ToolChoice.Any)
10621061
}
10631062

1063+
func TestAnthropicToAWSBedrockTranslator_RequestBody_ToolResultMessagesWithSystemMessages(t *testing.T) {
1064+
// Regression test: when system messages are present in the messages array,
1065+
// promoteAnthropicSystemMessagesToParam filters them out, creating a shorter
1066+
// `messages` slice. The coalescing loop must use the filtered `messages` slice
1067+
// (not body.Messages) to avoid index misalignment and potential out-of-bounds access.
1068+
translator := NewAnthropicToAWSBedrockTranslator("")
1069+
req := &anthropicschema.MessagesRequest{
1070+
Model: "test-model",
1071+
MaxTokens: 100,
1072+
Tools: []anthropicschema.ToolUnion{
1073+
{Tool: &anthropicschema.Tool{
1074+
Type: "custom",
1075+
Name: "get_weather",
1076+
Description: "Get weather",
1077+
InputSchema: anthropicschema.ToolInputSchema{Type: "object"},
1078+
}},
1079+
},
1080+
Messages: []anthropicschema.MessageParam{
1081+
// System message at the start — will be promoted out of the messages slice.
1082+
{
1083+
Role: "system",
1084+
Content: anthropicschema.MessageContent{Text: "You are a helpful assistant."},
1085+
},
1086+
// Another system message — also promoted.
1087+
{
1088+
Role: "system",
1089+
Content: anthropicschema.MessageContent{Text: "Always be concise."},
1090+
},
1091+
{
1092+
Role: anthropicschema.MessageRoleUser,
1093+
Content: anthropicschema.MessageContent{Text: "What's the weather?"},
1094+
},
1095+
{
1096+
Role: anthropicschema.MessageRoleAssistant,
1097+
Content: anthropicschema.MessageContent{
1098+
Array: []anthropicschema.ContentBlockParam{
1099+
{ToolUse: &anthropicschema.ToolUseBlockParam{
1100+
Type: "tool_use",
1101+
ID: "tu_abc",
1102+
Name: "get_weather",
1103+
Input: map[string]any{"city": "NYC"},
1104+
}},
1105+
},
1106+
},
1107+
},
1108+
// First tool result message.
1109+
{
1110+
Role: anthropicschema.MessageRoleUser,
1111+
Content: anthropicschema.MessageContent{
1112+
Array: []anthropicschema.ContentBlockParam{
1113+
{ToolResult: &anthropicschema.ToolResultBlockParam{
1114+
Type: "tool_result",
1115+
ToolUseID: "tu_abc",
1116+
Content: &anthropicschema.ToolResultContent{Text: "72°F and sunny"},
1117+
}},
1118+
},
1119+
},
1120+
},
1121+
// Second consecutive tool result message (should be coalesced with the first).
1122+
{
1123+
Role: anthropicschema.MessageRoleUser,
1124+
Content: anthropicschema.MessageContent{
1125+
Array: []anthropicschema.ContentBlockParam{
1126+
{ToolResult: &anthropicschema.ToolResultBlockParam{
1127+
Type: "tool_result",
1128+
ToolUseID: "tu_def",
1129+
IsError: true,
1130+
Content: &anthropicschema.ToolResultContent{
1131+
Array: []anthropicschema.ToolResultContentItem{
1132+
{Text: &anthropicschema.TextBlockParam{Type: "text", Text: "Error: city not found"}},
1133+
},
1134+
},
1135+
}},
1136+
},
1137+
},
1138+
},
1139+
},
1140+
}
1141+
rawBody, err := json.Marshal(req)
1142+
require.NoError(t, err)
1143+
1144+
_, body, err := translator.RequestBody(rawBody, req, false)
1145+
require.NoError(t, err)
1146+
1147+
var bedrockReq awsbedrock.ConverseInput
1148+
err = json.Unmarshal(body, &bedrockReq)
1149+
require.NoError(t, err)
1150+
1151+
// After promoting 2 system messages, messages has 4 entries:
1152+
// user, assistant, user(tool_result), user(tool_result)
1153+
// The two consecutive tool_result messages should be coalesced into one,
1154+
// yielding 3 Bedrock messages: user, assistant, user(coalesced tool results).
1155+
require.Len(t, bedrockReq.Messages, 3)
1156+
1157+
// First message: user text.
1158+
assert.Equal(t, awsbedrock.ConversationRoleUser, bedrockReq.Messages[0].Role)
1159+
require.Len(t, bedrockReq.Messages[0].Content, 1)
1160+
assert.NotNil(t, bedrockReq.Messages[0].Content[0].Text)
1161+
1162+
// Second message: assistant tool_use.
1163+
assert.Equal(t, awsbedrock.ConversationRoleAssistant, bedrockReq.Messages[1].Role)
1164+
1165+
// Third message: coalesced tool results (both tool_use_ids in one message).
1166+
toolResultMsg := bedrockReq.Messages[2]
1167+
assert.Equal(t, awsbedrock.ConversationRoleUser, toolResultMsg.Role)
1168+
require.Len(t, toolResultMsg.Content, 2)
1169+
require.NotNil(t, toolResultMsg.Content[0].ToolResult)
1170+
assert.Equal(t, "tu_abc", *toolResultMsg.Content[0].ToolResult.ToolUseID)
1171+
assert.Equal(t, "72°F and sunny", *toolResultMsg.Content[0].ToolResult.Content[0].Text)
1172+
require.NotNil(t, toolResultMsg.Content[1].ToolResult)
1173+
assert.Equal(t, "tu_def", *toolResultMsg.Content[1].ToolResult.ToolUseID)
1174+
assert.Equal(t, "error", *toolResultMsg.Content[1].ToolResult.Status)
1175+
assert.Equal(t, "Error: city not found", *toolResultMsg.Content[1].ToolResult.Content[0].Text)
1176+
1177+
// Verify system prompt was promoted.
1178+
require.NotNil(t, bedrockReq.System)
1179+
require.Len(t, bedrockReq.System, 1)
1180+
assert.Contains(t, *bedrockReq.System[0].Text, "You are a helpful assistant.")
1181+
assert.Contains(t, *bedrockReq.System[0].Text, "Always be concise.")
1182+
}
1183+
10641184
func TestAnthropicToAWSBedrockTranslator_ResponseBody_NonStreamingToolUseAndReasoning(t *testing.T) {
10651185
translator := NewAnthropicToAWSBedrockTranslator("")
10661186
req := &anthropicschema.MessagesRequest{

0 commit comments

Comments
 (0)