diff --git a/src/errors.ts b/src/errors.ts index 224610fb..ae91c3a0 100644 --- a/src/errors.ts +++ b/src/errors.ts @@ -1,6 +1,6 @@ export enum ErrorCodes { NotConnectedToMongoDB = 1_000_000, - InvalidParams = 1_000_001, + MisconfiguredConnectionString = 1_000_001, } export class MongoDBError extends Error { diff --git a/src/index.ts b/src/index.ts index 268c4803..ef496e5d 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,8 +1,7 @@ #!/usr/bin/env node import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"; -import logger from "./logger.js"; -import { mongoLogId } from "mongodb-log-writer"; +import logger, { LogId } from "./logger.js"; import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import { config } from "./config.js"; import { Session } from "./session.js"; @@ -29,6 +28,6 @@ try { await server.connect(transport); } catch (error: unknown) { - logger.emergency(mongoLogId(1_000_004), "server", `Fatal error running server: ${error as string}`); + logger.emergency(LogId.serverStartFailure, "server", `Fatal error running server: ${error as string}`); process.exit(1); } diff --git a/src/logger.ts b/src/logger.ts index 0ff292cc..cb127568 100644 --- a/src/logger.ts +++ b/src/logger.ts @@ -1,11 +1,29 @@ import fs from "fs/promises"; -import { MongoLogId, MongoLogManager, MongoLogWriter } from "mongodb-log-writer"; +import { mongoLogId, MongoLogId, MongoLogManager, MongoLogWriter } from "mongodb-log-writer"; import redact from "mongodb-redact"; import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import { LoggingMessageNotification } from "@modelcontextprotocol/sdk/types.js"; export type LogLevel = LoggingMessageNotification["params"]["level"]; +export const LogId = { + serverStartFailure: mongoLogId(1_000_001), + serverInitialized: mongoLogId(1_000_002), + + atlasCheckCredentials: mongoLogId(1_001_001), + + telemetryDisabled: mongoLogId(1_002_001), + telemetryEmitFailure: mongoLogId(1_002_002), + telemetryEmitStart: mongoLogId(1_002_003), + telemetryEmitSuccess: mongoLogId(1_002_004), + + toolExecute: mongoLogId(1_003_001), + toolExecuteFailure: mongoLogId(1_003_002), + toolDisabled: mongoLogId(1_003_003), + + mongodbConnectFailure: mongoLogId(1_004_001), +} as const; + abstract class LoggerBase { abstract log(level: LogLevel, id: MongoLogId, context: string, message: string): void; diff --git a/src/server.ts b/src/server.ts index 1bec50b0..46786164 100644 --- a/src/server.ts +++ b/src/server.ts @@ -3,8 +3,7 @@ import { Session } from "./session.js"; import { Transport } from "@modelcontextprotocol/sdk/shared/transport.js"; import { AtlasTools } from "./tools/atlas/tools.js"; import { MongoDbTools } from "./tools/mongodb/tools.js"; -import logger, { initializeLogger } from "./logger.js"; -import { mongoLogId } from "mongodb-log-writer"; +import logger, { initializeLogger, LogId } from "./logger.js"; import { ObjectId } from "mongodb"; import { Telemetry } from "./telemetry/telemetry.js"; import { UserConfig } from "./config.js"; @@ -23,7 +22,7 @@ export class Server { public readonly session: Session; private readonly mcpServer: McpServer; private readonly telemetry: Telemetry; - private readonly userConfig: UserConfig; + public readonly userConfig: UserConfig; private readonly startTime: number; constructor({ session, mcpServer, userConfig }: ServerOptions) { @@ -71,7 +70,7 @@ export class Server { this.session.sessionId = new ObjectId().toString(); logger.info( - mongoLogId(1_000_004), + LogId.serverInitialized, "server", `Server started with transport ${transport.constructor.name} and agent runner ${this.session.agentRunner?.name}` ); @@ -135,6 +134,32 @@ export class Server { } private registerResources() { + this.mcpServer.resource( + "config", + "config://config", + { + description: + "Server configuration, supplied by the user either as environment variables or as startup arguments", + }, + (uri) => { + const result = { + telemetry: this.userConfig.telemetry, + logPath: this.userConfig.logPath, + connectionString: this.userConfig.connectionString + ? "set; no explicit connect needed, use switch-connection tool to connect to a different connection if necessary" + : "not set; before using any mongodb tool, you need to call the connect tool with a connection string", + connectOptions: this.userConfig.connectOptions, + }; + return { + contents: [ + { + text: JSON.stringify(result), + uri: uri.href, + }, + ], + }; + } + ); if (this.userConfig.connectionString) { this.mcpServer.resource( "connection-string", diff --git a/src/session.ts b/src/session.ts index 2c5267ce..17357d6c 100644 --- a/src/session.ts +++ b/src/session.ts @@ -1,6 +1,7 @@ import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver"; import { ApiClient, ApiClientCredentials } from "./common/atlas/apiClient.js"; import { Implementation } from "@modelcontextprotocol/sdk/types.js"; +import EventEmitter from "events"; export interface SessionOptions { apiBaseUrl?: string; @@ -8,7 +9,9 @@ export interface SessionOptions { apiClientSecret?: string; } -export class Session { +export class Session extends EventEmitter<{ + close: []; +}> { sessionId?: string; serviceProvider?: NodeDriverServiceProvider; apiClient: ApiClient; @@ -18,6 +21,8 @@ export class Session { }; constructor({ apiBaseUrl, apiClientId, apiClientSecret }: SessionOptions = {}) { + super(); + const credentials: ApiClientCredentials | undefined = apiClientId && apiClientSecret ? { @@ -49,6 +54,8 @@ export class Session { console.error("Error closing service provider:", error); } this.serviceProvider = undefined; + + this.emit("close"); } } } diff --git a/src/telemetry/telemetry.ts b/src/telemetry/telemetry.ts index 518fc0d0..9b5986af 100644 --- a/src/telemetry/telemetry.ts +++ b/src/telemetry/telemetry.ts @@ -1,8 +1,7 @@ import { Session } from "../session.js"; import { BaseEvent, CommonProperties } from "./types.js"; import { config } from "../config.js"; -import logger from "../logger.js"; -import { mongoLogId } from "mongodb-log-writer"; +import logger, { LogId } from "../logger.js"; import { ApiClient } from "../common/atlas/apiClient.js"; import { MACHINE_METADATA } from "./constants.js"; import { EventCache } from "./eventCache.js"; @@ -61,7 +60,7 @@ export class Telemetry { await this.emit(events); } catch { - logger.debug(mongoLogId(1_000_002), "telemetry", `Error emitting telemetry events.`); + logger.debug(LogId.telemetryEmitFailure, "telemetry", `Error emitting telemetry events.`); } } @@ -89,7 +88,7 @@ export class Telemetry { const allEvents = [...cachedEvents, ...events]; logger.debug( - mongoLogId(1_000_003), + LogId.telemetryEmitStart, "telemetry", `Attempting to send ${allEvents.length} events (${cachedEvents.length} cached)` ); @@ -97,12 +96,12 @@ export class Telemetry { const result = await this.sendEvents(this.session.apiClient, allEvents); if (result.success) { this.eventCache.clearEvents(); - logger.debug(mongoLogId(1_000_004), "telemetry", `Sent ${allEvents.length} events successfully`); + logger.debug(LogId.telemetryEmitSuccess, "telemetry", `Sent ${allEvents.length} events successfully`); return; } - logger.warning( - mongoLogId(1_000_005), + logger.debug( + LogId.telemetryEmitFailure, "telemetry", `Error sending event to client: ${result.error instanceof Error ? result.error.message : String(result.error)}` ); diff --git a/src/tools/mongodb/metadata/connect.ts b/src/tools/mongodb/metadata/connect.ts index 746da9b3..defbf47f 100644 --- a/src/tools/mongodb/metadata/connect.ts +++ b/src/tools/mongodb/metadata/connect.ts @@ -2,92 +2,95 @@ import { z } from "zod"; import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { MongoDBToolBase } from "../mongodbTool.js"; import { ToolArgs, OperationType } from "../../tool.js"; -import { MongoError as DriverError } from "mongodb"; +import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; +import assert from "assert"; +import { UserConfig } from "../../../config.js"; +import { Telemetry } from "../../../telemetry/telemetry.js"; +import { Session } from "../../../session.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"; export class ConnectTool extends MongoDBToolBase { - protected name = "connect"; - protected description = "Connect to a MongoDB instance"; + protected name: typeof connectedName | typeof disconnectedName = disconnectedName; + protected description: typeof connectedDescription | typeof disconnectedDescription = disconnectedDescription; + + // 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 = { - options: z - .array( - z - .union([ - z.object({ - connectionString: z - .string() - .describe("MongoDB connection string (in the mongodb:// or mongodb+srv:// format)"), - }), - z.object({ - clusterName: z.string().describe("MongoDB cluster name"), - }), - ]) - .optional() - ) - .optional() - .describe( - "Options for connecting to MongoDB. If not provided, the connection string from the config://connection-string resource will be used. If the user hasn't specified Atlas cluster name or a connection string explicitly and the `config://connection-string` resource is present, always invoke this with no arguments." - ), + connectionString: z.string().optional(), }; protected operationType: OperationType = "metadata"; - protected async execute({ options: optionsArr }: ToolArgs): Promise { - const options = optionsArr?.[0]; - let connectionString: string; - if (!options && !this.config.connectionString) { - return { - content: [ - { type: "text", text: "No connection details provided." }, - { type: "text", text: "Please provide either a connection string or a cluster name" }, - ], - }; - } + constructor(session: Session, config: UserConfig, telemetry: Telemetry) { + super(session, config, telemetry); + session.on("close", () => { + this.updateMetadata(); + }); + } - if (!options) { - // eslint-disable-next-line @typescript-eslint/no-non-null-assertion - connectionString = this.config.connectionString!; - } else if ("connectionString" in options) { - connectionString = options.connectionString; - } else { - // TODO: https://github.com/mongodb-js/mongodb-mcp-server/issues/19 - // We don't support connecting via cluster name since we'd need to obtain the user credentials - // and fill in the connection string. - return { - content: [ - { - type: "text", - text: `Connecting via cluster name not supported yet. Please provide a connection string.`, - }, - ], - }; + 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; } - try { - await this.connectToMongoDB(connectionString); - return { - content: [{ type: "text", text: `Successfully connected to ${connectionString}.` }], - }; - } catch (error) { - // Sometimes the model will supply an incorrect connection string. If the user has configured - // a different one as environment variable or a cli argument, suggest using that one instead. - if ( - this.config.connectionString && - error instanceof DriverError && - this.config.connectionString !== connectionString - ) { - return { - content: [ - { - type: "text", - text: - `Failed to connect to MongoDB at '${connectionString}' due to error: '${error.message}.` + - `Your config lists a different connection string: '${this.config.connectionString}' - do you want to try connecting to it instead?`, - }, - ], - }; - } + await this.connectToMongoDB(connectionString); + this.updateMetadata(); + return { + content: [{ type: "text", text: "Successfully connected to MongoDB." }], + }; + } + + public register(server: McpServer): void { + super.register(server); - throw error; + this.updateMetadata(); + } + + private updateMetadata(): void { + if (this.config.connectionString || this.session.serviceProvider) { + this.update?.({ + name: connectedName, + description: connectedDescription, + inputSchema: connectedSchema, + }); + } else { + this.update?.({ + name: disconnectedName, + description: disconnectedDescription, + inputSchema: disconnectedSchema, + }); } } } diff --git a/src/tools/mongodb/mongodbTool.ts b/src/tools/mongodb/mongodbTool.ts index d818c7ab..7e067e0f 100644 --- a/src/tools/mongodb/mongodbTool.ts +++ b/src/tools/mongodb/mongodbTool.ts @@ -3,6 +3,7 @@ import { ToolArgs, ToolBase, ToolCategory } from "../tool.js"; import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver"; import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { ErrorCodes, MongoDBError } from "../../errors.js"; +import logger, { LogId } from "../../logger.js"; export const DbOperationArgs = { database: z.string().describe("Database name"), @@ -14,7 +15,16 @@ export abstract class MongoDBToolBase extends ToolBase { protected async ensureConnected(): Promise { if (!this.session.serviceProvider && this.config.connectionString) { - await this.connectToMongoDB(this.config.connectionString); + try { + await this.connectToMongoDB(this.config.connectionString); + } catch (error) { + logger.error( + LogId.mongodbConnectFailure, + "mongodbTool", + `Failed to connect to MongoDB instance using the connection string from the config: ${error as string}` + ); + throw new MongoDBError(ErrorCodes.MisconfiguredConnectionString, "Not connected to MongoDB."); + } } if (!this.session.serviceProvider) { @@ -28,20 +38,33 @@ export abstract class MongoDBToolBase extends ToolBase { error: unknown, args: ToolArgs ): Promise | CallToolResult { - if (error instanceof MongoDBError && error.code === ErrorCodes.NotConnectedToMongoDB) { - return { - content: [ - { - type: "text", - text: "You need to connect to a MongoDB instance before you can access its data.", - }, - { - type: "text", - text: "Please use the 'connect' tool to connect to a MongoDB instance.", - }, - ], - isError: true, - }; + if (error instanceof MongoDBError) { + switch (error.code) { + case ErrorCodes.NotConnectedToMongoDB: + return { + content: [ + { + type: "text", + text: "You need to connect to a MongoDB instance before you can access its data.", + }, + { + type: "text", + text: "Please use the 'connect' or 'switch-connection' tool to connect to a MongoDB instance.", + }, + ], + isError: true, + }; + case ErrorCodes.MisconfiguredConnectionString: + return { + content: [ + { + type: "text", + text: "The configured connection string is not valid. Please check the connection string and confirm it points to a valid MongoDB instance. Alternatively, use the 'switch-connection' tool to connect to a different instance.", + }, + ], + isError: true, + }; + } } return super.handleError(error, args); diff --git a/src/tools/tool.ts b/src/tools/tool.ts index a37c7224..9066f02f 100644 --- a/src/tools/tool.ts +++ b/src/tools/tool.ts @@ -1,9 +1,8 @@ -import { z, type ZodRawShape, type ZodNever } from "zod"; -import type { McpServer, ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js"; +import { z, type ZodRawShape, type ZodNever, AnyZodObject } from "zod"; +import type { McpServer, RegisteredTool, ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js"; import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { Session } from "../session.js"; -import logger from "../logger.js"; -import { mongoLogId } from "mongodb-log-writer"; +import logger, { LogId } from "../logger.js"; import { Telemetry } from "../telemetry/telemetry.js"; import { type ToolEvent } from "../telemetry/types.js"; import { UserConfig } from "../config.js"; @@ -63,17 +62,13 @@ export abstract class ToolBase { const callback: ToolCallback = async (...args) => { const startTime = Date.now(); try { - logger.debug( - mongoLogId(1_000_006), - "tool", - `Executing ${this.name} with args: ${JSON.stringify(args)}` - ); + logger.debug(LogId.toolExecute, "tool", `Executing ${this.name} with args: ${JSON.stringify(args)}`); const result = await this.execute(...args); await this.emitToolEvent(startTime, result); return result; } catch (error: unknown) { - logger.error(mongoLogId(1_000_000), "tool", `Error executing ${this.name}: ${error as string}`); + logger.error(LogId.toolExecuteFailure, "tool", `Error executing ${this.name}: ${error as string}`); const toolResult = await this.handleError(error, args[0] as ToolArgs); await this.emitToolEvent(startTime, toolResult).catch(() => {}); return toolResult; @@ -81,8 +76,36 @@ export abstract class ToolBase { }; server.tool(this.name, this.description, this.argsShape, 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. + this.update = (updates: { name?: string; description?: string; inputSchema?: AnyZodObject }) => { + const tools = server["_registeredTools"] as { [toolName: string]: RegisteredTool }; + const existingTool = tools[this.name]; + + if (updates.name && updates.name !== this.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.sendToolListChanged(); + }; } + protected update?: (updates: { name?: string; description?: string; inputSchema?: AnyZodObject }) => void; + // Checks if a tool is allowed to run based on the config protected verifyAllowed(): boolean { let errorClarification: string | undefined; @@ -96,7 +119,7 @@ export abstract class ToolBase { if (errorClarification) { logger.debug( - mongoLogId(1_000_010), + LogId.toolDisabled, "tool", `Prevented registration of ${this.name} because ${errorClarification} is disabled in the config` ); diff --git a/tests/integration/helpers.ts b/tests/integration/helpers.ts index 829fce73..98c8b970 100644 --- a/tests/integration/helpers.ts +++ b/tests/integration/helpers.ts @@ -20,11 +20,12 @@ export interface IntegrationTest { mcpServer: () => Server; } -export function setupIntegrationTest(userConfig: UserConfig = config): IntegrationTest { +export function setupIntegrationTest(getUserConfig: () => UserConfig = () => config): IntegrationTest { let mcpClient: Client | undefined; let mcpServer: Server | undefined; beforeAll(async () => { + const userConfig = getUserConfig(); const clientTransport = new InMemoryTransport(); const serverTransport = new InMemoryTransport(); @@ -50,6 +51,7 @@ export function setupIntegrationTest(userConfig: UserConfig = config): Integrati apiClientSecret: userConfig.apiClientSecret, }); + userConfig.telemetry = "disabled"; mcpServer = new Server({ session, userConfig, @@ -63,7 +65,15 @@ export function setupIntegrationTest(userConfig: UserConfig = config): Integrati }); beforeEach(() => { - config.telemetry = "disabled"; + if (mcpServer) { + mcpServer.userConfig.telemetry = "disabled"; + } + }); + + afterEach(async () => { + if (mcpServer) { + await mcpServer.session.close(); + } }); afterAll(async () => { diff --git a/tests/integration/server.test.ts b/tests/integration/server.test.ts index 3d12f129..341502c5 100644 --- a/tests/integration/server.test.ts +++ b/tests/integration/server.test.ts @@ -3,11 +3,11 @@ import { config } from "../../src/config.js"; describe("Server integration test", () => { describe("without atlas", () => { - const integration = setupIntegrationTest({ + const integration = setupIntegrationTest(() => ({ ...config, apiClientId: undefined, apiClientSecret: undefined, - }); + })); it("should return positive number of tools and have no atlas tools", async () => { const tools = await integration.mcpClient().listTools(); @@ -19,11 +19,11 @@ describe("Server integration test", () => { }); }); describe("with atlas", () => { - const integration = setupIntegrationTest({ + const integration = setupIntegrationTest(() => ({ ...config, apiClientId: "test", apiClientSecret: "test", - }); + })); describe("list capabilities", () => { it("should return positive number of tools and have some atlas tools", async () => { @@ -35,12 +35,6 @@ describe("Server integration test", () => { expect(atlasTools.length).toBeGreaterThan(0); }); - it("should return no resources", async () => { - await expect(() => integration.mcpClient().listResources()).rejects.toMatchObject({ - message: "MCP error -32601: Method not found", - }); - }); - it("should return no prompts", async () => { await expect(() => integration.mcpClient().listPrompts()).rejects.toMatchObject({ message: "MCP error -32601: Method not found", diff --git a/tests/integration/tools/mongodb/metadata/connect.test.ts b/tests/integration/tools/mongodb/metadata/connect.test.ts index 7992c5f2..ca138944 100644 --- a/tests/integration/tools/mongodb/metadata/connect.test.ts +++ b/tests/integration/tools/mongodb/metadata/connect.test.ts @@ -1,147 +1,124 @@ import { describeWithMongoDB } from "../mongodbHelpers.js"; - -import { getResponseContent, validateToolMetadata } from "../../../helpers.js"; - +import { getResponseContent, validateThrowsForInvalidArguments, validateToolMetadata } from "../../../helpers.js"; import { config } from "../../../../../src/config.js"; -describeWithMongoDB("Connect tool", (integration) => { - validateToolMetadata(integration, "connect", "Connect to a MongoDB instance", [ - { - name: "options", - description: - "Options for connecting to MongoDB. If not provided, the connection string from the config://connection-string resource will be used. If the user hasn't specified Atlas cluster name or a connection string explicitly and the `config://connection-string` resource is present, always invoke this with no arguments.", - type: "array", - required: false, - }, - ]); - - describe("without arguments", () => { - it("prompts for connection string if not set", async () => { - const response = await integration.mcpClient().callTool({ name: "connect" }); - const content = getResponseContent(response.content); - expect(content).toContain("No connection details provided"); +describeWithMongoDB( + "switchConnection tool", + (integration) => { + beforeEach(() => { + integration.mcpServer().userConfig.connectionString = integration.connectionString(); }); - it("connects to the database if connection string is set", async () => { - config.connectionString = integration.connectionString(); - - const response = await integration.mcpClient().callTool({ name: "connect" }); - const content = getResponseContent(response.content); - expect(content).toContain("Successfully connected"); - expect(content).toContain(integration.connectionString()); - }); - }); + validateToolMetadata( + integration, + "switch-connection", + "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.", + [ + { + name: "connectionString", + description: "MongoDB connection string to switch to (in the mongodb:// or mongodb+srv:// format)", + type: "string", + required: false, + }, + ] + ); - describe("with default config", () => { - describe("without connection string", () => { - it("prompts for connection string", async () => { - const response = await integration.mcpClient().callTool({ name: "connect", arguments: {} }); - const content = getResponseContent(response.content); - expect(content).toContain("No connection details provided"); - }); - }); + validateThrowsForInvalidArguments(integration, "switch-connection", [{ connectionString: 123 }]); - describe("with connection string", () => { + describe("without arguments", () => { it("connects to the database", async () => { - const response = await integration.mcpClient().callTool({ - name: "connect", - arguments: { - options: [ - { - connectionString: integration.connectionString(), - }, - ], - }, - }); + const response = await integration.mcpClient().callTool({ name: "switch-connection" }); const content = getResponseContent(response.content); expect(content).toContain("Successfully connected"); - expect(content).toContain(integration.connectionString()); }); }); - describe("with invalid connection string", () => { - it("returns error message", async () => { - const response = await integration.mcpClient().callTool({ - name: "connect", - arguments: { options: [{ connectionString: "mongodb://localhost:12345" }] }, - }); - const content = getResponseContent(response.content); - expect(content).toContain("Error running connect"); - - // Should not suggest using the config connection string (because we don't have one) - expect(content).not.toContain("Your config lists a different connection string"); - }); + it("doesn't have the connect tool registered", async () => { + const { tools } = await integration.mcpClient().listTools(); + const tool = tools.find((tool) => tool.name === "connect"); + expect(tool).toBeUndefined(); }); - }); - describe("with connection string in config", () => { - beforeEach(() => { - config.connectionString = integration.connectionString(); - }); - - it("uses the connection string from config", async () => { - const response = await integration.mcpClient().callTool({ name: "connect", arguments: {} }); + it("defaults to the connection string from config", async () => { + const response = await integration.mcpClient().callTool({ name: "switch-connection", arguments: {} }); const content = getResponseContent(response.content); expect(content).toContain("Successfully connected"); - expect(content).toContain(integration.connectionString()); }); - it("prefers connection string from arguments", async () => { + it("switches to the connection string from the arguments", async () => { const newConnectionString = `${integration.connectionString()}?appName=foo-bar`; const response = await integration.mcpClient().callTool({ - name: "connect", + name: "switch-connection", arguments: { - options: [ - { - connectionString: newConnectionString, - }, - ], + connectionString: newConnectionString, }, }); const content = getResponseContent(response.content); expect(content).toContain("Successfully connected"); - expect(content).toContain(newConnectionString); }); describe("when the arugment connection string is invalid", () => { - it("suggests the config connection string if set", async () => { + it("returns error message", async () => { const response = await integration.mcpClient().callTool({ - name: "connect", + name: "switch-connection", arguments: { - options: [ - { - connectionString: "mongodb://localhost:12345", - }, - ], + connectionString: "mongodb://localhost:12345", }, }); + const content = getResponseContent(response.content); - expect(content).toContain("Failed to connect to MongoDB at 'mongodb://localhost:12345'"); - expect(content).toContain( - `Your config lists a different connection string: '${config.connectionString}' - do you want to try connecting to it instead?` - ); + + expect(content).toContain("Error running switch-connection"); }); + }); + }, + (mdbIntegration) => ({ + ...config, + connectionString: mdbIntegration.connectionString(), + }) +); +describeWithMongoDB("Connect tool", (integration) => { + validateToolMetadata(integration, "connect", "Connect to a MongoDB instance", [ + { + name: "connectionString", + description: "MongoDB connection string (in the mongodb:// or mongodb+srv:// format)", + type: "string", + required: true, + }, + ]); - it("returns error message if the config connection string matches the argument", async () => { - config.connectionString = "mongodb://localhost:12345"; - const response = await integration.mcpClient().callTool({ - name: "connect", - arguments: { - options: [ - { - connectionString: "mongodb://localhost:12345", - }, - ], - }, - }); + validateThrowsForInvalidArguments(integration, "connect", [{}, { connectionString: 123 }]); - const content = getResponseContent(response.content); + it("doesn't have the switch-connection tool registered", async () => { + const { tools } = await integration.mcpClient().listTools(); + const tool = tools.find((tool) => tool.name === "switch-connection"); + expect(tool).toBeUndefined(); + }); - // Should be handled by default error handler and not suggest the config connection string - // because it matches the argument connection string - expect(content).toContain("Error running connect"); - expect(content).not.toContain("Your config lists a different connection string"); + describe("with connection string", () => { + it("connects to the database", async () => { + const response = await integration.mcpClient().callTool({ + name: "connect", + arguments: { + connectionString: integration.connectionString(), + }, }); + const content = getResponseContent(response.content); + expect(content).toContain("Successfully connected"); + }); + }); + + describe("with invalid connection string", () => { + it("returns error message", async () => { + const response = await integration.mcpClient().callTool({ + name: "connect", + arguments: { connectionString: "mongodb://localhost:12345" }, + }); + const content = getResponseContent(response.content); + expect(content).toContain("Error running connect"); + + // Should not suggest using the config connection string (because we don't have one) + expect(content).not.toContain("Your config lists a different connection string"); }); }); }); diff --git a/tests/integration/tools/mongodb/mongodbHelpers.ts b/tests/integration/tools/mongodb/mongodbHelpers.ts index b6bd47d7..807b0d20 100644 --- a/tests/integration/tools/mongodb/mongodbHelpers.ts +++ b/tests/integration/tools/mongodb/mongodbHelpers.ts @@ -3,29 +3,41 @@ import path from "path"; import fs from "fs/promises"; import { MongoClient, ObjectId } from "mongodb"; import { getResponseContent, IntegrationTest, setupIntegrationTest } from "../../helpers.js"; -import { config } from "../../../../src/config.js"; +import { config, UserConfig } from "../../../../src/config.js"; interface MongoDBIntegrationTest { mongoClient: () => MongoClient; connectionString: () => string; - connectMcpClient: () => Promise; randomDbName: () => string; } export function describeWithMongoDB( name: string, - fn: (integration: IntegrationTest & MongoDBIntegrationTest) => void -): void { - describe("mongodb", () => { - const integration = setupIntegrationTest(); - const mdbIntegration = setupMongoDBIntegrationTest(integration); - describe(name, () => { - fn({ ...integration, ...mdbIntegration }); + fn: (integration: IntegrationTest & MongoDBIntegrationTest & { connectMcpClient: () => Promise }) => void, + getUserConfig: (mdbIntegration: MongoDBIntegrationTest) => UserConfig = () => config +) { + describe(name, () => { + const mdbIntegration = setupMongoDBIntegrationTest(); + const integration = setupIntegrationTest(() => getUserConfig(mdbIntegration)); + + afterEach(() => { + integration.mcpServer().userConfig.connectionString = undefined; + }); + + fn({ + ...integration, + ...mdbIntegration, + connectMcpClient: async () => { + await integration.mcpClient().callTool({ + name: "connect", + arguments: { connectionString: mdbIntegration.connectionString() }, + }); + }, }); }); } -export function setupMongoDBIntegrationTest(integration: IntegrationTest): MongoDBIntegrationTest { +export function setupMongoDBIntegrationTest(): MongoDBIntegrationTest { let mongoCluster: // TODO: Fix this type once mongodb-runner is updated. | { connectionString: string; @@ -40,9 +52,6 @@ export function setupMongoDBIntegrationTest(integration: IntegrationTest): Mongo }); afterEach(async () => { - await integration.mcpServer().session.close(); - config.connectionString = undefined; - await mongoClient?.close(); mongoClient = undefined; }); @@ -108,12 +117,7 @@ export function setupMongoDBIntegrationTest(integration: IntegrationTest): Mongo return mongoClient; }, connectionString: getConnectionString, - connectMcpClient: async () => { - await integration.mcpClient().callTool({ - name: "connect", - arguments: { options: [{ connectionString: getConnectionString() }] }, - }); - }, + randomDbName: () => randomDbName, }; } @@ -134,7 +138,7 @@ export function validateAutoConnectBehavior( } it("connects automatically if connection string is configured", async () => { - config.connectionString = integration.connectionString(); + integration.mcpServer().userConfig.connectionString = integration.connectionString(); const validationInfo = validation();