Make MCPAgent more configurable

Author: RussellLuo

coagent/agents/mcp_agent.py

@@ -21,36 +21,53 @@ class Prompt:
 
 
 class MCPAgent(ChatAgent):
-    """An agent that can use tools provided by MCP (Model Context Protocol) servers."""
+    """An agent that can use prompts and tools provided by MCP (Model Context Protocol) servers."""
 
     def __init__(
         self,
-        system: Prompt | str = "",
         mcp_server_base_url: str = "",
-        selected_tools: list[str] | None = None,
+        mcp_server_headers: dict[str, Any] | None = None,
+        system: Prompt | str = "",
+        tools: list[str] | None = None,
         client: ModelClient = default_model_client,
     ) -> None:
         super().__init__(system="", client=client)
 
         self._mcp_server_base_url: str = mcp_server_base_url
+        self._mcp_server_headers: dict[str, Any] | None = mcp_server_headers
+
         self._mcp_client_transport: AsyncContextManager[tuple] | None = None
         self._mcp_client_session: ClientSession | None = None
 
         self._mcp_swarm_agent: SwarmAgent | None = None
         self._mcp_system_prompt_config: Prompt | str = system
         # The selected tools to use. If None, all available tools will be used.
-        self._mcp_selected_tools: list[str] | None = selected_tools
+        self._mcp_selected_tools: list[str] | None = tools
 
     @property
     def mcp_server_base_url(self) -> str:
         if not self._mcp_server_base_url:
             raise ValueError("MCP server base URL is empty")
         return self._mcp_server_base_url
 
+    @property
+    def mcp_server_headers(self) -> dict[str, Any] | None:
+        return self._mcp_server_headers
+
+    @property
+    def system(self) -> Prompt | str:
+        """Note that this property is different from the `system` property in ChatAgent."""
+        return self._mcp_system_prompt_config
+
+    @property
+    def tools(self) -> list[str] | None:
+        """Note that this property is different from the `tools` property in ChatAgent."""
+        return self._mcp_selected_tools
+
     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)
+            return sse_client(url=url, headers=self.mcp_server_headers)
         else:
             # Mainly for testing purposes.
             command, arg = self.mcp_server_base_url.split(" ", 1)
@@ -88,8 +105,8 @@ async def _handle_data(self) -> None:
 
     async def get_swarm_agent(self) -> SwarmAgent:
         if not self._mcp_swarm_agent:
-            system = await self._get_prompt(self._mcp_system_prompt_config)
-            tools = await self._get_tools()
+            system = await self._get_prompt(self.system)
+            tools = await self._get_tools(self.tools)
             self._mcp_swarm_agent = SwarmAgent(
                 name=self.name,
                 model=self.client.model,
@@ -117,13 +134,13 @@ async def _get_prompt(self, prompt_config: Prompt | str) -> str:
             case _:  # ImageContent() or EmbeddedResource() or other types
                 return ""
 
-    async def _get_tools(self) -> list[Callable]:
+    async def _get_tools(self, selected_tools: list[str] | None) -> list[Callable]:
         result = await self._mcp_client_session.list_tools()
 
         def filter_tool(t: Tool) -> bool:
-            if self._mcp_selected_tools is None:
+            if selected_tools is None:
                 return True
-            return t.name in self._mcp_selected_tools
+            return t.name in selected_tools
 
         tools = [self._make_tool(t) for t in result.tools if filter_tool(t)]
 

examples/mcp/agent.py

@@ -10,8 +10,8 @@
     "mcp",
     new(
         MCPAgent,
-        system=Prompt(name="system_prompt", arguments={"role": "Weather Reporter"}),
         mcp_server_base_url="http://localhost:8080",
+        system=Prompt(name="system_prompt", arguments={"role": "Weather Reporter"}),
     ),
 )
 

poetry.lock

@@ -35,7 +35,7 @@ blinker = "1.9.0"
 loguru = "0.7.3"
 jq = "1.8.0"
 litellm = "1.60.4"
-mcp = "1.2.0"
+mcp = "1.3.0"
 jinja2 = "3.1.5"
 
 [tool.pyright]

pyproject.toml

@@ -37,7 +37,7 @@ 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()
+        tools = await agent._get_tools(None)
         assert len(tools) == 2
 
         # Get query_weather
@@ -99,13 +99,14 @@ async def test_get_tools(self):
     @pytest.mark.skipif(sys.platform == "win32", reason="Does not run on Windows.")
     @pytest.mark.asyncio
     async def test_get_tools_with_selection(self):
+        selected_tools = ["query_weather"]
         agent = MCPAgent(
             mcp_server_base_url="python tests/agents/mcp_server.py",
-            selected_tools=["query_weather"],
+            tools=selected_tools,
         )
         await agent.started()
 
-        tools = await agent._get_tools()
+        tools = await agent._get_tools(selected_tools)
         assert len(tools) == 1
 
         tool = tools[0]