From 445aacf92ef885bfea1ce1db0b69f3f996ab6a39 Mon Sep 17 00:00:00 2001 From: Shoubhit Dash Date: Fri, 14 Aug 2026 22:06:29 +0530 Subject: [PATCH] feat(ai): add Responses websocket continuation (#42596) --- .../src/protocols/open-responses-channel.ts | 29 ++- .../src/protocols/openai-responses-channel.ts | 159 ++++++++++++ packages/ai/src/protocols/openai-responses.ts | 2 + .../ai/test/provider/openai-responses.test.ts | 234 +++++++++++++++++- packages/core/src/session/model-transport.ts | 43 +++- packages/core/src/session/runner/llm.ts | 15 +- .../core/test/session-model-transport.test.ts | 167 +++++++++++++ packages/core/test/session-runner.test.ts | 44 ++++ 8 files changed, 676 insertions(+), 17 deletions(-) create mode 100644 packages/ai/src/protocols/openai-responses-channel.ts diff --git a/packages/ai/src/protocols/open-responses-channel.ts b/packages/ai/src/protocols/open-responses-channel.ts index 52f957eb0e9..898b97e4760 100644 --- a/packages/ai/src/protocols/open-responses-channel.ts +++ b/packages/ai/src/protocols/open-responses-channel.ts @@ -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> + 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 = (options: Options): Transport { diff --git a/packages/ai/src/protocols/openai-responses-channel.ts b/packages/ai/src/protocols/openai-responses-channel.ts new file mode 100644 index 00000000000..f8287d7def6 --- /dev/null +++ b/packages/ai/src/protocols/openai-responses-channel.ts @@ -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> + readonly output: ReadonlyArray +} + +export interface DriverInput { + readonly id: string + readonly name: string + readonly request: Readonly> + 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>) => { + const { type: _type, input: _input, previous_response_id: _previousResponseID, ...rest } = request + return rest +} + +const incremental = ( + request: Readonly>, + checkpoint: CheckpointValue, +): ReadonlyArray | 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, + 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 diff --git a/packages/ai/src/protocols/openai-responses.ts b/packages/ai/src/protocols/openai-responses.ts index 278392c215a..5069874cb31 100644 --- a/packages/ai/src/protocols/openai-responses.ts +++ b/packages/ai/src/protocols/openai-responses.ts @@ -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({ 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({ diff --git a/packages/ai/test/provider/openai-responses.test.ts b/packages/ai/test/provider/openai-responses.test.ts index dc01f108ee5..6910dd2a2a2 100644 --- a/packages/ai/test/provider/openai-responses.test.ts +++ b/packages/ai/test/provider/openai-responses.test.ts @@ -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 => { + 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>) => { + 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) diff --git a/packages/core/src/session/model-transport.ts b/packages/core/src/session/model-transport.ts index dfeb6f37ccd..5865ed25ba6 100644 --- a/packages/core/src/session/model-transport.ts +++ b/packages/core/src/session/model-transport.ts @@ -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 } @@ -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(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), }) }, }) diff --git a/packages/core/src/session/runner/llm.ts b/packages/core/src/session/runner/llm.ts index d32568884ad..94ac3baf0e5 100644 --- a/packages/core/src/session/runner/llm.ts +++ b/packages/core/src/session/runner/llm.ts @@ -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() @@ -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 && diff --git a/packages/core/test/session-model-transport.test.ts b/packages/core/test/session-model-transport.test.ts index 1b684173fb1..876fd89dfce 100644 --- a/packages/core/test/session-model-transport.test.ts +++ b/packages/core/test/session-model-transport.test.ts @@ -66,6 +66,15 @@ const collect = (executor: ReturnType, return Array.from(yield* Stream.runCollect(execution.frames)) }).pipe(Effect.scoped) +const collectComplete = ( + executor: ReturnType, + 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 @@ -96,6 +105,164 @@ const automatic = () => { } describe("SessionModelTransport", () => { + test("commits checkpoints only after successful outer completion", async () => { + const messages = queue() + const checkpoints: Array = [] + 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 = [] + 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 = [] + 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() diff --git a/packages/core/test/session-runner.test.ts b/packages/core/test/session-runner.test.ts index e1cf377af0e..43ee1debfc9 100644 --- a/packages/core/test/session-runner.test.ts +++ b/packages/core/test/session-runner.test.ts @@ -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