fix: resolve entity areas through device registry in get_system_overview (#729)

* fix: resolve entity areas through device registry in get_system_overview

* fix: remove unused area_name_map and add regression test for device area resolution

Author: knowald

src/ha_mcp/tools/smart_search.py

@@ -374,12 +374,16 @@ async def get_system_overview(
             entity_registry_task = self.client.send_websocket_message(
                 {"type": "config/entity_registry/list"}
             )
+            device_registry_task = self.client.send_websocket_message(
+                {"type": "config/device_registry/list"}
+            )
 
             results = await asyncio.gather(
                 entities_task,
                 services_task,
                 area_registry_task,
                 entity_registry_task,
+                device_registry_task,
                 return_exceptions=True,
             )
 
@@ -401,11 +405,26 @@ async def get_system_overview(
             elif isinstance(results[3], dict) and results[3].get("success"):
                 entity_registry = results[3].get("result", [])
 
-            # Build entity_id -> area_id mapping from entity registry
+            # Handle device registry result
+            device_area_map: dict[str, str | None] = {}
+            if isinstance(results[4], Exception):
+                logger.debug(f"Could not fetch device registry: {results[4]}")
+            elif isinstance(results[4], dict) and results[4].get("success"):
+                for device in results[4].get("result", []):
+                    device_id = device.get("id", "")
+                    if device_id:
+                        device_area_map[device_id] = device.get("area_id")
+
+            # Build entity_id -> area_id mapping from entity + device registries
+            # Priority: entity direct area_id > device area_id
             entity_area_map: dict[str, str | None] = {}
             for entry in entity_registry:
                 entity_id = entry.get("entity_id")
                 area_id = entry.get("area_id")
+                if not area_id:
+                    device_id = entry.get("device_id")
+                    if device_id:
+                        area_id = device_area_map.get(device_id)
                 if entity_id:
                     entity_area_map[entity_id] = area_id
 
@@ -417,9 +436,19 @@ async def get_system_overview(
             if include_entity_id is None:
                 include_entity_id = detail_level == "full"
 
+            # Pre-populate area_stats to include empty areas
+            area_stats: dict[str, dict[str, Any]] = {}
+            for area in area_registry:
+                area_id = area.get("area_id", "")
+                if area_id:
+                    area_stats[area_id] = {
+                        "name": area.get("name", area_id),
+                        "count": 0,
+                        "domains": {},
+                    }
+
             # Analyze entities by domain
             domain_stats: dict[str, dict[str, Any]] = {}
-            area_stats: dict[str, dict[str, Any]] = {}
             device_types: dict[str, int] = {}
 
             for entity in entities:
@@ -454,11 +483,9 @@ async def get_system_overview(
 
                 domain_stats[domain]["all_entities"].append(entity_data)
 
-                # Area analysis - use entity registry mapping, not state attributes
+                # Area analysis - use entity + device registry mapping
                 area_id = entity_area_map.get(entity_id)
-                if area_id:
-                    if area_id not in area_stats:
-                        area_stats[area_id] = {"count": 0, "domains": {}}
+                if area_id and area_id in area_stats:
                     area_stats[area_id]["count"] += 1
                     if domain not in area_stats[area_id]["domains"]:
                         area_stats[area_id]["domains"][domain] = 0

tests/src/unit/test_performance_parallelization.py

@@ -235,7 +235,7 @@ async def test_parallel_calls_faster_than_sequential(
             elapsed = time.time() - start
 
         assert result["success"] is True
-        # 4 data sources at 0.05s each: sequential = 0.2s, parallel ≈ 0.05s
+        # 5 data sources at 0.05s each: sequential = 0.25s, parallel ≈ 0.05s
         assert elapsed < 0.15, f"Expected parallel speedup, took {elapsed:.2f}s"
 
     @pytest.mark.asyncio
@@ -258,6 +258,57 @@ async def failing_websocket(message: dict) -> dict:
         assert result["system_summary"]["total_entities"] == 2
         assert result["system_summary"]["total_areas"] == 0
 
+    @pytest.mark.asyncio
+    async def test_resolves_area_through_device_registry(self, sample_services):
+        """Entities with no direct area_id inherit area from their parent device."""
+        entities = [
+            {
+                "entity_id": "light.kitchen",
+                "attributes": {"friendly_name": "Kitchen Light"},
+                "state": "on",
+            },
+        ]
+        client = MockClient(entities=entities, services=sample_services)
+
+        # Override websocket to return entity without area_id but with device_id,
+        # and a device registry that maps that device to an area.
+        original_ws = client.send_websocket_message
+
+        async def ws_with_device_registry(message: dict) -> dict:
+            msg_type = message.get("type", "")
+            if msg_type == "config/entity_registry/list":
+                return {
+                    "success": True,
+                    "result": [
+                        {
+                            "entity_id": "light.kitchen",
+                            "area_id": None,
+                            "device_id": "device_1",
+                        },
+                    ],
+                }
+            if msg_type == "config/device_registry/list":
+                return {
+                    "success": True,
+                    "result": [
+                        {"id": "device_1", "area_id": "living_room"},
+                    ],
+                }
+            return await original_ws(message)
+
+        client.send_websocket_message = ws_with_device_registry
+
+        with patch("ha_mcp.tools.smart_search.get_global_settings") as mock_settings:
+            mock_settings.return_value.fuzzy_threshold = 60
+            tools = SmartSearchTools(client=client)
+            result = await tools.get_system_overview(detail_level="full")
+
+        assert result["success"] is True
+        assert result["system_summary"]["total_areas"] == 2
+        area = result["area_analysis"]["living_room"]
+        assert area["count"] == 1
+        assert area["domains"]["light"] == 1
+
 
 # ---------------------------------------------------------------------------
 # deep_search – outcome-based tests