fix: surface WebSocket close to tool calls instead of hanging (#1773)

* fix: surface WebSocket close to tool calls instead of hanging (#1721)

When the WebSocket to Home Assistant closed while a command awaited its
response, the pending future was cancelled; the resulting CancelledError
escaped every except-Exception handler and reached the MCP SDK, which
suppressed the response entirely. The MCP client then hung until its own
timeout (4 minutes in Claude clients) with only a bare 'WebSocket
connection closed' log line to go on.

- Fail pending request/event futures with HomeAssistantConnectionError
  (carrying the close code/reason) instead of cancelling them, so tools
  raise a structured error immediately.
- Log the close code and reason (e.g. 1009 message-too-big) instead of
  discarding them; WARNING for abnormal closes.
- Raise the client max_size 20MB -> 64MB, matching the Supervisor's own
  Core-connection limit: config/entity_registry/list arrives as one
  frame that scales with entity count (no pagination in the HA WS API)
  and overflowed 20MB on a ~6.4k-entity instance.
- Chunk config/entity_registry/get_entries alias enrichment (500 ids per
  call) so its response frame stays bounded on large instances.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* test: use future.result() to re-raise, satisfying py/ineffectual-statement

Bare 'await request_future' statements tripped CodeQL's
py/ineffectual-statement (await of a Name reads as an expression with no
effect). The futures are already done after reset_connection sets the
exception, so result() re-raises the stored exception directly -- a call
expression the rule does not flag, and a more direct assertion of the
state contract.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* fix: retrieve orphaned event-future exception; surface alias-fetch warnings

Review-driven hardening on top of the #1721 fix:

- send_command_with_event: a connection drop during the result phase now
  retrieves event_future's exception (reset_connection fails BOTH futures;
  only result_future was awaited), preventing asyncio's ERROR-level
  'Future exception was never retrieved' on the exact path this PR makes
  diagnosable. Regression-tested red/green via asyncio's retrieved-flag.
- _fetch_entity_aliases now returns (aliases_map, warnings): failed
  get_entries chunks surface as a top-level warning in search results
  instead of silently returning success with missing aliases.
- Close the untested branches flagged by review: self-initiated close
  (sent 1009 -- the literal #1721 trigger), WARNING-vs-INFO level split,
  abrupt no-close-frame drop, generic handler exception teardown,
  max_size wiring into websockets.connect, CancelledError propagation in
  the chunk loop, malformed success payload tolerance.
- Docstring: drop hardcoded 500 duplicating _GET_ENTRIES_CHUNK_SIZE.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

---------

Co-authored-by: kingpanther13 <kingpanther13@users.noreply.github.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: kingpanther13

src/ha_mcp/client/websocket_client.py

@@ -30,6 +30,15 @@
 
 logger = logging.getLogger(__name__)
 
+# Matches the Supervisor's own Core-connection receive limit
+# (MAX_MESSAGE_SIZE_FROM_CORE in home-assistant/supervisor, see supervisor
+# issue #4392). On the add-on path frames above that limit die at the
+# Supervisor proxy anyway, so a larger client-side cap adds nothing there.
+# Registry list responses scale with entity count and arrive as ONE frame
+# (the HA WebSocket API has no pagination); a ~6.4k-entity instance
+# overflowed the previous 20MB cap (#1721).
+MAX_WS_MESSAGE_BYTES = 64 * 1024 * 1024
+
 
 class WebSocketConnectionState:
     """Encapsulates mutable state used by the WebSocket client."""
@@ -109,20 +118,43 @@ def consume_auth_message(self, message_type: str) -> dict[str, Any] | None:
         """Retrieve and remove an authentication message if present."""
         return self._auth_messages.pop(message_type, None)
 
-    def reset_connection(self) -> None:
-        """Reset connection-specific state while preserving handlers."""
+    def reset_connection(self, close_reason: str | None = None) -> None:
+        """Reset connection-specific state while preserving handlers.
+
+        Args:
+            close_reason: Human-readable description of why the connection
+                went away (e.g. a close code/reason pair), included in the
+                error surfaced to any request still awaiting a response.
+        """
         self.connected = False
         self.authenticated = False
         self._message_id = 0
 
+        # ``future.cancel()`` makes awaiters see ``asyncio.CancelledError``,
+        # a BaseException that skips every ``except Exception`` handler in
+        # the tool layer and reaches the MCP SDK, which treats it as a
+        # client-initiated cancellation, suppresses the response entirely,
+        # and leaves the MCP client hanging until its own timeout (#1721).
+        # ``set_exception`` with a normal exception propagates through
+        # ``except Exception`` as expected instead.
+        message = (
+            "WebSocket connection to Home Assistant closed while waiting for a response"
+        )
+        if close_reason:
+            message = f"{message} ({close_reason})"
+
         for future in self._pending_requests.values():
             if not future.done():
-                future.cancel()
+                # A fresh instance per future: sharing one exception object
+                # across multiple ``set_exception`` calls means the second
+                # raise attaches a traceback to an exception already
+                # associated with another future's stack.
+                future.set_exception(HomeAssistantConnectionError(message))
         self._pending_requests.clear()
 
         for future in self._event_responses.values():
             if not future.done():
-                future.cancel()
+                future.set_exception(HomeAssistantConnectionError(message))
         self._event_responses.clear()
 
         # Drop any subscription queues — readers wake on the close signal
@@ -144,9 +176,9 @@ def mark_authenticated(self) -> None:
         """Mark the socket as authenticated and ready for commands."""
         self.authenticated = True
 
-    def mark_disconnected(self) -> None:
+    def mark_disconnected(self, close_reason: str | None = None) -> None:
         """Reset connection state when the socket is closed."""
-        self.reset_connection()
+        self.reset_connection(close_reason)
 
     @property
     def is_ready(self) -> bool:
@@ -293,9 +325,7 @@ async def connect(self) -> bool:
                 ping_timeout=10,
                 additional_headers={"Authorization": f"Bearer {self.token}"},
                 ssl=ssl_ctx,
-                # Increase max message size to 20MB for large responses
-                # (e.g., HACS repository list can be 2MB+)
-                max_size=20 * 1024 * 1024,
+                max_size=MAX_WS_MESSAGE_BYTES,
             )
             self._state.mark_connected()
 
@@ -390,6 +420,11 @@ async def _message_handler(self) -> None:
         """Background task to handle incoming WebSocket messages."""
         if not self.websocket:
             raise Exception("WebSocket not connected")
+        # None means a clean exit (the async-for loop simply ended); the
+        # except blocks below fill this in so pending futures — and the
+        # log — carry *why* the connection went away instead of a bare
+        # "WebSocket connection closed" that gave no lead on #1721.
+        close_reason: str | None = None
         try:
             async for message in self.websocket:
                 try:
@@ -400,12 +435,30 @@ async def _message_handler(self) -> None:
                     logger.error(f"Invalid JSON received: {e}")
                 except Exception as e:
                     logger.error(f"Error processing message: {e}")
-        except websockets.exceptions.ConnectionClosed:
-            logger.info("WebSocket connection closed")
+        except websockets.exceptions.ConnectionClosed as e:
+            # Prefer the frame we received (the peer closed on us); fall
+            # back to the frame we sent (we failed the connection
+            # ourselves, e.g. an over-max_size frame produces a sent
+            # Close(1009, ...) with no frame received at all).
+            close = e.rcvd if e.rcvd is not None else e.sent
+            if close is not None:
+                verb = "received" if e.rcvd is not None else "sent"
+                close_reason = f"{verb} close code {close.code}"
+                if close.reason:
+                    close_reason = f"{close_reason} ({close.reason})"
+            else:
+                close_reason = "connection dropped without a close frame"
+
+            log_message = f"WebSocket connection closed ({close_reason})"
+            if close is None or close.code not in (1000, 1001):
+                logger.warning(log_message)
+            else:
+                logger.info(log_message)
         except Exception as e:
+            close_reason = str(e)
             logger.error(f"WebSocket message handler error: {e}")
         finally:
-            self._state.mark_disconnected()
+            self._state.mark_disconnected(close_reason)
 
     async def _process_message(self, data: dict[str, Any]) -> None:
         """Process incoming WebSocket message."""
@@ -649,9 +702,15 @@ async def send_command_with_event(
             result_response = await asyncio.wait_for(
                 result_future, timeout=wait_timeout
             )
-        except TimeoutError:
+        except BaseException:
             self.cancel_pending_response(message_id)
             self.cancel_event_response(message_id)
+            # A connection drop fails BOTH futures via reset_connection;
+            # only result_future gets awaited on this path, so retrieve
+            # event_future's exception too or asyncio logs an ERROR-level
+            # "Future exception was never retrieved" when it is GC'd.
+            if event_future.done() and not event_future.cancelled():
+                event_future.exception()
             raise
 
         if not result_response.get("success"):

src/ha_mcp/tools/smart_search/_entities.py

@@ -12,6 +12,11 @@
 
 logger = logging.getLogger(__name__)
 
+# Bounds the per-frame response size of config/entity_registry/get_entries
+# (extended entries, ~1KB typical each) so alias enrichment can't produce an
+# over-cap WebSocket frame on large instances (#1721).
+_GET_ENTRIES_CHUNK_SIZE = 500
+
 
 class EntitySearchMixin(_SearchBase):
     """``smart_entity_search`` and ``get_entities_by_area`` plus helpers."""
@@ -144,41 +149,89 @@ def _filter_hidden_entities(
 
     async def _fetch_entity_aliases(
         self, survivor_ids: list[str]
-    ) -> dict[str, list[str]]:
+    ) -> tuple[dict[str, list[str]], list[str]]:
         """Batch-fetch full registry entries for aliases.
 
         ``config/entity_registry/list`` deliberately omits ``aliases``;
-        ``get_entries`` includes them. One extra round-trip enriches the
-        survivor set without N+1 fan-out.
+        ``get_entries`` includes them. Survivors are split into
+        ``_GET_ENTRIES_CHUNK_SIZE``-sized chunks fetched concurrently: a
+        bounded number of round-trips (one per chunk), not N+1 fan-out, and
+        each response frame stays bounded regardless of instance size.
+        Enrichment is best-effort per chunk — one chunk failing does not
+        drop aliases fetched by the others; failures are reported in the
+        returned warnings so the caller can tell "no aliases exist" apart
+        from "the alias fetch failed".
+
+        Returns:
+            ``(aliases_map, warnings)`` — warnings has one entry when any
+            chunk failed, empty otherwise.
         """
         aliases_map: dict[str, list[str]] = {}
         if not survivor_ids:
-            return aliases_map
-        try:
-            entries_resp = await self.client.send_websocket_message(
-                {
-                    "type": "config/entity_registry/get_entries",
-                    "entity_ids": survivor_ids,
-                }
-            )
-            if isinstance(entries_resp, dict) and entries_resp.get("success"):
-                for eid, entry in (entries_resp.get("result", {}) or {}).items():
-                    if isinstance(entry, dict):
-                        aliases_map[eid] = entry.get("aliases", []) or []
-            else:
+            return aliases_map, []
+
+        chunks: list[list[str]] = [
+            survivor_ids[i : i + _GET_ENTRIES_CHUNK_SIZE]
+            for i in range(0, len(survivor_ids), _GET_ENTRIES_CHUNK_SIZE)
+        ]
+        responses: list[Any] = await asyncio.gather(
+            *(
+                self.client.send_websocket_message(
+                    {
+                        "type": "config/entity_registry/get_entries",
+                        "entity_ids": chunk,
+                    }
+                )
+                for chunk in chunks
+            ),
+            return_exceptions=True,
+        )
+
+        failed_chunks = 0
+        for chunk, entries_resp in zip(chunks, responses, strict=True):
+            # Same convention as _fetch_search_entities: a captured
+            # CancelledError means the surrounding task is being torn
+            # down — propagate it instead of degrading to a warning.
+            if isinstance(entries_resp, asyncio.CancelledError):
+                raise entries_resp
+            if isinstance(entries_resp, BaseException):
+                failed_chunks += 1
                 logger.warning(
-                    "alias_enrichment_failed: get_entries returned non-success "
-                    "for %d entities (resp=%r)",
-                    len(survivor_ids),
+                    "alias_enrichment_failed: get_entries chunk of %d entities "
+                    "raised (err=%r)",
+                    len(chunk),
                     entries_resp,
                 )
-        except (KeyError, TypeError, AttributeError) as alias_err:
-            logger.warning(
-                "alias_enrichment_failed: malformed payload for %d entities (err=%r)",
-                len(survivor_ids),
-                alias_err,
+                continue
+            try:
+                if isinstance(entries_resp, dict) and entries_resp.get("success"):
+                    for eid, entry in (entries_resp.get("result", {}) or {}).items():
+                        if isinstance(entry, dict):
+                            aliases_map[eid] = entry.get("aliases", []) or []
+                else:
+                    failed_chunks += 1
+                    logger.warning(
+                        "alias_enrichment_failed: get_entries returned non-success "
+                        "for a chunk of %d entities (resp=%r)",
+                        len(chunk),
+                        entries_resp,
+                    )
+            except (KeyError, TypeError, AttributeError) as alias_err:
+                failed_chunks += 1
+                logger.warning(
+                    "alias_enrichment_failed: malformed payload for a chunk of "
+                    "%d entities (err=%r)",
+                    len(chunk),
+                    alias_err,
+                )
+        warnings: list[str] = []
+        if failed_chunks:
+            warnings.append(
+                f"Alias enrichment incomplete: {failed_chunks} of {len(chunks)} "
+                "entity-registry lookups failed; alias-based matches may be "
+                "missing from these results."
             )
-        return aliases_map
+        return aliases_map, warnings
 
     async def _fetch_search_entities(
         self, domain_filter: str | None, include_hidden: bool
@@ -223,7 +276,8 @@ async def _fetch_search_entities(
         survivor_ids, survivor_states = self._filter_hidden_entities(
             entities, registry_slim, include_hidden, visibility_hidden
         )
-        aliases_map = await self._fetch_entity_aliases(survivor_ids)
+        aliases_map, alias_warnings = await self._fetch_entity_aliases(survivor_ids)
+        visibility_warnings = [*visibility_warnings, *alias_warnings]
 
         # Enrich with aliases + hidden_by for the fuzzy layer. Shallow copy +
         # private-prefixed keys so downstream consumers that round-trip these