diff --git a/packages/ai/src/route/transport/websocket.ts b/packages/ai/src/route/transport/websocket.ts index bce2d28a6bc4..dfb2a0a30468 100644 --- a/packages/ai/src/route/transport/websocket.ts +++ b/packages/ai/src/route/transport/websocket.ts @@ -39,6 +39,12 @@ type WebSocketConstructorWithHeaders = ( ) => globalThis.WebSocket const MAX_FRAME_BYTES = 16 * 1024 * 1024 +const MAX_INBOUND_BYTES = 64 * 1024 * 1024 +const MAX_INBOUND_FRAMES = 4_096 +// Charge UTF-16 text or the binary backing buffer, plus an entry allowance. +// The separate frame ceiling also bounds object overhead for empty/tiny frames. +const retainedBytes = (message: string | Uint8Array) => + (typeof message === "string" ? message.length * 2 : message.buffer.byteLength) + 64 const transportError = ( message: string, input: { @@ -215,12 +221,17 @@ export const fromWebSocket = ( ): Effect.Effect => Effect.gen(function* () { yield* waitOpen(ws, input) - const messages = yield* Queue.bounded>(128) + // The native message callback cannot suspend the sender. Bound both retained payload + // and per-frame overhead, allowing small delta bursts without an unbounded queue. + const messages = yield* Queue.bounded>(MAX_INBOUND_FRAMES) + let bufferedBytes = 0 + let stopped = false const oversized = (message: string | Uint8Array) => typeof message === "string" ? new Blob([message]).size > MAX_FRAME_BYTES : message.byteLength > MAX_FRAME_BYTES const rejectOversized = (message: string | Uint8Array) => { if (!oversized(message)) return false + stopped = true Queue.failCauseUnsafe( messages, Cause.fail( @@ -237,8 +248,14 @@ export const fromWebSocket = ( return true } const offer = (message: string | Uint8Array) => { + if (stopped) return if (rejectOversized(message)) return - if (Queue.offerUnsafe(messages, message)) return + const bytes = retainedBytes(message) + if (bufferedBytes + bytes <= MAX_INBOUND_BYTES && Queue.offerUnsafe(messages, message)) { + bufferedBytes += bytes + return + } + stopped = true Queue.failCauseUnsafe( messages, Cause.fail( @@ -251,12 +268,15 @@ export const fromWebSocket = ( }), ), ) + if (ws.readyState === globalThis.WebSocket.OPEN) ws.close(1009, "Inbound queue overflow") } const onMessage = (event: MessageEvent) => { + if (stopped) return if (typeof event.data === "string") return offer(event.data) const binary = binaryMessage(event.data) if (binary) return offer(binary) + stopped = true Queue.failCauseUnsafe( messages, Cause.fail( @@ -271,6 +291,7 @@ export const fromWebSocket = ( ) } const onError = (event: Event) => { + stopped = true Queue.failCauseUnsafe( messages, Cause.fail( @@ -285,6 +306,7 @@ export const fromWebSocket = ( ) } const onClose = (event: CloseEvent) => { + stopped = true Queue.failCauseUnsafe( messages, Cause.fail( @@ -300,6 +322,8 @@ export const fromWebSocket = ( ) } const cleanup = Effect.sync(() => { + stopped = true + bufferedBytes = 0 ws.removeEventListener("message", onMessage) ws.removeEventListener("error", onError) ws.removeEventListener("close", onClose) @@ -333,7 +357,15 @@ export const fromWebSocket = ( }), }) }), - messages: Stream.fromQueue(messages), + // Taking a whole batch would release its budget while downstream still retains it. + messages: Stream.fromEffectRepeat( + Queue.take(messages).pipe( + Effect.map((message) => { + bufferedBytes -= retainedBytes(message) + return message + }), + ), + ), close: cleanup.pipe( Effect.andThen( Effect.sync(() => { diff --git a/packages/ai/test/websocket-buffer.test.ts b/packages/ai/test/websocket-buffer.test.ts new file mode 100644 index 000000000000..b466f3c2a4c4 --- /dev/null +++ b/packages/ai/test/websocket-buffer.test.ts @@ -0,0 +1,208 @@ +import { describe, expect } from "bun:test" +import { Deferred, Effect, Fiber, Stream } from "effect" +import { Headers } from "effect/unstable/http" +import { WebSocketTransport } from "../src/route.js" +import { it } from "./lib/effect.js" + +class TestSocket extends EventTarget { + readyState: number = WebSocket.OPEN + closes: number[] = [] + listeners = new Set() + send() {} + close(code: number) { + this.closes.push(code) + this.readyState = WebSocket.CLOSING + } + override addEventListener(...args: Parameters) { + this.listeners.add(args[0]) + super.addEventListener(...args) + } + override removeEventListener(...args: Parameters) { + this.listeners.delete(args[0]) + super.removeEventListener(...args) + } + receive(data: string | Uint8Array) { + this.dispatchEvent(new MessageEvent("message", { data })) + } +} + +const open = (socket: TestSocket) => + Effect.acquireRelease( + WebSocketTransport.fromWebSocket( + // oxlint-disable-next-line typescript-eslint/no-unsafe-type-assertion -- deterministic native-event adapter fixture. + socket as unknown as WebSocket, + { url: "wss://provider.test/responses", headers: Headers.empty }, + ), + (connection) => connection.close, + ) + +describe("WebSocket inbound buffering", () => { + it.live("preserves a burst while the consumer is paused", () => + Effect.gen(function* () { + const frames = Array.from({ length: 1_024 }, (_, index) => `frame:${index}`) + const arrived = Deferred.makeUnsafe() + const paused = Deferred.makeUnsafe() + const resume = Deferred.makeUnsafe() + const server = yield* Effect.acquireRelease( + Effect.sync(() => + Bun.serve({ + hostname: "127.0.0.1", + port: 0, + fetch(request, server) { + if (server.upgrade(request)) return + return new Response("WebSocket required", { status: 400 }) + }, + websocket: { + message(socket, message) { + if (message === "start") return void socket.send("start") + frames.forEach((frame) => socket.send(frame)) + }, + }, + }), + ), + (server) => Effect.promise(() => server.stop(true)), + ) + const socket = new WebSocket(`ws://127.0.0.1:${server.port}`) + const connection = yield* Effect.acquireRelease( + WebSocketTransport.fromWebSocket(socket, { url: socket.url, headers: Headers.empty }), + (connection) => connection.close, + ) + socket.addEventListener("message", (event) => { + if (event.data === frames.at(-1)) Effect.runSync(Deferred.succeed(arrived, undefined)) + }) + const consumer = yield* connection.messages.pipe( + Stream.mapEffect((frame) => + frame === "start" + ? Deferred.succeed(paused, undefined).pipe(Effect.andThen(Deferred.await(resume)), Effect.as(frame)) + : Effect.succeed(frame), + ), + Stream.take(frames.length + 1), + Stream.runCollect, + Effect.forkChild({ startImmediately: true }), + ) + yield* connection.sendText("start") + yield* Deferred.await(paused) + yield* connection.sendText("burst") + yield* Deferred.await(arrived) + yield* Deferred.succeed(resume, undefined) + expect(Array.from(yield* Fiber.join(consumer))).toEqual(["start", ...frames]) + yield* connection.sendText("burst") + expect(Array.from(yield* connection.messages.pipe(Stream.take(frames.length), Stream.runCollect))).toEqual(frames) + }), + ) + + it.effect("bounds retained text bytes before the frame count limit and preserves the accepted prefix", () => + Effect.gen(function* () { + const socket = new TestSocket() + const connection = yield* open(socket) + const frame = "x".repeat(1024 * 1024) + Array.from({ length: 32 }, () => socket.receive(frame)) + socket.receive("late") + let received = 0 + const error = yield* connection.messages.pipe( + Stream.tap(() => Effect.sync(() => received++)), + Stream.runDrain, + Effect.flip, + ) + expect(received).toBe(31) + expect(error.reason).toMatchObject({ _tag: "Transport", code: "queue-overflow", body: frame }) + expect(socket.closes).toEqual([1009]) + }), + ) + + it.effect("bounds empty-frame overhead and ignores further frames after overflow", () => + Effect.gen(function* () { + const socket = new TestSocket() + const connection = yield* open(socket) + Array.from({ length: 4_097 }, () => socket.receive("")) + socket.receive("late") + let received = 0 + const error = yield* connection.messages.pipe( + Stream.tap(() => Effect.sync(() => received++)), + Stream.runDrain, + Effect.flip, + ) + expect(received).toBe(4_096) + expect(error.reason).toMatchObject({ _tag: "Transport", code: "queue-overflow", body: "" }) + expect(socket.closes).toEqual([1009]) + yield* connection.close + expect(socket.listeners.size).toBe(0) + expect(socket.closes).toEqual([1009]) + }), + ) + + it.effect("reclaims byte capacity as frames are consumed", () => + Effect.gen(function* () { + const socket = new TestSocket() + const connection = yield* open(socket) + const frame = "x".repeat(1024 * 1024) + for (const _ of [0, 1]) { + Array.from({ length: 31 }, () => socket.receive(frame)) + const received = yield* connection.messages.pipe(Stream.take(31), Stream.runCollect) + expect(received).toHaveLength(31) + expect(received.every((message) => message === frame)).toBe(true) + } + expect(socket.closes).toEqual([]) + }), + ) + + it.effect("preserves mixed text and binary frame boundaries, including a maximum-sized frame", () => + Effect.gen(function* () { + const socket = new TestSocket() + const connection = yield* open(socket) + const frames = ["first", new Uint8Array([1, 2, 3]), "x".repeat(16 * 1024 * 1024), new Uint8Array([4])] + frames.forEach((frame) => socket.receive(frame)) + expect(Array.from(yield* connection.messages.pipe(Stream.take(frames.length), Stream.runCollect))).toEqual(frames) + expect(socket.closes).toEqual([]) + }), + ) + + it.effect("charges a binary view for the backing buffer it retains", () => + Effect.gen(function* () { + const socket = new TestSocket() + const connection = yield* open(socket) + socket.receive(new Uint8Array(new ArrayBuffer(64 * 1024 * 1024), 0, 1)) + const error = yield* connection.messages.pipe(Stream.runDrain, Effect.flip) + expect(error.reason).toMatchObject({ _tag: "Transport", code: "queue-overflow" }) + expect(socket.closes).toEqual([1009]) + }), + ) + + it.effect("retains the UTF-8 frame-size guard for text and the binary size guard", () => + Effect.gen(function* () { + for (const frame of [ + "\u0800".repeat(Math.floor((16 * 1024 * 1024) / 3) + 1), + new Uint8Array(16 * 1024 * 1024 + 1), + ]) { + const socket = new TestSocket() + const connection = yield* open(socket) + socket.receive(frame) + const error = yield* connection.messages.pipe(Stream.runDrain, Effect.flip) + expect(error.reason).toMatchObject({ _tag: "Transport", code: "message-too-large" }) + expect(socket.closes).toEqual([1009]) + } + }), + ) + + it.effect("cleans up a cancelled consumer with a buffered burst", () => + Effect.gen(function* () { + const socket = new TestSocket() + const paused = Deferred.makeUnsafe() + const resume = Deferred.makeUnsafe() + const consumer = yield* Effect.gen(function* () { + const connection = yield* open(socket) + yield* connection.messages.pipe( + Stream.runForEach(() => Deferred.succeed(paused, undefined).pipe(Effect.andThen(Deferred.await(resume)))), + ) + }).pipe(Effect.scoped, Effect.forkChild({ startImmediately: true })) + socket.receive("first") + yield* Deferred.await(paused) + Array.from({ length: 1_024 }, (_, index) => socket.receive(`frame:${index}`)) + yield* Fiber.interrupt(consumer) + expect(socket.listeners.size).toBe(0) + expect(socket.closes).toEqual([1000]) + socket.receive("late") + expect(socket.closes).toEqual([1000]) + }), + ) +}) diff --git a/packages/core/src/session/model-transport.ts b/packages/core/src/session/model-transport.ts index 590cabf3c157..3ef0a01f0e5b 100644 --- a/packages/core/src/session/model-transport.ts +++ b/packages/core/src/session/model-transport.ts @@ -20,7 +20,7 @@ import { SessionSchema } from "./schema.js" import { webSocketConstructor } from "../effect/app-node-platform.js" const ROTATE_AFTER_MS = 55 * 60 * 1000 -const INBOUND_CAPACITY = 128 +const INBOUND_CAPACITY = 1 const CONNECT_TIMEOUT = "10 seconds" const IDLE_TIMEOUT = "5 minutes" const events = Metric.counter("opencode_session_websocket_events_total", { @@ -235,11 +235,12 @@ export const makeLayer = (connector: WebSocketConnector) => code: "message", phase: "receive", }) - if (Queue.offerUnsafe(active.queue, message)) return undefined - return yield* transportError("Session WebSocket inbound queue overflow", { + // This reader can suspend; the socket adapter owns bounded burst buffering. + if (yield* Queue.offer(active.queue, message)) return undefined + return yield* transportError("WebSocket data arrived after the exchange stopped", { url: exchange.connect.url, operation: "read", - code: "queue-overflow", + code: "idle-data", phase: "receive", delivery: "accepted", }) @@ -401,7 +402,8 @@ export const makeLayer = (connector: WebSocketConnector) => let terminal: ChannelObservation | undefined const token = {} - const frames = Stream.fromQueue(active.queue).pipe( + // Keep frames after a terminal event in the queue so the finalizer can reject them. + const frames = Stream.fromEffectRepeat(Queue.take(active.queue)).pipe( Stream.timeoutOrElse({ duration: IDLE_TIMEOUT, orElse: () => diff --git a/packages/core/test/session-model-transport.test.ts b/packages/core/test/session-model-transport.test.ts index 6933d49957f8..f0d0534910cc 100644 --- a/packages/core/test/session-model-transport.test.ts +++ b/packages/core/test/session-model-transport.test.ts @@ -1,5 +1,6 @@ import { describe, expect, test } from "bun:test" import { AIError, HttpContext, InvalidRequestError, TransportError } from "@opencode/ai" +import { WebSocketTransport } from "@opencode/ai/route" import type { ChannelObservation, WebSocketChannelExchange, @@ -915,24 +916,95 @@ describe("SessionModelTransport", () => { ) }) - test("poisons instead of dropping data when the inbound queue overflows", async () => { + test("drains a real socket burst through a paused exchange and reuses the connection", async () => { + const frames = Array.from({ length: 1_024 }, (_, index) => `frame:${index}`) + const arrived = Deferred.makeUnsafe() + const paused = Deferred.makeUnsafe() + const resume = Deferred.makeUnsafe() + const sockets: WebSocket[] = [] + const server = Bun.serve({ + hostname: "127.0.0.1", + port: 0, + fetch(request, server) { + if (server.upgrade(request)) return + return new Response("WebSocket required", { status: 400 }) + }, + websocket: { + message(socket, message) { + if (message === "first") return void socket.send("paused") + if (message === "burst") frames.forEach((frame) => socket.send(frame)) + socket.send(`completed:${message}`) + }, + }, + }) + const connector: WebSocketConnector = { + open: () => + Effect.suspend(() => { + const socket = new WebSocket(`ws://127.0.0.1:${server.port}`) + sockets.push(socket) + socket.addEventListener("message", (event) => { + if (event.data === "completed:burst") Effect.runSync(Deferred.succeed(arrived, undefined)) + }) + return WebSocketTransport.fromWebSocket(socket, { url: socket.url, headers: Headers.empty }) + }), + } + + await run( + connector, + Effect.gen(function* () { + const transport = yield* SessionModelTransport.Service + const executor = transport.bind(session) + const item = exchange("first") + const consumer = yield* collectComplete(executor, { + ...item, + driver: { + create: item.driver.create, + observe: (_create, frame) => + frame === "paused" + ? Deferred.succeed(paused, undefined).pipe( + Effect.andThen(Deferred.await(resume)), + Effect.as({ type: "frame" as const, frame }), + ) + : Effect.succeed({ type: frame.startsWith("completed:") ? "completed" : "frame", frame }), + }, + }).pipe(Effect.forkChild({ startImmediately: true })) + yield* Deferred.await(paused) + sockets[0].send("burst") + yield* Deferred.await(arrived) + yield* Deferred.succeed(resume, undefined) + expect(yield* Fiber.join(consumer)).toEqual(["paused", ...frames, "completed:burst"]) + expect(yield* collect(executor, exchange("second"))).toEqual(["completed:second"]) + expect(sockets).toHaveLength(1) + }).pipe(Effect.ensuring(Effect.promise(() => server.stop(true)))), + ) + }) + + test("poisons when the socket buffer overflows", async () => { const messages = queue() - const poisoned = Deferred.makeUnsafe() let closed = 0 const connector: WebSocketConnector = { open: () => Effect.succeed({ sendText: () => - // Hold consumption at the send boundary until the reader fills and poisons the inbound queue. Effect.sync(() => { - for (let index = 0; index <= 129; index++) Queue.offerUnsafe(messages, `frame:${index}`) - }).pipe(Effect.andThen(Deferred.await(poisoned))), + Queue.offerUnsafe(messages, "first") + Queue.failCauseUnsafe( + messages, + Cause.fail( + new AIError({ + reason: new TransportError({ + message: "WebSocket inbound queue overflow", + transport: "websocket", + operation: "read", + code: "queue-overflow", + phase: "receive", + }), + }), + ), + ) + }), messages: Stream.fromQueue(messages).pipe(Stream.tap(() => Effect.yieldNow)), - close: Effect.sync(() => closed++).pipe( - Effect.andThen(Deferred.succeed(poisoned, undefined)), - Effect.andThen(Queue.shutdown(messages)), - Effect.asVoid, - ), + close: Effect.sync(() => closed++).pipe(Effect.andThen(Queue.shutdown(messages)), Effect.asVoid), }), } @@ -959,6 +1031,90 @@ describe("SessionModelTransport", () => { ) }) + test("cancels a backpressured reader and reconnects for the next exchange", async () => { + const paused = Deferred.makeUnsafe() + const pending = Deferred.makeUnsafe() + const resume = Deferred.makeUnsafe() + const connections: Array<{ closed: number }> = [] + const connector: WebSocketConnector = { + open: () => + Effect.sync(() => { + const messages = queue() + const record = { closed: 0 } + connections.push(record) + return { + sendText: (message) => + Effect.sync(() => { + Queue.offerAllUnsafe(messages, message === "first" ? ["first", "second", "third"] : ["completed"]) + }), + messages: Stream.fromQueue(messages).pipe( + Stream.tap((message) => (message === "third" ? Deferred.succeed(pending, undefined) : Effect.void)), + ), + close: Effect.sync(() => record.closed++).pipe(Effect.andThen(Queue.shutdown(messages)), Effect.asVoid), + } + }), + } + + await run( + connector, + Effect.gen(function* () { + const transport = yield* SessionModelTransport.Service + const executor = transport.bind(session) + const item = exchange("first") + const consumer = yield* collect(executor, { + ...item, + driver: { + create: item.driver.create, + observe: (_create, frame) => + Deferred.succeed(paused, undefined).pipe( + Effect.andThen(Deferred.await(resume)), + Effect.as({ type: "frame" as const, frame }), + ), + }, + }).pipe(Effect.forkChild({ startImmediately: true })) + yield* Deferred.await(paused) + yield* Deferred.await(pending) + yield* Effect.yieldNow + yield* Fiber.interrupt(consumer) + expect(connections[0].closed).toBe(1) + expect(yield* collect(executor, exchange("second"))).toEqual(["completed"]) + expect(connections).toHaveLength(2) + }), + ) + }) + + test("poisons a terminal event followed by data in the same burst", async () => { + const connections: Array<{ closed: number }> = [] + const connector: WebSocketConnector = { + open: () => + Effect.sync(() => { + const messages = queue() + const record = { closed: 0 } + connections.push(record) + return { + sendText: (message) => + Effect.sync(() => { + Queue.offerAllUnsafe(messages, message === "first" ? ["completed", "late"] : ["completed"]) + }).pipe(Effect.andThen(Effect.yieldNow)), + messages: Stream.fromQueue(messages), + close: Effect.sync(() => record.closed++).pipe(Effect.andThen(Queue.shutdown(messages)), Effect.asVoid), + } + }), + } + + await run( + connector, + Effect.gen(function* () { + const transport = yield* SessionModelTransport.Service + const executor = transport.bind(session) + expect(yield* collect(executor, exchange("first"))).toEqual(["completed"]) + expect(connections[0].closed).toBe(1) + expect(yield* collect(executor, exchange("second"))).toEqual(["completed"]) + expect(connections).toHaveLength(2) + }), + ) + }) + test("poisons unsupported binary frames after provider observation", async () => { const messages = queue() const closed = Deferred.makeUnsafe()