WebSocket Subscription Filtering for Project-Scoped Broadcasts

codeframe/ui/routers/websocket.py

@@ -33,7 +33,22 @@ async def websocket_endpoint(websocket: WebSocket):
 
     Message Types:
         - ping: Client heartbeat (responds with pong)
-        - subscribe: Subscribe to specific project updates
+        - subscribe: Subscribe to specific project updates (requires integer project_id)
+        - unsubscribe: Unsubscribe from specific project updates (requires integer project_id)
+
+    Message Format:
+        All messages must be valid JSON. Example:
+        - Ping: {"type": "ping"}
+        - Subscribe: {"type": "subscribe", "project_id": 1}
+        - Unsubscribe: {"type": "unsubscribe", "project_id": 1}
+
+    Error Handling:
+        Invalid messages receive error responses with type "error".
+        Examples of invalid messages:
+        - Missing project_id in subscribe/unsubscribe
+        - Non-integer project_id (e.g., string or float)
+        - Non-positive project_id (≤ 0)
+        - Malformed JSON
 
     Broadcasts:
         - agent_started: When an agent starts
@@ -68,8 +83,93 @@ async def websocket_endpoint(websocket: WebSocket):
             elif message.get("type") == "subscribe":
                 # Subscribe to specific project updates
                 project_id = message.get("project_id")
-                # TODO: Track subscriptions
-                await websocket.send_json({"type": "subscribed", "project_id": project_id})
+
+                # Validate project_id is present
+                if project_id is None:
+                    logger.warning("Subscribe message missing project_id")
+                    await websocket.send_json({
+                        "type": "error",
+                        "error": "Subscribe message requires project_id"
+                    })
+                    continue
+
+                # Validate project_id is an integer
+                if not isinstance(project_id, int):
+                    logger.warning(f"Invalid project_id type: {type(project_id).__name__}")
+                    await websocket.send_json({
+                        "type": "error",
+                        "error": f"project_id must be an integer, got {type(project_id).__name__}"
+                    })
+                    continue
+
+                # Validate project_id is positive
+                if project_id <= 0:
+                    logger.warning(f"Invalid project_id: {project_id}")
+                    await websocket.send_json({
+                        "type": "error",
+                        "error": "project_id must be a positive integer"
+                    })
+                    continue
+
+                # Track subscription
+                try:
+                    await manager.subscription_manager.subscribe(websocket, project_id)
+                    logger.info(f"WebSocket subscribed to project {project_id}")
+                    await websocket.send_json({
+                        "type": "subscribed",
+                        "project_id": project_id
+                    })
+                except Exception as e:
+                    logger.error(f"Error subscribing to project {project_id}: {e}")
+                    await websocket.send_json({
+                        "type": "error",
+                        "error": "Failed to subscribe to project"
+                    })
+            elif message.get("type") == "unsubscribe":
+                # Unsubscribe from specific project updates
+                project_id = message.get("project_id")
+
+                # Validate project_id is present
+                if project_id is None:
+                    logger.warning("Unsubscribe message missing project_id")
+                    await websocket.send_json({
+                        "type": "error",
+                        "error": "Unsubscribe message requires project_id"
+                    })
+                    continue
+
+                # Validate project_id is an integer
+                if not isinstance(project_id, int):
+                    logger.warning(f"Invalid project_id type: {type(project_id).__name__}")
+                    await websocket.send_json({
+                        "type": "error",
+                        "error": f"project_id must be an integer, got {type(project_id).__name__}"
+                    })
+                    continue
+
+                # Validate project_id is positive
+                if project_id <= 0:
+                    logger.warning(f"Invalid project_id: {project_id}")
+                    await websocket.send_json({
+                        "type": "error",
+                        "error": "project_id must be a positive integer"
+                    })
+                    continue
+
+                # Remove subscription
+                try:
+                    await manager.subscription_manager.unsubscribe(websocket, project_id)
+                    logger.info(f"WebSocket unsubscribed from project {project_id}")
+                    await websocket.send_json({
+                        "type": "unsubscribed",
+                        "project_id": project_id
+                    })
+                except Exception as e:
+                    logger.error(f"Error unsubscribing from project {project_id}: {e}")
+                    await websocket.send_json({
+                        "type": "error",
+                        "error": "Failed to unsubscribe from project"
+                    })
 
     except WebSocketDisconnect:
         # Normal client disconnect - no error logging needed
@@ -79,7 +179,7 @@ async def websocket_endpoint(websocket: WebSocket):
         logger.error(f"WebSocket error: {type(e).__name__} - {str(e)}", exc_info=True)
     finally:
         # Always disconnect and clean up, regardless of how we exited
-        manager.disconnect(websocket)
+        await manager.disconnect(websocket)
         try:
             await websocket.close()
         except Exception:

codeframe/ui/shared.py

@@ -4,39 +4,155 @@
 preventing circular import issues.
 """
 
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Set
 from fastapi import WebSocket
 import asyncio
 import time
+import logging
 
 from codeframe.core.models import ProjectStatus
 from codeframe.persistence.database import Database
 from codeframe.agents.lead_agent import LeadAgent
 
+logger = logging.getLogger(__name__)
+
+
+class WebSocketSubscriptionManager:
+    """Manage WebSocket subscriptions for project-filtered broadcasts.
+
+    This manager tracks which WebSocket connections are subscribed to which
+    projects, enabling filtered broadcasts that only send events to clients
+    subscribed to the relevant project.
+
+    Thread Safety:
+        All methods use asyncio.Lock for thread-safe operations, consistent
+        with the ConnectionManager pattern.
+
+    Data Structure:
+        subscriptions: Dict[WebSocket, Set[int]]
+            Maps each websocket to a set of project_ids it's subscribed to.
+            This allows a single client to subscribe to multiple projects.
+    """
+
+    def __init__(self):
+        self._subscriptions: Dict[WebSocket, Set[int]] = {}
+        self._subscriptions_lock = asyncio.Lock()
+
+    async def subscribe(self, websocket: WebSocket, project_id: int) -> None:
+        """Add a project subscription for a websocket.
+
+        Args:
+            websocket: WebSocket connection to subscribe
+            project_id: Project ID to subscribe to
+        """
+        async with self._subscriptions_lock:
+            if websocket not in self._subscriptions:
+                self._subscriptions[websocket] = set()
+
+            if project_id not in self._subscriptions[websocket]:
+                self._subscriptions[websocket].add(project_id)
+                logger.debug(f"WebSocket subscribed to project {project_id}")
+            else:
+                logger.debug(f"WebSocket already subscribed to project {project_id}")
+
+    async def unsubscribe(self, websocket: WebSocket, project_id: int) -> None:
+        """Remove a project subscription for a websocket.
+
+        Args:
+            websocket: WebSocket connection to unsubscribe
+            project_id: Project ID to unsubscribe from
+        """
+        async with self._subscriptions_lock:
+            if websocket in self._subscriptions:
+                self._subscriptions[websocket].discard(project_id)
+                logger.debug(f"WebSocket unsubscribed from project {project_id}")
+
+                # Clean up empty subscription sets
+                if not self._subscriptions[websocket]:
+                    del self._subscriptions[websocket]
+
+    async def get_subscribers(self, project_id: int) -> List[WebSocket]:
+        """Get list of websockets subscribed to a project.
+
+        Args:
+            project_id: Project ID to get subscribers for
+
+        Returns:
+            List of WebSocket connections subscribed to the project
+        """
+        async with self._subscriptions_lock:
+            subscribers = [
+                ws for ws, projects in self._subscriptions.items()
+                if project_id in projects
+            ]
+            return subscribers
+
+    async def cleanup(self, websocket: WebSocket) -> None:
+        """Remove all subscriptions for a websocket (called on disconnect).
+
+        Args:
+            websocket: WebSocket connection to clean up
+        """
+        async with self._subscriptions_lock:
+            if websocket in self._subscriptions:
+                project_count = len(self._subscriptions[websocket])
+                del self._subscriptions[websocket]
+                logger.debug(f"Cleaned up {project_count} subscriptions for disconnected WebSocket")
+
+    async def get_subscriptions(self, websocket: WebSocket) -> Set[int]:
+        """Get all project_ids a websocket is subscribed to.
+
+        Args:
+            websocket: WebSocket connection to check
+
+        Returns:
+            Set of project_ids the websocket is subscribed to (empty set if none)
+        """
+        async with self._subscriptions_lock:
+            return self._subscriptions.get(websocket, set()).copy()
+
 
 class ConnectionManager:
-    """Manage WebSocket connections for real-time updates."""
+    """Manage WebSocket connections for real-time updates with project-based filtering."""
 
     def __init__(self):
         self.active_connections: List[WebSocket] = []
         self._connections_lock = asyncio.Lock()
+        self.subscription_manager = WebSocketSubscriptionManager()
 
     async def connect(self, websocket: WebSocket):
         await websocket.accept()
         async with self._connections_lock:
             self.active_connections.append(websocket)
 
     async def disconnect(self, websocket: WebSocket):
+        # Clean up subscriptions first
+        await self.subscription_manager.cleanup(websocket)
+
+        # Then remove from active connections
         async with self._connections_lock:
             if websocket in self.active_connections:
                 self.active_connections.remove(websocket)
 
-    async def broadcast(self, message: dict):
-        """Broadcast message to all connected clients."""
-        # Get snapshot of connections to avoid holding lock during I/O
-        async with self._connections_lock:
-            connections = self.active_connections.copy()
-
+    async def broadcast(self, message: dict, project_id: Optional[int] = None):
+        """Broadcast message to connected clients.
+
+        Args:
+            message: Message dict to broadcast
+            project_id: Optional project ID for filtered broadcasts.
+                       If None, broadcasts to all connected clients (backward compatible).
+                       If provided, only broadcasts to clients subscribed to that project.
+        """
+        # Determine which connections should receive the message
+        if project_id is None:
+            # Backward compatible: broadcast to all connections
+            async with self._connections_lock:
+                connections = self.active_connections.copy()
+        else:
+            # Filtered broadcast: only to subscribers of this project
+            connections = await self.subscription_manager.get_subscribers(project_id)
+
+        # Send to target connections (no lock held during I/O)
         for connection in connections:
             try:
                 await connection.send_json(message)
@@ -168,7 +284,8 @@ async def start_agent(
                     "project_id": project_id,
                     "agent_type": "lead",
                     "timestamp": time.time(),  # Wall-clock timestamp (seconds since epoch)
-                }
+                },
+                project_id=project_id
             )
         except Exception:
             # Continue even if broadcast fails
@@ -177,7 +294,8 @@ async def start_agent(
         # cf-10.4: Broadcast status_update message
         try:
             await manager.broadcast(
-                {"type": "status_update", "project_id": project_id, "status": "running"}
+                {"type": "status_update", "project_id": project_id, "status": "running"},
+                project_id=project_id
             )
         except Exception:
             pass
@@ -202,7 +320,8 @@ async def start_agent(
                     "project_id": project_id,
                     "role": "assistant",
                     "content": greeting,
-                }
+                },
+                project_id=project_id
             )
         except Exception:
             pass