@schlessera/brain-ui-react 0.33.1 → 0.34.1

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

@@ -0,0 +1,59 @@
+import type { ServerMessage } from "@schlessera/brain-ui-sdk/protocol";
+import { activityFrameHandlers } from "./activity.js";
+import { chatFrameHandlers } from "./chat.js";
+import { connectionFrameHandlers } from "./connection.js";
+import { maskFrameHandlers } from "./mask.js";
+import { providerFrameHandlers } from "./provider.js";
+import type {
+  DispatchContext,
+  ServerMessageHandlerMap,
+  ServerMessageType,
+} from "./types.js";
+
+type HandlerKeys<T> = T extends unknown ? keyof T : never;
+type ClaimedFrame = HandlerKeys<
+  | typeof chatFrameHandlers
+  | typeof activityFrameHandlers
+  | typeof maskFrameHandlers
+  | typeof providerFrameHandlers
+  | typeof connectionFrameHandlers
+>;
+type DuplicateFrame =
+  | Extract<
+      keyof typeof chatFrameHandlers,
+      | keyof typeof activityFrameHandlers
+      | keyof typeof maskFrameHandlers
+      | keyof typeof providerFrameHandlers
+      | keyof typeof connectionFrameHandlers
+    >
+  | Extract<
+      keyof typeof activityFrameHandlers,
+      | keyof typeof maskFrameHandlers
+      | keyof typeof providerFrameHandlers
+      | keyof typeof connectionFrameHandlers
+    >
+  | Extract<
+      keyof typeof maskFrameHandlers,
+      keyof typeof providerFrameHandlers | keyof typeof connectionFrameHandlers
+    >
+  | Extract<keyof typeof providerFrameHandlers, keyof typeof connectionFrameHandlers>;
+
+type Assert<T extends true> = T;
+type Equal<A, B> = [A, B] extends [B, A] ? true : false;
+type _EveryFrameClaimed = Assert<Equal<ClaimedFrame, ServerMessageType>>;
+type _NoFrameClaimedTwice = Assert<Equal<DuplicateFrame, never>>;
+
+const serverMessageHandlers = {
+  ...chatFrameHandlers,
+  ...activityFrameHandlers,
+  ...maskFrameHandlers,
+  ...providerFrameHandlers,
+  ...connectionFrameHandlers,
+} satisfies ServerMessageHandlerMap<ServerMessageType>;
+
+export function dispatchServerMessage(msg: ServerMessage, context: DispatchContext): void {
+  // The mapped type preserves the message/handler correlation. TypeScript
+  // loses that correlation at an indexed union access, so restore it only at
+  // this single composition boundary.
+  serverMessageHandlers[msg.type](msg as never, context);
+}

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

@@ -0,0 +1,17 @@
+import { useMaskStore } from "../../stores/mask-store.js";
+import type { ServerMessageHandlerMap } from "./types.js";
+
+type MaskFrame = "mask_request";
+
+export const maskFrameHandlers = {
+  mask_request: (msg) => {
+    // Opens the editor; the answer travels back from the component, because
+    // only the user can say which part of the picture they meant.
+    useMaskStore.getState().open({
+      requestId: msg.requestId,
+      imagePath: msg.imagePath,
+      ...(msg.instruction ? { instruction: msg.instruction } : {}),
+      ...(msg.turnId ? { turnId: msg.turnId } : {}),
+    });
+  },
+} satisfies ServerMessageHandlerMap<MaskFrame>;

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

@@ -0,0 +1,35 @@
+import { useChatStore } from "../../stores/chat-store.js";
+import { useProviderStore } from "../../stores/provider-store.js";
+import type { ServerMessageHandlerMap } from "./types.js";
+
+type ProviderFrame = "session_info";
+
+export const providerFrameHandlers = {
+  session_info: (msg, context) => {
+    context.ensureActivitySubscription(msg.sessionId);
+    // Record backend ownership for renderer scoping — for ANY session, since
+    // background sessions keep their own transcript buffers. Older servers
+    // omit backendId; derive it from the pinned profile when still possible
+    // (fails only for since-hidden profiles, which then use the default).
+    const ownerBackendId =
+      msg.backendId ??
+      useProviderStore
+        .getState()
+        .available.find((p) => p.id === msg.providerId)?.backendId;
+    if (ownerBackendId) {
+      useChatStore.getState().setSessionBackend(msg.sessionId, ownerBackendId);
+    }
+    // bindDraftSession in the dispatcher handled draft adoption; an info frame
+    // may still re-pin the provider picker when it concerns the session in view.
+    const current = useChatStore.getState();
+    if (context.frameSessionId && context.frameSessionId === current.activeSessionId) {
+      useProviderStore.getState().setPinned(msg.providerId ?? null);
+      // Reattachment on a fresh connection: enter streaming mode so deltas
+      // append to a live assistant message instead of vanishing.
+      const chat = current.buffers[context.frameSessionId];
+      if (chat && !chat.isStreaming && chat.messages.length === 0) {
+        current.startAssistantMessage(context.frameSessionId);
+      }
+    }
+  },
+} satisfies ServerMessageHandlerMap<ProviderFrame>;

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

