butler/mcp/async_runner.py

Rama main · réplica de lectura del repositorio original.

"""Dedicated asyncio loop for MCP SDK calls from sync Butler code."""

from __future__ import annotations

import asyncio
import atexit
import logging
import signal
import threading
from typing import Any, Coroutine, TypeVar

T = TypeVar("T")

logger = logging.getLogger(__name__)

_loop: asyncio.AbstractEventLoop | None = None
_thread: threading.Thread | None = None
_lock = threading.RLock()
_atexit_registered = False
_signal_registered = False
_prev_signal_handlers: dict[int, Any] = {}
_shutdown_once = threading.Lock()
_shutdown_done = False


def _ensure_loop() -> asyncio.AbstractEventLoop:
    global _loop, _thread, _atexit_registered
    with _lock:
        if _loop is not None and _loop.is_running():
            return _loop
        loop = asyncio.new_event_loop()

        def _run() -> None:
            from butler.mcp.async_runner_ops import close_event_loop_safe

            asyncio.set_event_loop(loop)
            try:
                loop.run_forever()
            finally:
                close_event_loop_safe(loop)

        _thread = threading.Thread(target=_run, daemon=True, name="butler-mcp-loop")
        _thread.start()
        _loop = loop
        if not _atexit_registered:
            atexit.register(_atexit_shutdown)
            _atexit_registered = True
        _register_signal_shutdown()
        return loop


def run_mcp_async(coro: Coroutine[Any, Any, T], *, timeout: float = 120.0) -> T:
    loop = _ensure_loop()
    future = asyncio.run_coroutine_threadsafe(coro, loop)
    return future.result(timeout=timeout)


def shutdown_async_runner(*, timeout: float = 5.0) -> bool:
    """Stop the MCP async loop, cancel pending tasks, and join the thread.

    Returns True if the thread was cleanly joined within ``timeout``.
    Idempotent: safe to call multiple times; safe to call before
    ``_ensure_loop()`` was ever invoked.
    """
    from butler.mcp.async_runner_ops import (
        close_loop_if_open_safe,
        drain_pending_tasks_safe,
    )

    global _loop, _thread
    with _lock:
        if _loop is None or _thread is None:
            return True
        loop = _loop
        thread = _thread

    try:
        loop.call_soon_threadsafe(loop.stop)
    except RuntimeError:
        pass

    thread.join(timeout=timeout)
    joined = not thread.is_alive()
    if not joined:
        logger.warning(
            "async_runner shutdown timeout: 守护线程未在 %.1fs 内退出, "
            "残留资源可能未被清理",
            timeout,
        )

    drain_pending_tasks_safe(loop)
    close_loop_if_open_safe(loop)

    if joined:
        with _lock:
            if _loop is loop and _thread is thread:
                _loop = None
                _thread = None
    return joined


def _disconnect_mcp_connections() -> None:
    """Run MCP transport cleanup while the dedicated loop is still alive."""
    from butler.mcp.async_runner_ops import disconnect_mcp_connections_safe

    disconnect_mcp_connections_safe()


def graceful_shutdown_mcp_stack(*, timeout: float = 5.0) -> bool:
    """Disconnect MCP stdio/http sessions, then stop the MCP asyncio loop."""
    global _shutdown_done
    with _shutdown_once:
        if _shutdown_done:
            return True
        _shutdown_done = True
    _disconnect_mcp_connections()
    return shutdown_async_runner(timeout=timeout)


def _register_signal_shutdown() -> None:
    """Register SIGTERM/SIGINT to drain MCP before loop teardown (R5-13)."""
    global _signal_registered
    with _lock:
        if _signal_registered:
            return
        _signal_registered = True

    def _handler(signum: int, frame: Any) -> None:
        from butler.mcp.async_runner_ops import (
            chain_signal_handler_safe,
            signal_shutdown_safe,
        )

        signal_shutdown_safe()
        prev = _prev_signal_handlers.get(signum)
        chain_signal_handler_safe(prev, signum, frame)

    for sig_name in ("SIGTERM", "SIGINT"):
        sig = getattr(signal, sig_name, None)
        if sig is None:
            continue
        try:
            _prev_signal_handlers[sig] = signal.getsignal(sig)
            signal.signal(sig, _handler)
        except (ValueError, OSError) as exc:
            logger.debug("async_runner signal register %s skipped: %s", sig_name, exc)


def _atexit_shutdown() -> None:
    """atexit 钩子: 进程退出前清理 MCP 异步线程。

    用较短的 timeout (2s) 因为 atexit 阶段不应长时间阻塞进程退出。
    """
    from butler.mcp.async_runner_ops import atexit_shutdown_safe

    atexit_shutdown_safe()