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

contributing/samples/gepa/experiment.py

@@ -43,7 +43,6 @@
 from tau_bench.types import EnvRunResult
 from tau_bench.types import RunConfig
 import tau_bench_agent as tau_bench_agent_lib
-
 import utils
 
 

contributing/samples/gepa/run_experiment.py

@@ -25,7 +25,6 @@
 from absl import flags
 import experiment
 from google.genai import types
-
 import utils
 
 _OUTPUT_DIR = flags.DEFINE_string(

src/google/adk/agents/invocation_context.py

@@ -168,6 +168,14 @@ class InvocationContext(BaseModel):
   agent_states: dict[str, dict[str, Any]] = Field(default_factory=dict)
   """The state of the agent for this invocation."""
 
+  request_state: dict[str, Any] = Field(default_factory=dict)
+  """The ephemeral state of the request.
+
+  This state is not persisted to the session and is only available for the
+  current invocation. It is used to pass sensitive information like tokens
+  that should not be stored in the session state.
+  """
+
   end_of_agents: dict[str, bool] = Field(default_factory=dict)
   """The end of agent status for each agent in this invocation."""
 

src/google/adk/agents/readonly_context.py

@@ -14,8 +14,10 @@
 
 from __future__ import annotations
 
+from collections import ChainMap
 from types import MappingProxyType
 from typing import Any
+from typing import Mapping
 from typing import Optional
 from typing import TYPE_CHECKING
 
@@ -51,9 +53,14 @@ def agent_name(self) -> str:
     return self._invocation_context.agent.name
 
   @property
-  def state(self) -> MappingProxyType[str, Any]:
+  def state(self) -> Mapping[str, Any]:
     """The state of the current session. READONLY field."""
-    return MappingProxyType(self._invocation_context.session.state)
+    return MappingProxyType(
+        ChainMap(
+            self._invocation_context.request_state,
+            self._invocation_context.session.state,
+        )
+    )
 
   @property
   def session(self) -> Session:

src/google/adk/cli/adk_web_server.py

@@ -175,6 +175,7 @@ class RunAgentRequest(common.BaseModel):
   new_message: types.Content
   streaming: bool = False
   state_delta: Optional[dict[str, Any]] = None
+  request_state: Optional[dict[str, Any]] = None
   # for resume long running functions
   invocation_id: Optional[str] = None
 
@@ -1463,6 +1464,7 @@ async def run_agent(req: RunAgentRequest) -> list[Event]:
               session_id=req.session_id,
               new_message=req.new_message,
               state_delta=req.state_delta,
+              request_state=req.request_state,
           )
       ) as agen:
         events = [event async for event in agen]
@@ -1492,6 +1494,7 @@ async def event_generator():
                   session_id=req.session_id,
                   new_message=req.new_message,
                   state_delta=req.state_delta,
+                  request_state=req.request_state,
                   run_config=RunConfig(streaming_mode=stream_mode),
                   invocation_id=req.invocation_id,
               )

src/google/adk/runners.py

@@ -398,6 +398,7 @@ async def run_async(
       invocation_id: Optional[str] = None,
       new_message: Optional[types.Content] = None,
       state_delta: Optional[dict[str, Any]] = None,
+      request_state: Optional[dict[str, Any]] = None,
       run_config: Optional[RunConfig] = None,
   ) -> AsyncGenerator[Event, None]:
     """Main entry method to run the agent in this runner.
@@ -415,6 +416,7 @@ async def run_async(
         interrupted invocation.
       new_message: A new message to append to the session.
       state_delta: Optional state changes to apply to the session.
+      request_state: Optional ephemeral state for the request.
       run_config: The run config for the agent.
 
     Yields:
@@ -462,6 +464,7 @@ async def _run_with_trace(
               invocation_id=invocation_id,
               run_config=run_config,
               state_delta=state_delta,
+              request_state=request_state,
           )
           if invocation_context.end_of_agents.get(
               invocation_context.agent.name
@@ -475,6 +478,7 @@ async def _run_with_trace(
               new_message=new_message,  # new_message is not None.
               run_config=run_config,
               state_delta=state_delta,
+              request_state=request_state,
           )
 
         async def execute(ctx: InvocationContext) -> AsyncGenerator[Event]:
@@ -1185,6 +1189,7 @@ async def _setup_context_for_new_invocation(
       new_message: types.Content,
       run_config: RunConfig,
       state_delta: Optional[dict[str, Any]],
+      request_state: Optional[dict[str, Any]] = None,
   ) -> InvocationContext:
     """Sets up the context for a new invocation.
 
@@ -1193,6 +1198,7 @@ async def _setup_context_for_new_invocation(
       new_message: The new message to process and append to the session.
       run_config: The run config of the agent.
       state_delta: Optional state changes to apply to the session.
+      request_state: Optional ephemeral state for the request.
 
     Returns:
       The invocation context for the new invocation.
@@ -1202,6 +1208,7 @@ async def _setup_context_for_new_invocation(
         session,
         new_message=new_message,
         run_config=run_config,
+        request_state=request_state,
     )
     # Step 2: Handle new message, by running callbacks and appending to
     # session.
@@ -1224,6 +1231,7 @@ async def _setup_context_for_resumed_invocation(
       invocation_id: Optional[str],
       run_config: RunConfig,
       state_delta: Optional[dict[str, Any]],
+      request_state: Optional[dict[str, Any]] = None,
   ) -> InvocationContext:
     """Sets up the context for a resumed invocation.
 
@@ -1233,6 +1241,7 @@ async def _setup_context_for_resumed_invocation(
       invocation_id: The invocation id to resume.
       run_config: The run config of the agent.
       state_delta: Optional state changes to apply to the session.
+      request_state: Optional ephemeral state for the request.
 
     Returns:
       The invocation context for the resumed invocation.
@@ -1258,6 +1267,7 @@ async def _setup_context_for_resumed_invocation(
         new_message=user_message,
         run_config=run_config,
         invocation_id=invocation_id,
+        request_state=request_state,
     )
     # Step 3: Maybe handle new message.
     if new_message:
@@ -1302,6 +1312,7 @@ def _new_invocation_context(
       new_message: Optional[types.Content] = None,
       live_request_queue: Optional[LiveRequestQueue] = None,
       run_config: Optional[RunConfig] = None,
+      request_state: Optional[dict[str, Any]] = None,
   ) -> InvocationContext:
     """Creates a new invocation context.
 
@@ -1311,6 +1322,7 @@ def _new_invocation_context(
         new_message: The new message for the context.
         live_request_queue: The live request queue for the context.
         run_config: The run config for the context.
+        request_state: The ephemeral state for the request.
 
     Returns:
         The new invocation context.
@@ -1342,6 +1354,7 @@ def _new_invocation_context(
         live_request_queue=live_request_queue,
         run_config=run_config,
         resumability_config=self.resumability_config,
+        request_state=request_state or {},
     )
 
   def _new_invocation_context_for_live(

src/google/adk/tools/mcp_tool/__init__.py

@@ -22,6 +22,7 @@
   from .mcp_session_manager import StreamableHTTPConnectionParams
   from .mcp_tool import MCPTool
   from .mcp_tool import McpTool
+  from .mcp_toolset import create_session_state_header_provider
   from .mcp_toolset import MCPToolset
   from .mcp_toolset import McpToolset
 
@@ -32,6 +33,7 @@
       'MCPTool',
       'McpToolset',
       'MCPToolset',
+      'create_session_state_header_provider',
       'SseConnectionParams',
       'StdioConnectionParams',
       'StreamableHTTPConnectionParams',