Merge branch 'main' of github.com:mongodb-js/mongodb-mcp-server into gagik/add-depcheck

Author: gagik

README.md

@@ -1,7 +1,7 @@
 {
   "name": "mongodb-mcp-server",
   "description": "MongoDB Model Context Protocol Server",
-  "version": "1.0.0",
+  "version": "1.0.1",
   "type": "module",
   "exports": {
     ".": {

package-lock.json

@@ -17,7 +17,7 @@ export MDB_ACCURACY_RUN_ID=$(npx uuid v4)
 # specified in the command line. Such as:
 # npm run test:accuracy -- tests/accuracy/some-test.test.ts
 echo "Running accuracy tests with MDB_ACCURACY_RUN_ID '$MDB_ACCURACY_RUN_ID'"
-vitest --config vitest.config.ts --project=accuracy --coverage=false --run "$@"
+vitest --config vitest.config.ts --project=accuracy --coverage=false --no-file-parallelism --run "$@"
 
 # Preserving the exit code from test run to correctly notify in the CI
 # environments when the tests fail.

package.json

@@ -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))
@@ -175,6 +178,8 @@ export interface UserConfig extends CliOptions {
     loggers: Array<"stderr" | "disk" | "mcp">;
     idleTimeoutMs: number;
     notificationTimeoutMs: number;
+    maxDocumentsPerQuery: number;
+    maxBytesPerQuery: number;
     atlasTemporaryDatabaseUserLifetimeMs: number;
 }
 
@@ -202,6 +207,8 @@ export const defaultUserConfig: UserConfig = {
     idleTimeoutMs: 10 * 60 * 1000, // 10 minutes
     notificationTimeoutMs: 9 * 60 * 1000, // 9 minutes
     httpHeaders: {},
+    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
 };
 

scripts/accuracy/runAccuracyTests.sh

@@ -44,6 +44,7 @@ export const LogId = {
     mongodbConnectFailure: mongoLogId(1_004_001),
     mongodbDisconnectFailure: mongoLogId(1_004_002),
     mongodbConnectTry: mongoLogId(1_004_003),
+    mongodbCursorCloseError: mongoLogId(1_004_004),
 
     toolUpdateFailure: mongoLogId(1_005_001),
     resourceUpdateFailure: mongoLogId(1_005_002),

src/common/config.ts

@@ -1,5 +1,5 @@
 // This file was generated by scripts/updatePackageVersion.ts - Do not edit it manually.
 export const packageInfo = {
-    version: "1.0.0",
+    version: "1.0.1",
     mcpServerName: "MongoDB MCP Server",
 };

src/common/logger.ts

@@ -0,0 +1,103 @@
+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;
+    const bufferedDocuments: T[] = [];
+    while (true) {
+        if (abortSignal?.aborted) {
+            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;
+        bufferedDocuments.push(nextDocument);
+    }
+
+    return {
+        cappedBy: wasCapped ? cappedBy : undefined,
+        documents: bufferedDocuments,
+    };
+}

src/common/packageInfo.ts

@@ -0,0 +1,26 @@
+/**
+ * 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.
+ */
+export const QUERY_COUNT_MAX_TIME_MS_CAP: number = 10_000;
+
+/**
+ * A cap for the maxTimeMS used for counting resulting documents of an
+ * 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/collectCursorUntilMaxBytes.ts

@@ -0,0 +1,15 @@
+type EmptyObject = { [x: string]: never } | null | undefined;
+
+export function isObjectEmpty(value: object | null | undefined): value is EmptyObject {
+    if (!value) {
+        return true;
+    }
+
+    for (const prop in value) {
+        if (Object.prototype.hasOwnProperty.call(value, prop)) {
+            return false;
+        }
+    }
+
+    return true;
+}