feat(websocket): implement proactive messaging system#224
Conversation
Transform WebSocket from passive (only responding to client messages) to proactive (actively sending connection health and state information). New proactive messages: - connection_ack: Sent immediately after subscription to confirm real-time connectivity - project_status: Initial project state snapshot sent on subscription - heartbeat: Periodic messages every 30 seconds to keep connections alive Implementation: - Add HEARTBEAT_INTERVAL_SECONDS constant (30s default) - Add send_heartbeats() async function for background heartbeat tasks - Track heartbeat tasks per project, cancel on disconnect - Add broadcast_heartbeat() helper function - Update WebSocketMessageType union with new message types Test updates: - Add comprehensive test suite (test_websocket_proactive.py) with 17 tests - Add drain_proactive_messages() helper for integration tests - Update existing integration tests to handle new message sequence - Update E2E test expectations for proactive messaging This ensures clients always receive messages to confirm the real-time system is operational, fixing the issue where E2E tests received 0 messages from the WebSocket.
WalkthroughSends proactive WebSocket messages: Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Router as WebSocket Router
participant Manager as Connection Manager
participant DB as Database
rect rgb(200,220,255)
Note over Client,Router: Subscribe → register & proactive messages
Client->>Router: subscribe(project_id)
Router->>Manager: register_subscription(project_id, client)
Manager-->>Router: subscribed
Router-->>Client: {"type":"subscribed", ...}
Router-->>Client: {"type":"connection_ack","project_id":project_id,"timestamp":...}
Router->>DB: fetch project_status(project_id)
DB-->>Router: project data / not found
Router-->>Client: {"type":"project_status", ...}
end
rect rgb(220,240,220)
Note over Router,Manager: Per-subscription heartbeat loop (every HEARTBEAT_INTERVAL_SECONDS)
loop every HEARTBEAT_INTERVAL_SECONDS
Router->>Manager: broadcast_heartbeat(project_id)
Manager-->>Client: {"type":"heartbeat","project_id":project_id,"timestamp":...}
end
end
rect rgb(240,220,220)
Note over Client,Router: Unsubscribe / Disconnect → cleanup
Client->>Router: unsubscribe(project_id) / disconnect
Router->>Router: cancel heartbeat task(s)
Router->>Manager: remove_subscription(project_id, client)
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Comment |
Code Review: Proactive WebSocket Messaging SystemThis is an excellent enhancement that transforms the WebSocket from passive to proactive messaging. The implementation is well-structured, thoroughly tested, and addresses a real usability issue (E2E tests receiving 0 messages). Here's my detailed feedback: ✅ Strengths1. Excellent Design & Architecture
2. Comprehensive Test Coverage
3. Type Safety
4. Backward Compatibility
🔍 Issues & Recommendations1. Memory Leak Risk: Heartbeat Task Cleanup
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In @codeframe/ui/routers/websocket.py:
- Around line 288-294: When handling an unsubscribe for a project, cancel and
cleanup the associated heartbeat task stored in heartbeat_tasks so heartbeats
stop after unsubscribe: look up heartbeat_tasks[project_id], call task.cancel(),
await it (e.g., with asyncio.wait_for or try/except catching
asyncio.CancelledError/asyncio.TimeoutError) to suppress warnings, and then del
heartbeat_tasks[project_id]; do this in the unsubscribe handling code path that
currently removes subscriptions (the same block referencing project_id and
heartbeat_tasks), ensuring you only cancel if the task exists and is not done.
In @tests/ui/test_websocket_proactive.py:
- Around line 270-279: Replace the silent conditional with an assertion so the
test fails if no heartbeat was sent: assert that the computed heartbeat_calls
(from mock_websocket.send_json.call_args_list filtered by call[0][0].get("type")
== "heartbeat") has length > 0, then extract heartbeat =
heartbeat_calls[0][0][0] and assert the fields ("type" == "heartbeat",
"project_id" == 1, and "timestamp" in heartbeat) as before; this ensures
test_heartbeat_contains_required_fields fails when no heartbeat is present.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
codeframe/ui/routers/websocket.pycodeframe/ui/websocket_broadcasts.pytests/e2e/test_dashboard.spec.tstests/ui/test_websocket_integration.pytests/ui/test_websocket_proactive.pyweb-ui/src/types/index.ts
🧰 Additional context used
📓 Path-based instructions (3)
web-ui/src/**/*.{ts,tsx}
📄 CodeRabbit inference engine (CLAUDE.md)
web-ui/src/**/*.{ts,tsx}: Use TypeScript 5.3+ for frontend development with React 18, Next.js 14, Tailwind CSS, and Hugeicons for icons
Use TypeScript strict mode and ensure 100% type safety with no hardcoded types or any usage in frontend code
Files:
web-ui/src/types/index.ts
codeframe/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
codeframe/**/*.py: Use Python 3.11+ for backend development with FastAPI, AsyncAnthropic, SQLite with async support (aiosqlite), and tiktoken for token counting
Use token counting via tiktoken library for token budget management with ~50,000 token limit per conversation
Use asyncio patterns with AsyncAnthropic for async/await in Python backend for concurrent operations
Implement quality gates with multi-stage pre-completion checks (tests → type → coverage → review) and Git + SQLite + context snapshots for project state rollback
Use tiered memory system (HOT/WARM/COLD) with importance scoring using hybrid exponential decay algorithm for context management with 30-50% token reduction
Implement session lifecycle management with auto-save/restore using file-based storage at .codeframe/session_state.json
Files:
codeframe/ui/websocket_broadcasts.pycodeframe/ui/routers/websocket.py
tests/e2e/**/*.ts
📄 CodeRabbit inference engine (CLAUDE.md)
Implement E2E tests using Playwright + TestSprite with loginUser() helper from tests/e2e/test-utils.ts for authentication
Files:
tests/e2e/test_dashboard.spec.ts
🧠 Learnings (4)
📓 Common learnings
Learnt from: CR
Repo: frankbria/codeframe PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-01-04T06:26:12.870Z
Learning: Applies to web-ui/src/lib/websocketMessageMapper.ts : Implement WebSocket event mapping with 9 event types (agent_created, task_assigned, etc.) with automatic reconnection using exponential backoff (1s → 30s)
Learnt from: CR
Repo: frankbria/codeframe PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-01-04T06:26:12.870Z
Learning: Applies to web-ui/src/lib/websocket.ts : Implement WebSocket connections with authentication token passed as query parameter (?token=TOKEN)
📚 Learning: 2026-01-04T06:26:12.870Z
Learnt from: CR
Repo: frankbria/codeframe PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-01-04T06:26:12.870Z
Learning: Applies to web-ui/src/lib/websocketMessageMapper.ts : Implement WebSocket event mapping with 9 event types (agent_created, task_assigned, etc.) with automatic reconnection using exponential backoff (1s → 30s)
Applied to files:
web-ui/src/types/index.tscodeframe/ui/routers/websocket.pytests/e2e/test_dashboard.spec.ts
📚 Learning: 2026-01-04T06:26:12.870Z
Learnt from: CR
Repo: frankbria/codeframe PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-01-04T06:26:12.870Z
Learning: Applies to web-ui/src/lib/websocket.ts : Implement WebSocket connections with authentication token passed as query parameter (?token=TOKEN)
Applied to files:
web-ui/src/types/index.ts
📚 Learning: 2025-11-25T19:08:37.203Z
Learnt from: CR
Repo: frankbria/codeframe PR: 0
File: docs/CLAUDE.md:0-0
Timestamp: 2025-11-25T19:08:37.203Z
Learning: Applies to docs/codeframe/ui/**/*.py : Use FastAPI with Uvicorn for the async API backend and WebSockets for real-time communication
Applied to files:
codeframe/ui/routers/websocket.py
🧬 Code graph analysis (2)
codeframe/ui/websocket_broadcasts.py (3)
tests/workspace/test_workspace_manager_comprehensive.py (1)
manager(39-41)codeframe/core/models.py (1)
project_id(234-235)codeframe/ui/shared.py (1)
broadcast(154-185)
codeframe/ui/routers/websocket.py (3)
codeframe/core/models.py (2)
project_id(234-235)Task(257-305)codeframe/ui/routers/projects.py (1)
get_project(210-247)codeframe/persistence/database.py (2)
get_project(278-280)create_task(358-360)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: Frontend Unit Tests
- GitHub Check: Backend Unit Tests
- GitHub Check: claude-review
- GitHub Check: E2E Smoke Tests (Chromium)
🔇 Additional comments (13)
web-ui/src/types/index.ts (1)
103-105: LGTM! New proactive message types correctly added.The three new WebSocketMessageType union members align with the backend proactive messaging implementation. The comments clearly explain the purpose of each type.
Consider adding optional fields to
WebSocketMessageinterface for type safety on these new message types (e.g.,message?: stringforconnection_ack,status?: stringandphase?: stringforproject_status), though the existingdata?: anyfield provides flexibility.codeframe/ui/websocket_broadcasts.py (1)
39-60: LGTM! Clean implementation following established patterns.The
broadcast_heartbeatfunction correctly follows the same structure as other broadcast helpers in this module:
- Consistent message format with
type,project_id, andtimestamp- Proper error handling that logs but doesn't raise
- Uses the established UTC timestamp format
tests/e2e/test_dashboard.spec.ts (1)
443-475: LGTM! E2E test correctly updated for proactive messaging.The test expectations are properly aligned with the new backend behavior:
- Comments and error messages clearly explain the proactive message flow
- Message verification includes all three new types (
connection_ack,project_status,heartbeat)- The guard path with warning at lines 471-474 provides good debugging info
tests/ui/test_websocket_integration.py (2)
38-78: Well-designed helper for draining proactive messages.The
drain_proactive_messageshelper cleanly abstracts the new message flow:
- Correctly handles up to 3 proactive messages with appropriate timeout
- Returns message types for flexible assertions
- Early break on unexpected message types prevents test hangs
One minor note: The 0.5s timeout per message (line 63) should be sufficient for local integration tests but may need tuning for CI environments with higher latency.
102-105: Consistent updates across test methods.The test updates correctly integrate the
drain_proactive_messageshelper while maintaining the original test intent. The assertions on"subscribed" in drainedverify the subscription succeeded before testing broadcast behavior.Also applies to: 126-128
codeframe/ui/routers/websocket.py (3)
47-77: LGTM! Clean heartbeat implementation with proper cancellation handling.The
send_heartbeatsfunction correctly:
- Runs an infinite loop with the configured interval
- Catches send exceptions and exits gracefully
- Re-raises
CancelledErrorfor proper task cancellation semantics
264-286: LGTM! Proactive messaging sequence implemented correctly.The subscription flow properly sends messages in the documented order:
subscribed- confirmationconnection_ack- proactive acknowledgmentproject_status- initial state snapshot (with graceful error handling)Good defensive coding: the
project_statusfetch is wrapped in try/except so database errors don't break the subscription.
350-358: LGTM! Proper cleanup of heartbeat tasks on disconnect.The finally block correctly:
- Iterates over all heartbeat tasks and cancels them
- Awaits each task to ensure clean cancellation
- Suppresses expected
CancelledError- Clears the task dictionary
tests/ui/test_websocket_proactive.py (5)
25-35: LGTM! Well-structured mock fixtures.The mock fixtures properly set up:
- WebSocket with async methods and query params for auth
- ConnectionManager with subscription handling
- Database with project access and retrieval
92-102: Clean auto-applied authentication patches.The
autouse=Truefixture elegantly patches JWT authentication for all tests in this module, reducing boilerplate while keeping tests focused on proactive messaging behavior.
108-129: LGTM! Comprehensive connection acknowledgment tests.The tests thoroughly cover:
- Message presence and content validation
- Correct ordering relative to
subscribed- Timestamp format validation (ISO8601)
- Access denied scenario (no ack sent)
Also applies to: 131-153, 176-194
311-332: LGTM! Robust initial state snapshot tests.Good coverage including:
- Normal project_status delivery
- Data verification from mock database
- Graceful handling of
Noneproject (line 364)- Graceful handling of DB exceptions (line 385)
The tests correctly verify that subscription succeeds even when project_status cannot be sent.
Also applies to: 361-401
461-493: LGTM! Broadcast helper tests validate integration point.These tests ensure the
broadcast_heartbeathelper:
- Is exported and callable
- Sends correct message format
- Handles broadcast exceptions gracefully
Address CodeRabbit review feedback: 1. Cancel heartbeat task when client unsubscribes from a project - Previously heartbeats continued until disconnect - Now properly cancelled and cleaned up on unsubscribe 2. Fix test assertion in test_heartbeat_contains_required_fields - Changed conditional `if len(heartbeat_calls) > 0` to assertion - Test now fails properly if no heartbeat is received 3. Add new test: test_heartbeat_task_cancelled_on_unsubscribe - Verifies heartbeat task cleanup on unsubscribe works correctly
Updated Code Review: Proactive WebSocket Messaging SystemThank you for addressing issue #2 from my previous review! I can see the second commit (fc8cdb1) now properly cancels heartbeat tasks on unsubscribe. Excellent response time! ✅ What's Been FixedIssue #2: Unsubscribe Now Cancels Heartbeat ✅Location: The unsubscribe handler now properly cancels the heartbeat task: # Cancel heartbeat task for this project
if project_id in heartbeat_tasks:
heartbeat_tasks[project_id].cancel()
try:
await heartbeat_tasks[project_id]
except asyncio.CancelledError:
pass
del heartbeat_tasks[project_id]
logger.debug(f"Cancelled heartbeat task for unsubscribed project {project_id}")This prevents unnecessary network traffic and CPU usage after unsubscription. Great fix!
|
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
codeframe/ui/routers/websocket.py (1)
264-294: Consider extracting a timestamp helper for consistency.The ISO8601 timestamp formatting pattern
datetime.now(UTC).isoformat().replace("+00:00", "Z")is repeated at lines 66, 269, and 282. A small helper would reduce duplication and ensure consistent formatting.♻️ Optional: Extract timestamp helper
+def utc_timestamp_iso() -> str: + """Return current UTC time as ISO8601 string with Z suffix.""" + return datetime.now(UTC).isoformat().replace("+00:00", "Z") + async def send_heartbeats(websocket: WebSocket, project_id: int) -> None: ... message = { "type": "heartbeat", "project_id": project_id, - "timestamp": datetime.now(UTC).isoformat().replace("+00:00", "Z"), + "timestamp": utc_timestamp_iso(), }Apply similarly to lines 269 and 282.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
codeframe/ui/routers/websocket.pytests/ui/test_websocket_proactive.py
🧰 Additional context used
📓 Path-based instructions (1)
codeframe/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
codeframe/**/*.py: Use Python 3.11+ for backend development with FastAPI, AsyncAnthropic, SQLite with async support (aiosqlite), and tiktoken for token counting
Use token counting via tiktoken library for token budget management with ~50,000 token limit per conversation
Use asyncio patterns with AsyncAnthropic for async/await in Python backend for concurrent operations
Implement quality gates with multi-stage pre-completion checks (tests → type → coverage → review) and Git + SQLite + context snapshots for project state rollback
Use tiered memory system (HOT/WARM/COLD) with importance scoring using hybrid exponential decay algorithm for context management with 30-50% token reduction
Implement session lifecycle management with auto-save/restore using file-based storage at .codeframe/session_state.json
Files:
codeframe/ui/routers/websocket.py
🧠 Learnings (4)
📓 Common learnings
Learnt from: CR
Repo: frankbria/codeframe PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-01-04T06:26:12.870Z
Learning: Applies to web-ui/src/lib/websocketMessageMapper.ts : Implement WebSocket event mapping with 9 event types (agent_created, task_assigned, etc.) with automatic reconnection using exponential backoff (1s → 30s)
📚 Learning: 2025-11-25T19:08:37.203Z
Learnt from: CR
Repo: frankbria/codeframe PR: 0
File: docs/CLAUDE.md:0-0
Timestamp: 2025-11-25T19:08:37.203Z
Learning: Applies to docs/codeframe/ui/**/*.py : Use FastAPI with Uvicorn for the async API backend and WebSockets for real-time communication
Applied to files:
codeframe/ui/routers/websocket.py
📚 Learning: 2026-01-04T06:26:12.870Z
Learnt from: CR
Repo: frankbria/codeframe PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-01-04T06:26:12.870Z
Learning: Applies to web-ui/src/lib/websocketMessageMapper.ts : Implement WebSocket event mapping with 9 event types (agent_created, task_assigned, etc.) with automatic reconnection using exponential backoff (1s → 30s)
Applied to files:
codeframe/ui/routers/websocket.py
📚 Learning: 2025-12-17T19:21:40.014Z
Learnt from: frankbria
Repo: frankbria/codeframe PR: 128
File: tests/agents/test_bottleneck_detection.py:486-500
Timestamp: 2025-12-17T19:21:40.014Z
Learning: In tests/agents/test_bottleneck_detection.py, test_skip_agents_below_threshold should be async and mock _get_agent_workload to return workload below AGENT_OVERLOAD_THRESHOLD (5) while providing non-empty tasks list to prevent early return in detect_bottlenecks().
Applied to files:
tests/ui/test_websocket_proactive.py
🧬 Code graph analysis (1)
codeframe/ui/routers/websocket.py (2)
codeframe/ui/routers/projects.py (1)
get_project(210-247)codeframe/persistence/database.py (2)
get_project(278-280)create_task(358-360)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: Frontend Unit Tests
- GitHub Check: Backend Unit Tests
- GitHub Check: E2E Smoke Tests (Chromium)
- GitHub Check: claude-review
🔇 Additional comments (9)
codeframe/ui/routers/websocket.py (3)
47-77: LGTM!The heartbeat function correctly handles both normal send failures (breaks loop) and task cancellation (re-raises
CancelledError). The sleep-before-send pattern ensures no immediate heartbeat on subscription, which is appropriate sinceconnection_ackalready confirms connectivity.
337-346: Heartbeat cancellation on unsubscribe correctly implemented.The cancellation logic properly awaits the cancelled task to suppress warnings, handles
CancelledError, and cleans up the dictionary entry. This addresses the prior review feedback.
361-369: LGTM!The cleanup logic in the
finallyblock correctly cancels all heartbeat tasks and awaits their completion. The pattern of iterating first and then clearing is safe, and the exception handling aligns withsend_heartbeats's behavior.tests/ui/test_websocket_proactive.py (6)
25-102: LGTM!The fixtures are well-designed with clear separation of concerns. The
apply_jwt_auth_patchesautouse fixture cleanly handles authentication mocking for all tests, reducing boilerplate.
105-194: LGTM!The connection acknowledgment tests thoroughly validate the message content, ordering relative to
subscribed, timestamp format, and the absence of acknowledgment on failed access. Good coverage of edge cases.
270-279: Past review concern addressed.The assertion at line 275 (
assert len(heartbeat_calls) > 0) correctly replaces the previous silent conditional, ensuring the test fails if no heartbeat is received.
342-435: LGTM!The initial state snapshot tests properly validate both happy path data retrieval and graceful handling of missing projects or database errors—ensuring subscription reliability is independent of the optional status snapshot.
438-489: LGTM!The sequence tests ensure the contract
subscribed → connection_ack → project_statusis maintained, and verify that multiple subscriptions each receive their own set of proactive messages with correct project IDs.
492-527: LGTM!The broadcast helper tests validate existence, message format, and graceful exception handling. The exception suppression test at lines 519-527 appropriately ensures the helper doesn't propagate broadcast failures.
Address 5 recommendations from claude bot code review: 1. Fix memory leak on re-subscription - Cancel existing heartbeat task before starting new one - Prevents orphaned tasks when client re-subscribes to same project 2. Cancel heartbeat on unsubscribe (already done in previous commit) - Verified heartbeat task is cancelled when client unsubscribes 3. Make heartbeat interval configurable - Added WEBSOCKET_HEARTBEAT_INTERVAL env var (default: 30 seconds) - Allows tuning for different environments and faster testing 4. Improve project status error handling specificity - Distinguish between project not found, missing fields, and unexpected errors - Use appropriate log levels (debug vs warning) 5. Add test for unsubscribe cancellation (already done in previous commit) - test_heartbeat_task_cancelled_on_unsubscribe verifies cleanup
✅ Final Review: LGTM - Ready to Merge!Excellent work! All feedback from my previous reviews has been thoroughly addressed. This PR is now in excellent shape. 🎉 What's Been Fixed (Latest Commit)The third commit (8686581) addresses all remaining issues from my previous reviews: 1. ✅ Heartbeat Interval Now ConfigurableLocation: HEARTBEAT_INTERVAL_SECONDS = int(os.getenv("WEBSOCKET_HEARTBEAT_INTERVAL", "30"))Impact:
2. ✅ Improved Error Handling in Project StatusLocation: Now distinguishes between:
else:
logger.debug(f"Project {project_id} not found for status snapshot")
except (KeyError, AttributeError) as e:
logger.debug(f"Project {project_id} missing status fields: {e}")
except Exception as e:
logger.warning(f"Unexpected error sending project status for {project_id}: {e}")This provides much better observability for debugging production issues. 3. ✅ Memory Leak Prevention - Re-subscription HandlingLocation: # Cancel existing task if present (handles re-subscription case to prevent memory leak)
if project_id in heartbeat_tasks:
old_task = heartbeat_tasks[project_id]
old_task.cancel()
try:
await old_task
except asyncio.CancelledError:
pass
logger.debug(f"Cancelled old heartbeat task for project {project_id}")Impact: No more orphaned background tasks when clients re-subscribe to the same project. 4. ✅ Heartbeat Cancellation on UnsubscribeLocation: # Cancel heartbeat task for this project
if project_id in heartbeat_tasks:
heartbeat_tasks[project_id].cancel()
try:
await heartbeat_tasks[project_id]
except asyncio.CancelledError:
pass
del heartbeat_tasks[project_id]
logger.debug(f"Cancelled heartbeat task for unsubscribed project {project_id}")Impact: No unnecessary network traffic or CPU usage after unsubscribe. 5. ✅ Proper CancelledError HandlingLocation: except asyncio.CancelledError:
# Normal cancellation on disconnect
logger.debug(f"Heartbeat task cancelled for project {project_id}")
raiseImpact: Follows asyncio best practices - re-raises CancelledError to propagate cancellation properly. 📊 Code Quality AssessmentStrengths
Test Coverage
Documentation
🔍 Final Verification Checklist
🚀 ImpactThis PR solves a real usability problem:
💡 RecommendationAPPROVE & MERGE ✅ This PR is production-ready. All code quality concerns have been addressed, test coverage is comprehensive, and the implementation follows best practices for asyncio task management. Great work on the quick iteration and thorough responses to feedback! 🎉 |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In @codeframe/ui/routers/websocket.py:
- Line 49: The module-level parse of WEBSOCKET_HEARTBEAT_INTERVAL can raise
ValueError and crash app; change the HEARTBEAT_INTERVAL_SECONDS assignment to
parse the env var inside a try/except (use
os.getenv("WEBSOCKET_HEARTBEAT_INTERVAL") or default "30"), catching ValueError
and falling back to a safe default (e.g., 30), and emit a clear warning via the
module logger (e.g., logger = logging.getLogger(__name__)) that includes the
invalid value and the fallback used; update any code relying on
HEARTBEAT_INTERVAL_SECONDS to use the corrected variable.
🧹 Nitpick comments (1)
codeframe/ui/routers/websocket.py (1)
52-83: Consider sending first heartbeat immediately for better UX.The current implementation sleeps for
HEARTBEAT_INTERVAL_SECONDSbefore sending the first heartbeat, meaning clients wait 30 seconds for the first heartbeat message. Whileconnection_ackis sent immediately on subscription, sending the first heartbeat right away would allow clients to verify the heartbeat mechanism is working without delay.♻️ Proposed refactor to send first heartbeat immediately
async def send_heartbeats(websocket: WebSocket, project_id: int) -> None: """ Send periodic heartbeat messages to a WebSocket client. This function runs in a background task and sends heartbeat messages every HEARTBEAT_INTERVAL_SECONDS to keep the connection alive and verify real-time functionality. Args: websocket: WebSocket connection to send heartbeats to project_id: Project ID to include in heartbeat messages """ try: while True: - await asyncio.sleep(HEARTBEAT_INTERVAL_SECONDS) try: message = { "type": "heartbeat", "project_id": project_id, "timestamp": datetime.now(UTC).isoformat().replace("+00:00", "Z"), } await websocket.send_json(message) logger.debug(f"Sent heartbeat for project {project_id}") except Exception as e: # Connection may have closed - exit the loop logger.debug(f"Heartbeat send failed (connection likely closed): {e}") break + await asyncio.sleep(HEARTBEAT_INTERVAL_SECONDS) except asyncio.CancelledError: # Normal cancellation on disconnect logger.debug(f"Heartbeat task cancelled for project {project_id}") raise
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
codeframe/ui/routers/websocket.py
🧰 Additional context used
📓 Path-based instructions (1)
codeframe/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
codeframe/**/*.py: Use Python 3.11+ for backend development with FastAPI, AsyncAnthropic, SQLite with async support (aiosqlite), and tiktoken for token counting
Use token counting via tiktoken library for token budget management with ~50,000 token limit per conversation
Use asyncio patterns with AsyncAnthropic for async/await in Python backend for concurrent operations
Implement quality gates with multi-stage pre-completion checks (tests → type → coverage → review) and Git + SQLite + context snapshots for project state rollback
Use tiered memory system (HOT/WARM/COLD) with importance scoring using hybrid exponential decay algorithm for context management with 30-50% token reduction
Implement session lifecycle management with auto-save/restore using file-based storage at .codeframe/session_state.json
Files:
codeframe/ui/routers/websocket.py
🧠 Learnings (3)
📓 Common learnings
Learnt from: CR
Repo: frankbria/codeframe PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-01-04T06:26:12.870Z
Learning: Applies to web-ui/src/lib/websocketMessageMapper.ts : Implement WebSocket event mapping with 9 event types (agent_created, task_assigned, etc.) with automatic reconnection using exponential backoff (1s → 30s)
📚 Learning: 2025-11-25T19:08:37.203Z
Learnt from: CR
Repo: frankbria/codeframe PR: 0
File: docs/CLAUDE.md:0-0
Timestamp: 2025-11-25T19:08:37.203Z
Learning: Applies to docs/codeframe/ui/**/*.py : Use FastAPI with Uvicorn for the async API backend and WebSockets for real-time communication
Applied to files:
codeframe/ui/routers/websocket.py
📚 Learning: 2026-01-04T06:26:12.870Z
Learnt from: CR
Repo: frankbria/codeframe PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-01-04T06:26:12.870Z
Learning: Applies to web-ui/src/lib/websocketMessageMapper.ts : Implement WebSocket event mapping with 9 event types (agent_created, task_assigned, etc.) with automatic reconnection using exponential backoff (1s → 30s)
Applied to files:
codeframe/ui/routers/websocket.py
🧬 Code graph analysis (1)
codeframe/ui/routers/websocket.py (2)
codeframe/ui/routers/projects.py (1)
get_project(210-247)codeframe/persistence/database.py (2)
get_project(278-280)create_task(358-360)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: Backend Unit Tests
- GitHub Check: Frontend Unit Tests
- GitHub Check: claude-review
- GitHub Check: E2E Smoke Tests (Chromium)
🔇 Additional comments (3)
codeframe/ui/routers/websocket.py (3)
270-314: Well-implemented proactive messaging sequence.The implementation correctly:
- Sends connection acknowledgment to confirm real-time connectivity
- Provides initial project status snapshot with appropriate error handling that distinguishes between not-found, missing-fields, and unexpected errors
- Prevents memory leak on re-subscription by cancelling existing heartbeat tasks before starting new ones (lines 300-307)
- Includes proper metadata (timestamp, project_id) in all proactive messages
The error handling granularity for project status is particularly good—using debug level for expected cases (not found, missing fields) and warning for unexpected errors.
358-366: Proper heartbeat cancellation on unsubscribe.The implementation correctly cancels the heartbeat task when a client unsubscribes, ensuring that heartbeat messages stop immediately and resources are released. The cancellation pattern (cancel → await → catch CancelledError → delete) is the correct approach for cleaning up asyncio tasks.
381-389: Thorough cleanup of heartbeat tasks on disconnect.The finally block correctly cancels all active heartbeat tasks and clears the tracking dictionary, ensuring no resource leaks occur when the WebSocket connection closes. The pattern of cancelling, awaiting, and catching
CancelledErrorfor each task is the proper cleanup approach.
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
✅ Final Review: LGTM - Ready to Merge!Excellent work! All feedback from my previous reviews has been thoroughly addressed. This PR is now production-ready. 🎉 What's Been AddressedI've reviewed all three commits in this PR, and I'm pleased to confirm that all critical issues and recommendations from my previous reviews have been implemented: 1. ✅ Memory Leak Prevention - Re-subscription HandlingLocation: The code now properly cancels existing heartbeat tasks before creating new ones on re-subscription: # Cancel existing task if present (handles re-subscription case to prevent memory leak)
if project_id in heartbeat_tasks:
old_task = heartbeat_tasks[project_id]
old_task.cancel()
try:
await old_task
except asyncio.CancelledError:
pass
logger.debug(f"Cancelled old heartbeat task for project {project_id}")Impact: No more orphaned background tasks when clients rapidly re-subscribe. 2. ✅ Heartbeat Cancellation on UnsubscribeLocation: Heartbeat tasks are now properly cancelled when clients unsubscribe: # Cancel heartbeat task for this project
if project_id in heartbeat_tasks:
heartbeat_tasks[project_id].cancel()
try:
await heartbeat_tasks[project_id]
except asyncio.CancelledError:
pass
del heartbeat_tasks[project_id]Impact: No unnecessary network traffic or CPU usage after unsubscribe. 3. ✅ Configurable Heartbeat IntervalLocation: try:
HEARTBEAT_INTERVAL_SECONDS = int(os.getenv("WEBSOCKET_HEARTBEAT_INTERVAL", "30"))
except ValueError:
logger.warning("Invalid WEBSOCKET_HEARTBEAT_INTERVAL value, using default 30 seconds")
HEARTBEAT_INTERVAL_SECONDS = 30Impact:
4. ✅ Improved Error Handling for Project StatusLocation: Now distinguishes between expected and unexpected errors: else:
logger.debug(f"Project {project_id} not found for status snapshot")
except (KeyError, AttributeError) as e:
# Expected - project may not have status/phase fields
logger.debug(f"Project {project_id} missing status fields: {e}")
except Exception as e:
# Unexpected error - log as warning for investigation
logger.warning(f"Unexpected error sending project status for {project_id}: {e}")Impact: Better observability for debugging production issues. 5. ✅ Proper CancelledError HandlingLocation: except asyncio.CancelledError:
# Normal cancellation on disconnect
logger.debug(f"Heartbeat task cancelled for project {project_id}")
raiseImpact: Follows asyncio best practices - re-raises CancelledError to propagate cancellation properly. 📊 Code Quality AssessmentStrengths
Test Coverage ✅
Documentation ✅
🚀 ImpactThis PR elegantly solves a real usability problem: Before:
After:
🎯 Final RecommendationAPPROVE & MERGE ✅ This PR is production-ready. All code quality concerns have been thoroughly addressed:
The implementation follows asyncio best practices and addresses all edge cases. Great work on the quick iteration and thorough responses to feedback! 🚀 |
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
codeframe/ui/routers/websocket.py (3)
68-82: Consider sending first heartbeat immediately after subscription.The initial
await asyncio.sleep(HEARTBEAT_INTERVAL_SECONDS)at line 70 means the first heartbeat arrives 30 seconds after subscription. Whileconnection_ackprovides immediate feedback, sending the first heartbeat without delay would provide quicker verification of the heartbeat mechanism itself, which could be valuable for testing and monitoring.♻️ Proposed refactor to send first heartbeat immediately
async def send_heartbeats(websocket: WebSocket, project_id: int) -> None: """ Send periodic heartbeat messages to a WebSocket client. This function runs in a background task and sends heartbeat messages every HEARTBEAT_INTERVAL_SECONDS to keep the connection alive and verify real-time functionality. Args: websocket: WebSocket connection to send heartbeats to project_id: Project ID to include in heartbeat messages """ try: + # Send first heartbeat immediately + try: + message = { + "type": "heartbeat", + "project_id": project_id, + "timestamp": datetime.now(UTC).isoformat().replace("+00:00", "Z"), + } + await websocket.send_json(message) + logger.debug(f"Sent initial heartbeat for project {project_id}") + except Exception as e: + logger.debug(f"Initial heartbeat send failed: {e}") + return + while True: await asyncio.sleep(HEARTBEAT_INTERVAL_SECONDS) try: message = { "type": "heartbeat", "project_id": project_id, "timestamp": datetime.now(UTC).isoformat().replace("+00:00", "Z"), } await websocket.send_json(message) logger.debug(f"Sent heartbeat for project {project_id}") except Exception as e: # Connection may have closed - exit the loop logger.debug(f"Heartbeat send failed (connection likely closed): {e}") break except asyncio.CancelledError: # Normal cancellation on disconnect logger.debug(f"Heartbeat task cancelled for project {project_id}") raise
75-75: Prefer.isoformat()with timezone parameter for cleaner timestamp formatting.The pattern
datetime.now(UTC).isoformat().replace("+00:00", "Z")appears multiple times and relies on string replacement to convert the timezone suffix. While functional, usingisoformat(timespec='milliseconds')or strftime would be more explicit and maintainable.♻️ Proposed refactor for timestamp formatting
Option 1: Using strftime (most explicit):
-"timestamp": datetime.now(UTC).isoformat().replace("+00:00", "Z"), +"timestamp": datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S.%fZ")[:-4] + "Z",Option 2: Define a helper function (if used frequently):
def utc_timestamp() -> str: """Return current UTC timestamp in ISO format with Z suffix.""" return datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")Then use:
-"timestamp": datetime.now(UTC).isoformat().replace("+00:00", "Z"), +"timestamp": utc_timestamp(),Also applies to: 278-278, 291-291
295-297: Clarify whether missing status fields are truly "expected" behavior.The comment states "Expected - project may not have status/phase fields" but logging this at debug level suggests uncertainty. If these fields are optional by design, the comment is accurate. However, if they should always be present, this might indicate a data integrity issue that should be logged at a higher level.
Consider either:
- Confirming this is expected behavior and keeping debug logging
- If these fields should always exist, elevating to warning level or adding validation
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
codeframe/ui/routers/websocket.py
🧰 Additional context used
📓 Path-based instructions (1)
codeframe/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
codeframe/**/*.py: Use Python 3.11+ for backend development with FastAPI, AsyncAnthropic, SQLite with async support (aiosqlite), and tiktoken for token counting
Use token counting via tiktoken library for token budget management with ~50,000 token limit per conversation
Use asyncio patterns with AsyncAnthropic for async/await in Python backend for concurrent operations
Implement quality gates with multi-stage pre-completion checks (tests → type → coverage → review) and Git + SQLite + context snapshots for project state rollback
Use tiered memory system (HOT/WARM/COLD) with importance scoring using hybrid exponential decay algorithm for context management with 30-50% token reduction
Implement session lifecycle management with auto-save/restore using file-based storage at .codeframe/session_state.json
Files:
codeframe/ui/routers/websocket.py
🧠 Learnings (3)
📓 Common learnings
Learnt from: CR
Repo: frankbria/codeframe PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-01-04T06:26:12.870Z
Learning: Applies to web-ui/src/lib/websocketMessageMapper.ts : Implement WebSocket event mapping with 9 event types (agent_created, task_assigned, etc.) with automatic reconnection using exponential backoff (1s → 30s)
📚 Learning: 2025-11-25T19:08:37.203Z
Learnt from: CR
Repo: frankbria/codeframe PR: 0
File: docs/CLAUDE.md:0-0
Timestamp: 2025-11-25T19:08:37.203Z
Learning: Applies to docs/codeframe/ui/**/*.py : Use FastAPI with Uvicorn for the async API backend and WebSockets for real-time communication
Applied to files:
codeframe/ui/routers/websocket.py
📚 Learning: 2026-01-04T06:26:12.870Z
Learnt from: CR
Repo: frankbria/codeframe PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-01-04T06:26:12.870Z
Learning: Applies to web-ui/src/lib/websocketMessageMapper.ts : Implement WebSocket event mapping with 9 event types (agent_created, task_assigned, etc.) with automatic reconnection using exponential backoff (1s → 30s)
Applied to files:
codeframe/ui/routers/websocket.py
🧬 Code graph analysis (1)
codeframe/ui/routers/websocket.py (2)
codeframe/persistence/database.py (2)
get_project(278-280)create_task(358-360)codeframe/ui/routers/tasks.py (1)
create_task(46-112)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)
- GitHub Check: Backend Unit Tests
- GitHub Check: Frontend Unit Tests
- GitHub Check: E2E Smoke Tests (Chromium)
🔇 Additional comments (6)
codeframe/ui/routers/websocket.py (6)
49-53: LGTM: Robust configuration parsing.The environment variable parsing properly handles invalid values with a try/except block, logs a warning, and falls back to a safe default. This addresses the potential ValueError crash noted in previous reviews.
202-203: LGTM: Clean task tracking initialization.The heartbeat_tasks dictionary is properly typed and scoped to the WebSocket connection lifecycle.
304-318: LGTM: Proper heartbeat task lifecycle management.The re-subscription handling correctly cancels existing heartbeat tasks before starting new ones (lines 304-311), preventing memory leaks and duplicate heartbeats. The task creation and tracking (lines 314-318) is clean and properly logged.
362-370: LGTM: Complete heartbeat cleanup on unsubscribe.The heartbeat task cancellation properly handles the cancellation flow: cancels the task, awaits it, catches CancelledError, removes from tracking dict, and logs. This ensures no heartbeats are sent after unsubscription and prevents resource leaks.
385-393: LGTM: Comprehensive cleanup in finally block.All heartbeat tasks are properly cancelled on disconnect, with each task awaited and CancelledError caught. The dict is cleared at the end, ensuring complete cleanup regardless of how the connection terminated.
284-292: The code is correct. Verified thatdb.get_project()returnsOptional[dict]with both "status" and "phase" fields guaranteed by the database schema. The defensive usage of.get()with "unknown" defaults is appropriate and consistent with project-wide patterns.
1 participant
Summary
connection_ack,project_status,heartbeatChanges
Backend (
codeframe/ui/routers/websocket.py)HEARTBEAT_INTERVAL_SECONDSconstant (30 seconds)send_heartbeats()async function for background heartbeat tasksconnection_ackimmediately after subscriptionproject_statuswith current project state on subscriptionBroadcast Helper (
codeframe/ui/websocket_broadcasts.py)broadcast_heartbeat()helper functionFrontend Types (
web-ui/src/types/index.ts)heartbeat,connection_ack,project_statustoWebSocketMessageTypeunionTests
tests/ui/test_websocket_proactive.pywith 17 comprehensive testsdrain_proactive_messages()helper for integration testsMessage Flow After Subscription
Test plan
uv run pytest tests/ui/test_websocket_proactive.py -v- 17 tests passuv run pytest tests/ui/test_websocket_integration.py -v- 30 tests passuv run pytest tests/ui/ -v- 238 tests pass, 1 skippeduv run ruff check .- All checks passednpm run type-checkin web-ui - No TypeScript errorstests/e2e/test_dashboard.spec.tsWebSocket test should now receive messagesSummary by CodeRabbit
New Features
Tests
Docs / Types
✏️ Tip: You can customize this high-level summary in your review settings.