PR feedback - mpangrazzi

Author: vblagoje

integrations/mcp/src/haystack_integrations/tools/mcp/mcp_toolset.py

@@ -29,13 +29,17 @@
 logger = logging.getLogger(__name__)
 
 
-def _serialize_state_config(config: dict[str, dict[str, Any]]) -> dict[str, dict[str, Any]] | None:
+def _serialize_state_config(config: dict[str, dict[str, Any]] | None) -> dict[str, dict[str, Any]] | None:
     """
     Serialize a state configuration dictionary, converting any callable handlers to their string representation.
 
     Works for both outputs_to_state (tool_name -> {state_key -> {source, handler}})
     and outputs_to_string (tool_name -> {source, handler}).
 
+    Note: The keys "source" and "handler" are reserved and used internally to distinguish between
+    outputs_to_string format and outputs_to_state format. Do not use these as state keys in
+    outputs_to_state configurations.
+
     :param config: The state configuration dictionary to serialize
     :returns: The serialized configuration dictionary, or None if empty
     """

integrations/mcp/tests/test_mcp_toolset.py

@@ -292,7 +292,7 @@ async def test_toolset_with_state_config(self, calculator_toolset_with_state_con
         assert add_tool.outputs_to_string is not None
         assert subtract_tool.outputs_to_string is None
 
-    async def test_toolset_state_config_serde(self, calculator_toolset_with_state_config):
+    async def test_toolset_state_config_serde(self, calculator_toolset_with_state_config, mcp_tool_cleanup):
         """Test serialization and deserialization of MCPToolset with state configuration."""
         toolset = calculator_toolset_with_state_config
 
@@ -308,20 +308,19 @@ async def test_toolset_state_config_serde(self, calculator_toolset_with_state_co
         # Handler should be serialized as a string
         assert isinstance(toolset_dict["data"]["outputs_to_string"]["add"]["handler"], str)
 
-        # Test deserialization
-        with patch("haystack_integrations.tools.mcp.mcp_toolset.MCPToolset.__init__", return_value=None) as mock_init:
-            MCPToolset.from_dict(toolset_dict)
+        # Test deserialization with full roundtrip
+        new_toolset = MCPToolset.from_dict(toolset_dict)
+        mcp_tool_cleanup(new_toolset)
 
-            mock_init.assert_called_once()
-            _, kwargs = mock_init.call_args
-            assert kwargs["inputs_from_state"] == {
-                "add": {"first_number": "a"},
-                "subtract": {"first_number": "a", "second_number": "b"},
-            }
-            assert "add" in kwargs["outputs_to_state"]
-            assert "add" in kwargs["outputs_to_string"]
-            # Handler should be deserialized back to a callable
-            assert callable(kwargs["outputs_to_string"]["add"]["handler"])
+        # Verify state configs are correctly deserialized
+        assert new_toolset.inputs_from_state == {
+            "add": {"first_number": "a"},
+            "subtract": {"first_number": "a", "second_number": "b"},
+        }
+        assert "add" in new_toolset.outputs_to_state
+        assert "add" in new_toolset.outputs_to_string
+        # Handler should be deserialized back to a callable
+        assert callable(new_toolset.outputs_to_string["add"]["handler"])
 
     async def test_toolset_state_config_unknown_tool_warning(self, caplog):
         """Test that a warning is logged when state config references unknown tools.