Update tests for MCPAgent

Author: RussellLuo

coagent/agents/mcp_agent.py

@@ -36,7 +36,7 @@ def __init__(
         self._mcp_client_session: ClientSession | None = None
 
         self._mcp_swarm_agent: SwarmAgent | None = None
-        self._mcp_system_prompt: Prompt | None = system
+        self._mcp_system_prompt_config: Prompt | None = system
 
     @property
     def mcp_server_base_url(self) -> str:
@@ -85,7 +85,7 @@ async def _handle_data(self) -> None:
 
     async def get_swarm_agent(self) -> SwarmAgent:
         if not self._mcp_swarm_agent:
-            system = await self._get_system_prompt()
+            system = await self._get_prompt(self._mcp_system_prompt_config)
             tools = await self._get_tools()
             self._mcp_swarm_agent = SwarmAgent(
                 name=self.name,
@@ -95,13 +95,13 @@ async def get_swarm_agent(self) -> SwarmAgent:
             )
         return self._mcp_swarm_agent
 
-    async def _get_system_prompt(self) -> str:
-        if not self._mcp_system_prompt:
+    async def _get_prompt(self, prompt_config: Prompt | None) -> str:
+        if not prompt_config:
             return ""
 
         try:
             prompt = await self._mcp_client_session.get_prompt(
-                **dataclasses.asdict(self._mcp_system_prompt),
+                **dataclasses.asdict(prompt_config),
             )
         except McpError as exc:
             raise InternalError(str(exc))

tests/agents/mcp_server.py

@@ -3,6 +3,12 @@
 mcp = FastMCP("Weather")
 
 
+@mcp.prompt()
+def system_prompt(role: str) -> str:
+    """Create a system prompt."""
+    return f"You are a helpful {role}."
+
+
 @mcp.tool()
 def query_weather(city: str) -> str:
     """Query the weather in the given city."""

tests/agents/test_mcp_agent.py

@@ -2,10 +2,30 @@
 
 import pytest
 
-from coagent.agents.mcp_agent import MCPAgent
+from coagent.agents.mcp_agent import MCPAgent, Prompt
+from coagent.core.exceptions import BaseError
 
 
 class TestMCPAgent:
+    @pytest.mark.skipif(sys.platform == "win32", reason="Does not run on Windows.")
+    @pytest.mark.asyncio
+    async def test_get_prompt(self):
+        agent = MCPAgent(mcp_server_base_url="python tests/agents/mcp_server.py")
+        await agent.started()
+
+        # Success
+        config = Prompt(name="system_prompt", arguments={"role": "Weather Reporter"})
+        prompt = await agent._get_prompt(config)
+        assert prompt == "You are a helpful Weather Reporter."
+
+        # Error
+        config = Prompt(name="x", arguments={"role": "Weather Reporter"})
+        with pytest.raises(BaseError) as exc:
+            _ = await agent._get_prompt(config)
+        assert str(exc.value).startswith("Unknown prompt: x")
+
+        await agent.stopped()
+
     @pytest.mark.skipif(sys.platform == "win32", reason="Does not run on Windows.")
     @pytest.mark.asyncio
     async def test_get_tools(self):