From 9fb2aa8d7ac21239210a2d0a639ad6a2c1deb66c Mon Sep 17 00:00:00 2001
From: NhatNguyen <ah2631@student.jamk.fi>
Date: Sun, 12 Jul 2026 23:17:43 +0700
Subject: [PATCH] fix(dispatcher): clean up closed stream tasks

---
 crawl4ai/async_dispatcher.py    | 22 ++++++++++++++---
 tests/async/test_dispatchers.py | 44 ++++++++++++++++++++++++++++++++-
 2 files changed, 62 insertions(+), 4 deletions(-)

diff --git a/crawl4ai/async_dispatcher.py b/crawl4ai/async_dispatcher.py
index 1d6a236b7..ebbe4cb92 100644
--- a/crawl4ai/async_dispatcher.py
+++ b/crawl4ai/async_dispatcher.py
@@ -534,6 +534,7 @@ async def run_urls_stream(
         config: Union[CrawlerRunConfig, List[CrawlerRunConfig]],
     ) -> AsyncGenerator[CrawlerTaskResult, None]:
         self.crawler = crawler
+        active_tasks = []
         
         # Start the memory monitor task
         memory_monitor = asyncio.create_task(self._memory_monitor_task())
@@ -550,7 +551,6 @@ async def run_urls_stream(
                 # Add to queue with initial priority 0, retry count 0, and current time
                 await self.task_queue.put((0, (url, task_id, 0, time.time())))
                 
-            active_tasks = []
             completed_count = 0
             total_urls = len(urls)
 
@@ -614,8 +614,24 @@ async def run_urls_stream(
                 await self._update_queue_priorities()
                 
         finally:
-            # Clean up
+            # Cancel and await every task owned by this stream before returning
+            # control to the caller. Otherwise a closed stream can leave crawls
+            # using browser pages and contexts in the background.
+            for task in active_tasks:
+                if not task.done():
+                    task.cancel()
+            if active_tasks:
+                await asyncio.gather(*active_tasks, return_exceptions=True)
+
+            # Discard URLs that were queued by this stream but never started.
+            while True:
+                try:
+                    self.task_queue.get_nowait()
+                except asyncio.QueueEmpty:
+                    break
+
             memory_monitor.cancel()
+            await asyncio.gather(memory_monitor, return_exceptions=True)
             if self.monitor:
                 self.monitor.stop()
                 
@@ -769,4 +785,4 @@ async def run_urls(
             return await asyncio.gather(*tasks, return_exceptions=True)
         finally:
             if self.monitor:
-                self.monitor.stop()
\ No newline at end of file
+                self.monitor.stop()
diff --git a/tests/async/test_dispatchers.py b/tests/async/test_dispatchers.py
index 99cf4a989..0ed9fad10 100644
--- a/tests/async/test_dispatchers.py
+++ b/tests/async/test_dispatchers.py
@@ -1,5 +1,8 @@
-import pytest
+import asyncio
 import time
+from types import SimpleNamespace
+
+import pytest
 from crawl4ai import (
     AsyncWebCrawler,
     BrowserConfig,
@@ -63,6 +66,45 @@ 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 arun(self, url, config=None, session_id=None):
+                if url == "fast":
+                    return SimpleNamespace(
+                        success=True,
+                        status_code=200,
+                        error_message="",
+                    )
+                await asyncio.Event().wait()
+
+        class TrackingDispatcher(MemoryAdaptiveDispatcher):
+            def __init__(self):
+                super().__init__(max_session_permit=3)
+                self.tasks = {}
+
+            async def crawl_url(self, url, config, task_id, retry_count=0):
+                self.tasks[url] = asyncio.current_task()
+                return await super().crawl_url(url, config, task_id, retry_count)
+
+        dispatcher = TrackingDispatcher()
+        stream = dispatcher.run_urls_stream(
+            ["fast", "blocked-1", "blocked-2", "queued"],
+            BlockingCrawler(),
+            run_config,
+        )
+
+        first = await asyncio.wait_for(stream.__anext__(), timeout=1)
+        assert first.url == "fast"
+
+        await asyncio.wait_for(stream.aclose(), timeout=1)
+
+        assert set(dispatcher.tasks) == {"fast", "blocked-1", "blocked-2"}
+        assert all(task.done() for task in dispatcher.tasks.values())
+        assert dispatcher.tasks["blocked-1"].cancelled()
+        assert dispatcher.tasks["blocked-2"].cancelled()
+        assert dispatcher.concurrent_sessions == 0
+        assert dispatcher.task_queue.empty()
+
     async def test_semaphore_basic(self, browser_config, run_config, test_urls):
         async with AsyncWebCrawler(config=browser_config) as crawler:
             dispatcher = SemaphoreDispatcher(semaphore_count=2)