fix: dispose engine after import-time table creation (#1255) (#1256)

when `ensure_db_tables_exist()` runs on import, it creates an async
engine via `asyncio.run()`. the engine's aiosqlite worker thread would
remain cached after the event loop closed, potentially preventing
python from exiting cleanly.

this adds a `dispose_engine` parameter to `create_db_and_tables()` which
disposes the engine and removes it from the cache after table creation.
`ensure_db_tables_exist()` now uses this to clean up after itself.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>

Author: zzstoatzz

docs/api-reference/marvin-database.mdx

@@ -120,15 +120,16 @@ Custom type for Usage objects that stores them as JSON in the database.
 
 ### `create_db_and_tables`
 ```python
-def create_db_and_tables(force: bool = False)
+def create_db_and_tables(force: bool = False, dispose_engine: bool = False)
 ```
-Create all database tables synchronously.
-
-This is a synchronous alternative to create_db_and_tables() that can be used
-in contexts where asyncio.run() cannot be called.
+Create all database tables.
 
 Args:
     force: If True, drops all existing tables before creating new ones.
+    dispose_engine: If True, dispose the engine after creating tables.
+        This is useful when called from asyncio.run() to ensure the
+        aiosqlite worker thread is cleaned up and doesn't prevent
+        Python from exiting.
 
 ### `ensure_db_tables_exist`
 ```python

src/marvin/database.py

@@ -416,14 +416,17 @@ def _run_migrations(alembic_log_level: str = "WARNING") -> bool:
         return False
 
 
-async def create_db_and_tables(*, force: bool = False) -> None:
-    """Create all database tables synchronously.
-
-    This is a synchronous alternative to create_db_and_tables() that can be used
-    in contexts where asyncio.run() cannot be called.
+async def create_db_and_tables(
+    *, force: bool = False, dispose_engine: bool = False
+) -> None:
+    """Create all database tables.
 
     Args:
         force: If True, drops all existing tables before creating new ones.
+        dispose_engine: If True, dispose the engine after creating tables.
+            This is useful when called from asyncio.run() to ensure the
+            aiosqlite worker thread is cleaned up and doesn't prevent
+            Python from exiting.
     """
     engine = get_async_engine()
 
@@ -435,6 +438,15 @@ async def create_db_and_tables(*, force: bool = False) -> None:
         await conn.run_sync(Base.metadata.create_all)
         logger.debug("Database tables created.")
 
+    if dispose_engine:
+        await engine.dispose()
+        # Remove the engine from the cache since it's disposed
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            loop = None
+        _async_engine_cache.pop(loop, None)
+
 
 def init_database_if_necessary():
     """Initialize the database file if necessary.
@@ -543,7 +555,10 @@ def ensure_db_tables_exist():
     # Create tables synchronously using asyncio.run() if we're not in an async context
     try:
         with _create_tables_lock:
-            asyncio.run(create_db_and_tables())
+            # Use dispose_engine=True to clean up the aiosqlite worker thread
+            # after table creation. Without this, the cached engine's worker
+            # thread could prevent Python from exiting cleanly.
+            asyncio.run(create_db_and_tables(dispose_engine=True))
             logger.debug("Database tables created successfully.")
     except Exception as e:
         logger.error(f"Failed to create database tables: {e}")

tests/basic/engine/test_database.py

@@ -1,9 +1,13 @@
+import asyncio
+import threading
+
 from sqlalchemy import select
 from sqlalchemy.orm import selectinload
 
 from marvin.database import (
     DBMessage,
     DBThread,
+    _async_engine_cache,
     create_db_and_tables,
 )
 
@@ -104,3 +108,40 @@ async def test_relationship_operations(session):
 
     result = await session.execute(select(DBThread).where(DBThread.id == "test-thread"))
     assert result.scalars().first() is None
+
+
+def test_create_db_and_tables_with_dispose_cleans_up_engine():
+    """Test that create_db_and_tables with dispose_engine=True cleans up resources.
+
+    Regression test for issue #1255: process hangs on exit after importing marvin.
+
+    When create_db_and_tables is called with dispose_engine=True (as it is from
+    ensure_db_tables_exist), the engine should be disposed and removed from
+    the cache. This ensures the aiosqlite worker thread is cleaned up and
+    doesn't prevent Python from exiting.
+    """
+
+    def get_non_daemon_threads():
+        """Return non-daemon threads excluding the main thread."""
+        return [
+            t
+            for t in threading.enumerate()
+            if t is not threading.main_thread() and not t.daemon
+        ]
+
+    # record initial state
+    initial_cache_size = len(_async_engine_cache)
+    initial_threads = len(get_non_daemon_threads())
+
+    # run create_db_and_tables with dispose_engine=True
+    asyncio.run(create_db_and_tables(dispose_engine=True))
+
+    # verify engine was removed from cache
+    assert len(_async_engine_cache) == initial_cache_size
+
+    # verify no new non-daemon threads were left behind
+    current_threads = get_non_daemon_threads()
+    assert len(current_threads) == initial_threads, (
+        f"expected {initial_threads} non-daemon threads, "
+        f"found {len(current_threads)}: {[t.name for t in current_threads]}"
+    )