Package com.google.genkit.plugins.mongodb.session

Class MongoSessionStore<S>

java.lang.Object
com.google.genkit.plugins.mongodb.session.MongoSessionStore<S>
Type Parameters:
S - the type of custom session state
All Implemented Interfaces:
SessionStore<S>, SnapshotReader<S>, SnapshotSubscriber, SnapshotWriter<S>
public final class MongoSessionStore<S> extends Object implements SessionStore<S>, SnapshotSubscriber
MongoDB-backed implementation of SessionStore and SnapshotSubscriber.

Persists session snapshots with the same sharded checkpoint + diff + pointer layout as the Firestore, DynamoDB, Cosmos DB, and PostgreSQL backends (see FirestoreSessionStore), sharing the pure logic in SnapshotSharding.

Storage layout (single collection)

All documents live in one collection (default database genkit, collection genkit_sessions). Each document's _id is <prefix>::<recordId> where prefix is the per-tenant prefix (default "global") and recordId discriminates the record kind:

  • SNAP_<snapshotId> — one metadata document per snapshot. kind is "checkpoint" or "diff"; carries checkpointId, checkpointShardCount, segmentPath, statePatch (RFC-6902 as a JSON string for diffs) and error (JSON string).
  • SHARD_<checkpointId>_<index> — a shard of the checkpoint state JSON.
  • PTR_<sessionId> — the current leaf pointer for a session.

Concurrency

Shard and snapshot documents are idempotent by _id (upserted); updating an existing snapshot id uses a version conditional replace, re-applying the (pure) mutator on conflict. The session pointer is advanced monotonically (never backward) under the same version concurrency. Shards are written before the snapshot document, which is written before the pointer flips, so a reader following the pointer always sees complete data.

  • Constructor Details

    • MongoSessionStore

      public MongoSessionStore(com.mongodb.client.MongoClient client)
      Creates a store with default options.
      Parameters:
      client - the MongoDB client

    • MongoSessionStore

      public MongoSessionStore(com.mongodb.client.MongoClient client, MongoSessionStoreOptions options)
      Creates a store.
      Parameters:
      client - the MongoDB client
      options - the store options

  • Method Details

    • saveSnapshot

      public String saveSnapshot(String snapshotId, SnapshotMutator<S> mutator, SessionStoreOptions storeOpts)
      Atomically reads the snapshot identified by snapshotId (or creates a new one), applies mutator, and persists the result.

      Implements the same identity/sessionId/status defaulting contract as the reference stores, then writes the snapshot (checkpoint shards + metadata document) and advances the session pointer.

      Specified by:
      saveSnapshot in interface SnapshotWriter<S>
      Parameters:
      snapshotId - the ID of the snapshot to write; if null the store mints a UUID
      mutator - the pure read-modify-write function; receives the existing snapshot (or null) and returns the snapshot to persist, or null to decline
      storeOpts - store options
      Returns:
      the snapshot ID that was written, or null if the mutator declined

    • getSnapshot

      public SessionSnapshot<S> getSnapshot(GetSnapshotOptions opts)
      Retrieves a session snapshot.

      When GetSnapshotOptions.getSnapshotId() is set the store returns that specific snapshot. When GetSnapshotOptions.getSessionId() is set the store resolves and returns the latest leaf snapshot for that session. Returns null if no matching snapshot exists.

      By snapshotId: loads the snapshot document and reconstructs its state from the checkpoint shards + ordered segmentPath diffs. By sessionId: reads the pointer, then loads and reconstructs the pointed snapshot.

      Specified by:
      getSnapshot in interface SnapshotReader<S>
      Parameters:
      opts - options specifying which snapshot to retrieve; exactly one of snapshotId/sessionId must be set
      Returns:
      the matching snapshot, or null if none exists

    • onSnapshotStateChange

      public AutoCloseable onSnapshotStateChange(String snapshotId, Consumer<SessionSnapshot<?>> cb, SessionStoreOptions storeOpts)
      Subscribes to status changes for the snapshot identified by snapshotId.

      Invokes cb immediately with the current snapshot and again on every subsequent change. The returned AutoCloseable unsubscribes when closed.

      The subscription polls getSnapshot(com.google.genkit.ai.agent.GetSnapshotOptions) on a shared daemon scheduler and fires the callback whenever the serialized snapshot content changes. The callback also fires immediately if the snapshot already exists.

      Specified by:
      onSnapshotStateChange in interface SnapshotSubscriber
      Parameters:
      snapshotId - the snapshot to observe
      cb - callback invoked with the updated snapshot on each change
      storeOpts - store options
      Returns:
      a handle that cancels the subscription when AutoCloseable.close() is called