Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

Commit 94dfc08

Browse files
chore: extend library interfaces to allow injecting a custom connection error handler MCP-132 (#502)
1 parent 47c0a09 commit 94dfc08

File tree

12 files changed

+450
-88
lines changed

12 files changed

+450
-88
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<ErrorCodeextendsErrorCodes=ErrorCodes> extends Error {
88
constructor(
9-
public code: ErrorCodes,
9+
public code: ErrorCode,
1010
message: string
1111
) {
1212
super(message);

‎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: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,12 +21,14 @@ import assert from "assert";
2121
import type { ToolBase } from "./tools/tool.js";
2222
import { validateConnectionString } from "./helpers/connectionOptions.js";
2323
import { packageInfo } from "./common/packageInfo.js";
24+
import { type ConnectionErrorHandler } from "./common/connectionErrorHandler.js";
2425

2526
export interface ServerOptions {
2627
session: Session;
2728
userConfig: UserConfig;
2829
mcpServer: McpServer;
2930
telemetry: Telemetry;
31+
connectionErrorHandler: ConnectionErrorHandler;
3032
}
3133

3234
export class Server {
@@ -35,6 +37,7 @@ export class Server {
3537
private readonly telemetry: Telemetry;
3638
public readonly userConfig: UserConfig;
3739
public readonly tools: ToolBase[] = [];
40+
public readonly connectionErrorHandler: ConnectionErrorHandler;
3841

3942
private _mcpLogLevel: LogLevel = "debug";
4043

@@ -45,12 +48,13 @@ export class Server {
4548
private readonly startTime: number;
4649
private readonly subscriptions = new Set<string>();
4750

48-
constructor({ session, mcpServer, userConfig, telemetry }: ServerOptions) {
51+
constructor({ session, mcpServer, userConfig, telemetry, connectionErrorHandler }: ServerOptions) {
4952
this.startTime = Date.now();
5053
this.session = session;
5154
this.telemetry = telemetry;
5255
this.mcpServer = mcpServer;
5356
this.userConfig = userConfig;
57+
this.connectionErrorHandler = connectionErrorHandler;
5458
}
5559

5660
async connect(transport: Transport): Promise<void> {

‎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 { typeLoggerBase,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> {

‎src/transports/streamableHttp.ts‎

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,9 @@ import type http from "http";
33
import { randomUUID } from "crypto";
44
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
55
import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
6-
import { LogId, type LoggerBase } from "../common/logger.js";
7-
import { type UserConfig } from "../common/config.js";
6+
import { LogId } from "../common/logger.js";
87
import { SessionStore } from "../common/sessionStore.js";
9-
import { TransportRunnerBase } from "./base.js";
10-
import { createMCPConnectionManager, type ConnectionManagerFactoryFn } from "../common/connectionManager.js";
8+
import { TransportRunnerBase, type TransportRunnerConfig } from "./base.js";
119

1210
const JSON_RPC_ERROR_CODE_PROCESSING_REQUEST_FAILED = -32000;
1311
const JSON_RPC_ERROR_CODE_SESSION_ID_REQUIRED = -32001;
@@ -19,12 +17,8 @@ export class StreamableHttpRunner extends TransportRunnerBase {
1917
private httpServer: http.Server | undefined;
2018
private sessionStore!: SessionStore;
2119

22-
constructor(
23-
userConfig: UserConfig,
24-
createConnectionManager: ConnectionManagerFactoryFn = createMCPConnectionManager,
25-
additionalLoggers: LoggerBase[] = []
26-
) {
27-
super(userConfig, createConnectionManager, additionalLoggers);
20+
constructor(config: TransportRunnerConfig) {
21+
super(config);
2822
}
2923

3024
public get serverAddress(): string {

‎tests/integration/helpers.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "vitest
1313
import type { ConnectionManager, ConnectionState } from "../../src/common/connectionManager.js";
1414
import { MCPConnectionManager } from "../../src/common/connectionManager.js";
1515
import { DeviceId } from "../../src/helpers/deviceId.js";
16+
import { connectionErrorHandler } from "../../src/common/connectionErrorHandler.js";
1617

1718
interface ParameterInfo {
1819
name: string;
@@ -101,6 +102,7 @@ export function setupIntegrationTest(
101102
name: "test-server",
102103
version: "5.2.3",
103104
}),
105+
connectionErrorHandler,
104106
});
105107

106108
await mcpServer.connect(serverTransport);

0 commit comments

Comments
(0)

AltStyle によって変換されたページ (->オリジナル) /