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

package provider

import (
	"net/http"
	"strings"
	"testing"

	"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-proxy/util"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/tidwall/gjson"
)

func TestAppendOrReplaceAPIKey(t *testing.T) {
	t.Run("empty apiKey returns path unchanged", func(t *testing.T) {
		path := "/v1/publishers/google/models/gemini:generateContent"
		assert.Equal(t, path, appendOrReplaceAPIKey(path, ""))
	})

	t.Run("path without query appends ?key=", func(t *testing.T) {
		result := appendOrReplaceAPIKey("/v1/models/gemini:generateContent", "my-key")
		assert.Equal(t, "/v1/models/gemini:generateContent?key=my-key", result)
	})

	t.Run("path with existing query appends &key=", func(t *testing.T) {
		result := appendOrReplaceAPIKey("/v1/models/gemini:streamGenerateContent?alt=sse", "my-key")
		assert.Contains(t, result, "alt=sse")
		assert.Contains(t, result, "key=my-key")
	})

	t.Run("existing key parameter is replaced", func(t *testing.T) {
		result := appendOrReplaceAPIKey("/v1/models/gemini:generateContent?key=old-key&trace=1", "new-key")
		assert.Contains(t, result, "key=new-key")
		assert.NotContains(t, result, "old-key")
		assert.Contains(t, result, "trace=1")
	})

	t.Run("unparseable path without query falls back to ?key= append", func(t *testing.T) {
		// A bare string with no leading slash is not a valid RequestURI
		result := appendOrReplaceAPIKey("not-a-valid-uri", "my-key")
		assert.Equal(t, "not-a-valid-uri?key=my-key", result)
	})

	t.Run("unparseable path with query falls back to &key= append", func(t *testing.T) {
		result := appendOrReplaceAPIKey("not-a-valid-uri?foo=bar", "my-key")
		assert.Equal(t, "not-a-valid-uri?foo=bar&key=my-key", result)
	})
}

