agent-mcp-gateway 0.1.5__py3-none-any.whl

src/metrics.py

@@ -0,0 +1,299 @@
+"""Metrics collection for Agent MCP Gateway."""
+
+import asyncio
+from dataclasses import dataclass, field
+from typing import Dict, List
+
+
+@dataclass
+class OperationMetrics:
+    """Metrics for a specific operation.
+
+    Attributes:
+        count: Total number of operations recorded
+        total_latency_ms: Cumulative latency in milliseconds
+        latencies: List of individual latency measurements
+        errors: Number of operations that resulted in errors
+    """
+    count: int = 0
+    total_latency_ms: float = 0.0
+    latencies: List[float] = field(default_factory=list)
+    errors: int = 0
+
+    def record(self, latency_ms: float, is_error: bool = False):
+        """Record a single operation.
+
+        Args:
+            latency_ms: Operation latency in milliseconds
+            is_error: Whether the operation resulted in an error
+        """
+        self.count += 1
+        self.total_latency_ms += latency_ms
+        self.latencies.append(latency_ms)
+        if is_error:
+            self.errors += 1
+
+    def get_summary(self) -> dict:
+        """Generate summary statistics for this operation.
+
+        Returns:
+            Dictionary containing count, avg, percentiles, and error_rate
+        """
+        if self.count == 0:
+            return {
+                "count": 0,
+                "avg_latency_ms": 0.0,
+                "p50_latency_ms": 0.0,
+                "p95_latency_ms": 0.0,
+                "p99_latency_ms": 0.0,
+                "error_rate": 0.0
+            }
+
+        avg_latency = self.total_latency_ms / self.count
+        error_rate = self.errors / self.count
+
+        # Calculate percentiles
+        sorted_latencies = sorted(self.latencies)
+        p50 = self._percentile(sorted_latencies, 50)
+        p95 = self._percentile(sorted_latencies, 95)
+        p99 = self._percentile(sorted_latencies, 99)
+
+        return {
+            "count": self.count,
+            "avg_latency_ms": round(avg_latency, 2),
+            "p50_latency_ms": round(p50, 2),
+            "p95_latency_ms": round(p95, 2),
+            "p99_latency_ms": round(p99, 2),
+            "error_rate": round(error_rate, 4)
+        }
+
+    @staticmethod
+    def _percentile(sorted_values: List[float], percentile: int) -> float:
+        """Calculate percentile from sorted values.
+
+        Args:
+            sorted_values: List of values sorted in ascending order
+            percentile: Percentile to calculate (0-100)
+
+        Returns:
+            Value at the specified percentile
+        """
+        if not sorted_values:
+            return 0.0
+
+        if len(sorted_values) == 1:
+            return sorted_values[0]
+
+        # Use linear interpolation method
+        k = (len(sorted_values) - 1) * (percentile / 100.0)
+        f = int(k)
+        c = f + 1
+
+        if c >= len(sorted_values):
+            return sorted_values[-1]
+
+        # Interpolate between floor and ceiling
+        d0 = sorted_values[f] * (c - k)
+        d1 = sorted_values[c] * (k - f)
+
+        return d0 + d1
+
+
+class MetricsCollector:
+    """Collects and aggregates metrics for gateway operations.
+
+    Thread-safe metrics collection with per-agent and per-operation tracking.
+    Metrics are stored in memory and can be aggregated for monitoring.
+    """
+
+    def __init__(self):
+        """Initialize metrics collector with empty storage."""
+        # Overall metrics: operation -> OperationMetrics
+        self._metrics: Dict[str, OperationMetrics] = {}
+
+        # Per-agent metrics: agent_id -> operation -> OperationMetrics
+        self._agent_metrics: Dict[str, Dict[str, OperationMetrics]] = {}
+
+        # Lock for thread-safe operations
+        self._lock = asyncio.Lock()
+
+    async def record(
+        self,
+        agent_id: str,
+        operation: str,
+        latency_ms: float,
+        is_error: bool = False
+    ):
+        """Record a single operation metric.
+
+        Args:
+            agent_id: Agent identifier
+            operation: Operation name (list_servers, execute_tool, etc.)
+            latency_ms: Operation latency in milliseconds
+            is_error: Whether the operation resulted in an error
+        """
+        async with self._lock:
+            # Record overall metrics
+            if operation not in self._metrics:
+                self._metrics[operation] = OperationMetrics()
+            self._metrics[operation].record(latency_ms, is_error)
+
+            # Record per-agent metrics
+            if agent_id not in self._agent_metrics:
+                self._agent_metrics[agent_id] = {}
+            if operation not in self._agent_metrics[agent_id]:
+                self._agent_metrics[agent_id][operation] = OperationMetrics()
+            self._agent_metrics[agent_id][operation].record(latency_ms, is_error)
+
+    def record_sync(
+        self,
+        agent_id: str,
+        operation: str,
+        latency_ms: float,
+        is_error: bool = False
+    ):
+        """Record a single operation metric (synchronous version).
+
+        Note: This is not thread-safe. Use async record() for concurrent access.
+
+        Args:
+            agent_id: Agent identifier
+            operation: Operation name (list_servers, execute_tool, etc.)
+            latency_ms: Operation latency in milliseconds
+            is_error: Whether the operation resulted in an error
+        """
+        # Record overall metrics
+        if operation not in self._metrics:
+            self._metrics[operation] = OperationMetrics()
+        self._metrics[operation].record(latency_ms, is_error)
+
+        # Record per-agent metrics
+        if agent_id not in self._agent_metrics:
+            self._agent_metrics[agent_id] = {}
+        if operation not in self._agent_metrics[agent_id]:
+            self._agent_metrics[agent_id][operation] = OperationMetrics()
+        self._agent_metrics[agent_id][operation].record(latency_ms, is_error)
+
+    async def get_summary(self) -> dict:
+        """Get overall summary of all operations.
+
+        Returns:
+            Dictionary mapping operation names to their summary statistics
+        """
+        async with self._lock:
+            return self._get_summary_internal()
+
+    def get_summary_sync(self) -> dict:
+        """Get overall summary of all operations (synchronous version).
+
+        Returns:
+            Dictionary mapping operation names to their summary statistics
+        """
+        return self._get_summary_internal()
+
+    def _get_summary_internal(self) -> dict:
+        """Internal method to get summary without locking."""
+        return {
+            operation: metrics.get_summary()
+            for operation, metrics in self._metrics.items()
+        }
+
+    async def get_agent_summary(self, agent_id: str) -> dict:
+        """Get summary for a specific agent.
+
+        Args:
+            agent_id: Agent identifier
+
+        Returns:
+            Dictionary mapping operation names to summary statistics for this agent,
+            or empty dict if agent has no recorded metrics
+        """
+        async with self._lock:
+            return self._get_agent_summary_internal(agent_id)
+
+    def get_agent_summary_sync(self, agent_id: str) -> dict:
+        """Get summary for a specific agent (synchronous version).
+
+        Args:
+            agent_id: Agent identifier
+
+        Returns:
+            Dictionary mapping operation names to summary statistics for this agent,
+            or empty dict if agent has no recorded metrics
+        """
+        return self._get_agent_summary_internal(agent_id)
+
+    def _get_agent_summary_internal(self, agent_id: str) -> dict:
+        """Internal method to get agent summary without locking."""
+        if agent_id not in self._agent_metrics:
+            return {}
+
+        return {
+            operation: metrics.get_summary()
+            for operation, metrics in self._agent_metrics[agent_id].items()
+        }
+
+    async def get_operation_summary(self, operation: str) -> dict:
+        """Get summary for a specific operation.
+
+        Args:
+            operation: Operation name
+
+        Returns:
+            Summary statistics for this operation, or empty metrics if not found
+        """
+        async with self._lock:
+            return self._get_operation_summary_internal(operation)
+
+    def get_operation_summary_sync(self, operation: str) -> dict:
+        """Get summary for a specific operation (synchronous version).
+
+        Args:
+            operation: Operation name
+
+        Returns:
+            Summary statistics for this operation, or empty metrics if not found
+        """
+        return self._get_operation_summary_internal(operation)
+
+    def _get_operation_summary_internal(self, operation: str) -> dict:
+        """Internal method to get operation summary without locking."""
+        if operation not in self._metrics:
+            return {
+                "count": 0,
+                "avg_latency_ms": 0.0,
+                "p50_latency_ms": 0.0,
+                "p95_latency_ms": 0.0,
+                "p99_latency_ms": 0.0,
+                "error_rate": 0.0
+            }
+
+        return self._metrics[operation].get_summary()
+
+    async def get_all_agents(self) -> List[str]:
+        """Get list of all agents with recorded metrics.
+
+        Returns:
+            List of agent identifiers
+        """
+        async with self._lock:
+            return list(self._agent_metrics.keys())
+
+    def get_all_agents_sync(self) -> List[str]:
+        """Get list of all agents with recorded metrics (synchronous version).
+
+        Returns:
+            List of agent identifiers
+        """
+        return list(self._agent_metrics.keys())
+
+    async def reset(self):
+        """Reset all metrics (useful for testing)."""
+        async with self._lock:
+            self._metrics.clear()
+            self._agent_metrics.clear()
+
+    def reset_sync(self):
+        """Reset all metrics (synchronous version, useful for testing)."""
+        self._metrics.clear()
+        self._agent_metrics.clear()

