diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 7ba445a2b..cd5413e84 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed repeated `/mcp reauth` commands getting stuck behind the previous unfinished MCP OAuth login; a new reauthorization now cancels and cleans up the prior flow before starting its replacement. + ## [17.2.9] - 2026-08-05 ### Breaking Changes diff --git a/packages/coding-agent/src/modes/controllers/mcp-command-controller.ts b/packages/coding-agent/src/modes/controllers/mcp-command-controller.ts index fa5e1f5d1..b741c5b29 100644 --- a/packages/coding-agent/src/modes/controllers/mcp-command-controller.ts +++ b/packages/coding-agent/src/modes/controllers/mcp-command-controller.ts @@ -89,6 +89,66 @@ function raceAbortSignal(promise: Promise, signal: AbortSignal, createErro }); } +type ActiveMCPOAuthFlow = { + cancel: (reason: string) => void; + completion: Promise; + complete: () => void; +}; + +type MCPOAuthFlowCoordinator = { + active?: ActiveMCPOAuthFlow; + transition: Promise; +}; + +const mcpOAuthFlowCoordinators = new WeakMap(); +const MCP_OAUTH_SUPERSEDED_REASON = "MCP OAuth flow superseded by a new login"; + +/** + * Serialize MCP OAuth ownership across slash-command controller instances. + * Interactive mode creates a new controller for every command, while the + * manual-input manager remains stable for the session and is therefore the + * lifecycle key. + */ +async function claimMCPOAuthFlow(owner: object, cancel: (reason: string) => void): Promise<{ release: () => void }> { + let coordinator = mcpOAuthFlowCoordinators.get(owner); + if (!coordinator) { + coordinator = { transition: Promise.resolve() }; + mcpOAuthFlowCoordinators.set(owner, coordinator); + } + + const precedingTransition = coordinator.transition; + const transition = Promise.withResolvers(); + coordinator.transition = transition.promise; + await precedingTransition; + + try { + const active = coordinator.active; + if (active) { + active.cancel(MCP_OAUTH_SUPERSEDED_REASON); + await active.completion; + } + + const completion = Promise.withResolvers(); + const flow: ActiveMCPOAuthFlow = { + cancel, + completion: completion.promise, + complete: () => completion.resolve(), + }; + coordinator.active = flow; + let released = false; + return { + release: () => { + if (released) return; + released = true; + if (coordinator.active === flow) coordinator.active = undefined; + flow.complete(); + }, + }; + } finally { + transition.resolve(); + } +} + /** * Minimum column budget for URL wrapping. Below this the terminal is * effectively unusable, but we still emit chunks so no character is silently @@ -782,28 +842,23 @@ export class MCPCommandController { const resolvedClientSecret = clientSecret.trim() || undefined; const manualInput = this.ctx.oauthManualInput; - if (manualInput.hasPending()) { - const pendingProvider = manualInput.pendingProviderId ?? "another provider"; - throw new Error( - `OAuth login already in progress for ${pendingProvider}. Complete or cancel it before starting MCP OAuth.`, - ); - } let manualInputClaim: { promise: Promise; clear: (reason?: string) => void } | undefined; const oauthTimeout = new AbortController(); - // User Esc and external aborts route through here; the timeout path sets - // its own reason and leaves this flag false so the catch can distinguish - // "user cancelled" (status) from "deadline elapsed" (error). - let userCancelled = false; - const requestUserCancel = (reason: string): void => { - userCancelled = true; + // Esc, external aborts, and a replacement MCP flow route through here; + // the timeout path sets its own reason and leaves this flag false so the + // catch can distinguish cancellation (status) from deadline failure. + let cancellationRequested = false; + const requestCancellation = (reason: string): void => { + cancellationRequested = true; if (!oauthTimeout.signal.aborted) oauthTimeout.abort(reason); }; + const flowClaim = await claimMCPOAuthFlow(manualInput, requestCancellation); const originalOnEscape = this.ctx.editor.onEscape; - this.ctx.editor.onEscape = () => requestUserCancel(MCP_OAUTH_USER_CANCEL_REASON); + this.ctx.editor.onEscape = () => requestCancellation(MCP_OAUTH_USER_CANCEL_REASON); const externalSignal = opts?.abortSignal; const onExternalAbort = (): void => { const reason = externalSignal?.reason; - requestUserCancel(typeof reason === "string" ? reason : MCP_OAUTH_USER_CANCEL_REASON); + requestCancellation(typeof reason === "string" ? reason : MCP_OAUTH_USER_CANCEL_REASON); }; if (externalSignal?.aborted) { onExternalAbort(); @@ -811,6 +866,12 @@ export class MCPCommandController { externalSignal?.addEventListener("abort", onExternalAbort, { once: true }); } try { + if (manualInput.hasPending()) { + const pendingProvider = manualInput.pendingProviderId ?? "another provider"; + throw new Error( + `OAuth login already in progress for ${pendingProvider}. Complete or cancel it before starting MCP OAuth.`, + ); + } // Create OAuth flow const flow = new MCPOAuthFlow( { @@ -887,7 +948,7 @@ export class MCPCommandController { const createAbortError = (): Error => { const reason = String(oauthTimeout.signal.reason ?? "MCP OAuth flow aborted"); - return userCancelled ? new MCPOAuthCancelledError() : new Error(reason); + return cancellationRequested ? new MCPOAuthCancelledError() : new Error(reason); }; if (oauthTimeout.signal.aborted) throw createAbortError(); @@ -934,11 +995,10 @@ export class MCPCommandController { resource: flow.resource, }; } catch (error) { - // User-initiated cancel (Esc or external signal) → neutral status, not - // a failure. Check the flag we set in `requestUserCancel`, not the - // abort reason: the timeout path also aborts but with a different - // reason, and we want it to surface as a timeout error below. - if (userCancelled) { + // Esc, an external abort, or a newer MCP flow are neutral + // cancellations. The timeout path also aborts the controller but does + // not set this flag, so it remains a surfaced error. + if (cancellationRequested) { throw new MCPOAuthCancelledError(); } @@ -960,6 +1020,7 @@ export class MCPCommandController { this.ctx.editor.onEscape = originalOnEscape; externalSignal?.removeEventListener("abort", onExternalAbort); manualInputClaim?.clear("Manual MCP OAuth input cleared"); + flowClaim.release(); } } diff --git a/packages/coding-agent/test/mcp-command-reauth.test.ts b/packages/coding-agent/test/mcp-command-reauth.test.ts index 3f3f98e2f..8f137291a 100644 --- a/packages/coding-agent/test/mcp-command-reauth.test.ts +++ b/packages/coding-agent/test/mcp-command-reauth.test.ts @@ -8,6 +8,7 @@ import * as mcpClient from "@oh-my-pi/pi-coding-agent/mcp/client"; import * as oauthFlow from "@oh-my-pi/pi-coding-agent/mcp/oauth-flow"; import type { MCPServerConfig } from "@oh-my-pi/pi-coding-agent/mcp/types"; import { MCPCommandController } from "@oh-my-pi/pi-coding-agent/modes/controllers/mcp-command-controller"; +import { OAuthManualInputManager } from "@oh-my-pi/pi-coding-agent/modes/oauth-manual-input"; import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; import { getConfigRootDir, @@ -56,7 +57,8 @@ function createController(authStorage: AuthStorage, mcpManagerOverrides: Record< getConnectionStatus: vi.fn(() => "connected"), ...mcpManagerOverrides, }; - const controller = new MCPCommandController({ + const oauthManualInput = new OAuthManualInputManager(); + const ctx = { chatContainer: { addChild: vi.fn() }, present, presentCommandOutput: present, @@ -64,11 +66,7 @@ function createController(authStorage: AuthStorage, mcpManagerOverrides: Record< editor, showError, showStatus, - oauthManualInput: { - hasPending: vi.fn(() => false), - pendingProviderId: undefined, - tryClaimInput: vi.fn(), - }, + oauthManualInput, settings: { get: vi.fn((_key: string): unknown => undefined), }, @@ -78,9 +76,10 @@ function createController(authStorage: AuthStorage, mcpManagerOverrides: Record< modelRegistry: { authStorage }, }, mcpManager, - } as never); + } as never; + const controller = new MCPCommandController(ctx); - return { controller, showError, showStatus, present, editor, prepareConfig, mcpManager }; + return { controller, ctx, showError, showStatus, present, editor, oauthManualInput, prepareConfig, mcpManager }; } describe("/mcp auth commands", () => { @@ -446,6 +445,65 @@ describe("/mcp auth commands", () => { expect(editor.onEscape).not.toBe(installedEscape); }); + test("reauth supersedes an unfinished MCP OAuth flow", async () => { + const authStorage = freshAuthStorage(); + await authStorage.reload(); + vi.spyOn(mcpClient, "connectToServer").mockRejectedValue(AUTH_ERROR); + let loginAttempt = 0; + vi.spyOn(oauthFlow.MCPOAuthFlow.prototype, "login").mockImplementation(function (this: oauthFlow.MCPOAuthFlow) { + loginAttempt += 1; + if (loginAttempt > 1) { + return Promise.resolve({ + access: "replacement-access", + refresh: "replacement-refresh", + expires: Date.now() + 3_600_000, + }); + } + + const manualInputPromise = this.ctrl.onManualCodeInput?.(); + void manualInputPromise?.catch(() => {}); + const pending = Promise.withResolvers(); + this.ctrl.signal?.addEventListener( + "abort", + () => { + pending.reject(new Error(`OAuth callback cancelled: ${String(this.ctrl.signal?.reason ?? "aborted")}`)); + }, + { once: true }, + ); + return pending.promise; + }); + + const { controller, ctx, showError, showStatus, editor, oauthManualInput } = createController(authStorage); + const firstReauth = controller.handle("/mcp reauth envserver"); + const claimDeadline = Date.now() + 1_000; + while (!oauthManualInput.hasPending() && Date.now() < claimDeadline) { + await Bun.sleep(10); + } + expect(oauthManualInput.pendingProviderId).toBe("mcp"); + + const replacementReauth = new MCPCommandController(ctx).handle("/mcp reauth envserver"); + await Promise.race([ + replacementReauth, + Bun.sleep(2_000).then(() => { + throw new Error("replacement reauth did not resolve within 2s"); + }), + ]); + if (oauthManualInput.hasPending()) editor.onEscape?.(); + await Promise.race([ + firstReauth, + Bun.sleep(2_000).then(() => { + throw new Error("superseded reauth did not resolve within 2s"); + }), + ]); + + expect(loginAttempt).toBe(2); + expect(showError).not.toHaveBeenCalled(); + expect(showStatus).toHaveBeenCalledWith(expect.stringMatching(/cancel/i)); + expect(authStorage.get(oauthFlow.mcpOAuthCredentialId(EXPANDED_SERVER_URL))).toMatchObject({ + access: "replacement-access", + }); + }); + test("Esc cancels even when OAuth login has not registered its signal listener yet", async () => { const authStorage = freshAuthStorage(); await authStorage.reload();