romtuck/openclaw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add multi-account load balancing for OAuth providers

Browse Source 
- Add AccountManager with health scoring, rate limit tracking, and quota awareness
- Implement 3 selection strategies: hybrid, sticky, round-robin
- Add CLI commands: clawdbot accounts list/status/reset
- Wire into pi-embedded-runner for automatic failover
- Add config schema for auth.multiAccount

Enables distributing requests across multiple OAuth accounts for higher throughput
and automatic failover on rate limits.
This commit is contained in:
PREETHAM1590 2026-01-27 16:12:34 +05:30
parent 3e07bd8b48
commit a16abbd3e1
20 changed files with 2019 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

280
src/agents/multi-account/account-manager.ts Normal file
View File

@ -0,0 +1,280 @@
/**
* Account Manager
*
* Central manager for multi-account load balancing.
*/
import { RateLimitTracker } from "./rate-limit-tracker.js";
import { HealthScorer } from "./health-scorer.js";
import { QuotaTracker } from "./quota-tracker.js";
import { createStrategy, STRATEGY_NAMES, type Strategy } from "./strategies.js";
import { DEFAULT_COOLDOWN_MS, MAX_WAIT_BEFORE_ERROR_MS } from "./constants.js";
export interface Account {
profileId: string;
provider: string;
email: string;
isInvalid: boolean;
invalidReason: string | null;
lastUsed: number | null;
}
export interface AccountManagerConfig {
strategy?: string;
provider?: string;
defaultCooldownMs?: number;
maxWaitBeforeErrorMs?: number;
}
export interface AuthProfileStore {
profiles: Record<string, {
type: string;
provider: string;
email?: string;
access?: string;
refresh?: string;
expires?: number;
projectId?: string;
}>;
}
export class AccountManager {
private provider: string;
private authStore: AuthProfileStore | null = null;
private accounts = new Map<string, Account>();
private strategy: Strategy;
private defaultCooldownMs: number;
private invalidProfiles = new Set<string>();
private tokenCache = new Map<string, string>();
private projectCache = new Map<string, string>();
public rateLimitTracker = new RateLimitTracker();
public healthScorer = new HealthScorer();
public quotaTracker = new QuotaTracker();
constructor(config: AccountManagerConfig = {}) {
this.provider = config.provider ?? "google-antigravity";
this.defaultCooldownMs = config.defaultCooldownMs ?? DEFAULT_COOLDOWN_MS;
this.strategy = createStrategy(config.strategy ?? "hybrid");
this.strategy.init(this);
}
async initialize(
authStore: AuthProfileStore,
strategyOverride?: string
): Promise<this> {
this.authStore = authStore;
if (strategyOverride && STRATEGY_NAMES.includes(strategyOverride.toLowerCase() as any)) {
this.strategy = createStrategy(strategyOverride);
this.strategy.init(this);
}
this.loadAccountsFromStore();
return this;
}
private loadAccountsFromStore(): void {
if (!this.authStore?.profiles) return;
for (const [profileId, profile] of Object.entries(this.authStore.profiles)) {
if (profile.provider !== this.provider) continue;
if (profile.type !== "oauth") continue;
this.accounts.set(profileId, {
profileId,
provider: profile.provider,
email: profile.email ?? profileId,
isInvalid: this.invalidProfiles.has(profileId),
invalidReason: null,
lastUsed: null,
});
}
}
getProfilesForProvider(): string[] {
return Array.from(this.accounts.keys()).filter(
(id) => !this.invalidProfiles.has(id)
);
}
getAccountByProfileId(profileId: string): Account | null {
return this.accounts.get(profileId) ?? null;
}
getAllAccounts(): Account[] {
return Array.from(this.accounts.values());
}
getAccountCount(): number {
return this.accounts.size;
}
getAvailableAccounts(modelId: string): Account[] {
return this.getProfilesForProvider()
.filter((id) => !this.rateLimitTracker.isRateLimited(id, modelId))
.map((id) => this.accounts.get(id)!)
.filter(Boolean);
}
selectAccount(modelId: string): { account: Account | null; waitMs: number } {
this.rateLimitTracker.clearExpired();
return this.strategy.select(modelId) as { account: Account | null; waitMs: number };
}
notifySuccess(profileId: string, modelId: string): void {
this.rateLimitTracker.clearRateLimit(profileId, modelId);
this.healthScorer.recordSuccess(profileId);
this.strategy.notifySuccess(profileId, modelId);
const account = this.accounts.get(profileId);
if (account) account.lastUsed = Date.now();
}
notifyRateLimit(profileId: string, modelId: string, cooldownMs?: number): void {
this.markRateLimited(profileId, modelId, cooldownMs);
this.healthScorer.recordRateLimit(profileId);
this.strategy.notifyRateLimit(profileId, modelId);
}
notifyFailure(profileId: string, _modelId: string): void {
this.healthScorer.recordFailure(profileId);
this.rateLimitTracker.incrementFailures(profileId);
if (this.rateLimitTracker.shouldExtendCooldown(profileId)) {
this.rateLimitTracker.applyExtendedCooldown(profileId);
}
}
markRateLimited(profileId: string, modelId: string, cooldownMs?: number): void {
this.rateLimitTracker.markRateLimited(
profileId,
modelId,
cooldownMs ?? this.defaultCooldownMs
);
}
markInvalid(profileId: string, reason: string): void {
this.invalidProfiles.add(profileId);
const account = this.accounts.get(profileId);
if (account) {
account.isInvalid = true;
account.invalidReason = reason;
}
}
isAllRateLimited(modelId: string): boolean {
const profiles = this.getProfilesForProvider();
return this.rateLimitTracker.areAllRateLimited(profiles, modelId);
}
getMinWaitTimeMs(modelId: string): number {
const profiles = this.getProfilesForProvider();
return this.rateLimitTracker.getMinWaitTime(profiles, modelId);
}
resetAllRateLimits(): void {
this.rateLimitTracker.clearExpired();
}
clearExpiredLimits(): void {
this.rateLimitTracker.clearExpired();
}
getConsecutiveFailures(profileId: string): number {
return this.rateLimitTracker.getFailureCount(profileId);
}
async getTokenForAccount(account: Account): Promise<string> {
const cached = this.tokenCache.get(account.profileId);
if (cached) return cached;
const profile = this.authStore?.profiles?.[account.profileId];
if (!profile || profile.type !== "oauth") {
throw new Error(`No OAuth credentials for ${account.profileId}`);
}
if (profile.expires && Date.now() >= profile.expires) {
throw new Error(`Token expired for ${account.profileId}`);
}
const token = profile.access;
if (!token) throw new Error(`No access token for ${account.profileId}`);
this.tokenCache.set(account.profileId, token);
return token;
}
async getProjectForAccount(account: Account): Promise<string> {
const cached = this.projectCache.get(account.profileId);
if (cached) return cached;
const profile = this.authStore?.profiles?.[account.profileId];
const projectId = profile?.projectId;
if (projectId) {
this.projectCache.set(account.profileId, projectId);
return projectId;
}
throw new Error(`No project ID for ${account.profileId}`);
}
clearTokenCache(profileId?: string): void {
if (profileId) this.tokenCache.delete(profileId);
else this.tokenCache.clear();
}
clearProjectCache(profileId?: string): void {
if (profileId) this.projectCache.delete(profileId);
else this.projectCache.clear();
}
getStrategyLabel(): string {
return this.strategy.getLabel();
}
getStatus() {
const profiles = this.getProfilesForProvider();
const invalid = Array.from(this.invalidProfiles);
const rateLimited = profiles.filter(
(id) => this.rateLimitTracker.getSoonestReset(id) !== null
);
const available = profiles.filter(
(id) => !rateLimited.includes(id) && !invalid.includes(id)
);
return {
total: profiles.length + invalid.length,
available: available.length,
rateLimited: rateLimited.length,
invalid: invalid.length,
summary: `${available.length}/${profiles.length + invalid.length} available`,
accounts: this.getAllAccounts().map((acc) => ({
email: acc.email,
isInvalid: acc.isInvalid,
invalidReason: acc.invalidReason,
lastUsed: acc.lastUsed,
healthScore: this.healthScorer.getScore(acc.profileId),
})),
};
}
toJSON() {
return {
provider: this.provider,
invalidProfiles: Array.from(this.invalidProfiles),
rateLimits: this.rateLimitTracker.toJSON(),
health: this.healthScorer.toJSON(),
quotas: this.quotaTracker.toJSON(),
};
}
fromJSON(data: ReturnType<typeof this.toJSON>): void {
if (data?.invalidProfiles) {
for (const id of data.invalidProfiles) this.invalidProfiles.add(id);
}
if (data?.rateLimits) this.rateLimitTracker.fromJSON(data.rateLimits);
if (data?.health) this.healthScorer.fromJSON(data.health);
if (data?.quotas) this.quotaTracker.fromJSON(data.quotas);
}
}

