fix: key STS token cache by acting subject, not session alone - #2317

…-dev#2459)

Supersedes kagent-dev#2317 (auto-closed by stale-bot; a rebase onto main
conflicted with kagent-dev#2106, so GitHub would not let it be reopened).

Closes kagent-dev#2181. The Python counterpart is kagent-dev#2460.

## Problem

The STS token cache was keyed by session ID alone. When several subjects
share one session (for example several users on the same A2A session),
every caller runs its tool calls with whichever exchanged token was
cached first.

## Change

Key the cache by `(sessionID, subject)`, where `subject` is a SHA-256
hash of the caller's bearer token. `BeforeRunCallback` and
`HeaderProvider` both resolve the caller's bearer before the cache
lookup, so each caller gets its own exchanged token.

The bearer resolution moves to `models.BearerTokenFromContext`. It
prefers the value the executor stored under `BearerTokenKey` and falls
back to the A2A call context, which is what reaches the MCP transport
layer when `BearerTokenKey` was not threaded into the request context.
`a2a/executor.go` now calls the same helper, so there is one copy of the
parser instead of two.

`GetSubjectTokenFunc` documents that it must be a pure function of the
bearer. The hash of the bearer stands in for the subject token only
because the hook derives one from the other; an implementation that
mints or fetches a token would have its first result served for the
entry's lifetime.

## Why the key is a token 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.

An empty subject identifies no principal, so it yields no key at all. An
entry stored under it would be shared by every credential-less caller in
the session.

## Cache lifetime

An entry never outlives the credential that keyed it. The expiry is the
earliest of the exchange's own lifetime (`expires_in`, or the exchanged
token's `exp`), the subject token's `exp`, and the caller bearer's
`exp`. Without that cap, a caller replaying an expired credential keeps
hitting the cached delegated token instead of reaching the STS that
would reject it.

When no token states an `exp` there is no expiry, and nothing may invent
one: `BeforeRunCallback` runs once per run, so a run outliving a made-up
deadline would lose its credential part way through. The entry carries a
second field, `evictAfter`, which bounds memory alone. `Expiry` decides
whether the token may still be injected; `evictAfter` decides only when
the entry may be dropped, and it comes due 5 minutes after the last use
when nothing states an expiry.

A cache hit renews `evictAfter` and records a use sequence, so neither
the sweep nor either capacity bound can drop an entry a run is still
using, and a bound drops the least recently used entries first.
`getCachedToken` therefore takes the write lock, and returns a copy so
that no caller reads a cached entry's fields once the lock is released.
`AfterRunCallback` sweeps every entry whose `evictAfter` has come due.

## Cache capacity

The cache is bounded per session (`maxEntriesPerSession`) and across
sessions (`maxCacheEntries`). The per-session bound is what matters:
without it, one session presenting many distinct credentials evicts the
entries of unrelated sessions, and a victim run that is between tool
calls has nothing more recent to show. `BeforeRunCallback` has already
run by then, so that run would reach its backends with no delegated
identity. The whole-cache bound backstops it so that many sessions
cannot grow the cache without limit.

Unlike the sweep, a capacity eviction drops entries whose credentials
are still valid, so it is reported above debug.

## API change

`GetTokenForSession` is removed. A session no longer has a single token,
so the signature cannot answer the question it asks. It had no callers
outside the plugin.

## Out of scope

kagent-dev#2181 also describes session state being keyed by `userID`. A shared
session still works: `userID` defaults to `A2A_USER_<contextID>`, so
callers on one context share a session and this fix gives each of them
its own exchanged token. What remains is that `userID` is part of the
session address, so sending a per-sender `x-user-id` resolves those
senders onto separate sessions instead of one conversation. That affects
attribution, not authorization, and is untouched here.

## Header path

`getCachedToken` checks the entry's expiry, and both `BeforeRunCallback`
and `HeaderProvider` go through it, so an entry past the caller
credential's expiry is never injected and a rejected exchange leaves
nothing usable behind. Two regression tests pin both, matching the cases
the Python plugin needed a fix for.

A request whose caller is identified but has no usable entry is reported
above debug. It loses its delegated identity, and nothing else says so.

---------

Signed-off-by: QuentinBisson <[email protected]>
Signed-off-by: Quentin Bisson <[email protected]>
Signed-off-by: Eitan Yarmush <[email protected]>
Co-authored-by: Eitan Yarmush <[email protected]>

