EFF-717 OpenAiSocket should cancel the request on interrupt (#1759)

Author: tim-smart

.changeset/eff-717-openai-socket-cancel.md

@@ -0,0 +1,5 @@
+---
+"@effect/ai-openai": patch
+---
+
+Ensure OpenAiSocket sends a `{"type":"response.cancel"}` websocket event when a response stream is interrupted.

packages/ai/openai/src/OpenAiClient.ts

@@ -10,6 +10,7 @@ import * as Array from "effect/Array"
 import * as Cause from "effect/Cause"
 import type * as Config from "effect/Config"
 import * as Effect from "effect/Effect"
+import * as Exit from "effect/Exit"
 import { identity } from "effect/Function"
 import * as Function from "effect/Function"
 import * as Layer from "effect/Layer"
@@ -420,6 +421,10 @@ const makeSocket = Effect.gen(function*() {
           )
         )
 
+      const cancel = Effect.suspend(() => write(JSON.stringify({ type: "response.cancel" }))).pipe(
+        Effect.ignore
+      )
+
       const decoder = new TextDecoder()
       const decode = Schema.decodeUnknownSync(Schema.fromJsonString(Generated.ResponseStreamEvent))
       yield* socket.runRaw((msg) => {
@@ -462,7 +467,7 @@ const makeSocket = Effect.gen(function*() {
         Effect.forkScoped
       )
 
-      return { send } as const
+      return { send, cancel } as const
     })
   })
 
@@ -476,11 +481,12 @@ const makeSocket = Effect.gen(function*() {
           semaphore.take(1),
           () => semaphore.release(1)
         )
-        const { send } = yield* RcRef.get(queueRef)
+        const { send, cancel } = yield* RcRef.get(queueRef)
         const incoming = yield* Queue.unbounded<ResponseStreamEvent, AiError.AiError | Cause.Done>()
         yield* Effect.forkScoped(send(incoming, options), { startImmediately: true })
         return Stream.fromQueue(incoming).pipe(
-          Stream.takeUntil((e) => e.type === "response.completed" || e.type === "response.incomplete")
+          Stream.takeUntil((e) => e.type === "response.completed" || e.type === "response.incomplete"),
+          Stream.onExit((exit) => Exit.hasInterrupts(exit) ? cancel : Effect.void)
         )
       }).pipe(Stream.unwrap)
 

packages/ai/openai/test/OpenAiClient.test.ts

@@ -3,11 +3,14 @@ import * as Errors from "@effect/ai-openai/internal/errors"
 import * as OpenAiClient from "@effect/ai-openai/OpenAiClient"
 import { assert, describe, it } from "@effect/vitest"
 import { Config, ConfigProvider, Effect, Layer, Redacted, Schema, Stream } from "effect"
+import * as Fiber from "effect/Fiber"
 import type * as AiError from "effect/unstable/ai/AiError"
 import * as HttpClient from "effect/unstable/http/HttpClient"
 import * as HttpClientError from "effect/unstable/http/HttpClientError"
 import * as HttpClientRequest from "effect/unstable/http/HttpClientRequest"
 import * as HttpClientResponse from "effect/unstable/http/HttpClientResponse"
+import * as Socket from "effect/unstable/socket/Socket"
+import { WS } from "vitest-websocket-mock"
 
 // =============================================================================
 // Mock Helpers
@@ -631,5 +634,56 @@ describe("OpenAiClient", () => {
         assert.strictEqual(result.reason._tag, "InternalProviderError")
       }).pipe(Effect.provide(MainLayer))
     })
+
+    it.effect("sends response.cancel on interrupt in websocket mode", () => {
+      const port = 42345
+      const apiUrl = `http://localhost:${port}/v1`
+      const serverUrl = `ws://localhost:${port}/v1/responses`
+
+      const HttpClientLayer = Layer.succeed(
+        HttpClient.HttpClient,
+        makeMockHttpClient((request) => Effect.succeed(makeMockResponse({ status: 200, body: {}, request })))
+      )
+
+      const MainLayer = OpenAiClient.layer({
+        apiKey: Redacted.make("test-key"),
+        apiUrl
+      }).pipe(Layer.provide(HttpClientLayer))
+
+      return Effect.gen(function*() {
+        const server = yield* Effect.acquireRelease(
+          Effect.sync(() => new WS(serverUrl, { jsonProtocol: true })),
+          (server) =>
+            Effect.sync(() => {
+              server.close()
+              WS.clean()
+            })
+        )
+
+        const client = yield* OpenAiClient.OpenAiClient
+        const fiber = yield* Effect.forkScoped(
+          OpenAiClient.withWebSocketMode(
+            client.createResponseStream({
+              model: "gpt-4o",
+              input: "test"
+            }).pipe(
+              Effect.andThen(([_, stream]) => Stream.runDrain(stream))
+            )
+          ),
+          { startImmediately: true }
+        )
+
+        const createEvent = yield* Effect.promise(() => server.nextMessage as Promise<any>)
+        assert.strictEqual(createEvent.type, "response.create")
+
+        yield* Fiber.interrupt(fiber)
+
+        const cancelEvent = yield* Effect.promise(() => server.nextMessage as Promise<any>)
+        assert.deepStrictEqual(cancelEvent, { type: "response.cancel" })
+      }).pipe(
+        Effect.provide(MainLayer),
+        Effect.provideService(Socket.WebSocketConstructor, (url) => new globalThis.WebSocket(url))
+      )
+    })
   })
 })