music-assistant-server

1"""AI Radio Plugin Provider for Music Assistant."""
2
3from __future__ import annotations
4
5import asyncio
6from collections.abc import Callable
7from copy import deepcopy
8from pathlib import Path
9from typing import TYPE_CHECKING, Any
10from uuid import uuid4
11
12from music_assistant_models.auth import Scope
13from music_assistant_models.enums import EventType
14from music_assistant_models.errors import InvalidDataError, SetupFailedError
15
16from music_assistant.helpers.plugin_engines import (
17    select_ai_engine,
18    select_tts_engine,
19)
20from music_assistant.models.plugin import PluginProvider
21
22from .constants import (
23    CONF_AI_ENGINE,
24    CONF_TTS_ENGINE,
25    DEFAULT_MAX_CONCURRENT_RUNS,
26    ENGINE_DISCOVERY_TIMEOUT,
27    ENGINE_RECHECK_GRACE,
28    ENGINE_RETRY_DELAY,
29    MAX_FINISHED_SESSIONS,
30    SUPPORTED_FEATURES,
31    TRANSLATION_OWNER,
32)
33from .helpers import utc_now_iso
34from .models import SessionState
35from .rendering import AIRadioRenderMixin
36from .runtime import AIRadioRuntimeMixin
37from .storage import AIRadioStorageMixin
38
39if TYPE_CHECKING:
40    from music_assistant_models.config_entries import ConfigEntry, ProviderConfig
41    from music_assistant_models.event import MassEvent
42    from music_assistant_models.provider import ProviderManifest
43
44    from music_assistant.mass import MusicAssistant
45    from music_assistant.models import ProviderInstanceType
46
47
48async def setup(
49    mass: MusicAssistant, manifest: ProviderManifest, config: ProviderConfig
50) -> ProviderInstanceType:
51    """Initialize provider(instance) with given configuration."""
52    return AIRadioProvider(mass, manifest, config, SUPPORTED_FEATURES)
53
54
55class AIRadioProvider(AIRadioRuntimeMixin, AIRadioRenderMixin, AIRadioStorageMixin, PluginProvider):
56    """Implementation of the AI Radio plugin provider."""
57
58    def __init__(
59        self,
60        mass: MusicAssistant,
61        manifest: ProviderManifest,
62        config: ProviderConfig,
63        supported_features: set[Any],
64    ) -> None:
65        """Initialize the AI Radio provider."""
66        super().__init__(mass, manifest, config, supported_features)
67        self._station_lock = asyncio.Lock()
68        self._session_lock = asyncio.Lock()
69        self._unregister_handles: list[Callable[[], None]] = []
70        self._unloading = False
71        self._engine_recheck_task: asyncio.Task[None] | None = None
72        self._sessions: dict[str, SessionState] = {}
73        self._stations: dict[str, dict[str, Any]] = {}
74        self._sections: dict[str, dict[str, Any]] = {}
75        self._storage_dir = Path(self.mass.storage_path) / "ai_radio" / self.instance_id
76        self._stations_file = self._storage_dir / "stations.json"
77        self._sections_file = self._storage_dir / "sections.json"
78
79    async def get_config_entries(self) -> tuple[ConfigEntry, ...]:
80        """Return Config entries to configure this provider."""
81        from .config import get_config_entries as build_config_entries  # noqa: PLC0415
82
83        return await build_config_entries(self.mass, self.instance_id)
84
85    async def handle_async_init(self) -> None:
86        """Handle async initialization of the provider."""
87        await asyncio.to_thread(self._storage_dir.mkdir, parents=True, exist_ok=True)
88        await self._load_sections()
89        await self._load_stations()
90        await self._wait_for_engines()
91        self.logger.info(
92            "AI Radio initialized for instance '%s' with %d stations and %d sections",
93            self.instance_id,
94            len(self._stations),
95            len(self._sections),
96        )
97
98    async def loaded_in_mass(self) -> None:
99        """Call after the provider has been loaded."""
100        api_handlers = (
101            ("ai_radio/stations/list", self.list_stations),
102            ("ai_radio/stations/get", self.get_station),
103            ("ai_radio/stations/save", self.save_station),
104            ("ai_radio/stations/delete", self.delete_station),
105            ("ai_radio/stations/validate", self.validate_station),
106            ("ai_radio/stations/template", self.station_template),
107            ("ai_radio/sections/list", self.list_sections),
108            ("ai_radio/sections/get", self.get_section),
109            ("ai_radio/sections/save", self.save_section),
110            ("ai_radio/sections/delete", self.delete_section),
111            ("ai_radio/sections/template", self.section_template),
112            ("ai_radio/start", self.start_run),
113            ("ai_radio/stop", self.stop_run),
114            ("ai_radio/status", self.get_status),
115        )
116        for command, handler in api_handlers:
117            required_scope = (
118                Scope.CONFIG_PROVIDERS_READ
119                if command.endswith(("/list", "/get", "/template", "/validate", "/status"))
120                else Scope.CONFIG_PROVIDERS_WRITE
121            )
122            self._unregister_handles.append(
123                self.mass.register_api_command(command, handler, required_scope=required_scope)
124            )
125        self._unregister_handles.append(
126            self.mass.subscribe(self._on_providers_updated, EventType.PROVIDERS_UPDATED)
127        )
128        self.logger.info(
129            "AI Radio API routes registered (%d handlers)",
130            len(api_handlers),
131        )
132
133    async def unload(self, is_removed: bool = False) -> None:
134        """Handle close/cleanup of the provider."""
135        self._unloading = True
136        if self._engine_recheck_task and not self._engine_recheck_task.done():
137            self._engine_recheck_task.cancel()
138        cancelled = 0
139        for session in self._sessions.values():
140            if session.task and not session.task.done():
141                session.task.cancel()
142                cancelled += 1
143        for handle in self._unregister_handles:
144            handle()
145        self._unregister_handles.clear()
146        self.logger.info(
147            "AI Radio unloaded (removed=%s, cancelled_sessions=%d)",
148            is_removed,
149            cancelled,
150        )
151        await super().unload(is_removed)
152
153    async def list_stations(self) -> list[dict[str, Any]]:
154        """Return all configured AI Radio stations."""
155        return sorted(
156            (deepcopy(station) for station in self._stations.values()),
157            key=lambda station: station["name"],
158        )
159
160    async def get_station(self, station_id: str) -> dict[str, Any]:
161        """Return one station by id."""
162        if station_id not in self._stations:
163            raise KeyError(f"Unknown station id: {station_id}")
164        return deepcopy(self._stations[station_id])
165
166    async def save_station(self, station: dict[str, Any]) -> dict[str, Any]:
167        """Create or update a station."""
168        station_payload = deepcopy(station)
169        async with self._station_lock:
170            # validate against a scratch copy so a rejected station
171            # cannot leave partial section edits behind
172            sections_scratch = deepcopy(self._sections)
173            sections_changed = self._upsert_embedded_sections_from_station(
174                station_payload, sections_scratch
175            )
176            normalized = self._normalize_station(station_payload, sections_scratch)
177            self._stations[normalized["id"]] = normalized
178            if sections_changed:
179                self._sections = sections_scratch
180                await self._write_sections()
181            await self._write_stations()
182        self.logger.info("AI Radio station saved: %s (%s)", normalized["id"], normalized["name"])
183        return deepcopy(normalized)
184
185    async def delete_station(self, station_id: str) -> None:
186        """Delete a station."""
187        async with self._station_lock:
188            if station_id not in self._stations:
189                raise KeyError(f"Unknown station id: {station_id}")
190            self._stations.pop(station_id)
191            await self._write_stations()
192        self.logger.info("AI Radio station deleted: %s", station_id)
193
194    async def validate_station(self, station: dict[str, Any]) -> dict[str, Any]:
195        """Validate station payload and return the normalized profile."""
196        station_payload = deepcopy(station)
197        # validation must be side-effect free: resolve embedded sections
198        # against a scratch copy instead of the shared section store
199        sections_scratch = deepcopy(self._sections)
200        self._upsert_embedded_sections_from_station(station_payload, sections_scratch)
201        return self._normalize_station(station_payload, sections_scratch)
202
203    async def station_template(self) -> dict[str, Any]:
204        """Return a default station template."""
205        return self._default_station_template()
206
207    async def list_sections(self) -> list[dict[str, Any]]:
208        """Return all shared section definitions."""
209        return sorted(
210            (deepcopy(section) for section in self._sections.values()),
211            key=lambda section: section["id"].lower(),
212        )
213
214    async def get_section(self, section_id: str) -> dict[str, Any]:
215        """Return one shared section by id."""
216        if section_id not in self._sections:
217            raise KeyError(f"Unknown section id: {section_id}")
218        return deepcopy(self._sections[section_id])
219
220    async def save_section(self, section: dict[str, Any]) -> dict[str, Any]:
221        """Create or update a shared section."""
222        normalized = self._normalize_section(section)
223        async with self._station_lock:
224            self._sections[normalized["id"]] = normalized
225            self._refresh_station_sections()
226            await self._write_sections()
227            await self._write_stations()
228        self.logger.info("AI Radio section saved: %s", normalized["id"])
229        return deepcopy(normalized)
230
231    async def delete_section(self, section_id: str) -> None:
232        """Delete a shared section when unused by stations."""
233        async with self._station_lock:
234            if section_id not in self._sections:
235                raise KeyError(f"Unknown section id: {section_id}")
236            used_by = [
237                station["id"]
238                for station in self._stations.values()
239                if section_id in station.get("section_ids", [])
240            ]
241            if used_by:
242                used_list = ", ".join(sorted(used_by))
243                raise InvalidDataError(
244                    f"Section '{section_id}' is used by stations: {used_list}. "
245                    "Remove it from those stations first."
246                )
247            self._sections.pop(section_id)
248            await self._write_sections()
249        self.logger.info("AI Radio section deleted: %s", section_id)
250
251    async def section_template(self) -> dict[str, Any]:
252        """Return default section template."""
253        defaults = self._default_sections_template()
254        return deepcopy(defaults[0])
255
256    async def start_run(
257        self,
258        station_id: str,
259        source_playlist_id_override: str | None = None,
260        source_playlist_provider_override: str | None = None,
261        player_id_override: str | None = None,
262        dynamic_source_playtime_cap_override: int | float | None = None,  # noqa: PYI041
263    ) -> dict[str, Any]:
264        """Start a new AI Radio run."""
265        if station_id not in self._stations:
266            raise KeyError(f"Unknown station id: {station_id}")
267
268        station = deepcopy(self._stations[station_id])
269        overrides = {
270            "source_playlist_id": source_playlist_id_override,
271            "source_playlist_provider": source_playlist_provider_override,
272            "default_player_id": player_id_override,
273        }
274        for key, value in overrides.items():
275            if value:
276                station[key] = value
277        if dynamic_source_playtime_cap_override is not None:
278            if float(dynamic_source_playtime_cap_override) < 0:
279                raise InvalidDataError("dynamic_source_playtime_cap_override must be >= 0")
280            station["max_duration_minutes"] = float(dynamic_source_playtime_cap_override)
281        player_id = str(station.get("default_player_id") or "").strip()
282        if not player_id:
283            raise InvalidDataError("AI Radio requires a target player")
284        player = self.mass.players.get_player(player_id)
285        if player is None:
286            raise InvalidDataError(f"Unknown target player: {player_id}")
287        if player.available is False:
288            raise InvalidDataError(f"Target player is unavailable: {player_id}")
289        if player.enabled is False:
290            raise InvalidDataError(f"Target player is disabled: {player_id}")
291
292        # the run guards and the session insert must stay one critical section, or a future
293        # await between them would let concurrent callers slip past the concurrency limits
294        async with self._session_lock:
295            max_runs = DEFAULT_MAX_CONCURRENT_RUNS
296            running = [
297                session for session in self._sessions.values() if session.status == "running"
298            ]
299            if len(running) >= max_runs:
300                raise InvalidDataError(
301                    f"Max concurrent runs reached ({max_runs}). Stop an active run first."
302                )
303            if any(
304                session.status == "running" and session.station_id == station_id
305                for session in self._sessions.values()
306            ):
307                raise InvalidDataError(f"Station {station_id} already has an active run")
308
309            session_id = uuid4().hex
310            session = SessionState(
311                session_id=session_id,
312                station_id=station_id,
313            )
314            self._sessions[session_id] = session
315            self._prune_finished_sessions()
316            session.task = self.mass.create_task(
317                self._run_session(session_id, station),
318                task_id=f"ai_radio_session_{session_id}",
319            )
320        self.logger.debug(
321            "AI Radio session started: session=%s station=%s",
322            session_id,
323            station_id,
324        )
325        return session.as_dict()
326
327    async def stop_run(
328        self,
329        session_id: str | None = None,
330        station_id: str | None = None,
331    ) -> dict[str, Any]:
332        """Stop an active run."""
333        selected = self._resolve_session_for_stop(session_id=session_id, station_id=station_id)
334
335        # cancel first so the run cannot queue another batch after playback stopped
336        if selected.task and not selected.task.done():
337            selected.task.cancel()
338        selected.status = "stopped"
339        selected.ended_at = utc_now_iso()
340        await self._stop_session_queue(selected)
341        self.logger.info(
342            "AI Radio session stopped: session=%s station=%s",
343            selected.session_id,
344            selected.station_id,
345        )
346        return selected.as_dict()
347
348    async def get_status(self, session_id: str | None = None) -> dict[str, Any]:
349        """Return run status information."""
350        if session_id:
351            if session_id not in self._sessions:
352                raise KeyError(f"Unknown session id: {session_id}")
353            return {"sessions": [self._sessions[session_id].as_dict()]}
354        sessions = sorted(self._sessions.values(), key=lambda item: item.created_at, reverse=True)
355        return {"sessions": [session.as_dict() for session in sessions]}
356
357    async def _wait_for_engines(self, timeout: float | None = None) -> None:
358        """
359        Wait (bounded) until a concrete AI and TTS engine are selected for this instance.
360
361        :param timeout: How long to wait, defaulting to the engine discovery timeout.
362        :raises SetupFailedError: When either engine is still unavailable at the deadline.
363        """
364        engines_changed = asyncio.Event()
365        unsubscribe = self.mass.subscribe(
366            lambda _event: engines_changed.set(), EventType.PROVIDERS_UPDATED
367        )
368        try:
369            async with asyncio.timeout(ENGINE_DISCOVERY_TIMEOUT if timeout is None else timeout):
370                while (error := await self._engine_selection_error()) is not None:
371                    await engines_changed.wait()
372                    # clearing only after the wait keeps an update that lands during the
373                    # probe above signalled, so that wakeup is never lost
374                    engines_changed.clear()
375        except TimeoutError:
376            error = await self._engine_selection_error()
377        finally:
378            unsubscribe()
379        if error is not None:
380            raise error
381
382    async def _engine_selection_error(self) -> SetupFailedError | None:
383        """
384        Seed a concrete engine selection where none is stored yet.
385
386        :return: The error for the first engine that cannot be selected or no longer
387            resolves, or None when both engines are settled.
388        """
389        if await select_ai_engine(self, CONF_AI_ENGINE, in_setup_data=True) is None:
390            return SetupFailedError(
391                "AI Radio has no AI engine available",
392                translation_key="ai_radio_no_ai_engine",
393                translation_owner=TRANSLATION_OWNER,
394            )
395        if await select_tts_engine(self, CONF_TTS_ENGINE, in_setup_data=True) is None:
396            return SetupFailedError(
397                "AI Radio has no text-to-speech engine available",
398                translation_key="ai_radio_no_tts_engine",
399                translation_owner=TRANSLATION_OWNER,
400            )
401        return None
402
403    async def _on_providers_updated(self, _event: MassEvent) -> None:
404        """Re-check the engine selection whenever the set of loaded providers changes."""
405        # nothing to watch when this instance, or the whole server, is shutting down anyway
406        if self._unloading or self.mass.closing:
407            return
408        if self._engine_recheck_task and not self._engine_recheck_task.done():
409            return
410        if await self._engine_selection_error() is None:
411            return
412        self._engine_recheck_task = self.mass.create_task(self._unload_when_engines_stay_missing())
413
414    async def _unload_when_engines_stay_missing(self) -> None:
415        """Unload with an error when a vanished engine does not come back in time."""
416        # a plugin reload or a Home Assistant restart takes its engines with it for a
417        # while, so wait that out instead of tearing the provider down right away
418        try:
419            await self._wait_for_engines(ENGINE_RECHECK_GRACE)
420        except SetupFailedError as err:
421            # a shutdown (or our own unload) landing during the wait can surface as the
422            # timeout instead of a cancellation, and needs no error for the user
423            if self._unloading or self.mass.closing:
424                return
425            self.logger.warning("%s - unloading the provider", err)
426            self.unload_with_error(err)
427            # unloading records the error but arms no retry of its own, so schedule the
428            # reload that picks the provider back up once the engines return. Armed under
429            # the load path's task id, so any (re)load starting before it fires cancels it.
430            self.mass.call_later(
431                ENGINE_RETRY_DELAY,
432                self.mass.load_provider,
433                self.instance_id,
434                allow_retry=True,
435                task_id=f"load_provider_{self.instance_id}",
436            )
437
438    def _prune_finished_sessions(self) -> None:
439        """Drop the oldest finished sessions beyond the retention limit."""
440        finished = sorted(
441            (session for session in self._sessions.values() if session.status != "running"),
442            key=lambda item: item.created_at,
443            reverse=True,
444        )
445        for session in finished[MAX_FINISHED_SESSIONS:]:
446            self._sessions.pop(session.session_id, None)
447
448    def _resolve_session_for_stop(
449        self,
450        session_id: str | None,
451        station_id: str | None,
452    ) -> SessionState:
453        """Resolve which running session should be stopped."""
454        if session_id:
455            selected = self._sessions.get(session_id)
456            if selected is None:
457                raise KeyError(f"Unknown session id: {session_id}")
458            if selected.status != "running":
459                raise InvalidDataError(
460                    f"Session {session_id} is not running (status={selected.status})"
461                )
462            return selected
463
464        running = [session for session in self._sessions.values() if session.status == "running"]
465        if station_id:
466            running = [session for session in running if session.station_id == station_id]
467            if not running:
468                raise KeyError(f"No active run found for station: {station_id}")
469        elif not running:
470            raise KeyError("No active AI Radio run found")
471
472        return max(running, key=lambda item: item.created_at)
473