Skip to content

Commit e23bd01

Browse files
committed
feat: Introduce state_header_strict config for strict state header type validation and add a test.
1 parent f5ca9e2 commit e23bd01

File tree

2 files changed

+40
-0
lines changed

2 files changed

+40
-0
lines changed

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -367,6 +367,7 @@ def from_config(
367367
header_name=header_name,
368368
header_format=state_format.get(header_name, "{value}"),
369369
default_value=None,
370+
strict=mcp_toolset_config.state_header_strict,
370371
)
371372
for state_key, header_name in state_mapping.items()
372373
]
@@ -450,6 +451,21 @@ class McpToolsetConfig(BaseToolConfig):
450451
used as-is.
451452
"""
452453

454+
state_header_strict: bool = False
455+
"""Enable strict type validation for state header values.
456+
457+
When True, raises ValueError if state values are non-primitive types
458+
(not str, int, float, or bool). This helps catch configuration errors
459+
early by preventing accidental serialization of complex objects into headers.
460+
461+
When False (default), non-primitive types trigger a warning but are still
462+
formatted into headers.
463+
464+
Example::
465+
466+
state_header_strict: true # Raises error on non-primitive types
467+
"""
468+
453469
@model_validator(mode="after")
454470
def _check_only_one_params_field(self):
455471
param_fields = [

tests/unittests/tools/mcp_tool/test_jwt_token_propagation.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -318,3 +318,27 @@ def test_from_config_no_state_mapping_no_provider(self):
318318

319319
# No header provider should be created
320320
assert toolset._header_provider is None
321+
322+
def test_from_config_with_strict_mode(self):
323+
"""Test that from_config respects state_header_strict setting."""
324+
from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
325+
from google.adk.tools.tool_configs import ToolArgsConfig
326+
327+
config = ToolArgsConfig(
328+
stdio_server_params={"command": "test_command", "args": []},
329+
state_header_mapping={"data": "X-Data"},
330+
state_header_strict=True, # Enable strict mode
331+
)
332+
333+
toolset = McpToolset.from_config(config, "/fake/path")
334+
335+
# Test with non-primitive type - should raise ValueError
336+
mock_context = Mock(spec=ReadonlyContext)
337+
mock_context.state = {"data": {"nested": "object"}}
338+
339+
with pytest.raises(ValueError) as exc_info:
340+
toolset._header_provider(mock_context)
341+
342+
assert "data" in str(exc_info.value)
343+
assert "dict" in str(exc_info.value)
344+

0 commit comments

Comments
 (0)