feat: enhance MCP integration in Swarms with full testing support

.replit

@@ -2,3 +2,14 @@ modules = ["python-3.10", "bash"]
 
 [nix]
 channel = "stable-24_05"
+
+[workflows]
+
+[[workflows.workflow]]
+name = "Run MCP Tests"
+author = 13983571
+mode = "sequential"
+
+[[workflows.workflow.tasks]]
+task = "shell.exec"
+args = "python -m pytest tests/tools/test_mcp_integration.py -v"

swarms/structs/agent.py

@@ -695,6 +695,7 @@ class Agent:
             exists(self.tools)
             or exists(self.list_base_models)
             or exists(self.tool_schema)
+            or exists(self.mcp_servers)
         ):
 
             self.tool_struct = BaseTool(
@@ -1107,11 +1108,12 @@ class Agent:
                         ) as executor:
                             executor.map(lambda f: f(), update_tasks)
 
-                        # Check and execute tools
+                        # Check and execute tools (including MCP)
-                        if self.tools is not None:
+                        if self.tools is not None or hasattr(self, 'mcp_servers'):
-                            out = self.parse_and_execute_tools(
+                            if self.tools:
-                                response
+                                out = self.parse_and_execute_tools(response)
-                            )
+                            if hasattr(self, 'mcp_servers') and self.mcp_servers:
+                                out = self.mcp_execution_flow(response)
 
                             self.short_memory.add(
                                 role="Tool Executor", content=out

swarms/tools/mcp_integration.py

@@ -363,20 +363,30 @@ def mcp_flow(
     params: MCPServerSseParams,
     function_call: dict[str, Any],
 ) -> MCPServer:
-    server = MCPServerSse(params, cache_tools_list=True)
+    try:
-
+        server = MCPServerSse(params, cache_tools_list=True)
-    # Connect the server
+        
-    asyncio.run(server.connect())
+        # Connect the server
-
+        asyncio.run(server.connect())
-    # Return the server
+        
-    output = asyncio.run(server.call_tool(function_call))
+        # Extract tool name and args from function call
-
+        tool_name = function_call.get("tool_name") or function_call.get("name")
-    output = output.model_dump()
+        if not tool_name:
-
+            raise ValueError("No tool name provided in function call")
-    # Cleanup the server
+
-    asyncio.run(server.cleanup())
+        # Call the tool
-
+        output = asyncio.run(server.call_tool(function_call))
-    return any_to_str(output)
+        
+        # Convert to serializable format
+        output = output.model_dump() if hasattr(output, "model_dump") else output
+        
+        # Cleanup the server
+        asyncio.run(server.cleanup())
+        
+        return any_to_str(output)
+    except Exception as e:
+        logger.error(f"Error in MCP flow: {e}")
+        raise
 
 
 def batch_mcp_flow(

tests/tools/test_mcp_integration.py

@@ -0,0 +1,24 @@
+
+import pytest
+from swarms.tools.mcp_integration import MCPServerSseParams, mcp_flow
+
+def test_mcp_flow():
+    params = MCPServerSseParams(
+        url="http://localhost:6274",
+        headers={"Content-Type": "application/json"}
+    )
+    
+    function_call = {
+        "tool_name": "test_tool",
+        "args": {"param1": "value1"}
+    }
+    
+    try:
+        result = mcp_flow(params, function_call)
+        assert isinstance(result, str)
+    except Exception as e:
+        pytest.fail(f"MCP flow failed: {e}")
+
+def test_mcp_invalid_params():
+    with pytest.raises(Exception):
+        mcp_flow(None, {})