Skip to content

Commit 06da4c5

Browse files
adressed PR remarks
1 parent e6b0d15 commit 06da4c5

File tree

6 files changed

+50
-51
lines changed

6 files changed

+50
-51
lines changed

src/common/session.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import type {
1515
import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
1616
import { ErrorCodes, MongoDBError } from "./errors.js";
1717
import type { ExportsManager } from "./exportsManager.js";
18+
import type { Client } from "@mongodb-js-preview/atlas-local";
1819

1920
export interface SessionOptions {
2021
apiBaseUrl: string;
@@ -42,6 +43,7 @@ export class Session extends EventEmitter<SessionEvents> {
4243
version?: string;
4344
title?: string;
4445
};
46+
atlasLocalClient?: Client;
4547

4648
public logger: CompositeLogger;
4749

@@ -93,6 +95,10 @@ export class Session extends EventEmitter<SessionEvents> {
9395
this.connectionManager.setClientName(this.mcpClient.name || "unknown");
9496
}
9597

98+
setAtlasLocalClient(atlasLocalClient: Client): void {
99+
this.atlasLocalClient = atlasLocalClient;
100+
}
101+
96102
async disconnect(): Promise<void> {
97103
const atlasCluster = this.connectedAtlasCluster;
98104

src/server.ts

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
22
import type { Session } from "./common/session.js";
33
import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
44
import { AtlasTools } from "./tools/atlas/tools.js";
5-
import { BuildAtlasLocalTools } from "./tools/atlasLocal/tools.js";
5+
import { AtlasLocalTools } from "./tools/atlasLocal/tools.js";
66
import { MongoDbTools } from "./tools/mongodb/tools.js";
77
import { Resources } from "./resources/resources.js";
88
import type { LogLevel } from "./common/logger.js";
@@ -61,8 +61,10 @@ export class Server {
6161

6262
this.mcpServer.server.registerCapabilities({ logging: {}, resources: { listChanged: true, subscribe: true } });
6363

64+
await this.tryInitializeAtlasLocalClient();
65+
6466
// TODO: Eventually we might want to make tools reactive too instead of relying on custom logic.
65-
await this.registerTools();
67+
this.registerTools();
6668

6769
// This is a workaround for an issue we've seen with some models, where they'll see that everything in the `arguments`
6870
// object is optional, and then not pass it at all. However, the MCP server expects the `arguments` object to be if
@@ -193,9 +195,23 @@ export class Server {
193195
this.telemetry.emitEvents([event]).catch(() => {});
194196
}
195197

196-
private async registerTools(): Promise<void> {
197-
const atlasLocalTools = await BuildAtlasLocalTools();
198-
for (const toolConstructor of [...AtlasTools, ...atlasLocalTools, ...MongoDbTools]) {
198+
private async tryInitializeAtlasLocalClient(): Promise<void> {
199+
try {
200+
const { Client: AtlasLocalClient } = await import("@mongodb-js-preview/atlas-local");
201+
202+
const client = AtlasLocalClient.connect();
203+
this.session.setAtlasLocalClient(client);
204+
} catch (error) {
205+
console.warn(
206+
"Failed to initialize Atlas Local client, atlas-local tools will be disabled (error: ",
207+
error,
208+
")"
209+
);
210+
}
211+
}
212+
213+
private registerTools(): void {
214+
for (const toolConstructor of [...AtlasTools, ...AtlasLocalTools, ...MongoDbTools]) {
199215
const tool = new toolConstructor(this.session, this.userConfig, this.telemetry);
200216
if (tool.register(this)) {
201217
this.tools.push(tool);

src/tools/atlasLocal/atlasLocalTool.ts

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,17 +2,31 @@ import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
22
import type { TelemetryToolMetadata, ToolArgs, ToolCategory } from "../tool.js";
33
import { ToolBase } from "../tool.js";
44
import type { ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js";
5-
import type AtlasLocal from "@mongodb-js-preview/atlas-local";
5+
import type { Client } from "@mongodb-js-preview/atlas-local";
66

77
export abstract class AtlasLocalToolBase extends ToolBase {
88
public category: ToolCategory = "atlas-local";
9-
// Will be injected by BuildAtlasLocalTools() in atlasLocal/tools.ts
10-
public client?: AtlasLocal.Client;
119

1210
protected verifyAllowed(): boolean {
13-
return this.client !== undefined && super.verifyAllowed();
11+
return this.session.atlasLocalClient !== undefined && super.verifyAllowed();
1412
}
1513

14+
protected async execute(): Promise<CallToolResult> {
15+
// Get the client
16+
const client = this.session.atlasLocalClient;
17+
18+
// If the client is not found, throw an error
19+
// This should never happen, because the tool should have been disabled.
20+
// verifyAllowed in the base class returns false if the client is not found
21+
if (!client) {
22+
throw new Error("Atlas Local client not found, tool should have been disabled.");
23+
}
24+
25+
return this.executeWithAtlasLocalClient(client);
26+
}
27+
28+
protected abstract executeWithAtlasLocalClient(client: Client): Promise<CallToolResult>;
29+
1630
protected handleError(
1731
error: unknown,
1832
args: ToolArgs<typeof this.argsShape>

src/tools/atlasLocal/read/listDeployments.ts

Lines changed: 2 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,24 +3,15 @@ import { AtlasLocalToolBase } from "../atlasLocalTool.js";
33
import type { OperationType } from "../../tool.js";
44
import { formatUntrustedData } from "../../tool.js";
55
import type { Deployment } from "@mongodb-js-preview/atlas-local";
6+
import type { Client } from "@mongodb-js-preview/atlas-local";
67

78
export class ListDeploymentsTool extends AtlasLocalToolBase {
89
public name = "atlas-local-list-deployments";
910
protected description = "List MongoDB Atlas local deployments";
1011
public operationType: OperationType = "read";
1112
protected argsShape = {};
1213

13-
protected async execute(): Promise<CallToolResult> {
14-
// Get the client
15-
const client = this.client;
16-
17-
// If the client is not found, throw an error
18-
// This should never happen, because the tool should have been disabled.
19-
// verifyAllowed in the base class returns false if the client is not found
20-
if (!client) {
21-
throw new Error("Atlas Local client not found, tool should have been disabled.");
22-
}
23-
14+
protected async executeWithAtlasLocalClient(client: Client): Promise<CallToolResult> {
2415
// List the deployments
2516
const deployments = await client.listDeployments();
2617

src/tools/atlasLocal/tools.ts

Lines changed: 1 addition & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,33 +1,3 @@
11
import { ListDeploymentsTool } from "./read/listDeployments.js";
2-
import type AtlasLocal from "@mongodb-js-preview/atlas-local";
32

4-
// Don't use this directly, use BuildAtlasLocalTools instead
5-
const atlasLocalTools = [ListDeploymentsTool];
6-
7-
// Build the Atlas Local tools
8-
export const BuildAtlasLocalTools = async (): Promise<typeof atlasLocalTools> => {
9-
// Initialize the Atlas Local client
10-
const client = await GetAtlasLocalClient();
11-
12-
// If the client is found, set it on the tools
13-
// On unsupported platforms, the client will be undefined
14-
if (client) {
15-
// Set the client on the tools
16-
atlasLocalTools.forEach((tool) => {
17-
tool.prototype.client = client;
18-
});
19-
}
20-
21-
return atlasLocalTools;
22-
};
23-
24-
export const GetAtlasLocalClient = async (): Promise<AtlasLocal.Client | undefined> => {
25-
try {
26-
const { Client: AtlasLocalClient } = await import("@mongodb-js-preview/atlas-local");
27-
return AtlasLocalClient.connect();
28-
} catch (error) {
29-
// We only get here if the user is running atlas-local on a unsupported platform
30-
console.warn("Atlas Local native binding not available:", error);
31-
return undefined;
32-
}
33-
};
3+
export const AtlasLocalTools = [ListDeploymentsTool];

tests/integration/tools/atlas-local/listDeployments.test.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ import { describe, expect, it } from "vitest";
99

1010
const isMacOSInGitHubActions = process.platform === "darwin" && process.env.GITHUB_ACTIONS === "true";
1111

12+
// Docker is not available on macOS in GitHub Actions
13+
// That's why we skip the tests on macOS in GitHub Actions
1214
describe("atlas-local-list-deployments", () => {
1315
const integration = setupIntegrationTest(
1416
() => defaultTestConfig,

0 commit comments

Comments
 (0)