65
src/agents/multi-account/constants.ts Normal file
View File

@ -0,0 +1,65 @@
/**
* Constants for multi-account management
* Aligned with antigravity-claude-proxy defaults
*/
// Default cooldown when rate-limited (10 seconds)
export const DEFAULT_COOLDOWN_MS = 10_000;
// Max time to wait before throwing error (2 minutes)
export const MAX_WAIT_BEFORE_ERROR_MS = 120_000;
// Deduplication window for rate limits (prevents thundering herd)
export const RATE_LIMIT_DEDUP_WINDOW_MS = 2_000;
// Reset rate limit state after this inactivity period
export const RATE_LIMIT_STATE_RESET_MS = 60_000;
// First retry delay on 429 (quick retry)
export const FIRST_RETRY_DELAY_MS = 1_000;
// Delay before switching accounts
export const SWITCH_ACCOUNT_DELAY_MS = 500;
// Consecutive failures before extended cooldown
export const MAX_CONSECUTIVE_FAILURES = 3;
// Extended cooldown for consistently failing accounts (5 minutes)
export const EXTENDED_COOLDOWN_MS = 300_000;
// Exponential backoff tiers for capacity exhaustion
export const BACKOFF_TIERS_MS = [1_000, 2_000, 5_000, 10_000];
// Quota exhausted backoff tiers (progressive: 1min, 5min, 30min, 2h)
export const QUOTA_EXHAUSTED_BACKOFF_TIERS_MS = [60_000, 300_000, 1_800_000, 7_200_000];
// Minimum backoff to prevent loops
export const MIN_BACKOFF_MS = 500;
// Backoff by error type
export const BACKOFF_BY_ERROR_TYPE = {
RATE_LIMIT_EXCEEDED: 5_000,
QUOTA_EXHAUSTED: 60_000,
MODEL_CAPACITY_EXHAUSTED: 10_000,
SERVER_ERROR: 2_000,
UNKNOWN: 5_000,
};
// Health score defaults
export const HEALTH_SCORE = {
initial: 100,
maxScore: 100,
minScore: 0,
successBonus: 5,
failurePenalty: 15,
rateLimitPenalty: 10,
recoveryPerMinute: 2,
};
// Subscription tier weights (higher = more capacity expected)
export const TIER_WEIGHTS: Record<string, number> = {
ultra: 3,
pro: 2,
free: 1,
unknown: 1,
};

97
src/agents/multi-account/health-scorer.ts Normal file
View File

@ -0,0 +1,97 @@
/**
* Health Scorer
*
* Tracks health scores for accounts based on:
* - Success/failure rates
* - Rate limit frequency
* - Recovery over time
*/
import { HEALTH_SCORE } from "./constants.js";
export class HealthScorer {
private scores = new Map<string, number>();
private lastUpdate = new Map<string, number>();
getScore(profileId: string): number {
const baseScore = this.scores.get(profileId) ?? HEALTH_SCORE.initial;
const lastUpdateTime = this.lastUpdate.get(profileId);
if (!lastUpdateTime) return baseScore;
// Apply time-based recovery
const minutesSinceUpdate = (Date.now() - lastUpdateTime) / 60_000;
const recovery = Math.floor(
minutesSinceUpdate * HEALTH_SCORE.recoveryPerMinute
);
return Math.min(HEALTH_SCORE.maxScore, baseScore + recovery);
}
recordSuccess(profileId: string): void {
const current = this.getScore(profileId);
const newScore = Math.min(
HEALTH_SCORE.maxScore,
current + HEALTH_SCORE.successBonus
);
this.scores.set(profileId, newScore);
this.lastUpdate.set(profileId, Date.now());
}
recordFailure(profileId: string): void {
const current = this.getScore(profileId);
const newScore = Math.max(
HEALTH_SCORE.minScore,
current - HEALTH_SCORE.failurePenalty
);
this.scores.set(profileId, newScore);
this.lastUpdate.set(profileId, Date.now());
}
recordRateLimit(profileId: string): void {
const current = this.getScore(profileId);
const newScore = Math.max(
HEALTH_SCORE.minScore,
current - HEALTH_SCORE.rateLimitPenalty
);
this.scores.set(profileId, newScore);
this.lastUpdate.set(profileId, Date.now());
}
getSortedByHealth(
profileIds: string[]
): Array<{ profileId: string; score: number }> {
return profileIds
.map((profileId) => ({
profileId,
score: this.getScore(profileId),
}))
.sort((a, b) => b.score - a.score);
}
toJSON(): {
scores: Record<string, number>;
lastUpdate: Record<string, number>;
} {
return {
scores: Object.fromEntries(this.scores),
lastUpdate: Object.fromEntries(this.lastUpdate),
};
}
fromJSON(data: {
scores?: Record<string, number>;
lastUpdate?: Record<string, number>;
}): void {
if (data?.scores) {
for (const [key, value] of Object.entries(data.scores)) {
this.scores.set(key, value);
}
}
if (data?.lastUpdate) {
for (const [key, value] of Object.entries(data.lastUpdate)) {
this.lastUpdate.set(key, value);
}
}
}
}

50
src/agents/multi-account/index.ts Normal file
View File

@ -0,0 +1,50 @@
/**
* Multi-Account Module
*
* Provides multi-account load balancing for Clawdbot providers.
*/
export { AccountManager } from "./account-manager.js";
export type { Account, AccountManagerConfig, AuthProfileStore } from "./account-manager.js";
export { RateLimitTracker } from "./rate-limit-tracker.js";
export type { RateLimitState, BackoffResult } from "./rate-limit-tracker.js";
export { HealthScorer } from "./health-scorer.js";
export { QuotaTracker } from "./quota-tracker.js";
export type { ModelQuota, AccountQuota } from "./quota-tracker.js";
export {
HybridStrategy,
StickyStrategy,
RoundRobinStrategy,
createStrategy,
STRATEGY_NAMES,
} from "./strategies.js";
export type { Strategy, SelectionResult, StrategyName } from "./strategies.js";
export * from "./constants.js";
// Integration helpers
export {
isMultiAccountEnabled,
getOrCreateManager,
getManager,
selectAccountForModel,
notifyMultiAccountSuccess,
notifyMultiAccountRateLimit,
notifyMultiAccountFailure,
notifyMultiAccountInvalid,
getMultiAccountStatus,
clearAllManagers,
resetManagerState,
type MultiAccountConfig,
type MultiAccountSelection,
} from "./model-auth-integration.js";
// Profile order integration
export {
resolveProfileOrderWithMultiAccount,
resolveProfileOrderWithMultiAccountSync,
} from "./profile-order.js";

175
src/agents/multi-account/integration.ts Normal file
View File

@ -0,0 +1,175 @@
/**
* Multi-Account Integration for Pi Runner
*
* This module provides hooks to integrate multi-account load balancing
* with Clawdbot's pi-embedded-runner.
*/
import { AccountManager, type AuthProfileStore } from "./account-manager.js";
// eslint-disable-next-line @typescript-eslint/no-explicit-any
type ClawdbotConfig = any;
// Singleton managers per provider
const managers = new Map<string, AccountManager>();
/**
* Get or create an AccountManager for a provider
*/
export async function getAccountManager(
provider: string,
authStore: AuthProfileStore,
cfg?: ClawdbotConfig
): Promise<AccountManager> {
if (managers.has(provider)) {
return managers.get(provider)!;
}
const multiAccountConfig = (cfg as any)?.auth?.multiAccount;
const manager = new AccountManager({
provider,
strategy: multiAccountConfig?.strategy ?? "hybrid",
defaultCooldownMs: multiAccountConfig?.defaultCooldownMs,
});
await manager.initialize(authStore, multiAccountConfig?.strategy);
managers.set(provider, manager);
return manager;
}
/**
* Check if multi-account is enabled for a provider
*/
export function isMultiAccountEnabled(
provider: string,
cfg?: ClawdbotConfig
): boolean {
const multiAccountConfig = (cfg as any)?.auth?.multiAccount;
if (!multiAccountConfig?.enabled) return false;
const providerList = multiAccountConfig?.providers ?? ["google-antigravity"];
return providerList.includes(provider);
}
/**
* Resolve API key using multi-account load balancing
*
* Returns null if multi-account is not available, signaling to use default resolver.
*/
export async function resolveWithMultiAccount(params: {
provider: string;
modelId: string;
cfg?: ClawdbotConfig;
authStore: AuthProfileStore;
}): Promise<{
apiKey: string;
profileId: string;
source: string;
mode: "oauth";
_manager: AccountManager;
_profileId: string;
} | null> {
const { provider, modelId, cfg, authStore } = params;
// Check if enabled
if (!isMultiAccountEnabled(provider, cfg)) {
return null;
}
const manager = await getAccountManager(provider, authStore, cfg);
// Check if we have multiple accounts
if (manager.getAccountCount() <= 1) {
return null;
}
// Select account
const { account, waitMs } = manager.selectAccount(modelId);
if (!account) {
if (waitMs > 0) {
const minWait = manager.getMinWaitTimeMs(modelId);
throw new Error(
`MULTI_ACCOUNT_EXHAUSTED: All accounts rate limited for ${modelId}. ` +
`Wait ${Math.ceil(minWait / 1000)}s.`
);
}
return null;
}
// Get credentials
const token = await manager.getTokenForAccount(account);
const projectId = await manager.getProjectForAccount(account);
return {
apiKey: JSON.stringify({ token, projectId }),
profileId: account.profileId,
source: `multi-account:${account.email}`,
mode: "oauth",
_manager: manager,
_profileId: account.profileId,
};
}
/**
* Report success to account manager
*/
export function reportMultiAccountSuccess(
result: { _manager?: AccountManager; _profileId?: string } | null,
modelId: string
): void {
if (!result?._manager || !result?._profileId) return;
result._manager.notifySuccess(result._profileId, modelId);
}
/**
* Report rate limit to account manager
*/
export function reportMultiAccountRateLimit(
result: { _manager?: AccountManager; _profileId?: string } | null,
modelId: string,
cooldownMs?: number
): void {
if (!result?._manager || !result?._profileId) return;
result._manager.notifyRateLimit(result._profileId, modelId, cooldownMs);
}
/**
* Report failure to account manager
*/
export function reportMultiAccountFailure(
result: { _manager?: AccountManager; _profileId?: string } | null,
modelId: string
): void {
if (!result?._manager || !result?._profileId) return;
result._manager.notifyFailure(result._profileId, modelId);
}
/**
* Report auth invalid to account manager
*/
export function reportMultiAccountAuthInvalid(
result: { _manager?: AccountManager; _profileId?: string } | null,
reason: string
): void {
if (!result?._manager || !result?._profileId) return;
result._manager.markInvalid(result._profileId, reason);
}
/**
* Get status of all multi-account managers
*/
export function getMultiAccountStatus(): Record<string, ReturnType<AccountManager["getStatus"]>> {
const status: Record<string, ReturnType<AccountManager["getStatus"]>> = {};
for (const [provider, manager] of managers) {
status[provider] = manager.getStatus();
}
return status;
}
/**
* Clear all managers (for testing)
*/
export function clearManagers(): void {
managers.clear();
}

166
src/agents/multi-account/model-auth-integration.ts Normal file
View File

@ -0,0 +1,166 @@
/**
* Multi-Account Integration for Model Auth
*
* Drop-in enhancement for resolveApiKeyForProvider that enables
* multi-account load balancing when configured.
*/
import type { AuthProfileStore } from "../auth-profiles.js";
import { AccountManager, type Account } from "./account-manager.js";
import { normalizeProviderId } from "../model-selection.js";
// Singleton managers per provider
const managers = new Map<string, AccountManager>();
// eslint-disable-next-line @typescript-eslint/no-explicit-any
type ClawdbotConfig = any;
export interface MultiAccountConfig {
enabled?: boolean;
strategy?: "hybrid" | "sticky" | "round-robin";
providers?: string[];
defaultCooldownMs?: number;
}
export function getMultiAccountConfig(cfg?: ClawdbotConfig): MultiAccountConfig | null {
const multiAccount = cfg?.auth?.multiAccount;
if (!multiAccount?.enabled) return null;
return multiAccount as MultiAccountConfig;
}
export function isMultiAccountEnabled(provider: string, cfg?: ClawdbotConfig): boolean {
const config = getMultiAccountConfig(cfg);
if (!config) return false;
const providers = config.providers ?? ["google-antigravity"];
const normalized = normalizeProviderId(provider);
return providers.some((p) => normalizeProviderId(p) === normalized);
}
export async function getOrCreateManager(
provider: string,
store: AuthProfileStore,
cfg?: ClawdbotConfig
): Promise<AccountManager> {
const normalized = normalizeProviderId(provider);
if (managers.has(normalized)) {
return managers.get(normalized)!;
}
const config = getMultiAccountConfig(cfg);
const manager = new AccountManager({
provider: normalized,
strategy: config?.strategy ?? "hybrid",
defaultCooldownMs: config?.defaultCooldownMs,
});
await manager.initialize(store, config?.strategy);
managers.set(normalized, manager);
return manager;
}
export function getManager(provider: string): AccountManager | undefined {
return managers.get(normalizeProviderId(provider));
}
export interface MultiAccountSelection {
account: Account;
manager: AccountManager;
profileId: string;
}
export async function selectAccountForModel(params: {
provider: string;
modelId: string;
cfg?: ClawdbotConfig;
store: AuthProfileStore;
}): Promise<MultiAccountSelection | null> {
const { provider, modelId, cfg, store } = params;
if (!isMultiAccountEnabled(provider, cfg)) {
return null;
}
const manager = await getOrCreateManager(provider, store, cfg);
// Need at least 2 accounts for multi-account to be useful
if (manager.getAccountCount() < 2) {
return null;
}
const { account, waitMs } = manager.selectAccount(modelId);
if (!account) {
if (waitMs > 0) {
const waitSec = Math.ceil(waitMs / 1000);
throw new Error(
`All ${manager.getAccountCount()} accounts rate-limited for ${modelId}. ` +
`Shortest wait: ${waitSec}s. Try again later.`
);
}
return null;
}
return {
account,
manager,
profileId: account.profileId,
};
}
export function notifyMultiAccountSuccess(
selection: MultiAccountSelection | null | undefined,
modelId: string
): void {
if (!selection) return;
selection.manager.notifySuccess(selection.profileId, modelId);
}
export function notifyMultiAccountRateLimit(
selection: MultiAccountSelection | null | undefined,
modelId: string,
cooldownMs?: number
): void {
if (!selection) return;
selection.manager.notifyRateLimit(selection.profileId, modelId, cooldownMs);
}
export function notifyMultiAccountFailure(
selection: MultiAccountSelection | null | undefined,
modelId: string
): void {
if (!selection) return;
selection.manager.notifyFailure(selection.profileId, modelId);
}
export function notifyMultiAccountInvalid(
selection: MultiAccountSelection | null | undefined,
reason: string
): void {
if (!selection) return;
selection.manager.markInvalid(selection.profileId, reason);
}
export function getMultiAccountStatus(): Record<
string,
ReturnType<AccountManager["getStatus"]>
> {
const status: Record<string, ReturnType<AccountManager["getStatus"]>> = {};
for (const [provider, manager] of managers) {
status[provider] = manager.getStatus();
}
return status;
}
export function clearAllManagers(): void {
managers.clear();
}
export function resetManagerState(provider: string): void {
const manager = managers.get(normalizeProviderId(provider));
if (manager) {
manager.resetAllRateLimits();
}
}