@@ -0,0 +1,27 @@
+import type { ServerMessage } from "@schlessera/brain-ui-sdk/protocol";
+import type { ChatKey, SessionChat } from "../../stores/chat-store.js";
+import { useChatStore } from "../../stores/chat-store.js";
+
+export type ServerMessageType = ServerMessage["type"];
+
+export interface DispatchContext {
+  state: ReturnType<typeof useChatStore.getState>;
+  frameSessionId: string | undefined;
+  key: ChatKey;
+  buffer: () => SessionChat | null | undefined;
+  enqueueDelta: (key: ChatKey, kind: "text" | "thinking", text: string) => void;
+  ensureActivitySubscription: (sessionId: string | null | undefined) => void;
+  requestBrowserLocation: (
+    msg: Extract<ServerMessage, { type: "location_request" }>
+  ) => void;
+  resyncIfNeeded: (sessionId: string | null) => void;
+  coldResumeIfNeeded: (sessionId: string | null, messageCount: number) => void;
+  markHistoryReplaced: (key: ChatKey) => void;
+}
+
+export type ServerMessageHandlerMap<T extends ServerMessageType> = {
+  [K in T]: (
+    msg: Extract<ServerMessage, { type: K }>,
+    context: DispatchContext
+  ) => void;
+};

package/src/stores/activity-store-helpers.ts

@@ -0,0 +1,164 @@
+import type { ActivitySpan, ActivitySpanEvent } from "@schlessera/brain-ui-sdk/protocol";
+import type { ActivityState, MirrorMaps } from "./activity-store-types.js";
+
+/** Bound on the mirror: how many TERMINAL runs stay held. Live runs never count. */
+const MAX_TERMINAL_RUNS = 50;
+
+export function rootSpanOf(byId: Record<string, ActivitySpan>): ActivitySpan | undefined {
+  return Object.values(byId).find((span) => span.parentSpanId === undefined);
+}
+
+/** Drop a whole run from the mirror — spans, events, high-water, reverse index. */
+export function evictRun(maps: MirrorMaps, runId: string): void {
+  for (const spanId of Object.keys(maps.spans[runId] ?? {})) {
+    delete maps.events[spanId];
+    delete maps.spanRun[spanId];
+  }
+  delete maps.spans[runId];
+  delete maps.highWater[runId];
+  delete maps.deltaSeq[runId];
+}
+
+/**
+ * Enforce the terminal-run bound in place. Returns whether anything was
+ * evicted. Runs whose root span is still open (or unknown) are never touched.
+ */
+export function pruneTerminalRuns(maps: MirrorMaps): boolean {
+  const terminal: Array<{ runId: string; startedAt: number }> = [];
+  for (const [runId, byId] of Object.entries(maps.spans)) {
+    const root = rootSpanOf(byId);
+    if (root && root.outcome !== undefined) {
+      terminal.push({ runId, startedAt: root.startedAt });
+    }
+  }
+  if (terminal.length <= MAX_TERMINAL_RUNS) return false;
+  terminal.sort((a, b) => a.startedAt - b.startedAt);
+  for (const { runId } of terminal.slice(0, terminal.length - MAX_TERMINAL_RUNS)) {
+    evictRun(maps, runId);
+  }
+  return true;
+}
+
+export function insertEvent(
+  list: ActivitySpanEvent[] | undefined,
+  event: ActivitySpanEvent
+): ActivitySpanEvent[] {
+  const existing = list ?? [];
+  if (existing.some((e) => e.eventIndex === event.eventIndex)) return existing;
+  const next = [...existing, event];
+  next.sort((a, b) => a.eventIndex - b.eventIndex);
+  return next;
+}
+
+// --- Selectors (plain functions over the state, like activeChat) ---
+
+/** All spans of a run, roots first, then by start time. */
+export function runSpans(state: ActivityState, runId: string): ActivitySpan[] {
+  const byId = state.spans[runId];
+  if (!byId) return [];
+  return Object.values(byId).sort(
+    (a, b) =>
+      (a.parentSpanId ? 1 : 0) - (b.parentSpanId ? 1 : 0) || a.startedAt - b.startedAt
+  );
+}
+
+/** Direct children of a span (e.g. the tool calls inside a subagent). */
+export function childSpans(state: ActivityState, parentSpanId: string): ActivitySpan[] {
+  const runId = state.spanRun[parentSpanId];
+  if (!runId) return [];
+  return Object.values(state.spans[runId] ?? {})
+    .filter((s) => s.parentSpanId === parentSpanId)
+    .sort((a, b) => a.startedAt - b.startedAt);
+}
+
+/** The subagent spans spawned by a turn's Agent tool calls, live status included. */
+export function subagentSpans(state: ActivityState, sessionId: string): ActivitySpan[] {
+  const out: ActivitySpan[] = [];
+  for (const byId of Object.values(state.spans)) {
+    for (const span of Object.values(byId)) {
+      if (span.kind === "subagent" && span.sessionId === sessionId) out.push(span);
+    }
+  }
+  return out.sort((a, b) => a.startedAt - b.startedAt);
+}
+
+/** Stable empty result for `eventsFor` — a fresh `[]` per call would defeat
+ *  reference-equality checks in zustand selectors. */
+export const EMPTY_EVENTS: ActivitySpanEvent[] = Object.freeze(
+  [] as ActivitySpanEvent[]
+) as ActivitySpanEvent[];
+
+/** Ordered events of a span; the shared frozen empty array when none. */
+export function eventsFor(state: ActivityState, spanId: string): ActivitySpanEvent[] {
+  return state.events[spanId] ?? EMPTY_EVENTS;
+}
+
+/** Payload event types the tool expander owns — every OTHER type belongs to
+ *  the narrative stream (`narrativeEventsFor`) so nothing recorded is
+ *  rendered nowhere. */
+const TOOL_PAYLOAD_TYPES = new Set(["tool_input", "tool_output"]);
+
+/**
+ * The recorded input/output payload events of a tool span (AE7). Returns a
+ * fresh array per call — subscribe through `useShallow` (like `childSpans`).
+ */
+export function payloadEventsFor(state: ActivityState, spanId: string): ActivitySpanEvent[] {
+  const events = state.events[spanId];
+  if (!events) return EMPTY_EVENTS;
+  const payloads = events.filter((e) => TOOL_PAYLOAD_TYPES.has(e.eventType));
+  return payloads.length > 0 ? payloads : EMPTY_EVENTS;
+}
+
+/**
+ * Everything recorded against a span that is NOT a tool input/output payload:
+ * transcript excerpts, job output, and any span-sink event type a producer
+ * invents. Rendered as the span's narrative so an unknown type degrades to a
+ * labelled block rather than to invisibility.
+ */
+export function narrativeEventsFor(
+  state: ActivityState,
+  spanId: string
+): ActivitySpanEvent[] {
+  const events = state.events[spanId];
+  if (!events) return EMPTY_EVENTS;
+  const rest = events.filter((e) => !TOOL_PAYLOAD_TYPES.has(e.eventType));
+  return rest.length > 0 ? rest : EMPTY_EVENTS;
+}
+
+/** Every event recorded under a run, ordered by time — the run detail's
+ *  narrative stream and the raw-trace dump both read through this. */
+export function runEvents(state: ActivityState, runId: string): ActivitySpanEvent[] {
+  const byId = state.spans[runId];
+  if (!byId) return EMPTY_EVENTS;
+  const out: ActivitySpanEvent[] = [];
+  for (const spanId of Object.keys(byId)) {
+    const events = state.events[spanId];
+    if (events) out.push(...events);
+  }
+  if (out.length === 0) return EMPTY_EVENTS;
+  return out.sort((a, b) => a.ts - b.ts || a.eventIndex - b.eventIndex);
+}
+
+/** The span behind one tool call (span ids ARE toolUseIds), if streamed. */
+export function spanForTool(state: ActivityState, toolUseId: string): ActivitySpan | null {
+  const runId = state.spanRun[toolUseId];
+  return runId ? (state.spans[runId]?.[toolUseId] ?? null) : null;
+}
+
+/**
+ * Server-stamped timing for a tool call. One clock for live and reloaded
+ * views: when the activity stream carries the span, its timestamps win over
+ * client stamps (AE2 — live and revisit must render identical timings).
+ * `waitUntil` marks the approval-wait/execution boundary.
+ */
+export function timingFor(
+  state: ActivityState,
+  toolUseId: string
+): { startedAt: number; endedAt?: number } | null {
+  const span = spanForTool(state, toolUseId);
+  if (!span) return null;
+  return {
+    startedAt: span.waitUntil ?? span.startedAt,
+    ...(span.endedAt !== undefined ? { endedAt: span.endedAt } : {}),
+  };
+}

