[glean-vnext] Fix intermittent re-auth from cross-process refresh-token rotation

shared/glean/mcp/src/auth-provider.ts

@@ -6,12 +6,21 @@ import type {
 } from "@modelcontextprotocol/sdk/shared/auth.js";
 import { execFile, spawn } from "node:child_process";
 import { randomUUID } from "node:crypto";
+import { setTimeout as sleep } from "node:timers/promises";
 import { platform } from "node:os";
 import { getCallbackUrl, setExpectedState } from "./auth-callback-server.js";
-import { clearCredentials, loadCredentials, saveCredentials } from "./token-store.js";
+import {
+  clearCredentials,
+  loadCredentials,
+  saveCredentials,
+} from "./token-store.js";
 
 export type InvalidationScope = "all" | "client" | "tokens" | "verifier";
 
+// Grace window for a sibling's in-flight refresh to land on disk.
+const ROTATION_GRACE_MS = 2000;
+const ROTATION_POLL_MS = 100;
+
 /**
  * Open `url` in the user's default browser. Used for the self-open sign-in
  * path when the client does not support URL-mode elicitation (where the client
@@ -48,7 +57,6 @@ export class GleanOAuthClientProvider implements OAuthClientProvider {
   // explicitly invalidating. Used to detect when a previous auth URL didn't
   // complete — likely because the server rejected the (stale) client_id.
   private _authUrlPending = false;
-
   authorizationUrl: string | undefined;
 
   /**
@@ -67,6 +75,33 @@ export class GleanOAuthClientProvider implements OAuthClientProvider {
     }
   }
 
+  // Re-read the shared store on every token access so a sibling's rotated
+  // grant is used instead of a stale in-memory copy.
+  private syncTokensFromDisk(): void {
+    const stored = loadCredentials();
+    if (!stored) return;
+    if (stored.tokens) {
+      this._tokens = stored.tokens as OAuthTokens;
+    }
+    if (stored.clientInfo) {
+      this._clientInfo = stored.clientInfo as OAuthClientInformationMixed;
+    }
+  }
+
+  // Wait for a sibling's refresh to land on disk. Returns true once a
+  // different access token is available for adoption/retry.
+  async waitForSiblingRefresh(
+    previousAccessToken: string | undefined,
+  ): Promise<boolean> {
+    const deadline = Date.now() + ROTATION_GRACE_MS;
+    for (;;) {
+      const current = this.tokens()?.access_token;
+      if (current && current !== previousAccessToken) return true;
+      if (Date.now() >= deadline) return false;
+      await sleep(ROTATION_POLL_MS);
+    }
+  }
+
   get redirectUrl(): string {
     return getCallbackUrl();
   }
@@ -88,6 +123,7 @@ export class GleanOAuthClientProvider implements OAuthClientProvider {
   }
 
   tokens(): OAuthTokens | undefined {
+    this.syncTokensFromDisk();
     return this._tokens;
   }
 
@@ -113,10 +149,19 @@ export class GleanOAuthClientProvider implements OAuthClientProvider {
         this._clientInfo = undefined;
         saveCredentials(this._tokens, undefined);
         break;
-      case "tokens":
+      case "tokens": {
+        // Usually a sibling's rotation — try adopting before clearing.
+        const previousAccessToken = this._tokens?.access_token;
+        if (
+          this._tokens?.refresh_token &&
+          (await this.waitForSiblingRefresh(previousAccessToken))
+        ) {
+          return;
+        }
         this._tokens = undefined;
         saveCredentials(undefined, this._clientInfo);
         break;
+      }
       case "verifier":
         this._codeVerifier = "";
         break;

shared/glean/mcp/src/remote-client.ts

@@ -1,6 +1,7 @@
 import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
 import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js";
+import { OAuthError } from "@modelcontextprotocol/sdk/server/auth/errors.js";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import type { GleanOAuthClientProvider } from "./auth-provider.js";
 import { PLUGIN_VERSION } from "./version.js";
@@ -167,6 +168,7 @@ export async function createRemoteClient(
   serverUrl: string,
   opts: RemoteClientOptions,
   chatSessionId?: string,
+  authRetry = false,
 ): Promise<Client> {
   const authProvider = opts.authProvider;
 
@@ -214,21 +216,62 @@ export async function createRemoteClient(
     { capabilities: {} },
   );
 
+  // Snapshot to detect a sibling's refresh between connect and failure.
+  const accessTokenAtConnect = authProvider?.tokens()?.access_token;
+
   const transport = buildTransport(serverUrl, opts, chatSessionId);
 
   try {
     await withConnectLock(() => client.connect(transport));
   } catch (error) {
-    if (error instanceof UnauthorizedError && authProvider?.authorizationUrl) {
-      pendingTransport = transport;
-      throw new AuthRequiredError(authProvider.authorizationUrl);
+    if (error instanceof UnauthorizedError && authProvider) {
+      const refreshedAccessToken = authProvider.tokens()?.access_token;
+      if (
+        !authRetry &&
+        refreshedAccessToken &&
+        refreshedAccessToken !== accessTokenAtConnect
+      ) {
+        console.error(
+          "[auth] Auth failed but a newer token is on disk " +
+            "(sibling refresh) — retrying once",
+        );
+        return createRemoteClient(serverUrl, opts, chatSessionId, true);
+      }
+      if (authProvider.authorizationUrl) {
+        pendingTransport = transport;
+        throw new AuthRequiredError(authProvider.authorizationUrl);
+      }
+    }
+    // Concurrent-refresh losers are reported with structured OAuth errors
+    // (typically invalid_request); retry once if a sibling's grant lands in the
+    // grace window.
+    if (
+      authProvider &&
+      !authRetry &&
+      isRefreshOAuthError(error) &&
+      (await authProvider.waitForSiblingRefresh(accessTokenAtConnect))
+    ) {
+      console.error(
+        "[auth] Refresh failed but a sibling refreshed — retrying with its token",
+      );
+      return createRemoteClient(serverUrl, opts, chatSessionId, true);
     }
     throw error;
   }
 
   return client;
 }
 
+// Restrict recovery to OAuth errors that can indicate a refresh race. The SDK
+// preserves the response's machine-readable error code.
+function isRefreshOAuthError(error: unknown): boolean {
+  return (
+    error instanceof OAuthError &&
+    (error.errorCode === "invalid_request" ||
+      error.errorCode === "invalid_grant")
+  );
+}
+
 export async function callRemoteTool(
   client: Client,
   name: string,

shared/glean/mcp/src/token-store.ts

@@ -35,11 +35,14 @@ export function saveCredentials(tokens: unknown, clientInfo: unknown): void {
     fs.mkdirSync(dir, { recursive: true, mode: DIR_MODE });
     fs.chmodSync(dir, DIR_MODE);
     const data: StoredCredentials = { tokens, clientInfo };
-    fs.writeFileSync(filePath, JSON.stringify(data, null, 2), {
+    // Temp-file + rename: concurrent readers never see a half-written store.
+    const tmpPath = `${filePath}.${process.pid}.tmp`;
+    fs.writeFileSync(tmpPath, JSON.stringify(data, null, 2), {
       encoding: "utf-8",
       mode: FILE_MODE,
     });
-    fs.chmodSync(filePath, FILE_MODE);
+    fs.chmodSync(tmpPath, FILE_MODE);
+    fs.renameSync(tmpPath, filePath);
   } catch (err) {
     const msg = err instanceof Error ? err.message : String(err);
     console.error(`[auth] Failed to persist credentials: ${msg}`);

shared/glean/mcp/tests/auth-provider.test.ts

@@ -33,6 +33,7 @@ describe("GleanOAuthClientProvider", () => {
   });
 
   afterEach(() => {
+    vi.useRealTimers();
     fs.rmSync(gleanDir, { recursive: true, force: true });
   });
 
@@ -74,6 +75,143 @@ describe("GleanOAuthClientProvider", () => {
     expect(raw.tokens.access_token).toBe("new_tok");
   });
 
+  // --- Cross-process sync: tokens() must pick up a sibling's rewrite. ---
+
+  const credFile = path.join(gleanDir, "mcp-credentials.json");
+
+  function writeCredFile(tokens: unknown, clientInfo?: unknown): void {
+    fs.mkdirSync(gleanDir, { recursive: true });
+    fs.writeFileSync(credFile, JSON.stringify({ tokens, clientInfo }));
+  }
+
+  it("tokens() adopts a token written by another process", () => {
+    fs.mkdirSync(gleanDir, { recursive: true });
+    fs.writeFileSync(
+      credFile,
+      JSON.stringify({
+        tokens: { access_token: "T0", refresh_token: "R0" },
+        clientInfo: { client_id: "cid" },
+      }),
+    );
+    const provider = new GleanOAuthClientProvider();
+    expect(provider.tokens()?.access_token).toBe("T0");
+    const originalMtime = fs.statSync(credFile).mtime;
+
+    // Sibling refreshes: new access + rotated refresh token on disk.
+    writeCredFile(
+      { access_token: "T1", refresh_token: "R1" },
+      { client_id: "cid" },
+    );
+    // The provider must not rely on mtime to observe this rewrite.
+    fs.utimesSync(credFile, originalMtime, originalMtime);
+
+    expect(provider.tokens()?.access_token).toBe("T1");
+    expect(provider.tokens()?.refresh_token).toBe("R1");
+  });
+
+  it("tokens() keeps the in-memory token when the file is deleted", () => {
+    fs.mkdirSync(gleanDir, { recursive: true });
+    fs.writeFileSync(
+      credFile,
+      JSON.stringify({ tokens: { access_token: "T0" }, clientInfo: {} }),
+    );
+    const provider = new GleanOAuthClientProvider();
+    expect(provider.tokens()?.access_token).toBe("T0");
+
+    // Transient disappearance / another process mid-write — don't self-evict.
+    fs.rmSync(credFile, { force: true });
+    expect(provider.tokens()?.access_token).toBe("T0");
+  });
+
+  it("tokens() does not adopt a rewrite that carries no tokens", () => {
+    fs.mkdirSync(gleanDir, { recursive: true });
+    fs.writeFileSync(
+      credFile,
+      JSON.stringify({ tokens: { access_token: "T0" }, clientInfo: {} }),
+    );
+    const provider = new GleanOAuthClientProvider();
+    expect(provider.tokens()?.access_token).toBe("T0");
+
+    // A client-only rewrite (tokens dropped) must not log us out in-memory.
+    writeCredFile(undefined, { client_id: "cid" });
+    expect(provider.tokens()?.access_token).toBe("T0");
+  });
+
+  it("invalidateCredentials('tokens') adopts a sibling's token instead of wiping the store", async () => {
+    fs.mkdirSync(gleanDir, { recursive: true });
+    fs.writeFileSync(
+      credFile,
+      JSON.stringify({
+        tokens: { access_token: "T0", refresh_token: "R0" },
+        clientInfo: { client_id: "cid" },
+      }),
+    );
+    const provider = new GleanOAuthClientProvider();
+    expect(provider.tokens()?.access_token).toBe("T0");
+
+    // A sibling refreshed + rotated: fresh grant is now on disk.
+    writeCredFile(
+      { access_token: "T1", refresh_token: "R1" },
+      { client_id: "cid" },
+    );
+
+    // The SDK calls this on invalid_grant. It must NOT clear — the failure was
+    // just our stale token; adopt the sibling's fresh one and leave it on disk.
+    await provider.invalidateCredentials("tokens");
+
+    expect(provider.tokens()?.access_token).toBe("T1");
+    expect(provider.tokens()?.refresh_token).toBe("R1");
+    const raw = JSON.parse(fs.readFileSync(credFile, "utf-8"));
+    expect(raw.tokens.access_token).toBe("T1"); // not clobbered with undefined
+  });
+
+  it("invalidateCredentials('tokens') clears when there is no newer token on disk", async () => {
+    const provider = new GleanOAuthClientProvider();
+    provider.saveTokens({ access_token: "T0", refresh_token: "R0" } as any);
+    expect(provider.tokens()?.access_token).toBe("T0");
+
+    // No sibling write since our snapshot → a genuine invalidation → clear.
+    vi.useFakeTimers();
+    const invalidation = provider.invalidateCredentials("tokens");
+    await vi.advanceTimersByTimeAsync(2000);
+    await invalidation;
+
+    expect(provider.tokens()).toBeUndefined();
+    const raw = JSON.parse(fs.readFileSync(credFile, "utf-8"));
+    expect(raw.tokens).toBeUndefined();
+  });
+
+  it("invalidateCredentials('tokens') adopts a token that lands during the grace window", async () => {
+    // The winner's write lands just after the loser's invalid_grant.
+    const provider = new GleanOAuthClientProvider();
+    provider.saveTokens({ access_token: "T0", refresh_token: "R0" } as any);
+
+    const invalidation = provider.invalidateCredentials("tokens");
+    // Sibling's write lands mid-window.
+    setTimeout(() => {
+      writeCredFile(
+        { access_token: "T1", refresh_token: "R1" },
+        { client_id: "cid" },
+      );
+    }, 150);
+    await invalidation;
+
+    expect(provider.tokens()?.access_token).toBe("T1");
+    const raw = JSON.parse(fs.readFileSync(credFile, "utf-8"));
+    expect(raw.tokens.access_token).toBe("T1"); // not clobbered with undefined
+  });
+
+  it("skips the grace window when no refresh token was held (no race possible)", async () => {
+    const provider = new GleanOAuthClientProvider();
+    provider.saveTokens({ access_token: "T0" } as any); // no refresh_token
+
+    const start = Date.now();
+    await provider.invalidateCredentials("tokens");
+
+    expect(Date.now() - start).toBeLessThan(1000); // no 5s poll
+    expect(provider.tokens()).toBeUndefined();
+  });
+
   it("saveClientInformation persists to disk", () => {
     const provider = new GleanOAuthClientProvider();
     const info = { client_id: "cid", client_secret: "sec" } as any;