rootly-mcp-server 2.0.6__py3-none-any.whl → 2.0.8__py3-none-any.whl

rootly_mcp_server/server.py

@@ -7,21 +7,18 @@ the Rootly API's OpenAPI (Swagger) specification using FastMCP's OpenAPI integra
 
 import json
 import os
-import re
 import logging
+from copy import deepcopy
 from pathlib import Path
 import requests
 import httpx
-from typing import Any, Dict, List, Optional, Tuple, Union, Callable, Annotated, Literal
-from enum import Enum
+from typing import Any, Dict, List, Optional, Annotated
 
 from fastmcp import FastMCP
 
-from fastmcp.server.dependencies import get_http_request
-from starlette.requests import Request
-from pydantic import BaseModel, Field
+from pydantic import Field
 
-from .client import RootlyClient
+from .utils import sanitize_parameters_in_spec
 
 # Set up logger
 logger = logging.getLogger(__name__)
@@ -29,29 +26,68 @@ logger = logging.getLogger(__name__)
 # Default Swagger URL
 SWAGGER_URL = "https://rootly-heroku.s3.amazonaws.com/swagger/v1/swagger.json"
 
+# Default allowed API paths
+DEFAULT_ALLOWED_PATHS = [
+    "/incidents/{incident_id}/alerts",
+    "/alerts",
+    "/alerts/{alert_id}",
+    "/severities",
+    "/severities/{severity_id}",
+    "/teams",
+    "/teams/{team_id}",
+    "/services",
+    "/services/{service_id}",
+    "/functionalities",
+    "/functionalities/{functionality_id}",
+    # Incident types
+    "/incident_types",
+    "/incident_types/{incident_type_id}",
+    # Action items (all, by id, by incident)
+    "/incident_action_items",
+    "/incident_action_items/{incident_action_item_id}",
+    "/incidents/{incident_id}/action_items",
+    # Workflows
+    "/workflows",
+    "/workflows/{workflow_id}",
+    # Workflow runs
+    "/workflow_runs",
+    "/workflow_runs/{workflow_run_id}",
+    # Environments
+    "/environments",
+    "/environments/{environment_id}",
+    # Users
+    "/users",
+    "/users/{user_id}",
+    "/users/me",
+    # Status pages
+    "/status_pages",
+    "/status_pages/{status_page_id}",
+]
+
 
 class AuthenticatedHTTPXClient:
-    """An HTTPX client wrapper that handles Rootly API authentication."""
-    
-    def __init__(self, base_url: str = "https://api.rootly.com", hosted: bool = False):
+    """An HTTPX client wrapper that handles Rootly API authentication and parameter transformation."""
+
+    def __init__(self, base_url: str = "https://api.rootly.com", hosted: bool = False, parameter_mapping: Optional[Dict[str, str]] = None):
         self.base_url = base_url
         self.hosted = hosted
         self._api_token = None
-        
+        self.parameter_mapping = parameter_mapping or {}
+
         if not self.hosted:
             self._api_token = self._get_api_token()
-        
+
         # Create the HTTPX client
-        headers = {"Content-Type": "application/json", "Accept": "application/json"}
+        headers = {"Content-Type": "application/vnd.api+json", "Accept": "application/vnd.api+json"}
         if self._api_token:
             headers["Authorization"] = f"Bearer {self._api_token}"
-            
+
         self.client = httpx.AsyncClient(
             base_url=base_url,
             headers=headers,
             timeout=30.0
         )
-    
+
     def _get_api_token(self) -> Optional[str]:
         """Get the API token from environment variables."""
         api_token = os.getenv("ROOTLY_API_TOKEN")
@@ -59,13 +95,56 @@ class AuthenticatedHTTPXClient:
             logger.warning("ROOTLY_API_TOKEN environment variable is not set")
             return None
         return api_token
-    
+
+    def _transform_params(self, params: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
+        """Transform sanitized parameter names back to original names."""
+        if not params or not self.parameter_mapping:
+            return params
+
+        transformed = {}
+        for key, value in params.items():
+            # Use the original name if we have a mapping, otherwise keep the sanitized name
+            original_key = self.parameter_mapping.get(key, key)
+            transformed[original_key] = value
+            if original_key != key:
+                logger.debug(f"Transformed parameter: '{key}' -> '{original_key}'")
+        return transformed
+
+    async def request(self, method: str, url: str, **kwargs):
+        """Override request to transform parameters."""
+        # Transform query parameters
+        if 'params' in kwargs:
+            kwargs['params'] = self._transform_params(kwargs['params'])
+
+        # Call the underlying client's request method
+        return await self.client.request(method, url, **kwargs)
+
+    async def get(self, url: str, **kwargs):
+        """Proxy to request with GET method."""
+        return await self.request('GET', url, **kwargs)
+
+    async def post(self, url: str, **kwargs):
+        """Proxy to request with POST method."""
+        return await self.request('POST', url, **kwargs)
+
+    async def put(self, url: str, **kwargs):
+        """Proxy to request with PUT method."""
+        return await self.request('PUT', url, **kwargs)
+
+    async def patch(self, url: str, **kwargs):
+        """Proxy to request with PATCH method."""
+        return await self.request('PATCH', url, **kwargs)
+
+    async def delete(self, url: str, **kwargs):
+        """Proxy to request with DELETE method."""
+        return await self.request('DELETE', url, **kwargs)
+
     async def __aenter__(self):
-        return self.client
-    
+        return self
+
     async def __aexit__(self, exc_type, exc_val, exc_tb):
-        await self.client.aclose()
-    
+        pass
+
     def __getattr__(self, name):
         # Delegate all other attributes to the underlying client
         return getattr(self.client, name)
@@ -93,44 +172,8 @@ def create_rootly_mcp_server(
     """
     # Set default allowed paths if none provided
     if allowed_paths is None:
-        allowed_paths = [
-            "/incidents",
-            "/incidents/{incident_id}/alerts",
-            "/alerts",
-            "/alerts/{alert_id}",
-            "/severities",
-            "/severities/{severity_id}",
-            "/teams",
-            "/teams/{team_id}",
-            "/services",
-            "/services/{service_id}",
-            "/functionalities",
-            "/functionalities/{functionality_id}",
-            # Incident types
-            "/incident_types",
-            "/incident_types/{incident_type_id}",
-            # Action items (all, by id, by incident)
-            "/incident_action_items",
-            "/incident_action_items/{incident_action_item_id}",
-            "/incidents/{incident_id}/action_items",
-            # Workflows
-            "/workflows",
-            "/workflows/{workflow_id}",
-            # Workflow runs
-            "/workflow_runs",
-            "/workflow_runs/{workflow_run_id}",
-            # Environments
-            "/environments",
-            "/environments/{environment_id}",
-            # Users
-            "/users",
-            "/users/{user_id}",
-            "/users/me",
-            # Status pages
-            "/status_pages",
-            "/status_pages/{status_page_id}",
-        ]
-    
+        allowed_paths = DEFAULT_ALLOWED_PATHS
+
     # Add /v1 prefix to paths if not present
     allowed_paths_v1 = [
         f"/v1{path}" if not path.startswith("/v1") else path
@@ -147,22 +190,23 @@ def create_rootly_mcp_server(
     filtered_spec = _filter_openapi_spec(swagger_spec, allowed_paths_v1)
     logger.info(f"Filtered spec to {len(filtered_spec.get('paths', {}))} allowed paths")
 
+    # Sanitize all parameter names in the filtered spec to be MCP-compliant
+    parameter_mapping = sanitize_parameters_in_spec(filtered_spec)
+    logger.info(f"Sanitized parameter names for MCP compatibility (mapped {len(parameter_mapping)} parameters)")
+
     # Determine the base URL
     if base_url is None:
         base_url = os.getenv("ROOTLY_BASE_URL", "https://api.rootly.com")
-    
+
     logger.info(f"Using Rootly API base URL: {base_url}")
 
-    # Create the authenticated HTTP client
-    try:
-        http_client = AuthenticatedHTTPXClient(
-            base_url=base_url,
-            hosted=hosted
-        )
-    except Exception as e:
-        logger.warning(f"Failed to create authenticated client: {e}")
-        # Create a mock client for testing
-        http_client = httpx.AsyncClient(base_url=base_url)
+    # Create the authenticated HTTP client with parameter mapping
+
+    http_client = AuthenticatedHTTPXClient(
+        base_url=base_url,
+        hosted=hosted,
+        parameter_mapping=parameter_mapping
+    )
 
     # Create the MCP server using OpenAPI integration
     # By default, all routes become tools which is what we want
@@ -173,10 +217,16 @@ def create_rootly_mcp_server(
         timeout=30.0,
         tags={"rootly", "incident-management"},
     )
-
+    
+    @mcp.custom_route("/healthz", methods=["GET"])
+    @mcp.custom_route("/health", methods=["GET"])
+    async def health_check(request):
+        from starlette.responses import PlainTextResponse
+        return PlainTextResponse("OK")
+    
     # Add some custom tools for enhanced functionality
     @mcp.tool()
-    def list_endpoints() -> str:
+    def list_endpoints() -> list:
         """List all available Rootly API endpoints with their descriptions."""
         endpoints = []
         for path, path_item in filtered_spec.get("paths", {}).items():
@@ -186,7 +236,7 @@ def create_rootly_mcp_server(
 
                 summary = operation.get("summary", "")
                 description = operation.get("description", "")
-                
+
                 endpoints.append({
                     "path": path,
                     "method": method.upper(),
@@ -194,106 +244,121 @@ def create_rootly_mcp_server(
                     "description": description,
                 })
 
-        return json.dumps(endpoints, indent=2)
+        return endpoints
+
+    async def make_authenticated_request(method: str, url: str, **kwargs):
+        """Make an authenticated request, extracting token from MCP headers in hosted mode."""
+        # In hosted mode, get token from MCP request headers
+        if hosted:
+            try:
+                from fastmcp.server.dependencies import get_http_headers
+                request_headers = get_http_headers()
+                auth_header = request_headers.get("authorization", "")
+                if auth_header:
+                    # Add authorization header to the request
+                    if "headers" not in kwargs:
+                        kwargs["headers"] = {}
+                    kwargs["headers"]["Authorization"] = auth_header
+            except Exception:
+                pass  # Fallback to default client behavior
+        
+        # Make the request using the underlying httpx client
+        return await http_client.client.request(method, url, **kwargs)
 
     @mcp.tool()
-    async def search_incidents_paginated(
+    async def search_incidents(
         query: Annotated[str, Field(description="Search query to filter incidents by title/summary")] = "",
-        page_size: Annotated[int, Field(description="Number of results per page (max: 100)", ge=1, le=100)] = 100,
-        page_number: Annotated[int, Field(description="Page number to retrieve", ge=1)] = 1,
-    ) -> str:
+        page_size: Annotated[int, Field(description="Number of results per page (max: 20)", ge=1, le=20)] = 10,
+        page_number: Annotated[int, Field(description="Page number to retrieve (use 0 for all pages)", ge=0)] = 1,
+        max_results: Annotated[int, Field(description="Maximum total results when fetching all pages (ignored if page_number > 0)", ge=1, le=100)] = 20,
+    ) -> dict:
         """
-        Search incidents with enhanced pagination control.
-        
-        This tool provides better pagination handling than the standard API endpoint.
+        Search incidents with flexible pagination control.
+
+        Use page_number=0 to fetch all matching results across multiple pages up to max_results.
+        Use page_number>0 to fetch a specific page.
         """
-        params = {
-            "page[size]": min(page_size, 100),
-            "page[number]": page_number,
-        }
-        if query:
-            params["filter[search]"] = query
-        
-        try:
-            async with http_client as client:
-                response = await client.get("/v1/incidents", params=params)
+        # Single page mode
+        if page_number > 0:
+            params = {
+                "page[size]": min(page_size, 20),
+                "page[number]": page_number,
+                "include": "",
+            }
+            if query:
+                params["filter[search]"] = query
+
+            try:
+                response = await make_authenticated_request("GET", "/v1/incidents", params=params)
                 response.raise_for_status()
-                result = response.json()
-        except Exception as e:
-            result = {"error": str(e)}
-        
-        return json.dumps(result, indent=2)
+                return response.json()
+            except Exception as e:
+                return {"error": str(e)}
 
-    @mcp.tool()
-    async def get_all_incidents_matching(
-        query: Annotated[str, Field(description="Search query to filter incidents by title/summary")] = "",
-        max_results: Annotated[int, Field(description="Maximum number of results to return", ge=1, le=1000)] = 500,
-    ) -> str:
-        """
-        Get all incidents matching a query by automatically fetching multiple pages.
-        
-        This tool automatically handles pagination to fetch multiple pages of results.
-        """
+        # Multi-page mode (page_number = 0)
         all_incidents = []
-        page_number = 1
-        page_size = 100
-        
+        current_page = 1
+        effective_page_size = min(page_size, 10)
+
         try:
-            async with http_client as client:
-                while len(all_incidents) < max_results:
-                    params = {
-                        "page[size]": page_size,
-                        "page[number]": page_number,
-                    }
-                    if query:
-                        params["filter[search]"] = query
-                    
-                    try:
-                        response = await client.get("/v1/incidents", params=params)
-                        response.raise_for_status()
-                        response_data = response.json()
+            while len(all_incidents) < max_results:
+                params = {
+                    "page[size]": effective_page_size,
+                    "page[number]": current_page,
+                    "include": "",
+                }
+                if query:
+                    params["filter[search]"] = query
+
+                try:
+                    response = await make_authenticated_request("GET", "/v1/incidents", params=params)
+                    response.raise_for_status()
+                    response_data = response.json()
+
+                    if "data" in response_data:
+                        incidents = response_data["data"]
+                        if not incidents:
+                            break
                         
-                        if "data" in response_data:
-                            incidents = response_data["data"]
-                            if not incidents:  # No more results
-                                break
-                            all_incidents.extend(incidents)
-                            
-                            # Check if we have more pages
-                            meta = response_data.get("meta", {})
-                            current_page = meta.get("current_page", page_number)
-                            total_pages = meta.get("total_pages", 1)
-                            
-                            if current_page >= total_pages:
-                                break  # No more pages
-                                
-                            page_number += 1
-                        else:
-                            break  # Unexpected response format
-                            
-                    except Exception as e:
-                        logger.error(f"Error fetching incidents page {page_number}: {e}")
+                        all_incidents.extend(incidents)
+
+                        # Check if we have more pages
+                        meta = response_data.get("meta", {})
+                        current_page_meta = meta.get("current_page", current_page)
+                        total_pages = meta.get("total_pages", 1)
+
+                        if current_page_meta >= total_pages:
+                            break
+
+                        current_page += 1
+                    else:
                         break
-                
-                # Limit to max_results
-                if len(all_incidents) > max_results:
-                    all_incidents = all_incidents[:max_results]
-                
-                result = {
-                    "data": all_incidents,
-                    "meta": {
-                        "total_fetched": len(all_incidents),
-                        "max_results": max_results,
-                        "query": query
-                    }
+
+                except Exception as e:
+                    # Re-raise authentication or critical errors
+                    if "401" in str(e) or "Unauthorized" in str(e) or "authentication" in str(e).lower():
+                        raise e
+                    break
+
+            # Limit to max_results
+            if len(all_incidents) > max_results:
+                all_incidents = all_incidents[:max_results]
+
+            return {
+                "data": all_incidents,
+                "meta": {
+                    "total_fetched": len(all_incidents),
+                    "max_results": max_results,
+                    "query": query,
+                    "pages_fetched": current_page - 1,
+                    "page_size": effective_page_size
                 }
+            }
         except Exception as e:
-            result = {"error": str(e)}
-        
-        return json.dumps(result, indent=2)
+            return {"error": str(e)}
 
     # Log server creation (tool count will be shown when tools are accessed)
-    logger.info(f"Created Rootly MCP Server successfully")
+    logger.info("Created Rootly MCP Server successfully")
     return mcp
 
 
@@ -395,25 +460,26 @@ def _filter_openapi_spec(spec: Dict[str, Any], allowed_paths: List[str]) -> Dict
     Returns:
         A filtered OpenAPI specification with cleaned schema references.
     """
-    filtered_spec = spec.copy()
-    
+    # Use deepcopy to ensure all nested structures are properly copied
+    filtered_spec = deepcopy(spec)
+
     # Filter paths
-    original_paths = spec.get("paths", {})
+    original_paths = filtered_spec.get("paths", {})
     filtered_paths = {
         path: path_item
         for path, path_item in original_paths.items()
         if path in allowed_paths
     }
-    
+
     filtered_spec["paths"] = filtered_paths
-    
+
     # Clean up schema references that might be broken
     # Remove problematic schema references from request bodies and parameters
     for path, path_item in filtered_paths.items():
         for method, operation in path_item.items():
             if method.lower() not in ["get", "post", "put", "delete", "patch"]:
                 continue
-                
+
             # Clean request body schemas
             if "requestBody" in operation:
                 request_body = operation["requestBody"]
@@ -429,8 +495,21 @@ def _filter_openapi_spec(spec: Dict[str, Any], allowed_paths: List[str]) -> Dict
                                     "description": "Request parameters for this endpoint",
                                     "additionalProperties": True
                                 }
-            
-            # Clean parameter schemas
+
+            # Remove response schemas to avoid validation issues
+            # FastMCP will still return the data, just without strict validation
+            if "responses" in operation:
+                for status_code, response in operation["responses"].items():
+                    if "content" in response:
+                        for content_type, content_info in response["content"].items():
+                            if "schema" in content_info:
+                                # Replace with a simple schema that accepts any response
+                                content_info["schema"] = {
+                                    "type": "object",
+                                    "additionalProperties": True
+                                }
+
+            # Clean parameter schemas (parameter names are already sanitized)
             if "parameters" in operation:
                 for param in operation["parameters"]:
                     if "schema" in param and "$ref" in param["schema"]:
@@ -441,7 +520,7 @@ def _filter_openapi_spec(spec: Dict[str, Any], allowed_paths: List[str]) -> Dict
                                 "type": "string",
                                 "description": param.get("description", "Parameter value")
                             }
-    
+
     # Also clean up any remaining broken references in components
     if "components" in filtered_spec and "schemas" in filtered_spec["components"]:
         schemas = filtered_spec["components"]["schemas"]
@@ -450,11 +529,11 @@ def _filter_openapi_spec(spec: Dict[str, Any], allowed_paths: List[str]) -> Dict
         for schema_name, schema_def in schemas.items():
             if isinstance(schema_def, dict) and _has_broken_references(schema_def):
                 schemas_to_remove.append(schema_name)
-        
+
         for schema_name in schemas_to_remove:
             logger.warning(f"Removing schema with broken references: {schema_name}")
             del schemas[schema_name]
-    
+
     return filtered_spec
 
 
@@ -465,13 +544,13 @@ def _has_broken_references(schema_def: Dict[str, Any]) -> bool:
         # List of known broken references in the Rootly API spec
         broken_refs = [
             "incident_trigger_params",
-            "new_workflow", 
+            "new_workflow",
             "update_workflow",
             "workflow"
         ]
         if any(broken_ref in ref_path for broken_ref in broken_refs):
             return True
-    
+
     # Recursively check nested schemas
     for key, value in schema_def.items():
         if isinstance(value, dict):
@@ -481,7 +560,7 @@ def _has_broken_references(schema_def: Dict[str, Any]) -> bool:
             for item in value:
                 if isinstance(item, dict) and _has_broken_references(item):
                     return True
-    
+
     return False
 
 
@@ -489,10 +568,10 @@ def _has_broken_references(schema_def: Dict[str, Any]) -> bool:
 class RootlyMCPServer(FastMCP):
     """
     Legacy Rootly MCP Server class for backward compatibility.
-    
+
     This class is deprecated. Use create_rootly_mcp_server() instead.
     """
-    
+
     def __init__(
         self,
         swagger_path: Optional[str] = None,
@@ -506,7 +585,7 @@ class RootlyMCPServer(FastMCP):
         logger.warning(
             "RootlyMCPServer class is deprecated. Use create_rootly_mcp_server() function instead."
         )
-        
+
         # Create the server using the new function
         server = create_rootly_mcp_server(
             swagger_path=swagger_path,
@@ -514,7 +593,7 @@ class RootlyMCPServer(FastMCP):
             allowed_paths=allowed_paths,
             hosted=hosted
         )
-        
+
         # Copy the server's state to this instance
         super().__init__(name, *args, **kwargs)
         # For compatibility, store reference to the new server

rootly_mcp_server/test_client.py

@@ -33,7 +33,7 @@ async def test_server():
             hosted=False  # Use local API token
         )
         
-        print(f"✅ Server created successfully")
+        print("✅ Server created successfully")
         print(f"Server type: {type(server)}")
         
         # Use the get_tools method to access tools
@@ -78,11 +78,11 @@ async def test_server():
                         if props:
                             print(f"    Parameters: {', '.join(props.keys())}")
             else:
-                print(f"\n⚠️  No tools found")
+                print("\n⚠️  No tools found")
                 
             # Test accessing a specific tool
             if tool_count > 0:
-                print(f"\n🔍 Testing tool access...")
+                print("\n🔍 Testing tool access...")
                 if isinstance(tools, dict):
                     first_tool_name = tools_names[0]
                     first_tool = tools[first_tool_name]
@@ -93,15 +93,17 @@ async def test_server():
                 print(f"  ✅ First tool: {first_tool_name}")
                 print(f"  Tool details: {first_tool}")
                 
-                # Try to get tool by name
+                # Try to get tools and find the specific tool
                 try:
-                    retrieved_tool = await server.get_tool(first_tool_name)
-                    if retrieved_tool:
+                    all_tools = await server.get_tools()
+                    if first_tool_name in all_tools:
+                        retrieved_tool = all_tools[first_tool_name]
                         print(f"  ✅ Successfully retrieved tool by name: {first_tool_name}")
+                        print(f"  Retrieved tool type: {type(retrieved_tool)}")
                     else:
-                        print(f"  ❌ Could not retrieve tool by name: {first_tool_name}")
+                        print(f"  ❌ Could not find tool by name: {first_tool_name}")
                 except Exception as e:
-                    print(f"  ❌ Error retrieving tool: {e}")
+                    print(f"  ❌ Error retrieving tools: {e}")
             
         except Exception as e:
             print(f"❌ Error accessing tools: {e}")
@@ -109,7 +111,7 @@ async def test_server():
             traceback.print_exc()
             tool_count = 0
         
-        print(f"\n🎉 Test completed successfully!")
+        print("\n🎉 Test completed successfully!")
         print(f"Total tools found: {tool_count}")
         
     except Exception as e: