anomalyco/opencode
1
0
Fork 0
mirror of https://github.com/anomalyco/opencode.git synced 2026-08-15 07:48:24 -04:00
Code Issues 5.9k Packages Projects Releases 19 Wiki Activity

feat(ai): add Responses websocket continuation (#42596)

Browse Source
This commit is contained in:
Shoubhit Dash
2026-08-14 22:06:29 +05:30
committed by GitHub
parent a2a53032d4
commit 445aacf92e
8 changed files with 676 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ai/src/protocols/open-responses-channel.ts
+20 -9
View File
@@ -23,6 +23,11 @@ export interface Options {
readonly name: string
readonly rotateAfterMs?: number
readonly headers?: (headers: Headers.Headers) => Headers.Headers
readonly driver?: (input: {
readonly request: Readonly<Record<string, unknown>>
readonly message: string
readonly base: WebSocketChannelDriver
}) => WebSocketChannelDriver
}
export interface Prepared {
@@ -40,7 +45,8 @@ const message = (body: unknown) =>
if (!ProviderShared.isRecord(body))
return yield* ProviderShared.invalidRequest("Open Responses WebSocket body must be a JSON object")
const { stream: _stream, stream_options: _streamOptions, background: _background, ...request } = body
return encodeMessage(yield* decodeMessage({ ...request, type: "response.create" }))
const decoded = yield* decodeMessage({ ...request, type: "response.create" })
return { request: decoded, message: encodeMessage(decoded) }
})
const driver = (options: Options, body: string): WebSocketChannelDriver => {
@@ -141,20 +147,25 @@ export const transport = <Body>(options: Options): Transport<Body, Prepared, str
Effect.gen(function* () {
const parts = yield* HttpTransport.jsonRequestParts(input)
const headers = Headers.remove(options.headers?.(parts.headers) ?? parts.headers, "content-length")
const channel = input.webSocket
? yield* Effect.gen(function* () {
const create = yield* message(parts.jsonBody)
const base = driver(options, create.message)
return {
url: yield* WebSocketTransport.toWebSocketUrl(parts.url),
headers,
rotateAfterMs: options.rotateAfterMs,
driver: options.driver?.({ request: create.request, message: create.message, base }) ?? base,
}
})
: undefined
return {
http: {
request: ProviderShared.jsonPost({ url: parts.url, body: parts.bodyText, headers: parts.headers }),
framing: Framing.sse,
middleware: input.middleware,
},
channel: input.webSocket
? {
url: yield* WebSocketTransport.toWebSocketUrl(parts.url),
headers,
rotateAfterMs: options.rotateAfterMs,
driver: driver(options, yield* message(parts.jsonBody)),
}
: undefined,
channel,
}
}),
execute: (prepared, request, runtime, executeOptions) => {
packages/ai/src/protocols/openai-responses-channel.ts
+159
View File
@@ -0,0 +1,159 @@
import { AIError, TransportReason } from "../schema/index.js"
import type { ChannelCheckpoint, ChannelObservation, WebSocketChannelDriver } from "../route/transport/index.js"
import { Effect, Schema } from "effect"
import * as ProviderShared from "./shared.js"
import { OpenResponses } from "./open-responses.js"
const PROTOCOL = "openai-responses.websocket.v1"
const VERSION = 1
const decodeEvent = Schema.decodeUnknownEffect(OpenResponses.protocol.stream.event)
interface CheckpointValue {
readonly version: typeof VERSION
readonly responseID: string
readonly request: Readonly<Record<string, unknown>>
readonly output: ReadonlyArray<unknown>
}
export interface DriverInput {
readonly id: string
readonly name: string
readonly request: Readonly<Record<string, unknown>>
readonly message: string
readonly base: WebSocketChannelDriver
}
const checkpointValue = (checkpoint: ChannelCheckpoint | undefined): CheckpointValue | undefined => {
if (checkpoint?.protocol !== PROTOCOL || !ProviderShared.isRecord(checkpoint.value)) return undefined
if (checkpoint.value.version !== VERSION) return undefined
if (typeof checkpoint.value.responseID !== "string" || checkpoint.value.responseID.trim().length === 0)
return undefined
if (!ProviderShared.isRecord(checkpoint.value.request) || !Array.isArray(checkpoint.value.output)) return undefined
return {
version: VERSION,
responseID: checkpoint.value.responseID,
request: checkpoint.value.request,
output: checkpoint.value.output,
}
}
const canonical = (value: unknown): string => {
if (value === undefined) return "undefined"
if (Array.isArray(value)) return `[${value.map(canonical).join(",")}]`
if (!ProviderShared.isRecord(value)) return ProviderShared.encodeJson(value)
return `{${Object.keys(value)
.sort()
.map((key) => `${ProviderShared.encodeJson(key)}:${canonical(value[key])}`)
.join(",")}}`
}
const comparable = (value: unknown) => {
if (!ProviderShared.isRecord(value)) return value
if (value.type === "message" && value.role === "assistant")
return {
role: "assistant",
content: value.content,
...(value.phase === undefined ? {} : { phase: value.phase }),
}
if (value.type === "function_call")
return {
type: value.type,
call_id: value.call_id,
name: value.name,
arguments: value.arguments,
}
if (value.type === "reasoning")
return {
type: value.type,
...(value.id === undefined ? {} : { id: value.id }),
summary: value.summary,
encrypted_content: value.encrypted_content,
}
return value
}
const invariant = (request: Readonly<Record<string, unknown>>) => {
const { type: _type, input: _input, previous_response_id: _previousResponseID, ...rest } = request
return rest
}
const incremental = (
request: Readonly<Record<string, unknown>>,
checkpoint: CheckpointValue,
): ReadonlyArray<unknown> | undefined => {
const input = request.input
const previousInput = checkpoint.request.input
if (!Array.isArray(input) || !Array.isArray(previousInput)) return undefined
if (canonical(invariant(request)) !== canonical(invariant(checkpoint.request))) return undefined
const baseline = [...previousInput, ...checkpoint.output]
if (input.length <= baseline.length) return undefined
if (!baseline.every((item, index) => canonical(comparable(item)) === canonical(input[index]))) return undefined
return input.slice(baseline.length)
}
const code = (event: OpenResponses.Event) => event.code || event.error?.code || event.response?.error?.code || undefined
const rejected = (
input: DriverInput,
observation: Extract<ChannelObservation, { readonly type: "provider-failure" }>,
recovery: "retry-full" | "rotate-and-retry-full",
): ChannelObservation => ({
type: "rejected",
recovery,
error: new AIError({
module: input.id,
method: "stream",
reason: new TransportReason({
message: observation.error.message,
transport: "websocket",
operation: "read",
phase: "receive",
delivery: "rejected",
recovery,
}),
}),
})
export const driver = (input: DriverInput): WebSocketChannelDriver => {
const { previous_response_id: _previousResponseID, ...request } = input.request
let output: unknown[] = []
return {
create: (checkpoint) =>
Effect.sync(() => {
output = []
const previous = checkpointValue(checkpoint)
const delta = previous ? incremental(request, previous) : undefined
if (!previous || !delta) return { message: ProviderShared.encodeJson(request), mode: "full" as const }
return {
message: ProviderShared.encodeJson({ ...request, input: delta, previous_response_id: previous.responseID }),
mode: "incremental" as const,
}
}),
observe: (create, frame) =>
Effect.gen(function* () {
const event = yield* decodeEvent(frame).pipe(
Effect.mapError(() => ProviderShared.eventError(input.id, `Invalid ${input.name} WebSocket event`, frame)),
)
const observation = yield* input.base.observe(create, frame)
if (event.type === "response.output_item.done" && event.item) output.push(event.item)
if (observation.type === "provider-failure") {
const rejection = code(event)
if (rejection === "previous_response_not_found") return rejected(input, observation, "retry-full")
if (rejection === "websocket_connection_limit_reached")
return rejected(input, observation, "rotate-and-retry-full")
}
if (observation.type !== "completed") return observation
const responseID = event.response?.id
if (!responseID || responseID.trim().length === 0) return observation
return {
...observation,
checkpoint: {
protocol: PROTOCOL,
value: { version: VERSION, responseID, request, output: output.slice() } satisfies CheckpointValue,
},
}
}),
}
}
export const OpenAIResponsesChannel = { driver } as const
packages/ai/src/protocols/openai-responses.ts
+2
View File
@@ -12,6 +12,7 @@ import { Lifecycle } from "./utils/lifecycle.js"
import { OpenAIImage } from "./utils/openai-image.js"
import { ToolSchemaProjection } from "./utils/tool-schema.js"
import { OpenResponsesChannel } from "./open-responses-channel.js"
import { OpenAIResponsesChannel } from "./openai-responses-channel.js"
const ADAPTER = "openai-responses"
const NAME = "OpenAI Responses"
@@ -248,6 +249,7 @@ export const transport = OpenResponsesChannel.transport<OpenAIResponsesBody>({
name: NAME,
rotateAfterMs: WEBSOCKET_ROTATE_AFTER_MS,
headers: (headers) => Headers.set(headers, "openai-beta", headers["openai-beta"] ?? WEBSOCKET_PROTOCOL_HEADER),
driver: (input) => OpenAIResponsesChannel.driver({ id: ADAPTER, name: NAME, ...input }),
})
export const route = Route.make({
packages/ai/test/provider/openai-responses.test.ts
+233 -1
View File
@@ -15,12 +15,20 @@ import {
TransportReason,
Usage,
} from "../../src/index.js"
import { Auth, LLMClient, RequestExecutor, WebSocketTransport } from "../../src/route.js"
import {
Auth,
LLMClient,
RequestExecutor,
WebSocketTransport,
type ChannelObservation,
type WebSocketChannelDriver,
} from "../../src/route.js"
import { compileRequest } from "../../src/route/client.js"
import * as Azure from "../../src/providers/azure.js"
import * as OpenAI from "../../src/providers/openai.js"
import * as XAI from "../../src/providers/xai.js"
import * as OpenAIResponses from "../../src/protocols/openai-responses.js"
import { OpenAIResponsesChannel } from "../../src/protocols/openai-responses-channel.js"
import * as ProviderShared from "../../src/protocols/shared.js"
import { continuationRequest, nativeOpenAIResponsesContinuation } from "../continuation-scenarios.js"
import { it } from "../lib/effect.js"
@@ -33,6 +41,47 @@ const model = OpenAIResponses.route
const xaiModel = XAI.configure({ apiKey: "test", baseURL: "https://api.x.ai/v1" }).responses("grok-4.5")
const baseChannelDriver = (message: string): WebSocketChannelDriver => ({
create: () => Effect.succeed({ message, mode: "full" }),
observe: (_create, frame): Effect.Effect<ChannelObservation, AIError> => {
const event = ProviderShared.decodeJson(frame)
if (!ProviderShared.isRecord(event)) return Effect.die("Expected event")
if (event.type === "response.completed") return Effect.succeed({ type: "completed", frame })
if (event.type === "response.incomplete") return Effect.succeed({ type: "incomplete", frame })
if (event.type === "error" || event.type === "response.failed")
return Effect.succeed({
type: "provider-failure",
error: new AIError({
module: "test",
method: "stream",
reason: new TransportReason({
message: "provider rejected request",
transport: "websocket",
operation: "read",
phase: "receive",
}),
}),
})
return Effect.succeed({ type: "frame", frame })
},
})
const continuationDriver = (request: Readonly<Record<string, unknown>>) => {
const message = ProviderShared.encodeJson(request)
return OpenAIResponsesChannel.driver({
id: "openai-responses",
name: "OpenAI Responses",
request,
message,
base: baseChannelDriver(message),
})
}
const checkpoint = (observation: ChannelObservation) => {
if (observation.type !== "completed" || !observation.checkpoint) throw new Error("Expected checkpoint")
return observation.checkpoint
}
const request = LLM.request({
id: "req_1",
model,
@@ -344,6 +393,189 @@ describe("OpenAI Responses route", () => {
}),
)
it.effect("continues a tool call with only the new tool output", () =>
Effect.gen(function* () {
const firstRequest = {
type: "response.create",
model: "gpt-5.2",
store: false,
input: [{ role: "user", content: [{ type: "input_text", text: "Weather?" }] }],
}
const first = continuationDriver(firstRequest)
const firstCreate = yield* first.create(undefined)
yield* first.observe(
firstCreate,
ProviderShared.encodeJson({
type: "response.output_item.done",
item: {
type: "function_call",
id: "fc_1",
status: "completed",
call_id: "call_1",
name: "weather",
arguments: '{"city":"Paris"}',
},
}),
)
const saved = checkpoint(
yield* first.observe(
firstCreate,
ProviderShared.encodeJson({ type: "response.completed", response: { id: "resp_1" } }),
),
)
const second = continuationDriver({
...firstRequest,
input: [
...firstRequest.input,
{ type: "function_call", call_id: "call_1", name: "weather", arguments: '{"city":"Paris"}' },
{ type: "function_call_output", call_id: "call_1", output: '{"temperature":22}' },
],
})
const create = yield* second.create(saved)
expect(create.mode).toBe("incremental")
expect(ProviderShared.decodeJson(create.message)).toMatchObject({
previous_response_id: "resp_1",
input: [{ type: "function_call_output", call_id: "call_1", output: '{"temperature":22}' }],
})
}),
)
it.effect("continues a promoted steer after the completed assistant output", () =>
Effect.gen(function* () {
const firstInput = [{ role: "user", content: [{ type: "input_text", text: "First" }] }]
const first = continuationDriver({ type: "response.create", model: "gpt-5.2", store: false, input: firstInput })
const create = yield* first.create(undefined)
yield* first.observe(
create,
ProviderShared.encodeJson({
type: "response.output_item.done",
item: {
type: "message",
id: "msg_1",
status: "completed",
role: "assistant",
content: [{ type: "output_text", text: "Hello" }],
},
}),
)
const saved = checkpoint(
yield* first.observe(
create,
ProviderShared.encodeJson({ type: "response.completed", response: { id: "resp_1" } }),
),
)
const steer = { role: "user", content: [{ type: "input_text", text: "Actually, be brief" }] }
const next = continuationDriver({
type: "response.create",
model: "gpt-5.2",
store: false,
input: [...firstInput, { role: "assistant", content: [{ type: "output_text", text: "Hello" }] }, steer],
})
const continued = yield* next.create(saved)
expect(continued.mode).toBe("incremental")
expect(ProviderShared.decodeJson(continued.message)).toMatchObject({
previous_response_id: "resp_1",
input: [steer],
})
}),
)
it.effect("uses a full request when any non-input invariant changes", () =>
Effect.gen(function* () {
const request = {
type: "response.create",
model: "gpt-5.2",
store: false,
metadata: { source: "one" },
input: [{ role: "user", content: [{ type: "input_text", text: "First" }] }],
}
const first = continuationDriver(request)
const create = yield* first.create(undefined)
const saved = checkpoint(
yield* first.observe(
create,
ProviderShared.encodeJson({ type: "response.completed", response: { id: "resp_1" } }),
),
)
const appended = [...request.input, { role: "user", content: [{ type: "input_text", text: "Second" }] }]
const changes = [
{ ...request, model: "gpt-5.3", input: appended },
{ ...request, instructions: "Changed", input: appended },
{ ...request, tools: [{ type: "function", name: "other" }], input: appended },
{ ...request, temperature: 0.5, input: appended },
{ ...request, metadata: { source: "two" }, input: appended },
{
...request,
input: [{ role: "user", content: [{ type: "input_text", text: "Rewritten history" }] }, appended[1]],
},
]
const creates = yield* Effect.forEach(changes, (changed) => continuationDriver(changed).create(saved))
expect(creates.map((item) => item.mode)).toEqual(changes.map(() => "full"))
expect(
creates
.map((item) => ProviderShared.decodeJson(item.message))
.every((item) => ProviderShared.isRecord(item) && !("previous_response_id" in item)),
).toBe(true)
}),
)
it.effect("stages no checkpoint for incomplete or ID-less completion", () =>
Effect.gen(function* () {
const driver = continuationDriver({ type: "response.create", model: "gpt-5.2", input: [] })
const create = yield* driver.create(undefined)
const completed = yield* driver.observe(
create,
ProviderShared.encodeJson({ type: "response.completed", response: {} }),
)
expect(completed).toMatchObject({ type: "completed" })
expect(completed).not.toHaveProperty("checkpoint")
expect(
yield* driver.observe(create, ProviderShared.encodeJson({ type: "response.incomplete", response: {} })),
).toMatchObject({ type: "incomplete" })
}),
)
it.effect("classifies explicit continuation rejection for runner-owned recovery", () =>
Effect.gen(function* () {
const driver = continuationDriver({ type: "response.create", model: "gpt-5.2", input: [] })
const create = yield* driver.create(undefined)
const missing = yield* driver.observe(
create,
ProviderShared.encodeJson({
type: "error",
error: { code: "previous_response_not_found", message: "Missing response" },
}),
)
const limit = yield* driver.observe(
create,
ProviderShared.encodeJson({
type: "error",
error: { code: "websocket_connection_limit_reached", message: "Rotate" },
}),
)
expect(missing).toMatchObject({
type: "rejected",
recovery: "retry-full",
error: { reason: { _tag: "Transport", delivery: "rejected", recovery: "retry-full" } },
})
expect(limit).toMatchObject({
type: "rejected",
recovery: "rotate-and-retry-full",
error: {
reason: { _tag: "Transport", delivery: "rejected", recovery: "rotate-and-retry-full" },
},
})
}),
)
it.effect("builds WebSocket and HTTP fallback from the same final request", () =>
Effect.gen(function* () {
const attempts = yield* Ref.make(0)
packages/core/src/session/model-transport.ts
+37 -6
View File
@@ -3,6 +3,7 @@ export * as SessionModelTransport from "./model-transport.js"
import {
WebSocketTransport,
type ChannelObservation,
type ChannelCheckpoint,
type WebSocketChannelExchange,
type WebSocketChannelExecution,
type WebSocketChannelExecutor,
@@ -35,6 +36,8 @@ interface Channel {
active?: Active
closing: boolean
poisoned: boolean
checkpoint?: ChannelCheckpoint
pending?: { readonly token: object; readonly checkpoint: ChannelCheckpoint }
reader?: Fiber.Fiber<unknown, unknown>
}
@@ -301,10 +304,16 @@ export const makeLayer = (connector: WebSocketConnector) =>
if (!channel) return fallback(exchange)
lifecycle.delivery = "ready"
const create = yield* exchange.driver.create(undefined).pipe(
if (channel.pending) {
channel.pending = undefined
channel.checkpoint = undefined
}
const create = yield* exchange.driver.create(channel.checkpoint).pipe(
Effect.tapError(() => closeChannel(owner, channel)),
Effect.onInterrupt(() => closeChannel(owner, channel)),
)
if (create.mode === "full") channel.checkpoint = undefined
const active: Active = { queue: yield* Queue.bounded<string, AIError>(INBOUND_CAPACITY), lifecycle }
channel.active = active
lifecycle.delivery = "send-attempted"
@@ -320,7 +329,9 @@ export const makeLayer = (connector: WebSocketConnector) =>
return yield* annotate(failure, { phase: "send", delivery: "ambiguous" })
}
let terminal = false
let terminal: ChannelObservation | undefined
const token = {}
let staged: ChannelCheckpoint | undefined
const frames = Stream.fromQueue(active.queue).pipe(
Stream.timeoutOrElse({
duration: IDLE_TIMEOUT,
@@ -339,8 +350,10 @@ export const makeLayer = (connector: WebSocketConnector) =>
Stream.tap((observation) =>
Effect.sync(() => {
if (!observationTerminal(observation)) return
terminal = true
terminal = observation
lifecycle.delivery = "terminal"
staged = observation.type === "completed" ? observation.checkpoint : undefined
if (observation.type !== "completed" || !staged) channel.checkpoint = undefined
}),
),
Stream.takeUntil(observationTerminal),
@@ -350,7 +363,14 @@ export const makeLayer = (connector: WebSocketConnector) =>
if (channel.active === active) channel.active = undefined
const pending = yield* Queue.size(active.queue)
yield* Queue.shutdown(active.queue)
if (terminal && pending === 0) return
if (terminal && pending === 0) {
if (staged) channel.pending = { token, checkpoint: staged }
if (terminal.type === "rejected" && terminal.recovery === "rotate-and-retry-full")
yield* closeChannel(owner, channel)
return
}
channel.checkpoint = undefined
channel.pending = undefined
const error = terminal
? transportError("receive", "WebSocket data arrived after the terminal event", {
url: exchange.connect.url,
@@ -370,21 +390,32 @@ export const makeLayer = (connector: WebSocketConnector) =>
}),
),
)
return { frames, complete: Effect.void }
const complete = Effect.sync(() => {
if (owner.channel !== channel || channel.pending?.token !== token) return
channel.checkpoint = channel.pending.checkpoint
channel.pending = undefined
})
return { frames, complete }
})
const bind = (sessionID: SessionSchema.ID): WebSocketChannelExecutor => ({
execute: (exchange) => {
const owner = state(sessionID)
const lifecycle = { delivery: "queued" as Delivery }
let complete = Effect.void
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.tap((execution) =>
Effect.sync(() => {
complete = execution.complete
}),
),
Effect.map((execution) => execution.frames),
),
),
complete: Effect.void,
complete: Effect.suspend(() => complete),
})
},
})
packages/core/src/session/runner/llm.ts
+14 -1
View File
@@ -35,7 +35,7 @@ import { SessionRunnerRetry } from "./retry.js"
import { SessionUsage } from "../usage.js"
import { ToolOutput } from "../../tool-output.js"
/** How one model call ended: settled, awaiting a scheduled retry, or restarted by compaction. */
/** How one model call ended: settled, awaiting retry/recovery, or restarted by compaction. */
type CallOutcome = Data.TaggedEnum<{
Completed: { readonly needsContinuation: boolean; readonly step: number }
Retry: { readonly step: number }
@@ -44,6 +44,7 @@ type CallOutcome = Data.TaggedEnum<{
readonly error: SessionRunnerRetry.RetryableFailure["error"]
readonly step: number
}
RecoverFull: { readonly step: number }
Restart: { readonly step: number; readonly recoveredOverflow: boolean }
}>
const CallOutcome = Data.taggedEnum<CallOutcome>()
@@ -207,12 +208,15 @@ const layer = Layer.effect(
let currentStep = step
// Overflow recovery is one-shot: a call after recovery must not recover another overflow.
let recoverOverflow = true
// Continuation rejection permits one immediate full-context Physical Attempt without generic backoff.
let recoverContinuation = true
while (true) {
const outcome = yield* callModel(
sessionID,
currentPromotable,
currentStep,
recoverOverflow,
recoverContinuation,
assistantMessageID,
).pipe(Effect.catchTag("SessionRunner.RetryableFailure", waitForRetry))
if (outcome._tag === "Completed") return { needsContinuation: outcome.needsContinuation, step: outcome.step }
@@ -234,6 +238,7 @@ const layer = Layer.effect(
if (outcome.recoveredOverflow) recoverOverflow = false
assistantMessageID = SessionMessage.ID.create()
}
if (outcome._tag === "RecoverFull") recoverContinuation = false
// Neither a retry nor a compaction restart re-promotes input.
currentPromotable = undefined
currentStep = outcome.step
@@ -249,6 +254,7 @@ const layer = Layer.effect(
promotable: SessionInbox.Promotable | undefined,
step: number,
recoverOverflow: boolean,
recoverContinuation: boolean,
assistantMessageID: SessionMessage.ID,
) {
const selected = yield* context.select(sessionID)
@@ -414,6 +420,13 @@ const layer = Layer.effect(
// escapes as a scheduled retry or fails the assistant durably.
const llmFailure = streamFailure instanceof AIError ? streamFailure : undefined
const llmError = llmFailure && !publisher.record().providerFailed ? toSessionError(llmFailure) : undefined
if (
recoverContinuation &&
llmFailure?.reason._tag === "Transport" &&
(llmFailure.reason.recovery === "retry-full" || llmFailure.reason.recovery === "rotate-and-retry-full") &&
!publisher.record().outputStarted
)
return CallOutcome.RecoverFull({ step: currentStep })
if (
llmFailure &&
llmError &&
packages/core/test/session-model-transport.test.ts
+167
View File
@@ -66,6 +66,15 @@ const collect = (executor: ReturnType<SessionModelTransport.Interface["bind"]>,
return Array.from(yield* Stream.runCollect(execution.frames))
}).pipe(Effect.scoped)
const collectComplete = (
executor: ReturnType<SessionModelTransport.Interface["bind"]>,
item: WebSocketChannelExchange,
) =>
Effect.gen(function* () {
const execution = yield* executor.execute(item)
return Array.from(yield* Stream.runCollect(execution.frames.pipe(Stream.onEnd(execution.complete))))
}).pipe(Effect.scoped)
const automatic = () => {
const connections: Array<{
readonly messages: Queue.Queue<string | Uint8Array, AIError>
@@ -96,6 +105,164 @@ const automatic = () => {
}
describe("SessionModelTransport", () => {
test("commits checkpoints only after successful outer completion", async () => {
const messages = queue<string | Uint8Array, AIError>()
const checkpoints: Array<unknown> = []
const candidate = { protocol: "test", value: { response: "one" } }
const connector: WebSocketConnector = {
open: () =>
Effect.succeed({
sendText: (message) =>
Effect.sync(() => Queue.offerUnsafe(messages, `completed:${message}`)).pipe(Effect.asVoid),
messages: Stream.fromQueue(messages),
close: Queue.shutdown(messages).pipe(Effect.asVoid),
}),
}
const item = (id: string): WebSocketChannelExchange => ({
...exchange(id),
driver: {
create: (checkpoint) =>
Effect.sync(() => {
checkpoints.push(checkpoint)
return { message: id, mode: checkpoint ? "incremental" : "full" }
}),
observe: (_create, frame) => Effect.succeed({ type: "completed", frame, checkpoint: candidate }),
},
})
await run(
connector,
Effect.gen(function* () {
const transport = yield* SessionModelTransport.Service
const executor = transport.bind(session)
yield* collectComplete(executor, item("first"))
yield* collect(executor, item("second"))
yield* collect(executor, item("third"))
expect(checkpoints).toEqual([undefined, candidate, undefined])
}),
)
})
test("does not carry a checkpoint across physical connection rotation", async () => {
const fixture = automatic()
const checkpoints: Array<unknown> = []
const candidate = { protocol: "test", value: { response: "one" } }
const item = (id: string, authorization: string): WebSocketChannelExchange => ({
...exchange(id, { headers: { authorization } }),
driver: {
create: (checkpoint) =>
Effect.sync(() => {
checkpoints.push(checkpoint)
return { message: id, mode: checkpoint ? "incremental" : "full" }
}),
observe: (_create, frame) => Effect.succeed({ type: "completed", frame, checkpoint: candidate }),
},
})
await run(
fixture.connector,
Effect.gen(function* () {
const transport = yield* SessionModelTransport.Service
const executor = transport.bind(session)
yield* collectComplete(executor, item("first", "one"))
yield* collect(executor, item("second", "two"))
expect(checkpoints).toEqual([undefined, undefined])
expect(fixture.connections).toHaveLength(2)
}),
)
})
test("clears a rejected checkpoint before the runner retries full", async () => {
const fixture = automatic()
const checkpoints: Array<unknown> = []
const candidate = { protocol: "test", value: { response: "one" } }
const item = (id: string): WebSocketChannelExchange => ({
...exchange(id),
driver: {
create: (checkpoint) =>
Effect.sync(() => {
checkpoints.push(checkpoint)
return { message: id, mode: checkpoint ? "incremental" : "full" }
}),
observe: (_create, frame) =>
id === "rejected"
? Effect.succeed({
type: "rejected",
recovery: "retry-full",
error: new AIError({
module: "test",
method: "stream",
reason: new TransportReason({
message: "missing response",
transport: "websocket",
operation: "read",
phase: "receive",
delivery: "rejected",
recovery: "retry-full",
}),
}),
})
: Effect.succeed({ type: "completed", frame, checkpoint: candidate }),
},
})
await run(
fixture.connector,
Effect.gen(function* () {
const transport = yield* SessionModelTransport.Service
const executor = transport.bind(session)
yield* collectComplete(executor, item("first"))
yield* Effect.result(collect(executor, item("rejected")))
yield* collect(executor, item("retry"))
expect(checkpoints).toEqual([undefined, candidate, undefined])
expect(fixture.connections).toHaveLength(1)
}),
)
})
test("rotates after the provider rejects the connection generation", async () => {
const fixture = automatic()
const rejected: WebSocketChannelExchange = {
...exchange("rejected"),
driver: {
create: () => Effect.succeed({ message: "rejected", mode: "incremental" }),
observe: () =>
Effect.succeed({
type: "rejected",
recovery: "rotate-and-retry-full",
error: new AIError({
module: "test",
method: "stream",
reason: new TransportReason({
message: "connection limit",
transport: "websocket",
operation: "read",
phase: "receive",
delivery: "rejected",
recovery: "rotate-and-retry-full",
}),
}),
}),
},
}
await run(
fixture.connector,
Effect.gen(function* () {
const transport = yield* SessionModelTransport.Service
const executor = transport.bind(session)
yield* Effect.result(collect(executor, rejected))
yield* collect(executor, exchange("retry"))
expect(fixture.connections).toHaveLength(2)
expect(fixture.connections[0]?.closed).toBe(1)
}),
)
})
test("reuses one physical connection for sequential Session calls", async () => {
const fixture = automatic()
packages/core/test/session-runner.test.ts
+44
View File
@@ -526,6 +526,20 @@ const providerUnavailable = () =>
}),
})
const continuationRejected = (recovery: "retry-full" | "rotate-and-retry-full") =>
new AIError({
module: "test",
method: "stream",
reason: new TransportReason({
message: "Continuation rejected",
transport: "websocket",
operation: "read",
phase: "receive",
delivery: "rejected",
recovery,
}),
})
const incompleteStream = () =>
new AIError({
module: "test",
@@ -4093,6 +4107,36 @@ describe("SessionRunnerLLM", () => {
}),
)
it.effect("immediately rebuilds once after explicit continuation rejection", () =>
Effect.gen(function* () {
const session = yield* setup
yield* TestLLM.push(Stream.fail(continuationRejected("retry-full")))
yield* TestLLM.push(TestLLM.text("Recovered", "continuation-recovery"))
yield* runPrompt(session, "Recover continuation")
expect(requests).toHaveLength(2)
expect(yield* recordedEventTypes(sessionID)).not.toContain("session.retry.scheduled.1")
expect(yield* session.context(sessionID)).toMatchObject([
{ type: "user" },
{ type: "assistant", finish: "stop", content: [{ type: "text", text: "Recovered" }] },
])
}),
)
it.effect("bounds repeated continuation rejection to one immediate recovery", () =>
Effect.gen(function* () {
const session = yield* setup
const failure = continuationRejected("rotate-and-retry-full")
yield* TestLLM.push(Stream.fail(failure), Stream.fail(failure))
expect(yield* runPrompt(session, "Reject continuation twice").pipe(Effect.flip)).toBe(failure)
expect(requests).toHaveLength(2)
expect(yield* recordedEventTypes(sessionID)).not.toContain("session.retry.scheduled.1")
}),
)
it.effect("retries an incomplete stream before output", () =>
Effect.gen(function* () {
const session = yield* setup