saveSnapshot method
Atomically reads the current snapshot (if snapshotId is provided),
passes it to mutator, and persists the result.
context carries the ambient request/action context (e.g. the
authenticated user) so multi-tenant stores can route writes.
Returns the snapshotId that was used, or null when the mutator
returned null.
Implementation
@override
Future<String?> saveSnapshot(
String? snapshotId,
SnapshotMutator mutator, {
Map<String, dynamic>? context,
}) {
return db.runTransaction((tx) async {
final reader = _Reader(tx);
// Reads phase 1: load the existing snapshot (if any) so the mutator can
// inspect the current full state.
_Reconstructed? existing;
if (snapshotId != null && snapshotId.isNotEmpty) {
existing = await _reconstruct(reader, snapshotId, context);
}
final current = existing != null
? _toSnapshot(existing.doc, existing.state)
: null;
final result = mutator(current);
if (result == null) return null;
final id = _resolveSnapshotId(
snapshotId != null && snapshotId.isNotEmpty ? snapshotId : null,
result,
);
// Prefer the snapshot's top-level `sessionId`; fall back to the id
// carried in its state for rows written before snapshot-level ids
// existed.
final sessionId = _snapshotSessionId(result);
if (sessionId == null) {
throw GenkitException(
"FirestoreSessionStore requires 'sessionId' to be set on the "
'snapshot.',
status: StatusCodes.INVALID_ARGUMENT,
);
}
final newState = result.state?.toJson() ?? <String, dynamic>{};
// Reads phase 2: the per-session pointer (current leaf metadata).
final pointerRef = _pointersCol(context).doc(sessionId);
final pointerSnap = await tx.get(pointerRef);
final pointer = pointerSnap.exists
? _PointerDoc.fromData(pointerSnap.data()!)
: null;
_ChainWrite chain;
if (existing != null) {
// Upsert: preserve the document's role and chain position; only the
// state/metadata change. Callers must only upsert the *leaf* -
// rewriting a non-leaf snapshot's state would invalidate its
// descendants' diffs (re-checkpointing prunes trailing shards and
// promotion nulls a diff's statePatch, either of which corrupts a
// descendant that still depends on this document's chain position).
//
// Nothing enforces "leaf only" at the type level, so enforce the
// invariant that guarantees it: only a *terminal* snapshot can have
// descendants (the runtime only ever resumes - and thus parents a child
// onto - a `completed` snapshot), so a snapshot in a terminal state may
// already have descendants and must never be rewritten in place. Every
// legitimate upsert (a detached `pending -> completed` upgrade, an
// `abort`'s `pending -> aborted`) targets a non-terminal snapshot, so
// this rejects only genuine misuse - loudly, before any shard is
// pruned - rather than silently corrupting the chain.
if (existing.doc.status != null &&
_terminalStatuses.contains(existing.doc.status)) {
throw GenkitException(
"FirestoreSessionStore: cannot upsert snapshot '$id' because it is "
"in a terminal state ('${existing.doc.status}'). Terminal snapshots "
'are immutable and may have descendants; write a new child snapshot '
'instead.',
status: StatusCodes.FAILED_PRECONDITION,
);
}
if (existing.doc.kind == 'checkpoint') {
chain = _writeCheckpoint(
tx,
id,
newState,
context,
existing.doc.checkpointShardCount,
);
} else {
// Reads phase 3 (diff upsert): resolve parent state for the patch.
final parentState = existing.doc.parentId != null
? (await _reconstruct(
reader,
existing.doc.parentId!,
context,
))?.state
: null;
final candidatePatch = diff(parentState, newState);
// Promote an oversized diff to a (sharded) checkpoint so even an
// in-place leaf rewrite can never push the document past the 1 MiB
// limit. Safe because callers only upsert the leaf, which has no
// descendants depending on its chain position.
if (_byteLength(candidatePatch) > shardSize) {
chain = _writeCheckpoint(tx, id, newState, context);
} else {
chain = (
kind: 'diff',
checkpointId: existing.doc.checkpointId,
checkpointShardCount: existing.doc.checkpointShardCount,
segmentPath: existing.doc.segmentPath,
statePatch: candidatePatch,
);
}
}
} else {
// New snapshot: resolve the parent's *chain metadata* (no state) to
// decide diff vs checkpoint. Materializing the parent's full state is
// deferred until we know we actually need a diff - so the expensive
// reconstruction is skipped on every checkpoint-boundary turn (which
// would rewrite the whole state regardless).
_ChainMeta? parentMeta;
if (result.parentId != null) {
parentMeta = await _loadParentChainMeta(
reader,
result.parentId!,
pointer,
context,
);
}
if (result.parentId == null ||
parentMeta == null ||
parentMeta.segmentPath.length + 1 >= checkpointInterval) {
// Write a full checkpoint without ever reconstructing the parent's
// state, for any of: a session root, an orphaned parent, or reaching
// the checkpoint interval (whose final segment is exactly the
// longest, costliest one we'd otherwise pay to reconstruct here).
chain = _writeCheckpoint(tx, id, newState, context);
} else {
// Diff candidate: now we must materialize the parent's state to
// compute the patch.
final parentState = (await _reconstructFrom(
reader,
parentMeta.checkpointId,
parentMeta.checkpointShardCount,
parentMeta.segmentPath,
result.parentId!,
context,
))?.state;
final candidatePatch = diff(parentState, newState);
// Promote oversized diffs to checkpoints so a single large turn is
// sharded rather than rejected by the 1 MiB limit.
if (_byteLength(candidatePatch) > shardSize) {
chain = _writeCheckpoint(tx, id, newState, context);
} else {
chain = (
kind: 'diff',
checkpointId: parentMeta.checkpointId,
checkpointShardCount: parentMeta.checkpointShardCount,
segmentPath: [...parentMeta.segmentPath, id],
statePatch: candidatePatch,
);
}
}
}
// Writes phase.
final doc = _SnapshotDoc(
snapshotId: id,
sessionId: sessionId,
parentId: result.parentId,
createdAt: result.createdAt,
updatedAt: result.updatedAt ?? result.createdAt,
status: result.status?.value,
heartbeatAt: result.heartbeatAt,
finishReason: result.finishReason?.value,
error: result.error?.toJson(),
kind: chain.kind,
checkpointId: chain.checkpointId,
checkpointShardCount: chain.checkpointShardCount,
segmentPath: chain.segmentPath,
statePatch: chain.statePatch,
);
tx.set(
_snapshotsCol(context).doc(id),
(_sanitize(doc.toData()) as Map).cast<String, Object?>(),
);
// Update the pointer in one of two cases:
//
// - Refresh: we just rewrote the snapshot the pointer already tracks
// (an in-place leaf upsert), so re-point at its (possibly changed)
// chain metadata under the same id.
// - Advance: a brand-new leaf. When there is no pointer yet it wins
// unconditionally; otherwise it must be strictly newer - by
// `(createdAt, snapshotId)`, the same ordering `selectLeafSnapshot` and
// the Go store use - than the leaf the pointer currently tracks. Gating
// on recency (rather than always advancing) means a backdated or
// concurrently-committed *older* leaf can't clobber a newer one and
// stick, which would otherwise violate the 'most recently created leaf'
// contract under clock skew / concurrency. We rely on `(createdAt, id)`
// here rather than a full-collection scan fallback (as InMemory/File
// do) so `sessionId` resolution stays a single pointer read - the whole
// point of the pointer design.
final isNew = existing == null;
final isRefresh = pointer != null && pointer.currentSnapshotId == id;
final advances =
isNew &&
(pointer == null ||
_isNewerLeaf(
result.createdAt,
id,
pointer.createdAt,
pointer.currentSnapshotId,
));
if (isRefresh || advances) {
tx.set(
pointerRef,
_PointerDoc(
currentSnapshotId: advances ? id : pointer!.currentSnapshotId,
checkpointId: chain.checkpointId,
checkpointShardCount: chain.checkpointShardCount,
segmentPath: chain.segmentPath,
createdAt: result.createdAt,
updatedAt: DateTime.now().toUtc().toIso8601String(),
).toData(),
);
}
return id;
});
}