connect.ts•4.48 kB
import { z } from "zod";
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
import { MongoDBToolBase } from "../mongodbTool.js";
import type { ToolArgs, OperationType, ToolConstructorParams } from "../../tool.js";
import assert from "assert";
import type { Server } from "../../../server.js";
import { LogId } from "../../../common/logger.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. The config resource captures if the server is already connected to a MongoDB cluster. 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 MongoDB cluster.";
export class ConnectTool extends MongoDBToolBase {
public 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 = {
connectionString: z.string().optional(),
};
public operationType: OperationType = "connect";
constructor({ session, config, telemetry, elicitation }: ToolConstructorParams) {
super({ session, config, telemetry, elicitation });
session.on("connect", () => {
this.updateMetadata();
});
session.on("disconnect", () => {
this.updateMetadata();
});
}
protected async execute({ connectionString }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
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;
}
await this.connectToMongoDB(connectionString);
this.updateMetadata();
return {
content: [{ type: "text", text: "Successfully connected to MongoDB." }],
};
}
public register(server: Server): boolean {
if (super.register(server)) {
this.updateMetadata();
return true;
}
return false;
}
private updateMetadata(): void {
let name: string;
let description: string;
let inputSchema: z.ZodObject<z.ZodRawShape>;
if (this.session.isConnectedToMongoDB) {
name = connectedName;
description = connectedDescription;
inputSchema = connectedSchema;
} else {
name = disconnectedName;
description = disconnectedDescription;
inputSchema = disconnectedSchema;
}
this.session.logger.info({
id: LogId.updateToolMetadata,
context: "tool",
message: `Updating tool metadata to ${name}`,
});
this.update?.({
name,
description,
inputSchema,
});
}
}