From 0761165684006bc64ebf79b476c8c85accc4da06 Mon Sep 17 00:00:00 2001
From: Pavan Kumar <66913595+ascender1729@users.noreply.github.com>
Date: Thu, 17 Apr 2025 17:09:19 +0000
Subject: [PATCH] fix: enhance MCP integration with robust error handling and
 task coverage
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Improved `math_server.py` and `test_integration.py` for stability
- Added structured error handling across MCP components
- Aligned implementation with Kye’s task requirements for Swarms MCP integration
---
 examples/mcp_example/math_server.py      | 69 ++++++++++++++++++------
 examples/mcp_example/test_integration.py | 16 +++++-
 2 files changed, 68 insertions(+), 17 deletions(-)

diff --git a/examples/mcp_example/math_server.py b/examples/mcp_example/math_server.py
index 0e3dab68..9fcfcafd 100644
--- a/examples/mcp_example/math_server.py
+++ b/examples/mcp_example/math_server.py
@@ -1,33 +1,70 @@
-
+import logging
 from fastmcp import FastMCP
 from litellm import LiteLLM
 
+# Configure logging
+logging.basicConfig(level=logging.ERROR)
+logger = logging.getLogger(__name__)
+
 # Initialize MCP server for math operations
 mcp = FastMCP("Math-Server")
 
-@mcp.tool()
+@mcp.tool(name="add", description="Add two numbers") 
 def add(a: float, b: float) -> float:
-    """Add two numbers together"""
-    return a + b
+    try:
+        result = float(a) + float(b)
+        return result
+    except (ValueError, TypeError) as e:
+        logger.error(f"Invalid input types for addition: {e}")
+        raise ValueError("Inputs must be valid numbers")
+    except Exception as e:
+        logger.error(f"Unexpected error in add operation: {e}")
+        raise
 
-@mcp.tool()
+@mcp.tool(name="subtract", description="Subtract b from a")
 def subtract(a: float, b: float) -> float:
-    """Subtract b from a"""
-    return a - b
+    try:
+        result = float(a) - float(b)
+        return result
+    except (ValueError, TypeError) as e:
+        logger.error(f"Invalid input types for subtraction: {e}")
+        raise ValueError("Inputs must be valid numbers")
+    except Exception as e:
+        logger.error(f"Unexpected error in subtract operation: {e}")
+        raise
 
-@mcp.tool()
+@mcp.tool(name="multiply", description="Multiply two numbers together")
 def multiply(a: float, b: float) -> float:
-    """Multiply two numbers together"""
-    return a * b
+    try:
+        result = float(a) * float(b)
+        return result
+    except (ValueError, TypeError) as e:
+        logger.error(f"Invalid input types for multiplication: {e}")
+        raise ValueError("Inputs must be valid numbers")
+    except Exception as e:
+        logger.error(f"Unexpected error in multiply operation: {e}")
+        raise
 
-@mcp.tool()
+@mcp.tool(name="divide", description="Divide a by b")
 def divide(a: float, b: float) -> float:
-    """Divide a by b"""
-    if b == 0:
-        return {"error": "Cannot divide by zero"}
-    return a / b
+    try:
+        if float(b) == 0:
+            raise ZeroDivisionError("Cannot divide by zero")
+        result = float(a) / float(b)
+        return result
+    except (ValueError, TypeError) as e:
+        logger.error(f"Invalid input types for division: {e}")
+        raise ValueError("Inputs must be valid numbers")
+    except ZeroDivisionError as e:
+        logger.error(f"ZeroDivisionError: {e}")
+        raise
+    except Exception as e:
+        logger.error(f"Unexpected error in divide operation: {e}")
+        raise
+
+
 
 if __name__ == "__main__":
     print("Starting Math Server on port 6274...")
     llm = LiteLLM(model_name="gpt-4o-mini")
-    mcp.run(transport="sse", host="0.0.0.0", port=6274)
+    mcp.run(transport="sse", host="0.0.0.0", port=6274)
\ No newline at end of file
diff --git a/examples/mcp_example/test_integration.py b/examples/mcp_example/test_integration.py
index e057c9b0..21f96272 100644
--- a/examples/mcp_example/test_integration.py
+++ b/examples/mcp_example/test_integration.py
@@ -16,9 +16,23 @@ if not api_key:
 math_server = MCPServerSseParams(
     url="http://0.0.0.0:6274",
     headers={"Content-Type": "application/json"},
-    timeout=10.0
+    timeout=10.0,
+    sse_read_timeout=300.0  # 5 minute timeout for long-running operations
 )
 
+# Error handling wrapper
+def setup_mcp_server(server_params: MCPServerSseParams):
+    try:
+        math_agent = setup_agent(
+            "Math-Agent",
+            "Handles mathematical calculations",
+            [server_params]
+        )
+        return math_agent
+    except Exception as e:
+        logger.error(f"Failed to setup MCP server: {e}")
+        raise
+
 def setup_agent(name: str, description: str, servers: list) -> Agent:
     """Setup an agent with MCP server connections"""
     return Agent(