feat(mcp): Add session state-based JWT token propagation for MCP tools#3675
feat(mcp): Add session state-based JWT token propagation for MCP tools#3675timof1308 wants to merge 22 commits intogoogle:mainfrom
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces a secure and flexible mechanism for propagating ephemeral data, like JWT tokens, from the client to MCP tools. The use of a non-persistent request_state that overrides the session state is a great design for handling sensitive, short-lived data. The addition of header_provider and declarative configuration options (state_header_mapping, state_header_format) in McpToolset makes this feature powerful and easy to use. The code is well-structured and thoroughly tested. My review includes a couple of suggestions to enhance the robustness of the new header provider logic by adding warnings for when non-primitive data types are used from the state, which could prevent silent misconfigurations.
There was a problem hiding this comment.
Code Review
This pull request introduces a secure mechanism for propagating ephemeral, request-specific data like JWT tokens to MCP tools. The addition of request_state to the InvocationContext and its integration into the ReadonlyContext is a clean solution to avoid persisting sensitive data. The new configuration options state_header_mapping and state_header_format in McpToolset provide a flexible, declarative way to manage headers. The changes are well-tested with new unit tests that cover the new functionality thoroughly. I have one suggestion to improve maintainability by reducing code duplication.
|
+1 |
c6d2245 to
d211763
Compare
|
Hi @timof1308 , Thank you for your work on this pull request. We appreciate the effort you've invested. |
|
thank you @ryanaiagent the CI is now pending on the latest commit (d0c7829). Once the workflow finishes, all tests should pass. Let me know if anything else is needed! |
|
+1 |
1 similar comment
|
+1 |
|
Hi @timof1308 ,This PR now has merge conflicts that require changes from your end. Could you please rebase your branch with the latest main branch to address these? Once this is complete, please let us know so we can proceed with the review. |
d0c7829 to
f5e20b0
Compare
|
hi @ryanaiagent, rebased with upstream main and reapplied |
02be37c to
6de52f9
Compare
02be37c to
087ee78
Compare
94038a3 to
fc287ca
Compare
381f70f to
6e513d2
Compare
e24335e to
d7f1652
Compare
|
Hi @timof1308 , can you fix the failing unit tests |
|
hi @ryanaiagent, I've pushed a fix for the unit tests The failures in however I noticed other failures in the logs (related to pubsub and litellm), but those appear to be present in the upstream adk-python repo |
5bf9574 to
038a129
Compare
|
Hi @timof1308 , can you also fix the failing formatting test. You can use autoformat.sh |
289687e to
5c7d479
Compare
|
@ryanaiagent file formatting has been applied |
|
Hi @timof1308 , Your PR has been received by the team and is currently under review. We will provide feedback as soon as we have an update to share. |
|
Hi @GWeale , can you please review this. |
…ze correctly into headers.
…ation logic in MCP tools.
…ise errors for non-primitive types and skip empty string values.
…type validation and add a test.
… and update header value sanitization to target control characters.
- Remove duplicate function from _internal.py (kept in mcp_toolset.py) - Inline _sanitize_header_value into sanitize_header_value - Update test imports to use public API Addresses review comment
…with `ChainMap` details, refine `request_state` initialization, and remove an unused header regex
…n improved docstring
…base - Restore missing __init__ method body that was lost during rebase - Add proper instance variable initialization (_mcp_tool, _mcp_session_manager, _require_confirmation, _header_provider, _progress_callback) - Rename get_function_declaration to _get_declaration with @OverRide decorator to match upstream convention This fixes test failures in test_mcp_toolset.py where tests expect _progress_callback attribute on McpTool instances.
012f8bc to
bc4e197
Compare
Please ensure you have read the contribution guide before creating a pull request.
Link to Issue or Description of Change
1. Main Tracking Issue:
This PR is a core component of the main MCP support effort tracked in #3449
2. Issues Closed/Addressed:
3. Related Pull Requests:
Related: #3675
Problem:
Currently, MCP tools in ADK lack a secure, standardized way to propagate per-user authentication tokens (like JWTs) from the client to the MCP server. Developers often have to hardcode credentials or modify core FastAPI endpoints to pass these headers, which is not scalable or secure for multi-user environments. Additionally, storing short-lived, sensitive tokens in the persistent session.state is a security risk as they may be logged or stored in the database.
Solution:
I implemented a mechanism to propagate ephemeral state (
request_state) from theRunAgentRequestthrough to theInvocationContext. I also added aheader_providertoMcpToolsetthat can dynamically generate headers (e.g.,Authorization: Bearer ...) from this state.Key changes:
request_state: Added toInvocationContextfor ephemeral data that overridessession.statebut is not persistedMcpToolsetConfig: Addedstate_header_mappingto declaratively map state keys to HTTP headerscreate_session_state_header_provider: A utility to generate the header provider functionThis allows clients to pass a JWT in the request payload, have it available to the agent for that request only, and automatically attach it to MCP tool calls.
Testing Plan
Unit Tests:
Summary of pytest results:
tests/unittests/tools/mcp_tool/test_jwt_token_propagation.py: PASSED. Verified header generation, precedence ofrequest_state, and configuration parsing.tests/unittests/agents/test_readonly_context_state.py: PASSED. VerifiedReadonlyContext.statecorrectly merges ephemeral and persistent state.tests/unittests/tools/mcp_tool/test_mcp_toolset.py: PASSED. Verified no regressions in existing toolset functionality.Manual End-to-End (E2E) Tests:
I performed a live verification using a local FastMCP server and a mock LLM agent.
McpToolsetandstate_header_mappingrun_agentrequest withrequest_state={"jwt_token": "test-token-123"}Bearer test-token-123Checklist
Additional context
This feature was designed to prioritize security by ensuring sensitive tokens are not persisted in the session history.
Usage Example:
request_state["jwt_token"](orsession.state["jwt_token"]) as anAuthorization: Bearer <token>header:request_state(recommended for security) orstate_delta: