fix: enhance MCP integration with robust error handling and task coverage

- 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

@@ -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"""
+    try:
-    return a + b
+        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"""
+    try:
-    return a - b
+        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"""
+    try:
-    return a * b
+        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"""
+    try:
-    if b == 0:
+        if float(b) == 0:
-        return {"error": "Cannot divide by zero"}
+            raise ZeroDivisionError("Cannot divide by zero")
-    return a / b
+        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)

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(