mixdog 0.9.52 → 0.9.53

package/src/runtime/agent/orchestrator/providers/openai-compat-stream.mjs

@@ -11,6 +11,20 @@ import { customToolCallFromResponseItem } from './custom-tool-wire.mjs';
 import { createLeakGuard, createToolCallDedupe, dedupeToolCallList } from './anthropic-leaked-toolcall.mjs';
 import { randomBytes } from 'crypto';
 import { createActiveToolItemTracker } from './tool-stream-state.mjs';
+import {
+    truncatedCompatStreamError,
+    makeInvalidToolArgsMarker,
+    isInvalidToolArgsMarker,
+    formatInvalidToolArgsResult,
+    parseCompletedToolCallArgumentsJson,
+} from './lib/openai-tool-args.mjs';
+
+export {
+    makeInvalidToolArgsMarker,
+    isInvalidToolArgsMarker,
+    formatInvalidToolArgsResult,
+    parseCompletedToolCallArgumentsJson,
+};
 
 // Synthesize a native-shaped OpenAI tool call from a recovered leaked call.
 // Matches the `call_...` id scheme the native Responses/Chat paths use so the
@@ -25,77 +39,6 @@ function synthLeakedOpenAICall(recovered) {
     };
 }
 
-function truncatedCompatStreamError(label, detail) {
-    return Object.assign(
-        new Error(`${label} SSE stream truncated${detail ? `: ${detail}` : ''}`),
-        { name: 'TruncatedStreamError', code: 'TRUNCATED_STREAM', truncatedStream: true },
-    );
-}
-
-// Invalid-tool-args marker: completed-but-malformed tool_call arguments JSON must NOT throw
-// (kills the turn) NOR be silently swallowed to `{}`. Instead the parse
-// failure is carried as data on the tool call's `arguments` slot so the
-// dispatch loop can turn it into an is_error tool_result and let the model
-// re-issue the call with valid JSON in the SAME turn (follow-up retry).
-//   { __invalidToolArgs: true, __rawArguments: <raw string>, __parseError: <msg> }
-export function makeInvalidToolArgsMarker(rawArguments, parseError) {
-    return {
-        __invalidToolArgs: true,
-        __rawArguments: typeof rawArguments === 'string' ? rawArguments : String(rawArguments ?? ''),
-        __parseError: typeof parseError === 'string' ? parseError : String(parseError ?? 'parse error'),
-    };
-}
-export function isInvalidToolArgsMarker(value) {
-    return !!value && typeof value === 'object' && value.__invalidToolArgs === true;
-}
-/** Model-facing tool_result text for a tool call whose arguments failed to
- * parse; instructs the model to retry with valid JSON in the same turn. */
-export function formatInvalidToolArgsResult(call) {
-    const name = call?.name || 'tool';
-    const detail = call?.arguments?.__parseError || 'arguments were not valid JSON';
-    return `The arguments provided to \`${name}\` are invalid JSON and could not be parsed: ${detail}. Re-issue this tool call with valid JSON arguments.`;
-}
-
-/** Completed tool_call.arguments must be valid JSON; empty/missing → {}.
- * @param {any} raw - raw arguments value (string or object)
- * @param {string} label - provider label for error messages
- * @param {{id?:string,name?:string,index?:number,finishReason?:string}} [meta] - optional tool-call identity for diagnostics.
- *   When `meta.finishReason` is set, a completion/finish signal was observed for
- *   the call: a JSON.parse failure is then deterministic bad JSON (permanent),
- *   not a mid-stream truncation (retryable). */
-export function parseCompletedToolCallArgumentsJson(raw, label, meta) {
-    const text = typeof raw === 'string' ? raw : (raw == null ? '' : String(raw));
-    const src = text === '' ? '{}' : text;
-    try {
-        return JSON.parse(src);
-    } catch (err) {
-        const preview = text.length <= 64
-            ? text
-            : text.slice(0, 32) + '...' + text.slice(-32);
-        const detailParts = [`invalid tool_call arguments JSON: len=${text.length} preview=${JSON.stringify(preview)}`];
-        if (meta) {
-            const m = {};
-            if (meta.id) m.id = meta.id;
-            if (meta.name) m.name = meta.name;
-            if (meta.index != null) m.index = meta.index;
-            if (meta.finishReason) m.finishReason = meta.finishReason;
-            detailParts.push(`tool=${JSON.stringify(m)}`);
-        }
-        // Invariant: a completion/finish signal was observed for this tool call
-        // (finish_reason present, or a per-call/response "done" event fired), so
-        // the arguments are NOT mid-stream-truncated — they are complete but
-        // malformed. Return an invalid-args MARKER (not a throw) so the
-        // dispatch loop feeds the parse error back to the model as a
-        // tool_result and the model self-corrects in the same turn. Only an
-        // unfinished stream (no finishReason) stays the retryable truncation
-        // case — that transient behavior is deliberately preserved.
-        if (meta?.finishReason) {
-            return makeInvalidToolArgsMarker(text, err instanceof Error ? err.message : String(err));
-        }
-        throw truncatedCompatStreamError(label, detailParts.join(' '));
-    }
-}
-
 function firstByteCompatStreamError(label) {
     const err = providerTimeoutError(`${label} first byte`, PROVIDER_FIRST_BYTE_TIMEOUT_MS);
     err.firstByteTimeout = true;

package/src/runtime/agent/orchestrator/providers/openai-compat.mjs

@@ -3,6 +3,7 @@ import { getAgentApiKey } from '../../../shared/provider-api-key.mjs';
 import { withRetry } from './retry-classifier.mjs';
 import { getLlmDispatcher, preconnect } from '../../../shared/llm/http-agent.mjs';
 import { sendViaWebSocket } from './openai-oauth-ws.mjs';
+import { _combineUsageWithWarmup } from './openai-ws-events.mjs';
 import {
     consumeCompatChatCompletionStream,
     consumeCompatResponsesStream,
@@ -67,6 +68,36 @@ function loadOpenAI() {
     }
     return _OpenAI;
 }
+
+function attachCompletedWarmup(err, warmup) {
+    if (!err || !warmup?.usage) return err;
+    try {
+        Object.defineProperty(err, '__warmup', {
+            value: warmup,
+            configurable: true,
+            enumerable: false,
+        });
+    } catch {}
+    return err;
+}
+
+function includeCompletedXaiWarmup(result, warmup) {
+    if (!result || !warmup?.usage) return result;
+    const usage = _combineUsageWithWarmup(result.usage, warmup.usage, {
+        separateMainContext: true,
+    });
+    const ticks = usage?.raw?.cost_in_usd_ticks;
+    const costUsd = typeof ticks === 'number' && ticks >= 0
+        ? Number((ticks * 1e-10).toFixed(8))
+        : undefined;
+    return {
+        ...result,
+        usage: usage ? {
+            ...usage,
+            ...(costUsd != null ? { costUsd } : {}),
+        } : usage,
+    };
+}
 export { OPENAI_COMPAT_PRESETS } from './openai-compat-presets.mjs';
 const PRESETS = OPENAI_COMPAT_PRESETS;
 const MODEL_LIST_TIMEOUT_MS = resolveTimeoutMs(
@@ -257,7 +288,10 @@ export class OpenAICompatProvider {
                 }
                 process.stderr.write(`[provider] Auth error, re-reading provider authentication...\n`);
                 this.reloadApiKey();
-                return await this._doSend(messages, model, tools, sendOpts);
+                const retryOpts = this.name === 'xai' && err?.__warmup?.usage
+                    ? { ...(sendOpts || {}), _carriedWarmup: err.__warmup }
+                    : sendOpts;
+                return await this._doSend(messages, model, tools, retryOpts);
             }
             throw err;
         }
