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
This commit is contained in:
Bruno Guidolim 2026-01-29 00:59:55 +01:00
parent 9d0841a86c
commit e3caf006bb

View File

@ -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<string, ProfileUsageStats>,
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);
}