refactor(ScopeMiddleware): use a proper Middleware class for ScopeFiltering

Author: KNedelec

packages/gg_api_core/src/gg_api_core/mcp_server.py

@@ -2,18 +2,19 @@
 
 import logging
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterator, Callable
+from collections.abc import AsyncIterator, Callable, Sequence
 from contextlib import asynccontextmanager
 from enum import Enum
 from typing import Any
 
 from fastmcp import FastMCP
 from fastmcp.exceptions import FastMCPError, ValidationError
 from fastmcp.server.dependencies import get_http_headers
-from fastmcp.server.middleware import MiddlewareContext
+from fastmcp.server.middleware import Middleware
+from fastmcp.tools import Tool
 from mcp.types import Tool as MCPTool
 
-from gg_api_core.client import get_personal_access_token_from_env, is_oauth_enabled
+from gg_api_core.client import GitGuardianClient, get_personal_access_token_from_env, is_oauth_enabled
 from gg_api_core.utils import get_client
 
 # Configure logger
@@ -32,7 +33,11 @@ class AuthenticationMode(Enum):
 
 
 class CachedTokenInfoMixin:
-    """Mixin for MCP servers that are mono-tenant (only one authenticated identity from startup to close of the server)"""
+    """Mixin for MCP servers that are mono-tenant (only one authenticated identity from startup to close of the server)
+
+    Note: This mixin expects to be used with AbstractGitGuardianFastMCP which provides
+    _fetch_token_scopes_from_api() and _fetch_token_info_from_api() methods.
+    """
 
     _token_scopes: set[str] = set()
     _token_info: dict[str, Any] | None = None
@@ -65,7 +70,7 @@ async def token_scope_lifespan(fastmcp) -> AsyncIterator[dict]:
 
             # Cache scopes at startup (single token throughout lifespan)
             try:
-                self._token_scopes = await self._fetch_token_scopes_from_api()
+                self._token_scopes = await self._fetch_token_scopes_from_api()  # type: ignore[attr-defined]
                 logger.debug(f"Retrieved token scopes: {self._token_scopes}")
             except Exception as e:
                 logger.warning(f"Failed to fetch token scopes during startup: {str(e)}")
@@ -82,11 +87,41 @@ async def get_token_info(self) -> dict[str, Any]:
         if self._token_info is not None:
             return self._token_info
 
-        # Type ignore because this method will be available via MRO from AbstractGitGuardianFastMCP
         self._token_info = await self._fetch_token_info_from_api()  # type: ignore[attr-defined]
         return self._token_info
 
 
+class ScopeFilteringMiddleware(Middleware):
+    """Middleware to filter tools based on token scopes."""
+
+    def __init__(self, mcp_server: "AbstractGitGuardianFastMCP"):
+        self._mcp_server = mcp_server
+
+    async def on_list_tools(
+        self,
+        context,
+        call_next,
+    ) -> Sequence[Tool]:
+        """Filter tools based on the user's API token scopes."""
+        # Get all tools from the next middleware/handler
+        all_tools = await call_next(context)
+
+        # Filter tools by scopes
+        scopes = await self._mcp_server.get_scopes()
+        filtered_tools: list[Tool] = []
+        for tool in all_tools:
+            tool_name = tool.name
+            required_scopes = self._mcp_server._tool_scopes.get(tool_name, set())
+
+            if not required_scopes or required_scopes.issubset(scopes):
+                filtered_tools.append(tool)
+            else:
+                missing_scopes = required_scopes - scopes
+                logger.info(f"Removing tool '{tool_name}' due to missing scopes: {', '.join(missing_scopes)}")
+
+        return filtered_tools
+
+
 class AbstractGitGuardianFastMCP(FastMCP, ABC):
     """Abstract base class for GitGuardian MCP servers with scope-based tool filtering.
 
@@ -106,7 +141,7 @@ def __init__(self, *args, default_scopes: list[str] | None = None, **kwargs):
         # Map each tool to its required scopes (instance attribute)
         self._tool_scopes: dict[str, set[str]] = {}
 
-        self.add_middleware(self._scope_filtering_middleware)
+        self.add_middleware(ScopeFilteringMiddleware(self))
 
     def clear_cache(self) -> None:
         """Clear cached data. Override in subclasses that cache."""
@@ -122,7 +157,7 @@ async def get_token_info(self) -> dict[str, Any]:
         """Return the token info dictionary."""
         pass
 
-    def get_client(self):
+    def get_client(self) -> GitGuardianClient:
         return get_client(personal_access_token=self.get_personal_access_token())
 
     async def revoke_current_token(self) -> dict[str, Any]:
@@ -199,9 +234,7 @@ async def _fetch_token_scopes_from_api(self, client=None) -> set[str]:
     async def _fetch_token_info_from_api(self) -> dict[str, Any]:
         try:
             client = self.get_client()
-            token_info = await client.get_current_token_info()
-            self._token_info = token_info
-            return token_info
+            return await client.get_current_token_info()
         except Exception as e:
             raise FastMCPError("Error fetching token info from /api_tokens/self endpoint") from e
 
@@ -215,37 +248,6 @@ async def get_scopes(self) -> set[str]:
         logger.debug(f"scopes: {scopes}")
         return scopes
 
-    async def _scope_filtering_middleware(self, context: MiddlewareContext, call_next: Callable) -> Any:
-        """Middleware to filter tools based on token scopes.
-
-        This middleware intercepts tools/list requests and filters the tools
-        based on the user's API token scopes.
-
-        The authentication strategy determines whether to use cached scopes
-        or fetch them per-request.
-        """
-        # Only apply filtering to tools/list requests
-        if context.method != "tools/list":
-            return await call_next(context)
-
-        # Get all tools from the next middleware/handler
-        all_tools = await call_next(context)
-
-        # Filter tools by scopes
-        scopes = await self.get_scopes()
-        filtered_tools = []
-        for tool in all_tools:
-            tool_name = tool.name
-            required_scopes = self._tool_scopes.get(tool_name, set())
-
-            if not required_scopes or required_scopes.issubset(scopes):
-                filtered_tools.append(tool)
-            else:
-                missing_scopes = required_scopes - scopes
-                logger.info(f"Removing tool '{tool_name}' due to missing scopes: {', '.join(missing_scopes)}")
-
-        return filtered_tools
-
     async def list_tools(self) -> list[MCPTool]:
         """
         Public method to list tools (for compatibility with tests and external code).