Skip to content

Commit a4b0751

Browse files
nirinchevnirinchev
committed
rework state
1 parent eecbdcf commit a4b0751

File tree

6 files changed

+53
-60
lines changed

6 files changed

+53
-60
lines changed

‎src/common/atlas/auth.ts‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,21 +8,21 @@ export async function ensureAuthenticated(state: State, apiClient: ApiClient): P
88
}
99

1010
export async function isAuthenticated(state: State, apiClient: ApiClient): Promise<boolean> {
11-
switch (state.persistent.auth.status) {
11+
switch (state.credentials.auth.status) {
1212
case "not_auth":
1313
return false;
1414
case "requested":
1515
try {
16-
if (!state.persistent.auth.code) {
16+
if (!state.credentials.auth.code) {
1717
return false;
1818
}
19-
await apiClient.retrieveToken(state.persistent.auth.code.device_code);
20-
return !!state.persistent.auth.token;
19+
await apiClient.retrieveToken(state.credentials.auth.code.device_code);
20+
return !!state.credentials.auth.token;
2121
} catch {
2222
return false;
2323
}
2424
case "issued":
25-
if (!state.persistent.auth.token) {
25+
if (!state.credentials.auth.token) {
2626
return false;
2727
}
2828
return await apiClient.validateToken();

‎src/server.ts‎

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
22
import { ApiClient } from "./common/atlas/apiClient.js";
3-
import { State, saveState, loadState } from "./state.js";
3+
import defaultState, { State } from "./state.js";
44
import { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
55
import { registerAtlasTools } from "./tools/atlas/tools.js";
66
import { registerMongoDBTools } from "./tools/mongodb/index.js";
@@ -9,26 +9,27 @@ import logger, { initializeLogger } from "./logger.js";
99
import { mongoLogId } from "mongodb-log-writer";
1010

1111
export class Server {
12-
state: State | undefined = undefined;
12+
state: State = defaultState;
1313
apiClient: ApiClient | undefined = undefined;
1414
initialized: boolean = false;
1515

1616
private async init() {
1717
if (this.initialized) {
1818
return;
1919
}
20-
this.state = await loadState();
20+
21+
await this.state.loadCredentials();
2122

2223
this.apiClient = new ApiClient({
23-
token: this.state.persistent.auth.token,
24+
token: this.state.credentials.auth.token,
2425
saveToken: async (token) => {
2526
if (!this.state) {
2627
throw new Error("State is not initialized");
2728
}
28-
this.state.persistent.auth.code = undefined;
29-
this.state.persistent.auth.token = token;
30-
this.state.persistent.auth.status = "issued";
31-
await saveState(this.state);
29+
this.state.credentials.auth.code = undefined;
30+
this.state.credentials.auth.token = token;
31+
this.state.credentials.auth.status = "issued";
32+
await this.state.persistCredentials();
3233
},
3334
});
3435

@@ -43,8 +44,8 @@ export class Server {
4344

4445
server.server.registerCapabilities({ logging: {} });
4546

46-
registerAtlasTools(server, this.state!, this.apiClient!);
47-
registerMongoDBTools(server, this.state!);
47+
registerAtlasTools(server, this.state, this.apiClient!);
48+
registerMongoDBTools(server, this.state);
4849

4950
return server;
5051
}

‎src/state.ts‎

Lines changed: 28 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -4,48 +4,42 @@ import { AsyncEntry } from "@napi-rs/keyring";
44
import logger from "./logger.js";
55
import { mongoLogId } from "mongodb-log-writer";
66

7-
const entry = new AsyncEntry("mongodb-mcp", "credentials");
8-
9-
export interface State {
10-
persistent: {
11-
auth: {
12-
status: "not_auth" | "requested" | "issued";
13-
code?: OauthDeviceCode;
14-
token?: OAuthToken;
15-
};
16-
connectionString?: string;
17-
};
18-
session: {
19-
serviceProvider?: NodeDriverServiceProvider;
7+
interface Credentials {
8+
auth: {
9+
status: "not_auth" | "requested" | "issued";
10+
code?: OauthDeviceCode;
11+
token?: OAuthToken;
2012
};
13+
connectionString?: string;
2114
}
2215

23-
const defaultState: State = {
24-
persistent: {
16+
export class State {
17+
private entry = new AsyncEntry("mongodb-mcp", "credentials");
18+
credentials: Credentials = {
2519
auth: {
2620
status: "not_auth",
2721
},
28-
},
29-
session: {},
30-
};
22+
};
23+
serviceProvider?: NodeDriverServiceProvider;
3124

32-
export async function saveState(state: State): Promise<void> {
33-
await entry.setPassword(JSON.stringify(state.persistent));
34-
}
25+
public async persistCredentials(): Promise<void> {
26+
await this.entry.setPassword(JSON.stringify(this.credentials));
27+
}
3528

36-
export async function loadState(): Promise<State> {
37-
try {
38-
const data = await entry.getPassword();
39-
if (!data) {
40-
return defaultState;
41-
}
29+
public async loadCredentials(): Promise<boolean> {
30+
try {
31+
const data = await this.entry.getPassword();
32+
if (data) {
33+
this.credentials = JSON.parse(data);
34+
}
4235

43-
return {
44-
persistent: JSON.parse(data),
45-
session: {},
46-
};
47-
} catch (err: unknown) {
48-
logger.error(mongoLogId(1_000_007), "state", `Failed to load state: ${err}`);
49-
return defaultState;
36+
return true;
37+
} catch (err: unknown) {
38+
logger.error(mongoLogId(1_000_007), "state", `Failed to load state: ${err}`);
39+
return false;
40+
}
5041
}
5142
}
43+
44+
const defaultState = new State();
45+
export default defaultState;

‎src/tools/atlas/auth.ts‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
2-
import { saveState } from "../../state.js";
32
import { AtlasToolBase } from "./atlasTool.js";
43
import { isAuthenticated } from "../../common/atlas/auth.js";
54
import logger from "../../logger.js";
@@ -25,11 +24,11 @@ export class AuthTool extends AtlasToolBase {
2524
try {
2625
const code = await this.apiClient.authenticate();
2726

28-
this.state.persistent.auth.status = "requested";
29-
this.state.persistent.auth.code = code;
30-
this.state.persistent.auth.token = undefined;
27+
this.state.credentials.auth.status = "requested";
28+
this.state.credentials.auth.code = code;
29+
this.state.credentials.auth.token = undefined;
3130

32-
await saveState(this.state);
31+
await this.state.persistCredentials();
3332

3433
return {
3534
content: [

‎src/tools/mongodb/connect.ts‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver
44
import { DbOperationType, MongoDBToolBase } from "./mongodbTool.js";
55
import { ToolArgs } from "../tool.js";
66
import { ErrorCodes, MongoDBError } from "../../errors.js";
7-
import { saveState } from "../../state.js";
87

98
export class ConnectTool extends MongoDBToolBase {
109
protected name = "connect";
@@ -21,7 +20,7 @@ export class ConnectTool extends MongoDBToolBase {
2120
protected async execute({
2221
connectionStringOrClusterName,
2322
}: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
24-
connectionStringOrClusterName ??= this.state.persistent.connectionString;
23+
connectionStringOrClusterName ??= this.state.credentials.connectionString;
2524
if (!connectionStringOrClusterName) {
2625
return {
2726
content: [
@@ -71,8 +70,8 @@ export class ConnectTool extends MongoDBToolBase {
7170
productName: "MongoDB MCP",
7271
});
7372

74-
this.state.session.serviceProvider = provider;
75-
this.state.persistent.connectionString = connectionString;
76-
await saveState(this.state);
73+
this.state.serviceProvider = provider;
74+
this.state.credentials.connectionString = connectionString;
75+
await this.state.persistCredentials();
7776
}
7877
}

‎src/tools/mongodb/mongodbTool.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ export abstract class MongoDBToolBase extends ToolBase {
2020
protected abstract operationType: DbOperationType;
2121

2222
protected ensureConnected(): NodeDriverServiceProvider {
23-
const provider = this.state.session.serviceProvider;
23+
const provider = this.state.serviceProvider;
2424
if (!provider) {
2525
throw new MongoDBError(ErrorCodes.NotConnectedToMongoDB, "Not connected to MongoDB");
2626
}

0 commit comments

Comments
 (0)