Files
basicmachines-co-basic-memory/src/basic_memory/cli/commands/command_utils.py
T
Paul Hernandez 863e0a4e24 fix: prevent CLI commands from hanging on exit (Python 3.14) (#505)
Signed-off-by: phernandez <paul@basicmachines.co>
Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
2026-01-11 16:47:35 -06:00

104 lines
3.5 KiB
Python

"""utility functions for commands"""
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
from basic_memory.mcp.project_context import get_active_project
from basic_memory.schemas import ProjectInfoResponse
from basic_memory.telemetry import shutdown_telemetry
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 and telemetry threads 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()
# Shutdown telemetry to stop the OpenPanel background thread
# This prevents hangs on Python 3.14+ during thread shutdown
shutdown_telemetry()
return asyncio.run(_with_cleanup())
async def run_sync(
project: Optional[str] = None,
force_full: bool = False,
run_in_background: bool = True,
):
"""Run sync operation via API endpoint.
Args:
project: Optional project name
force_full: If True, force a full scan bypassing watermark optimization
run_in_background: If True, return immediately; if False, wait for completion
"""
try:
async with get_client() as client:
project_item = await get_active_project(client, project, None)
url = f"{project_item.project_url}/project/sync"
params = []
if force_full:
params.append("force_full=true")
if not run_in_background:
params.append("run_in_background=false")
if params:
url += "?" + "&".join(params)
response = await call_post(client, url)
data = response.json()
# Background mode returns {"message": "..."}, foreground returns SyncReportResponse
if "message" in data:
console.print(f"[green]{data['message']}[/green]")
else:
# Foreground mode - show summary of sync results
total = data.get("total", 0)
new_count = len(data.get("new", []))
modified_count = len(data.get("modified", []))
deleted_count = len(data.get("deleted", []))
console.print(
f"[green]Synced {total} files[/green] "
f"(new: {new_count}, modified: {modified_count}, deleted: {deleted_count})"
)
except (ToolError, ValueError) as e:
console.print(f"[red]Sync failed: {e}[/red]")
raise typer.Exit(1)
async def get_project_info(project: str):
"""Get project information via API endpoint."""
try:
async with get_client() as client:
project_item = await get_active_project(client, project, None)
response = await call_get(client, f"{project_item.project_url}/project/info")
return ProjectInfoResponse.model_validate(response.json())
except (ToolError, ValueError) as e:
console.print(f"[red]Sync failed: {e}[/red]")
raise typer.Exit(1)