chore: implementing OIDC callbacks and flow detection

kmruiz · kmruiz · commit c7856f0477ef · 2025-08-12T14:17:53.000+02:00
diff --git a/src/common/connectionManager.ts b/src/common/connectionManager.ts
@@ -1,11 +1,23 @@
-import { driverOptions } from "./config.js";
+import { config, UserConfig, driverOptions } from "./config.js";
 import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
 import EventEmitter from "events";
 import { setAppNameParamIfMissing } from "../helpers/connectionOptions.js";
 import { packageInfo } from "./packageInfo.js";
 import ConnectionString from "mongodb-connection-string-url";
-import { MongoClientOptions } from "mongodb";
+import { MongoClientOptions, OIDCCallbackParams } from "mongodb";
 import { ErrorCodes, MongoDBError } from "./errors.js";
+import type { MongoshBus } from "@mongosh/types";
+import { CompositeLogger, LogId } from "./logger.js";
+
+// https://github.com/mongodb-js/oidc-plugin/blob/main/src/types.ts
+const MONGODB_OIDC_CLIENT_ERROR_EVENTS = [
+    "mongodb-oidc-plugin:deserialization-failed",
+    "mongodb-oidc-plugin:local-listen-failed",
+    "mongodb-oidc-plugin:auth-attempt-failed",
+    "mongodb-oidc-plugin:refresh-failed",
+    "mongodb-oidc-plugin:auth-failed",
+    "mongodb-oidc-plugin:outbound-http-request-failed",
+] as const;
 
 export interface AtlasClusterConnectionInfo {
     username: string;
@@ -67,11 +79,25 @@ export interface ConnectionManagerEvents {
 
 export class ConnectionManager extends EventEmitter<ConnectionManagerEvents> {
     private state: AnyConnectionState;
+    private bus: MongoshBus;
 
-    constructor() {
+    constructor(
+        private logger: CompositeLogger,
+        bus?: MongoshBus
+    ) {
         super();
 
+        this.bus = bus ?? new EventEmitter();
         this.state = { tag: "disconnected" };
+
+        for (const clientErrorEvent of MONGODB_OIDC_CLIENT_ERROR_EVENTS) {
+            this.bus.on(clientErrorEvent, this.onOidcClientErrorCallback.bind(this, clientErrorEvent));
+        }
+
+        this.bus.on("mongodb-oidc-plugin:received-server-params", this.onOidcReceivedServerParameters.bind(this));
+        this.bus.on("mongodb-oidc-plugin:auth-failed", this.onOidcAuthenticationFailed.bind(this));
+        this.bus.on("mongodb-oidc-plugin:auth-succeeded", this.onOidcAuthSucceeded.bind(this));
+        this.bus.on("mongodb-oidc-plugin:notify-device-flow", this.onOidcNotifyDeviceFlow.bind(this));
     }
 
     async connect(settings: ConnectionSettings): Promise<AnyConnectionState> {
@@ -89,11 +115,16 @@ export class ConnectionManager extends EventEmitter<ConnectionManagerEvents> {
                 defaultAppName: `${packageInfo.mcpServerName} ${packageInfo.version}`,
             });
 
-            serviceProvider = await NodeDriverServiceProvider.connect(settings.connectionString, {
-                productDocsLink: "https://github.com/mongodb-js/mongodb-mcp-server/",
-                productName: "MongoDB MCP",
-                ...driverOptions,
-            });
+            serviceProvider = await NodeDriverServiceProvider.connect(
+                settings.connectionString,
+                {
+                    productDocsLink: "https://github.com/mongodb-js/mongodb-mcp-server/",
+                    productName: "MongoDB MCP",
+                    ...driverOptions,
+                },
+                undefined,
+                this.bus
+            );
         } catch (error: unknown) {
             const errorReason = error instanceof Error ? error.message : `${error as string}`;
             this.changeState("connection-errored", {
@@ -111,7 +142,7 @@ export class ConnectionManager extends EventEmitter<ConnectionManagerEvents> {
                 tag: "connected",
                 connectedAtlasCluster: settings.atlas,
                 serviceProvider,
-                connectionStringAuthType: ConnectionManager.inferConnectionTypeFromSettings(settings),
+                connectionStringAuthType: ConnectionManager.inferConnectionTypeFromSettings(config, settings),
             });
         } catch (error: unknown) {
             const errorReason = error instanceof Error ? error.message : `${error as string}`;
@@ -157,13 +188,92 @@ export class ConnectionManager extends EventEmitter<ConnectionManagerEvents> {
         return newState;
     }
 
-    static inferConnectionTypeFromSettings(settings: ConnectionSettings): ConnectionStringAuthType {
+    /**
+     * This handles OIDC errors happening at the client side. There is a difference between errors happening
+     * here and rejections from the server, so we will treat them differently.
+     */
+    private onOidcClientErrorCallback(event: string, eventData: { error: string }): void {
+        if (this.state.tag === "connecting" && this.state.connectionStringAuthType?.startsWith("oidc")) {
+            this.changeState("connection-errored", { tag: "errored", errorReason: eventData.error });
+        }
+
+        // If we are not in the connecting state, we really can't do much with this. It might be a
+        // transitional error (connectivity issues), a refresh token issue (server side configuration),
+        // anything. Log the error to get some context later for debugging.
+        this.logger.error({
+            id: LogId.oidcClientError,
+            context: event,
+            message: `Error during OIDC flow, at event ${event}: ${eventData.error}`,
+        });
+    }
+
+    private onOidcReceivedServerParameters(eventData: OIDCCallbackParams): void {
+        if (this.state.tag === "connecting" && this.state.connectionStringAuthType?.startsWith("oidc")) {
+            this.changeState("connection-succeeded", { ...this.state, tag: "connected" });
+        }
+
+        this.logger.info({
+            id: LogId.oidcFlow,
+            context: "mongodb-oidc-plugin:received-server-parameters",
+            message: `Received server parameters successfully: ${JSON.stringify(eventData)}`,
+        });
+    }
+
+    private onOidcAuthenticationFailed(eventData: { error: string }): void {
+        if (this.state.tag === "connecting" && this.state.connectionStringAuthType?.startsWith("oidc")) {
+            this.changeState("connection-succeeded", { ...this.state, tag: "connected" });
+        }
+
+        this.logger.error({
+            id: LogId.oidcAuthFailed,
+            context: "mongodb-oidc-plugin:authentication-failed",
+            message: `Could not authenticate: ${eventData.error}`,
+        });
+    }
+
+    private onOidcAuthSucceeded(): void {
+        if (this.state.tag === "connecting" && this.state.connectionStringAuthType?.startsWith("oidc")) {
+            this.changeState("connection-succeeded", { ...this.state, tag: "connected" });
+        }
+
+        this.logger.info({
+            id: LogId.oidcFlow,
+            context: "mongodb-oidc-plugin:auth-succeeded",
+            message: "Authenticated successfully.",
+        });
+    }
+
+    private onOidcNotifyDeviceFlow(): void {
+        if (this.state.tag === "connecting" && this.state.connectionStringAuthType?.startsWith("oidc")) {
+            this.changeState("connection-requested", {
+                ...this.state,
+                tag: "connecting",
+                connectionStringAuthType: "oidc-device-flow",
+            });
+        }
+
+        this.logger.info({
+            id: LogId.oidcFlow,
+            context: "mongodb-oidc-plugin:notify-device-flow",
+            message: "OIDC Flow changed automatically to device flow.",
+        });
+    }
+
+    static inferConnectionTypeFromSettings(config: UserConfig, settings: ConnectionSettings): ConnectionStringAuthType {
         const connString = new ConnectionString(settings.connectionString);
         const searchParams = connString.typedSearchParams<MongoClientOptions>();
 
         switch (searchParams.get("authMechanism")) {
             case "MONGODB-OIDC": {
-                return "oidc-auth-flow"; // TODO: depending on if we don't have a --browser later it can be oidc-device-flow
+                if (config.transport === "stdio" && config.browser) {
+                    return "oidc-auth-flow";
+                }
+
+                if (config.transport === "http" && config.httpHost === "127.0.0.1" && config.browser) {
+                    return "oidc-auth-flow";
+                }
+
+                return "oidc-device-flow";
             }
             case "MONGODB-X509":
                 return "x.509";
diff --git a/src/common/logger.ts b/src/common/logger.ts
@@ -58,6 +58,10 @@ export const LogId = {
     exportedDataListError: mongoLogId(1_007_006),
     exportedDataAutoCompleteError: mongoLogId(1_007_007),
     exportLockError: mongoLogId(1_007_008),
+
+    oidcFlow: mongoLogId(1_008_001),
+    oidcClientError: mongoLogId(1_008_002),
+    oidcAuthFailed: mongoLogId(1_008_003),
 } as const;
 
 interface LogPayload {
diff --git a/tests/integration/common/connectionManager.test.ts b/tests/integration/common/connectionManager.test.ts
@@ -4,6 +4,7 @@ import {
     ConnectionStateConnected,
     ConnectionStringAuthType,
 } from "../../../src/common/connectionManager.js";
+import type { UserConfig } from "../../../src/common/config.js";
 import { describeWithMongoDB } from "../tools/mongodb/mongodbHelpers.js";
 import { describe, beforeEach, expect, it, vi, afterEach } from "vitest";
 
@@ -136,23 +137,52 @@ describeWithMongoDB("Connection Manager", (integration) => {
 
 describe("Connection Manager connection type inference", () => {
     const testCases = [
-        { connectionString: "mongodb://localhost:27017", connectionType: "scram" },
-        { connectionString: "mongodb://localhost:27017?authMechanism=MONGODB-X509", connectionType: "x.509" },
-        { connectionString: "mongodb://localhost:27017?authMechanism=GSSAPI", connectionType: "kerberos" },
+        { userConfig: {}, connectionString: "mongodb://localhost:27017", connectionType: "scram" },
         {
+            userConfig: {},
+            connectionString: "mongodb://localhost:27017?authMechanism=MONGODB-X509",
+            connectionType: "x.509",
+        },
+        {
+            userConfig: {},
+            connectionString: "mongodb://localhost:27017?authMechanism=GSSAPI",
+            connectionType: "kerberos",
+        },
+        {
+            userConfig: {},
             connectionString: "mongodb://localhost:27017?authMechanism=PLAIN&authSource=$external",
             connectionType: "ldap",
         },
-        { connectionString: "mongodb://localhost:27017?authMechanism=PLAIN", connectionType: "scram" },
-        { connectionString: "mongodb://localhost:27017?authMechanism=MONGODB-OIDC", connectionType: "oidc-auth-flow" },
+        { userConfig: {}, connectionString: "mongodb://localhost:27017?authMechanism=PLAIN", connectionType: "scram" },
+        {
+            userConfig: { transport: "stdio", browser: "firefox" },
+            connectionString: "mongodb://localhost:27017?authMechanism=MONGODB-OIDC",
+            connectionType: "oidc-auth-flow",
+        },
+        {
+            userConfig: { transport: "http", httpHost: "127.0.0.1", browser: "ie6" },
+            connectionString: "mongodb://localhost:27017?authMechanism=MONGODB-OIDC",
+            connectionType: "oidc-auth-flow",
+        },
+        {
+            userConfig: { transport: "http", httpHost: "0.0.0.0", browser: "ie6" },
+            connectionString: "mongodb://localhost:27017?authMechanism=MONGODB-OIDC",
+            connectionType: "oidc-device-flow",
+        },
+        {
+            userConfig: { transport: "stdio" },
+            connectionString: "mongodb://localhost:27017?authMechanism=MONGODB-OIDC",
+            connectionType: "oidc-device-flow",
+        },
     ] as {
+        userConfig: Partial<UserConfig>;
         connectionString: string;
         connectionType: ConnectionStringAuthType;
     }[];
 
-    for (const { connectionString, connectionType } of testCases) {
+    for (const { userConfig, connectionString, connectionType } of testCases) {
         it(`infers ${connectionType} from ${connectionString}`, () => {
-            const actualConnectionType = ConnectionManager.inferConnectionTypeFromSettings({
+            const actualConnectionType = ConnectionManager.inferConnectionTypeFromSettings(userConfig as UserConfig, {
                 connectionString,
             });
 
