Skip to content

Commit f895766

Browse files
gagikgagik
committed
chore: add createSessionConfig hook MCP-294
Add a new async hook for library consumers to provide a configuration right before establishing an MCP session.
1 parent 40cb62e commit f895766

File tree

2 files changed

+174
-7
lines changed

2 files changed

+174
-7
lines changed

‎src/transports/base.ts‎

Lines changed: 21 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@ import type { Client } from "@mongodb-js/atlas-local";
2222
import { VectorSearchEmbeddingsManager } from "../common/search/vectorSearchEmbeddingsManager.js";
2323
import type { ToolBase, ToolConstructorParams } from "../tools/tool.js";
2424

25+
type CreateSessionConfigFn = (userConfig: UserConfig) => Promise<UserConfig>;
26+
2527
export type TransportRunnerConfig = {
2628
userConfig: UserConfig;
2729
createConnectionManager?: ConnectionManagerFactoryFn;
@@ -30,6 +32,11 @@ export type TransportRunnerConfig = {
3032
additionalLoggers?: LoggerBase[];
3133
telemetryProperties?: Partial<CommonProperties>;
3234
tools?: (new (params: ToolConstructorParams) => ToolBase)[];
35+
/**
36+
* Hook which allows library consumers to fetch configuration from external sources (e.g., secrets managers, APIs)
37+
* or modify the existing configuration before the session is created.
38+
*/
39+
createSessionConfig?: CreateSessionConfigFn;
3340
};
3441

3542
export abstract class TransportRunnerBase {
@@ -41,6 +48,7 @@ export abstract class TransportRunnerBase {
4148
private readonly atlasLocalClient: Promise<Client | undefined>;
4249
private readonly telemetryProperties: Partial<CommonProperties>;
4350
private readonly tools?: (new (params: ToolConstructorParams) => ToolBase)[];
51+
private readonly createSessionConfig?: CreateSessionConfigFn;
4452

4553
protected constructor({
4654
userConfig,
@@ -50,13 +58,15 @@ export abstract class TransportRunnerBase {
5058
additionalLoggers = [],
5159
telemetryProperties = {},
5260
tools,
61+
createSessionConfig,
5362
}: TransportRunnerConfig) {
5463
this.userConfig = userConfig;
5564
this.createConnectionManager = createConnectionManager;
5665
this.connectionErrorHandler = connectionErrorHandler;
5766
this.atlasLocalClient = createAtlasLocalClient();
5867
this.telemetryProperties = telemetryProperties;
5968
this.tools = tools;
69+
this.createSessionConfig = createSessionConfig;
6070
const loggers: LoggerBase[] = [...additionalLoggers];
6171
if (this.userConfig.loggers.includes("stderr")) {
6272
loggers.push(new ConsoleLogger(Keychain.root));
@@ -81,30 +91,34 @@ export abstract class TransportRunnerBase {
8191
}
8292

8393
protected async setupServer(): Promise<Server> {
94+
// Call the config provider hook if provided, allowing consumers to
95+
// fetch or modify configuration before session initialization
96+
const userConfig = this.createSessionConfig ? await this.createSessionConfig(this.userConfig) : this.userConfig;
97+
8498
const mcpServer = new McpServer({
8599
name: packageInfo.mcpServerName,
86100
version: packageInfo.version,
87101
});
88102

89103
const logger = new CompositeLogger(this.logger);
90-
const exportsManager = ExportsManager.init(this.userConfig, logger);
104+
const exportsManager = ExportsManager.init(userConfig, logger);
91105
const connectionManager = await this.createConnectionManager({
92106
logger,
93-
userConfig: this.userConfig,
107+
userConfig,
94108
deviceId: this.deviceId,
95109
});
96110

97111
const session = new Session({
98-
userConfig: this.userConfig,
112+
userConfig,
99113
atlasLocalClient: await this.atlasLocalClient,
100114
logger,
101115
exportsManager,
102116
connectionManager,
103117
keychain: Keychain.root,
104-
vectorSearchEmbeddingsManager: new VectorSearchEmbeddingsManager(this.userConfig, connectionManager),
118+
vectorSearchEmbeddingsManager: new VectorSearchEmbeddingsManager(userConfig, connectionManager),
105119
});
106120

107-
const telemetry = Telemetry.create(session, this.userConfig, this.deviceId, {
121+
const telemetry = Telemetry.create(session, userConfig, this.deviceId, {
108122
commonProperties: this.telemetryProperties,
109123
});
110124

@@ -114,15 +128,15 @@ export abstract class TransportRunnerBase {
114128
mcpServer,
115129
session,
116130
telemetry,
117-
userConfig: this.userConfig,
131+
userConfig,
118132
connectionErrorHandler: this.connectionErrorHandler,
119133
elicitation,
120134
tools: this.tools,
121135
});
122136

123137
// We need to create the MCP logger after the server is constructed
124138
// because it needs the server instance
125-
if (this.userConfig.loggers.includes("mcp")) {
139+
if (userConfig.loggers.includes("mcp")) {
126140
logger.addLogger(new McpLogger(result, Keychain.root));
127141
}
128142

‎tests/integration/transports/createSessionConfig.test.ts‎

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
import { StreamableHttpRunner } from "../../../src/transports/streamableHttp.js";
2+
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
3+
import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
4+
import { describe, expect, it, beforeAll, afterAll } from "vitest";
5+
import type { UserConfig } from "../../../src/common/config.js";
6+
import { TransportRunnerConfig } from "../../../src/lib.js";
7+
import { defaultTestConfig } from "../helpers.js";
8+
9+
describe("createSessionConfig", () => {
10+
const userConfig = defaultTestConfig;
11+
let runner: StreamableHttpRunner;
12+
13+
describe("basic functionality", () => {
14+
it("should use the modified config from createSessionConfig", async () => {
15+
const createSessionConfig: TransportRunnerConfig["createSessionConfig"] = async (userConfig) => {
16+
return {
17+
...userConfig,
18+
apiBaseUrl: "https://test-api.mongodb.com/",
19+
};
20+
};
21+
22+
userConfig.httpPort = 0; // Use a random port
23+
runner = new StreamableHttpRunner({
24+
userConfig,
25+
createSessionConfig,
26+
});
27+
await runner.start();
28+
29+
const server = await runner["setupServer"]();
30+
expect(server.userConfig.apiBaseUrl).toBe("https://test-api.mongodb.com/");
31+
32+
await runner.close();
33+
});
34+
35+
it("should work without a createSessionConfig", async () => {
36+
userConfig.httpPort = 0; // Use a random port
37+
runner = new StreamableHttpRunner({
38+
userConfig,
39+
});
40+
await runner.start();
41+
42+
const server = await runner["setupServer"]();
43+
expect(server.userConfig.apiBaseUrl).toBe(userConfig.apiBaseUrl);
44+
45+
await runner.close();
46+
});
47+
});
48+
49+
describe("connection string modification", () => {
50+
it("should allow modifying connection string via createSessionConfig", async () => {
51+
const createSessionConfig: TransportRunnerConfig["createSessionConfig"] = async (userConfig) => {
52+
// Simulate fetching connection string from environment or secrets
53+
await new Promise((resolve) => setTimeout(resolve, 10));
54+
55+
return {
56+
...userConfig,
57+
connectionString: "mongodb://test-server:27017/test-db",
58+
};
59+
};
60+
61+
userConfig.httpPort = 0; // Use a random port
62+
runner = new StreamableHttpRunner({
63+
userConfig: { ...userConfig, connectionString: undefined },
64+
createSessionConfig,
65+
});
66+
await runner.start();
67+
68+
const server = await runner["setupServer"]();
69+
expect(server.userConfig.connectionString).toBe("mongodb://test-server:27017/test-db");
70+
71+
await runner.close();
72+
});
73+
});
74+
75+
describe("server integration", () => {
76+
let client: Client;
77+
let transport: StreamableHTTPClientTransport;
78+
79+
it("should successfully initialize server with createSessionConfig and serve requests", async () => {
80+
const createSessionConfig: TransportRunnerConfig["createSessionConfig"] = async (userConfig) => {
81+
// Simulate async config modification
82+
await new Promise((resolve) => setTimeout(resolve, 10));
83+
return {
84+
...userConfig,
85+
readOnly: true, // Enable read-only mode
86+
};
87+
};
88+
89+
userConfig.httpPort = 0; // Use a random port
90+
runner = new StreamableHttpRunner({
91+
userConfig,
92+
createSessionConfig,
93+
});
94+
await runner.start();
95+
96+
client = new Client({
97+
name: "test-client",
98+
version: "1.0.0",
99+
});
100+
transport = new StreamableHTTPClientTransport(new URL(`${runner.serverAddress}/mcp`));
101+
102+
await client.connect(transport);
103+
const response = await client.listTools();
104+
105+
expect(response).toBeDefined();
106+
expect(response.tools).toBeDefined();
107+
expect(response.tools.length).toBeGreaterThan(0);
108+
109+
// Verify read-only mode is applied - insert-one should not be available
110+
const writeTools = response.tools.filter((tool) => tool.name === "insert-one");
111+
expect(writeTools.length).toBe(0);
112+
113+
// Verify read tools are available
114+
const readTools = response.tools.filter((tool) => tool.name === "find");
115+
expect(readTools.length).toBe(1);
116+
117+
await client.close();
118+
await transport.close();
119+
await runner.close();
120+
});
121+
});
122+
123+
describe("error handling", () => {
124+
it("should propagate errors from configProvider on client connection", async () => {
125+
const createSessionConfig: TransportRunnerConfig["createSessionConfig"] = async () => {
126+
throw new Error("Failed to fetch config");
127+
};
128+
129+
userConfig.httpPort = 0; // Use a random port
130+
runner = new StreamableHttpRunner({
131+
userConfig,
132+
createSessionConfig,
133+
});
134+
135+
// Start succeeds because setupServer is only called when a client connects
136+
await runner.start();
137+
138+
// Error should occur when a client tries to connect
139+
const testClient = new Client({
140+
name: "test-client",
141+
version: "1.0.0",
142+
});
143+
const testTransport = new StreamableHTTPClientTransport(new URL(`${runner.serverAddress}/mcp`));
144+
145+
await expect(testClient.connect(testTransport)).rejects.toThrow();
146+
147+
await testClient.close();
148+
await testTransport.close();
149+
150+
await runner.close();
151+
});
152+
});
153+
});

0 commit comments

Comments
 (0)