diff --git a/devlog/_fin/260906_manual_account_selection/000_plan.md b/devlog/_fin/260906_manual_account_selection/000_plan.md new file mode 100644 index 0000000000..3ee8ea01cd --- /dev/null +++ b/devlog/_fin/260906_manual_account_selection/000_plan.md @@ -0,0 +1,33 @@ +# Manual account selection must control dispatch + +Loop: single-cycle satisfy-spec, C4 (credential/account allocation). Trigger: the user selected a healthy OAuth account in the dashboard, but every request was silently assigned to another account. Goal: disabled pools do not proactively reassign healthy requests; explicit selection wins; permitted automatic assignment is reflected in active-account state and dashboard. One cohesive PR targeting dev, pushed with `--no-verify` and merged as explicitly authorized. No stacks. + +Boundaries: existing account/key owners, request credential pairing, dashboard account synchronization, regression coverage, matching public docs. No provider API changes, real-account configuration changes, paid inference probes, releases, service restarts, or unrelated refactors. Use existing dependencies and temp credentials. No token, cost, or wall-clock budget was set; no paid oracle is required. Completion requires fresh direct verification and actual remote merge; a pending PR is not DONE. Escalate only a genuine tool/access block or a necessary authority not already granted. Main reclaims a lane after two distinct worker failures; new worker scope requires a P amendment. + +Memory: this numbered unit, `.tmp/manual-account-selection/` evidence, session-bound goalplan. The implementation and test map is in 010. One PABCD work-phase covers this single contract across its existing owners; frontend/runtime lanes are subtasks rather than separate deliverables. + +## Evidence and rival hypotheses + +GUI `useProviderAccountPools.ts:247` sends selected account to PUT `/api/oauth/accounts/active`. `oauth-account-routes.ts:322` calls `setActiveAccount`, which saves the selected id. `generic-account-failover.ts:185-196` defaults healthy proactive selection on when two accounts exist. `preferredInitialAccount` ranks by quota and can replace the selected healthy account. `responses/core.ts` resolves that other credential and logs it without updating stored selection. Provider quota reads stored active selection separately. Exact user-provided request IDs matched another account on attempt1, sendCount1, no retry. This is not an upstream429 recovery or a failed GUI save. + +H1 failed persistence is disproved by stored selected id and successful route semantics. H2 only a2s roster cache is insufficient: quota ranking would choose the other account after the cache expires. H3 manual authority absent from the selector is supported by the complete route-to-store-to-request chain. We will prove H2/H3 with isolated real owner functions before implementation. + +## Contract + +- Presence of multiple stored accounts does not enable proactive healthy-request allocation. Absent generic OAuth proactive enablement is off; explicit false wins for proactive allocation. REACTIVE429 recovery remains mandatory whenever another usable account exists, regardless of the pool switch, per the latest explicit user correction. +- With a pool enabled, the healthy active/manual account stays eligible and preferred. A quota percentage alone below exhaustion must not replace it. On a real account-scoped refusal or known exhaustion, enabled rotation may choose another usable account. +- Committing an automatic account selection must be conditional on the selection generation that produced the request. A newer manual selection (including A→B→A) wins. +- The selected account and its access token/project/origin metadata travel together; unsupported/unknown quota is not exhaustion. +- Codex, Anthropic, and API keys keep their own established contracts, but any contradiction with these user requirements is repaired in the same PR. Their specific findings must be folded into010 before their edits. + +Enforcement: runtime selectors plus guarded persisted active-account transition; execution surface covers first dispatch and every reactive replay. Known bypass: external callers can intentionally route exact account-targeting selectors, which remain their own explicit contract. Residual: concurrent requests can already be in flight on different credentials; dashboard reports the latest committed allocation, never retroactively cancels an already sent request. Wording: no claim that a UI highlight can reassign a request already upstream. + +Verification: focused Bun store/management/selection/retry tests first; full `bun run typecheck` and `bun run test` before PR review-ready; relevant GUI tests/lint/build and a rendered state transition if frontend changes. Regression tests use synthetic identities only. Public SoT: `structure/05_gui-and-management-api.md` and existing account/pool guide pages. Final record includes rejected hypotheses, unmodified owners with proof, security/concurrency audit, and remote merge. + +Latest steering: the user reports Codex works correctly. Preserve Codex routing/controller semantics and use its existing manual-selection behavior as the reference; include regression-only checks for Codex. Concentrate changes on the other quota-aware account paths where a concrete mismatch is established. + +Latest explicit design instruction: GUI selection and pool selection must share one selection owner, as Codex does. Both must commit the same authoritative selection BEFORE dispatch; requests use the committed account. Do not bolt on a separate UI-only mirror or route around manual selection while pretending the old active account remains selected. A concurrent newer user selection wins. This strengthens the existing planned store-owned selection transaction and applies with pool on or off. + +Latest correction: 429 automatic account switching is ALWAYS allowed, including poolOFF. Withdraw the planned reactive disable gate. Manual selection wins ordinary dispatch; a real429 may replace it using the same guarded selection owner, and the GUI follows that committed replacement. + +Latest explicit verification restriction: do not run repository-wide tests. The earlier full-suite requirement is superseded. One attempted full run was interrupted by user at exit130; it is not completion proof. Resolve observed failures with their specific test files, run focused affected checks and typecheck, then push --no-verify and merge the single PR. diff --git a/devlog/_fin/260906_manual_account_selection/010_implementation.md b/devlog/_fin/260906_manual_account_selection/010_implementation.md new file mode 100644 index 0000000000..85bf466a3c --- /dev/null +++ b/devlog/_fin/260906_manual_account_selection/010_implementation.md @@ -0,0 +1,71 @@ +# Implementation — one account-selection contract + +Depends on000. One PABCD work-phase, one PR. Existing subsystem owners stay intact. + +## Shared OAuth store and generic allocation (main) + +MODIFY `src/oauth/types.ts`, `src/oauth/store.ts`: add optional non-secret `ProviderAccountSet.selectionRevision` and a typed selection snapshot `{ accountId, revision? }`. Legacy files without the field remain valid. Normalize/persist/copy it through the existing auth-store boundary; every active-selection change (including re-selecting the same account) advances it. Add a store-owned capture function and conditional active-selection commit that compares both original active id and revision inside `mutateStore` before writing. Credential-only refresh must preserve the selection revision. Consumers: generic and Anthropic request admission/promotion. Management DTO need not expose the revision: existing active id remains the public selection. No credential value is logged. + +MODIFY `src/oauth/generic-account-failover.ts`: proactive allocation requires effective pool `enabled === true` (provider override then global). Merely storing2 accounts does not enable healthy-request steering. Preserve presence-based REACTIVE429 switching even when disabled, as the user expressly requires. Keep a healthy manually active account first; use quota ranking only when the chosen account is ineligible/exhausted or when an enabled pool recovers a real refusal. Unknown quota never implies exhaustion. Preserve existing cooldowns and bounded attempts. Clear roster on manual selection so an old2s cache cannot dispatch the previously selected account. + +MODIFY `src/server/management/oauth-account-routes.ts`: manual selection retains successful persistence and cache invalidation, and invalidates relevant generic selection state. Existing Anthropic manual handler remains its owner. Update outdated pool-settings contract comments/DTO docs in `src/oauth/pool-settings-capability.ts`, `src/types/provider.ts`, `src/types/config.ts` to match effective enablement; do not introduce a new UI toggle simply to repair a default. + +MODIFY `src/server/responses/core.ts`: capture the OAuth selection snapshot before awaited token materialization; preserve token/project/origin pairing. For actual automatic initial/retry selection, await the guarded active-account commit before publishing that allocation. A rejected promotion caused by a newer user selection must not overwrite it; continue via current valid selected account or return the original request failure as appropriate. Apply consistently to direct upstream errors, runTurn on429 callback, passthrough/combo retries and downstream stream recovery call sites. The same shared core serves Responses/Chat/Messages surfaces. + +## Anthropic and API keys (backend lane) + +MODIFY `src/oauth/anthropic-routing.ts` and its existing tests: preserve reactive429 rotation even when the pool is off (latest user correction). Manual selection must seed a preference that wins the next eligible dispatch, including quota strategy, clearing stale affinities. Guard automatic active-account promotion with the same store selection snapshot; main owns core call-site integration. Maintain account-scoped refresh restrictions. + +MODIFY `src/providers/key-failover.ts` and relevant caller/rotation types only if the isolated repro confirms env/keychain reference identity mismatch: match failed attempts by stable pool-entry identity/reference, never by comparing a resolved secret with a stored reference. Preserve newer manual key selection and committed config-to-dashboard mapping. API-key pools have no separate enable boolean: an explicitly configured multi-key pool remains their existing enable contract. No speculative new mode is added. Codex is unchanged and regression-only, per latest user steering. + +## Dashboard (frontend lane) + +MODIFY existing `gui/src/hooks/useProviderAccountPools.ts` / `gui/src/pages/Providers.tsx` runtime-read integration and existing tests as needed: periodically reconcile cheap local OAuth/key roster active state through the shared scheduler, without forcing upstream quota probes on every tick. Reuse quota rows and reject stale read results around manual mutation. Do not alter Codex controller behavior. Keep layout and existing localized labels. Render a synthetic selected-account transition and retain a screenshot in the unit for PR evidence. + +## Proof / reachable cases + +- Pool absent/false with2 accounts: manual A30%, B11%; ordinary dispatch staysA. Simulated429 MUST auto-switch to another usable account even with pool off, and persist that selection. Pool enabled: manual A remains preferred; known exhaustedA or actual refusal chooses usableB and active DTO becomesB. +- A delayed token refresh/429 starts onA; manual choiceB or A→B→A occurs before completion; older selection revision cannot change active state. Removal/reauth during candidate resolution cannot promote an invalid target. +- Generic matrix runs xAI, Cursor, Kimi, Copilot, Antigravity, Nous and representative passive-quota provider using synthetic snapshots. Copilot regional origin and Antigravity project remain paired to the selected bearer. +- Anthropic off/on, quota/RR manual priority, stale affinity, failed token resolution, and promotion races; Codex direct/manual/pin existing regressions stay green. +- Literal/env/keychain-supported API-key identities: rejected attempted key rotates to another distinct key, newer manual key wins, chosen key is the persisted active key. +- Dashboard backendA→B read updates highlighted selection and current quota association; an older quota/roster poll cannot revert a newer manual choice; a roster-only poll never initiates a paid/upstream quota read. + +Reuse existing tests: `tests/oauth/generic-oauth-failover.test.ts`, `tests/oauth/oauth-store-multi.test.ts`, `tests/oauth/adapter-event-oauth-failover.test.ts`, `tests/server/account-pool-management-api.test.ts`, provider quota/Anthropic/key failover suites discovered by owner search, and existing `gui/tests/provider-account-quota-loading.test.tsx` / provider revalidation tests. Prefer these files to new layout entries. Verify focused red before repair, then green; run `bun run typecheck`, `bun run test`, `bun run privacy:scan`, and relevant GUI tests/lint/build. New failures outside scope are diagnosed and recorded, not ignored. Fresh independent security/concurrency review before delivery. + +SoT sync: `structure/05_gui-and-management-api.md`, existing English configuration/account pool guide plus translated statements that would otherwise contradict changed enablement. Record provider coverage and limitations in011 evidence. Push single branch with `git push --no-verify`; create one PR using repository template againstdev, attach rendered UI evidence if GUI changed, verify remote head/CI, then merge as authorized and verify integration SHA. + +Known design risk for A audit: making selection persistence part of dispatch must not serialize all independent successful requests; commit only actual account changes, and use the existing guarded store writer. Manual selection while an upstream request is already running applies to subsequent allocation; no retroactive cancellation claim. + +## P clarification from frontend owner + +Current scheduler is `useKeyedClientResource`/`client-resource.ts`, not `useRuntimeRead`. Exact frontend writes: `gui/src/hooks/useProviderAccountPools.ts`, `gui/src/pages/Providers.tsx`, and `gui/src/pages/use-providers-oauth.ts`; tests: existing `provider-account-quota-loading.test.tsx` and `provider-revalidation-policy.test.tsx`. Register one local-roster refresh through App's existing30s shared scheduler, key by server+sorted provider list, not active ids. Preserve initial quota enrichment; never add quota=1 to periodic reads. Invalidate per-provider read generation at manual PUT start; apply successful response active id; late login-status hydration may seed only a missing roster. Codex controller stays unchanged. Existing relevant GUI baseline52 tests passed in separate file processes; grouped globals can collide, so run each file separately. + +## Shared-selection requirement (latest steering) + +GUI PUT and pool choice both use the store-owned active-selection transaction. Make the common operation return/confirm the committed selection, and dispatch from its matching credential snapshot; a pool proposal is not authoritative until it commits. Do not mutate per-request bearer first and asynchronously update GUI afterward. The generation guard is a concurrency condition inside that same shared operation, not a competing selection state. Keep the public active-id DTO unchanged. Main and backend lane must align on this seam before writing callers. + +Authoritative429 exception: poolOFF suppresses only proactive steering. Every generic/Anthropic/key recovery test must preserve automatic429 failover. Any earlier statement blocking429 whileOFF is superseded. + +## Immediate synchronization amendment + +Latest user rejects waiting for a poll after automatic selection. The repository has no dashboard EventSource subscription to reuse. Add a narrow authenticated management SSE invalidation channel for committed account/key selection (`/api/accounts/events`) with bounded subscriber count, lightweight heartbeat, disconnect cleanup, and no credentials/account identifiers in events (provider plus kind/revision only). A dependency-leaf `src/lib/account-selection-events.ts` owns subscription/publication; it must not import server or Lab. Shared authoritative OAuth/key selection writers publish only after successful persistence. `src/server/management/oauth-account-routes.ts` serves the channel behind existing management auth; close it through existing optional shutdown hooks if necessary. Frontend lane adds a single lifecycle-owned EventSource for this screen, invalidates cheap roster via current generation guards, reconnects with a full local refresh, and keeps30s scheduler as recovery only. Existing test files cover event arrival→highlight change without advancing poll clock, blocked/failed writes emitting no selection event, and subscription cleanup. New endpoint is authenticated and carries no authority to select; data-plane keys cannot subscribe. This replaces the earlier30s-only plan. + +## A synthesis — accepted bounded corrections + +Independent reviewer verdict: GO-WITH-FIXES(blockers=5). All five are folded into the implementation, none rebutted: +1. Revision lifecycle covers manual reselect, new activation, removal promotion, replacement/recreation; rollback replacement receives a new revision, never resurrects an old one. Credential-only writes preserve it. +2. Common store operation `commitOAuthAccountSelection(provider, accountId, {expectedSelection?, expectedCredentialGeneration?, requireUsableAccount?})` returns committed `{accountId,revision?}` ornull. GUI's existing `setActiveAccount` boolean API wraps this same operation. `captureOAuthAccountSelection` supplies the expected snapshot. Validate unchanged-account admission too; retry current selection after a failed CAS, never send the rejected candidate. Cover generic core sites4868/5753/6100/6834/7244 and initial selection. Failed CAS emits nothing. GUI invalidates reads at both PUT start and settle and preserves settled quota state. +3. Anthropic affinity/rotation success bookkeeping occurs only after selection commits. All four promotion callers await it; background local-CLI token restrictions remain checked before commit. +4. API-key attempt carries stable pool identity/reference plus selection generation across all callers including nativeChat; common manual/automatic selection commit guards ABA and notifies only after persistence. +5. Quota eligibility explicitly distinguishes known exhaustion from unknown, including Kiro overage rules. Parameterized provider coverage includes Kiro and passive providers. + +B lane allocation (approved plan): main owns core.ts, OAuth management route/SSE route registration, generic selector/rank and integration proof/docs; store lane owns oauth/types.ts+store.ts, leaf account-selection event bus, and oauth-store-multi.test.ts; backend lane owns Anthropic routing+tests and API-key source/router/transport+tests including types/provider.ts; frontend lane owns the3GUI sourcefiles and2testfiles above. No worker changes maincore or another lane's files. Independent context review follows integration. + +Latest explicit verification restriction: do not run repository-wide tests. The earlier full-suite requirement is superseded. One attempted full run was interrupted by user at exit130; it is not completion proof. Resolve observed failures with their specific test files, run focused affected checks and typecheck, then push --no-verify and merge the single PR. + +## C corrective review amendment + +Accepted independent review findings: dispatch must revalidate after pacing/build waits;401 replay must use the common selection owner; CCA project must always come from the admitted account; Anthropic initial manual choice must survive restart; selection SSE must stop on session revocation/expiry; hub relay must not apply its15s total deadline to an established selection stream; late initial quota data must survive manual selection without restoring old active flags. Main owns physical dispatch,401,CCA; backend lane owns Anthropic/API-key corrections; frontend lane owns reconnect/quota fixes. For bounded parallel C repair, the completed store worker is reassigned to SSE/management session liveness and hub-relay fixes only; no concurrent write ownership overlaps. All verification remains focused; full suite is prohibited. Draft PR3768 is open and CI runs asynchronously. + +C second-review correction: a rebuilt adapter must replace the active adapter/cache, and physical admission must be bound to the particular wire request's originating credential, not merely shared request state. Image/search model loops need the same request-specific executor. The runtime reviewer is reassigned as an exclusive repair worker for core/fetch-helpers and those two loops plus focused regression tests; main pauses edits there and independently verifies the returned delta. API-key helper/native Chat remains the backend lane; runtime worker integrates its exported helpers. Codex forward path remains unchanged. No full local tests. diff --git a/devlog/_fin/260906_manual_account_selection/011_verification.md b/devlog/_fin/260906_manual_account_selection/011_verification.md new file mode 100644 index 0000000000..ac944beeb8 --- /dev/null +++ b/devlog/_fin/260906_manual_account_selection/011_verification.md @@ -0,0 +1,50 @@ +# Verification and delivery record + +The fix uses a common committed selection for manual and automatic OAuth/API-key allocation. +A healthy manual selection has priority; reactive429 recovery remains enabled with poolOFF. +The physical request carries the binding of the adapter that built it. A stale binding is rebuilt, +and the new adapter and request cache remain authoritative for later retries and continuations. +Image/search loops share the request-specific executor. Codex routing/controller semantics are unchanged. + +## Focused evidence + +| Surface | Evidence | +| --- | --- | +| Generic OAuth | Parameterized xAI, Cursor, Kimi, Copilot, Antigravity, Nous, Kiro, Meta-Muse manual priority;36 focused checks passed | +| Store |13 failing regression cases before repair;36 store checks passed, including ABA/removal/recreation and refresh-only preservation | +| Actual dispatch | Copilot build/pacing races,413 follow-up, image/search pacing, and runTurn first-send coverage;31 checks passed with3 final boundary cases demonstrated RED→GREEN | +| API keys | Literal/env/keychain identity and newer manual selection; native Chat pacing revalidation; focused12+59 checks passed | +| Anthropic | Manual selection, guarded promotion, restart bootstrap, and always-on429; focused96-test group passed | +| Antigravity |20 OAuth401/project tests passed; a project-less account is refused before dispatch; every admitted request uses its account's project | +| Image/search | Image loops31, search61, timeout contract7 passed in separate processes | +| Management/relay |40 focused checks passed; authenticated invalidation stream, client cancellation, byte/subscriber bounds, expiration/revocation, and established SSE lifetime | +| Dashboard | Probe/passive quota hydration, stale selection guards and immediate event/reconnect behavior;38 roster+8 page checks passed | +| CI fixes | Upsert fixtures now persist like the real login flow and verify disk; GUI source binding check updated; React Doctor0.9.11 changed-file scan has0 errors/0 warnings | +| Static/privacy | Typecheck, privacy scan and diff check passed at integration checkpoints | + +Counts identify each recorded check group; they overlap and must not be added into a unique-test total. +The user prohibited repository-wide local tests. An earlier full run was interrupted with exit130; +it is not completion evidence and was not repeated. CI runs asynchronously on PR3768. +At dd5aec571, all23 applicable CI checks passed, with2 intentional skips. + +## Current browser proof + +Aside opened a local synthetic fixture rendering the real Providers component and styles at1440×900. +The current GUI moved ChoiceA→ChoiceB from a selection event without advancing the poll clock; +upstream quota-read count stayed2→2. Both screenshots were inspected by main; no Korean clipping or +incorrect active indicator was observed. The fixture and owned browser tabs were stopped afterward. +The screenshots contain only synthetic account names and masked IDs. + +- [Before](evidence/011_selection-before.png) +- [After](evidence/012_selection-after.png) + +Independent C reviews identified and drove repairs for cached adapter reuse, sidecar dispatch, +credential refresh priority, account/project pairing, restart priority, stream lifetime and quota +hydration. All identified findings were implemented and the repaired cases were exercised. + +## Delivery scope + +One PR: https://github.com/lidge-jun/opencodex/pull/3768 . Every push uses`git push --no-verify` as +explicitly requested. The maintainer explicitly authorized an administrator merge. Remote merge +state and its final SHA are verified separately from local implementation proof; no runtime service +restart or real-account configuration mutation is part of this change. diff --git a/devlog/_fin/260906_manual_account_selection/evidence/011_selection-before.png b/devlog/_fin/260906_manual_account_selection/evidence/011_selection-before.png new file mode 100644 index 0000000000..e48fdf946f Binary files /dev/null and b/devlog/_fin/260906_manual_account_selection/evidence/011_selection-before.png differ diff --git a/devlog/_fin/260906_manual_account_selection/evidence/012_selection-after.png b/devlog/_fin/260906_manual_account_selection/evidence/012_selection-after.png new file mode 100644 index 0000000000..1806056ab7 Binary files /dev/null and b/devlog/_fin/260906_manual_account_selection/evidence/012_selection-after.png differ diff --git a/docs-site/src/content/docs/guides/web-dashboard.md b/docs-site/src/content/docs/guides/web-dashboard.md index 3c46748a31..aef01e2fec 100644 --- a/docs-site/src/content/docs/guides/web-dashboard.md +++ b/docs-site/src/content/docs/guides/web-dashboard.md @@ -89,6 +89,15 @@ host and port over a LAN IP or an alias. | **Storage** | Read-only CODEX_HOME disk breakdown (sessions, archives, DBs, attachments). Optional archived cleanup: preview the oldest N%, then quarantine to `CODEX_HOME/.trash` (default) or permanently delete behind an explicit checkbox. **Auto-cleanup policy** is opt-in and **default OFF** (`storageCleanupPolicy.enabled`); configure threshold/target/schedule/mode on the Storage page, or trigger **Run now**. Quarantined entries can be restored from the Storage page (JSONL + threads). Active sessions stay read-only. Cleanup and restore are refused while Codex holds the newest/active `state_*.sqlite` locked. | | **Stop** | Gracefully stop the proxy and installed background service, restore native Codex, and exit (`POST /api/stop`). On Windows with the Task Scheduler backend the dashboard refuses and asks you to run `ocx stop` instead: that wrapper can respawn the proxy after the task ends, and only a stop running outside this process can verify the restart window before restoring your client config. Nothing is changed when it refuses. | +### Account selection + +Account selection is shared with request routing. Selecting an OAuth account takes effect on the +next request even when a pool is enabled. A healthy selection is not replaced merely because +another generic OAuth account has more unused quota. If the account returns 429, automatic +failover can still select another usable account with the pool off. A committed automatic +selection updates the dashboard immediately; account changes do not wait for the quota refresh +timer. Requests already sent upstream retain their original credentials. + ### Filtering request logs Logs filters combine surface, intercepted requests, provider, exact model, status, time, diff --git a/docs-site/src/content/docs/ko/guides/web-dashboard.md b/docs-site/src/content/docs/ko/guides/web-dashboard.md index dee6fc2e2a..efdd80179f 100644 --- a/docs-site/src/content/docs/ko/guides/web-dashboard.md +++ b/docs-site/src/content/docs/ko/guides/web-dashboard.md @@ -221,3 +221,11 @@ GUI는 프록시의 JSON 관리 API를 사용하는 얇은 클라이언트입니 프로바이더 설정에 복사됩니다. 별도 분류 작업 없이도 [비전 사이드카](/ko/guides/sidecars/)가 올바른 조건에서만 실행됩니다. ::: + +### 계정 선택과 자동 전환 + +GUI에서 OAuth 계정을 선택하면 풀 모드에서도 다음 요청에 반영돼요. 일반 OAuth 계정은 +정상적으로 사용할 수 있는 선택 계정을 유지하며, 다른 계정의 남은 할당량이 더 많다는 +이유만으로 바꾸지 않아요. 선택 계정이 429를 반환하면 풀이 꺼져 있어도 사용 가능한 다른 +계정으로 자동 전환해요. 자동 선택이 저장되면 GUI의 활성 표시도 즉시 바뀌어요. +이미 서버로 보낸 요청의 인증 정보는 바꾸지 않아요. diff --git a/gui/src/hooks/useProviderAccountPools.ts b/gui/src/hooks/useProviderAccountPools.ts index 255ad4dfbe..0197a547d5 100644 --- a/gui/src/hooks/useProviderAccountPools.ts +++ b/gui/src/hooks/useProviderAccountPools.ts @@ -24,6 +24,32 @@ export interface OAuthAccount extends AccountQuotaReading { healthAction?: string; } export interface ApiKeyEntry extends AccountQuotaReading { id: string; label?: string; masked: string; active: boolean } +export interface AccountSelectionTarget { provider: string; kind: "oauth" | "api-key" } + +function selectionRows(rows: T[], id: string | null | undefined): T[] { + return id === undefined ? rows : rows.map(row => ({ ...row, active: row.id === id })); +} + +/** An invalidation read changes membership/selection, not quota probe state. */ +function mergeRosterRows(rows: T[], previous: T[]): T[] { + const prior = new Map(previous.map(row => [row.id, row])); + return mergeQuotaRows(rows, previous, false).map(row => supportsQuotaRead(row) ? { + ...row, + quotaPending: prior.get(row.id)?.quotaPending ?? false, + quotaUnavailable: prior.get(row.id)?.quotaUnavailable ?? false, + } : row); +} + +/** A probe started before a newer roster may update quota only on surviving IDs. */ +function mergeLateQuotaRows(rows: T[], enriched: T[]): T[] { + const byId = new Map(enriched.map(row => [row.id, row])); + return rows.map(row => { + const incoming = byId.get(row.id); + if (!incoming || incoming.quotaMode !== row.quotaMode) return row; + const quota = mergeQuotaRows([incoming], [row], true)[0]; + return { ...row, quota: quota.quota, quotaPending: quota.quotaPending, quotaUnavailable: quota.quotaUnavailable }; + }); +} type QuotaRow = AccountQuotaReading & { id: string }; const supportsQuotaRead = (row: AccountQuotaReading) => row.quotaMode === "probe" || row.quotaMode === "passive"; @@ -47,8 +73,9 @@ function mergeQuotaRows(rows: T[], previous: T[], enriched: }); } -function unavailableQuotaRows(rows: T[]): T[] { - return rows.map(row => supportsQuotaRead(row) +function unavailableQuotaRows(rows: T[], attempted?: T[]): T[] { + const attemptedModes = attempted && new Map(attempted.map(row => [row.id, row.quotaMode])); + return rows.map(row => supportsQuotaRead(row) && (!attemptedModes || attemptedModes.get(row.id) === row.quotaMode) ? { ...row, quotaUnavailable: true, quotaPending: false } : row); } @@ -91,12 +118,18 @@ export function useProviderAccountPools(deps: { const [addingKeyFor, setAddingKeyFor] = useState(null); const [newKeyValue, setNewKeyValue] = useState(""); const accountRequestGenerationRef = useRef>({}); + const rosterGenerationRef = useRef>({}); + const quotaGenerationRef = useRef>({}); + const selectionMutationsRef = useRef(new Map()); const requestsRef = useRef(new Set()); const mountedRef = useRef(true); const serverRef = useRef(apiBase); useEffect(() => { const generations = accountRequestGenerationRef.current; const requests = requestsRef.current; + const rosterGenerations = rosterGenerationRef.current; + const quotaGenerations = quotaGenerationRef.current; + const mutations = selectionMutationsRef.current; mountedRef.current = true; const serverChanged = serverRef.current !== apiBase; serverRef.current = apiBase; @@ -109,6 +142,9 @@ export function useProviderAccountPools(deps: { return () => { mountedRef.current = false; for (const key of Object.keys(generations)) generations[key] += 1; + for (const key of Object.keys(rosterGenerations)) rosterGenerations[key] += 1; + for (const key of Object.keys(quotaGenerations)) quotaGenerations[key] += 1; + mutations.clear(); for (const controller of requests) controller.abort(); requests.clear(); }; @@ -119,8 +155,11 @@ export function useProviderAccountPools(deps: { const keyPoolsKeyRef = useRef(null); const switchingAccountRef = useRef<{ provider: string; accountId: string } | null>(null); - const readRoster = useCallback(async (url: string): Promise => { + const readRoster = useCallback(async (url: string, signal?: AbortSignal): Promise => { const bounded = createBoundedFetch(20_000); + const abort = () => bounded.controller.abort(); + if (signal?.aborted) abort(); + signal?.addEventListener("abort", abort, { once: true }); requestsRef.current.add(bounded.controller); try { const response = await fetch(url, { signal: bounded.signal }); @@ -130,6 +169,7 @@ export function useProviderAccountPools(deps: { return data; } finally { bounded.clear(); + signal?.removeEventListener("abort", abort); requestsRef.current.delete(bounded.controller); } }, []); @@ -146,7 +186,10 @@ export function useProviderAccountPools(deps: { const key = `oauth:${provider}`; const generation = (accountRequestGenerationRef.current[key] ?? 0) + 1; accountRequestGenerationRef.current[key] = generation; + const rosterGeneration = (rosterGenerationRef.current[key] ?? 0) + 1; + rosterGenerationRef.current[key] = rosterGeneration; const currentRequest = () => aliveRef.current && mountedRef.current && serverRef.current === apiBase && accountRequestGenerationRef.current[key] === generation; + const currentRoster = () => currentRequest() && rosterGenerationRef.current[key] === rosterGeneration; const url = `${apiBase}/api/oauth/accounts?provider=${encodeURIComponent(provider)}`; try { // Cheap local read first so account switch / reauth / remove controls appear @@ -154,32 +197,39 @@ export function useProviderAccountPools(deps: { const data = await readRoster<{ activeAccountId?: string | null; accounts?: OAuthAccount[] }>(url); if (!Array.isArray(data.accounts)) throw new Error("Invalid account roster"); if (!currentRequest()) return false; - const rows = data.accounts; - setAccountSets(current => currentRequest() ? { ...current, [provider]: { + const rows = selectionRows(data.accounts, data.activeAccountId); + setAccountSets(current => currentRoster() ? { ...current, [provider]: { activeAccountId: data.activeAccountId ?? null, accounts: mergeQuotaRows(rows, current[provider]?.accounts ?? [], false), } } : current); - setAccountLoadStates(current => currentRequest() ? { ...current, [provider]: "ready" } : current); + setAccountLoadStates(current => currentRoster() ? { ...current, [provider]: "ready" } : current); if (!rows.some(supportsQuotaRead)) return true; const enrich = async (): Promise => { + // Manual selection invalidates roster reads, not a per-ID quota probe already sent. + const quotaGeneration = (quotaGenerationRef.current[key] ?? 0) + 1; + quotaGenerationRef.current[key] = quotaGeneration; + const currentQuota = () => aliveRef.current && mountedRef.current && serverRef.current === apiBase + && quotaGenerationRef.current[key] === quotaGeneration; try { const quotaData = await readRoster<{ activeAccountId?: string | null; accounts?: OAuthAccount[] }>(`${url}"a=1${refresh ? "&refresh=1" : ""}`); if (!Array.isArray(quotaData.accounts)) throw new Error("Invalid account quota roster"); - if (!currentRequest()) return false; - const enriched = quotaData.accounts; - setAccountSets(current => currentRequest() ? { + if (!currentQuota()) return false; + const enriched = selectionRows(quotaData.accounts, quotaData.activeAccountId); + setAccountSets(current => !currentQuota() ? current : !currentRoster() ? { + ...current, [provider]: { ...current[provider], accounts: mergeLateQuotaRows(current[provider]?.accounts ?? [], enriched) }, + } : { ...current, [provider]: { - activeAccountId: quotaData.activeAccountId ?? data.activeAccountId ?? null, + activeAccountId: quotaData.activeAccountId === undefined ? data.activeAccountId ?? null : quotaData.activeAccountId, accounts: mergeQuotaRows(enriched, current[provider]?.accounts ?? [], true), }, - } : current); + }); return !enriched.some(row => row.quotaUnavailable === true); } catch { - if (!currentRequest()) return false; - setAccountSets(current => currentRequest() && current[provider] ? { - ...current, [provider]: { ...current[provider], accounts: unavailableQuotaRows(current[provider].accounts) }, + if (!currentQuota()) return false; + setAccountSets(current => currentQuota() && current[provider] ? { + ...current, [provider]: { ...current[provider], accounts: unavailableQuotaRows(current[provider].accounts, rows) }, } : current); return false; } @@ -188,9 +238,9 @@ export function useProviderAccountPools(deps: { void enrich(); return true; } catch { - if (!currentRequest()) return false; - setAccountLoadStates(current => currentRequest() ? { ...current, [provider]: "error" } : current); - setAccountSets(current => currentRequest() && current[provider] ? { + if (!currentRoster()) return false; + setAccountLoadStates(current => currentRoster() ? { ...current, [provider]: "error" } : current); + setAccountSets(current => currentRoster() && current[provider] ? { ...current, [provider]: { ...current[provider], accounts: unavailableQuotaRows(current[provider].accounts) }, } : current); return false; @@ -205,29 +255,42 @@ export function useProviderAccountPools(deps: { const key = `key:${name}`; const generation = (accountRequestGenerationRef.current[key] ?? 0) + 1; accountRequestGenerationRef.current[key] = generation; + const rosterGeneration = (rosterGenerationRef.current[key] ?? 0) + 1; + rosterGenerationRef.current[key] = rosterGeneration; const currentRequest = () => aliveRef.current && mountedRef.current && serverRef.current === apiBase && accountRequestGenerationRef.current[key] === generation; + const currentRoster = () => currentRequest() && rosterGenerationRef.current[key] === rosterGeneration; const url = `${apiBase}/api/providers/keys?name=${encodeURIComponent(name)}`; const failed = () => { - if (currentRequest()) setKeyPools(current => currentRequest() + if (currentRoster()) setKeyPools(current => currentRoster() ? { ...current, [name]: unavailableQuotaRows(current[name] ?? []) } : current); return false; }; try { - const data = await readRoster<{ keys?: ApiKeyEntry[] }>(url); + const data = await readRoster<{ activeId?: string | null; keys?: ApiKeyEntry[] }>(url); if (!Array.isArray(data.keys)) throw new Error("Invalid key roster"); if (!currentRequest()) return false; - const rows = data.keys; - setKeyPools(current => currentRequest() ? { ...current, [name]: mergeQuotaRows(rows, current[name] ?? [], false) } : current); + const rows = selectionRows(data.keys, data.activeId); + setKeyPools(current => currentRoster() ? { ...current, [name]: mergeQuotaRows(rows, current[name] ?? [], false) } : current); if (!rows.some(supportsQuotaRead)) return true; const enrich = async (): Promise => { + const quotaGeneration = (quotaGenerationRef.current[key] ?? 0) + 1; + quotaGenerationRef.current[key] = quotaGeneration; + const currentQuota = () => aliveRef.current && mountedRef.current && serverRef.current === apiBase + && quotaGenerationRef.current[key] === quotaGeneration; try { - const data = await readRoster<{ keys?: ApiKeyEntry[] }>(`${url}"a=1${refresh ? "&refresh=1" : ""}`); + const data = await readRoster<{ activeId?: string | null; keys?: ApiKeyEntry[] }>(`${url}"a=1${refresh ? "&refresh=1" : ""}`); if (!Array.isArray(data.keys)) throw new Error("Invalid key quota roster"); - if (!currentRequest()) return false; - const enriched = data.keys; - setKeyPools(current => currentRequest() ? { ...current, [name]: mergeQuotaRows(enriched, current[name] ?? [], true) } : current); + if (!currentQuota()) return false; + const enriched = selectionRows(data.keys, data.activeId); + setKeyPools(current => currentQuota() ? { ...current, [name]: currentRoster() + ? mergeQuotaRows(enriched, current[name] ?? [], true) + : mergeLateQuotaRows(current[name] ?? [], enriched) } : current); return !enriched.some(row => row.quotaUnavailable === true); - } catch { return failed(); } + } catch { + if (currentQuota()) setKeyPools(current => currentQuota() + ? { ...current, [name]: unavailableQuotaRows(current[name] ?? [], rows) } : current); + return false; + } }; if (refresh) return await enrich(); void enrich(); @@ -237,22 +300,88 @@ export function useProviderAccountPools(deps: { return results.every(Boolean); }, [apiBase, aliveRef, readRoster]); + const refreshAccountRosters = useCallback(async (target?: AccountSelectionTarget, signal?: AbortSignal): Promise => { + if (!aliveRef.current || !mountedRef.current || serverRef.current !== apiBase || signal?.aborted) return false; + const targets: AccountSelectionTarget[] = target ? [target] : Object.entries(config?.providers ?? {}).flatMap(([provider, p]) => + p.authMode === "oauth" && provider !== "openai" ? [{ provider, kind: "oauth" as const }] + : p.hasApiKey && p.authMode !== "oauth" && p.authMode !== "forward" ? [{ provider, kind: "api-key" as const }] : []); + const results = await Promise.all(targets.map(async ({ provider, kind }) => { + const key = `${kind === "oauth" ? "oauth" : "key"}:${provider}`; + // PUT settlement always reconciles, including events received while it is pending. + if (selectionMutationsRef.current.has(key)) return false; + const generation = (rosterGenerationRef.current[key] ?? 0) + 1; + rosterGenerationRef.current[key] = generation; + const currentRequest = () => aliveRef.current && mountedRef.current && serverRef.current === apiBase + && !signal?.aborted && rosterGenerationRef.current[key] === generation && !selectionMutationsRef.current.has(key); + try { + if (kind === "oauth") { + const data = await readRoster<{ activeAccountId?: string | null; accounts?: OAuthAccount[] }>( + `${apiBase}/api/oauth/accounts?provider=${encodeURIComponent(provider)}`, signal); + if (!Array.isArray(data.accounts) || !currentRequest()) return false; + const rows = selectionRows(data.accounts, data.activeAccountId); + setAccountSets(current => currentRequest() ? { ...current, [provider]: { + activeAccountId: data.activeAccountId === undefined ? rows.find(row => row.active)?.id ?? null : data.activeAccountId, + accounts: mergeRosterRows(rows, current[provider]?.accounts ?? []), + } } : current); + setAccountLoadStates(current => currentRequest() ? { ...current, [provider]: "ready" } : current); + } else { + const data = await readRoster<{ activeId?: string | null; keys?: ApiKeyEntry[] }>( + `${apiBase}/api/providers/keys?name=${encodeURIComponent(provider)}`, signal); + if (!Array.isArray(data.keys) || !currentRequest()) return false; + const rows = selectionRows(data.keys, data.activeId); + setKeyPools(current => currentRequest() ? { ...current, [provider]: mergeRosterRows(rows, current[provider] ?? []) } : current); + } + return true; + } catch { + // A missed invalidation read does not change quota health; recovery retries it. + return false; + } + })); + return results.every(Boolean); + }, [aliveRef, apiBase, config, readRoster]); + + const invalidateSelectionReads = (provider: string, kind: AccountSelectionTarget["kind"]) => { + const key = `${kind === "oauth" ? "oauth" : "key"}:${provider}`; + accountRequestGenerationRef.current[key] = (accountRequestGenerationRef.current[key] ?? 0) + 1; + rosterGenerationRef.current[key] = (rosterGenerationRef.current[key] ?? 0) + 1; + // The independent quota generation still owns pending/error flags for surviving IDs. + return key; + }; + const switchAccount = async (provider: string, account: OAuthAccount) => { if (account.active || account.needsReauth || switchingAccountRef.current) return; const target = { provider, accountId: account.id }; switchingAccountRef.current = target; setSwitchingAccount(target); + const key = invalidateSelectionReads(provider, "oauth"); + const mutation = Symbol(); + selectionMutationsRef.current.set(key, mutation); + const currentMutation = () => aliveRef.current && mountedRef.current && serverRef.current === apiBase && selectionMutationsRef.current.get(key) === mutation; const label = oauthAccountDisplayLabel(accountSets[provider]?.accounts ?? [account], account, t); try { const res = await fetch(`${apiBase}/api/oauth/accounts/active`, { method: "PUT", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ provider, accountId: account.id }) }); + if (!currentMutation()) return; if (!res.ok) { notify(t("prov.accountSwitchFail"), false); return; } - const refreshed = await fetchAccountSets([provider]); + const result = await res.json().catch(() => ({})) as { activeAccountId?: string | null }; + if (!currentMutation()) return; + invalidateSelectionReads(provider, "oauth"); + const selected = result.activeAccountId === undefined ? account.id : result.activeAccountId; + setAccountSets(current => current[provider] ? { ...current, [provider]: { + ...current[provider], activeAccountId: selected, accounts: selectionRows(current[provider].accounts, selected), + } } : current); + selectionMutationsRef.current.delete(key); + const refreshed = await refreshAccountRosters({ provider, kind: "oauth" }); await Promise.all([fetchOauth(), fetchProviderQuotas(true)]); if (!refreshed) { notify(t("pws.accountsLoadFailed"), false); return; } notify(t("prov.accountSwitched", { email: label }), true); } catch { - notify(t("prov.accountSwitchFail"), false); + if (currentMutation()) notify(t("prov.accountSwitchFail"), false); } finally { + if (currentMutation()) { + invalidateSelectionReads(provider, "oauth"); + selectionMutationsRef.current.delete(key); + void refreshAccountRosters({ provider, kind: "oauth" }); + } if (switchingAccountRef.current?.provider === target.provider && switchingAccountRef.current.accountId === target.accountId) { switchingAccountRef.current = null; if (aliveRef.current) setSwitchingAccount(null); @@ -261,15 +390,34 @@ export function useProviderAccountPools(deps: { }; const switchApiKey = async (provider: string, entry: ApiKeyEntry) => { - if (entry.active) return; - const res = await fetch(`${apiBase}/api/providers/keys/active`, { method: "PUT", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ name: provider, id: entry.id }) }); - if (res.ok) { + if (entry.active || selectionMutationsRef.current.has(`key:${provider}`)) return; + const key = invalidateSelectionReads(provider, "api-key"); + const mutation = Symbol(); + selectionMutationsRef.current.set(key, mutation); + const currentMutation = () => aliveRef.current && mountedRef.current && serverRef.current === apiBase && selectionMutationsRef.current.get(key) === mutation; + try { + const res = await fetch(`${apiBase}/api/providers/keys/active`, { method: "PUT", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ name: provider, id: entry.id }) }); + if (!currentMutation()) return; + if (!res.ok) { + const failure = await res.json().catch(() => ({})) as { error?: string }; + if (currentMutation()) notify(failure.error || t("prov.keySwitchFail"), false); + return; + } + const data = await res.json().catch(() => ({})) as { activeId?: string | null }; + if (!currentMutation()) return; + invalidateSelectionReads(provider, "api-key"); + const selected = data.activeId === undefined ? entry.id : data.activeId; + setKeyPools(current => current[provider] ? { ...current, [provider]: selectionRows(current[provider], selected) } : current); notify(t("prov.keySwitched", { key: entry.label ?? entry.masked }), true); - void fetchKeyPools(Object.keys(keyPools)); void fetchProviderQuotas(true); - } else { - const data = await res.json().catch(() => ({})); - notify(data.error || t("prov.keySwitchFail"), false); + } catch { + if (currentMutation()) notify(t("prov.keySwitchFail"), false); + } finally { + if (currentMutation()) { + invalidateSelectionReads(provider, "api-key"); + selectionMutationsRef.current.delete(key); + void refreshAccountRosters({ provider, kind: "api-key" }); + } } }; @@ -382,7 +530,7 @@ export function useProviderAccountPools(deps: { return { accountSets, accountLoadStates, switchingAccount, openAccounts, keyPools, addingKeyFor, newKeyValue, setAccountSets, setAccountLoadStates, setSwitchingAccount, setOpenAccounts, setKeyPools, setAddingKeyFor, setNewKeyValue, - fetchAccountSets, fetchKeyPools, switchAccount, switchApiKey, removeApiKey, addApiKeyValue, addApiKey, editCredentialAlias, removeAccount, + fetchAccountSets, fetchKeyPools, refreshAccountRosters, switchAccount, switchApiKey, removeApiKey, addApiKeyValue, addApiKey, editCredentialAlias, removeAccount, oauthCardProviders, keyCardProviders, activeAccountNeedsReauth, }; } diff --git a/gui/src/pages/Providers.tsx b/gui/src/pages/Providers.tsx index 779bc99289..78b2096d95 100644 --- a/gui/src/pages/Providers.tsx +++ b/gui/src/pages/Providers.tsx @@ -1,5 +1,5 @@ import { usageSummary30dResourceKey } from "../usage-summary-resource"; -import { useCallback, useEffect, useMemo, useRef, useState } from "react"; +import { useCallback, useEffect, useLayoutEffect, useMemo, useRef, useState } from "react"; import ProviderWorkspaceShell, { type AddProviderIntent } from "../components/provider-workspace/ProviderWorkspaceShell"; import ProviderDetails from "../components/provider-workspace/ProviderDetails"; import { isAccountProvider, type WorkspaceProvider } from "../provider-workspace/catalog"; @@ -8,7 +8,7 @@ import { oauthTosRisk } from "../oauth-tos-risk"; import { ToastNotice, type NoticeTone } from "../ui"; import { IconPlus } from "../icons"; import { useT } from "../i18n/shared"; -import { useProviderAccountPools } from "../hooks/useProviderAccountPools"; +import { useProviderAccountPools, type AccountSelectionTarget } from "../hooks/useProviderAccountPools"; import { useCodexAccountPool } from "../hooks/useCodexAccountPool"; import { useJsonConfigEditor } from "../hooks/useJsonConfigEditor"; import { useKeyedClientResource } from "../client-resource"; @@ -75,6 +75,137 @@ export function useQuotaRefreshCoordinator(apiBase: string) { return { quotaRefresh, invalidateProviderQuotas, settleQuotaRefresh, beginQuotaRefresh }; } +/** One authenticated SSE connection, with bounded reconnect backoff and scheduler recovery. */ +function useAccountSelectionEvents( + apiBase: string, + enabled: boolean, + refresh: (target?: AccountSelectionTarget) => Promise, +) { + const refreshRef = useRef(refresh); + useLayoutEffect(() => { refreshRef.current = refresh; }); + const recoverRef = useRef<() => void>(() => {}); + useEffect(() => { + if (!enabled) return; + let stopped = false; + let retryTimer: ReturnType | null = null; + let retryDelay = 250; + type Connection = { controller: AbortController; reader?: ReadableStreamDefaultReader; lastActivity: number; openedAt?: number }; + let connection: Connection | null = null; + const clearRetry = () => { + if (retryTimer !== null) window.clearTimeout(retryTimer); + retryTimer = null; + }; + const close = (current: Connection) => { + current.controller.abort(); + void current.reader?.cancel().catch(() => {}); + }; + const connect = () => { + if (stopped || connection) return; + clearRetry(); + const current: Connection = { controller: new AbortController(), lastActivity: Date.now() }; + connection = current; + void (async () => { + try { + // Native EventSource cannot send the session/relay headers installed by api.ts. + const response = await fetch(`${apiBase}/api/accounts/events`, { + signal: current.controller.signal, credentials: "same-origin", headers: { Accept: "text/event-stream" }, + }); + if (!response.ok || !response.headers.get("content-type")?.includes("text/event-stream") || !response.body) { + throw new Error("Account selection stream unavailable"); + } + if (stopped || current.controller.signal.aborted) { await response.body.cancel(); return; } + const reader = response.body.getReader(); + current.reader = reader; + current.openedAt = Date.now(); + const decoder = new TextDecoder(); + const revisions = new Map(); + const pending = new Map(); + let refreshAll = false; + let queued = false; + const flush = () => { + if (queued) return; + queued = true; + void Promise.resolve().then(async () => { + queued = false; + if (stopped || current.controller.signal.aborted) return; + const targets = refreshAll ? [undefined] : [...pending.values()]; + refreshAll = false; + pending.clear(); + await Promise.all(targets.map(target => refreshRef.current(target))); + }).catch(() => { /* The recovery tick retries failed invalidation reads. */ }); + }; + let buffer = ""; + let event = ""; + let data: string[] = []; + let frameSize = 0; + const dispatch = () => { + let value: { provider?: unknown; kind?: unknown; revision?: unknown }; + try { value = JSON.parse(data.join("\n")) as typeof value; } catch { return; } + if (!value || typeof value !== "object" || typeof value.revision !== "number" || !Number.isSafeInteger(value.revision) || value.revision < 0) return; + if (event === "ready") { + revisions.clear(); + refreshAll = true; + flush(); + } else if (event === "account-selection" && typeof value.provider === "string" && value.provider + && (value.kind === "oauth" || value.kind === "api-key")) { + const key = `${value.kind}:${value.provider}`; + if (value.revision <= (revisions.get(key) ?? -1)) return; + revisions.set(key, value.revision); + pending.set(key, { provider: value.provider, kind: value.kind }); + flush(); + } + }; + while (!stopped && !current.controller.signal.aborted) { + const chunk = await reader.read(); + if (chunk.done) break; + current.lastActivity = Date.now(); + buffer += decoder.decode(chunk.value, { stream: true }); + let end: number; + while ((end = buffer.indexOf("\n")) !== -1) { + const line = buffer.slice(0, end).replace(/\r$/, ""); + buffer = buffer.slice(end + 1); + frameSize += line.length; + if (frameSize > 16_384) throw new Error("Account selection event too large"); + if (!line) { dispatch(); event = ""; data = []; frameSize = 0; } + else if (line.startsWith("event:")) event = line.slice(6).replace(/^ /, ""); + else if (line.startsWith("data:")) data.push(line.slice(5).replace(/^ /, "")); + } + if (buffer.length + frameSize > 16_384) throw new Error("Account selection event too large"); + } + } catch { + // api.ts owns authentication. Transport failures follow the same retry path as EOF. + } finally { + close(current); + if (connection === current) { + connection = null; + if (!stopped) { + // Only a stable connection resets backoff; repeated ready-then-EOF cannot spin. + if (current.openedAt !== undefined && Date.now() - current.openedAt >= 10_000) retryDelay = 250; + const delay = retryDelay; + retryDelay = Math.min(retryDelay * 2, 5_000); + retryTimer = window.setTimeout(() => { retryTimer = null; connect(); }, delay); + } + } + } + })(); + }; + const recover = () => { + // Also retire a hung handshake or a silent connection that lost its heartbeat. + if (connection && Date.now() - connection.lastActivity > 60_000) { close(connection); connection = null; } + connect(); + }; + recoverRef.current = recover; + connect(); + return () => { + stopped = true; + recoverRef.current = () => {}; + clearRetry(); + if (connection) close(connection); + }; + }, [apiBase, enabled]); + return useCallback(() => recoverRef.current(), []); +} + export default function Providers({ apiBase }: { apiBase: string }) { const t = useT(); const configCacheKey = `ocx.providers.config.v1:${apiBase}`; @@ -243,9 +374,23 @@ export default function Providers({ apiBase }: { apiBase: string }) { }); const { accountSets, setAccountSets, accountLoadStates, switchingAccount, keyPools, fetchAccountSets, fetchKeyPools, + refreshAccountRosters, oauthCardProviders, keyCardProviders, switchAccount, switchApiKey, removeApiKey, addApiKeyValue, editCredentialAlias, removeAccount, activeAccountNeedsReauth, } = pools; + const refreshSelection = useCallback((target?: AccountSelectionTarget) => { + if (target && !(target.kind === "oauth" ? oauthCardProviders : keyCardProviders).includes(target.provider)) return Promise.resolve(true); + return refreshAccountRosters(target); + }, [refreshAccountRosters, oauthCardProviders, keyCardProviders]); + const recoverSelectionStream = useAccountSelectionEvents(apiBase, config !== null, refreshSelection); + const rosterKey = JSON.stringify([apiBase, oauthCardProviders.toSorted(), keyCardProviders.toSorted()]); + const rosterRecoveryKeyRef = useRef(null); + useKeyedClientResource(`provider-rosters:${rosterKey}`, [rosterKey], async signal => { + // Existing bootstrap effects own the first enriched reads. This resource is recovery only. + if (rosterRecoveryKeyRef.current !== rosterKey) { rosterRecoveryKeyRef.current = rosterKey; return true; } + recoverSelectionStream(); + return refreshAccountRosters(undefined, signal); + }, { enabled: config !== null, pollMs: 30_000 }); const jsonEditor = useJsonConfigEditor({ apiBase, config, notify, diff --git a/gui/src/pages/use-providers-oauth.ts b/gui/src/pages/use-providers-oauth.ts index d804ec3739..3440939ef1 100644 --- a/gui/src/pages/use-providers-oauth.ts +++ b/gui/src/pages/use-providers-oauth.ts @@ -143,11 +143,11 @@ export function useProvidersOAuth({ finished = true; break; } - // Seed the account list from the status poll immediately so Accounts does not - // briefly render empty while the follow-up /api/oauth/accounts round-trip runs. + // Seed only a missing roster. This status read can predate a manual selection; + // an existing roster is reconciled by the guarded /accounts read below. if (s.accounts) { const activeFromRow = s.accounts.find(a => a.active)?.id ?? null; - setAccountSets(current => ({ + setAccountSets(current => current[provider] ? current : ({ ...current, [provider]: { activeAccountId: s.activeAccountId ?? activeFromRow, diff --git a/gui/tests/provider-account-quota-loading.test.tsx b/gui/tests/provider-account-quota-loading.test.tsx index 5b856945b4..bfd811ae7f 100644 --- a/gui/tests/provider-account-quota-loading.test.tsx +++ b/gui/tests/provider-account-quota-loading.test.tsx @@ -5,16 +5,103 @@ import { createRoot, type Root } from "react-dom/client"; import { useProviderAccountPools, type OAuthAccount, type ApiKeyEntry } from "../src/hooks/useProviderAccountPools"; const globals = ["document", "window", "navigator", "fetch", "IS_REACT_ACT_ENVIRONMENT"] as const; -let previous: Record<(typeof globals)[number], unknown>; +let previous: Record<(typeof globals)[number], PropertyDescriptor | undefined>; let win: Window; let root: Root | null; let host: HTMLElement; let pools: ReturnType; let requests: Array<{ url: string; signal?: AbortSignal | null }>; -let respond: (url: string, signal?: AbortSignal | null) => Promise; +let respond: (url: string, signal?: AbortSignal | null, init?: RequestInit) => Promise; const noop = async () => {}; const reading = { fiveHourPercent: 21, weeklyPercent: 34, updatedAt: 1_700_000_000_000 }; +for (const kind of ["oauth", "api-key"] as const) { + for (const quotaMode of ["probe", "passive"] as const) { + test(`${kind} ${quotaMode} newer quota failure cannot be cleared by a pre-selection probe`, async () => { + let selected = "a"; + const oldQuota = deferred(); + const oldStarted = deferred(); + let probes = 0; + const rows = () => ["a", "b"].map(id => ({ id, masked: id, active: id === selected, quotaMode })); + respond = async (url, _signal, init) => { + if (init?.method === "PUT") { selected = "b"; return Response.json({ ok: true, activeAccountId: "b", activeId: "b" }); } + if (url.includes("quota=1")) { + if (++probes === 1) { oldStarted.resolve(); return oldQuota.promise; } + return new Response(null, { status: 503 }); + } + return Response.json({ activeAccountId: selected, activeId: selected, accounts: rows(), keys: rows() }); + }; + const load = () => kind === "oauth" ? pools.fetchAccountSets(["fixture"], true) : pools.fetchKeyPools(["fixture"], true); + let old!: Promise; + await act(async () => { old = load(); await oldStarted.promise; }); + const before = rows(); + await act(async () => { + if (kind === "oauth") await pools.switchAccount("fixture", before[1]); + else await pools.switchApiKey("fixture", before[1]); + }); + await act(async () => { expect(await load()).toBe(false); }); + await act(async () => { + const enriched = before.map(row => ({ ...row, quota: reading })); + oldQuota.resolve(Response.json({ activeAccountId: "a", activeId: "a", accounts: enriched, keys: enriched })); + expect(await old).toBe(false); + }); + const current = kind === "oauth" ? pools.accountSets.fixture.accounts : pools.keyPools.fixture; + expect(current.find(row => row.active)?.id).toBe("b"); + expect(current[0]).toMatchObject({ quotaPending: false, quotaUnavailable: true }); + expect(current[0].quota).toBeUndefined(); + }); + + for (const outcome of ["success", "unavailable", "null", "http-error"] as const) { + test(`${kind} ${quotaMode} manual selection preserves initial quota ownership (${outcome})`, async () => { + let selected = "a"; + let ids = ["a", "b", "removed"]; + const quota = deferred(); + const started = deferred(); + const roster = () => ids.map(id => ({ id, masked: id, active: id === selected, quotaMode })); + respond = async (url, _signal, init) => { + if (init?.method === "PUT") { + selected = "b"; + ids = ["a", "b", "new"]; + return Response.json({ ok: true, activeAccountId: selected, activeId: selected }); + } + if (url.includes("quota=1")) { started.resolve(); return quota.promise; } + return Response.json({ activeAccountId: selected, activeId: selected, accounts: roster(), keys: roster() }); + }; + let full!: Promise; + await act(async () => { + full = kind === "oauth" ? pools.fetchAccountSets(["fixture"], true) : pools.fetchKeyPools(["fixture"], true); + await started.promise; + }); + const before = roster(); + const current = () => kind === "oauth" ? pools.accountSets.fixture.accounts : pools.keyPools.fixture; + expect(current()[0].quotaPending).toBe(quotaMode === "probe"); + await act(async () => { + if (kind === "oauth") await pools.switchAccount("fixture", before[1]); + else await pools.switchApiKey("fixture", before[1]); + }); + expect(current().find(row => row.active)?.id).toBe("b"); + expect(current()[0].quotaPending).toBe(quotaMode === "probe"); + const enriched = before.map(row => ({ ...row, + quota: outcome === "null" ? null : reading, + quotaUnavailable: outcome === "unavailable" || outcome === "null", + })); + await act(async () => { + quota.resolve(outcome === "http-error" ? new Response(null, { status: 503 }) + : Response.json({ activeAccountId: "a", activeId: "a", accounts: enriched, keys: enriched })); + expect(await full).toBe(outcome === "success"); + }); + expect(current().map(row => row.id)).toEqual(["a", "b", "new"]); + expect(current().find(row => row.active)?.id).toBe("b"); + expect(current()[0]).toMatchObject({ quotaPending: false, quotaUnavailable: outcome !== "success" }); + expect(current()[0].quota).toEqual(outcome === "null" ? null : outcome === "http-error" ? undefined : reading); + expect(current()[2].quota).toBeUndefined(); + expect(current()[2].quotaUnavailable).toBe(false); + expect(requests.filter(request => request.url.includes("quota=1"))).toHaveLength(1); + }); + } + } +} + function deferred() { let resolve!: (value: T) => void; const promise = new Promise(done => { resolve = done; }); @@ -29,7 +116,7 @@ function Harness({ apiBase = "/quota-hook" }: { apiBase?: string }) { return null; } beforeEach(async () => { - previous = Object.fromEntries(globals.map(key => [key, Reflect.get(globalThis, key)])) as typeof previous; + previous = Object.fromEntries(globals.map(key => [key, Object.getOwnPropertyDescriptor(globalThis, key)])) as typeof previous; win = new Window({ url: "http://localhost" }); Object.defineProperties(globalThis, { document: { configurable: true, value: win.document }, window: { configurable: true, value: win }, @@ -39,18 +126,148 @@ beforeEach(async () => { respond = async () => Response.json({ accounts: [], keys: [] }); Object.defineProperty(globalThis, "fetch", { configurable: true, value: (input: RequestInfo | URL, init?: RequestInit) => { requests.push({ url: String(input), signal: init?.signal }); - return respond(String(input), init?.signal); + return respond(String(input), init?.signal, init); } }); host = win.document.createElement("div") as unknown as HTMLElement; win.document.body.appendChild(host as never); await act(async () => { root = createRoot(host); root.render(); }); }); + +for (const kind of ["oauth", "api-key"] as const) { + test(`${kind} stale base failure cannot mark a newer successful selection read unavailable`, async () => { + const rows = ["a", "b"].map(id => ({ id, masked: id, active: id === "a", quotaMode: "probe" as const, + quota: reading, quotaPending: false, quotaUnavailable: false })); + await act(async () => { + if (kind === "oauth") pools.setAccountSets({ fixture: { activeAccountId: "a", accounts: rows } }); + else pools.setKeyPools({ fixture: rows }); + }); + const old = deferred(); + respond = async () => old.promise; + let full!: Promise; + await act(async () => { full = kind === "oauth" ? pools.fetchAccountSets(["fixture"]) : pools.fetchKeyPools(["fixture"]); }); + respond = async () => Response.json({ activeAccountId: "b", activeId: "b", accounts: rows, keys: rows }); + await act(async () => { await pools.refreshAccountRosters({ provider: "fixture", kind }); }); + await act(async () => { old.resolve(new Response(null, { status: 503 })); expect(await full).toBe(false); }); + const result = kind === "oauth" ? pools.accountSets.fixture.accounts : pools.keyPools.fixture; + expect(result.find(row => row.active)?.id).toBe("b"); + expect(result[0]).toMatchObject({ quota: reading, quotaUnavailable: false, quotaPending: false }); + }); + + test(`${kind} selection reads preserve settled quota flags without enrichment`, async () => { + const rows = ["a", "b"].map(id => ({ id, active: id === "a", masked: id, + quotaMode: "probe" as const, quota: reading, quotaPending: false, quotaUnavailable: true })); + await act(async () => { + if (kind === "oauth") pools.setAccountSets({ fixture: { activeAccountId: "a", accounts: rows } }); + else pools.setKeyPools({ fixture: rows }); + }); + respond = async () => Response.json({ activeAccountId: "b", activeId: "b", + accounts: rows.map(row => ({ id: row.id, active: row.id === "b", quotaMode: "probe" })), + keys: rows.map(row => ({ id: row.id, masked: row.id, active: row.id === "b", quotaMode: "probe" })), + }); + await act(async () => { await pools.refreshAccountRosters({ provider: "fixture", kind }); }); + const result = kind === "oauth" ? pools.accountSets.fixture.accounts : pools.keyPools.fixture; + expect(result.find(row => row.active)?.id).toBe("b"); + expect(result[0]).toMatchObject({ quota: reading, quotaPending: false, quotaUnavailable: true }); + expect(requests).toHaveLength(1); + expect(requests[0].url).not.toContain("quota="); + }); + + test(`${kind} PUT invalidates old reads at start and settle and publishes the accepted selection immediately`, async () => { + const rows = ["a", "b"].map(id => ({ id, active: id === "a", masked: id, quotaMode: "unsupported" as const })); + await act(async () => { + if (kind === "oauth") pools.setAccountSets({ fixture: { activeAccountId: "a", accounts: rows } }); + else pools.setKeyPools({ fixture: rows }); + }); + const beforePut = deferred(); + const duringPut = deferred(); + const afterPut = deferred(); + const put = deferred(); + let reads = 0; + respond = async (_url, _signal, init) => init?.method === "PUT" ? put.promise + : ++reads === 1 ? beforePut.promise : reads === 2 ? duringPut.promise : afterPut.promise; + const read = () => kind === "oauth" ? pools.fetchAccountSets(["fixture"]) : pools.fetchKeyPools(["fixture"]); + let old!: Promise; + let during!: Promise; + let switched!: Promise; + await act(async () => { old = read(); }); + await act(async () => { switched = kind === "oauth" ? pools.switchAccount("fixture", rows[1]) : pools.switchApiKey("fixture", rows[1]); }); + const roster = (id: string) => Response.json({ activeAccountId: id, activeId: id, + accounts: rows.map(row => ({ ...row, active: row.id === id })), keys: rows.map(row => ({ ...row, active: row.id === id })) }); + await act(async () => { beforePut.resolve(roster("b")); expect(await old).toBe(false); }); + expect((kind === "oauth" ? pools.accountSets.fixture.accounts : pools.keyPools.fixture).find(row => row.active)?.id).toBe("a"); + await act(async () => { during = read(); }); + await act(async () => { put.resolve(Response.json({ ok: true, activeAccountId: "b", activeId: "b" })); }); + expect((kind === "oauth" ? pools.accountSets.fixture.accounts : pools.keyPools.fixture).find(row => row.active)?.id).toBe("b"); + await act(async () => { duringPut.resolve(roster("a")); expect(await during).toBe(false); }); + expect((kind === "oauth" ? pools.accountSets.fixture.accounts : pools.keyPools.fixture).find(row => row.active)?.id).toBe("b"); + await act(async () => { afterPut.resolve(roster("b")); await switched; }); + }); +} afterEach(async () => { if (root) await act(async () => { root!.unmount(); root = null; }); - for (const key of globals) Object.defineProperty(globalThis, key, { configurable: true, value: previous[key] }); + for (const key of globals) { + const descriptor = previous[key]; + if (descriptor) Object.defineProperty(globalThis, key, descriptor); + else Reflect.deleteProperty(globalThis, key); + } await win.happyDOM.close(); }); +for (const kind of ["oauth", "api-key"] as const) { + test(`${kind} push refresh overtakes a slow quota probe while late quota cannot change selection or resurrect rows`, async () => { + const quota = deferred(); + const started = deferred(); + const original = ["a", "b", "removed"].map(id => ({ id, masked: id, active: id === "a", quotaMode: "probe" as const })); + respond = async url => { + if (url.includes("quota=1")) { started.resolve(); return quota.promise; } + return Response.json({ activeAccountId: "a", activeId: "a", accounts: original, keys: original }); + }; + let full!: Promise; + await act(async () => { + full = kind === "oauth" ? pools.fetchAccountSets(["fixture"], true) : pools.fetchKeyPools(["fixture"], true); + await started.promise; + }); + const latest = original.filter(row => row.id !== "removed").map(row => ({ ...row, active: row.id === "b" })); + respond = async () => Response.json({ activeAccountId: "b", activeId: "b", accounts: latest, keys: latest }); + await act(async () => { expect(await pools.refreshAccountRosters({ provider: "fixture", kind })).toBe(true); }); + const current = () => kind === "oauth" ? pools.accountSets.fixture.accounts : pools.keyPools.fixture; + expect(current().find(row => row.active)?.id).toBe("b"); + expect(current()[0].quotaPending).toBe(true); + const enriched = original.map(row => ({ ...row, quota: reading })); + await act(async () => { + quota.resolve(Response.json({ activeAccountId: "a", activeId: "a", accounts: enriched, keys: enriched })); + expect(await full).toBe(true); + }); + expect(current().map(row => row.id)).toEqual(["a", "b"]); + expect(current().find(row => row.active)?.id).toBe("b"); + expect(current()[0]).toMatchObject({ quota: reading, quotaPending: false }); + expect(requests.filter(request => request.url.includes("quota=1"))).toHaveLength(1); + }); + + test(`${kind} rejected manual PUT retires reads begun during the write and reconciles without changing quota health`, async () => { + const rows = ["a", "b"].map(id => ({ id, masked: id, active: id === "a", quotaMode: "probe" as const, + quota: reading, quotaUnavailable: true, quotaPending: false })); + await act(async () => { + if (kind === "oauth") pools.setAccountSets({ fixture: { activeAccountId: "a", accounts: rows } }); + else pools.setKeyPools({ fixture: rows }); + }); + const put = deferred(); + const during = deferred(); + let reads = 0; + respond = async (_url, _signal, init) => init?.method === "PUT" ? put.promise : ++reads === 1 ? during.promise + : Response.json({ activeAccountId: "a", activeId: "a", accounts: rows, keys: rows }); + let changed!: Promise; + let old!: Promise; + await act(async () => { changed = kind === "oauth" ? pools.switchAccount("fixture", rows[1]) : pools.switchApiKey("fixture", rows[1]); }); + await act(async () => { old = kind === "oauth" ? pools.fetchAccountSets(["fixture"]) : pools.fetchKeyPools(["fixture"]); }); + await act(async () => { put.resolve(new Response(null, { status: 409 })); await changed; }); + await act(async () => { during.resolve(Response.json({ activeAccountId: "b", activeId: "b", accounts: [], keys: [] })); expect(await old).toBe(false); }); + const current = kind === "oauth" ? pools.accountSets.fixture.accounts : pools.keyPools.fixture; + expect(current.find(row => row.active)?.id).toBe("a"); + expect(current[0]).toMatchObject({ quota: reading, quotaUnavailable: true, quotaPending: false }); + }); +} + test("cheap probe rows paint with same-ID last-good; forced enrichment awaits and HTTP failure clears pending", async () => { const account: OAuthAccount = { id: "same", active: true, quotaMode: "probe", quota: reading }; await act(async () => { pools.setAccountSets({ oauth: { activeAccountId: "same", accounts: [account, { ...account, id: "removed" }] } }); }); diff --git a/gui/tests/provider-revalidation-policy.test.tsx b/gui/tests/provider-revalidation-policy.test.tsx index df5a983d80..e6be1936e1 100644 --- a/gui/tests/provider-revalidation-policy.test.tsx +++ b/gui/tests/provider-revalidation-policy.test.tsx @@ -4,7 +4,7 @@ import { act } from "react"; import { createRoot, type Root } from "react-dom/client"; import Providers from "../src/pages/Providers"; import { LanguageProvider } from "../src/i18n/provider"; -import { clearClientResourceStoresForTests } from "../src/client-resource"; +import { clearClientResourceStoresForTests, pollBucketCountForTests } from "../src/client-resource"; /** * Quota revalidation policy. @@ -19,7 +19,7 @@ import { clearClientResourceStoresForTests } from "../src/client-resource"; */ const globals = ["document", "window", "navigator", "localStorage", "sessionStorage", "fetch", "IS_REACT_ACT_ENVIRONMENT"] as const; -let previousGlobals: Record<(typeof globals)[number], unknown>; +let previousGlobals: Record<(typeof globals)[number], PropertyDescriptor | undefined>; let testWindow: Window; let container: HTMLElement; let root: Root | null = null; @@ -28,7 +28,7 @@ let quotaCalls: string[] = []; const PROVIDERS = ["anthropic", "cursor", "kimi"]; beforeEach(() => { - previousGlobals = Object.fromEntries(globals.map(k => [k, Reflect.get(globalThis, k)])) as typeof previousGlobals; + previousGlobals = Object.fromEntries(globals.map(k => [k, Object.getOwnPropertyDescriptor(globalThis, k)])) as typeof previousGlobals; clearClientResourceStoresForTests(); testWindow = new Window({ url: "http://localhost/#providers" }); Object.defineProperty(testWindow.navigator, "language", { configurable: true, value: "en-US" }); @@ -100,7 +100,9 @@ afterEach(async () => { } clearClientResourceStoresForTests(); for (const key of globals) { - Object.defineProperty(globalThis, key, { configurable: true, value: previousGlobals[key] }); + const descriptor = previousGlobals[key]; + if (descriptor) Object.defineProperty(globalThis, key, descriptor); + else Reflect.deleteProperty(globalThis, key); } }); @@ -221,3 +223,143 @@ for (const kind of ["oauth", "key", "codex"] as const) { expect(container.textContent).toContain("Quota check completed"); }); } + +for (const kind of ["oauth", "api-key"] as const) { + test(`${kind} selection event updates the rendered active row without a poll tick or quota probe`, async () => { + const name = kind === "oauth" ? "cursor" : "deepseek"; + let selected = "a"; + let stream!: ReadableStreamDefaultController; + let streamSignal: AbortSignal | null | undefined; + let subscriptions = 0; + let rejectConnection = false; + let cancelled = false; + const reads: string[] = []; + const encoder = new TextEncoder(); + const send = (event: string, data: unknown) => { + const bytes = encoder.encode(`event: ${event}\r\ndata: ${JSON.stringify(data)}\r\n\r\n`); + // Network chunks need not end on either a line or an event boundary. + stream.enqueue(bytes.slice(0, 11)); + stream.enqueue(bytes.slice(11)); + }; + Object.defineProperty(globalThis, "fetch", { configurable: true, value: async (input: RequestInfo | URL, init?: RequestInit) => { + const url = new URL(String(input), "http://localhost"); + reads.push(url.pathname + url.search); + if (url.pathname === "/api/accounts/events") { + subscriptions += 1; + streamSignal = init?.signal; + if (rejectConnection) return new Response(null, { status: 503 }); + return new Response(new ReadableStream({ + start(controller) { stream = controller; }, cancel() { cancelled = true; }, + }), { headers: { "Content-Type": "text/event-stream" } }); + } + if (url.pathname === "/api/config") return Response.json({ port: 10100, defaultProvider: name, providers: { + [name]: { adapter: "openai-chat", authMode: kind === "oauth" ? "oauth" : "key", hasApiKey: kind === "api-key", baseUrl: "https://fixture.test/v1" }, + } }); + if (url.pathname === "/api/oauth/providers") return Response.json({ providers: kind === "oauth" ? [name] : [] }); + if (url.pathname === "/api/oauth/status") return Response.json({ loggedIn: true }); + if (url.pathname === "/api/oauth/accounts" || url.pathname === "/api/providers/keys") { + const rows = ["a", "b"].map(id => ({ id, alias: `Choice ${id.toUpperCase()}`, label: `Choice ${id.toUpperCase()}`, masked: id, + active: selected === id, quotaMode: "probe", ...(url.searchParams.has("quota") ? { + quota: { fiveHourPercent: id === "a" ? 30 : 11, updatedAt: 1_700_000_000_000 }, + } : {}), + })); + return Response.json({ activeAccountId: selected, activeId: selected, accounts: rows, keys: rows }); + } + if (url.pathname === "/api/codex-auth/accounts") return Response.json({ accounts: [] }); + if (url.pathname === "/api/codex-auth/active") return Response.json({ activeCodexAccountId: null }); + if (url.pathname === "/api/provider-quotas") return Response.json({ reports: [] }); + if (url.pathname === "/api/selected-models") return Response.json({ models: {} }); + if (url.pathname === "/api/usage") return Response.json({ providers: [], models: [] }); + return Response.json({}); + } }); + await act(async () => { root = createRoot(container); root.render(); }); + await act(async () => { container.querySelector(".providers-workspace-rail-row")!.click(); }); + const accountsTab = Array.from(container.querySelectorAll("[role=tab]")) + .find(button => /Accounts|Keys/.test(button.textContent ?? "")); + expect(accountsTab).toBeDefined(); + await act(async () => { accountsTab!.click(); }); + expect(container.querySelector(".pwi-auth-acct--active")?.textContent).toContain("Choice A"); + expect(subscriptions).toBe(1); + const quotaReads = () => reads.filter(path => path.includes("quota=1") || path.startsWith("/api/provider-quotas")).length; + const initialQuotas = quotaReads(); + const initialReads = reads.length; + await act(async () => { send("ready", { revision: 0 }); }); + expect(reads.length).toBeGreaterThan(initialReads); + selected = "b"; + await act(async () => { send("account-selection", { provider: name, kind, revision: 1 }); }); + expect(container.querySelector(".pwi-auth-acct--active")?.textContent).toContain("Choice B"); + expect(container.querySelector(".pwi-auth-acct--active")?.textContent).not.toContain("Choice A"); + expect(quotaReads()).toBe(initialQuotas); + expect(subscriptions).toBe(1); + expect(pollBucketCountForTests()).toBe(1); + const afterEvent = reads.length; + await act(async () => { + send("account-selection", { provider: name, kind, revision: 1 }); + send("account-selection", { provider: "unknown-provider", kind, revision: 2 }); + }); + expect(reads).toHaveLength(afterEvent); + // Only the reconnect timeout advances. No 30-second scheduler tick or visibility wake. + const retries = new Map void; delay: number }>(); + let timerId = 10_000; + const setTimeoutBefore = testWindow.setTimeout.bind(testWindow); + const clearTimeoutBefore = testWindow.clearTimeout.bind(testWindow); + Object.defineProperty(testWindow, "setTimeout", { configurable: true, value: (run: () => void, delay: number) => { + if (delay > 5_000) return setTimeoutBefore(run, delay); + const id = ++timerId; + retries.set(id, { run, delay }); + return id; + } }); + Object.defineProperty(testWindow, "clearTimeout", { configurable: true, value: (id: number) => { + if (!retries.delete(id)) clearTimeoutBefore(id); + } }); + const retry = async (delay: number) => { + expect(retries.size).toBe(1); + const [id, pending] = [...retries][0]; + expect(pending.delay).toBe(delay); + retries.delete(id); + await act(async () => { pending.run(); }); + }; + await act(async () => { stream.close(); }); + selected = "a"; + await retry(250); + expect(subscriptions).toBe(2); + await act(async () => { send("ready", { revision: 0 }); }); + expect(container.querySelector(".pwi-auth-acct--active")?.textContent).toContain("Choice A"); + expect(quotaReads()).toBe(initialQuotas); + + await act(async () => { stream.error(new Error("connection interrupted")); }); + rejectConnection = true; + await retry(500); + expect(subscriptions).toBe(3); + for (const delay of [1_000, 2_000, 4_000, 5_000]) await retry(delay); + expect(subscriptions).toBe(7); + rejectConnection = false; + await retry(5_000); + expect(subscriptions).toBe(8); + selected = "b"; + await act(async () => { send("ready", { revision: 0 }); }); + expect(container.querySelector(".pwi-auth-acct--active")?.textContent).toContain("Choice B"); + expect(quotaReads()).toBe(initialQuotas); + + // Cleanup cancels scheduled retries, and a callback already dequeued cannot orphan a loop. + await act(async () => { stream.close(); }); + expect(retries.size).toBe(1); + const queuedRetry = [...retries.values()][0].run; + await act(async () => { root!.unmount(); root = null; }); + expect(retries.size).toBe(0); + await act(async () => { queuedRetry(); }); + expect(subscriptions).toBe(8); + expect(streamSignal?.aborted).toBe(true); + // Closed/error streams need no cancel callback; also prove cleanup of a live reader. + await act(async () => { + root = createRoot(container); + root.render(); + }); + expect(subscriptions).toBe(9); + expect(streamSignal?.aborted).toBe(false); + await act(async () => { root!.unmount(); root = null; }); + expect(streamSignal?.aborted).toBe(true); + expect(cancelled).toBe(true); + expect(retries.size).toBe(0); + }); +} diff --git a/src/client/hub-relay.ts b/src/client/hub-relay.ts index 820ffd3845..95ab7b0c0d 100644 --- a/src/client/hub-relay.ts +++ b/src/client/hub-relay.ts @@ -164,6 +164,7 @@ function boundedRelayResponseStream( body: ReadableStream, limit: number, signal: AbortSignal, + cleanup: () => void, ): ReadableStream { const reader = body.getReader(); let bytes = 0; @@ -172,6 +173,7 @@ function boundedRelayResponseStream( if (finished) return; finished = true; signal.removeEventListener("abort", onAbort); + cleanup(); try { reader.releaseLock(); } catch { /* a pending read may still own it */ } }; const onAbort = () => { @@ -245,9 +247,19 @@ export async function relayHubManagementRequest( ? Math.min(Math.floor(deps.timeoutMs), 120_000) : HUB_RELAY_DEFAULT_TIMEOUT_MS; const timeoutSignal = AbortSignal.timeout(timeoutMs); - const signal = req.signal - ? AbortSignal.any([req.signal, timeoutSignal]) - : timeoutSignal; + const relayAbort = new AbortController(); + const signal = relayAbort.signal; + const stopDeadline = () => timeoutSignal.removeEventListener("abort", onTimeout); + const cleanup = () => { + stopDeadline(); + req.signal.removeEventListener("abort", onClientAbort); + }; + const onTimeout = () => { relayAbort.abort(timeoutSignal.reason); cleanup(); }; + const onClientAbort = () => { relayAbort.abort(req.signal.reason); cleanup(); }; + timeoutSignal.addEventListener("abort", onTimeout, { once: true }); + req.signal.addEventListener("abort", onClientAbort, { once: true }); + if (req.signal.aborted) onClientAbort(); + else if (timeoutSignal.aborted) onTimeout(); let upstream: Response; try { upstream = await (deps.fetchImpl ?? fetch)(destination, { @@ -258,9 +270,16 @@ export async function relayHubManagementRequest( signal, }); } catch { + cleanup(); + return relayError(502, "hub relay unavailable"); + } + if (signal.aborted) { + cleanup(); + try { await upstream.body?.cancel(); } catch { /* best effort */ } return relayError(502, "hub relay unavailable"); } if (upstream.status >= 300 && upstream.status < 400) { + cleanup(); try { await upstream.body?.cancel(); } catch { /* best effort */ } return relayError(502, "hub relay redirect refused"); } @@ -268,18 +287,31 @@ export async function relayHubManagementRequest( const responseConnectionNamed = new Set((upstream.headers.get("connection") ?? "").split(",").map(value => value.trim().toLowerCase()).filter(Boolean)); const responseHeaders = filteredHeaders(upstream.headers, RESPONSE_HEADERS, responseConnectionNamed); if (!headersWithinLimit(responseHeaders)) { + cleanup(); try { await upstream.body?.cancel(); } catch { /* best effort */ } return relayError(502, "hub relay response headers too large"); } const declaredResponseLength = upstream.headers.get("content-length"); if (declaredResponseLength !== null && (!/^\d+$/.test(declaredResponseLength) || Number(declaredResponseLength) > HUB_RELAY_RESPONSE_BODY_MAX_BYTES)) { + cleanup(); try { await upstream.body?.cancel(); } catch { /* best effort */ } return relayError(502, "hub relay response body too large"); } - const responseBody = method === "HEAD" || !upstream.body - ? null - : boundedRelayResponseStream(upstream.body, HUB_RELAY_RESPONSE_BODY_MAX_BYTES, signal); + // Only this known, successfully established SSE endpoint outlives the handshake. + // Its body remains byte-bounded and connected to the browser's abort signal. + if (method === "GET" && destination.pathname === "/api/accounts/events" && !destination.search + && upstream.status === 200 + && responseHeaders.get("content-type")?.split(";", 1)[0]?.trim().toLowerCase() === "text/event-stream") { + stopDeadline(); + } + let responseBody: ReadableStream | null = null; + if (method === "HEAD" || !upstream.body) { + cleanup(); + try { await upstream.body?.cancel(); } catch { /* best effort */ } + } else { + responseBody = boundedRelayResponseStream(upstream.body, HUB_RELAY_RESPONSE_BODY_MAX_BYTES, signal, cleanup); + } return new Response(responseBody, { status: upstream.status, statusText: upstream.statusText, diff --git a/src/images/loop.ts b/src/images/loop.ts index 0191215960..e3a7f8252f 100644 --- a/src/images/loop.ts +++ b/src/images/loop.ts @@ -253,6 +253,8 @@ export interface ImageBridgeDeps { stallTimeoutSec?: number; /** Provider-specific fetch (e.g. xAI transport wrapper). Falls back to global fetch. */ fetchImpl?: typeof globalThis.fetch; + /** Bind physical dispatch to this iteration's built request; pacing remains owned by the loop. */ + fetchForRequest?: (request: AdapterRequest, parsed: OcxParsedRequest) => typeof globalThis.fetch; /** Reserve the routed provider's next request-start slot before each adapter dispatch. */ waitForRequestSlot?: (signal?: AbortSignal) => Promise; /** Raw adapter usage at the terminal event, pre wire-normalization (see bridgeToResponsesSSE onUsage). */ @@ -499,6 +501,7 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise void>(); +let revision = 0; + +/** Call only after the authoritative selection has been persisted. */ +export function publishAccountSelection(provider: string, kind: AccountSelectionEvent["kind"]): void { + const event: AccountSelectionEvent = Object.freeze({ provider, kind, revision: ++revision }); + for (const listener of [...listeners]) { + try { + listener(event); + } catch { + // A disconnected consumer must not turn a committed write into a reported failure. + } + } +} + +export function subscribeAccountSelections(listener: (event: AccountSelectionEvent) => void): () => void { + // Give each subscription its own lifetime, even when a callback is reused. + const subscription = (event: AccountSelectionEvent) => listener(event); + listeners.add(subscription); + return () => { listeners.delete(subscription); }; +} + +export function currentAccountSelectionRevision(): number { + return revision; +} diff --git a/src/oauth/account-quota-rank.ts b/src/oauth/account-quota-rank.ts index 36360c655b..978f73de9d 100644 --- a/src/oauth/account-quota-rank.ts +++ b/src/oauth/account-quota-rank.ts @@ -67,6 +67,14 @@ function headroomOf(provider: string, accountId: string): number | null { return 100 - Math.max(...percents); } +/** Unknown usage is not exhaustion; Kiro's explicit overage verdict is authoritative. */ +export function isAccountQuotaExhausted(provider: string, accountId: string): boolean { + const exhaustion = provider === "kiro" ? getKiroAccountExhaustion(`${provider}\u0000${accountId}`) : null; + if (exhaustion !== null) return exhaustion.exhausted; + const headroom = headroomOf(provider, accountId); + return headroom !== null && headroom <= 0; +} + /** * Order candidates best-first. * @@ -90,7 +98,7 @@ export function rankAccountsByHeadroom(provider: string, ring: readonly string[] 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 (isAccountQuotaExhausted(provider, id)) 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 }; }); diff --git a/src/oauth/anthropic-routing.ts b/src/oauth/anthropic-routing.ts index 95197ef1f1..a029207be5 100644 --- a/src/oauth/anthropic-routing.ts +++ b/src/oauth/anthropic-routing.ts @@ -15,7 +15,8 @@ * store (existing OAuth path) so the account is excluded from eligibility. */ import { createHash } from "node:crypto"; -import { setActiveAccount, getAccountSet, getAccountCredential } from "./store"; +import { captureOAuthAccountSelection, commitOAuthAccountSelection, credentialGeneration, getAccountSet, getAccountCredential, getAccountCredentialWithStatus } from "./store"; +import type { OAuthAccessSnapshot } from "./index"; import { getCachedProviderAccountQuota } from "../providers/quota"; import { fallbackCodexAccountLogLabel } from "../codex/account-label"; import { @@ -24,6 +25,7 @@ import { notePoolRotationFailure, notePoolRotationSuccess, pickRoundRobinAccount, + peekRoundRobinAccount, POOL_KEY_ANTHROPIC, seedPoolRotationAccount, } from "../codex/pool-rotation"; @@ -68,6 +70,10 @@ interface AffinityEntry { const upstreamHealth = new Map(); const sessionAffinity = new Map(); +type OAuthAccountSelection = NonNullable>; +// Undefined means this runtime has not admitted a selection yet; null means consumed. +// The startup baseline comes from the authoritative store, never a second persisted pin. +let manualPreference: OAuthAccountSelection | null | undefined; function normalizeAffinityComponent(value: string | null | undefined): string { const normalized = value?.trim() ?? ""; @@ -152,6 +158,7 @@ export function sweepExpiredAnthropicRoutingHealth(now = Date.now()): number { export function clearAnthropicAccountPoolState(): void { upstreamHealth.clear(); sessionAffinity.clear(); + manualPreference = undefined; quorumCache = null; } @@ -409,7 +416,7 @@ function pickAlternateAnthropicAccount( const strategy = anthropicPoolStrategy(config); const eligible = getEligibleAnthropicAccounts(now).filter(id => id !== excludeId); if (strategy === "round-robin") { - return pickRoundRobinAccount(POOL_KEY_ANTHROPIC, eligible, stickyLimitForPool(config)); + return peekRoundRobinAccount(POOL_KEY_ANTHROPIC, eligible, stickyLimitForPool(config)); } if (strategy === "fill-first") { return pickNextFillFirstAnthropicAccount(config, excludeId, eligible); @@ -434,6 +441,7 @@ export type AnthropicAccountSelectionReason = | "lowest-usage" | "only-eligible" | "round-robin" + | "manual" | "fill-first" | "none" | "all-cooled"; @@ -500,9 +508,8 @@ function pickUnboundStrategyAccount( if (strategy === "round-robin") { const eligible = getEligibleAnthropicAccounts(now); const limit = stickyLimitForPool(config); - const picked = pickRoundRobinAccount(POOL_KEY_ANTHROPIC, eligible, limit); + const picked = peekRoundRobinAccount(POOL_KEY_ANTHROPIC, eligible, limit); if (!picked) return null; - notePoolRotationSuccess(POOL_KEY_ANTHROPIC, picked, limit); return { accountId: picked, reason: "round-robin" }; } @@ -528,17 +535,39 @@ export function resolveAnthropicAccountForSession( const set = getAccountSet(PROVIDER); if (!set || set.accounts.length === 0) return { accountId: null, reason: "none" }; + if (manualPreference === undefined) { + manualPreference = set.selectionRevision !== undefined + ? { accountId: set.activeAccountId, revision: set.selectionRevision } + : null; + } + if (!isAnthropicAccountPoolEnabled(config)) { return { accountId: set.activeAccountId, reason: "pool-disabled" }; } + // A manual choice is a one-dispatch preference, not a lower-priority quota hint. + // Consume it only after admission commits, so a failed token lookup cannot spend it. + if (manualPreference) { + if (manualPreference.accountId !== set.activeAccountId || manualPreference.revision !== set.selectionRevision) { + manualPreference = null; + } else { + const chosen = manualPreference.accountId; + const quota = getCachedProviderAccountQuota(PROVIDER, chosen); + const exhausted = [quota?.fiveHourPercent, quota?.weeklyPercent, quota?.monthlyPercent, + ...(quota?.customWindows ?? []).map(window => window.percent)] + .some(percent => typeof percent === "number" && percent >= 100); + if (!exhausted && getEligibleAnthropicAccounts(now).includes(chosen)) { + return { accountId: chosen, reason: "manual" }; + } + } + } + const key = normalizeAffinityComponent(sessionKey); if (key) { const affined = sessionAffinity.get(key); if (affined && now - affined.lastUsedAt <= AFFINITY_IDLE_TTL_MS) { const stillThere = set.accounts.some(a => a.id === affined.accountId && a.needsReauth !== true); if (stillThere && !isCooled(affined.accountId, now) && isPoolCredentialUsable(affined.accountId, now)) { - affined.lastUsedAt = now; return { accountId: affined.accountId, reason: "affinity" }; } sessionAffinity.delete(key); @@ -560,12 +589,6 @@ export function resolveAnthropicAccountForSession( const strategyPick = pickUnboundStrategyAccount(config, now); if (strategyPick) { - // Do not promote active here — token validation may still fail. Callers - // (responses/core) promote after getAnthropicPoolAccessToken succeeds. - if (key && normalizeAffinityComponent(strategyPick.accountId)) { - sessionAffinity.set(key, { accountId: strategyPick.accountId, lastUsedAt: now }); - pruneExpiredAffinity(now); - } return { accountId: strategyPick.accountId, reason: strategyPick.reason }; } @@ -611,10 +634,6 @@ export function resolveAnthropicAccountForSession( return { accountId: null, reason: anyCooled ? "all-cooled" : "none" }; } - if (key && normalizeAffinityComponent(accountId)) { - sessionAffinity.set(key, { accountId, lastUsedAt: now }); - pruneExpiredAffinity(now); - } return { accountId, reason }; } @@ -684,20 +703,57 @@ export function rotateAnthropicAccountOn429( return null; } - const affinityKey = normalizeAffinityComponent(sessionKey); - if (affinityKey && normalizeAffinityComponent(next)) { - sessionAffinity.set(affinityKey, { accountId: next, lastUsedAt: now }); - pruneExpiredAffinity(now); - } console.warn( `[anthropic-pool] 429 on ${formatAnthropicAccountOrdinal(failedAccountId)}; failing over to ${formatAnthropicAccountOrdinal(next)}`, ); return next; } -/** Promote dashboard active account after a validated failover target is usable. */ -export function promoteAnthropicActiveAccount(accountId: string): void { - void setActiveAccount(PROVIDER, accountId).catch(() => { /* best-effort */ }); +export interface AnthropicSelectionRoutingOptions { + config: OcxConfig; + sessionKey?: string | null; + reason?: AnthropicAccountSelectionReason; + expectedCredentialGeneration?: string; +} + +/** Commit the selected account before dispatch; rejected proposals have no routing side effects. */ +export async function promoteAnthropicActiveAccount( + accountId: string, + expectedSelection: OAuthAccountSelection | null, + options: AnthropicSelectionRoutingOptions, +): Promise { + if (!expectedSelection || !isPoolCredentialUsable(accountId, Date.now()) || isCooled(accountId, Date.now())) return null; + const committed = await commitOAuthAccountSelection(PROVIDER, accountId, { + expectedSelection, + expectedCredentialGeneration: options.expectedCredentialGeneration, + requireUsableAccount: true, + }); + if (!committed) return null; + return commitAnthropicSelectionRouting(accountId, expectedSelection, committed, options) ? committed : null; +} + +/** Main's shared selection owner calls this only after its authoritative commit succeeds. */ +export function commitAnthropicSelectionRouting( + accountId: string, + expectedSelection: OAuthAccountSelection, + committed: OAuthAccountSelection, + options: AnthropicSelectionRoutingOptions, +): boolean { + if (committed.accountId !== accountId) return false; + const current = captureOAuthAccountSelection(PROVIDER); + if (current?.accountId !== committed.accountId || current.revision !== committed.revision) return false; + if (isAnthropicAccountPoolEnabled(options.config)) { + if (anthropicPoolStrategy(options.config) === "round-robin" && options.reason !== "affinity") { + const limit = stickyLimitForPool(options.config); + const picked = pickRoundRobinAccount(POOL_KEY_ANTHROPIC, getEligibleAnthropicAccounts(), limit); + if (picked !== accountId) seedPoolRotationAccount(POOL_KEY_ANTHROPIC, accountId); + notePoolRotationSuccess(POOL_KEY_ANTHROPIC, accountId, limit); + } + bindAnthropicSessionAffinity(options.sessionKey, accountId); + } + if (manualPreference === undefined || (manualPreference?.accountId === expectedSelection.accountId + && manualPreference.revision === expectedSelection.revision)) manualPreference = null; + return true; } /** @@ -706,6 +762,7 @@ export function promoteAnthropicActiveAccount(accountId: string): void { */ export function resetAnthropicRoutingForManualSelection(accountId: string): void { sessionAffinity.clear(); + manualPreference = captureOAuthAccountSelection(PROVIDER); seedPoolRotationAccount(POOL_KEY_ANTHROPIC, accountId); // A manual account selection is an operator statement about the roster; do not answer the // next activation question from a count read before it. @@ -731,6 +788,17 @@ export async function getAnthropicPoolAccessToken(accountId: string): Promise { + const accessToken = await getAnthropicPoolAccessToken(accountId); + const row = getAccountCredentialWithStatus(PROVIDER, accountId); + if (!row || row.needsReauth || row.credential.access !== accessToken + || row.credential.expires <= Date.now()) { + throw new Error("Anthropic pool credential changed during account selection"); + } + return { provider: PROVIDER, accountId, accessToken, generation: credentialGeneration(row.credential) }; +} + /** * Whether the pool may refresh this account's token. Background `local-cli` slots must not * adopt the global Claude CLI credential (same fail-closed rule as quota probes). diff --git a/src/oauth/generic-account-failover.ts b/src/oauth/generic-account-failover.ts index 321b5f92a1..1ccfaf71df 100644 --- a/src/oauth/generic-account-failover.ts +++ b/src/oauth/generic-account-failover.ts @@ -16,7 +16,7 @@ */ import { getAccountSet } from "./store"; import { getValidAccessSnapshotForAccount, type OAuthAccessSnapshot } from "./index"; -import { exhaustedCooldownMs, hasHeadroomEvidence, rankAccountsByHeadroom } from "./account-quota-rank"; +import { exhaustedCooldownMs, hasHeadroomEvidence, isAccountQuotaExhausted, rankAccountsByHeadroom } from "./account-quota-rank"; import { parseRetryAfterMs } from "../combos/failover"; import { sweepExpiredOnWrite } from "../lib/state-store-sweeper"; import type { OcxConfig, OcxProviderConfig } from "../types"; @@ -61,29 +61,12 @@ interface PresenceEntry { readAt: number; } -/** - * Ordered roster plus the active id, for the pre-dispatch preference. - * - * Same reasoning as the presence cache: `getAccountSet` reads through `loadAuthStore`, - * which chmods and re-parses the whole credential file on every call. Selection needs the - * ORDER and the active id, which the presence count cannot supply, so it gets its own - * TTL-bounded row. Ids and an active pointer only — never a credential. - */ -interface RosterEntry { - ids: string[]; - activeId: string | null; - readAt: number; -} - /** Process-local, like the Anthropic pool's: a restart is allowed to forget a cooldown. */ const health = new Map(); /** Provider -> recent eligible-account count. TTL-bounded; never holds credential material. */ const presence = new Map(); -/** Provider -> recently read roster. TTL-bounded; never holds credential material. */ -const roster = new Map(); - const healthKey = (provider: string, accountId: string) => `${provider}\u0000${accountId}`; function isCooled(provider: string, accountId: string, now: number): boolean { @@ -118,24 +101,6 @@ function eligibleAccountCount(providerName: string, now: number): number { return eligible; } -/** - * Roster ids and the active pointer, read at most once per TTL window. - * - * `needsReauth` accounts are excluded for the same reason the presence count excludes - * them: a revoked credential cannot serve the request we are about to send. - */ -function cachedRoster(providerName: string, now: number): { ids: string[]; activeId: string | null } { - const cached = roster.get(providerName); - if (cached && now >= cached.readAt && now - cached.readAt < PRESENCE_CACHE_TTL_MS) { - return { ids: cached.ids, activeId: cached.activeId }; - } - const set = getAccountSet(providerName); - const ids = set ? set.accounts.filter(a => a.needsReauth !== true).map(a => a.id) : []; - const activeId = set?.activeAccountId ?? null; - roster.set(providerName, { ids, activeId, readAt: now }); - return { ids, activeId }; -} - /** * Presence IS consent (#2568d). * @@ -192,8 +157,7 @@ function isProactivePreferenceEnabled(config: OcxConfig, providerName: string, n if (typeof perProvider === "boolean") { return perProvider && hasFailoverAccountQuorum(providerName, now); } - if (config.oauthAccountFailover?.enabled === false) return false; - return hasFailoverAccountQuorum(providerName, now); + return config.oauthAccountFailover?.enabled === true && hasFailoverAccountQuorum(providerName, now); } /** Accounts that may serve traffic right now: not cooled, not flagged for reauth. */ @@ -240,8 +204,6 @@ export function rotateGenericOAuthAccountOn429( // A rotation means the roster in use just changed; do not answer the next activation question // from a count read before the failure. presence.delete(providerName); - // Same for the selection roster: the next request must not pick from a pre-failure read. - roster.delete(providerName); // Deterministic: start after the failed account so repeated 429s walk the roster instead of // 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. @@ -289,13 +251,19 @@ export function preferredInitialAccount( // The PROACTIVE predicate, not the reactive one: this steers a request upstream has not // refused, so `oauthAccountFailover.enabled: false` must still be able to refuse it. if (!isProactivePreferenceEnabled(config, providerName, now)) return null; - // This runs on the initial resolution of EVERY request, and `loadAuthStore` has no - // cache: each call chmods the config dir, chmods the secret, reads the whole file and - // normalizes it (store.ts:136-151). So the store is consulted at most ONCE here, behind - // the same TTL the presence check uses, and never at all for a single-account provider. - const { ids: order, activeId: active } = cachedRoster(providerName, now); + // Read the same authoritative selection the management writer commits. Caching the + // active id separately would delay manual selection and account removal. + const selected = getAccountSet(providerName); + if (!selected) return null; + const active = selected.activeAccountId; + const order = selected.accounts.filter(account => account.needsReauth !== true).map(account => account.id); if (order.length < 2) return null; + const activeRow = selected.accounts.find(account => account.id === active); + if (activeRow && activeRow.needsReauth !== true + && !isCooled(providerName, activeRow.id, now) + && !isAccountQuotaExhausted(providerName, activeRow.id)) return null; + // Evidence is required BEFORE eligibility narrows the field. Without this, a provider // with no quota data at all could still be redirected: cool the active account with a // 429 and the eligible list collapses to one candidate, which any ranking returns @@ -318,11 +286,8 @@ export function preferredInitialAccount( const best = rankAccountsByHeadroom(providerName, candidates)[0] ?? null; // Nothing to do when the ranking agrees with the account we would have used anyway. // - // The roster may be up to PRESENCE_CACHE_TTL_MS old, so this answer is a PREFERENCE the - // caller must be able to abandon: it resolves the account with `requireUsableAccount`, - // which rejects a removed or reauth-flagged account inside the store read it was already - // performing, and falls back to the active account. Validating here instead would mean a - // second uncached read of the credential file on every redirected request. + // A proposal still needs guarded selection commit after credential resolution: a + // removal, reauth verdict, or manual choice can arrive during that await. return best && best !== active ? best : null; } @@ -341,7 +306,6 @@ export function genericFailoverRetryAfterSeconds(providerName: string, now = Dat /** Test seam and manual-recovery hook. */ export function forgetGenericFailoverRoster(providerName: string): void { - roster.delete(providerName); presence.delete(providerName); } @@ -350,11 +314,9 @@ export function clearGenericFailoverHealth(providerName?: string): void { if (!providerName) { health.clear(); presence.clear(); - roster.clear(); return; } presence.delete(providerName); - roster.delete(providerName); for (const key of [...health.keys()]) { if (key.startsWith(`${providerName}\u0000`)) health.delete(key); } diff --git a/src/oauth/pool-settings-capability.ts b/src/oauth/pool-settings-capability.ts index 82f8c06863..210870b718 100644 --- a/src/oauth/pool-settings-capability.ts +++ b/src/oauth/pool-settings-capability.ts @@ -10,7 +10,8 @@ import type { OcxProviderConfig } from "../types"; * * `strategy` and `autoSwitchThreshold` are still a declared contract the selector does not * consume — that is what `inert` reports. `enabled` is NOT inert any more: an explicit - * `false` refuses the pre-dispatch account preference (`preferredInitialAccount`). What it can + * `true` enables pre-dispatch exhaustion avoidance (`preferredInitialAccount`); absence is off. + * Healthy manual selections remain authoritative. What the switch can * no longer do is refuse reactive 429 rotation, which activates on account presence and is not * disableable. */ diff --git a/src/oauth/store.ts b/src/oauth/store.ts index bc9e0e11c1..e641b81be7 100644 --- a/src/oauth/store.ts +++ b/src/oauth/store.ts @@ -23,12 +23,13 @@ import { getConfigDir, atomicWriteFile, backupInvalidConfig, hardenConfigDir, ha import { assertNotRealHomeUnderTest } from "../lib/test-home-guard"; import { recordOwnedConfigPath } from "../lib/config-ownership"; import { MAX_PENDING_OAUTH_MUTATIONS } from "../lib/translator-budget"; +import { publishAccountSelection } from "../lib/account-selection-events"; import { captureConfigGeneration, type GenerationContext, } from "../lib/state-store-sweeper"; import { validateCopilotApiBaseUrl } from "./github-copilot"; -import type { OAuthCredentialSource, OAuthCredentials, ProviderAccount, ProviderAccountSet } from "./types"; +import type { OAuthAccountSelection, OAuthCredentialSource, OAuthCredentials, ProviderAccount, ProviderAccountSet } from "./types"; export type AuthStore = Record; @@ -536,7 +537,15 @@ function normalizeAccountSet(raw: unknown): { set: ProviderAccountSet | null; wa const active = typeof candidate.activeAccountId === "string" && accounts.some(a => a.id === candidate.activeAccountId) ? candidate.activeAccountId : accounts[0]!.id; - return { set: { activeAccountId: active, accounts }, wasLegacy: false }; + const set: ProviderAccountSet = { activeAccountId: active, accounts }; + // Healing a dangling active id invalidates its old selection generation. Reads stay + // deterministic; only the serialized writer creates new revisions. + if (active === candidate.activeAccountId + && typeof candidate.selectionRevision === "string" + && /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i.test(candidate.selectionRevision)) { + set.selectionRevision = candidate.selectionRevision; + } + return { set, wasLegacy: false }; } // Legacy single-credential value. const cred = normalizeCredential(raw); @@ -670,9 +679,35 @@ function serializeMutation(work: () => Promise, retainedValues: readonly u export function mutateStore(fn:(store:AuthStore)=>T|Promise, retainedValues: readonly unknown[] = [], options?: { waitMs?: number; assertBeforePersist?: () => void }):Promise{return serializeMutation(async()=>{const guard=await createOAuthFileLock({path:getAuthStoreLockPath(),staleAfterMs:30000}).acquire();try{ const { store, hadLegacy } = loadAuthStoreInternal(); if (hadLegacy) backupLegacyOnce(); + const selections = new Map(Object.entries(store).map(([provider, set]) => [provider, { + set, + accountId: set.activeAccountId, + revision: set.selectionRevision, + accountIds: set.accounts.map(account => account.id), + }])); const result = await fn(store); options?.assertBeforePersist?.(); + const changedProviders: string[] = []; + for (const provider of new Set([...selections.keys(), ...Object.keys(store)])) { + const before = selections.get(provider); + const after = store[provider]; + if (!after) { + if (before) changedProviders.push(provider); + continue; + } + const replaced = before?.set !== after; + const removed = before?.accountIds.some(id => !after.accounts.some(account => account.id === id)); + if (replaced || removed || before?.accountId !== after.activeAccountId) { + // Replacement/rollback must never restore a previous generation. A common + // selection commit already assigned its revision before forming its result. + if (replaced || before?.revision === after.selectionRevision) after.selectionRevision = randomUUID(); + } + if (!before || before.accountId !== after.activeAccountId || before.revision !== after.selectionRevision) { + changedProviders.push(provider); + } + } persist(store); + for (const provider of changedProviders) publishAccountSelection(provider, "oauth"); return result; }finally{guard.release();}}, retainedValues, options?.waitMs); } @@ -700,6 +735,8 @@ export async function saveCredential( if (!safe) return; await mutateStore(store => { const set = store[provider]; + // Login explicitly selects an account, including a re-login to the same slot. + if (set) set.selectionRevision = randomUUID(); const identity = safe.accountId ?? safe.email; if (!set || SINGLE_SLOT_PROVIDERS.has(provider)) { const id = newAccountId(safe); @@ -876,13 +913,60 @@ export async function saveAccountCredential( }, [provider, accountId, safe], { assertBeforePersist: opts.assertBeforePersist }); } -export async function setActiveAccount(provider: string, accountId: string): Promise { +function accountSelection(set: ProviderAccountSet): OAuthAccountSelection { + return { + accountId: set.activeAccountId, + ...(set.selectionRevision !== undefined ? { revision: set.selectionRevision } : {}), + }; +} + +export function captureOAuthAccountSelection(provider: string): OAuthAccountSelection | null { + const set = getAccountSet(provider); + return set ? accountSelection(set) : null; +} + +/** + * Shared manual/automatic selection owner. An expected snapshot marks an automatic + * proposal: validating an unchanged selection preserves its revision. An unconditional + * (manual) selection always advances it, even when reselecting the current account. + */ +export async function commitOAuthAccountSelection( + provider: string, + accountId: string, + options: { + expectedSelection?: OAuthAccountSelection; + expectedCredentialGeneration?: string; + requireUsableAccount?: boolean; + } = {}, +): Promise { + // Snapshot caller-owned options before waiting for the serialized writer. + const expected = options.expectedSelection ? { ...options.expectedSelection } : undefined; + const { expectedCredentialGeneration, requireUsableAccount } = options; + const valid = (set: ProviderAccountSet): boolean => { + if (expected && (set.activeAccountId !== expected.accountId || set.selectionRevision !== expected.revision)) return false; + const account = set.accounts.find(account => account.id === accountId); + if (!account || (requireUsableAccount && account.needsReauth === true)) return false; + return expectedCredentialGeneration === undefined || credentialGeneration(account.credential) === expectedCredentialGeneration; + }; + if (expected?.accountId === accountId) { + // Ordinary admission is one synchronous read/validation, with no await, writer + // queue, or persistence. A changed selection must still take the guarded writer. + const set = getAccountSet(provider); + return set && valid(set) ? accountSelection(set) : null; + } return await mutateStore(store => { const set = store[provider]; - if (!set || !set.accounts.some(a => a.id === accountId)) return false; - set.activeAccountId = accountId; - return true; - }, [provider, accountId]); + if (!set || !valid(set)) return null; + if (!expected || set.activeAccountId !== accountId) { + set.activeAccountId = accountId; + set.selectionRevision = randomUUID(); + } + return accountSelection(set); + }, [provider, accountId, expected, expectedCredentialGeneration]); +} + +export async function setActiveAccount(provider: string, accountId: string): Promise { + return (await commitOAuthAccountSelection(provider, accountId)) !== null; } export async function setAccountAlias(provider: string, accountId: string, alias: string | undefined): Promise { diff --git a/src/oauth/types.ts b/src/oauth/types.ts index 5c2dac541b..19cf435d54 100644 --- a/src/oauth/types.ts +++ b/src/oauth/types.ts @@ -59,9 +59,17 @@ export interface ProviderAccount { /** auth.json value per provider: N accounts + which one requests use. */ export interface ProviderAccountSet { activeAccountId: string; + /** Opaque selection generation; absent in legacy stores, independent of token refresh. */ + selectionRevision?: string; accounts: ProviderAccount[]; } +/** Non-secret snapshot used to condition a selection on the choice that started a request. */ +export interface OAuthAccountSelection { + accountId: string; + revision?: string; +} + export interface OAuthController { onAuth?(info: { url: string; instructions?: string; deviceCode?: string }): void; onProgress?(message: string): void; diff --git a/src/providers/api-key-selection.ts b/src/providers/api-key-selection.ts new file mode 100644 index 0000000000..8cf14cfcb3 --- /dev/null +++ b/src/providers/api-key-selection.ts @@ -0,0 +1,110 @@ +import { randomUUID } from "node:crypto"; +import { mutatePersistedConfig } from "../config"; +import { publishAccountSelection } from "../lib/account-selection-events"; +import type { OcxConfig, OcxProviderConfig } from "../types"; +import type { ProviderApiKeySelection } from "../types/provider"; +import { routedProviderConfig } from "../router"; +import { OPENCODE_GO_SESSION_HEADER } from "./opencode-go-transport"; +import { resolveProviderTransport, XAI_GROK_COMPATIBILITY, type OcxProviderTransport } from "./xai-transport"; + +export function captureProviderApiKeySelection(provider: OcxProviderConfig): ProviderApiKeySelection { + return { + entryId: provider.apiKeyPool?.find(entry => entry.key === provider.apiKey)?.id, + reference: provider.apiKey, + revision: provider.apiKeySelectionRevision, + }; +} + +function matchesSelection(provider: OcxProviderConfig, expected: ProviderApiKeySelection): boolean { + const current = captureProviderApiKeySelection(provider); + return current.entryId === expected.entryId && current.reference === expected.reference + && current.revision === expected.revision; +} + +function currentKeyProvider(config: OcxConfig, name: string): OcxProviderConfig | null { + const configured = config.providers[name]; + if (!configured || configured.disabled) return null; + const current = routedProviderConfig(name, { ...configured, _apiKeyAttempt: undefined }); + if (current.authMode === "oauth" || current.authMode === "forward") return null; + if (current.authMode === "key" && !current.keyOptional && !current.apiKey?.trim()) return null; + return current; +} + +/** Physical-send check; stored references alone do not detect a changed env/keychain value. */ +export function providerApiKeySelectionIsCurrent( + config: OcxConfig, + name: string, + routedProvider: OcxProviderConfig, +): boolean { + const current = currentKeyProvider(config, name); + const expected = routedProvider._apiKeyAttempt; + return current !== null && expected !== undefined + && matchesSelection(config.providers[name]!, expected) + && current.apiKey === routedProvider.apiKey + && current.authMode === routedProvider.authMode + && current.baseUrl === routedProvider.baseUrl; +} + +/** Rebuild transport from the already committed choice; never allocate or publish a selection. */ +export function resolveCurrentProviderApiKeyTransport( + config: OcxConfig, + name: string, + routedProvider: OcxProviderConfig, +): OcxProviderConfig | null { + const current = currentKeyProvider(config, name); + if (!current) return null; + const runtime = routedProvider as OcxProviderTransport; + const headers = { ...current.headers }; + const affinityHeaders = name === "xai" + ? [XAI_GROK_COMPATIBILITY.headers.conversationId, XAI_GROK_COMPATIBILITY.headers.sessionId] + : [OPENCODE_GO_SESSION_HEADER]; + for (const header of affinityHeaders) { + const configured = Object.keys(headers).some(key => key.toLowerCase() === header.toLowerCase()); + const value = Object.entries(runtime.headers ?? {}).find(([key]) => key.toLowerCase() === header.toLowerCase())?.[1]; + if (!configured && value !== undefined) headers[header] = value; + } + const fetch = (current as OcxProviderTransport).fetch ?? runtime.fetch; + return resolveProviderTransport(name, { + ...current, + ...(Object.keys(headers).length ? { headers } : {}), + ...(fetch ? { fetch } : {}), + }); +} + +type SelectionMutation = { changed: boolean; value: T; selectionChanged?: boolean }; +export type ProviderApiKeyCommit = + | { status: "committed"; provider: OcxProviderConfig; value: T } + | { status: "superseded"; provider: OcxProviderConfig } + | { status: "unavailable" }; + +/** GUI and recovery share one persisted selection transaction and post-commit notification. */ +export function commitProviderApiKeySelection( + config: OcxConfig, + name: string, + mutation: (provider: OcxProviderConfig) => SelectionMutation, + expectedSelection?: ProviderApiKeySelection, +): ProviderApiKeyCommit { + const outcome = mutatePersistedConfig & { notify?: boolean }>(fresh => { + const provider = fresh.providers[name]; + if (!provider || provider.authMode === "oauth" || provider.authMode === "forward") { + return { changed: false, value: { status: "unavailable" } }; + } + if (expectedSelection && !matchesSelection(provider, expectedSelection)) { + return { changed: false, value: { status: "superseded", provider: structuredClone(provider) } }; + } + const before = provider.apiKey; + const result = mutation(provider); + const notify = result.selectionChanged === true || before !== provider.apiKey; + if (notify) provider.apiKeySelectionRevision = randomUUID(); + delete provider._apiKeyAttempt; + return { + changed: result.changed || notify, + value: { status: "committed", provider: structuredClone(provider), value: result.value, notify }, + }; + }); + if (outcome.status === "unavailable") return { status: "unavailable" }; + const committed = outcome.value; + if (committed.status !== "unavailable") config.providers[name] = structuredClone(committed.provider); + if (committed.status === "committed" && committed.notify) publishAccountSelection(name, "api-key"); + return committed; +} diff --git a/src/providers/api-keys.ts b/src/providers/api-keys.ts index d138cdbf50..ceb5522df2 100644 --- a/src/providers/api-keys.ts +++ b/src/providers/api-keys.ts @@ -10,6 +10,7 @@ import { createHash } from "node:crypto"; import { saveConfigPreservingClaudeCode } from "../config"; import type { OcxConfig, OcxProviderConfig } from "../types"; import type { AccountQuotaFields } from "./quota-types"; +import { commitProviderApiKeySelection } from "./api-key-selection"; export interface ProviderApiKeyInfo extends AccountQuotaFields { id: string; @@ -85,28 +86,33 @@ export function addProviderApiKey(config: OcxConfig, name: string, key: string, if (typeof key !== "string" || !key.trim()) return { error: "key is required" }; const trimmed = sanitizeApiKeyValue(key); if (!trimmed) return { error: "key must not include line breaks" }; - const pool = ensurePool(provider); const id = apiKeyPoolEntryId(trimmed); - const existing = pool.find(e => e.id === id); - if (existing) { - if (label?.trim()) existing.label = label.trim(); - } else { - pool.push({ id, key: trimmed, ...(label?.trim() ? { label: label.trim() } : {}), addedAt: Date.now() }); - } - provider.apiKey = trimmed; - saveConfigPreservingClaudeCode(config); - return { id }; + const committed = commitProviderApiKeySelection(config, name, fresh => { + const pool = ensurePool(fresh); + const existing = pool.find(e => e.id === id); + if (existing) { + if (label?.trim()) existing.label = label.trim(); + } else { + pool.push({ id, key: trimmed, ...(label?.trim() ? { label: label.trim() } : {}), addedAt: Date.now() }); + } + fresh.apiKey = trimmed; + return { changed: true, selectionChanged: true, value: id }; + }); + return committed.status === "committed" ? { id } : { error: "provider selection unavailable" }; } /** Switch the ACTIVE key (mirrors into `provider.apiKey`). Persists config. */ export function setActiveProviderApiKey(config: OcxConfig, name: string, id: string): boolean { - const provider = config.providers[name]; - if (!provider || !isKeyAuthProvider(provider)) return false; - const entry = ensurePool(provider).find(e => e.id === id); - if (!entry) return false; - provider.apiKey = entry.key; - saveConfigPreservingClaudeCode(config); - return true; + const committed = commitProviderApiKeySelection(config, name, provider => { + const entry = provider.apiKeyPool?.find(e => e.id === id) + ?? (!provider.apiKeyPool?.length && provider.apiKey && apiKeyPoolEntryId(provider.apiKey) === id + ? { id, key: provider.apiKey } : undefined); + if (!entry) return { changed: false, value: false }; + ensurePool(provider); + provider.apiKey = entry.key; + return { changed: true, selectionChanged: true, value: true }; + }); + return committed.status === "committed" && committed.value; } /** Rename a key slot without changing its id, secret, or active routing state. */ @@ -123,18 +129,19 @@ export function setProviderApiKeyLabel(config: OcxConfig, name: string, id: stri /** Remove one key; removing the active one promotes the first remaining. Persists config. */ export function removeProviderApiKey(config: OcxConfig, name: string, id: string): boolean { - const provider = config.providers[name]; - if (!provider || !isKeyAuthProvider(provider)) return false; - const pool = ensurePool(provider); - const entry = pool.find(e => e.id === id); - if (!entry) return false; - provider.apiKeyPool = pool.filter(e => e.id !== id); - if (provider.apiKey === entry.key) { - const next = provider.apiKeyPool[0]; - if (next) provider.apiKey = next.key; - else delete provider.apiKey; - } - if (provider.apiKeyPool.length === 0) delete provider.apiKeyPool; - saveConfigPreservingClaudeCode(config); - return true; + const committed = commitProviderApiKeySelection(config, name, provider => { + const pool = provider.apiKeyPool?.length ? provider.apiKeyPool + : provider.apiKey ? [{ id: apiKeyPoolEntryId(provider.apiKey), key: provider.apiKey }] : []; + const entry = pool.find(e => e.id === id); + if (!entry) return { changed: false, value: false }; + provider.apiKeyPool = pool.filter(e => e.id !== id); + if (provider.apiKey === entry.key) { + const next = provider.apiKeyPool[0]; + if (next) provider.apiKey = next.key; + else delete provider.apiKey; + } + if (provider.apiKeyPool.length === 0) delete provider.apiKeyPool; + return { changed: true, value: true }; + }); + return committed.status === "committed" && committed.value; } diff --git a/src/providers/key-failover.ts b/src/providers/key-failover.ts index 9c7a42e11b..ad4d61ba8d 100644 --- a/src/providers/key-failover.ts +++ b/src/providers/key-failover.ts @@ -8,7 +8,8 @@ * * Modelled after src/codex/routing.ts cooldown logic but scoped to plain API-key pools. */ -import { mutatePersistedConfig } from "../config"; +import { commitProviderApiKeySelection } from "./api-key-selection"; +import type { ProviderApiKeySelection } from "../types/provider"; import { routedProviderConfig } from "../router"; import type { OcxConfig, OcxProviderConfig, RateLimitRetryPolicy, TransientRetryPolicy } from "../types"; import { OPENCODE_GO_SESSION_HEADER } from "./opencode-go-transport"; @@ -186,33 +187,32 @@ function rotateKeyAfterFailure( retryAfterHeader: string | null | undefined, now = Date.now(), attemptedKey?: string, + attemptedSelection?: ProviderApiKeySelection, ): OcxProviderConfig | null { const provider = config.providers[providerName]; if (!provider) return null; if (provider.authMode === "oauth" || provider.authMode === "forward") return null; - const failedKey = attemptedKey ?? provider.apiKey; + const failedKey = attemptedSelection?.reference ?? attemptedKey ?? provider.apiKey; type Rotation = - | { provider: OcxProviderConfig; failedId?: string; candidateId?: string } + | { failedId?: string; candidateId?: string } | { exhaustedCount: number; failedId?: string }; - const outcome = mutatePersistedConfig(fresh => { - const freshProvider = fresh.providers[providerName]; - if (!freshProvider || freshProvider.authMode === "oauth" || freshProvider.authMode === "forward") { - return { changed: false, value: null }; - } + const outcome = commitProviderApiKeySelection(config, providerName, freshProvider => { const pool = freshProvider.apiKeyPool; if (!pool || pool.length < 2) return { changed: false, value: null }; // The callback can be rerun after rebasing, so identify the failed key here but // defer the in-memory cooldown side effect until persistence has succeeded. - const failedEntry = pool.find(entry => entry.key === failedKey); + const failedEntry = attemptedSelection?.entryId + ? pool.find(entry => entry.id === attemptedSelection.entryId && entry.key === failedKey) + : pool.find(entry => entry.key === failedKey); if (freshProvider.apiKey !== failedKey) { const activeEntry = pool.find(entry => entry.key === freshProvider.apiKey); if (activeEntry && !isKeyInCooldown(providerName, activeEntry.id, now)) { return { changed: false, - value: { provider: structuredClone(freshProvider), failedId: failedEntry?.id }, + value: { failedId: failedEntry?.id }, }; } } @@ -226,15 +226,20 @@ function rotateKeyAfterFailure( return { changed: true, value: { - provider: structuredClone(freshProvider), failedId: failedEntry?.id, candidateId: candidate.id, }, }; } return { changed: false, value: { exhaustedCount: pool.length, failedId: failedEntry?.id } }; - }); - if (outcome.status === "unavailable" || outcome.value === null) return null; + }, attemptedSelection); + if (outcome.status === "unavailable") return null; + if (outcome.status === "superseded") { + // A newer manual selection (including A→B→A) owns subsequent dispatch. Reusing the + // same failed key here would loop forever; preserve its original failure instead. + return outcome.provider.apiKey !== failedKey ? structuredClone(outcome.provider) : null; + } + if (outcome.value === null) return null; if (outcome.value.failedId) { // A 401 is a verdict about the credential itself, not a timing signal: the key is rejected // until an operator replaces it, and upstreams send no Retry-After for it. Hold it for the @@ -250,7 +255,7 @@ function rotateKeyAfterFailure( return null; } - const committed = structuredClone(outcome.value.provider); + const committed = structuredClone(outcome.provider); config.providers[providerName] = committed; if (outcome.value.candidateId) { console.warn( @@ -267,8 +272,9 @@ export function rotateKeyOn429( retryAfterHeader: string | null | undefined, now = Date.now(), attemptedKey?: string, + attemptedSelection?: ProviderApiKeySelection, ): OcxProviderConfig | null { - return rotateKeyAfterFailure(config, providerName, 429, retryAfterHeader, now, attemptedKey); + return rotateKeyAfterFailure(config, providerName, 429, retryAfterHeader, now, attemptedKey, attemptedSelection); } /** @@ -284,8 +290,9 @@ export function rotateKeyOn401( providerName: string, now = Date.now(), attemptedKey?: string, + attemptedSelection?: ProviderApiKeySelection, ): OcxProviderConfig | null { - return rotateKeyAfterFailure(config, providerName, 401, null, now, attemptedKey); + return rotateKeyAfterFailure(config, providerName, 401, null, now, attemptedKey, attemptedSelection); } export function sweepExpiredApiKeyCooldowns(now = Date.now()): number { @@ -302,6 +309,7 @@ interface RotateProviderTransportOptions { retryAfter?: string | null; now?: number; attemptedKey?: string; + attemptedSelection?: ProviderApiKeySelection; promptCacheKey?: string; } @@ -323,6 +331,7 @@ export function rotateProviderTransportOn429( options.retryAfter, options.now, options.attemptedKey, + options.attemptedSelection ?? routedProvider._apiKeyAttempt, ); if (!rotated) return null; return applyRotatedTransport(providerName, routedProvider, rotated, options.promptCacheKey); @@ -335,7 +344,8 @@ export function rotateProviderTransportOn401( routedProvider: OcxProviderTransport, options: Omit = {}, ): OcxProviderTransport | null { - const rotated = rotateKeyOn401(config, providerName, options.now, options.attemptedKey); + const rotated = rotateKeyOn401(config, providerName, options.now, options.attemptedKey, + options.attemptedSelection ?? routedProvider._apiKeyAttempt); if (!rotated) return null; return applyRotatedTransport(providerName, routedProvider, rotated, options.promptCacheKey); } diff --git a/src/router.ts b/src/router.ts index 758f34e751..b2f887f0a9 100644 --- a/src/router.ts +++ b/src/router.ts @@ -10,6 +10,7 @@ import { import type { NormalizedComboConfig } from "./combos/types"; import { hasOwnProvider } from "./config/provider-name"; import { providerUsesKeyAuthOverride, resolveProviderApiKey } from "./providers/key-store"; +import { captureProviderApiKeySelection } from "./providers/api-key-selection"; import { assertProviderDestinationAllowed } from "./lib/destination-policy"; import { redactSecretString, redactUrlForLog } from "./lib/redact"; import { @@ -297,6 +298,7 @@ function usableResolvedApiKey(apiKey: string | undefined): string | undefined { } export function routedProviderConfig(providerName: string, provider: OcxProviderConfig): OcxProviderConfig { + provider = { ...provider, _apiKeyAttempt: provider._apiKeyAttempt ?? captureProviderApiKeySelection(provider) }; const registryEntry = PROVIDER_REGISTRY.find(entry => entry.id === providerName); if (!registryEntry || !providerMatchesRegistryTransportWithStaticGuards(providerName, provider)) { assertProviderDestinationAllowed(providerName, provider); diff --git a/src/server/auth-cors.ts b/src/server/auth-cors.ts index 0677722979..0dd49910fb 100644 --- a/src/server/auth-cors.ts +++ b/src/server/auth-cors.ts @@ -530,6 +530,7 @@ export function requireResponsesApiAuth(req: Request, config: RequestPolicyView) const FORBIDDEN_PROVIDER_RUNTIME_FIELDS = [ "virtualModels", "codexAuthContext", "selectedForwardHeaders", "sidecarOutcomeRecorder", "_codexAccountOverride", "_codexAccountRequired", + "_apiKeyAttempt", ] as const; function sameCanonicalProviderSeed(actual: Record, expected: OcxProviderConfig): boolean { @@ -795,6 +796,8 @@ const PROVIDER_CONFIG_FIELD_POLICY = { apiKey: "redacted", apiKeyTransport: "editor", apiKeyPool: "redacted", + apiKeySelectionRevision: "runtime", + _apiKeyAttempt: "runtime", defaultModel: "editor", models: "editor", liveModels: "editor", diff --git a/src/server/chat-native.ts b/src/server/chat-native.ts index 513a9d0913..30ff39ff44 100644 --- a/src/server/chat-native.ts +++ b/src/server/chat-native.ts @@ -37,6 +37,8 @@ import { transientRetryPolicyFor, } from "../providers/key-failover"; import { fastPolicyForModel } from "../providers/service-tier"; +import { providerApiKeySelectionIsCurrent, resolveCurrentProviderApiKeyTransport } from "../providers/api-key-selection"; +import type { OcxProviderTransport } from "../providers/xai-transport"; import type { RouteResult } from "../router"; import type { OcxConfig, OcxProviderConfig } from "../types"; import { fetchWithHeaderTimeout, providerFetch, safeHostLabel } from "./responses/fetch-helpers"; @@ -228,7 +230,6 @@ export async function handleNativeChatCompletions(options: HandleNativeChatOptio const fetchWithPolicy = requestTransientPolicy ? fetchWithTransientRetry : fetchWithResetRetry; return await fetchWithPolicy( (transportRecovery?: UpstreamSendRecovery) => { - noteAttemptSend(attempt, logCtx.usageLogInputTokens, transportRecovery ?? recovery); return fetchWithHeaderTimeout( request.url, applyUpstreamRecoveryInit({ @@ -242,6 +243,32 @@ export async function handleNativeChatCompletions(options: HandleNativeChatOptio providerFetch(activeProvider, undefined, { providerName: route.providerName, modelId: route.modelId, + dispatchOverride: async (_input, init, execute) => { + if (!providerApiKeySelectionIsCurrent(config, route.providerName, activeProvider)) { + const current = resolveCurrentProviderApiKeyTransport(config, route.providerName, activeProvider); + if (!current || !isNativeChatRouteEligible({ ...route, provider: current }, options.chatBody)) { + throw new Error("Provider key selection is no longer available for native Chat"); + } + activeProvider = current; + activeAdapter = createOpenAIChatAdapter(current); + activeRequest.releaseBodyObservation?.(); + releaseRetainedRequest(); + activeRequest = buildActiveRequest(); + try { retainRequest(activeRequest); } + catch (error) { activeRequest.releaseBodyObservation?.(); throw error; } + } + // The retry closure may still hold a pre-pacing request. Replace its entire + // wire shape, not just Authorization, and retain transport recovery flags. + request = activeRequest; + const headers = new Headers(request.headers); + const encoding = new Headers(init.headers).get("accept-encoding"); + if (!headers.has("accept-encoding") && encoding) headers.set("accept-encoding", encoding); + if (init.signal?.aborted) throw init.signal.reason; + noteAttemptSend(attempt, logCtx.usageLogInputTokens, transportRecovery ?? recovery); + return ((activeProvider as OcxProviderTransport).fetch ?? execute)(request.url, applyUpstreamRecoveryInit({ + ...init, method: request.method, headers, body: request.body, + }, transportRecovery)); + }, }), ); }, @@ -286,6 +313,7 @@ export async function handleNativeChatCompletions(options: HandleNativeChatOptio retryAfter: response.headers.get("retry-after"), now: Date.now(), attemptedKey: activeProvider.apiKey, + attemptedSelection: activeProvider._apiKeyAttempt, promptCacheKey: typeof options.chatBody.prompt_cache_key === "string" ? options.chatBody.prompt_cache_key : undefined, }); if (!rotated) break; @@ -306,6 +334,9 @@ export async function handleNativeChatCompletions(options: HandleNativeChatOptio cleanupAbort(); upstream.abort(); if (req.signal.aborted) return fail(499, "Client cancelled request", "client_cancelled"); + if (isTranslatorBudgetExceededError(error)) { + return fail(413, "request translation buffer exceeded the safe limit", "request_too_large", "translation_buffer_limit"); + } return fail(502, error instanceof Error ? error.message : String(error), "server_error"); } releaseRetainedRequest(); diff --git a/src/server/management-auth.ts b/src/server/management-auth.ts index 09127b0b58..58ae3f7d60 100644 --- a/src/server/management-auth.ts +++ b/src/server/management-auth.ts @@ -251,10 +251,26 @@ export function issueGuiSession( export interface ManagementSessionControl { revokeCurrent(req: Request): boolean; + /** Revalidate a long-lived request against current authority, without cached admission or renewal. */ + isCurrent(req: Request, config: OcxConfig): boolean; } export function createManagementSessionControl(state: ManagementAuthState): ManagementSessionControl { return { + isCurrent(req: Request, config: OcxConfig): boolean { + if (!state.available) return false; + const credential = requestManagementCredential(req); + if (!credential) return false; + if (equalSecret(credential, state.token)) return true; + const session = state.sessions.get(credential); + if (!session) return false; + // Reuse the full origin/expiry/CSRF predicate against the current record, but + // isolate its sliding-expiry mutation: SSE heartbeats are not browser activity. + return authorizeGuiSessionRequest(req, config, { + sessions: new Map([[credential, { ...session }]]), + pairingGrants: state.pairingGrants, + }).ok; + }, revokeCurrent(req: Request): boolean { if (!state.available) return false; const credential = requestManagementCredential(req); diff --git a/src/server/management/account-selection-stream.ts b/src/server/management/account-selection-stream.ts new file mode 100644 index 0000000000..e04e8a3e14 --- /dev/null +++ b/src/server/management/account-selection-stream.ts @@ -0,0 +1,70 @@ +import { currentAccountSelectionRevision, subscribeAccountSelections } from "../../lib/account-selection-events"; +import { registerOptionalShutdownHook } from "../../lib/optional-shutdown-hooks"; + +const MAX_SELECTION_STREAMS = 64; +const HEARTBEAT_MS = 15_000; +const encoder = new TextEncoder(); +const connections = new Set<() => void>(); + +/** The management boundary admits the request; every frame revalidates current authority. */ +export function accountSelectionStream(request: Request, validate: () => boolean): Response { + const authorized = () => { + try { return validate() === true; } catch { return false; } + }; + if (!authorized()) return Response.json({ error: "Management session is no longer authorized" }, { status: 401 }); + if (connections.size >= MAX_SELECTION_STREAMS) { + return Response.json({ error: "Too many account selection streams" }, { + status: 429, headers: { "Retry-After": "15" }, + }); + } + let cleanup = () => {}; + const body = new ReadableStream({ + start(controller) { + let closed = false; + let unsubscribe = () => {}; + let heartbeat: ReturnType | undefined; + const close = () => { + if (closed) return; + closed = true; + unsubscribe(); + if (heartbeat) clearInterval(heartbeat); + request.signal.removeEventListener("abort", close); + connections.delete(close); + try { controller.close(); } catch { /* The consumer may already have cancelled. */ } + }; + cleanup = close; + const send = (frame: string) => { + if (closed) return; + if (!authorized()) { + // Error clears queued frames as well, so a revoked consumer cannot drain them. + try { controller.error(new DOMException("Management session is no longer authorized", "NotAllowedError")); } + finally { close(); } + return; + } + // Reconnection sends a ready event, so a slow reader can reconcile without an + // unbounded queue or silently dropping a provider's latest invalidation. + if (controller.desiredSize !== null && controller.desiredSize <= 0) { close(); return; } + try { controller.enqueue(encoder.encode(frame)); } catch { close(); } + }; + connections.add(close); + registerOptionalShutdownHook("account-selection-streams", () => { + for (const finish of [...connections]) finish(); + }); + if (request.signal.aborted) { close(); return; } + request.signal.addEventListener("abort", close, { once: true }); + unsubscribe = subscribeAccountSelections(event => { + send(`event: account-selection\ndata: ${JSON.stringify(event)}\n\n`); + }); + send(`event: ready\ndata: ${JSON.stringify({ revision: currentAccountSelectionRevision() })}\n\n`); + if (closed) return; + heartbeat = setInterval(() => send(": heartbeat\n\n"), HEARTBEAT_MS); + heartbeat.unref?.(); + }, + cancel() { cleanup(); }, + }, { highWaterMark: 16 }); + return new Response(body, { headers: { + "Content-Type": "text/event-stream; charset=utf-8", + "Cache-Control": "no-cache, no-transform", + "X-Accel-Buffering": "no", + } }); +} diff --git a/src/server/management/oauth-account-routes.ts b/src/server/management/oauth-account-routes.ts index 07884d5f4f..874e3e24fa 100644 --- a/src/server/management/oauth-account-routes.ts +++ b/src/server/management/oauth-account-routes.ts @@ -129,6 +129,11 @@ function validateKeyName( export async function handleOauthAccountRoutes(ctx: ManagementContext): Promise { const { req, url, config, deps, syncClaudeAgentDefsBestEffort } = ctx; + if (url.pathname === "/api/accounts/events" && req.method === "GET") { + const { accountSelectionStream } = await import("./account-selection-stream"); + return accountSelectionStream(req, () => ctx.sessionControl?.isCurrent(req, config) === true); + } + // Which providers support real OAuth login (drives the GUI's "Log in with …" buttons). if (url.pathname === "/api/oauth/providers" && req.method === "GET") { return jsonResponse({ providers: listOAuthProviders() }); @@ -320,6 +325,8 @@ export async function handleOauthAccountRoutes(ctx: ManagementContext): Promise< if (!body.accountId) return jsonResponse({ error: "missing accountId" }, 400); const { setActiveAccount } = await import("../../oauth/store"); if (!(await setActiveAccount(provider, body.accountId))) return jsonResponse({ error: "account not found" }, 404); + const { forgetGenericFailoverRoster } = await import("../../oauth/generic-account-failover"); + forgetGenericFailoverRoster(provider); if (provider === "anthropic") { const { resetAnthropicRoutingForManualSelection } = await import("../../oauth/anthropic-routing"); resetAnthropicRoutingForManualSelection(body.accountId); diff --git a/src/server/management/route-registry.ts b/src/server/management/route-registry.ts index 6bd0a95a52..421ead31e0 100644 --- a/src/server/management/route-registry.ts +++ b/src/server/management/route-registry.ts @@ -26,6 +26,8 @@ export type HttpMethod = "GET" | "POST" | "PUT" | "PATCH" | "DELETE" | "HEAD"; * can check, not a way to quiet the parity test. */ export type ExemptionReason = + /** Invalidates dashboard state; the CLI reads the underlying resource directly. */ + | "gui-invalidation" /** Requires a dashboard browser session. Includes the user-consent star boundary. */ | "session-only" /** Deliberately returns 405; there is nothing to drive. */ @@ -256,6 +258,7 @@ export const MANAGEMENT_ROUTES: readonly ManagementRoute[] = [ { method: "GET", path: "/api/key-providers", module: "server/management/oauth-account-routes", mutates: false }, { method: "GET", path: "/api/keys", module: "server/management/oauth-account-routes", mutates: false }, { method: "GET", path: "/api/oauth/accounts", module: "server/management/oauth-account-routes", mutates: false }, + { method: "GET", path: "/api/accounts/events", module: "server/management/oauth-account-routes", mutates: false, exempt: { reason: "gui-invalidation", why: "Dashboard selection invalidation stream; CLI account commands read the authoritative account/key resources directly rather than subscribing to browser refresh notifications." } }, { method: "GET", path: "/api/oauth/accounts/pool", module: "server/management/oauth-account-routes", mutates: false }, { method: "GET", path: "/api/oauth/providers", module: "server/management/oauth-account-routes", mutates: false }, { method: "GET", path: "/api/oauth/status", module: "server/management/oauth-account-routes", mutates: false }, diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index b48438deb5..debba5c707 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -118,18 +118,19 @@ import { type OAuthAccessSnapshot, UnsupportedOAuthProviderError, } from "../../oauth"; +import { captureOAuthAccountSelection, commitOAuthAccountSelection, credentialGeneration, getAccountCredentialWithStatus } from "../../oauth/store"; import { ANTHROPIC_POOL_MAX_FAILOVERS_PER_REQUEST, anthropicSessionKeyFromParts, - bindAnthropicSessionAffinity, + commitAnthropicSelectionRouting, formatAnthropicProviderForLog, - getAnthropicPoolAccessToken, + getAnthropicPoolAccessSnapshot, getAnthropicPoolRetryAfterSeconds, isAnthropicAccountPoolEnabled, hasAnthropicFailoverQuorum, - promoteAnthropicActiveAccount, resolveAnthropicAccountForSession, rotateAnthropicAccountOn429, + type AnthropicAccountSelectionReason, } from "../../oauth/anthropic-routing"; import { stampOAuthAccountLabel } from "../../providers/label"; import { @@ -241,6 +242,7 @@ import { type InboundWire, } from "../../providers/registry"; import type { AdapterRequest, ProviderAdapter } from "../../adapters/base"; +import { providerApiKeySelectionIsCurrent, resolveCurrentProviderApiKeyTransport } from "../../providers/api-key-selection"; import { hasKeyPoolFailover, rateLimitRetryDelayMs, @@ -358,7 +360,7 @@ import type { EffectiveSubagentRoster, SpawnAgentSurface } from "../../codex/cat import { buildToolBridgeMaps, collabSurface, injectDeveloperMessage, multiAgentGuidanceText } from "./collaboration"; import { mapCodexAuthContextErrorToResponse, nativeMainRefreshFailureResponse } from "./codex-auth-error"; import { hasUnreadableEncryptedAgentTask, looksLikeBackendCiphertext, sanitizeEncryptedContentInPlace } from "./encrypted-payload"; -import { fetchWithHeaderTimeout, providerFetch, safeHostLabel, safeOriginLabel, storedPoolReplayDispatchNotifier } from "./fetch-helpers"; +import { fetchWithHeaderTimeout, providerFetch, safeHostLabel, safeOriginLabel, storedPoolReplayDispatchNotifier, type ProviderFetchOptions } from "./fetch-helpers"; import { classifyTransportFailureKind, transportErrorCode } from "../../lib/upstream-reachability"; import { acquireUpstreamHostAdmission, @@ -3673,6 +3675,82 @@ async function handleResponsesInner( // the request actually used, so a concurrent rotation cannot cool an innocent replacement. let genericFailoverAccountId: string | null = null; let genericFailovers = 0; + let oauthSelection = route.provider.authMode === "oauth" + ? captureOAuthAccountSelection(route.providerName) : null; + let servingOAuthSnapshot: OAuthAccessSnapshot | undefined; + // These owners also serve early passthrough and sidecar sends. A dispatch-time + // rebuild must update every later builder, without entering a later block's TDZ. + let adapter: ProviderAdapter; + let activeAdapter: ProviderAdapter; + let runTurnAdapter: ProviderAdapter; + let sameTargetRequest: AdapterRequest | undefined; + let sameTargetParsed: OcxParsedRequest | undefined; + let sameTargetToken = 0; + let transportToken = 0; + let imageTierBias = 0; + const invalidateSameTargetRequest = (): void => { transportToken += 1; }; + type DispatchBinding = + | { kind: "oauth"; selection: NonNullable; snapshot: OAuthAccessSnapshot } + | { kind: "api-key"; provider: OcxProviderConfig }; + const requestBindings = new WeakMap(); + const adapterBindings = new WeakMap(); + const rawRunTurns = new WeakMap>(); + const commitResolvedOAuthSelection = async ( + candidate: OAuthAccessSnapshot, + proactive = false, + anthropicReason?: AnthropicAccountSelectionReason, + ): Promise => { + const maxSelectionAttempts = 3; + for (let attempt = 0; attempt < maxSelectionAttempts; attempt++) { + if (!oauthSelection) return null; + const proactiveEnabled = route.providerName === "anthropic" + ? isAnthropicAccountPoolEnabled(config) + : (config.providers[route.providerName]?.oauthAccountFailover?.enabled + ?? config.oauthAccountFailover?.enabled) === true; + if (proactive && candidate.accountId !== oauthSelection.accountId && !proactiveEnabled) { + oauthSelection = captureOAuthAccountSelection(route.providerName); + if (!oauthSelection) return null; + candidate = route.providerName === "anthropic" + ? await getAnthropicPoolAccessSnapshot(oauthSelection.accountId) + : await getValidAccessSnapshotForAccount(route.providerName, oauthSelection.accountId, { requireUsableAccount: true }); + } + const committed = await commitOAuthAccountSelection(route.providerName, candidate.accountId, { + expectedSelection: oauthSelection, + expectedCredentialGeneration: candidate.generation, + requireUsableAccount: true, + }); + if (committed) { + if (route.providerName === "anthropic" && !commitAnthropicSelectionRouting( + candidate.accountId, oauthSelection, committed, + { config, sessionKey: anthropicSessionKey, reason: anthropicReason, expectedCredentialGeneration: candidate.generation }, + )) return null; + oauthSelection = committed; + servingOAuthSnapshot = candidate; + forgetGenericFailoverRoster(route.providerName); + return candidate; + } + // A newer manual choice wins over this request's old proposal, including A→B→A. + // Resolve that choice, not the rejected candidate, before trying admission again. + oauthSelection = captureOAuthAccountSelection(route.providerName); + if (!oauthSelection) return null; + candidate = route.providerName === "anthropic" + ? await getAnthropicPoolAccessSnapshot(oauthSelection.accountId) + : await getValidAccessSnapshotForAccount(route.providerName, oauthSelection.accountId, { requireUsableAccount: true }); + if (route.provider.googleMode === "cloud-code-assist" && !candidate.projectId) return null; + } + return null; + }; + const refreshResolvedOAuthSelection = async (sent: OAuthAccessSnapshot): Promise => { + const current = captureOAuthAccountSelection(route.providerName); + const unchanged = current?.accountId === oauthSelection?.accountId + && current?.revision === oauthSelection?.revision; + const candidate = unchanged ? await forceRefreshOAuthAccessSnapshot(sent) : sent; + const admitted = await commitResolvedOAuthSelection(candidate); + if (!admitted) throw new Error("OAuth selection changed during credential recovery"); + genericFailoverAccountId = admitted.accountId; + stampOAuthAccountLabel(logCtx, route.providerName, route.provider, admitted.accountId); + return admitted; + }; /** * Config generation captured where the serving credential is RESOLVED, not where the * quota is written. A streaming turn is a long await, so a generation captured at write @@ -3701,11 +3779,14 @@ async function handleResponsesInner( * tolerates project discovery failing, so a stored account can legitimately have no project; * sending that account's bearer with the FAILED account's project is worse than not rotating. */ - const applyFailoverSnapshot = ( + const applyFailoverSnapshot = async ( snapshot: OAuthAccessSnapshot, retryParsed: OcxParsedRequest = parsed, - ): boolean => { + ): Promise => { if (route.provider.googleMode === "cloud-code-assist" && !snapshot.projectId) return false; + const committed = await commitResolvedOAuthSelection(snapshot); + if (!committed) return false; + snapshot = committed; let rotatedProvider: OcxProviderConfig = { ...route.provider, apiKey: snapshot.accessToken }; if (route.providerName === "github-copilot") { rotatedProvider = resolveProviderTransport( @@ -3729,8 +3810,149 @@ async function handleResponsesInner( // served it. All three rotation sites funnel through here, so this is the only re-stamp // needed -- and putting it anywhere else would let one of the three drift. stampOAuthAccountLabel(logCtx, route.providerName, route.provider, snapshot.accountId); + if (route.providerName === "anthropic") { + anthropicPoolAccountId = snapshot.accountId; + logCtx.provider = formatAnthropicProviderForLog("anthropic", snapshot.accountId, config); + } else { + genericFailoverAccountId = snapshot.accountId; + } + sentOAuthSnapshot = snapshot; + replayOAuthCredentialSnapshot = { accountId: snapshot.accountId, generation: snapshot.generation }; return true; }; + const selectionIsCurrent = (binding: DispatchBinding | undefined): boolean => { + if (route.provider.authMode === "forward") return true; + if (!binding) return false; + if (binding.kind === "api-key") return providerApiKeySelectionIsCurrent(config, route.providerName, binding.provider); + const selected = captureOAuthAccountSelection(route.providerName); + const row = getAccountCredentialWithStatus(route.providerName, binding.snapshot.accountId); + return selected?.accountId === binding.selection.accountId && selected?.revision === binding.selection.revision + && !!row && !row.needsReauth && row.credential.expires > Date.now() + && credentialGeneration(row.credential) === binding.snapshot.generation; + }; + const resolveSelectionAdapter = (provider: OcxProviderConfig, retention = config.cacheRetention): ProviderAdapter => { + const resolved = resolveAdapter(provider, retention); + if (route.provider.authMode === "forward") return resolved; + const binding: DispatchBinding | undefined = route.provider.authMode === "oauth" + ? oauthSelection && servingOAuthSnapshot + ? { kind: "oauth", selection: { ...oauthSelection }, snapshot: servingOAuthSnapshot } + : undefined + : { kind: "api-key", provider: { ...route.provider } }; + if (binding) adapterBindings.set(resolved, binding); + const build = resolved.buildRequest.bind(resolved); + resolved.buildRequest = async (requestParsed, incoming) => { + const request = await build(requestParsed, incoming); + // Capture at adapter creation, never from mutable serving state after an await. + if (binding) requestBindings.set(request, binding); + return request; + }; + if (resolved.runTurn) { + rawRunTurns.set(resolved, resolved.runTurn.bind(resolved)); + resolved.runTurn = (requestParsed, incoming, emit) => runSelectedTurn(resolved, requestParsed, incoming, emit); + } + return resolved; + }; + const refreshDispatchAdapter = async (requestParsed: OcxParsedRequest): Promise => { + if (route.provider.authMode === "oauth") { + if (!servingOAuthSnapshot || !await applyFailoverSnapshot(servingOAuthSnapshot, requestParsed)) { + throw new Error("OAuth account selection changed before dispatch"); + } + } else { + const current = resolveCurrentProviderApiKeyTransport(config, route.providerName, route.provider); + if (!current) throw new Error("API key selection is unavailable before dispatch"); + route.provider = current; + } + adapter = activeAdapter = runTurnAdapter = resolveSelectionAdapter( + resolveWireProtocolOverride(route.providerName, route.modelId, route.provider, inboundWire), + ); + invalidateSameTargetRequest(); + return adapter; + }; + const refreshRunTurnAdapter = async (requestParsed: OcxParsedRequest): Promise => { + requestParsed._cursorIdentityScope = undefined; + requestParsed._cursorConversationId = undefined; + if (requestParsed._providerContinuation?.cursor) { + const { cursor: _oldCursor, ...rest } = requestParsed._providerContinuation; + requestParsed._providerContinuation = rest; + } + return refreshDispatchAdapter(requestParsed); + }; + const runSelectedTurn = async ( + selectedAdapter: ProviderAdapter, + ...[requestParsed, incoming, emit]: Parameters> + ): Promise => { + for (let attempt = 0; attempt < 3; attempt++) { + if (!selectionIsCurrent(adapterBindings.get(selectedAdapter))) selectedAdapter = await refreshRunTurnAdapter(requestParsed); + const binding = adapterBindings.get(selectedAdapter); + const run = rawRunTurns.get(selectedAdapter); + if (!run) throw new Error("Selected provider no longer supports this turn transport"); + let sent = false; + let refused = false; + // Both main and image-loop callers already acquired the initial pacing slot. + // Subsequent physical messages retain this adapter/credential and are paced normally. + const fetch = providerFetch(route.provider, options.codexWsRuntimeIdentity, { + providerName: route.providerName, modelId: route.modelId, pacingSlotAcquired: true, + beforeDispatch: () => { + if (sent) return; + if (!selectionIsCurrent(binding)) { + refused = true; + throw new Error("Account selection changed before the first turn dispatch"); + } + sent = true; + }, + }); + try { + await run(requestParsed, { ...incoming, providerFetch: fetch }, event => { if (!refused) emit(event); }); + } catch (error) { + if (!refused) throw error; + } + if (!refused) return; + // The adapter may map the guard's exception to an error event. Neither that + // event nor a refused send may escape before retrying the newly selected account. + selectedAdapter = await refreshRunTurnAdapter(requestParsed); + } + throw new Error("Account selection changed repeatedly before turn dispatch"); + }; + const oauthDispatch = (wireRequest: AdapterRequest, requestParsed = parsed): ProviderFetchOptions["dispatchOverride"] => { + if (route.provider.authMode === "forward") return undefined; + return async (input, init, execute) => { + let destination = input; + let dispatchInit = init; + for (let attempt = 0; attempt < 3; attempt++) { + if (selectionIsCurrent(requestBindings.get(wireRequest))) { + const fetchImpl = (route.provider as OcxProviderConfig & { fetch?: typeof globalThis.fetch }).fetch ?? execute; + return fetchImpl(destination, dispatchInit); + } + const nextAdapter = await refreshDispatchAdapter(requestParsed); + const rebuilt = await nextAdapter.buildRequest(requestParsed, { + headers: selectedForwardHeaders, translatorBudget, + ...(imageTierBias > 0 ? { imageTierBias } : {}), + }); + const bodySize = checkOutboundBodySize(rebuilt.body, config.maxUpstreamBodyBytes); + if (!bodySize.admitted) { + rebuilt.releaseBodyObservation?.(); + return formatErrorResponse(413, "outbound_body_too_large", describeOutboundBodyRefusal(bodySize)); + } + const headers = new Headers(dispatchInit.headers); + for (const name of Object.keys(wireRequest.headers)) headers.delete(name); + for (const [name, value] of Object.entries(rebuilt.headers)) headers.set(name, value); + wireRequest.releaseBodyObservation?.(); + Object.assign(wireRequest, rebuilt); + const binding = requestBindings.get(rebuilt); + if (binding) requestBindings.set(wireRequest, binding); + else requestBindings.delete(wireRequest); + sameTargetRequest = wireRequest; + sameTargetParsed = requestParsed; + sameTargetToken = transportToken; + destination = rebuilt.url; + dispatchInit = { ...dispatchInit, method: rebuilt.method, headers, body: rebuilt.body }; + bindRouteReasoningReplayScope({ parsed: requestParsed, providerName: route.providerName, provider: route.provider, + adapterName: nextAdapter.name, oauthCredentialSnapshot: replayOAuthCredentialSnapshot }); + // The next iteration validates synchronously and calls fetch in that same turn. + } + throw new Error("OAuth account selection changed repeatedly before dispatch"); + }; + }; const anthropicSessionKey = route.providerName === "anthropic" && route.provider.authMode === "oauth" ? anthropicSessionKeyFromParts({ sessionIdHeader: sessionIdHeaderFromRequest(req.headers), @@ -3756,12 +3978,11 @@ async function handleResponsesInner( } return formatErrorResponse(401, "authentication_error", "No eligible Anthropic OAuth account available"); } - const accessToken = await getAnthropicPoolAccessToken(selection.accountId); - anthropicPoolAccountId = selection.accountId; - bindAnthropicSessionAffinity(anthropicSessionKey, selection.accountId); - promoteAnthropicActiveAccount(selection.accountId); - route.provider = { ...route.provider, apiKey: accessToken }; - logCtx.provider = formatAnthropicProviderForLog("anthropic", selection.accountId, config); + const admitted = await commitResolvedOAuthSelection(await getAnthropicPoolAccessSnapshot(selection.accountId), true, selection.reason); + if (!admitted) return formatErrorResponse(409, "conflict_error", "OAuth account selection changed; retry the request"); + anthropicPoolAccountId = admitted.accountId; + route.provider = { ...route.provider, apiKey: admitted.accessToken }; + logCtx.provider = formatAnthropicProviderForLog("anthropic", admitted.accountId, config); } else { // Prefer the account with known headroom BEFORE the first attempt. Rotation alone // only reacts to a 429, so a turn could open on an account a previous probe already @@ -3810,6 +4031,10 @@ async function handleResponsesInner( resolved = await getValidAccessTokenSnapshot(route.providerName); usedPreferredAccount = false; } + const admitted = await commitResolvedOAuthSelection(resolved, true); + if (!admitted) return formatErrorResponse(409, "conflict_error", "OAuth account selection changed; retry the request"); + if (admitted.accountId !== resolved.accountId) usedPreferredAccount = true; + resolved = admitted; replayOAuthCredentialSnapshot = { accountId: resolved.accountId, generation: resolved.generation, @@ -3842,22 +4067,11 @@ async function handleResponsesInner( // Only genuinely accountless adapter calls leave the context undefined and use local/env fallback. parsed._kiroAuthContext = { ...(resolved.kiro ?? {}) }; } - // Antigravity (cloud-code-assist) needs the discovered Cloud Code Assist project id in the - // CCA envelope. Keep it paired with the token snapshot so an account rotation cannot mix - // a fresh token with project metadata re-read from a different credential generation. + // Project identity belongs to the admitted account on EVERY request, including + // the request after a pool transition made that account the persisted active one. if (route.provider.googleMode === "cloud-code-assist") { - // When pre-dispatch chose a DIFFERENT account, the configured project belongs to - // the account we did not use, and `!route.provider.project` would skip right past - // it — installing B's bearer alongside A's project. That is the #2841 pairing bug - // in its original shape, so the preferred-account path replaces the project - // unconditionally and refuses to dispatch at all if the chosen account has none. - // A project-less preferred account already fell back above, so by here the - // preferred path always has one. - if (usedPreferredAccount && resolved.projectId) { - route.provider = { ...route.provider, project: resolved.projectId }; - } else if (!route.provider.project && resolved.projectId) { - route.provider = { ...route.provider, project: resolved.projectId }; - } + if (!resolved.projectId) return formatErrorResponse(401, "authentication_error", publicOAuthAuthenticationErrorMessage(new Error("Cloud Code Assist account project is unavailable"))); + route.provider = { ...route.provider, project: resolved.projectId }; } } } catch (err) { @@ -3899,7 +4113,7 @@ async function handleResponsesInner( logCtx.provider = route.providerName; delete logCtx.accountLogLabel; } - const adapter = resolveAdapter(adapterProvider, config.cacheRetention); + adapter = resolveSelectionAdapter(adapterProvider, config.cacheRetention); bindRouteReasoningReplayScope({ parsed, providerName: route.providerName, @@ -3928,7 +4142,7 @@ async function handleResponsesInner( } sealRequestAttemptIdentity(logCtx.activeAttempt, logCtx.provider, adapter.name, logCtx.accountLogLabel); recordAttemptCredentialSource(logCtx.activeAttempt, route.providerName, adapterProvider, adapter.name); - let runTurnAdapter = adapter; + runTurnAdapter = adapter; if (adapter.runTurn) { recordAdapterTierMetadata(logCtx, adapter.tierLogForRunTurn?.(parsed)); } @@ -4529,6 +4743,7 @@ async function handleResponsesInner( body: request.body, }, recovery), upstream.signal, connectMs, parsed.stream, providerFetch(route.provider, options.codexWsRuntimeIdentity, { + dispatchOverride: oauthDispatch(request), providerName: route.providerName, modelId: route.modelId, onCodexWsQuota: codexWsQuotaObserver(authCtx, route.provider), @@ -4559,7 +4774,7 @@ async function handleResponsesInner( const rebuildAndRefetch = async ( recovery: AttemptRecoveryKind, ): Promise => { - const retryAdapter = resolveAdapter( + const retryAdapter = resolveSelectionAdapter( resolveWireProtocolOverride(route.providerName, route.modelId, route.provider, inboundWire), config.cacheRetention, ); @@ -4606,6 +4821,7 @@ async function handleResponsesInner( body: request.body, }, innerRecovery), upstream.signal, connectMs, parsed.stream, providerFetch(route.provider, options.codexWsRuntimeIdentity, { + dispatchOverride: oauthDispatch(request), providerName: route.providerName, modelId: route.modelId, onCodexWsQuota: codexWsQuotaObserver(authCtx, route.provider), @@ -4664,7 +4880,7 @@ async function handleResponsesInner( authCtx = replay.authCtx; route.provider = replay.provider; selectedForwardHeaders = replay.headers; - const replayAdapter = resolveAdapter( + const replayAdapter = resolveSelectionAdapter( resolveWireProtocolOverride(route.providerName, route.modelId, replay.provider, inboundWire), config.cacheRetention, ); @@ -4711,6 +4927,7 @@ async function handleResponsesInner( // here on is a genuine transport attempt. storedPoolReplayDispatchNotifier( providerFetch(route.provider, options.codexWsRuntimeIdentity, { + dispatchOverride: oauthDispatch(request), providerName: route.providerName, modelId: route.modelId, onCodexWsQuota: codexWsQuotaObserver(authCtx, route.provider), @@ -4748,7 +4965,7 @@ async function handleResponsesInner( try { void upstreamResponse.body?.cancel().catch(() => {}); } catch { /* already consumed/closed */ } let refreshed: OAuthAccessSnapshot; try { - refreshed = await forceRefreshOAuthAccessSnapshot(sentOAuthSnapshot); + refreshed = await refreshResolvedOAuthSelection(sentOAuthSnapshot); } catch (err) { upstream.abort(); releaseCodexAuthContextProbeLease(authCtx); @@ -4780,7 +4997,7 @@ async function handleResponsesInner( : undefined, ); route.provider = refreshedProvider; - const refreshedAdapter = resolveAdapter( + const refreshedAdapter = resolveSelectionAdapter( resolveWireProtocolOverride(route.providerName, route.modelId, refreshedProvider, inboundWire), config.cacheRetention, ); @@ -4830,6 +5047,7 @@ async function handleResponsesInner( body: request.body, }, recovery), upstream.signal, connectMs, parsed.stream, providerFetch(route.provider, options.codexWsRuntimeIdentity, { + dispatchOverride: oauthDispatch(request), providerName: route.providerName, modelId: route.modelId, onCodexWsQuota: codexWsQuotaObserver(authCtx, route.provider), @@ -4868,13 +5086,10 @@ async function handleResponsesInner( try { snapshot = await failoverAccountSnapshot(route.providerName, nextAccountId); } catch { /* Keep the original 429 body readable when the next credential is unavailable. */ } } - if (snapshot && applyFailoverSnapshot(snapshot)) { - genericFailoverAccountId = snapshot.accountId; + if (snapshot && await applyFailoverSnapshot(snapshot)) { genericFailovers += 1; - sentOAuthSnapshot = snapshot; - replayOAuthCredentialSnapshot = { accountId: snapshot.accountId, generation: snapshot.generation }; route.provider = resolveProviderTransport( - route.providerName, route.provider, parsed.options.promptCacheKey, snapshot.apiBaseUrl, + route.providerName, route.provider, parsed.options.promptCacheKey, sentOAuthSnapshot?.apiBaseUrl, ); bindRouteReasoningReplayScope({ parsed, providerName: route.providerName, provider: route.provider, @@ -4932,6 +5147,7 @@ async function handleResponsesInner( body: request.body, }, recovery), upstream.signal, connectMs, parsed.stream, providerFetch(route.provider, options.codexWsRuntimeIdentity, { + dispatchOverride: oauthDispatch(request), providerName: route.providerName, modelId: route.modelId, onCodexWsQuota: codexWsQuotaObserver(authCtx, route.provider), @@ -5751,9 +5967,8 @@ async function handleResponsesInner( if (!nextAccountId) return null; try { const snapshot = await failoverAccountSnapshot(route.providerName, nextAccountId); - genericFailoverAccountId = nextAccountId; genericFailovers += 1; - if (!applyFailoverSnapshot(snapshot)) return null; + if (!await applyFailoverSnapshot(snapshot)) return null; } catch { return null; } @@ -5777,12 +5992,12 @@ async function handleResponsesInner( // carries none, and getAnthropicPoolAccessToken is what enforces its fail-closed // local-cli credential rule. Both existing Anthropic rotation sites apply the token the // same way. - const accessToken = await getAnthropicPoolAccessToken(nextAccountId); - anthropicPoolAccountId = nextAccountId; + const admitted = await commitResolvedOAuthSelection(await getAnthropicPoolAccessSnapshot(nextAccountId)); + if (!admitted) throw new Error("OAuth selection changed during recovery"); + anthropicPoolAccountId = admitted.accountId; anthropicPoolFailovers += 1; - route.provider = { ...route.provider, apiKey: accessToken }; - promoteAnthropicActiveAccount(nextAccountId); - logCtx.provider = formatAnthropicProviderForLog("anthropic", nextAccountId, config); + route.provider = { ...route.provider, apiKey: admitted.accessToken }; + logCtx.provider = formatAnthropicProviderForLog("anthropic", admitted.accountId, config); } catch { return null; } @@ -5791,7 +6006,7 @@ async function handleResponsesInner( // credential. The 429 is terminal for this sidecar turn. return null; } - const rotatedAdapter = resolveAdapter( + const rotatedAdapter = resolveSelectionAdapter( resolveWireProtocolOverride(route.providerName, route.modelId, route.provider, inboundWire), config.cacheRetention, ); @@ -5875,6 +6090,13 @@ async function handleResponsesInner( stallTimeoutSec: config.stallTimeoutSec, waitForRequestSlot: imageProviderFetch.waitForPacing, fetchImpl: imageProviderFetch.unpacedFetch ?? imageProviderFetch, + fetchForRequest: (request, iterParsed) => { + const fetch = providerFetch(route.provider, options.codexWsRuntimeIdentity, { + dispatchOverride: oauthDispatch(request, iterParsed), + providerName: route.providerName, modelId: route.modelId, + }); + return fetch.unpacedFetch ?? fetch; + }, onRequestBuilt: request => { recordAdapterReasoning(logCtx, request); recordAdapterTier(logCtx, request); @@ -5935,6 +6157,10 @@ async function handleResponsesInner( })(input, init)) as typeof globalThis.fetch; const wsResponse = await runWithWebSearch({ parsed, adapter, + fetchForRequest: (request, iterParsed) => providerFetch(route.provider, options.codexWsRuntimeIdentity, { + dispatchOverride: oauthDispatch(request, iterParsed), + providerName: route.providerName, modelId: route.modelId, + }), incomingMeta: { headers: selectedForwardHeaders, abortSignal: options.abortSignal, @@ -6010,6 +6236,13 @@ async function handleResponsesInner( const queue = createAdapterEventQueue({ onBacklogExceeded: () => runTurnAbort.abort(), }); + const refreshRunTurnSelection = async (): Promise => { + if (selectionIsCurrent(adapterBindings.get(runTurnAdapter))) return; + await refreshRunTurnAdapter(parsed); + bindRouteReasoningReplayScope({ parsed, providerName: route.providerName, provider: route.provider, + adapterName: runTurnAdapter.name, oauthCredentialSnapshot: replayOAuthCredentialSnapshot }); + sealRequestAttemptIdentity(logCtx.activeAttempt, logCtx.provider, runTurnAdapter.name, logCtx.accountLogLabel); + }; // Initial admission must settle before the streaming Response commits HTTP 200. // Let the outer Responses facade preserve the local retryable-429 contract. try { @@ -6033,6 +6266,7 @@ async function handleResponsesInner( if (!pacingSlotAcquired) { await waitForProviderRequestSlot(route.providerName, route.provider, route.modelId, runTurnAbort.signal); } + await refreshRunTurnSelection(); noteAttemptSend(logCtx.activeAttempt, logCtx.usageLogInputTokens, recovery); const runTurnProviderFetch = providerFetch( route.provider, @@ -6098,9 +6332,8 @@ async function handleResponsesInner( if (!nextAccountId) return false; try { const snapshot = await failoverAccountSnapshot(route.providerName, nextAccountId); - genericFailoverAccountId = nextAccountId; genericFailovers += 1; - if (!applyFailoverSnapshot(snapshot)) return false; + if (!await applyFailoverSnapshot(snapshot)) return false; // A Cursor conversation/checkpoint is credential-scoped. The failed attempt emitted no // client-visible bytes, so replay is safe, but carrying its account identity into the next // account would not be. Let the rotated adapter derive a fresh identity and conversation. @@ -6116,7 +6349,7 @@ async function handleResponsesInner( route.provider, inboundWire, ); - const rotatedAdapter = resolveAdapter(rotatedProvider, config.cacheRetention); + const rotatedAdapter = resolveSelectionAdapter(rotatedProvider, config.cacheRetention); if (!rotatedAdapter.runTurn) return false; runTurnAdapter = rotatedAdapter; bindRouteReasoningReplayScope({ @@ -6166,7 +6399,7 @@ async function handleResponsesInner( if (parsed.stream) { void runTurn(); let eventSource: AsyncIterable = queue.stream(); - if (genericFailoverAccountId && isGenericOAuthFailoverEnabled(config, route.providerName)) { + if (route.provider.authMode === "oauth" || (genericFailoverAccountId && isGenericOAuthFailoverEnabled(config, route.providerName))) { // Preflight holds only heartbeats and the first meaningful event. A first-event 429 can be // replayed transparently; after any output reaches the bridge, a later error stays terminal. eventSource = await preflightRunTurnFailover(eventSource); @@ -6249,7 +6482,7 @@ async function handleResponsesInner( await runTurn(); const firstAttemptEvents = await queue.collect(); let runTurnEvents: AdapterEvent[] = firstAttemptEvents; - if (genericFailoverAccountId && isGenericOAuthFailoverEnabled(config, route.providerName)) { + if (route.provider.authMode === "oauth" || (genericFailoverAccountId && isGenericOAuthFailoverEnabled(config, route.providerName))) { runTurnEvents = []; for await (const event of await preflightRunTurnFailover( (async function* () { yield* firstAttemptEvents; })(), @@ -6321,7 +6554,7 @@ async function handleResponsesInner( const stallTimeoutMs = typeof config.stallTimeoutSec === "number" && Number.isFinite(config.stallTimeoutSec) && config.stallTimeoutSec > 0 ? Math.floor(config.stallTimeoutSec * 1000) : 300_000; - let activeAdapter = adapter; + activeAdapter = adapter; // One immutable, body-safe outbound request per same-target sequence (URL, serialized body, // auth headers, generated compat headers). Same-target 429 replays reuse it verbatim; the @@ -6420,16 +6653,15 @@ async function handleResponsesInner( // Capture it in a const so the fetch callbacks read a narrowed, immutable value // (TypeScript drops narrowing for a `let` captured by a nested function). const builtInitialRequest = initialRequest; - let sameTargetRequest: AdapterRequest | undefined = builtInitialRequest; - let sameTargetParsed: OcxParsedRequest | undefined = parsed; - let sameTargetToken = 0; - let transportToken = 0; + sameTargetRequest = builtInitialRequest; + sameTargetParsed = parsed; + sameTargetToken = transportToken; /** * Invalidate the same-target request cache. Every credential/adapter/parsed mutation MUST * go through here: the cache keys on `parsed` REFERENCE identity, so an in-place mutation * is invisible to it and a missed bump would replay a request built with a stale key. */ - const invalidateSameTargetRequest = (): void => { transportToken += 1; }; + let upstreamResponse: Response; try { if (activeAdapter.fetchResponse) { @@ -6440,6 +6672,7 @@ async function handleResponsesInner( timeoutMs: connectMs, stream: parsed.stream, executor: providerFetch(route.provider, options.codexWsRuntimeIdentity, { + dispatchOverride: oauthDispatch(builtInitialRequest), providerName: route.providerName, modelId: route.modelId, }), @@ -6465,6 +6698,7 @@ async function handleResponsesInner( body: builtInitialRequest.body, }, recovery), upstream.signal, connectMs, parsed.stream, providerFetch(route.provider, options.codexWsRuntimeIdentity, { + dispatchOverride: oauthDispatch(builtInitialRequest), providerName: route.providerName, modelId: route.modelId, })); @@ -6496,7 +6730,6 @@ async function handleResponsesInner( let rateLimitRetries = 0; // Shared with the terminal-guard continuation below: an image-tier reduction that let the // main request clear a 413 must not be forgotten on the very next continuation build. - let imageTierBias = 0; if (!upstreamResponse.ok) { // Recovery loop: multi-key 429 failover + at most ONE opaque-state rebuild and ONE // anthropic 413 tightened retry @@ -6560,6 +6793,7 @@ async function handleResponsesInner( timeoutMs: connectMs, stream: parsed.stream, executor: providerFetch(route.provider, options.codexWsRuntimeIdentity, { + dispatchOverride: oauthDispatch(retryRequest), providerName: route.providerName, modelId: route.modelId, }), @@ -6581,6 +6815,7 @@ async function handleResponsesInner( method: retryRequest.method, headers: retryRequest.headers, body: retryRequest.body, }, recoveryKind), upstream.signal, connectMs, parsed.stream, providerFetch(route.provider, options.codexWsRuntimeIdentity, { + dispatchOverride: oauthDispatch(retryRequest), providerName: route.providerName, modelId: route.modelId, })), @@ -6620,7 +6855,7 @@ async function handleResponsesInner( try { void upstreamResponse.body?.cancel().catch(() => {}); } catch { /* already consumed/closed */ } let refreshed: OAuthAccessSnapshot; try { - refreshed = await forceRefreshOAuthAccessSnapshot(sentOAuthSnapshot); + refreshed = await refreshResolvedOAuthSelection(sentOAuthSnapshot); } catch (err) { cleanupUpstreamAbort(); return formatErrorResponse(401, "authentication_error", publicOAuthAuthenticationErrorMessage(err)); @@ -6651,7 +6886,7 @@ async function handleResponsesInner( ); route.provider = refreshedProvider; invalidateSameTargetRequest(); - activeAdapter = resolveAdapter( + activeAdapter = resolveSelectionAdapter( resolveWireProtocolOverride(route.providerName, route.modelId, refreshedProvider, inboundWire), config.cacheRetention, ); @@ -6684,7 +6919,7 @@ async function handleResponsesInner( try { void upstreamResponse.body?.cancel().catch(() => {}); } catch { /* already consumed/closed */ } route.provider = rotated; invalidateSameTargetRequest(); - activeAdapter = resolveAdapter( + activeAdapter = resolveSelectionAdapter( resolveWireProtocolOverride(route.providerName, route.modelId, route.provider, inboundWire), config.cacheRetention, ); @@ -6754,7 +6989,7 @@ async function handleResponsesInner( try { void upstreamResponse.body?.cancel().catch(() => {}); } catch { /* already consumed/closed */ } route.provider = rotated; invalidateSameTargetRequest(); - activeAdapter = resolveAdapter( + activeAdapter = resolveSelectionAdapter( resolveWireProtocolOverride(route.providerName, route.modelId, route.provider, inboundWire), config.cacheRetention, ); @@ -6785,14 +7020,14 @@ async function handleResponsesInner( if (!nextAccountId) break; try { void upstreamResponse.body?.cancel().catch(() => {}); } catch { /* already consumed/closed */ } try { - const accessToken = await getAnthropicPoolAccessToken(nextAccountId); - anthropicPoolAccountId = nextAccountId; + const admitted = await commitResolvedOAuthSelection(await getAnthropicPoolAccessSnapshot(nextAccountId)); + if (!admitted) throw new Error("OAuth selection changed during recovery"); + anthropicPoolAccountId = admitted.accountId; anthropicPoolFailovers += 1; - route.provider = { ...route.provider, apiKey: accessToken }; + route.provider = { ...route.provider, apiKey: admitted.accessToken }; invalidateSameTargetRequest(); - promoteAnthropicActiveAccount(nextAccountId); - logCtx.provider = formatAnthropicProviderForLog("anthropic", nextAccountId, config); - activeAdapter = resolveAdapter( + logCtx.provider = formatAnthropicProviderForLog("anthropic", admitted.accountId, config); + activeAdapter = resolveSelectionAdapter( resolveWireProtocolOverride(route.providerName, route.modelId, route.provider, inboundWire), config.cacheRetention, ); @@ -6832,11 +7067,10 @@ async function handleResponsesInner( // projectId with its token and Kiro carries routing metadata, so a token-only swap // would mix one account's credential with another's routing data. const snapshot = await failoverAccountSnapshot(route.providerName, nextAccountId); - genericFailoverAccountId = nextAccountId; - genericFailovers += 1; - if (!applyFailoverSnapshot(snapshot)) break; + genericFailovers += 1; + if (!await applyFailoverSnapshot(snapshot)) break; invalidateSameTargetRequest(); - activeAdapter = resolveAdapter( + activeAdapter = resolveSelectionAdapter( resolveWireProtocolOverride(route.providerName, route.modelId, route.provider, inboundWire), config.cacheRetention, ); @@ -7041,6 +7275,7 @@ async function handleResponsesInner( timeoutMs: connectMs, stream: nextParsed.stream, executor: providerFetch(route.provider, options.codexWsRuntimeIdentity, { + dispatchOverride: oauthDispatch(builtContinuationRequest, nextParsed), providerName: route.providerName, modelId: nextParsed.modelId, }), @@ -7066,6 +7301,7 @@ async function handleResponsesInner( connectMs, nextParsed.stream, providerFetch(route.provider, options.codexWsRuntimeIdentity, { + dispatchOverride: oauthDispatch(builtContinuationRequest, nextParsed), providerName: route.providerName, modelId: nextParsed.modelId, }), @@ -7161,7 +7397,7 @@ async function handleResponsesInner( try { void response.body?.cancel().catch(() => {}); } catch { /* already closed */ } route.provider = rotated; invalidateSameTargetRequest(); - activeAdapter = resolveAdapter( + activeAdapter = resolveSelectionAdapter( resolveWireProtocolOverride(route.providerName, route.modelId, route.provider, inboundWire), config.cacheRetention, ); @@ -7197,14 +7433,14 @@ async function handleResponsesInner( if (nextAccountId) { try { void response.body?.cancel().catch(() => {}); } catch { /* already closed */ } try { - const accessToken = await getAnthropicPoolAccessToken(nextAccountId); - anthropicPoolAccountId = nextAccountId; + const admitted = await commitResolvedOAuthSelection(await getAnthropicPoolAccessSnapshot(nextAccountId)); + if (!admitted) throw new Error("OAuth selection changed during recovery"); + anthropicPoolAccountId = admitted.accountId; anthropicPoolFailovers += 1; - route.provider = { ...route.provider, apiKey: accessToken }; + route.provider = { ...route.provider, apiKey: admitted.accessToken }; invalidateSameTargetRequest(); - promoteAnthropicActiveAccount(nextAccountId); - logCtx.provider = formatAnthropicProviderForLog("anthropic", nextAccountId, config); - activeAdapter = resolveAdapter( + logCtx.provider = formatAnthropicProviderForLog("anthropic", admitted.accountId, config); + activeAdapter = resolveSelectionAdapter( resolveWireProtocolOverride(route.providerName, route.modelId, route.provider, inboundWire), config.cacheRetention, ); @@ -7242,11 +7478,10 @@ async function handleResponsesInner( // metadata, so a token-only swap would mix one account's credential with another's // routing data. const snapshot = await failoverAccountSnapshot(route.providerName, nextAccountId); - genericFailoverAccountId = nextAccountId; - genericFailovers += 1; - if (applyFailoverSnapshot(snapshot, nextParsed)) { + genericFailovers += 1; + if (await applyFailoverSnapshot(snapshot, nextParsed)) { invalidateSameTargetRequest(); - activeAdapter = resolveAdapter( + activeAdapter = resolveSelectionAdapter( resolveWireProtocolOverride(route.providerName, route.modelId, route.provider, inboundWire), config.cacheRetention, ); diff --git a/src/server/responses/fetch-helpers.ts b/src/server/responses/fetch-helpers.ts index e1423d3311..b6365be4be 100644 --- a/src/server/responses/fetch-helpers.ts +++ b/src/server/responses/fetch-helpers.ts @@ -61,6 +61,8 @@ export interface ProviderFetchOptions { onCodexWsQuota?: CodexWsQuotaObserver; /** Synchronous admission at actual credential dispatch, after pacing/backoff. */ beforeDispatch?: (headers: Headers) => void; + /** Revalidate/rebuild a queued request at its physical send boundary, after pacing. */ + dispatchOverride?: (input: Parameters[0], init: RequestInit, execute: typeof globalThis.fetch) => Promise; } export function providerFetch( @@ -75,7 +77,10 @@ export function providerFetch( const httpFetch = Object.assign( async (input: Parameters[0], init?: RequestInit) => { options.beforeDispatch?.(new Headers(init?.headers ?? (input instanceof Request ? input.headers : undefined))); - return base(input, { ...withUpstreamHttpVersion(input, init, provider), timeout: 0 }); + const dispatchInit = { ...withUpstreamHttpVersion(input, init, provider), timeout: 0 }; + return options.dispatchOverride + ? options.dispatchOverride(input, dispatchInit, base) + : base(input, dispatchInit); }, { preconnect }, ) as typeof globalThis.fetch; diff --git a/src/types/config.ts b/src/types/config.ts index 27b1a8ec81..c6fdfb063a 100644 --- a/src/types/config.ts +++ b/src/types/config.ts @@ -815,10 +815,10 @@ export interface OcxConfig { * provider has 2 or more eligible stored accounts, the same consent rule an `apiKeyPool` of * two keys already applies, and a single account remains a strict no-op. * - * What `enabled: false` still refuses is the PRE-DISPATCH preference: steering a request - * upstream has not refused toward the account with more known headroom. That moves a healthy - * request, so it stays a real choice. `providers..oauthAccountFailover` overrides this - * per provider in either direction; reactive 429 rotation remains presence-driven. + * Proactive avoidance of an exhausted selected account requires `enabled: true`. + * A healthy selected account retains priority; an unknown quota is not exhaustion. + * `providers..oauthAccountFailover` overrides this per provider in either direction. + * Reactive 429 rotation remains presence-driven even when proactive routing is disabled. */ oauthAccountFailover?: { enabled?: boolean; diff --git a/src/types/provider.ts b/src/types/provider.ts index cb2abc1f0c..97a359506a 100644 --- a/src/types/provider.ts +++ b/src/types/provider.ts @@ -8,6 +8,13 @@ import type { UpstreamHttpVersion, ReasoningSummaryDelivery, CodexAccountMode } */ export type RefreshPolicy = "proactive" | "lazy-only" | "disabled"; +/** Request-owned identity of the configured key, before env/keychain resolution. */ +export interface ProviderApiKeySelection { + entryId?: string; + reference?: string; + revision?: string; +} + export interface OpenRouterProviderRouting { /** OpenRouter provider slugs to try first, in priority order. */ order?: string[]; @@ -329,6 +336,10 @@ export interface OcxProviderConfig { * `apiKey` seeds a one-entry pool on first management touch. */ apiKeyPool?: Array<{ id: string; key: string; label?: string; addedAt?: number }>; + /** Changes on manual selection (including re-selection) and committed automatic allocation. */ + apiKeySelectionRevision?: string; + /** Runtime only. Never expose in management responses or persist a routed provider. */ + _apiKeyAttempt?: ProviderApiKeySelection; defaultModel?: string; models?: string[]; /** @@ -441,9 +452,9 @@ export interface OcxProviderConfig { * * Reactive 429 rotation is presence-driven and cannot be refused here — 2+ logged-in accounts * activate it, and a 429 with an idle second account is a defect rather than a preference. - * What an explicit `false` still refuses is the pre-dispatch preference that steers a HEALTHY - * request toward the account with more known headroom. It beats the global - * `oauthAccountFailover` in either direction; reactive 429 rotation remains presence-driven. + * Proactive exhaustion avoidance requires explicit `true`; a healthy selected account + * retains priority. This overrides global `oauthAccountFailover` in either direction. + * Reactive 429 rotation remains available even when proactive routing is disabled. */ oauthAccountFailover?: { enabled?: boolean; diff --git a/src/web-search/loop.ts b/src/web-search/loop.ts index 682e482eea..3a2c5e99b4 100644 --- a/src/web-search/loop.ts +++ b/src/web-search/loop.ts @@ -301,6 +301,8 @@ export interface WebSearchLoopDeps { onUsage?: (usage: OcxUsage | undefined) => void; /** Observe the exact adapter request selected for each routed-model iteration. */ onRequestBuilt?: (request: AdapterRequest) => void; + /** Request-scoped executor retains the core's selection binding across loop retries. */ + fetchForRequest?: (request: AdapterRequest, parsed: OcxParsedRequest) => typeof globalThis.fetch; /** Called before each routed-model dispatch in the loop, for attempt telemetry. Same-target 429 replays pass the `rate-limit-429` recovery kind. */ onAttemptSend?: (recovery?: AttemptRecoveryKind) => void; /** @@ -440,6 +442,7 @@ export async function runWithWebSearch(deps: WebSearchLoopDeps): Promise ` (`rename` is accepted as a synonym). +OAuth manual and automatic selection share `commitOAuthAccountSelection` in the auth store. +The caller resolves a usable credential, commits its matching selection, then dispatches it; +request-local token replacement must not leave a different dashboard account selected. +Opaque selection revisions protect manual reselection and A→B→A changes from older requests. +Credential-only refresh preserves the revision. Generic proactive routing is opt-in and retains +a healthy selected account; reactive 429 recovery remains available even when the pool is off. +API-key manual selection and failover similarly share `commitProviderApiKeySelection`, carrying +stable entry identity and selection revision instead of comparing a resolved secret with an env reference. + +The authenticated `GET /api/accounts/events` stream invalidates account/key selection after +successful persistence. Events contain provider/kind/revision only. The dashboard immediately +reconciles the cheap local roster and preserves its quota rows; no upstream quota probe is caused +by an event. One screen-owned stream has disconnect cleanup and bounded server subscribers; +reconnection and the existing shared scheduler provide recovery. Codex retains its own established +selection controller. These events cannot change credentials or select an account. + Selection order is the opposite case and must not be folded into the alias route. `codexAccountPriorities` is routing metadata that Pool selection consults, it lives in config rather than on `CodexAccount` so the `__main__` Desktop login can carry one, and the alias route's rejection of `__main__` would be wrong for diff --git a/tests/adapters/anthropic/anthropic-account-pool.test.ts b/tests/adapters/anthropic/anthropic-account-pool.test.ts index 7817fd7104..61396bd71e 100644 --- a/tests/adapters/anthropic/anthropic-account-pool.test.ts +++ b/tests/adapters/anthropic/anthropic-account-pool.test.ts @@ -16,8 +16,12 @@ import { resolveAnthropicAccountForSession, resetAnthropicRoutingForManualSelection, rotateAnthropicAccountOn429, + getAnthropicPoolAccessSnapshot, + promoteAnthropicActiveAccount, + anthropicSessionAffinitySizeForTests, } from "../../../src/oauth/anthropic-routing"; -import { getAccountSet, saveCredential, setActiveAccount } from "../../../src/oauth/store"; +import { captureOAuthAccountSelection, getAccountSet, markAccountNeedsReauth, saveCredential, saveAccountCredential, setActiveAccount } from "../../../src/oauth/store"; +import { subscribeAccountSelections } from "../../../src/lib/account-selection-events"; import { clearAccountQuotaCache, setCachedProviderAccountQuotaForTests } from "../../../src/providers/quota"; import type { OcxAccountPoolQuotaWindow, OcxAccountPoolRotationStrategy, OcxConfig } from "../../../src/types"; import { removeTreeWithRetry } from "../../helpers/remove-tree"; @@ -25,6 +29,18 @@ import { removeTreeWithRetry } from "../../helpers/remove-tree"; const originalHome = process.env.OPENCODEX_HOME; let home: string; +async function admitAnthropic(sessionKey: string, config: OcxConfig) { + const expected = captureOAuthAccountSelection("anthropic"); + const choice = resolveAnthropicAccountForSession(sessionKey, config); + if (choice.accountId) { + const snapshot = await getAnthropicPoolAccessSnapshot(choice.accountId); + expect(await promoteAnthropicActiveAccount(choice.accountId, expected, { + config, sessionKey, reason: choice.reason, expectedCredentialGeneration: snapshot.generation, + })).not.toBeNull(); + } + return choice; +} + beforeEach(() => { home = mkdtempSync(join(tmpdir(), "ocx-anthropic-pool-")); process.env.OPENCODEX_HOME = home; @@ -63,6 +79,9 @@ async function seedTwoAccounts() { const a = set.accounts.find(acc => acc.credential.accountId === "uuid-aaaa")!; const b = set.accounts.find(acc => acc.credential.accountId === "uuid-bbbb")!; await setActiveAccount("anthropic", a.id); + // Ordinary policy cases start after the initial stored selection has been admitted. + // Restart cases explicitly clear runtime state below to exercise first admission again. + await admitAnthropic("", cfg(false)); return { aId: a.id, bId: b.id }; } @@ -117,10 +136,122 @@ async function seedThreeAccounts() { const b = set.accounts.find(acc => acc.credential.accountId === "uuid-bbbb")!; const c = set.accounts.find(acc => acc.credential.accountId === "uuid-cccc")!; await setActiveAccount("anthropic", a.id); + await admitAnthropic("", cfg(false)); return { aId: a.id, bId: b.id, cId: c.id }; } describe("anthropic account pool", () => { + test.each(["round-robin", "quota"] as const)("persisted manual choice survives restart before the first %s dispatch", async strategy => { + const { aId, bId } = await seedTwoAccounts(); + setCachedProviderAccountQuotaForTests("anthropic", aId, { fiveHourPercent: 11 }); + setCachedProviderAccountQuotaForTests("anthropic", bId, { fiveHourPercent: 30 }); + await setActiveAccount("anthropic", bId); + resetAnthropicRoutingForManualSelection(bId); + const persisted = captureOAuthAccountSelection("anthropic"); + // A process restart loses both local preference and the shared RR cursor. + clearAnthropicAccountPoolState(); + clearPoolRotationState(); + const config = cfg(true, 20, { strategy, stickyLimit: 1 }); + expect(resolveAnthropicAccountForSession("restart-first", config).accountId).toBe(bId); + expect(resolveAnthropicAccountForSession("restart-proposal", config).accountId).toBe(bId); + expect(captureOAuthAccountSelection("anthropic")).toEqual(persisted); + expect((await admitAnthropic("restart-first", config)).accountId).toBe(bId); + // Once the authoritative first selection commits, the ordinary algorithm resumes. + expect((await admitAnthropic("restart-next", config)).accountId).toBe(aId); + expect(resolveAnthropicAccountForSession("restart-first", config).accountId).toBe(bId); + }); + + test("pool-off 429 recovery commits the replacement and notifies before dispatch", async () => { + const { aId, bId } = await seedTwoAccounts(); + const config = cfg(false); + const expected = captureOAuthAccountSelection("anthropic"); + const next = rotateAnthropicAccountOn429(config, aId, "30", "off-retry"); + expect(next).toBe(bId); + const snapshot = await getAnthropicPoolAccessSnapshot(bId); + let notifications = 0; + const unsubscribe = subscribeAccountSelections(event => { + if (event.provider === "anthropic") { + notifications++; + expect(captureOAuthAccountSelection("anthropic")?.accountId).toBe(bId); + } + }); + try { + expect(await promoteAnthropicActiveAccount(bId, expected, { + config, sessionKey: "off-retry", expectedCredentialGeneration: snapshot.generation, + })).not.toBeNull(); + expect(captureOAuthAccountSelection("anthropic")?.accountId).toBe(bId); + expect(notifications).toBe(1); + expect(anthropicSessionAffinitySizeForTests()).toBe(0); + } finally { unsubscribe(); } + }); + + test.each([false, true])("stale promotion cannot replace a newer manual choice (ABA=%s)", async aba => { + const { aId, bId } = await seedTwoAccounts(); + const expected = captureOAuthAccountSelection("anthropic"); + const snapshot = await getAnthropicPoolAccessSnapshot(bId); + await setActiveAccount("anthropic", bId); + if (aba) await setActiveAccount("anthropic", aId); + const manual = captureOAuthAccountSelection("anthropic"); + resetAnthropicRoutingForManualSelection(manual!.accountId); + let notifications = 0; + const unsubscribe = subscribeAccountSelections(() => { notifications++; }); + try { + expect(await promoteAnthropicActiveAccount(bId, expected, { + config: cfg(true), sessionKey: "stale", expectedCredentialGeneration: snapshot.generation, + })).toBeNull(); + expect(captureOAuthAccountSelection("anthropic")).toEqual(manual); + expect(anthropicSessionAffinitySizeForTests()).toBe(0); + expect(notifications).toBe(0); + expect(resolveAnthropicAccountForSession("next", cfg(true)).accountId).toBe(manual!.accountId); + } finally { unsubscribe(); } + }); + + test("token rejection does not consume the manual preference or install affinity", async () => { + const { aId, bId } = await seedTwoAccounts(); + await setActiveAccount("anthropic", aId); + resetAnthropicRoutingForManualSelection(aId); + const expected = captureOAuthAccountSelection("anthropic"); + const snapshot = await getAnthropicPoolAccessSnapshot(aId); + await markAccountNeedsReauth("anthropic", aId, true); + expect(await promoteAnthropicActiveAccount(aId, expected, { + config: cfg(true), sessionKey: "rejected", expectedCredentialGeneration: snapshot.generation, + })).toBeNull(); + expect(anthropicSessionAffinitySizeForTests()).toBe(0); + await markAccountNeedsReauth("anthropic", aId, false); + setCachedProviderAccountQuotaForTests("anthropic", aId, { fiveHourPercent: 30 }); + setCachedProviderAccountQuotaForTests("anthropic", bId, { fiveHourPercent: 11 }); + expect(resolveAnthropicAccountForSession("recovered", cfg(true, 20)).accountId).toBe(aId); + }); + + test("account snapshot refuses expired background local-CLI credentials without refreshing", async () => { + const { aId, bId } = await seedTwoAccounts(); + const account = getAccountSet("anthropic")!.accounts.find(account => account.id === bId)!; + await saveAccountCredential("anthropic", bId, { ...account.credential, source: "local-cli", expires: 1 }); + await expect(getAnthropicPoolAccessSnapshot(bId)).rejects.toThrow("background local-cli token expired"); + expect(captureOAuthAccountSelection("anthropic")?.accountId).toBe(aId); + }); + + test("manual choice wins the next healthy quota dispatch above the automatic threshold", async () => { + const { aId, bId } = await seedTwoAccounts(); + setCachedProviderAccountQuotaForTests("anthropic", aId, { fiveHourPercent: 30 }); + setCachedProviderAccountQuotaForTests("anthropic", bId, { fiveHourPercent: 11 }); + await setActiveAccount("anthropic", aId); + resetAnthropicRoutingForManualSelection(aId); + expect(resolveAnthropicAccountForSession("manual-quota", cfg(true, 20)).accountId).toBe(aId); + }); + + test("uncommitted proposals neither bind affinity nor advance round-robin", async () => { + const { aId, bId } = await seedTwoAccounts(); + const config = cfg(true, 80, { strategy: "round-robin", stickyLimit: 1 }); + const first = resolveAnthropicAccountForSession("uncommitted", config); + expect(resolveAnthropicAccountForSession("another-uncommitted", config).accountId).toBe(first.accountId); + // A failed candidate must not capture the task's affinity before its selection commits. + setCachedProviderAccountQuotaForTests("anthropic", aId, { fiveHourPercent: 90 }); + setCachedProviderAccountQuotaForTests("anthropic", bId, { fiveHourPercent: 5 }); + const quota = cfg(true); + expect(resolveAnthropicAccountForSession("uncommitted", quota).accountId).toBe(bId); + }); + test("default off always returns the active account", async () => { const { aId, bId } = await seedTwoAccounts(); expect(isAnthropicAccountPoolEnabled(cfg(false))).toBe(false); @@ -135,7 +266,7 @@ describe("anthropic account pool", () => { // Force lowest-usage toward B for a cold start with high active usage. setCachedProviderAccountQuotaForTests("anthropic", aId, { fiveHourPercent: 95 }); setCachedProviderAccountQuotaForTests("anthropic", bId, { fiveHourPercent: 10 }); - const first = resolveAnthropicAccountForSession("sess-sticky", cfg(true)); + const first = await admitAnthropic("sess-sticky", cfg(true)); expect(first.accountId).toBe(bId); // Even if A becomes "better", affinity keeps B. setCachedProviderAccountQuotaForTests("anthropic", aId, { fiveHourPercent: 1 }); @@ -213,9 +344,9 @@ describe("anthropic account pool", () => { const config = cfg(true, 80, { strategy: "round-robin" }); const picks = [ - resolveAnthropicAccountForSession("sess-1", config).accountId, - resolveAnthropicAccountForSession("sess-2", config).accountId, - resolveAnthropicAccountForSession("sess-3", config).accountId, + (await admitAnthropic("sess-1", config)).accountId, + (await admitAnthropic("sess-2", config)).accountId, + (await admitAnthropic("sess-3", config)).accountId, ]; expect(new Set(picks).size).toBe(3); }); @@ -239,7 +370,7 @@ describe("anthropic account pool", () => { setCachedProviderAccountQuotaForTests("anthropic", cId, { fiveHourPercent: 10 }); const config = cfg(true, 80, { strategy: "round-robin" }); - const first = resolveAnthropicAccountForSession("T", config); + const first = await admitAnthropic("T", config); expect(first.accountId).toBeTruthy(); const pinned = first.accountId!; await setActiveAccount("anthropic", pinned === aId ? bId : aId); @@ -293,11 +424,11 @@ describe("anthropic account pool", () => { setCachedProviderAccountQuotaForTests("anthropic", cId, { fiveHourPercent: 10 }); const config = cfg(true, 80, { strategy: "round-robin", stickyLimit: 3 }); - const first = resolveAnthropicAccountForSession("s1", config).accountId; + const first = (await admitAnthropic("s1", config)).accountId; expect(first).toBeTruthy(); - expect(resolveAnthropicAccountForSession("s2", config).accountId).toBe(first); - expect(resolveAnthropicAccountForSession("s3", config).accountId).toBe(first); - const fourth = resolveAnthropicAccountForSession("s4", config).accountId; + expect((await admitAnthropic("s2", config)).accountId).toBe(first); + expect((await admitAnthropic("s3", config)).accountId).toBe(first); + const fourth = (await admitAnthropic("s4", config)).accountId; expect(fourth).not.toBe(first); }); @@ -308,17 +439,17 @@ describe("anthropic account pool", () => { setCachedProviderAccountQuotaForTests("anthropic", cId, { fiveHourPercent: 10 }); const config = cfg(true, 80, { strategy: "round-robin", stickyLimit: 10 }); - const sticky = resolveAnthropicAccountForSession("sticky-1", config).accountId!; + const sticky = (await admitAnthropic("sticky-1", config)).accountId!; expect(resolveAnthropicAccountForSession("sticky-2", config).accountId).toBe(sticky); notePoolRotationFailure(POOL_KEY_ANTHROPIC, sticky); - const afterClear = resolveAnthropicAccountForSession("sticky-3", config).accountId; + const afterClear = (await admitAnthropic("sticky-3", config)).accountId; expect(afterClear).toBeTruthy(); expect(afterClear).not.toBe(sticky); // Re-establish sticky, then 429-cool the sticky account — failover + ring must leave it. clearPoolRotationState(); - const again = resolveAnthropicAccountForSession("again-1", config).accountId!; + const again = (await admitAnthropic("again-1", config)).accountId!; expect(resolveAnthropicAccountForSession("again-2", config).accountId).toBe(again); const failover = rotateAnthropicAccountOn429(config, again, "30"); expect(failover).toBeTruthy(); @@ -369,7 +500,7 @@ describe("anthropic account pool", () => { const before = getAccountSet("anthropic")!.activeAccountId; const picks = Array.from({ length: 3 }, (_, i) => resolveAnthropicAccountForSession(`promo-${i}`, config)); - expect(new Set(picks.map(p => p.accountId)).size).toBe(3); + expect(new Set(picks.map(p => p.accountId)).size).toBe(1); expect(getAccountSet("anthropic")!.activeAccountId).toBe(before); }); diff --git a/tests/adapters/key-failover.test.ts b/tests/adapters/key-failover.test.ts index 7f4a77ef55..8efb505aa6 100644 --- a/tests/adapters/key-failover.test.ts +++ b/tests/adapters/key-failover.test.ts @@ -20,7 +20,11 @@ import { } from "../../src/providers/key-failover"; import { resolveOpenCodeGoTransport } from "../../src/providers/opencode-go-transport"; import { deriveXaiConvId } from "../../src/providers/xai-transport"; -import { routeModel } from "../../src/router"; +import { routeModel, routedProviderConfig } from "../../src/router"; +import { setProviderKeychainEntryFactoryForTests } from "../../src/providers/key-store"; +import { setActiveProviderApiKey } from "../../src/providers/api-keys"; +import { subscribeAccountSelections } from "../../src/lib/account-selection-events"; +import { providerManagementConfigError, safeConfigDTO } from "../../src/server/auth-cors"; import type { OcxConfig, OcxParsedRequest, OcxProviderConfig } from "../../src/types"; import { removeTreeWithRetry } from "../helpers/remove-tree"; @@ -63,6 +67,17 @@ afterEach(() => { }); describe("hasKeyPoolFailover", () => { + test("request key identity is rejected by management and stripped from the public config", () => { + const config = makeConfig({ apiKey: "synthetic-first", apiKeyPool: [{ id: "first", key: "synthetic-first" }] }); + const routed = routedProviderConfig("p", config.providers.p); + expect(providerManagementConfigError("p", routed)).toContain("runtime field"); + const dto = JSON.stringify(safeConfigDTO({ ...config, providers: { p: { + ...routed, apiKeySelectionRevision: "internal-revision", + } } })); + expect(dto).not.toContain("_apiKeyAttempt"); + expect(dto).not.toContain("apiKeySelectionRevision"); + expect(dto).not.toContain("synthetic-first"); + }); test("true only for key-auth providers with 2+ pool entries", () => { expect(hasKeyPoolFailover({ adapter: "openai-chat", baseUrl: "x", apiKeyPool: pool3() } as OcxProviderConfig)).toBe(true); expect(hasKeyPoolFailover({ adapter: "openai-chat", baseUrl: "x", apiKeyPool: [pool3()![0]] } as OcxProviderConfig)).toBe(false); @@ -73,6 +88,70 @@ describe("hasKeyPoolFailover", () => { }); describe("rotateKeyOn429", () => { + test("an old attempt cannot overwrite a newer manual key selection or its ABA revision", () => { + const config = makeConfig({ apiKey: "key-alpha-000111222333", apiKeyPool: pool3() }); + const routed = routedProviderConfig("p", config.providers.p); + const events: string[] = []; + const unsubscribe = subscribeAccountSelections(event => { + if (event.provider === "p") events.push(loadConfig().providers.p.apiKey!); + }); + try { + expect(setActiveProviderApiKey(config, "p", "k2")).toBe(true); + expect(rotateProviderTransportOn429(config, "p", routed, { attemptedKey: routed.apiKey })?.apiKey) + .toBe("key-beta-444555666777"); + expect(events).toEqual(["key-beta-444555666777"]); + expect(setActiveProviderApiKey(config, "p", "k1")).toBe(true); + expect(rotateProviderTransportOn429(config, "p", routed, { attemptedKey: routed.apiKey })).toBeNull(); + expect(loadConfig().providers.p.apiKey).toBe("key-alpha-000111222333"); + expect(getKeyCooldownUntil("p", "k1")).toBeNull(); + expect(events).toEqual(["key-beta-444555666777", "key-alpha-000111222333"]); + } finally { unsubscribe(); } + }); + + test("manual and automatic selection events observe committed disk state", () => { + const config = makeConfig({ apiKey: "key-alpha-000111222333", apiKeyPool: pool3() }); + const events: string[] = []; + const unsubscribe = subscribeAccountSelections(event => { + if (event.provider === "p") { + expect(event.kind).toBe("api-key"); + expect(Object.keys(event).sort()).toEqual(["kind", "provider", "revision"]); + events.push(loadConfig().providers.p.apiKey!); + } + }); + try { + const routed = routedProviderConfig("p", config.providers.p); + expect(rotateProviderTransportOn429(config, "p", routed)?.apiKey).toBe("key-beta-444555666777"); + expect(events).toEqual(["key-beta-444555666777"]); + unlinkSync(getConfigPath()); + expect(rotateKeyOn429(config, "p", null)).toBeNull(); + expect(events).toHaveLength(1); + } finally { unsubscribe(); } + }); + + test.each(["env", "keychain"])("rotates a rejected %s reference instead of reusing its resolved credential", kind => { + const reference = kind === "env" ? "${OCX_SELECTION_TEST_KEY}" : "keychain:p/k1"; + process.env.OCX_SELECTION_TEST_KEY = "synthetic-resolved-first"; + setProviderKeychainEntryFactoryForTests(() => ({ + getPassword: () => "synthetic-resolved-first", + setPassword: () => {}, + deletePassword: () => true, + })); + try { + const config = makeConfig({ apiKey: reference, apiKeyPool: [ + { id: "k1", key: reference }, { id: "k2", key: "synthetic-second" }, + ] }); + const routed = routedProviderConfig("p", config.providers.p); + expect(routed.apiKey).toBe("synthetic-resolved-first"); + const rotated = rotateProviderTransportOn429(config, "p", routed, { attemptedKey: routed.apiKey }); + expect(rotated?.apiKey).toBe("synthetic-second"); + expect(loadConfig().providers.p.apiKey).toBe("synthetic-second"); + expect(getKeyCooldownUntil("p", "k1")).not.toBeNull(); + } finally { + delete process.env.OCX_SELECTION_TEST_KEY; + setProviderKeychainEntryFactoryForTests(null); + } + }); + test("rotates to the next key and cools down the exhausted one", () => { const config = makeConfig({ apiKey: "key-alpha-000111222333", apiKeyPool: pool3() }); const now = 1_000_000; diff --git a/tests/cli/cli-headless-parity.test.ts b/tests/cli/cli-headless-parity.test.ts index 44302ea884..dab85a132c 100644 --- a/tests/cli/cli-headless-parity.test.ts +++ b/tests/cli/cli-headless-parity.test.ts @@ -232,6 +232,7 @@ describe("headless GUI parity CLI", () => { // skipping the endpoint. ["/api/github/star", "(none — GUI-only)"], ["/api/oauth", "ocx account"], + ["/api/accounts/events", "(none — dashboard invalidation; ocx account reads current selection)"], ["/api/providers/keys", "ocx account"], ["/api/providers", "ocx provider"], ["/api/provider-", "ocx provider/models"], diff --git a/tests/clients/client-hub-relay.test.ts b/tests/clients/client-hub-relay.test.ts index 442bc91810..9baab569fc 100644 --- a/tests/clients/client-hub-relay.test.ts +++ b/tests/clients/client-hub-relay.test.ts @@ -1,7 +1,8 @@ -import { describe, expect, test } from "bun:test"; +import { describe, expect, spyOn, test } from "bun:test"; import { HUB_RELAY_REQUEST_BODY_MAX_BYTES, HUB_RELAY_RESPONSE_BODY_MAX_BYTES, + HUB_RELAY_DEFAULT_TIMEOUT_MS, relayHubManagementRequest, validateHubRelayRequestHeaders, } from "../../src/client/hub-relay"; @@ -144,4 +145,128 @@ describe("fixed-target hub management relay", () => { await reader.cancel(); expect(cancelled).toBe(true); }); + + test("established account SSE outlives the handshake deadline and still cancels on client abort", async () => { + const deadline = new AbortController(); + const timeout = spyOn(AbortSignal, "timeout").mockReturnValue(deadline.signal); + const browser = new AbortController(); + let upstreamSignal!: AbortSignal; + let upstreamController!: ReadableStreamDefaultController; + let cancelled = false; + let reader: ReadableStreamDefaultReader | undefined; + try { + const response = await relayHubManagementRequest(relayRequest("/api/accounts/events", { signal: browser.signal }), "/api/accounts/events", target, { + fetchImpl: (async (_input, init) => { + upstreamSignal = init!.signal!; + return new Response(new ReadableStream({ + start(controller) { upstreamController = controller; controller.enqueue(new TextEncoder().encode("event: ready\n\n")); }, + cancel() { cancelled = true; }, + }), { headers: { "Content-Type": "text/event-stream; charset=utf-8" } }); + }) as typeof fetch, + }); + expect(timeout).toHaveBeenCalledWith(HUB_RELAY_DEFAULT_TIMEOUT_MS); + reader = response.body!.getReader(); + expect(new TextDecoder().decode((await reader.read()).value)).toContain("ready"); + deadline.abort(new DOMException("Handshake deadline", "TimeoutError")); + expect(upstreamSignal.aborted).toBe(false); + expect(cancelled).toBe(false); + upstreamController.enqueue(new TextEncoder().encode("event: account-selection\n\n")); + expect(new TextDecoder().decode((await reader.read()).value)).toContain("account-selection"); + const pending = reader.read(); + browser.abort(); + await pending.catch(() => undefined); + expect(upstreamSignal.aborted).toBe(true); + expect(cancelled).toBe(true); + } finally { + await reader?.cancel().catch(() => undefined); + browser.abort(); + timeout.mockRestore(); + } + }); + + test.each([ + ["/api/config", "GET", 200, "application/json"], + ["/api/config", "GET", 200, "text/event-stream"], + ["/api/accounts/events", "POST", 200, "text/event-stream"], + ["/api/accounts/events", "GET", 201, "text/event-stream"], + ["/api/accounts/events", "GET", 401, "text/event-stream"], + ["/api/accounts/events", "GET", 200, "application/json"], + ["/api/accounts/events", "GET", 200, "text/event-streamish"], + ["/api/accounts/events?other=1", "GET", 200, "text/event-stream"], + ] as const)("relay keeps the total deadline for %s %s %i %s", async (path, method, status, contentType) => { + const deadline = new AbortController(); + const timeout = spyOn(AbortSignal, "timeout").mockReturnValue(deadline.signal); + let upstreamSignal!: AbortSignal; + let cancelled = false; + let reader: ReadableStreamDefaultReader | undefined; + try { + const response = await relayHubManagementRequest(relayRequest(path, { method }), path, target, { + fetchImpl: (async (_input, init) => { + upstreamSignal = init!.signal!; + return new Response(new ReadableStream({ cancel() { cancelled = true; } }), { + status, headers: { "Content-Type": contentType }, + }); + }) as typeof fetch, + }); + reader = response.body!.getReader(); + const pending = reader.read(); + deadline.abort(new DOMException("Total deadline", "TimeoutError")); + await pending.catch(() => undefined); + expect(upstreamSignal.aborted).toBe(true); + expect(cancelled).toBe(true); + } finally { + await reader?.cancel().catch(() => undefined); + timeout.mockRestore(); + } + }); + + test("selection SSE still has a handshake deadline and a response body cap", async () => { + const deadline = new AbortController(); + const timeout = spyOn(AbortSignal, "timeout").mockReturnValue(deadline.signal); + let handshakeStarted!: () => void; + const started = new Promise(resolve => { handshakeStarted = resolve; }); + try { + const pending = relayHubManagementRequest(relayRequest("/api/accounts/events"), "/api/accounts/events", target, { + fetchImpl: (async (_input, init) => new Promise((_resolve, reject) => { + init!.signal!.addEventListener("abort", () => reject(init!.signal!.reason), { once: true }); + handshakeStarted(); + })) as typeof fetch, + }); + await started; + deadline.abort(new DOMException("Handshake deadline", "TimeoutError")); + expect((await pending).status).toBe(502); + } finally { timeout.mockRestore(); } + + let cancelled = false; + const response = await relayHubManagementRequest(relayRequest("/api/accounts/events"), "/api/accounts/events", target, { + fetchImpl: (async () => new Response(new ReadableStream({ + start(controller) { controller.enqueue(new Uint8Array(HUB_RELAY_RESPONSE_BODY_MAX_BYTES + 1)); }, + cancel() { cancelled = true; }, + }), { headers: { "Content-Type": "text/event-stream" } })) as typeof fetch, + }); + await expect(response.arrayBuffer()).rejects.toThrow("response body too large"); + expect(cancelled).toBe(true); + }); + + test.each(["complete", "cancel"] as const)("relay detaches deadline and client listeners after body %s", async disposition => { + const deadline = new AbortController(); + const timeout = spyOn(AbortSignal, "timeout").mockReturnValue(deadline.signal); + const browser = new AbortController(); + let upstreamSignal!: AbortSignal; + try { + const response = await relayHubManagementRequest(relayRequest("/api/config", { signal: browser.signal }), "/api/config", target, { + fetchImpl: (async (_input, init) => { + upstreamSignal = init!.signal!; + return new Response(new ReadableStream({ + start(controller) { if (disposition === "complete") controller.close(); }, + }), { headers: { "Content-Type": "application/json" } }); + }) as typeof fetch, + }); + if (disposition === "complete") await response.text(); + else await response.body!.cancel(); + deadline.abort(); + browser.abort(); + expect(upstreamSignal.aborted).toBe(false); + } finally { timeout.mockRestore(); } + }); }); diff --git a/tests/gui/provider-workspace-auth.test.ts b/tests/gui/provider-workspace-auth.test.ts index edc407906b..e121142c01 100644 --- a/tests/gui/provider-workspace-auth.test.ts +++ b/tests/gui/provider-workspace-auth.test.ts @@ -82,13 +82,18 @@ async function providersPageSeam(): Promise { describe("workspace account integration seam", () => { test("passes account state and handlers into provider details", async () => { const source = await providersPageSeam(); + const page = await Bun.file("gui/src/pages/Providers.tsx").text(); + // Additional type imports must not obscure the runtime hook binding and its caller. + const poolBindings = page.match(/import\s*\{([^}]+)\}\s*from\s*["']\.\.\/hooks\/useProviderAccountPools["']/)?.[1]; + expect(poolBindings?.split(",").map(binding => binding.trim())).toContain("useProviderAccountPools"); + expect(page).toContain("const pools = useProviderAccountPools({"); expect(source).toContain("accountLoadState={accountLoadStates[item.name]"); expect(source).toContain("switchingAccountId={switchingAccount?.provider === item.name"); expect(source).toContain("onRetryAccounts: async provider => { await fetchAccountSets([provider]); }"); expect(source).toContain("key={item.name}"); expect(source).toContain("switchingAccountRef.current"); - expect(source).toContain("const refreshed = await fetchAccountSets([provider])"); - expect(source).toContain("if (!refreshed)"); + expect(source).toContain('const refreshed = await refreshAccountRosters({ provider, kind: "oauth" })'); + expect(source).toContain('if (!refreshed) { notify(t("pws.accountsLoadFailed"), false); return; }'); }); test("owns an accessible dynamic account panel instead of nesting auth in Settings", async () => { diff --git a/tests/oauth/adapter-event-oauth-failover.test.ts b/tests/oauth/adapter-event-oauth-failover.test.ts index 8309d9ffd3..a96c948ed9 100644 --- a/tests/oauth/adapter-event-oauth-failover.test.ts +++ b/tests/oauth/adapter-event-oauth-failover.test.ts @@ -4,7 +4,7 @@ import { tmpdir } from "node:os"; import { join } from "node:path"; import type { ProviderAdapter } from "../../src/adapters/base"; import { clearGenericFailoverHealth } from "../../src/oauth/generic-account-failover"; -import { saveCredential } from "../../src/oauth/store"; +import { getAccountSet, getCredential, saveCredential, setActiveAccount } from "../../src/oauth/store"; import type { AdapterEvent, OcxConfig, OcxProviderConfig } from "../../src/types"; import { removeTreeWithRetry } from "../helpers/remove-tree"; @@ -12,8 +12,12 @@ const actualResolver = await import("../../src/server/adapter-resolve"); const actualResolveAdapter = actualResolver.resolveAdapter; let attempts: AdapterEvent[][] = []; let attemptKeys: string[] = []; +let attemptProjects: Array = []; /** Set by the delivery test: an attempt that emits, then blocks before completing the turn. */ let slowAttempt: ((emit: (event: AdapterEvent) => void) => Promise) | undefined; +let beforePhysicalSend: (() => Promise) | undefined; +let physicalSends = 0; +const originalFetch = globalThis.fetch; function fixtureAdapter(provider: OcxProviderConfig): ProviderAdapter { return { @@ -22,9 +26,18 @@ function fixtureAdapter(provider: OcxProviderConfig): ProviderAdapter { async *parseStream() { yield { type: "error", message: "fixture uses runTurn" } as AdapterEvent; }, - async runTurn(_parsed, _incoming, emit) { + async runTurn(_parsed, incoming, emit) { const index = attemptKeys.length; attemptKeys.push(provider.apiKey ?? ""); + attemptProjects.push(provider.project); + const gate = beforePhysicalSend; + beforePhysicalSend = undefined; + await gate?.(); + for (let send = 0; send < physicalSends; send++) { + await incoming.providerFetch!(provider.baseUrl, { + method: "POST", headers: { Authorization: `Bearer ${provider.apiKey}` }, body: "{}", + }); + } if (slowAttempt) return await slowAttempt(emit); for (const event of attempts[index] ?? []) emit(event); }, @@ -34,7 +47,7 @@ function fixtureAdapter(provider: OcxProviderConfig): ProviderAdapter { mock.module("../../src/server/adapter-resolve", () => ({ ...actualResolver, resolveAdapter(provider: OcxProviderConfig, cacheRetention?: "none" | "short" | "long") { - if (provider.adapter === "cursor") return fixtureAdapter(provider); + if (provider.adapter === "cursor" || provider.googleMode === "cloud-code-assist") return fixtureAdapter(provider); return actualResolveAdapter(provider, cacheRetention); }, })); @@ -88,10 +101,14 @@ beforeEach(() => { clearGenericFailoverHealth(); attempts = []; attemptKeys = []; + attemptProjects = []; slowAttempt = undefined; + beforePhysicalSend = undefined; + physicalSends = 0; }); afterEach(() => { + globalThis.fetch = originalFetch; clearGenericFailoverHealth(); if (originalHome === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = originalHome; @@ -99,6 +116,72 @@ afterEach(() => { }); describe("#2568 adapter-event OAuth failover", () => { + test.each([false, true])("runTurn first physical send follows a changed selection (image loop=%s)", async imageLoop => { + await seedAccounts(2); + const accounts = getAccountSet("cursor")!.accounts; + beforePhysicalSend = async () => { await setActiveAccount("cursor", accounts[0]!.id); }; + physicalSends = 1; + slowAttempt = async emit => { emit({ type: "text_delta", text: "selected answer" }); emit({ type: "done" }); }; + const sent: string[] = []; + globalThis.fetch = (async (_input, init) => { + sent.push(new Headers(init?.headers).get("authorization") ?? ""); + return new Response("{}"); + }) as typeof fetch; + const cfg = config(false); + if (imageLoop) { + cfg.images = { bridgeEnabled: true }; + cfg.providers.xai = { adapter: "openai-chat", baseUrl: "https://api.x.ai/v1", authMode: "key", apiKey: "synthetic-image-key" }; + } + const req = imageLoop ? new Request("http://localhost/v1/responses", { + method: "POST", headers: { "content-type": "application/json" }, + body: JSON.stringify({ model: "cursor/model", input: "answer", stream: true, tools: [{ type: "image_generation" }] }), + }) : request(true); + const response = await handleResponses(req, cfg, { model: "", provider: "" }); + expect(await response.text()).toContain("selected answer"); + expect(sent).toEqual(["Bearer cursor-access-0"]); + }); + + test("an already started multi-message turn keeps its original credential", async () => { + await seedAccounts(2); + const accounts = getAccountSet("cursor")!.accounts; + physicalSends = 2; + slowAttempt = async emit => { emit({ type: "text_delta", text: "same turn" }); emit({ type: "done" }); }; + const sent: string[] = []; + globalThis.fetch = (async (_input, init) => { + sent.push(new Headers(init?.headers).get("authorization") ?? ""); + if (sent.length === 1) await setActiveAccount("cursor", accounts[0]!.id); + return new Response("{}"); + }) as typeof fetch; + const response = await handleResponses(request(false), config(false), { model: "", provider: "" }); + expect(await response.text()).toContain("same turn"); + expect(sent).toEqual(["Bearer cursor-access-1", "Bearer cursor-access-1"]); + expect(getCredential("cursor")?.access).toBe("cursor-access-0"); + }); + + test("every CCA request pairs the persisted active account with its own project", async () => { + for (const id of ["a", "b"]) await saveCredential("google-antigravity", { + access: `ga-access-${id}`, refresh: `ga-refresh-${id}`, expires: Date.now() + 3_600_000, + accountId: id, projectId: `project-${id}`, + }); + const cfg = config(); + cfg.defaultProvider = "google-antigravity"; + cfg.providers = { "google-antigravity": { ...cfg.providers.cursor!, googleMode: "cloud-code-assist", project: "project-a" } }; + attempts = [ + [{ type: "text_delta", text: "first" }, { type: "done" }], + [{ type: "text_delta", text: "second" }, { type: "done" }], + ]; + for (let i = 0; i < 2; i++) { + const req = new Request("http://localhost/v1/responses", { + method: "POST", headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ model: "google-antigravity/model", input: "answer", stream: false }), + }); + const res = await handleResponses(req, cfg, { model: "", provider: "" }); + const responseBody = await res.text(); + expect(res.status, responseBody).toBe(200); + } + expect(attemptKeys).toEqual(["ga-access-b", "ga-access-b"]); + expect(attemptProjects).toEqual(["project-b", "project-b"]); + }); for (const stream of [true, false]) { test(`${stream ? "streaming" : "non-streaming"} first-event 429 rotates and replays`, async () => { await seedAccounts(2); @@ -113,9 +196,28 @@ describe("#2568 adapter-event OAuth failover", () => { expect(attemptKeys).toEqual(["cursor-access-1", "cursor-access-0"]); expect(body).toContain("alternate answer"); expect(body).not.toContain("Cursor rate limit exceeded"); + expect(getCredential("cursor")?.access).toBe("cursor-access-0"); }); } + test("a newer manual choice wins a pending request's 429 proposal", async () => { + await seedAccounts(3); + const accounts = getAccountSet("cursor")!.accounts; + slowAttempt = async emit => { + if (attemptKeys.length === 1) { + await setActiveAccount("cursor", accounts[1]!.id); + emit({ type: "error", message: "Cursor rate limit exceeded: resource_exhausted" }); + } else { + emit({ type: "text_delta", text: "manual choice answered" }); + emit({ type: "done" }); + } + }; + const response = await handleResponses(request(false), config(false), { model: "", provider: "" }); + expect(await response.text()).toContain("manual choice answered"); + expect(attemptKeys).toEqual(["cursor-access-2", "cursor-access-1"]); + expect(getCredential("cursor")?.access).toBe("cursor-access-1"); + }); + test("a single account is a strict no-op", async () => { await seedAccounts(1); attempts = [[{ type: "error", message: "Cursor rate limit exceeded: resource_exhausted" }]]; diff --git a/tests/oauth/generic-oauth-failover.test.ts b/tests/oauth/generic-oauth-failover.test.ts index 3761d65ce3..48fb2782ec 100644 --- a/tests/oauth/generic-oauth-failover.test.ts +++ b/tests/oauth/generic-oauth-failover.test.ts @@ -71,6 +71,45 @@ async function seed(count: number, offset = 0): Promise { } describe("#2568 generic OAuth account failover", () => { + for (const provider of ["xai", "cursor", "kimi", "github-copilot", "google-antigravity", "nous", "kiro", "meta-muse"]) { + test(`manual selection owns healthy dispatch for ${provider}, with pool off or on`, async () => { + for (const accountId of ["selected", "spare"]) { + await saveCredential(provider, { + access: `synthetic-${accountId}`, refresh: `refresh-${accountId}`, + expires: Date.now() + 3_600_000, accountId, + }); + } + const ids = getAccountSet(provider)!.accounts.map(a => a.id); + await setActiveAccount(provider, ids[0]!); + setCachedProviderAccountQuotaForTests(provider, ids[0]!, { weeklyPercent: 30, updatedAt: Date.now() }); + setCachedProviderAccountQuotaForTests(provider, ids[1]!, { weeklyPercent: 11, updatedAt: Date.now() }); + for (const enabled of [undefined, false, true]) { + const cfg = { providers: { [provider]: { ...OAUTH_PROVIDER, + ...(enabled === undefined ? {} : { oauthAccountFailover: { enabled } }), + } } } as OcxConfig; + expect(preferredInitialAccount(cfg, provider)).toBeNull(); + } + clearAccountQuotaCache(provider); + }); + } + + test("proactive exhaustion avoidance requires explicit pool enablement", async () => { + const [selected, spare] = await seed(2); + await setActiveAccount("xai", selected!); + setCachedProviderAccountQuotaForTests("xai", selected!, { weeklyPercent: 100, updatedAt: Date.now() }); + setCachedProviderAccountQuotaForTests("xai", spare!, { weeklyPercent: 11, updatedAt: Date.now() }); + expect(preferredInitialAccount(config(), "xai")).toBeNull(); + expect(preferredInitialAccount(config(false), "xai")).toBeNull(); + expect(preferredInitialAccount(config(true), "xai")).toBe(spare); + }); + + test("unknown selected quota is not permission to replace the account", async () => { + const [selected, spare] = await seed(2); + await setActiveAccount("xai", selected!); + setCachedProviderAccountQuotaForTests("xai", spare!, { weeklyPercent: 11, updatedAt: Date.now() }); + expect(preferredInitialAccount(config(true), "xai")).toBeNull(); + }); + test("two logged-in accounts rotate with NO configuration at all (#2568d)", async () => { // The reported workflow: three xAI accounts are logged in, the active one hits its limit, and // the operator never went looking for a toggle. Presence is the consent signal. @@ -145,7 +184,7 @@ describe("#2568 generic OAuth account failover", () => { const ids = await seed(2); await setActiveAccount("xai", ids[0]!); clearGenericFailoverHealth("xai"); - setCachedProviderAccountQuotaForTests("xai", ids[0]!, { fiveHourPercent: 99 }); + setCachedProviderAccountQuotaForTests("xai", ids[0]!, { fiveHourPercent: 100 }); setCachedProviderAccountQuotaForTests("xai", ids[1]!, { fiveHourPercent: 1 }); expect(preferredInitialAccount(config(false, true), "xai")).toBe(ids[1]); @@ -390,8 +429,11 @@ describe("sidecar on429 wiring", () => { // to the configured account's project — #2841 in its original shape. const start = coreSource.indexOf("const preferredAccountId ="); expect(start).toBeGreaterThan(-1); - const region = coreSource.slice(start, start + 6000); - expect(region).toContain("usedPreferredAccount && resolved.projectId"); + const end = coreSource.indexOf("\n route.provider = resolveProviderTransport(", start); + expect(end).toBeGreaterThan(start); + const region = coreSource.slice(start, end); + expect(region).toContain("project: resolved.projectId"); + expect(region).not.toContain("!route.provider.project"); // A project-less preferred account falls BACK to the ordinary active-account resolution // rather than erroring: a preference must never turn a working request into a failure, // and Antigravity tolerates project discovery failing, so an account with no project is diff --git a/tests/oauth/oauth-accounts-api.test.ts b/tests/oauth/oauth-accounts-api.test.ts index cbeb165e59..75e88d7dbb 100644 --- a/tests/oauth/oauth-accounts-api.test.ts +++ b/tests/oauth/oauth-accounts-api.test.ts @@ -1,4 +1,4 @@ -import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { afterEach, beforeEach, describe, expect, spyOn, test } from "bun:test"; import { managementFetch as fetch } from "../helpers/management-auth"; import { mkdtempSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; @@ -14,6 +14,35 @@ import { withStubbedProviderFetch } from "../helpers/catalog-provider-fetch"; import { getAccountSet } from "../../src/oauth/store"; import { ACCOUNT_IMPORT_DEADLINE_MS, ACCOUNT_IMPORT_MAX_BYTES, ACCOUNT_IMPORT_MAX_REQUEST_BYTES } from "../../src/oauth/account-import/types"; import { handleOauthAccountRoutes } from "../../src/server/management/oauth-account-routes"; +import { createManagementSessionControl, requireManagementAuth, type ManagementAuthState } from "../../src/server/management-auth"; +import { handleSessionRoutes } from "../../src/server/management/session-routes"; +import type { ManagementContext } from "../../src/server/management/context"; +import { publishAccountSelection } from "../../src/lib/account-selection-events"; + +function selectionSessionFixture() { + const origin = "http://127.0.0.1:10100"; + const token = "ocx_session_selection_liveness_test"; + const state: Extract = { + available: true, token: "ocx_admin_selection_test", source: "environment", + sessions: new Map([[token, { + serverOrigin: origin, browserOrigin: origin, csrfToken: "selection-csrf", + expiresAt: Date.now() + 60_000, issuance: "loopback", + }]]), pairingGrants: new Map(), + }; + const req = new Request(`${origin}/api/accounts/events`, { headers: { + Host: "127.0.0.1:10100", Origin: origin, "x-opencodex-gui-origin": origin, + "x-opencodex-api-key": token, "x-opencodex-csrf-token": "selection-csrf", + } }); + const ctx: ManagementContext = { + req, url: new URL(req.url), config: baseConfig(), deps: {}, version: "test", + principal: "gui-session", sessionControl: createManagementSessionControl(state), + convergeCodexCatalog: async () => ({ status: "failed", reason: "disk" }), + syncClaudeAgentDefsBestEffort: async () => {}, + }; + // Cache this Request's original admission: subsequent stream checks must ignore it. + expect(requireManagementAuth(req, state, ctx.config)).toBeNull(); + return { ctx, state, token }; +} let testDir = ""; let previousHome: string | undefined; @@ -66,6 +95,142 @@ afterEach(() => { }); describe("multiauth accounts API", () => { + test("selection events require management authentication", async () => { + const server = startServer(0); + try { + const response = await originalFetch(new URL("/api/accounts/events", server.url)); + expect(response.status).toBe(401); + await response.body?.cancel(); + } finally { await server.stop(true); } + }); + + test("selection event streams bound subscribers and release cancelled connections", async () => { + const { accountSelectionStream } = await import("../../src/server/management/account-selection-stream"); + const streams: Response[] = []; + try { + for (let i = 0; i < 64; i++) { + const response = accountSelectionStream(new Request("http://localhost/api/accounts/events"), () => true); + expect(response.status).toBe(200); + streams.push(response); + } + expect(accountSelectionStream(new Request("http://localhost/api/accounts/events"), () => true).status).toBe(429); + } finally { + await Promise.all(streams.map(response => response.body!.cancel())); + } + const response = accountSelectionStream(new Request("http://localhost/api/accounts/events"), () => true); + expect(response.status).toBe(200); + await response.body!.cancel(); + }); + + test("selection event route denies admission without a current session validator", async () => { + const { ctx } = selectionSessionFixture(); + const response = await handleOauthAccountRoutes({ ...ctx, sessionControl: undefined }); + try { expect(response?.status).toBe(401); } + finally { await response?.body?.cancel(); } + }); + + test.each(["false", "throw"] as const)("selection stream denies an initial validator result of %s", async result => { + const { accountSelectionStream } = await import("../../src/server/management/account-selection-stream"); + const response = accountSelectionStream(new Request("http://localhost/api/accounts/events"), () => { + if (result === "throw") throw new Error("validator unavailable"); + return false; + }); + try { expect(response.status).toBe(401); } + finally { await response.body?.cancel(); } + }); + + test.each(["logout", "expiry"] as const)("selection stream stops publishing after GUI %s despite cached request admission", async change => { + const { ctx, state, token } = selectionSessionFixture(); + const response = await handleOauthAccountRoutes(ctx); + expect(response?.status).toBe(200); + const reader = response!.body!.getReader(); + try { + expect(new TextDecoder().decode((await reader.read()).value)).toContain("event: ready"); + if (change === "logout") { + const req = new Request(new URL("/api/session/logout", ctx.req.url), { method: "POST", headers: ctx.req.headers }); + expect(requireManagementAuth(req, state, ctx.config)).toBeNull(); + expect(handleSessionRoutes({ ...ctx, req, url: new URL(req.url) })?.status).toBe(200); + } else { + state.sessions.get(token)!.expiresAt = Date.now() - 1; + } + expect(requireManagementAuth(ctx.req, state, ctx.config)).toBeNull(); // Deliberately memoized. + const pending = reader.read(); + publishAccountSelection("private-provider", "oauth"); + await expect(pending).rejects.toMatchObject({ name: "NotAllowedError" }); + } finally { await reader.cancel().catch(() => undefined); } + }); + + test("selection heartbeat revalidates expiry without extending a remote session", async () => { + const { ctx, state, token } = selectionSessionFixture(); + const session = state.sessions.get(token)!; + session.issuance = "pairing"; + const expiresAt = session.expiresAt; + const interval = spyOn(globalThis, "setInterval"); + let reader: ReadableStreamDefaultReader | undefined; + try { + const response = await handleOauthAccountRoutes(ctx); + reader = response!.body!.getReader(); + await reader.read(); + const tick = interval.mock.calls.find(call => call[1] === 15_000)?.[0]; + if (typeof tick !== "function") throw new Error("selection heartbeat not registered"); + tick(); + expect(new TextDecoder().decode((await reader.read()).value)).toContain(": heartbeat"); + expect(session.expiresAt).toBe(expiresAt); + session.expiresAt = Date.now() - 1; + const pending = reader.read(); + tick(); + await expect(pending).rejects.toMatchObject({ name: "NotAllowedError" }); + } finally { + await reader?.cancel().catch(() => undefined); + interval.mockRestore(); + } + }); + + test("management session liveness rereads revoked sessions and the current admin token", () => { + const { ctx, state, token } = selectionSessionFixture(); + const control = ctx.sessionControl!; + expect(control.isCurrent(ctx.req, ctx.config)).toBe(true); + expect(control.revokeCurrent(ctx.req)).toBe(true); + expect(control.isCurrent(ctx.req, ctx.config)).toBe(false); + expect(state.sessions.has(token)).toBe(false); + const adminReq = new Request(ctx.req.url, { headers: { "x-opencodex-api-key": state.token } }); + expect(requireManagementAuth(adminReq, state, ctx.config)).toBeNull(); + expect(control.isCurrent(adminReq, ctx.config)).toBe(true); + state.token = "ocx_admin_rotated_selection_test"; + expect(control.isCurrent(adminReq, ctx.config)).toBe(false); + expect(createManagementSessionControl({ available: false, reason: "test" }).isCurrent(adminReq, ctx.config)).toBe(false); + }); + + test("selection events notify only after the new active account is committed", async () => { + const server = startServer(0); + const abort = new AbortController(); + let reader: ReadableStreamDefaultReader | undefined; + try { + const response = await fetch(new URL("/api/accounts/events", server.url), { signal: abort.signal }); + expect(response.status).toBe(200); + expect(response.headers.get("content-type")).toContain("text/event-stream"); + reader = response.body!.getReader(); + const ready = new TextDecoder().decode((await reader.read()).value); + expect(ready).toContain("event: ready"); + const eventRead = reader.read(); + const selected = await fetch(new URL("/api/oauth/accounts/active", server.url), { + method: "PUT", headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ provider: "anthropic", accountId: "bbbb2222" }), + }); + expect(selected.status).toBe(200); + const notification = new TextDecoder().decode((await eventRead).value); + expect(notification).toContain("event: account-selection"); + expect(notification).toContain('"provider":"anthropic"'); + expect(notification).not.toContain("bbbb2222"); + expect(notification).not.toContain("t2"); + expect(getAccountSet("anthropic")?.activeAccountId).toBe("bbbb2222"); + } finally { + await reader?.cancel(); + abort.abort(); + await server.stop(true); + } + }); + test("GET lists masked accounts with active flag", async () => { const server = startServer(0); try { diff --git a/tests/oauth/oauth-store-multi.test.ts b/tests/oauth/oauth-store-multi.test.ts index 58791494de..6cd3f21dbe 100644 --- a/tests/oauth/oauth-store-multi.test.ts +++ b/tests/oauth/oauth-store-multi.test.ts @@ -1,7 +1,9 @@ -import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { afterEach, beforeEach, describe, expect, spyOn, test } from "bun:test"; import { INTERNAL_DEADLINE_MS, STORE_BUDGET_MS } from "../helpers/test-budget"; import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs"; import { join } from "node:path"; +import * as atomicWrite from "../../src/config/atomic-write"; +import * as oauthStore from "../../src/oauth/store"; import { resetHardenedStateForTests, setIcaclsRunnerForTests, @@ -14,16 +16,19 @@ import { listAccounts, markAccountNeedsReauth, markAccountNeedsReauthIfGeneration, + mergeAccountCredential, mutateStore, OAuthMutationBusyError, oauthMutationTailSnapshot, reconcileOAuthReauthState, removeAccount, removeCredential, + replaceProviderAccountSet, saveAccountCredential, saveCredential, setAccountAlias, setActiveAccount, + upsertCredentialByIdentity, } from "../../src/oauth/store"; import type { OAuthCredentials } from "../../src/oauth/types"; import { removeTreeWithRetry } from "../helpers/remove-tree"; @@ -38,6 +43,17 @@ const cred = (over: Partial = {}): OAuthCredentials => ({ ...over, }); +const SELECTION_UUID = /^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/i; + +async function selectionAccounts() { + await saveCredential("xai", cred({ accountId: "selection-a" })); + const idA = getAccountSet("xai")!.activeAccountId; + await saveCredential("xai", cred({ accountId: "selection-b", access: "access-b" })); + const idB = getAccountSet("xai")!.activeAccountId; + await setActiveAccount("xai", idA); + return { idA, idB }; +} + describe("multi-account auth store", () => { beforeEach(() => { previousOpencodexHome = process.env.OPENCODEX_HOME; @@ -306,6 +322,264 @@ describe("multi-account auth store", () => { expect(set.activeAccountId).toBe("ok"); // dangling active healed }); + test("selection revision rejects an automatic promotion after manual A-B-A", async () => { + const { idA, idB } = await selectionAccounts(); + const before = getAccountSet("xai")!.selectionRevision; + expect(before).toMatch(SELECTION_UUID); + const expectedSelection = oauthStore.captureOAuthAccountSelection("xai")!; + await setActiveAccount("xai", idB); + await setActiveAccount("xai", idA); + expect(getAccountSet("xai")!.selectionRevision).not.toBe(before); + expect(await oauthStore.commitOAuthAccountSelection("xai", idB, { expectedSelection })).toBeNull(); + expect(oauthStore.captureOAuthAccountSelection("xai")?.accountId).toBe(idA); + }); + + test("selection revision preserves credential-only refresh and unrelated account metadata", async () => { + const { idA, idB } = await selectionAccounts(); + // Seed a persisted revision independently to catch normalization dropping it. + const authPath = join(TEST_DIR, "auth.json"); + const raw = JSON.parse(readFileSync(authPath, "utf8")); + const revision = "f4abbddc-5c7c-4e87-bd8a-b5775a182860"; + raw.xai.selectionRevision = revision; + writeFileSync(authPath, JSON.stringify(raw)); + await saveAccountCredential("xai", idA, cred({ accountId: "selection-a", access: "refreshed-a" })); + expect(getAccountSet("xai")!.selectionRevision).toBe(revision); + await mergeAccountCredential("xai", idB, cred({ accountId: "selection-b", access: "refreshed-b" })); + await setAccountAlias("xai", idA, "Selection test"); + await markAccountNeedsReauth("xai", idB, true); + await upsertCredentialByIdentity("xai", cred({ accountId: "selection-a", access: "import-refreshed" })); + expect(getAccountSet("xai")!.selectionRevision).toBe(revision); + expect(JSON.parse(readFileSync(authPath, "utf8")).xai.selectionRevision).toBe(revision); + expect(oauthStore.captureOAuthAccountSelection("xai")).toEqual({ accountId: idA, revision }); + }); + + test("selection revision advances for removal, recreation, and rollback replacement", async () => { + const { idA, idB } = await selectionAccounts(); + const original = getAccountSet("xai")!; + expect(original.selectionRevision).toMatch(SELECTION_UUID); + const expectedSelection = oauthStore.captureOAuthAccountSelection("xai")!; + await removeAccount("xai", idA); + expect(getAccountSet("xai")!.activeAccountId).toBe(idB); + const promoted = getAccountSet("xai")!.selectionRevision; + expect(promoted).not.toBe(original.selectionRevision); + await removeCredential("xai"); + expect(oauthStore.captureOAuthAccountSelection("xai")).toBeNull(); + await saveCredential("xai", cred({ accountId: "selection-a" })); + const recreated = getAccountSet("xai")!; + expect(recreated.activeAccountId).toBe(idA); + expect(recreated.selectionRevision).not.toBe(original.selectionRevision); + await replaceProviderAccountSet("xai", original); + const restored = getAccountSet("xai")!; + expect(restored.selectionRevision).toMatch(SELECTION_UUID); + expect([original.selectionRevision, promoted, recreated.selectionRevision]).not.toContain(restored.selectionRevision); + expect(original.selectionRevision).toBe(expectedSelection.revision); + expect(await oauthStore.commitOAuthAccountSelection("xai", idB, { expectedSelection })).toBeNull(); + }); + + test("selection revision advances on same-id manual reselect but not automatic validation", async () => { + const { idA } = await selectionAccounts(); + const before = getAccountSet("xai")!.selectionRevision; + await setActiveAccount("xai", idA); + const after = getAccountSet("xai")!.selectionRevision; + expect(after).not.toBe(before); + expect(after).toMatch(SELECTION_UUID); + const expectedSelection = oauthStore.captureOAuthAccountSelection("xai")!; + expect(await oauthStore.commitOAuthAccountSelection("xai", idA, { + expectedSelection, + expectedCredentialGeneration: credentialGeneration(getAccountCredential("xai", idA)!), + requireUsableAccount: true, + })).toEqual(expectedSelection); + expect(oauthStore.captureOAuthAccountSelection("xai")).toEqual(expectedSelection); + }); + + test("selection commit supports revisionless legacy snapshots and guards the original id", async () => { + const authPath = join(TEST_DIR, "auth.json"); + writeFileSync(authPath, JSON.stringify({ xai: { + activeAccountId: "legacy-a", + accounts: [{ id: "legacy-a", credential: cred() }, { id: "legacy-b", credential: cred({ access: "b" }) }], + } })); + const expectedSelection = oauthStore.captureOAuthAccountSelection("xai")!; + expect(expectedSelection).toEqual({ accountId: "legacy-a" }); + expect(await oauthStore.commitOAuthAccountSelection("xai", "legacy-b", { + expectedSelection: { accountId: "wrong-id" }, + })).toBeNull(); + expect(await oauthStore.commitOAuthAccountSelection("xai", "legacy-a", { expectedSelection })).toEqual(expectedSelection); + const committed = await oauthStore.commitOAuthAccountSelection("xai", "legacy-b", { expectedSelection }); + expect(committed?.accountId).toBe("legacy-b"); + expect(committed?.revision).toMatch(SELECTION_UUID); + expect(oauthStore.captureOAuthAccountSelection("xai")).toEqual(committed); + }); + + test.each(["manual", "refresh", "reauth", "remove"] as const)("selection commit rechecks queued %s changes under the writer", async change => { + const { idA, idB } = await selectionAccounts(); + const expectedSelection = oauthStore.captureOAuthAccountSelection("xai")!; + const expectedCredentialGeneration = credentialGeneration(getAccountCredential("xai", idB)!); + let release!: () => void; + let entered!: () => void; + const gate = new Promise(resolve => { release = resolve; }); + const started = new Promise(resolve => { entered = resolve; }); + const blocker = mutateStore(async () => { entered(); await gate; }); + await started; + const mutation = change === "manual" ? setActiveAccount("xai", idA) + : change === "refresh" ? saveAccountCredential("xai", idB, cred({ accountId: "selection-b", access: "fresh-b" })) + : change === "reauth" ? markAccountNeedsReauth("xai", idB, true) + : removeAccount("xai", idB); + const pending = oauthStore.commitOAuthAccountSelection("xai", idB, { + expectedSelection, expectedCredentialGeneration, requireUsableAccount: true, + }); + try { + release(); + await blocker; + await mutation; + expect(await pending).toBeNull(); + expect(getAccountSet("xai")!.activeAccountId).toBe(idA); + } finally { + release(); + await Promise.allSettled([blocker, mutation, pending]); + } + }); + + test("selection commit checks same-account usability and refreshed credential generation", async () => { + const { idA, idB } = await selectionAccounts(); + const expectedSelection = oauthStore.captureOAuthAccountSelection("xai")!; + const oldGeneration = credentialGeneration(getAccountCredential("xai", idA)!); + await saveAccountCredential("xai", idA, cred({ accountId: "selection-a", access: "rotated-a" })); + expect(await oauthStore.commitOAuthAccountSelection("xai", idA, { + expectedSelection, expectedCredentialGeneration: oldGeneration, requireUsableAccount: true, + })).toBeNull(); + await markAccountNeedsReauth("xai", idA, true); + expect(await oauthStore.commitOAuthAccountSelection("xai", idA, { + expectedSelection, requireUsableAccount: true, + })).toBeNull(); + const committed = await oauthStore.commitOAuthAccountSelection("xai", idB, { + expectedSelection, + expectedCredentialGeneration: credentialGeneration(getAccountCredential("xai", idB)!), + requireUsableAccount: true, + }); + expect(committed?.accountId).toBe(idB); + expect(committed?.revision).not.toBe(expectedSelection.revision); + expect(oauthStore.captureOAuthAccountSelection("xai")).toEqual(committed); + }); + + test("unchanged selection admission neither joins a busy writer nor persists", async () => { + const { idA } = await selectionAccounts(); + const expectedSelection = oauthStore.captureOAuthAccountSelection("xai")!; + const expectedCredentialGeneration = credentialGeneration(getAccountCredential("xai", idA)!); + let release!: () => void; + let entered!: () => void; + const gate = new Promise(resolve => { release = resolve; }); + const started = new Promise(resolve => { entered = resolve; }); + const blocker = mutateStore(async () => { entered(); await gate; }); + await started; + const write = spyOn(atomicWrite, "atomicWriteFile"); + const pending = oauthStore.commitOAuthAccountSelection("xai", idA, { + expectedSelection, expectedCredentialGeneration, requireUsableAccount: true, + }); + try { + expect(oauthMutationTailSnapshot().active).toBe(1); + expect(await pending).toEqual(expectedSelection); + expect(write).not.toHaveBeenCalled(); + } finally { + write.mockRestore(); + release(); + await Promise.allSettled([blocker, pending]); + } + }); + + test("selection events follow persistence and omit failed commits, refreshes, and credentials", async () => { + const { idA, idB } = await selectionAccounts(); + const { subscribeAccountSelections, currentAccountSelectionRevision } = await import("../../src/lib/account-selection-events"); + const events: unknown[] = []; + const observedSelections: unknown[] = []; + const start = currentAccountSelectionRevision(); + const unsubscribe = subscribeAccountSelections(event => { + events.push(event); + observedSelections.push(oauthStore.captureOAuthAccountSelection("xai")); + }); + try { + const expectedSelection = oauthStore.captureOAuthAccountSelection("xai")!; + await setActiveAccount("xai", idA); + const manual = oauthStore.captureOAuthAccountSelection("xai")!; + expect(events).toEqual([{ provider: "xai", kind: "oauth", revision: start + 1 }]); + expect(observedSelections).toEqual([manual]); + expect(await oauthStore.commitOAuthAccountSelection("xai", idB, { expectedSelection })).toBeNull(); + expect(await oauthStore.commitOAuthAccountSelection("xai", "missing")).toBeNull(); + expect(await setActiveAccount("xai", "missing")).toBe(false); + await oauthStore.commitOAuthAccountSelection("xai", idA, { expectedSelection: manual }); + await saveAccountCredential("xai", idA, cred({ accountId: "selection-a", access: "event-refresh" })); + expect(events).toHaveLength(1); + + // Only the I/O boundary is faulted; the actual commit, locks, and store stay real. + const write = spyOn(atomicWrite, "atomicWriteFile").mockImplementation(() => { throw new Error("selection persist failed"); }); + try { + await expect(oauthStore.commitOAuthAccountSelection("xai", idB, { expectedSelection: manual })).rejects.toThrow("selection persist failed"); + } finally { + write.mockRestore(); + } + expect(oauthStore.captureOAuthAccountSelection("xai")).toEqual(manual); + expect(events).toHaveLength(1); + expect(currentAccountSelectionRevision()).toBe(start + 1); + await expect(saveCredential("xai", cred({ accountId: "blocked-login" }), { + assertBeforePersist: () => { throw new Error("selection pre-persist rejected"); }, + })).rejects.toThrow("selection pre-persist rejected"); + expect(events).toHaveLength(1); + + await oauthStore.commitOAuthAccountSelection("xai", idB, { expectedSelection: manual }); + expect(events).toEqual([ + { provider: "xai", kind: "oauth", revision: start + 1 }, + { provider: "xai", kind: "oauth", revision: start + 2 }, + ]); + unsubscribe(); + unsubscribe(); + await setActiveAccount("xai", idA); + expect(events).toHaveLength(2); + } finally { + unsubscribe(); + } + }); + + test("selection events cover create, inactive removal, replacement, clear, and recreate", async () => { + const { subscribeAccountSelections, currentAccountSelectionRevision, publishAccountSelection } = await import("../../src/lib/account-selection-events"); + const events: unknown[] = []; + const start = currentAccountSelectionRevision(); + const unsubscribe = subscribeAccountSelections(event => { events.push(event); }); + try { + await upsertCredentialByIdentity("xai", cred({ accountId: "selection-a" })); + const original = getAccountSet("xai")!; + await upsertCredentialByIdentity("xai", cred({ accountId: "selection-b" })); + expect(events).toHaveLength(1); // Importing an inactive account preserves the selection. + const inactive = listAccounts("xai").find(account => account.id !== original.activeAccountId)!; + await removeAccount("xai", inactive.id); + expect(getAccountSet("xai")!.selectionRevision).not.toBe(original.selectionRevision); + await replaceProviderAccountSet("xai", original); + await replaceProviderAccountSet("xai", null); + await replaceProviderAccountSet("xai", null); + await saveCredential("xai", cred({ accountId: "selection-a" })); + publishAccountSelection("key-provider", "api-key"); + expect(events).toEqual([ + ...Array.from({ length: 5 }, (_, index) => ({ provider: "xai", kind: "oauth", revision: start + index + 1 })), + { provider: "key-provider", kind: "api-key", revision: start + 6 }, + ]); + } finally { + unsubscribe(); + } + }); + + test("selection subscriber failure cannot fail a persisted selection or block other subscribers", async () => { + const { idA } = await selectionAccounts(); + const { subscribeAccountSelections } = await import("../../src/lib/account-selection-events"); + const stopBroken = subscribeAccountSelections(() => { throw new Error("disconnected consumer"); }); + const seen: unknown[] = []; + const stopHealthy = subscribeAccountSelections(event => { seen.push(event); }); + try { + expect(await setActiveAccount("xai", idA)).toBe(true); + expect(seen).toHaveLength(1); + } finally { + stopBroken(); + stopHealthy(); + } + }); + test("queued generation-checked reauth mutation rechecks liveness after reconciliation", async () => { await saveCredential("xai", cred({ email: "race@example.com", accountId: "race-account" })); const accountId = getAccountSet("xai")!.activeAccountId; diff --git a/tests/oauth/oauth-upsert-preserves-api-key.test.ts b/tests/oauth/oauth-upsert-preserves-api-key.test.ts index caec3363b6..a22239dd69 100644 --- a/tests/oauth/oauth-upsert-preserves-api-key.test.ts +++ b/tests/oauth/oauth-upsert-preserves-api-key.test.ts @@ -3,6 +3,7 @@ import { mkdtempSync} from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { OAUTH_PROVIDERS, upsertOAuthProvider } from "../../src/oauth"; +import { loadConfig, saveConfig } from "../../src/config"; import { migrateXaiResponsesDefault } from "../../src/providers/xai-responses-opt-in"; import { resolveWireProtocolOverride } from "../../src/server/adapter-resolve"; import { @@ -240,8 +241,13 @@ describe("upsertOAuthProvider credential preservation", () => { expect(listed.keys.find(entry => entry.id === activeId)?.active).toBe(true); expect(listed.keys.find(entry => entry.id === "pool-visible")?.active).toBe(false); + // runLogin persists the upsert before GUI key mutations. The shared selection + // transaction requires that authoritative file; it must not recreate missing config. + saveConfig(config); + expect(loadConfig().providers.xai!.apiKeyPool).toEqual(provider.apiKeyPool); expect(setActiveProviderApiKey(config, "xai", "pool-visible")).toBe(true); expect(config.providers.xai!.apiKey).toBe("pool-visible-key"); + expect(loadConfig().providers.xai!.apiKey).toBe("pool-visible-key"); expect(listProviderApiKeys(config, "xai").activeId).toBe("pool-visible"); expect(setActiveProviderApiKey(config, "xai", activeId)).toBe(true); @@ -250,6 +256,7 @@ describe("upsertOAuthProvider credential preservation", () => { expect(config.providers.xai!.apiKey).toBe("pool-visible-key"); expect(config.providers.xai!.apiKeyPool).toEqual([{ id: "pool-visible", key: "pool-visible-key" }]); expect(listProviderApiKeys(config, "xai").activeId).toBe("pool-visible"); + expect(loadConfig().providers.xai!.apiKeyPool).toEqual(config.providers.xai!.apiKeyPool); } finally { if (previousHome === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = previousHome; @@ -329,16 +336,21 @@ describe("upsertOAuthProvider credential preservation", () => { const testHome = mkdtempSync(join(tmpdir(), "ocx-oauth-upsert-")); process.env.OPENCODEX_HOME = testHome; try { + saveConfig(config); expect(removeProviderApiKey(config, "xai", "aaaaaaaa")).toBe(true); expect(config.providers.xai!.authMode).toBe("key"); expect(config.providers.xai!.apiKey).toBeUndefined(); expect(config.providers.xai!.apiKeyPool).toBeUndefined(); + expect(loadConfig().providers.xai!.apiKey).toBeUndefined(); + expect(loadConfig().providers.xai!.apiKeyPool).toBeUndefined(); upsertOAuthProvider(config, "xai"); const provider = config.providers.xai!; expect(provider.authMode).toBe("oauth"); expect(provider.apiKey).toBeUndefined(); expect(provider.apiKeyPool).toBeUndefined(); + saveConfig(config); + expect(loadConfig().providers.xai!.authMode).toBe("oauth"); } finally { if (previousHome === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = previousHome; diff --git a/tests/providers/github-copilot/github-copilot-account-origin.test.ts b/tests/providers/github-copilot/github-copilot-account-origin.test.ts index 5f464a1d7e..4da89c816e 100644 --- a/tests/providers/github-copilot/github-copilot-account-origin.test.ts +++ b/tests/providers/github-copilot/github-copilot-account-origin.test.ts @@ -1,10 +1,12 @@ -import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test"; import { mkdtempSync} from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { clearGenericFailoverHealth } from "../../../src/oauth/generic-account-failover"; import { getAccountSet, saveCredential, setActiveAccount } from "../../../src/oauth/store"; import { handleResponses } from "../../../src/server/responses"; +import { saveConfig } from "../../../src/config"; +import { setActiveProviderApiKey } from "../../../src/providers/api-keys"; import type { OcxConfig } from "../../../src/types"; import { removeTreeWithRetry } from "../../helpers/remove-tree"; @@ -16,6 +18,37 @@ const GITHUB_USER_URL = "https://api.github.com/user"; const originalFetch = globalThis.fetch; const originalHome = process.env.OPENCODEX_HOME; let home = ""; +let beforeBuildReturns: (() => Promise) | undefined; +let beforePacingReturns: (() => Promise) | undefined; +const actualPacing = await import("../../../src/providers/request-pacing"); +const originalWaitForSlot = actualPacing.waitForProviderRequestSlot; +mock.module("../../../src/providers/request-pacing", () => ({ + ...actualPacing, + waitForProviderRequestSlot: async (...args: Parameters) => { + const result = await originalWaitForSlot(...args); + const gate = beforePacingReturns; + beforePacingReturns = undefined; + await gate?.(); + return result; + }, +})); +const actualAdapterResolver = await import("../../../src/server/adapter-resolve"); +const originalResolveAdapter = actualAdapterResolver.resolveAdapter; +mock.module("../../../src/server/adapter-resolve", () => ({ + ...actualAdapterResolver, + resolveAdapter: (...args: Parameters) => { + const adapter = originalResolveAdapter(...args); + const build = adapter.buildRequest.bind(adapter); + adapter.buildRequest = async (...buildArgs) => { + const built = await build(...buildArgs); + const gate = beforeBuildReturns; + beforeBuildReturns = undefined; + await gate?.(); + return built; + }; + return adapter; + }, +})); type Wire = "chat" | "responses"; @@ -40,12 +73,12 @@ function config(wire: Wire): OcxConfig { } as OcxConfig; } -function request(wire: Wire): Request { +function request(wire: Wire, extra: Record = {}): Request { const model = wire === "chat" ? "gpt-4o" : "gpt-5.4"; return new Request("http://localhost/v1/responses", { method: "POST", headers: { "content-type": "application/json" }, - body: JSON.stringify({ model: `github-copilot/${model}`, input: "hello", stream: false }), + body: JSON.stringify({ model: `github-copilot/${model}`, input: "hello", stream: false, ...extra }), }); } @@ -104,6 +137,7 @@ function installFetch(options: { statuses: number[]; switchToAccountId?: string; switchOn: "refresh" | "first-dispatch" | "never"; + emptyFirst?: boolean; }): { dispatches: { origin: string; authorization: string }[] } { const dispatches: { origin: string; authorization: string }[] = []; let refreshSwitched = false; @@ -138,6 +172,14 @@ function installFetch(options: { headers: status === 429 ? { "retry-after": "1" } : undefined, }); } + if (options.emptyFirst && dispatches.length === 1) { + return Response.json({ choices: [{ index: 0, message: { role: "assistant", content: "" }, finish_reason: "stop" }] }); + } + if (JSON.parse(String(init?.body ?? "{}")).stream === true && options.wire === "chat") { + return new Response('data: {"choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":null}]}\n\ndata: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n', { + headers: { "content-type": "text/event-stream" }, + }); + } return successResponse(options.wire); } return originalFetch(input, init); @@ -146,6 +188,8 @@ function installFetch(options: { } beforeEach(() => { + beforeBuildReturns = undefined; + beforePacingReturns = undefined; home = mkdtempSync(join(tmpdir(), "ocx-copilot-origin-")); process.env.OPENCODEX_HOME = home; clearGenericFailoverHealth(); @@ -160,8 +204,154 @@ afterEach(() => { }); describe("GitHub Copilot bearer/origin snapshot atomicity", () => { + test.each(["chat", "responses", "image", "web-search"] as const)("%s API-key dispatch uses a selection committed during pacing", async path => { + const cfg = config("chat"); + cfg.defaultProvider = "fixture"; + cfg.providers = { fixture: { + adapter: path === "responses" ? "openai-responses" : "openai-chat", authMode: "key", + baseUrl: "https://fixture.invalid/v1", apiKey: "synthetic-a", models: ["model"], + apiKeyPool: [{ id: "a", key: "synthetic-a" }, { id: "b", key: "synthetic-b" }], + } }; + if (path === "image") { + cfg.images = { bridgeEnabled: true }; + cfg.providers.xai = { adapter: "openai-chat", authMode: "key", apiKey: "synthetic-image-key", baseUrl: "https://api.x.ai/v1" }; + } else if (path === "web-search") cfg.webSearchSidecar = { enabled: true, backend: "exa", exaApiKey: "synthetic-search-key" }; + saveConfig(cfg); + beforePacingReturns = async () => { expect(setActiveProviderApiKey(cfg, "fixture", "b")).toBe(true); }; + const sent: string[] = []; + globalThis.fetch = (async (_input, init) => { + sent.push(new Headers(init?.headers).get("authorization") ?? ""); + if (path === "image" || path === "web-search") return new Response('data: {"choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n', { headers: { "content-type": "text/event-stream" } }); + return successResponse(path); + }) as typeof fetch; + const response = await handleResponses(request("chat", { + model: "fixture/model", stream: path === "image" || path === "web-search", + ...(path === "image" || path === "web-search" ? { tools: [{ type: path === "image" ? "image_generation" : "web_search" }] } : {}), + }), cfg, { model: "", provider: "" }); + expect(await response.text()).toContain("ok"); + expect(sent).toEqual(["Bearer synthetic-b"]); + }); + + test("a pacing switch to B keeps B when Anthropic rebuilds an image after 413", async () => { + for (const id of ["a", "b"]) await saveCredential("anthropic", { + access: `synthetic-anthropic-${id}`, refresh: `synthetic-refresh-${id}`, + expires: Date.now() + 3_600_000, accountId: id, + }); + const rows = getAccountSet("anthropic")!.accounts; + await setActiveAccount("anthropic", rows[0]!.id); + beforePacingReturns = async () => { await setActiveAccount("anthropic", rows[1]!.id); }; + const sent: string[] = []; + globalThis.fetch = (async (_input, init) => { + sent.push(new Headers(init?.headers).get("authorization") ?? ""); + if (sent.length === 1) return Response.json({ error: { type: "request_too_large", message: "too large" } }, { status: 413 }); + return Response.json({ id: "message-selection", type: "message", role: "assistant", + content: [{ type: "text", text: "ok" }], stop_reason: "end_turn", usage: { input_tokens: 1, output_tokens: 1 } }); + }) as typeof fetch; + const cfg = config("chat"); + cfg.providers = { anthropic: { adapter: "anthropic", authMode: "oauth", baseUrl: "https://api.anthropic.com", models: ["claude-fable-5"] } }; + cfg.defaultProvider = "anthropic"; + const response = await handleResponses(request("chat", { + model: "anthropic/claude-fable-5", + input: [{ role: "user", content: [ + { type: "input_text", text: "look" }, + { type: "input_image", image_url: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==" }, + ] }], + }), cfg, { model: "", provider: "" }); + expect(await response.text()).toContain("ok"); + expect(sent).toEqual(["Bearer synthetic-anthropic-b", "Bearer synthetic-anthropic-b"]); + }); + + test("a pacing switch to B keeps B for an empty-completion continuation", async () => { + const accounts = await seedAccounts(); + beforePacingReturns = async () => { await setActiveAccount("github-copilot", accounts.b); }; + const observed = installFetch({ wire: "chat", statuses: [200, 200], switchOn: "never", emptyFirst: true }); + const cfg = config("chat"); + cfg.emptyCompletionRetry = true; + const response = await handleResponses(request("chat"), cfg, { model: "", provider: "" }); + expect(await response.text()).toContain("ok"); + expect(observed.dispatches).toEqual([ + { origin: ACCOUNT_B_ORIGIN, authorization: bearer("copilot-access-b") }, + { origin: ACCOUNT_B_ORIGIN, authorization: bearer("copilot-access-b") }, + ]); + }); + + test.each(["image", "web-search"] as const)("%s main-model dispatch follows a manual choice made during pacing", async path => { + const accounts = await seedAccounts(); + beforePacingReturns = async () => { await setActiveAccount("github-copilot", accounts.b); }; + const observed = installFetch({ wire: "chat", statuses: [200], switchOn: "never" }); + const cfg = config("chat"); + if (path === "image") { + cfg.images = { bridgeEnabled: true }; + cfg.providers.xai = { adapter: "openai-chat", authMode: "key", apiKey: "synthetic-image-key", baseUrl: "https://api.x.ai/v1" }; + } else { + cfg.webSearchSidecar = { enabled: true, backend: "exa", exaApiKey: "synthetic-search-key" }; + } + const response = await handleResponses(request("chat", { + stream: true, tools: [{ type: path === "image" ? "image_generation" : "web_search" }], + }), cfg, { model: "", provider: "" }); + expect(await response.text()).toContain("ok"); + expect(observed.dispatches).toEqual([{ origin: ACCOUNT_B_ORIGIN, authorization: bearer("copilot-access-b") }]); + }); + + test("a cached search-loop adapter cannot bless A wire with B's current snapshot", async () => { + const accounts = await seedAccounts(); + beforePacingReturns = async () => { await setActiveAccount("github-copilot", accounts.b); }; + const sent: string[] = []; + const bodies: Array<{ messages: Array<{ role: string }> }> = []; + globalThis.fetch = (async (input, init) => { + const url = new URL(input instanceof Request ? input.url : String(input)); + if (url.hostname === "api.exa.ai") return Response.json({ results: [] }); + if (!url.hostname.endsWith(".githubcopilot.com")) throw new Error("Unexpected fixture request"); + sent.push(new Headers(init?.headers).get("authorization") ?? ""); + bodies.push(JSON.parse(String(init?.body))); + const delta = sent.length === 1 + ? { tool_calls: [{ index: 0, id: "search-1", type: "function", function: { name: "web_search", arguments: '{"query":"fixture"}' } }] } + : { content: "ok after search" }; + return new Response(`data: ${JSON.stringify({ choices: [{ index: 0, delta, finish_reason: sent.length === 1 ? "tool_calls" : "stop" }] })}\n\ndata: [DONE]\n\n`, { + headers: { "content-type": "text/event-stream" }, + }); + }) as typeof fetch; + const cfg = config("chat"); + cfg.webSearchSidecar = { enabled: true, backend: "exa", exaApiKey: "synthetic-search-key" }; + const response = await handleResponses(request("chat", { stream: true, tools: [{ type: "web_search" }] }), cfg, { model: "", provider: "" }); + expect(await response.text()).toContain("ok after search"); + expect(sent).toEqual(["Bearer copilot-access-b", "Bearer copilot-access-b"]); + expect(bodies[1]!.messages.some(message => message.role === "tool")).toBe(true); + }); + + test("a key removed during pacing is never dispatched", async () => { + const cfg = config("chat"); + cfg.defaultProvider = "fixture"; + cfg.providers = { fixture: { adapter: "openai-chat", authMode: "key", baseUrl: "https://fixture.invalid/v1", apiKey: "synthetic-a" } }; + beforePacingReturns = async () => { delete cfg.providers.fixture; }; + let sends = 0; + globalThis.fetch = (async () => { sends++; return successResponse("chat"); }) as typeof fetch; + const response = await handleResponses(request("chat", { model: "fixture/model" }), cfg, { model: "", provider: "" }); + await response.text(); + expect(response.status).not.toBe(200); + expect(sends).toBe(0); + }); + for (const wire of ["chat", "responses"] as const) { - test(`${wire} initial refresh keeps account A's origin after B becomes active`, async () => { + test(`${wire} revalidates selection after pacing and before physical dispatch`, async () => { + const accounts = await seedAccounts(); + beforePacingReturns = async () => { await setActiveAccount("github-copilot", accounts.b); }; + const observed = installFetch({ wire, statuses: [200], switchOn: "never" }); + const response = await handleResponses(request(wire), config(wire), { model: "", provider: "" }); + await response.text(); + expect(response.status).toBe(200); + expect(observed.dispatches).toEqual([{ origin: ACCOUNT_B_ORIGIN, authorization: bearer("copilot-access-b") }]); + }); + test(`${wire} rebuilds after a manual selection during asynchronous request building`, async () => { + const accounts = await seedAccounts(); + beforeBuildReturns = async () => { await setActiveAccount("github-copilot", accounts.b); }; + const observed = installFetch({ wire, statuses: [200], switchOn: "never" }); + const response = await handleResponses(request(wire), config(wire), { model: "", provider: "" }); + await response.text(); + expect(response.status).toBe(200); + expect(observed.dispatches).toEqual([{ origin: ACCOUNT_B_ORIGIN, authorization: bearer("copilot-access-b") }]); + }); + test(`${wire} initial admission follows a newer manual selection with its matching origin`, async () => { const accounts = await seedAccounts(0); const observed = installFetch({ wire, @@ -175,12 +365,12 @@ describe("GitHub Copilot bearer/origin snapshot atomicity", () => { expect(response.status).toBe(200); expect(observed.dispatches).toEqual([{ - origin: ACCOUNT_A_ORIGIN, - authorization: bearer("copilot-access-a-refreshed"), + origin: ACCOUNT_B_ORIGIN, + authorization: bearer("copilot-access-b"), }]); }); - test(`${wire} 401 replay keeps refreshed account A's origin after B becomes active`, async () => { + test(`${wire} 401 replay follows a newer manual selection with its matching origin`, async () => { const accounts = await seedAccounts(); const observed = installFetch({ wire, @@ -195,7 +385,7 @@ describe("GitHub Copilot bearer/origin snapshot atomicity", () => { expect(response.status).toBe(200); expect(observed.dispatches).toEqual([ { origin: ACCOUNT_A_ORIGIN, authorization: bearer("copilot-access-a") }, - { origin: ACCOUNT_A_ORIGIN, authorization: bearer("copilot-access-a-refreshed") }, + { origin: ACCOUNT_B_ORIGIN, authorization: bearer("copilot-access-b") }, ]); }); } diff --git a/tests/providers/kiro/kiro-pool-rank.test.ts b/tests/providers/kiro/kiro-pool-rank.test.ts index b42eb969d0..554972704a 100644 --- a/tests/providers/kiro/kiro-pool-rank.test.ts +++ b/tests/providers/kiro/kiro-pool-rank.test.ts @@ -143,7 +143,7 @@ describe("pre-dispatch account preference", () => { authMode: "oauth", } as unknown as OcxProviderConfig; - const config = { providers: { xai: OAUTH_PROVIDER } } as unknown as OcxConfig; + const config = { providers: { xai: OAUTH_PROVIDER }, oauthAccountFailover: { enabled: true } } as unknown as OcxConfig; const originalHome = process.env.OPENCODEX_HOME; let home: string; @@ -159,7 +159,7 @@ describe("pre-dispatch account preference", () => { return getAccountSet(providerName)?.accounts.map(a => a.id) ?? []; } - test("the account with more headroom is chosen before the first request", async () => { + test("an enabled pool avoids a known-exhausted selected account", async () => { home = mkdtempSync(join(tmpdir(), "ocx-predispatch-")); process.env.OPENCODEX_HOME = home; clearGenericFailoverHealth(); @@ -167,7 +167,7 @@ describe("pre-dispatch account preference", () => { try { const ids = await seedAccounts(2); await setActiveAccount("xai", ids[0]!); - setCachedProviderAccountQuotaForTests("xai", ids[0]!, { monthlyPercent: 95, updatedAt: Date.now() }); + setCachedProviderAccountQuotaForTests("xai", ids[0]!, { monthlyPercent: 100, updatedAt: Date.now() }); setCachedProviderAccountQuotaForTests("xai", ids[1]!, { monthlyPercent: 5, updatedAt: Date.now() }); expect(preferredInitialAccount(config, "xai")).toBe(ids[1]); } finally { @@ -365,14 +365,7 @@ describe("pre-dispatch account preference", () => { } }); - test("neither a redirecting nor a non-redirecting selection touches the credential store", async () => { - // loadAuthStore chmods the config dir, chmods the secret, and re-parses the whole - // credential file on every call — and this runs on the initial resolution of EVERY - // request. The steady state of this feature is a pool where one account consistently - // ranks higher, so the REDIRECTING path must be cached too — validating the winner here - // would put a second uncached read in front of every such request. Deleting the store - // proves it: an uncached path could not answer at all. Staleness is caught at - // resolution instead, inside a store read the resolver already performs. + test("removing the credential store invalidates an earlier selection proposal", async () => { home = mkdtempSync(join(tmpdir(), "ocx-predispatch-")); process.env.OPENCODEX_HOME = home; clearGenericFailoverHealth(); @@ -381,11 +374,11 @@ describe("pre-dispatch account preference", () => { const ids = await seedAccounts(2); await setActiveAccount("xai", ids[0]!); // Redirecting: the other account holds more headroom on every call. - setCachedProviderAccountQuotaForTests("xai", ids[0]!, { monthlyPercent: 95, updatedAt: Date.now() }); + setCachedProviderAccountQuotaForTests("xai", ids[0]!, { monthlyPercent: 100, updatedAt: Date.now() }); setCachedProviderAccountQuotaForTests("xai", ids[1]!, { monthlyPercent: 5, updatedAt: Date.now() }); expect(preferredInitialAccount(config, "xai")).toBe(ids[1]); rmSync(join(home, "auth.json"), { force: true }); - for (let i = 0; i < 4; i++) expect(preferredInitialAccount(config, "xai")).toBe(ids[1]); + for (let i = 0; i < 4; i++) expect(preferredInitialAccount(config, "xai")).toBeNull(); // Non-redirecting: the active account already ranks best. setCachedProviderAccountQuotaForTests("xai", ids[0]!, { monthlyPercent: 5, updatedAt: Date.now() }); @@ -412,15 +405,14 @@ describe("pre-dispatch account preference", () => { try { const ids = await seedAccounts(2); await setActiveAccount("xai", ids[0]!); - setCachedProviderAccountQuotaForTests("xai", ids[0]!, { monthlyPercent: 95, updatedAt: Date.now() }); + setCachedProviderAccountQuotaForTests("xai", ids[0]!, { monthlyPercent: 100, updatedAt: Date.now() }); setCachedProviderAccountQuotaForTests("xai", ids[1]!, { monthlyPercent: 5, updatedAt: Date.now() }); expect(preferredInitialAccount(config, "xai")).toBe(ids[1]); await removeAccount("xai", ids[1]!); - // Selection is a cached PREFERENCE, so it may still name the removed account... - expect(preferredInitialAccount(config, "xai")).toBe(ids[1]); - // ...and resolution is where that is caught. The request path absorbs this throw and - // falls back to the active account. + // Selection reads the authoritative roster; a removed target is never proposed. + expect(preferredInitialAccount(config, "xai")).toBeNull(); + // The credential resolver independently rejects the removed identity. await expect( getValidAccessSnapshotForAccount("xai", ids[1]!, { requireUsableAccount: true }), ).rejects.toThrow(); @@ -446,11 +438,12 @@ describe("pre-dispatch account preference", () => { try { const ids = await seedAccounts(2); await setActiveAccount("xai", ids[0]!); - setCachedProviderAccountQuotaForTests("xai", ids[0]!, { monthlyPercent: 95, updatedAt: Date.now() }); + setCachedProviderAccountQuotaForTests("xai", ids[0]!, { monthlyPercent: 100, updatedAt: Date.now() }); setCachedProviderAccountQuotaForTests("xai", ids[1]!, { monthlyPercent: 5, updatedAt: Date.now() }); expect(preferredInitialAccount(config, "xai")).toBe(ids[1]); await markAccountNeedsReauth("xai", ids[1]!, true); + expect(preferredInitialAccount(config, "xai")).toBeNull(); // An ordinary resolve SUCCEEDS — the credential is still readable — which is exactly // why the flag must be checked inside the resolver rather than trusted to throw. await expect(getValidAccessSnapshotForAccount("xai", ids[1]!)).resolves.toBeDefined(); diff --git a/tests/server/server-google-antigravity-oauth-401-replay.test.ts b/tests/server/server-google-antigravity-oauth-401-replay.test.ts index 18c0162f9a..b1f55c238c 100644 --- a/tests/server/server-google-antigravity-oauth-401-replay.test.ts +++ b/tests/server/server-google-antigravity-oauth-401-replay.test.ts @@ -353,7 +353,7 @@ describe("Google Antigravity OAuth upstream 401 replay", () => { } }); - test.each([false, true])("401 stays pinned to rejected account A after active switches to B (newer A generation=%s)", async newerGeneration => { + test.each([false, true])("401 recovery follows the newly selected account and its project (newer A generation=%s)", async newerGeneration => { await seedOAuth(); const accountA = getAccountSet("google-antigravity")!.activeAccountId; const config = antigravityConfig(); @@ -382,17 +382,17 @@ describe("Google Antigravity OAuth upstream 401 replay", () => { const response = await postResponses(server); expect(response.status).toBe(200); expect(await response.text()).toContain("ok after google refresh"); - expect(observed.counts.refresh).toBe(newerGeneration ? 0 : 1); - expect(observed.chatAuth).toEqual(["Bearer rejected-access", newerGeneration ? "Bearer newer-access-a" : "Bearer fresh-access"]); - expect(observed.chatProjects).toEqual(["initial-project-id", newerGeneration ? "newer-project-a" : "refreshed-project-a"]); + expect(observed.counts.refresh).toBe(0); + expect(observed.chatAuth).toEqual(["Bearer rejected-access", "Bearer access-b"]); + expect(observed.chatProjects).toEqual(["initial-project-id", "project-b"]); const accounts = getAccountSet("google-antigravity")!; expect(accounts.activeAccountId).not.toBe(accountA); expect(accounts.accounts.find(account => account.id === accounts.activeAccountId)?.credential).toMatchObject({ access: "access-b", projectId: "project-b", }); expect(accounts.accounts.find(account => account.id === accountA)?.credential).toMatchObject({ - access: newerGeneration ? "newer-access-a" : "fresh-access", - projectId: newerGeneration ? "newer-project-a" : "refreshed-project-a", + access: newerGeneration ? "newer-access-a" : "rejected-access", + projectId: newerGeneration ? "newer-project-a" : "initial-project-id", }); } finally { await server.stop(true); @@ -594,45 +594,45 @@ describe("Google Antigravity OAuth upstream 401 replay", () => { } }); - test("negative project-less refresh rejects replay in native Responses passthrough", async () => { + test("project-less account is refused before dispatch in native Responses passthrough", async () => { await seedOAuth(undefined, null); saveConfig(antigravityPassthroughConfig()); const observed = installOAuthFetch([401], { refreshedProjectId: null }); const server = startServer(0); try { const response = await postResponses(server); - expect(observed.requestPaths).toEqual(["/v1/responses"]); + expect(observed.requestPaths).toEqual([]); const json = await response.json() as { error?: { code?: string; message?: string; type?: string } }; expect(response.status).toBe(401); expect(json.error?.type).toBe("authentication_error"); expect(json.error?.message).toBe(PUBLIC_OAUTH_AUTHENTICATION_ERROR); - expect(observed.counts.refresh).toBe(1); - expect(observed.chatAuth).toEqual(["Bearer rejected-access"]); + expect(observed.counts.refresh).toBe(0); + expect(observed.chatAuth).toEqual([]); } finally { await server.stop(true); } }); - test("negative project-less refresh rejects replay in generic adapter", async () => { + test("project-less account is refused before dispatch in generic adapter", async () => { await seedOAuth(undefined, null); saveConfig(antigravityConfig()); const observed = installOAuthFetch([401], { refreshedProjectId: null }); const server = startServer(0); try { const response = await postResponses(server); - expect(observed.requestPaths).toEqual(["/v1internal:generateContent"]); + expect(observed.requestPaths).toEqual([]); const json = await response.json() as { error?: { code?: string; message?: string; type?: string } }; expect(response.status).toBe(401); expect(json.error?.type).toBe("authentication_error"); expect(json.error?.message).toBe(PUBLIC_OAUTH_AUTHENTICATION_ERROR); - expect(observed.counts.refresh).toBe(1); - expect(observed.chatAuth).toEqual(["Bearer rejected-access"]); + expect(observed.counts.refresh).toBe(0); + expect(observed.chatAuth).toEqual([]); } finally { await server.stop(true); } }); - test("negative project-less refresh rejects replay in chat completions", async () => { + test("project-less account is refused before dispatch in chat completions", async () => { await seedOAuth(undefined, null); saveConfig(antigravityConfig()); const observed = installOAuthFetch([401], { refreshedProjectId: null }); @@ -643,8 +643,8 @@ describe("Google Antigravity OAuth upstream 401 replay", () => { expect(response.status).toBe(401); expect(json.error?.type).toBe("authentication_error"); expect(json.error?.message).toBe(PUBLIC_OAUTH_AUTHENTICATION_ERROR); - expect(observed.counts.refresh).toBe(1); - expect(observed.chatAuth).toEqual(["Bearer rejected-access"]); + expect(observed.counts.refresh).toBe(0); + expect(observed.chatAuth).toEqual([]); } finally { await server.stop(true); } diff --git a/tests/server/server-key-failover-e2e.test.ts b/tests/server/server-key-failover-e2e.test.ts index 57ab87b2cb..d23bf848dd 100644 --- a/tests/server/server-key-failover-e2e.test.ts +++ b/tests/server/server-key-failover-e2e.test.ts @@ -2,7 +2,7 @@ import { afterEach, beforeEach, describe, expect, test } from "bun:test"; import { mkdtempSync} from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; -import { saveConfig } from "../../src/config"; +import { loadConfig, saveConfig } from "../../src/config"; import { clearKeyCooldowns } from "../../src/providers/key-failover"; import { deriveXaiConvId } from "../../src/providers/xai-transport"; import { clearReasoningReplayCacheForTests } from "../../src/responses/reasoning-replay-cache"; @@ -10,6 +10,11 @@ import { startServer } from "../../src/server"; import type { OcxConfig } from "../../src/types"; import { installIsolatedCodexHome, type IsolatedCodexHome } from "../helpers/isolated-codex-home"; import { removeTreeWithRetry } from "../helpers/remove-tree"; +import { managementFetch } from "../helpers/management-auth"; +import { resetProviderRequestPacingForTest, setProviderRequestPacingRuntimeForTest, waitForProviderRequestSlot } from "../../src/providers/request-pacing"; +import { providerApiKeySelectionIsCurrent, resolveCurrentProviderApiKeyTransport } from "../../src/providers/api-key-selection"; +import { routedProviderConfig } from "../../src/router"; +import type { OcxProviderTransport } from "../../src/providers/xai-transport"; let testDir = ""; let previousHome: string | undefined; @@ -38,6 +43,137 @@ afterEach(() => { }); describe("server 429 key failover (end-to-end)", () => { + test("physical key selection rejects disabled, removed, and changed-auth providers", () => { + const provider = { adapter: "openai-chat", baseUrl: "https://example.test/v1", authMode: "key", apiKey: "synthetic-first" } as const; + const config = { providers: { current: { ...provider } } } as unknown as OcxConfig; + const routed = routedProviderConfig("current", config.providers.current); + expect(providerApiKeySelectionIsCurrent(config, "current", routed)).toBe(true); + // A model-level wire override does not change which key was selected. + expect(providerApiKeySelectionIsCurrent(config, "current", { ...routed, adapter: "openai-responses" })).toBe(true); + for (const replacement of [{ ...provider, disabled: true }, { ...provider, authMode: "oauth" }, { ...provider, apiKey: undefined }]) { + config.providers.current = replacement as OcxConfig["providers"][string]; + expect(providerApiKeySelectionIsCurrent(config, "current", routed)).toBe(false); + expect(resolveCurrentProviderApiKeyTransport(config, "current", routed)).toBeNull(); + } + delete config.providers.current; + expect(providerApiKeySelectionIsCurrent(config, "current", routed)).toBe(false); + expect(resolveCurrentProviderApiKeyTransport(config, "current", routed)).toBeNull(); + }); + + test("physical transport refresh keeps its executor and affinity but takes current static headers", () => { + const config = { providers: { current: { + adapter: "openai-chat", baseUrl: "https://example.test/v1", authMode: "key", apiKey: "synthetic-first", + headers: { "x-old-static": "old" }, apiKeySelectionRevision: "first-revision", + } } } as unknown as OcxConfig; + const executor = (async () => Response.json({})) as typeof fetch; + const routed: OcxProviderTransport = { + ...routedProviderConfig("current", config.providers.current), fetch: executor, + headers: { "x-old-static": "old", "x-opencode-session": "runtime-session" }, + }; + config.providers.current = { ...config.providers.current, apiKey: "synthetic-second", + apiKeySelectionRevision: "second-revision", headers: { "x-new-static": "new" }, + }; + expect(providerApiKeySelectionIsCurrent(config, "current", routed)).toBe(false); + const current = resolveCurrentProviderApiKeyTransport(config, "current", routed) as OcxProviderTransport; + expect(current.apiKey).toBe("synthetic-second"); + expect(current.fetch).toBe(executor); + expect(current.headers).toEqual({ "x-new-static": "new", "x-opencode-session": "runtime-session" }); + expect(providerApiKeySelectionIsCurrent(config, "current", current)).toBe(true); + }); + + test("native Chat rebuilds a queued request after a manual key selection during pacing", async () => { + let now = 0; + let resumePacing: (() => void) | undefined; + const queued = Promise.withResolvers(); + setProviderRequestPacingRuntimeForTest({ + now: () => now, + setTimer(callback, delayMs) { + resumePacing = () => { now += delayMs; callback(); }; + queued.resolve(); + return callback; + }, + clearTimer() {}, + enqueueMicrotask: queueMicrotask, + }); + const seen: Headers[] = []; + upstream = Bun.serve({ hostname: "127.0.0.1", port: 0, fetch(req) { + seen.push(new Headers(req.headers)); + return Response.json({ id: "chatcmpl-paced", object: "chat.completion", + choices: [{ index: 0, message: { role: "assistant", content: "current selection" }, finish_reason: "stop" }], + }); + } }); + const config = { port: 0, hostname: "127.0.0.1", defaultProvider: "paced", providers: { paced: { + adapter: "openai-chat", baseUrl: `http://127.0.0.1:${upstream.port}/v1`, allowPrivateNetwork: true, + authMode: "key", apiKey: "synthetic-first", headers: { "x-static-test": "retained" }, + apiKeyPool: [{ id: "first", key: "synthetic-first" }, { id: "second", key: "synthetic-second" }], + requestPacing: { enabled: true, minIntervalMs: 100 }, + } } } as OcxConfig; + saveConfig(config); + const server = startServer(0); + const abort = new AbortController(); + try { + await waitForProviderRequestSlot("paced", config.providers.paced); + const pending = fetch(new URL("/v1/chat/completions", server.url), { + method: "POST", headers: { "content-type": "application/json" }, signal: abort.signal, + body: JSON.stringify({ model: "paced/test", stream: false, messages: [{ role: "user", content: "hello" }] }), + }); + await queued.promise; + expect(seen).toHaveLength(0); + const selected = await managementFetch(new URL("/api/providers/keys/active", server.url), { + method: "PUT", headers: { "content-type": "application/json" }, + body: JSON.stringify({ name: "paced", id: "second" }), + }); + expect(selected.status).toBe(200); + await selected.text(); + resumePacing!(); + const response = await pending; + expect(response.status).toBe(200); + expect(await response.text()).toContain("current selection"); + expect(seen.map(headers => headers.get("authorization"))).toEqual(["Bearer synthetic-second"]); + expect(seen[0]!.get("x-static-test")).toBe("retained"); + } finally { + abort.abort(); + await server.stop(true); + resetProviderRequestPacingForTest(); + } + }); + + test.each(["responses", "chat/completions"])("%s carries the configured env-key identity through 429 recovery", async inbound => { + const seen: string[] = []; + process.env.OCX_SELECTION_E2E_KEY = "synthetic-env-first"; + upstream = Bun.serve({ hostname: "127.0.0.1", port: 0, fetch(req) { + seen.push(req.headers.get("authorization") ?? ""); + if (seen.length === 1) return Response.json({ error: { message: "rate limited" } }, { status: 429 }); + return Response.json({ id: "chatcmpl-env", object: "chat.completion", + choices: [{ index: 0, message: { role: "assistant", content: "recovered" }, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 1 }, + }); + } }); + saveConfig({ port: 0, hostname: "127.0.0.1", defaultProvider: "pooled", providers: { pooled: { + adapter: "openai-chat", baseUrl: `http://127.0.0.1:${upstream.port}/v1`, allowPrivateNetwork: true, + authMode: "key", apiKey: "${OCX_SELECTION_E2E_KEY}", apiKeyPool: [ + { id: "first", key: "${OCX_SELECTION_E2E_KEY}" }, { id: "second", key: "synthetic-second" }, + ], + } } } as OcxConfig); + const server = startServer(0); + try { + const response = await fetch(new URL(`/v1/${inbound}`, server.url), { + method: "POST", headers: { "content-type": "application/json" }, + body: JSON.stringify({ model: "pooled/test", stream: false, + ...(inbound === "responses" ? { input: "hello" } : { messages: [{ role: "user", content: "hello" }] }), + }), + }); + expect(response.status).toBe(200); + expect(await response.text()).toContain("recovered"); + expect(seen).toEqual(["Bearer synthetic-env-first", "Bearer synthetic-second"]); + expect(loadConfig().providers.pooled.apiKey).toBe("synthetic-second"); + expect(loadConfig().providers.pooled._apiKeyAttempt).toBeUndefined(); + } finally { + await server.stop(true); + delete process.env.OCX_SELECTION_E2E_KEY; + } + }); + test("xAI API-key rotation preserves cache affinity and never adds OAuth CLI headers", async () => { const originalFetch = globalThis.fetch; const promptCacheKey = "codex-session-high-entropy-429-e2e";