@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.
476 lines • 18.3 kB
JavaScript
"use strict";
var __defProp = Object.defineProperty;
var __getOwnPropDesc = Object.getOwnPropertyDescriptor;
var __getOwnPropNames = Object.getOwnPropertyNames;
var __hasOwnProp = Object.prototype.hasOwnProperty;
var __export = (target, all) => {
for (var name in all)
__defProp(target, name, { get: all[name], enumerable: true });
};
var __copyProps = (to, from, except, desc) => {
if (from && typeof from === "object" || typeof from === "function") {
for (let key of __getOwnPropNames(from))
if (!__hasOwnProp.call(to, key) && key !== except)
__defProp(to, key, { get: () => from[key], enumerable: !(desc = __getOwnPropDesc(from, key)) || desc.enumerable });
}
return to;
};
var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: true }), mod);
var firestore_exports = {};
__export(firestore_exports, {
FirestoreSessionStore: () => FirestoreSessionStore
});
module.exports = __toCommonJS(firestore_exports);
var import_firestore = require("@google-cloud/firestore");
var import_beta = require("genkit/beta");
var import_logging = require("genkit/logging");
const DEFAULT_CHECKPOINT_INTERVAL = 25;
const DEFAULT_SHARD_SIZE = 512 * 1024;
const DEFAULT_PREFIX = "global";
function sanitize(value) {
return JSON.parse(JSON.stringify(value ?? null));
}
class FirestoreSessionStore {
db;
collection;
checkpointInterval;
shardSize;
snapshotPathPrefix;
constructor(opts) {
this.db = opts?.db ?? new import_firestore.Firestore();
this.collection = opts?.collection ?? "genkit-sessions";
this.checkpointInterval = opts?.checkpointInterval ?? DEFAULT_CHECKPOINT_INTERVAL;
this.shardSize = opts?.shardSize ?? DEFAULT_SHARD_SIZE;
this.snapshotPathPrefix = opts?.snapshotPathPrefix;
}
/** Resolves the (per-tenant) prefix for the given call options. */
prefixFor(options) {
return this.snapshotPathPrefix ? this.snapshotPathPrefix(options) : DEFAULT_PREFIX;
}
/** The (per-tenant) snapshots subcollection. */
snapshotsCol(options) {
return this.db.collection(this.collection).doc(this.prefixFor(options)).collection("snapshots");
}
/** The (per-tenant) pointers subcollection. */
pointersCol(options) {
return this.db.collection(`${this.collection}-pointers`).doc(this.prefixFor(options)).collection("pointers");
}
/** The (per-tenant) shards subcollection. */
shardsCol(options) {
return this.db.collection(`${this.collection}-shards`).doc(this.prefixFor(options)).collection("shards");
}
async getSnapshot(opts) {
const { snapshotId, sessionId } = this.normalize(opts);
const options = { context: opts.context };
return this.db.runTransaction(
async (tx) => {
const reader = this.reader(tx);
if (sessionId) {
const pointerSnap = await tx.get(
this.pointersCol(options).doc(sessionId)
);
if (!pointerSnap.exists) return void 0;
const pointer = pointerSnap.data();
const reconstructed2 = await this.reconstructFrom(
reader,
pointer.checkpointId,
pointer.checkpointShardCount,
pointer.segmentPath,
pointer.currentSnapshotId,
options
);
if (!reconstructed2) return void 0;
return this.toSnapshot(reconstructed2.doc, reconstructed2.state);
}
const reconstructed = await this.reconstruct(
reader,
snapshotId,
options
);
if (!reconstructed) return void 0;
return this.toSnapshot(reconstructed.doc, reconstructed.state);
},
{ readOnly: true }
);
}
async saveSnapshot(snapshotId, mutator, options) {
return this.db.runTransaction(async (tx) => {
const reader = this.reader(tx);
let existing;
if (snapshotId) {
existing = await this.reconstruct(reader, snapshotId, options);
}
const current = existing ? this.toSnapshot(existing.doc, existing.state) : void 0;
const result = mutator(current);
if (result === null) return null;
const id = snapshotId || result.snapshotId || globalThis.crypto.randomUUID();
const sessionId = result.sessionId ?? result.state?.sessionId;
if (!sessionId) {
throw new import_beta.GenkitError({
status: "INVALID_ARGUMENT",
message: `FirestoreSessionStore requires 'sessionId' to be set on the snapshot.`
});
}
const newState = result.state ?? {};
const pointerRef = this.pointersCol(options).doc(sessionId);
const pointerSnap = await tx.get(pointerRef);
const pointer = pointerSnap.exists ? pointerSnap.data() : void 0;
let kind;
let checkpointId;
let checkpointShardCount;
let segmentPath;
let statePatch;
if (existing) {
if (existing.doc.kind === "checkpoint") {
({
kind,
checkpointId,
checkpointShardCount,
segmentPath,
statePatch
} = this.writeCheckpoint(
tx,
id,
newState,
options,
existing.doc.checkpointShardCount
));
} else {
const parentState = existing.doc.parentId ? (await this.reconstruct(reader, existing.doc.parentId, options))?.state : void 0;
const candidatePatch = (0, import_beta.diff)(parentState, newState);
if (this.byteLength(candidatePatch) > this.shardSize) {
({
kind,
checkpointId,
checkpointShardCount,
segmentPath,
statePatch
} = this.writeCheckpoint(tx, id, newState, options));
} else {
kind = "diff";
checkpointId = existing.doc.checkpointId;
checkpointShardCount = existing.doc.checkpointShardCount;
segmentPath = existing.doc.segmentPath;
statePatch = candidatePatch;
}
}
} else {
let parentMeta;
if (result.parentId) {
parentMeta = await this.loadParentChainMeta(
reader,
result.parentId,
pointer,
options
);
}
if (!result.parentId || !parentMeta || parentMeta.segmentPath.length + 1 >= this.checkpointInterval) {
({
kind,
checkpointId,
checkpointShardCount,
segmentPath,
statePatch
} = this.writeCheckpoint(tx, id, newState, options));
} else {
const parentState = (await this.reconstructFrom(
reader,
parentMeta.checkpointId,
parentMeta.checkpointShardCount,
parentMeta.segmentPath,
result.parentId,
options
))?.state;
const candidatePatch = (0, import_beta.diff)(parentState, newState);
if (this.byteLength(candidatePatch) > this.shardSize) {
({
kind,
checkpointId,
checkpointShardCount,
segmentPath,
statePatch
} = this.writeCheckpoint(tx, id, newState, options));
} else {
kind = "diff";
checkpointId = parentMeta.checkpointId;
checkpointShardCount = parentMeta.checkpointShardCount;
segmentPath = [...parentMeta.segmentPath, id];
statePatch = candidatePatch;
}
}
}
const doc = {
snapshotId: id,
sessionId,
parentId: result.parentId,
createdAt: result.createdAt,
updatedAt: result.updatedAt ?? result.createdAt,
status: result.status,
heartbeatAt: result.heartbeatAt,
finishReason: result.finishReason,
error: result.error,
kind,
checkpointId,
checkpointShardCount,
segmentPath,
statePatch
};
tx.set(this.snapshotsCol(options).doc(id), sanitize(doc));
const isNew = !existing;
if (isNew || !pointer || pointer.currentSnapshotId === id) {
tx.set(
pointerRef,
sanitize({
currentSnapshotId: isNew || !pointer ? id : pointer.currentSnapshotId,
checkpointId,
checkpointShardCount,
segmentPath,
updatedAt: (/* @__PURE__ */ new Date()).toISOString()
})
);
}
return id;
});
}
onSnapshotStateChange(snapshotId, callback, options) {
const ref = this.snapshotsCol(options).doc(snapshotId);
return ref.onSnapshot(async (docSnap) => {
if (!docSnap.exists) return;
try {
const snapshot = await this.getSnapshot({
snapshotId,
context: options?.context
});
if (snapshot) callback(snapshot);
} catch (err) {
import_logging.logger.error(
`FirestoreSessionStore.watch failed to load snapshot ${snapshotId}`,
err
);
}
});
}
/**
* 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.
*/
normalize(opts) {
const snapshotId = opts.snapshotId?.trim() ? opts.snapshotId : void 0;
const sessionId = opts.sessionId?.trim() ? opts.sessionId : void 0;
if (!!snapshotId === !!sessionId) {
throw new import_beta.GenkitError({
status: "INVALID_ARGUMENT",
message: `getSnapshot requires exactly one non-empty 'snapshotId' or 'sessionId' (got ${snapshotId ? "snapshotId" : opts.snapshotId !== void 0 ? "blank snapshotId" : "neither"}${sessionId ? " and sessionId" : opts.sessionId !== void 0 ? " and blank sessionId" : ""}).`
});
}
return { snapshotId, sessionId };
}
/** Builds a {@link Reader} bound to a transaction or the bare instance. */
reader(tx) {
if (tx) {
return {
get: (ref) => tx.get(ref),
getAll: (refs) => refs.length ? tx.getAll(...refs) : Promise.resolve([])
};
}
return {
get: (ref) => ref.get(),
getAll: (refs) => refs.length ? this.db.getAll(...refs) : Promise.resolve([])
};
}
/**
* 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.
*/
async loadParentChainMeta(reader, parentId, pointer, options) {
if (pointer && pointer.currentSnapshotId === parentId) {
return {
checkpointId: pointer.checkpointId,
checkpointShardCount: pointer.checkpointShardCount,
segmentPath: pointer.segmentPath
};
}
const snap = await reader.get(this.snapshotsCol(options).doc(parentId));
if (!snap.exists) return void 0;
const d = snap.data();
return {
checkpointId: d.checkpointId,
checkpointShardCount: d.checkpointShardCount,
segmentPath: d.segmentPath
};
}
/**
* 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.
*/
async reconstruct(reader, id, options) {
const snap = await reader.get(this.snapshotsCol(options).doc(id));
if (!snap.exists) return void 0;
const d = snap.data();
return this.reconstructFrom(
reader,
d.checkpointId,
d.checkpointShardCount,
d.segmentPath,
id,
options
);
}
/**
* 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.
*/
async reconstructFrom(reader, checkpointId, shardCount, segmentPath, targetId, options) {
const targetIsCheckpoint = segmentPath.length === 0;
const snapshotsCol = this.snapshotsCol(options);
const shardsCol = this.shardsCol(options);
const checkpointRef = snapshotsCol.doc(checkpointId);
const shardRefs = Array.from(
{ length: shardCount },
(_, i) => shardsCol.doc(`${checkpointId}_${i}`)
);
const segRefs = segmentPath.map((sid) => snapshotsCol.doc(sid));
const snaps = await reader.getAll([
// The checkpoint document is only needed when it is itself the target;
// otherwise the last segment document carries the target metadata.
...targetIsCheckpoint ? [checkpointRef] : [],
...shardRefs,
...segRefs
]);
const byPath = /* @__PURE__ */ new Map();
for (const s of snaps) byPath.set(s.ref.path, s);
const shardSnaps = shardRefs.map((ref) => byPath.get(ref.path));
let state = this.stitch(shardSnaps);
if (targetIsCheckpoint) {
const checkpointSnap = byPath.get(checkpointRef.path);
if (!checkpointSnap?.exists) return void 0;
const checkpointDoc = checkpointSnap.data();
if (checkpointDoc.snapshotId !== targetId) return void 0;
return { doc: checkpointDoc, state: state ?? {} };
}
let targetDoc;
for (const ref of segRefs) {
const segSnap = byPath.get(ref.path);
if (!segSnap?.exists) return void 0;
const segDoc = segSnap.data();
state = (0, import_beta.applyPatch)(state, segDoc.statePatch ?? []);
targetDoc = segDoc;
}
if (!targetDoc || targetDoc.snapshotId !== targetId) return void 0;
return { doc: targetDoc, state: state ?? {} };
}
/**
* 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.
*/
writeShards(tx, checkpointId, state, options, oldShardCount = 0) {
const shardsCol = this.shardsCol(options);
const buf = Buffer.from(JSON.stringify(state ?? null), "utf8");
const count = Math.max(1, Math.ceil(buf.length / this.shardSize));
for (let i = 0; i < count; i++) {
const chunk = Buffer.from(
buf.subarray(i * this.shardSize, (i + 1) * this.shardSize)
);
tx.set(shardsCol.doc(`${checkpointId}_${i}`), {
chunk
});
}
for (let i = count; i < oldShardCount; i++) {
tx.delete(shardsCol.doc(`${checkpointId}_${i}`));
}
return count;
}
/**
* 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.
*/
writeCheckpoint(tx, id, state, options, oldShardCount = 0) {
return {
kind: "checkpoint",
checkpointId: id,
checkpointShardCount: this.writeShards(
tx,
id,
state,
options,
oldShardCount
),
segmentPath: [],
statePatch: void 0
};
}
/** Concatenates ordered shard documents and parses the materialized state. */
stitch(shardSnaps) {
if (shardSnaps.length === 0) return void 0;
const buffers = [];
for (const s of shardSnaps) {
if (!s.exists) {
throw new import_beta.GenkitError({
status: "DATA_LOSS",
message: `FirestoreSessionStore: missing checkpoint shard '${s.id}'.`
});
}
buffers.push(s.data().chunk);
}
return JSON.parse(Buffer.concat(buffers).toString("utf8"));
}
/** UTF-8 byte length of a JSON-serializable value. */
byteLength(value) {
return Buffer.byteLength(JSON.stringify(value ?? null), "utf8");
}
/** Assembles a {@link SessionSnapshot} from a document and its state. */
toSnapshot(doc, state) {
const snapshot = {
snapshotId: doc.snapshotId,
sessionId: doc.sessionId,
createdAt: doc.createdAt,
// Normalize to plain objects: values reconstructed from Firestore
// documents (e.g. patch operands) can carry non-plain prototypes.
state: sanitize(state)
};
if (doc.parentId !== void 0) snapshot.parentId = doc.parentId;
if (doc.updatedAt !== void 0) snapshot.updatedAt = doc.updatedAt;
if (doc.status !== void 0) snapshot.status = doc.status;
if (doc.heartbeatAt !== void 0) snapshot.heartbeatAt = doc.heartbeatAt;
if (doc.finishReason !== void 0)
snapshot.finishReason = doc.finishReason;
if (doc.error !== void 0) snapshot.error = doc.error;
return snapshot;
}
}
// Annotate the CommonJS export names for ESM import in node:
0 && (module.exports = {
FirestoreSessionStore
});
//# sourceMappingURL=firestore.js.map