@schlessera/brain-ui-server 0.34.0 → 0.35.0

package/src/ws/turns.ts

@@ -11,12 +11,31 @@ import type {
   ChatImageAttachment,
   ClientEnvironment,
 } from "@schlessera/brain-ui-sdk/protocol";
+import type { TurnRecorder } from "../activity/recorder.js";
+
+const authorizationContextBrand: unique symbol = Symbol("AuthorizationContext");
+
+/** Mutable admission decision shared by every asynchronous path from one socket. */
+export interface AuthorizationContext {
+  readonly [authorizationContextBrand]: true;
+  readonly principalId: string;
+  readonly expiresAt: number;
+  valid: boolean;
+  /** Keep this context discoverable until the returned release is called. */
+  retain(): () => void;
+  /** Release the reference acquired when the coordinator opened this context. */
+  release(): void;
+}
 
 export interface QueuedFollowUp {
+  principalId: string;
+  authorization: AuthorizationContext;
   text: string;
   attachments: ChatImageAttachment[];
   /** Device snapshot taken when the message was sent, not when it runs. */
   client?: ClientEnvironment;
+  /** Queue-owned authorization lease, transferred to the runner on dequeue. */
+  releaseAuthorization: () => void;
 }
 
 /**
@@ -30,7 +49,9 @@ export interface QueuedFollowUp {
  * on the wire.) The client snapshot is a handful of short fields and is not
  * worth walking.
  */
-export function queuedFollowUpBytes(entry: QueuedFollowUp): number {
+export function queuedFollowUpBytes(
+  entry: Pick<QueuedFollowUp, "text" | "attachments">
+): number {
   let bytes = Buffer.byteLength(entry.text, "utf-8");
   for (const attachment of entry.attachments) {
     bytes += attachment.data.length;
@@ -39,7 +60,7 @@ export function queuedFollowUpBytes(entry: QueuedFollowUp): number {
 }
 
 /** Total bytes currently parked in a session's follow-up queue. */
-export function queuedBytes(turn: RunningTurn): number {
+export function queuedBytes(turn: Pick<RunningTurn, "queue">): number {
   let bytes = 0;
   for (const entry of turn.queue) bytes += queuedFollowUpBytes(entry);
   return bytes;
@@ -51,6 +72,16 @@ export function queuedBytes(turn: RunningTurn): number {
  * same slot (same sessionId) after the current one resolves.
  */
 export interface RunningTurn {
+  /** Principal responsible for the CURRENT turn in this session slot. */
+  principalId: string;
+  authorization: AuthorizationContext;
+  /** Activity recorder for the CURRENT turn, once startup reaches the backend. */
+  recorder?: TurnRecorder;
+  /**
+   * Cancelling actors accepted before the recorder exists (while billing is
+   * resolving). Drained into immutable activity events at recorder creation.
+   */
+  pendingCancellationPrincipalIds: string[];
   sessionId: string | null; // null until session_info resolves it (new session)
   /** Host-minted id of the CURRENT turn in this slot; re-minted per queued follow-up. */
   turnId: string;
@@ -88,7 +119,10 @@ export interface PendingApproval {
    * locks; an approval must survive that and reappear, not silently die.
    */
   request: PermissionRequest;
-  resolve: (decision: PermissionDecision) => void;
+  resolve: (
+    decision: PermissionDecision,
+    response?: { principalId: string; always?: boolean }
+  ) => void;
 }
 
 export interface PendingAskUser {
@@ -132,12 +166,17 @@ export class TurnCoordinator {
   readonly running = new Set<RunningTurn>();
   readonly bySession = new Map<string, RunningTurn>();
   startingSessions = 0;
+  /** Reserve a known session before asynchronous backend routing. */
+  readonly startingBySession = new Map<string, { queue: QueuedFollowUp[]; cancelled: boolean }>();
 
   readonly pendingApprovals = new Map<string, PendingApproval>();
   readonly pendingAskUser = new Map<string, PendingAskUser>();
   readonly pendingLocation = new Map<string, PendingLocation>();
   readonly pendingMask = new Map<string, PendingMask>();
 
+  /** Every authorization context with at least one live owner or async lease. */
+  readonly authorizationRegistry = new Map<AuthorizationContext, number>();
+
   private locationCounter = 0;
   private maskCounter = 0;
 
@@ -154,9 +193,89 @@ export class TurnCoordinator {
     return this.running.size > 0;
   }
 
+  /** Construct and synchronously register a context with its initial reference. */
+  openAuthorization(input: {
+    principalId: string;
+    expiresAt: number;
+    valid: boolean;
+  }): AuthorizationContext {
+    let initialReleased = false;
+    const releaseReference = () => {
+      const count = this.authorizationRegistry.get(authorization);
+      if (count === undefined) return;
+      if (count === 1) this.authorizationRegistry.delete(authorization);
+      else this.authorizationRegistry.set(authorization, count - 1);
+    };
+    const authorization: AuthorizationContext = {
+      [authorizationContextBrand]: true,
+      ...input,
+      retain: () => {
+        const count = this.authorizationRegistry.get(authorization);
+        if (count === undefined) {
+          throw new Error("Cannot retain a released authorization context");
+        }
+        this.authorizationRegistry.set(authorization, count + 1);
+        let released = false;
+        return () => {
+          if (released) return;
+          released = true;
+          releaseReference();
+        };
+      },
+      release: () => {
+        if (initialReleased) return;
+        initialReleased = true;
+        releaseReference();
+      },
+    };
+    this.authorizationRegistry.set(authorization, 1);
+    return authorization;
+  }
+
+  /** Add principals whose registered authorization context has expired. */
+  collectExpiredPrincipalIds(now: number, ids: Set<string>): void {
+    for (const authorization of this.authorizationRegistry.keys()) {
+      if (authorization.valid && authorization.expiresAt <= now) {
+        ids.add(authorization.principalId);
+      }
+    }
+  }
+
+  /** Invalidate every registered context belonging to one of these principals. */
+  invalidateAuthorizations(principalIds: ReadonlySet<string>): void {
+    for (const authorization of this.authorizationRegistry.keys()) {
+      if (principalIds.has(authorization.principalId)) authorization.valid = false;
+    }
+  }
+
+  /**
+   * Record revocation without aborting running work. Unstarted queued work is
+   * removed per sender, while a current turn merely carries the invalid marker.
+   */
+  revokePrincipals(principalIds: ReadonlySet<string>): RunningTurn[] {
+    const affectedRunning: RunningTurn[] = [];
+    this.invalidateAuthorizations(principalIds);
+    for (const turn of this.running) {
+      if (principalIds.has(turn.principalId)) {
+        affectedRunning.push(turn);
+      }
+      turn.queue = turn.queue.filter((entry) => {
+        if (!principalIds.has(entry.principalId)) return true;
+        entry.releaseAuthorization();
+        return false;
+      });
+    }
+    return affectedRunning;
+  }
+
   /** Cancel every running turn (used on shutdown). */
   cancelAll(reason: string): boolean {
-    if (this.running.size === 0) return false;
+    const hadWork = this.running.size > 0 || this.startingBySession.size > 0;
+    for (const starting of this.startingBySession.values()) {
+      starting.cancelled = true;
+      starting.queue.length = 0;
+    }
+    if (!hadWork) return false;
     for (const turn of [...this.running]) this.cancelTurn(turn, reason);
     return true;
   }
@@ -164,12 +283,18 @@ export class TurnCoordinator {
   /** Cancel a session's current turn and drop its queued follow-ups. */
   cancelTurn(turn: RunningTurn, reason: string): void {
     turn.cancelled = true;
-    turn.queue.length = 0;
+    for (const entry of turn.queue.splice(0)) entry.releaseAuthorization();
     clearTimeout(turn.timeoutHandle);
     turn.abortController.abort();
     this.drainPendingForTurn(turn, reason);
   }
 
+  /** Record now, or retain the actor across the pre-recorder startup window. */
+  recordCancellation(turn: RunningTurn, principalId: string): void {
+    if (turn.recorder) turn.recorder.recordCancellation(principalId);
+    else turn.pendingCancellationPrincipalIds.push(principalId);
+  }
+
   /** Resolve/reject every pending interactive request belonging to one turn. */
   drainPendingForTurn(turn: RunningTurn, reason: string): void {
     for (const [id, p] of this.pendingApprovals) {
@@ -238,7 +363,13 @@ export class TurnCoordinator {
     for (const turn of [...this.running]) {
       clearTimeout(turn.timeoutHandle);
       turn.abortController.abort();
+      for (const entry of turn.queue.splice(0)) entry.releaseAuthorization();
+    }
+    for (const starting of this.startingBySession.values()) {
+      starting.cancelled = true;
+      starting.queue.length = 0;
     }
+    this.startingBySession.clear();
     this.running.clear();
     this.bySession.clear();
     this.pendingApprovals.clear();