feat(mongodb): stream legacy agent exports via find cursors

agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java

@@ -10,10 +10,14 @@ import com.google.gson.JsonNull;
 import com.google.gson.JsonObject;
 import com.google.gson.JsonParser;
 import com.mongodb.ConnectionString;
+import com.mongodb.MongoBulkWriteException;
 import com.mongodb.MongoCredential;
 import com.mongodb.MongoClientSettings;
 import com.mongodb.ServerAddress;
+import com.mongodb.bulk.BulkWriteError;
+import com.mongodb.bulk.WriteConcernError;
 import com.mongodb.client.AggregateIterable;
+import com.mongodb.client.FindIterable;
 import com.mongodb.client.MongoClient;
 import com.mongodb.client.MongoClients;
 import com.mongodb.client.MongoCollection;
@@ -58,6 +62,7 @@ import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -81,11 +86,22 @@ public final class MongoAgent {
     private static final JsonWriterSettings EXTENDED_JSON_SETTINGS = JsonWriterSettings.builder()
         .outputMode(JsonMode.RELAXED)
         .build();
+    // Driver 3.x names the canonical Extended JSON v2 output "EXTENDED"; RELAXED is its lossy
+    // sibling (plain numbers, ISO dates), so exports must use this one to keep BSON types.
+    private static final JsonWriterSettings CANONICAL_JSON_SETTINGS = JsonWriterSettings.builder()
+        .outputMode(JsonMode.EXTENDED)
+        .build();
     private static final String LEGACY_SESSION_ID = "__legacy__";
     private static final String DEFAULT_ID_INDEX_NAME = "_id_";
     private static final int CLONE_INSERT_BATCH_SIZE = 1_000;
+    private static final int DEFAULT_FIND_CURSOR_BATCH_SIZE = 1_000;
+    private static final int MAX_FIND_CURSORS_PER_OWNER = 16;
+    private static final int MAX_FIND_CURSORS_TOTAL = 64;
     private static final int MAX_SESSIONS = 256;
     private static final ThreadLocal<MongoClient> CURRENT_CLIENT = new ThreadLocal<>();
+    private static final ThreadLocal<Set<String>> CURRENT_FIND_CURSOR_OWNER = new ThreadLocal<>();
+    private static final ConcurrentHashMap<String, MongoCursor<Document>> FIND_CURSORS = new ConcurrentHashMap<>();
+    private static final Set<String> LEGACY_FIND_CURSOR_IDS = ConcurrentHashMap.newKeySet();
     private static MongoClient legacyClient;
 
     private MongoAgent() {
@@ -1489,8 +1505,174 @@ public final class MongoAgent {
         MongoClient client = requireClient();
         String database = params.get("database").getAsString();
         String collection = params.get("collection").getAsString();
-        client.getDatabase(database).getCollection(collection).insertMany(documents);
-        return Collections.singletonMap("affected_rows", documents.size());
+        InsertManyOptions options = new InsertManyOptions();
+        if (params.has("ordered") && !params.get("ordered").isJsonNull()) {
+            JsonElement ordered = params.get("ordered");
+            if (!ordered.isJsonPrimitive() || !ordered.getAsJsonPrimitive().isBoolean()) {
+                throw new IllegalArgumentException("ordered must be a boolean");
+            }
+            options.ordered(ordered.getAsBoolean());
+        }
+        try {
+            client.getDatabase(database).getCollection(collection).insertMany(documents, options);
+            return Collections.singletonMap("affected_rows", documents.size());
+        } catch (MongoBulkWriteException error) {
+            return bulkInsertResult(error);
+        }
+    }
+
+    static Map<String, Object> bulkInsertResult(MongoBulkWriteException error) {
+        Map<String, Object> result = new LinkedHashMap<>();
+        result.put("affected_rows", error.getWriteResult().getInsertedCount());
+        List<Map<String, Object>> errors = new ArrayList<>();
+        for (BulkWriteError writeError : error.getWriteErrors()) {
+            Map<String, Object> item = new LinkedHashMap<>();
+            item.put("index", writeError.getIndex());
+            item.put("code", writeError.getCode());
+            item.put("message", writeError.getMessage());
+            errors.add(item);
+        }
+        // The driver raises the same exception for write concern failures only (for example a
+        // wtimeout): there are no per-document errors then, but the batch must not look
+        // successful, so report the write concern error as a batch-wide rejection without index.
+        WriteConcernError writeConcernError = error.getWriteConcernError();
+        if (errors.isEmpty() && writeConcernError != null) {
+            Map<String, Object> item = new LinkedHashMap<>();
+            item.put("code", writeConcernError.getCode());
+            item.put("message", writeConcernError.getMessage());
+            errors.add(item);
+        }
+        result.put("errors", errors);
+        return result;
+    }
+
+    private static Object startFindCursor(JsonObject params) {
+        Set<String> owner = CURRENT_FIND_CURSOR_OWNER.get();
+        if (owner != null && owner.size() >= MAX_FIND_CURSORS_PER_OWNER) {
+            throw new IllegalStateException("MongoDB find cursor limit reached: " + MAX_FIND_CURSORS_PER_OWNER);
+        }
+        // The per-owner budget keeps one session from starving the others; the process-wide cap
+        // only bounds how many server cursors this agent can hold in total.
+        if (FIND_CURSORS.size() >= MAX_FIND_CURSORS_TOTAL) {
+            throw new IllegalStateException("MongoDB find cursor limit reached: " + MAX_FIND_CURSORS_TOTAL);
+        }
+        MongoClient client = requireClient();
+        String database = params.get("database").getAsString();
+        String collection = params.get("collection").getAsString();
+        Document filterDoc = documentOrNull(params, "filter");
+        Document projectionDoc = documentOrNull(params, "projection");
+        Document sortDoc = documentOrNull(params, "sort");
+        Collation collation = collationOrNull(documentOrNull(params, "collation"));
+        int batchSize = findCursorBatchSize(params);
+
+        FindIterable<Document> iterable = client.getDatabase(database)
+            .getCollection(collection)
+            .find(filterDoc == null ? new Document() : filterDoc)
+            .batchSize(batchSize);
+        if (projectionDoc != null) {
+            iterable = iterable.projection(projectionDoc);
+        }
+        if (sortDoc != null) {
+            iterable = iterable.sort(sortDoc);
+        }
+        if (collation != null) {
+            iterable = iterable.collation(collation);
+        }
+
+        String cursorId = UUID.randomUUID().toString();
+        FIND_CURSORS.put(cursorId, iterable.iterator());
+        if (owner != null) {
+            owner.add(cursorId);
+        }
+        Map<String, Object> result = new LinkedHashMap<>();
+        result.put("cursor_id", cursorId);
+        result.put("batch_size", batchSize);
+        return result;
+    }
+
+    private static Object fetchFindCursor(JsonObject params) {
+        String cursorId = requiredCursorId(params);
+        MongoCursor<Document> cursor = FIND_CURSORS.get(cursorId);
+        if (cursor == null) {
+            throw new IllegalStateException("Find cursor not found");
+        }
+        int limit = params.has("limit") ? params.get("limit").getAsInt() : DEFAULT_FIND_CURSOR_BATCH_SIZE;
+        if (limit <= 0) {
+            throw new IllegalArgumentException("Find cursor limit must be a positive integer");
+        }
+        List<JsonObject> documents = new ArrayList<>();
+        try {
+            while (documents.size() < limit && cursor.hasNext()) {
+                documents.add(bsonToCanonicalExtendedJson(cursor.next()));
+            }
+        } catch (RuntimeException error) {
+            // A cursor that fails mid-iteration (server-side idle timeout, lost connection) can
+            // never be resumed, so drop it instead of leaking a slot in the cursor budgets.
+            closeFindCursorById(cursorId);
+            throw error;
+        }
+        boolean exhausted = !cursor.hasNext();
+        if (exhausted) {
+            closeFindCursorById(cursorId);
+        }
+        Map<String, Object> result = new LinkedHashMap<>();
+        result.put("documents", documents);
+        result.put("exhausted", exhausted);
+        return result;
+    }
+
+    private static Object closeFindCursor(JsonObject params) {
+        closeFindCursorById(requiredCursorId(params));
+        return Collections.singletonMap("ok", true);
+    }
+
+    private static String requiredCursorId(JsonObject params) {
+        String cursorId = stringOrNull(params, "cursor_id");
+        if (cursorId == null || cursorId.isBlank()) {
+            throw new IllegalArgumentException("cursor_id is required");
+        }
+        return cursorId;
+    }
+
+    private static int findCursorBatchSize(JsonObject params) {
+        if (!params.has("batch_size") || params.get("batch_size").isJsonNull()) {
+            return DEFAULT_FIND_CURSOR_BATCH_SIZE;
+        }
+        int batchSize = params.get("batch_size").getAsInt();
+        if (batchSize <= 0) {
+            throw new IllegalArgumentException("batch_size must be a positive integer");
+        }
+        return batchSize;
+    }
+
+    private static void closeFindCursorById(String cursorId) {
+        MongoCursor<Document> cursor = FIND_CURSORS.remove(cursorId);
+        if (cursor != null) {
+            try {
+                cursor.close();
+            } catch (RuntimeException ignored) {
+                // Closing a server cursor is best-effort; the client may already be gone.
+            }
+        }
+        Set<String> owner = CURRENT_FIND_CURSOR_OWNER.get();
+        if (owner != null) {
+            owner.remove(cursorId);
+        }
+        LEGACY_FIND_CURSOR_IDS.remove(cursorId);
+    }
+
+    private static void closeFindCursors(Set<String> cursorIds) {
+        for (String cursorId : new ArrayList<>(cursorIds)) {
+            closeFindCursorById(cursorId);
+        }
+        cursorIds.clear();
+    }
+
+    static void resetFindCursorsForTests() {
+        for (String cursorId : new ArrayList<>(FIND_CURSORS.keySet())) {
+            closeFindCursorById(cursorId);
+        }
+        LEGACY_FIND_CURSOR_IDS.clear();
     }
 
     static Object parseId(String id) {
@@ -1787,6 +1969,10 @@ public final class MongoAgent {
         return preserveUnsafeLongsForJsonClients(doc, relaxed).getAsJsonObject();
     }
 
+    static JsonObject bsonToCanonicalExtendedJson(Document doc) {
+        return JsonParser.parseString(doc.toJson(CANONICAL_JSON_SETTINGS)).getAsJsonObject();
+    }
+
     private static JsonElement preserveUnsafeLongsForJsonClients(Object bsonValue, JsonElement relaxedValue) {
         if (
             bsonValue instanceof Long longValue &&
@@ -1914,6 +2100,9 @@ public final class MongoAgent {
             case AgentProtocol.MONGO_METHOD_DROP_DATABASE -> dropDatabase(params);
             case AgentProtocol.MONGO_METHOD_INSERT_DOCUMENT -> insertDocument(params);
             case AgentProtocol.MONGO_METHOD_INSERT_DOCUMENTS -> insertDocuments(params);
+            case AgentProtocol.MONGO_METHOD_START_FIND_CURSOR -> startFindCursor(params);
+            case AgentProtocol.MONGO_METHOD_FETCH_FIND_CURSOR -> fetchFindCursor(params);
+            case AgentProtocol.MONGO_METHOD_CLOSE_FIND_CURSOR -> closeFindCursor(params);
             case AgentProtocol.MONGO_METHOD_UPDATE_DOCUMENT -> updateDocument(params);
             case AgentProtocol.MONGO_METHOD_UPDATE_DOCUMENTS -> updateDocuments(params);
             case AgentProtocol.MONGO_METHOD_REPLACE_DOCUMENT -> replaceDocument(params);
@@ -1921,6 +2110,7 @@ public final class MongoAgent {
             case AgentProtocol.MONGO_METHOD_DELETE_DOCUMENTS -> deleteDocuments(params);
             case AgentProtocol.MONGO_METHOD_RUN_COMMAND -> runCommand(params);
             case AgentProtocol.METHOD_DISCONNECT, AgentProtocol.METHOD_SHUTDOWN -> {
+                closeFindCursors(LEGACY_FIND_CURSOR_IDS);
                 closeLegacyClient();
                 if (AgentProtocol.METHOD_SHUTDOWN.equals(method)) {
                     System.exit(0);
@@ -1968,6 +2158,7 @@ public final class MongoAgent {
     static String handleRequest(String line, MongoClient client) {
         if (client != null) {
             CURRENT_CLIENT.set(client);
+            CURRENT_FIND_CURSOR_OWNER.set(LEGACY_FIND_CURSOR_IDS);
         }
         try {
             JsonObject req = JsonParser.parseString(line).getAsJsonObject();
@@ -1995,6 +2186,7 @@ public final class MongoAgent {
         } finally {
             if (client != null) {
                 CURRENT_CLIENT.remove();
+                CURRENT_FIND_CURSOR_OWNER.remove();
             }
         }
     }
@@ -2120,6 +2312,7 @@ public final class MongoAgent {
 
     private static final class Session {
         private final MongoClient client;
+        private final Set<String> findCursorIds = ConcurrentHashMap.newKeySet();
 
         private Session(MongoClient client) {
             this.client = client;
@@ -2127,10 +2320,12 @@ public final class MongoAgent {
 
         private synchronized Object handle(String method, JsonObject params) {
             CURRENT_CLIENT.set(client);
+            CURRENT_FIND_CURSOR_OWNER.set(findCursorIds);
             try {
                 return dispatch(method, params);
             } finally {
                 CURRENT_CLIENT.remove();
+                CURRENT_FIND_CURSOR_OWNER.remove();
             }
         }
 
@@ -2144,6 +2339,7 @@ public final class MongoAgent {
         }
 
         private synchronized void close() {
+            closeFindCursors(findCursorIds);
             client.close();
         }
     }

agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java

@@ -1,5 +1,6 @@
 package com.dbx.agent.mongodb;
 
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
@@ -12,7 +13,12 @@ import com.google.gson.Gson;
 import com.google.gson.JsonArray;
 import com.google.gson.JsonObject;
 import com.google.gson.JsonParser;
+import com.mongodb.MongoBulkWriteException;
 import com.mongodb.MongoClientSettings;
+import com.mongodb.ServerAddress;
+import com.mongodb.bulk.BulkWriteError;
+import com.mongodb.bulk.BulkWriteResult;
+import com.mongodb.bulk.WriteConcernError;
 import com.mongodb.client.AggregateIterable;
 import com.mongodb.client.FindIterable;
 import com.mongodb.client.ListCollectionsIterable;
@@ -40,8 +46,12 @@ import java.util.Date;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
+import org.bson.BsonDocument;
+import org.bson.BsonRegularExpression;
 import org.bson.Document;
+import org.bson.types.Binary;
 import org.bson.types.ObjectId;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
@@ -99,6 +109,11 @@ class MongoAgentTest {
         Files.writeString(clientKeyPath, pkcs8Pem);
     }
 
+    @AfterEach
+    void resetFindCursors() {
+        MongoAgent.resetFindCursorsForTests();
+    }
+
     // ─── existing tests ───
 
     @Test
@@ -128,6 +143,7 @@ class MongoAgentTest {
         assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MONGO_CLONE_COLLECTION));
         assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MONGO_RUN_COMMAND));
         assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MONGO_INSERT_DOCUMENTS));
+        assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MONGO_FIND_CURSOR));
     }
 
     @Test
@@ -150,6 +166,7 @@ class MongoAgentTest {
         assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MONGO_CLONE_COLLECTION));
         assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MONGO_RUN_COMMAND));
         assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MONGO_INSERT_DOCUMENTS));
+        assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MONGO_FIND_CURSOR));
     }
 
     @Test
