orynacode-ai 1.16.2

package/src/plugin/openai/ws-pool.ts

@@ -0,0 +1,290 @@
+import WebSocket from "ws"
+import * as Log from "@opencode-ai/core/util/log"
+import { ProviderError } from "@/provider/error"
+import { isRecord } from "@/util/record"
+import { OpenAIWebSocket } from "./ws"
+
+export const TITLE_HEADER = "x-opencode-title"
+
+const log = Log.create({ service: "plugin.openai.ws" })
+
+export interface CreateWebSocketFetchOptions {
+  httpFetch?: typeof globalThis.fetch
+  url?: string
+  connectTimeout?: number
+  idleTimeout?: number
+  maxConnectionAge?: number
+  streamRetries?: number
+}
+
+interface PoolEntry {
+  socket?: WebSocket
+  connectedAt?: number
+  lastUsedAt: number
+  busy: boolean
+  fallback: boolean
+  streamFailures: number
+}
+
+const DEFAULT_CONNECT_TIMEOUT = 15_000
+const DEFAULT_IDLE_TIMEOUT = 5 * 60 * 1000
+const DEFAULT_MAX_CONNECTION_AGE = 55 * 60 * 1000
+const CONNECTION_LIMIT_REACHED_CODE = "websocket_connection_limit_reached"
+
+export function createWebSocketFetch(options?: CreateWebSocketFetchOptions) {
+  const httpFetch = options?.httpFetch ?? globalThis.fetch
+  const pool = new Map<string, PoolEntry>()
+  const connectTimeout = options?.connectTimeout ?? DEFAULT_CONNECT_TIMEOUT
+  const idleTimeout = options?.idleTimeout ?? DEFAULT_IDLE_TIMEOUT
+  const maxConnectionAge = options?.maxConnectionAge ?? DEFAULT_MAX_CONNECTION_AGE
+  const streamRetries = options?.streamRetries ?? 5
+  const pruneTimer = setInterval(() => prune(), Math.min(idleTimeout, 60_000))
+  if (typeof pruneTimer === "object" && "unref" in pruneTimer && typeof pruneTimer.unref === "function") {
+    pruneTimer.unref()
+  }
+
+  async function websocketFetch(input: RequestInfo | URL, init?: RequestInit): Promise<Response> {
+    const url = input instanceof URL ? input.toString() : typeof input === "string" ? input : input.url
+    const internalHeaders = OpenAIWebSocket.normalizeHeaders(init?.headers)
+    const httpInit = withoutInternalHeaders(init)
+
+    if (init?.method !== "POST" || !new URL(url).pathname.endsWith("/responses")) {
+      return httpFetch(input, httpInit)
+    }
+
+    const body = (() => {
+      try {
+        if (typeof init?.body !== "string") return undefined
+        const parsed = JSON.parse(init.body)
+        return typeof parsed === "object" && parsed !== null ? parsed : undefined
+      } catch {
+        return undefined
+      }
+    })()
+    if (!body?.stream) return httpFetch(input, httpInit)
+    if (internalHeaders[TITLE_HEADER] === "true") {
+      log.debug("http fallback", { reason: "title" })
+      return httpFetch(input, httpInit)
+    }
+
+    const sessionID = internalHeaders["x-session-affinity"] ?? internalHeaders["session-id"]
+    if (!sessionID) {
+      log.debug("http fallback", { reason: "missing_session" })
+      return httpFetch(input, httpInit)
+    }
+    const key = `${sessionID}:conversation`
+
+    const entry = pool.get(key) ?? { lastUsedAt: Date.now(), busy: false, fallback: false, streamFailures: 0 }
+    pool.set(key, entry)
+
+    if (entry.fallback) {
+      log.debug("http fallback", { key, reason: "fallback_active" })
+      return httpFetch(input, httpInit)
+    }
+    if (entry.busy) {
+      log.debug("http fallback", { key, reason: "busy" })
+      return httpFetch(input, httpInit)
+    }
+
+    entry.busy = true
+    entry.lastUsedAt = Date.now()
+    try {
+      entry.socket = await socket(
+        entry,
+        options?.url ?? url,
+        OpenAIWebSocket.normalizeHeaders(httpInit?.headers),
+        connectTimeout,
+        maxConnectionAge,
+        init?.signal,
+      )
+      let resolveFirstEvent: (event: boolean | OpenAIWebSocket.WrappedError) => void = () => {}
+      let rejectFirstEvent: (error: Error) => void = () => {}
+      const firstEvent = new Promise<boolean | OpenAIWebSocket.WrappedError>((resolve, reject) => {
+        resolveFirstEvent = resolve
+        rejectFirstEvent = reject
+      })
+      const response = OpenAIWebSocket.streamResponsesWebSocket({
+        socket: entry.socket,
+        body,
+        idleTimeout,
+        signal: init?.signal ?? undefined,
+        onFirstEvent: (error) => resolveFirstEvent(error ?? true),
+        onTerminal: (event) => {
+          entry.busy = false
+          entry.lastUsedAt = Date.now()
+          entry.streamFailures = 0
+          if (event.type !== "response.completed" && event.type !== "response.done") {
+            log.warn("websocket terminal failure", { key, type: event.type })
+            invalidate(entry)
+          }
+        },
+        onConnectionInvalid: (error) => {
+          log.warn("websocket invalidated", { key, error: error.message })
+          entry.busy = false
+          entry.lastUsedAt = Date.now()
+          if (!entry.fallback) recordStreamFailure(entry)
+          invalidate(entry)
+          resolveFirstEvent(false)
+        },
+        onAbort: (error) => {
+          log.debug("websocket aborted", { key })
+          entry.busy = false
+          entry.lastUsedAt = Date.now()
+          entry.streamFailures = 0
+          invalidate(entry)
+          rejectFirstEvent(error)
+        },
+        onRetryableTerminal: async (event) => {
+          const error = connectionLimitError(event)
+          if (!error) return undefined
+          log.warn("websocket connection limit reached", { key })
+          throw error
+        },
+      })
+      const first = await firstEvent
+      if (first !== false) {
+        if (first === true || first.status < 200 || first.status > 599) return response
+        return new Response(first.body, {
+          status: first.status,
+          headers: { "content-type": "application/json", ...first.headers },
+        })
+      }
+      if (!entry.fallback) return response
+      log.debug("http fallback", { key, reason: "websocket_retries_exhausted" })
+      return httpFetch(input, httpInit)
+    } catch (error) {
+      entry.busy = false
+      entry.lastUsedAt = Date.now()
+      if (OpenAIWebSocket.isAbortError(error)) {
+        entry.streamFailures = 0
+        invalidate(entry)
+        throw error
+      }
+
+      recordStreamFailure(entry)
+      log.warn("websocket setup failed", {
+        key,
+        error: error instanceof Error ? error.message : String(error),
+        fallback: entry.fallback ? "http" : undefined,
+      })
+      invalidate(entry)
+      if (entry.fallback) return httpFetch(input, httpInit)
+      return failedResponse(
+        new ProviderError.ResponseStreamError(error instanceof Error ? error.message : String(error), {
+          cause: error,
+        }),
+      )
+    }
+  }
+
+  function recordStreamFailure(entry: PoolEntry) {
+    entry.streamFailures++
+    // Codex counts retries after the initial failed WebSocket attempt.
+    if (entry.streamFailures > streamRetries) entry.fallback = true
+  }
+
+  function prune() {
+    const now = Date.now()
+    for (const [key, entry] of pool) {
+      if (entry.busy) continue
+      if (entry.fallback) continue
+      if (now - entry.lastUsedAt < idleTimeout) continue
+      log.debug("websocket idle prune", { key })
+      invalidate(entry)
+      pool.delete(key)
+    }
+  }
+
+  function close() {
+    log.debug("websocket pool close", { count: pool.size })
+    clearInterval(pruneTimer)
+    for (const entry of pool.values()) invalidate(entry)
+    pool.clear()
+  }
+
+  function remove(sessionID: string) {
+    const key = `${sessionID}:conversation`
+    const entry = pool.get(key)
+    if (!entry) return
+    log.debug("websocket pool remove", { key })
+    invalidate(entry)
+    pool.delete(key)
+  }
+
+  return Object.assign(websocketFetch, { close, remove })
+}
+
+function connectionLimitError(event: Record<string, unknown>) {
+  if (event.type !== "error" || !isRecord(event.error) || event.error.code !== CONNECTION_LIMIT_REACHED_CODE) return
+  return new Error(typeof event.error.message === "string" ? event.error.message : CONNECTION_LIMIT_REACHED_CODE)
+}
+
+function failedResponse(error: ProviderError.ResponseStreamError) {
+  return new Response(
+    new ReadableStream({
+      start(controller) {
+        controller.error(error)
+      },
+    }),
+    {
+      status: 200,
+      headers: { "content-type": "text/event-stream" },
+    },
+  )
+}
+
+async function socket(
+  entry: PoolEntry,
+  url: string,
+  headers: Record<string, string>,
+  connectTimeout: number,
+  maxConnectionAge: number,
+  signal?: AbortSignal | null,
+) {
+  if (
+    entry.socket?.readyState === WebSocket.OPEN &&
+    entry.connectedAt &&
+    Date.now() - entry.connectedAt < maxConnectionAge
+  ) {
+    return entry.socket
+  }
+
+  invalidate(entry)
+  const next = await OpenAIWebSocket.connectResponsesWebSocket({
+    url: OpenAIWebSocket.toWebSocketUrl(url),
+    headers,
+    timeout: connectTimeout,
+    signal: signal ?? undefined,
+  })
+  entry.connectedAt = Date.now()
+  return next
+}
+
+function invalidate(entry: PoolEntry) {
+  if (entry.socket) {
+    entry.socket.on("error", () => {})
+    entry.socket.terminate()
+    entry.socket = undefined
+  }
+  entry.connectedAt = undefined
+}
+
+export function withoutInternalHeaders<T extends { headers?: HeadersInit }>(init: T | undefined): T | undefined {
+  if (!init?.headers) return init
+  if (init.headers instanceof Headers) {
+    const headers = new Headers(init.headers)
+    headers.delete(TITLE_HEADER)
+    return { ...init, headers }
+  }
+
+  if (Array.isArray(init.headers)) {
+    return { ...init, headers: init.headers.filter((item) => item[0].toLowerCase() !== TITLE_HEADER) }
+  }
+
+  return {
+    ...init,
+    headers: Object.fromEntries(Object.entries(init.headers).filter(([key]) => key.toLowerCase() !== TITLE_HEADER)),
+  }
+}
+
+export * as OpenAIWebSocketPool from "./ws-pool"

package/src/plugin/openai/ws.ts

@@ -0,0 +1,381 @@
+// Low-level OpenAI Responses WebSocket protocol helpers. Session pooling,
+// fallback, and continuation state intentionally live above this file.
+
+import WebSocket from "ws"
+import { APICallError } from "ai"
+import { ProviderError } from "@/provider/error"
+import { errorMessage } from "@/util/error"
+import { ProxyEnv } from "@/util/proxy-env"
+import { isRecord } from "@/util/record"
+
+export const PROTOCOL_HEADER = "responses_websockets=2026-02-06"
+
+export interface ConnectResponsesWebSocketOptions {
+  url: string
+  headers: Record<string, string>
+  timeout?: number
+  signal?: AbortSignal
+}
+
+export interface StreamResponsesWebSocketOptions {
+  socket: WebSocket
+  body: Record<string, unknown>
+  idleTimeout?: number
+  signal?: AbortSignal
+  onFirstEvent?: (error?: WrappedError) => void
+  onComplete?: (event: Record<string, unknown>) => void
+  onTerminal?: (event: Record<string, unknown>) => void
+  onRetryableTerminal?: (event: Record<string, unknown>) => Promise<WebSocket | undefined>
+  onConnectionInvalid?: (error: ProviderError.ResponseStreamError) => void
+  onAbort?: (error: Error) => void
+}
+
+export interface WrappedError {
+  status: number
+  headers?: Record<string, string>
+  body: string
+}
+
+export function toWebSocketUrl(url: string) {
+  return url.replace(/^http/, "ws")
+}
+
+export function normalizeHeaders(headers: HeadersInit | undefined): Record<string, string> {
+  const result: Record<string, string> = {}
+  if (!headers) return result
+
+  if (headers instanceof Headers) {
+    headers.forEach((value, key) => {
+      result[key.toLowerCase()] = value
+    })
+    return result
+  }
+
+  if (Array.isArray(headers)) {
+    for (const [key, value] of headers) {
+      result[key.toLowerCase()] = value
+    }
+    return result
+  }
+
+  for (const [key, value] of Object.entries(headers)) {
+    if (value != null) result[key.toLowerCase()] = value
+  }
+  return result
+}
+
+export function isAbortError(error: unknown): error is DOMException {
+  return error instanceof DOMException && error.name === "AbortError"
+}
+
+export function connectResponsesWebSocket(options: ConnectResponsesWebSocketOptions) {
+  return new Promise<WebSocket>((resolve, reject) => {
+    if (options.signal?.aborted) {
+      reject(abortError(options.signal))
+      return
+    }
+
+    const headers: Record<string, string> = {
+      ...options.headers,
+      "openai-beta": options.headers["openai-beta"] ?? PROTOCOL_HEADER,
+    }
+    delete headers["content-length"]
+
+    // Bun does not apply HTTP(S)_PROXY to WebSockets unless the proxy is supplied explicitly.
+    const proxy =
+      typeof Bun === "undefined"
+        ? undefined
+        : ProxyEnv.getProxyForUrl(options.url.replace(/^wss:/, "https:").replace(/^ws:/, "http:"))
+    const connect = { headers, ...(proxy ? { proxy } : {}) }
+    const socket = new WebSocket(options.url, connect)
+    const timeout = options.timeout
+      ? setTimeout(() => {
+          cleanup()
+          socket.on("error", () => {})
+          socket.terminate()
+          reject(new Error("WebSocket connect timed out"))
+        }, options.timeout)
+      : undefined
+
+    function cleanup() {
+      if (timeout) clearTimeout(timeout)
+      socket.off("open", onOpen)
+      socket.off("error", onError)
+      socket.off("close", onClose)
+      options.signal?.removeEventListener("abort", onAbort)
+    }
+
+    function onOpen() {
+      cleanup()
+      resolve(socket)
+    }
+
+    function onError(error: unknown) {
+      socket.on("error", () => {})
+      cleanup()
+      reject(error instanceof Error ? error : new Error(errorMessage(error), { cause: error }))
+    }
+
+    function onClose(code: number, reason: Buffer) {
+      cleanup()
+      reject(new Error(closeMessage("WebSocket closed before open", code, reason)))
+    }
+
+    function onAbort() {
+      cleanup()
+      socket.on("error", () => {})
+      socket.terminate()
+      reject(abortError(options.signal))
+    }
+
+    socket.once("open", onOpen)
+    socket.once("error", onError)
+    socket.once("close", onClose)
+    options.signal?.addEventListener("abort", onAbort, { once: true })
+  })
+}
+
+export function streamResponsesWebSocket(options: StreamResponsesWebSocketOptions) {
+  const encoder = new TextEncoder()
+
+  let socket = options.socket
+  let controller: ReadableStreamDefaultController<Uint8Array> | undefined
+  let cleanupSocket = () => {}
+  let completed = false
+  let emitted = false
+  let idleTimer: ReturnType<typeof setTimeout> | undefined
+
+  function cleanup() {
+    if (idleTimer) clearTimeout(idleTimer)
+    cleanupSocket()
+    options.signal?.removeEventListener("abort", onAbort)
+  }
+
+  function terminateSocket(target = socket) {
+    target.on("error", () => {})
+    target.terminate()
+  }
+
+  function closeCompleted() {
+    cleanup()
+    controller?.enqueue(encoder.encode("data: [DONE]\n\n"))
+    controller?.close()
+  }
+
+  function invalidate(error: ProviderError.ResponseStreamError) {
+    if (completed) return
+    completed = true
+    cleanup()
+    options.onConnectionInvalid?.(error)
+    controller?.error(error)
+  }
+
+  function resetIdleTimeout(message: string) {
+    if (completed) return
+    if (!options.idleTimeout) return
+    if (idleTimer) clearTimeout(idleTimer)
+    idleTimer = setTimeout(() => invalidate(new ProviderError.ResponseStreamError(message)), options.idleTimeout)
+  }
+
+  async function onMessage(data: WebSocket.RawData, isBinary: boolean) {
+    if (completed) return
+    if (isBinary) {
+      invalidate(new ProviderError.ResponseStreamError("Unexpected binary WebSocket frame"))
+      return
+    }
+
+    const text = data.toString()
+    const event = (() => {
+      try {
+        const parsed = JSON.parse(text)
+        return typeof parsed === "object" && parsed !== null ? parsed : undefined
+      } catch {
+        return undefined
+      }
+    })()
+
+    if (event?.type === "error" && options.onRetryableTerminal) {
+      cleanupSocket()
+      if (idleTimer) clearTimeout(idleTimer)
+      idleTimer = undefined
+      try {
+        const next = await options.onRetryableTerminal(event)
+        if (completed) {
+          if (next) terminateSocket(next)
+          return
+        }
+        if (next) {
+          attach(next)
+          return
+        }
+      } catch (error) {
+        invalidate(
+          new ProviderError.ResponseStreamError(error instanceof Error ? error.message : String(error), {
+            cause: error,
+          }),
+        )
+        return
+      }
+    }
+
+    const wrappedError = parseWrappedError(event, text)
+    if (wrappedError && event) {
+      if (!emitted) options.onFirstEvent?.(wrappedError)
+      completed = true
+      cleanup()
+      options.onTerminal?.(event)
+      controller?.error(
+        new APICallError({
+          message: wrappedError.message,
+          url: socket.url,
+          requestBodyValues: options.body,
+          statusCode: wrappedError.status,
+          responseHeaders: wrappedError.headers,
+          responseBody: wrappedError.body,
+        }),
+      )
+      return
+    }
+
+    if (!emitted) options.onFirstEvent?.()
+    controller?.enqueue(
+      encoder.encode(
+        `${text
+          .split(/\r?\n/)
+          .map((line) => `data: ${line}`)
+          .join("\n")}\n\n`,
+      ),
+    )
+    emitted = true
+    resetIdleTimeout("idle timeout waiting for websocket")
+
+    if (!event) return
+
+    if (event.type === "response.completed" || event.type === "response.done") {
+      completed = true
+      options.onComplete?.(event)
+      options.onTerminal?.(event)
+      closeCompleted()
+      return
+    }
+
+    if (event.type === "response.failed" || event.type === "response.incomplete" || event.type === "error") {
+      completed = true
+      options.onTerminal?.(event)
+      closeCompleted()
+    }
+  }
+
+  function onError(error: Error) {
+    invalidate(new ProviderError.ResponseStreamError(error.message, { cause: error }))
+  }
+
+  function onClose(code: number, reason: Buffer) {
+    if (completed) return
+    invalidate(
+      new ProviderError.ResponseStreamError(closeMessage("WebSocket closed before response.completed", code, reason)),
+    )
+  }
+
+  function onAbort() {
+    const error = abortError(options.signal)
+    if (completed) return
+    completed = true
+    cleanup()
+    terminateSocket()
+    options.onAbort?.(error)
+    controller?.error(error)
+  }
+
+  function onCancel(reason: unknown) {
+    if (completed) return
+    completed = true
+    cleanup()
+    terminateSocket()
+    options.onAbort?.(cancelError(reason))
+  }
+
+  function attach(next: WebSocket) {
+    cleanupSocket()
+    socket = next
+    socket.on("message", onMessage)
+    socket.once("error", onError)
+    socket.once("close", onClose)
+    cleanupSocket = () => {
+      socket.off("message", onMessage)
+      socket.off("error", onError)
+      socket.off("close", onClose)
+    }
+    const { stream: _stream, background: _background, ...payload } = options.body
+    resetIdleTimeout("idle timeout sending websocket request")
+    socket.send(JSON.stringify({ type: "response.create", ...payload }), (error) => {
+      if (completed) return
+      resetIdleTimeout("idle timeout waiting for websocket")
+      if (error) invalidate(new ProviderError.ResponseStreamError(error.message, { cause: error }))
+    })
+  }
+
+  return new Response(
+    new ReadableStream<Uint8Array>({
+      start(next) {
+        controller = next
+        options.signal?.addEventListener("abort", onAbort, { once: true })
+
+        if (options.signal?.aborted) {
+          onAbort()
+          return
+        }
+
+        attach(socket)
+      },
+      cancel(reason) {
+        onCancel(reason)
+      },
+    }),
+    {
+      status: 200,
+      headers: { "content-type": "text/event-stream" },
+    },
+  )
+}
+
+function parseWrappedError(event: Record<string, unknown> | undefined, body: string) {
+  if (event?.type !== "error") return
+  const status = event.status ?? event.status_code
+  if (typeof status !== "number" || (status >= 200 && status < 300)) return
+  return {
+    status,
+    headers: isRecord(event.headers)
+      ? Object.fromEntries(
+          Object.entries(event.headers).flatMap(([key, value]) =>
+            typeof value === "string" || typeof value === "number" || typeof value === "boolean"
+              ? [[key, String(value)]]
+              : [],
+          ),
+        )
+      : undefined,
+    body,
+    message: isRecord(event.error) && typeof event.error.message === "string" ? event.error.message : `${status}`,
+  }
+}
+
+function cancelError(reason: unknown) {
+  if (isAbortError(reason)) return reason
+  if (reason instanceof Error) return reason
+  return new DOMException(typeof reason === "string" ? reason : "Aborted", "AbortError")
+}
+
+function abortError(signal: AbortSignal | undefined) {
+  const reason = signal?.reason
+  if (isAbortError(reason)) return reason
+  return new DOMException(reason instanceof Error ? reason.message : "Aborted", "AbortError")
+}
+
+function closeMessage(message: string, code: number, reason: Buffer) {
+  const details = [`code ${code}`]
+  if (code === 1009) details.push("message too big")
+  if (reason.length > 0) details.push(reason.toString())
+  return `${message} (${details.join(": ")})`
+}
+
+export * as OpenAIWebSocket from "./ws"