package providers import ( "context" "encoding/json" "io" "net/http" "net/http/httptest" "os" "path/filepath" "strings" "testing" ) func TestVertexDefaultAPIBase(t *testing.T) { cases := []struct { name, project, region, want string }{ {"basic", "my-proj", "us-central1", "https://us-central1-aiplatform.googleapis.com/v1/projects/my-proj/locations/us-central1/endpoints/openapi"}, {"asia", "acme", "asia-southeast1", "https://asia-southeast1-aiplatform.googleapis.com/v1/projects/acme/locations/asia-southeast1/endpoints/openapi"}, {"empty_project", "", "us-central1", ""}, {"empty_region", "my-proj", "", ""}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { if got := VertexDefaultAPIBase(tc.project, tc.region); got != tc.want { t.Errorf("got %q, want %q", got, tc.want) } }) } } func TestNewVertexProviderMissingFields(t *testing.T) { cases := []struct { name string cfg VertexConfig wantSub string }{ {"no_project", VertexConfig{Region: "us-central1"}, "project_id"}, {"no_region", VertexConfig{ProjectID: "x"}, "region"}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { _, err := NewVertexProvider(context.Background(), tc.cfg) if err == nil { t.Fatal("expected error, got nil") } if !strings.Contains(err.Error(), tc.wantSub) { t.Errorf("error %q missing %q", err, tc.wantSub) } }) } } func TestNewVertexProviderInvalidInlineJSON(t *testing.T) { _, err := NewVertexProvider(context.Background(), VertexConfig{ CredentialsJSON: "not json", ProjectID: "my-proj", Region: "us-central1", }) if err == nil { t.Fatal("expected error parsing bad JSON") } if !strings.Contains(err.Error(), "credentials") { t.Errorf("error %q does not mention credentials", err) } } func TestNewVertexProviderCredentialsFileMissing(t *testing.T) { _, err := NewVertexProvider(context.Background(), VertexConfig{ CredentialsFile: filepath.Join(t.TempDir(), "does-not-exist.json"), ProjectID: "my-proj", Region: "us-central1", }) if err == nil { t.Fatal("expected error for missing file") } if !strings.Contains(err.Error(), "read credentials file") { t.Errorf("error %q missing expected prefix", err) } } func TestNewVertexProviderCredentialsFileInvalid(t *testing.T) { dir := t.TempDir() path := filepath.Join(dir, "bad.json") if err := os.WriteFile(path, []byte("{invalid"), 0o600); err != nil { t.Fatal(err) } _, err := NewVertexProvider(context.Background(), VertexConfig{ CredentialsFile: path, ProjectID: "my-proj", Region: "us-central1", }) if err == nil { t.Fatal("expected parse error") } if !strings.Contains(err.Error(), "credentials file") { t.Errorf("error %q missing expected phrase", err) } } // TestOpenAIProviderWithoutAuthHeaderSkipsAuthorization verifies the skip-auth path // added for Vertex — doRequest() must NOT set an Authorization header when skipAuthHeader is true. // This is the sole non-trivial code change in openai.go needed for Vertex to work. func TestOpenAIProviderWithoutAuthHeaderSkipsAuthorization(t *testing.T) { var gotAuth string server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { gotAuth = r.Header.Get("Authorization") // Minimal successful openai response _, _ = io.WriteString(w, `{"id":"1","choices":[{"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}]}`) })) defer server.Close() prov := NewOpenAIProvider("test", "sk-should-not-appear", server.URL, "x"). WithoutAuthHeader() resp, err := prov.Chat(context.Background(), ChatRequest{ Messages: []Message{{Role: "user", Content: "hi"}}, }) if err != nil { t.Fatalf("chat: %v", err) } if resp.Content != "ok" { t.Errorf("content=%q, want %q", resp.Content, "ok") } if gotAuth != "" { t.Errorf("unexpected Authorization header %q — WithoutAuthHeader() should skip it", gotAuth) } } // TestOpenAIProviderWithHTTPClientUsesCustomClient verifies WithHTTPClient() replaces the default. // A transport that tags outgoing requests with a sentinel header lets us confirm the custom client // is the one used for Vertex AI (so oauth2.Transport actually runs). func TestOpenAIProviderWithHTTPClientUsesCustomClient(t *testing.T) { var sawSentinel bool server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { sawSentinel = r.Header.Get("X-Test-Transport") == "custom" _, _ = io.WriteString(w, `{"id":"1","choices":[{"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}]}`) })) defer server.Close() customClient := &http.Client{Transport: &taggingTransport{Base: http.DefaultTransport, Header: "X-Test-Transport", Value: "custom"}} prov := NewOpenAIProvider("test", "ignored", server.URL, "x"). WithHTTPClient(customClient). WithoutAuthHeader() if _, err := prov.Chat(context.Background(), ChatRequest{Messages: []Message{{Role: "user", Content: "hi"}}}); err != nil { t.Fatalf("chat: %v", err) } if !sawSentinel { t.Error("custom transport did not run — WithHTTPClient() may not have replaced the client") } } // taggingTransport is a test-only RoundTripper that sets a fixed header on every outbound request. type taggingTransport struct { Base http.RoundTripper Header string Value string } func (t *taggingTransport) RoundTrip(req *http.Request) (*http.Response, error) { req.Header.Set(t.Header, t.Value) return t.Base.RoundTrip(req) } // Sanity check: ensure Vertex provider wires default model and endpoint correctly. // We cannot exercise real token refresh without a real SA — skipAuthHeader + endpoint // assertions cover the provider-specific wiring. func TestNewVertexProviderWiresEndpointAndModel(t *testing.T) { // Valid (but fake) SA JSON — CredentialsFromJSON parses structure without fetching tokens. fakeSA := map[string]any{ "type": "service_account", "project_id": "my-proj", "private_key": fakePEM, "client_email": "test@my-proj.iam.gserviceaccount.com", "token_uri": "https://oauth2.googleapis.com/token", } data, _ := json.Marshal(fakeSA) prov, err := NewVertexProvider(context.Background(), VertexConfig{ CredentialsJSON: string(data), ProjectID: "my-proj", Region: "us-central1", }) if err != nil { t.Fatalf("NewVertexProvider: %v", err) } wantBase := "https://us-central1-aiplatform.googleapis.com/v1/projects/my-proj/locations/us-central1/endpoints/openapi" if prov.APIBase() != wantBase { t.Errorf("APIBase=%q, want %q", prov.APIBase(), wantBase) } if prov.DefaultModel() != VertexDefaultModel { t.Errorf("DefaultModel=%q, want %q", prov.DefaultModel(), VertexDefaultModel) } if prov.Name() != "vertex" { t.Errorf("Name=%q, want %q", prov.Name(), "vertex") } if prov.ProviderType() != ProviderTypeVertex { t.Errorf("ProviderType=%q, want %q", prov.ProviderType(), ProviderTypeVertex) } } // Minimal valid-looking PKCS#8 PEM body — google.CredentialsFromJSON parses lazily // so it does NOT attempt real key validation; test just needs structurally-valid JSON. // The private_key field can be any non-empty string. const fakePEM = "-----BEGIN PRIVATE KEY-----\nAAAA\n-----END PRIVATE KEY-----\n" // Regression test for H1 from code review: thought_signature detection must recognize // providers whose providerType is "vertex" (or apiBase contains "aiplatform"), // even when the model string does NOT contain "gemini". Without this fix, tool-call // rounds against a fine-tuned Vertex endpoint ID would drop the signature on passback // and trigger HTTP 400 from the Vertex API. func TestVertexProviderForwardsThoughtSignatureOnToolCalls(t *testing.T) { var bodies []string server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { b, _ := io.ReadAll(r.Body) bodies = append(bodies, string(b)) // Return a tool call with a thought_signature so the next round would echo it. _, _ = io.WriteString(w, `{"id":"1","choices":[{"message":{"role":"assistant","tool_calls":[{"id":"t1","type":"function","function":{"name":"noop","arguments":"{}","thought_signature":"sig-xyz"}}]},"finish_reason":"tool_calls"}]}`) })) defer server.Close() // Build a Vertex-style OpenAIProvider manually (avoids oauth2 in tests). prov := NewOpenAIProvider("vertex", "", server.URL, "some-tuned-endpoint-id"). WithProviderType(ProviderTypeVertex). WithoutAuthHeader() // Round 1: assistant responds with tool_calls carrying thought_signature. r1, err := prov.Chat(context.Background(), ChatRequest{ Messages: []Message{{Role: "user", Content: "go"}}, Tools: []ToolDefinition{{Type: "function", Function: &ToolFunctionSchema{Name: "noop", Parameters: map[string]any{"type": "object"}}}}, }) if err != nil { t.Fatalf("round 1: %v", err) } if len(r1.ToolCalls) != 1 { t.Fatalf("round 1 tool_calls = %d, want 1", len(r1.ToolCalls)) } if r1.ToolCalls[0].Metadata["thought_signature"] != "sig-xyz" { t.Fatalf("thought_signature metadata missing on round 1 tool call") } // Round 2: pass the assistant's tool call + a tool-result message. Expect the // outbound request to INCLUDE thought_signature on the tool_calls entry. toolCall := r1.ToolCalls[0] toolCall.Arguments = map[string]any{} _, err = prov.Chat(context.Background(), ChatRequest{ Messages: []Message{ {Role: "user", Content: "go"}, {Role: "assistant", Content: "", ToolCalls: []ToolCall{toolCall}}, {Role: "tool", Content: "ok", ToolCallID: "t1"}, {Role: "user", Content: "next"}, }, Tools: []ToolDefinition{{Type: "function", Function: &ToolFunctionSchema{Name: "noop", Parameters: map[string]any{"type": "object"}}}}, }) if err != nil { t.Fatalf("round 2: %v", err) } if len(bodies) < 2 { t.Fatalf("expected 2 round-trips, got %d", len(bodies)) } if !strings.Contains(bodies[1], `"thought_signature":"sig-xyz"`) { t.Errorf("round 2 body missing thought_signature (H1 regression): %s", bodies[1]) } } // Sanity check the validation helpers surface clear errors on bad input (M1 / M2). func TestVertexValidationRejectsMalformedInput(t *testing.T) { cases := []struct { name, project, region, apiBase, wantSub string }{ {"region_host_escape", "my-proj", "evil.com/a?", "", "invalid region"}, {"region_with_slash", "my-proj", "us/central1", "", "invalid region"}, {"project_uppercase", "MY-PROJ", "us-central1", "", "invalid project_id"}, {"project_starts_with_digit", "1badproj", "us-central1", "", "invalid project_id"}, {"project_too_short", "abc", "us-central1", "", "invalid project_id"}, {"override_http", "my-proj", "us-central1", "http://evil.com", "https scheme"}, {"override_non_google", "my-proj", "us-central1", "https://evil.com/vertex", "googleapis.com"}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { _, err := NewVertexProvider(context.Background(), VertexConfig{ ProjectID: tc.project, Region: tc.region, APIBaseOverride: tc.apiBase, }) if err == nil { t.Fatalf("expected error, got nil") } if !strings.Contains(err.Error(), tc.wantSub) { t.Errorf("error %q missing %q", err.Error(), tc.wantSub) } }) } } // Confirm well-formed projects+regions plus a valid override URL still work. func TestVertexValidationAcceptsWellFormedInput(t *testing.T) { fakeSA := map[string]any{ "type": "service_account", "project_id": "my-proj", "private_key": fakePEM, "client_email": "test@my-proj.iam.gserviceaccount.com", "token_uri": "https://oauth2.googleapis.com/token", } data, _ := json.Marshal(fakeSA) _, err := NewVertexProvider(context.Background(), VertexConfig{ CredentialsJSON: string(data), ProjectID: "my-proj", Region: "asia-southeast1", APIBaseOverride: "https://asia-southeast1-aiplatform.googleapis.com/v1/projects/my-proj/locations/asia-southeast1/endpoints/openapi", }) if err != nil { t.Fatalf("well-formed input rejected: %v", err) } }