From 2636797c654888bc56dec7a1adb9f6653800a9aa Mon Sep 17 00:00:00 2001 From: Shoubhit Dash Date: Fri, 21 Aug 2026 19:58:01 +0530 Subject: [PATCH] feat(ai): enable xai responses websocket (#43889) --- packages/ai/src/protocols/openai-responses.ts | 10 ++++-- packages/ai/src/providers/xai.ts | 8 ++++- .../ai/test/provider/openai-responses.test.ts | 34 +++++++++++++++++++ packages/core/src/plugin/provider/xai.ts | 11 ++++++ packages/core/src/session/model-request.ts | 14 +++++--- .../core/test/plugin/provider-openai.test.ts | 13 +++++-- .../core/test/plugin/provider-xai.test.ts | 22 ++++++++++++ 7 files changed, 102 insertions(+), 10 deletions(-) diff --git a/packages/ai/src/protocols/openai-responses.ts b/packages/ai/src/protocols/openai-responses.ts index 915c5dfd85..ee96e9abfe 100644 --- a/packages/ai/src/protocols/openai-responses.ts +++ b/packages/ai/src/protocols/openai-responses.ts @@ -11,7 +11,7 @@ import { optionalArray, ProviderShared } from "./shared.js" 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 { OpenResponsesChannel, type Options } from "./open-responses-channel.js" import { OpenAIResponsesChannel } from "./openai-responses-channel.js" const ADAPTER = "openai-responses" @@ -247,12 +247,16 @@ const endpoint = Endpoint.path(PATH, { baseURL: DEFAULT_BAS const auth = Auth.none export const httpTransport = HttpTransport.sseJson.with() -export const transport = OpenResponsesChannel.transport({ +export const channelTransport = (options: Omit) => + OpenResponsesChannel.transport({ + ...options, + driver: (input) => OpenAIResponsesChannel.driver({ id: options.id, name: options.name, ...input }), + }) +export const transport = channelTransport({ id: ADAPTER, 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/src/providers/xai.ts b/packages/ai/src/providers/xai.ts index 47b55ef1bd..65b32ab521 100644 --- a/packages/ai/src/providers/xai.ts +++ b/packages/ai/src/providers/xai.ts @@ -28,13 +28,19 @@ export interface Settings extends ProviderPackage.Settings { export type { XAIImageOptions } from "../protocols/xai-images.js" +const RESPONSES_WEBSOCKET_ROTATE_AFTER_MS = 24 * 60 * 1000 + const responsesRoute = Route.make({ id: "openai-responses", provider: id, providerMetadataKey: "xai", protocol: OpenAIResponses.protocol, endpoint: Endpoint.path("/responses", { baseURL: OpenAICompatibleProfiles.profiles.xai.baseURL }), - transport: OpenAIResponses.httpTransport, + transport: OpenAIResponses.channelTransport({ + id: "openai-responses", + name: "xAI Responses", + rotateAfterMs: RESPONSES_WEBSOCKET_ROTATE_AFTER_MS, + }), defaults: { providerOptions: { store: false } }, }) diff --git a/packages/ai/test/provider/openai-responses.test.ts b/packages/ai/test/provider/openai-responses.test.ts index 4336e5ac4e..fa4d090e4b 100644 --- a/packages/ai/test/provider/openai-responses.test.ts +++ b/packages/ai/test/provider/openai-responses.test.ts @@ -691,6 +691,40 @@ describe("OpenAI Responses route", () => { }), ) + it.effect("builds xAI WebSocket requests without OpenAI handshake headers", () => + Effect.gen(function* () { + const deps = Layer.succeed( + RequestExecutor.Service, + RequestExecutor.Service.of({ execute: () => Effect.die("unexpected HTTP request") }), + ) + const response = yield* LLMClient.generate(LLM.request({ model: xaiModel, prompt: "Say hello." }), { + webSocket: { + execute: (exchange) => + Effect.gen(function* () { + expect(exchange.connect.url).toBe("wss://api.x.ai/v1/responses") + expect(exchange.connect.rotateAfterMs).toBe(24 * 60 * 1000) + expect(exchange.connect.headers.authorization).toBe("Bearer test") + expect(exchange.connect.headers["openai-beta"]).toBeUndefined() + expect(JSON.parse((yield* exchange.driver.create(undefined)).message)).toMatchObject({ + type: "response.create", + model: "grok-4.5", + store: false, + }) + return { + frames: Stream.make( + JSON.stringify({ type: "response.created", response: { id: "resp_xai" } }), + JSON.stringify({ type: "response.completed", response: { id: "resp_xai" } }), + ), + complete: Effect.void, + } + }), + }, + }).pipe(Effect.provide(LLMClient.layer.pipe(Layer.provide(deps)))) + + expect(response.finishReason.normalized).toBe("stop") + }), + ) + it.effect("uses exactly one HTTP request when no WebSocket executor is supplied", () => Effect.gen(function* () { const attempts = yield* Ref.make(0) diff --git a/packages/core/src/plugin/provider/xai.ts b/packages/core/src/plugin/provider/xai.ts index abf3bcfb3c..9ed1168f3a 100644 --- a/packages/core/src/plugin/provider/xai.ts +++ b/packages/core/src/plugin/provider/xai.ts @@ -4,6 +4,7 @@ import { Clock, Effect, Option, Schema } from "effect" import { App } from "../../app.js" import { Credential } from "../../credential.js" import { Integration } from "../../integration.js" +import { Provider } from "../../provider.js" const clientID = "b1a00492-073a-47ea-816f-4c329264a828" const issuer = "https://auth.x.ai/oauth2" @@ -12,6 +13,7 @@ const scope = "openid profile email offline_access grok-cli:access api:access" const pollingSafetyMargin = 3000 const browserMethodID = Integration.MethodID.make("browser") const deviceMethodID = Integration.MethodID.make("device") +const providerID = Provider.ID.make("xai") const Token = Schema.Struct({ access_token: Schema.String, @@ -93,6 +95,15 @@ export const XAIPlugin = define({ draft.method.update(device(ctx.app)) draft.method.update({ integrationID: "xai", method: { type: "key", label: "Manually enter API Key" } }) }) + yield* ctx.catalog.transform((catalog) => { + const provider = catalog.provider.get(providerID) + if (!provider) return + for (const model of provider.models.values()) { + catalog.model.update(providerID, model.id, (draft) => { + draft.capabilities.responsesWebsockets = true + }) + } + }) }), }) diff --git a/packages/core/src/session/model-request.ts b/packages/core/src/session/model-request.ts index 678ca9f234..722ac60231 100644 --- a/packages/core/src/session/model-request.ts +++ b/packages/core/src/session/model-request.ts @@ -28,6 +28,9 @@ const IMAGE_BYTES_TARGET = 15 * 1024 * 1024 // 15 MiB const IMAGE_REMOVED = "[This image was removed to reduce the request size and is no longer visible. Do not make claims about its contents from memory. If needed, retrieve it again with an available tool or ask the user to attach it again.]" +const responsesWebSocketFlag = (providerID: string) => + `OPENCODE_EXPERIMENTAL_${providerID.replace(/[^a-zA-Z0-9]+/g, "_").toUpperCase()}_RESPONSES_WEBSOCKET` + /** Failures a prepared execution can surface: infrastructure errors plus user declines resurfaced from the defect tunnel. */ export type ExecuteError = Tool.Error | Permission.DeclinedError | QuestionTool.CancelledError @@ -207,10 +210,6 @@ export const layer = Layer.effect( 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 prepare = Effect.fn("SessionModelRequest.prepare")(function* (input: PrepareInput) { const session = input.scope.session const resolved = input.scope.model @@ -270,6 +269,13 @@ export const layer = Layer.effect( const webSocketEligible = !(yield* hooks.has("session", "http.request", resolved.ref.providerID)) && !(yield* hooks.has("session", "http.response", resolved.ref.providerID)) + const webSocket = + resolved.capabilities.responsesWebsockets === true + ? yield* Config.boolean(responsesWebSocketFlag(resolved.ref.providerID)).pipe( + Config.withDefault(false), + Effect.orDie, + ) + : false const http = webSocketEligible ? undefined : SessionModelHttp.middleware(hooks, { diff --git a/packages/core/test/plugin/provider-openai.test.ts b/packages/core/test/plugin/provider-openai.test.ts index f0a60277ce..d44ec1cb3d 100644 --- a/packages/core/test/plugin/provider-openai.test.ts +++ b/packages/core/test/plugin/provider-openai.test.ts @@ -253,6 +253,16 @@ describe("OpenAIPlugin", () => { }).pipe( Effect.provide(SessionModelRequest.layer), Effect.provideService(SessionModelTransport.Service, transport), + ) + + const prepared = yield* program.pipe( + Effect.provide( + ConfigProvider.layer( + ConfigProvider.fromEnv({ env: { OPENCODE_EXPERIMENTAL_DEPLOYMENT_RESPONSES_WEBSOCKET: "true" } }), + ), + ), + ) + const otherProvider = yield* program.pipe( Effect.provide( ConfigProvider.layer( ConfigProvider.fromEnv({ env: { OPENCODE_EXPERIMENTAL_OPENAI_RESPONSES_WEBSOCKET: "true" } }), @@ -260,10 +270,9 @@ describe("OpenAIPlugin", () => { ), ) - const prepared = yield* program - expect(prepared.options.webSocket).toBe(executor) expect(prepared.options.http).toBeUndefined() + expect(otherProvider.options.webSocket).toBeUndefined() }), ) }) diff --git a/packages/core/test/plugin/provider-xai.test.ts b/packages/core/test/plugin/provider-xai.test.ts index 7eb4a20a30..c2302cb30e 100644 --- a/packages/core/test/plugin/provider-xai.test.ts +++ b/packages/core/test/plugin/provider-xai.test.ts @@ -1,10 +1,13 @@ import { describe, expect } from "bun:test" import { Effect } from "effect" +import { Catalog } from "@opencode-ai/core/catalog" import { Credential } from "@opencode-ai/core/credential" import { Integration } from "@opencode-ai/core/integration" import { Plugin } from "@opencode-ai/core/plugin" import { PluginHost } from "@opencode-ai/core/plugin/host" import { XAIPlugin } from "@opencode-ai/core/plugin/provider/xai" +import { Model } from "@opencode-ai/core/model" +import { Provider } from "@opencode-ai/core/provider" import { testEffect } from "../lib/effect" import { PluginTestLayer } from "./fixture" @@ -65,4 +68,23 @@ describe("XAIPlugin", () => { }) }), ) + + it.effect("marks xAI deployments as Responses WebSocket capable", () => + Effect.gen(function* () { + const catalog = yield* Catalog.Service + const providerID = Provider.ID.make("xai") + yield* catalog.transform((draft) => { + draft.provider.update(providerID, (provider) => { + provider.package = Provider.aisdk("@ai-sdk/xai") + }) + draft.model.update(providerID, Model.ID.make("grok-4.6"), () => {}) + }) + + yield* addPlugin() + + expect((yield* catalog.model.get(providerID, Model.ID.make("grok-4.6")))?.capabilities.responsesWebsockets).toBe( + true, + ) + }), + ) })