fix: fix mcp types + add py.typed (#2248)

* fix: fix mcp types + add py.typed

* Star expression is from 3.11, use alternative

* Update CI to use test:types

* PR feedback

* PR feedback

---------

Co-authored-by: David S. Batista <[email protected]>

Author: vblagoje

.github/workflows/mcp.yml

@@ -65,7 +65,7 @@ jobs:
 
       - name: Lint
         if: matrix.python-version == '3.10' && runner.os == 'Linux'
-        run: hatch run fmt-check && hatch run lint:typing
+        run: hatch run fmt-check && hatch run test:types
 
       - name: Generate docs
         if: matrix.python-version == '3.10' && runner.os == 'Linux'

integrations/mcp/pyproject.toml

@@ -78,18 +78,13 @@ unit = 'pytest -m "not integration" {args:tests}'
 integration = 'pytest -m "integration" {args:tests}'
 all = 'pytest {args:tests}'
 cov-retry = 'all --cov=haystack_integrations --reruns 3 --reruns-delay 30 -x'
-types = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
+types = """mypy -p haystack_integrations.tools.mcp {args}"""
 
-# TODO: remove lint environment once this integration is properly typed
-# test environment should be used instead
-# https://github.com/deepset-ai/haystack-core-integrations/issues/1771
-[tool.hatch.envs.lint]
-installer = "uv"
-detached = true
-dependencies = ["pip", "mypy>=1.0.0", "ruff>=0.0.243"]
-
-[tool.hatch.envs.lint.scripts]
-typing = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
+[tool.mypy]
+install_types = true
+non_interactive = true
+check_untyped_defs = true
+disallow_incomplete_defs = true
 
 
 [tool.ruff]
@@ -167,19 +162,6 @@ omit = ["*/tests/*", "*/__init__.py"]
 show_missing = true
 exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"]
 
-[[tool.mypy.overrides]]
-module = [
-    "haystack.*",
-    "haystack_integrations.*",
-    "pytest.*",
-    "pytest_asyncio",
-    "anyio.*",
-    "mcp.*",
-    "mcp",
-    "httpx",
-    "exceptiongroup"
-]
-ignore_missing_imports = true
 
 [tool.pytest.ini_options]
 addopts = "--strict-markers"

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

@@ -29,6 +29,7 @@
 from mcp.client.sse import sse_client
 from mcp.client.stdio import stdio_client
 from mcp.client.streamable_http import streamablehttp_client
+from mcp.shared.message import SessionMessage
 
 logger = logging.getLogger(__name__)
 
@@ -121,7 +122,7 @@ async def _coroutine_with_stop_event():
         # use it to control the coroutine.
         return future, stop_event_promise.result(timeout)
 
-    def shutdown(self, timeout: float = 2):
+    def shutdown(self, timeout: float = 2) -> None:
         """
         Shut down the background event loop and thread.
 
@@ -208,14 +209,14 @@ class MCPClient(ABC):
     def __init__(self, max_retries: int = 3, base_delay: float = 1.0, max_delay: float = 30.0) -> None:
         self.session: ClientSession | None = None
         self.exit_stack: AsyncExitStack = AsyncExitStack()
-        self.stdio: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception] | None = None
-        self.write: MemoryObjectSendStream[types.JSONRPCMessage] | None = None
+        self.stdio: MemoryObjectReceiveStream[SessionMessage | Exception] | None = None
+        self.write: MemoryObjectSendStream[SessionMessage] | None = None
         self.max_retries = max_retries
         self.base_delay = base_delay
         self.max_delay = max_delay
 
     @abstractmethod
-    async def connect(self) -> list[Tool]:
+    async def connect(self) -> list[types.Tool]:
         """
         Connect to an MCP server.
 
@@ -329,10 +330,16 @@ async def aclose(self) -> None:
     async def _initialize_session_with_transport(
         self,
         transport_tuple: tuple[
-            MemoryObjectReceiveStream[types.JSONRPCMessage | Exception], MemoryObjectSendStream[types.JSONRPCMessage]
+            MemoryObjectReceiveStream[SessionMessage | Exception],
+            MemoryObjectSendStream[SessionMessage],
+        ]
+        | tuple[
+            MemoryObjectReceiveStream[SessionMessage | Exception],
+            MemoryObjectSendStream[SessionMessage],
+            Any,
         ],
         connection_type: str,
-    ) -> list[Tool]:
+    ) -> list[types.Tool]:
         """
         Common session initialization logic for all transports.
 
@@ -390,11 +397,11 @@ def __init__(
         self.env: dict[str, str] | None = None
         if env:
             self.env = {
-                key: value.resolve_value() if isinstance(value, Secret) else value for key, value in env.items()
+                key: (value.resolve_value() if isinstance(value, Secret) else value) or "" for key, value in env.items()
             }
         logger.debug(f"PROCESS: Created StdioClient for command: {command} {' '.join(self.args or [])}")
 
-    async def connect(self) -> list[Tool]:
+    async def connect(self) -> list[types.Tool]:
         """
         Connect to an MCP server using stdio transport.
 
@@ -433,7 +440,7 @@ def __init__(
         )
         self.timeout: int = server_info.timeout
 
-    async def connect(self) -> list[Tool]:
+    async def connect(self) -> list[types.Tool]:
         """
         Connect to an MCP server using SSE transport.
 
@@ -481,7 +488,7 @@ def __init__(
         )
         self.timeout: int = server_info.timeout
 
-    async def connect(self) -> list[Tool]:
+    async def connect(self) -> list[types.Tool]:
         """
         Connect to an MCP server using streamable HTTP transport.
 