package/src/stores/activity-store-types.ts

@@ -0,0 +1,55 @@
+import type {
+  ActivitySpan,
+  ActivitySpanEvent,
+  ServerActivityDelta,
+  ServerActivitySnapshot,
+} from "@schlessera/brain-ui-sdk/protocol";
+import type { ActivityIntent } from "../lib/api-client.js";
+
+export interface ActivityState {
+  /** server_hello said this host records activity. */
+  supported: boolean;
+  /** Session ids this connection has an activity subscription for. */
+  subscribed: Record<string, true>;
+  /** Bumped on every new connection so subscribe effects re-send. */
+  connectionEpoch: number;
+  /** Live span state per run, keyed runId -> spanId. */
+  spans: Record<string, Record<string, ActivitySpan>>;
+  /** Ordered events per span (already sorted by eventIndex). */
+  events: Record<string, ActivitySpanEvent[]>;
+  /** Per-run high-water seq from the snapshot; deltas at or below drop. */
+  highWater: Record<string, number>;
+  /**
+   * Newest APPLIED delta seq per run. Kept apart from `highWater` on
+   * purpose: the snapshot staleness guard compares against the max of both
+   * (a snapshot older than an applied delta must not roll it back), while
+   * the delta gate keeps comparing against the snapshot floor alone so
+   * out-of-order delta delivery stays tolerated (insert-if-absent absorbs it).
+   */
+  deltaSeq: Record<string, number>;
+  /** spanId -> runId reverse index, for timing lookups by toolUseId. */
+  spanRun: Record<string, string>;
+
+  /** Unacknowledged notification intents — the guaranteed-tier inbox. */
+  inbox: ActivityIntent[];
+  loadInbox: () => Promise<void>;
+  acknowledgeIntent: (id: number) => Promise<void>;
+  acknowledgeAllIntents: () => Promise<void>;
+  setSupported: (supported: boolean) => void;
+  /** New connection: server-side subscriptions are gone; re-subscribe lazily. */
+  resetSubscriptions: () => void;
+  bumpConnectionEpoch: () => void;
+  markSubscribed: (sessionId: string) => void;
+  applySnapshot: (msg: ServerActivitySnapshot) => void;
+  applyDelta: (msg: ServerActivityDelta) => void;
+  clear: () => void;
+}
+
+/** The five indexes that make up the mirror, mutated together during merges. */
+export interface MirrorMaps {
+  spans: Record<string, Record<string, ActivitySpan>>;
+  events: Record<string, ActivitySpanEvent[]>;
+  highWater: Record<string, number>;
+  deltaSeq: Record<string, number>;
+  spanRun: Record<string, string>;
+}

package/src/stores/activity-store.ts

@@ -1,19 +1,30 @@
 import { create } from "zustand";
-import type {
-  ActivitySpan,
-  ActivitySpanEvent,
-  ServerActivityDelta,
-  ServerActivitySnapshot,
-} from "@schlessera/brain-ui-sdk/protocol";
-import { api, type ActivityIntent } from "../lib/api-client.js";
+import { api } from "../lib/api-client.js";
 import { registerDevHandle } from "../config.js";
+import type { ActivityState, MirrorMaps } from "./activity-store-types.js";
+import {
+  evictRun,
+  insertEvent,
+  pruneTerminalRuns,
+  rootSpanOf,
+} from "./activity-store-helpers.js";
+
+export {
+  EMPTY_EVENTS,
+  childSpans,
+  eventsFor,
+  narrativeEventsFor,
+  payloadEventsFor,
+  runEvents,
+  runSpans,
+  spanForTool,
+  subagentSpans,
+  timingFor,
+} from "./activity-store-helpers.js";
 
 /** How many finished runs' details are fetched when a session opens. */
 const HISTORY_RUN_LIMIT = 10;
 
