chore: PR feedback

1. implements disabling maxDocumentsPerQuery and maxBytesPerQuery
2. use correct numbers for bytes calculation

Author: himanshusinghs

src/common/config.ts

@@ -183,7 +183,7 @@ export const defaultUserConfig: UserConfig = {
     notificationTimeoutMs: 540000, // 9 minutes
     httpHeaders: {},
     maxDocumentsPerQuery: 10, // By default, we only fetch a maximum 10 documents per query / aggregation
-    maxBytesPerQuery: 1 * 1000 * 1000, // By default, we only return ~1 mb of data per query / aggregation
+    maxBytesPerQuery: 1 * 1024 * 1024, // By default, we only return ~1 mb of data per query / aggregation
 };
 
 export const config = setupUserConfig({

src/helpers/iterateCursor.ts

@@ -12,6 +12,12 @@ export async function iterateCursorUntilMaxBytes(
     cursor: FindCursor<unknown> | AggregationCursor<unknown>,
     maxBytesPerQuery: number
 ): Promise<unknown[]> {
+    // Setting configured limit to zero or negative is equivalent to disabling
+    // the max bytes limit applied on tool responses.
+    if (maxBytesPerQuery <= 0) {
+        return await cursor.toArray();
+    }
+
     let biggestDocSizeSoFar = 0;
     let totalBytes = 0;
     const bufferedDocuments: unknown[] = [];

src/tools/mongodb/read/aggregate.ts

@@ -45,10 +45,11 @@ export class AggregateTool extends MongoDBToolBase {
                 });
             }
 
-            const cappedResultsPipeline = [...pipeline, { $limit: this.config.maxDocumentsPerQuery }];
-            aggregationCursor = provider
-                .aggregate(database, collection, cappedResultsPipeline)
-                .batchSize(this.config.maxDocumentsPerQuery);
+            const cappedResultsPipeline = [...pipeline];
+            if (this.config.maxDocumentsPerQuery > 0) {
+                cappedResultsPipeline.push({ $limit: this.config.maxDocumentsPerQuery });
+            }
+            aggregationCursor = provider.aggregate(database, collection, cappedResultsPipeline);
 
             const [totalDocuments, documents] = await Promise.all([
                 this.countAggregationResultDocuments({ provider, database, collection, pipeline }),

src/tools/mongodb/read/find.ts

@@ -61,12 +61,12 @@ export class FindTool extends MongoDBToolBase {
                 });
             }
 
-            const limitOnFindCursor = Math.min(limit ?? Number.POSITIVE_INFINITY, this.config.maxDocumentsPerQuery);
+            const limitOnFindCursor = this.getLimitForFindCursor(limit);
+
             findCursor = provider.find(database, collection, filter, {
                 projection,
                 limit: limitOnFindCursor,
                 sort,
-                batchSize: limitOnFindCursor,
             });
 
             const [queryResultsCount, documents] = await Promise.all([
@@ -75,7 +75,8 @@ export class FindTool extends MongoDBToolBase {
                         provider.countDocuments(database, collection, filter, {
                             // We should be counting documents that the original
                             // query would have yielded which is why we don't
-                            // use `limitOnFindCursor` calculated above.
+                            // use `limitOnFindCursor` calculated above, only
+                            // the limit provided to the tool.
                             limit,
                             maxTimeMS: QUERY_COUNT_MAX_TIME_MS_CAP,
                         }),
@@ -104,4 +105,17 @@ Note to LLM: If entire query result is needed then use "export" tool to export t
             await findCursor?.close();
         }
     }
+
+    private getLimitForFindCursor(providedLimit: number | undefined): number | undefined {
+        const configuredLimit = this.config.maxDocumentsPerQuery;
+        // Setting configured limit to negative or zero is equivalent to
+        // disabling the max limit applied on documents
+        if (configuredLimit <= 0) {
+            return providedLimit;
+        }
+
+        return providedLimit === null || providedLimit === undefined
+            ? configuredLimit
+            : Math.min(providedLimit, configuredLimit);
+    }
 }

tests/integration/tools/mongodb/read/aggregate.test.ts

@@ -5,7 +5,7 @@ import {
     getResponseContent,
     defaultTestConfig,
 } from "../../../helpers.js";
-import { describe, expect, it, vi } from "vitest";
+import { beforeEach, describe, expect, it, vi, afterEach } from "vitest";
 import { describeWithMongoDB, getDocsFromUntrustedContent, validateAutoConnectBehavior } from "../mongodbHelpers.js";
 import * as constants from "../../../../../src/helpers/constants.js";
 import { freshInsertDocuments } from "./find.test.js";
@@ -142,90 +142,94 @@ describeWithMongoDB("aggregate tool", (integration) => {
             expectedResponse: "The aggregation resulted in 0 documents",
         };
     });
+
+    describe("when counting documents exceed the configured count maxTimeMS", () => {
+        beforeEach(async () => {
+            await freshInsertDocuments({
+                collection: integration.mongoClient().db(integration.randomDbName()).collection("people"),
+                count: 1000,
+                documentMapper(index) {
+                    return { name: `Person ${index}`, age: index };
+                },
+            });
+        });
+
+        afterEach(() => {
+            vi.resetAllMocks();
+        });
+
+        it("should abort discard count operation and respond with indeterminable count", async () => {
+            vi.spyOn(constants, "AGG_COUNT_MAX_TIME_MS_CAP", "get").mockReturnValue(0.1);
+            await integration.connectMcpClient();
+            const response = await integration.mcpClient().callTool({
+                name: "aggregate",
+                arguments: {
+                    database: integration.randomDbName(),
+                    collection: "people",
+                    pipeline: [{ $match: { age: { $gte: 10 } } }, { $sort: { name: -1 } }],
+                },
+            });
+            const content = getResponseContent(response);
+            expect(content).toContain("The aggregation resulted in indeterminable number of documents");
+            expect(content).toContain(`Returning 10 documents while respecting the applied limits.`);
+            const docs = getDocsFromUntrustedContent(content);
+            expect(docs[0]).toEqual(
+                expect.objectContaining({
+                    _id: expect.any(Object) as object,
+                    name: "Person 999",
+                    age: 999,
+                })
+            );
+            expect(docs[1]).toEqual(
+                expect.objectContaining({
+                    _id: expect.any(Object) as object,
+                    name: "Person 998",
+                    age: 998,
+                })
+            );
+        });
+    });
 });
 
 describeWithMongoDB(
     "aggregate tool with configured max documents per query",
     (integration) => {
-        describe("when the aggregation results are larger than the configured limit", () => {
-            it("should return documents limited to the configured limit", async () => {
-                await freshInsertDocuments({
-                    collection: integration.mongoClient().db(integration.randomDbName()).collection("people"),
-                    count: 1000,
-                    documentMapper(index) {
-                        return { name: `Person ${index}`, age: index };
-                    },
-                });
-                await integration.connectMcpClient();
-                const response = await integration.mcpClient().callTool({
-                    name: "aggregate",
-                    arguments: {
-                        database: integration.randomDbName(),
-                        collection: "people",
-                        pipeline: [{ $match: { age: { $gte: 10 } } }, { $sort: { name: -1 } }],
-                    },
-                });
-
-                const content = getResponseContent(response);
-                expect(content).toContain("The aggregation resulted in 990 documents");
-                expect(content).toContain(`Returning 20 documents while respecting the applied limits.`);
-                const docs = getDocsFromUntrustedContent(content);
-                expect(docs[0]).toEqual(
-                    expect.objectContaining({
-                        _id: expect.any(Object) as object,
-                        name: "Person 999",
-                        age: 999,
-                    })
-                );
-                expect(docs[1]).toEqual(
-                    expect.objectContaining({
-                        _id: expect.any(Object) as object,
-                        name: "Person 998",
-                        age: 998,
-                    })
-                );
+        it("should return documents limited to the configured limit", async () => {
+            await freshInsertDocuments({
+                collection: integration.mongoClient().db(integration.randomDbName()).collection("people"),
+                count: 1000,
+                documentMapper(index) {
+                    return { name: `Person ${index}`, age: index };
+                },
             });
-        });
-
-        describe("when counting documents exceed the configured count maxTimeMS", () => {
-            it("should abort discard count operation and respond with indeterminable count", async () => {
-                vi.spyOn(constants, "AGG_COUNT_MAX_TIME_MS_CAP", "get").mockReturnValue(0.1);
-                await freshInsertDocuments({
-                    collection: integration.mongoClient().db(integration.randomDbName()).collection("people"),
-                    count: 1000,
-                    documentMapper(index) {
-                        return { name: `Person ${index}`, age: index };
-                    },
-                });
-                await integration.connectMcpClient();
-                const response = await integration.mcpClient().callTool({
-                    name: "aggregate",
-                    arguments: {
-                        database: integration.randomDbName(),
-                        collection: "people",
-                        pipeline: [{ $match: { age: { $gte: 10 } } }, { $sort: { name: -1 } }],
-                    },
-                });
-                const content = getResponseContent(response);
-                expect(content).toContain("The aggregation resulted in indeterminable number of documents");
-                expect(content).toContain(`Returning 20 documents while respecting the applied limits.`);
-                const docs = getDocsFromUntrustedContent(content);
-                expect(docs[0]).toEqual(
-                    expect.objectContaining({
-                        _id: expect.any(Object) as object,
-                        name: "Person 999",
-                        age: 999,
-                    })
-                );
-                expect(docs[1]).toEqual(
-                    expect.objectContaining({
-                        _id: expect.any(Object) as object,
-                        name: "Person 998",
-                        age: 998,
-                    })
-                );
-                vi.resetAllMocks();
+            await integration.connectMcpClient();
+            const response = await integration.mcpClient().callTool({
+                name: "aggregate",
+                arguments: {
+                    database: integration.randomDbName(),
+                    collection: "people",
+                    pipeline: [{ $match: { age: { $gte: 10 } } }, { $sort: { name: -1 } }],
+                },
             });
+
+            const content = getResponseContent(response);
+            expect(content).toContain("The aggregation resulted in 990 documents");
+            expect(content).toContain(`Returning 20 documents while respecting the applied limits.`);
+            const docs = getDocsFromUntrustedContent(content);
+            expect(docs[0]).toEqual(
+                expect.objectContaining({
+                    _id: expect.any(Object) as object,
+                    name: "Person 999",
+                    age: 999,
+                })
+            );
+            expect(docs[1]).toEqual(
+                expect.objectContaining({
+                    _id: expect.any(Object) as object,
+                    name: "Person 998",
+                    age: 998,
+                })
+            );
         });
     },
     () => ({ ...defaultTestConfig, maxDocumentsPerQuery: 20 })
@@ -234,30 +238,57 @@ describeWithMongoDB(
 describeWithMongoDB(
     "aggregate tool with configured max bytes per query",
     (integration) => {
-        describe("when the provided maxBytesPerQuery is hit", () => {
-            it("should return only the documents that could fit in maxBytesPerQuery limit", async () => {
-                await freshInsertDocuments({
-                    collection: integration.mongoClient().db(integration.randomDbName()).collection("people"),
-                    count: 1000,
-                    documentMapper(index) {
-                        return { name: `Person ${index}`, age: index };
-                    },
-                });
-                await integration.connectMcpClient();
-                const response = await integration.mcpClient().callTool({
-                    name: "aggregate",
-                    arguments: {
-                        database: integration.randomDbName(),
-                        collection: "people",
-                        pipeline: [{ $match: { age: { $gte: 10 } } }, { $sort: { name: -1 } }],
-                    },
-                });
-
-                const content = getResponseContent(response);
-                expect(content).toContain("The aggregation resulted in 990 documents");
-                expect(content).toContain(`Returning 1 documents while respecting the applied limits.`);
+        it("should return only the documents that could fit in maxBytesPerQuery limit", async () => {
+            await freshInsertDocuments({
+                collection: integration.mongoClient().db(integration.randomDbName()).collection("people"),
+                count: 1000,
+                documentMapper(index) {
+                    return { name: `Person ${index}`, age: index };
+                },
+            });
+            await integration.connectMcpClient();
+            const response = await integration.mcpClient().callTool({
+                name: "aggregate",
+                arguments: {
+                    database: integration.randomDbName(),
+                    collection: "people",
+                    pipeline: [{ $match: { age: { $gte: 10 } } }, { $sort: { name: -1 } }],
+                },
             });
+
+            const content = getResponseContent(response);
+            expect(content).toContain("The aggregation resulted in 990 documents");
+            expect(content).toContain(`Returning 1 documents while respecting the applied limits.`);
         });
     },
     () => ({ ...defaultTestConfig, maxBytesPerQuery: 100 })
 );
+
+describeWithMongoDB(
+    "aggregate tool with disabled max documents and max bytes per query",
+    (integration) => {
+        it("should return all the documents", async () => {
+            await freshInsertDocuments({
+                collection: integration.mongoClient().db(integration.randomDbName()).collection("people"),
+                count: 1000,
+                documentMapper(index) {
+                    return { name: `Person ${index}`, age: index };
+                },
+            });
+            await integration.connectMcpClient();
+            const response = await integration.mcpClient().callTool({
+                name: "aggregate",
+                arguments: {
+                    database: integration.randomDbName(),
+                    collection: "people",
+                    pipeline: [{ $match: { age: { $gte: 10 } } }, { $sort: { name: -1 } }],
+                },
+            });
+
+            const content = getResponseContent(response);
+            expect(content).toContain("The aggregation resulted in 990 documents");
+            expect(content).toContain(`Returning 990 documents while respecting the applied limits.`);
+        });
+    },
+    () => ({ ...defaultTestConfig, maxDocumentsPerQuery: -1, maxBytesPerQuery: -1 })
+);