From f10993bc7f28145df1375f3ff39c7f5c4cfc05f0 Mon Sep 17 00:00:00 2001 From: Armin Ronacher Date: Wed, 7 Oct 2026 16:17:59 +0200 Subject: [PATCH] fix(coding-agent): cancel MCP OAuth sign-ins and stop refreshing on close The OAuth flow in pi-mcp takes a signal for every request, and an aborted refresh no longer falls back to a new authorization. The streamable HTTP transport closes its session with the last request's token instead of asking the auth provider, which could refresh first. The MCP extension owns one signal per session: it aborts sign-ins on Esc and on session shutdown, replaces the generation counter, and shutdown waits for aborted sign-ins to clean up. Requests to the authorization server time out after 15s, and pi mcp login --timeout limits the whole sign-in. closes #10565 --- packages/coding-agent/CHANGELOG.md | 3 + .../coding-agent/src/extensions/mcp/cli.ts | 45 +++---- .../coding-agent/src/extensions/mcp/index.ts | 118 +++++++++++------- .../coding-agent/src/extensions/mcp/oauth.ts | 48 +++++-- .../coding-agent/src/extensions/mcp/ui.ts | 15 ++- .../test/mcp-manager-view.test.ts | 40 ++++++ .../suite/agent-session-mcp-oauth.test.ts | 76 ++++++++++- .../test/suite/mcp-oauth-server.ts | 19 ++- packages/mcp/CHANGELOG.md | 8 ++ packages/mcp/src/auth-provider.ts | 4 + packages/mcp/src/oauth/discovery.ts | 42 ++++++- packages/mcp/src/oauth/flow.ts | 18 ++- .../mcp/src/transports/streamable-http.ts | 31 +++-- packages/mcp/test/oauth.test.ts | 40 ++++++ packages/mcp/test/streamable-http.test.ts | 20 +++ 15 files changed, 414 insertions(+), 113 deletions(-) create mode 100644 packages/coding-agent/test/mcp-manager-view.test.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index d887dbf94..a2d1c4143 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -13,6 +13,7 @@ ### Changed - Changed `outputPad` to also apply to `!` command output, tool output, and summary blocks ([#9946](https://github.com/earendil-works/pi/issues/9946), [#10557](https://github.com/earendil-works/pi/pull/10557) by [@rwachtler](https://github.com/rwachtler)) +- Changed `pi mcp login --timeout` to limit the whole sign-in, including requests to the authorization server, instead of only the wait for the browser ([#10565](https://github.com/earendil-works/pi/issues/10565)) ### Fixed @@ -26,6 +27,8 @@ - Fixed images being dropped as "could not be resized" when running under `node --watch` on Node 24.19+ and 26.x, where Node posts its own messages on the image resize worker channel ([#10527](https://github.com/earendil-works/pi/issues/10527)) - Fixed clipboard paste doing nothing in Termux, and failed copies there omitting the Termux:API install hint ([#10391](https://github.com/earendil-works/pi/issues/10391)) - Fixed `!` and RPC `bash` output keeping fragments of color codes, such as a stray `m`, when a code was split across output chunks ([#10504](https://github.com/earendil-works/pi/issues/10504)) +- Fixed MCP OAuth sign-ins that could not be cancelled while waiting on the authorization server and kept running after the session ended. The sign-in screen now cancels with Esc at every step, session shutdown aborts a running sign-in, and each request to the authorization server times out after 15 seconds ([#10565](https://github.com/earendil-works/pi/issues/10565)) +- Fixed shutdown waiting up to 15 seconds to refresh an MCP OAuth token that was about to expire, only to close the server's session ([#10565](https://github.com/earendil-works/pi/issues/10565)) ## [1.0.4] - 2026-10-05 diff --git a/packages/coding-agent/src/extensions/mcp/cli.ts b/packages/coding-agent/src/extensions/mcp/cli.ts index 9c7c38f63..951d77872 100644 --- a/packages/coding-agent/src/extensions/mcp/cli.ts +++ b/packages/coding-agent/src/extensions/mcp/cli.ts @@ -554,8 +554,9 @@ async function login( log(`Sign in to MCP server "${name}" in your browser:\n${authorizationUrl.href}`); openUrl(authorizationUrl.href); }, - promptForRedirectUrl: (signal) => waitForRedirectUrl(signal, timeoutMs, interactive), + promptForRedirectUrl: (signal) => waitForRedirectUrl(signal, interactive), }, + signal: AbortSignal.timeout(timeoutMs), }); } catch (signInError) { error( @@ -578,37 +579,21 @@ async function login( /** * The pasted redirect URL in a terminal; otherwise only the browser callback can finish the sign-in. - * Resolves to undefined (cancelling the sign-in) after `timeoutMs`, or when the callback arrived. + * Resolves to undefined when `signal` aborts: the callback arrived, or the sign-in timed out. */ -async function waitForRedirectUrl( - signal: AbortSignal, - timeoutMs: number, - interactive: boolean, -): Promise { - const controller = new AbortController(); - const abort = () => controller.abort(); - signal.addEventListener("abort", abort, { once: true }); - const timer = setTimeout(abort, timeoutMs); +async function waitForRedirectUrl(signal: AbortSignal, interactive: boolean): Promise { + if (!interactive) { + if (!signal.aborted) await new Promise((resolve) => signal.addEventListener("abort", resolve, { once: true })); + return undefined; + } + const readline = createInterface({ input: process.stdin, output: process.stderr }); try { - if (!interactive) { - await new Promise((resolve) => - controller.signal.addEventListener("abort", () => resolve(), { once: true }), - ); - return undefined; - } - const readline = createInterface({ input: process.stdin, output: process.stderr }); - try { - return await readline.question( - "If the browser cannot reach this machine, paste the URL it was redirected to: ", - { signal: controller.signal }, - ); - } catch { - return undefined; - } finally { - readline.close(); - } + return await readline.question("If the browser cannot reach this machine, paste the URL it was redirected to: ", { + signal, + }); + } catch { + return undefined; } finally { - clearTimeout(timer); - signal.removeEventListener("abort", abort); + readline.close(); } } diff --git a/packages/coding-agent/src/extensions/mcp/index.ts b/packages/coding-agent/src/extensions/mcp/index.ts index 391c271b4..43e2ecebe 100644 --- a/packages/coding-agent/src/extensions/mcp/index.ts +++ b/packages/coding-agent/src/extensions/mcp/index.ts @@ -287,8 +287,6 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension let projectConfig: string | undefined; /** Registered servers that `mcp.json` overrides, shown in `/mcp`. */ let overridden: string[] = []; - /** Between session_start and session_shutdown. Registrations before that are read on session_start. */ - let sessionActive = false; let autoEnableCodemode = true; /** Whether the "codemode tools unreachable" warning was shown since the session started. */ let warnedUnreachable = false; @@ -296,8 +294,6 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension /** Whether a prompt already waited for the startup connections since the session started. */ let waitedForStartup = false; const startupWaitMs = options.startupWaitMs ?? DEFAULT_STARTUP_WAIT_MS; - /** Bumped on every session start and shutdown so a runtime load that resolves late is dropped. */ - let generation = 0; /** Working directory of the session, for stdio servers. */ let sessionCwd = process.cwd(); let credentials = options.credentials; @@ -312,8 +308,26 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension override: entry.override !== undefined, })); - /** Manager actions still opening or closing connections; shutdown waits for their cleanup. */ + /** + * Lifetime of the current session: aborted on session_shutdown, and before the first session_start + * (registrations before that are read on session_start). Work captures `session.signal` when it + * starts and drops late results once it aborted, for example a runtime load that resolves after + * shutdown. Work that outlives the command that started it also stops on the signal and is + * registered with `track`, so shutdown can wait for its cleanup. + */ + let session = new AbortController(); + session.abort(); + /** Work registered with `track`: manager actions opening or closing connections, and sign-ins. */ const backgroundActions = new Set>(); + const track = (work: Promise): Promise => { + const task = work.then( + () => undefined, + () => undefined, + ); + backgroundActions.add(task); + void task.finally(() => backgroundActions.delete(task)); + return work; + }; const listeners = new Set<() => void>(); const emitChange = () => { for (const listener of listeners) listener(); @@ -574,11 +588,11 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension * meanwhile. The returned promise rejects when the MCP runtime cannot be loaded. */ const startConnection = (server: McpServer, after?: Promise): Promise => { - const current = generation; + const { signal } = session; const attempt = Symbol(); server.attempt = attempt; const isCurrent = () => - current === generation && server.attempt === attempt && isEnabled(server) && servers.includes(server); + !signal.aborted && server.attempt === attempt && isEnabled(server) && servers.includes(server); const ready = (async () => { await after; if (!isCurrent()) return; @@ -655,19 +669,28 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension return undefined; }; - const signIn = async (server: McpServer, prompt: McpSignInPrompt): Promise => { + /** Sign in, then reconnect. Returns an error message on failure. `cancel` or the session's end stops the sign-in. */ + const signIn = async ( + server: McpServer, + prompt: McpSignInPrompt, + cancel?: AbortSignal, + ): Promise => { const connection = server.connection; const url = connection?.oauthUrl; if (!connection || !url) return `MCP server "${server.entry.name}" does not use OAuth.`; const runtime = await loadMcpRuntime(); + const signal = cancel ? AbortSignal.any([session.signal, cancel]) : session.signal; try { - await runtime.signInMcpServer({ - serverUrl: url, - store: getCredentials(runtime).forServer(server.entry.name, url), - settings: connection.oauthSettings(), - challenge: connection.challenge, - prompt, - }); + await track( + runtime.signInMcpServer({ + serverUrl: url, + store: getCredentials(runtime).forServer(server.entry.name, url), + settings: connection.oauthSettings(), + challenge: connection.challenge, + prompt, + signal, + }), + ); } catch (error) { if (error instanceof runtime.McpSignInCancelledError) return "Sign-in cancelled."; return `Sign-in failed: ${errorMessage(error)}`; @@ -891,22 +914,28 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension return setExposure(server, choice as McpExposure); }; - /** Sign in with the manager view's sign-in screen, which shows the URL with a copy key. */ + /** Sign in with the manager view's sign-in screen, which shows the URL with a copy key and cancels with Esc. */ const signInWithUi = (ui: McpUi, server: McpServer): Promise => { const title = `Sign in to ${server.entry.name}`; + const cancel = new AbortController(); + const status = (message: string) => ui.status(title, message, () => cancel.abort()); let authorizationUrl = ""; - ui.status(title, "Contacting the authorization server…"); - return signIn(server, { - showAuthorizationUrl: (url) => { - authorizationUrl = url.href; - openUrl(url.href); + status("Contacting the authorization server…"); + return signIn( + server, + { + showAuthorizationUrl: (url) => { + authorizationUrl = url.href; + openUrl(url.href); + }, + promptForRedirectUrl: async (signal) => { + const value = await ui.redirectUrl(title, authorizationUrl, signal); + status("Connecting…"); + return value; + }, }, - promptForRedirectUrl: async (signal) => { - const value = await ui.redirectUrl(title, authorizationUrl, signal); - ui.status(title, "Connecting…"); - return value; - }, - }); + cancel.signal, + ); }; /** Keep the subscribed menu usable while a connection opens or closes. */ @@ -917,20 +946,18 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension ) => { const ready = operation(); const attempt = server.attempt; - const current = generation; + const { signal } = session; const finish = (message: string | undefined) => { - if (current !== generation || server.attempt !== attempt || !servers.includes(server)) return; + if (signal.aborted || server.attempt !== attempt || !servers.includes(server)) return; server.message = message; ensureDiscoveryActive(ctx); emitChange(); }; - const task = ready - .then(finish, (error: unknown) => finish(errorMessage(error))) - .finally(() => backgroundActions.delete(task)); - backgroundActions.add(task); + void track(ready.then(finish, (error: unknown) => finish(errorMessage(error)))); }; const runAction = async (ui: McpUi, ctx: ExtensionContext, server: McpServer, action: string) => { + const { signal } = session; let message: string | undefined; switch (action) { case "signin": @@ -961,6 +988,8 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension break; } } + // The session ended meanwhile (for example during a sign-in), which made ctx stale. + if (signal.aborted) return; server.message = message; ensureDiscoveryActive(ctx); emitChange(); @@ -1049,6 +1078,7 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension ctx.ui.notify(`Signing in to MCP server "${name}" requires interactive mode.`, "error"); return; } + const { signal } = session; let failure: string | undefined; if (ctx.mode === "tui") { await showMcpManager(ctx, async (ui) => { @@ -1068,6 +1098,8 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension ), }); } + // The session ended meanwhile, which cancelled the sign-in and made ctx stale. + if (signal.aborted) return; if (failure) { ctx.ui.notify(failure, failure === "Sign-in cancelled." ? "info" : "error"); return; @@ -1085,8 +1117,8 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension waitedForStartup = false; sessionCwd = ctx.cwd; modelRegistry = ctx.modelRegistry; - const current = ++generation; - sessionActive = true; + session = new AbortController(); + const { signal } = session; configuredEntries = loaded.servers; const registered = registeredServers(); overridden = registered.overridden; @@ -1103,10 +1135,9 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension // The MCP client loads only now, so sessions without servers never pay for it. Waiting one // event loop turn lets the first render happen before loading and connecting. const runtime = new Promise((resolve) => setImmediate(resolve)).then(() => loadMcpRuntime()); - const isCurrent = () => current === generation; pending = Promise.all(enabled.map((server) => startConnection(server, runtime))) .then(() => { - if (isCurrent()) reportProblems(ctx); + if (!signal.aborted) reportProblems(ctx); }) .catch((error: unknown) => { // The session may have been disposed meanwhile, which makes ctx stale. @@ -1181,8 +1212,8 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension // Servers registered or unregistered during the session connect or disconnect right away. pi.on("mcp_servers_change", async (_event, ctx) => { - if (!sessionActive) return; - const current = generation; + const { signal } = session; + if (signal.aborted) return; const registered = registeredServers(); overridden = registered.overridden; const next = new Map(registered.servers.map((server) => [server.entry.name, server])); @@ -1200,10 +1231,10 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension ensureDiscoveryActive(ctx); await Promise.all(removed.map((server) => server.connection?.close())); const connecting = added.filter(isEnabled); - if (current !== generation || connecting.length === 0) return; + if (signal.aborted || connecting.length === 0) return; try { await Promise.all(connecting.map((server) => startConnection(server))); - if (current !== generation) { + if (signal.aborted) { await Promise.all(connecting.map((server) => server.connection?.close())); return; } @@ -1211,13 +1242,12 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension ctx.ui.notify(`MCP failed to load: ${errorMessage(error)}`, "error"); return; } - if (current !== generation) return; + if (signal.aborted) return; reportProblems(ctx, connecting); }); pi.on("session_shutdown", async () => { - sessionActive = false; - generation++; + session.abort(); const closing = connections(); servers = []; emitChange(); diff --git a/packages/coding-agent/src/extensions/mcp/oauth.ts b/packages/coding-agent/src/extensions/mcp/oauth.ts index 96ff4143d..b1ca5601a 100644 --- a/packages/coding-agent/src/extensions/mcp/oauth.ts +++ b/packages/coding-agent/src/extensions/mcp/oauth.ts @@ -44,8 +44,11 @@ const CLIENT_METADATA_BASE_URL = "https://pi.dev/oauth"; const FALLBACK_REDIRECT_URL = `http://${CALLBACK_HOST}${CALLBACK_PATH}`; /** Access tokens this close to expiry are refreshed before they are sent. */ const REFRESH_SKEW_MS = 30_000; -/** Bounds each request of a refresh, so it cannot hold the refresh lock or delay shutdown for long. */ -const REFRESH_REQUEST_TIMEOUT_MS = 15_000; +/** + * Bounds each request to the authorization server, so an unresponsive one cannot hold the refresh lock, + * delay shutdown, or keep a sign-in waiting for long. + */ +const OAUTH_REQUEST_TIMEOUT_MS = 15_000; /** A refresh lock that its holder stopped renewing (the process was killed) is taken over after this. */ const REFRESH_LOCK_STALE_MS = 20_000; /** How long to wait for another process's refresh: longer than a stale lock lives. */ @@ -220,6 +223,14 @@ export class McpOAuthCredentialStore { } } +/** `fetch` with `OAUTH_REQUEST_TIMEOUT_MS` per request, on top of the request's own signal. */ +function timedFetch(fetch: McpFetch = globalThis.fetch): McpFetch { + return (input, init) => { + const timeout = AbortSignal.timeout(OAUTH_REQUEST_TIMEOUT_MS); + return fetch(input, { ...init, signal: init?.signal ? AbortSignal.any([init.signal, timeout]) : timeout }); + }; +} + function registeredRedirectUrls(client: OAuthClientInformationMixed | undefined): string[] { return client && "redirect_uris" in client ? client.redirect_uris : []; } @@ -321,14 +332,14 @@ export function createMcpAuthProvider(options: { registeredRedirectUrls(state.clientInformation)[0] ?? FALLBACK_REDIRECT_URL; const provider = createProvider(serverUrl, store, settings, redirectUrl, () => {}); - // Refreshes the tokens, or reports that a new sign-in is needed. + // Refreshes the tokens, or reports that a new sign-in is needed. Not cancellable: a refresh the + // server answered may have rotated the refresh token, so its answer must be saved. const result = await authorizeMcp(provider, { serverUrl, resourceMetadataUrl: challenge?.resourceMetadataUrl, authorizationServerMetadataUrl: settings.authServerMetadataUrl, scope: challenge?.scope, - fetch: (input, init) => - fetch(input, { ...init, signal: AbortSignal.timeout(REFRESH_REQUEST_TIMEOUT_MS) }), + fetch: timedFetch(fetch), }); if (result === "REDIRECT") throw new McpOAuthAuthorizationRequiredError(); }) @@ -367,8 +378,8 @@ export interface McpSignInPrompt { showAuthorizationUrl(url: URL): void; /** * Ask for the redirect URL from the browser address bar, for when the browser cannot reach the - * loopback callback (for example over SSH). Aborted once the callback arrives. Resolves to - * `undefined` or an empty string when the user cancels. + * loopback callback (for example over SSH). Aborted once the callback arrives or the sign-in is + * cancelled. Resolves to `undefined` or an empty string when the user cancels. */ promptForRedirectUrl(signal: AbortSignal): Promise; } @@ -401,16 +412,21 @@ function responseFromRedirectUrl(input: string, state: string, redirectUrl: URL) return { code, iss: url.searchParams.get("iss") ?? undefined }; } -/** Wait for the browser callback or a pasted redirect URL, whichever comes first. */ +/** + * Wait for the browser callback or a pasted redirect URL, whichever comes first. Aborting `signal` + * aborts the prompt, which then cancels the sign-in. + */ async function waitForAuthorizationResponse( callback: OAuthCallbackServer, state: string, redirectUrl: URL, prompt: McpSignInPrompt, + signal: AbortSignal | undefined, ): Promise { const controller = new AbortController(); + const promptSignal = signal ? AbortSignal.any([controller.signal, signal]) : controller.signal; const fromBrowser = callback.waitForCallback(state, redirectUrl.pathname); - const fromUser = prompt.promptForRedirectUrl(controller.signal).then((input) => { + const fromUser = prompt.promptForRedirectUrl(promptSignal).then((input) => { if (!input?.trim()) throw new McpSignInCancelledError(); return responseFromRedirectUrl(input, state, redirectUrl); }); @@ -451,7 +467,8 @@ async function listenForCallback( /** * Sign in to an MCP server. Uses the stored refresh token when possible; otherwise runs the browser - * authorization code flow. Tokens are saved to `store`. + * authorization code flow. Tokens are saved to `store`. Aborting `signal` stops the sign-in at any + * step with `McpSignInCancelledError`; requests to the authorization server are also time-limited. */ export async function signInMcpServer(options: { serverUrl: string; @@ -459,8 +476,10 @@ export async function signInMcpServer(options: { settings: McpOAuthSettings; challenge?: OAuthChallenge; prompt: McpSignInPrompt; + signal?: AbortSignal; }): Promise { - const { serverUrl, store, settings } = options; + const { serverUrl, store, settings, signal } = options; + if (signal?.aborted) throw new McpSignInCancelledError(); const stored = await store.load(); const stepUp = options.challenge?.error === "insufficient_scope"; const callbackOptions = callbackSettings(settings); @@ -501,6 +520,8 @@ export async function signInMcpServer(options: { }); const flow = { serverUrl, + fetch: timedFetch(), + signal, resourceMetadataUrl: options.challenge?.resourceMetadataUrl, authorizationServerMetadataUrl: settings.authServerMetadataUrl, // A server asking for more scope gets it on top of the configured scope and, since the challenge @@ -524,8 +545,13 @@ export async function signInMcpServer(options: { state, authorizationRedirectUrl, options.prompt, + signal, ); await authorizeMcp(provider, { ...flow, authorizationCode: code, iss }); + } catch (error) { + // Aborted requests fail with the signal's reason; report them as the cancellation they are. + if (signal?.aborted) throw new McpSignInCancelledError(); + throw error; } finally { await callback.close(); } diff --git a/packages/coding-agent/src/extensions/mcp/ui.ts b/packages/coding-agent/src/extensions/mcp/ui.ts index 80f3cb818..771e76b61 100644 --- a/packages/coding-agent/src/extensions/mcp/ui.ts +++ b/packages/coding-agent/src/extensions/mcp/ui.ts @@ -46,8 +46,8 @@ export interface McpUi { * rebuilds the menu on every change, keeping the selected item. */ menu(build: () => McpMenu, subscribe?: (listener: () => void) => () => void): Promise; - /** Show a message while an operation runs. */ - status(title: string, message: string): void; + /** Show a message while an operation runs. With `onCancel`, the cancel key calls it. */ + status(title: string, message: string, onCancel?: () => void): void; /** * Show the authorization URL and wait for a pasted redirect URL. Resolves to undefined when * cancelled or when `signal` aborts (the browser reached the callback). @@ -156,8 +156,15 @@ export class McpManagerView implements McpUi, Component, Focusable { }); } - status(title: string, message: string): void { - this.setContent(frame(this.theme, title, [new Spacer(1), new Text(this.theme.fg("muted", message), 1, 0)])); + status(title: string, message: string, onCancel?: () => void): void { + const body = [new Spacer(1), new Text(this.theme.fg("muted", message), 1, 0)]; + if (!onCancel) { + this.setContent(frame(this.theme, title, body)); + return; + } + this.setContent(frame(this.theme, title, body, keyHint("tui.select.cancel", "cancel")), (data) => { + if (this.keybindings.matches(data, "tui.select.cancel")) onCancel(); + }); } redirectUrl(title: string, authorizationUrl: string, signal: AbortSignal): Promise { diff --git a/packages/coding-agent/test/mcp-manager-view.test.ts b/packages/coding-agent/test/mcp-manager-view.test.ts new file mode 100644 index 000000000..ede09a8c9 --- /dev/null +++ b/packages/coding-agent/test/mcp-manager-view.test.ts @@ -0,0 +1,40 @@ +import { setKeybindings, type TUI } from "@earendil-works/pi-tui"; +import { beforeAll, beforeEach, describe, expect, test, vi } from "vitest"; +import { KeybindingsManager } from "../src/core/keybindings.ts"; +import { McpManagerView } from "../src/extensions/mcp/ui.ts"; +import { initTheme, theme } from "../src/modes/interactive/theme/theme.ts"; +import { stripAnsi } from "../src/utils/ansi.ts"; + +const ESCAPE = "\x1b"; +const tui = { requestRender: vi.fn() } as unknown as TUI; + +function rendered(view: McpManagerView): string { + return stripAnsi(view.render(80).join("\n")); +} + +describe("MCP manager status screen", () => { + beforeAll(() => { + initTheme("dark"); + }); + + beforeEach(() => { + setKeybindings(new KeybindingsManager()); + }); + + // #10565 + test("cancels the running operation with the cancel key", () => { + const view = new McpManagerView(tui, theme, new KeybindingsManager()); + const onCancel = vi.fn(); + view.status("Sign in to issues", "Contacting the authorization server…", onCancel); + expect(rendered(view)).toContain("cancel"); + view.handleInput(ESCAPE); + expect(onCancel).toHaveBeenCalledOnce(); + }); + + test("ignores the cancel key for operations that cannot be cancelled", () => { + const view = new McpManagerView(tui, theme, new KeybindingsManager()); + view.status("Sign in to issues", "Connecting…"); + expect(rendered(view)).not.toContain("cancel"); + view.handleInput(ESCAPE); + }); +}); diff --git a/packages/coding-agent/test/suite/agent-session-mcp-oauth.test.ts b/packages/coding-agent/test/suite/agent-session-mcp-oauth.test.ts index cb5b89a83..82c9c6b33 100644 --- a/packages/coding-agent/test/suite/agent-session-mcp-oauth.test.ts +++ b/packages/coding-agent/test/suite/agent-session-mcp-oauth.test.ts @@ -4,8 +4,9 @@ import { tmpdir } from "node:os"; import { join } from "node:path"; import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai"; import type { ToolResultMessage } from "@earendil-works/pi-ai/compat"; -import { afterEach, describe, expect, it } from "vitest"; +import { afterEach, describe, expect, it, vi } from "vitest"; import { InMemoryAuthStorageBackend } from "../../src/core/auth-storage.ts"; +import { emitSessionShutdownEvent } from "../../src/core/extensions/runner.ts"; import { runMcpCommand } from "../../src/extensions/mcp/cli.ts"; import type { McpOAuthConfig, McpServerEntry } from "../../src/extensions/mcp/config.ts"; import { createMcpExtension } from "../../src/extensions/mcp/index.ts"; @@ -128,6 +129,48 @@ describe("AgentSession MCP OAuth", () => { expect(server.log).toContain("token code"); }); + /** Emit session_shutdown like quitting does; resolves to how long the handlers took. */ + async function shutdown(harness: Harness): Promise { + const start = Date.now(); + await emitSessionShutdownEvent(harness.session.extensionRunner, { type: "session_shutdown", reason: "quit" }); + return Date.now() - start; + } + + // #10565 + for (const path of ["/.well-known/oauth-authorization-server", "/token"]) { + it(`cancels a sign-in waiting on ${path} when the session shuts down`, async () => { + const { harness, server, notifications } = await setup("follow"); + server.stall.add(path); + const login = harness.session.prompt("/mcp login issues"); + await vi.waitFor(() => expect(server.stalled).toHaveLength(1), { timeout: 5_000 }); + + // Shutdown waits for the sign-in to clean up, which aborting makes immediate. + expect(await shutdown(harness)).toBeLessThan(2_000); + await login; + // The request is aborted instead of left open, and the ended session reports nothing. + await vi.waitFor(() => expect(server.stalled[0]?.request.socket.destroyed).toBe(true)); + expect(notifications.filter((message) => message.startsWith("Sign-in"))).toEqual([]); + }); + } + + // #10565 + it("closes the session without refreshing an expiring token", async () => { + const { harness, server, notifications, backend } = await setup("follow"); + await harness.session.prompt("/mcp login issues"); + expect(notifications.at(-1)).toBe('Signed in to MCP server "issues" (1 tools).'); + // Still accepted, but close enough to expiry that the next request would refresh it first. + backend.withLock((current) => { + const states = JSON.parse(current ?? "{}") as Record; + for (const state of Object.values(states)) state.tokensExpireAt = Date.now() + 10_000; + return { result: undefined, next: JSON.stringify(states) }; + }); + server.stall.add("/token"); + + expect(await shutdown(harness)).toBeLessThan(2_000); + expect(server.stalled).toEqual([]); + expect(server.deletes).toEqual(["access-1"]); + }); + async function freePort(): Promise { return new Promise((resolve) => { const probe = createServer().listen(0, "127.0.0.1", () => { @@ -174,6 +217,37 @@ describe("AgentSession MCP OAuth", () => { expect(fixed.opened[0].searchParams.get("redirect_uri")).toBe(`http://127.0.0.1:${port}/oauth/done`); }); + it("gives up on pi mcp login after --timeout, also while the authorization server hangs", async () => { + const server = await startOAuthMcpServer(); + cleanups.push(server.close); + const agentDir = mkdtempSync(join(tmpdir(), "pi-mcp-login-")); + cleanups.push(() => rmSync(agentDir, { recursive: true, force: true })); + writeFileSync(join(agentDir, "mcp.json"), JSON.stringify({ mcpServers: { issues: { url: server.url } } })); + const login = async () => { + const output: string[] = []; + const exitCode = await runMcpCommand(["login", "issues", "--timeout", "0.5"], { + cwd: agentDir, + agentDir, + credentials: new McpOAuthCredentialStore(new InMemoryAuthStorageBackend()), + // The user never approves in the browser. + openUrl: () => {}, + log: (line) => output.push(line), + error: (line) => output.push(line), + }); + return { exitCode, output }; + }; + + const unapproved = await login(); + expect(unapproved.exitCode).toBe(1); + expect(unapproved.output.at(-1)).toContain("was cancelled or not completed within"); + + // #10565 + server.stall.add("/.well-known/oauth-authorization-server"); + const stalled = await login(); + expect(stalled.exitCode).toBe(1); + expect(stalled.output.at(-1)).toContain("was cancelled or not completed within"); + }); + it("uses credentials from pi mcp login on the next turn", async () => { const { harness, server, backend } = await setup("follow"); const agentDir = mkdtempSync(join(tmpdir(), "pi-mcp-login-")); diff --git a/packages/coding-agent/test/suite/mcp-oauth-server.ts b/packages/coding-agent/test/suite/mcp-oauth-server.ts index 3721462e9..cd2126cea 100644 --- a/packages/coding-agent/test/suite/mcp-oauth-server.ts +++ b/packages/coding-agent/test/suite/mcp-oauth-server.ts @@ -17,6 +17,8 @@ function json(response: ServerResponse, status: number, body: unknown, headers: * `iss` is sent as the `iss` parameter of authorization responses (RFC 9207). `issParameter` advertises * that parameter and sends the server's issuer. `cimd` advertises Client ID Metadata Documents. * `redirectPath` replaces the path of the redirect URI, like a mixed-up authorization server. + * The MCP endpoint assigns a session, so closing a connection sends a DELETE. Paths added to `stall` + * accept requests and never answer them, like an unresponsive server. */ export async function startOAuthMcpServer( options: { iss?: string; issParameter?: boolean; cimd?: boolean; redirectPath?: string } = {}, @@ -25,6 +27,9 @@ export async function startOAuthMcpServer( const registrations: Record[] = []; const authorizations: URLSearchParams[] = []; const tokenRequests: URLSearchParams[] = []; + const deletes: (string | undefined)[] = []; + const stall = new Set(); + const stalled: { path: string; request: IncomingMessage }[] = []; const validTokens = new Set(); const refreshTokens = new Set(); const challenges = new Map(); @@ -40,11 +45,12 @@ export async function startOAuthMcpServer( }; const handleMcp = async (request: IncomingMessage, response: ServerResponse) => { + const token = request.headers.authorization?.replace(/^Bearer /, ""); + if (request.method === "DELETE") deletes.push(token); if (request.method !== "POST") { response.writeHead(request.method === "GET" ? 405 : 200).end(); return; } - const token = request.headers.authorization?.replace(/^Bearer /, ""); if (!token || !validTokens.has(token)) { log.push(`401 ${token ?? "none"}`); response @@ -74,11 +80,15 @@ export async function startOAuthMcpServer( } else { result = {}; } - json(response, 200, { jsonrpc: "2.0", id: message.id, result }); + json(response, 200, { jsonrpc: "2.0", id: message.id, result }, { "mcp-session-id": "session-1" }); }; const handle = async (request: IncomingMessage, response: ServerResponse) => { const url = new URL(request.url ?? "/", origin); + if (stall.has(url.pathname)) { + stalled.push({ path: url.pathname, request }); + return; + } switch (url.pathname) { case "/mcp": return handleMcp(request, response); @@ -156,6 +166,11 @@ export async function startOAuthMcpServer( authorizations, /** Parameters of token requests. */ tokenRequests, + /** Access tokens of session DELETE requests. */ + deletes, + stall, + /** Requests to stalled paths, still open unless the client gave up. */ + stalled, /** Simulates access token expiry. */ expireAccessTokens: () => validTokens.clear(), close: () => diff --git a/packages/mcp/CHANGELOG.md b/packages/mcp/CHANGELOG.md index d85751eed..e24f27eae 100644 --- a/packages/mcp/CHANGELOG.md +++ b/packages/mcp/CHANGELOG.md @@ -2,6 +2,14 @@ ## [Unreleased] +### Added + +- Added a `signal` option to `authorizeMcp()`, `registerClient()`, the token request functions, and the discovery functions, which aborts their requests. An aborted refresh no longer falls back to a new authorization ([#10565](https://github.com/earendil-works/pi/issues/10565)) + +### Fixed + +- Fixed `StreamableHttpTransport.close()` calling `AuthProvider.token()`, which could refresh tokens over the network before the session DELETE and delay closing; the DELETE now reuses the token of the last request ([#10565](https://github.com/earendil-works/pi/issues/10565)) + ## [1.0.4] - 2026-10-05 ### Fixed diff --git a/packages/mcp/src/auth-provider.ts b/packages/mcp/src/auth-provider.ts index 2757d0377..c1d9e70ca 100644 --- a/packages/mcp/src/auth-provider.ts +++ b/packages/mcp/src/auth-provider.ts @@ -11,6 +11,10 @@ export interface UnauthorizedContext { /** Supplies bearer tokens to an MCP HTTP transport and may refresh them after a 401 response. */ export interface AuthProvider { + /** + * Token for the next request. May refresh it over the network first, so the transport calls it only + * before requests whose answer it waits for. Closing a session reuses the token of the last request. + */ token(): Promise; onUnauthorized?(context: UnauthorizedContext): Promise; } diff --git a/packages/mcp/src/oauth/discovery.ts b/packages/mcp/src/oauth/discovery.ts index e8d3030a4..1d348ca98 100644 --- a/packages/mcp/src/oauth/discovery.ts +++ b/packages/mcp/src/oauth/discovery.ts @@ -55,15 +55,26 @@ export function parseWwwAuthenticate(header: string | null): OAuthChallenge { }; } -async function fetchMetadata(url: URL, fetch: McpFetch, protocolVersion: string): Promise { +async function fetchMetadata( + url: URL, + fetch: McpFetch, + protocolVersion: string, + signal: AbortSignal | undefined, +): Promise { return fetch(url, { headers: { Accept: "application/json", "MCP-Protocol-Version": protocolVersion }, + signal, }); } export async function discoverProtectedResourceMetadata( serverUrl: string | URL, - options: { resourceMetadataUrl?: string | URL; protocolVersion?: string; fetch?: McpFetch } = {}, + options: { + resourceMetadataUrl?: string | URL; + protocolVersion?: string; + fetch?: McpFetch; + signal?: AbortSignal; + } = {}, ): Promise { const server = new URL(serverUrl); const fetch = options.fetch ?? globalThis.fetch; @@ -74,10 +85,16 @@ export async function discoverProtectedResourceMetadata( : new URL(`/.well-known/oauth-protected-resource${pathSuffix(server.pathname)}`, server.origin), fetch, version, + options.signal, ); if (!options.resourceMetadataUrl && server.pathname !== "/" && isDiscoveryMiss(response.status)) { discard(response); - response = await fetchMetadata(new URL("/.well-known/oauth-protected-resource", server.origin), fetch, version); + response = await fetchMetadata( + new URL("/.well-known/oauth-protected-resource", server.origin), + fetch, + version, + options.signal, + ); } if (!response.ok) { discard(response); @@ -101,11 +118,16 @@ export function buildAuthorizationServerDiscoveryUrls( export async function discoverAuthorizationServerMetadata( authorizationServerUrl: string | URL, - options: { fetch?: McpFetch; protocolVersion?: string; skipIssuerValidation?: boolean } = {}, + options: { fetch?: McpFetch; protocolVersion?: string; skipIssuerValidation?: boolean; signal?: AbortSignal } = {}, ): Promise { const fetch = options.fetch ?? globalThis.fetch; for (const { url } of buildAuthorizationServerDiscoveryUrls(authorizationServerUrl)) { - const response = await fetchMetadata(url, fetch, options.protocolVersion ?? LATEST_PROTOCOL_VERSION); + const response = await fetchMetadata( + url, + fetch, + options.protocolVersion ?? LATEST_PROTOCOL_VERSION, + options.signal, + ); if (!response.ok) { discard(response); if (isDiscoveryMiss(response.status)) continue; @@ -131,6 +153,7 @@ export async function discoverOAuthServerInfo( authorizationServerMetadataUrl?: URL; fetch?: McpFetch; skipIssuerValidation?: boolean; + signal?: AbortSignal; } = {}, ): Promise { let resourceMetadata: OAuthProtectedResourceMetadata | undefined; @@ -138,13 +161,19 @@ export async function discoverOAuthServerInfo( resourceMetadata = await discoverProtectedResourceMetadata(serverUrl, { resourceMetadataUrl: options.resourceMetadataUrl, fetch: options.fetch, + signal: options.signal, }); } catch (error) { if (error instanceof TypeError) throw error; } if (options.authorizationServerMetadataUrl) { const url = options.authorizationServerMetadataUrl; - const response = await fetchMetadata(url, options.fetch ?? globalThis.fetch, LATEST_PROTOCOL_VERSION); + const response = await fetchMetadata( + url, + options.fetch ?? globalThis.fetch, + LATEST_PROTOCOL_VERSION, + options.signal, + ); if (!response.ok) { discard(response); throw new Error(`HTTP ${response.status} loading authorization server metadata from ${url}`); @@ -158,6 +187,7 @@ export async function discoverOAuthServerInfo( authorizationServerMetadata: await discoverAuthorizationServerMetadata(authorizationServerUrl, { fetch: options.fetch, skipIssuerValidation: options.skipIssuerValidation, + signal: options.signal, }), resourceMetadata, }; diff --git a/packages/mcp/src/oauth/flow.ts b/packages/mcp/src/oauth/flow.ts index d8711a34b..e93f05d1f 100644 --- a/packages/mcp/src/oauth/flow.ts +++ b/packages/mcp/src/oauth/flow.ts @@ -80,6 +80,8 @@ export interface OAuthFlowOptions { */ authorizationServerMetadataUrl?: URL; fetch?: McpFetch; + /** Aborts every request of the flow. Requests have no time limit of their own; combine with a timeout as needed. */ + signal?: AbortSignal; skipIssuerValidation?: boolean; /** * Go straight to the authorization redirect instead of refreshing stored tokens, for example when the @@ -97,6 +99,7 @@ export interface TokenRequestOptions { resource?: string; addClientAuthentication?: AddClientAuthentication; fetch?: McpFetch; + signal?: AbortSignal; } function loopback(hostname: string): boolean { @@ -217,7 +220,12 @@ async function tokenRequest( params, ); } - const response = await (options.fetch ?? globalThis.fetch)(url, { method: "POST", headers, body: params }); + const response = await (options.fetch ?? globalThis.fetch)(url, { + method: "POST", + headers, + body: params, + signal: options.signal, + }); const text = await response.text(); let value: unknown; try { @@ -242,6 +250,7 @@ export async function registerClient( clientMetadata: OAuthClientMetadata; scope?: string; fetch?: McpFetch; + signal?: AbortSignal; }, ): Promise { const endpoint = options.metadata?.registration_endpoint; @@ -258,6 +267,7 @@ export async function registerClient( options.clientMetadata.application_type ?? applicationType(options.clientMetadata.redirect_uris), ...(options.scope ? { scope: options.scope } : {}), }), + signal: options.signal, }, ); if (!response.ok) throw new OAuthRegistrationError(response.status, await response.text()); @@ -319,6 +329,7 @@ async function runFlow(provider: OAuthClientProvider, options: OAuthFlowOptions) (await discoverAuthorizationServerMetadata(cached.authorizationServerUrl, { fetch: options.fetch, skipIssuerValidation: options.skipIssuerValidation, + signal: options.signal, })), resourceMetadata: cached.resourceMetadata, } @@ -327,6 +338,7 @@ async function runFlow(provider: OAuthClientProvider, options: OAuthFlowOptions) authorizationServerMetadataUrl: metadataUrl, fetch: options.fetch, skipIssuerValidation: options.skipIssuerValidation, + signal: options.signal, }); if (!metadataUrl) { await provider.saveDiscoveryState?.({ @@ -354,6 +366,7 @@ async function runFlow(provider: OAuthClientProvider, options: OAuthFlowOptions) clientMetadata: provider.clientMetadata, scope, fetch: options.fetch, + signal: options.signal, }); await provider.saveClientInformation(client); } @@ -365,6 +378,7 @@ async function runFlow(provider: OAuthClientProvider, options: OAuthFlowOptions) resource, addClientAuthentication: provider.addClientAuthentication, fetch: options.fetch, + signal: options.signal, }; if (options.authorizationCode) { // RFC 9207: never send a code from another authorization server to this one. @@ -394,7 +408,7 @@ async function runFlow(provider: OAuthClientProvider, options: OAuthFlowOptions) await provider.saveTokens(withScope(tokens, existing.scope)); return "AUTHORIZED"; } catch (error) { - if (error instanceof OAuthInsecureEndpointError) throw error; + if (options.signal?.aborted || error instanceof OAuthInsecureEndpointError) throw error; if (error instanceof OAuthError && error.code !== "server_error") throw error; } } diff --git a/packages/mcp/src/transports/streamable-http.ts b/packages/mcp/src/transports/streamable-http.ts index 8fa8f3a06..662465100 100644 --- a/packages/mcp/src/transports/streamable-http.ts +++ b/packages/mcp/src/transports/streamable-http.ts @@ -195,6 +195,8 @@ export class StreamableHttpTransport extends TransportEvents implements McpTrans private sessionIdValue: string | undefined; private protocolVersion: string | undefined; private getStreamStarted = false; + /** Access token of the latest request, which `close()` reuses instead of asking the auth provider. */ + private lastToken: string | undefined; constructor(options: StreamableHttpTransportOptions) { super(); @@ -258,18 +260,16 @@ export class StreamableHttpTransport extends TransportEvents implements McpTrans this.closed = true; this.controller.abort(); if (this.started && this.sessionIdValue) { - const controller = new AbortController(); - const timeout = setTimeout(() => controller.abort(), 1_000); + // Best effort: the session expires on the server anyway. The auth provider may refresh tokens + // over the network, so it is not asked here and closing never waits for a refresh. try { - const { headers } = await this.headers(); - await this.fetch(this.url, { method: "DELETE", headers, signal: controller.signal }) - .then(discard) - .catch(() => undefined); - } catch { - // Resolving auth headers failed; the session will expire on the server. - } finally { - clearTimeout(timeout); - } + const response = await this.fetch(this.url, { + method: "DELETE", + headers: this.buildHeaders({}, this.lastToken), + signal: AbortSignal.timeout(1_000), + }); + await discard(response); + } catch {} } this.emitClose(); } @@ -301,13 +301,18 @@ export class StreamableHttpTransport extends TransportEvents implements McpTrans } private async headers(extra: Record = {}): Promise<{ headers: Headers; token?: string }> { + const token = await this.options.authProvider?.token(); + this.lastToken = token; + return { headers: this.buildHeaders(extra, token), ...(token ? { token } : {}) }; + } + + private buildHeaders(extra: Record, token: string | undefined): Headers { const headers = new Headers(this.options.headers); for (const [name, value] of Object.entries(extra)) headers.set(name, value); if (this.sessionIdValue) headers.set("Mcp-Session-Id", this.sessionIdValue); if (this.protocolVersion) headers.set("MCP-Protocol-Version", this.protocolVersion); - const token = await this.options.authProvider?.token(); if (token) headers.set("Authorization", `Bearer ${token}`); - return { headers, ...(token ? { token } : {}) }; + return headers; } private captureSession(response: Response): void { diff --git a/packages/mcp/test/oauth.test.ts b/packages/mcp/test/oauth.test.ts index 2415c7346..57612f6af 100644 --- a/packages/mcp/test/oauth.test.ts +++ b/packages/mcp/test/oauth.test.ts @@ -520,6 +520,46 @@ describe("MCP OAuth", () => { expect(await exchange("omitted", undefined, false)).toBe("AUTHORIZED"); expect(codes).toEqual(["matching", "omitted"]); }); + + // #10565 + it("stops when its signal aborts, without falling back to a redirect", async () => { + const stalled: string[] = []; + // Accepts every request and never answers. + const origin = await listen(async (request, _response, serverOrigin) => { + stalled.push(new URL(request.url ?? "/", serverOrigin).pathname); + }); + const run = async (provider: TestOAuthProvider) => { + const controller = new AbortController(); + const count = stalled.length; + const flow = authorizeMcp(provider, { serverUrl: `${origin}/mcp`, signal: controller.signal }); + const settled = flow.catch((error: unknown) => error); + await expect.poll(() => stalled.length).toBe(count + 1); + controller.abort(); + expect(await settled).toMatchObject({ name: "AbortError" }); + expect(provider.authorizationUrl).toBeUndefined(); + return stalled.at(-1); + }; + + // Discovery. + expect(await run(new TestOAuthProvider("http://127.0.0.1/callback"))).toBe( + "/.well-known/oauth-protected-resource/mcp", + ); + + // A failed refresh otherwise falls back to a new authorization. + const refreshing = new TestOAuthProvider("http://127.0.0.1/callback"); + refreshing.client = { client_id: "client" }; + refreshing.tokenSet = { access_token: "a1", refresh_token: "r1", token_type: "Bearer" }; + refreshing.discovery = { + authorizationServerUrl: origin, + authorizationServerMetadata: { + issuer: origin, + authorization_endpoint: `${origin}/authorize`, + token_endpoint: `${origin}/token`, + response_types_supported: ["code"], + }, + }; + expect(await run(refreshing)).toBe("/token"); + }); }); describe("OAuthCallbackServer pages", () => { diff --git a/packages/mcp/test/streamable-http.test.ts b/packages/mcp/test/streamable-http.test.ts index d4d260d0a..e0c0b25ad 100644 --- a/packages/mcp/test/streamable-http.test.ts +++ b/packages/mcp/test/streamable-http.test.ts @@ -342,6 +342,26 @@ describe("StreamableHttpTransport", () => { await client.close(); }); + // #10565: the auth provider may refresh over the network, which closing must not wait for. + it("closes the session with the last request's token, without asking the auth provider", async () => { + const { url, requests } = await startServer(protocolHandler); + let calls = 0; + const client = new McpClient({ name: "http-test", version: "1.0.0" }); + await client.connect( + new StreamableHttpTransport({ + url, + openGetStream: false, + authProvider: { token: async () => `token-${++calls}` }, + }), + ); + const before = calls; + await client.close(); + expect(calls).toBe(before); + const deletes = requests.filter((entry) => entry.method === "DELETE"); + expect(deletes.map((entry) => entry.headers.authorization)).toEqual([`Bearer token-${before}`]); + expect(deletes[0]?.headers["mcp-session-id"]).toBe("session-1"); + }); + // #10188: Cloudflare Workers reject the platform fetch when called with a receiver other than globalThis. it("calls fetch without a receiver", async () => { const realFetch = globalThis.fetch;