feat: Add strict mode to create_session_state_header_provider to raise errors for non-primitive types and skip empty string values.

Author: timof1308

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

@@ -64,61 +64,97 @@
 logger = logging.getLogger("google_adk." + __name__)
 
 
-def _validate_header_value(state_key: str, value: Any) -> None:
-  """Validates that a state value is suitable for use in a header."""
+def _validate_header_value(
+    state_key: str, value: Any, strict: bool = False
+) -> None:
+  """Validates that a state value is suitable for use in a header.
+
+  Args:
+      state_key: The key being validated.
+      value: The value to validate.
+      strict: If True, raises ValueError for non-primitive types.
+
+  Raises:
+      ValueError: If strict=True and value is not a primitive type.
+  """
   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__,
+    msg = (
+        f'Value for state key "{state_key}" is of type'
+        f' {type(value).__name__}, which may not serialize correctly into a'
+        ' header. Consider pre-serializing complex values or using'
+        ' state_header_format.'
     )
+    if strict:
+      raise ValueError(msg)
+    logger.warning(msg)
 
 
 def create_session_state_header_provider(
     state_key: str,
     header_name: str = "Authorization",
     header_format: str = "Bearer {value}",
     default_value: Optional[str] = None,
+    strict: bool = False,
 ) -> HeaderProvider:
   """Creates a header provider that extracts values from session state.
 
   This utility function generates a header_provider callable that can be used
   with McpToolset to automatically extract values from the session state and
   format them as HTTP headers for MCP server connections.
 
+  .. warning::
+      **Security Best Practice**: For sensitive, short-lived tokens like JWTs,
+      use ``request_state`` instead of ``session.state`` to avoid persisting
+      sensitive data to the database. Pass tokens via
+      ``RunAgentRequest.request_state``, which will override ``session.state``
+      for the duration of the request without being persisted.
+
   Args:
-      state_key: The key to look up in session.state.
+      state_key: The key to look up in session.state (or request_state).
       header_name: The HTTP header name to set (default: 'Authorization').
       header_format: Format string for the header value. Use {value} as a
         placeholder for the state value (default: 'Bearer {value}').
       default_value: Default value if state_key is not found in session state.
         If None, the header is omitted when the key is missing.
+      strict: If True, raises ValueError when non-primitive types are
+        encountered. If False (default), logs a warning instead.
 
   Returns:
       A callable that takes a ReadonlyContext and returns a dictionary of
       headers to be used for the MCP session.
 
+  Raises:
+      ValueError: If strict=True and a non-primitive type is found in state.
+
   Example::
 
+      # Example 1: Using request_state for JWT tokens (recommended)
       toolset = McpToolset(
           connection_params=StreamableHTTPConnectionParams(
               url="http://api.example.com/mcp"
           ),
           header_provider=create_session_state_header_provider(
-              state_key="jwt_token",
+              state_key="jwt_token",  # Will read from request_state first
               header_name="Authorization",
               header_format="Bearer {value}"
           )
       )
+
+      # Client sends request with ephemeral JWT
+      response = await agent.run(
+          RunAgentRequest(
+              session_id="user-123",
+              request_state={"jwt_token": "eyJhbG..."}  # Ephemeral, not persisted
+          )
+      )
   """
 
   def provider(ctx: ReadonlyContext) -> Dict[str, str]:
     value = ctx.state.get(state_key, default_value)
-    if value is None:
+    # Skip header if value is None or empty string
+    if value is None or value == "":
       return {}
-    _validate_header_value(state_key, value)
+    _validate_header_value(state_key, value, strict=strict)
     formatted_value = header_format.format(value=value)
     return {header_name: formatted_value}
 

tests/unittests/tools/mcp_tool/test_jwt_token_propagation.py

@@ -129,6 +129,54 @@ def test_none_value_in_state_returns_empty(self):
 
     assert headers == {}
 
+  def test_empty_string_value_returns_empty(self):
+    """Test that empty string value in state returns empty dict."""
+    mock_context = Mock(spec=ReadonlyContext)
+    mock_context.state = {"jwt_token": ""}
+
+    provider = create_session_state_header_provider(state_key="jwt_token")
+
+    headers = provider(mock_context)
+
+    assert headers == {}
+
+  def test_strict_mode_with_primitive_types(self):
+    """Test that strict mode works properly with primitive types."""
+    mock_context = Mock(spec=ReadonlyContext)
+    
+    # Test with string
+    mock_context.state = {"token": "my-token"}
+    provider = create_session_state_header_provider(
+        state_key="token", strict=True
+    )
+    headers = provider(mock_context)
+    assert headers == {"Authorization": "Bearer my-token"}
+    
+    # Test with int
+    mock_context.state = {"count": 42}
+    provider = create_session_state_header_provider(
+        state_key="count", header_name="X-Count", header_format="{value}", strict=True
+    )
+    headers = provider(mock_context)
+    assert headers == {"X-Count": "42"}
+
+  def test_strict_mode_raises_on_non_primitive_types(self):
+    """Test that strict mode raises ValueError for non-primitive types."""
+    mock_context = Mock(spec=ReadonlyContext)
+    mock_context.state = {"complex_data": {"nested": "dict"}}
+
+    provider = create_session_state_header_provider(
+        state_key="complex_data", strict=True
+    )
+
+    with pytest.raises(ValueError) as exc_info:
+      provider(mock_context)
+    
+    assert "complex_data" in str(exc_info.value)
+    assert "dict" in str(exc_info.value)
+    assert "may not serialize correctly" in str(exc_info.value)
+
+
 
 class TestMcpToolsetConfigStateHeaderMapping:
   """Test suite for state_header_mapping configuration."""