fix state-saving

Author: nirinchev

src/config.ts

@@ -5,6 +5,8 @@ import os from "os";
 
 import argv from "yargs-parser";
 
+const localDataPath = getLocalDataPath();
+
 // If we decide to support non-string config options, we'll need to extend the mechanism for parsing
 // env variables.
 interface UserConfig extends Record<string, string> {
@@ -19,7 +21,7 @@ const cliConfig = argv(process.argv.slice(2)) as unknown as Partial<UserConfig>;
 const defaults: UserConfig = {
     apiBaseUrl: "https://cloud.mongodb.com/",
     clientId: "0oabtxactgS3gHIR0297",
-    stateFile: path.join(getLocalDataPath(), "state.json"),
+    stateFile: path.join(localDataPath, "state.json"),
     projectId: "",
 };
 
@@ -36,21 +38,27 @@ const config = {
     atlasApiVersion: `2025-03-12`,
     version: packageJson.version,
     userAgent: `AtlasMCP/${packageJson.version} (${process.platform}; ${process.arch}; ${process.env.HOSTNAME || "unknown"})`,
-    localDataPath: getLocalDataPath(),
+    localDataPath,
 };
 
 export default config;
 
 function getLocalDataPath(): string {
+    let result: string | undefined;
+
     if (process.platform === "win32") {
         const appData = process.env.APPDATA;
         const localAppData = process.env.LOCALAPPDATA ?? process.env.APPDATA;
         if (localAppData && appData) {
-            return path.join(localAppData, "mongodb", "mongodb-mcp");
+            result = path.join(localAppData, "mongodb", "mongodb-mcp");
         }
     }
 
-    return path.join(os.homedir(), ".mongodb", "mongodb-mcp");
+    result ??= path.join(os.homedir(), ".mongodb", "mongodb-mcp");
+
+    fs.mkdirSync(result, { recursive: true });
+
+    return result;
 }
 
 // Gets the config supplied by the user as environment variables. The variable names
@@ -75,7 +83,7 @@ function getEnvConfig(): Partial<UserConfig> {
 // Gets the config supplied by the user as a JSON file. The file is expected to be located in the local data path
 // and named `config.json`.
 function getFileConfig(): Partial<UserConfig> {
-    const configPath = path.join(getLocalDataPath(), "config.json");
+    const configPath = path.join(localDataPath, "config.json");
 
     try {
         const config = fs.readFileSync(configPath, "utf8");

src/server.ts

@@ -21,14 +21,14 @@ export class Server {
 
         this.apiClient = new ApiClient({
             token: this.state?.auth.token,
-            saveToken: (token) => {
+            saveToken: async (token) => {
                 if (!this.state) {
                     throw new Error("State is not initialized");
                 }
                 this.state.auth.code = undefined;
                 this.state.auth.token = token;
                 this.state.auth.status = "issued";
-                saveState(this.state);
+                await saveState(this.state);
             },
         });
 

src/state.ts

@@ -1,4 +1,4 @@
-import fs from "fs";
+import fs from "fs/promises";
 import config from "./config.js";
 import { OauthDeviceCode, OAuthToken } from "./common/atlas/apiClient.js";
 
@@ -8,37 +8,26 @@ export interface State {
         code?: OauthDeviceCode;
         token?: OAuthToken;
     };
+    connectionString?: string;
 }
 
 export async function saveState(state: State): Promise<void> {
-    return new Promise((resolve, reject) => {
-        fs.writeFile(config.stateFile, JSON.stringify(state), function (err) {
-            if (err) {
-                return reject(err);
-            }
-
-            return resolve();
-        });
-    });
+    await fs.writeFile(config.stateFile, JSON.stringify(state), { encoding: "utf-8" });
 }
 
-export async function loadState() {
-    return new Promise<State>((resolve, reject) => {
-        fs.readFile(config.stateFile, "utf-8", (err, data) => {
-            if (err) {
-                if (err.code === "ENOENT") {
-                    // File does not exist, return default state
-                    const defaultState: State = {
-                        auth: {
-                            status: "not_auth",
-                        },
-                    };
-                    return resolve(defaultState);
-                } else {
-                    return reject(err);
-                }
-            }
-            return resolve(JSON.parse(data) as State);
-        });
-    });
+export async function loadState(): Promise<State> {
+    try {
+        const data = await fs.readFile(config.stateFile, "utf-8");
+        return JSON.parse(data) as State;
+    } catch (err: unknown) {
+        if (err && typeof err === "object" && "code" in err && err.code === "ENOENT") {
+            return {
+                auth: {
+                    status: "not_auth",
+                },
+            };
+        }
+
+        throw err;
+    }
 }

src/tools/atlas/auth.ts

@@ -11,7 +11,7 @@ export class AuthTool extends AtlasToolBase {
     protected argsShape = {};
 
     private async isAuthenticated(): Promise<boolean> {
-        return isAuthenticated(this.state!, this.apiClient);
+        return isAuthenticated(this.state, this.apiClient);
     }
 
     async execute(): Promise<CallToolResult> {
@@ -25,11 +25,11 @@ export class AuthTool extends AtlasToolBase {
         try {
             const code = await this.apiClient.authenticate();
 
-            this.state!.auth.status = "requested";
-            this.state!.auth.code = code;
-            this.state!.auth.token = undefined;
+            this.state.auth.status = "requested";
+            this.state.auth.code = code;
+            this.state.auth.token = undefined;
 
-            await saveState(this.state!);
+            await saveState(this.state);
 
             return {
                 content: [

src/tools/mongodb/connect.ts

@@ -4,6 +4,7 @@ import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver
 import { DbOperationType, MongoDBToolBase } from "./mongodbTool.js";
 import { ToolArgs } from "../tool";
 import { ErrorCodes, MongoDBError } from "../../errors.js";
+import { saveState } from "../../state.js";
 
 export class ConnectTool extends MongoDBToolBase {
     protected name = "connect";
@@ -20,8 +21,8 @@ export class ConnectTool extends MongoDBToolBase {
     protected async execute({
         connectionStringOrClusterName,
     }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
+        connectionStringOrClusterName ??= this.state.connectionString;
         if (!connectionStringOrClusterName) {
-            // TODO: try reconnecting to the default connection
             return {
                 content: [
                     { type: "text", text: "No connection details provided." },
@@ -71,5 +72,7 @@ export class ConnectTool extends MongoDBToolBase {
         });
 
         this.mongodbState.serviceProvider = provider;
+        this.state.connectionString = connectionString;
+        await saveState(this.state);
     }
 }