Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -60,8 +60,8 @@
*
* <pre>
* CREATE TABLE GRAPH_THREAD (
* thread_id VARCHAR(36) PRIMARY KEY,
* thread_name VARCHAR(255) NOT NULL,
* thread_id VARCHAR(36) PRIMARY KEY, -- internal surrogate id, not the user-facing thread id
* thread_name VARCHAR(255) NOT NULL, -- user-facing thread id accepted by the API
* is_released BOOLEAN DEFAULT FALSE NOT NULL,
* active_thread_name VARCHAR(255) GENERATED ALWAYS AS (
* CASE WHEN is_released = FALSE THEN thread_name ELSE NULL END
Expand All @@ -73,7 +73,7 @@
* CREATE TABLE GRAPH_CHECKPOINT (
* checkpoint_seq BIGINT GENERATED BY DEFAULT AS IDENTITY UNIQUE,
* checkpoint_id VARCHAR(36) PRIMARY KEY,
* thread_id VARCHAR(36) NOT NULL,
* thread_id VARCHAR(36) NOT NULL, -- references the internal surrogate id
* node_id VARCHAR(255),
* next_node_id VARCHAR(255),
* state_data CLOB NOT NULL,
Expand All @@ -88,6 +88,14 @@
* </pre>
* </p>
* <p>
* Thread identity: GRAPH_THREAD.thread_name stores the thread id supplied
* through {@code RunnableConfig}, while GRAPH_THREAD.thread_id stores an
* internally generated UUID that identifies one activation of that thread
* between a release and the next reuse of the same id. This column split is
* what allows a released thread id to be reused without orphaning the released
* checkpoint history.
* </p>
* <p>
* A builder can be used to create an instance of H2Saver. The builder allows
* configuring a DataSource or JDBC URL, CreateOption, StateSerializer, and the
* maximum number of latest checkpoints retained in memory.
Expand Down Expand Up @@ -388,30 +396,34 @@ protected void insertCheckpoint(String threadId, Checkpoint checkpoint) throws E
}
}

private String activeThreadId(Connection conn, String threadName) throws SQLException {
Optional<String> activeThreadId = selectActiveThreadId(conn, threadName);
/**
* Returns the internal surrogate thread id of the active row for the given
* user-facing thread id, inserting a new row when none is active.
*/
private String activeThreadId(Connection conn, String threadId) throws SQLException {
Optional<String> activeThreadId = selectActiveThreadId(conn, threadId);
if (activeThreadId.isPresent()) {
return activeThreadId.get();
}

String persistedThreadId = UUID.randomUUID().toString();
String newThreadId = UUID.randomUUID().toString();
try (PreparedStatement ps = conn.prepareStatement(INSERT_THREAD)) {
ps.setString(1, persistedThreadId);
ps.setString(2, threadName);
ps.setString(1, newThreadId);
ps.setString(2, threadId);
ps.executeUpdate();
return persistedThreadId;
return newThreadId;
}
catch (SQLException ex) {
if (isUniqueConstraintViolation(ex)) {
return selectActiveThreadId(conn, threadName).orElseThrow(() -> ex);
return selectActiveThreadId(conn, threadId).orElseThrow(() -> ex);
}
throw ex;
}
}

private Optional<String> selectActiveThreadId(Connection conn, String threadName) throws SQLException {
private Optional<String> selectActiveThreadId(Connection conn, String threadId) throws SQLException {
try (PreparedStatement ps = conn.prepareStatement(SELECT_ACTIVE_THREAD)) {
ps.setString(1, threadName);
ps.setString(1, threadId);
try (ResultSet rs = ps.executeQuery()) {
if (rs.next()) {
return Optional.of(rs.getString(1));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,15 @@
* This class owns the common saver lifecycle and latest-checkpoint cache behavior.
* Subclasses keep database-specific SQL, transaction details and row mapping logic.
* <p>
* Thread identity model: every saver operation is keyed on the user-facing thread
* id resolved from {@link RunnableConfig} (see
* {@link BaseCheckpointSaver#checkpointThreadId(RunnableConfig)}), which is the
* only thread identifier the public API exposes. Concrete schemas persist that
* id in their {@code thread_name} column, while their {@code thread_id} column
* holds an internally generated surrogate UUID identifying one activation of the
* thread between a release and the next reuse of the same id. The protected
* methods below always receive the user-facing thread id.
* <p>
* Replacement: add artifact
* {@code io.github.agentic-spring-ai:agentic-spring-ai-graph-persistence-jdbc}
* and use
Expand Down Expand Up @@ -209,7 +218,7 @@ public final void latestCheckpointCacheEnabled(boolean enabled) {
/**
* Selects the active checkpoint history for a thread from the backing database.
*
* @param threadId thread name/id used by the concrete saver schema
* @param threadId user-facing thread id; persisted by the concrete schema as its thread name
* @return checkpoint history in latest-first order
* @throws Exception when the concrete saver cannot read checkpoint history
*/
Expand All @@ -218,7 +227,7 @@ public final void latestCheckpointCacheEnabled(boolean enabled) {
/**
* Selects only the latest active checkpoint for a thread.
*
* @param threadId thread name/id used by the concrete saver schema
* @param threadId user-facing thread id; persisted by the concrete schema as its thread name
* @return latest checkpoint when one exists
* @throws Exception when the concrete saver cannot read the latest checkpoint
*/
Expand All @@ -227,7 +236,7 @@ public final void latestCheckpointCacheEnabled(boolean enabled) {
/**
* Selects an active checkpoint by id for a thread.
*
* @param threadId thread name/id used by the concrete saver schema
* @param threadId user-facing thread id; persisted by the concrete schema as its thread name
* @param checkpointId checkpoint id to look up
* @return matching checkpoint when one exists
* @throws Exception when the concrete saver cannot read the checkpoint
Expand All @@ -237,7 +246,7 @@ public final void latestCheckpointCacheEnabled(boolean enabled) {
/**
* Inserts a new active checkpoint for a thread.
*
* @param threadId thread name/id used by the concrete saver schema
* @param threadId user-facing thread id; persisted by the concrete schema as its thread name
* @param checkpoint checkpoint to persist
* @throws Exception when the concrete saver cannot insert the checkpoint
*/
Expand All @@ -246,7 +255,7 @@ public final void latestCheckpointCacheEnabled(boolean enabled) {
/**
* Replaces an existing active checkpoint for a thread.
*
* @param threadId thread name/id used by the concrete saver schema
* @param threadId user-facing thread id; persisted by the concrete schema as its thread name
* @param checkpointId checkpoint id to replace
* @param checkpoint replacement checkpoint data
* @throws Exception when the concrete saver cannot update the checkpoint
Expand All @@ -256,7 +265,7 @@ public final void latestCheckpointCacheEnabled(boolean enabled) {
/**
* Deletes active checkpoints by id for a thread.
*
* @param threadId thread name/id used by the concrete saver schema
* @param threadId user-facing thread id; persisted by the concrete schema as its thread name
* @param checkpointIds checkpoint ids to delete
* @throws Exception when the concrete saver cannot delete checkpoints
*/
Expand All @@ -281,7 +290,7 @@ private void deleteRetainedCheckpoints(String threadId, RunnableConfig config) t
/**
* Marks the active thread as released in the backing database.
*
* @param threadId thread name/id used by the concrete saver schema
* @param threadId user-facing thread id; persisted by the concrete schema as its thread name
* @throws Exception when the concrete saver cannot release the thread
*/
protected abstract void releaseThread(String threadId) throws Exception;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,16 @@
/**
* MongoDB checkpoint saver.
* <p>
* Thread identity: every operation is keyed on the user-facing thread id
* resolved from {@code RunnableConfig}. That id is used verbatim as the
* {@code _id} of the {@code thread_meta} document
* ({@code mongo:thread:meta:<user-thread-id>}) and stored in its
* {@code thread_name} field, while the {@code thread_id} field holds an
* internally generated UUID identifying one activation of the thread between a
* release and the next reuse of the same id. Checkpoint documents are stored
* under that internal id so a released thread id can be reused without
* orphaning the released checkpoint history.
* <p>
* Replacement: add artifact
* {@code io.github.agentic-spring-ai:agentic-spring-ai-graph-persistence-mongodb}
* and use {@code io.github.agentic.spring.ai.graph.persistence.mongodb.MongoSaver}.
Expand Down Expand Up @@ -155,20 +165,20 @@ private LinkedList<Checkpoint> deserializeCheckpoints(String content) throws IOE
}

/**
* Gets or creates a thread_id for the given thread_name.
* If an active thread exists, returns its thread_id.
* If no active thread exists or the thread is released, creates a new thread_id.
* Returns the internal surrogate thread id of the active entry for the given
* user-facing thread id, creating a new one when no active entry exists or
* the previous one was released.
*
* This method uses atomic operations to prevent race conditions in concurrent scenarios.
* Uses findOneAndUpdate with conditional logic to ensure thread-safe creation.
*
* @param threadName the thread name
* @param threadId the user-facing thread id
* @param clientSession the MongoDB client session for transaction
* @return the thread_id (UUID string)
* @return the internal thread id (UUID string)
*/
private String getOrCreateThreadId(String threadName, ClientSession clientSession) {
private String getOrCreateThreadId(String threadId, ClientSession clientSession) {
MongoCollection<Document> threadMetaCollection = database.getCollection(THREAD_META_COLLECTION);
String metaId = THREAD_META_PREFIX + threadName;
String metaId = THREAD_META_PREFIX + threadId;

// Step 1: Try to atomically get an active thread
// Filter: _id matches AND is_released != true
Expand All @@ -187,10 +197,10 @@ private String getOrCreateThreadId(String threadName, ClientSession clientSessio
);

if (existingDoc != null) {
String threadId = existingDoc.getString(FIELD_THREAD_ID);
if (threadId != null) {
// Active thread exists, return its thread_id
return threadId;
String persistedThreadId = existingDoc.getString(FIELD_THREAD_ID);
if (persistedThreadId != null) {
// Active thread exists, return its internal thread id
return persistedThreadId;
}
}

Expand Down Expand Up @@ -271,53 +281,57 @@ private String getOrCreateThreadId(String threadName, ClientSession clientSessio
}

/**
* Gets the active thread_id for the given thread_name.
* Gets the internal surrogate thread id of the active entry for the given
* user-facing thread id.
* Returns null if no active thread exists.
*
* @param threadName the thread name
* @param threadId the user-facing thread id
* @param clientSession the MongoDB client session for transaction
* @return the active thread_id, or null if not found
* @return the active internal thread id, or null if not found
*/
private String getActiveThreadId(String threadName, ClientSession clientSession) {
private String getActiveThreadId(String threadId, ClientSession clientSession) {
MongoCollection<Document> threadMetaCollection = database.getCollection(THREAD_META_COLLECTION);
String metaId = THREAD_META_PREFIX + threadName;
String metaId = THREAD_META_PREFIX + threadId;

Document metaDoc = threadMetaCollection.find(clientSession, new BasicDBObject("_id", metaId)).first();

if (metaDoc != null) {
String threadId = metaDoc.getString(FIELD_THREAD_ID);
String persistedThreadId = metaDoc.getString(FIELD_THREAD_ID);
Boolean isReleased = metaDoc.getBoolean(FIELD_IS_RELEASED, false);

if (threadId != null && !Boolean.TRUE.equals(isReleased)) {
return threadId;
if (persistedThreadId != null && !Boolean.TRUE.equals(isReleased)) {
return persistedThreadId;
}
}

return null; // No active thread exists
}

private String threadName(RunnableConfig config) {
/**
* Resolves the user-facing thread id every saver operation is keyed on.
*/
private String threadId(RunnableConfig config) {
return checkpointThreadId(config);
}

@Override
public Collection<Checkpoint> list(RunnableConfig config) {
String threadName = threadName(config);
String threadId = threadId(config);
ClientSession clientSession = this.client
.startSession(ClientSessionOptions.builder().defaultTransactionOptions(txnOptions).build());
clientSession.startTransaction();
List<Checkpoint> checkpoints = null;
try {
// Get active thread_id for the thread_name
String threadId = getActiveThreadId(threadName, clientSession);
if (threadId == null) {
// Get the internal thread id of the active entry
String persistedThreadId = getActiveThreadId(threadId, clientSession);
if (persistedThreadId == null) {
clientSession.commitTransaction();
return Collections.emptyList();
}

// Use thread_id to query checkpoints
// Use the internal thread id to query checkpoints
MongoCollection<Document> collection = database.getCollection(CHECKPOINT_COLLECTION);
String checkpointId = CHECKPOINT_PREFIX + threadId;
String checkpointId = CHECKPOINT_PREFIX + persistedThreadId;
Document document = collection.find(clientSession, new BasicDBObject("_id", checkpointId)).first();
if (document == null) {
clientSession.commitTransaction();
Expand All @@ -339,23 +353,23 @@ public Collection<Checkpoint> list(RunnableConfig config) {

@Override
public Optional<Checkpoint> get(RunnableConfig config) {
String threadName = threadName(config);
String threadId = threadId(config);
ClientSession clientSession = this.client
.startSession(ClientSessionOptions.builder().defaultTransactionOptions(txnOptions).build());
LinkedList<Checkpoint> checkpoints = null;
try {
clientSession.startTransaction();

// Get active thread_id for the thread_name
String threadId = getActiveThreadId(threadName, clientSession);
if (threadId == null) {
// Get the internal thread id of the active entry
String persistedThreadId = getActiveThreadId(threadId, clientSession);
if (persistedThreadId == null) {
clientSession.commitTransaction();
return Optional.empty();
}

// Use thread_id to query checkpoints
// Use the internal thread id to query checkpoints
MongoCollection<Document> collection = database.getCollection(CHECKPOINT_COLLECTION);
String checkpointId = CHECKPOINT_PREFIX + threadId;
String checkpointId = CHECKPOINT_PREFIX + persistedThreadId;
Document document = collection.find(clientSession, new BasicDBObject("_id", checkpointId)).first();
if (document == null) {
clientSession.commitTransaction();
Expand Down Expand Up @@ -385,17 +399,17 @@ public Optional<Checkpoint> get(RunnableConfig config) {

@Override
public RunnableConfig put(RunnableConfig config, Checkpoint checkpoint) throws Exception {
String threadName = threadName(config);
String threadId = threadId(config);
ClientSession clientSession = this.client
.startSession(ClientSessionOptions.builder().defaultTransactionOptions(txnOptions).build());
clientSession.startTransaction();
try {
// Get or create thread_id
String threadId = getOrCreateThreadId(threadName, clientSession);
// Get or create the internal thread id
String persistedThreadId = getOrCreateThreadId(threadId, clientSession);

// Use thread_id as key for checkpoint storage
// Use the internal thread id as key for checkpoint storage
MongoCollection<Document> collection = database.getCollection(CHECKPOINT_COLLECTION);
String checkpointDocId = CHECKPOINT_PREFIX + threadId;
String checkpointDocId = CHECKPOINT_PREFIX + persistedThreadId;
Document document = collection.find(clientSession, new BasicDBObject("_id", checkpointDocId)).first();
LinkedList<Checkpoint> checkpointLinkedList = null;

Expand Down Expand Up @@ -449,24 +463,24 @@ public RunnableConfig put(RunnableConfig config, Checkpoint checkpoint) throws E

@Override
public Tag release(RunnableConfig config) throws Exception {
String threadName = threadName(config);
String threadId = threadId(config);
ClientSession clientSession = this.client
.startSession(ClientSessionOptions.builder().defaultTransactionOptions(txnOptions).build());
clientSession.startTransaction();
try {
MongoCollection<Document> threadMetaCollection = database.getCollection(THREAD_META_COLLECTION);
String metaId = THREAD_META_PREFIX + threadName;
String metaId = THREAD_META_PREFIX + threadId;

Document metaDoc = threadMetaCollection.find(clientSession, new BasicDBObject("_id", metaId)).first();
if (metaDoc == null) {
clientSession.abortTransaction();
throw new IllegalStateException("Thread not found: " + threadName);
throw new IllegalStateException("Thread not found: " + threadId);
}

String threadId = metaDoc.getString(FIELD_THREAD_ID);
if (threadId == null) {
String persistedThreadId = metaDoc.getString(FIELD_THREAD_ID);
if (persistedThreadId == null) {
clientSession.abortTransaction();
throw new IllegalStateException("Thread not found: " + threadName);
throw new IllegalStateException("Thread not found: " + threadId);
}

// Mark thread as released atomically
Expand All @@ -484,12 +498,12 @@ public Tag release(RunnableConfig config) throws Exception {
if (updatedDoc == null) {
// Thread was already released or doesn't exist
clientSession.abortTransaction();
throw new IllegalStateException("Thread is not active or already released: " + threadName);
throw new IllegalStateException("Thread is not active or already released: " + threadId);
}

// Get checkpoints for Tag (using thread_id)
// Get checkpoints for Tag (using the internal thread id)
MongoCollection<Document> checkpointCollection = database.getCollection(CHECKPOINT_COLLECTION);
String checkpointDocId = CHECKPOINT_PREFIX + threadId;
String checkpointDocId = CHECKPOINT_PREFIX + persistedThreadId;
Document checkpointDoc = checkpointCollection.find(clientSession, new BasicDBObject("_id", checkpointDocId))
.first();

Expand All @@ -502,7 +516,7 @@ public Tag release(RunnableConfig config) throws Exception {
}

clientSession.commitTransaction();
return new Tag(threadName, checkpoints);
return new Tag(threadId, checkpoints);

}
catch (Exception e) {
Expand Down
Loading