feat(mcp): Add session state-based JWT token propagation for MCP tools

[timof1308]

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):

• Closes: MCP Tools: No mechanism to pass user JWT token through ADK session.state to MCP server context #1402
• Related: feat: Dynamic authentication handling in MCPToolset from ADK Readonly… #2208

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 the RunAgentRequest through to the
InvocationContext. I also added a header_provider to McpToolset that can dynamically generate headers (e.g., Authorization: Bearer ...) from this state.

Key changes:

• request_state: Added to InvocationContext for ephemeral data that overrides session.state but is not persisted
• McpToolset Config: Added state_header_mapping to declaratively map state keys to HTTP headers
• create_session_state_header_provider: A utility to generate the header provider function

This 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:

• I have added or updated unit tests for my change.
• All unit tests pass locally.

Summary of pytest results:

• tests/unittests/tools/mcp_tool/test_jwt_token_propagation.py: PASSED. Verified header generation, precedence of request_state, and configuration parsing.
• tests/unittests/agents/test_readonly_context_state.py: PASSED. Verified ReadonlyContext.state correctly 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.

• Setup:
  • Started a local FastMCP server hat echoes the Authorization header
  • Ran an ADK agent configured with McpToolset and state_header_mapping
• Test:
  • Sent a run_agent request with request_state={"jwt_token": "test-token-123"}
  • The agent called the MCP tool
• Result:
  • The MCP server received the header Authorization: Bearer test-token-123
  • The agent successfully retrieved this value from the tool, confirming propagation

Checklist

• I have read the CONTRIBUTING.md document.
• I have performed a self-review of my own code.
• I have commented my code, particularly in hard-to-understand areas.
• I have added tests that prove my fix is effective or that my feature works.
• New and existing unit tests pass locally with my changes.
• I have manually tested my changes end-to-end.
• Any dependent changes have been merged and published in downstream modules.

Additional context

This feature was designed to prioritize security by ensuring sensitive tokens are not persisted in the session history.

Usage Example:

1. Configuration (YAML): To propagate a JWT token stored in request_state["jwt_token"] (or session.state["jwt_token"]) as an Authorization: Bearer <token> header:

tools:
  - name: google.adk.tools.mcp_tool.McpToolset
    args:
      streamable_http_connection_params:
        url: http://api.example.com/mcp
      state_header_mapping:
        jwt_token: Authorization
      state_header_format:
        Authorization: "Bearer {value}"

2. Client-Side Usage: When running an agent, pass the JWT token in request_state (recommended for security) or state_delta:

# Option A: Ephemeral (Secure) - Token NOT persisted
requests.post(
    "/run",
    json={
        "app_name": "my_app",
        "user_id": "user123",
        "session_id": "session_id",
        "new_message": {"parts": [{"text": "query"}]},
        "request_state": {
            "jwt_token": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9..."
        }
    }
)

# Option B: Persistent - Token saved to session history
requests.post(
    "/run",
    json={
        "app_name": "my_app",
        "user_id": "user123",
        "session_id": "session_id",
        "new_message": {"parts": [{"text": "query"}]},
        "state_delta": {
            "jwt_token": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9..."
        }
    }
)

3. Configuration (Python): Alternatively, you can configure the toolset programmatically in your agent definition:

from google.adk.tools.mcp_tool import create_session_state_header_provider, McpToolset

# Using helper function
toolset = McpToolset(
    connection_params=StreamableHTTPConnectionParams(
        url='http://api.example.com/mcp'
    ),
    header_provider=create_session_state_header_provider(
        state_key="jwt_token",
        header_name="Authorization",
        header_format="Bearer {value}"
    )
)

[gemini-code-assist]

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.

[gemini-code-assist]

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.

[chenvaltzer-boop]

+1

[timof1308]

All bot review comments have been addressed through the following commits:

• f5ca9e2 - Added strict mode for type validation
• e23bd01 - Added state_header_strict config option
• 124cc14 - Added RFC 7230 compliant validation and sanitization

[ryanaiagent]

Hi @timof1308 , Thank you for your work on this pull request. We appreciate the effort you've invested.
Before we can proceed with the review can you please fix the failing unit tests and lint errors

[timof1308]

thank you @ryanaiagent
I have fixed the imports and the failing unit tests

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!

[hagaic-ship-it]

+1

[eranc-google]

+1

[ryanaiagent]

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.

[timof1308]

hi @ryanaiagent, rebased with upstream main and reapplied ./autoformat.sh ✅