diff --git a/src/agent/tools.ts b/src/agent/tools.ts index f04a92e9a..341d85fe3 100644 --- a/src/agent/tools.ts +++ b/src/agent/tools.ts @@ -353,6 +353,18 @@ export async function createAgentToolset(args: AgentToolsetArgs): Promise callbacks.onStatus({ name, state: "needs-auth", url }), + // Mid-session re-auth fires needs-auth again without a later connected + // event. Re-emit connected only when tools are already registered so + // first-connect still waits for the real post-connect status. + onAuthorized: (name) => { + const client = connectedClients.find((c) => c.serverName === name); + if (client === undefined) return; + callbacks.onStatus({ + name, + state: "connected", + tools: client.tools.map((t) => t.name), + }); + }, ...(signal !== undefined ? { signal } : {}), }); if (!result.ok) { diff --git a/src/mcp/client-auth-retry.test.ts b/src/mcp/client-auth-retry.test.ts new file mode 100644 index 000000000..98b63152f --- /dev/null +++ b/src/mcp/client-auth-retry.test.ts @@ -0,0 +1,51 @@ +import { describe, expect, test } from "bun:test"; +import { retryAfterInteractiveAuth } from "./client.js"; + +describe("retryAfterInteractiveAuth", () => { + test("notifies only after the retried operation succeeds", async () => { + let notified = false; + await expect( + retryAfterInteractiveAuth( + async () => undefined, + async () => { + throw new Error("still unauthorized"); + }, + () => { + notified = true; + }, + ), + ).rejects.toThrow("still unauthorized"); + expect(notified).toBe(false); + + const value = await retryAfterInteractiveAuth( + async () => undefined, + async () => "ok", + () => { + notified = true; + }, + ); + expect(value).toBe("ok"); + expect(notified).toBe(true); + }); + + test("propagates auth completion failures without notifying", async () => { + let notified = false; + let operationRan = false; + await expect( + retryAfterInteractiveAuth( + async () => { + throw new Error("auth aborted"); + }, + async () => { + operationRan = true; + return "ok"; + }, + () => { + notified = true; + }, + ), + ).rejects.toThrow("auth aborted"); + expect(operationRan).toBe(false); + expect(notified).toBe(false); + }); +}); diff --git a/src/mcp/client.ts b/src/mcp/client.ts index 8f7bafe37..00e05c39d 100644 --- a/src/mcp/client.ts +++ b/src/mcp/client.ts @@ -22,6 +22,12 @@ export type MCPConnectResult = { ok: true; client: MCPClient } | { ok: false; se export type MCPConnectOptions = { stderr?: "inherit" | "ignore" | "pipe"; onAuthURL?: (serverName: string, authorizationUrl: string) => void; + /** + * Interactive OAuth finished and the retried operation succeeded. Callers + * that already registered tools for this server can re-emit a connected + * status so standing "needs auth" chrome clears mid-session. + */ + onAuthorized?: (serverName: string) => void; signal?: AbortSignal; }; @@ -39,7 +45,15 @@ export function unwrapToolContent(content: unknown): string { }).join("\n"); } -type HTTPAuthContext = { url: URL; authProvider: CorbitsOAuthProvider; callback: CallbackServer; signal?: AbortSignal; interactive: boolean }; +type HTTPAuthContext = { + url: URL; + authProvider: CorbitsOAuthProvider; + callback: CallbackServer; + signal?: AbortSignal; + interactive: boolean; + serverName: string; + onAuthorized?: (serverName: string) => void; +}; function isRecoverableAuthError(err: unknown): boolean { return err instanceof UnauthorizedError || err instanceof OAuthError; @@ -51,14 +65,35 @@ async function completeInteractiveAuth(context: HTTPAuthContext): Promise await new StreamableHTTPClientTransport(context.url, { authProvider: context.authProvider }).finishAuth(code); } +/** + * Run interactive OAuth, retry the failed operation, and notify only when the + * retry itself succeeded — a failed re-auth must leave standing "needs auth" + * chrome alone. + */ +export async function retryAfterInteractiveAuth( + completeAuth: () => Promise, + operation: () => Promise, + onAuthorized: (() => void) | undefined, +): Promise { + await completeAuth(); + const value = await operation(); + onAuthorized?.(); + return value; +} + async function recoverHTTPAuthorization(err: unknown, context: HTTPAuthContext | undefined, operation: () => Promise): Promise { if (context === undefined || !isRecoverableAuthError(err)) throw err; let lastErr: unknown = err; for (let attempt = 0; attempt < 2; attempt += 1) { if (lastErr instanceof OAuthError) await context.authProvider.resetAuthorization(); if (lastErr instanceof UnauthorizedError) { - await completeInteractiveAuth(context); - return operation(); + return retryAfterInteractiveAuth( + () => completeInteractiveAuth(context), + operation, + context.onAuthorized === undefined + ? undefined + : () => context.onAuthorized?.(context.serverName), + ); } try { return await operation(); @@ -128,7 +163,15 @@ async function connectHttp(config: MCPServerConfig, options: MCPConnectOptions): }); const makeTransport = (): Transport => new StreamableHTTPClientTransport(url, { authProvider }) as unknown as Transport; const client = new Client({ name: MCP_CLIENT_NAME, version: "1.0.0" }); - const authContext: HTTPAuthContext = { url, authProvider, callback, interactive: options.onAuthURL !== undefined, ...(options.signal !== undefined ? { signal: options.signal } : {}) }; + const authContext: HTTPAuthContext = { + url, + authProvider, + callback, + interactive: options.onAuthURL !== undefined, + serverName: config.name, + ...(options.onAuthorized !== undefined ? { onAuthorized: options.onAuthorized } : {}), + ...(options.signal !== undefined ? { signal: options.signal } : {}), + }; try { await withHTTPAuthorizationRecovery(authContext, () => client.connect(makeTransport())); return { ok: true, client: await finishClient(client, config.name, authContext) }; diff --git a/src/tui-opentui/runtime-channels.test.ts b/src/tui-opentui/runtime-channels.test.ts index f38de2155..a554fbe99 100644 --- a/src/tui-opentui/runtime-channels.test.ts +++ b/src/tui-opentui/runtime-channels.test.ts @@ -115,13 +115,15 @@ describe("mcp.status channel", () => { } }) - test("connecting clears the standing auth segment", async () => { + test("connected clears the standing auth segment from state and the painted frame", async () => { const { host, emitter, frame, cleanup } = await mountHeadless() try { emitter.emit("mcp.status", { name: "linear", state: "needs-auth", url: "https://x/a" }) + expect(await frame()).toContain("needs auth") emitter.emit("mcp.status", { name: "linear", state: "connected", tools: ["a"] }) - await frame() + const painted = await frame() expect(host.shell.mcpNeedsAuth).toEqual([]) + expect(painted).not.toContain("needs auth") } finally { cleanup() }