fix(ai-proxy): preserve Bedrock Claude reasoning blocks (#3788)

Signed-off-by: Betula-L <6059935+Betula-L@users.noreply.github.com> Co-authored-by: Betula-L <6059935+Betula-L@users.noreply.github.com>
2026-05-28 22:57:31 +08:00 · 2026-05-07 19:27:48 -07:00
parent 6199fe414d
commit b77a074831
8 changed files with 1420 additions and 142 deletions
--- a/plugins/wasm-go/extensions/ai-proxy/provider/claude_test.go
+++ b/plugins/wasm-go/extensions/ai-proxy/provider/claude_test.go
@@ -184,6 +184,48 @@ func TestClaudeProvider_BuildClaudeTextGenRequest_StandardMode(t *testing.T) {
 		assert.Equal(t, "You are a helpful assistant.", claudeReq.System.StringValue)
 	})

+	t.Run("preserves_bridge_thinking_blocks_and_output_config", func(t *testing.T) {
+		request := &chatCompletionRequest{
+			Model:          "claude-sonnet-4-5-20250929",
+			MaxTokens:      8192,
+			ClaudeThinking: &claudeThinkingConfig{Type: "adaptive", Display: "omitted"},
+			ClaudeOutputConfig: &claudeOutputConfig{
+				Effort: "high",
+				Format: json.RawMessage(`{
+					"type":"json_schema",
+					"schema":{"type":"object","properties":{"answer":{"type":"string"}}}
+				}`),
+			},
+			Messages: []chatMessage{{
+				Role: roleAssistant,
+				ClaudeContentBlocks: []claudeChatMessageContent{
+					{Type: "thinking", Thinking: "", Signature: "sig"},
+					{Type: "redacted_thinking", Data: "opaque-base64"},
+					{Type: "text", Text: "answer"},
+				},
+			}},
+		}
+
+		claudeReq := provider.buildClaudeTextGenRequest(request)
+
+		require.NotNil(t, claudeReq.Thinking)
+		assert.Equal(t, "adaptive", claudeReq.Thinking.Type)
+		assert.Equal(t, "omitted", claudeReq.Thinking.Display)
+		require.NotNil(t, claudeReq.OutputConfig)
+		assert.Equal(t, "high", claudeReq.OutputConfig.Effort)
+		require.NotEmpty(t, claudeReq.OutputConfig.Format)
+		assert.Contains(t, string(claudeReq.OutputConfig.Format), "json_schema")
+		require.Len(t, claudeReq.Messages, 1)
+		blocks := claudeReq.Messages[0].Content.GetArrayValue()
+		require.Len(t, blocks, 3)
+		assert.Equal(t, "thinking", blocks[0].Type)
+		assert.Equal(t, "sig", blocks[0].Signature)
+		assert.Equal(t, "redacted_thinking", blocks[1].Type)
+		assert.Equal(t, "opaque-base64", blocks[1].Data)
+		assert.Equal(t, "text", blocks[2].Type)
+		assert.Equal(t, "answer", blocks[2].Text)
+	})
+
 	t.Run("maps_openai_function_tool_choice_to_claude_tool_choice", func(t *testing.T) {
 		request := &chatCompletionRequest{
 			Model:     "claude-sonnet-4-5-20250929",
@@ -241,6 +283,30 @@ func TestClaudeProvider_BuildClaudeTextGenRequest_StandardMode(t *testing.T) {
 		assert.True(t, claudeReq.ToolChoice.DisableParallelToolUse)
 	})

+	t.Run("downgrades_forced_tool_choice_to_auto_when_thinking_enabled", func(t *testing.T) {
+		request := &chatCompletionRequest{
+			Model:          "claude-sonnet-4-5-20250929",
+			MaxTokens:      8192,
+			ClaudeThinking: &claudeThinkingConfig{Type: "enabled", BudgetTokens: 8192},
+			Messages: []chatMessage{
+				{Role: roleUser, Content: "Search."},
+			},
+			Tools: []tool{{
+				Type: "function",
+				Function: function{
+					Name:       "web_search",
+					Parameters: map[string]interface{}{"type": "object"},
+				},
+			}},
+			ToolChoice: "required",
+		}
+
+		claudeReq := provider.buildClaudeTextGenRequest(request)
+
+		require.NotNil(t, claudeReq.ToolChoice)
+		assert.Equal(t, "auto", claudeReq.ToolChoice.Type)
+	})
+
 	t.Run("maps_openai_string_none_tool_choice_to_claude_none", func(t *testing.T) {
 		request := &chatCompletionRequest{
 			Model:     "claude-sonnet-4-5-20250929",
@@ -264,6 +330,117 @@ func TestClaudeProvider_BuildClaudeTextGenRequest_StandardMode(t *testing.T) {
 		assert.Equal(t, "none", claudeReq.ToolChoice.Type)
 		assert.Empty(t, claudeReq.ToolChoice.Name)
 	})
+
+	t.Run("preserves_bridge_tool_result_blocks_before_role_tool_fallback", func(t *testing.T) {
+		request := &chatCompletionRequest{
+			Model:     "claude-sonnet-4-5-20250929",
+			MaxTokens: 8192,
+			Messages: []chatMessage{{
+				Role: roleTool,
+				ClaudeContentBlocks: []claudeChatMessageContent{{
+					Type:      "tool_result",
+					ToolUseId: "toolu_1",
+					IsError:   true,
+					Content: &claudeChatMessageContentWr{
+						ArrayValue: []claudeChatMessageContent{{
+							Type: "image",
+							Source: &claudeChatMessageContentSource{
+								Type:      "base64",
+								MediaType: "image/png",
+								Data:      "AAAA",
+							},
+						}},
+					},
+				}},
+			}},
+		}
+
+		claudeReq := provider.buildClaudeTextGenRequest(request)
+
+		require.Len(t, claudeReq.Messages, 1)
+		assert.Equal(t, roleUser, claudeReq.Messages[0].Role)
+		blocks := claudeReq.Messages[0].Content.GetArrayValue()
+		require.Len(t, blocks, 1)
+		assert.Equal(t, "tool_result", blocks[0].Type)
+		assert.True(t, blocks[0].IsError)
+		require.NotNil(t, blocks[0].Content)
+		require.Len(t, blocks[0].Content.ArrayValue, 1)
+		assert.Equal(t, "image", blocks[0].Content.ArrayValue[0].Type)
+	})
+}
+
+func TestClaudeProvider_ResponsePreservesNativeThinkingBlocksForClaudeConversion(t *testing.T) {
+	provider := &claudeProvider{}
+	ctx := newMockMultipartHttpContext()
+	ctx.SetContext("needClaudeResponseConversion", true)
+	thinking := "reasoning"
+	signature := "sig"
+	text := "answer"
+
+	response := provider.responseClaude2OpenAI(ctx, &claudeTextGenResponse{
+		Id:    "msg_1",
+		Model: "claude-sonnet-4-5-20250929",
+		Type:  "message",
+		Role:  roleAssistant,
+		Content: []claudeTextGenContent{
+			{Type: "thinking", Thinking: &thinking, Signature: &signature},
+			{Type: "redacted_thinking", Data: "opaque-base64"},
+			{Type: "text", Text: &text},
+		},
+	})
+	body, err := json.Marshal(response)
+	require.NoError(t, err)
+
+	converted, err := (&ClaudeToOpenAIConverter{}).ConvertOpenAIResponseToClaude(ctx, body)
+	require.NoError(t, err)
+
+	var claudeResponse claudeTextGenResponse
+	require.NoError(t, json.Unmarshal(converted, &claudeResponse))
+	require.Len(t, claudeResponse.Content, 3)
+	assert.Equal(t, "thinking", claudeResponse.Content[0].Type)
+	require.NotNil(t, claudeResponse.Content[0].Signature)
+	assert.Equal(t, "sig", *claudeResponse.Content[0].Signature)
+	assert.Equal(t, "redacted_thinking", claudeResponse.Content[1].Type)
+	assert.Equal(t, "opaque-base64", claudeResponse.Content[1].Data)
+	assert.Equal(t, "text", claudeResponse.Content[2].Type)
+}
+
+func TestClaudeProvider_StreamPreservesNativeSignatureAndStopsForClaudeConversion(t *testing.T) {
+	provider := &claudeProvider{}
+	ctx := newMockMultipartHttpContext()
+	ctx.SetContext("needClaudeResponseConversion", true)
+	converter := &ClaudeToOpenAIConverter{}
+	index := 1
+
+	signatureResponse := provider.streamResponseClaude2OpenAI(ctx, &claudeTextGenStreamResponse{
+		Type:  "content_block_delta",
+		Index: &index,
+		Delta: &claudeTextGenDelta{
+			Type:      "signature_delta",
+			Signature: "sig",
+		},
+	})
+	signatureBody, err := json.Marshal(signatureResponse)
+	require.NoError(t, err)
+	converted, err := converter.ConvertOpenAIStreamResponseToClaude(ctx, []byte("data: "+string(signatureBody)+"\n\n"))
+	require.NoError(t, err)
+	events := parseClaudeSSEEvents(t, converted)
+	require.Len(t, events, 2)
+	assert.Equal(t, "content_block_start", events[0].Name)
+	assert.Equal(t, "content_block_delta", events[1].Name)
+	assert.Equal(t, "signature_delta", events[1].Payload.Delta.Type)
+
+	stopResponse := provider.streamResponseClaude2OpenAI(ctx, &claudeTextGenStreamResponse{
+		Type:  "content_block_stop",
+		Index: &index,
+	})
+	stopBody, err := json.Marshal(stopResponse)
+	require.NoError(t, err)
+	converted, err = converter.ConvertOpenAIStreamResponseToClaude(ctx, []byte("data: "+string(stopBody)+"\n\n"))
+	require.NoError(t, err)
+	events = parseClaudeSSEEvents(t, converted)
+	require.Len(t, events, 1)
+	assert.Equal(t, "content_block_stop", events[0].Name)
 }

 func TestClaudeProvider_BuildClaudeTextGenRequest_ClaudeCodeMode(t *testing.T) {