@schlessera/brain-ui-server 0.34.1 → 0.35.0

package/src/ws/run-session.ts

@@ -9,7 +9,7 @@ import { withSessionId, withTurnScope } from "./frames.js";
 import { makeBridge, emitTurnError } from "./bridge.js";
 import { createTurnRecorder, type TurnRecorder } from "../activity/recorder.js";
 import { resolveTurnTarget } from "./routing.js";
-import type { QueuedFollowUp, RunningTurn } from "./turns.js";
+import type { AuthorizationContext, QueuedFollowUp, RunningTurn } from "./turns.js";
 import { queuedBytes, queuedFollowUpBytes } from "./turns.js";
 import {
   MAX_SESSION_QUEUE,
@@ -24,19 +24,71 @@ import {
  * order. The host owns per-turn cancellation (AbortController + timeout). Never
  * throws — runtime failures are emitted as error frames.
  */
-export async function runSession(
+/**
+ * Write out cancellations that were accepted before a recorder existed.
+ *
+ * A cancel can land while routing or billing is still in flight, when there is
+ * nothing to record it on yet, so the actor is buffered on the turn. Both exits
+ * from that window have to drain it — the ordinary one, which then has a
+ * recorder, and the revoked one, which must build a throwaway recorder rather
+ * than drop a decision someone actually made.
+ */
+function recordPendingCancellations(
   host: WsHost,
-  initial: {
-    text: string;
-    sessionId?: string;
-    attachments: ChatImageAttachment[];
-    providerId?: string;
-    client?: ClientEnvironment;
-    /** Client correlation id for a new conversation; echoed on session_info. */
-    draftId?: string;
+  turn: RunningTurn,
+  billing: Awaited<ReturnType<typeof resolveRunBilling>> | undefined
+): void {
+  if (turn.pendingCancellationPrincipalIds.length === 0) return;
+  const recorder =
+    turn.recorder ??
+    (host.activity
+      ? createTurnRecorder(
+          {
+            store: host.activity.store,
+            onWrite: () => host.activity!.stream.pump(),
+            log: host.log,
+          },
+          {
+            turnId: turn.turnId,
+            sessionId: turn.sessionId,
+            principalId: turn.principalId,
+            ...billing,
+          }
+        )
+      : undefined);
+  for (const principalId of turn.pendingCancellationPrincipalIds) {
+    recorder?.recordCancellation(principalId);
+  }
+  turn.pendingCancellationPrincipalIds.length = 0;
+}
+
+type RunSessionInput = {
+  authorization: AuthorizationContext;
+  text: string;
+  sessionId?: string;
+  attachments: ChatImageAttachment[];
+  providerId?: string;
+  client?: ClientEnvironment;
+  /** Client correlation id for a new conversation; echoed on session_info. */
+  draftId?: string;
+};
+
+export async function runSession(host: WsHost, initial: RunSessionInput): Promise<void> {
+  const releaseAuthorization = initial.authorization.retain();
+  try {
+    await runRetainedSession(host, initial);
+  } finally {
+    releaseAuthorization();
   }
+}
+
+async function runRetainedSession(
+  host: WsHost,
+  initial: RunSessionInput
 ): Promise<void> {
   const { coordinator } = host;
+  const starting = { queue: [] as QueuedFollowUp[], cancelled: false };
+  if (initial.sessionId) coordinator.startingBySession.set(initial.sessionId, starting);
   coordinator.startingSessions += 1;
   let target: Awaited<ReturnType<typeof resolveTurnTarget>>;
   try {
@@ -55,8 +107,14 @@ export async function runSession(
     return;
   } finally {
     coordinator.startingSessions = Math.max(0, coordinator.startingSessions - 1);
+    if (initial.sessionId && coordinator.startingBySession.get(initial.sessionId) === starting) {
+      coordinator.startingBySession.delete(initial.sessionId);
+    }
   }
 
+  if (starting.cancelled) return;
+  if (!initial.authorization.valid) return;
+
   const { backend, profileId: initialProfileId } = target;
   if (target.droppedPin) {
     host.log.emit({
@@ -72,6 +130,8 @@ export async function runSession(
     });
   }
   const turn: RunningTurn = {
+    principalId: initial.authorization.principalId,
+    authorization: initial.authorization,
     sessionId: initial.sessionId ?? null,
     turnId: crypto.randomUUID(),
     // Only meaningful when the client had no session id to send.
@@ -80,7 +140,8 @@ export async function runSession(
     backend,
     abortController: new AbortController(),
     timeoutHandle: setTimeout(() => {}, 0),
-    queue: [],
+    queue: starting.queue,
+    pendingCancellationPrincipalIds: [],
     cancelled: false,
     lastResult: null,
   };
@@ -91,13 +152,30 @@ export async function runSession(
   let profileId = initialProfileId;
   let resumeId = initial.sessionId;
   let next: QueuedFollowUp | null = {
+    principalId: initial.authorization.principalId,
+    authorization: initial.authorization,
     text: initial.text,
     attachments: initial.attachments,
     ...(initial.client ? { client: initial.client } : {}),
+    releaseAuthorization: () => {},
   };
+  let releaseActiveAuthorization: (() => void) | undefined;
 
   try {
     while (next && !turn.cancelled) {
+      releaseActiveAuthorization = next.releaseAuthorization;
+      turn.principalId = next.principalId;
+      turn.authorization = next.authorization;
+      if (!turn.authorization.valid) {
+        releaseActiveAuthorization();
+        releaseActiveAuthorization = undefined;
+        next = turn.queue.shift() ?? null;
+        if (next) {
+          turn.turnId = crypto.randomUUID();
+          turn.lastResult = null;
+        }
+        continue;
+      }
       const { text, attachments, client } = next;
       next = null;
 
@@ -129,6 +207,32 @@ export async function runSession(
       const billing = host.activity
         ? await resolveRunBilling(host.registry, backend.id, profileId)
         : undefined;
+      // Revocation can land while routing or billing is in flight. This is the
+      // final await boundary before startTurn, so an invalid principal never
+      // reaches the backend while an already-running turn remains untouched.
+      // Expiry is authoritative, not advisory: the sweep that turns an expired
+      // principal into a revocation runs on a timer, and a follow-up dequeued
+      // inside that window would otherwise start — reviewer measured one
+      // starting 12ms past expiry — and then survive, because a started turn is
+      // deliberately never aborted. Enforce it at the same boundary that
+      // enforces revocation, so the timer only has to close the socket.
+      host.expireAuthorizationContexts();
+      if (!turn.authorization.valid) {
+        clearTimeout(timeoutHandle);
+        releaseActiveAuthorization();
+        releaseActiveAuthorization = undefined;
+        // A cancellation accepted while billing was in flight is already a
+        // decision someone made; losing it because the turn's own principal
+        // was revoked a moment later would erase the actor, not the turn. Drain
+        // the buffer before leaving.
+        recordPendingCancellations(host, turn, billing);
+        if (turn.queue.length > 0) {
+          next = turn.queue.shift()!;
+          turn.turnId = crypto.randomUUID();
+          turn.lastResult = null;
+        }
+        continue;
+      }
       const recorder: TurnRecorder | undefined = host.activity
         ? createTurnRecorder(
             {
@@ -136,9 +240,19 @@ export async function runSession(
               onWrite: () => host.activity!.stream.pump(),
               log: host.log,
             },
-            { turnId: turn.turnId, sessionId: turn.sessionId, ...billing }
+            {
+              turnId: turn.turnId,
+              sessionId: turn.sessionId,
+              principalId: turn.principalId,
+              ...billing,
+            }
           )
         : undefined;
+      turn.recorder = recorder;
+      for (const principalId of turn.pendingCancellationPrincipalIds) {
+        recorder?.recordCancellation(principalId);
+      }
+      turn.pendingCancellationPrincipalIds.length = 0;
       const bridge = makeBridge(host, turn, text, backend.id, recorder);
       const startedAt = Date.now();
       host.reportTurnStarted(turn);
@@ -182,6 +296,9 @@ export async function runSession(
         );
       } finally {
         clearTimeout(timeoutHandle);
+        turn.recorder = undefined;
+        releaseActiveAuthorization();
+        releaseActiveAuthorization = undefined;
       }
 
       // Subsequent (queued) turns resume the now-known session and keep its
@@ -200,7 +317,12 @@ export async function runSession(
       }
     }
   } finally {
-    if (turn.sessionId) coordinator.bySession.delete(turn.sessionId);
+    releaseActiveAuthorization?.();
+    next?.releaseAuthorization();
+    for (const entry of turn.queue.splice(0)) entry.releaseAuthorization();
+    if (turn.sessionId && coordinator.bySession.get(turn.sessionId) === turn) {
+      coordinator.bySession.delete(turn.sessionId);
+    }
     coordinator.running.delete(turn);
     coordinator.drainPendingForTurn(turn, "Session ended");
   }
@@ -241,11 +363,70 @@ function formatMb(bytes: number): string {
   return `${(bytes / 1024 / 1024).toFixed(1)} MB`;
 }
 
+/** Apply the same queue budgets during routing and during a running turn. */
+function queueFollowUp(host: WsHost, ws: WSContext, sessionId: string, slot: { queue: QueuedFollowUp[] }, entry: QueuedFollowUp): void {
+  const parked = queuedBytes(slot);
+  const incoming = queuedFollowUpBytes(entry);
+  // A single message can never exceed the budget on its own: the frame cap
+  // is 12 MB and a message's attachments are capped well below that, so an
+  // empty queue always has room and this cannot wedge.
+  const overBudget = parked + incoming > QUEUE_MAX_BYTES;
+  const overDepth = slot.queue.length >= MAX_SESSION_QUEUE;
+
+  if (overBudget || overDepth) {
+    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: overBudget
+            ? `This session's queue is full (${formatMb(parked)} of ${formatMb(QUEUE_MAX_BYTES)}; this message needs ${formatMb(incoming)}). Wait for it to catch up.`
+            : `This session's queue is full (${MAX_SESSION_QUEUE} messages). Wait for it to catch up.`,
+        },
+        sessionId
+      )
+    );
+    // The entry never joins a queue, so nothing else will ever release its
+    // lease: a refused follow-up would otherwise hold its principal
+    // discoverable forever and the registry would never empty.
+    entry.releaseAuthorization();
+    return;
+  }
+
+  // Queue it as the session's next turn; report queued immediately.
+  slot.queue.push(entry);
+  const total = parked + incoming;
+  // Accepted, but heavy enough that the sender should know before they hit
+  // the wall — every queued byte is held in this process until its turn runs.
+  const detail =
+    total >= QUEUE_WARN_BYTES
+      ? `Queue is holding ${formatMb(total)} across ${slot.queue.length} messages (limit ${formatMb(QUEUE_MAX_BYTES)}).`
+      : undefined;
+  if (detail) {
+    host.log.emit({
+      severityText: "WARN",
+      body: "session follow-up queue is heavy",
+      attributes: {
+        "session.id": sessionId,
+        size: formatMb(total),
+        queued: slot.queue.length,
+      },
+    });
+  }
+  host.sendToClients(
+    withSessionId({ type: "status", status: "queued", ...(detail ? { detail } : {}) }, sessionId)
+  );
+}
+
 /** Dispatch a chat_message: follow-up to a running session, or a new session. */
 export async function handleChatMessage(
   host: WsHost,
   ws: WSContext,
   msg: {
+    authorization: AuthorizationContext;
     text: string;
     sessionId?: string;
     attachments: ChatImageAttachment[];
@@ -254,76 +435,67 @@ export async function handleChatMessage(
     draftId?: string;
   }
 ): Promise<void> {
-  const { text, attachments, sessionId, providerId: requestedProviderId, client, draftId } = msg;
+  const {
+    authorization,
+    text,
+    attachments,
+    sessionId,
+    providerId: requestedProviderId,
+    client,
+    draftId,
+  } = msg;
   const { coordinator } = host;
   const runningTurn = sessionId ? coordinator.bySession.get(sessionId) : undefined;
 
+  const starting = sessionId ? coordinator.startingBySession.get(sessionId) : undefined;
+  if (starting && sessionId) {
+    if (starting.cancelled) {
+      host.sendMessage(ws, { type: "error", code: "SESSION_BUSY", sessionId, message: "This session is cancelling. Wait before sending again." });
+    } else {
+      queueFollowUp(host, ws, sessionId, starting, {
+        principalId: authorization.principalId,
+        authorization,
+        text,
+        attachments,
+        ...(client ? { client } : {}),
+        releaseAuthorization: authorization.retain(),
+      });
+    }
+    return;
+  }
+
   if (runningTurn && sessionId) {
     const backend = runningTurn.backend;
-    // Message to a session whose turn is running = a follow-up.
-    if (backend.capabilities.followUp && backend.followUp) {
+    // Drain older messages queued during routing before allowing native
+    // injection to overtake them.
+    if (backend.capabilities.followUp && backend.followUp && runningTurn.queue.length === 0) {
       // Inject into the running turn; frames flow through its bridge. The
       // device snapshot is deliberately not forwarded: a follow-up joins a
       // turn whose system prompt was already built and cannot be revised.
-      backend.followUp({ sessionId, prompt: text, attachments }).catch((err) => {
-        const message = err instanceof Error ? err.message : String(err);
-        host.reportTurnFailed("FOLLOWUP_FAILED", runningTurn, message);
-        host.sendToClients(
-          withTurnScope({ type: "error", code: "FOLLOWUP_FAILED", message }, runningTurn)
-        );
-      });
+      runningTurn.recorder?.recordFollowUp(authorization.principalId);
+      void (async () => {
+        const releaseFollowUp = authorization.retain();
+        try {
+          await backend.followUp!({ sessionId, prompt: text, attachments });
+        } catch (err) {
+          const message = err instanceof Error ? err.message : String(err);
+          host.reportTurnFailed("FOLLOWUP_FAILED", runningTurn, message);
+          host.sendToClients(
+            withTurnScope({ type: "error", code: "FOLLOWUP_FAILED", message }, runningTurn)
+          );
+        } finally {
+          releaseFollowUp();
+        }
+      })();
     } else {
-      const entry: QueuedFollowUp = { text, attachments, ...(client ? { client } : {}) };
-      const parked = queuedBytes(runningTurn);
-      const incoming = queuedFollowUpBytes(entry);
-      // A single message can never exceed the budget on its own: the frame cap
-      // is 12 MB and a message's attachments are capped well below that, so an
-      // empty queue always has room and this cannot wedge.
-      const overBudget = parked + incoming > QUEUE_MAX_BYTES;
-      const overDepth = runningTurn.queue.length >= MAX_SESSION_QUEUE;
-
-      if (overBudget || overDepth) {
-        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: overBudget
-                ? `This session's queue is full (${formatMb(parked)} of ${formatMb(QUEUE_MAX_BYTES)}; this message needs ${formatMb(incoming)}). Wait for it to catch up.`
-                : `This session's queue is full (${MAX_SESSION_QUEUE} messages). Wait for it to catch up.`,
-            },
-            sessionId
-          )
-        );
-        return;
-      }
-
-      // Queue it as the session's next turn; report queued immediately.
-      runningTurn.queue.push(entry);
-      const total = parked + incoming;
-      // Accepted, but heavy enough that the sender should know before they hit
-      // the wall — every queued byte is held in this process until its turn runs.
-      const detail =
-        total >= QUEUE_WARN_BYTES
-          ? `Queue is holding ${formatMb(total)} across ${runningTurn.queue.length} messages (limit ${formatMb(QUEUE_MAX_BYTES)}).`
-          : undefined;
-      if (detail) {
-        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)
-      );
+      queueFollowUp(host, ws, sessionId, runningTurn, {
+        principalId: authorization.principalId,
+        authorization,
+        text,
+        attachments,
+        ...(client ? { client } : {}),
+        releaseAuthorization: authorization.retain(),
+      });
     }
     return;
   }
@@ -344,6 +516,7 @@ export async function handleChatMessage(
     host.sendToClients(withSessionId({ type: "status", status: "thinking" }, sessionId));
   }
   void runSession(host, {
+    authorization,
     text,
     sessionId,
     attachments,

package/src/ws/shrink.ts

@@ -32,10 +32,9 @@ interface ShrinkPass {
 }
 
 /**
- * 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.
+ * Progressive shrink passes. Each tightens display strings and array tails;
+ * the loop stops when the serialized UTF-8 payload fits. Identity strings
+ * remain intact, and typed arrays never receive an untyped elision marker.
  */
 const SHRINK_PASSES: readonly ShrinkPass[] = [
   { stringCap: 64 * 1024, arrayLimit: 256 },
@@ -49,29 +48,64 @@ const SHRINK_PASSES: readonly ShrinkPass[] = [
 
 const STRING_ELISION_RESERVE = 80;
 
+/** JSON's UTF-8 wire size, including escaped control characters. */
+export function jsonBytes(value: unknown): number {
+  return Buffer.byteLength(JSON.stringify(value), "utf8");
+}
+
+// These strings identify protocol objects rather than display text.
+const IDENTITY_FIELDS = new Set([
+  "id", "type", "kind", "role", "name", "sessionId", "turnId", "draftId",
+  "requestId", "toolUseId", "parentToolUseId", "providerId", "backendId",
+]);
+
 /**
  * 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 {
+function shrinkWalk(value: unknown, stringCap: number, arrayLimit: number, opaque = false): unknown {
   if (typeof value === "string") {
     if (value.length <= stringCap) return value;
-    const headLen = Math.max(0, stringCap - STRING_ELISION_RESERVE);
+    let headLen = Math.max(0, stringCap - STRING_ELISION_RESERVE);
+    // Never split a UTF-16 surrogate pair at the truncation boundary.
+    const last = value.charCodeAt(headLen - 1);
+    if (last >= 0xd800 && last <= 0xdbff) headLen--;
     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]`;
+    const out: unknown[] = new Array(keep);
+    for (let i = 0; i < keep; i++) out[i] = shrinkWalk(value[i], stringCap, arrayLimit, opaque);
     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);
+    // Arbitrary tool input can exceed the cap through key count alone. Clip
+    // its records recursively, while preserving fields on protocol objects.
+    const keys = Object.keys(src);
+    for (const k of opaque ? keys.slice(0, arrayLimit) : keys) {
+      out[k] = !opaque && IDENTITY_FIELDS.has(k) && typeof src[k] === "string"
+        ? src[k]
+        : shrinkWalk(src[k], stringCap, arrayLimit, opaque || k === "input");
+    }
+    // History parts refer to toolCalls by array index. Head clipping preserves
+    // retained indices, but references to removed tools must go with them.
+    if (!opaque && (src.role === "assistant" || src.role === "user") && Array.isArray(out.toolCalls)) {
+      const toolCount = out.toolCalls.length;
+      if (Array.isArray(out.parts)) {
+        out.parts = out.parts.filter((part) =>
+          !part || part.kind !== "tool" ||
+          (Number.isInteger(part.toolIndex) && part.toolIndex >= 0 && part.toolIndex < toolCount)
+        );
+      }
+      if (Array.isArray(src.toolCalls) && src.toolCalls.length > toolCount && typeof out.content === "string") {
+        const notice = `\n…[${src.toolCalls.length - toolCount} tool calls elided]`;
+        out.content += notice;
+        if (Array.isArray(out.parts)) out.parts.push({ kind: "text", text: notice });
+      }
+    }
     return out;
   }
   return value;
@@ -84,12 +118,12 @@ function shrinkWalk(value: unknown, stringCap: number, arrayLimit: number): unkn
  * 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;
+export function shrinkForReplication<T>(value: T, maxBytes = MAX_WS_MESSAGE_BYTES): T {
+  if (jsonBytes(value) <= maxBytes) 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;
+    if (jsonBytes(shrunk) <= maxBytes) return shrunk as T;
   }
   return shrunk as T;
 }