Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 26 additions & 2 deletions packages/core/src/session/model-transport.ts
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ interface Channel {
interface State {
readonly lock: Semaphore.Semaphore
closed: boolean
httpFallback: boolean
channel?: Channel
}

Expand Down Expand Up @@ -119,7 +120,7 @@ export const makeLayer = (connector: WebSocketConnector) =>
const state = (sessionID: SessionSchema.ID) => {
const current = states.get(sessionID)
if (current) return current
const created = { lock: Semaphore.makeUnsafe(1), closed: false }
const created = { lock: Semaphore.makeUnsafe(1), closed: false, httpFallback: false }
states.set(sessionID, created)
return created
}
Expand Down Expand Up @@ -241,7 +242,9 @@ export const makeLayer = (connector: WebSocketConnector) =>
channel.active?.lifecycle.delivery === "terminal" ||
(error.reason._tag === "Transport" && error.reason.code === "queue-overflow")
? "accepted"
: "ambiguous",
: error.reason._tag === "Transport" && error.reason.code === "1009"
? "rejected"
: "ambiguous",
}),
),
),
Expand Down Expand Up @@ -274,6 +277,7 @@ export const makeLayer = (connector: WebSocketConnector) =>
phase: "queue",
delivery: "not-sent",
})
if (owner.httpFallback) return fallback(exchange)
const key = affinity(exchange)
const now = yield* Clock.currentTimeMillis
const current = owner.channel
Expand Down Expand Up @@ -420,6 +424,26 @@ export const makeLayer = (connector: WebSocketConnector) =>
yield* poison(owner, channel, error)
}),
),
Stream.catch((error) => {
if (
error.reason._tag !== "Transport" ||
error.reason.code !== "1009" ||
error.reason.delivery !== "rejected"
)
return Stream.fail(error)
owner.httpFallback = true
return Stream.unwrap(
Effect.logWarning("session websocket request too large; using http", {
sessionTransport: "websocket",
phase: "close",
delivery: "rejected",
code: error.reason.code,
}).pipe(
Effect.andThen(metric("fallback", { reason: "message_too_large" })),
Effect.as(exchange.fallback()),
),
)
}),
)
const complete = Effect.sync(() => {
if (owner.channel !== channel || channel.pending?.token !== token) return
Expand Down
59 changes: 58 additions & 1 deletion packages/core/test/session-model-transport.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ import type {
} 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, Metric, Queue, Stream } from "effect"
import { Cause, Deferred, Effect, Fiber, Metric, Queue, Stream } from "effect"
import { TestClock } from "effect/testing"
import { Headers } from "effect/unstable/http"

Expand Down Expand Up @@ -545,6 +545,63 @@ describe("SessionModelTransport", () => {
)
})

test("falls back to HTTP after close code 1009 and keeps the Session on HTTP", async () => {
const messages = queue<string | Uint8Array, AIError>()
let opened = 0
let fallbacks = 0
let closed = 0
const connector: WebSocketConnector = {
open: () =>
Effect.sync(() => {
opened++
return {
sendText: () =>
Effect.sync(() => {
Queue.failCauseUnsafe(
messages,
Cause.fail(
new AIError({
module: "test",
method: "websocket",
reason: new TransportReason({
message: "message too big",
transport: "websocket",
operation: "read",
code: "1009",
phase: "close",
}),
}),
),
)
}),
messages: Stream.fromQueue(messages),
close: Effect.sync(() => closed++).pipe(Effect.andThen(Queue.shutdown(messages)), Effect.asVoid),
}
}),
}
const item = (id: string) =>
exchange(id, {
fallback: () => {
fallbacks++
return Stream.make(`http:${id}`)
},
})

await run(
connector,
Effect.gen(function* () {
const transport = yield* SessionModelTransport.Service
const executor = transport.bind(session)

expect(yield* collect(executor, item("first"))).toEqual(["http:first"])
expect(yield* collect(executor, item("second"))).toEqual(["http:second"])
expect(opened).toBe(1)
expect(fallbacks).toBe(2)
expect(closed).toBe(1)
}),
)
})

test("does not fall back after an ambiguous send failure", async () => {
const messages = queue<string | Uint8Array, AIError>()
let fallbacks = 0
Expand Down
Loading