diff --git a/headroom/memory/backends/direct_mem0.py b/headroom/memory/backends/direct_mem0.py index c086ef012..20c454620 100644 --- a/headroom/memory/backends/direct_mem0.py +++ b/headroom/memory/backends/direct_mem0.py @@ -49,6 +49,7 @@ from __future__ import annotations import asyncio import hashlib +import inspect import logging import uuid from dataclasses import dataclass, field @@ -150,6 +151,7 @@ class DirectMem0Adapter: """ self._config = config or Mem0Config() self._mem0_client: Any = None + self._openai_client: Any = None self._embedder: Any = None self._neo4j_graph: Any = None self._neo4j_driver: Any = None @@ -159,6 +161,7 @@ class DirectMem0Adapter: # Background task tracking self._background_tasks: dict[str, asyncio.Task] = {} self._task_results: dict[str, dict[str, Any]] = {} + self._close_lock = asyncio.Lock() async def _ensure_initialized(self) -> None: """Ensure all clients are initialized.""" @@ -389,7 +392,10 @@ class DirectMem0Adapter: task = self._background_tasks[task_id] try: - await asyncio.wait_for(task, timeout=timeout) + # A timeout must not cancel a save that may be awaiting + # ``asyncio.to_thread``. Cancelling the asyncio task does not stop + # the underlying worker, and would hide that worker from close(). + await asyncio.wait_for(asyncio.shield(task), timeout=timeout) return self.get_task_status(task_id) except asyncio.TimeoutError: return {"status": "timeout", "task_id": task_id} @@ -949,9 +955,71 @@ class DirectMem0Adapter: """Whether this backend supports vector search.""" return True - async def close(self) -> None: - """Close connections and release resources.""" - if self._neo4j_driver: - self._neo4j_driver.close() - self._mem0_client = None - self._initialized = False + async def close(self, timeout: float = 60.0) -> None: + """Drain background writes and close all initialized resources. + + Args: + timeout: Maximum seconds to wait for background writes to finish. + + Raises: + TimeoutError: If background writes have not quiesced within + ``timeout``. Tasks remain tracked and resources remain open so + callers can retry after the writes finish. + """ + # Concurrent shutdown callers must observe one lifecycle transition. + # In particular, a second caller must not detach resources while the + # first is still waiting for executor-backed writes to quiesce. + async with self._close_lock: + if self._background_tasks: + task_items = list(self._background_tasks.items()) + tasks = [task for _, task in task_items] + _, pending = await asyncio.wait(tasks, timeout=timeout) + + if pending: + pending_ids = [task_id for task_id, task in task_items if task in pending] + raise TimeoutError( + "Timed out waiting for DirectMem0 background writes: " + + ", ".join(pending_ids) + ) + + for task_id, task in task_items: + try: + self._task_results[task_id] = { + "status": "completed", + "result": task.result(), + } + except Exception as e: + self._task_results[task_id] = { + "status": "failed", + "error": str(e), + } + self._background_tasks.clear() + + resources = [ + ("Mem0 client", self._mem0_client), + ("OpenAI client", self._openai_client), + ("Qdrant client", self._qdrant_client), + ("Neo4j driver", self._neo4j_driver), + ("embedder", self._embedder), + ("Neo4j graph", self._neo4j_graph), + ] + self._mem0_client = None + self._openai_client = None + self._qdrant_client = None + self._neo4j_driver = None + self._embedder = None + self._neo4j_graph = None + self._initialized = False + + for name, resource in resources: + if resource is None: + continue + close = getattr(resource, "close", None) or getattr(resource, "aclose", None) + if close is None: + continue + try: + result = close() + if inspect.isawaitable(result): + await result + except Exception as e: + logger.warning("Failed to close %s: %s", name, e) diff --git a/tests/test_memory/test_direct_mem0.py b/tests/test_memory/test_direct_mem0.py new file mode 100644 index 000000000..2a00e229a --- /dev/null +++ b/tests/test_memory/test_direct_mem0.py @@ -0,0 +1,184 @@ +"""Tests for the Direct Mem0 adapter lifecycle.""" + +from __future__ import annotations + +import asyncio +import threading +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from headroom.memory.backends.direct_mem0 import DirectMem0Adapter, Mem0Config + + +def _adapter() -> DirectMem0Adapter: + return DirectMem0Adapter(Mem0Config(enable_graph=True)) + + +@pytest.mark.asyncio +async def test_close_drains_tasks_and_closes_initialized_resources() -> None: + adapter = _adapter() + resources = { + "_mem0_client": MagicMock(), + "_openai_client": MagicMock(), + "_qdrant_client": MagicMock(), + "_neo4j_driver": MagicMock(), + } + resources["_mem0_client"].close = AsyncMock() + for name, resource in resources.items(): + setattr(adapter, name, resource) + + task = asyncio.create_task(asyncio.sleep(0, result="saved")) + adapter._background_tasks["task_1"] = task + + await adapter.close(timeout=1.0) + + assert adapter.get_pending_tasks() == [] + assert adapter.get_task_status("task_1") == { + "status": "completed", + "result": "saved", + } + resources["_mem0_client"].close.assert_awaited_once_with() + for resource in resources.values(): + resource.close.assert_called_once_with() + assert adapter._initialized is False + assert adapter._mem0_client is None + assert adapter._openai_client is None + assert adapter._qdrant_client is None + assert adapter._neo4j_driver is None + + await adapter.close(timeout=0.01) + resources["_mem0_client"].close.assert_awaited_once_with() + for resource in resources.values(): + resource.close.assert_called_once_with() + + +@pytest.mark.asyncio +async def test_close_timeout_keeps_tasks_and_resources_attached() -> None: + adapter = _adapter() + release = asyncio.Event() + resource = MagicMock() + adapter._mem0_client = resource + task = asyncio.create_task(release.wait()) + adapter._background_tasks["task_1"] = task + await asyncio.sleep(0) + + with pytest.raises(TimeoutError, match="task_1"): + await adapter.close(timeout=0.01) + + assert not task.done() + assert adapter.get_pending_tasks() == ["task_1"] + assert adapter._mem0_client is resource + resource.close.assert_not_called() + + release.set() + await adapter.close(timeout=1.0) + + assert adapter.get_pending_tasks() == [] + assert adapter.get_task_status("task_1") == { + "status": "completed", + "result": True, + } + resource.close.assert_called_once_with() + + +@pytest.mark.asyncio +async def test_close_does_not_close_resources_while_sync_worker_is_running() -> None: + adapter = _adapter() + worker_started = threading.Event() + release_worker = threading.Event() + + class BlockingMem0Client: + def __init__(self) -> None: + self.close_calls = 0 + + def add(self, *_args: object, **_kwargs: object) -> dict[str, list[dict[str, str]]]: + worker_started.set() + assert release_worker.wait(timeout=5.0), "test did not release worker" + return {"results": [{"id": "memory-1", "memory": "saved"}]} + + def close(self) -> None: + self.close_calls += 1 + + client = BlockingMem0Client() + adapter._mem0_client = client + adapter._initialized = True + + memory = await adapter.save_memory( + content="saved", + user_id="user-1", + importance=0.5, + background=True, + ) + task_id = memory.metadata["_task_id"] + + for _ in range(100): + if worker_started.is_set(): + break + await asyncio.sleep(0.01) + assert worker_started.is_set() + + with pytest.raises(TimeoutError, match=task_id): + await adapter.close(timeout=0.01) + + assert client.close_calls == 0 + assert adapter._mem0_client is client + assert adapter.get_pending_tasks() == [task_id] + + release_worker.set() + await adapter.close(timeout=1.0) + + assert client.close_calls == 1 + assert adapter._mem0_client is None + assert adapter.get_pending_tasks() == [] + assert adapter.get_task_status(task_id)["status"] == "completed" + + +@pytest.mark.asyncio +async def test_concurrent_close_calls_serialize_resource_cleanup() -> None: + adapter = _adapter() + close_started = asyncio.Event() + release_close = asyncio.Event() + resource = MagicMock() + + async def slow_close() -> None: + close_started.set() + await release_close.wait() + + resource.close = AsyncMock(side_effect=slow_close) + adapter._mem0_client = resource + + first = asyncio.create_task(adapter.close()) + await close_started.wait() + second = asyncio.create_task(adapter.close()) + await asyncio.sleep(0) + + resource.close.assert_awaited_once_with() + assert not first.done() + assert not second.done() + + release_close.set() + await asyncio.gather(first, second) + + resource.close.assert_awaited_once_with() + + +@pytest.mark.asyncio +async def test_wait_for_task_timeout_does_not_cancel_background_write() -> None: + adapter = _adapter() + release = asyncio.Event() + task = asyncio.create_task(release.wait()) + adapter._background_tasks["task_1"] = task + + assert await adapter.wait_for_task("task_1", timeout=0.01) == { + "status": "timeout", + "task_id": "task_1", + } + assert not task.done() + assert adapter.get_pending_tasks() == ["task_1"] + + release.set() + assert await adapter.wait_for_task("task_1", timeout=1.0) == { + "status": "completed", + "result": True, + }