@genkit-ai/google-cloud
Version:
Genkit AI framework plugin for Google Cloud Platform including Firestore trace/state store and deployment helpers for Cloud Functions for Firebase.
196 lines (192 loc) • 9.77 kB
text/typescript
import { Firestore } from '@google-cloud/firestore';
import { SessionStore, SessionStoreOptions, SessionSnapshot, SnapshotMutator } from 'genkit/beta';
/**
* Copyright 2026 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
/**
* Options for {@link FirestoreSessionStore}.
*/
interface FirestoreSessionStoreOptions {
/**
* An explicit Firestore instance. Defaults to a new {@link Firestore}
* instance (which picks up Application Default Credentials and the
* `FIRESTORE_EMULATOR_HOST` environment variable).
*/
db?: Firestore;
/**
* The collection where snapshot documents are stored (keyed by
* `snapshotId`). Defaults to `"genkit-sessions"`. Two companion collections
* are derived from it: `"<collection>-pointers"` holds one pointer document
* per session and `"<collection>-shards"` holds the sharded checkpoint
* state.
*/
collection?: string;
/**
* Number of turns between full-state checkpoints. A larger value stores
* fewer (but reconstructs over more) diffs; a smaller value reconstructs
* faster at the cost of more frequent full-state writes. Defaults to
* {@link DEFAULT_CHECKPOINT_INTERVAL}.
*
* Cost tuning: per-save reconstruction reads grow ~linearly with this value
* (so per-interval read work is ~quadratic in it), while checkpoint write and
* storage cost shrink with it. Lower it (e.g. 10) for small-state, read-heavy
* sessions; raise it (e.g. 50-100) for large per-turn state retained for a
* long time, where checkpoint write/storage amplification dominates.
*/
checkpointInterval?: number;
/**
* Maximum size in bytes of a single shard / diff document. Checkpoint state
* is split into chunks of this size, and any diff exceeding it is promoted to
* a (sharded) checkpoint so that no document approaches Firestore's 1 MiB
* limit. Defaults to {@link DEFAULT_SHARD_SIZE}.
*/
shardSize?: number;
/**
* Returns a per-tenant prefix derived from the call's
* {@link SessionStoreOptions} (e.g. the authenticated user id from
* `options.context`). When set, all snapshot, pointer and shard documents are
* nested under a tenant-scoped subcollection keyed by this prefix, so reads
* and writes are isolated per tenant: one tenant can never see (or even
* address) another's snapshots, even if they get hold of a `snapshotId` -
* resolving it still requires the matching, auth-derived prefix. Defaults to
* `"global"`.
*
* The prefix is used only to build document paths; the stored ids
* (`snapshotId`, `parentId`, `checkpointId`, ...) remain prefix-agnostic.
*/
snapshotPathPrefix?: (options?: SessionStoreOptions) => string;
}
/**
* A Firestore-backed {@link SessionStore} that persists session snapshots as
* incremental JSON Patch diffs anchored to periodic, sharded full-state
* checkpoints.
*
* Storage layout (the `<prefix>` segment is the per-tenant prefix returned by
* {@link FirestoreSessionStoreOptions.snapshotPathPrefix}, or `"global"` when
* none is configured):
*
* - `<collection>/<prefix>/snapshots/<snapshotId>` - one document per snapshot.
* A `diff` document holds the patch from its parent (`statePatch`); a
* `checkpoint` document holds a full-state materialization (sharded out of
* band).
* - `<collection>-shards/<prefix>/shards/<checkpointId>_<index>` - the sharded
* full state for a checkpoint.
* - `<collection>-pointers/<prefix>/pointers/<sessionId>` - one document per
* session pointing at the latest leaf snapshot and the metadata needed to
* reconstruct it.
*
* Reconstruction uses only document-ID lookups (`getAll`), so it needs no
* secondary indexes and is strongly consistent. No single document approaches
* the 1 MiB limit (state is sharded by `shardSize`), and the number of *diff*
* documents touched per read/write is bounded by `checkpointInterval` rather
* than total session length - so the store scales to arbitrarily long sessions
* (e.g. coding agents, long-lived chatbots). Note that checkpoints still store
* the full accumulated state, so checkpoint shard count (and the bytes written
* per checkpoint) grow with the state's size; tune `checkpointInterval` to
* trade per-save diff reads against checkpoint write amplification.
*/
declare class FirestoreSessionStore<S = unknown> implements SessionStore<S> {
protected db: Firestore;
protected collection: string;
protected checkpointInterval: number;
protected shardSize: number;
protected snapshotPathPrefix?: (options?: SessionStoreOptions) => string;
constructor(opts?: FirestoreSessionStoreOptions);
/** Resolves the (per-tenant) prefix for the given call options. */
private prefixFor;
/** The (per-tenant) snapshots subcollection. */
private snapshotsCol;
/** The (per-tenant) pointers subcollection. */
private pointersCol;
/** The (per-tenant) shards subcollection. */
private shardsCol;
getSnapshot(opts: {
snapshotId?: string;
sessionId?: string;
context?: SessionStoreOptions['context'];
}): Promise<SessionSnapshot<S> | undefined>;
saveSnapshot(snapshotId: string | undefined, mutator: SnapshotMutator<S>, options?: SessionStoreOptions): Promise<string | null>;
onSnapshotStateChange(snapshotId: string, callback: (snapshot: SessionSnapshot<S>) => void, options?: SessionStoreOptions): void | (() => void);
/**
* Validates that exactly one of `snapshotId` / `sessionId` is provided, and
* that the provided one is a non-blank string. Blank / whitespace-only ids
* are rejected up front (rather than silently treated as "absent") so callers
* get a clear `INVALID_ARGUMENT` instead of an unusable document key.
*/
private normalize;
/** Builds a {@link Reader} bound to a transaction or the bare instance. */
private reader;
/**
* Resolves a parent's chain metadata (nearest checkpoint, shard count and
* segment path) *without* materializing its - potentially large - state.
*
* In the common linear case the parent is the session's current leaf, so the
* metadata is read straight off the pointer and this performs *zero*
* document reads. Otherwise it reads the single parent document. Crucially,
* resolving only the metadata lets `saveSnapshot` decide diff-vs-checkpoint
* (which only needs `segmentPath.length`) before paying for a full state
* reconstruction - so checkpoint-boundary turns, which would rewrite the
* whole state anyway, skip reconstruction entirely.
*/
private loadParentChainMeta;
/**
* Reconstructs the state of `id` by reading its document to learn its
* checkpoint and segment path, then materializing from that checkpoint.
* Returns `undefined` when the snapshot does not exist.
*/
private reconstruct;
/**
* Materializes the state of `targetId` from a known checkpoint using a single
* batched, strongly-consistent `getAll`: the checkpoint's shards, the
* (bounded) segment of diff documents along `segmentPath`, and - only when
* the target *is* the checkpoint - the checkpoint document itself. The diffs
* are then applied in order onto the checkpoint's state. Cost is bounded by
* `checkpointInterval` + shard count, independent of total session length.
*
* Note: when `segmentPath` is non-empty the state comes entirely from the
* shards and the target's metadata from the last segment document, so the
* checkpoint *document* is not read - saving one read on the common path.
*/
private reconstructFrom;
/**
* Serializes `state` to UTF-8, splits it into `shardSize`-byte chunks, and
* writes them at `<checkpointId>_<index>`. When `oldShardCount` exceeds the
* new count (a shrinking re-checkpoint), the now-stale trailing shards are
* deleted. Returns the number of shards written.
*/
private writeShards;
/**
* Writes a full-state checkpoint at `id` (sharding the state via
* {@link writeShards}) and returns the snapshot metadata describing it: a
* checkpoint anchors itself (`checkpointId === id`), has an empty
* `segmentPath`, and carries no `statePatch`.
*
* This is the shared promotion path used whenever a snapshot must be a
* checkpoint rather than a diff - the session root, an orphaned parent, a
* checkpoint-interval boundary, an in-place checkpoint rewrite, and the
* promotion of an oversized diff (whether new turn or leaf upsert) so that no
* single document approaches Firestore's 1 MiB limit. Pass `oldShardCount`
* when re-checkpointing an existing checkpoint so stale trailing shards are
* pruned.
*/
private writeCheckpoint;
/** Concatenates ordered shard documents and parses the materialized state. */
private stitch;
/** UTF-8 byte length of a JSON-serializable value. */
private byteLength;
/** Assembles a {@link SessionSnapshot} from a document and its state. */
private toSnapshot;
}
export { FirestoreSessionStore, type FirestoreSessionStoreOptions };