Skip to content

Commit 916baf8

Browse files
phernandezclaude
phernandez
and
claude
authored
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>
1 parent 95937c6 commit 916baf8

17 files changed

Lines changed: 403 additions & 273 deletions

‎src/basic_memory/alembic/versions/f8a9b2c3d4e5_add_pg_trgm_for_fuzzy_link_resolution.py‎

Lines changed: 116 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,42 @@
1010

1111
import sqlalchemy as sa
1212
from alembic import op
13+
from sqlalchemy import text
14+
15+
16+
def column_exists(connection, table: str, column: str) -> bool:
17+
"""Check if a column exists in a table (idempotent migration support)."""
18+
if connection.dialect.name == "postgresql":
19+
result = connection.execute(
20+
text(
21+
"SELECT 1 FROM information_schema.columns "
22+
"WHERE table_name = :table AND column_name = :column"
23+
),
24+
{"table": table, "column": column},
25+
)
26+
return result.fetchone() is not None
27+
else:
28+
# SQLite
29+
result = connection.execute(text(f"PRAGMA table_info({table})"))
30+
columns = [row[1] for row in result]
31+
return column in columns
32+
33+
34+
def index_exists(connection, index_name: str) -> bool:
35+
"""Check if an index exists (idempotent migration support)."""
36+
if connection.dialect.name == "postgresql":
37+
result = connection.execute(
38+
text("SELECT 1 FROM pg_indexes WHERE indexname = :index_name"),
39+
{"index_name": index_name},
40+
)
41+
return result.fetchone() is not None
42+
else:
43+
# SQLite
44+
result = connection.execute(
45+
text("SELECT 1 FROM sqlite_master WHERE type='index' AND name = :index_name"),
46+
{"index_name": index_name},
47+
)
48+
return result.fetchone() is not None
1349

1450

1551
# revision identifiers, used by Alembic.
@@ -36,101 +72,105 @@ def upgrade() -> None:
3672
# Add project_id to relation table
3773
# -------------------------------------------------------------------------
3874

39-
# Step 1: Add project_id column as nullable first
40-
op.add_column("relation", sa.Column("project_id", sa.Integer(), nullable=True))
75+
# Step 1: Add project_id column as nullable first (idempotent)
76+
if not column_exists(connection, "relation", "project_id"):
77+
op.add_column("relation", sa.Column("project_id", sa.Integer(), nullable=True))
4178

42-
# Step 2: Backfill project_id from entity.project_id via from_id
43-
if dialect == "postgresql":
44-
op.execute("""
45-
UPDATE relation
46-
SET project_id = entity.project_id
47-
FROM entity
48-
WHERE relation.from_id = entity.id
49-
""")
50-
else:
51-
# SQLite syntax
52-
op.execute("""
53-
UPDATE relation
54-
SET project_id = (
55-
SELECT entity.project_id
79+
# Step 2: Backfill project_id from entity.project_id via from_id
80+
if dialect == "postgresql":
81+
op.execute("""
82+
UPDATE relation
83+
SET project_id = entity.project_id
5684
FROM entity
57-
WHERE entity.id = relation.from_id
58-
)
59-
""")
60-
61-
# Step 3: Make project_id NOT NULL and add foreign key
62-
if dialect == "postgresql":
63-
op.alter_column("relation", "project_id", nullable=False)
64-
op.create_foreign_key(
65-
"fk_relation_project_id",
66-
"relation",
67-
"project",
68-
["project_id"],
69-
["id"],
70-
)
71-
else:
72-
# SQLite requires batch operations for ALTER COLUMN
73-
with op.batch_alter_table("relation") as batch_op:
74-
batch_op.alter_column("project_id", nullable=False)
75-
batch_op.create_foreign_key(
85+
WHERE relation.from_id = entity.id
86+
""")
87+
else:
88+
# SQLite syntax
89+
op.execute("""
90+
UPDATE relation
91+
SET project_id = (
92+
SELECT entity.project_id
93+
FROM entity
94+
WHERE entity.id = relation.from_id
95+
)
96+
""")
97+
98+
# Step 3: Make project_id NOT NULL and add foreign key
99+
if dialect == "postgresql":
100+
op.alter_column("relation", "project_id", nullable=False)
101+
op.create_foreign_key(
76102
"fk_relation_project_id",
103+
"relation",
77104
"project",
78105
["project_id"],
79106
["id"],
80107
)
81-
82-
# Step 4: Create index on relation.project_id
83-
op.create_index("ix_relation_project_id", "relation", ["project_id"])
108+
else:
109+
# SQLite requires batch operations for ALTER COLUMN
110+
with op.batch_alter_table("relation") as batch_op:
111+
batch_op.alter_column("project_id", nullable=False)
112+
batch_op.create_foreign_key(
113+
"fk_relation_project_id",
114+
"project",
115+
["project_id"],
116+
["id"],
117+
)
118+
119+
# Step 4: Create index on relation.project_id (idempotent)
120+
if not index_exists(connection, "ix_relation_project_id"):
121+
op.create_index("ix_relation_project_id", "relation", ["project_id"])
84122

