From 471ee2a2559aa3bd243c6ead2a30016711fb4a18 Mon Sep 17 00:00:00 2001
From: Dinh Le <dinwwwh@gmail.com>
Date: Sun, 2 Aug 2026 15:59:47 +0700
Subject: [PATCH] fix(peer): finalize hibernation iterators instead of throwing
 on cleanup

HibernationAsyncIteratorClass no longer throws when its cleanup runs
(return/throw/asyncDispose), and ServerPeer now finalizes the iterator
right after invoking the hibernation callback, even when the callback
rejects.
---
 packages/peer/src/hibernation.test.ts | 4 ++--
 packages/peer/src/hibernation.ts      | 6 ++----
 packages/peer/src/server.test.ts      | 2 ++
 packages/peer/src/server.ts           | 7 ++++++-
 4 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/packages/peer/src/hibernation.test.ts b/packages/peer/src/hibernation.test.ts
index 46b9f0b..1268ff6 100644
--- a/packages/peer/src/hibernation.test.ts
+++ b/packages/peer/src/hibernation.test.ts
@@ -12,9 +12,9 @@ describe('hibernationAsyncIteratorClass', () => {
     await expect(iterator.next()).rejects.toThrow('Cannot use hibernating iterator directly')
   })
 
-  it('return() throws', async () => {
+  it('return() does not throw', async () => {
     const iterator = new HibernationAsyncIteratorClass(vi.fn())
-    await expect(iterator.return()).rejects.toThrow('Cannot use hibernating iterator directly')
+    await expect(iterator.return()).resolves.toEqual({ done: true, value: undefined })
   })
 
   it('invokes callback with correct id', () => {
diff --git a/packages/peer/src/hibernation.ts b/packages/peer/src/hibernation.ts
index 47e15af..1436f25 100644
--- a/packages/peer/src/hibernation.ts
+++ b/packages/peer/src/hibernation.ts
@@ -16,10 +16,8 @@ export class HibernationAsyncIteratorClass<T, TReturn = unknown, TNext = unknown
   ) {
     super(async () => {
       throw new Error('Cannot use hibernating iterator directly')
-    }, async ({ kind }) => {
-      if (kind === 'cancelled') {
-        throw new Error('Cannot use hibernating iterator directly')
-      }
+    }, async () => {
+      // nothing to clean up
     })
 
     this['~callback'] = callback
diff --git a/packages/peer/src/server.test.ts b/packages/peer/src/server.test.ts
index 4d51037..0f926b0 100644
--- a/packages/peer/src/server.test.ts
+++ b/packages/peer/src/server.test.ts
@@ -324,6 +324,7 @@ describe('serverPeer', () => {
         expect(callback).toHaveBeenCalledWith('1')
         expect(send).toHaveBeenCalledTimes(1)
         expect(send).toHaveBeenNthCalledWith(1, expect.objectContaining({ kind: 'response' }))
+        await expect(hibernationIter.next()).resolves.toEqual({ done: true, value: undefined }) // already cleaned up
       })
 
       it('reject if HibernationAsyncIteratorClassCallback reject', async () => {
@@ -340,6 +341,7 @@ describe('serverPeer', () => {
         expect(send).toHaveBeenCalledTimes(2)
         expect(send).toHaveBeenNthCalledWith(1, expect.objectContaining({ kind: 'response' }))
         expect(send).toHaveBeenNthCalledWith(2, expect.objectContaining({ kind: 'cancel' }))
+        await expect(hibernationIter.next()).resolves.toEqual({ done: true, value: undefined }) // already cleaned up
       })
 
       it('cancels active transmitter on close', async () => {
diff --git a/packages/peer/src/server.ts b/packages/peer/src/server.ts
index ac05d10..e24df9e 100644
--- a/packages/peer/src/server.ts
+++ b/packages/peer/src/server.ts
@@ -116,7 +116,12 @@ export class ServerPeer {
 
       if (isAsyncIteratorObject(response.body)) {
         if (response.body instanceof HibernationAsyncIteratorClass) {
-          await response.body['~callback']?.(id)
+          try {
+            await response.body['~callback']?.(id)
+          }
+          finally {
+            await response.body.return()
+          }
         }
         else {
           const transmitter = new EventStreamTransmitter(response.body, id, this.send)