adk-python icon indicating copy to clipboard operation
adk-python copied to clipboard

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

Open timof1308 opened this issue 1 month ago • 6 comments

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: #1402
  • Related: #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:

  • [x] I have added or updated unit tests for my change.
  • [x] 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

  • [x] I have read the CONTRIBUTING.md document.
  • [x] I have performed a self-review of my own code.
  • [x] I have commented my code, particularly in hard-to-understand areas.
  • [x] I have added tests that prove my fix is effective or that my feature works.
  • [x] New and existing unit tests pass locally with my changes.
  • [x] I have manually tested my changes end-to-end.
  • [x] 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}"
  1. 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..."
        }
    }
)
  1. 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}"
    )
)

timof1308 avatar Nov 23 '25 13:11 timof1308

+1

chenvaltzer-boop avatar Nov 25 '25 12:11 chenvaltzer-boop

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

  • f5ca9e24 - Added strict mode for type validation
  • e23bd018 - Added state_header_strict config option
  • 124cc14c - Added RFC 7230 compliant validation and sanitization

timof1308 avatar Nov 28 '25 11:11 timof1308

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

ryanaiagent avatar Nov 30 '25 04:11 ryanaiagent

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

the CI is now pending on the latest commit (d0c78291). Once the workflow finishes, all tests should pass. Let me know if anything else is needed!

timof1308 avatar Dec 01 '25 07:12 timof1308

+1

hagaic-ship-it avatar Dec 01 '25 09:12 hagaic-ship-it

+1

eranc-google avatar Dec 01 '25 09:12 eranc-google