chore: PR feedback about generous config defaults

This commit implements the PR feedback about being generous on the
config defaults and applying recommended restrictions on the tool
parameters for capping the memory usage.

Author: himanshusinghs

src/common/config.ts

@@ -9,6 +9,7 @@ import levenshtein from "ts-levenshtein";
 
 // From: https://github.com/mongodb-js/mongosh/blob/main/packages/cli-repl/src/arg-parser.ts
 const OPTIONS = {
+    number: ["maxDocumentsPerQuery", "maxBytesPerQuery"],
     string: [
         "apiBaseUrl",
         "apiClientId",
@@ -98,6 +99,7 @@ const OPTIONS = {
 
 interface Options {
     string: string[];
+    number: string[];
     boolean: string[];
     array: string[];
     alias: Record<string, string>;
@@ -106,6 +108,7 @@ interface Options {
 
 export const ALL_CONFIG_KEYS = new Set(
     (OPTIONS.string as readonly string[])
+        .concat(OPTIONS.number)
         .concat(OPTIONS.array)
         .concat(OPTIONS.boolean)
         .concat(Object.keys(OPTIONS.alias))
@@ -204,8 +207,8 @@ export const defaultUserConfig: UserConfig = {
     idleTimeoutMs: 10 * 60 * 1000, // 10 minutes
     notificationTimeoutMs: 9 * 60 * 1000, // 9 minutes
     httpHeaders: {},
-    maxDocumentsPerQuery: 10, // By default, we only fetch a maximum 10 documents per query / aggregation
-    maxBytesPerQuery: 1 * 1024 * 1024, // By default, we only return ~1 mb of data per query / aggregation
+    maxDocumentsPerQuery: 100, // By default, we only fetch a maximum 100 documents per query / aggregation
+    maxBytesPerQuery: 16 * 1024 * 1024, // By default, we only return ~16 mb of data per query / aggregation
     atlasTemporaryDatabaseUserLifetimeMs: 4 * 60 * 60 * 1000, // 4 hours
 };
 

src/helpers/collectCursorUntilMaxBytes.ts

@@ -0,0 +1,112 @@
+import { calculateObjectSize } from "bson";
+import type { AggregationCursor, FindCursor } from "mongodb";
+
+export function getResponseBytesLimit(
+    toolResponseBytesLimit: number | undefined | null,
+    configuredMaxBytesPerQuery: unknown
+): {
+    cappedBy: "config.maxBytesPerQuery" | "tool.responseBytesLimit" | undefined;
+    limit: number;
+} {
+    const configuredLimit: number = parseInt(String(configuredMaxBytesPerQuery), 10);
+
+    // Setting configured maxBytesPerQuery to negative, zero or nullish is
+    // equivalent to disabling the max limit applied on documents
+    const configuredLimitIsNotApplicable = Number.isNaN(configuredLimit) || configuredLimit <= 0;
+
+    // It's possible to have tool parameter responseBytesLimit as null or
+    // negative values in which case we consider that no limit is to be
+    // applied from tool call perspective unless we have a maxBytesPerQuery
+    // configured.
+    const toolResponseLimitIsNotApplicable = typeof toolResponseBytesLimit !== "number" || toolResponseBytesLimit <= 0;
+
+    if (configuredLimitIsNotApplicable) {
+        return {
+            cappedBy: toolResponseLimitIsNotApplicable ? undefined : "tool.responseBytesLimit",
+            limit: toolResponseLimitIsNotApplicable ? 0 : toolResponseBytesLimit,
+        };
+    }
+
+    if (toolResponseLimitIsNotApplicable) {
+        return { cappedBy: "config.maxBytesPerQuery", limit: configuredLimit };
+    }
+
+    return {
+        cappedBy: configuredLimit < toolResponseBytesLimit ? "config.maxBytesPerQuery" : "tool.responseBytesLimit",
+        limit: Math.min(toolResponseBytesLimit, configuredLimit),
+    };
+}
+
+/**
+ * This function attempts to put a guard rail against accidental memory overflow
+ * on the MCP server.
+ *
+ * The cursor is iterated until we can predict that fetching next doc won't
+ * exceed the derived limit on number of bytes for the tool call. The derived
+ * limit takes into account the limit provided from the Tool's interface and the
+ * configured maxBytesPerQuery for the server.
+ */
+export async function collectCursorUntilMaxBytesLimit<T = unknown>({
+    cursor,
+    toolResponseBytesLimit,
+    configuredMaxBytesPerQuery,
+    abortSignal,
+}: {
+    cursor: FindCursor<T> | AggregationCursor<T>;
+    toolResponseBytesLimit: number | undefined | null;
+    configuredMaxBytesPerQuery: unknown;
+    abortSignal?: AbortSignal;
+}): Promise<{ cappedBy: "config.maxBytesPerQuery" | "tool.responseBytesLimit" | undefined; documents: T[] }> {
+    const { limit: maxBytesPerQuery, cappedBy } = getResponseBytesLimit(
+        toolResponseBytesLimit,
+        configuredMaxBytesPerQuery
+    );
+
+    // It's possible to have no limit on the cursor response by setting both the
+    // config.maxBytesPerQuery and tool.responseBytesLimit to nullish or
+    // negative values.
+    if (maxBytesPerQuery <= 0) {
+        return {
+            cappedBy,
+            documents: await cursor.toArray(),
+        };
+    }
+
+    let wasCapped: boolean = false;
+    let totalBytes = 0;
+    let biggestDocSizeSoFar = 0;
+    const bufferedDocuments: T[] = [];
+    while (true) {
+        if (abortSignal?.aborted) {
+            break;
+        }
+
+        // This is an eager attempt to validate that fetching the next document
+        // won't exceed the applicable limit.
+        if (totalBytes + biggestDocSizeSoFar >= maxBytesPerQuery) {
+            wasCapped = true;
+            break;
+        }
+
+        // If the cursor is empty then there is nothing for us to do anymore.
+        const nextDocument = await cursor.tryNext();
+        if (!nextDocument) {
+            break;
+        }
+
+        const nextDocumentSize = calculateObjectSize(nextDocument);
+        if (totalBytes + nextDocumentSize >= maxBytesPerQuery) {
+            wasCapped = true;
+            break;
+        }
+
+        totalBytes += nextDocumentSize;
+        biggestDocSizeSoFar = Math.max(biggestDocSizeSoFar, nextDocumentSize);
+        bufferedDocuments.push(nextDocument);
+    }
+
+    return {
+        cappedBy: wasCapped ? cappedBy : undefined,
+        documents: bufferedDocuments,
+    };
+}

src/helpers/constants.ts

@@ -12,3 +12,15 @@ export const QUERY_COUNT_MAX_TIME_MS_CAP: number = 10_000;
  * aggregation.
  */
 export const AGG_COUNT_MAX_TIME_MS_CAP: number = 60_000;
+
+export const ONE_MB: number = 1 * 1024 * 1024;
+
+/**
+ * A map of applied limit on cursors to a text that is supposed to be sent as
+ * response to LLM
+ */
+export const CURSOR_LIMITS_TO_LLM_TEXT = {
+    "config.maxDocumentsPerQuery": "server's configured - maxDocumentsPerQuery",
+    "config.maxBytesPerQuery": "server's configured - maxBytesPerQuery",
+    "tool.responseBytesLimit": "tool's parameter - responseBytesLimit",
+} as const;

src/helpers/iterateCursor.ts

@@ -8,12 +8,16 @@ import { formatUntrustedData } from "../../tool.js";
 import { checkIndexUsage } from "../../../helpers/indexCheck.js";
 import { type Document, EJSON } from "bson";
 import { ErrorCodes, MongoDBError } from "../../../common/errors.js";
-import { iterateCursorUntilMaxBytes } from "../../../helpers/iterateCursor.js";
+import { collectCursorUntilMaxBytesLimit } from "../../../helpers/collectCursorUntilMaxBytes.js";
 import { operationWithFallback } from "../../../helpers/operationWithFallback.js";
-import { AGG_COUNT_MAX_TIME_MS_CAP } from "../../../helpers/constants.js";
+import { AGG_COUNT_MAX_TIME_MS_CAP, ONE_MB, CURSOR_LIMITS_TO_LLM_TEXT } from "../../../helpers/constants.js";
 
 export const AggregateArgs = {
     pipeline: z.array(z.object({}).passthrough()).describe("An array of aggregation stages to execute"),
+    responseBytesLimit: z.number().optional().default(ONE_MB).describe(`\
+The maximum number of bytes to return in the response. This value is capped by the server’s configured maxBytesPerQuery and cannot be exceeded. \
+Note to LLM: If the entire aggregation result is required, use the "export" tool instead of increasing this limit.\
+`),
 };
 
 export class AggregateTool extends MongoDBToolBase {
@@ -26,7 +30,7 @@ export class AggregateTool extends MongoDBToolBase {
     public operationType: OperationType = "read";
 
     protected async execute(
-        { database, collection, pipeline }: ToolArgs<typeof this.argsShape>,
+        { database, collection, pipeline, responseBytesLimit }: ToolArgs<typeof this.argsShape>,
         { signal }: ToolExecutionContext
     ): Promise<CallToolResult> {
         let aggregationCursor: AggregationCursor | undefined;
@@ -50,29 +54,36 @@ export class AggregateTool extends MongoDBToolBase {
             }
             aggregationCursor = provider.aggregate(database, collection, cappedResultsPipeline);
 
-            const [totalDocuments, documents] = await Promise.all([
+            const [totalDocuments, cursorResults] = await Promise.all([
                 this.countAggregationResultDocuments({ provider, database, collection, pipeline }),
-                iterateCursorUntilMaxBytes({
+                collectCursorUntilMaxBytesLimit({
                     cursor: aggregationCursor,
-                    maxBytesPerQuery: this.config.maxBytesPerQuery,
+                    configuredMaxBytesPerQuery: this.config.maxBytesPerQuery,
+                    toolResponseBytesLimit: responseBytesLimit,
                     abortSignal: signal,
                 }),
             ]);
 
-            let messageDescription = `\
-The aggregation resulted in ${totalDocuments === undefined ? "indeterminable number of" : totalDocuments} documents.\
-`;
-            if (documents.length) {
-                messageDescription += ` \
-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.\
-`;
-            }
+            // If the total number of documents that the aggregation would've
+            // resulted in would be greater than the configured
+            // maxDocumentsPerQuery then we know for sure that the results were
+            // capped.
+            const aggregationResultsCappedByMaxDocumentsLimit =
+                this.config.maxDocumentsPerQuery > 0 &&
+                !!totalDocuments &&
+                totalDocuments > this.config.maxDocumentsPerQuery;
 
             return {
                 content: formatUntrustedData(
-                    messageDescription,
-                    documents.length > 0 ? EJSON.stringify(documents) : undefined
+                    this.generateMessage({
+                        aggResultsCount: totalDocuments,
+                        documents: cursorResults.documents,
+                        appliedLimits: [
+                            aggregationResultsCappedByMaxDocumentsLimit ? "config.maxDocumentsPerQuery" : undefined,
+                            cursorResults.cappedBy,
+                        ].filter((limit): limit is keyof typeof CURSOR_LIMITS_TO_LLM_TEXT => !!limit),
+                    }),
+                    cursorResults.documents.length > 0 ? EJSON.stringify(cursorResults.documents) : undefined
                 ),
             };
         } finally {
@@ -132,4 +143,26 @@ Note to LLM: If entire aggregation result is needed then use "export" tool to ex
             return totalDocuments;
         }, undefined);
     }
+
+    private generateMessage({
+        aggResultsCount,
+        documents,
+        appliedLimits,
+    }: {
+        aggResultsCount: number | undefined;
+        documents: unknown[];
+        appliedLimits: (keyof typeof CURSOR_LIMITS_TO_LLM_TEXT)[];
+    }): string {
+        const appliedLimitText = appliedLimits.length
+            ? `\
+while respecting the applied limits of ${appliedLimits.map((limit) => CURSOR_LIMITS_TO_LLM_TEXT[limit]).join(", ")}. \
+Note to LLM: If the entire query result is required then use "export" tool to export the query results.\
+`
+            : "";
+
+        return `\
+The aggregation resulted in ${aggResultsCount === undefined ? "indeterminable number of" : aggResultsCount} documents. \
+Returning ${documents.length} documents${appliedLimitText ? ` ${appliedLimitText}` : "."}\
+`;
+    }
 }