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

[QuentinBisson]

Supersedes #2333 (auto-closed by stale-bot; rebase onto current main hit real conflicts with the resource/audience support, so GitHub would not let the original be reopened).

What

Python counterpart to the Go STS fix: a session shared by multiple subjects reused whichever caller's exchanged token was cached first, because the cache was keyed only by session ID.

Keys the cache by (session_id, subject_key(subject_token)), where subject_key derives a stable per-principal discriminator from the acting bearer's issuer-scoped sub claim (falling back to a hash of the raw token for opaque/sub-less tokens). The subject token is now resolved before the cache lookup in before_run_callback, and after_run_callback sweeps all expired entries (not just the acting subject's) since a session can now hold one entry per subject.

Closes #2181 (Python side; the Go counterpart is #2459).

Cache lifetime

Keying by (session_id, subject) holds one entry per caller instead of one per
session, so an entry whose token carries no exp claim would pin a slot per
caller for the lifetime of the process. Those entries now get a bounded
lifetime. That also arms the sweep gate, which a None expiry left unset.

header_provider runs on every tool call and dereferenced a private attribute
of an optional argument. It now returns no headers instead of raising into the
invocation, which is the same fail-closed handling _read_subject_token already
applies to a raising get_subject_token.

An absent iss leaves sub unqualified, so issuer-less tokens partition by
token hash rather than by a key that two issuers omitting iss could both
produce.

Scope

Propagate-only mode (no STS configured) forwards the caller's own token, so 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 one is not configured, and a second caller in a session is now a
normal case in both modes rather than something either mode has to refuse.

go-unit-tests is red on main as well (CRD validation in go/api/v1alpha2
and go/api/v1alpha3); it is not caused by this 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.