fix: preserve websocket turn cancellation hooks

Author: lidge-jun

src/ws-bridge.ts

@@ -140,9 +140,10 @@ export async function pumpResponsesSseToWebSocket(
 ): Promise<void> {
   const reader = sseStream.getReader();
   const isCurrent = options.isCurrent ?? (() => true);
-  ws.data.cancel = () => {
+  const cancel = () => {
     void reader.cancel().catch(() => {});
   };
+  ws.data.cancel = cancel;
 
   const decoder = new TextDecoder();
   let buffer = "";
@@ -154,6 +155,8 @@ export async function pumpResponsesSseToWebSocket(
     const type = payloadType(payload);
     if (!type) {
       sendProtocolError(ws, 502, "Invalid JSON payload in upstream SSE frame");
+      terminalSeen = true;
+      void reader.cancel().catch(() => {});
       return true;
     }
     if (terminalSeen) return true;
@@ -191,7 +194,7 @@ export async function pumpResponsesSseToWebSocket(
       sendProtocolError(ws, 502, err instanceof Error ? err.message : String(err));
     }
   } finally {
-    if (ws.data.cancel) ws.data.cancel = undefined;
+    if (ws.data.cancel === cancel) ws.data.cancel = undefined;
   }
 }
 

tests/ws-endpoint.test.ts

@@ -131,6 +131,34 @@ describe("WS endpoint re-framer (120/132)", () => {
     expect(sent).toEqual([]);
   });
 
+  test("stale pump cleanup does not erase the replacement turn cancel hook", async () => {
+    const { ws } = mockWs();
+    let current = false;
+    const stalePump = pumpResponsesSseToWebSocket(ws, sseStream([
+      'event: response.created\ndata: {"type":"response.created"}\n\n',
+    ]), { isCurrent: () => current });
+    const replacementCancel = () => {};
+    ws.data.cancel = replacementCancel;
+    await stalePump;
+    expect(ws.data.cancel).toBe(replacementCancel);
+  });
+
+  test("invalid upstream SSE JSON emits one standalone protocol error and cancels", async () => {
+    const { ws, sent } = mockWs();
+    let cancelled = false;
+    await pumpResponsesSseToWebSocket(ws, sseStream([
+      "event: response.created\ndata: {not-json}\n\n",
+      'event: response.completed\ndata: {"type":"response.completed"}\n\n',
+    ], () => { cancelled = true; }));
+    expect(sent).toHaveLength(1);
+    expect(JSON.parse(sent[0])).toMatchObject({
+      type: "error",
+      status: 502,
+      error: { code: "websocket_protocol_error" },
+    });
+    expect(cancelled).toBe(true);
+  });
+
   test("converts successful Responses JSON into output_item.done plus response.completed frames", () => {
     const { ws, sent } = mockWs();
     sendResponsesJsonAsEvents(ws, {