Add rotation strategies + key health engine for multi-key providers - #1191
Add rotation strategies + key health engine for multi-key providers#1191FiredMosquito831 wants to merge 5 commits into
Conversation
…tion Providers now accept a comma-separated list of keys in their credential env var (e.g. NVIDIA_NIM_API_KEY=k1,k2) plus an optional *_ROTATION policy (single|round_robin|on_error). A RotatingProvider runtime wrapper builds one sub-provider per key, spreads requests across keys (round_robin), and fails over to the next key when the first attempt errors before any output is streamed (on_error, with per-key backoff). Plain 400s are not rotated. Admin UI gains a per-provider Key Rotation select field.
Every provider credential field in the admin dashboard now has a Manage keys panel: masked key list, add-key input, and per-key removal.
Backed by /admin/api/credentials/{env_key}/keys endpoints (list/add/delete) that persist through the existing admin apply pipeline, so changes take effect immediately without restart. Keys are always masked in API responses; process-env-sourced credentials are read-only.
Backports the multi-provider proxy's rotation system: - Policies: single, round_robin, least_used, failover (on_error kept as an alias of failover), selectable per provider via *_ROTATION. - Health engine per key: tiered cooldowns (10s-120s), circuit breaker after 3 consecutive failures with single-probe half-open recovery, escalating auth lockouts for 401/403 (5min-1h-24h), request/failure metrics, and a 503-style error with retry wait when every key is benched. - Dashboard: rotation policy select next to Manage keys, and live per-key health badges (state + 'back in Ns' + request/failure counts) in the key manager, fed by the keys endpoint's new health field.
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if index in attempted: | ||
| remaining = [ | ||
| i for i in range(len(self._providers)) if i not in attempted | ||
| ] | ||
| if not remaining: | ||
| break | ||
| index = remaining[0] | ||
| attempted.add(index) |
There was a problem hiding this comment.
With multiple configured keys and the default single policy, acquire() always returns key 0. After a retryable failure, this branch sees key 0 in attempted and directly selects key 1, so a policy documented as using only the first key silently performs failover.
Artifacts
Repro: executable focused async pytest harness
- Contains supporting evidence from the run (text/x-python; charset=utf-8).
Repro: verbose pytest output showing retryable key 0 failure and key 1 failover
- Keeps the command output available without making the summary code-heavy.
| try: | ||
| yield first_chunk | ||
| async for chunk in iterator: | ||
| yield chunk | ||
| finally: | ||
| await maybe_await_aclose(iterator) | ||
| await self._state.report_success(index) |
There was a problem hiding this comment.
When a half-open stream is cancelled, fails after its first chunk, or raises during cleanup, control leaves through finally without calling either health reporter. The key remains HALF_OPEN with is_probing=True, so every later acquisition skips it permanently and the provider can remain unavailable after its cooldown has expired.
Artifacts
Repro: focused async harness for a post-first-chunk HALF_OPEN probe failure
- Contains supporting evidence from the run (text/x-python; charset=utf-8).
Repro: verbose pytest output showing the leaked reservation and subsequent acquisition failure
- Keeps the command output available without making the summary code-heavy.
| except StopAsyncIteration: | ||
| return | ||
| except Exception as error: | ||
| last_error = error |
There was a problem hiding this comment.
If a sub-provider ends before producing its first chunk, this branch returns without reporting success or failure. The request therefore completes as an empty response instead of surfacing an upstream failure or trying another eligible key, and a half-open key also remains stuck in its probing state.
Artifacts
Repro: focused executable async RotatingProvider test harness
- Contains supporting evidence from the run (text/x-python; charset=utf-8).
Stack trace captured during the T-Rex run
- Keeps the raw stack trace available without making the summary code-heavy.
| snapshots = provider.key_health() | ||
| for i in range(min(len(keys), len(snapshots))): | ||
| health[i] = snapshots[i] |
There was a problem hiding this comment.
After deleting a non-final key, the persisted key list changes before an existing provider generation is necessarily gone. Deleting B from [A, B, C] leaves [A, C], but this index copy maps B's old snapshot onto C, so the dashboard shows the deleted credential's health and counters beside the wrong surviving key.
| :class:`RotatingProvider` that applies the configured rotation policy. | ||
| """ | ||
| descriptor = PROVIDER_CATALOG.get(provider_id) | ||
| if descriptor is None: | ||
| raise UnknownProviderError.for_provider(provider_id, PROVIDER_CATALOG) | ||
|
|
||
| config = build_provider_config(descriptor, settings) |
There was a problem hiding this comment.
This adds a user-visible provider and configuration capability, but the change does not update pyproject.toml or regenerate uv.lock. The repository's release rules require both updates in the same production change, with a minor version bump for a new capability.
Context Used: CLAUDE.md (source)
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
Artifacts
Repro: executable base/head release metadata harness
- Contains supporting evidence from the run (text/x-python; charset=utf-8).
Repro: failing base/head version and lockfile comparison
- Keeps the command output available without making the summary code-heavy.
Repro: head uv lock consistency command output
- Keeps the command output available without making the summary code-heavy.
|
This is a good idea, ive thought of doing this for nvidia nim keys, by editing the nvidia nim default base to my multi key proxy. It works pretty good. |
2 participants
Split out from the too-large #1164 — this PR contains only the rotation strategies and health engine.
Depends on / includes #1190 (multi-key support) — only the last commit (
feat(rotation): three rotation strategies with full key health engine) is new; please merge after #1190.What it adds
round_robin(spread across healthy keys),least_used(fewest requests first),failover(stick to one key until it fails;on_errorkept as an alias). Selected via*_ROTATIONor the dashboard select next to Manage keys.healthfield on the keys endpoint (read from the liveRotatingProvider).1167 tests pass, including strategy selection, cooldown/lockout escalation, circuit recovery, and all-benched behavior.
Greptile Summary
This PR adds credential rotation and per-key health tracking for multi-key providers. The main changes are:
Confidence Score: 4/5
The streaming rotation lifecycle can use unintended credentials and permanently strand half-open keys. These paths need fixes before merging.
The default single policy can retry with a secondary key. Cancellation, mid-stream failure, and cleanup failure can leave a probe reserved forever. An empty upstream stream is silently treated as a completed response. Key deletion can temporarily display another credential's health. The required package version and lockfile updates are missing.
src/free_claude_code/providers/runtime/rotating.py, src/free_claude_code/api/admin_routes.py, pyproject.toml, and uv.lock
What T-Rex did
Important Files Changed
Sequence Diagram
%%{init: {'theme': 'neutral'}}%% sequenceDiagram participant Client participant Wrapper as RotatingProvider participant Health as RotationState participant KeyA as Provider A participant KeyB as Provider B Client->>Wrapper: stream request Wrapper->>Health: acquire() Health-->>Wrapper: key index Wrapper->>KeyA: start stream alt failure before first chunk KeyA-->>Wrapper: retryable error Wrapper->>Health: report_failure() Wrapper->>Health: acquire() Health-->>Wrapper: next key Wrapper->>KeyB: retry request else stream starts KeyA-->>Wrapper: chunks Wrapper-->>Client: chunks Wrapper->>Health: report_success() else all keys benched Health-->>Wrapper: -1 Wrapper-->>Client: unavailable with retry delay end%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%% sequenceDiagram participant Client participant Wrapper as RotatingProvider participant Health as RotationState participant KeyA as Provider A participant KeyB as Provider B Client->>Wrapper: stream request Wrapper->>Health: acquire() Health-->>Wrapper: key index Wrapper->>KeyA: start stream alt failure before first chunk KeyA-->>Wrapper: retryable error Wrapper->>Health: report_failure() Wrapper->>Health: acquire() Health-->>Wrapper: next key Wrapper->>KeyB: retry request else stream starts KeyA-->>Wrapper: chunks Wrapper-->>Client: chunks Wrapper->>Health: report_success() else all keys benched Health-->>Wrapper: -1 Wrapper-->>Client: unavailable with retry delay endReviews (1): Last reviewed commit: "feat(rotation): three rotation strategie..." | Re-trigger Greptile
Context used: