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

control_plane_api/app/routers/executions/streaming/live_source.py

@@ -0,0 +1,693 @@
+"""
+LiveEventSource for real-time event streaming from Redis.
+
+This module provides the LiveEventSource class that handles real-time event streaming
+from Redis with proper polling, deduplication, and workflow completion detection.
+
+Key Features:
+- Redis polling with explicit 50ms sleep interval to prevent CPU spinning
+- Integration with MessageDeduplicator to filter duplicate events
+- Workflow completion detection via Temporal status monitoring
+- Graceful degradation when Redis or Temporal unavailable
+- Keepalive and timeout handling for long-running streams
+- Support for both Upstash REST API and standard redis-py clients
+
+Test Strategy:
+- Test Redis polling with mock events at 50ms intervals
+- Test deduplication of overlapping events (history + live)
+- Test completion detection stops streaming when workflow reaches terminal state
+- Test timeout behavior (0 = no timeout, streams until task completes)
+- Test keepalive events sent every 15 seconds
+- Test graceful degradation when Redis unavailable
+- Test graceful degradation when Temporal unavailable
+- Test various event types: message, status, tool_started, tool_completed, etc.
+- Test Redis LLEN and LRANGE operations with both client types
+- Test terminal states: COMPLETED, FAILED, TERMINATED, CANCELLED
+"""
+
+import asyncio
+import json
+import time
+from typing import Any, AsyncGenerator, Dict, List, Optional
+
+import structlog
+
+from .deduplication import MessageDeduplicator
+
+logger = structlog.get_logger(__name__)
+
+
+class LiveEventSource:
+    """
+    Handles real-time event streaming from Redis with polling and completion detection.
+
+    This class polls Redis for new events at 50ms intervals, deduplicates events
+    against historical data, and monitors workflow status to detect completion.
+
+    The streaming stops when:
+    1. Workflow reaches terminal state (COMPLETED, FAILED, CANCELLED, TERMINATED)
+    2. Timeout is reached (default 0 = no timeout, streams until task completes)
+    3. An exception occurs that cannot be recovered
+
+    Events are yielded as dictionaries with the following structure:
+    {
+        "event_type": "message" | "status" | "tool_started" | "tool_completed" | ...,
+        "data": {...},  # Event-specific data
+        "timestamp": "2024-01-15T10:30:00Z"
+    }
+    """
+
+    # Terminal workflow states that indicate streaming should stop
+    TERMINAL_STATES = {"COMPLETED", "FAILED", "CANCELLED", "TERMINATED"}
+
+    # Terminal database execution states (lowercase, as stored in DB)
+    TERMINAL_DB_STATES = {"completed", "failed", "cancelled", "terminated", "interrupted"}
+
+    # Temporal status cache TTL (seconds) to reduce API load
+    TEMPORAL_STATUS_CACHE_TTL = 1.0
+
+    # Database status poll interval (seconds) - check for status changes
+    DB_STATUS_POLL_INTERVAL = 2.0
+
+    def __init__(
+        self,
+        execution_id: str,
+        organization_id: str,
+        redis_client,
+        workflow_handle,
+        deduplicator: MessageDeduplicator,
+        timeout_seconds: int = 0,  # 0 = no timeout, stream until task completes
+        keepalive_interval: int = 15,
+        db_session=None,  # SQLAlchemy session for status polling
+    ):
+        """
+        Initialize LiveEventSource.
+
+        Args:
+            execution_id: Execution ID to stream events for
+            organization_id: Organization ID for authorization
+            redis_client: Redis client instance (UpstashRedisClient or StandardRedisClient)
+            workflow_handle: Temporal workflow handle for status checks (can be None)
+            deduplicator: MessageDeduplicator instance for filtering duplicates
+            timeout_seconds: Maximum streaming duration in seconds (default: 0 = no timeout)
+            keepalive_interval: Seconds between keepalive messages (default: 15)
+            db_session: Database session for polling execution status (optional)
+        """
+        self.execution_id = execution_id
+        self.organization_id = organization_id
+        self.redis_client = redis_client
+        self.workflow_handle = workflow_handle
+        self.deduplicator = deduplicator
+        self.timeout_seconds = timeout_seconds
+        self.keepalive_interval = keepalive_interval
+        self.db_session = db_session
+
+        # Streaming state
+        self._start_time = None
+        self._last_keepalive = None
+        self._last_redis_index = -1  # Track last processed Redis event index
+        self._stopped = False
+        self._is_workflow_running = True  # Assume running until proven otherwise
+
+        # Temporal status caching
+        self._cached_temporal_status = None
+        self._cached_workflow_description = None
+        self._last_temporal_check = 0
+
+        # Database status caching
+        self._cached_db_status = None
+        self._last_db_status_check = 0
+
+        # Redis key for events
+        self._redis_key = f"execution:{execution_id}:events"
+
+        logger.info(
+            "live_event_source_initialized",
+            execution_id=execution_id[:8],
+            timeout_seconds=timeout_seconds,
+            keepalive_interval=keepalive_interval,
+            has_workflow_handle=workflow_handle is not None,
+            has_redis_client=redis_client is not None,
+        )
+
+    async def stream(self) -> AsyncGenerator[Dict[str, Any], None]:
+        """
+        Stream live events from Redis until workflow completes or timeout.
+
+        This is the main entry point for streaming. It polls Redis every 50ms
+        for new events, deduplicates them, and yields them as dictionaries.
+
+        Yields:
+            Event dictionaries with event_type, data, and timestamp
+
+        Example:
+            ```python
+            async for event in live_source.stream():
+                print(f"Event: {event['event_type']}")
+                # Process event...
+            ```
+        """
+        self._start_time = time.time()
+        self._last_keepalive = self._start_time
+
+        logger.info(
+            "live_streaming_started",
+            execution_id=self.execution_id[:8],
+            timeout_seconds=self.timeout_seconds,
+        )
+
+        try:
+            while not self.should_stop:
+                current_time = time.time()
+
+                # Check timeout (skip if timeout_seconds is 0 = no timeout)
+                elapsed = current_time - self._start_time
+                if self.timeout_seconds > 0 and elapsed >= self.timeout_seconds:
+                    logger.warning(
+                        "live_streaming_timeout",
+                        execution_id=self.execution_id[:8],
+                        elapsed_seconds=int(elapsed),
+                        timeout_seconds=self.timeout_seconds,
+                    )
+                    self._stopped = True
+                    break
+
+                # Send keepalive if needed
+                if current_time - self._last_keepalive >= self.keepalive_interval:
+                    yield self._create_keepalive_event()
+                    self._last_keepalive = current_time
+
+                # Check workflow completion (if handle available)
+                if self.workflow_handle is not None:
+                    try:
+                        is_complete = await self._check_completion()
+                        if is_complete:
+                            logger.info(
+                                "workflow_completed_stopping_stream",
+                                execution_id=self.execution_id[:8],
+                                status=self._cached_temporal_status,
+                            )
+                            self._stopped = True
+                            break
+                    except Exception as e:
+                        # Don't fail streaming if status check fails
+                        logger.warning(
+                            "workflow_status_check_failed",
+                            execution_id=self.execution_id[:8],
+                            error=str(e),
+                        )
+
+                # Check database status for changes (every 2 seconds)
+                db_status_event = await self._check_db_status()
+                if db_status_event:
+                    # Emit status change event
+                    yield db_status_event
+
+                # Poll Redis for new events
+                if self.redis_client:
+                    try:
+                        new_events = await self._poll_redis(self._last_redis_index)
+                        for event in new_events:
+                            # Deduplicate event
+                            if self.deduplicator.is_sent(event):
+                                logger.debug(
+                                    "duplicate_event_skipped",
+                                    execution_id=self.execution_id[:8],
+                                    event_type=event.get("event_type"),
+                                    message_id=event.get("message_id"),
+                                )
+                                continue
+
+                            # Mark as sent and yield
+                            self.deduplicator.mark_sent(event)
+                            yield event
+                    except Exception as redis_error:
+                        # Log error but don't fail streaming
+                        logger.error(
+                            "redis_poll_error",
+                            execution_id=self.execution_id[:8],
+                            error=str(redis_error),
+                        )
+                        # Yield degraded state notification
+                        yield self._create_degraded_event(str(redis_error))
+
+                # Sleep 50ms before next poll - reduced from 200ms for faster event delivery
+                # This 4x improvement significantly reduces perceived latency
+                await asyncio.sleep(0.05)
+
+        except Exception as e:
+            logger.error(
+                "live_streaming_error",
+                execution_id=self.execution_id[:8],
+                error=str(e),
+            )
+            raise
+        finally:
+            elapsed = time.time() - self._start_time
+            logger.info(
+                "live_streaming_stopped",
+                execution_id=self.execution_id[:8],
+                elapsed_seconds=int(elapsed),
+                events_processed=self._last_redis_index + 1,
+            )
+
+    async def _poll_redis(self, last_index: int) -> List[Dict[str, Any]]:
+        """
+        Poll Redis for new events since last_index.
+
+        This method:
+        1. Uses LLEN to get total event count
+        2. Uses LRANGE to get new events since last index
+        3. Parses events from JSON
+        4. Reverses them to chronological order (Redis uses LPUSH, newest first)
+        5. Extracts message_id for deduplication
+
+        Args:
+            last_index: Last processed event index (0-based)
+
+        Returns:
+            List of new event dictionaries in chronological order
+
+        Note:
+            Redis stores events in reverse chronological order (LPUSH adds to head).
+            We reverse them here to get chronological order (oldest first).
+        """
+        if not self.redis_client:
+            return []
+
+        try:
+            # Get total event count
+            total_events = await self.redis_client.llen(self._redis_key)
+
+            if total_events is None or total_events == 0:
+                return []
+
+            # Check if there are new events
+            if total_events <= (last_index + 1):
+                return []
+
+            logger.debug(
+                "redis_new_events_found",
+                execution_id=self.execution_id[:8],
+                total=total_events,
+                last_index=last_index,
+                new_count=total_events - (last_index + 1),
+            )
+
+            # Get all events (they're in reverse chronological order from LPUSH)
+            all_redis_events = await self.redis_client.lrange(self._redis_key, 0, -1)
+
+            if not all_redis_events:
+                return []
+
+            # Reverse to get chronological order (oldest first)
+            chronological_events = list(reversed(all_redis_events))
+
+            # Extract only NEW events we haven't processed yet
+            new_events = []
+            for i in range(last_index + 1, len(chronological_events)):
+                event_json = chronological_events[i]
+
+                try:
+                    event_data = json.loads(event_json)
+                    event_type = event_data.get("event_type", "message")
+
+                    # Extract payload based on event structure
+                    # Two formats:
+                    # 1. Message events: flat structure {role, content, timestamp, message_id}
+                    # 2. Other events: nested {event_type, data: {...}, timestamp}
+                    if "data" in event_data and isinstance(event_data["data"], dict):
+                        if event_type == "message" and "role" in event_data["data"]:
+                            # Message events expect flat structure
+                            payload = event_data["data"].copy()
+                        else:
+                            # Chunk events and others expect nested structure
+                            payload = {
+                                "data": event_data["data"],
+                                "timestamp": event_data.get("timestamp"),
+                            }
+                    else:
+                        # Fallback for legacy format
+                        payload = event_data.copy()
+
+                    # Ensure message_id exists for deduplication
+                    if event_type == "message" and isinstance(payload, dict):
+                        if not payload.get("message_id"):
+                            # Generate stable message_id
+                            timestamp = payload.get("timestamp") or event_data.get("timestamp")
+                            role = payload.get("role", "unknown")
+                            if timestamp:
+                                try:
+                                    from datetime import datetime
+                                    timestamp_micros = int(
+                                        datetime.fromisoformat(timestamp.replace("Z", "+00:00")).timestamp()
+                                        * 1000000
+                                    )
+                                except Exception:
+                                    timestamp_micros = int(time.time() * 1000000)
+                            else:
+                                timestamp_micros = int(time.time() * 1000000)
+
+                            generated_id = f"{self.execution_id}_{role}_{timestamp_micros}"
+                            payload["message_id"] = generated_id
+
+                            logger.debug(
+                                "generated_message_id_for_redis_event",
+                                execution_id=self.execution_id[:8],
+                                role=role,
+                                generated_id=generated_id,
+                            )
+
+                    # Store event type and payload
+                    event = {
+                        "event_type": event_type,
+                        **payload,  # Merge payload into event
+                    }
+
+                    new_events.append(event)
+
+                    # Update last processed index
+                    self._last_redis_index = i
+
+                    logger.debug(
+                        "redis_event_parsed",
+                        execution_id=self.execution_id[:8],
+                        event_type=event_type,
+                        index=i,
+                    )
+
+                except json.JSONDecodeError as e:
+                    logger.warning(
+                        "invalid_redis_event_json",
+                        execution_id=self.execution_id[:8],
+                        event=event_json[:100],
+                        error=str(e),
+                    )
+                    # Update index even for invalid events
+                    self._last_redis_index = i
+                    continue
+                except Exception as e:
+                    logger.error(
+                        "redis_event_processing_error",
+                        execution_id=self.execution_id[:8],
+                        event=event_json[:100],
+                        error=str(e),
+                    )
+                    # Update index even for failed events
+                    self._last_redis_index = i
+                    continue
+
+            return new_events
+
+        except Exception as e:
+            logger.error(
+                "redis_poll_failed",
+                execution_id=self.execution_id[:8],
+                error=str(e),
+            )
+            return []
+
+    async def _check_completion(self) -> bool:
+        """
+        Check if workflow is in terminal state (completed, failed, cancelled).
+
+        This method uses cached status if within TTL (1 second) to reduce
+        Temporal API load. Fresh status is fetched if cache expired.
+
+        Returns:
+            True if workflow is in terminal state, False otherwise
+
+        Note:
+            Temporal execution status enum values:
+            - RUNNING: Workflow is actively processing
+            - COMPLETED: Workflow completed successfully
+            - FAILED: Workflow failed with error
+            - CANCELLED: Workflow was cancelled by user
+            - TERMINATED: Workflow was terminated
+            - TIMED_OUT: Workflow exceeded timeout
+            - CONTINUED_AS_NEW: Workflow continued as new execution
+        """
+        if not self.workflow_handle:
+            # No workflow handle - can't check completion
+            return False
+
+        try:
+            current_time = time.time()
+
+            # Use cached status if within TTL
+            if (
+                self._cached_temporal_status
+                and (current_time - self._last_temporal_check) < self.TEMPORAL_STATUS_CACHE_TTL
+            ):
+                temporal_status = self._cached_temporal_status
+                logger.debug(
+                    "using_cached_temporal_status",
+                    execution_id=self.execution_id[:8],
+                    status=temporal_status,
+                )
+            else:
+                # Cache expired or not set - fetch fresh status
+                t0 = time.time()
+                description = await self.workflow_handle.describe()
+                temporal_status = description.status.name  # Get enum name (e.g., "RUNNING")
+                describe_duration = int((time.time() - t0) * 1000)
+
+                # Update cache
+                self._cached_temporal_status = temporal_status
+                self._cached_workflow_description = description
+                self._last_temporal_check = t0
+
+                logger.debug(
+                    "temporal_status_fetched",
+                    execution_id=self.execution_id[:8],
+                    status=temporal_status,
+                    duration_ms=describe_duration,
+                )
+
+                # Log slow describe calls (>100ms)
+                if describe_duration > 100:
+                    logger.warning(
+                        "slow_temporal_describe",
+                        execution_id=self.execution_id[:8],
+                        duration_ms=describe_duration,
+                    )
+
+            # Update running state
+            previous_running = self._is_workflow_running
+            self._is_workflow_running = temporal_status == "RUNNING"
+
+            # Log state changes
+            if previous_running != self._is_workflow_running:
+                logger.info(
+                    "workflow_running_state_changed",
+                    execution_id=self.execution_id[:8],
+                    temporal_status=temporal_status,
+                    is_running=self._is_workflow_running,
+                )
+
+            # Check if terminal state reached
+            return temporal_status in self.TERMINAL_STATES
+
+        except Exception as e:
+            error_msg = str(e).lower()
+
+            # Check if workflow is not found - could be:
+            # 1. Not started yet (race condition) - keep waiting
+            # 2. Already completed and cleaned up - check DB
+            if "workflow not found" in error_msg or "not found" in error_msg:
+                # Track how many times we've seen "not found"
+                if not hasattr(self, '_not_found_count'):
+                    self._not_found_count = 0
+                self._not_found_count += 1
+
+                # First few times: might not be started yet, wait and retry
+                if self._not_found_count <= 10:  # ~10 seconds of retries
+                    logger.debug(
+                        "workflow_not_found_waiting",
+                        execution_id=self.execution_id[:8],
+                        attempt=self._not_found_count,
+                    )
+                    return False  # Keep streaming, check again later
+
+                # After 10 attempts: workflow probably doesn't exist
+                logger.info(
+                    "workflow_not_found_stopping_temporal_polling",
+                    execution_id=self.execution_id[:8],
+                    attempts=self._not_found_count,
+                )
+                # Disable further workflow checks by clearing the handle
+                # Stream will continue from Redis/DB only
+                self.workflow_handle = None
+                return False
+
+            logger.warning(
+                "temporal_status_check_error",
+                execution_id=self.execution_id[:8],
+                error=str(e),
+            )
+            # Don't treat transient errors as completion - keep streaming
+            return False
+
+    async def _check_db_status(self) -> Optional[Dict[str, Any]]:
+        """
+        Check database execution status and emit status event if changed.
+
+        This polls the database every 2 seconds to detect status changes
+        (running, waiting_for_input, completed, failed, etc.) and returns
+        a status event if the status has changed since last check.
+
+        Returns:
+            Status event dict if status changed, None otherwise
+        """
+        if not self.db_session:
+            return None
+
+        try:
+            current_time = time.time()
+
+            # Check if we should poll (respect poll interval)
+            if (current_time - self._last_db_status_check) < self.DB_STATUS_POLL_INTERVAL:
+                return None
+
+            # Query execution status from database
+            from control_plane_api.app.models.execution import Execution
+            import uuid as uuid_module
+
+            execution = self.db_session.query(Execution).filter(
+                Execution.id == uuid_module.UUID(self.execution_id),
+                Execution.organization_id == self.organization_id
+            ).first()
+
+            if not execution:
+                return None
+
+            db_status = execution.status
+            self._last_db_status_check = current_time
+
+            # Check if status changed
+            status_event = None
+            if db_status != self._cached_db_status:
+                logger.info(
+                    "execution_status_changed",
+                    execution_id=self.execution_id[:8],
+                    old_status=self._cached_db_status,
+                    new_status=db_status,
+                )
+
+                self._cached_db_status = db_status
+
+                # Return status event
+                status_event = {
+                    "event_type": "status",
+                    "status": db_status,
+                    "data": {
+                        "execution_id": self.execution_id,
+                        "status": db_status,
+                        "source": "database_poll",
+                    }
+                }
+
+            # Check if DB status is a terminal state - signal completion
+            # This is a fallback when Temporal workflow handle is unavailable
+            if db_status in self.TERMINAL_DB_STATES:
+                logger.info(
+                    "db_terminal_state_detected",
+                    execution_id=self.execution_id[:8],
+                    db_status=db_status,
+                )
+                self._stopped = True
+
+            return status_event
+
+        except Exception as e:
+            logger.warning(
+                "db_status_check_error",
+                execution_id=self.execution_id[:8],
+                error=str(e),
+            )
+            return None
+
+    @property
+    def should_stop(self) -> bool:
+        """
+        Whether streaming should stop.
+
+        Returns:
+            True if streaming should stop (completed or timeout), False otherwise
+        """
+        return self._stopped
+
+    def _create_keepalive_event(self) -> Dict[str, Any]:
+        """
+        Create a keepalive event.
+
+        Keepalive events are sent periodically to prevent connection timeout
+        and to inform the client that streaming is still active.
+
+        Returns:
+            Keepalive event dictionary
+        """
+        elapsed = time.time() - self._start_time
+        remaining = max(0, self.timeout_seconds - elapsed)
+
+        event = {
+            "event_type": "keepalive",
+            "data": {
+                "execution_id": self.execution_id,
+                "elapsed_seconds": int(elapsed),
+                "remaining_seconds": int(remaining),
+            },
+            "timestamp": self._current_timestamp(),
+        }
+
+        logger.debug(
+            "keepalive_sent",
+            execution_id=self.execution_id[:8],
+            elapsed_seconds=int(elapsed),
+            remaining_seconds=int(remaining),
+        )
+
+        return event
+
+    def _create_degraded_event(self, reason: str) -> Dict[str, Any]:
+        """
+        Create a degraded state event.
+
+        Degraded events inform the client that streaming quality is reduced
+        due to Redis unavailability or other issues.
+
+        Args:
+            reason: Reason for degraded state
+
+        Returns:
+            Degraded event dictionary
+        """
+        event = {
+            "event_type": "degraded",
+            "data": {
+                "reason": "redis_unavailable",
+                "fallback": "temporal_polling",
+                "message": f"Real-time events unavailable: {reason}",
+                "execution_id": self.execution_id,
+            },
+            "timestamp": self._current_timestamp(),
+        }
+
+        logger.warning(
+            "degraded_state_notification",
+            execution_id=self.execution_id[:8],
+            reason=reason,
+        )
+
+        return event
+
+    def _current_timestamp(self) -> str:
+        """
+        Get current timestamp in ISO format.
+
+        Returns:
+            ISO format timestamp string (e.g., "2024-01-15T10:30:00Z")
+        """
+        from datetime import datetime, timezone
+
+        return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")