diff --git a/package-lock.json b/package-lock.json index 1385e8ca6..229e3081b 100644 --- a/package-lock.json +++ b/package-lock.json @@ -75,6 +75,7 @@ "node": "^20.19.0 || ^22.12.0 || >= 23.0.0" }, "optionalDependencies": { + "@mongodb-js-preview/atlas-local": "^0.0.0-preview.1", "kerberos": "^2.2.2" } }, @@ -2048,6 +2049,103 @@ "node": ">=16.20.0" } }, + "node_modules/@mongodb-js-preview/atlas-local": { + "version": "0.0.0-preview.1", + "resolved": "https://registry.npmjs.org/@mongodb-js-preview/atlas-local/-/atlas-local-0.0.0-preview.1.tgz", + "integrity": "sha512-py3roloK+dyq9bCU139f3JdFykige1kWwUli9qWE4daODFdJ0mvQPN1EChw3lzI4rv53cX1CvApbh20liOukoQ==", + "license": "Apache-2.0", + "optional": true, + "engines": { + "node": ">= 12.22.0 < 13 || >= 14.17.0 < 15 || >= 15.12.0 < 16 || >= 16.0.0" + }, + "optionalDependencies": { + "@mongodb-js-preview/atlas-local-darwin-arm64": "0.0.0-preview.1", + "@mongodb-js-preview/atlas-local-darwin-x64": "0.0.0-preview.1", + "@mongodb-js-preview/atlas-local-linux-arm64-gnu": "0.0.0-preview.1", + "@mongodb-js-preview/atlas-local-linux-x64-gnu": "0.0.0-preview.1", + "@mongodb-js-preview/atlas-local-win32-x64-msvc": "0.0.0-preview.1" + } + }, + "node_modules/@mongodb-js-preview/atlas-local-darwin-arm64": { + "version": "0.0.0-preview.1", + "resolved": "https://registry.npmjs.org/@mongodb-js-preview/atlas-local-darwin-arm64/-/atlas-local-darwin-arm64-0.0.0-preview.1.tgz", + "integrity": "sha512-TcH7CFCg6pAx0KPhTUOyaZRwXOOTb5WCo9on12GqEk/oM+vERwfK5ztGSZns45IvxgN6zUsERoV+O6SEoK1gsA==", + "cpu": [ + "arm64" + ], + "license": "Apache-2.0", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": ">= 12.22.0 < 13 || >= 14.17.0 < 15 || >= 15.12.0 < 16 || >= 16.0.0" + } + }, + "node_modules/@mongodb-js-preview/atlas-local-darwin-x64": { + "version": "0.0.0-preview.1", + "resolved": "https://registry.npmjs.org/@mongodb-js-preview/atlas-local-darwin-x64/-/atlas-local-darwin-x64-0.0.0-preview.1.tgz", + "integrity": "sha512-KmG+xKCS5f3adhznYH569mq0PHrFoGuqsGN5XEtVtUEYgv/gQicgJ0voWMrwTHu3jIFFQeGEvMgKsciyXAlVaQ==", + "cpu": [ + "x64" + ], + "license": "Apache-2.0", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": ">= 12.22.0 < 13 || >= 14.17.0 < 15 || >= 15.12.0 < 16 || >= 16.0.0" + } + }, + "node_modules/@mongodb-js-preview/atlas-local-linux-arm64-gnu": { + "version": "0.0.0-preview.1", + "resolved": "https://registry.npmjs.org/@mongodb-js-preview/atlas-local-linux-arm64-gnu/-/atlas-local-linux-arm64-gnu-0.0.0-preview.1.tgz", + "integrity": "sha512-0ImE3RUdWiO38JWXiG6xAZzpz3CA2MHfpQdcwIomF0ldw/14ofRYkH31KX0M444j2rS2/AHBa+zdswYqFZCQbg==", + "cpu": [ + "arm64" + ], + "license": "Apache-2.0", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 12.22.0 < 13 || >= 14.17.0 < 15 || >= 15.12.0 < 16 || >= 16.0.0" + } + }, + "node_modules/@mongodb-js-preview/atlas-local-linux-x64-gnu": { + "version": "0.0.0-preview.1", + "resolved": "https://registry.npmjs.org/@mongodb-js-preview/atlas-local-linux-x64-gnu/-/atlas-local-linux-x64-gnu-0.0.0-preview.1.tgz", + "integrity": "sha512-AM+s8ygWU5gkNm6rDkLnWueOIon9T3kaeSleo4qgeAt4rgA7C9f2XUkQmFsv8b1E9g6CbNYrrbAAcnc0xVMtLQ==", + "cpu": [ + "x64" + ], + "license": "Apache-2.0", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 12.22.0 < 13 || >= 14.17.0 < 15 || >= 15.12.0 < 16 || >= 16.0.0" + } + }, + "node_modules/@mongodb-js-preview/atlas-local-win32-x64-msvc": { + "version": "0.0.0-preview.1", + "resolved": "https://registry.npmjs.org/@mongodb-js-preview/atlas-local-win32-x64-msvc/-/atlas-local-win32-x64-msvc-0.0.0-preview.1.tgz", + "integrity": "sha512-1oEsFgKz4Hatp+lD4pIB5EcrCqSwAx3p4FEJqwCKHbYGIsSVuvjhvN28hZHPHqkQhM9l8ZWFa8Z87/LR/Hy8FQ==", + "cpu": [ + "x64" + ], + "license": "Apache-2.0", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": ">= 12.22.0 < 13 || >= 14.17.0 < 15 || >= 15.12.0 < 16 || >= 16.0.0" + } + }, "node_modules/@mongodb-js/device-id": { "version": "0.3.1", "resolved": "https://registry.npmjs.org/@mongodb-js/device-id/-/device-id-0.3.1.tgz", diff --git a/package.json b/package.json index cc7e68cc9..7246ec1de 100644 --- a/package.json +++ b/package.json @@ -121,6 +121,7 @@ "node": "^20.19.0 || ^22.12.0 || >= 23.0.0" }, "optionalDependencies": { + "@mongodb-js-preview/atlas-local": "^0.0.0-preview.1", "kerberos": "^2.2.2" } } diff --git a/src/common/session.ts b/src/common/session.ts index 24946b171..92d9e4f4f 100644 --- a/src/common/session.ts +++ b/src/common/session.ts @@ -15,6 +15,7 @@ import type { import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver"; import { ErrorCodes, MongoDBError } from "./errors.js"; import type { ExportsManager } from "./exportsManager.js"; +import type { Client } from "@mongodb-js-preview/atlas-local"; import type { Keychain } from "./keychain.js"; export interface SessionOptions { @@ -46,6 +47,7 @@ export class Session extends EventEmitter { version?: string; title?: string; }; + atlasLocalClient?: Client; public logger: CompositeLogger; @@ -99,6 +101,10 @@ export class Session extends EventEmitter { this.connectionManager.setClientName(this.mcpClient.name || "unknown"); } + setAtlasLocalClient(atlasLocalClient: Client): void { + this.atlasLocalClient = atlasLocalClient; + } + async disconnect(): Promise { const atlasCluster = this.connectedAtlasCluster; diff --git a/src/server.ts b/src/server.ts index 5db421dbb..74afc93b2 100644 --- a/src/server.ts +++ b/src/server.ts @@ -19,7 +19,7 @@ import { UnsubscribeRequestSchema, } from "@modelcontextprotocol/sdk/types.js"; import assert from "assert"; -import type { ToolBase } from "./tools/tool.js"; +import type { ToolBase, ToolConstructor } from "./tools/tool.js"; import { validateConnectionString } from "./helpers/connectionOptions.js"; import { packageInfo } from "./common/packageInfo.js"; import { type ConnectionErrorHandler } from "./common/connectionErrorHandler.js"; @@ -69,6 +69,9 @@ export class Server { // TODO: Eventually we might want to make tools reactive too instead of relying on custom logic. this.registerTools(); + // Atlas Local tools are optional and require async initialization + void this.registerAtlasLocalTools(); + // This is a workaround for an issue we've seen with some models, where they'll see that everything in the `arguments` // object is optional, and then not pass it at all. However, the MCP server expects the `arguments` object to be if // the tool accepts any arguments, even if they're all optional. @@ -197,8 +200,41 @@ export class Server { this.telemetry.emitEvents([event]).catch(() => {}); } + private async registerAtlasLocalTools(): Promise { + // If Atlas Local tools are disabled, don't attempt to connect to the client + if (this.userConfig.disabledTools.includes("atlas-local")) { + return; + } + + try { + // Import Atlas Local client asyncronously + // This will fail on unsupported platforms + const { Client: AtlasLocalClient } = await import("@mongodb-js-preview/atlas-local"); + + // Connect to Atlas Local client + // This will fail if docker is not running + const client = AtlasLocalClient.connect(); + + // Set Atlas Local client + this.session.setAtlasLocalClient(client); + + // Register Atlas Local tools + this.registerToolInstances(AtlasLocalTools); + } catch (error) { + console.warn( + "Failed to initialize Atlas Local client, atlas-local tools will be disabled (error: ", + error, + ")" + ); + } + } + private registerTools(): void { - for (const toolConstructor of [...AtlasTools, ...AtlasLocalTools, ...MongoDbTools]) { + this.registerToolInstances([...AtlasTools, ...MongoDbTools]); + } + + private registerToolInstances(tools: Array): void { + for (const toolConstructor of tools) { const tool = new toolConstructor(this.session, this.userConfig, this.telemetry); if (tool.register(this)) { this.tools.push(tool); diff --git a/src/tools/atlasLocal/atlasLocalTool.ts b/src/tools/atlasLocal/atlasLocalTool.ts index b5c7899fa..8aca9d551 100644 --- a/src/tools/atlasLocal/atlasLocalTool.ts +++ b/src/tools/atlasLocal/atlasLocalTool.ts @@ -1,10 +1,45 @@ import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; -import type { ToolArgs, ToolCategory } from "../tool.js"; +import type { TelemetryToolMetadata, ToolArgs, ToolCategory } from "../tool.js"; import { ToolBase } from "../tool.js"; +import type { ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js"; +import type { Client } from "@mongodb-js-preview/atlas-local"; export abstract class AtlasLocalToolBase extends ToolBase { public category: ToolCategory = "atlas-local"; + protected verifyAllowed(): boolean { + return this.session.atlasLocalClient !== undefined && super.verifyAllowed(); + } + + protected async execute(): Promise { + // Get the client + const client = this.session.atlasLocalClient; + + // If the client is not found, throw an error + // This should never happen: + // - atlas-local tools are only added after the client is set + // this means that if we were unable to get the client, the tool will not be registered + // - in case the tool was registered by accident + // verifyAllowed in the base class would still return false preventing the tool from being registered, + // preventing the tool from being executed + if (!client) { + return { + content: [ + { + type: "text", + text: `Something went wrong on our end, this tool should have been disabled but it was not. +please log a ticket here: https://github.com/mongodb-js/mongodb-mcp-server/issues/new?template=bug_report.yml`, + }, + ], + isError: true, + }; + } + + return this.executeWithAtlasLocalClient(client); + } + + protected abstract executeWithAtlasLocalClient(client: Client): Promise; + protected handleError( error: unknown, args: ToolArgs @@ -14,4 +49,12 @@ export abstract class AtlasLocalToolBase extends ToolBase { // For other types of errors, use the default error handling from the base class return super.handleError(error, args); } + + protected resolveTelemetryMetadata( + ...args: Parameters> + ): TelemetryToolMetadata { + // TODO: include deployment id in the metadata where possible + void args; // this shuts up the eslint rule until we implement the TODO above + return {}; + } } diff --git a/src/tools/atlasLocal/read/listDeployments.ts b/src/tools/atlasLocal/read/listDeployments.ts new file mode 100644 index 000000000..3716efd71 --- /dev/null +++ b/src/tools/atlasLocal/read/listDeployments.ts @@ -0,0 +1,46 @@ +import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; +import { AtlasLocalToolBase } from "../atlasLocalTool.js"; +import type { OperationType } from "../../tool.js"; +import { formatUntrustedData } from "../../tool.js"; +import type { Deployment } from "@mongodb-js-preview/atlas-local"; +import type { Client } from "@mongodb-js-preview/atlas-local"; + +export class ListDeploymentsTool extends AtlasLocalToolBase { + public name = "atlas-local-list-deployments"; + protected description = "List MongoDB Atlas local deployments"; + public operationType: OperationType = "read"; + protected argsShape = {}; + + protected async executeWithAtlasLocalClient(client: Client): Promise { + // List the deployments + const deployments = await client.listDeployments(); + + // Format the deployments + return this.formatDeploymentsTable(deployments); + } + + private formatDeploymentsTable(deployments: Deployment[]): CallToolResult { + // Check if deployments are absent + if (!deployments?.length) { + return { + content: [{ type: "text", text: "No deployments found." }], + }; + } + + // Turn the deployments into a markdown table + const rows = deployments + .map((deployment) => { + return `${deployment.name || "Unknown"} | ${deployment.state} | ${deployment.mongodbVersion}`; + }) + .join("\n"); + + return { + content: formatUntrustedData( + `Found ${deployments.length} deployments:`, + `Deployment Name | State | MongoDB Version +----------------|----------------|---------------- +${rows}` + ), + }; + } +} diff --git a/src/tools/atlasLocal/tools.ts b/src/tools/atlasLocal/tools.ts index be1a95529..6d8cf7a5b 100644 --- a/src/tools/atlasLocal/tools.ts +++ b/src/tools/atlasLocal/tools.ts @@ -1 +1,3 @@ -export const AtlasLocalTools = []; +import { ListDeploymentsTool } from "./read/listDeployments.js"; + +export const AtlasLocalTools = [ListDeploymentsTool]; diff --git a/src/tools/tool.ts b/src/tools/tool.ts index bfedf0c5c..85f166b3a 100644 --- a/src/tools/tool.ts +++ b/src/tools/tool.ts @@ -18,6 +18,8 @@ export type TelemetryToolMetadata = { orgId?: string; }; +export type ToolConstructor = new (session: Session, config: UserConfig, telemetry: Telemetry) => ToolBase; + export abstract class ToolBase { public abstract name: string; diff --git a/tests/integration/helpers.ts b/tests/integration/helpers.ts index 1f28995dd..e2a7bce94 100644 --- a/tests/integration/helpers.ts +++ b/tests/integration/helpers.ts @@ -15,6 +15,7 @@ import { MCPConnectionManager } from "../../src/common/connectionManager.js"; import { DeviceId } from "../../src/helpers/deviceId.js"; import { connectionErrorHandler } from "../../src/common/connectionErrorHandler.js"; import { Keychain } from "../../src/common/keychain.js"; +import type { Client as AtlasLocalClient } from "@mongodb-js-preview/atlas-local"; interface ParameterInfo { name: string; @@ -345,6 +346,37 @@ export function waitUntil( }); } +export function waitUntilMcpClientIsSet( + mcpServer: Server, + signal: AbortSignal, + timeout: number = 5000 +): Promise { + let ts: NodeJS.Timeout | undefined; + + const timeoutSignal = AbortSignal.timeout(timeout); + const combinedSignal = AbortSignal.any([signal, timeoutSignal]); + + return new Promise((resolve, reject) => { + ts = setInterval(() => { + if (combinedSignal.aborted) { + return reject(new Error(`Aborted: ${combinedSignal.reason}`)); + } + + // wait until session.client != undefined + // do not wait more than 1 second, should take a few milliseconds at most + // try every 50ms to see if the client is set, if it's not set after 1 second, throw an error + const client = mcpServer.session.atlasLocalClient; + if (client) { + return resolve(client); + } + }, 100); + }).finally(() => { + if (ts !== undefined) { + clearInterval(ts); + } + }); +} + export function getDataFromUntrustedContent(content: string): string { const regex = /^[ \t]*(?.*)^[ \t]*<\/untrusted-user-data-[0-9a-f\\-]*>/gms; const match = regex.exec(content); diff --git a/tests/integration/server.test.ts b/tests/integration/server.test.ts index ef98075a7..446c31215 100644 --- a/tests/integration/server.test.ts +++ b/tests/integration/server.test.ts @@ -11,7 +11,9 @@ describe("Server integration test", () => { expectDefined(tools); expect(tools.tools.length).toBeGreaterThan(0); - const atlasTools = tools.tools.filter((tool) => tool.name.startsWith("atlas-")); + const atlasTools = tools.tools.filter( + (tool) => tool.name.startsWith("atlas-") && !tool.name.startsWith("atlas-local-") + ); expect(atlasTools.length).toBeLessThanOrEqual(0); }); }, diff --git a/tests/integration/tools/atlas-local/listDeployments.test.ts b/tests/integration/tools/atlas-local/listDeployments.test.ts new file mode 100644 index 000000000..38d2446c7 --- /dev/null +++ b/tests/integration/tools/atlas-local/listDeployments.test.ts @@ -0,0 +1,72 @@ +import { + defaultDriverOptions, + defaultTestConfig, + expectDefined, + getResponseElements, + setupIntegrationTest, + waitUntilMcpClientIsSet, +} from "../../helpers.js"; +import { describe, expect, it } from "vitest"; + +const isMacOSInGitHubActions = process.platform === "darwin" && process.env.GITHUB_ACTIONS === "true"; + +// Docker is not available on macOS in GitHub Actions +// That's why we skip the tests on macOS in GitHub Actions +describe("atlas-local-list-deployments", () => { + const integration = setupIntegrationTest( + () => defaultTestConfig, + () => defaultDriverOptions + ); + + it.skipIf(isMacOSInGitHubActions)("should have the atlas-local-list-deployments tool", async ({ signal }) => { + await waitUntilMcpClientIsSet(integration.mcpServer(), signal); + + const { tools } = await integration.mcpClient().listTools(); + const listDeployments = tools.find((tool) => tool.name === "atlas-local-list-deployments"); + expectDefined(listDeployments); + }); + + it.skipIf(!isMacOSInGitHubActions)( + "[MacOS in GitHub Actions] should not have the atlas-local-list-deployments tool", + async ({ signal }) => { + // This should throw an error because the client is not set within the timeout of 5 seconds (default) + await expect(waitUntilMcpClientIsSet(integration.mcpServer(), signal)).rejects.toThrow(); + + const { tools } = await integration.mcpClient().listTools(); + const listDeployments = tools.find((tool) => tool.name === "atlas-local-list-deployments"); + expect(listDeployments).toBeUndefined(); + } + ); + + it.skipIf(isMacOSInGitHubActions)("should have correct metadata", async ({ signal }) => { + await waitUntilMcpClientIsSet(integration.mcpServer(), signal); + const { tools } = await integration.mcpClient().listTools(); + const listDeployments = tools.find((tool) => tool.name === "atlas-local-list-deployments"); + expectDefined(listDeployments); + expect(listDeployments.inputSchema.type).toBe("object"); + expectDefined(listDeployments.inputSchema.properties); + expect(listDeployments.inputSchema.properties).toEqual({}); + }); + + it.skipIf(isMacOSInGitHubActions)("should not crash when calling the tool", async ({ signal }) => { + await waitUntilMcpClientIsSet(integration.mcpServer(), signal); + + const response = await integration.mcpClient().callTool({ + name: "atlas-local-list-deployments", + arguments: {}, + }); + const elements = getResponseElements(response.content); + expect(elements.length).toBeGreaterThanOrEqual(1); + + if (elements.length === 1) { + expect(elements[0]?.text).toContain("No deployments found."); + } + + if (elements.length > 1) { + expect(elements[0]?.text).toMatch(/Found \d+ deployments/); + expect(elements[1]?.text).toContain( + "Deployment Name | State | MongoDB Version\n----------------|----------------|----------------\n" + ); + } + }); +}); diff --git a/tests/integration/transports/stdio.test.ts b/tests/integration/transports/stdio.test.ts index aaa61d638..a5ba4259b 100644 --- a/tests/integration/transports/stdio.test.ts +++ b/tests/integration/transports/stdio.test.ts @@ -10,7 +10,7 @@ describeWithMongoDB("StdioRunner", (integration) => { beforeAll(async () => { transport = new StdioClientTransport({ command: "node", - args: ["dist/index.js"], + args: ["dist/index.js", "--disabledTools", "atlas-local"], env: { MDB_MCP_TRANSPORT: "stdio", MDB_MCP_CONNECTION_STRING: integration.connectionString(),