@@ -894,6 +911,10 @@ class MongoAgentTest {
         assertEquals("12", batches.get(0).get(1).getString("refid"));
         assertTrue(AgentProtocol.MONGO_LEGACY_METHODS.contains(AgentProtocol.MONGO_METHOD_INSERT_DOCUMENTS));
         assertTrue(AgentProtocol.MONGO_LEGACY_CAPABILITIES.contains(AgentProtocol.CAPABILITY_MONGO_INSERT_DOCUMENTS));
+        assertTrue(AgentProtocol.MONGO_LEGACY_CAPABILITIES.contains(AgentProtocol.CAPABILITY_MONGO_FIND_CURSOR));
+        assertTrue(AgentProtocol.MONGO_LEGACY_METHODS.contains(AgentProtocol.MONGO_METHOD_START_FIND_CURSOR));
+        assertTrue(AgentProtocol.MONGO_LEGACY_METHODS.contains(AgentProtocol.MONGO_METHOD_FETCH_FIND_CURSOR));
+        assertTrue(AgentProtocol.MONGO_LEGACY_METHODS.contains(AgentProtocol.MONGO_METHOD_CLOSE_FIND_CURSOR));
     }
 
     @Test
@@ -939,6 +960,293 @@ class MongoAgentTest {
         }
     }
 
