fix: tweak the arg shapes to improve tool accuracy (#381)

Author: nirinchev

src/tools/mongodb/metadata/explain.ts

@@ -16,7 +16,7 @@ export class ExplainTool extends MongoDBToolBase {
         ...DbOperationArgs,
         method: z
             .array(
-                z.union([
+                z.discriminatedUnion("name", [
                     z.object({
                         name: z.literal("aggregate"),
                         arguments: z.object(AggregateArgs),

src/tools/mongodb/read/aggregate.ts

@@ -6,7 +6,7 @@ import { EJSON } from "bson";
 import { checkIndexUsage } from "../../../helpers/indexCheck.js";
 
 export const AggregateArgs = {
-    pipeline: z.array(z.record(z.string(), z.unknown())).describe("An array of aggregation stages to execute"),
+    pipeline: z.array(z.object({}).passthrough()).describe("An array of aggregation stages to execute"),
 };
 
 export class AggregateTool extends MongoDBToolBase {

src/tools/mongodb/read/count.ts

@@ -6,7 +6,8 @@ import { checkIndexUsage } from "../../../helpers/indexCheck.js";
 
 export const CountArgs = {
     query: z
-        .record(z.string(), z.unknown())
+        .object({})
+        .passthrough()
         .optional()
         .describe(
             "A filter/query parameter. Allows users to filter the documents to count. Matches the syntax of the filter argument of db.collection.count()."

src/tools/mongodb/read/find.ts

@@ -8,18 +8,23 @@ import { checkIndexUsage } from "../../../helpers/indexCheck.js";
 
 export const FindArgs = {
     filter: z
-        .record(z.string(), z.unknown())
+        .object({})
+        .passthrough()
         .optional()
         .describe("The query filter, matching the syntax of the query argument of db.collection.find()"),
     projection: z
-        .record(z.string(), z.unknown())
+        .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
-        .record(z.string(), z.custom<SortDirection>())
+        .object({})
+        .catchall(z.custom<SortDirection>())
         .optional()
-        .describe("A document, describing the sort order, matching the syntax of the sort argument of cursor.sort()"),
+        .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)."
+        ),
 };
 
 export class FindTool extends MongoDBToolBase {