From 0f04084e9230cfc75853fc461e40532b73843a04 Mon Sep 17 00:00:00 2001
From: Matthew Zhou <mattzh1314@gmail.com>
Date: Mon, 5 May 2025 00:26:19 -0700
Subject: [PATCH] test: Write unit tests for stdio MCP server integration
 (#2001)

---
 tests/mcp/__init__.py              |   0
 tests/mcp/mcp_config.json          |  11 +++
 tests/mcp/test_mcp.py              | 125 +++++++++++++++++++++++++++++
 tests/mcp/weather/requirements.txt |  27 +++++++
 tests/mcp/weather/weather.py       |  97 ++++++++++++++++++++++
 5 files changed, 260 insertions(+)
 create mode 100644 tests/mcp/__init__.py
 create mode 100644 tests/mcp/mcp_config.json
 create mode 100644 tests/mcp/test_mcp.py
 create mode 100644 tests/mcp/weather/requirements.txt
 create mode 100644 tests/mcp/weather/weather.py

diff --git a/tests/mcp/__init__.py b/tests/mcp/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/mcp/mcp_config.json b/tests/mcp/mcp_config.json
new file mode 100644
index 00000000..3804eb5d
--- /dev/null
+++ b/tests/mcp/mcp_config.json
@@ -0,0 +1,11 @@
+{
+  "mcpServers": {
+    "weather": {
+      "transport": "stdio",
+      "command": "/Users/mattzhou/letta-cloud/apps/core/tests/mcp/weather/venv/bin/python3",
+      "args": [
+        "/Users/mattzhou/letta-cloud/apps/core/tests/mcp/weather/weather.py"
+      ]
+    }
+  }
+}
diff --git a/tests/mcp/test_mcp.py b/tests/mcp/test_mcp.py
new file mode 100644
index 00000000..2b0e0610
--- /dev/null
+++ b/tests/mcp/test_mcp.py
@@ -0,0 +1,125 @@
+import json
+import os
+import subprocess
+import venv
+from pathlib import Path
+
+import pytest
+from mcp import Tool as MCPTool
+
+import letta.constants as constants
+from letta.config import LettaConfig
+from letta.functions.mcp_client.types import MCPServerType, StdioServerConfig
+from letta.schemas.tool import ToolCreate
+from letta.server.server import SyncServer
+
+
+def create_virtualenv_and_install_requirements(requirements_path: Path, name="venv") -> Path:
+    requirements_path = requirements_path.resolve()
+
+    if not requirements_path.exists():
+        raise FileNotFoundError(f"Requirements file not found: {requirements_path}")
+    if requirements_path.name != "requirements.txt":
+        raise ValueError(f"Expected file named 'requirements.txt', got: {requirements_path.name}")
+
+    venv_dir = requirements_path.parent / name
+
+    if not venv_dir.exists():
+        venv.EnvBuilder(with_pip=True).create(venv_dir)
+
+    pip_path = venv_dir / ("Scripts" if os.name == "nt" else "bin") / "pip"
+    if not pip_path.exists():
+        raise FileNotFoundError(f"pip executable not found at: {pip_path}")
+
+    try:
+        subprocess.check_call([str(pip_path), "install", "-r", str(requirements_path)])
+    except subprocess.CalledProcessError as exc:
+        raise RuntimeError(f"pip install failed with exit code {exc.returncode}")
+
+    return venv_dir
+
+
+@pytest.fixture
+def empty_mcp_config(tmp_path):
+    path = Path(__file__).parent / "mcp_config.json"
+    path.write_text(json.dumps({}))  # writes "{}"
+
+    create_virtualenv_and_install_requirements(Path(__file__).parent / "weather" / "requirements.txt")
+    return path
+
+
+@pytest.fixture
+def server(empty_mcp_config):
+    config = LettaConfig.load()
+    print("CONFIG PATH", config.config_path)
+
+    config.save()
+
+    old_dir = constants.LETTA_DIR
+    constants.LETTA_DIR = str(Path(__file__).parent)
+
+    server = SyncServer()
+    yield server
+    constants.LETTA_DIR = old_dir
+
+
+@pytest.fixture
+def default_user(server):
+    user = server.user_manager.get_user_or_default()
+    yield user
+
+
+def test_stdio_mcp_server(server, default_user):
+    assert server.mcp_clients == {}
+
+    mcp_server_name = "weather"
+    command = str(Path(__file__).parent / "weather" / "venv" / "bin" / "python3")
+    args = [str(Path(__file__).parent / "weather" / "weather.py")]
+    stdio_mcp_config = StdioServerConfig(server_name=mcp_server_name, command=command, args=args)
+    server.add_mcp_server_to_config(stdio_mcp_config)
+
+    # Check that it's in clients
+    assert mcp_server_name in server.mcp_clients
+
+    # Check that it's in the server mapping
+    mcp_server_mapping = server.get_mcp_servers()
+    assert mcp_server_name in mcp_server_mapping
+    assert mcp_server_mapping[mcp_server_name] == StdioServerConfig(
+        server_name=mcp_server_name, type=MCPServerType.STDIO, command=command, args=args, env=None
+    )
+
+    # Check that it can return valid tools
+    tools = server.get_tools_from_mcp_server(mcp_server_name)
+    assert tools == [
+        MCPTool(
+            name="get_alerts",
+            description="Get weather alerts for a US state.\n\n    Args:\n        state: Two-letter US state code (e.g. CA, NY)\n    ",
+            inputSchema={
+                "properties": {"state": {"title": "State", "type": "string"}},
+                "required": ["state"],
+                "title": "get_alertsArguments",
+                "type": "object",
+            },
+        ),
+        MCPTool(
+            name="get_forecast",
+            description="Get weather forecast for a location.\n\n    Args:\n        latitude: Latitude of the location\n        longitude: Longitude of the location\n    ",
+            inputSchema={
+                "properties": {"latitude": {"title": "Latitude", "type": "number"}, "longitude": {"title": "Longitude", "type": "number"}},
+                "required": ["latitude", "longitude"],
+                "title": "get_forecastArguments",
+                "type": "object",
+            },
+        ),
+    ]
+    get_alerts_mcp_tool = tools[0]
+
+    tool_create = ToolCreate.from_mcp(mcp_server_name=mcp_server_name, mcp_tool=get_alerts_mcp_tool)
+    get_alerts_tool = server.tool_manager.create_or_update_mcp_tool(
+        tool_create=tool_create, mcp_server_name=mcp_server_name, actor=default_user
+    )
+
+    # Attempt running the tool
+    function_response, is_error = server.mcp_clients[mcp_server_name].execute_tool(tool_name="get_alerts", tool_args={"state": "CA"})
+    assert not is_error
+    assert len(function_response) > 1000, function_response  # Crude heuristic for an expected result
diff --git a/tests/mcp/weather/requirements.txt b/tests/mcp/weather/requirements.txt
new file mode 100644
index 00000000..5de03968
--- /dev/null
+++ b/tests/mcp/weather/requirements.txt
@@ -0,0 +1,27 @@
+annotated-types==0.7.0
+anyio==4.9.0
+certifi==2025.4.26
+click==8.1.8
+h11==0.16.0
+httpcore==1.0.9
+httpx==0.28.1
+httpx-sse==0.4.0
+idna==3.10
+markdown-it-py==3.0.0
+mcp==1.7.1
+mdurl==0.1.2
+pydantic==2.11.4
+pydantic-settings==2.9.1
+pydantic_core==2.33.2
+Pygments==2.19.1
+python-dotenv==1.1.0
+python-multipart==0.0.20
+rich==14.0.0
+shellingham==1.5.4
+sniffio==1.3.1
+sse-starlette==2.3.3
+starlette==0.46.2
+typer==0.15.3
+typing-inspection==0.4.0
+typing_extensions==4.13.2
+uvicorn==0.34.2
diff --git a/tests/mcp/weather/weather.py b/tests/mcp/weather/weather.py
new file mode 100644
index 00000000..599b8fd1
--- /dev/null
+++ b/tests/mcp/weather/weather.py
@@ -0,0 +1,97 @@
+from typing import Any
+
+import httpx
+from mcp.server.fastmcp import FastMCP
+
+# Initialize FastMCP server
+mcp = FastMCP("weather")
+
+# Constants
+NWS_API_BASE = "https://api.weather.gov"
+USER_AGENT = "weather-app/1.0"
+
+
+async def make_nws_request(url: str) -> dict[str, Any] | None:
+    """Make a request to the NWS API with proper error handling."""
+    headers = {"User-Agent": USER_AGENT, "Accept": "application/geo+json"}
+    async with httpx.AsyncClient() as client:
+        try:
+            response = await client.get(url, headers=headers, timeout=30.0)
+            response.raise_for_status()
+            return response.json()
+        except Exception:
+            return None
+
+
+def format_alert(feature: dict) -> str:
+    """Format an alert feature into a readable string."""
+    props = feature["properties"]
+    return f"""
+Event: {props.get('event', 'Unknown')}
+Area: {props.get('areaDesc', 'Unknown')}
+Severity: {props.get('severity', 'Unknown')}
+Description: {props.get('description', 'No description available')}
+Instructions: {props.get('instruction', 'No specific instructions provided')}
+"""
+
+
+@mcp.tool()
+async def get_alerts(state: str) -> str:
+    """Get weather alerts for a US state.
+
+    Args:
+        state: Two-letter US state code (e.g. CA, NY)
+    """
+    url = f"{NWS_API_BASE}/alerts/active/area/{state}"
+    data = await make_nws_request(url)
+
+    if not data or "features" not in data:
+        return "Unable to fetch alerts or no alerts found."
+
+    if not data["features"]:
+        return "No active alerts for this state."
+
+    alerts = [format_alert(feature) for feature in data["features"]]
+    return "\n---\n".join(alerts)
+
+
+@mcp.tool()
+async def get_forecast(latitude: float, longitude: float) -> str:
+    """Get weather forecast for a location.
+
+    Args:
+        latitude: Latitude of the location
+        longitude: Longitude of the location
+    """
+    # First get the forecast grid endpoint
+    points_url = f"{NWS_API_BASE}/points/{latitude},{longitude}"
+    points_data = await make_nws_request(points_url)
+
+    if not points_data:
+        return "Unable to fetch forecast data for this location."
+
+    # Get the forecast URL from the points response
+    forecast_url = points_data["properties"]["forecast"]
+    forecast_data = await make_nws_request(forecast_url)
+
+    if not forecast_data:
+        return "Unable to fetch detailed forecast."
+
+    # Format the periods into a readable forecast
+    periods = forecast_data["properties"]["periods"]
+    forecasts = []
+    for period in periods[:5]:  # Only show next 5 periods
+        forecast = f"""
+{period['name']}:
+Temperature: {period['temperature']}°{period['temperatureUnit']}
+Wind: {period['windSpeed']} {period['windDirection']}
+Forecast: {period['detailedForecast']}
+"""
+        forecasts.append(forecast)
+
+    return "\n---\n".join(forecasts)
+
+
+if __name__ == "__main__":
+    # Initialize and run the server
+    mcp.run(transport="stdio")