Add lifecycle management for beads-mcp processes (bd-148)
- Register atexit handler to close daemon connections - Add signal handlers for SIGTERM/SIGINT for graceful shutdown - Implement cleanup() to close all daemon client connections - Track daemon clients globally for cleanup - Add close() method to BdDaemonClient (no-op since connections are per-request) - Register client on first use via _get_client() - Add comprehensive lifecycle tests This prevents MCP server processes from accumulating without cleanup. Each tool invocation will now properly clean up on exit. Amp-Thread-ID: https://ampcode.com/threads/T-05d76b8e-dac9-472b-bfd0-afe10e3457cd Co-authored-by: Amp <amp@ampcode.com>
This commit is contained in:
@@ -1,8 +1,12 @@
|
||||
"""FastMCP server for beads issue tracker."""
|
||||
|
||||
import asyncio
|
||||
import atexit
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
from functools import wraps
|
||||
from typing import Callable, TypeVar
|
||||
|
||||
@@ -24,8 +28,19 @@ from beads_mcp.tools import (
|
||||
beads_update_issue,
|
||||
)
|
||||
|
||||
# Setup logging for lifecycle events
|
||||
logger = logging.getLogger(__name__)
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
)
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
# Global state for cleanup
|
||||
_daemon_clients: list = []
|
||||
_cleanup_done = False
|
||||
|
||||
# Create FastMCP server
|
||||
mcp = FastMCP(
|
||||
name="Beads",
|
||||
@@ -38,6 +53,49 @@ IMPORTANT: Call set_context with your workspace root before any write operations
|
||||
)
|
||||
|
||||
|
||||
def cleanup() -> None:
|
||||
"""Clean up resources on exit.
|
||||
|
||||
Closes daemon connections and removes temp files.
|
||||
Safe to call multiple times.
|
||||
"""
|
||||
global _cleanup_done
|
||||
|
||||
if _cleanup_done:
|
||||
return
|
||||
|
||||
_cleanup_done = True
|
||||
logger.info("Cleaning up beads-mcp resources...")
|
||||
|
||||
# Close all daemon client connections
|
||||
for client in _daemon_clients:
|
||||
try:
|
||||
if hasattr(client, 'close'):
|
||||
client.close()
|
||||
logger.debug(f"Closed daemon client: {client}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Error closing daemon client: {e}")
|
||||
|
||||
_daemon_clients.clear()
|
||||
logger.info("Cleanup complete")
|
||||
|
||||
|
||||
def signal_handler(signum: int, frame) -> None:
|
||||
"""Handle termination signals gracefully."""
|
||||
sig_name = signal.Signals(signum).name
|
||||
logger.info(f"Received {sig_name}, shutting down gracefully...")
|
||||
cleanup()
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
# Register cleanup handlers
|
||||
atexit.register(cleanup)
|
||||
signal.signal(signal.SIGTERM, signal_handler)
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
|
||||
logger.info("beads-mcp server initialized with lifecycle management")
|
||||
|
||||
|
||||
def require_context(func: Callable[..., T]) -> Callable[..., T]:
|
||||
"""Decorator to enforce context has been set before write operations.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user