85123
# -------------------------------------------------------------------------
86124
# Add project_id to observation table
87125
# -------------------------------------------------------------------------
88126

89-
# Step 1: Add project_id column as nullable first
90-
op.add_column("observation", sa.Column("project_id", sa.Integer(), nullable=True))
127+
# Step 1: Add project_id column as nullable first (idempotent)
128+
if not column_exists(connection, "observation", "project_id"):
129+
op.add_column("observation", sa.Column("project_id", sa.Integer(), nullable=True))
91130

92-
# Step 2: Backfill project_id from entity.project_id via entity_id
93-
if dialect == "postgresql":
94-
op.execute("""
95-
UPDATE observation
96-
SET project_id = entity.project_id
97-
FROM entity
98-
WHERE observation.entity_id = entity.id
99-
""")
100-
else:
101-
# SQLite syntax
102-
op.execute("""
103-
UPDATE observation
104-
SET project_id = (
105-
SELECT entity.project_id
131+
# Step 2: Backfill project_id from entity.project_id via entity_id
132+
if dialect == "postgresql":
133+
op.execute("""
134+
UPDATE observation
135+
SET project_id = entity.project_id
106136
FROM entity
107-
WHERE entity.id = observation.entity_id
108-
)
109-
""")
110-
111-
# Step 3: Make project_id NOT NULL and add foreign key
112-
if dialect == "postgresql":
113-
op.alter_column("observation", "project_id", nullable=False)
114-
op.create_foreign_key(
115-
"fk_observation_project_id",
116-
"observation",
117-
"project",
118-
["project_id"],
119-
["id"],
120-
)
121-
else:
122-
# SQLite requires batch operations for ALTER COLUMN
123-
with op.batch_alter_table("observation") as batch_op:
124-
batch_op.alter_column("project_id", nullable=False)
125-
batch_op.create_foreign_key(
137+
WHERE observation.entity_id = entity.id
138+
""")
139+
else:
140+
# SQLite syntax
141+
op.execute("""
142+
UPDATE observation
143+
SET project_id = (
144+
SELECT entity.project_id
145+
FROM entity
146+
WHERE entity.id = observation.entity_id
147+
)
148+
""")
149+
150+
# Step 3: Make project_id NOT NULL and add foreign key
151+
if dialect == "postgresql":
152+
op.alter_column("observation", "project_id", nullable=False)
153+
op.create_foreign_key(
126154
"fk_observation_project_id",
155+
"observation",
127156
"project",
128157
["project_id"],
129158
["id"],
130159
)
131-
132-
# Step 4: Create index on observation.project_id
133-
op.create_index("ix_observation_project_id", "observation", ["project_id"])
160+
else:
161+
# SQLite requires batch operations for ALTER COLUMN
162+
with op.batch_alter_table("observation") as batch_op:
163+
batch_op.alter_column("project_id", nullable=False)
164+
batch_op.create_foreign_key(
165+
"fk_observation_project_id",
166+
"project",
167+
["project_id"],
168+
["id"],
169+
)
170+
171+
# Step 4: Create index on observation.project_id (idempotent)
172+
if not index_exists(connection, "ix_observation_project_id"):
173+
op.create_index("ix_observation_project_id", "observation", ["project_id"])
134174

135175
# Postgres-specific: pg_trgm and GIN indexes
136176
if dialect == "postgresql":

‎src/basic_memory/api/app.py‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
"""FastAPI application for basic-memory knowledge graph API."""
22

33
import asyncio
4+
import os
45
from contextlib import asynccontextmanager
56

67
from fastapi import FastAPI, HTTPException
@@ -53,12 +54,25 @@ async def lifespan(app: FastAPI): # pragma: no cover
5354
app.state.session_maker = session_maker
5455
logger.info("Database connections cached in app state")
5556

56-
logger.info(f"Sync changes enabled: {app_config.sync_changes}")
57-
if app_config.sync_changes:
57+
# Start file sync if enabled
58+
is_test_env = (
59+
app_config.env == "test"
60+
or os.getenv("BASIC_MEMORY_ENV", "").lower() == "test"
61+
or os.getenv("PYTEST_CURRENT_TEST") is not None
62+
)
63+
if app_config.sync_changes and not is_test_env:
64+
logger.info(f"Sync changes enabled: {app_config.sync_changes}")
65+
5866
# start file sync task in background
59-
app.state.sync_task = asyncio.create_task(initialize_file_sync(app_config))
67+
async def _file_sync_runner() -> None:
68+
await initialize_file_sync(app_config)
69+
70+
app.state.sync_task = asyncio.create_task(_file_sync_runner())
6071
else:
61-
logger.info("Sync changes disabled. Skipping file sync service.")
72+
if is_test_env:
73+
logger.info("Test environment detected. Skipping file sync service.")
74+
else:
75+
logger.info("Sync changes disabled. Skipping file sync service.")
6276
app.state.sync_task = None
6377

6478
# proceed with startup

‎src/basic_memory/cli/app.py‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,8 +34,15 @@ def app_callback(
3434
# Initialize logging for CLI (file only, no stdout)
3535
init_cli_logging()
3636

37-
# Run initialization for every command unless --version was specified
38-
if not version and ctx.invoked_subcommand is not None:
37+
# Run initialization for commands that don't use the API
38+
# Skip for 'mcp' command - it has its own lifespan that handles initialization
39+
# Skip for API-using commands (status, sync, etc.) - they handle initialization via deps.py
40+
api_commands = {"mcp", "status", "sync", "project", "tool"}
41+
if (
42+
not version
43+
and ctx.invoked_subcommand is not None
44+
and ctx.invoked_subcommand not in api_commands
45+
):
3946
from basic_memory.services.initialization import ensure_initialization
4047

4148
app_config = ConfigManager().config

‎src/basic_memory/cli/commands/command_utils.py‎

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,14 @@
11
"""utility functions for commands"""
22

3-
from typing import Optional
3+
import asyncio
4+
from typing import Optional, TypeVar, Coroutine, Any
45

56
from mcp.server.fastmcp.exceptions import ToolError
67
import typer
78

89
from rich.console import Console
910

11+
from basic_memory import db
1012
from basic_memory.mcp.async_client import get_client
1113

1214
from basic_memory.mcp.tools.utils import call_post, call_get
@@ -15,6 +17,30 @@
1517

1618
console = Console()
1719

20+
T = TypeVar("T")
21+
22+
23+
def run_with_cleanup(coro: Coroutine[Any, Any, T]) -> T:
24+
"""Run an async coroutine with proper database cleanup.
25+
26+
This helper ensures database connections are cleaned up before the event
27+
loop closes, preventing process hangs in CLI commands.
28+
29+
Args:
30+
coro: The coroutine to run
31+
32+
Returns:
33+
The result of the coroutine
34+
"""
35+
36+
async def _with_cleanup() -> T:
37+
try:
38+
return await coro
39+
finally:
40+
await db.shutdown_db()
41+
42+
return asyncio.run(_with_cleanup())
43+
1844

1945
async def run_sync(project: Optional[str] = None, force_full: bool = False):
2046
"""Run sync operation via API endpoint.

‎src/basic_memory/cli/commands/mcp.py‎

Lines changed: 5 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,13 @@
11
"""MCP server command with streamable HTTP transport."""
22

3-
import asyncio
43
import os
54
import typer
65
from typing import Optional
76

87
from basic_memory.cli.app import app
98
from basic_memory.config import ConfigManager, init_mcp_logging
109

11-
# Import mcp instance
10+
# Import mcp instance (has lifespan that handles initialization and file sync)
1211
from basic_memory.mcp.server import mcp as mcp_server # pragma: no cover
1312

1413
# Import mcp tools to register them
@@ -17,8 +16,6 @@
1716
# Import prompts to register them
1817
import basic_memory.mcp.prompts # noqa: F401 # pragma: no cover
1918
from loguru import logger
20-
import threading
21-
from basic_memory.services.initialization import initialize_file_sync
2219

2320
config = ConfigManager().config
2421

@@ -43,6 +40,8 @@ def mcp(
4340
- stdio: Standard I/O (good for local usage)
4441
- streamable-http: Recommended for web deployments (default)
4542
- sse: Server-Sent Events (for compatibility with existing clients)
43+
44+
Initialization, file sync, and cleanup are handled by the MCP server's lifespan.
4645
"""
4746
# Initialize logging for MCP (file only, stdout breaks protocol)
4847
init_mcp_logging()
@@ -59,27 +58,8 @@ def mcp(
5958
os.environ["BASIC_MEMORY_MCP_PROJECT"] = project_name
6059
logger.info(f"MCP server constrained to project: {project_name}")
6160

62-
app_config = ConfigManager().config
63-
64-
def run_file_sync():
65-
"""Run file sync in a separate thread with its own event loop."""
66-
loop = asyncio.new_event_loop()
67-
asyncio.set_event_loop(loop)
68-
try:
69-
loop.run_until_complete(initialize_file_sync(app_config))
70-
except Exception as e:
71-
logger.error(f"File sync error: {e}", err=True)
72-
finally:
73-
loop.close()
74-
75-
logger.info(f"Sync changes enabled: {app_config.sync_changes}")
76-
if app_config.sync_changes:
77-
# Start the sync thread
78-
sync_thread = threading.Thread(target=run_file_sync, daemon=True)
79-
sync_thread.start()
80-
logger.info("Started file sync in background")
81-
82-
# Now run the MCP server (blocks)
61+
# Run the MCP server (blocks)
62+
# Lifespan handles: initialization, migrations, file sync, cleanup
8363
logger.info(f"Starting MCP server with {transport.upper()} transport")
8464

8565
if transport == "stdio":

0 commit comments

Comments
 (0)