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

control_plane_api/worker/utils/environment.py

@@ -0,0 +1,65 @@
+"""
+Environment detection utilities for worker deployment context.
+
+Provides functions to detect serverless environments and determine
+whether WebSocket connections should be used.
+
+Serverless environments detected:
+- AWS Lambda
+- Vercel Functions
+- Google Cloud Functions
+- Azure Functions
+"""
+
+from typing import Literal
+import os
+
+
+def detect_environment() -> Literal["standard", "serverless"]:
+    """
+    Detect if running in a serverless environment.
+
+    Returns:
+        "serverless" if running in a serverless environment (Lambda, Vercel, GCF, Azure Functions)
+        "standard" for traditional server environments
+    """
+    # AWS Lambda
+    if os.environ.get("AWS_LAMBDA_FUNCTION_NAME"):
+        return "serverless"
+
+    # Vercel Functions
+    if os.environ.get("VERCEL"):
+        return "serverless"
+
+    # Google Cloud Functions
+    if os.environ.get("FUNCTION_TARGET") or os.environ.get("K_SERVICE"):
+        return "serverless"
+
+    # Azure Functions
+    if os.environ.get("AZURE_FUNCTIONS_ENVIRONMENT"):
+        return "serverless"
+
+    return "standard"
+
+
+def should_use_websocket() -> bool:
+    """
+    Determine if WebSocket should be used based on environment.
+
+    WebSocket is disabled in serverless environments due to limitations:
+    - Short-lived execution contexts
+    - No persistent connections
+    - Function timeout constraints
+
+    Returns:
+        True if WebSocket should be used, False otherwise
+    """
+    # Never use WebSocket in serverless
+    if detect_environment() == "serverless":
+        return False
+
+    # Check explicit disable
+    if os.environ.get("WEBSOCKET_ENABLED", "true").lower() == "false":
+        return False
+
+    return True

control_plane_api/worker/utils/error_publisher.py

@@ -0,0 +1,260 @@
+"""
+Error event publisher for standardized error handling.
+
+This module provides utilities for publishing error events to the Control Plane
+with consistent structure, user-friendly messages, and recovery suggestions.
+"""
+
+import structlog
+import traceback
+from typing import Optional, List, Dict, Any
+from datetime import datetime, timezone
+
+from control_plane_api.worker.models.error_events import (
+    ErrorEvent, ErrorSeverity, ErrorCategory, ErrorContext,
+    ErrorDetails, ErrorRecovery
+)
+from control_plane_api.worker.control_plane_client import ControlPlaneClient
+
+logger = structlog.get_logger()
+
+
+class ErrorEventPublisher:
+    """Helper class for publishing standardized error events."""
+
+    def __init__(self, control_plane: ControlPlaneClient):
+        self.control_plane = control_plane
+
+    async def publish_error(
+        self,
+        execution_id: str,
+        exception: Exception,
+        severity: ErrorSeverity,
+        category: ErrorCategory,
+        stage: str,
+        component: str,
+        operation: Optional[str] = None,
+        recovery_actions: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        include_stack_trace: bool = True,
+    ) -> bool:
+        """
+        Publish a standardized error event to Control Plane.
+
+        Args:
+            execution_id: Execution ID
+            exception: The exception that occurred
+            severity: Error severity level
+            category: Error category
+            stage: Execution stage where error occurred
+            component: Component that generated error
+            operation: Specific operation that failed
+            recovery_actions: List of user-actionable recovery steps
+            metadata: Additional metadata
+            include_stack_trace: Whether to include stack trace
+
+        Returns:
+            True if published successfully
+        """
+        try:
+            # Extract stack trace
+            stack_trace = None
+            code_location = None
+            if include_stack_trace:
+                tb_lines = traceback.format_exception(
+                    type(exception), exception, exception.__traceback__
+                )
+                stack_trace = ''.join(tb_lines)[-2000:]  # Last 2000 chars
+
+                # Extract code location from traceback
+                tb = exception.__traceback__
+                if tb:
+                    frame = traceback.extract_tb(tb)[-1]
+                    code_location = f"{frame.filename}:{frame.lineno}"
+
+            # Create user-friendly error message
+            user_message = self._create_user_friendly_message(
+                exception, category, component
+            )
+
+            # Determine if retryable
+            is_retryable = self._is_retryable_error(exception, category)
+
+            # Get default recovery actions if none provided
+            if not recovery_actions:
+                recovery_actions = self._get_default_recovery_actions(
+                    category, is_retryable
+                )
+
+            # Get documentation URL
+            doc_url = self._get_documentation_url(category)
+
+            # Build error event
+            error_event = ErrorEvent(
+                severity=severity,
+                category=category,
+                context=ErrorContext(
+                    execution_id=execution_id,
+                    timestamp=datetime.now(timezone.utc).isoformat(),
+                    stage=stage,
+                    component=component,
+                    operation=operation,
+                ),
+                details=ErrorDetails(
+                    error_type=type(exception).__name__,
+                    error_message=user_message,
+                    technical_message=str(exception),
+                    stack_trace=stack_trace,
+                    code_location=code_location,
+                ),
+                recovery=ErrorRecovery(
+                    is_retryable=is_retryable,
+                    retry_suggested=is_retryable and severity != ErrorSeverity.CRITICAL,
+                    recovery_actions=recovery_actions,
+                    documentation_url=doc_url,
+                ),
+                metadata=metadata,
+            )
+
+            # Publish error event
+            success = await self.control_plane.publish_event_async(
+                execution_id=execution_id,
+                event_type="error",
+                data=error_event.to_event_data(),
+            )
+
+            if success:
+                logger.info(
+                    "error_event_published",
+                    execution_id=execution_id[:8] if len(execution_id) >= 8 else execution_id,
+                    severity=severity.value,
+                    category=category.value,
+                    error_type=type(exception).__name__,
+                )
+            else:
+                logger.warning(
+                    "error_event_publish_failed",
+                    execution_id=execution_id[:8] if len(execution_id) >= 8 else execution_id,
+                    error=str(exception)[:200],
+                )
+
+            return success
+
+        except Exception as publish_error:
+            # Never let error publishing break the main execution
+            logger.error(
+                "failed_to_publish_error_event",
+                execution_id=execution_id[:8] if len(execution_id) >= 8 else execution_id,
+                error=str(publish_error),
+            )
+            return False
+
+    def _create_user_friendly_message(
+        self, exception: Exception, category: ErrorCategory, component: str
+    ) -> str:
+        """Create user-friendly error message."""
+
+        messages = {
+            ErrorCategory.RUNTIME_INIT: f"Failed to initialize {component}",
+            ErrorCategory.SKILL_LOADING: f"Could not load skill: {str(exception)[:100]}",
+            ErrorCategory.MCP_CONNECTION: f"Failed to connect to MCP server",
+            ErrorCategory.TOOL_EXECUTION: f"Tool execution failed: {str(exception)[:100]}",
+            ErrorCategory.TIMEOUT: f"Operation timed out in {component}",
+            ErrorCategory.MODEL_ERROR: f"LLM model error: {str(exception)[:100]}",
+            ErrorCategory.AUTHENTICATION: "Authentication failed - please check credentials",
+            ErrorCategory.NETWORK: "Network connection error",
+            ErrorCategory.VALIDATION: f"Validation error: {str(exception)[:100]}",
+        }
+
+        return messages.get(category, str(exception)[:200])
+
+    def _is_retryable_error(
+        self, exception: Exception, category: ErrorCategory
+    ) -> bool:
+        """Determine if error is retryable."""
+
+        # Network and timeout errors are generally retryable
+        if category in [ErrorCategory.NETWORK, ErrorCategory.TIMEOUT]:
+            return True
+
+        # Authentication errors not retryable without credential changes
+        if category == ErrorCategory.AUTHENTICATION:
+            return False
+
+        # Validation errors not retryable without input changes
+        if category == ErrorCategory.VALIDATION:
+            return False
+
+        # Check exception type
+        retryable_types = [
+            "TimeoutError",
+            "ConnectionError",
+            "HTTPError",
+            "TemporaryFailure",
+        ]
+
+        return type(exception).__name__ in retryable_types
+
+    def _get_default_recovery_actions(
+        self, category: ErrorCategory, is_retryable: bool
+    ) -> List[str]:
+        """Get default recovery actions for error category."""
+
+        actions = {
+            ErrorCategory.RUNTIME_INIT: [
+                "Check that all required dependencies are installed",
+                "Verify runtime configuration is correct",
+                "Check worker logs for detailed error information",
+            ],
+            ErrorCategory.SKILL_LOADING: [
+                "Verify skill configuration is correct",
+                "Check that skill dependencies are available",
+                "Review skill permissions and access",
+            ],
+            ErrorCategory.MCP_CONNECTION: [
+                "Verify MCP server is running and accessible",
+                "Check network connectivity",
+                "Review MCP server configuration and credentials",
+            ],
+            ErrorCategory.TIMEOUT: [
+                "Consider increasing timeout settings",
+                "Check system resources and load",
+                "Simplify the operation if possible",
+            ],
+            ErrorCategory.MODEL_ERROR: [
+                "Verify API credentials are valid",
+                "Check API rate limits and quotas",
+                "Try a different model if available",
+            ],
+            ErrorCategory.AUTHENTICATION: [
+                "Verify API credentials are correct",
+                "Check that credentials have required permissions",
+                "Regenerate credentials if expired",
+            ],
+        }
+
+        default_actions = actions.get(category, [
+            "Review error details above",
+            "Check system logs for more information",
+            "Contact support if issue persists",
+        ])
+
+        if is_retryable:
+            default_actions.insert(0, "Retry the operation")
+
+        return default_actions
+
+    def _get_documentation_url(self, category: ErrorCategory) -> Optional[str]:
+        """Get documentation URL for error category."""
+
+        base_url = "https://docs.kubiya.ai/troubleshooting"
+
+        urls = {
+            ErrorCategory.RUNTIME_INIT: f"{base_url}/runtime-initialization",
+            ErrorCategory.SKILL_LOADING: f"{base_url}/skill-configuration",
+            ErrorCategory.MCP_CONNECTION: f"{base_url}/mcp-servers",
+            ErrorCategory.MODEL_ERROR: f"{base_url}/llm-configuration",
+            ErrorCategory.AUTHENTICATION: f"{base_url}/authentication",
+        }
+
+        return urls.get(category, base_url)

control_plane_api/worker/utils/event_batcher.py

@@ -0,0 +1,256 @@
+"""
+Event Batcher for efficient HTTP event publishing.
+
+Batches multiple events into a single HTTP request to reduce overhead,
+especially important for serverless deployments like Vercel where you
+pay per request.
+
+Features:
+- Time-based flushing (default 50ms)
+- Size-based flushing (default 10 events)
+- Priority-based immediate flushing for critical events
+- Automatic background flushing task
+- Graceful shutdown with pending event flush
+
+Usage:
+    batcher = EventBatcher(
+        flush_callback=send_batch_to_server,
+        max_batch_size=10,
+        max_batch_time_ms=50
+    )
+
+    await batcher.start()
+    await batcher.add_event(event_data, priority="normal")
+    await batcher.add_event(critical_event, priority="critical")  # Flushes immediately
+    await batcher.stop()
+"""
+
+import asyncio
+import time
+from typing import Dict, Any, Callable, Literal, List, Optional
+import structlog
+
+logger = structlog.get_logger()
+
+
+class EventBatcher:
+    """
+    Batches events for efficient HTTP publishing.
+
+    Reduces HTTP request count by batching multiple events together,
+    with smart flushing based on time, size, and priority.
+    """
+
+    def __init__(
+        self,
+        flush_callback: Callable[[List[Dict[str, Any]]], Any],
+        max_batch_size: int = 10,
+        max_batch_time_ms: int = 50,
+        enabled: bool = True
+    ):
+        """
+        Initialize event batcher.
+
+        Args:
+            flush_callback: Async callback to send batch (receives list of events)
+            max_batch_size: Maximum events per batch before auto-flush
+            max_batch_time_ms: Maximum milliseconds to wait before auto-flush
+            enabled: Whether batching is enabled (if False, events flush immediately)
+        """
+        self.flush_callback = flush_callback
+        self.max_batch_size = max_batch_size
+        self.max_batch_time_ms = max_batch_time_ms
+        self.enabled = enabled
+
+        # Batch state
+        self.batch: List[Dict[str, Any]] = []
+        self.batch_lock = asyncio.Lock()
+        self.last_flush_time = time.time()
+
+        # Background flush task
+        self._flush_task: Optional[asyncio.Task] = None
+        self._running = False
+
+        # Statistics
+        self.stats = {
+            "total_events": 0,
+            "total_batches": 0,
+            "total_flushes": 0,
+            "time_based_flushes": 0,
+            "size_based_flushes": 0,
+            "critical_flushes": 0,
+            "events_per_batch_avg": 0.0
+        }
+
+    async def start(self):
+        """Start the background flush task."""
+        if not self.enabled:
+            logger.info("event_batching_disabled")
+            return
+
+        if self._running:
+            logger.warning("event_batcher_already_running")
+            return
+
+        self._running = True
+        self._flush_task = asyncio.create_task(self._flush_loop())
+        logger.info(
+            "event_batcher_started",
+            max_batch_size=self.max_batch_size,
+            max_batch_time_ms=self.max_batch_time_ms
+        )
+
+    async def stop(self):
+        """Stop the batcher and flush any pending events."""
+        self._running = False
+
+        # Cancel background task
+        if self._flush_task and not self._flush_task.done():
+            self._flush_task.cancel()
+            try:
+                await self._flush_task
+            except asyncio.CancelledError:
+                pass
+
+        # Flush any remaining events
+        await self._flush_batch("shutdown")
+
+        logger.info(
+            "event_batcher_stopped",
+            stats=self.stats
+        )
+
+    async def add_event(
+        self,
+        event: Dict[str, Any],
+        priority: Literal["normal", "critical"] = "normal"
+    ) -> bool:
+        """
+        Add an event to the batch.
+
+        Args:
+            event: Event data to batch
+            priority: "normal" for regular events, "critical" for immediate flush
+
+        Returns:
+            True if event was added successfully
+        """
+        if not self.enabled:
+            # If batching disabled, send immediately
+            try:
+                result = self.flush_callback([event])
+                if asyncio.iscoroutine(result):
+                    await result
+                return True
+            except Exception as e:
+                logger.error("event_send_failed", error=str(e))
+                return False
+
+        async with self.batch_lock:
+            self.batch.append(event)
+            self.stats["total_events"] += 1
+
+            # Critical events flush immediately
+            if priority == "critical":
+                await self._flush_batch("critical")
+                return True
+
+            # Size-based flush
+            if len(self.batch) >= self.max_batch_size:
+                await self._flush_batch("size")
+                return True
+
+        return True
+
+    async def _flush_loop(self):
+        """Background task that flushes batches based on time."""
+        while self._running:
+            try:
+                # Check every 10ms (more responsive than max_batch_time_ms)
+                await asyncio.sleep(0.01)
+
+                # Check if we need time-based flush
+                async with self.batch_lock:
+                    if not self.batch:
+                        continue
+
+                    elapsed_ms = (time.time() - self.last_flush_time) * 1000
+                    if elapsed_ms >= self.max_batch_time_ms:
+                        await self._flush_batch("time")
+
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error("flush_loop_error", error=str(e))
+                # Continue running even if flush fails
+
+    async def _flush_batch(self, reason: str):
+        """
+        Flush the current batch.
+
+        Args:
+            reason: Reason for flush (time/size/critical/shutdown)
+
+        Note: Caller must hold batch_lock
+        """
+        if not self.batch:
+            return
+
+        batch_to_send = self.batch.copy()
+        batch_size = len(batch_to_send)
+
+        # Clear batch immediately (unlock for new events)
+        self.batch.clear()
+        self.last_flush_time = time.time()
+
+        # Update statistics
+        self.stats["total_batches"] += 1
+        self.stats["total_flushes"] += 1
+
+        if reason == "time":
+            self.stats["time_based_flushes"] += 1
+        elif reason == "size":
+            self.stats["size_based_flushes"] += 1
+        elif reason == "critical":
+            self.stats["critical_flushes"] += 1
+
+        # Calculate average batch size
+        self.stats["events_per_batch_avg"] = (
+            self.stats["total_events"] / self.stats["total_batches"]
+        )
+
+        # Send batch (outside lock)
+        try:
+            logger.debug(
+                "flushing_event_batch",
+                batch_size=batch_size,
+                reason=reason,
+                avg_batch_size=f"{self.stats['events_per_batch_avg']:.1f}"
+            )
+
+            result = self.flush_callback(batch_to_send)
+            if asyncio.iscoroutine(result):
+                await result
+
+            logger.debug(
+                "event_batch_flushed",
+                batch_size=batch_size,
+                reason=reason
+            )
+
+        except Exception as e:
+            logger.error(
+                "batch_flush_failed",
+                error=str(e),
+                batch_size=batch_size,
+                reason=reason
+            )
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Get batching statistics."""
+        return self.stats.copy()
+
+    async def force_flush(self):
+        """Force flush any pending events (useful for testing or critical sections)."""
+        async with self.batch_lock:
+            await self._flush_batch("manual")