ha-mcp-dev 9.1.1.dev3370__tar.gz → 9.1.1.dev3378__tar.gz

{ha_mcp_dev-9.1.1.dev3370 → ha_mcp_dev-9.1.1.dev3378}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ha-mcp-dev
-Version: 9.1.1.dev3370
+Version: 9.1.1.dev3378
 Summary: Home Assistant MCP Server - Complete control of Home Assistant through MCP
 Keywords: mcp,home-assistant,ai,automation,smart-home
 Author: Julien
@@ -442,7 +442,7 @@ A `ha_manage_*` tool combines several operations, so it is reachable from more t
 
 Leave it off in clients with deferred tool loading (claude.ai, Claude Desktop, Claude Code); the full catalog has no idle cost there, direct calls skip the search step, and the client's built-in tool search is the better choice — there is no benefit to running ha-mcp's on top of it. Whether tools are deferred depends on the client and model combination: the same model can behave differently per client — GitHub Copilot CLI running Claude Sonnet/Opus inlines the full catalog and still benefits from tool search here. Some Codex models and ChatGPT include deferred tools too — check your client/model directly to confirm its features so you don't leave this enabled unnecessarily.
 
-> 🔄 **Refresh your client's tool list after changing this (or any) setting.** Toggling `ENABLE_TOOL_SEARCH` (or changing pinned/disabled tools, Read Only Mode, etc.) changes the tools the server exposes, but your AI client keeps serving its **cached** tool list until it re-fetches. Restarting the app or Home Assistant does **not** refresh the client — reconnect or refresh the MCP server in your client (e.g. re-add/refresh the connector in ChatGPT, or close and reopen Claude Desktop). If you skip this, newly enabled tools won't appear in the client at all, and tools the server no longer exposes still show as available but fail when called (`Unknown tool`, or the Read Only Mode error for a write tool that mode hides). ChatGPT sometimes keeps serving the stale list even after the connector is removed and re-added under the same name — if tools are still missing after re-adding, delete the connector and create a new one with a **different name**.
+> 🔄 **Refresh your client's tool list after changing this (or any) setting.** Toggling `ENABLE_TOOL_SEARCH` (or changing pinned/disabled tools, Read Only Mode, etc.) changes the tools the server exposes, but your AI client keeps serving its **cached** tool list until it re-fetches. Restarting the app or Home Assistant does **not** refresh the client — reconnect or refresh the MCP server in your client (e.g. in ChatGPT open **Settings**, then the plugin settings, select the connector and click **Refresh tools**; or close and reopen Claude Desktop). If you skip this, newly enabled tools won't appear in the client at all, and tools the server no longer exposes still show as available but fail when called (`Unknown tool`, or the Read Only Mode error for a write tool that mode hides). ChatGPT sometimes keeps serving the stale list even after the connector is removed and re-added under the same name — if tools are still missing after a refresh, delete the connector and create a new one with a **different name**. The [FAQ](https://homeassistant-ai.github.io/ha-mcp/faq/#chatgpt-stale-tools) has the full steps.
 
 For the HA app, the same option is documented in [`homeassistant-addon/DOCS.md`](homeassistant-addon/DOCS.md#enable_tool_search) along with the in-app settings UI for fine-grained tool enable/disable/pin.
 

{ha_mcp_dev-9.1.1.dev3370 → ha_mcp_dev-9.1.1.dev3378}/README.md

@@ -381,7 +381,7 @@ A `ha_manage_*` tool combines several operations, so it is reachable from more t
 
 Leave it off in clients with deferred tool loading (claude.ai, Claude Desktop, Claude Code); the full catalog has no idle cost there, direct calls skip the search step, and the client's built-in tool search is the better choice — there is no benefit to running ha-mcp's on top of it. Whether tools are deferred depends on the client and model combination: the same model can behave differently per client — GitHub Copilot CLI running Claude Sonnet/Opus inlines the full catalog and still benefits from tool search here. Some Codex models and ChatGPT include deferred tools too — check your client/model directly to confirm its features so you don't leave this enabled unnecessarily.
 
-> 🔄 **Refresh your client's tool list after changing this (or any) setting.** Toggling `ENABLE_TOOL_SEARCH` (or changing pinned/disabled tools, Read Only Mode, etc.) changes the tools the server exposes, but your AI client keeps serving its **cached** tool list until it re-fetches. Restarting the app or Home Assistant does **not** refresh the client — reconnect or refresh the MCP server in your client (e.g. re-add/refresh the connector in ChatGPT, or close and reopen Claude Desktop). If you skip this, newly enabled tools won't appear in the client at all, and tools the server no longer exposes still show as available but fail when called (`Unknown tool`, or the Read Only Mode error for a write tool that mode hides). ChatGPT sometimes keeps serving the stale list even after the connector is removed and re-added under the same name — if tools are still missing after re-adding, delete the connector and create a new one with a **different name**.
+> 🔄 **Refresh your client's tool list after changing this (or any) setting.** Toggling `ENABLE_TOOL_SEARCH` (or changing pinned/disabled tools, Read Only Mode, etc.) changes the tools the server exposes, but your AI client keeps serving its **cached** tool list until it re-fetches. Restarting the app or Home Assistant does **not** refresh the client — reconnect or refresh the MCP server in your client (e.g. in ChatGPT open **Settings**, then the plugin settings, select the connector and click **Refresh tools**; or close and reopen Claude Desktop). If you skip this, newly enabled tools won't appear in the client at all, and tools the server no longer exposes still show as available but fail when called (`Unknown tool`, or the Read Only Mode error for a write tool that mode hides). ChatGPT sometimes keeps serving the stale list even after the connector is removed and re-added under the same name — if tools are still missing after a refresh, delete the connector and create a new one with a **different name**. The [FAQ](https://homeassistant-ai.github.io/ha-mcp/faq/#chatgpt-stale-tools) has the full steps.
 
 For the HA app, the same option is documented in [`homeassistant-addon/DOCS.md`](homeassistant-addon/DOCS.md#enable_tool_search) along with the in-app settings UI for fine-grained tool enable/disable/pin.
 

{ha_mcp_dev-9.1.1.dev3370 → ha_mcp_dev-9.1.1.dev3378}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "uv_build"
 
 [project]
 name = "ha-mcp-dev"
-version = "9.1.1.dev3370"
+version = "9.1.1.dev3378"
 description = "Home Assistant MCP Server - Complete control of Home Assistant through MCP"
 readme = "README.md"
 requires-python = ">=3.13,<3.15"

{ha_mcp_dev-9.1.1.dev3370 → ha_mcp_dev-9.1.1.dev3378}/pyproject.toml.orig

@@ -4,7 +4,7 @@ build-backend = "uv_build"
 
 [project]
 name = "ha-mcp-dev"
-version = "9.1.1.dev3370"
+version = "9.1.1.dev3378"
 description = "Home Assistant MCP Server - Complete control of Home Assistant through MCP"
 readme = "README.md"
 # Upper bound widened to <3.15 so the package installs on Home Assistant's

{ha_mcp_dev-9.1.1.dev3370 → ha_mcp_dev-9.1.1.dev3378}/src/ha_mcp/__main__.py

@@ -65,8 +65,8 @@ class OAuthProxyClient:
         self._oauth_clients: dict[str, HomeAssistantClient] = {}
         self._lock = threading.Lock()
 
-    def _get_oauth_client(self) -> "HomeAssistantClient":
-        """Get the OAuth client for the current request context."""
+    def current_client(self) -> "HomeAssistantClient":
+        """The requesting user's client; per-user caches key on it, not the proxy."""
         from ha_mcp._vendor.fastmcp.server.dependencies import get_access_token
         from ha_mcp.client.rest_client import (
             HomeAssistantAuthError,
@@ -116,7 +116,7 @@ class OAuthProxyClient:
 
     def __getattr__(self, name: str) -> Any:
         """Forward all attribute access to the OAuth client."""
-        client = self._get_oauth_client()
+        client = self.current_client()
         return getattr(client, name)
 
 

{ha_mcp_dev-9.1.1.dev3370 → ha_mcp_dev-9.1.1.dev3378}/src/ha_mcp/client/websocket_client.py

@@ -14,7 +14,6 @@ import json
 import logging
 import ssl
 import time
-from collections import defaultdict
 from collections.abc import Awaitable, Callable, Coroutine
 from typing import Any
 from urllib.parse import urlparse
@@ -37,6 +36,7 @@ from .rest_client import (
     HomeAssistantConnectionError,
     _is_ssl_error,
 )
+from .websocket_state import WebSocketConnectionState
 
 logger = logging.getLogger(__name__)
 
@@ -136,229 +136,6 @@ def _extract_ws_error(error: Any) -> tuple[str, str | None]:
     return str(error), None
 
 
-class WebSocketConnectionState:
-    """Encapsulates mutable state used by the WebSocket client."""
-
-    def __init__(self) -> None:
-        self.connected = False
-        self.authenticated = False
-        self._message_id = 0
-        self._pending_requests: dict[int, asyncio.Future[dict[str, Any]]] = {}
-        self._auth_messages: dict[str, dict[str, Any]] = {}
-        self._event_responses: dict[int, asyncio.Future[dict[str, Any]]] = {}
-        self._event_handlers: dict[
-            str, set[Callable[[dict[str, Any]], Awaitable[None]]]
-        ] = defaultdict(set)
-        # Continuous-subscription queues keyed by the subscribe command's
-        # message_id. Long-lived subscriptions (e.g. HACS' ``hacs/subscribe``)
-        # deliver many events sharing one id; the one-shot
-        # ``_event_responses`` future can't handle that. When a queue is
-        # registered for a given id, every event with that id is pushed
-        # into it instead of going to ``event_type``-keyed handlers.
-        self._subscription_queues: dict[int, asyncio.Queue[dict[str, Any]]] = {}
-        # Ids of subscribe commands whose caller gave up before Home Assistant
-        # acknowledged them. See ``take_abandoned_subscription``.
-        self._abandoned_subscriptions: set[int] = set()
-
-    def next_message_id(self) -> int:
-        """Reserve the next available WebSocket message identifier."""
-        self._message_id += 1
-        return self._message_id
-
-    def register_pending_request(
-        self, message_id: int
-    ) -> asyncio.Future[dict[str, Any]]:
-        """Create and register a future for a pending command response."""
-        future: asyncio.Future[dict[str, Any]] = (
-            asyncio.get_running_loop().create_future()
-        )
-        self._pending_requests[message_id] = future
-        return future
-
-    def resolve_pending_request(
-        self, message_id: int
-    ) -> asyncio.Future[dict[str, Any]] | None:
-        """Resolve and remove a pending request future."""
-        return self._pending_requests.pop(message_id, None)
-
-    def cancel_pending_request(self, message_id: int) -> None:
-        """Remove a pending request future, cancelling it or draining its exception."""
-        future = self._pending_requests.pop(message_id, None)
-        if future and not future.done():
-            future.cancel()
-        elif future and not future.cancelled():
-            # A done future may hold the exception reset_connection set on it;
-            # retrieve it here so dropping the last reference doesn't log an
-            # ERROR-level "Future exception was never retrieved" at GC.
-            future.exception()
-
-    def register_event_response(
-        self, message_id: int
-    ) -> asyncio.Future[dict[str, Any]]:
-        """Create and register a future for a follow-up event."""
-        future: asyncio.Future[dict[str, Any]] = (
-            asyncio.get_running_loop().create_future()
-        )
-        self._event_responses[message_id] = future
-        return future
-
-    def resolve_event_response(
-        self, message_id: int
-    ) -> asyncio.Future[dict[str, Any]] | None:
-        """Resolve a stored event future."""
-        return self._event_responses.pop(message_id, None)
-
-    def cancel_event_response(self, message_id: int) -> None:
-        """Remove a pending event future, cancelling it or draining its exception."""
-        future = self._event_responses.pop(message_id, None)
-        if future and not future.done():
-            future.cancel()
-        elif future and not future.cancelled():
-            # Same GC guard as cancel_pending_request above.
-            future.exception()
-
-    def mark_abandoned_subscription(self, message_id: int) -> None:
-        """Record a subscribe command whose acknowledgement nobody awaits."""
-        self._abandoned_subscriptions.add(message_id)
-
-    def take_abandoned_subscription(self, message_id: int) -> bool:
-        """Forget an abandoned subscribe command; True if it was one.
-
-        Home Assistant may register a subscription only after the caller has
-        given up and the immediate release answered ``not_found`` — a
-        ``render_template`` setup waits out the template's own timeout first.
-        Its late acknowledgement then arrives with no pending request, and
-        this is how the dispatch recognises that it still has to release it.
-        """
-        if message_id in self._abandoned_subscriptions:
-            self._abandoned_subscriptions.discard(message_id)
-            return True
-        return False
-
-    def store_auth_message(self, message_type: str, data: dict[str, Any]) -> None:
-        """Store an authentication handshake message."""
-        self._auth_messages[message_type] = data
-
-    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, 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():
-                # 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.set_exception(HomeAssistantConnectionError(message))
-        self._event_responses.clear()
-
-        # Drop any subscription queues. ``shutdown`` rather than just clearing
-        # the dict, so blocked ``queue.get()`` awaiters unblock instead of
-        # hanging; not ``immediate``, so a reader still gets the events that
-        # arrived before the close, then ``QueueShutDown`` (3.13).
-        for queue in self._subscription_queues.values():
-            queue.shutdown()
-        self._subscription_queues.clear()
-        # Subscriptions die with the socket, and ids restart on the next one.
-        self._abandoned_subscriptions.clear()
-
-        self._auth_messages.clear()
-
-    def mark_connected(self) -> None:
-        """Mark the socket as connected but not yet authenticated."""
-        self.connected = True
-        self.authenticated = False
-
-    def mark_authenticated(self) -> None:
-        """Mark the socket as authenticated and ready for commands."""
-        self.authenticated = True
-
-    def mark_disconnected(self, close_reason: str | None = None) -> None:
-        """Reset connection state when the socket is closed."""
-        self.reset_connection(close_reason)
-
-    @property
-    def is_ready(self) -> bool:
-        """Whether the connection is active and authenticated."""
-        return self.connected and self.authenticated
-
-    def add_event_handler(
-        self, event_type: str, handler: Callable[[dict[str, Any]], Awaitable[None]]
-    ) -> None:
-        """Register an async handler for a Home Assistant event type."""
-        self._event_handlers[event_type].add(handler)
-
-    def remove_event_handler(
-        self, event_type: str, handler: Callable[[dict[str, Any]], Awaitable[None]]
-    ) -> None:
-        """Remove an event handler and prune empty handler sets."""
-        if event_type in self._event_handlers:
-            self._event_handlers[event_type].discard(handler)
-            if not self._event_handlers[event_type]:
-                self._event_handlers.pop(event_type, None)
-
-    def get_event_handlers(
-        self, event_type: str
-    ) -> tuple[Callable[[dict[str, Any]], Awaitable[None]], ...]:
-        """Return registered handlers for a given event type."""
-        if event_type not in self._event_handlers:
-            return ()
-        return tuple(self._event_handlers[event_type])
-
-    def register_subscription_queue(
-        self, message_id: int
-    ) -> asyncio.Queue[dict[str, Any]]:
-        """Register a queue for continuous-subscription event delivery."""
-        queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
-        self._subscription_queues[message_id] = queue
-        return queue
-
-    def get_subscription_queue(
-        self, message_id: int
-    ) -> asyncio.Queue[dict[str, Any]] | None:
-        """Return the queue for a subscription id, or None."""
-        return self._subscription_queues.get(message_id)
-
-    def unregister_subscription_queue(self, message_id: int) -> None:
-        """Drop a subscription queue and wake any blocked readers."""
-        queue = self._subscription_queues.pop(message_id, None)
-        if queue is not None:
-            # ``shutdown(immediate=True)`` raises ``QueueShutDown`` in any
-            # pending ``get()`` so a waiter doesn't deadlock when the
-            # caller decides to stop listening.
-            queue.shutdown(immediate=True)
-
-
 class HomeAssistantWebSocketClient:
     """WebSocket client for Home Assistant real-time communication."""
 
@@ -1541,6 +1318,40 @@ class WebSocketManager:
         """Create a cache key from credentials."""
         return hashlib.sha256(f"{url.rstrip('/')}:{token}".encode()).hexdigest()
 
+    def _pool_key(
+        self, url: str | None, token: str | None, verify_ssl: bool | None
+    ) -> tuple[str, str, str]:
+        """The credentials a caller connects with, and their pool key."""
+        if url and token:
+            ws_url, ws_token = url, token
+        else:
+            settings = get_global_settings()
+            ws_url = settings.homeassistant_url
+            ws_token = settings.homeassistant_token
+        # Key on the EFFECTIVE verification mode: a caller passing the
+        # resolved settings default (send_websocket_message) must share
+        # the pooled connection with callers that omit the argument
+        # (listener, HACS, installer) — only a genuine override such as
+        # verify_ssl=False gets its own isolated connection.
+        key = (
+            f"{self._client_key(ws_url, ws_token)}"
+            f"|verify_ssl={self._effective_verify_ssl(verify_ssl)}"
+        )
+        return ws_url, ws_token, key
+
+    def pooled_client(
+        self,
+        url: str | None = None,
+        token: str | None = None,
+        verify_ssl: bool | None = None,
+    ) -> HomeAssistantWebSocketClient | None:
+        """The pool's client for these credentials, without connecting.
+
+        It may be disconnected: ``get_client`` replaces such a client only on
+        its next call. None when the pool holds none for them.
+        """
+        return self._clients.get(self._pool_key(url, token, verify_ssl)[2])
+
     @staticmethod
     def _effective_verify_ssl(verify_ssl: bool | None) -> bool:
         """Resolve the effective TLS-verification mode for the pool key."""
@@ -1653,25 +1464,7 @@ class WebSocketManager:
             self._current_loop = current_loop
             self._release_stale_clients(stale_clients, previous_loop)
 
-            # Determine credentials to use
-            if url and token:
-                ws_url = url
-                ws_token = token
-            else:
-                settings = get_global_settings()
-                ws_url = settings.homeassistant_url
-                ws_token = settings.homeassistant_token
-
-            # Key on the EFFECTIVE verification mode: a caller passing the
-            # resolved settings default (send_websocket_message) must share
-            # the pooled connection with callers that omit the argument
-            # (listener, HACS, installer) — only a genuine override such as
-            # verify_ssl=False gets its own isolated connection.
-            effective_verify_ssl = self._effective_verify_ssl(verify_ssl)
-            key = (
-                f"{self._client_key(ws_url, ws_token)}"
-                f"|verify_ssl={effective_verify_ssl}"
-            )
+            ws_url, ws_token, key = self._pool_key(url, token, verify_ssl)
 
             # Return existing connected client for these credentials
             existing = self._clients.get(key)

ha_mcp_dev-9.1.1.dev3378/src/ha_mcp/client/websocket_state.py

@@ -0,0 +1,231 @@
+"""Request, event and subscription bookkeeping for one WebSocket connection."""
+
+import asyncio
+from collections import defaultdict
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+from .rest_client import HomeAssistantConnectionError
+
+
+class WebSocketConnectionState:
+    """Encapsulates mutable state used by the WebSocket client."""
+
+    def __init__(self) -> None:
+        self.connected = False
+        self.authenticated = False
+        self._message_id = 0
+        self._pending_requests: dict[int, asyncio.Future[dict[str, Any]]] = {}
+        self._auth_messages: dict[str, dict[str, Any]] = {}
+        self._event_responses: dict[int, asyncio.Future[dict[str, Any]]] = {}
+        self._event_handlers: dict[
+            str, set[Callable[[dict[str, Any]], Awaitable[None]]]
+        ] = defaultdict(set)
+        # Continuous-subscription queues keyed by the subscribe command's
+        # message_id. Long-lived subscriptions (e.g. HACS' ``hacs/subscribe``)
+        # deliver many events sharing one id; the one-shot
+        # ``_event_responses`` future can't handle that. When a queue is
+        # registered for a given id, every event with that id is pushed
+        # into it instead of going to ``event_type``-keyed handlers.
+        self._subscription_queues: dict[int, asyncio.Queue[dict[str, Any]]] = {}
+        # Ids of subscribe commands whose caller gave up before Home Assistant
+        # acknowledged them. See ``take_abandoned_subscription``.
+        self._abandoned_subscriptions: set[int] = set()
+
+    def next_message_id(self) -> int:
+        """Reserve the next available WebSocket message identifier."""
+        self._message_id += 1
+        return self._message_id
+
+    def register_pending_request(
+        self, message_id: int
+    ) -> asyncio.Future[dict[str, Any]]:
+        """Create and register a future for a pending command response."""
+        future: asyncio.Future[dict[str, Any]] = (
+            asyncio.get_running_loop().create_future()
+        )
+        self._pending_requests[message_id] = future
+        return future
+
+    def resolve_pending_request(
+        self, message_id: int
+    ) -> asyncio.Future[dict[str, Any]] | None:
+        """Resolve and remove a pending request future."""
+        return self._pending_requests.pop(message_id, None)
+
+    def cancel_pending_request(self, message_id: int) -> None:
+        """Remove a pending request future, cancelling it or draining its exception."""
+        future = self._pending_requests.pop(message_id, None)
+        if future and not future.done():
+            future.cancel()
+        elif future and not future.cancelled():
+            # A done future may hold the exception reset_connection set on it;
+            # retrieve it here so dropping the last reference doesn't log an
+            # ERROR-level "Future exception was never retrieved" at GC.
+            future.exception()
+
+    def register_event_response(
+        self, message_id: int
+    ) -> asyncio.Future[dict[str, Any]]:
+        """Create and register a future for a follow-up event."""
+        future: asyncio.Future[dict[str, Any]] = (
+            asyncio.get_running_loop().create_future()
+        )
+        self._event_responses[message_id] = future
+        return future
+
+    def resolve_event_response(
+        self, message_id: int
+    ) -> asyncio.Future[dict[str, Any]] | None:
+        """Resolve a stored event future."""
+        return self._event_responses.pop(message_id, None)
+
+    def cancel_event_response(self, message_id: int) -> None:
+        """Remove a pending event future, cancelling it or draining its exception."""
+        future = self._event_responses.pop(message_id, None)
+        if future and not future.done():
+            future.cancel()
+        elif future and not future.cancelled():
+            # Same GC guard as cancel_pending_request above.
+            future.exception()
+
+    def mark_abandoned_subscription(self, message_id: int) -> None:
+        """Record a subscribe command whose acknowledgement nobody awaits."""
+        self._abandoned_subscriptions.add(message_id)
+
+    def take_abandoned_subscription(self, message_id: int) -> bool:
+        """Forget an abandoned subscribe command; True if it was one.
+
+        Home Assistant may register a subscription only after the caller has
+        given up and the immediate release answered ``not_found`` — a
+        ``render_template`` setup waits out the template's own timeout first.
+        Its late acknowledgement then arrives with no pending request, and
+        this is how the dispatch recognises that it still has to release it.
+        """
+        if message_id in self._abandoned_subscriptions:
+            self._abandoned_subscriptions.discard(message_id)
+            return True
+        return False
+
+    def store_auth_message(self, message_type: str, data: dict[str, Any]) -> None:
+        """Store an authentication handshake message."""
+        self._auth_messages[message_type] = data
+
+    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, 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():
+                # 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.set_exception(HomeAssistantConnectionError(message))
+        self._event_responses.clear()
+
+        # Drop any subscription queues. ``shutdown`` rather than just clearing
+        # the dict, so blocked ``queue.get()`` awaiters unblock instead of
+        # hanging; not ``immediate``, so a reader still gets the events that
+        # arrived before the close, then ``QueueShutDown`` (3.13).
+        for queue in self._subscription_queues.values():
+            queue.shutdown()
+        self._subscription_queues.clear()
+        # Subscriptions die with the socket, and ids restart on the next one.
+        self._abandoned_subscriptions.clear()
+
+        self._auth_messages.clear()
+
+    def mark_connected(self) -> None:
+        """Mark the socket as connected but not yet authenticated."""
+        self.connected = True
+        self.authenticated = False
+
+    def mark_authenticated(self) -> None:
+        """Mark the socket as authenticated and ready for commands."""
+        self.authenticated = True
+
+    def mark_disconnected(self, close_reason: str | None = None) -> None:
+        """Reset connection state when the socket is closed."""
+        self.reset_connection(close_reason)
+
+    @property
+    def is_ready(self) -> bool:
+        """Whether the connection is active and authenticated."""
+        return self.connected and self.authenticated
+
+    def add_event_handler(
+        self, event_type: str, handler: Callable[[dict[str, Any]], Awaitable[None]]
+    ) -> None:
+        """Register an async handler for a Home Assistant event type."""
+        self._event_handlers[event_type].add(handler)
+
+    def remove_event_handler(
+        self, event_type: str, handler: Callable[[dict[str, Any]], Awaitable[None]]
+    ) -> None:
+        """Remove an event handler and prune empty handler sets."""
+        if event_type in self._event_handlers:
+            self._event_handlers[event_type].discard(handler)
+            if not self._event_handlers[event_type]:
+                self._event_handlers.pop(event_type, None)
+
+    def get_event_handlers(
+        self, event_type: str
+    ) -> tuple[Callable[[dict[str, Any]], Awaitable[None]], ...]:
+        """Return registered handlers for a given event type."""
+        if event_type not in self._event_handlers:
+            return ()
+        return tuple(self._event_handlers[event_type])
+
+    def register_subscription_queue(
+        self, message_id: int
+    ) -> asyncio.Queue[dict[str, Any]]:
+        """Register a queue for continuous-subscription event delivery."""
+        queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
+        self._subscription_queues[message_id] = queue
+        return queue
+
+    def get_subscription_queue(
+        self, message_id: int
+    ) -> asyncio.Queue[dict[str, Any]] | None:
+        """Return the queue for a subscription id, or None."""
+        return self._subscription_queues.get(message_id)
+
+    def unregister_subscription_queue(self, message_id: int) -> None:
+        """Drop a subscription queue and wake any blocked readers."""
+        queue = self._subscription_queues.pop(message_id, None)
+        if queue is not None:
+            # ``shutdown(immediate=True)`` raises ``QueueShutDown`` in any
+            # pending ``get()`` so a waiter doesn't deadlock when the
+            # caller decides to stop listening.
+            queue.shutdown(immediate=True)