-
Notifications
You must be signed in to change notification settings - Fork 424
Description
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:
-
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).
-
Single audience:
before_run_callbackdoesn'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
BaseToolinstance, enabling per-audience exchange)
The flow becomes:
before_run_callback: Extract and store the subject token (no exchange)before_tool_callback: Exchange on firstMCPToolcall per (session, audience), cache result; skip non-MCP toolsheader_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 audienceMcpServerTool.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_providerintypes.pycreates per-toolset closures, plugin handles exchange lifecycle
Metadata
Metadata
Assignees
Labels
Type
Projects
Status