Skip to content

feat: replace index status polling with WebSocket updates#138

Merged
JohnRDOrazio merged 7 commits into
devfrom
feat/index-status-websocket
Apr 13, 2026
Merged

feat: replace index status polling with WebSocket updates#138
JohnRDOrazio merged 7 commits into
devfrom
feat/index-status-websocket

Conversation

@JohnRDOrazio
Copy link
Copy Markdown
Member

@JohnRDOrazio JohnRDOrazio commented Apr 13, 2026

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

  • New lib/api/indexStatus.tscreateIndexWebSocket() and IndexWebSocketManager (mirrors lib/api/lint.ts pattern)
  • Settings page — replaced setInterval polling with WebSocket subscription to /api/v1/projects/{id}/ontology/index-ws; receives index_started, index_complete, index_failed events in real-time
  • Fixes the issue where polling stopped after a 404 response before the worker created the index record

WebSocket authentication

  • createIndexWebSocket() and createLintWebSocket() now accept an optional token parameter, forwarded as ?token=... query string
  • IndexWebSocketManager and LintWebSocketManager store and pass the token on connect/reconnect
  • Settings page passes accessTokenRef.current to the index WebSocket
  • HealthCheckPanel passes accessToken prop to the lint WebSocket

Dependencies

Closes #137

Test plan

  • Open project settings, click "Rebuild Index", verify status updates in real-time via WebSocket (no polling)
  • Verify status shows "indexing" during build, then "ready" with entity count on completion
  • Verify failed index shows error message
  • Verify lint WebSocket still works with token auth in HealthCheckPanel
  • Verify unauthenticated WebSocket connections are rejected
  • TypeScript type-check passes

🤖 Generated with Claude Code

Summary by CodeRabbit

Release Notes

  • New Features

    • Added real-time ontology index status updates via WebSocket with automatic reconnection and exponential backoff.
    • Improved WebSocket authentication to support secure token-based connections.
  • Refactor

    • Replaced polling-based index status checks with efficient WebSocket streaming.
    • Removed connection status indicator from project viewer interface.
  • Tests

    • Added comprehensive test coverage for WebSocket index status management and hook behavior.

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>
@coderabbitai
Copy link
Copy Markdown

coderabbitai Bot commented Apr 13, 2026

📝 Walkthrough

Walkthrough

This PR introduces WebSocket-based real-time ontology index status updates, replacing polling. It adds createIndexWebSocket and IndexWebSocketManager in a new lib/api/indexStatus.ts module, integrates WebSocket subscriptions into the settings page, extends hook infrastructure to support authentication tokens, and includes comprehensive test coverage.

Changes

Cohort / File(s) Summary
WebSocket Index Status Infrastructure
lib/api/indexStatus.ts
New module introducing IndexWebSocketMessage interface, createIndexWebSocket() factory function with URL fallback logic (NEXT_PUBLIC_WS_URLNEXT_PUBLIC_API_URL conversion → default), and IndexWebSocketManager class with auto-reconnect support (exponential backoff, 5 max attempts).
Index Status Tests
__tests__/lib/api/indexStatus.test.ts, __tests__/lib/hooks/useIndexStatus.test.ts
Comprehensive test suites for WebSocket creation (URL construction, message/error/close handling), manager lifecycle (connect/disconnect, duplicate prevention, reconnection), and useIndexStatus hook behavior with mocked API and React Query integration.
Settings Page WebSocket Integration
app/projects/[id]/settings/page.tsx
Replaced interval-based polling with WebSocket subscription; added effect that opens socket and updates isReindexing/indexStatus state on index_started/index_complete/index_failed messages; invalidates React Query cache on completion to sync server state.
Viewer and Project Hook Updates
app/projects/[id]/editor/page.tsx, lib/hooks/useProjectViewer.ts
Added enableWebSocket?: boolean option to useProjectViewer; updated editor page to pass enableWebSocket: true; gated collaboration WebSocket to require authentication (enableWebSocket && projectId && accessToken && authenticated).
Token Support for WebSocket Authentication
lib/api/lint.ts, lib/hooks/useCollaborationStatus.ts, components/editor/HealthCheckPanel.tsx
Extended createLintWebSocket and LintWebSocketManager to accept optional token; updated useCollaborationStatus to accept and append token to WebSocket URL; modified HealthCheckPanel to pass token and depend on it for reconnection.
Connection Status UI Removal
app/projects/[id]/page.tsx
Removed ConnectionStatus import and UI rendering of WebSocket connection status from viewer header; eliminated destructuring of connectionStatus, wsEndpoint, and wsPurpose from hook output.

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
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

Possibly related PRs

Suggested reviewers

  • damienriehl

Poem

🐰 WebSocket whispers now guide the way,
No more polls through the long delay,
Real-time messages flow so free,
Index status, swift as can be!
Reconnection keeps the thread alive, 🧵✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 37.50% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title accurately and concisely summarizes the main change: replacing index status polling with WebSocket updates, which is the primary objective of the PR.
Linked Issues check ✅ Passed All coding requirements from issue #137 are met: createIndexWebSocket/IndexWebSocketManager implemented, settings page uses WebSocket instead of polling, index events handled with state/cache updates, and authentication added to WebSocket connections.
Out of Scope Changes check ✅ Passed All changes are directly related to the PR objective of replacing polling with WebSocket updates. UI changes to editor page and removal of connection status display are within scope for enabling WebSocket authentication.
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch feat/index-status-websocket

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@codecov
Copy link
Copy Markdown

codecov Bot commented Apr 13, 2026

Codecov Report

❌ Patch coverage is 85.48387% with 9 lines in your changes missing coverage. Please review.
✅ All tests successful. No failed tests found.

Files with missing lines Patch % Lines
components/editor/HealthCheckPanel.tsx 28.57% 5 Missing ⚠️
lib/api/indexStatus.ts 97.91% 0 Missing and 1 partial ⚠️
lib/api/lint.ts 66.66% 0 Missing and 1 partial ⚠️
lib/hooks/useCollaborationStatus.ts 50.00% 0 Missing and 1 partial ⚠️
lib/hooks/useProjectViewer.ts 50.00% 0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

JohnRDOrazio and others added 5 commits April 13, 2026 16:32
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>
Copy link
Copy Markdown

@coderabbitai coderabbitai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 | 🟠 Major

Handle WebSocket failure callbacks to prevent stale “Running…” state.

At Line 153, onError/onClose are undefined. If the socket fails after a lint run starts, isRunning may 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.ts which 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 as CONNECTING, not OPEN.

The current mock makes every socket look fully open immediately, which is not how browser WebSocket behaves. That means these tests never exercise the CONNECTING window where duplicate reconnects are most likely to happen, so they'll miss the overlap bug in IndexWebSocketManager.

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

📥 Commits

Reviewing files that changed from the base of the PR and between e050cd7 and 6e42a38.

📒 Files selected for processing (10)
  • __tests__/lib/api/indexStatus.test.ts
  • __tests__/lib/hooks/useIndexStatus.test.ts
  • app/projects/[id]/editor/page.tsx
  • app/projects/[id]/page.tsx
  • app/projects/[id]/settings/page.tsx
  • components/editor/HealthCheckPanel.tsx
  • lib/api/indexStatus.ts
  • lib/api/lint.ts
  • lib/hooks/useCollaborationStatus.ts
  • lib/hooks/useProjectViewer.ts
💤 Files with no reviewable changes (1)
  • app/projects/[id]/page.tsx

Comment thread app/projects/[id]/settings/page.tsx Outdated
Comment thread lib/api/indexStatus.ts
- 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>
@JohnRDOrazio JohnRDOrazio merged commit 2d02cd5 into dev Apr 13, 2026
12 checks passed
@JohnRDOrazio JohnRDOrazio deleted the feat/index-status-websocket branch April 13, 2026 19:28
@JohnRDOrazio JohnRDOrazio added this to the v0.4.0 milestone Apr 13, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat: use WebSocket for real-time ontology index status updates

1 participant