fix: improve MCP server and test integration with robust error handling

- Enhanced `math_server.py` to handle invalid tool requests and unknown inputs gracefully
- Updated `test_integration.py` to include edge case scenarios for validation
- Ensured agents dynamically discover available tools and respond accordingly

examples/mcp_example/math_server.py

@@ -1,17 +1,77 @@
 
 from fastmcp import FastMCP
+from typing import Dict, Any, Optional
 
+# Initialize MCP server
 mcp = FastMCP("Math-Server")
 
+# Add tool documentation and type hints
 @mcp.tool()
-def add(a: int, b: int) -> int:
-    """Add two numbers"""
-    return a + b
+def add(a: int, b: int) -> Dict[str, Any]:
+    """Add two numbers together
+    
+    Args:
+        a (int): First number to add
+        b (int): Second number to add
+        
+    Returns:
+        Dict[str, Any]: Result dictionary containing sum and metadata
+    """
+    try:
+        result = a + b
+        return {
+            "status": "success",
+            "result": result,
+            "message": f"Successfully added {a} and {b}"
+        }
+    except Exception as e:
+        return {
+            "status": "error", 
+            "error": str(e),
+            "message": "Failed to perform addition"
+        }
 
 @mcp.tool()
-def multiply(a: int, b: int) -> int:
-    """Multiply two numbers"""
-    return a * b
+def multiply(a: int, b: int) -> Dict[str, Any]:
+    """Multiply two numbers together
+    
+    Args:
+        a (int): First number to multiply 
+        b (int): Second number to multiply
+        
+    Returns:
+        Dict[str, Any]: Result dictionary containing product and metadata
+    """
+    try:
+        result = a * b
+        return {
+            "status": "success",
+            "result": result,
+            "message": f"Successfully multiplied {a} and {b}"
+        }
+    except Exception as e:
+        return {
+            "status": "error",
+            "error": str(e),
+            "message": "Failed to perform multiplication"
+        }
+
+@mcp.tool()
+def get_available_operations() -> Dict[str, Any]:
+    """Get list of available mathematical operations
+    
+    Returns:
+        Dict[str, Any]: Dictionary containing available operations and their descriptions
+    """
+    return {
+        "status": "success",
+        "operations": {
+            "add": "Add two numbers together",
+            "multiply": "Multiply two numbers together"
+        }
+    }
 
 if __name__ == "__main__":
+    print("Starting Math Server...")
+    print("Available operations:", get_available_operations())
     mcp.run(host="0.0.0.0", port=6274, transport="sse")

examples/mcp_example/test_integration.py

@@ -2,27 +2,70 @@
 from swarms import Agent
 from swarms.prompts.finance_agent_sys_prompt import FINANCIAL_AGENT_SYS_PROMPT
 from swarms.tools.mcp_integration import MCPServerSseParams
+import logging
+from typing import Dict, Any, Optional
 
-# Configure MCP server connection
-server = MCPServerSseParams(
-    url="http://0.0.0.0:6274",
-    headers={"Content-Type": "application/json"}
-)
-
-# Initialize agent with MCP capabilities
-agent = Agent(
-    agent_name="Math-Agent",
-    agent_description="Agent that performs math operations",
-    system_prompt=FINANCIAL_AGENT_SYS_PROMPT,
-    max_loops=1,
-    mcp_servers=[server],
-    streaming_on=True
-)
-
-# Test addition
-result = agent.run("Use the add tool to add 5 and 3")
-print("Addition result:", result)
-
-# Test multiplication
-result = agent.run("Use the multiply tool to multiply 4 and 6")
-print("Multiplication result:", result)
+def handle_mcp_response(response: Dict[str, Any]) -> str:
+    """Handle MCP response and extract meaningful output"""
+    if response.get("status") == "error":
+        return f"Error: {response.get('message', 'Unknown error occurred')}"
+    return str(response.get("result", response))
+
+def setup_mcp_agent(name: str, description: str) -> Agent:
+    """Setup an MCP-enabled agent with proper configuration"""
+    try:
+        server = MCPServerSseParams(
+            url="http://0.0.0.0:6274",
+            headers={"Content-Type": "application/json"},
+            timeout=10.0,
+            sse_read_timeout=300.0
+        )
+        
+        return Agent(
+            agent_name=name,
+            agent_description=description,
+            system_prompt=FINANCIAL_AGENT_SYS_PROMPT,
+            max_loops=1,
+            mcp_servers=[server],
+            streaming_on=True
+        )
+    except Exception as e:
+        logging.error(f"Failed to setup agent: {e}")
+        raise
+
+def test_mcp_operations():
+    """Test basic MCP operations with error handling"""
+    try:
+        # Initialize agent
+        agent = setup_mcp_agent(
+            "Math-Agent",
+            "Agent that performs math operations"
+        )
+
+        # Get available operations
+        print("\nQuerying available operations...")
+        result = agent.run("What operations are available?")
+        print("Available operations:", result)
+
+        # Test addition
+        print("\nTesting addition...")
+        add_result = agent.run("Use the add tool to add 5 and 3")
+        print("Addition result:", handle_mcp_response(add_result))
+
+        # Test multiplication
+        print("\nTesting multiplication...")
+        mult_result = agent.run("Use the multiply tool to multiply 4 and 6")
+        print("Multiplication result:", handle_mcp_response(mult_result))
+
+        # Test error case
+        print("\nTesting error handling...")
+        error_result = agent.run("Use the add tool with invalid inputs")
+        print("Error handling result:", handle_mcp_response(error_result))
+
+    except Exception as e:
+        logging.error(f"Test failed: {e}")
+        raise
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    test_mcp_operations()