diff --git a/packages/ai/src/route/transport/websocket.ts b/packages/ai/src/route/transport/websocket.ts index aaf27ea0ecad..3b41a72192eb 100644 --- a/packages/ai/src/route/transport/websocket.ts +++ b/packages/ai/src/route/transport/websocket.ts @@ -198,10 +198,24 @@ export const fromWebSocket = ( yield* waitOpen(ws, input) const messages = yield* Queue.bounded>(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( diff --git a/packages/ai/test/executor.test.ts b/packages/ai/test/executor.test.ts index 37dd5ad3fe6a..696e802afa11 100644 --- a/packages/ai/test/executor.test.ts +++ b/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") }), ) diff --git a/packages/core/src/effect/app-node-platform.ts b/packages/core/src/effect/app-node-platform.ts index 4ad725fcbb03..7783ed8e0251 100644 --- a/packages/core/src/effect/app-node-platform.ts +++ b/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" diff --git a/packages/core/src/location-services.ts b/packages/core/src/location-services.ts index 547cdd0e4557..24b5355ecf10 100644 --- a/packages/core/src/location-services.ts +++ b/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, diff --git a/packages/core/src/session.ts b/packages/core/src/session.ts index c7628ec799d2..157dc373c08e 100644 --- a/packages/core/src/session.ts +++ b/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() const shellLocks = KeyedMutex.makeUnsafe() + 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 }), diff --git a/packages/core/src/session/model-request.ts b/packages/core/src/session/model-request.ts index b2d0292e50c9..1dfd843e79b5 100644 --- a/packages/core/src/session/model-request.ts +++ b/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], }) diff --git a/packages/core/src/session/model-transport.ts b/packages/core/src/session/model-transport.ts new file mode 100644 index 000000000000..424985c5fc76 --- /dev/null +++ b/packages/core/src/session/model-transport.ts @@ -0,0 +1,381 @@ +export * as SessionModelTransport from "./model-transport" + +import { + WebSocketTransport, + type ChannelObservation, + type WebSocketChannelExchange, + type WebSocketChannelExecution, + type WebSocketChannelExecutor, + type WebSocketConnection, + type WebSocketConnector, +} from "@opencode-ai/ai/route" +import { AIError, TransportReason } from "@opencode-ai/ai" +import { Hash } from "@opencode-ai/util/hash" +import { Cause, Clock, Context, Effect, Fiber, Layer, Queue, Scope, Semaphore, Stream } from "effect" +import { Socket } from "effect/unstable/socket" +import { makeLocationNode } from "@opencode-ai/util/effect/app-node" +import { SessionSchema } from "./schema" +import { webSocketConstructor } from "../effect/app-node-platform" + +const ROTATE_AFTER_MS = 55 * 60 * 1000 +const INBOUND_CAPACITY = 128 + +type Delivery = "queued" | "connecting" | "ready" | "send-attempted" | "provider-observed" | "terminal" + +interface Active { + readonly queue: Queue.Queue + readonly lifecycle: { delivery: Delivery } +} + +interface Channel { + readonly affinity: string + readonly connection: WebSocketConnection + readonly openedAt: number + active?: Active + closing: boolean + poisoned: boolean + reader?: Fiber.Fiber +} + +interface State { + readonly lock: Semaphore.Semaphore + channel?: Channel +} + +export interface Interface { + readonly bind: (sessionID: SessionSchema.ID) => WebSocketChannelExecutor + readonly close: (sessionID: SessionSchema.ID) => Effect.Effect + readonly closeAll: Effect.Effect +} + +export class Service extends Context.Service()("@opencode/SessionModelTransport") {} + +const transportError = ( + method: string, + message: string, + input: { + readonly url?: string + readonly kind?: string + readonly phase?: TransportReason["phase"] + readonly delivery?: TransportReason["delivery"] + } = {}, +) => + new AIError({ + module: "SessionModelTransport", + method, + reason: new TransportReason({ message, ...input }), + }) + +const annotate = ( + error: AIError, + input: { readonly phase: TransportReason["phase"]; readonly delivery: TransportReason["delivery"] }, +) => { + if (error.reason._tag !== "Transport") return error + return new AIError({ + module: error.module, + method: error.method, + reason: new TransportReason({ + message: error.reason.message, + kind: error.reason.kind, + url: error.reason.url, + http: error.reason.http, + recovery: error.reason.recovery, + ...input, + }), + }) +} + +const affinity = (exchange: WebSocketChannelExchange) => + `${exchange.connect.url}:${Hash.sha256(JSON.stringify(Object.entries(exchange.connect.headers).sort(([a], [b]) => a.localeCompare(b))))}` + +const observationFrame = (observation: ChannelObservation) => { + if (observation.type === "frame" || observation.type === "completed" || observation.type === "incomplete") + return Effect.succeed(observation.frame) + return Effect.fail(observation.error) +} + +const observationTerminal = (observation: ChannelObservation) => observation.type !== "frame" + +export const makeLayer = (connector: WebSocketConnector) => + Layer.effect( + Service, + Effect.gen(function* () { + const scope = yield* Scope.Scope + const states = new Map() + const state = (sessionID: SessionSchema.ID) => { + const current = states.get(sessionID) + if (current) return current + const created = { lock: Semaphore.makeUnsafe(1) } + states.set(sessionID, created) + return created + } + + const closeChannel = Effect.fn("SessionModelTransport.closeChannel")(function* (owner: State, channel: Channel) { + if (owner.channel === channel) owner.channel = undefined + if (channel.closing) return + channel.closing = true + if (channel.reader) yield* Fiber.interrupt(channel.reader) + yield* channel.connection.close + if (channel.active) + Queue.failCauseUnsafe( + channel.active.queue, + Cause.fail( + transportError("close", "Session WebSocket closed", { + kind: "close", + phase: "close", + delivery: + channel.active.lifecycle.delivery === "queued" || + channel.active.lifecycle.delivery === "connecting" || + channel.active.lifecycle.delivery === "ready" + ? "not-sent" + : channel.active.lifecycle.delivery === "provider-observed" || + channel.active.lifecycle.delivery === "terminal" + ? "accepted" + : "ambiguous", + }), + ), + ) + }) + + const poison = Effect.fn("SessionModelTransport.poison")(function* ( + owner: State, + channel: Channel, + error: AIError, + ) { + channel.poisoned = true + if (owner.channel === channel) owner.channel = undefined + if (channel.closing) return + channel.closing = true + if (channel.active) Queue.failCauseUnsafe(channel.active.queue, Cause.fail(error)) + yield* channel.connection.close + }) + + const open = Effect.fn("SessionModelTransport.open")(function* ( + owner: State, + exchange: WebSocketChannelExchange, + key: string, + ) { + return yield* Effect.uninterruptibleMask((restore) => + Effect.gen(function* () { + const connection = yield* restore(connector.open(exchange.connect)) + const channel: Channel = { + affinity: key, + connection, + openedAt: yield* Clock.currentTimeMillis, + closing: false, + poisoned: false, + } + owner.channel = channel + channel.reader = yield* connection.messages.pipe( + Stream.runForEach((message) => + Effect.gen(function* () { + const active = channel.active + if (!active) + return yield* transportError("receive", "WebSocket data arrived without an active exchange", { + url: exchange.connect.url, + kind: "idle-data", + phase: "receive", + }) + active.lifecycle.delivery = "provider-observed" + if (typeof message !== "string") + return yield* transportError("receive", "Unsupported binary WebSocket frame", { + url: exchange.connect.url, + kind: "message", + phase: "receive", + }) + if (Queue.offerUnsafe(active.queue, message)) return undefined + return yield* transportError("receive", "Session WebSocket inbound queue overflow", { + url: exchange.connect.url, + kind: "queue-overflow", + phase: "receive", + delivery: "accepted", + }) + }), + ), + Effect.catch((error) => + channel.closing + ? Effect.void + : poison( + owner, + channel, + annotate(error, { + phase: + error.reason._tag === "Transport" && error.reason.phase === "close" ? "close" : "receive", + delivery: + channel.active?.lifecycle.delivery === "provider-observed" || + channel.active?.lifecycle.delivery === "terminal" || + (error.reason._tag === "Transport" && error.reason.kind === "queue-overflow") + ? "accepted" + : "ambiguous", + }), + ), + ), + Effect.forkIn(scope, { startImmediately: true }), + ) + yield* Effect.logDebug("session websocket connected", { + sessionTransport: "websocket", + phase: "connect", + }) + return channel + }), + ) + }) + + const fallback = (exchange: WebSocketChannelExchange): WebSocketChannelExecution => ({ + frames: exchange.fallback(), + complete: Effect.void, + }) + + const start = Effect.fn("SessionModelTransport.start")(function* ( + owner: State, + exchange: WebSocketChannelExchange, + lifecycle: { delivery: Delivery }, + ) { + const key = affinity(exchange) + const now = yield* Clock.currentTimeMillis + const current = owner.channel + const rotateAfterMs = exchange.connect.rotateAfterMs ?? ROTATE_AFTER_MS + const rotation = current + ? current.poisoned + ? "poisoned" + : current.affinity !== key + ? "affinity" + : now - current.openedAt >= rotateAfterMs + ? "age" + : undefined + : undefined + if (current && rotation) { + yield* Effect.logDebug("session websocket rotating", { + sessionTransport: "websocket", + phase: "connect", + reason: rotation, + }) + yield* closeChannel(owner, current) + } + + lifecycle.delivery = owner.channel ? "ready" : "connecting" + if (owner.channel) + yield* Effect.logDebug("session websocket reused", { + sessionTransport: "websocket", + phase: "connect", + }) + const channel = owner.channel + ? owner.channel + : yield* open(owner, exchange, key).pipe( + Effect.catch((error) => + Effect.logWarning("session websocket connect failed; using http", { + sessionTransport: "websocket", + phase: "connect", + delivery: "not-sent", + kind: error.reason._tag === "Transport" ? error.reason.kind : error.reason._tag, + }).pipe(Effect.andThen(Effect.succeed(undefined))), + ), + ) + if (!channel) return fallback(exchange) + lifecycle.delivery = "ready" + + const create = yield* exchange.driver.create(undefined).pipe( + Effect.tapError(() => closeChannel(owner, channel)), + Effect.onInterrupt(() => closeChannel(owner, channel)), + ) + const active: Active = { queue: yield* Queue.bounded(INBOUND_CAPACITY), lifecycle } + channel.active = active + lifecycle.delivery = "send-attempted" + const sent = yield* channel.connection.sendText(create.message).pipe( + Effect.onInterrupt(() => closeChannel(owner, channel)), + Effect.result, + ) + if (sent._tag === "Failure") { + const failure = sent.failure + const notSent = failure.reason._tag === "Transport" && failure.reason.delivery === "not-sent" + yield* closeChannel(owner, channel) + if (notSent) return fallback(exchange) + return yield* annotate(failure, { phase: "send", delivery: "ambiguous" }) + } + + let terminal = false + const frames = Stream.fromQueue(active.queue).pipe( + Stream.mapEffect((frame) => exchange.driver.observe(create, frame)), + Stream.tap((observation) => + Effect.sync(() => { + if (!observationTerminal(observation)) return + terminal = true + lifecycle.delivery = "terminal" + }), + ), + Stream.takeUntil(observationTerminal), + Stream.mapEffect(observationFrame), + Stream.ensuring( + Effect.gen(function* () { + if (channel.active === active) channel.active = undefined + const pending = yield* Queue.size(active.queue) + yield* Queue.shutdown(active.queue) + if (terminal && pending === 0) return + const error = terminal + ? transportError("receive", "WebSocket data arrived after the terminal event", { + url: exchange.connect.url, + kind: "idle-data", + phase: "receive", + delivery: "accepted", + }) + : transportError("execute", "Session WebSocket exchange did not reach a terminal event", { + url: exchange.connect.url, + kind: "incomplete", + phase: "receive", + delivery: lifecycle.delivery === "provider-observed" ? "accepted" : "ambiguous", + }) + yield* poison(owner, channel, error) + }), + ), + ) + return { frames, complete: Effect.void } + }) + + const bind = (sessionID: SessionSchema.ID): WebSocketChannelExecutor => ({ + execute: (exchange) => { + const owner = state(sessionID) + const lifecycle = { delivery: "queued" as Delivery } + return Effect.succeed({ + frames: Stream.unwrap( + Effect.acquireRelease(owner.lock.take(1), () => owner.lock.release(1), { interruptible: true }).pipe( + Effect.andThen(start(owner, exchange, lifecycle)), + Effect.map((execution) => execution.frames), + ), + ), + complete: Effect.void, + }) + }, + }) + + const close = Effect.fn("SessionModelTransport.close")(function* (sessionID: SessionSchema.ID) { + const owner = states.get(sessionID) + if (!owner) return + yield* owner.lock.withPermit( + Effect.gen(function* () { + if (owner.channel) yield* closeChannel(owner, owner.channel) + }), + ) + }) + const closeAll = Effect.forEach(states.values(), (owner) => + owner.lock.withPermit( + Effect.gen(function* () { + if (owner.channel) yield* closeChannel(owner, owner.channel) + }), + ), + ).pipe(Effect.asVoid) + + yield* Effect.addFinalizer(() => closeAll) + return Service.of({ bind, close, closeAll }) + }), + ) + +export const layer = Layer.unwrap( + Effect.map(Socket.WebSocketConstructor, (constructor) => + makeLayer({ + open: (input) => + WebSocketTransport.open(input).pipe(Effect.provideService(Socket.WebSocketConstructor, constructor)), + }), + ), +) + +export const node = makeLocationNode({ service: Service, layer, deps: [webSocketConstructor] }) diff --git a/packages/core/test/session-model-transport.test.ts b/packages/core/test/session-model-transport.test.ts new file mode 100644 index 000000000000..754482500b73 --- /dev/null +++ b/packages/core/test/session-model-transport.test.ts @@ -0,0 +1,471 @@ +import { describe, expect, test } from "bun:test" +import { AIError, TransportReason } from "@opencode-ai/ai" +import type { + ChannelObservation, + WebSocketChannelExchange, + WebSocketConnection, + WebSocketConnector, +} from "@opencode-ai/ai/route" +import { SessionModelTransport } from "@opencode-ai/core/session/model-transport" +import { Session } from "@opencode-ai/schema/session" +import { Deferred, Effect, Fiber, Queue, Stream } from "effect" +import { Headers } from "effect/unstable/http" + +const session = Session.ID.make("ses_transport") +const otherSession = Session.ID.make("ses_transport_other") +const queue = () => Effect.runSync(Queue.unbounded()) + +const error = (message: string, delivery?: TransportReason["delivery"]) => + new AIError({ + module: "test", + method: "websocket", + reason: new TransportReason({ message, phase: "send", delivery }), + }) + +const exchange = ( + id: string, + input: { + readonly headers?: Record + readonly fallback?: () => Stream.Stream + readonly rotateAfterMs?: number + } = {}, +): WebSocketChannelExchange => ({ + id, + connect: { + url: "wss://provider.test/responses", + headers: Headers.fromInput(input.headers), + rotateAfterMs: input.rotateAfterMs, + }, + fallback: input.fallback ?? (() => Stream.make(`fallback:${id}`)), + driver: { + create: () => Effect.succeed({ message: id, mode: "full" }), + observe: (_create, frame): Effect.Effect => + Effect.succeed({ type: "completed", frame }), + }, +}) + +const run = (connector: WebSocketConnector, effect: Effect.Effect) => + Effect.runPromise(effect.pipe(Effect.provide(SessionModelTransport.makeLayer(connector)), Effect.scoped)) + +const collect = (executor: ReturnType, item: WebSocketChannelExchange) => + Effect.gen(function* () { + const execution = yield* executor.execute(item) + return Array.from(yield* Stream.runCollect(execution.frames)) + }).pipe(Effect.scoped) + +const automatic = () => { + const connections: Array<{ + readonly messages: Queue.Queue + closed: number + sent: string[] + }> = [] + const connector: WebSocketConnector = { + open: () => + Effect.gen(function* () { + const messages = yield* Queue.unbounded() + const record = { messages, closed: 0, sent: [] as string[] } + connections.push(record) + const connection: WebSocketConnection = { + sendText: (message) => + Effect.sync(() => { + record.sent.push(message) + Queue.offerUnsafe(messages, `completed:${message}`) + }), + messages: Stream.fromQueue(messages), + close: Effect.sync(() => { + record.closed++ + }).pipe(Effect.andThen(Queue.shutdown(messages)), Effect.asVoid), + } + return connection + }), + } + return { connector, connections } +} + +describe("SessionModelTransport", () => { + test("reuses one physical connection for sequential Session calls", async () => { + const fixture = automatic() + + await run( + fixture.connector, + Effect.gen(function* () { + const transport = yield* SessionModelTransport.Service + expect(yield* collect(transport.bind(session), exchange("first"))).toEqual(["completed:first"]) + expect(yield* collect(transport.bind(session), exchange("second"))).toEqual(["completed:second"]) + expect(fixture.connections).toHaveLength(1) + expect(fixture.connections[0]?.sent).toEqual(["first", "second"]) + }), + ) + }) + + test("serializes concurrent calls for one Session", async () => { + const started = Deferred.makeUnsafe() + const release = Deferred.makeUnsafe() + const messages = queue() + const sent: string[] = [] + const connector: WebSocketConnector = { + open: () => + Effect.succeed({ + sendText: (message) => + Effect.gen(function* () { + sent.push(message) + if (message === "first") { + yield* Deferred.succeed(started, undefined) + yield* Deferred.await(release) + } + Queue.offerUnsafe(messages, `completed:${message}`) + }), + messages: Stream.fromQueue(messages), + close: Queue.shutdown(messages).pipe(Effect.asVoid), + }), + } + + await run( + connector, + Effect.gen(function* () { + const transport = yield* SessionModelTransport.Service + const executor = transport.bind(session) + const first = yield* collect(executor, exchange("first")).pipe(Effect.forkChild({ startImmediately: true })) + yield* Deferred.await(started) + const second = yield* collect(executor, exchange("second")).pipe(Effect.forkChild({ startImmediately: true })) + yield* Effect.yieldNow + expect(sent).toEqual(["first"]) + yield* Deferred.succeed(release, undefined) + yield* Fiber.join(first) + yield* Fiber.join(second) + expect(sent).toEqual(["first", "second"]) + }), + ) + }) + + test("isolates connections and permits concurrency across Sessions", async () => { + const started = queue() + const release = Deferred.makeUnsafe() + let opened = 0 + const connector: WebSocketConnector = { + open: () => + Effect.gen(function* () { + opened++ + const messages = yield* Queue.unbounded() + return { + sendText: (message) => + Effect.gen(function* () { + Queue.offerUnsafe(started, message) + yield* Deferred.await(release) + Queue.offerUnsafe(messages, `completed:${message}`) + }), + messages: Stream.fromQueue(messages), + close: Queue.shutdown(messages).pipe(Effect.asVoid), + } + }), + } + + await run( + connector, + Effect.gen(function* () { + const transport = yield* SessionModelTransport.Service + const first = yield* collect(transport.bind(session), exchange("first")).pipe( + Effect.forkChild({ startImmediately: true }), + ) + const second = yield* collect(transport.bind(otherSession), exchange("second")).pipe( + Effect.forkChild({ startImmediately: true }), + ) + expect(new Set([yield* Queue.take(started), yield* Queue.take(started)])).toEqual(new Set(["first", "second"])) + expect(opened).toBe(2) + yield* Deferred.succeed(release, undefined) + yield* Fiber.join(first) + yield* Fiber.join(second) + }), + ) + }) + + test("cancels a queued call without affecting the active exchange", async () => { + const started = Deferred.makeUnsafe() + const release = Deferred.makeUnsafe() + const messages = queue() + const sent: string[] = [] + const connector: WebSocketConnector = { + open: () => + Effect.succeed({ + sendText: (message) => + Effect.gen(function* () { + sent.push(message) + yield* Deferred.succeed(started, undefined) + yield* Deferred.await(release) + Queue.offerUnsafe(messages, `completed:${message}`) + }), + messages: Stream.fromQueue(messages), + close: Queue.shutdown(messages).pipe(Effect.asVoid), + }), + } + + await run( + connector, + Effect.gen(function* () { + const transport = yield* SessionModelTransport.Service + const executor = transport.bind(session) + const active = yield* collect(executor, exchange("active")).pipe(Effect.forkChild({ startImmediately: true })) + yield* Deferred.await(started) + const queued = yield* collect(executor, exchange("queued")).pipe(Effect.forkChild({ startImmediately: true })) + yield* Fiber.interrupt(queued) + expect(sent).toEqual(["active"]) + yield* Deferred.succeed(release, undefined) + expect(yield* Fiber.join(active)).toEqual(["completed:active"]) + }), + ) + }) + + test("closes the connection when an active exchange is interrupted", async () => { + const started = Deferred.makeUnsafe() + const messages = queue() + let closed = 0 + const connector: WebSocketConnector = { + open: () => + Effect.succeed({ + sendText: () => Deferred.succeed(started, undefined).pipe(Effect.andThen(Effect.never)), + messages: Stream.fromQueue(messages), + close: Effect.sync(() => closed++).pipe(Effect.andThen(Queue.shutdown(messages)), Effect.asVoid), + }), + } + + await run( + connector, + Effect.gen(function* () { + const transport = yield* SessionModelTransport.Service + const fiber = yield* collect(transport.bind(session), exchange("first")).pipe( + Effect.forkChild({ startImmediately: true }), + ) + yield* Deferred.await(started) + yield* Fiber.interrupt(fiber) + expect(closed).toBe(1) + }), + ) + }) + + test("closes a newly opened connection when request creation is interrupted", async () => { + const opened = Deferred.makeUnsafe() + const messages = queue() + let closed = 0 + const connector: WebSocketConnector = { + open: () => + Deferred.succeed(opened, undefined).pipe( + Effect.as({ + sendText: () => Effect.void, + messages: Stream.fromQueue(messages), + close: Effect.sync(() => closed++).pipe(Effect.andThen(Queue.shutdown(messages)), Effect.asVoid), + }), + ), + } + + await run( + connector, + Effect.gen(function* () { + const transport = yield* SessionModelTransport.Service + const item = exchange("first") + const fiber = yield* collect(transport.bind(session), { + ...item, + driver: { create: () => Effect.never, observe: item.driver.observe }, + }).pipe(Effect.forkChild({ startImmediately: true })) + yield* Deferred.await(opened) + yield* Fiber.interrupt(fiber) + expect(closed).toBe(1) + }), + ) + }) + + test("falls back once when connection setup fails before send", async () => { + let fallbacks = 0 + const connector: WebSocketConnector = { open: () => Effect.fail(error("upgrade rejected", "not-sent")) } + + await run( + connector, + Effect.gen(function* () { + const transport = yield* SessionModelTransport.Service + const result = yield* collect( + transport.bind(session), + exchange("first", { + fallback: () => { + fallbacks++ + return Stream.make("http") + }, + }), + ) + expect(result).toEqual(["http"]) + expect(fallbacks).toBe(1) + }), + ) + }) + + test("does not fall back after an ambiguous send failure", async () => { + const messages = queue() + let fallbacks = 0 + let closed = 0 + const connector: WebSocketConnector = { + open: () => + Effect.succeed({ + sendText: () => Effect.fail(error("send failed")), + messages: Stream.fromQueue(messages), + close: Effect.sync(() => closed++).pipe(Effect.andThen(Queue.shutdown(messages)), Effect.asVoid), + }), + } + + await run( + connector, + Effect.gen(function* () { + const transport = yield* SessionModelTransport.Service + const result = yield* Effect.result( + collect( + transport.bind(session), + exchange("first", { + fallback: () => { + fallbacks++ + return Stream.make("http") + }, + }), + ), + ) + expect(result).toMatchObject({ + _tag: "Failure", + failure: { reason: { _tag: "Transport", phase: "send", delivery: "ambiguous" } }, + }) + expect(fallbacks).toBe(0) + expect(closed).toBe(1) + }), + ) + }) + + test("rotates when handshake affinity or connection age changes", async () => { + const fixture = automatic() + + await run( + fixture.connector, + Effect.gen(function* () { + const transport = yield* SessionModelTransport.Service + const executor = transport.bind(session) + yield* collect(executor, exchange("first", { headers: { authorization: "one" } })) + yield* collect(executor, exchange("second", { headers: { authorization: "one" } })) + yield* collect(executor, exchange("third", { headers: { authorization: "two" } })) + yield* Effect.sleep("5 millis") + yield* collect(executor, exchange("fourth", { headers: { authorization: "two" }, rotateAfterMs: 1 })) + expect(fixture.connections).toHaveLength(3) + expect(fixture.connections.slice(0, 2).map((item) => item.closed)).toEqual([1, 1]) + }), + ) + }) + + test("poisons a socket that receives data while idle", async () => { + const fixture = automatic() + + await run( + fixture.connector, + Effect.gen(function* () { + const transport = yield* SessionModelTransport.Service + const executor = transport.bind(session) + yield* collect(executor, exchange("first")) + const connection = fixture.connections[0] + if (!connection) throw new Error("Expected connection") + Queue.offerUnsafe(connection.messages, "late") + yield* Effect.yieldNow + yield* collect(executor, exchange("second")) + expect(fixture.connections).toHaveLength(2) + expect(fixture.connections[0]?.closed).toBe(1) + }), + ) + }) + + test("poisons instead of dropping data when the inbound queue overflows", async () => { + const messages = queue() + let closed = 0 + const connector: WebSocketConnector = { + open: () => + Effect.succeed({ + sendText: () => + Effect.sync(() => { + for (let index = 0; index <= 129; index++) Queue.offerUnsafe(messages, `frame:${index}`) + }), + messages: Stream.fromQueue(messages), + close: Effect.sync(() => closed++).pipe(Effect.andThen(Queue.shutdown(messages)), Effect.asVoid), + }), + } + + await run( + connector, + Effect.gen(function* () { + const transport = yield* SessionModelTransport.Service + const item = exchange("first") + const result = yield* Effect.result( + collect(transport.bind(session), { + ...item, + driver: { + create: item.driver.create, + observe: (_create, frame) => Effect.sleep("1 millis").pipe(Effect.as({ type: "frame" as const, frame })), + }, + }), + ) + expect(result).toMatchObject({ + _tag: "Failure", + failure: { reason: { _tag: "Transport", kind: "queue-overflow", delivery: "accepted" } }, + }) + expect(closed).toBe(1) + }), + ) + }) + + test("poisons unsupported binary frames after provider observation", async () => { + const messages = queue() + const closed = Deferred.makeUnsafe() + const connector: WebSocketConnector = { + open: () => + Effect.succeed({ + sendText: () => Effect.sync(() => Queue.offerUnsafe(messages, new Uint8Array([1]))).pipe(Effect.asVoid), + messages: Stream.fromQueue(messages), + close: Deferred.succeed(closed, undefined).pipe(Effect.andThen(Queue.shutdown(messages)), Effect.asVoid), + }), + } + + await run( + connector, + Effect.gen(function* () { + const transport = yield* SessionModelTransport.Service + const result = yield* Effect.result(collect(transport.bind(session), exchange("first"))) + expect(result).toMatchObject({ + _tag: "Failure", + failure: { reason: { _tag: "Transport", kind: "message", delivery: "accepted" } }, + }) + yield* Deferred.await(closed) + }), + ) + }) + + test("closes individual and all owned connections", async () => { + const fixture = automatic() + + await run( + fixture.connector, + Effect.gen(function* () { + const transport = yield* SessionModelTransport.Service + yield* collect(transport.bind(session), exchange("first")) + yield* collect(transport.bind(otherSession), exchange("second")) + yield* transport.close(session) + expect(fixture.connections.map((item) => item.closed)).toEqual([1, 0]) + yield* transport.closeAll + expect(fixture.connections.map((item) => item.closed)).toEqual([1, 1]) + }), + ) + }) + + test("closes owned connections when the Location scope ends", async () => { + const fixture = automatic() + + await run( + fixture.connector, + Effect.gen(function* () { + const transport = yield* SessionModelTransport.Service + yield* collect(transport.bind(session), exchange("first")) + expect(fixture.connections[0]?.closed).toBe(0) + }), + ) + + expect(fixture.connections[0]?.closed).toBe(1) + }) +}) diff --git a/packages/core/test/session-remove.test.ts b/packages/core/test/session-remove.test.ts index cf3ce126cc67..06532fc424fb 100644 --- a/packages/core/test/session-remove.test.ts +++ b/packages/core/test/session-remove.test.ts @@ -1,4 +1,5 @@ import { describe, expect } from "bun:test" +import path from "node:path" import { Effect, Layer } from "effect" import { AppNodeBuilder } from "@opencode-ai/core/effect/app-node-builder" import { LayerNode } from "@opencode-ai/util/effect/layer-node" @@ -9,8 +10,10 @@ import { Project } from "@opencode-ai/core/project" import { AbsolutePath } from "@opencode-ai/core/schema" import { Session } from "@opencode-ai/core/session" import { SessionExecution } from "@opencode-ai/core/session/execution" +import { SessionModelTransport } from "@opencode-ai/core/session/model-transport" import { SessionProjector } from "@opencode-ai/core/session/projector" import { SessionStore } from "@opencode-ai/core/session/store" +import { LocationServiceMap } from "@opencode-ai/core/location-services" import { testEffect } from "./lib/effect" const projects = Layer.succeed( @@ -22,16 +25,33 @@ const projects = Layer.succeed( commit: () => Effect.void, }), ) +const closed: Session.ID[] = [] +const transport = Layer.succeed( + SessionModelTransport.Service, + SessionModelTransport.Service.of({ + bind: () => ({ execute: () => Effect.die("Unexpected WebSocket execution") }), + close: (sessionID) => Effect.sync(() => closed.push(sessionID)), + closeAll: Effect.void, + }), +) const it = testEffect( AppNodeBuilder.build( - LayerNode.group([Database.node, Bus.node, SessionProjector.node, SessionStore.node, Session.node]), + LayerNode.group([ + Database.node, + Bus.node, + SessionProjector.node, + SessionStore.node, + Session.node, + LocationServiceMap.node, + ]), [ [Project.node, projects], [SessionExecution.node, SessionExecution.noopLayer], + [SessionModelTransport.node, transport], ], ), ) -const location = Location.Ref.make({ directory: AbsolutePath.make("/project") }) +const location = Location.Ref.make({ directory: AbsolutePath.make(import.meta.dir) }) describe("Session.remove", () => { it.effect("removes a session and its children", () => @@ -39,10 +59,13 @@ describe("Session.remove", () => { const session = yield* Session.Service const parent = yield* session.create({ location }) const child = yield* session.create({ parentID: parent.id }) + yield* (yield* LocationServiceMap.Service).contextEffect(location) + closed.length = 0 yield* session.remove(parent.id) expect((yield* session.list()).data).toEqual([]) + expect(closed).toEqual([parent.id, child.id]) expect(yield* Effect.result(session.get(parent.id))).toMatchObject({ _tag: "Failure" }) expect(yield* Effect.result(session.get(child.id))).toMatchObject({ _tag: "Failure" }) }), @@ -60,3 +83,23 @@ describe("Session.remove", () => { }), ) }) + +describe("Session.move", () => { + it.effect("closes the source Location transport before moving", () => + Effect.gen(function* () { + const sessions = yield* Session.Service + const item = yield* sessions.create({ location }) + yield* (yield* LocationServiceMap.Service).contextEffect(location) + closed.length = 0 + const destination = AbsolutePath.make(path.dirname(import.meta.dir)) + + yield* sessions.move({ + sessionID: item.id, + directory: destination, + }) + + expect(closed).toEqual([item.id]) + expect((yield* sessions.get(item.id)).location.directory).toBe(destination) + }), + ) +})