fix: rerank integration

core/providers/bedrock/parallel_tool_result_ordering_test.go

@@ -18,8 +18,11 @@ package bedrock
 
 import (
 	"context"
+	"regexp"
+	"strings"
 	"testing"
 
+	"github.com/cespare/xxhash/v2"
 	"github.com/maximhq/bifrost/core/schemas"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -160,3 +163,178 @@ func TestParallelToolResultOrdering(t *testing.T) {
 			i, toolResultIDs, toolUseIDs)
 	}
 }
+
+// reportedGeminiToolUseID is the id from the OpenCode/Bedrock report: a Gemini-minted call id
+// with an embedded thought signature, replayed onto a glm-5 / kimi-2.5 turn.
+const reportedGeminiToolUseID = "d556q31u_ts_AY89a18eH3FucoBHPCdX6w7jgIhjSnIj7hU_mGofiw3SE2HL8kuMRcnfV3pGV7iUXRi_YQVYUgCH4bmcbYrS03yxWRtsPFb7KHPbp_iRinWZKPHtiyGVRlXfTsqeDBZOt3YNzKL-ycnZh0WeoLthSqnjkeZKT6idSNfd"
+
+var bedrockToolUseIDPattern = regexp.MustCompile(`^[a-zA-Z0-9_.:-]{1,64}$`)
+
+// TestBedrockAliasToolUseID verifies the alias only rewrites ids Bedrock would reject, so
+// every id that works today stays byte-identical on the wire.
+func TestBedrockAliasToolUseID(t *testing.T) {
+	tests := []struct {
+		name       string
+		id         string
+		wantSameID bool
+	}{
+		{"anthropic id", "tooluse_RwHN0v2n5kuNuZ2qoMV3SN", true},
+		{"openai id", "call_9v2n5kuNuZ2qoMV3SNRwHN0", true},
+		{"kimi id keeps dots and colons", "functions.read:0", true},
+		{"exactly 64 chars", strings.Repeat("a", 64), true},
+		{"65 chars", strings.Repeat("a", 65), false},
+		{"gemini thought signature", reportedGeminiToolUseID, false},
+		{"unsafe characters", "call/abc 123", false},
+		{"empty", "", false},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			got := bedrockAliasToolUseID(tc.id)
+			assert.Regexp(t, bedrockToolUseIDPattern, got, "alias violates Bedrock's toolUseId constraint")
+			if tc.wantSameID {
+				assert.Equal(t, tc.id, got, "id Bedrock already accepts must pass through unchanged")
+			} else {
+				assert.Equal(t, got, bedrockAliasToolUseID(tc.id), "alias must be deterministic")
+			}
+		})
+	}
+
+	// The hash covers the full id, so ids sharing a truncated head still alias apart.
+	head := strings.Repeat("z", 70)
+	assert.NotEqual(t, bedrockAliasToolUseID(head+"_one"), bedrockAliasToolUseID(head+"_two"),
+		"ids sharing a truncated head must not collide")
+}
+
+// TestBedrockToolUseIDPairingResponsesPath verifies an over-long id is aliased identically on
+// the tool_use and its tool_result, so Bedrock can still pair them.
+func TestBedrockToolUseIDPairingResponsesPath(t *testing.T) {
+	ptr := func(s string) *string { return &s }
+	msgType := func(t schemas.ResponsesMessageType) *schemas.ResponsesMessageType { return &t }
+
+	req := &schemas.BifrostResponsesRequest{
+		Model: "zai.glm-5",
+		Input: []schemas.ResponsesMessage{
+			{
+				Type: msgType(schemas.ResponsesMessageTypeMessage),
+				Role: func(r schemas.ResponsesMessageRoleType) *schemas.ResponsesMessageRoleType { return &r }(schemas.ResponsesInputMessageRoleUser),
+				Content: &schemas.ResponsesMessageContent{
+					ContentStr: ptr("What did we do so far?"),
+				},
+			},
+			{
+				Type: msgType(schemas.ResponsesMessageTypeFunctionCall),
+				ResponsesToolMessage: &schemas.ResponsesToolMessage{
+					CallID:    ptr(reportedGeminiToolUseID),
+					Name:      ptr("bash"),
+					Arguments: ptr(`{"command":"ls"}`),
+				},
+			},
+			{
+				Type: msgType(schemas.ResponsesMessageTypeFunctionCallOutput),
+				ResponsesToolMessage: &schemas.ResponsesToolMessage{
+					CallID: ptr(reportedGeminiToolUseID),
+					Output: &schemas.ResponsesToolMessageOutputStruct{ResponsesToolCallOutputStr: ptr("build/ src/")},
+				},
+			},
+		},
+	}
+
+	bedrockReq, err := ToBedrockResponsesRequest(schemas.NewBifrostContext(context.Background(), schemas.NoDeadline), req)
+	require.NoError(t, err)
+
+	var toolUseIDs, toolResultIDs []string
+	for _, msg := range bedrockReq.Messages {
+		for _, block := range msg.Content {
+			if block.ToolUse != nil {
+				toolUseIDs = append(toolUseIDs, block.ToolUse.ToolUseID)
+			}
+			if block.ToolResult != nil {
+				toolResultIDs = append(toolResultIDs, block.ToolResult.ToolUseID)
+			}
+		}
+	}
+
+	require.Len(t, toolUseIDs, 1)
+	require.Len(t, toolResultIDs, 1)
+	assert.Regexp(t, bedrockToolUseIDPattern, toolUseIDs[0], "Bedrock rejects toolUse.toolUseId over 64 chars")
+	assert.Regexp(t, bedrockToolUseIDPattern, toolResultIDs[0], "Bedrock rejects toolResult.toolUseId over 64 chars")
+	assert.Equal(t, toolUseIDs[0], toolResultIDs[0], "tool_use and tool_result must alias to the same id")
+}
+
+// TestBedrockToolUseIDPairingChatPath is TestBedrockToolUseIDPairingResponsesPath for the
+// chat completions path, which builds toolUse/toolResult blocks through a separate converter.
+func TestBedrockToolUseIDPairingChatPath(t *testing.T) {
+	ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
+	req := &schemas.BifrostChatRequest{
+		Model: "zai.glm-5",
+		Input: []schemas.ChatMessage{
+			{
+				Role:    schemas.ChatMessageRoleUser,
+				Content: &schemas.ChatMessageContent{ContentStr: schemas.Ptr("What did we do so far?")},
+			},
+			{
+				Role: schemas.ChatMessageRoleAssistant,
+				ChatAssistantMessage: &schemas.ChatAssistantMessage{
+					ToolCalls: []schemas.ChatAssistantMessageToolCall{
+						{
+							ID:   schemas.Ptr(reportedGeminiToolUseID),
+							Type: schemas.Ptr(string(schemas.ChatToolChoiceTypeFunction)),
+							Function: schemas.ChatAssistantMessageToolCallFunction{
+								Name:      schemas.Ptr("bash"),
+								Arguments: `{"command":"ls"}`,
+							},
+						},
+					},
+				},
+			},
+			{
+				Role:            schemas.ChatMessageRoleTool,
+				ChatToolMessage: &schemas.ChatToolMessage{ToolCallID: schemas.Ptr(reportedGeminiToolUseID)},
+				Content:         &schemas.ChatMessageContent{ContentStr: schemas.Ptr("build/ src/")},
+			},
+		},
+	}
+
+	bedrockReq, err := ToBedrockChatCompletionRequest(ctx, req)
+	require.NoError(t, err)
+
+	var toolUseID, toolResultID string
+	for _, msg := range bedrockReq.Messages {
+		for _, block := range msg.Content {
+			if block.ToolUse != nil {
+				toolUseID = block.ToolUse.ToolUseID
+			}
+			if block.ToolResult != nil {
+				toolResultID = block.ToolResult.ToolUseID
+			}
+		}
+	}
+
+	assert.Regexp(t, bedrockToolUseIDPattern, toolUseID, "Bedrock rejects toolUse.toolUseId over 64 chars")
+	assert.Regexp(t, bedrockToolUseIDPattern, toolResultID, "Bedrock rejects toolResult.toolUseId over 64 chars")
+	assert.Equal(t, toolUseID, toolResultID, "tool_use and tool_result must alias to the same id")
+}
+
+// TestBedrockAliasToolUseIDFullHashAvoidsCollision uses a real pair of 68-char ids found by
+// brute-force search: their uint32(xxhash.Sum64String(...)) values collide (the bug in the
+// previous 32-bit-truncated hash), but the full 64-bit hashes don't. If bedrockAliasToolUseID
+// ever regresses to hashing only a uint32-truncated slice, this reproduces a real duplicate
+// toolUseId — which Bedrock rejects outright, or worse, silently misattributes a tool_result.
+func TestBedrockAliasToolUseIDFullHashAvoidsCollision(t *testing.T) {
+	idA := "d556q31u_ts_AY89a18eH3FucoBHPCdX6w7jgIhjSnIj7hU_mGofiw3SE2HL8_n38390"
+	idB := "d556q31u_ts_AY89a18eH3FucoBHPCdX6w7jgIhjSnIj7hU_mGofiw3SE2HL8_n63004"
+	require.Greater(t, len(idA), 64)
+	require.Greater(t, len(idB), 64)
+
+	// Document the collision this pair exercises: same 32-bit-truncated hash, different ids.
+	require.Equal(t, uint32(xxhash.Sum64String(idA)), uint32(xxhash.Sum64String(idB)),
+		"fixture no longer demonstrates a 32-bit hash collision")
+	require.NotEqual(t, idA, idB)
+
+	aliasA := bedrockAliasToolUseID(idA)
+	aliasB := bedrockAliasToolUseID(idB)
+	assert.NotEqual(t, aliasA, aliasB, "ids with a colliding 32-bit hash must still alias apart")
+	assert.Regexp(t, bedrockToolUseIDPattern, aliasA)
+	assert.Regexp(t, bedrockToolUseIDPattern, aliasB)
+}

