refactor(llm): move websocket recorder into http-recorder

packages/http-recorder/src/diff.ts

@@ -2,7 +2,7 @@ import { Option } from "effect"
 import { Headers, HttpBody, HttpClientRequest, UrlParams } from "effect/unstable/http"
 import { decodeJson } from "./matching"
 import { REDACTED, redactUrl, secretFindings } from "./redaction"
-import { isHttpInteraction, type Cassette, type RequestSnapshot } from "./schema"
+import { httpInteractions, type Cassette, type RequestSnapshot } from "./schema"
 
 const safeText = (value: unknown) => {
   if (value === undefined) return "undefined"
@@ -75,7 +75,7 @@ export const requestDiff = (expected: RequestSnapshot, received: RequestSnapshot
 }
 
 export const mismatchDetail = (cassette: Cassette, incoming: RequestSnapshot) => {
-  const interactions = cassette.interactions.filter(isHttpInteraction)
+  const interactions = httpInteractions(cassette)
   if (interactions.length === 0) return "cassette has no recorded HTTP interactions"
   const ranked = interactions
     .map((interaction, index) => ({ index, lines: requestDiff(interaction.request, incoming) }))

packages/http-recorder/src/effect.ts

@@ -11,7 +11,7 @@ import { redactedErrorRequest, mismatchDetail, requestDiff } from "./diff"
 import { defaultMatcher, decodeJson, type RequestMatcher } from "./matching"
 import { redactHeaders, redactUrl, type SecretFinding } from "./redaction"
 import {
-  isHttpInteraction,
+  httpInteractions,
   type Cassette,
   type CassetteMetadata,
   type HttpInteraction,
@@ -138,7 +138,7 @@ export const recordingLayer = (
 
       const selectInteraction = (cassette: Cassette, incoming: HttpInteraction["request"]) =>
         Effect.gen(function* () {
-          const interactions = cassette.interactions.filter(isHttpInteraction)
+          const interactions = httpInteractions(cassette)
           if (sequential) {
             const index = yield* Ref.get(cursor)
             const interaction = interactions[index]

packages/http-recorder/src/index.ts

@@ -3,6 +3,7 @@ export * from "./redaction"
 export * from "./matching"
 export * from "./diff"
 export * from "./storage"
+export * from "./websocket"
 export * from "./effect"
 export * as Cassette from "./cassette"
 

packages/http-recorder/src/matching.ts

@@ -4,13 +4,16 @@ import type { RequestSnapshot } from "./schema"
 const JsonValue = Schema.fromJsonString(Schema.Unknown)
 export const decodeJson = Schema.decodeUnknownOption(JsonValue)
 
-const canonicalize = (value: unknown): unknown => {
-  if (Array.isArray(value)) return value.map(canonicalize)
-  if (value !== null && typeof value === "object") {
+const isRecord = (value: unknown): value is Record<string, unknown> =>
+  value !== null && typeof value === "object" && !Array.isArray(value)
+
+export const canonicalizeJson = (value: unknown): unknown => {
+  if (Array.isArray(value)) return value.map(canonicalizeJson)
+  if (isRecord(value)) {
     return Object.fromEntries(
-      Object.keys(value as Record<string, unknown>)
+      Object.keys(value)
         .toSorted()
-        .map((key) => [key, canonicalize((value as Record<string, unknown>)[key])]),
+        .map((key) => [key, canonicalizeJson(value[key])]),
     )
   }
   return value
@@ -22,10 +25,10 @@ export const canonicalSnapshot = (snapshot: RequestSnapshot): string =>
   JSON.stringify({
     method: snapshot.method,
     url: snapshot.url,
-    headers: canonicalize(snapshot.headers),
+    headers: canonicalizeJson(snapshot.headers),
     body: Option.match(decodeJson(snapshot.body), {
       onNone: () => snapshot.body,
-      onSome: canonicalize,
+      onSome: canonicalizeJson,
     }),
   })
 

packages/http-recorder/src/schema.ts

@@ -52,6 +52,10 @@ export const isHttpInteraction = (interaction: Interaction): interaction is Http
 export const isWebSocketInteraction = (interaction: Interaction): interaction is WebSocketInteraction =>
   interaction.transport === "websocket"
 
+export const httpInteractions = (cassette: Cassette) => cassette.interactions.filter(isHttpInteraction)
+
+export const webSocketInteractions = (cassette: Cassette) => cassette.interactions.filter(isWebSocketInteraction)
+
 export const CassetteSchema = Schema.Struct({
   version: Schema.Literal(1),
   metadata: Schema.optional(CassetteMetadataSchema),

packages/http-recorder/src/websocket.ts

@@ -0,0 +1,203 @@
+import { Effect, Option, Ref, Scope, Stream } from "effect"
+import type { Headers } from "effect/unstable/http"
+import * as CassetteService from "./cassette"
+import { canonicalizeJson, decodeJson } from "./matching"
+import { redactHeaders, redactUrl, type SecretFinding } from "./redaction"
+import { webSocketInteractions, type CassetteMetadata, type WebSocketFrame, type WebSocketInteraction } from "./schema"
+
+export const DEFAULT_WEBSOCKET_REQUEST_HEADERS: ReadonlyArray<string> = ["content-type", "accept", "openai-beta"]
+
+export interface WebSocketRequest {
+  readonly url: string
+  readonly headers: Headers.Headers
+}
+
+export interface WebSocketConnection<E> {
+  readonly sendText: (message: string) => Effect.Effect<void, E>
+  readonly messages: Stream.Stream<string | Uint8Array, E>
+  readonly close: Effect.Effect<void>
+}
+
+export interface WebSocketExecutor<E> {
+  readonly open: (request: WebSocketRequest) => Effect.Effect<WebSocketConnection<E>, E>
+}
+
+export interface WebSocketRecordReplayOptions<E> {
+  readonly name: string
+  readonly mode?: "record" | "replay" | "passthrough"
+  readonly metadata?: CassetteMetadata
+  readonly cassette: CassetteService.Interface
+  readonly live: WebSocketExecutor<E>
+  readonly redact?: {
+    readonly headers?: ReadonlyArray<string>
+    readonly query?: ReadonlyArray<string>
+  }
+  readonly requestHeaders?: ReadonlyArray<string>
+  readonly compareClientMessagesAsJson?: boolean
+}
+
+const headersRecord = (headers: Headers.Headers) =>
+  Object.fromEntries(
+    Object.entries(headers as Record<string, unknown>)
+      .filter((entry): entry is [string, string] => typeof entry[1] === "string")
+      .toSorted(([a], [b]) => a.localeCompare(b)),
+  )
+
+const openSnapshot = (
+  request: WebSocketRequest,
+  options: Pick<WebSocketRecordReplayOptions<never>, "redact" | "requestHeaders"> = {},
+) => ({
+  url: redactUrl(request.url, options.redact?.query),
+  headers: redactHeaders(
+    headersRecord(request.headers),
+    options.requestHeaders ?? DEFAULT_WEBSOCKET_REQUEST_HEADERS,
+    options.redact?.headers,
+  ),
+})
+
+const textFrame = (body: string): WebSocketFrame => ({ kind: "text", body })
+
+const frameText = (frame: WebSocketFrame) => {
+  if (frame.kind === "text") return frame.body
+  return new TextDecoder().decode(Buffer.from(frame.body, "base64"))
+}
+
+const frameMessage = (frame: WebSocketFrame) =>
+  frame.kind === "text" ? frame.body : new Uint8Array(Buffer.from(frame.body, "base64"))
+
+const receivedFrame = (message: string | Uint8Array): WebSocketFrame =>
+  typeof message === "string"
+    ? textFrame(message)
+    : { kind: "binary", body: Buffer.from(message).toString("base64"), bodyEncoding: "base64" }
+
+const unsafeCassette = (name: string, findings: ReadonlyArray<SecretFinding>) =>
+  new Error(
+    `Refusing to write WebSocket cassette "${name}" because it contains possible secrets: ${findings
+      .map((item) => `${item.path} (${item.reason})`)
+      .join(", ")}`,
+  )
+
+const mismatch = (message: string, actual: unknown, expected: unknown) =>
+  new Error(`${message}: expected ${JSON.stringify(expected)}, received ${JSON.stringify(actual)}`)
+
+const assertEqual = (message: string, actual: unknown, expected: unknown) =>
+  Effect.sync(() => {
+    if (JSON.stringify(actual) === JSON.stringify(expected)) return
+    throw mismatch(message, actual, expected)
+  })
+
+const jsonOrText = (value: string) => Option.match(decodeJson(value), { onNone: () => value, onSome: canonicalizeJson })
+
+const compareClientMessage = (actual: string, expected: WebSocketFrame | undefined, index: number, asJson: boolean) => {
+  if (!expected)
+    return Effect.sync(() => {
+      throw new Error(`Unexpected WebSocket client frame ${index + 1}: ${actual}`)
+    })
+  const expectedText = frameText(expected)
+  if (!asJson) return assertEqual(`WebSocket client frame ${index + 1}`, actual, expectedText)
+  return assertEqual(`WebSocket client JSON frame ${index + 1}`, jsonOrText(actual), jsonOrText(expectedText))
+}
+
+export const makeWebSocketExecutor = <E>(
+  options: WebSocketRecordReplayOptions<E>,
+): Effect.Effect<WebSocketExecutor<E>, never, Scope.Scope> =>
+  Effect.gen(function* () {
+    const mode = options.mode ?? "replay"
+
+    if (mode === "passthrough") return options.live
+
+    if (mode === "record") {
+      return {
+        open: (request) =>
+          Effect.gen(function* () {
+            const client: WebSocketFrame[] = []
+            const server: WebSocketFrame[] = []
+            const connection = yield* options.live.open(request)
+            const closed = yield* Ref.make(false)
+            const closeOnce = Effect.gen(function* () {
+              if (yield* Ref.getAndSet(closed, true)) return
+              yield* connection.close
+              const result = yield* options.cassette
+                .append(
+                  options.name,
+                  { transport: "websocket", open: openSnapshot(request, options), client, server },
+                  options.metadata,
+                )
+                .pipe(Effect.orDie)
+              if (result.findings.length > 0) yield* Effect.die(unsafeCassette(options.name, result.findings))
+            })
+            return {
+              sendText: (message: string) =>
+                connection.sendText(message).pipe(Effect.tap(() => Effect.sync(() => client.push(textFrame(message))))),
+              messages: connection.messages.pipe(
+                Stream.map((message) => {
+                  server.push(receivedFrame(message))
+                  return message
+                }),
+              ),
+              close: closeOnce,
+            }
+          }),
+      }
+    }
+
+    const replay = yield* Ref.make<{ readonly interactions: ReadonlyArray<WebSocketInteraction> } | undefined>(
+      undefined,
+    )
+    const cursor = yield* Ref.make(0)
+
+    yield* Effect.addFinalizer(() =>
+      Effect.gen(function* () {
+        const input = yield* Ref.get(replay)
+        if (!input) return
+        yield* assertEqual(
+          `Unused recorded WebSocket interactions in ${options.name}`,
+          yield* Ref.get(cursor),
+          input.interactions.length,
+        )
+      }),
+    )
+
+    const loadReplay = Effect.fn("WebSocketRecorder.loadReplay")(function* () {
+      const cached = yield* Ref.get(replay)
+      if (cached) return cached
+      const input = {
+        interactions: webSocketInteractions(yield* options.cassette.read(options.name).pipe(Effect.orDie)),
+      }
+      yield* Ref.set(replay, input)
+      return input
+    })
+
+    return {
+      open: (request) => {
+        return Effect.gen(function* () {
+          const input = yield* loadReplay()
+          const index = yield* Ref.getAndUpdate(cursor, (value) => value + 1)
+          const interaction = input.interactions[index]
+          if (!interaction) return yield* Effect.die(new Error(`No recorded WebSocket interaction for ${request.url}`))
+          yield* assertEqual(`WebSocket open frame ${index + 1}`, openSnapshot(request, options), interaction.open)
+          const messageIndex = yield* Ref.make(0)
+          return {
+            sendText: (message: string) =>
+              Effect.gen(function* () {
+                const current = yield* Ref.getAndUpdate(messageIndex, (value) => value + 1)
+                yield* compareClientMessage(
+                  message,
+                  interaction.client[current],
+                  current,
+                  options.compareClientMessagesAsJson === true,
+                )
+              }),
+            messages: Stream.fromIterable(interaction.server).pipe(Stream.map(frameMessage)),
+            close: Effect.gen(function* () {
+              yield* assertEqual(
+                `WebSocket client frame count for interaction ${index + 1}`,
+                yield* Ref.get(messageIndex),
+                interaction.client.length,
+              )
+            }),
+          }
+        })
+      },
+    }
+  })

packages/http-recorder/test/record-replay.test.ts

@@ -1,6 +1,10 @@
+import { NodeFileSystem } from "@effect/platform-node"
 import { describe, expect, test } from "bun:test"
-import { Cause, Effect, Exit } from "effect"
-import { HttpBody, HttpClient, HttpClientRequest } from "effect/unstable/http"
+import { Cause, Effect, Exit, Scope, Stream } from "effect"
+import { Headers, HttpBody, HttpClient, HttpClientRequest } from "effect/unstable/http"
+import * as fs from "node:fs"
+import * as os from "node:os"
+import * as path from "node:path"
 import { HttpRecorder } from "../src"
 import { redactedErrorRequest } from "../src/diff"
 
@@ -24,6 +28,18 @@ const runWith = <A, E>(
   effect: Effect.Effect<A, E, HttpClient.HttpClient>,
 ) => Effect.runPromise(effect.pipe(Effect.provide(HttpRecorder.cassetteLayer(name, options))))
 
+const runRecorder = <A, E>(effect: Effect.Effect<A, E, HttpRecorder.Cassette.Service | Scope.Scope>) =>
+  Effect.runPromise(
+    Effect.scoped(
+      effect.pipe(
+        Effect.provide(
+          HttpRecorder.Cassette.layer({ directory: fs.mkdtempSync(path.join(os.tmpdir(), "http-recorder-")) }),
+        ),
+        Effect.provide(NodeFileSystem.layer),
+      ),
+    ),
+  )
+
 const failureText = (exit: Exit.Exit<unknown, unknown>) => {
   if (Exit.isSuccess(exit)) return ""
   return Cause.prettyErrors(exit.cause).join("\n")
@@ -138,6 +154,86 @@ describe("http-recorder", () => {
     expect(HttpRecorder.parseCassette(HttpRecorder.formatCassette(cassette))).toEqual(cassette)
   })
 
+  test("replays websocket interactions from the shared cassette service", async () => {
+    await runRecorder(
+      Effect.gen(function* () {
+        const cassette = yield* HttpRecorder.Cassette.Service
+        yield* cassette.write(
+          "websocket/replay",
+          HttpRecorder.cassetteFor(
+            "websocket/replay",
+            [
+              {
+                transport: "websocket",
+                open: { url: "wss://example.test/realtime", headers: { "content-type": "application/json" } },
+                client: [{ kind: "text", body: JSON.stringify({ type: "response.create" }) }],
+                server: [{ kind: "text", body: JSON.stringify({ type: "response.completed" }) }],
+              },
+            ],
+            undefined,
+          ),
+        )
+        const executor = yield* HttpRecorder.makeWebSocketExecutor({
+          name: "websocket/replay",
+          cassette,
+          compareClientMessagesAsJson: true,
+          live: { open: () => Effect.die(new Error("unexpected live WebSocket open")) },
+        })
+        const connection = yield* executor.open({
+          url: "wss://example.test/realtime",
+          headers: Headers.fromInput({ "content-type": "application/json" }),
+        })
+        yield* connection.sendText(JSON.stringify({ type: "response.create" }))
+        const messages: Array<string | Uint8Array> = []
+        yield* connection.messages.pipe(Stream.runForEach((message) => Effect.sync(() => messages.push(message))))
+        yield* connection.close
+
+        expect(messages).toEqual([JSON.stringify({ type: "response.completed" })])
+      }),
+    )
+  })
+
+  test("records websocket interactions into the shared cassette service", async () => {
+    await runRecorder(
+      Effect.gen(function* () {
+        const cassette = yield* HttpRecorder.Cassette.Service
+        const executor = yield* HttpRecorder.makeWebSocketExecutor({
+          name: "websocket/record",
+          mode: "record",
+          metadata: { provider: "test" },
+          cassette,
+          live: {
+            open: () =>
+              Effect.succeed({
+                sendText: () => Effect.void,
+                messages: Stream.fromIterable([JSON.stringify({ type: "response.completed" })]),
+                close: Effect.void,
+              }),
+          },
+        })
+        const connection = yield* executor.open({
+          url: "wss://example.test/realtime",
+          headers: Headers.fromInput({ "content-type": "application/json" }),
+        })
+        yield* connection.sendText(JSON.stringify({ type: "response.create" }))
+        yield* connection.messages.pipe(Stream.runDrain)
+        yield* connection.close
+
+        expect(yield* cassette.read("websocket/record")).toMatchObject({
+          metadata: { name: "websocket/record", provider: "test" },
+          interactions: [
+            {
+              transport: "websocket",
+              open: { url: "wss://example.test/realtime", headers: { "content-type": "application/json" } },
+              client: [{ kind: "text", body: JSON.stringify({ type: "response.create" }) }],
+              server: [{ kind: "text", body: JSON.stringify({ type: "response.completed" }) }],
+            },
+          ],
+        })
+      }),
+    )
+  })
+
   test("default matcher dispatches multi-interaction cassettes by request shape", async () => {
     await run(
       Effect.gen(function* () {

packages/llm/test/recorded-websocket.ts

@@ -1,157 +1,26 @@
-import { expect } from "bun:test"
-import {
-  Cassette,
-  redactHeaders,
-  redactUrl,
-  isWebSocketInteraction,
-  type WebSocketFrame,
-  type WebSocketInteraction,
-} from "@opencode-ai/http-recorder"
-import { Effect, Layer, Stream } from "effect"
-import type { Headers } from "effect/unstable/http"
+import { Cassette, makeWebSocketExecutor } from "@opencode-ai/http-recorder"
+import { Effect, Layer } from "effect"
 import { WebSocketExecutor } from "../src/route"
-import type { Service as WebSocketExecutorService, WebSocketRequest } from "../src/route/transport/websocket"
+import type { Service as WebSocketExecutorService } from "../src/route/transport/websocket"
 
 const liveWebSocket = WebSocketExecutor.open
-const WEBSOCKET_REQUEST_HEADERS = ["content-type", "accept", "openai-beta"]
-
-const headersRecord = (headers: Headers.Headers) =>
-  Object.fromEntries(
-    Object.entries(headers as Record<string, unknown>)
-      .filter((entry): entry is [string, string] => typeof entry[1] === "string")
-      .toSorted(([a], [b]) => a.localeCompare(b)),
-  )
-
-const openSnapshot = (request: WebSocketRequest) => {
-  const headers = headersRecord(request.headers)
-  return {
-    url: redactUrl(request.url),
-    headers: redactHeaders(headers, WEBSOCKET_REQUEST_HEADERS),
-  }
-}
-
-const textFrame = (body: string): WebSocketFrame => ({ kind: "text", body })
-
-const frameText = (frame: WebSocketFrame) => {
-  if (frame.kind === "text") return frame.body
-  return new TextDecoder().decode(Buffer.from(frame.body, "base64"))
-}
-
-const frameMessage = (frame: WebSocketFrame) =>
-  frame.kind === "text" ? frame.body : new Uint8Array(Buffer.from(frame.body, "base64"))
-
-const receivedFrame = (message: string | Uint8Array): WebSocketFrame =>
-  typeof message === "string"
-    ? textFrame(message)
-    : { kind: "binary", body: Buffer.from(message).toString("base64"), bodyEncoding: "base64" }
-
-const unsafeCassette = (
-  cassette: string,
-  findings: ReadonlyArray<{ readonly path: string; readonly reason: string }>,
-) =>
-  new Error(
-    `Refusing to write WebSocket cassette "${cassette}" because it contains possible secrets: ${findings
-      .map((item) => `${item.path} (${item.reason})`)
-      .join(", ")}`,
-  )
 
 export const webSocketCassetteLayer = (
   cassette: string,
   input: { readonly metadata?: Record<string, unknown>; readonly recording: boolean },
 ): Layer.Layer<WebSocketExecutorService, never, Cassette.Service> =>
-  input.recording ? recordingLayer(cassette, input.metadata) : replayLayer(cassette)
-
-const replayLayer = (cassette: string): Layer.Layer<WebSocketExecutorService, never, Cassette.Service> => {
-  let input: { readonly interactions: ReadonlyArray<WebSocketInteraction> } | undefined
-  let interactionIndex = 0
-  return Layer.effect(
+  Layer.effect(
     WebSocketExecutor.Service,
     Effect.gen(function* () {
       const cassetteService = yield* Cassette.Service
-      yield* Effect.addFinalizer(() =>
-        Effect.sync(() => {
-          if (!input) return
-          expect(interactionIndex, `Unused recorded WebSocket interactions in ${cassette}`).toBe(
-            input.interactions.length,
-          )
-        }),
-      )
-      return WebSocketExecutor.Service.of({
-        open: (request) =>
-          Effect.gen(function* () {
-            input = input ?? {
-              interactions: (yield* cassetteService.read(cassette).pipe(Effect.orDie)).interactions.filter(
-                isWebSocketInteraction,
-              ),
-            }
-            const interaction = input.interactions[interactionIndex]
-            interactionIndex++
-            if (!interaction) throw new Error(`No recorded WebSocket interaction for ${request.url}`)
-            expect(openSnapshot(request)).toEqual(interaction.open)
-            let index = 0
-            return {
-              sendText: (message: string) =>
-                Effect.sync(() => {
-                  expect(JSON.parse(message)).toEqual(
-                    JSON.parse(frameText(interaction.client[index] ?? textFrame("null"))),
-                  )
-                  index++
-                }),
-              messages: Stream.fromIterable(interaction.server).pipe(Stream.map(frameMessage)),
-              close: Effect.sync(() => {
-                expect(index).toBe(interaction.client.length)
-              }),
-            }
-          }),
+      const executor = yield* makeWebSocketExecutor({
+        name: cassette,
+        mode: input.recording ? "record" : "replay",
+        metadata: input.metadata,
+        cassette: cassetteService,
+        live: { open: liveWebSocket },
+        compareClientMessagesAsJson: true,
       })
+      return WebSocketExecutor.Service.of(executor)
     }),
   )
-}
-
-const recordingLayer = (
-  cassette: string,
-  metadata: Record<string, unknown> | undefined,
-): Layer.Layer<WebSocketExecutorService, never, Cassette.Service> => {
-  const webSocket = Layer.effect(
-    WebSocketExecutor.Service,
-    Effect.gen(function* () {
-      const cassetteService = yield* Cassette.Service
-      return WebSocketExecutor.Service.of({
-        open: (request) =>
-          Effect.gen(function* () {
-            const client: WebSocketFrame[] = []
-            const server: WebSocketFrame[] = []
-            const connection = yield* liveWebSocket(request)
-            const decoder = new TextDecoder()
-            return {
-              sendText: (message: string) =>
-                connection.sendText(message).pipe(Effect.tap(() => Effect.sync(() => client.push(textFrame(message))))),
-              messages: connection.messages.pipe(
-                Stream.map((message) => {
-                  const text = WebSocketExecutor.messageText(message, decoder)
-                  server.push(receivedFrame(message))
-                  return text
-                }),
-              ),
-              close: connection.close.pipe(
-                Effect.andThen(
-                  Effect.gen(function* () {
-                    const result = yield* cassetteService
-                      .append(
-                        cassette,
-                        { transport: "websocket", open: openSnapshot(request), client, server },
-                        metadata,
-                      )
-                      .pipe(Effect.orDie)
-                    if (result.findings.length > 0) return yield* Effect.die(unsafeCassette(cassette, result.findings))
-                    return yield* Effect.void
-                  }),
-                ),
-              ),
-            }
-          }),
-      })
-    }),
-  )
-  return webSocket
-}