fix(ai-proxy): preserve Vertex thoughtSignature in OpenAI tool calls (#3973)

Signed-off-by: DENG <33118163+XinhhD@users.noreply.github.com>

plugins/wasm-go/extensions/ai-proxy/provider/model.go

@@ -492,10 +492,49 @@ func (m *chatMessage) ParseContent() []chatMessageContent {
 }
 
 type toolCall struct {
-	Index    int          `json:"index"`
-	Id       string       `json:"id,omitempty"`
-	Type     string       `json:"type"`
-	Function functionCall `json:"function"`
+	Index            int            `json:"index"`
+	Id               string         `json:"id,omitempty"`
+	Type             string         `json:"type"`
+	Function         functionCall   `json:"function"`
+	ThoughtSignature string         `json:"thought_signature,omitempty"`
+	ExtraContent     map[string]any `json:"extra_content,omitempty"`
+}
+
+func (t *toolCall) getThoughtSignature() string {
+	if t == nil {
+		return ""
+	}
+	if t.ThoughtSignature != "" {
+		return t.ThoughtSignature
+	}
+	return getNestedString(t.ExtraContent, "google", "thought_signature")
+}
+
+func buildGoogleThoughtSignatureExtraContent(signature string) map[string]any {
+	if signature == "" {
+		return nil
+	}
+	return map[string]any{
+		"google": map[string]any{
+			"thought_signature": signature,
+		},
+	}
+}
+
+func getNestedString(data map[string]any, path ...string) string {
+	var current any = data
+	for _, key := range path {
+		currentMap, ok := current.(map[string]any)
+		if !ok {
+			return ""
+		}
+		current, ok = currentMap[key]
+		if !ok {
+			return ""
+		}
+	}
+	value, _ := current.(string)
+	return value
 }
 
 type functionCall struct {

plugins/wasm-go/extensions/ai-proxy/provider/vertex.go

@@ -871,7 +871,9 @@ func (v *vertexProvider) buildChatCompletionResponse(ctx wrapper.HttpContext, re
 				args, _ := json.Marshal(part.FunctionCall.Args)
 				choice.Message.ToolCalls = []toolCall{
 					{
-						Type: "function",
+						Type:             "function",
+						ThoughtSignature: part.ThoughtSignature,
+						ExtraContent:     buildGoogleThoughtSignatureExtraContent(part.ThoughtSignature),
 						Function: functionCall{
 							Name:      part.FunctionCall.Name,
 							Arguments: string(args),
@@ -979,7 +981,9 @@ func (v *vertexProvider) buildChatCompletionStreamResponse(ctx wrapper.HttpConte
 			choice.Delta = &chatMessage{
 				ToolCalls: []toolCall{
 					{
-						Type: "function",
+						Type:             "function",
+						ThoughtSignature: part.ThoughtSignature,
+						ExtraContent:     buildGoogleThoughtSignatureExtraContent(part.ThoughtSignature),
 						Function: functionCall{
 							Name:      part.FunctionCall.Name,
 							Arguments: string(args),
@@ -1143,6 +1147,7 @@ func (v *vertexProvider) buildVertexChatRequest(request *chatCompletionRequest)
 					Name: lastFunctionName,
 					Args: args,
 				},
+				ThoughtSignature: message.ToolCalls[0].getThoughtSignature(),
 			})
 		} else {
 			for _, part := range message.ParseContent() {
@@ -1359,6 +1364,7 @@ type vertexPart struct {
 	FunctionCall     *vertexFunctionCall     `json:"functionCall,omitempty"`
 	FunctionResponse *vertexFunctionResponse `json:"functionResponse,omitempty"`
 	Thounght         *bool                   `json:"thought,omitempty"`
+	ThoughtSignature string                  `json:"thoughtSignature,omitempty"`
 }
 
 type blob struct {

plugins/wasm-go/extensions/ai-proxy/provider/vertex_test.go

@@ -578,3 +578,365 @@ func TestVertexAnthropicPassthrough_MaxTokensDefault(t *testing.T) {
 			"client-supplied max_tokens must not be overwritten by the default")
 	})
 }
+
+func TestVertexProviderPreservesFunctionCallThoughtSignature(t *testing.T) {
+	v := &vertexProvider{}
+	ctx := newMockMultipartHttpContext()
+	ctx.SetContext(ctxKeyFinalRequestModel, "gemini-3.1-pro-preview")
+
+	response := v.buildChatCompletionResponse(ctx, &vertexChatResponse{
+		ResponseId: "vertex-response-id",
+		Candidates: []vertexChatCandidate{
+			{
+				Index: 0,
+				Content: vertexChatContent{
+					Role: "model",
+					Parts: []vertexPart{
+						{
+							FunctionCall: &vertexFunctionCall{
+								Name: "Skill",
+								Args: map[string]interface{}{"query": "intelligentization"},
+							},
+							ThoughtSignature: "thought-signature-from-vertex",
+						},
+					},
+				},
+				FinishReason: "STOP",
+			},
+		},
+	})
+
+	require.Len(t, response.Choices, 1)
+	require.NotNil(t, response.Choices[0].Message)
+	require.Len(t, response.Choices[0].Message.ToolCalls, 1)
+	assert.Equal(t, "thought-signature-from-vertex", response.Choices[0].Message.ToolCalls[0].ThoughtSignature)
+	assert.Equal(
+		t,
+		"thought-signature-from-vertex",
+		getNestedString(response.Choices[0].Message.ToolCalls[0].ExtraContent, "google", "thought_signature"),
+	)
+}
+
+func TestVertexProviderRestoresFunctionCallThoughtSignature(t *testing.T) {
+	v := &vertexProvider{}
+	req := &chatCompletionRequest{
+		Model: "gemini-3.1-pro-preview",
+		Messages: []chatMessage{
+			{Role: roleUser, Content: "search docs"},
+			{
+				Role: roleAssistant,
+				ToolCalls: []toolCall{
+					{
+						Type:             "function",
+						ThoughtSignature: "thought-signature-from-client",
+						Function: functionCall{
+							Name:      "Skill",
+							Arguments: `{"query":"intelligentization"}`,
+						},
+					},
+				},
+			},
+			{Role: roleTool, Content: "tool result"},
+		},
+	}
+
+	vertexReq, err := v.buildVertexChatRequest(req)
+	require.NoError(t, err)
+	require.NotNil(t, vertexReq)
+	require.Len(t, vertexReq.Contents, 3)
+	require.Len(t, vertexReq.Contents[1].Parts, 1)
+	require.NotNil(t, vertexReq.Contents[1].Parts[0].FunctionCall)
+	assert.Equal(t, "thought-signature-from-client", vertexReq.Contents[1].Parts[0].ThoughtSignature)
+}
+
+func TestVertexProviderRestoresFunctionCallThoughtSignatureFromGoogleExtraContent(t *testing.T) {
+	v := &vertexProvider{}
+	req := &chatCompletionRequest{
+		Model: "gemini-3.1-pro-preview",
+		Messages: []chatMessage{
+			{Role: roleUser, Content: "search docs"},
+			{
+				Role: roleAssistant,
+				ToolCalls: []toolCall{
+					{
+						Type: "function",
+						ExtraContent: map[string]any{
+							"google": map[string]any{
+								"thought_signature": "thought-signature-from-extra-content",
+							},
+						},
+						Function: functionCall{
+							Name:      "Skill",
+							Arguments: `{"query":"intelligentization"}`,
+						},
+					},
+				},
+			},
+			{Role: roleTool, Content: "tool result"},
+		},
+	}
+
+	vertexReq, err := v.buildVertexChatRequest(req)
+	require.NoError(t, err)
+	require.NotNil(t, vertexReq)
+	require.Len(t, vertexReq.Contents, 3)
+	require.Len(t, vertexReq.Contents[1].Parts, 1)
+	require.NotNil(t, vertexReq.Contents[1].Parts[0].FunctionCall)
+	assert.Equal(t, "thought-signature-from-extra-content", vertexReq.Contents[1].Parts[0].ThoughtSignature)
+}
+
+func TestVertexProviderRestoresFunctionCallThoughtSignatureInvalidArguments(t *testing.T) {
+	v := &vertexProvider{}
+	req := &chatCompletionRequest{
+		Model: "gemini-3.1-pro-preview",
+		Messages: []chatMessage{
+			{
+				Role: roleAssistant,
+				ToolCalls: []toolCall{
+					{
+						Type:             "function",
+						ThoughtSignature: "thought-signature-from-client",
+						Function: functionCall{
+							Name:      "Skill",
+							Arguments: `invalid-json`,
+						},
+					},
+				},
+			},
+		},
+	}
+
+	vertexReq, err := v.buildVertexChatRequest(req)
+	require.NoError(t, err)
+	require.NotNil(t, vertexReq)
+	require.Len(t, vertexReq.Contents, 1)
+	require.Len(t, vertexReq.Contents[0].Parts, 1)
+	require.NotNil(t, vertexReq.Contents[0].Parts[0].FunctionCall)
+	assert.Equal(t, "thought-signature-from-client", vertexReq.Contents[0].Parts[0].ThoughtSignature)
+}
+
+func TestVertexProviderStreamPreservesFunctionCallThoughtSignature(t *testing.T) {
+	v := &vertexProvider{}
+	ctx := newMockMultipartHttpContext()
+	ctx.SetContext(ctxKeyFinalRequestModel, "gemini-3.1-pro-preview")
+
+	response := v.buildChatCompletionStreamResponse(ctx, &vertexChatResponse{
+		ResponseId: "vertex-response-id",
+		Candidates: []vertexChatCandidate{
+			{
+				Index: 0,
+				Content: vertexChatContent{
+					Role: "model",
+					Parts: []vertexPart{
+						{
+							FunctionCall: &vertexFunctionCall{
+								Name: "Skill",
+								Args: map[string]interface{}{"query": "intelligentization"},
+							},
+							ThoughtSignature: "thought-signature-from-vertex-stream",
+						},
+					},
+				},
+				FinishReason: "STOP",
+			},
+		},
+	})
+
+	require.NotNil(t, response)
+	require.NotNil(t, response.Choices)
+	require.Len(t, response.Choices, 1)
+	require.NotNil(t, response.Choices[0].Delta)
+	require.Len(t, response.Choices[0].Delta.ToolCalls, 1)
+	assert.Equal(t, "thought-signature-from-vertex-stream", response.Choices[0].Delta.ToolCalls[0].ThoughtSignature)
+	assert.Equal(
+		t,
+		"thought-signature-from-vertex-stream",
+		getNestedString(response.Choices[0].Delta.ToolCalls[0].ExtraContent, "google", "thought_signature"),
+	)
+}
+
+func TestToolCallGetThoughtSignatureAllPaths(t *testing.T) {
+	// 1. nil toolCall pointer
+	var nilTc *toolCall
+	assert.Equal(t, "", nilTc.getThoughtSignature())
+
+	// 2. ThoughtSignature is directly set
+	tc1 := &toolCall{
+		ThoughtSignature: "direct-sig",
+	}
+	assert.Equal(t, "direct-sig", tc1.getThoughtSignature())
+
+	// 3. ExtraContent is nil
+	tc2 := &toolCall{}
+	assert.Equal(t, "", tc2.getThoughtSignature())
+
+	// 4. ExtraContent does not contain "google"
+	tc3 := &toolCall{
+		ExtraContent: map[string]any{
+			"other": "val",
+		},
+	}
+	assert.Equal(t, "", tc3.getThoughtSignature())
+
+	// 5. ExtraContent contains "google" but not as map[string]any
+	tc4 := &toolCall{
+		ExtraContent: map[string]any{
+			"google": "not-a-map",
+		},
+	}
+	assert.Equal(t, "", tc4.getThoughtSignature())
+
+	// 6. ExtraContent contains "google" map but not "thought_signature"
+	tc5 := &toolCall{
+		ExtraContent: map[string]any{
+			"google": map[string]any{
+				"other": "val",
+			},
+		},
+	}
+	assert.Equal(t, "", tc5.getThoughtSignature())
+
+	// 7. ExtraContent contains "google" map with "thought_signature" but not string
+	tc6 := &toolCall{
+		ExtraContent: map[string]any{
+			"google": map[string]any{
+				"thought_signature": 12345,
+			},
+		},
+	}
+	assert.Equal(t, "", tc6.getThoughtSignature())
+
+	// 8. ExtraContent contains "google" map with "thought_signature" as string
+	tc7 := &toolCall{
+		ExtraContent: map[string]any{
+			"google": map[string]any{
+				"thought_signature": "google-extra-sig",
+			},
+		},
+	}
+	assert.Equal(t, "google-extra-sig", tc7.getThoughtSignature())
+}
+
+func TestBuildGoogleThoughtSignatureExtraContentEmpty(t *testing.T) {
+	assert.Nil(t, buildGoogleThoughtSignatureExtraContent(""))
+}
+
+func TestVertexProviderStreamThoughtAndText(t *testing.T) {
+	v := &vertexProvider{}
+
+	t.Run("thinking start and continue", func(t *testing.T) {
+		ctx := newMockMultipartHttpContext()
+
+		// 1. Thinking start
+		resp1 := v.buildChatCompletionStreamResponse(ctx, &vertexChatResponse{
+			Candidates: []vertexChatCandidate{
+				{
+					Content: vertexChatContent{
+						Parts: []vertexPart{
+							{
+								Text:     "thinking process",
+								Thounght: util.Ptr(true),
+							},
+						},
+					},
+				},
+			},
+		})
+		require.NotNil(t, resp1)
+		assert.Equal(t, reasoningStartTag+"thinking process", resp1.Choices[0].Delta.Content)
+		assert.Equal(t, true, ctx.GetContext("thinking_start"))
+
+		// 2. Thinking continue
+		resp2 := v.buildChatCompletionStreamResponse(ctx, &vertexChatResponse{
+			Candidates: []vertexChatCandidate{
+				{
+					Content: vertexChatContent{
+						Parts: []vertexPart{
+							{
+								Text:     " more thinking",
+								Thounght: util.Ptr(true),
+							},
+						},
+					},
+				},
+			},
+		})
+		require.NotNil(t, resp2)
+		assert.Equal(t, " more thinking", resp2.Choices[0].Delta.Content)
+	})
+
+	t.Run("thinking end and text continue", func(t *testing.T) {
+		ctx := newMockMultipartHttpContext()
+		ctx.SetContext("thinking_start", true)
+
+		// 1. Thinking end
+		resp1 := v.buildChatCompletionStreamResponse(ctx, &vertexChatResponse{
+			Candidates: []vertexChatCandidate{
+				{
+					Content: vertexChatContent{
+						Parts: []vertexPart{
+							{
+								Text: "final answer",
+							},
+						},
+					},
+				},
+			},
+		})
+		require.NotNil(t, resp1)
+		assert.Equal(t, reasoningEndTag+"final answer", resp1.Choices[0].Delta.Content)
+		assert.Equal(t, true, ctx.GetContext("thinking_end"))
+
+		// 2. Text continue
+		resp2 := v.buildChatCompletionStreamResponse(ctx, &vertexChatResponse{
+			Candidates: []vertexChatCandidate{
+				{
+					Content: vertexChatContent{
+						Parts: []vertexPart{
+							{
+								Text: " suffix",
+							},
+						},
+					},
+				},
+			},
+		})
+		require.NotNil(t, resp2)
+		assert.Equal(t, " suffix", resp2.Choices[0].Delta.Content)
+	})
+}
+
+func TestGetNestedString(t *testing.T) {
+	// 1. empty path
+	assert.Equal(t, "", getNestedString(map[string]any{"a": 1}))
+
+	// 2. nested string exists
+	data := map[string]any{
+		"google": map[string]any{
+			"thought_signature": "sig",
+		},
+	}
+	assert.Equal(t, "sig", getNestedString(data, "google", "thought_signature"))
+
+	// 3. nested key not a map
+	data2 := map[string]any{
+		"google": "not-a-map",
+	}
+	assert.Equal(t, "", getNestedString(data2, "google", "thought_signature"))
+
+	// 4. nested key is missing
+	data3 := map[string]any{
+		"google": map[string]any{
+			"other": "val",
+		},
+	}
+	assert.Equal(t, "", getNestedString(data3, "google", "thought_signature"))
+
+	// 5. nested value is not a string
+	data4 := map[string]any{
+		"google": map[string]any{
+			"thought_signature": 1234,
+		},
+	}
+	assert.Equal(t, "", getNestedString(data4, "google", "thought_signature"))
+}