155
src/agents/multi-account/multi-account.test.ts Normal file
View File

@ -0,0 +1,155 @@
/**
* Tests for Multi-Account Module
*/
import { describe, it, expect, beforeEach } from "vitest";
import { AccountManager } from "./account-manager.js";
import { RateLimitTracker } from "./rate-limit-tracker.js";
const mockAuthStore = {
profiles: {
"google-antigravity:user1@gmail.com": {
type: "oauth",
provider: "google-antigravity",
email: "user1@gmail.com",
access: "token1",
refresh: "refresh1",
expires: Date.now() + 3600_000,
projectId: "project1",
},
"google-antigravity:user2@gmail.com": {
type: "oauth",
provider: "google-antigravity",
email: "user2@gmail.com",
access: "token2",
refresh: "refresh2",
expires: Date.now() + 3600_000,
projectId: "project2",
},
"google-antigravity:user3@gmail.com": {
type: "oauth",
provider: "google-antigravity",
email: "user3@gmail.com",
access: "token3",
refresh: "refresh3",
expires: Date.now() + 3600_000,
projectId: "project3",
},
},
};
describe("RateLimitTracker", () => {
it("initially not rate limited", () => {
const tracker = new RateLimitTracker();
expect(tracker.isRateLimited("p1", "m1")).toBe(false);
});
it("marks rate limited", () => {
const tracker = new RateLimitTracker();
tracker.markRateLimited("p1", "m1", 5000);
expect(tracker.isRateLimited("p1", "m1")).toBe(true);
});
it("clears rate limit", () => {
const tracker = new RateLimitTracker();
tracker.markRateLimited("p1", "m1", 5000);
tracker.clearRateLimit("p1", "m1");
expect(tracker.isRateLimited("p1", "m1")).toBe(false);
});
it("checks all rate limited", () => {
const tracker = new RateLimitTracker();
const profiles = ["p1", "p2", "p3"];
expect(tracker.areAllRateLimited(profiles, "m1")).toBe(false);
tracker.markRateLimited("p1", "m1", 5000);
tracker.markRateLimited("p2", "m1", 5000);
tracker.markRateLimited("p3", "m1", 5000);
expect(tracker.areAllRateLimited(profiles, "m1")).toBe(true);
});
});
describe("AccountManager", () => {
let manager: AccountManager;
beforeEach(async () => {
manager = new AccountManager({ provider: "google-antigravity" });
await manager.initialize(mockAuthStore);
});
it("initializes with auth store", () => {
expect(manager.getAccountCount()).toBe(3);
});
it("selects account", () => {
const { account, waitMs } = manager.selectAccount("claude-opus-4-5");
expect(account).not.toBeNull();
expect(waitMs).toBe(0);
});
it("handles rate limit failover", () => {
const { account: first } = manager.selectAccount("m1");
manager.markRateLimited(first!.profileId, "m1", 60_000);
const { account: second } = manager.selectAccount("m1");
expect(second!.profileId).not.toBe(first!.profileId);
});
it("returns wait time when all limited", () => {
for (const acc of manager.getAllAccounts()) {
manager.markRateLimited(acc.profileId, "m1", 30_000);
}
expect(manager.isAllRateLimited("m1")).toBe(true);
const { account, waitMs } = manager.selectAccount("m1");
expect(account).toBeNull();
expect(waitMs).toBeGreaterThan(0);
});
it("updates health on success", () => {
const { account } = manager.selectAccount("m1");
const before = manager.healthScorer.getScore(account!.profileId);
manager.notifySuccess(account!.profileId, "m1");
const after = manager.healthScorer.getScore(account!.profileId);
expect(after).toBeGreaterThanOrEqual(before);
});
it("marks invalid accounts", () => {
const profiles = manager.getProfilesForProvider();
const first = profiles[0];
manager.markInvalid(first, "Token revoked");
const updated = manager.getProfilesForProvider();
expect(updated).not.toContain(first);
});
it("sticky strategy keeps same account", async () => {
const stickyManager = new AccountManager({
provider: "google-antigravity",
strategy: "sticky",
});
await stickyManager.initialize(mockAuthStore);
const { account: a1 } = stickyManager.selectAccount("m1");
const { account: a2 } = stickyManager.selectAccount("m1");
expect(a1!.profileId).toBe(a2!.profileId);
});
it("round-robin cycles accounts", async () => {
const rrManager = new AccountManager({
provider: "google-antigravity",
strategy: "round-robin",
});
await rrManager.initialize(mockAuthStore);
const seen = new Set<string>();
for (let i = 0; i < 6; i++) {
const { account } = rrManager.selectAccount("m1");
seen.add(account!.profileId);
}
expect(seen.size).toBe(3);
});
});

112
src/agents/multi-account/profile-order.ts Normal file
View File

@ -0,0 +1,112 @@
/**
* Multi-Account Profile Order Integration
*
* Wraps resolveAuthProfileOrder to use intelligent multi-account selection
* when enabled. Falls back to default behavior otherwise.
*/
import type { ClawdbotConfig } from "../../config/config.js";
import type { AuthProfileStore } from "../auth-profiles/types.js";
import { resolveAuthProfileOrder as baseResolveAuthProfileOrder } from "../auth-profiles/order.js";
import { isMultiAccountEnabled, getOrCreateManager } from "./model-auth-integration.js";
/**
* Enhanced profile order resolution with multi-account intelligence.
*
* When multi-account is enabled for the provider:
* - Uses health scores, quota tracking, and strategy-based selection
* - Falls back to base order if multi-account unavailable
*
* When multi-account is disabled:
* - Uses standard resolveAuthProfileOrder (round-robin + cooldown)
*/
export async function resolveProfileOrderWithMultiAccount(params: {
cfg?: ClawdbotConfig;
store: AuthProfileStore;
provider: string;
modelId?: string;
preferredProfile?: string;
}): Promise<string[]> {
const { cfg, store, provider, modelId, preferredProfile } = params;
// Get base order first
const baseOrder = baseResolveAuthProfileOrder({
cfg,
store,
provider,
preferredProfile,
});
// If multi-account not enabled or not enough profiles, use base order
if (!isMultiAccountEnabled(provider, cfg) || baseOrder.length < 2) {
return baseOrder;
}
// If no modelId provided, can't do intelligent selection
if (!modelId) {
return baseOrder;
}
try {
const manager = await getOrCreateManager(provider, store, cfg);
// Get intelligent ordering from manager
const healthSorted = manager.healthScorer.getSortedByHealth(baseOrder);
// Filter out rate-limited profiles and sort by health
const available: string[] = [];
const rateLimited: Array<{ profileId: string; score: number }> = [];
for (const { profileId, score } of healthSorted) {
if (manager.rateLimitTracker.isRateLimited(profileId, modelId)) {
rateLimited.push({ profileId, score });
} else {
available.push(profileId);
}
}
// Sort rate-limited by soonest cooldown expiry
const rateLimitedSorted = rateLimited
.map(({ profileId }) => ({
profileId,
cooldownRemaining: manager.rateLimitTracker.getCooldownRemaining(profileId, modelId),
}))
.sort((a, b) => a.cooldownRemaining - b.cooldownRemaining)
.map(({ profileId }) => profileId);
const result = [...available, ...rateLimitedSorted];
// Ensure preferredProfile is first if specified and in list
if (preferredProfile && result.includes(preferredProfile)) {
return [preferredProfile, ...result.filter((p) => p !== preferredProfile)];
}
return result;
} catch {
// Fall back to base order on any error
return baseOrder;
}
}
/**
* Synchronous version that uses cached manager state.
* Falls back to base order if manager not initialized.
*/
export function resolveProfileOrderWithMultiAccountSync(params: {
cfg?: ClawdbotConfig;
store: AuthProfileStore;
provider: string;
modelId?: string;
preferredProfile?: string;
}): string[] {
const { cfg, store, provider, preferredProfile } = params;
// Always return base order for sync version
// Multi-account benefits come from the manager callbacks in run.ts
return baseResolveAuthProfileOrder({
cfg,
store,
provider,
preferredProfile,
});
}

136
src/agents/multi-account/quota-tracker.ts Normal file
View File

@ -0,0 +1,136 @@
/**
* Quota Tracker
*
* Tracks quota information per account:
* - Subscription tier (free/pro/ultra)
* - Per-model remaining quota
* - Reset times
*/
import { TIER_WEIGHTS } from "./constants.js";
export interface ModelQuota {
remainingFraction: number | null;
resetTime: number | null;
}
export interface AccountQuota {
tier: "free" | "pro" | "ultra" | "unknown";
projectId: string | null;
models: Record<string, ModelQuota>;
lastChecked: number;
}
export class QuotaTracker {
private quotas = new Map<string, AccountQuota>();
private cacheTtlMs = 5 * 60 * 1000;
setCacheTtl(ms: number): void {
this.cacheTtlMs = ms;
}
getQuota(profileId: string): AccountQuota | null {
const quota = this.quotas.get(profileId);
if (!quota) return null;
if (Date.now() - quota.lastChecked > this.cacheTtlMs) {
return null;
}
return quota;
}
updateQuota(profileId: string, data: Partial<AccountQuota>): void {
const existing = this.quotas.get(profileId) ?? {
tier: "unknown" as const,
projectId: null,
models: {},
lastChecked: 0,
};
this.quotas.set(profileId, {
...existing,
...data,
lastChecked: Date.now(),
});
}
updateModelQuota(
profileId: string,
modelId: string,
quota: ModelQuota
): void {
const existing = this.getQuota(profileId) ?? {
tier: "unknown" as const,
projectId: null,
models: {},
lastChecked: 0,
};
existing.models[modelId] = quota;
existing.lastChecked = Date.now();
this.quotas.set(profileId, existing);
}
getTierWeight(profileId: string): number {
const quota = this.quotas.get(profileId);
const tier = quota?.tier ?? "unknown";
return TIER_WEIGHTS[tier] ?? TIER_WEIGHTS.unknown;
}
getModelRemaining(profileId: string, modelId: string): number | null {
const quota = this.getQuota(profileId);
if (!quota) return null;
return quota.models[modelId]?.remainingFraction ?? null;
}
isQuotaExhausted(profileId: string, modelId: string): boolean {
const remaining = this.getModelRemaining(profileId, modelId);
return remaining === 0;
}
isQuotaCritical(profileId: string, modelId: string): boolean {
const remaining = this.getModelRemaining(profileId, modelId);
if (remaining === null) return false;
return remaining < 0.05;
}
getSortedByQuota(profileIds: string[], modelId: string): string[] {
return [...profileIds].sort((a, b) => {
const quotaA = this.getModelRemaining(a, modelId) ?? 0.5;
const quotaB = this.getModelRemaining(b, modelId) ?? 0.5;
if (quotaA === 0 && quotaB !== 0) return 1;
if (quotaB === 0 && quotaA !== 0) return -1;
return quotaB - quotaA;
});
}
getSortedByTier(profileIds: string[]): string[] {
return [...profileIds].sort((a, b) => {
return this.getTierWeight(b) - this.getTierWeight(a);
});
}
markExhausted(profileId: string, modelId: string, resetTime: number): void {
this.updateModelQuota(profileId, modelId, {
remainingFraction: 0,
resetTime,
});
}
toJSON(): { quotas: Record<string, AccountQuota> } {
return {
quotas: Object.fromEntries(this.quotas),
};
}
fromJSON(data: { quotas?: Record<string, AccountQuota> }): void {
if (data?.quotas) {
for (const [key, value] of Object.entries(data.quotas)) {
this.quotas.set(key, value);
}
}
}
}

238
src/agents/multi-account/rate-limit-tracker.ts Normal file
View File

@ -0,0 +1,238 @@
/**
* Rate Limit Tracker
*
* Tracks per-account, per-model rate limits with:
* - Cooldown periods with expiry
* - Exponential backoff state
* - Deduplication to prevent thundering herd
* - Consecutive failure tracking
*/
import {
DEFAULT_COOLDOWN_MS,
RATE_LIMIT_DEDUP_WINDOW_MS,
RATE_LIMIT_STATE_RESET_MS,
FIRST_RETRY_DELAY_MS,
MAX_CONSECUTIVE_FAILURES,
EXTENDED_COOLDOWN_MS,
MIN_BACKOFF_MS,
} from "./constants.js";
export interface RateLimitState {
cooldownUntil: number;
consecutive429: number;
lastAt: number;
consecutiveFailures: number;
}
export interface BackoffResult {
attempt: number;
delayMs: number;
isDuplicate: boolean;
}
export class RateLimitTracker {
private stateMap = new Map<string, RateLimitState>();
private failureCount = new Map<string, number>();
private getKey(profileId: string, modelId: string): string {
return `${profileId}:${modelId}`;
}
isRateLimited(profileId: string, modelId: string): boolean {
const key = this.getKey(profileId, modelId);
const state = this.stateMap.get(key);
if (!state) return false;
return Date.now() < state.cooldownUntil;
}
getCooldownRemaining(profileId: string, modelId: string): number {
const key = this.getKey(profileId, modelId);
const state = this.stateMap.get(key);
if (!state) return 0;
return Math.max(0, state.cooldownUntil - Date.now());
}
getSoonestReset(profileId: string): number | null {
const now = Date.now();
let soonest: number | null = null;
for (const [key, state] of this.stateMap.entries()) {
if (!key.startsWith(profileId + ":")) continue;
if (state.cooldownUntil <= now) continue;
if (soonest === null || state.cooldownUntil < soonest) {
soonest = state.cooldownUntil;
}
}
return soonest;
}
markRateLimited(
profileId: string,
modelId: string,
cooldownMs: number = DEFAULT_COOLDOWN_MS
): BackoffResult {
const key = this.getKey(profileId, modelId);
const now = Date.now();
const previous = this.stateMap.get(key);
// Check deduplication window
if (previous && now - previous.lastAt < RATE_LIMIT_DEDUP_WINDOW_MS) {
const backoffDelay = Math.min(
FIRST_RETRY_DELAY_MS * Math.pow(2, previous.consecutive429 - 1),
60_000
);
return {
attempt: previous.consecutive429,
delayMs: Math.max(cooldownMs, backoffDelay),
isDuplicate: true,
};
}
// Determine attempt number
const attempt =
previous && now - previous.lastAt < RATE_LIMIT_STATE_RESET_MS
? previous.consecutive429 + 1
: 1;
// Calculate exponential backoff
const backoffDelay = Math.min(
FIRST_RETRY_DELAY_MS * Math.pow(2, attempt - 1),
60_000
);
const effectiveCooldown = Math.max(cooldownMs, backoffDelay, MIN_BACKOFF_MS);
// Update state
this.stateMap.set(key, {
cooldownUntil: now + effectiveCooldown,
consecutive429: attempt,
lastAt: now,
consecutiveFailures: (previous?.consecutiveFailures ?? 0) + 1,
});
this.incrementFailures(profileId);
return {
attempt,
delayMs: effectiveCooldown,
isDuplicate: false,
};
}
clearRateLimit(profileId: string, modelId: string): void {
const key = this.getKey(profileId, modelId);
this.stateMap.delete(key);
this.failureCount.set(profileId, 0);
}
clearExpired(): void {
const now = Date.now();
for (const [key, state] of this.stateMap.entries()) {
if (state.cooldownUntil <= now) {
this.stateMap.delete(key);
}
}
}
getFailureCount(profileId: string): number {
return this.failureCount.get(profileId) ?? 0;
}
incrementFailures(profileId: string): number {
const current = this.failureCount.get(profileId) ?? 0;
const newCount = current + 1;
this.failureCount.set(profileId, newCount);
return newCount;
}
shouldExtendCooldown(profileId: string): boolean {
return this.getFailureCount(profileId) >= MAX_CONSECUTIVE_FAILURES;
}
applyExtendedCooldown(profileId: string): void {
const now = Date.now();
const cooldownUntil = now + EXTENDED_COOLDOWN_MS;
for (const [key, state] of this.stateMap.entries()) {
if (key.startsWith(profileId + ":")) {
state.cooldownUntil = Math.max(state.cooldownUntil, cooldownUntil);
}
}
const defaultKey = `${profileId}:*`;
if (!this.stateMap.has(defaultKey)) {
this.stateMap.set(defaultKey, {
cooldownUntil,
consecutive429: 0,
lastAt: now,
consecutiveFailures: this.getFailureCount(profileId),
});
}
}
getRateLimitedProfiles(modelId: string): string[] {
const now = Date.now();
const limited: string[] = [];
for (const [key, state] of this.stateMap.entries()) {
if (!key.endsWith(":" + modelId) && !key.endsWith(":*")) continue;
if (state.cooldownUntil <= now) continue;
const profileId = key.split(":").slice(0, -1).join(":");
if (!limited.includes(profileId)) {
limited.push(profileId);
}
}
return limited;
}
areAllRateLimited(profileIds: string[], modelId: string): boolean {
if (profileIds.length === 0) return false;
return profileIds.every((id) => this.isRateLimited(id, modelId));
}
getMinWaitTime(profileIds: string[], modelId: string): number {
let minWait = Infinity;
for (const profileId of profileIds) {
const remaining = this.getCooldownRemaining(profileId, modelId);
if (remaining > 0 && remaining < minWait) {
minWait = remaining;
}
}
return minWait === Infinity ? 0 : minWait;
}
toJSON(): {
rateState: Record<string, RateLimitState>;
failureCounts: Record<string, number>;
} {
const state: Record<string, RateLimitState> = {};
for (const [key, value] of this.stateMap.entries()) {
state[key] = value;
}
return {
rateState: state,
failureCounts: Object.fromEntries(this.failureCount),
};
}
fromJSON(data: {
rateState?: Record<string, RateLimitState>;
failureCounts?: Record<string, number>;
}): void {
if (data?.rateState) {
for (const [key, value] of Object.entries(data.rateState)) {
this.stateMap.set(key, value);
}
}
if (data?.failureCounts) {
for (const [key, value] of Object.entries(data.failureCounts)) {
this.failureCount.set(key, value);
}
}
this.clearExpired();
}
}

188
src/agents/multi-account/strategies.ts Normal file
View File

@ -0,0 +1,188 @@
/**
* Account Selection Strategies
*/
import type { AccountManager } from "./account-manager.js";
export interface SelectionResult {
account: { profileId: string; email: string } | null;
waitMs: number;
}
export interface Strategy {
init(manager: AccountManager): void;
select(modelId: string): SelectionResult;
notifySuccess(profileId: string, modelId: string): void;
notifyRateLimit(profileId: string, modelId: string): void;
getLabel(): string;
}
export class HybridStrategy implements Strategy {
private manager!: AccountManager;
private lastUsed = new Map<string, number>();
init(manager: AccountManager): void {
this.manager = manager;
}
select(modelId: string): SelectionResult {
const { rateLimitTracker, healthScorer, quotaTracker } = this.manager;
const profiles = this.manager.getProfilesForProvider();
const available = profiles.filter(
(id) => !rateLimitTracker.isRateLimited(id, modelId)
);
if (available.length === 0) {
const minWait = rateLimitTracker.getMinWaitTime(profiles, modelId);
return { account: null, waitMs: minWait };
}
const scored = available.map((profileId) => {
const healthScore = healthScorer.getScore(profileId);
const quotaRemaining = quotaTracker.getModelRemaining(profileId, modelId);
const tierWeight = quotaTracker.getTierWeight(profileId);
const lastUsedTime = this.lastUsed.get(profileId) ?? 0;
const lruScore = Math.max(0, 100 - (Date.now() - lastUsedTime) / 60_000);
const quotaPenalty = quotaTracker.isQuotaCritical(profileId, modelId) ? 30 : 0;
const score =
healthScore * 0.4 +
(quotaRemaining ?? 0.5) * 100 * 0.3 +
tierWeight * 20 * 0.2 +
(100 - lruScore) * 0.1 -
quotaPenalty;
return { profileId, score };
});
scored.sort((a, b) => b.score - a.score);
const selected = scored[0];
if (!selected) return { account: null, waitMs: 0 };
this.lastUsed.set(selected.profileId, Date.now());
return {
account: this.manager.getAccountByProfileId(selected.profileId),
waitMs: 0,
};
}
notifySuccess(profileId: string): void {
this.lastUsed.set(profileId, Date.now());
}
notifyRateLimit(): void {}
getLabel(): string {
return "Hybrid (health + quota + LRU)";
}
}
export class StickyStrategy implements Strategy {
private manager!: AccountManager;
private sticky = new Map<string, string>();
init(manager: AccountManager): void {
this.manager = manager;
}
select(modelId: string): SelectionResult {
const { rateLimitTracker } = this.manager;
const profiles = this.manager.getProfilesForProvider();
const current = this.sticky.get(modelId);
if (current && !rateLimitTracker.isRateLimited(current, modelId)) {
return {
account: this.manager.getAccountByProfileId(current),
waitMs: 0,
};
}
const available = profiles.filter(
(id) => !rateLimitTracker.isRateLimited(id, modelId)
);
if (available.length === 0) {
const minWait = rateLimitTracker.getMinWaitTime(profiles, modelId);
return { account: null, waitMs: minWait };
}
const selected = available[0];
this.sticky.set(modelId, selected);
return {
account: this.manager.getAccountByProfileId(selected),
waitMs: 0,
};
}
notifySuccess(): void {}
notifyRateLimit(profileId: string, modelId: string): void {
if (this.sticky.get(modelId) === profileId) {
this.sticky.delete(modelId);
}
}
getLabel(): string {
return "Sticky (cache-optimized)";
}
}
export class RoundRobinStrategy implements Strategy {
private manager!: AccountManager;
private indices = new Map<string, number>();
init(manager: AccountManager): void {
this.manager = manager;
}
select(modelId: string): SelectionResult {
const { rateLimitTracker } = this.manager;
const profiles = this.manager.getProfilesForProvider();
if (profiles.length === 0) return { account: null, waitMs: 0 };
const startIndex = this.indices.get(modelId) ?? 0;
let attempts = 0;
while (attempts < profiles.length) {
const index = (startIndex + attempts) % profiles.length;
const profileId = profiles[index];
if (!rateLimitTracker.isRateLimited(profileId, modelId)) {
this.indices.set(modelId, (index + 1) % profiles.length);
return {
account: this.manager.getAccountByProfileId(profileId),
waitMs: 0,
};
}
attempts++;
}
const minWait = rateLimitTracker.getMinWaitTime(profiles, modelId);
return { account: null, waitMs: minWait };
}
notifySuccess(): void {}
notifyRateLimit(): void {}
getLabel(): string {
return "Round-Robin (even distribution)";
}
}
export const STRATEGY_NAMES = ["hybrid", "sticky", "round-robin"] as const;
export type StrategyName = (typeof STRATEGY_NAMES)[number];
export function createStrategy(name?: string): Strategy {
switch (name?.toLowerCase()) {
case "sticky":
return new StickyStrategy();
case "round-robin":
case "roundrobin":
return new RoundRobinStrategy();
case "hybrid":
default:
return new HybridStrategy();
}
}

91
src/agents/pi-embedded-runner/run.ts
View File

@ -41,6 +41,11 @@ import {
type FailoverReason,
} from "../pi-embedded-helpers.js";
import { normalizeUsage, type UsageLike } from "../usage.js";
import {
isMultiAccountEnabled,
getManager,
getOrCreateManager,
} from "../multi-account/index.js";
import { compactEmbeddedPiSessionDirect } from "./compact.js";
import { resolveGlobalLane, resolveSessionLane } from "./lanes.js";
@ -136,6 +141,12 @@ export async function runEmbeddedPiAgent(
}
const authStore = ensureAuthProfileStore(agentDir, { allowKeychainPrompt: false });
// Initialize multi-account manager if enabled for this provider
if (isMultiAccountEnabled(provider, params.config)) {
await getOrCreateManager(provider, authStore, params.config);
}
const preferredProfileId = params.authProfileId?.trim();
let lockedProfileId = params.authProfileIdSource === "user" ? preferredProfileId : undefined;
if (lockedProfileId) {
@ -147,12 +158,57 @@ export async function runEmbeddedPiAgent(
lockedProfileId = undefined;
}
}
const profileOrder = resolveAuthProfileOrder({
cfg: params.config,
store: authStore,
provider,
preferredProfile: preferredProfileId,
});
// Use multi-account intelligent ordering when enabled
let profileOrder: string[];
const manager = getManager(provider);
if (manager && manager.getAccountCount() >= 2) {
// Get base order
const baseOrder = resolveAuthProfileOrder({
cfg: params.config,
store: authStore,
provider,
preferredProfile: preferredProfileId,
});
// Re-order based on multi-account intelligence (health + rate limits)
const available: string[] = [];
const rateLimited: Array<{ profileId: string; cooldown: number }> = [];
for (const profileId of baseOrder) {
if (manager.rateLimitTracker.isRateLimited(profileId, modelId)) {
rateLimited.push({
profileId,
cooldown: manager.rateLimitTracker.getCooldownRemaining(profileId, modelId),
});
} else {
available.push(profileId);
}
}
// Sort available by health score (highest first)
const healthSorted = manager.healthScorer.getSortedByHealth(available);
const sortedAvailable = healthSorted.map(h => h.profileId);
// Sort rate-limited by soonest cooldown expiry
rateLimited.sort((a, b) => a.cooldown - b.cooldown);
const sortedRateLimited = rateLimited.map(r => r.profileId);
profileOrder = [...sortedAvailable, ...sortedRateLimited];
// Ensure preferred profile is first if specified
if (preferredProfileId && profileOrder.includes(preferredProfileId)) {
profileOrder = [preferredProfileId, ...profileOrder.filter(p => p !== preferredProfileId)];
}
} else {
// Fall back to standard ordering
profileOrder = resolveAuthProfileOrder({
cfg: params.config,
store: authStore,
provider,
preferredProfile: preferredProfileId,
});
}
if (lockedProfileId && !profileOrder.includes(lockedProfileId)) {
throw new Error(`Auth profile "${lockedProfileId}" is not configured for ${provider}.`);
}
@ -449,6 +505,15 @@ export async function runEmbeddedPiAgent(
cfg: params.config,
agentDir: params.agentDir,
});
// Multi-account: notify rate limit or failure
const manager = getManager(provider);
if (manager) {
if (promptFailoverReason === "billing") {
manager.notifyRateLimit(lastProfileId, modelId);
} else {
manager.notifyFailure(lastProfileId, modelId);
}
}
}
if (
isFailoverErrorMessage(errorText) &&
@ -536,6 +601,15 @@ export async function runEmbeddedPiAgent(
cfg: params.config,
agentDir: params.agentDir,
});
// Multi-account: notify rate limit or failure
const manager = getManager(provider);
if (manager) {
if (rateLimitFailure || timedOut) {
manager.notifyRateLimit(lastProfileId, modelId);
} else {
manager.notifyFailure(lastProfileId, modelId);
}
}
if (timedOut && !isProbeSession) {
log.warn(
`Profile ${lastProfileId} timed out (possible rate limit). Trying next account...`,
@ -617,6 +691,11 @@ export async function runEmbeddedPiAgent(
profileId: lastProfileId,
agentDir: params.agentDir,
});
// Multi-account: notify success
const manager = getManager(provider);
if (manager) {
manager.notifySuccess(lastProfileId, modelId);
}
}
return {
payloads: payloads.length ? payloads : undefined,

57
src/cli/accounts-cli.ts Normal file
View File

@ -0,0 +1,57 @@
/**
* Accounts CLI - Multi-account management commands
*/
import type { Command } from "commander";
import { defaultRuntime } from "../runtime.js";
import { formatDocsLink } from "../terminal/links.js";
import { theme } from "../terminal/theme.js";
import { runCommandWithRuntime } from "./cli-utils.js";
import {
accountsListCommand,
accountsStatusCommand,
accountsResetCommand,
} from "../commands/accounts/index.js";
function runAccountsCommand(action: () => Promise<void>) {
return runCommandWithRuntime(defaultRuntime, action);
}
export function registerAccountsCli(program: Command) {
const accounts = program
.command("accounts")
.description("Manage multi-account load balancing")
.addHelpText(
"after",
() =>
`\n${theme.muted("Multi-account distributes requests across multiple OAuth accounts for higher throughput.")}\n`
);
accounts
.command("list")
.description("List all accounts for a provider")
.option("-p, --provider <provider>", "Provider name", "google-antigravity")
.option("--json", "Output as JSON")
.action((options) =>
runAccountsCommand(() => accountsListCommand(defaultRuntime, options))
);
accounts
.command("status")
.description("Show multi-account status (health scores, rate limits)")
.option("-p, --provider <provider>", "Provider name", "google-antigravity")
.option("--json", "Output as JSON")
.action((options) =>
runAccountsCommand(() => accountsStatusCommand(defaultRuntime, options))
);
accounts
.command("reset")
.description("Reset rate limits for all accounts")
.option("-p, --provider <provider>", "Provider name", "google-antigravity")
.action((options) =>
runAccountsCommand(() => accountsResetCommand(defaultRuntime, options))
);
return accounts;
}

8
src/cli/program/register.subclis.ts
View File

@ -190,6 +190,14 @@ const entries: SubCliEntry[] = [
mod.registerChannelsCli(program);
},
},
{
name: "accounts",
description: "Multi-account management",
register: async (program) => {
const mod = await import("../accounts-cli.js");
mod.registerAccountsCli(program);
},
},
{
name: "directory",
description: "Directory commands",

8
src/commands/accounts/index.ts Normal file
View File

@ -0,0 +1,8 @@
export type { AccountsStatusOptions } from "./status.js";
export { accountsStatusCommand } from "./status.js";
export type { AccountsListOptions } from "./list.js";
export { accountsListCommand } from "./list.js";
export type { AccountsResetOptions } from "./reset.js";
export { accountsResetCommand } from "./reset.js";

60
src/commands/accounts/list.ts Normal file
View File

@ -0,0 +1,60 @@
/**
* Multi-Account List Command
*/
import type { RuntimeEnv } from "../../runtime.js";
import { ensureAuthProfileStore, listProfilesForProvider } from "../../agents/auth-profiles.js";
import { resolveClawdbotAgentDir } from "../../agents/agent-paths.js";
import { loadConfig } from "../../config/config.js";
import { isMultiAccountEnabled } from "../../agents/multi-account/index.js";
export interface AccountsListOptions {
provider?: string;
json?: boolean;
}
export async function accountsListCommand(
runtime: RuntimeEnv,
options: AccountsListOptions
): Promise<void> {
const cfg = loadConfig();
const agentDir = resolveClawdbotAgentDir();
const authStore = ensureAuthProfileStore(agentDir);
const provider = options.provider ?? "google-antigravity";
const profiles = listProfilesForProvider(authStore, provider);
if (profiles.length === 0) {
runtime.log(`No accounts found for provider: ${provider}`);
return;
}
const multiAccountEnabled = isMultiAccountEnabled(provider, cfg);
if (options.json) {
const data = profiles.map((profileId) => {
const profile = authStore.profiles[profileId];
return {
profileId,
email: profile?.email ?? profileId,
type: profile?.type,
provider: profile?.provider,
};
});
runtime.log(JSON.stringify(data, null, 2));
return;
}
runtime.log(`\n📋 Accounts for ${provider}`);
runtime.log(`${"─".repeat(40)}`);
runtime.log(`Total: ${profiles.length}`);
runtime.log(`Multi-account: ${multiAccountEnabled ? "✅ Enabled" : "❌ Disabled"}`);
runtime.log("");
for (const profileId of profiles) {
const profile = authStore.profiles[profileId];
const email = profile?.email ?? profileId;
const type = profile?.type ?? "unknown";
runtime.log(` • ${email} (${type})`);
}
}

28
src/commands/accounts/reset.ts Normal file
View File

@ -0,0 +1,28 @@
/**
* Multi-Account Reset Command
*/
import type { RuntimeEnv } from "../../runtime.js";
import { resetManagerState, getManager } from "../../agents/multi-account/index.js";
export interface AccountsResetOptions {
provider?: string;
}
export async function accountsResetCommand(
runtime: RuntimeEnv,
options: AccountsResetOptions
): Promise<void> {
const provider = options.provider ?? "google-antigravity";
const manager = getManager(provider);
if (!manager) {
runtime.log(`No active multi-account manager for ${provider}.`);
runtime.log(`The manager is created when the first request is made.`);
return;
}
resetManagerState(provider);
runtime.log(`✅ Reset rate limits for ${provider}`);
runtime.log(`All accounts are now available for selection.`);
}

80
src/commands/accounts/status.ts Normal file
View File

@ -0,0 +1,80 @@
/**
* Multi-Account Status Command
*/
import type { RuntimeEnv } from "../../runtime.js";
import { ensureAuthProfileStore } from "../../agents/auth-profiles.js";
import { resolveClawdbotAgentDir } from "../../agents/agent-paths.js";
import {
getMultiAccountStatus,
getOrCreateManager,
isMultiAccountEnabled,
} from "../../agents/multi-account/index.js";
import { loadConfig } from "../../config/config.js";
export interface AccountsStatusOptions {
provider?: string;
json?: boolean;
}
export async function accountsStatusCommand(
runtime: RuntimeEnv,
options: AccountsStatusOptions
): Promise<void> {
const cfg = loadConfig();
const agentDir = resolveClawdbotAgentDir();
const authStore = ensureAuthProfileStore(agentDir);
const provider = options.provider ?? "google-antigravity";
if (!isMultiAccountEnabled(provider, cfg)) {
runtime.log(
`Multi-account is not enabled for ${provider}.\n` +
`Enable it in config:\n` +
` auth:\n` +
` multiAccount:\n` +
` enabled: true\n` +
` providers:\n` +
` - ${provider}`
);
return;
}
// Initialize manager to get status
await getOrCreateManager(provider, authStore, cfg);
const status = getMultiAccountStatus();
const providerStatus = status[provider];
if (!providerStatus) {
runtime.log(`No multi-account data for ${provider}.`);
return;
}
if (options.json) {
runtime.log(JSON.stringify(providerStatus, null, 2));
return;
}
// Pretty print
runtime.log(`\n📊 Multi-Account Status: ${provider}`);
runtime.log(`${"─".repeat(40)}`);
runtime.log(`Accounts: ${providerStatus.summary}`);
runtime.log(` Available: ${providerStatus.available}`);
runtime.log(` Rate-limited: ${providerStatus.rateLimited}`);
runtime.log(` Invalid: ${providerStatus.invalid}`);
runtime.log("");
if (providerStatus.accounts.length > 0) {
runtime.log("Account Details:");
for (const acc of providerStatus.accounts) {
const status = acc.isInvalid
? `❌ Invalid (${acc.invalidReason})`
: `✅ Health: ${acc.healthScore}`;
const lastUsed = acc.lastUsed
? new Date(acc.lastUsed).toLocaleTimeString()
: "never";
runtime.log(` • ${acc.email}`);
runtime.log(` Status: ${status}`);
runtime.log(` Last used: ${lastUsed}`);
}
}
}

18
src/config/types.auth.ts
View File

@ -10,6 +10,22 @@ export type AuthProfileConfig = {
email?: string;
};
export type MultiAccountConfig = {
/** Enable multi-account load balancing. Default: false */
enabled?: boolean;
/**
* Selection strategy:
* - hybrid: health + quota + LRU weighted scoring (default)
* - sticky: prefer same account per model (cache-friendly)
* - round-robin: even distribution across accounts
*/
strategy?: "hybrid" | "sticky" | "round-robin";
/** Providers to enable multi-account for. Default: ["google-antigravity"] */
providers?: string[];
/** Default cooldown when rate-limited (ms). Default: 10000 */
defaultCooldownMs?: number;
};
export type AuthConfig = {
profiles?: Record<string, AuthProfileConfig>;
order?: Record<string, string[]>;
@ -26,4 +42,6 @@ export type AuthConfig = {
*/
failureWindowHours?: number;
};
/** Multi-account load balancing configuration */
multiAccount?: MultiAccountConfig;
};

13
src/config/zod-schema.ts
View File

@ -204,6 +204,19 @@ export const ClawdbotSchema = z
})
.strict()
.optional(),
multiAccount: z
.object({
enabled: z.boolean().optional(),
strategy: z.union([
z.literal("hybrid"),
z.literal("sticky"),
z.literal("round-robin"),
]).optional(),
providers: z.array(z.string()).optional(),
defaultCooldownMs: z.number().positive().optional(),
})
.strict()
.optional(),
})
.strict()
.optional(),