diff --git a/crawl4ai/async_webcrawler.py b/crawl4ai/async_webcrawler.py index a2aabcc41..e89abc9bb 100644 --- a/crawl4ai/async_webcrawler.py +++ b/crawl4ai/async_webcrawler.py @@ -1108,14 +1108,19 @@ async def maybe_release_session(): if stream: async def result_transformer(): + inner = 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 inner: yield transform_result(task_result) finally: - # Auto-release session after streaming completes - await maybe_release_session() + try: + await inner.aclose() + finally: + # Auto-release only after dispatcher cleanup completes, + # even if closing the inner stream raises. + await maybe_release_session() return result_transformer() else: @@ -1247,4 +1252,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..7136c9511 100644 --- a/tests/async/test_dispatchers.py +++ b/tests/async/test_dispatchers.py @@ -105,6 +105,47 @@ async def crawl_url(self, url, config, task_id, retry_count=0): assert dispatcher.concurrent_sessions == 0 assert dispatcher.task_queue.empty() + async def test_arun_many_stream_closure_cleans_up_dispatcher_tasks(self): + 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) + + async def controlled_arun(url, config=None, session_id=None): + if url == "fast": + return SimpleNamespace( + url=url, + success=True, + status_code=200, + error_message="", + ) + await asyncio.Event().wait() + + crawler = AsyncWebCrawler() + crawler.arun = controlled_arun + dispatcher = TrackingDispatcher() + stream = await crawler.arun_many( + ["fast", "blocked-1", "blocked-2"], + config=CrawlerRunConfig(stream=True), + dispatcher=dispatcher, + ) + + 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)