fix(mcp): Mongo maxTimeMS and a total response size cap (#1776)

lib/core/database/mongodb_service.dart

@@ -157,6 +157,7 @@ class MongoService {
     Map<String, dynamic>? sort,
     int? limit,
     int? skip,
+    int? maxTimeMs,
   }) async {
     if (!connection.isConnected) {
       throw StateError('Not connected to MongoDB');
@@ -183,7 +184,17 @@ class MongoService {
       // Never unbounded: `stream.toList()` below holds every document.
       selector.limit(effectiveMongoFindLimit(limit));
 
-      final stream = coll.find(selector);
+      // `maxTimeMS` makes the server stop a runaway query (`$regex`, `$expr`)
+      // itself; dropping the client future would leave it running. It needs
+      // the find command (MongoDB 3.6+); older servers get the plain find.
+      final stream = maxTimeMs != null &&
+              maxTimeMs > 0 &&
+              db.masterConnection.serverCapabilities.supportsOpMsg
+          ? coll.modernFind(
+              selector: selector,
+              findOptions: FindOptions(maxTimeMS: maxTimeMs),
+            )
+          : coll.find(selector);
       return await stream.toList();
     });
   }
@@ -243,11 +254,13 @@ class MongoService {
     String database,
     String collection, {
     Map<String, dynamic>? filter,
+    int? maxTimeMs,
   }) async {
     return _withDb(connection, database, (db) async {
       final result = await db.runCommand(<String, Object>{
         'count': collection,
         if (filter != null && filter.isNotEmpty) 'query': filter,
+        if (maxTimeMs != null && maxTimeMs > 0) 'maxTimeMS': maxTimeMs,
       });
       final n = result['n'];
       if (n is int) return n;

lib/core/mcp/mcp_mongo_service.dart

@@ -1,3 +1,4 @@
+import 'dart:convert';
 import 'dart:async';
 
 import 'package:querya_desktop/core/database/mongodb_connection.dart';
@@ -81,14 +82,20 @@ abstract interface class McpMongoSession {
 
   Future<List<String>> listCollections();
 
+  /// [maxTimeMs]: the server gives up on the query after this long.
   Future<List<Map<String, dynamic>>> find(
     String collection, {
     Map<String, dynamic>? filter,
     Map<String, dynamic>? sort,
     required int limit,
+    int? maxTimeMs,
   });
 
-  Future<int> count(String collection, {Map<String, dynamic>? filter});
+  Future<int> count(
+    String collection, {
+    Map<String, dynamic>? filter,
+    int? maxTimeMs,
+  });
 
   Future<void> close();
 }
@@ -129,14 +136,19 @@ class _LiveMongoSession implements McpMongoSession {
     Map<String, dynamic>? filter,
     Map<String, dynamic>? sort,
     required int limit,
+    int? maxTimeMs,
   }) =>
       _service.find(_connection, database, collection,
-          filter: filter, sort: sort, limit: limit);
+          filter: filter, sort: sort, limit: limit, maxTimeMs: maxTimeMs);
 
   @override
-  Future<int> count(String collection, {Map<String, dynamic>? filter}) =>
+  Future<int> count(
+    String collection, {
+    Map<String, dynamic>? filter,
+    int? maxTimeMs,
+  }) =>
       _service.countDocuments(_connection, database, collection,
-          filter: filter);
+          filter: filter, maxTimeMs: maxTimeMs);
 
   @override
   Future<void> close() => _connection.disconnect();
@@ -155,6 +167,7 @@ class McpMongoService {
     this.maxDocuments = 100,
     this.timeout = const Duration(seconds: 15),
     this.maxFieldChars = 4096,
+    this.maxTotalChars = 200000,
   })  : _access = access ?? McpAccessStore.instance,
         _loadConnections =
             loadConnections ?? (() => LocalDb.instance.getConnections()),
@@ -168,6 +181,10 @@ class McpMongoService {
   final Duration timeout;
   final int maxFieldChars;
 
+  /// Characters of all returned documents together (as JSON); more are cut
+  /// off and the result is marked truncated.
+  final int maxTotalChars;
+
   /// Shared MongoDB connections, without credentials.
   Future<List<McpConnectionInfo>> listConnections() async {
     final out = <McpConnectionInfo>[];
@@ -212,15 +229,28 @@ class McpMongoService {
         filter: filter,
         sort: sort,
         limit: n + 1,
+        maxTimeMs: timeout.inMilliseconds,
       );
       final cut = docs.length > n;
+      // Each field is capped, but 100 documents of many capped fields are
+      // still a lot: stop at a total size too.
+      final documents = <Map<String, Object?>>[];
+      var total = 0;
+      var cutBySize = false;
+      for (final d in (cut ? docs.sublist(0, n) : docs)) {
+        final shortened = _shorten(d) as Map<String, Object?>;
+        final size = jsonEncode(shortened, toEncodable: (o) => '$o').length;
+        if (documents.isNotEmpty && total + size > maxTotalChars) {
+          cutBySize = true;
+          break;
+        }
+        total += size;
+        documents.add(shortened);
+      }
       return McpMongoDocuments(
         collection: name,
-        documents: [
-          for (final d in (cut ? docs.sublist(0, n) : docs))
-            _shorten(d) as Map<String, Object?>,
-        ],
-        truncated: cut,
+        documents: documents,
+        truncated: cut || cutBySize,
       );
     });
   }
