@schlessera/brain-ui-server 0.4.0

package/src/ws/attachments.ts

@@ -0,0 +1,86 @@
+import type { ChatImageAttachment } from "@schlessera/brain-ui-sdk/protocol";
+import {
+  ALLOWED_IMAGE_MEDIA_TYPES,
+  MAX_IMAGE_BYTES,
+  MAX_IMAGES_PER_MESSAGE,
+  MAX_TOTAL_IMAGE_BYTES,
+} from "@schlessera/brain-ui-sdk/protocol";
+
+export function estimateDecodedBase64Bytes(data: string): number {
+  const normalized = data.replace(/\s/g, "");
+  const padding = normalized.endsWith("==")
+    ? 2
+    : normalized.endsWith("=")
+      ? 1
+      : 0;
+  return Math.floor((normalized.length * 3) / 4) - padding;
+}
+
+export function validateAttachments(
+  attachments: unknown
+): { ok: true; attachments: ChatImageAttachment[] } | { ok: false; reason: string } {
+  if (attachments === undefined) {
+    return { ok: true, attachments: [] };
+  }
+
+  if (!Array.isArray(attachments)) {
+    return { ok: false, reason: "Attachments must be an array" };
+  }
+
+  if (attachments.length > MAX_IMAGES_PER_MESSAGE) {
+    return {
+      ok: false,
+      reason: `Too many images attached; maximum is ${MAX_IMAGES_PER_MESSAGE}`,
+    };
+  }
+
+  let totalBytes = 0;
+  const validated: ChatImageAttachment[] = [];
+
+  for (const attachment of attachments) {
+    if (
+      !attachment ||
+      typeof attachment !== "object" ||
+      typeof (attachment as ChatImageAttachment).data !== "string"
+    ) {
+      return { ok: false, reason: "Each attachment must include base64 image data" };
+    }
+
+    const data = (attachment as ChatImageAttachment).data;
+    if (
+      data.length === 0 ||
+      data.length % 4 !== 0 ||
+      !/^[A-Za-z0-9+/]+={0,2}$/.test(data)
+    ) {
+      return { ok: false, reason: "Attachment data is not valid base64" };
+    }
+
+    const mediaType = (attachment as ChatImageAttachment).mediaType;
+    if (!ALLOWED_IMAGE_MEDIA_TYPES.includes(mediaType)) {
+      return {
+        ok: false,
+        reason: `Unsupported image type: ${String(mediaType)}`,
+      };
+    }
+
+    const bytes = estimateDecodedBase64Bytes((attachment as ChatImageAttachment).data);
+    if (bytes > MAX_IMAGE_BYTES) {
+      return {
+        ok: false,
+        reason: `Image is too large; maximum decoded size is ${MAX_IMAGE_BYTES} bytes`,
+      };
+    }
+
+    totalBytes += bytes;
+    if (totalBytes > MAX_TOTAL_IMAGE_BYTES) {
+      return {
+        ok: false,
+        reason: `Attached images are too large; maximum total decoded size is ${MAX_TOTAL_IMAGE_BYTES} bytes`,
+      };
+    }
+
+    validated.push(attachment as ChatImageAttachment);
+  }
+
+  return { ok: true, attachments: validated };
+}

package/src/ws/bridge.ts

@@ -0,0 +1,139 @@
+import type {
+  BackendBridge,
+  PermissionDecision,
+  AskUserResult,
+  LocationFix,
+} from "@schlessera/brain-ui-sdk/server";
+import { BackendBusyError, BackendRequestError } from "@schlessera/brain-ui-sdk/server";
+import { hasClients } from "./clients.js";
+import { withTurnScope } from "./frames.js";
+import type { RunningTurn } from "./turns.js";
+import type { WsHost } from "./host.js";
+
+/** Build the per-turn bridge the backend drives. */
+export function makeBridge(
+  host: WsHost,
+  turn: RunningTurn,
+  promptText: string,
+  backendId: string
+): BackendBridge {
+  const { coordinator, catalog } = host;
+  // Capture the turn identity at construction: the slot's turnId is re-minted
+  // for each queued follow-up, and a backend can still emit late frames
+  // through this bridge after its startTurn resolved. Stamping from the live
+  // field would attribute those to the NEXT turn.
+  const turnId = turn.turnId;
+  return {
+    emit: (msg) => {
+      if (msg.type === "session_info") {
+        turn.sessionId = msg.sessionId;
+        if (msg.providerId) turn.providerId = msg.providerId;
+        coordinator.bySession.set(msg.sessionId, turn);
+        // Persist ownership the moment the identity exists — a turn that
+        // later fails or is cancelled must not leave an unowned transcript.
+        catalog.persistSessionStub(msg.sessionId, promptText, turn.providerId, backendId);
+      }
+      host.sendToClients(withTurnScope(msg, turn, turnId));
+      if (msg.type === "result") {
+        catalog.persistSession(msg, promptText, turn.providerId, backendId);
+      }
+    },
+    requestPermission: (req) => {
+      host.sendToClients(
+        withTurnScope(
+          {
+            type: "tool_approval_request",
+            toolUseId: req.toolUseId,
+            toolName: req.toolName,
+            input: req.input,
+            description: req.description,
+          },
+          turn,
+          turnId
+        )
+      );
+      host.sendToClients(
+        withTurnScope(
+          { type: "status", status: "tool_executing", detail: `Waiting for approval: ${req.toolName}` },
+          turn,
+          turnId
+        )
+      );
+      return new Promise<PermissionDecision>((resolve) => {
+        if (coordinator.collidesAcrossTurns(coordinator.pendingApprovals, req.toolUseId, turn)) {
+          resolve({ behavior: "deny", message: "Duplicate tool-approval id" });
+          return;
+        }
+        coordinator.pendingApprovals.set(req.toolUseId, { turn, turnId, resolve });
+      });
+    },
+    askUser: (requestId, questions) => {
+      host.sendToClients(
+        withTurnScope({ type: "ask_user_request", requestId, questions }, turn, turnId)
+      );
+      host.sendToClients(
+        withTurnScope(
+          { type: "status", status: "tool_executing", detail: "Waiting for your input" },
+          turn,
+          turnId
+        )
+      );
+      return new Promise<AskUserResult>((resolve, reject) => {
+        if (coordinator.collidesAcrossTurns(coordinator.pendingAskUser, requestId, turn)) {
+          reject(new Error("Duplicate ask-user request id"));
+          return;
+        }
+        coordinator.pendingAskUser.set(requestId, { turn, turnId, resolve, reject });
+      });
+    },
+    getLocation: (options) => {
+      if (!hasClients()) {
+        return Promise.reject(
+          new Error(`No ${host.appName} client is connected to read the location from.`)
+        );
+      }
+      const requestId = coordinator.nextLocationRequestId();
+      host.sendToClients(
+        withTurnScope({ type: "location_request", requestId, options }, turn, turnId)
+      );
+      host.sendToClients(
+        withTurnScope(
+          { type: "status", status: "tool_executing", detail: "Requesting your location" },
+          turn,
+          turnId
+        )
+      );
+      return new Promise<LocationFix>((resolve, reject) => {
+        if (coordinator.collidesAcrossTurns(coordinator.pendingLocation, requestId, turn)) {
+          reject(new Error("Duplicate location request id"));
+          return;
+        }
+        coordinator.pendingLocation.set(requestId, { turn, turnId, resolve, reject });
+      });
+    },
+  };
+}
+
+export function emitTurnError(host: WsHost, turn: RunningTurn, err: unknown): void {
+  // startTurn resolves for runtime failures (it emits its own error frame); it
+  // only rejects for caller errors.
+  if (err instanceof BackendBusyError) {
+    host.sendToClients(
+      withTurnScope(
+        { type: "error", code: "SESSION_BUSY", message: "That session already has a running turn." },
+        turn
+      )
+    );
+  } else if (err instanceof BackendRequestError) {
+    host.sendToClients(
+      withTurnScope({ type: "error", code: "BACKEND_REQUEST_ERROR", message: err.message }, turn)
+    );
+  } else {
+    host.sendToClients(
+      withTurnScope(
+        { type: "error", code: "BACKEND_ERROR", message: err instanceof Error ? err.message : String(err) },
+        turn
+      )
+    );
+  }
+}

