@genkit-ai/ai
Version:
Genkit AI framework generative AI APIs.
505 lines • 19.4 kB
JavaScript
import { GenkitError } from "@genkit-ai/core";
import * as fs from "fs";
import * as fsp from "fs/promises";
import * as path from "path";
import {
assertValidSessionId
} from "./session.mjs";
function normalizeGetSnapshotOptions(opts) {
const { snapshotId, sessionId } = opts;
if (!!snapshotId === !!sessionId) {
throw new GenkitError({
status: "INVALID_ARGUMENT",
message: `getSnapshot requires exactly one of 'snapshotId' or 'sessionId' (got ${snapshotId ? "snapshotId" : "neither"}${sessionId ? " and sessionId" : ""}).`
});
}
if (sessionId) {
assertValidSessionId(sessionId);
}
return { snapshotId, sessionId };
}
function selectLeafSnapshot(snapshots, sessionId, rejectBranching = false) {
if (snapshots.length === 0) return void 0;
const parentIds = /* @__PURE__ */ new Set();
for (const snap of snapshots) {
if (snap.parentId) parentIds.add(snap.parentId);
}
const leaves = snapshots.filter((s) => !parentIds.has(s.snapshotId));
if (leaves.length === 1) return leaves[0];
if (leaves.length === 0) {
throw new GenkitError({
status: "FAILED_PRECONDITION",
message: `Session '${sessionId}' has no leaf snapshot (corrupt or cyclic history). Resume by snapshotId instead.`
});
}
if (rejectBranching) {
throw new GenkitError({
status: "FAILED_PRECONDITION",
message: `Session '${sessionId}' has branching snapshots (${leaves.length} leaves), so there is no single latest snapshot. This happens when a conversation is branched (e.g. regenerate). Resume by snapshotId instead.`
});
}
return leaves.reduce(
(latest, snap) => snap.createdAt > latest.createdAt ? snap : latest
);
}
class InMemorySessionStore {
snapshots = /* @__PURE__ */ new Map();
listeners = /* @__PURE__ */ new Map();
rejectBranchingSessions;
/**
* @param options.rejectBranchingSessions When `true`, a `sessionId` lookup
* that resolves to a branched history (more than one leaf) throws
* `FAILED_PRECONDITION` instead of returning the latest leaf. Defaults to
* `false`; opt in (e.g. in dev) to surface accidental branching early.
*/
constructor(options) {
this.rejectBranchingSessions = options?.rejectBranchingSessions ?? false;
}
async getSnapshot(opts) {
const { snapshotId, sessionId } = normalizeGetSnapshotOptions(opts);
if (snapshotId) {
const snap = this.snapshots.get(snapshotId);
if (!snap) return void 0;
return structuredClone(snap);
}
const owned = [];
for (const snap of this.snapshots.values()) {
if ((snap.sessionId ?? snap.state?.sessionId) === sessionId) {
owned.push(snap);
}
}
const leaf = selectLeafSnapshot(
owned,
sessionId,
this.rejectBranchingSessions
);
return leaf ? structuredClone(leaf) : void 0;
}
async saveSnapshot(snapshotId, mutator, options) {
const current = snapshotId ? this.snapshots.get(snapshotId) : void 0;
const result = mutator(current ? structuredClone(current) : void 0);
if (result === null) return null;
const id = snapshotId || result.snapshotId || globalThis.crypto.randomUUID();
const full = {
...result,
snapshotId: id
};
this.snapshots.set(id, structuredClone(full));
const snapshotListeners = this.listeners.get(id);
if (snapshotListeners) {
for (const listener of snapshotListeners) {
listener(structuredClone(full));
}
}
return id;
}
onSnapshotStateChange(snapshotId, callback, options) {
if (!this.listeners.has(snapshotId)) {
this.listeners.set(snapshotId, []);
}
this.listeners.get(snapshotId).push(callback);
return () => {
const list = this.listeners.get(snapshotId);
if (list) {
const index = list.indexOf(callback);
if (index >= 0) list.splice(index, 1);
}
};
}
}
const DEFAULT_SNAPSHOT_WATCH_POLL_INTERVAL_MS = 2e3;
function assertSafeId(id, label) {
if (!id || id.includes("/") || id.includes("\\") || id.includes("\0") || id === "." || id === ".." || path.basename(id) !== id) {
throw new GenkitError({
status: "INVALID_ARGUMENT",
message: `Invalid ${label}: "${id}". A ${label} must be a plain file name (no path separators or "..").`
});
}
}
function assertSafeSnapshotId(snapshotId) {
assertSafeId(snapshotId, "snapshotId");
}
const POINTERS_SUBDIR = ".pointers";
class FileSessionStore {
dirPath;
maxPersistedChainLength;
snapshotPathPrefix;
rejectBranchingSessions;
snapshotWatchPollIntervalMs;
/**
* Per-file write locks. The {@link SessionStore} contract (and the
* abort-aware mutator that branches on `current.status`) assumes
* read-modify-write is atomic, but on the file system a read and the
* `writeFile` below it are not. Without a lock two concurrent saves can
* read the same `current` and the later write clobbers the earlier one
* (e.g. a `completed` write overwriting a concurrent `aborted`). We
* serialize saves per resolved file path with a simple promise chain.
*/
writeLocks = /* @__PURE__ */ new Map();
/**
* @param dirPath Directory where snapshot JSON files are stored.
* @param options.maxPersistedChainLength When set, snapshots older than this
* many entries in a chain are automatically deleted on each save.
* @param options.snapshotPathPrefix Returns a sub-directory prefix derived
* from the call's {@link SessionStoreOptions} (e.g. the authenticated user
* id from `options.context`), useful for multi-tenant isolation: all reads
* and writes are scoped to that prefix, so one tenant can never see
* another's snapshots. Defaults to `"global"`.
* @param options.rejectBranchingSessions When `true`, a `sessionId` lookup
* that resolves to a branched history (more than one leaf) throws
* `FAILED_PRECONDITION` instead of returning the latest leaf. Defaults to
* `false`; opt in (e.g. in dev) to surface accidental branching early.
* @param options.snapshotWatchPollIntervalMs Polling interval (ms) for the
* {@link FileSessionStore.onSnapshotStateChange} fallback that backstops
* `fs.watch` (which can miss events on some filesystems, e.g. network
* mounts). Defaults to {@link DEFAULT_SNAPSHOT_WATCH_POLL_INTERVAL_MS}.
*/
constructor(dirPath, options) {
this.dirPath = path.resolve(dirPath);
fs.mkdirSync(this.dirPath, { recursive: true });
this.maxPersistedChainLength = options?.maxPersistedChainLength;
this.snapshotPathPrefix = options?.snapshotPathPrefix;
this.rejectBranchingSessions = options?.rejectBranchingSessions ?? false;
this.snapshotWatchPollIntervalMs = options?.snapshotWatchPollIntervalMs ?? DEFAULT_SNAPSHOT_WATCH_POLL_INTERVAL_MS;
}
async ensureDir(dir) {
await fsp.mkdir(dir, { recursive: true });
}
/** Resolves the (per-tenant) directory snapshots are stored under. */
prefixDir(options) {
const prefix = this.snapshotPathPrefix ? this.snapshotPathPrefix(options) : "global";
return path.join(this.dirPath, prefix);
}
/**
* Resolves the file path for a given snapshotId: `<prefix>/<snapshotId>.json`.
*/
async getFilePath(snapshotId, options) {
assertSafeSnapshotId(snapshotId);
const dir = this.prefixDir(options);
await this.ensureDir(dir);
const filePath = path.join(dir, `${snapshotId}.json`);
const resolvedDir = path.resolve(dir);
const resolved = path.resolve(filePath);
if (resolved !== path.join(resolvedDir, `${snapshotId}.json`) || !resolved.startsWith(resolvedDir + path.sep)) {
throw new GenkitError({
status: "INVALID_ARGUMENT",
message: `Invalid snapshotId: "${snapshotId}". Resolved path escapes the snapshot directory.`
});
}
return filePath;
}
/** Resolves the (per-tenant) directory holding per-session pointer files. */
pointersDir(options) {
return path.join(this.prefixDir(options), POINTERS_SUBDIR);
}
/**
* Resolves the pointer file path for a session, validating `sessionId` is a
* plain basename so it can never escape the pointers directory. Pure: it does
* not create the directory, so the read path stays side-effect free. The
* write path calls {@link ensureDir} before writing.
*/
getPointerPath(sessionId, options) {
assertSafeId(sessionId, "sessionId");
return path.join(this.pointersDir(options), `${sessionId}.json`);
}
/**
* Reads the per-session {@link PointerDoc}, or `undefined` when it is missing
* (legacy store / not yet written) or unreadable / corrupt - callers fall
* back to a full directory scan in that case. Best-effort: any IO/parse error
* resolves to `undefined` so the optimization can never make a lookup (or
* save) fail where the scan-only baseline would have succeeded. An invalid
* `sessionId` still throws (path validation is resolved outside the try) so it
* fails fast rather than silently being ignored.
*/
async readPointer(sessionId, options) {
const filePath = this.getPointerPath(sessionId, options);
let contents;
try {
contents = await fsp.readFile(filePath, "utf-8");
} catch {
return void 0;
}
try {
const parsed = JSON.parse(contents);
return typeof parsed?.currentSnapshotId === "string" ? parsed : void 0;
} catch {
return void 0;
}
}
/**
* Atomically writes the per-session {@link PointerDoc}. Best-effort: a
* pointer write failure is swallowed since the pointer is only an
* optimization - `sessionId` lookups still self-heal via the full scan. An
* invalid `sessionId` still throws (path validation is resolved outside the
* try) so it fails fast rather than silently being ignored.
*/
async writePointer(sessionId, currentSnapshotId, options) {
const filePath = this.getPointerPath(sessionId, options);
const dir = this.pointersDir(options);
try {
await this.ensureDir(dir);
const pointer = {
currentSnapshotId,
updatedAt: (/* @__PURE__ */ new Date()).toISOString()
};
await this.atomicWrite(filePath, JSON.stringify(pointer, null, 2));
} catch {
}
}
/**
* Serializes async work per resolved file path so a read-modify-write in
* {@link saveSnapshot} is not interleaved with a concurrent one for the same
* snapshot (see {@link writeLocks}).
*/
async withFileLock(filePath, fn) {
const prev = this.writeLocks.get(filePath) ?? Promise.resolve();
const gate = prev.then(
() => void 0,
() => void 0
);
this.writeLocks.set(filePath, gate);
await gate;
try {
return await fn();
} finally {
if (this.writeLocks.get(filePath) === gate) {
this.writeLocks.delete(filePath);
}
}
}
async getSnapshot(opts) {
const { snapshotId, sessionId } = normalizeGetSnapshotOptions(opts);
if (sessionId) {
return this.getLatestSnapshotForSession(sessionId, opts);
}
const filePath = await this.getFilePath(snapshotId, opts);
try {
const fileContents = await fsp.readFile(filePath, "utf-8");
return JSON.parse(fileContents);
} catch (e) {
if (e.code === "ENOENT") return void 0;
throw e;
}
}
/**
* Loads a single snapshot file by its id (no sessionId branch). Used by
* internal traversal (parent chains) where we always have a concrete id.
*/
async getSnapshotById(snapshotId, options) {
const filePath = await this.getFilePath(snapshotId, options);
try {
const fileContents = await fsp.readFile(filePath, "utf-8");
return JSON.parse(fileContents);
} catch (e) {
if (e.code === "ENOENT") return void 0;
throw e;
}
}
/**
* Resolves the latest (leaf) snapshot for a session.
*
* Fast path: read the per-session pointer file and load the leaf it names -
* one pointer read plus one snapshot read, independent of session count /
* length. The pointer is skipped (and the scan used) when
* `rejectBranchingSessions` is set, since detecting branches requires seeing
* every leaf.
*
* Fallback (no/stale/corrupt pointer, or branch detection): scan every
* snapshot file in the prefix directory, keep those whose `sessionId`
* matches, select the single leaf, and refresh the pointer so later lookups
* take the fast path.
*
* Known limitation: the fast path trusts the pointer when the snapshot it
* names still exists and belongs to the session - it does not re-verify that
* it is the actual leaf. So if a save succeeds but the subsequent (best-effort)
* `writePointer` does not (crash/disk error), or two new saves for the same
* session race and the older one writes the pointer last, the pointer can
* linger on a valid-but-older same-session snapshot and lookups return it
* until the next save advances the pointer. This is the accepted trade-off for
* a best-effort cache: verifying leaf-ness on every read would reintroduce the
* full scan the pointer exists to avoid. Callers needing strict guarantees can
* resume by `snapshotId`, or set `rejectBranchingSessions` (which always
* scans).
*/
async getLatestSnapshotForSession(sessionId, options) {
if (!this.rejectBranchingSessions) {
const pointer = await this.readPointer(sessionId, options);
if (pointer) {
const snap = await this.getSnapshotById(
pointer.currentSnapshotId,
options
);
if (snap && (snap.sessionId ?? snap.state?.sessionId) === sessionId) {
return snap;
}
}
}
const dir = this.prefixDir(options);
let files;
try {
files = await fsp.readdir(dir);
} catch (e) {
if (e.code === "ENOENT") return void 0;
throw e;
}
const snapshots = [];
for (const file of files) {
if (!file.endsWith(".json")) continue;
try {
const contents = await fsp.readFile(path.join(dir, file), "utf-8");
const snap = JSON.parse(contents);
if ((snap.sessionId ?? snap.state?.sessionId) === sessionId) {
snapshots.push(snap);
}
} catch (e) {
if (e.code === "ENOENT") continue;
throw e;
}
}
const leaf = selectLeafSnapshot(
snapshots,
sessionId,
this.rejectBranchingSessions
);
if (!leaf) return void 0;
await this.writePointer(sessionId, leaf.snapshotId, options);
return leaf;
}
async saveSnapshot(snapshotId, mutator, options) {
if (snapshotId) {
const lockPath = await this.getFilePath(snapshotId, options);
return this.withFileLock(
lockPath,
() => this.saveSnapshotUnlocked(snapshotId, mutator, options)
);
}
return this.saveSnapshotUnlocked(snapshotId, mutator, options);
}
async saveSnapshotUnlocked(snapshotId, mutator, options) {
const current = snapshotId ? await this.getSnapshotById(snapshotId, options) : void 0;
const snapshot = mutator(current);
if (snapshot === null) return null;
const id = snapshotId || snapshot.snapshotId || globalThis.crypto.randomUUID();
const full = {
...snapshot,
snapshotId: id
};
const filePath = await this.getFilePath(id, options);
await this.atomicWrite(filePath, JSON.stringify(full, null, 2));
const sessionId = full.sessionId ?? full.state?.sessionId;
if (sessionId && !current) {
await this.writePointer(sessionId, id, options);
}
if (this.maxPersistedChainLength && this.maxPersistedChainLength > 0) {
let cur = full;
const chain = [];
while (cur) {
chain.push(cur.snapshotId);
if (cur.parentId) {
cur = await this.getSnapshotById(cur.parentId, options);
} else {
break;
}
}
if (chain.length > this.maxPersistedChainLength) {
for (let i = this.maxPersistedChainLength; i < chain.length; i++) {
const pathToDelete = await this.getFilePath(chain[i], options);
await fsp.unlink(pathToDelete).catch((e) => {
if (e.code !== "ENOENT") throw e;
});
}
}
}
return id;
}
/**
* Writes `contents` to `filePath` atomically: write to a temp file in the
* same directory, then rename over the target. `rename` is atomic on POSIX
* and Windows, so a concurrent reader in {@link getSnapshot} never observes a
* half-written (torn) file.
*/
async atomicWrite(filePath, contents) {
const tmpPath = `${filePath}.${process.pid}.${globalThis.crypto.randomUUID()}.tmp`;
try {
await fsp.writeFile(tmpPath, contents, "utf-8");
await fsp.rename(tmpPath, filePath);
} catch (e) {
await fsp.unlink(tmpPath).catch(() => {
});
throw e;
}
}
/**
* Watches a single snapshot file for changes and invokes `callback` with the
* parsed snapshot whenever it changes.
*
* Unlike {@link InMemorySessionStore}, file-backed snapshots are frequently
* mutated by a *different* process (e.g. the request handler that received an
* abort writes `status: 'aborted'`, while a detached background worker is the
* one watching). Detecting that requires observing the filesystem rather than
* in-process `saveSnapshot` calls.
*
* Reliability comes from two layers:
* - `fs.watch` on the (per-tenant) prefix directory, filtered to the target
* `<snapshotId>.json`. This is low latency but can miss events on some
* filesystems (network mounts, certain container volumes).
* - A polling fallback (`snapshotWatchPollIntervalMs`) that re-reads the file
* on an interval, backstopping any events `fs.watch` drops. Its timer is
* `unref`'d so it never keeps the process alive on its own.
*
* Callbacks are de-duplicated by serialized content, so the noisy/duplicate
* events `fs.watch` emits collapse into one callback per real change.
* Transient read errors (e.g. a partially written file mid-rewrite, or a
* not-yet-created file) are swallowed; the next event/poll re-reads.
*
* @returns An unsubscribe function that stops watching and polling.
*/
onSnapshotStateChange(snapshotId, callback, options) {
const dir = this.prefixDir(options);
fs.mkdirSync(dir, { recursive: true });
const fileName = `${snapshotId}.json`;
const filePath = path.join(dir, fileName);
let closed = false;
let lastSerialized;
const emitIfChanged = async () => {
if (closed) return;
let contents;
try {
contents = await fsp.readFile(filePath, "utf-8");
} catch (e) {
return;
}
if (closed || contents === lastSerialized) return;
let snapshot;
try {
snapshot = JSON.parse(contents);
} catch {
return;
}
lastSerialized = contents;
callback(snapshot);
};
let watcher;
try {
watcher = fs.watch(dir, (_event, changed) => {
if (!changed || changed === fileName) void emitIfChanged();
});
} catch {
}
const pollTimer = setInterval(() => {
void emitIfChanged();
}, this.snapshotWatchPollIntervalMs);
pollTimer.unref?.();
void emitIfChanged();
return () => {
closed = true;
watcher?.close();
clearInterval(pollTimer);
};
}
}
export {
FileSessionStore,
InMemorySessionStore
};
//# sourceMappingURL=session-stores.mjs.map