From 961ddc0920f516762292fb5a93245d931c496938 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Mon, 24 Nov 2025 16:50:58 +0100 Subject: [PATCH 1/3] chore: split connect and switch-connection tool --- src/common/logger.ts | 1 + src/tools/mongodb/connect/connect.ts | 104 ++++-------------- src/tools/mongodb/tools.ts | 2 + src/tools/tool.ts | 76 ++++++------- .../tools/mongodb/connect/connect.test.ts | 3 +- 5 files changed, 57 insertions(+), 129 deletions(-) diff --git a/src/common/logger.ts b/src/common/logger.ts index ec2343e3b..ad481a141 100644 --- a/src/common/logger.ts +++ b/src/common/logger.ts @@ -40,6 +40,7 @@ export const LogId = { toolExecute: mongoLogId(1_003_001), toolExecuteFailure: mongoLogId(1_003_002), toolDisabled: mongoLogId(1_003_003), + toolMetadataChange: mongoLogId(1_003_004), mongodbConnectFailure: mongoLogId(1_004_001), mongodbDisconnectFailure: mongoLogId(1_004_002), diff --git a/src/tools/mongodb/connect/connect.ts b/src/tools/mongodb/connect/connect.ts index 87219d0db..fd5dc0caf 100644 --- a/src/tools/mongodb/connect/connect.ts +++ b/src/tools/mongodb/connect/connect.ts @@ -2,114 +2,48 @@ import { z } from "zod"; import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { MongoDBToolBase } from "../mongodbTool.js"; import type { ToolArgs, OperationType, ToolConstructorParams } from "../../tool.js"; -import assert from "assert"; import type { Server } from "../../../server.js"; -import { LogId } from "../../../common/logger.js"; - -const disconnectedSchema = z - .object({ - connectionString: z.string().describe("MongoDB connection string (in the mongodb:// or mongodb+srv:// format)"), - }) - .describe("Options for connecting to MongoDB."); - -const connectedSchema = z - .object({ - connectionString: z - .string() - .optional() - .describe("MongoDB connection string to switch to (in the mongodb:// or mongodb+srv:// format)"), - }) - .describe( - "Options for switching the current MongoDB connection. If a connection string is not provided, the connection string from the config will be used." - ); - -const connectedName = "switch-connection" as const; -const disconnectedName = "connect" as const; - -const connectedDescription = - "Switch to a different MongoDB connection. If the user has configured a connection string or has previously called the connect tool, a connection is already established and there's no need to call this tool unless the user has explicitly requested to switch to a new instance."; -const disconnectedDescription = - "Connect to a MongoDB instance. The config resource captures if the server is already connected to a MongoDB cluster. If the user has configured a connection string or has previously called the connect tool, a connection is already established and there's no need to call this tool unless the user has explicitly requested to switch to a new MongoDB cluster."; - export class ConnectTool extends MongoDBToolBase { - public name: typeof connectedName | typeof disconnectedName = disconnectedName; - protected description: typeof connectedDescription | typeof disconnectedDescription = disconnectedDescription; + public override name = "connect"; + protected override description = + "Connect to a MongoDB instance. The config resource captures if the server is already connected to a MongoDB cluster. If the user has configured a connection string or has previously called the connect tool, a connection is already established and there's no need to call this tool unless the user has explicitly requested to switch to a new MongoDB cluster."; // Here the default is empty just to trigger registration, but we're going to override it with the correct // schema in the register method. - protected argsShape = { - connectionString: z.string().optional(), + protected override argsShape = { + connectionString: z.string().describe("MongoDB connection string (in the mongodb:// or mongodb+srv:// format)"), }; - public operationType: OperationType = "connect"; + public override operationType: OperationType = "connect"; constructor({ session, config, telemetry, elicitation }: ToolConstructorParams) { super({ session, config, telemetry, elicitation }); session.on("connect", () => { - this.updateMetadata(); + this.disable(); }); session.on("disconnect", () => { - this.updateMetadata(); + this.enable(); }); } - protected async execute({ connectionString }: ToolArgs): Promise { - switch (this.name) { - case disconnectedName: - assert(connectionString, "Connection string is required"); - break; - case connectedName: - connectionString ??= this.config.connectionString; - assert( - connectionString, - "Cannot switch to a new connection because no connection string was provided and no default connection string is configured." - ); - break; + public override register(server: Server): boolean { + const registrationSuccessful = super.register(server); + /** + * When connected to mongodb we want to swap connect with + * switch-connection tool. + */ + if (registrationSuccessful && this.session.isConnectedToMongoDB) { + this.disable(); } + return registrationSuccessful; + } + protected override async execute({ connectionString }: ToolArgs): Promise { await this.session.connectToMongoDB({ connectionString }); - this.updateMetadata(); return { content: [{ type: "text", text: "Successfully connected to MongoDB." }], }; } - - public register(server: Server): boolean { - if (super.register(server)) { - this.updateMetadata(); - return true; - } - - return false; - } - - private updateMetadata(): void { - let name: string; - let description: string; - let inputSchema: z.ZodObject; - - if (this.session.isConnectedToMongoDB) { - name = connectedName; - description = connectedDescription; - inputSchema = connectedSchema; - } else { - name = disconnectedName; - description = disconnectedDescription; - inputSchema = disconnectedSchema; - } - - this.session.logger.info({ - id: LogId.updateToolMetadata, - context: "tool", - message: `Updating tool metadata to ${name}`, - }); - - this.update?.({ - name, - description, - inputSchema, - }); - } } diff --git a/src/tools/mongodb/tools.ts b/src/tools/mongodb/tools.ts index c4498c805..ffbb71d80 100644 --- a/src/tools/mongodb/tools.ts +++ b/src/tools/mongodb/tools.ts @@ -20,9 +20,11 @@ import { CreateCollectionTool } from "./create/createCollection.js"; import { LogsTool } from "./metadata/logs.js"; import { ExportTool } from "./read/export.js"; import { DropIndexTool } from "./delete/dropIndex.js"; +import { SwitchConnectionTool } from "./connect/switchConnection.js"; export const MongoDbTools = [ ConnectTool, + SwitchConnectionTool, ListCollectionsTool, ListDatabasesTool, CollectionIndexesTool, diff --git a/src/tools/tool.ts b/src/tools/tool.ts index 6917020eb..f24afa555 100644 --- a/src/tools/tool.ts +++ b/src/tools/tool.ts @@ -1,4 +1,4 @@ -import type { z, AnyZodObject } from "zod"; +import type { z } from "zod"; import { type ZodRawShape, type ZodNever } from "zod"; import type { RegisteredTool, ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js"; import type { CallToolResult, ToolAnnotations } from "@modelcontextprotocol/sdk/types.js"; @@ -57,6 +57,8 @@ export abstract class ToolBase { protected abstract argsShape: ZodRawShape; + private registeredTool: RegisteredTool | undefined; + protected get annotations(): ToolAnnotations { const annotations: ToolAnnotations = { title: this.name, @@ -168,52 +170,40 @@ export abstract class ToolBase { } }; - server.mcpServer.tool(this.name, this.description, this.argsShape, this.annotations, callback); - - // This is very similar to RegisteredTool.update, but without the bugs around the name. - // In the upstream update method, the name is captured in the closure and not updated when - // the tool name changes. This means that you only get one name update before things end up - // in a broken state. - // See https://github.com/modelcontextprotocol/typescript-sdk/issues/414 for more details. - this.update = (updates: { name?: string; description?: string; inputSchema?: AnyZodObject }): void => { - const tools = server.mcpServer["_registeredTools"] as { [toolName: string]: RegisteredTool }; - const existingTool = tools[this.name]; - - if (!existingTool) { - this.session.logger.warning({ - id: LogId.toolUpdateFailure, - context: "tool", - message: `Tool ${this.name} not found in update`, - noRedaction: true, - }); - return; - } - - existingTool.annotations = this.annotations; - - if (updates.name && updates.name !== this.name) { - existingTool.annotations.title = updates.name; - delete tools[this.name]; - this.name = updates.name; - tools[this.name] = existingTool; - } - - if (updates.description) { - existingTool.description = updates.description; - this.description = updates.description; - } - - if (updates.inputSchema) { - existingTool.inputSchema = updates.inputSchema; - } - - server.mcpServer.sendToolListChanged(); - }; + this.registeredTool = server.mcpServer.tool( + this.name, + this.description, + this.argsShape, + this.annotations, + callback + ); return true; } - protected update?: (updates: { name?: string; description?: string; inputSchema?: AnyZodObject }) => void; + protected disable(): void { + if (!this.registeredTool) { + this.session.logger.warning({ + id: LogId.toolMetadataChange, + context: `tool - ${this.name}`, + message: "Requested disabling of tool but it was never registered", + }); + return; + } + this.registeredTool.disable(); + } + + protected enable(): void { + if (!this.registeredTool) { + this.session.logger.warning({ + id: LogId.toolMetadataChange, + context: `tool - ${this.name}`, + message: "Requested enabling of tool but it was never registered", + }); + return; + } + this.registeredTool.enable(); + } // Checks if a tool is allowed to run based on the config protected verifyAllowed(): boolean { diff --git a/tests/integration/tools/mongodb/connect/connect.test.ts b/tests/integration/tools/mongodb/connect/connect.test.ts index c197c63a9..e2b1f06a8 100644 --- a/tests/integration/tools/mongodb/connect/connect.test.ts +++ b/tests/integration/tools/mongodb/connect/connect.test.ts @@ -26,7 +26,8 @@ describeWithMongoDB( [ { name: "connectionString", - description: "MongoDB connection string to switch to (in the mongodb:// or mongodb+srv:// format)", + description: + "MongoDB connection string to switch to (in the mongodb:// or mongodb+srv:// format). If a connection string is not provided, the connection string from the config will be used.", type: "string", required: false, }, From e75d5538f91b1ec49847277784d0ca49be39be22 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Mon, 24 Nov 2025 16:55:37 +0100 Subject: [PATCH 2/3] chore: add missing tool file --- src/tools/mongodb/connect/switchConnection.ts | 58 +++++++++++++++++++ 1 file changed, 58 insertions(+) create mode 100644 src/tools/mongodb/connect/switchConnection.ts diff --git a/src/tools/mongodb/connect/switchConnection.ts b/src/tools/mongodb/connect/switchConnection.ts new file mode 100644 index 000000000..53ce500c2 --- /dev/null +++ b/src/tools/mongodb/connect/switchConnection.ts @@ -0,0 +1,58 @@ +import z from "zod"; +import { type CallToolResult } from "@modelcontextprotocol/sdk/types.js"; + +import { MongoDBToolBase } from "../mongodbTool.js"; +import { type ToolArgs, type OperationType, type ToolConstructorParams } from "../../tool.js"; +import type { Server } from "../../../server.js"; + +export class SwitchConnectionTool extends MongoDBToolBase { + public override name = "switch-connection"; + protected override description = + "Switch to a different MongoDB connection. If the user has configured a connection string or has previously called the connect tool, a connection is already established and there's no need to call this tool unless the user has explicitly requested to switch to a new instance."; + + protected override argsShape = { + connectionString: z + .string() + .optional() + .describe( + "MongoDB connection string to switch to (in the mongodb:// or mongodb+srv:// format). If a connection string is not provided, the connection string from the config will be used." + ), + }; + + public override operationType: OperationType = "connect"; + + constructor({ session, config, telemetry, elicitation }: ToolConstructorParams) { + super({ session, config, telemetry, elicitation }); + session.on("connect", () => { + this.enable(); + }); + + session.on("disconnect", () => { + this.disable(); + }); + } + + public override register(server: Server): boolean { + const registrationSuccessful = super.register(server); + /** + * When connected to mongodb we want to swap connect with + * switch-connection tool. + */ + if (registrationSuccessful && !this.session.isConnectedToMongoDB) { + this.disable(); + } + return registrationSuccessful; + } + + protected override async execute({ connectionString }: ToolArgs): Promise { + if (typeof connectionString !== "string") { + await this.session.connectToConfiguredConnection(); + } else { + await this.session.connectToMongoDB({ connectionString }); + } + + return { + content: [{ type: "text", text: "Successfully connected to MongoDB." }], + }; + } +} From aa9c57febae642993558968b06aca73c1ffb97c4 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Mon, 24 Nov 2025 17:12:17 +0100 Subject: [PATCH 3/3] chore: suggest only enabled tools --- src/common/connectionErrorHandler.ts | 2 +- src/tools/tool.ts | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/src/common/connectionErrorHandler.ts b/src/common/connectionErrorHandler.ts index 30b637963..41cc89639 100644 --- a/src/common/connectionErrorHandler.ts +++ b/src/common/connectionErrorHandler.ts @@ -14,7 +14,7 @@ export type ConnectionErrorHandled = { errorHandled: true; result: CallToolResul export const connectionErrorHandler: ConnectionErrorHandler = (error, { availableTools, connectionState }) => { const connectTools = availableTools - .filter((t) => t.operationType === "connect") + .filter((t) => t.operationType === "connect" && t.isEnabled()) .sort((a, b) => a.category.localeCompare(b.category)); // Sort Atlas tools before MongoDB tools // Find what Atlas connect tools are available and suggest when the LLM should to use each. If no Atlas tools are found, return a suggestion for the MongoDB connect tool. diff --git a/src/tools/tool.ts b/src/tools/tool.ts index f24afa555..8173b8d20 100644 --- a/src/tools/tool.ts +++ b/src/tools/tool.ts @@ -181,6 +181,10 @@ export abstract class ToolBase { return true; } + public isEnabled(): boolean { + return this.registeredTool?.enabled ?? false; + } + protected disable(): void { if (!this.registeredTool) { this.session.logger.warning({