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