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

control_plane_api/app/routers/websocket_executions_status.py

@@ -0,0 +1,437 @@
+"""
+WebSocket endpoint for real-time execution status updates across all executions.
+
+Allows the kanban board to receive instant status changes without polling.
+This is a lightweight endpoint that only streams status changes, not full
+execution content (messages, tool calls, etc.).
+
+Features:
+- Long-running connection (no timeout)
+- Authentication via JWT tokens (same as websocket_client.py)
+- Subscribes to organization-wide status updates
+- Heartbeat/keepalive mechanism
+
+Architecture:
+    Browser → WebSocket → Control Plane API → Redis Pub/Sub → Status Events
+
+Events sent:
+- execution_status: { execution_id, status, updated_at }
+- execution_created: { execution_id, status, entity_type, entity_name, prompt_preview }
+- execution_deleted: { execution_id }
+"""
+
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, status
+from typing import Optional, Dict, Any, Set
+import structlog
+import json
+import asyncio
+from datetime import datetime, timezone
+
+from control_plane_api.app.lib.redis_client import get_redis_client
+
+logger = structlog.get_logger()
+
+router = APIRouter()
+
+
+class StatusConnectionManager:
+    """
+    Manages WebSocket connections for organization-wide status updates.
+
+    Features:
+    - Per-organization connection tracking
+    - Connection statistics
+    - Cleanup on disconnect
+    """
+
+    def __init__(self):
+        # connection_id -> WebSocket
+        self._connections: Dict[str, WebSocket] = {}
+
+        # organization_id -> Set[connection_id]
+        self._org_connections: Dict[str, Set[str]] = {}
+
+        # Connection statistics
+        self._stats = {
+            "total_connections": 0,
+            "active_connections": 0,
+            "messages_sent": 0,
+            "errors": 0,
+        }
+
+        # Connection counter for unique IDs
+        self._connection_counter = 0
+
+    def _generate_connection_id(self) -> str:
+        """Generate unique connection ID."""
+        self._connection_counter += 1
+        return f"status_conn_{self._connection_counter}_{int(datetime.now(timezone.utc).timestamp() * 1000)}"
+
+    async def connect(
+        self,
+        organization_id: str,
+        websocket: WebSocket,
+    ) -> str:
+        """Register a new status WebSocket connection."""
+        connection_id = self._generate_connection_id()
+
+        # Track connection
+        self._connections[connection_id] = websocket
+
+        if organization_id not in self._org_connections:
+            self._org_connections[organization_id] = set()
+        self._org_connections[organization_id].add(connection_id)
+
+        # Update stats
+        self._stats["total_connections"] += 1
+        self._stats["active_connections"] = len(self._connections)
+
+        logger.info(
+            "status_websocket_connected",
+            connection_id=connection_id,
+            organization_id=organization_id[:8] if len(organization_id) > 8 else organization_id,
+            active_connections=self._stats["active_connections"],
+        )
+
+        return connection_id
+
+    async def disconnect(self, connection_id: str, organization_id: str) -> None:
+        """Unregister a status WebSocket connection."""
+        if connection_id in self._connections:
+            del self._connections[connection_id]
+
+        if organization_id in self._org_connections:
+            self._org_connections[organization_id].discard(connection_id)
+            if not self._org_connections[organization_id]:
+                del self._org_connections[organization_id]
+
+        self._stats["active_connections"] = len(self._connections)
+
+        logger.info(
+            "status_websocket_disconnected",
+            connection_id=connection_id,
+            organization_id=organization_id[:8] if len(organization_id) > 8 else organization_id,
+            active_connections=self._stats["active_connections"],
+        )
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Get connection statistics."""
+        return {
+            **self._stats,
+            "connections_by_org": {
+                org_id[:8] if len(org_id) > 8 else org_id: len(conns)
+                for org_id, conns in self._org_connections.items()
+            },
+        }
+
+
+# Global connection manager
+status_manager = StatusConnectionManager()
+
+
+async def send_json(websocket: WebSocket, event_type: str, data: Dict[str, Any]) -> None:
+    """
+    Send JSON message via WebSocket.
+
+    Args:
+        websocket: WebSocket connection
+        event_type: Event type (e.g., 'execution_status', 'connected')
+        data: Event data payload
+    """
+    try:
+        message = {
+            "type": event_type,
+            **data
+        }
+        await websocket.send_text(json.dumps(message))
+        status_manager._stats["messages_sent"] += 1
+    except Exception as e:
+        logger.error("failed_to_send_status_message", error=str(e), event_type=event_type)
+        status_manager._stats["errors"] += 1
+        raise
+
+
+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,
+        )
+
+        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("status_ws_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("status_websocket_auth_via_kubiya_api", org=organization_id)
+
+        if not organization_id:
+            logger.error("status_ws_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(
+            "status_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("status_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/status")
+async def websocket_execution_status_stream(
+    websocket: WebSocket,
+):
+    """
+    WebSocket endpoint for organization-wide execution status updates.
+
+    Streams lightweight status events for ALL executions in the organization,
+    enabling real-time kanban board updates without polling.
+
+    Protocol:
+        1. Client connects
+        2. Client sends auth message: {"type": "auth", "token": "jwt..."}
+        3. Server authenticates and sends 'auth_success'
+        4. Server sends 'connected' event
+        5. Server subscribes to Redis pub/sub for organization status events
+        6. Server streams events until client disconnects
+        7. Client sends ping, server responds with pong (keepalive)
+
+    Events sent to client:
+        - auth_success: Authentication successful
+        - auth_error: Authentication failed
+        - connected: Connection established
+        - execution_status: { execution_id, status, updated_at }
+        - pong: Response to ping (keepalive)
+        - error: Error occurred
+
+    Client messages:
+        - {"type": "auth", "token": "..."}: Authenticate
+        - {"type": "ping"}: Keepalive ping
+    """
+    organization_id: Optional[str] = None
+    connection_id: Optional[str] = None
+    pubsub = None
+
+    try:
+        # Accept connection
+        await websocket.accept()
+
+        logger.info("status_websocket_connection_started")
+
+        # 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
+
+                # Register connection
+                connection_id = await status_manager.connect(organization_id, websocket)
+
+        except asyncio.TimeoutError:
+            logger.error("status_ws_auth_timeout")
+            await websocket.close(code=4002, reason="Authentication timeout")
+            return
+        except json.JSONDecodeError:
+            logger.error("status_ws_invalid_auth_message")
+            await websocket.close(code=4003, reason="Invalid authentication message")
+            return
+
+        # Send connected event
+        await send_json(websocket, "connected", {
+            "connection_id": connection_id,
+            "organization_id": organization_id,
+            "connected_at": datetime.now(timezone.utc).isoformat(),
+        })
+
+        # Get Redis client and subscribe to status channel
+        redis_client = get_redis_client()
+        if not redis_client:
+            logger.error("status_ws_redis_not_available")
+            await send_json(websocket, "error", {"error": "Redis not available"})
+            await websocket.close(code=1011, reason="Redis not available")
+            return
+
+        # Subscribe to organization-wide status channel
+        status_channel = f"organization:{organization_id}:execution_status"
+        pubsub = redis_client.pubsub()
+        await pubsub.subscribe(status_channel)
+
+        logger.info(
+            "status_websocket_subscribed",
+            channel=status_channel,
+            organization_id=organization_id[:8] if len(organization_id) > 8 else organization_id,
+        )
+
+        # Listen for both Redis events and client messages
+        async def listen_redis():
+            """Listen for Redis pub/sub status events."""
+            try:
+                async for message in pubsub.listen():
+                    if message["type"] == "message":
+                        try:
+                            data = json.loads(message["data"])
+                            event_type = data.get("event_type", "execution_status")
+
+                            # Forward the event to WebSocket
+                            await send_json(websocket, event_type, {
+                                "execution_id": data.get("execution_id"),
+                                "status": data.get("status"),
+                                "updated_at": data.get("updated_at") or datetime.now(timezone.utc).isoformat(),
+                                "entity_type": data.get("entity_type"),
+                                "entity_name": data.get("entity_name"),
+                            })
+
+                            logger.debug(
+                                "status_event_forwarded",
+                                execution_id=data.get("execution_id", "")[:8],
+                                status=data.get("status"),
+                            )
+
+                        except json.JSONDecodeError as e:
+                            logger.warning("status_ws_invalid_redis_message", error=str(e))
+
+            except asyncio.CancelledError:
+                pass
+            except Exception as e:
+                logger.error("status_redis_listener_error", error=str(e))
+
+        async def listen_client():
+            """Listen for client messages (ping, etc.)."""
+            try:
+                while True:
+                    message = await websocket.receive_text()
+                    data = json.loads(message)
+
+                    if data.get("type") == "ping":
+                        await send_json(websocket, "pong", {
+                            "timestamp": int(datetime.now(timezone.utc).timestamp() * 1000)
+                        })
+
+            except WebSocketDisconnect:
+                pass
+            except asyncio.CancelledError:
+                pass
+            except Exception as e:
+                logger.error("status_client_listener_error", error=str(e))
+
+        # Run both listeners concurrently
+        redis_task = asyncio.create_task(listen_redis())
+        client_task = asyncio.create_task(listen_client())
+
+        # Wait for either task to complete (client disconnect or error)
+        done, pending = await asyncio.wait(
+            {redis_task, client_task},
+            return_when=asyncio.FIRST_COMPLETED
+        )
+
+        # Cancel pending tasks
+        for task in pending:
+            task.cancel()
+            try:
+                await task
+            except asyncio.CancelledError:
+                pass
+
+    except WebSocketDisconnect:
+        logger.info("status_websocket_client_disconnected", connection_id=connection_id)
+
+    except Exception as e:
+        logger.error(
+            "status_websocket_error",
+            error=str(e),
+            error_type=type(e).__name__,
+            connection_id=connection_id,
+        )
+        try:
+            await send_json(websocket, "error", {"error": str(e)})
+        except:
+            pass
+
+    finally:
+        # Cleanup
+        if pubsub:
+            try:
+                await pubsub.unsubscribe()
+                await pubsub.close()
+            except:
+                pass
+
+        if connection_id and organization_id:
+            await status_manager.disconnect(connection_id, organization_id)
+
+        logger.info("status_websocket_closed", connection_id=connection_id)
+
+
+@router.get("/ws/executions/status/stats")
+async def websocket_status_stats():
+    """
+    Get WebSocket status connection statistics.
+
+    Returns statistics about active connections and message throughput.
+    """
+    return status_manager.get_stats()