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]

Directly formatting a non-primitive type can lead to unexpected header values. For robustness, it's good practice to check if the value from the state is a simple scalar type. If not, logging a warning can help the user debug why their headers might not be what they expect, guiding them to pre-serialize complex data.

    if not isinstance(value, (str, int, float, bool)):
        logger.warning(
            'Value for state key "%s" is of type %s, which may not serialize correctly into a header. '
            'Consider pre-serializing complex values or using a different header_format.',
            state_key, type(value).__name__
        )
    formatted_value = header_format.format(value=value)
    return {header_name: formatted_value}

[gemini-code-assist]

The default serialization of non-primitive types to a header value using str(value) can lead to unexpected behavior. For example, if a user stores a dictionary or a custom object in the state, str(value) will produce a representation like '<...>' which is not a useful HTTP header. To improve robustness and prevent silent misconfiguration, I suggest adding a check for the value's type and logging a warning if it's not a simple scalar type. This will guide users to either pre-serialize complex data or use the state_header_format for custom formatting.

            if not isinstance(value, (str, int, float, bool)):
                logger.warning(
                    'Value for state key "%s" is of type %s, which may not serialize correctly into a header. '
                    'Consider pre-serializing complex values or using state_header_format.',
                    state_key, type(value).__name__
                )
            # Apply formatting if specified for this header
            if header_name in state_format:
              formatted_value = state_format[header_name].format(value=value)
            else:
              formatted_value = str(value)
            headers[header_name] = formatted_value

[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.

[gemini-code-assist]

The logic within config_based_header_provider duplicates functionality already present in the new create_session_state_header_provider helper function, particularly around checking value types and logging warnings. To improve maintainability and reduce code duplication, you could refactor this to build a list of individual header providers using create_session_state_header_provider and then create a combined provider that merges their results.

    if mcp_toolset_config.state_header_mapping:
      state_mapping = mcp_toolset_config.state_header_mapping
      state_format = mcp_toolset_config.state_header_format or {}

      providers = []
      for state_key, header_name in state_mapping.items():
        # The default format should be just the value, equivalent to str(value).
        header_format_str = state_format.get(header_name, "{value}")
        providers.append(
            create_session_state_header_provider(
                state_key=state_key,
                header_name=header_name,
                header_format=header_format_str,
                default_value=None,  # Omit header if key is missing
            )
        )

      def combined_header_provider(
          ctx: ReadonlyContext,
      ) -> Dict[str, str]:
        headers = {}
        for provider_func in providers:
          headers.update(provider_func(ctx))
        return headers

      header_provider = combined_header_provider

[chenvaltzer-boop]

+1