fix(auth): improve multi-account round-robin rotation and 429 handling

This commit fixes several issues with multi-account OAuth rotation that
were causing slow responses and inefficient account cycling.

## Changes

### 1. Fix usageStats race condition (auth-profiles.ts)

The `markAuthProfileUsed`, `markAuthProfileCooldown`, `markAuthProfileGood`,
and `clearAuthProfileCooldown` functions were using a stale in-memory store
passed as a parameter. Long-running sessions would overwrite usageStats
updates from concurrent sessions when saving.

**Fix:** Re-read the store from disk before each update to get fresh
usageStats from other sessions, then merge the update.

### 2. Capture AbortError from waitForCompactionRetry (pi-embedded-runner.ts)

When a request timed out, `session.abort()` was called which throws an
`AbortError`. The code structure was:

```javascript
try {
  await session.prompt(params.prompt);
} catch (err) {
  promptError = err;  // Catches AbortError here
}
await waitForCompactionRetry();  // But THIS also throws AbortError!
```

The second `AbortError` from `waitForCompactionRetry()` escaped and
bypassed the rotation/fallback logic entirely.

**Fix:** Wrap `waitForCompactionRetry()` in its own try/catch to capture
the error as `promptError`, enabling proper timeout handling.

Root cause analysis and fix proposed by @erikpr1994 in #313.

Fixes #313

### 3. Fail fast on 429 rate limits (pi-ai patch)

The pi-ai library was retrying 429 errors up to 3 times with exponential
backoff before throwing. This meant a rate-limited account would waste
30+ seconds retrying before our rotation code could try the next account.

**Fix:** Patch google-gemini-cli.js to:
- Throw immediately on first 429 (no retries)
- Not catch and retry 429 errors in the network error handler

This allows the caller to rotate to the next account instantly on rate limit.

Note: We submitted this fix upstream (https://github.com/badlogic/pi-mono/pull/504)
but it was closed without merging. Keeping as a local patch for now.

## Testing

With 6 Antigravity accounts configured:
- Accounts rotate properly based on lastUsed (round-robin)
- 429s trigger immediate rotation to next account
- usageStats persist correctly across concurrent sessions
- Cooldown tracking works as expected

## Before/After

**Before:** Multiple 429 retries on same account, 30-90s delays
**After:** Instant rotation on 429, responses in seconds
This commit is contained in:
Muhammed Mukhthar CM
2026-01-06 22:44:19 +00:00
committed by Peter Steinberger
parent 2871657ebe
commit eb5f758f6b
4 changed files with 71 additions and 27 deletions

10
pnpm-lock.yaml generated
View File

@@ -9,7 +9,7 @@ overrides:
patchedDependencies:
'@mariozechner/pi-ai':
hash: b49275c3e2023970d8248ababef6df60e093e58a3ba3127c2ba4de1df387d06a
hash: 31ccee9eb1fc5053d766ccbf3abc8340a5fb3bea1c52b491a385b53468b81819
path: patches/@mariozechner__pi-ai.patch
playwright-core@1.57.0:
hash: 66f1f266424dbe354068aaa5bba87bfb0e1d7d834a938c25dd70d43cdf1c1b02
@@ -39,7 +39,7 @@ importers:
version: 0.37.2(ws@8.19.0)(zod@4.3.5)
'@mariozechner/pi-ai':
specifier: ^0.37.2
version: 0.37.2(patch_hash=b49275c3e2023970d8248ababef6df60e093e58a3ba3127c2ba4de1df387d06a)(ws@8.19.0)(zod@4.3.5)
version: 0.37.2(patch_hash=31ccee9eb1fc5053d766ccbf3abc8340a5fb3bea1c52b491a385b53468b81819)(ws@8.19.0)(zod@4.3.5)
'@mariozechner/pi-coding-agent':
specifier: ^0.37.2
version: 0.37.2(ws@8.19.0)(zod@4.3.5)
@@ -3570,7 +3570,7 @@ snapshots:
'@mariozechner/pi-agent-core@0.37.2(ws@8.19.0)(zod@4.3.5)':
dependencies:
'@mariozechner/pi-ai': 0.37.2(patch_hash=b49275c3e2023970d8248ababef6df60e093e58a3ba3127c2ba4de1df387d06a)(ws@8.19.0)(zod@4.3.5)
'@mariozechner/pi-ai': 0.37.2(patch_hash=31ccee9eb1fc5053d766ccbf3abc8340a5fb3bea1c52b491a385b53468b81819)(ws@8.19.0)(zod@4.3.5)
'@mariozechner/pi-tui': 0.37.2
transitivePeerDependencies:
- '@modelcontextprotocol/sdk'
@@ -3580,7 +3580,7 @@ snapshots:
- ws
- zod
'@mariozechner/pi-ai@0.37.2(patch_hash=b49275c3e2023970d8248ababef6df60e093e58a3ba3127c2ba4de1df387d06a)(ws@8.19.0)(zod@4.3.5)':
'@mariozechner/pi-ai@0.37.2(patch_hash=31ccee9eb1fc5053d766ccbf3abc8340a5fb3bea1c52b491a385b53468b81819)(ws@8.19.0)(zod@4.3.5)':
dependencies:
'@anthropic-ai/sdk': 0.71.2(zod@4.3.5)
'@google/genai': 1.34.0
@@ -3604,7 +3604,7 @@ snapshots:
dependencies:
'@crosscopy/clipboard': 0.2.8
'@mariozechner/pi-agent-core': 0.37.2(ws@8.19.0)(zod@4.3.5)
'@mariozechner/pi-ai': 0.37.2(patch_hash=b49275c3e2023970d8248ababef6df60e093e58a3ba3127c2ba4de1df387d06a)(ws@8.19.0)(zod@4.3.5)
'@mariozechner/pi-ai': 0.37.2(patch_hash=31ccee9eb1fc5053d766ccbf3abc8340a5fb3bea1c52b491a385b53468b81819)(ws@8.19.0)(zod@4.3.5)
'@mariozechner/pi-tui': 0.37.2
chalk: 5.6.2
cli-highlight: 2.1.11