func TestVertexProviderBuildChatRequestStructuredOutputMapping(t *testing.T) {
	t.Run("json_object response format", func(t *testing.T) {
		v := &vertexProvider{}
		req := &chatCompletionRequest{
			Model: "gemini-2.5-flash",
			Messages: []chatMessage{
				{Role: roleUser, Content: "hello"},
			},
			ResponseFormat: map[string]interface{}{
				"type": "json_object",
			},
		}

		vertexReq, err := v.buildVertexChatRequest(req)
		require.NoError(t, err)
		require.NotNil(t, vertexReq)

		assert.Equal(t, util.MimeTypeApplicationJson, vertexReq.GenerationConfig.ResponseMimeType)
		assert.Nil(t, vertexReq.GenerationConfig.ResponseSchema)
	})

	t.Run("json_schema response format with nested schema", func(t *testing.T) {
		v := &vertexProvider{}
		schema := map[string]interface{}{
			"type": "object",
			"properties": map[string]interface{}{
				"answer": map[string]interface{}{
					"type": "string",
				},
			},
			"required": []interface{}{"answer"},
		}
		req := &chatCompletionRequest{
			Model: "gemini-2.5-flash",
			Messages: []chatMessage{
				{Role: roleUser, Content: "hello"},
			},
			ResponseFormat: map[string]interface{}{
				"type": "json_schema",
				"json_schema": map[string]interface{}{
					"name":   "response",
					"strict": true,
					"schema": schema,
				},
			},
		}

		vertexReq, err := v.buildVertexChatRequest(req)
		require.NoError(t, err)
		require.NotNil(t, vertexReq)

		assert.Equal(t, util.MimeTypeApplicationJson, vertexReq.GenerationConfig.ResponseMimeType)
		assert.Equal(t, schema, vertexReq.GenerationConfig.ResponseSchema)
	})

	t.Run("json_schema response format with direct schema object", func(t *testing.T) {
		v := &vertexProvider{}
		schema := map[string]interface{}{
			"type": "object",
			"properties": map[string]interface{}{
				"city": map[string]interface{}{
					"type": "string",
				},
			},
			"required": []interface{}{"city"},
		}
		req := &chatCompletionRequest{
			Model: "gemini-2.5-flash",
			Messages: []chatMessage{
				{Role: roleUser, Content: "hello"},
			},
			ResponseFormat: map[string]interface{}{
				"type":        "json_schema",
				"json_schema": schema,
			},
		}

		vertexReq, err := v.buildVertexChatRequest(req)
		require.NoError(t, err)
		require.NotNil(t, vertexReq)

		assert.Equal(t, util.MimeTypeApplicationJson, vertexReq.GenerationConfig.ResponseMimeType)
		assert.Equal(t, schema, vertexReq.GenerationConfig.ResponseSchema)
	})

	t.Run("json_schema response format without valid schema should return error", func(t *testing.T) {
		v := &vertexProvider{}
		req := &chatCompletionRequest{
			Model: "gemini-2.5-flash",
			Messages: []chatMessage{
				{Role: roleUser, Content: "hello"},
			},
			ResponseFormat: map[string]interface{}{
				"type":        "json_schema",
				"json_schema": "invalid",
			},
		}

		vertexReq, err := v.buildVertexChatRequest(req)
		require.Error(t, err)
		assert.Nil(t, vertexReq)
		assert.Contains(t, err.Error(), "invalid response_format.json_schema")
	})

	t.Run("direct schema in response_format for compatibility", func(t *testing.T) {
		v := &vertexProvider{}
		schema := map[string]interface{}{
			"type": "object",
			"properties": map[string]interface{}{
				"result": map[string]interface{}{
					"type": "string",
				},
			},
		}
		req := &chatCompletionRequest{
			Model: "gemini-2.5-flash",
			Messages: []chatMessage{
				{Role: roleUser, Content: "hello"},
			},
			ResponseFormat: schema,
		}

		vertexReq, err := v.buildVertexChatRequest(req)
		require.NoError(t, err)
		require.NotNil(t, vertexReq)

		assert.Equal(t, util.MimeTypeApplicationJson, vertexReq.GenerationConfig.ResponseMimeType)
		assert.Equal(t, schema, vertexReq.GenerationConfig.ResponseSchema)
	})

	t.Run("text response format keeps default text output", func(t *testing.T) {
		v := &vertexProvider{}
		req := &chatCompletionRequest{
			Model: "gemini-2.5-flash",
			Messages: []chatMessage{
				{Role: roleUser, Content: "hello"},
			},
			ResponseFormat: map[string]interface{}{
				"type": "text",
			},
		}

		vertexReq, err := v.buildVertexChatRequest(req)
		require.NoError(t, err)
		require.NotNil(t, vertexReq)

		assert.Empty(t, vertexReq.GenerationConfig.ResponseMimeType)
		assert.Nil(t, vertexReq.GenerationConfig.ResponseSchema)
	})

	t.Run("unknown response format does not inject schema config", func(t *testing.T) {
		v := &vertexProvider{}
		req := &chatCompletionRequest{
			Model: "gemini-2.5-flash",
			Messages: []chatMessage{
				{Role: roleUser, Content: "hello"},
			},
			ResponseFormat: map[string]interface{}{
				"type": "xml",
			},
		}

		vertexReq, err := v.buildVertexChatRequest(req)
		require.NoError(t, err)
		require.NotNil(t, vertexReq)

		assert.Empty(t, vertexReq.GenerationConfig.ResponseMimeType)
		assert.Nil(t, vertexReq.GenerationConfig.ResponseSchema)
	})

	t.Run("gemini 2.0 json_schema is ignored for stability", func(t *testing.T) {
		v := &vertexProvider{}
		schema := map[string]interface{}{
			"type": "object",
			"properties": map[string]interface{}{
				"answer": map[string]interface{}{
					"type": "string",
				},
			},
		}
		req := &chatCompletionRequest{
			Model: "gemini-2.0-flash",
			Messages: []chatMessage{
				{Role: roleUser, Content: "hello"},
			},
			ResponseFormat: map[string]interface{}{
				"type": "json_schema",
				"json_schema": map[string]interface{}{
					"name":   "response",
					"strict": true,
					"schema": schema,
				},
			},
		}

		vertexReq, err := v.buildVertexChatRequest(req)
		require.NoError(t, err)
		require.NotNil(t, vertexReq)
		assert.Empty(t, vertexReq.GenerationConfig.ResponseMimeType)
		assert.Nil(t, vertexReq.GenerationConfig.ResponseSchema)
	})

	t.Run("gemini 2.0 malformed json_schema is also ignored", func(t *testing.T) {
		v := &vertexProvider{}
		req := &chatCompletionRequest{
			Model: "gemini-2.0-flash",
			Messages: []chatMessage{
				{Role: roleUser, Content: "hello"},
			},
			ResponseFormat: map[string]interface{}{
				"type":        "json_schema",
				"json_schema": "invalid",
			},
		}

		vertexReq, err := v.buildVertexChatRequest(req)
		require.NoError(t, err)
		require.NotNil(t, vertexReq)
		assert.Empty(t, vertexReq.GenerationConfig.ResponseMimeType)
		assert.Nil(t, vertexReq.GenerationConfig.ResponseSchema)
	})

	t.Run("gemini 2.0 json_object is ignored", func(t *testing.T) {
		v := &vertexProvider{}
		req := &chatCompletionRequest{
			Model: "gemini-2.0-flash",
			Messages: []chatMessage{
				{Role: roleUser, Content: "hello"},
			},
			ResponseFormat: map[string]interface{}{
				"type": "json_object",
			},
		}

		vertexReq, err := v.buildVertexChatRequest(req)
		require.NoError(t, err)
		require.NotNil(t, vertexReq)
		assert.Empty(t, vertexReq.GenerationConfig.ResponseMimeType)
		assert.Nil(t, vertexReq.GenerationConfig.ResponseSchema)
	})
}

func TestVertexProviderApplyResponseFormatNilSafety(t *testing.T) {
	v := &vertexProvider{}
	require.NoError(t, v.applyResponseFormatToGenerationConfig(map[string]interface{}{"type": "json_object"}, nil, "gemini-2.5-flash"))
	require.NoError(t, v.applyResponseFormatToGenerationConfig(nil, &vertexChatGenerationConfig{}, "gemini-2.5-flash"))
	require.NoError(t, v.applyResponseFormatToGenerationConfig(map[string]interface{}{}, &vertexChatGenerationConfig{}, "gemini-2.5-flash"))
}

// newAnthropicVertexProvider builds a vertexProvider with project/region/modelMapping
// suitable for exercising onAnthropicMessagesRequestBody without OAuth or wasm runtime.
func newAnthropicVertexProvider(openAICompat bool) *vertexProvider {
	cfg := ProviderConfig{
		vertexProjectId:        "test-proj",
		vertexRegion:           "us-east5",
		vertexOpenAICompatible: openAICompat,
		modelMapping: map[string]string{
			"claude-sonnet-4":   "claude-sonnet-4@20250514",
			"claude-sonnet-4-5": "claude-sonnet-4-5@20250929",
		},
	}
	return &vertexProvider{config: cfg}
}

// TestVertexAnthropicPassthrough_BuiltinTool_TypePreserved is the core regression test
// for the original bug: builtin Anthropic tools (e.g. web_search_20250305) carry only
// a `type` discriminator and no `name`. The previous Anthropic→OpenAI→Anthropic round
// trip lost the type field, producing `tools.0.custom.name: String should have at least
// 1 character` from vertex. After the fix, the body is passthrough — `type` survives.
func TestVertexAnthropicPassthrough_BuiltinTool_TypePreserved(t *testing.T) {
	v := newAnthropicVertexProvider(false)
	ctx := newMapCtx()
	headers := http.Header{}
	body := []byte(`{
		"model": "claude-sonnet-4",
		"max_tokens": 4096,
		"messages": [{"role": "user", "content": "search the web"}],
		"tools": [
			{"type": "web_search_20250305"},
			{"type": "bash_20250124"},
			{"type": "text_editor_20250124"}
		]
	}`)

	out, err := v.onAnthropicMessagesRequestBody(ctx, body, headers)
	require.NoError(t, err)

	// Path: non-stream → :rawPredict, model fully-qualified via modelMapping.
	assert.Equal(t,
		"/v1/projects/test-proj/locations/us-east5/publishers/anthropic/models/claude-sonnet-4@20250514:rawPredict",
		headers.Get(":path"))

	// Body: model stripped, anthropic_version injected.
	assert.False(t, gjson.GetBytes(out, "model").Exists(), "model must be stripped (vertex :rawPredict rejects it)")
	assert.Equal(t, vertexAnthropicVersion, gjson.GetBytes(out, "anthropic_version").String())

	// The bug-defining assertion: builtin tool `type` survives verbatim, and we did
	// NOT manufacture a `name` for it. If a future change re-introduces the lossy
	// conversion, the type will disappear or a synthetic name will appear and this
	// test will fail.
	tools := gjson.GetBytes(out, "tools").Array()
	require.Len(t, tools, 3)
	assert.Equal(t, "web_search_20250305", tools[0].Get("type").String())
	assert.False(t, tools[0].Get("name").Exists(), "builtin tool must not have a synthetic name")
	assert.Equal(t, "bash_20250124", tools[1].Get("type").String())
	assert.Equal(t, "text_editor_20250124", tools[2].Get("type").String())
}

// TestVertexAnthropicPassthrough_StreamPath verifies stream=true routes to
// :streamRawPredict and stream=false routes to :rawPredict.
func TestVertexAnthropicPassthrough_StreamPath(t *testing.T) {
	t.Run("stream true → streamRawPredict", func(t *testing.T) {
		v := newAnthropicVertexProvider(false)
		ctx := newMapCtx()
		headers := http.Header{}
		body := []byte(`{"model":"claude-sonnet-4","max_tokens":16,"stream":true,"messages":[{"role":"user","content":"hi"}]}`)

		_, err := v.onAnthropicMessagesRequestBody(ctx, body, headers)
		require.NoError(t, err)

		assert.Equal(t,
			"/v1/projects/test-proj/locations/us-east5/publishers/anthropic/models/claude-sonnet-4@20250514:streamRawPredict",
			headers.Get(":path"))
	})

	t.Run("stream false → rawPredict", func(t *testing.T) {
		v := newAnthropicVertexProvider(false)
		ctx := newMapCtx()
		headers := http.Header{}
		body := []byte(`{"model":"claude-sonnet-4","max_tokens":16,"messages":[{"role":"user","content":"hi"}]}`)

		_, err := v.onAnthropicMessagesRequestBody(ctx, body, headers)
		require.NoError(t, err)

		assert.Equal(t,
			"/v1/projects/test-proj/locations/us-east5/publishers/anthropic/models/claude-sonnet-4@20250514:rawPredict",
			headers.Get(":path"))
	})
}

// TestVertexAnthropicPassthrough_ModelMappingUnconfigured verifies that when no
// mapping entry matches, the model name is left untouched (vertex will 404 — we
// don't second-guess the user's config here).
func TestVertexAnthropicPassthrough_ModelMappingUnconfigured(t *testing.T) {
	v := &vertexProvider{config: ProviderConfig{
		vertexProjectId: "test-proj",
		vertexRegion:    "us-east5",
		// no modelMapping configured
	}}
	ctx := newMapCtx()
	headers := http.Header{}
	body := []byte(`{"model":"claude-sonnet-4","max_tokens":16,"messages":[{"role":"user","content":"hi"}]}`)

	_, err := v.onAnthropicMessagesRequestBody(ctx, body, headers)
	require.NoError(t, err)

	// model name passes through as-is (no @date suffix)
	assert.Equal(t,
		"/v1/projects/test-proj/locations/us-east5/publishers/anthropic/models/claude-sonnet-4:rawPredict",
		headers.Get(":path"))
}

