diff --git a/agent/common_context.go b/agent/common_context.go index a78268178..f3af44c3e 100644 --- a/agent/common_context.go +++ b/agent/common_context.go @@ -23,12 +23,38 @@ import ( "google.golang.org/genai" "google.golang.org/adk/v2/artifact" + "google.golang.org/adk/v2/internal/adkcontext" "google.golang.org/adk/v2/memory" "google.golang.org/adk/v2/platform" "google.golang.org/adk/v2/session" "google.golang.org/adk/v2/tool/toolconfirmation" ) +// Identity is the read-only identity of an ADK invocation: the acting user, app +// name, and session a call belongs to. It is recovered from a plain +// context.Context via [IdentityFromContext]. +type Identity struct { + UserID string + AppName string + SessionID string +} + +// IdentityFromContext returns the ADK invocation [Identity] carried by ctx, if +// present. +// +// ADK contexts embed context.Context and register their identity under a private +// key, so code that only holds a context.Context — for example an +// http.RoundTripper running deep beneath a tool call, past intermediaries that +// wrap the context — can recover the acting identity without threading a typed +// context through every layer. +// +// It returns (zero, false) for a context that does not descend from an ADK +// context (a non-agent caller). +func IdentityFromContext(ctx context.Context) (Identity, bool) { + id, ok := ctx.Value(adkcontext.IdentityKey).(Identity) + return id, ok +} + // In general CommonContext should not be wrapped with contexts not providing agent.Context. // It allows to copy&modify context instead of building chains. @@ -344,6 +370,19 @@ func (c *commonContext) UserID() string { return c.invocationContext.Session().UserID() } +// Value implements context.Context. For the ADK self key it returns this +// invocation's [Identity] (so [IdentityFromContext] can recover it from a +// derived context); every other key delegates to the embedded context, +// preserving existing behavior. A context without a usable session also +// delegates, so a wrapped ADK context can still supply the identity and Value +// never panics. +func (c *commonContext) Value(key any) any { + if key == adkcontext.IdentityKey && c.invocationContext != nil && c.invocationContext.Session() != nil { + return Identity{UserID: c.UserID(), AppName: c.AppName(), SessionID: c.SessionID()} + } + return c.Context.Value(key) +} + var ( _ Context = (*commonContext)(nil) _ InvocationContext = (*commonContext)(nil) diff --git a/auth/gcp/provider.go b/auth/gcp/provider.go new file mode 100644 index 000000000..5198e2d23 --- /dev/null +++ b/auth/gcp/provider.go @@ -0,0 +1,142 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package gcp + +import ( + "context" + "errors" + "slices" + "sync" + "time" + + "golang.org/x/sync/singleflight" + + "google.golang.org/adk/v2/agent" + "google.golang.org/adk/v2/auth" +) + +// Scheme identifies a GCP auth resource and the access it requests. It mirrors +// adk-python's GcpAuthProviderScheme. +type Scheme struct { + // Name is the full resource name, routed by [Client]: either + // projects/*/locations/*/connectors/* (IAM Connector) or + // projects/*/locations/*/authProviders/* (Agent Identity). + Name string + // Scopes are the OAuth scopes requested for the credential. + Scopes []string + // ContinueURI is the developer-hosted URI used to finalize managed-OAuth + // (3-legged) flows; unused by non-interactive flows. + ContinueURI string +} + +// ProviderConfig configures a provider built by [NewProvider]. A nil +// *ProviderConfig, or any zero-valued field, uses the corresponding default. +type ProviderConfig struct { + // Client reaches the credential services. When nil, a default client (backed + // by Application Default Credentials) is created lazily on first use. + Client *Client +} + +// NewProvider returns an [auth.CredentialProvider] that resolves credentials for +// the given GCP resource via the Agent Identity / IAM Connector services. +// +// The acting user is taken from the ADK context ([agent.IdentityFromContext]) at +// resolve time, so the provider must run within an agent invocation (e.g. wired +// into mcptoolset or remoteagent). +func NewProvider(scheme Scheme, cfg *ProviderConfig) (auth.CredentialProvider, error) { + if scheme.Name == "" { + return nil, errors.New("gcp: NewProvider requires a scheme Name") + } + if cfg == nil { + cfg = &ProviderConfig{} + } + // Defensive copy: the provider outlives this call and re-reads Scopes per request, so it must not alias a caller-mutable slice. + scheme.Scopes = slices.Clone(scheme.Scopes) + return &provider{scheme: scheme, client: cfg.Client}, nil +} + +type provider struct { + scheme Scheme + + mu sync.Mutex + client *Client + clientInit singleflight.Group // coalesces concurrent first-time client init +} + +var _ auth.CredentialProvider = (*provider)(nil) + +// Credential implements [auth.CredentialProvider]. +func (p *provider) Credential(ctx context.Context) (auth.Credential, error) { + id, ok := agent.IdentityFromContext(ctx) + if !ok || id.UserID == "" { + return nil, errors.New("gcp: no acting user in ADK context; provider must run within an agent invocation") + } + + client, err := p.resolveClient(ctx) + if err != nil { + return nil, err + } + return client.RetrieveCredential(ctx, Request{ + Resource: p.scheme.Name, + UserID: id.UserID, + Scopes: p.scheme.Scopes, + ContinueURI: p.scheme.ContinueURI, + }) +} + +// clientInitTimeout bounds the lazy ADC lookup so a hung probe fails fast +// instead of blocking callers indefinitely. +const clientInitTimeout = 30 * time.Second + +// resolveClient returns the configured client, creating a default one (backed by +// Application Default Credentials) on first use. +// +// Concurrent first callers are coalesced via singleflight so the (up to +// clientInitTimeout) ADC lookup runs once, off the mutex. Construction uses +// [context.WithoutCancel] so the cached client isn't tied to one request's +// cancellation (but keeps its values). A failed init is not cached; the next +// call retries. +func (p *provider) resolveClient(ctx context.Context) (*Client, error) { + p.mu.Lock() + c := p.client + p.mu.Unlock() + if c != nil { + return c, nil + } + + v, err, _ := p.clientInit.Do("client", func() (any, error) { + // A prior winner may have set the client while we waited on Do. + p.mu.Lock() + c := p.client + p.mu.Unlock() + if c != nil { + return c, nil + } + dctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), clientInitTimeout) + defer cancel() + nc, err := NewClient(dctx, nil) + if err != nil { + return nil, err + } + p.mu.Lock() + p.client = nc + p.mu.Unlock() + return nc, nil + }) + if err != nil { + return nil, err + } + return v.(*Client), nil +} diff --git a/auth/gcp/provider_test.go b/auth/gcp/provider_test.go new file mode 100644 index 000000000..290e10a9d --- /dev/null +++ b/auth/gcp/provider_test.go @@ -0,0 +1,100 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package gcp_test + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "testing" + + "google.golang.org/adk/v2/auth" + "google.golang.org/adk/v2/auth/gcp" + icontext "google.golang.org/adk/v2/internal/context" + "google.golang.org/adk/v2/session" +) + +func TestProviderCredential(t *testing.T) { + var gotUserID string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body struct { + UserID string `json:"userId"` + } + _ = json.NewDecoder(r.Body).Decode(&body) + gotUserID = body.UserID + _, _ = io.WriteString(w, `{"success":{"token":"tok","header":"Authorization: Bearer"}}`) + })) + defer srv.Close() + + client, err := gcp.NewClient(t.Context(), &gcp.Config{ + HTTPClient: srv.Client(), + AgentIdentityEndpoint: srv.URL, + }) + if err != nil { + t.Fatalf("NewClient() error = %v", err) + } + p, err := gcp.NewProvider( + gcp.Scheme{Name: "projects/p/locations/l/authProviders/ap", Scopes: []string{"s1"}}, + &gcp.ProviderConfig{Client: client}, + ) + if err != nil { + t.Fatalf("NewProvider() error = %v", err) + } + + cred, err := p.Credential(adkContext(t, "user-1")) + if err != nil { + t.Fatalf("Credential() error = %v", err) + } + if bc, ok := cred.(auth.BearerCredential); !ok || bc.Token != "tok" { + t.Fatalf("credential = %+v, want bearer token %q", cred, "tok") + } + if gotUserID != "user-1" { + t.Errorf("service saw userId = %q, want %q (identity from agent.IdentityFromContext)", gotUserID, "user-1") + } +} + +func TestProviderRequiresADKContext(t *testing.T) { + p, err := gcp.NewProvider( + gcp.Scheme{Name: "projects/p/locations/l/authProviders/ap"}, + &gcp.ProviderConfig{Client: &gcp.Client{}}, + ) + if err != nil { + t.Fatalf("NewProvider() error = %v", err) + } + // Plain context: no ADK identity to resolve, so no service call is made. + if _, err := p.Credential(t.Context()); err == nil { + t.Fatal("Credential() = nil error, want error for missing ADK context") + } +} + +func TestNewProviderValidatesScheme(t *testing.T) { + if _, err := gcp.NewProvider(gcp.Scheme{}, nil); err == nil { + t.Fatal("NewProvider() = nil error, want error for empty scheme Name") + } +} + +// adkContext returns an ADK invocation context (recoverable via agent.IdentityFromContext) +// for the given user. +func adkContext(t *testing.T, userID string) context.Context { + t.Helper() + svc := session.InMemoryService() + resp, err := svc.Create(t.Context(), &session.CreateRequest{AppName: "app", UserID: userID}) + if err != nil { + t.Fatalf("session Create() error = %v", err) + } + return icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{Session: resp.Session}) +} diff --git a/auth/providers.go b/auth/providers.go index 9c9866c3e..53cb43a68 100644 --- a/auth/providers.go +++ b/auth/providers.go @@ -33,9 +33,9 @@ type CredentialProvider interface { // and deadlines. // // A provider that needs the acting user's identity (for example the GCP - // provider, which keys on the user) recovers the ADK context from ctx via a - // shared helper introduced together with the first provider that needs it — - // so identity rides on the one ADK context rather than an auth-specific key. + // provider, which keys on the user) recovers it from ctx via + // [agent.IdentityFromContext] — so identity rides on the one ADK context + // rather than an auth-specific key. // // When interactive (3-legged) consent is required and cannot be completed // non-interactively, Credential returns a *ConsentRequiredError carrying the diff --git a/internal/adkcontext/adkcontext.go b/internal/adkcontext/adkcontext.go new file mode 100644 index 000000000..71a5cb068 --- /dev/null +++ b/internal/adkcontext/adkcontext.go @@ -0,0 +1,26 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package adkcontext holds the private context key under which an ADK context +// registers its invocation identity, so it can be recovered from a derived +// context.Context via agent.IdentityFromContext. It is a tiny leaf package shared +// by the agent and internal/context packages to avoid an import cycle. +package adkcontext + +type ctxKey int + +// IdentityKey is the context value key for the agent.Identity of an ADK context. +// It lives in an internal package with an unexported type, so no code outside +// the module can name or forge it. +const IdentityKey ctxKey = 0 diff --git a/internal/context/from_context_test.go b/internal/context/from_context_test.go new file mode 100644 index 000000000..c8f1ffee8 --- /dev/null +++ b/internal/context/from_context_test.go @@ -0,0 +1,95 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package context_test + +import ( + "context" + "testing" + + "google.golang.org/adk/v2/agent" + icontext "google.golang.org/adk/v2/internal/context" + "google.golang.org/adk/v2/session" +) + +type wrapKey struct{} + +// TestIdentityFromContextRecoversIdentity verifies that agent.IdentityFromContext +// recovers the ADK identity from a context that has been wrapped by non-ADK +// intermediaries (as jsonrpc2 / net/http do), across the base invocation context, +// a promoted common context, and a tool context. +func TestIdentityFromContextRecoversIdentity(t *testing.T) { + svc := session.InMemoryService() + resp, err := svc.Create(t.Context(), &session.CreateRequest{AppName: "app-1", UserID: "user-42"}) + if err != nil { + t.Fatalf("session Create() error = %v", err) + } + sessionID := resp.Session.ID() + ic := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{Session: resp.Session}) + + cases := []struct { + name string + ctx context.Context + }{ + {"invocation context", ic}, + {"promoted common context", agent.Promote(ic)}, + {"tool context", agent.NewToolContext(ic, "fc-1", nil, nil)}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + // Wrap in non-ADK children so a plain type-assert is erased but the + // Value lookup still resolves up the chain. + wrapped := context.WithValue(tc.ctx, wrapKey{}, "x") + wrapped, cancel := context.WithCancel(wrapped) + defer cancel() + + id, ok := agent.IdentityFromContext(wrapped) + if !ok { + t.Fatal("IdentityFromContext() ok = false, want true") + } + want := agent.Identity{UserID: "user-42", AppName: "app-1", SessionID: sessionID} + if id != want { + t.Errorf("IdentityFromContext() = %+v, want %+v", id, want) + } + }) + } +} + +func TestIdentityFromContextAbsent(t *testing.T) { + if _, ok := agent.IdentityFromContext(t.Context()); ok { + t.Error("IdentityFromContext() ok = true for a plain context, want false") + } +} + +// TestIdentityFromContextNoSession pins that an ADK context with no session +// yields (zero, false) instead of panicking — Value is a context.Context method +// and must never panic — across both the invocation and promoted common context +// Value paths. +func TestIdentityFromContextNoSession(t *testing.T) { + ic := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{}) // Session nil + cases := []struct { + name string + ctx context.Context + }{ + {"invocation context", ic}, + {"promoted common context", agent.Promote(ic)}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if id, ok := agent.IdentityFromContext(tc.ctx); ok { + t.Errorf("IdentityFromContext() = %+v, ok = true; want zero, false", id) + } + }) + } +} diff --git a/internal/context/invocation_context.go b/internal/context/invocation_context.go index c31d1e30d..df219b227 100644 --- a/internal/context/invocation_context.go +++ b/internal/context/invocation_context.go @@ -20,6 +20,7 @@ import ( "google.golang.org/genai" "google.golang.org/adk/v2/agent" + "google.golang.org/adk/v2/internal/adkcontext" "google.golang.org/adk/v2/platform" "google.golang.org/adk/v2/session" ) @@ -121,6 +122,21 @@ func (c *InvocationContext) WithContext(ctx context.Context) agent.InvocationCon return &newCtx } +// Value implements context.Context. It returns this invocation's [agent.Identity] +// for the ADK self key (so agent.IdentityFromContext can recover it); every other +// key, and a context with no session, delegates to the embedded context — +// preserving existing behavior and never panicking. +func (c *InvocationContext) Value(key any) any { + if key == adkcontext.IdentityKey && c.params.Session != nil { + return agent.Identity{ + UserID: c.params.Session.UserID(), + AppName: c.params.Session.AppName(), + SessionID: c.params.Session.ID(), + } + } + return c.Context.Value(key) +} + // ResumedInput always returns (nil, false) for the base // invocation context. Implementations that carry a resume payload // override this method.