fix(ai): route Z.AI API keys by usage#5275
Conversation
- Added Z.AI credential ranking so API-key accounts use request-window usage before dispatch. - Extended API-key selection to prefetch usage reports and block exhausted siblings until reset. - Covered the regression with an AuthStorage Z.AI API-key selection test. Fixes #5265
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2010f01cd5
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| return { | ||
| type: "api_key", | ||
| apiKey: credential.key, |
There was a problem hiding this comment.
When an API-key credential stores a config reference, the same getApiKey path later resolves it through #configValueResolver, but the new usage-ranking path passes the raw stored string to the usage provider here. In coding-agent configurations that use env/!command values, the Z.AI quota probe runs with Authorization: ZAI_API_KEY or !cmd, returns no report, and the exhausted referenced key is ranked as unknown even though the actual request uses the resolved secret. Resolve the key before building the UsageCredential, or skip usage ranking when it cannot be resolved.
Useful? React with 👍 / 👎.
| function rankZaiRequestLimits(report: UsageReport): UsageLimit[] { | ||
| const requestLimits = report.limits.filter(limit => limit.id.startsWith("zai:requests:")); | ||
| const limits = requestLimits.length > 0 ? requestLimits : report.limits; |
There was a problem hiding this comment.
When the quota payload includes the Web Search / Reader / Zread TIME_LIMIT that this provider labels separately, this prefix match treats that feature counter as the general request window. Because zaiRankingStrategy does not scope limits, #rankApiKeySelections can mark the entire credential blocked when only that feature quota is exhausted, causing ordinary Z.AI model requests to avoid an otherwise healthy key. Filter or scope those feature-specific limits separately before using them for credential-wide ranking/blocking.
Useful? React with 👍 / 👎.
- Resolved stored API-key config references before Z.AI usage ranking probes. - Scoped Z.AI ranking to model request/token quotas so feature-only exhaustion does not block model routing. - Added regressions for config references and feature-specific quota rows. Fixes #5265
|
Resolved stored API-key config references before usage-ranking probes so Z.AI quota checks use the same secret as request dispatch. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8d4e4fde57
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| const featureLimit = isZaiFeatureRequestLimit(parsed); | ||
| limits.push({ | ||
| id: `zai:requests:${window.id}`, | ||
| id: featureLimit ? `zai:features:web-search-reader-zread:${window.id}` : `zai:requests:${window.id}`, |
There was a problem hiding this comment.
With the existing fixture in packages/ai/test/zai-usage.test.ts, this branch now emits zai:features:web-search-reader-zread:1mo for the Web Search / Reader / Zread limit, while that committed test still asserts the first id is zai:requests:1mo. Once the package test suite runs with workspace aliases available, the parser contract test will fail unless the expectation is updated to match this intentional feature-quota split (or this id change is reverted).
Useful? React with 👍 / 👎.
- Updated the Z.AI usage parser fixture to expect feature-specific Web Search / Reader / Zread quota ids. - Asserted the feature quota is not marked as a shared model-routing quota. Fixes #5265
|
Updated |
1 participant
Repro
A focused AuthStorage reproduction configured two stored Z.AI login API keys with a custom usage provider: the first key reported an exhausted shared
zai:requests:5hwindow and the second reported request headroom. Before the fix,bun test tmp-repro/zai-api-key-ranking-repro.test.tsfailed becausegetApiKey("zai")returned"zai-exhausted"instead of"zai-healthy".Cause
packages/ai/src/auth-storage.tsonly appliedCredentialRankingStrategyusage reports in the OAuth path (#rankOAuthSelections/#resolveOAuthSelection). Stored API-key credentials still used the synchronous round-robin selector, andpackages/ai/src/usage/zai.tsexported onlyzaiUsageProvider, so Z.AI usage data populated/usagebut never affected request-time credential selection.Fix
Verification
bun test packages/ai/test/auth-storage-zai-api-key-selection.test.ts packages/ai/test/auth-storage-codex-selection.test.ts packages/ai/test/auth-storage-antigravity-selection.test.tspassed: 47 tests, 112 expects.bun checkpassed. Fixes #5265.