From 715728c989a997dd101a15031e2d93a8018b0ab2 Mon Sep 17 00:00:00 2001 From: Bruno Guidolim Date: Wed, 28 Jan 2026 18:01:21 +0100 Subject: [PATCH 1/8] feat(auth): track cooldown per (auth profile + model) Allow different models from the same provider to have independent cooldowns. When a model hits rate limit, only that specific model is blocked, not all models using the same auth profile. - Add cooldownKey() helper for composite key generation - Update isProfileInCooldown to check both per-model and profile-level - Update markAuthProfileFailure/Cooldown with optional model param - Pass model to cooldown checks in model-fallback and agent runner - Add comprehensive tests for per-model cooldown behavior Ref: #3417 --- ...th-profiles.auth-profile-cooldowns.test.ts | 97 ++++++++++++++++++- src/agents/auth-profiles.ts | 1 + src/agents/auth-profiles/usage.ts | 93 ++++++++++++++---- src/agents/model-fallback.test.ts | 60 ++++++++++++ src/agents/model-fallback.ts | 8 +- src/agents/pi-embedded-runner/run.ts | 2 + 6 files changed, 240 insertions(+), 21 deletions(-) diff --git a/src/agents/auth-profiles.auth-profile-cooldowns.test.ts b/src/agents/auth-profiles.auth-profile-cooldowns.test.ts index e5fe3900a..8eab8d3d7 100644 --- a/src/agents/auth-profiles.auth-profile-cooldowns.test.ts +++ b/src/agents/auth-profiles.auth-profile-cooldowns.test.ts @@ -1,5 +1,11 @@ import { describe, expect, it } from "vitest"; -import { calculateAuthProfileCooldownMs } from "./auth-profiles.js"; +import { + calculateAuthProfileCooldownMs, + cooldownKey, + isProfileInCooldown, +} from "./auth-profiles.js"; +import type { AuthProfileStore } from "./auth-profiles.js"; +import { AUTH_STORE_VERSION } from "./auth-profiles/constants.js"; describe("auth profile cooldowns", () => { it("applies exponential backoff with a 1h cap", () => { @@ -10,3 +16,92 @@ describe("auth profile cooldowns", () => { expect(calculateAuthProfileCooldownMs(5)).toBe(60 * 60_000); }); }); + +describe("cooldownKey", () => { + it("returns profileId when model is not provided", () => { + expect(cooldownKey("openai:default")).toBe("openai:default"); + expect(cooldownKey("openai:default", undefined)).toBe("openai:default"); + }); + + it("returns composite key when model is provided", () => { + expect(cooldownKey("openai:default", "gpt-4")).toBe("openai:default:gpt-4"); + expect(cooldownKey("github-copilot:default", "gpt-5.2")).toBe("github-copilot:default:gpt-5.2"); + }); +}); + +describe("isProfileInCooldown with per-model support", () => { + it("returns false when no cooldown exists", () => { + const store: AuthProfileStore = { + version: AUTH_STORE_VERSION, + profiles: { + "openai:default": { type: "api_key", provider: "openai", key: "test" }, + }, + }; + expect(isProfileInCooldown(store, "openai:default")).toBe(false); + expect(isProfileInCooldown(store, "openai:default", "gpt-4")).toBe(false); + }); + + it("checks profile-level cooldown when model not provided", () => { + const store: AuthProfileStore = { + version: AUTH_STORE_VERSION, + profiles: { + "openai:default": { type: "api_key", provider: "openai", key: "test" }, + }, + usageStats: { + "openai:default": { cooldownUntil: Date.now() + 60_000 }, + }, + }; + expect(isProfileInCooldown(store, "openai:default")).toBe(true); + }); + + it("checks per-model cooldown when model is provided", () => { + const store: AuthProfileStore = { + version: AUTH_STORE_VERSION, + profiles: { + "openai:default": { type: "api_key", provider: "openai", key: "test" }, + }, + usageStats: { + "openai:default:gpt-4": { cooldownUntil: Date.now() + 60_000 }, + }, + }; + // model-specific cooldown exists + expect(isProfileInCooldown(store, "openai:default", "gpt-4")).toBe(true); + // different model is not in cooldown + expect(isProfileInCooldown(store, "openai:default", "gpt-3.5")).toBe(false); + // profile-level is not in cooldown + expect(isProfileInCooldown(store, "openai:default")).toBe(false); + }); + + it("allows independent cooldowns per model", () => { + const store: AuthProfileStore = { + version: AUTH_STORE_VERSION, + profiles: { + "github-copilot:default": { + type: "api_key", + provider: "github-copilot", + key: "test", + }, + }, + usageStats: { + // gpt-5.2 is in cooldown (rate limited) + "github-copilot:default:gpt-5.2": { cooldownUntil: Date.now() + 60_000 }, + // gpt-5-mini has no cooldown (unlimited quota) + }, + }; + expect(isProfileInCooldown(store, "github-copilot:default", "gpt-5.2")).toBe(true); + expect(isProfileInCooldown(store, "github-copilot:default", "gpt-5-mini")).toBe(false); + }); + + it("returns false when cooldown has expired", () => { + const store: AuthProfileStore = { + version: AUTH_STORE_VERSION, + profiles: { + "openai:default": { type: "api_key", provider: "openai", key: "test" }, + }, + usageStats: { + "openai:default:gpt-4": { cooldownUntil: Date.now() - 1000 }, // expired + }, + }; + expect(isProfileInCooldown(store, "openai:default", "gpt-4")).toBe(false); + }); +}); diff --git a/src/agents/auth-profiles.ts b/src/agents/auth-profiles.ts index 9a6c75b10..e68a27f1c 100644 --- a/src/agents/auth-profiles.ts +++ b/src/agents/auth-profiles.ts @@ -32,6 +32,7 @@ export type { export { calculateAuthProfileCooldownMs, clearAuthProfileCooldown, + cooldownKey, isProfileInCooldown, markAuthProfileCooldown, markAuthProfileFailure, diff --git a/src/agents/auth-profiles/usage.ts b/src/agents/auth-profiles/usage.ts index 2bf79b299..43bb660b6 100644 --- a/src/agents/auth-profiles/usage.ts +++ b/src/agents/auth-profiles/usage.ts @@ -3,6 +3,15 @@ import { normalizeProviderId } from "../model-selection.js"; import { saveAuthProfileStore, updateAuthProfileStoreWithLock } from "./store.js"; import type { AuthProfileFailureReason, AuthProfileStore, ProfileUsageStats } from "./types.js"; +/** + * Generate a cooldown key that optionally includes the model. + * When model is provided, cooldowns are tracked per (profile + model) combination. + * This allows different models from the same provider to have independent cooldowns. + */ +export function cooldownKey(profileId: string, model?: string): string { + return model ? `${profileId}:${model}` : profileId; +} + function resolveProfileUnusableUntil(stats: ProfileUsageStats): number | null { const values = [stats.cooldownUntil, stats.disabledUntil] .filter((value): value is number => typeof value === "number") @@ -14,11 +23,39 @@ function resolveProfileUnusableUntil(stats: ProfileUsageStats): number | null { /** * Check if a profile is currently in cooldown (due to rate limiting or errors). */ -export function isProfileInCooldown(store: AuthProfileStore, profileId: string): boolean { - const stats = store.usageStats?.[profileId]; - if (!stats) return false; - const unusableUntil = resolveProfileUnusableUntil(stats); - return unusableUntil ? Date.now() < unusableUntil : false; +/** + * Check if a profile is currently in cooldown (due to rate limiting or errors). + * When model is provided, checks the per-model cooldown key. + */ +/** + * Check if a profile is currently in cooldown (due to rate limiting or errors). + * When model is provided, checks both the per-model cooldown key and the profile-level key. + * A profile-level cooldown applies to all models (backward compatibility). + */ +export function isProfileInCooldown( + store: AuthProfileStore, + profileId: string, + model?: string, +): boolean { + const now = Date.now(); + + // Check per-model cooldown first (if model provided) + if (model) { + const modelKey = cooldownKey(profileId, model); + const modelStats = store.usageStats?.[modelKey]; + if (modelStats) { + const modelUnusableUntil = resolveProfileUnusableUntil(modelStats); + if (modelUnusableUntil && now < modelUnusableUntil) { + return true; + } + } + } + + // Also check profile-level cooldown (applies to all models) + const profileStats = store.usageStats?.[profileId]; + if (!profileStats) return false; + const profileUnusableUntil = resolveProfileUnusableUntil(profileStats); + return profileUnusableUntil ? now < profileUnusableUntil : false; } /** @@ -188,21 +225,28 @@ function computeNextProfileUsageStats(params: { * Mark a profile as failed for a specific reason. Billing failures are treated * as "disabled" (longer backoff) vs the regular cooldown window. */ +/** + * Mark a profile as failed for a specific reason. Billing failures are treated + * as "disabled" (longer backoff) vs the regular cooldown window. + * When model is provided, cooldown is tracked per (profile + model) combination. + */ export async function markAuthProfileFailure(params: { store: AuthProfileStore; profileId: string; + model?: string; reason: AuthProfileFailureReason; cfg?: MoltbotConfig; agentDir?: string; }): Promise { - const { store, profileId, reason, agentDir, cfg } = params; + const { store, profileId, model, reason, agentDir, cfg } = params; + const key = cooldownKey(profileId, model); const updated = await updateAuthProfileStoreWithLock({ agentDir, updater: (freshStore) => { const profile = freshStore.profiles[profileId]; if (!profile) return false; freshStore.usageStats = freshStore.usageStats ?? {}; - const existing = freshStore.usageStats[profileId] ?? {}; + const existing = freshStore.usageStats[key] ?? {}; const now = Date.now(); const providerKey = normalizeProviderId(profile.provider); @@ -211,7 +255,7 @@ export async function markAuthProfileFailure(params: { providerId: providerKey, }); - freshStore.usageStats[profileId] = computeNextProfileUsageStats({ + freshStore.usageStats[key] = computeNextProfileUsageStats({ existing, now, reason, @@ -227,7 +271,7 @@ export async function markAuthProfileFailure(params: { if (!store.profiles[profileId]) return; store.usageStats = store.usageStats ?? {}; - const existing = store.usageStats[profileId] ?? {}; + const existing = store.usageStats[key] ?? {}; const now = Date.now(); const providerKey = normalizeProviderId(store.profiles[profileId]?.provider ?? ""); const cfgResolved = resolveAuthCooldownConfig({ @@ -235,7 +279,7 @@ export async function markAuthProfileFailure(params: { providerId: providerKey, }); - store.usageStats[profileId] = computeNextProfileUsageStats({ + store.usageStats[key] = computeNextProfileUsageStats({ existing, now, reason, @@ -249,14 +293,22 @@ export async function markAuthProfileFailure(params: { * Cooldown times: 1min, 5min, 25min, max 1 hour. * Uses store lock to avoid overwriting concurrent usage updates. */ +/** + * Mark a profile as failed/rate-limited. Applies exponential backoff cooldown. + * Cooldown times: 1min, 5min, 25min, max 1 hour. + * Uses store lock to avoid overwriting concurrent usage updates. + * When model is provided, cooldown is tracked per (profile + model) combination. + */ export async function markAuthProfileCooldown(params: { store: AuthProfileStore; profileId: string; + model?: string; agentDir?: string; }): Promise { await markAuthProfileFailure({ store: params.store, profileId: params.profileId, + model: params.model, reason: "unknown", agentDir: params.agentDir, }); @@ -266,19 +318,26 @@ export async function markAuthProfileCooldown(params: { * Clear cooldown for a profile (e.g., manual reset). * Uses store lock to avoid overwriting concurrent usage updates. */ +/** + * Clear cooldown for a profile (e.g., manual reset). + * Uses store lock to avoid overwriting concurrent usage updates. + * When model is provided, clears the per-model cooldown key. + */ export async function clearAuthProfileCooldown(params: { store: AuthProfileStore; profileId: string; + model?: string; agentDir?: string; }): Promise { - const { store, profileId, agentDir } = params; + const { store, profileId, model, agentDir } = params; + const key = cooldownKey(profileId, model); const updated = await updateAuthProfileStoreWithLock({ agentDir, updater: (freshStore) => { - if (!freshStore.usageStats?.[profileId]) return false; + if (!freshStore.usageStats?.[key]) return false; - freshStore.usageStats[profileId] = { - ...freshStore.usageStats[profileId], + freshStore.usageStats[key] = { + ...freshStore.usageStats[key], errorCount: 0, cooldownUntil: undefined, }; @@ -289,10 +348,10 @@ export async function clearAuthProfileCooldown(params: { store.usageStats = updated.usageStats; return; } - if (!store.usageStats?.[profileId]) return; + if (!store.usageStats?.[key]) return; - store.usageStats[profileId] = { - ...store.usageStats[profileId], + store.usageStats[key] = { + ...store.usageStats[key], errorCount: 0, cooldownUntil: undefined, }; diff --git a/src/agents/model-fallback.test.ts b/src/agents/model-fallback.test.ts index f8a94560c..25c311ae7 100644 --- a/src/agents/model-fallback.test.ts +++ b/src/agents/model-fallback.test.ts @@ -240,6 +240,66 @@ describe("runWithModelFallback", () => { } }); + it("allows different models from same provider when only one model is in cooldown", async () => { + const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "moltbot-auth-")); + const provider = `per-model-cooldown-${crypto.randomUUID()}`; + const profileId = `${provider}:default`; + + const store: AuthProfileStore = { + version: AUTH_STORE_VERSION, + profiles: { + [profileId]: { + type: "api_key", + provider, + key: "test-key", + }, + }, + usageStats: { + // Only model-a is in cooldown (per-model key) + [`${profileId}:model-a`]: { + cooldownUntil: Date.now() + 60_000, + }, + // model-b has no cooldown + }, + }; + + saveAuthProfileStore(store, tempDir); + + const cfg = makeCfg({ + agents: { + defaults: { + model: { + primary: `${provider}/model-a`, + fallbacks: [`${provider}/model-b`], + }, + }, + }, + }); + + const run = vi.fn().mockImplementation(async (providerId, modelId) => { + if (modelId === "model-b") return "ok"; + throw new Error(`unexpected model: ${providerId}/${modelId}`); + }); + + try { + const result = await runWithModelFallback({ + cfg, + provider, + model: "model-a", + agentDir: tempDir, + run, + }); + + expect(result.result).toBe("ok"); + // model-a should be skipped (in cooldown), model-b should be tried + expect(run.mock.calls).toEqual([[provider, "model-b"]]); + expect(result.attempts[0]?.reason).toBe("rate_limit"); + expect(result.attempts[0]?.model).toBe("model-a"); + } finally { + await fs.rm(tempDir, { recursive: true, force: true }); + } + }); + it("does not append configured primary when fallbacksOverride is set", async () => { const cfg = makeCfg({ agents: { diff --git a/src/agents/model-fallback.ts b/src/agents/model-fallback.ts index b87135f9f..aa27a5776 100644 --- a/src/agents/model-fallback.ts +++ b/src/agents/model-fallback.ts @@ -231,14 +231,16 @@ export async function runWithModelFallback(params: { store: authStore, provider: candidate.provider, }); - const isAnyProfileAvailable = profileIds.some((id) => !isProfileInCooldown(authStore, id)); + const isAnyProfileAvailable = profileIds.some( + (id) => !isProfileInCooldown(authStore, id, candidate.model), + ); if (profileIds.length > 0 && !isAnyProfileAvailable) { - // All profiles for this provider are in cooldown; skip without attempting + // All profiles for this provider+model are in cooldown; skip without attempting attempts.push({ provider: candidate.provider, model: candidate.model, - error: `Provider ${candidate.provider} is in cooldown (all profiles unavailable)`, + error: `${candidate.provider}/${candidate.model} is in cooldown (all profiles unavailable)`, reason: "rate_limit", }); continue; diff --git a/src/agents/pi-embedded-runner/run.ts b/src/agents/pi-embedded-runner/run.ts index 870453f38..746398217 100644 --- a/src/agents/pi-embedded-runner/run.ts +++ b/src/agents/pi-embedded-runner/run.ts @@ -474,6 +474,7 @@ export async function runEmbeddedPiAgent( await markAuthProfileFailure({ store: authStore, profileId: lastProfileId, + model: modelId, reason: promptFailoverReason, cfg: params.config, agentDir: params.agentDir, @@ -561,6 +562,7 @@ export async function runEmbeddedPiAgent( await markAuthProfileFailure({ store: authStore, profileId: lastProfileId, + model: modelId, reason, cfg: params.config, agentDir: params.agentDir, From 33f9bcc3ce7f4b080843b2d1875c528cfd770a4d Mon Sep 17 00:00:00 2001 From: Bruno Guidolim Date: Wed, 28 Jan 2026 18:45:26 +0100 Subject: [PATCH 2/8] fix(auth): clear per-model cooldown on success When a model succeeds, also clear its per-model cooldown key so the system doesn't think it's still rate-limited. - Add optional `model` param to markAuthProfileUsed - Pass modelId when marking profile used in agent runner - Add tests for per-model cooldown clearing behavior --- ...th-profiles.auth-profile-cooldowns.test.ts | 72 +++++++++++++++++++ src/agents/auth-profiles/usage.ts | 38 +++++++++- src/agents/pi-embedded-runner/run.ts | 1 + 3 files changed, 110 insertions(+), 1 deletion(-) diff --git a/src/agents/auth-profiles.auth-profile-cooldowns.test.ts b/src/agents/auth-profiles.auth-profile-cooldowns.test.ts index 8eab8d3d7..7ed3df38b 100644 --- a/src/agents/auth-profiles.auth-profile-cooldowns.test.ts +++ b/src/agents/auth-profiles.auth-profile-cooldowns.test.ts @@ -1,8 +1,13 @@ +import fs from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; import { describe, expect, it } from "vitest"; import { calculateAuthProfileCooldownMs, cooldownKey, isProfileInCooldown, + markAuthProfileUsed, + saveAuthProfileStore, } from "./auth-profiles.js"; import type { AuthProfileStore } from "./auth-profiles.js"; import { AUTH_STORE_VERSION } from "./auth-profiles/constants.js"; @@ -105,3 +110,70 @@ describe("isProfileInCooldown with per-model support", () => { expect(isProfileInCooldown(store, "openai:default", "gpt-4")).toBe(false); }); }); + +describe("markAuthProfileUsed with per-model support", () => { + it("clears per-model cooldown when model is provided", async () => { + const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "moltbot-auth-")); + const cooldownTime = Date.now() + 60_000; + const store: AuthProfileStore = { + version: AUTH_STORE_VERSION, + profiles: { + "openai:default": { type: "api_key", provider: "openai", key: "test" }, + }, + usageStats: { + "openai:default": { cooldownUntil: cooldownTime }, + "openai:default:gpt-4": { cooldownUntil: cooldownTime, errorCount: 3 }, + "openai:default:gpt-3.5": { cooldownUntil: cooldownTime }, + }, + }; + saveAuthProfileStore(store, tempDir); + + try { + // Mark gpt-4 as used (successful) + await markAuthProfileUsed({ + store, + profileId: "openai:default", + model: "gpt-4", + agentDir: tempDir, + }); + + // Profile-level cooldown should be cleared + expect(store.usageStats?.["openai:default"]?.cooldownUntil).toBeUndefined(); + // Per-model cooldown for gpt-4 should be cleared + expect(store.usageStats?.["openai:default:gpt-4"]?.cooldownUntil).toBeUndefined(); + expect(store.usageStats?.["openai:default:gpt-4"]?.errorCount).toBe(0); + // Per-model cooldown for gpt-3.5 should remain (different model) + expect(store.usageStats?.["openai:default:gpt-3.5"]?.cooldownUntil).toBe(cooldownTime); + } finally { + await fs.rm(tempDir, { recursive: true, force: true }); + } + }); + + it("only clears profile-level cooldown when model is not provided", async () => { + const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "moltbot-auth-")); + const cooldownTime = Date.now() + 60_000; + const store: AuthProfileStore = { + version: AUTH_STORE_VERSION, + profiles: { + "openai:default": { type: "api_key", provider: "openai", key: "test" }, + }, + usageStats: { + "openai:default": { cooldownUntil: cooldownTime }, + "openai:default:gpt-4": { cooldownUntil: cooldownTime }, + }, + }; + saveAuthProfileStore(store, tempDir); + + try { + // Mark profile as used without specifying model + await markAuthProfileUsed({ store, profileId: "openai:default", agentDir: tempDir }); + + // Profile-level cooldown should be cleared + expect(store.usageStats?.["openai:default"]?.cooldownUntil).toBeUndefined(); + // Per-model cooldown should remain (no model specified) + expect(store.usageStats?.["openai:default:gpt-4"]?.cooldownUntil).toBe(cooldownTime); + } finally { + await fs.rm(tempDir, { recursive: true, force: true }); + } + }); +}); diff --git a/src/agents/auth-profiles/usage.ts b/src/agents/auth-profiles/usage.ts index 43bb660b6..1681cdad2 100644 --- a/src/agents/auth-profiles/usage.ts +++ b/src/agents/auth-profiles/usage.ts @@ -62,17 +62,24 @@ export function isProfileInCooldown( * Mark a profile as successfully used. Resets error count and updates lastUsed. * Uses store lock to avoid overwriting concurrent usage updates. */ +/** + * Mark a profile as successfully used. Resets error count and updates lastUsed. + * Uses store lock to avoid overwriting concurrent usage updates. + * When model is provided, also clears the per-model cooldown. + */ export async function markAuthProfileUsed(params: { store: AuthProfileStore; profileId: string; + model?: string; agentDir?: string; }): Promise { - const { store, profileId, agentDir } = params; + const { store, profileId, model, agentDir } = params; const updated = await updateAuthProfileStoreWithLock({ agentDir, updater: (freshStore) => { if (!freshStore.profiles[profileId]) return false; freshStore.usageStats = freshStore.usageStats ?? {}; + // Clear profile-level cooldown freshStore.usageStats[profileId] = { ...freshStore.usageStats[profileId], lastUsed: Date.now(), @@ -82,6 +89,20 @@ export async function markAuthProfileUsed(params: { disabledReason: undefined, failureCounts: undefined, }; + // Also clear per-model cooldown if model provided + if (model) { + const modelKey = cooldownKey(profileId, model); + if (freshStore.usageStats[modelKey]) { + freshStore.usageStats[modelKey] = { + ...freshStore.usageStats[modelKey], + errorCount: 0, + cooldownUntil: undefined, + disabledUntil: undefined, + disabledReason: undefined, + failureCounts: undefined, + }; + } + } return true; }, }); @@ -92,6 +113,7 @@ export async function markAuthProfileUsed(params: { if (!store.profiles[profileId]) return; store.usageStats = store.usageStats ?? {}; + // Clear profile-level cooldown store.usageStats[profileId] = { ...store.usageStats[profileId], lastUsed: Date.now(), @@ -101,6 +123,20 @@ export async function markAuthProfileUsed(params: { disabledReason: undefined, failureCounts: undefined, }; + // Also clear per-model cooldown if model provided + if (model) { + const modelKey = cooldownKey(profileId, model); + if (store.usageStats[modelKey]) { + store.usageStats[modelKey] = { + ...store.usageStats[modelKey], + errorCount: 0, + cooldownUntil: undefined, + disabledUntil: undefined, + disabledReason: undefined, + failureCounts: undefined, + }; + } + } saveAuthProfileStore(store, agentDir); } diff --git a/src/agents/pi-embedded-runner/run.ts b/src/agents/pi-embedded-runner/run.ts index 746398217..b6b002319 100644 --- a/src/agents/pi-embedded-runner/run.ts +++ b/src/agents/pi-embedded-runner/run.ts @@ -646,6 +646,7 @@ export async function runEmbeddedPiAgent( await markAuthProfileUsed({ store: authStore, profileId: lastProfileId, + model: modelId, agentDir: params.agentDir, }); } From 7f0c09866527c0508d5b10738aa42da92ba48312 Mon Sep 17 00:00:00 2001 From: Bruno Guidolim Date: Wed, 28 Jan 2026 23:43:56 +0100 Subject: [PATCH 3/8] chore(auth): clean up duplicate JSDoc comments and improve cooldownKey - Remove duplicate JSDoc blocks left from iterative editing - Add @example annotations to cooldownKey() for clarity - Handle empty/whitespace model strings in cooldownKey() - Improve isProfileInCooldown() documentation to explain dual-check behavior - Clarify markAuthProfileCooldown() is a convenience wrapper --- src/agents/auth-profiles/usage.ts | 46 +++++++++++-------------------- 1 file changed, 16 insertions(+), 30 deletions(-) diff --git a/src/agents/auth-profiles/usage.ts b/src/agents/auth-profiles/usage.ts index 1681cdad2..a425bd78b 100644 --- a/src/agents/auth-profiles/usage.ts +++ b/src/agents/auth-profiles/usage.ts @@ -7,9 +7,14 @@ import type { AuthProfileFailureReason, AuthProfileStore, ProfileUsageStats } fr * Generate a cooldown key that optionally includes the model. * When model is provided, cooldowns are tracked per (profile + model) combination. * This allows different models from the same provider to have independent cooldowns. + * + * @example cooldownKey("openai:default", "gpt-4") => "openai:default:gpt-4" + * @example cooldownKey("openai:default") => "openai:default" */ export function cooldownKey(profileId: string, model?: string): string { - return model ? `${profileId}:${model}` : profileId; + // Treat empty/whitespace-only string as "no model" to avoid trailing colon in key + const normalizedModel = model?.trim() || undefined; + return normalizedModel ? `${profileId}:${normalizedModel}` : profileId; } function resolveProfileUnusableUntil(stats: ProfileUsageStats): number | null { @@ -22,15 +27,13 @@ function resolveProfileUnusableUntil(stats: ProfileUsageStats): number | null { /** * Check if a profile is currently in cooldown (due to rate limiting or errors). - */ -/** - * Check if a profile is currently in cooldown (due to rate limiting or errors). - * When model is provided, checks the per-model cooldown key. - */ -/** - * Check if a profile is currently in cooldown (due to rate limiting or errors). - * When model is provided, checks both the per-model cooldown key and the profile-level key. - * A profile-level cooldown applies to all models (backward compatibility). + * + * When model is provided, checks both: + * 1. The per-model cooldown key (e.g., "openai:default:gpt-4") + * 2. The profile-level cooldown key (e.g., "openai:default") + * + * Profile-level cooldowns apply to all models under that profile, supporting + * legacy entries and scenarios where failures affect all models (e.g., auth errors). */ export function isProfileInCooldown( store: AuthProfileStore, @@ -58,10 +61,6 @@ export function isProfileInCooldown( return profileUnusableUntil ? now < profileUnusableUntil : false; } -/** - * Mark a profile as successfully used. Resets error count and updates lastUsed. - * Uses store lock to avoid overwriting concurrent usage updates. - */ /** * Mark a profile as successfully used. Resets error count and updates lastUsed. * Uses store lock to avoid overwriting concurrent usage updates. @@ -257,10 +256,6 @@ function computeNextProfileUsageStats(params: { return updatedStats; } -/** - * Mark a profile as failed for a specific reason. Billing failures are treated - * as "disabled" (longer backoff) vs the regular cooldown window. - */ /** * Mark a profile as failed for a specific reason. Billing failures are treated * as "disabled" (longer backoff) vs the regular cooldown window. @@ -325,14 +320,9 @@ export async function markAuthProfileFailure(params: { } /** - * Mark a profile as failed/rate-limited. Applies exponential backoff cooldown. - * Cooldown times: 1min, 5min, 25min, max 1 hour. - * Uses store lock to avoid overwriting concurrent usage updates. - */ -/** - * Mark a profile as failed/rate-limited. Applies exponential backoff cooldown. - * Cooldown times: 1min, 5min, 25min, max 1 hour. - * Uses store lock to avoid overwriting concurrent usage updates. + * Mark a profile as failed/rate-limited with "unknown" reason. + * Convenience wrapper around markAuthProfileFailure() for generic failures. + * Applies exponential backoff cooldown: 1min, 5min, 25min, max 1 hour. * When model is provided, cooldown is tracked per (profile + model) combination. */ export async function markAuthProfileCooldown(params: { @@ -350,10 +340,6 @@ export async function markAuthProfileCooldown(params: { }); } -/** - * Clear cooldown for a profile (e.g., manual reset). - * Uses store lock to avoid overwriting concurrent usage updates. - */ /** * Clear cooldown for a profile (e.g., manual reset). * Uses store lock to avoid overwriting concurrent usage updates. From e9e6689c9c024d04917131cb10a0083874bf7daf Mon Sep 17 00:00:00 2001 From: Bruno Guidolim Date: Wed, 28 Jan 2026 23:44:06 +0100 Subject: [PATCH 4/8] test(auth): add tests for per-model cooldown functions - Add tests for markAuthProfileFailure with model param - Add tests for markAuthProfileCooldown with model param - Add tests for clearAuthProfileCooldown with model param - Add backward compatibility test (profile-level blocks all models) - Add test for disabledUntil handling in per-model cooldowns - Add edge case test for empty string model parameter --- ...th-profiles.auth-profile-cooldowns.test.ts | 230 ++++++++++++++++++ 1 file changed, 230 insertions(+) diff --git a/src/agents/auth-profiles.auth-profile-cooldowns.test.ts b/src/agents/auth-profiles.auth-profile-cooldowns.test.ts index 7ed3df38b..dcb68315f 100644 --- a/src/agents/auth-profiles.auth-profile-cooldowns.test.ts +++ b/src/agents/auth-profiles.auth-profile-cooldowns.test.ts @@ -4,8 +4,11 @@ import path from "node:path"; import { describe, expect, it } from "vitest"; import { calculateAuthProfileCooldownMs, + clearAuthProfileCooldown, cooldownKey, isProfileInCooldown, + markAuthProfileCooldown, + markAuthProfileFailure, markAuthProfileUsed, saveAuthProfileStore, } from "./auth-profiles.js"; @@ -177,3 +180,230 @@ describe("markAuthProfileUsed with per-model support", () => { } }); }); + +describe("cooldownKey edge cases", () => { + it("treats empty string model the same as undefined", () => { + // Empty string should be treated as "no model" to avoid trailing colon + expect(cooldownKey("openai:default", "")).toBe("openai:default"); + expect(cooldownKey("openai:default", " ")).toBe("openai:default"); + }); +}); + +describe("isProfileInCooldown backward compatibility", () => { + it("returns true for any model when profile-level cooldown exists", () => { + const store: AuthProfileStore = { + version: AUTH_STORE_VERSION, + profiles: { + "openai:default": { type: "api_key", provider: "openai", key: "test" }, + }, + usageStats: { + "openai:default": { cooldownUntil: Date.now() + 60_000 }, // profile-level only + }, + }; + // Any model should be blocked when profile-level cooldown exists + expect(isProfileInCooldown(store, "openai:default", "gpt-4")).toBe(true); + expect(isProfileInCooldown(store, "openai:default", "gpt-3.5")).toBe(true); + expect(isProfileInCooldown(store, "openai:default", "o1-preview")).toBe(true); + // Profile-level check also works + expect(isProfileInCooldown(store, "openai:default")).toBe(true); + }); + + it("checks disabledUntil for per-model cooldowns (billing failures)", () => { + const store: AuthProfileStore = { + version: AUTH_STORE_VERSION, + profiles: { + "openai:default": { type: "api_key", provider: "openai", key: "test" }, + }, + usageStats: { + "openai:default:gpt-4": { disabledUntil: Date.now() + 60_000 }, // billing failure + }, + }; + expect(isProfileInCooldown(store, "openai:default", "gpt-4")).toBe(true); + expect(isProfileInCooldown(store, "openai:default", "gpt-3.5")).toBe(false); + }); +}); + +describe("markAuthProfileFailure with per-model support", () => { + it("tracks failure per model when model is provided", async () => { + const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "moltbot-auth-")); + const store: AuthProfileStore = { + version: AUTH_STORE_VERSION, + profiles: { + "openai:default": { type: "api_key", provider: "openai", key: "test" }, + }, + }; + saveAuthProfileStore(store, tempDir); + + try { + await markAuthProfileFailure({ + store, + profileId: "openai:default", + model: "gpt-4", + reason: "rate_limit", + agentDir: tempDir, + }); + + // Per-model key should have cooldown + expect(store.usageStats?.["openai:default:gpt-4"]?.cooldownUntil).toBeGreaterThan(Date.now()); + expect(store.usageStats?.["openai:default:gpt-4"]?.errorCount).toBe(1); + // Profile-level should NOT have cooldown (only model-specific) + expect(store.usageStats?.["openai:default"]?.cooldownUntil).toBeUndefined(); + // Other models should not be affected + expect(store.usageStats?.["openai:default:gpt-3.5"]).toBeUndefined(); + } finally { + await fs.rm(tempDir, { recursive: true, force: true }); + } + }); + + it("tracks failure at profile level when model is not provided", async () => { + const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "moltbot-auth-")); + const store: AuthProfileStore = { + version: AUTH_STORE_VERSION, + profiles: { + "openai:default": { type: "api_key", provider: "openai", key: "test" }, + }, + }; + saveAuthProfileStore(store, tempDir); + + try { + await markAuthProfileFailure({ + store, + profileId: "openai:default", + reason: "auth", + agentDir: tempDir, + }); + + // Profile-level key should have cooldown + expect(store.usageStats?.["openai:default"]?.cooldownUntil).toBeGreaterThan(Date.now()); + expect(store.usageStats?.["openai:default"]?.errorCount).toBe(1); + } finally { + await fs.rm(tempDir, { recursive: true, force: true }); + } + }); + + it("tracks billing failures with disabledUntil per model", async () => { + const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "moltbot-auth-")); + const store: AuthProfileStore = { + version: AUTH_STORE_VERSION, + profiles: { + "openai:default": { type: "api_key", provider: "openai", key: "test" }, + }, + }; + saveAuthProfileStore(store, tempDir); + + try { + await markAuthProfileFailure({ + store, + profileId: "openai:default", + model: "gpt-4", + reason: "billing", + agentDir: tempDir, + }); + + // Billing failures use disabledUntil instead of cooldownUntil + expect(store.usageStats?.["openai:default:gpt-4"]?.disabledUntil).toBeGreaterThan(Date.now()); + expect(store.usageStats?.["openai:default:gpt-4"]?.disabledReason).toBe("billing"); + } finally { + await fs.rm(tempDir, { recursive: true, force: true }); + } + }); +}); + +describe("markAuthProfileCooldown with per-model support", () => { + it("marks cooldown per model when model is provided", async () => { + const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "moltbot-auth-")); + const store: AuthProfileStore = { + version: AUTH_STORE_VERSION, + profiles: { + "openai:default": { type: "api_key", provider: "openai", key: "test" }, + }, + }; + saveAuthProfileStore(store, tempDir); + + try { + await markAuthProfileCooldown({ + store, + profileId: "openai:default", + model: "gpt-4", + agentDir: tempDir, + }); + + // Per-model key should have cooldown + expect(store.usageStats?.["openai:default:gpt-4"]?.cooldownUntil).toBeGreaterThan(Date.now()); + // Profile-level should NOT have cooldown + expect(store.usageStats?.["openai:default"]?.cooldownUntil).toBeUndefined(); + } finally { + await fs.rm(tempDir, { recursive: true, force: true }); + } + }); +}); + +describe("clearAuthProfileCooldown with per-model support", () => { + it("clears per-model cooldown when model is provided", async () => { + const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "moltbot-auth-")); + const cooldownTime = Date.now() + 60_000; + const store: AuthProfileStore = { + version: AUTH_STORE_VERSION, + profiles: { + "openai:default": { type: "api_key", provider: "openai", key: "test" }, + }, + usageStats: { + "openai:default": { cooldownUntil: cooldownTime }, + "openai:default:gpt-4": { cooldownUntil: cooldownTime, errorCount: 3 }, + "openai:default:gpt-3.5": { cooldownUntil: cooldownTime }, + }, + }; + saveAuthProfileStore(store, tempDir); + + try { + await clearAuthProfileCooldown({ + store, + profileId: "openai:default", + model: "gpt-4", + agentDir: tempDir, + }); + + // Per-model cooldown for gpt-4 should be cleared + expect(store.usageStats?.["openai:default:gpt-4"]?.cooldownUntil).toBeUndefined(); + expect(store.usageStats?.["openai:default:gpt-4"]?.errorCount).toBe(0); + // Profile-level cooldown should remain (different key) + expect(store.usageStats?.["openai:default"]?.cooldownUntil).toBe(cooldownTime); + // Other model cooldown should remain + expect(store.usageStats?.["openai:default:gpt-3.5"]?.cooldownUntil).toBe(cooldownTime); + } finally { + await fs.rm(tempDir, { recursive: true, force: true }); + } + }); + + it("clears profile-level cooldown when model is not provided", async () => { + const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "moltbot-auth-")); + const cooldownTime = Date.now() + 60_000; + const store: AuthProfileStore = { + version: AUTH_STORE_VERSION, + profiles: { + "openai:default": { type: "api_key", provider: "openai", key: "test" }, + }, + usageStats: { + "openai:default": { cooldownUntil: cooldownTime, errorCount: 2 }, + "openai:default:gpt-4": { cooldownUntil: cooldownTime }, + }, + }; + saveAuthProfileStore(store, tempDir); + + try { + await clearAuthProfileCooldown({ + store, + profileId: "openai:default", + agentDir: tempDir, + }); + + // Profile-level cooldown should be cleared + expect(store.usageStats?.["openai:default"]?.cooldownUntil).toBeUndefined(); + expect(store.usageStats?.["openai:default"]?.errorCount).toBe(0); + // Per-model cooldown should remain (different key) + expect(store.usageStats?.["openai:default:gpt-4"]?.cooldownUntil).toBe(cooldownTime); + } finally { + await fs.rm(tempDir, { recursive: true, force: true }); + } + }); +}); From 9edf0306af85978e260968be3df9800274cf408b Mon Sep 17 00:00:00 2001 From: Bruno Guidolim Date: Wed, 28 Jan 2026 23:51:40 +0100 Subject: [PATCH 5/8] docs(auth): document per-model cooldown design decision - Add module-level comment in usage.ts explaining the key asymmetry: failures create per-model keys, successes update profile-level keys - Add explanatory comment at top of cooldown test file - Create Serena memory (decision_auth_permodel_cooldown_design) for future reference This documents the design from discussion #3417 where per-model cooldowns allow independent rate limits while keeping the store clean. --- ...th-profiles.auth-profile-cooldowns.test.ts | 13 +++++++++++ src/agents/auth-profiles/usage.ts | 23 +++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/src/agents/auth-profiles.auth-profile-cooldowns.test.ts b/src/agents/auth-profiles.auth-profile-cooldowns.test.ts index dcb68315f..e63910250 100644 --- a/src/agents/auth-profiles.auth-profile-cooldowns.test.ts +++ b/src/agents/auth-profiles.auth-profile-cooldowns.test.ts @@ -1,3 +1,16 @@ +/* + * Per-Model Cooldown Tests + * ──────────────────────── + * These tests verify the per-model cooldown feature (discussion #3417). + * + * Key design asymmetry: + * - Failures CREATE per-model keys (e.g., "openai:default:gpt-4") + * - Successes UPDATE profile-level keys AND clear per-model keys (if they exist) + * - Per-model keys are ephemeral "penalty boxes" that only exist during cooldowns + * + * This allows independent rate limits per model while keeping the store clean. + * See: src/agents/auth-profiles/usage.ts for implementation details. + */ import fs from "node:fs/promises"; import os from "node:os"; import path from "node:path"; diff --git a/src/agents/auth-profiles/usage.ts b/src/agents/auth-profiles/usage.ts index a425bd78b..0514ccc8d 100644 --- a/src/agents/auth-profiles/usage.ts +++ b/src/agents/auth-profiles/usage.ts @@ -3,6 +3,29 @@ import { normalizeProviderId } from "../model-selection.js"; import { saveAuthProfileStore, updateAuthProfileStoreWithLock } from "./store.js"; import type { AuthProfileFailureReason, AuthProfileStore, ProfileUsageStats } from "./types.js"; +/* + * Per-Model Cooldown Design + * ───────────────────────── + * Cooldowns can be tracked at two granularities: + * + * 1. Profile-level keys (e.g., "github-copilot:github") + * - Track success metrics: lastUsed, lastGood, errorCount reset + * - Used by `lastGood` to remember which profile worked for a provider + * - Auth failures (wrong API key) should use profile-level cooldowns + * + * 2. Per-model keys (e.g., "github-copilot:github:gpt-5.2") + * - Created ONLY on failure to track rate limits + * - Act as ephemeral "penalty boxes" for specific models + * - Naturally disappear when cooldown expires or model recovers + * + * Key asymmetry: + * - Failures → create per-model key (if model provided) + * - Successes → update profile-level key + clear per-model key (if it exists) + * + * This keeps the store clean and allows independent rate limits per model + * while maintaining backward compatibility with profile-level cooldowns. + */ + /** * Generate a cooldown key that optionally includes the model. * When model is provided, cooldowns are tracked per (profile + model) combination. From 9d0841a86c571f725d4e929fd6b340b580dc3a57 Mon Sep 17 00:00:00 2001 From: Bruno Guidolim Date: Thu, 29 Jan 2026 00:04:37 +0100 Subject: [PATCH 6/8] docs(auth): document intentional model-agnostic design in order.ts and usage.ts - Add module comment in order.ts explaining profile ordering is intentionally model-agnostic (per-model filtering happens downstream) - Expand markAuthProfileUsed JSDoc explaining why success metrics are profile-level (if ANY model works, credentials are valid) --- src/agents/auth-profiles/order.ts | 20 ++++++++++++++++++++ src/agents/auth-profiles/usage.ts | 8 +++++++- 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/src/agents/auth-profiles/order.ts b/src/agents/auth-profiles/order.ts index 677c29069..d5c954147 100644 --- a/src/agents/auth-profiles/order.ts +++ b/src/agents/auth-profiles/order.ts @@ -4,6 +4,26 @@ import { listProfilesForProvider } from "./profiles.js"; import type { AuthProfileStore } from "./types.js"; import { isProfileInCooldown } from "./usage.js"; +/* + * Profile Ordering Design + * ─────────────────────── + * Profile ordering is intentionally MODEL-AGNOSTIC. This module answers: + * "Which auth profiles should we try for this provider?" + * + * Per-model cooldown filtering happens DOWNSTREAM in model-fallback.ts, + * which calls isProfileInCooldown(store, profileId, model) for each candidate. + * + * Why two layers? + * 1. Profile layer (here): Selects credentials based on provider, type preference, + * round-robin (lastUsed), and profile-level cooldowns (auth failures). + * 2. Model layer (model-fallback.ts): Filters by per-model rate limits before + * each API call attempt. + * + * This separation exists because: + * - Auth credentials (API keys/tokens) are profile-level + * - Rate limits are often model-level (e.g., gpt-5.2 has quota, gpt-5-mini unlimited) + */ + function resolveProfileUnusableUntil(stats: { cooldownUntil?: number; disabledUntil?: number; diff --git a/src/agents/auth-profiles/usage.ts b/src/agents/auth-profiles/usage.ts index 0514ccc8d..2cab9a729 100644 --- a/src/agents/auth-profiles/usage.ts +++ b/src/agents/auth-profiles/usage.ts @@ -87,7 +87,13 @@ export function isProfileInCooldown( /** * Mark a profile as successfully used. Resets error count and updates lastUsed. * Uses store lock to avoid overwriting concurrent usage updates. - * When model is provided, also clears the per-model cooldown. + * + * Success metrics (lastUsed, lastGood) are ALWAYS updated at the profile level, + * regardless of which model was used. This is intentional: if ANY model works, + * the credentials are valid and the profile should be remembered as "good". + * + * When model is provided, also clears the per-model cooldown (if one exists). + * This allows a recovered model to be used immediately without waiting for expiry. */ export async function markAuthProfileUsed(params: { store: AuthProfileStore; From e3caf006bbd09b9c454c70298670eaa33fb6aeb6 Mon Sep 17 00:00:00 2001 From: Bruno Guidolim Date: Thu, 29 Jan 2026 00:59:55 +0100 Subject: [PATCH 7/8] refactor(auth): extract helpers to simplify usage.ts cooldown logic - Add isKeyInCooldown() to reduce duplication in cooldown checking - Add clearCooldownFields() to centralize cooldown field clearing - Add applySuccessUpdates() to eliminate duplicated success update logic - Reduce markAuthProfileUsed from ~66 to ~26 lines --- src/agents/auth-profiles/usage.ts | 105 +++++++++++++----------------- 1 file changed, 44 insertions(+), 61 deletions(-) diff --git a/src/agents/auth-profiles/usage.ts b/src/agents/auth-profiles/usage.ts index 2cab9a729..ec867bb66 100644 --- a/src/agents/auth-profiles/usage.ts +++ b/src/agents/auth-profiles/usage.ts @@ -48,6 +48,14 @@ function resolveProfileUnusableUntil(stats: ProfileUsageStats): number | null { return Math.max(...values); } +/** Checks if a key is currently in cooldown. */ +function isKeyInCooldown(store: AuthProfileStore, key: string, now: number): boolean { + const stats = store.usageStats?.[key]; + if (!stats) return false; + const unusableUntil = resolveProfileUnusableUntil(stats); + return unusableUntil !== null && now < unusableUntil; +} + /** * Check if a profile is currently in cooldown (due to rate limiting or errors). * @@ -66,22 +74,43 @@ export function isProfileInCooldown( const now = Date.now(); // Check per-model cooldown first (if model provided) - if (model) { - const modelKey = cooldownKey(profileId, model); - const modelStats = store.usageStats?.[modelKey]; - if (modelStats) { - const modelUnusableUntil = resolveProfileUnusableUntil(modelStats); - if (modelUnusableUntil && now < modelUnusableUntil) { - return true; - } - } + if (model && isKeyInCooldown(store, cooldownKey(profileId, model), now)) { + return true; } // Also check profile-level cooldown (applies to all models) - const profileStats = store.usageStats?.[profileId]; - if (!profileStats) return false; - const profileUnusableUntil = resolveProfileUnusableUntil(profileStats); - return profileUnusableUntil ? now < profileUnusableUntil : false; + return isKeyInCooldown(store, profileId, now); +} + +/** Clears cooldown fields from usage stats, preserving other fields. */ +function clearCooldownFields( + stats: ProfileUsageStats | undefined, + options?: { setLastUsed?: boolean }, +): ProfileUsageStats { + return { + ...stats, + ...(options?.setLastUsed ? { lastUsed: Date.now() } : {}), + errorCount: 0, + cooldownUntil: undefined, + disabledUntil: undefined, + disabledReason: undefined, + failureCounts: undefined, + }; +} + +/** Applies success updates to usage stats in-place. */ +function applySuccessUpdates( + usageStats: Record, + profileId: string, + model?: string, +): void { + usageStats[profileId] = clearCooldownFields(usageStats[profileId], { setLastUsed: true }); + if (model) { + const modelKey = cooldownKey(profileId, model); + if (usageStats[modelKey]) { + usageStats[modelKey] = clearCooldownFields(usageStats[modelKey]); + } + } } /** @@ -107,30 +136,7 @@ export async function markAuthProfileUsed(params: { updater: (freshStore) => { if (!freshStore.profiles[profileId]) return false; freshStore.usageStats = freshStore.usageStats ?? {}; - // Clear profile-level cooldown - freshStore.usageStats[profileId] = { - ...freshStore.usageStats[profileId], - lastUsed: Date.now(), - errorCount: 0, - cooldownUntil: undefined, - disabledUntil: undefined, - disabledReason: undefined, - failureCounts: undefined, - }; - // Also clear per-model cooldown if model provided - if (model) { - const modelKey = cooldownKey(profileId, model); - if (freshStore.usageStats[modelKey]) { - freshStore.usageStats[modelKey] = { - ...freshStore.usageStats[modelKey], - errorCount: 0, - cooldownUntil: undefined, - disabledUntil: undefined, - disabledReason: undefined, - failureCounts: undefined, - }; - } - } + applySuccessUpdates(freshStore.usageStats, profileId, model); return true; }, }); @@ -141,30 +147,7 @@ export async function markAuthProfileUsed(params: { if (!store.profiles[profileId]) return; store.usageStats = store.usageStats ?? {}; - // Clear profile-level cooldown - store.usageStats[profileId] = { - ...store.usageStats[profileId], - lastUsed: Date.now(), - errorCount: 0, - cooldownUntil: undefined, - disabledUntil: undefined, - disabledReason: undefined, - failureCounts: undefined, - }; - // Also clear per-model cooldown if model provided - if (model) { - const modelKey = cooldownKey(profileId, model); - if (store.usageStats[modelKey]) { - store.usageStats[modelKey] = { - ...store.usageStats[modelKey], - errorCount: 0, - cooldownUntil: undefined, - disabledUntil: undefined, - disabledReason: undefined, - failureCounts: undefined, - }; - } - } + applySuccessUpdates(store.usageStats, profileId, model); saveAuthProfileStore(store, agentDir); } From e954d29291f49d21b0a3e994097d105b2fcf8a82 Mon Sep 17 00:00:00 2001 From: Bruno Guidolim Date: Thu, 29 Jan 2026 01:04:26 +0100 Subject: [PATCH 8/8] refactor(test): simplify auth-profile-cooldowns tests with helpers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add makeStore() helper to eliminate repeated store object literals - Add withTempDir() helper to eliminate try/finally boilerplate - Merge cooldownKey edge cases into main describe block - Remove redundant markAuthProfileCooldown test (covered by markAuthProfileFailure) - Net reduction: 245 lines removed (57% smaller), 19→17 tests --- ...th-profiles.auth-profile-cooldowns.test.ts | 289 ++++-------------- 1 file changed, 66 insertions(+), 223 deletions(-) diff --git a/src/agents/auth-profiles.auth-profile-cooldowns.test.ts b/src/agents/auth-profiles.auth-profile-cooldowns.test.ts index e63910250..01f0722f5 100644 --- a/src/agents/auth-profiles.auth-profile-cooldowns.test.ts +++ b/src/agents/auth-profiles.auth-profile-cooldowns.test.ts @@ -20,7 +20,6 @@ import { clearAuthProfileCooldown, cooldownKey, isProfileInCooldown, - markAuthProfileCooldown, markAuthProfileFailure, markAuthProfileUsed, saveAuthProfileStore, @@ -28,6 +27,24 @@ import { import type { AuthProfileStore } from "./auth-profiles.js"; import { AUTH_STORE_VERSION } from "./auth-profiles/constants.js"; +// Test helpers +const makeStore = (usageStats?: AuthProfileStore["usageStats"]): AuthProfileStore => ({ + version: AUTH_STORE_VERSION, + profiles: { + "openai:default": { type: "api_key", provider: "openai", key: "test" }, + }, + ...(usageStats && { usageStats }), +}); + +async function withTempDir(fn: (tempDir: string) => Promise): Promise { + const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "moltbot-auth-")); + try { + return await fn(tempDir); + } finally { + await fs.rm(tempDir, { recursive: true, force: true }); + } +} + describe("auth profile cooldowns", () => { it("applies exponential backoff with a 1h cap", () => { expect(calculateAuthProfileCooldownMs(1)).toBe(60_000); @@ -39,9 +56,11 @@ describe("auth profile cooldowns", () => { }); describe("cooldownKey", () => { - it("returns profileId when model is not provided", () => { + it("returns profileId when model is not provided or empty", () => { expect(cooldownKey("openai:default")).toBe("openai:default"); expect(cooldownKey("openai:default", undefined)).toBe("openai:default"); + expect(cooldownKey("openai:default", "")).toBe("openai:default"); + expect(cooldownKey("openai:default", " ")).toBe("openai:default"); }); it("returns composite key when model is provided", () => { @@ -52,44 +71,20 @@ describe("cooldownKey", () => { describe("isProfileInCooldown with per-model support", () => { it("returns false when no cooldown exists", () => { - const store: AuthProfileStore = { - version: AUTH_STORE_VERSION, - profiles: { - "openai:default": { type: "api_key", provider: "openai", key: "test" }, - }, - }; + const store = makeStore(); expect(isProfileInCooldown(store, "openai:default")).toBe(false); expect(isProfileInCooldown(store, "openai:default", "gpt-4")).toBe(false); }); it("checks profile-level cooldown when model not provided", () => { - const store: AuthProfileStore = { - version: AUTH_STORE_VERSION, - profiles: { - "openai:default": { type: "api_key", provider: "openai", key: "test" }, - }, - usageStats: { - "openai:default": { cooldownUntil: Date.now() + 60_000 }, - }, - }; + const store = makeStore({ "openai:default": { cooldownUntil: Date.now() + 60_000 } }); expect(isProfileInCooldown(store, "openai:default")).toBe(true); }); it("checks per-model cooldown when model is provided", () => { - const store: AuthProfileStore = { - version: AUTH_STORE_VERSION, - profiles: { - "openai:default": { type: "api_key", provider: "openai", key: "test" }, - }, - usageStats: { - "openai:default:gpt-4": { cooldownUntil: Date.now() + 60_000 }, - }, - }; - // model-specific cooldown exists + const store = makeStore({ "openai:default:gpt-4": { cooldownUntil: Date.now() + 60_000 } }); expect(isProfileInCooldown(store, "openai:default", "gpt-4")).toBe(true); - // different model is not in cooldown expect(isProfileInCooldown(store, "openai:default", "gpt-3.5")).toBe(false); - // profile-level is not in cooldown expect(isProfileInCooldown(store, "openai:default")).toBe(false); }); @@ -97,55 +92,31 @@ describe("isProfileInCooldown with per-model support", () => { const store: AuthProfileStore = { version: AUTH_STORE_VERSION, profiles: { - "github-copilot:default": { - type: "api_key", - provider: "github-copilot", - key: "test", - }, - }, - usageStats: { - // gpt-5.2 is in cooldown (rate limited) - "github-copilot:default:gpt-5.2": { cooldownUntil: Date.now() + 60_000 }, - // gpt-5-mini has no cooldown (unlimited quota) + "github-copilot:default": { type: "api_key", provider: "github-copilot", key: "test" }, }, + usageStats: { "github-copilot:default:gpt-5.2": { cooldownUntil: Date.now() + 60_000 } }, }; expect(isProfileInCooldown(store, "github-copilot:default", "gpt-5.2")).toBe(true); expect(isProfileInCooldown(store, "github-copilot:default", "gpt-5-mini")).toBe(false); }); it("returns false when cooldown has expired", () => { - const store: AuthProfileStore = { - version: AUTH_STORE_VERSION, - profiles: { - "openai:default": { type: "api_key", provider: "openai", key: "test" }, - }, - usageStats: { - "openai:default:gpt-4": { cooldownUntil: Date.now() - 1000 }, // expired - }, - }; + const store = makeStore({ "openai:default:gpt-4": { cooldownUntil: Date.now() - 1000 } }); expect(isProfileInCooldown(store, "openai:default", "gpt-4")).toBe(false); }); }); describe("markAuthProfileUsed with per-model support", () => { it("clears per-model cooldown when model is provided", async () => { - const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "moltbot-auth-")); - const cooldownTime = Date.now() + 60_000; - const store: AuthProfileStore = { - version: AUTH_STORE_VERSION, - profiles: { - "openai:default": { type: "api_key", provider: "openai", key: "test" }, - }, - usageStats: { + await withTempDir(async (tempDir) => { + const cooldownTime = Date.now() + 60_000; + const store = makeStore({ "openai:default": { cooldownUntil: cooldownTime }, "openai:default:gpt-4": { cooldownUntil: cooldownTime, errorCount: 3 }, "openai:default:gpt-3.5": { cooldownUntil: cooldownTime }, - }, - }; - saveAuthProfileStore(store, tempDir); + }); + saveAuthProfileStore(store, tempDir); - try { - // Mark gpt-4 as used (successful) await markAuthProfileUsed({ store, profileId: "openai:default", @@ -153,84 +124,41 @@ describe("markAuthProfileUsed with per-model support", () => { agentDir: tempDir, }); - // Profile-level cooldown should be cleared expect(store.usageStats?.["openai:default"]?.cooldownUntil).toBeUndefined(); - // Per-model cooldown for gpt-4 should be cleared expect(store.usageStats?.["openai:default:gpt-4"]?.cooldownUntil).toBeUndefined(); expect(store.usageStats?.["openai:default:gpt-4"]?.errorCount).toBe(0); - // Per-model cooldown for gpt-3.5 should remain (different model) expect(store.usageStats?.["openai:default:gpt-3.5"]?.cooldownUntil).toBe(cooldownTime); - } finally { - await fs.rm(tempDir, { recursive: true, force: true }); - } + }); }); it("only clears profile-level cooldown when model is not provided", async () => { - const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "moltbot-auth-")); - const cooldownTime = Date.now() + 60_000; - const store: AuthProfileStore = { - version: AUTH_STORE_VERSION, - profiles: { - "openai:default": { type: "api_key", provider: "openai", key: "test" }, - }, - usageStats: { + await withTempDir(async (tempDir) => { + const cooldownTime = Date.now() + 60_000; + const store = makeStore({ "openai:default": { cooldownUntil: cooldownTime }, "openai:default:gpt-4": { cooldownUntil: cooldownTime }, - }, - }; - saveAuthProfileStore(store, tempDir); + }); + saveAuthProfileStore(store, tempDir); - try { - // Mark profile as used without specifying model await markAuthProfileUsed({ store, profileId: "openai:default", agentDir: tempDir }); - // Profile-level cooldown should be cleared expect(store.usageStats?.["openai:default"]?.cooldownUntil).toBeUndefined(); - // Per-model cooldown should remain (no model specified) expect(store.usageStats?.["openai:default:gpt-4"]?.cooldownUntil).toBe(cooldownTime); - } finally { - await fs.rm(tempDir, { recursive: true, force: true }); - } - }); -}); - -describe("cooldownKey edge cases", () => { - it("treats empty string model the same as undefined", () => { - // Empty string should be treated as "no model" to avoid trailing colon - expect(cooldownKey("openai:default", "")).toBe("openai:default"); - expect(cooldownKey("openai:default", " ")).toBe("openai:default"); + }); }); }); describe("isProfileInCooldown backward compatibility", () => { it("returns true for any model when profile-level cooldown exists", () => { - const store: AuthProfileStore = { - version: AUTH_STORE_VERSION, - profiles: { - "openai:default": { type: "api_key", provider: "openai", key: "test" }, - }, - usageStats: { - "openai:default": { cooldownUntil: Date.now() + 60_000 }, // profile-level only - }, - }; - // Any model should be blocked when profile-level cooldown exists + const store = makeStore({ "openai:default": { cooldownUntil: Date.now() + 60_000 } }); expect(isProfileInCooldown(store, "openai:default", "gpt-4")).toBe(true); expect(isProfileInCooldown(store, "openai:default", "gpt-3.5")).toBe(true); expect(isProfileInCooldown(store, "openai:default", "o1-preview")).toBe(true); - // Profile-level check also works expect(isProfileInCooldown(store, "openai:default")).toBe(true); }); it("checks disabledUntil for per-model cooldowns (billing failures)", () => { - const store: AuthProfileStore = { - version: AUTH_STORE_VERSION, - profiles: { - "openai:default": { type: "api_key", provider: "openai", key: "test" }, - }, - usageStats: { - "openai:default:gpt-4": { disabledUntil: Date.now() + 60_000 }, // billing failure - }, - }; + const store = makeStore({ "openai:default:gpt-4": { disabledUntil: Date.now() + 60_000 } }); expect(isProfileInCooldown(store, "openai:default", "gpt-4")).toBe(true); expect(isProfileInCooldown(store, "openai:default", "gpt-3.5")).toBe(false); }); @@ -238,16 +166,10 @@ describe("isProfileInCooldown backward compatibility", () => { describe("markAuthProfileFailure with per-model support", () => { it("tracks failure per model when model is provided", async () => { - const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "moltbot-auth-")); - const store: AuthProfileStore = { - version: AUTH_STORE_VERSION, - profiles: { - "openai:default": { type: "api_key", provider: "openai", key: "test" }, - }, - }; - saveAuthProfileStore(store, tempDir); + await withTempDir(async (tempDir) => { + const store = makeStore(); + saveAuthProfileStore(store, tempDir); - try { await markAuthProfileFailure({ store, profileId: "openai:default", @@ -256,29 +178,18 @@ describe("markAuthProfileFailure with per-model support", () => { agentDir: tempDir, }); - // Per-model key should have cooldown expect(store.usageStats?.["openai:default:gpt-4"]?.cooldownUntil).toBeGreaterThan(Date.now()); expect(store.usageStats?.["openai:default:gpt-4"]?.errorCount).toBe(1); - // Profile-level should NOT have cooldown (only model-specific) expect(store.usageStats?.["openai:default"]?.cooldownUntil).toBeUndefined(); - // Other models should not be affected expect(store.usageStats?.["openai:default:gpt-3.5"]).toBeUndefined(); - } finally { - await fs.rm(tempDir, { recursive: true, force: true }); - } + }); }); it("tracks failure at profile level when model is not provided", async () => { - const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "moltbot-auth-")); - const store: AuthProfileStore = { - version: AUTH_STORE_VERSION, - profiles: { - "openai:default": { type: "api_key", provider: "openai", key: "test" }, - }, - }; - saveAuthProfileStore(store, tempDir); + await withTempDir(async (tempDir) => { + const store = makeStore(); + saveAuthProfileStore(store, tempDir); - try { await markAuthProfileFailure({ store, profileId: "openai:default", @@ -286,25 +197,16 @@ describe("markAuthProfileFailure with per-model support", () => { agentDir: tempDir, }); - // Profile-level key should have cooldown expect(store.usageStats?.["openai:default"]?.cooldownUntil).toBeGreaterThan(Date.now()); expect(store.usageStats?.["openai:default"]?.errorCount).toBe(1); - } finally { - await fs.rm(tempDir, { recursive: true, force: true }); - } + }); }); it("tracks billing failures with disabledUntil per model", async () => { - const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "moltbot-auth-")); - const store: AuthProfileStore = { - version: AUTH_STORE_VERSION, - profiles: { - "openai:default": { type: "api_key", provider: "openai", key: "test" }, - }, - }; - saveAuthProfileStore(store, tempDir); + await withTempDir(async (tempDir) => { + const store = makeStore(); + saveAuthProfileStore(store, tempDir); - try { await markAuthProfileFailure({ store, profileId: "openai:default", @@ -313,62 +215,23 @@ describe("markAuthProfileFailure with per-model support", () => { agentDir: tempDir, }); - // Billing failures use disabledUntil instead of cooldownUntil expect(store.usageStats?.["openai:default:gpt-4"]?.disabledUntil).toBeGreaterThan(Date.now()); expect(store.usageStats?.["openai:default:gpt-4"]?.disabledReason).toBe("billing"); - } finally { - await fs.rm(tempDir, { recursive: true, force: true }); - } - }); -}); - -describe("markAuthProfileCooldown with per-model support", () => { - it("marks cooldown per model when model is provided", async () => { - const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "moltbot-auth-")); - const store: AuthProfileStore = { - version: AUTH_STORE_VERSION, - profiles: { - "openai:default": { type: "api_key", provider: "openai", key: "test" }, - }, - }; - saveAuthProfileStore(store, tempDir); - - try { - await markAuthProfileCooldown({ - store, - profileId: "openai:default", - model: "gpt-4", - agentDir: tempDir, - }); - - // Per-model key should have cooldown - expect(store.usageStats?.["openai:default:gpt-4"]?.cooldownUntil).toBeGreaterThan(Date.now()); - // Profile-level should NOT have cooldown - expect(store.usageStats?.["openai:default"]?.cooldownUntil).toBeUndefined(); - } finally { - await fs.rm(tempDir, { recursive: true, force: true }); - } + }); }); }); describe("clearAuthProfileCooldown with per-model support", () => { it("clears per-model cooldown when model is provided", async () => { - const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "moltbot-auth-")); - const cooldownTime = Date.now() + 60_000; - const store: AuthProfileStore = { - version: AUTH_STORE_VERSION, - profiles: { - "openai:default": { type: "api_key", provider: "openai", key: "test" }, - }, - usageStats: { + await withTempDir(async (tempDir) => { + const cooldownTime = Date.now() + 60_000; + const store = makeStore({ "openai:default": { cooldownUntil: cooldownTime }, "openai:default:gpt-4": { cooldownUntil: cooldownTime, errorCount: 3 }, "openai:default:gpt-3.5": { cooldownUntil: cooldownTime }, - }, - }; - saveAuthProfileStore(store, tempDir); + }); + saveAuthProfileStore(store, tempDir); - try { await clearAuthProfileCooldown({ store, profileId: "openai:default", @@ -376,47 +239,27 @@ describe("clearAuthProfileCooldown with per-model support", () => { agentDir: tempDir, }); - // Per-model cooldown for gpt-4 should be cleared expect(store.usageStats?.["openai:default:gpt-4"]?.cooldownUntil).toBeUndefined(); expect(store.usageStats?.["openai:default:gpt-4"]?.errorCount).toBe(0); - // Profile-level cooldown should remain (different key) expect(store.usageStats?.["openai:default"]?.cooldownUntil).toBe(cooldownTime); - // Other model cooldown should remain expect(store.usageStats?.["openai:default:gpt-3.5"]?.cooldownUntil).toBe(cooldownTime); - } finally { - await fs.rm(tempDir, { recursive: true, force: true }); - } + }); }); it("clears profile-level cooldown when model is not provided", async () => { - const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "moltbot-auth-")); - const cooldownTime = Date.now() + 60_000; - const store: AuthProfileStore = { - version: AUTH_STORE_VERSION, - profiles: { - "openai:default": { type: "api_key", provider: "openai", key: "test" }, - }, - usageStats: { + await withTempDir(async (tempDir) => { + const cooldownTime = Date.now() + 60_000; + const store = makeStore({ "openai:default": { cooldownUntil: cooldownTime, errorCount: 2 }, "openai:default:gpt-4": { cooldownUntil: cooldownTime }, - }, - }; - saveAuthProfileStore(store, tempDir); - - try { - await clearAuthProfileCooldown({ - store, - profileId: "openai:default", - agentDir: tempDir, }); + saveAuthProfileStore(store, tempDir); + + await clearAuthProfileCooldown({ store, profileId: "openai:default", agentDir: tempDir }); - // Profile-level cooldown should be cleared expect(store.usageStats?.["openai:default"]?.cooldownUntil).toBeUndefined(); expect(store.usageStats?.["openai:default"]?.errorCount).toBe(0); - // Per-model cooldown should remain (different key) expect(store.usageStats?.["openai:default:gpt-4"]?.cooldownUntil).toBe(cooldownTime); - } finally { - await fs.rm(tempDir, { recursive: true, force: true }); - } + }); }); });