[P1] Implement WebSocket subscription tracking in websocket router

[frankbria]

Overview

The WebSocket router has a TODO for tracking client subscriptions to project events.

Current State

File: codeframe/ui/routers/websocket.py lines 70-72

project_id = message.get("project_id")
# TODO: Track subscriptions
await websocket.send_json({"type": "subscribed", "project_id": project_id})

Expected Behavior

• Maintain in-memory mapping of websocket → subscribed projects
• Support subscribing to multiple projects per connection
• Support unsubscribe messages
• Broadcast events only to subscribed clients
• Clean up subscriptions on disconnect

Implementation Requirements

1. Create WebSocketSubscriptionManager class
2. Store subscriptions in dict: {websocket_id: set(project_ids)}
3. Implement subscribe/unsubscribe message handlers
4. Filter broadcasts by subscription list
5. Clean up on connection close

Data Structure

class WebSocketSubscriptionManager:
    subscriptions: Dict[str, Set[int]] = {}  # websocket_id -> project_ids
    
    def subscribe(self, websocket_id: str, project_id: int)
    def unsubscribe(self, websocket_id: str, project_id: int)
    def get_subscribers(self, project_id: int) -> List[str]
    def cleanup(self, websocket_id: str)

Acceptance Criteria

• Clients can subscribe to specific projects
• Clients can unsubscribe from projects
• Events broadcast only to subscribed clients
• Subscriptions cleaned up on disconnect
• Multiple subscriptions per client supported

Priority

P1 - Critical: Required for real-time dashboard updates.

References

• Issues Analysis: CODEFRAME_ISSUES_ANALYSIS.md lines 68-73
• User Documentation: CODEFRAME_USER_DOCUMENTATION.md lines 37-39

[traycerai]

Plan

Observations

The WebSocket system currently broadcasts all events to all connected clients without filtering. The ConnectionManager in file:codeframe/ui/shared.py maintains a simple list of connections, and the WebSocket router at file:codeframe/ui/routers/websocket.py has a TODO at lines 70-72 for subscription tracking. All broadcast helper functions in file:codeframe/ui/websocket_broadcasts.py include project_id in messages, indicating the infrastructure is ready for project-based filtering. The frontend already sends subscribe messages and filters client-side, but the backend doesn't track or honor these subscriptions.

Approach

Implement a WebSocketSubscriptionManager class following the existing manager patterns in the codebase (similar to AgentPoolManager). Integrate it into the existing ConnectionManager to enable filtered broadcasts based on project subscriptions. Modify the WebSocket router to handle subscribe/unsubscribe messages and clean up subscriptions on disconnect. This approach maintains backward compatibility while adding subscription filtering, ensuring broadcasts only reach clients subscribed to specific projects.

Implementation Steps

1. Create WebSocketSubscriptionManager Class

Create a new class in file:codeframe/ui/shared.py (add after ConnectionManager class):

Class Structure:

• Use asyncio.Lock for thread-safe operations (consistent with ConnectionManager)
• Store subscriptions as Dict[WebSocket, Set[int]] mapping websocket → set of project_ids
• Implement methods: subscribe(), unsubscribe(), get_subscribers(), cleanup(), get_subscriptions()

Key Methods:

• async def subscribe(websocket: WebSocket, project_id: int): Add project_id to websocket's subscription set
• async def unsubscribe(websocket: WebSocket, project_id: int): Remove project_id from websocket's subscription set
• async def get_subscribers(project_id: int) -> List[WebSocket]: Return list of websockets subscribed to project_id
• async def cleanup(websocket: WebSocket): Remove all subscriptions for a websocket (called on disconnect)
• async def get_subscriptions(websocket: WebSocket) -> Set[int]: Get all project_ids a websocket is subscribed to

Implementation Details:

• Initialize with empty dict and lock in __init__
• Use async with self._subscriptions_lock for all operations
• Log subscription changes at DEBUG level
• Handle edge cases: websocket not in dict, project_id already subscribed, etc.

2. Integrate Subscription Manager into ConnectionManager

Modify ConnectionManager class in file:codeframe/ui/shared.py:

Add Subscription Manager:

• Initialize self.subscription_manager = WebSocketSubscriptionManager() in __init__
• Keep existing active_connections list for backward compatibility

Modify broadcast() Method:

• Add optional parameter: project_id: Optional[int] = None
• If project_id is None, broadcast to all connections (backward compatible)
• If project_id is provided, get subscribers via subscription_manager.get_subscribers(project_id)
• Only send to subscribed connections
• Keep existing error handling and disconnect logic

Update disconnect() Method:

• Call await self.subscription_manager.cleanup(websocket) before removing from active_connections
• Ensures subscriptions are cleaned up when client disconnects

3. Update WebSocket Router Message Handlers

Modify file:codeframe/ui/routers/websocket.py:

Handle Subscribe Messages (lines 68-72):

• Replace TODO with actual subscription tracking
• Extract project_id from message
• Validate project_id is an integer
• Call await manager.subscription_manager.subscribe(websocket, project_id)
• Send confirmation: {"type": "subscribed", "project_id": project_id}
• Add error handling for invalid project_id

Add Unsubscribe Message Handler:

• Add new elif block after subscribe handler
• Check for message.get("type") == "unsubscribe"
• Extract and validate project_id
• Call await manager.subscription_manager.unsubscribe(websocket, project_id)
• Send confirmation: {"type": "unsubscribed", "project_id": project_id}

Update Disconnect Cleanup:

• The finally block already calls manager.disconnect(websocket) at line 82
• No changes needed since disconnect() will call subscription cleanup

4. Update Broadcast Helper Functions

Modify all broadcast functions in file:codeframe/ui/websocket_broadcasts.py:

Pattern for All Functions:

• Each function already receives project_id parameter
• Change await manager.broadcast(message) to await manager.broadcast(message, project_id=project_id)
• This enables filtered broadcasting based on subscriptions

Functions to Update:

• broadcast_task_status() - line 70
• broadcast_agent_status() - line 115
• broadcast_test_result() - line 160
• broadcast_commit_created() - line 201
• broadcast_activity_update() - line 241
• broadcast_progress_update() - line 279
• broadcast_correction_attempt() - line 320
• broadcast_agent_created() - line 358
• broadcast_agent_retired() - line 385
• broadcast_task_assigned() - line 416
• broadcast_task_blocked() - line 448
• broadcast_task_unblocked() - line 487
• broadcast_blocker_created() - line 542
• broadcast_blocker_resolved() - line 569
• broadcast_agent_resumed() - line 598
• broadcast_blocker_expired() - line 629
• broadcast_discovery_answer_submitted() - line 682
• broadcast_discovery_question_presented() - line 718
• broadcast_discovery_progress_updated() - line 753
• broadcast_discovery_completed() - line 783

5. Update start_agent Function

Modify start_agent() function in file:codeframe/ui/shared.py:

Update Broadcast Calls:

• Line 165-172: agent_started broadcast - add project_id=project_id parameter
• Line 179-181: status_update broadcast - add project_id=project_id parameter
• Line 199-207: chat_message broadcast - add project_id=project_id parameter

This ensures agent startup events are only sent to subscribed clients.

Data Flow Diagram

sequenceDiagram
    participant Client as WebSocket Client
    participant Router as WebSocket Router
    participant Manager as ConnectionManager
    participant SubMgr as SubscriptionManager
    participant Broadcast as Broadcast Functions
    
    Client->>Router: Connect (WS /ws)
    Router->>Manager: connect(websocket)
    Manager->>Manager: Add to active_connections
    
    Client->>Router: {type: "subscribe", project_id: 1}
    Router->>SubMgr: subscribe(websocket, 1)
    SubMgr->>SubMgr: Add to subscriptions[ws] = {1}
    Router->>Client: {type: "subscribed", project_id: 1}
    
    Note over Broadcast: Event occurs for project 1
    Broadcast->>Manager: broadcast(message, project_id=1)
    Manager->>SubMgr: get_subscribers(1)
    SubMgr-->>Manager: [websocket]
    Manager->>Client: Send message (filtered)
    
    Client->>Router: {type: "unsubscribe", project_id: 1}
    Router->>SubMgr: unsubscribe(websocket, 1)
    SubMgr->>SubMgr: Remove from subscriptions[ws]
    Router->>Client: {type: "unsubscribed", project_id: 1}
    
    Client->>Router: Disconnect
    Router->>Manager: disconnect(websocket)
    Manager->>SubMgr: cleanup(websocket)
    SubMgr->>SubMgr: Remove all subscriptions
    Manager->>Manager: Remove from active_connections

Loading

Subscription Data Structure

Component	Data Structure	Purpose
WebSocketSubscriptionManager	Dict[WebSocket, Set[int]]	Maps each websocket to set of subscribed project_ids
ConnectionManager	List[WebSocket]	Maintains list of all active connections (unchanged)
Message Format	{"type": "subscribe", "project_id": int}	Client subscription request
Response Format	{"type": "subscribed", "project_id": int}	Server confirmation

Key Implementation Notes

• Thread Safety: Use asyncio.Lock consistently with existing ConnectionManager pattern
• Backward Compatibility: broadcast() without project_id still broadcasts to all clients
• Multiple Subscriptions: Each websocket can subscribe to multiple projects (stored as Set)
• Automatic Cleanup: Subscriptions automatically cleaned up on disconnect via finally block
• Error Handling: Validate project_id is integer, handle malformed messages gracefully
• Logging: Add DEBUG-level logs for subscription operations to aid debugging
• No Database Changes: All subscription tracking is in-memory (matches ConnectionManager pattern)

Import In IDE

🤖 Prompt for AI Agents

## Observations

The WebSocket system currently broadcasts all events to all connected clients without filtering. The `ConnectionManager` in `file:codeframe/ui/shared.py` maintains a simple list of connections, and the WebSocket router at `file:codeframe/ui/routers/websocket.py` has a TODO at lines 70-72 for subscription tracking. All broadcast helper functions in `file:codeframe/ui/websocket_broadcasts.py` include `project_id` in messages, indicating the infrastructure is ready for project-based filtering. The frontend already sends subscribe messages and filters client-side, but the backend doesn't track or honor these subscriptions.

## Approach

Implement a `WebSocketSubscriptionManager` class following the existing manager patterns in the codebase (similar to `AgentPoolManager`). Integrate it into the existing `ConnectionManager` to enable filtered broadcasts based on project subscriptions. Modify the WebSocket router to handle subscribe/unsubscribe messages and clean up subscriptions on disconnect. This approach maintains backward compatibility while adding subscription filtering, ensuring broadcasts only reach clients subscribed to specific projects.

## Implementation Steps

### 1. Create WebSocketSubscriptionManager Class

Create a new class in `file:codeframe/ui/shared.py` (add after `ConnectionManager` class):

**Class Structure:**
- Use `asyncio.Lock` for thread-safe operations (consistent with `ConnectionManager`)
- Store subscriptions as `Dict[WebSocket, Set[int]]` mapping websocket → set of project_ids
- Implement methods: `subscribe()`, `unsubscribe()`, `get_subscribers()`, `cleanup()`, `get_subscriptions()`

**Key Methods:**
- `async def subscribe(websocket: WebSocket, project_id: int)`: Add project_id to websocket's subscription set
- `async def unsubscribe(websocket: WebSocket, project_id: int)`: Remove project_id from websocket's subscription set
- `async def get_subscribers(project_id: int) -> List[WebSocket]`: Return list of websockets subscribed to project_id
- `async def cleanup(websocket: WebSocket)`: Remove all subscriptions for a websocket (called on disconnect)
- `async def get_subscriptions(websocket: WebSocket) -> Set[int]`: Get all project_ids a websocket is subscribed to

**Implementation Details:**
- Initialize with empty dict and lock in `__init__`
- Use `async with self._subscriptions_lock` for all operations
- Log subscription changes at DEBUG level
- Handle edge cases: websocket not in dict, project_id already subscribed, etc.

### 2. Integrate Subscription Manager into ConnectionManager

Modify `ConnectionManager` class in `file:codeframe/ui/shared.py`:

**Add Subscription Manager:**
- Initialize `self.subscription_manager = WebSocketSubscriptionManager()` in `__init__`
- Keep existing `active_connections` list for backward compatibility

**Modify broadcast() Method:**
- Add optional parameter: `project_id: Optional[int] = None`
- If `project_id` is None, broadcast to all connections (backward compatible)
- If `project_id` is provided, get subscribers via `subscription_manager.get_subscribers(project_id)`
- Only send to subscribed connections
- Keep existing error handling and disconnect logic

**Update disconnect() Method:**
- Call `await self.subscription_manager.cleanup(websocket)` before removing from active_connections
- Ensures subscriptions are cleaned up when client disconnects

### 3. Update WebSocket Router Message Handlers

Modify `file:codeframe/ui/routers/websocket.py`:

**Handle Subscribe Messages (lines 68-72):**
- Replace TODO with actual subscription tracking
- Extract `project_id` from message
- Validate `project_id` is an integer
- Call `await manager.subscription_manager.subscribe(websocket, project_id)`
- Send confirmation: `{"type": "subscribed", "project_id": project_id}`
- Add error handling for invalid project_id

**Add Unsubscribe Message Handler:**
- Add new `elif` block after subscribe handler
- Check for `message.get("type") == "unsubscribe"`
- Extract and validate `project_id`
- Call `await manager.subscription_manager.unsubscribe(websocket, project_id)`
- Send confirmation: `{"type": "unsubscribed", "project_id": project_id}`

**Update Disconnect Cleanup:**
- The `finally` block already calls `manager.disconnect(websocket)` at line 82
- No changes needed since disconnect() will call subscription cleanup

### 4. Update Broadcast Helper Functions

Modify all broadcast functions in `file:codeframe/ui/websocket_broadcasts.py`:

**Pattern for All Functions:**
- Each function already receives `project_id` parameter
- Change `await manager.broadcast(message)` to `await manager.broadcast(message, project_id=project_id)`
- This enables filtered broadcasting based on subscriptions

**Functions to Update:**
- `broadcast_task_status()` - line 70
- `broadcast_agent_status()` - line 115
- `broadcast_test_result()` - line 160
- `broadcast_commit_created()` - line 201
- `broadcast_activity_update()` - line 241
- `broadcast_progress_update()` - line 279
- `broadcast_correction_attempt()` - line 320
- `broadcast_agent_created()` - line 358
- `broadcast_agent_retired()` - line 385
- `broadcast_task_assigned()` - line 416
- `broadcast_task_blocked()` - line 448
- `broadcast_task_unblocked()` - line 487
- `broadcast_blocker_created()` - line 542
- `broadcast_blocker_resolved()` - line 569
- `broadcast_agent_resumed()` - line 598
- `broadcast_blocker_expired()` - line 629
- `broadcast_discovery_answer_submitted()` - line 682
- `broadcast_discovery_question_presented()` - line 718
- `broadcast_discovery_progress_updated()` - line 753
- `broadcast_discovery_completed()` - line 783

### 5. Update start_agent Function

Modify `start_agent()` function in `file:codeframe/ui/shared.py`:

**Update Broadcast Calls:**
- Line 165-172: `agent_started` broadcast - add `project_id=project_id` parameter
- Line 179-181: `status_update` broadcast - add `project_id=project_id` parameter
- Line 199-207: `chat_message` broadcast - add `project_id=project_id` parameter

This ensures agent startup events are only sent to subscribed clients.

## Data Flow Diagram

```mermaid
sequenceDiagram
    participant Client as WebSocket Client
    participant Router as WebSocket Router
    participant Manager as ConnectionManager
    participant SubMgr as SubscriptionManager
    participant Broadcast as Broadcast Functions
    
    Client->>Router: Connect (WS /ws)
    Router->>Manager: connect(websocket)
    Manager->>Manager: Add to active_connections
    
    Client->>Router: {type: "subscribe", project_id: 1}
    Router->>SubMgr: subscribe(websocket, 1)
    SubMgr->>SubMgr: Add to subscriptions[ws] = {1}
    Router->>Client: {type: "subscribed", project_id: 1}
    
    Note over Broadcast: Event occurs for project 1
    Broadcast->>Manager: broadcast(message, project_id=1)
    Manager->>SubMgr: get_subscribers(1)
    SubMgr-->>Manager: [websocket]
    Manager->>Client: Send message (filtered)
    
    Client->>Router: {type: "unsubscribe", project_id: 1}
    Router->>SubMgr: unsubscribe(websocket, 1)
    SubMgr->>SubMgr: Remove from subscriptions[ws]
    Router->>Client: {type: "unsubscribed", project_id: 1}
    
    Client->>Router: Disconnect
    Router->>Manager: disconnect(websocket)
    Manager->>SubMgr: cleanup(websocket)
    SubMgr->>SubMgr: Remove all subscriptions
    Manager->>Manager: Remove from active_connections
```

## Subscription Data Structure

| Component | Data Structure | Purpose |
|-----------|---------------|---------|
| `WebSocketSubscriptionManager` | `Dict[WebSocket, Set[int]]` | Maps each websocket to set of subscribed project_ids |
| `ConnectionManager` | `List[WebSocket]` | Maintains list of all active connections (unchanged) |
| Message Format | `{"type": "subscribe", "project_id": int}` | Client subscription request |
| Response Format | `{"type": "subscribed", "project_id": int}` | Server confirmation |

## Key Implementation Notes

- **Thread Safety**: Use `asyncio.Lock` consistently with existing `ConnectionManager` pattern
- **Backward Compatibility**: `broadcast()` without `project_id` still broadcasts to all clients
- **Multiple Subscriptions**: Each websocket can subscribe to multiple projects (stored as Set)
- **Automatic Cleanup**: Subscriptions automatically cleaned up on disconnect via `finally` block
- **Error Handling**: Validate project_id is integer, handle malformed messages gracefully
- **Logging**: Add DEBUG-level logs for subscription operations to aid debugging
- **No Database Changes**: All subscription tracking is in-memory (matches ConnectionManager pattern)

---

Execution Information

Branch: main
Commit: 1a87e4e

---

💡 Tips

Supported Commands (Inside Comments)

• Use @traycerai generate to iterate on the previous version of the implementation plan.

Supported Commands (Inside Description)

• Add @traycerai ignore anywhere in the ticket description to prevent this ticket from being processed.
• Add @traycerai branch:<branch-name> anywhere in the ticket description to specify the target branch for the implementation plan.

Community

• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[coderabbitai]

📝 CodeRabbit Plan Mode

Generate an implementation plan and prompts that you can use with your favorite coding agent.

• Create Plan

Examples

• Example 1
• Example 2

_{You can configure auto-planning by selecting labels in the issue_enrichment configuration.}

🔗 Similar Issues

Related Issues

• Fix WebSocket real-time updates - Connection refused errors #87
• [P1] Implement Project.pause() with flash save mechanism #103
• [P1] Replace mock data with database queries in GET /api/projects/{id}/tasks #104

👤 Suggested Assignees

• frankbria

---

🧪 Issue enrichment is currently in open beta.

To disable automatic issue enrichment, add the following to your .coderabbit.yaml:

issue_enrichment:
  auto_enrich:
    enabled: false

💬 Have feedback or questions? Drop into our discord or schedule a call!

[frankbria]

Closed by PR #134