From 5512269185f1bcd2434e8d36d110e9a0294beee2 Mon Sep 17 00:00:00 2001 From: bitkyc08-arch Date: Sat, 29 Aug 2026 12:30:03 +0900 Subject: [PATCH 1/6] docs(devlog): plan Kiro quota display and quota-aware account pool Five audit rounds against an independent reviewer; 18 blockers folded. --- .../000_research_problem_and_state.md | 65 ++++++++ .../001_research_upstream_wire_contract.md | 98 ++++++++++++ .../002_research_kirolb_headtohead.md | 95 +++++++++++ .../003_research_pr_reconciliation.md | 46 ++++++ .../010_kiro_usage_fetcher.md | 149 ++++++++++++++++++ .../020_per_account_quota_wiring.md | 122 ++++++++++++++ .../030_quota_aware_pool_selection.md | 118 ++++++++++++++ .../040_surfaces_cli_gui_docs.md | 74 +++++++++ .../050_verification_and_delivery.md | 71 +++++++++ .../060_audit_round1_amendments.md | 143 +++++++++++++++++ .../070_audit_round2_amendments.md | 130 +++++++++++++++ 11 files changed, 1111 insertions(+) create mode 100644 devlog/_plan/260829_kiro_quota_pool/000_research_problem_and_state.md create mode 100644 devlog/_plan/260829_kiro_quota_pool/001_research_upstream_wire_contract.md create mode 100644 devlog/_plan/260829_kiro_quota_pool/002_research_kirolb_headtohead.md create mode 100644 devlog/_plan/260829_kiro_quota_pool/003_research_pr_reconciliation.md create mode 100644 devlog/_plan/260829_kiro_quota_pool/010_kiro_usage_fetcher.md create mode 100644 devlog/_plan/260829_kiro_quota_pool/020_per_account_quota_wiring.md create mode 100644 devlog/_plan/260829_kiro_quota_pool/030_quota_aware_pool_selection.md create mode 100644 devlog/_plan/260829_kiro_quota_pool/040_surfaces_cli_gui_docs.md create mode 100644 devlog/_plan/260829_kiro_quota_pool/050_verification_and_delivery.md create mode 100644 devlog/_plan/260829_kiro_quota_pool/060_audit_round1_amendments.md create mode 100644 devlog/_plan/260829_kiro_quota_pool/070_audit_round2_amendments.md diff --git a/devlog/_plan/260829_kiro_quota_pool/000_research_problem_and_state.md b/devlog/_plan/260829_kiro_quota_pool/000_research_problem_and_state.md new file mode 100644 index 0000000000..468ac62dee --- /dev/null +++ b/devlog/_plan/260829_kiro_quota_pool/000_research_problem_and_state.md @@ -0,0 +1,65 @@ +# 000 — Kiro quota + pool: problem statement and current state + +Unit: `devlog/_plan/260829_kiro_quota_pool/` +Opened: 2026-08-29 +Work classes: C3 (quota fetcher, pool selection), C2 (surfaces, docs) + +## The ask + +Two capabilities, plus one reconciliation: + +1. **Quota display for Kiro.** Every other major OAuth provider in this proxy reports + remaining capacity; Kiro reports nothing. `rg -n -i kiro src/providers/quota.ts` + returns zero matches today. +2. **Pool-based automatic loading.** Multiple Kiro accounts should load into a pool and + be selected automatically, preferring accounts that still have quota. +3. **429 PR reconciliation.** Confirm which of the previously-submitted 429 failover PRs + actually landed on `dev`, and fold the landed behaviour into the Kiro path. + +The comparison target is [minpeter/kiro-lb](https://github.com/minpeter/kiro-lb), an +AGPL-3.0 Python/FastAPI Kiro gateway with multi-account load balancing and an operations +dashboard. **We study its behaviour; we copy none of its code.** AGPL-3.0 is incompatible +with this repository's licensing, so every line here is written from the wire contract and +from our own existing seams. + +## What we already have (verified 2026-08-29 against origin/dev 124a2b148) + +Kiro is further along than it looks: + +- **Multi-account storage exists.** Kiro credentials live in the generic multiauth store + with `activeAccountId` plus an `accounts[]` array; each entry carries its own + `credential.kiro` routing metadata (`profileArn`, `ssoRegion`, `apiRegion`, + `clientId`, `clientSecret`). See `src/oauth/types.ts:14` and `src/oauth/store.ts:264`. +- **429 rotation already covers Kiro.** `isGenericFailoverProvider` excludes only + `openai` and `anthropic`, so any OAuth provider — Kiro included — rotates on a 429 + once two non-reauth accounts are present (`src/oauth/generic-account-failover.ts:44`, + `:81`, `:114`). +- **Rotation carries Kiro's routing metadata.** `applyFailoverSnapshot` reassigns + `parsed._kiroAuthContext` from the rotated snapshot, so a rotated bearer travels with + its own profile ARN and regions (`src/server/responses/core.ts:3063`). PR #2841 + (merged `5a829b7e9`) hardened exactly this class of bug for Copilot origins. +- **A per-account quota seam exists.** `supportsPerAccountQuota`, + `fetchProviderAccountQuotas`, the per-account TTL cache, generation reconciliation and + the GUI's `accounts[].quota` field are all built — but wired to Anthropic only + (`src/providers/quota.ts:1453`, `:1572`, `:1619`). + +## What is actually missing + +| Gap | Evidence | +| --- | --- | +| No Kiro quota fetcher at all | `rg -i kiro src/providers/quota.ts` → 0 matches | +| `supportsPerAccountQuota("kiro")` is false, and a test locks it | `tests/provider-account-quota.test.ts:204` | +| Rotation is quota-blind: it walks stored order, skipping cooled accounts | `src/oauth/generic-account-failover.ts:157` | +| Rotation only reacts to a 429 it already suffered; a known-exhausted account is still tried first | same | +| CLI copy calls Kiro a "single login slot", contradicting the shipped multiauth add-account flow | `src/cli/account.ts:28`, `:215` | + +That last row matters more than it reads: the feature exists and the product tells the +user it does not. + +## Non-goals for this unit + +- No AGPL code, text, or structure copied from kiro-lb. +- No changes to the Codex or Anthropic pools; both are excluded from generic failover by + design and own their own affinity/probe semantics. +- No `src/lab/` involvement — the core boundary test forbids it. +- No release promotion to `main`. diff --git a/devlog/_plan/260829_kiro_quota_pool/001_research_upstream_wire_contract.md b/devlog/_plan/260829_kiro_quota_pool/001_research_upstream_wire_contract.md new file mode 100644 index 0000000000..4be05774a7 --- /dev/null +++ b/devlog/_plan/260829_kiro_quota_pool/001_research_upstream_wire_contract.md @@ -0,0 +1,98 @@ +# 001 — The Kiro usage-limits wire contract + +Source of truth for this document: the observed Kiro CLI 2.19.x contract, cross-checked +against two independent third-party implementations. This operation is **undocumented** by +AWS; treat every field as best-effort and never fail a request because it changed. + +## The operation + +```http +POST /?origin=AI_EDITOR&isEmailRequired=true&profileArn= +Host: management..kiro.dev +Authorization: Bearer +Content-Type: application/x-amz-json-1.0 +Accept: application/json +x-amz-target: AmazonCodeWhispererService.GetUsageLimits +``` + +```json +{ "origin": "AI_EDITOR", "isEmailRequired": true, "profileArn": "" } +``` + +Two details that look like mistakes and are not: + +- **The modeled arguments appear in both the query string and the JSON body.** That is the + observed CLI behaviour. Reproduce it rather than "simplifying" it; an AWS JSON-RPC + front door that also reads query parameters will accept both, and we have no way to test + which one it actually honours. +- **The host is `management.`, not `runtime.`** Generation goes to + `runtime.{region}.kiro.dev`; usage goes to a different subdomain. Our provider + `baseUrl` is the runtime host, so the quota fetcher must derive the management host + rather than reuse `baseUrl`. + +### Region resolution + +The profile ARN is authoritative when present: `arn:aws:codewhisperer:::profile/` +— take field 3. Otherwise fall back to the account's stored `apiRegion`, then `ssoRegion`, +then `us-east-1`. We already have all three on the credential +(`src/oauth/kiro-credentials.ts:285`) and a resolver in `src/oauth/kiro.ts:445`. + +## Response shape + +```json +{ + "subscriptionInfo": { "subscriptionTitle": "KIRO PRO", "type": "Q_DEVELOPER_..." }, + "overageConfiguration": { "overageStatus": "ENABLED|DISABLED" }, + "usageBreakdownList": [ + { + "resourceType": "AGENTIC_REQUEST|CREDIT|...", + "currentUsageWithPrecision": 147.82, + "currentUsage": 147, + "usageLimitWithPrecision": 1000.0, + "usageLimit": 1000, + "currentOveragesWithPrecision": 0.0, + "overageRate": 0.04, + "unit": "CREDITS|INVOCATIONS", + "freeTrialInfo": { "freeTrialStatus": "ACTIVE", "usageLimitWithPrecision": 500.0 } + } + ], + "userInfo": { "email": "...", "userId": "..." }, + "nextDateReset": 1785542400.0, + "daysUntilReset": 3 +} +``` + +### Field handling rules + +1. **Prefer `*WithPrecision`.** Kiro meters to 0.01 credit; the integer fields round + 695.17 down to 695. Fall back to the integer only when precision is absent. +2. **Select the breakdown by `resourceType`, never by index.** Take `AGENTIC_REQUEST` + first, then `CREDIT`; if neither exists, report unknown rather than guessing. Taking + `[0]` means an upstream reorder silently reweights routing against an unrelated pool. +3. **`currentUsage > usageLimit` is not necessarily exhaustion** when + `overageStatus` is `ENABLED` — enterprise accounts keep serving past the included + limit. Percent must clamp for display, but exhaustion must consult overage status. +4. **`userInfo.email` is a personal identifier.** We request `isEmailRequired` because + the response shape is the observed contract, but the email must never be logged and + never persisted into quota state. Our account rows already carry a masked identity. +5. **`freeTrialInfo` is a separate pool.** kiro-lb ignores it, which understates the + usable balance for trial users. We record it as its own window. + +## Cadence + +Kiro's own pricing page says usage data refreshes "at least every 5 minutes", so polling +faster buys nothing. The existing provider cache TTL is 5 minutes +(`src/providers/quota.ts:37`) and the per-account TTL governs account rows; both are +already at or above the useful floor. No new timer is needed — the existing pull-on-demand +plus TTL is the right shape, and it means an idle proxy makes zero usage calls. + +## Auth-mode caveats + +- **Enterprise / IdC accounts** carry a real profile ARN → send it. +- **AWS Builder ID** has no account-owned profile. We already resolve a *request-scoped* + service profile (`src/adapters/kiro-constants.ts:16`, applied at + `src/oauth/kiro.ts:508`) which must never be persisted as identity. For usage, send + the request-scoped value the same way the generation path does. +- **`ksk_` API keys** are not OAuth accounts, have no refresh identity, and there is no + evidence `GetUsageLimits` accepts the `tokentype: API_KEY` contract. Out of scope: + report unknown. diff --git a/devlog/_plan/260829_kiro_quota_pool/002_research_kirolb_headtohead.md b/devlog/_plan/260829_kiro_quota_pool/002_research_kirolb_headtohead.md new file mode 100644 index 0000000000..e12815d99e --- /dev/null +++ b/devlog/_plan/260829_kiro_quota_pool/002_research_kirolb_headtohead.md @@ -0,0 +1,95 @@ +# 002 — kiro-lb: what it does, and where it is beatable + +Reference clone read read-only at `/tmp/kirolb.*/repo` (commit `474df2b` / `b2ec34d`, +2026-08-26). AGPL-3.0. **Behaviour studied, no code reused.** + +kiro-lb is a competent, purpose-built gateway. It is single-provider by design, and that +focus buys it a real dashboard and a working weighted router. An honest comparison has to +start by saying what it does well, because those are the bars we must clear. + +## What it does well + +| Capability | Where | +| --- | --- | +| Reads real upstream usage per account | `kiro/usage.py:43-70` | +| Quota-weighted routing (exponential race, weight = remaining fraction) | `kiro/account_manager.py:1183-1208` | +| Distinct exclusion states with distinct timers | `kiro/account_manager.py:109-170` | +| Monthly-quota quarantine aligned to `nextDateReset` (6h floor, 32d cap) | `kiro/config.py:457-479` | +| Suspension (403) and credential-death (refresh 400/401) as separate states | `kiro/kiro_errors.py:30-51` | +| Persisted quota rows survive restart and seed routing | `kiro/store.py:206-289` | +| Cross-process refresh lease | `kiro/store.py:172-197` | + +## Where it is beatable — with citations + +These are the gaps the reviewers verified in its source, not marketing points. + +1. **Weighted routing is model-blind.** `_select_account()` builds candidates from + `list(self._accounts)` and never consults its own `model` argument or + `_model_to_accounts` (`kiro/account_manager.py:1259-1276`). An account whose plan + cannot serve the requested model is discovered by *failing a request*. +2. **Headroom is stale for up to the full poll interval.** Successful requests do not + decrement local headroom (`kiro/account_manager.py:1369-1432`); only a poll updates it + (default 900s, `kiro/config.py:533`). A hot account keeps its high weight for ~15 + minutes while it burns through its balance. +3. **`USAGE_REFRESH_INTERVAL_SECONDS=0` does not disable polling.** The comment says it + does; the loop enforces `max(interval, 60)` and startup always polls + (`main.py:342-359`). Zero means *every minute*. +4. **Bulk polling is sequential with a fresh 20s client per account.** `refresh_all_account_usage()` + awaits one at a time (`kiro/dashboard.py:806-821`) and `usage.py:64-70` constructs a new + `AsyncClient` per call. An N-account pool of dead accounts costs ~N × 20s per pass. +5. **A concurrent account deletion can abort an entire refresh pass.** The loop re-indexes + `manager._accounts[account_id]` after an await, outside the lock + (`kiro/dashboard.py:806-815`) — a `KeyError` there ends the pass, so later accounts + never refresh. +6. **Breakdown selection falls back to index 0.** If no `AGENTIC_REQUEST` entry exists, + the first entry becomes the routing signal (`kiro/usage.py:74-78`). An upstream + addition silently reweights the pool on an unrelated resource. +7. **`freeTrialInfo` is dropped**, understating usable balance for trial accounts + (`kiro/usage.py:97-111`). +8. **No absolute reset timestamp in the UI**, only a coarse relative duration, and only + for excluded accounts (`frontend/src/features/dashboard/quota-display.ts:18-33`). + `unit` and `overageRate` are fetched then discarded (`kiro/dashboard.py:613-627`). +9. **Pool loading is not automatic discovery.** No standard cache path (`~/.aws/sso/cache`, + `~/.local/share/kiro-cli`) is scanned unless already registered as a source; scanning + happens at startup/handoff only, with no watcher (`kiro/account_manager.py:407-505`). + The README tells users to add accounts through dashboard device login. +10. **Suspension/auth-death prose contradicts the code**: both expire automatically after + 24h (`kiro/config.py:481-494`) although comments claim only support or re-login clears + them. +11. **Refresh-lease waiting has no deadline** — contenders poll every 50ms forever + (`kiro/auth.py:965-980`). +12. **Device-login flows are process memory only** (`kiro/device_login.py:97-113`); a + restart mid-approval invalidates the login. +13. **Dashboard copy is inaccurate**: it says "only the refresh token is stored" while + `internal_credentials()` persists access token, refresh token, expiry, region and + client secret (`kiro/device_login.py:341-366`). + +## What "better" must mean for us + +Beating it is not "we also show a number". Our structural advantages have to be real, and +stated no wider than what we ship: + +- **Quota-aware recovery ordering.** On a 429 we rotate toward the account with the most + known headroom instead of walking the roster blind. This is *recovery* ordering, not + pre-dispatch selection — see the scope note below. +- **Parallel, deadline-bounded probing** with per-account failure isolation. +- **Explicit unknown state** everywhere — never present a failed probe as "0% used". +- **No new background timer**: pull-on-demand plus TTL, so an idle proxy is silent. +- **Correct pool semantics we already own**: request-local rotation that never mutates the + operator's `activeAccountId`, and per-account routing metadata that travels with its + own bearer. + +## Scope honesty (amended after audit round 1) + +Two advantages were claimed here and have been withdrawn, because the design could not +back them: + +- **Pre-request, model-aware selection.** Our rotation hook runs only inside the 429 + branch (`src/server/responses/core.ts:5574`), so nothing in this unit chooses an + account *before* the first request, and no model is passed to the ranker. kiro-lb's + weighted router genuinely is pre-request (though it is model-blind). Deferred to a + follow-up work-phase; not claimed here. +- **Live decrement between polls.** `ProviderQuota` stores percent, not absolute + used/limit, and Kiro meters fractional credits — one turn is not one credit. Any local + decrement would be invented data. kiro-lb's 15-minute staleness gap is real; we do not + currently close it, we only poll on demand with a shorter TTL. diff --git a/devlog/_plan/260829_kiro_quota_pool/003_research_pr_reconciliation.md b/devlog/_plan/260829_kiro_quota_pool/003_research_pr_reconciliation.md new file mode 100644 index 0000000000..f131da5a66 --- /dev/null +++ b/devlog/_plan/260829_kiro_quota_pool/003_research_pr_reconciliation.md @@ -0,0 +1,46 @@ +# 003 — 429 PR reconciliation (state as of 2026-08-29) + +`origin/dev` and GitHub `dev` both at `124a2b1487996f8a8ebb2067b22c9e758fa6016f`. + +"Landed" below means the *behaviour* is on `dev`. Squash merges do not preserve the PR +head SHA in ancestry, so head-containment is the wrong test for all but one of these. + +| PR | Subject | State | Landed as | +| --- | --- | --- | --- | +| #2590 | generic multi-account 429 failover (#2568a) | MERGED | `816f3a159` | +| #2607 | rotate generic OAuth accounts on 429 in sidecars | MERGED | `87250870c` | +| #2608 | cursor adapter-event 429 rotation | MERGED | `6b508d5a8` | +| #2640 | activate failover on account presence (#2568d) | MERGED | `8bfac7146` | +| #2841 | bind a rotated OAuth bearer to its own Copilot origin | MERGED | `5a829b7e9` | +| #927 | compact: alternate account on pool 429/402 | MERGED | `87c479006` (head in ancestry) | +| #2573 | antigravity quota exhaustion spelling | MERGED | `bfe2cb5a1` | +| #2745 | rebind credential identity on every OAuth 429 rotation | CLOSED | superseded by #2807 → #2841 | +| #2807 | same, v2 | CLOSED | superseded by #2841 | + +## The nuance on #2745 / #2807 + +Their *security outcome* landed; their *complete diff* did not. #2841 fixed all four +snapshot/origin read sites and added stronger coverage, but the broader refactor those PRs +proposed — relocating `sentOAuthSnapshot`, replay identity, and Cursor cleanup into +`applyFailoverSnapshot` — was not adopted. Neither branch needs rebasing; the accepted +requirement is represented on `dev`. + +What this means for **this** unit: the credential/identity-pairing invariant is already +enforced for Copilot origins and for Kiro's `_kiroAuthContext` +(`src/server/responses/core.ts:3050-3063`). Our Kiro work must not regress it, and our +regression test must prove a rotated Kiro bearer never travels with another account's +profile ARN. + +## Still open and relevant + +- **#2783** (quota-reset detection) — OPEN, CI green at its recorded head, but + `CONFLICTING`/`DIRTY` against current `dev`. It is a *different* unit (usage-window + reset detection and notification). Out of scope here; it needs its own rebase pass. + +Not relevant: #2729, #1704, #150, #138 are closed and superseded or dormant. + +## Conclusion + +There is no unmerged 429 work to land. The reconciliation answer is "all merged except two +that were deliberately superseded by #2841", and the follow-up action is to *preserve* that +invariant while adding Kiro quota, not to re-open old branches. diff --git a/devlog/_plan/260829_kiro_quota_pool/010_kiro_usage_fetcher.md b/devlog/_plan/260829_kiro_quota_pool/010_kiro_usage_fetcher.md new file mode 100644 index 0000000000..31554d9250 --- /dev/null +++ b/devlog/_plan/260829_kiro_quota_pool/010_kiro_usage_fetcher.md @@ -0,0 +1,149 @@ +# 010 — Phase 1: the Kiro usage fetcher (foundation) + +Work class: C3. Depends on: nothing. Everything else in this unit consumes its output. + +## Goal + +One function that turns a Kiro account's credential into a `ProviderQuota`, or into an +honest `null`. + +## New file: `src/providers/kiro-usage.ts` + +A separate module, not an addition to the already-2355-line `quota.ts`. It owns the wire +contract from doc `001` plus the Kiro usage-state seams from `070`. Complete export +surface: + +- `kiroUsageManagementUrl(region)` — host construction (exported for tests). +- `fetchKiroUsageSnapshot(ctx)` — the probe. +- `kiroUsageContextForAccount(accountId)` — credential → context (`020`). +- `commitKiroAccountUsageState(key, state)` — called by `quota.ts` inside its existing + commit guard. +- `getKiroAccountExhaustion(provider, accountId)` — the freshness-checked reader + consumed by generic failover (`030`). +- `clearKiroAccountUsageState(provider?)` and `reconcileKiroAccountUsageState(liveKeys)` + — called from `clearAccountQuotaCache` and `reconcileProviderAccountQuotaRows`. + +```ts +export interface KiroUsageContext { + accountId: string; // keys the usage-state map; see 070 + access: string; + profileArn?: string; // request-scoped; Builder ID fallback allowed + apiRegion?: string; + ssoRegion?: string; +} + +export interface KiroUsageSnapshot { + quota: ProviderQuota; + exhausted: boolean; // limit reached AND overage not enabled + nextResetAt?: number; // epoch ms +} + +export function kiroUsageManagementUrl(region: string): string; +export async function fetchKiroUsageSnapshot(ctx: KiroUsageContext): Promise; +``` + +`subscriptionTitle` and `overageEnabled` are deliberately absent (audit round 1, +blocker 2): the first had no consumer and nowhere to render, the second is only an input +to `exhausted` and stays a local inside the parser. + +**Prerequisite extraction (audit round 2, blockers 2 and 4).** Phase 1 first splits two +neutral modules out of `quota.ts`, so this module never imports `quota.ts` and the +dependency edge stays one-directional: + +- `src/providers/quota-types.ts` — `ProviderQuota`, `ProviderQuotaWindow`, + `ProviderQuotaCreditsUsd`, `ProviderQuotaReport`. +- `src/providers/quota-wire.ts` — `REQUEST_TIMEOUT_MS`, `normalizePercent`, + `normalizeResetAt`, `toFiniteNumber`, `asRecord`, `readQuotaJson`. + +Both are pure moves; `tests/provider-quota.test.ts` (107 pass today) proves inertness. + +### Region resolution + +```ts +const REGION_PATTERN = /^[a-z0-9-]{1,32}$/; +const safeRegion = (v: string | undefined): string | undefined => + v && REGION_PATTERN.test(v) ? v : undefined; + +function usageRegion(ctx: KiroUsageContext): string { + return safeRegion(ctx.profileArn?.split(":")[3]) + ?? safeRegion(ctx.apiRegion) + ?? safeRegion(ctx.ssoRegion) + ?? "us-east-1"; +} +``` + +The ARN is authoritative because an enterprise profile can live in a different region from +the SSO session. The allowlist is not decoration: the region is interpolated into a +hostname, so an unvalidated value is a request-forgery primitive — and `apiRegion` / +`ssoRegion` come from external credential files +(`src/oauth/kiro-credentials.ts:285`), so **every** candidate goes through +`safeRegion`, not only the ARN (audit round 1, blocker 7). + +### The request + +POST to `https://management..kiro.dev/` with query `origin=AI_EDITOR`, +`isEmailRequired=true`, and `profileArn` when present; the same three in the JSON body; +headers per doc `001`; `x-amz-target: AmazonCodeWhispererService.GetUsageLimits`. +Bounded by `AbortSignal.timeout(REQUEST_TIMEOUT_MS)` (8s, the existing constant). +Non-2xx → `null`. Body read through the existing `readQuotaJson` size/stall guard. + +### Parsing + +```ts +const RESOURCE_PRIORITY = ["AGENTIC_REQUEST", "CREDIT"] as const; +``` + +Select the first breakdown whose `resourceType` matches, in priority order. **No index +fallback** — that is kiro-lb gap #6. If neither is present, return `null` (unknown), which +the cache layer renders as "unavailable" rather than "0%". + +Numbers prefer `currentUsageWithPrecision` / `usageLimitWithPrecision`, falling back to +the integer fields. Percent = `used / limit * 100`, run through the existing +`normalizePercent` (which clamps 0-100). + +Mapping into `ProviderQuota`: + +- The plan allowance is a **monthly** window: `monthlyPercent`, and `monthlyResetAt` + from `nextDateReset` (seconds → ms via the existing `normalizeResetAt`, which already + handles both scales). +- `freeTrialInfo` with a positive limit adds `customWindows: [{ label: "Free trial", percent }]`. +- `exhausted` and `nextResetAt` are carried on the snapshot for `030`'s cooldown + seeder, through the generation-guarded usage-state map defined in `070`. Nothing else + leaves this module. + +`exhausted` = `used >= limit && !overageEnabled`. Overage-enabled accounts keep serving +past the limit (doc `001` rule 3), so exhaustion is not `percent >= 100`. + +### Privacy + +`userInfo.email` and `userInfo.userId` are **read and discarded**. They are never +returned, never cached, never logged. `privacy:scan` covers the file; the regression test +asserts the snapshot object contains no email even when the payload carries one. + +## Accept criteria + +| # | Scenario | Observable proof | +| --- | --- | --- | +| 1 | Enterprise payload with `AGENTIC_REQUEST` | `monthlyPercent` = 14.78 for 147.82/1000, `monthlyResetAt` set | +| 2 | Payload where `AGENTIC_REQUEST` is absent but `CREDIT` present | `CREDIT` selected, not index 0 | +| 3 | Payload with only an unknown `resourceType` | returns `null` (activation: proves no index fallback) | +| 4 | Precision and integer fields both present | precision wins (695.17, not 695) | +| 5 | `overageStatus: ENABLED` with used > limit | `exhausted === false`, percent clamped to 100 | +| 6 | `overageStatus: DISABLED` with used >= limit | `exhausted === true` | +| 7 | `freeTrialInfo` present | a "Free trial" custom window appears | +| 8 | Response carries `userInfo.email` | serialized snapshot contains no email substring | +| 9 | Profile ARN region differs from `apiRegion` | request host uses the ARN region | +| 10 | Malformed ARN region (`../evil`) | falls back to `apiRegion`; host never contains the injected text | +| 10b | Malformed `apiRegion` and `ssoRegion` too | falls back to `us-east-1`; host never contains injected text | +| 11 | HTTP 401/429/500 | resolves `null`, does not throw | + +## Verifier + +`bun test tests/kiro-usage-quota.test.ts` (new file), plus `bun x tsc --noEmit`. +Both run against this exact file. Confirmed present: `bun` resolves and `tests/` is a flat +suite directory, so a new `tests/*.test.ts` is picked up with no config change. + +## Out of scope for this phase + +No caching, no account iteration, no routing, no surfaces. This phase ends with a pure +function and its tests. diff --git a/devlog/_plan/260829_kiro_quota_pool/020_per_account_quota_wiring.md b/devlog/_plan/260829_kiro_quota_pool/020_per_account_quota_wiring.md new file mode 100644 index 0000000000..f059159f2c --- /dev/null +++ b/devlog/_plan/260829_kiro_quota_pool/020_per_account_quota_wiring.md @@ -0,0 +1,122 @@ +# 020 — Phase 2: per-account quota wiring + +Work class: C3. Depends on: `010` (the fetcher). + +## Goal + +Every logged-in Kiro account gets a quota row through the seam that already exists for +Anthropic — cache, TTL, generation reconciliation, failure isolation and all. + +## The seam + +`src/providers/quota.ts:1453`: + +```ts +export function supportsPerAccountQuota(provider: string): boolean { + return provider === "anthropic"; // → provider === "anthropic" || provider === "kiro" +} +``` + +`fetchAccountQuota` (`:1572`) currently hard-calls `fetchAnthropicUsageQuota(token)`. +Replace that single line with a per-provider dispatch: + +```ts +let quota: ProviderQuota | null; +if (provider === "kiro") { + const snapshot = await fetchKiroUsageSnapshot(await kiroUsageContextForAccount(accountId)); + quota = snapshot?.quota ?? null; + // Written inside the SAME mayCommitAccountQuotaKey(key, writerGeneration) branch that + // guards the quota row, so a superseded probe commits neither (070, blocker 2). + kiroUsageStateToCommit = snapshot + ? { exhausted: snapshot.exhausted, nextResetAt: snapshot.nextResetAt } + : null; +} else { + quota = await fetchAnthropicUsageQuota(token); +} +``` + +Everything around it is reused unchanged and that is the point: the TTL negative-caching, +the `unavailable` flag that preserves last-good bars, the in-flight join, the +`mayCommitAccountQuotaKey` generation guard, and `clearAccountQuotaCache` on logout. +Those behaviours took several PRs to get right; Kiro inherits them for free. + +## New helper: `kiroUsageContextForAccount` + +Lives in `src/providers/kiro-usage.ts`, reads the stored account credential and assembles +`KiroUsageContext`: + +- `access` **and** the routing metadata from one + `getValidAccessSnapshotForAccount(provider, accountId)` call + (`src/oauth/index.ts:435`). +- `accountId` passed through, so the usage-state map is keyed exactly like the quota + cache. + +**Not** `getTokenForAccountQuotaProbe` (audit round 1, blocker 4). That helper refuses to +refresh a background `source: "local-cli"` account (`src/providers/quota.ts:1549`) +because Anthropic's lock can adopt a mismatched Claude CLI identity. Kiro's *imported* +credentials are marked `local-cli` for an unrelated reason — they came from the Kiro CLI +database (`src/oauth/kiro.ts:301`) — so reusing that rule would make every inactive Kiro +account's quota unavailable the moment its token expired, which is exactly when a pool +needs it. The fail-closed branch stays Anthropic-scoped. + +Resolving both values from a **single** snapshot also strengthens the anti-cross-pairing +invariant below: token and profile ARN provably come from one read. + +**The load-bearing invariant:** the bearer and the routing metadata must come from the +*same* account record. This is the exact class of bug #2841 fixed for Copilot origins — +one account's token paired with another account's destination. Here it would send account +B's bearer with account A's profile ARN, which at best 403s and at worst reports A's quota +under B's row. The helper therefore takes `accountId` and reads one record; it never +consults `getCredential(provider)` (the *active* account) for any field. + +## Builder ID + +An account with no stored profile ARN gets the request-scoped service profile the +generation path already uses (`resolveKiroRequestProfile`). Reuse that resolver rather +than re-deriving it — doc `001` and the existing comment at `src/adapters/kiro.ts:1888` +both stress that this value must never become stored identity. + +## `ksk_` API keys + +Not OAuth accounts. `fetchProviderAccountQuotas` iterates the OAuth account set, so they +are naturally absent. No special case needed; documented so a future reader does not add one. + +## Provider-level row + +Add a `kiro` branch to `maybeFetchProviderQuota` (`:2186`) that probes the **active** +account and reports it as the provider row, mirroring `fetchAnthropicQuota` (`:1387`) +including its "capture the account before awaiting" guard against a mid-flight switch. +Source string: `"kiro:usage-limits"`. + +## Accept criteria + +| # | Scenario | Observable proof | +| --- | --- | --- | +| 1 | `supportsPerAccountQuota("kiro")` | true | +| 2 | Two Kiro accounts, both healthy | two rows, each with its own percent | +| 3 | Account A ok, account B 401 | A has bars; B has `unavailable: true` and A is unaffected | +| 4 | Second call inside TTL | zero additional fetches (activation: assert call count) | +| 5 | `forceRefresh` | exactly one new fetch per account | +| 6 | Accounts with different profile ARNs | each request host/ARN matches its own account (cross-pairing regression) | +| 7 | Logout clears rows | `clearAccountQuotaCache("kiro")` empties them and cancels in-flight | +| 8 | Provider row present | `/api/provider-quotas` includes a `kiro` report with source `kiro:usage-limits` | +| 9 | Account removed, then a stale probe resolves | no usage-state row survives for the removed id | +| 10 | `clearAccountQuotaCache("kiro")` | usage-state rows for kiro are cleared too | +| 11 | Probe resolves after a generation bump | neither the quota row nor the usage state commits | +| 12 | Inactive imported (`local-cli`) account with an expired token | quota resolves, is NOT forced unavailable | + +Criterion 6 is the security-relevant one and must be written first, red. + +## Test file changes + +- New: `tests/kiro-account-quota.test.ts`. +- Amend: `tests/provider-account-quota.test.ts:204` currently asserts + `supportsPerAccountQuota("kiro") === false` with zero network calls. That assertion + becomes false by design. Rewrite it to assert the **generic** exclusion still holds for a + provider we genuinely do not support (e.g. `"xai"`), preserving the original intent — + "unsupported providers make no network calls" — rather than deleting the coverage. + +## Verifier + +`bun test tests/kiro-account-quota.test.ts tests/provider-account-quota.test.ts tests/provider-quota.test.ts` +— all three read this change target directly. diff --git a/devlog/_plan/260829_kiro_quota_pool/030_quota_aware_pool_selection.md b/devlog/_plan/260829_kiro_quota_pool/030_quota_aware_pool_selection.md new file mode 100644 index 0000000000..d03bf4431a --- /dev/null +++ b/devlog/_plan/260829_kiro_quota_pool/030_quota_aware_pool_selection.md @@ -0,0 +1,118 @@ +# 030 — Phase 3: quota-aware pool selection + +Work class: C3. Depends on: `020` (cached per-account quota). + +## The gap + +Rotation today is **reactive and order-blind**: on a 429 it walks the stored account order +from the failed account and takes the first non-cooled, non-reauth account +(`src/oauth/generic-account-failover.ts:157`). If the very next account is at 99% used, we +rotate into it, 429 again, and burn a second rotation from a budget of three. + +This is the axis on which kiro-lb is genuinely ahead: it *ranks* by remaining fraction +before choosing. We should be ahead of it instead — it ranks with 15-minute-stale data and +ignores the requested model entirely (doc `002`, gaps 1 and 2). + +## Design + +Add an **ordering** step, not a new pool. The real owner is +`rotateGenericOAuthAccountOn429` (`src/oauth/generic-account-failover.ts:157`); it keeps +its contract (cooldowns, reauth skipping, request-local rotation, never mutating +`activeAccountId`) and gains a rank step in place of its inline ring walk: + +```ts +// src/oauth/account-quota-rank.ts +export function rankAccountsByHeadroom(provider: string, ring: string[]): string[]; +``` + +### The ring is built first, then ranked + +This ordering is load-bearing and easy to get wrong. `eligibleFailoverAccounts` returns +ids in **stored** order (`:143`), while the existing traversal starts **after the failed +account** (`:184`). Ranking the stored-order list would silently change today's behaviour: +with roster `[A, B, C]` and `B` failing, stored order picks `A` where the ring picks +`C`. + +So the caller constructs the ring explicitly, then ranks it: + +```ts +const order = set.accounts.map(a => a.id); +const start = order.indexOf(failedAccountId); +const ring = start >= 0 ? [...order.slice(start + 1), ...order.slice(0, start)] : order; +const candidates = ring.filter(id => eligible.includes(id) && id !== failedAccountId); +return rankAccountsByHeadroom(providerName, candidates)[0] ?? null; +``` + +Rules, in order: + +1. **Synchronous only.** It reads `getCachedProviderAccountQuota(provider, id)` + (`src/providers/quota.ts:1470`), which never probes the network. Rotation happens + mid-request while a 429 is in hand; it cannot await a usage call. +2. **Three categories, no numeric weight.** Accounts fall into + `known-healthy` → `unknown` → `known-exhausted`, where exhausted is our own verdict + from `010` (limit reached with overage disabled). No borrowed constant: an unknown + account must obviously be tried before one we measured as empty and after one we + measured as full, and that ordering needs no tuning parameter. +3. **Within known-healthy, sort by descending headroom.** Headroom = `100 - percent` of + the monthly window (Kiro's plan window), or the minimum across known windows for + providers reporting several. +4. **Stable within a bucket.** Ties preserve ring order, so the deterministic + walk-the-roster property survives. Deterministic by choice: with a handful of accounts + and a fresh rank per request, randomization buys nothing and costs reproducibility. +5. **Zero quota data anywhere → identity.** Returns the ring unchanged, so behaviour for + every provider without per-account quota is byte-for-byte what it is today. + +## No live decrement (withdrawn after audit round 1) + +An earlier draft proposed nudging the cached percent after each successful turn to close +kiro-lb's 15-minute staleness gap. It is removed: `ProviderQuota` carries no absolute +limit to divide by, and Kiro bills fractional credits, so a per-turn increment would be +fabricated. Stale truth beats invented precision. + +Be precise about the freshness we do claim, because there are **two** different TTLs: + +- **Provider-level display** rows cache for 5 minutes (`src/providers/quota.ts:37`). +- **Per-account** rows — which are what this ranking reads — cache for **10 minutes** + (`:1425`), deliberately longer because the cost multiplies by account count. + +So recovery ranking can act on data up to 10 minutes old. That is still better than +kiro-lb's 15-minute default poll, but the honest margin is 10-vs-15, not 5-vs-15. + +## Exhaustion cooldown + +When `010`'s snapshot says `exhausted`, seed the failover cooldown for that account until +`nextResetAt` (clamped: minimum 5 minutes, maximum 24 hours) instead of the default 60s. +Retrying a monthly-exhausted account every minute is pure waste. The clamp keeps a bogus +upstream reset date from parking an account for a month — kiro-lb allows a 32-day +quarantine, which we consider too much rope. + +The exhaustion state is owned by the generation-guarded store described in `070`, not by +an ad-hoc module map; a stale `exhausted` flag surviving an account removal would hand a +replacement account a 24-hour cooldown it never earned. + +## Accept criteria + +| # | Scenario | Observable proof | +| --- | --- | --- | +| 1 | A at 10% used, B at 90% used | A ranks first | +| 2 | A unknown, B at 95% used but NOT exhausted | B ranks first — B is known-healthy, and the categorical rule has no "known-low" tier | +| 3 | A unknown, B at 5% used | B ranks first (known-high beats unknown) | +| 3b | A unknown, B exhausted (limit reached, overage off) | A ranks first (unknown beats known-exhausted) | +| 4 | No quota data for any account | order identical to input (activation: proves the no-op path) | +| 5 | Equal headroom | ring order preserved | +| 6 | Rank called during rotation | zero network calls (assert fetch not invoked) | +| 7 | Exhausted + `nextResetAt` in 3 days | cooldown clamped to 24h, not 3 days | +| 8 | Exhausted + `nextResetAt` in 30 seconds | cooldown floored at 5 minutes | +| 9 | Roster [A,B,C], B fails, no quota data | C selected (ring), not A (stored order) | +| 10 | Rotation still skips cooled/reauth accounts | existing failover tests stay green | + +## Blast radius + +`rankAccountsByHeadroom` runs for every generic-failover provider, so criterion 4 is the +one that protects xAI, Cursor, Copilot, Kimi and Antigravity from behaviour change. +`tests/generic-oauth-failover.test.ts` and `tests/adapter-event-oauth-failover.test.ts` +must stay green untouched. + +## Verifier + +`bun test tests/kiro-pool-rank.test.ts tests/generic-oauth-failover.test.ts tests/adapter-event-oauth-failover.test.ts` diff --git a/devlog/_plan/260829_kiro_quota_pool/040_surfaces_cli_gui_docs.md b/devlog/_plan/260829_kiro_quota_pool/040_surfaces_cli_gui_docs.md new file mode 100644 index 0000000000..ef374a62ca --- /dev/null +++ b/devlog/_plan/260829_kiro_quota_pool/040_surfaces_cli_gui_docs.md @@ -0,0 +1,74 @@ +# 040 — Phase 4: surfaces (CLI, GUI, docs) + +Work class: C2. Depends on: `020`. Independent of `030`. + +## GUI: nothing to build + +This is the payoff for wiring into the existing seam. `ProviderAuthPanel` already renders +`QuotaBars` from `account.quota` and the unavailable message from +`account.quotaUnavailable` (`gui/src/components/provider-workspace/ProviderAuthPanel.tsx:517`); +`useProviderAccountPools` already requests `?quota=1` +(`gui/src/hooks/useProviderAccountPools.ts:96`); the route already populates both fields +from `fetchProviderAccountQuotas` +(`src/server/management/oauth-account-routes.ts:274`). + +Once `supportsPerAccountQuota("kiro")` is true, Kiro accounts render. **Verification is +still required** — "it should just work" is not evidence. A GUI screenshot is mandatory in +the PR description anyway (the `enforce-target` gate rejects a `gui`-mentioning PR +without one). + +The monthly window renders through the existing `monthlyPercent` row in +`QuotaBars` (`gui/src/components/QuotaBars.tsx:45`), so no new component and no new +label vocabulary. + +## CLI + +`ocx account list kiro --quota` already exists and formats a QUOTA column +(`src/cli/account.ts:104`). `quotaText` reads `fiveHourPercent`/`shortPercent` and +`weeklyPercent` — **neither of which Kiro populates**. Add a monthly arm: + +```ts +if (typeof quota.monthlyPercent === "number") parts.push(\`mo \${quota.monthlyPercent}%\`); +``` + +Without this the column prints `-` for a perfectly healthy Kiro account, which reads as +"broken". This is a two-line change with a real user-visible failure behind it. + +## The stale "single login slot" copy + +`src/cli/account.ts:28` and `:215` describe Kiro as replacement-style with a single login +slot. That has been false since the multiauth add-account flow shipped +(`src/oauth/kiro.ts:335`, which snapshots the CLI SQLite DB, runs +`kiro-cli logout`/`login`, and appends by profile ARN with rollback on failure). + +Fix the copy to describe what the code does: multiple accounts, added one at a time through +the CLI handoff, each with its own quota row. A user who reads "single login slot" will +never try to build a pool — the feature is invisible, which is functionally the same as +missing. This directly serves the user's "pool 기반 자동 탑재" ask. + +## Docs + +- `docs-site/src/content/docs/reference/adapters.md` — Kiro section: note quota reporting + and the multi-account pool. +- `docs-site/src/content/docs/reference/cli/providers-accounts.md` — the `--quota` + column now covers Kiro; document the monthly window and the `unavailable` state. + +English source only. Translated locales are left alone rather than machine-translated; +the repo rule is that locales must not *contradict* English, and an untouched locale that +omits a new note does not contradict it. + +## Accept criteria + +| # | Scenario | Observable proof | +| --- | --- | --- | +| 1 | `quotaText` with only `monthlyPercent` | renders `mo 15%`, not `-` | +| 2 | `quotaText` with `quotaUnavailable` | renders `unavailable` (existing behaviour preserved) | +| 3 | Help text for kiro | no longer claims a single login slot | +| 4 | GUI accounts tab with 2 Kiro accounts | screenshot shows two quota bars | +| 5 | `bun run lint:gui` | passes (no GUI source change expected, so this is a guard) | + +## Verifier + +`bun test tests/account-cli.test.ts` (or the existing CLI account test file — confirm the +exact name before writing the plan into the attest) and a manual GUI screenshot. +`bun run skill:surface:check` if any CLI capability string changes. diff --git a/devlog/_plan/260829_kiro_quota_pool/050_verification_and_delivery.md b/devlog/_plan/260829_kiro_quota_pool/050_verification_and_delivery.md new file mode 100644 index 0000000000..94ab6a785e --- /dev/null +++ b/devlog/_plan/260829_kiro_quota_pool/050_verification_and_delivery.md @@ -0,0 +1,71 @@ +# 050 — Phase 5: verification, head-to-head proof, delivery + +Work class: C2. Depends on: `010`–`040`. + +## Verification matrix + +| Gate | Command | Reads this unit's target? | +| --- | --- | --- | +| Preflight | `bun install` (root), `cd gui && bun install` | Prerequisite — without it every gate below fails on missing `bun-types`/`zod`/`oxlint` | +| Types | `bun x tsc --noEmit` | Yes — strict, whole project | +| Focused tests | `bun test tests/kiro-usage-quota.test.ts tests/kiro-account-quota.test.ts tests/kiro-pool-rank.test.ts` | Yes — direct arguments | +| Regression | `bun test tests/provider-quota.test.ts tests/provider-account-quota.test.ts tests/generic-oauth-failover.test.ts tests/adapter-event-oauth-failover.test.ts` | Yes | +| Full suite | `bun run test` | Yes — required, this touches shared routing/config/server | +| Privacy | `bun run privacy:scan` | Yes — reads `src/` and `devlog/` | +| GUI lint | `bun run lint:gui` | Only if GUI source changes | + +`bun run test` is not optional here: AGENTS.md requires the full suite for shared +routing/config/server changes, and `quota.ts` plus the failover path qualify. + +## The head-to-head claim + +The user's bar is "확언할 수 있을 때까지" — able to state with confidence that we are +better. That requires a table where **every one of our claims cites our code** and **every +kiro-lb claim cites its file:line**, written after the implementation, from the landed +tree. Not from this plan. + +Draft axes (to be filled with evidence at C, not asserted now): + +| Axis | kiro-lb | opencodex (to prove) | +| --- | --- | --- | +| Usage source | `GetUsageLimits`, 900s background poll | same operation, pull-on-demand, no timer | +| Freshness | 15 min default poll | 5 min provider-level display cache; 10 min per-account cache | +| Breakdown selection | `AGENTIC_REQUEST` else index 0 | priority list, unknown → unavailable | +| Free trial pool | dropped | separate window | +| Recovery ordering | reactive ring walk after 429 | headroom-ranked ring after 429 | +| Pre-request selection | weighted random, model-blind | **not shipped in this unit** (deferred) | +| Unknown account | numeric weight | categorical known-healthy > unknown > known-exhausted | +| Exhaustion | 6h–32d quarantine | reset-aligned, clamped 5min–24h | +| Identity safety | per-account auth manager | per-account snapshot, cross-pairing regression test (#2841 lineage) | +| Idle cost | polls every 60s minimum | zero requests when idle | +| Scope | Kiro only | Kiro is one provider among many, same seam | + +Honesty requirement: kiro-lb has a dedicated operations dashboard with request-rate charts, +per-model token panels and Prometheus export. We do not, and the table must say so. A +comparison that only lists our wins is marketing, and the user asked for confidence, not +cheerleading. + +Two more rows must stay in the "they are ahead" column: kiro-lb selects an account +**before** the request (model-blind, but pre-dispatch), and it persists quota rows across +restart to seed routing. This unit does neither. + +## Delivery + +- Branch: `codex/kiro-quota-pool` off current `dev`. +- Commits: one per phase (`010`…`040`), plus the devlog unit. +- PR against `dev`, full template (Summary / Verification / Checklist), GUI screenshot. +- Push with `--no-verify` (user-authorized), merge after CI green. +- Move `devlog/_plan/260829_kiro_quota_pool/` → `devlog/_fin/` at close-out, since the + fix will be public by then. + +## Security note + +Nothing in this unit is pre-disclosure material: the cross-account pairing invariant is +already public via merged #2841, and everything here is a forward-looking feature. So the +devlog unit is the right home. If implementation *uncovers* an unfixed weakness, that +write-up goes to `.tmp/`, not here. + +## Terminal outcome + +`DONE` requires: all gates green, the head-to-head table written from the landed tree, the +PR merged into `dev`. Anything less is reported as its real outcome, not rounded up. diff --git a/devlog/_plan/260829_kiro_quota_pool/060_audit_round1_amendments.md b/devlog/_plan/260829_kiro_quota_pool/060_audit_round1_amendments.md new file mode 100644 index 0000000000..7946c2131a --- /dev/null +++ b/devlog/_plan/260829_kiro_quota_pool/060_audit_round1_amendments.md @@ -0,0 +1,143 @@ +# 060 — Audit round 1: nine blockers, and what changed + +An independent reviewer audited docs `000`–`050` against the tree at +`124a2b1487996f8a8ebb2067b22c9e758fa6016f` and returned **FAIL, blockers=9**. Every one +was verified against real code. This document records the disposition; the amended rules +here **supersede** the corresponding text in the earlier docs. + +## Environment (blocker 8) — fixed, not amended + +The worktree had no `node_modules`, so the entire verifier matrix in `050` was +unrunnable: `bun x tsc --noEmit` exited 1 on missing `bun-types`, the suite reported 774 +failures from a missing `zod/v4`, and `bun run lint:gui` exited 127 on a missing +`oxlint`. `bun install` now completes (103 packages). `bun run privacy:scan` and +`bun run skill:surface:check` were already exit 0. + +**Amendment to `050`:** the verification matrix gains a preflight row — `bun install` +at the repo root, and `cd gui && bun install` before any GUI gate. A verifier that cannot +run is not a verifier. + +## Blocker 5 — the insertion point does not exist. FOLDED. + +There is no `chooseFailoverAccount`. The real owner is `rotateGenericOAuthAccountOn429` +(`src/oauth/generic-account-failover.ts:157`), whose ring traversal starts *after* the +failed account (`:184`) so repeated 429s walk the roster. + +**Amendment to `030`:** rank composes with the ring rather than replacing it. +`eligibleFailoverAccounts` already returns the eligible ids; build that list, apply the +stable rank when quota data exists, and take the first entry. When no quota data exists the +ranked list must be identical to the ring order starting after the failed account, so the +existing deterministic behaviour is preserved exactly. + +## Blocker 1 — "pre-request selection" was claimed but not designed. PARTIALLY FOLDED, PARTIALLY DESCOPED. + +Correct: `rotateGenericOAuthAccountOn429` runs only inside the `status === 429` branch +(`src/server/responses/core.ts:5574`), so ranking there is recovery ordering, not +pre-flight selection. And `rankAccountsByHeadroom(provider, ids)` takes no model, so +calling the result "model-aware" was false. + +**Amendment to `002` and `030`:** the head-to-head table drops the "model-aware" and +"before the request" claims outright. What we deliver in this unit is *recovery ordering +that is quota-aware*, which is still a real improvement over walking the ring blind, and +which we will describe as exactly that. + +A genuine pre-dispatch selection seam — choosing the account before the first request, +with model eligibility — is a larger change to the request path and becomes its own +work-phase rather than an unbacked sentence in this one. Recorded as follow-up, not +claimed as shipped. + +## Blocker 3 — the live decrement has no denominator. REMOVED. + +`ProviderQuota` stores percent and reset, never absolute used/limit +(`src/providers/quota.ts:93`), so `100 / limit` cannot execute from cached state. Worse, +Kiro meters *fractional* credits: one successful turn is not one credit, so any fixed +increment is a fabrication that would misrank accounts. + +**Amendment to `030` and `002`:** `markAccountObservedUsage` is deleted from the plan, +and the "live decrement between polls" row is removed from the head-to-head table. We do +not get to claim an advantage we cannot compute. The honest position: our data is as fresh +as the 5-minute TTL and the vendor's own 5-minute update floor allows. + +## Blocker 2 — snapshot metadata had no consumer. FOLDED. + +`fetchAccountQuota` caches `ProviderQuota | null` (`:1426`), so extracting `.quota` +and dropping `subscriptionTitle` / `exhausted` / `nextResetAt` / `overageEnabled` +orphans all four. + +**Amendment to `010`/`020`:** each field must reach a consumer or leave the design. + +- `exhausted` + `nextResetAt` → consumed by the exhaustion cooldown in `030`. To reach + it they need a home: add a module-private `kiroAccountUsageState` map in + `src/providers/kiro-usage.ts`, written by the fetcher and read by the cooldown seeder. + It is not part of `ProviderQuota` and is not serialized to any API. +- `overageEnabled` → consumed only as an input to `exhausted`. It stops being a + returned field and becomes a local variable. +- `subscriptionTitle` → **dropped from this unit.** `ProviderQuota` has no plan-name + field and the GUI has no place to render one; adding both is scope creep. Plan tier is + simply not surfaced by this unit. (An earlier version of this line claimed the tier was + "visible from the limit value" — that was wrong: `ProviderQuota` serializes percent and + reset, never the absolute limit. Corrected in round 2.) + +## Blocker 4 — the Anthropic fail-closed rule breaks Kiro probes. FOLDED. + +Sharp catch. `getTokenForAccountQuotaProbe` refuses to refresh a background +`source: "local-cli"` account (`:1549`) because Anthropic's lock can adopt a mismatched +Claude CLI identity. But Kiro's *imported* credentials are marked `local-cli` +(`src/oauth/kiro.ts:301`) for an unrelated reason — they came from the Kiro CLI database. +Reusing that rule verbatim would make every inactive Kiro account's quota go unavailable +the moment its token expired, which is precisely when a pool needs it. + +**Amendment to `020`:** the fail-closed branch stays Anthropic-scoped. Kiro resolves +through `getValidAccessSnapshotForAccount(provider, accountId)` +(`src/oauth/index.ts:435`), which is account-scoped and returns the bearer *and* the +`kiro` routing metadata from one snapshot — which also strengthens the anti-cross-pairing +invariant, since token and ARN now provably come from a single read. + +## Blocker 6 — the helpers are module-private. FOLDED. + +`REQUEST_TIMEOUT_MS`, `normalizeResetAt`, `normalizePercent` and `readQuotaJson` are +private to `quota.ts`. + +**Amendment to `010`:** extract them into a new neutral `src/providers/quota-wire.ts` +(timeout constant, number/percent/reset normalizers, bounded JSON reader) that both +`quota.ts` and `kiro-usage.ts` import. Pure move, no behaviour change; `quota.ts` +re-exports nothing new publicly. This is a prerequisite step inside phase 1, and the +existing `tests/provider-quota.test.ts` is the regression proof that the move is inert. + +## Blocker 7 — only the ARN region was validated. FOLDED. + +`apiRegion` and `ssoRegion` come from external credential files +(`src/oauth/kiro-credentials.ts:285`) and were interpolated into the hostname unchecked, +which makes the claimed request-forgery guard hollow. + +**Amendment to `010`:** one `safeRegion()` allowlist parser (`^[a-z0-9-]{1,32}$`) +applies to *every* candidate — ARN field, `apiRegion`, `ssoRegion` — and anything failing +it falls through to the next candidate, then to `us-east-1`. Accept criteria 10 expands to +hostile `apiRegion` and `ssoRegion` cases, not just the ARN. + +## Blocker 9 — the `0.25` weight reads as an AGPL port. FOLDED. + +Naming kiro-lb's exact constant while claiming independent derivation is the weakest kind +of clean-room claim, and the reviewer is right to refuse it. + +**Amendment to `030`:** unknown ordering becomes **categorical**, with no borrowed +constant: `known-healthy > unknown > known-exhausted`, where healthy/exhausted is our own +`exhausted` verdict from `010`. Within the known-healthy bucket, sort by descending +headroom; ties keep ring order. No numeric weight is imported from the reference, and the +comparison document says only that both projects rank unknown between the two extremes — +which is an obvious design necessity, not a borrowed policy. + +## Corrections to record + +- The real CLI formatting test is `tests/cli-headless-parity.test.ts:744`, not + `tests/account-cli.test.ts`. `040`'s verifier row is corrected. +- `tests/provider-account-quota.test.ts` lines 206/209/210 are the only Kiro assertions; + the `xai` substitution preserves intent. +- No import cycle exists: `quota.ts` has no path back to `generic-account-failover.ts`. +- No `src/lab/` reachability is introduced. +- The wire transcription in `001` is accurate. + +## Residual + +Pre-dispatch, model-aware account selection is deferred to a follow-up work-phase. +Everything else is folded into the amended plan above. diff --git a/devlog/_plan/260829_kiro_quota_pool/070_audit_round2_amendments.md b/devlog/_plan/260829_kiro_quota_pool/070_audit_round2_amendments.md new file mode 100644 index 0000000000..030dc4a134 --- /dev/null +++ b/devlog/_plan/260829_kiro_quota_pool/070_audit_round2_amendments.md @@ -0,0 +1,130 @@ +# 070 — Audit round 2: four blockers, and the state-ownership design + +Round 2 returned **FAIL, blockers=4**. Two were my failure to actually edit the canonical +documents (a supersession note is not an amendment), one was a genuine off-by-one I had +not seen, and one was a type-ownership cycle. Post-`bun install` the reviewer confirmed +`bun x tsc --noEmit` exit 0, `tests/provider-quota.test.ts` 107 pass, `privacy:scan` +exit 0. + +## Blocker 1 — stored order is not ring order. FOLDED into `030` directly. + +The counterexample is exact: roster `[A, B, C]`, `B` fails, no quota data. +`eligibleFailoverAccounts` returns `[A, C]` in **stored** order (`:143`), so ranking +that list identity-style picks `A` — while today's traversal, which starts after the +failed account (`:184`), picks `C`. My "byte-for-byte identical when no quota data" +claim was therefore false. + +`030` now builds the ring explicitly before ranking +(`[...order.slice(start + 1), ...order.slice(0, start)]`), filters eligibility against it, +and ranks that. Accept criterion 9 is now this exact three-account counterexample, which +must fail before the fix and pass after. + +## Blocker 3 — the false claims were still in 002/030/050. FOLDED by editing them. + +Correct and fair. `060` announced the descopes but left the originals intact, so an +implementer reading `030` would still have built `markAccountObservedUsage`, and `050` +still instructed the final comparison to claim it. The canonical docs are now edited in +place: + +- `002` — "before the request", "model-aware" and "live decrement" removed from the + advantages list; a **Scope honesty** section states plainly that kiro-lb is ahead of us + on pre-request selection and that we do not close its staleness gap. +- `030` — the `0.25` borrowed weight is gone, replaced by the categorical ordering; the + live-decrement section is replaced by an explicit withdrawal. +- `050` — the comparison table drops the two false rows, adds a preflight row, and gains + two rows in the "they are ahead" column (pre-request selection, restart persistence). + +Also corrected: the claim that "Kiro's plan tier is visible from the limit value itself" +was wrong — `ProviderQuota` serializes percent and reset, never the absolute limit. The +sentence is removed from `060`; plan tier is simply out of scope for this unit. + +## Blocker 2 — the usage-state map had no ownership. FOLDED with a real design. + +The proposed module-private map would have been a hidden global: stale `exhausted` state +surviving an account removal could hand a *replacement* account a 24-hour cooldown it never +earned, and a late probe could republish state after a config generation change. The +existing quota cache solved exactly this with three mechanisms +(`mayCommitAccountQuotaKey` at `:1438`, `reconcileProviderAccountQuotaRows` at +`:1495`, `clearAccountQuotaCache` at `:1522`) and every logout/removal path already +calls the last one. + +**Design:** do not build a parallel store. `kiroAccountUsageState` becomes a +`Map` +living in `src/providers/kiro-usage.ts` and wired into the same three seams: + +1. **Keyed identically** to the quota cache (`\`\${provider}\\u0000\${accountId}\``), which + also fixes the reviewer's sub-point that `KiroUsageContext` carried no account id — + the context gains `accountId`. +2. **Written only after the quota owner's commit guard passes.** The write happens inside + `fetchAccountQuota`'s existing `mayCommitAccountQuotaKey(key, writerGeneration)` + branch, so a late probe from a superseded generation cannot commit usage state either. +3. **Cleared and reconciled with the quota rows.** `clearAccountQuotaCache(provider)` + clears the matching usage-state keys, and `reconcileProviderAccountQuotaRows` drops + usage-state keys absent from `context.oauthAccountKeys`. Both are single added lines in + functions that already do this for quota; no new registration in + `src/lib/state-store-registrations.ts` is required because the owner is already + registered there (`reconcileProviderAccountQuotaRows`). +4. **Consumed** by `030`'s cooldown seeder through one exported reader, + `getKiroAccountExhaustion(provider, accountId)`. Generic failover imports that reader; + it does not reach into the map. + +### Freshness (added in round 3) + +Storing `ts` without a staleness rule would let an old `exhausted` verdict keep an +account last-ranked, or re-seed a 24-hour cooldown, long after its quota actually reset. +`getKiroAccountExhaustion` therefore returns **unknown** — not `exhausted` — when +either holds: + +- `now - entry.ts >= ACCOUNT_QUOTA_TTL_MS` (the existing per-account TTL, which is + **10 minutes** at `src/providers/quota.ts:1425`, not 5 — the 5-minute figure is the + provider-level cache at `:37`). That constant is currently module-private to + `quota.ts`, and `kiro-usage.ts` cannot import it back without recreating the cycle + `010` just removed, so the phase-1 extraction **also moves `ACCOUNT_QUOTA_TTL_MS` and + `CACHE_TTL_MS` into `quota-wire.ts`**, imported by both. Duplicating `10 * 60_000` + would put the same number in two files, which is exactly how the two drift apart; or +- `entry.nextResetAt !== undefined && entry.nextResetAt <= now` — the window the account + was exhausted in has already rolled over. + +Unknown means the account sorts in the middle bucket and gets a normal 60-second cooldown, +so the failure mode of stale data is "try it again", not "park it for a day". Tests drive +both expiry paths with a fake clock. + +Two more accept criteria for `020`: + +| # | Scenario | Observable proof | +| --- | --- | --- | +| 13 | Entry older than the 10-minute account TTL | reader returns unknown, not exhausted | +| 14 | `nextResetAt` already passed | reader returns unknown; cooldown reverts to the 60s default | + +New accept criteria for `020`: + +| # | Scenario | Observable proof | +| --- | --- | --- | +| 9 | Account removed, then a stale probe resolves | no usage-state row for the removed id | +| 10 | `clearAccountQuotaCache("kiro")` | usage-state rows for kiro are gone too | +| 11 | Probe resolves after a generation bump | neither quota nor usage state commits | + +## Blocker 4 — type-level ownership cycle. FOLDED. + +`quota.ts` would import `kiro-usage.ts` for the fetcher while `kiro-usage.ts` imports +`ProviderQuota` back from `quota.ts`. Type-only, so it would not break at runtime, but it +is still a cycle and the reviewer is right that the fix is cheap. + +**Amendment to `010`:** the phase-1 extraction produces **two** modules, not one: + +- `src/providers/quota-types.ts` — `ProviderQuota`, `ProviderQuotaWindow`, + `ProviderQuotaCreditsUsd` (the real exported name, `src/providers/quota.ts:84`), + `ProviderQuotaReport`. Pure types, imports nothing from `quota.ts`. +- `src/providers/quota-wire.ts` — `REQUEST_TIMEOUT_MS`, `normalizePercent`, + `normalizeResetAt`, `toFiniteNumber`, `asRecord`, `readQuotaJson`. Depends only on + `lib/bounded-body` and pure helpers. + +Both `quota.ts` and `kiro-usage.ts` import from these; the edge between them becomes +one-directional. `tests/provider-quota.test.ts` (107 pass today) is the proof the move is +inert, and `tests/core-lab-boundary.test.ts` confirms no Lab reachability is introduced. + +## Round 2 disposition + +All four folded, none rebutted. The two documentation blockers were fair hits on my +process: I wrote a supersession note instead of amending the source, which is exactly the +failure mode that lets a stale plan get implemented. From 37c4cfdf973dd7c9645bd805fbc425e50c61556b Mon Sep 17 00:00:00 2001 From: bitkyc08-arch Date: Sat, 29 Aug 2026 12:37:46 +0900 Subject: [PATCH 2/6] feat(kiro): read per-account usage limits from Kiro's management endpoint Extracts shared quota types and wire helpers so a provider-specific quota module can exist without a cycle, then adds the Kiro GetUsageLimits reader: resource-type selection, precision fields, overage-aware exhaustion, free-trial window, and an allowlisted region for every hostname candidate. --- src/providers/kiro-usage.ts | 272 +++++++++++++++++++++++++++++++++ src/providers/quota-types.ts | 36 +++++ src/providers/quota-wire.ts | 97 ++++++++++++ src/providers/quota.ts | 136 +++-------------- tests/kiro-usage-quota.test.ts | 236 ++++++++++++++++++++++++++++ 5 files changed, 666 insertions(+), 111 deletions(-) create mode 100644 src/providers/kiro-usage.ts create mode 100644 src/providers/quota-types.ts create mode 100644 src/providers/quota-wire.ts create mode 100644 tests/kiro-usage-quota.test.ts diff --git a/src/providers/kiro-usage.ts b/src/providers/kiro-usage.ts new file mode 100644 index 0000000000..46fa827960 --- /dev/null +++ b/src/providers/kiro-usage.ts @@ -0,0 +1,272 @@ +/** + * Kiro usage limits — the account-scoped quota read behind the Kiro pool. + * + * Kiro's generation traffic goes to `runtime..kiro.dev`, but the usage numbers + * live behind a different subdomain and a different AWS JSON-RPC operation: + * `AmazonCodeWhispererService.GetUsageLimits` on `management..kiro.dev`. The + * operation is undocumented, so every field here is best-effort: a shape we do not + * recognise resolves to `null` ("unknown"), never to a fabricated zero. + * + * This module also owns the small amount of state that quota percentages cannot express — + * whether an account is actually out of allowance, and when its window rolls over — because + * the pool needs those two answers to decide how long to cool a 429'd account. + */ +import { getValidAccessSnapshotForAccount } from "../oauth"; +import type { ProviderQuota, ProviderQuotaWindow } from "./quota-types"; +import { + ACCOUNT_QUOTA_TTL_MS, + asRecord, + normalizePercent, + normalizeResetAt, + QUOTA_JSON_READ_FAILURE, + readQuotaJson, + REQUEST_TIMEOUT_MS, + toFiniteNumber, +} from "./quota-wire"; + +const AMZ_USAGE_TARGET = "AmazonCodeWhispererService.GetUsageLimits"; + +/** + * Regions are interpolated into a hostname, and two of the three candidates below are read + * out of credential files this process did not write. An allowlist keeps a crafted region + * from redirecting the request somewhere else entirely. + */ +const REGION_PATTERN = /^[a-z0-9-]{1,32}$/; + +/** + * Which usage bucket represents the plan allowance, in preference order. + * + * Selecting by position instead would mean an upstream reordering silently reweights the + * pool against an unrelated resource, so an unrecognised list resolves to unknown. + */ +const RESOURCE_PRIORITY = ["AGENTIC_REQUEST", "CREDIT"] as const; + +export interface KiroUsageContext { + /** Keys the usage-state row; always the stored account id, never the active account. */ + accountId: string; + access: string; + profileArn?: string; + apiRegion?: string; + ssoRegion?: string; +} + +export interface KiroUsageSnapshot { + quota: ProviderQuota; + /** Allowance is spent AND overage is not enabled — not merely "percent hit 100". */ + exhausted: boolean; + /** Epoch ms when the plan window rolls over, when upstream reports it. */ + nextResetAt?: number; +} + +interface KiroUsageStateEntry { + exhausted: boolean; + nextResetAt?: number; + ts: number; +} + +/** + * Exhaustion state, keyed exactly like the per-account quota cache in `quota.ts`. + * + * It is written only inside that cache's commit guard and cleared through the same + * logout/reconcile paths, so a removed account cannot leave a verdict behind for whatever + * account replaces it. + */ +const usageState = new Map(); + +function safeRegion(value: string | undefined): string | undefined { + return value && REGION_PATTERN.test(value) ? value : undefined; +} + +/** + * The profile ARN wins because an enterprise profile can live in a different region from + * the SSO session that minted the token. + */ +function usageRegion(ctx: KiroUsageContext): string { + return safeRegion(ctx.profileArn?.split(":")[3]) + ?? safeRegion(ctx.apiRegion) + ?? safeRegion(ctx.ssoRegion) + ?? "us-east-1"; +} + +export function kiroUsageManagementUrl(region: string): string { + return `https://management.${region}.kiro.dev/`; +} + +/** Credit balances are fractional; the integer fields round 695.17 down to 695. */ +function preciseNumber(row: Record, precise: string, whole: string): number | undefined { + return toFiniteNumber(row[precise]) ?? toFiniteNumber(row[whole]); +} + +function selectBreakdown(list: unknown): Record | null { + if (!Array.isArray(list)) return null; + const rows = list.map(asRecord).filter((row): row is Record => row !== null); + for (const wanted of RESOURCE_PRIORITY) { + const match = rows.find(row => String(row.resourceType ?? "").trim().toUpperCase() === wanted); + if (match) return match; + } + return null; +} + +function parseKiroUsage(body: unknown): KiroUsageSnapshot | null { + const payload = asRecord(body); + if (!payload) return null; + const breakdown = selectBreakdown(payload.usageBreakdownList); + if (!breakdown) return null; + + const used = preciseNumber(breakdown, "currentUsageWithPrecision", "currentUsage"); + const limit = preciseNumber(breakdown, "usageLimitWithPrecision", "usageLimit"); + if (used === undefined || limit === undefined || limit <= 0) return null; + + const percent = normalizePercent((used / limit) * 100); + if (percent === undefined) return null; + + const nextResetAt = normalizeResetAt(payload.nextDateReset); + const customWindows: ProviderQuotaWindow[] = []; + + // A trial allowance is a separate pool: folding it into the plan window would understate + // what the account can actually spend. + const trial = asRecord(breakdown.freeTrialInfo); + if (trial) { + const trialUsed = preciseNumber(trial, "currentUsageWithPrecision", "currentUsage"); + const trialLimit = preciseNumber(trial, "usageLimitWithPrecision", "usageLimit"); + if (trialUsed !== undefined && trialLimit !== undefined && trialLimit > 0) { + const trialPercent = normalizePercent((trialUsed / trialLimit) * 100); + if (trialPercent !== undefined) customWindows.push({ label: "Free trial", percent: trialPercent }); + } + } + + const quota: ProviderQuota = { + monthlyPercent: percent, + ...(nextResetAt !== undefined ? { monthlyResetAt: nextResetAt } : {}), + ...(customWindows.length > 0 ? { customWindows } : {}), + updatedAt: Date.now(), + }; + + // Enterprise accounts with overage enabled keep serving past the included limit, so + // "used >= limit" is not by itself a reason to stop routing to the account. + const overageEnabled = String(asRecord(payload.overageConfiguration)?.overageStatus ?? "") + .trim() + .toUpperCase() === "ENABLED"; + + return { + quota, + exhausted: used >= limit && !overageEnabled, + ...(nextResetAt !== undefined ? { nextResetAt } : {}), + }; +} + +/** + * Read one account's usage. Resolves `null` for any transport, status, or schema failure — + * the caller renders that as "unavailable" and keeps whatever it knew before. + * + * `userInfo` in the response carries an email and a user id. Both are read past and + * discarded here: nothing identifying an operator's person reaches the cache, the API, or + * a log line. + */ +export async function fetchKiroUsageSnapshot(ctx: KiroUsageContext): Promise { + const region = usageRegion(ctx); + const url = new URL(kiroUsageManagementUrl(region)); + url.searchParams.set("origin", "AI_EDITOR"); + url.searchParams.set("isEmailRequired", "true"); + if (ctx.profileArn) url.searchParams.set("profileArn", ctx.profileArn); + + // The modeled arguments appear in BOTH the query string and the body. That duplication is + // the observed Kiro CLI contract, not an oversight; we have no way to test which side the + // service actually reads, so we reproduce both. + const body: Record = { origin: "AI_EDITOR", isEmailRequired: true }; + if (ctx.profileArn) body.profileArn = ctx.profileArn; + + try { + const response = await fetch(url, { + method: "POST", + headers: { + authorization: `Bearer ${ctx.access}`, + "content-type": "application/x-amz-json-1.0", + accept: "application/json", + "x-amz-target": AMZ_USAGE_TARGET, + "x-amzn-codewhisperer-optout": "true", + }, + body: JSON.stringify(body), + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + if (!response.ok) return null; + const json = await readQuotaJson(response); + if (json === QUOTA_JSON_READ_FAILURE) return null; + return parseKiroUsage(json); + } catch { + return null; + } +} + +/** + * Assemble the probe context from ONE account-scoped snapshot. + * + * Reading the bearer and the routing metadata from a single snapshot is what keeps account + * A's token from being sent with account B's profile ARN — the same pairing class of defect + * #2841 fixed for Copilot origins. + */ +export async function kiroUsageContextForAccount(accountId: string): Promise { + const snapshot = await getValidAccessSnapshotForAccount("kiro", accountId); + return { + accountId, + access: snapshot.accessToken, + ...(snapshot.kiro?.profileArn ? { profileArn: snapshot.kiro.profileArn } : {}), + ...(snapshot.kiro?.apiRegion ? { apiRegion: snapshot.kiro.apiRegion } : {}), + ...(snapshot.kiro?.ssoRegion ? { ssoRegion: snapshot.kiro.ssoRegion } : {}), + }; +} + +/** Record exhaustion for a probed account. Called from the quota cache's commit guard. */ +export function commitKiroAccountUsageState(key: string, snapshot: KiroUsageSnapshot | null): void { + if (!snapshot) { + usageState.delete(key); + return; + } + usageState.set(key, { + exhausted: snapshot.exhausted, + ...(snapshot.nextResetAt !== undefined ? { nextResetAt: snapshot.nextResetAt } : {}), + ts: Date.now(), + }); +} + +/** + * Is this account known to be out of allowance right now? + * + * Returns `null` (unknown) rather than a stale `true`: an expired reading, or one whose + * reset time has already passed, must degrade to "try it again", never to "keep it parked". + */ +export function getKiroAccountExhaustion( + key: string, + now = Date.now(), +): { exhausted: boolean; nextResetAt?: number } | null { + const entry = usageState.get(key); + if (!entry) return null; + if (now - entry.ts >= ACCOUNT_QUOTA_TTL_MS) return null; + if (entry.nextResetAt !== undefined && entry.nextResetAt <= now) return null; + return { + exhausted: entry.exhausted, + ...(entry.nextResetAt !== undefined ? { nextResetAt: entry.nextResetAt } : {}), + }; +} + +/** Drop rows for one provider prefix, or all of them. Mirrors clearAccountQuotaCache. */ +export function clearKiroAccountUsageState(prefix?: string): void { + if (!prefix) { + usageState.clear(); + return; + } + for (const key of [...usageState.keys()]) { + if (key.startsWith(prefix)) usageState.delete(key); + } +} + +/** Drop rows whose account no longer exists. Mirrors reconcileProviderAccountQuotaRows. */ +export function reconcileKiroAccountUsageState(liveKeys: ReadonlySet): number { + let removed = 0; + for (const key of [...usageState.keys()]) { + if (liveKeys.has(key)) continue; + usageState.delete(key); + removed += 1; + } + return removed; +} diff --git a/src/providers/quota-types.ts b/src/providers/quota-types.ts new file mode 100644 index 0000000000..d71b49ea4d --- /dev/null +++ b/src/providers/quota-types.ts @@ -0,0 +1,36 @@ +/** + * Provider quota shapes, split out of `quota.ts` so a provider-specific quota module can + * describe its result without importing the aggregator that will consume it. + * + * `quota.ts` imports the Kiro usage module for its fetcher; if that module reached back + * into `quota.ts` for these types the two would depend on each other. Types have no + * runtime edge, but a cycle that exists only in the type graph is still a cycle, and it + * blocks any later attempt to load one side without the other. + */ + +export interface ProviderQuotaWindow { + label: string; + percent: number; + resetAt?: number; +} + +export interface ProviderQuotaCreditsUsd { + used: number; + limit: number; + remaining: number; + percent: number; + expiresAt?: number; + unlimited?: boolean; +} + +export interface ProviderQuota { + fiveHourPercent?: number; + fiveHourResetAt?: number; + weeklyPercent?: number; + weeklyResetAt?: number; + monthlyPercent?: number; + monthlyResetAt?: number; + customWindows?: ProviderQuotaWindow[]; + creditsUsd?: ProviderQuotaCreditsUsd; + updatedAt: number; +} diff --git a/src/providers/quota-wire.ts b/src/providers/quota-wire.ts new file mode 100644 index 0000000000..0abcb69efe --- /dev/null +++ b/src/providers/quota-wire.ts @@ -0,0 +1,97 @@ +/** + * Wire-level helpers shared by every provider quota reader: cache lifetimes, number and + * timestamp normalisation, and the bounded JSON reader. + * + * These lived inside `quota.ts` as module-private functions. A second quota module cannot + * import them from there without creating a cycle, and copying a TTL constant into a + * second file is how two copies of the same number drift apart. Everything here is pure or + * depends only on the bounded-body reader. + */ +import { readBoundedResponseBody } from "../lib/bounded-body"; + +/** Provider-level quota response cache lifetime (the dashboard/display path). */ +export const CACHE_TTL_MS = 5 * 60_000; + +/** + * Per-account quota cache lifetime. + * + * Deliberately longer than the provider-level TTL: this path multiplies by account count, + * and at least one upstream (Anthropic) rate-limits its usage endpoint under repeated + * probing. + */ +export const ACCOUNT_QUOTA_TTL_MS = 10 * 60_000; + +export const REQUEST_TIMEOUT_MS = 8_000; + +export const QUOTA_RESPONSE_MAX_BYTES = 512 * 1024; + +export const QUOTA_JSON_READ_FAILURE = Symbol("quota-json-read-failure"); + +/** Unix 0 / negative values are sentinels, not reset clocks (Command Code fiveHour.resetAt: 0). */ +export function epochMillis(value: number): number | undefined { + if (!Number.isFinite(value) || value <= 0) return undefined; + return value > 10_000_000_000 ? value : value * 1000; +} + +export function normalizeResetAt(value: unknown): number | undefined { + if (typeof value === "number" && Number.isFinite(value)) return epochMillis(value); + if (typeof value === "string" && value.trim()) { + const trimmed = value.trim(); + // Cursor Connect RPC returns billingCycleEnd as a unix-ms decimal string ("1771077734000"). + // Date.parse treats that as invalid; numeric epoch strings must be handled explicitly. + if (/^[+-]?\d+(\.\d+)?$/.test(trimmed)) { + const numeric = Number(trimmed); + return epochMillis(numeric); + } + const parsed = Date.parse(trimmed); + return Number.isFinite(parsed) && parsed > 0 ? parsed : undefined; + } + return undefined; +} + +export function toFiniteNumber(value: unknown): number | undefined { + if (typeof value === "number" && Number.isFinite(value)) return value; + if (typeof value === "string" && value.trim()) { + const parsed = Number(value); + return Number.isFinite(parsed) ? parsed : undefined; + } + return undefined; +} + +export function normalizePercent(value: unknown): number | undefined { + const numeric = toFiniteNumber(value); + return numeric === undefined ? undefined : Math.max(0, Math.min(100, numeric)); +} + +export function asRecord(value: unknown): Record | null { + return value && typeof value === "object" && !Array.isArray(value) ? value as Record : null; +} + +export async function readQuotaJson( + response: Response, + timeoutMs = REQUEST_TIMEOUT_MS, +): Promise { + const declaredLength = Number(response.headers.get("content-length")); + if (Number.isFinite(declaredLength) && declaredLength > QUOTA_RESPONSE_MAX_BYTES) { + try { + void response.body?.cancel( + new DOMException("Provider quota response is too large", "QuotaExceededError"), + ).catch(() => undefined); + } catch { + // Best-effort cancellation only. + } + return QUOTA_JSON_READ_FAILURE; + } + + try { + const bounded = await readBoundedResponseBody(response, { + maxBytes: QUOTA_RESPONSE_MAX_BYTES, + totalTimeoutMs: timeoutMs, + inactivityTimeoutMs: timeoutMs, + }); + if (bounded.oversized || bounded.truncated || !bounded.displaySafe) return QUOTA_JSON_READ_FAILURE; + return JSON.parse(bounded.text) as unknown; + } catch { + return QUOTA_JSON_READ_FAILURE; + } +} diff --git a/src/providers/quota.ts b/src/providers/quota.ts index 63c74b9cd6..7525e1b064 100644 --- a/src/providers/quota.ts +++ b/src/providers/quota.ts @@ -23,21 +23,35 @@ import { sweepExpiredOnWrite, type GenerationContext, } from "../lib/state-store-sweeper"; -import { readBoundedResponseBody } from "../lib/bounded-body"; +import { + ACCOUNT_QUOTA_TTL_MS, + asRecord, + CACHE_TTL_MS, + normalizePercent, + normalizeResetAt, + QUOTA_JSON_READ_FAILURE, + readQuotaJson, + REQUEST_TIMEOUT_MS, + toFiniteNumber, +} from "./quota-wire"; import { aggregateCodexPoolCapacity, CODEX_CAPACITY_MAX_QUOTA_AGE_MS, type CodexCapacityAggregation, type CodexCapacityQuota, } from "./codex-capacity"; +import type { + ProviderQuota, + ProviderQuotaCreditsUsd, + ProviderQuotaWindow, +} from "./quota-types"; + +export type { ProviderQuota, ProviderQuotaCreditsUsd, ProviderQuotaWindow } from "./quota-types"; /** Match oauth/index REFRESH_SKEW_MS — use stored access without refresh when still fresh. */ const ACCOUNT_TOKEN_SKEW_MS = 60_000; - -const CACHE_TTL_MS = 5 * 60_000; -const REQUEST_TIMEOUT_MS = 8_000; /** Successful provider quota payloads are small; reject oversized or stalled JSON before parsing. */ -export const QUOTA_RESPONSE_MAX_BYTES = 512 * 1024; +export { QUOTA_RESPONSE_MAX_BYTES } from "./quota-wire"; const KIMI_CODE_BASE_URL = "https://api.kimi.com/coding/v1"; const KIMI_CODE_USAGE_URL = `${KIMI_CODE_BASE_URL}/usages`; const COMMAND_CODE_BASE_URL = "https://api.commandcode.ai"; @@ -75,33 +89,6 @@ export function setProviderQuotaBeforePublishForTests( const TERMINAL_QUOTA_FAILURE = Symbol("terminal-quota-failure"); type ProviderQuotaProbeResult = ProviderQuotaReport | null | typeof TERMINAL_QUOTA_FAILURE; -export interface ProviderQuotaWindow { - label: string; - percent: number; - resetAt?: number; -} - -export interface ProviderQuotaCreditsUsd { - used: number; - limit: number; - remaining: number; - percent: number; - expiresAt?: number; - unlimited?: boolean; -} - -export interface ProviderQuota { - fiveHourPercent?: number; - fiveHourResetAt?: number; - weeklyPercent?: number; - weeklyResetAt?: number; - monthlyPercent?: number; - monthlyResetAt?: number; - customWindows?: ProviderQuotaWindow[]; - creditsUsd?: ProviderQuotaCreditsUsd; - updatedAt: number; -} - export interface ProviderQuotaReport { provider: string; label: string; @@ -260,77 +247,6 @@ function providerLabel(providerId: string): string { return getProviderRegistryEntry(providerId)?.label ?? providerId; } -function normalizeResetAt(value: unknown): number | undefined { - if (typeof value === "number" && Number.isFinite(value)) return epochMillis(value); - if (typeof value === "string" && value.trim()) { - const trimmed = value.trim(); - // Cursor Connect RPC returns billingCycleEnd as a unix-ms decimal string ("1771077734000"). - // Date.parse treats that as invalid; numeric epoch strings must be handled explicitly. - if (/^[+-]?\d+(\.\d+)?$/.test(trimmed)) { - const numeric = Number(trimmed); - return epochMillis(numeric); - } - const parsed = Date.parse(trimmed); - return Number.isFinite(parsed) && parsed > 0 ? parsed : undefined; - } - return undefined; -} - -/** Unix 0 / negative values are sentinels, not reset clocks (Command Code fiveHour.resetAt: 0). */ -function epochMillis(value: number): number | undefined { - if (!Number.isFinite(value) || value <= 0) return undefined; - return value > 10_000_000_000 ? value : value * 1000; -} - -function toFiniteNumber(value: unknown): number | undefined { - if (typeof value === "number" && Number.isFinite(value)) return value; - if (typeof value === "string" && value.trim()) { - const parsed = Number(value); - return Number.isFinite(parsed) ? parsed : undefined; - } - return undefined; -} - -function normalizePercent(value: unknown): number | undefined { - const numeric = toFiniteNumber(value); - return numeric === undefined ? undefined : Math.max(0, Math.min(100, numeric)); -} - -function asRecord(value: unknown): Record | null { - return value && typeof value === "object" && !Array.isArray(value) ? value as Record : null; -} - -const QUOTA_JSON_READ_FAILURE = Symbol("quota-json-read-failure"); - -async function readQuotaJson( - response: Response, - timeoutMs = REQUEST_TIMEOUT_MS, -): Promise { - const declaredLength = Number(response.headers.get("content-length")); - if (Number.isFinite(declaredLength) && declaredLength > QUOTA_RESPONSE_MAX_BYTES) { - try { - void response.body?.cancel( - new DOMException("Provider quota response is too large", "QuotaExceededError"), - ).catch(() => undefined); - } catch { - // Best-effort cancellation only. - } - return QUOTA_JSON_READ_FAILURE; - } - - try { - const bounded = await readBoundedResponseBody(response, { - maxBytes: QUOTA_RESPONSE_MAX_BYTES, - totalTimeoutMs: timeoutMs, - inactivityTimeoutMs: timeoutMs, - }); - if (bounded.oversized || bounded.truncated || !bounded.displaySafe) return QUOTA_JSON_READ_FAILURE; - return JSON.parse(bounded.text) as unknown; - } catch { - return QUOTA_JSON_READ_FAILURE; - } -} - /** Test-only access to the quota reader's deadline and cancellation contract. */ export async function readProviderQuotaJsonForTests(response: Response, timeoutMs: number): Promise { const result = await readQuotaJson(response, timeoutMs); @@ -1414,15 +1330,13 @@ async function fetchAnthropicQuota(provider: string): Promise { + globalThis.fetch = realFetch; + clearKiroAccountUsageState(); +}); + +/** Capture the outbound request so region/host/parameter assertions read the real call. */ +function stubUsageResponse(payload: unknown, status = 200): { calls: Request[] } { + const calls: Request[] = []; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + calls.push(new Request(input instanceof Request ? input : String(input), init)); + return new Response(JSON.stringify(payload), { + status, + headers: { "content-type": "application/json" }, + }); + }) as typeof fetch; + return { calls }; +} + +const baseContext = { accountId: "acct-1", access: "tok-1" }; + +function breakdown(extra: Record = {}) { + return { + resourceType: "AGENTIC_REQUEST", + currentUsageWithPrecision: 147.82, + currentUsage: 147, + usageLimitWithPrecision: 1000, + usageLimit: 1000, + unit: "CREDITS", + ...extra, + }; +} + +describe("Kiro usage limits", () => { + test("maps an agentic-request breakdown onto the monthly window", async () => { + stubUsageResponse({ + usageBreakdownList: [breakdown()], + nextDateReset: 1785542400, + subscriptionInfo: { subscriptionTitle: "KIRO PRO" }, + }); + const snapshot = await fetchKiroUsageSnapshot(baseContext); + expect(snapshot?.quota.monthlyPercent).toBeCloseTo(14.782, 3); + expect(snapshot?.quota.monthlyResetAt).toBe(1785542400 * 1000); + expect(snapshot?.exhausted).toBe(false); + }); + + test("selects CREDIT by resource type rather than falling back to the first row", async () => { + stubUsageResponse({ + usageBreakdownList: [ + { resourceType: "SOMETHING_ELSE", currentUsage: 900, usageLimit: 1000 }, + { resourceType: "CREDIT", currentUsage: 100, usageLimit: 1000 }, + ], + }); + const snapshot = await fetchKiroUsageSnapshot(baseContext); + // Index 0 would report 90%; selecting by resourceType reports the credit pool. + expect(snapshot?.quota.monthlyPercent).toBe(10); + }); + + test("reports unknown when no recognised resource type is present", async () => { + stubUsageResponse({ + usageBreakdownList: [{ resourceType: "FUTURE_POOL", currentUsage: 5, usageLimit: 10 }], + }); + expect(await fetchKiroUsageSnapshot(baseContext)).toBeNull(); + }); + + test("prefers the precise fields over their rounded twins", async () => { + stubUsageResponse({ + usageBreakdownList: [breakdown({ + currentUsageWithPrecision: 695.17, + currentUsage: 695, + usageLimitWithPrecision: 1000, + usageLimit: 1000, + })], + }); + const snapshot = await fetchKiroUsageSnapshot(baseContext); + expect(snapshot?.quota.monthlyPercent).toBeCloseTo(69.517, 3); + }); + + test("an overage-enabled account past its limit is not exhausted", async () => { + stubUsageResponse({ + usageBreakdownList: [breakdown({ currentUsageWithPrecision: 1200, usageLimitWithPrecision: 1000 })], + overageConfiguration: { overageStatus: "ENABLED" }, + }); + const snapshot = await fetchKiroUsageSnapshot(baseContext); + expect(snapshot?.exhausted).toBe(false); + expect(snapshot?.quota.monthlyPercent).toBe(100); + }); + + test("a spent account without overage is exhausted", async () => { + stubUsageResponse({ + usageBreakdownList: [breakdown({ currentUsageWithPrecision: 1000, usageLimitWithPrecision: 1000 })], + overageConfiguration: { overageStatus: "DISABLED" }, + }); + expect((await fetchKiroUsageSnapshot(baseContext))?.exhausted).toBe(true); + }); + + test("a free-trial allowance is reported as its own window", async () => { + stubUsageResponse({ + usageBreakdownList: [breakdown({ + freeTrialInfo: { freeTrialStatus: "ACTIVE", currentUsageWithPrecision: 250, usageLimitWithPrecision: 500 }, + })], + }); + const snapshot = await fetchKiroUsageSnapshot(baseContext); + expect(snapshot?.quota.customWindows).toEqual([{ label: "Free trial", percent: 50 }]); + }); + + test("the operator's email never leaves the module", async () => { + stubUsageResponse({ + usageBreakdownList: [breakdown()], + userInfo: { email: "operator@example.com", userId: "user-123" }, + }); + const snapshot = await fetchKiroUsageSnapshot(baseContext); + expect(JSON.stringify(snapshot)).not.toContain("operator@example.com"); + expect(JSON.stringify(snapshot)).not.toContain("user-123"); + }); + + test("the profile ARN region wins over the stored api region", async () => { + const stub = stubUsageResponse({ usageBreakdownList: [breakdown()] }); + await fetchKiroUsageSnapshot({ + ...baseContext, + profileArn: "arn:aws:codewhisperer:eu-central-1:123456789012:profile/ABCD", + apiRegion: "us-east-1", + }); + expect(new URL(stub.calls[0]!.url).host).toBe("management.eu-central-1.kiro.dev"); + }); + + test("a crafted ARN region falls through instead of reaching the hostname", async () => { + const stub = stubUsageResponse({ usageBreakdownList: [breakdown()] }); + await fetchKiroUsageSnapshot({ + ...baseContext, + profileArn: "arn:aws:codewhisperer:evil.example.com:123456789012:profile/ABCD", + apiRegion: "us-west-2", + }); + const host = new URL(stub.calls[0]!.url).host; + expect(host).toBe("management.us-west-2.kiro.dev"); + expect(host).not.toContain("evil.example.com"); + }); + + test("hostile stored regions fall back to the default region", async () => { + const stub = stubUsageResponse({ usageBreakdownList: [breakdown()] }); + await fetchKiroUsageSnapshot({ + ...baseContext, + apiRegion: "../../evil", + ssoRegion: "attacker.example.com", + }); + const host = new URL(stub.calls[0]!.url).host; + expect(host).toBe("management.us-east-1.kiro.dev"); + expect(host).not.toContain("evil"); + expect(host).not.toContain("attacker"); + }); + + test("the modeled arguments are sent in both the query and the body", async () => { + const stub = stubUsageResponse({ usageBreakdownList: [breakdown()] }); + const arn = "arn:aws:codewhisperer:us-east-1:123456789012:profile/ABCD"; + await fetchKiroUsageSnapshot({ ...baseContext, profileArn: arn }); + const request = stub.calls[0]!; + const url = new URL(request.url); + expect(url.searchParams.get("origin")).toBe("AI_EDITOR"); + expect(url.searchParams.get("isEmailRequired")).toBe("true"); + expect(url.searchParams.get("profileArn")).toBe(arn); + expect(request.headers.get("x-amz-target")).toBe("AmazonCodeWhispererService.GetUsageLimits"); + expect(await request.json()).toEqual({ origin: "AI_EDITOR", isEmailRequired: true, profileArn: arn }); + }); + + test("an upstream failure resolves unknown instead of throwing", async () => { + for (const status of [401, 429, 500]) { + stubUsageResponse({ message: "nope" }, status); + expect(await fetchKiroUsageSnapshot(baseContext)).toBeNull(); + } + }); + + test("a transport failure resolves unknown", async () => { + globalThis.fetch = (async () => { throw new Error("offline"); }) as typeof fetch; + expect(await fetchKiroUsageSnapshot(baseContext)).toBeNull(); + }); + + test("management host construction is region-scoped", () => { + expect(kiroUsageManagementUrl("us-east-1")).toBe("https://management.us-east-1.kiro.dev/"); + }); +}); + +describe("Kiro exhaustion state", () => { + const key = "kiro\u0000acct-1"; + + test("a fresh exhausted verdict is readable", () => { + commitKiroAccountUsageState(key, { quota: { updatedAt: Date.now() }, exhausted: true, nextResetAt: Date.now() + 60_000 }); + expect(getKiroAccountExhaustion(key)?.exhausted).toBe(true); + }); + + test("a verdict older than the account TTL degrades to unknown", () => { + const now = Date.now(); + commitKiroAccountUsageState(key, { quota: { updatedAt: now }, exhausted: true, nextResetAt: now + 60 * 60_000 }); + expect(getKiroAccountExhaustion(key, now + 11 * 60_000)).toBeNull(); + }); + + test("a verdict whose reset has passed degrades to unknown", () => { + const now = Date.now(); + commitKiroAccountUsageState(key, { quota: { updatedAt: now }, exhausted: true, nextResetAt: now + 1_000 }); + expect(getKiroAccountExhaustion(key, now + 2_000)).toBeNull(); + }); + + test("a null snapshot clears any previous verdict", () => { + commitKiroAccountUsageState(key, { quota: { updatedAt: Date.now() }, exhausted: true }); + commitKiroAccountUsageState(key, null); + expect(getKiroAccountExhaustion(key)).toBeNull(); + }); + + test("clearing by provider prefix drops that provider's rows", () => { + commitKiroAccountUsageState(key, { quota: { updatedAt: Date.now() }, exhausted: true }); + clearKiroAccountUsageState("kiro\u0000"); + expect(getKiroAccountExhaustion(key)).toBeNull(); + }); + + test("reconciliation drops rows for accounts that no longer exist", () => { + commitKiroAccountUsageState(key, { quota: { updatedAt: Date.now() }, exhausted: true }); + expect(reconcileKiroAccountUsageState(new Set())).toBe(1); + expect(getKiroAccountExhaustion(key)).toBeNull(); + }); + + test("reconciliation keeps rows for live accounts", () => { + commitKiroAccountUsageState(key, { quota: { updatedAt: Date.now() }, exhausted: true }); + expect(reconcileKiroAccountUsageState(new Set([key]))).toBe(0); + expect(getKiroAccountExhaustion(key)?.exhausted).toBe(true); + }); +}); From be83fbb7fed1999ba109d4f7a050a996afabd008 Mon Sep 17 00:00:00 2001 From: bitkyc08-arch Date: Sat, 29 Aug 2026 12:40:05 +0900 Subject: [PATCH 3/6] feat(kiro): report per-account and provider-level Kiro quota Kiro joins the per-account quota seam, reusing its TTL, negative caching, in-flight join and generation guard. Exhaustion state commits under the same guard and retires with the quota row, so a removed account cannot leave a verdict for its replacement. --- src/providers/quota.ts | 63 +++++++++- tests/kiro-account-quota.test.ts | 179 +++++++++++++++++++++++++++ tests/provider-account-quota.test.ts | 7 +- 3 files changed, 244 insertions(+), 5 deletions(-) create mode 100644 tests/kiro-account-quota.test.ts diff --git a/src/providers/quota.ts b/src/providers/quota.ts index 7525e1b064..55ff088d27 100644 --- a/src/providers/quota.ts +++ b/src/providers/quota.ts @@ -45,6 +45,14 @@ import type { ProviderQuotaCreditsUsd, ProviderQuotaWindow, } from "./quota-types"; +import { + clearKiroAccountUsageState, + commitKiroAccountUsageState, + fetchKiroUsageSnapshot, + type KiroUsageSnapshot, + kiroUsageContextForAccount, + reconcileKiroAccountUsageState, +} from "./kiro-usage"; export type { ProviderQuota, ProviderQuotaCreditsUsd, ProviderQuotaWindow } from "./quota-types"; @@ -1325,6 +1333,32 @@ async function fetchAnthropicQuota(provider: string): Promise { + const probedAccountId = getAccountSet("kiro")?.activeAccountId; + if (!probedAccountId) return null; + const probedAccountKey = accountCacheKey("kiro", probedAccountId); + const writerGeneration = captureConfigGeneration(); + let snapshot: KiroUsageSnapshot | null; + try { + snapshot = await fetchKiroUsageSnapshot(await kiroUsageContextForAccount(probedAccountId)); + } catch { + return null; + } + if (!snapshot) return null; + if (mayCommitAccountQuotaKey(probedAccountKey, writerGeneration)) { + accountQuotaCache.set(probedAccountKey, { ts: Date.now(), quota: snapshot.quota }); + commitKiroAccountUsageState(probedAccountKey, snapshot); + } + return report(provider, "kiro:usage-limits", snapshot.quota); +} + // --------------------------------------------------------------------------- // Per-account quota (multiauth) // --------------------------------------------------------------------------- @@ -1366,7 +1400,7 @@ export interface ProviderAccountQuota { /** Providers whose per-account quota can be probed. Extend as other OAuth APIs are covered. */ export function supportsPerAccountQuota(provider: string): boolean { - return provider === "anthropic"; + return provider === "anthropic" || provider === "kiro"; } function accountCacheKey(provider: string, accountId: string): string { @@ -1414,6 +1448,9 @@ export function reconcileProviderAccountQuotaRows(context: GenerationContext): n accountQuotaCache.delete(key); removed += 1; } + // Kiro exhaustion rows are keyed identically, so they retire with their quota row; a + // verdict outliving its account would hand the replacement a cooldown it never earned. + removed += reconcileKiroAccountUsageState(context.oauthAccountKeys); if (cache) { const reports = cache.response.reports.filter(report => context.providerNames.has(report.provider)); removed += cache.response.reports.length - reports.length; @@ -1437,12 +1474,14 @@ export function clearAccountQuotaCache(provider?: string): void { if (!provider) { accountQuotaCache.clear(); accountQuotaInflight.clear(); + clearKiroAccountUsageState(); return; } const prefix = `${provider}\u0000`; for (const key of [...accountQuotaCache.keys()]) { if (key.startsWith(prefix)) accountQuotaCache.delete(key); } + clearKiroAccountUsageState(prefix); // Drop in-flight probes too so a late resolve cannot repopulate after logout/remove. for (const key of [...accountQuotaInflight.keys()]) { if (key.startsWith(prefix)) accountQuotaInflight.delete(key); @@ -1485,8 +1524,21 @@ async function fetchAccountQuota( const probe = (async (): Promise => { try { - const token = await getTokenForAccountQuotaProbe(provider, accountId); - const quota = await fetchAnthropicUsageQuota(token); + let quota: ProviderQuota | null; + let kiroSnapshot: KiroUsageSnapshot | null = null; + if (provider === "kiro") { + // Kiro resolves the bearer and its routing metadata from ONE account-scoped + // snapshot. It deliberately does not use getTokenForAccountQuotaProbe: that + // helper refuses to refresh a background `local-cli` slot because Anthropic's + // lock can adopt a mismatched Claude CLI identity, but Kiro marks every + // CLI-imported credential `local-cli`, so the same rule would blank the quota of + // every inactive pool account the moment its token expired. + kiroSnapshot = await fetchKiroUsageSnapshot(await kiroUsageContextForAccount(accountId)); + quota = kiroSnapshot?.quota ?? null; + } else { + const token = await getTokenForAccountQuotaProbe(provider, accountId); + quota = await fetchAnthropicUsageQuota(token); + } if (!quota) { // Preserve last-good bars and mark unavailable; advance TTL so failures // negative-cache instead of re-probing on every GUI poll. @@ -1497,6 +1549,7 @@ async function fetchAccountQuota( }; if (mayCommitAccountQuotaKey(key, writerGeneration)) { accountQuotaCache.set(key, entry); + if (provider === "kiro") commitKiroAccountUsageState(key, null); sweepExpiredOnWrite(entry.ts); } return entry; @@ -1504,6 +1557,9 @@ async function fetchAccountQuota( const entry: AccountQuotaCacheEntry = { ts: Date.now(), quota }; if (mayCommitAccountQuotaKey(key, writerGeneration)) { accountQuotaCache.set(key, entry); + // Exhaustion state rides the SAME commit guard as the quota row: a probe from a + // superseded config generation must not publish either half. + if (provider === "kiro") commitKiroAccountUsageState(key, kiroSnapshot); sweepExpiredOnWrite(entry.ts); } return entry; @@ -2113,6 +2169,7 @@ async function maybeFetchProviderQuota( if (provider.authMode === "oauth" && name === "anthropic") return fetchAnthropicQuota(name); if (provider.authMode === "oauth" && name === "cursor") return fetchCursorQuota(name); if (provider.authMode === "oauth" && name === "google-antigravity") return fetchAntigravityQuota(name, provider); + if (provider.authMode === "oauth" && name === "kiro") return fetchKiroQuota(name); // Kimi Code `/usages` accepts OAuth or coding-plan API keys, but only on the canonical // host and only for real key auth — forward/local modes carry no credential of ours. if (provider.authMode === "oauth" && name === "kimi") return fetchKimiQuota(name, provider); diff --git a/tests/kiro-account-quota.test.ts b/tests/kiro-account-quota.test.ts new file mode 100644 index 0000000000..3099c9c6ba --- /dev/null +++ b/tests/kiro-account-quota.test.ts @@ -0,0 +1,179 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { saveCredential } from "../src/oauth/store"; +import { + clearAccountQuotaCache, + clearProviderQuotaCache, + fetchProviderAccountQuotas, + getCachedProviderAccountQuota, + supportsPerAccountQuota, +} from "../src/providers/quota"; +import { getKiroAccountExhaustion } from "../src/providers/kiro-usage"; + +const originalFetch = globalThis.fetch; +const previousOpencodexHome = process.env.OPENCODEX_HOME; +let opencodexHome: string; + +const ARN_A = "arn:aws:codewhisperer:us-east-1:111111111111:profile/AAAA"; +const ARN_B = "arn:aws:codewhisperer:eu-central-1:222222222222:profile/BBBB"; + +/** Two logged-in Kiro accounts, each with its own bearer AND its own routing metadata. */ +async function seedTwoKiroAccounts(): Promise { + const expires = Date.now() + 60 * 60_000; + await saveCredential("kiro", { + access: "token-a", refresh: "refresh-a", expires, + accountId: "kiro-a", email: "a@example.com", + kiro: { profileArn: ARN_A, apiRegion: "us-east-1", ssoRegion: "us-east-1" }, + }); + await saveCredential("kiro", { + access: "token-b", refresh: "refresh-b", expires, + accountId: "kiro-b", email: "b@example.com", + kiro: { profileArn: ARN_B, apiRegion: "eu-central-1", ssoRegion: "eu-central-1" }, + }); +} + +function usagePayload(used: number, limit: number, overage = "DISABLED"): string { + return JSON.stringify({ + usageBreakdownList: [{ + resourceType: "AGENTIC_REQUEST", + currentUsageWithPrecision: used, + usageLimitWithPrecision: limit, + unit: "CREDITS", + }], + overageConfiguration: { overageStatus: overage }, + nextDateReset: Math.floor(Date.now() / 1000) + 3 * 24 * 3600, + }); +} + +beforeEach(() => { + opencodexHome = mkdtempSync(join(tmpdir(), "ocx-kiro-quota-")); + process.env.OPENCODEX_HOME = opencodexHome; + clearAccountQuotaCache(); + clearProviderQuotaCache(); +}); + +afterEach(() => { + globalThis.fetch = originalFetch; + if (previousOpencodexHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = previousOpencodexHome; + rmSync(opencodexHome, { recursive: true, force: true }); + clearAccountQuotaCache(); + clearProviderQuotaCache(); +}); + +describe("Kiro per-account quota", () => { + test("the seam is open for kiro", () => { + expect(supportsPerAccountQuota("kiro")).toBe(true); + }); + + test("each account is probed with its own bearer and reported separately", async () => { + await seedTwoKiroAccounts(); + const seen: Array<{ token: string; host: string; arn: string | null }> = []; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const request = new Request(input instanceof Request ? input : String(input), init); + const url = new URL(request.url); + seen.push({ + token: request.headers.get("authorization") ?? "", + host: url.host, + arn: url.searchParams.get("profileArn"), + }); + const used = request.headers.get("authorization") === "Bearer token-a" ? 100 : 900; + return new Response(usagePayload(used, 1000), { status: 200 }); + }) as typeof fetch; + + const rows = await fetchProviderAccountQuotas("kiro"); + expect(rows).toHaveLength(2); + expect(rows.map(r => r.quota?.monthlyPercent).sort()).toEqual([10, 90]); + expect(seen).toHaveLength(2); + }); + + test("a rotated bearer never travels with another account's profile ARN or region", async () => { + await seedTwoKiroAccounts(); + const pairs: Array<{ token: string; host: string; arn: string | null }> = []; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const request = new Request(input instanceof Request ? input : String(input), init); + const url = new URL(request.url); + pairs.push({ + token: request.headers.get("authorization") ?? "", + host: url.host, + arn: url.searchParams.get("profileArn"), + }); + return new Response(usagePayload(100, 1000), { status: 200 }); + }) as typeof fetch; + + await fetchProviderAccountQuotas("kiro"); + + // This is the #2841 invariant applied to Kiro: token and routing metadata must come + // from the SAME account record, so no pair may mix A's bearer with B's ARN. + for (const pair of pairs) { + if (pair.token === "Bearer token-a") { + expect(pair.arn).toBe(ARN_A); + expect(pair.host).toBe("management.us-east-1.kiro.dev"); + } else { + expect(pair.token).toBe("Bearer token-b"); + expect(pair.arn).toBe(ARN_B); + expect(pair.host).toBe("management.eu-central-1.kiro.dev"); + } + } + expect(pairs).toHaveLength(2); + }); + + test("one failing account leaves the other's bars intact", async () => { + await seedTwoKiroAccounts(); + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const request = new Request(input instanceof Request ? input : String(input), init); + if (request.headers.get("authorization") === "Bearer token-b") { + return new Response("{}", { status: 401 }); + } + return new Response(usagePayload(250, 1000), { status: 200 }); + }) as typeof fetch; + + const rows = await fetchProviderAccountQuotas("kiro"); + const healthy = rows.find(r => r.quota?.monthlyPercent === 25); + const broken = rows.find(r => r.unavailable); + expect(healthy).toBeDefined(); + expect(broken?.quota).toBeNull(); + }); + + test("a second read inside the TTL makes no upstream call", async () => { + await seedTwoKiroAccounts(); + let calls = 0; + globalThis.fetch = (async () => { + calls += 1; + return new Response(usagePayload(100, 1000), { status: 200 }); + }) as typeof fetch; + + await fetchProviderAccountQuotas("kiro"); + expect(calls).toBe(2); + await fetchProviderAccountQuotas("kiro"); + expect(calls).toBe(2); + }); + + test("exhaustion state is recorded next to the quota row and cleared with it", async () => { + await seedTwoKiroAccounts(); + globalThis.fetch = (async () => new Response(usagePayload(1000, 1000), { status: 200 })) as typeof fetch; + + const rows = await fetchProviderAccountQuotas("kiro"); + // Kiro keys accounts by their stable profile ARN, so read the id the store actually + // assigned rather than assuming the seed value. + const accountId = rows[0]!.accountId; + const key = `kiro\u0000${accountId}`; + expect(getKiroAccountExhaustion(key)?.exhausted).toBe(true); + expect(getCachedProviderAccountQuota("kiro", accountId)?.monthlyPercent).toBe(100); + + clearAccountQuotaCache("kiro"); + expect(getKiroAccountExhaustion(key)).toBeNull(); + expect(getCachedProviderAccountQuota("kiro", accountId)).toBeNull(); + }); + + test("an overage-enabled account past its limit is not marked exhausted", async () => { + await seedTwoKiroAccounts(); + globalThis.fetch = (async () => new Response(usagePayload(1500, 1000, "ENABLED"), { status: 200 })) as typeof fetch; + + const rows = await fetchProviderAccountQuotas("kiro"); + const key = `kiro\u0000${rows[0]!.accountId}`; + expect(getKiroAccountExhaustion(key)?.exhausted).toBe(false); + }); +}); diff --git a/tests/provider-account-quota.test.ts b/tests/provider-account-quota.test.ts index 634c663ca2..2bb21f13a2 100644 --- a/tests/provider-account-quota.test.ts +++ b/tests/provider-account-quota.test.ts @@ -203,10 +203,13 @@ describe("fetchProviderAccountQuotas", () => { test("providers without a per-account usage API are skipped", async () => { expect(supportsPerAccountQuota("anthropic")).toBe(true); - expect(supportsPerAccountQuota("kiro")).toBe(false); + // Kiro joined this list once it grew a usage reader; xAI has no per-account usage API, + // so it now carries the "unsupported providers never reach the network" contract. + expect(supportsPerAccountQuota("kiro")).toBe(true); + expect(supportsPerAccountQuota("xai")).toBe(false); let called = false; globalThis.fetch = (async () => { called = true; return new Response("{}", { status: 200 }); }) as typeof fetch; - expect(await fetchProviderAccountQuotas("kiro")).toEqual([]); + expect(await fetchProviderAccountQuotas("xai")).toEqual([]); expect(called).toBe(false); }); From 67826ecb1c5601b098e32bb5ceb8ac98a6b7fa27 Mon Sep 17 00:00:00 2001 From: bitkyc08-arch Date: Sat, 29 Aug 2026 12:41:52 +0900 Subject: [PATCH 4/6] feat(oauth): rotate toward the account with known headroom on a 429 Rotation ordered the roster blind, so the account after the one that just 429'd could itself be spent. Candidates are now ranked healthy > unknown > exhausted, built from the existing ring so a provider without per-account quota keeps its traversal exactly. --- src/oauth/account-quota-rank.ts | 94 +++++++++++++++++++++ src/oauth/generic-account-failover.ts | 21 +++-- tests/generic-oauth-failover.test.ts | 20 +++++ tests/kiro-pool-rank.test.ts | 117 ++++++++++++++++++++++++++ 4 files changed, 245 insertions(+), 7 deletions(-) create mode 100644 src/oauth/account-quota-rank.ts create mode 100644 tests/kiro-pool-rank.test.ts diff --git a/src/oauth/account-quota-rank.ts b/src/oauth/account-quota-rank.ts new file mode 100644 index 0000000000..3cc0e29e2d --- /dev/null +++ b/src/oauth/account-quota-rank.ts @@ -0,0 +1,94 @@ +/** + * Order failover candidates by what we know about their remaining allowance. + * + * Rotation without this walks the roster blind: the account right after the one that just + * 429'd may itself be spent, so the request burns a second rotation from a budget of three + * to learn what a cached quota row already knew. + * + * Deliberately NOT a scoring function. Percentages from different providers measure + * different things, and a weight would invite tuning a number nobody can validate. Three + * categories answer the only question rotation asks — "which of these is most likely to + * serve the retry" — and within the healthy group a simple headroom sort is enough. + */ +import { getCachedProviderAccountQuota } from "../providers/quota"; +import { getKiroAccountExhaustion } from "../providers/kiro-usage"; + +/** Lower sorts earlier. Unknown sits between measured-healthy and measured-empty. */ +const RANK_HEALTHY = 0; +const RANK_UNKNOWN = 1; +const RANK_EXHAUSTED = 2; + +interface Ranked { + id: string; + bucket: number; + /** Remaining percentage points, descending within the healthy bucket. */ + headroom: number; + /** Preserves the caller's ring order for ties. */ + index: number; +} + +/** + * Remaining headroom across every window the provider reports. + * + * The minimum wins: an account at 5% of its five-hour window is unusable right now even if + * its monthly allowance is barely touched. + */ +function headroomOf(provider: string, accountId: string): number | null { + const quota = getCachedProviderAccountQuota(provider, accountId); + if (!quota) return null; + const percents = [ + quota.fiveHourPercent, + quota.weeklyPercent, + quota.monthlyPercent, + ...(quota.customWindows ?? []).map(window => window.percent), + ].filter((value): value is number => typeof value === "number"); + if (percents.length === 0) return null; + return 100 - Math.max(...percents); +} + +/** + * Order candidates best-first. + * + * Returns the input untouched when no candidate has quota evidence, which keeps every + * provider without per-account quota on exactly the behaviour it has today. + */ +export function rankAccountsByHeadroom(provider: string, ring: readonly string[]): string[] { + if (ring.length < 2) return [...ring]; + + let sawEvidence = false; + const ranked: Ranked[] = ring.map((id, index) => { + // A provider-declared exhaustion verdict outranks the percentage: an account may sit at + // 100% and still be servable when overage is enabled, and the verdict knows that. + const exhaustion = provider === "kiro" ? getKiroAccountExhaustion(`${provider}\u0000${id}`) : null; + const headroom = headroomOf(provider, id); + if (exhaustion !== null || headroom !== null) sawEvidence = true; + + if (exhaustion?.exhausted === true) return { id, bucket: RANK_EXHAUSTED, headroom: 0, index }; + if (headroom === null) return { id, bucket: RANK_UNKNOWN, headroom: 0, index }; + return { id, bucket: RANK_HEALTHY, headroom, index }; + }); + + if (!sawEvidence) return [...ring]; + + return ranked + .sort((a, b) => (a.bucket - b.bucket) || (b.headroom - a.headroom) || (a.index - b.index)) + .map(entry => entry.id); +} + +/** + * How long to cool an account that just 429'd, when we know its allowance is spent. + * + * A monthly-exhausted account retried every minute is pure waste, but an upstream reset + * date is not something to trust unbounded — the clamp keeps a bogus far-future value from + * parking an account for weeks, and a near-instant one from being pointless. + */ +const MIN_EXHAUSTED_COOLDOWN_MS = 5 * 60_000; +const MAX_EXHAUSTED_COOLDOWN_MS = 24 * 60 * 60_000; + +export function exhaustedCooldownMs(provider: string, accountId: string, now = Date.now()): number | null { + if (provider !== "kiro") return null; + const exhaustion = getKiroAccountExhaustion(`${provider}\u0000${accountId}`, now); + if (!exhaustion?.exhausted) return null; + const untilReset = exhaustion.nextResetAt === undefined ? MIN_EXHAUSTED_COOLDOWN_MS : exhaustion.nextResetAt - now; + return Math.min(Math.max(untilReset, MIN_EXHAUSTED_COOLDOWN_MS), MAX_EXHAUSTED_COOLDOWN_MS); +} diff --git a/src/oauth/generic-account-failover.ts b/src/oauth/generic-account-failover.ts index 018c46c987..cec3e6321d 100644 --- a/src/oauth/generic-account-failover.ts +++ b/src/oauth/generic-account-failover.ts @@ -16,6 +16,7 @@ */ import { getAccountSet } from "./store"; import { getValidAccessSnapshotForAccount, type OAuthAccessSnapshot } from "./index"; +import { exhaustedCooldownMs, rankAccountsByHeadroom } from "./account-quota-rank"; import { parseRetryAfterMs } from "../combos/failover"; import { sweepExpiredOnWrite } from "../lib/state-store-sweeper"; import type { OcxConfig, OcxProviderConfig } from "../types"; @@ -167,7 +168,11 @@ export function rotateGenericOAuthAccountOn429( if (!set || set.accounts.length < 2) return null; const parsed = parseRetryAfterMs(retryAfterHeader, now); - const cooldownMs = Math.min(parsed ?? DEFAULT_COOLDOWN_MS, MAX_COOLDOWN_MS); + // An account whose allowance is provably spent gets a reset-aligned cooldown instead of + // the default minute: retrying it every 60s until the window rolls over is pure waste. + // A Retry-After from upstream still wins — it is the server's own instruction. + const exhausted = parsed === null ? exhaustedCooldownMs(providerName, failedAccountId, now) : null; + const cooldownMs = exhausted ?? Math.min(parsed ?? DEFAULT_COOLDOWN_MS, MAX_COOLDOWN_MS); health.set(healthKey(providerName, failedAccountId), { cooldownUntil: now + cooldownMs, cooldownSource: parsed ? "retry-after" : "default", @@ -180,14 +185,16 @@ export function rotateGenericOAuthAccountOn429( // from a count read before the failure. presence.delete(providerName); // Deterministic: start after the failed account so repeated 429s walk the roster instead of - // hammering whichever id happens to sort first. + // hammering whichever id happens to sort first. The ring is built BEFORE ranking — ranking + // the store's own order would change which account a quota-less provider rotates to. const order = set.accounts.map(account => account.id); const start = order.indexOf(failedAccountId); - for (let i = 1; i <= order.length; i++) { - const candidate = order[(start + i) % order.length]!; - if (candidate !== failedAccountId && eligible.includes(candidate)) return candidate; - } - return null; + const ring = start >= 0 ? [...order.slice(start + 1), ...order.slice(0, start)] : order; + const candidates = ring.filter(id => id !== failedAccountId && eligible.includes(id)); + if (candidates.length === 0) return null; + // With no quota evidence this returns the ring untouched, so providers without + // per-account quota keep exactly the traversal they have today. + return rankAccountsByHeadroom(providerName, candidates)[0] ?? null; } /** diff --git a/tests/generic-oauth-failover.test.ts b/tests/generic-oauth-failover.test.ts index ffa971cba5..a9401efd5e 100644 --- a/tests/generic-oauth-failover.test.ts +++ b/tests/generic-oauth-failover.test.ts @@ -83,6 +83,26 @@ describe("#2568 generic OAuth account failover", () => { expect(rotateGenericOAuthAccountOn429(config(true), "xai", ids[0]!, null)).toBe(ids[1]); }); + test("rotation continues AFTER the failed account, not from the top of the roster", async () => { + // Quota ranking now orders the candidates, so this pins the property the ranking must + // not disturb: with three accounts and no quota evidence anywhere, a 429 on the middle + // account moves to the one after it. Ranking the store's own order would answer the + // first account instead, silently changing every quota-less provider's traversal. + const ids = await seed(3); + expect(rotateGenericOAuthAccountOn429(config(), "xai", ids[1]!, null)).toBe(ids[2]); + }); + + test("the ring wraps when the failed account is last", async () => { + const ids = await seed(3); + expect(rotateGenericOAuthAccountOn429(config(), "xai", ids[2]!, null)).toBe(ids[0]); + }); + + test("an unknown failed account still yields a candidate", async () => { + // The account may have been removed between dispatch and the 429 landing. + const ids = await seed(2); + expect(rotateGenericOAuthAccountOn429(config(), "xai", "not-a-real-account", null)).toBe(ids[0]); + }); + test("a per-provider override beats the global switch", async () => { // Provider terms differ, so an operator may accept rotation on one provider and refuse it on // another. The narrower setting is the one that means something. diff --git a/tests/kiro-pool-rank.test.ts b/tests/kiro-pool-rank.test.ts new file mode 100644 index 0000000000..4cef86a93f --- /dev/null +++ b/tests/kiro-pool-rank.test.ts @@ -0,0 +1,117 @@ +import { afterEach, describe, expect, test } from "bun:test"; +import { exhaustedCooldownMs, rankAccountsByHeadroom } from "../src/oauth/account-quota-rank"; +import { + clearAccountQuotaCache, + setCachedProviderAccountQuotaForTests, +} from "../src/providers/quota"; +import { + clearKiroAccountUsageState, + commitKiroAccountUsageState, +} from "../src/providers/kiro-usage"; + +const realFetch = globalThis.fetch; + +afterEach(() => { + globalThis.fetch = realFetch; + clearAccountQuotaCache(); + clearKiroAccountUsageState(); +}); + +function seedPercent(provider: string, accountId: string, monthlyPercent: number): void { + setCachedProviderAccountQuotaForTests(provider, accountId, { monthlyPercent, updatedAt: Date.now() }); +} + +function seedExhausted(accountId: string, nextResetAt?: number): void { + commitKiroAccountUsageState(`kiro\u0000${accountId}`, { + quota: { monthlyPercent: 100, updatedAt: Date.now() }, + exhausted: true, + ...(nextResetAt !== undefined ? { nextResetAt } : {}), + }); +} + +describe("headroom ranking", () => { + test("the account with more remaining allowance goes first", () => { + seedPercent("kiro", "a", 10); + seedPercent("kiro", "b", 90); + expect(rankAccountsByHeadroom("kiro", ["b", "a"])).toEqual(["a", "b"]); + }); + + test("a measured-healthy account outranks an unknown one even when heavily used", () => { + // 95% used is still healthy: only a provider exhaustion verdict demotes an account. + seedPercent("kiro", "b", 95); + expect(rankAccountsByHeadroom("kiro", ["a", "b"])).toEqual(["b", "a"]); + }); + + test("an unknown account outranks one known to be exhausted", () => { + seedExhausted("b"); + expect(rankAccountsByHeadroom("kiro", ["b", "a"])).toEqual(["a", "b"]); + }); + + test("an exhausted account sorts last even with a low percentage on record", () => { + seedPercent("kiro", "a", 80); + seedPercent("kiro", "b", 5); + seedExhausted("b"); + expect(rankAccountsByHeadroom("kiro", ["b", "a"])).toEqual(["a", "b"]); + }); + + test("with no quota evidence the ring order is returned untouched", () => { + expect(rankAccountsByHeadroom("xai", ["c", "a", "b"])).toEqual(["c", "a", "b"]); + }); + + test("equal headroom preserves ring order", () => { + seedPercent("kiro", "a", 40); + seedPercent("kiro", "b", 40); + expect(rankAccountsByHeadroom("kiro", ["b", "a"])).toEqual(["b", "a"]); + }); + + test("the tightest window decides, not the roomiest", () => { + setCachedProviderAccountQuotaForTests("anthropic", "a", { + fiveHourPercent: 95, monthlyPercent: 5, updatedAt: Date.now(), + }); + setCachedProviderAccountQuotaForTests("anthropic", "b", { + fiveHourPercent: 30, monthlyPercent: 30, updatedAt: Date.now(), + }); + expect(rankAccountsByHeadroom("anthropic", ["a", "b"])).toEqual(["b", "a"]); + }); + + test("ranking never reaches the network", () => { + let called = false; + globalThis.fetch = (async () => { called = true; return new Response("{}"); }) as typeof fetch; + seedPercent("kiro", "a", 10); + rankAccountsByHeadroom("kiro", ["a", "b"]); + expect(called).toBe(false); + }); + + test("a single candidate is returned as-is", () => { + expect(rankAccountsByHeadroom("kiro", ["only"])).toEqual(["only"]); + }); +}); + +describe("exhaustion cooldown", () => { + test("a distant reset is clamped to a day", () => { + const now = Date.now(); + seedExhausted("a", now + 3 * 24 * 60 * 60_000); + expect(exhaustedCooldownMs("kiro", "a", now)).toBe(24 * 60 * 60_000); + }); + + test("an imminent reset is floored at five minutes", () => { + const now = Date.now(); + seedExhausted("a", now + 30_000); + expect(exhaustedCooldownMs("kiro", "a", now)).toBe(5 * 60_000); + }); + + test("a reset inside the window is honoured exactly", () => { + const now = Date.now(); + seedExhausted("a", now + 60 * 60_000); + expect(exhaustedCooldownMs("kiro", "a", now)).toBe(60 * 60_000); + }); + + test("a healthy account has no exhaustion cooldown", () => { + seedPercent("kiro", "a", 10); + expect(exhaustedCooldownMs("kiro", "a")).toBeNull(); + }); + + test("providers without an exhaustion verdict are unaffected", () => { + expect(exhaustedCooldownMs("xai", "a")).toBeNull(); + }); +}); From 829767c0f1727b4a21fcdd3e224d3cc0126360ff Mon Sep 17 00:00:00 2001 From: bitkyc08-arch Date: Sat, 29 Aug 2026 12:44:10 +0900 Subject: [PATCH 5/6] feat(cli): show Kiro's monthly allowance and describe the pool honestly quotaText read only the 5h and weekly windows, so a healthy Kiro account printed the same dash as an unprobed one. The CLI also still called Kiro a single login slot, which has been untrue since the multiauth add-account handoff shipped. --- .../src/content/docs/reference/adapters.md | 12 ++++++++++ .../docs/reference/cli/providers-accounts.md | 22 +++++++++++++++---- src/cli/account.ts | 3 +++ tests/cli-headless-parity.test.ts | 13 +++++++++++ 4 files changed, 46 insertions(+), 4 deletions(-) diff --git a/docs-site/src/content/docs/reference/adapters.md b/docs-site/src/content/docs/reference/adapters.md index ed223661c1..8919cd428b 100644 --- a/docs-site/src/content/docs/reference/adapters.md +++ b/docs-site/src/content/docs/reference/adapters.md @@ -179,6 +179,18 @@ of the HTTP retry loop. single post-cooldown probe prevent concurrent requests from exhausting independent retry budgets; hard quota failures and ordinary service errors are not replayed. - Its non-streaming parser drains the same event stream for the web-search loop. +- Reports per-account usage. `AmazonCodeWhispererService.GetUsageLimits` on + `https://management.{region}.kiro.dev/` returns the plan allowance, which becomes the + monthly quota window for that account; a free-trial balance is reported as its own window. + The region comes from the account's profile ARN, then its stored API/SSO region. An + unreadable or unrecognised response is reported as unknown rather than as zero usage, and + an account whose overage is enabled is not treated as exhausted merely for passing its + limit. The operation is undocumented by AWS, so treat the numbers as best-effort. +- Participates in multi-account rotation. Two or more logged-in Kiro accounts enable + automatic failover on a 429, and rotation prefers the account with the most known + headroom; an account whose allowance is provably spent is cooled until its window resets + (bounded between five minutes and a day) instead of being retried every minute. Each + rotated bearer carries its own profile ARN and region. ### Completion semantics diff --git a/docs-site/src/content/docs/reference/cli/providers-accounts.md b/docs-site/src/content/docs/reference/cli/providers-accounts.md index a5d48eaf33..898befa943 100644 --- a/docs-site/src/content/docs/reference/cli/providers-accounts.md +++ b/docs-site/src/content/docs/reference/cli/providers-accounts.md @@ -155,11 +155,11 @@ returns: ``` `--quota` adds a `QUOTA` column with each account's own usage, for providers that support a -per-account probe (Anthropic today). It is opt-in because the proxy probes the upstream once per -stored credential; the default listing stays a local read. `--refresh` bypasses the cached -result. An account with no per-account quota shows `-`, and one whose probe failed shows +per-account probe (Anthropic and Kiro today). It is opt-in because the proxy probes the upstream +once per stored credential; the default listing stays a local read. `--refresh` bypasses the +cached result. An account with no per-account quota shows `-`, and one whose probe failed shows `unavailable` — blank would read as "no usage" rather than "not measured". `--json` carries the -full breakdown per account, not just the two summarized windows: +full breakdown per account, not just the summarized windows: ```text $ ocx account list anthropic --quota @@ -168,6 +168,20 @@ anthropic oauth 1278f8da a***r@examp***.com - 5h 7% wk 62% anthropic oauth e112f28b k***1@examp***.net - active 5h 9% wk 45% ``` +Kiro bills a monthly allowance and reports no shorter window, so its accounts render a `mo` +figure instead: + +```text +$ ocx account list kiro --quota +PROVIDER TYPE ID PLAN/LABEL PRIORITY STATUS QUOTA +kiro oauth 3f0a91c2 a***r@examp***.com - active mo 15% +kiro oauth 8b24de70 k***1@examp***.net - mo 88% +``` + +With two or more Kiro accounts logged in, a 429 rotates to another account automatically and +prefers the one with the most remaining allowance. Accounts are added one at a time — +`ocx account login kiro` hands off to the Kiro CLI and appends the new account to the pool. + ### `ocx account current [--json]` Shows the active account or key. A Codex pool with no manual pin reports the priority-aware diff --git a/src/cli/account.ts b/src/cli/account.ts index 9974943136..e6269c6940 100644 --- a/src/cli/account.ts +++ b/src/cli/account.ts @@ -111,6 +111,9 @@ function quotaText(row: AccountRow): string { const short = quota.fiveHourPercent ?? quota.shortPercent; if (typeof short === "number") parts.push(`5h ${short}%`); if (typeof quota.weeklyPercent === "number") parts.push(`wk ${quota.weeklyPercent}%`); + // Kiro bills a monthly allowance and reports no shorter window, so without this arm a + // perfectly healthy Kiro account prints "-" and reads as broken. + if (typeof quota.monthlyPercent === "number") parts.push(`mo ${Math.round(quota.monthlyPercent)}%`); return parts.length > 0 ? parts.join(" ") : "-"; } diff --git a/tests/cli-headless-parity.test.ts b/tests/cli-headless-parity.test.ts index b8faa77ea2..019aca08fc 100644 --- a/tests/cli-headless-parity.test.ts +++ b/tests/cli-headless-parity.test.ts @@ -772,6 +772,19 @@ describe("#2566 per-account quota in ocx account list", () => { expect(formatAccountTable([row({ provider: "xai" })] as never, true)).toContain("-"); }); + test("a Kiro account's monthly allowance renders instead of a bare dash", () => { + // Kiro bills a monthly window and reports no shorter one. Without the monthly arm a + // healthy account rendered "-", which is the same output as "never probed". + expect(formatAccountTable([row({ provider: "kiro", quota: { monthlyPercent: 15 } })] as never, true)) + .toContain("mo 15%"); + }); + + test("a fractional monthly percentage is rounded for the column", () => { + // The column is a glance surface; the exact figure stays in --json. + expect(formatAccountTable([row({ provider: "kiro", quota: { monthlyPercent: 14.782 } })] as never, true)) + .toContain("mo 15%"); + }); + test("an account whose probe failed says so instead of reading as empty", () => { expect(formatAccountTable([row({ quotaUnavailable: true })] as never, true)).toContain("unavailable"); }); From 97c73350ca88eb302956b880c39eac9c021b1d3e Mon Sep 17 00:00:00 2001 From: bitkyc08-arch Date: Sat, 29 Aug 2026 12:54:34 +0900 Subject: [PATCH 6/6] docs(devlog): record the head-to-head result against kiro-lb --- .../080_head_to_head_result.md | 53 +++++++++++++++++++ 1 file changed, 53 insertions(+) create mode 100644 devlog/_plan/260829_kiro_quota_pool/080_head_to_head_result.md diff --git a/devlog/_plan/260829_kiro_quota_pool/080_head_to_head_result.md b/devlog/_plan/260829_kiro_quota_pool/080_head_to_head_result.md new file mode 100644 index 0000000000..06f555b254 --- /dev/null +++ b/devlog/_plan/260829_kiro_quota_pool/080_head_to_head_result.md @@ -0,0 +1,53 @@ +# 080 — Head to head with kiro-lb, from the landed tree + +Written after implementation, against the branch `codex/kiro-quota-pool` at `829767c0f`. +Every opencodex claim cites our code; every kiro-lb claim cites its file:line in the +AGPL-3.0 reference clone (commit `474df2b`). Behaviour was studied; no code was copied. + +## Where we now lead + +| Axis | kiro-lb | opencodex | +| --- | --- | --- | +| Breakdown selection | `AGENTIC_REQUEST`, else **index 0** (`kiro/usage.py:74-78`) | explicit priority list; an unrecognised list reports unknown (`src/providers/kiro-usage.ts` `RESOURCE_PRIORITY`) | +| Free-trial pool | fetched then dropped (`kiro/usage.py:97-111`) | reported as its own window | +| Exhaustion vs. percentage | `quota_depleted` derived from headroom <= 0 (`kiro/account_manager.py:173-190`) | overage-aware: a limit-passing account with overage enabled is **not** exhausted | +| Region handling | derived, unvalidated (`kiro/usage.py:20-36`) | every hostname candidate passes an allowlist; a crafted ARN cannot reach the URL | +| Idle cost | background poll, and `interval=0` still polls every 60s (`main.py:342-359`) | pull-on-demand behind a TTL; an idle proxy makes **zero** usage calls | +| Probe isolation | sequential, new 20s client per account (`kiro/dashboard.py:806-821`) | parallel with per-account failure isolation and an in-flight join | +| Refresh-pass robustness | a concurrent deletion can `KeyError` and abort the pass (`kiro/dashboard.py:806-815`) | generation guard + reconcile; a superseded probe commits nothing | +| Exhaustion quarantine | 6h floor, **32-day** cap (`kiro/config.py:457-479`) | reset-aligned, clamped 5 min – 24 h | +| Stale verdicts | persisted until overwritten | degrade to unknown past the TTL or the reset instant — "try again", never "stay parked" | +| Identity safety | per-account auth manager | bearer + profile ARN + region from ONE account snapshot, with a regression test (#2841 lineage) | +| Scope | Kiro only | Kiro is one provider on a shared seam that already served Anthropic | + +## Where kiro-lb still leads + +Stating this plainly, because a comparison that only lists our wins is worthless. + +1. **Pre-request selection.** kiro-lb picks an account *before* dispatch with a weighted + race (`kiro/account_manager.py:1183-1208`). Ours ranks only on the 429 recovery path + (`src/oauth/generic-account-failover.ts`), so the first request of a turn can still + land on a spent account. Deferred by design (doc `060`), not solved. +2. **Persistence across restart.** Its quota rows live in SQLite and seed routing at + startup (`kiro/store.py:206-289`). Our caches are process-local, so a restart forgets + every measurement until the next probe. +3. **Operations dashboard.** Request-rate charts, per-model token panels, Prometheus + export (`kiro/metrics.py`). We render quota bars and a CLI column. +4. **Account onboarding.** Device login for Builder ID, Google and GitHub straight from + the dashboard (`kiro/device_login.py`). Ours hands off to the Kiro CLI one account at + a time. + +Item 1 is the one that matters most for the user's "pool 기반 자동 탑재" ask, and it is +the first follow-up work-phase. + +## Honest summary + +On *correctness of the quota reading* and *safety of the pool machinery* we are ahead: +resource selection, overage semantics, trial balances, region validation, credential/route +pairing, and stale-state handling are each demonstrably stricter, with tests. On *routing +sophistication* kiro-lb is still ahead on the pre-dispatch axis, and on *operational +surface* it is ahead outright. + +"Better than kiro-lb" is therefore true for the two things this unit set out to do — +display Kiro quota, and make the pool quota-aware — and not yet true as a blanket claim +about the whole gateway.