core: streaming: cancel inner task when client disconnects

Author: nicoschmdt

core/libs/commonwealth/src/commonwealth/utils/streaming.py

@@ -73,15 +73,24 @@ async def generator_wrapper(gen: AsyncGenerator[str | bytes, None], queue: async
         finally:
             if heartbeat_task:
                 heartbeat_task.cancel()
+            await gen.aclose()
             await queue.put(None)
 
-    asyncio.create_task(generator_wrapper(gen, queue))
+    wrapper_task = asyncio.create_task(generator_wrapper(gen, queue))
 
-    while True:
-        item = await queue.get()
-        if item is None:
-            break
-        yield item
+    try:
+        while True:
+            item = await queue.get()
+            if item is None:
+                break
+            yield item
+    finally:
+        if not wrapper_task.done():
+            wrapper_task.cancel()
+            try:
+                await wrapper_task
+            except asyncio.CancelledError:
+                pass
 
 
 async def _fetch_stream(
@@ -93,6 +102,7 @@ async def _fetch_stream(
     except Exception as e:
         await queue.put((None, e))
     finally:
+        await gen.aclose()
         await queue.put((None, None))