From eadd5a6b1b4900e0bede2fee7f4cbd85dca71f56 Mon Sep 17 00:00:00 2001 From: wolo Date: Thu, 9 Jul 2026 08:15:58 +0000 Subject: [PATCH 1/2] feat(auth/gcp): add GCP credential provider Add gcp.NewProvider, an auth.CredentialProvider that resolves per-user credentials from the Agent Identity / IAM Connector services (via the REST client) and maps them to an auth.Credential. It takes the acting user from the ADK context at resolve time, so it runs inside an agent invocation and needs no per-user configuration. To recover the user from an http.RoundTripper that only sees a context.Context (deep beneath a tool call, past jsonrpc2/net/http wrapping), this also adds agent.FromContext(ctx) (ReadonlyContext, bool): ADK contexts register a read-only view of themselves under a private key, and FromContext returns it. Identity still lives on the typed context via Session(); nothing is stored under a key. Additive and non-breaking: the only new public API in agent is FromContext; the Value overrides are on unexported/internal context types and change behavior only for the new private key. No new module dependencies (the GCP client is hand-rolled over net/http + ADC). --- agent/common_context.go | 63 +++++++++++++ auth/gcp/provider.go | 116 ++++++++++++++++++++++++ auth/gcp/provider_test.go | 100 ++++++++++++++++++++ internal/adkcontext/adkcontext.go | 25 +++++ internal/context/from_context_test.go | 121 +++++++++++++++++++++++++ internal/context/invocation_context.go | 11 +++ 6 files changed, 436 insertions(+) create mode 100644 auth/gcp/provider.go create mode 100644 auth/gcp/provider_test.go create mode 100644 internal/adkcontext/adkcontext.go create mode 100644 internal/context/from_context_test.go diff --git a/agent/common_context.go b/agent/common_context.go index a78268178..d0b8bec5c 100644 --- a/agent/common_context.go +++ b/agent/common_context.go @@ -16,6 +16,7 @@ package agent import ( "context" + "errors" "fmt" "iter" "time" @@ -23,12 +24,63 @@ 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" ) +// FromContext returns the ADK [ReadonlyContext] carried by ctx, if present. +// +// ADK contexts embed context.Context and register a read-only view of themselves +// 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 invocation's identity (UserID, AppName, +// SessionID) without threading a typed context through every layer. +// +// It returns (nil, false) for a context that does not descend from an ADK +// context (a non-agent caller). +func FromContext(ctx context.Context) (ReadonlyContext, bool) { + rc, ok := ctx.Value(adkcontext.SelfKey).(ReadonlyContext) + return rc, ok +} + +// RequireContext is like [FromContext] but returns an error instead of a boolean +// when ctx does not descend from an ADK context. It is a convenience for callers +// (for example credential providers) that require the invocation identity. +func RequireContext(ctx context.Context) (ReadonlyContext, error) { + rc, ok := FromContext(ctx) + if !ok { + return nil, errors.New("agent: context does not belong to an ADK invocation") + } + return rc, nil +} + +// readonlyView exposes only the [ReadonlyContext] surface of an ADK context, so +// a context recovered via [FromContext] cannot be widened back to a mutable +// Context/InvocationContext. The wrapped context is held in an unexported field +// (not embedded), so it is reachable neither by a type assertion nor by ordinary +// reflection — reflect.Value.Interface panics on an unexported field. +type readonlyView struct { + rc ReadonlyContext +} + +func (v readonlyView) Deadline() (time.Time, bool) { return v.rc.Deadline() } +func (v readonlyView) Done() <-chan struct{} { return v.rc.Done() } +func (v readonlyView) Err() error { return v.rc.Err() } +func (v readonlyView) Value(key any) any { return v.rc.Value(key) } +func (v readonlyView) UserContent() *genai.Content { return v.rc.UserContent() } +func (v readonlyView) InvocationID() string { return v.rc.InvocationID() } +func (v readonlyView) AgentName() string { return v.rc.AgentName() } +func (v readonlyView) ReadonlyState() session.ReadonlyState { return v.rc.ReadonlyState() } +func (v readonlyView) UserID() string { return v.rc.UserID() } +func (v readonlyView) AppName() string { return v.rc.AppName() } +func (v readonlyView) SessionID() string { return v.rc.SessionID() } +func (v readonlyView) Branch() string { return v.rc.Branch() } + +var _ ReadonlyContext = readonlyView{} + // 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 +396,17 @@ func (c *commonContext) UserID() string { return c.invocationContext.Session().UserID() } +// Value implements context.Context. For the ADK self key it returns a read-only +// view of this context (so [FromContext] can recover the identity without the +// result being widened back to a mutable context); every other key delegates to +// the embedded context, preserving existing behavior. +func (c *commonContext) Value(key any) any { + if key == adkcontext.SelfKey { + return readonlyView{rc: c} + } + 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..009119072 --- /dev/null +++ b/auth/gcp/provider.go @@ -0,0 +1,116 @@ +// 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" + "fmt" + "slices" + "sync" + + "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.FromContext]) 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 reads Scopes on every + // request, so it must not alias a slice the caller can mutate later. + scheme.Scopes = slices.Clone(scheme.Scopes) + return &provider{scheme: scheme, client: cfg.Client}, nil +} + +type provider struct { + scheme Scheme + + mu sync.Mutex + client *Client +} + +var _ auth.CredentialProvider = (*provider)(nil) + +// Credential implements [auth.CredentialProvider]. +func (p *provider) Credential(ctx context.Context) (auth.Credential, error) { + rc, err := agent.RequireContext(ctx) + if err != nil { + return nil, fmt.Errorf("gcp: %w", err) + } + userID := rc.UserID() + if userID == "" { + return nil, errors.New("gcp: ADK context has no user id") + } + + client, err := p.resolveClient(ctx) + if err != nil { + return nil, err + } + return client.RetrieveCredential(ctx, Request{ + Resource: p.scheme.Name, + UserID: userID, + Scopes: p.scheme.Scopes, + ContinueURI: p.scheme.ContinueURI, + }) +} + +// resolveClient returns the configured client, creating a default one (backed by +// Application Default Credentials) on first use. The client's lifetime is +// detached from this one request with [context.WithoutCancel] (it is cached and +// reused) while keeping ctx's values. +func (p *provider) resolveClient(ctx context.Context) (*Client, error) { + p.mu.Lock() + defer p.mu.Unlock() + if p.client == nil { + c, err := NewClient(context.WithoutCancel(ctx), nil) + if err != nil { + return nil, err + } + p.client = c + } + return p.client, nil +} diff --git a/auth/gcp/provider_test.go b/auth/gcp/provider_test.go new file mode 100644 index 000000000..7b7d6f5fe --- /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.FromContext)", 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.FromContext) +// 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/internal/adkcontext/adkcontext.go b/internal/adkcontext/adkcontext.go new file mode 100644 index 000000000..3d2873c61 --- /dev/null +++ b/internal/adkcontext/adkcontext.go @@ -0,0 +1,25 @@ +// 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 a read-only view of itself, so it can be recovered from a derived +// context.Context via agent.FromContext. It is a tiny leaf package shared by the +// agent and internal/context packages to avoid an import cycle. +package adkcontext + +type ctxKey int + +// SelfKey is the context value key for the live agent.ReadonlyContext view of an +// ADK context. Its unexported type prevents other packages from forging it. +const SelfKey ctxKey = 0 diff --git a/internal/context/from_context_test.go b/internal/context/from_context_test.go new file mode 100644 index 000000000..5d1d5dc52 --- /dev/null +++ b/internal/context/from_context_test.go @@ -0,0 +1,121 @@ +// 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" + "reflect" + "testing" + + "google.golang.org/adk/v2/agent" + icontext "google.golang.org/adk/v2/internal/context" + "google.golang.org/adk/v2/session" +) + +type wrapKey struct{} + +// TestFromContextRecoversIdentity verifies that agent.FromContext 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 TestFromContextRecoversIdentity(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) + } + 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() + + if _, ok := wrapped.(agent.ReadonlyContext); ok { + t.Fatal("wrapped context unexpectedly type-asserts to ReadonlyContext") + } + + rc, ok := agent.FromContext(wrapped) + if !ok { + t.Fatal("FromContext() ok = false, want true") + } + if got := rc.UserID(); got != "user-42" { + t.Errorf("UserID() = %q, want %q", got, "user-42") + } + if got := rc.AppName(); got != "app-1" { + t.Errorf("AppName() = %q, want %q", got, "app-1") + } + // The recovered context must stay read-only: it must not widen back + // to a mutable Context/InvocationContext via a type assertion. + if _, ok := rc.(agent.Context); ok { + t.Error("recovered context widened to agent.Context, want read-only only") + } + if _, ok := rc.(agent.InvocationContext); ok { + t.Error("recovered context widened to agent.InvocationContext, want read-only only") + } + }) + } +} + +func TestFromContextAbsent(t *testing.T) { + if _, ok := agent.FromContext(t.Context()); ok { + t.Error("FromContext() ok = true for a plain context, want false") + } +} + +// TestFromContextNoReflectionWiden pins that the recovered read-only view does +// not expose the underlying mutable context through any exported struct field, +// i.e. it cannot be widened back to an agent.InvocationContext by ordinary +// reflection (reflect.Value.Interface panics on unexported fields). +func TestFromContextNoReflectionWiden(t *testing.T) { + svc := session.InMemoryService() + resp, err := svc.Create(t.Context(), &session.CreateRequest{AppName: "app", UserID: "u"}) + if err != nil { + t.Fatalf("session Create() error = %v", err) + } + ic := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{Session: resp.Session}) + rc, ok := agent.FromContext(agent.Promote(ic)) + if !ok { + t.Fatal("FromContext() ok = false, want true") + } + + v := reflect.ValueOf(rc) + if v.Kind() == reflect.Pointer { + v = v.Elem() + } + if v.Kind() != reflect.Struct { + return + } + for i := range v.NumField() { + f := v.Field(i) + if !f.CanInterface() { + continue // unexported: reflection cannot extract it + } + if _, ok := f.Interface().(agent.InvocationContext); ok { + t.Errorf("field %q exposes a widenable agent.InvocationContext", v.Type().Field(i).Name) + } + } +} diff --git a/internal/context/invocation_context.go b/internal/context/invocation_context.go index c31d1e30d..f8510368a 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,16 @@ func (c *InvocationContext) WithContext(ctx context.Context) agent.InvocationCon return &newCtx } +// Value implements context.Context. It returns a read-only view of this context +// for the ADK self key (so agent.FromContext can recover it); every other key +// delegates to the embedded context, preserving existing behavior. +func (c *InvocationContext) Value(key any) any { + if key == adkcontext.SelfKey { + return NewReadonlyContext(c) + } + return c.Context.Value(key) +} + // ResumedInput always returns (nil, false) for the base // invocation context. Implementations that carry a resume payload // override this method. From 9186dc4d02a4ff41c19a73541dea8a06767b7c13 Mon Sep 17 00:00:00 2001 From: wolo Date: Sun, 19 Jul 2026 22:22:17 +0000 Subject: [PATCH 2/2] refactor(agent): recover invocation Identity, not a wrapped ReadonlyContext FromContext/RequireContext returned a full agent.ReadonlyContext just so the GCP credential provider could read UserID. That forced a readonlyView laundering type (13 method forwarders) plus anti-widening + reflection tests, all existing only to stop the recovered context from being widened back to a mutable Context/InvocationContext. Return a small immutable agent.Identity{UserID,AppName,SessionID} via IdentityFromContext/RequireIdentity instead. Returning a value erases the widening concern entirely: readonlyView and the widening tests are gone. Rename the context key SelfKey -> IdentityKey to match what it now carries. All of this API is new in this PR (absent from main), so nothing released changes; ReadonlyContext/Context/InvocationContext are untouched. --- agent/common_context.go | 78 ++++++++-------------- auth/gcp/provider.go | 37 +++++------ auth/gcp/provider_test.go | 4 +- auth/providers.go | 6 +- internal/adkcontext/adkcontext.go | 13 ++-- internal/context/from_context_test.go | 90 +++++++++----------------- internal/context/invocation_context.go | 15 +++-- 7 files changed, 100 insertions(+), 143 deletions(-) diff --git a/agent/common_context.go b/agent/common_context.go index d0b8bec5c..f3af44c3e 100644 --- a/agent/common_context.go +++ b/agent/common_context.go @@ -16,7 +16,6 @@ package agent import ( "context" - "errors" "fmt" "iter" "time" @@ -31,56 +30,31 @@ import ( "google.golang.org/adk/v2/tool/toolconfirmation" ) -// FromContext returns the ADK [ReadonlyContext] carried by ctx, if present. +// 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 a read-only view of themselves -// under a private key, so code that only holds a context.Context — for example an +// 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 invocation's identity (UserID, AppName, -// SessionID) without threading a typed context through every layer. +// wrap the context — can recover the acting identity without threading a typed +// context through every layer. // -// It returns (nil, false) for a context that does not descend from an ADK +// It returns (zero, false) for a context that does not descend from an ADK // context (a non-agent caller). -func FromContext(ctx context.Context) (ReadonlyContext, bool) { - rc, ok := ctx.Value(adkcontext.SelfKey).(ReadonlyContext) - return rc, ok -} - -// RequireContext is like [FromContext] but returns an error instead of a boolean -// when ctx does not descend from an ADK context. It is a convenience for callers -// (for example credential providers) that require the invocation identity. -func RequireContext(ctx context.Context) (ReadonlyContext, error) { - rc, ok := FromContext(ctx) - if !ok { - return nil, errors.New("agent: context does not belong to an ADK invocation") - } - return rc, nil -} - -// readonlyView exposes only the [ReadonlyContext] surface of an ADK context, so -// a context recovered via [FromContext] cannot be widened back to a mutable -// Context/InvocationContext. The wrapped context is held in an unexported field -// (not embedded), so it is reachable neither by a type assertion nor by ordinary -// reflection — reflect.Value.Interface panics on an unexported field. -type readonlyView struct { - rc ReadonlyContext +func IdentityFromContext(ctx context.Context) (Identity, bool) { + id, ok := ctx.Value(adkcontext.IdentityKey).(Identity) + return id, ok } -func (v readonlyView) Deadline() (time.Time, bool) { return v.rc.Deadline() } -func (v readonlyView) Done() <-chan struct{} { return v.rc.Done() } -func (v readonlyView) Err() error { return v.rc.Err() } -func (v readonlyView) Value(key any) any { return v.rc.Value(key) } -func (v readonlyView) UserContent() *genai.Content { return v.rc.UserContent() } -func (v readonlyView) InvocationID() string { return v.rc.InvocationID() } -func (v readonlyView) AgentName() string { return v.rc.AgentName() } -func (v readonlyView) ReadonlyState() session.ReadonlyState { return v.rc.ReadonlyState() } -func (v readonlyView) UserID() string { return v.rc.UserID() } -func (v readonlyView) AppName() string { return v.rc.AppName() } -func (v readonlyView) SessionID() string { return v.rc.SessionID() } -func (v readonlyView) Branch() string { return v.rc.Branch() } - -var _ ReadonlyContext = readonlyView{} - // In general CommonContext should not be wrapped with contexts not providing agent.Context. // It allows to copy&modify context instead of building chains. @@ -396,13 +370,15 @@ func (c *commonContext) UserID() string { return c.invocationContext.Session().UserID() } -// Value implements context.Context. For the ADK self key it returns a read-only -// view of this context (so [FromContext] can recover the identity without the -// result being widened back to a mutable context); every other key delegates to -// the embedded context, preserving existing behavior. +// 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.SelfKey { - return readonlyView{rc: c} + 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) } diff --git a/auth/gcp/provider.go b/auth/gcp/provider.go index 009119072..632c5056b 100644 --- a/auth/gcp/provider.go +++ b/auth/gcp/provider.go @@ -17,9 +17,9 @@ package gcp import ( "context" "errors" - "fmt" "slices" "sync" + "time" "google.golang.org/adk/v2/agent" "google.golang.org/adk/v2/auth" @@ -50,9 +50,9 @@ type ProviderConfig struct { // 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.FromContext]) at resolve -// time, so the provider must run within an agent invocation (e.g. wired into -// mcptoolset or remoteagent). +// 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") @@ -60,8 +60,7 @@ func NewProvider(scheme Scheme, cfg *ProviderConfig) (auth.CredentialProvider, e if cfg == nil { cfg = &ProviderConfig{} } - // Defensive copy: the provider outlives this call and reads Scopes on every - // request, so it must not alias a slice the caller can mutate later. + // 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 } @@ -77,13 +76,9 @@ var _ auth.CredentialProvider = (*provider)(nil) // Credential implements [auth.CredentialProvider]. func (p *provider) Credential(ctx context.Context) (auth.Credential, error) { - rc, err := agent.RequireContext(ctx) - if err != nil { - return nil, fmt.Errorf("gcp: %w", err) - } - userID := rc.UserID() - if userID == "" { - return nil, errors.New("gcp: ADK context has no user id") + 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) @@ -92,21 +87,27 @@ func (p *provider) Credential(ctx context.Context) (auth.Credential, error) { } return client.RetrieveCredential(ctx, Request{ Resource: p.scheme.Name, - UserID: userID, + UserID: id.UserID, Scopes: p.scheme.Scopes, ContinueURI: p.scheme.ContinueURI, }) } +// clientInitTimeout bounds the lazy ADC lookup so a hung probe fails instead of +// wedging every caller on the mutex. +const clientInitTimeout = 30 * time.Second + // resolveClient returns the configured client, creating a default one (backed by -// Application Default Credentials) on first use. The client's lifetime is -// detached from this one request with [context.WithoutCancel] (it is cached and -// reused) while keeping ctx's values. +// Application Default Credentials) on first use. Construction runs under +// [context.WithoutCancel] so the cached, reused client isn't bound to one +// request's cancellation, while keeping ctx's values. func (p *provider) resolveClient(ctx context.Context) (*Client, error) { p.mu.Lock() defer p.mu.Unlock() if p.client == nil { - c, err := NewClient(context.WithoutCancel(ctx), nil) + dctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), clientInitTimeout) + defer cancel() + c, err := NewClient(dctx, nil) if err != nil { return nil, err } diff --git a/auth/gcp/provider_test.go b/auth/gcp/provider_test.go index 7b7d6f5fe..290e10a9d 100644 --- a/auth/gcp/provider_test.go +++ b/auth/gcp/provider_test.go @@ -63,7 +63,7 @@ func TestProviderCredential(t *testing.T) { 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.FromContext)", gotUserID, "user-1") + t.Errorf("service saw userId = %q, want %q (identity from agent.IdentityFromContext)", gotUserID, "user-1") } } @@ -87,7 +87,7 @@ func TestNewProviderValidatesScheme(t *testing.T) { } } -// adkContext returns an ADK invocation context (recoverable via agent.FromContext) +// 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() 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 index 3d2873c61..71a5cb068 100644 --- a/internal/adkcontext/adkcontext.go +++ b/internal/adkcontext/adkcontext.go @@ -13,13 +13,14 @@ // limitations under the License. // Package adkcontext holds the private context key under which an ADK context -// registers a read-only view of itself, so it can be recovered from a derived -// context.Context via agent.FromContext. It is a tiny leaf package shared by the -// agent and internal/context packages to avoid an import cycle. +// 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 -// SelfKey is the context value key for the live agent.ReadonlyContext view of an -// ADK context. Its unexported type prevents other packages from forging it. -const SelfKey ctxKey = 0 +// 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 index 5d1d5dc52..c8f1ffee8 100644 --- a/internal/context/from_context_test.go +++ b/internal/context/from_context_test.go @@ -16,7 +16,6 @@ package context_test import ( "context" - "reflect" "testing" "google.golang.org/adk/v2/agent" @@ -26,16 +25,17 @@ import ( type wrapKey struct{} -// TestFromContextRecoversIdentity verifies that agent.FromContext 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 TestFromContextRecoversIdentity(t *testing.T) { +// 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 { @@ -54,68 +54,42 @@ func TestFromContextRecoversIdentity(t *testing.T) { wrapped, cancel := context.WithCancel(wrapped) defer cancel() - if _, ok := wrapped.(agent.ReadonlyContext); ok { - t.Fatal("wrapped context unexpectedly type-asserts to ReadonlyContext") - } - - rc, ok := agent.FromContext(wrapped) + id, ok := agent.IdentityFromContext(wrapped) if !ok { - t.Fatal("FromContext() ok = false, want true") - } - if got := rc.UserID(); got != "user-42" { - t.Errorf("UserID() = %q, want %q", got, "user-42") + t.Fatal("IdentityFromContext() ok = false, want true") } - if got := rc.AppName(); got != "app-1" { - t.Errorf("AppName() = %q, want %q", got, "app-1") - } - // The recovered context must stay read-only: it must not widen back - // to a mutable Context/InvocationContext via a type assertion. - if _, ok := rc.(agent.Context); ok { - t.Error("recovered context widened to agent.Context, want read-only only") - } - if _, ok := rc.(agent.InvocationContext); ok { - t.Error("recovered context widened to agent.InvocationContext, want read-only only") + want := agent.Identity{UserID: "user-42", AppName: "app-1", SessionID: sessionID} + if id != want { + t.Errorf("IdentityFromContext() = %+v, want %+v", id, want) } }) } } -func TestFromContextAbsent(t *testing.T) { - if _, ok := agent.FromContext(t.Context()); ok { - t.Error("FromContext() ok = true for a plain context, want false") +func TestIdentityFromContextAbsent(t *testing.T) { + if _, ok := agent.IdentityFromContext(t.Context()); ok { + t.Error("IdentityFromContext() ok = true for a plain context, want false") } } -// TestFromContextNoReflectionWiden pins that the recovered read-only view does -// not expose the underlying mutable context through any exported struct field, -// i.e. it cannot be widened back to an agent.InvocationContext by ordinary -// reflection (reflect.Value.Interface panics on unexported fields). -func TestFromContextNoReflectionWiden(t *testing.T) { - svc := session.InMemoryService() - resp, err := svc.Create(t.Context(), &session.CreateRequest{AppName: "app", UserID: "u"}) - if err != nil { - t.Fatalf("session Create() error = %v", err) - } - ic := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{Session: resp.Session}) - rc, ok := agent.FromContext(agent.Promote(ic)) - if !ok { - t.Fatal("FromContext() ok = false, want true") - } - - v := reflect.ValueOf(rc) - if v.Kind() == reflect.Pointer { - v = v.Elem() - } - if v.Kind() != reflect.Struct { - return +// 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 i := range v.NumField() { - f := v.Field(i) - if !f.CanInterface() { - continue // unexported: reflection cannot extract it - } - if _, ok := f.Interface().(agent.InvocationContext); ok { - t.Errorf("field %q exposes a widenable agent.InvocationContext", v.Type().Field(i).Name) - } + 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 f8510368a..df219b227 100644 --- a/internal/context/invocation_context.go +++ b/internal/context/invocation_context.go @@ -122,12 +122,17 @@ func (c *InvocationContext) WithContext(ctx context.Context) agent.InvocationCon return &newCtx } -// Value implements context.Context. It returns a read-only view of this context -// for the ADK self key (so agent.FromContext can recover it); every other key -// delegates to the embedded context, preserving existing behavior. +// 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.SelfKey { - return NewReadonlyContext(c) + 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) }