fix(mcp): preserve ASGI test database state

Signed-off-by: phernandez <paul@basicmachines.co>

Author: phernandez

src/basic_memory/mcp/async_client.py

@@ -1,6 +1,7 @@
 import os
 from contextlib import AbstractAsyncContextManager, asynccontextmanager
-from typing import AsyncIterator, Callable, Optional
+from inspect import isawaitable
+from typing import TYPE_CHECKING, AsyncIterator, Callable, Optional, cast
 
 from fastapi import FastAPI
 from httpx import ASGITransport, AsyncClient, Timeout
@@ -9,6 +10,12 @@
 import logfire
 from basic_memory.config import ConfigManager, ProjectMode
 
+if TYPE_CHECKING:
+    from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
+
+LocalDatabaseState = tuple["AsyncEngine", "async_sessionmaker[AsyncSession]"]
+_MISSING_STATE_VALUE = object()
+
 
 def _force_local_mode() -> bool:
     """Check if local mode is forced via environment variable."""
@@ -49,15 +56,48 @@ def _build_asgi_client(app: FastAPI, timeout: Timeout) -> AsyncClient:
     )
 
 
-async def _prepare_local_asgi_database(app: FastAPI) -> None:
-    """Initialize local ASGI database state before the first request."""
+async def _resolve_local_asgi_database(app: FastAPI) -> LocalDatabaseState:
+    """Resolve database state for a local ASGI request."""
+    from basic_memory.deps import get_engine_factory
+
+    override = app.dependency_overrides.get(get_engine_factory)
+    if override is not None:
+        result = override()
+        if isawaitable(result):
+            result = await result
+        return cast(LocalDatabaseState, result)
+
     from basic_memory import db
 
     config = ConfigManager().config
-    engine, session_maker = await db.get_or_create_db(config.database_path)
+    return await db.get_or_create_db(config.database_path)
+
+
+@asynccontextmanager
+async def _prepared_local_asgi_database(app: FastAPI) -> AsyncIterator[None]:
+    """Initialize local ASGI database state before the first request."""
+    previous_engine = getattr(app.state, "engine", _MISSING_STATE_VALUE)
+    previous_session_maker = getattr(app.state, "session_maker", _MISSING_STATE_VALUE)
+
+    engine, session_maker = await _resolve_local_asgi_database(app)
     app.state.engine = engine
     app.state.session_maker = session_maker
 
+    try:
+        yield
+    finally:
+        if previous_engine is _MISSING_STATE_VALUE:
+            if hasattr(app.state, "engine"):
+                delattr(app.state, "engine")
+        else:
+            app.state.engine = previous_engine
+
+        if previous_session_maker is _MISSING_STATE_VALUE:
+            if hasattr(app.state, "session_maker"):
+                delattr(app.state, "session_maker")
+        else:
+            app.state.session_maker = previous_session_maker
+
 
 @asynccontextmanager
 async def _asgi_client(timeout: Timeout) -> AsyncIterator[AsyncClient]:
@@ -71,9 +111,9 @@ async def _asgi_client(timeout: Timeout) -> AsyncIterator[AsyncClient]:
     # under Starlette's request loop and trigger CPython's empty-ready-queue race.
     # Outcome: request handling sees the same app.state database objects as API
     # lifespan startup would have provided.
-    await _prepare_local_asgi_database(fastapi_app)
-    async with _build_asgi_client(fastapi_app, timeout) as client:
-        yield client
+    async with _prepared_local_asgi_database(fastapi_app):
+        async with _build_asgi_client(fastapi_app, timeout) as client:
+            yield client
 
 
 async def _resolve_cloud_token(config) -> str:

tests/mcp/test_async_client_modes.py

@@ -81,6 +81,8 @@ async def fake_get_or_create_db(db_path):
             assert calls == [cfg.database_path]
             assert fastapi_app.state.engine is engine
             assert fastapi_app.state.session_maker is session_maker
+        assert not hasattr(fastapi_app.state, "engine")
+        assert not hasattr(fastapi_app.state, "session_maker")
     finally:
         if previous_engine is None:
             fastapi_app.state._state.pop("engine", None)  # pyright: ignore[reportPrivateUsage]
@@ -92,6 +94,51 @@ async def fake_get_or_create_db(db_path):
             fastapi_app.state.session_maker = previous_session_maker
 
 
+@pytest.mark.asyncio
+async def test_get_client_uses_existing_local_asgi_database_override(config_manager):
+    """Local ASGI routing honors FastAPI test dependency overrides."""
+    from basic_memory.api.app import app as fastapi_app
+    from basic_memory.deps import get_engine_factory
+
+    cfg = config_manager.load_config()
+    config_manager.save_config(cfg)
+
+    previous_overrides = dict(fastapi_app.dependency_overrides)
+    previous_engine = getattr(fastapi_app.state, "engine", None)
+    previous_session_maker = getattr(fastapi_app.state, "session_maker", None)
+    fastapi_app.state._state.pop("engine", None)  # pyright: ignore[reportPrivateUsage]
+    fastapi_app.state._state.pop("session_maker", None)  # pyright: ignore[reportPrivateUsage]
+
+    engine = object()
+    session_maker = object()
+    calls = []
+
+    def override_engine_factory():
+        calls.append("override")
+        return engine, session_maker
+
+    fastapi_app.dependency_overrides[get_engine_factory] = override_engine_factory
+
+    try:
+        async with get_client() as client:
+            assert isinstance(client._transport, httpx.ASGITransport)  # pyright: ignore[reportPrivateUsage]
+            assert calls == ["override"]
+            assert fastapi_app.state.engine is engine
+            assert fastapi_app.state.session_maker is session_maker
+        assert not hasattr(fastapi_app.state, "engine")
+        assert not hasattr(fastapi_app.state, "session_maker")
+    finally:
+        fastapi_app.dependency_overrides = previous_overrides
+        if previous_engine is None:
+            fastapi_app.state._state.pop("engine", None)  # pyright: ignore[reportPrivateUsage]
+        else:
+            fastapi_app.state.engine = previous_engine
+        if previous_session_maker is None:
+            fastapi_app.state._state.pop("session_maker", None)  # pyright: ignore[reportPrivateUsage]
+        else:
+            fastapi_app.state.session_maker = previous_session_maker
+
+
 @pytest.mark.asyncio
 async def test_get_client_explicit_cloud_uses_api_key(config_manager, monkeypatch):
     cfg = config_manager.load_config()