feat: Add Agent state-mapping parameters to MCPTool (#2501)

* Add Agent state-mapping parameters to MCPTool

* PR feedback, add Agent state MCPTool integration test

Author: vblagoje

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

@@ -845,6 +845,11 @@ class MCPTool(Tool):
     - The JSON contains the structured response from the MCP server
     - Use json.loads() to parse the response into a dictionary
 
+    State-mapping support:
+    - MCPTool supports state-mapping parameters (`outputs_to_string`, `inputs_from_state`, `outputs_to_state`)
+    - These enable integration with Agent state for automatic parameter injection and output handling
+    - See the `__init__` method documentation for details on each parameter
+
     Example using Streamable HTTP:
     ```python
     import json
@@ -902,6 +907,9 @@ def __init__(
         connection_timeout: int = 30,
         invocation_timeout: int = 30,
         eager_connect: bool = False,
+        outputs_to_string: dict[str, Any] | None = None,
+        inputs_from_state: dict[str, str] | None = None,
+        outputs_to_state: dict[str, dict[str, Any]] | None = None,
     ):
         """
         Initialize the MCP tool.
@@ -914,6 +922,17 @@ def __init__(
         :param eager_connect: If True, connect to server during initialization.
                              If False (default), defer connection until warm_up or first tool use,
                              whichever comes first.
+        :param outputs_to_string: Optional dictionary defining how tool outputs should be converted into a string.
+                                 If the source is provided only the specified output key is sent to the handler.
+                                 If the source is omitted the whole tool result is sent to the handler.
+                                 Example: `{"source": "docs", "handler": my_custom_function}`
+        :param inputs_from_state: Optional dictionary mapping state keys to tool parameter names.
+                                 Example: `{"repository": "repo"}` maps state's "repository" to tool's "repo" parameter.
+        :param outputs_to_state: Optional dictionary defining how tool outputs map to keys within state as well as
+                                optional handlers. If the source is provided only the specified output key is sent
+                                to the handler.
+                                Example with source: `{"documents": {"source": "docs", "handler": custom_handler}}`
+                                Example without source: `{"documents": {"handler": custom_handler}}`
         :raises MCPConnectionError: If connection to the server fails
         :raises MCPToolNotFoundError: If no tools are available or the requested tool is not found
         :raises TimeoutError: If connection times out
@@ -924,6 +943,9 @@ def __init__(
         self._connection_timeout = connection_timeout
         self._invocation_timeout = invocation_timeout
         self._eager_connect = eager_connect
+        self._outputs_to_string = outputs_to_string
+        self._inputs_from_state = inputs_from_state
+        self._outputs_to_state = outputs_to_state
         self._client: MCPClient | None = None
         self._worker: _MCPClientSessionManager | None = None
         self._lock = threading.RLock()
@@ -934,7 +956,15 @@ def __init__(
             # without discovering the remote schema during validation.
             # Tool parameters/schema will be replaced with the correct schema (from the MCP server) on first use.
             params = {"type": "object", "properties": {}, "additionalProperties": True}
-            super().__init__(name=name, description=description or "", parameters=params, function=self._invoke_tool)
+            super().__init__(
+                name=name,
+                description=description or "",
+                parameters=params,
+                function=self._invoke_tool,
+                outputs_to_string=outputs_to_string,
+                inputs_from_state=inputs_from_state,
+                outputs_to_state=outputs_to_state,
+            )
             return
 
         logger.debug(f"TOOL: Initializing MCPTool '{name}'")
@@ -950,7 +980,19 @@ def __init__(
                 description=description or tool_info.description or "",
                 parameters=tool_info.inputSchema,
                 function=self._invoke_tool,
+                outputs_to_string=outputs_to_string,
+                inputs_from_state=inputs_from_state,
+                outputs_to_state=outputs_to_state,
             )
+
+            # Remove inputs_from_state keys from parameters schema if present
+            # This matches the behavior of ComponentTool
+            if inputs_from_state and "properties" in self.parameters:
+                for key in inputs_from_state.values():
+                    self.parameters["properties"].pop(key, None)
+                    if "required" in self.parameters and key in self.parameters["required"]:
+                        self.parameters["required"].remove(key)
+
             logger.debug(f"TOOL: Initialization complete for '{name}'")
 
         except Exception as e:
@@ -1069,13 +1111,21 @@ def warm_up(self) -> None:
             tool = self._connect_and_initialize(self.name)
             self.parameters = tool.inputSchema
 
+            # Remove inputs_from_state keys from parameters schema if present
+            # This matches the behavior of ComponentTool
+            if self._inputs_from_state and "properties" in self.parameters:
+                for key in self._inputs_from_state.values():
+                    self.parameters["properties"].pop(key, None)
+                    if "required" in self.parameters and key in self.parameters["required"]:
+                        self.parameters["required"].remove(key)
+
     def to_dict(self) -> dict[str, Any]:
         """
         Serializes the MCPTool to a dictionary.
 
         The serialization preserves all information needed to recreate the tool,
-        including server connection parameters and timeout settings. Note that the
-        active connection is not maintained.
+        including server connection parameters, timeout settings, and state-mapping parameters.
+        Note that the active connection is not maintained.
 
         :returns: Dictionary with serialized data in the format:
                   `{"type": fully_qualified_class_name, "data": {parameters}}`
@@ -1087,6 +1137,9 @@ def to_dict(self) -> dict[str, Any]:
             "connection_timeout": self._connection_timeout,
             "invocation_timeout": self._invocation_timeout,
             "eager_connect": self._eager_connect,
+            "outputs_to_string": self._outputs_to_string,
+            "inputs_from_state": self._inputs_from_state,
+            "outputs_to_state": self._outputs_to_state,
         }
         return {
             "type": generate_qualified_class_name(type(self)),
@@ -1099,8 +1152,8 @@ def from_dict(cls, data: dict[str, Any]) -> "Tool":
         Deserializes the MCPTool from a dictionary.
 
         This method reconstructs an MCPTool instance from a serialized dictionary,
-        including recreating the server_info object. A new connection will be established
-        to the MCP server during initialization.
+        including recreating the server_info object and state-mapping parameters.
+        A new connection will be established to the MCP server during initialization.
 
         :param data: Dictionary containing serialized tool data
         :returns: A fully initialized MCPTool instance
@@ -1121,6 +1174,11 @@ def from_dict(cls, data: dict[str, Any]) -> "Tool":
         invocation_timeout = inner_data.get("invocation_timeout", 30)
         eager_connect = inner_data.get("eager_connect", False)  # because False is the default
 
+        # Handle state-mapping parameters
+        outputs_to_string = inner_data.get("outputs_to_string")
+        inputs_from_state = inner_data.get("inputs_from_state")
+        outputs_to_state = inner_data.get("outputs_to_state")
+
         # Create a new MCPTool instance with the deserialized parameters
         # This will establish a new connection to the MCP server
         return cls(
@@ -1130,6 +1188,9 @@ def from_dict(cls, data: dict[str, Any]) -> "Tool":
             connection_timeout=connection_timeout,
             invocation_timeout=invocation_timeout,
             eager_connect=eager_connect,
+            outputs_to_string=outputs_to_string,
+            inputs_from_state=inputs_from_state,
+            outputs_to_state=outputs_to_state,
         )
 
     def close(self):

integrations/mcp/tests/test_mcp_tool.py

@@ -133,6 +133,69 @@ def test_mcp_tool_serde(self, mcp_tool_cleanup):
 
         assert isinstance(new_tool._server_info, InMemoryServerInfo)
 
+    def test_mcp_tool_state_mapping_parameters(self, mcp_tool_cleanup):
+        """Test that MCPTool correctly initializes with state-mapping parameters."""
+        server_info = InMemoryServerInfo(server=calculator_mcp._mcp_server)
+
+        # Create tool with state-mapping parameters
+        # Map state key "state_a" to tool parameter "a"
+        tool = MCPTool(
+            name="add",
+            server_info=server_info,
+            eager_connect=False,
+            outputs_to_string={"source": "result", "handler": str},
+            inputs_from_state={"state_a": "a"},
+            outputs_to_state={"result": {"source": "output", "handler": str}},
+        )
+        mcp_tool_cleanup(tool)
+
+        # Verify the parameters are stored correctly
+        assert tool._outputs_to_string == {"source": "result", "handler": str}
+        assert tool._inputs_from_state == {"state_a": "a"}
+        assert tool._outputs_to_state == {"result": {"source": "output", "handler": str}}
+
+        # Warm up the tool to trigger schema adjustment
+        tool.warm_up()
+
+        # Verify that "a" was removed from parameters since it's in inputs_from_state
+        assert "a" not in tool.parameters["properties"]
+        assert "a" not in tool.parameters.get("required", [])
+        # Verify that "b" is still present (not removed)
+        assert "b" in tool.parameters["properties"]
+        assert "b" in tool.parameters["required"]
+
+    def test_mcp_tool_serde_with_state_mapping(self, mcp_tool_cleanup):
+        """Test serialization and deserialization of MCPTool with state-mapping parameters."""
+        server_info = InMemoryServerInfo(server=calculator_mcp._mcp_server)
+
+        # Create tool with state-mapping parameters
+        tool = MCPTool(
+            name="add",
+            server_info=server_info,
+            eager_connect=False,
+            outputs_to_string={"source": "result"},
+            inputs_from_state={"filter": "query_filter"},
+            outputs_to_state={"result": {"source": "output"}},
+        )
+        mcp_tool_cleanup(tool)
+
+        # Test serialization (to_dict)
+        tool_dict = tool.to_dict()
+
+        # Verify state-mapping parameters are serialized
+        assert tool_dict["data"]["outputs_to_string"] == {"source": "result"}
+        assert tool_dict["data"]["inputs_from_state"] == {"filter": "query_filter"}
+        assert tool_dict["data"]["outputs_to_state"] == {"result": {"source": "output"}}
+
+        # Test deserialization (from_dict)
+        new_tool = MCPTool.from_dict(tool_dict)
+        mcp_tool_cleanup(new_tool)
+
+        # Verify state-mapping parameters are restored
+        assert new_tool._outputs_to_string == {"source": "result"}
+        assert new_tool._inputs_from_state == {"filter": "query_filter"}
+        assert new_tool._outputs_to_state == {"result": {"source": "output"}}
+
     @pytest.mark.skipif("OPENAI_API_KEY" not in os.environ, reason="OPENAI_API_KEY not set")
     @pytest.mark.integration
     def test_pipeline_warmup_with_mcp_tool(self):
@@ -155,3 +218,54 @@ def test_pipeline_warmup_with_mcp_tool(self):
         finally:
             if tool:
                 tool.close()
+
+    @pytest.mark.skipif("OPENAI_API_KEY" not in os.environ, reason="OPENAI_API_KEY not set")
+    @pytest.mark.integration
+    def test_agent_with_state_mapping(self):
+        """Test Agent with MCPTool using state-mapping to inject location from state."""
+
+        # Create MCPTool with state-mapping that injects home_city from state as timezone parameter
+        server_info = StdioServerInfo(command="uvx", args=["mcp-server-time", "--local-timezone=Europe/Berlin"])
+        tool = MCPTool(
+            name="get_current_time",
+            server_info=server_info,
+            inputs_from_state={"home_city": "timezone"},  # Inject home_city from state as timezone
+        )
+
+        try:
+            # Build Agent with state schema that includes home_city
+            agent = Agent(
+                chat_generator=OpenAIChatGenerator(model="gpt-4o-mini"),
+                tools=[tool],
+                state_schema={"home_city": {"type": str}},
+            )
+            pipeline = Pipeline()
+            pipeline.add_component("agent", agent)
+
+            # Ask for time without mentioning the location - it should use home_city from state
+            user_input_msg = ChatMessage.from_user(text="What time is it at home?")
+            result = pipeline.run(
+                {
+                    "agent": {
+                        "messages": [user_input_msg],
+                        "home_city": "America/New_York",  # Inject New York as home city
+                    }
+                }
+            )
+
+            # Verify the agent got the time for New York
+            final_message = result["agent"]["messages"][-1].text
+
+            # The response should mention time
+            assert any(keyword in final_message.lower() for keyword in ["time", "o'clock", "am", "pm"]), (
+                f"Expected time in response: {final_message}"
+            )
+
+            # Verify the response mentions New York or Eastern timezone (proving state-mapping injected it)
+            # The user never mentioned location, but timezone info should appear in the response
+            assert any(keyword in final_message for keyword in ["New York", "New_York"]), (
+                f"Expected timezone reference (New York) to confirm state-mapping: {final_message}"
+            )
+        finally:
+            if tool:
+                tool.close()