core/providers/bedrock/rerank.go

@@ -151,6 +151,36 @@ func (response *BedrockRerankResponse) ToBifrostRerankResponse(documents []schem
 	return bifrostResponse
 }
 
+// ToBedrockRerankResponse converts a Bifrost rerank response into Bedrock Agent Runtime format.
+// Bedrock echoes the ranked document, which is only present when the rerank was requested
+// with ReturnDocuments enabled.
+func ToBedrockRerankResponse(bifrostResp *schemas.BifrostRerankResponse) *BedrockRerankResponse {
+	if bifrostResp == nil {
+		return nil
+	}
+
+	bedrockResp := &BedrockRerankResponse{
+		Results: make([]BedrockRerankResult, 0, len(bifrostResp.Results)),
+	}
+
+	for _, result := range bifrostResp.Results {
+		bedrockResult := BedrockRerankResult{
+			Index:          result.Index,
+			RelevanceScore: result.RelevanceScore,
+		}
+		if result.Document != nil {
+			bedrockResult.Document = &BedrockRerankResponseDocument{
+				Type:         bedrockRerankInlineDocumentTypeText,
+				TextDocument: &BedrockRerankTextValue{Text: result.Document.Text},
+			}
+		}
+
+		bedrockResp.Results = append(bedrockResp.Results, bedrockResult)
+	}
+
+	return bedrockResp
+}
+
 // ToBifrostRerankRequest converts a Bedrock Agent Runtime rerank request to Bifrost format.
 func (req *BedrockRerankRequest) ToBifrostRerankRequest(ctx *schemas.BifrostContext) *schemas.BifrostRerankRequest {
 	if req == nil {

core/providers/bedrock/rerank_test.go

@@ -333,3 +333,36 @@ func TestBedrockRerankRejectsEmptyModelIdentifier(t *testing.T) {
 	require.NotNil(t, bifrostErr.Error)
 	assert.Contains(t, bifrostErr.Error.Message, "model identifier")
 }
+
+func TestToBedrockRerankResponse(t *testing.T) {
+	response := ToBedrockRerankResponse(&schemas.BifrostRerankResponse{
+		Results: []schemas.RerankResult{
+			{Index: 1, RelevanceScore: 0.92, Document: &schemas.RerankDocument{Text: "Paris is capital of France"}},
+			{Index: 0, RelevanceScore: 0.11, Document: &schemas.RerankDocument{Text: "Berlin is capital of Germany"}},
+		},
+		Model: "arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0",
+	})
+
+	require.NotNil(t, response)
+	require.Len(t, response.Results, 2)
+	assert.Equal(t, 1, response.Results[0].Index)
+	assert.InDelta(t, 0.92, response.Results[0].RelevanceScore, 1e-9)
+	require.NotNil(t, response.Results[0].Document)
+	assert.Equal(t, bedrockRerankInlineDocumentTypeText, response.Results[0].Document.Type)
+	require.NotNil(t, response.Results[0].Document.TextDocument)
+	assert.Equal(t, "Paris is capital of France", response.Results[0].Document.TextDocument.Text)
+	require.NotNil(t, response.Results[1].Document)
+	assert.Equal(t, "Berlin is capital of Germany", response.Results[1].Document.TextDocument.Text)
+}
+
+func TestToBedrockRerankResponseOmitsMissingDocument(t *testing.T) {
+	response := ToBedrockRerankResponse(&schemas.BifrostRerankResponse{
+		Results: []schemas.RerankResult{{Index: 0, RelevanceScore: 0.7}},
+	})
+
+	require.NotNil(t, response)
+	require.Len(t, response.Results, 1)
+	assert.Nil(t, response.Results[0].Document)
+
+	assert.Nil(t, ToBedrockRerankResponse(nil))
+}

core/providers/bedrock/responses.go

@@ -3345,7 +3345,7 @@ func ConvertBifrostMessagesToBedrockMessages(ctx context.Context, bifrostMessage
 				result := pendingResults[callID]
 				resultBlocks = append(resultBlocks, BedrockContentBlock{
 					ToolResult: &BedrockToolResult{
-						ToolUseID: callID,
+						ToolUseID: bedrockAliasToolUseID(callID),
 						Content:   result.Content,
 						Status:    schemas.Ptr(result.Status),
 					},
@@ -3385,7 +3385,7 @@ func ConvertBifrostMessagesToBedrockMessages(ctx context.Context, bifrostMessage
 				if toolCall, exists := stateManager.toolCalls[callID]; exists {
 					toolUseBlock := &BedrockContentBlock{
 						ToolUse: &BedrockToolUse{
-							ToolUseID: toolCall.CallID,
+							ToolUseID: bedrockAliasToolUseID(toolCall.CallID),
 							Name:      toolCall.ToolName,
 						},
 					}
@@ -3539,7 +3539,7 @@ func ConvertBifrostMessagesToBedrockMessages(ctx context.Context, bifrostMessage
 						if toolCall, exists := stateManager.toolCalls[callID]; exists {
 							toolUseBlock := &BedrockContentBlock{
 								ToolUse: &BedrockToolUse{
-									ToolUseID: toolCall.CallID,
+									ToolUseID: bedrockAliasToolUseID(toolCall.CallID),
 									Name:      toolCall.ToolName,
 								},
 							}
@@ -3579,7 +3579,7 @@ func ConvertBifrostMessagesToBedrockMessages(ctx context.Context, bifrostMessage
 						result := pendingResults[callID]
 						resultBlocks = append(resultBlocks, BedrockContentBlock{
 							ToolResult: &BedrockToolResult{
-								ToolUseID: callID,
+								ToolUseID: bedrockAliasToolUseID(callID),
 								Content:   result.Content,
 								Status:    schemas.Ptr(result.Status),
 							},
@@ -3620,7 +3620,7 @@ func ConvertBifrostMessagesToBedrockMessages(ctx context.Context, bifrostMessage
 					if toolCall, exists := stateManager.toolCalls[callID]; exists {
 						toolUseBlock := &BedrockContentBlock{
 							ToolUse: &BedrockToolUse{
-								ToolUseID: toolCall.CallID,
+								ToolUseID: bedrockAliasToolUseID(toolCall.CallID),
 								Name:      toolCall.ToolName,
 							},
 						}
@@ -3662,7 +3662,7 @@ func ConvertBifrostMessagesToBedrockMessages(ctx context.Context, bifrostMessage
 					result := pendingResults[callID]
 					resultBlocks = append(resultBlocks, BedrockContentBlock{
 						ToolResult: &BedrockToolResult{
-							ToolUseID: callID,
+							ToolUseID: bedrockAliasToolUseID(callID),
 							Content:   result.Content,
 							Status:    schemas.Ptr(result.Status),
 						},
@@ -3739,7 +3739,7 @@ func ConvertBifrostMessagesToBedrockMessages(ctx context.Context, bifrostMessage
 			inputBytes, _ := json.Marshal(inputMap)
 			toolUseBlock := BedrockContentBlock{
 				ToolUse: &BedrockToolUse{
-					ToolUseID: callID,
+					ToolUseID: bedrockAliasToolUseID(callID),
 					Name:      string(BedrockSystemToolNovaGrounding),
 					Input:     json.RawMessage(inputBytes),
 					Type:      "server_tool_use",
@@ -3759,7 +3759,7 @@ func ConvertBifrostMessagesToBedrockMessages(ctx context.Context, bifrostMessage
 			resultType := BedrockNovaGroundingResultType
 			toolResultBlock := BedrockContentBlock{
 				ToolResult: &BedrockToolResult{
-					ToolUseID: callID,
+					ToolUseID: bedrockAliasToolUseID(callID),
 					Type:      &resultType,
 					Status:    schemas.Ptr("success"),
 					Content:   []BedrockContentBlock{{Text: &sourcesText}},
@@ -3785,7 +3785,7 @@ func ConvertBifrostMessagesToBedrockMessages(ctx context.Context, bifrostMessage
 			inputBytes, _ := json.Marshal(map[string]string{"snippet": code})
 			toolUseBlock := BedrockContentBlock{
 				ToolUse: &BedrockToolUse{
-					ToolUseID: toolUseID,
+					ToolUseID: bedrockAliasToolUseID(toolUseID),
 					Name:      string(BedrockSystemToolNovaCodeInterpreter),
 					Input:     json.RawMessage(inputBytes),
 					Type:      "server_tool_use",
@@ -3806,7 +3806,7 @@ func ConvertBifrostMessagesToBedrockMessages(ctx context.Context, bifrostMessage
 			resultType := BedrockNovaCodeInterpreterResultType
 			toolResultBlock := BedrockContentBlock{
 				ToolResult: &BedrockToolResult{
-					ToolUseID: toolUseID,
+					ToolUseID: bedrockAliasToolUseID(toolUseID),
 					Type:      &resultType,
 					Content:   []BedrockContentBlock{{Text: &execResultStr}},
 				},