@@ -274,6 +308,15 @@ export class OpenAICompatProvider {
             this._preconnectFn(this.baseURL);
         }
         if (this.name === 'xai' && useXaiResponsesApi(opts, this.config)) {
+            const carriedWarmup = opts._carriedWarmup?.usage ? opts._carriedWarmup : null;
+            const sendHttpWithWarmup = async (warmup = carriedWarmup) => {
+                try {
+                    const result = await this._doSendXaiResponses(messages, useModel, tools, opts);
+                    return includeCompletedXaiWarmup(result, warmup);
+                } catch (err) {
+                    throw attachCompletedWarmup(err, warmup);
+                }
+            };
             // Shared Responses transport switch (MIXDOG_OAI_TRANSPORT), capability-
             // gated for xAI/Grok. Provider-local HTTP pins still win: Grok
             // proxy-only models set responsesTransport:'http' because the WS
@@ -315,12 +358,12 @@ export class OpenAICompatProvider {
                                 },
                             });
                         } catch {}
-                        return await this._doSendXaiResponses(messages, useModel, tools, opts);
+                        return await sendHttpWithWarmup(err?.__warmup || carriedWarmup);
                     }
                     throw err;
                 }
             }
-            return await this._doSendXaiResponses(messages, useModel, tools, opts);
+            return await sendHttpWithWarmup();
         }
         const signal = opts.signal || null;
         if (signal?.aborted) {
@@ -788,6 +831,7 @@ export class OpenAICompatProvider {
                 traceProvider: 'xai',
                 logSuppressedReasoningDeltas: false,
                 warmupBody,
+                _carriedWarmup: opts._carriedWarmup || null,
                 // Mirror openai-oauth fast fallback: when the HTTP fallback is
                 // enabled (outer catch → _shouldFallbackXaiWsToHttp), a first
                 // acquire/first-byte failure should skip remaining WS

package/src/runtime/agent/orchestrator/providers/openai-oauth-http-sse.mjs

@@ -33,6 +33,7 @@ import { createLeakGuard, createToolCallDedupe, dedupeToolCallList } from './ant
 import { customToolCallFromResponseItem } from './custom-tool-wire.mjs';
 import { CODEX_OAUTH_ORIGINATOR, CODEX_RESPONSES_URL, _displayCodexModel } from './openai-oauth.mjs';
 import { createActiveToolItemTracker } from './tool-stream-state.mjs';
+export { envPositiveInt as _envPositiveInt } from './lib/env-utils.mjs';
 
 // Public OpenAI Responses API endpoint for the api-key `openai` provider.
 // The openai-direct WS transport hits the same origin (openai-ws-pool
@@ -50,13 +51,6 @@ export function _envFlag(name, fallback = true) {
     return !['0', 'false', 'off', 'no'].includes(String(raw).toLowerCase());
 }
 
-export function _envPositiveInt(name, fallback) {
-    const raw = process.env[name];
-    if (raw == null || raw === '') return fallback;
-    const n = Number(raw);
-    return Number.isFinite(n) && n > 0 ? Math.floor(n) : fallback;
-}
-
 // Completed function_call.arguments parse for the OpenAI Responses stream.
 // A function_call item arrives only on a completion/done signal, so a
 // non-empty-but-malformed

package/src/runtime/agent/orchestrator/providers/openai-oauth-ws.mjs

@@ -63,6 +63,7 @@ import {
 } from './openai-ws-stream.mjs';
 import { _buildResponseCreateFrame } from './openai-ws-delta.mjs';
 import { resolveOpenAiTransportPolicy } from './openai-transport-policy.mjs';
+import { envPositiveInt } from './lib/env-utils.mjs';
 
 // Legacy import paths for scripts/tool-smoke.mjs, mixdog-session-runtime.mjs
 // (drainOpenaiWsPool), scripts/provider-toolcall-test.mjs (parseToolSearchArgs,
@@ -80,7 +81,6 @@ export {
 
 globalThis.__mixdogOpenaiWsRuntimeLoaded = true;
 
-// --- WS_PRE_RESPONSE_CREATED_MS / WS_INTER_CHUNK_MS: extracted to openai-ws-stream.mjs ---
 // The official Codex Responses policy has one five-retry stream budget shared
 // by connect/handshake and pre-output stream failures.
 const MIDSTREAM_WS_TRANSIENT_RETRY_LIMIT = MIDSTREAM_RETRY_POLICY.ws.transientCloseRetries;
@@ -103,10 +103,6 @@ const HANDSHAKE_MAX_ATTEMPTS = MIDSTREAM_WS_TRANSIENT_RETRY_LIMIT + 1;
 const HANDSHAKE_BACKOFF_BASE_MS = 200;
 const HANDSHAKE_BACKOFF_CAP_MS = 3200;
 
-// --- WS pool/handshake/acquire/release: extracted to openai-ws-pool.mjs ---
-
-// --- Delta/matching helpers + parseToolSearchArgs: extracted to openai-ws-stream.mjs ---
-// --- Stream consumer (_streamResponse): extracted to openai-ws-stream.mjs ---
 export function _classifyHandshakeError(err) {
     return classifyHandshakeError(err);
 }
@@ -189,11 +185,6 @@ function _num(value, fallback = 0) {
     return Number.isFinite(n) ? n : fallback;
 }
 
-function _envPositiveInt(name, fallback) {
-    const n = Number(process.env[name]);
-    return Number.isFinite(n) && n > 0 ? Math.floor(n) : fallback;
-}
-
 function _envRatio(name, fallback) {
     const n = Number(process.env[name]);
     return Number.isFinite(n) && n >= 0 && n <= 1 ? n : fallback;
@@ -409,8 +400,8 @@ function _cacheObservation({ entry, result, continuityResetReason = null }) {
     const promptTokens = _num(result?.usage?.promptTokens, 0) || inputTokens;
     const cachedTokens = _num(result?.usage?.cachedTokens, 0);
     const previousMaxCached = _num(entry?.promptCacheMaxCachedTokens, 0);
-    const warmThreshold = _envPositiveInt('MIXDOG_OAI_CACHE_MISS_WARM_TOKENS', 2048);
-    const promptThreshold = _envPositiveInt('MIXDOG_OAI_CACHE_MISS_PROMPT_TOKENS', 4096);
+    const warmThreshold = envPositiveInt('MIXDOG_OAI_CACHE_MISS_WARM_TOKENS', 2048);
+    const promptThreshold = envPositiveInt('MIXDOG_OAI_CACHE_MISS_PROMPT_TOKENS', 4096);
     const dropRatio = _envRatio('MIXDOG_OAI_CACHE_MISS_DROP_RATIO', 0.6);
     const dropThreshold = Math.floor(previousMaxCached * dropRatio);
     // A full-frame chain break (most commonly compaction/input-prefix rewrite)
@@ -456,7 +447,7 @@ function _requestInputExtends(previousInput, currentInput) {
     );
 }
 
-function _cacheContinuityResetReason({ mode, deltaReason, entry, body }) {
+function _cacheContinuityResetReason({ mode, deltaReason, entry, body, traceProvider }) {
     if (mode === 'delta') return null;
     if (deltaReason && !['no_anchor', 'full_forced', 'full_default'].includes(deltaReason)) {
         return deltaReason;
@@ -465,7 +456,9 @@ function _cacheContinuityResetReason({ mode, deltaReason, entry, body }) {
     // full_default. Re-run the two cheap snapshot checks so compaction or any
     // other prompt rewrite still retires the old prompt's cache high-water.
     if (deltaReason !== 'full_default' || !entry?.lastResponseId) return null;
-    const currentSansInput = _stableStringify(_sansInput(body));
+    const currentSansInput = _stableStringify(_sansInput(body, {
+        normalizeWarmupGenerate: traceProvider === 'openai-oauth',
+    }));
     if (entry.lastRequestSansInput && currentSansInput !== entry.lastRequestSansInput) {
         return 'request_properties_changed';
     }
@@ -652,6 +645,7 @@ export async function sendViaWebSocket({
     _sleepFn = _defaultSleep,
     _sendSpanTraceFn = appendAgentTrace,
     _agentTraceFn = appendAgentTrace,
+    _carriedWarmup = null,
 }) {
     // One bounded Codex stream retry budget covers transient handshake and
     // pre-output stream failures. Every retry acquires a fresh connection.
@@ -659,6 +653,21 @@ export async function sendViaWebSocket({
     const MAX_MIDSTREAM_RETRIES = MIDSTREAM_WS_TRANSIENT_RETRY_LIMIT;
     let firstAttemptError = null;
     let firstAttemptClassifier = null;
+    // A generate:false prewarm is billable even if its main request later
+    // retries on a fresh socket or falls back to HTTP. Retain one completed
+    // result across the whole logical send and attach it to terminal errors.
+    let completedWarmup = _carriedWarmup?.usage ? _carriedWarmup : null;
+    const _stampWarmup = (err) => {
+        if (!err || !completedWarmup?.usage) return err;
+        try {
+            Object.defineProperty(err, '__warmup', {
+                value: completedWarmup,
+                configurable: true,
+                enumerable: false,
+            });
+        } catch {}
+        return err;
+    };
     // Known tool names for the leaked-tool-call guard in _streamResponse.
     // Derived from the exact request body so a recovered leaked call only
     // synthesizes when it names a tool actually offered to this request.
@@ -785,6 +794,7 @@ export async function sendViaWebSocket({
                 onBackoffSlept: (ms) => { sendSpan.retryBackoffMs += ms; },
             });
         } catch (err) {
+            _stampWarmup(err);
             sendSpan.poolAcquireMs += performance.now() - handshakeStart;
             // Provenance only; policy remains provider-owned below. This lets
             // the direct wrapper distinguish an upgrade rejection from an
@@ -850,7 +860,7 @@ export async function sendViaWebSocket({
                 } catch (sleepErr) {
                     sendSpan.retryBackoffMs += performance.now() - sleepStart;
                     emitSendSpan('error');
-                    throw sleepErr;
+                    throw _stampWarmup(sleepErr);
                 }
                 sendSpan.retryBackoffMs += performance.now() - sleepStart;
                 continue;
@@ -951,28 +961,24 @@ export async function sendViaWebSocket({
             // codex prewarm gate (client.rs:1686-1688): only when the session
             // has no prior request state. A reused pooled socket with a live
             // chain must go straight to the real request.
-            if (warmupBody && typeof warmupBody === 'object' && attemptIndex === 0 && !entry.lastResponseId) {
+            if (warmupBody && typeof warmupBody === 'object' && !completedWarmup
+                && attemptIndex === 0 && !entry.lastResponseId) {
                 const warmupBuildStart = performance.now();
-                // Codex WS prewarm parity (opt-in): codex's prewarm frame is a
-                // minimal generate:false request that omits transport-only
-                // fields (stream/background) on the wire (prewarm_websocket,
-                // client.rs:1673-1705). Under MIXDOG_OAI_CODEX_WIRE_PARITY force
-                // that omission for the warmup frame too; default is unchanged.
-                const warmupWireParity = process.env.MIXDOG_OAI_CODEX_WIRE_PARITY === '1';
-                const parityWarmupBody = warmupWireParity
-                    ? { ...warmupBody, input: [], generate: false }
-                    : warmupBody;
-                const warmupFrame = _buildResponseCreateFrame(parityWarmupBody, { omitTransportFields: warmupWireParity });
-                const warmupMetadataContext = warmupWireParity
-                    ? {
-                        ...codexMetadataContext,
-                        sendOpts: {
-                            ...(codexMetadataContext?.sendOpts || {}),
-                            requestKind: 'prewarm',
-                            codexRequestKind: 'prewarm',
-                        },
-                    }
-                    : codexMetadataContext;
+                // Codex startup prewarm contains stable instructions/tools but
+                // no live user/transcript input. Enforce that at the transport
+                // boundary as well as in the OAuth caller so a future caller
+                // cannot accidentally duplicate the transcript. WS-only fields
+                // are omitted exactly as in Codex's prewarm request.
+                const parityWarmupBody = { ...warmupBody, input: [], generate: false };
+                const warmupFrame = _buildResponseCreateFrame(parityWarmupBody, { omitTransportFields: true });
+                const warmupMetadataContext = {
+                    ...codexMetadataContext,
+                    sendOpts: {
+                        ...(codexMetadataContext?.sendOpts || {}),
+                        requestKind: 'prewarm',
+                        codexRequestKind: 'prewarm',
+                    },
+                };
                 const wireWarmupFrame = _withCodexWsClientMetadata(warmupFrame, entry, useCodexWsClientMetadata, warmupMetadataContext);
                 wireFrameHadTurnState = !!wireWarmupFrame?.client_metadata?.['x-codex-turn-state'];
                 wireFrameMetadataTrace = _metadataTrace(wireWarmupFrame?.client_metadata);
@@ -1010,8 +1016,15 @@ export async function sendViaWebSocket({
                 if (!warmupResult?.responseId) {
                     throw new Error('Responses WS warmup completed without response id');
                 }
+                completedWarmup = {
+                    requestBody: parityWarmupBody,
+                    responseId: warmupResult.responseId,
+                    usage: warmupResult.usage,
+                };
                 entry.lastResponseId = warmupResult.responseId;
-                entry.lastRequestSansInput = _stableStringify(_sansInput(parityWarmupBody));
+                entry.lastRequestSansInput = _stableStringify(_sansInput(parityWarmupBody, {
+                    normalizeWarmupGenerate: useCodexWsClientMetadata,
+                }));
                 const warmupInputArr = Array.isArray(parityWarmupBody.input) ? parityWarmupBody.input : [];
                 entry.lastRequestInput = _cloneJson(warmupInputArr);
                 entry.lastResponseItems = _cloneJson(Array.isArray(warmupResult.responseItems) ? warmupResult.responseItems : []);
@@ -1039,38 +1052,15 @@ export async function sendViaWebSocket({
                         payload: warmupPayload,
                     });
                 } catch {}
-                // Do NOT rewrite the main request after warmup (R23 finding).
-                // The old prev_id+no-instructions rewrite made it=1's frame a
-                // different shape from it=2+ full frames, so the first real
-                // full-frame cache write only landed at it=2 and it=3 raced
-                // its propagation (cached=0 early misses). Keeping it=1 as a
-                // normal full frame makes it byte-identical to the warmup's
-                // prefix (instant full hit on the cache warmup just wrote)
-                // and keeps every subsequent frame one consistent shape.
-                // Delta opt-in still chains via entry.lastResponseId above.
             }
 
-            // Warmup writes the same prefix with generate:false, but the first
-            // real response must still be a FULL generating frame. Reusing the
-            // warmup response_id here would make _computeDelta reduce the frame
-            // input to [] when the warmup input matches, and a generate:false
-            // warmup is not a chainable response to continue from — so the first
-            // real turn would generate from an empty frame. Keep the warmup state
-            // for cache/trace, but compute the main frame as cold (full input +
-            // instructions).
-            const deltaEntry = warmupResult
-                ? {
-                    ...entry,
-                    lastResponseId: null,
-                    lastRequestSansInput: null,
-                    lastRequestInput: null,
-                    lastResponseItems: null,
-                    lastInputLen: 0,
-                    lastInputPrefixHash: null,
-                }
-                : entry;
+            // A completed generate:false prewarm is a valid continuation
+            // anchor. Compute against its retained empty-input snapshot so the
+            // first real request sends previous_response_id plus exactly the
+            // real incremental input. _computeDelta still retreats to a full
+            // frame on every missing anchor/property/prefix/output mismatch.
             const requestBuildStart = performance.now();
-            const delta = _computeDelta({ entry: deltaEntry, body: requestBody, traceProvider });
+            const delta = _computeDelta({ entry, body: requestBody, traceProvider });
             ({ mode, frame } = delta);
             deltaReason = delta.reason || null;
             strippedResponseItems = delta.strippedResponseItems || 0;
@@ -1138,6 +1128,7 @@ export async function sendViaWebSocket({
                 knownToolNames,
             });
         } catch (err) {
+            _stampWarmup(err);
             // Preserve failure provenance for the direct provider. This marker
             // is observational for OAuth and does not change its classifier or
             // retry budget.
@@ -1210,7 +1201,7 @@ export async function sendViaWebSocket({
                 } catch (sleepErr) {
                     sendSpan.retryBackoffMs += performance.now() - sleepStart;
                     emitSendSpan('error');
-                    throw sleepErr;
+                    throw _stampWarmup(sleepErr);
                 }
                 sendSpan.retryBackoffMs += performance.now() - sleepStart;
                 continue;
@@ -1298,10 +1289,13 @@ export async function sendViaWebSocket({
             deltaReason,
             entry,
             body: requestBody,
+            traceProvider,
         });
         if (result.responseId && keepResponseChain) {
             entry.lastResponseId = result.responseId;
-            entry.lastRequestSansInput = _stableStringify(_sansInput(requestBody));
+            entry.lastRequestSansInput = _stableStringify(_sansInput(requestBody, {
+                normalizeWarmupGenerate: useCodexWsClientMetadata,
+            }));
             const inputArr = Array.isArray(requestBody.input) ? requestBody.input : [];
             entry.lastRequestInput = _cloneJson(inputArr);
             entry.lastResponseItems = _cloneJson(Array.isArray(result.responseItems) ? result.responseItems : []);
@@ -1329,8 +1323,13 @@ export async function sendViaWebSocket({
             result,
             continuityResetReason: cacheContinuityResetReason,
         });
-        if (warmupResult?.usage) {
-            result.usage = _combineUsageWithWarmup(result.usage, warmupResult.usage);
+        if (completedWarmup?.usage) {
+            result.usage = _combineUsageWithWarmup(result.usage, completedWarmup.usage, {
+                // xAI/Grok prewarm is billable just like Codex prewarm, but it
+                // is not part of the real request's context footprint. Direct
+                // OpenAI intentionally retains its existing usage shape.
+                separateMainContext: useCodexWsClientMetadata || traceProvider === 'xai',
+            });
         }
 
         const requestedServiceTier = body?.service_tier || null;
@@ -1539,14 +1538,10 @@ export async function sendViaWebSocket({
         releaseWebSocket({ entry, poolKey, keep: keepSocket });
         const { responseId: _ignored, responseItems: _responseItemsIgnored, closeSocket: _closeSocketIgnored, ...out } = result;
         if (includeResponseId && result.responseId) out.responseId = result.responseId;
-        if (warmupResult) {
+        if (completedWarmup) {
             try {
                 Object.defineProperty(out, '__warmup', {
-                    value: {
-                        requestBody,
-                        responseId: warmupResult.responseId,
-                        usage: warmupResult.usage,
-                    },
+                    value: completedWarmup,
                     enumerable: false,
                 });
             } catch {}
@@ -1558,5 +1553,5 @@ export async function sendViaWebSocket({
         return out;
     }
     // Unreachable — the loop either returns or throws above.
-    throw _stampTool(_stampLiveText(firstAttemptError || new Error('sendViaWebSocket: unreachable')));
+    throw _stampWarmup(_stampTool(_stampLiveText(firstAttemptError || new Error('sendViaWebSocket: unreachable'))));
 }

package/src/runtime/agent/orchestrator/providers/openai-oauth.mjs

@@ -18,6 +18,7 @@ import { boundProviderAuthPath } from '../../../shared/provider-auth-binding.mjs
 import { makeModelCache } from './model-cache.mjs';
 
 import { sendViaWebSocket } from './openai-oauth-ws.mjs';
+import { _combineUsageWithWarmup } from './openai-ws-events.mjs';
 import { resolveOpenAiTransportPolicy } from './openai-transport-policy.mjs';
 import {
     buildStableProviderPromptCacheKey,
@@ -72,7 +73,7 @@ import {
 } from './openai-codex-model.mjs';
 export { _displayCodexModel };
 
-// Legacy import path for scripts/tool-smoke.mjs (single-emit SSE smoke).
+// Public test/integration entry retained alongside the transport module export.
 export { sendViaHttpSse };
 // --- Constants ---
 const CLIENT_ID = 'app_EMoamEEZ73f0CkXaXp7hrann';
@@ -351,10 +352,6 @@ async function refreshTokens(refreshToken) {
     }
 }
 // --- Build Responses API request ---
-function _cloneJson(value) {
-    try { return JSON.parse(JSON.stringify(value)); } catch { return value; }
-}
-
 function _contentTextParts(content, type = 'input_text') {
     if (typeof content === 'string') return content ? [{ type, text: content }] : [];
     if (!Array.isArray(content)) {
@@ -631,8 +628,6 @@ export function buildRequestBody(messages, model, tools, sendOpts) {
     return ordered;
 }
 
-// --- HTTP/SSE fallback transport: extracted to openai-oauth-http-sse.mjs ---
-
 // --- Provider ---
 export class OpenAIOAuthProvider {
     // OpenAI input_tokens already INCLUDES cached_tokens (cached is a subset),
@@ -895,12 +890,17 @@ export class OpenAIOAuthProvider {
                 useModel,
                 fetchFn: opts._fetchFn,
             });
+            if (originalErr?.__warmup?.usage) {
+                result.usage = _combineUsageWithWarmup(result.usage, originalErr.__warmup.usage, {
+                    separateMainContext: true,
+                });
+            }
             if (process.env.MIXDOG_DEBUG_AGENT) {
                 process.stderr.write(`[agent-trace] provider-send-end elapsed=${Date.now() - _t1}ms result=ok transport=http-fallback\n`);
             }
             return recordLiveModel(result);
         };
-        const dispatchWs = (forceFresh = false) => sendWs({
+        const dispatchWs = (forceFresh = false, carriedWarmup = null) => sendWs({
             auth,
             body,
             sendOpts: opts,
@@ -922,18 +922,14 @@ export class OpenAIOAuthProvider {
             // the hot WS/cache path for temporary blips while still preventing
             // TUI-level hangs. Sticky HTTP fallback is only armed after this
             // bounded reconnect budget is exhausted.
-            // codex-parity prewarm (generate:false full frame on a fresh
-            // socket, wire-verified 2026-07-03). DEFAULT ON: R19(off) vs
-            // R20(on) A/B shows prewarm removes ALL early-session zero-cache
-            // misses (4 -> 0 at it<=3; zero-miss 3 -> 0) by writing the cache
-            // and waiting for completion before the first real turn, exactly
-            // like codex prewarm_websocket (client.rs:1673-1705). Cost is one
-            // small (~5k tok) generate:false request per fresh socket, far
-            // below the 5-12k uncached tokens each early miss burned.
-            // Mid-session partial drops are server-side and unaffected.
+            // Codex startup-prewarm parity: build from the stable request
+            // properties (instructions/tools/etc.) but never send the live
+            // transcript/user input. The completed generate:false response is
+            // retained by the WS transport and anchors the first real request.
             warmupBody: _envFlag('MIXDOG_OPENAI_OAUTH_WS_WARMUP', true)
-                ? { ...body, generate: false }
+                ? { ...body, input: [], generate: false }
                 : null,
+            _carriedWarmup: carriedWarmup,
         });
         const mustSurfaceFallbackError = (fallbackErr) => externalSignal?.aborted
             || fallbackErr?.name === 'AbortError'
@@ -974,7 +970,7 @@ export class OpenAIOAuthProvider {
                 this._refreshFallbackUntil = 0;
                 auth = await this.ensureAuth({ forceRefresh: true, reason: String(status) });
                 try {
-                    const result = await dispatchWs(true);
+                    const result = await dispatchWs(true, err?.__warmup || null);
                     if (process.env.MIXDOG_DEBUG_AGENT) { process.stderr.write(`[agent-trace] provider-send-end elapsed=${Date.now() - _t1}ms result=ok\n`); }
                     return recordLiveModel(result);
                 } catch (retryErr) {
@@ -1121,7 +1117,6 @@ export class OpenAIOAuthProvider {
     }
 }
 
-// --- OAuth PKCE login flow: extracted to openai-oauth-login.mjs ---
 const { beginOAuthLogin, loginOAuth } = createOpenAIOAuthLogin({
     clientId: CLIENT_ID,
     originator: CODEX_OAUTH_ORIGINATOR,

package/src/runtime/agent/orchestrator/providers/openai-ws-delta.mjs

@@ -20,16 +20,14 @@ import { createHash } from 'node:crypto';
 // If the cached request (sans input) matches the current one and the current
 // input starts with the cached input, return only the tail. Otherwise return
 // the full input (fresh turn).
-export function _sansInput(body) {
+export function _sansInput(body, { normalizeWarmupGenerate = false } = {}) {
     const { input: _ignored, previous_response_id: _prevIgnored, generate, ...rest } = body;
-    // Warmup/prewarm frames carry generate:false on the wire (codex prewarm
-    // marker, openai-oauth.mjs warmupBody). That flag must NOT count as a
-    // request-property change on the warmup->first-real comparison, or the
-    // first real turn always retreats to a full frame. Normalize away ONLY the
-    // warmup-only generate:false; any other generate value stays in the
-    // comparison snapshot so genuine differences still break the delta. The
-    // wire body is untouched — frames are built from the raw body, not this.
-    if (generate !== false && generate !== undefined) rest.generate = generate;
+    // Only OpenAI OAuth/Codex startup prewarm treats generate:false as a
+    // transport marker. Shared xAI callers retain it as a real request
+    // property so their existing property guard still falls back to full.
+    if (!(normalizeWarmupGenerate && generate === false) && generate !== undefined) {
+        rest.generate = generate;
+    }
     return rest;
 }
 
@@ -348,7 +346,9 @@ export function _computeDelta({ entry, body, traceProvider }) {
     if (!Array.isArray(entry.lastRequestInput)) {
         return { mode: 'full', reason: 'no_input_snapshot', frame: buildFrame(body) };
     }
-    const curSans = _stableStringify(_sansInput(body));
+    const curSans = _stableStringify(_sansInput(body, {
+        normalizeWarmupGenerate: traceProvider === 'openai-oauth',
+    }));
     if (curSans !== entry.lastRequestSansInput) {
         return { mode: 'full', reason: 'request_properties_changed', frame: buildFrame(body) };
     }