kailash 0.6.2__py3-none-any.whl → 0.6.4__py3-none-any.whl

kailash/core/resilience/circuit_breaker.py

@@ -23,6 +23,7 @@ Example:
 
 import asyncio
 import logging
+import random
 import time
 from collections import deque
 from dataclasses import dataclass, field
@@ -61,6 +62,15 @@ class CircuitBreakerConfig:
     window_size: int = 100  # Rolling window for error rate
     excluded_exceptions: List[type] = field(default_factory=list)  # Don't count these
 
+    # Enhanced configurable thresholds
+    min_calls_before_evaluation: int = 10  # Min calls before evaluating error rate
+    slow_call_threshold: float = 5.0  # Seconds to consider a call slow
+    slow_call_rate_threshold: float = 0.8  # Rate of slow calls to trigger open
+    max_wait_duration_in_half_open: int = 60  # Max wait in half-open state
+    exponential_backoff_multiplier: float = 2.0  # Backoff multiplier for recovery
+    jitter_enabled: bool = True  # Add jitter to recovery timeout
+    max_jitter_percentage: float = 0.1  # Maximum jitter as percentage of timeout
+
 
 @dataclass
 class CircuitBreakerMetrics:
@@ -70,36 +80,59 @@ class CircuitBreakerMetrics:
     successful_calls: int = 0
     failed_calls: int = 0
     rejected_calls: int = 0
+    slow_calls: int = 0  # New: Track slow calls
     state_transitions: List[Dict[str, Any]] = field(default_factory=list)
     last_failure_time: Optional[float] = None
     consecutive_failures: int = 0
     consecutive_successes: int = 0
+    avg_call_duration: float = 0.0  # New: Average call duration
+    total_call_duration: float = 0.0  # New: Total duration for average calculation
 
-    def record_success(self):
+    def record_success(self, duration: float = 0.0):
         """Record successful call."""
         self.total_calls += 1
         self.successful_calls += 1
         self.consecutive_successes += 1
         self.consecutive_failures = 0
+        self._update_duration(duration)
 
-    def record_failure(self):
+    def record_failure(self, duration: float = 0.0):
         """Record failed call."""
         self.total_calls += 1
         self.failed_calls += 1
         self.consecutive_failures += 1
         self.consecutive_successes = 0
         self.last_failure_time = time.time()
+        self._update_duration(duration)
 
     def record_rejection(self):
         """Record rejected call (circuit open)."""
         self.rejected_calls += 1
 
+    def record_slow_call(self):
+        """Record slow call."""
+        self.slow_calls += 1
+
+    def _update_duration(self, duration: float):
+        """Update duration metrics."""
+        if duration > 0:
+            self.total_call_duration += duration
+            # Update rolling average
+            if self.total_calls > 0:
+                self.avg_call_duration = self.total_call_duration / self.total_calls
+
     def get_error_rate(self) -> float:
         """Calculate current error rate."""
         if self.total_calls == 0:
             return 0.0
         return self.failed_calls / self.total_calls
 
