diff --git a/agent/agent.go b/agent/agent.go index c523e8755..a7c83e224 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -24,6 +24,7 @@ import ( "google.golang.org/genai" "google.golang.org/adk/v2/artifact" + "google.golang.org/adk/v2/internal/adkcontext" agentinternal "google.golang.org/adk/v2/internal/agent" "google.golang.org/adk/v2/internal/plugininternal/plugincontext" "google.golang.org/adk/v2/internal/telemetry" @@ -412,6 +413,23 @@ func (c *invocationContext) Session() session.Session { return c.session } +// Value implements context.Context, answering the ADK identity key like every +// other invocation context so a promoted copy and this one cannot disagree. It +// owns its session, so no session means no identity — never the enclosing +// invocation's, whose user made no such call. +func (c *invocationContext) Value(key any) any { + if key == adkcontext.IdentityKey { + if id, ok := identityOf(func() session.Session { return c.session }); ok { + return id + } + return nil + } + if c.Context == nil { + return nil + } + return c.Context.Value(key) +} + func (c *invocationContext) InvocationID() string { return c.invocationID } diff --git a/agent/common_context.go b/agent/common_context.go index 7bd19e3c9..316c58843 100644 --- a/agent/common_context.go +++ b/agent/common_context.go @@ -23,12 +23,67 @@ 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 an ADK invocation's identity: the acting user, app name, and +// session a call belongs to. It is recovered from a plain context.Context via +// [IdentityFromContext], which reads it off the live session each time, so a +// session mutated mid-invocation is reflected in the next lookup. +type Identity struct { + // UserID is the acting end user, as the embedding server put it on the + // session. ADK does not authenticate it: anything acting on behalf of this + // user — minting a per-user credential, for instance — is trusting the server + // to have bound session.UserID to an authenticated principal. ADK's own REST + // server takes it from the request body. + UserID string + // AppName is the app the invocation belongs to. + AppName string + // SessionID identifies the conversation the invocation belongs to. + SessionID string +} + +// identityOf reads an invocation identity from the session getSession returns. +// It is the one place package agent turns a session into an [Identity], so every +// context type here answers the identity key the same way. +// +// getSession is called inside the recover, not before it: Session() is itself a +// method on caller-supplied code and can panic on its own. +func identityOf(getSession func() session.Session) (Identity, bool) { + return adkcontext.Recovered(func() Identity { + // One Session value, then one call per field: re-reading Session() per + // field risks a torn identity, and some context wrappers log on every read. + s := getSession() + return Identity{UserID: s.UserID(), AppName: s.AppName(), SessionID: s.ID()} + }) +} + +// 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) both for a context that does not descend from an ADK +// context and for an invocation with no readable session; the two are not +// distinguishable here. The invocation types in this module never report an +// enclosing invocation's user in place of their own — one implemented elsewhere +// cannot make that promise, since the key is unnameable outside the module and +// its embedded parent answers instead. ok does not imply a populated Identity +// either: an invocation whose session carries no user yields an empty UserID, so +// a caller that needs one must check. +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 +399,79 @@ func (c *commonContext) UserID() string { return c.invocationContext.Session().UserID() } +// Value implements context.Context. For the ADK identity key it returns the +// [Identity] of the invocation this context speaks for (so [IdentityFromContext] +// can recover it from a derived context); every other key delegates to the +// embedded context, preserving existing behavior. +// +// Only the identity key touches the invocation, so no other key is affected by +// its state, and a session that panics costs the identity, not the process. +func (c *commonContext) Value(key any) any { + if key == adkcontext.IdentityKey { + return c.identity() + } + if c.Context == nil { + return nil + } + return c.Context.Value(key) +} + +// identity answers the ADK identity key, as an any so it can report "none". +// +// A commonContext owns no session, so it speaks for the invocation it wraps, and +// reads that invocation's own session first. Asking the invocation's Value first +// looks equivalent and is not: an InvocationContext written outside this module +// embeds the context it was derived from, to inherit cancellation, and cannot +// override a key it cannot name — so its Value answers with the *enclosing* +// invocation's identity even when it has a session of its own naming a different +// user. Reading the session first is what makes an invocation report itself. +// +// Only an invocation with no readable session of its own delegates, by asking the +// invocation it wraps. A tool or callback context is exactly that shape by design +// and hands the key to the context underneath. An invocation that owns a session +// field and has none — a nested invocation built without one — gets nil back from +// that ask and reports no identity, rather than inheriting a user who made no +// such call. +// +// A commonContext speaking for no invocation at all is the one case that consults +// its own parent, since there is nothing else it could answer for. +func (c *commonContext) identity() any { + if c.invocationContext == nil { + if c.Context == nil { + return nil + } + return c.Context.Value(adkcontext.IdentityKey) + } + // Method value and call both inside the recover: Session() is caller-supplied + // code and invocationContext can be a typed-nil pointer. + s, read := adkcontext.Recovered(func() session.Session { return c.invocationContext.Session() }) + if !read { + // The invocation cannot say who it is. It must not inherit an answer from + // what it embeds: a broken invocation is not an absent one, and the context + // it was derived from belongs to a different call. + return nil + } + if s != nil { + if id, ok := identityOf(func() session.Session { return s }); ok { + return id + } + // Present but unreadable is broken too, and delegating would inherit. + return nil + } + // No session of its own, which a tool or callback context is by design, so ask + // the invocation. Type-asserted, not merely checked against nil: one with a + // permissive Value that answers every key would otherwise hand back something + // that is not an Identity and be taken for one. + if v, ok := adkcontext.Recovered(func() any { + return c.invocationContext.Value(adkcontext.IdentityKey) + }); ok { + if id, isIdentity := v.(Identity); isIdentity { + return id + } + } + return nil +} + var ( _ Context = (*commonContext)(nil) _ InvocationContext = (*commonContext)(nil) diff --git a/agent/identity_matrix_test.go b/agent/identity_matrix_test.go new file mode 100644 index 000000000..2c2bf0ac4 --- /dev/null +++ b/agent/identity_matrix_test.go @@ -0,0 +1,214 @@ +// 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 agent + +import ( + "context" + "testing" + "time" + + "google.golang.org/adk/v2/session" +) + +// The identity key is answered by a decision procedure, so it is pinned by a +// table rather than by cases. Rows are the shapes a session or an invocation can +// legally take, columns the ways a context is derived from one. Reviewing this +// procedure a diff at a time found one failing shape per round over five rounds, +// each fix moving the failure to a neighbouring shape. +// +// The policy the table encodes: +// - An invocation reports the user of its OWN session, never one it inherited. +// - An invocation with no session of its own delegates, which is what a tool or +// callback context needs — theirs is nil by design. +// - An invocation that cannot answer at all reports nothing. Broken is not +// absent, and delegating would hand back a different call's user. +// - No shape panics out of Value, and no shape disturbs an unrelated key. + +type matrixSession struct { + session.Session + id, app, user string +} + +func (s *matrixSession) ID() string { return s.id } +func (s *matrixSession) AppName() string { return s.app } +func (s *matrixSession) UserID() string { return s.user } +func (s *matrixSession) State() session.State { return nil } +func (s *matrixSession) Events() session.Events { return nil } +func (s *matrixSession) LastUpdateTime() time.Time { return time.Time{} } + +func matrixOwner(user string) session.Session { + return &matrixSession{id: "sid-" + user, app: "app", user: user} +} + +// structSession is a value type, which a six-accessor interface invites and which +// a reflect-based nil check cannot inspect. +type structSession struct{ session.Session } + +func (structSession) ID() string { return "sid" } +func (structSession) AppName() string { return "app" } +func (structSession) UserID() string { return "owner" } + +// safeNilSession answers without touching its receiver, so a typed-nil one works. +type safeNilPtrSession struct{ session.Session } + +func (*safeNilPtrSession) ID() string { return "sid" } +func (*safeNilPtrSession) AppName() string { return "app" } +func (*safeNilPtrSession) UserID() string { return "owner" } + +// nilWrappingSession promotes its accessors from a nil embedded session, the +// shape llmagent.newWrappedSession produces for a nil original. +type nilWrappingSession struct{ session.Session } + +// panickingAccessorSession is broken in its own code, not in its nil-ness. +type panickingAccessorSession struct{ session.Session } + +func (panickingAccessorSession) ID() string { return "sid" } +func (panickingAccessorSession) AppName() string { return "app" } +func (panickingAccessorSession) UserID() string { panic("accessor is not available") } + +// panickingSessionInvocation declines to hand over a session at all, as the +// exported StrictContextMock does. +type panickingSessionInvocation struct{ InvocationContext } + +func (panickingSessionInvocation) Session() session.Session { panic("Session is not available") } + +// permissiveInvocationValue answers every key, as a decorator or double might. +type permissiveInvocationValue struct{ InvocationContext } + +func (permissiveInvocationValue) Value(any) any { return "something that is not an Identity" } + +// decoratedInvocationValue is the shape written outside this module: embed the +// invocation you were derived from to inherit cancellation, carry your own +// session. It cannot override the identity key, because it cannot name it. +type decoratedInvocationValue struct { + InvocationContext + own session.Session +} + +func (d decoratedInvocationValue) Session() session.Session { return d.own } + +func TestIdentityDecisionMatrix(t *testing.T) { + // Every row is nested under this one, so any cell that reports "enclosing" is + // serving a user who made no such call. + enclosing := &invocationContext{Context: t.Context(), session: matrixOwner("enclosing")} + + rows := []struct { + name string + ic func() InvocationContext + want string // "" means no identity + // outsideModule marks an invocation that cannot override the identity key, + // so asking it directly answers with whatever it embeds. That is the + // documented limit of the mechanism; the derivations are what fix it. + outsideModule bool + }{ + {name: "pointer session", want: "u", ic: func() InvocationContext { + return &invocationContext{Context: enclosing, session: matrixOwner("u")} + }}, + {name: "struct-value session", want: "owner", ic: func() InvocationContext { + return &invocationContext{Context: enclosing, session: structSession{}} + }}, + {name: "typed-nil with safe accessors", want: "owner", ic: func() InvocationContext { + return &invocationContext{Context: enclosing, session: (*safeNilPtrSession)(nil)} + }}, + {name: "no session", ic: func() InvocationContext { + return &invocationContext{Context: enclosing} + }}, + {name: "typed-nil session", ic: func() InvocationContext { + return &invocationContext{Context: enclosing, session: (*matrixSession)(nil)} + }}, + {name: "session wrapping a nil session", ic: func() InvocationContext { + return &invocationContext{Context: enclosing, session: &nilWrappingSession{}} + }}, + {name: "session accessor panics", ic: func() InvocationContext { + return &invocationContext{Context: enclosing, session: panickingAccessorSession{}} + }}, + {name: "Session() panics", outsideModule: true, ic: func() InvocationContext { + return panickingSessionInvocation{InvocationContext: enclosing} + }}, + {name: "permissive Value, own session", want: "u", outsideModule: true, ic: func() InvocationContext { + return permissiveInvocationValue{InvocationContext: &invocationContext{Context: enclosing, session: matrixOwner("u")}} + }}, + {name: "decorated outside the module", want: "u", outsideModule: true, ic: func() InvocationContext { + return decoratedInvocationValue{InvocationContext: enclosing, own: matrixOwner("u")} + }}, + // The two axes have to be crossed, not just walked. An invocation that owns + // a session fails closed on its own, so an unreadable session only reaches + // the delegation through a decorator — where inheriting is a live user's + // credential minted for someone else's call. + {name: "decorated, typed-nil session", outsideModule: true, ic: func() InvocationContext { + return decoratedInvocationValue{InvocationContext: enclosing, own: (*matrixSession)(nil)} + }}, + {name: "decorated, session accessor panics", outsideModule: true, ic: func() InvocationContext { + return decoratedInvocationValue{InvocationContext: enclosing, own: panickingAccessorSession{}} + }}, + } + + // Two columns deliberately put a value on the chain, so the probe for "an + // unrelated key is undisturbed" must use a key nothing injects. + type unrelatedKey struct{} + type probeKey struct{} + cols := []struct { + name string + of func(InvocationContext) context.Context + }{ + {"the invocation itself", func(ic InvocationContext) context.Context { return ic }}, + {"Promote", func(ic InvocationContext) context.Context { return Promote(ic) }}, + {"NewContext", func(ic InvocationContext) context.Context { return NewContext(ic) }}, + {"NewToolContext", func(ic InvocationContext) context.Context { return NewToolContext(ic, "fc", nil, nil) }}, + {"NewCallbackContext", func(ic InvocationContext) context.Context { return NewCallbackContext(ic, nil) }}, + {"NewCallbackContextWithArtifactTracking", func(ic InvocationContext) context.Context { + return NewCallbackContextWithArtifactTracking(ic, nil) + }}, + {"reparented onto a carrier of the enclosing invocation", func(ic InvocationContext) context.Context { + return Promote(ic).WithContext(context.WithValue(context.Context(enclosing), unrelatedKey{}, "x")) + }}, + {"tool context of a tool context", func(ic InvocationContext) context.Context { + return NewToolContext(NewToolContext(ic, "a", nil, nil), "b", nil, nil) + }}, + {"behind a non-ADK wrapper", func(ic InvocationContext) context.Context { + return context.WithValue(Promote(ic), unrelatedKey{}, "x") + }}, + } + + for _, r := range rows { + for _, c := range cols { + if c.name == "the invocation itself" && r.outsideModule { + continue + } + t.Run(r.name+" / "+c.name, func(t *testing.T) { + defer func() { + if p := recover(); p != nil { + t.Fatalf("Value panicked: %v", p) + } + }() + ctx := c.of(r.ic()) + var got string + if id, ok := IdentityFromContext(ctx); ok { + got = id.UserID + } + if got != r.want { + t.Errorf("IdentityFromContext() user = %q, want %q", got, r.want) + } + // An invocation that hijacks every key is answering for itself. + if _, hijacks := r.ic().(permissiveInvocationValue); hijacks { + return + } + if v := ctx.Value(probeKey{}); v != nil { + t.Errorf("Value(probeKey{}) = %v, want nil: only the identity key reads the session", v) + } + }) + } + } +} diff --git a/agent/identity_test.go b/agent/identity_test.go new file mode 100644 index 000000000..110fca0ed --- /dev/null +++ b/agent/identity_test.go @@ -0,0 +1,150 @@ +// 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 agent + +import ( + "context" + "testing" + + "google.golang.org/adk/v2/internal/adkcontext" + "google.golang.org/adk/v2/session" +) + +// TestInvocationContextIdentityIsOwned pins the guard on the third invocation +// implementation, the one in this package: it owns its session, so no session +// means no identity — never the enclosing invocation's, whose user made no such +// call and whose credential would otherwise be minted for it. +func TestInvocationContextIdentityIsOwned(t *testing.T) { + outer := &invocationContext{Context: t.Context(), session: &identityTestSession{}} + if got, ok := IdentityFromContext(outer); !ok || got.UserID != "alice" { + t.Fatalf("outer IdentityFromContext() = %+v, %v; want alice", got, ok) + } + + nested := &invocationContext{Context: outer} // no session of its own + if got, ok := IdentityFromContext(nested); ok { + t.Errorf("nested IdentityFromContext() = %+v, true; want no identity", got) + } + if got := nested.Value(wrapKey{}); got != nil { + t.Errorf("nested Value(wrapKey{}) = %v, want nil (unrelated keys still delegate)", got) + } + // A nil embedded parent must not panic either: Value is a context.Context + // method and runs inside an http.RoundTripper. + if got := (&invocationContext{}).Value(wrapKey{}); got != nil { + t.Errorf("Value(wrapKey{}) with no parent = %v, want nil", got) + } +} + +// TestCommonContextWithoutInvocation pins the guard for a commonContext that +// speaks for no invocation: it has nothing to answer the identity key with, so it +// delegates, and a nil parent on top of that must not panic — Value is a +// context.Context method and runs inside an http.RoundTripper. +func TestCommonContextWithoutInvocation(t *testing.T) { + owner := &invocationContext{Context: t.Context(), session: &identityTestSession{}} + c := &commonContext{Context: owner} // no invocationContext + if got, ok := IdentityFromContext(c); !ok || got.UserID != "alice" { + t.Errorf("IdentityFromContext() = %+v, %v; want the parent's identity", got, ok) + } + if got := (&commonContext{}).Value(adkcontext.IdentityKey); got != nil { + t.Errorf("Value(IdentityKey) with no invocation and no parent = %v, want nil", got) + } + if got := (&commonContext{}).Value(wrapKey{}); got != nil { + t.Errorf("Value(wrapKey{}) with no invocation and no parent = %v, want nil", got) + } +} + +// TestReadIdentityRecoversPanickingAccessor pins that a session accessor which +// panics costs the identity and not the process. +func TestReadIdentityRecoversPanickingAccessor(t *testing.T) { + c := &invocationContext{Context: t.Context(), session: panickingSession{}} + if got := c.Value(adkcontext.IdentityKey); got != nil { + t.Errorf("Value(IdentityKey) = %v, want nil for a session that panics", got) + } +} + +// identityTestSession answers the three identity accessors; panickingSession +// panics on the first one, the shape a broken third-party session takes. +type identityTestSession struct{ session.Session } + +func (identityTestSession) ID() string { return "sid-1" } +func (identityTestSession) AppName() string { return "app-1" } +func (identityTestSession) UserID() string { return "alice" } + +type panickingSession struct{ session.Session } + +func (panickingSession) UserID() string { panic("UserID is not available") } + +type wrapKey struct{} + +var _ context.Context = (*invocationContext)(nil) + +// TestIdentityFromPermissiveInvocation pins that an invocation answering every +// key with something that is not an [Identity] does not swallow the fallback: a +// decorator or test double that returns a placeholder for any key would +// otherwise cost the identity on every outbound request. +func TestIdentityFromPermissiveInvocation(t *testing.T) { + owner := &invocationContext{Context: t.Context(), session: &identityTestSession{}} + c := &commonContext{Context: t.Context(), invocationContext: permissiveInvocation{InvocationContext: owner}} + got, ok := IdentityFromContext(c) + if !ok || got.UserID != "alice" { + t.Errorf("IdentityFromContext() = %+v, %v; want the session read to be reached", got, ok) + } +} + +// TestIdentityFromDecoratedInvocation pins that an invocation reports its OWN +// user, not the one it inherited. An InvocationContext written outside this +// module embeds the context it was derived from, to inherit cancellation, and +// cannot override a key it cannot name — so its Value answers with the enclosing +// invocation's identity. Reading its session first is what stops one user's +// credential being minted for another's call. +func TestIdentityFromDecoratedInvocation(t *testing.T) { + enclosing := &invocationContext{Context: t.Context(), session: &identityTestSession{}} // alice + decorated := decoratedInvocation{ + InvocationContext: enclosing, + own: &otherUserSession{}, + } + for _, tc := range []struct { + name string + ctx context.Context + }{ + {"promoted", Promote(decorated)}, + {"tool context", NewToolContext(decorated, "fc-1", nil, nil)}, + {"callback context", NewCallbackContext(decorated, nil)}, + } { + id, ok := IdentityFromContext(tc.ctx) + if !ok || id.UserID != "bob" { + t.Errorf("%s IdentityFromContext() = %+v, %v; want bob, the decorated invocation's own user", tc.name, id, ok) + } + } +} + +// decoratedInvocation is how an invocation is wrapped outside this module: embed +// the enclosing one, override the accessors that differ. +type decoratedInvocation struct { + InvocationContext + own session.Session +} + +func (d decoratedInvocation) Session() session.Session { return d.own } + +type otherUserSession struct{ session.Session } + +func (otherUserSession) ID() string { return "sid-2" } +func (otherUserSession) AppName() string { return "app-1" } +func (otherUserSession) UserID() string { return "bob" } + +// permissiveInvocation answers every key, as a decorator or a test double might. +type permissiveInvocation struct{ InvocationContext } + +func (permissiveInvocation) Value(any) any { return "something that is not an Identity" } diff --git a/auth/gcp/client.go b/auth/gcp/client.go index 1851bb6f4..07aaf0359 100644 --- a/auth/gcp/client.go +++ b/auth/gcp/client.go @@ -49,11 +49,37 @@ const ( // routed to the Agent Identity service (same split as adk-python). var connectorResourceRE = regexp.MustCompile(`^projects/[^/]+/locations/[^/]+/connectors/[^/]+$`) +// authProviderResourceRE matches an Agent Identity resource name. Together with +// connectorResourceRE it is the full set [NewProvider] accepts; the client +// itself is looser, routing any non-connector name to Agent Identity. +var authProviderResourceRE = regexp.MustCompile(`^projects/[^/]+/locations/[^/]+/authProviders/[^/]+$`) + // resourceNameRE bounds a resource name to the characters GCP resource names -// use, so it can't inject extra path segments, a query, or a fragment into the -// request URL it is interpolated into. A separate ".." check blocks path -// traversal (dots are allowed so domain-style ids still pass). -var resourceNameRE = regexp.MustCompile(`^[A-Za-z0-9._~/-]+$`) +// use. It cannot inject a query, a fragment, an authority or a percent-escape +// into the request URL the name is interpolated into; extra path segments are +// allowed, since a resource name is itself a path. The colon is allowed for +// domain-scoped project ids (projects/example.com:my-project/...) — the name is +// always appended after the endpoint and a /v1 segment, so it can never be read +// as a scheme. +var resourceNameRE = regexp.MustCompile(`^[A-Za-z0-9._~:/-]+$`) + +// validateResource rejects a resource name that cannot be safely interpolated +// into a request URL, or that would not survive path normalization — an empty, +// "." or ".." segment blocks traversal, and also keeps the name the caller +// validated identical to the one connectorResourceRE routes on. [NewProvider] +// applies it at wiring time too, so a malformed name fails once rather than on +// every request. +func validateResource(name string) error { + if !resourceNameRE.MatchString(name) { + return fmt.Errorf("resource %q has invalid characters", name) + } + for seg := range strings.SplitSeq(name, "/") { + if seg == "" || seg == "." || seg == ".." { + return fmt.Errorf("resource %q has an empty or relative path segment", name) + } + } + return nil +} // Sentinel errors from [Client.RetrieveCredential]; callers test with errors.Is. var ( @@ -192,8 +218,8 @@ func (c *Client) RetrieveCredential(ctx context.Context, req Request) (auth.Cred if req.UserID == "" { return nil, errors.New("gcp: RetrieveCredential requires a UserID") } - if !resourceNameRE.MatchString(req.Resource) || strings.Contains(req.Resource, "..") { - return nil, fmt.Errorf("gcp: RetrieveCredential resource %q has invalid characters", req.Resource) + if err := validateResource(req.Resource); err != nil { + return nil, fmt.Errorf("gcp: RetrieveCredential: %w", err) } retrieve := c.retrieveAgentIdentity @@ -214,11 +240,11 @@ func (c *Client) RetrieveCredential(ctx context.Context, req Request) (auth.Cred case consentOutcome: return nil, &auth.ConsentRequiredError{AuthURI: o.authURI, Nonce: o.nonce} case rejectedOutcome: - return nil, fmt.Errorf("%w for %q", ErrConsentRejected, req.Resource) + return nil, ErrConsentRejected case pendingOutcome: remaining := time.Until(deadline) if remaining <= 0 { - return nil, fmt.Errorf("%w for %q", ErrPollTimeout, req.Resource) + return nil, ErrPollTimeout } wait := min(backoff, remaining) select { diff --git a/auth/gcp/client_test.go b/auth/gcp/client_test.go index 119c40a10..6c1a9ad11 100644 --- a/auth/gcp/client_test.go +++ b/auth/gcp/client_test.go @@ -282,6 +282,17 @@ func TestRetrieveValidatesRequest(t *testing.T) { {name: "resource path traversal", req: Request{Resource: "projects/p/../q/authProviders/a", UserID: "u"}}, {name: "resource query injection", req: Request{Resource: "projects/p/authProviders/a?x=1", UserID: "u"}}, {name: "resource with space", req: Request{Resource: "projects/p/authProviders/a b", UserID: "u"}}, + // A name that normalizes to a different one routes to a different service + // than the one validateResource inspected. + {name: "resource empty segment", req: Request{Resource: "projects/p//authProviders/a", UserID: "u"}}, + {name: "resource trailing slash", req: Request{Resource: "projects/p/locations/l/connectors/c/", UserID: "u"}}, + {name: "resource dot segment", req: Request{Resource: "projects/p/locations/l/connectors/c/.", UserID: "u"}}, + // Percent-escapes are rejected by the charset, not decoded: the name is + // interpolated into a URL, so an escape that survives becomes traversal or + // a segment break once the server decodes it. + {name: "resource percent-escaped dot", req: Request{Resource: "projects/p/authProviders/a%2e%2e", UserID: "u"}}, + {name: "resource percent-escaped slash", req: Request{Resource: "projects/p%2flocations/authProviders/a", UserID: "u"}}, + {name: "resource bare percent", req: Request{Resource: "projects/p/authProviders/a%", UserID: "u"}}, } // Point at a live server: a client with no endpoint fails at transport for // every input, which cannot tell a rejected request from an unreachable one. @@ -298,7 +309,7 @@ func TestRetrieveValidatesRequest(t *testing.T) { if err == nil { t.Fatalf("RetrieveCredential(%+v) = nil error, want error", tc.req) } - if !strings.Contains(err.Error(), "requires a") && !strings.Contains(err.Error(), "invalid characters") { + if !strings.Contains(err.Error(), "requires a") && !strings.Contains(err.Error(), "resource ") { t.Errorf("error = %v, want a request-validation error", err) } if got := hits.Load(); got != 0 { diff --git a/auth/gcp/doc.go b/auth/gcp/doc.go index 11aec2d70..b7f4bd3c7 100644 --- a/auth/gcp/doc.go +++ b/auth/gcp/doc.go @@ -12,19 +12,21 @@ // See the License for the specific language governing permissions and // limitations under the License. -// Package gcp is a hand-rolled REST client for the Google Cloud Agent Identity +// Package gcp resolves per-user Google Cloud credentials from the Agent Identity // and IAM Connector credential services. Given a resource name it retrieves an // end-user credential and maps it to an [auth.Credential], polling while the // service reports a non-interactive "pending" state and surfacing interactive // consent as an [auth.ConsentRequiredError]. // +// [NewProvider] is the entry point for an agent: it returns an +// [auth.CredentialProvider] that takes the acting user from the invocation +// context, so a tool's outbound requests are authenticated as the end user. +// [NewClient] is the transport underneath, usable on its own where the caller +// already knows the user. +// // No generated Go client libraries exist for these (preview) services and their // surface is a single RPC, so the client is hand-rolled over net/http to keep // dependencies light. Calls to the credential services are authenticated with // Application Default Credentials (cloud-platform scope) unless a custom // *http.Client is supplied. -// -// This package holds only the transport-level client. The [auth.CredentialProvider] -// that resolves the acting user from the invocation context is a separate, -// higher layer. package gcp diff --git a/auth/gcp/provider.go b/auth/gcp/provider.go new file mode 100644 index 000000000..b7c7d4a3e --- /dev/null +++ b/auth/gcp/provider.go @@ -0,0 +1,317 @@ +// 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" + "sync/atomic" + "time" + + "google.golang.org/adk/v2/agent" + "google.golang.org/adk/v2/auth" +) + +// ProviderScheme identifies a GCP auth resource and the access it requests. It +// mirrors adk-python's GcpAuthProviderScheme. +type ProviderScheme struct { + // Name is the full resource name, routed by [Client]: either + // projects/*/locations/*/connectors/* (IAM Connector) or + // projects/*/locations/*/authProviders/* (Agent Identity). + // + // [NewProvider] accepts only those two shapes. That is stricter than + // [Client.RetrieveCredential] and than adk-python, both of which send any + // non-connector name to Agent Identity: at wiring time a name outside the two + // is a typo, and this type's name invites passing an HTTP auth scheme. + 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]. +type ProviderConfig struct { + // Scheme is the resource to mint credentials for. Required. + Scheme ProviderScheme + // Client reaches the credential services. When nil, a default client backed + // by Application Default Credentials is built lazily on first use, against + // the production endpoints; that adds two failure modes to Credential, since + // the build can fail or exceed its 30s bound. Pass a [Client] from + // [NewClient] to reach another endpoint or to tune the poll timeout — note + // that this trades the lazy path away: unless the Client carries its own + // HTTPClient, [NewClient] discovers Application Default Credentials + // synchronously, so the cost and any failure move to startup, where they can + // at least be reported. + Client *Client +} + +// ErrClientUnavailable means the default Application Default Credentials client +// is not available: discovery failed, or it did not finish inside the bound. The +// lookup is not cancellable, so that bound is on the wait rather than on the +// attempt, which keeps running — a later call may well succeed. +var ErrClientUnavailable = errors.New("gcp: default credentials client unavailable") + +// ErrNoActingUser means the provider could not determine the acting end user, +// either because the context is not an ADK context or because the invocation +// carries no user. Unlike adk-python, which degrades such a turn into an auth +// request, the Go provider fails the request: no user, no credential. +var ErrNoActingUser = errors.New("gcp: no acting user") + +// defaultInitTimeout bounds how long a caller waits for the default client. The +// build itself cannot be bounded by a context — FindDefaultCredentials reads the +// credentials file with os.ReadFile and probes with the context-free +// metadata.OnGCE(), neither of which observes cancellation — so the bound lives +// on the waiting side. +const defaultInitTimeout = 30 * time.Second + +// NewProvider returns an [auth.CredentialProvider] that resolves credentials for +// cfg.Scheme 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. Two +// requirements then fall on the transport that carries the authenticated +// requests, neither of which this package can enforce: +// +// - Every request must descend from the invoking user's context. A transport +// that shares one connection across invocations does not qualify — +// mcptoolset included: its per-call POSTs are per-user, but the MCP session +// and the standalone server-to-client stream are opened on the context of +// whichever user connected first and stay bound to it. +// - The transport must not follow a cross-host redirect. net/http strips +// Authorization above the RoundTripper, so [auth.Transport] re-resolves and +// re-applies the end user's credential to the redirect target. Set +// CheckRedirect on the http.Client that carries the transport; the +// ADC-backed client [NewClient] builds for itself does the same. +// +// Wiring this up also means trusting the embedding server: ADK does not +// authenticate session.UserID, and it now decides whose credential is minted. +// +// ctx is used only to build the default client, and only for its values; its +// cancellation is not honored, because that client outlives any one request. +// Pass the process-scoped context the rest of the app is wired with, not a +// request's. It is ignored when cfg.Client is set. +func NewProvider(ctx context.Context, cfg ProviderConfig) (auth.CredentialProvider, error) { + if cfg.Scheme.Name == "" { + return nil, errors.New("gcp: NewProvider requires a scheme Name") + } + // A malformed name is a wiring mistake; catch it here rather than on every + // request from inside a transport. Stricter than RetrieveCredential, which + // routes any non-connector name to Agent Identity: at wiring time a name that + // is not one of the two known shapes is a typo, not a new collection. + if err := validateResource(cfg.Scheme.Name); err != nil { + return nil, fmt.Errorf("gcp: NewProvider: %w", err) + } + if !connectorResourceRE.MatchString(cfg.Scheme.Name) && !authProviderResourceRE.MatchString(cfg.Scheme.Name) { + return nil, fmt.Errorf("gcp: NewProvider: scheme Name %q is neither projects/*/locations/*/connectors/* nor projects/*/locations/*/authProviders/*", cfg.Scheme.Name) + } + // A zero Client would nil-deref deep inside net/http on first use. + if cfg.Client != nil && cfg.Client.httpClient == nil { + return nil, errors.New("gcp: ProviderConfig.Client must come from NewClient") + } + p := &provider{ + scheme: cfg.Scheme, + client: cfg.Client, + newClient: func(ctx context.Context) (*Client, error) { return NewClient(ctx, nil) }, + initTimeout: defaultInitTimeout, + } + // The provider outlives this call and re-reads Scopes per request, so it must + // not alias a caller-mutable slice. + p.scheme.Scopes = slices.Clone(cfg.Scheme.Scopes) + if cfg.Client == nil { + // Captured only where it will be used: it is retained for the life of the + // provider, and pinning a caller's context graph for nothing is a leak. + p.initCtx = context.WithoutCancel(ctx) + } + return p, nil +} + +type provider struct { + scheme ProviderScheme + // initCtx roots the lazily built default client, which is why NewProvider + // asks for a process-scoped context. Nil when a Client was supplied. + initCtx context.Context + // newClient and initTimeout are fields, not package constants, so tests can + // drive the failure and hang paths a real ADC lookup cannot be made to hit. + newClient func(context.Context) (*Client, error) + initTimeout time.Duration + + mu sync.Mutex + client *Client + pending *clientInit // in-flight lazy init, shared by concurrent callers +} + +// clientInit is one attempt at building the default client. Its fields are +// written by the attempt's own goroutine and read by waiters only after done is +// closed. +type clientInit struct { + done chan struct{} + client *Client + err error + // blown is set by the first waiter whose bound expires. Later waiters fail + // fast instead of each paying the bound again: the attempt is kept running, + // so without this a stuck lookup costs every outbound request its full + // initTimeout for as long as it is stuck. + blown atomic.Bool +} + +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 { + return nil, fmt.Errorf("%w: no ADK invocation identity on the context — not an agent invocation, or its session is unset", ErrNoActingUser) + } + if id.UserID == "" { + // No ids in the message: this text is fed to the model and persisted in + // the session, and every id here comes off the request. + return nil, fmt.Errorf("%w: the invocation's session carries no user", ErrNoActingUser) + } + + client, err := p.resolveClient(ctx) + if err != nil { + // Not attributed: a client-init failure is about this process's own + // credentials, not about the resource, and every provider in the process + // fails it identically. Attributing it would also stack the package prefix. + return nil, err + } + cred, err := client.RetrieveCredential(ctx, Request{ + Resource: p.scheme.Name, + UserID: id.UserID, + Scopes: p.scheme.Scopes, + ContinueURI: p.scheme.ContinueURI, + }) + if err != nil { + return nil, p.attribute(err) + } + return cred, nil +} + +// attribute names the resource a retrieval failure belongs to: several providers +// can be wired into one process, and an unattributed error says nothing about +// which. +// +// It names nothing else. This error becomes the tool's error, which is fed to +// the model and persisted in the session, and every id available here is +// supplied by the caller — a user id is commonly an email, and a session id +// arrives unvalidated from the request path. The invocation is already +// identified by the trace and the session the error is stored in. +func (p *provider) attribute(err error) error { + return fmt.Errorf("gcp: resource %q: %w", p.scheme.Name, err) +} + +// resolveClient returns the configured client, building a default one (backed by +// Application Default Credentials) on first use. +// +// Concurrent callers share one attempt and each waits on the earlier of its own +// context and initTimeout: auth.Transport resolves a credential per outbound +// request, so a slow cold start must not outlive the request that triggered it. +// A failed attempt is not cached; the next call retries. +// +// A hung attempt is not abandoned. The lookup cannot be cancelled, so retiring +// it would start a fresh one every initTimeout, each parked in a syscall pinning +// an OS thread. Waiters get a prompt error instead, and the moment the stuck +// lookup returns its client is published and callers recover. +func (p *provider) resolveClient(ctx context.Context) (*Client, error) { + p.mu.Lock() + if c := p.client; c != nil { + p.mu.Unlock() + return c, nil + } + in := p.pending + if in == nil { + in = &clientInit{done: make(chan struct{})} + p.pending = in + go p.runInit(in) + } + p.mu.Unlock() + + if in.blown.Load() { + select { + case <-in.done: // landed after the bound blew; fall through to the result + default: + return nil, fmt.Errorf("%w: an earlier attempt exceeded %v and is still running", ErrClientUnavailable, p.initTimeout) + } + } + + timer := time.NewTimer(p.initTimeout) + defer timer.Stop() + select { + case <-in.done: + if in.err != nil { + return nil, in.err + } + return in.client, nil + case <-timer.C: + // A sentinel of its own, not context.DeadlineExceeded: that is what the + // caller-deadline arm below returns, and the two mean different things. + in.blown.Store(true) + return nil, fmt.Errorf("%w after %v", ErrClientUnavailable, p.initTimeout) + case <-ctx.Done(): + return nil, fmt.Errorf("gcp: waiting for the default credentials client: %w", ctx.Err()) + } +} + +// runInit builds the default client once and publishes it to the waiters on in. +func (p *provider) runInit(in *clientInit) { + // This runs on a goroutine the provider owns, so nothing above can recover a + // panic here and it would take the process down — where an eagerly built + // client would merely have panicked in the caller's own frame. Report it as + // this attempt's failure instead, panic value and all, and release the + // waiters: without this, an abrupt exit leaves pending set with its goroutine + // dead and every later caller waits out initTimeout forever. + published := false + defer func() { + if published { + return + } + if r := recover(); r != nil { + in.err = fmt.Errorf("gcp: building the default credentials client panicked: %v", r) + } else if in.err == nil { + in.err = errors.New("gcp: building the default credentials client did not complete") + } + p.publish(in) + }() + + c, err := p.newClient(p.initCtx) + switch { + case err != nil: + in.err = fmt.Errorf("%w: %w", ErrClientUnavailable, err) + case c == nil: + // Caching a nil client would only move the failure to the next retrieval. + in.err = errors.New("gcp: default credentials client builder returned no client") + default: + in.client = c + } + published = true + p.publish(in) +} + +// publish caches a successful client, frees the in-flight slot so a failed +// attempt is retried rather than cached, and releases the waiters. +func (p *provider) publish(in *clientInit) { + p.mu.Lock() + if in.err == nil { + p.client = in.client + } + p.pending = nil + p.mu.Unlock() + close(in.done) +} diff --git a/auth/gcp/provider_internal_test.go b/auth/gcp/provider_internal_test.go new file mode 100644 index 000000000..b183e1af6 --- /dev/null +++ b/auth/gcp/provider_internal_test.go @@ -0,0 +1,392 @@ +// 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" + "net/http" + "runtime" + "strings" + "sync" + "sync/atomic" + "testing" + "time" +) + +// TestResolveClientBuildsDefaultClient drives the lazy ADC path end to end: +// discovery, the cached client, and a retrieval whose token is minted after the +// request that triggered construction is already gone. +func TestResolveClientBuildsDefaultClient(t *testing.T) { + fakeADC(t) + srv, calls := sequenceServer(`{"success":{"token":"tok","header":"Authorization: Bearer"}}`) + defer srv.Close() + + p := newTestProvider(t) + ctx, cancel := context.WithCancel(t.Context()) + c, err := p.resolveClient(ctx) + if err != nil { + t.Fatalf("resolveClient() error = %v", err) + } + cancel() + // The default client targets the production endpoint; retarget it so the + // retrieval below stays offline. + c.agentIdentityURL = srv.URL + + if _, err := c.RetrieveCredential(t.Context(), Request{Resource: authProviderResource, UserID: "u"}); err != nil { + t.Fatalf("RetrieveCredential() error = %v", err) + } + if got := atomic.LoadInt32(calls); got != 1 { + t.Errorf("service calls = %d, want 1", got) + } + if again, err := p.resolveClient(t.Context()); err != nil || again != c { + t.Errorf("resolveClient() = %v, %v; want the cached client", again, err) + } +} + +// TestResolveClientSingleFlight pins the concurrency design: many first callers +// share one build and all get the same client. Nothing exercises the mutex +// unless a test runs callers together — the race detector included, which is +// what catches a lock that stops locking. +func TestResolveClientSingleFlight(t *testing.T) { + var attempts atomic.Int32 + built := &Client{httpClient: http.DefaultClient} + release := make(chan struct{}) + p := newTestProvider(t) + p.newClient = func(context.Context) (*Client, error) { + attempts.Add(1) + <-release // hold the flight open so every caller piles up on it + return built, nil + } + + const callers = 16 + var wg sync.WaitGroup + got := make([]*Client, callers) + errs := make([]error, callers) + for i := range callers { + wg.Add(1) + go func() { + defer wg.Done() + got[i], errs[i] = p.resolveClient(t.Context()) + }() + } + time.Sleep(20 * time.Millisecond) + close(release) + wg.Wait() + + if n := attempts.Load(); n != 1 { + t.Errorf("init attempts = %d, want 1 (concurrent callers must share one flight)", n) + } + for i := range callers { + if errs[i] != nil || got[i] != built { + t.Fatalf("caller %d = %v, %v; want the single shared client", i, got[i], errs[i]) + } + } +} + +// TestResolveClientHonorsCallerDeadline pins that a caller waiting on a slow +// init is bounded by its own context: auth.Transport resolves a credential per +// outbound request, so one cold start must not stall every concurrent request. +func TestResolveClientHonorsCallerDeadline(t *testing.T) { + p := newTestProvider(t) + p.newClient = blockingInit(t) + + const deadline = 50 * time.Millisecond + ctx, cancel := context.WithTimeout(t.Context(), deadline) + defer cancel() + start := time.Now() + _, err := p.resolveClient(ctx) + elapsed := time.Since(start) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("resolveClient() error = %v, want the caller's deadline", err) + } + // Loose enough not to flake on a loaded machine, tight enough that a waiter + // which ignored the caller's ctx — and so sat out the 30s initTimeout — + // fails here. + if elapsed >= p.initTimeout { + t.Errorf("resolveClient() returned after %v, want the caller's %v deadline, not the %v init bound", elapsed, deadline, p.initTimeout) + } +} + +// TestResolveClientBoundsWaitOnHungInit pins the initTimeout: an ADC lookup that +// never returns cannot be cancelled, so a caller with no deadline of its own must +// still be released — while the attempt itself is kept. +func TestResolveClientBoundsWaitOnHungInit(t *testing.T) { + p := newTestProvider(t) + p.newClient = blockingInit(t) + p.initTimeout = 50 * time.Millisecond + + // A sentinel of its own: the caller-deadline arm returns + // context.DeadlineExceeded, and a caller must be able to tell them apart. + _, err := p.resolveClient(t.Context()) + if !errors.Is(err, ErrClientUnavailable) { + t.Fatalf("resolveClient() error = %v, want ErrClientUnavailable", err) + } + if errors.Is(err, context.DeadlineExceeded) { + t.Errorf("resolveClient() error = %v, want it distinguishable from a caller deadline", err) + } + // Retiring the attempt would start a fresh lookup — and leak a fresh + // thread-pinning goroutine — every initTimeout, forever. + p.mu.Lock() + pending := p.pending + p.mu.Unlock() + if pending == nil { + t.Error("provider dropped the hung attempt; the next caller would start another lookup") + } +} + +// TestResolveClientPublishesLateClient pins that a lookup which lands after every +// waiter gave up is not wasted: the attempt is kept rather than retired precisely +// so the next caller gets its client instead of starting another lookup. +func TestResolveClientPublishesLateClient(t *testing.T) { + built := &Client{httpClient: http.DefaultClient} + release := make(chan struct{}) + var attempts atomic.Int32 + p := newTestProvider(t) + p.initTimeout = 20 * time.Millisecond + p.newClient = func(context.Context) (*Client, error) { + attempts.Add(1) + <-release + return built, nil + } + + if _, err := p.resolveClient(t.Context()); !errors.Is(err, ErrClientUnavailable) { + t.Fatalf("resolveClient() error = %v, want ErrClientUnavailable", err) + } + close(release) + + // The attempt finishes on its own goroutine; wait for it to publish. + deadline := time.Now().Add(2 * time.Second) + var got *Client + for time.Now().Before(deadline) { + if c, err := p.resolveClient(t.Context()); err == nil { + got = c + break + } + time.Sleep(time.Millisecond) + } + if got != built { + t.Fatalf("resolveClient() = %v, want the late-landing client", got) + } + if n := attempts.Load(); n != 1 { + t.Errorf("init attempts = %d, want 1 (the late client must be used, not rebuilt)", n) + } +} + +// TestRunInitSurvivesAbruptBuilder pins that a builder which does not return +// normally still releases the waiters and the in-flight slot. Without the +// deferred publish, pending stays set with its goroutine dead and every later +// caller waits out initTimeout, forever. +func TestRunInitSurvivesAbruptBuilder(t *testing.T) { + for _, tc := range []struct { + name string + build func(context.Context) (*Client, error) + wantMsg string + }{ + { + name: "panic", + build: func(context.Context) (*Client, error) { panic("credentials discovery blew up") }, + wantMsg: "blew up", // surfaced, not swallowed + }, + { + // What a t.Fatal inside a caller-supplied builder does. + name: "runtime.Goexit", + build: func(context.Context) (*Client, error) { runtime.Goexit(); return nil, nil }, + wantMsg: "did not complete", + }, + } { + t.Run(tc.name, func(t *testing.T) { + p := newTestProvider(t) + p.initTimeout = 5 * time.Second // long enough that a wedge shows up as one + p.newClient = tc.build + + done := make(chan error, 1) + go func() { _, err := p.resolveClient(t.Context()); done <- err }() + select { + case err := <-done: + if err == nil || !strings.Contains(err.Error(), tc.wantMsg) { + t.Fatalf("resolveClient() error = %v, want it to report %q", err, tc.wantMsg) + } + case <-time.After(2 * time.Second): + t.Fatal("resolveClient() blocked; an abrupt builder wedged the provider") + } + p.mu.Lock() + pending := p.pending + p.mu.Unlock() + if pending != nil { + t.Error("provider kept the dead attempt; the next caller would wait on it forever") + } + }) + } +} + +// TestResolveClientFailFastAfterBlownBound pins that a stuck lookup costs the +// bound once, not once per request. The attempt is deliberately kept running, so +// without the latch every outbound request would re-pay the full initTimeout for +// as long as the lookup is stuck. +func TestResolveClientFailFastAfterBlownBound(t *testing.T) { + p := newTestProvider(t) + p.newClient = blockingInit(t) + p.initTimeout = 200 * time.Millisecond + + start := time.Now() + if _, err := p.resolveClient(t.Context()); !errors.Is(err, ErrClientUnavailable) { + t.Fatalf("first resolveClient() error = %v, want ErrClientUnavailable", err) + } + first := time.Since(start) + + start = time.Now() + _, err := p.resolveClient(t.Context()) + second := time.Since(start) + if !errors.Is(err, ErrClientUnavailable) { + t.Fatalf("second resolveClient() error = %v, want ErrClientUnavailable", err) + } + // The fail-fast path does not wait at all, so a quarter of the bound is a + // generous ceiling and still far under the ~200ms a re-paid bound costs. + if second > p.initTimeout/4 { + t.Errorf("second caller waited %v against a %v bound (first paid %v): the blown bound must be latched", second, p.initTimeout, first) + } +} + +// TestResolveClientDiscoveryFailureIsMatchable pins that the common failure — +// discovery not working at all — carries the same sentinel as the timeout. A +// caller behind a RoundTripper cannot switch on a message. +func TestResolveClientDiscoveryFailureIsMatchable(t *testing.T) { + p := newTestProvider(t) + p.newClient = func(context.Context) (*Client, error) { return nil, errors.New("no credentials on this host") } + + _, err := p.resolveClient(t.Context()) + if !errors.Is(err, ErrClientUnavailable) { + t.Fatalf("resolveClient() error = %v, want it to wrap ErrClientUnavailable", err) + } + if !strings.Contains(err.Error(), "no credentials on this host") { + t.Errorf("resolveClient() error = %v, want the underlying cause kept", err) + } +} + +// TestResolveClientRetriesFailedInit pins that a failed ADC discovery is not +// cached: the doc promises the next call retries, and a provider that wedged on +// a transient environment failure would never recover. +func TestResolveClientRetriesFailedInit(t *testing.T) { + var attempts atomic.Int32 + p := newTestProvider(t) + p.newClient = func(context.Context) (*Client, error) { + if attempts.Add(1) == 1 { + return nil, errors.New("no credentials") + } + return &Client{httpClient: http.DefaultClient}, nil + } + + if _, err := p.resolveClient(t.Context()); err == nil { + t.Fatal("resolveClient() = nil error, want the discovery failure") + } + if _, err := p.resolveClient(t.Context()); err != nil { + t.Errorf("resolveClient() after a failed attempt: error = %v, want a retry", err) + } + if got := attempts.Load(); got != 2 { + t.Errorf("init attempts = %d, want 2 (a failure must not be cached)", got) + } +} + +// TestResolveClientRejectsNilClient pins that a builder returning (nil, nil) is +// an error rather than a cached nil that nil-derefs on the next retrieval, inside +// an http.RoundTripper. +func TestResolveClientRejectsNilClient(t *testing.T) { + p := newTestProvider(t) + p.newClient = func(context.Context) (*Client, error) { return nil, nil } + + if _, err := p.resolveClient(t.Context()); err == nil { + t.Fatal("resolveClient() = nil error, want a nil client rejected") + } + p.mu.Lock() + cached := p.client + p.mu.Unlock() + if cached != nil { + t.Errorf("provider cached %v, want nothing", cached) + } +} + +// TestNewProviderIgnoresWiringContextCancellation pins the documented contract +// that NewProvider's ctx supplies values only: the default client outlives any +// one request, so cancelling what was passed at wiring time must not stop it +// being built. +func TestNewProviderIgnoresWiringContextCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + p, err := NewProvider(ctx, ProviderConfig{Scheme: ProviderScheme{Name: authProviderResource}}) + if err != nil { + t.Fatalf("NewProvider() error = %v", err) + } + cancel() + + built := &Client{httpClient: http.DefaultClient} + prov := p.(*provider) + var sawErr error + prov.newClient = func(ctx context.Context) (*Client, error) { + sawErr = ctx.Err() + return built, nil + } + got, err := prov.resolveClient(t.Context()) + if err != nil || got != built { + t.Fatalf("resolveClient() = %v, %v; want the client despite the cancelled wiring context", got, err) + } + if sawErr != nil { + t.Errorf("the builder saw ctx.Err() = %v, want the cancellation stripped", sawErr) + } +} + +// TestNewProviderKeepsWiringContextOnlyWhenLazy pins that a provider given a +// Client does not pin its caller's context — and with it the caller's whole +// session and event graph — for the life of the process. +func TestNewProviderKeepsWiringContextOnlyWhenLazy(t *testing.T) { + client, err := NewClient(t.Context(), &Config{HTTPClient: http.DefaultClient}) + if err != nil { + t.Fatalf("NewClient() error = %v", err) + } + cfg := ProviderConfig{Scheme: ProviderScheme{Name: authProviderResource}, Client: client} + p, err := NewProvider(t.Context(), cfg) + if err != nil { + t.Fatalf("NewProvider() error = %v", err) + } + if got := p.(*provider).initCtx; got != nil { + t.Errorf("initCtx = %v, want nil when a Client is supplied", got) + } + if got := newTestProvider(t).initCtx; got == nil { + t.Error("initCtx = nil on the lazy path, want the wiring context") + } +} + +// blockingInit returns a client init that never returns until the test ends — +// the shape the init timeout exists for (os.ReadFile on a stalled mount, a hung +// resolver inside metadata.OnGCE). +func blockingInit(t *testing.T) func(context.Context) (*Client, error) { + t.Helper() + release := make(chan struct{}) + t.Cleanup(func() { close(release) }) + return func(context.Context) (*Client, error) { + <-release + return nil, errors.New("released") + } +} + +// newTestProvider builds a provider with no configured client, so resolveClient +// takes the lazy default-client path. +func newTestProvider(t *testing.T) *provider { + t.Helper() + p, err := NewProvider(t.Context(), ProviderConfig{Scheme: ProviderScheme{Name: authProviderResource}}) + if err != nil { + t.Fatalf("NewProvider() error = %v", err) + } + return p.(*provider) +} diff --git a/auth/gcp/provider_test.go b/auth/gcp/provider_test.go new file mode 100644 index 000000000..56fa5cba6 --- /dev/null +++ b/auth/gcp/provider_test.go @@ -0,0 +1,294 @@ +// 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" + "errors" + "io" + "net/http" + "net/http/httptest" + "slices" + "strings" + "sync" + "testing" + + "google.golang.org/adk/v2/agent" + "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" +) + +const testResource = "projects/p/locations/l/authProviders/ap" + +// TestProviderCredential drives two users through one shared provider: the +// provider is long-lived, so serving one user's credential to another is the +// failure that matters. It also pins scopes and continueUri on the wire. +func TestProviderCredential(t *testing.T) { + var gotUsers []string + var gotScopes []string + var gotContinueURI string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body struct { + UserID string `json:"userId"` + Scopes []string `json:"scopes"` + ContinueURI string `json:"continueUri"` + } + _ = json.NewDecoder(r.Body).Decode(&body) + gotUsers = append(gotUsers, body.UserID) + gotScopes, gotContinueURI = body.Scopes, body.ContinueURI + // Echo the caller back, so a credential served to the wrong user shows up. + _, _ = io.WriteString(w, `{"success":{"token":"tok-`+body.UserID+`","header":"Authorization: Bearer"}}`) + })) + defer srv.Close() + + scopes := []string{"s1", "s2"} + p := newProvider(t, srv, gcp.ProviderScheme{ + Name: testResource, + Scopes: scopes, + ContinueURI: "https://example.test/continue", + }) + scopes[0] = "mutated" // the provider must have cloned this + + for _, user := range []string{"alice", "bob"} { + cred, err := p.Credential(adkContext(t, user)) + if err != nil { + t.Fatalf("Credential(%q) error = %v", user, err) + } + if bc, ok := cred.(auth.BearerCredential); !ok || bc.Token != "tok-"+user { + t.Errorf("credential for %q = %+v, want bearer %q", user, cred, "tok-"+user) + } + } + if !slices.Equal(gotUsers, []string{"alice", "bob"}) { + t.Errorf("service saw users %q, want [alice bob]", gotUsers) + } + if !slices.Equal(gotScopes, []string{"s1", "s2"}) { + t.Errorf("body scopes = %q, want [s1 s2] (caller's later mutation must not leak)", gotScopes) + } + if gotContinueURI != "https://example.test/continue" { + t.Errorf("body continueUri = %q, want the scheme's", gotContinueURI) + } +} + +// TestProviderCredentialConcurrent drives two users through one shared provider +// at the same time. The sequential case above pins the wire contract; this one +// pins that concurrency cannot cross the streams. +func TestProviderCredentialConcurrent(t *testing.T) { + 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) + _, _ = io.WriteString(w, `{"success":{"token":"tok-`+body.UserID+`","header":"Authorization: Bearer"}}`) + })) + defer srv.Close() + + p := newProvider(t, srv, gcp.ProviderScheme{Name: testResource}) + users := []string{"alice", "bob", "carol", "dave"} + ctxs := make([]context.Context, len(users)) + for i, u := range users { + ctxs[i] = adkContext(t, u) + } + + var wg sync.WaitGroup + got := make([]auth.Credential, len(users)*8) + errs := make([]error, len(users)*8) + for i := range got { + wg.Add(1) + go func() { + defer wg.Done() + got[i], errs[i] = p.Credential(ctxs[i%len(users)]) + }() + } + wg.Wait() + + for i, cred := range got { + want := "tok-" + users[i%len(users)] + if errs[i] != nil { + t.Fatalf("Credential(%s) error = %v", users[i%len(users)], errs[i]) + } + if bc, ok := cred.(auth.BearerCredential); !ok || bc.Token != want { + t.Errorf("credential %d = %+v, want bearer %q", i, cred, want) + } + } +} + +// TestProviderErrorAttribution pins that a failed retrieval says which resource +// failed — several providers can be wired into one process — and names no +// caller-supplied id, since this text is fed to the model and persisted in the +// session. Sentinels must stay matchable through the wrap. +func TestProviderErrorAttribution(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusForbidden) + _, _ = io.WriteString(w, `denied`) + })) + defer srv.Close() + + p := newProvider(t, srv, gcp.ProviderScheme{Name: testResource}) + ctx := adkContext(t, "alice@example.test") + id, _ := agent.IdentityFromContext(ctx) + _, err := p.Credential(ctx) + if err == nil { + t.Fatal("Credential() = nil error, want the service failure") + } + if !strings.Contains(err.Error(), testResource) { + t.Errorf("Credential() error = %v, want it to name the resource", err) + } + for _, unwanted := range []string{"alice@example.test", id.SessionID} { + if strings.Contains(err.Error(), unwanted) { + t.Errorf("Credential() error = %v, want the caller-supplied %q kept out of it", err, unwanted) + } + } + var apiErr *gcp.APIError + if !errors.As(err, &apiErr) || apiErr.StatusCode != http.StatusForbidden { + t.Errorf("Credential() error = %v, want a wrapped *gcp.APIError with status 403", err) + } +} + +// TestProviderNoActingUser covers both identity failures: the guard must reject +// before any service call, the two cases must stay distinguishable, and neither +// message may carry a caller-supplied id — this text reaches the model and is +// persisted in the session. +func TestProviderNoActingUser(t *testing.T) { + tests := []struct { + name string + ctx func(t *testing.T) context.Context + wantMsg string + }{ + { + name: "not an ADK context", + ctx: func(t *testing.T) context.Context { return t.Context() }, + wantMsg: "no ADK invocation identity", + }, + { + name: "invocation without a user", + ctx: userlessADKContext, + wantMsg: "carries no user", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Fails the test if reached: no identity means no service call. + srv := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + t.Error("credentials service must not be called without an ADK identity") + })) + defer srv.Close() + + p := newProvider(t, srv, gcp.ProviderScheme{Name: testResource}) + _, err := p.Credential(tt.ctx(t)) + if !errors.Is(err, gcp.ErrNoActingUser) { + t.Fatalf("Credential() error = %v, want gcp.ErrNoActingUser", err) + } + if !strings.Contains(err.Error(), tt.wantMsg) { + t.Errorf("Credential() error = %v, want it to mention %q", err, tt.wantMsg) + } + for _, unwanted := range []string{"app", "sid"} { + if strings.Contains(err.Error(), unwanted) { + t.Errorf("Credential() error = %v, want the caller-supplied %q kept out of it", err, unwanted) + } + } + }) + } +} + +func TestNewProviderValidatesScheme(t *testing.T) { + // Everything here is a wiring mistake, and each must fail at construction + // rather than on every request from inside an http.RoundTripper. + bad := []struct { + name string + cfg gcp.ProviderConfig + }{ + {"empty name", gcp.ProviderConfig{}}, + {"path traversal", cfgFor("projects/p/locations/l/authProviders/../../secret")}, + {"empty path segment", cfgFor("projects/p/locations/l/authProviders//ap")}, + {"trailing slash routes differently after normalization", cfgFor("projects/p/locations/l/connectors/c/")}, + {"not a resource name at all", cfgFor("Bearer")}, + {"unknown collection", cfgFor("projects/p/locations/l/authProvidrs/ap")}, + {"truncated", cfgFor("projects/p")}, + {"unconstructed client", gcp.ProviderConfig{Scheme: gcp.ProviderScheme{Name: testResource}, Client: &gcp.Client{}}}, + } + for _, tt := range bad { + t.Run(tt.name, func(t *testing.T) { + if _, err := gcp.NewProvider(t.Context(), tt.cfg); err == nil { + t.Fatal("NewProvider() = nil error, want the config rejected") + } + }) + } + + good := []string{ + testResource, + "projects/p/locations/l/connectors/c", + // Domain-scoped project ids carry a colon. + "projects/example.com:my-project/locations/l/authProviders/ap", + } + for _, name := range good { + t.Run("accepts "+name, func(t *testing.T) { + if _, err := gcp.NewProvider(t.Context(), cfgFor(name)); err != nil { + t.Fatalf("NewProvider(%q) error = %v", name, err) + } + }) + } +} + +func cfgFor(name string) gcp.ProviderConfig { + return gcp.ProviderConfig{Scheme: gcp.ProviderScheme{Name: name}} +} + +// newProvider builds a provider whose client targets srv. +func newProvider(t *testing.T, srv *httptest.Server, scheme gcp.ProviderScheme) auth.CredentialProvider { + t.Helper() + 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(t.Context(), gcp.ProviderConfig{Scheme: scheme, Client: client}) + if err != nil { + t.Fatalf("NewProvider() error = %v", err) + } + return p +} + +// 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}) +} + +// userlessADKContext returns an invocation whose session carries no user — the +// shape session.InMemoryService refuses to create, but that a custom session +// service can produce. +func userlessADKContext(t *testing.T) context.Context { + t.Helper() + return icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{Session: userlessSession{}}) +} + +// userlessSession embeds a nil session.Session for the accessors the identity +// path never reaches. +type userlessSession struct{ session.Session } + +func (userlessSession) ID() string { return "sid" } +func (userlessSession) AppName() string { return "app" } +func (userlessSession) UserID() string { return "" } diff --git a/auth/providers.go b/auth/providers.go index 9c9866c3e..4b5d248c1 100644 --- a/auth/providers.go +++ b/auth/providers.go @@ -33,13 +33,15 @@ 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 + // IdentityFromContext in the agent package — 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 - // authorization URI; the tool layer turns that into a human-in-the-loop + // non-interactively, Credential returns an error wrapping a + // *ConsentRequiredError that carries the authorization URI — find it with + // errors.As, since a provider may wrap it. The tool layer turns that into a + // human-in-the-loop // consent round-trip. Non-interactive providers never return it. Credential(ctx context.Context) (Credential, error) } @@ -63,9 +65,12 @@ type ConsentRequiredError struct { Key string } -// Error implements error. +// Error implements error. It deliberately omits AuthURI: this error becomes a +// tool's error, which is fed to the model and persisted in the session, and the +// consent URI carries the state and nonce that bind the credential. Consumers +// read it off the field. func (e *ConsentRequiredError) Error() string { - return fmt.Sprintf("auth: interactive consent required (auth_uri=%q)", e.AuthURI) + return "auth: interactive consent required" } // StaticToken returns a provider that always yields the given bearer token. diff --git a/auth/providers_test.go b/auth/providers_test.go index a9d655d58..8dc3c2252 100644 --- a/auth/providers_test.go +++ b/auth/providers_test.go @@ -110,8 +110,11 @@ func TestConsentRequiredError(t *testing.T) { if consent.AuthURI != "https://consent.example" { t.Errorf("AuthURI = %q, want %q", consent.AuthURI, "https://consent.example") } - if !strings.Contains(err.Error(), "consent.example") { - t.Errorf("Error() = %q, want it to mention the auth URI", err.Error()) + // The message must not carry the URI: it becomes a tool's error, which reaches + // the model and the session store, and the URI carries the state and nonce + // that bind the credential. + if strings.Contains(err.Error(), "consent.example") { + t.Errorf("Error() = %q, want the auth URI kept on the field only", err.Error()) } } diff --git a/internal/adkcontext/adkcontext.go b/internal/adkcontext/adkcontext.go new file mode 100644 index 000000000..e9b83d0a1 --- /dev/null +++ b/internal/adkcontext/adkcontext.go @@ -0,0 +1,58 @@ +// 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 the key. The key is therefore unforgeable, but the value +// it addresses is not: any in-process context wrapper can read the exported +// agent.Identity on its way past and hand back a different one, without ever +// naming the key. Treat the identity as trusted only as far as every wrapper in +// the chain is. +const IdentityKey ctxKey = 0 + +// Recovered returns what read produced, and whether it returned at all. +// +// It exists for the ADK Value implementations, which read the invocation +// identity off session.Session — a public interface whose implementations are +// arbitrary code. A nil or typed-nil session, a session wrapping a nil one (the +// shape llmagent.newWrappedSession produces for a nil original), or a Session() +// accessor that declines, all panic on the way. Value runs inside +// http.RoundTripper on the caller's goroutine, where net/http does not recover, +// so a broken session would take the process down; reporting no identity instead +// fails the credential path closed. +// +// Nothing partially built escapes: on a panic the zero value is returned. The +// cost is that a bug inside a caller's accessor surfaces as a missing identity +// rather than a stack trace — deliberate, since the alternative here is killing +// the process, and the credential path fails closed either way. +// +// It contains a panic, not every way out: an accessor that calls runtime.Goexit +// ends the calling goroutine, which no recover can undo. +func Recovered[T any](read func() T) (v T, ok bool) { + defer func() { + if recover() != nil { + var zero T + v, ok = zero, false + } + }() + return read(), true +} diff --git a/internal/context/from_context_test.go b/internal/context/from_context_test.go new file mode 100644 index 000000000..d52f437d6 --- /dev/null +++ b/internal/context/from_context_test.go @@ -0,0 +1,320 @@ +// 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") + } +} + +// TestIdentityFromContextSessionShapes covers the session shapes a +// [session.Session] implementation can legally take. Several of them used to +// panic inside Value — a struct value tripped reflect.Value.IsNil, a typed-nil +// pointer passed an interface-nil check and then dereferenced, and a session +// wrapping a nil one (llmagent.newWrappedSession's shape for a nil original) +// panicked in the accessor. Value runs inside an http.RoundTripper, where +// net/http does not recover. Every context implementation must also agree: two +// of them answering the identity key differently is its own bug. +func TestIdentityFromContextSessionShapes(t *testing.T) { + cases := []struct { + name string + session session.Session + want agent.Identity + wantOK bool + }{ + { + name: "pointer", + session: &ptrSession{id: "sid-1", app: "app-1", user: "user-42"}, + want: agent.Identity{UserID: "user-42", AppName: "app-1", SessionID: "sid-1"}, + wantOK: true, + }, + { + name: "struct value", + session: valueSession{id: "sid-1", app: "app-1", user: "user-42"}, + want: agent.Identity{UserID: "user-42", AppName: "app-1", SessionID: "sid-1"}, + wantOK: true, + }, + { + // A typed-nil pointer whose accessors do not dereference is usable. + name: "typed-nil pointer with safe accessors", + session: (*safeNilSession)(nil), + want: agent.Identity{UserID: "user-nil", AppName: "app-nil", SessionID: "sid-nil"}, + wantOK: true, + }, + {name: "nil", session: nil}, + {name: "typed-nil pointer", session: (*ptrSession)(nil)}, + {name: "wrapper over a nil session", session: &wrapperSession{}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + ic := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{Session: tc.session}) + for _, ctx := range []struct { + name string + ctx context.Context + }{ + {"invocation context", ic}, + {"promoted common context", agent.Promote(ic)}, + {"tool context", agent.NewToolContext(ic, "fc-1", nil, nil)}, + } { + id, ok := agent.IdentityFromContext(ctx.ctx) + if ok != tc.wantOK || id != tc.want { + t.Errorf("IdentityFromContext(%s) = %+v, %v; want %+v, %v", ctx.name, id, ok, tc.want, tc.wantOK) + } + } + }) + } +} + +// TestIdentityAfterWithContext pins the promoted context's own identity branch. +// WithContext replaces the embedded parent with a non-ADK context while keeping +// the invocation — the one shape where delegating to the parent cannot recover +// the identity, and the reason the branch exists. agent.go does exactly this +// around a tracing span. +func TestIdentityAfterWithContext(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}) + + detached := agent.Promote(ic).WithContext(context.Background()) + id, ok := agent.IdentityFromContext(detached) + want := agent.Identity{UserID: "user-42", AppName: "app-1", SessionID: resp.Session.ID()} + if !ok || id != want { + t.Errorf("IdentityFromContext(WithContext) = %+v, %v; want %+v, true", id, ok, want) + } +} + +// TestValueWithNilEmbeddedContext pins the nil-parent guard: +// NewCleanToolContextTestOnly builds a context with no embedded parent, so +// without it every non-identity key is a nil-interface method call. +func TestValueWithNilEmbeddedContext(t *testing.T) { + ic := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{}) + clean, err := agent.NewCleanToolContextTestOnly(agent.Promote(ic), "fc-1", nil, nil) + if err != nil { + t.Fatalf("NewCleanToolContextTestOnly() error = %v", err) + } + if got := clean.Value(wrapKey{}); got != nil { + t.Errorf("Value(wrapKey{}) = %v, want nil", got) + } +} + +// TestIdentityDoesNotInheritEnclosingInvocation pins that identity resolution +// fails closed. A nested invocation with no session of its own must not report +// the enclosing invocation's user: that user's credential would then be minted +// for a call they never made. +func TestIdentityDoesNotInheritEnclosingInvocation(t *testing.T) { + svc := session.InMemoryService() + resp, err := svc.Create(t.Context(), &session.CreateRequest{AppName: "app-1", UserID: "alice"}) + if err != nil { + t.Fatalf("session Create() error = %v", err) + } + outer := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{Session: resp.Session}) + if id, ok := agent.IdentityFromContext(outer); !ok || id.UserID != "alice" { + t.Fatalf("outer IdentityFromContext() = %+v, %v; want alice", id, ok) + } + + nested := icontext.NewInvocationContext(outer, icontext.InvocationContextParams{}) + for _, tc := range []struct { + name string + ctx context.Context + }{ + {"nested invocation", nested}, + // Promote copies the invocation, so this would resolve through the + // nested context's own guard even if the promoted one leaked. + {"nested and promoted", agent.Promote(nested)}, + {"nested tool context", agent.NewToolContext(nested, "fc-1", nil, nil)}, + // A non-ADK wrapper in between must not restore what the guard refused. + {"nested behind a non-ADK wrapper", context.WithValue(nested, wrapKey{}, "x")}, + // Reparented onto a plain context that happens to carry the enclosing + // invocation. The derived context still speaks for the nested invocation, + // so the parent must not supply a user that invocation refused. + // (WithContext given an InvocationContext is different on the two + // implementations that rebind on it — agent.commonContext and + // agent.callbackContextWrapper — where it deliberately changes which + // invocation the context speaks for.) + {"nested, reparented onto a plain carrier of the enclosing invocation", agent.Promote(nested).WithContext(context.WithValue(outer, wrapKey{}, "x"))}, + } { + if id, ok := agent.IdentityFromContext(tc.ctx); ok { + t.Errorf("%s IdentityFromContext() = %+v, true; want no identity, not the enclosing user", tc.name, id) + } + } +} + +// TestIdentityThroughSessionlessWrappers pins the other half of that rule: a +// context that does not own a session — a tool or callback context, whose +// Session() returns nil by design — must delegate rather than report no +// identity, or every outbound request from a re-derived tool context fails with +// ErrNoActingUser. +func TestIdentityThroughSessionlessWrappers(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}) + want := agent.Identity{UserID: "user-42", AppName: "app-1", SessionID: resp.Session.ID()} + + toolCtx := agent.NewToolContext(ic, "fc-1", nil, nil) + callbackCtx := agent.NewCallbackContext(ic, nil) + // Each of these has a session-less wrapper as the invocation it speaks for, so + // every one exercises the delegation. A context derived directly from ic does + // not, and would pass whatever the delegation did. + for _, tc := range []struct { + name string + ctx context.Context + }{ + {"tool context", toolCtx}, + {"promoted tool context", agent.Promote(toolCtx)}, + {"tool context re-derived from a tool context", agent.NewToolContext(toolCtx, "fc-2", nil, nil)}, + {"callback context re-derived from a tool context", agent.NewCallbackContext(toolCtx, nil)}, + {"callback context", callbackCtx}, + {"promoted callback context", agent.Promote(callbackCtx)}, + {"tool context re-derived from a callback context", agent.NewToolContext(callbackCtx, "fc-3", nil, nil)}, + } { + id, ok := agent.IdentityFromContext(tc.ctx) + if !ok || id != want { + t.Errorf("%s IdentityFromContext() = %+v, %v; want %+v, true", tc.name, id, ok, want) + } + } +} + +// TestIdentityFromPanickingSession pins that a session whose Session() accessor +// itself panics — not only its field accessors — costs the identity and not the +// process: this runs inside an http.RoundTripper, where net/http does not +// recover. +func TestIdentityFromPanickingSession(t *testing.T) { + parent := context.WithValue(t.Context(), wrapKey{}, "x") + inner := icontext.NewInvocationContext(parent, icontext.InvocationContextParams{}) + ic := agent.Promote(panickingInvocation{InvocationContext: inner}) + if id, ok := agent.IdentityFromContext(ic); ok { + t.Errorf("IdentityFromContext() = %+v, true; want no identity", id) + } + // The panic costs the identity and nothing else. + if got := ic.Value(wrapKey{}); got != "x" { + t.Errorf("Value(wrapKey{}) = %v, want %q", got, "x") + } +} + +type panickingInvocation struct{ agent.InvocationContext } + +func (panickingInvocation) Session() session.Session { panic("Session() is not supported here") } + +// TestValueDelegatesUnknownKeys pins that a session shape that stops the +// identity lookup does not stop every other key: Value is on the hot path for +// net/http, tracing and logging keys that have nothing to do with the session. +func TestValueDelegatesUnknownKeys(t *testing.T) { + parent := context.WithValue(t.Context(), wrapKey{}, "x") + ic := icontext.NewInvocationContext(parent, icontext.InvocationContextParams{Session: (*ptrSession)(nil)}) + for _, tc := range []struct { + name string + ctx context.Context + }{ + {"invocation context", ic}, + {"promoted common context", agent.Promote(ic)}, + } { + if got := tc.ctx.Value(wrapKey{}); got != "x" { + t.Errorf("%s Value(wrapKey{}) = %v, want %q", tc.name, got, "x") + } + } +} + +// valueSession and ptrSession embed a nil session.Session for the accessors the +// identity path never reaches, and read their own fields for the ones it does — +// so a typed-nil *ptrSession panics on use, as a real session would. +type valueSession struct { + session.Session + id, app, user string +} + +func (s valueSession) ID() string { return s.id } +func (s valueSession) AppName() string { return s.app } +func (s valueSession) UserID() string { return s.user } + +type ptrSession struct { + session.Session + id, app, user string +} + +// wrapperSession is the shape llmagent.newWrappedSession produces for a nil +// original: non-nil, but every identity accessor is promoted from the nil +// embedded interface and panics on the first call. +type wrapperSession struct{ session.Session } + +// safeNilSession answers without touching its receiver, so a typed-nil one is +// still a working session. +type safeNilSession struct{ session.Session } + +func (*safeNilSession) ID() string { return "sid-nil" } +func (*safeNilSession) AppName() string { return "app-nil" } +func (*safeNilSession) UserID() string { return "user-nil" } + +func (s *ptrSession) ID() string { return s.id } +func (s *ptrSession) AppName() string { return s.app } +func (s *ptrSession) UserID() string { return s.user } diff --git a/internal/context/invocation_context.go b/internal/context/invocation_context.go index c31d1e30d..7eff54487 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,30 @@ 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 identity key (so agent.IdentityFromContext can recover it); every +// other key delegates to the embedded context, preserving existing behavior. +// +// This type owns its session, so the identity key is answered here even when the +// session cannot be read — with no identity rather than the enclosing +// invocation's, whose user made no such call and whose credential would +// otherwise be minted for it. +func (c *InvocationContext) Value(key any) any { + if key == adkcontext.IdentityKey { + if id, ok := adkcontext.Recovered(func() agent.Identity { + s := c.params.Session + return agent.Identity{UserID: s.UserID(), AppName: s.AppName(), SessionID: s.ID()} + }); ok { + return id + } + return nil + } + if c.Context == nil { + return nil + } + return c.Context.Value(key) +} + // ResumedInput always returns (nil, false) for the base // invocation context. Implementations that carry a resume payload // override this method.