feat(core): manage session websocket channels

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

@@ -198,10 +198,24 @@ export const fromWebSocket = (
     yield* waitOpen(ws, input)
     const messages = yield* Queue.bounded<string | Uint8Array, AIError | Cause.Done<void>>(128)
 
+    const offer = (message: string | Uint8Array) => {
+      if (Queue.offerUnsafe(messages, message)) return
+      Queue.failCauseUnsafe(
+        messages,
+        Cause.fail(
+          transportError("message", "WebSocket inbound queue overflow", {
+            url: input.url,
+            kind: "queue-overflow",
+            phase: "receive",
+          }),
+        ),
+      )
+    }
+
     const onMessage = (event: MessageEvent) => {
-      if (typeof event.data === "string") return Queue.offerUnsafe(messages, event.data)
+      if (typeof event.data === "string") return offer(event.data)
       const binary = binaryMessage(event.data)
-      if (binary) return Queue.offerUnsafe(messages, binary)
+      if (binary) return offer(binary)
       Queue.failCauseUnsafe(
         messages,
         Cause.fail(
@@ -249,15 +263,26 @@ export const fromWebSocket = (
 
     return {
       sendText: (message) =>
-        Effect.try({
-          try: () => ws.send(message),
-          catch: (error) =>
-            transportError("sendText", error instanceof Error ? error.message : "Failed to send WebSocket message", {
-              url: input.url,
-              kind: "write",
-              phase: "send",
-              delivery: "not-sent",
-            }),
+        Effect.suspend(() => {
+          if (ws.readyState !== globalThis.WebSocket.OPEN)
+            return Effect.fail(
+              transportError("sendText", `WebSocket is not open (state ${ws.readyState})`, {
+                url: input.url,
+                kind: "write",
+                phase: "send",
+                delivery: "not-sent",
+              }),
+            )
+          return Effect.try({
+            try: () => ws.send(message),
+            catch: (error) =>
+              transportError("sendText", error instanceof Error ? error.message : "Failed to send WebSocket message", {
+                url: input.url,
+                kind: "write",
+                phase: "send",
+                delivery: "not-sent",
+              }),
+          })
         }),
       messages: Stream.fromQueue(messages),
       close: cleanup.pipe(

packages/ai/test/executor.test.ts

@@ -7,7 +7,7 @@ import * as OpenAIChat from "../src/protocols/openai-chat"
 import * as OpenAI from "../src/providers/openai"
 import { dynamicResponse, fixedResponse } from "./lib/http"
 import { deltaChunk } from "./lib/openai-chunks"
-import { sseRaw } from "./lib/sse"
+import { sseEvents, sseRaw } from "./lib/sse"
 import { it } from "./lib/effect"
 
 const request = HttpClientRequest.post("https://provider.test/v1/chat?api_key=secret&key=secret&debug=1").pipe(
@@ -463,16 +463,37 @@ describe("WebSocket channel execution", () => {
     }),
   )
 
-  it.effect("requires a per-call WebSocket executor", () =>
+  it.effect("rejects a closed socket before attempting to send", () =>
     Effect.gen(function* () {
-      const error = yield* LLMClient.generate(request).pipe(Effect.provide(fixedResponse("")), Effect.flip)
+      class ClosedBeforeSend extends EventTarget {
+        readyState = globalThis.WebSocket.OPEN
+        sends = 0
+        send() {
+          this.sends++
+        }
+        close() {}
+      }
+      const socket = new ClosedBeforeSend()
+      const connection = yield* WebSocketTransport.fromWebSocket(
+        // oxlint-disable-next-line typescript-eslint/no-unsafe-type-assertion
+        socket as unknown as globalThis.WebSocket,
+        { url: "wss://api.openai.test/v1/responses", headers: Headers.empty },
+      )
+      socket.readyState = globalThis.WebSocket.CLOSED
 
-      expect(error.reason).toMatchObject({
-        _tag: "Transport",
-        phase: "prepare",
-        delivery: "not-sent",
-      })
-      expect(error.message).toContain("StreamOptions.webSocket")
+      const error = yield* connection.sendText("create").pipe(Effect.flip)
+
+      expect(error.reason).toMatchObject({ _tag: "Transport", phase: "send", delivery: "not-sent" })
+      expect(socket.sends).toBe(0)
+      yield* connection.close
+    }),
+  )
+
+  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))))
+
+      expect(response.text).toBe("Hi")
     }),
   )
 

packages/core/src/effect/app-node-platform.ts

@@ -1,4 +1,6 @@
 import { LLMClient, RequestExecutor } from "@opencode-ai/ai/route"
+import { NodeSocket } from "@effect/platform-node"
+import { Socket } from "effect/unstable/socket"
 import { makeGlobalNode } from "@opencode-ai/util/effect/app-node"
 import { httpClient } from "@opencode-ai/util/effect/app-node-platform"
 
@@ -10,4 +12,10 @@ export const requestExecutor = makeGlobalNode({
 
 export const llmClient = makeGlobalNode({ service: LLMClient.Service, layer: LLMClient.layer, deps: [requestExecutor] })
 
+export const webSocketConstructor = makeGlobalNode({
+  service: Socket.WebSocketConstructor,
+  layer: NodeSocket.layerWebSocketConstructorWS,
+  deps: [],
+})
+
 export * as LayerNodePlatform from "./app-node-platform"

packages/core/src/location-services.ts

@@ -33,6 +33,7 @@ import { WebSearch } from "./websearch"
 import { ReferenceInstructions } from "./reference/instructions"
 import { SessionRunnerLLM } from "./session/runner/llm"
 import { SessionRunnerModel } from "./session/runner/model"
+import { SessionModelTransport } from "./session/model-transport"
 import { SessionCompaction } from "./session/compaction"
 import { SessionTitle } from "./session/title"
 import { Skill } from "./skill"
@@ -90,6 +91,7 @@ const locationServiceNodes = [
   McpTool.node,
   SessionInstructions.node,
   SessionRunnerModel.node,
+  SessionModelTransport.node,
   SessionCompaction.node,
   SessionTitle.node,
   Snapshot.node,

packages/core/src/session.ts

@@ -1,7 +1,7 @@
 export * as Session from "./session"
 export * from "./session/schema"
 
-import { Effect, Layer, Schema, Context, Stream, Scope } from "effect"
+import { Effect, Layer, Schema, Context, RcMap, Stream, Scope } from "effect"
 import { ListAnchor } from "@opencode-ai/schema/session"
 import { and, asc, desc, eq, gt, isNotNull, isNull, like, lt, ne, or, type SQL } from "drizzle-orm"
 import { Project } from "./project"
@@ -28,6 +28,7 @@ import { fromRow } from "./session/info"
 import { SessionRunner } from "./session/runner/index"
 import { SessionStore } from "./session/store"
 import { SessionExecution } from "./session/execution"
+import { SessionModelTransport } from "./session/model-transport"
 import { ForkEmptyError, MessageDecodeError, NotFoundError } from "./session/error"
 import { makeGlobalNode } from "@opencode-ai/util/effect/app-node"
 import { LocationServiceMap } from "./location-service-map"
@@ -323,6 +324,16 @@ const layer = Layer.effect(
     const scope = yield* Scope.Scope
     const activeShells = new Set<SessionSchema.ID>()
     const shellLocks = KeyedMutex.makeUnsafe<SessionSchema.ID>()
+    const closeTransport = Effect.fn("Session.closeTransport")(function* (session: SessionSchema.Info) {
+      const location = Location.Ref.make({
+        directory: session.location.directory,
+        workspaceID: session.location.workspaceID,
+      })
+      if (!(yield* RcMap.has(locations.rcMap, location))) return
+      yield* SessionModelTransport.Service.use((transport) => transport.close(session.id)).pipe(
+        Effect.provide(locations.get(location)),
+      )
+    })
     const decodeMessage = Schema.decodeUnknownEffect(SessionMessage.Info)
     const isDurableSessionEvent = Schema.is(SessionEvent.Durable)
     const persistProject = (project: Project.Resolved) => {
@@ -446,9 +457,10 @@ const layer = Layer.effect(
         return session
       }),
       remove: Effect.fn("Session.remove")(function* (sessionID) {
-        yield* result.get(sessionID)
+        const session = yield* result.get(sessionID)
         yield* execution.interrupt(sessionID)
         yield* execution.awaitIdle(sessionID)
+        yield* closeTransport(session)
         const children = yield* result.list({ parentID: sessionID })
         yield* Effect.forEach(children.data, (child) => result.remove(child.id), { concurrency: 1, discard: true })
         yield* bus.publish(SessionEvent.Deleted, { sessionID })
@@ -748,8 +760,9 @@ const layer = Layer.effect(
         yield* persistProject(project)
         if ((yield* execution.active).has(input.sessionID)) {
           yield* execution.interrupt(input.sessionID)
-          yield* execution.awaitIdle(input.sessionID)
         }
+        yield* execution.awaitIdle(input.sessionID)
+        yield* closeTransport(current)
         yield* bus.publish(SessionEvent.Moved, {
           sessionID: input.sessionID,
           location: Location.Ref.make({ directory, workspaceID: input.workspaceID }),

packages/core/src/session/model-request.ts

@@ -10,12 +10,14 @@ import { HttpClientRequest, HttpClientResponse } from "effect/unstable/http"
 import { makeLocationNode } from "@opencode-ai/util/effect/app-node"
 import { App } from "../app"
 import { Model } from "../model"
+import { Provider } from "../provider"
 import { Permission } from "../permission"
 import { PluginHooks } from "../plugin/hooks"
 import { QuestionTool } from "../tool/plugin/question"
 import { Tool } from "../tool"
 import { SessionContext } from "./context"
 import { SessionModelHeaders } from "./model-headers"
+import { SessionModelTransport } from "./model-transport"
 import { PromptCacheDiagnostics } from "./prompt-cache-diagnostics"
 import { MAX_STEPS_PROMPT } from "./runner/max-steps"
 import PROMPT_DEFAULT from "./runner/prompt/base.txt"
@@ -201,7 +203,12 @@ export const layer = Layer.effect(
   Service,
   Effect.gen(function* () {
     const hooks = yield* PluginHooks.Service
+    const transport = yield* SessionModelTransport.Service
     const app = yield* App.Metadata
+    const webSocket = yield* Config.boolean("OPENCODE_EXPERIMENTAL_OPENAI_RESPONSES_WEBSOCKET").pipe(
+      Config.withDefault(false),
+      Effect.orDie,
+    )
     const diagnostics = yield* Config.boolean("OPENCODE_PROMPT_CACHE_DIAGNOSTICS").pipe(
       Config.withDefault(false),
       Effect.orDie,
@@ -274,7 +281,16 @@ export const layer = Layer.effect(
           }),
       })
       const http = composeHttpMiddleware(middlewares)
-      const options: StreamOptions = http ? { http } : {}
+      const webSocketEligible = middlewares.length === 0
+      const options: StreamOptions = {
+        ...(http ? { http } : {}),
+        ...(webSocket &&
+        webSocketEligible &&
+        resolved.ref.providerID === Provider.ID.openai &&
+        model.route.id === "openai-responses"
+          ? { webSocket: transport.bind(session.id) }
+          : {}),
+      }
       if (promptCacheSnapshots) {
         const current = PromptCacheDiagnostics.snapshot(request)
         const comparison = PromptCacheDiagnostics.compare(promptCacheSnapshots.get(session.id), current)
@@ -306,7 +322,7 @@ export const layer = Layer.effect(
       return {
         request,
         options,
-        webSocketEligible: middlewares.length === 0,
+        webSocketEligible,
         executeTool,
         stepLimitReached,
       }
@@ -319,5 +335,5 @@ export const layer = Layer.effect(
 export const node = makeLocationNode({
   service: Service,
   layer,
-  deps: [PluginHooks.node, App.node],
+  deps: [PluginHooks.node, SessionModelTransport.node, App.node],
 })