diff --git a/src/agents/auth-profiles.ensureauthprofilestore.test.ts b/src/agents/auth-profiles.ensureauthprofilestore.test.ts index 6fdf30b31d2..ca6a062d747 100644 --- a/src/agents/auth-profiles.ensureauthprofilestore.test.ts +++ b/src/agents/auth-profiles.ensureauthprofilestore.test.ts @@ -3,13 +3,13 @@ import os from "node:os"; import path from "node:path"; import { afterEach, describe, expect, it, vi } from "vitest"; import type { ProviderExternalAuthProfile } from "../plugins/provider-external-auth.types.js"; +import { AUTH_STORE_VERSION, log } from "./auth-profiles/constants.js"; import { clearRuntimeAuthProfileStoreSnapshots, ensureAuthProfileStore, loadAuthProfileStoreForRuntime, saveAuthProfileStore, -} from "./auth-profiles.js"; -import { AUTH_STORE_VERSION, log } from "./auth-profiles/constants.js"; +} from "./auth-profiles/store.js"; import type { AuthProfileCredential } from "./auth-profiles/types.js"; const resolveExternalAuthProfilesWithPluginsMock = vi.hoisted(() => diff --git a/src/agents/auth-profiles.getsoonestcooldownexpiry.test.ts b/src/agents/auth-profiles.getsoonestcooldownexpiry.test.ts index 98653bc89fc..51aee236697 100644 --- a/src/agents/auth-profiles.getsoonestcooldownexpiry.test.ts +++ b/src/agents/auth-profiles.getsoonestcooldownexpiry.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it } from "vitest"; -import type { AuthProfileStore } from "./auth-profiles.js"; -import { getSoonestCooldownExpiry } from "./auth-profiles.js"; +import type { AuthProfileStore } from "./auth-profiles/types.js"; +import { getSoonestCooldownExpiry } from "./auth-profiles/usage-state.js"; function makeStore(usageStats?: AuthProfileStore["usageStats"]): AuthProfileStore { return { diff --git a/src/agents/auth-profiles.store-cache.test.ts b/src/agents/auth-profiles.store-cache.test.ts index ea23d9c490b..e16cef9f295 100644 --- a/src/agents/auth-profiles.store-cache.test.ts +++ b/src/agents/auth-profiles.store-cache.test.ts @@ -2,8 +2,11 @@ import fs from "node:fs"; import os from "node:os"; import path from "node:path"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import { clearRuntimeAuthProfileStoreSnapshots, ensureAuthProfileStore } from "./auth-profiles.js"; import { AUTH_STORE_VERSION } from "./auth-profiles/constants.js"; +import { + clearRuntimeAuthProfileStoreSnapshots, + ensureAuthProfileStore, +} from "./auth-profiles/store.js"; import type { OAuthCredential } from "./auth-profiles/types.js"; type RuntimeOnlyOverlay = { profileId: string; credential: OAuthCredential }; diff --git a/src/agents/auth-profiles/usage-state.ts b/src/agents/auth-profiles/usage-state.ts index 0bf0841d5cd..90774f52b52 100644 --- a/src/agents/auth-profiles/usage-state.ts +++ b/src/agents/auth-profiles/usage-state.ts @@ -64,6 +64,54 @@ export function isProfileInCooldown( return unusableUntil ? ts < unusableUntil : false; } +/** + * Return the soonest `unusableUntil` timestamp (ms epoch) among the given + * profiles, or `null` when no profile has a recorded cooldown. Note: the + * returned timestamp may be in the past if the cooldown has already expired. + */ +export function getSoonestCooldownExpiry( + store: AuthProfileStore, + profileIds: string[], + options?: { now?: number; forModel?: string }, +): number | null { + const ts = options?.now ?? Date.now(); + let soonest: number | null = null; + let latestMatchingModelCooldown: number | null = null; + for (const id of profileIds) { + const stats = store.usageStats?.[id]; + if (!stats) { + continue; + } + if (shouldBypassModelScopedCooldown(stats, ts, options?.forModel)) { + continue; + } + const until = resolveProfileUnusableUntil(stats); + if (typeof until !== "number" || !Number.isFinite(until) || until <= 0) { + continue; + } + const matchingModelScopedCooldown = + options?.forModel && + stats.cooldownReason === "rate_limit" && + stats.cooldownModel === options.forModel && + !isActiveUnusableWindow(stats.disabledUntil, ts); + if (matchingModelScopedCooldown) { + latestMatchingModelCooldown = + latestMatchingModelCooldown === null ? until : Math.max(latestMatchingModelCooldown, until); + continue; + } + if (soonest === null || until < soonest) { + soonest = until; + } + } + if (soonest === null) { + return latestMatchingModelCooldown; + } + if (latestMatchingModelCooldown === null) { + return soonest; + } + return Math.min(soonest, latestMatchingModelCooldown); +} + /** * Clear expired cooldowns from all profiles in the store. * diff --git a/src/agents/auth-profiles/usage.ts b/src/agents/auth-profiles/usage.ts index 06d96fc70db..6d292f0e1c8 100644 --- a/src/agents/auth-profiles/usage.ts +++ b/src/agents/auth-profiles/usage.ts @@ -9,10 +9,10 @@ import { isAuthCooldownBypassedForProvider, isProfileInCooldown, resolveProfileUnusableUntil, - shouldBypassModelScopedCooldown, } from "./usage-state.js"; export { clearExpiredCooldowns, + getSoonestCooldownExpiry, isProfileInCooldown, resolveProfileUnusableUntil, } from "./usage-state.js"; @@ -309,54 +309,6 @@ export function resolveProfilesUnavailableReason(params: { return best; } -/** - * Return the soonest `unusableUntil` timestamp (ms epoch) among the given - * profiles, or `null` when no profile has a recorded cooldown. Note: the - * returned timestamp may be in the past if the cooldown has already expired. - */ -export function getSoonestCooldownExpiry( - store: AuthProfileStore, - profileIds: string[], - options?: { now?: number; forModel?: string }, -): number | null { - const ts = options?.now ?? Date.now(); - let soonest: number | null = null; - let latestMatchingModelCooldown: number | null = null; - for (const id of profileIds) { - const stats = store.usageStats?.[id]; - if (!stats) { - continue; - } - if (shouldBypassModelScopedCooldown(stats, ts, options?.forModel)) { - continue; - } - const until = resolveProfileUnusableUntil(stats); - if (typeof until !== "number" || !Number.isFinite(until) || until <= 0) { - continue; - } - const matchingModelScopedCooldown = - options?.forModel && - stats.cooldownReason === "rate_limit" && - stats.cooldownModel === options.forModel && - !isActiveUnusableWindow(stats.disabledUntil, ts); - if (matchingModelScopedCooldown) { - latestMatchingModelCooldown = - latestMatchingModelCooldown === null ? until : Math.max(latestMatchingModelCooldown, until); - continue; - } - if (soonest === null || until < soonest) { - soonest = until; - } - } - if (soonest === null) { - return latestMatchingModelCooldown; - } - if (latestMatchingModelCooldown === null) { - return soonest; - } - return Math.min(soonest, latestMatchingModelCooldown); -} - /** * Mark a profile as successfully used. Resets error count and updates lastUsed. * Uses store lock to avoid overwriting concurrent usage updates. diff --git a/src/agents/pi-model-discovery.synthetic-auth.test.ts b/src/agents/pi-model-discovery.synthetic-auth.test.ts index d2634a6f455..8208fbaee38 100644 --- a/src/agents/pi-model-discovery.synthetic-auth.test.ts +++ b/src/agents/pi-model-discovery.synthetic-auth.test.ts @@ -2,7 +2,7 @@ import fs from "node:fs/promises"; import os from "node:os"; import path from "node:path"; import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; -import { saveAuthProfileStore } from "./auth-profiles.js"; +import { saveAuthProfileStore } from "./auth-profiles/store.js"; const resolveRuntimeSyntheticAuthProviderRefs = vi.hoisted(() => vi.fn(() => ["claude-cli"]));