Skip to content

Commit 8ca583a

Browse files
gagikfmenezes
gagik
and
fmenezes
authored
refactor: rename state to session, combine tool registration, and clearer dependencies (#55)
Co-authored-by: Filipe Constantinov Menezes <filipe.menezes@mongodb.com>
1 parent 3f1dece commit 8ca583a

23 files changed

+142
-153
lines changed

‎eslint.config.js‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ export default defineConfig([
2222
files,
2323
rules: {
2424
"@typescript-eslint/switch-exhaustiveness-check": "error",
25+
"@typescript-eslint/no-non-null-assertion": "error",
2526
},
2627
},
2728
// Ignore features specific to TypeScript resolved rules

‎src/common/atlas/apiClient.ts‎

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -55,30 +55,31 @@ export class ApiClient {
5555
return this.accessToken?.token.access_token as string | undefined;
5656
};
5757

58-
private authMiddleware = (apiClient: ApiClient): Middleware => ({
59-
async onRequest({ request, schemaPath }) {
58+
private authMiddleware: Middleware = {
59+
onRequest: async ({ request, schemaPath }) => {
6060
if (schemaPath.startsWith("/api/private/unauth") || schemaPath.startsWith("/api/oauth")) {
6161
return undefined;
6262
}
6363

6464
try {
65-
const accessToken = await apiClient.getAccessToken();
65+
const accessToken = await this.getAccessToken();
6666
request.headers.set("Authorization", `Bearer ${accessToken}`);
6767
return request;
6868
} catch {
6969
// ignore not availble tokens, API will return 401
7070
}
7171
},
72-
});
73-
private errorMiddleware = (): Middleware => ({
72+
};
73+
74+
private readonly errorMiddleware: Middleware = {
7475
async onResponse({ response }) {
7576
if (!response.ok) {
7677
throw await ApiClientError.fromResponse(response);
7778
}
7879
},
79-
});
80+
};
8081

81-
constructor(options?: ApiClientOptions) {
82+
constructor(options: ApiClientOptions) {
8283
const defaultOptions = {
8384
baseUrl: "https://cloud.mongodb.com/",
8485
userAgent: `AtlasMCP/${config.version} (${process.platform}; ${process.arch}; ${process.env.HOSTNAME || "unknown"})`,
@@ -107,9 +108,9 @@ export class ApiClient {
107108
tokenPath: "/api/oauth/token",
108109
},
109110
});
110-
this.client.use(this.authMiddleware(this));
111+
this.client.use(this.authMiddleware);
111112
}
112-
this.client.use(this.errorMiddleware());
113+
this.client.use(this.errorMiddleware);
113114
}
114115

115116
public async getIpInfo(): Promise<{

‎src/config.ts‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,10 @@ function getLogPath(): string {
6060
// to SNAKE_UPPER_CASE.
6161
function getEnvConfig(): Partial<UserConfig> {
6262
function setValue(obj: Record<string, unknown>, path: string[], value: string): void {
63-
const currentField = path.shift()!;
63+
const currentField = path.shift();
64+
if (!currentField) {
65+
return;
66+
}
6467
if (path.length === 0) {
6568
const numberValue = Number(value);
6669
if (!isNaN(numberValue)) {

‎src/index.ts‎

Lines changed: 19 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,30 @@
11
#!/usr/bin/env node
22

33
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
4-
import { Server } from "./server.js";
54
import logger from "./logger.js";
65
import { mongoLogId } from "mongodb-log-writer";
6+
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
7+
import config from "./config.js";
8+
import { Session } from "./session.js";
9+
import { Server } from "./server.js";
710

8-
export async function runServer() {
9-
const server = new Server();
11+
try {
12+
const session = new Session();
13+
const mcpServer = new McpServer({
14+
name: "MongoDB Atlas",
15+
version: config.version,
16+
});
17+
18+
const server = new Server({
19+
mcpServer,
20+
session,
21+
});
1022

1123
const transport = new StdioServerTransport();
12-
await server.connect(transport);
13-
}
1424

15-
runServer().catch((error) => {
16-
logger.emergency(mongoLogId(1_000_004), "server", `Fatal error running server: ${error}`);
25+
await server.connect(transport);
26+
} catch (error: unknown) {
27+
logger.emergency(mongoLogId(1_000_004), "server", `Fatal error running server: ${error as string}`);
1728

1829
process.exit(1);
19-
});
30+
}

‎src/server.ts‎

Lines changed: 21 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,39 +1,40 @@
11
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
2-
import defaultState, { State } from "./state.js";
2+
import { Session } from "./session.js";
33
import { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
4-
import { registerAtlasTools } from "./tools/atlas/tools.js";
5-
import { registerMongoDBTools } from "./tools/mongodb/index.js";
6-
import config from "./config.js";
4+
import { AtlasTools } from "./tools/atlas/tools.js";
5+
import { MongoDbTools } from "./tools/mongodb/tools.js";
76
import logger, { initializeLogger } from "./logger.js";
87
import { mongoLogId } from "mongodb-log-writer";
98

109
export class Server {
11-
state: State = defaultState;
12-
private server?: McpServer;
10+
public readonly session: Session;
11+
private readonly mcpServer: McpServer;
12+
13+
constructor({ mcpServer, session }: { mcpServer: McpServer; session: Session }) {
14+
this.mcpServer = mcpServer;
15+
this.session = session;
16+
}
1317

1418
async connect(transport: Transport) {
15-
this.server = new McpServer({
16-
name: "MongoDB Atlas",
17-
version: config.version,
18-
});
19+
this.mcpServer.server.registerCapabilities({ logging: {} });
1920

20-
this.server.server.registerCapabilities({ logging: {} });
21+
this.registerTools();
2122

22-
registerAtlasTools(this.server, this.state);
23-
registerMongoDBTools(this.server, this.state);
23+
await initializeLogger(this.mcpServer);
2424

25-
await initializeLogger(this.server);
26-
await this.server.connect(transport);
25+
await this.mcpServer.connect(transport);
2726

2827
logger.info(mongoLogId(1_000_004), "server", `Server started with transport ${transport.constructor.name}`);
2928
}
3029

3130
async close(): Promise<void> {
32-
try {
33-
await this.state.serviceProvider?.close(true);
34-
} catch {
35-
// Ignore errors during service provider close
31+
await this.session.close();
32+
await this.mcpServer.close();
33+
}
34+
35+
private registerTools() {
36+
for (const tool of [...AtlasTools, ...MongoDbTools]) {
37+
new tool(this.session).register(this.mcpServer);
3638
}
37-
await this.server?.close();
3839
}
3940
}

‎src/state.ts‎ renamed to ‎src/session.ts‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,11 @@ import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver
22
import { ApiClient } from "./common/atlas/apiClient.js";
33
import config from "./config.js";
44

5-
export class State {
5+
export class Session {
66
serviceProvider?: NodeDriverServiceProvider;
77
apiClient?: ApiClient;
88

9-
ensureApiClient(): asserts this is { apiClient: ApiClient } {
9+
ensureAuthenticated(): asserts this is { apiClient: ApiClient } {
1010
if (!this.apiClient) {
1111
if (!config.apiClientId || !config.apiClientSecret) {
1212
throw new Error(
@@ -23,7 +23,15 @@ export class State {
2323
});
2424
}
2525
}
26-
}
2726

28-
const defaultState = new State();
29-
export default defaultState;
27+
async close(): Promise<void> {
28+
if (this.serviceProvider) {
29+
try {
30+
await this.serviceProvider.close(true);
31+
} catch (error) {
32+
console.error("Error closing service provider:", error);
33+
}
34+
this.serviceProvider = undefined;
35+
}
36+
}
37+
}

‎src/tools/atlas/atlasTool.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
import { ToolBase } from "../tool.js";
2-
import { State } from "../../state.js";
2+
import { Session } from "../../session.js";
33

44
export abstract class AtlasToolBase extends ToolBase {
5-
constructor(state: State) {
6-
super(state);
5+
constructor(protected readonly session: Session) {
6+
super(session);
77
}
88
}

‎src/tools/atlas/createAccessList.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ export class CreateAccessListTool extends AtlasToolBase {
2626
comment,
2727
currentIpAddress,
2828
}: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
29-
this.state.ensureApiClient();
29+
this.session.ensureAuthenticated();
3030

3131
if (!ipAddresses?.length && !cidrBlocks?.length && !currentIpAddress) {
3232
throw new Error("One of ipAddresses, cidrBlocks, currentIpAddress must be provided.");
@@ -39,7 +39,7 @@ export class CreateAccessListTool extends AtlasToolBase {
3939
}));
4040

4141
if (currentIpAddress) {
42-
const currentIp = await this.state.apiClient.getIpInfo();
42+
const currentIp = await this.session.apiClient.getIpInfo();
4343
const input = {
4444
groupId: projectId,
4545
ipAddress: currentIp.currentIpv4Address,
@@ -56,7 +56,7 @@ export class CreateAccessListTool extends AtlasToolBase {
5656

5757
const inputs = [...ipInputs, ...cidrInputs];
5858

59-
await this.state.apiClient.createProjectIpAccessList({
59+
await this.session.apiClient.createProjectIpAccessList({
6060
params: {
6161
path: {
6262
groupId: projectId,

‎src/tools/atlas/createDBUser.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ export class CreateDBUserTool extends AtlasToolBase {
3333
roles,
3434
clusters,
3535
}: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
36-
this.state.ensureApiClient();
36+
this.session.ensureAuthenticated();
3737

3838
const input = {
3939
groupId: projectId,
@@ -53,7 +53,7 @@ export class CreateDBUserTool extends AtlasToolBase {
5353
: undefined,
5454
} as CloudDatabaseUser;
5555

56-
await this.state.apiClient.createDatabaseUser({
56+
await this.session.apiClient.createDatabaseUser({
5757
params: {
5858
path: {
5959
groupId: projectId,

‎src/tools/atlas/createFreeCluster.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ export class CreateFreeClusterTool extends AtlasToolBase {
1414
};
1515

1616
protected async execute({ projectId, name, region }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
17-
this.state.ensureApiClient();
17+
this.session.ensureAuthenticated();
1818

1919
const input = {
2020
groupId: projectId,
@@ -38,7 +38,7 @@ export class CreateFreeClusterTool extends AtlasToolBase {
3838
terminationProtectionEnabled: false,
3939
} as unknown as ClusterDescription20240805;
4040

41-
await this.state.apiClient.createCluster({
41+
await this.session.apiClient.createCluster({
4242
params: {
4343
path: {
4444
groupId: projectId,

0 commit comments

Comments
 (0)