Merge bebe209785 into 4583f88626

CHANGELOG.md

@@ -14,6 +14,7 @@ Status: beta.
 - Memory Search: allow extra paths for memory indexing. (#3600) Thanks @kira-ariaki.
 
 ### Changes
+- Auth: fix OAuth refresh only refreshing one profile per provider, leaving other accounts expired. (#3803)
 - Providers: add Venice AI integration; update Moonshot Kimi references to kimi-k2.5; update MiniMax API endpoint/format. (#2762, #3064)
 - Providers: add Xiaomi MiMo (mimo-v2-flash) support and onboarding flow. (#3454) Thanks @WqyJh.
 - Telegram: quote replies, edit-message action, silent sends, sticker support + vision caching, linkPreview toggle, plugin sendPayload support. (#2900, #2394, #2382, #2548, #1700, #1917)

src/agents/auth-profiles/oauth-refresh.test.ts

@@ -0,0 +1,262 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { mkdtempSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+
+// Mock pi-ai's getOAuthApiKey
+const mockGetOAuthApiKey = vi.fn();
+vi.mock("@mariozechner/pi-ai", () => ({
+  getOAuthApiKey: mockGetOAuthApiKey,
+}));
+
+describe("OAuth refresh for Google providers", () => {
+  let tempDir: string;
+
+  beforeEach(() => {
+    tempDir = mkdtempSync(join(tmpdir(), "moltbot-oauth-test-"));
+    vi.resetModules();
+    mockGetOAuthApiKey.mockReset();
+  });
+
+  afterEach(() => {
+    try {
+      rmSync(tempDir, { recursive: true, force: true });
+    } catch {
+      // ignore cleanup errors
+    }
+  });
+
+  it("refreshes google-gemini-cli tokens via pi-ai when expired", async () => {
+    // Setup: expired google-gemini-cli credentials
+    const now = Date.now();
+    const expiredCreds = {
+      type: "oauth" as const,
+      provider: "google-gemini-cli",
+      access: "old-access-token",
+      refresh: "valid-refresh-token",
+      expires: now - 1000, // expired
+      email: "user@gmail.com",
+      projectId: "test-project-123",
+    };
+
+    // Mock successful refresh
+    mockGetOAuthApiKey.mockResolvedValue({
+      apiKey: JSON.stringify({
+        token: "new-access-token",
+        projectId: "test-project-123",
+      }),
+      newCredentials: {
+        access: "new-access-token",
+        refresh: "new-refresh-token",
+        expires: now + 3600_000,
+        projectId: "test-project-123",
+      },
+    });
+
+    const { ensureAuthProfileStore, saveAuthProfileStore } = await import("./store.js");
+    const { resolveApiKeyForProfile } = await import("./oauth.js");
+
+    // Create auth store with expired credentials
+    const store = ensureAuthProfileStore(tempDir);
+    store.profiles["google-gemini-cli:user@gmail.com"] = expiredCreds;
+    saveAuthProfileStore(store, tempDir);
+
+    // Attempt to resolve API key (should trigger refresh)
+    const result = await resolveApiKeyForProfile({
+      store,
+      profileId: "google-gemini-cli:user@gmail.com",
+      agentDir: tempDir,
+    });
+
+    expect(result).toBeTruthy();
+    expect(mockGetOAuthApiKey).toHaveBeenCalledWith("google-gemini-cli", {
+      "google-gemini-cli": expect.objectContaining({
+        refresh: "valid-refresh-token",
+        projectId: "test-project-123",
+      }),
+    });
+
+    // Verify credentials were updated in store
+    const updatedStore = ensureAuthProfileStore(tempDir);
+    const updatedCred = updatedStore.profiles["google-gemini-cli:user@gmail.com"];
+    expect(updatedCred).toMatchObject({
+      type: "oauth",
+      access: "new-access-token",
+      refresh: "new-refresh-token",
+    });
+    expect(updatedCred.expires).toBeGreaterThan(now);
+  });
+
+  it("refreshes google-antigravity tokens via pi-ai when expired", async () => {
+    const now = Date.now();
+    const expiredCreds = {
+      type: "oauth" as const,
+      provider: "google-antigravity",
+      access: "old-access-token",
+      refresh: "valid-refresh-token",
+      expires: now - 1000,
+      email: "user@gmail.com",
+      projectId: "antigravity-project",
+    };
+
+    mockGetOAuthApiKey.mockResolvedValue({
+      apiKey: JSON.stringify({
+        token: "new-access-token",
+        projectId: "antigravity-project",
+      }),
+      newCredentials: {
+        access: "new-access-token",
+        refresh: "new-refresh-token",
+        expires: now + 3600_000,
+        projectId: "antigravity-project",
+      },
+    });
+
+    const { ensureAuthProfileStore, saveAuthProfileStore } = await import("./store.js");
+    const { resolveApiKeyForProfile } = await import("./oauth.js");
+
+    const store = ensureAuthProfileStore(tempDir);
+    store.profiles["google-antigravity:user@gmail.com"] = expiredCreds;
+    saveAuthProfileStore(store, tempDir);
+
+    const result = await resolveApiKeyForProfile({
+      store,
+      profileId: "google-antigravity:user@gmail.com",
+      agentDir: tempDir,
+    });
+
+    expect(result).toBeTruthy();
+    expect(mockGetOAuthApiKey).toHaveBeenCalledWith("google-antigravity", {
+      "google-antigravity": expect.objectContaining({
+        refresh: "valid-refresh-token",
+        projectId: "antigravity-project",
+      }),
+    });
+  });
+
+  it("handles refresh failure and throws meaningful error", async () => {
+    const now = Date.now();
+    const expiredCreds = {
+      type: "oauth" as const,
+      provider: "google-gemini-cli",
+      access: "old-access-token",
+      refresh: "invalid-refresh-token",
+      expires: now - 1000,
+      email: "user@gmail.com",
+      projectId: "test-project",
+    };
+
+    // Mock refresh failure
+    mockGetOAuthApiKey.mockRejectedValue(new Error("Invalid refresh token"));
+
+    const { ensureAuthProfileStore, saveAuthProfileStore } = await import("./store.js");
+    const { resolveApiKeyForProfile } = await import("./oauth.js");
+
+    const store = ensureAuthProfileStore(tempDir);
+    store.profiles["google-gemini-cli:user@gmail.com"] = expiredCreds;
+    saveAuthProfileStore(store, tempDir);
+
+    await expect(
+      resolveApiKeyForProfile({
+        store,
+        profileId: "google-gemini-cli:user@gmail.com",
+        agentDir: tempDir,
+      }),
+    ).rejects.toThrow(/OAuth token refresh failed for google-gemini-cli/);
+  });
+
+  it("works with multiple accounts independently", async () => {
+    const now = Date.now();
+
+    // Setup: Two expired accounts
+    const account1 = {
+      type: "oauth" as const,
+      provider: "google-gemini-cli",
+      access: "old-access-1",
+      refresh: "refresh-1",
+      expires: now - 1000,
+      email: "user1@gmail.com",
+      projectId: "project-1",
+    };
+
+    const account2 = {
+      type: "oauth" as const,
+      provider: "google-gemini-cli",
+      access: "old-access-2",
+      refresh: "refresh-2",
+      expires: now - 1000,
+      email: "user2@gmail.com",
+      projectId: "project-2",
+    };
+
+    // Mock refresh for each account
+    mockGetOAuthApiKey
+      .mockResolvedValueOnce({
+        apiKey: JSON.stringify({ token: "new-access-1", projectId: "project-1" }),
+        newCredentials: {
+          access: "new-access-1",
+          refresh: "new-refresh-1",
+          expires: now + 3600_000,
+          projectId: "project-1",
+        },
+      })
+      .mockResolvedValueOnce({
+        apiKey: JSON.stringify({ token: "new-access-2", projectId: "project-2" }),
+        newCredentials: {
+          access: "new-access-2",
+          refresh: "new-refresh-2",
+          expires: now + 3600_000,
+          projectId: "project-2",
+        },
+      });
+
+    const { ensureAuthProfileStore, saveAuthProfileStore } = await import("./store.js");
+    const { resolveApiKeyForProfile } = await import("./oauth.js");
+
+    const store = ensureAuthProfileStore(tempDir);
+    store.profiles["google-gemini-cli:user1@gmail.com"] = account1;
+    store.profiles["google-gemini-cli:user2@gmail.com"] = account2;
+    saveAuthProfileStore(store, tempDir);
+
+    // Refresh account 1
+    const result1 = await resolveApiKeyForProfile({
+      store: ensureAuthProfileStore(tempDir),
+      profileId: "google-gemini-cli:user1@gmail.com",
+      agentDir: tempDir,
+    });
+
+    expect(result1).toBeTruthy();
+    expect(mockGetOAuthApiKey).toHaveBeenNthCalledWith(1, "google-gemini-cli", {
+      "google-gemini-cli": expect.objectContaining({
+        refresh: "refresh-1",
+        projectId: "project-1",
+      }),
+    });
+
+    // Refresh account 2
+    const result2 = await resolveApiKeyForProfile({
+      store: ensureAuthProfileStore(tempDir),
+      profileId: "google-gemini-cli:user2@gmail.com",
+      agentDir: tempDir,
+    });
+
+    expect(result2).toBeTruthy();
+    expect(mockGetOAuthApiKey).toHaveBeenNthCalledWith(2, "google-gemini-cli", {
+      "google-gemini-cli": expect.objectContaining({
+        refresh: "refresh-2",
+        projectId: "project-2",
+      }),
+    });
+
+    // Verify both accounts were updated independently
+    const updatedStore = ensureAuthProfileStore(tempDir);
+    expect(updatedStore.profiles["google-gemini-cli:user1@gmail.com"]).toMatchObject({
+      access: "new-access-1",
+      refresh: "new-refresh-1",
+    });
+    expect(updatedStore.profiles["google-gemini-cli:user2@gmail.com"]).toMatchObject({
+      access: "new-access-2",
+      refresh: "new-refresh-2",
+    });
+  });
+});

src/commands/models/list.status-command.ts

@@ -528,6 +528,18 @@ export async function modelsStatusCommand(
       }
     }
 
+    // Re-read auth health after usage fetch, which may have refreshed expired tokens
+    const refreshedStore = ensureAuthProfileStore();
+    const refreshedHealth = buildAuthHealthSummary({
+      store: refreshedStore,
+      cfg,
+      warnAfterMs: DEFAULT_OAUTH_WARN_MS,
+      providers,
+    });
+    const refreshedOauthProfiles = refreshedHealth.profiles.filter(
+      (profile) => profile.type === "oauth" || profile.type === "token",
+    );
+
     const formatStatus = (status: string) => {
       if (status === "ok") return colorize(rich, theme.success, "ok");
       if (status === "static") return colorize(rich, theme.muted, "static");
@@ -536,8 +548,8 @@ export async function modelsStatusCommand(
       return colorize(rich, theme.error, "expired");
     };
 
-    const profilesByProvider = new Map<string, typeof oauthProfiles>();
-    for (const profile of oauthProfiles) {
+    const profilesByProvider = new Map<string, typeof refreshedOauthProfiles>();
+    for (const profile of refreshedOauthProfiles) {
       const current = profilesByProvider.get(profile.provider);
       if (current) current.push(profile);
       else profilesByProvider.set(profile.provider, [profile]);

src/infra/provider-usage.auth.ts

@@ -123,10 +123,10 @@ function resolveXiaomiApiKey(): string | undefined {
   return undefined;
 }
 
-async function resolveOAuthToken(params: {
+async function resolveOAuthTokens(params: {
   provider: UsageProviderId;
   agentDir?: string;
-}): Promise<ProviderAuth | null> {
+}): Promise<ProviderAuth[]> {
   const cfg = loadConfig();
   const store = ensureAuthProfileStore(params.agentDir, {
     allowKeychainPrompt: false,
@@ -143,6 +143,7 @@ async function resolveOAuthToken(params: {
     if (!deduped.includes(entry)) deduped.push(entry);
   }
 
+  const results: ProviderAuth[] = [];
   for (const profileId of deduped) {
     const cred = store.profiles[profileId];
     if (!cred || (cred.type !== "oauth" && cred.type !== "token")) continue;
@@ -161,20 +162,20 @@ async function resolveOAuthToken(params: {
         const parsed = parseGoogleToken(resolved.apiKey);
         token = parsed?.token ?? resolved.apiKey;
       }
-      return {
+      results.push({
         provider: params.provider,
         token,
         accountId:
           cred.type === "oauth" && "accountId" in cred
             ? (cred as { accountId?: string }).accountId
             : undefined,
-      };
+      });
     } catch {
       // ignore
     }
   }
 
-  return null;
+  return results;
 }
 
 function resolveOAuthProviders(agentDir?: string): UsageProviderId[] {
@@ -233,11 +234,11 @@ export async function resolveProviderAuths(params: {
     }
 
     if (!oauthProviders.includes(provider)) continue;
-    const auth = await resolveOAuthToken({
+    const resolved = await resolveOAuthTokens({
       provider,
       agentDir: params.agentDir,
     });
-    if (auth) auths.push(auth);
+    auths.push(...resolved);
   }
 
   return auths;