WebSocket Migration - Replace SSE with real-time updates

[circus-chimp]

Replace Server-Sent Events (SSE) with WebSocket for real-time dashboard updates. Reference existing WebSocket pattern in packages/shared/src/ws.ts

[circus-chimp]

Thanks for the enhancement request! This is a significant architectural change that involves replacing SSE with WebSocket for dashboard real-time updates.

I'm handing this off to our architect to design an implementation plan that considers:

• Current SSE implementation analysis
• Existing WebSocket pattern integration
• Impact on the dashboard
• Testing strategy

Working on it now. 🚀

[circus-chimp]

Thanks for adding the dashboard label — confirmed this is a dashboard-related enhancement. Our architect is already working on designing the WebSocket migration plan. 🏗️

[circus-chimp]

Implementation Plan: WebSocket Migration from SSE

Current State

The dashboard currently uses 3 SSE (Server-Sent Events) endpoints for real-time updates:

Endpoint	Service	Frontend Consumer	Purpose
GET /api/chimps/live	ChimpService.createLiveStream()	useChimps() hook	Initial chimp list + live status/profile/topics updates
GET /api/chimp/:chimpId/activity	createActivityStream()	useSSE() hook	Per-chimp event/output/meta activity feed
GET /api/meta/events	MessageService.createMetaEventStream()	useSSE() hook	Raw meta events stream

Each SSE endpoint creates a ReadableStream<Uint8Array> that subscribes to NATS subjects, formats events as text/event-stream, and sends keep-alive pings every 3 seconds. The frontend uses browser EventSource with exponential backoff retry.

Why WebSocket

• Bidirectional: WebSocket allows the client to send subscribe/unsubscribe messages, enabling dynamic channel management without opening new connections.
• Connection efficiency: 3 SSE connections per browser tab → 1 WebSocket. The activity page is particularly expensive since it creates JetStream consumers per SSE connection.
• Native Bun support: Bun.serve() already supports WebSocket with built-in pub/sub — this is the idiomatic pattern per AGENTS.md.
• No ping hacks: SSE requires :ping\n\n keep-alive workarounds. WebSocket has native ping/pong frames.

Architecture Decision: Single Multiplexed WebSocket

Chosen approach: One WebSocket endpoint (/ws) with channel-based multiplexing.

Rationale: Bun.serve has a single global websocket handler — even "multiple WS endpoints" route through the same open/message/close callbacks. A single endpoint with explicit subscribe/unsubscribe messages is the cleanest mapping and reduces per-tab connections from 3 to 1.

Rejected alternatives:

• Multiple WS endpoints: Same handler complexity, more connections, no benefit.
• Bun pub/sub only: Good for broadcast but doesn't handle per-connection initial data delivery (the chimps-live init event requires fetching from Redis).

Wire Protocol

Define in packages/shared/src/protocol.ts (alongside existing schemas):

// Client → Server
type WsClientMessage =
  | { action: "subscribe"; channel: "chimps-live" }
  | { action: "subscribe"; channel: "meta-events" }
  | { action: "subscribe"; channel: "activity"; chimpId: string }
  | { action: "unsubscribe"; channel: "chimps-live" }
  | { action: "unsubscribe"; channel: "meta-events" }
  | { action: "unsubscribe"; channel: "activity"; chimpId: string }
  | { action: "ping" }

// Server → Client
type WsServerMessage =
  | { channel: "chimps-live"; event: "init"; data: { chimps: ChimpState[] } }
  | { channel: "chimps-live"; event: "status"; data: { chimpId, status, timestamp } }
  | { channel: "chimps-live"; event: "profile"; data: { chimpId, profile, timestamp } }
  | { channel: "chimps-live"; event: "topics"; data: { chimpId, topics, timestamp } }
  | { channel: "activity"; chimpId: string; data: ActivityEvent }
  | { channel: "meta-events"; data: MetaEvent }
  | { channel: "system"; event: "subscribed"; data: { channel: string } }
  | { channel: "system"; event: "error"; data: { message: string } }
  | { channel: "system"; event: "pong" }

All messages are Zod-validated on both sides — consistent with the project's end-to-end validation pattern.

Per-Connection State

Each WebSocket connection tracks:

interface WsConnectionState {
  chimpsLive: { sub: NatsSubscription } | null;
  metaEvents: { sub: NatsSubscription } | null;
  activity: Map<string, {  // keyed by chimpId
    consumers: JetStreamConsumer[];
    messages: ConsumerMessages[];
    metaSub: NatsSubscription;
  }>;
}

On close, all NATS subscriptions and JetStream consumers are cleaned up — matching the existing resource ownership pattern.

Implementation Phases

---

Phase 1: Protocol Schemas (packages/shared)

Files to modify:

• packages/shared/src/protocol.ts — Add Zod schemas for WsClientMessage and WsServerMessage

What to do:

• Add WsClientMessageSchema (discriminated union on action + channel)
• Add WsServerMessageSchema (discriminated union on channel + event)
• Export types derived from schemas
• Keep alongside existing protocol schemas — these are wire protocol definitions

---

Phase 2: Server-Side WebSocket Handler (packages/dashboard)

Files to create:

• packages/dashboard/src/services/ws-handler.ts — WebSocket connection manager

Files to modify:

• packages/dashboard/src/dashboard.ts — Add websocket config + fetch fallback for upgrade

What to do:

1. Create WsHandler class that receives dependencies (NATS connection, Redis status source, topic registry, chimp profile store, logger) via constructor injection — same pattern as existing services.

2. Implement channel handlers — Extract the NATS subscription logic from existing SSE services into channel-specific methods:

   • subscribeChimpsLive(ws) — Fetch initial chimps from Redis, subscribe to meta.>, forward parsed events
   • subscribeActivity(ws, chimpId) — Create JetStream consumers + meta sub (mirrors createActivityStream)
   • subscribeMetaEvents(ws) — Subscribe to meta.>, forward parsed events
   • unsubscribe*(ws) — Clean up subscriptions for that channel

3. Wire into Bun.serve:

   this.server = serve({
     routes: {
       "/healthz": new Response("OK"),
       "/*": index,
       ...chimpRouter.routes,
       ...profileRouter.routes,
       ...messageRouter.routes,
       // Activity routes removed — served via WS now
     },
     fetch(req, server) {
       const url = new URL(req.url);
       if (url.pathname === "/ws") {
         server.upgrade(req);
         return;
       }
       return new Response("Not Found", { status: 404 });
     },
     websocket: wsHandler.handlers(),
   });

4. Keep REST endpoints (/api/chimps, /api/topics, etc.) — only SSE streaming endpoints are replaced.

---

Phase 3: Frontend WebSocket Hook (packages/dashboard)

Files to create:

• packages/dashboard/src/hooks/useWs.ts — Shared WebSocket connection manager (singleton per page)

Files to modify:

• packages/dashboard/src/hooks/useChimps.ts — Replace EventSource with WS subscribe
• packages/dashboard/src/pages/ChimpActivity.tsx — Replace useSSE with WS activity channel

What to do:

1. Create useWs() singleton hook — Manages one WebSocket connection per browser tab:

   • Connect to ws://host/ws with exponential backoff retry
   • Expose subscribe(channel, callback) / unsubscribe(channel)
   • Parse incoming messages, route to registered callbacks by channel
   • Return { connected, error, subscribe, unsubscribe, send }

2. Migrate useChimps():

   • On mount: subscribe("chimps-live")
   • Handle init, status, profile, topics events (same state update logic)
   • On unmount: unsubscribe("chimps-live")

3. Create useActivity(chimpId) hook (replaces useSSE usage in ChimpActivity):

   • On mount/chimpId change: subscribe("activity", { chimpId })
   • Forward messages to state
   • On unmount: unsubscribe("activity", { chimpId })

---

Phase 4: Cleanup

Files to delete/gut:

• packages/dashboard/src/hooks/useSSE.ts — Remove entirely (no more SSE consumers)
• packages/dashboard/src/services/activity-service.ts — Remove (logic moved to ws-handler)
• packages/dashboard/src/routes/activity.ts — Remove (no more SSE endpoint)

Files to modify:

• packages/dashboard/src/services/chimp-service.ts — Remove createLiveStream() method
• packages/dashboard/src/services/message-service.ts — Remove createMetaEventStream() method
• packages/dashboard/src/routes/chimps.ts — Remove /api/chimps/live SSE route
• packages/dashboard/src/routes/messages.ts — Remove /api/meta/events SSE route
• packages/dashboard/src/dashboard.ts — Remove activity router import, clean up service wiring

---

Key Constraints

1. No ws library — Use Bun's built-in WebSocket (per AGENTS.md)
2. No type assertions — All Zod-validated, types derived from schemas
3. Resource ownership — WsHandler owns per-connection NATS subscriptions; close cleans up everything
4. Explicit routes object — Don't spread into routes (per AGENTS.md: "causes type error with websocket required")
5. Constructor injection — No process.env in components; WsHandler receives dependencies via constructor

Risk & Mitigation

Risk	Mitigation
JetStream consumer leak on abrupt disconnect	close handler runs cleanup; add timeout-based safety sweep
Memory growth from many concurrent WS connections	Per-connection state is bounded; NATS subs are lightweight
Frontend reconnect during navigation	Singleton WS connection survives React re-renders; only unmount kills it
Breaking change for any external SSE consumers	Keep /api/chimps REST endpoint; SSE was internal-only

Estimated Scope

• ~8 files modified, ~3 files created, ~3 files deleted
• Touches only packages/shared (protocol) and packages/dashboard (server + frontend)
• No changes to ringmaster, usher, bullhorn, or chimp packages
• NATS subject topology unchanged — only the dashboard↔browser transport changes