fix: add guards against possible memory overflow

himanshusinghs · himanshusinghs · commit 042aa1d9589a · 2025-09-09T16:47:56.000+02:00
Targets find and aggregate tool and does the following to avoid the
memory overflow possibility:
1. Adds a configurable limit to restrict the number of documents fetched
   per query / aggregation.
2. Adds an iterator that keeps track of bytes consumed in memory by the
   retrieved documents and cuts off the iteration when there is a
   possibility of overflow. The overflow is based on configured
   maxBytesPerQuery parameter which defaults to 1MB.
diff --git a/src/common/config.ts b/src/common/config.ts
@@ -161,6 +161,8 @@ export interface UserConfig extends CliOptions {
     loggers: Array<"stderr" | "disk" | "mcp">;
     idleTimeoutMs: number;
     notificationTimeoutMs: number;
+    maxDocumentsPerQuery: number;
+    maxBytesPerQuery: number;
 }
 
 export const defaultUserConfig: UserConfig = {
@@ -180,6 +182,8 @@ export const defaultUserConfig: UserConfig = {
     idleTimeoutMs: 600000, // 10 minutes
     notificationTimeoutMs: 540000, // 9 minutes
     httpHeaders: {},
+    maxDocumentsPerQuery: 50,
+    maxBytesPerQuery: 1 * 1000 * 1000, // 1 mb
 };
 
 export const config = setupUserConfig({
diff --git a/src/helpers/iterateCursor.ts b/src/helpers/iterateCursor.ts
@@ -0,0 +1,39 @@
+import { calculateObjectSize } from "bson";
+import type { AggregationCursor, FindCursor } from "mongodb";
+
+/**
+ * This function attempts to put a guard rail against accidental memory over
+ * flow on the MCP server.
+ *
+ * The cursor is iterated until we can predict that fetching next doc won't
+ * exceed the maxBytesPerQuery limit.
+ */
+export async function iterateCursorUntilMaxBytes(
+    cursor: FindCursor<unknown> | AggregationCursor<unknown>,
+    maxBytesPerQuery: number
+): Promise<unknown[]> {
+    let biggestDocSizeSoFar = 0;
+    let totalBytes = 0;
+    const bufferedDocuments: unknown[] = [];
+    while (true) {
+        if (totalBytes + biggestDocSizeSoFar >= maxBytesPerQuery) {
+            break;
+        }
+
+        const nextDocument = await cursor.tryNext();
+        if (!nextDocument) {
+            break;
+        }
+
+        const nextDocumentSize = calculateObjectSize(nextDocument);
+        if (totalBytes + nextDocumentSize >= maxBytesPerQuery) {
+            break;
+        }
+
+        totalBytes += nextDocumentSize;
+        biggestDocSizeSoFar = Math.max(biggestDocSizeSoFar, nextDocumentSize);
+        bufferedDocuments.push(nextDocument);
+    }
+
+    return bufferedDocuments;
+}
diff --git a/src/helpers/operationWithFallback.ts b/src/helpers/operationWithFallback.ts
@@ -0,0 +1,12 @@
+type OperationCallback<OperationResult> = () => Promise<OperationResult>;
+
+export async function operationWithFallback<OperationResult, FallbackValue>(
+    performOperation: OperationCallback<OperationResult>,
+    fallback: FallbackValue
+): Promise<OperationResult | FallbackValue> {
+    try {
+        return await performOperation();
+    } catch {
+        return fallback;
+    }
+}
diff --git a/src/tools/mongodb/read/aggregate.ts b/src/tools/mongodb/read/aggregate.ts
@@ -1,11 +1,21 @@
 import { z } from "zod";
+import type { AggregationCursor } from "mongodb";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
 import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
 import type { ToolArgs, OperationType } from "../../tool.js";
 import { formatUntrustedData } from "../../tool.js";
 import { checkIndexUsage } from "../../../helpers/indexCheck.js";
-import { EJSON } from "bson";
+import { type Document, EJSON } from "bson";
 import { ErrorCodes, MongoDBError } from "../../../common/errors.js";
+import { iterateCursorUntilMaxBytes } from "../../../helpers/iterateCursor.js";
+import { operationWithFallback } from "../../../helpers/operationWithFallback.js";
+
+/**
+ * A cap for the maxTimeMS used for counting resulting documents of an
+ * aggregation.
+ */
+const AGG_COUNT_MAX_TIME_MS_CAP = 60_000;
 
 export const AggregateArgs = {
     pipeline: z.array(z.object({}).passthrough()).describe("An array of aggregation stages to execute"),
@@ -25,27 +35,43 @@ export class AggregateTool extends MongoDBToolBase {
         collection,
         pipeline,
     }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
-        const provider = await this.ensureConnected();
+        let aggregationCursor: AggregationCursor | undefined;
+        try {
+            const provider = await this.ensureConnected();
 
-        this.assertOnlyUsesPermittedStages(pipeline);
+            this.assertOnlyUsesPermittedStages(pipeline);
 
-        // Check if aggregate operation uses an index if enabled
-        if (this.config.indexCheck) {
-            await checkIndexUsage(provider, database, collection, "aggregate", async () => {
-                return provider
-                    .aggregate(database, collection, pipeline, {}, { writeConcern: undefined })
-                    .explain("queryPlanner");
-            });
-        }
+            // Check if aggregate operation uses an index if enabled
+            if (this.config.indexCheck) {
+                await checkIndexUsage(provider, database, collection, "aggregate", async () => {
+                    return provider
+                        .aggregate(database, collection, pipeline, {}, { writeConcern: undefined })
+                        .explain("queryPlanner");
+                });
+            }
+
+            const cappedResultsPipeline = [...pipeline, { $limit: this.config.maxDocumentsPerQuery }];
+            aggregationCursor = provider
+                .aggregate(database, collection, cappedResultsPipeline)
+                .batchSize(this.config.maxDocumentsPerQuery);
 
-        const documents = await provider.aggregate(database, collection, pipeline).toArray();
+            const [totalDocuments, documents] = await Promise.all([
+                this.countAggregationResultDocuments({ provider, database, collection, pipeline }),
+                iterateCursorUntilMaxBytes(aggregationCursor, this.config.maxBytesPerQuery),
+            ]);
 
-        return {
-            content: formatUntrustedData(
-                `The aggregation resulted in ${documents.length} documents.`,
-                documents.length > 0 ? EJSON.stringify(documents) : undefined
-            ),
-        };
+            const messageDescription = `\
+    The aggregation resulted in ${totalDocuments === undefined ? "indeterminable number of" : totalDocuments} documents. \
+    Returning ${documents.length} documents while respecting the applied limits. \
+    Note to LLM: If entire aggregation result is needed then use "export" tool to export the aggregation results.\
+    `;
+
+            return {
+                content: formatUntrustedData(messageDescription, EJSON.stringify(documents)),
+            };
+        } finally {
+            await aggregationCursor?.close();
+        }
     }
 
     private assertOnlyUsesPermittedStages(pipeline: Record<string, unknown>[]): void {
@@ -62,4 +88,35 @@ export class AggregateTool extends MongoDBToolBase {
             }
         }
     }
+
+    private async countAggregationResultDocuments({
+        provider,
+        database,
+        collection,
+        pipeline,
+    }: {
+        provider: NodeDriverServiceProvider;
+        database: string;
+        collection: string;
+        pipeline: Document[];
+    }): Promise<number | undefined> {
+        const resultsCountAggregation = [...pipeline, { $count: "totalDocuments" }];
+        return await operationWithFallback(async (): Promise<number | undefined> => {
+            const aggregationResults = await provider
+                .aggregate(database, collection, resultsCountAggregation)
+                .maxTimeMS(AGG_COUNT_MAX_TIME_MS_CAP)
+                .toArray();
+
+            const documentWithCount: unknown = aggregationResults.length === 1 ? aggregationResults[0] : undefined;
+            const totalDocuments =
+                documentWithCount &&
+                typeof documentWithCount === "object" &&
+                "totalDocuments" in documentWithCount &&
+                typeof documentWithCount.totalDocuments === "number"
+                    ? documentWithCount.totalDocuments
+                    : undefined;
+
+            return totalDocuments;
+        }, undefined);
+    }
 }
diff --git a/src/tools/mongodb/read/find.ts b/src/tools/mongodb/read/find.ts
@@ -3,9 +3,20 @@ import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
 import type { ToolArgs, OperationType } from "../../tool.js";
 import { formatUntrustedData } from "../../tool.js";
-import type { SortDirection } from "mongodb";
+import type { FindCursor, SortDirection } from "mongodb";
 import { checkIndexUsage } from "../../../helpers/indexCheck.js";
 import { EJSON } from "bson";
+import { iterateCursorUntilMaxBytes } from "../../../helpers/iterateCursor.js";
+import { operationWithFallback } from "../../../helpers/operationWithFallback.js";
+
+/**
+ * A cap for the maxTimeMS used for FindCursor.countDocuments.
+ *
+ * The number is relatively smaller because we expect the count documents query
+ * to be finished sooner if not by the time the batch of documents is retrieved
+ * so that count documents query don't hold the final response back.
+ */
+const QUERY_COUNT_MAX_TIME_MS_CAP = 10_000;
 
 export const FindArgs = {
     filter: z
@@ -45,22 +56,50 @@ export class FindTool extends MongoDBToolBase {
         limit,
         sort,
     }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
-        const provider = await this.ensureConnected();
+        let findCursor: FindCursor<unknown> | undefined;
+        try {
+            const provider = await this.ensureConnected();
+
+            // Check if find operation uses an index if enabled
+            if (this.config.indexCheck) {
+                await checkIndexUsage(provider, database, collection, "find", async () => {
+                    return provider
+                        .find(database, collection, filter, { projection, limit, sort })
+                        .explain("queryPlanner");
+                });
+            }
 
-        // Check if find operation uses an index if enabled
-        if (this.config.indexCheck) {
-            await checkIndexUsage(provider, database, collection, "find", async () => {
-                return provider.find(database, collection, filter, { projection, limit, sort }).explain("queryPlanner");
+            const appliedLimit = Math.min(limit, this.config.maxDocumentsPerQuery);
+            findCursor = provider.find(database, collection, filter, {
+                projection,
+                limit: appliedLimit,
+                sort,
+                batchSize: appliedLimit,
             });
-        }
 
-        const documents = await provider.find(database, collection, filter, { projection, limit, sort }).toArray();
+            const [queryResultsCount, documents] = await Promise.all([
+                operationWithFallback(
+                    () =>
+                        provider.countDocuments(database, collection, filter, {
+                            limit,
+                            maxTimeMS: QUERY_COUNT_MAX_TIME_MS_CAP,
+                        }),
+                    undefined
+                ),
+                iterateCursorUntilMaxBytes(findCursor, this.config.maxBytesPerQuery),
+            ]);
 
-        return {
-            content: formatUntrustedData(
-                `Found ${documents.length} documents in the collection "${collection}".`,
-                documents.length > 0 ? EJSON.stringify(documents) : undefined
-            ),
-        };
+            const messageDescription = `\
+Query on collection "${collection}" resulted in ${queryResultsCount === undefined ? "indeterminable number of" : queryResultsCount} documents. \
+Returning ${documents.length} documents while respecting the applied limits. \
+Note to LLM: If entire query result is needed then use "export" tool to export the query results.\
+`;
+
+            return {
+                content: formatUntrustedData(messageDescription, EJSON.stringify(documents)),
+            };
+        } finally {
+            await findCursor?.close();
+        }
     }
 }
diff --git a/tests/unit/helpers/iterateCursor.test.ts b/tests/unit/helpers/iterateCursor.test.ts
@@ -0,0 +1,62 @@
+import { describe, it, expect, vi } from "vitest";
+import type { FindCursor } from "mongodb";
+import { calculateObjectSize } from "bson";
+import { iterateCursorUntilMaxBytes } from "../../../src/helpers/iterateCursor.js";
+
+describe("iterateCursorUntilMaxBytes", () => {
+    function createMockCursor(docs: unknown[]): FindCursor<unknown> {
+        let idx = 0;
+        return {
+            tryNext: vi.fn(() => {
+                if (idx < docs.length) {
+                    return Promise.resolve(docs[idx++]);
+                }
+                return Promise.resolve(null);
+            }),
+        } as unknown as FindCursor<unknown>;
+    }
+
+    it("returns all docs if under maxBytesPerQuery", async () => {
+        const docs = [{ a: 1 }, { b: 2 }];
+        const cursor = createMockCursor(docs);
+        const maxBytes = 10000;
+        const result = await iterateCursorUntilMaxBytes(cursor, maxBytes);
+        console.log("test result", result);
+        expect(result).toEqual(docs);
+    });
+
+    it("returns only docs that fit under maxBytesPerQuery", async () => {
+        const doc1 = { a: "x".repeat(100) };
+        const doc2 = { b: "y".repeat(1000) };
+        const docs = [doc1, doc2];
+        const cursor = createMockCursor(docs);
+        const maxBytes = calculateObjectSize(doc1) + 10;
+        const result = await iterateCursorUntilMaxBytes(cursor, maxBytes);
+        expect(result).toEqual([doc1]);
+    });
+
+    it("returns empty array if maxBytesPerQuery is smaller than even the first doc", async () => {
+        const docs = [{ a: "x".repeat(100) }];
+        const cursor = createMockCursor(docs);
+        const result = await iterateCursorUntilMaxBytes(cursor, 10);
+        expect(result).toEqual([]);
+    });
+
+    it("handles empty cursor", async () => {
+        const cursor = createMockCursor([]);
+        const result = await iterateCursorUntilMaxBytes(cursor, 1000);
+        expect(result).toEqual([]);
+    });
+
+    it("does not include a doc that would overflow the max bytes allowed", async () => {
+        const doc1 = { a: "x".repeat(10) };
+        const doc2 = { b: "y".repeat(1000) };
+        const docs = [doc1, doc2];
+        const cursor = createMockCursor(docs);
+        // Set maxBytes so that after doc1, biggestDocSizeSoFar would prevent fetching doc2
+        const maxBytes = calculateObjectSize(doc1) + calculateObjectSize(doc2) - 1;
+        const result = await iterateCursorUntilMaxBytes(cursor, maxBytes);
+        // Should only include doc1, not doc2
+        expect(result).toEqual([doc1]);
+    });
+});
diff --git a/tests/unit/helpers/operationWithFallback.test.ts b/tests/unit/helpers/operationWithFallback.test.ts
@@ -0,0 +1,24 @@
+import { describe, it, expect, vi } from "vitest";
+import { operationWithFallback } from "../../../src/helpers/operationWithFallback.js";
+
+describe("operationWithFallback", () => {
+    it("returns operation result when operation succeeds", async () => {
+        const successfulOperation = vi.fn().mockResolvedValue("success");
+        const fallbackValue = "fallback";
+
+        const result = await operationWithFallback(successfulOperation, fallbackValue);
+
+        expect(result).toBe("success");
+        expect(successfulOperation).toHaveBeenCalledOnce();
+    });
+
+    it("returns fallback value when operation throws an error", async () => {
+        const failingOperation = vi.fn().mockRejectedValue(new Error("Operation failed"));
+        const fallbackValue = "fallback";
+
+        const result = await operationWithFallback(failingOperation, fallbackValue);
+
+        expect(result).toBe("fallback");
+        expect(failingOperation).toHaveBeenCalledOnce();
+    });
+});
