@schlessera/brain-ui-react 0.35.0 → 0.37.0

package/src/hooks/use-share-intake.ts

@@ -1,4 +1,5 @@
-import { useCallback, useEffect } from "react";
+import { useBrainUiRoot } from "../root-context.js";
+import { useCallback, useEffect, useRef } from "react";
 import {
   SHARE_STASH_TTL_MS,
   type ShareIntakeResult,
@@ -9,8 +10,6 @@ import {
   readShareLaunchParams,
   type StoredShare,
 } from "@schlessera/brain-ui-sdk/share-target";
-import { useChatStore } from "../stores/chat-store.js";
-import { useConnectionStore } from "../stores/connection-store.js";
 import { useShareStore } from "../stores/share-store.js";
 import { detectClientEnvironment } from "../lib/client-environment.js";
 import {
@@ -23,7 +22,6 @@ import {
   shareImagesToAttachments,
   uploadShare,
 } from "../lib/share-intake.js";
-import { sendClientMessage } from "./use-websocket.js";
 
 /**
  * Pick up shares the service worker stashed, and file them once confirmed.
@@ -66,6 +64,13 @@ export function useShareIntake(): {
   confirm: (record: StoredShare) => Promise<void>;
   dismiss: (record: StoredShare) => void;
 } {
+  const root = useBrainUiRoot();
+  const lifetime = useRef(new AbortController());
+  useEffect(() => {
+    const controller = new AbortController();
+    lifetime.current = controller;
+    return () => controller.abort();
+  }, [root]);
   const enqueue = useShareStore((s) => s.enqueue);
   const remove = useShareStore((s) => s.remove);
   const setBusy = useShareStore((s) => s.setBusy);
@@ -82,8 +87,12 @@ export function useShareIntake(): {
       try {
         for (const id of claimed) {
           const record = await store.take(id);
+          if (disposed) {
+            if (record) await store.put(record);
+            return;
+          }
           dropShareClaim(id);
-          if (record && !disposed) enqueue(record);
+          if (record) enqueue(record);
         }
 
         // Orphans: a share whose landing page never ran (the app opened
@@ -91,8 +100,13 @@ export function useShareIntake(): {
         // parameter was read) leaves a record nobody holds the id for. The
         // stash is the only place it still exists.
         for (const record of await store.list()) {
+          if (disposed) return;
           const taken = await store.take(record.id);
-          if (taken && !disposed) enqueue(taken);
+          if (disposed) {
+            if (taken) await store.put(taken);
+            return;
+          }
+          if (taken) enqueue(taken);
         }
 
         // The service worker's own prune runs after a stash, so it can never
@@ -112,9 +126,10 @@ export function useShareIntake(): {
 
   const confirm = useCallback(
     async (record: StoredShare) => {
-      const share = useShareStore.getState();
+      const signal = lifetime.current.signal;
+      const share = root.stores.share.getState();
       if (share.busy) return;
-      if (useConnectionStore.getState().wsStatus !== "connected") {
+      if (root.stores.connection.getState().wsStatus !== "connected") {
         setError("Not connected yet — this will work as soon as the app reconnects.");
         return;
       }
@@ -123,17 +138,22 @@ export function useShareIntake(): {
       setError(null);
       setNotes([]);
       try {
-        const outcome = await uploadShare(record);
+        const outcome = await uploadShare(record, root.request, root.apiBase(), signal);
+        if (signal.aborted) return;
         if (!outcome.ok) {
           setError(describeUploadError(outcome));
           return;
         }
 
         const { attachments, errors } = await shareImagesToAttachments(record.files);
+        if (signal.aborted) {
+          for (const attachment of attachments) URL.revokeObjectURL(attachment.previewUrl);
+          return;
+        }
         if (errors.length) setNotes(errors);
 
         const prompt = buildSharePrompt(outcome.result as ShareIntakeResult);
-        const chat = useChatStore.getState();
+        const chat = root.stores.chat.getState();
         // Always start from a fresh draft: `activeSessionId` is restored from
         // localStorage at store creation, so on a cold boot it already points
         // at whatever session was last open.
@@ -153,7 +173,7 @@ export function useShareIntake(): {
         );
         chat.startAssistantMessage(null);
 
-        const sent = sendClientMessage({
+        const sent = root.connection.send({
           type: "chat_message",
           text: prompt,
           ...(attachments.length
@@ -172,7 +192,7 @@ export function useShareIntake(): {
         setBusy(false);
       }
     },
-    [remove, setBusy, setError, setNotes]
+    [remove, setBusy, setError, setNotes, root]
   );
 
   const dismiss = useCallback(

package/src/hooks/use-vpn-status.ts

@@ -1,21 +1,21 @@
+import { defaultRoot } from "../default-root.js";
+import { useBrainUiRoot } from "../root-context.js";
+import type { BrainUiRoot } from "../root.js";
 import { useEffect, useState } from "react";
 import {
-  useConnectionStore,
   type VpnStatus,
 } from "../stores/connection-store.js";
-import { getBackendUrl } from "../lib/backend.js";
 
 const POLL_INTERVAL = 15_000;
 // Tighter cadence while offline so recovery is near-immediate
 const OFFLINE_POLL_INTERVAL = 3_000;
 const TIMEOUT_MS = 5_000;
-let activeRecheck: (() => void) | null = null;
 
-async function fetchVpnStatus(): Promise<VpnStatus> {
+async function fetchVpnStatus(root: BrainUiRoot): Promise<VpnStatus> {
   const controller = new AbortController();
   const timeout = setTimeout(() => controller.abort(), TIMEOUT_MS);
   try {
-    const res = await fetch(getBackendUrl("/api/vpn-check"), {
+    const res = await root.request(root.backendUrl("/api/vpn-check"), {
       signal: controller.signal,
     });
 
@@ -35,11 +35,12 @@ async function fetchVpnStatus(): Promise<VpnStatus> {
  * the only writer of vpnStatus and serializes this with its scheduled checks.
  */
 export function recheckVpnStatus(): void {
-  activeRecheck?.();
+  defaultRoot.recheckVpn?.();
 }
 
 export function useVpnStatus() {
-  const [successfulProbeCount, setSuccessfulProbeCount] = useState(0);
+  const root = useBrainUiRoot();
+  const [successfulProbe, setSuccessfulProbe] = useState({ root, count: 0 });
 
   useEffect(() => {
     let timer: ReturnType<typeof setTimeout> | undefined;
@@ -54,20 +55,20 @@ export function useVpnStatus() {
         return;
       }
       inFlight = true;
-      const opensAtStart = useConnectionStore.getState().socketOpens;
+      const opensAtStart = root.stores.connection.getState().socketOpens;
       try {
-        const status = await fetchVpnStatus();
+        const status = await fetchVpnStatus(root);
         // A socket that reached `open` while this probe was in flight is newer
         // reachability evidence than a failure the probe was started for.
         // Counting opens rather than sampling wsStatus twice also catches a
         // drop and a replacement open inside one probe, where both samples
         // read "connected".
         const socketOpenedDuringCheck =
-          useConnectionStore.getState().socketOpens !== opensAtStart;
+          root.stores.connection.getState().socketOpens !== opensAtStart;
         if (!disposed && !(socketOpenedDuringCheck && status !== "connected")) {
-          useConnectionStore.getState().setVpnStatus(status);
+          root.stores.connection.getState().setVpnStatus(status);
           if (status === "connected") {
-            setSuccessfulProbeCount((count) => count + 1);
+            setSuccessfulProbe((previous) => ({ root, count: previous.root === root ? previous.count + 1 : 1 }));
           }
         }
       } finally {
@@ -78,7 +79,7 @@ export function useVpnStatus() {
             recheckRequested = false;
             void check();
           } else {
-            const status = useConnectionStore.getState().vpnStatus;
+            const status = root.stores.connection.getState().vpnStatus;
             // Auth required but no valid session is still a reachable server.
             const reachable = status === "connected" || status === "unauthorized";
             timer = setTimeout(
@@ -93,7 +94,7 @@ export function useVpnStatus() {
     // Re-check immediately when the network returns or the PWA is foregrounded
     // instead of waiting out the poll interval.
     const recheckNow = () => void check();
-    activeRecheck = recheckNow;
+    const removeRecheck = root.registerVpnRecheck(recheckNow);
     const onVisible = () => {
       if (document.visibilityState === "visible") void check();
     };
@@ -105,11 +106,11 @@ export function useVpnStatus() {
     return () => {
       disposed = true;
       clearTimeout(timer);
-      if (activeRecheck === recheckNow) activeRecheck = null;
+      removeRecheck();
       window.removeEventListener("online", recheckNow);
       document.removeEventListener("visibilitychange", onVisible);
     };
-  }, []);
+  }, [root]);
 
-  return successfulProbeCount;
+  return successfulProbe.root === root ? successfulProbe.count : 0;
 }

package/src/hooks/use-websocket.ts

@@ -1,434 +1,18 @@
-import { useEffect, useRef, useCallback } from "react";
-import {
-  BrainUiClient,
-  type WebSocketClose,
-} from "@schlessera/brain-ui-sdk/client";
-import { useConnectionStore } from "../stores/connection-store.js";
-import { useChatStore, activeChat, type ChatKey } from "../stores/chat-store.js";
-import { useActivityStore, loadSessionActivityHistory } from "../stores/activity-store.js";
-import { getWsUrl } from "../lib/backend.js";
-import type {
-  ServerMessage,
-  ClientMessage,
-  ServerLocationRequest,
-} from "@schlessera/brain-ui-sdk/protocol";
-import { dispatchServerMessage } from "./websocket-handlers/index.js";
-import { runStateForFrame } from "./websocket-handlers/chat.js";
-import { recheckVpnStatus } from "./use-vpn-status.js";
-import { REFUSAL_ATTEMPTS } from "../components/connectivity/connection-state.js";
-
+import { useEffect, useCallback } from "react";
+import type { ClientMessage, ServerMessage } from "@schlessera/brain-ui-sdk/protocol";
+import { defaultRoot } from "../default-root.js";
+import { useBrainUiRoot } from "../root-context.js";
 export { runStateForFrame } from "./websocket-handlers/chat.js";
 
-/**
- * Does this frame announce the identity of the conversation THIS client just
- * started?
- *
- * The old test was `session_info || result` for any unknown session, which
- * adopted whichever arrived first — an older background turn finishing, or
- * another client's new session, would capture the user's draft and bind it to
- * a transcript they never wrote. `draftId` is minted per draft turn and echoed
- * on `session_info`, so a match is proof.
- *
- * A server too old to echo it sends no `draftId`, and the pre-existing
- * behaviour applies rather than the draft never binding at all: correctness
- * where the information exists, compatibility where it does not.
- */
-function isOurDraftAnnouncement(
-  state: ReturnType<typeof useChatStore.getState>,
-  msg: ServerMessage
-): boolean {
-  // A terminal result has no draft correlation. Another client's turn may
-  // finish before our session_info arrives, so it can never claim our draft.
-  if (msg.type !== "session_info") return false;
-  const pending = state.pendingDraftId;
-  if (msg.draftId) return msg.draftId === pending;
-  // No echo to compare against.
-  return true;
-}
-
-/**
- * Lazily open a session-scoped activity subscription (live subagent rows,
- * server-stamped tool timing). Idempotent per connection; the subscribed set
- * resets when a new socket's server_hello arrives.
- */
-function ensureActivitySubscription(sessionId: string | null | undefined): void {
-  if (!sessionId) return;
-  const activity = useActivityStore.getState();
-  if (!activity.supported || activity.subscribed[sessionId]) return;
-  activity.markSubscribed(sessionId);
-  wsClient?.send({ type: "activity_subscribe", view: "session", sessionId });
-  // The subscription snapshot carries OPEN runs; finished runs (history
-  // timing, past fan-outs) come from the REST side.
-  void loadSessionActivityHistory(sessionId);
-}
-
-/**
- * Text and thinking deltas waiting to be applied, per buffer.
- *
- * A model streams tokens far faster than the display refreshes, and each socket
- * frame arrives in its own macrotask, so React cannot batch them: one token was
- * one store write, one render and one markdown re-parse of the growing message.
- * Coalescing a frame's worth of tokens into a single write costs nothing
- * visually — the screen could not have shown the intermediate states anyway.
- *
- * Consecutive same-kind chunks merge, so a burst of text becomes ONE appendText
- * rather than one per token. Kind changes start a new chunk, which is what keeps
- * interleaved thinking and text in their true chronological order.
- */
-type PendingDelta = { kind: "text" | "thinking"; text: string };
-const pendingDeltas = new Map<ChatKey, PendingDelta[]>();
-let flushScheduled = false;
-
-/**
- * Apply every buffered delta, in arrival order, and clear the buffer.
- *
- * Exported so tests can force the frame rather than wait for one; production
- * code never needs to call it, because every path that reads the transcript
- * flushes first.
- */
-export function flushChatDeltas(): void {
-  flushDeltas();
-}
-
-function flushDeltas(): void {
-  flushScheduled = false;
-  if (pendingDeltas.size === 0) return;
-  const batches = [...pendingDeltas.entries()];
-  pendingDeltas.clear();
-  const state = useChatStore.getState();
-  for (const [key, chunks] of batches) {
-    for (const chunk of chunks) {
-      if (chunk.kind === "text") state.appendText(key, chunk.text);
-      else state.appendThinking(key, chunk.text);
-    }
-  }
-}
-
-function enqueueDelta(key: ChatKey, kind: PendingDelta["kind"], text: string): void {
-  const chunks = pendingDeltas.get(key);
-  if (!chunks) {
-    pendingDeltas.set(key, [{ kind, text }]);
-  } else {
-    const last = chunks[chunks.length - 1];
-    if (last.kind === kind) last.text += text;
-    else chunks.push({ kind, text });
-  }
-  // Outside a browser (unit tests, SSR) there is no frame to wait for, so apply
-  // straight away and keep the handler's behaviour synchronous.
-  if (typeof requestAnimationFrame !== "function") {
-    flushDeltas();
-    return;
-  }
-  if (!flushScheduled) {
-    flushScheduled = true;
-    requestAnimationFrame(flushDeltas);
-  }
-}
-
-export function handleServerMessage(msg: ServerMessage) {
-  // Every frame that is not itself a delta must see the transcript fully
-  // applied: the demux below reads buffer state, and the store records parts in
-  // arrival order, so a tool call landing ahead of buffered text would reorder
-  // the message.
-  if (msg.type !== "text_delta" && msg.type !== "thinking_delta") flushDeltas();
-
-  const state = useChatStore.getState();
-
-  // Activity stream frames feed their own store and never touch chat state.
-  if (
-    msg.type === "server_hello" ||
-    msg.type === "activity_snapshot" ||
-    msg.type === "activity_delta"
-  ) {
-    const key = state.activeSessionId;
-    dispatchServerMessage(msg, {
-      state,
-      frameSessionId: undefined,
-      key,
-      buffer: () => (key === null ? state.draft : state.buffers[key]),
-      enqueueDelta,
-      ensureActivitySubscription,
-      requestBrowserLocation,
-      resyncIfNeeded,
-      coldResumeIfNeeded,
-      markHistoryReplaced,
-    });
-    return;
-  }
-
-  // Parallel-session demux (per-session buffers). Every frame is scoped by
-  // sessionId; a scoped frame lands in ITS session's buffer, so a background
-  // session's transcript keeps accumulating while another one is in view.
-  // Frames without a sessionId (legacy single-session servers, or a new
-  // session's pre-binding frames) apply to the buffer in view.
-  const frameSessionId = (msg as { sessionId?: string }).sessionId;
-  if (frameSessionId) {
-    // `detail` on a queued status is the host's queue-pressure note; it rides
-    // along so the session's badge can show WHY it is queued deep.
-    const queueNote =
-      msg.type === "status" && msg.status === "queued" ? msg.detail : undefined;
-    state.setRunState(frameSessionId, runStateForFrame(msg), queueNote);
-  }
-
-  // Resolve the target buffer key for this frame.
-  let key: ChatKey;
-  if (!frameSessionId) {
-    key = state.activeSessionId; // null = the draft view
-  } else if (state.buffers[frameSessionId]) {
-    key = frameSessionId;
-  } else if (state.draft && isOurDraftAnnouncement(state, msg)) {
-    // A draft run just got its server identity: adopt the draft buffer.
-    state.bindDraftSession(frameSessionId);
-    key = frameSessionId;
-  } else if (msg.type === "session_info" && frameSessionId === state.activeSessionId) {
-    // The ACTIVE session announced itself but has no buffer yet — the cold
-    // reattach case (PWA relaunch mid-turn; the server may skip the history
-    // frame when a first turn has no history yet). Materialize the buffer so
-    // this session_info and every following delta land instead of vanishing.
-    state.setActiveSession(frameSessionId);
-    key = frameSessionId;
-  } else if (msg.type === "session_history") {
-    // History replay may create a buffer (cold resume / reattach).
-    key = frameSessionId;
-  } else {
-    // Scoped frame for a session with no buffer (never opened, or evicted):
-    // the badge above is all we track. The transcript heals via
-    // session_resume when the user opens it.
-    return;
-  }
-
-  const buffer = () => {
-    const s = useChatStore.getState();
-    return key === null ? s.draft : s.buffers[key];
-  };
-
-  dispatchServerMessage(msg, {
-    state,
-    frameSessionId,
-    key,
-    buffer,
-    enqueueDelta,
-    ensureActivitySubscription,
-    requestBrowserLocation,
-    resyncIfNeeded,
-    coldResumeIfNeeded,
-    markHistoryReplaced,
-  });
-}
-
-/**
- * The server asked the browser for its location (the agent called
- * `get_current_location`). Read `navigator.geolocation` and reply over the same
- * requestId. The browser shows its own permission prompt on first use; a denial,
- * timeout, or unavailable position comes back as a `location_error` that the
- * server turns into a tool error for the agent.
- */
-function requestBrowserLocation(msg: ServerLocationRequest) {
-  const { requestId } = msg;
-  if (typeof navigator === "undefined" || !navigator.geolocation) {
-    wsClient?.send({
-      type: "location_error",
-      requestId,
-      code: 0,
-      message: "Geolocation is not available in this browser.",
-    });
-    return;
-  }
-  const opts = msg.options ?? {};
-  navigator.geolocation.getCurrentPosition(
-    (pos) => {
-      wsClient?.send({
-        type: "location_response",
-        requestId,
-        coords: {
-          latitude: pos.coords.latitude,
-          longitude: pos.coords.longitude,
-          accuracy: pos.coords.accuracy,
-          altitude: pos.coords.altitude,
-          altitudeAccuracy: pos.coords.altitudeAccuracy,
-          heading: pos.coords.heading,
-          speed: pos.coords.speed,
-        },
-        timestamp: pos.timestamp,
-      });
-    },
-    (err) => {
-      wsClient?.send({
-        type: "location_error",
-        requestId,
-        code: err.code,
-        message: err.message || "Failed to retrieve location.",
-      });
-    },
-    {
-      enableHighAccuracy: opts.enableHighAccuracy ?? false,
-      timeout: opts.timeoutMs ?? 15000,
-      maximumAge: opts.maximumAgeMs ?? 60000,
-    }
-  );
-}
-
-// Reconnection bookkeeping: after a WS drop, the rendered transcript may have
-// gaps (deltas streamed while offline are lost). We heal by replaying the
-// authoritative session history — immediately if nothing is streaming, or
-// once the in-flight turn settles (result / idle status). The pending resync
-// is keyed to the session that needs healing (the one that was active at the
-// reconnect), so a background session settling first cannot consume it.
-let wasDisconnected = false;
-let resyncSessionId: string | null = null;
-
-function handleSocketClose(close: WebSocketClose): void {
-  const connection = useConnectionStore.getState();
-  connection.recordWsClose(close.opened, close.code);
-  if (
-    !close.opened &&
-    useConnectionStore.getState().handshakeFailures >= REFUSAL_ATTEMPTS
-  ) {
-    recheckVpnStatus();
-  }
-}
-
-function resyncIfNeeded(sessionId: string | null) {
-  if (!resyncSessionId || !sessionId || sessionId !== resyncSessionId) return;
-  resyncSessionId = null;
-  wsClient?.send({ type: "session_resume", sessionId });
-}
-
-// A cold page / PWA relaunch restores only the stored session id, not the
-// transcript. The server pushes a snapshot on connect only while a session is
-// actively running; when it's idle (the common "agent finished while the phone
-// slept" case) nothing arrives, so the client must fetch the finished
-// session's history itself. Guarded to fire at most once per session id per
-// connection — an empty session would otherwise re-trigger on every idle
-// status and loop.
-let coldResumedSessionId: string | null = null;
-
-function markHistoryReplaced(key: ChatKey): void {
-  if (key !== null && key === resyncSessionId) resyncSessionId = null;
-  coldResumedSessionId = key;
-}
-
-function coldResumeIfNeeded(sessionId: string | null, messageCount: number) {
-  if (!sessionId || messageCount > 0) return;
-  if (coldResumedSessionId === sessionId) return;
-  coldResumedSessionId = sessionId;
-  wsClient?.send({ type: "session_resume", sessionId });
-}
-
-function handleStatusChange(status: "connecting" | "connected" | "disconnected") {
-  useConnectionStore.getState().setWsStatus(status);
-
-  if (status === "connected") {
-    useConnectionStore.getState().noteSocketOpen();
-  }
-
-  // A status change can be followed by a history replay that rewrites the
-  // buffer, so land whatever is still buffered before anything reads it. In a
-  // hidden tab requestAnimationFrame does not run at all, which is exactly the
-  // case where deltas could otherwise sit unapplied across a reconnect.
-  flushDeltas();
-
-  if (status === "disconnected") {
-    wasDisconnected = true;
-    // Allow a fresh cold-resume attempt after we reconnect.
-    coldResumedSessionId = null;
-  } else if (status === "connected" && wasDisconnected) {
-    wasDisconnected = false;
-    const chat = useChatStore.getState();
-    const active = activeChat(chat);
-    if (chat.activeSessionId && active.messages.length > 0) {
-      resyncSessionId = chat.activeSessionId;
-      // Not mid-stream: replay right away. Mid-stream: the flag holds until
-      // the running turn finishes (result or idle status).
-      if (!active.isStreaming) {
-        resyncIfNeeded(chat.activeSessionId);
-      }
-    }
-  }
-}
-
-// Singleton client - survives React re-renders
-let wsClient: BrainUiClient | null = null;
-
-/**
- * Send on the live socket from outside a component.
- *
- * `useWebSocket()` OWNS the socket — its construction guard is a per-instance
- * ref, so a second caller would build a second client, overwrite this
- * singleton, and orphan the first; unmounting either one would then close the
- * socket for both. Anything that needs to send but not to own (the share
- * intake) goes through here instead of calling the hook again.
- *
- * Returns false when there is no open socket: `send` drops silently in that
- * case and a caller that just staged an upload needs to know.
- */
-export function sendClientMessage(msg: ClientMessage): boolean {
-  if (!wsClient) return false;
-  if (useConnectionStore.getState().wsStatus !== "connected") return false;
-  wsClient.send(msg);
-  return true;
-}
-
-/** Retry the singleton socket immediately, skipping any remaining backoff. */
-export function reconnectWebSocketNow(): void {
-  wsClient?.reconnectNow();
-}
+/** Default-app entry points for callers outside React. */
+export function handleServerMessage(message: ServerMessage) { defaultRoot.connection.handleServerMessage(message); }
+export function flushChatDeltas() { defaultRoot.connection.flushChatDeltas(); }
+export function sendClientMessage(message: ClientMessage) { return defaultRoot.connection.send(message); }
+export function reconnectWebSocketNow() { defaultRoot.connection.reconnectNow(); }
 
 export function useWebSocket() {
-  const initialized = useRef(false);
-
-  useEffect(() => {
-    if (initialized.current) return;
-    initialized.current = true;
-    let disposed = false;
-    useConnectionStore.getState().noteSocketOpen();
-
-    wsClient = new BrainUiClient({
-      url: getWsUrl(),
-      // One handler with a shared preamble, rather than sixteen copies of the
-      // session-buffer demux — see handleServerMessage.
-      handlers: { onAny: handleServerMessage },
-      onStatusChange: handleStatusChange,
-      onClose: (close) => {
-        if (!disposed) handleSocketClose(close);
-      },
-      // A frame the SDK refused. Surfaced rather than logged into a console
-      // nobody is attached to; the server-side counterpart is the
-      // ws.frames.dropped counter.
-      onProtocolError: (err) => {
-        useConnectionStore
-          .getState()
-          .reportError(
-            "PROTOCOL_ERROR",
-            err.frameType
-              ? `Dropped a ${err.frameType} frame: ${err.detail}`
-              : `Dropped an unreadable frame: ${err.detail}`
-          );
-      },
-    });
-    wsClient.connect();
-
-    // Skip the exponential backoff when the network demonstrably returns.
-    const reconnectNow = () => wsClient?.reconnectNow();
-    const onVisible = () => {
-      if (document.visibilityState === "visible") wsClient?.reconnectNow();
-    };
-    window.addEventListener("online", reconnectNow);
-    document.addEventListener("visibilitychange", onVisible);
-
-    return () => {
-      disposed = true;
-      window.removeEventListener("online", reconnectNow);
-      document.removeEventListener("visibilitychange", onVisible);
-      wsClient?.close();
-      wsClient = null;
-      initialized.current = false;
-    };
-  }, []);
-
-  const send = useCallback((msg: ClientMessage) => {
-    wsClient?.send(msg);
-  }, []);
-
+  const root = useBrainUiRoot();
+  useEffect(() => root.connection.connect(), [root]);
+  const send = useCallback((message: ClientMessage) => root.connection.send(message), [root]);
   return { send };
 }

package/src/hooks/websocket-handlers/activity.ts

@@ -1,13 +1,12 @@
-import { useActivityStore } from "../../stores/activity-store.js";
 import type { ServerMessageHandlerMap } from "./types.js";
 
 type ActivityFrame = "activity_snapshot" | "activity_delta";
 
 export const activityFrameHandlers = {
-  activity_snapshot: (msg) => {
-    useActivityStore.getState().applySnapshot(msg);
+  activity_snapshot: (msg, context) => {
+    context.stores.activity.getState().applySnapshot(msg);
   },
-  activity_delta: (msg) => {
-    useActivityStore.getState().applyDelta(msg);
+  activity_delta: (msg, context) => {
+    context.stores.activity.getState().applyDelta(msg);
   },
 } satisfies ServerMessageHandlerMap<ActivityFrame>;