Allow extracting non-primary entities in websocket command

homeassistant/components/websocket_api/commands.py

@@ -865,6 +865,7 @@ def handle_entity_source(
         vol.Required("type"): "extract_from_target",
         vol.Required("target"): cv.TARGET_FIELDS,
         vol.Optional("expand_group", default=False): bool,
+        vol.Optional("primary_entities_only", default=True): bool,
     }
 )
 def handle_extract_from_target(
@@ -874,7 +875,10 @@ def handle_extract_from_target(
 
     target_selection = target_helpers.TargetSelection(msg["target"])
     extracted = target_helpers.async_extract_referenced_entity_ids(
-        hass, target_selection, expand_group=msg["expand_group"]
+        hass,
+        target_selection,
+        expand_group=msg["expand_group"],
+        primary_entities_only=msg["primary_entities_only"],
     )
 
     extracted_dict = {

homeassistant/helpers/target.py

@@ -147,9 +147,22 @@ def log_missing(self, missing_entities: set[str], logger: Logger) -> None:
 
 
 def async_extract_referenced_entity_ids(
-    hass: HomeAssistant, target_selection: TargetSelection, expand_group: bool = True
+    hass: HomeAssistant,
+    target_selection: TargetSelection,
+    expand_group: bool = True,
+    *,
+    primary_entities_only: bool = True,
 ) -> SelectedEntities:
-    """Extract referenced entity IDs from a target selection."""
+    """Extract referenced entity IDs from a target selection.
+
+    When `primary_entities_only` is True (the default), entities with an
+    `entity_category` (i.e. config or diagnostic entities) are excluded from
+    indirect expansion via device, area, and floor. When False, those entities
+    are included. Direct label-to-entity expansion is unaffected by this flag.
+    Label targeting via labeled devices or areas follows the same filtering
+    rules as other indirect device/area expansion paths: filtered when
+    `primary_entities_only` is True, and included when it is False.
+    """
     selected = SelectedEntities()
 
     if not target_selection.has_any_target:
@@ -217,14 +230,18 @@ def async_extract_referenced_entity_ids(
     if not selected.referenced_areas and not selected.referenced_devices:
         return selected
 
+    def _include_entry(entry: er.RegistryEntry) -> bool:
+        """Return True if the entry should be included in indirect expansion."""
+        if entry.hidden_by is not None:
+            return False
+        return not primary_entities_only or entry.entity_category is None
+
     # Add indirectly referenced by device
     selected.indirectly_referenced.update(
         entry.entity_id
         for device_id in selected.referenced_devices
         for entry in entities.get_entries_for_device_id(device_id)
-        # Do not add entities which are hidden or which are config
-        # or diagnostic entities.
-        if (entry.entity_category is None and entry.hidden_by is None)
+        if _include_entry(entry)
     )
 
     # Find devices for targeted areas
@@ -243,27 +260,16 @@ def async_extract_referenced_entity_ids(
         for area_id in selected.referenced_areas
         # The entity's area matches a targeted area
         for entry in entities.get_entries_for_area_id(area_id)
-        # Do not add entities which are hidden or which are config
-        # or diagnostic entities.
-        if entry.entity_category is None and entry.hidden_by is None
+        if _include_entry(entry)
     )
     # Add indirectly referenced by area through device
     selected.indirectly_referenced.update(
         entry.entity_id
         for device_id in referenced_devices_by_area
         for entry in entities.get_entries_for_device_id(device_id)
-        # Do not add entities which are hidden or which are config
-        # or diagnostic entities.
-        if (
-            entry.entity_category is None
-            and entry.hidden_by is None
-            and (
-                # The entity's device matches a device referenced
-                # by an area and the entity
-                # has no explicitly set area
-                not entry.area_id
-            )
-        )
+        # The entity's device matches a device referenced by an area and the
+        # entity has no explicitly set area.
+        if _include_entry(entry) and not entry.area_id
     )
 
     return selected
@@ -277,11 +283,14 @@ def __init__(
         hass: HomeAssistant,
         target_selection: TargetSelection,
         entity_filter: Callable[[set[str]], set[str]],
+        *,
+        primary_entities_only: bool = True,
     ) -> None:
         """Initialize the state change tracker."""
         self._hass = hass
         self._target_selection = target_selection
         self._entity_filter = entity_filter
+        self._primary_entities_only = primary_entities_only
 
         self._registry_unsubs: list[CALLBACK_TYPE] = []
 
@@ -300,7 +309,10 @@ def _handle_entities_update(self, tracked_entities: set[str]) -> None:
     def _handle_target_update(self, event: Event[Any] | None = None) -> None:
         """Handle updates in the tracked targets."""
         selected = async_extract_referenced_entity_ids(
-            self._hass, self._target_selection, expand_group=False
+            self._hass,
+            self._target_selection,
+            expand_group=False,
+            primary_entities_only=self._primary_entities_only,
         )
         filtered_entities = self._entity_filter(
             selected.referenced | selected.indirectly_referenced
@@ -345,9 +357,16 @@ def __init__(
         target_selection: TargetSelection,
         action: Callable[[TargetStateChangedData], Any],
         entity_filter: Callable[[set[str]], set[str]],
+        *,
+        primary_entities_only: bool = True,
     ) -> None:
         """Initialize the state change tracker."""
-        super().__init__(hass, target_selection, entity_filter)
+        super().__init__(
+            hass,
+            target_selection,
+            entity_filter,
+            primary_entities_only=primary_entities_only,
+        )
         self._action = action
         self._state_change_unsub: CALLBACK_TYPE | None = None
 
@@ -380,12 +399,24 @@ def async_track_target_selector_state_change_event(
     target_selector_config: ConfigType,
     action: Callable[[TargetStateChangedData], Any],
     entity_filter: Callable[[set[str]], set[str]] = lambda x: x,
+    *,
+    primary_entities_only: bool = True,
 ) -> CALLBACK_TYPE:
-    """Track state changes for entities referenced directly or indirectly in a target selector."""
+    """Track state changes for entities referenced directly or indirectly in a target selector.
+
+    When `primary_entities_only` is True, indirect target expansion (via device, area,
+    and floor) skips entities with an `entity_category` (i.e. config or diagnostic entities).
+    """
     target_selection = TargetSelection(target_selector_config)
     if not target_selection.has_any_target:
         raise HomeAssistantError(
             f"Target selector {target_selector_config} does not have any selectors defined"
         )
-    tracker = TargetStateChangeTracker(hass, target_selection, action, entity_filter)
+    tracker = TargetStateChangeTracker(
+        hass,
+        target_selection,
+        action,
+        entity_filter,
+        primary_entities_only=primary_entities_only,
+    )
     return tracker.async_setup()

tests/components/websocket_api/test_commands.py

@@ -34,7 +34,11 @@
 )
 from homeassistant.components.websocket_api.const import FEATURE_COALESCE_MESSAGES, URL
 from homeassistant.config_entries import ConfigEntryState
-from homeassistant.const import CONF_EXTERNAL_URL, SIGNAL_BOOTSTRAP_INTEGRATIONS
+from homeassistant.const import (
+    CONF_EXTERNAL_URL,
+    SIGNAL_BOOTSTRAP_INTEGRATIONS,
+    EntityCategory,
+)
 from homeassistant.core import Context, HomeAssistant, State, SupportsResponse, callback
 from homeassistant.exceptions import HomeAssistantError, ServiceValidationError
 from homeassistant.helpers import (
@@ -3604,6 +3608,88 @@ async def test_extract_from_target_expand_group(
     )
 
 
+async def test_extract_from_target_primary_entities_only(
+    hass: HomeAssistant,
+    websocket_client: MockHAClientWebSocket,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+) -> None:
+    """Test extract_from_target command with primary_entities_only parameter."""
+    config_entry = MockConfigEntry(domain="test")
+    config_entry.add_to_hass(hass)
+
+    device = device_registry.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        identifiers={("test", "device1")},
+    )
+
+    primary_entity = entity_registry.async_get_or_create(
+        "light", "test", "unique1", device_id=device.id
+    )
+    diagnostic_entity = entity_registry.async_get_or_create(
+        "sensor",
+        "test",
+        "unique2",
+        device_id=device.id,
+        entity_category=EntityCategory.DIAGNOSTIC,
+    )
+    config_entity = entity_registry.async_get_or_create(
+        "switch",
+        "test",
+        "unique3",
+        device_id=device.id,
+        entity_category=EntityCategory.CONFIG,
+    )
+
+    # Default (primary_entities_only=True): config/diagnostic entities excluded
+    await websocket_client.send_json_auto_id(
+        {
+            "type": "extract_from_target",
+            "target": {"device_id": [device.id]},
+        }
+    )
+    msg = await websocket_client.receive_json()
+    _assert_extract_from_target_command_result(
+        msg,
+        entities={primary_entity.entity_id},
+        devices={device.id},
+    )
+
+    # Explicit primary_entities_only=True
+    await websocket_client.send_json_auto_id(
+        {
+            "type": "extract_from_target",
+            "target": {"device_id": [device.id]},
+            "primary_entities_only": True,
+        }
+    )
+    msg = await websocket_client.receive_json()
+    _assert_extract_from_target_command_result(
+        msg,
+        entities={primary_entity.entity_id},
+        devices={device.id},
+    )
+
+    # primary_entities_only=False: config/diagnostic entities included
+    await websocket_client.send_json_auto_id(
+        {
+            "type": "extract_from_target",
+            "target": {"device_id": [device.id]},
+            "primary_entities_only": False,
+        }
+    )
+    msg = await websocket_client.receive_json()
+    _assert_extract_from_target_command_result(
+        msg,
+        entities={
+            primary_entity.entity_id,
+            diagnostic_entity.entity_id,
+            config_entity.entity_id,
+        },
+        devices={device.id},
+    )
+
+
 async def test_extract_from_target_missing_entities(
     hass: HomeAssistant, websocket_client: MockHAClientWebSocket
 ) -> None:

tests/helpers/test_target.py

@@ -499,6 +499,37 @@ async def test_extract_referenced_entity_ids(
     )
 
 
+@pytest.mark.parametrize(
+    ("selector_config", "non_primary_entities"),
+    [
+        ({ATTR_AREA_ID: "own-area"}, {"light.config_in_own_area"}),
+        ({ATTR_DEVICE_ID: "device-no-area-id"}, {"light.config_no_area"}),
+        ({ATTR_AREA_ID: "test-area"}, {"light.config_in_area"}),
+    ],
+)
+@pytest.mark.usefixtures("registries_mock")
+async def test_extract_referenced_entity_ids_primary_entities_only(
+    hass: HomeAssistant,
+    selector_config: ConfigType,
+    non_primary_entities: set[str],
+) -> None:
+    """Test that primary_entities_only controls inclusion of config/diagnostic entities."""
+    target_selection = target.TargetSelection(selector_config)
+
+    selected_primary = target.async_extract_referenced_entity_ids(
+        hass, target_selection, expand_group=False, primary_entities_only=True
+    )
+    selected_all = target.async_extract_referenced_entity_ids(
+        hass, target_selection, expand_group=False, primary_entities_only=False
+    )
+
+    assert (
+        selected_all.indirectly_referenced
+        == selected_primary.indirectly_referenced | non_primary_entities
+    )
+    assert non_primary_entities.isdisjoint(selected_primary.indirectly_referenced)
+
+
 async def test_async_track_target_selector_state_change_event_empty_selector(
     hass: HomeAssistant, caplog: pytest.LogCaptureFixture
 ) -> None: