feat(plugin): add experimental WebSocket handshake hook

packages/ai/src/providers/azure.ts

@@ -1,4 +1,3 @@
-import { Headers } from "effect/unstable/http"
 import { Auth } from "../route/auth.js"
 import { type AtLeastOne, type ProviderAuthOption } from "../route/auth-options.js"
 import type { Route, RouteDefaultsInput, CompactionOperations } from "../route/client.js"
@@ -61,11 +60,6 @@ const responsesRoute = OpenAIResponses.route.with({
       url.searchParams.delete("api-version")
       return url.toString()
     },
-    headers: (headers) => {
-      const apiKey = headers["api-key"]
-      if (!apiKey) return headers
-      return Headers.remove(Headers.set(headers, "authorization", `Bearer ${apiKey}`), "api-key")
-    },
   }),
 })
 

packages/ai/test/provider/openai-responses.test.ts

@@ -760,7 +760,11 @@ describe("OpenAI Responses route", () => {
       )
       const next = continuationDriver({
         ...request,
-        input: [...firstInput, reasoning, { type: "message", role: "user", content: [{ type: "input_text", text: "Continue" }] }],
+        input: [
+          ...firstInput,
+          reasoning,
+          { type: "message", role: "user", content: [{ type: "input_text", text: "Continue" }] },
+        ],
       })
 
       const continued = yield* next.create(saved)
@@ -795,7 +799,10 @@ describe("OpenAI Responses route", () => {
           ProviderShared.encodeJson({ type: "response.completed", response: { id: "resp_1" } }),
         ),
       )
-      const appended = [...request.input, { type: "message", role: "user", content: [{ type: "input_text", text: "Second" }] }]
+      const appended = [
+        ...request.input,
+        { type: "message", role: "user", content: [{ type: "input_text", text: "Second" }] },
+      ]
       const changes = [
         { ...request, model: "gpt-5.3", input: appended },
         { ...request, instructions: "Changed", input: appended },
@@ -804,7 +811,10 @@ describe("OpenAI Responses route", () => {
         { ...request, metadata: { source: "two" }, input: appended },
         {
           ...request,
-          input: [{ type: "message", role: "user", content: [{ type: "input_text", text: "Rewritten history" }] }, appended[1]],
+          input: [
+            { type: "message", role: "user", content: [{ type: "input_text", text: "Rewritten history" }] },
+            appended[1],
+          ],
         },
       ]
 
@@ -900,7 +910,10 @@ describe("OpenAI Responses route", () => {
       const second = continuationDriver(
         {
           ...firstRequest,
-          input: [...firstRequest.input, { type: "message", role: "user", content: [{ type: "input_text", text: "Second" }] }],
+          input: [
+            ...firstRequest.input,
+            { type: "message", role: "user", content: [{ type: "input_text", text: "Second" }] },
+          ],
         },
         classifyingChannelDriver,
       )
@@ -951,7 +964,10 @@ describe("OpenAI Responses route", () => {
       }
       const secondRequest = {
         ...firstRequest,
-        input: [...firstRequest.input, { type: "message", role: "user", content: [{ type: "input_text", text: "Second" }] }],
+        input: [
+          ...firstRequest.input,
+          { type: "message", role: "user", content: [{ type: "input_text", text: "Second" }] },
+        ],
       }
       const saved = checkpoint(
         yield* continuationDriver(firstRequest).observe(
@@ -1082,7 +1098,9 @@ describe("OpenAI Responses route", () => {
       )
 
       const expected = {
-        input: [{ type: "message", role: "user", content: [{ type: "input_text", text: "Say \uFFFDhello \u{1F600}." }] }],
+        input: [
+          { type: "message", role: "user", content: [{ type: "input_text", text: "Say \uFFFDhello \u{1F600}." }] },
+        ],
         metadata: { source: "overlay\uFFFD" },
       }
       expect(JSON.parse(yield* Ref.get(message))).toMatchObject(expected)
@@ -1124,7 +1142,9 @@ describe("OpenAI Responses route", () => {
     }),
   )
 
-  it.effect("builds Azure WebSocket requests with v1 URLs and bearer auth", () =>
+  // Azure's WebSocket upgrade accepts the same credential header as HTTP: `api-key` for keys,
+  // `Authorization: Bearer` for Entra tokens. Rewriting a key into a bearer is rejected.
+  it.effect("builds Azure WebSocket requests with v1 URLs and the route's auth header", () =>
     Effect.gen(function* () {
       const deps = Layer.succeed(
         RequestExecutor.Service,
@@ -1133,13 +1153,13 @@ describe("OpenAI Responses route", () => {
       const cases = [
         {
           model: Azure.configure({ resourceName: "opencode-test", apiKey: "azure-key" }).responses("deployment"),
-          authorization: "Bearer azure-key",
+          headers: { "api-key": "azure-key", authorization: undefined },
         },
         {
           model: Azure.configure({ resourceName: "opencode-test", auth: Auth.bearer("entra-token") }).responses(
             "deployment",
           ),
-          authorization: "Bearer entra-token",
+          headers: { "api-key": undefined, authorization: "Bearer entra-token" },
         },
       ]
 
@@ -1150,8 +1170,8 @@ describe("OpenAI Responses route", () => {
               Effect.gen(function* () {
                 expect(exchange.connect.url).toBe("wss://opencode-test.openai.azure.com/openai/v1/responses")
                 expect(exchange.connect.rotateAfterMs).toBe(55 * 60 * 1000)
-                expect(exchange.connect.headers.authorization).toBe(item.authorization)
-                expect(exchange.connect.headers["api-key"]).toBeUndefined()
+                expect(exchange.connect.headers.authorization).toBe(item.headers.authorization)
+                expect(exchange.connect.headers["api-key"]).toBe(item.headers["api-key"])
                 expect(exchange.connect.headers["openai-beta"]).toBeUndefined()
                 expect(JSON.parse((yield* exchange.driver.create(undefined)).message)).toMatchObject({
                   type: "response.create",
@@ -1850,7 +1870,11 @@ describe("OpenAI Responses route", () => {
             summary: [{ type: "summary_text", text: "I inspected the previous turn." }],
           },
           { role: "assistant", content: [{ type: "output_text", text: "It shows a small test image." }] },
-          { type: "message", role: "user", content: [{ type: "input_text", text: "Check the weather in Paris before continuing." }] },
+          {
+            type: "message",
+            role: "user",
+            content: [{ type: "input_text", text: "Check the weather in Paris before continuing." }],
+          },
           { type: "function_call", call_id: "call_weather_1", name: "get_weather", arguments: '{"city":"Paris"}' },
           { type: "function_call_output", call_id: "call_weather_1", output: '{"temperature":22}' },
           { role: "assistant", content: [{ type: "output_text", text: "Paris is 22 degrees." }] },
@@ -4607,8 +4631,20 @@ describe("OpenAI Responses route", () => {
         compileRequest(LLM.request({ model, messages: [response.message], providerOptions: { store } })),
       )
       expect(prepared.map((request) => request.body.input)).toEqual([
-        [{ type: "message", role: "user", content: [{ type: "input_image", image_url: "data:image/png;base64,AQID" }] }],
-        [{ type: "message", role: "user", content: [{ type: "input_image", image_url: "data:image/png;base64,AQID" }] }],
+        [
+          {
+            type: "message",
+            role: "user",
+            content: [{ type: "input_image", image_url: "data:image/png;base64,AQID" }],
+          },
+        ],
+        [
+          {
+            type: "message",
+            role: "user",
+            content: [{ type: "input_image", image_url: "data:image/png;base64,AQID" }],
+          },
+        ],
       ])
     }),
   )

packages/core/src/plugin/provider/azure.ts

@@ -180,29 +180,50 @@ export const AzurePlugin = define({
       Effect.forkScoped({ startImmediately: true }),
     )
 
+    // Entra bearer tokens are minted per request from the target URL's scope, so they are injected
+    // at the transport hooks rather than stored as a credential.
+    const bearer = Effect.fn("AzurePlugin.bearer")(function* (url: string) {
+      const connection = yield* ctx.integration.connection.active(Provider.ID.azure)
+      const credential = connection
+        ? yield* ctx.integration.connection.resolve(connection).pipe(Effect.orElseSucceed(() => undefined))
+        : undefined
+      if (credential?.type !== "oauth" || credential.methodID !== methodID) return
+      const target = new URL(url)
+      const scope =
+        target.hostname.endsWith(".services.ai.azure.com") && !target.pathname.startsWith("/models")
+          ? foundryScope
+          : cognitiveScope
+      const current = yield* token(scope).pipe(Effect.orDie)
+      return `Bearer ${current.access}`
+    })
     yield* ctx.session.hook(
       "http.request",
       (evt) =>
         Effect.gen(function* () {
           if (evt.model.providerID !== Provider.ID.azure) return
-          const connection = yield* ctx.integration.connection.active(Provider.ID.azure)
-          const credential = connection
-            ? yield* ctx.integration.connection.resolve(connection).pipe(Effect.orElseSucceed(() => undefined))
-            : undefined
-          if (credential?.type !== "oauth" || credential.methodID !== methodID) return
-          const url = new URL(evt.request.url)
-          const scope =
-            url.hostname.endsWith(".services.ai.azure.com") && !url.pathname.startsWith("/models")
-              ? foundryScope
-              : cognitiveScope
-          const current = yield* token(scope).pipe(Effect.orDie)
+          const authorization = yield* bearer(evt.request.url)
+          if (!authorization) return
           evt.request.headers.delete("api-key")
           evt.request.headers.delete("x-api-key")
-          evt.request.headers.set("authorization", `Bearer ${current.access}`)
+          evt.request.headers.set("authorization", authorization)
           evt.request.headers.set("user-agent", App.useragent(ctx.app))
         }),
       { providerID: Provider.ID.azure },
     )
+    yield* ctx.session.hook(
+      "experimental.ws.handshake",
+      (evt) =>
+        Effect.gen(function* () {
+          if (evt.model.providerID !== Provider.ID.azure) return
+          const authorization = yield* bearer(evt.url)
+          if (!authorization) return
+          delete evt.headers["api-key"]
+          delete evt.headers["x-api-key"]
+          evt.headers.authorization = authorization
+          evt.headers["user-agent"] = App.useragent(ctx.app)
+        }),
+      { providerID: Provider.ID.azure },
+    )
   }),
 })
 

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

@@ -316,14 +316,25 @@ export const layer = Layer.effect(
               return HttpClientResponse.fromWeb(sent, after.response)
             }).pipe(Effect.mapError((cause) => (cause instanceof Error ? cause : new Error(String(cause)))))
         : undefined
-      // HTTP hooks must observe every request, so they keep the provider on HTTP. The route decides
-      // whether it actually has a WebSocket channel and falls back to HTTP otherwise.
-      const webSocket = input.webSocket === "session" && !hasHttpHooks && model.transport === "websocket"
+      // HTTP hooks wrap every HTTP request, including the WebSocket fallback path. The route decides
+      // which transport actually carries the request, so both hook families are always offered.
+      const webSocket =
+        input.webSocket === "session" && model.transport === "websocket"
+          ? transport.bind(session.id, (connect) =>
+              hooks
+                .trigger("session", "experimental.ws.handshake", {
+                  ...scope,
+                  url: connect.url,
+                  headers: connect.headers,
+                })
+                .pipe(Effect.map((event) => ({ url: event.url, headers: event.headers }))),
+            )
+          : undefined
 
       return {
         event: shaped,
         request,
-        options: { ...(http ? { http } : {}), ...(webSocket ? { webSocket: transport.bind(session.id) } : {}) },
+        options: { ...(http ? { http } : {}), ...(webSocket ? { webSocket } : {}) },
         retry: (event: Parameters<Prepared["retry"]>[0]) =>
           hooks.trigger("session", "retry", event).pipe(Effect.asVoid),
         // Permission.assert and the question tool throw declines as defects so tools cannot

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

@@ -13,6 +13,7 @@ import {
 import { AIError, AIErrorReason, TransportError, type TransportOperation } from "@opencode/ai"
 import { Hash } from "@opencode/util/hash"
 import { Cause, Clock, Context, Effect, Fiber, Layer, Metric, Queue, Scope, Semaphore, Stream } from "effect"
+import { Headers } from "effect/unstable/http"
 import { Socket } from "effect/unstable/socket"
 import { makeGlobalNode } from "@opencode/util/effect/app-node"
 import { SessionSchema } from "./schema.js"
@@ -52,8 +53,17 @@ interface State {
   channel?: Channel
 }
 
+/** Selects the connection for one exchange. Its output feeds the affinity key, so changed headers reopen the socket. */
+export interface Handshake {
+  readonly url: string
+  readonly headers: Record<string, string>
+}
+
 export interface Interface {
-  readonly bind: (sessionID: SessionSchema.ID) => WebSocketChannelExecutor
+  readonly bind: (
+    sessionID: SessionSchema.ID,
+    handshake?: (connect: Handshake) => Effect.Effect<Handshake>,
+  ) => WebSocketChannelExecutor
   readonly close: (sessionID: SessionSchema.ID) => Effect.Effect<void>
   readonly closeAll: Effect.Effect<void>
 }
@@ -267,7 +277,8 @@ export const makeLayer = (connector: WebSocketConnector) =>
 
       const start = Effect.fn("SessionModelTransport.start")(function* (
         owner: State,
-        exchange: WebSocketChannelExchange,
+        input: WebSocketChannelExchange,
+        handshake?: (connect: Handshake) => Effect.Effect<Handshake>,
       ) {
         if (owner.closed)
           return yield* transportError("Session WebSocket owner is closed", {
@@ -276,7 +287,13 @@ export const makeLayer = (connector: WebSocketConnector) =>
             phase: "queue",
             delivery: "not-sent",
           })
-        if (owner.httpFallback) return fallback(exchange)
+        if (owner.httpFallback) return fallback(input)
+        const selected = handshake
+          ? yield* handshake({ url: input.connect.url, headers: { ...input.connect.headers } })
+          : undefined
+        const exchange: WebSocketChannelExchange = selected
+          ? { ...input, connect: { ...input.connect, url: selected.url, headers: Headers.fromInput(selected.headers) } }
+          : input
         const key = affinity(exchange)
         const now = yield* Clock.currentTimeMillis
         const current = owner.channel
@@ -465,7 +482,10 @@ export const makeLayer = (connector: WebSocketConnector) =>
         return { frames, complete, http: channel.connection.http }
       })
 
-      const bind = (sessionID: SessionSchema.ID): WebSocketChannelExecutor => ({
+      const bind = (
+        sessionID: SessionSchema.ID,
+        handshake?: (connect: Handshake) => Effect.Effect<Handshake>,
+      ): WebSocketChannelExecutor => ({
         execute: (exchange) => {
           const owner = state(sessionID)
           let execution: WebSocketChannelExecution | undefined
@@ -475,7 +495,7 @@ export const makeLayer = (connector: WebSocketConnector) =>
             },
             frames: Stream.unwrap(
               Effect.acquireRelease(owner.lock.take(1), () => owner.lock.release(1), { interruptible: true }).pipe(
-                Effect.andThen(start(owner, exchange)),
+                Effect.andThen(start(owner, exchange, handshake)),
                 Effect.tap((started) =>
                   Effect.sync(() => {
                     execution = started

packages/core/test/plugin/provider-azure.test.ts

@@ -301,6 +301,20 @@ describe("AzurePlugin", () => {
           })
           expect(foundry.request.headers.get("authorization")).toBe("Bearer https://ai.azure.com/.default-token")
           expect(foundry.request.headers.has("x-api-key")).toBe(false)
+
+          const handshake = yield* hooks.trigger("session", "experimental.ws.handshake", {
+            sessionID: Session.ID.make("ses_azure_ws"),
+            agent: Agent.ID.make("build"),
+            model,
+            kind: "primary",
+            url: "wss://test-resource.openai.azure.com/openai/v1/responses",
+            headers: { "api-key": "stored-token", "x-keep": "yes" },
+          })
+          expect(handshake.headers).toMatchObject({
+            authorization: "Bearer https://cognitiveservices.azure.com/.default-token",
+            "x-keep": "yes",
+          })
+          expect(handshake.headers).not.toHaveProperty("api-key")
         }),
     ),
   )
@@ -362,9 +376,7 @@ describe("AzurePlugin", () => {
           resourceName: "from-env",
           baseURL: "https://from-env.cognitiveservices.azure.com/openai",
         })
-        expect(
-          required(yield* models.get(Provider.ID.azure, Model.ID.make("anthropic"))).settings,
-        ).toMatchObject({
+        expect(required(yield* models.get(Provider.ID.azure, Model.ID.make("anthropic"))).settings).toMatchObject({
           resourceName: "model-resource",
           baseURL: "https://model-resource.services.ai.azure.com/anthropic/v1",
         })
@@ -469,5 +481,4 @@ describe("AzurePlugin", () => {
       }),
     ),
   )
-
 })