From 631ad71e0338de47794abe222995ea46a1f5f0e1 Mon Sep 17 00:00:00 2001
From: Pavan Kumar <66913595+ascender1729@users.noreply.github.com>
Date: Thu, 17 Apr 2025 16:06:00 +0000
Subject: [PATCH] feat: enhance MCP integration in Swarms with full testing
 support

---
 .replit                             | 11 +++++++++
 swarms/structs/agent.py             | 12 +++++----
 swarms/tools/mcp_integration.py     | 38 ++++++++++++++++++-----------
 tests/tools/test_mcp_integration.py | 24 ++++++++++++++++++
 4 files changed, 66 insertions(+), 19 deletions(-)
 create mode 100644 tests/tools/test_mcp_integration.py

diff --git a/.replit b/.replit
index d288edfe..c70f9765 100644
--- a/.replit
+++ b/.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"
diff --git a/swarms/structs/agent.py b/swarms/structs/agent.py
index c45d039b..582d298b 100644
--- a/swarms/structs/agent.py
+++ b/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
-                        if self.tools is not None:
-                            out = self.parse_and_execute_tools(
-                                response
-                            )
+                        # Check and execute tools (including MCP)
+                        if self.tools is not None or hasattr(self, 'mcp_servers'):
+                            if self.tools:
+                                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
diff --git a/swarms/tools/mcp_integration.py b/swarms/tools/mcp_integration.py
index 7a74dbaa..6880527d 100644
--- a/swarms/tools/mcp_integration.py
+++ b/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)
-
-    # Connect the server
-    asyncio.run(server.connect())
-
-    # Return the server
-    output = asyncio.run(server.call_tool(function_call))
-
-    output = output.model_dump()
-
-    # Cleanup the server
-    asyncio.run(server.cleanup())
-
-    return any_to_str(output)
+    try:
+        server = MCPServerSse(params, cache_tools_list=True)
+        
+        # Connect the server
+        asyncio.run(server.connect())
+        
+        # Extract tool name and args from function call
+        tool_name = function_call.get("tool_name") or function_call.get("name")
+        if not tool_name:
+            raise ValueError("No tool name provided in function call")
+
+        # Call the tool
+        output = asyncio.run(server.call_tool(function_call))
+        
+        # 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(
diff --git a/tests/tools/test_mcp_integration.py b/tests/tools/test_mcp_integration.py
new file mode 100644
index 00000000..f67fcf51
--- /dev/null
+++ b/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, {})