@schlessera/brain-ui-server 0.4.0

package/src/ws/history.ts

@@ -0,0 +1,51 @@
+import type { SessionHistoryMessage } from "@schlessera/brain-ui-sdk/protocol";
+import { type WSContext, sendTo } from "./clients.js";
+import { HISTORY_CHUNK_BYTES, shrinkForReplication } from "./shrink.js";
+
+/**
+ * Send a structured history to one socket as one or more `session_history`
+ * frames. Each message is size-bounded, then batched into byte-limited chunks
+ * so no single frame exceeds the socket's per-message limit (which would drop
+ * the connection and wedge the reconnect loop). The first frame replaces the
+ * client transcript; continuation frames carry `append: true`. An empty
+ * history still emits one (empty, replacing) frame so the client settles.
+ *
+ * Every frame carries `sessionId` so a client with multiple sessions open can
+ * demux it: a history load for a background session must never replace the
+ * transcript in view (a reconnect or an out-of-order getHistory otherwise
+ * silently overwrites the wrong pane).
+ */
+export function sendSessionHistory(
+  ws: WSContext,
+  sessionId: string,
+  messages: SessionHistoryMessage[]
+): void {
+  const bounded = messages.map((m) => shrinkForReplication(m));
+  if (bounded.length === 0) {
+    sendTo(ws, { type: "session_history", sessionId, messages: [] });
+    return;
+  }
+
+  let batch: SessionHistoryMessage[] = [];
+  let batchBytes = 0;
+  let isFirst = true;
+  const flush = () => {
+    sendTo(ws, {
+      type: "session_history",
+      sessionId,
+      messages: batch,
+      ...(isFirst ? {} : { append: true }),
+    });
+    isFirst = false;
+    batch = [];
+    batchBytes = 0;
+  };
+
+  for (const message of bounded) {
+    const size = JSON.stringify(message).length;
+    if (batch.length > 0 && batchBytes + size > HISTORY_CHUNK_BYTES) flush();
+    batch.push(message);
+    batchBytes += size;
+  }
+  flush();
+}

package/src/ws/host.ts

@@ -0,0 +1,69 @@
+import type { ServerMessage } from "@schlessera/brain-ui-sdk/protocol";
+import { broadcast, sendTo, type WSContext } from "./clients.js";
+import { TurnCoordinator } from "./turns.js";
+import { createSessionCatalog, type SessionCatalog } from "./session-catalog.js";
+
+/** Host-side turn timeout. The backend no longer times out — the host owns it. */
+const DEFAULT_TIMEOUT_MS = 10 * 60 * 1000; // 10 minutes
+
+/** Cap on host-side follow-up queue depth per session (backends without followUp). */
+export const MAX_SESSION_QUEUE = 5;
+
+export interface WsHostOptions {
+  /** Session persistence seam; defaults to the package's SQLite catalog. */
+  catalog?: SessionCatalog;
+  /** Display name used in connection/status copy. */
+  appName?: string;
+  /** Per-turn timeout in ms (default 10 minutes). */
+  turnTimeoutMs?: number;
+  /**
+   * Deployment-time cap on concurrent RUNNING sessions. Each running turn is
+   * roughly one CLI subprocess, so this bounds memory/CPU. Read per request so
+   * the deploy-time env is honored without a restart.
+   */
+  maxConcurrentSessions?: () => number;
+}
+
+function envMaxConcurrentSessions(): number {
+  return Math.max(1, Number(process.env.MAX_CONCURRENT_SESSIONS) || 3);
+}
+
+/**
+ * Everything one WebSocket coordinator instance owns: turn state, the session
+ * catalog, branding copy, and the frame senders. Handlers receive this host
+ * explicitly instead of reaching for module globals.
+ */
+export class WsHost {
+  readonly coordinator = new TurnCoordinator();
+  catalog: SessionCatalog;
+  appName: string;
+  turnTimeoutMs: number;
+  maxConcurrentSessions: () => number;
+
+  constructor(options: WsHostOptions = {}) {
+    this.catalog = options.catalog ?? createSessionCatalog();
+    this.appName = options.appName ?? "Brain UI";
+    this.turnTimeoutMs = options.turnTimeoutMs ?? DEFAULT_TIMEOUT_MS;
+    this.maxConcurrentSessions = options.maxConcurrentSessions ?? envMaxConcurrentSessions;
+  }
+
+  /** Re-apply embedder options (createApp configures the default host). */
+  configure(options: WsHostOptions): void {
+    if (options.catalog) this.catalog = options.catalog;
+    if (options.appName) this.appName = options.appName;
+    if (options.turnTimeoutMs) this.turnTimeoutMs = options.turnTimeoutMs;
+    if (options.maxConcurrentSessions) {
+      this.maxConcurrentSessions = options.maxConcurrentSessions;
+    }
+  }
+
+  /** Fan a frame out to every attached client (size-bounded per frame). */
+  sendToClients(msg: ServerMessage): void {
+    broadcast(msg);
+  }
+
+  /** Send a frame to one specific socket (size-bounded). */
+  sendMessage(ws: WSContext, msg: ServerMessage): void {
+    sendTo(ws, msg);
+  }
+}

package/src/ws/routing.ts

@@ -0,0 +1,43 @@
+import type { AgentBackend } from "@schlessera/brain-ui-sdk/server";
+import {
+  getBackendForProfile,
+  getBackendForSession,
+  getDefaultBackend,
+  listAllProviders,
+} from "../agent/backend.js";
+import type { SessionCatalog } from "./session-catalog.js";
+
+/**
+ * Resolve one turn slot's pinned backend and profile. Resumed sessions route by
+ * their stored backend id; new sessions route by the requested profile owner.
+ */
+export async function resolveTurnTarget(
+  catalog: SessionCatalog,
+  sessionId: string | undefined,
+  requested: string | undefined
+): Promise<{ profileId?: string; backend: AgentBackend; droppedPin?: string }> {
+  if (sessionId) {
+    const backend = await getBackendForSession(catalog.getStoredBackendId(sessionId));
+    const stored = catalog.getStoredProviderId(sessionId);
+    const available = await backend.listProfiles();
+    if (stored && available.some((profile) => profile.id === stored)) {
+      return { backend, profileId: stored };
+    }
+    // A stored pin that no longer resolves: fall back to the default profile
+    // but report it (droppedPin) so the caller can surface the switch instead
+    // of silently changing model/endpoint/billing with nothing in the logs.
+    return { backend, ...(stored ? { droppedPin: stored } : {}) };
+  }
+
+  if (requested) {
+    const available = await listAllProviders();
+    if (available.some((profile) => profile.id === requested)) {
+      const backend = await getBackendForProfile(requested);
+      if (backend) return { backend, profileId: requested };
+    }
+  }
+
+  const backend = await getDefaultBackend();
+  const profileId = (await backend.listProfiles())[0]?.id;
+  return { backend, ...(profileId ? { profileId } : {}) };
+}

package/src/ws/run-session.ts

@@ -0,0 +1,192 @@
+import type { ChatImageAttachment } from "@schlessera/brain-ui-sdk/protocol";
+import type { WSContext } from "./clients.js";
+import { withSessionId, withTurnScope } from "./frames.js";
+import { makeBridge, emitTurnError } from "./bridge.js";
+import { resolveTurnTarget } from "./routing.js";
+import type { QueuedFollowUp, RunningTurn } from "./turns.js";
+import { MAX_SESSION_QUEUE, type WsHost } from "./host.js";
+
+/**
+ * Run one session slot: the initial turn, then any queued follow-up turns in
+ * order. The host owns per-turn cancellation (AbortController + timeout). Never
+ * throws — runtime failures are emitted as error frames.
+ */
+export async function runSession(
+  host: WsHost,
+  initial: {
+    text: string;
+    sessionId?: string;
+    attachments: ChatImageAttachment[];
+    providerId?: string;
+  }
+): Promise<void> {
+  const { coordinator } = host;
+  coordinator.startingSessions += 1;
+  let target: Awaited<ReturnType<typeof resolveTurnTarget>>;
+  try {
+    target = await resolveTurnTarget(host.catalog, initial.sessionId, initial.providerId);
+  } catch (err) {
+    host.sendToClients({
+      type: "error",
+      code: "BACKEND_ERROR",
+      message: err instanceof Error ? err.message : String(err),
+      ...(initial.sessionId ? { sessionId: initial.sessionId } : {}),
+    });
+    return;
+  } finally {
+    coordinator.startingSessions = Math.max(0, coordinator.startingSessions - 1);
+  }
+
+  const { backend, profileId: initialProfileId } = target;
+  if (target.droppedPin) {
+    console.warn(
+      `[agent] session ${initial.sessionId} pinned profile "${target.droppedPin}" ` +
+        `is unavailable; running on the default profile.`
+    );
+    host.sendToClients({
+      type: "status",
+      status: "thinking",
+      detail: `Pinned model "${target.droppedPin}" is unavailable — running on the default.`,
+      ...(initial.sessionId ? { sessionId: initial.sessionId } : {}),
+    });
+  }
+  const turn: RunningTurn = {
+    sessionId: initial.sessionId ?? null,
+    turnId: crypto.randomUUID(),
+    providerId: initialProfileId ?? null,
+    backend,
+    abortController: new AbortController(),
+    timeoutHandle: setTimeout(() => {}, 0),
+    queue: [],
+    cancelled: false,
+  };
+  clearTimeout(turn.timeoutHandle);
+  coordinator.running.add(turn);
+  if (turn.sessionId) coordinator.bySession.set(turn.sessionId, turn);
+
+  let profileId = initialProfileId;
+  let resumeId = initial.sessionId;
+  let next: QueuedFollowUp | null = {
+    text: initial.text,
+    attachments: initial.attachments,
+  };
+
+  try {
+    while (next && !turn.cancelled) {
+      const { text, attachments } = next;
+      next = null;
+
+      const abortController = new AbortController();
+      turn.abortController = abortController;
+      const timeoutHandle = setTimeout(() => {
+        console.log("[agent] Turn timed out after", host.turnTimeoutMs, "ms");
+        abortController.abort();
+        // Reject any pending interactive request for this turn too. A bridge
+        // that awaits askUser/approval without racing the abort signal would
+        // otherwise leave startTurn parked forever after abort(), leaking a
+        // concurrency slot until restart. Draining unblocks the unwind.
+        coordinator.drainPendingForTurn(turn, "Turn timed out");
+      }, host.turnTimeoutMs);
+      turn.timeoutHandle = timeoutHandle;
+
+      const bridge = makeBridge(host, turn, text, backend.id);
+      try {
+        await backend.startTurn({
+          prompt: text,
+          attachments,
+          sessionId: resumeId,
+          profileId,
+          signal: abortController.signal,
+          bridge,
+        });
+      } catch (err) {
+        emitTurnError(host, turn, err);
+      } finally {
+        clearTimeout(timeoutHandle);
+      }
+
+      // Subsequent (queued) turns resume the now-known session and keep its
+      // pin. turn.providerId is the resolved profile (set from initial routing
+      // and re-pinned by session_info); carrying it forward stops a queued
+      // follow-up from silently dropping onto the default profile/model/billing.
+      resumeId = turn.sessionId ?? resumeId;
+      profileId = turn.providerId ?? profileId;
+
+      if (!turn.cancelled && turn.queue.length > 0) {
+        next = turn.queue.shift()!;
+        // A queued follow-up is its own turn — give it a fresh identity.
+        turn.turnId = crypto.randomUUID();
+      }
+    }
+  } finally {
+    if (turn.sessionId) coordinator.bySession.delete(turn.sessionId);
+    coordinator.running.delete(turn);
+    coordinator.drainPendingForTurn(turn, "Session ended");
+  }
+}
+
+/** Dispatch a chat_message: follow-up to a running session, or a new session. */
+export async function handleChatMessage(
+  host: WsHost,
+  ws: WSContext,
+  text: string,
+  sessionId: string | undefined,
+  attachments: ChatImageAttachment[],
+  requestedProviderId: string | undefined
+): Promise<void> {
+  const { coordinator } = host;
+  const runningTurn = sessionId ? coordinator.bySession.get(sessionId) : undefined;
+
+  if (runningTurn && sessionId) {
+    const backend = runningTurn.backend;
+    // Message to a session whose turn is running = a follow-up.
+    if (backend.capabilities.followUp && backend.followUp) {
+      // Inject into the running turn; frames flow through its bridge.
+      backend.followUp({ sessionId, prompt: text, attachments }).catch((err) => {
+        host.sendToClients(
+          withTurnScope(
+            { type: "error", code: "FOLLOWUP_FAILED", message: err instanceof Error ? err.message : String(err) },
+            runningTurn
+          )
+        );
+      });
+    } else if (runningTurn.queue.length >= MAX_SESSION_QUEUE) {
+      host.sendMessage(
+        ws,
+        // Session-scoped only: the rejected message would have become a
+        // FUTURE turn in this slot, not the one currently running.
+        withSessionId(
+          { type: "error", code: "SESSION_QUEUE_FULL", message: `This session's queue is full (max ${MAX_SESSION_QUEUE}). Wait for it to catch up.` },
+          sessionId
+        )
+      );
+    } else {
+      // Queue it as the session's next turn; report queued immediately.
+      runningTurn.queue.push({ text, attachments });
+      host.sendToClients(withSessionId({ type: "status", status: "queued" }, sessionId));
+    }
+    return;
+  }
+
+  // New session (or resume of an idle one) — gated by the concurrency cap.
+  const cap = host.maxConcurrentSessions();
+  if (coordinator.running.size + coordinator.startingSessions >= cap) {
+    host.sendMessage(ws, {
+      type: "error",
+      code: "SESSION_LIMIT",
+      message: `Too many concurrent sessions (max ${cap}). Wait for one to finish.`,
+      ...(sessionId ? { sessionId } : {}),
+    });
+    return;
+  }
+
+  if (sessionId) {
+    host.sendToClients(withSessionId({ type: "status", status: "thinking" }, sessionId));
+  }
+  void runSession(host, {
+    text,
+    sessionId,
+    attachments,
+    providerId: requestedProviderId,
+  });
+}

package/src/ws/session-catalog.ts

@@ -0,0 +1,97 @@
+import type { Database } from "bun:sqlite";
+import type { ServerResultMessage } from "@schlessera/brain-ui-sdk/protocol";
+import { getDb } from "../db/client.js";
+
+/**
+ * Persistence seam for session ownership + accounting. The ws coordinator only
+ * talks to this interface — swapping the store means implementing four
+ * methods, not editing the turn loop.
+ */
+export interface SessionCatalog {
+  /** The provider/profile a stored session is pinned to, or null. */
+  getStoredProviderId(sessionId: string): string | null;
+  /** The backend that owns a stored session, or null for legacy/unknown rows. */
+  getStoredBackendId(sessionId: string): string | null;
+  /**
+   * Persist session ownership as soon as `session_info` names the session,
+   * with zero cost/turn deltas — the terminal `result` adds the accounting.
+   * Without this, a failed or cancelled first turn left the transcript
+   * unowned.
+   */
+  persistSessionStub(
+    sessionId: string,
+    promptText: string,
+    providerId: string | null,
+    backendId: string
+  ): void;
+  /** Persist session accounting on the terminal `result` frame. */
+  persistSession(
+    msg: ServerResultMessage,
+    promptText: string,
+    providerId: string | null,
+    backendId: string
+  ): void;
+}
+
+const UPSERT_SESSION_SQL = `INSERT INTO sessions (id, title, created_at, last_active_at, total_cost_usd, num_turns, provider_id, backend_id)
+   VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+   ON CONFLICT(id) DO UPDATE SET
+     last_active_at = excluded.last_active_at,
+     total_cost_usd = total_cost_usd + excluded.total_cost_usd,
+     num_turns = num_turns + excluded.num_turns,
+     provider_id = COALESCE(excluded.provider_id, provider_id),
+     backend_id = COALESCE(excluded.backend_id, backend_id)`;
+
+/** SQLite-backed catalog over the package's session database. */
+export function createSessionCatalog(db: () => Database = getDb): SessionCatalog {
+  return {
+    getStoredProviderId(sessionId) {
+      const row = db()
+        .query("SELECT provider_id AS providerId FROM sessions WHERE id = ?")
+        .get(sessionId) as { providerId: string | null } | null;
+      return row?.providerId ?? null;
+    },
+
+    getStoredBackendId(sessionId) {
+      const row = db()
+        .query("SELECT backend_id AS backendId FROM sessions WHERE id = ?")
+        .get(sessionId) as { backendId: string | null } | null;
+      return row?.backendId ?? null;
+    },
+
+    persistSessionStub(sessionId, promptText, providerId, backendId) {
+      try {
+        db()
+          .prepare(UPSERT_SESSION_SQL)
+          .run(
+            sessionId,
+            promptText.slice(0, 100),
+            Date.now(),
+            Date.now(),
+            0,
+            0,
+            providerId,
+            backendId
+          );
+      } catch {}
+    },
+
+    persistSession(msg, promptText, providerId, backendId) {
+      try {
+        db()
+          .prepare(UPSERT_SESSION_SQL)
+          .run(
+            msg.sessionId,
+            promptText.slice(0, 100),
+            Date.now(),
+            Date.now(),
+            // costUsd is optional (rev 2): absent = unknown, accounted as 0.
+            msg.costUsd ?? 0,
+            msg.numTurns,
+            providerId,
+            backendId
+          );
+      } catch {}
+    },
+  };
+}

package/src/ws/shrink.ts

