experimental-a2 0.7.0 → 0.8.0

package/src/server.ts

@@ -38,6 +38,7 @@ import type {
   Event,
   EventCause,
   StoreClaimAvailableResult,
+  StoreSnapshotWrite,
   StoreStateRead,
   PresenceRow,
   ReturnedEvent,
@@ -134,6 +135,10 @@ export type SessionSchedule<D extends EventDefs> = (
   ...events: AppendInput<D>[]
 ) => Promise<void>
 
+export type StateOptions = {
+  through?: number | 'latest'
+}
+
 /**
  * The presence members of a session — intersected in via
  * `WithPresence`, so they exist exactly when the contract declares
@@ -177,7 +182,10 @@ export type Session<
   append: Append
   schedule: SessionSchedule<D>
   history(options?: { gte?: number; lte?: number }): Promise<ContractEvent<D>[]>
-  state<S>(reducer: Reducer<D, S>): Promise<{ state: S; index: number }>
+  state<S>(
+    reducer: Reducer<D, S>,
+    options?: StateOptions,
+  ): Promise<{ state: S; index: number }>
   /**
    * A live feed of this session's events, starting after `startAfter`
    * (exclusive). Server-side only. `server.fetch` exposes it over SSE.
@@ -965,26 +973,85 @@ export function createServer<
     store: A2Store,
     sessionId: string,
     reducerName: string,
+    throughIndex?: number,
+    snapshotThroughIndex?: number,
   ): Promise<StoreStateRead> => {
     try {
-      return await store.readState(nsId(sessionId), reducerName)
+      return await store.readState(
+        nsId(sessionId),
+        reducerName,
+        throughIndex === undefined && snapshotThroughIndex === undefined
+          ? undefined
+          : {
+              ...(throughIndex === undefined ? {} : { throughIndex }),
+              ...(snapshotThroughIndex === undefined
+                ? {}
+                : { snapshotThroughIndex }),
+            },
+      )
     } catch {
       try {
-        return { snapshot: null, events: await store.read(nsId(sessionId)) }
+        return {
+          headIndex: null,
+          snapshot: null,
+          events: await store.read(
+            nsId(sessionId),
+            throughIndex === undefined ? undefined : { throughIndex },
+          ),
+        }
       } catch (err) {
         throw asStoreUnavailable(err)
       }
     }
   }
 
+  type PendingStateRead = {
+    sessionId: string
+    throughIndex?: number
+    pinEventIndex?: number
+    span: A2SpanHandle
+    resolve: (result: { state: unknown; index: number }) => void
+    reject: (reason: unknown) => void
+  }
+
+  type StateReadPlan = {
+    sessionId: string
+    entries: PendingStateRead[]
+    throughIndex?: number
+    snapshotThroughIndex?: number
+  }
+
+  const planStateReads = (batch: PendingStateRead[]): StateReadPlan[] => {
+    const bySession = new Map<string, PendingStateRead[]>()
+    for (const entry of batch) {
+      const entries = bySession.get(entry.sessionId)
+      if (entries) entries.push(entry)
+      else bySession.set(entry.sessionId, [entry])
+    }
+    const plans: StateReadPlan[] = []
+    for (const [sessionId, entries] of bySession) {
+      const boundaries = entries.flatMap((entry) =>
+        entry.throughIndex === undefined ? [] : [entry.throughIndex],
+      )
+      const plan: StateReadPlan = { sessionId, entries }
+      if (boundaries.length === entries.length) {
+        plan.throughIndex = Math.max(...boundaries)
+      }
+      if (boundaries.length > 0) {
+        plan.snapshotThroughIndex = Math.min(...boundaries)
+      }
+      plans.push(plan)
+    }
+    return plans
+  }
+
   // The cache is untrusted: schema rejection refolds from the raw store.
-  const foldStateRead = async <S>(
+  const foldStatePlan = async (
     store: A2Store,
-    sessionId: string,
-    reducer: Reducer<D, S>,
+    reducer: Reducer<D, unknown>,
+    plan: StateReadPlan,
     stateRead: StoreStateRead,
-    span: A2SpanHandle,
-  ): Promise<{ state: S; index: number; folded: number }> => {
+  ): Promise<void> => {
     let state = cloneInitial(reducer.initialState)
     let index = 0
     let snapshotOutcome = 'miss'
@@ -1000,44 +1067,101 @@ export function createServer<
         if (result.issues) {
           snapshotOutcome = 'rejected'
         } else {
-          state = result.value as S
+          state = result.value
           index = snap.index
           snapshotOutcome = 'hit'
         }
       } else {
-        state = snap.state as S
+        state = snap.state
         index = snap.index
         snapshotOutcome = 'hit'
       }
     }
-    span.setAttribute('a2.state.snapshot', snapshotOutcome)
-
     if (snapshotOutcome === 'rejected') {
       try {
-        rows = await store.read(nsId(sessionId))
+        rows = await store.read(
+          nsId(plan.sessionId),
+          plan.throughIndex === undefined
+            ? undefined
+            : { throughIndex: plan.throughIndex },
+        )
       } catch (err) {
         throw asStoreUnavailable(err)
       }
     }
-    for (const row of rows) {
-      state = reducer.fold(state, toPublic(row) as never)
-      index = row.index
+    const ordered = plan.entries.toSorted(
+      (a, b) =>
+        (a.throughIndex ?? Number.POSITIVE_INFINITY) -
+        (b.throughIndex ?? Number.POSITIVE_INFINITY),
+    )
+    const writes = new Map<
+      number,
+      { state: unknown; pinEventIndexes: Set<number> }
+    >()
+    let rowPosition = 0
+    for (const entry of ordered) {
+      const target = entry.throughIndex ?? Number.POSITIVE_INFINITY
+      while (rowPosition < rows.length && rows[rowPosition]!.index <= target) {
+        const row = rows[rowPosition]!
+        state = reducer.fold(state, toPublic(row) as never)
+        index = row.index
+        rowPosition += 1
+      }
+      const resultState = cloneInitial(state)
+      const folded = rowPosition
+      entry.span.setAttribute('a2.state.snapshot', snapshotOutcome)
+      entry.span.setAttribute('a2.state.folded', folded)
+      entry.span.setAttribute('a2.state.index', index)
+      if (index > 0 && (folded > 0 || entry.pinEventIndex !== undefined)) {
+        const advancesHead =
+          stateRead.headIndex === null || index > stateRead.headIndex
+        const repairsHead =
+          snapshotOutcome === 'rejected' && index === stateRead.headIndex
+        if (
+          !advancesHead &&
+          !repairsHead &&
+          entry.pinEventIndex === undefined
+        ) {
+          entry.resolve({ state: resultState, index })
+          continue
+        }
+        let write = writes.get(index)
+        if (!write) {
+          write = {
+            state: cloneInitial(resultState),
+            pinEventIndexes: new Set(),
+          }
+          writes.set(index, write)
+        }
+        if (entry.pinEventIndex !== undefined) {
+          write.pinEventIndexes.add(entry.pinEventIndex)
+        }
+      }
+      entry.resolve({ state: resultState, index })
+    }
+    if (writes.size > 0) {
+      const snapshots: StoreSnapshotWrite[] = []
+      for (const [snapshotIndex, write] of writes) {
+        const snapshot: StoreSnapshotWrite = {
+          index: snapshotIndex,
+          state: write.state,
+        }
+        if (write.pinEventIndexes.size > 0) {
+          snapshot.pinEventIndexes = [...write.pinEventIndexes]
+        }
+        snapshots.push(snapshot)
+      }
+      const persistence = Promise.resolve()
+        .then(() =>
+          store.putSnapshots(nsId(plan.sessionId), reducer.name, snapshots),
+        )
+        .catch(() => {})
+      track(persistence)
+      platformWaitUntil(persistence)
     }
-    span.setAttribute('a2.state.folded', rows.length)
-    span.setAttribute('a2.state.index', index)
-    return { state, index, folded: rows.length }
   }
 
   // ── same-tick state-read coalescing ──────────────────────────────
-  // Concurrent state() calls for one reducer name coalesce into a
-  // single store.readStates() round trip when the adapter has one.
-  // Unlike DataLoader there is deliberately NO result cache: a batch
-  // exists only between enqueue and flush, and nothing is shared after
-  // distribution — every state() call still observes a fresh frontier.
-  type PendingStateRead = {
-    sessionId: string
-    resolve: (read: StoreStateRead | Promise<StoreStateRead>) => void
-  }
   const pendingStateReads = new Map<string, PendingStateRead[]>()
 
   // The flush microtask voids this promise, so the function must be
@@ -1045,40 +1169,62 @@ export function createServer<
   // synchronously throwing adapter or a read structuredClone rejects.
   const flushStateReads = async (
     store: A2Store,
-    reducerName: string,
+    reducer: Reducer<D, unknown>,
     batch: PendingStateRead[],
   ): Promise<void> => {
+    const plans = planStateReads(batch)
     let reads: StoreStateRead[] | null
     try {
       reads = store.readStates
         ? await store.readStates(
-            batch.map((entry) => nsId(entry.sessionId)),
-            reducerName,
+            plans.map((plan) => ({
+              sessionId: nsId(plan.sessionId),
+              ...(plan.throughIndex === undefined
+                ? {}
+                : { throughIndex: plan.throughIndex }),
+              ...(plan.snapshotThroughIndex === undefined
+                ? {}
+                : { snapshotThroughIndex: plan.snapshotThroughIndex }),
+            })),
+            reducer.name,
+          )
+        : await Promise.all(
+            plans.map((plan) =>
+              readCachedState(
+                store,
+                plan.sessionId,
+                reducer.name,
+                plan.throughIndex,
+                plan.snapshotThroughIndex,
+              ),
+            ),
           )
-        : null
     } catch {
       reads = null
     }
-    // A failed or misaligned batch is an ill-behaved cache path, never a
-    // batch-shaped error: each caller falls back to its own per-call read
-    // (promise adoption carries a real storage failure to that caller).
-    if (!reads || reads.length !== batch.length) {
-      for (const entry of batch) {
-        entry.resolve(readCachedState(store, entry.sessionId, reducerName))
+    if (!reads || reads.length !== plans.length) {
+      try {
+        reads = await Promise.all(
+          plans.map((plan) =>
+            readCachedState(
+              store,
+              plan.sessionId,
+              reducer.name,
+              plan.throughIndex,
+              plan.snapshotThroughIndex,
+            ),
+          ),
+        )
+      } catch (err) {
+        for (const entry of batch) entry.reject(err)
+        return
       }
-      return
     }
-    // Duplicate ids share one store read; their folds must not share
-    // payload objects, exactly as two separate state() calls would not.
-    const seen = new Set<string>()
-    for (const [position, entry] of batch.entries()) {
-      const shared = seen.has(entry.sessionId)
-      seen.add(entry.sessionId)
+    for (const [position, plan] of plans.entries()) {
       try {
-        const read = reads[position]!
-        entry.resolve(shared ? structuredClone(read) : read)
-      } catch {
-        entry.resolve(readCachedState(store, entry.sessionId, reducerName))
+        await foldStatePlan(store, reducer, plan, reads[position]!)
+      } catch (err) {
+        for (const entry of plan.entries) entry.reject(err)
       }
     }
   }
@@ -1086,9 +1232,13 @@ export function createServer<
   const enqueueStateRead = (
     store: A2Store,
     sessionId: string,
-    reducerName: string,
-  ): Promise<StoreStateRead> =>
-    new Promise((resolve) => {
+    reducer: Reducer<D, unknown>,
+    throughIndex?: number,
+    pinEventIndex?: number,
+    span?: A2SpanHandle,
+  ): Promise<{ state: unknown; index: number }> =>
+    new Promise((resolve, reject) => {
+      const reducerName = reducer.name
       let batch = pendingStateReads.get(reducerName)
       if (!batch) {
         const opened: PendingStateRead[] = []
@@ -1104,16 +1254,25 @@ export function createServer<
         // the last. A deeper-staggered caller just opens the next batch.
         queueMicrotask(() => {
           pendingStateReads.delete(reducerName)
-          void flushStateReads(store, reducerName, opened)
+          void flushStateReads(store, reducer, opened)
         })
         batch = opened
       }
-      batch.push({ sessionId, resolve })
+      batch.push({
+        sessionId,
+        ...(throughIndex === undefined ? {} : { throughIndex }),
+        ...(pinEventIndex === undefined ? {} : { pinEventIndex }),
+        span: span!,
+        resolve,
+        reject,
+      })
     })
 
   const readState = async <S>(
     sessionId: string,
     reducer: Reducer<D, S>,
+    throughIndex?: number,
+    pinEventIndex?: number,
   ): Promise<{ state: S; index: number }> => {
     const store = await resolveStore()
     return telemetry.span(
@@ -1123,36 +1282,15 @@ export function createServer<
         'a2.session_id': sessionId,
         'a2.state.reducer': reducer.name,
       },
-      async (span) => {
-        const stateRead = store.readStates
-          ? await enqueueStateRead(store, sessionId, reducer.name)
-          : await readCachedState(store, sessionId, reducer.name)
-        const { state, index, folded } = await foldStateRead(
+      (span) =>
+        enqueueStateRead(
           store,
           sessionId,
-          reducer,
-          stateRead,
+          reducer as Reducer<D, unknown>,
+          throughIndex,
+          pinEventIndex,
           span,
-        )
-        // Write-back is a disposable cache. Copy before returning so caller
-        // mutations cannot race the background persistence.
-        if (folded > 0) {
-          const snapshotState = cloneInitial(state)
-          const write = Promise.resolve()
-            .then(() =>
-              store.putSnapshot(
-                nsId(sessionId),
-                reducer.name,
-                index,
-                snapshotState,
-              ),
-            )
-            .catch(() => {})
-          track(write)
-          platformWaitUntil(write)
-        }
-        return { state, index }
-      },
+        ) as Promise<{ state: S; index: number }>,
     )
   }
 
@@ -1176,11 +1314,11 @@ export function createServer<
       const payloads = events.map((event) =>
         cloneScheduledPayload(event.payload),
       )
-      const snapshottedEvents = events.map((event, index) => ({
-        type: event.type,
-        payload: payloads[index],
-        ...(event.id === undefined ? {} : { id: event.id }),
-      })) as AppendInput<D>[]
+      const snapshottedEvents = events.map((event, index) =>
+        event.id === undefined
+          ? { type: event.type, payload: payloads[index] }
+          : { type: event.type, payload: payloads[index], id: event.id },
+      ) as AppendInput<D>[]
       const validated = validateEvents(
         sessionId,
         structuredClone(snapshottedEvents),
@@ -1591,7 +1729,18 @@ export function createServer<
       schedule: makeSchedule(id, trigger),
       history: async (bounds?: { gte?: number; lte?: number }) =>
         (await readHistory(id, bounds)).map(toPublic),
-      state: <S>(reducer: Reducer<D, S>) => readState(id, reducer),
+      state: <S>(reducer: Reducer<D, S>, stateOptions?: StateOptions) => {
+        const through = stateOptions?.through ?? trigger?.index
+        if (through !== undefined && through !== 'latest') {
+          assertStateIndex(through)
+        }
+        return readState(
+          id,
+          reducer,
+          through === 'latest' ? undefined : through,
+          trigger && through !== 'latest' ? trigger.index : undefined,
+        )
+      },
       stream: (opts?: { startAfter?: number; presence?: boolean }) => {
         const startAfter = opts?.startAfter ?? 0
         assertStreamIndex(startAfter)
@@ -2359,6 +2508,12 @@ function assertStreamIndex(value: number): void {
   }
 }
 
+function assertStateIndex(value: number): void {
+  if (!Number.isSafeInteger(value) || value < 0) {
+    throw new TypeError('state.through must be a non-negative safe integer')
+  }
+}
+
 function assertAppendName(name: string): void {
   if (typeof name !== 'string' || name.length === 0) {
     throw new TypeError('handler append name must be a non-empty string')

package/src/store-memory.ts

@@ -101,6 +101,12 @@ const toEvent = (row: Row): Event => ({
   createdAt: new Date(row.createdAt),
 })
 
+const snapshotKey = (sessionId: string, reducerName: string): string =>
+  `${sessionId}\u0000${reducerName}`
+
+const eventPinKey = (sessionId: string, index: number): string =>
+  `${sessionId}\u0000${index}`
+
 export function memory(options: MemoryStoreOptions = {}): A2Store {
   const clock = options.clock ?? SYSTEM_CLOCK
   const generateId = options.ids ?? RANDOM_IDS
@@ -114,6 +120,12 @@ export function memory(options: MemoryStoreOptions = {}): A2Store {
     string,
     { index: number; state: unknown; updatedAt: Date }
   >()
+  const historicalSnapshots = new Map<
+    string,
+    { index: number; state: unknown; updatedAt: Date }
+  >()
+  const snapshotPins = new Map<string, Set<string>>()
+  const snapshotPinCounts = new Map<string, number>()
   const streamSubscribers = new Map<string, Set<(row: Row) => void>>()
   /** ns → participant → field → latest surviving write. */
   const presenceRows = new Map<
@@ -142,6 +154,28 @@ export function memory(options: MemoryStoreOptions = {}): A2Store {
     return rows
   }
 
+  const checkpointKey = (
+    sessionId: string,
+    reducerName: string,
+    index: number,
+  ): string => `${snapshotKey(sessionId, reducerName)}\u0000${index}`
+
+  const releaseSnapshotPins = (sessionId: string, eventIndex: number): void => {
+    const pinKey = eventPinKey(sessionId, eventIndex)
+    const checkpoints = snapshotPins.get(pinKey)
+    if (!checkpoints) return
+    for (const key of checkpoints) {
+      const next = (snapshotPinCounts.get(key) ?? 1) - 1
+      if (next === 0) {
+        snapshotPinCounts.delete(key)
+        historicalSnapshots.delete(key)
+      } else {
+        snapshotPinCounts.set(key, next)
+      }
+    }
+    snapshotPins.delete(pinKey)
+  }
+
   const dispatchOf = (sessionId: string): SessionDispatch => {
     let state = dispatch.get(sessionId)
     if (!state) {
@@ -295,16 +329,41 @@ export function memory(options: MemoryStoreOptions = {}): A2Store {
   const readStateOf = (
     sessionId: string,
     reducerName: string,
+    throughIndex?: number,
+    snapshotThroughIndex?: number,
   ): StoreStateRead => {
-    const snap = snapshots.get(`${sessionId}\u0000${reducerName}`)
-    const snapshot = snap
-      ? { index: snap.index, state: structuredClone(snap.state) }
+    const key = snapshotKey(sessionId, reducerName)
+    const snapshotFrontier = snapshotThroughIndex ?? throughIndex
+    const head = snapshots.get(key)
+    let candidate =
+      head && (snapshotFrontier === undefined || head.index <= snapshotFrontier)
+        ? head
+        : null
+    if (snapshotFrontier !== undefined) {
+      const prefix = `${key}\u0000`
+      for (const [historicalKey, snapshot] of historicalSnapshots) {
+        if (
+          historicalKey.startsWith(prefix) &&
+          snapshot.index <= snapshotFrontier &&
+          (!candidate || snapshot.index > candidate.index)
+        ) {
+          candidate = snapshot
+        }
+      }
+    }
+    const snapshot = candidate
+      ? { index: candidate.index, state: structuredClone(candidate.state) }
       : null
     const afterIndex = snapshot?.index ?? 0
     return {
+      headIndex: head?.index ?? null,
       snapshot,
       events: (sessions.get(sessionId) ?? [])
-        .filter((row) => row.index > afterIndex)
+        .filter(
+          (row) =>
+            row.index > afterIndex &&
+            (throughIndex === undefined || row.index <= throughIndex),
+        )
         .map(toEvent),
     }
   }
@@ -488,6 +547,7 @@ export function memory(options: MemoryStoreOptions = {}): A2Store {
       parent.claimHolder = null
       parent.claimExpiresAt = null
       settle(sessionId, parent)
+      releaseSnapshotPins(sessionId, index)
       notifyStreams(sessionId, inserted)
       return { outcome: 'completed', events: inserted.map(toStored) }
     },
@@ -520,6 +580,7 @@ export function memory(options: MemoryStoreOptions = {}): A2Store {
       if (row.failureCount >= maxFailures) {
         row.failedAt = new Date(failedAt)
         dispatchOf(sessionId).ready.delete(row.index)
+        releaseSnapshotPins(sessionId, index)
         return {
           outcome: 'dead_lettered',
           failureCount: row.failureCount,
@@ -528,26 +589,83 @@ export function memory(options: MemoryStoreOptions = {}): A2Store {
       return { outcome: 'failed', failureCount: row.failureCount }
     },
 
-    async readState(sessionId, reducerName) {
-      return readStateOf(sessionId, reducerName)
+    async readState(sessionId, reducerName, stateOptions) {
+      return readStateOf(
+        sessionId,
+        reducerName,
+        stateOptions?.throughIndex,
+        stateOptions?.snapshotThroughIndex,
+      )
     },
 
-    async readStates(sessionIds, reducerName) {
-      return sessionIds.map((sessionId) => readStateOf(sessionId, reducerName))
+    async readStates(requests, reducerName) {
+      return requests.map((request) =>
+        readStateOf(
+          request.sessionId,
+          reducerName,
+          request.throughIndex,
+          request.snapshotThroughIndex,
+        ),
+      )
     },
 
-    async putSnapshot(sessionId, reducerName, index, state) {
-      const key = `${sessionId}\u0000${reducerName}`
-      const existing = snapshots.get(key)
-      // Guarded upsert: a slower concurrent writer must never clobber a
-      // further-along snapshot. A lost race costs a few refolded events
-      // next read — never correctness.
-      if (existing && existing.index >= index) return
-      snapshots.set(key, {
-        index,
-        state: structuredClone(state),
-        updatedAt: clock.now(),
-      })
+    async putSnapshots(sessionId, reducerName, writes) {
+      const key = snapshotKey(sessionId, reducerName)
+      for (const write of writes.toSorted((a, b) => a.index - b.index)) {
+        const checkpoint = checkpointKey(sessionId, reducerName, write.index)
+        for (const eventIndex of write.pinEventIndexes ?? []) {
+          const event = sessions.get(sessionId)?.[eventIndex - 1]
+          if (!event || event.processedAt !== null || event.failedAt !== null) {
+            continue
+          }
+          const pinKey = eventPinKey(sessionId, eventIndex)
+          let pinned = snapshotPins.get(pinKey)
+          if (!pinned) {
+            pinned = new Set()
+            snapshotPins.set(pinKey, pinned)
+          }
+          if (!pinned.has(checkpoint)) {
+            pinned.add(checkpoint)
+            snapshotPinCounts.set(
+              checkpoint,
+              (snapshotPinCounts.get(checkpoint) ?? 0) + 1,
+            )
+          }
+        }
+
+        const current = snapshots.get(key)
+        if (!current || write.index >= current.index) {
+          if (current && write.index > current.index) {
+            const currentCheckpoint = checkpointKey(
+              sessionId,
+              reducerName,
+              current.index,
+            )
+            if ((snapshotPinCounts.get(currentCheckpoint) ?? 0) > 0) {
+              historicalSnapshots.set(currentCheckpoint, {
+                index: current.index,
+                state: structuredClone(current.state),
+                updatedAt: new Date(current.updatedAt),
+              })
+            }
+          }
+          snapshots.set(key, {
+            index: write.index,
+            state: structuredClone(write.state),
+            updatedAt: clock.now(),
+          })
+          historicalSnapshots.delete(checkpoint)
+        } else if (
+          write.index < current.index &&
+          (snapshotPinCounts.get(checkpoint) ?? 0) > 0
+        ) {
+          historicalSnapshots.set(checkpoint, {
+            index: write.index,
+            state: structuredClone(write.state),
+            updatedAt: clock.now(),
+          })
+        }
+      }
     },
 
     presence: {