-
Notifications
You must be signed in to change notification settings - Fork 2.5k
feat(mcp): Add session state-based JWT token propagation for MCP tools #3675
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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 |
…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.
c6d2245 to
d211763
Compare
Please ensure you have read the contribution guide before creating a pull request.
Link to Issue or Description of Change
1. Link to an existing issue (if applicable):
2. Or, if no issue exists, describe the change:
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: