feat: replace index status polling with WebSocket updates#138
Conversation
Replace setInterval polling in the settings page with a WebSocket
connection to /api/v1/projects/{id}/ontology/index-ws. The WebSocket
receives index_started, index_complete, and index_failed events from
the backend worker in real-time, avoiding the issue where polling
stopped after a 404 response before the index record was created.
Closes #137
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
📝 WalkthroughWalkthroughThis PR introduces WebSocket-based real-time ontology index status updates, replacing polling. It adds Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant SettingsPage as Settings Page
participant WebSocket as WebSocket Manager
participant Backend as Backend (Index Worker)
participant ReactQuery as React Query Cache
User->>SettingsPage: Trigger Reindex
SettingsPage->>Backend: POST /api/v1/projects/:id/ontology/reindex
Backend-->>SettingsPage: Reindex queued
SettingsPage->>SettingsPage: Set isReindexing = true
SettingsPage->>WebSocket: connect() via useEffect
WebSocket->>Backend: Open WS /api/v1/projects/:id/ontology/index-ws
Backend-->>WebSocket: Connection established
Backend->>WebSocket: Send index_started message
WebSocket->>SettingsPage: onMessage (index_started)
SettingsPage->>SettingsPage: Update indexStatus, isReindexing = true
Backend->>Backend: Process indexing...
Backend->>WebSocket: Send index_complete message
WebSocket->>SettingsPage: onMessage (index_complete)
SettingsPage->>SettingsPage: Set isReindexing = false
SettingsPage->>ReactQuery: Invalidate indexStatus cache
ReactQuery->>Backend: Fetch fresh index status
Backend-->>ReactQuery: Return updated status
ReactQuery-->>SettingsPage: Display final status
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
Covers createIndexWebSocket URL resolution, message parsing, error handling, and IndexWebSocketManager reconnection logic. Also tests useIndexStatus hook query enablement and error states. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
The backend now requires a token query parameter for WebSocket auth. Updated createIndexWebSocket() and IndexWebSocketManager to accept and forward the access token, and the settings page passes accessTokenRef. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Updated createLintWebSocket() and LintWebSocketManager to accept and forward the access token via query parameter, matching the index WebSocket auth pattern. HealthCheckPanel now passes accessToken. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
The useEffect that creates the lint WebSocket was missing accessToken in its dependency array, capturing undefined on first render. Added accessToken to deps and early-return when not yet available. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
The lint WebSocket connection and status icons were appearing on the project viewer page, causing a flood of failed unauthenticated connections with auto-reconnect. - Remove ConnectionStatus icons from viewer page (page.tsx) - Add enableWebSocket flag to useProjectViewer (defaults to false) - Only the editor page passes enableWebSocket: true - useCollaborationStatus now accepts and forwards auth token - WebSocket only connects when authenticated with a valid token Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
components/editor/HealthCheckPanel.tsx (1)
151-164:⚠️ Potential issue | 🟠 MajorHandle WebSocket failure callbacks to prevent stale “Running…” state.
At Line 153,
onError/onCloseareundefined. If the socket fails after a lint run starts,isRunningmay never recover and users get no actionable error.💡 Suggested fix
- ws = createLintWebSocket(projectId, handleMessage, undefined, undefined, accessToken); + ws = createLintWebSocket( + projectId, + handleMessage, + () => { + if (!isActive) return; + setIsRunning(false); + setError("Real-time lint connection failed"); + }, + (event) => { + if (!isActive) return; + if (event.code !== 1000) { + setIsRunning(false); + setError("Real-time lint connection closed"); + } + }, + accessToken + );🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@components/editor/HealthCheckPanel.tsx` around lines 151 - 164, The WebSocket created via createLintWebSocket is instantiated without onError/onClose handlers so a socket failure can leave isRunning true and the UI stuck; update the call in the useEffect that creates ws (where isActive and handleMessage are used) to supply onError and onClose callbacks that set the lint run state (isRunning) to false, surface an error message (or call the existing error handler), and perform cleanup (close ws and clear any timers) to ensure the UI recovers; also ensure the cleanup return still nullifies isActive, clears timeoutId, and closes ws to avoid leaks.
🧹 Nitpick comments (4)
lib/hooks/useCollaborationStatus.ts (1)
29-34: Prefer reusing the lint API WebSocket factory instead of rebuilding URL logic here.This duplicates endpoint/auth URL composition and can drift from
lib/api/lint.ts(host fallback + query behavior). Route this through the domain API helper for one source of truth.As per coding guidelines: "All backend communication should go through
lib/api/client.tswhich provides type-safe API methods ... domain-specific APIs ...".🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@lib/hooks/useCollaborationStatus.ts` around lines 29 - 34, The getWsUrl function in useCollaborationStatus.ts reconstructs the lint WebSocket URL and auth query params, duplicating logic from lib/api/lint.ts and lib/api/client.ts; replace this manual URL composition by calling the lint WebSocket factory or client helper exported from lib/api/lint.ts (or the domain API in lib/api/client.ts) instead of building wsUrl/token yourself, preserving the token/query behavior and projectId argument; update the useCollaborationStatus import to use that factory (e.g., createLintWebSocket/getLintWebSocketUrl or the client method) and pass projectId and token through so the single source of truth in lib/api/* controls host fallback and query composition.lib/api/lint.ts (1)
185-186: Consider hardening token handling for WebSocket query auth.Using bearer tokens in URL query params can leak through proxy/access logs. Consider switching to short-lived WS tickets (or enforce log redaction for query strings) to reduce exposure.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@lib/api/lint.ts` around lines 185 - 186, The current WebSocket creation embeds the bearer token in the query string (token, wsUrl, projectId, new WebSocket(...)) which risks leaking it in logs; fix by obtaining a short-lived WS ticket from the server (add/fetch via a new/getLintWSTicket API for the given projectId) and use that ticket instead of the bearer token when opening the socket, or alternatively send the token via a WebSocket subprotocol (pass the token as the protocol string) so it is not in the URL; update the WebSocket creation site to use the ticket or subprotocol and remove the token query param and adjust server-side auth to accept the ticket/subprotocol.__tests__/lib/api/indexStatus.test.ts (2)
9-25: Model the initial WebSocket state asCONNECTING, notOPEN.The current mock makes every socket look fully open immediately, which is not how browser
WebSocketbehaves. That means these tests never exercise theCONNECTINGwindow where duplicate reconnects are most likely to happen, so they'll miss the overlap bug inIndexWebSocketManager.Also applies to: 159-167
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@__tests__/lib/api/indexStatus.test.ts` around lines 9 - 25, The MockWebSocket currently sets readyState = MockWebSocket.OPEN causing tests to skip the CONNECTING phase; change the mock so readyState is initialized to MockWebSocket.CONNECTING (and ensure tests simulate transition to OPEN by updating readyState and invoking onopen where needed) so the CONNECTING window is exercised — update the MockWebSocket class (readyState field and any places in its constructor or test helpers that assume OPEN) and adjust usages that expect an immediate OPEN to instead trigger the onopen callback after simulating connection.
42-66: Add a URL test for the authenticated socket path.Token forwarding is one of the core behaviors in this PR, but the suite never asserts that
?token=is appended and encoded. A regression there would break authenticated WebSocket connections while all current URL tests still pass.🧪 Suggested test
it("falls back to ws://localhost:8000 when no env vars set", () => { const onMessage = vi.fn(); const ws = createIndexWebSocket("p1", onMessage); expect(ws.url).toBe("ws://localhost:8000/api/v1/projects/p1/ontology/index-ws"); }); + + it("appends an encoded token when provided", () => { + const ws = createIndexWebSocket( + "p1", + vi.fn(), + undefined, + undefined, + "a+b/=" + ); + + expect(ws.url).toBe( + "ws://localhost:8000/api/v1/projects/p1/ontology/index-ws?token=a%2Bb%2F%3D" + ); + });🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@__tests__/lib/api/indexStatus.test.ts` around lines 42 - 66, Add a test in the same suite that verifies createIndexWebSocket appends an encoded token query param to the WebSocket URL: set the environment token (e.g., process.env.NEXT_PUBLIC_INDEX_SOCKET_TOKEN or whatever token source your implementation reads), call createIndexWebSocket("p1", onMessage) and assert ws.url ends with "?token=<encoded-token>" (or contains "&token=" if other query params exist), ensuring the token is encoded via encodeURIComponent; reference createIndexWebSocket to locate where the URL is built and add the assertion similar to the existing URL tests so token forwarding is covered.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@app/projects/`[id]/settings/page.tsx:
- Around line 185-226: The effect opens a raw WebSocket (createIndexWebSocket)
as soon as project?.id is present which can occur before session?.accessToken is
available, causing an unauthenticated socket that never recovers; replace the
direct createIndexWebSocket usage with the reconnecting IndexWebSocketManager
(or its manager API) and key the effect on session?.accessToken (or
accessTokenRef.current) so the socket is only opened once a token exists, and
use the manager to create/close sockets and handle reconnects when token or
project changes; update the useEffect dependencies to include
session?.accessToken (or token ref) and ensure cleanup calls manager.close() or
equivalent instead of raw ws.close().
In `@lib/api/indexStatus.ts`:
- Around line 79-115: The connect()/handleReconnect() logic allows overlapping
connects and double-retries and never resets the retry budget; modify connect()
to treat WebSocket.CONNECTING the same as OPEN (skip creating a new socket if
ws.readyState is OPEN or CONNECTING) and ensure the createIndexWebSocket
callbacks (onopen/onerror/onclose) deduplicate reconnect triggers by using a
single reconnection path: reset reconnectAttempts to 0 inside the socket onopen
handler (so a successful reopen restores the budget) and guard onerror/onclose
so they call handleReconnect() only if a reconnection hasn't already been
scheduled (e.g., a boolean like isReconnecting or checking ws.readyState before
calling handleReconnect()); keep disconnect() behavior but ensure any scheduled
reconnect timers are cleared when explicitly closing.
---
Outside diff comments:
In `@components/editor/HealthCheckPanel.tsx`:
- Around line 151-164: The WebSocket created via createLintWebSocket is
instantiated without onError/onClose handlers so a socket failure can leave
isRunning true and the UI stuck; update the call in the useEffect that creates
ws (where isActive and handleMessage are used) to supply onError and onClose
callbacks that set the lint run state (isRunning) to false, surface an error
message (or call the existing error handler), and perform cleanup (close ws and
clear any timers) to ensure the UI recovers; also ensure the cleanup return
still nullifies isActive, clears timeoutId, and closes ws to avoid leaks.
---
Nitpick comments:
In `@__tests__/lib/api/indexStatus.test.ts`:
- Around line 9-25: The MockWebSocket currently sets readyState =
MockWebSocket.OPEN causing tests to skip the CONNECTING phase; change the mock
so readyState is initialized to MockWebSocket.CONNECTING (and ensure tests
simulate transition to OPEN by updating readyState and invoking onopen where
needed) so the CONNECTING window is exercised — update the MockWebSocket class
(readyState field and any places in its constructor or test helpers that assume
OPEN) and adjust usages that expect an immediate OPEN to instead trigger the
onopen callback after simulating connection.
- Around line 42-66: Add a test in the same suite that verifies
createIndexWebSocket appends an encoded token query param to the WebSocket URL:
set the environment token (e.g., process.env.NEXT_PUBLIC_INDEX_SOCKET_TOKEN or
whatever token source your implementation reads), call
createIndexWebSocket("p1", onMessage) and assert ws.url ends with
"?token=<encoded-token>" (or contains "&token=" if other query params exist),
ensuring the token is encoded via encodeURIComponent; reference
createIndexWebSocket to locate where the URL is built and add the assertion
similar to the existing URL tests so token forwarding is covered.
In `@lib/api/lint.ts`:
- Around line 185-186: The current WebSocket creation embeds the bearer token in
the query string (token, wsUrl, projectId, new WebSocket(...)) which risks
leaking it in logs; fix by obtaining a short-lived WS ticket from the server
(add/fetch via a new/getLintWSTicket API for the given projectId) and use that
ticket instead of the bearer token when opening the socket, or alternatively
send the token via a WebSocket subprotocol (pass the token as the protocol
string) so it is not in the URL; update the WebSocket creation site to use the
ticket or subprotocol and remove the token query param and adjust server-side
auth to accept the ticket/subprotocol.
In `@lib/hooks/useCollaborationStatus.ts`:
- Around line 29-34: The getWsUrl function in useCollaborationStatus.ts
reconstructs the lint WebSocket URL and auth query params, duplicating logic
from lib/api/lint.ts and lib/api/client.ts; replace this manual URL composition
by calling the lint WebSocket factory or client helper exported from
lib/api/lint.ts (or the domain API in lib/api/client.ts) instead of building
wsUrl/token yourself, preserving the token/query behavior and projectId
argument; update the useCollaborationStatus import to use that factory (e.g.,
createLintWebSocket/getLintWebSocketUrl or the client method) and pass projectId
and token through so the single source of truth in lib/api/* controls host
fallback and query composition.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: eea372d7-2d1c-4e30-afb2-f0db87ba500e
📒 Files selected for processing (10)
__tests__/lib/api/indexStatus.test.ts__tests__/lib/hooks/useIndexStatus.test.tsapp/projects/[id]/editor/page.tsxapp/projects/[id]/page.tsxapp/projects/[id]/settings/page.tsxcomponents/editor/HealthCheckPanel.tsxlib/api/indexStatus.tslib/api/lint.tslib/hooks/useCollaborationStatus.tslib/hooks/useProjectViewer.ts
💤 Files with no reviewable changes (1)
- app/projects/[id]/page.tsx
- IndexWebSocketManager: guard CONNECTING state to prevent duplicate sockets, reset retry budget on successful open, use tracked timer to prevent double-reconnect from onerror+onclose, clear pending timer on disconnect - settings/page.tsx: switch from raw createIndexWebSocket to IndexWebSocketManager and gate effect on session.accessToken - HealthCheckPanel: add onError/onClose handlers to lint WebSocket so isRunning resets on socket failure instead of leaving UI stuck - Tests: MockWebSocket starts at CONNECTING, added token URL tests, CONNECTING guard, retry budget reset, timer cleanup, and double-reconnect prevention tests Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Summary
This PR replaces the polling-based index status mechanism with real-time WebSocket updates and adds authentication to both WebSocket connections.
WebSocket for index status
lib/api/indexStatus.ts—createIndexWebSocket()andIndexWebSocketManager(mirrorslib/api/lint.tspattern)setIntervalpolling with WebSocket subscription to/api/v1/projects/{id}/ontology/index-ws; receivesindex_started,index_complete,index_failedevents in real-timeWebSocket authentication
createIndexWebSocket()andcreateLintWebSocket()now accept an optionaltokenparameter, forwarded as?token=...query stringIndexWebSocketManagerandLintWebSocketManagerstore and pass the token on connect/reconnectaccessTokenRef.currentto the index WebSocketaccessTokenprop to the lint WebSocketDependencies
Closes #137
Test plan
🤖 Generated with Claude Code
Summary by CodeRabbit
Release Notes
New Features
Refactor
Tests