Skip to content

Commit f220e86

Browse files
gagikgagik
authored
Merge a49fad5 into 40cb62e 
2 parents 40cb62e + a49fad5 commit f220e86

File tree

2 files changed

+172
-7
lines changed

2 files changed

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

0 commit comments

Comments
 (0)