Package com.google.genkit.plugins.azurefoundry.session

Class CosmosSessionStore<S>

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

Persists session snapshots with the same sharded checkpoint + diff + pointer layout as the Firestore backend (see FirestoreSessionStore), sharing the pure logic in SnapshotSharding.

Storage layout (single container)

All documents live in one container (default database genkit, container genkit-sessions) partitioned by /pk (the per-tenant prefix, default "global"). Each document's id 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.

Cosmos document ids may not contain #, so the layout uses _ separators; ids are only ever constructed and looked up (never parsed back into components).

Concurrency

Snapshot and shard documents are idempotent by id (upserted); the snapshot document uses ETag (If-Match) optimistic concurrency when updating an existing id, re-applying the (pure) mutator on 412. The session pointer is advanced monotonically (never backward) under ETag 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

    • CosmosSessionStore

      public CosmosSessionStore(com.azure.cosmos.CosmosClient client)
      Creates a store with default options.
      Parameters:
      client - the Cosmos DB client

    • CosmosSessionStore

      public CosmosSessionStore(com.azure.cosmos.CosmosClient client, CosmosSessionStoreOptions options)
      Creates a store.
      Parameters:
      client - the Cosmos DB 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