23 files

2026-08-05 23:27:04 +00:00
parent 32a0b6c7f3
commit be7e50610a
23 changed files with 229 additions and 77 deletions
@@ -38,7 +38,7 @@ import threading
 from contextlib import suppress
 from datetime import timedelta
 from functools import partial
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING, Any, Literal

 from homeassistant.auth.const import GROUP_ID_ADMIN
 from homeassistant.auth.models import TOKEN_TYPE_LONG_LIVED_ACCESS_TOKEN
@@ -116,6 +116,16 @@ _READY_POLL_INTERVAL_SECONDS = 0.5
 # leaking it rather than blocking HA shutdown.
 _STOP_JOIN_TIMEOUT_SECONDS = 10.0

+# Budget for each teardown phase (the _serve resource cleanup and the
+# worker-loop pending-task sweep). Mirrors the CLI runner's
+# SHUTDOWN_TIMEOUT_SECONDS: both phases together must finish inside
+# _STOP_JOIN_TIMEOUT_SECONDS, or async_stop declares the worker orphaned
+# while the old thread is still executing shared ha_mcp modules. The budget
+# bounds only the phases that accept one — asyncgen finalization and
+# uvicorn's post-drain lifespan shutdown remain unbounded — so it buys
+# headroom, not a hard ceiling on the join.
+_TEARDOWN_TIMEOUT_SECONDS = 2.0
+
 # Per-download HTTP timeout for a forced reinstall. The first install pulls the
 # whole fastmcp tree, well beyond HA's 60s requirements default.
 _PIP_INSTALL_TIMEOUT_SECONDS = 300
@@ -1201,24 +1211,7 @@ class EmbeddedServerManager:
        finally:
            with _IMPORTING_WORKERS_LOCK:
                _IMPORTING_WORKERS.discard(threading.current_thread())
-            # Teardown is best-effort but never SILENT (review finding): a
-            # raise here must not mask the primary outcome, yet a recurring
-            # cleanup failure (leaking executor threads across reloads) has
-            # to be visible in the logs. Each call gets its own guard so one
-            # failure cannot skip the other.
-            for _label, _coro_factory in (
-                ("asyncgen", loop.shutdown_asyncgens),
-                ("executor", loop.shutdown_default_executor),
-            ):
-                try:
-                    loop.run_until_complete(_coro_factory())
-                except Exception:
-                    _LOGGER.warning(
-                        "Worker-loop %s shutdown failed during teardown",
-                        _label,
-                        exc_info=True,
-                    )
-            loop.close()
+            _teardown_worker_loop(loop)

    async def _serve(self, access_token: str, stop_event: asyncio.Event) -> None:
        """Build the ha-mcp server and run it until a stop is signaled.
@@ -1396,23 +1389,33 @@ class EmbeddedServerManager:

        self._note_startup_phase("starting the HTTP listener")
        stop_task = asyncio.create_task(stop_event.wait())
-        async with server.mcp._lifespan_manager():
-            serve_task = asyncio.create_task(uv_server.serve())
-            done, _pending = await asyncio.wait(
-                {serve_task, stop_task}, return_when=asyncio.FIRST_COMPLETED
-            )
-            if stop_task in done:
-                # Graceful shutdown through uvicorn's own path: waits out
-                # in-flight requests (2s cap), runs lifespan shutdown, and
-                # deterministically releases the socket for the next bring-up.
-                uv_server.should_exit = True
-                await serve_task
-            else:
-                stop_task.cancel()
-                with suppress(asyncio.CancelledError):
-                    await stop_task
-                # Surface a server that exited on its own (bind failure, etc.).
-                serve_task.result()
+        try:
+            async with server.mcp._lifespan_manager():
+                serve_task = asyncio.create_task(uv_server.serve())
+                done, _pending = await asyncio.wait(
+                    {serve_task, stop_task}, return_when=asyncio.FIRST_COMPLETED
+                )
+                if stop_task in done:
+                    # Graceful shutdown through uvicorn's own path: waits out
+                    # in-flight requests (2s cap), runs lifespan shutdown, and
+                    # deterministically releases the socket for the next
+                    # bring-up.
+                    uv_server.should_exit = True
+                    await serve_task
+                else:
+                    stop_task.cancel()
+                    with suppress(asyncio.CancelledError):
+                        await stop_task
+                    # Surface a server that exited on its own (bind failure,
+                    # etc.).
+                    serve_task.result()
+        finally:
+            # CLI parity: the HTTP runner's shutdown path releases the served
+            # stack's HA connections; this in-process runner must too, on this
+            # loop, while it still runs — otherwise the reader tasks are only
+            # cancelled by the thread's loop teardown and their sockets are
+            # abandoned to garbage collection (issue #2127).
+            await _shutdown_server_resources_bounded(server)

    def _progress_signature(self) -> tuple[int, str]:
        """Snapshot the observable startup progress of the worker thread.
@@ -1526,6 +1529,149 @@ _IMPORTING_WORKERS_LOCK = threading.Lock()
 _IMPORTING_WORKERS: set[threading.Thread] = set()