src/middleware.py

@@ -0,0 +1,166 @@
+"""Access control middleware for Agent MCP Gateway.
+
+This module implements the AgentAccessControl middleware that enforces
+per-agent access rules for gateway tools. It extracts agent identity from
+tool call arguments, validates permissions, and manages agent context state.
+"""
+
+from fastmcp.server.middleware import Middleware, MiddlewareContext
+from fastmcp.exceptions import ToolError
+from .policy import PolicyEngine
+
+
+class AgentAccessControl(Middleware):
+    """Enforces per-agent access rules for gateway operations.
+
+    This middleware intercepts tool calls to:
+    1. Extract agent_id from arguments
+    2. Validate agent identity (handle missing agent_id based on default policy)
+    3. Store agent in context state for downstream use
+    4. Remove agent_id from arguments before forwarding to tools
+    5. Allow gateway tools to perform their own authorization
+
+    Gateway tools (list_servers, get_server_tools, execute_tool) handle their
+    own permission checks, so the middleware just extracts and cleans agent_id
+    without blocking them.
+    """
+
+    def __init__(self, policy_engine: PolicyEngine):
+        """Initialize middleware with policy engine.
+
+        Args:
+            policy_engine: PolicyEngine instance for evaluating access rules
+        """
+        self.policy_engine = policy_engine
+
+    async def on_call_tool(self, context: MiddlewareContext, call_next):
+        """Intercept tool calls to extract and validate agent identity.
+
+        This hook:
+        - Extracts agent_id from tool arguments
+        - Validates agent identity based on default policy
+        - Applies fallback chain if agent_id is missing (when deny_on_missing_agent: false)
+        - Stores agent in context state for downstream tools
+        - Keeps agent_id in arguments (gateway tools need it)
+        - Allows gateway tools to pass through (they do own auth)
+
+        Args:
+            context: Middleware context containing the tool call message
+            call_next: Callable to invoke next middleware/handler in chain
+
+        Returns:
+            Result from downstream handler
+
+        Raises:
+            ToolError: If agent_id is missing and default policy denies access,
+                      or if fallback chain fails to find a valid agent
+        """
+        # Extract the tool call message
+        tool_call = context.message
+        arguments = tool_call.arguments or {}
+
+        # Extract agent_id from arguments
+        agent_id = arguments.get("agent_id")
+
+        # Handle missing agent_id based on default policy
+        if not agent_id:
+            # Check if default policy allows missing agents
+            deny_on_missing = self.policy_engine.defaults.get("deny_on_missing_agent", True)
+            if deny_on_missing:
+                raise ToolError(
+                    "Missing required parameter 'agent_id'. "
+                    "All tool calls must include agent identity."
+                )
+
+            # Apply fallback chain (deny_on_missing_agent: false)
+            # Priority: 1. GATEWAY_DEFAULT_AGENT env var, 2. "default" agent in rules
+            agent_id = self._resolve_fallback_agent(context)
+
+            # Inject the resolved agent_id back into arguments for gateway tools
+            if agent_id:
+                arguments["agent_id"] = agent_id
+                tool_call.arguments = arguments
+
+        # Store agent in context state for downstream tools
+        # This allows gateway tools to access the current agent
+        if context.fastmcp_context:
+            context.fastmcp_context.set_state("current_agent", agent_id)
+
+        # NOTE: We do NOT remove agent_id from arguments because the gateway
+        # tools (list_servers, get_server_tools, execute_tool) need it as
+        # a parameter to perform their authorization checks.
+        # If we ever add direct proxying to downstream servers in the future,
+        # we would need to remove it at that point.
+
+        # Gateway tools (list_servers, get_server_tools, execute_tool) are
+        # allowed through - they perform their own permission checks using
+        # the agent_id parameter
+        return await call_next(context)
+
+    def _resolve_fallback_agent(self, context: MiddlewareContext) -> str:
+        """Resolve fallback agent when agent_id is missing and deny_on_missing_agent: false.
+
+        Fallback priority order:
+        1. GATEWAY_DEFAULT_AGENT environment variable (read from gateway module)
+        2. Agent named "default" in gateway rules configuration
+        3. Raise helpful error if neither is configured
+
+        Args:
+            context: Middleware context to access gateway state
+
+        Returns:
+            Resolved agent_id from fallback chain
+
+        Raises:
+            ToolError: If no fallback agent is configured or if fallback agent doesn't exist in rules
+        """
+        # Import here to avoid circular dependency
+        from .gateway import get_default_agent_id
+
+        # Try to get default agent from environment variable (stored in gateway module)
+        default_agent_from_env = get_default_agent_id()
+
+        # Priority 1: GATEWAY_DEFAULT_AGENT environment variable
+        if default_agent_from_env:
+            # Validate that this agent exists in policy rules
+            if default_agent_from_env in self.policy_engine.agents:
+                return default_agent_from_env
+            else:
+                raise ToolError(
+                    f"Missing 'agent_id' parameter and fallback agent '{default_agent_from_env}' "
+                    f"is not configured in gateway rules.\n"
+                    f"Either provide 'agent_id' in your tool calls, or ask the user to configure "
+                    f"the gateway fallback settings. See gateway documentation for configuration options."
+                )
+
+        # Priority 2: Agent named "default" in gateway rules
+        if "default" in self.policy_engine.agents:
+            return "default"
+
+        # Priority 3: No fallback configured - provide helpful error
+        raise ToolError(
+            "Missing 'agent_id' parameter and no fallback agent configured.\n"
+            "Either provide 'agent_id' in your tool calls, or ask the user to configure "
+            "the gateway fallback settings. See gateway documentation for configuration options."
+        )
+
+    async def on_list_tools(self, context: MiddlewareContext, call_next):
+        """Pass through list_tools requests without filtering.
+
+        Gateway tools (list_servers, get_server_tools, execute_tool) should
+        always be visible to all agents since they perform their own
+        authorization based on agent_id passed in arguments.
+
+        This differs from a traditional MCP proxy that might filter downstream
+        tools at the middleware level. Our gateway exposes only 3 gateway tools
+        that act as an API for discovering and executing downstream tools.
+
+        Args:
+            context: Middleware context containing the list request
+            call_next: Callable to invoke next middleware/handler in chain
+
+        Returns:
+            Full list of gateway tools from downstream handler
+        """
+        # No filtering needed - gateway tools handle their own authorization
+        return await call_next(context)