Set MCPServer id and tool_prefix in load_mcp_servers (#3052)

Co-authored-by: Marcelo Trylesinski <[email protected]>

Author: DouweM

pydantic_ai_slim/pydantic_ai/mcp.py

@@ -167,6 +167,10 @@ async def client_streams(
     def id(self) -> str | None:
         return self._id
 
+    @id.setter
+    def id(self, value: str | None):
+        self._id = value
+
     @property
     def label(self) -> str:
         if self.id:
@@ -414,6 +418,9 @@ def _get_content(
         else:
             assert_never(resource)
 
+    def __eq__(self, value: object, /) -> bool:
+        return isinstance(value, MCPServer) and self.id == value.id and self.tool_prefix == value.tool_prefix
+
 
 class MCPServerStdio(MCPServer):
     """Runs an MCP server in a subprocess and communicates with it over stdin/stdout.
@@ -568,10 +575,10 @@ def __repr__(self) -> str:
         return f'{self.__class__.__name__}({", ".join(repr_args)})'
 
     def __eq__(self, value: object, /) -> bool:
-        if not isinstance(value, MCPServerStdio):
-            return False  # pragma: no cover
         return (
-            self.command == value.command
+            super().__eq__(value)
+            and isinstance(value, MCPServerStdio)
+            and self.command == value.command
             and self.args == value.args
             and self.env == value.env
             and self.cwd == value.cwd
@@ -809,9 +816,7 @@ def _transport_client(self):
         return sse_client  # pragma: no cover
 
     def __eq__(self, value: object, /) -> bool:
-        if not isinstance(value, MCPServerSSE):
-            return False  # pragma: no cover
-        return self.url == value.url
+        return super().__eq__(value) and isinstance(value, MCPServerSSE) and self.url == value.url
 
 
 @deprecated('The `MCPServerHTTP` class is deprecated, use `MCPServerSSE` instead.')
@@ -885,9 +890,7 @@ def _transport_client(self):
         return streamablehttp_client  # pragma: no cover
 
     def __eq__(self, value: object, /) -> bool:
-        if not isinstance(value, MCPServerStreamableHTTP):
-            return False  # pragma: no cover
-        return self.url == value.url
+        return super().__eq__(value) and isinstance(value, MCPServerStreamableHTTP) and self.url == value.url
 
 
 ToolResult = (
@@ -964,4 +967,11 @@ def load_mcp_servers(config_path: str | Path) -> list[MCPServerStdio | MCPServer
         raise FileNotFoundError(f'Config file {config_path} not found')
 
     config = MCPServerConfig.model_validate_json(config_path.read_bytes())
-    return list(config.mcp_servers.values())
+
+    servers: list[MCPServerStdio | MCPServerStreamableHTTP | MCPServerSSE] = []
+    for name, server in config.mcp_servers.items():
+        server.id = name
+        server.tool_prefix = name
+        servers.append(server)
+
+    return servers

pyproject.toml

@@ -87,7 +87,7 @@ dev = [
     "anyio>=4.5.0",
     "asgi-lifespan>=2.1.0",
     "devtools>=0.12.2",
-    "coverage[toml]>=7.10.3",
+    "coverage[toml]>=7.10.7",
     "dirty-equals>=0.9.0",
     "duckduckgo-search>=7.0.0",
     "inline-snapshot>=0.19.3",

tests/test_mcp.py

@@ -1464,13 +1464,18 @@ def test_load_mcp_servers(tmp_path: Path):
     config = tmp_path / 'mcp.json'
 
     config.write_text('{"mcpServers": {"potato": {"url": "https://example.com/mcp"}}}')
-    assert load_mcp_servers(config) == snapshot([MCPServerStreamableHTTP(url='https://example.com/mcp')])
+    server = load_mcp_servers(config)[0]
+    assert server == MCPServerStreamableHTTP(url='https://example.com/mcp', id='potato', tool_prefix='potato')
 
     config.write_text('{"mcpServers": {"potato": {"command": "python", "args": ["-m", "tests.mcp_server"]}}}')
-    assert load_mcp_servers(config) == snapshot([MCPServerStdio(command='python', args=['-m', 'tests.mcp_server'])])
+    server = load_mcp_servers(config)[0]
+    assert server == MCPServerStdio(
+        command='python', args=['-m', 'tests.mcp_server'], id='potato', tool_prefix='potato'
+    )
 
     config.write_text('{"mcpServers": {"potato": {"url": "https://example.com/sse"}}}')
-    assert load_mcp_servers(config) == snapshot([MCPServerSSE(url='https://example.com/sse')])
+    server = load_mcp_servers(config)[0]
+    assert server == MCPServerSSE(url='https://example.com/sse', id='potato', tool_prefix='potato')
 
     with pytest.raises(FileNotFoundError):
         load_mcp_servers(tmp_path / 'does_not_exist.json')