From 441e2d3dc056b197b22d836b5ca1094e3209fde5 Mon Sep 17 00:00:00 2001 From: AvivGuiser Date: Sun, 21 Jun 2026 22:59:32 +0300 Subject: [PATCH] fix(connect): surface OIDC device-flow URL and code from connect tool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit For OIDC connection strings, MCPConnectionManager.connect() returns a `connecting` state immediately — before the device-flow callback has populated the verification URL and user code. The connect tool ignored that state and always reported "Successfully connected to MongoDB.", so the device-flow URL+code only surfaced later when the first data operation failed with NotConnectedToMongoDB. Make the connect tool wait for the OIDC attempt to make progress and return the verification URL+code directly when a device flow is reported, so the user is asked to authenticate up front instead of after a failing data op. - Add waitForConnectResult: a platform-independent helper that resolves on the first of device-flow notification / success / error / timeout. Kept dependency-free so it is unit-testable on any platform (the real OIDC integration path is Linux-only). - Extract the device-flow message into oidcDeviceFlowMessage so the connect tool and connectionErrorHandler share one wording. - Add unit tests for the wait helper and assert the connect tool's own response carries the URL+code in the Linux device-flow integration test. Fixes #1269 Signed-off-by: AvivGuiser --- src/common/connectionErrorHandler.ts | 3 +- src/common/oidcDeviceFlowMessage.ts | 11 ++ src/common/waitForConnectResult.ts | 116 +++++++++++++ src/tools/mongodb/connect/connect.ts | 23 +++ .../common/connectionManager.oidc.test.ts | 35 ++++ .../unit/common/waitForConnectResult.test.ts | 156 ++++++++++++++++++ 6 files changed, 343 insertions(+), 1 deletion(-) create mode 100644 src/common/oidcDeviceFlowMessage.ts create mode 100644 src/common/waitForConnectResult.ts create mode 100644 tests/unit/common/waitForConnectResult.test.ts diff --git a/src/common/connectionErrorHandler.ts b/src/common/connectionErrorHandler.ts index 9810d47dc..58ca65c5d 100644 --- a/src/common/connectionErrorHandler.ts +++ b/src/common/connectionErrorHandler.ts @@ -2,6 +2,7 @@ import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { ErrorCodes, type MongoDBError } from "./errors.js"; import type { AnyConnectionState } from "./connectionManager.js"; import type { AnyToolBase } from "../tools/tool.js"; +import { oidcDeviceFlowMessage } from "./oidcDeviceFlowMessage.js"; export type ConnectionErrorHandler = ( error: MongoDBError, @@ -47,7 +48,7 @@ export const connectionErrorHandler: ConnectionErrorHandler = (error, { availabl if (connectionState.tag === "connecting" && connectionState.oidcConnectionType) { additionalPromptForConnectivity.push({ type: "text", - text: `The user needs to finish their OIDC connection by opening '${connectionState.oidcLoginUrl}' in the browser and use the following user code: '${connectionState.oidcUserCode}'`, + text: oidcDeviceFlowMessage(connectionState.oidcLoginUrl, connectionState.oidcUserCode), }); } else { additionalPromptForConnectivity.push({ diff --git a/src/common/oidcDeviceFlowMessage.ts b/src/common/oidcDeviceFlowMessage.ts new file mode 100644 index 000000000..c1e0fe11a --- /dev/null +++ b/src/common/oidcDeviceFlowMessage.ts @@ -0,0 +1,11 @@ +/** + * Builds the user-facing instruction shown when an OIDC connection needs to be + * finished via the device-authorization flow. + * + * Shared between the `connect` tool (which surfaces it directly in its response) + * and the connection-error handler (which appends it when a data operation runs + * before authentication completed) so the two surfacings cannot drift. + */ +export function oidcDeviceFlowMessage(oidcLoginUrl: string | undefined, oidcUserCode: string | undefined): string { + return `The user needs to finish their OIDC connection by opening '${oidcLoginUrl}' in the browser and use the following user code: '${oidcUserCode}'`; +} diff --git a/src/common/waitForConnectResult.ts b/src/common/waitForConnectResult.ts new file mode 100644 index 000000000..658209e14 --- /dev/null +++ b/src/common/waitForConnectResult.ts @@ -0,0 +1,116 @@ +import type { EventEmitter } from "events"; +import type { AnyConnectionState, ConnectionManagerEvents } from "./connectionManager.js"; + +/** + * Outcome of waiting for an OIDC connection attempt to make progress. + * + * - `connected`: authentication completed (e.g. a cached/refreshed token), the + * connection is fully established. + * - `device-flow`: the OIDC plugin reported a device-authorization flow; the + * user must visit `oidcLoginUrl` and enter `oidcUserCode` to finish auth. + * - `errored`: the connection attempt failed. + * - `timed-out`: nothing terminal happened within the allotted time; the + * attempt is still in-flight. + */ +export type WaitForConnectResult = + | { kind: "connected"; state: AnyConnectionState } + | { kind: "device-flow"; state: AnyConnectionState; oidcLoginUrl: string; oidcUserCode: string } + | { kind: "errored"; state: AnyConnectionState } + | { kind: "timed-out"; state: AnyConnectionState }; + +const DEFAULT_TIMEOUT_MS = 10_000; + +/** + * Waits for an in-flight OIDC connection attempt to reach a point where the + * `connect` tool can report something useful back to the user. + * + * For OIDC, {@link ConnectionManager.connect} returns a `connecting` state + * immediately — before the device-flow callback has populated the verification + * URL and user code. This helper bridges that gap: it subscribes to the + * connection manager's events and resolves on the first of: + * + * 1. a device-flow notification (`connecting` state carrying `oidcLoginUrl`), + * 2. a successful connection, + * 3. an error / time-out, or + * 4. the local timeout below. + * + * It deliberately resolves as soon as the device-flow URL+code are known rather + * than waiting for `connection-success` — at that point the user has not yet + * authenticated, so blocking for success would hang until they finish the + * browser step. + * + * Kept free of any MongoDB/transport dependency (it only needs an event emitter + * and a current-state getter) so it can be unit-tested on any platform with a + * plain `EventEmitter`. + */ +export function waitForConnectResult({ + events, + getCurrentState, + timeoutMs = DEFAULT_TIMEOUT_MS, +}: { + events: Pick, "on" | "off">; + getCurrentState: () => AnyConnectionState; + timeoutMs?: number; +}): Promise { + return new Promise((resolve) => { + let settled = false; + + const classify = (state: AnyConnectionState): WaitForConnectResult | undefined => { + switch (state.tag) { + case "connected": + return { kind: "connected", state }; + case "errored": + return { kind: "errored", state }; + case "connecting": + if (state.oidcLoginUrl && state.oidcUserCode) { + return { + kind: "device-flow", + state, + oidcLoginUrl: state.oidcLoginUrl, + oidcUserCode: state.oidcUserCode, + }; + } + return undefined; + default: + return undefined; + } + }; + + const finish = (result: WaitForConnectResult): void => { + if (settled) { + return; + } + settled = true; + clearTimeout(timer); + events.off("connection-request", onState); + events.off("connection-success", onState); + events.off("connection-error", onState); + events.off("connection-time-out", onState); + resolve(result); + }; + + const onState = (state: AnyConnectionState): void => { + const result = classify(state); + if (result) { + finish(result); + } + }; + + events.on("connection-request", onState); + events.on("connection-success", onState); + events.on("connection-error", onState); + events.on("connection-time-out", onState); + + const timer = setTimeout(() => finish({ kind: "timed-out", state: getCurrentState() }), timeoutMs); + // Don't let a pending timer keep the process alive. + timer.unref?.(); + + // Re-check the current state immediately after subscribing to close the + // race where a terminal event fires between the caller's state read and + // our subscription. + const current = classify(getCurrentState()); + if (current) { + finish(current); + } + }); +} diff --git a/src/tools/mongodb/connect/connect.ts b/src/tools/mongodb/connect/connect.ts index 3e2fa78b8..bef8023d4 100644 --- a/src/tools/mongodb/connect/connect.ts +++ b/src/tools/mongodb/connect/connect.ts @@ -2,6 +2,8 @@ import { z } from "zod"; import { MongoDBToolBase } from "../mongodbTool.js"; import type { ToolArgs, OperationType, ToolConstructorParams, ToolResult } from "../../tool.js"; import type { Server } from "../../../server.js"; +import { waitForConnectResult } from "../../../common/waitForConnectResult.js"; +import { oidcDeviceFlowMessage } from "../../../common/oidcDeviceFlowMessage.js"; const ConnectOutputSchema = { connected: z.boolean(), @@ -50,6 +52,27 @@ export class ConnectTool extends MongoDBToolBase { }: ToolArgs): Promise> { await this.session.connectToMongoDB({ connectionString }); + const connectionManager = this.session.connectionManager; + + // For OIDC, `connectToMongoDB` resolves while still in the `connecting` + // state — before the device-flow callback has populated the verification + // URL and user code. Wait for the attempt to make progress so we can + // surface those to the user directly, instead of reporting success and + // having the next data operation fail asking them to authenticate. + if (connectionManager.currentConnectionState.tag === "connecting") { + const result = await waitForConnectResult({ + events: connectionManager.events, + getCurrentState: () => connectionManager.currentConnectionState, + }); + + if (result.kind === "device-flow") { + return { + content: [{ type: "text", text: oidcDeviceFlowMessage(result.oidcLoginUrl, result.oidcUserCode) }], + structuredContent: { connected: false }, + }; + } + } + return { content: [{ type: "text", text: "Successfully connected to MongoDB." }], structuredContent: { connected: true }, diff --git a/tests/integration/common/connectionManager.oidc.test.ts b/tests/integration/common/connectionManager.oidc.test.ts index 4723c412f..72df6396b 100644 --- a/tests/integration/common/connectionManager.oidc.test.ts +++ b/tests/integration/common/connectionManager.oidc.test.ts @@ -291,6 +291,41 @@ describe.skipIf(process.platform !== "linux")("ConnectionManager OIDC Tests", as expect(connectedResponse).toContain("config"); expect(connectedResponse).toContain("local"); }); + + it("surfaces the device-flow URL and code directly from the connect tool", async ({ + signal, + skip, + }, integration) => { + // Reset so we can drive a fresh connect through the tool and inspect its response. + const connectionManager = integration.mcpServer().session + .connectionManager as TestConnectionManager; + await connectionManager.disconnect(); + connectionManager.changeState("connection-close", { tag: "disconnected" }); + + const connectResponse = responseAsText( + await integration.mcpClient().callTool({ + name: "connect", + arguments: { connectionString: integration.connectionString() }, + }) + ); + + // If the IdP completed auth without ever falling back to device flow, there is + // nothing to assert here; the dedicated test above covers the device-flow path. + if (!connectResponse.includes("The user needs to finish their OIDC connection by opening")) { + return skip(); + } + + const state = await waitUntil( + "connecting", + connectionManager, + signal, + (state) => !!state.oidcLoginUrl && !!state.oidcUserCode + ); + + expect(connectResponse).toContain(state.oidcLoginUrl); + expect(connectResponse).toContain(state.oidcUserCode); + expect(connectResponse).not.toContain("Successfully connected to MongoDB."); + }); } ); } diff --git a/tests/unit/common/waitForConnectResult.test.ts b/tests/unit/common/waitForConnectResult.test.ts new file mode 100644 index 000000000..1e530579b --- /dev/null +++ b/tests/unit/common/waitForConnectResult.test.ts @@ -0,0 +1,156 @@ +import { EventEmitter } from "events"; +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { waitForConnectResult } from "../../../src/common/waitForConnectResult.js"; +import type { + AnyConnectionState, + ConnectionManagerEvents, + ConnectionStateConnecting, +} from "../../../src/common/connectionManager.js"; + +/** + * Minimal harness mimicking the connection manager's event bus + current-state + * getter, so we can exercise `waitForConnectResult` without a real MongoDB + * connection (the real OIDC path is Linux-only and skipped on other platforms). + */ +function makeHarness(initial: AnyConnectionState): { + events: EventEmitter; + getCurrentState: () => AnyConnectionState; + changeState: (event: keyof ConnectionManagerEvents, state: AnyConnectionState) => void; +} { + const events = new EventEmitter(); + let state = initial; + const emit = events.emit.bind(events) as ( + event: keyof ConnectionManagerEvents, + payload: AnyConnectionState + ) => void; + return { + events, + getCurrentState: (): AnyConnectionState => state, + changeState: (event: keyof ConnectionManagerEvents, next: AnyConnectionState): void => { + state = next; + emit(event, next); + }, + }; +} + +const connectingState: ConnectionStateConnecting = { + tag: "connecting", + serviceProvider: Promise.resolve() as unknown as ConnectionStateConnecting["serviceProvider"], + oidcConnectionType: "oidc-auth-flow", + connectionStringInfo: { authType: "oidc-auth-flow", hostType: "unknown" }, +}; + +describe("waitForConnectResult", () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + it("resolves with the device-flow URL and code when the OIDC plugin notifies", async () => { + const harness = makeHarness(connectingState); + + const pending = waitForConnectResult({ + events: harness.events, + getCurrentState: harness.getCurrentState, + }); + + harness.changeState("connection-request", { + ...connectingState, + oidcLoginUrl: "https://idp.example/device", + oidcUserCode: "ABCD-EFGH", + }); + + const result = await pending; + expect(result).toMatchObject({ + kind: "device-flow", + oidcLoginUrl: "https://idp.example/device", + oidcUserCode: "ABCD-EFGH", + }); + expect(result.state.tag).toBe("connecting"); + }); + + it("resolves as connected when authentication succeeds without a device flow", async () => { + const harness = makeHarness(connectingState); + + const pending = waitForConnectResult({ + events: harness.events, + getCurrentState: harness.getCurrentState, + }); + + const connected = { tag: "connected" } as unknown as AnyConnectionState; + harness.changeState("connection-success", connected); + + const result = await pending; + expect(result.kind).toBe("connected"); + }); + + it("resolves as errored when the connection attempt fails", async () => { + const harness = makeHarness(connectingState); + + const pending = waitForConnectResult({ + events: harness.events, + getCurrentState: harness.getCurrentState, + }); + + harness.changeState("connection-error", { + tag: "errored", + errorReason: "boom", + } as AnyConnectionState); + + const result = await pending; + expect(result.kind).toBe("errored"); + }); + + it("times out when nothing terminal happens", async () => { + const harness = makeHarness(connectingState); + + const pending = waitForConnectResult({ + events: harness.events, + getCurrentState: harness.getCurrentState, + timeoutMs: 5_000, + }); + + await vi.advanceTimersByTimeAsync(5_000); + + const result = await pending; + expect(result.kind).toBe("timed-out"); + }); + + it("short-circuits on an already-terminal current state (subscribe race)", async () => { + const harness = makeHarness({ + ...connectingState, + oidcLoginUrl: "https://idp.example/device", + oidcUserCode: "WXYZ-1234", + }); + + const result = await waitForConnectResult({ + events: harness.events, + getCurrentState: harness.getCurrentState, + }); + + expect(result).toMatchObject({ kind: "device-flow", oidcUserCode: "WXYZ-1234" }); + }); + + it("ignores intermediate connecting states without a URL yet", async () => { + const harness = makeHarness(connectingState); + + const pending = waitForConnectResult({ + events: harness.events, + getCurrentState: harness.getCurrentState, + }); + + // A connecting event with no URL yet should not resolve the promise. + harness.changeState("connection-request", connectingState); + harness.changeState("connection-request", { + ...connectingState, + oidcLoginUrl: "https://idp.example/device", + oidcUserCode: "LATE-CODE", + }); + + const result = await pending; + expect(result).toMatchObject({ kind: "device-flow", oidcUserCode: "LATE-CODE" }); + }); +});