fix: prevent CLI commands from hanging on exit (#471)

Signed-off-by: phernandez <paul@basicmachines.co>
Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
Paul Hernandez
2025-12-24 11:27:11 -06:00
committed by GitHub
parent 95937c6d0a
commit 916baf8971
17 changed files with 399 additions and 269 deletions
@@ -10,6 +10,42 @@ from typing import Sequence, Union
import sqlalchemy as sa
from alembic import op
from sqlalchemy import text
def column_exists(connection, table: str, column: str) -> bool:
"""Check if a column exists in a table (idempotent migration support)."""
if connection.dialect.name == "postgresql":
result = connection.execute(
text(
"SELECT 1 FROM information_schema.columns "
"WHERE table_name = :table AND column_name = :column"
),
{"table": table, "column": column},
)
return result.fetchone() is not None
else:
# SQLite
result = connection.execute(text(f"PRAGMA table_info({table})"))
columns = [row[1] for row in result]
return column in columns
def index_exists(connection, index_name: str) -> bool:
"""Check if an index exists (idempotent migration support)."""
if connection.dialect.name == "postgresql":
result = connection.execute(
text("SELECT 1 FROM pg_indexes WHERE indexname = :index_name"),
{"index_name": index_name},
)
return result.fetchone() is not None
else:
# SQLite
result = connection.execute(
text("SELECT 1 FROM sqlite_master WHERE type='index' AND name = :index_name"),
{"index_name": index_name},
)
return result.fetchone() is not None
# revision identifiers, used by Alembic.
@@ -36,101 +72,105 @@ def upgrade() -> None:
# Add project_id to relation table
# -------------------------------------------------------------------------
# Step 1: Add project_id column as nullable first
op.add_column("relation", sa.Column("project_id", sa.Integer(), nullable=True))
# Step 1: Add project_id column as nullable first (idempotent)
if not column_exists(connection, "relation", "project_id"):
op.add_column("relation", sa.Column("project_id", sa.Integer(), nullable=True))
# Step 2: Backfill project_id from entity.project_id via from_id
if dialect == "postgresql":
op.execute("""
UPDATE relation
SET project_id = entity.project_id
FROM entity
WHERE relation.from_id = entity.id
""")
else:
# SQLite syntax
op.execute("""
UPDATE relation
SET project_id = (
SELECT entity.project_id
# Step 2: Backfill project_id from entity.project_id via from_id
if dialect == "postgresql":
op.execute("""
UPDATE relation
SET project_id = entity.project_id
FROM entity
WHERE entity.id = relation.from_id
)
""")
WHERE relation.from_id = entity.id
""")
else:
# SQLite syntax
op.execute("""
UPDATE relation
SET project_id = (
SELECT entity.project_id
FROM entity
WHERE entity.id = relation.from_id
)
""")
# Step 3: Make project_id NOT NULL and add foreign key
if dialect == "postgresql":
op.alter_column("relation", "project_id", nullable=False)
op.create_foreign_key(
"fk_relation_project_id",
"relation",
"project",
["project_id"],
["id"],
)
else:
# SQLite requires batch operations for ALTER COLUMN
with op.batch_alter_table("relation") as batch_op:
batch_op.alter_column("project_id", nullable=False)
batch_op.create_foreign_key(
# Step 3: Make project_id NOT NULL and add foreign key
if dialect == "postgresql":
op.alter_column("relation", "project_id", nullable=False)
op.create_foreign_key(
"fk_relation_project_id",
"relation",
"project",
["project_id"],
["id"],
)
else:
# SQLite requires batch operations for ALTER COLUMN
with op.batch_alter_table("relation") as batch_op:
batch_op.alter_column("project_id", nullable=False)
batch_op.create_foreign_key(
"fk_relation_project_id",
"project",
["project_id"],
["id"],
)
# Step 4: Create index on relation.project_id
op.create_index("ix_relation_project_id", "relation", ["project_id"])
# Step 4: Create index on relation.project_id (idempotent)
if not index_exists(connection, "ix_relation_project_id"):
op.create_index("ix_relation_project_id", "relation", ["project_id"])
# -------------------------------------------------------------------------
# Add project_id to observation table
# -------------------------------------------------------------------------
# Step 1: Add project_id column as nullable first
op.add_column("observation", sa.Column("project_id", sa.Integer(), nullable=True))
# Step 1: Add project_id column as nullable first (idempotent)
if not column_exists(connection, "observation", "project_id"):
op.add_column("observation", sa.Column("project_id", sa.Integer(), nullable=True))
# Step 2: Backfill project_id from entity.project_id via entity_id
if dialect == "postgresql":
op.execute("""
UPDATE observation
SET project_id = entity.project_id
FROM entity
WHERE observation.entity_id = entity.id
""")
else:
# SQLite syntax
op.execute("""
UPDATE observation
SET project_id = (
SELECT entity.project_id
# Step 2: Backfill project_id from entity.project_id via entity_id
if dialect == "postgresql":
op.execute("""
UPDATE observation
SET project_id = entity.project_id
FROM entity
WHERE entity.id = observation.entity_id
)
""")
WHERE observation.entity_id = entity.id
""")
else:
# SQLite syntax
op.execute("""
UPDATE observation
SET project_id = (
SELECT entity.project_id
FROM entity
WHERE entity.id = observation.entity_id
)
""")
# Step 3: Make project_id NOT NULL and add foreign key
if dialect == "postgresql":
op.alter_column("observation", "project_id", nullable=False)
op.create_foreign_key(
"fk_observation_project_id",
"observation",
"project",
["project_id"],
["id"],
)
else:
# SQLite requires batch operations for ALTER COLUMN
with op.batch_alter_table("observation") as batch_op:
batch_op.alter_column("project_id", nullable=False)
batch_op.create_foreign_key(
# Step 3: Make project_id NOT NULL and add foreign key
if dialect == "postgresql":
op.alter_column("observation", "project_id", nullable=False)
op.create_foreign_key(
"fk_observation_project_id",
"observation",
"project",
["project_id"],
["id"],
)
else:
# SQLite requires batch operations for ALTER COLUMN
with op.batch_alter_table("observation") as batch_op:
batch_op.alter_column("project_id", nullable=False)
batch_op.create_foreign_key(
"fk_observation_project_id",
"project",
["project_id"],
["id"],
)
# Step 4: Create index on observation.project_id
op.create_index("ix_observation_project_id", "observation", ["project_id"])
# Step 4: Create index on observation.project_id (idempotent)
if not index_exists(connection, "ix_observation_project_id"):
op.create_index("ix_observation_project_id", "observation", ["project_id"])
# Postgres-specific: pg_trgm and GIN indexes
if dialect == "postgresql":
+18 -4
View File
@@ -1,6 +1,7 @@
"""FastAPI application for basic-memory knowledge graph API."""
import asyncio
import os
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException
@@ -53,12 +54,25 @@ async def lifespan(app: FastAPI): # pragma: no cover
app.state.session_maker = session_maker
logger.info("Database connections cached in app state")
logger.info(f"Sync changes enabled: {app_config.sync_changes}")
if app_config.sync_changes:
# Start file sync if enabled
is_test_env = (
app_config.env == "test"
or os.getenv("BASIC_MEMORY_ENV", "").lower() == "test"
or os.getenv("PYTEST_CURRENT_TEST") is not None
)
if app_config.sync_changes and not is_test_env:
logger.info(f"Sync changes enabled: {app_config.sync_changes}")
# start file sync task in background
app.state.sync_task = asyncio.create_task(initialize_file_sync(app_config))
async def _file_sync_runner() -> None:
await initialize_file_sync(app_config)
app.state.sync_task = asyncio.create_task(_file_sync_runner())
else:
logger.info("Sync changes disabled. Skipping file sync service.")
if is_test_env:
logger.info("Test environment detected. Skipping file sync service.")
else:
logger.info("Sync changes disabled. Skipping file sync service.")
app.state.sync_task = None
# proceed with startup
+9 -2
View File
@@ -34,8 +34,15 @@ def app_callback(
# Initialize logging for CLI (file only, no stdout)
init_cli_logging()
# Run initialization for every command unless --version was specified
if not version and ctx.invoked_subcommand is not None:
# Run initialization for commands that don't use the API
# Skip for 'mcp' command - it has its own lifespan that handles initialization
# Skip for API-using commands (status, sync, etc.) - they handle initialization via deps.py
api_commands = {"mcp", "status", "sync", "project", "tool"}
if (
not version
and ctx.invoked_subcommand is not None
and ctx.invoked_subcommand not in api_commands
):
from basic_memory.services.initialization import ensure_initialization
app_config = ConfigManager().config
+27 -1
View File
@@ -1,12 +1,14 @@
"""utility functions for commands"""
from typing import Optional
import asyncio
from typing import Optional, TypeVar, Coroutine, Any
from mcp.server.fastmcp.exceptions import ToolError
import typer
from rich.console import Console
from basic_memory import db
from basic_memory.mcp.async_client import get_client
from basic_memory.mcp.tools.utils import call_post, call_get
@@ -15,6 +17,30 @@ from basic_memory.schemas import ProjectInfoResponse
console = Console()
T = TypeVar("T")
def run_with_cleanup(coro: Coroutine[Any, Any, T]) -> T:
"""Run an async coroutine with proper database cleanup.
This helper ensures database connections are cleaned up before the event
loop closes, preventing process hangs in CLI commands.
Args:
coro: The coroutine to run
Returns:
The result of the coroutine
"""
async def _with_cleanup() -> T:
try:
return await coro
finally:
await db.shutdown_db()
return asyncio.run(_with_cleanup())
async def run_sync(project: Optional[str] = None, force_full: bool = False):
"""Run sync operation via API endpoint.
+5 -25
View File
@@ -1,6 +1,5 @@
"""MCP server command with streamable HTTP transport."""
import asyncio
import os
import typer
from typing import Optional
@@ -8,7 +7,7 @@ from typing import Optional
from basic_memory.cli.app import app
from basic_memory.config import ConfigManager, init_mcp_logging
# Import mcp instance
# Import mcp instance (has lifespan that handles initialization and file sync)
from basic_memory.mcp.server import mcp as mcp_server # pragma: no cover
# Import mcp tools to register them
@@ -17,8 +16,6 @@ import basic_memory.mcp.tools # noqa: F401 # pragma: no cover
# Import prompts to register them
import basic_memory.mcp.prompts # noqa: F401 # pragma: no cover
from loguru import logger
import threading
from basic_memory.services.initialization import initialize_file_sync
config = ConfigManager().config
@@ -43,6 +40,8 @@ if not config.cloud_mode_enabled:
- stdio: Standard I/O (good for local usage)
- streamable-http: Recommended for web deployments (default)
- sse: Server-Sent Events (for compatibility with existing clients)
Initialization, file sync, and cleanup are handled by the MCP server's lifespan.
"""
# Initialize logging for MCP (file only, stdout breaks protocol)
init_mcp_logging()
@@ -59,27 +58,8 @@ if not config.cloud_mode_enabled:
os.environ["BASIC_MEMORY_MCP_PROJECT"] = project_name
logger.info(f"MCP server constrained to project: {project_name}")
app_config = ConfigManager().config
def run_file_sync():
"""Run file sync in a separate thread with its own event loop."""
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
loop.run_until_complete(initialize_file_sync(app_config))
except Exception as e:
logger.error(f"File sync error: {e}", err=True)
finally:
loop.close()
logger.info(f"Sync changes enabled: {app_config.sync_changes}")
if app_config.sync_changes:
# Start the sync thread
sync_thread = threading.Thread(target=run_file_sync, daemon=True)
sync_thread.start()
logger.info("Started file sync in background")
# Now run the MCP server (blocks)
# Run the MCP server (blocks)
# Lifespan handles: initialization, migrations, file sync, cleanup
logger.info(f"Starting MCP server with {transport.upper()} transport")
if transport == "stdio":
+3 -2
View File
@@ -1,6 +1,5 @@
"""Status command for basic-memory CLI."""
import asyncio
from typing import Set, Dict
from typing import Annotated, Optional
@@ -165,8 +164,10 @@ def status(
verbose: bool = typer.Option(False, "--verbose", "-v", help="Show detailed file information"),
):
"""Show sync status between files and database."""
from basic_memory.cli.commands.command_utils import run_with_cleanup
try:
asyncio.run(run_status(project, verbose)) # pragma: no cover
run_with_cleanup(run_status(project, verbose)) # pragma: no cover
except Exception as e:
logger.error(f"Error checking status: {e}")
typer.echo(f"Error checking status: {e}", err=True)
+16
View File
@@ -1,5 +1,6 @@
import asyncio
import os
import sys
from contextlib import asynccontextmanager
from enum import Enum, auto
from pathlib import Path
@@ -23,6 +24,21 @@ from sqlalchemy.pool import NullPool
from basic_memory.repository.postgres_search_repository import PostgresSearchRepository
from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository
# -----------------------------------------------------------------------------
# Windows event loop policy
# -----------------------------------------------------------------------------
# On Windows, the default ProactorEventLoop has known rough edges with aiosqlite
# during shutdown/teardown (threads posting results to a loop that's closing),
# which can manifest as:
# - "RuntimeError: Event loop is closed"
# - "IndexError: pop from an empty deque"
#
# The SelectorEventLoop doesn't support subprocess operations, so code that uses
# asyncio.create_subprocess_shell() (like sync_service._quick_count_files) must
# detect Windows and use fallback implementations.
if sys.platform == "win32": # pragma: no cover
asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
# Module level state
_engine: Optional[AsyncEngine] = None
_session_maker: Optional[async_sessionmaker[AsyncSession]] = None
+1
View File
@@ -95,6 +95,7 @@ async def get_client() -> AsyncIterator[AsyncClient]:
yield client
else:
# Local mode: ASGI transport for in-process calls
# Note: ASGI transport does NOT trigger FastAPI lifespan, so no special handling needed
logger.info("Creating ASGI client for local Basic Memory API")
async with AsyncClient(
transport=ASGITransport(app=fastapi_app), base_url="http://test", timeout=timeout
+72
View File
@@ -2,8 +2,80 @@
Basic Memory FastMCP server.
"""
import asyncio
import os
from contextlib import asynccontextmanager
from fastmcp import FastMCP
from loguru import logger
from basic_memory import db
from basic_memory.config import ConfigManager
from basic_memory.services.initialization import initialize_app, initialize_file_sync
@asynccontextmanager
async def lifespan(app: FastMCP):
"""Lifecycle manager for the MCP server.
Handles:
- Database initialization and migrations
- File sync in background (if enabled and not in cloud mode)
- Proper cleanup on shutdown
"""
app_config = ConfigManager().config
logger.info("Starting Basic Memory MCP server")
# Track if we created the engine (vs test fixtures providing it)
# This prevents disposing an engine provided by test fixtures when
# multiple Client connections are made in the same test
engine_was_none = db._engine is None
# Initialize app (runs migrations, reconciles projects)
await initialize_app(app_config)
# Start file sync as background task (if enabled and not in cloud mode)
sync_task = None
is_test_env = (
app_config.env == "test"
or os.getenv("BASIC_MEMORY_ENV", "").lower() == "test"
or os.getenv("PYTEST_CURRENT_TEST") is not None
)
if is_test_env:
logger.info("Test environment detected - skipping local file sync")
elif app_config.sync_changes and not app_config.cloud_mode_enabled:
logger.info("Starting file sync in background")
async def _file_sync_runner() -> None:
await initialize_file_sync(app_config)
sync_task = asyncio.create_task(_file_sync_runner())
elif app_config.cloud_mode_enabled:
logger.info("Cloud mode enabled - skipping local file sync")
else:
logger.info("Sync changes disabled - skipping file sync")
try:
yield
finally:
# Shutdown
logger.info("Shutting down Basic Memory MCP server")
if sync_task:
sync_task.cancel()
try:
await sync_task
except asyncio.CancelledError:
logger.info("File sync task cancelled")
# Only shutdown DB if we created it (not if test fixture provided it)
if engine_was_none:
await db.shutdown_db()
logger.info("Database connections closed")
else:
logger.debug("Skipping DB shutdown - engine provided externally")
mcp = FastMCP(
name="Basic Memory",
lifespan=lifespan,
)
+44 -26
View File
@@ -6,6 +6,7 @@ to ensure consistent application startup across all entry points.
import asyncio
import os
import sys
from pathlib import Path
@@ -29,15 +30,12 @@ async def initialize_database(app_config: BasicMemoryConfig) -> None:
Database migrations are now handled automatically when the database
connection is first established via get_or_create_db().
"""
# Trigger database initialization and migrations by getting the database connection
try:
await db.get_or_create_db(app_config.database_path)
logger.info("Database initialization completed")
except Exception as e:
logger.error(f"Error initializing database: {e}")
# Allow application to continue - it might still work
# depending on what the error was, and will fail with a
# more specific error if the database is actually unusable
logger.error(f"Error during database initialization: {e}")
raise
async def reconcile_projects_with_config(app_config: BasicMemoryConfig):
@@ -51,31 +49,29 @@ async def reconcile_projects_with_config(app_config: BasicMemoryConfig):
"""
logger.info("Reconciling projects from config with database...")
# Get database session - migrations handled centrally
# Get database session (engine already created by initialize_database)
_, session_maker = await db.get_or_create_db(
db_path=app_config.database_path,
db_type=db.DatabaseType.FILESYSTEM,
ensure_migrations=False,
)
project_repository = ProjectRepository(session_maker)
# Import ProjectService here to avoid circular imports
from basic_memory.services.project_service import ProjectService
# Create project service and synchronize projects
project_service = ProjectService(repository=project_repository)
try:
# Create project service and synchronize projects
project_service = ProjectService(repository=project_repository)
await project_service.synchronize_projects()
logger.info("Projects successfully reconciled between config and database")
except Exception as e:
# Log the error but continue with initialization
logger.error(f"Error during project synchronization: {e}")
logger.info("Continuing with initialization despite synchronization error")
async def initialize_file_sync(
app_config: BasicMemoryConfig,
):
) -> None:
"""Initialize file synchronization services. This function starts the watch service and does not return
Args:
@@ -84,15 +80,26 @@ async def initialize_file_sync(
Returns:
The watch service task that's monitoring file changes
"""
# Never start file watching during tests. Even "background" watchers add tasks/threads
# and can interact badly with strict asyncio teardown (especially on Windows/aiosqlite).
#
# Note: Some tests patch ConfigManager.config with a minimal BasicMemoryConfig that
# may not set env="test". So also detect pytest via env vars.
if (
app_config.env == "test"
or os.getenv("BASIC_MEMORY_ENV", "").lower() == "test"
or os.getenv("PYTEST_CURRENT_TEST") is not None
):
logger.info("Test environment detected - skipping file sync initialization")
return None
# delay import
from basic_memory.sync import WatchService
# Load app configuration - migrations handled centrally
# Get database session (migrations already run if needed)
_, session_maker = await db.get_or_create_db(
db_path=app_config.database_path,
db_type=db.DatabaseType.FILESYSTEM,
ensure_migrations=False,
)
project_repository = ProjectRepository(session_maker)
@@ -139,12 +146,10 @@ async def initialize_file_sync(
# Then start the watch service in the background
logger.info("Starting watch service for all projects")
# run the watch service
try:
await watch_service.run()
logger.info("Watch service started")
except Exception as e: # pragma: no cover
logger.error(f"Error starting watch service: {e}")
await watch_service.run()
logger.info("Watch service started")
return None
@@ -189,11 +194,24 @@ def ensure_initialization(app_config: BasicMemoryConfig) -> None:
logger.debug("Skipping initialization in cloud mode - projects managed by cloud")
return
try:
result = asyncio.run(initialize_app(app_config))
logger.info(f"Initialization completed successfully: result={result}")
except Exception as e: # pragma: no cover
logger.exception(f"Error during initialization: {e}")
# Continue execution even if initialization fails
# The command might still work, or will fail with a
# more specific error message
async def _init_and_cleanup():
"""Initialize app and clean up database connections.
Database connections created during initialization must be cleaned up
before the event loop closes, otherwise the process will hang indefinitely.
"""
try:
await initialize_app(app_config)
finally:
# Always cleanup database connections to prevent process hang
await db.shutdown_db()
# On Windows, use SelectorEventLoop to avoid ProactorEventLoop cleanup issues
# The ProactorEventLoop can raise "IndexError: pop from an empty deque" during
# event loop cleanup when there are pending handles. SelectorEventLoop is more
# stable for our use case (no subprocess pipes or named pipes needed).
if sys.platform == "win32":
asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
asyncio.run(_init_and_cleanup())
logger.info("Initialization completed successfully")
+24
View File
@@ -2,6 +2,7 @@
import asyncio
import os
import sys
import time
from collections import OrderedDict
from dataclasses import dataclass, field
@@ -1027,12 +1028,22 @@ class SyncService:
Uses subprocess to leverage OS-level file counting which is much faster
than Python iteration, especially on network filesystems like TigrisFS.
On Windows, subprocess is not supported with SelectorEventLoop (which we use
to avoid aiosqlite cleanup issues), so we fall back to Python-based counting.
Args:
directory: Directory to count files in
Returns:
Number of files in directory (recursive)
"""
# Windows with SelectorEventLoop doesn't support subprocess
if sys.platform == "win32":
count = 0
async for _ in self.scan_directory(directory):
count += 1
return count
process = await asyncio.create_subprocess_shell(
f'find "{directory}" -type f | wc -l',
stdout=asyncio.subprocess.PIPE,
@@ -1063,6 +1074,9 @@ class SyncService:
This is dramatically faster than scanning all files and comparing mtimes,
especially on network filesystems like TigrisFS where stat operations are expensive.
On Windows, subprocess is not supported with SelectorEventLoop (which we use
to avoid aiosqlite cleanup issues), so we implement mtime filtering in Python.
Args:
directory: Directory to scan
since_timestamp: Unix timestamp to find files newer than
@@ -1070,6 +1084,16 @@ class SyncService:
Returns:
List of relative file paths modified since the timestamp (respects .bmignore)
"""
# Windows with SelectorEventLoop doesn't support subprocess
# Implement mtime filtering in Python to preserve watermark optimization
if sys.platform == "win32":
file_paths = []
async for file_path_str, stat_info in self.scan_directory(directory):
if stat_info.st_mtime > since_timestamp:
rel_path = Path(file_path_str).relative_to(directory).as_posix()
file_paths.append(rel_path)
return file_paths
# Convert timestamp to find-compatible format
since_date = datetime.fromtimestamp(since_timestamp).strftime("%Y-%m-%d %H:%M:%S")
+3 -1
View File
@@ -232,6 +232,8 @@ def setup_logging(
if log_to_file:
log_path = Path.home() / ".basic-memory" / "basic-memory.log"
log_path.parent.mkdir(parents=True, exist_ok=True)
# Keep logging synchronous (enqueue=False) to avoid background logging threads.
# Background threads are a common source of "hang on exit" issues in CLI/test runs.
logger.add(
str(log_path),
level=log_level,
@@ -239,7 +241,7 @@ def setup_logging(
retention="10 days",
backtrace=True,
diagnose=True,
enqueue=True, # Thread-safe async logging
enqueue=False,
colorize=False,
)