Skip to content

Commit 0e79eab

Browse files
nirinchevnirinchev
committed
feat: update the connect tool based on connectivity status
1 parent 77513d9 commit 0e79eab

File tree

8 files changed

+277
-221
lines changed

8 files changed

+277
-221
lines changed

‎src/server.ts‎

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ export class Server {
2121
public readonly session: Session;
2222
private readonly mcpServer: McpServer;
2323
private readonly telemetry: Telemetry;
24-
private readonly userConfig: UserConfig;
24+
public readonly userConfig: UserConfig;
2525

2626
constructor({ session, mcpServer, userConfig }: ServerOptions) {
2727
this.session = session;
@@ -32,6 +32,7 @@ export class Server {
3232

3333
async connect(transport: Transport) {
3434
this.mcpServer.server.registerCapabilities({ logging: {} });
35+
this.mcpServer.sendToolListChanged;
3536
this.registerTools();
3637
this.registerResources();
3738

@@ -86,6 +87,32 @@ export class Server {
8687
}
8788

8889
private registerResources() {
90+
this.mcpServer.resource(
91+
"config",
92+
"config://config",
93+
{
94+
description:
95+
"Server configuration, supplied by the user either as environment variables or as startup arguments",
96+
},
97+
(uri) => {
98+
const result = {
99+
telemetry: this.userConfig.telemetry,
100+
logPath: this.userConfig.logPath,
101+
connectionString: this.userConfig.connectionString
102+
? "set; no explicit connect needed, use switch-connection tool to connect to a different connection if necessary"
103+
: "not set; before using any mongodb tool, you need to call the connect tool with a connection string",
104+
connectOptions: this.userConfig.connectOptions,
105+
};
106+
return {
107+
contents: [
108+
{
109+
text: JSON.stringify(result),
110+
uri: uri.href,
111+
},
112+
],
113+
};
114+
}
115+
);
89116
if (this.userConfig.connectionString) {
90117
this.mcpServer.resource(
91118
"connection-string",

‎src/session.ts‎

Lines changed: 8 additions & 1 deletion
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/tools/mongodb/metadata/connect.ts‎

Lines changed: 79 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -1,93 +1,96 @@
1-
import { z } from "zod";
1+
import { AnyZodObject, z, ZodRawShape } 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
}

‎src/tools/tool.ts‎

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
import { z, type ZodRawShape, type ZodNever } from "zod";
2-
import type { McpServer, ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js";
1+
import { z, type ZodRawShape, type ZodNever, AnyZodObject } from "zod";
2+
import type { McpServer, RegisteredTool, ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js";
33
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
44
import { Session } from "../session.js";
55
import logger from "../logger.js";
@@ -81,8 +81,36 @@ export abstract class ToolBase {
8181
};
8282

8383
server.tool(this.name, this.description, this.argsShape, callback);
84+
85+
// This is very similar to RegisteredTool.update, but without the bugs around the name.
86+
// In the upstream update method, the name is captured in the closure and not updated when
87+
// the tool name changes. This means that you only get one name update before things end up
88+
// in a broken state.
89+
this.update = (updates: { name?: string; description?: string; inputSchema?: AnyZodObject }) => {
90+
const tools = server["_registeredTools"] as { [toolName: string]: RegisteredTool };
91+
const existingTool = tools[this.name];
92+
93+
if (updates.name && updates.name !== this.name) {
94+
delete tools[this.name];
95+
this.name = updates.name;
96+
tools[this.name] = existingTool;
97+
}
98+
99+
if (updates.description) {
100+
existingTool.description = updates.description;
101+
this.description = updates.description;
102+
}
103+
104+
if (updates.inputSchema) {
105+
existingTool.inputSchema = updates.inputSchema;
106+
}
107+
108+
server.sendToolListChanged();
109+
};
84110
}
85111

112+
protected update?: (updates: { name?: string; description?: string; inputSchema?: AnyZodObject }) => void;
113+
86114
// Checks if a tool is allowed to run based on the config
87115
protected verifyAllowed(): boolean {
88116
let errorClarification: string | undefined;

‎tests/integration/helpers.ts‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import { InMemoryTransport } from "./inMemoryTransport.js";
33
import { Server } from "../../src/server.js";
44
import { ObjectId } from "mongodb";
55
import { config, UserConfig } from "../../src/config.js";
6-
import { McpError } from "@modelcontextprotocol/sdk/types.js";
6+
import { McpError, ToolListChangedNotificationSchema } from "@modelcontextprotocol/sdk/types.js";
77
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
88
import { Session } from "../../src/session.js";
99
import { toIncludeAllMembers } from "jest-extended";
@@ -22,13 +22,14 @@ export interface IntegrationTest {
2222
mcpServer: () => Server;
2323
}
2424

25-
export function setupIntegrationTest(userConfig: UserConfig = config): IntegrationTest {
25+
export function setupIntegrationTest(userConfigGetter: () => UserConfig = () => config): IntegrationTest {
2626
let mcpClient: Client | undefined;
2727
let mcpServer: Server | undefined;
2828

2929
let randomDbName: string;
3030

3131
beforeAll(async () => {
32+
const userConfig = userConfigGetter();
3233
const clientTransport = new InMemoryTransport();
3334
const serverTransport = new InMemoryTransport();
3435

@@ -54,6 +55,7 @@ export function setupIntegrationTest(userConfig: UserConfig = config): Integrati
5455
apiClientSecret: userConfig.apiClientSecret,
5556
});
5657

58+
userConfig.telemetry = "disabled";
5759
mcpServer = new Server({
5860
session,
5961
userConfig,
@@ -67,10 +69,19 @@ export function setupIntegrationTest(userConfig: UserConfig = config): Integrati
6769
});
6870

6971
beforeEach(async () => {
70-
config.telemetry = "disabled";
72+
if (mcpServer) {
73+
mcpServer.userConfig.telemetry = "disabled";
74+
}
7175
randomDbName = new ObjectId().toString();
7276
});
7377

78+
afterEach(async () => {
79+
if (mcpServer) {
80+
await mcpServer.session.close();
81+
mcpServer.userConfig.connectionString = undefined;
82+
}
83+
});
84+
7485
afterAll(async () => {
7586
await mcpClient?.close();
7687
mcpClient = undefined;

‎tests/integration/server.test.ts‎

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,11 @@ import { config } from "../../src/config.js";
33

44
describe("Server integration test", () => {
55
describe("without atlas", () => {
6-
const integration = setupIntegrationTest({
6+
const integration = setupIntegrationTest(() => ({
77
...config,
88
apiClientId: undefined,
99
apiClientSecret: undefined,
10-
});
10+
}));
1111

1212
it("should return positive number of tools and have no atlas tools", async () => {
1313
const tools = await integration.mcpClient().listTools();
@@ -19,11 +19,11 @@ describe("Server integration test", () => {
1919
});
2020
});
2121
describe("with atlas", () => {
22-
const integration = setupIntegrationTest({
22+
const integration = setupIntegrationTest(() => ({
2323
...config,
2424
apiClientId: "test",
2525
apiClientSecret: "test",
26-
});
26+
}));
2727

2828
describe("list capabilities", () => {
2929
it("should return positive number of tools and have some atlas tools", async () => {
@@ -35,12 +35,6 @@ describe("Server integration test", () => {
3535
expect(atlasTools.length).toBeGreaterThan(0);
3636
});
3737

38-
it("should return no resources", async () => {
39-
await expect(() => integration.mcpClient().listResources()).rejects.toMatchObject({
40-
message: "MCP error -32601: Method not found",
41-
});
42-
});
43-
4438
it("should return no prompts", async () => {
4539
await expect(() => integration.mcpClient().listPrompts()).rejects.toMatchObject({
4640
message: "MCP error -32601: Method not found",

0 commit comments

Comments
 (0)