diff --git a/go/adk/pkg/models/anthropic.go b/go/adk/pkg/models/anthropic.go index 0dee13adc..2c08ccfdf 100644 --- a/go/adk/pkg/models/anthropic.go +++ b/go/adk/pkg/models/anthropic.go @@ -3,6 +3,7 @@ package models import ( "context" "fmt" + "net/http" "os" "github.com/anthropics/anthropic-sdk-go" @@ -11,8 +12,16 @@ import ( "github.com/anthropics/anthropic-sdk-go/vertex" awsconfig "github.com/aws/aws-sdk-go-v2/config" "github.com/go-logr/logr" + "golang.org/x/oauth2" + "golang.org/x/oauth2/google" ) +// vertexAIScope is the OAuth2 scope required for Vertex AI. This mirrors the +// scope used internally by anthropic-sdk-go's vertex.WithGoogleAuth when no +// scopes are supplied, so behavior stays consistent when we resolve +// credentials ourselves via composeVertexHTTPClient. +const vertexAIScope = "https://www.googleapis.com/auth/cloud-platform" + // anthropicPassthroughOpts returns a per-request option that sets the Anthropic API key // from the bearer token in ctx when APIKeyPassthrough is enabled. The Anthropic SDK sends // this as the x-api-key header, which is the correct auth mechanism for Anthropic. @@ -88,20 +97,68 @@ func newAnthropicModelFromConfig(config *AnthropicConfig, apiKey string, logger }, nil } +// composeVertexHTTPClient returns an *http.Client that layers OAuth2 token +// injection on top of the caller-supplied base transport (custom TLS, custom +// headers, connect timeout, etc.). It preserves base.Timeout so the request +// timeout from TransportConfig is honored. +// +// The composed client is what makes it safe to hand a fully-customized +// *http.Client to the Anthropic SDK for Vertex: without this, either the +// Vertex option would wipe our custom transport stack, or a naive +// WithHTTPClient would wipe the SDK's OAuth2-wrapped client. +func composeVertexHTTPClient(base *http.Client, tokenSource oauth2.TokenSource) *http.Client { + baseTransport := base.Transport + if baseTransport == nil { + baseTransport = http.DefaultTransport + } + return &http.Client{ + Timeout: base.Timeout, + Transport: &oauth2.Transport{ + Base: baseTransport, + Source: tokenSource, + }, + } +} + // NewAnthropicVertexAIModelWithLogger creates an Anthropic model that authenticates // via Google Cloud Vertex AI using Application Default Credentials (ADC). // This is used for the GeminiAnthropic / AnthropicVertexAI provider type. +// +// Composition contract: +// - vertex.WithCredentials is applied first so its base URL and its +// URL-rewrite/anthropic_version middleware are registered. +// - option.WithHTTPClient is applied last with a client we build ourselves +// from the user's TransportConfig, wrapped by an oauth2.Transport so the +// custom TLS / headers / timeout stack AND the Google OAuth2 token are +// both applied on the wire. Base URL and middleware are stored on +// separate fields on the SDK's request config and middleware is additive, +// so only the HTTP client field is overridden. func NewAnthropicVertexAIModelWithLogger(ctx context.Context, config *AnthropicConfig, region, projectID string, logger logr.Logger) (*AnthropicModel, error) { - opts := []option.RequestOption{ - vertex.WithGoogleAuth(ctx, region, projectID), - } - - // Create HTTP client with timeout, custom headers, TLS, and passthrough + // Build the caller's HTTP client (TLS, headers, timeout). httpClient, err := BuildHTTPClient(config.TransportConfig) if err != nil { return nil, err } - opts = append(opts, option.WithHTTPClient(httpClient)) + + // Resolve Google Application Default Credentials ourselves so we can + // wrap the token source into our transport stack. This mirrors what + // vertex.WithGoogleAuth does internally, but avoids its panic-on-error + // behavior and avoids the double credential lookup that would happen + // if we let WithGoogleAuth build its own HTTP client only to discard it. + creds, err := google.FindDefaultCredentials(ctx, vertexAIScope) + if err != nil { + return nil, fmt.Errorf("failed to find Google default credentials for Vertex AI: %w", err) + } + + composedClient := composeVertexHTTPClient(httpClient, creds.TokenSource) + + opts := []option.RequestOption{ + // Registers Vertex base URL + URL-rewrite middleware. + vertex.WithCredentials(ctx, region, projectID, creds), + // Overrides only the HTTP client field. Base URL and middleware + // registered above survive. + option.WithHTTPClient(composedClient), + } client := anthropic.NewClient(opts...) logger.Info("Initialized Anthropic Vertex AI model", "model", config.Model, "region", region, "project", projectID) diff --git a/go/adk/pkg/models/anthropic_vertex_test.go b/go/adk/pkg/models/anthropic_vertex_test.go new file mode 100644 index 000000000..cb352fc8d --- /dev/null +++ b/go/adk/pkg/models/anthropic_vertex_test.go @@ -0,0 +1,137 @@ +package models + +import ( + "io" + "net/http" + "net/http/httptest" + "testing" + "time" + + "golang.org/x/oauth2" +) + +// TestComposeVertexHTTPClient_PreservesTimeout asserts the composed client +// inherits the caller's timeout so TransportConfig.Timeout is honored. +func TestComposeVertexHTTPClient_PreservesTimeout(t *testing.T) { + base := &http.Client{Timeout: 42 * time.Second} + ts := oauth2.StaticTokenSource(&oauth2.Token{AccessToken: "unused"}) + + got := composeVertexHTTPClient(base, ts) + if got.Timeout != 42*time.Second { + t.Errorf("Timeout = %v, want 42s", got.Timeout) + } + if _, ok := got.Transport.(*oauth2.Transport); !ok { + t.Errorf("Transport = %T, want *oauth2.Transport", got.Transport) + } +} + +// TestComposeVertexHTTPClient_NilBaseTransport asserts the composed client +// falls back to http.DefaultTransport when base.Transport is nil, matching +// the standard library's behavior. +func TestComposeVertexHTTPClient_NilBaseTransport(t *testing.T) { + base := &http.Client{} + ts := oauth2.StaticTokenSource(&oauth2.Token{AccessToken: "unused"}) + + got := composeVertexHTTPClient(base, ts) + oauthTransport, ok := got.Transport.(*oauth2.Transport) + if !ok { + t.Fatalf("Transport = %T, want *oauth2.Transport", got.Transport) + } + if oauthTransport.Base == nil { + t.Errorf("oauth2.Transport.Base is nil, want http.DefaultTransport fallback") + } +} + +// TestComposeVertexHTTPClient_AttachesBearerToken verifies that requests made +// through the composed client carry an Authorization: Bearer header +// sourced from the supplied oauth2.TokenSource. This is the regression guard +// for the original 401 bug. +func TestComposeVertexHTTPClient_AttachesBearerToken(t *testing.T) { + var gotAuth string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth = r.Header.Get("Authorization") + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + base := &http.Client{Transport: http.DefaultTransport, Timeout: 5 * time.Second} + ts := oauth2.StaticTokenSource(&oauth2.Token{AccessToken: "test-token"}) + client := composeVertexHTTPClient(base, ts) + + resp, err := client.Get(server.URL) + if err != nil { + t.Fatalf("Get: %v", err) + } + defer resp.Body.Close() + _, _ = io.Copy(io.Discard, resp.Body) + + if want := "Bearer test-token"; gotAuth != want { + t.Errorf("Authorization = %q, want %q", gotAuth, want) + } +} + +// TestComposeVertexHTTPClient_PreservesCustomHeaders verifies that +// TransportConfig.Headers set via the headerTransport survive composition, +// so operators' custom headers (e.g. proxy tokens) still reach the wire +// alongside the OAuth2 Authorization header. +func TestComposeVertexHTTPClient_PreservesCustomHeaders(t *testing.T) { + var gotAuth, gotCustom string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth = r.Header.Get("Authorization") + gotCustom = r.Header.Get("X-Custom-Header") + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + base, err := BuildHTTPClient(TransportConfig{ + Headers: map[string]string{"X-Custom-Header": "custom-value"}, + }) + if err != nil { + t.Fatalf("BuildHTTPClient: %v", err) + } + + ts := oauth2.StaticTokenSource(&oauth2.Token{AccessToken: "test-token"}) + client := composeVertexHTTPClient(base, ts) + + resp, err := client.Get(server.URL) + if err != nil { + t.Fatalf("Get: %v", err) + } + defer resp.Body.Close() + _, _ = io.Copy(io.Discard, resp.Body) + + if want := "Bearer test-token"; gotAuth != want { + t.Errorf("Authorization = %q, want %q", gotAuth, want) + } + if want := "custom-value"; gotCustom != want { + t.Errorf("X-Custom-Header = %q, want %q", gotCustom, want) + } +} + +// TestComposeVertexHTTPClient_PreservesBaseTransportType is a lightweight +// guard that the composed client's oauth2.Transport wraps the caller's +// transport chain (not a bare http.DefaultTransport). This gives us +// confidence that custom TLS config on the base transport survives +// composition, without needing a full mTLS harness. +func TestComposeVertexHTTPClient_PreservesBaseTransportType(t *testing.T) { + insecure := true + base, err := BuildHTTPClient(TransportConfig{TLSInsecureSkipVerify: &insecure}) + if err != nil { + t.Fatalf("BuildHTTPClient: %v", err) + } + if base.Transport == nil { + t.Fatal("expected BuildHTTPClient to produce a non-nil Transport when TLS is customized") + } + originalTransport := base.Transport + + ts := oauth2.StaticTokenSource(&oauth2.Token{AccessToken: "unused"}) + client := composeVertexHTTPClient(base, ts) + + oauthTransport, ok := client.Transport.(*oauth2.Transport) + if !ok { + t.Fatalf("Transport = %T, want *oauth2.Transport", client.Transport) + } + if oauthTransport.Base != originalTransport { + t.Errorf("oauth2.Transport.Base = %p, want caller's transport %p", oauthTransport.Base, originalTransport) + } +} diff --git a/go/go.mod b/go/go.mod index ffb95ef58..597318c3e 100644 --- a/go/go.mod +++ b/go/go.mod @@ -79,6 +79,7 @@ require ( go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.44.0 go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.44.0 go.opentelemetry.io/otel/sdk/log v0.20.0 + golang.org/x/oauth2 v0.36.0 google.golang.org/grpc v1.82.1 k8s.io/apiextensions-apiserver v0.36.3 ) @@ -423,7 +424,6 @@ require ( golang.org/x/exp/typeparams v0.0.0-20260209203927-2842357ff358 // indirect golang.org/x/mod v0.37.0 // indirect golang.org/x/net v0.57.0 // indirect - golang.org/x/oauth2 v0.36.0 // indirect golang.org/x/sync v0.22.0 // indirect golang.org/x/sys v0.47.0 // indirect golang.org/x/term v0.45.0 // indirect