@data-club/ai-hub-server 0.0.1

package/src/state/turnRegistry.ts

@@ -0,0 +1,210 @@
+/**
+ * In-memory broadcaster for live turn streams.
+ *
+ * Every active turn has at most one writer (the runner) and zero-or-more
+ * subscribers (the original `POST /chat` connection + every
+ * `GET /chat/:turn_id/events?from_seq=N` resumption connection still
+ * running while the turn is live).
+ *
+ * Persistence lives in `turn_events` (§11) — the registry only handles the
+ * **live attach** part of resumption. The first half of a resumption read
+ * (replay from `from_seq`) is a plain DB query.
+ *
+ * The registry is process-local — this matches the single-customer
+ * deployment model (§1). When ai-hub scales to multiple replicas in a
+ * v1.x release, this gets a Redis-backed implementation behind the same
+ * interface; the wire protocol (§9) and the resumption contract (§6.3.2)
+ * stay unchanged.
+ *
+ * @see SPECIFICATIONS.md §6.3.2 — resumption endpoint
+ * @see SPECIFICATIONS.md §6.6 — server-authoritative state machine
+ * @see SPECIFICATIONS.md §11 — `turn_events` table
+ */
+
+import type { AIHubEvent } from "@data-club/ai-hub";
+
+/** Callback signature for a live subscriber. */
+export type TurnEventListener = (event: AIHubEvent) => void;
+
+/** Returned by `subscribe()` so the caller can unsubscribe. */
+export type Unsubscribe = () => void;
+
+/**
+ * Per-turn live channel. The runner calls `publish()` once per event it
+ * emits; subscribers receive every event from the point of subscription
+ * onward. When the runner calls `close()`, every remaining subscriber's
+ * stream gets `null` so it knows to flush + close its SSE writer.
+ */
+interface ActiveTurn {
+  /** Highest-seen seq on this channel (informational; the runner is the source of truth). */
+  lastSeq: number;
+  /** Whether the runner has called `close()`. */
+  closed: boolean;
+  /** Active listeners. */
+  listeners: Set<TurnEventListener>;
+  /** Listeners notified once when the turn closes. */
+  closeListeners: Set<() => void>;
+  /**
+   * AbortController exposed to the runner — fired by
+   * `requestTurnCancellation(turnId)` so a separate `POST /chat/:id/cancel`
+   * request can short-circuit the in-process runner without having to
+   * abort the runner's originating request signal (which lives on a
+   * different `Request` instance). The runner translates this into a
+   * `handle.abort()` call on the spawned Letta Code subprocess (which
+   * sends an `interrupt` control to the CLI).
+   *
+   * Pre-migration this also drove an `agents.messages.cancel({run_ids})`
+   * call on the Letta server side via the captured `lettaRunId`. Post-
+   * migration the run-id tracking is gone — the SDK abstracts it — and
+   * cancel is `handle.abort()` only. See migration plan §4.2.
+   */
+  cancelController: AbortController;
+}
+
+/**
+ * Process-local map of live turn channels. The map prunes a channel on
+ * close as soon as every subscriber has unsubscribed (or immediately if
+ * no subscribers are attached).
+ */
+const channels = new Map<string, ActiveTurn>();
+
+/**
+ * Begins a new live turn channel.
+ *
+ * The runner calls this when it inserts the `turn` row, before it starts
+ * pumping events. Returns an opaque `TurnPublisher` the runner uses to
+ * publish + close.
+ *
+ * @param turnId - the freshly-created turn id.
+ * @returns the publisher handle.
+ */
+export function openTurnChannel(turnId: string): TurnPublisher {
+  if (channels.has(turnId)) {
+    throw new Error(`Turn ${turnId} already has an open channel.`);
+  }
+  const channel: ActiveTurn = {
+    lastSeq: -1,
+    closed: false,
+    listeners: new Set(),
+    closeListeners: new Set(),
+    cancelController: new AbortController(),
+  };
+  channels.set(turnId, channel);
+  return {
+    cancelSignal: channel.cancelController.signal,
+    publish(event: AIHubEvent) {
+      channel.lastSeq = event.seq;
+      for (const listener of channel.listeners) {
+        try {
+          listener(event);
+        } catch {
+          /* swallow — a misbehaving listener cannot break the channel */
+        }
+      }
+    },
+    close() {
+      if (channel.closed) return;
+      channel.closed = true;
+      for (const cb of channel.closeListeners) {
+        try {
+          cb();
+        } catch {
+          /* ignored */
+        }
+      }
+      channel.closeListeners.clear();
+      // Drop the channel once every subscriber unsubscribed. If nobody is
+      // listening at close time, prune immediately.
+      if (channel.listeners.size === 0) {
+        channels.delete(turnId);
+      }
+    },
+  };
+}
+
+/**
+ * Subscribes to a live turn channel.
+ *
+ * Returns `null` when the turn is not active (either never opened or
+ * already closed) — the caller should fall back to replaying persisted
+ * events from `turn_events` instead.
+ *
+ * @param turnId - the turn id to attach to.
+ * @param listener - called once per event (with `seq` already assigned).
+ * @param onClose - called once when the channel closes (after final event).
+ * @returns an unsubscribe function, or `null` when the channel is closed.
+ */
+export function subscribeToTurn(
+  turnId: string,
+  listener: TurnEventListener,
+  onClose: () => void,
+): Unsubscribe | null {
+  const channel = channels.get(turnId);
+  if (!channel || channel.closed) return null;
+  channel.listeners.add(listener);
+  channel.closeListeners.add(onClose);
+  return () => {
+    channel.listeners.delete(listener);
+    channel.closeListeners.delete(onClose);
+    if (channel.closed && channel.listeners.size === 0) {
+      channels.delete(turnId);
+    }
+  };
+}
+
+/**
+ * Returns `true` when the turn has an open channel right now.
+ *
+ * The chat handler uses this to disambiguate "turn is in flight" from
+ * "turn completed long ago" without having to query the DB twice.
+ */
+export function isTurnLive(turnId: string): boolean {
+  const channel = channels.get(turnId);
+  return !!channel && !channel.closed;
+}
+
+/**
+ * Test-only utility — wipes the registry so tests can isolate from each
+ * other. Production code never calls this.
+ */
+export function _resetTurnRegistryForTests(): void {
+  channels.clear();
+}
+
+/**
+ * The runner's handle on a live channel. Opaque on purpose — the runner
+ * is the only writer.
+ */
+export interface TurnPublisher {
+  /**
+   * AbortSignal that fires when a separate `POST /chat/:id/cancel`
+   * request calls `requestTurnCancellation(turnId)`. The runner listens
+   * to this in addition to the originating request's signal.
+   */
+  cancelSignal: AbortSignal;
+  /** Push one event to every current subscriber. */
+  publish(event: AIHubEvent): void;
+  /** Mark the channel closed. After this, `subscribeToTurn` returns null. */
+  close(): void;
+}
+
+/**
+ * Triggers cancellation of the named turn. The runner is expected to
+ * have wired its run to the publisher's `cancelSignal`; firing it from
+ * here aborts the upstream Letta call and short-circuits the loop.
+ *
+ * Idempotent: calling on a closed / unknown turn is a no-op. The
+ * `POST /chat/:id/cancel` route's user-facing 200 response covers
+ * race-conditions where the turn terminated between the request and
+ * this call.
+ *
+ * @param turnId - the turn to cancel.
+ * @returns `true` if the signal was fired, `false` if the channel is
+ *   closed or unknown.
+ */
+export function requestTurnCancellation(turnId: string): boolean {
+  const channel = channels.get(turnId);
+  if (!channel || channel.closed) return false;
+  channel.cancelController.abort();
+  return true;
+}