@cursor/july 0.1.1 → 0.1.2

package/src/internal/session-engine.ts

@@ -7,8 +7,8 @@
  * event handlers and hooks.
  */
 
-import { mkdir, rm } from "node:fs/promises";
-import { join } from "node:path";
+import { mkdir, rm, writeFile } from "node:fs/promises";
+import { dirname, join } from "node:path";
 import type { SDKCustomTool } from "@cursor/sdk";
 import {
   type ABDefinition,
@@ -68,6 +68,7 @@ import { McpHost } from "./mcp-host.js";
 import { assertCloudCanHonorAccountServersFilters } from "./resolved-connections.js";
 import type { AgentRunner } from "./sdk-runner.js";
 import { SessionStore } from "./session-store.js";
+import { StorageCoordinator } from "./storage-coordinator.js";
 import { normalizeToolResult, toolCallErrorMessage } from "./tool-result.js";
 import {
   buildAgentToolsCatalog,
@@ -157,6 +158,7 @@ export interface SessionEngineOptions {
   stateRoot: string;
   runner: AgentRunner;
   logger?: (line: string) => void;
+  platforms?: HostPlatforms;
 }
 
 /**
@@ -192,6 +194,21 @@ export class SessionEngine {
   readonly project: AgentProject;
   readonly stateRoot: string;
   readonly sessions: SessionStore;
+  /**
+   * Runtime for `agent/storage.ts` (`defineStorage`), when
+   * authored. All durable state changes funnel through it (session
+   * records, event chunks, A/B samples/snapshots); the eval run store
+   * attaches via {@link StorageCoordinator.asEvalRunPersistence} at
+   * serve start.
+   */
+  readonly storage: StorageCoordinator | undefined;
+  /** In-flight lazy restores, deduped per channel + continuation key. */
+  private readonly restoreInFlight = new Map<
+    string,
+    Promise<SessionRecord | undefined>
+  >();
+  /** Once-per-process cache for the A/B snapshot backfill (see abSnapshot). */
+  private restoredAbSnapshot: Promise<ABSnapshot | undefined> | undefined;
 
   private readonly runner: AgentRunner;
   private readonly logger: (line: string) => void;
@@ -229,15 +246,26 @@ export class SessionEngine {
           `[agent-serve] failed to persist events for ${sessionId}: ${describeError(error)}`
         )
     );
+    this.storage =
+      options.project.storage === undefined
+        ? undefined
+        : new StorageCoordinator({
+            definition: options.project.storage,
+            agentName: options.project.name,
+            projectRoot: options.project.rootDir,
+            logger: this.logger,
+          });
+    // Domain-specific hooks win; the defineStorage sink is the fallback.
+    const persistSamples =
+      options.project.abConfig?.persistSamples ??
+      this.storage?.asABSamplePersistence();
     this.abCollector = new ABCollector(
       options.project.abs,
       this.logger,
       async (sessionId) => (await this.logs.get(sessionId)).snapshot(0),
       {
         projectRoot: options.project.rootDir,
-        ...(options.project.abConfig?.persistSamples === undefined
-          ? {}
-          : { persistSamples: options.project.abConfig.persistSamples }),
+        persistSamples,
       }
     );
     // Peer and Cursor-account transports are symbolic until the serve host
@@ -248,7 +276,7 @@ export class SessionEngine {
         (connection) => !isSymbolicConnectionTransport(connection.transport)
       )
     );
-    this.platforms = createHostPlatforms();
+    this.platforms = options.platforms ?? createHostPlatforms();
     for (const channel of options.project.channels) {
       this.channelsById.set(channel.id, channel.definition);
     }
@@ -305,13 +333,21 @@ export class SessionEngine {
     options: EngineSendOptions = {}
   ): Promise<ChannelSession> {
     const auth = options.auth ?? null;
-    const existing =
+    let existing =
       options.continuationToken === undefined
         ? undefined
         : await this.sessions.findByContinuation(
             channelId,
             options.continuationToken
           );
+    if (existing === undefined && options.continuationToken !== undefined) {
+      // Local miss: the session may live only in the storage sink (it
+      // fell outside the startup restore window, or restore is lazy-only).
+      existing = await this.restoreSessionByContinuation(
+        channelId,
+        options.continuationToken
+      );
+    }
 
     if (existing !== undefined) {
       if (!samePrincipal(existing.auth, auth)) {
@@ -361,6 +397,7 @@ export class SessionEngine {
       updatedAt: now,
     };
     await this.sessions.save(record);
+    this.storage?.sessionRecord(record);
     await this.appendEvent(record.sessionId, {
       type: "session.started",
       data: { channelId },
@@ -1337,6 +1374,9 @@ export class SessionEngine {
   }
 
   private queueDispatch(sessionId: string, event: SessionEvent): void {
+    // Single funnel for both append paths (appendEvent + the held-log emit),
+    // so the storage sink sees every durable event exactly once.
+    this.storage?.event(event);
     const chain = this.dispatchChains.get(sessionId) ?? Promise.resolve();
     const next = chain
       .then(() => this.dispatchEvent(sessionId, event))
@@ -1344,6 +1384,12 @@ export class SessionEngine {
         this.logger(
           `[agent-serve] event dispatch failed for ${sessionId}: ${describeError(error)}`
         );
+      })
+      .then(() => {
+        // Turn-end flush happens here — after the boundary event's handlers
+        // and their record updates settled — so the persisted record carries
+        // the turn's final channel state and continuation token.
+        this.storage?.eventDispatched(event);
       });
     this.dispatchChains.set(sessionId, next);
   }
@@ -1508,6 +1554,25 @@ export class SessionEngine {
       abConfig?.maxPlaygroundSessions
     );
     const foldedSessions = sessions.slice(0, maxPlaygroundSessions);
+    // Backfill for a replacement host with nothing local to fold (e.g.
+    // lazy-only restore before any session resumed): serve the persisted
+    // aggregate instead of an empty surface. Gated on the host having no
+    // sessions AT ALL — the persisted snapshot spans every principal, so a
+    // caller who merely owns no sessions must not see it. Loaded once per
+    // process (playground polls this route ~every 4s); skip the persist
+    // side effect below — re-saving a restored snapshot would echo stale
+    // data back into the store.
+    if (
+      foldedSessions.length === 0 &&
+      this.storage !== undefined &&
+      (await this.sessions.list()).length === 0
+    ) {
+      this.restoredAbSnapshot ??= this.storage.getLatestAbSnapshot();
+      const restoredSnapshot = await this.restoredAbSnapshot;
+      if (restoredSnapshot !== undefined) {
+        return restoredSnapshot;
+      }
+    }
     const snapshot = await buildABSnapshot({
       experiments: this.project.abs,
       agentName: this.project.name,
@@ -1535,6 +1600,9 @@ export class SessionEngine {
           `[agent-serve] ab persistSnapshots.save threw: ${describeError(error)}`
         );
       }
+    } else {
+      // defineStorage fallback (throttled inside the coordinator).
+      this.storage?.abSnapshot(snapshot);
     }
     return snapshot;
   }
@@ -1700,6 +1768,7 @@ export class SessionEngine {
         updatedAt: new Date().toISOString(),
       };
       await this.sessions.save(updated);
+      this.storage?.sessionRecord(updated);
       return updated;
     });
     this.updateChains.set(
@@ -1713,6 +1782,203 @@ export class SessionEngine {
   // Lifecycle
   // ==========================================================================
 
+  /**
+   * Restore sessions from the authored storage sink. Called once at
+   * serve start, before channels accept traffic: any persisted session
+   * missing from the local `--state-root` is materialized there — events
+   * as `events.ndjson`, record as `session.json` — after which every
+   * existing read path (continuation lookup, event replay/streams, A/B
+   * fold hydration) works unchanged. Local state always wins over the
+   * sink, so restarts on the same host are no-ops.
+   *
+   * Failures are logged and skipped; a broken store (including rows that
+   * violate the SessionRecord shape) never blocks serve start.
+   */
+  async restoreStoredSessions(): Promise<void> {
+    const storage = this.storage;
+    if (storage === undefined || !storage.canBulkRestore) {
+      return;
+    }
+    const policy = storage.policy.restore;
+    if (policy === "off") {
+      return;
+    }
+    try {
+      await this.runBulkRestore(storage, policy);
+    } catch (error) {
+      this.logger(
+        `[agent-serve] session restore failed: ${describeError(error)}`
+      );
+    }
+  }
+
+  private async runBulkRestore(
+    storage: StorageCoordinator,
+    policy: { maxSessions: number; maxAgeMs: number; maxTotalBytes: number }
+  ): Promise<void> {
+    const loaded = await storage.listSessions();
+    // Guardrails on whatever the sink returned: drop invalid rows and
+    // sessions older than the age window, then restore newest-first under
+    // the count cap. Anything filtered out here stays reachable through
+    // the lazy continuation-lookup path.
+    const cutoff = Date.now() - policy.maxAgeMs;
+    const candidates = loaded
+      .map((record) => ({
+        record,
+        // Numeric timestamp for both the age filter and the sort; a
+        // non-string updatedAt (e.g. a raw DB Date from an untyped sink)
+        // parses to NaN and is dropped instead of crashing serve start.
+        updatedAt:
+          typeof record?.updatedAt === "string"
+            ? Date.parse(record.updatedAt)
+            : Number.NaN,
+      }))
+      .filter(
+        ({ record, updatedAt }) =>
+          typeof record?.sessionId === "string" &&
+          record.sessionId !== "" &&
+          updatedAt >= cutoff
+      )
+      .sort((a, b) => b.updatedAt - a.updatedAt)
+      .slice(0, policy.maxSessions)
+      .map(({ record }) => record);
+
+    let restored = 0;
+    let budget = policy.maxTotalBytes;
+    // Sink round trips dominate restore latency, so event streams are
+    // fetched with bounded parallelism; writes stay sequential, newest
+    // first, so the byte budget is enforced deterministically.
+    outer: for (
+      let i = 0;
+      i < candidates.length;
+      i += RESTORE_FETCH_CONCURRENCY
+    ) {
+      const chunk = candidates.slice(i, i + RESTORE_FETCH_CONCURRENCY);
+      const fetched = await Promise.all(
+        chunk.map(async (record) => {
+          try {
+            // Local state always wins over the sink.
+            if ((await this.sessions.get(record.sessionId)) !== undefined) {
+              return undefined;
+            }
+            const events = await storage.listSessionEvents(record.sessionId);
+            return { record, events };
+          } catch (error) {
+            this.logger(
+              `[agent-serve] failed to restore session ${record.sessionId}: ${describeError(error)}`
+            );
+            return undefined;
+          }
+        })
+      );
+      for (const item of fetched) {
+        if (item === undefined) {
+          continue;
+        }
+        const ndjson = eventsNdjson(item.events);
+        // Budget is checked BEFORE anything is written: one oversized
+        // session cannot blow past maxTotalBytes on disk. Skipped sessions
+        // stay reachable through the lazy path.
+        const bytes =
+          Buffer.byteLength(JSON.stringify(item.record), "utf8") +
+          (ndjson === undefined ? 0 : Buffer.byteLength(ndjson, "utf8"));
+        if (bytes > budget) {
+          this.logger(
+            `[agent-serve] session restore stopped at the ${policy.maxTotalBytes}-byte budget after ${restored} session(s); remaining sessions restore lazily on resume`
+          );
+          break outer;
+        }
+        try {
+          await this.materializeSession(item.record, ndjson);
+          budget -= bytes;
+          restored += 1;
+        } catch (error) {
+          this.logger(
+            `[agent-serve] failed to restore session ${item.record.sessionId}: ${describeError(error)}`
+          );
+        }
+      }
+    }
+    if (restored > 0) {
+      this.logger(`[agent-serve] restored ${restored} session(s) from storage`);
+    }
+  }
+
+  /**
+   * Lazy restore: a continuation lookup missed locally, so ask the sink
+   * for that one session and materialize it before resuming. Disk grows
+   * with sessions that are actually resumed, not with history.
+   *
+   * Concurrent misses for the same key (webhook redeliveries, retries)
+   * share one in-flight restore — a second materialization racing the
+   * first resumed turn could truncate events the turn already appended.
+   * Sink errors propagate and fail the send (see
+   * {@link StorageCoordinator.getSessionByContinuation}).
+   */
+  private restoreSessionByContinuation(
+    channelId: string,
+    continuationKey: string
+  ): Promise<SessionRecord | undefined> {
+    const storage = this.storage;
+    if (storage === undefined) {
+      return Promise.resolve(undefined);
+    }
+    const key = `${channelId}\u0000${continuationKey}`;
+    let pending = this.restoreInFlight.get(key);
+    if (pending === undefined) {
+      pending = this.loadAndMaterializeByContinuation(
+        storage,
+        channelId,
+        continuationKey
+      ).finally(() => this.restoreInFlight.delete(key));
+      this.restoreInFlight.set(key, pending);
+    }
+    return pending;
+  }
+
+  private async loadAndMaterializeByContinuation(
+    storage: StorageCoordinator,
+    channelId: string,
+    continuationKey: string
+  ): Promise<SessionRecord | undefined> {
+    const record = await storage.getSessionByContinuation(
+      channelId,
+      continuationKey
+    );
+    if (record === undefined) {
+      // Definitive miss: the caller starts a fresh session under this token.
+      return undefined;
+    }
+    const local = await this.sessions.get(record.sessionId);
+    if (local !== undefined) {
+      return local;
+    }
+    const events = await storage.listSessionEvents(record.sessionId);
+    await this.materializeSession(record, eventsNdjson(events));
+    this.logger(
+      `[agent-serve] lazily restored session ${record.sessionId} from storage`
+    );
+    return record;
+  }
+
+  /**
+   * Write one restored session to the local state root. Events land on
+   * disk before the record: a session becomes visible (and addressable via
+   * its continuation key) only with its full stream in place, and the
+   * lazily opened event log hydrates from this file.
+   */
+  private async materializeSession(
+    record: SessionRecord,
+    eventsNdjson: string | undefined
+  ): Promise<void> {
+    if (eventsNdjson !== undefined) {
+      const filePath = this.sessions.eventFilePath(record.sessionId);
+      await mkdir(dirname(filePath), { recursive: true });
+      await writeFile(filePath, eventsNdjson, "utf8");
+    }
+    await this.sessions.save(record);
+  }
+
   trackBackground(promise: Promise<unknown>): void {
     this.backgroundWork.add(promise);
     void promise
@@ -1730,11 +1996,23 @@ export class SessionEngine {
     await Promise.allSettled([...this.dispatchChains.values()]);
     await Promise.allSettled([...this.updateChains.values()]);
     await this.logs.flushAll();
+    await this.storage?.close();
     await this.mcpHost.close();
     await this.runner.dispose?.();
   }
 }
 
+/** Parallel event-stream loads during the startup bulk restore. */
+const RESTORE_FETCH_CONCURRENCY = 8;
+
+/** NDJSON payload for a restored event stream; undefined when empty. */
+function eventsNdjson(events: SessionEvent[]): string | undefined {
+  if (events.length === 0) {
+    return undefined;
+  }
+  return `${events.map((event) => JSON.stringify(event)).join("\n")}\n`;
+}
+
 function subagentPrompt(subagent: ResolvedAgent): string {
   if (subagent.instructions !== undefined && subagent.instructions !== "") {
     return subagent.instructions;