-/** Bound on the mirror: how many TERMINAL runs stay held. Live runs never count. */
-const MAX_TERMINAL_RUNS = 50;
-
 /**
  * Client mirror of the server's activity record, fed by the snapshot-then-
  * delta stream. The server computes; this store only holds and indexes what
@@ -34,45 +45,6 @@ const MAX_TERMINAL_RUNS = 50;
  * evicted wholesale (spans, events, high-water, reverse index). Live runs
  * are never evicted; evicted history remains fetchable over REST.
  */
-interface ActivityState {
-  /** server_hello said this host records activity. */
-  supported: boolean;
-  /** Session ids this connection has an activity subscription for. */
-  subscribed: Record<string, true>;
-  /** Bumped on every new connection so subscribe effects re-send. */
-  connectionEpoch: number;
-  /** Live span state per run, keyed runId -> spanId. */
-  spans: Record<string, Record<string, ActivitySpan>>;
-  /** Ordered events per span (already sorted by eventIndex). */
-  events: Record<string, ActivitySpanEvent[]>;
-  /** Per-run high-water seq from the snapshot; deltas at or below drop. */
-  highWater: Record<string, number>;
-  /**
-   * Newest APPLIED delta seq per run. Kept apart from `highWater` on
-   * purpose: the snapshot staleness guard compares against the max of both
-   * (a snapshot older than an applied delta must not roll it back), while
-   * the delta gate keeps comparing against the snapshot floor alone so
-   * out-of-order delta delivery stays tolerated (insert-if-absent absorbs it).
-   */
-  deltaSeq: Record<string, number>;
-  /** spanId -> runId reverse index, for timing lookups by toolUseId. */
-  spanRun: Record<string, string>;
-
-  /** Unacknowledged notification intents — the guaranteed-tier inbox. */
-  inbox: ActivityIntent[];
-  loadInbox: () => Promise<void>;
-  acknowledgeIntent: (id: number) => Promise<void>;
-  acknowledgeAllIntents: () => Promise<void>;
-  setSupported: (supported: boolean) => void;
-  /** New connection: server-side subscriptions are gone; re-subscribe lazily. */
-  resetSubscriptions: () => void;
-  bumpConnectionEpoch: () => void;
-  markSubscribed: (sessionId: string) => void;
-  applySnapshot: (msg: ServerActivitySnapshot) => void;
-  applyDelta: (msg: ServerActivityDelta) => void;
-  clear: () => void;
-}
-
 export const useActivityStore = create<ActivityState>((set, get) => ({
   supported: false,
   subscribed: {},
@@ -309,174 +281,6 @@ export async function loadSpanPayloads(spanId: string): Promise<void> {
   }
 }
 
-/** The four indexes that make up the mirror, mutated together during merges. */
-interface MirrorMaps {
-  spans: Record<string, Record<string, ActivitySpan>>;
-  events: Record<string, ActivitySpanEvent[]>;
-  highWater: Record<string, number>;
-  deltaSeq: Record<string, number>;
-  spanRun: Record<string, string>;
-}
-
-function rootSpanOf(byId: Record<string, ActivitySpan>): ActivitySpan | undefined {
-  return Object.values(byId).find((span) => span.parentSpanId === undefined);
-}
-
-/** Drop a whole run from the mirror — spans, events, high-water, reverse index. */
-function evictRun(maps: MirrorMaps, runId: string): void {
-  for (const spanId of Object.keys(maps.spans[runId] ?? {})) {
-    delete maps.events[spanId];
-    delete maps.spanRun[spanId];
-  }
-  delete maps.spans[runId];
-  delete maps.highWater[runId];
-  delete maps.deltaSeq[runId];
-}
-
-/**
- * Enforce the terminal-run bound in place. Returns whether anything was
- * evicted. Runs whose root span is still open (or unknown) are never touched.
- */
-function pruneTerminalRuns(maps: MirrorMaps): boolean {
-  const terminal: Array<{ runId: string; startedAt: number }> = [];
-  for (const [runId, byId] of Object.entries(maps.spans)) {
-    const root = rootSpanOf(byId);
-    if (root && root.outcome !== undefined) {
-      terminal.push({ runId, startedAt: root.startedAt });
-    }
-  }
-  if (terminal.length <= MAX_TERMINAL_RUNS) return false;
-  terminal.sort((a, b) => a.startedAt - b.startedAt);
-  for (const { runId } of terminal.slice(0, terminal.length - MAX_TERMINAL_RUNS)) {
-    evictRun(maps, runId);
-  }
-  return true;
-}
-
-function insertEvent(
-  list: ActivitySpanEvent[] | undefined,
-  event: ActivitySpanEvent
-): ActivitySpanEvent[] {
-  const existing = list ?? [];
-  if (existing.some((e) => e.eventIndex === event.eventIndex)) return existing;
-  const next = [...existing, event];
-  next.sort((a, b) => a.eventIndex - b.eventIndex);
-  return next;
-}
-
-// --- Selectors (plain functions over the state, like activeChat) ---
-
-/** All spans of a run, roots first, then by start time. */
-export function runSpans(state: ActivityState, runId: string): ActivitySpan[] {
-  const byId = state.spans[runId];
-  if (!byId) return [];
-  return Object.values(byId).sort(
-    (a, b) =>
-      (a.parentSpanId ? 1 : 0) - (b.parentSpanId ? 1 : 0) || a.startedAt - b.startedAt
-  );
-}
-
-/** Direct children of a span (e.g. the tool calls inside a subagent). */
-export function childSpans(state: ActivityState, parentSpanId: string): ActivitySpan[] {
-  const runId = state.spanRun[parentSpanId];
-  if (!runId) return [];
-  return Object.values(state.spans[runId] ?? {})
-    .filter((s) => s.parentSpanId === parentSpanId)
-    .sort((a, b) => a.startedAt - b.startedAt);
-}
-
-/** The subagent spans spawned by a turn's Agent tool calls, live status included. */
-export function subagentSpans(state: ActivityState, sessionId: string): ActivitySpan[] {
-  const out: ActivitySpan[] = [];
-  for (const byId of Object.values(state.spans)) {
-    for (const span of Object.values(byId)) {
-      if (span.kind === "subagent" && span.sessionId === sessionId) out.push(span);
-    }
-  }
-  return out.sort((a, b) => a.startedAt - b.startedAt);
-}
-
-/** Stable empty result for `eventsFor` — a fresh `[]` per call would defeat
- *  reference-equality checks in zustand selectors. */
-export const EMPTY_EVENTS: ActivitySpanEvent[] = Object.freeze(
-  [] as ActivitySpanEvent[]
-) as ActivitySpanEvent[];
-
-/** Ordered events of a span; the shared frozen empty array when none. */
-export function eventsFor(state: ActivityState, spanId: string): ActivitySpanEvent[] {
-  return state.events[spanId] ?? EMPTY_EVENTS;
-}
-
-/** Payload event types the tool expander owns — every OTHER type belongs to
- *  the narrative stream (`narrativeEventsFor`) so nothing recorded is
- *  rendered nowhere. */
-const TOOL_PAYLOAD_TYPES = new Set(["tool_input", "tool_output"]);
-
-/**
- * The recorded input/output payload events of a tool span (AE7). Returns a
- * fresh array per call — subscribe through `useShallow` (like `childSpans`).
- */
-export function payloadEventsFor(state: ActivityState, spanId: string): ActivitySpanEvent[] {
-  const events = state.events[spanId];
-  if (!events) return EMPTY_EVENTS;
-  const payloads = events.filter((e) => TOOL_PAYLOAD_TYPES.has(e.eventType));
-  return payloads.length > 0 ? payloads : EMPTY_EVENTS;
-}
-
-/**
- * Everything recorded against a span that is NOT a tool input/output payload:
- * transcript excerpts, job output, and any span-sink event type a producer
- * invents. Rendered as the span's narrative so an unknown type degrades to a
- * labelled block rather than to invisibility.
- */
-export function narrativeEventsFor(
-  state: ActivityState,
-  spanId: string
-): ActivitySpanEvent[] {
-  const events = state.events[spanId];
-  if (!events) return EMPTY_EVENTS;
-  const rest = events.filter((e) => !TOOL_PAYLOAD_TYPES.has(e.eventType));
-  return rest.length > 0 ? rest : EMPTY_EVENTS;
-}
-
-/** Every event recorded under a run, ordered by time — the run detail's
- *  narrative stream and the raw-trace dump both read through this. */
-export function runEvents(state: ActivityState, runId: string): ActivitySpanEvent[] {
-  const byId = state.spans[runId];
-  if (!byId) return EMPTY_EVENTS;
-  const out: ActivitySpanEvent[] = [];
-  for (const spanId of Object.keys(byId)) {
-    const events = state.events[spanId];
-    if (events) out.push(...events);
-  }
-  if (out.length === 0) return EMPTY_EVENTS;
-  return out.sort((a, b) => a.ts - b.ts || a.eventIndex - b.eventIndex);
-}
-
-/** The span behind one tool call (span ids ARE toolUseIds), if streamed. */
-export function spanForTool(state: ActivityState, toolUseId: string): ActivitySpan | null {
-  const runId = state.spanRun[toolUseId];
-  return runId ? (state.spans[runId]?.[toolUseId] ?? null) : null;
-}
-
-/**
- * Server-stamped timing for a tool call. One clock for live and reloaded
- * views: when the activity stream carries the span, its timestamps win over
- * client stamps (AE2 — live and revisit must render identical timings).
- * `waitUntil` marks the approval-wait/execution boundary.
- */
-export function timingFor(
-  state: ActivityState,
-  toolUseId: string
-): { startedAt: number; endedAt?: number } | null {
-  const span = spanForTool(state, toolUseId);
-  if (!span) return null;
-  return {
-    startedAt: span.waitUntil ?? span.startedAt,
-    ...(span.endedAt !== undefined ? { endedAt: span.endedAt } : {}),
-  };
-}
-
 registerDevHandle(() => {
   if (typeof window === "undefined") return;
   (window as unknown as { __activityStore?: typeof useActivityStore }).__activityStore =

package/src/stores/connection-store.ts

@@ -1,7 +1,7 @@
 import { create } from "zustand";
 
-type WsStatus = "connecting" | "connected" | "disconnected";
-type VpnStatus =
+export type WsStatus = "connecting" | "connected" | "disconnected";
+export type VpnStatus =
   | "connected"
   | "forbidden"
   | "unauthorized"
@@ -27,9 +27,21 @@ export interface ConnectionError {
 interface ConnectionState {
   wsStatus: WsStatus;
   vpnStatus: VpnStatus;
+  handshakeFailures: number;
+  lastCloseCode: number | null;
+  /**
+   * Monotonic count of sockets that reached `open`. A probe started before a
+   * socket opened compares this to decide whether its answer is still the
+   * newest evidence — comparing wsStatus twice cannot see a drop and a
+   * replacement open that both land inside one probe.
+   */
+  socketOpens: number;
   lastError: ConnectionError | null;
   setWsStatus: (status: WsStatus) => void;
   setVpnStatus: (status: VpnStatus) => void;
+  recordWsClose: (opened: boolean, code: number) => void;
+  /** A socket reached `open`: the refusal evidence so far is spent. */
+  noteSocketOpen: () => void;
   reportError: (code: string, message: string) => void;
   clearError: () => void;
 }
@@ -37,9 +49,28 @@ interface ConnectionState {
 export const useConnectionStore = create<ConnectionState>((set) => ({
   wsStatus: "disconnected",
   vpnStatus: "checking",
+  handshakeFailures: 0,
+  lastCloseCode: null,
+  socketOpens: 0,
   lastError: null,
   setWsStatus: (wsStatus) => set({ wsStatus }),
   setVpnStatus: (vpnStatus) => set({ vpnStatus }),
+  recordWsClose: (opened, lastCloseCode) =>
+    set((state) => ({
+      lastCloseCode,
+      handshakeFailures: opened
+        ? state.handshakeFailures
+        : state.handshakeFailures + 1,
+    })),
+  noteSocketOpen: () =>
+    set((state) => ({
+      handshakeFailures: 0,
+      // Cleared here, not on the next close: the server accepts the upgrade and
+      // THEN closes with 4008 when it is at its connection cap, so the code
+      // must survive an open to be read, and must not outlive the next one.
+      lastCloseCode: null,
+      socketOpens: state.socketOpens + 1,
+    })),
   reportError: (code, message) => set({ lastError: { code, message, at: Date.now() } }),
   clearError: () => set({ lastError: null }),
 }));