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;