Package com.google.genkit.ai.agent

Class InMemorySessionStore<S>

java.lang.Object
com.google.genkit.ai.agent.InMemorySessionStore<S>
Type Parameters:
S - the type of custom session state
All Implemented Interfaces:
SessionStore<S>, SnapshotReader<S>, SnapshotSubscriber, SnapshotWriter<S>
public final class InMemorySessionStore<S> extends Object implements SessionStore<S>, SnapshotSubscriber
Reference in-memory implementation of SessionStore and SnapshotSubscriber.

Used by conformance tests and as the simplest full implementation of the saveSnapshot read-modify-write contract.

Deep-copy strategy

Snapshots are deep-copied on both write (store) and read (return) using a JSON round-trip via the shared ObjectMapper (valueToTree → treeToValue(node, SessionSnapshot.class)). Because of Java's type erasure, the round-trip uses the raw type SessionSnapshot (without the <S> parameter), so the custom field of SessionState will be deserialized as a LinkedHashMap rather than as S. This is acceptable for the in-memory store used dynamically (conformance tests use Map<String,Object> state).

Subscriber notification

Callbacks are collected under the lock and invoked after releasing the lock to prevent deadlock if a callback re-enters the store (e.g. to read the new snapshot).

Subscriber semantics (Go-compatible)

If the snapshot is already present when onSnapshotStateChange(java.lang.String, java.util.function.Consumer<com.google.genkit.ai.agent.SessionSnapshot<?>>, com.google.genkit.ai.agent.SessionStoreOptions) is called, the callback is not invoked immediately. The callback fires only when a subsequent saveSnapshot(java.lang.String, com.google.genkit.ai.agent.SnapshotMutator<S>, com.google.genkit.ai.agent.SessionStoreOptions) causes a status change (including the first save when existing is null).

  • Constructor Details

  • Method Details

    • 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.

      When GetSnapshotOptions.getSnapshotId() is set, returns a deep copy of the stored snapshot (or null if absent). When GetSnapshotOptions.getSessionId() is set, gathers all snapshots whose sessionId equals it, applies LeafSelection, and returns a deep copy of the result. If neither id is set, returns null.

      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

    • saveSnapshot

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

      Implements the full saveSnapshot contract:

      1. Under lock: read existing snapshot (deep copy) identified by snapshotId.
      2. Invoke mutator.apply(existing); if result is null, return null (no write).
      3. Determine final id: snapshotId != null ? snapshotId : (result.snapshotId != null ? result.snapshotId : UUID.randomUUID()). Force result.snapshotId = finalId.
      4. Preserve sessionId from existing snapshot when updating.
      5. Validate sessionId is non-null and non-empty; throw GenkitException with INVALID_ARGUMENT otherwise.
      6. Default null status to SnapshotStatus.COMPLETED.
      7. Store deep copy. Collect subscriber callbacks if status changed (existing was null or status differs).
      8. Invoke collected callbacks after releasing the lock.
      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
      options - store options
      Returns:
      the snapshot ID that was written, or null if the mutator declined

    • onSnapshotStateChange

      public AutoCloseable onSnapshotStateChange(String snapshotId, Consumer<SessionSnapshot<?>> cb, SessionStoreOptions options)
      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.

      Registers cb to be invoked on subsequent status changes for snapshotId. If the snapshot already exists, the callback is invoked immediately with a deep copy of the current snapshot (before returning). The callback will also fire on any subsequent status changes. If the snapshot does not currently exist, the callback is registered and will fire on the first save (status change from null). Returns an AutoCloseable that unregisters the callback.

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