From fca2e830e07becd53bd43b3980c2fc13b488b5d5 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Thu, 16 Jul 2026 00:39:28 -0700 Subject: [PATCH] fix(auth): session refresh preserves OAuth providers (#108680) * fix(auth): isolate OAuth provider registries * fix(auth): remove OAuth registry back edge --- .../sessions/auth-storage-oauth-registry.ts | 14 +++ src/agents/sessions/auth-storage.ts | 21 ++-- src/agents/sessions/model-registry.test.ts | 76 +++++++++++- src/agents/sessions/model-registry.ts | 6 +- src/llm/utils/oauth/index.ts | 108 +++++++++++------- test/setup.shared.ts | 1 - 6 files changed, 168 insertions(+), 58 deletions(-) create mode 100644 src/agents/sessions/auth-storage-oauth-registry.ts diff --git a/src/agents/sessions/auth-storage-oauth-registry.ts b/src/agents/sessions/auth-storage-oauth-registry.ts new file mode 100644 index 000000000000..88799b7a0132 --- /dev/null +++ b/src/agents/sessions/auth-storage-oauth-registry.ts @@ -0,0 +1,14 @@ +import { OAuthProviderRegistry } from "../../llm/utils/oauth/index.js"; + +// Values belong to one AuthStorage object. The weak attachment keeps ModelRegistry +// on the same registry without adding lifecycle methods to the public SDK class. +const registries = new WeakMap(); + +export function getAuthStorageOAuthProviderRegistry(authStorage: object): OAuthProviderRegistry { + let registry = registries.get(authStorage); + if (!registry) { + registry = new OAuthProviderRegistry(); + registries.set(authStorage, registry); + } + return registry; +} diff --git a/src/agents/sessions/auth-storage.ts b/src/agents/sessions/auth-storage.ts index 3eb7b88c5b62..4a50887318a3 100644 --- a/src/agents/sessions/auth-storage.ts +++ b/src/agents/sessions/auth-storage.ts @@ -11,17 +11,13 @@ import { dirname, join } from "node:path"; import { findEnvKeys, getEnvApiKey } from "@openclaw/ai/internal/runtime"; import lockfile from "proper-lockfile"; import { replaceFileAtomicSync } from "../../infra/replace-file.js"; -import { - getOAuthApiKey, - getOAuthProvider, - getOAuthProviders, -} from "../../llm/utils/oauth/index.js"; import type { OAuthCredentials, OAuthLoginCallbacks, OAuthProviderId, } from "../../llm/utils/oauth/types.js"; import { getAgentDir } from "../config.js"; +import { getAuthStorageOAuthProviderRegistry } from "./auth-storage-oauth-registry.js"; import { resolveConfigValue } from "./resolve-config-value.js"; import { acquireLockSyncWithRetry } from "./storage-lock.js"; @@ -369,7 +365,7 @@ export class AuthStorage { * Login to an OAuth provider. */ async login(providerId: OAuthProviderId, callbacks: OAuthLoginCallbacks): Promise { - const provider = getOAuthProvider(providerId); + const provider = getAuthStorageOAuthProviderRegistry(this).get(providerId); if (!provider) { throw new Error(`Unknown OAuth provider: ${providerId}`); } @@ -392,7 +388,7 @@ export class AuthStorage { private async refreshOAuthTokenWithLock( providerId: OAuthProviderId, ): Promise<{ apiKey: string; newCredentials: OAuthCredentials } | null> { - const provider = getOAuthProvider(providerId); + const provider = getAuthStorageOAuthProviderRegistry(this).get(providerId); if (!provider) { return null; } @@ -418,7 +414,10 @@ export class AuthStorage { } } - const refreshed = await getOAuthApiKey(providerId, oauthCreds); + const refreshed = await getAuthStorageOAuthProviderRegistry(this).getApiKey( + providerId, + oauthCreds, + ); if (!refreshed) { return { result: null }; } @@ -461,7 +460,7 @@ export class AuthStorage { } if (cred?.type === "oauth") { - const provider = getOAuthProvider(providerId); + const provider = getAuthStorageOAuthProviderRegistry(this).get(providerId); if (!provider) { // Unknown OAuth provider, can't get API key return undefined; @@ -513,9 +512,9 @@ export class AuthStorage { } /** - * Get all registered OAuth providers + * Get all OAuth providers registered for this auth/session runtime. */ getOAuthProviders() { - return getOAuthProviders(); + return getAuthStorageOAuthProviderRegistry(this).getAll(); } } diff --git a/src/agents/sessions/model-registry.test.ts b/src/agents/sessions/model-registry.test.ts index 122152d34635..386615a40dd3 100644 --- a/src/agents/sessions/model-registry.test.ts +++ b/src/agents/sessions/model-registry.test.ts @@ -6,7 +6,7 @@ import { dirname, join } from "node:path"; import { afterEach, describe, expect, it } from "vitest"; import { PLUGIN_MODEL_CATALOG_GENERATED_BY } from "../plugin-model-catalog.js"; import { AuthStorage } from "./auth-storage.js"; -import { ModelRegistry } from "./model-registry.js"; +import { ModelRegistry, type ProviderConfigInput } from "./model-registry.js"; const PLUGIN_MODEL_CATALOG_FILE = "catalog.json"; @@ -85,6 +85,27 @@ function pluginOwnerSnapshotEntries( }; } +function oauthProviderConfig(name: string, apiKeyPrefix: string): ProviderConfigInput { + return { + oauth: { + name, + login: async () => ({ + access: "test-token-placeholder", + refresh: "test-token-placeholder", + expires: Date.now() + 60_000, + }), + async refreshToken(credentials) { + return { + ...credentials, + access: "test-token-placeholder", + expires: Date.now() + 60_000, + }; + }, + getApiKey: (credentials) => `${apiKeyPrefix}:${credentials.access}`, + }, + }; +} + afterEach(() => { for (const dir of tempDirs.splice(0)) { rmSync(dir, { recursive: true, force: true }); @@ -451,3 +472,56 @@ describe("ModelRegistry models.json auth", () => { expect(registry.find("zai", "glm-5.1")).toBeUndefined(); }); }); + +describe("ModelRegistry OAuth provider ownership", () => { + it("keeps providers isolated when another registry refreshes", async () => { + const sessionAAuth = AuthStorage.inMemory({ + "corporate-ai": { + type: "oauth", + access: "test-token-placeholder", + refresh: "test-token-placeholder", + expires: 0, + }, + }); + const sessionA = ModelRegistry.inMemory(sessionAAuth); + sessionA.registerProvider("corporate-ai", oauthProviderConfig("Corporate AI", "corporate")); + + const sessionBAuth = AuthStorage.inMemory(); + const sessionB = ModelRegistry.inMemory(sessionBAuth); + sessionB.registerProvider("team-proxy", oauthProviderConfig("Team Proxy", "team")); + + expect(sessionAAuth.getOAuthProviders().map((provider) => provider.id)).toContain( + "corporate-ai", + ); + await expect(sessionA.getApiKeyForProvider("corporate-ai")).resolves.toBe( + "corporate:test-token-placeholder", + ); + + sessionB.unregisterProvider("team-proxy"); + + expect(sessionBAuth.getOAuthProviders().map((provider) => provider.id)).not.toContain( + "team-proxy", + ); + expect(sessionAAuth.getOAuthProviders().map((provider) => provider.id)).toContain( + "corporate-ai", + ); + await expect(sessionA.getApiKeyForProvider("corporate-ai")).resolves.toBe( + "corporate:test-token-placeholder", + ); + }); + + it("keeps a built-in override local to its registry", () => { + const sessionAAuth = AuthStorage.inMemory(); + const sessionA = ModelRegistry.inMemory(sessionAAuth); + sessionA.registerProvider("anthropic", oauthProviderConfig("Corporate Anthropic", "corp")); + + const sessionBAuth = AuthStorage.inMemory(); + + expect( + sessionAAuth.getOAuthProviders().find((provider) => provider.id === "anthropic")?.name, + ).toBe("Corporate Anthropic"); + expect( + sessionBAuth.getOAuthProviders().find((provider) => provider.id === "anthropic")?.name, + ).toBe("Anthropic (Claude Pro/Max)"); + }); +}); diff --git a/src/agents/sessions/model-registry.ts b/src/agents/sessions/model-registry.ts index badf975b5481..0730802f880b 100644 --- a/src/agents/sessions/model-registry.ts +++ b/src/agents/sessions/model-registry.ts @@ -19,7 +19,6 @@ import type { OpenAIResponsesCompat, SimpleStreamOptions, } from "../../llm/types.js"; -import { registerOAuthProvider, resetOAuthProviders } from "../../llm/utils/oauth/index.js"; import type { OAuthProviderInterface } from "../../llm/utils/oauth/types.js"; import { createSubsystemLogger } from "../../logging/subsystem.js"; import { getAgentDir } from "../config.js"; @@ -30,6 +29,7 @@ import { listPluginModelCatalogFiles, type PluginModelCatalogMetadataSnapshot, } from "../plugin-model-catalog.js"; +import { getAuthStorageOAuthProviderRegistry } from "./auth-storage-oauth-registry.js"; import type { AuthStatus, AuthStorage } from "./auth-storage.js"; import { BUILT_IN_PROVIDER_DISPLAY_NAMES } from "./provider-display-names.js"; import { @@ -359,7 +359,7 @@ export class ModelRegistry { // Ensure dynamic API/OAuth registrations are rebuilt from current provider state. resetApiProviders(defaultApiRegistry); - resetOAuthProviders(); + getAuthStorageOAuthProviderRegistry(this.authStorage).reset(); this.loadModels(); @@ -863,7 +863,7 @@ export class ModelRegistry { ...config.oauth, id: providerName, }; - registerOAuthProvider(oauthProvider); + getAuthStorageOAuthProviderRegistry(this.authStorage).register(oauthProvider); } if (config.streamSimple) { diff --git a/src/llm/utils/oauth/index.ts b/src/llm/utils/oauth/index.ts index a878c2b9db7d..b9db1a6e8084 100644 --- a/src/llm/utils/oauth/index.ts +++ b/src/llm/utils/oauth/index.ts @@ -14,7 +14,7 @@ export * from "./types.js"; // ============================================================================ -// Provider Registry +// Built-in providers and instance-owned registries // ============================================================================ import { anthropicOAuthProvider } from "./anthropic.js"; @@ -28,43 +28,83 @@ const BUILT_IN_OAUTH_PROVIDERS: OAuthProviderInterface[] = [ openaiCodexOAuthProvider, ]; -const oauthProviderRegistry = new Map( - BUILT_IN_OAUTH_PROVIDERS.map((provider) => [provider.id, provider]), -); +type OAuthApiKeyResult = { newCredentials: OAuthCredentials; apiKey: string } | null; -/** - * Get an OAuth provider by ID - */ -export function getOAuthProvider(id: OAuthProviderId): OAuthProviderInterface | undefined { - return oauthProviderRegistry.get(id); +async function resolveOAuthApiKey( + provider: OAuthProviderInterface, + credentials: Record, +): Promise { + let creds = credentials[provider.id]; + if (!creds) { + return null; + } + + if (Date.now() >= creds.expires) { + try { + creds = await provider.refreshToken(creds); + } catch (error) { + throw new Error(`Failed to refresh OAuth token for ${provider.id}`, { cause: error }); + } + } + + return { newCredentials: creds, apiKey: provider.getApiKey(creds) }; } -/** - * Register a custom OAuth provider - */ -export function registerOAuthProvider(provider: OAuthProviderInterface): void { - oauthProviderRegistry.set(provider.id, provider); -} +/** Mutable OAuth provider registrations owned by one auth/session runtime. */ +export class OAuthProviderRegistry { + private providers = new Map(); -/** - * Reset OAuth providers to built-ins. - */ -export function resetOAuthProviders(): void { - oauthProviderRegistry.clear(); - for (const provider of BUILT_IN_OAUTH_PROVIDERS) { - oauthProviderRegistry.set(provider.id, provider); + constructor() { + this.reset(); + } + + get(id: OAuthProviderId): OAuthProviderInterface | undefined { + return this.providers.get(id); + } + + register(provider: OAuthProviderInterface): void { + this.providers.set(provider.id, provider); + } + + reset(): void { + this.providers.clear(); + for (const provider of BUILT_IN_OAUTH_PROVIDERS) { + this.providers.set(provider.id, provider); + } + } + + getAll(): OAuthProviderInterface[] { + return Array.from(this.providers.values()); + } + + async getApiKey( + providerId: OAuthProviderId, + credentials: Record, + ): Promise { + const provider = this.get(providerId); + if (!provider) { + throw new Error(`Unknown OAuth provider: ${providerId}`); + } + return resolveOAuthApiKey(provider, credentials); } } /** - * Get all registered OAuth providers + * Get a built-in OAuth provider by ID. + */ +function getOAuthProvider(id: OAuthProviderId): OAuthProviderInterface | undefined { + return BUILT_IN_OAUTH_PROVIDERS.find((provider) => provider.id === id); +} + +/** + * Get all built-in OAuth providers. */ export function getOAuthProviders(): OAuthProviderInterface[] { - return Array.from(oauthProviderRegistry.values()); + return [...BUILT_IN_OAUTH_PROVIDERS]; } // ============================================================================ -// High-level API (uses provider registry) +// High-level built-in provider API // ============================================================================ /** @@ -82,21 +122,5 @@ export async function getOAuthApiKey( if (!provider) { throw new Error(`Unknown OAuth provider: ${providerId}`); } - - let creds = credentials[providerId]; - if (!creds) { - return null; - } - - // Refresh if expired - if (Date.now() >= creds.expires) { - try { - creds = await provider.refreshToken(creds); - } catch (error) { - throw new Error(`Failed to refresh OAuth token for ${providerId}`, { cause: error }); - } - } - - const apiKey = provider.getApiKey(creds); - return { newCredentials: creds, apiKey }; + return resolveOAuthApiKey(provider, credentials); } diff --git a/test/setup.shared.ts b/test/setup.shared.ts index fe75fcbf0f9a..f72269f2363b 100644 --- a/test/setup.shared.ts +++ b/test/setup.shared.ts @@ -15,7 +15,6 @@ vi.mock("../src/llm/oauth.js", () => ({ ...args, ), ), - resetOAuthProviders: vi.fn(), })); vi.mock("@mariozechner/clipboard", () => ({