fix: await background sync task cancellation in lifespan shutdown (#456)

Signed-off-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com>
Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com>
Co-authored-by: Paul Hernandez <phernandez@users.noreply.github.com>
This commit is contained in:
Paul Hernandez
2025-12-17 09:53:34 -06:00
committed by GitHub
parent a0f20eb102
commit efbc758325
2 changed files with 75 additions and 0 deletions
+5
View File
@@ -59,6 +59,7 @@ async def lifespan(app: FastAPI): # pragma: no cover
app.state.sync_task = asyncio.create_task(initialize_file_sync(app_config))
else:
logger.info("Sync changes disabled. Skipping file sync service.")
app.state.sync_task = None
# proceed with startup
yield
@@ -67,6 +68,10 @@ async def lifespan(app: FastAPI): # pragma: no cover
if app.state.sync_task:
logger.info("Stopping sync...")
app.state.sync_task.cancel() # pyright: ignore
try:
await app.state.sync_task
except asyncio.CancelledError:
logger.info("Sync task cancelled successfully")
await db.shutdown_db()
@@ -0,0 +1,70 @@
"""
Integration test for FastAPI lifespan shutdown behavior.
This test verifies the asyncio cancellation pattern used by the API lifespan:
when the background sync task is cancelled during shutdown, it must be *awaited*
before database shutdown begins. This prevents "hang on exit" scenarios in
`asyncio.run(...)` callers (e.g. CLI/MCP clients using httpx ASGITransport).
"""
import asyncio
from httpx import ASGITransport, AsyncClient
def test_lifespan_shutdown_awaits_sync_task_cancellation(app, monkeypatch):
"""
Ensure lifespan shutdown awaits the cancelled background sync task.
Why this is deterministic:
- Cancelling a task does not make it "done" immediately; it becomes done only
once the event loop schedules it and it processes the CancelledError.
- In the buggy version, shutdown proceeded directly to db.shutdown_db()
immediately after calling cancel(), so at *entry* to shutdown_db the task
is still not done.
- In the fixed version, lifespan does `await sync_task` before shutdown_db,
so by the time shutdown_db is called, the task is done (cancelled).
"""
# Import the *module* (not the package-level FastAPI `basic_memory.api.app` export)
# so monkeypatching affects the exact symbols referenced inside lifespan().
#
# Note: `basic_memory/api/__init__.py` re-exports `app`, so `import basic_memory.api.app`
# can resolve to the FastAPI instance rather than the `basic_memory.api.app` module.
import importlib
api_app_module = importlib.import_module("basic_memory.api.app")
# Keep startup cheap: we don't need real DB init for this ordering test.
async def _noop_initialize_app(_app_config):
return None
async def _fake_get_or_create_db(*_args, **_kwargs):
return object(), object()
monkeypatch.setattr(api_app_module, "initialize_app", _noop_initialize_app)
monkeypatch.setattr(api_app_module.db, "get_or_create_db", _fake_get_or_create_db)
# Make the sync task long-lived so it must be cancelled on shutdown.
async def _fake_initialize_file_sync(_app_config):
await asyncio.Event().wait()
monkeypatch.setattr(api_app_module, "initialize_file_sync", _fake_initialize_file_sync)
# Assert ordering: shutdown_db must be called only after the sync_task is done.
async def _assert_sync_task_done_before_db_shutdown():
assert api_app_module.app.state.sync_task is not None
assert api_app_module.app.state.sync_task.done()
monkeypatch.setattr(api_app_module.db, "shutdown_db", _assert_sync_task_done_before_db_shutdown)
async def _run_client_once():
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
# Any request is sufficient to trigger lifespan startup/shutdown.
await client.get("/__nonexistent__")
# Use asyncio.run to match the CLI/MCP execution model where loop teardown
# would hang if a background task is left running.
asyncio.run(_run_client_once())