feat: normalize http responses and types

Author: KNedelec

packages/gg_api_core/pyproject.toml

@@ -37,5 +37,5 @@ sentry = [
 ]
 
 [build-system]
-requires = ["hatchling"]
-build-backend = "hatchling.build"
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"

packages/gg_api_core/src/gg_api_core/client.py

@@ -35,14 +35,19 @@ class CachedTokenInfoMixin:
     """Mixin for MCP servers that are mono-tenant (only one authenticated identity from startup to close of the server)"""
 
     _token_scopes: set[str] = set()
-    _token_info = None
+    _token_info: dict[str, Any] | None = None
 
     def __init__(self, *args, **kwargs):
         # Add a custom lifespan contextmanager that fetches and cache token scopes and infos
         original_lifespan = kwargs.get("lifespan")
         kwargs["lifespan"] = self._create_token_scope_lifespan(original_lifespan)
         # Call parent __init__ in the MRO chain
         super().__init__(*args, **kwargs)
+    
+    def clear_cache(self) -> None:
+        """Clear cached token information and scopes."""
+        self._token_scopes = set()
+        self._token_info = None
 
     def _create_token_scope_lifespan(self, original_lifespan=None):
         """Create a lifespan context manager that fetches token scopes."""
@@ -72,12 +77,14 @@ async def token_scope_lifespan(fastmcp) -> AsyncIterator[dict]:
 
         return token_scope_lifespan
 
-    async def get_token_info(self) -> dict:
+    async def get_token_info(self) -> dict[str, Any]:
         """Return the token info dictionary."""
-        if token_info := getattr(self, "_token_info", None):
-            return token_info
+        if self._token_info is not None:
+            return self._token_info
 
-        return await self._fetch_token_info_from_api()
+        # 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 AbstractGitGuardianFastMCP(FastMCP, ABC):
@@ -87,6 +94,8 @@ class AbstractGitGuardianFastMCP(FastMCP, ABC):
     Subclasses implement authentication-specific behavior.
     """
 
+    authentication_mode: AuthenticationMode
+
     def __init__(self, *args, default_scopes: list[str] | None = None, **kwargs):
         """
         Initialize the GitGuardian MCP server.
@@ -97,27 +106,31 @@ 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(self._scope_filtering_middleware)  
+    
+    def clear_cache(self) -> None:
+        """Clear cached data. Override in subclasses that cache."""
+        pass
 
     @abstractmethod
-    def get_personal_access_token(self):
+    def get_personal_access_token(self) -> str | None:
         """Get the personal access token for the current request"""
         pass
 
     @abstractmethod
-    async def get_token_info(self):
+    async def get_token_info(self) -> dict[str, Any]:
         """Return the token info dictionary."""
         pass
 
     def get_client(self):
         return get_client(personal_access_token=self.get_personal_access_token())
 
-    async def revoke_current_token(self) -> dict:
+    async def revoke_current_token(self) -> dict[str, Any]:
         """Revoke the current API token via GitGuardian API."""
         try:
             logger.debug("Revoking current API token")
             # Call the DELETE /api_tokens/self endpoint
-            result = await self.get_client()._request("DELETE", "/api_tokens/self")
+            result = await self.get_client().revoke_current_token()
             logger.debug("API token revoked")
             return result
         except Exception as e:
@@ -157,7 +170,7 @@ def wrapper(fn):
 
         return result
 
-    async def _fetch_token_scopes_from_api(self, client=None):
+    async def _fetch_token_scopes_from_api(self, client=None) -> set[str]:
         """Fetch token scopes from the GitGuardian API.
 
         Args:
@@ -183,19 +196,20 @@ async def _fetch_token_scopes_from_api(self, client=None):
         except Exception as e:
             raise FastMCPError("Error fetching token scopes from /api_tokens/self endpoint") from e
 
-    async def _fetch_token_info_from_api(self) -> dict:
+    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
         except Exception as e:
-            raise FastMCPError("Error fetching token scopes from /api_tokens/self endpoint") from e
+            raise FastMCPError("Error fetching token info from /api_tokens/self endpoint") from e
 
-    async def get_scopes(self):
-        if scopes := getattr(self, "_token_scopes", None):
+    async def get_scopes(self) -> set[str]:
+        cached_scopes: set[str] | None = getattr(self, "_token_scopes", None)
+        if cached_scopes:
             logger.debug("reading from cached scopes")
-            return scopes
+            return cached_scopes
 
         scopes = await self._fetch_token_scopes_from_api()
         logger.debug(f"scopes: {scopes}")
@@ -242,7 +256,7 @@ async def list_tools(self) -> list[MCPTool]:
 
 
 # Common MCP tools for user information and token management
-def register_common_tools(mcp_instance: "AbstractGitGuardianFastMCP"):
+def register_common_tools(mcp_instance: AbstractGitGuardianFastMCP):
     """Register common MCP tools for user information and token management."""
 
     logger.debug("Registering common MCP tools...")
@@ -251,42 +265,37 @@ def register_common_tools(mcp_instance: "AbstractGitGuardianFastMCP"):
         name="get_authenticated_user_info",
         description="Get comprehensive information about the authenticated user and current API token including scopes and authentication method",
     )
-    async def get_authenticated_user_info() -> dict:
+    async def get_authenticated_user_info() -> dict[str, Any]:
         """Get information about the authenticated user and current API token."""
         logger.debug("Getting authenticated user information")
 
         token_info = await mcp_instance.get_token_info()
         scopes = await mcp_instance.get_scopes()
         return {
             "token_info": token_info,
-            "authentication_mode": mcp_instance.authentication_mode,
+            "authentication_mode": mcp_instance.authentication_mode.value,
             "available_scopes": list(scopes),
         }
 
     @mcp_instance.tool(
         name="revoke_current_token",
         description="Revoke the current API token and clean up stored credentials",
     )
-    async def revoke_current_token() -> dict:
+    async def revoke_current_token() -> dict[str, Any]:
         """Revoke the current API token and clean up stored credentials."""
         logger.debug("Starting token revocation process")
 
         try:
-            client = mcp_instance.client
-            await client._request("DELETE", "/api_tokens/self")
+            await mcp_instance.revoke_current_token()
             logger.debug("Token revoked via API")
 
-            # Clear cached client and token info
-            mcp_instance._client = None
-            mcp_instance._token_info = None
-            # Only clear _token_scopes if it exists (cached modes only)
-            if hasattr(mcp_instance, "_token_scopes"):
-                mcp_instance._token_scopes = set()
+            # Clear cached data
+            mcp_instance.clear_cache()
 
             return {
                 "success": True,
                 "message": "Token revoked and credentials cleaned up",
-                "authentication_method": mcp_instance._get_auth_method(),
+                "authentication_method": mcp_instance.authentication_mode.value,
             }
 
         except Exception as e:
@@ -304,7 +313,7 @@ class GitGuardianLocalOAuthMCP(CachedTokenInfoMixin, AbstractGitGuardianFastMCP)
 
     authentication_mode = AuthenticationMode.LOCAL_OAUTH_FLOW
 
-    def get_personal_access_token(self):
+    def get_personal_access_token(self) -> str | None:
         # It will be actually provided within the client by the OAuth flow, or from the filesystem storage
         return None
 
@@ -318,16 +327,16 @@ def __init__(self, *args, personal_access_token: str, **kwargs):
         super().__init__(*args, **kwargs)
         self.personal_access_token = personal_access_token
 
-    def get_personal_access_token(self):
+    def get_personal_access_token(self) -> str:
         return self.personal_access_token
 
 
 class GitGuardianAuthorizationHeaderMCP(AbstractGitGuardianFastMCP):
     """GitGuardian MCP server using per-request Authorization header (HTTP/SSE mode)."""
 
-    authentication_mode = (AuthenticationMode.AUTHORIZATION_HEADER,)
+    authentication_mode = AuthenticationMode.AUTHORIZATION_HEADER
 
-    def get_personal_access_token(self):
+    def get_personal_access_token(self) -> str:
         headers = get_http_headers()
         if not headers:
             raise ValidationError("No HTTP headers available - Authorization header required")
@@ -364,7 +373,7 @@ def _default_extract_token(auth_header: str) -> str | None:
 
         return None
 
-    async def get_token_info(self):
+    async def get_token_info(self) -> dict[str, Any]:
         return await self._fetch_token_info_from_api()
 
 

packages/gg_api_core/src/gg_api_core/mcp_server.py

@@ -95,15 +95,8 @@ async def assign_incident(params: AssignIncidentParams) -> AssignIncidentResult:
         logger.debug(f"Looking up member ID for email: {params.email}")
         try:
             # Use the /members endpoint to search by email
-            result, _ = await client._request("GET", "/members", params={"search": params.email}, return_headers=True)
-
-            # Handle response format
-            if isinstance(result, dict):
-                members = cast(list[dict[str, Any]], result.get("results", result.get("data", [])))
-            elif isinstance(result, list):
-                members = result
-            else:
-                raise ToolError(f"Unexpected response format when searching for member: {type(result).__name__}")
+            result = await client._request_list("/members", params={"search": params.email})
+            members = result["data"]
 
             # Find exact email match
             matching_member: dict[str, Any] | None = None
@@ -150,14 +143,14 @@ async def assign_incident(params: AssignIncidentParams) -> AssignIncidentResult:
 
     try:
         # Call the client method
-        result = await client.assign_incident(incident_id=str(params.incident_id), assignee_id=str(assignee_id))
+        api_result = await client.assign_incident(incident_id=str(params.incident_id), assignee_id=str(assignee_id))
 
         logger.debug(f"Successfully assigned incident {params.incident_id} to member {assignee_id}")
 
         # Parse the response
-        if isinstance(result, dict):
+        if isinstance(api_result, dict):
             # Remove assignee_id from result dict to avoid conflict with our explicit parameter
-            result_copy = result.copy()
+            result_copy = api_result.copy()
             result_copy.pop("assignee_id", None)
             return AssignIncidentResult(
                 incident_id=params.incident_id, assignee_id=assignee_id, success=True, **result_copy

packages/gg_api_core/src/gg_api_core/tools/assign_incident.py

@@ -83,10 +83,7 @@ async def generate_honeytoken(params: GenerateHoneytokenParams) -> GenerateHoney
                 result = await client.list_honeytokens(**filters)
 
                 # Process the result to get the list of tokens
-                if isinstance(result, dict):
-                    honeytokens = result.get("honeytokens", [])
-                else:
-                    honeytokens = result
+                honeytokens = result.get("data", [])
 
                 # Find the most recent active token
                 if honeytokens:
@@ -111,22 +108,22 @@ async def generate_honeytoken(params: GenerateHoneytokenParams) -> GenerateHoney
             {"key": "source", "value": "auto-generated"},
             {"key": "type", "value": "aws"},
         ]
-        result = await client.create_honeytoken(
+        creation_result = await client.create_honeytoken(
             name=params.name, description=params.description, custom_tags=custom_tags
         )
 
         # Validate that we got an ID in the response
-        if not result.get("id"):
+        if not creation_result.get("id"):
             raise ToolError("Failed to get honeytoken ID from GitGuardian API")
 
-        logger.debug(f"Generated new honeytoken with ID: {result.get('id')}")
+        logger.debug(f"Generated new honeytoken with ID: {creation_result.get('id')}")
 
         # Add injection recommendations to the response
-        result["injection_recommendations"] = {
+        creation_result["injection_recommendations"] = {
             "instructions": "Add the honeytoken to your codebase in configuration files, environment variables, or code comments to detect unauthorized access."
         }
 
-        return GenerateHoneytokenResult(**result)
+        return GenerateHoneytokenResult(**creation_result)
     except Exception as e:
         logger.error(f"Error generating honeytoken: {str(e)}")
         raise ToolError(f"Failed to generate honeytoken: {str(e)}")

packages/gg_api_core/src/gg_api_core/tools/generate_honey_token.py

@@ -26,13 +26,15 @@ class ListHoneytokensParams(BaseModel):
     creator_id: str | int | None = Field(default=None, description="Filter by creator ID")
     creator_api_token_id: str | int | None = Field(default=None, description="Filter by creator API token ID")
     per_page: int = Field(default=20, description="Number of results per page (default: 20, min: 1, max: 100)")
+    cursor: str | None = Field(default=None, description="Pagination cursor from a previous response")
     get_all: bool = Field(default=False, description="If True, fetch all results using cursor-based pagination")
 
 
 class ListHoneytokensResult(BaseModel):
     """Result from listing honeytokens."""
 
     honeytokens: list[dict[str, Any]] = Field(description="List of honeytoken objects")
+    next_cursor: str | None = Field(default=None, description="Cursor for fetching the next page (null if no more results)")
 
 
 async def list_honeytokens(params: ListHoneytokensParams) -> ListHoneytokensResult:
@@ -73,48 +75,25 @@ async def list_honeytokens(params: ListHoneytokensParams) -> ListHoneytokensResu
         except Exception as e:
             logger.warning(f"Failed to get current user info for 'mine' filter: {str(e)}")
 
-    # Build filters dictionary with parameters supported by the client API
-    filters: dict[str, Any] = {}
-    if params.status is not None:
-        filters["status"] = params.status
-    if params.search is not None:
-        filters["search"] = params.search
-    if params.ordering is not None:
-        filters["ordering"] = params.ordering
-    if params.show_token is not None:
-        filters["show_token"] = params.show_token
-    if creator_id is not None:
-        filters["creator_id"] = creator_id
-    if params.creator_api_token_id is not None:
-        filters["creator_api_token_id"] = params.creator_api_token_id
-    if params.per_page is not None:
-        filters["per_page"] = params.per_page
-    if params.get_all is not None:
-        filters["get_all"] = params.get_all
-
-    logger.debug(f"Filters: {json.dumps({k: v for k, v in filters.items() if v is not None})}")
 
     try:
-        result = await client.list_honeytokens(
+        response = await client.list_honeytokens(
             status=params.status,
             search=params.search,
             ordering=params.ordering,
             show_token=params.show_token,
-            creator_id=creator_id,
-            creator_api_token_id=params.creator_api_token_id,
+            creator_id=str(creator_id) if creator_id is not None else None,
+            creator_api_token_id=str(params.creator_api_token_id) if params.creator_api_token_id is not None else None,
             per_page=params.per_page,
+            cursor=params.cursor,
             get_all=params.get_all,
         )
 
-        # Handle both response formats: either a dict with 'honeytokens' key or a list directly
-        if isinstance(result, dict):
-            honeytokens = result.get("honeytokens", [])
-        else:
-            # If the result is already a list, use it directly
-            honeytokens = result
-
-        logger.debug(f"Found {len(honeytokens)} honeytokens")
-        return ListHoneytokensResult(honeytokens=honeytokens)
+        honeytokens_data = response["data"]
+        next_cursor = response["cursor"]
+        
+        logger.debug(f"Found {len(honeytokens_data)} honeytokens")
+        return ListHoneytokensResult(honeytokens=honeytokens_data, next_cursor=next_cursor)
     except Exception as e:
         logger.error(f"Error listing honeytokens: {str(e)}")
         raise ToolError(str(e))