diff --git a/haiku_rag_slim/haiku/rag/ingester/app.py b/haiku_rag_slim/haiku/rag/ingester/app.py index 7b64556f..70c5a673 100644 --- a/haiku_rag_slim/haiku/rag/ingester/app.py +++ b/haiku_rag_slim/haiku/rag/ingester/app.py @@ -181,6 +181,7 @@ class IngesterApp: await asyncio.gather(api_task, return_exceptions=True) await self._pollers.stop() await self._stop_pool() + await self._pollers.close_sources() async def run_batch(self) -> BatchReport: """Run one discover() sweep across every configured source, drain the @@ -213,6 +214,7 @@ class IngesterApp: ) finally: await self._stop_pool() + await self._pollers.close_sources() async def _maybe_start_api(self, api: bool): """Spin up the FastAPI control plane on an asyncio task. Returns diff --git a/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py b/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py index bd7638d5..0e390693 100644 --- a/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py +++ b/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py @@ -96,9 +96,13 @@ class PollerManager: await asyncio.gather(*self._tasks, return_exceptions=True) self._tasks.clear() self._started = False + + async def close_sources(self) -> None: + """Close all source adapters (e.g. HTTP connection pools). Must be + called after the worker pool has fully stopped so in-flight fetches + don't hit a closed client.""" for source in self.sources: - if hasattr(source, "aclose"): - await source.aclose() + await source.aclose() @property def pollers(self) -> list[BasePoller]: diff --git a/haiku_rag_slim/haiku/rag/ingester/sources/base.py b/haiku_rag_slim/haiku/rag/ingester/sources/base.py index 32099b35..c282ece9 100644 --- a/haiku_rag_slim/haiku/rag/ingester/sources/base.py +++ b/haiku_rag_slim/haiku/rag/ingester/sources/base.py @@ -61,6 +61,11 @@ class Source(Protocol): """ ... + async def aclose(self) -> None: + """Release any resources held by the source (e.g. HTTP connection + pools). Called once during shutdown, after all workers have stopped.""" + ... + async def fetch(self, uri: str) -> FetchResult: ... def discover( diff --git a/haiku_rag_slim/haiku/rag/ingester/sources/fs.py b/haiku_rag_slim/haiku/rag/ingester/sources/fs.py index 8b944e04..ba97f120 100644 --- a/haiku_rag_slim/haiku/rag/ingester/sources/fs.py +++ b/haiku_rag_slim/haiku/rag/ingester/sources/fs.py @@ -75,6 +75,9 @@ class FSSource: return False return self._resolve_within_root(uri) is not None + async def aclose(self) -> None: # pragma: no cover - no resources to release + pass + async def head(self, uri: str) -> str | None: path = self._resolve_within_root(uri) if path is None or not path.exists(): diff --git a/haiku_rag_slim/haiku/rag/ingester/sources/s3.py b/haiku_rag_slim/haiku/rag/ingester/sources/s3.py index 0b0b229f..f3c8a446 100644 --- a/haiku_rag_slim/haiku/rag/ingester/sources/s3.py +++ b/haiku_rag_slim/haiku/rag/ingester/sources/s3.py @@ -63,6 +63,9 @@ class S3Source: def supports(self, uri: str) -> bool: return uri.startswith(self.uri_prefix) + async def aclose(self) -> None: # pragma: no cover - no resources to release + pass + async def head(self, uri: str) -> str | None: import obstore # type: ignore[import-not-found] diff --git a/tests/ingester/test_sources_base.py b/tests/ingester/test_sources_base.py index b2ffe995..ebaaebae 100644 --- a/tests/ingester/test_sources_base.py +++ b/tests/ingester/test_sources_base.py @@ -59,6 +59,9 @@ def test_source_protocol_runtime_checkable(): def supports(self, uri: str) -> bool: return True + async def aclose(self) -> None: + pass + async def head(self, uri: str): return None