From d89b409add4412ca594ab6f4e9bf77bc917d7aa8 Mon Sep 17 00:00:00 2001 From: yushen Date: Tue, 3 Feb 2026 16:55:42 +0800 Subject: [PATCH] feat(agent): add auth profile rotation and cooldown module Implement multi-profile auth system with exponential-backoff cooldowns, modeled after OpenClaw's auth-profiles system. Includes types, constants, persistent store (auth-profiles.json), usage tracking with cooldown calculation, and profile ordering with round-robin and explicit modes. Co-Authored-By: Claude Opus 4.5 --- src/agent/auth-profiles/constants.ts | 45 +++++++ src/agent/auth-profiles/index.ts | 45 +++++++ src/agent/auth-profiles/order.test.ts | 176 ++++++++++++++++++++++++++ src/agent/auth-profiles/order.ts | 94 ++++++++++++++ src/agent/auth-profiles/store.ts | 84 ++++++++++++ src/agent/auth-profiles/types.ts | 47 +++++++ src/agent/auth-profiles/usage.test.ts | 147 +++++++++++++++++++++ src/agent/auth-profiles/usage.ts | 169 +++++++++++++++++++++++++ 8 files changed, 807 insertions(+) create mode 100644 src/agent/auth-profiles/constants.ts create mode 100644 src/agent/auth-profiles/index.ts create mode 100644 src/agent/auth-profiles/order.test.ts create mode 100644 src/agent/auth-profiles/order.ts create mode 100644 src/agent/auth-profiles/store.ts create mode 100644 src/agent/auth-profiles/types.ts create mode 100644 src/agent/auth-profiles/usage.test.ts create mode 100644 src/agent/auth-profiles/usage.ts diff --git a/src/agent/auth-profiles/constants.ts b/src/agent/auth-profiles/constants.ts new file mode 100644 index 00000000..d519a246 --- /dev/null +++ b/src/agent/auth-profiles/constants.ts @@ -0,0 +1,45 @@ +/** + * Auth Profile Constants + * + * Cooldown timings, store version, and file names. + */ + +/** Store format version */ +export const AUTH_STORE_VERSION = 1; + +/** Runtime store filename (inside ~/.super-multica/) */ +export const AUTH_PROFILE_STORE_FILENAME = "auth-profiles.json"; + +// ============================================================ +// Non-billing cooldown (rate_limit, auth, timeout, unknown) +// Progression: 1min -> 5min -> 25min -> 1hr (cap) +// Formula: min(MAX, BASE * FACTOR ^ min(errorCount - 1, 3)) +// ============================================================ + +/** Base cooldown duration in milliseconds (1 minute) */ +export const COOLDOWN_BASE_MS = 60_000; + +/** Exponential factor for cooldown progression */ +export const COOLDOWN_FACTOR = 5; + +/** Maximum cooldown duration in milliseconds (1 hour) */ +export const COOLDOWN_MAX_MS = 3_600_000; + +// ============================================================ +// Billing disable (longer backoff for payment/quota issues) +// Progression: 5h -> 10h -> 20h -> 24h (cap) +// Formula: min(MAX_HOURS, BASE_HOURS * 2 ^ (count - 1)) +// ============================================================ + +/** Base billing disable duration in hours */ +export const BILLING_BACKOFF_HOURS = 5; + +/** Maximum billing disable duration in hours */ +export const BILLING_MAX_HOURS = 24; + +// ============================================================ +// Failure window +// ============================================================ + +/** Failure window in milliseconds (24 hours) — errors older than this are forgotten */ +export const FAILURE_WINDOW_MS = 24 * 60 * 60 * 1000; diff --git a/src/agent/auth-profiles/index.ts b/src/agent/auth-profiles/index.ts new file mode 100644 index 00000000..bd87cbb7 --- /dev/null +++ b/src/agent/auth-profiles/index.ts @@ -0,0 +1,45 @@ +/** + * Auth Profiles — barrel export + */ + +export type { + AuthProfileFailureReason, + AuthProfileStore, + ProfileUsageStats, + ResolvedProfileAuth, +} from "./types.js"; + +export { + AUTH_STORE_VERSION, + AUTH_PROFILE_STORE_FILENAME, + COOLDOWN_BASE_MS, + COOLDOWN_FACTOR, + COOLDOWN_MAX_MS, + BILLING_BACKOFF_HOURS, + BILLING_MAX_HOURS, + FAILURE_WINDOW_MS, +} from "./constants.js"; + +export { + resolveAuthStorePath, + loadAuthProfileStore, + saveAuthProfileStore, + updateAuthProfileStore, +} from "./store.js"; + +export { + listProfilesForProvider, + resolveAuthProfileOrder, +} from "./order.js"; + +export { + isProfileInCooldown, + resolveProfileUnusableUntil, + calculateCooldownMs, + calculateBillingDisableMs, + computeNextProfileUsageStats, + markAuthProfileFailure, + markAuthProfileUsed, + markAuthProfileGood, + clearAuthProfileCooldown, +} from "./usage.js"; diff --git a/src/agent/auth-profiles/order.test.ts b/src/agent/auth-profiles/order.test.ts new file mode 100644 index 00000000..dd729719 --- /dev/null +++ b/src/agent/auth-profiles/order.test.ts @@ -0,0 +1,176 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { resolveAuthProfileOrder, listProfilesForProvider } from "./order.js"; +import type { AuthProfileStore } from "./types.js"; + +// Mock credentialManager +vi.mock("../credentials.js", () => { + let _profiles: Record = {}; + let _order: Record = {}; + + return { + credentialManager: { + listProfileIdsForProvider(provider: string): string[] { + return Object.keys(_profiles).filter( + (key) => key === provider || key.startsWith(`${provider}:`), + ); + }, + getLlmOrder(provider: string): string[] | undefined { + return _order[provider]; + }, + // Test helpers + __setProfiles(profiles: Record) { + _profiles = profiles; + }, + __setOrder(order: Record) { + _order = order; + }, + }, + }; +}); + +// Import the mock to use test helpers +import { credentialManager } from "../credentials.js"; +const mock = credentialManager as unknown as { + __setProfiles: (p: Record) => void; + __setOrder: (o: Record) => void; +}; + +beforeEach(() => { + mock.__setProfiles({}); + mock.__setOrder({}); +}); + +// ============================================================ +// listProfilesForProvider +// ============================================================ + +describe("listProfilesForProvider", () => { + it("returns profiles matching the provider", () => { + mock.__setProfiles({ + anthropic: { apiKey: "sk-1" }, + "anthropic:backup": { apiKey: "sk-2" }, + openai: { apiKey: "sk-3" }, + }); + expect(listProfilesForProvider("anthropic")).toEqual([ + "anthropic", + "anthropic:backup", + ]); + }); + + it("returns empty array when no profiles match", () => { + mock.__setProfiles({ openai: { apiKey: "sk-1" } }); + expect(listProfilesForProvider("anthropic")).toEqual([]); + }); +}); + +// ============================================================ +// resolveAuthProfileOrder +// ============================================================ + +describe("resolveAuthProfileOrder", () => { + const now = 1_000_000; + + it("returns round-robin order by lastUsed when no explicit order", () => { + mock.__setProfiles({ + "anthropic": { apiKey: "sk-1" }, + "anthropic:b": { apiKey: "sk-2" }, + "anthropic:c": { apiKey: "sk-3" }, + }); + const store: AuthProfileStore = { + version: 1, + usageStats: { + "anthropic": { lastUsed: 300 }, + "anthropic:b": { lastUsed: 100 }, + "anthropic:c": { lastUsed: 200 }, + }, + }; + + const order = resolveAuthProfileOrder("anthropic", store, now); + // Sorted by lastUsed ascending: b(100) -> c(200) -> default(300) + expect(order).toEqual(["anthropic:b", "anthropic:c", "anthropic"]); + }); + + it("respects explicit order from config", () => { + mock.__setProfiles({ + "anthropic": { apiKey: "sk-1" }, + "anthropic:b": { apiKey: "sk-2" }, + "anthropic:c": { apiKey: "sk-3" }, + }); + mock.__setOrder({ anthropic: ["anthropic:c", "anthropic", "anthropic:b"] }); + + const store: AuthProfileStore = { version: 1 }; + const order = resolveAuthProfileOrder("anthropic", store, now); + expect(order).toEqual(["anthropic:c", "anthropic", "anthropic:b"]); + }); + + it("pushes cooldown profiles to the end", () => { + mock.__setProfiles({ + "anthropic": { apiKey: "sk-1" }, + "anthropic:b": { apiKey: "sk-2" }, + "anthropic:c": { apiKey: "sk-3" }, + }); + const store: AuthProfileStore = { + version: 1, + usageStats: { + "anthropic": { lastUsed: 100 }, + "anthropic:b": { lastUsed: 200, cooldownUntil: now + 5000 }, + "anthropic:c": { lastUsed: 300 }, + }, + }; + + const order = resolveAuthProfileOrder("anthropic", store, now); + // anthropic and anthropic:c are available; anthropic:b is in cooldown -> pushed to end + expect(order).toEqual(["anthropic", "anthropic:c", "anthropic:b"]); + }); + + it("sorts cooldown profiles by earliest recovery", () => { + mock.__setProfiles({ + "anthropic": { apiKey: "sk-1" }, + "anthropic:b": { apiKey: "sk-2" }, + "anthropic:c": { apiKey: "sk-3" }, + }); + const store: AuthProfileStore = { + version: 1, + usageStats: { + "anthropic": { cooldownUntil: now + 10_000 }, + "anthropic:b": { cooldownUntil: now + 1_000 }, + "anthropic:c": { cooldownUntil: now + 5_000 }, + }, + }; + + const order = resolveAuthProfileOrder("anthropic", store, now); + // All in cooldown, sorted by soonest recovery + expect(order).toEqual(["anthropic:b", "anthropic:c", "anthropic"]); + }); + + it("deduplicates profile IDs", () => { + mock.__setProfiles({ + "anthropic": { apiKey: "sk-1" }, + "anthropic:b": { apiKey: "sk-2" }, + }); + // Explicit order has duplicate + mock.__setOrder({ anthropic: ["anthropic", "anthropic", "anthropic:b"] }); + + const store: AuthProfileStore = { version: 1 }; + const order = resolveAuthProfileOrder("anthropic", store, now); + expect(order).toEqual(["anthropic", "anthropic:b"]); + }); + + it("appends unlisted profiles to explicit order", () => { + mock.__setProfiles({ + "anthropic": { apiKey: "sk-1" }, + "anthropic:b": { apiKey: "sk-2" }, + "anthropic:c": { apiKey: "sk-3" }, + }); + // Only lists one profile in explicit order + mock.__setOrder({ anthropic: ["anthropic:b"] }); + + const store: AuthProfileStore = { version: 1 }; + const order = resolveAuthProfileOrder("anthropic", store, now); + // anthropic:b first (explicit), then the rest + expect(order[0]).toBe("anthropic:b"); + expect(order).toHaveLength(3); + expect(order).toContain("anthropic"); + expect(order).toContain("anthropic:c"); + }); +}); diff --git a/src/agent/auth-profiles/order.ts b/src/agent/auth-profiles/order.ts new file mode 100644 index 00000000..16055e81 --- /dev/null +++ b/src/agent/auth-profiles/order.ts @@ -0,0 +1,94 @@ +/** + * Auth Profile Ordering + * + * Determines the order in which auth profiles are tried for a given provider. + * Supports explicit ordering (from credentials.json5) and automatic round-robin + * based on lastUsed time. Profiles in cooldown are pushed to the end. + */ + +import { credentialManager } from "../credentials.js"; +import type { AuthProfileStore, ProfileUsageStats } from "./types.js"; +import { isProfileInCooldown, resolveProfileUnusableUntil } from "./usage.js"; + +// ============================================================ +// Profile discovery +// ============================================================ + +/** + * List all profile IDs from credentials.json5 that belong to a given provider. + * A profile matches if its key equals the provider exactly or starts with "provider:". + */ +export function listProfilesForProvider(provider: string): string[] { + return credentialManager.listProfileIdsForProvider(provider); +} + +// ============================================================ +// Ordering +// ============================================================ + +/** + * Resolve the ordered list of profile IDs to try for a given provider. + * + * Strategy: + * 1. If credentials.json5 has `llm.order[provider]`, use that explicit order. + * 2. Otherwise, use round-robin ordered by `lastUsed` ascending (oldest first). + * + * In both cases, profiles currently in cooldown are pushed to the end, + * sorted by earliest cooldown expiry (soonest-to-recover first). + */ +export function resolveAuthProfileOrder( + provider: string, + store: AuthProfileStore, + now?: number, +): string[] { + const ts = now ?? Date.now(); + + // Gather candidates + const explicitOrder = credentialManager.getLlmOrder(provider); + const allProfiles = listProfilesForProvider(provider); + + let candidates: string[]; + if (explicitOrder && explicitOrder.length > 0) { + // Use explicit order, filter to only existing profiles + const profileSet = new Set(allProfiles); + candidates = explicitOrder.filter((id) => profileSet.has(id)); + // Append any profiles not in the explicit order + for (const id of allProfiles) { + if (!candidates.includes(id)) { + candidates.push(id); + } + } + } else { + // Round-robin by lastUsed (oldest first) + candidates = [...allProfiles].sort((a, b) => { + const statsA = store.usageStats?.[a]; + const statsB = store.usageStats?.[b]; + return (statsA?.lastUsed ?? 0) - (statsB?.lastUsed ?? 0); + }); + } + + // Deduplicate + candidates = [...new Set(candidates)]; + + // Partition into available and in-cooldown + const available: string[] = []; + const inCooldown: string[] = []; + + for (const id of candidates) { + const stats = store.usageStats?.[id]; + if (stats && isProfileInCooldown(stats, ts)) { + inCooldown.push(id); + } else { + available.push(id); + } + } + + // Sort cooldown profiles by soonest recovery + inCooldown.sort((a, b) => { + const statsA = store.usageStats?.[a] ?? {}; + const statsB = store.usageStats?.[b] ?? {}; + return resolveProfileUnusableUntil(statsA) - resolveProfileUnusableUntil(statsB); + }); + + return [...available, ...inCooldown]; +} diff --git a/src/agent/auth-profiles/store.ts b/src/agent/auth-profiles/store.ts new file mode 100644 index 00000000..3bb0480f --- /dev/null +++ b/src/agent/auth-profiles/store.ts @@ -0,0 +1,84 @@ +/** + * Auth Profile Store + * + * Persistence layer for auth profile runtime state. + * Stores usage stats, cooldowns, and last-good info in ~/.super-multica/auth-profiles.json. + */ + +import { existsSync, readFileSync, writeFileSync, mkdirSync } from "node:fs"; +import { join, dirname } from "node:path"; +import { DATA_DIR } from "../../shared/paths.js"; +import { AUTH_STORE_VERSION, AUTH_PROFILE_STORE_FILENAME } from "./constants.js"; +import type { AuthProfileStore } from "./types.js"; + +// ============================================================ +// Paths +// ============================================================ + +/** Resolve the auth profile store file path */ +export function resolveAuthStorePath(): string { + return join(DATA_DIR, AUTH_PROFILE_STORE_FILENAME); +} + +// ============================================================ +// Load / Save +// ============================================================ + +function createEmptyStore(): AuthProfileStore { + return { version: AUTH_STORE_VERSION }; +} + +function coerceStore(raw: unknown): AuthProfileStore { + if (!raw || typeof raw !== "object") return createEmptyStore(); + + const obj = raw as Record; + const store: AuthProfileStore = { + version: typeof obj.version === "number" ? obj.version : AUTH_STORE_VERSION, + }; + + if (obj.lastGood && typeof obj.lastGood === "object") { + store.lastGood = obj.lastGood as Record; + } + if (obj.usageStats && typeof obj.usageStats === "object") { + store.usageStats = obj.usageStats as AuthProfileStore["usageStats"]; + } + + return store; +} + +/** Load auth profile store from disk. Returns empty store if file doesn't exist. */ +export function loadAuthProfileStore(): AuthProfileStore { + const storePath = resolveAuthStorePath(); + if (!existsSync(storePath)) return createEmptyStore(); + + try { + const raw = readFileSync(storePath, "utf8"); + return coerceStore(JSON.parse(raw)); + } catch { + return createEmptyStore(); + } +} + +/** Save auth profile store to disk */ +export function saveAuthProfileStore(store: AuthProfileStore): void { + const storePath = resolveAuthStorePath(); + const dir = dirname(storePath); + if (!existsSync(dir)) { + mkdirSync(dir, { recursive: true }); + } + writeFileSync(storePath, JSON.stringify(store, null, 2), "utf8"); +} + +/** + * Atomic load-update-save cycle. + * The updater receives the current store and should mutate it in place. + * Returns the updated store. + */ +export function updateAuthProfileStore( + updater: (store: AuthProfileStore) => void, +): AuthProfileStore { + const store = loadAuthProfileStore(); + updater(store); + saveAuthProfileStore(store); + return store; +} diff --git a/src/agent/auth-profiles/types.ts b/src/agent/auth-profiles/types.ts new file mode 100644 index 00000000..e4e34eb2 --- /dev/null +++ b/src/agent/auth-profiles/types.ts @@ -0,0 +1,47 @@ +/** + * Auth Profile Types + * + * Type definitions for the auth profile rotation and cooldown system. + */ + +/** Reason for an auth profile failure, determines cooldown behavior */ +export type AuthProfileFailureReason = + | "auth" + | "rate_limit" + | "billing" + | "timeout" + | "unknown"; + +/** Per-profile usage and cooldown state (persisted in auth-profiles.json) */ +export type ProfileUsageStats = { + /** Timestamp of last successful use */ + lastUsed?: number | undefined; + /** Cooldown expiry for non-billing failures (rate_limit, auth, timeout, unknown) */ + cooldownUntil?: number | undefined; + /** Disable expiry for billing failures (longer backoff) */ + disabledUntil?: number | undefined; + /** Reason for the current disable period */ + disabledReason?: AuthProfileFailureReason | undefined; + /** Consecutive error count (resets on success or after failure window) */ + errorCount?: number | undefined; + /** Per-reason failure counts within the failure window */ + failureCounts?: Partial> | undefined; + /** Timestamp of the last failure (used for failure window expiry) */ + lastFailureAt?: number | undefined; +}; + +/** Persisted runtime store for auth profile state */ +export type AuthProfileStore = { + version: number; + /** Last known good profile per provider */ + lastGood?: Record | undefined; + /** Per-profile usage/cooldown stats */ + usageStats?: Record | undefined; +}; + +/** Resolved auth info returned by profile-aware key resolution */ +export type ResolvedProfileAuth = { + apiKey: string; + profileId: string; + provider: string; +}; diff --git a/src/agent/auth-profiles/usage.test.ts b/src/agent/auth-profiles/usage.test.ts new file mode 100644 index 00000000..daaeed8e --- /dev/null +++ b/src/agent/auth-profiles/usage.test.ts @@ -0,0 +1,147 @@ +import { describe, it, expect } from "vitest"; +import { + calculateCooldownMs, + calculateBillingDisableMs, + computeNextProfileUsageStats, + isProfileInCooldown, + resolveProfileUnusableUntil, +} from "./usage.js"; +import { + COOLDOWN_BASE_MS, + COOLDOWN_MAX_MS, + FAILURE_WINDOW_MS, +} from "./constants.js"; +import type { ProfileUsageStats } from "./types.js"; + +// ============================================================ +// calculateCooldownMs +// ============================================================ + +describe("calculateCooldownMs", () => { + it("applies exponential backoff with a 1h cap", () => { + expect(calculateCooldownMs(1)).toBe(60_000); // 1 min + expect(calculateCooldownMs(2)).toBe(5 * 60_000); // 5 min + expect(calculateCooldownMs(3)).toBe(25 * 60_000); // 25 min + expect(calculateCooldownMs(4)).toBe(60 * 60_000); // 1 hour (cap) + expect(calculateCooldownMs(5)).toBe(60 * 60_000); // 1 hour (cap) + expect(calculateCooldownMs(100)).toBe(60 * 60_000); // still capped + }); + + it("returns 0 for errorCount <= 0", () => { + expect(calculateCooldownMs(0)).toBe(0); + expect(calculateCooldownMs(-1)).toBe(0); + }); +}); + +// ============================================================ +// calculateBillingDisableMs +// ============================================================ + +describe("calculateBillingDisableMs", () => { + it("applies exponential backoff with a 24h cap", () => { + const h = 60 * 60 * 1000; + expect(calculateBillingDisableMs(1)).toBe(5 * h); // 5h + expect(calculateBillingDisableMs(2)).toBe(10 * h); // 10h + expect(calculateBillingDisableMs(3)).toBe(20 * h); // 20h + expect(calculateBillingDisableMs(4)).toBe(24 * h); // 24h (cap) + expect(calculateBillingDisableMs(5)).toBe(24 * h); // still capped + }); + + it("returns 0 for count <= 0", () => { + expect(calculateBillingDisableMs(0)).toBe(0); + expect(calculateBillingDisableMs(-1)).toBe(0); + }); +}); + +// ============================================================ +// isProfileInCooldown / resolveProfileUnusableUntil +// ============================================================ + +describe("isProfileInCooldown", () => { + const now = 1_000_000; + + it("returns false for empty stats", () => { + expect(isProfileInCooldown({}, now)).toBe(false); + }); + + it("returns true when cooldownUntil is in the future", () => { + expect(isProfileInCooldown({ cooldownUntil: now + 1000 }, now)).toBe(true); + }); + + it("returns false when cooldownUntil has passed", () => { + expect(isProfileInCooldown({ cooldownUntil: now - 1 }, now)).toBe(false); + }); + + it("returns true when disabledUntil is in the future", () => { + expect(isProfileInCooldown({ disabledUntil: now + 1000 }, now)).toBe(true); + }); + + it("uses max of cooldownUntil and disabledUntil", () => { + const stats: ProfileUsageStats = { + cooldownUntil: now - 1, + disabledUntil: now + 5000, + }; + expect(isProfileInCooldown(stats, now)).toBe(true); + expect(resolveProfileUnusableUntil(stats)).toBe(now + 5000); + }); +}); + +// ============================================================ +// computeNextProfileUsageStats +// ============================================================ + +describe("computeNextProfileUsageStats", () => { + const now = 1_700_000_000_000; + + it("increments errorCount and sets cooldown for non-billing failure", () => { + const next = computeNextProfileUsageStats({}, "rate_limit", now); + expect(next.errorCount).toBe(1); + expect(next.lastFailureAt).toBe(now); + expect(next.cooldownUntil).toBe(now + COOLDOWN_BASE_MS); + expect(next.failureCounts?.rate_limit).toBe(1); + expect(next.disabledUntil).toBeUndefined(); + }); + + it("applies exponential backoff on consecutive failures", () => { + const stats: ProfileUsageStats = { + errorCount: 2, + lastFailureAt: now - 1000, + failureCounts: { rate_limit: 2 }, + }; + const next = computeNextProfileUsageStats(stats, "rate_limit", now); + expect(next.errorCount).toBe(3); + // Error 3 -> 25 min cooldown + expect(next.cooldownUntil).toBe(now + 25 * 60_000); + }); + + it("sets disabledUntil for billing failures (~5h by default)", () => { + const next = computeNextProfileUsageStats({}, "billing", now); + expect(next.errorCount).toBe(1); + expect(next.disabledUntil).toBe(now + 5 * 60 * 60 * 1000); + expect(next.disabledReason).toBe("billing"); + expect(next.failureCounts?.billing).toBe(1); + }); + + it("resets counters when lastFailureAt is outside the failure window", () => { + const oldFailure = now - FAILURE_WINDOW_MS - 1000; + const stats: ProfileUsageStats = { + errorCount: 5, + lastFailureAt: oldFailure, + failureCounts: { auth: 3, rate_limit: 2 }, + }; + const next = computeNextProfileUsageStats(stats, "auth", now); + // Counters reset, so this is treated as error #1 + expect(next.errorCount).toBe(1); + expect(next.failureCounts?.auth).toBe(1); + expect(next.cooldownUntil).toBe(now + COOLDOWN_BASE_MS); + }); + + it("caps cooldown at COOLDOWN_MAX_MS", () => { + const stats: ProfileUsageStats = { + errorCount: 10, + lastFailureAt: now - 1000, + }; + const next = computeNextProfileUsageStats(stats, "unknown", now); + expect(next.cooldownUntil).toBe(now + COOLDOWN_MAX_MS); + }); +}); diff --git a/src/agent/auth-profiles/usage.ts b/src/agent/auth-profiles/usage.ts new file mode 100644 index 00000000..181c78de --- /dev/null +++ b/src/agent/auth-profiles/usage.ts @@ -0,0 +1,169 @@ +/** + * Auth Profile Usage Tracking + * + * Tracks per-profile usage, computes cooldown durations with exponential backoff, + * and manages failure/success state transitions. + */ + +import { + COOLDOWN_BASE_MS, + COOLDOWN_FACTOR, + COOLDOWN_MAX_MS, + BILLING_BACKOFF_HOURS, + BILLING_MAX_HOURS, + FAILURE_WINDOW_MS, +} from "./constants.js"; +import { updateAuthProfileStore } from "./store.js"; +import type { + AuthProfileFailureReason, + AuthProfileStore, + ProfileUsageStats, +} from "./types.js"; + +// ============================================================ +// Cooldown checks +// ============================================================ + +/** Returns the timestamp until which a profile is unusable (0 if available) */ +export function resolveProfileUnusableUntil(stats: ProfileUsageStats): number { + return Math.max(stats.cooldownUntil ?? 0, stats.disabledUntil ?? 0); +} + +/** Check if a profile is currently in cooldown or disabled */ +export function isProfileInCooldown(stats: ProfileUsageStats, now?: number): boolean { + return resolveProfileUnusableUntil(stats) > (now ?? Date.now()); +} + +// ============================================================ +// Cooldown duration calculation +// ============================================================ + +/** + * Calculate non-billing cooldown duration in milliseconds. + * Exponential backoff: 1min -> 5min -> 25min -> 1hr (cap). + * + * Formula: min(COOLDOWN_MAX_MS, COOLDOWN_BASE_MS * COOLDOWN_FACTOR ^ min(errorCount - 1, 3)) + */ +export function calculateCooldownMs(errorCount: number): number { + if (errorCount <= 0) return 0; + const exponent = Math.min(errorCount - 1, 3); + return Math.min(COOLDOWN_MAX_MS, COOLDOWN_BASE_MS * COOLDOWN_FACTOR ** exponent); +} + +/** + * Calculate billing disable duration in milliseconds. + * Exponential backoff: 5h -> 10h -> 20h -> 24h (cap). + * + * Formula: min(BILLING_MAX_HOURS, BILLING_BACKOFF_HOURS * 2 ^ (count - 1)) * hours_to_ms + */ +export function calculateBillingDisableMs(billingFailCount: number): number { + if (billingFailCount <= 0) return 0; + const hours = Math.min( + BILLING_MAX_HOURS, + BILLING_BACKOFF_HOURS * 2 ** (billingFailCount - 1), + ); + return hours * 60 * 60 * 1000; +} + +// ============================================================ +// State transitions +// ============================================================ + +function ensureUsageStats(store: AuthProfileStore, profileId: string): ProfileUsageStats { + if (!store.usageStats) store.usageStats = {}; + if (!store.usageStats[profileId]) store.usageStats[profileId] = {}; + return store.usageStats[profileId]; +} + +/** + * Compute updated usage stats after a failure. + * Pure function — does not mutate the input stats. + */ +export function computeNextProfileUsageStats( + stats: ProfileUsageStats, + reason: AuthProfileFailureReason, + now?: number, +): ProfileUsageStats { + const ts = now ?? Date.now(); + const next = { ...stats }; + + // Reset counters if last failure is outside the failure window + if (next.lastFailureAt && ts - next.lastFailureAt > FAILURE_WINDOW_MS) { + next.errorCount = 0; + next.failureCounts = {}; + } + + // Increment counters + next.errorCount = (next.errorCount ?? 0) + 1; + next.lastFailureAt = ts; + + if (!next.failureCounts) next.failureCounts = {}; + next.failureCounts = { + ...next.failureCounts, + [reason]: (next.failureCounts[reason] ?? 0) + 1, + }; + + // Apply cooldown based on failure reason + if (reason === "billing") { + const billingCount = next.failureCounts.billing ?? 1; + const disableMs = calculateBillingDisableMs(billingCount); + next.disabledUntil = ts + disableMs; + next.disabledReason = "billing"; + } else { + const cooldownMs = calculateCooldownMs(next.errorCount); + next.cooldownUntil = ts + cooldownMs; + } + + return next; +} + +/** + * Mark a profile as having failed. Persists updated stats to disk. + */ +export function markAuthProfileFailure( + profileId: string, + reason: AuthProfileFailureReason, + now?: number, +): void { + updateAuthProfileStore((store) => { + const current = ensureUsageStats(store, profileId); + const next = computeNextProfileUsageStats(current, reason, now); + store.usageStats![profileId] = next; + }); +} + +/** + * Mark a profile as successfully used. Resets all cooldown/error state. + */ +export function markAuthProfileUsed(profileId: string, now?: number): void { + updateAuthProfileStore((store) => { + const stats = ensureUsageStats(store, profileId); + stats.lastUsed = now ?? Date.now(); + stats.errorCount = 0; + stats.cooldownUntil = undefined; + stats.disabledUntil = undefined; + stats.disabledReason = undefined; + stats.failureCounts = undefined; + }); +} + +/** + * Mark a profile as the last known good for a provider. + */ +export function markAuthProfileGood(provider: string, profileId: string): void { + updateAuthProfileStore((store) => { + if (!store.lastGood) store.lastGood = {}; + store.lastGood[provider] = profileId; + }); +} + +/** + * Clear cooldown for a specific profile. + */ +export function clearAuthProfileCooldown(profileId: string): void { + updateAuthProfileStore((store) => { + const stats = ensureUsageStats(store, profileId); + stats.errorCount = 0; + stats.cooldownUntil = undefined; + }); +}