diff --git a/packages/ai/src/protocols/open-responses-channel.ts b/packages/ai/src/protocols/open-responses-channel.ts index 68ed3f295ae..07f2ec3ff66 100644 --- a/packages/ai/src/protocols/open-responses-channel.ts +++ b/packages/ai/src/protocols/open-responses-channel.ts @@ -18,7 +18,6 @@ const WebSocketResponseCreate = Schema.StructWithRest(Schema.Struct({ type: Sche ]) const decodeMessage = ProviderShared.validateWith(Schema.decodeUnknownEffect(WebSocketResponseCreate)) const encodeMessage = Schema.encodeSync(Schema.fromJsonString(WebSocketResponseCreate)) -const decodeEvent = Schema.decodeUnknownEffect(OpenResponses.protocol.stream.event) export interface Options { readonly id: string @@ -27,6 +26,7 @@ export interface Options { readonly enabled?: (url: string) => boolean readonly url?: (url: string) => string readonly headers?: (headers: Headers.Headers) => Headers.Headers + readonly continuation?: OpenResponsesContinuation.Shape } export interface Prepared { @@ -60,7 +60,7 @@ const driver = (options: Options, body: string): WebSocketChannelDriver => { }), observe: (_create, frame) => Effect.gen(function* () { - const event = yield* decodeEvent(frame).pipe( + const event = yield* OpenResponses.decodeChannelEvent(frame).pipe( Effect.mapError((cause) => ProviderShared.eventError(options.id, `Invalid ${options.name} WebSocket event`, frame, cause), ), @@ -163,6 +163,7 @@ export const transport = (options: Options): Transport } +/** + * Fields to send next to `previous_response_id` on an incremental step, or undefined to send the step in full. + * Whether omitted fields carry over from the continued response is provider behavior the route must know. + */ +export type Shape = (request: Readonly>) => Readonly> | undefined + export interface DriverInput { readonly id: string readonly name: string readonly request: Readonly> readonly message: string readonly base: WebSocketChannelDriver + readonly continuation?: Shape } const checkpointValue = (checkpoint: ChannelCheckpoint | undefined): CheckpointValue | undefined => { @@ -127,22 +133,26 @@ const rejected = ( export const driver = (input: DriverInput): WebSocketChannelDriver => { const { previous_response_id: _previousResponseID, ...request } = input.request + const shape = input.continuation ?? ((fields: Readonly>) => fields) let output: OpenResponses.StreamItem[] = [] 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 } + // Ask the route first: diffing the whole history is wasted when it declines the continuation. + const fields = previous ? shape(request) : undefined + const delta = previous && fields ? incremental(request, previous) : undefined + if (!previous || !fields || !delta) + return { message: ProviderShared.encodeJson(request), mode: "full" as const } return { - message: ProviderShared.encodeJson({ ...request, input: delta, previous_response_id: previous.responseID }), + message: ProviderShared.encodeJson({ ...fields, input: delta, previous_response_id: previous.responseID }), mode: "incremental" as const, } }), observe: (create, frame) => Effect.gen(function* () { - const event = yield* decodeEvent(frame).pipe( + const event = yield* OpenResponses.decodeChannelEvent(frame).pipe( Effect.mapError((cause) => ProviderShared.eventError(input.id, `Invalid ${input.name} WebSocket event`, frame, cause), ), @@ -195,4 +205,4 @@ export const driver = (input: DriverInput): WebSocketChannelDriver => { } } -export const OpenResponsesContinuation = { driver } as const +export * as OpenResponsesContinuation from "./open-responses-continuation.js" diff --git a/packages/ai/src/protocols/open-responses.ts b/packages/ai/src/protocols/open-responses.ts index e9f3628e557..23ca283faa9 100644 --- a/packages/ai/src/protocols/open-responses.ts +++ b/packages/ai/src/protocols/open-responses.ts @@ -405,6 +405,24 @@ export const Event = Schema.StructWithRest( export type Event = Schema.Schema.Type export type NormalizedEvent = Event & { readonly item?: OutputItem | null } +const decodeEventValue = Schema.decodeUnknownEffect(Event) +const decodeFrame = Schema.decodeUnknownEffect(ProviderShared.Json) + +/** + * Decodes one WebSocket frame. xAI answers a rejected `response.create` with `{ "error": { "message", "type" } }` and no + * event type; that envelope reads as an error event so the failure classifies instead of failing decoding. + */ +export const decodeChannelEvent = (frame: string) => + decodeFrame(frame).pipe( + Effect.flatMap((value) => + decodeEventValue( + ProviderShared.isRecord(value) && value.type === undefined && ProviderShared.isRecord(value.error) + ? { ...value, type: "error" } + : value, + ), + ), + ) + export interface ProviderAdapter { readonly id: string readonly name: string diff --git a/packages/ai/src/providers/xai.ts b/packages/ai/src/providers/xai.ts index d639d4d8ae8..43a528325ff 100644 --- a/packages/ai/src/providers/xai.ts +++ b/packages/ai/src/providers/xai.ts @@ -41,6 +41,10 @@ const responsesRoute = Route.make({ id: "openai-responses", name: "xAI Responses", rotateAfterMs: RESPONSES_WEBSOCKET_ROTATE_AFTER_MS, + // xAI continues a chain only from stored responses: with `store: false` (the route default) `previous_response_id` + // fails with "Response with id=… not found", so those steps are sent in full over the reused connection. It also + // rejects `instructions` next to `previous_response_id` and keeps the instructions of the response it continues. + continuation: ({ instructions: _instructions, ...request }) => (request.store === false ? undefined : request), }), defaults: { providerOptions: { store: false, include: ["reasoning.encrypted_content"] } }, }) diff --git a/packages/ai/test/provider/openai-responses.test.ts b/packages/ai/test/provider/openai-responses.test.ts index 7cefd28f483..4027de00d4d 100644 --- a/packages/ai/test/provider/openai-responses.test.ts +++ b/packages/ai/test/provider/openai-responses.test.ts @@ -90,7 +90,11 @@ const classifyingChannelDriver = (message: string): WebSocketChannelDriver => { } } -const continuationDriver = (request: Readonly>, base = baseChannelDriver) => { +const continuationDriver = ( + request: Readonly>, + base = baseChannelDriver, + continuation?: OpenResponsesContinuation.Shape, +) => { const message = ProviderShared.encodeJson(request) return OpenResponsesContinuation.driver({ id: "openai-responses", @@ -98,6 +102,7 @@ const continuationDriver = (request: Readonly>, base = b request, message, base: base(message), + continuation, }) } @@ -921,6 +926,58 @@ describe("OpenAI Responses route", () => { type: "provider-failure", error: { reason: { _tag: "InvalidRequest", classification: "context-overflow" } }, }) + + // A retryable failure stays one: the runner retries it, and the transport has already dropped the + // checkpoint, so that retry is a full send. xAI reports every rejection this way. + const internal = ProviderShared.encodeJson({ + type: "error", + error: { type: "api_error", message: "gRPC error: Response with id=resp_1 not found" }, + }) + expect(yield* second.observe(yield* second.create(saved), internal)).toMatchObject({ + type: "provider-failure", + error: { reason: { _tag: "ProviderInternal" } }, + }) + }), + ) + + it.effect("shapes the incremental send with the route continuation", () => + Effect.gen(function* () { + const firstRequest = { + type: "response.create", + model: "grok-4.6", + store: true, + instructions: "You are terse.", + input: [{ role: "user", content: [{ type: "input_text", text: "First" }] }], + } + const secondRequest = { + ...firstRequest, + input: [...firstRequest.input, { role: "user", content: [{ type: "input_text", text: "Second" }] }], + } + const saved = checkpoint( + yield* continuationDriver(firstRequest).observe( + yield* continuationDriver(firstRequest).create(undefined), + ProviderShared.encodeJson({ type: "response.completed", response: { id: "resp_1" } }), + ), + ) + + const trimmed = yield* continuationDriver( + secondRequest, + baseChannelDriver, + ({ instructions: _, ...rest }) => rest, + ).create(saved) + expect(trimmed.mode).toBe("incremental") + expect(JSON.parse(trimmed.message)).toEqual({ + type: "response.create", + model: "grok-4.6", + store: true, + previous_response_id: "resp_1", + input: [{ role: "user", content: [{ type: "input_text", text: "Second" }] }], + }) + + // Declining the continuation sends the step in full and never sends a previous_response_id. + const declined = yield* continuationDriver(secondRequest, baseChannelDriver, () => undefined).create(saved) + expect(declined.mode).toBe("full") + expect(JSON.parse(declined.message)).toEqual(secondRequest) }), ) diff --git a/packages/ai/test/provider/xai-responses.test.ts b/packages/ai/test/provider/xai-responses.test.ts index a567ba90375..bbf9d92f7d7 100644 --- a/packages/ai/test/provider/xai-responses.test.ts +++ b/packages/ai/test/provider/xai-responses.test.ts @@ -1,11 +1,18 @@ import { describe, expect } from "bun:test" -import { Effect } from "effect" +import { Effect, Layer, Stream } from "effect" import { LLM, LLMEvent, Message } from "../../src/index.js" import { XAI } from "../../src/providers.js" import { OpenResponses } from "../../src/protocols/open-responses.js" import { OpenAIResponses } from "../../src/protocols/openai-responses.js" +import * as ProviderShared from "../../src/protocols/shared.js" import { XAIResponses } from "../../src/protocols/xai-responses.js" -import { LLMClient } from "../../src/route.js" +import { + LLMClient, + RequestExecutor, + WebSocketTransport, + type ChannelCheckpoint, + type WebSocketChannelDriver, +} from "../../src/route.js" import { compileRequest } from "../../src/route/client.js" import { it } from "../lib/effect.js" import { fixedResponse } from "../lib/http.js" @@ -13,6 +20,35 @@ import { sseEvents } from "../lib/sse.js" const model = XAI.configure({ apiKey: "test", baseURL: "https://api.x.ai/v1" }).responses("grok-4.6") +/** Runs a request through the WebSocket transport and hands back its channel driver; the HTTP fallback answers. */ +const channelDriver = (request: ReturnType) => + Effect.gen(function* () { + let driver: WebSocketChannelDriver | undefined + yield* LLMClient.generate(request, { + webSocket: { + execute: (exchange) => + Effect.sync(() => { + driver = exchange.driver + return { frames: exchange.fallback(), complete: Effect.void } + }), + }, + }).pipe(Effect.provide(fixedResponse(sseEvents({ type: "response.completed", response: { id: "http" } })))) + if (!driver) throw new Error("Expected a WebSocket channel driver") + return driver + }) + +const completed = (driver: WebSocketChannelDriver, id: string) => + Effect.gen(function* () { + const create = yield* driver.create(undefined) + yield* driver.observe(create, ProviderShared.encodeJson({ type: "response.created", response: { id } })) + const observation = yield* driver.observe( + create, + ProviderShared.encodeJson({ type: "response.completed", response: { id } }), + ) + if (observation.type !== "completed" || !observation.checkpoint) throw new Error("Expected a checkpoint") + return observation.checkpoint + }) + describe("xAI Responses route", () => { it.effect("composes the Open Responses baseline with xAI extensions", () => Effect.gen(function* () { @@ -162,6 +198,78 @@ describe("xAI Responses route", () => { }), ) + it.effect("classifies xAI's untyped WebSocket error envelope", () => + Effect.gen(function* () { + // xAI answers a rejected response.create with an error envelope that carries no event type. + const envelope = ProviderShared.encodeJson({ + error: { + message: + 'Request validation error: {"code":"400","error":"Argument not supported: instructions and previous_response_id together"}', + type: "api_error", + }, + }) + const webSocket = WebSocketTransport.makeDirect({ + open: () => + Effect.succeed({ sendText: () => Effect.void, messages: Stream.make(envelope), close: Effect.void }), + }) + const error = yield* LLMClient.generate(LLM.request({ model, prompt: "Hello" }), { webSocket }).pipe( + Effect.provide( + LLMClient.layer.pipe( + Layer.provide( + Layer.succeed( + RequestExecutor.Service, + RequestExecutor.Service.of({ execute: () => Effect.die("unexpected HTTP request") }), + ), + ), + ), + ), + Effect.flip, + ) + + expect(error.reason._tag).toBe("ProviderInternal") + expect(error.message).toContain("Argument not supported: instructions and previous_response_id together") + expect(error.reason.body).toBe(envelope) + }), + ) + + it.effect("continues stored responses without instructions and sends unstored steps in full", () => + Effect.gen(function* () { + const step = (store: boolean, ...prompts: string[]) => + LLM.request({ + model, + system: "You are terse.", + messages: prompts.map((prompt) => Message.user(prompt)), + providerOptions: { store }, + }) + const send = (store: boolean, checkpoint: ChannelCheckpoint) => + channelDriver(step(store, "First", "Second")).pipe(Effect.flatMap((driver) => driver.create(checkpoint))) + + const stored = yield* send(true, yield* completed(yield* channelDriver(step(true, "First")), "resp_1")) + expect(stored.mode).toBe("incremental") + expect(JSON.parse(stored.message)).toEqual({ + type: "response.create", + model: "grok-4.6", + store: true, + include: ["reasoning.encrypted_content"], + previous_response_id: "resp_1", + input: [{ role: "user", content: [{ type: "input_text", text: "Second" }] }], + }) + + // The connection cache only serves stored responses, so the default store: false never chains. + const unstored = yield* send(false, yield* completed(yield* channelDriver(step(false, "First")), "resp_1")) + expect(unstored.mode).toBe("full") + expect(JSON.parse(unstored.message)).toMatchObject({ + instructions: "You are terse.", + store: false, + input: [ + { role: "user", content: [{ type: "input_text", text: "First" }] }, + { role: "user", content: [{ type: "input_text", text: "Second" }] }, + ], + }) + expect(JSON.parse(unstored.message).previous_response_id).toBeUndefined() + }), + ) + it.effect("parses xAI hosted tool items", () => Effect.gen(function* () { const item = { type: "x_search_call", id: "x_search_1", status: "completed", action: { query: "news" } } diff --git a/packages/core/src/plugin/provider/xai.ts b/packages/core/src/plugin/provider/xai.ts index 99d21c56ace..b5da131526a 100644 --- a/packages/core/src/plugin/provider/xai.ts +++ b/packages/core/src/plugin/provider/xai.ts @@ -101,6 +101,7 @@ export const XAIPlugin = define({ for (const model of provider.models.values()) { catalog.model.update(providerID, model.id, (draft) => { draft.capabilities.responsesWebsockets = true + draft.websocket = true }) } }) diff --git a/packages/core/test/plugin/provider-xai.test.ts b/packages/core/test/plugin/provider-xai.test.ts index e420771b2a5..88cf1dd3357 100644 --- a/packages/core/test/plugin/provider-xai.test.ts +++ b/packages/core/test/plugin/provider-xai.test.ts @@ -69,7 +69,7 @@ describe("XAIPlugin", () => { }), ) - it.effect("keeps xAI Responses WebSockets opt-in", () => + it.effect("enables xAI Responses WebSockets", () => Effect.gen(function* () { const catalog = yield* Catalog.Service const providerID = Provider.ID.make("xai") @@ -84,7 +84,7 @@ describe("XAIPlugin", () => { const model = yield* catalog.model.get(providerID, Model.ID.make("grok-4.6")) expect(model?.capabilities.responsesWebsockets).toBe(true) - expect(model?.websocket).toBeUndefined() + expect(model?.websocket).toBe(true) }), ) }) diff --git a/services/www/src/docs/content/providers.mdx b/services/www/src/docs/content/providers.mdx index 0dcdf714eec..acee8ad62d6 100644 --- a/services/www/src/docs/content/providers.mdx +++ b/services/www/src/docs/content/providers.mdx @@ -110,10 +110,11 @@ role for other Foundry models. If a request fails because the token belongs to a ## WebSocket transport -OpenAI and supported Azure Responses models keep one WebSocket connection open per session and send each step over it -instead of opening a new HTTP request. While the request prefix is unchanged, consecutive steps only transmit what was -added since the previous response, which cuts upload volume on long sessions. Provider compaction runs over the same -connection. +OpenAI, xAI, and supported Azure Responses models keep one WebSocket connection open per session and send each step +over it instead of opening a new HTTP request. While the request prefix is unchanged, consecutive steps only transmit +what was added since the previous response, which cuts upload volume on long sessions. OpenAI provider compaction runs +over the same connection. xAI continues a chain only from stored responses, so with its default `store: false` each step +is sent in full over the reused connection. The connection is transparent. When the provider closes the socket, the next step reconnects; when a connection cannot be opened at all, the session continues over HTTP. Plugins that register `http.request` or `http.response` hooks for a