mixdog 0.9.3 → 0.9.4

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

@@ -0,0 +1,436 @@
+/**
+ * openai-ws-pool.mjs — WebSocket connection pool for the OpenAI OAuth provider.
+ *
+ * Extracted from openai-oauth-ws.mjs. Owns the socket pool singleton
+ * (_wsPool), handshake/open/acquire/release lifecycle, idle-close timers and
+ * the process-exit drain fence. openai-oauth-ws.mjs imports acquire/release/
+ * _sendFrame and re-exports the drain hooks for legacy import paths.
+ */
+import WebSocket from 'ws';
+import { errText } from '../../../shared/err-text.mjs';
+import { createHash, randomBytes } from 'crypto';
+import { codexUserAgent, codexVersionHeader } from './codex-client-meta.mjs';
+import {
+    PROVIDER_WS_ACQUIRE_TIMEOUT_MS,
+    PROVIDER_WS_HANDSHAKE_TIMEOUT_MS,
+    resolveTimeoutMs,
+} from '../stall-policy.mjs';
+
+// Human-readable transport label for handshake/acquire error messages. Shared
+// with openai-oauth-ws.mjs (stream-side errors use the same labels).
+export function _wsErrLabel(p) {
+    if (p === 'xai') return 'xAI WS';
+    if (p === 'openai-direct' || p === 'openai') return 'OpenAI WS';
+    return 'OpenAI OAuth WS';
+}
+
+const CODEX_WS_URL = 'wss://chatgpt.com/backend-api/codex/responses';
+const CODEX_OAUTH_ORIGINATOR = 'codex_cli_rs';
+const OPENAI_WS_URL = 'wss://api.openai.com/v1/responses';
+const XAI_WS_URL = 'wss://api.x.ai/v1/responses';
+export const WS_IDLE_MS = resolveTimeoutMs(
+    'MIXDOG_PROVIDER_WS_IDLE_MS',
+    20 * 60_000,
+    { minMs: 60_000, maxMs: 60 * 60_000 },
+);
+const WS_HANDSHAKE_TIMEOUT_MS = PROVIDER_WS_HANDSHAKE_TIMEOUT_MS;
+const WS_ACQUIRE_TIMEOUT_MS = PROVIDER_WS_ACQUIRE_TIMEOUT_MS;
+
+// WS socket pool buckets are keyed by `poolKey` (the per-call sessionId)
+// to isolate parallel agent invocations — each gets its own socket so
+// a second caller cannot grab a sibling's mid-turn entry (openai-oauth would
+// otherwise reject the new response.create with "No tool output found
+// for function call ..."). The handshake `session_id` header/URL
+// uses `cacheKey` — for OpenAI OAuth this mirrors Codex's thread-scoped
+// prompt_cache_key by default. The backend dedupes cache by handshake
+// session_id, not by body.prompt_cache_key alone (measured 2026-04-19 after the
+// v0.6.151 regression).
+const MAX_POOLED_SOCKETS_PER_KEY = 8;
+
+// poolKey -> Entry[]
+// Entry: { socket, busy, idleTimer, lastResponseId, lastRequestSansInput,
+//          lastRequestInput, lastResponseItems, lastInputLen, turnState,
+//          closing, ephemeral }
+const _wsPool = new Map();
+
+function _getPoolArr(poolKey) {
+    if (!poolKey) return null;
+    let arr = _wsPool.get(poolKey);
+    if (!arr) {
+        arr = [];
+        _wsPool.set(poolKey, arr);
+    }
+    return arr;
+}
+
+function _removeFromPool(poolKey, entry) {
+    if (!poolKey) return;
+    const arr = _wsPool.get(poolKey);
+    if (!arr) return;
+    const idx = arr.indexOf(entry);
+    if (idx >= 0) arr.splice(idx, 1);
+    if (arr.length === 0) _wsPool.delete(poolKey);
+}
+
+function _scheduleIdleClose(poolKey, entry) {
+    if (!entry) return;
+    if (entry.idleTimer) clearTimeout(entry.idleTimer);
+    entry.idleTimer = setTimeout(() => {
+        if (entry.busy) return;
+        try { entry.socket.close(1000, 'idle_timeout'); } catch {}
+        _removeFromPool(poolKey, entry);
+    }, WS_IDLE_MS);
+    try { entry.idleTimer.unref?.(); } catch {}
+}
+
+function _clearIdle(entry) {
+    if (entry?.idleTimer) {
+        clearTimeout(entry.idleTimer);
+        entry.idleTimer = null;
+    }
+}
+
+function _isOpen(entry) {
+    return entry?.socket?.readyState === WebSocket.OPEN;
+}
+
+// Awaited frame send. Asserts the socket is OPEN and resolves only after
+// the underlying transport reports the buffered write succeeded (or fails)
+// via the WebSocket send callback. Raw `socket.send(JSON.stringify(...))`
+// is fire-and-forget — a wedged or half-closed socket silently queues the
+// payload and the caller assumes it landed, then later times out waiting
+// for a server event that will never arrive. Tag any failure with
+// `wsSendFailed=true` so _classifyMidstreamError routes the next attempt
+// through a fresh socket.
+export function _sendFrame(entry, frame) {
+    return new Promise((resolve, reject) => {
+        const socket = entry?.socket;
+        if (!socket || socket.readyState !== WebSocket.OPEN) {
+            const err = new Error(`WS send: socket not OPEN (readyState=${socket?.readyState ?? 'n/a'})`);
+            err.wsSendFailed = true;
+            reject(err);
+            return;
+        }
+        let payload;
+        try { payload = JSON.stringify(frame); }
+        catch (e) {
+            const err = e instanceof Error ? e : new Error(String(e));
+            err.wsSendFailed = true;
+            reject(err);
+            return;
+        }
+        try {
+            // Do NOT await the send callback: on a wedged-but-OPEN socket the
+            // ws write callback may never fire, which would hang this Promise
+            // before _streamResponse arms its first-byte watchdog. Fire and
+            // resolve immediately; transport failures surface via the socket
+            // 'error'/'close' handlers and the first-byte watchdog.
+            socket.send(payload, () => {});
+            resolve();
+        } catch (e) {
+            const err = e instanceof Error ? e : new Error(String(e));
+            err.wsSendFailed = true;
+            reject(err);
+        }
+    });
+}
+
+function _buildHandshakeHeaders({ auth, sessionToken, turnState, cacheKey: _cacheKey, codexHeaders }) {
+    // xAI WS: do NOT pin x-grok-conv-id. Measured parallel runs show that
+    // forcing a routing shard via that header alternates cold caches across
+    // parallel workers; the automatic prompt-prefix cache holds up better
+    // when each handshake is unpinned. Reference: vercel/ai xai provider.
+    const headers = auth.type === 'xai'
+        ? {
+            'Authorization': `Bearer ${auth.apiKey}`,
+        }
+        : auth.type === 'openai-direct'
+        ? {
+            'Authorization': `Bearer ${auth.apiKey}`,
+            'OpenAI-Beta': 'responses_websockets=2026-02-06',
+        }
+        : {
+            'Authorization': `Bearer ${auth.access_token}`,
+            'chatgpt-account-id': auth.account_id || '',
+            'originator': CODEX_OAUTH_ORIGINATOR,
+            'OpenAI-Beta': 'responses_websockets=2026-02-06',
+            // codex-rs merges provider http_headers ("version") plus
+            // default_headers (User-Agent) into the WS handshake
+            // (client.rs:970-975, model-provider-info/src/lib.rs:339-343,
+            // login default_client.rs:289-305). The backend fingerprints
+            // clients on these; missing them can route us onto a different
+            // (colder) cache-node class than codex.
+            'User-Agent': codexUserAgent(),
+            'version': codexVersionHeader(),
+        };
+    const isOpenAiOauth = auth.type !== 'xai' && auth.type !== 'openai-direct';
+    if (sessionToken && !isOpenAiOauth) {
+        // openai-direct keeps the legacy underscore header; codex-rs sends no
+        // underscore session_id on the OAuth WS handshake (client.rs:1033-1057),
+        // only the dashed session-id/thread-id pair below.
+        headers['session_id'] = String(sessionToken);
+    }
+    if (isOpenAiOauth && (sessionToken || _cacheKey)) {
+        const threadId = String(_cacheKey || sessionToken);
+        const sessionId = String(sessionToken || _cacheKey);
+        headers['session-id'] = sessionId;
+        headers['thread-id'] = threadId;
+        headers['x-client-request-id'] = threadId;
+        if (codexHeaders && typeof codexHeaders === 'object') {
+            for (const [key, value] of Object.entries(codexHeaders)) {
+                if (typeof key === 'string' && typeof value === 'string' && value) {
+                    headers[key] = value;
+                }
+            }
+        }
+    } else {
+        // xAI/direct keep a per-request value so their server-side traces stay
+        // distinguishable across reconnects.
+        headers['x-client-request-id'] = randomBytes(16).toString('hex');
+    }
+    if (turnState) headers['x-codex-turn-state'] = turnState;
+    return headers;
+}
+
+// handshake session_id is the conversation slot openai-oauth uses for in-memory
+// prefix state. OpenAI OAuth uses the Codex-style thread cache key by default;
+// xAI leaves routing unpinned.
+function _mintSessionToken(cacheKey, auth) {
+    // xAI's public WebSocket endpoint uses the open connection plus
+    // response ids for continuation; unlike openai-oauth, it does not need the
+    // OAuth-specific session_id handshake shard.
+    if (auth?.type === 'xai') return null;
+    return cacheKey || 'mixdog-default';
+}
+
+function _openSocket({ auth, sessionToken, turnState, externalSignal, cacheKey, codexHeaders }) {
+    const headers = _buildHandshakeHeaders({ auth, sessionToken, turnState, cacheKey, codexHeaders });
+    const baseUrl = auth.type === 'xai'
+        ? XAI_WS_URL
+        : auth.type === 'openai-direct'
+            ? OPENAI_WS_URL
+            : CODEX_WS_URL;
+    const _wsOpenStart = Date.now();
+    if (process.env.MIXDOG_DEBUG_AGENT) {
+        process.stderr.write(`[agent-trace] ws-open-start url=${baseUrl} tokenHash=${createHash('sha256').update(String(sessionToken)).digest('hex').slice(0, 8)} ts=${_wsOpenStart}\n`);
+    }
+    const url = baseUrl + (sessionToken ? `?session_id=${encodeURIComponent(String(sessionToken))}` : '');
+    return new Promise((resolve, reject) => {
+        let settled = false;
+        let abortListener = null;
+        let acquireTimer = null;
+        const settle = (ok, val) => {
+            if (settled) return;
+            settled = true;
+            if (acquireTimer) {
+                clearTimeout(acquireTimer);
+                acquireTimer = null;
+            }
+            if (abortListener && externalSignal) {
+                try { externalSignal.removeEventListener('abort', abortListener); } catch {}
+            }
+            (ok ? resolve : reject)(val);
+        };
+        const socket = new WebSocket(url, { headers, handshakeTimeout: WS_HANDSHAKE_TIMEOUT_MS });
+        acquireTimer = setTimeout(() => {
+            if (settled) return;
+            if (process.env.MIXDOG_DEBUG_AGENT) {
+                process.stderr.write(`[agent-trace] ws-open-fail kind=acquire_timeout timeoutMs=${WS_ACQUIRE_TIMEOUT_MS} elapsed=${Date.now() - _wsOpenStart}ms\n`);
+            }
+            try { socket.terminate(); } catch {}
+            settle(false, Object.assign(
+                new Error(`${_wsErrLabel(auth?.type === 'xai' ? 'xai' : auth?.type === 'openai-direct' ? 'openai-direct' : 'openai-oauth')} acquire timed out before open (${WS_ACQUIRE_TIMEOUT_MS}ms)`),
+                { code: 'EWSACQUIRETIMEOUT', acquireTimeoutMs: WS_ACQUIRE_TIMEOUT_MS },
+            ));
+        }, WS_ACQUIRE_TIMEOUT_MS);
+        try { acquireTimer.unref?.(); } catch {}
+        const capturedHeaders = { turnState: null };
+        socket.once('upgrade', (res) => {
+            try {
+                const ts = res?.headers?.['x-codex-turn-state'];
+                if (typeof ts === 'string' && ts.length) capturedHeaders.turnState = ts;
+            } catch {}
+        });
+        socket.once('open', () => {
+            if (process.env.MIXDOG_DEBUG_AGENT) {
+                process.stderr.write(`[agent-trace] ws-open-ok elapsed=${Date.now() - _wsOpenStart}ms\n`);
+            }
+            settle(true, { socket, turnState: capturedHeaders.turnState });
+        });
+        socket.once('error', (err) => {
+            if (process.env.MIXDOG_DEBUG_AGENT) {
+                process.stderr.write(`[agent-trace] ws-open-fail kind=error msg=${String(err?.message || err).slice(0, 120)} elapsed=${Date.now() - _wsOpenStart}ms\n`);
+            }
+            try { socket.terminate(); } catch {}
+            settle(false, err instanceof Error ? err : Object.assign(new Error(errText(err) || 'openai-oauth WS error'), { wsErrorEvent: true, original: err }));
+        });
+        socket.once('close', (code, reason) => {
+            // Half-open handshake: the peer closed before 'open'/'error' fired
+            // (TCP RST / TLS edge). Without this the connect Promise never
+            // settles and only the 600s outer watchdog can break the stall
+            // (observed stage=requesting 601s hang). Open-path closes are
+            // no-ops here because settle() has already flipped `settled`.
+            if (settled) return;
+            try { socket.terminate(); } catch {}
+            settle(false, Object.assign(
+                new Error(`${_wsErrLabel(auth?.type === 'xai' ? 'xai' : auth?.type === 'openai-direct' ? 'openai-direct' : 'openai-oauth')} handshake closed before open (code=${code})`),
+                { wsCloseCode: code, wsCloseReason: (reason && reason.toString) ? reason.toString('utf-8') : '' }));
+        });
+        socket.once('unexpected-response', (_req, res) => {
+            if (settled) return;
+            const status = res?.statusCode || 0;
+            let body = '';
+            res.on('data', c => { if (body.length < 2048) body += c.toString('utf-8'); });
+            res.on('end', () => {
+                if (process.env.MIXDOG_DEBUG_AGENT) {
+                    process.stderr.write(`[agent-trace] ws-open-fail kind=http status=${status} body=${body.slice(0, 120)} elapsed=${Date.now() - _wsOpenStart}ms\n`);
+                }
+                try { socket.terminate(); } catch {}
+                settle(false, Object.assign(new Error(`${_wsErrLabel(auth?.type === 'xai' ? 'xai' : auth?.type === 'openai-direct' ? 'openai-direct' : 'openai-oauth')} handshake ${status}: ${body.slice(0, 200)}`), { httpStatus: status, httpBody: body }));
+            });
+        });
+        if (externalSignal) {
+            const onAbort = () => {
+                try { socket.terminate(); } catch {}
+                const reason = externalSignal.reason;
+                settle(false, reason instanceof Error ? reason : new Error(`${_wsErrLabel(auth?.type === 'xai' ? 'xai' : auth?.type === 'openai-direct' ? 'openai-direct' : 'openai-oauth')} handshake aborted`));
+            };
+            if (externalSignal.aborted) { onAbort(); return; }
+            abortListener = onAbort;
+            externalSignal.addEventListener('abort', onAbort, { once: true });
+        }
+    });
+}
+
+export async function acquireWebSocket({ auth, poolKey, cacheKey, codexHeaders, forceFresh, externalSignal }) {
+    const _acqStart = Date.now();
+    if (process.env.MIXDOG_DEBUG_AGENT) {
+        process.stderr.write(`[agent-trace] acquire-start poolKey=${poolKey} cacheKey=${cacheKey} forceFresh=${forceFresh} externalAborted=${!!externalSignal?.aborted} ts=${_acqStart}\n`);
+    }
+    if (externalSignal?.aborted) {
+        const reason = externalSignal.reason;
+        throw reason instanceof Error ? reason : new Error('OpenAI OAuth WS acquire aborted');
+    }
+    if (poolKey && !forceFresh) {
+        const arr = _wsPool.get(poolKey) || [];
+        // Prune dead entries first.
+        for (let i = arr.length - 1; i >= 0; i--) {
+            if (!_isOpen(arr[i]) || arr[i].closing) {
+                _clearIdle(arr[i]);
+                arr.splice(i, 1);
+            }
+        }
+        if (arr.length === 0) _wsPool.delete(poolKey);
+        // Reuse any idle open entry (cache-warm path).
+        const idle = arr.find(e => !e.busy);
+        if (idle) {
+            _clearIdle(idle);
+            idle.busy = true;
+            // Defensive: pre-existing pooled entries created before the
+            // prefix-hash field was introduced may not have it set. Normalize
+            // to null so the first delta check reads a deterministic value
+            // (and falls back to full-create instead of silently passing).
+            if (idle.lastInputPrefixHash === undefined) idle.lastInputPrefixHash = null;
+            if (idle.lastRequestInput === undefined) idle.lastRequestInput = null;
+            if (idle.lastResponseItems === undefined) idle.lastResponseItems = null;
+            if (process.env.MIXDOG_DEBUG_AGENT) {
+                process.stderr.write(`[agent-trace] acquire-reuse poolKey=${poolKey} openSockets=${arr.length} elapsed=${Date.now() - _acqStart}ms\n`);
+            }
+            return { entry: idle, reused: true };
+        }
+        // All entries busy and bucket at cap: fall through to ephemeral socket.
+        if (arr.length >= MAX_POOLED_SOCKETS_PER_KEY) {
+            if (process.env.MIXDOG_DEBUG_AGENT) {
+                process.stderr.write(`[agent-trace] acquire-ephemeral cacheKey=${cacheKey} reason=cap elapsed=${Date.now() - _acqStart}ms\n`);
+            }
+            const ephSessionToken = _mintSessionToken(cacheKey, auth);
+            const { socket, turnState } = await _openSocket({ auth, sessionToken: ephSessionToken, turnState: null, externalSignal, cacheKey, codexHeaders });
+            // Drain-complete fence: same invariant as the normal acquire path —
+            // if drain fired during the await, do NOT push an ephemeral entry
+            // back into the pool.
+            if (_drainComplete) {
+                try { socket.close(1000, 'drain-complete'); } catch {}
+                throw new Error('WS pool drained — process exiting');
+            }
+            const entry = {
+                socket,
+                busy: true,
+                idleTimer: null,
+                lastResponseId: null,
+                lastRequestSansInput: null,
+                lastRequestInput: null,
+                lastResponseItems: null,
+                lastInputLen: 0,
+                lastInputPrefixHash: null,
+                turnState: turnState || null,
+                closing: false,
+                ephemeral: true,
+                sessionToken: ephSessionToken,
+            };
+            socket.on('close', () => { entry.closing = true; });
+            return { entry, reused: false };
+        }
+    }
+    // Parallel sockets must not inherit sibling turnState or the openai-oauth server
+    // treats the new request as a continuation of another in-flight turn and
+    // returns "No tool output found for function call …". turnState only
+    // propagates within a single entry across its own iterations.
+    const sessionToken = _mintSessionToken(cacheKey, auth);
+    if (process.env.MIXDOG_DEBUG_AGENT) {
+        process.stderr.write(`[agent-trace] acquire-new tokenHash=${createHash('sha256').update(String(sessionToken)).digest('hex').slice(0, 8)} elapsed=${Date.now() - _acqStart}ms\n`);
+    }
+    const { socket, turnState } = await _openSocket({ auth, sessionToken, turnState: null, externalSignal, cacheKey, codexHeaders });
+    const entry = {
+        socket,
+        busy: true,
+        idleTimer: null,
+        lastResponseId: null,
+        lastRequestSansInput: null,
+        lastRequestInput: null,
+        lastResponseItems: null,
+        lastInputLen: 0,
+        lastInputPrefixHash: null,
+        turnState: turnState || null,
+        closing: false,
+        ephemeral: false,
+        sessionToken,
+    };
+    if (poolKey && !forceFresh) _getPoolArr(poolKey).push(entry);
+    socket.on('close', () => {
+        entry.closing = true;
+        _removeFromPool(poolKey, entry);
+    });
+    return { entry, reused: false };
+}
+
+export function releaseWebSocket({ entry, poolKey, keep }) {
+    if (!entry) return;
+    entry.busy = false;
+    if (!keep || !_isOpen(entry) || !poolKey || entry.ephemeral) {
+        try { entry.socket.close(1000, keep ? 'no_session' : 'release_no_keep'); } catch {}
+        _removeFromPool(poolKey, entry);
+        return;
+    }
+    _scheduleIdleClose(poolKey, entry);
+}
+
+// Drain-complete fence — set true once _closeAllPooledSockets runs so any
+// in-flight acquire that resumes after drain throws instead of pushing a
+// fresh socket into the cleared pool. Single-set, process-lifetime invariant.
+let _drainComplete = false;
+
+// Drain hook — self-registered exit drain.
+// Force-closes pooled sockets and fences subsequent acquires.
+// `drainOpenaiWsPool` alias matches the registry's `drain*` naming convention;
+// `_closeAllPooledSockets` kept for backward compat with existing call sites.
+export function _closeAllPooledSockets(reason = 'shutdown') {
+    _drainComplete = true;
+    for (const arr of _wsPool.values()) {
+        for (const entry of arr) {
+            try { entry.socket.close(1000, reason); } catch {}
+        }
+    }
+    _wsPool.clear();
+}
+export const drainOpenaiWsPool = _closeAllPooledSockets;
+process.on('exit', drainOpenaiWsPool);