@earendil-works/pi-ai 0.82.1 → 0.84.0

package/dist/api/openai-codex-responses.js

@@ -44,6 +44,14 @@ const CODEX_RESPONSE_STATUSES = new Set([
     "queued",
     "in_progress",
 ]);
+function assertSuccessfulOutput(output) {
+    if (output.stopReason === "pending") {
+        throw new Error("Codex stream ended without a stop reason");
+    }
+    if (output.stopReason === "error" || output.stopReason === "aborted") {
+        throw new Error(output.errorMessage || "An unknown error occurred");
+    }
+}
 // ============================================================================
 // Retry Helpers
 // ============================================================================
@@ -155,7 +163,7 @@ export const stream = (model, context, options) => {
                 totalTokens: 0,
                 cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
             },
-            stopReason: "stop",
+            stopReason: "pending",
             timestamp: Date.now(),
         };
         try {
@@ -197,10 +205,11 @@ export const stream = (model, context, options) => {
                                 startEmitted = true;
                                 stream.push({ type: "start", partial: output });
                             }
-                        }, httpTimeoutMs, websocketConnectTimeoutMs, cacheSessionId, grammarToolInputProperties, options);
+                        }, httpTimeoutMs, websocketConnectTimeoutMs, cacheSessionId, accountId, grammarToolInputProperties, options);
                         if (options?.signal?.aborted) {
                             throw new Error("Request was aborted");
                         }
+                        assertSuccessfulOutput(output);
                         stream.push({
                             type: "done",
                             reason: output.stopReason,
@@ -260,7 +269,7 @@ export const stream = (model, context, options) => {
                     const headerTimeoutSignal = httpTimeoutMs !== undefined && httpTimeoutMs > 0 ? AbortSignal.timeout(httpTimeoutMs) : undefined;
                     const combinedSignal = combineAbortSignals([options?.signal, headerTimeoutSignal]);
                     try {
-                        response = await fetch(resolveCodexUrl(model.baseUrl), {
+                        response = await (options?.fetch ?? globalThis.fetch)(resolveCodexUrl(model.baseUrl), {
                             method: "POST",
                             headers: sseHeaders,
                             body: sseBody,
@@ -329,6 +338,7 @@ export const stream = (model, context, options) => {
             if (options?.signal?.aborted) {
                 throw new Error("Request was aborted");
             }
+            assertSuccessfulOutput(output);
             stream.push({ type: "done", reason: output.stopReason, message: output });
             stream.end();
         }
@@ -653,14 +663,14 @@ export function closeOpenAICodexWebSocketSessions(sessionId) {
         closeWebSocketSilently(entry.socket, 1000, "debug_close");
     };
     if (sessionId) {
-        const entry = websocketSessionCache.get(sessionId);
-        if (entry)
+        for (const entry of websocketSessionCache.get(sessionId)?.values() ?? [])
             closeEntry(entry);
         websocketSessionCache.delete(sessionId);
         return;
     }
-    for (const entry of websocketSessionCache.values()) {
-        closeEntry(entry);
+    for (const accountEntries of websocketSessionCache.values()) {
+        for (const entry of accountEntries.values())
+            closeEntry(entry);
     }
     websocketSessionCache.clear();
 }
@@ -744,7 +754,7 @@ function closeWebSocketSilently(socket, code = 1000, reason = "done") {
     }
     catch { }
 }
-function scheduleSessionWebSocketExpiry(sessionId, entry) {
+function scheduleSessionWebSocketExpiry(sessionId, accountId, entry) {
     if (entry.idleTimer) {
         clearTimeout(entry.idleTimer);
     }
@@ -752,7 +762,11 @@ function scheduleSessionWebSocketExpiry(sessionId, entry) {
         if (entry.busy)
             return;
         closeWebSocketSilently(entry.socket, 1000, "idle_timeout");
-        websocketSessionCache.delete(sessionId);
+        const accountEntries = websocketSessionCache.get(sessionId);
+        if (accountEntries?.get(accountId) === entry)
+            accountEntries.delete(accountId);
+        if (accountEntries?.size === 0)
+            websocketSessionCache.delete(sessionId);
     }, SESSION_WEBSOCKET_CACHE_TTL_MS);
 }
 async function connectWebSocket(url, headers, signal, connectTimeoutMs = DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS, env) {
@@ -823,7 +837,7 @@ async function connectWebSocket(url, headers, signal, connectTimeoutMs = DEFAULT
         }
     });
 }
-async function acquireWebSocket(url, headers, sessionId, signal, connectTimeoutMs, env) {
+async function acquireWebSocket(url, headers, sessionId, accountId, signal, connectTimeoutMs, env) {
     if (!sessionId) {
         const socket = await connectWebSocket(url, headers, signal, connectTimeoutMs, env);
         return {
@@ -832,7 +846,8 @@ async function acquireWebSocket(url, headers, sessionId, signal, connectTimeoutM
             release: () => closeWebSocketSilently(socket),
         };
     }
-    const cached = websocketSessionCache.get(sessionId);
+    let accountEntries = websocketSessionCache.get(sessionId);
+    const cached = accountEntries?.get(accountId);
     if (cached) {
         if (cached.idleTimer) {
             clearTimeout(cached.idleTimer);
@@ -840,7 +855,9 @@ async function acquireWebSocket(url, headers, sessionId, signal, connectTimeoutM
         }
         if (!cached.busy && isWebSocketSessionExpired(cached)) {
             closeWebSocketSilently(cached.socket, 1000, "connection_age_limit");
-            websocketSessionCache.delete(sessionId);
+            accountEntries?.delete(accountId);
+            if (accountEntries?.size === 0)
+                websocketSessionCache.delete(sessionId);
         }
         else if (!cached.busy && isWebSocketReusable(cached.socket)) {
             cached.busy = true;
@@ -851,11 +868,15 @@ async function acquireWebSocket(url, headers, sessionId, signal, connectTimeoutM
                 release: ({ keep } = {}) => {
                     if (!keep || !isWebSocketReusable(cached.socket)) {
                         closeWebSocketSilently(cached.socket);
-                        websocketSessionCache.delete(sessionId);
+                        const currentEntries = websocketSessionCache.get(sessionId);
+                        if (currentEntries?.get(accountId) === cached)
+                            currentEntries.delete(accountId);
+                        if (currentEntries?.size === 0)
+                            websocketSessionCache.delete(sessionId);
                         return;
                     }
                     cached.busy = false;
-                    scheduleSessionWebSocketExpiry(sessionId, cached);
+                    scheduleSessionWebSocketExpiry(sessionId, accountId, cached);
                 },
             };
         }
@@ -871,12 +892,19 @@ async function acquireWebSocket(url, headers, sessionId, signal, connectTimeoutM
         }
         if (!isWebSocketReusable(cached.socket)) {
             closeWebSocketSilently(cached.socket);
-            websocketSessionCache.delete(sessionId);
+            accountEntries?.delete(accountId);
+            if (accountEntries?.size === 0)
+                websocketSessionCache.delete(sessionId);
         }
     }
     const socket = await connectWebSocket(url, headers, signal, connectTimeoutMs, env);
     const entry = { socket, busy: true, createdAt: Date.now() };
-    websocketSessionCache.set(sessionId, entry);
+    accountEntries = websocketSessionCache.get(sessionId);
+    if (!accountEntries) {
+        accountEntries = new Map();
+        websocketSessionCache.set(sessionId, accountEntries);
+    }
+    accountEntries.set(accountId, entry);
     return {
         socket,
         entry,
@@ -886,13 +914,15 @@ async function acquireWebSocket(url, headers, sessionId, signal, connectTimeoutM
                 closeWebSocketSilently(entry.socket);
                 if (entry.idleTimer)
                     clearTimeout(entry.idleTimer);
-                if (websocketSessionCache.get(sessionId) === entry) {
+                const currentEntries = websocketSessionCache.get(sessionId);
+                if (currentEntries?.get(accountId) === entry)
+                    currentEntries.delete(accountId);
+                if (currentEntries?.size === 0)
                     websocketSessionCache.delete(sessionId);
-                }
                 return;
             }
             entry.busy = false;
-            scheduleSessionWebSocketExpiry(sessionId, entry);
+            scheduleSessionWebSocketExpiry(sessionId, accountId, entry);
         },
     };
 }
@@ -1112,8 +1142,8 @@ async function* startWebSocketOutputOnFirstEvent(events, onStart) {
         yield event;
     }
 }
-async function processWebSocketStream(url, body, headers, output, stream, model, onStart, idleTimeoutMs, websocketConnectTimeoutMs, cacheSessionId, grammarToolInputProperties, options) {
-    const { socket, entry, reused, release } = await acquireWebSocket(url, headers, cacheSessionId, options?.signal, websocketConnectTimeoutMs, options?.env);
+async function processWebSocketStream(url, body, headers, output, stream, model, onStart, idleTimeoutMs, websocketConnectTimeoutMs, cacheSessionId, accountId, grammarToolInputProperties, options) {
+    const { socket, entry, reused, release } = await acquireWebSocket(url, headers, cacheSessionId, accountId, options?.signal, websocketConnectTimeoutMs, options?.env);
     let keepConnection = true;
     const useCachedContext = options?.transport === "websocket-cached" || options?.transport === "auto";
     // ChatGPT Codex Responses rejects `store: true` ("Store must be set to false").