diff --git a/src/tools/mongodb/read/count.ts b/src/tools/mongodb/read/count.ts index bd86169b..19782ce1 100644 --- a/src/tools/mongodb/read/count.ts +++ b/src/tools/mongodb/read/count.ts @@ -4,17 +4,21 @@ import { ToolArgs, OperationType } from "../../tool.js"; import { z } from "zod"; export const CountArgs = { - query: z + filter: z .record(z.string(), z.unknown()) .optional() .describe( - "The query filter to count documents. Matches the syntax of the filter argument of db.collection.count()" + "The query filter to count documents. Matches the syntax of the filter argument of db.collection.countDocuments()" ), + query: z + .record(z.string(), z.unknown()) + .optional() + .describe("Alternative old name for filter. Will be used in db.collection.countDocuments()"), }; export class CountTool extends MongoDBToolBase { protected name = "count"; - protected description = "Gets the number of documents in a MongoDB collection"; + protected description = "Gets the number of documents in a MongoDB collection using countDocuments()"; protected argsShape = { ...DbOperationArgs, ...CountArgs, @@ -22,9 +26,16 @@ export class CountTool extends MongoDBToolBase { protected operationType: OperationType = "read"; - protected async execute({ database, collection, query }: ToolArgs): Promise { + protected async execute({ + database, + collection, + query, + filter, + }: ToolArgs): Promise { const provider = await this.ensureConnected(); - const count = await provider.count(database, collection, query); + // use either filter or query, since we're using countDocuments, prefer filter + const queryFilter = filter || query; + const count = await provider.countDocuments(database, collection, queryFilter); return { content: [ diff --git a/tests/integration/tools/mongodb/read/count.test.ts b/tests/integration/tools/mongodb/read/count.test.ts index 5b288448..1f9e942a 100644 --- a/tests/integration/tools/mongodb/read/count.test.ts +++ b/tests/integration/tools/mongodb/read/count.test.ts @@ -8,16 +8,27 @@ import { } from "../../../helpers.js"; describeWithMongoDB("count tool", (integration) => { - validateToolMetadata(integration, "count", "Gets the number of documents in a MongoDB collection", [ - { - name: "query", - description: - "The query filter to count documents. Matches the syntax of the filter argument of db.collection.count()", - type: "object", - required: false, - }, - ...databaseCollectionParameters, - ]); + validateToolMetadata( + integration, + "count", + "Gets the number of documents in a MongoDB collection using countDocuments()", + [ + { + name: "filter", + description: + "The query filter to count documents. Matches the syntax of the filter argument of db.collection.countDocuments()", + type: "object", + required: false, + }, + { + name: "query", + description: "Alternative old name for filter. Will be used in db.collection.countDocuments()", + type: "object", + required: false, + }, + ...databaseCollectionParameters, + ] + ); validateThrowsForInvalidArguments(integration, "count", [ {}, @@ -79,6 +90,43 @@ describeWithMongoDB("count tool", (integration) => { expect(content).toEqual(`Found ${testCase.expectedCount} documents in the collection "foo"`); }); } + + it("correctly filters documents when using 'filter' parameter", async () => { + await integration.connectMcpClient(); + + // Using 'filter' parameter - should work correctly after the fix + const response = await integration.mcpClient().callTool({ + name: "count", + arguments: { + database: integration.randomDbName(), + collection: "foo", + filter: { age: { $lt: 15 } }, + }, + }); + + const content = getResponseContent(response.content); + expect(content).toEqual('Found 2 documents in the collection "foo"'); + }); + + it("prioritizes filter over query when both are provided", async () => { + await integration.connectMcpClient(); + + // Using both 'filter' and 'query' parameters + const response = await integration.mcpClient().callTool({ + name: "count", + arguments: { + database: integration.randomDbName(), + collection: "foo", + filter: { age: { $lt: 15 } }, + query: { age: { $gt: 10 } }, + }, + }); + + const content = getResponseContent(response.content); + // Filter takes precedence over query + // Filter is { age: { $lt: 15 } } which matches 2 documents + expect(content).toEqual('Found 2 documents in the collection "foo"'); + }); }); validateAutoConnectBehavior(integration, "count", () => {