@cursor/july 0.1.1 → 0.1.2

package/src/internal/storage-coordinator.ts

@@ -0,0 +1,615 @@
+/**
+ * Runtime for `defineStorage` (`agent/storage.ts`).
+ *
+ * The coordinator is the single funnel between the engine's hot paths and
+ * the author's key-value sink. The sink is four functions (`put` / `get` /
+ * `delete` / `list`); the coordinator owns everything else: it mints every
+ * key from the versioned scheme ({@link storageKeys}), coalesces
+ * session records and batches event chunks per turn (or debounce window),
+ * serializes delivery on one bounded queue per agent, and isolates
+ * failures — a throwing sink is logged and its write dropped; storage
+ * must never stall or fail a turn.
+ *
+ * Scaling shape: sessions, reminders, evals, and A/Bs all reduce to keyed
+ * puts on the same queue, so one author-owned sink covers every domain,
+ * and new domains are new key prefixes — not new config surface.
+ */
+
+import type { ABMetricSample, ABSamplePersistence } from "../ab.js";
+import type { EvalRunPersistence, EvalRunSnapshot } from "../evals.js";
+import {
+  type ResolvedStoragePolicy,
+  resolveStoragePolicy,
+  type StorageContext,
+  type StorageDefinition,
+  storageKeys,
+} from "../storage.js";
+import type { JsonValue, SessionEvent, SessionRecord } from "../types.js";
+import type { ABSnapshot } from "./ab-snapshot.js";
+import { describeError } from "./describe-error.js";
+import { isReminderRecord, type ReminderRecord } from "./reminder-store.js";
+
+/** Event types that close a unit of work — flush point for turn-end batching. */
+const TURN_BOUNDARY_EVENTS: ReadonlySet<SessionEvent["type"]> = new Set([
+  "turn.completed",
+  "turn.failed",
+  "session.waiting",
+  "session.completed",
+  "session.failed",
+]);
+
+/**
+ * Max writes waiting on the delivery queue before new ones are dropped.
+ * At-most-once delivery already tolerates loss; bounding the queue keeps a
+ * slow sink from growing memory without limit under sustained traffic.
+ */
+const MAX_PENDING_OPS = 1_000;
+
+/**
+ * How long {@link StorageCoordinator.close} waits for the sink to
+ * drain before giving up. Keeps a hung sink from stalling shutdown past
+ * the pod grace period (where a SIGKILL would lose the flush anyway).
+ */
+const CLOSE_DRAIN_TIMEOUT_MS = 10_000;
+
+/** Debounced event batching flushes early past this many buffered events. */
+const DEBOUNCE_MAX_BATCH = 200;
+
+/**
+ * Aggregate A/B snapshots are refreshed at most this often — the
+ * playground polls `GET /v1/abs` every few seconds, and the snapshot is a
+ * convenience backfill, not a source of truth.
+ */
+const AB_SNAPSHOT_MIN_INTERVAL_MS = 60_000;
+
+/** One queued sink call. */
+type StorageOp =
+  | { op: "put"; key: string; value: JsonValue }
+  | { op: "delete"; key: string };
+
+export interface StorageCoordinatorOptions {
+  definition: StorageDefinition;
+  agentName: string;
+  projectRoot: string;
+  logger: (line: string) => void;
+}
+
+interface SessionBuffer {
+  events: SessionEvent[];
+  /** Coalesced latest record; only the newest at flush time is written. */
+  record?: SessionRecord;
+  debounceTimer?: ReturnType<typeof setTimeout>;
+}
+
+export class StorageCoordinator {
+  readonly policy: ResolvedStoragePolicy;
+
+  private readonly definition: StorageDefinition;
+  private readonly agentName: string;
+  private readonly projectRoot: string;
+  private readonly logger: (line: string) => void;
+  /** Per-session pending state (events + coalesced record). */
+  private readonly buffers = new Map<string, SessionBuffer>();
+  /**
+   * Last continuation key written to the sink per session, so a token
+   * change emits a delete for the stale index entry alongside the new put.
+   */
+  private readonly continuationIndex = new Map<string, string | null>();
+  /** Single delivery chain: writes reach the sink serialized, in order. */
+  private queue: Promise<void> = Promise.resolve();
+  /** Ops on the queue not yet delivered (bounded by MAX_PENDING_OPS). */
+  private pending = 0;
+  /** Ops dropped because the queue was full. */
+  private dropped = 0;
+  private lastAbSnapshotAt = 0;
+  private closed = false;
+
+  constructor(options: StorageCoordinatorOptions) {
+    this.definition = options.definition;
+    this.policy = resolveStoragePolicy(options.definition.policy);
+    this.agentName = options.agentName;
+    this.projectRoot = options.projectRoot;
+    this.logger = options.logger;
+  }
+
+  // ==========================================================================
+  // Sessions (writes)
+  // ==========================================================================
+
+  /** Durable session-record update (engine calls after every store save). */
+  sessionRecord(record: SessionRecord): void {
+    if (this.closed) {
+      return;
+    }
+    // Coalesced: only the newest record at flush time reaches the sink.
+    this.buffer(record.sessionId).record = record;
+  }
+
+  /** One appended session event (engine calls from its dispatch funnel). */
+  event(event: SessionEvent): void {
+    if (this.closed) {
+      return;
+    }
+    const buffer = this.buffer(event.sessionId);
+    buffer.events.push(event);
+    const events = this.policy.events;
+    if (events !== "turnEnd") {
+      // Debounce spans turn boundaries by design ("flush when the session
+      // goes quiet", not once per turn); the coalesced record rides the
+      // same debounced flush.
+      if (buffer.events.length >= DEBOUNCE_MAX_BATCH) {
+        this.flushSession(event.sessionId);
+      } else {
+        this.armDebounce(event.sessionId, buffer, events.debounceMs);
+      }
+    }
+    // Turn boundaries do NOT flush here: the boundary event's own handlers
+    // (and their record updates) haven't run yet. The engine calls
+    // eventDispatched() once dispatch settles — that's the flush point.
+  }
+
+  /**
+   * A boundary event's dispatch settled: channel/hook handlers ran and
+   * their record updates (channel state, continuation tokens) are
+   * buffered. This — not the append — is the turn-end flush point, so the
+   * persisted session record includes the boundary event's own handler
+   * mutations. Debounced batching intentionally ignores boundaries.
+   */
+  eventDispatched(event: SessionEvent): void {
+    if (
+      this.closed ||
+      this.policy.events !== "turnEnd" ||
+      !TURN_BOUNDARY_EVENTS.has(event.type)
+    ) {
+      return;
+    }
+    this.flushSession(event.sessionId);
+  }
+
+  /** Flush everything buffered for one session (ordered: events, record). */
+  flushSession(
+    sessionId: string,
+    reason: StorageContext["reason"] = "policy"
+  ): void {
+    const buffer = this.buffers.get(sessionId);
+    if (buffer === undefined) {
+      return;
+    }
+    this.buffers.delete(sessionId);
+    if (buffer.debounceTimer !== undefined) {
+      clearTimeout(buffer.debounceTimer);
+    }
+    const ops: StorageOp[] = [];
+    const first = buffer.events[0];
+    if (first !== undefined) {
+      ops.push({
+        op: "put",
+        key: storageKeys.sessionEvents(this.agentName, sessionId, first.index),
+        value: buffer.events as unknown as JsonValue,
+      });
+    }
+    if (buffer.record !== undefined) {
+      ops.push({
+        op: "put",
+        key: storageKeys.session(this.agentName, sessionId),
+        value: buffer.record as unknown as JsonValue,
+      });
+      ops.push(...this.continuationOps(buffer.record));
+    }
+    if (ops.length > 0) {
+      this.enqueue(ops, reason);
+    }
+  }
+
+  /**
+   * Index maintenance for `continuation/{channelId}/{token}` → sessionId:
+   * put the current token, delete the previous one when it changed.
+   */
+  private continuationOps(record: SessionRecord): StorageOp[] {
+    const previous = this.continuationIndex.get(record.sessionId);
+    const next = record.continuationKey;
+    if (previous === next) {
+      return [];
+    }
+    this.continuationIndex.set(record.sessionId, next);
+    const ops: StorageOp[] = [];
+    if (previous != null) {
+      ops.push({
+        op: "delete",
+        key: storageKeys.continuation(
+          this.agentName,
+          record.channelId,
+          previous
+        ),
+      });
+    }
+    if (next != null) {
+      ops.push({
+        op: "put",
+        key: storageKeys.continuation(this.agentName, record.channelId, next),
+        value: { sessionId: record.sessionId },
+      });
+    }
+    return ops;
+  }
+
+  // ==========================================================================
+  // Sessions (reads / restore)
+  // ==========================================================================
+
+  /** Whether the sink can serve the startup bulk restore (`list`). */
+  get canBulkRestore(): boolean {
+    return this.definition.list !== undefined;
+  }
+
+  /**
+   * Saved session records for this agent (startup restore). Errors are
+   * logged and read as "nothing saved" — a broken store must never block
+   * serve start.
+   */
+  async listSessions(): Promise<SessionRecord[]> {
+    const entries = await this.tryList(
+      storageKeys.sessionPrefix(this.agentName)
+    );
+    return entries
+      .map((entry) => entry.value as unknown as SessionRecord)
+      .filter((record) => record != null);
+  }
+
+  /**
+   * Saved event stream for one session: chunk values concatenated in key
+   * order (chunks are keyed by their first event's index, zero-padded, so
+   * ascending key order is append order). Errors read as an empty stream.
+   */
+  async listSessionEvents(sessionId: string): Promise<SessionEvent[]> {
+    const entries = await this.tryList(
+      storageKeys.sessionEventsPrefix(this.agentName, sessionId)
+    );
+    const events: SessionEvent[] = [];
+    for (const entry of entries) {
+      if (Array.isArray(entry.value)) {
+        events.push(...(entry.value as unknown as SessionEvent[]));
+      }
+    }
+    return events;
+  }
+
+  /**
+   * One saved session by channel continuation key (lazy restore). Unlike
+   * the startup reads, sink errors **propagate**: a failing store must
+   * fail the follow-up (which the caller can retry) rather than read as
+   * "unknown token" — that would mint a new session under the same
+   * continuation key and permanently shadow the real one on this host.
+   * A definitive miss (no index entry, or a write-only sink without
+   * `get`) resolves undefined.
+   */
+  async getSessionByContinuation(
+    channelId: string,
+    continuationKey: string
+  ): Promise<SessionRecord | undefined> {
+    const get = this.definition.get;
+    if (get === undefined) {
+      return undefined;
+    }
+    const ctx = this.context("restore");
+    const pointer = await get(
+      storageKeys.continuation(this.agentName, channelId, continuationKey),
+      ctx
+    );
+    if (pointer == null) {
+      return undefined;
+    }
+    const sessionId =
+      typeof pointer === "object" && !Array.isArray(pointer)
+        ? pointer.sessionId
+        : undefined;
+    if (typeof sessionId !== "string" || sessionId === "") {
+      throw new Error(
+        `storage returned a malformed continuation entry for channel ${channelId}`
+      );
+    }
+    const record = (await get(
+      storageKeys.session(this.agentName, sessionId),
+      ctx
+    )) as SessionRecord | null | undefined;
+    if (record == null) {
+      // Index entry without its record (e.g. the record write was shed
+      // under at-most-once). The session is unrecoverable; treat as a miss
+      // so the conversation can continue rather than failing forever.
+      this.logger(
+        `[agent-serve] storage continuation entry points at missing session ${sessionId}; treating as a miss`
+      );
+      return undefined;
+    }
+    // The sink must answer for the exact key it was asked about — a
+    // mismatched row is a broken store, not a miss; treating it as a miss
+    // would fork the continuation token onto a new session.
+    if (
+      record.sessionId !== sessionId ||
+      record.channelId !== channelId ||
+      record.continuationKey !== continuationKey
+    ) {
+      throw new Error(
+        `storage returned a session that does not match the requested continuation (channel ${channelId})`
+      );
+    }
+    return record;
+  }
+
+  // ==========================================================================
+  // Reminders
+  // ==========================================================================
+
+  /**
+   * Mirror one reminder record (create / fire / cancel / disarm). Local
+   * disk under `--state-root/reminders` remains the live source of truth;
+   * this is the durable copy for host replacement.
+   */
+  reminder(record: ReminderRecord): void {
+    if (this.closed) {
+      return;
+    }
+    this.enqueue(
+      [
+        {
+          op: "put",
+          key: storageKeys.reminder(this.agentName, record.id),
+          value: record as unknown as JsonValue,
+        },
+      ],
+      "policy"
+    );
+  }
+
+  /**
+   * Saved reminder records for this agent (startup hydrate into local
+   * disk). Errors read as empty — same posture as session bulk restore.
+   */
+  async listReminders(): Promise<ReminderRecord[]> {
+    const entries = await this.tryList(
+      storageKeys.reminderPrefix(this.agentName)
+    );
+    return entries.flatMap((entry) =>
+      isReminderRecord(entry.value) ? [entry.value] : []
+    );
+  }
+
+  // ==========================================================================
+  // Evals
+  // ==========================================================================
+
+  /**
+   * Adapter for {@link EvalRunPersistence} so the playground eval store
+   * can fall back to this sink when `evals.config.ts` sets no
+   * `persistRuns`.
+   */
+  asEvalRunPersistence(): EvalRunPersistence {
+    return {
+      load: async () => {
+        const entries = await this.tryList(
+          storageKeys.evalRunPrefix(this.agentName)
+        );
+        return entries
+          .map((entry) => entry.value as unknown as EvalRunSnapshot)
+          .filter((run) => run != null);
+      },
+      save: (run: EvalRunSnapshot) => {
+        this.enqueue(
+          [
+            {
+              op: "put",
+              key: storageKeys.evalRun(this.agentName, run.runId),
+              value: run as unknown as JsonValue,
+            },
+          ],
+          "policy"
+        );
+      },
+      delete: (runId: string) => {
+        this.enqueue(
+          [
+            {
+              op: "delete",
+              key: storageKeys.evalRun(this.agentName, runId),
+            },
+          ],
+          "policy"
+        );
+      },
+    };
+  }
+
+  // ==========================================================================
+  // A/Bs
+  // ==========================================================================
+
+  /**
+   * Adapter for {@link ABSamplePersistence} so the AB collector can fall
+   * back to this sink when `ab.config.ts` sets no `persistSamples`.
+   */
+  asABSamplePersistence(): ABSamplePersistence {
+    return {
+      save: (sample: ABMetricSample) => {
+        this.enqueue(
+          [
+            {
+              op: "put",
+              key: storageKeys.abSample(
+                this.agentName,
+                sample.sessionId,
+                sample.at
+              ),
+              value: sample as unknown as JsonValue,
+            },
+          ],
+          "policy"
+        );
+      },
+    };
+  }
+
+  /**
+   * Refresh the persisted aggregate A/B snapshot, throttled to once per
+   * {@link AB_SNAPSHOT_MIN_INTERVAL_MS} (the playground recomputes the
+   * fold on every `GET /v1/abs` poll).
+   */
+  abSnapshot(snapshot: ABSnapshot): void {
+    const now = Date.now();
+    if (
+      this.closed ||
+      now - this.lastAbSnapshotAt < AB_SNAPSHOT_MIN_INTERVAL_MS
+    ) {
+      return;
+    }
+    this.lastAbSnapshotAt = now;
+    this.enqueue(
+      [
+        {
+          op: "put",
+          key: storageKeys.abSnapshot(this.agentName),
+          value: snapshot as unknown as JsonValue,
+        },
+      ],
+      "policy"
+    );
+  }
+
+  /** Latest persisted aggregate A/B snapshot; errors read as absent. */
+  async getLatestAbSnapshot(): Promise<ABSnapshot | undefined> {
+    const get = this.definition.get;
+    if (get === undefined) {
+      return undefined;
+    }
+    try {
+      const value = await get(
+        storageKeys.abSnapshot(this.agentName),
+        this.context("restore")
+      );
+      return value == null ? undefined : (value as unknown as ABSnapshot);
+    } catch (error) {
+      this.logger(
+        `[agent-serve] storage get(ab-snapshot) failed: ${describeError(error)}`
+      );
+      return undefined;
+    }
+  }
+
+  // ==========================================================================
+  // Delivery
+  // ==========================================================================
+
+  /**
+   * Flush all buffers and drain the queue, giving the sink at most
+   * {@link CLOSE_DRAIN_TIMEOUT_MS} — a hung sink must not stall shutdown
+   * past the pod grace period. Called from engine close.
+   */
+  async close(): Promise<void> {
+    if (!this.closed) {
+      for (const sessionId of [...this.buffers.keys()]) {
+        this.flushSession(sessionId, "shutdown");
+      }
+      this.closed = true;
+    }
+    let timer: ReturnType<typeof setTimeout> | undefined;
+    const drained = await Promise.race([
+      this.queue.then(() => true),
+      new Promise<boolean>((resolve) => {
+        timer = setTimeout(() => resolve(false), CLOSE_DRAIN_TIMEOUT_MS);
+        timer.unref?.();
+      }),
+    ]);
+    clearTimeout(timer);
+    if (!drained) {
+      this.logger(
+        `[agent-serve] storage shutdown drain timed out after ${CLOSE_DRAIN_TIMEOUT_MS}ms with ${this.pending} write(s) undelivered`
+      );
+    }
+  }
+
+  /** Pending deliveries (exposed for tests and drain instrumentation). */
+  whenIdle(): Promise<void> {
+    return this.queue;
+  }
+
+  private buffer(sessionId: string): SessionBuffer {
+    let buffer = this.buffers.get(sessionId);
+    if (buffer === undefined) {
+      buffer = { events: [] };
+      this.buffers.set(sessionId, buffer);
+    }
+    return buffer;
+  }
+
+  private armDebounce(
+    sessionId: string,
+    buffer: SessionBuffer,
+    debounceMs: number
+  ): void {
+    if (buffer.debounceTimer !== undefined) {
+      clearTimeout(buffer.debounceTimer);
+    }
+    buffer.debounceTimer = setTimeout(() => {
+      buffer.debounceTimer = undefined;
+      this.flushSession(sessionId);
+    }, debounceMs);
+    // Never keep the process alive just for a pending storage flush.
+    buffer.debounceTimer.unref?.();
+  }
+
+  /** `list` wrapper for the startup reads: missing hook or throw ⇒ empty. */
+  private async tryList(
+    prefix: string
+  ): Promise<Array<{ key: string; value: JsonValue }>> {
+    const list = this.definition.list;
+    if (list === undefined) {
+      return [];
+    }
+    try {
+      const entries = await list(prefix, this.context("restore"));
+      return Array.isArray(entries) ? entries : [];
+    } catch (error) {
+      this.logger(
+        `[agent-serve] storage list(${prefix}) failed: ${describeError(error)}`
+      );
+      return [];
+    }
+  }
+
+  private enqueue(ops: StorageOp[], reason: StorageContext["reason"]): void {
+    if (this.closed && reason !== "shutdown") {
+      return;
+    }
+    if (this.pending + ops.length > MAX_PENDING_OPS) {
+      this.dropped += ops.length;
+      if (this.dropped <= ops.length || this.dropped % 1_000 < ops.length) {
+        this.logger(
+          `[agent-serve] storage queue full (${MAX_PENDING_OPS} pending writes); dropped ${this.dropped} write(s) so far — the sink is too slow for the traffic`
+        );
+      }
+      return;
+    }
+    this.pending += ops.length;
+    const ctx = this.context(reason);
+    this.queue = this.queue.then(async () => {
+      for (const op of ops) {
+        try {
+          if (op.op === "put") {
+            await this.definition.put(op.key, op.value, ctx);
+          } else {
+            await this.definition.delete?.(op.key, ctx);
+          }
+        } catch (error) {
+          this.logger(
+            `[agent-serve] storage ${op.op}(${op.key}) failed: ${describeError(error)}`
+          );
+        } finally {
+          this.pending -= 1;
+        }
+      }
+    });
+  }
+
+  private context(reason: StorageContext["reason"]): StorageContext {
+    return { agentName: this.agentName, projectRoot: this.projectRoot, reason };
+  }
+}