fix(ai): preserve websocket upgrade diagnostics

packages/ai/src/route/executor.ts

@@ -188,12 +188,14 @@ const statusError =
 
 // Classifies an HTTP failure captured outside the executor (for example by the
 // AI SDK's own fetch) onto the same reason types and HttpContext that
-// executor-driven requests produce. The originating request is not available on
-// that path, so the method is assumed (language model calls are always POST),
-// request headers are empty.
+// executor-driven requests produce. Callers may supply originating request
+// details when their transport exposes them; language model calls otherwise
+// default to POST with no request headers.
 export const classifyHttpFailure = (input: {
   readonly message: string
   readonly url: string
+  readonly method?: string | undefined
+  readonly requestHeaders?: Record<string, string> | undefined
   readonly status?: number | undefined
   readonly code?: string | undefined
   readonly responseHeaders?: Record<string, string> | undefined
@@ -210,7 +212,11 @@ export const classifyHttpFailure = (input: {
     retryAfterMs: retryAfter,
     rateLimit,
     http: new HttpContext({
-      request: new HttpRequestDetails({ method: "POST", url: input.url, headers: {} }),
+      request: new HttpRequestDetails({
+        method: input.method ?? "POST",
+        url: input.url,
+        headers: input.requestHeaders ?? {},
+      }),
       response:
         input.status === undefined
           ? undefined

packages/ai/src/route/transport/websocket.ts

@@ -2,6 +2,7 @@ import { Cause, Effect, Queue, Stream } from "effect"
 import { Headers } from "effect/unstable/http"
 import { Socket } from "effect/unstable/socket"
 import { AIError, TransportReason, type TransportOperation } from "../../schema/index.js"
+import { classifyHttpFailure } from "../executor.js"
 import * as HttpTransport from "./http.js"
 import type { Transport } from "./index.js"
 import type {
@@ -31,6 +32,18 @@ type WebSocketConstructorWithHeaders = (
   options?: { readonly headers?: Headers.Headers },
 ) => globalThis.WebSocket
 
+interface UpgradeResponse extends AsyncIterable<unknown> {
+  readonly statusCode?: number
+  readonly headers: Readonly<Record<string, string | ReadonlyArray<string> | undefined>>
+}
+
+type UpgradeListener = (request: unknown, response: UpgradeResponse) => void
+
+interface UpgradeAwareWebSocket extends globalThis.WebSocket {
+  readonly once?: (event: "unexpected-response", listener: UpgradeListener) => void
+  readonly off?: (event: "unexpected-response", listener: UpgradeListener) => void
+}
+
 const MAX_FRAME_BYTES = 16 * 1024 * 1024
 const transportError = (
   method: string,
@@ -91,6 +104,59 @@ const binaryMessage = (data: unknown) => {
   return undefined
 }
 
+const upgradeHeaders = (headers: UpgradeResponse["headers"]) =>
+  Object.fromEntries(
+    Object.entries(headers).flatMap(([name, value]) => {
+      if (value === undefined) return []
+      return [[name, Array.isArray(value) ? value.join(", ") : String(value)]]
+    }),
+  )
+
+const upgradeBody = (input: WebSocketRequest, response: UpgradeResponse) =>
+  Stream.fromAsyncIterable(response, (error) =>
+    transportError("open", error instanceof Error ? error.message : "Failed to read WebSocket upgrade response", {
+      url: input.url,
+      operation: "read",
+      phase: "connect",
+      delivery: "not-sent",
+    }),
+  ).pipe(
+    Stream.mapEffect((chunk) => {
+      if (typeof chunk === "string") return Effect.succeed(new TextEncoder().encode(chunk))
+      const binary = binaryMessage(chunk)
+      if (binary) return Effect.succeed(binary)
+      return Effect.fail(
+        transportError("open", "Unsupported WebSocket upgrade response body", {
+          url: input.url,
+          operation: "read",
+          phase: "connect",
+          delivery: "not-sent",
+        }),
+      )
+    }),
+    Stream.decodeText(),
+    Stream.runFold(
+      () => "",
+      (body, chunk) => body + chunk,
+    ),
+  )
+
+const rejectedUpgrade = (input: WebSocketRequest, response: UpgradeResponse, body: string | undefined) =>
+  new AIError({
+    module: "WebSocketConnector",
+    method: "open",
+    reason: classifyHttpFailure({
+      message: `WebSocket upgrade rejected with HTTP ${response.statusCode ?? "unknown"}`,
+      method: "GET",
+      url: input.url,
+      requestHeaders: { ...input.headers },
+      status: response.statusCode,
+      code: "UnexpectedServerResponse",
+      responseHeaders: upgradeHeaders(response.headers),
+      responseBody: body,
+    }),
+  })
+
 const waitOpen = (ws: globalThis.WebSocket, input: WebSocketRequest) => {
   if (ws.readyState === globalThis.WebSocket.OPEN) return Effect.void
   if (ws.readyState === globalThis.WebSocket.CLOSING || ws.readyState === globalThis.WebSocket.CLOSED) {
@@ -105,10 +171,13 @@ const waitOpen = (ws: globalThis.WebSocket, input: WebSocketRequest) => {
     )
   }
   return Effect.callback<void, AIError>((resume, signal) => {
+    const upgrade: UpgradeAwareWebSocket = ws
+    const unexpectedResponse = upgrade.once && upgrade.off ? upgrade : undefined
     const cleanup = () => {
       ws.removeEventListener("open", onOpen)
       ws.removeEventListener("error", onError)
       ws.removeEventListener("close", onClose)
+      unexpectedResponse?.off?.("unexpected-response", onUnexpectedResponse)
       signal.removeEventListener("abort", onAbort)
     }
     const onAbort = () => {
@@ -147,9 +216,25 @@ const waitOpen = (ws: globalThis.WebSocket, input: WebSocketRequest) => {
         ),
       )
     }
+    const onUnexpectedResponse: UpgradeListener = (_request, response) => {
+      cleanup()
+      resume(
+        upgradeBody(input, response).pipe(
+          Effect.catch(() => Effect.succeed(undefined)),
+          Effect.flatMap((body) => Effect.fail(rejectedUpgrade(input, response, body))),
+          Effect.ensuring(
+            Effect.sync(() => {
+              ws.addEventListener("error", () => {}, { once: true })
+              if (ws.readyState === globalThis.WebSocket.CONNECTING) ws.close()
+            }),
+          ),
+        ),
+      )
+    }
     ws.addEventListener("open", onOpen, { once: true })
     ws.addEventListener("error", onError, { once: true })
     ws.addEventListener("close", onClose, { once: true })
+    unexpectedResponse?.once?.("unexpected-response", onUnexpectedResponse)
     signal.addEventListener("abort", onAbort, { once: true })
   })
 }

packages/ai/test/executor.test.ts

@@ -1,6 +1,7 @@
 import { describe, expect } from "bun:test"
 import { Deferred, Effect, Fiber, Layer, Ref, Stream } from "effect"
 import { Headers, HttpClient, HttpClientError, HttpClientRequest, HttpClientResponse } from "effect/unstable/http"
+import { Socket } from "effect/unstable/socket"
 import { LLM, AIError } from "../src/index.js"
 import { LLMClient, RequestExecutor, WebSocketTransport, type WebSocketChannelExecutor } from "../src/route.js"
 import * as OpenAIChat from "../src/protocols/openai-chat.js"
@@ -66,6 +67,53 @@ const expectAIError = (error: unknown) => {
 }
 
 const errorHttp = (error: AIError) => ("http" in error.reason ? error.reason.http : undefined)
+
+interface RejectedUpgradeResponse extends AsyncIterable<Uint8Array> {
+  readonly statusCode: number
+  readonly headers: Readonly<Record<string, string | undefined>>
+}
+
+type RejectedUpgradeListener = (request: unknown, response: RejectedUpgradeResponse) => void
+
+class RejectedUpgradeSocket extends EventTarget {
+  readyState = globalThis.WebSocket.CONNECTING
+  listener?: RejectedUpgradeListener
+
+  constructor(
+    readonly response: { status: number; headers: Readonly<Record<string, string | undefined>>; body: string },
+  ) {
+    super()
+  }
+
+  once(_event: "unexpected-response", listener: RejectedUpgradeListener) {
+    this.listener = listener
+    const body = this.response.body
+    queueMicrotask(() =>
+      listener(
+        {},
+        {
+          statusCode: this.response.status,
+          headers: this.response.headers,
+          async *[Symbol.asyncIterator]() {
+            const split = Math.floor(body.length / 2)
+            yield new TextEncoder().encode(body.slice(0, split))
+            yield new TextEncoder().encode(body.slice(split))
+          },
+        },
+      ),
+    )
+  }
+
+  off(_event: "unexpected-response", listener: RejectedUpgradeListener) {
+    if (this.listener === listener) this.listener = undefined
+  }
+
+  close() {
+    this.readyState = globalThis.WebSocket.CLOSED
+  }
+
+  send() {}
+}
 const largeProviderMessage = `Upstream request failed: ${"validation failed; ".repeat(1_000)}`
 
 describe("RequestExecutor", () => {
@@ -622,6 +670,78 @@ describe("WebSocket channel execution", () => {
     }),
   )
 
+  it.effect("preserves rejected upgrade diagnostics", () =>
+    Effect.gen(function* () {
+      const cases = [
+        {
+          status: 401,
+          headers: { "x-request-id": "req_401" },
+          body: '{"error":{"message":"invalid key"}}',
+          reason: { _tag: "Authentication", kind: "invalid" },
+        },
+        {
+          status: 403,
+          headers: { "x-request-id": "req_403" },
+          body: '{"error":{"message":"forbidden"}}',
+          reason: { _tag: "Authentication", kind: "insufficient-permissions" },
+        },
+        {
+          status: 429,
+          headers: {
+            "retry-after": "2",
+            "x-request-id": "req_429",
+            "x-ratelimit-limit-requests": "500",
+          },
+          body: '{"error":{"message":"rate limited"}}',
+          reason: { _tag: "RateLimit", retryAfterMs: 2_000 },
+        },
+        {
+          status: 503,
+          headers: { "retry-after-ms": "250", "x-request-id": "req_503" },
+          body: '{"error":{"message":"overloaded"}}',
+          reason: { _tag: "ProviderInternal", status: 503, retryAfterMs: 250 },
+        },
+      ] as const
+
+      yield* Effect.forEach(
+        cases,
+        (item) =>
+          Effect.gen(function* () {
+            const socket = new RejectedUpgradeSocket(item)
+            const error = yield* WebSocketTransport.open({
+              url: "wss://api.openai.test/v1/responses",
+              headers: Headers.fromInput({ authorization: "Bearer secret", "x-client": "visible" }),
+            }).pipe(
+              Effect.provideService(Socket.WebSocketConstructor, () => {
+                // oxlint-disable-next-line typescript-eslint/no-unsafe-type-assertion -- fixture implements the socket surface used by the connector.
+                return socket as unknown as globalThis.WebSocket
+              }),
+              Effect.flip,
+            )
+
+            expectAIError(error)
+            expect(error.reason).toMatchObject(item.reason)
+            expect(errorHttp(error)).toMatchObject({
+              request: {
+                method: "GET",
+                url: "wss://api.openai.test/v1/responses",
+                headers: { authorization: "Bearer secret", "x-client": "visible" },
+              },
+              response: { status: item.status, headers: item.headers },
+              body: item.body,
+              requestId: `req_${item.status}`,
+              ...(item.status === 429
+                ? { rateLimit: { retryAfterMs: 2_000, limit: { requests: "500" } } }
+                : item.status === 503
+                  ? { rateLimit: { retryAfterMs: 250 } }
+                  : {}),
+            })
+          }),
+        { discard: true },
+      )
+    }),
+  )
+
   it.effect("uses HTTP when no per-call WebSocket executor is provided", () =>
     Effect.gen(function* () {
       const response = yield* LLMClient.generate(request).pipe(Effect.provide(fixedResponse(sseEvents(...frames))))

packages/core/test/session-model-transport-live.test.ts

@@ -34,13 +34,16 @@ const exchange = (server: WebSocketServerFixture, id: string): WebSocketChannelE
 
 const withServer = <A>(
   options: WebSocketServerOptions,
-  effect: (server: WebSocketServerFixture) => Effect.Effect<A, unknown, SessionModelTransport.Service>,
+  effect: (
+    server: WebSocketServerFixture,
+    constructor: Socket.WebSocketConstructor["Service"],
+  ) => Effect.Effect<A, unknown, SessionModelTransport.Service>,
 ) =>
   Effect.runPromise(
     Effect.gen(function* () {
       const constructor = yield* Socket.WebSocketConstructor
       const server = yield* makeWebSocketServer(options)
-      return yield* effect(server).pipe(
+      return yield* effect(server, constructor).pipe(
         Effect.provide(
           SessionModelTransport.makeLayer({
             open: (input) =>
@@ -293,8 +296,27 @@ describe("SessionModelTransport local WebSocket server", () => {
     )
   })
 
-  // Effect's browser-compatible constructor does not expose upgrade response bodies or headers.
-  // The real 426 fixture therefore pins the observable contract: a not-sent connect failure and one HTTP fallback.
+  test("preserves a real rejected upgrade response", async () => {
+    await withServer({ upgrade: () => false }, (server, constructor) =>
+      Effect.gen(function* () {
+        const error = yield* WebSocketTransport.open({ url: server.url, headers: Headers.empty }).pipe(
+          Effect.provideService(Socket.WebSocketConstructor, constructor),
+          Effect.flip,
+        )
+
+        expect(error.reason).toMatchObject({
+          _tag: "UnknownProvider",
+          status: 426,
+          http: {
+            request: { method: "GET", url: server.url },
+            response: { status: 426, headers: { "x-upgrade-rejected": "true" } },
+            body: "WebSocket upgrade required",
+          },
+        })
+      }),
+    )
+  })
+
   test("falls back once after a real rejected upgrade", async () => {
     let fallbacks = 0
     await withServer({ upgrade: () => false }, (server) =>