Skip to content

fix(sts): deliver the exchanged STS token on MCP tool calls - #2795

Open
ricardomiguel-os wants to merge 1 commit into
kagent-dev:mainfrom
ricardomiguel-os:fix/sts-mcp-session-id-context-value
Open

fix(sts): deliver the exchanged STS token on MCP tool calls#2795
ricardomiguel-os wants to merge 1 commit into
kagent-dev:mainfrom
ricardomiguel-os:fix/sts-mcp-session-id-context-value

Conversation

@ricardomiguel-os

Copy link
Copy Markdown

Problem

The STS token exchange succeeds and the exchanged token is cached, but it never leaves the pod. MCP tool calls go out with no Authorization header at all, so the upstream server sees no identity instead of the user delegated to the agent.

sessionIDFromContext recovered the session with a direct type assertion:

ctx.(interface{ SessionID() string })

A type assertion inspects only the context object in hand. createTransport always builds its http.Client with a non-zero Timeout (defaultTimeout is 30m), and http.Client re-wraps the request context in a deadline context whenever Timeout > 0 (net/http.setRequestCancelcontext.WithDeadline). By the time RoundTrip runs, the context in hand is that wrapper, not ADK's ToolContext, so the assertion never matches and HeaderProvider returns nil.

No configuration avoids this: the timeout is non-zero even when unset.

The neighbouring propagateToken path keeps working because it uses a2asrv.CallContextFrom, a Value lookup, and Value lookups traverse the parent chain.

Fix

  1. a2a/executor.go stamps the session ID as a context value, beside the bearer token it already stamps.
  2. sts/plugin.go prefers that Value lookup in sessionIDFromContext, keeping the type assertion as a fallback so callers still holding ADK's ToolContext are unaffected.
  3. models/base.go gives contextKey a name. BearerTokenKey and SessionIDKey would otherwise both be pointers to a zero-size struct, and Go may place those at the same address, so the two keys compare equal and the session stamp overwrites the bearer token. The end-to-end test caught this by seeing the session ID arrive at the STS as subject_token.

The token cache key, its write side and its eviction are untouched. NewTokenPropagationPlugin's signature is unchanged.

Requests with no user in context, such as the toolset discovery performed at startup, still get no header. That behaviour is asserted in the new test rather than changed.

Evidence

Red/green on the end-to-end test, reverting only the Value lookup in sessionIDFromContext:

--- FAIL: TestSTSExchangedTokenReachesMCPTool
    invocation[0..5] Authorization = ""
    want: "Bearer EXCHANGED-STS-TOKEN"

With the fix: go test -race -skip 'TestE2E.*' ./adk/... green across all packages, golangci-lint run reports 0 issues, go build ./... and gofmt clean.

Tests

  • mcp/sts_injection_test.go — end-to-end, nothing stubbed but the LLM: real runner.Runner, llmagent, mcptoolset, the real CreateToolsets/createTransport (so a real http.Client with the deadline wrapper), real sts.TokenPropagationPlugin performing a real exchange against an httptest STS, and a real MCP server recording every Authorization it receives.
  • sts/plugin_test.go — session recovery precedence: session as a value, the same wrapped in a deadline context, session only via SessionID(), no session, session present with no cached token, and value-beats-method.
  • a2a/executor_test.go — the executor stamps the session ID, and the two context keys coexist without overwriting each other.

@github-actions github-actions Bot added the bug Something isn't working label Sep 10, 2026
The STS token exchange succeeds and the exchanged token is cached, but it never
leaves the pod: MCP tool calls go out carrying no Authorization header at all, so
the upstream server never sees the user delegated to the agent.

HeaderProvider recovered the session with a direct type assertion:

    ctx.(interface{ SessionID() string })

That inspects only the context object in hand, without consulting its parents.
createTransport always builds its http.Client with a non-zero Timeout
(defaultTimeout is 30m), and http.Client re-wraps the request context in a
deadline context whenever Timeout > 0 (net/http.setRequestCancel ->
context.WithDeadline). By the time RoundTrip runs, the context in hand is that
wrapper, not ADK's ToolContext, so the assertion never matched, HeaderProvider
returned nil, and no header was injected. No configuration avoids it: the
timeout is non-zero even when unset. The neighbouring propagateToken path keeps
working because it uses a2asrv.CallContextFrom, a Value lookup, and Value
lookups traverse the parent chain.

Store the session ID as a context value in the A2A executor, beside the bearer
token it already stamps, and prefer that lookup in sessionIDFromContext. The
type assertion is retained as a fallback so callers still holding ADK's
ToolContext are unaffected.

Behaviour for requests with no user in context, such as the toolset discovery
performed at startup, is unchanged: no header is injected.

One supporting change: models.contextKey now carries a name. BearerTokenKey and
SessionIDKey are both pointers to it, and pointers to zero-size structs may
share an address, which made the two keys compare equal and let the session
stamp overwrite the bearer token. Caught by the new end-to-end test, which saw
the session ID arrive at the STS as the subject token.

Tests:

- mcp: end-to-end regression driving a real runner, llmagent, mcptoolset,
  CreateToolsets/createTransport, TokenPropagationPlugin and MCP server,
  asserting the exchanged token reaches the tool. With session recovery by type
  assertion alone, every invocation request goes out with an empty Authorization
  header.
- sts: HeaderProvider session recovery, including the deadline-wrapped case.
- a2a: the executor stamps the session ID, and the two context keys coexist.

Signed-off-by: Ricardo Gonçalves (rng) <ricardo.goncalves@outsystems.com>
@ricardomiguel-os
ricardomiguel-os force-pushed the fix/sts-mcp-session-id-context-value branch from 31cbb5c to 61ce6e7 Compare September 10, 2026 16:35
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