+    def get_slow_call_rate(self) -> float:
+        """Calculate current slow call rate."""
+        if self.total_calls == 0:
+            return 0.0
+        return self.slow_calls / self.total_calls
+
 
 class ConnectionCircuitBreaker(Generic[T]):
     """Circuit breaker for database connections and operations.
@@ -159,14 +192,21 @@ class ConnectionCircuitBreaker(Generic[T]):
         start_time = time.time()
         try:
             result = await func(*args, **kwargs)
-            await self._record_success()
+            execution_time = time.time() - start_time
+
+            # Check if this was a slow call
+            is_slow = execution_time > self.config.slow_call_threshold
+
+            await self._record_success(execution_time, is_slow)
             return result
         except Exception as e:
+            execution_time = time.time() - start_time
+
             # Check if this exception should be counted
             if not any(
                 isinstance(e, exc_type) for exc_type in self.config.excluded_exceptions
             ):
-                await self._record_failure(e)
+                await self._record_failure(e, execution_time)
             raise
 
     async def _check_state_transition(self):
@@ -191,6 +231,10 @@ class ConnectionCircuitBreaker(Generic[T]):
 
     def _should_open(self) -> bool:
         """Determine if circuit should open based on failures."""
+        # Only evaluate if we have minimum number of calls
+        if self.metrics.total_calls < self.config.min_calls_before_evaluation:
+            return False
+
         # Check consecutive failures
         if self.metrics.consecutive_failures >= self.config.failure_threshold:
             return True
@@ -202,22 +246,32 @@ class ConnectionCircuitBreaker(Generic[T]):
             if error_rate >= self.config.error_rate_threshold:
                 return True
 
+        # Check slow call rate
+        slow_call_rate = self.metrics.get_slow_call_rate()
+        if slow_call_rate >= self.config.slow_call_rate_threshold:
+            return True
+
         return False
 
-    async def _record_success(self):
+    async def _record_success(self, duration: float = 0.0, is_slow: bool = False):
         """Record successful execution."""
         async with self._lock:
-            self.metrics.record_success()
+            self.metrics.record_success(duration)
+            if is_slow:
+                self.metrics.record_slow_call()
             self._rolling_window.append(True)
 
             if self.state == CircuitState.HALF_OPEN:
                 if self.metrics.consecutive_successes >= self.config.success_threshold:
                     await self._transition_to(CircuitState.CLOSED)
 
-    async def _record_failure(self, error: Exception):
+    async def _record_failure(self, error: Exception, duration: float = 0.0):
         """Record failed execution."""
         async with self._lock:
-            self.metrics.record_failure()
+            self.metrics.record_failure(duration)
+            # Consider slow failures as additional burden
+            if duration > self.config.slow_call_threshold:
+                self.metrics.record_slow_call()
             self._rolling_window.append(False)
 
             if self.state == CircuitState.HALF_OPEN:
@@ -280,11 +334,31 @@ class ConnectionCircuitBreaker(Generic[T]):
         return "Unknown reason"
 
     def _time_until_recovery(self) -> float:
-        """Calculate seconds until recovery attempt."""
+        """Calculate seconds until recovery attempt with jitter and backoff."""
         if self.state != CircuitState.OPEN:
             return 0.0
+
         elapsed = time.time() - self._last_state_change
-        remaining = self.config.recovery_timeout - elapsed
+
+        # Apply exponential backoff based on number of state transitions to OPEN
+        open_transitions = sum(
+            1
+            for t in self.metrics.state_transitions
+            if t.get("to") == CircuitState.OPEN.value
+        )
+        backoff_multiplier = self.config.exponential_backoff_multiplier ** max(
+            0, open_transitions - 1
+        )
+
+        base_timeout = self.config.recovery_timeout * backoff_multiplier
+
+        # Add jitter if enabled
+        if self.config.jitter_enabled:
+            jitter_range = base_timeout * self.config.max_jitter_percentage
+            jitter = random.uniform(-jitter_range, jitter_range)
+            base_timeout += jitter
+
+        remaining = base_timeout - elapsed
         return max(0.0, remaining)
 
     async def force_open(self, reason: str = "Manual override"):
@@ -331,7 +405,10 @@ class ConnectionCircuitBreaker(Generic[T]):
                 "successful_calls": self.metrics.successful_calls,
                 "failed_calls": self.metrics.failed_calls,
                 "rejected_calls": self.metrics.rejected_calls,
+                "slow_calls": self.metrics.slow_calls,
                 "error_rate": self.metrics.get_error_rate(),
+                "slow_call_rate": self.metrics.get_slow_call_rate(),
+                "avg_call_duration": self.metrics.avg_call_duration,
                 "consecutive_failures": self.metrics.consecutive_failures,
                 "consecutive_successes": self.metrics.consecutive_successes,
             },
@@ -340,6 +417,11 @@ class ConnectionCircuitBreaker(Generic[T]):
                 "success_threshold": self.config.success_threshold,
                 "recovery_timeout": self.config.recovery_timeout,
                 "error_rate_threshold": self.config.error_rate_threshold,
+                "slow_call_threshold": self.config.slow_call_threshold,
+                "slow_call_rate_threshold": self.config.slow_call_rate_threshold,
+                "min_calls_before_evaluation": self.config.min_calls_before_evaluation,
+                "exponential_backoff_multiplier": self.config.exponential_backoff_multiplier,
+                "jitter_enabled": self.config.jitter_enabled,
             },
             "time_until_recovery": (
                 self._time_until_recovery() if self.state == CircuitState.OPEN else None

kailash/edge/discovery.py

@@ -199,6 +199,92 @@ class EdgeDiscovery:
         self._last_health_check[location.location_id] = datetime.now(UTC)
         logger.info(f"Added edge location: {location.name}")
 
+    async def register_edge(self, edge_config: Dict[str, Any]):
+        """Register an edge location from configuration dictionary.
+
+        Args:
+            edge_config: Dictionary containing edge location configuration
+        """
+        from .location import (
+            ComplianceZone,
+            EdgeCapabilities,
+            EdgeLocation,
+            EdgeRegion,
+            GeographicCoordinates,
+        )
+
+        # Extract basic info
+        location_id = edge_config["id"]
+        region_str = edge_config.get("region", "us-east")
+
+        # Map region string to enum
+        region_map = {
+            "us-east-1": EdgeRegion.US_EAST,
+            "us-west-1": EdgeRegion.US_WEST,
+            "eu-west-1": EdgeRegion.EU_WEST,
+            "eu-central-1": EdgeRegion.EU_CENTRAL,
+            "asia-southeast-1": EdgeRegion.ASIA_SOUTHEAST,
+        }
+        region = region_map.get(region_str, EdgeRegion.US_EAST)
+
+        # Default coordinates based on region
+        coord_map = {
+            EdgeRegion.US_EAST: GeographicCoordinates(39.0458, -76.6413),  # Virginia
+            EdgeRegion.US_WEST: GeographicCoordinates(37.7749, -122.4194),  # California
+            EdgeRegion.EU_WEST: GeographicCoordinates(53.3498, -6.2603),  # Ireland
+            EdgeRegion.EU_CENTRAL: GeographicCoordinates(50.1109, 8.6821),  # Frankfurt
+            EdgeRegion.ASIA_SOUTHEAST: GeographicCoordinates(
+                1.3521, 103.8198
+            ),  # Singapore
+        }
+        coordinates = coord_map.get(region, GeographicCoordinates(39.0458, -76.6413))
+
+        # Create capabilities
+        capabilities = EdgeCapabilities(
+            cpu_cores=edge_config.get("capacity", 1000) // 100,  # Rough mapping
+            memory_gb=edge_config.get("capacity", 1000) // 50,
+            storage_gb=edge_config.get("capacity", 1000) * 2,
+            bandwidth_gbps=10.0,
+            database_support=["postgresql", "redis"],
+            ai_models_available=["llama", "claude"],
+        )
+
+        # Create edge location
+        location = EdgeLocation(
+            location_id=location_id,
+            name=f"Edge {region_str.title()}",
+            region=region,
+            coordinates=coordinates,
+            capabilities=capabilities,
+            endpoint_url=edge_config.get(
+                "endpoint", f"http://{location_id}.edge.local:8080"
+            ),
+        )
+
+        # Set health status
+        from .location import EdgeStatus
+
+        if edge_config.get("healthy", True):
+            location.status = EdgeStatus.ACTIVE
+            self._health_results[location_id] = HealthCheckResult.HEALTHY
+        else:
+            location.status = EdgeStatus.OFFLINE
+            self._health_results[location_id] = HealthCheckResult.UNHEALTHY
+
+        # Update metrics
+        location.metrics.latency_p50_ms = edge_config.get("latency_ms", 10)
+        location.metrics.cpu_utilization = edge_config.get(
+            "current_load", 0
+        ) / edge_config.get("capacity", 1000)
+
+        # Add to locations
+        self.locations[location_id] = location
+        self._last_health_check[location_id] = datetime.now(UTC)
+
+        logger.info(f"Registered edge location: {location_id} in {region_str}")
+
+        return location
+
     def remove_location(self, location_id: str):
         """Remove an edge location from the discovery pool."""
         if location_id in self.locations:

kailash/mcp_server/__init__.py

@@ -0,0 +1,334 @@
+"""Enhanced Model Context Protocol (MCP) Service Layer with Service Discovery.
+
+This module provides production-ready MCP client and server functionality with
+comprehensive service discovery, authentication, monitoring, and resilience
+features. Built on top of the official Anthropic MCP Python SDK.
+
+Enhanced Features:
+    - Service Discovery: Automatic server registration and discovery
+    - Authentication: Multiple auth providers (API Key, JWT, OAuth)
+    - Load Balancing: Intelligent server selection and failover
+    - Health Monitoring: Automatic health checking and status tracking
+    - Circuit Breaker: Failure detection and recovery patterns
+    - Metrics Collection: Comprehensive performance monitoring
+    - Network Discovery: UDP broadcast/multicast server discovery
+    - Error Handling: Structured error codes and retry strategies
+
+Design Philosophy:
+    Provides production-ready distributed systems infrastructure for MCP
+    while maintaining compatibility with the official SDK. Enhances the
+    basic protocol implementation with enterprise-grade features.
+
+Key Components:
+    - MCPClient: Enhanced client with auth, retry, and multi-transport support
+    - MCPServer: Production-ready server with all enterprise features
+    - ServiceRegistry: Central registry for server discovery and management
+    - ServiceMesh: Intelligent routing and load balancing
+    - ServerRegistrar: Automatic server registration and lifecycle management
+
+Service Discovery Features:
+    - File-based registry with JSON storage
+    - Network discovery via UDP broadcast/multicast
+    - Health checking and automatic status updates
+    - Capability-based server filtering
+    - Load balancing with priority scoring
+    - Automatic failover and circuit breaker patterns
+
+Examples:
+    Enhanced MCP client with discovery:
+
+    >>> from kailash.mcp_server import get_mcp_client, discover_mcp_servers
+    >>> # Discover servers with specific capability
+    >>> servers = await discover_mcp_servers(capability="search")
+    >>> # Get best client for capability
+    >>> client = await get_mcp_client("search")
+    >>> result = await client.call_tool(server_config, "search", {"query": "AI"})
+
+    Production MCP server with auto-discovery:
+
+    >>> from kailash.mcp_server import MCPServer, enable_auto_discovery
+    >>> from kailash.mcp_server.auth import APIKeyAuth
+    >>>
+    >>> # Create server with authentication
+    >>> auth = APIKeyAuth({"user1": "secret-key"})
+    >>> server = MCPServer(
+    ...     "my-tools",
+    ...     auth_provider=auth,
+    ...     enable_metrics=True,
+    ...     circuit_breaker_config={"failure_threshold": 5}
+    ... )
+    >>>
+    >>> @server.tool(required_permission="tools.calculate")
+    ... def calculate(a: int, b: int) -> int:
+    ...     return a + b
+    >>>
+    >>> # Enable auto-discovery and start
+    >>> registrar = enable_auto_discovery(server, enable_network_discovery=True)
+    >>> registrar.start_with_registration()
+
+    Service mesh with failover:
+
+    >>> from kailash.mcp_server import ServiceRegistry, ServiceMesh
+    >>> registry = ServiceRegistry()
+    >>> mesh = ServiceMesh(registry)
+    >>>
+    >>> # Call with automatic failover
+    >>> result = await mesh.call_with_failover(
+    ...     "search", "web_search", {"query": "Python"}, max_retries=3
+    ... )
+"""
+
+# Advanced Features
+from .advanced_features import (
+    BinaryResourceHandler,
+    CancellationContext,
+    ChangeType,
+    Content,
+    ContentType,
+    ElicitationSystem,
+    MultiModalContent,
+    ProgressReporter,
+    ResourceChange,
+)
+from .advanced_features import ResourceTemplate as AdvancedResourceTemplate
+from .advanced_features import (
+    SchemaValidator,
+    StreamingHandler,
+    StructuredTool,
+    ToolAnnotation,
+    create_cancellation_context,
+    create_progress_reporter,
+    structured_tool,
+)
+
+# Authentication framework
+from .auth import (
+    APIKeyAuth,
+    AuthManager,
+    AuthProvider,
+    BasicAuth,
+    BearerTokenAuth,
+    JWTAuth,
+    PermissionManager,
+    RateLimiter,
+)
+from .client import MCPClient
+
+# Service Discovery System
+from .discovery import (
+    DiscoveryBackend,
+    FileBasedDiscovery,
+    HealthChecker,
+    LoadBalancer,
+    NetworkDiscovery,
+    ServerInfo,
+    ServiceMesh,
+    ServiceRegistry,
+    create_default_registry,
+    discover_mcp_servers,
+    get_mcp_client,
+)
+
+# Enhanced error handling
+from .errors import (
+    AuthenticationError,
+    AuthorizationError,
+    CircuitBreakerRetry,
+    ErrorAggregator,
+    ExponentialBackoffRetry,
+    MCPError,
+    MCPErrorCode,
+    RateLimitError,
+    ResourceError,
+    RetryableOperation,
+    RetryStrategy,
+    ServiceDiscoveryError,
+    ToolError,
+    TransportError,
+    ValidationError,
+)
+
+# OAuth 2.1 Authentication
+from .oauth import (
+    AccessToken,
+    AuthorizationCode,
+    AuthorizationServer,
+    ClientStore,
+    ClientType,
+    GrantType,
+    InMemoryClientStore,
+    InMemoryTokenStore,
+    JWTManager,
+    OAuth2Client,
+    OAuthClient,
+    RefreshToken,
+    ResourceServer,
+    TokenStore,
+    TokenType,
+)
+
+# Complete Protocol Implementation
+from .protocol import (
+    CancellationManager,
+    CancelledNotification,
+    CompletionManager,
+    CompletionRequest,
+    CompletionResult,
+    MessageType,
+    MetaData,
+    ProgressManager,
+    ProgressNotification,
+    ProgressToken,
+    ProtocolManager,
+    ResourceTemplate,
+    RootsManager,
+    SamplingManager,
+    SamplingRequest,
+    ToolResult,
+    cancel_request,
+    complete_progress,
+    get_protocol_manager,
+    is_cancelled,
+    start_progress,
+    update_progress,
+)
+
+# Registry Integration
+from .registry_integration import (
+    NetworkAnnouncer,
+    ServerRegistrar,
+    enable_auto_discovery,
+    register_server_manually,
+)
+
+# Enhanced server with production features
+from .server import MCPServer, MCPServerBase, SimpleMCPServer
+
+# Enhanced Transport Layer
+from .transports import (
+    BaseTransport,
+    EnhancedStdioTransport,
+    SSETransport,
+    StreamableHTTPTransport,
+    TransportManager,
+    TransportSecurity,
+    WebSocketTransport,
+    get_transport_manager,
+)
+
+__all__ = [
+    # Core MCP Components
+    "MCPClient",
+    "MCPServer",
+    "MCPServerBase",
+    # Prototyping server
+    "SimpleMCPServer",
+    # Service Discovery System
+    "ServiceRegistry",
+    "ServerInfo",
+    "DiscoveryBackend",
+    "FileBasedDiscovery",
+    "NetworkDiscovery",
+    "HealthChecker",
+    "LoadBalancer",
+    "ServiceMesh",
+    "create_default_registry",
+    "discover_mcp_servers",
+    "get_mcp_client",
+    # Registry Integration
+    "ServerRegistrar",
+    "NetworkAnnouncer",
+    "enable_auto_discovery",
+    "register_server_manually",
+    # Authentication
+    "AuthProvider",
+    "APIKeyAuth",
+    "BearerTokenAuth",
+    "JWTAuth",
+    "BasicAuth",
+    "AuthManager",
+    "PermissionManager",
+    "RateLimiter",
+    # Enhanced Error Handling
+    "MCPError",
+    "MCPErrorCode",
+    "AuthenticationError",
+    "AuthorizationError",
+    "RateLimitError",
+    "ToolError",
+    "ResourceError",
+    "TransportError",
+    "ServiceDiscoveryError",
+    "ValidationError",
+    "RetryStrategy",
+    "RetryableOperation",
+    "ExponentialBackoffRetry",
+    "CircuitBreakerRetry",
+    "ErrorAggregator",
+    # Complete Protocol Implementation
+    "MessageType",
+    "ProgressToken",
+    "MetaData",
+    "ProgressNotification",
+    "CancelledNotification",
+    "CompletionRequest",
+    "CompletionResult",
+    "SamplingRequest",
+    "ResourceTemplate",
+    "ToolResult",
+    "ProgressManager",
+    "CancellationManager",
+    "CompletionManager",
+    "SamplingManager",
+    "RootsManager",
+    "ProtocolManager",
+    "get_protocol_manager",
+    "start_progress",
+    "update_progress",
+    "complete_progress",
+    "is_cancelled",
+    "cancel_request",
+    # Enhanced Transport Layer
+    "BaseTransport",
+    "EnhancedStdioTransport",
+    "SSETransport",
+    "StreamableHTTPTransport",
+    "WebSocketTransport",
+    "TransportSecurity",
+    "TransportManager",
+    "get_transport_manager",
+    # OAuth 2.1 Authentication
+    "GrantType",
+    "TokenType",
+    "ClientType",
+    "OAuthClient",
+    "AccessToken",
+    "RefreshToken",
+    "AuthorizationCode",
+    "ClientStore",
+    "InMemoryClientStore",
+    "TokenStore",
+    "InMemoryTokenStore",
+    "JWTManager",
+    "AuthorizationServer",
+    "ResourceServer",
+    "OAuth2Client",
+    # Advanced Features
+    "ContentType",
+    "ChangeType",
+    "Content",
+    "ResourceChange",
+    "ToolAnnotation",
+    "MultiModalContent",
+    "SchemaValidator",
+    "StructuredTool",
+    "AdvancedResourceTemplate",
+    "BinaryResourceHandler",
+    "StreamingHandler",
+    "ElicitationSystem",
+    "ProgressReporter",
+    "CancellationContext",
+    "structured_tool",
+    "create_progress_reporter",
+    "create_cancellation_context",
+]