@schlessera/brain-ui-server 0.15.0 → 0.16.0

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 });
       }

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,19 @@ 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 };
 }
 
 /**
@@ -66,6 +81,13 @@ 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"]
+  >;
 
   constructor(options: WsHostOptions) {
     this.registry = options.registry;
@@ -74,6 +96,38 @@ 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");
+    this.framesDropped = this.observability
+      .meter("ws")
+      .createCounter("ws.frames.dropped", {
+        description: "Inbound frames refused before reaching a handler",
+      });
+    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 },
+    });
   }
 
   /** Fan a frame out to every attached client (size-bounded per frame). */

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;
+  }
+}

package/src/ws/run-session.ts

@@ -20,6 +20,8 @@ export async function runSession(
     attachments: ChatImageAttachment[];
     providerId?: string;
     client?: ClientEnvironment;
+    /** Client correlation id for a new conversation; echoed on session_info. */
+    draftId?: string;
   }
 ): Promise<void> {
   const { coordinator } = host;
@@ -41,10 +43,11 @@ export async function runSession(
 
   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.log.emit({
+      severityText: "WARN",
+      body: "pinned profile unavailable; running on the default profile",
+      attributes: { "session.id": initial.sessionId, profile: target.droppedPin },
+    });
     host.sendToClients({
       type: "status",
       status: "thinking",
@@ -55,6 +58,8 @@ export async function runSession(
   const turn: RunningTurn = {
     sessionId: initial.sessionId ?? null,
     turnId: crypto.randomUUID(),
+    // Only meaningful when the client had no session id to send.
+    draftId: initial.sessionId ? null : (initial.draftId ?? null),
     providerId: initialProfileId ?? null,
     backend,
     abortController: new AbortController(),
@@ -82,7 +87,11 @@ export async function runSession(
       const abortController = new AbortController();
       turn.abortController = abortController;
       const timeoutHandle = setTimeout(() => {
-        console.log("[agent] Turn timed out after", host.turnTimeoutMs, "ms");
+        host.log.emit({
+          severityText: "WARN",
+          body: "turn timed out",
+          attributes: { "timeout.ms": host.turnTimeoutMs },
+        });
         abortController.abort();
         // Reject any pending interactive request for this turn too. A bridge
         // that awaits askUser/approval without racing the abort signal would
@@ -144,9 +153,10 @@ export async function handleChatMessage(
     attachments: ChatImageAttachment[];
     providerId?: string;
     client?: ClientEnvironment;
+    draftId?: string;
   }
 ): Promise<void> {
-  const { text, attachments, sessionId, providerId: requestedProviderId, client } = msg;
+  const { text, attachments, sessionId, providerId: requestedProviderId, client, draftId } = msg;
   const { coordinator } = host;
   const runningTurn = sessionId ? coordinator.bySession.get(sessionId) : undefined;
 
@@ -204,7 +214,15 @@ export async function handleChatMessage(
           ? `Queue is holding ${formatMb(total)} across ${runningTurn.queue.length} messages (limit ${formatMb(QUEUE_MAX_BYTES)}).`
           : undefined;
       if (detail) {
-        console.warn(`[ws] session ${sessionId} queue at ${formatMb(total)} across ${runningTurn.queue.length} messages`);
+        host.log.emit({
+          severityText: "WARN",
+          body: "session follow-up queue is heavy",
+          attributes: {
+            "session.id": sessionId,
+            size: formatMb(total),
+            queued: runningTurn.queue.length,
+          },
+        });
       }
       host.sendToClients(
         withSessionId({ type: "status", status: "queued", ...(detail ? { detail } : {}) }, sessionId)
@@ -234,5 +252,6 @@ export async function handleChatMessage(
     attachments,
     providerId: requestedProviderId,
     ...(client ? { client } : {}),
+    ...(draftId ? { draftId } : {}),
   });
 }

package/src/ws/turns.ts

@@ -1,3 +1,4 @@
+import type { Logger } from "@opentelemetry/api-logs";
 import type {
   AgentBackend,
   PermissionDecision,
@@ -54,6 +55,12 @@ export interface RunningTurn {
   timeoutHandle: ReturnType<typeof setTimeout>;
   queue: QueuedFollowUp[];
   cancelled: boolean;
+  /**
+   * Correlation id the client minted for this NEW conversation, echoed back on
+   * `session_info` so the client can tell its own turn's identity from a
+   * background turn's. Null on a resumed session or a client that sent none.
+   */
+  draftId: string | null;
 }
 
 // Pending interactive requests are tagged with their turn so a per-session
@@ -93,6 +100,13 @@ export interface PendingMask {
  * explicit object with a defined lifecycle.
  */
 export class TurnCoordinator {
+  /**
+   * Where collisions are reported. Assigned by the owning WsHost so the
+   * coordinator reports through the same consumer as everything else; absent
+   * (a bare coordinator in a test) means silence.
+   */
+  log?: Logger;
+
   readonly running = new Set<RunningTurn>();
   readonly bySession = new Map<string, RunningTurn>();
   startingSessions = 0;
@@ -173,9 +187,11 @@ export class TurnCoordinator {
   ): boolean {
     const existing = map.get(id);
     if (existing && existing.turn !== turn) {
-      console.error(
-        `[ws] interactive id collision: "${id}" is already pending for another turn`
-      );
+      this.log?.emit({
+        severityText: "ERROR",
+        body: "interactive id collision: already pending for another turn",
+        attributes: { "request.id": id },
+      });
       return true;
     }
     return false;