Skip to content

Commit a1a36fb

Browse files
nirinchevnirinchev
authored
feat: update the connect tool based on connectivity status (#118)
1 parent 8423389 commit a1a36fb

File tree

13 files changed

+337
-255
lines changed

13 files changed

+337
-255
lines changed

‎src/errors.ts

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
export enum ErrorCodes {
22
NotConnectedToMongoDB = 1_000_000,
3-
InvalidParams = 1_000_001,
3+
MisconfiguredConnectionString = 1_000_001,
44
}
55

66
export class MongoDBError extends Error {

‎src/index.ts

+2-3
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,7 @@
11
#!/usr/bin/env node
22

33
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
4-
import logger from "./logger.js";
5-
import { mongoLogId } from "mongodb-log-writer";
4+
import logger, { LogId } from "./logger.js";
65
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
76
import { config } from "./config.js";
87
import { Session } from "./session.js";
@@ -29,6 +28,6 @@ try {
2928

3029
await server.connect(transport);
3130
} catch (error: unknown) {
32-
logger.emergency(mongoLogId(1_000_004), "server", `Fatal error running server: ${error as string}`);
31+
logger.emergency(LogId.serverStartFailure, "server", `Fatal error running server: ${error as string}`);
3332
process.exit(1);
3433
}

‎src/logger.ts

+19-1
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,29 @@
11
import fs from "fs/promises";
2-
import { MongoLogId, MongoLogManager, MongoLogWriter } from "mongodb-log-writer";
2+
import { mongoLogId, MongoLogId, MongoLogManager, MongoLogWriter } from "mongodb-log-writer";
33
import redact from "mongodb-redact";
44
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
55
import { LoggingMessageNotification } from "@modelcontextprotocol/sdk/types.js";
66

77
export type LogLevel = LoggingMessageNotification["params"]["level"];
88

9+
export const LogId = {
10+
serverStartFailure: mongoLogId(1_000_001),
11+
serverInitialized: mongoLogId(1_000_002),
12+
13+
atlasCheckCredentials: mongoLogId(1_001_001),
14+
15+
telemetryDisabled: mongoLogId(1_002_001),
16+
telemetryEmitFailure: mongoLogId(1_002_002),
17+
telemetryEmitStart: mongoLogId(1_002_003),
18+
telemetryEmitSuccess: mongoLogId(1_002_004),
19+
20+
toolExecute: mongoLogId(1_003_001),
21+
toolExecuteFailure: mongoLogId(1_003_002),
22+
toolDisabled: mongoLogId(1_003_003),
23+
24+
mongodbConnectFailure: mongoLogId(1_004_001),
25+
} as const;
26+
927
abstract class LoggerBase {
1028
abstract log(level: LogLevel, id: MongoLogId, context: string, message: string): void;
1129

‎src/server.ts

+29-4
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,7 @@ import { Session } from "./session.js";
33
import { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
44
import { AtlasTools } from "./tools/atlas/tools.js";
55
import { MongoDbTools } from "./tools/mongodb/tools.js";
6-
import logger, { initializeLogger } from "./logger.js";
7-
import { mongoLogId } from "mongodb-log-writer";
6+
import logger, { initializeLogger, LogId } from "./logger.js";
87
import { ObjectId } from "mongodb";
98
import { Telemetry } from "./telemetry/telemetry.js";
109
import { UserConfig } from "./config.js";
@@ -23,7 +22,7 @@ export class Server {
2322
public readonly session: Session;
2423
private readonly mcpServer: McpServer;
2524
private readonly telemetry: Telemetry;
26-
private readonly userConfig: UserConfig;
25+
public readonly userConfig: UserConfig;
2726
private readonly startTime: number;
2827

2928
constructor({ session, mcpServer, userConfig }: ServerOptions) {
@@ -71,7 +70,7 @@ export class Server {
7170
this.session.sessionId = new ObjectId().toString();
7271

7372
logger.info(
74-
mongoLogId(1_000_004),
73+
LogId.serverInitialized,
7574
"server",
7675
`Server started with transport ${transport.constructor.name} and agent runner ${this.session.agentRunner?.name}`
7776
);
@@ -135,6 +134,32 @@ export class Server {
135134
}
136135

137136
private registerResources() {
137+
this.mcpServer.resource(
138+
"config",
139+
"config://config",
140+
{
141+
description:
142+
"Server configuration, supplied by the user either as environment variables or as startup arguments",
143+
},
144+
(uri) => {
145+
const result = {
146+
telemetry: this.userConfig.telemetry,
147+
logPath: this.userConfig.logPath,
148+
connectionString: this.userConfig.connectionString
149+
? "set; no explicit connect needed, use switch-connection tool to connect to a different connection if necessary"
150+
: "not set; before using any mongodb tool, you need to call the connect tool with a connection string",
151+
connectOptions: this.userConfig.connectOptions,
152+
};
153+
return {
154+
contents: [
155+
{
156+
text: JSON.stringify(result),
157+
uri: uri.href,
158+
},
159+
],
160+
};
161+
}
162+
);
138163
if (this.userConfig.connectionString) {
139164
this.mcpServer.resource(
140165
"connection-string",

‎src/session.ts

+8-1
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,17 @@
11
import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
22
import { ApiClient, ApiClientCredentials } from "./common/atlas/apiClient.js";
33
import { Implementation } from "@modelcontextprotocol/sdk/types.js";
4+
import EventEmitter from "events";
45

56
export interface SessionOptions {
67
apiBaseUrl?: string;
78
apiClientId?: string;
89
apiClientSecret?: string;
910
}
1011

11-
export class Session {
12+
export class Session extends EventEmitter<{
13+
close: [];
14+
}> {
1215
sessionId?: string;
1316
serviceProvider?: NodeDriverServiceProvider;
1417
apiClient: ApiClient;
@@ -18,6 +21,8 @@ export class Session {
1821
};
1922

2023
constructor({ apiBaseUrl, apiClientId, apiClientSecret }: SessionOptions = {}) {
24+
super();
25+
2126
const credentials: ApiClientCredentials | undefined =
2227
apiClientId && apiClientSecret
2328
? {
@@ -49,6 +54,8 @@ export class Session {
4954
console.error("Error closing service provider:", error);
5055
}
5156
this.serviceProvider = undefined;
57+
58+
this.emit("close");
5259
}
5360
}
5461
}

‎src/telemetry/telemetry.ts

+6-7
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,7 @@
11
import { Session } from "../session.js";
22
import { BaseEvent, CommonProperties } from "./types.js";
33
import { config } from "../config.js";
4-
import logger from "../logger.js";
5-
import { mongoLogId } from "mongodb-log-writer";
4+
import logger, { LogId } from "../logger.js";
65
import { ApiClient } from "../common/atlas/apiClient.js";
76
import { MACHINE_METADATA } from "./constants.js";
87
import { EventCache } from "./eventCache.js";
@@ -61,7 +60,7 @@ export class Telemetry {
6160

6261
await this.emit(events);
6362
} catch {
64-
logger.debug(mongoLogId(1_000_002), "telemetry", `Error emitting telemetry events.`);
63+
logger.debug(LogId.telemetryEmitFailure, "telemetry", `Error emitting telemetry events.`);
6564
}
6665
}
6766

@@ -89,20 +88,20 @@ export class Telemetry {
8988
const allEvents = [...cachedEvents, ...events];
9089

9190
logger.debug(
92-
mongoLogId(1_000_003),
91+
LogId.telemetryEmitStart,
9392
"telemetry",
9493
`Attempting to send ${allEvents.length} events (${cachedEvents.length} cached)`
9594
);
9695

9796
const result = await this.sendEvents(this.session.apiClient, allEvents);
9897
if (result.success) {
9998
this.eventCache.clearEvents();
100-
logger.debug(mongoLogId(1_000_004), "telemetry", `Sent ${allEvents.length} events successfully`);
99+
logger.debug(LogId.telemetryEmitSuccess, "telemetry", `Sent ${allEvents.length} events successfully`);
101100
return;
102101
}
103102

104-
logger.warning(
105-
mongoLogId(1_000_005),
103+
logger.debug(
104+
LogId.telemetryEmitFailure,
106105
"telemetry",
107106
`Error sending event to client: ${result.error instanceof Error ? result.error.message : String(result.error)}`
108107
);

‎src/tools/mongodb/metadata/connect.ts

+78-75
Original file line numberDiff line numberDiff line change
@@ -2,92 +2,95 @@ import { z } from "zod";
22
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
33
import { MongoDBToolBase } from "../mongodbTool.js";
44
import { ToolArgs, OperationType } from "../../tool.js";
5-
import { MongoError as DriverError } from "mongodb";
5+
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
6+
import assert from "assert";
7+
import { UserConfig } from "../../../config.js";
8+
import { Telemetry } from "../../../telemetry/telemetry.js";
9+
import { Session } from "../../../session.js";
10+
11+
const disconnectedSchema = z
12+
.object({
13+
connectionString: z.string().describe("MongoDB connection string (in the mongodb:// or mongodb+srv:// format)"),
14+
})
15+
.describe("Options for connecting to MongoDB.");
16+
17+
const connectedSchema = z
18+
.object({
19+
connectionString: z
20+
.string()
21+
.optional()
22+
.describe("MongoDB connection string to switch to (in the mongodb:// or mongodb+srv:// format)"),
23+
})
24+
.describe(
25+
"Options for switching the current MongoDB connection. If a connection string is not provided, the connection string from the config will be used."
26+
);
27+
28+
const connectedName = "switch-connection" as const;
29+
const disconnectedName = "connect" as const;
30+
31+
const connectedDescription =
32+
"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.";
33+
const disconnectedDescription = "Connect to a MongoDB instance";
634

735
export class ConnectTool extends MongoDBToolBase {
8-
protected name = "connect";
9-
protected description = "Connect to a MongoDB instance";
36+
protected name: typeof connectedName | typeof disconnectedName = disconnectedName;
37+
protected description: typeof connectedDescription | typeof disconnectedDescription = disconnectedDescription;
38+
39+
// Here the default is empty just to trigger registration, but we're going to override it with the correct
40+
// schema in the register method.
1041
protected argsShape = {
11-
options: z
12-
.array(
13-
z
14-
.union([
15-
z.object({
16-
connectionString: z
17-
.string()
18-
.describe("MongoDB connection string (in the mongodb:// or mongodb+srv:// format)"),
19-
}),
20-
z.object({
21-
clusterName: z.string().describe("MongoDB cluster name"),
22-
}),
23-
])
24-
.optional()
25-
)
26-
.optional()
27-
.describe(
28-
"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."
29-
),
42+
connectionString: z.string().optional(),
3043
};
3144

3245
protected operationType: OperationType = "metadata";
3346

34-
protected async execute({ options: optionsArr }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
35-
const options = optionsArr?.[0];
36-
let connectionString: string;
37-
if (!options && !this.config.connectionString) {
38-
return {
39-
content: [
40-
{ type: "text", text: "No connection details provided." },
41-
{ type: "text", text: "Please provide either a connection string or a cluster name" },
42-
],
43-
};
44-
}
47+
constructor(session: Session, config: UserConfig, telemetry: Telemetry) {
48+
super(session, config, telemetry);
49+
session.on("close", () => {
50+
this.updateMetadata();
51+
});
52+
}
4553

46-
if (!options) {
47-
// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
48-
connectionString = this.config.connectionString!;
49-
} else if ("connectionString" in options) {
50-
connectionString = options.connectionString;
51-
} else {
52-
// TODO: https://github.com/mongodb-js/mongodb-mcp-server/issues/19
53-
// We don't support connecting via cluster name since we'd need to obtain the user credentials
54-
// and fill in the connection string.
55-
return {
56-
content: [
57-
{
58-
type: "text",
59-
text: `Connecting via cluster name not supported yet. Please provide a connection string.`,
60-
},
61-
],
62-
};
54+
protected async execute({ connectionString }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
55+
switch (this.name) {
56+
case disconnectedName:
57+
assert(connectionString, "Connection string is required");
58+
break;
59+
case connectedName:
60+
connectionString ??= this.config.connectionString;
61+
assert(
62+
connectionString,
63+
"Cannot switch to a new connection because no connection string was provided and no default connection string is configured."
64+
);
65+
break;
6366
}
6467

65-
try {
66-
await this.connectToMongoDB(connectionString);
67-
return {
68-
content: [{ type: "text", text: `Successfully connected to ${connectionString}.` }],
69-
};
70-
} catch (error) {
71-
// Sometimes the model will supply an incorrect connection string. If the user has configured
72-
// a different one as environment variable or a cli argument, suggest using that one instead.
73-
if (
74-
this.config.connectionString &&
75-
error instanceof DriverError &&
76-
this.config.connectionString !== connectionString
77-
) {
78-
return {
79-
content: [
80-
{
81-
type: "text",
82-
text:
83-
`Failed to connect to MongoDB at '${connectionString}' due to error: '${error.message}.` +
84-
`Your config lists a different connection string: '${this.config.connectionString}' - do you want to try connecting to it instead?`,
85-
},
86-
],
87-
};
88-
}
68+
await this.connectToMongoDB(connectionString);
69+
this.updateMetadata();
70+
return {
71+
content: [{ type: "text", text: "Successfully connected to MongoDB." }],
72+
};
73+
}
74+
75+
public register(server: McpServer): void {
76+
super.register(server);
8977

90-
throw error;
78+
this.updateMetadata();
79+
}
80+
81+
private updateMetadata(): void {
82+
if (this.config.connectionString || this.session.serviceProvider) {
83+
this.update?.({
84+
name: connectedName,
85+
description: connectedDescription,
86+
inputSchema: connectedSchema,
87+
});
88+
} else {
89+
this.update?.({
90+
name: disconnectedName,
91+
description: disconnectedDescription,
92+
inputSchema: disconnectedSchema,
93+
});
9194
}
9295
}
9396
}

0 commit comments

Comments
 (0)