// TestVertexAnthropicPassthrough_CustomToolFieldsPreserved verifies that
// custom tool fields not in the OpenAI schema (cache_control, thinking config,
// arbitrary input_schema shapes) survive the passthrough — they were silently
// dropped by the old double-conversion path.
func TestVertexAnthropicPassthrough_CustomToolFieldsPreserved(t *testing.T) {
	v := newAnthropicVertexProvider(false)
	ctx := newMapCtx()
	headers := http.Header{}
	body := []byte(`{
		"model": "claude-sonnet-4",
		"max_tokens": 1024,
		"messages": [{"role": "user", "content": "list files"}],
		"tools": [{
			"name": "Bash",
			"description": "run a shell command",
			"input_schema": {
				"type": "object",
				"properties": {"command": {"type": "string"}},
				"required": ["command"]
			},
			"cache_control": {"type": "ephemeral"}
		}],
		"thinking": {"type": "enabled", "budget_tokens": 1024}
	}`)

	out, err := v.onAnthropicMessagesRequestBody(ctx, body, headers)
	require.NoError(t, err)

	tool := gjson.GetBytes(out, "tools.0")
	assert.Equal(t, "Bash", tool.Get("name").String())
	assert.Equal(t, "ephemeral", tool.Get("cache_control.type").String(), "cache_control must survive passthrough")
	assert.Equal(t, "object", tool.Get("input_schema.type").String())
	assert.Equal(t, "command", tool.Get("input_schema.required.0").String())

	thinking := gjson.GetBytes(out, "thinking")
	assert.Equal(t, "enabled", thinking.Get("type").String(), "thinking config must survive passthrough")
	assert.Equal(t, int64(1024), thinking.Get("budget_tokens").Int())
}

// TestVertexAnthropicPassthrough_OpenAICompatibleConfigDoesNotInterfere verifies
// the contract from the plan: vertexOpenAICompatible: true affects ONLY
// chat/completions; /v1/messages still goes to the Anthropic native endpoint.
// (We exercise the handler that TransformRequestBodyHeaders dispatches to;
// see OnRequestBody:302 for the bypass condition itself.)
func TestVertexAnthropicPassthrough_OpenAICompatibleConfigDoesNotInterfere(t *testing.T) {
	v := newAnthropicVertexProvider(true) // vertexOpenAICompatible: true
	ctx := newMapCtx()
	headers := http.Header{}
	body := []byte(`{"model":"claude-sonnet-4","max_tokens":16,"messages":[{"role":"user","content":"hi"}]}`)

	out, err := v.onAnthropicMessagesRequestBody(ctx, body, headers)
	require.NoError(t, err)

	// Must be Anthropic native path, NOT /v1beta1/.../openai/chat/completions.
	assert.Contains(t, headers.Get(":path"), "publishers/anthropic/models/")
	assert.Contains(t, headers.Get(":path"), ":rawPredict")
	assert.NotContains(t, headers.Get(":path"), "/openai/")
	assert.Equal(t, vertexAnthropicVersion, gjson.GetBytes(out, "anthropic_version").String())
}

// TestVertexAnthropicPassthrough_DispatchedFromTransformRequestBodyHeaders covers
// the wiring step: TransformRequestBodyHeaders sees ApiNameAnthropicMessages and
// routes to the passthrough handler. Guards against accidental removal of the
// case branch.
func TestVertexAnthropicPassthrough_DispatchedFromTransformRequestBodyHeaders(t *testing.T) {
	v := newAnthropicVertexProvider(false)
	ctx := newMapCtx()
	headers := http.Header{}
	body := []byte(`{"model":"claude-sonnet-4","max_tokens":16,"messages":[{"role":"user","content":"hi"}],"tools":[{"type":"web_search_20250305"}]}`)

	out, err := v.TransformRequestBodyHeaders(ctx, ApiNameAnthropicMessages, body, headers)
	require.NoError(t, err)

	assert.Contains(t, headers.Get(":path"), ":rawPredict")
	assert.Equal(t, "web_search_20250305", gjson.GetBytes(out, "tools.0.type").String())
	assert.False(t, gjson.GetBytes(out, "model").Exists())
}

// TestVertexAnthropicPassthrough_ResponseBodyUnchanged verifies the non-stream
// response branch: TransformResponseBody returns the body verbatim for
// ApiNameAnthropicMessages, so vertex's native Anthropic JSON reaches the client
// without OpenAI→Anthropic re-translation.
func TestVertexAnthropicPassthrough_ResponseBodyUnchanged(t *testing.T) {
	v := newAnthropicVertexProvider(false)
	ctx := newMapCtx()
	body := []byte(`{"id":"msg_01","type":"message","role":"assistant","content":[{"type":"text","text":"hi"}],"model":"claude-sonnet-4@20250514","stop_reason":"end_turn","usage":{"input_tokens":3,"output_tokens":1}}`)

	out, err := v.TransformResponseBody(ctx, ApiNameAnthropicMessages, body)
	require.NoError(t, err)
	assert.Equal(t, body, out, "vertex Anthropic response must be returned byte-for-byte")
}

// TestVertexAnthropicPassthrough_StreamingChunkUnchanged verifies the streaming
// counterpart: each SSE chunk is forwarded verbatim because vertex's
// :streamRawPredict already emits standard Anthropic SSE events.
func TestVertexAnthropicPassthrough_StreamingChunkUnchanged(t *testing.T) {
	v := newAnthropicVertexProvider(false)
	ctx := newMapCtx()
	chunk := []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_01\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-sonnet-4@20250514\",\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":3,\"output_tokens\":1}}}\n\n")

	out, err := v.OnStreamingResponseBody(ctx, ApiNameAnthropicMessages, chunk, false)
	require.NoError(t, err)
	assert.Equal(t, chunk, out, "vertex Anthropic SSE chunk must be returned byte-for-byte")
}

// TestVertexTransformRequestHeaders_StripsAnthropicHeaders ensures that
// Anthropic-specific headers (credentials + protocol) are NOT forwarded to
// vertex. The regular Anthropic SDK sends these but vertex's Anthropic
// endpoint rejects or misinterprets them:
//   - x-api-key / anthropic-api-key: credential leak to Google logs
//   - anthropic-beta: vertex 400 "Unexpected value(s) ... for the anthropic-beta header"
//   - anthropic-version: conflicts with body-level anthropic_version "vertex-2023-10-16"
func TestVertexTransformRequestHeaders_StripsAnthropicHeaders(t *testing.T) {
	v := newAnthropicVertexProvider(false)
	ctx := newMapCtx()
	headers := http.Header{}
	headers.Set("x-api-key", "sk-ant-api03-secret")
	headers.Set("anthropic-api-key", "sk-ant-api03-secret")
	headers.Set("anthropic-beta", "advanced-tool-use-2025-11-20,prompt-caching-scope-2026-01-05")
	headers.Set("anthropic-version", "2023-06-01")
	headers.Set("content-type", "application/json")

	v.TransformRequestHeaders(ctx, ApiNameAnthropicMessages, headers)

	assert.Empty(t, headers.Get("x-api-key"), "x-api-key must be stripped before forwarding to vertex")
	assert.Empty(t, headers.Get("anthropic-api-key"), "anthropic-api-key must be stripped before forwarding to vertex")
	assert.Empty(t, headers.Get("anthropic-beta"), "anthropic-beta must be stripped — vertex rejects unknown beta flags with 400")
	assert.Empty(t, headers.Get("anthropic-version"), "anthropic-version must be stripped — vertex uses body-level anthropic_version instead")
	// Sanity: unrelated headers untouched.
	assert.Equal(t, "application/json", headers.Get("content-type"))
}

