Skip to content

Commit 9073c52

Browse files
rekram1-noderekram1-node
authored
feat(plugin): add experimental WebSocket send and receive hooks (#49136)
1 parent 04c2963 commit 9073c52

8 files changed

Lines changed: 224 additions & 34 deletions

File tree

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

Lines changed: 18 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -320,15 +320,24 @@ export const layer = Layer.effect(
320320
// which transport actually carries the request, so both hook families are always offered.
321321
const webSocket =
322322
input.webSocket === "session" && model.transport === "websocket"
323-
? transport.bind(session.id, (connect) =>
324-
hooks
325-
.trigger("session", "experimental.ws.handshake", {
326-
...scope,
327-
url: connect.url,
328-
headers: connect.headers,
329-
})
330-
.pipe(Effect.map((event) => ({ url: event.url, headers: event.headers }))),
331-
)
323+
? transport.bind(session.id, {
324+
handshake: (connect) =>
325+
hooks
326+
.trigger("session", "experimental.ws.handshake", {
327+
...scope,
328+
url: connect.url,
329+
headers: connect.headers,
330+
})
331+
.pipe(Effect.map((event) => ({ url: event.url, headers: event.headers }))),
332+
send: (frame) =>
333+
hooks
334+
.trigger("session", "experimental.ws.send", { ...scope, frame })
335+
.pipe(Effect.map((event) => event.frame)),
336+
receive: (frame) =>
337+
hooks
338+
.trigger("session", "experimental.ws.receive", { ...scope, frame })
339+
.pipe(Effect.map((event) => event.frame)),
340+
})
332341
: undefined
333342

334343
return {

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

Lines changed: 21 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -59,11 +59,18 @@ export interface Handshake {
5959
readonly headers: Record<string, string>
6060
}
6161

62+
/**
63+
* Per-exchange taps. `handshake` runs before the connection is selected; `send` sees each outbound
64+
* frame after the driver builds it; `receive` sees each inbound frame before the driver observes it.
65+
*/
66+
export interface Interceptor {
67+
readonly handshake?: (connect: Handshake) => Effect.Effect<Handshake>
68+
readonly send?: (frame: string) => Effect.Effect<string>
69+
readonly receive?: (frame: string) => Effect.Effect<string>
70+
}
71+
6272
export interface Interface {
63-
readonly bind: (
64-
sessionID: SessionSchema.ID,
65-
handshake?: (connect: Handshake) => Effect.Effect<Handshake>,
66-
) => WebSocketChannelExecutor
73+
readonly bind: (sessionID: SessionSchema.ID, interceptor?: Interceptor) => WebSocketChannelExecutor
6774
readonly close: (sessionID: SessionSchema.ID) => Effect.Effect<void>
6875
readonly closeAll: Effect.Effect<void>
6976
}
@@ -278,7 +285,7 @@ export const makeLayer = (connector: WebSocketConnector) =>
278285
const start = Effect.fn("SessionModelTransport.start")(function* (
279286
owner: State,
280287
input: WebSocketChannelExchange,
281-
handshake?: (connect: Handshake) => Effect.Effect<Handshake>,
288+
interceptor?: Interceptor,
282289
) {
283290
if (owner.closed)
284291
return yield* transportError("Session WebSocket owner is closed", {
@@ -288,8 +295,8 @@ export const makeLayer = (connector: WebSocketConnector) =>
288295
delivery: "not-sent",
289296
})
290297
if (owner.httpFallback) return fallback(input)
291-
const selected = handshake
292-
? yield* handshake({ url: input.connect.url, headers: { ...input.connect.headers } })
298+
const selected = interceptor?.handshake
299+
? yield* interceptor.handshake({ url: input.connect.url, headers: { ...input.connect.headers } })
293300
: undefined
294301
const exchange: WebSocketChannelExchange = selected
295302
? { ...input, connect: { ...input.connect, url: selected.url, headers: Headers.fromInput(selected.headers) } }
@@ -354,6 +361,9 @@ export const makeLayer = (connector: WebSocketConnector) =>
354361
Effect.onInterrupt(() => closeChannel(owner, channel)),
355362
)
356363
if (create.mode === "full") channel.checkpoint = undefined
364+
const message = interceptor?.send
365+
? yield* interceptor.send(create.message).pipe(Effect.onInterrupt(() => closeChannel(owner, channel)))
366+
: create.message
357367
yield* Effect.logDebug("session websocket sending", {
358368
sessionTransport: "websocket",
359369
phase: "send",
@@ -364,7 +374,7 @@ export const makeLayer = (connector: WebSocketConnector) =>
364374
delivery: "send-attempted",
365375
}
366376
channel.active = active
367-
const sent = yield* channel.connection.sendText(create.message).pipe(
377+
const sent = yield* channel.connection.sendText(message).pipe(
368378
Effect.withSpan("SessionModelTransport.send"),
369379
Effect.onInterrupt(() => closeChannel(owner, channel)),
370380
Effect.result,
@@ -405,6 +415,7 @@ export const makeLayer = (connector: WebSocketConnector) =>
405415
}),
406416
),
407417
}),
418+
Stream.mapEffect((frame) => (interceptor?.receive ? interceptor.receive(frame) : Effect.succeed(frame))),
408419
Stream.mapEffect((frame) => exchange.driver.observe(create, frame)),
409420
Stream.tap((observation) =>
410421
Effect.sync(() => {
@@ -482,10 +493,7 @@ export const makeLayer = (connector: WebSocketConnector) =>
482493
return { frames, complete, http: channel.connection.http }
483494
})
484495

485-
const bind = (
486-
sessionID: SessionSchema.ID,
487-
handshake?: (connect: Handshake) => Effect.Effect<Handshake>,
488-
): WebSocketChannelExecutor => ({
496+
const bind = (sessionID: SessionSchema.ID, interceptor?: Interceptor): WebSocketChannelExecutor => ({
489497
execute: (exchange) => {
490498
const owner = state(sessionID)
491499
let execution: WebSocketChannelExecution | undefined
@@ -495,7 +503,7 @@ export const makeLayer = (connector: WebSocketConnector) =>
495503
},
496504
frames: Stream.unwrap(
497505
Effect.acquireRelease(owner.lock.take(1), () => owner.lock.release(1), { interruptible: true }).pipe(
498-
Effect.andThen(start(owner, exchange, handshake)),
506+
Effect.andThen(start(owner, exchange, interceptor)),
499507
Effect.tap((started) =>
500508
Effect.sync(() => {
501509
execution = started

‎packages/core/test/session-model-request-hooks.test.ts‎

Lines changed: 28 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ describe("SessionModelRequest HTTP hooks", () => {
8080
}).pipe(Effect.provideService(SessionModelTransport.Service, transport)),
8181
)
8282

83-
it.effect("offers the WebSocket executor alongside HTTP hooks and routes the handshake hook", () =>
83+
it.effect("offers the WebSocket executor alongside HTTP hooks and routes the WebSocket hooks", () =>
8484
Effect.gen(function* () {
8585
const hooks = yield* PluginHooks.Service
8686
const seen: string[] = []
@@ -92,13 +92,31 @@ describe("SessionModelRequest HTTP hooks", () => {
9292
delete event.headers["api-key"]
9393
}),
9494
)
95+
yield* hooks.register("session", "experimental.ws.send", (event) =>
96+
Effect.sync(() => {
97+
seen.push(`send:${event.kind}:${event.frame}`)
98+
event.frame = `${event.frame}+plugin`
99+
}),
100+
)
101+
yield* hooks.register("session", "experimental.ws.receive", (event) =>
102+
Effect.sync(() => {
103+
seen.push(`receive:${event.kind}:${event.frame}`)
104+
event.frame = event.frame.toUpperCase()
105+
}),
106+
)
95107
const bound: Array<{ url: string; headers: Record<string, string> }> = []
108+
const frames: string[] = []
96109
const websocketTransport = SessionModelTransport.Service.of({
97-
bind: (_sessionID, handshake) => ({
110+
bind: (_sessionID, interceptor) => ({
98111
execute: () =>
99112
Effect.gen(function* () {
100-
if (!handshake) throw new Error("Expected a handshake interceptor")
101-
bound.push(yield* handshake({ url: "wss://example.test/v1/responses", headers: { "api-key": "k" } }))
113+
if (!interceptor?.handshake || !interceptor.send || !interceptor.receive)
114+
throw new Error("Expected a full WebSocket interceptor")
115+
bound.push(
116+
yield* interceptor.handshake({ url: "wss://example.test/v1/responses", headers: { "api-key": "k" } }),
117+
)
118+
frames.push(yield* interceptor.send("create"))
119+
frames.push(yield* interceptor.receive("created"))
102120
return { frames: Stream.empty, complete: Effect.void }
103121
}),
104122
}),
@@ -127,7 +145,12 @@ describe("SessionModelRequest HTTP hooks", () => {
127145
expect(prepared.options.webSocket).toBeDefined()
128146
yield* prepared.options.webSocket!.execute({} as never)
129147
expect(bound).toEqual([{ url: "wss://example.test/v1/responses", headers: { authorization: "Bearer minted" } }])
130-
expect(seen).toEqual(["handshake:primary:wss://example.test/v1/responses"])
148+
expect(frames).toEqual(["create+plugin", "CREATED"])
149+
expect(seen).toEqual([
150+
"handshake:primary:wss://example.test/v1/responses",
151+
"send:primary:create",
152+
"receive:primary:created",
153+
])
131154
}),
132155
)
133156
})

‎packages/core/test/session-model-transport.test.ts‎

Lines changed: 37 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -178,12 +178,13 @@ describe("SessionModelTransport", () => {
178178
fixture.connector,
179179
Effect.gen(function* () {
180180
const transport = yield* SessionModelTransport.Service
181-
const executor = transport.bind(session, (connect) =>
182-
Effect.succeed({
183-
url: connect.url,
184-
headers: { ...connect.headers, authorization: `Bearer ${tokens.shift()}` },
185-
}),
186-
)
181+
const executor = transport.bind(session, {
182+
handshake: (connect) =>
183+
Effect.succeed({
184+
url: connect.url,
185+
headers: { ...connect.headers, authorization: `Bearer ${tokens.shift()}` },
186+
}),
187+
})
187188
yield* collect(executor, exchange("first", { headers: { "api-key": "k" } }))
188189
yield* collect(executor, exchange("second", { headers: { "api-key": "k" } }))
189190
yield* collect(executor, exchange("third", { headers: { "api-key": "k" } }))
@@ -196,6 +197,36 @@ describe("SessionModelTransport", () => {
196197
)
197198
})
198199

200+
test("sends the frame the send tap returns and observes the frame the receive tap returns", async () => {
201+
const fixture = automatic()
202+
const seen: Array<{ tap: "send" | "receive"; frame: string }> = []
203+
await run(
204+
fixture.connector,
205+
Effect.gen(function* () {
206+
const transport = yield* SessionModelTransport.Service
207+
const executor = transport.bind(session, {
208+
send: (frame) => {
209+
seen.push({ tap: "send", frame })
210+
return Effect.succeed(`${frame}:rewritten`)
211+
},
212+
receive: (frame) => {
213+
seen.push({ tap: "receive", frame })
214+
return Effect.succeed(`${frame}:observed`)
215+
},
216+
})
217+
const frames = yield* collect(executor, exchange("first"))
218+
219+
// The wire carries the rewritten outbound frame; the driver sees the rewritten inbound frame.
220+
expect(fixture.connections.map((item) => item.sent)).toEqual([["first:rewritten"]])
221+
expect(frames).toEqual(["completed:first:rewritten:observed"])
222+
expect(seen).toEqual([
223+
{ tap: "send", frame: "first" },
224+
{ tap: "receive", frame: "completed:first:rewritten" },
225+
])
226+
}),
227+
)
228+
})
229+
199230
test("does not carry a checkpoint across physical connection rotation", async () => {
200231
const fixture = automatic()
201232
const checkpoints: Array<unknown> = []

‎packages/plugin/src/effect/session.ts‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,31 @@ export interface SessionWebSocketHandshake {
9999
headers: Record<string, string>
100100
}
101101

102+
/**
103+
* Outbound frame about to be written to the Session's socket, after the provider driver has built
104+
* it. Replacing `frame` sends the replacement verbatim; the driver still tracks state from the
105+
* provider's replies, so a rewrite that changes protocol meaning is on the plugin. Experimental.
106+
*/
107+
export interface SessionWebSocketSend {
108+
readonly sessionID: Session.ID
109+
readonly agent: Agent.ID
110+
readonly model: Model.Ref
111+
readonly kind: SessionRequestKind
112+
frame: string
113+
}
114+
115+
/**
116+
* Inbound frame read from the Session's socket, before the provider driver observes it. Replacing
117+
* `frame` hands the replacement to the driver verbatim. Experimental.
118+
*/
119+
export interface SessionWebSocketReceive {
120+
readonly sessionID: Session.ID
121+
readonly agent: Agent.ID
122+
readonly model: Model.Ref
123+
readonly kind: SessionRequestKind
124+
frame: string
125+
}
126+
102127
export type SessionRetryDecision = { retry: false } | { retry: true; delay: number }
103128

104129
export interface SessionRetry {
@@ -120,6 +145,8 @@ export interface SessionHooks {
120145
readonly "http.request": SessionHttpRequest
121146
readonly "http.response": SessionHttpResponse
122147
readonly "experimental.ws.handshake": SessionWebSocketHandshake
148+
readonly "experimental.ws.send": SessionWebSocketSend
149+
readonly "experimental.ws.receive": SessionWebSocketReceive
123150
readonly retry: SessionRetry
124151
}
125152

‎packages/plugin/src/promise/session.ts‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,31 @@ export interface SessionWebSocketHandshake {
9999
headers: Record<string, string>
100100
}
101101

102+
/**
103+
* Outbound frame about to be written to the Session's socket, after the provider driver has built
104+
* it. Replacing `frame` sends the replacement verbatim; the driver still tracks state from the
105+
* provider's replies, so a rewrite that changes protocol meaning is on the plugin. Experimental.
106+
*/
107+
export interface SessionWebSocketSend {
108+
readonly sessionID: Session.ID
109+
readonly agent: Agent.ID
110+
readonly model: Model.Ref
111+
readonly kind: SessionRequestKind
112+
frame: string
113+
}
114+
115+
/**
116+
* Inbound frame read from the Session's socket, before the provider driver observes it. Replacing
117+
* `frame` hands the replacement to the driver verbatim. Experimental.
118+
*/
119+
export interface SessionWebSocketReceive {
120+
readonly sessionID: Session.ID
121+
readonly agent: Agent.ID
122+
readonly model: Model.Ref
123+
readonly kind: SessionRequestKind
124+
frame: string
125+
}
126+
102127
export type SessionRetryDecision = { retry: false } | { retry: true; delay: number }
103128

104129
export interface SessionRetry {
@@ -120,6 +145,8 @@ export interface SessionHooks {
120145
readonly "http.request": SessionHttpRequest
121146
readonly "http.response": SessionHttpResponse
122147
readonly "experimental.ws.handshake": SessionWebSocketHandshake
148+
readonly "experimental.ws.send": SessionWebSocketSend
149+
readonly "experimental.ws.receive": SessionWebSocketReceive
123150
readonly retry: SessionRetry
124151
}
125152

‎services/www/src/docs/content/build/plugins/effect.mdx‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1277,6 +1277,26 @@ effect: (ctx) =>
12771277
}),
12781278
```
12791279

1280+
`experimental.ws.send` and `experimental.ws.receive` expose the frames themselves: `send` runs after the provider
1281+
driver builds an outbound frame, `receive` runs on each inbound frame before the driver observes it. Whatever `frame` holds when the hook returns is what crosses the wire or reaches the driver;
1282+
OpenCode does not validate it.
1283+
1284+
```ts
1285+
effect: (ctx) =>
1286+
Effect.gen(function* () {
1287+
yield* ctx.session.hook(
1288+
"experimental.ws.send",
1289+
(event) =>
1290+
Effect.sync(() => {
1291+
const body = JSON.parse(event.frame)
1292+
if (body.type === "response.create") body.metadata = { ...body.metadata, session: event.sessionID }
1293+
event.frame = JSON.stringify(body)
1294+
}),
1295+
{ providerID: "openai" },
1296+
)
1297+
}),
1298+
```
1299+
12801300
Override the retry decision for a provider failure or replace its delay in milliseconds. The hook runs after OpenCode
12811301
classifies the failure and proposes its policy, but before any retry is scheduled. It does not expose how OpenCode
12821302
internally performs the next attempt.
@@ -1317,6 +1337,9 @@ interface SessionHooks {
13171337
readonly "model.request": SessionModelRequest
13181338
readonly "http.request": SessionHttpRequest
13191339
readonly "http.response": SessionHttpResponse
1340+
readonly "experimental.ws.handshake": SessionWebSocketHandshake
1341+
readonly "experimental.ws.send": SessionWebSocketSend
1342+
readonly "experimental.ws.receive": SessionWebSocketReceive
13201343
readonly retry: SessionRetry
13211344
}
13221345

0 commit comments

Comments
 (0)