From c04985c1ee65ab32050a7d198087105998f65e00 Mon Sep 17 00:00:00 2001 From: Rahul Sethuram Date: Thu, 27 Aug 2026 16:46:31 +0400 Subject: [PATCH] feat: spill Codex traffic to a fallback API key --- docs/GROK_MODELS.md | 4 + docs/HOW_IT_WORKS.md | 4 +- docs/OPENAI_MODELS.md | 59 +++ docs/SECRETS.md | 22 +- .../routes/model-provider-accounts.test.ts | 23 + .../src/routes/model-provider-accounts.ts | 21 +- .../src/sandbox_runtime/opencode_server.py | 22 +- .../plugins/codex-auth-plugin.js | 354 ++++++++++++- .../plugins/provider-token-broker.js | 22 +- .../tests/codex-auth-plugin.test.mjs | 498 +++++++++++++++++- .../tests/provider-token-broker.test.mjs | 18 +- .../tests/test_openai_oauth_setup.py | 28 + 12 files changed, 1026 insertions(+), 49 deletions(-) create mode 100644 packages/control-plane/src/routes/model-provider-accounts.test.ts diff --git a/docs/GROK_MODELS.md b/docs/GROK_MODELS.md index 8890655ff0..7bfcf95a6f 100644 --- a/docs/GROK_MODELS.md +++ b/docs/GROK_MODELS.md @@ -86,6 +86,10 @@ If the resolved secrets contain no legacy xAI refresh token, sandbox preparation The sandbox never receives the refresh token. Broker responses use `Cache-Control: no-store`, and the endpoint rejects user and service credentials in favor of the matching session's sandbox token. +A session that can see an `XAI_API_KEY` secret skips this path entirely: no marker, no sentinel, no +plugin. Grok calls then go straight to xAI's metered API with that key. Delete the secret to return +to the SuperGrok subscription. + --- ## Deployment and Rollout diff --git a/docs/HOW_IT_WORKS.md b/docs/HOW_IT_WORKS.md index 82ccb5a8e6..28292bb28d 100644 --- a/docs/HOW_IT_WORKS.md +++ b/docs/HOW_IT_WORKS.md @@ -631,7 +631,9 @@ scoped OAuth. New sessions use an explicit choice, then a provider-account defau retain legacy scoped OAuth or API-key behavior. Setting a default affects only future sessions; operators may remove legacy keys after legacy-bound sessions are no longer needed. See [Using OpenAI Models](./OPENAI_MODELS.md) and -[Using Grok with a SuperGrok Subscription](./GROK_MODELS.md). +[Using Grok with a SuperGrok Subscription](./GROK_MODELS.md). A plain provider API key visible to +the session (`OPENAI_API_KEY`, `XAI_API_KEY`) disables broker mode for that provider and is injected +into the sandbox like any other secret, because OpenCode reads those variables directly. > **LLM API keys** (e.g., `ANTHROPIC_API_KEY` for Claude models) are added as global secrets. A > deployment can instead configure `anthropic_api_key` in Terraform to inject one fleet-wide key diff --git a/docs/OPENAI_MODELS.md b/docs/OPENAI_MODELS.md index 1cc7b57bd5..7f43c92308 100644 --- a/docs/OPENAI_MODELS.md +++ b/docs/OPENAI_MODELS.md @@ -66,6 +66,58 @@ choice for future runs. --- +## Spilling over before the subscription runs out + +A ChatGPT subscription that hits its Codex quota fails the session outright: +`Execution failed: The usage limit has been reached...`. Two optional secrets let a deployment keep +working on a per-token key, and cap how much of the subscription sandboxes may take in the first +place: + +| Secret Name | Value | +| --------------------------------- | ------------------------------------------------------------------------- | +| `OPENAI_API_KEY_FALLBACK` | A platform API key, used only as a spillover | +| `OPENAI_SUBSCRIPTION_MAX_PERCENT` | Optional share of a rate-limit window sandboxes may consume (default 100) | + +`OPENAI_API_KEY_FALLBACK` is deliberately a separate name from `OPENAI_API_KEY`: the latter selects +metered billing for the whole session and is stripped from sessions routed to a subscription, while +this one rides along unused until the subscription cannot answer. Set +`OPENAI_SUBSCRIPTION_MAX_PERCENT` to `80` to reserve the last fifth of each window for whoever else +uses that ChatGPT account. + +A sandbox sends OpenAI traffic to the subscription until one of these happens, then latches to a +successful fallback path for the rest of its life: + +- when `OPENAI_SUBSCRIPTION_MAX_PERCENT` is below `100`, usage is already at or above that ceiling + before the first turn. The percentage comes from `GET /backend-api/wham/usage`, which reports both + windows without consuming either +- a Codex response reports either window at or above the ceiling. On a successful response the + in-flight reply is kept and only the next request moves over, because a started stream cannot be + replayed +- Codex answers `429` with a quota signal: a recognized primary/secondary + `x-codex-rate-limit-reached-type`, a specific usage-limit-reached message, or a window at or above + the ceiling. That request is retried on the fallback key immediately +- the control plane reports that the subscription credential is unusable or requires reconnection. + Sandbox-auth, transient broker, network, timeout, configuration, and storage failures do not spend + the fallback key + +A plain `429` with no quota signal is passed through untouched, so short-window throttling does not +spend money. Codex tracks a short (roughly 5-hour) and a weekly window, and the higher usage of the +two decides. An unparseable ceiling is ignored with a log line and treated as 100. If the usage +probe fails, the sandbox stays on the subscription and relies on response headers instead. + +`gpt-5.3-codex-spark` is subscription-only, so its platform fallback uses `gpt-5.3-codex`. Other +allowed models are sent unchanged. If the platform rejects a fallback request, the latch is cleared +and the next turn retries the subscription instead of remaining on a permanently failing paid path. + +Every switch is logged in the sandbox logs as +`[codex-auth-plugin] spilling OpenAI traffic over to OPENAI_API_KEY_FALLBACK: `. + +One caveat: after a successful spillover, the latch lasts as long as the sandbox even if the +subscription window resets under it. OpenCode also reports OpenAI token costs as `0` because the +Codex proxy zeroes them at startup. + +--- + ## How It Works The OpenAI device authorization result is encrypted with `PROVIDER_ACCOUNTS_ENCRYPTION_KEY` in the @@ -103,6 +155,13 @@ Open-Inspect. Confirm that the selected/default OpenAI account is active and the account is verified. If the session explicitly uses API-key mode, confirm `OPENAI_API_KEY` is available in its secret scope. +### "The usage limit has been reached" + +The ChatGPT subscription hit its Codex quota. Wait for the window to reset, switch the session to +another provider's model, or configure a spillover key +([Spilling over](#spilling-over-before-the-subscription-runs-out)) so sessions continue on metered +billing. + ### "Token refresh failed" errors The OAuth grant may have been revoked, expired, or rotated elsewhere. Use **Reconnect** on the diff --git a/docs/SECRETS.md b/docs/SECRETS.md index fcfb4e959d..063e53d470 100644 --- a/docs/SECRETS.md +++ b/docs/SECRETS.md @@ -198,16 +198,18 @@ from it, even after you rotate the secret. Two guidelines: ## Common Examples -| Key | Scope | Purpose | -| ------------------- | ------ | ----------------------------------------------------- | -| `ANTHROPIC_API_KEY` | Global | Claude API access | -| `OPENAI_API_KEY` | Global | OpenAI API access when a session selects API-key mode | -| `XAI_API_KEY` | Global | xAI API access when a session selects API-key mode | -| `DEEPSEEK_API_KEY` | Global | DeepSeek API access | -| `ZHIPU_API_KEY` | Global | Z.AI Coding Plan GLM access | -| `DATABASE_URL` | Repo | Database connection string | -| `AWS_ACCESS_KEY_ID` | Repo | AWS credentials for a specific project | -| `STRIPE_SECRET_KEY` | Repo | Stripe API key for a specific project | +| Key | Scope | Purpose | +| --------------------------------- | ------ | -------------------------------------------------------------------------------------------------------------------------------------- | +| `ANTHROPIC_API_KEY` | Global | Claude API access | +| `OPENAI_API_KEY` | Global | OpenAI API access when a session selects API-key mode | +| `XAI_API_KEY` | Global | xAI API access when a session selects API-key mode | +| `DEEPSEEK_API_KEY` | Global | DeepSeek API access | +| `ZHIPU_API_KEY` | Global | Z.AI Coding Plan GLM access | +| `OPENAI_API_KEY_FALLBACK` | Any | Spillover once the ChatGPT subscription reaches its ceiling ([guide](OPENAI_MODELS.md#spilling-over-before-the-subscription-runs-out)) | +| `OPENAI_SUBSCRIPTION_MAX_PERCENT` | Any | Share of a Codex rate-limit window sandboxes may consume (default 100) | +| `DATABASE_URL` | Repo | Database connection string | +| `AWS_ACCESS_KEY_ID` | Repo | AWS credentials for a specific project | +| `STRIPE_SECRET_KEY` | Repo | Stripe API key for a specific project | --- diff --git a/packages/control-plane/src/routes/model-provider-accounts.test.ts b/packages/control-plane/src/routes/model-provider-accounts.test.ts new file mode 100644 index 0000000000..1e7af08388 --- /dev/null +++ b/packages/control-plane/src/routes/model-provider-accounts.test.ts @@ -0,0 +1,23 @@ +import { describe, expect, it } from "vitest"; +import type { ModelProviderAccountBrokerErrorCode } from "../auth/model-provider-account-broker"; +import { modelProviderBrokerHttpStatus } from "./model-provider-accounts"; + +describe("modelProviderBrokerHttpStatus", () => { + it.each([ + ["account_not_found", 404], + ["upstream_retry_safe", 502], + ["provider_unavailable", 503], + ["exchange_busy", 503], + ["account_inactive", 409], + ["account_archived", 409], + ["provider_mismatch", 409], + ["credential_not_found", 409], + ["credential_invalid", 409], + ["reconnect_required", 409], + ] satisfies Array<[ModelProviderAccountBrokerErrorCode, 404 | 409 | 502 | 503]>)( + "maps %s to HTTP %i", + (code, status) => { + expect(modelProviderBrokerHttpStatus(code)).toBe(status); + } + ); +}); diff --git a/packages/control-plane/src/routes/model-provider-accounts.ts b/packages/control-plane/src/routes/model-provider-accounts.ts index 0d50c97fcd..0cbcbcb3d4 100644 --- a/packages/control-plane/src/routes/model-provider-accounts.ts +++ b/packages/control-plane/src/routes/model-provider-accounts.ts @@ -19,6 +19,7 @@ import { modelProviderAccountAdapterRegistry } from "../auth/model-provider-acco import { ModelProviderAccountBroker, ModelProviderAccountBrokerError, + type ModelProviderAccountBrokerErrorCode, } from "../auth/model-provider-account-broker"; import { ModelProviderAccountStore } from "../db/model-provider-accounts"; import { D1ModelProviderAccountAtomicWriter } from "../db/model-provider-account-atomic-writer"; @@ -440,6 +441,22 @@ async function handleLegacyProviderAccess( }); } +export function modelProviderBrokerHttpStatus( + code: ModelProviderAccountBrokerErrorCode +): 404 | 409 | 502 | 503 { + switch (code) { + case "account_not_found": + return 404; + case "upstream_retry_safe": + return 502; + case "provider_unavailable": + case "exchange_busy": + return 503; + default: + return 409; + } +} + async function handleProviderAccess( _request: Request, env: Env, @@ -489,9 +506,7 @@ async function handleProviderAccess( return json(await broker.getAccess(binding.providerAccountId, parsedProvider)); } catch (cause) { if (cause instanceof ModelProviderAccountBrokerError) { - const status = - cause.code === "account_not_found" ? 404 : cause.code === "upstream_retry_safe" ? 502 : 409; - return error(cause.message, status); + return error(cause.message, modelProviderBrokerHttpStatus(cause.code)); } return error("Provider access unavailable", 503); } diff --git a/packages/sandbox-runtime/src/sandbox_runtime/opencode_server.py b/packages/sandbox-runtime/src/sandbox_runtime/opencode_server.py index 1a8dc49ae8..e83e650e0c 100644 --- a/packages/sandbox-runtime/src/sandbox_runtime/opencode_server.py +++ b/packages/sandbox-runtime/src/sandbox_runtime/opencode_server.py @@ -320,15 +320,19 @@ def _install_skills(self, workdir: Path) -> set[str]: return installed def _setup_managed_oauth(self) -> None: - """Write OpenCode OAuth sentinels for control-plane-managed providers.""" + """Sync OpenCode auth entries with the control-plane-managed providers. + + Writes an OAuth sentinel for every provider the control plane brokers, + and drops sentinels for providers it no longer brokers. The removal + matters on snapshot restores: OpenCode prefers an OAuth entry over a + provider's API key env var, so a leftover sentinel would shadow a key + the operator installed to replace the subscription. + """ openai_managed = os.environ.get("OPENAI_OAUTH_MANAGED") xai_managed = os.environ.get("XAI_OAUTH_MANAGED") - if not openai_managed and not xai_managed: - return try: auth_dir = Path.home() / ".local" / "share" / "opencode" - auth_dir.mkdir(parents=True, exist_ok=True) oauth_entry = { "type": "oauth", @@ -345,7 +349,7 @@ def _setup_managed_oauth(self) -> None: auth_file = auth_dir / "auth.json" tmp_file = auth_dir / ".auth.json.tmp" - existing_entries = {} + existing_entries: dict[str, Any] = {} if auth_file.exists(): try: existing = json.loads(auth_file.read_text()) @@ -353,7 +357,7 @@ def _setup_managed_oauth(self) -> None: existing_entries = existing except (OSError, json.JSONDecodeError): self.log.warn("managed_oauth.existing_auth_invalid") - existing_entries = { + retained = { key: value for key, value in existing_entries.items() if not ( @@ -362,7 +366,11 @@ def _setup_managed_oauth(self) -> None: and key not in entries ) } - entries = {**existing_entries, **entries} + merged = {**retained, **entries} + if not entries and merged == existing_entries: + return + entries = merged + auth_dir.mkdir(parents=True, exist_ok=True) # Write to a temp file created with 0o600 from the start, then # atomically rename so the target is never world-readable. diff --git a/packages/sandbox-runtime/src/sandbox_runtime/plugins/codex-auth-plugin.js b/packages/sandbox-runtime/src/sandbox_runtime/plugins/codex-auth-plugin.js index eae2e60cf5..22b008b10a 100644 --- a/packages/sandbox-runtime/src/sandbox_runtime/plugins/codex-auth-plugin.js +++ b/packages/sandbox-runtime/src/sandbox_runtime/plugins/codex-auth-plugin.js @@ -13,9 +13,37 @@ import { createProviderTokenBroker } from "./provider-token-broker.js"; const CODEX_API_ENDPOINT = "https://chatgpt.com/backend-api/codex/responses"; +const OPENAI_API_ENDPOINT = "https://api.openai.com/v1/responses"; +const OPENAI_CHAT_COMPLETIONS_ENDPOINT = "https://api.openai.com/v1/chat/completions"; const OAUTH_DUMMY_KEY = "opencode-oauth-dummy-key"; const tokenBroker = createProviderTokenBroker({ provider: "openai", providerLabel: "OpenAI" }); +/** + * Optional per-token key used only once the ChatGPT subscription cannot serve a + * request. Deliberately not named OPENAI_API_KEY: prepareManagedProviderEnv + * strips that variable from sessions routed to a subscription, because it + * selects metered billing outright. This one rides along and stays unused + * until the subscription cannot answer. + */ +const FALLBACK_KEY_ENV = "OPENAI_API_KEY_FALLBACK"; + +/** + * Percentage of a subscription rate-limit window this sandbox may consume before + * spilling over. Defaults to 100 (spend the window, then switch). Lower values + * reserve headroom for whoever else uses the same ChatGPT account. + */ +const MAX_PERCENT_ENV = "OPENAI_SUBSCRIPTION_MAX_PERCENT"; + +/** Reads window usage without consuming any of it. */ +const USAGE_STATUS_ENDPOINT = "https://chatgpt.com/backend-api/wham/usage"; +const USAGE_PROBE_TIMEOUT_MS = 5000; + +/** Headers the ChatGPT backend expects that api.openai.com has no use for. */ +const CHATGPT_ONLY_HEADERS = ["chatgpt-account-id", "originator", "session_id"]; + +/** Response headers that describe the transport, not the payload. */ +const TRANSPORT_HEADERS = ["content-encoding", "content-length"]; + const ALLOWED_MODELS = new Set([ "gpt-5.1-codex-max", "gpt-5.1-codex-mini", @@ -30,6 +58,16 @@ const ALLOWED_MODELS = new Set([ "gpt-5.1-codex", ]); +const PLATFORM_MODEL_ALIASES = new Map([["gpt-5.3-codex-spark", "gpt-5.3-codex"]]); + +// Latched for the rest of the sandbox's life once the subscription is spent, so +// a doomed Codex call is not repeated on every later turn. +let spilloverLatched = false; + +// One usage probe per sandbox: afterwards every Codex response carries the +// numbers in its headers for free. +let usageProbed = false; + async function ensureAccessToken(getAuth, setAuth) { const result = await tokenBroker.getAccessToken(async (refreshed) => { // Update OpenCode's auth state for consistency. The broker cache remains @@ -54,6 +92,223 @@ async function ensureAccessToken(getAuth, setAuth) { }; } +function headersFrom(init) { + const headers = new Headers(); + if (!init?.headers) return headers; + if (init.headers instanceof Headers) { + init.headers.forEach((value, key) => headers.set(key, value)); + } else if (Array.isArray(init.headers)) { + for (const [key, value] of init.headers) { + if (value !== undefined) headers.set(key, String(value)); + } + } else { + for (const [key, value] of Object.entries(init.headers)) { + if (value !== undefined) headers.set(key, String(value)); + } + } + return headers; +} + +/** + * Fold both fetch shapes — `(url, init)` and a `Request` — into one plain init. + * opencode's provider client passes an init today, but a `Request` carries its + * own method, headers and body, and spreading an absent init would send the + * subscription call and every spillover retry as a bodiless GET. Buffering the + * body to a string here is also what lets a 429 be retried at all. + */ +async function normalizeRequest(requestInput, init) { + const request = requestInput instanceof Request ? requestInput : null; + const url = request + ? new URL(request.url) + : requestInput instanceof URL + ? requestInput + : new URL(String(requestInput)); + + const headers = new Headers(); + if (request) request.headers.forEach((value, key) => headers.set(key, value)); + for (const [key, value] of headersFrom(init)) headers.set(key, value); + + let body = init?.body; + if (body === undefined && request?.body) body = await request.text(); + + return { + url, + headers, + method: init?.method ?? request?.method, + body, + // Without the source Request's signal, a cancelled turn would leave the + // subscription or platform call running. + signal: init?.signal ?? request?.signal, + }; +} + +function isChatCompletionsRequest(url) { + return url.pathname.includes("/chat/completions"); +} + +function isModelRequest(url) { + return url.pathname.includes("/v1/responses") || isChatCompletionsRequest(url); +} + +/** + * Platform endpoint that keeps the request contract the caller chose: Chat + * Completions and Responses payloads are not interchangeable, so a + * /chat/completions body must not be replayed against /v1/responses. Origin and + * path are fixed rather than forwarded from the request, so a proxied base URL + * cannot steer spillover traffic somewhere else. + */ +function fallbackEndpoint(url) { + return isChatCompletionsRequest(url) ? OPENAI_CHAT_COMPLETIONS_ENDPOINT : OPENAI_API_ENDPOINT; +} + +function toPercent(value) { + if (typeof value === "number") return Number.isFinite(value) ? value : null; + if (typeof value !== "string" || value.trim() === "") return null; + const percent = Number(value); + return Number.isFinite(percent) ? percent : null; +} + +/** The configured share of a subscription window this sandbox may consume. */ +function subscriptionMaxPercent() { + const raw = process.env[MAX_PERCENT_ENV]; + if (!raw) return 100; + const percent = toPercent(raw); + if (percent === null || percent <= 0 || percent > 100) { + console.error( + `[codex-auth-plugin] ignoring ${MAX_PERCENT_ENV}="${raw}": expected a percentage in (0, 100]` + ); + return 100; + } + return percent; +} + +/** Highest window usage Codex reported on a response, or null when absent. */ +function usedPercentFromHeaders(headers) { + let highest = null; + for (const window of ["primary", "secondary"]) { + const used = toPercent(headers.get(`x-codex-${window}-used-percent`)); + if (used !== null) highest = Math.max(highest ?? 0, used); + } + return highest; +} + +/** + * Why the subscription can no longer serve this request, or null to keep using + * it. Codex reports usage through its own header family (x-codex-*) rather than + * the standard x-ratelimit-* headers. + */ +function spentReason(response, { maxPercent = 100, bodyText = "" } = {}) { + const reached = response.headers.get("x-codex-rate-limit-reached-type")?.toLowerCase(); + if (reached === "primary" || reached === "secondary") { + return `Codex reported the ${reached} limit reached`; + } + const used = usedPercentFromHeaders(response.headers); + if (used !== null && used >= maxPercent) { + return `subscription usage at ${used}% of the ${maxPercent}% ceiling`; + } + if (/\busage limit(?: has been)? reached\b/i.test(bodyText)) { + return "the ChatGPT subscription reported its usage limit"; + } + return null; +} + +/** + * Reads the account's window usage from the ChatGPT usage endpoint, which does + * not consume any of it. Returns the highest window, or null when the payload + * carries no usage at all. + */ +async function probeUsedPercent(accessToken, accountId, callerSignal) { + const headers = new Headers({ + authorization: `Bearer ${accessToken}`, + originator: "opencode", + }); + if (accountId) headers.set("ChatGPT-Account-Id", accountId); + + // The probe runs before the turn's own request, so a cancelled turn must not + // wait out the probe timeout. + const timeout = AbortSignal.timeout(USAGE_PROBE_TIMEOUT_MS); + const response = await fetch(USAGE_STATUS_ENDPOINT, { + headers, + signal: callerSignal ? AbortSignal.any([callerSignal, timeout]) : timeout, + }); + if (!response.ok) throw new Error(`usage status ${response.status}`); + + const rateLimit = (await response.json())?.rate_limit; + if (rateLimit?.limit_reached) return 100; + let highest = null; + for (const window of [rateLimit?.primary_window, rateLimit?.secondary_window]) { + const used = toPercent(window?.used_percent); + if (used !== null) highest = Math.max(highest ?? 0, used); + } + return highest; +} + +function spilloverHeaders(headers, apiKey) { + const next = new Headers(headers); + for (const name of CHATGPT_ONLY_HEADERS) next.delete(name); + next.set("authorization", `Bearer ${apiKey}`); + return next; +} + +function platformFallbackBody(body) { + if (typeof body !== "string") return body; + try { + const parsed = JSON.parse(body); + const platformModel = PLATFORM_MODEL_ALIASES.get(parsed?.model); + return platformModel ? JSON.stringify({ ...parsed, model: platformModel }) : body; + } catch { + return body; + } +} + +/** Re-materialize a response whose body was read to classify a 429. */ +function replayResponse(response, bodyText) { + const headers = new Headers(response.headers); + for (const name of TRANSPORT_HEADERS) headers.delete(name); + return new Response(bodyText, { + status: response.status, + statusText: response.statusText, + headers, + }); +} + +function latchSpillover(reason) { + if (spilloverLatched) return; + spilloverLatched = true; + console.error( + `[codex-auth-plugin] spilling OpenAI traffic over to ${FALLBACK_KEY_ENV}: ${reason}` + ); +} + +async function fetchFallback(fallbackUrl, baseInit, headers, apiKey, reason = null) { + const response = await fetch(fallbackUrl, { + ...baseInit, + body: platformFallbackBody(baseInit.body), + headers: spilloverHeaders(headers, apiKey), + }); + if (response.ok) { + if (reason) latchSpillover(reason); + return response; + } + + // A platform outage or unsupported model must not strand the sandbox on a + // permanently failing paid path. Retry the subscription on the next turn. + spilloverLatched = false; + console.error( + `[codex-auth-plugin] ${FALLBACK_KEY_ENV} request failed with status ${response.status}; retrying the subscription on the next turn` + ); + return response; +} + +function isPermanentSubscriptionTokenFailure(error) { + // Provider-account credentials that are invalid or require reconnection are + // reported as 409. A 401 means the router rejected this sandbox's own token + // and says nothing about the ChatGPT subscription. + return ( + error?.name === "ProviderTokenBrokerError" && error.kind === "http" && error.status === 409 + ); +} + export const CodexAuthProxy = async (input) => { return { auth: { @@ -113,29 +368,96 @@ export const CodexAuthProxy = async (input) => { return { apiKey: OAUTH_DUMMY_KEY, async fetch(requestInput, init) { - const request = new Request(requestInput, init); + const { + url: parsed, + headers, + method, + body, + signal, + } = await normalizeRequest(requestInput, init); + const { headers: _discardedHeaders, ...restInit } = init ?? {}; + const baseInit = { ...restInit, method, body, signal }; const currentAuth = await getAuth(); - if (currentAuth.type !== "oauth") return fetch(request); + if (currentAuth.type !== "oauth") return fetch(parsed, { ...baseInit, headers }); + + // opencode signs the request with a placeholder API key; this proxy + // supplies the real credential instead. A caller that has switched + // away from OAuth keeps its own authorization, hence the early + // return above. + headers.delete("authorization"); + + const modelRequest = isModelRequest(parsed); + const fallbackKey = (modelRequest && process.env[FALLBACK_KEY_ENV]) || ""; + const fallbackUrl = fallbackEndpoint(parsed); + + if (fallbackKey && spilloverLatched) { + return fetchFallback(fallbackUrl, baseInit, headers, fallbackKey); + } + + let accessToken; + let accountId; + try { + ({ accessToken, accountId } = await ensureAccessToken(getAuth, setAuth)); + } catch (error) { + if (!fallbackKey || !isPermanentSubscriptionTokenFailure(error)) throw error; + return fetchFallback( + fallbackUrl, + baseInit, + headers, + fallbackKey, + `subscription token unavailable (${error.message})` + ); + } + + headers.set("authorization", `Bearer ${accessToken}`); + if (accountId) headers.set("ChatGPT-Account-Id", accountId); - request.headers.delete("authorization"); + const maxPercent = fallbackKey ? subscriptionMaxPercent() : 100; - // Ensure we have a valid access token - const { accessToken, accountId } = await ensureAccessToken(getAuth, setAuth); + // With a ceiling below 100 the first request of a sandbox must not + // discover the ceiling by consuming a turn past it, so ask the usage + // endpoint first. A failed probe simply leaves the header path to it. + if (fallbackKey && maxPercent < 100 && !usageProbed) { + usageProbed = true; + try { + const used = await probeUsedPercent(accessToken, accountId, signal); + if (used !== null && used >= maxPercent) { + return fetchFallback( + fallbackUrl, + baseInit, + headers, + fallbackKey, + `subscription usage at ${used}% of the ${maxPercent}% ceiling` + ); + } + } catch (error) { + console.error( + `[codex-auth-plugin] usage probe failed, staying on the subscription: ${error.message}` + ); + } + } - const parsed = new URL(request.url); - const url = - parsed.pathname.includes("/v1/responses") || - parsed.pathname.includes("/chat/completions") - ? new URL(CODEX_API_ENDPOINT) - : parsed; - const proxiedRequest = new Request(url, request); + const response = await fetch(modelRequest ? CODEX_API_ENDPOINT : parsed, { + ...baseInit, + headers, + }); + if (!fallbackKey) return response; - // Replace the dummy API key without discarding source Request options. - proxiedRequest.headers.set("authorization", `Bearer ${accessToken}`); - if (accountId) proxiedRequest.headers.set("ChatGPT-Account-Id", accountId); + // A stream that has already started cannot be replayed, so a spent + // window observed on a successful call only redirects the next one. + if (response.status !== 429) { + const reason = spentReason(response, { maxPercent }); + if (reason) latchSpillover(reason); + return response; + } - return fetch(proxiedRequest); + const bodyText = await response.text().catch(() => ""); + const reason = spentReason(response, { maxPercent, bodyText }); + if (!reason || typeof body !== "string") { + return replayResponse(response, bodyText); + } + return fetchFallback(fallbackUrl, baseInit, headers, fallbackKey, reason); }, }; }, diff --git a/packages/sandbox-runtime/src/sandbox_runtime/plugins/provider-token-broker.js b/packages/sandbox-runtime/src/sandbox_runtime/plugins/provider-token-broker.js index 2cefbc7b34..f75b4ef067 100644 --- a/packages/sandbox-runtime/src/sandbox_runtime/plugins/provider-token-broker.js +++ b/packages/sandbox-runtime/src/sandbox_runtime/plugins/provider-token-broker.js @@ -1,3 +1,12 @@ +export class ProviderTokenBrokerError extends Error { + constructor(message, { kind, status = null }) { + super(message); + this.name = "ProviderTokenBrokerError"; + this.kind = kind; + this.status = status; + } +} + const REFRESH_BUFFER_MS = 5 * 60 * 1000; const TOKEN_REQUEST_TIMEOUT_MS = 30_000; const DEFAULT_EXPIRES_IN_SECONDS = 3600; @@ -21,7 +30,9 @@ function validateBrokerResponse(result, providerLabel) { !Number.isFinite(result.expiresIn) || result.expiresIn <= 0)) ) { - throw new Error(`Invalid ${providerLabel} token broker response`); + throw new ProviderTokenBrokerError(`Invalid ${providerLabel} token broker response`, { + kind: "invalid_response", + }); } } @@ -39,7 +50,9 @@ export function createProviderTokenBroker({ provider, providerLabel }) { const authToken = process.env.SANDBOX_AUTH_TOKEN; const sessionId = getSessionId(); if (!controlPlaneUrl || !authToken || !sessionId) { - throw new Error(`Missing environment for ${providerLabel} token refresh`); + throw new ProviderTokenBrokerError(`Missing environment for ${providerLabel} token refresh`, { + kind: "configuration", + }); } const response = await fetch( @@ -52,7 +65,10 @@ export function createProviderTokenBroker({ provider, providerLabel }) { ); if (!response.ok) { const body = (await response.text()).slice(0, 200); - throw new Error(`${providerLabel} token refresh failed (${response.status}): ${body}`); + throw new ProviderTokenBrokerError( + `${providerLabel} token refresh failed (${response.status}): ${body}`, + { kind: "http", status: response.status } + ); } const result = await response.json(); diff --git a/packages/sandbox-runtime/tests/codex-auth-plugin.test.mjs b/packages/sandbox-runtime/tests/codex-auth-plugin.test.mjs index d480fea435..b15006b515 100644 --- a/packages/sandbox-runtime/tests/codex-auth-plugin.test.mjs +++ b/packages/sandbox-runtime/tests/codex-auth-plugin.test.mjs @@ -1,12 +1,495 @@ import assert from "node:assert/strict"; import test from "node:test"; -import { CodexAuthProxy } from "../src/sandbox_runtime/plugins/codex-auth-plugin.js"; +const PLUGIN_PATH = "../src/sandbox_runtime/plugins/codex-auth-plugin.js"; +const MODEL_REQUEST_URL = "https://api.openai.com/v1/responses"; +const REQUEST_INIT = { + method: "POST", + body: JSON.stringify({ model: "gpt-5.4", input: "hi" }), + headers: { authorization: "Bearer opencode-oauth-dummy-key", originator: "opencode" }, +}; + +process.env.CONTROL_PLANE_URL = "https://control.test"; +process.env.SANDBOX_AUTH_TOKEN = "sandbox-token"; +process.env.SESSION_CONFIG = JSON.stringify({ sessionId: "session-1" }); +delete process.env.OPENAI_SUBSCRIPTION_MAX_PERCENT; + +/** + * Load a fresh copy of the plugin. The spillover latch is module state, so each + * case needs its own instance. + */ +async function loadProxy(tag) { + const { CodexAuthProxy } = await import(`${PLUGIN_PATH}?case=${tag}`); + const plugin = await CodexAuthProxy({ client: { auth: { set: async () => {} } } }); + return plugin.auth.loader(async () => ({ type: "oauth", refresh: "managed-by-control-plane" }), { + models: { "gpt-5.4": { cost: {} } }, + }); +} + +/** + * Route stubbed traffic by path: the control-plane broker always mints a token + * unless `broker` overrides it, the usage endpoint answers with `usage`, the + * Codex backend with `codex`, and the platform API always succeeds. + */ +function stubFetch({ codex, broker, usage, platform } = {}) { + const calls = []; + globalThis.fetch = async (url, init) => { + const target = String(url); + calls.push({ + url: target, + method: init?.method, + headers: new Headers(init?.headers), + body: init?.body, + signal: init?.signal, + }); + if (target.includes("/provider-auth/openai/access-token")) { + return ( + broker?.() ?? + Response.json({ + accessToken: "cp-access", + expiresIn: 3600, + providerMetadata: { accountId: "acct-1" }, + }) + ); + } + if (target.includes("/wham/usage")) { + return usage?.(init) ?? new Response("no usage stub", { status: 404 }); + } + if (target.startsWith("https://chatgpt.com/")) return codex(calls.length); + return platform?.(target, init) ?? new Response("platform-ok", { status: 200 }); + }; + return calls; +} + +const usageResponse = (primary, secondary) => + Response.json({ + plan_type: "pro", + rate_limit: { + allowed: true, + limit_reached: false, + primary_window: { used_percent: primary, limit_window_seconds: 18000, reset_at: 1 }, + secondary_window: { used_percent: secondary, limit_window_seconds: 604800, reset_at: 2 }, + }, + }); + +const usageLimitResponse = () => + new Response(JSON.stringify({ error: { message: "The usage limit has been reached" } }), { + status: 429, + headers: { "x-codex-rate-limit-reached-type": "secondary" }, + }); + +test("spills over to the platform API on a usage-limit 429, then latches", async () => { + process.env.OPENAI_API_KEY_FALLBACK = "sk-fallback"; + const calls = stubFetch({ codex: () => usageLimitResponse() }); + const loaded = await loadProxy("latch"); + + const first = await loaded.fetch(MODEL_REQUEST_URL, REQUEST_INIT); + assert.equal(first.status, 200); + assert.equal(await first.text(), "platform-ok"); + + const subscriptionCall = calls.find((call) => call.url.startsWith("https://chatgpt.com/")); + assert.equal(subscriptionCall.headers.get("authorization"), "Bearer cp-access"); + assert.equal(subscriptionCall.headers.get("chatgpt-account-id"), "acct-1"); + + const spilloverCall = calls.at(-1); + assert.equal(spilloverCall.url, MODEL_REQUEST_URL); + assert.equal(spilloverCall.headers.get("authorization"), "Bearer sk-fallback"); + assert.equal(spilloverCall.headers.get("chatgpt-account-id"), null); + assert.equal(spilloverCall.headers.get("originator"), null); + assert.equal(spilloverCall.body, REQUEST_INIT.body); + + // Latched: the second turn must not retry the exhausted subscription. + const before = calls.length; + await loaded.fetch(MODEL_REQUEST_URL, REQUEST_INIT); + assert.deepEqual( + calls.slice(before).map((call) => call.url), + [MODEL_REQUEST_URL] + ); +}); + +test("retries the subscription after a failed platform spillover", async () => { + process.env.OPENAI_API_KEY_FALLBACK = "sk-fallback"; + let platformCalls = 0; + const calls = stubFetch({ + codex: () => usageLimitResponse(), + platform: () => + ++platformCalls === 1 + ? new Response("platform unavailable", { status: 503 }) + : new Response("platform-ok", { status: 200 }), + }); + const loaded = await loadProxy("fallback-failure"); + + assert.equal((await loaded.fetch(MODEL_REQUEST_URL, REQUEST_INIT)).status, 503); + const before = calls.length; + assert.equal((await loaded.fetch(MODEL_REQUEST_URL, REQUEST_INIT)).status, 200); + assert.ok( + calls.slice(before).some((call) => call.url.startsWith("https://chatgpt.com/")), + "the subscription is retried after a failed platform request" + ); +}); + +test("keeps the Chat Completions contract when spilling over", async () => { + process.env.OPENAI_API_KEY_FALLBACK = "sk-fallback"; + const calls = stubFetch({ codex: () => usageLimitResponse() }); + const loaded = await loadProxy("chat-completions"); + + const chatUrl = "https://api.openai.com/v1/chat/completions"; + const chatInit = { + ...REQUEST_INIT, + body: JSON.stringify({ model: "gpt-5.4", messages: [{ role: "user", content: "hi" }] }), + }; + const response = await loaded.fetch(chatUrl, chatInit); + assert.equal(response.status, 200); + + // A Chat Completions body is not a Responses body, so the spillover must not + // rewrite the path to /v1/responses. + const spilloverCall = calls.at(-1); + assert.equal(spilloverCall.url, chatUrl); + assert.equal(spilloverCall.headers.get("authorization"), "Bearer sk-fallback"); + assert.equal(spilloverCall.body, chatInit.body); +}); + +test("maps a subscription-only Spark model to its platform fallback", async () => { + process.env.OPENAI_API_KEY_FALLBACK = "sk-fallback"; + const calls = stubFetch({ codex: () => usageLimitResponse() }); + const loaded = await loadProxy("spark-model"); + const sparkInit = { + ...REQUEST_INIT, + body: JSON.stringify({ model: "gpt-5.3-codex-spark", input: "hi" }), + }; + + const response = await loaded.fetch(MODEL_REQUEST_URL, sparkInit); + + assert.equal(response.status, 200); + assert.equal(JSON.parse(calls.at(-1).body).model, "gpt-5.3-codex"); +}); + +test("passes a throttling 429 through without spending the fallback key", async () => { + process.env.OPENAI_API_KEY_FALLBACK = "sk-fallback"; + const calls = stubFetch({ + codex: () => new Response("slow down", { status: 429 }), + }); + const loaded = await loadProxy("throttle"); + + const response = await loaded.fetch(MODEL_REQUEST_URL, REQUEST_INIT); + assert.equal(response.status, 429); + assert.equal(await response.text(), "slow down"); + assert.equal( + calls.filter((call) => call.url === MODEL_REQUEST_URL).length, + 0, + "no platform-API call for a transient throttle" + ); +}); + +test("does not treat generic quota wording as subscription exhaustion", async () => { + process.env.OPENAI_API_KEY_FALLBACK = "sk-fallback"; + const calls = stubFetch({ + codex: () => new Response("temporary quota throttle", { status: 429 }), + }); + const loaded = await loadProxy("quota-wording"); + + const response = await loaded.fetch(MODEL_REQUEST_URL, REQUEST_INIT); + + assert.equal(response.status, 429); + assert.equal(await response.text(), "temporary quota throttle"); + assert.equal(calls.filter((call) => call.url === MODEL_REQUEST_URL).length, 0); +}); + +test("leaves a usage-limit 429 alone when no fallback key is configured", async () => { + delete process.env.OPENAI_API_KEY_FALLBACK; + const calls = stubFetch({ codex: () => usageLimitResponse() }); + const loaded = await loadProxy("no-key"); + + const response = await loaded.fetch(MODEL_REQUEST_URL, REQUEST_INIT); + assert.equal(response.status, 429); + assert.equal(calls.filter((call) => call.url === MODEL_REQUEST_URL).length, 0); +}); + +test("spills over when the control plane cannot mint a subscription token", async () => { + process.env.OPENAI_API_KEY_FALLBACK = "sk-fallback"; + const calls = stubFetch({ + codex: () => new Response("unreachable", { status: 500 }), + broker: () => new Response("reconnect required", { status: 409 }), + }); + const loaded = await loadProxy("broker-down"); + + const response = await loaded.fetch(MODEL_REQUEST_URL, REQUEST_INIT); + assert.equal(await response.text(), "platform-ok"); + assert.equal(calls.filter((call) => call.url.startsWith("https://chatgpt.com/")).length, 0); + assert.equal(calls.at(-1).headers.get("authorization"), "Bearer sk-fallback"); +}); + +test("does not spend the fallback key on a transient broker failure", async () => { + process.env.OPENAI_API_KEY_FALLBACK = "sk-fallback"; + const calls = stubFetch({ + codex: () => new Response("unreachable", { status: 500 }), + broker: () => new Response("temporarily unavailable", { status: 503 }), + }); + const loaded = await loadProxy("broker-transient"); + + await assert.rejects( + loaded.fetch(MODEL_REQUEST_URL, REQUEST_INIT), + /OpenAI token refresh failed \(503\)/ + ); + assert.equal(calls.filter((call) => call.url === MODEL_REQUEST_URL).length, 0); +}); + +test("does not spend the fallback key when sandbox authentication is rejected", async () => { + process.env.OPENAI_API_KEY_FALLBACK = "sk-fallback"; + const calls = stubFetch({ + codex: () => new Response("unreachable", { status: 500 }), + broker: () => new Response("invalid sandbox token", { status: 401 }), + }); + const loaded = await loadProxy("broker-sandbox-auth"); + + await assert.rejects( + loaded.fetch(MODEL_REQUEST_URL, REQUEST_INIT), + /OpenAI token refresh failed \(401\)/ + ); + assert.equal(calls.filter((call) => call.url === MODEL_REQUEST_URL).length, 0); +}); + +test("keeps a Request-shaped call intact when the subscription token fails", async () => { + process.env.OPENAI_API_KEY_FALLBACK = "sk-fallback"; + const calls = stubFetch({ + codex: () => new Response("unreachable", { status: 500 }), + broker: () => new Response("reconnect required", { status: 409 }), + }); + const loaded = await loadProxy("request-input-token"); + + // A Request carries its own method and body; an absent init must not turn the + // spillover retry into a bodiless GET. + const response = await loaded.fetch(new Request(MODEL_REQUEST_URL, REQUEST_INIT)); + assert.equal(await response.text(), "platform-ok"); + + const spilloverCall = calls.at(-1); + assert.equal(spilloverCall.url, MODEL_REQUEST_URL); + assert.equal(spilloverCall.method, "POST"); + assert.equal(spilloverCall.body, REQUEST_INIT.body); + assert.equal(spilloverCall.headers.get("authorization"), "Bearer sk-fallback"); +}); + +test("spills over a Request-shaped call on a usage-limit 429", async () => { + process.env.OPENAI_API_KEY_FALLBACK = "sk-fallback"; + const calls = stubFetch({ codex: () => usageLimitResponse() }); + const loaded = await loadProxy("request-input-429"); + + const response = await loaded.fetch(new Request(MODEL_REQUEST_URL, REQUEST_INIT)); + assert.equal(await response.text(), "platform-ok"); + + const subscriptionCall = calls.find((call) => call.url.startsWith("https://chatgpt.com/")); + assert.equal(subscriptionCall.method, "POST"); + assert.equal(subscriptionCall.body, REQUEST_INIT.body); + + // Retrying a 429 needs the body in hand, which a live Request stream cannot give. + const spilloverCall = calls.at(-1); + assert.equal(spilloverCall.url, MODEL_REQUEST_URL); + assert.equal(spilloverCall.body, REQUEST_INIT.body); + assert.equal(spilloverCall.headers.get("authorization"), "Bearer sk-fallback"); +}); + +test("forwards a Request's abort signal to both legs", async () => { + process.env.OPENAI_API_KEY_FALLBACK = "sk-fallback"; + const calls = stubFetch({ codex: () => usageLimitResponse() }); + const loaded = await loadProxy("request-input-signal"); + + const controller = new AbortController(); + const request = new Request(MODEL_REQUEST_URL, { ...REQUEST_INIT, signal: controller.signal }); + await loaded.fetch(request); + + // `new Request(url, { signal })` exposes a dependent signal rather than the + // one passed in, so cancellation, not identity, is what must survive. + const subscriptionCall = calls.find((call) => call.url.startsWith("https://chatgpt.com/")); + const spilloverCall = calls.at(-1); + assert.ok(subscriptionCall.signal, "subscription call carries a signal"); + assert.ok(spilloverCall.signal, "spillover call carries a signal"); + assert.equal(subscriptionCall.signal.aborted, false); + controller.abort(); + assert.equal(subscriptionCall.signal.aborted, true); + assert.equal(spilloverCall.signal.aborted, true); +}); + +test("latches on exhausted usage headers reported by a successful call", async () => { + process.env.OPENAI_API_KEY_FALLBACK = "sk-fallback"; + const calls = stubFetch({ + codex: () => + new Response("codex-ok", { + status: 200, + headers: { "x-codex-secondary-used-percent": "100" }, + }), + }); + const loaded = await loadProxy("headers"); + + const first = await loaded.fetch(MODEL_REQUEST_URL, REQUEST_INIT); + assert.equal(await first.text(), "codex-ok", "the in-flight call is never discarded"); + + const before = calls.length; + const second = await loaded.fetch(MODEL_REQUEST_URL, REQUEST_INIT); + assert.equal(await second.text(), "platform-ok"); + assert.deepEqual( + calls.slice(before).map((call) => call.url), + [MODEL_REQUEST_URL] + ); +}); + +test("spills over before touching the subscription when usage is over the ceiling", async () => { + process.env.OPENAI_API_KEY_FALLBACK = "sk-fallback"; + process.env.OPENAI_SUBSCRIPTION_MAX_PERCENT = "80"; + const calls = stubFetch({ + codex: () => new Response("codex-should-not-be-called", { status: 200 }), + usage: () => usageResponse(42, 85), + }); + const loaded = await loadProxy("ceiling-over"); + + const response = await loaded.fetch(MODEL_REQUEST_URL, REQUEST_INIT); + assert.equal(await response.text(), "platform-ok"); + + const probe = calls.find((call) => call.url.includes("/wham/usage")); + assert.equal(probe.headers.get("authorization"), "Bearer cp-access"); + assert.equal(probe.headers.get("chatgpt-account-id"), "acct-1"); + assert.equal( + calls.filter((call) => call.url.includes("/codex/responses")).length, + 0, + "no subscription turn is spent past the ceiling" + ); +}); + +test("keeps the subscription while usage is under the ceiling", async () => { + process.env.OPENAI_API_KEY_FALLBACK = "sk-fallback"; + process.env.OPENAI_SUBSCRIPTION_MAX_PERCENT = "80"; + const calls = stubFetch({ + codex: () => + new Response("codex-ok", { + status: 200, + headers: { "x-codex-secondary-used-percent": "50" }, + }), + usage: () => usageResponse(40, 50), + }); + const loaded = await loadProxy("ceiling-under"); + + assert.equal(await (await loaded.fetch(MODEL_REQUEST_URL, REQUEST_INIT)).text(), "codex-ok"); + assert.equal(await (await loaded.fetch(MODEL_REQUEST_URL, REQUEST_INIT)).text(), "codex-ok"); + assert.equal( + calls.filter((call) => call.url.includes("/wham/usage")).length, + 1, + "the usage endpoint is probed once per sandbox" + ); + assert.equal(calls.filter((call) => call.url === MODEL_REQUEST_URL).length, 0); +}); + +test("abandons a usage probe still in flight when the caller aborts", async () => { + process.env.OPENAI_API_KEY_FALLBACK = "sk-fallback"; + process.env.OPENAI_SUBSCRIPTION_MAX_PERCENT = "80"; + + // A probe that only settles on cancellation: without the caller's signal it + // would hang until USAGE_PROBE_TIMEOUT_MS, which is exactly the wait under test. + let probeStarted; + const probeSignal = new Promise((resolve) => (probeStarted = resolve)); + stubFetch({ + codex: () => new Response("codex-ok", { status: 200 }), + usage: (init) => + new Promise((_, reject) => { + probeStarted(init.signal); + init.signal.addEventListener("abort", () => + reject(new DOMException("aborted", "AbortError")) + ); + }), + }); + const loaded = await loadProxy("probe-abort"); + + const controller = new AbortController(); + const request = new Request(MODEL_REQUEST_URL, { ...REQUEST_INIT, signal: controller.signal }); + const pending = loaded.fetch(request); + + const signal = await probeSignal; + assert.equal(signal.aborted, false, "the probe is in flight"); + controller.abort(); + + let stall; + const outcome = await Promise.race([ + pending.then( + () => "settled", + () => "settled" + ), + new Promise((resolve) => (stall = setTimeout(() => resolve("blocked"), 500))), + ]); + clearTimeout(stall); + assert.equal(signal.aborted, true); + assert.equal(outcome, "settled", "the turn does not wait out the probe timeout"); +}); + +test("latches at the ceiling from a successful response's headers", async () => { + process.env.OPENAI_API_KEY_FALLBACK = "sk-fallback"; + process.env.OPENAI_SUBSCRIPTION_MAX_PERCENT = "80"; + const calls = stubFetch({ + codex: () => + new Response("codex-ok", { + status: 200, + headers: { "x-codex-primary-used-percent": "80.4" }, + }), + usage: () => usageResponse(10, 10), + }); + const loaded = await loadProxy("ceiling-headers"); + + assert.equal(await (await loaded.fetch(MODEL_REQUEST_URL, REQUEST_INIT)).text(), "codex-ok"); + const before = calls.length; + assert.equal(await (await loaded.fetch(MODEL_REQUEST_URL, REQUEST_INIT)).text(), "platform-ok"); + assert.deepEqual( + calls.slice(before).map((call) => call.url), + [MODEL_REQUEST_URL] + ); +}); + +test("ignores a malformed ceiling and spends the whole subscription", async () => { + process.env.OPENAI_API_KEY_FALLBACK = "sk-fallback"; + process.env.OPENAI_SUBSCRIPTION_MAX_PERCENT = "eighty"; + const calls = stubFetch({ + codex: () => + new Response("codex-ok", { + status: 200, + headers: { "x-codex-secondary-used-percent": "85" }, + }), + usage: () => usageResponse(85, 85), + }); + const loaded = await loadProxy("ceiling-invalid"); + + assert.equal(await (await loaded.fetch(MODEL_REQUEST_URL, REQUEST_INIT)).text(), "codex-ok"); + assert.equal(calls.filter((call) => call.url.includes("/wham/usage")).length, 0); + assert.equal(await (await loaded.fetch(MODEL_REQUEST_URL, REQUEST_INIT)).text(), "codex-ok"); +}); + +test("rejects a partially numeric ceiling", async () => { + process.env.OPENAI_API_KEY_FALLBACK = "sk-fallback"; + process.env.OPENAI_SUBSCRIPTION_MAX_PERCENT = "80garbage"; + const calls = stubFetch({ + codex: () => + new Response("codex-ok", { + status: 200, + headers: { "x-codex-secondary-used-percent": "85" }, + }), + usage: () => usageResponse(85, 85), + }); + const loaded = await loadProxy("ceiling-suffix"); + + assert.equal(await (await loaded.fetch(MODEL_REQUEST_URL, REQUEST_INIT)).text(), "codex-ok"); + assert.equal(calls.filter((call) => call.url.includes("/wham/usage")).length, 0); + assert.equal(await (await loaded.fetch(MODEL_REQUEST_URL, REQUEST_INIT)).text(), "codex-ok"); +}); + +test("stays on the subscription when the usage probe fails", async () => { + process.env.OPENAI_API_KEY_FALLBACK = "sk-fallback"; + process.env.OPENAI_SUBSCRIPTION_MAX_PERCENT = "80"; + const calls = stubFetch({ + codex: () => new Response("codex-ok", { status: 200 }), + usage: () => new Response("boom", { status: 500 }), + }); + const loaded = await loadProxy("probe-failure"); + + assert.equal(await (await loaded.fetch(MODEL_REQUEST_URL, REQUEST_INIT)).text(), "codex-ok"); + assert.equal(calls.filter((call) => call.url.includes("/codex/responses")).length, 1); +}); test("preserves a source Request while proxying Codex authentication", async () => { - process.env.CONTROL_PLANE_URL = "https://control.test"; - process.env.SANDBOX_AUTH_TOKEN = "sandbox-token"; - process.env.SESSION_CONFIG = JSON.stringify({ sessionId: "session-1" }); let upstreamRequest; globalThis.fetch = async (input, init) => { const request = input instanceof Request ? input : new Request(input, init); @@ -20,10 +503,7 @@ test("preserves a source Request while proxying Codex authentication", async () upstreamRequest = request; return new Response(null, { status: 200 }); }; - const plugin = await CodexAuthProxy({ client: { auth: { set: async () => undefined } } }); - const loaded = await plugin.auth.loader(async () => ({ type: "oauth", refresh: "managed" }), { - models: {}, - }); + const loaded = await loadProxy("preserve-request"); await loaded.fetch( new Request("https://api.openai.com/v1/responses", { @@ -50,6 +530,7 @@ test("preserves caller authorization after switching away from OAuth", async () let authReadCount = 0; const getAuth = async () => authReadCount++ === 0 ? { type: "oauth", refresh: "managed" } : { type: "api" }; + const { CodexAuthProxy } = await import(`${PLUGIN_PATH}?case=switch-away`); const plugin = await CodexAuthProxy({ client: { auth: { set: async () => undefined } } }); const loaded = await plugin.auth.loader(getAuth, { models: {} }); @@ -68,6 +549,7 @@ test("keeps GPT-6 Astra available for Codex subscriptions", async () => { cost: { input: 1, output: 1 }, }; const provider = { models: { "gpt-6-astra": astra, "unsupported-model": {} } }; + const { CodexAuthProxy } = await import(`${PLUGIN_PATH}?case=astra-entitlement`); const plugin = await CodexAuthProxy({ client: { auth: { set: async () => undefined } } }); await plugin.auth.loader(async () => ({ type: "oauth", refresh: "managed" }), provider); diff --git a/packages/sandbox-runtime/tests/provider-token-broker.test.mjs b/packages/sandbox-runtime/tests/provider-token-broker.test.mjs index 0c9981de0c..87250bdf9b 100644 --- a/packages/sandbox-runtime/tests/provider-token-broker.test.mjs +++ b/packages/sandbox-runtime/tests/provider-token-broker.test.mjs @@ -1,7 +1,10 @@ import assert from "node:assert/strict"; import test from "node:test"; -import { createProviderTokenBroker } from "../src/sandbox_runtime/plugins/provider-token-broker.js"; +import { + ProviderTokenBrokerError, + createProviderTokenBroker, +} from "../src/sandbox_runtime/plugins/provider-token-broker.js"; function configureSession() { process.env.CONTROL_PLANE_URL = "https://control.test"; @@ -70,3 +73,16 @@ test("clears a failed in-flight refresh so a later request can retry", async () assert.equal((await broker.getAccessToken()).accessToken, "recovered"); assert.equal(requestCount, 2); }); + +test("preserves HTTP status on broker rejections", async () => { + configureSession(); + globalThis.fetch = async () => new Response("temporarily unavailable", { status: 503 }); + const broker = createProviderTokenBroker({ provider: "openai", providerLabel: "OpenAI" }); + + await assert.rejects(broker.getAccessToken(), (error) => { + assert.ok(error instanceof ProviderTokenBrokerError); + assert.equal(error.kind, "http"); + assert.equal(error.status, 503); + return true; + }); +}); diff --git a/packages/sandbox-runtime/tests/test_openai_oauth_setup.py b/packages/sandbox-runtime/tests/test_openai_oauth_setup.py index 1aab50de41..b721667e26 100644 --- a/packages/sandbox-runtime/tests/test_openai_oauth_setup.py +++ b/packages/sandbox-runtime/tests/test_openai_oauth_setup.py @@ -81,6 +81,34 @@ def test_skips_when_no_refresh_token(self, tmp_path, monkeypatch): assert not _auth_file(tmp_path).exists() + def test_drops_stale_sentinel_when_provider_is_no_longer_managed(self, tmp_path, monkeypatch): + """A snapshot-carried sentinel must not shadow OPENAI_API_KEY.""" + sup = _make_opencode_server() + auth_file = _auth_file(tmp_path) + auth_file.parent.mkdir(parents=True) + auth_file.write_text( + json.dumps( + { + "openai": { + "type": "oauth", + "refresh": "managed-by-control-plane", + "access": "stale", + "expires": 0, + }, + "anthropic": {"type": "api", "key": "user-owned"}, + } + ) + ) + monkeypatch.delenv("OPENAI_OAUTH_MANAGED", raising=False) + monkeypatch.delenv("XAI_OAUTH_MANAGED", raising=False) + + with patch("pathlib.Path.home", return_value=tmp_path): + sup._setup_managed_oauth() + + assert json.loads(auth_file.read_text()) == { + "anthropic": {"type": "api", "key": "user-owned"} + } + def test_sets_secure_permissions(self, tmp_path): sup = _make_opencode_server()