From 123b6b687e8107bc816d99f6782823bffe199bdd Mon Sep 17 00:00:00 2001
From: Michael Hansen <mike@rhasspy.org>
Date: Fri, 20 Sep 2024 12:57:55 -0500
Subject: [PATCH] Route non-TTS media through ESPHome ffmpeg proxy (#126287)

* Route non-TTS media through proxy

* Use media_id_source
---
 .../components/esphome/assist_satellite.py    | 30 +++++++++++++-
 .../esphome/test_assist_satellite.py          | 40 ++++++++++++++++++-
 2 files changed, 67 insertions(+), 3 deletions(-)

diff --git a/homeassistant/components/esphome/assist_satellite.py b/homeassistant/components/esphome/assist_satellite.py
index a0e05a6c565..1485d88a7d2 100644
--- a/homeassistant/components/esphome/assist_satellite.py
+++ b/homeassistant/components/esphome/assist_satellite.py
@@ -14,6 +14,7 @@ import wave
 
 from aioesphomeapi import (
     MediaPlayerFormatPurpose,
+    MediaPlayerSupportedFormat,
     VoiceAssistantAnnounceFinished,
     VoiceAssistantAudioSettings,
     VoiceAssistantCommandFlag,
@@ -44,6 +45,7 @@ from .const import DOMAIN
 from .entity import EsphomeAssistEntity
 from .entry_data import ESPHomeConfigEntry, RuntimeEntryData
 from .enum_mapper import EsphomeEnumMapper
+from .ffmpeg_proxy import async_create_proxy_url
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -325,8 +327,34 @@ class EsphomeAssistSatellite(
             announcement.message,
             announcement.media_id,
         )
+        media_id = announcement.media_id
+        if announcement.media_id_source != "tts":
+            # Route non-TTS media through the proxy
+            format_to_use: MediaPlayerSupportedFormat | None = None
+            for supported_format in chain(
+                *self.entry_data.media_player_formats.values()
+            ):
+                if supported_format.purpose == MediaPlayerFormatPurpose.ANNOUNCEMENT:
+                    format_to_use = supported_format
+                    break
+
+            if format_to_use is not None:
+                assert (self.registry_entry is not None) and (
+                    self.registry_entry.device_id is not None
+                )
+                proxy_url = async_create_proxy_url(
+                    self.hass,
+                    self.registry_entry.device_id,
+                    media_id,
+                    media_format=format_to_use.format,
+                    rate=format_to_use.sample_rate or None,
+                    channels=format_to_use.num_channels or None,
+                    width=format_to_use.sample_bytes or None,
+                )
+                media_id = async_process_play_media_url(self.hass, proxy_url)
+
         await self.cli.send_voice_assistant_announcement_await_response(
-            announcement.media_id, _ANNOUNCEMENT_TIMEOUT_SEC, announcement.message
+            media_id, _ANNOUNCEMENT_TIMEOUT_SEC, announcement.message
         )
 
     async def handle_pipeline_start(
diff --git a/tests/components/esphome/test_assist_satellite.py b/tests/components/esphome/test_assist_satellite.py
index 03111c0d8d8..71bae989daf 100644
--- a/tests/components/esphome/test_assist_satellite.py
+++ b/tests/components/esphome/test_assist_satellite.py
@@ -1222,11 +1222,29 @@ async def test_announce_media_id(
         [APIClient, list[EntityInfo], list[UserService], list[EntityState]],
         Awaitable[MockESPHomeDevice],
     ],
+    device_registry: dr.DeviceRegistry,
 ) -> None:
     """Test announcement with media id."""
     mock_device: MockESPHomeDevice = await mock_esphome_device(
         mock_client=mock_client,
-        entity_info=[],
+        entity_info=[
+            MediaPlayerInfo(
+                object_id="mymedia_player",
+                key=1,
+                name="my media_player",
+                unique_id="my_media_player",
+                supports_pause=True,
+                supported_formats=[
+                    MediaPlayerSupportedFormat(
+                        format="flac",
+                        sample_rate=48000,
+                        num_channels=2,
+                        purpose=MediaPlayerFormatPurpose.ANNOUNCEMENT,
+                        sample_bytes=2,
+                    ),
+                ],
+            )
+        ],
         user_service=[],
         states=[],
         device_info={
@@ -1238,6 +1256,10 @@ async def test_announce_media_id(
     )
     await hass.async_block_till_done()
 
+    dev = device_registry.async_get_device(
+        connections={(dr.CONNECTION_NETWORK_MAC, mock_device.entry.unique_id)}
+    )
+
     satellite = get_satellite_entity(hass, mock_device.device_info.mac_address)
     assert satellite is not None
 
@@ -1247,7 +1269,7 @@ async def test_announce_media_id(
         media_id: str, timeout: float, text: str
     ):
         assert satellite.state == AssistSatelliteState.RESPONDING
-        assert media_id == "https://www.home-assistant.io/resolved.mp3"
+        assert media_id == "https://www.home-assistant.io/proxied.flac"
 
         done.set()
 
@@ -1257,6 +1279,10 @@ async def test_announce_media_id(
             "send_voice_assistant_announcement_await_response",
             new=send_voice_assistant_announcement_await_response,
         ),
+        patch(
+            "homeassistant.components.esphome.assist_satellite.async_create_proxy_url",
+            return_value="https://www.home-assistant.io/proxied.flac",
+        ) as mock_async_create_proxy_url,
     ):
         async with asyncio.timeout(1):
             await hass.services.async_call(
@@ -1271,6 +1297,16 @@ async def test_announce_media_id(
             await done.wait()
             assert satellite.state == AssistSatelliteState.LISTENING_WAKE_WORD
 
+        mock_async_create_proxy_url.assert_called_once_with(
+            hass,
+            dev.id,
+            "https://www.home-assistant.io/resolved.mp3",
+            media_format="flac",
+            rate=48000,
+            channels=2,
+            width=2,
+        )
+
 
 async def test_satellite_unloaded_on_disconnect(
     hass: HomeAssistant,