From 6d470d26be1122451f12552d45910cd9b741307b Mon Sep 17 00:00:00 2001 From: Mykyta Yaromenko Date: Fri, 17 Jul 2026 20:54:42 +0200 Subject: [PATCH 1/4] fix(migrations): improve async migration handling by adapting to existing event loops Signed-off-by: Mykyta Yaromenko --- src/basic_memory/alembic/env.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/src/basic_memory/alembic/env.py b/src/basic_memory/alembic/env.py index 4b9bb4b2e..e4ed251ee 100644 --- a/src/basic_memory/alembic/env.py +++ b/src/basic_memory/alembic/env.py @@ -168,16 +168,15 @@ def run_in_thread(): def _run_async_engine_migrations(connectable) -> None: - """Run async-engine migrations with a running-loop fallback.""" + """Run async migrations, adapting to whether an event loop is already running.""" try: + asyncio.get_running_loop() + except RuntimeError: + # No running loop -> safe to spin up our own via asyncio.run() _run_async_migrations_with_asyncio_run(connectable) - except RuntimeError as e: - if is_running_loop_error(e): - # We're in a running event loop (likely uvloop or Python 3.14+ tests). - # Switch to a dedicated thread so Alembic can finish without nesting loops. - _run_async_migrations_in_thread(connectable) - else: - raise + else: + # A loop is already running -> can't nest asyncio.run(), so offload to a thread + _run_async_migrations_in_thread(connectable) def run_migrations_online() -> None: From 6d6e2e1cf126497fc6a640ec4310e55b9a72e83e Mon Sep 17 00:00:00 2001 From: Mykyta Yaromenko Date: Sat, 18 Jul 2026 13:21:36 +0200 Subject: [PATCH 2/4] fix(migrations): enhance async migration handling to support running event loops Signed-off-by: Mykyta Yaromenko --- src/basic_memory/alembic/env.py | 28 ++++++++++++++++++++++------ tests/test_alembic_env.py | 11 +++++++++-- 2 files changed, 31 insertions(+), 8 deletions(-) diff --git a/src/basic_memory/alembic/env.py b/src/basic_memory/alembic/env.py index e4ed251ee..5fb1a4ffb 100644 --- a/src/basic_memory/alembic/env.py +++ b/src/basic_memory/alembic/env.py @@ -29,7 +29,9 @@ # running loop (ValueError). # Outcome: log at DEBUG (observable, not noisy) and fall through to the # thread-based migration fallback. - logger.debug(f"nest_asyncio not applied ({exc!r}); using thread-based migration fallback") + logger.debug( + f"nest_asyncio not applied ({exc!r}); using thread-based migration fallback" + ) # Trigger: only set test env when actually running under pytest # Why: alembic/env.py is imported during normal operations (MCP server startup, migrations) @@ -167,15 +169,29 @@ def run_in_thread(): future.result() # Wait for completion and re-raise any exceptions -def _run_async_engine_migrations(connectable) -> None: - """Run async migrations, adapting to whether an event loop is already running.""" +def _loop_is_running() -> bool: + """Check if an event loop is currently running.""" try: asyncio.get_running_loop() except RuntimeError: - # No running loop -> safe to spin up our own via asyncio.run() + return False + return True + + +def _run_async_engine_migrations(connectable) -> None: + """Run async migrations, adapting to whether an event loop is already running.""" + if _loop_is_running(): + # Can't nest asyncio.run() inside a running loop -> offload to a thread. + _run_async_migrations_in_thread(connectable) + return + + # No running loop -> use asyncio.run(), falling back to a thread if it still + # reports a loop (covers races / odd environments). + try: _run_async_migrations_with_asyncio_run(connectable) - else: - # A loop is already running -> can't nest asyncio.run(), so offload to a thread + except RuntimeError as exc: + if not is_running_loop_error(exc): + raise _run_async_migrations_in_thread(connectable) diff --git a/tests/test_alembic_env.py b/tests/test_alembic_env.py index 2dd3a6ae1..48a7dae41 100644 --- a/tests/test_alembic_env.py +++ b/tests/test_alembic_env.py @@ -79,14 +79,21 @@ def raising_asyncio_run(coro): assert fake_coro.closed is True -def test_running_loop_error_uses_thread_fallback(monkeypatch, tmp_path): +@pytest.mark.parametrize( + "message", + [ + "asyncio.run() cannot be called from a running event loop", + "this event loop is already running", + ], +) +def test_running_loop_error_uses_thread_fallback(monkeypatch, tmp_path, message): """Async-engine helper should switch to the thread fallback for running-loop errors.""" env_module = load_alembic_env_module(monkeypatch, tmp_path) connectable = object() fallback_calls: list[object] = [] def raising_run(connectable): - raise RuntimeError("asyncio.run() cannot be called from a running event loop") + raise RuntimeError(message) def record_fallback(target): fallback_calls.append(target) From 878d5d61dbf367e0e6629c21cab6c6b5139822cf Mon Sep 17 00:00:00 2001 From: Mykyta Yaromenko Date: Sat, 18 Jul 2026 13:36:56 +0200 Subject: [PATCH 3/4] fix(migrations): streamline async migration handling by removing redundant error checks Signed-off-by: Mykyta Yaromenko --- src/basic_memory/alembic/env.py | 20 +++---- src/basic_memory/migration_loop.py | 15 ------ tests/test_alembic_env.py | 87 +++++++++++++++++++++--------- tests/test_migration_loop.py | 28 ++++------ 4 files changed, 81 insertions(+), 69 deletions(-) diff --git a/src/basic_memory/alembic/env.py b/src/basic_memory/alembic/env.py index 5fb1a4ffb..cda7be893 100644 --- a/src/basic_memory/alembic/env.py +++ b/src/basic_memory/alembic/env.py @@ -13,7 +13,7 @@ from alembic import context from basic_memory.config import ConfigManager -from basic_memory.migration_loop import is_running_loop_error, running_on_uvloop +from basic_memory.migration_loop import running_on_uvloop # Allow nested event loops (needed for pytest-asyncio and other async contexts). # nest_asyncio cannot patch a uvloop loop or Python 3.14+; in those cases we skip @@ -174,25 +174,25 @@ def _loop_is_running() -> bool: try: asyncio.get_running_loop() except RuntimeError: + # Trigger: no running loop (get_running_loop raises RuntimeError) + # Outcome: return False to indicate no loop is running return False return True def _run_async_engine_migrations(connectable) -> None: """Run async migrations, adapting to whether an event loop is already running.""" + + # Trigger: caller may be inside a running event loop (e.g. pytest-asyncio, uvicorn). + # Why: asyncio.run() raises RuntimeError when nested inside a running loop, so we + # detect the condition up front via _loop_is_running() instead of catching the error. + # Outcome: running-loop context -> thread fallback; no loop -> asyncio.run() directly. if _loop_is_running(): # Can't nest asyncio.run() inside a running loop -> offload to a thread. _run_async_migrations_in_thread(connectable) - return - - # No running loop -> use asyncio.run(), falling back to a thread if it still - # reports a loop (covers races / odd environments). - try: + else: + # No running loop: asyncio.run() is safe; let any unexpected errors propagate. _run_async_migrations_with_asyncio_run(connectable) - except RuntimeError as exc: - if not is_running_loop_error(exc): - raise - _run_async_migrations_in_thread(connectable) def run_migrations_online() -> None: diff --git a/src/basic_memory/migration_loop.py b/src/basic_memory/migration_loop.py index 38e01a83e..d9179b1d9 100644 --- a/src/basic_memory/migration_loop.py +++ b/src/basic_memory/migration_loop.py @@ -28,18 +28,3 @@ def running_on_uvloop() -> bool: except ImportError: return False return isinstance(asyncio.get_event_loop_policy(), uvloop.EventLoopPolicy) - - -def is_running_loop_error(exc: BaseException) -> bool: - """Whether an error means ``asyncio.run`` hit an already-running event loop. - - Two spellings reach the migration fallback: the stdlib "cannot be called - from a running event loop", and "this event loop is already running" when - nest_asyncio mis-patched a uvloop loop. Both mean the same thing — retry the - migration in a dedicated thread rather than re-raising. - """ - msg = str(exc) - return ( - "cannot be called from a running event loop" in msg - or "this event loop is already running" in msg - ) diff --git a/tests/test_alembic_env.py b/tests/test_alembic_env.py index 48a7dae41..5176244a7 100644 --- a/tests/test_alembic_env.py +++ b/tests/test_alembic_env.py @@ -46,8 +46,12 @@ def load_alembic_env_module(monkeypatch, tmp_path): fake_config = FakeAlembicConfig() monkeypatch.setattr(alembic_context, "config", fake_config, raising=False) - monkeypatch.setattr(alembic_context, "configure", lambda *args, **kwargs: None, raising=False) - monkeypatch.setattr(alembic_context, "begin_transaction", lambda: nullcontext(), raising=False) + monkeypatch.setattr( + alembic_context, "configure", lambda *args, **kwargs: None, raising=False + ) + monkeypatch.setattr( + alembic_context, "begin_transaction", lambda: nullcontext(), raising=False + ) monkeypatch.setattr(alembic_context, "run_migrations", lambda: None, raising=False) monkeypatch.setattr(alembic_context, "is_offline_mode", lambda: True, raising=False) @@ -66,7 +70,9 @@ def test_asyncio_run_failure_closes_migration_coroutine(monkeypatch, tmp_path): env_module = load_alembic_env_module(monkeypatch, tmp_path) fake_coro = FakeCoroutine() - monkeypatch.setattr(env_module, "run_async_migrations", lambda connectable: fake_coro) + monkeypatch.setattr( + env_module, "run_async_migrations", lambda connectable: fake_coro + ) def raising_asyncio_run(coro): raise RuntimeError("asyncio.run() cannot be called from a running event loop") @@ -79,41 +85,72 @@ def raising_asyncio_run(coro): assert fake_coro.closed is True -@pytest.mark.parametrize( - "message", - [ - "asyncio.run() cannot be called from a running event loop", - "this event loop is already running", - ], -) -def test_running_loop_error_uses_thread_fallback(monkeypatch, tmp_path, message): - """Async-engine helper should switch to the thread fallback for running-loop errors.""" +def test_running_loop_uses_thread_path(monkeypatch, tmp_path): + """When a loop is already running, _run_async_engine_migrations must use the thread path.""" + import asyncio + env_module = load_alembic_env_module(monkeypatch, tmp_path) connectable = object() - fallback_calls: list[object] = [] + thread_calls: list[object] = [] + asyncio_run_calls: list[object] = [] + + monkeypatch.setattr(env_module, "_loop_is_running", lambda: True) + monkeypatch.setattr( + env_module, "_run_async_migrations_in_thread", lambda c: thread_calls.append(c) + ) + monkeypatch.setattr( + env_module, + "_run_async_migrations_with_asyncio_run", + lambda c: asyncio_run_calls.append(c), + ) - def raising_run(connectable): - raise RuntimeError(message) + env_module._run_async_engine_migrations(connectable) - def record_fallback(target): - fallback_calls.append(target) + assert thread_calls == [connectable], "thread fallback should have been called" + assert ( + asyncio_run_calls == [] + ), "asyncio.run path must not be called when loop is running" - monkeypatch.setattr(env_module, "_run_async_migrations_with_asyncio_run", raising_run) - monkeypatch.setattr(env_module, "_run_async_migrations_in_thread", record_fallback) + +def test_no_running_loop_uses_asyncio_run_path(monkeypatch, tmp_path): + """When no loop is running, _run_async_engine_migrations must use asyncio.run directly.""" + env_module = load_alembic_env_module(monkeypatch, tmp_path) + connectable = object() + asyncio_run_calls: list[object] = [] + thread_calls: list[object] = [] + + monkeypatch.setattr(env_module, "_loop_is_running", lambda: False) + monkeypatch.setattr( + env_module, + "_run_async_migrations_with_asyncio_run", + lambda c: asyncio_run_calls.append(c), + ) + monkeypatch.setattr( + env_module, "_run_async_migrations_in_thread", lambda c: thread_calls.append(c) + ) env_module._run_async_engine_migrations(connectable) - assert fallback_calls == [connectable] + assert asyncio_run_calls == [ + connectable + ], "asyncio.run path should have been called" + assert ( + thread_calls == [] + ), "thread fallback must not be called when no loop is running" -def test_non_loop_runtime_error_is_re_raised(monkeypatch, tmp_path): - """Unexpected RuntimeError values should not be swallowed by the fallback path.""" +def test_runtime_error_from_asyncio_run_propagates(monkeypatch, tmp_path): + """RuntimeErrors from _run_async_migrations_with_asyncio_run should propagate unchanged.""" env_module = load_alembic_env_module(monkeypatch, tmp_path) + monkeypatch.setattr(env_module, "_loop_is_running", lambda: False) + def raising_run(connectable): - raise RuntimeError("different runtime failure") + raise RuntimeError("unexpected failure") - monkeypatch.setattr(env_module, "_run_async_migrations_with_asyncio_run", raising_run) + monkeypatch.setattr( + env_module, "_run_async_migrations_with_asyncio_run", raising_run + ) - with pytest.raises(RuntimeError, match="different runtime failure"): + with pytest.raises(RuntimeError, match="unexpected failure"): env_module._run_async_engine_migrations(object()) diff --git a/tests/test_migration_loop.py b/tests/test_migration_loop.py index 668def3c4..d3df8a092 100644 --- a/tests/test_migration_loop.py +++ b/tests/test_migration_loop.py @@ -14,32 +14,22 @@ from basic_memory import migration_loop -def test_is_running_loop_error_matches_stdlib_message(): - err = RuntimeError("asyncio.run() cannot be called from a running event loop") - assert migration_loop.is_running_loop_error(err) is True - - -def test_is_running_loop_error_matches_nest_asyncio_uvloop_message(): - # nest_asyncio mis-patched a uvloop loop -> different wording, same meaning. - err = RuntimeError("this event loop is already running") - assert migration_loop.is_running_loop_error(err) is True - - -def test_is_running_loop_error_rejects_unrelated_runtime_error(): - err = RuntimeError("migration failed: column already exists") - assert migration_loop.is_running_loop_error(err) is False - - -@pytest.mark.skipif(sys.platform == "win32", reason="uvloop is not available on Windows") +@pytest.mark.skipif( + sys.platform == "win32", reason="uvloop is not available on Windows" +) def test_running_on_uvloop_true_when_policy_is_uvloop(monkeypatch): import uvloop - monkeypatch.setattr(asyncio, "get_event_loop_policy", lambda: uvloop.EventLoopPolicy()) + monkeypatch.setattr( + asyncio, "get_event_loop_policy", lambda: uvloop.EventLoopPolicy() + ) assert migration_loop.running_on_uvloop() is True def test_running_on_uvloop_false_for_default_policy(monkeypatch): - monkeypatch.setattr(asyncio, "get_event_loop_policy", lambda: asyncio.DefaultEventLoopPolicy()) + monkeypatch.setattr( + asyncio, "get_event_loop_policy", lambda: asyncio.DefaultEventLoopPolicy() + ) assert migration_loop.running_on_uvloop() is False From d0707852eeefaf4d7ab57d24885e89c948ccf249 Mon Sep 17 00:00:00 2001 From: phernandez Date: Mon, 20 Jul 2026 21:09:40 -0500 Subject: [PATCH 4/4] style: format migration fallback changes to repo standard ruff format + one unused-import fix so the fork branch passes the static checks that don't run on fork PRs. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01A2bgrGfWiL4izcjj66u9R8 Signed-off-by: phernandez --- src/basic_memory/alembic/env.py | 4 +--- tests/test_alembic_env.py | 29 +++++++---------------------- tests/test_migration_loop.py | 12 +++--------- 3 files changed, 11 insertions(+), 34 deletions(-) diff --git a/src/basic_memory/alembic/env.py b/src/basic_memory/alembic/env.py index cda7be893..8649da590 100644 --- a/src/basic_memory/alembic/env.py +++ b/src/basic_memory/alembic/env.py @@ -29,9 +29,7 @@ # running loop (ValueError). # Outcome: log at DEBUG (observable, not noisy) and fall through to the # thread-based migration fallback. - logger.debug( - f"nest_asyncio not applied ({exc!r}); using thread-based migration fallback" - ) + logger.debug(f"nest_asyncio not applied ({exc!r}); using thread-based migration fallback") # Trigger: only set test env when actually running under pytest # Why: alembic/env.py is imported during normal operations (MCP server startup, migrations) diff --git a/tests/test_alembic_env.py b/tests/test_alembic_env.py index 5176244a7..64e20ec65 100644 --- a/tests/test_alembic_env.py +++ b/tests/test_alembic_env.py @@ -46,12 +46,8 @@ def load_alembic_env_module(monkeypatch, tmp_path): fake_config = FakeAlembicConfig() monkeypatch.setattr(alembic_context, "config", fake_config, raising=False) - monkeypatch.setattr( - alembic_context, "configure", lambda *args, **kwargs: None, raising=False - ) - monkeypatch.setattr( - alembic_context, "begin_transaction", lambda: nullcontext(), raising=False - ) + monkeypatch.setattr(alembic_context, "configure", lambda *args, **kwargs: None, raising=False) + monkeypatch.setattr(alembic_context, "begin_transaction", lambda: nullcontext(), raising=False) monkeypatch.setattr(alembic_context, "run_migrations", lambda: None, raising=False) monkeypatch.setattr(alembic_context, "is_offline_mode", lambda: True, raising=False) @@ -70,9 +66,7 @@ def test_asyncio_run_failure_closes_migration_coroutine(monkeypatch, tmp_path): env_module = load_alembic_env_module(monkeypatch, tmp_path) fake_coro = FakeCoroutine() - monkeypatch.setattr( - env_module, "run_async_migrations", lambda connectable: fake_coro - ) + monkeypatch.setattr(env_module, "run_async_migrations", lambda connectable: fake_coro) def raising_asyncio_run(coro): raise RuntimeError("asyncio.run() cannot be called from a running event loop") @@ -87,7 +81,6 @@ def raising_asyncio_run(coro): def test_running_loop_uses_thread_path(monkeypatch, tmp_path): """When a loop is already running, _run_async_engine_migrations must use the thread path.""" - import asyncio env_module = load_alembic_env_module(monkeypatch, tmp_path) connectable = object() @@ -107,9 +100,7 @@ def test_running_loop_uses_thread_path(monkeypatch, tmp_path): env_module._run_async_engine_migrations(connectable) assert thread_calls == [connectable], "thread fallback should have been called" - assert ( - asyncio_run_calls == [] - ), "asyncio.run path must not be called when loop is running" + assert asyncio_run_calls == [], "asyncio.run path must not be called when loop is running" def test_no_running_loop_uses_asyncio_run_path(monkeypatch, tmp_path): @@ -131,12 +122,8 @@ def test_no_running_loop_uses_asyncio_run_path(monkeypatch, tmp_path): env_module._run_async_engine_migrations(connectable) - assert asyncio_run_calls == [ - connectable - ], "asyncio.run path should have been called" - assert ( - thread_calls == [] - ), "thread fallback must not be called when no loop is running" + assert asyncio_run_calls == [connectable], "asyncio.run path should have been called" + assert thread_calls == [], "thread fallback must not be called when no loop is running" def test_runtime_error_from_asyncio_run_propagates(monkeypatch, tmp_path): @@ -148,9 +135,7 @@ def test_runtime_error_from_asyncio_run_propagates(monkeypatch, tmp_path): def raising_run(connectable): raise RuntimeError("unexpected failure") - monkeypatch.setattr( - env_module, "_run_async_migrations_with_asyncio_run", raising_run - ) + monkeypatch.setattr(env_module, "_run_async_migrations_with_asyncio_run", raising_run) with pytest.raises(RuntimeError, match="unexpected failure"): env_module._run_async_engine_migrations(object()) diff --git a/tests/test_migration_loop.py b/tests/test_migration_loop.py index d3df8a092..98d4c77d2 100644 --- a/tests/test_migration_loop.py +++ b/tests/test_migration_loop.py @@ -14,22 +14,16 @@ from basic_memory import migration_loop -@pytest.mark.skipif( - sys.platform == "win32", reason="uvloop is not available on Windows" -) +@pytest.mark.skipif(sys.platform == "win32", reason="uvloop is not available on Windows") def test_running_on_uvloop_true_when_policy_is_uvloop(monkeypatch): import uvloop - monkeypatch.setattr( - asyncio, "get_event_loop_policy", lambda: uvloop.EventLoopPolicy() - ) + monkeypatch.setattr(asyncio, "get_event_loop_policy", lambda: uvloop.EventLoopPolicy()) assert migration_loop.running_on_uvloop() is True def test_running_on_uvloop_false_for_default_policy(monkeypatch): - monkeypatch.setattr( - asyncio, "get_event_loop_policy", lambda: asyncio.DefaultEventLoopPolicy() - ) + monkeypatch.setattr(asyncio, "get_event_loop_policy", lambda: asyncio.DefaultEventLoopPolicy()) assert migration_loop.running_on_uvloop() is False