@@ -526,7 +533,7 @@ def to_dict(self) -> dict[str, Any]:
         :returns: Dictionary representation of this server info
         """
         # Store the fully qualified class name for deserialization
-        result = {"type": generate_qualified_class_name(type(self))}
+        result: dict[str, Any] = {"type": generate_qualified_class_name(type(self))}
 
         # Add all fields from the dataclass
         for dataclass_field in fields(self):
@@ -629,7 +636,7 @@ def __post_init__(self):
             # from now on only use url for the lifetime of the SSEServerInfo instance, never base_url
             self.url = f"{self.base_url.rstrip('/')}/sse"
 
-        elif not is_valid_http_url(self.url):
+        elif self.url and not is_valid_http_url(self.url):
             message = f"Invalid url: {self.url}"
             raise ValueError(message)
 
@@ -834,7 +841,7 @@ def __init__(
             tool_dict = {t.name: t for t in tools}
             logger.debug(f"TOOL: Available tools: {list(tool_dict.keys())}")
 
-            tool_info = tool_dict.get(name)
+            tool_info: types.Tool | None = tool_dict.get(name)
 
             if not tool_info:
                 available = list(tool_dict.keys())
@@ -846,7 +853,7 @@ def __init__(
             # Initialize the parent class
             super().__init__(
                 name=name,
-                description=description or tool_info.description,
+                description=description or tool_info.description or "",
                 parameters=tool_info.inputSchema,
                 function=self._invoke_tool,
             )
@@ -971,7 +978,7 @@ def from_dict(cls, data: dict[str, Any]) -> "Tool":
         # First get the appropriate class by name
         server_info_class = import_class_by_name(server_info_dict["type"])
         # Then deserialize using that class's from_dict method
-        server_info = server_info_class.from_dict(server_info_dict)
+        server_info = cast(MCPServerInfo, server_info_class).from_dict(server_info_dict)
 
         # Handle backward compatibility for timeout parameters
         connection_timeout = inner_data.get("connection_timeout", 30)
@@ -1027,7 +1034,7 @@ def __init__(self, client: "MCPClient", *, timeout: float | None = None):
         self.executor = AsyncExecutor.get_instance()
 
         # Where the tool list (or an exception) will be delivered.
-        self._tools_promise: Future[list[Tool]] = Future()
+        self._tools_promise: Future[list[types.Tool]] = Future()
 
         # Kick off the worker coroutine in the background loop
         self._worker_future, self._stop_event = self.executor.run_background(self._run, timeout=None)
@@ -1040,15 +1047,15 @@ def __init__(self, client: "MCPClient", *, timeout: float | None = None):
             self.stop()
             raise
 
-    def tools(self) -> list[Tool]:
+    def tools(self) -> list[types.Tool]:
         """Return the tool list already collected during startup."""
 
         return self._tools_promise.result()
 
     def stop(self) -> None:
         """Request the worker to shut down and block until done."""
 
-        def _set(ev: asyncio.Event):
+        def _set(ev: asyncio.Event) -> None:
             if not ev.is_set():
                 ev.set()
 
@@ -1067,7 +1074,7 @@ def _set(ev: asyncio.Event):
             logger.debug(f"Error during worker future result: {e}")
             pass
 
-    async def _run(self, stop_event: asyncio.Event):
+    async def _run(self, stop_event: asyncio.Event) -> None:
         """Background coroutine living in AsyncExecutor's loop."""
 
         try:

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

@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from collections.abc import Callable
-from typing import Any
+from typing import Any, cast
 from urllib.parse import urlparse
 
 import httpx
@@ -155,7 +155,7 @@ def create_invoke_tool(
             ) -> Callable[..., Any]:
                 """Return a closure that keeps a strong reference to *owner_toolset* alive."""
 
-                def invoke_tool(**kwargs) -> Any:
+                def invoke_tool(**kwargs: Any) -> Any:
                     _ = owner_toolset  # strong reference so GC can't collect the toolset too early
                     return AsyncExecutor.get_instance().run(
                         mcp_client.call_tool(tool_name, kwargs), timeout=tool_timeout
@@ -176,7 +176,7 @@ def invoke_tool(**kwargs) -> Any:
                 # Use the helper function to create the invoke_tool function
                 tool = Tool(
                     name=tool_info.name,
-                    description=tool_info.description,
+                    description=tool_info.description or "",
                     parameters=tool_info.inputSchema,
                     function=create_invoke_tool(self, client, tool_info.name, self.invocation_timeout),
                 )
@@ -273,7 +273,7 @@ def from_dict(cls, data: dict[str, Any]) -> "MCPToolset":
         # Reconstruct the server_info object
         server_info_dict = inner_data.get("server_info", {})
         server_info_class = import_class_by_name(server_info_dict["type"])
-        server_info = server_info_class.from_dict(server_info_dict)
+        server_info = cast(MCPServerInfo, server_info_class).from_dict(server_info_dict)
 
         # Create a new MCPToolset instance
         return cls(

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

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 from typing import Any
 
-from haystack.tools import Tool
+from mcp import types
 from mcp.server import Server
 from mcp.shared.memory import create_connected_server_and_client_session
 
@@ -17,7 +17,7 @@ def __init__(self, server: Server) -> None:
         super().__init__()
         self.server: Server = server
 
-    async def connect(self) -> list[Tool]:
+    async def connect(self) -> list[types.Tool]:
         """
         Connect to an MCP server using stdio transport.