fix: lint and format issues for multi-account module

This commit is contained in:
PREETHAM1590 2026-01-27 17:36:25 +05:30
parent c20b27f5e6
commit c5a28b28d4
15 changed files with 77 additions and 132 deletions

View File

@ -8,7 +8,7 @@ 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";
import { DEFAULT_COOLDOWN_MS } from "./constants.js";
export interface Account {
profileId: string;
@ -27,15 +27,18 @@ export interface AccountManagerConfig {
}
export interface AuthProfileStore {
profiles: Record<string, {
type: string;
provider: string;
email?: string;
access?: string;
refresh?: string;
expires?: number;
projectId?: string;
}>;
profiles: Record<
string,
{
type: string;
provider: string;
email?: string;
access?: string;
refresh?: string;
expires?: number;
projectId?: string;
}
>;
}
export class AccountManager {
@ -59,10 +62,7 @@ export class AccountManager {
this.strategy.init(this);
}
async initialize(
authStore: AuthProfileStore,
strategyOverride?: string
): Promise<this> {
async initialize(authStore: AuthProfileStore, strategyOverride?: string): Promise<this> {
this.authStore = authStore;
if (strategyOverride && STRATEGY_NAMES.includes(strategyOverride.toLowerCase() as any)) {
@ -93,9 +93,7 @@ export class AccountManager {
}
getProfilesForProvider(): string[] {
return Array.from(this.accounts.keys()).filter(
(id) => !this.invalidProfiles.has(id)
);
return Array.from(this.accounts.keys()).filter((id) => !this.invalidProfiles.has(id));
}
getAccountByProfileId(profileId: string): Account | null {
@ -147,11 +145,7 @@ export class AccountManager {
}
markRateLimited(profileId: string, modelId: string, cooldownMs?: number): void {
this.rateLimitTracker.markRateLimited(
profileId,
modelId,
cooldownMs ?? this.defaultCooldownMs
);
this.rateLimitTracker.markRateLimited(profileId, modelId, cooldownMs ?? this.defaultCooldownMs);
}
markInvalid(profileId: string, reason: string): void {
@ -236,12 +230,8 @@ export class AccountManager {
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)
);
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,

View File

@ -21,46 +21,33 @@ export class HealthScorer {
// Apply time-based recovery
const minutesSinceUpdate = (Date.now() - lastUpdateTime) / 60_000;
const recovery = Math.floor(
minutesSinceUpdate * HEALTH_SCORE.recoveryPerMinute
);
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
);
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
);
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
);
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 }> {
getSortedByHealth(profileIds: string[]): Array<{ profileId: string; score: number }> {
return profileIds
.map((profileId) => ({
profileId,
@ -79,10 +66,7 @@ export class HealthScorer {
};
}
fromJSON(data: {
scores?: Record<string, number>;
lastUpdate?: Record<string, number>;
}): void {
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);

View File

@ -1,6 +1,6 @@
/**
* Multi-Account Module
*
*
* Provides multi-account load balancing for Clawdbot providers.
*/

View File

@ -18,7 +18,7 @@ const managers = new Map<string, AccountManager>();
export async function getAccountManager(
provider: string,
authStore: AuthProfileStore,
cfg?: ClawdbotConfig
cfg?: ClawdbotConfig,
): Promise<AccountManager> {
if (managers.has(provider)) {
return managers.get(provider)!;
@ -40,10 +40,7 @@ export async function getAccountManager(
/**
* Check if multi-account is enabled for a provider
*/
export function isMultiAccountEnabled(
provider: string,
cfg?: ClawdbotConfig
): boolean {
export function isMultiAccountEnabled(provider: string, cfg?: ClawdbotConfig): boolean {
const multiAccountConfig = (cfg as any)?.auth?.multiAccount;
if (!multiAccountConfig?.enabled) return false;
@ -91,7 +88,7 @@ export async function resolveWithMultiAccount(params: {
const minWait = manager.getMinWaitTimeMs(modelId);
throw new Error(
`MULTI_ACCOUNT_EXHAUSTED: All accounts rate limited for ${modelId}. ` +
`Wait ${Math.ceil(minWait / 1000)}s.`
`Wait ${Math.ceil(minWait / 1000)}s.`,
);
}
return null;
@ -116,7 +113,7 @@ export async function resolveWithMultiAccount(params: {
*/
export function reportMultiAccountSuccess(
result: { _manager?: AccountManager; _profileId?: string } | null,
modelId: string
modelId: string,
): void {
if (!result?._manager || !result?._profileId) return;
result._manager.notifySuccess(result._profileId, modelId);
@ -128,7 +125,7 @@ export function reportMultiAccountSuccess(
export function reportMultiAccountRateLimit(
result: { _manager?: AccountManager; _profileId?: string } | null,
modelId: string,
cooldownMs?: number
cooldownMs?: number,
): void {
if (!result?._manager || !result?._profileId) return;
result._manager.notifyRateLimit(result._profileId, modelId, cooldownMs);
@ -139,7 +136,7 @@ export function reportMultiAccountRateLimit(
*/
export function reportMultiAccountFailure(
result: { _manager?: AccountManager; _profileId?: string } | null,
modelId: string
modelId: string,
): void {
if (!result?._manager || !result?._profileId) return;
result._manager.notifyFailure(result._profileId, modelId);
@ -150,7 +147,7 @@ export function reportMultiAccountFailure(
*/
export function reportMultiAccountAuthInvalid(
result: { _manager?: AccountManager; _profileId?: string } | null,
reason: string
reason: string,
): void {
if (!result?._manager || !result?._profileId) return;
result._manager.markInvalid(result._profileId, reason);

View File

@ -40,7 +40,7 @@ export function isMultiAccountEnabled(provider: string, cfg?: ClawdbotConfig): b
export async function getOrCreateManager(
provider: string,
store: AuthProfileStore,
cfg?: ClawdbotConfig
cfg?: ClawdbotConfig,
): Promise<AccountManager> {
const normalized = normalizeProviderId(provider);
@ -97,7 +97,7 @@ export async function selectAccountForModel(params: {
const waitSec = Math.ceil(waitMs / 1000);
throw new Error(
`All ${manager.getAccountCount()} accounts rate-limited for ${modelId}. ` +
`Shortest wait: ${waitSec}s. Try again later.`
`Shortest wait: ${waitSec}s. Try again later.`,
);
}
return null;
@ -112,7 +112,7 @@ export async function selectAccountForModel(params: {
export function notifyMultiAccountSuccess(
selection: MultiAccountSelection | null | undefined,
modelId: string
modelId: string,
): void {
if (!selection) return;
selection.manager.notifySuccess(selection.profileId, modelId);
@ -121,7 +121,7 @@ export function notifyMultiAccountSuccess(
export function notifyMultiAccountRateLimit(
selection: MultiAccountSelection | null | undefined,
modelId: string,
cooldownMs?: number
cooldownMs?: number,
): void {
if (!selection) return;
selection.manager.notifyRateLimit(selection.profileId, modelId, cooldownMs);
@ -129,7 +129,7 @@ export function notifyMultiAccountRateLimit(
export function notifyMultiAccountFailure(
selection: MultiAccountSelection | null | undefined,
modelId: string
modelId: string,
): void {
if (!selection) return;
selection.manager.notifyFailure(selection.profileId, modelId);
@ -137,16 +137,13 @@ export function notifyMultiAccountFailure(
export function notifyMultiAccountInvalid(
selection: MultiAccountSelection | null | undefined,
reason: string
reason: string,
): void {
if (!selection) return;
selection.manager.markInvalid(selection.profileId, reason);
}
export function getMultiAccountStatus(): Record<
string,
ReturnType<AccountManager["getStatus"]>
> {
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();

View File

@ -60,13 +60,13 @@ describe("RateLimitTracker", () => {
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);
});
});
@ -121,7 +121,7 @@ describe("AccountManager", () => {
const profiles = manager.getProfilesForProvider();
const first = profiles[0];
manager.markInvalid(first, "Token revoked");
const updated = manager.getProfilesForProvider();
expect(updated).not.toContain(first);
});

View File

@ -55,11 +55,7 @@ export class QuotaTracker {
});
}
updateModelQuota(
profileId: string,
modelId: string,
quota: ModelQuota
): void {
updateModelQuota(profileId: string, modelId: string, quota: ModelQuota): void {
const existing = this.getQuota(profileId) ?? {
tier: "unknown" as const,
projectId: null,

View File

@ -71,7 +71,7 @@ export class RateLimitTracker {
markRateLimited(
profileId: string,
modelId: string,
cooldownMs: number = DEFAULT_COOLDOWN_MS
cooldownMs: number = DEFAULT_COOLDOWN_MS,
): BackoffResult {
const key = this.getKey(profileId, modelId);
const now = Date.now();
@ -81,7 +81,7 @@ export class RateLimitTracker {
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
60_000,
);
return {
attempt: previous.consecutive429,
@ -97,10 +97,7 @@ export class RateLimitTracker {
: 1;
// Calculate exponential backoff
const backoffDelay = Math.min(
FIRST_RETRY_DELAY_MS * Math.pow(2, attempt - 1),
60_000
);
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

View File

@ -29,9 +29,7 @@ export class HybridStrategy implements Strategy {
const { rateLimitTracker, healthScorer, quotaTracker } = this.manager;
const profiles = this.manager.getProfilesForProvider();
const available = profiles.filter(
(id) => !rateLimitTracker.isRateLimited(id, modelId)
);
const available = profiles.filter((id) => !rateLimitTracker.isRateLimited(id, modelId));
if (available.length === 0) {
const minWait = rateLimitTracker.getMinWaitTime(profiles, modelId);
@ -98,9 +96,7 @@ export class StickyStrategy implements Strategy {
};
}
const available = profiles.filter(
(id) => !rateLimitTracker.isRateLimited(id, modelId)
);
const available = profiles.filter((id) => !rateLimitTracker.isRateLimited(id, modelId));
if (available.length === 0) {
const minWait = rateLimitTracker.getMinWaitTime(profiles, modelId);

View File

@ -41,11 +41,7 @@ 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 { isMultiAccountEnabled, getManager, getOrCreateManager } from "../multi-account/index.js";
import { compactEmbeddedPiSessionDirect } from "./compact.js";
import { resolveGlobalLane, resolveSessionLane } from "./lanes.js";
@ -141,12 +137,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) {
@ -158,7 +154,7 @@ export async function runEmbeddedPiAgent(
lockedProfileId = undefined;
}
}
// Use multi-account intelligent ordering when enabled
let profileOrder: string[];
const manager = getManager(provider);
@ -170,11 +166,11 @@ export async function runEmbeddedPiAgent(
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({
@ -185,20 +181,23 @@ export async function runEmbeddedPiAgent(
available.push(profileId);
}
}
// Sort available by health score (highest first)
const healthSorted = manager.healthScorer.getSortedByHealth(available);
const sortedAvailable = healthSorted.map(h => h.profileId);
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);
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)];
profileOrder = [
preferredProfileId,
...profileOrder.filter((p) => p !== preferredProfileId),
];
}
} else {
// Fall back to standard ordering

View File

@ -4,7 +4,6 @@
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 {
@ -24,7 +23,7 @@ export function registerAccountsCli(program: Command) {
.addHelpText(
"after",
() =>
`\n${theme.muted("Multi-account distributes requests across multiple OAuth accounts for higher throughput.")}\n`
`\n${theme.muted("Multi-account distributes requests across multiple OAuth accounts for higher throughput.")}\n`,
);
accounts
@ -32,26 +31,20 @@ export function registerAccountsCli(program: Command) {
.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))
);
.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))
);
.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))
);
.action((options) => runAccountsCommand(() => accountsResetCommand(defaultRuntime, options)));
return accounts;
}

View File

@ -15,7 +15,7 @@ export interface AccountsListOptions {
export async function accountsListCommand(
runtime: RuntimeEnv,
options: AccountsListOptions
options: AccountsListOptions,
): Promise<void> {
const cfg = loadConfig();
const agentDir = resolveClawdbotAgentDir();

View File

@ -11,7 +11,7 @@ export interface AccountsResetOptions {
export async function accountsResetCommand(
runtime: RuntimeEnv,
options: AccountsResetOptions
options: AccountsResetOptions,
): Promise<void> {
const provider = options.provider ?? "google-antigravity";

View File

@ -19,7 +19,7 @@ export interface AccountsStatusOptions {
export async function accountsStatusCommand(
runtime: RuntimeEnv,
options: AccountsStatusOptions
options: AccountsStatusOptions,
): Promise<void> {
const cfg = loadConfig();
const agentDir = resolveClawdbotAgentDir();
@ -34,7 +34,7 @@ export async function accountsStatusCommand(
` multiAccount:\n` +
` enabled: true\n` +
` providers:\n` +
` - ${provider}`
` - ${provider}`,
);
return;
}
@ -69,9 +69,7 @@ export async function accountsStatusCommand(
const status = acc.isInvalid
? `❌ Invalid (${acc.invalidReason})`
: `✅ Health: ${acc.healthScore}`;
const lastUsed = acc.lastUsed
? new Date(acc.lastUsed).toLocaleTimeString()
: "never";
const lastUsed = acc.lastUsed ? new Date(acc.lastUsed).toLocaleTimeString() : "never";
runtime.log(`${acc.email}`);
runtime.log(` Status: ${status}`);
runtime.log(` Last used: ${lastUsed}`);

View File

@ -206,11 +206,9 @@ export const ClawdbotSchema = z
multiAccount: z
.object({
enabled: z.boolean().optional(),
strategy: z.union([
z.literal("hybrid"),
z.literal("sticky"),
z.literal("round-robin"),
]).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(),
})