@schlessera/brain-ui-react 0.34.1 → 0.36.0

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

@@ -1,21 +1,21 @@
-import { useEffect } from "react";
+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,10 +35,13 @@ 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 root = useBrainUiRoot();
+  const [successfulProbe, setSuccessfulProbe] = useState({ root, count: 0 });
+
   useEffect(() => {
     let timer: ReturnType<typeof setTimeout> | undefined;
     let inFlight = false;
@@ -52,18 +55,21 @@ 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") {
+            setSuccessfulProbe((previous) => ({ root, count: previous.root === root ? previous.count + 1 : 1 }));
+          }
         }
       } finally {
         inFlight = false;
@@ -73,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(
@@ -88,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();
     };
@@ -100,9 +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 successfulProbe.root === root ? successfulProbe.count : 0;
 }

package/src/hooks/use-websocket.ts

@@ -1,432 +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 {
-  if (msg.type !== "session_info" && msg.type !== "result") return false;
-  const pending = state.pendingDraftId;
-  if (msg.type === "session_info" && 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>;

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

@@ -1,4 +1,4 @@
-import { activeChat, useChatStore } from "../../stores/chat-store.js";
+import { activeChat } from "../../stores/chat-state.js";
 import type { ChatMessage, ToolCall, AskUserExchange } from "../../stores/chat-store.js";
 import type {
   AskUserQuestion,
@@ -18,6 +18,7 @@ type ChatFrame =
   | "tool_result"
   | "ask_user_request"
   | "result"
+  | "message_blocks"
   | "session_history"
   | "status";
 
@@ -56,10 +57,20 @@ function convertHistoryMessage(msg: SessionHistoryMessage): ChatMessage {
     timestamp: Date.now(),
     ...(askUserExchanges ? { askUserExchanges } : {}),
     ...(msg.attachmentCount ? { attachmentCount: msg.attachmentCount } : {}),
+    ...(msg.blocks && msg.blocks.length > 0 ? { blocks: msg.blocks } : {}),
   };
 }
 
-/** Rebuild ask_user exchanges from a resumed message's tool calls. */
+/**
+ * Rebuild ask_user exchanges from a resumed message's tool calls.
+ *
+ * The persisted tool output is `{answers, annotations}` (Claude) or the bare
+ * answers map (pi). Neither carries whether the answer was chosen from the
+ * options or typed into the composer (`AskUserExchange.typed`), nor when it
+ * was given, so a resumed exchange comes back as a plain `answered` one with
+ * no time. Carrying `typed` would mean widening the `ask_user` response —
+ * a `CONTRACT:` change — not guessing here.
+ */
 function isBareAnswersMap(v: unknown): v is Record<string, string> {
   return (
     !!v &&
@@ -117,14 +128,16 @@ export const chatFrameHandlers = {
     // Opening the bubble stays immediate — the first token should show a
     // message starting, and every later delta needs isStreaming to be true.
     if (!context.buffer()?.isStreaming) {
-      context.state.startAssistantMessage(context.key);
+      context.state.startAssistantMessage(context.key, context.frameTurnId);
     }
+    if (context.frameTurnId) context.state.stampTurn(context.key, context.frameTurnId);
     context.enqueueDelta(context.key, "text", msg.text);
   },
   thinking_delta: (msg, context) => {
     if (!context.buffer()?.isStreaming) {
-      context.state.startAssistantMessage(context.key);
+      context.state.startAssistantMessage(context.key, context.frameTurnId);
     }
+    if (context.frameTurnId) context.state.stampTurn(context.key, context.frameTurnId);
     context.enqueueDelta(context.key, "thinking", msg.text);
   },
   tool_use_start: (msg, context) => {
@@ -169,6 +182,14 @@ export const chatFrameHandlers = {
     context.state.finishAssistantMessage(context.key);
     context.resyncIfNeeded(msg.sessionId);
   },
+  message_blocks: (msg, context) => {
+    // Arrives after `result`, for the turn's own assistant message. An
+    // answer never waits on it: the markdown is already on screen, and this
+    // only says which spans to draw as blocks (D42). Targeted by the turn
+    // the frame belongs to: a queued follow-up may already have opened a
+    // newer assistant message by the time the pass returns.
+    context.state.setMessageBlocks(context.key, msg.blocks, context.frameTurnId);
+  },
   session_history: (msg, context) => {
     const converted = msg.messages.map(convertHistoryMessage);
     if (msg.append) {
@@ -194,7 +215,7 @@ export const chatFrameHandlers = {
     // THIS frame's session: a background session settling must not consume
     // (or trigger) the active session's pending resync.
     if (msg.status === "idle" || msg.status === "cancelled") {
-      const current = useChatStore.getState();
+      const current = context.stores.chat.getState();
       context.resyncIfNeeded(context.frameSessionId ?? current.activeSessionId);
       // Cold load: we hold a stored session id but an empty transcript and
       // the server is idle (so it sent no snapshot). Fetch the history.

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

@@ -1,16 +1,14 @@
-import { useActivityStore } from "../../stores/activity-store.js";
-import { useConnectionStore } from "../../stores/connection-store.js";
 import type { ServerMessageHandlerMap } from "./types.js";
 
 type ConnectionFrame = "server_hello" | "location_request" | "error";
 
 export const connectionFrameHandlers = {
   server_hello: (msg, context) => {
-    useActivityStore.getState().setSupported(msg.capabilities?.activity === true);
+    context.stores.activity.getState().setSupported(msg.capabilities?.activity === true);
     // A new hello means a new connection: server-side subscriptions are gone.
-    useActivityStore.getState().resetSubscriptions();
+    context.stores.activity.getState().resetSubscriptions();
     // …and view-owned subscriptions (the Activity index) must re-send too.
-    useActivityStore.getState().bumpConnectionEpoch();
+    context.stores.activity.getState().bumpConnectionEpoch();
     context.ensureActivitySubscription(context.state.activeSessionId);
   },
   location_request: (msg, context) => {
@@ -21,7 +19,7 @@ export const connectionFrameHandlers = {
     // message is streaming, and this used to be the whole handler — so an
     // error arriving between turns (a rejected frame, a failed resume) was
     // dropped as silently on the client as it was on the server.
-    useConnectionStore.getState().reportError(msg.code, msg.message);
+    context.stores.connection.getState().reportError(msg.code, msg.message);
     if (context.buffer()?.isStreaming) {
       context.state.appendText(context.key, `\n\n**Error:** ${msg.message}`);
       context.state.finishAssistantMessage(context.key);