+async def _shutdown_server_resources_bounded(server: Any) -> None:
+    """Run :func:`_shutdown_server_resources` inside the teardown budget.
+
+    An unresponsive peer's close handshake (websockets' 10s default
+    close_timeout) must not eat the whole ``_STOP_JOIN_TIMEOUT_SECONDS`` join
+    budget. Cancel-and-abandon, not ``wait_for``: ``wait_for`` awaits the
+    cancelled coroutine before raising, and the cleanup stack swallows
+    ``CancelledError`` at several layers (per-client in
+    ``WebSocketManager.disconnect``, in ``client.disconnect``'s own
+    task-cancel guard), so a straggler must be left to the thread's loop
+    teardown sweep instead of being joined here.
+    """
+    task = asyncio.ensure_future(_shutdown_server_resources(server))
+    _done, pending = await asyncio.wait({task}, timeout=_TEARDOWN_TIMEOUT_SECONDS)
+    if pending:
+        task.cancel()
+        _LOGGER.warning("Embedded resource cleanup timed out")
+
+
+async def _shutdown_server_resources(server: Any) -> None:
+    """Release the served stack's Home Assistant connections on its own loop.
+
+    Mirrors the CLI runner's ``_cleanup_resources`` (``ha_mcp.__main__``)
+    without importing it: stop the WebSocket listener service, disconnect the
+    pooled WebSocket clients, and close the server's HTTP client. Every step
+    guards independently — a failing step must not keep the next one from
+    running, and no failure here may mask the serve outcome.
+    """
+    try:
+        from ha_mcp.client.websocket_listener import stop_websocket_listener
+
+        await stop_websocket_listener()
+    except ImportError:
+        _LOGGER.debug("WebSocket listener module not available")
+    except Exception as err:
+        _LOGGER.warning("WebSocket listener cleanup failed: %s", err)
+
+    try:
+        from ha_mcp.client.websocket_client import websocket_manager
+
+        await websocket_manager.disconnect()
+    except ImportError:
+        _LOGGER.debug("WebSocket manager module not available")
+    except Exception as err:
+        _LOGGER.warning("WebSocket manager cleanup failed: %s", err)
+
+    try:
+        await server.close()
+    except Exception as err:
+        _LOGGER.warning("Server cleanup failed: %s", err)
+
+
+def _cancel_pending_tasks(loop: asyncio.AbstractEventLoop) -> None:
+    """Cancel every task still pending on ``loop`` and wait them out.
+
+    Mirrors ``asyncio.runners._cancel_all_tasks`` — the step ``asyncio.run``
+    performs between the main coroutine returning and asyncgen finalization,
+    which this worker's hand-rolled loop lifecycle skipped (issue #2127).
+    Without it, tasks the served stack leaves behind — WebSocket reader tasks
+    parked in ``Connection.__aiter__``, sse_starlette's ``_shutdown_watcher``
+    poll (unreachable by its uvicorn signal hooks on a non-main thread) — are
+    still pending at teardown: ``shutdown_asyncgens()`` then acloses
+    generators mid-``__anext__`` (``RuntimeError: aclose(): asynchronous
+    generator is already running``) and ``loop.close()`` destroys the
+    survivors ("Task was destroyed but it is pending!"), one error pair per
+    entry reload.
+
+    Abandoning is inherently partial: a task that ignores cancellation past
+    the budget and still drives an async generator leaves that generator
+    running, and ``shutdown_asyncgens()`` then reports the same ``aclose()``
+    error this sweep exists to remove. The ignored-cancellation warning
+    below is the tell when that residual fires.
+    """
+    pending = asyncio.all_tasks(loop)
+    if not pending:
+        return
+    for task in pending:
+        task.cancel()
+    # Bounded, unlike asyncio.runners: async_stop joins this worker for only
+    # _STOP_JOIN_TIMEOUT_SECONDS, so a task that ignores cancellation gets
+    # logged and abandoned rather than hanging the join (the CLI's
+    # _cancel_tasks does the same, issue #2027 precedent).
+    done, still_pending = loop.run_until_complete(
+        asyncio.wait(pending, timeout=_TEARDOWN_TIMEOUT_SECONDS)
+    )
+    if still_pending:
+        _LOGGER.warning(
+            "%d task(s) ignored cancellation during worker-loop teardown",
+            len(still_pending),
+        )
+    for task in done:
+        if task.cancelled():
+            continue
+        exc = task.exception()
+        if exc is not None:
+            _LOGGER.warning(
+                "Task %r raised during worker-loop teardown: %r",
+                task.get_name(),
+                exc,
+            )
+
+
+def _teardown_worker_loop(loop: asyncio.AbstractEventLoop) -> None:
+    """Drain and close the worker loop with ``asyncio.run`` teardown parity.
+
+    Teardown is best-effort but never SILENT (review finding): a raise here
+    must not mask the primary outcome, yet a recurring cleanup failure
+    (leaking executor threads across reloads) has to be visible in the logs.
+    Each step gets its own guard so one failure cannot skip the others.
+    """
+    try:
+        _cancel_pending_tasks(loop)
+    except Exception:
+        _LOGGER.warning(
+            "Worker-loop task cancellation failed during teardown",
+            exc_info=True,
+        )
+    for _label, _coro_factory in (
+        ("asyncgen", loop.shutdown_asyncgens),
+        # The executor join is bounded too: a stuck executor thread must not
+        # keep the worker alive past the join deadline (abandoning it emits
+        # a RuntimeWarning instead of hanging). The runtime has accepted
+        # timeout= since Python 3.12; typeshed's AbstractEventLoop signature
+        # lags behind, hence the scoped ignore.
+        (
+            "executor",
+            partial(
+                loop.shutdown_default_executor,
+                timeout=_TEARDOWN_TIMEOUT_SECONDS,  # type: ignore[call-arg]
+            ),
+        ),
+    ):
+        try:
+            loop.run_until_complete(_coro_factory())
+        except Exception:
+            _LOGGER.warning(
+                "Worker-loop %s shutdown failed during teardown",
+                _label,
+                exc_info=True,
+            )
+    loop.close()
+
+
 def _prune_and_check_importing_workers() -> bool:
    """Drop dead workers from the registry; return True if any live one remains."""
    with _IMPORTING_WORKERS_LOCK: