@schlessera/brain-ui-server 0.15.0 → 0.17.0

package/src/ws/connection.ts

@@ -3,16 +3,32 @@ import { PROTOCOL_REV } from "@schlessera/brain-ui-sdk/protocol";
 import { parseClientMessage } from "@schlessera/brain-ui-sdk/schemas";
 import { withTurnScope } from "./frames.js";
 import { sendSessionHistory } from "./history.js";
-import { handleClientMessage } from "./dispatch.js";
+import { handleClientMessage, type ConnectionState } from "./dispatch.js";
 import type { WsHost } from "./host.js";
+import type { WSContext } from "./clients.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");
+/**
+ * The socket lifecycle handlers for one host, separate from the Hono upgrade
+ * that wraps them.
+ *
+ * Split out so the real frame path is reachable from a test without standing
+ * up an HTTP server: `createWsHandlers(host).onMessage(...)` runs exactly what
+ * production runs, which is what makes an assertion about a dropped frame an
+ * assertion about the shipped code rather than about a re-implementation.
+ */
+export function createWsHandlers(host: WsHost) {
+  // One bucket per connection, created here so it lives and dies with the
+  // socket rather than in a map keyed by something a peer controls.
+  const limiter = host.newRateLimiter();
+  // Per-connection negotiation state: what revision this client declared.
+  // Lives with the socket, like the limiter.
+  const connection: ConnectionState = {};
+
+  return {
+    async onOpen(_evt: Event, ws: WSContext) {
+      host.log.emit({ severityText: "INFO", body: "client connected" });
       const { coordinator, catalog } = host;
 
       // Handshake first (rev 2, additive): protocol revision + coarse
@@ -53,7 +69,11 @@ export function createWsUpgrade(host: WsHost) {
           const history = await backend.getHistory(sid);
           if (history.length > 0) sendSessionHistory(ws, sid, history);
         } catch (err) {
-          console.error("[ws] snapshot-on-connect failed:", err);
+          host.log.emit({
+            severityText: "ERROR",
+            body: "snapshot-on-connect failed",
+            attributes: { error: err instanceof Error ? err.message : String(err) },
+          });
         } finally {
           host.clients.add(ws);
           host.sendMessage(
@@ -80,7 +100,7 @@ export function createWsUpgrade(host: WsHost) {
       });
     },
 
-    onMessage(evt, ws) {
+    onMessage(evt: MessageEvent, ws: WSContext) {
       // Boundary validation (rev 2): byte cap + JSON decode + schema, in one
       // place. No more casting client JSON to ClientMessage.
       //
@@ -88,8 +108,21 @@ export function createWsUpgrade(host: WsHost) {
       // 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.
+      // Metered BEFORE parsing: the point is to bound work an unmetered peer
+      // can make this process do, and parsing is most of that work.
+      if (limiter && !limiter.take().allowed) {
+        host.reportDroppedFrame("rate_limited");
+        host.sendMessage(ws, {
+          type: "error",
+          code: "RATE_LIMITED",
+          message: "Too many frames; slow down.",
+        });
+        return;
+      }
+
       const raw = evt.data;
       if (typeof raw !== "string") {
+        host.reportDroppedFrame("binary_frame");
         host.sendMessage(ws, {
           type: "error",
           code: "PARSE_ERROR",
@@ -99,14 +132,27 @@ export function createWsUpgrade(host: WsHost) {
       }
       const parsed = parseClientMessage(raw);
       if (!parsed.ok) {
+        // The reason is the parser's own bounded message, never the frame:
+        // the payload is caller-supplied and capped at 12 MB.
+        host.reportDroppedFrame("parse_error", parsed.error);
         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(() => {
+        .then(() => handleClientMessage(host, ws, parsed.message, connection))
+        .catch((err) => {
+          // This used to swallow the cause entirely: the client got a generic
+          // frame and the server kept no record of what threw.
+          host.log.emit({
+            severityText: "ERROR",
+            body: "client message handler failed",
+            attributes: {
+              "frame.type": parsed.message.type,
+              error: err instanceof Error ? err.message : String(err),
+            },
+          });
           host.sendMessage(ws, {
             type: "error",
             code: "INTERNAL_ERROR",
@@ -115,8 +161,17 @@ export function createWsUpgrade(host: WsHost) {
         });
     },
 
-    onClose(_evt, ws) {
-      console.log("[ws] Client disconnected");
+    onClose(evt: CloseEvent, ws: WSContext) {
+      host.log.emit({ severityText: "INFO", body: "client disconnected" });
+      // No onError here on purpose: hono's Bun adapter never dispatches it
+      // (only open/message/close reach these handlers), so a transport failure
+      // is only visible as an abnormal close code. 1000/1001 are the two
+      // clean endings (normal closure, going away); anything else — before it
+      // was recorded — looked exactly like a clean disconnect.
+      const code = (evt as { code?: unknown }).code;
+      if (typeof code === "number" && code !== 1000 && code !== 1001) {
+        host.reportAbnormalClose(code);
+      }
       host.clients.remove(ws);
       // Turns keep running in the background. Only reject pending interactive
       // requests once the LAST client leaves — while another client remains it
@@ -126,5 +181,10 @@ export function createWsUpgrade(host: WsHost) {
         host.coordinator.drainPendingForTurn(turn, "Client disconnected");
       }
     },
-  }));
+  };
+}
+
+/** Build the Hono WebSocket upgrade handler bound to one host. */
+export function createWsUpgrade(host: WsHost) {
+  return upgradeWebSocket(() => createWsHandlers(host));
 }

package/src/ws/dispatch.ts

@@ -1,3 +1,4 @@
+import { PROTOCOL_REV_CLIENT_ECHO } from "@schlessera/brain-ui-sdk/protocol";
 import type { ClientMessage } from "@schlessera/brain-ui-sdk/protocol";
 import type { WSContext } from "./clients.js";
 import { locationErrorText } from "./frames.js";
@@ -6,23 +7,51 @@ import { validateAttachments } from "./attachments.js";
 import { handleChatMessage } from "./run-session.js";
 import type { WsHost } from "./host.js";
 
+/** Per-connection negotiation state, owned by the socket handler. */
+export interface ConnectionState {
+  /** Revision the client declared via `client_hello`; absent means rev 2. */
+  protocolRev?: number;
+}
+
 /**
- * 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.
+ * Does this reply's echoed turnId identify the turn that raised the request?
+ *
+ * A wrong id is always refused: a stale echo from before a reconnect or a
+ * follow-up would resolve a different turn's pending promise.
+ *
+ * A MISSING id depends on who is speaking. A client that declared rev 3
+ * promised to echo, so silence means the reply cannot be correlated and is
+ * refused. A client that declared nothing is rev 2, where the field is
+ * optional, and is still tolerated — that tolerance is the deprecation window,
+ * and it is what lets this be enforced at all without breaking clients that
+ * predate `client_hello`.
  */
-function turnIdMatches(pending: { turnId: string }, echoed: string | undefined): boolean {
-  return echoed === undefined || echoed === pending.turnId;
+export function turnIdMatches(
+  pending: { turnId: string },
+  echoed: string | undefined,
+  requireEcho: boolean
+): boolean {
+  if (echoed === undefined) return !requireEcho;
+  return echoed === pending.turnId;
 }
 
 export async function handleClientMessage(
   host: WsHost,
   ws: WSContext,
-  msg: ClientMessage
+  msg: ClientMessage,
+  connection: ConnectionState = {}
 ): Promise<void> {
   const { coordinator, catalog } = host;
+  const requireEcho = (connection.protocolRev ?? 2) >= PROTOCOL_REV_CLIENT_ECHO;
   switch (msg.type) {
+    case "client_hello": {
+      // Record what this connection speaks. Never rejected on version: a
+      // future client declaring rev 9 is simply held to the rules this host
+      // knows, and an unknown capability flag is ignored.
+      connection.protocolRev = msg.protocolRev;
+      return;
+    }
+
     case "chat_message": {
       const attachmentResult = validateAttachments(msg.attachments);
       if (!attachmentResult.ok) {
@@ -46,7 +75,7 @@ export async function handleClientMessage(
 
     case "ask_user_response": {
       const pending = coordinator.pendingAskUser.get(msg.requestId);
-      if (pending && turnIdMatches(pending, msg.turnId)) {
+      if (pending && turnIdMatches(pending, msg.turnId, requireEcho)) {
         coordinator.pendingAskUser.delete(msg.requestId);
         pending.resolve({ answers: msg.answers, annotations: msg.annotations });
       }
@@ -55,7 +84,7 @@ export async function handleClientMessage(
 
     case "ask_user_cancel": {
       const pending = coordinator.pendingAskUser.get(msg.requestId);
-      if (pending && turnIdMatches(pending, msg.turnId)) {
+      if (pending && turnIdMatches(pending, msg.turnId, requireEcho)) {
         coordinator.pendingAskUser.delete(msg.requestId);
         pending.reject(new Error(msg.reason || "User cancelled the question"));
       }
@@ -64,7 +93,7 @@ export async function handleClientMessage(
 
     case "location_response": {
       const pending = coordinator.pendingLocation.get(msg.requestId);
-      if (pending && turnIdMatches(pending, msg.turnId)) {
+      if (pending && turnIdMatches(pending, msg.turnId, requireEcho)) {
         coordinator.pendingLocation.delete(msg.requestId);
         pending.resolve({ coords: msg.coords, timestamp: msg.timestamp });
       }
@@ -73,7 +102,7 @@ export async function handleClientMessage(
 
     case "location_error": {
       const pending = coordinator.pendingLocation.get(msg.requestId);
-      if (pending && turnIdMatches(pending, msg.turnId)) {
+      if (pending && turnIdMatches(pending, msg.turnId, requireEcho)) {
         coordinator.pendingLocation.delete(msg.requestId);
         pending.reject(new Error(locationErrorText(msg.code, msg.message)));
       }
@@ -82,7 +111,7 @@ export async function handleClientMessage(
 
     case "mask_response": {
       const pending = coordinator.pendingMask.get(msg.requestId);
-      if (pending && turnIdMatches(pending, msg.turnId)) {
+      if (pending && turnIdMatches(pending, msg.turnId, requireEcho)) {
         coordinator.pendingMask.delete(msg.requestId);
         // Decoded here rather than in the tool: the boundary already validated
         // the base64 and its size, so the backend gets bytes it can trust.
@@ -93,7 +122,7 @@ export async function handleClientMessage(
 
     case "mask_error": {
       const pending = coordinator.pendingMask.get(msg.requestId);
-      if (pending && turnIdMatches(pending, msg.turnId)) {
+      if (pending && turnIdMatches(pending, msg.turnId, requireEcho)) {
         coordinator.pendingMask.delete(msg.requestId);
         pending.reject(
           new Error(
@@ -108,7 +137,7 @@ export async function handleClientMessage(
 
     case "tool_approval": {
       const pending = coordinator.pendingApprovals.get(msg.toolUseId);
-      if (pending && turnIdMatches(pending, msg.turnId)) {
+      if (pending && turnIdMatches(pending, msg.turnId, requireEcho)) {
         coordinator.pendingApprovals.delete(msg.toolUseId);
         pending.resolve(
           msg.updatedInput
@@ -121,7 +150,7 @@ export async function handleClientMessage(
 
     case "tool_denial": {
       const pending = coordinator.pendingApprovals.get(msg.toolUseId);
-      if (pending && turnIdMatches(pending, msg.turnId)) {
+      if (pending && turnIdMatches(pending, msg.turnId, requireEcho)) {
         coordinator.pendingApprovals.delete(msg.toolUseId);
         pending.resolve({ behavior: "deny", message: msg.message });
       }
@@ -173,10 +202,12 @@ export async function handleClientMessage(
           ...(runningTurn ? { turnId: runningTurn.turnId } : {}),
         });
       } catch (err) {
+        const message = err instanceof Error ? err.message : "Failed to load session";
+        host.reportTurnFailed("SESSION_LOAD_ERROR", { sessionId: msg.sessionId }, message);
         host.sendMessage(ws, {
           type: "error",
           code: "SESSION_LOAD_ERROR",
-          message: err instanceof Error ? err.message : "Failed to load session",
+          message,
           sessionId: msg.sessionId,
         });
       }

package/src/ws/host.ts

@@ -3,6 +3,8 @@ import { ClientSet, sendTo, type WSContext } from "./clients.js";
 import { TurnCoordinator } from "./turns.js";
 import type { SessionCatalog } from "./session-catalog.js";
 import type { BackendRegistry } from "../agent/backend.js";
+import { createSilentObservability, type Observability } from "../observability/index.js";
+import { FrameRateLimiter } from "./rate-limit.js";
 
 /** Host-side turn timeout. The backend no longer times out — the host owns it. */
 const DEFAULT_TIMEOUT_MS = 10 * 60 * 1000; // 10 minutes
@@ -50,6 +52,38 @@ export interface WsHostOptions {
    * it dynamic; createApp passes the resolved config value.
    */
   maxConcurrentSessions?: () => number;
+  /**
+   * Where this coordinator reports. Injected rather than reached for, so two
+   * apps in one process report separately and a test can assert on what the
+   * socket layer actually said. Defaults to silence: an embedder that never
+   * passes one gets no output, not a surprise stream on stdout.
+   */
+  observability?: Observability;
+  /**
+   * Inbound frame metering policy, per connection. Omitted (or a rate of 0)
+   * means no metering — which is what a test wants, and what an embedder
+   * fronting the socket with its own limiter wants.
+   */
+  wsRate?: { ratePerSecond: number; burst: number };
+}
+
+/** Identity of one turn, as it appears on a log record. */
+export interface TurnLogContext {
+  sessionId?: string | null;
+  turnId?: string | null;
+  providerId?: string | null;
+}
+
+/**
+ * Log attributes for one turn. Null fields are omitted rather than stringified
+ * — a new session has no sessionId until `session_info` names it.
+ */
+export function turnLogAttributes(turn: TurnLogContext): Record<string, string> {
+  return {
+    ...(turn.sessionId ? { "session.id": turn.sessionId } : {}),
+    ...(turn.turnId ? { "turn.id": turn.turnId } : {}),
+    ...(turn.providerId ? { profile: turn.providerId } : {}),
+  };
 }
 
 /**
@@ -66,6 +100,25 @@ export class WsHost {
   appName: string;
   turnTimeoutMs: number;
   maxConcurrentSessions: () => number;
+  readonly observability: Observability;
+  readonly wsRate: { ratePerSecond: number; burst: number } | null;
+  /** Scoped instruments, resolved once — `[ws]` is the existing log prefix. */
+  readonly log: ReturnType<Observability["logger"]>;
+  private readonly framesDropped: ReturnType<
+    ReturnType<Observability["meter"]>["createCounter"]
+  >;
+  private readonly turnsStarted: ReturnType<
+    ReturnType<Observability["meter"]>["createCounter"]
+  >;
+  private readonly turnsCompleted: ReturnType<
+    ReturnType<Observability["meter"]>["createCounter"]
+  >;
+  private readonly turnsFailed: ReturnType<
+    ReturnType<Observability["meter"]>["createCounter"]
+  >;
+  private readonly wsErrors: ReturnType<
+    ReturnType<Observability["meter"]>["createCounter"]
+  >;
 
   constructor(options: WsHostOptions) {
     this.registry = options.registry;
@@ -74,11 +127,121 @@ export class WsHost {
     this.turnTimeoutMs = options.turnTimeoutMs ?? DEFAULT_TIMEOUT_MS;
     this.maxConcurrentSessions =
       options.maxConcurrentSessions ?? (() => DEFAULT_MAX_CONCURRENT_SESSIONS);
+    this.observability = options.observability ?? createSilentObservability();
+    this.wsRate =
+      options.wsRate && options.wsRate.ratePerSecond > 0 ? options.wsRate : null;
+    this.log = this.observability.logger("ws");
+    const meter = this.observability.meter("ws");
+    this.framesDropped = meter.createCounter("ws.frames.dropped", {
+      description: "Inbound frames refused before reaching a handler",
+    });
+    this.turnsStarted = meter.createCounter("turns.started", {
+      description: "Turns handed to a backend",
+    });
+    this.turnsCompleted = meter.createCounter("turns.completed", {
+      description: "Turns whose backend call resolved without throwing",
+    });
+    this.turnsFailed = meter.createCounter("turns.failed", {
+      description: "Turn failures surfaced to the client, by error code",
+    });
+    this.wsErrors = meter.createCounter("ws.errors", {
+      description: "Transport errors reported by the socket layer",
+    });
+    this.coordinator.log = this.log;
+  }
+
+  /** A metering bucket for one new connection, or null when metering is off. */
+  newRateLimiter(): FrameRateLimiter | null {
+    return this.wsRate ? new FrameRateLimiter(this.wsRate) : null;
+  }
+
+  /**
+   * Record a frame this server refused, and say so.
+   *
+   * Before this existed, every rejection answered the client and vanished:
+   * `parseClientMessage` failures were never logged, so the validation we
+   * already shipped had no observability at all. `reason` is a bounded token,
+   * never the frame body — the payload is caller-supplied and can be 12 MB.
+   */
+  reportDroppedFrame(reason: string, detail?: string): void {
+    this.framesDropped.add(1, { reason, direction: "inbound" });
+    this.log.emit({
+      severityText: "WARN",
+      body: "inbound frame rejected",
+      attributes: detail ? { reason, detail } : { reason },
+    });
+  }
+
+  /** A turn began executing: counted, and logged with its correlation ids. */
+  reportTurnStarted(turn: TurnLogContext): void {
+    this.turnsStarted.add(1);
+    this.log.emit({
+      severityText: "INFO",
+      body: "turn started",
+      attributes: turnLogAttributes(turn),
+    });
+  }
+
+  /** A turn's backend call resolved: counted, and logged with its duration. */
+  reportTurnCompleted(turn: TurnLogContext, durationMs: number): void {
+    this.turnsCompleted.add(1);
+    this.log.emit({
+      severityText: "INFO",
+      body: "turn completed",
+      attributes: { ...turnLogAttributes(turn), "duration.ms": durationMs },
+    });
+  }
+
+  /**
+   * Record a turn failure the client is being told about. Before this existed
+   * every such failure was an error FRAME only — visible on one phone screen,
+   * absent from the server's own record.
+   *
+   * `code` is the bounded error-frame code (it feeds a counter attribute);
+   * `error` is the thrown message and rides only on the log record — never a
+   * caller-supplied frame body, per the reportDroppedFrame model.
+   */
+  reportTurnFailed(code: string, turn: TurnLogContext, error?: string): void {
+    this.turnsFailed.add(1, { code });
+    this.log.emit({
+      // A busy session is the client racing itself; everything else is a
+      // failure the operator should see.
+      severityText: code === "SESSION_BUSY" ? "WARN" : "ERROR",
+      body: "turn failed",
+      attributes: {
+        code,
+        ...turnLogAttributes(turn),
+        ...(error ? { error } : {}),
+      },
+    });
+  }
+
+  /**
+   * Record a socket that closed abnormally. This is the transport-error signal
+   * available on Bun: hono's Bun adapter dispatches only open/message/close
+   * (never WSEvents.onError, and Bun's ServerWebSocket has no error callback),
+   * so a transport failure surfaces as a close with an abnormal code.
+   */
+  reportAbnormalClose(code: number): void {
+    this.wsErrors.add(1, { "close.code": code });
+    this.log.emit({
+      severityText: "WARN",
+      body: "websocket closed abnormally",
+      attributes: { "close.code": code },
+    });
   }
 
   /** Fan a frame out to every attached client (size-bounded per frame). */
   sendToClients(msg: ServerMessage): void {
-    this.clients.broadcast(msg);
+    this.clients.broadcast(msg, () => {
+      // Outbound counterpart of reportDroppedFrame: the peer never saw this
+      // frame. Counted only — a dead socket would otherwise WARN per frame
+      // until its onClose prunes it.
+      this.framesDropped.add(1, {
+        reason: "broadcast_send_failed",
+        direction: "outbound",
+      });
+    });
   }
 
   /** Send a frame to one specific socket (size-bounded). */

package/src/ws/rate-limit.ts

@@ -0,0 +1,86 @@
+/**
+ * Per-connection inbound frame metering.
+ *
+ * Frames are already size-, cardinality- and depth-bounded, so the cheap CPU
+ * amplification is closed. What was left unmetered is VOLUME: a flood of
+ * individually valid frames behind the auth guard. That is a smaller threat
+ * than an unauthenticated one — it needs a session — but "the attacker had to
+ * log in first" is not a bound.
+ *
+ * A token bucket rather than a fixed window, because the traffic this protects
+ * is bursty by nature and a fixed window punishes the wrong thing. Opening the
+ * app fires a handful of frames at once (resume, environment, a first
+ * message); an approval storm during a busy turn is a dozen in a second. Both
+ * are legitimate and both fit in the burst. What does not fit is a sustained
+ * rate no human produces.
+ *
+ * Inbound traffic is naturally low-rate — deltas flow the other way — so the
+ * defaults are generous enough that no real client should ever see one of
+ * these, which is the property that makes it safe to turn on by default.
+ */
+
+/** Injected so tests need no real time. */
+export interface Clock {
+  now(): number;
+}
+
+const realClock: Clock = { now: () => Date.now() };
+
+export interface RateLimitOptions {
+  /** Sustained frames per second. */
+  ratePerSecond: number;
+  /** Frames absorbable in one burst before the sustained rate applies. */
+  burst: number;
+  clock?: Clock;
+}
+
+export interface RateLimitDecision {
+  allowed: boolean;
+  /** Tokens left after this decision, for reporting. */
+  remaining: number;
+}
+
+/**
+ * One bucket per connection. Not shared, and not keyed by anything
+ * user-supplied: a map keyed by session or IP is itself an unbounded
+ * allocation an attacker controls, which is the classic way a rate limiter
+ * becomes the memory-exhaustion bug it was added to prevent. This lives on the
+ * socket and dies with it.
+ */
+export class FrameRateLimiter {
+  private tokens: number;
+  private lastRefill: number;
+  private readonly ratePerMs: number;
+  private readonly burst: number;
+  private readonly clock: Clock;
+  /** Frames refused since the connection opened. */
+  private rejected = 0;
+
+  constructor(options: RateLimitOptions) {
+    this.burst = Math.max(1, options.burst);
+    this.ratePerMs = Math.max(0, options.ratePerSecond) / 1000;
+    this.clock = options.clock ?? realClock;
+    this.tokens = this.burst;
+    this.lastRefill = this.clock.now();
+  }
+
+  /** Take one token. False means this frame should be refused. */
+  take(): RateLimitDecision {
+    const now = this.clock.now();
+    const elapsed = Math.max(0, now - this.lastRefill);
+    this.lastRefill = now;
+    this.tokens = Math.min(this.burst, this.tokens + elapsed * this.ratePerMs);
+
+    if (this.tokens < 1) {
+      this.rejected++;
+      return { allowed: false, remaining: 0 };
+    }
+    this.tokens -= 1;
+    return { allowed: true, remaining: Math.floor(this.tokens) };
+  }
+
+  /** How many frames this connection has had refused. */
+  get rejectedCount(): number {
+    return this.rejected;
+  }
+}