Skip to content

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

Draft
QuentinBisson wants to merge 2 commits into
kagent-dev:mainfrom
QuentinBisson:fix/py-sts-cache-per-subject
Draft

fix(sts): key the ADK token cache by acting subject, not session alone (Python)#2333
QuentinBisson wants to merge 2 commits into
kagent-dev:mainfrom
QuentinBisson:fix/py-sts-cache-per-subject

Conversation

@QuentinBisson

@QuentinBisson QuentinBisson commented Jul 25, 2026

Copy link
Copy Markdown
Contributor

What

The Python ADK token-propagation plugin cached the exchanged (delegated) token
keyed by session ID alone. A session that carries messages from more than one
subject therefore reused whichever subject's token seeded the cache first, so
every later caller's MCP tool calls ran under that first subject's delegated
identity.

This is the Python counterpart of #2317, which fixes the same defect in the Go
ADK runtime. The two runtimes had the same cache shape and the same bug.

Changes

  • cache_key is now session.id plus the acting subject. _subject_key()
    derives the discriminator from the caller's token as iss + sub (sub is
    only unique within an issuer), or a hash of the raw token when there is no
    sub. The token is parsed unverified and used only to partition the cache,
    never to gate a security decision; it is validated server-side on exchange.
  • before_run_callback resolves the subject token before the cache lookup
    instead of after it, so the lookup can be keyed by subject at all. The
    resolution is centralized in _read_subject_token and reused rather than
    repeated.
  • after_run_callback sweeps every expired entry. A session now holds one entry
    per subject, and only the acting subject's key is derivable there, so the
    previous single-key eviction would leave other subjects' expired entries
    behind.

Both the write path (before_run_callback) and the read path
(header_provider) resolve the subject from session.state, which
_agent_executor.py rewrites from the current request on every normal run. No
equivalent of the Go fix's CallContext fallback is needed here.

Behavior change

Same shape as #2317: in a read-write shared session, a visitor's tool calls move
from the owner's exchanged token to the visitor's own. That corrects an identity
collapse rather than introducing one, but it is externally observable.

A request that carries no resolvable subject token no longer reuses a token the
session cached earlier; it skips propagation. That is the fail-closed half of
keying by subject.

Known gap, not addressed here

_agent_executor.py refreshes session.state["headers"] only on the normal run
branch. The HITL-resume branch skips it, so a resume keeps the previous request's
headers and the subject resolves to the original caller. If a different user
approves an interrupted run, they get the original caller's token. Per-subject
keying cannot see that, because the approver's identity never reaches session
state. Fixing it means writing headers on the HITL branch too, in kagent-adk
rather than agentsts-adk.

Tests

Adds six cases: two subjects in one session keep separate tokens and each gets
its own back from header_provider; the same subject twice exchanges once; the
same sub at two issuers does not collide; opaque and sub-less tokens still
partition via the hash branch; an expired entry belonging to a subject other
than the acting one is still evicted; and a tokenless caller in a session that
already holds a cached token gets nothing back from header_provider (the
fail-closed read path).

Verified the new tests catch the original defect: with the key forced back to
session.id alone, two subjects collapse to one cache entry and one exchange,
with the second caller inheriting the first's token.

The 57 existing plugin tests assert on cache contents keyed by session id and
are updated to key through plugin.cache_key(...) so they no longer pin the key
format.

The Python token-propagation plugin cached the exchanged token under the
session id, so a session carrying messages from more than one subject reused
whichever subject's token seeded the cache first. Every later caller's MCP tool
calls then ran under that first subject's delegated identity.

The acting subject is now resolved from session state before the cache lookup
and folded into the key as the issuer-scoped sub claim, falling back to a hash
of the raw token for opaque or sub-less tokens. after_run_callback sweeps every
expired entry, since a session now holds one per subject and only the acting
subject's key is derivable there.

Signed-off-by: QuentinBisson <quentin@giantswarm.io>
@github-actions github-actions Bot added bug Something isn't working and removed bug Something isn't working labels Jul 25, 2026
…y note

Signed-off-by: QuentinBisson <quentin@giantswarm.io>
@QuentinBisson
QuentinBisson force-pushed the fix/py-sts-cache-per-subject branch from 3b591e5 to be7a504 Compare July 25, 2026 13:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant