Feat/better account rotation #180
Conversation
add more rotation method like hybrid Signed-off-by: CypherpunkSamurai <CypherpunkSamurai@users.noreply.github.com>
add tests for rotation Signed-off-by: CypherpunkSamurai <CypherpunkSamurai@users.noreply.github.com>
|
Warning Rate limit exceeded@CypherpunkSamurai has exceeded the limit for the number of commits that can be reviewed per hour. Please wait 15 minutes and 5 seconds before requesting another review. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. 📒 Files selected for processing (1)
WalkthroughAdds a rotation subsystem (health-score and token-bucket trackers), integrates priority-queue selection into account manager and request flow, initializes trackers on plugin startup, consumes/refunds tokens during request lifecycle, and records successes, rate-limits, and failures for account health. Changes
Sequence DiagramssequenceDiagram
participant Client
participant Plugin
participant AccountManager
participant Rotation
participant Upstream
Client->>Plugin: Send request
Plugin->>Rotation: evaluate candidates (health, tokens, rate-limit)
Rotation-->>Plugin: selected account index
Plugin->>Rotation: consume(accountIndex, cost)
Rotation-->>Plugin: token consumed
Plugin->>AccountManager: set current account
Plugin->>Upstream: execute request using selected account
alt Success
Upstream-->>Plugin: 200 OK
Plugin->>Rotation: recordSuccess(accountIndex)
else Rate Limit (429)
Upstream-->>Plugin: 429
Plugin->>Rotation: recordRateLimit(accountIndex)
Plugin->>Rotation: refund(accountIndex, cost)
Plugin->>Plugin: apply backoff / quota switching
else Network/API Error
Upstream-->>Plugin: Error
Plugin->>Rotation: recordFailure(accountIndex)
Plugin->>Rotation: refund(accountIndex, cost)
end
Plugin-->>Client: return response
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In @src/plugin/rotation.ts:
- Around line 416-419: The score calculation is using a hardcoded 50 instead of
the tracker's configured max tokens; update the call in the candidates mapping
so calculatePriorityScore(acc, ...) uses the TokenBucketTracker's actual
maxTokens (e.g., call tracker.getMaxTokens() or access the
tracker.config.maxTokens) rather than 50; if maxTokens is private, add a
TokenBucketTracker.getMaxTokens() accessor and use that in the mapping to
compute each candidate's score.
🧹 Nitpick comments (4)
src/plugin/accounts.ts (1)
287-313: Priority-queue strategy lacks fallback to default sticky behavior.When
selectPriorityQueueAccountreturnsnull(all accounts filtered out), the code silently falls through without returning. This means the function continues to the sticky strategy code below (line 341+), which may not be intentional behavior. Consider either:
- Returning
nullexplicitly to indicate no account available- Adding a comment clarifying this is intentional fallback behavior
Also, the comment on line 306 mentions token consumption is deferred to
plugin.ts, which aligns with the implementation in that file.💡 Suggestion: Make fallback behavior explicit
const selectedIndex = selectPriorityQueueAccount(accountsWithMetrics, tokenTracker); if (selectedIndex !== null) { const selected = this.accounts[selectedIndex]; if (selected) { // Token consumption is deferred to request execution in plugin.ts selected.lastUsed = nowMs(); this.markTouchedForQuota(selected, quotaKey); this.currentAccountIndexByFamily[family] = selected.index; return selected; } } + // Fall through to sticky strategy when priority-queue finds no candidates }src/plugin/rotation.test.ts (1)
397-474: selectPriorityQueueAccount tests cover key scenarios.Tests verify:
- Empty input handling
- Filtering of rate-limited accounts
- Filtering of accounts without tokens
- Filtering of unhealthy accounts
- Single candidate scenarios
- Valid index returned from multiple candidates
Consider adding a test for the weighted random distribution to verify that higher-scored accounts are selected more frequently over many iterations.
💡 Optional: Add weighted distribution test
it("favors higher-scored candidates in weighted selection", () => { const tokenTracker = new TokenBucketTracker(); const accounts: AccountWithMetrics[] = [ { index: 0, lastUsed: 1000, healthScore: 100, isRateLimited: false, isCoolingDown: false }, { index: 1, lastUsed: 1000, healthScore: 50, isRateLimited: false, isCoolingDown: false }, ]; const counts = new Map<number, number>(); for (let i = 0; i < 1000; i++) { const result = selectPriorityQueueAccount(accounts, tokenTracker); counts.set(result!, (counts.get(result!) ?? 0) + 1); } // Higher health score should be selected more often expect(counts.get(0)!).toBeGreaterThan(counts.get(1)!); });src/plugin/rotation.ts (2)
438-440: Fallback chain returns 0 as last resort.The fallback
?? 0on line 439 would return index 0 even if no candidates exist. However, this code path should be unreachable since we already returnnullon line 412 whencandidates.length === 0.The fallback is defensive but could mask bugs if somehow reached.
💡 Consider removing unreachable fallback
- // Fallback to top score - return topCandidates[0]?.acc.index ?? candidates[0]?.index ?? 0; + // Fallback to top score (topCandidates is guaranteed non-empty at this point) + return topCandidates[0]!.acc.index;
446-480: Singleton pattern for global trackers.The
init*functions overwrite existing trackers without warning. This is fine for plugin startup but could cause issues if called multiple times. Theget*functions safely create with defaults if not initialized.Consider adding a warning log if re-initializing an existing tracker to help debug unexpected behavior.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Cache: Disabled due to data retention organization setting
Knowledge base: Disabled due to data retention organization setting
📒 Files selected for processing (5)
src/plugin.tssrc/plugin/accounts.tssrc/plugin/config/schema.tssrc/plugin/rotation.test.tssrc/plugin/rotation.ts
🧰 Additional context used
🧬 Code graph analysis (3)
src/plugin.ts (1)
src/plugin/rotation.ts (4)
initHealthTracker(477-480)initTokenTracker(455-458)getHealthTracker(466-471)getTokenTracker(448-453)
src/plugin/rotation.test.ts (1)
src/plugin/rotation.ts (8)
HealthScoreTracker(54-163)TokenBucketTracker(283-345)addJitter(177-181)randomDelay(190-192)AccountWithMetrics(198-204)sortByLruWithHealth(215-229)selectHybridAccount(241-253)selectPriorityQueueAccount(393-440)
src/plugin/accounts.ts (1)
src/plugin/rotation.ts (5)
getHealthTracker(466-471)getTokenTracker(448-453)AccountWithMetrics(198-204)selectPriorityQueueAccount(393-440)sortByLruWithHealth(215-229)
🔇 Additional comments (23)
src/plugin/accounts.ts (2)
5-5: LGTM on imports.The new imports from
./rotationalign with the strategy implementations below and follow the module's public API.
315-339: Hybrid strategy also lacks explicit fallback handling.Same observation as priority-queue: when
sortByLruWithHealthreturns an empty array, the code falls through to sticky behavior. The logic is correct, but documenting the intentional fallback would improve maintainability.The implementation correctly uses
sortByLruWithHealthand selects the first candidate (LRU with health tiebreaker).src/plugin.ts (6)
42-42: LGTM on imports.The imports correctly include both getter and initializer functions for the trackers.
666-686: Tracker initialization depends on optional config blocks.The trackers are only explicitly initialized when their config sections exist. However,
getHealthTracker()andgetTokenTracker()are called unconditionally in the request flow (lines 988, 1037, 1063, 1229, 1242, 1280, 1437, 1539). These getter functions create trackers with default config if not initialized, which is safe but could lead to unexpected default behavior.This is acceptable since the defaults are sensible, but consider documenting this behavior or always initializing with defaults to make it explicit.
1226-1230: Token consumption only for priority-queue strategy - verify alignment with accounts.ts.Token is consumed here before the request, but in
accounts.tsthe comment says "Token consumption is deferred to request execution in plugin.ts". This aligns correctly.However, ensure that the
consumecall here and thehasTokenscheck inselectPriorityQueueAccount(rotation.ts line 404) are consistent - currently both use default cost of 1.
1240-1243: Token refund correctly handles 429 and network errors.The refund logic properly returns the token when:
- Rate limit (429) is encountered
- Network/API errors occur in the catch block
This ensures tokens aren't lost due to failed requests. Good defensive pattern.
Also applies to: 1537-1540
988-988: Health tracker records failures on auth and project errors.These
recordFailurecalls appropriately penalize accounts that encounter authentication or project context errors. The tracker is obtained viagetHealthTracker()which creates with defaults if not initialized.Also applies to: 1037-1037, 1063-1063
1280-1280: Health events recorded appropriately on rate limit and success.
recordRateLimiton 429 applies moderate penaltyrecordSuccesson successful response improves healthThis aligns with the health score design in rotation.ts.
Also applies to: 1437-1437
src/plugin/rotation.test.ts (5)
1-12: LGTM on test setup and imports.Clean import of all public exports from the rotation module. Using vitest with real timers reset is appropriate for most tests.
14-154: Good coverage for HealthScoreTracker.Tests cover:
- Initial state and custom config
- Success recording with score cap
- Rate limit and failure penalties with floor at 0
- Consecutive failures tracking
- Usability threshold
- Reset and snapshot functionality
156-242: Good coverage for TokenBucketTracker.Tests cover:
- Initial state and custom config
- Token consumption with success/failure cases
- Exact balance consumption
- Multiple consumes
- hasTokens with custom cost
- Refund with cap at maxTokens
244-301: Good statistical tests for jitter utilities.Using 100 iterations to verify range bounds is appropriate for testing random functions. Tests correctly verify:
- Range bounds for addJitter
- Default jitter factor
- Non-negative results
- Integer results
- randomDelay range and edge case (min === max)
303-395: Comprehensive tests for selection functions.Tests for
sortByLruWithHealthandselectHybridAccountcover:
- Filtering rate-limited, cooling-down, and unhealthy accounts
- LRU sorting (oldest first)
- Health score as tiebreaker
- Empty results handling
- Deterministic selection of best candidate
src/plugin/config/schema.ts (4)
20-22: LGTM on strategy schema extension.Good documentation for the new
priority-queuestrategy. The enum extension is backward compatible.
242-242: Breaking change: Default strategy changed from 'sticky' to 'hybrid'.This changes the default behavior for all users who haven't explicitly configured
account_selection_strategy. The hybrid strategy will now touch all fresh accounts to sync reset timers before becoming sticky.While hybrid may be better for multi-account setups, this could surprise users expecting sticky behavior (preserves prompt cache). Consider:
- Documenting this change prominently in release notes
- Or keeping 'sticky' as default for backward compatibility
266-288: Well-structured config blocks for health score and token bucket.Good validation ranges:
- Health scores: 0-100 with appropriate sub-ranges for penalties/rewards
- Token bucket: reasonable defaults (50 tokens, 6/min regeneration)
The
optional()wrapper allows these to be omitted, letting the code use defaults.
333-346: DEFAULT_CONFIG includes full health_score and token_bucket objects.This ensures consistent defaults whether config is loaded from file or not. The values match the Zod schema defaults.
src/plugin/rotation.ts (6)
1-11: Good module documentation.Clear overview of the rotation system's purpose and strategies.
54-163: HealthScoreTracker implementation is sound.Key observations:
- Time-based recovery in
getScore()correctly applies passive healingrecordSuccessresets consecutive failuresrecordRateLimitandrecordFailureproperly increment consecutive failures- Score is clamped between 0 and maxScore
getSnapshot()provides debugging capability
177-192: Jitter utilities are correctly implemented.
addJitter: Properly generates value in ±jitterFactor range, ensures non-negativerandomDelay: Simple uniform distribution between min and max
215-253: LRU selection functions are well-designed.
sortByLruWithHealth:
- Correctly filters unavailable accounts
- Primary sort by lastUsed (oldest first)
- Tiebreaker by health score (higher wins)
selectHybridAccount:
- Simple wrapper that returns first sorted candidate
- Properly returns null for empty results
283-345: TokenBucketTracker implements standard token bucket algorithm.
getTokens: Correctly applies time-based regenerationhasTokens: Simple threshold check with default costconsume: Atomic check-and-decrementrefund: Properly caps at maxTokens
363-377: Priority score formula weights tokens heavily over health.The formula:
health*2 + (tokens/max)*500 + freshness*0.1With default 50 tokens and 100 health:
- Health component: 0-200 points
- Token component: 0-500 points
- Freshness component: 0-360 points (capped at 1 hour)
Tokens dominate the score, which may cause high-token accounts to always win even with lower health. This could be intentional (prioritize accounts with capacity) but worth verifying this matches the design intent.
There was a problem hiding this comment.
Pull request overview
This PR implements an advanced account rotation system to prevent account limiting and bans. The system introduces health scoring to track account wellness, a token bucket algorithm for client-side rate limiting, and sophisticated selection strategies (hybrid and priority-queue) that consider account health, token availability, and LRU (Least Recently Used) metrics.
Changes:
- Added rotation.ts with HealthScoreTracker, TokenBucketTracker, LRU selection, and weighted random selection algorithms
- Added comprehensive test suite in rotation.test.ts covering core functionality
- Updated config schema to support new 'priority-queue' strategy and changed default from 'sticky' to 'hybrid'
- Integrated health tracking and token management into accounts.ts and plugin.ts
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 10 comments.
Show a summary per file
| File | Description |
|---|---|
| src/plugin/rotation.ts | Core rotation system with health scoring, token bucket, LRU selection, and priority queue algorithms |
| src/plugin/rotation.test.ts | Comprehensive test suite for rotation functionality |
| src/plugin/config/schema.ts | Schema updates to support new strategies, health score config, and token bucket config with default changed to 'hybrid' |
| src/plugin/accounts.ts | Integration of hybrid and priority-queue strategies into account selection logic |
| src/plugin.ts | Initialization of trackers and integration of health/token tracking throughout request lifecycle |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
use max tokens instead of hardcoded values Signed-off-by: CypherpunkSamurai <CypherpunkSamurai@users.noreply.github.com>
update default value Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
update default value for rotation Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In @src/plugin/rotation.ts:
- Around line 436-444: The fallback returns a hardcoded 0 which can be invalid;
update the final return in the random selection block so it never falls back to
literal 0: return topCandidates[0].acc.index ?? candidates[0].index (or simply
return topCandidates[0].acc.index since topCandidates is guaranteed non-empty).
Change the expression in rotation.ts (the loop using topCandidates and
randomValue) to remove "?? 0" and rely on the existing candidate objects
(topCandidates[0] or candidates[0]) as the safe fallback.
🧹 Nitpick comments (3)
src/plugin/rotation.ts (3)
294-308: Minor inconsistency: Token regeneration produces fractional values.Unlike
HealthScoreTracker.getScore()which usesMath.floorfor recovered points (line 74),getTokens()doesn't floor the recovered tokens, resulting in fractional token counts. This works correctly but is inconsistent between the two trackers.Consider whether fractional tokens are intentional or if
Math.floorshould be applied for consistency:const minutesSinceUpdate = (now - state.lastUpdated) / (1000 * 60); - const recoveredTokens = minutesSinceUpdate * this.config.regenerationRatePerMinute; + const recoveredTokens = Math.floor(minutesSinceUpdate * this.config.regenerationRatePerMinute);
453-476: Config parameter is silently ignored after first initialization.In
getTokenTrackerandgetHealthTracker, theconfigparameter is only used when creating a new instance. If a tracker already exists, subsequent calls with different configs will silently use the existing tracker, potentially leading to unexpected behavior.This is a common singleton pattern trade-off. The current design relies on calling
initTokenTracker/initHealthTrackerat startup with the desired config. Consider adding a warning or documentation note for callers who might expect config to be applied on every call.Optional: Add runtime warning when config is ignored
export function getTokenTracker(config?: Partial<TokenBucketConfig>): TokenBucketTracker { if (!globalTokenTracker) { globalTokenTracker = new TokenBucketTracker(config); + } else if (config !== undefined) { + console.warn('getTokenTracker: config ignored, tracker already initialized. Use initTokenTracker to reconfigure.'); } return globalTokenTracker; }
54-56: Consider adding cleanup for long-running processes.Both
HealthScoreTrackerandTokenBucketTrackeruseMapstructures that grow as accounts are tracked but only shrink via explicitreset()calls (health tracker only). For long-running processes with account churn, consider:
- Adding a
reset(accountIndex)method toTokenBucketTrackerfor parity.- Optionally, implementing a cleanup mechanism for stale entries (e.g., accounts not used for N hours).
This is unlikely to be an issue for typical usage patterns but worth considering for operational resilience.
Also applies to: 283-285
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Cache: Disabled due to data retention organization setting
Knowledge base: Disabled due to data retention organization setting
📒 Files selected for processing (2)
src/plugin/rotation.test.tssrc/plugin/rotation.ts
🚧 Files skipped from review as they are similar to previous changes (1)
- src/plugin/rotation.test.ts
🔇 Additional comments (4)
src/plugin/rotation.ts (4)
16-41: LGTM!The configuration interface and default values are well-structured with clear documentation. The scoring parameters provide a reasonable balance between penalizing failures and allowing recovery.
54-163: LGTM!The
HealthScoreTrackerclass is well-implemented:
- Time-based recovery in
getScore()is correctly calculated.- The state management properly handles both known and unknown accounts.
- Consecutive failure tracking provides useful insights for monitoring.
177-192: LGTM!The jitter utilities are correctly implemented.
addJitterproperly handles the symmetric jitter distribution and clamps to non-negative values.randomDelayprovides a clean uniform distribution.
215-253: LGTM!The LRU selection logic is well-implemented:
- Filtering correctly excludes rate-limited, cooling-down, and unhealthy accounts.
- The sort order (oldest first, then healthiest) aligns with the documented priority.
- The null-safety handling is robust.
ref: update default value for rotation Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
ref: update default value for rotation Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Signed-off-by: CypherpunkSamurai <CypherpunkSamurai@users.noreply.github.com>
Signed-off-by: CypherpunkSamurai <CypherpunkSamurai@users.noreply.github.com>
Signed-off-by: CypherpunkSamurai <CypherpunkSamurai@users.noreply.github.com>
update hybrid doc Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
update default value for rotation Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
clamp using `Math.Max` Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Signed-off-by: CypherpunkSamurai <CypherpunkSamurai@users.noreply.github.com>
Signed-off-by: CypherpunkSamurai <CypherpunkSamurai@users.noreply.github.com>
Signed-off-by: CypherpunkSamurai <CypherpunkSamurai@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In @src/plugin/accounts.test.ts:
- Around line 828-844: The test named "falls back to sticky when all accounts
unavailable" never marks the single account as unavailable, so it doesn't
exercise the fallback; modify the stored AccountStorageV3 used to construct
AccountManager so the sole account is actually unavailable at time 0 (e.g., add
a rate-limiting/cooldown field such as rateLimitedUntil or cooldownUntil set >
0, e.g., 1000, or use the AccountManager/API helper that marks accounts
unavailable) and then call manager.getCurrentOrNextForFamily("claude", null,
"priority-queue") and assert the fallback behavior (sticky index returned or
null depending on expected semantics).
🧹 Nitpick comments (4)
src/plugin/rotation.ts (2)
367-381: Priority score formula favors freshness (older = higher score).The
freshnessComponentadds to the score for accounts that haven't been used recently (capped at 1 hour = 360 points). This is the inverse of typical "freshness" naming (where fresh usually means recent). The naming is slightly counterintuitive but the comment on line 374-375 clarifies the intent.Consider renaming
freshnessComponenttorestPeriodBonusorlruBonusto better reflect that older/more-rested accounts get higher scores.
460-494: Singleton pattern for trackers lacks reset capability for testing.The global singletons work correctly for production use, but there's no way to reset them between tests. This could lead to test pollution if tests share the same process.
Consider adding reset functions for testing purposes:
♻️ Suggested addition for test isolation
// For testing only - resets singleton state export function resetTokenTracker(): void { globalTokenTracker = null; } export function resetHealthTracker(): void { globalHealthTracker = null; }src/plugin/config/schema.ts (1)
323-346: Inconsistent indentation in DEFAULT_CONFIG.Lines 324-326 have inconsistent indentation compared to surrounding lines. This is a minor formatting issue.
🔧 Fix indentation
account_selection_strategy: 'hybrid', -pid_offset_enabled: false, - switch_on_first_rate_limit: true, - auto_update: true, + pid_offset_enabled: false, + switch_on_first_rate_limit: true, + auto_update: true,src/plugin/accounts.test.ts (1)
773-790: Test doesn't verify health-based or token-based selection.The test name suggests it validates selection based on "health and token availability," but all accounts have identical metrics (same
lastUsed: 0, no rate limits, no cooldown). The assertion only confirms some account is returned, not that the priority-queue algorithm correctly prioritizes by health score or token availability.Consider mocking or initializing the health/token trackers with different values per account and asserting that the account with the best metrics is selected.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Cache: Disabled due to data retention organization setting
Knowledge base: Disabled due to data retention organization setting
📒 Files selected for processing (6)
src/plugin.tssrc/plugin/accounts.test.tssrc/plugin/accounts.tssrc/plugin/config/schema.tssrc/plugin/rotation.test.tssrc/plugin/rotation.ts
🧰 Additional context used
🧬 Code graph analysis (3)
src/plugin/accounts.test.ts (2)
src/plugin/storage.ts (1)
AccountStorageV3(76-84)src/plugin/accounts.ts (1)
AccountManager(107-586)
src/plugin.ts (1)
src/plugin/rotation.ts (4)
initHealthTracker(491-494)initTokenTracker(469-472)getHealthTracker(480-485)getTokenTracker(462-467)
src/plugin/rotation.test.ts (1)
src/plugin/rotation.ts (8)
HealthScoreTracker(54-163)TokenBucketTracker(283-349)addJitter(177-181)randomDelay(190-192)AccountWithMetrics(198-204)sortByLruWithHealth(215-229)selectHybridAccount(241-253)selectPriorityQueueAccount(397-454)
🔇 Additional comments (33)
src/plugin/rotation.ts (8)
16-41: LGTM - Health score configuration is well-designed.The configuration interface provides sensible defaults and clear documentation. The penalty values being negative is intuitive for the additive score model.
54-80: Time-based recovery calculation looks correct.The
getScoremethod properly applies passive recovery based on elapsed time and caps atmaxScore. The use ofMath.floorfor recovery points prevents fractional score accumulation.
85-127: Consider potential double-counting with time-based recovery.When
recordSuccess,recordRateLimit, orrecordFailureare called, they fetch the current score viagetScore()which includes time-based recovery, then store the adjusted score. This is correct behavior. However, be aware that if these methods are called in rapid succession, the time-based recovery between calls will be minimal (essentially zero), which is the intended behavior.
150-162:getSnapshotonly returns accounts that have been explicitly tracked.The snapshot won't include accounts that have never had any health events recorded (they'd return the initial score via
getScorebut won't appear in the Map). This is likely intentional for debugging purposes but worth noting.
177-192: Jitter utilities are correctly implemented.Both functions handle edge cases appropriately -
addJitterclamps to non-negative withMath.max(0, ...)and both return rounded integers. The randomness distribution is uniform which is appropriate for this use case.
215-253: LRU selection and hybrid account selection look correct.The filtering and sorting logic properly handles rate-limited, cooling-down, and unhealthy accounts. The defensive
sorted[0]?.index ?? nullhandles the edge case where the array might be empty after filtering.
283-349: Token bucket implementation is correct.The regeneration calculation and clamping logic are sound. Note that
getTokensreturns a floating-point value (regenerated tokens aren't floored), which is fine sincehasTokensandconsumeuse>=and<comparisons.
397-454: Priority queue selection has sound logic with good defensive coding.The weighted random selection properly handles edge cases (empty candidates, zero total score). The top 30% selection with min/max bounds (1-5) prevents both degenerate cases and excessive randomization.
src/plugin/rotation.test.ts (6)
1-17: Test setup imports and timer management look correct.Using
vi.useRealTimers()inbeforeEachensures clean state. The imports cover all public APIs being tested.
129-190: Time-based recovery tests properly mock Date.now.The tests correctly use
vi.spyOn(Date, 'now')and restore mocks withvi.restoreAllMocks(). The test at line 169-189 validates the floor behavior for partial recovery points.
203-216:getSnapshottest validates account tracking.The test confirms that accounts with recorded events appear in the snapshot with correct scores and failure counts.
316-354: Token regeneration tests have good coverage.Tests verify both regeneration over time and the cap at
maxTokens. The time mocking is handled correctly.
356-413: Jitter and delay tests use statistical approach.Running 100 iterations to verify bounds is appropriate for testing randomized functions. The tests cover edge cases like small base values with high jitter factors and min === max.
509-586:selectPriorityQueueAccounttests cover filtering scenarios but could test weighted selection.The tests verify filtering behavior well but don't explicitly test the weighted random distribution. Given the randomness, this is acceptable, but the test at lines 573-585 running 100 iterations provides some confidence.
src/plugin/config/schema.ts (2)
17-22: Strategy enum extended correctly.The new
priority-queueoption and updated documentation forhybridas the default are consistent with the implementation.
266-288: Health score and token bucket schemas have appropriate validation bounds.The constraints are sensible:
- Health scores bounded 0-100
- Penalties are negative (min -100/-50 to max 0)
- Token bucket allows up to 1000 tokens with regeneration 0.1-60/min
Both schemas are optional, which aligns with the initialization logic in plugin.ts that only initializes trackers when config is present.
src/plugin.ts (9)
42-42: Import statement for rotation module utilities.The imports are correctly scoped to what's needed in this file.
666-686: Tracker initialization maps config keys correctly.The snake_case config keys are properly mapped to camelCase tracker config. Note that trackers are only initialized when config blocks are present, otherwise defaults will be used via the lazy-init in
getHealthTracker()/getTokenTracker().
988-988: Health failures recorded on auth/project errors.The
recordFailurecalls are placed at appropriate error handling points. These degrade account health when authentication or project context fails.Also applies to: 1037-1037, 1063-1063
1188-1234: Token consumption tracking initialized correctly.The
tokenConsumedflag is properly scoped to the endpoint retry loop and only set when strategy ispriority-queue. This ensures tokens are only consumed once per request attempt.
1244-1248: Token refund on 429 is correctly implemented.The refund happens before rate-limit handling, and the flag is reset to prevent double-refunds.
1285-1286: Rate limit recorded to health tracker.This appropriately degrades account health when rate-limited, influencing future selection.
1440-1443: Success recorded after successful response.The health tracker records success only when
response.okis true, which correctly rewards healthy accounts.
1541-1546: Token refund on network/API errors is correctly handled.The refund only occurs if
tokenConsumedis true, preventing refunds for requests that never consumed a token. The flag is reset after refund.
1229-1234: Token consumption/refund appears balanced across all paths.Tokens are consumed once (line 1232), and refunded in two error scenarios:
- On 429 rate limit (line 1246)
- On network/API catch block (line 1544)
However, there's a potential edge case: if the request succeeds (line 1440-1442), the token is never "confirmed consumed" - it just stays consumed. This is the intended behavior (successful requests cost tokens).
Also applies to: 1244-1248, 1541-1546
src/plugin/accounts.ts (5)
5-5: Import statement correctly scoped.Imports the necessary rotation utilities and types.
287-313: Priority-queue selection correctly integrates with rotation module.The implementation:
- Builds metrics for all accounts with health scores from the tracker
- Delegates selection to
selectPriorityQueueAccount- Updates account state on selection
The comment on line 306 correctly notes that token consumption is deferred to plugin.ts request execution.
315-339: Hybrid strategy refactored to use health-aware LRU sorting.The logic correctly uses
sortByLruWithHealthfrom the rotation module, ensuring consistent filtering and sorting behavior. The null-safety check onsorted[0]andselectedis appropriate.
341-367: Fallback to sticky selection preserves backward compatibility.When priority-queue or hybrid strategies find no candidates (all filtered out), the code falls through to the sticky selection logic. This provides graceful degradation.
287-339: Both strategies updatelastUsedtimestamp consistently.Lines 307 and 333 both update
lastUsed = nowMs()on selection, ensuring the LRU tracking remains accurate across strategy changes.src/plugin/accounts.test.ts (3)
792-808: LGTM!The test correctly verifies that rate-limited accounts are skipped by the priority-queue strategy.
810-826: LGTM!The test correctly verifies that cooling-down accounts are skipped by the priority-queue strategy.
846-865: LGTM!The test properly verifies that the
lastUsedtimestamp andcurrentAccountIndexByFamilyare updated when an account is selected via the priority-queue strategy.
| it("falls back to sticky when all accounts unavailable", () => { | ||
| vi.useFakeTimers(); | ||
| vi.setSystemTime(new Date(0)); | ||
|
|
||
| const stored: AccountStorageV3 = { | ||
| version: 3, | ||
| accounts: [ | ||
| { refreshToken: "r1", projectId: "p1", addedAt: 1, lastUsed: 0 }, | ||
| ], | ||
| activeIndex: 0, | ||
| }; | ||
|
|
||
| const manager = new AccountManager(undefined, stored); | ||
|
|
||
| const selected = manager.getCurrentOrNextForFamily("claude", null, "priority-queue"); | ||
| expect(selected?.index).toBe(0); | ||
| }); |
There was a problem hiding this comment.
Test doesn't actually verify fallback behavior.
The test name states "falls back to sticky when all accounts unavailable," but the single account is never marked as unavailable (not rate-limited, not cooling down). This tests normal selection of an available account, not the fallback path.
To properly test fallback behavior, mark the account as rate-limited or cooling down, then verify the sticky fallback logic activates (or returns null if that's the expected behavior when all accounts are truly unavailable).
Suggested fix to actually test fallback
it("falls back to sticky when all accounts unavailable", () => {
vi.useFakeTimers();
vi.setSystemTime(new Date(0));
const stored: AccountStorageV3 = {
version: 3,
accounts: [
{ refreshToken: "r1", projectId: "p1", addedAt: 1, lastUsed: 0 },
+ { refreshToken: "r2", projectId: "p2", addedAt: 1, lastUsed: 0 },
],
activeIndex: 0,
};
const manager = new AccountManager(undefined, stored);
+ const accounts = manager.getAccounts();
+
+ // Mark all accounts as rate-limited to trigger fallback
+ manager.markRateLimited(accounts[0]!, 60000, "claude");
+ manager.markRateLimited(accounts[1]!, 60000, "claude");
const selected = manager.getCurrentOrNextForFamily("claude", null, "priority-queue");
- expect(selected?.index).toBe(0);
+ // Verify fallback behavior - either returns null or falls back to sticky
+ expect(selected).toBeNull();
});🤖 Prompt for AI Agents
In @src/plugin/accounts.test.ts around lines 828 - 844, The test named "falls
back to sticky when all accounts unavailable" never marks the single account as
unavailable, so it doesn't exercise the fallback; modify the stored
AccountStorageV3 used to construct AccountManager so the sole account is
actually unavailable at time 0 (e.g., add a rate-limiting/cooldown field such as
rateLimitedUntil or cooldownUntil set > 0, e.g., 1000, or use the
AccountManager/API helper that marks accounts unavailable) and then call
manager.getCurrentOrNextForFamily("claude", null, "priority-queue") and assert
the fallback behavior (sticky index returned or null depending on expected
semantics).
|
merged, thanks you for the help |
|
this wasn't fully tested locally. i opened it hoping for a manual code review 😅 ah i noticed it's on dev. ok, will patch it if any bugs come forth. |
deployed on beta.3 You can test the release beta version, I'm testing it on my side as well |
2 participants
feature: add more rotation to account
added more rotation techniques for better account rotation
why
Summary by CodeRabbit
New Features
Improvements
Tests
✏️ Tip: You can customize this high-level summary in your review settings.