17863910702

Committed 19 Sep 2025 04:19PM UTC coverage: 82.391% (+0.5%) from 81.843%

Build # 17863910702

Build Type

Pull #536

github

Committed by

web-flow

Commit Message

Merge e1c95bda3 into c10955af6

Pull Request Pull Request #536: fix: add guards against possible memory overflow in find and aggregate tools MCP-21

Run Details

1083 of 1425 branches covered (76.0%)

Branch coverage included in aggregate %.

275 of 289 new or added lines in 7 files covered. (95.16%)

2 existing lines in 1 file now uncovered.

5238 of 6247 relevant lines covered (83.85%)

64.29 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

92.26

/src/tools/mongodb/read/find.ts

import { z } from "zod";
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
import type { ToolArgs, OperationType, ToolExecutionContext } from "../../tool.js";
import { formatUntrustedData } from "../../tool.js";
import type { FindCursor, SortDirection } from "mongodb";
import { checkIndexUsage } from "../../../helpers/indexCheck.js";
import { EJSON } from "bson";
import { collectCursorUntilMaxBytesLimit } from "../../../helpers/collectCursorUntilMaxBytes.js";
import { operationWithFallback } from "../../../helpers/operationWithFallback.js";
import { ONE_MB, QUERY_COUNT_MAX_TIME_MS_CAP, CURSOR_LIMITS_TO_LLM_TEXT } from "../../../helpers/constants.js";
import { zEJSON } from "../../args.js";
import { LogId } from "../../../common/logger.js";

