From dba915858676646df4d75b65a6fcbb72cc2ffb1f Mon Sep 17 00:00:00 2001
From: nightcityblade <nightcityblade@gmail.com>
Date: Thu, 16 Jul 2026 23:32:27 +0800
Subject: [PATCH] fix: await dispatcher stream cleanup

---
 crawl4ai/async_webcrawler.py    | 16 ++++++++++------
 tests/async/test_dispatchers.py | 15 ++++++++++-----
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/crawl4ai/async_webcrawler.py b/crawl4ai/async_webcrawler.py
index 8216d19bc..233a66a92 100644
--- a/crawl4ai/async_webcrawler.py
+++ b/crawl4ai/async_webcrawler.py
@@ -1107,14 +1107,18 @@ async def maybe_release_session():
 
         if stream:
             async def result_transformer():
+                dispatcher_stream = dispatcher.run_urls_stream(
+                    crawler=self, urls=urls, config=config
+                )
                 try:
-                    async for task_result in dispatcher.run_urls_stream(
-                        crawler=self, urls=urls, config=config
-                    ):
+                    async for task_result in dispatcher_stream:
                         yield transform_result(task_result)
                 finally:
-                    # Auto-release session after streaming completes
-                    await maybe_release_session()
+                    try:
+                        await dispatcher_stream.aclose()
+                    finally:
+                        # Auto-release session after streaming completes
+                        await maybe_release_session()
 
             return result_transformer()
         else:
@@ -1246,4 +1250,4 @@ async def amap_domain(
             config or DomainMapperConfig(**kwargs) if kwargs else DomainMapperConfig()
         )
 
-        return await self._domain_mapper.scan(domain, mapper_config)
\ No newline at end of file
+        return await self._domain_mapper.scan(domain, mapper_config)
diff --git a/tests/async/test_dispatchers.py b/tests/async/test_dispatchers.py
index 0ed9fad10..a7c7f4ad2 100644
--- a/tests/async/test_dispatchers.py
+++ b/tests/async/test_dispatchers.py
@@ -66,11 +66,15 @@ async def test_memory_adaptive_with_rate_limit(
             assert len(results) == len(test_urls)
             assert all(r.success for r in results)
 
-    async def test_memory_adaptive_stream_closure_cleans_up_tasks(self, run_config):
-        class BlockingCrawler:
+    async def test_arun_many_stream_closure_cleans_up_tasks(self, run_config):
+        class BlockingCrawler(AsyncWebCrawler):
+            def __init__(self):
+                pass
+
             async def arun(self, url, config=None, session_id=None):
                 if url == "fast":
                     return SimpleNamespace(
+                        url=url,
                         success=True,
                         status_code=200,
                         error_message="",
@@ -87,10 +91,11 @@ async def crawl_url(self, url, config, task_id, retry_count=0):
                 return await super().crawl_url(url, config, task_id, retry_count)
 
         dispatcher = TrackingDispatcher()
-        stream = dispatcher.run_urls_stream(
+        run_config.stream = True
+        stream = await BlockingCrawler().arun_many(
             ["fast", "blocked-1", "blocked-2", "queued"],
-            BlockingCrawler(),
-            run_config,
+            config=run_config,
+            dispatcher=dispatcher,
         )
 
         first = await asyncio.wait_for(stream.__anext__(), timeout=1)