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

control_plane_api/app/routers/websocket_client.py

@@ -0,0 +1,479 @@
+"""
+WebSocket Client Endpoint for Execution Streaming.
+
+Provides persistent WebSocket connections for frontend clients to receive
+real-time execution updates, eliminating the 300-second HTTP timeout limitation.
+
+Features:
+- Long-running execution streaming (no timeout)
+- Authentication via JWT tokens
+- Automatic reconnection support with Last-Event-ID
+- Heartbeat/keepalive mechanism (ping/pong)
+- Gap detection and recovery
+- Graceful degradation when services unavailable
+
+Architecture:
+    Browser → WebSocket → Control Plane API → Redis → Execution Events
+                                            ↓
+                                        PostgreSQL (historical messages)
+
+This replaces the SSE-based streaming with a more robust WebSocket solution.
+"""
+
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, status
+from typing import Optional, Dict, Any
+import structlog
+import json
+import asyncio
+from datetime import datetime, timezone
+
+from control_plane_api.app.middleware.auth import get_current_organization
+from control_plane_api.app.lib.redis_client import get_redis_client
+from control_plane_api.app.routers.executions.streaming.streamer import ExecutionStreamer
+from control_plane_api.app.models.execution import Execution
+from control_plane_api.app.database import get_db
+from control_plane_api.app.lib.temporal_client import get_temporal_client
+from sqlalchemy.orm import Session
+
+logger = structlog.get_logger()
+
+router = APIRouter()
+
+
+class ClientConnectionManager:
+    """
+    Manages WebSocket connections for frontend clients.
+
+    Features:
+    - Per-organization connection limits
+    - Connection tracking and cleanup
+    - Statistics and monitoring
+    """
+
+    def __init__(self):
+        # execution_id -> WebSocket
+        self._connections: Dict[str, WebSocket] = {}
+
+        # organization_id -> Set[execution_id]
+        self._org_connections: Dict[str, set] = {}
+
+        # Connection statistics
+        self._stats = {
+            "total_connections": 0,
+            "active_connections": 0,
+            "messages_sent": 0,
+            "errors": 0,
+        }
+
+    async def connect(
+        self,
+        execution_id: str,
+        organization_id: str,
+        websocket: WebSocket,
+    ) -> None:
+        """Register a new client WebSocket connection."""
+        await websocket.accept()
+
+        # Track connection
+        self._connections[execution_id] = websocket
+
+        if organization_id not in self._org_connections:
+            self._org_connections[organization_id] = set()
+        self._org_connections[organization_id].add(execution_id)
+
+        # Update stats
+        self._stats["total_connections"] += 1
+        self._stats["active_connections"] = len(self._connections)
+
+        logger.info(
+            "client_websocket_connected",
+            execution_id=execution_id[:8],
+            organization_id=organization_id[:8],
+            active_connections=self._stats["active_connections"],
+        )
+
+    async def disconnect(self, execution_id: str, organization_id: str) -> None:
+        """Unregister a client WebSocket connection."""
+        if execution_id in self._connections:
+            del self._connections[execution_id]
+
+        if organization_id in self._org_connections:
+            self._org_connections[organization_id].discard(execution_id)
+            if not self._org_connections[organization_id]:
+                del self._org_connections[organization_id]
+
+        self._stats["active_connections"] = len(self._connections)
+
+        logger.info(
+            "client_websocket_disconnected",
+            execution_id=execution_id[:8],
+            organization_id=organization_id[:8],
+            active_connections=self._stats["active_connections"],
+        )
+
+    def get_stats(self) -> Dict[str, int]:
+        """Get connection statistics."""
+        return self._stats.copy()
+
+
+# Global connection manager
+client_manager = ClientConnectionManager()
+
+
+async def send_json(websocket: WebSocket, event_type: str, data: Any) -> None:
+    """
+    Send JSON message via WebSocket.
+
+    Args:
+        websocket: WebSocket connection
+        event_type: Event type (e.g., 'message', 'status', 'connected')
+        data: Event data payload
+    """
+    try:
+        message = {
+            "type": event_type,
+            **data
+        }
+        await websocket.send_text(json.dumps(message))
+        client_manager._stats["messages_sent"] += 1
+        logger.info("websocket_message_sent", event_type=event_type, data_keys=list(data.keys()))
+    except Exception as e:
+        logger.error("failed_to_send_websocket_message", error=str(e), event_type=event_type)
+        client_manager._stats["errors"] += 1
+        raise
+
+
+async def handle_ping_pong(websocket: WebSocket) -> None:
+    """
+    Handle ping/pong heartbeat messages.
+
+    Responds to ping messages with pong to keep connection alive.
+    """
+    try:
+        await send_json(websocket, "pong", {"timestamp": int(datetime.now(timezone.utc).timestamp() * 1000)})
+    except Exception as e:
+        logger.error("failed_to_send_pong", error=str(e))
+
+
+async def handle_auth(websocket: WebSocket, token: str) -> Optional[str]:
+    """
+    Handle authentication message.
+
+    Supports multiple authentication methods:
+    1. JWT tokens with org_id claim (Auth0 tokens)
+    2. Kubiya API keys (validated via Kubiya API)
+
+    Args:
+        websocket: WebSocket connection
+        token: Authentication token (JWT or API key)
+
+    Returns:
+        organization_id if authentication successful, None otherwise
+    """
+    try:
+        # Import auth utilities
+        from control_plane_api.app.middleware.auth import (
+            decode_jwt_token,
+            validate_kubiya_api_key,
+            is_kubiya_api_key_jwt,
+        )
+
+        organization_id = None
+        user_id = None
+
+        # First, try to decode as JWT
+        decoded = decode_jwt_token(token)
+
+        if decoded:
+            # Check for org_id in various claim locations
+            organization_id = (
+                decoded.get('https://kubiya.ai/org_id') or
+                decoded.get('org_id') or
+                decoded.get('organization_id') or
+                decoded.get('organization')  # Kubiya API keys use 'organization' field
+            )
+            user_id = decoded.get('sub') or decoded.get('user_id') or decoded.get('email')
+
+        # If JWT decode failed or org_id not found, try Kubiya API validation
+        if not organization_id:
+            logger.debug("jwt_org_id_not_found_trying_api_validation")
+
+            # Try UserKey format first (for CLI API keys)
+            user_data = await validate_kubiya_api_key(f"UserKey {token}", use_userkey=True)
+
+            if not user_data:
+                # Try Bearer format
+                user_data = await validate_kubiya_api_key(f"Bearer {token}", use_userkey=False)
+
+            if user_data:
+                organization_id = user_data.get('org')
+                user_id = user_data.get('uuid') or user_data.get('email')
+                logger.info("websocket_auth_via_kubiya_api", org=organization_id)
+
+        if not organization_id:
+            logger.error("org_id_missing", decoded_claims=list(decoded.keys()) if decoded else [])
+            await send_json(websocket, "auth_error", {
+                "error": "Organization ID not found in token",
+                "code": "ORG_ID_MISSING",
+            })
+            return None
+
+        logger.info(
+            "websocket_authenticated",
+            organization_id=organization_id[:8] if len(organization_id) > 8 else organization_id,
+            user_id=user_id[:8] if user_id and len(user_id) > 8 else user_id,
+        )
+
+        # Send auth success
+        await send_json(websocket, "auth_success", {
+            "organization_id": organization_id,
+            "user_id": user_id,
+            "authenticated_at": datetime.now(timezone.utc).isoformat(),
+        })
+
+        return organization_id
+
+    except Exception as e:
+        logger.error("authentication_failed", error=str(e), error_type=type(e).__name__)
+        await send_json(websocket, "auth_error", {
+            "error": "Authentication failed",
+            "code": "AUTH_FAILED",
+        })
+        return None
+
+
+@router.websocket("/ws/executions/{execution_id}")
+async def websocket_execution_stream(
+    websocket: WebSocket,
+    execution_id: str,
+    last_event_id: Optional[str] = None,
+    db: Session = Depends(get_db),
+):
+    """
+    WebSocket endpoint for execution streaming.
+
+    Streams execution events to frontend clients with automatic reconnection support.
+
+    Args:
+        websocket: WebSocket connection
+        execution_id: Execution ID to stream
+        last_event_id: Last received event ID (for resumption)
+        db: Database session
+
+    Flow:
+        1. Accept WebSocket connection
+        2. Wait for auth message with JWT token
+        3. Validate token and extract organization_id
+        4. Send 'connected' event
+        5. Load and stream historical messages (last 200)
+        6. Subscribe to Redis for live events
+        7. Stream events until execution completes or client disconnects
+        8. Handle ping/pong for keepalive
+        9. Support resumption via last_event_id
+    """
+    organization_id: Optional[str] = None
+    authenticated = False
+
+    try:
+        # Accept connection first (before manager tracks it)
+        await websocket.accept()
+
+        # Track connection in manager
+        client_manager._connections[execution_id] = websocket
+        client_manager._org_connections.setdefault("pending", set()).add(execution_id)
+        client_manager._stats["total_connections"] += 1
+        client_manager._stats["active_connections"] = len(client_manager._connections)
+
+        logger.info(
+            "client_websocket_connection_started",
+            execution_id=execution_id[:8],
+            last_event_id=last_event_id,
+        )
+
+        # Wait for authentication message (timeout after 5 seconds)
+        try:
+            auth_message = await asyncio.wait_for(websocket.receive_text(), timeout=5.0)
+            auth_data = json.loads(auth_message)
+
+            if auth_data.get("type") == "auth" and "token" in auth_data:
+                organization_id = await handle_auth(websocket, auth_data["token"])
+                if not organization_id:
+                    await websocket.close(code=4001, reason="Authentication failed")
+                    return
+
+                authenticated = True
+
+                # Update connection with actual org_id
+                # Remove from pending
+                if "pending" in client_manager._org_connections:
+                    client_manager._org_connections["pending"].discard(execution_id)
+                    if not client_manager._org_connections["pending"]:
+                        del client_manager._org_connections["pending"]
+
+                # Add to actual org
+                client_manager._org_connections.setdefault(organization_id, set()).add(execution_id)
+
+        except asyncio.TimeoutError:
+            logger.error("auth_timeout", execution_id=execution_id[:8])
+            await websocket.close(code=4002, reason="Authentication timeout")
+            return
+        except json.JSONDecodeError:
+            logger.error("invalid_auth_message", execution_id=execution_id[:8])
+            await websocket.close(code=4003, reason="Invalid authentication message")
+            return
+
+        # Send connected event
+        await send_json(websocket, "connected", {
+            "execution_id": execution_id,
+            "organization_id": organization_id,
+            "connected_at": datetime.now(timezone.utc).isoformat(),
+        })
+
+        # Initialize streaming
+        redis_client = get_redis_client()
+        temporal_client = await get_temporal_client()
+
+        # Get execution type from database
+        execution = db.query(Execution).filter(Execution.id == execution_id).first()
+        execution_type = execution.execution_type if execution else "AGENT"
+
+        streamer = ExecutionStreamer(
+            execution_id=execution_id,
+            organization_id=organization_id,
+            db_session=db,
+            redis_client=redis_client,
+            temporal_client=temporal_client,
+            last_event_id=last_event_id,
+            timeout_seconds=0,  # No timeout for WebSocket
+            execution_type=execution_type,
+        )
+
+        logger.info("streaming_initialized", execution_id=execution_id[:8])
+
+        # Helper to parse SSE events and convert to WebSocket JSON
+        def parse_sse_event(sse_text: str) -> tuple[str, dict]:
+            """Parse SSE format to (event_type, data_dict)."""
+            lines = sse_text.strip().split('\n')
+            event_type = "message"
+            data_lines = []
+
+            # Check for keepalive (just comment lines, no event/data)
+            is_keepalive = all(line.startswith(':') or not line.strip() for line in lines)
+            if is_keepalive:
+                return "keepalive", {}
+
+            for line in lines:
+                if line.startswith("event: "):
+                    event_type = line[7:].strip()
+                elif line.startswith("data: "):
+                    # Collect data line (could be multiline)
+                    data_lines.append(line[6:])
+                elif line and not line.startswith("id:") and not line.startswith(':') and data_lines:
+                    # Continuation of data line
+                    data_lines.append(line)
+
+            # Join all data lines
+            data_str = ''.join(data_lines).strip()
+
+            try:
+                data = json.loads(data_str) if data_str else {}
+            except json.JSONDecodeError as e:
+                logger.warning("json_decode_error", error=str(e), data_str=data_str[:200])
+                data = {"raw": data_str}
+
+            return event_type, data
+
+        # Listen for both streaming events and client messages (ping/pong)
+        async def listen_stream():
+            """Listen for streaming events from ExecutionStreamer."""
+            try:
+                async for sse_event in streamer.stream():
+                    # Parse SSE event and convert to WebSocket JSON
+                    event_type, data = parse_sse_event(sse_event)
+
+                    # Skip keepalive messages
+                    if event_type == "keepalive":
+                        continue
+
+                    # DEBUG: Log what we're parsing
+                    logger.info(
+                        "parsed_sse_event",
+                        event_type=event_type,
+                        data_keys=list(data.keys()) if isinstance(data, dict) else [],
+                        sse_preview=sse_event[:200] if len(sse_event) > 200 else sse_event
+                    )
+
+                    await send_json(websocket, event_type, data)
+            except Exception as e:
+                logger.error("streaming_error", error=str(e), error_type=type(e).__name__)
+                raise
+
+        async def listen_client():
+            """Listen for client messages (ping, resume, etc.)."""
+            try:
+                while True:
+                    message = await websocket.receive_text()
+                    data = json.loads(message)
+
+                    if data.get("type") == "ping":
+                        await handle_ping_pong(websocket)
+
+                    elif data.get("type") == "resume":
+                        # Handle resume request
+                        logger.info("resume_requested", execution_id=execution_id[:8], last_event_id=data.get("last_event_id"))
+
+            except WebSocketDisconnect:
+                pass
+            except Exception as e:
+                logger.error("client_message_error", error=str(e))
+
+        # Run both listeners concurrently
+        # When one finishes (e.g., client disconnect), cancel the other
+        stream_task = asyncio.create_task(listen_stream())
+        client_task = asyncio.create_task(listen_client())
+
+        # Wait for either task to complete
+        done, pending = await asyncio.wait(
+            {stream_task, client_task},
+            return_when=asyncio.FIRST_COMPLETED
+        )
+
+        # Cancel any pending tasks
+        for task in pending:
+            task.cancel()
+            try:
+                await task
+            except asyncio.CancelledError:
+                pass
+
+    except WebSocketDisconnect:
+        logger.info("client_disconnected", execution_id=execution_id[:8])
+
+    except Exception as e:
+        logger.error(
+            "websocket_error",
+            execution_id=execution_id[:8],
+            error=str(e),
+            error_type=type(e).__name__,
+        )
+        try:
+            await send_json(websocket, "error", {"error": str(e)})
+        except:
+            pass
+
+    finally:
+        # Cleanup
+        if organization_id:
+            await client_manager.disconnect(execution_id, organization_id)
+
+        logger.info("client_websocket_closed", execution_id=execution_id[:8])
+
+
+@router.get("/ws/stats")
+async def websocket_stats():
+    """
+    Get WebSocket connection statistics.
+
+    Returns statistics about active connections, messages sent, etc.
+    """
+    return client_manager.get_stats()