// TestVertexAnthropicPassthrough_MaxTokensDefault ensures that when the client
// omits max_tokens, the passthrough handler injects claudeDefaultMaxTokens.
// Vertex's Anthropic endpoint rejects requests without max_tokens with a 400 —
// some SDKs (and lenient clients) leave it unset, expecting the upstream to
// default. Matches buildClaudeTextGenRequest's behavior in claude.go.
func TestVertexAnthropicPassthrough_MaxTokensDefault(t *testing.T) {
	t.Run("missing max_tokens gets defaulted", func(t *testing.T) {
		v := newAnthropicVertexProvider(false)
		ctx := newMapCtx()
		headers := http.Header{}
		body := []byte(`{
			"model": "claude-sonnet-4",
			"messages": [{"role": "user", "content": "hi"}]
		}`)

		out, err := v.onAnthropicMessagesRequestBody(ctx, body, headers)
		require.NoError(t, err)

		assert.Equal(t, int64(claudeDefaultMaxTokens), gjson.GetBytes(out, "max_tokens").Int())
	})

	t.Run("client-supplied max_tokens preserved", func(t *testing.T) {
		v := newAnthropicVertexProvider(false)
		ctx := newMapCtx()
		headers := http.Header{}
		body := []byte(`{
			"model": "claude-sonnet-4",
			"max_tokens": 1024,
			"messages": [{"role": "user", "content": "hi"}]
		}`)

		out, err := v.onAnthropicMessagesRequestBody(ctx, body, headers)
		require.NoError(t, err)

		assert.Equal(t, int64(1024), gjson.GetBytes(out, "max_tokens").Int(),
			"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.NotEmpty(t, response.Choices[0].Message.ToolCalls[0].Id)
	assert.True(t, strings.HasPrefix(response.Choices[0].Message.ToolCalls[0].Id, "call_"))
	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 TestVertexProviderStreamToolCallIncludesStableID(t *testing.T) {
	v := &vertexProvider{}
	ctx := newMockMultipartHttpContext()
	ctx.SetContext(ctxKeyFinalRequestModel, "gemini-3.1-pro-preview")
	vertexResp := &vertexChatResponse{
		ResponseId: "vertex-response-id",
		Candidates: []vertexChatCandidate{
			{
				Index: 0,
				Content: vertexChatContent{
					Role: "model",
					Parts: []vertexPart{
						{
							FunctionCall: &vertexFunctionCall{
								Name: "lookup",
								Args: map[string]interface{}{"query": "test"},
							},
							ThoughtSignature: "thought-signature-from-vertex",
						},
					},
				},
				FinishReason: "STOP",
			},
		},
	}

	first := v.buildChatCompletionStreamResponse(ctx, vertexResp)
	second := v.buildChatCompletionStreamResponse(ctx, vertexResp)

	require.Len(t, first.Choices, 1)
	require.NotNil(t, first.Choices[0].Delta)
	require.Len(t, first.Choices[0].Delta.ToolCalls, 1)
	firstID := first.Choices[0].Delta.ToolCalls[0].Id
	assert.NotEmpty(t, firstID)
	assert.True(t, strings.HasPrefix(firstID, "call_"))

	require.Len(t, second.Choices, 1)
	require.NotNil(t, second.Choices[0].Delta)
	require.Len(t, second.Choices[0].Delta.ToolCalls, 1)
	assert.Equal(t, firstID, second.Choices[0].Delta.ToolCalls[0].Id)
}

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"))
}