Add tests for MCPAgent

Author: RussellLuo

coagent/agents/mcp_agent.py

@@ -4,6 +4,7 @@
 from mcp import ClientSession, Tool
 from mcp.types import ImageContent, TextContent
 from mcp.client.sse import sse_client
+from mcp.client.stdio import stdio_client, StdioServerParameters
 import jsonschema
 
 from .aswarm import Agent as SwarmAgent
@@ -23,7 +24,7 @@ def __init__(
         super().__init__(system=system, client=client)
 
         self._mcp_server_base_url: str = mcp_server_base_url
-        self._mcp_sse_client: AsyncContextManager[tuple] | None = None
+        self._mcp_client_transport: AsyncContextManager[tuple] | None = None
         self._mcp_client_session: ClientSession | None = None
 
         self._mcp_swarm_agent: SwarmAgent | None = None
@@ -34,7 +35,62 @@ def mcp_server_base_url(self) -> str:
             raise ValueError("MCP server base URL is empty")
         return self._mcp_server_base_url
 
-    def make_tool(self, t: Tool) -> Callable:
+    def _make_mcp_client_transport(self) -> AsyncContextManager[tuple]:
+        if self.mcp_server_base_url.startswith(("http://", "https://")):
+            url = urljoin(self.mcp_server_base_url, "sse")
+            return sse_client(url=url)
+        else:
+            # Mainly for testing purposes.
+            command, arg = self.mcp_server_base_url.split(" ", 1)
+            params = StdioServerParameters(command=command, args=[arg])
+            return stdio_client(params)
+
+    async def started(self) -> None:
+        """
+        Combining `started` and `stopped` to achieve the following behavior:
+
+            async with sse_client(url=url) as (read, write):
+                async with ClientSession(read, write) as session:
+                    pass
+        """
+        self._mcp_client_transport = self._make_mcp_client_transport()
+        read, write = await self._mcp_client_transport.__aenter__()
+
+        self._mcp_client_session = ClientSession(read, write)
+        await self._mcp_client_session.__aenter__()
+
+        # Initialize the connection
+        await self._mcp_client_session.initialize()
+
+    async def stopped(self) -> None:
+        await self._mcp_client_session.__aexit__(None, None, None)
+        await self._mcp_client_transport.__aexit__(None, None, None)
+
+    async def _handle_data(self) -> None:
+        """Override the method to handle exceptions properly."""
+        try:
+            await super()._handle_data()
+        finally:
+            # Ensure the resources created in `started` are properly cleaned up.
+            await self.stopped()
+
+    async def get_swarm_agent(self) -> SwarmAgent:
+        if not self._mcp_swarm_agent:
+            tools = await self._get_tools()
+            self._mcp_swarm_agent = SwarmAgent(
+                name=self.name,
+                model=self.client.model,
+                instructions=self.system,
+                functions=[wrap_error(t) for t in tools],
+            )
+        return self._mcp_swarm_agent
+
+    async def _get_tools(self) -> list[Callable]:
+        result = await self._mcp_client_session.list_tools()
+        tools = [self._make_tool(t) for t in result.tools]
+        return tools
+
+    def _make_tool(self, t: Tool) -> Callable:
         async def tool(**kwargs) -> Any:
             # Validate the input against the schema
             jsonschema.validate(instance=kwargs, schema=t.inputSchema)
@@ -64,51 +120,5 @@ async def tool(**kwargs) -> Any:
             description=t.description,
             parameters=t.inputSchema,
         )
-        tool.__mcp_tool_args__ = t.inputSchema["properties"].keys()
+        tool.__mcp_tool_args__ = tuple(t.inputSchema["properties"].keys())
         return tool
-
-    async def get_tools(self) -> list[Callable]:
-        result = await self._mcp_client_session.list_tools()
-        tools = [self.make_tool(t) for t in result.tools]
-        return tools
-
-    async def get_swarm_agent(self) -> SwarmAgent:
-        if not self._mcp_swarm_agent:
-            tools = await self.get_tools()
-            self._mcp_swarm_agent = SwarmAgent(
-                name=self.name,
-                model=self.client.model,
-                instructions=self.system,
-                functions=[wrap_error(t) for t in tools],
-            )
-        return self._mcp_swarm_agent
-
-    async def started(self) -> None:
-        """
-        Combining `started` and `stopped` to achieve the following behavior:
-
-            async with sse_client(url=url) as (read, write):
-                async with ClientSession(read, write) as session:
-                    pass
-        """
-        url = urljoin(self.mcp_server_base_url, "sse")
-        self._mcp_sse_client = sse_client(url=url)
-        read, write = await self._mcp_sse_client.__aenter__()
-
-        self._mcp_client_session = ClientSession(read, write)
-        await self._mcp_client_session.__aenter__()
-
-        # Initialize the connection
-        await self._mcp_client_session.initialize()
-
-    async def stopped(self) -> None:
-        await self._mcp_client_session.__aexit__(None, None, None)
-        await self._mcp_sse_client.__aexit__(None, None, None)
-
-    async def _handle_data(self) -> None:
-        """Override the method to handle exceptions properly."""
-        try:
-            await super()._handle_data()
-        finally:
-            # Ensure the resources created in `started` are properly cleaned up.
-            await self.stopped()

tests/agents/mcp_server.py

@@ -0,0 +1,13 @@
+from mcp.server.fastmcp import FastMCP
+
+mcp = FastMCP("Weather")
+
+
+@mcp.tool()
+def query_weather(city: str) -> str:
+    """Query the weather in the given city."""
+    return f"The weather in {city} is sunny."
+
+
+if __name__ == "__main__":
+    mcp.run()

tests/agents/test_mcp_agent.py

@@ -0,0 +1,37 @@
+import pytest
+
+from coagent.agents.mcp_agent import MCPAgent
+
+
+class TestMCPAgent:
+    @pytest.mark.asyncio
+    async def test_get_tools(self):
+        agent = MCPAgent(mcp_server_base_url="python tests/agents/mcp_server.py")
+        await agent.started()
+
+        tools = await agent._get_tools()
+        tool = tools[0]
+
+        assert tool.__name__ == "query_weather"
+        assert tool.__doc__ == "Query the weather in the given city."
+        assert tool.__mcp_tool_schema__ == {
+            "description": "Query the weather in the given city.",
+            "name": "query_weather",
+            "parameters": {
+                "properties": {
+                    "city": {
+                        "title": "City",
+                        "type": "string",
+                    }
+                },
+                "required": ["city"],
+                "title": "query_weatherArguments",
+                "type": "object",
+            },
+        }
+        assert tool.__mcp_tool_args__ == ("city",)
+
+        result = await tool(city="Beijing")
+        assert result == "The weather in Beijing is sunny."
+
+        await agent.stopped()