package/src/ws/clients.ts

@@ -0,0 +1,57 @@
+import type { ServerMessage } from "@schlessera/brain-ui-sdk/protocol";
+import { shrinkForReplication } from "./shrink.js";
+
+/** Minimal structural view of a live socket — all we need to write to it. */
+export type WSContext = { send: (data: string) => void };
+
+/**
+ * Every currently-attached client socket. The active agent turn is a
+ * process-global singleton (server/src/ws/handler.ts), so all sockets observe
+ * the same conversation; streamed output fans out to all of them.
+ * Previously the server held a single `activeWs`, so a second connection
+ * silently orphaned the first (it stayed open but never received output).
+ */
+const clients = new Set<WSContext>();
+
+export function addClient(ws: WSContext): void {
+  clients.add(ws);
+}
+
+export function removeClient(ws: WSContext): void {
+  clients.delete(ws);
+}
+
+export function clientCount(): number {
+  return clients.size;
+}
+
+export function hasClients(): boolean {
+  return clients.size > 0;
+}
+
+/** Test-only: drop all registered sockets. */
+export function resetClientsForTests(): void {
+  clients.clear();
+}
+
+/** Serialize + size-bound a frame, then send it to one specific socket. */
+export function sendTo(ws: WSContext, msg: ServerMessage): void {
+  ws.send(JSON.stringify(shrinkForReplication(msg)));
+}
+
+/**
+ * Broadcast a frame to every attached client. Serializes once. A failing
+ * socket is skipped (its `onClose` will prune it) so one dead peer can't
+ * block delivery to the others.
+ */
+export function broadcast(msg: ServerMessage): void {
+  if (clients.size === 0) return;
+  const payload = JSON.stringify(shrinkForReplication(msg));
+  for (const ws of clients) {
+    try {
+      ws.send(payload);
+    } catch {
+      // Drop; the socket's onClose handler removes it from the set.
+    }
+  }
+}

package/src/ws/connection.ts

@@ -0,0 +1,132 @@
+import { upgradeWebSocket, websocket } from "hono/bun";
+import { PROTOCOL_REV } from "@schlessera/brain-ui-sdk/protocol";
+import { parseClientMessage } from "@schlessera/brain-ui-sdk/schemas";
+import { getBackendForSession } from "../agent/backend.js";
+import { addClient, removeClient, hasClients } from "./clients.js";
+import { withTurnScope } from "./frames.js";
+import { sendSessionHistory } from "./history.js";
+import { handleClientMessage } from "./dispatch.js";
+import type { WsHost } from "./host.js";
+
+export { websocket };
+
+/** Build the Hono WebSocket upgrade handler bound to one host. */
+export function createWsUpgrade(host: WsHost) {
+  return upgradeWebSocket(() => ({
+    async onOpen(_evt, ws) {
+      console.log("[ws] Client connected");
+      const { coordinator, catalog } = host;
+
+      // Handshake first (rev 2, additive): protocol revision + coarse
+      // capabilities, so the client can gate behavior instead of sniffing.
+      host.sendMessage(ws, {
+        type: "server_hello",
+        protocolRev: PROTOCOL_REV,
+        capabilities: { multiSession: true, askUser: true, location: true },
+      });
+
+      // Snapshot-on-connect only for the single-running-session case (backward
+      // compatible). With zero or several running sessions the client rehydrates
+      // itself per-session via session_resume, and live frames (sessionId-scoped)
+      // fan out to it once it joins the broadcast set.
+      const runningTurns = [...coordinator.running].filter((t) => t.sessionId);
+      if (runningTurns.length === 1) {
+        const turn = runningTurns[0];
+        const sid = turn.sessionId!;
+        host.sendMessage(
+          ws,
+          // Stamp the live turn identity: a reconnecting client must learn the
+          // current turnId, or the correlation guarantee dies at the reconnect
+          // it exists for.
+          withTurnScope(
+            {
+              type: "session_info",
+              sessionId: sid,
+              isNew: false,
+              providerId: turn.providerId ?? catalog.getStoredProviderId(sid) ?? undefined,
+            },
+            turn
+          )
+        );
+        try {
+          const backend = await getBackendForSession(
+            catalog.getStoredBackendId(sid) ?? turn.backend.id
+          );
+          const history = await backend.getHistory(sid);
+          if (history.length > 0) sendSessionHistory(ws, sid, history);
+        } catch (err) {
+          console.error("[ws] snapshot-on-connect failed:", err);
+        } finally {
+          addClient(ws);
+          host.sendMessage(
+            ws,
+            withTurnScope(
+              {
+                type: "status",
+                status: "thinking",
+                detail: "Session in progress",
+                sessionId: sid,
+              },
+              turn
+            )
+          );
+        }
+        return;
+      }
+
+      addClient(ws);
+      host.sendMessage(ws, {
+        type: "status",
+        status: "idle",
+        detail: `Connected to ${host.appName}`,
+      });
+    },
+
+    onMessage(evt, ws) {
+      // Boundary validation (rev 2): byte cap + JSON decode + schema, in one
+      // place. No more casting client JSON to ClientMessage.
+      //
+      // Text frames only: hono's Bun adapter hands binary frames over as the
+      // underlying POOLED ArrayBuffer (byteOffset/byteLength discarded), so a
+      // binary frame cannot be decoded correctly here. The protocol is JSON
+      // text; reject anything else rather than parse a slab.
+      const raw = evt.data;
+      if (typeof raw !== "string") {
+        host.sendMessage(ws, {
+          type: "error",
+          code: "PARSE_ERROR",
+          message: "Binary frames are not supported; send JSON text",
+        });
+        return;
+      }
+      const parsed = parseClientMessage(raw);
+      if (!parsed.ok) {
+        host.sendMessage(ws, { type: "error", code: "PARSE_ERROR", message: parsed.error });
+        return;
+      }
+      // handleClientMessage is async — a rejection must not escape as an
+      // unhandled rejection with no frame sent.
+      void Promise.resolve()
+        .then(() => handleClientMessage(host, ws, parsed.message))
+        .catch(() => {
+          host.sendMessage(ws, {
+            type: "error",
+            code: "INTERNAL_ERROR",
+            message: "Failed to handle message",
+          });
+        });
+    },
+
+    onClose(_evt, ws) {
+      console.log("[ws] Client disconnected");
+      removeClient(ws);
+      // Turns keep running in the background. Only reject pending interactive
+      // requests once the LAST client leaves — while another client remains it
+      // can still answer them.
+      if (hasClients()) return;
+      for (const turn of host.coordinator.running) {
+        host.coordinator.drainPendingForTurn(turn, "Client disconnected");
+      }
+    },
+  }));
+}

package/src/ws/dispatch.ts

@@ -0,0 +1,162 @@
+import type { ClientMessage } from "@schlessera/brain-ui-sdk/protocol";
+import { getBackendForSession } from "../agent/backend.js";
+import type { WSContext } from "./clients.js";
+import { locationErrorText } from "./frames.js";
+import { sendSessionHistory } from "./history.js";
+import { validateAttachments } from "./attachments.js";
+import { handleChatMessage } from "./run-session.js";
+import type { WsHost } from "./host.js";
+
+/**
+ * A client MAY echo the request's turnId (rev 2). When it does, it must match
+ * the turn that raised the request — a stale echo from before a reconnect or
+ * a follow-up would otherwise resolve the wrong turn's pending promise.
+ * Absent turnId stays valid: the field is optional on the wire.
+ */
+function turnIdMatches(pending: { turnId: string }, echoed: string | undefined): boolean {
+  return echoed === undefined || echoed === pending.turnId;
+}
+
+export async function handleClientMessage(
+  host: WsHost,
+  ws: WSContext,
+  msg: ClientMessage
+): Promise<void> {
+  const { coordinator, catalog } = host;
+  switch (msg.type) {
+    case "chat_message": {
+      const attachmentResult = validateAttachments(msg.attachments);
+      if (!attachmentResult.ok) {
+        host.sendMessage(ws, {
+          type: "error",
+          code: "ATTACHMENT_REJECTED",
+          message: attachmentResult.reason,
+          ...(msg.sessionId ? { sessionId: msg.sessionId } : {}),
+        });
+        return;
+      }
+      await handleChatMessage(
+        host,
+        ws,
+        msg.text,
+        msg.sessionId,
+        attachmentResult.attachments,
+        msg.providerId
+      );
+      break;
+    }
+
+    case "ask_user_response": {
+      const pending = coordinator.pendingAskUser.get(msg.requestId);
+      if (pending && turnIdMatches(pending, msg.turnId)) {
+        coordinator.pendingAskUser.delete(msg.requestId);
+        pending.resolve({ answers: msg.answers, annotations: msg.annotations });
+      }
+      break;
+    }
+
+    case "ask_user_cancel": {
+      const pending = coordinator.pendingAskUser.get(msg.requestId);
+      if (pending && turnIdMatches(pending, msg.turnId)) {
+        coordinator.pendingAskUser.delete(msg.requestId);
+        pending.reject(new Error(msg.reason || "User cancelled the question"));
+      }
+      break;
+    }
+
+    case "location_response": {
+      const pending = coordinator.pendingLocation.get(msg.requestId);
+      if (pending && turnIdMatches(pending, msg.turnId)) {
+        coordinator.pendingLocation.delete(msg.requestId);
+        pending.resolve({ coords: msg.coords, timestamp: msg.timestamp });
+      }
+      break;
+    }
+
+    case "location_error": {
+      const pending = coordinator.pendingLocation.get(msg.requestId);
+      if (pending && turnIdMatches(pending, msg.turnId)) {
+        coordinator.pendingLocation.delete(msg.requestId);
+        pending.reject(new Error(locationErrorText(msg.code, msg.message)));
+      }
+      break;
+    }
+
+    case "tool_approval": {
+      const pending = coordinator.pendingApprovals.get(msg.toolUseId);
+      if (pending && turnIdMatches(pending, msg.turnId)) {
+        coordinator.pendingApprovals.delete(msg.toolUseId);
+        pending.resolve(
+          msg.updatedInput
+            ? { behavior: "allow", updatedInput: msg.updatedInput }
+            : { behavior: "allow" }
+        );
+      }
+      break;
+    }
+
+    case "tool_denial": {
+      const pending = coordinator.pendingApprovals.get(msg.toolUseId);
+      if (pending && turnIdMatches(pending, msg.turnId)) {
+        coordinator.pendingApprovals.delete(msg.toolUseId);
+        pending.resolve({ behavior: "deny", message: msg.message });
+      }
+      break;
+    }
+
+    case "cancel": {
+      if (msg.sessionId) {
+        const turn = coordinator.bySession.get(msg.sessionId);
+        if (turn) coordinator.cancelTurn(turn, "Cancelled by user");
+        return;
+      }
+      // No sessionId: cancel the sole running session; ambiguous if several run.
+      if (coordinator.running.size === 0) return;
+      if (coordinator.running.size > 1) {
+        host.sendMessage(ws, {
+          type: "error",
+          code: "CANCEL_AMBIGUOUS",
+          message: "Several sessions are running — specify which to cancel.",
+        });
+        return;
+      }
+      coordinator.cancelTurn([...coordinator.running][0], "Cancelled by user");
+      break;
+    }
+
+    case "session_resume": {
+      host.sendMessage(ws, {
+        type: "session_info",
+        sessionId: msg.sessionId,
+        isNew: false,
+        providerId: catalog.getStoredProviderId(msg.sessionId) ?? undefined,
+      });
+
+      try {
+        const backend = await getBackendForSession(catalog.getStoredBackendId(msg.sessionId));
+        const messages = await backend.getHistory(msg.sessionId);
+        sendSessionHistory(ws, msg.sessionId, messages);
+        // A resume of a RUNNING session (reattach) must not report idle: idle
+        // would clear the client's running badge and finish its streaming
+        // message mid-turn. Mirror the snapshot-on-connect status instead.
+        const runningTurn = coordinator.bySession.get(msg.sessionId);
+        host.sendMessage(ws, {
+          type: "status",
+          ...(runningTurn
+            ? { status: "thinking" as const, detail: "Session in progress" }
+            : { status: "idle" as const, detail: "Session loaded" }),
+          sessionId: msg.sessionId,
+          ...(runningTurn ? { turnId: runningTurn.turnId } : {}),
+        });
+      } catch (err) {
+        host.sendMessage(ws, {
+          type: "error",
+          code: "SESSION_LOAD_ERROR",
+          message: err instanceof Error ? err.message : "Failed to load session",
+          sessionId: msg.sessionId,
+        });
+      }
+      break;
+    }
+  }
+}

package/src/ws/frames.ts

@@ -0,0 +1,42 @@
+import type { ServerMessage } from "@schlessera/brain-ui-sdk/protocol";
+import type { RunningTurn } from "./turns.js";
+
+/**
+ * Stamp a session-scoped frame with a sessionId for client demux. Every
+ * ServerMessage is session-scoped (or carries its own sessionId), so this only
+ * fills in a missing id and never overwrites one the backend already set.
+ */
+export function withSessionId(msg: ServerMessage, sessionId: string | null): ServerMessage {
+  if (!sessionId) return msg;
+  if ((msg as { sessionId?: string }).sessionId) return msg;
+  return { ...msg, sessionId } as ServerMessage;
+}
+
+/**
+ * Stamp a frame with the turn's session AND turn identity (rev 2). turnId lets
+ * the client correlate interactive round-trips (approvals / ask-user /
+ * location) with the exact turn that raised them, even across a reconnect.
+ */
+export function withTurnScope(
+  msg: ServerMessage,
+  turn: RunningTurn,
+  turnId: string = turn.turnId
+): ServerMessage {
+  const scoped = withSessionId(msg, turn.sessionId);
+  if ((scoped as { turnId?: string }).turnId) return scoped;
+  return { ...scoped, turnId } as ServerMessage;
+}
+
+/** Human-readable reason for a GeolocationPositionError code. */
+export function locationErrorText(code: number, fallback: string): string {
+  switch (code) {
+    case 1:
+      return "The user denied permission to share their location.";
+    case 2:
+      return "The user's location is currently unavailable.";
+    case 3:
+      return "Timed out while retrieving the user's location.";
+    default:
+      return fallback || "Failed to retrieve the user's location.";
+  }
+}

package/src/ws/handler.ts

@@ -0,0 +1,50 @@
+/**
+ * Composition root for the WebSocket side: one process-wide host instance
+ * wired to the default SQLite session catalog and the backend registry. The
+ * actual behavior lives in the focused modules next to this file —
+ * turns/bridge/run-session/dispatch/connection — which all take the host
+ * explicitly. createApp() configures this host from its options before
+ * building the upgrade handler.
+ */
+import type { ClientMessage } from "@schlessera/brain-ui-sdk/protocol";
+import type { WSContext } from "./clients.js";
+import { WsHost, type WsHostOptions } from "./host.js";
+import { createWsUpgrade, websocket } from "./connection.js";
+import { handleClientMessage as dispatchClientMessage } from "./dispatch.js";
+
+export { websocket };
+export { WsHost, type WsHostOptions } from "./host.js";
+export { resolveTurnTarget } from "./routing.js";
+export { validateAttachments, estimateDecodedBase64Bytes } from "./attachments.js";
+export { createSessionCatalog, type SessionCatalog } from "./session-catalog.js";
+
+/** The process-default host. */
+export const defaultWsHost = new WsHost();
+
+/** Apply embedder options to the default host (called by createApp). */
+export function configureWsHost(options: WsHostOptions): void {
+  defaultWsHost.configure(options);
+}
+
+/** Hono upgrade handler bound to the default host. */
+export const wsUpgrade = createWsUpgrade(defaultWsHost);
+
+/** True while any session has a running turn. */
+export function isTurnActive(): boolean {
+  return defaultWsHost.coordinator.isTurnActive();
+}
+
+/** Cancel every running turn (used on shutdown). */
+export function cancelActiveTurn(): boolean {
+  return defaultWsHost.coordinator.cancelAll("Server shutting down");
+}
+
+/** Drive a client message against the default host (also the test seam). */
+export async function handleClientMessage(ws: WSContext, msg: ClientMessage): Promise<void> {
+  return dispatchClientMessage(defaultWsHost, ws, msg);
+}
+
+/** Test seam: clear all turn/pending state on the default host. */
+export function resetForTests(): void {
+  defaultWsHost.coordinator.reset();
+}