…-dev#2459)

Supersedes kagent-dev#2317 (auto-closed by stale-bot; a rebase onto main
conflicted with kagent-dev#2106, so GitHub would not let it be reopened).

Closes kagent-dev#2181. The Python counterpart is kagent-dev#2460.

## Problem

The STS token cache was keyed by session ID alone. When several subjects
share one session (for example several users on the same A2A session),
every caller runs its tool calls with whichever exchanged token was
cached first.

## Change

Key the cache by `(sessionID, subject)`, where `subject` is a SHA-256
hash of the caller's bearer token. `BeforeRunCallback` and
`HeaderProvider` both resolve the caller's bearer before the cache
lookup, so each caller gets its own exchanged token.

The bearer resolution moves to `models.BearerTokenFromContext`. It
prefers the value the executor stored under `BearerTokenKey` and falls
back to the A2A call context, which is what reaches the MCP transport
layer when `BearerTokenKey` was not threaded into the request context.
`a2a/executor.go` now calls the same helper, so there is one copy of the
parser instead of two.

`GetSubjectTokenFunc` documents that it must be a pure function of the
bearer. The hash of the bearer stands in for the subject token only
because the hook derives one from the other; an implementation that
mints or fetches a token would have its first result served for the
entry's lifetime.

## Why the key is a token 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.

An empty subject identifies no principal, so it yields no key at all. An
entry stored under it would be shared by every credential-less caller in
the session.

## Cache lifetime

An entry never outlives the credential that keyed it. The expiry is the
earliest of the exchange's own lifetime (`expires_in`, or the exchanged
token's `exp`), the subject token's `exp`, and the caller bearer's
`exp`. Without that cap, a caller replaying an expired credential keeps
hitting the cached delegated token instead of reaching the STS that
would reject it.

When no token states an `exp` there is no expiry, and nothing may invent
one: `BeforeRunCallback` runs once per run, so a run outliving a made-up
deadline would lose its credential part way through. The entry carries a
second field, `evictAfter`, which bounds memory alone. `Expiry` decides
whether the token may still be injected; `evictAfter` decides only when
the entry may be dropped, and it comes due 5 minutes after the last use
when nothing states an expiry.

A cache hit renews `evictAfter` and records a use sequence, so neither
the sweep nor either capacity bound can drop an entry a run is still
using, and a bound drops the least recently used entries first.
`getCachedToken` therefore takes the write lock, and returns a copy so
that no caller reads a cached entry's fields once the lock is released.
`AfterRunCallback` sweeps every entry whose `evictAfter` has come due.

## Cache capacity

The cache is bounded per session (`maxEntriesPerSession`) and across
sessions (`maxCacheEntries`). The per-session bound is what matters:
without it, one session presenting many distinct credentials evicts the
entries of unrelated sessions, and a victim run that is between tool
calls has nothing more recent to show. `BeforeRunCallback` has already
run by then, so that run would reach its backends with no delegated
identity. The whole-cache bound backstops it so that many sessions
cannot grow the cache without limit.

Unlike the sweep, a capacity eviction drops entries whose credentials
are still valid, so it is reported above debug.

## API change

`GetTokenForSession` is removed. A session no longer has a single token,
so the signature cannot answer the question it asks. It had no callers
outside the plugin.

## Out of scope

kagent-dev#2181 also describes session state being keyed by `userID`. A shared
session still works: `userID` defaults to `A2A_USER_<contextID>`, so
callers on one context share a session and this fix gives each of them
its own exchanged token. What remains is that `userID` is part of the
session address, so sending a per-sender `x-user-id` resolves those
senders onto separate sessions instead of one conversation. That affects
attribution, not authorization, and is untouched here.

## Header path

`getCachedToken` checks the entry's expiry, and both `BeforeRunCallback`
and `HeaderProvider` go through it, so an entry past the caller
credential's expiry is never injected and a rejected exchange leaves
nothing usable behind. Two regression tests pin both, matching the cases
the Python plugin needed a fix for.

A request whose caller is identified but has no usable entry is reported
above debug. It loses its delegated identity, and nothing else says so.

---------

Signed-off-by: QuentinBisson <[email protected]>
Signed-off-by: Quentin Bisson <[email protected]>
Signed-off-by: Eitan Yarmush <[email protected]>
Co-authored-by: Eitan Yarmush <[email protected]>