Skip to content

Commit cb39ea1

Browse files
nexxelnnexxeln
authored
fix(core): prefer websocket with request hooks (#43234)
1 parent ff9452b commit cb39ea1

21 files changed

Lines changed: 389 additions & 127 deletions

‎.changeset/websocket-request-hooks.md‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
---
2+
"@opencode-ai/plugin": patch
3+
"@opencode-ai/core": patch
4+
---
5+
6+
Add transport-neutral Session model request hooks and provider-scoped hook registration so eligible OpenAI Responses requests can prefer WebSocket without bypassing HTTP-only middleware.

‎packages/core/src/plugin/hooks.ts‎

Lines changed: 43 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import type { AISDKHooks } from "@opencode-ai/plugin/effect/aisdk"
44
import type { SessionHooks } from "@opencode-ai/plugin/effect/session"
55
import type { ShellHooks } from "@opencode-ai/plugin/effect/shell"
66
import type { ToolFailures, ToolHooks } from "@opencode-ai/plugin/effect/tool"
7+
import type { ModelHookOptions } from "@opencode-ai/plugin/effect/registration"
78
import { Context, Effect, Layer, Scope } from "effect"
89
import { makeLocationNode } from "@opencode-ai/util/effect/app-node"
910
import { State } from "../state.js"
@@ -26,16 +27,26 @@ interface Failures extends Record<keyof Domains, unknown> {
2627
}
2728

2829
type Callback<Event, Error> = (event: Event) => Effect.Effect<void, Error>
30+
type Entry = { readonly callback: Function; readonly options?: ModelHookOptions }
31+
32+
const eventProviderID = (event: unknown) => {
33+
if (typeof event !== "object" || event === null || !("model" in event)) return undefined
34+
const model = event.model
35+
if (typeof model !== "object" || model === null || !("providerID" in model)) return undefined
36+
return typeof model.providerID === "string" ? model.providerID : undefined
37+
}
2938

3039
export interface Interface {
3140
readonly has: <Domain extends keyof Domains>(
3241
domain: Domain,
3342
name: keyof Domains[Domain] & keyof Failures[Domain],
43+
providerID?: string,
3444
) => Effect.Effect<boolean>
3545
readonly register: <Domain extends keyof Domains, Name extends keyof Domains[Domain] & keyof Failures[Domain]>(
3646
domain: Domain,
3747
name: Name,
3848
callback: Callback<Domains[Domain][Name], Failures[Domain][Name]>,
49+
options?: ModelHookOptions,
3950
) => Effect.Effect<State.Registration, never, Scope.Scope>
4051
readonly trigger: <Domain extends keyof Domains, Name extends keyof Domains[Domain] & keyof Failures[Domain]>(
4152
domain: Domain,
@@ -49,36 +60,47 @@ export class Service extends Context.Service<Service, Interface>()("@opencode/Pl
4960
const layer = Layer.effect(
5061
Service,
5162
Effect.gen(function* () {
52-
const callbacks = new Map<string, Function[]>()
63+
const callbacks = new Map<string, Entry[]>()
5364
const key = (domain: keyof Domains, name: PropertyKey) => `${domain}.${String(name)}`
5465

55-
const register: Interface["register"] = Effect.fn("PluginHooks.register")(function* (domain, name, callback) {
56-
const scope = yield* Scope.Scope
57-
const id = key(domain, name)
58-
let active = true
59-
callbacks.set(id, [...(callbacks.get(id) ?? []), callback])
60-
const dispose = Effect.sync(() => {
61-
if (!active) return
62-
active = false
63-
const next = (callbacks.get(id) ?? []).filter((item) => item !== callback)
64-
if (next.length === 0) callbacks.delete(id)
65-
else callbacks.set(id, next)
66-
})
67-
yield* Scope.addFinalizer(scope, dispose)
68-
return { dispose }
69-
})
66+
const register: Interface["register"] = Effect.fn("PluginHooks.register")(
67+
function* (domain, name, callback, options) {
68+
const scope = yield* Scope.Scope
69+
const id = key(domain, name)
70+
let active = true
71+
const entry = { callback, options }
72+
callbacks.set(id, [...(callbacks.get(id) ?? []), entry])
73+
const dispose = Effect.sync(() => {
74+
if (!active) return
75+
active = false
76+
const next = (callbacks.get(id) ?? []).filter((item) => item !== entry)
77+
if (next.length === 0) callbacks.delete(id)
78+
else callbacks.set(id, next)
79+
})
80+
yield* Scope.addFinalizer(scope, dispose)
81+
return { dispose }
82+
},
83+
)
7084

7185
const trigger: Interface["trigger"] = Effect.fnUntraced(function* (domain, name, event) {
72-
for (const callback of callbacks.get(key(domain, name)) ?? []) {
73-
const result: Effect.Effect<void, Failures[typeof domain][typeof name]> = Reflect.apply(callback, undefined, [
74-
event,
75-
])
86+
for (const entry of callbacks.get(key(domain, name)) ?? []) {
87+
if (entry.options?.providerID !== undefined && entry.options.providerID !== eventProviderID(event)) continue
88+
const result: Effect.Effect<void, Failures[typeof domain][typeof name]> = Reflect.apply(
89+
entry.callback,
90+
undefined,
91+
[event],
92+
)
7693
yield* result
7794
}
7895
return event
7996
})
8097

81-
const has: Interface["has"] = (domain, name) => Effect.sync(() => callbacks.has(key(domain, name)))
98+
const has: Interface["has"] = (domain, name, providerID) =>
99+
Effect.sync(() =>
100+
(callbacks.get(key(domain, name)) ?? []).some(
101+
(entry) => entry.options?.providerID === undefined || entry.options.providerID === providerID,
102+
),
103+
)
82104

83105
return Service.of({ has, register, trigger })
84106
}),

‎packages/core/src/plugin/host.ts‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -104,9 +104,10 @@ export const make = Effect.fn("PluginHost.make")(function* (plugin: import("../p
104104
}),
105105
},
106106
aisdk: {
107-
hook: (name, callback) => {
107+
hook: (name, callback, options) => {
108108
if (name === "sdk") {
109109
return aisdk.hook.sdk((event) => {
110+
if (options?.providerID !== undefined && options.providerID !== event.model.providerID) return Effect.void
110111
const output = {
111112
model: mutable(event.model),
112113
package: event.package,
@@ -119,6 +120,7 @@ export const make = Effect.fn("PluginHost.make")(function* (plugin: import("../p
119120
})
120121
}
121122
return aisdk.hook.language((event) => {
123+
if (options?.providerID !== undefined && options.providerID !== event.model.providerID) return Effect.void
122124
const output = {
123125
model: mutable(event.model),
124126
options: event.options,
@@ -382,7 +384,7 @@ export const make = Effect.fn("PluginHost.make")(function* (plugin: import("../p
382384
}),
383385
},
384386
session: {
385-
hook: (name, callback) => hooks.register("session", name, callback),
387+
hook: (name, callback, options) => hooks.register("session", name, callback, options),
386388
create: (input) =>
387389
runtime.session.create({
388390
id: input?.id,

‎packages/core/src/plugin/provider/github-copilot.ts‎

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -241,19 +241,22 @@ export const GithubCopilotPlugin = define({
241241
evt.sdk = mod.createOpenaiCompatible(evt.options)
242242
}),
243243
)
244-
yield* ctx.session.hook("http.request", (evt) =>
245-
Effect.gen(function* () {
246-
if (evt.model.providerID !== Provider.ID.githubCopilot) return
247-
if (evt.agent === Agent.ID.make("title"))
248-
evt.request.headers.set("X-Interaction-Type", "conversation-background")
249-
if (evt.agent === Agent.ID.make("compaction"))
250-
evt.request.headers.set("X-Interaction-Type", "conversation-compaction")
251-
const token = evt.request.headers.get("x-api-key")
252-
if (!token) return
253-
const text = yield* Effect.promise(() => evt.request.clone().text())
254-
const body = Option.getOrUndefined(decodeBody(text))
255-
applyHeaders(evt.request.headers, token, ctx.app, requestMetadata(evt.request.url, body), true)
256-
}),
244+
yield* ctx.session.hook(
245+
"http.request",
246+
(evt) =>
247+
Effect.gen(function* () {
248+
if (evt.model.providerID !== Provider.ID.githubCopilot) return
249+
if (evt.agent === Agent.ID.make("title"))
250+
evt.request.headers.set("X-Interaction-Type", "conversation-background")
251+
if (evt.agent === Agent.ID.make("compaction"))
252+
evt.request.headers.set("X-Interaction-Type", "conversation-compaction")
253+
const token = evt.request.headers.get("x-api-key")
254+
if (!token) return
255+
const text = yield* Effect.promise(() => evt.request.clone().text())
256+
const body = Option.getOrUndefined(decodeBody(text))
257+
applyHeaders(evt.request.headers, token, ctx.app, requestMetadata(evt.request.url, body), true)
258+
}),
259+
{ providerID: Provider.ID.githubCopilot },
257260
)
258261
yield* ctx.aisdk.hook(
259262
"language",

‎packages/core/src/plugin/provider/openai.ts‎

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ import { App } from "../../app.js"
55
import { Credential } from "../../credential.js"
66
import { Bus } from "../../bus.js"
77
import { Integration } from "../../integration.js"
8-
import { Model } from "../../model.js"
98
import { OauthCallbackPage } from "../../oauth/page.js"
109
import { Provider } from "../../provider.js"
1110
import type { PluginInternal } from "../internal.js"
@@ -230,15 +229,17 @@ export const OpenAIPlugin = define({
230229
})
231230
}
232231
})
233-
yield* ctx.session.hook("http.request", (evt) =>
234-
Effect.sync(() => {
235-
if (!chatgpt || evt.model.providerID !== Provider.ID.openai) return
236-
const url = new URL(evt.request.url)
237-
evt.request.headers.set("originator", "opencode")
238-
evt.request.headers.set("session-id", evt.sessionID)
239-
if (url.origin !== "https://api.openai.com") return
240-
evt.request = new Request(`${codexBaseURL}${url.pathname.replace(/^\/v1/, "")}${url.search}`, evt.request)
241-
}),
232+
yield* ctx.session.hook(
233+
"model.request",
234+
(evt) =>
235+
Effect.sync(() => {
236+
if (!chatgpt) return
237+
if (evt.baseURL && URL.canParse(evt.baseURL) && new URL(evt.baseURL).origin === "https://api.openai.com")
238+
evt.baseURL = codexBaseURL
239+
evt.headers.originator = "opencode"
240+
evt.headers["session-id"] = evt.sessionID
241+
}),
242+
{ providerID: Provider.ID.openai },
242243
)
243244
const refresh = () => loading.withPermit(load().pipe(Effect.andThen(ctx.catalog.reload())))
244245
yield* bus.subscribe(Integration.Event.ConnectionUpdated).pipe(

‎packages/core/src/session/compaction.ts‎

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import { llmClient } from "../effect/app-node-platform.js"
1212
import { SessionEvent } from "./event.js"
1313
import type { SessionMessage } from "./message.js"
1414
import { SessionModelHeaders } from "./model-headers.js"
15+
import { SessionModelHook } from "./model-hook.js"
1516
import { SessionModelHttp } from "./model-http.js"
1617
import { SessionPromptCacheKey } from "./prompt-cache-key.js"
1718
import { App } from "../app.js"
@@ -270,23 +271,25 @@ const make = (dependencies: Dependencies) => {
270271
})
271272
: Effect.void,
272273
)
274+
const request = yield* SessionModelHook.apply(
275+
dependencies.hooks,
276+
{ sessionID: plan.session.id, agent: Agent.ID.make("compaction"), model: plan.ref },
277+
LLM.request({
278+
model: plan.model,
279+
promptCacheKey: SessionPromptCacheKey.make(plan.session.id),
280+
http: { headers: SessionModelHeaders.make(plan.session, dependencies.app) },
281+
messages: [Message.user(plan.prompt)],
282+
tools: [],
283+
}),
284+
)
273285
yield* dependencies.llm
274-
.stream(
275-
LLM.request({
276-
model: plan.model,
277-
promptCacheKey: SessionPromptCacheKey.make(plan.session.id),
278-
http: { headers: SessionModelHeaders.make(plan.session, dependencies.app) },
279-
messages: [Message.user(plan.prompt)],
280-
tools: [],
286+
.stream(request, {
287+
http: SessionModelHttp.middleware(dependencies.hooks, {
288+
sessionID: plan.session.id,
289+
agent: Agent.ID.make("compaction"),
290+
model: plan.ref,
281291
}),
282-
{
283-
http: SessionModelHttp.middleware(dependencies.hooks, {
284-
sessionID: plan.session.id,
285-
agent: Agent.ID.make("compaction"),
286-
model: plan.ref,
287-
}),
288-
},
289-
)
292+
})
290293
.pipe(
291294
Stream.runForEach((event) => {
292295
if (LLMEvent.is.providerError(event))

‎packages/core/src/session/generate-node.ts‎

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import { SessionContext } from "./context.js"
1111
import { SessionGenerate } from "./generate.js"
1212
import { SessionHistory } from "./history.js"
1313
import { SessionModelHeaders } from "./model-headers.js"
14+
import { SessionModelHook } from "./model-hook.js"
1415
import { SessionModelHttp } from "./model-http.js"
1516
import { SessionPromptCacheKey } from "./prompt-cache-key.js"
1617
import { SessionRunnerModel } from "./runner/model.js"
@@ -71,7 +72,9 @@ export const layer = Layer.effect(
7172
providerID: model.ref.providerID,
7273
modelID: model.ref.id,
7374
})
74-
const response = yield* llm.generate(
75+
const request = yield* SessionModelHook.apply(
76+
hooks,
77+
{ sessionID: selection.session.id, agent: selection.agent.id, model: model.ref },
7578
LLM.request({
7679
model: model.model,
7780
http: { headers: SessionModelHeaders.make(selection.session, app) },
@@ -80,14 +83,14 @@ export const layer = Layer.effect(
8083
messages: contextEvent.messages,
8184
tools: hookedTools,
8285
}),
83-
{
84-
http: SessionModelHttp.middleware(hooks, {
85-
sessionID: selection.session.id,
86-
agent: selection.agent.id,
87-
model: model.ref,
88-
}),
89-
},
9086
)
87+
const response = yield* llm.generate(request, {
88+
http: SessionModelHttp.middleware(hooks, {
89+
sessionID: selection.session.id,
90+
agent: selection.agent.id,
91+
model: model.ref,
92+
}),
93+
})
9194
yield* Effect.logInfo("session generation usage diagnostic", { usage: response.usage })
9295
return response.text
9396
}),

‎packages/core/src/session/model-hook.ts‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
export * as SessionModelHook from "./model-hook.js"
2+
3+
import { HttpOptions, LanguageModel, LLMRequest } from "@opencode-ai/ai"
4+
import type { Agent } from "@opencode-ai/schema/agent"
5+
import type { Model } from "@opencode-ai/schema/model"
6+
import type { Session } from "@opencode-ai/schema/session"
7+
import { Effect } from "effect"
8+
import { PluginHooks } from "../plugin/hooks.js"
9+
10+
export const apply = (
11+
hooks: PluginHooks.Interface,
12+
input: { readonly sessionID: Session.ID; readonly agent: Agent.ID; readonly model: Model.Ref },
13+
request: LLMRequest,
14+
) =>
15+
Effect.gen(function* () {
16+
const currentBaseURL = request.model.route.endpoint.baseURL
17+
const event = yield* hooks.trigger("session", "model.request", {
18+
...input,
19+
baseURL: typeof currentBaseURL === "string" ? currentBaseURL : undefined,
20+
headers: { ...request.http?.headers },
21+
})
22+
const route =
23+
event.baseURL !== undefined && event.baseURL !== currentBaseURL
24+
? request.model.route.with({ endpoint: { baseURL: event.baseURL } })
25+
: request.model.route
26+
return LLMRequest.update(request, {
27+
model: route === request.model.route ? request.model : LanguageModel.update(request.model, { route }),
28+
http: new HttpOptions({
29+
body: request.http?.body,
30+
headers: Object.keys(event.headers).length === 0 ? undefined : event.headers,
31+
query: request.http?.query,
32+
}),
33+
})
34+
})

‎packages/core/src/session/model-request.ts‎

Lines changed: 20 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import { QuestionTool } from "../tool/plugin/question.js"
1414
import { Tool } from "../tool.js"
1515
import { SessionContext } from "./context.js"
1616
import { SessionModelHeaders } from "./model-headers.js"
17+
import { SessionModelHook } from "./model-hook.js"
1718
import { SessionModelHttp } from "./model-http.js"
1819
import { SessionModelTransport } from "./model-transport.js"
1920
import { SessionPromptCacheKey } from "./prompt-cache-key.js"
@@ -226,20 +227,25 @@ export const layer = Layer.effect(
226227
return [[name, { ...tool, description: definition.description, inputSchema: definition.input }] as const]
227228
}),
228229
)
229-
const request = LLM.request({
230-
model,
231-
http: {
232-
headers: SessionModelHeaders.make(session, app),
233-
},
234-
// TODO: Persist cache lineage so nested forks reuse the root session's cache key.
235-
promptCacheKey: SessionPromptCacheKey.make(session.fork?.sessionID ?? session.id),
236-
system: context.system,
237-
messages: boundImages(unsupportedParts(context.messages, resolved.capabilities)),
238-
tools: Array.from(hooked, ([name, tool]) => ({ ...tool, name })),
239-
toolChoice: stepLimitReached ? "none" : undefined,
240-
})
230+
const request = yield* SessionModelHook.apply(
231+
hooks,
232+
{ sessionID: session.id, agent: agent.id, model: resolved.ref },
233+
LLM.request({
234+
model,
235+
http: {
236+
headers: SessionModelHeaders.make(session, app),
237+
},
238+
// TODO: Persist cache lineage so nested forks reuse the root session's cache key.
239+
promptCacheKey: SessionPromptCacheKey.make(session.fork?.sessionID ?? session.id),
240+
system: context.system,
241+
messages: boundImages(unsupportedParts(context.messages, resolved.capabilities)),
242+
tools: Array.from(hooked, ([name, tool]) => ({ ...tool, name })),
243+
toolChoice: stepLimitReached ? "none" : undefined,
244+
}),
245+
)
241246
const webSocketEligible =
242-
!(yield* hooks.has("session", "http.request")) && !(yield* hooks.has("session", "http.response"))
247+
!(yield* hooks.has("session", "http.request", resolved.ref.providerID)) &&
248+
!(yield* hooks.has("session", "http.response", resolved.ref.providerID))
243249
const http = webSocketEligible
244250
? undefined
245251
: SessionModelHttp.middleware(hooks, {
@@ -252,7 +258,7 @@ export const layer = Layer.effect(
252258
...(webSocket &&
253259
webSocketEligible &&
254260
resolved.ref.providerID === Provider.ID.openai &&
255-
model.route.id === "openai-responses"
261+
request.model.route.id === "openai-responses"
256262
? { webSocket: transport.bind(session.id) }
257263
: {}),
258264
}

0 commit comments

Comments
 (0)