simplify tool types

Author: nirinchev

src/tools/atlas/atlasTool.ts

@@ -1,9 +1,8 @@
-import { ZodRawShape } from "zod";
 import { ToolBase } from "../tool.js";
 import { ApiClient } from "../../client.js";
 import { State } from "../../state.js";
 
-export abstract class AtlasToolBase<Args extends ZodRawShape = ZodRawShape> extends ToolBase<Args> {
+export abstract class AtlasToolBase extends ToolBase {
     constructor(
         state: State,
         protected apiClient: ApiClient

src/tools/atlas/listClusters.ts

@@ -7,9 +7,7 @@ import { AtlasToolBase } from "./atlasTool.js";
 import { State } from "../../state.js";
 import { ToolArgs } from "../tool.js";
 
-export class ListClustersTool extends AtlasToolBase<{
-    projectId: ZodString | ZodOptional<ZodString>;
-}> {
+export class ListClustersTool extends AtlasToolBase {
     protected name = "listClusters";
     protected description = "List MongoDB Atlas clusters";
     protected argsShape;

src/tools/atlas/tools.ts

@@ -1,4 +1,3 @@
-import { ZodRawShape } from "zod";
 import { ToolBase } from "../tool.js";
 import { ApiClient } from "../../client.js";
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
@@ -8,7 +7,7 @@ import { ListClustersTool } from "./listClusters.js";
 import { ListProjectsTool } from "./listProjects.js";
 
 export function registerAtlasTools(server: McpServer, state: State, apiClient: ApiClient) {
-    const tools: ToolBase<ZodRawShape>[] = [
+    const tools: ToolBase[] = [
         new AuthTool(state, apiClient),
         new ListClustersTool(state, apiClient),
         new ListProjectsTool(state, apiClient),

src/tools/mongodb/collectionIndexes.ts

@@ -2,7 +2,7 @@ import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { DbOperationArgs, MongoDBToolBase } from "./mongodbTool.js";
 import { ToolArgs } from "../tool.js";
 
-export class CollectionIndexesTool extends MongoDBToolBase<typeof DbOperationArgs> {
+export class CollectionIndexesTool extends MongoDBToolBase {
     protected name = "collection-indexes";
     protected description = "Describe the indexes for a collection";
     protected argsShape = DbOperationArgs;

src/tools/mongodb/collectionSchema.ts

@@ -3,7 +3,7 @@ import { DbOperationArgs, MongoDBToolBase } from "./mongodbTool.js";
 import { ToolArgs } from "../tool.js";
 import { parseSchema, SchemaField } from "mongodb-schema";
 
-export class CollectionSchemaTool extends MongoDBToolBase<typeof DbOperationArgs> {
+export class CollectionSchemaTool extends MongoDBToolBase {
     protected name = "collection-schema";
     protected description = "Describe the schema for a collection";
     protected argsShape = DbOperationArgs;

src/tools/mongodb/connect.ts

@@ -5,19 +5,19 @@ import { MongoDBToolBase } from "./mongodbTool.js";
 import { ToolArgs } from "../tool";
 import { ErrorCodes } from "../../errors.js";
 
-const argsShape = {
-    connectionStringOrClusterName: z
-        .string()
-        .optional()
-        .describe("MongoDB connection string (in the mongodb:// or mongodb+srv:// format) or cluster name"),
-};
-
-export class ConnectTool extends MongoDBToolBase<typeof argsShape> {
+export class ConnectTool extends MongoDBToolBase {
     protected name = "connect";
     protected description = "Connect to a MongoDB instance";
-    protected argsShape = argsShape;
+    protected argsShape = {
+        connectionStringOrClusterName: z
+            .string()
+            .optional()
+            .describe("MongoDB connection string (in the mongodb:// or mongodb+srv:// format) or cluster name"),
+    };
 
-    protected async execute({ connectionStringOrClusterName }: ToolArgs<typeof argsShape>): Promise<CallToolResult> {
+    protected async execute({
+        connectionStringOrClusterName,
+    }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
         if (!connectionStringOrClusterName) {
             // TODO: try reconnecting to the default connection
             return {

src/tools/mongodb/createIndex.ts

@@ -4,17 +4,15 @@ import { DbOperationArgs, MongoDBToolBase } from "./mongodbTool.js";
 import { ToolArgs } from "../tool.js";
 import { IndexDirection } from "mongodb";
 
-const argsShape = {
-    ...DbOperationArgs,
-    keys: z.record(z.string(), z.custom<IndexDirection>()).describe("The index definition"),
-};
-
-export class CreateIndexTool extends MongoDBToolBase<typeof argsShape> {
+export class CreateIndexTool extends MongoDBToolBase {
     protected name = "create-index";
     protected description = "Create an index for a collection";
-    protected argsShape = argsShape;
+    protected argsShape = {
+        ...DbOperationArgs,
+        keys: z.record(z.string(), z.custom<IndexDirection>()).describe("The index definition"),
+    };
 
-    protected async execute({ database, collection, keys }: ToolArgs<typeof argsShape>): Promise<CallToolResult> {
+    protected async execute({ database, collection, keys }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
         const provider = this.ensureConnected();
         const indexes = await provider.createIndexes(database, collection, [
             {

src/tools/mongodb/find/find.ts

@@ -3,34 +3,32 @@ import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { MongoDBToolBase } from "../mongodbTool.js";
 import { ToolArgs } from "../../tool.js";
 
-const argsShape = {
-    collection: z.string().describe("Collection name"),
-    database: z.string().describe("Database name"),
-    filter: z
-        .object({})
-        .passthrough()
-        .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"),
-};
-
-export class FindTool extends MongoDBToolBase<typeof argsShape> {
+export class FindTool extends MongoDBToolBase {
     protected name = "find";
     protected description = "Run a find query against a MongoDB collection";
-    protected argsShape = argsShape;
+    protected argsShape = {
+        collection: z.string().describe("Collection name"),
+        database: z.string().describe("Database name"),
+        filter: z
+            .object({})
+            .passthrough()
+            .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"),
+    };
 
     protected async execute({
         database,
         collection,
         filter,
         projection,
         limit,
-    }: ToolArgs<typeof argsShape>): Promise<CallToolResult> {
+    }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
         const provider = this.ensureConnected();
         const documents = await provider.find(database, collection, filter, { projection, limit }).toArray();
 

src/tools/mongodb/insert/insertOne.ts

@@ -3,21 +3,25 @@ import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { MongoDBToolBase } from "../mongodbTool.js";
 import { ToolArgs } from "../../tool.js";
 
-const argsShape = {
-    collection: z.string().describe("Collection name"),
-    database: z.string().describe("Database name"),
-    document: z
-        .object({})
-        .passthrough()
-        .describe("The document to insert, matching the syntax of the document argument of db.collection.insertOne()"),
-};
-
-export class InsertOneTool extends MongoDBToolBase<typeof argsShape> {
+export class InsertOneTool extends MongoDBToolBase {
     protected name = "insert-one";
     protected description = "Insert a document into a MongoDB collection";
-    protected argsShape = argsShape;
+    protected argsShape = {
+        collection: z.string().describe("Collection name"),
+        database: z.string().describe("Database name"),
+        document: z
+            .object({})
+            .passthrough()
+            .describe(
+                "The document to insert, matching the syntax of the document argument of db.collection.insertOne()"
+            ),
+    };
 
-    protected async execute({ database, collection, document }: ToolArgs<typeof argsShape>): Promise<CallToolResult> {
+    protected async execute({
+        database,
+        collection,
+        document,
+    }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
         const provider = this.ensureConnected();
         const result = await provider.insertOne(database, collection, document);
 

src/tools/mongodb/listCollections.ts

@@ -2,16 +2,14 @@ import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { DbOperationArgs, MongoDBToolBase } from "./mongodbTool.js";
 import { ToolArgs } from "../tool.js";
 
-const argsShape = {
-    database: DbOperationArgs.database,
-};
-
-export class ListCollectionsTool extends MongoDBToolBase<typeof argsShape> {
+export class ListCollectionsTool extends MongoDBToolBase {
     protected name = "list-collections";
     protected description = "List all collections for a given database";
-    protected argsShape = argsShape;
+    protected argsShape = {
+        database: DbOperationArgs.database,
+    };
 
-    protected async execute({ database }: ToolArgs<typeof argsShape>): Promise<CallToolResult> {
+    protected async execute({ database }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
         const provider = this.ensureConnected();
         const collections = await provider.listCollections(database);