Skip to content

Lazy STS token exchange with per-audience scoping #1430

@syn-zhu

Description

@syn-zhu

Problem

The current ADKTokenPropagationPlugin performs STS token exchange eagerly in before_run_callback, which runs once at the start of every agent invocation. This has two issues:

  1. Unnecessary exchanges: The token is exchanged even if no MCP tools end up being called during the invocation (e.g., the agent only calls memory tools or decides it doesn't need tools at all).

  2. Single audience: before_run_callback doesn't have access to tool metadata, so there's no way to determine which MCP server (and therefore which audience) the token should be exchanged for. All MCP servers get the same unscoped token, which is insufficient when different downstream services require different audience claims.

Solution

Lazy exchange

Move the STS token exchange to before_tool_callback, which is:

  • Async (avoids the sync/async problem that would arise from doing it in header_provider)
  • Per-tool (only fires when a tool is actually called)
  • Tool-aware (has access to the BaseTool instance, enabling per-audience exchange)

The flow becomes:

  • before_run_callback: Extract and store the subject token (no exchange)
  • before_tool_callback: Exchange on first MCPTool call per (session, audience), cache result; skip non-MCP tools
  • header_provider: Read from cache (sync, as before)
  • after_run_callback: Clean up all caches (bare session key + all audience-scoped keys)

Per-audience scoping

Add STSAudience field at two levels:

  • RemoteMCPServerSpec.STSAudience -- server-level default audience
  • McpServerTool.STSAudience -- per-agent override (takes precedence)

The plugin maintains a register_toolset(toolset, audience) mapping from id(session_manager) to audience string. When before_tool_callback fires for an MCPTool, it resolves the audience from the tool's session manager and exchanges with that specific audience. Tokens are cached under session_id:audience keys so different MCP servers in the same invocation get independently scoped tokens.

Benefits

  • No wasted HTTP round-trips to the STS when MCP tools aren't called
  • Each MCP server can receive a token scoped to its own audience
  • Non-MCP tools (SaveMemoryTool, LoadMemoryTool, AgentTool, etc.) are correctly skipped
  • Clean separation: create_header_provider in types.py creates per-toolset closures, plugin handles exchange lifecycle

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    Status

    Backlog

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions