fix(smithery): implement SmitheryConfigMiddleware for query parameter handling and add tests

stickerdaniel · stickerdaniel · commit c0d6e371d7e9 · 2025-07-03T23:08:54.000-04:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -20,6 +20,7 @@ linkedin-scraper = { git = "https://github.com/joeyism/linkedin_scraper.git" }
 
 [dependency-groups]
 dev = [
+    "aiohttp>=3.12.13",
     "pre-commit>=4.2.0",
     "pytest>=8.3.5",
     "pytest-asyncio>=1.0.0",
diff --git a/smithery_main.py b/smithery_main.py
@@ -12,38 +12,87 @@
 import os
 import logging
 from urllib.parse import parse_qs
+from fastmcp.server.middleware import Middleware, MiddlewareContext
 
 from linkedin_mcp_server.config import get_config, reset_config
 from linkedin_mcp_server.drivers.chrome import initialize_driver
 from linkedin_mcp_server.server import create_mcp_server, shutdown_handler
 
 
-def setup_smithery_environment(query_string: str | None = None) -> None:
+class SmitheryConfigMiddleware(Middleware):
     """
-    Set up environment variables from Smithery query parameters.
+    FastMCP middleware to handle Smithery query parameter configuration.
 
-    Args:
-        query_string: Query parameters from Smithery configuration
+    Intercepts HTTP requests and extracts configuration from query parameters,
+    then temporarily sets environment variables for the duration of the request.
     """
-    if not query_string:
-        return
 
-    # Parse query parameters
-    parsed = parse_qs(query_string)
-
-    # Map Smithery parameters to environment variables
-    param_mapping = {
-        "linkedin_email": "LINKEDIN_EMAIL",
-        "linkedin_password": "LINKEDIN_PASSWORD",
-    }
-
-    for param, env_var in param_mapping.items():
-        if param in parsed and parsed[param]:
-            value = parsed[param][0]  # Take first value
-            os.environ[env_var] = value
-
-    # Reset config to pick up new environment variables
-    reset_config()
+    def __init__(self):
+        super().__init__()
+        self.param_mapping = {
+            "linkedin_email": "LINKEDIN_EMAIL",
+            "linkedin_password": "LINKEDIN_PASSWORD",
+        }
+
+    async def on_call_tool(self, context: MiddlewareContext, call_next):
+        """
+        Called before each tool execution.
+        Extract configuration from HTTP request query parameters.
+        """
+        # Store original environment variables
+        original_env = {}
+        for env_var in self.param_mapping.values():
+            original_env[env_var] = os.environ.get(env_var)
+
+        # Extract query parameters from the request context
+        query_params = self._extract_query_params(context)
+
+        if query_params:
+            # Apply configuration from query parameters
+            self._apply_config(query_params)
+
+            # Reset configuration to pick up new environment variables
+            reset_config()
+
+        try:
+            # Execute the tool with the new configuration
+            result = await call_next(context)
+            return result
+        finally:
+            # Restore original environment variables
+            self._restore_env(original_env)
+
+    def _extract_query_params(self, context: MiddlewareContext) -> dict:
+        """Extract query parameters from the request context."""
+        # Check if we can access FastMCP context for HTTP transport
+        if hasattr(context, "fastmcp_context") and context.fastmcp_context:
+            # Check if there's transport-specific information
+            if hasattr(context.fastmcp_context, "transport_info"):
+                transport_info = context.fastmcp_context.transport_info
+                if hasattr(transport_info, "query_params"):
+                    return dict(transport_info.query_params)
+
+        # Try to get from environment if set by HTTP server
+        query_string = os.environ.get("QUERY_STRING", "")
+        if query_string:
+            return {k: v[0] for k, v in parse_qs(query_string).items()}
+
+        return {}
+
+    def _apply_config(self, query_params: dict):
+        """Apply configuration from query parameters to environment variables."""
+        for param, env_var in self.param_mapping.items():
+            if param in query_params and query_params[param]:
+                os.environ[env_var] = query_params[param]
+                print(f"🔧 Applied config: {param} -> {env_var}")
+
+    def _restore_env(self, original_env: dict):
+        """Restore original environment variables."""
+        for env_var, original_value in original_env.items():
+            if original_value is not None:
+                os.environ[env_var] = original_value
+            elif env_var in os.environ:
+                del os.environ[env_var]
 
 
 def main() -> None:
@@ -85,6 +134,9 @@ def main() -> None:
     # Create MCP server (tools will be registered and available for discovery)
     mcp = create_mcp_server()
 
+    # Add Smithery configuration middleware
+    mcp.add_middleware(SmitheryConfigMiddleware())
+
     # Start HTTP server
     print("\n🚀 Running LinkedIn MCP server (Smithery HTTP mode)...")
     print(f"📡 HTTP server listening on http://0.0.0.0:{port}/mcp")
diff --git a/tests/test_smithery_config.py b/tests/test_smithery_config.py
@@ -0,0 +1,138 @@
+# tests/test_smithery_config.py
+"""
+Test Smithery configuration parameter passing.
+"""
+
+import pytest
+import os
+from unittest.mock import patch, MagicMock
+from fastmcp.client import Client
+from fastmcp.server.middleware import MiddlewareContext
+from linkedin_mcp_server.server import create_mcp_server
+from smithery_main import SmitheryConfigMiddleware
+
+
+@pytest.mark.asyncio
+async def test_smithery_middleware_extracts_config():
+    """Test that SmitheryConfigMiddleware correctly extracts configuration from query parameters."""
+    middleware = SmitheryConfigMiddleware()
+
+    # Mock MiddlewareContext with query parameters via environment
+    context = MagicMock(spec=MiddlewareContext)
+    context.fastmcp_context = None
+
+    # Set query string in environment to simulate HTTP request
+    os.environ["QUERY_STRING"] = (
+        "linkedin_email=test@example.com&linkedin_password=testpass123"
+    )
+
+    # Mock call_next
+    async def mock_call_next(ctx):
+        # During tool execution, check that env vars are set
+        assert os.environ.get("LINKEDIN_EMAIL") == "test@example.com"
+        assert os.environ.get("LINKEDIN_PASSWORD") == "testpass123"
+        return MagicMock()
+
+    # Store original env vars
+    original_email = os.environ.get("LINKEDIN_EMAIL")
+    original_password = os.environ.get("LINKEDIN_PASSWORD")
+    original_query_string = os.environ.get("QUERY_STRING")
+
+    try:
+        # Execute middleware
+        await middleware.on_call_tool(context, mock_call_next)
+
+        # After execution, env vars should be restored
+        assert os.environ.get("LINKEDIN_EMAIL") == original_email
+        assert os.environ.get("LINKEDIN_PASSWORD") == original_password
+
+        print("✅ Smithery middleware correctly handles configuration")
+
+    finally:
+        # Cleanup
+        if original_email is not None:
+            os.environ["LINKEDIN_EMAIL"] = original_email
+        elif "LINKEDIN_EMAIL" in os.environ:
+            del os.environ["LINKEDIN_EMAIL"]
+
+        if original_password is not None:
+            os.environ["LINKEDIN_PASSWORD"] = original_password
+        elif "LINKEDIN_PASSWORD" in os.environ:
+            del os.environ["LINKEDIN_PASSWORD"]
+
+        if original_query_string is not None:
+            os.environ["QUERY_STRING"] = original_query_string
+        elif "QUERY_STRING" in os.environ:
+            del os.environ["QUERY_STRING"]
+
+
+@pytest.mark.asyncio
+async def test_smithery_middleware_with_empty_config():
+    """Test that middleware works correctly with no configuration."""
+    middleware = SmitheryConfigMiddleware()
+
+    # Mock context with no query parameters
+    context = MagicMock(spec=MiddlewareContext)
+    context.fastmcp_context = None
+
+    # Mock call_next
+    async def mock_call_next(ctx):
+        return MagicMock()
+
+    # Should not raise any errors
+    result = await middleware.on_call_tool(context, mock_call_next)
+    assert result is not None
+
+    print("✅ Smithery middleware handles empty configuration")
+
+
+@pytest.mark.asyncio
+async def test_smithery_server_with_middleware():
+    """Test that MCP server with Smithery middleware can be created and tools discovered."""
+    with patch("sys.argv", ["smithery_main.py"]):
+        # Create server (simulate smithery_main.py)
+        mcp = create_mcp_server()
+
+        # Add middleware
+        mcp.add_middleware(SmitheryConfigMiddleware())
+
+        # Test that tools are discoverable
+        async with Client(mcp) as client:
+            tools = await client.list_tools()
+
+            tool_names = [tool.name for tool in tools]
+            expected_tools = [
+                "get_person_profile",
+                "get_company_profile",
+                "get_job_details",
+                "close_session",
+            ]
+
+            for expected_tool in expected_tools:
+                assert expected_tool in tool_names, f"Tool '{expected_tool}' not found"
+
+            print(f"✅ Smithery server with middleware: {len(tools)} tools discovered")
+
+
+def test_smithery_middleware_param_mapping():
+    """Test that SmitheryConfigMiddleware has correct parameter mapping."""
+    middleware = SmitheryConfigMiddleware()
+
+    expected_mapping = {
+        "linkedin_email": "LINKEDIN_EMAIL",
+        "linkedin_password": "LINKEDIN_PASSWORD",
+    }
+
+    assert middleware.param_mapping == expected_mapping
+    print("✅ Smithery middleware parameter mapping is correct")
+
+
+if __name__ == "__main__":
+    # Run tests manually if executed directly
+    import asyncio
+
+    asyncio.run(test_smithery_middleware_extracts_config())
+    asyncio.run(test_smithery_middleware_with_empty_config())
+    asyncio.run(test_smithery_server_with_middleware())
+    test_smithery_middleware_param_mapping()
+    print("🎉 All Smithery configuration tests passed!")
diff --git a/uv.lock b/uv.lock
