From 7cd270dc2cd4f5a4cb07b2f23b4eb0dfe869b0da Mon Sep 17 00:00:00 2001 From: bitkyc08-arch Date: Thu, 20 Aug 2026 05:10:47 +0900 Subject: [PATCH] fix(xai): stream OAuth Grok through Responses OAuth xai/grok-4.5 and grok-4.6 Codex /v1/responses traffic still used the provider-wide openai-chat adapter, while the official Grok CLI catalog declares api_backend: "responses". Chat Completions compatibility holds the stream until the reasoning turn finishes, so Codex sat blank until the turn was effectively done. Declare the Responses wire default for those two models, scoped to OAuth and to responses-shaped inbound traffic. API-key xAI, Chat/Anthropic translation, other Grok models, and any explicit modelAdapters override all stay on Chat. Native Responses returns before the generic recovery loop, so the OAuth 401 replay never ran on this path. Add the equivalent one-shot: refresh once, rebuild the provider and adapter, replay once. It is a single branch rather than a loop, so a second 401 cannot refresh again. Refresh failures go through the existing public OAuth error projector, which the tests pin against path canaries. Carries @olddonkey's #2104 unchanged. Closes #1886 --- src/providers/fastwire.ts | 47 +++- src/providers/registry.ts | 36 ++- src/providers/service-tier.ts | 24 +- src/server/responses/core.ts | 92 ++++++++ structure/04_transports-and-sidecars.md | 24 +- tests/adapter-resolve.test.ts | 32 +++ tests/fastwire-policy.test.ts | 42 ++++ tests/server-xai-oauth-401-replay.test.ts | 28 ++- tests/server-xai-responses-streaming.test.ts | 223 +++++++++++++++++++ 9 files changed, 515 insertions(+), 33 deletions(-) create mode 100644 tests/server-xai-responses-streaming.test.ts diff --git a/src/providers/fastwire.ts b/src/providers/fastwire.ts index d36a9cf19f..1972c4339f 100644 --- a/src/providers/fastwire.ts +++ b/src/providers/fastwire.ts @@ -7,7 +7,7 @@ import type { } from "../types"; import { MODEL_ADAPTER_OVERRIDE_ALLOWED } from "../types"; import { sanitizeLogMetadataString } from "../lib/redact"; -import type { InboundWire, ModelWireDefault } from "./registry"; +import type { InboundWire, ModelWireDefault, ProviderAuthKind } from "./registry"; const SERVICE_TIER_ADAPTERS = new Set(["openai-chat", "openai-responses"]); const FAST_WIRE_ADAPTERS: Readonly>> = { @@ -31,6 +31,7 @@ export type FastPolicyAuthTransport = export interface FastPolicyAuthority { readonly providerAdapter: string; + readonly providerAuthMode?: ProviderAuthKind; readonly fastWireDeclaration: FastWire | null | undefined; readonly modelWireOverrideAllowed: boolean; readonly authTransport: FastPolicyAuthTransport; @@ -114,21 +115,33 @@ function registryDefaultForModel( defaults: Readonly>, modelId: string, inbound: InboundWire, -): string | undefined { + authMode: ProviderAuthKind | undefined, +): { adapter: string; forwardCallerServiceTier?: boolean } | undefined { const normalizedModelId = modelId.trim().toLowerCase(); if (!Object.hasOwn(defaults, normalizedModelId)) return undefined; const declared = defaults[normalizedModelId]; if (declared === undefined) return undefined; - if (typeof declared !== "string" && !declared.inbound.includes(inbound)) return undefined; + if (typeof declared !== "string") { + if (!declared.inbound.includes(inbound)) return undefined; + if (declared.authModes && (authMode === undefined || !declared.authModes.includes(authMode))) { + return undefined; + } + } const wire = typeof declared === "string" ? declared : declared.wire; - return MODEL_ADAPTER_OVERRIDE_ALLOWED.has(wire) ? wire : undefined; + if (!MODEL_ADAPTER_OVERRIDE_ALLOWED.has(wire)) return undefined; + return { + adapter: wire, + ...(typeof declared !== "string" && declared.forwardCallerServiceTier !== undefined + ? { forwardCallerServiceTier: declared.forwardCallerServiceTier } + : {}), + }; } function resolvePolicyAdapter( authority: FastPolicyAuthority, modelId: string, inbound: InboundWire, -): { adapter: string; hardPinned: boolean } { +): { adapter: string; hardPinned: boolean; forwardCallerServiceTier?: boolean } { // Hard pins and configured overrides deliberately use the same exact-key semantics as // resolveWireProtocolOverride(). Registry defaults alone normalize ids at their boundary. const hardPin = Object.hasOwn(authority.hardPins, modelId) @@ -143,8 +156,21 @@ function resolvePolicyAdapter( return { adapter: configured, hardPinned: false }; } if (MODEL_ADAPTER_OVERRIDE_ALLOWED.has(authority.providerAdapter)) { - const registryDefault = registryDefaultForModel(authority.registryWireDefaults, modelId, inbound); - if (registryDefault !== undefined) return { adapter: registryDefault, hardPinned: false }; + const registryDefault = registryDefaultForModel( + authority.registryWireDefaults, + modelId, + inbound, + authority.providerAuthMode, + ); + if (registryDefault !== undefined) { + return { + adapter: registryDefault.adapter, + hardPinned: false, + ...(registryDefault.forwardCallerServiceTier !== undefined + ? { forwardCallerServiceTier: registryDefault.forwardCallerServiceTier } + : {}), + }; + } } } return { adapter: authority.providerAdapter, hardPinned: false }; @@ -155,7 +181,11 @@ export function resolveFastPolicy( modelId: string, inbound: InboundWire = "responses", ): ResolvedFastPolicy { - const { adapter, hardPinned } = resolvePolicyAdapter(authority, modelId, inbound); + const { adapter, hardPinned, forwardCallerServiceTier } = resolvePolicyAdapter( + authority, + modelId, + inbound, + ); const exactCapability = exactModelValue(authority.capability.models, modelId); const capability = authority.capability.provider === false ? false @@ -173,6 +203,7 @@ export function resolveFastPolicy( // tier still needs the final wire's forwarding permission. const forwardCallerTier = capability !== false && callerWireAvailable + && forwardCallerServiceTier !== false && (adapter !== "openai-chat" || authority.capability.chatServiceTier === true); let eligibility: ResolvedFastPolicy["eligibility"]; diff --git a/src/providers/registry.ts b/src/providers/registry.ts index 53922e75ad..3fd6ac7dd8 100644 --- a/src/providers/registry.ts +++ b/src/providers/registry.ts @@ -31,9 +31,15 @@ export type InboundWire = "responses" | "chat" | "anthropic"; /** * A per-model wire default: a bare string applies to every inbound, while the object - * form applies only to the listed inbound protocols. + * form may scope the default to listed inbound protocols and authentication modes. */ -export type ModelWireDefault = string | { wire: string; inbound: readonly InboundWire[] }; +export type ModelWireDefault = string | { + wire: string; + inbound: readonly InboundWire[]; + authModes?: readonly ProviderAuthKind[]; + /** Whether this registry-selected route may relay a caller-owned service_tier. */ + forwardCallerServiceTier?: boolean; +}; export interface ResponsesTerminalRepairPolicy { /** Quiet time after a structurally complete output graph before synthesizing completion. */ @@ -1017,6 +1023,24 @@ export const PROVIDER_REGISTRY: readonly ProviderRegistryEntry[] = [ // grok-4.5; the reasoning ladder does not — 4.6 adds the documented xhigh rung. models: ["grok-4.6", "grok-4.5", "grok-4.3", "grok-4.20-0309-reasoning", "grok-4.20-0309-non-reasoning", "grok-build-0.1", "grok-composer-2.5-fast"], defaultModel: "grok-4.5", + // The current Grok CLI catalog declares both subscription models as native Responses + // backends. Keep API-key and translated Chat/Anthropic callers on their existing wire; + // Codex Responses traffic can relay xAI's SSE as it arrives instead of waiting for the + // Chat Completions compatibility stream to flush at the end of a reasoning turn. + modelWireDefaults: { + "grok-4.6": { + wire: "openai-responses", + inbound: ["responses"], + authModes: ["oauth"], + forwardCallerServiceTier: false, + }, + "grok-4.5": { + wire: "openai-responses", + inbound: ["responses"], + authModes: ["oauth"], + forwardCallerServiceTier: false, + }, + }, // Vision lineup per docs.x.ai model-capabilities/images/understanding: the grok-4.x chat // models accept image input (JPEG/PNG, URL or base64). Without this the catalog leaves // inputModalities undefined, and deriveComboCatalogModel defaults an undefined member to @@ -2679,8 +2703,12 @@ export function providerModelWireDefault( if (!entry?.modelWireDefaults || !providerMatchesRegistryTransport(id, provider)) return undefined; const declared = entry.modelWireDefaults[modelId.trim().toLowerCase()]; if (declared === undefined) return undefined; - // A bare string applies to every inbound; the object form only to the listed ones. - if (typeof declared !== "string" && !declared.inbound.includes(inbound)) return undefined; + // A bare string applies to every inbound/auth mode; the object form may narrow either. + if (typeof declared !== "string") { + if (!declared.inbound.includes(inbound)) return undefined; + const authMode = provider.authMode ?? entry.authKind; + if (declared.authModes && !declared.authModes.includes(authMode)) return undefined; + } const wire = typeof declared === "string" ? declared : declared.wire; return wire !== undefined && allowedWires.has(wire) ? wire : undefined; } diff --git a/src/providers/service-tier.ts b/src/providers/service-tier.ts index 278f89394e..2a09530d23 100644 --- a/src/providers/service-tier.ts +++ b/src/providers/service-tier.ts @@ -45,7 +45,16 @@ function cloneRegistryWireDefaults( for (const [modelId, declaration] of Object.entries(defaults)) { clone[modelId.trim().toLowerCase()] = typeof declaration === "string" ? declaration - : Object.freeze({ wire: declaration.wire, inbound: Object.freeze([...declaration.inbound]) }); + : Object.freeze({ + wire: declaration.wire, + inbound: Object.freeze([...declaration.inbound]), + ...(declaration.authModes + ? { authModes: Object.freeze([...declaration.authModes]) } + : {}), + ...(declaration.forwardCallerServiceTier !== undefined + ? { forwardCallerServiceTier: declaration.forwardCallerServiceTier } + : {}), + }); } return Object.freeze(clone); } @@ -68,6 +77,7 @@ function buildFastPolicyAuthority( const providerCapability = capabilityProvider.supportsServiceTier ?? registry?.supportsServiceTier; const authority: FastPolicyAuthority = Object.freeze({ providerAdapter: provider.adapter, + providerAuthMode: provider.authMode ?? registry?.authKind ?? "key", fastWireDeclaration: cloneFastWire( provider.fastWire !== undefined ? provider.fastWire : registry?.fastWire, { freeze: true }, @@ -241,13 +251,11 @@ export function serviceTierSupportFromPolicy( ): boolean | undefined { if (policy.eligibility === "eligible") return true; if (policy.eligibility === "unclassified") { - // B1 regression guard: an unclassified chat-wire route whose final adapter will not - // forward any tier cannot serialize service_tier, so projecting "unknown" would let - // require.serviceTier: "unsupported" routing stop matching groq/ollama-class providers - // that main projected as false. Chat + no forwarding stays a definitive false; a - // chat route with chatServiceTier: true (forwarding allowed) keeps the historical - // unknown, as does every unclassified Responses-wire route. - if (policy.adapter === "openai-chat" && !policy.forwardCallerTier) return false; + // An unclassified route that cannot forward a caller tier has definitive negative + // evidence even when its adapter can normally serialize service_tier. This covers both + // Chat routes without chatServiceTier and a registry default that explicitly closes a + // subscription gateway. Generic unclassified Responses routes still project unknown. + if (!policy.forwardCallerTier) return false; return undefined; } return false; diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index 9e2813d0b5..d503a08c32 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -2622,6 +2622,98 @@ async function handleResponsesInner( request.releaseBodyObservation?.(); } + // Native Responses providers return before the generic adapter recovery loop below. Keep + // their OAuth contract identical: one pre-stream 401 forces a credential refresh and one + // rebuilt replay. xAI's current subscription models use this branch now that their official + // Grok CLI catalog declares the Responses backend. + if (upstreamResponse.status === 401 && isOAuth401ReplayProvider && sentOAuthSnapshot) { + try { void upstreamResponse.body?.cancel().catch(() => {}); } catch { /* already consumed/closed */ } + let refreshed: OAuthAccessSnapshot; + try { + refreshed = await forceRefreshOAuthAccessSnapshot(sentOAuthSnapshot); + } catch (err) { + upstream.abort(); + releaseCodexAuthContextProbeLease(authCtx); + return formatErrorResponse(401, "authentication_error", publicOAuthAuthenticationErrorMessage(err)); + } + sentOAuthSnapshot = refreshed; + replayOAuthCredentialSnapshot = { + accountId: refreshed.accountId, + generation: refreshed.generation, + }; + if (route.providerName === "kiro") { + parsed._kiroAuthContext = { ...(refreshed.kiro ?? {}) }; + } + const refreshedProvider = resolveProviderTransport( + route.providerName, + { ...route.provider, apiKey: refreshed.accessToken }, + parsed.options.promptCacheKey, + route.providerName === "github-copilot" ? getOAuthCredentialApiBaseUrl(route.providerName) : undefined, + ); + route.provider = refreshedProvider; + const refreshedAdapter = resolveAdapter( + resolveWireProtocolOverride(route.providerName, route.modelId, refreshedProvider, inboundWire), + config.cacheRetention, + ); + if (!("passthrough" in refreshedAdapter) || !refreshedAdapter.passthrough) { + upstream.abort(); + return formatErrorResponse(502, "upstream_error", "OAuth refresh changed the provider wire unexpectedly"); + } + bindRouteReasoningReplayScope({ + parsed, + providerName: route.providerName, + provider: refreshedProvider, + adapterName: refreshedAdapter.name, + oauthCredentialSnapshot: replayOAuthCredentialSnapshot, + }); + logCtx.providerAdapter = refreshedAdapter.name; + sealRequestAttemptIdentity( + logCtx.activeAttempt, + logCtx.provider, + refreshedAdapter.name, + logCtx.accountLogLabel, + ); + try { + request = await refreshedAdapter.buildRequest(parsed, { + headers: selectedForwardHeaders, + translatorBudget, + }); + recordAdapterReasoning(logCtx, request); + recordAdapterTier(logCtx, request); + } catch (err) { + upstream.abort(); + if (options.abortSignal?.aborted) return clientCancelledResponse(); + const msg = err instanceof Error ? err.message : String(err); + return formatErrorResponse(400, "invalid_request_error", redactSecretString(msg)); + } + try { + upstreamResponse = await fetchWithTransientRetry( + recovery => { + noteAttemptSend(logCtx.activeAttempt, passthroughEstimate, recovery ?? "oauth-401"); + return fetchWithHeaderTimeout(request.url, applyUpstreamRecoveryInit({ + method: request.method, + headers: request.headers, + body: request.body, + }, recovery), upstream.signal, connectMs, parsed.stream, + providerFetch(route.provider, options.codexWsRuntimeIdentity, { + providerName: route.providerName, + modelId: route.modelId, + }), + route.provider.authMode === "forward") + .then(res => { + settleObservedHostResponse(); + return res; + }); + }, + { abortSignal: upstream.signal, label: safeHostLabel(request.url) }, + ); + } catch (err) { + return transportFailureResponse(err); + } finally { + request.releaseBodyObservation?.(); + } + } + // Same-target 429 wait-and-retry (opt-in `retryOn429`) for key-auth providers on the // passthrough wire. This branch returns before the recovery loop below, so Responses-shaped // key-auth gateways (e.g. the built-in DeepSeek preset) would otherwise surface 429 diff --git a/structure/04_transports-and-sidecars.md b/structure/04_transports-and-sidecars.md index bdb0d15e0c..fbca919cde 100644 --- a/structure/04_transports-and-sidecars.md +++ b/structure/04_transports-and-sidecars.md @@ -51,7 +51,9 @@ requires a compatible FastWire mapping on the final adapter and an eligible poli `fastMode: false` drops it. On classified Chat routes, `chatServiceTier` separately authorizes foreign caller values; an exact-model `true` does not grant that forwarding permission. On unclassified Chat routes it gates every caller tier because no canonical Fast capability has been -validated. Exact `false` +validated. An object-form registry wire default may also set `forwardCallerServiceTier: false` to +close a known subscription gateway while leaving generic unclassified Responses passthrough +unchanged. Exact `false` narrows provider defaults, and provider-level `supportsServiceTier: false` cannot be reopened. Capability is namespaced by the selected provider and model; model-name similarity and adapter type alone never opt a gateway in. @@ -66,7 +68,18 @@ Registry `modelWireDefaults` select an evidence-backed upstream protocol for an changing the provider-wide adapter. Explicit, allowed `modelAdapters` configuration always wins, including an entry that opts the model back into the provider-wide wire. Defaults are applied only while the configured provider still matches the registry transport, so reusing a preset name for a -different custom destination does not inherit its upstream assumptions. +different custom destination does not inherit its upstream assumptions. Object-form defaults may +also narrow the decision by inbound protocol and authentication mode; an auth-scoped default must +not leak from a subscription transport into an API-key or forwarded-credential route. + +xAI keeps `openai-chat` as its provider-wide compatibility wire. The official Grok CLI catalog +declares the Grok 4.5 and 4.6 subscription models as Responses backends, so only OAuth-backed native +Responses traffic for those exact models selects `openai-responses`. API-key requests, translated +Chat/Anthropic callers, other Grok models, and explicit model adapter overrides retain their +existing wire. This lets Codex receive native xAI SSE deltas as they arrive without widening the +credential or compatibility boundary. These OAuth subscription defaults drop caller-owned +`service_tier`; they neither advertise nor inject Fast. The API-key transport remains governed by +its separate capability declaration. OpenCode Go documents `gpt-5.6-luna` on `/zen/go/v1/responses` while sibling models use its Chat or Anthropic endpoints. The built-in preset therefore selects `openai-responses` only for Luna and @@ -601,9 +614,10 @@ Grounded in the open-sourced official client (xai-org/grok-build); unit + eviden `auth.json` load-merge-persist (`src/oauth/store.ts`); generation-guarded persist (`expectedGeneration` → superseded adoption), conditional `needsReauth`, bounded jittered retry for transient token-endpoint failures. -- **Reactive 401 replay:** the serving recovery loop force-refreshes once (singleflight, - generation-checked) and replays OAuth-backed xAI requests exactly once with a re-resolved - transport; API-key/BYOK paths excluded (`src/server/responses.ts`). +- **Reactive 401 replay:** both the adapter recovery loop and native Responses passthrough branch + force-refresh once (singleflight, generation-checked) and replay OAuth-backed xAI requests + exactly once with a re-resolved transport; API-key/BYOK paths are excluded + (`src/server/responses/core.ts`). - **Header parity:** per-attempt `x-grok-req-id` (fresh UUID inside the transport fetch wrapper), stable session/conv affinity headers, always-set User-Agent, and a single compatibility profile const for the Grok client version (`src/providers/xai-transport.ts`); diff --git a/tests/adapter-resolve.test.ts b/tests/adapter-resolve.test.ts index 28b07091b5..ceffd92149 100644 --- a/tests/adapter-resolve.test.ts +++ b/tests/adapter-resolve.test.ts @@ -92,6 +92,38 @@ describe("per-model wire override (#404)", () => { }); describe("registry per-model wire defaults", () => { + function xai(authMode: "oauth" | "key", overrides: Partial = {}): OcxProviderConfig { + return gateway({ + baseUrl: "https://api.x.ai/v1", + authMode, + ...overrides, + }); + } + + test("routes current xAI subscription models through Responses for native Codex traffic", () => { + for (const model of ["grok-4.6", "grok-4.5"]) { + expect(resolveWireProtocolOverride("xai", model, xai("oauth"), "responses").adapter) + .toBe("openai-responses"); + } + }); + + test("keeps xAI key auth and translated callers on their existing Chat wire", () => { + expect(resolveWireProtocolOverride("xai", "grok-4.6", xai("key"), "responses").adapter) + .toBe("openai-chat"); + expect(resolveWireProtocolOverride("xai", "grok-4.6", xai("oauth"), "chat").adapter) + .toBe("openai-chat"); + expect(resolveWireProtocolOverride("xai", "grok-4.6", xai("oauth"), "anthropic").adapter) + .toBe("openai-chat"); + expect(resolveWireProtocolOverride("xai", "grok-4.3", xai("oauth"), "responses").adapter) + .toBe("openai-chat"); + }); + + test("an explicit xAI Chat override opts out of the subscription Responses default", () => { + const provider = xai("oauth", { modelAdapters: { "grok-4.6": "openai-chat" } }); + expect(resolveWireProtocolOverride("xai", "grok-4.6", provider, "responses").adapter) + .toBe("openai-chat"); + }); + function deepseek(overrides: Partial = {}): OcxProviderConfig { return gateway({ baseUrl: "https://api.deepseek.com", diff --git a/tests/fastwire-policy.test.ts b/tests/fastwire-policy.test.ts index 7aed643290..6cd7529a94 100644 --- a/tests/fastwire-policy.test.ts +++ b/tests/fastwire-policy.test.ts @@ -128,6 +128,27 @@ describe("resolveFastPolicy matrix", () => { expect(resolveFastPolicy(authority, MODEL, "responses").adapter).toBe("openai-responses"); }); + test("registry defaults retain their auth-mode constraint", () => { + const base: FastPolicyAuthority = { + ...authorityForMatrix({ + source: "provider-adapter", + declaration: "undefined", + overrideAllowed: true, + capability: "true", + chatForeignTierForward: true, + }), + providerAdapter: "openai-chat", + registryWireDefaults: { + [MODEL]: { wire: "openai-responses", inbound: ["responses"], authModes: ["oauth"] }, + }, + }; + expect(resolveFastPolicy({ ...base, providerAuthMode: "oauth" }, MODEL).adapter) + .toBe("openai-responses"); + expect(resolveFastPolicy({ ...base, providerAuthMode: "key" }, MODEL).adapter) + .toBe("openai-chat"); + expect(resolveFastPolicy(base, MODEL).adapter).toBe("openai-chat"); + }); + test("hard pins and configured overrides retain exact runtime model-key semantics", () => { const authority: FastPolicyAuthority = { ...authorityForMatrix({ @@ -235,6 +256,27 @@ describe("resolveFastPolicy matrix", () => { expect(fastPolicyForModel(provider, MODEL, "fixture").capability).toBe(false); }); + test("captured xAI registry defaults keep OAuth and key transports separate", () => { + const oauthProvider = Object.freeze({ + adapter: "openai-chat", + baseUrl: "https://api.x.ai/v1", + authMode: "oauth" as const, + }); + const keyProvider = Object.freeze({ + adapter: "openai-chat", + baseUrl: "https://api.x.ai/v1", + authMode: "key" as const, + }); + + expect(fastPolicyForModel(oauthProvider, "grok-4.6", "xai")).toMatchObject({ + adapter: "openai-responses", + eligibility: "unclassified", + forwardCallerTier: false, + }); + expect(fastPolicyForModel(keyProvider, "grok-4.6", "xai").adapter) + .toBe("openai-chat"); + }); + test("prototype-named providers and models use only own wire-policy rows", () => { expect(captureWireAdapterHardPins("toString")).toEqual({}); expect(isWirePinnedModel("toString", MODEL)).toBe(false); diff --git a/tests/server-xai-oauth-401-replay.test.ts b/tests/server-xai-oauth-401-replay.test.ts index 0d3e03d0a3..abe2e62673 100644 --- a/tests/server-xai-oauth-401-replay.test.ts +++ b/tests/server-xai-oauth-401-replay.test.ts @@ -11,7 +11,7 @@ import type { OcxConfig } from "../src/types"; import { installIsolatedCodexHome, type IsolatedCodexHome } from "./helpers/isolated-codex-home"; const TOKEN_ENDPOINT = "https://auth.x.ai/oauth/token"; -const CHAT_ENDPOINT = `${XAI_GROK_CLI_BASE_URL}/chat/completions`; +const OAUTH_RESPONSES_ENDPOINT = `${XAI_GROK_CLI_BASE_URL}/responses`; const PUBLIC_OAUTH_AUTHENTICATION_ERROR = "OAuth authentication failed. Check the OpenCodex account status and retry."; const WINDOWS_PATH_CANARY = "C:\\Users\\Alice\\.opencodex\\auth.json.ocx-tmp"; const UNC_PATH_CANARY = "\\\\server\\share\\opencodex\\auth.json.ocx-tmp"; @@ -68,10 +68,18 @@ function xaiConfig(authMode: "oauth" | "key" = "oauth"): OcxConfig { function successBody(text: string): string { return JSON.stringify({ - id: "chatcmpl-xai-401", - object: "chat.completion", - choices: [{ index: 0, message: { role: "assistant", content: text }, finish_reason: "stop" }], - usage: { prompt_tokens: 3, completion_tokens: 2, total_tokens: 5 }, + id: "resp-xai-401", + object: "response", + status: "completed", + model: "grok-4.5", + output: [{ + id: "msg-xai-401", + type: "message", + status: "completed", + role: "assistant", + content: [{ type: "output_text", text, annotations: [] }], + }], + usage: { input_tokens: 3, output_tokens: 2, total_tokens: 5 }, }); } @@ -114,7 +122,11 @@ function installOAuthFetch( expires_in: 3600, }), { headers: { "content-type": "application/json" } }); } - if (url === CHAT_ENDPOINT) { + if (url === OAUTH_RESPONSES_ENDPOINT) { + const body = JSON.parse(String(init?.body)) as Record; + expect(body.model).toBe("grok-4.5"); + expect(body.input).toBe("hello"); + expect(body.messages).toBeUndefined(); chatAuth.push(new Headers(init?.headers).get("authorization") ?? ""); const status = chatStatuses.shift() ?? 200; if (status === 401) { @@ -209,7 +221,7 @@ describe("xAI OAuth upstream 401 replay", () => { const response = await post(server); const json = await response.json() as { error?: { message?: string } }; expect(response.status).toBe(401); - expect(json.error?.message).toContain("Provider error 401"); + expect(json.error?.message).toBe("rejected"); expect(observed.counts.refresh).toBe(1); expect(observed.chatAuth).toEqual(["Bearer rejected-access", "Bearer fresh-access"]); } finally { @@ -277,7 +289,7 @@ describe("xAI OAuth upstream 401 replay", () => { expires_in: 3600, }), { headers: { "content-type": "application/json" } }); } - if (url === CHAT_ENDPOINT) { + if (url === OAUTH_RESPONSES_ENDPOINT) { const bearer = new Headers(init?.headers).get("authorization") ?? ""; attemptsByBearer.set(bearer, (attemptsByBearer.get(bearer) ?? 0) + 1); if (bearer === "Bearer rejected-access") { diff --git a/tests/server-xai-responses-streaming.test.ts b/tests/server-xai-responses-streaming.test.ts new file mode 100644 index 0000000000..7195436a60 --- /dev/null +++ b/tests/server-xai-responses-streaming.test.ts @@ -0,0 +1,223 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { saveConfig } from "../src/config"; +import { saveCredential } from "../src/oauth/store"; +import { + XAI_GROK_CLI_BASE_URL, + XAI_GROK_CLIENT_VERSION, +} from "../src/providers/xai-transport"; +import { startServer } from "../src/server"; +import type { OcxConfig } from "../src/types"; +import { installIsolatedCodexHome, type IsolatedCodexHome } from "./helpers/isolated-codex-home"; + +const RESPONSES_ENDPOINT = `${XAI_GROK_CLI_BASE_URL}/responses`; +const encoder = new TextEncoder(); + +let testDir = ""; +let previousHome: string | undefined; +let isolatedCodexHome: IsolatedCodexHome | null = null; +let originalFetch: typeof fetch; + +beforeEach(async () => { + originalFetch = globalThis.fetch; + previousHome = process.env.OPENCODEX_HOME; + isolatedCodexHome = installIsolatedCodexHome("ocx-xai-responses-codex-"); + testDir = mkdtempSync(join(tmpdir(), "ocx-xai-responses-")); + process.env.OPENCODEX_HOME = testDir; + await saveCredential("xai", { + access: "stream-access", + refresh: "stream-refresh", + expires: Date.now() + 3_600_000, + accountId: "xai-stream-account", + source: "oauth", + }); +}); + +afterEach(() => { + globalThis.fetch = originalFetch; + if (previousHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = previousHome; + isolatedCodexHome?.restore(); + isolatedCodexHome = null; + if (testDir) rmSync(testDir, { recursive: true, force: true }); +}); + +function config(): OcxConfig { + return { + port: 0, + hostname: "127.0.0.1", + defaultProvider: "xai", + fastMode: true, + providers: { + xai: { + adapter: "openai-chat", + baseUrl: "https://api.x.ai/v1", + authMode: "oauth", + models: ["grok-4.6"], + }, + }, + } as OcxConfig; +} + +function sse(payload: unknown): Uint8Array { + return encoder.encode(`data: ${JSON.stringify(payload)}\n\n`); +} + +describe("xAI OAuth Responses streaming", () => { + test("uses the native Responses wire and relays the first delta before completion", async () => { + let releaseCompletion!: () => void; + const completionGate = new Promise(resolve => { releaseCompletion = resolve; }); + let completionReleased = false; + let outboundBody: Record | undefined; + let outboundHeaders: Headers | undefined; + let upstreamCalls = 0; + + globalThis.fetch = (async (input, init) => { + const url = input instanceof Request ? input.url : String(input); + if (url !== RESPONSES_ENDPOINT) return originalFetch(input, init); + upstreamCalls += 1; + outboundHeaders = new Headers(init?.headers); + outboundBody = JSON.parse(String(init?.body)) as Record; + + const body = new ReadableStream({ + start(controller) { + controller.enqueue(sse({ + type: "response.created", + sequence_number: 0, + response: { + id: "resp_xai_stream", + object: "response", + status: "in_progress", + model: "grok-4.6", + output: [], + }, + })); + controller.enqueue(sse({ + type: "response.output_item.added", + sequence_number: 1, + output_index: 0, + item: { id: "msg_xai_stream", type: "message", status: "in_progress", role: "assistant", content: [] }, + })); + controller.enqueue(sse({ + type: "response.content_part.added", + sequence_number: 2, + item_id: "msg_xai_stream", + output_index: 0, + content_index: 0, + part: { type: "output_text", text: "", annotations: [] }, + })); + controller.enqueue(sse({ + type: "response.output_text.delta", + sequence_number: 3, + item_id: "msg_xai_stream", + output_index: 0, + content_index: 0, + delta: "first", + })); + void completionGate.then(() => { + completionReleased = true; + const message = { + id: "msg_xai_stream", + type: "message", + status: "completed", + role: "assistant", + content: [{ type: "output_text", text: "first second", annotations: [] }], + }; + controller.enqueue(sse({ + type: "response.output_text.delta", + sequence_number: 4, + item_id: "msg_xai_stream", + output_index: 0, + content_index: 0, + delta: " second", + })); + controller.enqueue(sse({ + type: "response.output_item.done", + sequence_number: 5, + output_index: 0, + item: message, + })); + controller.enqueue(sse({ + type: "response.completed", + sequence_number: 6, + response: { + id: "resp_xai_stream", + object: "response", + status: "completed", + model: "grok-4.6", + output: [message], + usage: { input_tokens: 1, output_tokens: 2, total_tokens: 3 }, + }, + })); + controller.close(); + }); + }, + }); + return new Response(body, { headers: { "content-type": "text/event-stream" } }); + }) as typeof fetch; + + saveConfig(config()); + const server = startServer(0); + let reader: ReadableStreamDefaultReader | undefined; + try { + const response = await originalFetch(new URL("/v1/responses", server.url), { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + model: "xai/grok-4.6", + input: "hello", + stream: true, + store: false, + service_tier: "priority", + reasoning: { effort: "xhigh", summary: "auto" }, + }), + }); + expect(response.status).toBe(200); + reader = response.body!.getReader(); + const decoder = new TextDecoder(); + let received = ""; + await Promise.race([ + (async () => { + while (!received.includes("response.output_text.delta")) { + const chunk = await reader!.read(); + if (chunk.done) throw new Error("stream ended before the first xAI delta"); + received += decoder.decode(chunk.value, { stream: true }); + } + })(), + new Promise((_, reject) => setTimeout( + () => reject(new Error("the first xAI delta was not relayed before completion")), + 1_500, + )), + ]); + + expect(received).toContain("first"); + expect(completionReleased).toBe(false); + expect(upstreamCalls).toBe(1); + expect(outboundBody?.model).toBe("grok-4.6"); + expect(outboundBody?.input).toBe("hello"); + expect(outboundBody?.stream).toBe(true); + expect(outboundBody?.service_tier).toBeUndefined(); + expect(outboundBody?.reasoning).toMatchObject({ effort: "xhigh" }); + expect(outboundBody?.messages).toBeUndefined(); + expect(outboundBody?.reasoning_effort).toBeUndefined(); + expect(outboundHeaders?.get("authorization")).toBe("Bearer stream-access"); + expect(outboundHeaders?.get("x-grok-client-identifier")).toBe("opencodex"); + expect(outboundHeaders?.get("x-grok-client-version")).toBe(XAI_GROK_CLIENT_VERSION); + + releaseCompletion(); + while (true) { + const chunk = await reader.read(); + if (chunk.done) break; + received += decoder.decode(chunk.value, { stream: true }); + } + expect(received).toContain("response.completed"); + expect(received).toContain(" second"); + } finally { + releaseCompletion(); + await reader?.cancel().catch(() => {}); + await server.stop(true); + } + }, 10_000); +});