@@ -233,7 +263,8 @@ class McpMongoService {
     _guard(filter, 'filter');
     return _withSession(connectionId, (session) async {
       final name = await _knownCollection(session, collection);
-      return session.count(name, filter: filter);
+      return session.count(name,
+          filter: filter, maxTimeMs: timeout.inMilliseconds);
     });
   }
 

test/core/mcp/mcp_mongo_service_test.dart

@@ -33,6 +33,7 @@ class _FakeSession implements McpMongoSession {
   Map<String, dynamic>? lastFilter;
   Map<String, dynamic>? lastSort;
   int? lastLimit;
+  int? lastMaxTimeMs;
   var closed = false;
 
   @override
@@ -47,8 +48,10 @@ class _FakeSession implements McpMongoSession {
     Map<String, dynamic>? filter,
     Map<String, dynamic>? sort,
     required int limit,
+    int? maxTimeMs,
   }) async {
     calls.add('find $collection');
+    lastMaxTimeMs = maxTimeMs;
     lastFilter = filter;
     lastSort = sort;
     lastLimit = limit;
@@ -58,8 +61,10 @@ class _FakeSession implements McpMongoSession {
   }
 
   @override
-  Future<int> count(String collection, {Map<String, dynamic>? filter}) async {
+  Future<int> count(String collection,
+      {Map<String, dynamic>? filter, int? maxTimeMs}) async {
     calls.add('count $collection');
+    lastMaxTimeMs = maxTimeMs;
     lastFilter = filter;
     return docs.length;
   }
@@ -88,6 +93,7 @@ void main() {
     Set<int> shared = const {1},
     Duration timeout = const Duration(seconds: 5),
     int maxField = 4096,
+    int maxTotal = 200000,
   }) =>
       McpMongoService(
         access: _Access(shared),
@@ -102,6 +108,7 @@ void main() {
         },
         timeout: timeout,
         maxFieldChars: maxField,
+        maxTotalChars: maxTotal,
       );
 
   setUp(() {
@@ -130,6 +137,38 @@ void main() {
     expect(session.calls, isEmpty);
   });
 
+  test('the server is told to stop a query after the timeout (#1776)', () async {
+    final svc = service(timeout: const Duration(seconds: 7));
+
+    await svc.find(1, 'users');
+    expect(session.lastMaxTimeMs, 7000);
+
+    session.lastMaxTimeMs = null;
+    await svc.count(1, 'users');
+    expect(session.lastMaxTimeMs, 7000);
+  });
+
+  test('the documents together are capped, not only each field (#1776)',
+      () async {
+    session = _FakeSession(docs: [
+      for (var i = 0; i < 10; i++) {'_id': 'id$i', 'text': 'x' * 100},
+    ]);
+    final svc = service(maxTotal: 350);
+
+    final result = await svc.find(1, 'users', limit: 10);
+
+    expect(result.documents.length, lessThan(10));
+    expect(result.documents, isNotEmpty);
+    expect(result.truncated, isTrue);
+  });
+
+  test('a result under the total cap is returned whole (#1776)', () async {
+    final result = await service().find(1, 'users');
+
+    expect(result.documents.length, 3);
+    expect(result.truncated, isFalse);
+  });
+
   group('McpMongoGuard', () {
     test('plain filters pass', () {
       expect(

test/e2e/04_nosql_workspaces/e2e_mongo_explorer_and_crud_test.dart

@@ -42,6 +42,7 @@ class _FakeMongo extends MongoService {
     Map<String, dynamic>? sort,
     int? limit,
     int? skip,
+    int? maxTimeMs,
   }) async =>
       [for (final d in docs) Map<String, dynamic>.of(d)];
 
@@ -51,6 +52,7 @@ class _FakeMongo extends MongoService {
     String database,
     String collection, {
     Map<String, dynamic>? filter,
+    int? maxTimeMs,
   }) async =>
       docs.length;