Skip to content

Commit 05541e1

Browse files
fix: removes default limit from find tool schema
This commit removes default limit from the find tool schema because now we have a configurable max limit of the documents allowed to be sent per query.
1 parent a55795e commit 05541e1

File tree

4 files changed

+34
-33
lines changed

4 files changed

+34
-33
lines changed

src/common/config.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -182,8 +182,8 @@ export const defaultUserConfig: UserConfig = {
182182
idleTimeoutMs: 600000, // 10 minutes
183183
notificationTimeoutMs: 540000, // 9 minutes
184184
httpHeaders: {},
185-
maxDocumentsPerQuery: 50,
186-
maxBytesPerQuery: 1 * 1000 * 1000, // 1 mb
185+
maxDocumentsPerQuery: 10, // By default, we only fetch a maximum 10 documents per query / aggregation
186+
maxBytesPerQuery: 1 * 1000 * 1000, // By default, we only return ~1 mb of data per query / aggregation
187187
};
188188

189189
export const config = setupUserConfig({

src/tools/mongodb/read/export.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ export class ExportTool extends MongoDBToolBase {
2424
arguments: z
2525
.object({
2626
...FindArgs,
27-
limit: FindArgs.limit.removeDefault(),
27+
limit: FindArgs.limit,
2828
})
2929
.describe("The arguments for 'find' operation."),
3030
}),

src/tools/mongodb/read/find.ts

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ export const FindArgs = {
2121
.passthrough()
2222
.optional()
2323
.describe("The projection, matching the syntax of the projection argument of db.collection.find()"),
24-
limit: z.number().optional().default(10).describe("The maximum number of documents to return"),
24+
limit: z.number().optional().describe("The maximum number of documents to return"),
2525
sort: z
2626
.object({})
2727
.catchall(z.custom<SortDirection>())
@@ -61,18 +61,21 @@ export class FindTool extends MongoDBToolBase {
6161
});
6262
}
6363

64-
const appliedLimit = Math.min(limit, this.config.maxDocumentsPerQuery);
64+
const limitOnFindCursor = Math.min(limit ?? Number.POSITIVE_INFINITY, this.config.maxDocumentsPerQuery);
6565
findCursor = provider.find(database, collection, filter, {
6666
projection,
67-
limit: appliedLimit,
67+
limit: limitOnFindCursor,
6868
sort,
69-
batchSize: appliedLimit,
69+
batchSize: limitOnFindCursor,
7070
});
7171

7272
const [queryResultsCount, documents] = await Promise.all([
7373
operationWithFallback(
7474
() =>
7575
provider.countDocuments(database, collection, filter, {
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.
7679
limit,
7780
maxTimeMS: QUERY_COUNT_MAX_TIME_MS_CAP,
7881
}),

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

Lines changed: 24 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { beforeEach, describe, expect, it, vi } from "vitest";
1+
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
22
import type { Document, Collection } from "mongodb";
33
import {
44
getResponseContent,
@@ -228,37 +228,38 @@ describeWithMongoDB("find tool", (integration) => {
228228
describeWithMongoDB(
229229
"find tool with configured max documents per query",
230230
(integration) => {
231+
beforeEach(async () => {
232+
await freshInsertDocuments({
233+
collection: integration.mongoClient().db(integration.randomDbName()).collection("foo"),
234+
count: 1000,
235+
});
236+
});
237+
238+
afterEach(() => {
239+
vi.resetAllMocks();
240+
});
241+
231242
describe("when the provided limit is lower than the configured max limit", () => {
232243
it("should return documents limited to the provided limit", async () => {
233-
await freshInsertDocuments({
234-
collection: integration.mongoClient().db(integration.randomDbName()).collection("foo"),
235-
count: 1000,
236-
});
237244
await integration.connectMcpClient();
238245
const response = await integration.mcpClient().callTool({
239246
name: "find",
240247
arguments: {
241248
database: integration.randomDbName(),
242249
collection: "foo",
243250
filter: {},
244-
// default is 10
245-
limit: undefined,
251+
limit: 8,
246252
},
247253
});
248254

249255
const content = getResponseContent(response);
250-
expect(content).toContain(`Query on collection "foo" resulted in 10 documents.`);
251-
expect(content).toContain(`Returning 10 documents while respecting the applied limits.`);
256+
expect(content).toContain(`Query on collection "foo" resulted in 8 documents.`);
257+
expect(content).toContain(`Returning 8 documents while respecting the applied limits.`);
252258
});
253259
});
254260

255261
describe("when the provided limit is larger than the configured max limit", () => {
256262
it("should return documents limited to the configured max limit", async () => {
257-
await freshInsertDocuments({
258-
collection: integration.mongoClient().db(integration.randomDbName()).collection("foo"),
259-
count: 1000,
260-
});
261-
262263
await integration.connectMcpClient();
263264
const response = await integration.mcpClient().callTool({
264265
name: "find",
@@ -272,17 +273,13 @@ describeWithMongoDB(
272273

273274
const content = getResponseContent(response);
274275
expect(content).toContain(`Query on collection "foo" resulted in 1000 documents.`);
275-
expect(content).toContain(`Returning 20 documents while respecting the applied limits.`);
276+
expect(content).toContain(`Returning 10 documents while respecting the applied limits.`);
276277
});
277278
});
278279

279280
describe("when counting documents exceed the configured count maxTimeMS", () => {
280-
it("should abort discard count operation and respond with indeterminable count", async () => {
281+
it("should abort count operation and respond with indeterminable count", async () => {
281282
vi.spyOn(constants, "QUERY_COUNT_MAX_TIME_MS_CAP", "get").mockReturnValue(0.1);
282-
await freshInsertDocuments({
283-
collection: integration.mongoClient().db(integration.randomDbName()).collection("foo"),
284-
count: 1000,
285-
});
286283
await integration.connectMcpClient();
287284
const response = await integration.mcpClient().callTool({
288285
name: "find",
@@ -293,22 +290,23 @@ describeWithMongoDB(
293290

294291
const docs = getDocsFromUntrustedContent(content);
295292
expect(docs.length).toEqual(10);
296-
vi.resetAllMocks();
297293
});
298294
});
299295
},
300-
() => ({ ...defaultTestConfig, maxDocumentsPerQuery: 20 })
296+
() => ({ ...defaultTestConfig, maxDocumentsPerQuery: 10 })
301297
);
302298

303299
describeWithMongoDB(
304300
"find tool with configured max bytes per query",
305301
(integration) => {
302+
beforeEach(async () => {
303+
await freshInsertDocuments({
304+
collection: integration.mongoClient().db(integration.randomDbName()).collection("foo"),
305+
count: 1000,
306+
});
307+
});
306308
describe("when the provided maxBytesPerQuery is hit", () => {
307309
it("should return only the documents that could fit in maxBytesPerQuery limit", async () => {
308-
await freshInsertDocuments({
309-
collection: integration.mongoClient().db(integration.randomDbName()).collection("foo"),
310-
count: 1000,
311-
});
312310
await integration.connectMcpClient();
313311
const response = await integration.mcpClient().callTool({
314312
name: "find",

0 commit comments

Comments
 (0)