rework state

Author: nirinchev

src/common/atlas/auth.ts

@@ -8,21 +8,21 @@ export async function ensureAuthenticated(state: State, apiClient: ApiClient): P
 }
 
 export async function isAuthenticated(state: State, apiClient: ApiClient): Promise<boolean> {
-    switch (state.persistent.auth.status) {
+    switch (state.credentials.auth.status) {
         case "not_auth":
             return false;
         case "requested":
             try {
-                if (!state.persistent.auth.code) {
+                if (!state.credentials.auth.code) {
                     return false;
                 }
-                await apiClient.retrieveToken(state.persistent.auth.code.device_code);
-                return !!state.persistent.auth.token;
+                await apiClient.retrieveToken(state.credentials.auth.code.device_code);
+                return !!state.credentials.auth.token;
             } catch {
                 return false;
             }
         case "issued":
-            if (!state.persistent.auth.token) {
+            if (!state.credentials.auth.token) {
                 return false;
             }
             return await apiClient.validateToken();

src/server.ts

@@ -1,6 +1,6 @@
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { ApiClient } from "./common/atlas/apiClient.js";
-import { State, saveState, loadState } from "./state.js";
+import defaultState, { State } from "./state.js";
 import { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
 import { registerAtlasTools } from "./tools/atlas/tools.js";
 import { registerMongoDBTools } from "./tools/mongodb/index.js";
@@ -9,26 +9,27 @@ import logger, { initializeLogger } from "./logger.js";
 import { mongoLogId } from "mongodb-log-writer";
 
 export class Server {
-    state: State | undefined = undefined;
+    state: State = defaultState;
     apiClient: ApiClient | undefined = undefined;
     initialized: boolean = false;
 
     private async init() {
         if (this.initialized) {
             return;
         }
-        this.state = await loadState();
+
+        await this.state.loadCredentials();
 
         this.apiClient = new ApiClient({
-            token: this.state.persistent.auth.token,
+            token: this.state.credentials.auth.token,
             saveToken: async (token) => {
                 if (!this.state) {
                     throw new Error("State is not initialized");
                 }
-                this.state.persistent.auth.code = undefined;
-                this.state.persistent.auth.token = token;
-                this.state.persistent.auth.status = "issued";
-                await saveState(this.state);
+                this.state.credentials.auth.code = undefined;
+                this.state.credentials.auth.token = token;
+                this.state.credentials.auth.status = "issued";
+                await this.state.persistCredentials();
             },
         });
 
@@ -43,8 +44,8 @@ export class Server {
 
         server.server.registerCapabilities({ logging: {} });
 
-        registerAtlasTools(server, this.state!, this.apiClient!);
-        registerMongoDBTools(server, this.state!);
+        registerAtlasTools(server, this.state, this.apiClient!);
+        registerMongoDBTools(server, this.state);
 
         return server;
     }

src/state.ts

@@ -4,48 +4,42 @@ import { AsyncEntry } from "@napi-rs/keyring";
 import logger from "./logger.js";
 import { mongoLogId } from "mongodb-log-writer";
 
-const entry = new AsyncEntry("mongodb-mcp", "credentials");
-
-export interface State {
-    persistent: {
-        auth: {
-            status: "not_auth" | "requested" | "issued";
-            code?: OauthDeviceCode;
-            token?: OAuthToken;
-        };
-        connectionString?: string;
-    };
-    session: {
-        serviceProvider?: NodeDriverServiceProvider;
+interface Credentials {
+    auth: {
+        status: "not_auth" | "requested" | "issued";
+        code?: OauthDeviceCode;
+        token?: OAuthToken;
     };
+    connectionString?: string;
 }
 
-const defaultState: State = {
-    persistent: {
+export class State {
+    private entry = new AsyncEntry("mongodb-mcp", "credentials");
+    credentials: Credentials = {
         auth: {
             status: "not_auth",
         },
-    },
-    session: {},
-};
+    };
+    serviceProvider?: NodeDriverServiceProvider;
 
-export async function saveState(state: State): Promise<void> {
-    await entry.setPassword(JSON.stringify(state.persistent));
-}
+    public async persistCredentials(): Promise<void> {
+        await this.entry.setPassword(JSON.stringify(this.credentials));
+    }
 
-export async function loadState(): Promise<State> {
-    try {
-        const data = await entry.getPassword();
-        if (!data) {
-            return defaultState;
-        }
+    public async loadCredentials(): Promise<boolean> {
+        try {
+            const data = await this.entry.getPassword();
+            if (data) {
+                this.credentials = JSON.parse(data);
+            }
 
-        return {
-            persistent: JSON.parse(data),
-            session: {},
-        };
-    } catch (err: unknown) {
-        logger.error(mongoLogId(1_000_007), "state", `Failed to load state: ${err}`);
-        return defaultState;
+            return true;
+        } catch (err: unknown) {
+            logger.error(mongoLogId(1_000_007), "state", `Failed to load state: ${err}`);
+            return false;
+        }
     }
 }
+
+const defaultState = new State();
+export default defaultState;

src/tools/atlas/auth.ts

@@ -1,5 +1,4 @@
 import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
-import { saveState } from "../../state.js";
 import { AtlasToolBase } from "./atlasTool.js";
 import { isAuthenticated } from "../../common/atlas/auth.js";
 import logger from "../../logger.js";
@@ -25,11 +24,11 @@ export class AuthTool extends AtlasToolBase {
         try {
             const code = await this.apiClient.authenticate();
 
-            this.state.persistent.auth.status = "requested";
-            this.state.persistent.auth.code = code;
-            this.state.persistent.auth.token = undefined;
+            this.state.credentials.auth.status = "requested";
+            this.state.credentials.auth.code = code;
+            this.state.credentials.auth.token = undefined;
 
-            await saveState(this.state);
+            await this.state.persistCredentials();
 
             return {
                 content: [

src/tools/mongodb/connect.ts

@@ -4,7 +4,6 @@ import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver
 import { DbOperationType, MongoDBToolBase } from "./mongodbTool.js";
 import { ToolArgs } from "../tool.js";
 import { ErrorCodes, MongoDBError } from "../../errors.js";
-import { saveState } from "../../state.js";
 
 export class ConnectTool extends MongoDBToolBase {
     protected name = "connect";
@@ -21,7 +20,7 @@ export class ConnectTool extends MongoDBToolBase {
     protected async execute({
         connectionStringOrClusterName,
     }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
-        connectionStringOrClusterName ??= this.state.persistent.connectionString;
+        connectionStringOrClusterName ??= this.state.credentials.connectionString;
         if (!connectionStringOrClusterName) {
             return {
                 content: [
@@ -71,8 +70,8 @@ export class ConnectTool extends MongoDBToolBase {
             productName: "MongoDB MCP",
         });
 
-        this.state.session.serviceProvider = provider;
-        this.state.persistent.connectionString = connectionString;
-        await saveState(this.state);
+        this.state.serviceProvider = provider;
+        this.state.credentials.connectionString = connectionString;
+        await this.state.persistCredentials();
     }
 }

src/tools/mongodb/mongodbTool.ts

@@ -20,7 +20,7 @@ export abstract class MongoDBToolBase extends ToolBase {
     protected abstract operationType: DbOperationType;
 
     protected ensureConnected(): NodeDriverServiceProvider {
-        const provider = this.state.session.serviceProvider;
+        const provider = this.state.serviceProvider;
         if (!provider) {
             throw new MongoDBError(ErrorCodes.NotConnectedToMongoDB, "Not connected to MongoDB");
         }