feat(ai): enable xai responses websocket (#43889)

This commit is contained in:
Shoubhit Dash
2026-08-21 19:58:01 +05:30
committed by GitHub
parent e673807e39
commit 2636797c65
7 changed files with 102 additions and 10 deletions
@@ -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<OpenAIResponsesBody>(PATH, { baseURL: DEFAULT_BAS
const auth = Auth.none
export const httpTransport = HttpTransport.sseJson.with<OpenAIResponsesBody>()
export const transport = OpenResponsesChannel.transport<OpenAIResponsesBody>({
export const channelTransport = (options: Omit<Options, "driver">) =>
OpenResponsesChannel.transport<OpenAIResponsesBody>({
...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({
+7 -1
View File
@@ -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 } },
})
@@ -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)
+11
View File
@@ -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
})
}
})
}),
})
+10 -4
View File
@@ -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, {
@@ -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()
}),
)
})
@@ -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,
)
}),
)
})