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:
@@ -32,10 +32,19 @@ export type OAuthCredential = OAuthCredentials & {
|
|||||||
|
|
||||||
export type AuthProfileCredential = ApiKeyCredential | OAuthCredential;
|
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 = {
|
export type AuthProfileStore = {
|
||||||
version: number;
|
version: number;
|
||||||
profiles: Record<string, AuthProfileCredential>;
|
profiles: Record<string, AuthProfileCredential>;
|
||||||
lastGood?: Record<string, string>;
|
lastGood?: Record<string, string>;
|
||||||
|
/** Usage statistics per profile for round-robin rotation */
|
||||||
|
usageStats?: Record<string, ProfileUsageStats>;
|
||||||
};
|
};
|
||||||
|
|
||||||
type LegacyAuthStore = Record<string, AuthProfileCredential>;
|
type LegacyAuthStore = Record<string, AuthProfileCredential>;
|
||||||
@@ -183,6 +192,10 @@ function coerceAuthStore(raw: unknown): AuthProfileStore | null {
|
|||||||
record.lastGood && typeof record.lastGood === "object"
|
record.lastGood && typeof record.lastGood === "object"
|
||||||
? (record.lastGood as Record<string, string>)
|
? (record.lastGood as Record<string, string>)
|
||||||
: undefined,
|
: 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,
|
version: AUTH_STORE_VERSION,
|
||||||
profiles: store.profiles,
|
profiles: store.profiles,
|
||||||
lastGood: store.lastGood ?? undefined,
|
lastGood: store.lastGood ?? undefined,
|
||||||
|
usageStats: store.usageStats ?? undefined,
|
||||||
} satisfies AuthProfileStore;
|
} satisfies AuthProfileStore;
|
||||||
saveJsonFile(authPath, payload);
|
saveJsonFile(authPath, payload);
|
||||||
}
|
}
|
||||||
@@ -322,6 +336,85 @@ export function listProfilesForProvider(
|
|||||||
.map(([id]) => id);
|
.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: {
|
export function resolveAuthProfileOrder(params: {
|
||||||
cfg?: ClawdbotConfig;
|
cfg?: ClawdbotConfig;
|
||||||
store: AuthProfileStore;
|
store: AuthProfileStore;
|
||||||
@@ -376,14 +469,50 @@ function orderProfilesByMode(
|
|||||||
order: string[],
|
order: string[],
|
||||||
store: AuthProfileStore,
|
store: AuthProfileStore,
|
||||||
): string[] {
|
): 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 type = store.profiles[profileId]?.type;
|
||||||
const score = type === "oauth" ? 0 : type === "api_key" ? 1 : 2;
|
const typeScore = type === "oauth" ? 0 : type === "api_key" ? 1 : 2;
|
||||||
return { profileId, score };
|
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);
|
.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: {
|
export async function resolveApiKeyForProfile(params: {
|
||||||
|
|||||||
@@ -24,7 +24,7 @@ import {
|
|||||||
} from "../process/command-queue.js";
|
} from "../process/command-queue.js";
|
||||||
import { resolveUserPath } from "../utils.js";
|
import { resolveUserPath } from "../utils.js";
|
||||||
import { resolveClawdbotAgentDir } from "./agent-paths.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 type { BashElevatedDefaults } from "./bash-tools.js";
|
||||||
import { DEFAULT_MODEL, DEFAULT_PROVIDER } from "./defaults.js";
|
import { DEFAULT_MODEL, DEFAULT_PROVIDER } from "./defaults.js";
|
||||||
import {
|
import {
|
||||||
@@ -954,6 +954,10 @@ export async function runEmbeddedPiAgent(params: {
|
|||||||
const authFailure = isAuthAssistantError(lastAssistant);
|
const authFailure = isAuthAssistantError(lastAssistant);
|
||||||
const rateLimitFailure = isRateLimitAssistantError(lastAssistant);
|
const rateLimitFailure = isRateLimitAssistantError(lastAssistant);
|
||||||
if (!aborted && (authFailure || rateLimitFailure)) {
|
if (!aborted && (authFailure || rateLimitFailure)) {
|
||||||
|
// Mark current profile for cooldown before rotating
|
||||||
|
if (lastProfileId) {
|
||||||
|
markAuthProfileCooldown({ store: authStore, profileId: lastProfileId });
|
||||||
|
}
|
||||||
const rotated = await advanceAuthProfile();
|
const rotated = await advanceAuthProfile();
|
||||||
if (rotated) {
|
if (rotated) {
|
||||||
continue;
|
continue;
|
||||||
@@ -1040,6 +1044,8 @@ export async function runEmbeddedPiAgent(params: {
|
|||||||
provider,
|
provider,
|
||||||
profileId: lastProfileId,
|
profileId: lastProfileId,
|
||||||
});
|
});
|
||||||
|
// Track usage for round-robin rotation
|
||||||
|
markAuthProfileUsed({ store: authStore, profileId: lastProfileId });
|
||||||
}
|
}
|
||||||
return {
|
return {
|
||||||
payloads: payloads.length ? payloads : undefined,
|
payloads: payloads.length ? payloads : undefined,
|
||||||
|
|||||||
Reference in New Issue
Block a user