Clear the MCP auth banner after mid-session re-authentication succeeds (#433)

Re-emit connected status once interactive OAuth and the retried MCP
operation both succeed, so standing needs-auth chrome does not stick
around after tools are already registered.

Author: TheGreatAxios

src/agent/tools.ts

@@ -353,6 +353,18 @@ export async function createAgentToolset(args: AgentToolsetArgs): Promise<AgentT
         const result = await connectMCPServer(config, {
           stderr: "ignore",
           onAuthURL: (name, url) => 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) {

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);
+  });
+});

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<void>
   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<T>(
+  completeAuth: () => Promise<void>,
+  operation: () => Promise<T>,
+  onAuthorized: (() => void) | undefined,
+): Promise<T> {
+  await completeAuth();
+  const value = await operation();
+  onAuthorized?.();
+  return value;
+}
+
 async function recoverHTTPAuthorization<T>(err: unknown, context: HTTPAuthContext | undefined, operation: () => Promise<T>): Promise<T> {
   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) };

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()
     }