feat: add round-robin rotation and cooldown for auth profiles

Adds usage tracking to auth profiles for automatic rotation:

- ProfileUsageStats type with lastUsed, cooldownUntil, errorCount
- markAuthProfileUsed(): tracks successful usage, resets errors
- markAuthProfileCooldown(): applies exponential backoff (1/5/25/60min)
- isProfileInCooldown(): checks if profile should be skipped
- orderProfilesByMode(): now sorts by lastUsed (oldest first)

On auth/rate-limit failures, profiles are marked for cooldown before
rotation. On success, usage is recorded for round-robin ordering.

This enables automatic load distribution across multiple accounts
(e.g., Antigravity 5-hour rate limit windows).
This commit is contained in:
Muhammed Mukhthar CM
2026-01-06 04:43:59 +00:00
parent 06df6a955a
commit ce6c7737c1
2 changed files with 141 additions and 6 deletions

View File

@@ -32,10 +32,19 @@ export type OAuthCredential = OAuthCredentials & {
export type AuthProfileCredential = ApiKeyCredential | OAuthCredential;
/** Per-profile usage statistics for round-robin and cooldown tracking */
export type ProfileUsageStats = {
lastUsed?: number;
cooldownUntil?: number;
errorCount?: number;
};
export type AuthProfileStore = {
version: number;
profiles: Record<string, AuthProfileCredential>;
lastGood?: Record<string, string>;
/** Usage statistics per profile for round-robin rotation */
usageStats?: Record<string, ProfileUsageStats>;
};
type LegacyAuthStore = Record<string, AuthProfileCredential>;
@@ -183,6 +192,10 @@ function coerceAuthStore(raw: unknown): AuthProfileStore | null {
record.lastGood && typeof record.lastGood === "object"
? (record.lastGood as Record<string, string>)
: undefined,
usageStats:
record.usageStats && typeof record.usageStats === "object"
? (record.usageStats as Record<string, ProfileUsageStats>)
: undefined,
};
}
@@ -300,6 +313,7 @@ export function saveAuthProfileStore(store: AuthProfileStore): void {
version: AUTH_STORE_VERSION,
profiles: store.profiles,
lastGood: store.lastGood ?? undefined,
usageStats: store.usageStats ?? undefined,
} satisfies AuthProfileStore;
saveJsonFile(authPath, payload);
}
@@ -322,6 +336,85 @@ export function listProfilesForProvider(
.map(([id]) => id);
}
/**
* 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?.cooldownUntil) return false;
return Date.now() < stats.cooldownUntil;
}
/**
* Mark a profile as successfully used. Resets error count and updates lastUsed.
*/
export function markAuthProfileUsed(params: {
store: AuthProfileStore;
profileId: string;
}): void {
const { store, profileId } = params;
if (!store.profiles[profileId]) return;
store.usageStats = store.usageStats ?? {};
store.usageStats[profileId] = {
...store.usageStats[profileId],
lastUsed: Date.now(),
errorCount: 0,
cooldownUntil: undefined,
};
saveAuthProfileStore(store);
}
/**
* Mark a profile as failed/rate-limited. Applies exponential backoff cooldown.
* Cooldown times: 1min, 5min, 25min, max 1 hour.
*/
export function markAuthProfileCooldown(params: {
store: AuthProfileStore;
profileId: string;
}): void {
const { store, profileId } = params;
if (!store.profiles[profileId]) return;
store.usageStats = store.usageStats ?? {};
const existing = store.usageStats[profileId] ?? {};
const errorCount = (existing.errorCount ?? 0) + 1;
// Exponential backoff: 1min, 5min, 25min, capped at 1h
const backoffMs = Math.min(
60 * 60 * 1000, // 1 hour max
60 * 1000 * Math.pow(5, Math.min(errorCount - 1, 3)),
);
store.usageStats[profileId] = {
...existing,
errorCount,
cooldownUntil: Date.now() + backoffMs,
};
saveAuthProfileStore(store);
}
/**
* Clear cooldown for a profile (e.g., manual reset).
*/
export function clearAuthProfileCooldown(params: {
store: AuthProfileStore;
profileId: string;
}): void {
const { store, profileId } = params;
if (!store.usageStats?.[profileId]) return;
store.usageStats[profileId] = {
...store.usageStats[profileId],
errorCount: 0,
cooldownUntil: undefined,
};
saveAuthProfileStore(store);
}
export function resolveAuthProfileOrder(params: {
cfg?: ClawdbotConfig;
store: AuthProfileStore;
@@ -376,14 +469,50 @@ function orderProfilesByMode(
order: string[],
store: AuthProfileStore,
): string[] {
const scored = order.map((profileId) => {
const now = Date.now();
// Partition into available and in-cooldown
const available: string[] = [];
const inCooldown: string[] = [];
for (const profileId of order) {
if (isProfileInCooldown(store, profileId)) {
inCooldown.push(profileId);
} else {
available.push(profileId);
}
}
// Sort available profiles by lastUsed (oldest first = round-robin)
// Then by type (oauth preferred over api_key)
const scored = available.map((profileId) => {
const type = store.profiles[profileId]?.type;
const score = type === "oauth" ? 0 : type === "api_key" ? 1 : 2;
return { profileId, score };
const typeScore = type === "oauth" ? 0 : type === "api_key" ? 1 : 2;
const lastUsed = store.usageStats?.[profileId]?.lastUsed ?? 0;
return { profileId, typeScore, lastUsed };
});
return scored
.sort((a, b) => a.score - b.score)
// Primary sort: lastUsed (oldest first for round-robin)
// Secondary sort: type preference (oauth > api_key)
const sorted = scored
.sort((a, b) => {
// First by lastUsed (oldest first)
if (a.lastUsed !== b.lastUsed) return a.lastUsed - b.lastUsed;
// Then by type
return a.typeScore - b.typeScore;
})
.map((entry) => entry.profileId);
// Append cooldown profiles at the end (sorted by cooldown expiry, soonest first)
const cooldownSorted = inCooldown
.map((profileId) => ({
profileId,
cooldownUntil: store.usageStats?.[profileId]?.cooldownUntil ?? now,
}))
.sort((a, b) => a.cooldownUntil - b.cooldownUntil)
.map((entry) => entry.profileId);
return [...sorted, ...cooldownSorted];
}
export async function resolveApiKeyForProfile(params: {

View File

@@ -24,7 +24,7 @@ import {
} from "../process/command-queue.js";
import { resolveUserPath } from "../utils.js";
import { resolveClawdbotAgentDir } from "./agent-paths.js";
import { markAuthProfileGood } from "./auth-profiles.js";
import { markAuthProfileGood, markAuthProfileUsed, markAuthProfileCooldown } from "./auth-profiles.js";
import type { BashElevatedDefaults } from "./bash-tools.js";
import { DEFAULT_MODEL, DEFAULT_PROVIDER } from "./defaults.js";
import {
@@ -954,6 +954,10 @@ export async function runEmbeddedPiAgent(params: {
const authFailure = isAuthAssistantError(lastAssistant);
const rateLimitFailure = isRateLimitAssistantError(lastAssistant);
if (!aborted && (authFailure || rateLimitFailure)) {
// Mark current profile for cooldown before rotating
if (lastProfileId) {
markAuthProfileCooldown({ store: authStore, profileId: lastProfileId });
}
const rotated = await advanceAuthProfile();
if (rotated) {
continue;
@@ -1040,6 +1044,8 @@ export async function runEmbeddedPiAgent(params: {
provider,
profileId: lastProfileId,
});
// Track usage for round-robin rotation
markAuthProfileUsed({ store: authStore, profileId: lastProfileId });
}
return {
payloads: payloads.length ? payloads : undefined,