M8-R15: fix HA discovery identities and thermal totals

app/integrations/homeassistant.py

@@ -2,10 +2,14 @@ from __future__ import annotations
 
 import json
 import logging
+from time import monotonic
 from dataclasses import dataclass, field
 from typing import Any
 from urllib import error, parse, request
 
+from websockets.exceptions import WebSocketException
+from websockets.sync.client import connect
+
 from app.config import Settings
 
 logger = logging.getLogger(__name__)
@@ -57,6 +61,81 @@ class HomeAssistantClient:
 
         self._post_json(f"/api/webhook/{webhook_id}", body, operation="trigger_webhook")
 
+    def discovery_registry_bindings(self, unique_ids: set[str]) -> dict[str, set[str]]:
+        """Return HA device identifiers currently bound to MQTT unique IDs.
+
+        This is deliberately a read-only WebSocket query.  MQTT only confirms
+        broker receipt; the entity/device registries are the authoritative HA
+        observation that a discovery unload/re-add was actually processed.
+        """
+        self._require_config()
+        if not unique_ids:
+            return {}
+        try:
+            deadline = monotonic() + self.timeout_seconds
+            with connect(self._websocket_url(), open_timeout=self.timeout_seconds,
+                         close_timeout=self.timeout_seconds) as websocket:
+                greeting = json.loads(self._websocket_recv(websocket, deadline))
+                if greeting.get("type") != "auth_required":
+                    raise HomeAssistantRequestError("Unexpected Home Assistant WebSocket greeting")
+                websocket.send(json.dumps({"type": "auth", "access_token": self.settings.home_assistant_auth_token}))
+                auth = json.loads(self._websocket_recv(websocket, deadline))
+                if auth.get("type") != "auth_ok":
+                    raise HomeAssistantRequestError("Home Assistant WebSocket authentication failed")
+                entities = self._websocket_command(websocket, 1, "config/entity_registry/list", deadline)
+                devices = self._websocket_command(websocket, 2, "config/device_registry/list", deadline)
+        except (OSError, WebSocketException, TimeoutError, ValueError, KeyError, TypeError) as exc:
+            raise HomeAssistantRequestError("Home Assistant registry query failed") from exc
+
+        devices_by_id = {
+            device["id"]: {
+                identifier[1]
+                for identifier in device.get("identifiers", [])
+                if (
+                    isinstance(identifier, (list, tuple))
+                    and len(identifier) == 2
+                    and identifier[0] == "mqtt"
+                    and isinstance(identifier[1], str)
+                    and identifier[1]
+                )
+            }
+            for device in devices
+            if isinstance(device, dict) and isinstance(device.get("id"), str)
+        }
+        return {
+            entity["unique_id"]: devices_by_id.get(entity.get("device_id"), set())
+            for entity in entities
+            if entity.get("platform") == "mqtt" and entity.get("unique_id") in unique_ids
+        }
+
+    @staticmethod
+    def _websocket_recv(websocket: Any, deadline: float) -> str:
+        remaining = deadline - monotonic()
+        if remaining <= 0:
+            raise TimeoutError("Home Assistant WebSocket registry query timed out")
+        return websocket.recv(timeout=remaining)
+
+    @classmethod
+    def _websocket_command(
+        cls, websocket: Any, message_id: int, command: str, deadline: float
+    ) -> list[dict[str, Any]]:
+        websocket.send(json.dumps({"id": message_id, "type": command}))
+        while True:
+            response = json.loads(cls._websocket_recv(websocket, deadline))
+            if response.get("id") != message_id:
+                continue
+            if not response.get("success"):
+                raise HomeAssistantRequestError(f"Home Assistant WebSocket {command} failed")
+            return response.get("result", [])
+
+    def _websocket_url(self) -> str:
+        parsed = parse.urlsplit(self.settings.home_assistant_base_url)
+        if parsed.scheme not in {"http", "https"} or not parsed.netloc:
+            raise HomeAssistantConfigError("HOME_ASSISTANT_BASE_URL must be an HTTP(S) URL")
+        scheme = "wss" if parsed.scheme == "https" else "ws"
+        path = f"{parsed.path.rstrip('/')}/api/websocket"
+        return parse.urlunsplit((scheme, parsed.netloc, path, "", ""))
+
     def _require_config(self) -> None:
         if self.is_configured():
             return