Skip to content

Commit bd94eae

Browse files
blvablva
authored
Merge branch 'main' into MCP-124 
2 parents c609434 + 94dfc08 commit bd94eae

File tree

13 files changed

+464
-95
lines changed

13 files changed

+464
-95
lines changed

‎src/common/connectionErrorHandler.ts‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
2+
import { ErrorCodes, type MongoDBError } from "./errors.js";
3+
import type { AnyConnectionState } from "./connectionManager.js";
4+
import type { ToolBase } from "../tools/tool.js";
5+
6+
export type ConnectionErrorHandler = (
7+
error: MongoDBError<ErrorCodes.NotConnectedToMongoDB | ErrorCodes.MisconfiguredConnectionString>,
8+
additionalContext: ConnectionErrorHandlerContext
9+
) => ConnectionErrorUnhandled | ConnectionErrorHandled;
10+
11+
export type ConnectionErrorHandlerContext = { availableTools: ToolBase[]; connectionState: AnyConnectionState };
12+
export type ConnectionErrorUnhandled = { errorHandled: false };
13+
export type ConnectionErrorHandled = { errorHandled: true; result: CallToolResult };
14+
15+
export const connectionErrorHandler: ConnectionErrorHandler = (error, { availableTools, connectionState }) => {
16+
const connectTools = availableTools
17+
.filter((t) => t.operationType === "connect")
18+
.sort((a, b) => a.category.localeCompare(b.category)); // Sort Atlas tools before MongoDB tools
19+
20+
// Find the first Atlas connect tool if available and suggest to the LLM to use it.
21+
// Note: if we ever have multiple Atlas connect tools, we may want to refine this logic to select the most appropriate one.
22+
const atlasConnectTool = connectTools?.find((t) => t.category === "atlas");
23+
const llmConnectHint = atlasConnectTool
24+
? `Note to LLM: prefer using the "${atlasConnectTool.name}" tool to connect to an Atlas cluster over using a connection string. Make sure to ask the user to specify a cluster name they want to connect to or ask them if they want to use the "list-clusters" tool to list all their clusters. Do not invent cluster names or connection strings unless the user has explicitly specified them. If they've previously connected to MongoDB using MCP, you can ask them if they want to reconnect using the same cluster/connection.`
25+
: "Note to LLM: do not invent connection strings and explicitly ask the user to provide one. If they have previously connected to MongoDB using MCP, you can ask them if they want to reconnect using the same connection string.";
26+
27+
const connectToolsNames = connectTools?.map((t) => `"${t.name}"`).join(", ");
28+
const additionalPromptForConnectivity: { type: "text"; text: string }[] = [];
29+
30+
if (connectionState.tag === "connecting" && connectionState.oidcConnectionType) {
31+
additionalPromptForConnectivity.push({
32+
type: "text",
33+
text: `The user needs to finish their OIDC connection by opening '${connectionState.oidcLoginUrl}' in the browser and use the following user code: '${connectionState.oidcUserCode}'`,
34+
});
35+
} else {
36+
additionalPromptForConnectivity.push({
37+
type: "text",
38+
text: connectToolsNames
39+
? `Please use one of the following tools: ${connectToolsNames} to connect to a MongoDB instance or update the MCP server configuration to include a connection string. ${llmConnectHint}`
40+
: "There are no tools available to connect. Please update the configuration to include a connection string and restart the server.",
41+
});
42+
}
43+
44+
switch (error.code) {
45+
case ErrorCodes.NotConnectedToMongoDB:
46+
return {
47+
errorHandled: true,
48+
result: {
49+
content: [
50+
{
51+
type: "text",
52+
text: "You need to connect to a MongoDB instance before you can access its data.",
53+
},
54+
...additionalPromptForConnectivity,
55+
],
56+
isError: true,
57+
},
58+
};
59+
case ErrorCodes.MisconfiguredConnectionString:
60+
return {
61+
errorHandled: true,
62+
result: {
63+
content: [
64+
{
65+
type: "text",
66+
text: "The configured connection string is not valid. Please check the connection string and confirm it points to a valid MongoDB instance.",
67+
},
68+
{
69+
type: "text",
70+
text: connectTools
71+
? `Alternatively, you can use one of the following tools: ${connectToolsNames} to connect to a MongoDB instance. ${llmConnectHint}`
72+
: "Please update the configuration to use a valid connection string and restart the server.",
73+
},
74+
],
75+
isError: true,
76+
},
77+
};
78+
79+
default:
80+
return { errorHandled: false };
81+
}
82+
};

‎src/common/errors.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,9 @@ export enum ErrorCodes {
44
ForbiddenCollscan = 1_000_002,
55
}
66

7-
export class MongoDBError extends Error {
7+
export class MongoDBError<ErrorCode extends ErrorCodes = ErrorCodes> extends Error {
88
constructor(
9-
public code: ErrorCodes,
9+
public code: ErrorCode,
1010
message: string
1111
) {
1212
super(message);

‎src/common/logger.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ export const LogId = {
4141

4242
mongodbConnectFailure: mongoLogId(1_004_001),
4343
mongodbDisconnectFailure: mongoLogId(1_004_002),
44+
mongodbConnectTry: mongoLogId(1_004_003),
4445

4546
toolUpdateFailure: mongoLogId(1_005_001),
4647
resourceUpdateFailure: mongoLogId(1_005_002),

‎src/index.ts‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,14 @@ async function main(): Promise<void> {
4949
assertHelpMode();
5050
assertVersionMode();
5151

52-
const transportRunner = config.transport === "stdio" ? new StdioRunner(config) : new StreamableHttpRunner(config);
52+
const transportRunner =
53+
config.transport === "stdio"
54+
? new StdioRunner({
55+
userConfig: config,
56+
})
57+
: new StreamableHttpRunner({
58+
userConfig: config,
59+
});
5360
const shutdown = (): void => {
5461
transportRunner.logger.info({
5562
id: LogId.serverCloseRequested,

‎src/lib.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,4 +11,11 @@ export {
1111
type ConnectionStateErrored,
1212
type ConnectionManagerFactoryFn,
1313
} from "./common/connectionManager.js";
14+
export type {
15+
ConnectionErrorHandler,
16+
ConnectionErrorHandled,
17+
ConnectionErrorUnhandled,
18+
ConnectionErrorHandlerContext,
19+
} from "./common/connectionErrorHandler.js";
20+
export { ErrorCodes } from "./common/errors.js";
1421
export { Telemetry } from "./telemetry/telemetry.js";

‎src/server.ts‎

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -20,12 +20,15 @@ import {
2020
import assert from "assert";
2121
import type { ToolBase } from "./tools/tool.js";
2222
import { validateConnectionString } from "./helpers/connectionOptions.js";
23+
import { packageInfo } from "./common/packageInfo.js";
24+
import { type ConnectionErrorHandler } from "./common/connectionErrorHandler.js";
2325

2426
export interface ServerOptions {
2527
session: Session;
2628
userConfig: UserConfig;
2729
mcpServer: McpServer;
2830
telemetry: Telemetry;
31+
connectionErrorHandler: ConnectionErrorHandler;
2932
}
3033

3134
export class Server {
@@ -34,6 +37,7 @@ export class Server {
3437
private readonly telemetry: Telemetry;
3538
public readonly userConfig: UserConfig;
3639
public readonly tools: ToolBase[] = [];
40+
public readonly connectionErrorHandler: ConnectionErrorHandler;
3741

3842
private _mcpLogLevel: LogLevel = "debug";
3943

@@ -44,12 +48,13 @@ export class Server {
4448
private readonly startTime: number;
4549
private readonly subscriptions = new Set<string>();
4650

47-
constructor({ session, mcpServer, userConfig, telemetry }: ServerOptions) {
51+
constructor({ session, mcpServer, userConfig, telemetry, connectionErrorHandler }: ServerOptions) {
4852
this.startTime = Date.now();
4953
this.session = session;
5054
this.telemetry = telemetry;
5155
this.mcpServer = mcpServer;
5256
this.userConfig = userConfig;
57+
this.connectionErrorHandler = connectionErrorHandler;
5358
}
5459

5560
async connect(transport: Transport): Promise<void> {
@@ -119,11 +124,10 @@ export class Server {
119124
this.session.setMcpClient(this.mcpServer.server.getClientVersion());
120125
// Placed here to start the connection to the config connection string as soon as the server is initialized.
121126
void this.connectToConfigConnectionString();
122-
123127
this.session.logger.info({
124128
id: LogId.serverInitialized,
125129
context: "server",
126-
message: `Server started with transport ${transport.constructor.name} and agent runner ${this.session.mcpClient?.name}`,
130+
message: `Server with version ${packageInfo.version} started with transport ${transport.constructor.name} and agent runner ${JSON.stringify(this.session.mcpClient)}`,
127131
});
128132

129133
this.emitServerEvent("start", Date.now() - this.startTime);
@@ -244,15 +248,21 @@ export class Server {
244248
private async connectToConfigConnectionString(): Promise<void> {
245249
if (this.userConfig.connectionString) {
246250
try {
251+
this.session.logger.info({
252+
id: LogId.mongodbConnectTry,
253+
context: "server",
254+
message: `Detected a MongoDB connection string in the configuration, trying to connect...`,
255+
});
247256
await this.session.connectToMongoDB({
248257
connectionString: this.userConfig.connectionString,
249258
});
250259
} catch (error) {
251-
console.error(
252-
"Failed to connect to MongoDB instance using the connection string from the config: ",
253-
error
254-
);
255-
throw new Error("Failed to connect to MongoDB instance using the connection string from the config");
260+
// We don't throw an error here because we want to allow the server to start even if the connection string is invalid.
261+
this.session.logger.error({
262+
id: LogId.mongodbConnectFailure,
263+
context: "server",
264+
message: `Failed to connect to MongoDB instance using the connection string from the config: ${error instanceof Error ? error.message : String(error)}`,
265+
});
256266
}
257267
}
258268
}

‎src/tools/mongodb/mongodbTool.ts‎

Lines changed: 14 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -56,63 +56,22 @@ export abstract class MongoDBToolBase extends ToolBase {
5656
args: ToolArgs<typeof this.argsShape>
5757
): Promise<CallToolResult> | CallToolResult {
5858
if (error instanceof MongoDBError) {
59-
const connectTools = this.server?.tools
60-
.filter((t) => t.operationType === "connect")
61-
.sort((a, b) => a.category.localeCompare(b.category)); // Sort Altas tools before MongoDB tools
62-
63-
// Find the first Atlas connect tool if available and suggest to the LLM to use it.
64-
// Note: if we ever have multiple Atlas connect tools, we may want to refine this logic to select the most appropriate one.
65-
const atlasConnectTool = connectTools?.find((t) => t.category === "atlas");
66-
const llmConnectHint = atlasConnectTool
67-
? `Note to LLM: prefer using the "${atlasConnectTool.name}" tool to connect to an Atlas cluster over using a connection string. Make sure to ask the user to specify a cluster name they want to connect to or ask them if they want to use the "list-clusters" tool to list all their clusters. Do not invent cluster names or connection strings unless the user has explicitly specified them. If they've previously connected to MongoDB using MCP, you can ask them if they want to reconnect using the same cluster/connection.`
68-
: "Note to LLM: do not invent connection strings and explicitly ask the user to provide one. If they have previously connected to MongoDB using MCP, you can ask them if they want to reconnect using the same connection string.";
69-
70-
const connectToolsNames = connectTools?.map((t) => `"${t.name}"`).join(", ");
71-
const connectionStatus = this.session.connectionManager.currentConnectionState;
72-
const additionalPromptForConnectivity: { type: "text"; text: string }[] = [];
73-
74-
if (connectionStatus.tag === "connecting" && connectionStatus.oidcConnectionType) {
75-
additionalPromptForConnectivity.push({
76-
type: "text",
77-
text: `The user needs to finish their OIDC connection by opening '${connectionStatus.oidcLoginUrl}' in the browser and use the following user code: '${connectionStatus.oidcUserCode}'`,
78-
});
79-
} else {
80-
additionalPromptForConnectivity.push({
81-
type: "text",
82-
text: connectToolsNames
83-
? `Please use one of the following tools: ${connectToolsNames} to connect to a MongoDB instance or update the MCP server configuration to include a connection string. ${llmConnectHint}`
84-
: "There are no tools available to connect. Please update the configuration to include a connection string and restart the server.",
85-
});
86-
}
87-
8859
switch (error.code) {
8960
case ErrorCodes.NotConnectedToMongoDB:
90-
return {
91-
content: [
92-
{
93-
type: "text",
94-
text: "You need to connect to a MongoDB instance before you can access its data.",
95-
},
96-
...additionalPromptForConnectivity,
97-
],
98-
isError: true,
99-
};
100-
case ErrorCodes.MisconfiguredConnectionString:
101-
return {
102-
content: [
103-
{
104-
type: "text",
105-
text: "The configured connection string is not valid. Please check the connection string and confirm it points to a valid MongoDB instance.",
106-
},
107-
{
108-
type: "text",
109-
text: connectTools
110-
? `Alternatively, you can use one of the following tools: ${connectToolsNames} to connect to a MongoDB instance. ${llmConnectHint}`
111-
: "Please update the configuration to use a valid connection string and restart the server.",
112-
},
113-
],
114-
isError: true,
115-
};
61+
case ErrorCodes.MisconfiguredConnectionString: {
62+
const connectionError = error as MongoDBError<
63+
ErrorCodes.NotConnectedToMongoDB | ErrorCodes.MisconfiguredConnectionString
64+
>;
65+
const outcome = this.server?.connectionErrorHandler(connectionError, {
66+
availableTools: this.server?.tools ?? [],
67+
connectionState: this.session.connectionManager.currentConnectionState,
68+
});
69+
if (outcome?.errorHandled) {
70+
return outcome.result;
71+
}
72+
73+
return super.handleError(error, args);
74+
}
11675
case ErrorCodes.ForbiddenCollscan:
11776
return {
11877
content: [

‎src/transports/base.ts‎

Lines changed: 25 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,17 +8,35 @@ import type { LoggerBase } from "../common/logger.js";
88
import { CompositeLogger, ConsoleLogger, DiskLogger, McpLogger } from "../common/logger.js";
99
import { ExportsManager } from "../common/exportsManager.js";
1010
import { DeviceId } from "../helpers/deviceId.js";
11-
import { type ConnectionManagerFactoryFn } from "../common/connectionManager.js";
11+
import { createMCPConnectionManager, type ConnectionManagerFactoryFn } from "../common/connectionManager.js";
12+
import {
13+
type ConnectionErrorHandler,
14+
connectionErrorHandler as defaultConnectionErrorHandler,
15+
} from "../common/connectionErrorHandler.js";
16+
17+
export type TransportRunnerConfig = {
18+
userConfig: UserConfig;
19+
createConnectionManager?: ConnectionManagerFactoryFn;
20+
connectionErrorHandler?: ConnectionErrorHandler;
21+
additionalLoggers?: LoggerBase[];
22+
};
1223

1324
export abstract class TransportRunnerBase {
1425
public logger: LoggerBase;
1526
public deviceId: DeviceId;
27+
protected readonly userConfig: UserConfig;
28+
private readonly createConnectionManager: ConnectionManagerFactoryFn;
29+
private readonly connectionErrorHandler: ConnectionErrorHandler;
1630

17-
protected constructor(
18-
protected readonly userConfig: UserConfig,
19-
private readonly createConnectionManager: ConnectionManagerFactoryFn,
20-
additionalLoggers: LoggerBase[]
21-
) {
31+
protected constructor({
32+
userConfig,
33+
createConnectionManager = createMCPConnectionManager,
34+
connectionErrorHandler = defaultConnectionErrorHandler,
35+
additionalLoggers = [],
36+
}: TransportRunnerConfig) {
37+
this.userConfig = userConfig;
38+
this.createConnectionManager = createConnectionManager;
39+
this.connectionErrorHandler = connectionErrorHandler;
2240
const loggers: LoggerBase[] = [...additionalLoggers];
2341
if (this.userConfig.loggers.includes("stderr")) {
2442
loggers.push(new ConsoleLogger());
@@ -68,6 +86,7 @@ export abstract class TransportRunnerBase {
6886
session,
6987
telemetry,
7088
userConfig: this.userConfig,
89+
connectionErrorHandler: this.connectionErrorHandler,
7190
});
7291

7392
// We need to create the MCP logger after the server is constructed

‎src/transports/stdio.ts‎

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,9 @@ import { EJSON } from "bson";
22
import type { JSONRPCMessage } from "@modelcontextprotocol/sdk/types.js";
33
import { JSONRPCMessageSchema } from "@modelcontextprotocol/sdk/types.js";
44
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
5-
import { type LoggerBase, LogId } from "../common/logger.js";
5+
import { LogId } from "../common/logger.js";
66
import type { Server } from "../server.js";
7-
import { TransportRunnerBase } from "./base.js";
8-
import { type UserConfig } from "../common/config.js";
9-
import { createMCPConnectionManager, type ConnectionManagerFactoryFn } from "../common/connectionManager.js";
7+
import { TransportRunnerBase, type TransportRunnerConfig } from "./base.js";
108

119
// This is almost a copy of ReadBuffer from @modelcontextprotocol/sdk
1210
// but it uses EJSON.parse instead of JSON.parse to handle BSON types
@@ -55,12 +53,8 @@ export function createStdioTransport(): StdioServerTransport {
5553
export class StdioRunner extends TransportRunnerBase {
5654
private server: Server | undefined;
5755

58-
constructor(
59-
userConfig: UserConfig,
60-
createConnectionManager: ConnectionManagerFactoryFn = createMCPConnectionManager,
61-
additionalLoggers: LoggerBase[] = []
62-
) {
63-
super(userConfig, createConnectionManager, additionalLoggers);
56+
constructor(config: TransportRunnerConfig) {
57+
super(config);
6458
}
6559

6660
async start(): Promise<void> {

0 commit comments

Comments
 (0)