+    @Test
+    void documentForWritePreservesExtendedJsonBsonTypes() {
+        Document document = MongoAgent.documentForWrite(
+            "{\"_id\":{\"$oid\":\"507f1f77bcf86cd799439011\"},"
+                + "\"count\":{\"$numberLong\":\"9007199254740993\"},"
+                + "\"when\":{\"$date\":\"2020-01-02T03:04:05.000Z\"}}"
+        );
+
+        assertEquals(new ObjectId("507f1f77bcf86cd799439011"), document.getObjectId("_id"));
+        assertEquals(9007199254740993L, document.get("count"));
+        assertTrue(document.get("when") instanceof Date);
+        assertEquals(Date.from(java.time.Instant.parse("2020-01-02T03:04:05.000Z")), document.getDate("when"));
+    }
+
+    @Test
+    void insertDocumentsRpcHonoursExplicitUnorderedInserts() {
+        List<List<Document>> batches = new ArrayList<>();
+        List<String> calls = new ArrayList<>();
+        List<Boolean> ordered = new ArrayList<>();
+        MongoClient client = recordingInsertMongoClient(calls, batches, ordered);
+
+        JsonObject params = new JsonObject();
+        params.addProperty("database", "app");
+        params.addProperty("collection", "users");
+        params.addProperty("docs_json", "[{\"name\":\"Ada\"}]");
+        params.addProperty("ordered", false);
+        JsonObject request = new JsonObject();
+        request.addProperty("jsonrpc", "2.0");
+        request.addProperty("id", 44);
+        request.addProperty("method", "insert_documents");
+        request.add("params", params);
+
+        JsonObject response = JsonParser.parseString(MongoAgent.handleRequest(request.toString(), client)).getAsJsonObject();
+
+        assertFalse(response.has("error"), response.toString());
+        assertEquals(List.of(false), ordered);
+        assertEquals(1, response.getAsJsonObject("result").get("affected_rows").getAsInt());
+    }
+
+    @Test
+    void insertDocumentsRpcPreservesExtendedJsonTypesInTheInsertManyBatch() {
+        List<List<Document>> batches = new ArrayList<>();
+        List<String> calls = new ArrayList<>();
+        MongoClient client = recordingInsertMongoClient(calls, batches);
+        String documents = "[{\"_id\":{\"$oid\":\"507f1f77bcf86cd799439011\"},"
+            + "\"count\":{\"$numberLong\":\"9007199254740993\"}}]";
+
+        JsonObject response = insertDocumentsRpc(client, 45, "users", documents);
+
+        assertFalse(response.has("error"), response.toString());
+        assertEquals(1, batches.size());
+        Document written = batches.get(0).get(0);
+        assertEquals(new ObjectId("507f1f77bcf86cd799439011"), written.getObjectId("_id"));
+        assertEquals(9007199254740993L, written.get("count"));
+        assertTrue(written.get("count") instanceof Long);
+    }
+
+    @Test
+    void insertDocumentsRpcParsesCanonicalBinaryAndRegularExpression() {
+        List<List<Document>> batches = new ArrayList<>();
+        MongoClient client = recordingInsertMongoClient(new ArrayList<>(), batches);
+        String documents = "[{\"bin\":{\"$binary\":{\"base64\":\"AQID\",\"subType\":\"00\"}},"
+            + "\"re\":{\"$regularExpression\":{\"pattern\":\"^test$\",\"options\":\"i\"}}}]";
+
+        JsonObject response = insertDocumentsRpc(client, 46, "users", documents);
+
+        assertFalse(response.has("error"), response.toString());
+        Document written = batches.get(0).get(0);
+        assertTrue(written.get("bin") instanceof Binary, String.valueOf(written.get("bin")));
+        Binary binary = (Binary) written.get("bin");
+        assertEquals(0, binary.getType());
+        assertArrayEquals(new byte[] {1, 2, 3}, binary.getData());
+        assertTrue(written.get("re") instanceof BsonRegularExpression, String.valueOf(written.get("re")));
+        BsonRegularExpression regex = (BsonRegularExpression) written.get("re");
+        assertEquals("^test$", regex.getPattern());
+        assertEquals("i", regex.getOptions());
+    }
+
+    @Test
+    void bulkInsertResultKeepsInsertedCountAndPerDocumentErrors() {
+        BulkWriteResult writeResult = BulkWriteResult.acknowledged(1, 0, 0, 0, List.of());
+        MongoBulkWriteException error = new MongoBulkWriteException(
+            writeResult,
+            List.of(new BulkWriteError(11000, "E11000 duplicate key", new BsonDocument(), 1)),
+            null,
+            new ServerAddress()
+        );
+
+        Map<String, Object> result = MongoAgent.bulkInsertResult(error);
+
+        assertEquals(1, result.get("affected_rows"));
+        @SuppressWarnings("unchecked")
+        List<Map<String, Object>> errors = (List<Map<String, Object>>) result.get("errors");
+        assertEquals(1, errors.size());
+        assertEquals(1, errors.get(0).get("index"));
+        assertEquals(11000, errors.get(0).get("code"));
+        assertEquals("E11000 duplicate key", errors.get(0).get("message"));
+    }
+
+    @Test
+    void bulkInsertResultReportsWriteConcernFailures() {
+        // The driver raises MongoBulkWriteException with no write errors when only the write
+        // concern failed, so the batch must not come back looking successful.
+        MongoBulkWriteException error = new MongoBulkWriteException(
+            BulkWriteResult.acknowledged(0, 0, 0, 0, List.of()),
+            List.of(),
+            new WriteConcernError(64, "WriteConcernFailed", "waiting for replication timed out", new BsonDocument()),
+            new ServerAddress()
+        );
+
+        Map<String, Object> result = MongoAgent.bulkInsertResult(error);
+
+        assertEquals(0, result.get("affected_rows"));
+        @SuppressWarnings("unchecked")
+        List<Map<String, Object>> errors = (List<Map<String, Object>>) result.get("errors");
+        assertEquals(1, errors.size());
+        assertEquals(64, errors.get(0).get("code"));
+        assertEquals("waiting for replication timed out", errors.get(0).get("message"));
+        assertFalse(errors.get(0).containsKey("index"));
+    }
+
+    @Test
+    void findCursorRpcPagesExtendedJsonDocumentsThenExhausts() {
+        List<String> calls = new ArrayList<>();
+        List<Document> documents = List.of(
+            new Document("_id", new ObjectId("507f1f77bcf86cd799439011"))
+                .append("n", 1L)
+                .append("bin", new Binary((byte) 0, new byte[] {1, 2, 3}))
+                .append("re", new BsonRegularExpression("^test$", "i")),
+            new Document("_id", new ObjectId("507f1f77bcf86cd799439012")).append("n", 2L)
+        );
+        MongoClient client = recordingFindMongoClient(calls, documents);
+
+        JsonObject start = findCursorRpc(client, 50, "start_find_cursor", cursorParams("users", 1, null));
+        assertFalse(start.has("error"), start.toString());
+        String cursorId = start.getAsJsonObject("result").get("cursor_id").getAsString();
+        assertEquals(1, start.getAsJsonObject("result").get("batch_size").getAsInt());
+
+        JsonObject first = findCursorRpc(client, 51, "fetch_find_cursor", cursorParams(null, 0, cursorId));
+        assertFalse(first.has("error"), first.toString());
+        assertEquals(1, first.getAsJsonObject("result").getAsJsonArray("documents").size());
+        assertFalse(first.getAsJsonObject("result").get("exhausted").getAsBoolean());
+        JsonObject firstDocument = first.getAsJsonObject("result").getAsJsonArray("documents").get(0).getAsJsonObject();
+        assertEquals(
+            "507f1f77bcf86cd799439011",
+            firstDocument.getAsJsonObject("_id").get("$oid").getAsString()
+        );
+        assertEquals("1", firstDocument.getAsJsonObject("n").get("$numberLong").getAsString());
+        JsonObject binary = firstDocument.getAsJsonObject("bin").getAsJsonObject("$binary");
+        assertEquals("AQID", binary.get("base64").getAsString(), firstDocument.toString());
+        assertEquals("00", binary.get("subType").getAsString(), firstDocument.toString());
+        JsonObject regex = firstDocument.getAsJsonObject("re").getAsJsonObject("$regularExpression");
+        assertEquals("^test$", regex.get("pattern").getAsString(), firstDocument.toString());
+        assertEquals("i", regex.get("options").getAsString(), firstDocument.toString());
+
+        JsonObject second = findCursorRpc(client, 52, "fetch_find_cursor", cursorParams(null, 0, cursorId));
+        assertFalse(second.has("error"), second.toString());
+        assertEquals(1, second.getAsJsonObject("result").getAsJsonArray("documents").size());
+        assertTrue(second.getAsJsonObject("result").get("exhausted").getAsBoolean());
+
+        JsonObject missing = findCursorRpc(client, 53, "fetch_find_cursor", cursorParams(null, 0, cursorId));
+        assertTrue(missing.has("error"), missing.toString());
+        assertEquals("Find cursor not found", missing.getAsJsonObject("error").get("message").getAsString());
+        assertTrue(calls.contains("find"));
+        assertTrue(calls.contains("batchSize:1"));
+        assertTrue(calls.contains("iterator"));
+        assertTrue(calls.contains("close"));
+    }
+
+    @Test
+    void closeFindCursorRpcIsIdempotent() {
+        List<String> calls = new ArrayList<>();
+        MongoClient client = recordingFindMongoClient(calls, List.of(new Document("n", 1)));
+        JsonObject start = findCursorRpc(client, 54, "start_find_cursor", cursorParams("users", 10, null));
+        String cursorId = start.getAsJsonObject("result").get("cursor_id").getAsString();
+
+        JsonObject closed = findCursorRpc(client, 55, "close_find_cursor", cursorParams(null, 0, cursorId));
+        JsonObject again = findCursorRpc(client, 56, "close_find_cursor", cursorParams(null, 0, cursorId));
+
+        assertFalse(closed.has("error"), closed.toString());
+        assertFalse(again.has("error"), again.toString());
+        assertEquals(1, calls.stream().filter("close"::equals).count());
+    }
+
+    @Test
+    void bsonToCanonicalExtendedJsonKeepsInt64DistinctFromJsonNumber() {
+        JsonObject json = MongoAgent.bsonToCanonicalExtendedJson(
+            new Document("n", 5L)
+                .append("bin", new Binary((byte) 0, new byte[] {1, 2, 3}))
+                .append("re", new BsonRegularExpression("^test$", "i"))
+        );
+
+        assertEquals("5", json.getAsJsonObject("n").get("$numberLong").getAsString());
+        assertFalse(json.get("n").isJsonPrimitive(), json.toString());
+        JsonObject binary = json.getAsJsonObject("bin").getAsJsonObject("$binary");
+        assertEquals("AQID", binary.get("base64").getAsString(), json.toString());
+        assertEquals("00", binary.get("subType").getAsString(), json.toString());
+        assertFalse(json.getAsJsonObject("bin").has("$type"), json.toString());
+        JsonObject regex = json.getAsJsonObject("re").getAsJsonObject("$regularExpression");
+        assertEquals("^test$", regex.get("pattern").getAsString(), json.toString());
+        assertEquals("i", regex.get("options").getAsString(), json.toString());
+        assertFalse(json.getAsJsonObject("re").has("$regex"), json.toString());
+    }
+
+    @Test
+    void disconnectClosesOpenFindCursors() {
+        List<String> calls = new ArrayList<>();
+        MongoClient client = recordingFindMongoClient(calls, List.of(new Document("n", 1L)));
+        JsonObject start = findCursorRpc(client, 57, "start_find_cursor", cursorParams("users", 10, null));
+        String cursorId = start.getAsJsonObject("result").get("cursor_id").getAsString();
+
+        JsonObject disconnect = JsonParser.parseString(MongoAgent.handleRequest(
+            "{\"jsonrpc\":\"2.0\",\"id\":58,\"method\":\"disconnect\",\"params\":{}}"
+        )).getAsJsonObject();
+        JsonObject fetch = findCursorRpc(client, 59, "fetch_find_cursor", cursorParams(null, 0, cursorId));
+
+        assertFalse(disconnect.has("error"), disconnect.toString());
+        assertTrue(fetch.has("error"), fetch.toString());
+        assertEquals("Find cursor not found", fetch.getAsJsonObject("error").get("message").getAsString());
+        assertTrue(calls.contains("close"));
+    }
+
+    @Test
+    void fetchFailureReleasesTheCursorSlot() {
+        List<String> calls = new ArrayList<>();
+        MongoCursor<Document> failing = new MongoCursor<>() {
+            @Override
+            public void close() {
+                calls.add("close");
+            }
+
+            @Override
+            public boolean hasNext() {
+                throw new IllegalStateException("cursor id 42 not found");
+            }
+
+            @Override
+            public Document next() {
+                throw new IllegalStateException("cursor id 42 not found");
+            }
+
+            @Override
+            public Document tryNext() {
+                throw new IllegalStateException("cursor id 42 not found");
+            }
+
+            @Override
+            public ServerAddress getServerAddress() {
+                return new ServerAddress();
+            }
+
+            @Override
+            public com.mongodb.ServerCursor getServerCursor() {
+                return null;
+            }
+        };
+        MongoClient client = findMongoClientWithCursor(calls, failing);
+
+        JsonObject start = findCursorRpc(client, 60, "start_find_cursor", cursorParams("users", 10, null));
+        String cursorId = start.getAsJsonObject("result").get("cursor_id").getAsString();
+        JsonObject fetch = findCursorRpc(client, 61, "fetch_find_cursor", cursorParams(null, 0, cursorId));
+        JsonObject again = findCursorRpc(client, 62, "fetch_find_cursor", cursorParams(null, 0, cursorId));
+
+        assertTrue(fetch.has("error"), fetch.toString());
+        assertTrue(calls.contains("close"));
+        // A dead cursor is dropped, so the slot is free again instead of being held until disconnect.
+        assertEquals("Find cursor not found", again.getAsJsonObject("error").get("message").getAsString());
+    }
+
+    @Test
+    void startFindCursorEnforcesPerOwnerBudget() {
+        List<String> calls = new ArrayList<>();
+        MongoClient client = recordingFindMongoClient(calls, List.of(new Document("n", 1)));
+
+        for (int index = 0; index < 16; index++) {
+            JsonObject started = findCursorRpc(client, 100 + index, "start_find_cursor", cursorParams("users", 100, null));
+            assertFalse(started.has("error"), started.toString());
+        }
+        JsonObject rejected = findCursorRpc(client, 200, "start_find_cursor", cursorParams("users", 100, null));
+
+        assertTrue(rejected.has("error"), rejected.toString());
+        assertEquals(
+            "MongoDB find cursor limit reached: 16",
+            rejected.getAsJsonObject("error").get("message").getAsString()
+        );
+    }
+
     @Test
     void updateDocumentsRpcUsesDocumentOverloads() {
         List<String> calls = new ArrayList<>();
@@ -1458,9 +1766,140 @@ class MongoAgentTest {
         return JsonParser.parseString(MongoAgent.handleRequest(request.toString(), client)).getAsJsonObject();
     }
 
+    private static JsonObject cursorParams(String collection, int batchSize, String cursorId) {
+        JsonObject params = new JsonObject();
+        params.addProperty("database", "app");
+        if (collection != null) {
+            params.addProperty("collection", collection);
+        }
+        if (batchSize > 0) {
+            params.addProperty("batch_size", batchSize);
+        }
+        if (cursorId != null) {
+            params.addProperty("cursor_id", cursorId);
+            params.addProperty("limit", 1);
+        }
+        return params;
+    }
+
+    private static JsonObject findCursorRpc(MongoClient client, int id, String method, JsonObject params) {
+        JsonObject request = new JsonObject();
+        request.addProperty("jsonrpc", "2.0");
+        request.addProperty("id", id);
+        request.addProperty("method", method);
+        request.add("params", params);
+        return JsonParser.parseString(MongoAgent.handleRequest(request.toString(), client)).getAsJsonObject();
+    }
+
+    @SuppressWarnings("unchecked")
+    private static MongoClient recordingFindMongoClient(List<String> calls, List<Document> documents) {
+        MongoCursor<Document> cursor = new MongoCursor<>() {
+            private int index;
+            private boolean closed;
+
+            @Override
+            public void close() {
+                closed = true;
+                calls.add("close");
+            }
+
+            @Override
+            public boolean hasNext() {
+                return !closed && index < documents.size();
+            }
+
+            @Override
+            public Document next() {
+                return documents.get(index++);
+            }
+
+            @Override
+            public Document tryNext() {
+                return hasNext() ? next() : null;
+            }
+
+            @Override
+            public com.mongodb.ServerAddress getServerAddress() {
+                return new com.mongodb.ServerAddress();
+            }
+
+            @Override
+            public com.mongodb.ServerCursor getServerCursor() {
+                return null;
+            }
+        };
+        return findMongoClientWithCursor(calls, cursor);
+    }
+
+    /** Same proxy chain as [recordingFindMongoClient], for cursors that fail mid-iteration. */
+    @SuppressWarnings("unchecked")
+    private static MongoClient findMongoClientWithCursor(List<String> calls, MongoCursor<Document> cursor) {
+        FindIterable<Document> iterable = (FindIterable<Document>) Proxy.newProxyInstance(
+            FindIterable.class.getClassLoader(),
+            new Class<?>[] {FindIterable.class},
+            (proxy, method, args) -> {
+                switch (method.getName()) {
+                    case "batchSize" -> {
+                        calls.add("batchSize:" + args[0]);
+                        return proxy;
+                    }
+                    case "projection", "sort", "collation", "filter", "skip", "limit" -> {
+                        calls.add(method.getName());
+                        return proxy;
+                    }
+                    case "iterator" -> {
+                        calls.add("iterator");
+                        return cursor;
+                    }
+                    default -> throw new UnsupportedOperationException(method.getName());
+                }
+            }
+        );
+        MongoCollection<Document> collection = (MongoCollection<Document>) Proxy.newProxyInstance(
+            MongoCollection.class.getClassLoader(),
+            new Class<?>[] {MongoCollection.class},
+            (proxy, method, args) -> {
+                if ("find".equals(method.getName())) {
+                    calls.add("find");
+                    return iterable;
+                }
+                throw new UnsupportedOperationException(method.getName());
+            }
+        );
+        MongoDatabase database = (MongoDatabase) Proxy.newProxyInstance(
+            MongoDatabase.class.getClassLoader(),
+            new Class<?>[] {MongoDatabase.class},
+            (proxy, method, args) -> {
+                if ("getCollection".equals(method.getName())) {
+                    calls.add("getCollection:" + args[0]);
+                    return collection;
+                }
+                throw new UnsupportedOperationException(method.getName());
+            }
+        );
+        return (MongoClient) Proxy.newProxyInstance(
+            MongoClient.class.getClassLoader(),
+            new Class<?>[] {MongoClient.class},
+            (proxy, method, args) -> {
+                if ("getDatabase".equals(method.getName())) {
+                    calls.add("getDatabase:" + args[0]);
+                    return database;
+                }
+                throw new UnsupportedOperationException(method.getName());
+            }
+        );
+    }
+
     @SuppressWarnings("unchecked")
     private static MongoClient recordingInsertMongoClient(
         List<String> calls, List<List<Document>> batches
+    ) {
+        return recordingInsertMongoClient(calls, batches, null);
+    }
+
+    @SuppressWarnings("unchecked")
+    private static MongoClient recordingInsertMongoClient(
+        List<String> calls, List<List<Document>> batches, List<Boolean> ordered
     ) {
         MongoCollection<Document> collection = (MongoCollection<Document>) Proxy.newProxyInstance(
             MongoCollection.class.getClassLoader(),
@@ -1470,6 +1909,13 @@ class MongoAgentTest {
                     List<Document> batch = ((List<Document>) args[0]).stream().map(Document::new).toList();
                     batches.add(batch);
                     calls.add("insertMany:" + batch.size());
+                    if (ordered != null) {
+                        boolean isOrdered = true;
+                        if (args.length > 1 && args[1] instanceof InsertManyOptions options) {
+                            isOrdered = options.isOrdered();
+                        }
+                        ordered.add(isOrdered);
+                    }
                     return null;
                 }
                 throw new UnsupportedOperationException(method.getName());