diff --git a/CHANGELOG.md b/CHANGELOG.md index a134359f5..68661b8e5 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/src/agents/auth-profiles/oauth-refresh.test.ts b/src/agents/auth-profiles/oauth-refresh.test.ts new file mode 100644 index 000000000..99a2fc532 --- /dev/null +++ b/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", + }); + }); +}); diff --git a/src/commands/models/list.status-command.ts b/src/commands/models/list.status-command.ts index 3e18dc41f..4ba7c041c 100644 --- a/src/commands/models/list.status-command.ts +++ b/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(); - for (const profile of oauthProfiles) { + const profilesByProvider = new Map(); + for (const profile of refreshedOauthProfiles) { const current = profilesByProvider.get(profile.provider); if (current) current.push(profile); else profilesByProvider.set(profile.provider, [profile]); diff --git a/src/infra/provider-usage.auth.ts b/src/infra/provider-usage.auth.ts index e0d9a6ef9..e22c2c3fa 100644 --- a/src/infra/provider-usage.auth.ts +++ b/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 { +}): Promise { 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;