fix(schema): Use sample instead of find for schema sampling

src/tools/mongodb/metadata/collectionSchema.ts

@@ -1,19 +1,35 @@
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
-import type { ToolArgs, OperationType } from "../../tool.js";
+import type { ToolArgs, OperationType, ToolExecutionContext } from "../../tool.js";
 import { formatUntrustedData } from "../../tool.js";
 import { getSimplifiedSchema } from "mongodb-schema";
+import z from "zod";
+import { ONE_MB } from "../../../helpers/constants.js";
+import { collectCursorUntilMaxBytesLimit } from "../../../helpers/collectCursorUntilMaxBytes.js";
 
 export class CollectionSchemaTool extends MongoDBToolBase {
     public name = "collection-schema";
     protected description = "Describe the schema for a collection";
-    protected argsShape = DbOperationArgs;
+    protected argsShape = {
+        ...DbOperationArgs,
+        sampleSize: z.number().optional().default(50).describe("Number of documents to sample for schema inference"),
+        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.`),
+    };
 
     public operationType: OperationType = "metadata";
 
-    protected async execute({ database, collection }: ToolArgs<typeof DbOperationArgs>): Promise<CallToolResult> {
+    protected async execute(
+        { database, collection, sampleSize, responseBytesLimit }: ToolArgs<typeof this.argsShape>,
+        { signal }: ToolExecutionContext
+    ): Promise<CallToolResult> {
         const provider = await this.ensureConnected();
-        const documents = await provider.find(database, collection, {}, { limit: 5 }).toArray();
+        const cursor = provider.aggregate(database, collection, [{ $sample: { size: Math.min(sampleSize, this.config.maxDocumentsPerQuery) } }]);
+        const { cappedBy, documents } = await collectCursorUntilMaxBytesLimit({
+            cursor,
+            configuredMaxBytesPerQuery: this.config.maxBytesPerQuery,
+            toolResponseBytesLimit: responseBytesLimit,
+            abortSignal: signal,
+        });
         const schema = await getSimplifiedSchema(documents);
 
         const fieldsCount = Object.entries(schema).length;
@@ -28,9 +44,12 @@
             };
         }
 
+        const header = `Found ${fieldsCount} fields in the schema for "${database}.${collection}"`;
+        const cappedWarning = cappedBy !== undefined ? `\nThe schema was inferred from a subset of documents due to the response size limit. (${cappedBy})` : "";
+
         return {
             content: formatUntrustedData(
-                `Found ${fieldsCount} fields in the schema for "${database}.${collection}"`,
+                `${header}${cappedWarning}`,
                 JSON.stringify(schema)
             ),
         };

tests/integration/tools/mongodb/metadata/collectionSchema.test.ts

@@ -15,11 +15,25 @@
 import { describe, expect, it } from "vitest";
 
 describeWithMongoDB("collectionSchema tool", (integration) => {
     validateToolMetadata(
         integration,
         "collection-schema",
         "Describe the schema for a collection",
-        databaseCollectionParameters
+        [
+            ...databaseCollectionParameters,
+            {
+                name: "sampleSize",
+                type: "number",
+                description: "Number of documents to sample for schema inference",
+                required: false, 
+            },
+            {
+                name: "responseBytesLimit",
+                type: "number",
+                description: `The maximum number of bytes to return in the response. This value is capped by the server’s configured maxBytesPerQuery and cannot be exceeded.`,
+                required: false, 
+            }
+        ]
     );
 
     validateThrowsForInvalidArguments(integration, "collection-schema", databaseCollectionInvalidArgs);