export const FindArgs = {
    filter: zEJSON()
        .optional()
        .describe("The query filter, matching the syntax of the query argument of db.collection.find()"),
    projection: z
        .object({})
        .passthrough()
        .optional()
        .describe("The projection, matching the syntax of the projection argument of db.collection.find()"),
    limit: z.number().optional().default(10).describe("The maximum number of documents to return"),
    sort: z
        .object({})
        .catchall(z.custom<SortDirection>())
        .optional()
        .describe(
            "A document, describing the sort order, matching the syntax of the sort argument of cursor.sort(). The keys of the object are the fields to sort on, while the values are the sort directions (1 for ascending, -1 for descending)."
        ),
    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 query result is required, use the "export" tool instead of increasing this limit.\
`),
};

export class FindTool extends MongoDBToolBase {
    public name = "find";
    protected description = "Run a find query against a MongoDB collection";
    protected argsShape = {
        ...DbOperationArgs,
        ...FindArgs,
    };
    public operationType: OperationType = "read";

    protected async execute(
        { database, collection, filter, projection, limit, sort, responseBytesLimit }: ToolArgs<typeof this.argsShape>,
        { signal }: ToolExecutionContext
    ): Promise<CallToolResult> {
        let findCursor: FindCursor<unknown> | undefined = 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");
                });
            }

            const limitOnFindCursor = this.getLimitForFindCursor(limit);

            findCursor = provider.find(database, collection, filter, {
                projection,
                limit: limitOnFindCursor.limit,
                sort,
            });

            const [queryResultsCount, cursorResults] = await Promise.all([
                operationWithFallback(
                    () =>
                        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, only
                            // the limit provided to the tool.
                            limit,
                            maxTimeMS: QUERY_COUNT_MAX_TIME_MS_CAP,
                        }),
                    undefined
                ),
                collectCursorUntilMaxBytesLimit({
                    cursor: findCursor,
                    configuredMaxBytesPerQuery: this.config.maxBytesPerQuery,
                    toolResponseBytesLimit: responseBytesLimit,
                    abortSignal: signal,
                }),
            ]);

            return {
                content: formatUntrustedData(
                    this.generateMessage({
                        collection,
                        queryResultsCount,
                        documents: cursorResults.documents,
                        appliedLimits: [limitOnFindCursor.cappedBy, cursorResults.cappedBy].filter((limit) => !!limit),
                    }),
                    cursorResults.documents.length > 0 ? EJSON.stringify(cursorResults.documents) : undefined
                ),
            };
        } finally {
            if (findCursor) {
                void this.safeCloseCursor(findCursor);
            }
        }
    }

    private async safeCloseCursor(cursor: FindCursor<unknown>): Promise<void> {
        try {
            await cursor.close();
        } catch (error) {
            this.session.logger.warning({
                id: LogId.mongodbCursorCloseError,
                context: "find tool",
                message: `Error when closing the cursor - ${error instanceof Error ? error.message : String(error)}`,
            });
        }
    }

    private generateMessage({
        collection,
        queryResultsCount,
        documents,
        appliedLimits,
    }: {
        collection: string;
        queryResultsCount: number | undefined;
        documents: unknown[];
        appliedLimits: (keyof typeof CURSOR_LIMITS_TO_LLM_TEXT)[];
    }): string {
        const appliedLimitsText = 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 `\
Query on collection "${collection}" resulted in ${queryResultsCount === undefined ? "indeterminable number of" : queryResultsCount} documents. \
Returning ${documents.length} documents${appliedLimitsText ? ` ${appliedLimitsText}` : "."}\
`;
    }

    private getLimitForFindCursor(providedLimit: number | undefined | null): {
        cappedBy: "config.maxDocumentsPerQuery" | undefined;
        limit: number | undefined;
    } {
        const configuredLimit: number = parseInt(String(this.config.maxDocumentsPerQuery), 10);

        // Setting configured maxDocumentsPerQuery to negative, zero or nullish
        // is equivalent to disabling the max limit applied on documents
        const configuredLimitIsNotApplicable = Number.isNaN(configuredLimit) || configuredLimit <= 0;
        if (configuredLimitIsNotApplicable) {
            return { cappedBy: undefined, limit: providedLimit ?? undefined };
        }

        const providedLimitIsNotApplicable = providedLimit === null || providedLimit === undefined;
        if (providedLimitIsNotApplicable) {
            return { cappedBy: "config.maxDocumentsPerQuery", limit: configuredLimit };
        }

        return {
            cappedBy: configuredLimit < providedLimit ? "config.maxDocumentsPerQuery" : undefined,
            limit: Math.min(providedLimit, configuredLimit),
        };
    }
}

1	import { z } from "zod";	2✔
2	import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
3	import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";	2✔
4	import type { ToolArgs, OperationType, ToolExecutionContext } from "../../tool.js";
5	import { formatUntrustedData } from "../../tool.js";	2✔
6	import type { FindCursor, SortDirection } from "mongodb";
7	import { checkIndexUsage } from "../../../helpers/indexCheck.js";	2✔
8	import { EJSON } from "bson";	2✔
9	import { collectCursorUntilMaxBytesLimit } from "../../../helpers/collectCursorUntilMaxBytes.js";	2✔
10	import { operationWithFallback } from "../../../helpers/operationWithFallback.js";	2✔
11	import { ONE_MB, QUERY_COUNT_MAX_TIME_MS_CAP, CURSOR_LIMITS_TO_LLM_TEXT } from "../../../helpers/constants.js";	2✔
12	import { zEJSON } from "../../args.js";	2✔
13	import { LogId } from "../../../common/logger.js";	2✔
14
15	export const FindArgs = {	2✔
16	filter: zEJSON()	2✔
17	.optional()	2✔
18	.describe("The query filter, matching the syntax of the query argument of db.collection.find()"),	2✔
19	projection: z	2✔
20	.object({})	2✔
21	.passthrough()	2✔
22	.optional()	2✔
23	.describe("The projection, matching the syntax of the projection argument of db.collection.find()"),	2✔
24	limit: z.number().optional().default(10).describe("The maximum number of documents to return"),	2✔
25	sort: z	2✔
26	.object({})	2✔
27	.catchall(z.custom<SortDirection>())	2✔
28	.optional()	2✔
29	.describe(	2✔
30	"A document, describing the sort order, matching the syntax of the sort argument of cursor.sort(). The keys of the object are the fields to sort on, while the values are the sort directions (1 for ascending, -1 for descending)."	2✔
31	),	2✔
32	responseBytesLimit: z.number().optional().default(ONE_MB).describe(`\	2✔
33	The maximum number of bytes to return in the response. This value is capped by the server’s configured maxBytesPerQuery and cannot be exceeded. \
34	Note to LLM: If the entire query result is required, use the "export" tool instead of increasing this limit.\
35	`),	2✔
36	};	2✔
37
38	export class FindTool extends MongoDBToolBase {	2✔
39	public name = "find";	76✔
40	protected description = "Run a find query against a MongoDB collection";	76✔
41	protected argsShape = {	76✔
42	...DbOperationArgs,	76✔
43	...FindArgs,	76✔
44	};	76✔
45	public operationType: OperationType = "read";	76✔
46
47	protected async execute(	2✔
48	{ database, collection, filter, projection, limit, sort, responseBytesLimit }: ToolArgs<typeof this.argsShape>,	48✔
49	{ signal }: ToolExecutionContext	48✔
50	): Promise<CallToolResult> {	48✔
51	let findCursor: FindCursor<unknown> \| undefined = undefined;	48✔
52	try {	48✔
53	const provider = await this.ensureConnected();	48✔
54
55	// Check if find operation uses an index if enabled
56	if (this.config.indexCheck) {	45✔
57	await checkIndexUsage(provider, database, collection, "find", async () => {	5✔
58	return provider	5✔
59	.find(database, collection, filter, { projection, limit, sort })	5✔
60	.explain("queryPlanner");	5✔
61	});	5✔
62	}	2✔
63
64	const limitOnFindCursor = this.getLimitForFindCursor(limit);	40✔
65
66	findCursor = provider.find(database, collection, filter, {	40✔
67	projection,	40✔
68	limit: limitOnFindCursor.limit,	40✔
69	sort,	40✔
70	});	40✔
71
72	const [queryResultsCount, cursorResults] = await Promise.all([	40✔
73	operationWithFallback(	40✔
74	() =>	40✔
75	provider.countDocuments(database, collection, filter, {	40✔
76	// We should be counting documents that the original
77	// query would have yielded which is why we don't
78	// use `limitOnFindCursor` calculated above, only
79	// the limit provided to the tool.
80	limit,	40✔
81	maxTimeMS: QUERY_COUNT_MAX_TIME_MS_CAP,	40✔
82	}),	40✔
83	undefined	40✔
84	),	40✔
85	collectCursorUntilMaxBytesLimit({	40✔
86	cursor: findCursor,	40✔
87	configuredMaxBytesPerQuery: this.config.maxBytesPerQuery,	40✔
88	toolResponseBytesLimit: responseBytesLimit,	40✔
89	abortSignal: signal,	40✔
90	}),	40✔
91	]);	40✔
92
93	return {	40✔
94	content: formatUntrustedData(	40✔
95	this.generateMessage({	40✔
96	collection,	40✔
97	queryResultsCount,	40✔
98	documents: cursorResults.documents,	40✔
99	appliedLimits: [limitOnFindCursor.cappedBy, cursorResults.cappedBy].filter((limit) => !!limit),	40✔
100	}),	40✔
101	cursorResults.documents.length > 0 ? EJSON.stringify(cursorResults.documents) : undefined	48!
102	),	48✔
103	};	48✔
104	} finally {	48✔
105	if (findCursor) {	48!
106	void this.safeCloseCursor(findCursor);	40✔
107	}	40✔
108	}	48✔
109	}	48✔
110
111	private async safeCloseCursor(cursor: FindCursor<unknown>): Promise<void> {	2✔
112	try {	40✔
113	await cursor.close();	40✔
114	} catch (error) {	40!
NEW 115	this.session.logger.warning({	×
NEW 116	id: LogId.mongodbCursorCloseError,	×
NEW 117	context: "find tool",	×
NEW 118	message: `Error when closing the cursor - ${error instanceof Error ? error.message : String(error)}`,	×
UNCOV 119	});	×
UNCOV 120	}	×
121	}	40✔
122
123	private generateMessage({	2✔
124	collection,	40✔
125	queryResultsCount,	40✔
126	documents,	40✔
127	appliedLimits,	40✔
128	}: {	40✔
129	collection: string;
130	queryResultsCount: number \| undefined;
131	documents: unknown[];
132	appliedLimits: (keyof typeof CURSOR_LIMITS_TO_LLM_TEXT)[];
133	}): string {	40✔
134	const appliedLimitsText = appliedLimits.length	40✔
135	? `\	8✔
136	while respecting the applied limits of ${appliedLimits.map((limit) => CURSOR_LIMITS_TO_LLM_TEXT[limit]).join(", ")}. \	8✔
137	Note to LLM: If the entire query result is required then use "export" tool to export the query results.\
138	`
139	: "";	32✔
140
141	return `\	40✔
142	Query on collection "${collection}" resulted in ${queryResultsCount === undefined ? "indeterminable number of" : queryResultsCount} documents. \	40✔
143	Returning ${documents.length} documents${appliedLimitsText ? ` ${appliedLimitsText}` : "."}\	40✔
144	`;
145	}	40✔
146
147	private getLimitForFindCursor(providedLimit: number \| undefined \| null): {	2✔
148	cappedBy: "config.maxDocumentsPerQuery" \| undefined;
149	limit: number \| undefined;
150	} {	40✔
151	const configuredLimit: number = parseInt(String(this.config.maxDocumentsPerQuery), 10);	40✔
152
153	// Setting configured maxDocumentsPerQuery to negative, zero or nullish
154	// is equivalent to disabling the max limit applied on documents
155	const configuredLimitIsNotApplicable = Number.isNaN(configuredLimit) \|\| configuredLimit <= 0;	40✔
156	if (configuredLimitIsNotApplicable) {	40✔
157	return { cappedBy: undefined, limit: providedLimit ?? undefined };	4!
158	}	4✔
159
160	const providedLimitIsNotApplicable = providedLimit === null \|\| providedLimit === undefined;	36✔
161	if (providedLimitIsNotApplicable) {	40!
NEW 162	return { cappedBy: "config.maxDocumentsPerQuery", limit: configuredLimit };	×
NEW 163	}	✔
164
165	return {	36✔
166	cappedBy: configuredLimit < providedLimit ? "config.maxDocumentsPerQuery" : undefined,	40✔
167	limit: Math.min(providedLimit, configuredLimit),	40✔
168	};	40✔
169	}	40✔
170	}	2✔

mongodb-js / mongodb-mcp-server / 17863910702

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous