kubiya-control-plane-api 0.9.15__py3-none-any.whl

control_plane_api/app/lib/event_bus/base.py

@@ -0,0 +1,136 @@
+"""Base classes and interfaces for event bus providers."""
+
+from abc import ABC, abstractmethod
+from typing import Dict, Any, Optional, Callable
+from pydantic import BaseModel, Field
+import structlog
+
+logger = structlog.get_logger(__name__)
+
+
+class EventBusConfig(BaseModel):
+    """Base configuration for event bus providers."""
+
+    enabled: bool = Field(default=True, description="Whether this provider is enabled")
+    timeout_seconds: int = Field(
+        default=30, description="Timeout for operations in seconds"
+    )
+    retry_attempts: int = Field(
+        default=3, description="Number of retry attempts on failure"
+    )
+    retry_backoff_seconds: float = Field(
+        default=0.5, description="Backoff time between retries in seconds"
+    )
+
+
+class EventBusProvider(ABC):
+    """Abstract base class for event bus providers."""
+
+    def __init__(self, config: EventBusConfig):
+        self.config = config
+        self.logger = structlog.get_logger(
+            __name__, provider=self.__class__.__name__
+        )
+
+    @abstractmethod
+    async def initialize(self) -> None:
+        """
+        Initialize provider connection/resources.
+
+        Raises:
+            Exception: If initialization fails
+        """
+        pass
+
+    @abstractmethod
+    async def publish_event(
+        self,
+        execution_id: str,
+        event_type: str,
+        data: Dict[str, Any],
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> bool:
+        """
+        Publish event to the bus.
+
+        Args:
+            execution_id: Unique execution identifier
+            event_type: Type of event (e.g., "message_chunk", "tool_started")
+            data: Event payload data
+            metadata: Optional metadata (organization_id, worker_id, etc.)
+
+        Returns:
+            True if successful, False otherwise
+        """
+        pass
+
+    @abstractmethod
+    async def subscribe(
+        self, pattern: str, callback: Callable[[Dict[str, Any]], None]
+    ) -> None:
+        """
+        Subscribe to events matching pattern.
+
+        Args:
+            pattern: Pattern to match events (provider-specific format)
+            callback: Async callback function to handle events
+        """
+        pass
+
+    @abstractmethod
+    async def health_check(self) -> Dict[str, Any]:
+        """
+        Check provider health status.
+
+        Returns:
+            Dict with health info: {"healthy": bool, "details": {...}}
+        """
+        pass
+
+    @abstractmethod
+    async def shutdown(self) -> None:
+        """Cleanup resources and close connections."""
+        pass
+
+    async def _retry_operation(
+        self, operation: Callable, operation_name: str
+    ) -> Any:
+        """
+        Retry an operation with exponential backoff.
+
+        Args:
+            operation: Async operation to retry
+            operation_name: Name for logging
+
+        Returns:
+            Result of operation
+
+        Raises:
+            Last exception if all retries fail
+        """
+        import asyncio
+
+        last_exception = None
+        for attempt in range(self.config.retry_attempts):
+            try:
+                return await operation()
+            except Exception as e:
+                last_exception = e
+                if attempt < self.config.retry_attempts - 1:
+                    backoff = self.config.retry_backoff_seconds * (2**attempt)
+                    self.logger.warning(
+                        f"{operation_name}_retry",
+                        attempt=attempt + 1,
+                        max_attempts=self.config.retry_attempts,
+                        backoff_seconds=backoff,
+                        error=str(e),
+                    )
+                    await asyncio.sleep(backoff)
+                else:
+                    self.logger.error(
+                        f"{operation_name}_failed_all_retries",
+                        attempts=self.config.retry_attempts,
+                        error=str(e),
+                    )
+
+        raise last_exception

control_plane_api/app/lib/event_bus/manager.py

@@ -0,0 +1,335 @@
+"""Event bus manager for orchestrating multiple providers."""
+
+from typing import Dict, Any, Optional
+from pydantic import BaseModel, Field
+import asyncio
+import structlog
+
+from control_plane_api.app.lib.event_bus.base import EventBusProvider
+
+logger = structlog.get_logger(__name__)
+
+
+class EventBusManagerConfig(BaseModel):
+    """Configuration for event bus manager with all providers."""
+
+    # Import provider configs (will be defined in provider modules)
+    http: Optional[Any] = Field(default=None, description="HTTP provider config")
+    websocket: Optional[Any] = Field(
+        default=None, description="WebSocket provider config"
+    )
+    redis: Optional[Any] = Field(default=None, description="Redis provider config")
+    nats: Optional[Any] = Field(default=None, description="NATS provider config")
+
+
+class EventBusManager:
+    """
+    Manages multiple event bus providers with fallback support.
+    Publishes to all enabled providers simultaneously.
+    """
+
+    def __init__(self, config: EventBusManagerConfig):
+        self.config = config
+        self.providers: Dict[str, EventBusProvider] = {}
+        self.logger = structlog.get_logger(__name__)
+        self._initialized = False
+
+    async def initialize(self) -> None:
+        """Initialize all enabled providers."""
+        if self._initialized:
+            self.logger.warning("event_bus_manager_already_initialized")
+            return
+
+        initialized_providers = []
+
+        # HTTP Provider
+        if self.config.http and self.config.http.enabled:
+            try:
+                from control_plane_api.app.lib.event_bus.providers.http_provider import (
+                    HTTPEventProvider,
+                )
+
+                self.providers["http"] = HTTPEventProvider(self.config.http)
+                await self.providers["http"].initialize()
+                initialized_providers.append("http")
+                self.logger.info("http_provider_initialized")
+            except Exception as e:
+                self.logger.error("http_provider_init_failed", error=str(e))
+                # HTTP is critical, re-raise
+                raise
+
+        # WebSocket Provider
+        if self.config.websocket and self.config.websocket.enabled:
+            try:
+                from control_plane_api.app.lib.event_bus.providers.websocket_provider import (
+                    WebSocketEventProvider,
+                )
+
+                self.providers["websocket"] = WebSocketEventProvider(
+                    self.config.websocket
+                )
+                await self.providers["websocket"].initialize()
+                initialized_providers.append("websocket")
+                self.logger.info("websocket_provider_initialized")
+            except Exception as e:
+                self.logger.warning(
+                    "websocket_provider_init_failed_continuing", error=str(e)
+                )
+                # WebSocket is optional, continue without it
+
+        # Redis Provider
+        if self.config.redis and self.config.redis.enabled:
+            try:
+                from control_plane_api.app.lib.event_bus.providers.redis_provider import (
+                    RedisEventProvider,
+                )
+
+                self.providers["redis"] = RedisEventProvider(self.config.redis)
+                await self.providers["redis"].initialize()
+                initialized_providers.append("redis")
+                self.logger.info("redis_provider_initialized")
+            except Exception as e:
+                self.logger.warning(
+                    "redis_provider_init_failed_continuing", error=str(e)
+                )
+                # Redis is optional, continue without it
+
+        # NATS Provider (Optional)
+        if self.config.nats and self.config.nats.enabled:
+            try:
+                from control_plane_api.app.lib.event_bus.providers.nats_provider import (
+                    NATSEventProvider,
+                )
+
+                self.providers["nats"] = NATSEventProvider(self.config.nats)
+                await self.providers["nats"].initialize()
+                initialized_providers.append("nats")
+                self.logger.info("nats_provider_initialized")
+            except Exception as e:
+                self.logger.warning(
+                    "nats_provider_init_failed_optional", error=str(e)
+                )
+                # NATS is optional, continue without it
+
+        if not self.providers:
+            raise RuntimeError("No event bus providers initialized")
+
+        self._initialized = True
+        self.logger.info(
+            "event_bus_manager_initialized",
+            providers=initialized_providers,
+            count=len(self.providers),
+        )
+
+    async def publish_event(
+        self,
+        execution_id: str,
+        event_type: str,
+        data: Dict[str, Any],
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, bool]:
+        """
+        Publish event to all enabled providers in parallel.
+
+        Args:
+            execution_id: Unique execution identifier
+            event_type: Type of event
+            data: Event payload
+            metadata: Optional metadata
+
+        Returns:
+            Dict mapping provider name to success status
+        """
+        if not self._initialized:
+            self.logger.error("event_bus_manager_not_initialized")
+            return {}
+
+        # Log start of publishing
+        self.logger.info(
+            "event_bus_publishing_start",
+            execution_id=execution_id,
+            event_type=event_type,
+            providers=list(self.providers.keys()),
+            provider_count=len(self.providers)
+        )
+
+        # PERFORMANCE OPTIMIZATION: Prioritize WebSocket for instant delivery
+        # Other providers (Redis, HTTP) run in background without blocking
+        websocket_task = None
+        websocket_name = None
+        background_tasks = []
+        background_names = []
+
+        for name, provider in self.providers.items():
+            task = asyncio.create_task(
+                provider.publish_event(execution_id, event_type, data, metadata)
+            )
+            if name == "websocket":
+                websocket_task = task
+                websocket_name = name
+            else:
+                background_tasks.append(task)
+                background_names.append(name)
+
+        results = {}
+
+        # Wait only for WebSocket (real-time delivery to client)
+        if websocket_task:
+            try:
+                ws_result = await websocket_task
+                results[websocket_name] = ws_result
+                if ws_result:
+                    self.logger.debug(
+                        "websocket_publish_success",
+                        execution_id=execution_id,
+                        event_type=event_type
+                    )
+                else:
+                    self.logger.warning(
+                        "websocket_publish_failed",
+                        execution_id=execution_id,
+                        event_type=event_type
+                    )
+            except Exception as e:
+                self.logger.error(
+                    "websocket_publish_exception",
+                    provider=websocket_name,
+                    error=str(e),
+                    execution_id=execution_id,
+                    event_type=event_type,
+                )
+                results[websocket_name] = False
+
+        # Background providers - don't await, let them complete asynchronously
+        # Use asyncio.gather with return_exceptions to prevent unhandled exception warnings
+        if background_tasks:
+            async def process_background_providers():
+                try:
+                    bg_results = await asyncio.gather(*background_tasks, return_exceptions=True)
+                    for name, result in zip(background_names, bg_results):
+                        if isinstance(result, Exception):
+                            self.logger.warning(
+                                "background_provider_publish_exception",
+                                provider=name,
+                                error=str(result),
+                                execution_id=execution_id,
+                                event_type=event_type,
+                            )
+                        elif not result:
+                            self.logger.warning(
+                                "background_provider_publish_failed",
+                                provider=name,
+                                execution_id=execution_id,
+                                event_type=event_type
+                            )
+                except Exception as e:
+                    self.logger.error(
+                        "background_providers_gather_failed",
+                        error=str(e),
+                        execution_id=execution_id,
+                    )
+
+            # Fire and forget - schedule but don't await
+            asyncio.create_task(process_background_providers())
+
+            # Mark background providers as pending (we won't wait for results)
+            for name in background_names:
+                results[name] = True  # Optimistically assume success
+
+        # Log overall status
+        success_count = sum(1 for success in results.values() if success)
+        if success_count == 0:
+            self.logger.error(
+                "event_bus_all_providers_failed",
+                execution_id=execution_id,
+                event_type=event_type,
+                results=results,
+            )
+        elif success_count < len(results):
+            self.logger.warning(
+                "event_bus_partial_success",
+                execution_id=execution_id,
+                event_type=event_type,
+                success_count=success_count,
+                total_count=len(results),
+                results=results,
+            )
+        else:
+            self.logger.info(
+                "event_bus_all_providers_success",
+                execution_id=execution_id,
+                event_type=event_type,
+                provider_count=len(results),
+                providers=list(results.keys())
+            )
+
+        return results
+
+    async def health_check(self) -> Dict[str, Any]:
+        """
+        Check health of all providers.
+
+        Returns:
+            Dict with health status for each provider
+        """
+        health = {}
+
+        for name, provider in self.providers.items():
+            try:
+                provider_health = await provider.health_check()
+                health[name] = provider_health
+            except Exception as e:
+                self.logger.error(
+                    "provider_health_check_failed", provider=name, error=str(e)
+                )
+                health[name] = {"healthy": False, "error": str(e)}
+
+        # Overall status
+        healthy_count = sum(
+            1 for h in health.values() if h.get("healthy", False)
+        )
+        health["_overall"] = {
+            "healthy": healthy_count > 0,  # At least one provider healthy
+            "total_providers": len(self.providers),
+            "healthy_providers": healthy_count,
+        }
+
+        return health
+
+    async def shutdown(self) -> None:
+        """Shutdown all providers gracefully."""
+        if not self._initialized:
+            return
+
+        self.logger.info("shutting_down_event_bus_manager")
+
+        # Shutdown all providers in parallel
+        tasks = []
+        for name, provider in self.providers.items():
+            tasks.append(self._shutdown_provider(name, provider))
+
+        await asyncio.gather(*tasks, return_exceptions=True)
+
+        self.providers.clear()
+        self._initialized = False
+        self.logger.info("event_bus_manager_shutdown_complete")
+
+    async def _shutdown_provider(
+        self, name: str, provider: EventBusProvider
+    ) -> None:
+        """Shutdown a single provider with error handling."""
+        try:
+            await provider.shutdown()
+            self.logger.info("provider_shutdown", provider=name)
+        except Exception as e:
+            self.logger.error(
+                "provider_shutdown_failed", provider=name, error=str(e)
+            )
+
+    def is_initialized(self) -> bool:
+        """Check if manager is initialized."""
+        return self._initialized
+
+    def get_provider_names(self) -> list[str]:
+        """Get list of initialized provider names."""
+        return list(self.providers.keys())

control_plane_api/app/lib/event_bus/providers/__init__.py

@@ -0,0 +1,6 @@
+"""Event bus provider implementations."""
+
+# Providers will be imported lazily by the manager to avoid circular dependencies
+# and to handle optional dependencies gracefully
+
+__all__ = []

control_plane_api/app/lib/event_bus/providers/http_provider.py

@@ -0,0 +1,166 @@
+"""HTTP-based event bus provider."""
+
+from typing import Dict, Any, Optional, Callable
+from datetime import datetime
+import httpx
+import structlog
+from pydantic import Field
+
+from control_plane_api.app.lib.event_bus.base import EventBusProvider, EventBusConfig
+
+logger = structlog.get_logger(__name__)
+
+
+class HTTPConfig(EventBusConfig):
+    """Configuration for HTTP event provider."""
+
+    base_url: str = Field(..., description="Base URL for control plane API")
+    batching_enabled: bool = Field(
+        default=True, description="Enable event batching"
+    )
+    batch_window_ms: int = Field(
+        default=100, description="Batch window in milliseconds"
+    )
+    batch_max_size_bytes: int = Field(
+        default=1000, description="Maximum batch size in bytes"
+    )
+    api_key: Optional[str] = Field(
+        default=None, description="API key for authentication"
+    )
+
+
+class HTTPEventProvider(EventBusProvider):
+    """HTTP-based event publishing (default provider)."""
+
+    def __init__(self, config: HTTPConfig):
+        super().__init__(config)
+        self.config: HTTPConfig = config
+        self.base_url = config.base_url.rstrip("/")
+        self.client: Optional[httpx.AsyncClient] = None
+
+    async def initialize(self) -> None:
+        """Initialize HTTP client."""
+        self.client = httpx.AsyncClient(
+            timeout=self.config.timeout_seconds,
+            follow_redirects=False,  # Fail fast on redirects
+        )
+        self.logger.info(
+            "http_provider_initialized",
+            base_url=self.base_url,
+            batching_enabled=self.config.batching_enabled,
+        )
+
+    async def publish_event(
+        self,
+        execution_id: str,
+        event_type: str,
+        data: Dict[str, Any],
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> bool:
+        """
+        Publish event via HTTP POST.
+
+        Args:
+            execution_id: Execution ID
+            event_type: Event type
+            data: Event data
+            metadata: Optional metadata
+
+        Returns:
+            True if successful (202 Accepted)
+        """
+        if not self.client:
+            self.logger.error("http_client_not_initialized")
+            return False
+
+        url = f"{self.base_url}/api/v1/executions/{execution_id}/events"
+
+        payload = {
+            "event_type": event_type,
+            "data": data,
+            "timestamp": datetime.utcnow().isoformat(),
+        }
+
+        headers = {}
+        if self.config.api_key:
+            headers["Authorization"] = f"Bearer {self.config.api_key}"
+
+        try:
+            response = await self.client.post(url, json=payload, headers=headers)
+
+            if response.status_code == 202:
+                self.logger.debug(
+                    "http_event_published",
+                    execution_id=execution_id,
+                    event_type=event_type,
+                )
+                return True
+            else:
+                self.logger.warning(
+                    "http_event_publish_failed",
+                    execution_id=execution_id,
+                    event_type=event_type,
+                    status_code=response.status_code,
+                    response=response.text[:200],
+                )
+                return False
+
+        except httpx.TimeoutException:
+            self.logger.error(
+                "http_event_publish_timeout",
+                execution_id=execution_id,
+                event_type=event_type,
+                timeout=self.config.timeout_seconds,
+            )
+            return False
+        except Exception as e:
+            self.logger.error(
+                "http_event_publish_exception",
+                execution_id=execution_id,
+                event_type=event_type,
+                error=str(e),
+            )
+            return False
+
+    async def subscribe(
+        self, pattern: str, callback: Callable[[Dict[str, Any]], None]
+    ) -> None:
+        """HTTP provider does not support subscriptions."""
+        self.logger.warning(
+            "http_provider_subscribe_not_supported", pattern=pattern
+        )
+        raise NotImplementedError(
+            "HTTP provider does not support subscriptions"
+        )
+
+    async def health_check(self) -> Dict[str, Any]:
+        """Check HTTP client health."""
+        if not self.client:
+            return {
+                "healthy": False,
+                "error": "client_not_initialized",
+            }
+
+        try:
+            # Try to make a simple request to health endpoint
+            health_url = f"{self.base_url}/api/health"
+            response = await self.client.get(health_url, timeout=5.0)
+
+            return {
+                "healthy": response.status_code == 200,
+                "base_url": self.base_url,
+                "status_code": response.status_code,
+            }
+        except Exception as e:
+            return {
+                "healthy": False,
+                "base_url": self.base_url,
+                "error": str(e),
+            }
+
+    async def shutdown(self) -> None:
+        """Close HTTP client."""
+        if self.client:
+            await self.client.aclose()
+            self.client = None
+            self.logger.info("http_provider_shutdown")