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

control_plane_api/app/routers/executions/__init__.py

@@ -0,0 +1,28 @@
+"""
+Execution streaming components.
+
+This package contains modular components for handling streaming execution data,
+extracted from the monolithic executions.py router.
+"""
+
+# Import main router and utilities from parent executions.py
+# This allows: from control_plane_api.app.routers.executions import router, validate_job_exists
+import sys
+from pathlib import Path
+
+# Import from the sibling executions.py file (not this package)
+_parent_dir = Path(__file__).parent.parent
+_executions_module_path = _parent_dir / "executions.py"
+
+# We need to import from ../executions.py, not this package
+# Python's import system sees this directory first, so we explicitly import the file
+import importlib.util
+spec = importlib.util.spec_from_file_location("_executions_module", _executions_module_path)
+_executions_module = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(_executions_module)
+
+# Re-export the main router and utilities
+router = _executions_module.router
+validate_job_exists = _executions_module.validate_job_exists
+
+__all__ = ["router", "validate_job_exists"]

control_plane_api/app/routers/executions/router.py

@@ -0,0 +1,286 @@
+"""
+Clean streaming router using ExecutionStreamer architecture.
+
+This module provides a clean FastAPI router for execution streaming
+that replaces the monolithic ~2,100 line inline implementation with
+the modular ExecutionStreamer class.
+
+Architecture:
+- Immediate connection acknowledgment (<50ms)
+- Progressive historical message loading
+- Live event streaming with gap recovery
+- Graceful degradation when services unavailable
+- Last-Event-ID resumption support
+
+The endpoint uses ExecutionStreamer which orchestrates:
+- HistoryLoader: Database message loading with Temporal fallback
+- LiveEventSource: Redis event streaming with workflow state polling
+- MessageDeduplicator: Deduplication across history and live phases
+- EventBuffer: Gap detection and replay for reconnection
+- EventFormatter: Consistent SSE formatting
+- WorkerHealthChecker: Graceful degradation mode detection
+
+Test Strategy:
+- Integration test with real execution IDs
+- Test all event types are emitted correctly
+- Test Last-Event-ID resumption works
+- Test graceful degradation modes
+- Test backward compatibility with existing clients
+"""
+
+import asyncio
+import uuid as uuid_module
+from typing import Optional
+
+import structlog
+from fastapi import APIRouter, Depends, HTTPException, Request
+from fastapi.responses import StreamingResponse
+from sqlalchemy.orm import Session
+
+from control_plane_api.app.database import get_db
+from control_plane_api.app.lib.redis_client import get_redis_client
+from control_plane_api.app.lib.temporal_client import get_temporal_client
+from control_plane_api.app.middleware.auth import get_current_organization
+from control_plane_api.app.models.execution import Execution
+
+from .streaming import ExecutionStreamer
+from .services.worker_health import WorkerHealthChecker
+
+logger = structlog.get_logger(__name__)
+
+router = APIRouter()
+
+
+@router.get("/{execution_id}/stream")
+async def stream_execution(
+    execution_id: str,
+    request: Request,
+    last_event_id: Optional[str] = None,
+    organization: dict = Depends(get_current_organization),
+    db: Session = Depends(get_db),
+):
+    """
+    Stream execution updates using Server-Sent Events (SSE).
+
+    This endpoint provides real-time execution updates by combining:
+    1. Historical messages from database (last 200 messages)
+    2. Live events from Redis (sub-second latency)
+    3. Workflow state from Temporal (200ms polling)
+
+    The endpoint supports:
+    - Progressive history loading (no blocking on full history)
+    - Last-Event-ID resumption for reconnection
+    - Gap detection and notification
+    - Graceful degradation when services unavailable
+    - Backward compatible with existing clients
+
+    Architecture:
+    This endpoint uses the ExecutionStreamer class which orchestrates
+    all phases of streaming:
+    1. Immediate connection (<50ms)
+    2. History loading (progressive, one message at a time)
+    3. History completion notification
+    4. Live event streaming (until workflow completes or timeout)
+
+    Gap Recovery:
+    - Supports Last-Event-ID pattern for reconnection
+    - Client sends last_event_id query param or Last-Event-ID header
+    - Server resumes from that point or detects gaps
+    - Buffered events replayed if available
+    - Client notified if gap unrecoverable
+
+    SSE Event Types:
+    - connected: Initial connection acknowledgment
+    - message: Chat message (history or live)
+    - message_chunk: Streaming message chunk
+    - member_message_chunk: Team member message chunk
+    - history_complete: All historical messages sent
+    - status: Execution status update
+    - tool_started: Tool execution started
+    - tool_completed: Tool execution completed
+    - degraded: Service degradation notification
+    - recovered: Service recovery notification
+    - gap_detected: Gap in event stream detected
+    - error: Error occurred
+    - keepalive: Connection keepalive (comment)
+
+    SSE Format:
+    - id: {execution_id}_{counter}_{timestamp_micros}
+    - event: {event_type}
+    - data: {json object}
+
+    Args:
+        execution_id: Execution ID to stream
+        request: FastAPI request object (for headers)
+        last_event_id: Last event ID received by client (for resumption)
+        organization: Authenticated organization (injected by auth middleware)
+        db: Database session (injected by dependency)
+
+    Returns:
+        StreamingResponse with text/event-stream content
+
+    Raises:
+        HTTPException 404: Execution not found or not authorized
+        HTTPException 500: Critical initialization error
+    """
+
+    # ========== AUTHENTICATION & AUTHORIZATION ==========
+    # Already handled by get_current_organization dependency
+    logger.info(
+        "stream_request_received",
+        execution_id=execution_id[:8],
+        organization_id=organization["id"][:8],
+        has_last_event_id=bool(last_event_id or request.headers.get("Last-Event-ID")),
+    )
+
+    # ========== GET EXECUTION TYPE FROM DATABASE ==========
+    # This determines the workflow ID format:
+    # - TEAM: "team-execution-{id}"
+    # - AGENT: "agent-execution-{id}" (default)
+    try:
+        execution_record = db.query(Execution).filter(
+            Execution.id == uuid_module.UUID(execution_id),
+            Execution.organization_id == organization["id"]
+        ).first()
+
+        if not execution_record:
+            raise HTTPException(
+                status_code=404,
+                detail=f"Execution {execution_id} not found or not authorized"
+            )
+
+        execution_type = execution_record.execution_type or "AGENT"
+
+        logger.info(
+            "execution_type_determined",
+            execution_id=execution_id[:8],
+            execution_type=execution_type,
+        )
+
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(
+            "failed_to_query_execution",
+            execution_id=execution_id[:8],
+            error=str(e),
+            error_type=type(e).__name__,
+        )
+        raise HTTPException(
+            status_code=500,
+            detail=f"Failed to query execution: {str(e)}"
+        )
+
+    # ========== INITIALIZE INFRASTRUCTURE CLIENTS ==========
+    # Get clients with graceful degradation support
+    # - Temporal client: 2s timeout to fail fast when worker down
+    # - Redis client: Can be None if not configured
+    # - Database session: Already available from dependency
+
+    # Redis client (optional, graceful degradation if None)
+    redis_client = get_redis_client()
+    if not redis_client:
+        logger.warning(
+            "redis_not_configured",
+            execution_id=execution_id[:8],
+            note="Live streaming will be unavailable"
+        )
+
+    # Temporal client with fast timeout (optional, graceful degradation)
+    temporal_client = None
+    try:
+        temporal_client = await asyncio.wait_for(
+            get_temporal_client(),
+            timeout=2.0  # Fail fast - don't block streaming
+        )
+        logger.info(
+            "temporal_client_connected",
+            execution_id=execution_id[:8],
+        )
+    except asyncio.TimeoutError:
+        logger.warning(
+            "temporal_connection_timeout",
+            execution_id=execution_id[:8],
+            timeout_seconds=2.0,
+            note="Workflow queries will be unavailable"
+        )
+    except Exception as temporal_error:
+        logger.warning(
+            "temporal_connection_failed",
+            execution_id=execution_id[:8],
+            error=str(temporal_error),
+            error_type=type(temporal_error).__name__,
+            note="Continuing without Temporal support"
+        )
+
+    # ========== PARSE LAST-EVENT-ID ==========
+    # Check both query param and header (EventSource sets header automatically)
+    parsed_last_event_id = last_event_id or request.headers.get("Last-Event-ID")
+
+    if parsed_last_event_id:
+        logger.info(
+            "resumption_requested",
+            execution_id=execution_id[:8],
+            last_event_id=parsed_last_event_id,
+        )
+
+    # ========== CREATE WORKER HEALTH CHECKER ==========
+    # This determines degradation mode based on service availability
+    health_checker = WorkerHealthChecker(
+        temporal_client=temporal_client,
+        redis_client=redis_client,
+        db_session=db,
+    )
+
+    # ========== INSTANTIATE EXECUTION STREAMER ==========
+    # This orchestrates all phases of streaming:
+    # 1. Immediate connection
+    # 2. Resumption (if Last-Event-ID provided)
+    # 3. Historical message loading
+    # 4. History completion notification
+    # 5. Live event streaming
+    try:
+        streamer = ExecutionStreamer(
+            execution_id=execution_id,
+            organization_id=organization["id"],
+            db_session=db,
+            redis_client=redis_client,
+            temporal_client=temporal_client,
+            last_event_id=parsed_last_event_id,
+            timeout_seconds=0,  # 0 = no timeout, stream until task completes
+            execution_type=execution_type,
+            health_checker=health_checker,
+        )
+
+        logger.info(
+            "streamer_initialized",
+            execution_id=execution_id[:8],
+            execution_type=execution_type,
+            has_redis=bool(redis_client),
+            has_temporal=bool(temporal_client),
+            has_last_event_id=bool(parsed_last_event_id),
+        )
+
+    except Exception as e:
+        logger.error(
+            "failed_to_initialize_streamer",
+            execution_id=execution_id[:8],
+            error=str(e),
+            error_type=type(e).__name__,
+        )
+        raise HTTPException(
+            status_code=500,
+            detail=f"Failed to initialize streaming: {str(e)}"
+        )
+
+    # ========== RETURN STREAMING RESPONSE ==========
+    # StreamingResponse will call streamer.stream() which yields SSE events
+    return StreamingResponse(
+        streamer.stream(),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+            "X-Accel-Buffering": "no",  # Disable nginx buffering for real-time
+        }
+    )

control_plane_api/app/routers/executions/services/__init__.py

@@ -0,0 +1,22 @@
+"""Execution services package
+
+This package contains service classes for execution-related functionality:
+- worker_health: Worker and service health checking with graceful degradation
+- status_service: Cached workflow status queries to reduce Temporal API load
+"""
+
+from control_plane_api.app.routers.executions.services.worker_health import (
+    WorkerHealthChecker,
+    DegradationMode,
+    CAPABILITIES,
+)
+from control_plane_api.app.routers.executions.services.status_service import (
+    StatusService,
+)
+
+__all__ = [
+    "WorkerHealthChecker",
+    "DegradationMode",
+    "CAPABILITIES",
+    "StatusService",
+]

control_plane_api/app/routers/executions/services/demo_worker_health.py

@@ -0,0 +1,156 @@
+#!/usr/bin/env python
+"""Demo script for WorkerHealthChecker service
+
+This script demonstrates the WorkerHealthChecker service in action.
+It shows health checks, degradation modes, and capabilities.
+
+Usage:
+    cd /Users/shaked/projects/agent-control-plane/control_plane_api/app/routers/executions/services
+    python demo_worker_health.py
+"""
+
+import asyncio
+import sys
+from pathlib import Path
+
+# Add parent directory to path for imports
+sys.path.insert(0, str(Path(__file__).parent))
+
+from worker_health import (
+    WorkerHealthChecker,
+    DegradationMode,
+    CAPABILITIES,
+)
+
+
+def print_section(title: str):
+    """Print a section header."""
+    print(f"\n{'=' * 60}")
+    print(f"  {title}")
+    print(f"{'=' * 60}\n")
+
+
+async def demo_no_services():
+    """Demo: No services available."""
+    print_section("Demo 1: No Services Available")
+
+    checker = WorkerHealthChecker(
+        temporal_client=None,
+        redis_client=None,
+        db_session=None,
+    )
+
+    print("Checking service health...")
+    results = await checker.check_all()
+    print(f"  Temporal: {'✓' if results['temporal'] else '✗'}")
+    print(f"  Redis: {'✓' if results['redis'] else '✗'}")
+    print(f"  Database: {'✓' if results['database'] else '✗'}")
+
+    print("\nDetermining degradation mode...")
+    mode = await checker.get_degradation_mode()
+    print(f"  Mode: {mode.value}")
+
+    print("\nAvailable capabilities:")
+    capabilities = checker.get_capabilities(mode)
+    if capabilities:
+        for cap in capabilities:
+            print(f"  - {cap}")
+    else:
+        print("  (none)")
+
+
+async def demo_cache_behavior():
+    """Demo: Cache behavior."""
+    print_section("Demo 2: Health Check Caching")
+
+    checker = WorkerHealthChecker(
+        temporal_client=None,
+        redis_client=None,
+        db_session=None,
+        cache_ttl=5  # 5 second cache
+    )
+
+    print("First health check (will hit services)...")
+    start = asyncio.get_event_loop().time()
+    await checker.check_all()
+    duration1 = asyncio.get_event_loop().time() - start
+    print(f"  Duration: {duration1:.3f}s")
+
+    print("\nSecond health check (should use cache)...")
+    start = asyncio.get_event_loop().time()
+    await checker.check_all()
+    duration2 = asyncio.get_event_loop().time() - start
+    print(f"  Duration: {duration2:.3f}s")
+
+    if duration2 < duration1:
+        print(f"  ✓ Cache is working! Second check was faster.")
+    else:
+        print(f"  ℹ Both checks were similar (expected if services are None)")
+
+    print("\nClearing cache...")
+    checker.clear_cache()
+    print("  ✓ Cache cleared")
+
+    print("\nThird health check (should hit services again)...")
+    start = asyncio.get_event_loop().time()
+    await checker.check_all()
+    duration3 = asyncio.get_event_loop().time() - start
+    print(f"  Duration: {duration3:.3f}s")
+
+
+async def demo_all_degradation_modes():
+    """Demo: All degradation modes."""
+    print_section("Demo 3: All Degradation Modes")
+
+    print("Degradation mode capabilities:\n")
+
+    for mode in DegradationMode:
+        capabilities = CAPABILITIES.get(mode, [])
+        print(f"{mode.value.upper()}")
+        print(f"  Capabilities: {', '.join(capabilities) if capabilities else 'none'}")
+        print()
+
+
+async def demo_timeout_behavior():
+    """Demo: Timeout behavior simulation."""
+    print_section("Demo 4: Timeout Behavior")
+
+    print("Health check timeouts:")
+    print(f"  Temporal: {WorkerHealthChecker.TEMPORAL_TIMEOUT}s")
+    print(f"  Redis: {WorkerHealthChecker.REDIS_TIMEOUT}s")
+    print(f"  Database: {WorkerHealthChecker.DATABASE_TIMEOUT}s")
+
+    print("\nThese timeouts ensure the API remains responsive even when")
+    print("services are down or slow. Without timeouts, a client request")
+    print("could hang for 30+ seconds waiting for services to fail.")
+
+
+async def main():
+    """Run all demos."""
+    print("\n" + "=" * 60)
+    print("  WorkerHealthChecker Service Demo")
+    print("=" * 60)
+    print("\nThis demo shows how the WorkerHealthChecker service works.")
+    print("It demonstrates health checks, degradation modes, and caching.")
+
+    try:
+        await demo_no_services()
+        await demo_cache_behavior()
+        await demo_all_degradation_modes()
+        await demo_timeout_behavior()
+
+        print_section("Demo Complete")
+        print("✓ All demos completed successfully")
+        print("\nFor integration examples, see:")
+        print("  - USAGE_EXAMPLE.md")
+        print("  - test_worker_health.py")
+        print()
+
+    except Exception as e:
+        print(f"\n✗ Demo failed: {e}")
+        import traceback
+        traceback.print_exc()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())