fix(sts): key the ADK token cache by acting subject, not session alone (Python)

[QuentinBisson]

Supersedes #2333 (auto-closed by stale-bot; a rebase onto main conflicted with the resource/audience support, so GitHub would not let it be reopened).

Ref #2181. The Go counterpart is #2459 and carries the closing reference.

Problem

Same bug as the Go plugin: the cache was keyed by session ID alone, so a session shared by several subjects reused whichever caller's exchanged token was cached first.

Change

Key the cache by (session_id, caller credential), the credential hashed. before_run_callback and header_provider both resolve the caller before the cache lookup, so each caller gets its own exchanged token.

What the key is derived from

The caller part of the key comes from the inbound Authorization header, not from get_subject_token. That hook receives the whole session state, so an implementation reading a session-scoped field returns the same value for every caller and collapses the key back to one entry per session, which is the bug being fixed here.

The header is read from the run's own session state, which _agent_executor.py writes once per run as a state_delta and the session service hands out per run. So header_provider rebuilds the key on every tool call without consulting the hook and without state held between callbacks.

With no inbound credential, the hook's output stands in for the caller part. That mode has no per-caller identity to preserve, so one entry per session is correct for it, and keying on the header alone would drop propagation entirely.

An empty subject identifies no principal, so it yields no key at all, and neither does a session with no id.

Why the key is a hash and not iss+sub

A cache hit returns a delegated token without performing an exchange, so the key decides who receives whose authority. iss and sub are not verified anywhere on this path: a forged, unsigned token carrying a victim's claims would select the victim's entry and never reach the STS that would have rejected it. Hashing the raw token makes a forged token a cache miss, so it goes to the STS and fails there.

The cost is one extra exchange when a caller's bearer rotates mid-session.

Cache lifetime

An entry never outlives the credential that keyed it. The expiry is the earlier of the exchanged token's exp and the caller bearer's exp. Without that cap, a caller replaying an expired bearer keeps hitting the cached delegated token instead of reaching the STS that would reject it.

One entry per caller instead of one per session also means a token with no exp at all would pin an entry per caller for the lifetime of the process. Those entries fall back to a 5 minute lifetime, which also arms the sweep gate that a None expiry left unset.

An entry outliving the value a get_subject_token hook has since rotated is bounded by that same expiry.

The sweep removes expired entries only, so the cache is sized by the number of distinct (session, subject) pairs seen within a token lifetime. No size cap or LRU is added; the bounded TTL is what bounds it.

Failing closed in header_provider

header_provider returns no headers instead of raising into the invocation when there is no invocation context, no key for the caller, or a non-dict headers value in session state. This matches how _read_subject_token handles a raising get_subject_token.

Scope

Propagate-only mode (no STS configured) forwards the caller's own token. This fix stops one caller's token reaching another caller, but does not add exchange, audience narrowing or an act claim to that mode: there is no STS to exchange with when none is configured.

Note

go-unit-tests fails for a reason unrelated to this diff, which touches no Go. go test ./... runs api/v1alpha2 and api/v1alpha3 in parallel and both download the envtest binaries into the same bin dir on demand, so on a cold runner cache one package execs a half-written etcd ("text file busy") while the other fails to unpack kubectl. It reproduces on main with an empty go/bin/k8s, and passes once the binaries are provisioned first.

The two prev-stable upgrade jobs fail in helm-install-provider with UPGRADE FAILED: context deadline exceeded, which is also unrelated to a Python-only diff.

[Copilot]

Pull request overview

Fixes a correctness/security bug in the Python ADK STS token-propagation plugin where delegated (exchanged) tokens were cached only by session ID, causing shared sessions with multiple acting subjects to reuse the first subject’s exchanged token. This aligns the Python runtime behavior with the Go-side fix and prevents identity collapse across callers within the same session.

Changes:

• Key the token cache by (session_id, subject_key(subject_token)), where subject_key prefers iss+sub and falls back to a hash for opaque/sub-less tokens.
• Resolve the subject token before cache lookup in before_run_callback, and sweep expired entries across the cache in after_run_callback to handle per-subject entries.
• Update/extend integration tests to validate multi-subject session behavior, issuer scoping, opaque token handling, and fail-closed behavior when no subject token is present.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
python/packages/agentsts-adk/src/agentsts/adk/_base.py	Implements per-subject cache keying, adds subject discriminator logic, adjusts cache lookup + eviction behavior.
python/packages/agentsts-adk/tests/test_adk_integration.py	Updates existing assertions to use cache_key(...) and adds new cases covering multi-subject behavior and cache eviction semantics.

Suppressed comments (1)

python/packages/agentsts-adk/src/agentsts/adk/_base.py:263

• cache_key() now calls the configured get_subject_token callback, which can raise. Since header_provider uses cache_key during tool invocation, an exception here would break tool calls; it’s safer to fail closed (return a key with an empty subject) when subject resolution fails.

        session = invocation_context.session
        return self._cache_key_for(session.id, self._read_subject_token(session.state))

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[QuentinBisson]

Partly fixed in bc252de. The sweep is now gated on the earliest expiry across the cache, so after_run_callback returns without walking anything until a cached token can actually be evicted; the scan recomputes the earliest expiry as it goes, and inserts keep it up to date.

I kept the sweep global rather than scoping it to the running session: a session-scoped sweep never reaches the entries of sessions that stop running, so those would be retained for the lifetime of the process. Trading unbounded growth for an occasional walk seemed like the wrong side of that.

[QuentinBisson]

Closed out in 6471f57. The gate had a hole: a token with no exp cached a None expiry, _earlier_expiry skipped it, and a cache holding only such entries left _earliest_expiry unset so the sweep never ran. Those entries now get a bounded lifetime, so the gate is always armed and every entry stays evictable.

[QuentinBisson]

Fixed in bc252de. _read_subject_token now catches and logs, returning None, so a raising caller-supplied get_subject_token skips propagation instead of aborting the run. One change covers both call sites (before_run_callback and cache_key), and it matches how exchange_token and fetch_actor_token failures are already handled here.

[QuentinBisson]

Also handled for the other caller of the same code path in 6471f57: header_provider reaches _read_subject_token via cache_key on every tool call, and it was dereferencing readonly_context._invocation_context without a guard. It now returns no headers instead of raising into the invocation.