feat(middleware): enhance middleware discovery and registration

Author: aschlean

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,39 @@ 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 +1595,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:

tests/core/test_middleware_build.py

@@ -302,4 +302,88 @@ def some_method(self):
         assert "mcp.add_middleware(DuckTypedLogging())" in server_content
         assert "mcp.add_middleware(DuckTypedAuth())" in server_content
         # Non-middleware class should not be included
-        assert "NotMiddleware" not in server_content
+        assert "NotMiddleware" not in server_content
+
+    def test_starlette_http_middleware_build(self, sample_project: Path, temp_dir: Path):
+        """Test Starlette HTTP middleware (e.g., CacheControlMiddleware) is registered correctly.
+
+        Starlette middleware uses the dispatch() method and must be passed to mcp.run(middleware=[])
+        instead of mcp.add_middleware() which is for FastMCP protocol-level middleware.
+        """
+        middleware_file = sample_project / "middleware.py"
+        middleware_file.write_text('''
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from starlette.responses import Response
+from typing import Callable, Any
+
+
+class CacheControlMiddleware(BaseHTTPMiddleware):
+    """Middleware to add Cache-Control headers to all responses."""
+
+    async def dispatch(self, request: Request, call_next: Callable[..., Any]) -> Response:
+        response = await call_next(request)
+        response.headers["Cache-Control"] = "no-store"
+        return response
+''')
+
+        settings = load_settings(sample_project)
+        output_dir = temp_dir / "build"
+        build_project(sample_project, settings, output_dir, build_env="dev", copy_env=False)
+
+        server_content = (output_dir / "server.py").read_text()
+
+        # Starlette middleware should be imported
+        assert "from middleware import CacheControlMiddleware" in server_content
+
+        # Should NOT be registered via mcp.add_middleware() - that would fail!
+        assert "mcp.add_middleware(CacheControlMiddleware())" not in server_content
+
+        # Should be added to the middleware list for mcp.run()
+        assert "Middleware(CacheControlMiddleware)" in server_content
+
+    def test_mixed_fastmcp_and_starlette_middleware_build(self, sample_project: Path, temp_dir: Path):
+        """Test that both FastMCP and Starlette middleware can be used together."""
+        middleware_file = sample_project / "middleware.py"
+        middleware_file.write_text('''
+from starlette.middleware.base import BaseHTTPMiddleware
+from golf.middleware import Middleware as FastMCPMiddleware
+from typing import Callable, Any
+
+
+class LoggingMiddleware(FastMCPMiddleware):
+    """FastMCP middleware for logging MCP operations."""
+
+    async def on_call_tool(self, context, call_next):
+        print(f"Calling tool: {context.message.params.name}")
+        return await call_next(context)
+
+
+class SecurityHeadersMiddleware(BaseHTTPMiddleware):
+    """Starlette middleware for adding security headers."""
+
+    async def dispatch(self, request, call_next: Callable[..., Any]):
+        response = await call_next(request)
+        response.headers["Cache-Control"] = "no-store"
+        response.headers["X-Content-Type-Options"] = "nosniff"
+        return response
+''')
+
+        settings = load_settings(sample_project)
+        output_dir = temp_dir / "build"
+        build_project(sample_project, settings, output_dir, build_env="dev", copy_env=False)
+
+        server_content = (output_dir / "server.py").read_text()
+
+        # Both middleware should be imported
+        assert "LoggingMiddleware" in server_content
+        assert "SecurityHeadersMiddleware" in server_content
+
+        # FastMCP middleware should use mcp.add_middleware()
+        assert "mcp.add_middleware(LoggingMiddleware())" in server_content
+
+        # Starlette middleware should NOT use mcp.add_middleware()
+        assert "mcp.add_middleware(SecurityHeadersMiddleware())" not in server_content
+
+        # Starlette middleware should be in the middleware list for mcp.run()
+        assert "Middleware(SecurityHeadersMiddleware)" in server_content