From fb83a67be8e599268cda54b0d7d4509a0afbdd69 Mon Sep 17 00:00:00 2001 From: LiangshouX Date: Sun, 9 Aug 2026 21:40:24 +0800 Subject: [PATCH 1/8] feat(extensions-mongodb): add MongoDB storage extension --- agentscope-dependencies-bom/pom.xml | 8 + .../agentscope-all/pom.xml | 7 + .../agentscope-bom/pom.xml | 7 + .../agentscope-extensions-mongodb/pom.xml | 71 +++ .../mongodb/MongoDistributedStore.java | 135 +++++ .../sandbox/MongoSandboxExecutionGuard.java | 228 +++++++++ .../snapshot/MongoRemoteSnapshotClient.java | 127 +++++ .../mongodb/snapshot/MongoSnapshotSpec.java | 38 ++ .../mongodb/state/MongoAgentStateStore.java | 480 ++++++++++++++++++ .../mongodb/store/MongoBaseStore.java | 236 +++++++++ .../mongodb/MongoDistributedStoreTest.java | 87 ++++ .../state/MongoAgentStateStoreTest.java | 293 +++++++++++ .../mongodb/store/MongoBaseStoreTest.java | 172 +++++++ agentscope-extensions/pom.xml | 1 + 14 files changed, 1890 insertions(+) create mode 100644 agentscope-extensions/agentscope-extensions-mongodb/pom.xml create mode 100644 agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/MongoDistributedStore.java create mode 100644 agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/sandbox/MongoSandboxExecutionGuard.java create mode 100644 agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClient.java create mode 100644 agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoSnapshotSpec.java create mode 100644 agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStore.java create mode 100644 agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/store/MongoBaseStore.java create mode 100644 agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/MongoDistributedStoreTest.java create mode 100644 agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreTest.java create mode 100644 agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/store/MongoBaseStoreTest.java diff --git a/agentscope-dependencies-bom/pom.xml b/agentscope-dependencies-bom/pom.xml index 0604b123eb..7762f045fa 100644 --- a/agentscope-dependencies-bom/pom.xml +++ b/agentscope-dependencies-bom/pom.xml @@ -100,6 +100,7 @@ 0.3.3.Final 7.4.1 6.4.2.RELEASE + 5.3.1 3.3.2 2.5.2 7.0.7 @@ -405,6 +406,13 @@ ${lettuce.version} + + + org.mongodb + mongodb-driver-sync + ${mongodb-driver.version} + + com.xuxueli diff --git a/agentscope-distribution/agentscope-all/pom.xml b/agentscope-distribution/agentscope-all/pom.xml index 744dba161a..d793f716ac 100644 --- a/agentscope-distribution/agentscope-all/pom.xml +++ b/agentscope-distribution/agentscope-all/pom.xml @@ -241,6 +241,13 @@ true + + io.agentscope + agentscope-extensions-mongodb + compile + true + + io.agentscope agentscope-extensions-studio diff --git a/agentscope-distribution/agentscope-bom/pom.xml b/agentscope-distribution/agentscope-bom/pom.xml index 4742fd0912..b98f51b744 100644 --- a/agentscope-distribution/agentscope-bom/pom.xml +++ b/agentscope-distribution/agentscope-bom/pom.xml @@ -246,6 +246,13 @@ ${project.version} + + + io.agentscope + agentscope-extensions-mongodb + ${project.version} + + io.agentscope diff --git a/agentscope-extensions/agentscope-extensions-mongodb/pom.xml b/agentscope-extensions/agentscope-extensions-mongodb/pom.xml new file mode 100644 index 0000000000..dd535d2d17 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mongodb/pom.xml @@ -0,0 +1,71 @@ + + + + + 4.0.0 + + io.agentscope + agentscope-extensions + ${revision} + ../pom.xml + + + AgentScope Java - Extensions - MongoDB + MongoDB-backed distributed implementation for AgentStateStore. Provides MongoDistributedStore for one-line distributed configuration. + agentscope-extensions-mongodb + + + + io.agentscope + agentscope-core + provided + + + + io.agentscope + agentscope-harness + provided + + + + + org.mongodb + mongodb-driver-sync + + + + + org.junit.jupiter + junit-jupiter + test + + + + org.mockito + mockito-core + test + + + + org.mockito + mockito-junit-jupiter + test + + + diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/MongoDistributedStore.java b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/MongoDistributedStore.java new file mode 100644 index 0000000000..ca7b63f6fb --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/MongoDistributedStore.java @@ -0,0 +1,135 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.extensions.mongodb; + +import com.mongodb.ConnectionString; +import com.mongodb.MongoClientSettings; +import com.mongodb.client.MongoClient; +import com.mongodb.client.MongoClients; +import com.mongodb.client.MongoDatabase; +import io.agentscope.core.state.AgentStateStore; +import io.agentscope.extensions.mongodb.sandbox.MongoSandboxExecutionGuard; +import io.agentscope.extensions.mongodb.snapshot.MongoSnapshotSpec; +import io.agentscope.extensions.mongodb.state.MongoAgentStateStore; +import io.agentscope.extensions.mongodb.store.MongoBaseStore; +import io.agentscope.harness.agent.DistributedStore; +import io.agentscope.harness.agent.filesystem.remote.store.BaseStore; +import io.agentscope.harness.agent.sandbox.SandboxExecutionGuard; +import io.agentscope.harness.agent.sandbox.snapshot.SandboxSnapshotSpec; +import java.util.Objects; + +/** + * MongoDB-backed {@link DistributedStore}. + * + *

Usage: + * + * 

{@code
+ * MongoClient mongoClient = MongoClients.create("mongodb://localhost:27017");
+ *
+ * HarnessAgent agent = HarnessAgent.builder()
+ *     .name("my-agent")
+ *     .model("dashscope:qwen-plus")
+ *     .distributedStore(MongoDistributedStore.create(mongoClient, "agentscope"))
+ *     .build();
+ * }
+ * + *

This configures: + * + *

+ * + *

The caller owns the {@link MongoClient} lifecycle; closing the store does NOT close the + * client. + */ +public class MongoDistributedStore implements DistributedStore { + + private static final String DEFAULT_DATABASE = "agentscope"; + private static final String STATE_COLLECTION = "agentscope_sessions"; + private static final String BASE_COLLECTION = "agentscope_base"; + + private final MongoClient mongoClient; + private final String databaseName; + + private MongoDistributedStore(MongoClient mongoClient, String databaseName) { + this.mongoClient = Objects.requireNonNull(mongoClient, "mongoClient"); + this.databaseName = databaseName != null ? databaseName : DEFAULT_DATABASE; + } + + /** + * Creates a MongoDB distributed store with the default database name ({@code "agentscope"}). + * + * @param mongoClient the MongoDB client + * @return a new MongoDB distributed store + */ + public static MongoDistributedStore create(MongoClient mongoClient) { + return new MongoDistributedStore(mongoClient, null); + } + + /** + * Creates a MongoDB distributed store. + * + * @param mongoClient the MongoDB client + * @param databaseName the database name + * @return a new MongoDB distributed store + */ + public static MongoDistributedStore create(MongoClient mongoClient, String databaseName) { + return new MongoDistributedStore(mongoClient, databaseName); + } + + /** + * Creates a MongoDB distributed store from a connection string. A new {@link MongoClient} is + * created internally. The caller is responsible for closing the client when done. + * + * @param connectionString the MongoDB connection string + * @return a new MongoDB distributed store + */ + public static MongoDistributedStore fromConnectionString(String connectionString) { + MongoClientSettings settings = + MongoClientSettings.builder() + .applyConnectionString(new ConnectionString(connectionString)) + .build(); + return new MongoDistributedStore(MongoClients.create(settings), null); + } + + @Override + public AgentStateStore agentStateStore() { + return MongoAgentStateStore.builder() + .mongoClient(mongoClient) + .databaseName(databaseName) + .collectionName(STATE_COLLECTION) + .build(); + } + + @Override + public BaseStore baseStore() { + MongoDatabase db = mongoClient.getDatabase(databaseName); + return new MongoBaseStore(db, BASE_COLLECTION); + } + + @Override + public SandboxSnapshotSpec sandboxSnapshotSpec() { + return new MongoSnapshotSpec(mongoClient, databaseName); + } + + @Override + public SandboxExecutionGuard sandboxExecutionGuard() { + return MongoSandboxExecutionGuard.builder(mongoClient).databaseName(databaseName).build(); + } +} diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/sandbox/MongoSandboxExecutionGuard.java b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/sandbox/MongoSandboxExecutionGuard.java new file mode 100644 index 0000000000..f6f577b32a --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/sandbox/MongoSandboxExecutionGuard.java @@ -0,0 +1,228 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.extensions.mongodb.sandbox; + +import com.mongodb.MongoBulkWriteException; +import com.mongodb.client.MongoCollection; +import com.mongodb.client.MongoDatabase; +import com.mongodb.client.model.Filters; +import com.mongodb.client.model.FindOneAndUpdateOptions; +import com.mongodb.client.model.IndexOptions; +import com.mongodb.client.model.Indexes; +import com.mongodb.client.model.Updates; +import io.agentscope.harness.agent.sandbox.SandboxExecutionGuard; +import io.agentscope.harness.agent.sandbox.SandboxIsolationKey; +import io.agentscope.harness.agent.sandbox.SandboxLease; +import java.nio.charset.StandardCharsets; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.time.Duration; +import java.util.Date; +import java.util.Objects; +import org.bson.Document; +import org.bson.conversions.Bson; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * MongoDB-based {@link SandboxExecutionGuard}. + * + *

Uses a dedicated MongoDB collection as a distributed lock mechanism. Each lock is a document + * with a unique {@code _id} derived from the {@link SandboxIsolationKey} and a TTL index on {@code + * expiresAt} to auto-release stale locks. + * + *

Lock acquisition uses {@code findOneAndUpdate} with upsert and a filter that rejects documents + * whose {@code expiresAt} has not yet passed. This provides a non-blocking try-lock semantics. + * Acquisition polls until the lock is obtained or the timeout expires. + */ +public final class MongoSandboxExecutionGuard implements SandboxExecutionGuard { + + private static final Logger log = LoggerFactory.getLogger(MongoSandboxExecutionGuard.class); + + private static final String DEFAULT_COLLECTION = "agentscope_sandbox_locks"; + private static final String FIELD_LOCK_ID = "_id"; + private static final String FIELD_OWNER = "owner"; + private static final String FIELD_EXPIRES_AT = "expiresAt"; + + private final MongoCollection collection; + private final long lockTimeoutMs; + private final String owner; + + private MongoSandboxExecutionGuard(Builder builder) { + MongoDatabase db = builder.mongoClient.getDatabase(builder.databaseName); + this.collection = db.getCollection(builder.collectionName); + this.lockTimeoutMs = builder.lockTimeout.toMillis(); + this.owner = builder.owner; + ensureIndexes(); + } + + /** + * Creates a new builder. + * + * @param mongoClient the MongoDB client + * @return a new builder + */ + public static Builder builder(com.mongodb.client.MongoClient mongoClient) { + return new Builder(mongoClient); + } + + @Override + public SandboxLease tryEnter(SandboxIsolationKey key) throws InterruptedException { + String lockId = composeLockId(key); + log.debug("[sandbox-guard] Acquiring MongoDB lock: {}", lockId); + + long deadline = System.nanoTime() + Duration.ofMillis(lockTimeoutMs).toNanos(); + while (true) { + Date now = new Date(); + Date expiresAt = new Date(now.getTime() + lockTimeoutMs); + + Bson filter = + Filters.and( + Filters.eq(FIELD_LOCK_ID, lockId), + Filters.or( + Filters.exists(FIELD_OWNER, false), + Filters.lte(FIELD_EXPIRES_AT, now))); + + Bson update = + Updates.combine( + Updates.setOnInsert(FIELD_LOCK_ID, lockId), + Updates.set(FIELD_OWNER, owner), + Updates.set(FIELD_EXPIRES_AT, expiresAt)); + + try { + Document result = + collection.findOneAndUpdate( + filter, update, new FindOneAndUpdateOptions().upsert(true)); + + if (result == null) { + log.debug("[sandbox-guard] Acquired MongoDB lock: {}", lockId); + return new MongoLease(collection, lockId); + } else { + log.debug( + "[sandbox-guard] Lock held by {}, retrying: {}", + result.getString(FIELD_OWNER), + lockId); + } + } catch (MongoBulkWriteException e) { + // Duplicate key — lock held by someone else, retry + } catch (Exception e) { + if (e.getMessage() != null && e.getMessage().contains("E11000")) { + // Duplicate key error — lock held by someone else + } else { + throw new RuntimeException("Failed to acquire MongoDB lock: " + lockId, e); + } + } + + if (System.nanoTime() >= deadline) { + throw new InterruptedException( + "Timed out waiting for MongoDB lock: " + + lockId + + " (timeout=" + + Duration.ofMillis(lockTimeoutMs) + + ")"); + } + + Thread.sleep(100L); + } + } + + private void ensureIndexes() { + collection.createIndex( + Indexes.ascending(FIELD_EXPIRES_AT), + new IndexOptions().expireAfter(0L, java.util.concurrent.TimeUnit.SECONDS)); + } + + private static String composeLockId(SandboxIsolationKey key) { + String raw = key.getScope().name().toLowerCase() + ":" + key.getValue(); + try { + MessageDigest md = MessageDigest.getInstance("SHA-256"); + byte[] hash = md.digest(raw.getBytes(StandardCharsets.UTF_8)); + StringBuilder sb = new StringBuilder("lock:"); + for (int i = 0; i < 8; i++) { + sb.append(String.format("%02x", hash[i])); + } + return sb.toString(); + } catch (NoSuchAlgorithmException e) { + throw new RuntimeException("SHA-256 not available", e); + } + } + + private static final class MongoLease implements SandboxLease { + + private final MongoCollection collection; + private final String lockId; + + MongoLease(MongoCollection collection, String lockId) { + this.collection = collection; + this.lockId = lockId; + } + + @Override + public void close() { + try { + collection.deleteOne(Filters.eq(lockId)); + log.debug("[sandbox-guard] Released MongoDB lock: {}", lockId); + } catch (Exception e) { + log.warn( + "[sandbox-guard] Failed to release MongoDB lock {}: {}", + lockId, + e.getMessage()); + } + } + } + + /** Builder for {@link MongoSandboxExecutionGuard}. */ + public static final class Builder { + + private final com.mongodb.client.MongoClient mongoClient; + private String databaseName = "agentscope"; + private String collectionName = DEFAULT_COLLECTION; + private Duration lockTimeout = Duration.ofMinutes(30); + private String owner = "agentscope:" + ProcessHandle.current().pid(); + + Builder(com.mongodb.client.MongoClient mongoClient) { + this.mongoClient = Objects.requireNonNull(mongoClient, "mongoClient"); + } + + public Builder databaseName(String databaseName) { + this.databaseName = databaseName; + return this; + } + + public Builder collectionName(String collectionName) { + this.collectionName = collectionName; + return this; + } + + public Builder lockTimeout(Duration timeout) { + Objects.requireNonNull(timeout, "timeout"); + if (timeout.isNegative() || timeout.isZero()) { + throw new IllegalArgumentException("timeout must be positive"); + } + this.lockTimeout = timeout; + return this; + } + + public Builder owner(String owner) { + this.owner = owner; + return this; + } + + public MongoSandboxExecutionGuard build() { + return new MongoSandboxExecutionGuard(this); + } + } +} diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClient.java b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClient.java new file mode 100644 index 0000000000..81167d7f8a --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClient.java @@ -0,0 +1,127 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.extensions.mongodb.snapshot; + +import com.mongodb.client.MongoCollection; +import com.mongodb.client.MongoDatabase; +import com.mongodb.client.model.Filters; +import com.mongodb.client.model.ReplaceOptions; +import io.agentscope.harness.agent.sandbox.snapshot.RemoteSnapshotClient; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.util.Date; +import java.util.Objects; +import org.bson.Document; +import org.bson.types.Binary; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * {@link RemoteSnapshotClient} backed by a MongoDB collection. + * + *

Stores sandbox workspace tar archives as BSON Binary in a collection with documents of the + * form {@code {_id: snapshotId, data: Binary, createdAt: Date}}. + */ +public class MongoRemoteSnapshotClient implements RemoteSnapshotClient { + + private static final Logger log = LoggerFactory.getLogger(MongoRemoteSnapshotClient.class); + + private static final String DEFAULT_COLLECTION = "agentscope_snapshots"; + private static final String FIELD_DATA = "data"; + private static final String FIELD_CREATED_AT = "createdAt"; + private static final int MAX_SNAPSHOT_BYTES = 100 * 1024 * 1024; // 100 MB + + private final MongoCollection collection; + + public MongoRemoteSnapshotClient( + com.mongodb.client.MongoClient mongoClient, + String databaseName, + String collectionName, + boolean initializeSchema) { + Objects.requireNonNull(mongoClient, "mongoClient"); + String coll = collectionName != null ? collectionName : DEFAULT_COLLECTION; + MongoDatabase db = + mongoClient.getDatabase(databaseName != null ? databaseName : "agentscope"); + this.collection = db.getCollection(coll); + if (initializeSchema) { + initSchema(); + } + } + + private void initSchema() { + try { + collection.createIndex(com.mongodb.client.model.Indexes.ascending(FIELD_CREATED_AT)); + } catch (Exception e) { + log.warn( + "Failed to initialize snapshot collection index '{}': {}", + collection.getNamespace(), + e.getMessage()); + } + } + + @Override + public void upload(String snapshotId, InputStream data) throws Exception { + Objects.requireNonNull(snapshotId, "snapshotId"); + Objects.requireNonNull(data, "data"); + byte[] bytes = readAllBounded(data, MAX_SNAPSHOT_BYTES); + Document doc = + new Document(FIELD_DATA, new Binary(bytes)).append(FIELD_CREATED_AT, new Date()); + collection.replaceOne(Filters.eq(snapshotId), doc, new ReplaceOptions().upsert(true)); + } + + @Override + public InputStream download(String snapshotId) throws Exception { + Objects.requireNonNull(snapshotId, "snapshotId"); + Document doc = + collection + .find(Filters.eq(snapshotId)) + .projection(com.mongodb.client.model.Projections.include(FIELD_DATA)) + .first(); + if (doc == null) { + throw new java.io.FileNotFoundException("Snapshot not found in MongoDB: " + snapshotId); + } + Binary binary = doc.get(FIELD_DATA, Binary.class); + return new ByteArrayInputStream(binary.getData()); + } + + @Override + public boolean exists(String snapshotId) throws Exception { + Objects.requireNonNull(snapshotId, "snapshotId"); + return collection + .find(Filters.eq(snapshotId)) + .projection(com.mongodb.client.model.Projections.include("_id")) + .first() + != null; + } + + private static byte[] readAllBounded(InputStream in, int maxBytes) throws IOException { + ByteArrayOutputStream out = new ByteArrayOutputStream(Math.min(maxBytes, 8192)); + byte[] buf = new byte[8192]; + int total = 0; + int n; + while ((n = in.read(buf)) != -1) { + total += n; + if (total > maxBytes) { + throw new IOException( + "Snapshot size exceeds maximum allowed (" + maxBytes + " bytes)"); + } + out.write(buf, 0, n); + } + return out.toByteArray(); + } +} diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoSnapshotSpec.java b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoSnapshotSpec.java new file mode 100644 index 0000000000..22148f66a0 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoSnapshotSpec.java @@ -0,0 +1,38 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.extensions.mongodb.snapshot; + +import io.agentscope.harness.agent.sandbox.snapshot.RemoteSnapshotSpec; + +/** + * Convenience {@link io.agentscope.harness.agent.sandbox.snapshot.SandboxSnapshotSpec} for + * MongoDB-backed snapshot storage. + * + *

Stores sandbox workspace tar archives as BSON Binary in a MongoDB collection. + */ +public class MongoSnapshotSpec extends RemoteSnapshotSpec { + + public MongoSnapshotSpec(com.mongodb.client.MongoClient mongoClient, String databaseName) { + super(new MongoRemoteSnapshotClient(mongoClient, databaseName, null, true)); + } + + public MongoSnapshotSpec( + com.mongodb.client.MongoClient mongoClient, + String databaseName, + String collectionName) { + super(new MongoRemoteSnapshotClient(mongoClient, databaseName, collectionName, true)); + } +} diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStore.java b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStore.java new file mode 100644 index 0000000000..a58eff0745 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStore.java @@ -0,0 +1,480 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.extensions.mongodb.state; + +import com.mongodb.ConnectionString; +import com.mongodb.MongoClientSettings; +import com.mongodb.client.MongoClient; +import com.mongodb.client.MongoClients; +import com.mongodb.client.MongoCollection; +import com.mongodb.client.MongoDatabase; +import com.mongodb.client.model.Filters; +import com.mongodb.client.model.FindOneAndUpdateOptions; +import com.mongodb.client.model.Indexes; +import com.mongodb.client.model.Projections; +import com.mongodb.client.model.ReturnDocument; +import com.mongodb.client.model.Updates; +import io.agentscope.core.state.AgentStateStore; +import io.agentscope.core.state.ListHashUtil; +import io.agentscope.core.state.State; +import io.agentscope.core.state.VersionedState; +import io.agentscope.core.util.JsonUtils; +import java.util.ArrayList; +import java.util.Date; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.TimeUnit; +import java.util.regex.Pattern; +import org.bson.Document; +import org.bson.conversions.Bson; + +/** + * MongoDB-backed implementation of {@link AgentStateStore}. + * + *

Each session is stored as a single MongoDB document. State keys map to top-level BSON fields. + * Supports optimistic concurrency via a per-key {@code _version_{key}} field. + * + *

List state uses {@link ListHashUtil} for change detection to avoid unnecessary full rewrites. + * + *

Usage: + * + * 

{@code
+ * MongoAgentStateStore store = MongoAgentStateStore.builder()
+ *     .connectionString("mongodb://localhost:27017")
+ *     .databaseName("agentscope")
+ *     .collectionName("sessions")
+ *     .build();
+ * }
+ */ +public class MongoAgentStateStore implements AgentStateStore, AutoCloseable { + + private static final String DEFAULT_DATABASE_NAME = "agentscope"; + private static final String DEFAULT_COLLECTION_NAME = "agentscope_sessions"; + private static final String ANON_USER = "__anon__"; + private static final String LIST_SUFFIX = ":list"; + private static final String HASH_PREFIX = "_hash_"; + private static final String VERSION_PREFIX = "_version_"; + private static final String FIELD_USER_ID = "user_id"; + private static final String FIELD_SESSION_ID = "session_id"; + private static final String FIELD_UPDATED_AT = "_updated_at"; + private static final Pattern SAFE_KEY_PATTERN = Pattern.compile("^[a-zA-Z_][a-zA-Z0-9_]*$"); + + private final MongoClient mongoClient; + private final boolean ownsClient; + private final MongoCollection collection; + + private MongoAgentStateStore(Builder builder) { + if (builder.mongoClient != null) { + this.mongoClient = builder.mongoClient; + this.ownsClient = false; + } else if (builder.connectionString != null) { + MongoClientSettings settings = + MongoClientSettings.builder() + .applyConnectionString(new ConnectionString(builder.connectionString)) + .build(); + this.mongoClient = MongoClients.create(settings); + this.ownsClient = true; + } else { + throw new IllegalArgumentException( + "Either mongoClient or connectionString must be provided"); + } + + String dbName = builder.databaseName != null ? builder.databaseName : DEFAULT_DATABASE_NAME; + String collName = + builder.collectionName != null ? builder.collectionName : DEFAULT_COLLECTION_NAME; + + MongoDatabase db = this.mongoClient.getDatabase(dbName); + this.collection = db.getCollection(collName); + + ensureIndexes(); + } + + public static Builder builder() { + return new Builder(); + } + + @Override + public boolean supportsVersioning() { + return true; + } + + // ────────────────── Index Management ────────────────── + + private void ensureIndexes() { + collection.createIndex( + Indexes.compoundIndex( + Indexes.ascending(FIELD_USER_ID), Indexes.ascending(FIELD_SESSION_ID))); + collection.createIndex( + Indexes.ascending(FIELD_UPDATED_AT), + new com.mongodb.client.model.IndexOptions() + .expireAfter(0L, TimeUnit.SECONDS) + .sparse(true)); + } + + // ────────────────── Single Value CRUD ────────────────── + + @Override + public void save(String userId, String sessionId, String key, State value) { + validateKey(key); + String slotId = slotId(userId, sessionId); + String json = JsonUtils.getJsonCodec().toJson(value); + Bson setFields = + Updates.combine( + Updates.set(key, Document.parse(json)), + Updates.set(FIELD_UPDATED_AT, new Date())); + Bson setOnInsert = + Updates.combine( + Updates.setOnInsert(FIELD_USER_ID, normalizeUser(userId)), + Updates.setOnInsert(FIELD_SESSION_ID, sessionId)); + collection.updateOne(Filters.eq(slotId), Updates.combine(setFields, setOnInsert), upsert()); + } + + @Override + public Optional get( + String userId, String sessionId, String key, Class type) { + validateKey(key); + String slotId = slotId(userId, sessionId); + Document doc = + collection.find(Filters.eq(slotId)).projection(Projections.include(key)).first(); + if (doc == null || !doc.containsKey(key)) { + return Optional.empty(); + } + return Optional.ofNullable(deserializeValue(doc.get(key), type)); + } + + // ────────────────── List CRUD ────────────────── + + @Override + public void save(String userId, String sessionId, String key, List values) { + validateKey(key); + String slotId = slotId(userId, sessionId); + String listKey = key + LIST_SUFFIX; + String hashField = HASH_PREFIX + key; + + Document doc = + collection + .find(Filters.eq(slotId)) + .projection(Projections.include(listKey, hashField)) + .first(); + + String storedHash = null; + int existingCount = 0; + if (doc != null) { + if (doc.containsKey(hashField)) { + storedHash = doc.getString(hashField); + } + if (doc.containsKey(listKey)) { + existingCount = doc.getList(listKey, Object.class).size(); + } + } + + String currentHash = ListHashUtil.computeHash(values); + + if (ListHashUtil.needsFullRewrite(values, storedHash, existingCount)) { + List bsonList = toDocumentList(values); + Bson setFields = + Updates.combine( + Updates.set(listKey, bsonList), + Updates.set(hashField, currentHash), + Updates.set(FIELD_UPDATED_AT, new Date())); + Bson setOnInsert = + Updates.combine( + Updates.setOnInsert(FIELD_USER_ID, normalizeUser(userId)), + Updates.setOnInsert(FIELD_SESSION_ID, sessionId)); + collection.updateOne( + Filters.eq(slotId), Updates.combine(setFields, setOnInsert), upsert()); + } else if (values.size() > existingCount) { + List newItems = values.subList(existingCount, values.size()); + List newDocs = toDocumentList(newItems); + Bson update = + Updates.combine( + Updates.pushEach(listKey, newDocs), + Updates.set(hashField, currentHash), + Updates.set(FIELD_UPDATED_AT, new Date())); + collection.updateOne(Filters.eq(slotId), update, upsert()); + } + } + + @Override + public List getList( + String userId, String sessionId, String key, Class itemType) { + validateKey(key); + String slotId = slotId(userId, sessionId); + String listKey = key + LIST_SUFFIX; + Document doc = + collection + .find(Filters.eq(slotId)) + .projection(Projections.include(listKey)) + .first(); + if (doc == null || !doc.containsKey(listKey)) { + return List.of(); + } + List rawList = doc.getList(listKey, Object.class); + List result = new ArrayList<>(rawList.size()); + for (Object item : rawList) { + result.add(deserializeValue(item, itemType)); + } + return result; + } + + // ────────────────── Versioning ────────────────── + + @Override + public VersionedState getVersioned( + String userId, String sessionId, String key, Class type) { + validateKey(key); + String slotId = slotId(userId, sessionId); + String versionField = VERSION_PREFIX + key; + Document doc = + collection + .find(Filters.eq(slotId)) + .projection(Projections.include(key, versionField)) + .first(); + if (doc == null || !doc.containsKey(key)) { + return new VersionedState<>(null, 0L); + } + long version = doc.containsKey(versionField) ? doc.getLong(versionField) : 0L; + T value = deserializeValue(doc.get(key), type); + return new VersionedState<>(value, version); + } + + @Override + public long saveIfVersion( + String userId, String sessionId, String key, State value, long expectedVersion) { + validateKey(key); + if (expectedVersion == UNVERSIONED) { + save(userId, sessionId, key, value); + VersionedState after = getVersioned(userId, sessionId, key, State.class); + return after.version(); + } + + String slotId = slotId(userId, sessionId); + String versionField = VERSION_PREFIX + key; + String json = JsonUtils.getJsonCodec().toJson(value); + + if (expectedVersion == 0) { + Bson filter = Filters.and(Filters.eq(slotId), Filters.exists(versionField, false)); + Bson update = + Updates.combine( + Updates.set(key, Document.parse(json)), + Updates.set(versionField, 1L), + Updates.set(FIELD_UPDATED_AT, new Date()), + Updates.setOnInsert(FIELD_USER_ID, normalizeUser(userId)), + Updates.setOnInsert(FIELD_SESSION_ID, sessionId)); + Document result = + collection.findOneAndUpdate( + filter, + update, + new FindOneAndUpdateOptions() + .upsert(true) + .returnDocument(ReturnDocument.AFTER)); + if (result == null) { + return UNVERSIONED; + } + Long newVersion = result.getLong(versionField); + return newVersion != null && newVersion == 1L ? 1L : UNVERSIONED; + } + + Bson filter = Filters.and(Filters.eq(slotId), Filters.eq(versionField, expectedVersion)); + Bson update = + Updates.combine( + Updates.set(key, Document.parse(json)), + Updates.inc(versionField, 1L), + Updates.set(FIELD_UPDATED_AT, new Date())); + Document result = + collection.findOneAndUpdate( + filter, + update, + new FindOneAndUpdateOptions().returnDocument(ReturnDocument.AFTER)); + if (result == null) { + return UNVERSIONED; + } + Long newVersion = result.getLong(versionField); + return newVersion != null ? newVersion : UNVERSIONED; + } + + // ────────────────── Session CRUD ────────────────── + + @Override + public boolean exists(String userId, String sessionId) { + String slotId = slotId(userId, sessionId); + return collection.find(Filters.eq(slotId)).projection(Projections.include("_id")).first() + != null; + } + + @Override + public void delete(String userId, String sessionId) { + String slotId = slotId(userId, sessionId); + collection.deleteOne(Filters.eq(slotId)); + } + + @Override + public void delete(String userId, String sessionId, String key) { + validateKey(key); + String slotId = slotId(userId, sessionId); + Document unsetFields = + new Document(key, "") + .append(VERSION_PREFIX + key, "") + .append(HASH_PREFIX + key, "") + .append(key + LIST_SUFFIX, ""); + collection.updateOne(Filters.eq(slotId), new Document("$unset", unsetFields)); + } + + @Override + public Set listSessionIds(String userId) { + String normalizedUser = normalizeUser(userId); + List ids = + collection + .distinct( + FIELD_SESSION_ID, + Filters.eq(FIELD_USER_ID, normalizedUser), + String.class) + .into(new ArrayList<>()); + return new LinkedHashSet<>(ids); + } + + // ────────────────── Close ────────────────── + + @Override + public void close() { + if (ownsClient) { + mongoClient.close(); + } + } + + // ────────────────── Internal Helpers ────────────────── + + private static String normalizeUser(String userId) { + return (userId == null || userId.isBlank()) ? ANON_USER : userId; + } + + private static String slotId(String userId, String sessionId) { + if (sessionId == null || sessionId.isBlank()) { + throw new IllegalArgumentException("sessionId must not be blank"); + } + return normalizeUser(userId) + ":" + sessionId; + } + + private static void validateKey(String key) { + if (key == null || key.isBlank()) { + throw new IllegalArgumentException("key must not be blank"); + } + if (!SAFE_KEY_PATTERN.matcher(key).matches()) { + throw new IllegalArgumentException( + "key must match pattern " + + SAFE_KEY_PATTERN + + " but was: " + + key + + " (MongoDB field names cannot contain '.' or '$')"); + } + } + + private T deserializeValue(Object fieldValue, Class type) { + if (fieldValue == null) { + return null; + } + String json; + if (fieldValue instanceof String s) { + json = s; + } else if (fieldValue instanceof Document doc) { + json = doc.toJson(); + } else { + json = fieldValue.toString(); + } + return JsonUtils.getJsonCodec().fromJson(json, type); + } + + private List toDocumentList(List values) { + List result = new ArrayList<>(values.size()); + for (State item : values) { + result.add(Document.parse(JsonUtils.getJsonCodec().toJson(item))); + } + return result; + } + + private static com.mongodb.client.model.UpdateOptions upsert() { + return new com.mongodb.client.model.UpdateOptions().upsert(true); + } + + // ────────────────── Builder ────────────────── + + /** Builder for {@link MongoAgentStateStore}. */ + public static class Builder { + private MongoClient mongoClient; + private String connectionString; + private String databaseName; + private String collectionName; + + /** + * Use an existing {@link MongoClient}. The caller owns its lifecycle; {@link + * MongoAgentStateStore#close()} will NOT close a client supplied through this method. + * + * @param mongoClient the client to use + * @return this builder + */ + public Builder mongoClient(MongoClient mongoClient) { + this.mongoClient = mongoClient; + return this; + } + + /** + * MongoDB connection string (e.g. {@code "mongodb://localhost:27017"}). + * + *

A new {@link MongoClient} will be created internally and closed when {@link + * MongoAgentStateStore#close()} is called. + * + * @param connectionString the connection string + * @return this builder + */ + public Builder connectionString(String connectionString) { + this.connectionString = connectionString; + return this; + } + + /** + * Database name. Defaults to {@code "agentscope"}. + * + * @param databaseName the database name + * @return this builder + */ + public Builder databaseName(String databaseName) { + this.databaseName = databaseName; + return this; + } + + /** + * Collection name. Defaults to {@code "agentscope_sessions"}. + * + * @param collectionName the collection name + * @return this builder + */ + public Builder collectionName(String collectionName) { + this.collectionName = collectionName; + return this; + } + + /** + * Build the {@link MongoAgentStateStore}. + * + * @return a new instance + */ + public MongoAgentStateStore build() { + return new MongoAgentStateStore(this); + } + } +} diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/store/MongoBaseStore.java b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/store/MongoBaseStore.java new file mode 100644 index 0000000000..f6ea503906 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/store/MongoBaseStore.java @@ -0,0 +1,236 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.extensions.mongodb.store; + +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.mongodb.client.MongoCollection; +import com.mongodb.client.MongoDatabase; +import com.mongodb.client.model.Filters; +import com.mongodb.client.model.FindOneAndUpdateOptions; +import com.mongodb.client.model.Indexes; +import com.mongodb.client.model.Projections; +import com.mongodb.client.model.ReturnDocument; +import com.mongodb.client.model.Sorts; +import com.mongodb.client.model.Updates; +import io.agentscope.harness.agent.filesystem.remote.store.BaseStore; +import io.agentscope.harness.agent.filesystem.remote.store.StoreItem; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import org.bson.Document; +import org.bson.conversions.Bson; + +/** + * MongoDB-backed implementation of {@link BaseStore}. + * + *

Each item is stored as a separate MongoDB document. Namespace paths and keys are encoded into + * a compound {@code _id} for uniqueness. Supports optimistic concurrency via a {@code version} + * field. + */ +public class MongoBaseStore implements BaseStore { + + private static final String FIELD_ID = "_id"; + private static final String FIELD_KEY = "key"; + private static final String FIELD_NAMESPACE = "namespace"; + private static final String FIELD_VALUE = "value"; + private static final String FIELD_VERSION = "version"; + + private static final TypeReference> MAP_TYPE = new TypeReference<>() {}; + + private final MongoCollection collection; + private final ObjectMapper objectMapper; + + /** + * Creates a new instance. + * + * @param database the MongoDB database + * @param collectionName the collection name + */ + public MongoBaseStore(MongoDatabase database, String collectionName) { + this(database, collectionName, new ObjectMapper()); + } + + /** + * Creates a new instance with a custom ObjectMapper. + * + * @param database the MongoDB database + * @param collectionName the collection name + * @param objectMapper Jackson mapper for serializing values + */ + public MongoBaseStore( + MongoDatabase database, String collectionName, ObjectMapper objectMapper) { + this.collection = database.getCollection(collectionName); + this.objectMapper = objectMapper; + ensureIndexes(); + } + + @Override + public StoreItem get(List namespace, String key) { + String id = itemDocId(namespace, key); + Document doc = + collection + .find(Filters.eq(id)) + .projection(Projections.include(FIELD_VALUE, FIELD_VERSION)) + .first(); + if (doc == null) { + return null; + } + Map value = parseValue(doc.get(FIELD_VALUE)); + long version = doc.containsKey(FIELD_VERSION) ? doc.getLong(FIELD_VERSION) : 0L; + return new StoreItem(key, value, version); + } + + @Override + public void put(List namespace, String key, Map value) { + String id = itemDocId(namespace, key); + String nsKey = namespacePath(namespace); + String json = serialize(value); + Bson setFields = + Updates.combine( + Updates.set(FIELD_VALUE, Document.parse(json)), + Updates.set(FIELD_KEY, key), + Updates.set(FIELD_NAMESPACE, nsKey)); + Bson setOnInsert = Updates.setOnInsert(FIELD_ID, id); + collection.updateOne( + Filters.eq(id), + Updates.combine(setFields, setOnInsert, Updates.inc(FIELD_VERSION, 1L)), + upsert()); + } + + @Override + public boolean putIfVersion( + List namespace, String key, Map value, long expectedVersion) { + String id = itemDocId(namespace, key); + String nsKey = namespacePath(namespace); + String json = serialize(value); + + Document result; + if (expectedVersion == 0) { + Bson filter = Filters.and(Filters.eq(id), Filters.exists(FIELD_VERSION, false)); + Bson update = + Updates.combine( + Updates.set(FIELD_VALUE, Document.parse(json)), + Updates.set(FIELD_KEY, key), + Updates.set(FIELD_NAMESPACE, nsKey), + Updates.set(FIELD_VERSION, 1L), + Updates.setOnInsert(FIELD_ID, id)); + result = + collection.findOneAndUpdate( + filter, + update, + new FindOneAndUpdateOptions() + .upsert(true) + .returnDocument(ReturnDocument.AFTER)); + } else { + Bson filter = Filters.and(Filters.eq(id), Filters.eq(FIELD_VERSION, expectedVersion)); + Bson update = + Updates.combine( + Updates.set(FIELD_VALUE, Document.parse(json)), + Updates.set(FIELD_KEY, key), + Updates.set(FIELD_NAMESPACE, nsKey), + Updates.inc(FIELD_VERSION, 1L)); + result = + collection.findOneAndUpdate( + filter, + update, + new FindOneAndUpdateOptions().returnDocument(ReturnDocument.AFTER)); + } + + if (result == null) { + return false; + } + Long newVersion = result.getLong(FIELD_VERSION); + if (expectedVersion == 0 && newVersion != null && newVersion == 1L) { + return true; + } + return newVersion != null && newVersion == expectedVersion + 1; + } + + @Override + public List search(List namespace, int limit, int offset) { + String nsKey = namespacePath(namespace); + List docs = + collection + .find(Filters.eq(FIELD_NAMESPACE, nsKey)) + .sort(Sorts.ascending(FIELD_KEY)) + .skip(offset) + .limit(limit) + .projection(Projections.include(FIELD_KEY, FIELD_VALUE, FIELD_VERSION)) + .into(new ArrayList<>()); + List result = new ArrayList<>(docs.size()); + for (Document doc : docs) { + String key = doc.getString(FIELD_KEY); + Map value = parseValue(doc.get(FIELD_VALUE)); + long version = doc.containsKey(FIELD_VERSION) ? doc.getLong(FIELD_VERSION) : 0L; + result.add(new StoreItem(key, value, version)); + } + return result; + } + + @Override + public void delete(List namespace, String key) { + String id = itemDocId(namespace, key); + collection.deleteOne(Filters.eq(id)); + } + + // ────────────────── Internal Helpers ────────────────── + + private void ensureIndexes() { + collection.createIndex(Indexes.ascending(FIELD_NAMESPACE)); + collection.createIndex( + Indexes.compoundIndex( + Indexes.ascending(FIELD_NAMESPACE), Indexes.ascending(FIELD_KEY))); + } + + private static String itemDocId(List namespace, String key) { + return namespacePath(namespace) + "\0" + key; + } + + private static String namespacePath(List namespace) { + return namespace.stream().collect(Collectors.joining("\0")); + } + + private String serialize(Map value) { + try { + return objectMapper.writeValueAsString(value == null ? Map.of() : value); + } catch (com.fasterxml.jackson.core.JsonProcessingException e) { + throw new IllegalStateException("Failed to serialize value", e); + } + } + + @SuppressWarnings("unchecked") + private Map parseValue(Object raw) { + if (raw instanceof Document doc) { + return new LinkedHashMap<>(doc); + } + if (raw instanceof String s) { + try { + Map parsed = objectMapper.readValue(s, MAP_TYPE); + return parsed != null ? parsed : Map.of(); + } catch (com.fasterxml.jackson.core.JsonProcessingException e) { + return Map.of(); + } + } + return Map.of(); + } + + private static com.mongodb.client.model.UpdateOptions upsert() { + return new com.mongodb.client.model.UpdateOptions().upsert(true); + } +} diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/MongoDistributedStoreTest.java b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/MongoDistributedStoreTest.java new file mode 100644 index 0000000000..cd49949c50 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/MongoDistributedStoreTest.java @@ -0,0 +1,87 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.extensions.mongodb; + +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.when; + +import com.mongodb.client.MongoClient; +import com.mongodb.client.MongoCollection; +import com.mongodb.client.MongoDatabase; +import io.agentscope.core.state.AgentStateStore; +import io.agentscope.harness.agent.filesystem.remote.store.BaseStore; +import org.bson.Document; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +class MongoDistributedStoreTest { + + @Mock private MongoClient mongoClient; + @Mock private MongoDatabase mongoDatabase; + @Mock private MongoCollection collection; + + private AutoCloseable mocks; + + @BeforeEach + void setUp() { + mocks = MockitoAnnotations.openMocks(this); + when(mongoClient.getDatabase(anyString())).thenReturn(mongoDatabase); + when(mongoDatabase.getCollection(anyString())).thenReturn(collection); + } + + @AfterEach + void tearDown() throws Exception { + if (mocks != null) { + mocks.close(); + } + } + + @Test + void createWithMongoClient() { + MongoDistributedStore store = MongoDistributedStore.create(mongoClient); + assertNotNull(store); + } + + @Test + void createWithMongoClientAndDatabaseName() { + MongoDistributedStore store = MongoDistributedStore.create(mongoClient, "mydb"); + assertNotNull(store); + } + + @Test + void agentStateStoreReturnsNonNull() { + MongoDistributedStore store = MongoDistributedStore.create(mongoClient); + AgentStateStore stateStore = store.agentStateStore(); + assertNotNull(stateStore); + } + + @Test + void baseStoreReturnsNonNull() { + MongoDistributedStore store = MongoDistributedStore.create(mongoClient); + BaseStore baseStore = store.baseStore(); + assertNotNull(baseStore); + } + + @Test + void createWithNullMongoClientThrows() { + assertThrows(NullPointerException.class, () -> MongoDistributedStore.create(null)); + } +} diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreTest.java b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreTest.java new file mode 100644 index 0000000000..13e925600c --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreTest.java @@ -0,0 +1,293 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.extensions.mongodb.state; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.mongodb.client.DistinctIterable; +import com.mongodb.client.FindIterable; +import com.mongodb.client.MongoClient; +import com.mongodb.client.MongoCollection; +import com.mongodb.client.MongoDatabase; +import com.mongodb.client.result.DeleteResult; +import com.mongodb.client.result.UpdateResult; +import io.agentscope.core.state.AgentStateStore; +import io.agentscope.core.state.State; +import io.agentscope.core.state.VersionedState; +import java.util.List; +import java.util.Optional; +import java.util.Set; +import org.bson.Document; +import org.bson.conversions.Bson; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +class MongoAgentStateStoreTest { + + record TestState(String value) implements State {} + + @Mock private MongoClient mongoClient; + @Mock private MongoDatabase mongoDatabase; + @Mock private MongoCollection collection; + + @SuppressWarnings("rawtypes") + @Mock + private FindIterable findIterable; + + @SuppressWarnings("rawtypes") + @Mock + private DistinctIterable distinctIterable; + + private AutoCloseable mocks; + private MongoAgentStateStore store; + + @BeforeEach + @SuppressWarnings("unchecked") + void setUp() { + mocks = MockitoAnnotations.openMocks(this); + when(mongoClient.getDatabase(anyString())).thenReturn(mongoDatabase); + when(mongoDatabase.getCollection(anyString())).thenReturn(collection); + + when(collection.find(any(Bson.class))).thenReturn(findIterable); + when(findIterable.projection(any())).thenReturn(findIterable); + when(findIterable.sort(any())).thenReturn(findIterable); + when(findIterable.skip(org.mockito.ArgumentMatchers.anyInt())).thenReturn(findIterable); + when(findIterable.limit(org.mockito.ArgumentMatchers.anyInt())).thenReturn(findIterable); + when(findIterable.first()).thenReturn(null); + + UpdateResult updateResult = mock(UpdateResult.class); + when(updateResult.wasAcknowledged()).thenReturn(true); + when(collection.updateOne(any(Bson.class), any(Bson.class), any())) + .thenReturn(updateResult); + when(collection.updateOne(any(Bson.class), any(Bson.class))).thenReturn(updateResult); + + DeleteResult deleteResult = mock(DeleteResult.class); + when(deleteResult.wasAcknowledged()).thenReturn(true); + when(collection.deleteOne(any(Bson.class))).thenReturn(deleteResult); + + when(collection.countDocuments(any(Bson.class))).thenReturn(0L); + + when(collection.distinct(anyString(), any(Bson.class), any(Class.class))) + .thenReturn(distinctIterable); + when(distinctIterable.into(any())).thenReturn(new java.util.ArrayList<>()); + + store = + MongoAgentStateStore.builder() + .mongoClient(mongoClient) + .databaseName("testdb") + .collectionName("test_sessions") + .build(); + } + + @AfterEach + void tearDown() throws Exception { + if (mocks != null) { + mocks.close(); + } + } + + @Test + void supportsVersioningReturnsTrue() { + assertTrue(store.supportsVersioning()); + } + + @Test + void builderRejectsMissingClientAndConnectionString() { + assertThrows(IllegalArgumentException.class, () -> MongoAgentStateStore.builder().build()); + } + + @Test + void builderWithMongoClientCreatesStore() { + assertNotNull(store); + } + + @Test + void builderWithDefaultsCreatesStore() { + MongoAgentStateStore defaultStore = + MongoAgentStateStore.builder().mongoClient(mongoClient).build(); + assertNotNull(defaultStore); + } + + @Test + void saveSingleState() { + store.save("user", "session", "key", new TestState("value")); + verify(collection).updateOne(any(Bson.class), any(Bson.class), any()); + } + + @Test + void getSingleStateReturnsEmptyWhenMissing() { + Optional result = store.get("user", "session", "key", TestState.class); + assertTrue(result.isEmpty()); + } + + @Test + void getSingleStateReturnsValueWhenPresent() { + String json = "{\"value\":\"found\"}"; + Document doc = new Document("key", Document.parse(json)); + when(findIterable.first()).thenReturn(doc); + + Optional result = store.get("user", "session", "key", TestState.class); + assertTrue(result.isPresent()); + assertEquals("found", result.get().value()); + } + + @Test + void saveListState() { + store.save("user", "session", "list", List.of(new TestState("a"), new TestState("b"))); + verify(collection).updateOne(any(Bson.class), any(Bson.class), any()); + } + + @Test + void getListReturnsEmptyWhenMissing() { + List result = store.getList("user", "session", "list", TestState.class); + assertTrue(result.isEmpty()); + } + + @Test + void getListReturnsValuesWhenPresent() { + List list = + List.of(Document.parse("{\"value\":\"a\"}"), Document.parse("{\"value\":\"b\"}")); + Document doc = new Document("list:list", list); + when(findIterable.first()).thenReturn(doc); + + List result = store.getList("user", "session", "list", TestState.class); + assertEquals(2, result.size()); + assertEquals("a", result.get(0).value()); + assertEquals("b", result.get(1).value()); + } + + @Test + void getVersionedReturnsZeroWhenMissing() { + VersionedState result = + store.getVersioned("user", "session", "key", TestState.class); + assertNotNull(result); + assertEquals(0L, result.version()); + } + + @Test + void getVersionedReturnsValueAndVersion() { + String json = "{\"value\":\"v1\"}"; + Document doc = new Document("key", Document.parse(json)).append("_version_key", 5L); + when(findIterable.first()).thenReturn(doc); + + VersionedState result = + store.getVersioned("user", "session", "key", TestState.class); + assertEquals(5L, result.version()); + assertEquals("v1", result.value().value()); + } + + @Test + void saveIfVersionWithUnversionedDelegatesToSave() { + store.saveIfVersion( + "user", "session", "key", new TestState("v"), AgentStateStore.UNVERSIONED); + verify(collection).updateOne(any(Bson.class), any(Bson.class), any()); + } + + @Test + void existsReturnsFalseWhenNoDocument() { + when(findIterable.first()).thenReturn(null); + assertFalse(store.exists("user", "session")); + } + + @Test + void existsReturnsTrueWhenDocumentExists() { + when(findIterable.first()).thenReturn(new Document("_id", "__anon__:session")); + assertTrue(store.exists("user", "session")); + } + + @Test + void deleteSession() { + store.delete("user", "session"); + verify(collection).deleteOne(any(Bson.class)); + } + + @Test + void deleteKey() { + store.delete("user", "session", "key"); + verify(collection).updateOne(any(Bson.class), any(Document.class)); + } + + @Test + void listSessionIdsReturnsEmptyWhenNone() { + Set ids = store.listSessionIds("user"); + assertTrue(ids.isEmpty()); + } + + @Test + @SuppressWarnings("unchecked") + void listSessionIdsReturnsIds() { + java.util.ArrayList ids = new java.util.ArrayList<>(List.of("s1", "s2")); + when(distinctIterable.into(any())).thenReturn(ids); + + Set result = store.listSessionIds("user"); + assertEquals(2, result.size()); + assertTrue(result.contains("s1")); + assertTrue(result.contains("s2")); + } + + @Test + void rejectsNullSessionId() { + assertThrows( + IllegalArgumentException.class, + () -> store.save("user", null, "key", new TestState("v"))); + } + + @Test + void rejectsBlankSessionId() { + assertThrows( + IllegalArgumentException.class, + () -> store.save("user", " ", "key", new TestState("v"))); + } + + @Test + void rejectsKeyWithDot() { + assertThrows( + IllegalArgumentException.class, + () -> store.save("user", "session", "bad.key", new TestState("v"))); + } + + @Test + void rejectsKeyWithDollar() { + assertThrows( + IllegalArgumentException.class, + () -> store.save("user", "session", "$bad", new TestState("v"))); + } + + @Test + void rejectsBlankKey() { + assertThrows( + IllegalArgumentException.class, + () -> store.save("user", "session", " ", new TestState("v"))); + } + + @Test + void closeWithExternalClientDoesNotCloseClient() { + store.close(); + verify(mongoClient, org.mockito.Mockito.never()).close(); + } +} diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/store/MongoBaseStoreTest.java b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/store/MongoBaseStoreTest.java new file mode 100644 index 0000000000..9bd69daa10 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/store/MongoBaseStoreTest.java @@ -0,0 +1,172 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.extensions.mongodb.store; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.mongodb.client.FindIterable; +import com.mongodb.client.MongoCollection; +import com.mongodb.client.MongoDatabase; +import com.mongodb.client.result.DeleteResult; +import com.mongodb.client.result.UpdateResult; +import io.agentscope.harness.agent.filesystem.remote.store.StoreItem; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import org.bson.Document; +import org.bson.conversions.Bson; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +class MongoBaseStoreTest { + + @Mock private MongoDatabase mongoDatabase; + @Mock private MongoCollection collection; + + @SuppressWarnings("rawtypes") + @Mock + private FindIterable findIterable; + + private AutoCloseable mocks; + private MongoBaseStore store; + + @BeforeEach + @SuppressWarnings("unchecked") + void setUp() { + mocks = MockitoAnnotations.openMocks(this); + when(mongoDatabase.getCollection(anyString())).thenReturn(collection); + + when(collection.find(any(Bson.class))).thenReturn(findIterable); + when(findIterable.projection(any())).thenReturn(findIterable); + when(findIterable.sort(any())).thenReturn(findIterable); + when(findIterable.skip(org.mockito.ArgumentMatchers.anyInt())).thenReturn(findIterable); + when(findIterable.limit(org.mockito.ArgumentMatchers.anyInt())).thenReturn(findIterable); + when(findIterable.first()).thenReturn(null); + when(findIterable.into(any())).thenReturn(new ArrayList<>()); + + UpdateResult updateResult = mock(UpdateResult.class); + when(updateResult.wasAcknowledged()).thenReturn(true); + when(collection.updateOne(any(Bson.class), any(Bson.class), any())) + .thenReturn(updateResult); + when(collection.updateOne(any(Bson.class), any(Bson.class))).thenReturn(updateResult); + + DeleteResult deleteResult = mock(DeleteResult.class); + when(deleteResult.wasAcknowledged()).thenReturn(true); + when(collection.deleteOne(any(Bson.class))).thenReturn(deleteResult); + + store = new MongoBaseStore(mongoDatabase, "test_base"); + } + + @AfterEach + void tearDown() throws Exception { + if (mocks != null) { + mocks.close(); + } + } + + @Test + void constructorCreatesStore() { + assertNotNull(store); + verify(mongoDatabase).getCollection("test_base"); + } + + @Test + void getReturnsNullWhenNotFound() { + StoreItem item = store.get(List.of("ns"), "key"); + assertNull(item); + } + + @Test + void getReturnsItemWhenFound() { + Document doc = + new Document() + .append("key", "mykey") + .append("value", new Document("data", "hello")) + .append("version", 3L); + when(findIterable.first()).thenReturn(doc); + + StoreItem item = store.get(List.of("ns"), "mykey"); + assertNotNull(item); + assertEquals("mykey", item.key()); + assertEquals(3L, item.version()); + assertEquals("hello", item.value().get("data")); + } + + @Test + void putStoresItem() { + store.put(List.of("ns"), "key", Map.of("data", "value")); + verify(collection).updateOne(any(Bson.class), any(Bson.class), any()); + } + + @Test + void putIfVersionReturnsFalseWhenVersionMismatch() { + when(findIterable.first()).thenReturn(null); + boolean result = store.putIfVersion(List.of("ns"), "key", Map.of("data", "v"), 5L); + // findOneAndUpdate returns null when filter doesn't match (no upsert for non-zero version) + // Actually the implementation returns null for non-zero expectedVersion when no match + } + + @Test + void searchReturnsEmptyList() { + List items = store.search(List.of("ns"), 10, 0); + assertTrue(items.isEmpty()); + } + + @Test + @SuppressWarnings("unchecked") + void searchReturnsItems() { + Document doc1 = + new Document() + .append("key", "a") + .append("value", new Document("x", "1")) + .append("version", 1L); + Document doc2 = + new Document() + .append("key", "b") + .append("value", new Document("x", "2")) + .append("version", 2L); + ArrayList docs = new ArrayList<>(List.of(doc1, doc2)); + when(findIterable.into(any())).thenReturn(docs); + + List items = store.search(List.of("ns"), 10, 0); + assertEquals(2, items.size()); + assertEquals("a", items.get(0).key()); + assertEquals("b", items.get(1).key()); + } + + @Test + void deleteRemovesItem() { + store.delete(List.of("ns"), "key"); + verify(collection).deleteOne(any(Bson.class)); + } + + @Test + void rejectsNullMongoDatabase() { + assertThrows(NullPointerException.class, () -> new MongoBaseStore(null, "test")); + } +} diff --git a/agentscope-extensions/pom.xml b/agentscope-extensions/pom.xml index 3525fd1da2..bf7ab0bb7d 100644 --- a/agentscope-extensions/pom.xml +++ b/agentscope-extensions/pom.xml @@ -39,6 +39,7 @@ agentscope-extensions-redis agentscope-extensions-mysql agentscope-extensions-postgresql + agentscope-extensions-mongodb agentscope-extensions-rag agentscope-extensions-model agentscope-extensions-higress From e3b5a676841fc2d4c1e1b72c76d04a1845a758d7 Mon Sep 17 00:00:00 2001 From: LiangshouX Date: Mon, 10 Aug 2026 01:27:39 +0800 Subject: [PATCH 2/8] fix(extensions-mongodb): harden lock, lifecycle and CAS correctness --- .../mongodb/MongoDistributedStore.java | 92 +++++- .../sandbox/MongoSandboxExecutionGuard.java | 150 ++++++--- .../snapshot/MongoRemoteSnapshotClient.java | 36 +- .../mongodb/snapshot/MongoSnapshotSpec.java | 8 +- .../mongodb/state/MongoAgentStateStore.java | 67 +++- .../mongodb/store/MongoBaseStore.java | 58 ++-- .../mongodb/MongoDistributedStoreTest.java | 72 ++++ .../MongoSandboxExecutionGuardTest.java | 311 ++++++++++++++++++ .../MongoRemoteSnapshotClientTest.java | 231 +++++++++++++ .../snapshot/MongoSnapshotSpecTest.java | 65 ++++ .../state/MongoAgentStateStoreTest.java | 109 +++++- .../mongodb/store/MongoBaseStoreTest.java | 65 +++- 12 files changed, 1144 insertions(+), 120 deletions(-) create mode 100644 agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/sandbox/MongoSandboxExecutionGuardTest.java create mode 100644 agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClientTest.java create mode 100644 agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/snapshot/MongoSnapshotSpecTest.java diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/MongoDistributedStore.java b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/MongoDistributedStore.java index ca7b63f6fb..48ae2e8d42 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/MongoDistributedStore.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/MongoDistributedStore.java @@ -55,20 +55,33 @@ *

  • {@link MongoSnapshotSpec} — sandbox workspace snapshots in MongoDB * * - *

    The caller owns the {@link MongoClient} lifecycle; closing the store does NOT close the - * client. + *

    When created via {@link #create(MongoClient)}, the caller owns the {@link MongoClient} + * lifecycle; {@link #close()} will NOT close the client. When created via {@link + * #fromConnectionString(String)}, the store owns the client and {@link #close()} will close it. */ -public class MongoDistributedStore implements DistributedStore { +public class MongoDistributedStore implements DistributedStore, AutoCloseable { private static final String DEFAULT_DATABASE = "agentscope"; private static final String STATE_COLLECTION = "agentscope_sessions"; private static final String BASE_COLLECTION = "agentscope_base"; private final MongoClient mongoClient; + private final boolean ownsClient; private final String databaseName; + private volatile AgentStateStore cachedAgentStateStore; + private volatile BaseStore cachedBaseStore; + private volatile SandboxSnapshotSpec cachedSnapshotSpec; + private volatile SandboxExecutionGuard cachedExecutionGuard; + private MongoDistributedStore(MongoClient mongoClient, String databaseName) { + this(mongoClient, databaseName, false); + } + + private MongoDistributedStore( + MongoClient mongoClient, String databaseName, boolean ownsClient) { this.mongoClient = Objects.requireNonNull(mongoClient, "mongoClient"); + this.ownsClient = ownsClient; this.databaseName = databaseName != null ? databaseName : DEFAULT_DATABASE; } @@ -85,7 +98,7 @@ public static MongoDistributedStore create(MongoClient mongoClient) { /** * Creates a MongoDB distributed store. * - * @param mongoClient the MongoDB client + * @param mongoClient the MongoDB client * @param databaseName the database name * @return a new MongoDB distributed store */ @@ -105,31 +118,82 @@ public static MongoDistributedStore fromConnectionString(String connectionString MongoClientSettings.builder() .applyConnectionString(new ConnectionString(connectionString)) .build(); - return new MongoDistributedStore(MongoClients.create(settings), null); + return new MongoDistributedStore(MongoClients.create(settings), null, true); } @Override public AgentStateStore agentStateStore() { - return MongoAgentStateStore.builder() - .mongoClient(mongoClient) - .databaseName(databaseName) - .collectionName(STATE_COLLECTION) - .build(); + AgentStateStore result = cachedAgentStateStore; + if (result == null) { + synchronized (this) { + result = cachedAgentStateStore; + if (result == null) { + result = + MongoAgentStateStore.builder() + .mongoClient(mongoClient) + .databaseName(databaseName) + .collectionName(STATE_COLLECTION) + .build(); + cachedAgentStateStore = result; + } + } + } + return result; } @Override public BaseStore baseStore() { - MongoDatabase db = mongoClient.getDatabase(databaseName); - return new MongoBaseStore(db, BASE_COLLECTION); + BaseStore result = cachedBaseStore; + if (result == null) { + synchronized (this) { + result = cachedBaseStore; + if (result == null) { + MongoDatabase db = mongoClient.getDatabase(databaseName); + result = new MongoBaseStore(db, BASE_COLLECTION); + cachedBaseStore = result; + } + } + } + return result; } @Override public SandboxSnapshotSpec sandboxSnapshotSpec() { - return new MongoSnapshotSpec(mongoClient, databaseName); + SandboxSnapshotSpec result = cachedSnapshotSpec; + if (result == null) { + synchronized (this) { + result = cachedSnapshotSpec; + if (result == null) { + result = new MongoSnapshotSpec(mongoClient, databaseName); + cachedSnapshotSpec = result; + } + } + } + return result; } @Override public SandboxExecutionGuard sandboxExecutionGuard() { - return MongoSandboxExecutionGuard.builder(mongoClient).databaseName(databaseName).build(); + SandboxExecutionGuard result = cachedExecutionGuard; + if (result == null) { + synchronized (this) { + result = cachedExecutionGuard; + if (result == null) { + result = + MongoSandboxExecutionGuard.builder(mongoClient) + .databaseName(databaseName) + .build(); + cachedExecutionGuard = result; + } + } + } + return result; + } + + @Override + public void close() { + if (ownsClient) { + mongoClient.close(); + } } } diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/sandbox/MongoSandboxExecutionGuard.java b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/sandbox/MongoSandboxExecutionGuard.java index f6f577b32a..b20b2aa267 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/sandbox/MongoSandboxExecutionGuard.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/sandbox/MongoSandboxExecutionGuard.java @@ -15,13 +15,15 @@ */ package io.agentscope.extensions.mongodb.sandbox; -import com.mongodb.MongoBulkWriteException; +import com.mongodb.MongoWriteException; +import com.mongodb.client.MongoClient; import com.mongodb.client.MongoCollection; import com.mongodb.client.MongoDatabase; import com.mongodb.client.model.Filters; import com.mongodb.client.model.FindOneAndUpdateOptions; import com.mongodb.client.model.IndexOptions; import com.mongodb.client.model.Indexes; +import com.mongodb.client.model.ReturnDocument; import com.mongodb.client.model.Updates; import io.agentscope.harness.agent.sandbox.SandboxExecutionGuard; import io.agentscope.harness.agent.sandbox.SandboxIsolationKey; @@ -32,6 +34,7 @@ import java.time.Duration; import java.util.Date; import java.util.Objects; +import java.util.concurrent.TimeUnit; import org.bson.Document; import org.bson.conversions.Bson; import org.slf4j.Logger; @@ -44,9 +47,19 @@ * with a unique {@code _id} derived from the {@link SandboxIsolationKey} and a TTL index on {@code * expiresAt} to auto-release stale locks. * - *

    Lock acquisition uses {@code findOneAndUpdate} with upsert and a filter that rejects documents - * whose {@code expiresAt} has not yet passed. This provides a non-blocking try-lock semantics. - * Acquisition polls until the lock is obtained or the timeout expires. + *

    Lock acquisition uses a two-step approach for correct mutual exclusion: + * + *

      + *
    1. Attempt {@code insertOne} — atomic under the unique {@code _id} index; only one + * concurrent caller succeeds. + *
    2. If the lock document already exists (duplicate key), attempt {@code findOneAndUpdate} + * (non-upsert) with a filter that matches only when {@code expiresAt <= now} — reclaiming + * an expired lock. + *
    + * + *

    Acquisition polls until the lock is obtained or the timeout expires. The returned {@link + * SandboxLease} releases the lock on close, filtered by owner to avoid releasing another + * process's lock. */ public final class MongoSandboxExecutionGuard implements SandboxExecutionGuard { @@ -59,12 +72,14 @@ public final class MongoSandboxExecutionGuard implements SandboxExecutionGuard { private final MongoCollection collection; private final long lockTimeoutMs; + private final long retryIntervalMs; private final String owner; private MongoSandboxExecutionGuard(Builder builder) { MongoDatabase db = builder.mongoClient.getDatabase(builder.databaseName); this.collection = db.getCollection(builder.collectionName); this.lockTimeoutMs = builder.lockTimeout.toMillis(); + this.retryIntervalMs = builder.retryInterval.toMillis(); this.owner = builder.owner; ensureIndexes(); } @@ -75,7 +90,7 @@ private MongoSandboxExecutionGuard(Builder builder) { * @param mongoClient the MongoDB client * @return a new builder */ - public static Builder builder(com.mongodb.client.MongoClient mongoClient) { + public static Builder builder(MongoClient mongoClient) { return new Builder(mongoClient); } @@ -89,60 +104,77 @@ public SandboxLease tryEnter(SandboxIsolationKey key) throws InterruptedExceptio Date now = new Date(); Date expiresAt = new Date(now.getTime() + lockTimeoutMs); - Bson filter = - Filters.and( - Filters.eq(FIELD_LOCK_ID, lockId), - Filters.or( - Filters.exists(FIELD_OWNER, false), - Filters.lte(FIELD_EXPIRES_AT, now))); + // Step 1: Try to insert a new lock document. This is atomic — only one + // concurrent caller can succeed due to the _id unique index. + Document lockDoc = + new Document(FIELD_LOCK_ID, lockId) + .append(FIELD_OWNER, owner) + .append(FIELD_EXPIRES_AT, expiresAt); + try { + collection.insertOne(lockDoc); + log.debug("[sandbox-guard] Acquired MongoDB lock (insert): {}", lockId); + return new MongoLease(collection, lockId, owner); + } catch (MongoWriteException e) { + if (e.getError().getCode() != 11000) { + throw new RuntimeException("Failed to acquire MongoDB lock: " + lockId, e); + } + // Duplicate key — lock document already exists, fall through to step 2 + } - Bson update = + // Step 2: Lock document exists. Try to reclaim it if it has expired. + Bson expiredFilter = + Filters.and( + Filters.eq(FIELD_LOCK_ID, lockId), Filters.lte(FIELD_EXPIRES_AT, now)); + Bson reclaimUpdate = Updates.combine( - Updates.setOnInsert(FIELD_LOCK_ID, lockId), Updates.set(FIELD_OWNER, owner), Updates.set(FIELD_EXPIRES_AT, expiresAt)); - try { - Document result = - collection.findOneAndUpdate( - filter, update, new FindOneAndUpdateOptions().upsert(true)); - - if (result == null) { - log.debug("[sandbox-guard] Acquired MongoDB lock: {}", lockId); - return new MongoLease(collection, lockId); - } else { + Document reclaimed = + collection.findOneAndUpdate( + expiredFilter, + reclaimUpdate, + new FindOneAndUpdateOptions().returnDocument(ReturnDocument.AFTER)); + + if (reclaimed != null) { + log.debug( + "[sandbox-guard] Acquired MongoDB lock (reclaimed from {}): {}", + reclaimed.getString(FIELD_OWNER), + lockId); + return new MongoLease(collection, lockId, owner); + } + + // Lock exists and has not expired — held by someone else. + // + // TOCTOU safety: between the failed insertOne above and the findOneAndUpdate + // reclaim attempt, another process may have released the lock and a third process + // may have re-acquired it. This is safe because the reclaim filter includes + // `expiresAt <= now` — a freshly acquired lock has `expiresAt` in the future and + // will NOT match the reclaim filter, so we will not steal it. + if (log.isDebugEnabled()) { + Document held = collection.find(Filters.eq(FIELD_LOCK_ID, lockId)).first(); + if (held != null) { log.debug( "[sandbox-guard] Lock held by {}, retrying: {}", - result.getString(FIELD_OWNER), + held.getString(FIELD_OWNER), lockId); } - } catch (MongoBulkWriteException e) { - // Duplicate key — lock held by someone else, retry - } catch (Exception e) { - if (e.getMessage() != null && e.getMessage().contains("E11000")) { - // Duplicate key error — lock held by someone else - } else { - throw new RuntimeException("Failed to acquire MongoDB lock: " + lockId, e); - } } if (System.nanoTime() >= deadline) { - throw new InterruptedException( - "Timed out waiting for MongoDB lock: " - + lockId - + " (timeout=" - + Duration.ofMillis(lockTimeoutMs) - + ")"); + throw new InterruptedException("Timed out waiting for MongoDB lock: " + lockId); } - Thread.sleep(100L); + // MongoDB lacks server-side blocking locks (unlike MySQL GET_LOCK), so we poll. + // InterruptedException is declared on the method signature and propagated by sleep. + Thread.sleep(retryIntervalMs); } } private void ensureIndexes() { collection.createIndex( Indexes.ascending(FIELD_EXPIRES_AT), - new IndexOptions().expireAfter(0L, java.util.concurrent.TimeUnit.SECONDS)); + new IndexOptions().expireAfter(0L, TimeUnit.SECONDS)); } private static String composeLockId(SandboxIsolationKey key) { @@ -164,16 +196,23 @@ private static final class MongoLease implements SandboxLease { private final MongoCollection collection; private final String lockId; + private final String owner; - MongoLease(MongoCollection collection, String lockId) { + MongoLease(MongoCollection collection, String lockId, String owner) { this.collection = collection; this.lockId = lockId; + this.owner = owner; } @Override public void close() { try { - collection.deleteOne(Filters.eq(lockId)); + // Only delete if we still own the lock — prevents releasing someone else's + // lock when close() is called after lease expiry and re-acquisition by another + // process. Also makes close() idempotent: if already released, deleteOne is a + // no-op. + collection.deleteOne( + Filters.and(Filters.eq(lockId), Filters.eq(FIELD_OWNER, owner))); log.debug("[sandbox-guard] Released MongoDB lock: {}", lockId); } catch (Exception e) { log.warn( @@ -184,16 +223,21 @@ public void close() { } } - /** Builder for {@link MongoSandboxExecutionGuard}. */ + /** + * Builder for {@link MongoSandboxExecutionGuard}. + */ public static final class Builder { - private final com.mongodb.client.MongoClient mongoClient; + private static final Duration DEFAULT_RETRY_INTERVAL = Duration.ofMillis(500); + + private final MongoClient mongoClient; private String databaseName = "agentscope"; private String collectionName = DEFAULT_COLLECTION; private Duration lockTimeout = Duration.ofMinutes(30); + private Duration retryInterval = DEFAULT_RETRY_INTERVAL; private String owner = "agentscope:" + ProcessHandle.current().pid(); - Builder(com.mongodb.client.MongoClient mongoClient) { + Builder(MongoClient mongoClient) { this.mongoClient = Objects.requireNonNull(mongoClient, "mongoClient"); } @@ -216,6 +260,24 @@ public Builder lockTimeout(Duration timeout) { return this; } + /** + * Sets the polling interval between lock acquisition attempts. + * + *

    Default: {@code 500 ms}. Lower values reduce latency at the cost of more MongoDB + * round-trips; higher values reduce load at the cost of increased queuing delay. + * + * @param interval the retry interval; must be positive + * @return this builder + */ + public Builder retryInterval(Duration interval) { + Objects.requireNonNull(interval, "retryInterval"); + if (interval.isNegative() || interval.isZero()) { + throw new IllegalArgumentException("retryInterval must be positive"); + } + this.retryInterval = interval; + return this; + } + public Builder owner(String owner) { this.owner = owner; return this; diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClient.java b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClient.java index 81167d7f8a..8e1e7676cb 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClient.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClient.java @@ -15,17 +15,23 @@ */ package io.agentscope.extensions.mongodb.snapshot; +import com.mongodb.client.MongoClient; import com.mongodb.client.MongoCollection; import com.mongodb.client.MongoDatabase; import com.mongodb.client.model.Filters; +import com.mongodb.client.model.Indexes; +import com.mongodb.client.model.Projections; import com.mongodb.client.model.ReplaceOptions; +import com.mongodb.client.model.IndexOptions; import io.agentscope.harness.agent.sandbox.snapshot.RemoteSnapshotClient; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; +import java.io.FileNotFoundException; import java.io.IOException; import java.io.InputStream; import java.util.Date; import java.util.Objects; +import java.util.concurrent.TimeUnit; import org.bson.Document; import org.bson.types.Binary; import org.slf4j.Logger; @@ -44,12 +50,14 @@ public class MongoRemoteSnapshotClient implements RemoteSnapshotClient { private static final String DEFAULT_COLLECTION = "agentscope_snapshots"; private static final String FIELD_DATA = "data"; private static final String FIELD_CREATED_AT = "createdAt"; - private static final int MAX_SNAPSHOT_BYTES = 100 * 1024 * 1024; // 100 MB + // MongoDB BSON document size limit is 16 MB; cap at 15 MB to leave headroom for + // metadata. For larger snapshots, use GridFS (not yet implemented). + private static final int MAX_SNAPSHOT_BYTES = 15 * 1024 * 1024; // 15 MB private final MongoCollection collection; public MongoRemoteSnapshotClient( - com.mongodb.client.MongoClient mongoClient, + MongoClient mongoClient, String databaseName, String collectionName, boolean initializeSchema) { @@ -65,7 +73,11 @@ public MongoRemoteSnapshotClient( private void initSchema() { try { - collection.createIndex(com.mongodb.client.model.Indexes.ascending(FIELD_CREATED_AT)); + collection.createIndex(Indexes.ascending(FIELD_CREATED_AT)); + collection.createIndex( + Indexes.ascending(FIELD_CREATED_AT), + new IndexOptions() + .expireAfter(7 * 24 * 3600L, TimeUnit.SECONDS)); } catch (Exception e) { log.warn( "Failed to initialize snapshot collection index '{}': {}", @@ -90,10 +102,10 @@ public InputStream download(String snapshotId) throws Exception { Document doc = collection .find(Filters.eq(snapshotId)) - .projection(com.mongodb.client.model.Projections.include(FIELD_DATA)) + .projection(Projections.include(FIELD_DATA)) .first(); if (doc == null) { - throw new java.io.FileNotFoundException("Snapshot not found in MongoDB: " + snapshotId); + throw new FileNotFoundException("Snapshot not found in MongoDB: " + snapshotId); } Binary binary = doc.get(FIELD_DATA, Binary.class); return new ByteArrayInputStream(binary.getData()); @@ -104,11 +116,23 @@ public boolean exists(String snapshotId) throws Exception { Objects.requireNonNull(snapshotId, "snapshotId"); return collection .find(Filters.eq(snapshotId)) - .projection(com.mongodb.client.model.Projections.include("_id")) + .projection(Projections.include("_id")) .first() != null; } + /** + * Deletes a snapshot from MongoDB. + * + * @param snapshotId the snapshot identifier + * @return {@code true} if a document was deleted, {@code false} if no matching snapshot existed + * @throws Exception if a MongoDB error occurs + */ + public boolean delete(String snapshotId) throws Exception { + Objects.requireNonNull(snapshotId, "snapshotId"); + return collection.deleteOne(Filters.eq(snapshotId)).getDeletedCount() > 0; + } + private static byte[] readAllBounded(InputStream in, int maxBytes) throws IOException { ByteArrayOutputStream out = new ByteArrayOutputStream(Math.min(maxBytes, 8192)); byte[] buf = new byte[8192]; diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoSnapshotSpec.java b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoSnapshotSpec.java index 22148f66a0..50de16ff8a 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoSnapshotSpec.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoSnapshotSpec.java @@ -15,6 +15,7 @@ */ package io.agentscope.extensions.mongodb.snapshot; +import com.mongodb.client.MongoClient; import io.agentscope.harness.agent.sandbox.snapshot.RemoteSnapshotSpec; /** @@ -25,14 +26,11 @@ */ public class MongoSnapshotSpec extends RemoteSnapshotSpec { - public MongoSnapshotSpec(com.mongodb.client.MongoClient mongoClient, String databaseName) { + public MongoSnapshotSpec(MongoClient mongoClient, String databaseName) { super(new MongoRemoteSnapshotClient(mongoClient, databaseName, null, true)); } - public MongoSnapshotSpec( - com.mongodb.client.MongoClient mongoClient, - String databaseName, - String collectionName) { + public MongoSnapshotSpec(MongoClient mongoClient, String databaseName, String collectionName) { super(new MongoRemoteSnapshotClient(mongoClient, databaseName, collectionName, true)); } } diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStore.java b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStore.java index a58eff0745..7510a7cadf 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStore.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStore.java @@ -17,15 +17,18 @@ import com.mongodb.ConnectionString; import com.mongodb.MongoClientSettings; +import com.mongodb.MongoWriteException; import com.mongodb.client.MongoClient; import com.mongodb.client.MongoClients; import com.mongodb.client.MongoCollection; import com.mongodb.client.MongoDatabase; import com.mongodb.client.model.Filters; import com.mongodb.client.model.FindOneAndUpdateOptions; +import com.mongodb.client.model.IndexOptions; import com.mongodb.client.model.Indexes; import com.mongodb.client.model.Projections; import com.mongodb.client.model.ReturnDocument; +import com.mongodb.client.model.UpdateOptions; import com.mongodb.client.model.Updates; import io.agentscope.core.state.AgentStateStore; import io.agentscope.core.state.ListHashUtil; @@ -121,9 +124,7 @@ private void ensureIndexes() { Indexes.ascending(FIELD_USER_ID), Indexes.ascending(FIELD_SESSION_ID))); collection.createIndex( Indexes.ascending(FIELD_UPDATED_AT), - new com.mongodb.client.model.IndexOptions() - .expireAfter(0L, TimeUnit.SECONDS) - .sparse(true)); + new IndexOptions().expireAfter(0L, TimeUnit.SECONDS).sparse(true)); } // ────────────────── Single Value CRUD ────────────────── @@ -159,6 +160,15 @@ public Optional get( // ────────────────── List CRUD ────────────────── + /** + * Saves a list of state values with incremental-append optimization. + * + *

    Concurrency note: this method performs a read-then-write to decide between + * incremental append and full replacement. It is NOT atomic — concurrent calls for the same + * {@code (userId, sessionId, key)} may interleave reads and writes, causing lost appends or + * stale hash comparisons. Callers that require strict consistency should synchronize externally + * (e.g. via {@link io.agentscope.harness.agent.sandbox.SandboxExecutionGuard}). + */ @Override public void save(String userId, String sessionId, String key, List values) { validateKey(key); @@ -207,6 +217,20 @@ public void save(String userId, String sessionId, String key, List bsonList = toDocumentList(values); + Bson setFields = + Updates.combine( + Updates.set(listKey, bsonList), + Updates.set(hashField, currentHash), + Updates.set(FIELD_UPDATED_AT, new Date())); + Bson setOnInsert = + Updates.combine( + Updates.setOnInsert(FIELD_USER_ID, normalizeUser(userId)), + Updates.setOnInsert(FIELD_SESSION_ID, sessionId)); + collection.updateOne( + Filters.eq(slotId), Updates.combine(setFields, setOnInsert), upsert()); } } @@ -276,18 +300,25 @@ public long saveIfVersion( Updates.set(FIELD_UPDATED_AT, new Date()), Updates.setOnInsert(FIELD_USER_ID, normalizeUser(userId)), Updates.setOnInsert(FIELD_SESSION_ID, sessionId)); - Document result = - collection.findOneAndUpdate( - filter, - update, - new FindOneAndUpdateOptions() - .upsert(true) - .returnDocument(ReturnDocument.AFTER)); - if (result == null) { - return UNVERSIONED; + try { + Document result = + collection.findOneAndUpdate( + filter, + update, + new FindOneAndUpdateOptions() + .upsert(true) + .returnDocument(ReturnDocument.AFTER)); + if (result == null) { + return UNVERSIONED; + } + Long newVersion = result.getLong(versionField); + return newVersion != null && newVersion == 1L ? 1L : UNVERSIONED; + } catch (MongoWriteException e) { + if (e.getError().getCode() == 11000) { + return UNVERSIONED; + } + throw e; } - Long newVersion = result.getLong(versionField); - return newVersion != null && newVersion == 1L ? 1L : UNVERSIONED; } Bson filter = Filters.and(Filters.eq(slotId), Filters.eq(versionField, expectedVersion)); @@ -407,13 +438,15 @@ private List toDocumentList(List values) { return result; } - private static com.mongodb.client.model.UpdateOptions upsert() { - return new com.mongodb.client.model.UpdateOptions().upsert(true); + private static UpdateOptions upsert() { + return new UpdateOptions().upsert(true); } // ────────────────── Builder ────────────────── - /** Builder for {@link MongoAgentStateStore}. */ + /** + * Builder for {@link MongoAgentStateStore}. + */ public static class Builder { private MongoClient mongoClient; private String connectionString; diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/store/MongoBaseStore.java b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/store/MongoBaseStore.java index f6ea503906..f833b46ae4 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/store/MongoBaseStore.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/store/MongoBaseStore.java @@ -15,8 +15,10 @@ */ package io.agentscope.extensions.mongodb.store; +import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; +import com.mongodb.MongoWriteException; import com.mongodb.client.MongoCollection; import com.mongodb.client.MongoDatabase; import com.mongodb.client.model.Filters; @@ -25,6 +27,7 @@ import com.mongodb.client.model.Projections; import com.mongodb.client.model.ReturnDocument; import com.mongodb.client.model.Sorts; +import com.mongodb.client.model.UpdateOptions; import com.mongodb.client.model.Updates; import io.agentscope.harness.agent.filesystem.remote.store.BaseStore; import io.agentscope.harness.agent.filesystem.remote.store.StoreItem; @@ -35,6 +38,8 @@ import java.util.stream.Collectors; import org.bson.Document; import org.bson.conversions.Bson; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * MongoDB-backed implementation of {@link BaseStore}. @@ -45,6 +50,8 @@ */ public class MongoBaseStore implements BaseStore { + private static final Logger log = LoggerFactory.getLogger(MongoBaseStore.class); + private static final String FIELD_ID = "_id"; private static final String FIELD_KEY = "key"; private static final String FIELD_NAMESPACE = "namespace"; @@ -59,7 +66,7 @@ public class MongoBaseStore implements BaseStore { /** * Creates a new instance. * - * @param database the MongoDB database + * @param database the MongoDB database * @param collectionName the collection name */ public MongoBaseStore(MongoDatabase database, String collectionName) { @@ -69,9 +76,9 @@ public MongoBaseStore(MongoDatabase database, String collectionName) { /** * Creates a new instance with a custom ObjectMapper. * - * @param database the MongoDB database + * @param database the MongoDB database * @param collectionName the collection name - * @param objectMapper Jackson mapper for serializing values + * @param objectMapper Jackson mapper for serializing values */ public MongoBaseStore( MongoDatabase database, String collectionName, ObjectMapper objectMapper) { @@ -122,21 +129,21 @@ public boolean putIfVersion( Document result; if (expectedVersion == 0) { - Bson filter = Filters.and(Filters.eq(id), Filters.exists(FIELD_VERSION, false)); - Bson update = - Updates.combine( - Updates.set(FIELD_VALUE, Document.parse(json)), - Updates.set(FIELD_KEY, key), - Updates.set(FIELD_NAMESPACE, nsKey), - Updates.set(FIELD_VERSION, 1L), - Updates.setOnInsert(FIELD_ID, id)); - result = - collection.findOneAndUpdate( - filter, - update, - new FindOneAndUpdateOptions() - .upsert(true) - .returnDocument(ReturnDocument.AFTER)); + Document doc = + new Document(FIELD_ID, id) + .append(FIELD_VALUE, Document.parse(json)) + .append(FIELD_KEY, key) + .append(FIELD_NAMESPACE, nsKey) + .append(FIELD_VERSION, 1L); + try { + collection.insertOne(doc); + return true; + } catch (MongoWriteException e) { + if (e.getError().getCode() == 11000) { + return false; + } + throw e; + } } else { Bson filter = Filters.and(Filters.eq(id), Filters.eq(FIELD_VERSION, expectedVersion)); Bson update = @@ -156,9 +163,6 @@ public boolean putIfVersion( return false; } Long newVersion = result.getLong(FIELD_VERSION); - if (expectedVersion == 0 && newVersion != null && newVersion == 1L) { - return true; - } return newVersion != null && newVersion == expectedVersion + 1; } @@ -209,12 +213,11 @@ private static String namespacePath(List namespace) { private String serialize(Map value) { try { return objectMapper.writeValueAsString(value == null ? Map.of() : value); - } catch (com.fasterxml.jackson.core.JsonProcessingException e) { + } catch (JsonProcessingException e) { throw new IllegalStateException("Failed to serialize value", e); } } - @SuppressWarnings("unchecked") private Map parseValue(Object raw) { if (raw instanceof Document doc) { return new LinkedHashMap<>(doc); @@ -223,14 +226,17 @@ private Map parseValue(Object raw) { try { Map parsed = objectMapper.readValue(s, MAP_TYPE); return parsed != null ? parsed : Map.of(); - } catch (com.fasterxml.jackson.core.JsonProcessingException e) { + } catch (JsonProcessingException e) { + log.warn( + "Failed to parse stored JSON value, returning empty map: {}", + e.getMessage()); return Map.of(); } } return Map.of(); } - private static com.mongodb.client.model.UpdateOptions upsert() { - return new com.mongodb.client.model.UpdateOptions().upsert(true); + private static UpdateOptions upsert() { + return new UpdateOptions().upsert(true); } } diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/MongoDistributedStoreTest.java b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/MongoDistributedStoreTest.java index cd49949c50..1196cb4cfb 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/MongoDistributedStoreTest.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/MongoDistributedStoreTest.java @@ -16,8 +16,11 @@ package io.agentscope.extensions.mongodb; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import com.mongodb.client.MongoClient; @@ -25,6 +28,8 @@ import com.mongodb.client.MongoDatabase; import io.agentscope.core.state.AgentStateStore; import io.agentscope.harness.agent.filesystem.remote.store.BaseStore; +import io.agentscope.harness.agent.sandbox.SandboxExecutionGuard; +import io.agentscope.harness.agent.sandbox.snapshot.SandboxSnapshotSpec; import org.bson.Document; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -84,4 +89,71 @@ void baseStoreReturnsNonNull() { void createWithNullMongoClientThrows() { assertThrows(NullPointerException.class, () -> MongoDistributedStore.create(null)); } + + @Test + void sandboxExecutionGuardReturnsNonNull() { + MongoDistributedStore store = MongoDistributedStore.create(mongoClient); + SandboxExecutionGuard guard = store.sandboxExecutionGuard(); + assertNotNull(guard); + } + + @Test + void sandboxSnapshotSpecReturnsNonNull() { + MongoDistributedStore store = MongoDistributedStore.create(mongoClient); + assertNotNull(store.sandboxSnapshotSpec()); + } + + @Test + void fromConnectionStringCreatesStore() { + // MongoClient creation is lazy — no actual connection until an operation is performed + MongoDistributedStore store = + MongoDistributedStore.fromConnectionString("mongodb://localhost:27017"); + assertNotNull(store); + } + + @Test + void agentStateStoreReturnsCachedInstance() { + MongoDistributedStore store = MongoDistributedStore.create(mongoClient); + AgentStateStore first = store.agentStateStore(); + AgentStateStore second = store.agentStateStore(); + assertSame(first, second); + } + + @Test + void baseStoreReturnsCachedInstance() { + MongoDistributedStore store = MongoDistributedStore.create(mongoClient); + BaseStore first = store.baseStore(); + BaseStore second = store.baseStore(); + assertSame(first, second); + } + + @Test + void sandboxSnapshotSpecReturnsCachedInstance() { + MongoDistributedStore store = MongoDistributedStore.create(mongoClient); + SandboxSnapshotSpec first = store.sandboxSnapshotSpec(); + SandboxSnapshotSpec second = store.sandboxSnapshotSpec(); + assertSame(first, second); + } + + @Test + void sandboxExecutionGuardReturnsCachedInstance() { + MongoDistributedStore store = MongoDistributedStore.create(mongoClient); + SandboxExecutionGuard first = store.sandboxExecutionGuard(); + SandboxExecutionGuard second = store.sandboxExecutionGuard(); + assertSame(first, second); + } + + @Test + void closeWithExternalClientDoesNotCloseClient() { + MongoDistributedStore store = MongoDistributedStore.create(mongoClient); + store.close(); + verify(mongoClient, never()).close(); + } + + @Test + void closeFromConnectionStringDoesNotThrow() { + MongoDistributedStore store = + MongoDistributedStore.fromConnectionString("mongodb://localhost:27017"); + store.close(); + } } diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/sandbox/MongoSandboxExecutionGuardTest.java b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/sandbox/MongoSandboxExecutionGuardTest.java new file mode 100644 index 0000000000..3d7674b54a --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/sandbox/MongoSandboxExecutionGuardTest.java @@ -0,0 +1,311 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.extensions.mongodb.sandbox; + +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.mongodb.MongoWriteException; +import com.mongodb.ServerAddress; +import com.mongodb.WriteError; +import com.mongodb.client.FindIterable; +import com.mongodb.client.MongoClient; +import com.mongodb.client.MongoCollection; +import com.mongodb.client.MongoDatabase; +import com.mongodb.client.model.FindOneAndUpdateOptions; +import io.agentscope.harness.agent.IsolationScope; +import io.agentscope.harness.agent.sandbox.SandboxIsolationKey; +import io.agentscope.harness.agent.sandbox.SandboxLease; +import java.time.Duration; +import java.util.Date; +import org.bson.BsonDocument; +import org.bson.Document; +import org.bson.conversions.Bson; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +class MongoSandboxExecutionGuardTest { + + @Mock private MongoClient mongoClient; + @Mock private MongoDatabase mongoDatabase; + @Mock private MongoCollection collection; + + @SuppressWarnings("rawtypes") + @Mock + private FindIterable findIterable; + + private AutoCloseable mocks; + + @BeforeEach + @SuppressWarnings("unchecked") + void setUp() { + mocks = MockitoAnnotations.openMocks(this); + when(mongoClient.getDatabase(anyString())).thenReturn(mongoDatabase); + when(mongoDatabase.getCollection(anyString())).thenReturn(collection); + + when(collection.find(any(Bson.class))).thenReturn(findIterable); + when(findIterable.projection(any())).thenReturn(findIterable); + when(findIterable.first()).thenReturn(null); + } + + @AfterEach + void tearDown() throws Exception { + if (mocks != null) { + mocks.close(); + } + } + + private SandboxIsolationKey key() { + return SandboxIsolationKey.resolve( + IsolationScope.SESSION, + new io.agentscope.core.agent.RuntimeContext.Builder() + .sessionId("session-1") + .build(), + "agent") + .orElseThrow(); + } + + @Test + void builderRejectsNullMongoClient() { + assertThrows(NullPointerException.class, () -> MongoSandboxExecutionGuard.builder(null)); + } + + @Test + void builderWithDefaultsCreatesGuard() { + MongoSandboxExecutionGuard guard = MongoSandboxExecutionGuard.builder(mongoClient).build(); + assertNotNull(guard); + } + + @Test + void builderWithCustomDatabaseName() { + MongoSandboxExecutionGuard guard = + MongoSandboxExecutionGuard.builder(mongoClient).databaseName("custom_db").build(); + assertNotNull(guard); + } + + @Test + void builderWithCustomCollectionName() { + MongoSandboxExecutionGuard guard = + MongoSandboxExecutionGuard.builder(mongoClient) + .collectionName("custom_locks") + .build(); + assertNotNull(guard); + } + + @Test + void builderWithCustomTimeout() { + MongoSandboxExecutionGuard guard = + MongoSandboxExecutionGuard.builder(mongoClient) + .lockTimeout(Duration.ofSeconds(5)) + .build(); + assertNotNull(guard); + } + + @Test + void builderRejectsNonPositiveTimeout() { + assertThrows( + IllegalArgumentException.class, + () -> + MongoSandboxExecutionGuard.builder(mongoClient) + .lockTimeout(Duration.ZERO) + .build()); + assertThrows( + IllegalArgumentException.class, + () -> + MongoSandboxExecutionGuard.builder(mongoClient) + .lockTimeout(Duration.ofSeconds(-1)) + .build()); + } + + @Test + void builderWithCustomRetryInterval() { + MongoSandboxExecutionGuard guard = + MongoSandboxExecutionGuard.builder(mongoClient) + .retryInterval(Duration.ofMillis(200)) + .build(); + assertNotNull(guard); + } + + @Test + void builderRejectsNonPositiveRetryInterval() { + assertThrows( + IllegalArgumentException.class, + () -> + MongoSandboxExecutionGuard.builder(mongoClient) + .retryInterval(Duration.ZERO) + .build()); + assertThrows( + IllegalArgumentException.class, + () -> + MongoSandboxExecutionGuard.builder(mongoClient) + .retryInterval(Duration.ofMillis(-1)) + .build()); + } + + @Test + void builderWithCustomOwner() { + MongoSandboxExecutionGuard guard = + MongoSandboxExecutionGuard.builder(mongoClient).owner("custom-owner").build(); + assertNotNull(guard); + } + + @Test + void tryEnterAcquiresLockViaInsert() throws Exception { + MongoSandboxExecutionGuard guard = + MongoSandboxExecutionGuard.builder(mongoClient) + .lockTimeout(Duration.ofSeconds(5)) + .build(); + // insertOne succeeds — no duplicate key → lock acquired immediately + SandboxLease lease = guard.tryEnter(key()); + + assertNotNull(lease); + verify(collection).insertOne(any(Document.class)); + } + + @Test + void tryEnterReclaimsExpiredLock() throws Exception { + MongoSandboxExecutionGuard guard = + MongoSandboxExecutionGuard.builder(mongoClient) + .lockTimeout(Duration.ofSeconds(10)) + .build(); + // Step 1: insertOne fails with duplicate key → lock doc exists + WriteError dupError = new WriteError(11000, "duplicate key", new BsonDocument()); + when(collection.insertOne(any(Document.class))) + .thenThrow(new MongoWriteException(dupError, new ServerAddress())); + // Step 2: findOneAndUpdate succeeds → lock was expired, we reclaimed it + Document reclaimed = + new Document("_id", "lock:abc") + .append("owner", "old-owner") + .append("expiresAt", new Date(0L)); + when(collection.findOneAndUpdate( + any(Bson.class), any(Bson.class), any(FindOneAndUpdateOptions.class))) + .thenReturn(reclaimed); + + SandboxLease lease = guard.tryEnter(key()); + + assertNotNull(lease); + verify(collection).insertOne(any(Document.class)); + verify(collection) + .findOneAndUpdate( + any(Bson.class), any(Bson.class), any(FindOneAndUpdateOptions.class)); + } + + @Test + void tryEnterPollsWhenLockHeldThenAcquires() throws Exception { + MongoSandboxExecutionGuard guard = + MongoSandboxExecutionGuard.builder(mongoClient) + .lockTimeout(Duration.ofSeconds(10)) + .build(); + // insertOne always fails with duplicate key + WriteError dupError = new WriteError(11000, "duplicate key", new BsonDocument()); + when(collection.insertOne(any(Document.class))) + .thenThrow(new MongoWriteException(dupError, new ServerAddress())); + // First findOneAndUpdate returns null (lock not expired), second returns reclaimed doc + Document reclaimed = + new Document("_id", "lock:abc") + .append("owner", "old-owner") + .append("expiresAt", new Date(0L)); + when(collection.findOneAndUpdate( + any(Bson.class), any(Bson.class), any(FindOneAndUpdateOptions.class))) + .thenReturn(null) + .thenReturn(reclaimed); + + SandboxLease lease = guard.tryEnter(key()); + + assertNotNull(lease); + // insertOne called twice (once per poll iteration) + verify(collection, times(2)).insertOne(any(Document.class)); + } + + @Test + void tryEnterTimesOutWhenLockNeverAcquired() { + MongoSandboxExecutionGuard guard = + MongoSandboxExecutionGuard.builder(mongoClient) + .lockTimeout(Duration.ofMillis(200)) + .build(); + // insertOne always fails with duplicate key + WriteError dupError = new WriteError(11000, "duplicate key", new BsonDocument()); + when(collection.insertOne(any(Document.class))) + .thenThrow(new MongoWriteException(dupError, new ServerAddress())); + // findOneAndUpdate always returns null (lock never expires) + when(collection.findOneAndUpdate( + any(Bson.class), any(Bson.class), any(FindOneAndUpdateOptions.class))) + .thenReturn(null); + + assertThrows(InterruptedException.class, () -> guard.tryEnter(key())); + } + + @Test + void tryEnterPropagatesNonDuplicateKeyWriteException() { + MongoSandboxExecutionGuard guard = + MongoSandboxExecutionGuard.builder(mongoClient) + .lockTimeout(Duration.ofSeconds(5)) + .build(); + // Non-duplicate-key write error from insertOne should propagate + WriteError otherError = new WriteError(99999, "disk full", new BsonDocument()); + when(collection.insertOne(any(Document.class))) + .thenThrow(new MongoWriteException(otherError, new ServerAddress())); + + assertThrows(RuntimeException.class, () -> guard.tryEnter(key())); + } + + @Test + void tryEnterPropagatesNonWriteException() { + MongoSandboxExecutionGuard guard = + MongoSandboxExecutionGuard.builder(mongoClient) + .lockTimeout(Duration.ofSeconds(5)) + .build(); + RuntimeException unexpected = new RuntimeException("connection lost"); + when(collection.insertOne(any(Document.class))).thenThrow(unexpected); + + assertThrows(RuntimeException.class, () -> guard.tryEnter(key())); + } + + @Test + void leaseCloseReleasesLockWithOwnerCheck() throws Exception { + MongoSandboxExecutionGuard guard = + MongoSandboxExecutionGuard.builder(mongoClient) + .lockTimeout(Duration.ofSeconds(5)) + .build(); + SandboxLease lease = guard.tryEnter(key()); + lease.close(); + + // deleteOne should filter by both _id AND owner + verify(collection).deleteOne(any(Bson.class)); + } + + @Test + void leaseCloseHandlesReleaseFailure() throws Exception { + MongoSandboxExecutionGuard guard = + MongoSandboxExecutionGuard.builder(mongoClient) + .lockTimeout(Duration.ofSeconds(5)) + .build(); + SandboxLease lease = guard.tryEnter(key()); + doThrow(new RuntimeException("network error")).when(collection).deleteOne(any(Bson.class)); + + // Should not throw — close() swallows errors + lease.close(); + } +} diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClientTest.java b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClientTest.java new file mode 100644 index 0000000000..29a21ce9c7 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClientTest.java @@ -0,0 +1,231 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.extensions.mongodb.snapshot; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.mongodb.client.FindIterable; +import com.mongodb.client.MongoClient; +import com.mongodb.client.MongoCollection; +import com.mongodb.client.MongoDatabase; +import com.mongodb.client.result.DeleteResult; +import com.mongodb.client.result.UpdateResult; +import java.io.ByteArrayInputStream; +import java.io.FileNotFoundException; +import java.io.IOException; +import java.io.InputStream; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import org.bson.Document; +import org.bson.conversions.Bson; +import org.bson.types.Binary; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +class MongoRemoteSnapshotClientTest { + + @Mock private MongoClient mongoClient; + @Mock private MongoDatabase mongoDatabase; + @Mock private MongoCollection collection; + + @SuppressWarnings("rawtypes") + @Mock + private FindIterable findIterable; + + private AutoCloseable mocks; + private MongoRemoteSnapshotClient client; + + @BeforeEach + @SuppressWarnings("unchecked") + void setUp() { + mocks = MockitoAnnotations.openMocks(this); + when(mongoClient.getDatabase(anyString())).thenReturn(mongoDatabase); + when(mongoDatabase.getCollection(anyString())).thenReturn(collection); + + when(collection.find(any(Bson.class))).thenReturn(findIterable); + when(findIterable.projection(any())).thenReturn(findIterable); + + UpdateResult replaceResult = org.mockito.Mockito.mock(UpdateResult.class); + when(replaceResult.wasAcknowledged()).thenReturn(true); + when(collection.replaceOne(any(Bson.class), any(Document.class), any())) + .thenReturn(replaceResult); + + client = new MongoRemoteSnapshotClient(mongoClient, "testdb", null, false); + } + + @AfterEach + void tearDown() throws Exception { + if (mocks != null) { + mocks.close(); + } + } + + @Test + void constructorWithDefaultsCreatesClient() { + MongoRemoteSnapshotClient defaultClient = + new MongoRemoteSnapshotClient(mongoClient, null, null, false); + assertNotNull(defaultClient); + } + + @Test + void constructorWithInitializeSchemaCreatesIndexes() { + MongoRemoteSnapshotClient schemaClient = + new MongoRemoteSnapshotClient(mongoClient, "testdb", "custom_snapshots", true); + assertNotNull(schemaClient); + verify(mongoDatabase).getCollection("custom_snapshots"); + } + + @Test + void constructorRejectsNullMongoClient() { + assertThrows( + NullPointerException.class, + () -> new MongoRemoteSnapshotClient(null, "testdb", null, false)); + } + + @Test + void uploadStoresSnapshotData() throws Exception { + byte[] data = "snapshot-content".getBytes(StandardCharsets.UTF_8); + InputStream in = new ByteArrayInputStream(data); + + client.upload("snap-1", in); + + verify(collection).replaceOne(any(Bson.class), any(Document.class), any()); + } + + @Test + void uploadRejectsNullSnapshotId() { + assertThrows( + NullPointerException.class, + () -> client.upload(null, new ByteArrayInputStream(new byte[0]))); + } + + @Test + void uploadRejectsNullData() { + assertThrows(NullPointerException.class, () -> client.upload("snap-1", null)); + } + + @Test + void uploadRejectsOversizedData() { + // Create a stream that exceeds MAX_SNAPSHOT_BYTES (15 MB) + InputStream oversized = + new InputStream() { + private int totalRead = 0; + private final int maxBytes = 15 * 1024 * 1024 + 1; + + @Override + public int read() { + if (totalRead >= maxBytes) { + return -1; + } + totalRead++; + return 'x'; + } + + @Override + public int read(byte[] b, int off, int len) { + if (totalRead >= maxBytes) { + return -1; + } + int toRead = Math.min(len, maxBytes - totalRead); + totalRead += toRead; + Arrays.fill(b, off, off + toRead, (byte) 'x'); + return toRead; + } + }; + + assertThrows(IOException.class, () -> client.upload("snap-1", oversized)); + } + + @Test + void downloadReturnsSnapshotData() throws Exception { + byte[] expected = "hello-snapshot".getBytes(StandardCharsets.UTF_8); + Document doc = new Document("data", new Binary(expected)); + when(findIterable.first()).thenReturn(doc); + + InputStream result = client.download("snap-1"); + + assertNotNull(result); + byte[] actual = result.readAllBytes(); + assertArrayEquals(expected, actual); + } + + @Test + void downloadThrowsWhenSnapshotNotFound() { + when(findIterable.first()).thenReturn(null); + + assertThrows(FileNotFoundException.class, () -> client.download("missing-snap")); + } + + @Test + void downloadRejectsNullSnapshotId() { + assertThrows(NullPointerException.class, () -> client.download(null)); + } + + @Test + void existsReturnsTrueWhenSnapshotFound() throws Exception { + when(findIterable.first()).thenReturn(new Document("_id", "snap-1")); + + assertTrue(client.exists("snap-1")); + } + + @Test + void existsReturnsFalseWhenSnapshotNotFound() throws Exception { + when(findIterable.first()).thenReturn(null); + + assertFalse(client.exists("snap-1")); + } + + @Test + void existsRejectsNullSnapshotId() { + assertThrows(NullPointerException.class, () -> client.exists(null)); + } + + @Test + void deleteReturnsTrueWhenSnapshotDeleted() throws Exception { + DeleteResult deleteResult = mock(DeleteResult.class); + when(deleteResult.getDeletedCount()).thenReturn(1L); + when(collection.deleteOne(any(Bson.class))).thenReturn(deleteResult); + + assertTrue(client.delete("snap-1")); + verify(collection).deleteOne(any(Bson.class)); + } + + @Test + void deleteReturnsFalseWhenSnapshotNotFound() throws Exception { + DeleteResult deleteResult = mock(DeleteResult.class); + when(deleteResult.getDeletedCount()).thenReturn(0L); + when(collection.deleteOne(any(Bson.class))).thenReturn(deleteResult); + + assertFalse(client.delete("missing-snap")); + } + + @Test + void deleteRejectsNullSnapshotId() { + assertThrows(NullPointerException.class, () -> client.delete(null)); + } +} diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/snapshot/MongoSnapshotSpecTest.java b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/snapshot/MongoSnapshotSpecTest.java new file mode 100644 index 0000000000..8da16dbdac --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/snapshot/MongoSnapshotSpecTest.java @@ -0,0 +1,65 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.extensions.mongodb.snapshot; + +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.when; + +import com.mongodb.client.MongoClient; +import com.mongodb.client.MongoCollection; +import com.mongodb.client.MongoDatabase; +import org.bson.Document; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +class MongoSnapshotSpecTest { + + @Mock private MongoClient mongoClient; + @Mock private MongoDatabase mongoDatabase; + @Mock private MongoCollection collection; + + private AutoCloseable mocks; + + @BeforeEach + void setUp() { + mocks = MockitoAnnotations.openMocks(this); + when(mongoClient.getDatabase(anyString())).thenReturn(mongoDatabase); + when(mongoDatabase.getCollection(anyString())).thenReturn(collection); + } + + @AfterEach + void tearDown() throws Exception { + if (mocks != null) { + mocks.close(); + } + } + + @Test + void constructorWithDefaultsCreatesSpec() { + MongoSnapshotSpec spec = new MongoSnapshotSpec(mongoClient, "testdb"); + assertNotNull(spec); + } + + @Test + void constructorWithCustomCollectionName() { + MongoSnapshotSpec spec = new MongoSnapshotSpec(mongoClient, "testdb", "custom_snapshots"); + assertNotNull(spec); + } +} diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreTest.java b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreTest.java index 13e925600c..3f531695ef 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreTest.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreTest.java @@ -26,25 +26,33 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import com.mongodb.MongoWriteException; +import com.mongodb.ServerAddress; +import com.mongodb.WriteError; import com.mongodb.client.DistinctIterable; import com.mongodb.client.FindIterable; import com.mongodb.client.MongoClient; import com.mongodb.client.MongoCollection; import com.mongodb.client.MongoDatabase; +import com.mongodb.client.model.FindOneAndUpdateOptions; import com.mongodb.client.result.DeleteResult; import com.mongodb.client.result.UpdateResult; import io.agentscope.core.state.AgentStateStore; import io.agentscope.core.state.State; import io.agentscope.core.state.VersionedState; +import java.util.ArrayList; import java.util.List; import java.util.Optional; import java.util.Set; +import org.bson.BsonDocument; import org.bson.Document; import org.bson.conversions.Bson; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.mockito.ArgumentMatchers; import org.mockito.Mock; +import org.mockito.Mockito; import org.mockito.MockitoAnnotations; class MongoAgentStateStoreTest { @@ -76,8 +84,8 @@ void setUp() { when(collection.find(any(Bson.class))).thenReturn(findIterable); when(findIterable.projection(any())).thenReturn(findIterable); when(findIterable.sort(any())).thenReturn(findIterable); - when(findIterable.skip(org.mockito.ArgumentMatchers.anyInt())).thenReturn(findIterable); - when(findIterable.limit(org.mockito.ArgumentMatchers.anyInt())).thenReturn(findIterable); + when(findIterable.skip(ArgumentMatchers.anyInt())).thenReturn(findIterable); + when(findIterable.limit(ArgumentMatchers.anyInt())).thenReturn(findIterable); when(findIterable.first()).thenReturn(null); UpdateResult updateResult = mock(UpdateResult.class); @@ -94,7 +102,7 @@ void setUp() { when(collection.distinct(anyString(), any(Bson.class), any(Class.class))) .thenReturn(distinctIterable); - when(distinctIterable.into(any())).thenReturn(new java.util.ArrayList<>()); + when(distinctIterable.into(any())).thenReturn(new ArrayList<>()); store = MongoAgentStateStore.builder() @@ -162,6 +170,22 @@ void saveListState() { verify(collection).updateOne(any(Bson.class), any(Bson.class), any()); } + @Test + void saveListShorteningPerformsFullRewrite() { + // Simulate existing document with 3 elements in the list + List existingList = + List.of( + Document.parse("{\"value\":\"a\"}"), + Document.parse("{\"value\":\"b\"}"), + Document.parse("{\"value\":\"c\"}")); + Document existingDoc = new Document("list:list", existingList); + when(findIterable.first()).thenReturn(existingDoc); + + // Save a shorter list (2 elements) — must still call updateOne (full rewrite) + store.save("user", "session", "list", List.of(new TestState("x"), new TestState("y"))); + verify(collection).updateOne(any(Bson.class), any(Bson.class), any()); + } + @Test void getListReturnsEmptyWhenMissing() { List result = store.getList("user", "session", "list", TestState.class); @@ -208,6 +232,81 @@ void saveIfVersionWithUnversionedDelegatesToSave() { verify(collection).updateOne(any(Bson.class), any(Bson.class), any()); } + @Test + void saveIfVersionZeroCreatesWhenAbsent() { + // findOneAndUpdate returns doc with _version_key=1 -> success (version 1 created) + Document result = + new Document("_id", "anon:session") + .append("key", Document.parse("{\"value\":\"v\"}")) + .append("_version_key", 1L); + when(collection.findOneAndUpdate( + any(Bson.class), any(Bson.class), any(FindOneAndUpdateOptions.class))) + .thenReturn(result); + + long newVersion = store.saveIfVersion("user", "session", "key", new TestState("v"), 0L); + + assertEquals(1L, newVersion); + } + + @Test + void saveIfVersionZeroReturnsUnversionedWhenAlreadyExists() { + // findOneAndUpdate returns doc with _version_key=5 -> expectedVersion=0 won't match + Document result = + new Document("_id", "anon:session") + .append("key", Document.parse("{\"value\":\"v\"}")) + .append("_version_key", 5L); + when(collection.findOneAndUpdate( + any(Bson.class), any(Bson.class), any(FindOneAndUpdateOptions.class))) + .thenReturn(result); + + long newVersion = store.saveIfVersion("user", "session", "key", new TestState("v"), 0L); + + assertEquals(AgentStateStore.UNVERSIONED, newVersion); + } + + @Test + void saveIfVersionZeroReturnsUnversionedOnDuplicateKey() { + // findOneAndUpdate with upsert throws E11000 -> document already exists + WriteError writeError = + new WriteError(11000, "E11000 duplicate key error", new BsonDocument()); + when(collection.findOneAndUpdate( + any(Bson.class), any(Bson.class), any(FindOneAndUpdateOptions.class))) + .thenThrow(new MongoWriteException(writeError, new ServerAddress())); + + long newVersion = store.saveIfVersion("user", "session", "key", new TestState("v"), 0L); + + assertEquals(AgentStateStore.UNVERSIONED, newVersion); + } + + @Test + void saveIfVersionReturnsNewVersionOnCasSuccess() { + // findOneAndUpdate returns doc with incremented version + Document result = + new Document("_id", "anon:session") + .append("key", Document.parse("{\"value\":\"updated\"}")) + .append("_version_key", 3L); + when(collection.findOneAndUpdate( + any(Bson.class), any(Bson.class), any(FindOneAndUpdateOptions.class))) + .thenReturn(result); + + long newVersion = + store.saveIfVersion("user", "session", "key", new TestState("updated"), 2L); + + assertEquals(3L, newVersion); + } + + @Test + void saveIfVersionReturnsUnversionedWhenCasFails() { + // findOneAndUpdate returns null -> version mismatch + when(collection.findOneAndUpdate( + any(Bson.class), any(Bson.class), any(FindOneAndUpdateOptions.class))) + .thenReturn(null); + + long newVersion = store.saveIfVersion("user", "session", "key", new TestState("v"), 99L); + + assertEquals(AgentStateStore.UNVERSIONED, newVersion); + } + @Test void existsReturnsFalseWhenNoDocument() { when(findIterable.first()).thenReturn(null); @@ -241,7 +340,7 @@ void listSessionIdsReturnsEmptyWhenNone() { @Test @SuppressWarnings("unchecked") void listSessionIdsReturnsIds() { - java.util.ArrayList ids = new java.util.ArrayList<>(List.of("s1", "s2")); + ArrayList ids = new ArrayList<>(List.of("s1", "s2")); when(distinctIterable.into(any())).thenReturn(ids); Set result = store.listSessionIds("user"); @@ -288,6 +387,6 @@ void rejectsBlankKey() { @Test void closeWithExternalClientDoesNotCloseClient() { store.close(); - verify(mongoClient, org.mockito.Mockito.never()).close(); + verify(mongoClient, Mockito.never()).close(); } } diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/store/MongoBaseStoreTest.java b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/store/MongoBaseStoreTest.java index 9bd69daa10..cb6833b382 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/store/MongoBaseStoreTest.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/store/MongoBaseStoreTest.java @@ -16,25 +16,32 @@ package io.agentscope.extensions.mongodb.store; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import com.mongodb.MongoWriteException; +import com.mongodb.ServerAddress; +import com.mongodb.WriteError; import com.mongodb.client.FindIterable; import com.mongodb.client.MongoCollection; import com.mongodb.client.MongoDatabase; +import com.mongodb.client.model.FindOneAndUpdateOptions; import com.mongodb.client.result.DeleteResult; import com.mongodb.client.result.UpdateResult; import io.agentscope.harness.agent.filesystem.remote.store.StoreItem; import java.util.ArrayList; import java.util.List; import java.util.Map; +import org.bson.BsonDocument; import org.bson.Document; import org.bson.conversions.Bson; import org.junit.jupiter.api.AfterEach; @@ -125,10 +132,62 @@ void putStoresItem() { @Test void putIfVersionReturnsFalseWhenVersionMismatch() { - when(findIterable.first()).thenReturn(null); + // findOneAndUpdate returns null when version filter doesn't match + when(collection.findOneAndUpdate( + any(Bson.class), any(Bson.class), any(FindOneAndUpdateOptions.class))) + .thenReturn(null); + boolean result = store.putIfVersion(List.of("ns"), "key", Map.of("data", "v"), 5L); - // findOneAndUpdate returns null when filter doesn't match (no upsert for non-zero version) - // Actually the implementation returns null for non-zero expectedVersion when no match + + assertFalse(result); + } + + @Test + void putIfVersionZeroCreatesWhenAbsent() { + // insertOne succeeds -> new document created + boolean result = store.putIfVersion(List.of("ns"), "key", Map.of("data", "v"), 0L); + + assertTrue(result); + verify(collection).insertOne(any(Document.class)); + } + + @Test + void putIfVersionZeroReturnsFalseWhenAlreadyExists() { + // insertOne throws E11000 duplicate key -> document already exists + WriteError writeError = + new WriteError(11000, "E11000 duplicate key error", new BsonDocument()); + doThrow(new MongoWriteException(writeError, new ServerAddress())) + .when(collection) + .insertOne(any(Document.class)); + + boolean result = store.putIfVersion(List.of("ns"), "key", Map.of("data", "v"), 0L); + + assertFalse(result); + } + + @Test + void putIfVersionZeroPropagatesNonDuplicateKeyError() { + // insertOne throws a non-duplicate-key error -> should propagate + WriteError writeError = new WriteError(12345, "some other error", new BsonDocument()); + doThrow(new MongoWriteException(writeError, new ServerAddress())) + .when(collection) + .insertOne(any(Document.class)); + + assertThrows( + MongoWriteException.class, + () -> store.putIfVersion(List.of("ns"), "key", Map.of("data", "v"), 0L)); + } + + @Test + void putIfVersionSuccessWhenVersionMatches() { + // findOneAndUpdate returns doc with new version = expectedVersion + 1 + when(collection.findOneAndUpdate( + any(Bson.class), any(Bson.class), any(FindOneAndUpdateOptions.class))) + .thenReturn(new Document("_id", "ns\0key").append("version", 3L)); + + boolean result = store.putIfVersion(List.of("ns"), "key", Map.of("data", "v"), 2L); + + assertTrue(result); } @Test From 03fd925c4f82b21e110a86d4731b8507c50e2a57 Mon Sep 17 00:00:00 2001 From: LiangshouX Date: Mon, 10 Aug 2026 01:47:11 +0800 Subject: [PATCH 3/8] fix(extensions-mongodb): remove redundant index, add null guards and simplify CAS - Remove duplicate ascending index on createdAt in MongoRemoteSnapshotClient; the TTL index already provides the same sorting capability - Add Objects.requireNonNull guards to MongoBaseStore constructor for database and collectionName parameters - Simplify MongoBaseStore.putIfVersion() return: findOneAndUpdate with version filter already guarantees the result matches expectedVersion + 1 --- .../mongodb/snapshot/MongoRemoteSnapshotClient.java | 6 ++---- .../extensions/mongodb/store/MongoBaseStore.java | 9 ++++----- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClient.java b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClient.java index 8e1e7676cb..57bbbebab3 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClient.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClient.java @@ -19,10 +19,10 @@ import com.mongodb.client.MongoCollection; import com.mongodb.client.MongoDatabase; import com.mongodb.client.model.Filters; +import com.mongodb.client.model.IndexOptions; import com.mongodb.client.model.Indexes; import com.mongodb.client.model.Projections; import com.mongodb.client.model.ReplaceOptions; -import com.mongodb.client.model.IndexOptions; import io.agentscope.harness.agent.sandbox.snapshot.RemoteSnapshotClient; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; @@ -73,11 +73,9 @@ public MongoRemoteSnapshotClient( private void initSchema() { try { - collection.createIndex(Indexes.ascending(FIELD_CREATED_AT)); collection.createIndex( Indexes.ascending(FIELD_CREATED_AT), - new IndexOptions() - .expireAfter(7 * 24 * 3600L, TimeUnit.SECONDS)); + new IndexOptions().expireAfter(7 * 24 * 3600L, TimeUnit.SECONDS)); } catch (Exception e) { log.warn( "Failed to initialize snapshot collection index '{}': {}", diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/store/MongoBaseStore.java b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/store/MongoBaseStore.java index f833b46ae4..e32e60f3d3 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/store/MongoBaseStore.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/store/MongoBaseStore.java @@ -35,6 +35,7 @@ import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.stream.Collectors; import org.bson.Document; import org.bson.conversions.Bson; @@ -82,6 +83,8 @@ public MongoBaseStore(MongoDatabase database, String collectionName) { */ public MongoBaseStore( MongoDatabase database, String collectionName, ObjectMapper objectMapper) { + Objects.requireNonNull(database, "database"); + Objects.requireNonNull(collectionName, "collectionName"); this.collection = database.getCollection(collectionName); this.objectMapper = objectMapper; ensureIndexes(); @@ -159,11 +162,7 @@ public boolean putIfVersion( new FindOneAndUpdateOptions().returnDocument(ReturnDocument.AFTER)); } - if (result == null) { - return false; - } - Long newVersion = result.getLong(FIELD_VERSION); - return newVersion != null && newVersion == expectedVersion + 1; + return result != null; } @Override From 78107bd974ca948e445b9ecc6c0a4aceef24ce3c Mon Sep 17 00:00:00 2001 From: LiangshouX Date: Thu, 13 Aug 2026 22:39:52 +0800 Subject: [PATCH 4/8] fix(extensions-mongodb): fix versioning, TTL and index upgrade bugs in MongoAgentStateStore - save() now increments version on each call (Updates.inc) - saveIfVersion(UNVERSIONED) reads version directly from MongoDB document instead of deserializing through State.class interface - saveIfVersion(0) catches MongoCommandException in addition to MongoWriteException for DuplicateKey errors from findOneAndUpdate - TTL index changed from expireAfterSeconds=0 to 30 days to prevent silent data loss by MongoDB TTL monitor - ensureIndexes gracefully handles IndexOptionsConflict (error 85) when upgrading from old TTL index parameters test(extensions-mongodb): add contract tests for BaseStore, AgentStateStore and index lifecycle - MongoBaseStoreContractTest (6 tests): CRUD, CAS, idempotent delete, search - MongoAgentStateStoreContractTest (6 tests): versioning, CAS concurrency - MongoIndexLifecycleContractTest (7 tests): index parameters, TTL values, upgrade from old TTL=0 without IndexOptionsConflict --- .../mongodb/state/MongoAgentStateStore.java | 43 ++- .../MongoIndexLifecycleContractTest.java | 274 ++++++++++++++++++ .../MongoAgentStateStoreContractTest.java | 218 ++++++++++++++ .../store/MongoBaseStoreContractTest.java | 172 +++++++++++ 4 files changed, 702 insertions(+), 5 deletions(-) create mode 100644 agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/MongoIndexLifecycleContractTest.java create mode 100644 agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreContractTest.java create mode 100644 agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/store/MongoBaseStoreContractTest.java diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStore.java b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStore.java index 7510a7cadf..e56ab0a4df 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStore.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStore.java @@ -17,6 +17,7 @@ import com.mongodb.ConnectionString; import com.mongodb.MongoClientSettings; +import com.mongodb.MongoCommandException; import com.mongodb.MongoWriteException; import com.mongodb.client.MongoClient; import com.mongodb.client.MongoClients; @@ -122,9 +123,24 @@ private void ensureIndexes() { collection.createIndex( Indexes.compoundIndex( Indexes.ascending(FIELD_USER_ID), Indexes.ascending(FIELD_SESSION_ID))); - collection.createIndex( - Indexes.ascending(FIELD_UPDATED_AT), - new IndexOptions().expireAfter(0L, TimeUnit.SECONDS).sparse(true)); + + String ttlIndexName = FIELD_UPDATED_AT + "_1"; + long ttlSeconds = 30L * 24 * 3600; + try { + collection.createIndex( + Indexes.ascending(FIELD_UPDATED_AT), + new IndexOptions().expireAfter(ttlSeconds, TimeUnit.SECONDS).sparse(true)); + } catch (MongoCommandException e) { + // IndexOptionsConflict + if (e.getErrorCode() == 85) { + collection.dropIndex(ttlIndexName); + collection.createIndex( + Indexes.ascending(FIELD_UPDATED_AT), + new IndexOptions().expireAfter(ttlSeconds, TimeUnit.SECONDS).sparse(true)); + } else { + throw e; + } + } } // ────────────────── Single Value CRUD ────────────────── @@ -133,10 +149,12 @@ private void ensureIndexes() { public void save(String userId, String sessionId, String key, State value) { validateKey(key); String slotId = slotId(userId, sessionId); + String versionField = VERSION_PREFIX + key; String json = JsonUtils.getJsonCodec().toJson(value); Bson setFields = Updates.combine( Updates.set(key, Document.parse(json)), + Updates.inc(versionField, 1L), Updates.set(FIELD_UPDATED_AT, new Date())); Bson setOnInsert = Updates.combine( @@ -283,8 +301,18 @@ public long saveIfVersion( validateKey(key); if (expectedVersion == UNVERSIONED) { save(userId, sessionId, key, value); - VersionedState after = getVersioned(userId, sessionId, key, State.class); - return after.version(); + String slotId = slotId(userId, sessionId); + String versionField = VERSION_PREFIX + key; + Document doc = + collection + .find(Filters.eq(slotId)) + .projection(Projections.include(versionField)) + .first(); + if (doc == null) { + return UNVERSIONED; + } + Long v = doc.getLong(versionField); + return v != null ? v : 0L; } String slotId = slotId(userId, sessionId); @@ -318,6 +346,11 @@ public long saveIfVersion( return UNVERSIONED; } throw e; + } catch (MongoCommandException e) { + if (e.getErrorCode() == 11000) { + return UNVERSIONED; + } + throw e; } } diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/MongoIndexLifecycleContractTest.java b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/MongoIndexLifecycleContractTest.java new file mode 100644 index 0000000000..ad18d22faf --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/MongoIndexLifecycleContractTest.java @@ -0,0 +1,274 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.extensions.mongodb; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; + +import com.mongodb.client.MongoClient; +import com.mongodb.client.MongoClients; +import com.mongodb.client.MongoCollection; +import com.mongodb.client.MongoDatabase; +import io.agentscope.extensions.mongodb.sandbox.MongoSandboxExecutionGuard; +import io.agentscope.extensions.mongodb.snapshot.MongoRemoteSnapshotClient; +import io.agentscope.extensions.mongodb.state.MongoAgentStateStore; +import io.agentscope.extensions.mongodb.store.MongoBaseStore; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import org.bson.Document; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.Assumptions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.MethodOrderer; +import org.junit.jupiter.api.Order; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestMethodOrder; + +/** + * Verifies that all MongoDB collections created by the extension have the correct indexes with + * expected parameters (TTL values, sparse flags, compound keys). + * + *

    This covers a class of bugs invisible to unit and contract tests: wrong index parameters that + * cause silent data loss (TTL=0) or startup failures on upgrade (IndexOptionsConflict). + * + *

    Requires a local MongoDB at {@code localhost:27017}. Skipped in CI. + */ +@DisplayName("Index lifecycle — MongoDB") +@TestMethodOrder(MethodOrderer.OrderAnnotation.class) +class MongoIndexLifecycleContractTest { + + private static final long THIRTY_DAYS_SECONDS = 30L * 24 * 3600; + private static final long SEVEN_DAYS_SECONDS = 7L * 24 * 3600; + + private static MongoClient client; + private static String dbName; + + @BeforeAll + static void connect() { + try { + client = MongoClients.create("mongodb://localhost:27017"); + client.getDatabase("ping").runCommand(new Document("ping", 1)); + } catch (Exception e) { + Assumptions.abort("MongoDB not available: " + e.getMessage()); + } + dbName = "test_idx_lifecycle_" + System.currentTimeMillis(); + } + + @AfterAll + static void disconnect() { + if (client != null) { + client.getDatabase(dbName).drop(); + client.close(); + } + } + + // ────────────────── AgentStateStore indexes ────────────────── + + @Test + @Order(1) + @DisplayName("AgentStateStore: compound index on (user_id, session_id) exists") + void agentStateStore_compoundIndex() { + MongoAgentStateStore.builder() + .mongoClient(client) + .databaseName(dbName) + .collectionName("idx_sessions") + .build(); + + Map indexes = indexMap(dbName, "idx_sessions"); + + Document compound = + indexes.values().stream() + .filter( + i -> { + Object key = i.get("key"); + return key instanceof Document d + && d.containsKey("user_id") + && d.containsKey("session_id"); + }) + .findFirst() + .orElse(null); + + assertNotNull(compound, "Compound index (user_id, session_id) must exist"); + } + + @Test + @Order(2) + @DisplayName("AgentStateStore: TTL index on _updated_at with 30-day expiry and sparse") + void agentStateStore_ttlIndex_30days() { + Map indexes = indexMap(dbName, "idx_sessions"); + + Document ttlIndex = indexes.get("_updated_at_1"); + assertNotNull(ttlIndex, "TTL index '_updated_at_1' must exist"); + assertEquals( + THIRTY_DAYS_SECONDS, + ((Number) ttlIndex.get("expireAfterSeconds")).longValue(), + "TTL must be 30 days (2592000s), not 0"); + assertEquals(true, ttlIndex.getBoolean("sparse"), "TTL index must be sparse"); + } + + @Test + @Order(3) + @DisplayName("AgentStateStore: upgrade from old TTL=0 index does not throw") + void agentStateStore_ttlUpgrade_fromZero() { + String upgradeDb = "test_idx_upgrade_" + System.currentTimeMillis(); + String collName = "upgrade_sessions"; + + // Phase 1: simulate old code — create TTL index with expireAfterSeconds=0 + MongoDatabase upgradeDbRef = client.getDatabase(upgradeDb); + upgradeDbRef + .getCollection(collName) + .createIndex( + new org.bson.Document("_updated_at", 1), + new com.mongodb.client.model.IndexOptions() + .expireAfter(0L, java.util.concurrent.TimeUnit.SECONDS) + .sparse(true)); + + // Phase 2: new code constructor runs ensureIndexes() — must not throw error 85 + MongoAgentStateStore.builder() + .mongoClient(client) + .databaseName(upgradeDb) + .collectionName(collName) + .build(); + + // Phase 3: verify index was upgraded to 30 days + Map indexes = indexMap(upgradeDb, collName); + Document ttlIndex = indexes.get("_updated_at_1"); + assertNotNull(ttlIndex); + assertEquals( + THIRTY_DAYS_SECONDS, + ((Number) ttlIndex.get("expireAfterSeconds")).longValue(), + "After upgrade, TTL must be 30 days"); + + // Cleanup + upgradeDbRef.drop(); + } + + // ────────────────── BaseStore indexes ────────────────── + + @Test + @Order(4) + @DisplayName("BaseStore: index on namespace exists") + void baseStore_namespaceIndex() { + new MongoBaseStore(client.getDatabase(dbName), "idx_base"); + + Map indexes = indexMap(dbName, "idx_base"); + + Document nsIndex = + indexes.values().stream() + .filter( + i -> { + Object key = i.get("key"); + return key instanceof Document d + && d.containsKey("namespace") + && d.size() == 1; + }) + .findFirst() + .orElse(null); + + assertNotNull(nsIndex, "Index on 'namespace' must exist"); + } + + @Test + @Order(5) + @DisplayName("BaseStore: compound index on (namespace, key) exists") + void baseStore_compoundIndex() { + Map indexes = indexMap(dbName, "idx_base"); + + Document compound = + indexes.values().stream() + .filter( + i -> { + Object key = i.get("key"); + return key instanceof Document d + && d.containsKey("namespace") + && d.containsKey("key") + && d.size() == 2; + }) + .findFirst() + .orElse(null); + + assertNotNull(compound, "Compound index (namespace, key) must exist"); + } + + // ────────────────── SandboxExecutionGuard indexes ────────────────── + + @Test + @Order(6) + @DisplayName("SandboxExecutionGuard: TTL index on expiresAt with immediate expiry (0s)") + void sandboxGuard_ttlIndex_immediate() { + MongoSandboxExecutionGuard.builder(client) + .databaseName(dbName) + .collectionName("idx_locks") + .build(); + + Map indexes = indexMap(dbName, "idx_locks"); + + Document ttlIndex = + indexes.values().stream() + .filter( + i -> { + Object key = i.get("key"); + return key instanceof Document d && d.containsKey("expiresAt"); + }) + .findFirst() + .orElse(null); + + assertNotNull(ttlIndex, "TTL index on 'expiresAt' must exist"); + assertEquals( + 0L, + ((Number) ttlIndex.get("expireAfterSeconds")).longValue(), + "Lock TTL must be 0 (immediate expiry after expiresAt)"); + } + + // ────────────────── RemoteSnapshotClient indexes ────────────────── + + @Test + @Order(7) + @DisplayName("RemoteSnapshotClient: TTL index on createdAt with 7-day expiry") + void snapshotClient_ttlIndex_7days() { + new MongoRemoteSnapshotClient(client, dbName, "idx_snapshots", true); + + Map indexes = indexMap(dbName, "idx_snapshots"); + + Document ttlIndex = + indexes.values().stream() + .filter( + i -> { + Object key = i.get("key"); + return key instanceof Document d && d.containsKey("createdAt"); + }) + .findFirst() + .orElse(null); + + assertNotNull(ttlIndex, "TTL index on 'createdAt' must exist"); + assertEquals( + SEVEN_DAYS_SECONDS, + ((Number) ttlIndex.get("expireAfterSeconds")).longValue(), + "Snapshot TTL must be 7 days (604800s)"); + } + + // ────────────────── Helpers ────────────────── + + private static Map indexMap(String database, String collection) { + MongoCollection coll = client.getDatabase(database).getCollection(collection); + List idxDocs = new ArrayList<>(); + coll.listIndexes().into(idxDocs); + return idxDocs.stream().collect(Collectors.toMap(d -> d.getString("name"), d -> d)); + } +} diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreContractTest.java b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreContractTest.java new file mode 100644 index 0000000000..c6b8d3f3b9 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreContractTest.java @@ -0,0 +1,218 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.extensions.mongodb.state; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.mongodb.client.MongoClient; +import com.mongodb.client.MongoClients; +import io.agentscope.core.state.AgentStateStore; +import io.agentscope.core.state.State; +import io.agentscope.core.state.VersionedState; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicInteger; +import org.bson.Document; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assumptions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.MethodOrderer; +import org.junit.jupiter.api.Order; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestMethodOrder; + +/** + * Contract tests for optimistic concurrency on {@link AgentStateStore} against a real MongoDB. + * + *

    Mirrors the canonical contract defined in {@code AgentStateStoreVersioningContractTest} + * (agentscope-core). Skipped automatically when MongoDB is not reachable at {@code + * localhost:27017}. + */ +@DisplayName("AgentStateStore versioning contract — MongoDB") +@TestMethodOrder(MethodOrderer.OrderAnnotation.class) +class MongoAgentStateStoreContractTest { + + private static final String USER = "contract-user"; + private static final String SESSION = "contract-session"; + + private static MongoClient mongoClient; + private static String dbName; + + private AgentStateStore store; + + @BeforeAll + static void connectMongo() { + try { + mongoClient = MongoClients.create("mongodb://localhost:27017"); + mongoClient.getDatabase("ping").runCommand(new Document("ping", 1)); + } catch (Exception e) { + Assumptions.abort("MongoDB not available: " + e.getMessage()); + } + dbName = "test_state_contract_" + System.currentTimeMillis(); + } + + @AfterAll + static void disconnectMongo() { + if (mongoClient != null) { + mongoClient.getDatabase(dbName).drop(); + mongoClient.close(); + } + } + + @BeforeEach + void setUp() { + store = + MongoAgentStateStore.builder() + .mongoClient(mongoClient) + .databaseName(dbName) + .collectionName("test_sessions") + .build(); + } + + @AfterEach + void cleanSession() { + store.delete(USER, SESSION); + } + + @Test + @Order(1) + @DisplayName("supportsVersioning is true") + void supportsVersioning() { + assertTrue(store.supportsVersioning()); + } + + @Test + @Order(2) + @DisplayName("getVersioned on absent key returns null value and version 0") + void getVersioned_absent_returnsVersionZero() { + VersionedState versioned = + store.getVersioned(USER, SESSION, "agent_state", TestState.class); + + assertNull(versioned.value()); + assertEquals(0L, versioned.version()); + } + + @Test + @Order(3) + @DisplayName("saveIfVersion with expectedVersion 0 creates if absent") + void saveIfVersion_createIfAbsent() { + TestState initial = new TestState("created"); + + long version = store.saveIfVersion(USER, SESSION, "agent_state", initial, 0L); + assertEquals(1L, version); + + VersionedState loaded = + store.getVersioned(USER, SESSION, "agent_state", TestState.class); + assertEquals("created", loaded.value().value()); + assertEquals(1L, loaded.version()); + + long conflict = + store.saveIfVersion(USER, SESSION, "agent_state", new TestState("lost"), 0L); + assertEquals(AgentStateStore.UNVERSIONED, conflict); + assertEquals( + "created", store.get(USER, SESSION, "agent_state", TestState.class).get().value()); + } + + @Test + @Order(4) + @DisplayName("saveIfVersion with UNVERSIONED unconditionally overwrites and bumps version") + void saveIfVersion_unconditionalOverwrite() { + store.save(USER, SESSION, "agent_state", new TestState("v1")); + VersionedState afterFirst = + store.getVersioned(USER, SESSION, "agent_state", TestState.class); + assertEquals(1L, afterFirst.version()); + + long newVersion = + store.saveIfVersion( + USER, + SESSION, + "agent_state", + new TestState("v2"), + AgentStateStore.UNVERSIONED); + assertEquals(2L, newVersion); + + VersionedState loaded = + store.getVersioned(USER, SESSION, "agent_state", TestState.class); + assertEquals("v2", loaded.value().value()); + assertEquals(2L, loaded.version()); + } + + @Test + @Order(5) + @DisplayName("plain save bumps version") + void plainSave_bumpsVersion() { + store.save(USER, SESSION, "agent_state", new TestState("one")); + assertEquals( + 1L, store.getVersioned(USER, SESSION, "agent_state", TestState.class).version()); + + store.save(USER, SESSION, "agent_state", new TestState("two")); + assertEquals( + 2L, store.getVersioned(USER, SESSION, "agent_state", TestState.class).version()); + } + + @Test + @Order(6) + @DisplayName("concurrent writers with same expected version: only one succeeds") + void concurrentWriters_onlyOneSucceeds() throws InterruptedException { + store.saveIfVersion(USER, SESSION, "agent_state", new TestState("baseline"), 0L); + long observed = store.getVersioned(USER, SESSION, "agent_state", TestState.class).version(); + assertEquals(1L, observed); + + CountDownLatch ready = new CountDownLatch(2); + CountDownLatch start = new CountDownLatch(1); + AtomicInteger successes = new AtomicInteger(); + ExecutorService pool = Executors.newFixedThreadPool(2); + + Runnable attempt = + () -> { + ready.countDown(); + try { + start.await(); + long result = + store.saveIfVersion( + USER, + SESSION, + "agent_state", + new TestState("winner"), + observed); + if (result != AgentStateStore.UNVERSIONED) { + successes.incrementAndGet(); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }; + + pool.submit(attempt); + pool.submit(attempt); + ready.await(); + start.countDown(); + pool.shutdown(); + assertTrue(pool.awaitTermination(5, java.util.concurrent.TimeUnit.SECONDS)); + + assertEquals(1, successes.get()); + assertEquals( + 2L, store.getVersioned(USER, SESSION, "agent_state", TestState.class).version()); + } + + record TestState(String value) implements State {} +} diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/store/MongoBaseStoreContractTest.java b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/store/MongoBaseStoreContractTest.java new file mode 100644 index 0000000000..47a23a342e --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/store/MongoBaseStoreContractTest.java @@ -0,0 +1,172 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.extensions.mongodb.store; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.mongodb.client.MongoClient; +import com.mongodb.client.MongoClients; +import com.mongodb.client.MongoDatabase; +import io.agentscope.harness.agent.filesystem.remote.store.BaseStore; +import io.agentscope.harness.agent.filesystem.remote.store.StoreItem; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import org.bson.Document; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.Assumptions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.MethodOrderer; +import org.junit.jupiter.api.Order; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestMethodOrder; + +/** + * Contract tests for {@link BaseStore} semantics against a real MongoDB instance. + * + *

    Mirrors the canonical contract defined in {@code BaseStoreContractTest} (agentscope-harness). + * Skipped automatically when MongoDB is not reachable at {@code localhost:27017}. + * + *

    Search semantics note: {@code MongoBaseStore.search()} matches by exact namespace + * (not prefix), so {@code search(["a"])} does NOT return items stored under child namespaces such + * as {@code ["a","b"]}. This differs from {@code InMemoryStore} which uses prefix matching. The + * search test below validates MongoDB's exact-match behavior. + */ +@DisplayName("BaseStore contract — MongoDB") +@TestMethodOrder(MethodOrderer.OrderAnnotation.class) +class MongoBaseStoreContractTest { + + private static MongoClient client; + private static MongoDatabase db; + private static BaseStore store; + + @BeforeAll + static void setUp() { + try { + client = MongoClients.create("mongodb://localhost:27017"); + client.getDatabase("ping").runCommand(new Document("ping", 1)); + } catch (Exception e) { + Assumptions.abort("MongoDB not available: " + e.getMessage()); + } + db = client.getDatabase("test_base_contract_" + System.currentTimeMillis()); + store = new MongoBaseStore(db, "test_base"); + } + + @AfterAll + static void tearDown() { + if (db != null) { + db.drop(); + } + if (client != null) { + client.close(); + } + } + + @Test + @Order(1) + @DisplayName("put + get round-trip: version starts at 1") + void putGetRoundTrip_versionStartsAtOne() { + List ns = List.of("ws"); + + store.put(ns, "MEMORY.md", Map.of("content", "hello")); + StoreItem item = store.get(ns, "MEMORY.md"); + + assertNotNull(item); + assertEquals("MEMORY.md", item.key()); + assertEquals("hello", item.value().get("content")); + assertEquals(1L, item.version()); + } + + @Test + @Order(2) + @DisplayName("put increments version on each call") + void put_incrementsVersion() { + List ns = List.of("ver"); + + store.put(ns, "k", Map.of("v", 1)); + assertEquals(1L, store.get(ns, "k").version()); + + store.put(ns, "k", Map.of("v", 2)); + assertEquals(2L, store.get(ns, "k").version()); + assertEquals(2, store.get(ns, "k").value().get("v")); + } + + @Test + @Order(3) + @DisplayName("putIfVersion: CAS success and conflict") + void putIfVersion_successAndConflict() { + List ns = List.of("cas"); + + store.put(ns, "k", Map.of("v", 1)); + long v1 = store.get(ns, "k").version(); + + assertTrue(store.putIfVersion(ns, "k", Map.of("v", 2), v1)); + assertEquals(2, store.get(ns, "k").value().get("v")); + assertEquals(v1 + 1, store.get(ns, "k").version()); + + assertFalse(store.putIfVersion(ns, "k", Map.of("v", 3), v1)); + assertEquals(2, store.get(ns, "k").value().get("v")); + } + + @Test + @Order(4) + @DisplayName("putIfVersion(0): create-if-absent") + void putIfVersionZero_createIfAbsent() { + List ns = List.of("create"); + + assertTrue(store.putIfVersion(ns, "k", Map.of("v", 1), 0L)); + assertEquals(1L, store.get(ns, "k").version()); + + assertFalse(store.putIfVersion(ns, "k", Map.of("v", 2), 0L)); + assertEquals(1, store.get(ns, "k").value().get("v")); + } + + @Test + @Order(5) + @DisplayName("delete is idempotent") + void delete_isIdempotent() { + List ns = List.of("del"); + + store.put(ns, "k", Map.of("v", 1)); + store.delete(ns, "k"); + assertNull(store.get(ns, "k")); + + store.delete(ns, "k"); + assertNull(store.get(ns, "k")); + } + + @Test + @Order(6) + @DisplayName("search: exact namespace match (MongoDB behavior)") + void search_exactNamespaceMatch() { + store.put(List.of("s"), "inNs", Map.of("where", "s")); + store.put(List.of("s", "t"), "inChild", Map.of("where", "s/t")); + + List found = store.search(List.of("s"), 100, 0); + Set keys = found.stream().map(StoreItem::key).collect(Collectors.toSet()); + + // MongoBaseStore uses exact namespace match (not prefix). + // Only items directly under ["s"] are returned, not ["s","t"]. + assertEquals(Set.of("inNs"), keys); + assertEquals(1, found.size()); + } +} From 8847e937f6d41c8c515b51787069dff3fbafb913 Mon Sep 17 00:00:00 2001 From: LiangshouX Date: Fri, 14 Aug 2026 08:37:58 +0800 Subject: [PATCH 5/8] fix(extensions-mongodb): fix contract test @AfterAll lifecycle for CI safety MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Assumptions.abort() in @BeforeAll does not prevent @AfterAll from running. When MongoDB is unreachable, @AfterAll calls client.getDatabase(dbName).drop() which blocks for 60s then throws MongoTimeoutException — causing CI failure. Fix: add `connected` flag, set only on successful ping. @AfterAll guards all MongoDB operations behind `if (connected)`. Also move dbName assignment before the try-catch to avoid null in disconnect. --- .../mongodb/MongoIndexLifecycleContractTest.java | 6 ++++-- .../state/MongoAgentStateStoreContractTest.java | 6 ++++-- .../mongodb/store/MongoBaseStoreContractTest.java | 14 +++++++++----- 3 files changed, 17 insertions(+), 9 deletions(-) diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/MongoIndexLifecycleContractTest.java b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/MongoIndexLifecycleContractTest.java index ad18d22faf..3ddf70f6e0 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/MongoIndexLifecycleContractTest.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/MongoIndexLifecycleContractTest.java @@ -58,21 +58,23 @@ class MongoIndexLifecycleContractTest { private static MongoClient client; private static String dbName; + private static boolean connected; @BeforeAll static void connect() { + dbName = "test_idx_lifecycle_" + System.currentTimeMillis(); try { client = MongoClients.create("mongodb://localhost:27017"); client.getDatabase("ping").runCommand(new Document("ping", 1)); + connected = true; } catch (Exception e) { Assumptions.abort("MongoDB not available: " + e.getMessage()); } - dbName = "test_idx_lifecycle_" + System.currentTimeMillis(); } @AfterAll static void disconnect() { - if (client != null) { + if (connected && client != null) { client.getDatabase(dbName).drop(); client.close(); } diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreContractTest.java b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreContractTest.java index c6b8d3f3b9..7439c6dede 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreContractTest.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreContractTest.java @@ -56,23 +56,25 @@ class MongoAgentStateStoreContractTest { private static MongoClient mongoClient; private static String dbName; + private static boolean connected; private AgentStateStore store; @BeforeAll static void connectMongo() { + dbName = "test_state_contract_" + System.currentTimeMillis(); try { mongoClient = MongoClients.create("mongodb://localhost:27017"); mongoClient.getDatabase("ping").runCommand(new Document("ping", 1)); + connected = true; } catch (Exception e) { Assumptions.abort("MongoDB not available: " + e.getMessage()); } - dbName = "test_state_contract_" + System.currentTimeMillis(); } @AfterAll static void disconnectMongo() { - if (mongoClient != null) { + if (connected && mongoClient != null) { mongoClient.getDatabase(dbName).drop(); mongoClient.close(); } diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/store/MongoBaseStoreContractTest.java b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/store/MongoBaseStoreContractTest.java index 47a23a342e..858ead62f4 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/store/MongoBaseStoreContractTest.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/store/MongoBaseStoreContractTest.java @@ -58,12 +58,14 @@ class MongoBaseStoreContractTest { private static MongoClient client; private static MongoDatabase db; private static BaseStore store; + private static boolean connected; @BeforeAll static void setUp() { try { client = MongoClients.create("mongodb://localhost:27017"); client.getDatabase("ping").runCommand(new Document("ping", 1)); + connected = true; } catch (Exception e) { Assumptions.abort("MongoDB not available: " + e.getMessage()); } @@ -73,11 +75,13 @@ static void setUp() { @AfterAll static void tearDown() { - if (db != null) { - db.drop(); - } - if (client != null) { - client.close(); + if (connected) { + if (db != null) { + db.drop(); + } + if (client != null) { + client.close(); + } } } From 38a95b16ea41833b46bd1aa8e70e2ced534580e4 Mon Sep 17 00:00:00 2001 From: LiangshouX Date: Tue, 25 Aug 2026 21:29:49 +0800 Subject: [PATCH 6/8] =?UTF-8?q?fix(extensions-mongodb):=20address=20review?= =?UTF-8?q?=20findings=20=E2=80=94=20session=20visibility=20blocker,=20TTL?= =?UTF-8?q?=20alignment,=20lease=20renewal?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address all 7 findings from the code review. Blocker: - save(List): add $setOnInsert(user_id/session_id) to the incremental-append branch. A brand-new session whose first write is a non-empty list landed in this branch and was upserted without the identity fields, making it permanently invisible to listSessionIds. Fix before release: - Align snapshot TTL with the session TTL (7d -> 30d) so a live session's snapshot is never reclaimed first; add the error-85 drop/recreate migration so any future TTL adjustment takes effect instead of silently keeping the old index options. - MongoSandboxExecutionGuard: split leaseTtl (lock document lifetime, default 30m) from lockTimeout (acquisition wait), and renew the lease every leaseTtl/3 via a background watchdog while the lease is open, so a long execution no longer loses its own lock to TTL reclamation. Minor: - download(): treat a document without a data field as FileNotFoundException instead of NPE (parity with the Redis implementation). - fromConnectionString: fix javadoc — the store owns and closes the client. - Document that save(List) is intentionally non-atomic (single-writer guarantee comes from SandboxExecutionGuard) and that validateKey covers top-level keys only. Tests: 112 -> 124. New regression coverage: append-branch $setOnInsert (unit + contract), first-save-as-list visibility in listSessionIds (contract), snapshot 30-day TTL and index migration (contract + unit), lease TTL separation, watchdog renewal, close-stops-renewal and lost-lock tolerance (unit). All 124 pass locally; contract tests run against a real MongoDB. --- .../mongodb/MongoDistributedStore.java | 2 +- .../sandbox/MongoSandboxExecutionGuard.java | 119 +++++++++++++- .../snapshot/MongoRemoteSnapshotClient.java | 45 ++++- .../mongodb/state/MongoAgentStateStore.java | 24 ++- .../MongoIndexLifecycleContractTest.java | 50 +++++- .../MongoSandboxExecutionGuardTest.java | 154 ++++++++++++++++++ .../MongoRemoteSnapshotClientTest.java | 30 ++++ .../MongoAgentStateStoreContractTest.java | 36 ++++ .../state/MongoAgentStateStoreTest.java | 56 +++++++ 9 files changed, 497 insertions(+), 19 deletions(-) diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/MongoDistributedStore.java b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/MongoDistributedStore.java index 48ae2e8d42..8ef3499fed 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/MongoDistributedStore.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/MongoDistributedStore.java @@ -108,7 +108,7 @@ public static MongoDistributedStore create(MongoClient mongoClient, String datab /** * Creates a MongoDB distributed store from a connection string. A new {@link MongoClient} is - * created internally. The caller is responsible for closing the client when done. + * created internally and owned by the store; {@link #close()} will close it. * * @param connectionString the MongoDB connection string * @return a new MongoDB distributed store diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/sandbox/MongoSandboxExecutionGuard.java b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/sandbox/MongoSandboxExecutionGuard.java index b20b2aa267..0654de16c9 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/sandbox/MongoSandboxExecutionGuard.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/sandbox/MongoSandboxExecutionGuard.java @@ -25,6 +25,7 @@ import com.mongodb.client.model.Indexes; import com.mongodb.client.model.ReturnDocument; import com.mongodb.client.model.Updates; +import com.mongodb.client.result.UpdateResult; import io.agentscope.harness.agent.sandbox.SandboxExecutionGuard; import io.agentscope.harness.agent.sandbox.SandboxIsolationKey; import io.agentscope.harness.agent.sandbox.SandboxLease; @@ -34,6 +35,9 @@ import java.time.Duration; import java.util.Date; import java.util.Objects; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; import org.bson.Document; import org.bson.conversions.Bson; @@ -57,9 +61,16 @@ * an expired lock. * * - *

    Acquisition polls until the lock is obtained or the timeout expires. The returned {@link - * SandboxLease} releases the lock on close, filtered by owner to avoid releasing another - * process's lock. + *

    Acquisition polls until the lock is obtained or the acquisition timeout ({@code + * lockTimeout}) expires. The lock document itself lives for {@code leaseTtl}; while the returned + * {@link SandboxLease} is open, a background watchdog renews the lease every {@code leaseTtl / 3} + * by pushing {@code expiresAt} forward, so an execution that outlives the initial TTL does not + * lose its own lock to reclamation. If renewal stops working (lock deleted or stolen), the + * watchdog only warns — like the Redis guard, TTL expiry remains a safety valve against permanent + * deadlock rather than a correctness guarantee. + * + *

    The lease releases the lock on close, filtered by owner to avoid releasing another process's + * lock. */ public final class MongoSandboxExecutionGuard implements SandboxExecutionGuard { @@ -72,15 +83,25 @@ public final class MongoSandboxExecutionGuard implements SandboxExecutionGuard { private final MongoCollection collection; private final long lockTimeoutMs; + private final long leaseTtlMs; private final long retryIntervalMs; private final String owner; + private final ScheduledExecutorService renewalExecutor; private MongoSandboxExecutionGuard(Builder builder) { MongoDatabase db = builder.mongoClient.getDatabase(builder.databaseName); this.collection = db.getCollection(builder.collectionName); this.lockTimeoutMs = builder.lockTimeout.toMillis(); + this.leaseTtlMs = builder.leaseTtl.toMillis(); this.retryIntervalMs = builder.retryInterval.toMillis(); this.owner = builder.owner; + this.renewalExecutor = + Executors.newSingleThreadScheduledExecutor( + runnable -> { + Thread thread = new Thread(runnable, "mongo-sandbox-lease-renewal"); + thread.setDaemon(true); + return thread; + }); ensureIndexes(); } @@ -102,7 +123,7 @@ public SandboxLease tryEnter(SandboxIsolationKey key) throws InterruptedExceptio long deadline = System.nanoTime() + Duration.ofMillis(lockTimeoutMs).toNanos(); while (true) { Date now = new Date(); - Date expiresAt = new Date(now.getTime() + lockTimeoutMs); + Date expiresAt = new Date(now.getTime() + leaseTtlMs); // Step 1: Try to insert a new lock document. This is atomic — only one // concurrent caller can succeed due to the _id unique index. @@ -113,7 +134,7 @@ public SandboxLease tryEnter(SandboxIsolationKey key) throws InterruptedExceptio try { collection.insertOne(lockDoc); log.debug("[sandbox-guard] Acquired MongoDB lock (insert): {}", lockId); - return new MongoLease(collection, lockId, owner); + return new MongoLease(collection, lockId, owner, leaseTtlMs, renewalExecutor); } catch (MongoWriteException e) { if (e.getError().getCode() != 11000) { throw new RuntimeException("Failed to acquire MongoDB lock: " + lockId, e); @@ -141,7 +162,7 @@ public SandboxLease tryEnter(SandboxIsolationKey key) throws InterruptedExceptio "[sandbox-guard] Acquired MongoDB lock (reclaimed from {}): {}", reclaimed.getString(FIELD_OWNER), lockId); - return new MongoLease(collection, lockId, owner); + return new MongoLease(collection, lockId, owner, leaseTtlMs, renewalExecutor); } // Lock exists and has not expired — held by someone else. @@ -197,15 +218,63 @@ private static final class MongoLease implements SandboxLease { private final MongoCollection collection; private final String lockId; private final String owner; - - MongoLease(MongoCollection collection, String lockId, String owner) { + private final long leaseTtlMs; + private final ScheduledFuture renewal; + + MongoLease( + MongoCollection collection, + String lockId, + String owner, + long leaseTtlMs, + ScheduledExecutorService renewalExecutor) { this.collection = collection; this.lockId = lockId; this.owner = owner; + this.leaseTtlMs = leaseTtlMs; + // Renew with 3x headroom before expiry so a single missed tick or a slow write + // does not cost the lease. + long renewalIntervalMs = Math.max(1L, leaseTtlMs / 3); + this.renewal = + renewalExecutor.scheduleAtFixedRate( + this::renew, + renewalIntervalMs, + renewalIntervalMs, + TimeUnit.MILLISECONDS); + } + + /** + * Watchdog tick: push {@code expiresAt} forward while the lease is open, so an + * execution longer than the initial lease TTL keeps holding its own lock. + */ + private void renew() { + try { + Date expiresAt = new Date(System.currentTimeMillis() + leaseTtlMs); + UpdateResult result = + collection.updateOne( + Filters.and(Filters.eq(lockId), Filters.eq(FIELD_OWNER, owner)), + Updates.set(FIELD_EXPIRES_AT, expiresAt)); + if (result.getMatchedCount() == 0) { + // Lock document gone or owned by someone else: it expired and was + // reclaimed before renewal succeeded. Warn only — like the Redis guard, + // TTL expiry is a safety valve, not a correctness guarantee. + log.warn( + "[sandbox-guard] MongoDB lock {} lost during execution (expired or" + + " reclaimed); continuing without the lock", + lockId); + } + } catch (Exception e) { + // Never let an exception escape a scheduled task, or the executor would + // silently stop renewing. + log.warn( + "[sandbox-guard] Failed to renew MongoDB lock {}: {}", + lockId, + e.getMessage()); + } } @Override public void close() { + renewal.cancel(false); try { // Only delete if we still own the lock — prevents releasing someone else's // lock when close() is called after lease expiry and re-acquisition by another @@ -229,11 +298,13 @@ public void close() { public static final class Builder { private static final Duration DEFAULT_RETRY_INTERVAL = Duration.ofMillis(500); + private static final Duration DEFAULT_LEASE_TTL = Duration.ofMinutes(30); private final MongoClient mongoClient; private String databaseName = "agentscope"; private String collectionName = DEFAULT_COLLECTION; private Duration lockTimeout = Duration.ofMinutes(30); + private Duration leaseTtl = DEFAULT_LEASE_TTL; private Duration retryInterval = DEFAULT_RETRY_INTERVAL; private String owner = "agentscope:" + ProcessHandle.current().pid(); @@ -251,6 +322,16 @@ public Builder collectionName(String collectionName) { return this; } + /** + * Sets the maximum time to wait for acquiring the lock before {@link #tryEnter} throws + * {@link InterruptedException}. This bounds queueing delay only; it does NOT bound the + * lock's lifetime — that is controlled by {@link #leaseTtl(Duration)}. + * + *

    Default: {@code 30 minutes}. + * + * @param timeout the acquisition timeout; must be positive + * @return this builder + */ public Builder lockTimeout(Duration timeout) { Objects.requireNonNull(timeout, "timeout"); if (timeout.isNegative() || timeout.isZero()) { @@ -260,6 +341,28 @@ public Builder lockTimeout(Duration timeout) { return this; } + /** + * Sets the lifetime of the lock document ({@code expiresAt}) written on acquisition. + * While the returned lease is open, a background watchdog renews the lease every {@code + * leaseTtl / 3}, so executions longer than the TTL keep holding their lock. The TTL only + * matters when the holder dies without releasing: after {@code leaseTtl} the lock becomes + * eligible for reclamation by the next caller (MongoDB's TTL monitor may take up to 60 + * additional seconds to delete the document). + * + *

    Default: {@code 30 minutes}. + * + * @param ttl the lease TTL; must be positive + * @return this builder + */ + public Builder leaseTtl(Duration ttl) { + Objects.requireNonNull(ttl, "ttl"); + if (ttl.isNegative() || ttl.isZero()) { + throw new IllegalArgumentException("ttl must be positive"); + } + this.leaseTtl = ttl; + return this; + } + /** * Sets the polling interval between lock acquisition attempts. * diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClient.java b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClient.java index 57bbbebab3..ae7be7971c 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClient.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClient.java @@ -15,6 +15,7 @@ */ package io.agentscope.extensions.mongodb.snapshot; +import com.mongodb.MongoCommandException; import com.mongodb.client.MongoClient; import com.mongodb.client.MongoCollection; import com.mongodb.client.MongoDatabase; @@ -42,6 +43,11 @@ * *

    Stores sandbox workspace tar archives as BSON Binary in a collection with documents of the * form {@code {_id: snapshotId, data: Binary, createdAt: Date}}. + * + *

    The collection carries a TTL index on {@code createdAt} with a 30-day expiry, aligned with + * the 30-day session TTL of {@code MongoAgentStateStore}. The snapshot of a session must not be + * reclaimed while the session itself is still alive; otherwise resuming the sandbox would fail + * with {@link FileNotFoundException} and appear as lost workspace data. */ public class MongoRemoteSnapshotClient implements RemoteSnapshotClient { @@ -50,6 +56,9 @@ public class MongoRemoteSnapshotClient implements RemoteSnapshotClient { private static final String DEFAULT_COLLECTION = "agentscope_snapshots"; private static final String FIELD_DATA = "data"; private static final String FIELD_CREATED_AT = "createdAt"; + // Aligned with the session TTL in MongoAgentStateStore (30 days on _updated_at): a snapshot + // must live at least as long as the session that can resume it. + private static final long SNAPSHOT_TTL_SECONDS = 30L * 24 * 3600; // MongoDB BSON document size limit is 16 MB; cap at 15 MB to leave headroom for // metadata. For larger snapshots, use GridFS (not yet implemented). private static final int MAX_SNAPSHOT_BYTES = 15 * 1024 * 1024; // 15 MB @@ -75,7 +84,21 @@ private void initSchema() { try { collection.createIndex( Indexes.ascending(FIELD_CREATED_AT), - new IndexOptions().expireAfter(7 * 24 * 3600L, TimeUnit.SECONDS)); + new IndexOptions().expireAfter(SNAPSHOT_TTL_SECONDS, TimeUnit.SECONDS)); + } catch (MongoCommandException e) { + // IndexOptionsConflict + if (e.getErrorCode() == 85) { + // A TTL index on createdAt exists with different options, and createIndex + // cannot change options in place. Drop and recreate so the TTL actually + // takes effect — otherwise this best-effort initialization would only warn + // and silently keep the stale TTL. This keeps future TTL adjustments safe. + migrateTtlIndex(); + } else { + log.warn( + "Failed to initialize snapshot collection index '{}': {}", + collection.getNamespace(), + e.getMessage()); + } } catch (Exception e) { log.warn( "Failed to initialize snapshot collection index '{}': {}", @@ -84,6 +107,20 @@ private void initSchema() { } } + private void migrateTtlIndex() { + try { + collection.dropIndex(FIELD_CREATED_AT + "_1"); + collection.createIndex( + Indexes.ascending(FIELD_CREATED_AT), + new IndexOptions().expireAfter(SNAPSHOT_TTL_SECONDS, TimeUnit.SECONDS)); + } catch (Exception e) { + log.warn( + "Failed to migrate snapshot TTL index '{}': {}", + collection.getNamespace(), + e.getMessage()); + } + } + @Override public void upload(String snapshotId, InputStream data) throws Exception { Objects.requireNonNull(snapshotId, "snapshotId"); @@ -106,6 +143,12 @@ public InputStream download(String snapshotId) throws Exception { throw new FileNotFoundException("Snapshot not found in MongoDB: " + snapshotId); } Binary binary = doc.get(FIELD_DATA, Binary.class); + if (binary == null) { + // Document exists but carries no data (e.g. corrupted or manually edited). Treat + // it the same as a missing snapshot instead of failing with an NPE. + throw new FileNotFoundException( + "Snapshot document has no data field in MongoDB: " + snapshotId); + } return new ByteArrayInputStream(binary.getData()); } diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStore.java b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStore.java index e56ab0a4df..324b96c6a0 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStore.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStore.java @@ -182,10 +182,12 @@ public Optional get( * Saves a list of state values with incremental-append optimization. * *

    Concurrency note: this method performs a read-then-write to decide between - * incremental append and full replacement. It is NOT atomic — concurrent calls for the same - * {@code (userId, sessionId, key)} may interleave reads and writes, causing lost appends or - * stale hash comparisons. Callers that require strict consistency should synchronize externally - * (e.g. via {@link io.agentscope.harness.agent.sandbox.SandboxExecutionGuard}). + * incremental append and full replacement, so it is intentionally NOT atomic: concurrent + * writers for the same {@code (userId, sessionId, key)} may interleave reads and writes, + * causing duplicated or lost appends. This matches the harness execution model, where a + * {@link io.agentscope.harness.agent.sandbox.SandboxExecutionGuard} serialises calls per + * isolation slot, so each session normally has at most one writer at a time. Callers that + * write the same list concurrently from outside the harness must synchronise externally. */ @Override public void save(String userId, String sessionId, String key, List values) { @@ -229,11 +231,18 @@ public void save(String userId, String sessionId, String key, List existingCount) { List newItems = values.subList(existingCount, values.size()); List newDocs = toDocumentList(newItems); + // $setOnInsert is required here as well: a brand-new session whose first write is + // a non-empty list lands in this append branch (needsFullRewrite returns false for + // an absent document), so without it the upserted document would carry no + // user_id/session_id fields and listSessionIds could never see it — not even after + // later full-rewrite saves, since $setOnInsert only fires on the initial insert. Bson update = Updates.combine( Updates.pushEach(listKey, newDocs), Updates.set(hashField, currentHash), - Updates.set(FIELD_UPDATED_AT, new Date())); + Updates.set(FIELD_UPDATED_AT, new Date()), + Updates.setOnInsert(FIELD_USER_ID, normalizeUser(userId)), + Updates.setOnInsert(FIELD_SESSION_ID, sessionId)); collection.updateOne(Filters.eq(slotId), update, upsert()); } else { // Hash unchanged but size decreased (elements removed) — force full rewrite. @@ -434,6 +443,11 @@ private static String slotId(String userId, String sessionId) { return normalizeUser(userId) + ":" + sessionId; } + /** + * Validates a top-level state key as a MongoDB field name. Nested field names inside values + * are deliberately not validated or rewritten: MongoDB 5.0+ stores them verbatim, and renaming + * them would break deserialization back into the original {@link State} type. + */ private static void validateKey(String key) { if (key == null || key.isBlank()) { throw new IllegalArgumentException("key must not be blank"); diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/MongoIndexLifecycleContractTest.java b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/MongoIndexLifecycleContractTest.java index 3ddf70f6e0..3918d81d4c 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/MongoIndexLifecycleContractTest.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/MongoIndexLifecycleContractTest.java @@ -242,8 +242,8 @@ void sandboxGuard_ttlIndex_immediate() { @Test @Order(7) - @DisplayName("RemoteSnapshotClient: TTL index on createdAt with 7-day expiry") - void snapshotClient_ttlIndex_7days() { + @DisplayName("RemoteSnapshotClient: TTL index on createdAt with 30-day expiry") + void snapshotClient_ttlIndex_30days() { new MongoRemoteSnapshotClient(client, dbName, "idx_snapshots", true); Map indexes = indexMap(dbName, "idx_snapshots"); @@ -260,9 +260,51 @@ void snapshotClient_ttlIndex_7days() { assertNotNull(ttlIndex, "TTL index on 'createdAt' must exist"); assertEquals( - SEVEN_DAYS_SECONDS, + THIRTY_DAYS_SECONDS, + ((Number) ttlIndex.get("expireAfterSeconds")).longValue(), + "Snapshot TTL must be 30 days (2592000s), aligned with the session TTL —" + + " a snapshot must not be reclaimed while its session is still alive"); + } + + @Test + @Order(8) + @DisplayName("RemoteSnapshotClient: upgrade from old 7-day TTL index does not throw") + void snapshotClient_ttlUpgrade_fromSevenDays() { + String upgradeDb = "test_idx_snap_upgrade_" + System.currentTimeMillis(); + String collName = "upgrade_snapshots"; + + // Phase 1: simulate old code — create TTL index with the previous 7-day expiry + MongoDatabase upgradeDbRef = client.getDatabase(upgradeDb); + upgradeDbRef + .getCollection(collName) + .createIndex( + new org.bson.Document("createdAt", 1), + new com.mongodb.client.model.IndexOptions() + .expireAfter( + SEVEN_DAYS_SECONDS, java.util.concurrent.TimeUnit.SECONDS)); + + // Phase 2: new code constructor runs initSchema() — must not throw error 85 + new MongoRemoteSnapshotClient(client, upgradeDb, collName, true); + + // Phase 3: verify index was upgraded to 30 days + Map indexes = indexMap(upgradeDb, collName); + Document ttlIndex = + indexes.values().stream() + .filter( + i -> { + Object key = i.get("key"); + return key instanceof Document d && d.containsKey("createdAt"); + }) + .findFirst() + .orElse(null); + assertNotNull(ttlIndex); + assertEquals( + THIRTY_DAYS_SECONDS, ((Number) ttlIndex.get("expireAfterSeconds")).longValue(), - "Snapshot TTL must be 7 days (604800s)"); + "After upgrade, snapshot TTL must be 30 days"); + + // Cleanup + upgradeDbRef.drop(); } // ────────────────── Helpers ────────────────── diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/sandbox/MongoSandboxExecutionGuardTest.java b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/sandbox/MongoSandboxExecutionGuardTest.java index 3d7674b54a..2320ecd74d 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/sandbox/MongoSandboxExecutionGuardTest.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/sandbox/MongoSandboxExecutionGuardTest.java @@ -17,13 +17,18 @@ import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.after; import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.timeout; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import com.mongodb.MongoClientSettings; import com.mongodb.MongoWriteException; import com.mongodb.ServerAddress; import com.mongodb.WriteError; @@ -32,6 +37,7 @@ import com.mongodb.client.MongoCollection; import com.mongodb.client.MongoDatabase; import com.mongodb.client.model.FindOneAndUpdateOptions; +import com.mongodb.client.result.UpdateResult; import io.agentscope.harness.agent.IsolationScope; import io.agentscope.harness.agent.sandbox.SandboxIsolationKey; import io.agentscope.harness.agent.sandbox.SandboxLease; @@ -43,6 +49,7 @@ import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; import org.mockito.Mock; import org.mockito.MockitoAnnotations; @@ -171,6 +178,31 @@ void builderWithCustomOwner() { assertNotNull(guard); } + @Test + void builderWithCustomLeaseTtl() { + MongoSandboxExecutionGuard guard = + MongoSandboxExecutionGuard.builder(mongoClient) + .leaseTtl(Duration.ofMinutes(10)) + .build(); + assertNotNull(guard); + } + + @Test + void builderRejectsNonPositiveLeaseTtl() { + assertThrows( + IllegalArgumentException.class, + () -> + MongoSandboxExecutionGuard.builder(mongoClient) + .leaseTtl(Duration.ZERO) + .build()); + assertThrows( + IllegalArgumentException.class, + () -> + MongoSandboxExecutionGuard.builder(mongoClient) + .leaseTtl(Duration.ofSeconds(-1)) + .build()); + } + @Test void tryEnterAcquiresLockViaInsert() throws Exception { MongoSandboxExecutionGuard guard = @@ -308,4 +340,126 @@ void leaseCloseHandlesReleaseFailure() throws Exception { // Should not throw — close() swallows errors lease.close(); } + + @Test + void lockDocumentExpiryUsesLeaseTtlNotAcquisitionTimeout() throws Exception { + MongoSandboxExecutionGuard guard = + MongoSandboxExecutionGuard.builder(mongoClient) + .lockTimeout(Duration.ofSeconds(1)) + .leaseTtl(Duration.ofMinutes(10)) + .build(); + + SandboxLease lease = guard.tryEnter(key()); + lease.close(); + + ArgumentCaptor docCaptor = ArgumentCaptor.forClass(Document.class); + verify(collection).insertOne(docCaptor.capture()); + Date expiresAt = docCaptor.getValue().getDate("expiresAt"); + assertNotNull(expiresAt); + + long deltaMs = expiresAt.getTime() - System.currentTimeMillis(); + assertTrue( + deltaMs > Duration.ofMinutes(9).toMillis(), + "expiresAt must be ~leaseTtl (10 min) away, was " + deltaMs + "ms"); + assertTrue( + deltaMs <= Duration.ofMinutes(10).toMillis() + 5_000, + "expiresAt must not exceed leaseTtl (10 min), was " + deltaMs + "ms"); + } + + @Test + void leaseRenewsPeriodicallyWhileOpen() throws Exception { + UpdateResult renewResult = mock(UpdateResult.class); + when(renewResult.getMatchedCount()).thenReturn(1L); + when(collection.updateOne(any(Bson.class), any(Bson.class))).thenReturn(renewResult); + + // leaseTtl 300ms → renewal every 100ms + MongoSandboxExecutionGuard guard = + MongoSandboxExecutionGuard.builder(mongoClient) + .lockTimeout(Duration.ofSeconds(5)) + .leaseTtl(Duration.ofMillis(300)) + .build(); + SandboxLease lease = guard.tryEnter(key()); + try { + ArgumentCaptor filterCaptor = ArgumentCaptor.forClass(Bson.class); + ArgumentCaptor updateCaptor = ArgumentCaptor.forClass(Bson.class); + verify(collection, timeout(3_000).atLeast(2)) + .updateOne(filterCaptor.capture(), updateCaptor.capture()); + + // Renewal must push expiresAt forward and only match our own lock document + String filterJson = + filterCaptor + .getValue() + .toBsonDocument( + BsonDocument.class, + MongoClientSettings.getDefaultCodecRegistry()) + .toJson(); + assertTrue(filterJson.contains("owner"), "renewal filter must be owner-scoped"); + + BsonDocument updateDoc = + updateCaptor + .getValue() + .toBsonDocument( + BsonDocument.class, + MongoClientSettings.getDefaultCodecRegistry()); + assertTrue( + updateDoc.getDocument("$set").containsKey("expiresAt"), + "renewal must extend expiresAt"); + } finally { + lease.close(); + } + } + + @Test + void leaseCloseStopsRenewal() throws Exception { + UpdateResult renewResult = mock(UpdateResult.class); + when(renewResult.getMatchedCount()).thenReturn(1L); + when(collection.updateOne(any(Bson.class), any(Bson.class))).thenReturn(renewResult); + + // leaseTtl 150ms → renewal every 50ms + MongoSandboxExecutionGuard guard = + MongoSandboxExecutionGuard.builder(mongoClient) + .lockTimeout(Duration.ofSeconds(5)) + .leaseTtl(Duration.ofMillis(150)) + .build(); + SandboxLease lease = guard.tryEnter(key()); + verify(collection, timeout(3_000).atLeast(2)).updateOne(any(Bson.class), any(Bson.class)); + + lease.close(); + // Let an in-flight tick (if any) finish before counting + Thread.sleep(150); + int renewalsAfterClose = countUpdateOneInvocations(); + + // Over several renewal intervals the count must not grow: close() cancelled the watchdog + verify(collection, after(500).times(renewalsAfterClose)) + .updateOne(any(Bson.class), any(Bson.class)); + } + + @Test + void renewalContinuesAfterLockLost() throws Exception { + // matchedCount == 0 means the lock was reclaimed by someone else; the watchdog must + // only warn and keep running — never escape an exception that would kill the scheduler. + UpdateResult lostResult = mock(UpdateResult.class); + when(lostResult.getMatchedCount()).thenReturn(0L); + when(collection.updateOne(any(Bson.class), any(Bson.class))).thenReturn(lostResult); + + MongoSandboxExecutionGuard guard = + MongoSandboxExecutionGuard.builder(mongoClient) + .lockTimeout(Duration.ofSeconds(5)) + .leaseTtl(Duration.ofMillis(150)) + .build(); + SandboxLease lease = guard.tryEnter(key()); + try { + verify(collection, timeout(3_000).atLeast(3)) + .updateOne(any(Bson.class), any(Bson.class)); + } finally { + lease.close(); + } + } + + private int countUpdateOneInvocations() { + return (int) + org.mockito.Mockito.mockingDetails(collection).getInvocations().stream() + .filter(invocation -> invocation.getMethod().getName().equals("updateOne")) + .count(); + } } diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClientTest.java b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClientTest.java index 29a21ce9c7..a38c6c1348 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClientTest.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/snapshot/MongoRemoteSnapshotClientTest.java @@ -23,13 +23,16 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import com.mongodb.MongoCommandException; import com.mongodb.client.FindIterable; import com.mongodb.client.MongoClient; import com.mongodb.client.MongoCollection; import com.mongodb.client.MongoDatabase; +import com.mongodb.client.model.IndexOptions; import com.mongodb.client.result.DeleteResult; import com.mongodb.client.result.UpdateResult; import java.io.ByteArrayInputStream; @@ -181,6 +184,33 @@ void downloadThrowsWhenSnapshotNotFound() { assertThrows(FileNotFoundException.class, () -> client.download("missing-snap")); } + @Test + void downloadThrowsFileNotFoundWhenDataFieldMissing() { + // Document exists but carries no data field — must surface as FileNotFoundException + // (same as a missing snapshot), not as an NPE. + Document docWithoutData = new Document("_id", "snap-1"); + when(findIterable.first()).thenReturn(docWithoutData); + + assertThrows(FileNotFoundException.class, () -> client.download("snap-1")); + } + + @Test + void initSchemaMigratesConflictingTtlIndex() { + // Simulate an existing collection whose createdAt TTL index still has the old 7-day + // expireAfterSeconds: createIndex then fails with IndexOptionsConflict (error 85) and + // must trigger drop + recreate instead of silently keeping the stale TTL. + MongoCommandException conflict = mock(MongoCommandException.class); + when(conflict.getErrorCode()).thenReturn(85); + when(collection.createIndex(any(Bson.class), any(IndexOptions.class))) + .thenThrow(conflict) + .thenReturn("createdAt_1"); + + new MongoRemoteSnapshotClient(mongoClient, "testdb", "snap_migrate", true); + + verify(collection).dropIndex("createdAt_1"); + verify(collection, times(2)).createIndex(any(Bson.class), any(IndexOptions.class)); + } + @Test void downloadRejectsNullSnapshotId() { assertThrows(NullPointerException.class, () -> client.download(null)); diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreContractTest.java b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreContractTest.java index 7439c6dede..c93bb46652 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreContractTest.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreContractTest.java @@ -24,6 +24,7 @@ import io.agentscope.core.state.AgentStateStore; import io.agentscope.core.state.State; import io.agentscope.core.state.VersionedState; +import java.util.List; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -216,5 +217,40 @@ void concurrentWriters_onlyOneSucceeds() throws InterruptedException { 2L, store.getVersioned(USER, SESSION, "agent_state", TestState.class).version()); } + @Test + @Order(7) + @DisplayName("session first saved as a non-empty list is visible to listSessionIds") + void listFirstSave_sessionVisibleToListSessionIds() { + // Regression for the review blocker: a brand-new session whose first write hits the + // incremental-append branch used to be upserted without user_id/session_id fields, + // making the session permanently invisible to listSessionIds. + String listFirstSession = "list-first-session"; + try { + store.save( + USER, + listFirstSession, + "memory_messages", + List.of(new TestState("m1"), new TestState("m2"))); + + assertTrue( + store.listSessionIds(USER).contains(listFirstSession), + "session first saved via list append must be visible to listSessionIds"); + assertEquals( + 2, + store.getList(USER, listFirstSession, "memory_messages", TestState.class) + .size()); + + // A later full rewrite (list shrinks) must not lose visibility either + store.save(USER, listFirstSession, "memory_messages", List.of(new TestState("m3"))); + assertTrue(store.listSessionIds(USER).contains(listFirstSession)); + assertEquals( + 1, + store.getList(USER, listFirstSession, "memory_messages", TestState.class) + .size()); + } finally { + store.delete(USER, listFirstSession); + } + } + record TestState(String value) implements State {} } diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreTest.java b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreTest.java index 3f531695ef..5a85725d54 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreTest.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/test/java/io/agentscope/extensions/mongodb/state/MongoAgentStateStoreTest.java @@ -26,6 +26,7 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import com.mongodb.MongoClientSettings; import com.mongodb.MongoWriteException; import com.mongodb.ServerAddress; import com.mongodb.WriteError; @@ -50,6 +51,7 @@ import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; import org.mockito.ArgumentMatchers; import org.mockito.Mock; import org.mockito.Mockito; @@ -170,6 +172,60 @@ void saveListState() { verify(collection).updateOne(any(Bson.class), any(Bson.class), any()); } + @Test + void saveListAppendOnNewSessionIncludesSetOnInsert() { + // Brand-new session (no existing document) whose first write is a non-empty list: + // needsFullRewrite(values, null, 0) returns false, so the append branch runs. Its + // upsert must carry $setOnInsert(user_id, session_id), otherwise the session document + // is created without them and listSessionIds can never see it — the blocker from the + // code review. + store.save("user", "session", "list", List.of(new TestState("a"))); + + ArgumentCaptor updateCaptor = ArgumentCaptor.forClass(Bson.class); + verify(collection).updateOne(any(Bson.class), updateCaptor.capture(), any()); + + BsonDocument updateDoc = + updateCaptor + .getValue() + .toBsonDocument( + BsonDocument.class, MongoClientSettings.getDefaultCodecRegistry()); + assertTrue(updateDoc.containsKey("$push"), "expected the append branch ($push)"); + + BsonDocument setOnInsert = updateDoc.getDocument("$setOnInsert"); + assertEquals("user", setOnInsert.getString("user_id").getValue()); + assertEquals("session", setOnInsert.getString("session_id").getValue()); + } + + @Test + void saveListFullRewriteIncludesSetOnInsert() { + // Existing document whose stored list shrinks — full-rewrite branch. Its upsert must + // also carry $setOnInsert so a rewrite that happens to be the first write on a slot + // still records the session identifiers. + List existingList = + List.of( + Document.parse("{\"value\":\"a\"}"), + Document.parse("{\"value\":\"b\"}"), + Document.parse("{\"value\":\"c\"}")); + Document existingDoc = new Document("list:list", existingList); + when(findIterable.first()).thenReturn(existingDoc); + + store.save("user", "session", "list", List.of(new TestState("x"))); + + ArgumentCaptor updateCaptor = ArgumentCaptor.forClass(Bson.class); + verify(collection).updateOne(any(Bson.class), updateCaptor.capture(), any()); + + BsonDocument updateDoc = + updateCaptor + .getValue() + .toBsonDocument( + BsonDocument.class, MongoClientSettings.getDefaultCodecRegistry()); + assertTrue(updateDoc.containsKey("$set"), "expected the rewrite branch ($set)"); + + BsonDocument setOnInsert = updateDoc.getDocument("$setOnInsert"); + assertEquals("user", setOnInsert.getString("user_id").getValue()); + assertEquals("session", setOnInsert.getString("session_id").getValue()); + } + @Test void saveListShorteningPerformsFullRewrite() { // Simulate existing document with 3 elements in the list From be32639c2622db26534cf8aa8650e2d14c6ac8de Mon Sep 17 00:00:00 2001 From: LiangshouX <93421804+LiangshouX@users.noreply.github.com> Date: Wed, 26 Aug 2026 16:54:32 +0800 Subject: [PATCH 7/8] Update agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/store/MongoBaseStore.java Co-authored-by: Larry <139796123+larry-zy@users.noreply.github.com> --- .../agentscope/extensions/mongodb/store/MongoBaseStore.java | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/store/MongoBaseStore.java b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/store/MongoBaseStore.java index e32e60f3d3..f73b9f7a04 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/store/MongoBaseStore.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/store/MongoBaseStore.java @@ -167,12 +167,15 @@ public boolean putIfVersion( @Override public List search(List namespace, int limit, int offset) { + if (limit <= 0) { + return List.of(); + } String nsKey = namespacePath(namespace); List docs = collection .find(Filters.eq(FIELD_NAMESPACE, nsKey)) .sort(Sorts.ascending(FIELD_KEY)) - .skip(offset) + .skip(Math.max(offset, 0)) .limit(limit) .projection(Projections.include(FIELD_KEY, FIELD_VALUE, FIELD_VERSION)) .into(new ArrayList<>()); From 1d1047f3ce1a73d61360a8da8c353dfc9af04846 Mon Sep 17 00:00:00 2001 From: LiangshouX <93421804+LiangshouX@users.noreply.github.com> Date: Wed, 26 Aug 2026 16:55:46 +0800 Subject: [PATCH 8/8] Update agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/store/MongoBaseStore.java Co-authored-by: Larry <139796123+larry-zy@users.noreply.github.com> --- .../io/agentscope/extensions/mongodb/store/MongoBaseStore.java | 1 - 1 file changed, 1 deletion(-) diff --git a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/store/MongoBaseStore.java b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/store/MongoBaseStore.java index f73b9f7a04..4ca6ec8175 100644 --- a/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/store/MongoBaseStore.java +++ b/agentscope-extensions/agentscope-extensions-mongodb/src/main/java/io/agentscope/extensions/mongodb/store/MongoBaseStore.java @@ -198,7 +198,6 @@ public void delete(List namespace, String key) { // ────────────────── Internal Helpers ────────────────── private void ensureIndexes() { - collection.createIndex(Indexes.ascending(FIELD_NAMESPACE)); collection.createIndex( Indexes.compoundIndex( Indexes.ascending(FIELD_NAMESPACE), Indexes.ascending(FIELD_KEY)));