@@ -0,0 +1,95 @@
+/**
+ * Size-bounding for outbound WebSocket frames.
+ *
+ * A single oversized frame — typically a multi-megabyte tool result (a big
+ * `Read`/`Bash`/`Grep` output), a tool_use with a large input, or a tool
+ * result whose content array holds thousands of blocks — can exceed the
+ * socket's per-frame limit, drop the connection, and (because the client
+ * reconnects and re-requests the same history) wedge into a reconnect loop.
+ *
+ * This helper bounds any JSON-serializable payload below
+ * {@link MAX_WS_MESSAGE_BYTES} by head-truncating long strings AND
+ * head-clipping long arrays, leaving small metadata (ids, types, statuses)
+ * untouched. Both axes matter: string truncation alone leaves the cap
+ * unenforced for a payload built of many short strings. The model already
+ * received the full tool output server-side — this only bounds what the
+ * browser renders.
+ *
+ * Ported from oh-my-pi's collab replication-shrink.
+ */
+
+/** Per-frame ceiling for server → client messages. Well under Bun's 16 MB
+ * default and typical reverse-proxy WebSocket limits. */
+export const MAX_WS_MESSAGE_BYTES = 512_000;
+
+/** Target byte size for one `session_history` chunk. Below the frame cap so a
+ * batch plus its wrapper stays comfortably under {@link MAX_WS_MESSAGE_BYTES}. */
+export const HISTORY_CHUNK_BYTES = 400_000;
+
+interface ShrinkPass {
+  stringCap: number;
+  arrayLimit: number;
+}
+
+/**
+ * Progressive shrink passes. Each tightens both the per-string cap and the
+ * per-array head limit; the loop stops at the first pass whose output fits
+ * {@link MAX_WS_MESSAGE_BYTES}. The final pass clamps every string to 64 B and
+ * every array to one element, so even pathological mixes converge.
+ */
+const SHRINK_PASSES: readonly ShrinkPass[] = [
+  { stringCap: 64 * 1024, arrayLimit: 256 },
+  { stringCap: 16 * 1024, arrayLimit: 128 },
+  { stringCap: 4 * 1024, arrayLimit: 64 },
+  { stringCap: 1 * 1024, arrayLimit: 32 },
+  { stringCap: 256, arrayLimit: 16 },
+  { stringCap: 256, arrayLimit: 4 },
+  { stringCap: 64, arrayLimit: 1 },
+];
+
+const STRING_ELISION_RESERVE = 80;
+
+/**
+ * Recursively walk `value`, head-truncating any string longer than
+ * `stringCap` and head-clipping any array longer than `arrayLimit`. Returns a
+ * freshly built deep clone; short strings/numbers/booleans pass through.
+ */
+function shrinkWalk(value: unknown, stringCap: number, arrayLimit: number): unknown {
+  if (typeof value === "string") {
+    if (value.length <= stringCap) return value;
+    const headLen = Math.max(0, stringCap - STRING_ELISION_RESERVE);
+    return `${value.slice(0, headLen)}\n…[${value.length - headLen} chars elided]`;
+  }
+  if (Array.isArray(value)) {
+    const keep = Math.min(value.length, arrayLimit);
+    const elided = value.length - keep;
+    const out: unknown[] = new Array(elided > 0 ? keep + 1 : keep);
+    for (let i = 0; i < keep; i++) out[i] = shrinkWalk(value[i], stringCap, arrayLimit);
+    if (elided > 0) out[keep] = `…[${elided} items elided]`;
+    return out;
+  }
+  if (value && typeof value === "object") {
+    const src = value as Record<string, unknown>;
+    const out: Record<string, unknown> = {};
+    for (const k in src) out[k] = shrinkWalk(src[k], stringCap, arrayLimit);
+    return out;
+  }
+  return value;
+}
+
+/**
+ * Return `value` unchanged when its JSON serialization already fits
+ * {@link MAX_WS_MESSAGE_BYTES}; otherwise return a deep-cloned shadow shrunk
+ * along both string and array axes until it fits. The wire shape is preserved:
+ * only string leaves and array tails change — discriminator fields, ids, and
+ * other small metadata pass through untouched.
+ */
+export function shrinkForReplication<T>(value: T): T {
+  if (JSON.stringify(value).length <= MAX_WS_MESSAGE_BYTES) return value;
+  let shrunk: unknown = value;
+  for (const pass of SHRINK_PASSES) {
+    shrunk = shrinkWalk(value, pass.stringCap, pass.arrayLimit);
+    if (JSON.stringify(shrunk).length <= MAX_WS_MESSAGE_BYTES) return shrunk as T;
+  }
+  return shrunk as T;
+}