diff --git a/src/agents/multi-account/account-manager.ts b/src/agents/multi-account/account-manager.ts new file mode 100644 index 000000000..40fabe3cb --- /dev/null +++ b/src/agents/multi-account/account-manager.ts @@ -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; +} + +export class AccountManager { + private provider: string; + private authStore: AuthProfileStore | null = null; + private accounts = new Map(); + private strategy: Strategy; + private defaultCooldownMs: number; + private invalidProfiles = new Set(); + private tokenCache = new Map(); + private projectCache = new Map(); + + 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.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 { + 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 { + 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): 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); + } +} diff --git a/src/agents/multi-account/constants.ts b/src/agents/multi-account/constants.ts new file mode 100644 index 000000000..46c96de10 --- /dev/null +++ b/src/agents/multi-account/constants.ts @@ -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 = { + ultra: 3, + pro: 2, + free: 1, + unknown: 1, +}; diff --git a/src/agents/multi-account/health-scorer.ts b/src/agents/multi-account/health-scorer.ts new file mode 100644 index 000000000..981346413 --- /dev/null +++ b/src/agents/multi-account/health-scorer.ts @@ -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(); + private lastUpdate = new Map(); + + 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; + lastUpdate: Record; + } { + return { + scores: Object.fromEntries(this.scores), + lastUpdate: Object.fromEntries(this.lastUpdate), + }; + } + + fromJSON(data: { + scores?: Record; + lastUpdate?: Record; + }): 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); + } + } + } +} diff --git a/src/agents/multi-account/index.ts b/src/agents/multi-account/index.ts new file mode 100644 index 000000000..95c5ca09d --- /dev/null +++ b/src/agents/multi-account/index.ts @@ -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"; diff --git a/src/agents/multi-account/integration.ts b/src/agents/multi-account/integration.ts new file mode 100644 index 000000000..26267fad2 --- /dev/null +++ b/src/agents/multi-account/integration.ts @@ -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(); + +/** + * Get or create an AccountManager for a provider + */ +export async function getAccountManager( + provider: string, + authStore: AuthProfileStore, + cfg?: ClawdbotConfig +): Promise { + 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> { + const status: Record> = {}; + for (const [provider, manager] of managers) { + status[provider] = manager.getStatus(); + } + return status; +} + +/** + * Clear all managers (for testing) + */ +export function clearManagers(): void { + managers.clear(); +} diff --git a/src/agents/multi-account/model-auth-integration.ts b/src/agents/multi-account/model-auth-integration.ts new file mode 100644 index 000000000..59111c6b9 --- /dev/null +++ b/src/agents/multi-account/model-auth-integration.ts @@ -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(); + +// 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 { + 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 { + 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 +> { + const status: Record> = {}; + 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(); + } +} diff --git a/src/agents/multi-account/multi-account.test.ts b/src/agents/multi-account/multi-account.test.ts new file mode 100644 index 000000000..0d397f752 --- /dev/null +++ b/src/agents/multi-account/multi-account.test.ts @@ -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(); + for (let i = 0; i < 6; i++) { + const { account } = rrManager.selectAccount("m1"); + seen.add(account!.profileId); + } + expect(seen.size).toBe(3); + }); +}); diff --git a/src/agents/multi-account/profile-order.ts b/src/agents/multi-account/profile-order.ts new file mode 100644 index 000000000..3b1c7b1cc --- /dev/null +++ b/src/agents/multi-account/profile-order.ts @@ -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 { + 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, + }); +} diff --git a/src/agents/multi-account/quota-tracker.ts b/src/agents/multi-account/quota-tracker.ts new file mode 100644 index 000000000..aeaa9d545 --- /dev/null +++ b/src/agents/multi-account/quota-tracker.ts @@ -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; + lastChecked: number; +} + +export class QuotaTracker { + private quotas = new Map(); + 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): 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 } { + return { + quotas: Object.fromEntries(this.quotas), + }; + } + + fromJSON(data: { quotas?: Record }): void { + if (data?.quotas) { + for (const [key, value] of Object.entries(data.quotas)) { + this.quotas.set(key, value); + } + } + } +} diff --git a/src/agents/multi-account/rate-limit-tracker.ts b/src/agents/multi-account/rate-limit-tracker.ts new file mode 100644 index 000000000..3a6426150 --- /dev/null +++ b/src/agents/multi-account/rate-limit-tracker.ts @@ -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(); + private failureCount = new Map(); + + 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; + failureCounts: Record; + } { + const state: Record = {}; + for (const [key, value] of this.stateMap.entries()) { + state[key] = value; + } + return { + rateState: state, + failureCounts: Object.fromEntries(this.failureCount), + }; + } + + fromJSON(data: { + rateState?: Record; + failureCounts?: Record; + }): 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(); + } +} diff --git a/src/agents/multi-account/strategies.ts b/src/agents/multi-account/strategies.ts new file mode 100644 index 000000000..93bd31451 --- /dev/null +++ b/src/agents/multi-account/strategies.ts @@ -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(); + + 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(); + + 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(); + + 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(); + } +} diff --git a/src/agents/pi-embedded-runner/run.ts b/src/agents/pi-embedded-runner/run.ts index adbe6ad49..853d5248f 100644 --- a/src/agents/pi-embedded-runner/run.ts +++ b/src/agents/pi-embedded-runner/run.ts @@ -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, diff --git a/src/cli/accounts-cli.ts b/src/cli/accounts-cli.ts new file mode 100644 index 000000000..1acbdbba4 --- /dev/null +++ b/src/cli/accounts-cli.ts @@ -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) { + 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 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 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 name", "google-antigravity") + .action((options) => + runAccountsCommand(() => accountsResetCommand(defaultRuntime, options)) + ); + + return accounts; +} diff --git a/src/cli/program/register.subclis.ts b/src/cli/program/register.subclis.ts index b4a65c794..a35de0174 100644 --- a/src/cli/program/register.subclis.ts +++ b/src/cli/program/register.subclis.ts @@ -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", diff --git a/src/commands/accounts/index.ts b/src/commands/accounts/index.ts new file mode 100644 index 000000000..2a13ad142 --- /dev/null +++ b/src/commands/accounts/index.ts @@ -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"; diff --git a/src/commands/accounts/list.ts b/src/commands/accounts/list.ts new file mode 100644 index 000000000..d05914c2d --- /dev/null +++ b/src/commands/accounts/list.ts @@ -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 { + 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})`); + } +} diff --git a/src/commands/accounts/reset.ts b/src/commands/accounts/reset.ts new file mode 100644 index 000000000..9ac5a46f7 --- /dev/null +++ b/src/commands/accounts/reset.ts @@ -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 { + 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.`); +} diff --git a/src/commands/accounts/status.ts b/src/commands/accounts/status.ts new file mode 100644 index 000000000..2a1b48f09 --- /dev/null +++ b/src/commands/accounts/status.ts @@ -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 { + 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}`); + } + } +} diff --git a/src/config/types.auth.ts b/src/config/types.auth.ts index 2995224a4..1d9e1573a 100644 --- a/src/config/types.auth.ts +++ b/src/config/types.auth.ts @@ -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; order?: Record; @@ -26,4 +42,6 @@ export type AuthConfig = { */ failureWindowHours?: number; }; + /** Multi-account load balancing configuration */ + multiAccount?: MultiAccountConfig; }; diff --git a/src/config/zod-schema.ts b/src/config/zod-schema.ts index 3c5bba8d7..e849696da 100644 --- a/src/config/zod-schema.ts +++ b/src/config/zod-schema.ts @@ -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(),