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