Merge pull request #132 from golf-mcp/aschlean/starlette-middleware

Aschlean/starlette middleware

Author: aschlean

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "golf-mcp"
-version = "0.3.0rc3"
+version = "0.3.0rc5"
 description = "Framework for building MCP servers"
 authors = [
     {name = "Antoni Gmitruk", email = "antoni@golf.dev"}
@@ -66,7 +66,7 @@ golf = ["examples/**/*"]
 
 [tool.poetry]
 name = "golf-mcp"
-version = "0.3.0rc3"
+version = "0.3.0rc5"
 description = "Framework for building MCP servers with zero boilerplate"
 authors = ["Antoni Gmitruk <antoni@golf.dev>"]
 license = "Apache-2.0"

src/golf/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.3.0rc3"
+__version__ = "0.3.0rc5"
 
 from golf.decorators import prompt, resource, tool
 

src/golf/core/builder.py

@@ -1227,15 +1227,22 @@ def _generate_server(self) -> None:
             component_registrations.append("")
 
         # Check for custom middleware.py file and register middleware classes
-        middleware_classes = self._discover_middleware_classes(self.project_path)
-        if middleware_classes:
+        discovered_middleware = self._discover_middleware_classes(self.project_path)
+        fastmcp_middleware = discovered_middleware.get("fastmcp", [])
+        starlette_middleware = discovered_middleware.get("starlette", [])
+
+        # Import all middleware classes
+        all_middleware = fastmcp_middleware + starlette_middleware
+        if all_middleware:
             imports.append("# Import custom middleware")
-            imports.append("from middleware import " + ", ".join(middleware_classes))
+            imports.append("from middleware import " + ", ".join(all_middleware))
             imports.append("")
 
-            for cls_name in middleware_classes:
-                # Add each discovered middleware class to FastMCP
-                component_registrations.append(f"# Register custom middleware: {cls_name}")
+        # Register only FastMCP middleware via mcp.add_middleware()
+        # Starlette HTTP middleware will be added to mcp.run(middleware=[...]) later
+        if fastmcp_middleware:
+            for cls_name in fastmcp_middleware:
+                component_registrations.append(f"# Register custom FastMCP middleware: {cls_name}")
                 component_registrations.append(f"mcp.add_middleware({cls_name}())")
             component_registrations.append("")
 
@@ -1343,7 +1350,14 @@ def _generate_server(self) -> None:
                 middleware_setup.append("    from golf.telemetry.instrumentation import SessionTracingMiddleware")
                 middleware_list.append("Middleware(SessionTracingMiddleware)")
 
-            if middleware_setup:
+            # Add custom Starlette HTTP middleware (e.g., CacheControlMiddleware)
+            # These are wrapped in Middleware() and passed to mcp.run(), not mcp.add_middleware()
+            if starlette_middleware:
+                middleware_setup.append("    from starlette.middleware import Middleware")
+                for cls_name in starlette_middleware:
+                    middleware_list.append(f"Middleware({cls_name})")
+
+            if middleware_setup or starlette_middleware:
                 main_code.extend(middleware_setup)
                 main_code.append(f"    middleware = [{', '.join(middleware_list)}]")
                 main_code.append("")
@@ -1405,7 +1419,14 @@ def _generate_server(self) -> None:
                 middleware_setup.append("    from golf.telemetry.instrumentation import SessionTracingMiddleware")
                 middleware_list.append("Middleware(SessionTracingMiddleware)")
 
-            if middleware_setup:
+            # Add custom Starlette HTTP middleware (e.g., CacheControlMiddleware)
+            # These are wrapped in Middleware() and passed to mcp.run(), not mcp.add_middleware()
+            if starlette_middleware:
+                middleware_setup.append("    from starlette.middleware import Middleware")
+                for cls_name in starlette_middleware:
+                    middleware_list.append(f"Middleware({cls_name})")
+
+            if middleware_setup or starlette_middleware:
                 main_code.extend(middleware_setup)
                 main_code.append(f"    middleware = [{', '.join(middleware_list)}]")
                 main_code.append("")
@@ -1493,11 +1514,16 @@ def _generate_server(self) -> None:
         with open(server_file, "w") as f:
             f.write(code)
 
-    def _discover_middleware_classes(self, project_path: Path) -> list[str]:
-        """Discover middleware classes from middleware.py file."""
+    def _discover_middleware_classes(self, project_path: Path) -> dict[str, list[str]]:
+        """Discover middleware classes from middleware.py file.
+
+        Returns a dict with two keys:
+        - 'fastmcp': List of FastMCP middleware class names (use mcp.add_middleware())
+        - 'starlette': List of Starlette HTTP middleware class names (use middleware=[])
+        """
         middleware_path = project_path / "middleware.py"
         if not middleware_path.exists():
-            return []
+            return {"fastmcp": [], "starlette": []}
 
         try:
             # Save current directory and path
@@ -1513,25 +1539,46 @@ def _discover_middleware_classes(self, project_path: Path) -> list[str]:
 
             spec = importlib.util.spec_from_file_location("middleware", middleware_path)
             if spec is None or spec.loader is None:
-                return []
+                return {"fastmcp": [], "starlette": []}
             middleware_module = importlib.util.module_from_spec(spec)
             spec.loader.exec_module(middleware_module)
 
-            # Auto-discover middleware classes using duck typing
-            middleware_classes = []
+            # Auto-discover middleware classes, distinguishing between FastMCP and Starlette
+            fastmcp_middleware = []
+            starlette_middleware = []
+
+            # FastMCP middleware methods (MCP protocol level)
+            fastmcp_methods = [
+                "on_message",
+                "on_request",
+                "on_call_tool",
+                "on_read_resource",
+                "on_get_prompt",
+                "on_initialize",
+            ]
+            # Starlette/ASGI middleware method (HTTP level)
+            starlette_method = "dispatch"
+
             for name, obj in inspect.getmembers(middleware_module, inspect.isclass):
                 # Skip classes that are not defined in this module (imported classes)
                 if obj.__module__ != middleware_module.__name__:
                     continue
 
-                # Check if class actually implements middleware methods (not just inherits them)
-                middleware_methods = ["on_message", "on_request", "on_call_tool", "dispatch"]
-                has_implemented_method = any(method in obj.__dict__ for method in middleware_methods)
-                if has_implemented_method:
-                    middleware_classes.append(name)
-                    console.print(f"[green]Discovered middleware: {name}[/green]")
+                # Check if class implements FastMCP middleware methods
+                has_fastmcp_method = any(method in obj.__dict__ for method in fastmcp_methods)
+                # Check if class implements Starlette dispatch method
+                has_dispatch_method = starlette_method in obj.__dict__
+
+                if has_fastmcp_method and not has_dispatch_method:
+                    # Pure FastMCP middleware
+                    fastmcp_middleware.append(name)
+                    console.print(f"[green]Discovered FastMCP middleware: {name}[/green]")
+                elif has_dispatch_method:
+                    # Starlette/ASGI HTTP middleware (dispatch method indicates HTTP-level)
+                    starlette_middleware.append(name)
+                    console.print(f"[green]Discovered Starlette HTTP middleware: {name}[/green]")
 
-            return middleware_classes
+            return {"fastmcp": fastmcp_middleware, "starlette": starlette_middleware}
 
         except Exception as e:
             console.print(f"[yellow]Warning: Could not load middleware.py: {e}[/yellow]")
@@ -1555,7 +1602,7 @@ def _discover_middleware_classes(self, project_path: Path) -> list[str]:
             except Exception:
                 pass
 
-            return []
+            return {"fastmcp": [], "starlette": []}
 
         finally:
             # Always restore original directory and path

tests/core/test_middleware.py

@@ -24,9 +24,11 @@ class NotMiddleware:
 
         settings = load_settings(sample_project)
         generator = CodeGenerator(sample_project, settings, temp_dir)
-        middleware_classes = generator._discover_middleware_classes(sample_project)
-        
-        assert middleware_classes == ["LoggingMiddleware"]
+        discovered = generator._discover_middleware_classes(sample_project)
+
+        # FastMCP middleware (on_call_tool method)
+        assert discovered["fastmcp"] == ["LoggingMiddleware"]
+        assert discovered["starlette"] == []
 
     def test_discovers_multiple_middleware_methods(self, sample_project: Path, temp_dir: Path):
         """Test discovery of middleware with different methods."""
@@ -49,18 +51,20 @@ def dispatch(self, context):
 
         settings = load_settings(sample_project)
         generator = CodeGenerator(sample_project, settings, temp_dir)
-        middleware_classes = generator._discover_middleware_classes(sample_project)
-        
-        # Should discover all three middleware classes
-        assert set(middleware_classes) == {"MessageMiddleware", "RequestMiddleware", "DispatchMiddleware"}
+        discovered = generator._discover_middleware_classes(sample_project)
+
+        # FastMCP middleware (on_message, on_request methods)
+        assert set(discovered["fastmcp"]) == {"MessageMiddleware", "RequestMiddleware"}
+        # Starlette HTTP middleware (dispatch method)
+        assert discovered["starlette"] == ["DispatchMiddleware"]
 
     def test_no_middleware_when_file_missing(self, sample_project: Path, temp_dir: Path):
         """Test graceful handling when middleware.py doesn't exist."""
         settings = load_settings(sample_project)
         generator = CodeGenerator(sample_project, settings, temp_dir)
-        middleware_classes = generator._discover_middleware_classes(sample_project)
-        
-        assert middleware_classes == []
+        discovered = generator._discover_middleware_classes(sample_project)
+
+        assert discovered == {"fastmcp": [], "starlette": []}
 
     def test_ignores_classes_without_middleware_methods(self, sample_project: Path, temp_dir: Path):
         """Test that classes without middleware methods are ignored."""
@@ -83,9 +87,10 @@ def regular_method(self):
 
         settings = load_settings(sample_project)
         generator = CodeGenerator(sample_project, settings, temp_dir)
-        middleware_classes = generator._discover_middleware_classes(sample_project)
-        
-        assert middleware_classes == ["ValidMiddleware"]
+        discovered = generator._discover_middleware_classes(sample_project)
+
+        assert discovered["fastmcp"] == ["ValidMiddleware"]
+        assert discovered["starlette"] == []
 
 
 class TestMiddlewareCodeGeneration:
@@ -168,9 +173,10 @@ async def on_message(self, context, call_next):
 
         settings = load_settings(sample_project)
         generator = CodeGenerator(sample_project, settings, temp_dir)
-        middleware_classes = generator._discover_middleware_classes(sample_project)
-        
-        assert middleware_classes == []  # Should return empty list on error
+        discovered = generator._discover_middleware_classes(sample_project)
+
+        # Should return empty dict on error
+        assert discovered == {"fastmcp": [], "starlette": []}
 
     def test_handles_import_error(self, sample_project: Path, temp_dir: Path):
         """Test graceful handling of import errors."""
@@ -185,9 +191,10 @@ async def on_message(self, context, call_next):
 
         settings = load_settings(sample_project)
         generator = CodeGenerator(sample_project, settings, temp_dir)
-        middleware_classes = generator._discover_middleware_classes(sample_project)
-        
-        assert middleware_classes == []  # Should return empty list on error
+        discovered = generator._discover_middleware_classes(sample_project)
+
+        # Should return empty dict on error
+        assert discovered == {"fastmcp": [], "starlette": []}
 
     def test_build_succeeds_with_broken_middleware(self, sample_project: Path, temp_dir: Path):
         """Test that broken middleware doesn't break the build process."""
@@ -226,9 +233,10 @@ async def on_message(self, context, call_next):
 
         settings = load_settings(sample_project)
         generator = CodeGenerator(sample_project, settings, temp_dir)
-        middleware_classes = generator._discover_middleware_classes(sample_project)
-        
-        assert middleware_classes == []  # Should return empty list on error
+        discovered = generator._discover_middleware_classes(sample_project)
+
+        # Should return empty dict on error
+        assert discovered == {"fastmcp": [], "starlette": []}
 
     def test_handles_empty_middleware_file(self, sample_project: Path, temp_dir: Path):
         """Test handling of empty middleware.py file."""
@@ -237,9 +245,9 @@ def test_handles_empty_middleware_file(self, sample_project: Path, temp_dir: Pat
 
         settings = load_settings(sample_project)
         generator = CodeGenerator(sample_project, settings, temp_dir)
-        middleware_classes = generator._discover_middleware_classes(sample_project)
-        
-        assert middleware_classes == []
+        discovered = generator._discover_middleware_classes(sample_project)
+
+        assert discovered == {"fastmcp": [], "starlette": []}
 
 
 class TestMiddlewareDuckTyping:
@@ -265,12 +273,15 @@ def some_other_method(self):
 
         settings = load_settings(sample_project)
         generator = CodeGenerator(sample_project, settings, temp_dir)
-        middleware_classes = generator._discover_middleware_classes(sample_project)
-        
-        assert len(middleware_classes) == 2
-        assert "DuckTypedMiddleware" in middleware_classes
-        assert "AlsoMiddleware" in middleware_classes
-        assert "NotMiddleware" not in middleware_classes
+        discovered = generator._discover_middleware_classes(sample_project)
+
+        # FastMCP middleware (on_call_tool method)
+        assert discovered["fastmcp"] == ["DuckTypedMiddleware"]
+        # Starlette HTTP middleware (dispatch method)
+        assert discovered["starlette"] == ["AlsoMiddleware"]
+        # NotMiddleware should not appear in either list
+        all_middleware = discovered["fastmcp"] + discovered["starlette"]
+        assert "NotMiddleware" not in all_middleware
 
     def test_discovers_mixed_middleware_types(self, sample_project: Path, temp_dir: Path):
         """Test discovery of middleware with and without base class."""
@@ -293,12 +304,14 @@ def regular_method(self):
 
         settings = load_settings(sample_project)
         generator = CodeGenerator(sample_project, settings, temp_dir)
-        middleware_classes = generator._discover_middleware_classes(sample_project)
-        
-        assert len(middleware_classes) == 2
-        assert "InheritedMiddleware" in middleware_classes
-        assert "DuckTypedMiddleware" in middleware_classes
-        assert "NoMiddlewareMethods" not in middleware_classes
+        discovered = generator._discover_middleware_classes(sample_project)
+
+        # Both are FastMCP middleware (on_message, on_call_tool methods)
+        assert set(discovered["fastmcp"]) == {"InheritedMiddleware", "DuckTypedMiddleware"}
+        assert discovered["starlette"] == []
+        # NoMiddlewareMethods should not appear
+        all_middleware = discovered["fastmcp"] + discovered["starlette"]
+        assert "NoMiddlewareMethods" not in all_middleware
 
 
 class TestMiddlewareRegistrationOrder: