feat(ai): enable xai responses websocket (#43889)
This commit is contained in:
@@ -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({
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
})
|
||||
}
|
||||
})
|
||||
}),
|
||||
})
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user