Make wake word interception a subscription

homeassistant/components/assist_satellite/entity.py

@@ -25,13 +25,13 @@
 from homeassistant.components.tts.media_source import (
     generate_media_source_id as tts_generate_media_source_id,
 )
-from homeassistant.core import Context, callback
+from homeassistant.core import CALLBACK_TYPE, Context, callback
 from homeassistant.helpers import entity
 from homeassistant.helpers.entity import EntityDescription
 from homeassistant.util import ulid
 
 from .const import AssistSatelliteEntityFeature
-from .errors import AssistSatelliteError, SatelliteBusyError
+from .errors import SatelliteBusyError
 
 _CONVERSATION_TIMEOUT_SEC: Final = 5 * 60  # 5 minutes
 
@@ -72,7 +72,7 @@ class AssistSatelliteEntity(entity.Entity):
 
     _run_has_tts: bool = False
     _is_announcing = False
-    _wake_word_intercept_future: asyncio.Future[str | None] | None = None
+    _wake_word_listener: Callable[[str | None, str], None] | None = None
     _attr_tts_options: dict[str, Any] | None = None
     _pipeline_task: asyncio.Task | None = None
 
@@ -99,33 +99,25 @@ def tts_options(self) -> dict[str, Any] | None:
         """Options passed for text-to-speech."""
         return self._attr_tts_options
 
-    async def async_intercept_wake_word(self) -> str | None:
-        """Intercept the next wake word from the satellite.
+    def async_intercept_wake_word(
+        self, wake_word_listener: Callable[[str | None, str], None]
+    ) -> CALLBACK_TYPE:
+        """Set a listener to intercept the next wake word from the satellite.
 
-        Returns the detected wake word phrase or None.
+        Returns a callback to remove the listener.
         """
-        if self._wake_word_intercept_future is not None:
+        if self._wake_word_listener is not None:
             # Cancel existing interception
-            self._wake_word_intercept_future.set_result(None)
-            self._wake_word_intercept_future = None
+            self._wake_word_listener(None, "Only one interception is allowed")
 
-        # Will cause next wake word to be intercepted in
-        # async_accept_pipeline_from_satellite
-        self._wake_word_intercept_future = asyncio.Future()
+        self._wake_word_listener = wake_word_listener
 
-        _LOGGER.debug("Next wake word will be intercepted: %s", self.entity_id)
+        @callback
+        def unsubscribe() -> None:
+            if self._wake_word_listener == wake_word_listener:
+                self._wake_word_listener = None
 
-        try:
-            return await self._wake_word_intercept_future
-        finally:
-            self._wake_word_intercept_future = None
-
-    def async_stop_intercept_wake_word(self) -> None:
-        """Stop intercepting wake words."""
-        if self._wake_word_intercept_future is not None:
-            # Cancel existing interception
-            self._wake_word_intercept_future.set_result(None)
-            self._wake_word_intercept_future = None
+        return unsubscribe
 
     async def async_internal_announce(
         self,
@@ -213,11 +205,10 @@ async def async_accept_pipeline_from_satellite(
             PipelineStage.STT,
         ):
             if start_stage == PipelineStage.WAKE_WORD:
-                self._wake_word_intercept_future.set_exception(
-                    AssistSatelliteError(
-                        "Only on-device wake words currently supported"
-                    )
+                self._wake_word_listener(
+                    None, "Only on-device wake words currently supported"
                 )
+                self._wake_word_listener = None
                 return
 
             # Intercepting wake word and immediately end pipeline
@@ -227,12 +218,14 @@ async def async_accept_pipeline_from_satellite(
                 self.entity_id,
             )
 
-            if wake_word_phrase is None:
-                self._wake_word_intercept_future.set_exception(
-                    AssistSatelliteError("No wake word phrase provided")
-                )
-            else:
-                self._wake_word_intercept_future.set_result(wake_word_phrase)
+            try:
+                if wake_word_phrase is None:
+                    self._wake_word_listener(None, "No wake word phrase provided")
+                else:
+                    self._wake_word_listener(wake_word_phrase, "")
+            finally:
+                self._wake_word_listener = None
+
             self._internal_on_pipeline_event(PipelineEvent(PipelineEventType.RUN_END))
             return
 

homeassistant/components/assist_satellite/websocket_api.py

@@ -16,20 +16,19 @@
 @callback
 def async_register_websocket_api(hass: HomeAssistant) -> None:
     """Register the websocket API."""
-    websocket_api.async_register_command(hass, websocket_intercept_wake_word_start)
-    websocket_api.async_register_command(hass, websocket_intercept_wake_word_stop)
+    websocket_api.async_register_command(hass, websocket_intercept_wake_word)
 
 
 @callback
 @websocket_api.websocket_command(
     {
-        vol.Required("type"): "assist_satellite/intercept_wake_word/start",
+        vol.Required("type"): "assist_satellite/intercept_wake_word",
         vol.Required("entity_id"): cv.entity_domain(DOMAIN),
     }
 )
 @websocket_api.require_admin
 @websocket_api.async_response
-async def websocket_intercept_wake_word_start(
+async def websocket_intercept_wake_word(
     hass: HomeAssistant,
     connection: websocket_api.connection.ActiveConnection,
     msg: dict[str, Any],
@@ -43,31 +42,20 @@ async def websocket_intercept_wake_word_start(
         )
         return
 
-    wake_word_phrase = await satellite.async_intercept_wake_word()
-    connection.send_result(msg["id"], {"wake_word_phrase": wake_word_phrase})
-
-
-@callback
-@websocket_api.websocket_command(
-    {
-        vol.Required("type"): "assist_satellite/intercept_wake_word/stop",
-        vol.Required("entity_id"): cv.entity_domain(DOMAIN),
-    }
-)
-@websocket_api.require_admin
-@websocket_api.async_response
-async def websocket_intercept_wake_word_stop(
-    hass: HomeAssistant,
-    connection: websocket_api.connection.ActiveConnection,
-    msg: dict[str, Any],
-) -> None:
-    """Intercept the next wake word from a satellite."""
-    component: EntityComponent[AssistSatelliteEntity] = hass.data[DOMAIN]
-    satellite = component.get_entity(msg["entity_id"])
-    if satellite is None:
-        connection.send_error(
-            msg["id"], websocket_api.ERR_NOT_FOUND, "Entity not found"
-        )
-        return
-
-    satellite.async_stop_intercept_wake_word()
+    @callback
+    def wake_word_listener(wake_word_phrase: str | None, message: str) -> None:
+        """Push an intercepted wake word to websocket."""
+        if wake_word_phrase is not None:
+            connection.send_message(
+                websocket_api.event_message(
+                    msg["id"],
+                    {"wake_word_phrase": wake_word_phrase},
+                )
+            )
+        else:
+            connection.send_error(msg["id"], "home_assistant_error", message)
+
+    connection.subscriptions[msg["id"]] = satellite.async_intercept_wake_word(
+        wake_word_listener
+    )
+    connection.send_message(websocket_api.result_message(msg["id"]))

tests/components/assist_satellite/test_websocket_api.py

@@ -24,24 +24,25 @@ async def test_intercept_wake_word(
 
     await ws_client.send_json_auto_id(
         {
-            "type": "assist_satellite/intercept_wake_word/start",
+            "type": "assist_satellite/intercept_wake_word",
             "entity_id": ENTITY_ID,
         }
     )
-
-    for _ in range(3):
-        await asyncio.sleep(0)
+    msg = await ws_client.receive_json()
+    assert msg["success"]
+    assert msg["result"] is None
+    subscription_id = msg["id"]
 
     await entity.async_accept_pipeline_from_satellite(
-        object(),
+        object(),  # type: ignore[arg-type]
         start_stage=PipelineStage.STT,
         wake_word_phrase="ok, nabu",
     )
 
-    response = await ws_client.receive_json()
-
-    assert response["success"]
-    assert response["result"] == {"wake_word_phrase": "ok, nabu"}
+    msg = await ws_client.receive_json()
+    assert msg["id"] == subscription_id
+    assert msg["type"] == "event"
+    assert msg["event"] == {"wake_word_phrase": "ok, nabu"}
 
 
 async def test_intercept_wake_word_requires_on_device_wake_word(
@@ -55,16 +56,17 @@ async def test_intercept_wake_word_requires_on_device_wake_word(
 
     await ws_client.send_json_auto_id(
         {
-            "type": "assist_satellite/intercept_wake_word/start",
+            "type": "assist_satellite/intercept_wake_word",
             "entity_id": ENTITY_ID,
         }
     )
 
-    for _ in range(3):
-        await asyncio.sleep(0)
+    msg = await ws_client.receive_json()
+    assert msg["success"]
+    assert msg["result"] is None
 
     await entity.async_accept_pipeline_from_satellite(
-        object(),
+        object(),  # type: ignore[arg-type]
         # Emulate wake word processing in Home Assistant
         start_stage=PipelineStage.WAKE_WORD,
     )
@@ -88,16 +90,17 @@ async def test_intercept_wake_word_requires_wake_word_phrase(
 
     await ws_client.send_json_auto_id(
         {
-            "type": "assist_satellite/intercept_wake_word/start",
+            "type": "assist_satellite/intercept_wake_word",
             "entity_id": ENTITY_ID,
         }
     )
 
-    for _ in range(3):
-        await asyncio.sleep(0)
+    msg = await ws_client.receive_json()
+    assert msg["success"]
+    assert msg["result"] is None
 
     await entity.async_accept_pipeline_from_satellite(
-        object(),
+        object(),  # type: ignore[arg-type]
         start_stage=PipelineStage.STT,
         # We are not passing wake word phrase
     )
@@ -124,7 +127,7 @@ async def test_intercept_wake_word_require_admin(
 
     await ws_client.send_json_auto_id(
         {
-            "type": "assist_satellite/intercept_wake_word/start",
+            "type": "assist_satellite/intercept_wake_word",
             "entity_id": ENTITY_ID,
         }
     )
@@ -148,14 +151,14 @@ async def test_intercept_wake_word_invalid_satellite(
 
     await ws_client.send_json_auto_id(
         {
-            "type": "assist_satellite/intercept_wake_word/start",
+            "type": "assist_satellite/intercept_wake_word",
             "entity_id": "assist_satellite.invalid",
         }
     )
-    response = await ws_client.receive_json()
+    msg = await ws_client.receive_json()
 
-    assert not response["success"]
-    assert response["error"] == {
+    assert not msg["success"]
+    assert msg["error"] == {
         "code": "not_found",
         "message": "Entity not found",
     }
@@ -172,50 +175,33 @@ async def test_intercept_wake_word_twice(
 
     await ws_client.send_json_auto_id(
         {
-            "type": "assist_satellite/intercept_wake_word/start",
-            "entity_id": ENTITY_ID,
-        }
-    )
-
-    await ws_client.send_json_auto_id(
-        {
-            "type": "assist_satellite/intercept_wake_word/start",
+            "type": "assist_satellite/intercept_wake_word",
             "entity_id": ENTITY_ID,
         }
     )
-    response = await ws_client.receive_json()
-
-    assert response["success"]
-    assert response["result"] == {"wake_word_phrase": None}
 
+    msg = await ws_client.receive_json()
+    assert msg["success"]
+    assert msg["result"] is None
 
-async def test_intercept_wake_word_stop(
-    hass: HomeAssistant,
-    init_components: ConfigEntry,
-    entity: MockAssistSatellite,
-    hass_ws_client: WebSocketGenerator,
-) -> None:
-    """Test that we can stop intercepting wake words."""
-    ws_client = await hass_ws_client(hass)
+    task = asyncio.create_task(ws_client.receive_json())
 
     await ws_client.send_json_auto_id(
         {
-            "type": "assist_satellite/intercept_wake_word/start",
+            "type": "assist_satellite/intercept_wake_word",
             "entity_id": ENTITY_ID,
         }
     )
 
-    task = asyncio.create_task(ws_client.receive_json())
+    # Should get an error from previous subscription
+    msg = await task
+    assert not msg["success"]
+    assert msg["error"] == {
+        "code": "home_assistant_error",
+        "message": "Only one interception is allowed",
+    }
 
-    async with asyncio.timeout(1):
-        # Will cancel the open request
-        await ws_client.send_json_auto_id(
-            {
-                "type": "assist_satellite/intercept_wake_word/stop",
-                "entity_id": ENTITY_ID,
-            }
-        )
-
-        response = await task
-        assert response["success"]
-        assert response["result"] == {"wake_word_phrase": None}
+    # Response to second subscription
+    msg = await ws_client.receive_json()
+    assert msg["success"]
+    assert msg["result"] is None