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 <[email protected]>
1 parent 3f1dece commit 8ca583a

23 files changed

+142
-153
lines changed

‎eslint.config.js

+1
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

+10-9
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

+4-1
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

+19-8
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

+21-20
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

+13-5
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

+3-3
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

+3-3
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

+2-2
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

+2-2
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,

‎src/tools/atlas/inspectAccessList.ts

+2-2
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,9 @@ export class InspectAccessListTool extends AtlasToolBase {
1111
};
1212

1313
protected async execute({ projectId }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
14-
this.state.ensureApiClient();
14+
this.session.ensureAuthenticated();
1515

16-
const accessList = await this.state.apiClient.listProjectIpAccessLists({
16+
const accessList = await this.session.apiClient.listProjectIpAccessLists({
1717
params: {
1818
path: {
1919
groupId: projectId,

‎src/tools/atlas/inspectCluster.ts

+2-2
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,9 @@ export class InspectClusterTool extends AtlasToolBase {
1313
};
1414

1515
protected async execute({ projectId, clusterName }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
16-
this.state.ensureApiClient();
16+
this.session.ensureAuthenticated();
1717

18-
const cluster = await this.state.apiClient.getCluster({
18+
const cluster = await this.session.apiClient.getCluster({
1919
params: {
2020
path: {
2121
groupId: projectId,

‎src/tools/atlas/listClusters.ts

+4-4
Original file line numberDiff line numberDiff line change
@@ -12,14 +12,14 @@ export class ListClustersTool extends AtlasToolBase {
1212
};
1313

1414
protected async execute({ projectId }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
15-
this.state.ensureApiClient();
15+
this.session.ensureAuthenticated();
1616

1717
if (!projectId) {
18-
const data = await this.state.apiClient.listClustersForAllProjects();
18+
const data = await this.session.apiClient.listClustersForAllProjects();
1919

2020
return this.formatAllClustersTable(data);
2121
} else {
22-
const project = await this.state.apiClient.getProject({
22+
const project = await this.session.apiClient.getProject({
2323
params: {
2424
path: {
2525
groupId: projectId,
@@ -31,7 +31,7 @@ export class ListClustersTool extends AtlasToolBase {
3131
throw new Error(`Project with ID "${projectId}" not found.`);
3232
}
3333

34-
const data = await this.state.apiClient.listClusters({
34+
const data = await this.session.apiClient.listClusters({
3535
params: {
3636
path: {
3737
groupId: project.id || "",

‎src/tools/atlas/listDBUsers.ts

+2-2
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,9 @@ export class ListDBUsersTool extends AtlasToolBase {
1212
};
1313

1414
protected async execute({ projectId }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
15-
this.state.ensureApiClient();
15+
this.session.ensureAuthenticated();
1616

17-
const data = await this.state.apiClient.listDatabaseUsers({
17+
const data = await this.session.apiClient.listDatabaseUsers({
1818
params: {
1919
path: {
2020
groupId: projectId,

‎src/tools/atlas/listProjects.ts

+2-2
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,9 @@ export class ListProjectsTool extends AtlasToolBase {
77
protected argsShape = {};
88

99
protected async execute(): Promise<CallToolResult> {
10-
this.state.ensureApiClient();
10+
this.session.ensureAuthenticated();
1111

12-
const data = await this.state.apiClient.listProjects();
12+
const data = await this.session.apiClient.listProjects();
1313

1414
if (!data?.results?.length) {
1515
throw new Error("No projects found in your MongoDB Atlas account.");

‎src/tools/atlas/tools.ts

+10-19
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,3 @@
1-
import { ToolBase } from "../tool.js";
2-
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
3-
import { State } from "../../state.js";
41
import { ListClustersTool } from "./listClusters.js";
52
import { ListProjectsTool } from "./listProjects.js";
63
import { InspectClusterTool } from "./inspectCluster.js";
@@ -10,19 +7,13 @@ import { InspectAccessListTool } from "./inspectAccessList.js";
107
import { ListDBUsersTool } from "./listDBUsers.js";
118
import { CreateDBUserTool } from "./createDBUser.js";
129

13-
export function registerAtlasTools(server: McpServer, state: State) {
14-
const tools: ToolBase[] = [
15-
new ListClustersTool(state),
16-
new ListProjectsTool(state),
17-
new InspectClusterTool(state),
18-
new CreateFreeClusterTool(state),
19-
new CreateAccessListTool(state),
20-
new InspectAccessListTool(state),
21-
new ListDBUsersTool(state),
22-
new CreateDBUserTool(state),
23-
];
24-
25-
for (const tool of tools) {
26-
tool.register(server);
27-
}
28-
}
10+
export const AtlasTools = [
11+
ListClustersTool,
12+
ListProjectsTool,
13+
InspectClusterTool,
14+
CreateFreeClusterTool,
15+
CreateAccessListTool,
16+
InspectAccessListTool,
17+
ListDBUsersTool,
18+
CreateDBUserTool,
19+
];

‎src/tools/mongodb/connect.ts

+1-1
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ export class ConnectTool extends MongoDBToolBase {
5757
throw new MongoDBError(ErrorCodes.InvalidParams, "Invalid connection options");
5858
}
5959

60-
await this.connectToMongoDB(connectionString, this.state);
60+
await this.connectToMongoDB(connectionString);
6161

6262
return {
6363
content: [{ type: "text", text: `Successfully connected to ${connectionString}.` }],

0 commit comments

Comments
 (0)