/
/
/
1"""Sticky per-queue AI DJ for AI Radio."""
2# mypy: disable-error-code=attr-defined
3
4from __future__ import annotations
5
6import asyncio
7import logging
8from typing import TYPE_CHECKING, Any, Literal
9from uuid import uuid4
10
11import aiofiles
12from music_assistant_models.enums import EventType
13from music_assistant_models.errors import InvalidDataError
14
15from music_assistant.controllers.player_queues.helpers import committed_index
16from music_assistant.helpers.json import async_json_loads
17
18from .constants import ATTR_GAP_NEXT_ID, ATTR_QUEUE_DJ, ATTR_SESSION_ID
19from .models import DJQueueState, PlannedSection, SessionState
20
21if TYPE_CHECKING:
22 from pathlib import Path
23
24 from music_assistant_models.event import MassEvent
25 from music_assistant_models.player_queue import PlayerQueue
26 from music_assistant_models.queue_item import QueueItem
27
28 from music_assistant.mass import MusicAssistant
29
30# a track without a known duration still has to count for something in the minute
31# bookkeeping, so it is billed as an average length song
32FALLBACK_TRACK_SECONDS = 210
33
34QUEUE_PAGE_SIZE = 500
35
36# per section history cap, generous for the widest guard window (60 minutes)
37HISTORY_EVENTS_PER_SECTION = 50
38
39# result of one splice attempt. only "gap_gone" leaves its gap unserved: a gap that already
40# holds a clip is served, and one that slipped behind the player can never be served again
41DJSpliceOutcome = Literal["injected", "gap_gone", "too_close", "occupied"]
42
43
44class AIRadioQueueDJMixin:
45 """Mixin managing sticky queue DJ state and clip injection."""
46
47 if TYPE_CHECKING:
48 mass: MusicAssistant
49 logger: logging.Logger
50 _hosts: dict[str, dict[str, Any]]
51 _dj_queues: dict[str, DJQueueState]
52 _sessions: dict[str, SessionState]
53 _dj_file: Path
54 _dj_lock: asyncio.Lock
55 _unloading: bool
56
57 async def set_queue_dj(self, queue_id: str, host_id: str | None) -> dict[str, str]:
58 """
59 Enable, switch or disable the sticky AI DJ on a queue.
60
61 :param queue_id: The queue to change.
62 :param host_id: The host to enable, or None to disable.
63 :return: The full queue-to-host mapping after the change.
64 """
65 queue_id = str(queue_id).strip()
66 if not queue_id:
67 raise InvalidDataError("queue_id is required")
68 armed: DJQueueState | None = None
69 try:
70 async with self._dj_lock:
71 if host_id is None:
72 self._dj_queues.pop(queue_id, None)
73 else:
74 host_id = str(host_id).strip()
75 if host_id not in self._hosts:
76 raise InvalidDataError(f"Unknown host id: {host_id}")
77 armed = self._arm_dj_state(queue_id, host_id)
78 await self._write_queue_dj()
79 # stale clips carry the old host's persona, so they must go before a replan
80 # can reuse the gaps they occupy
81 self._remove_pending_dj_clips(queue_id)
82 except Exception:
83 # an armed state that never reached its cleanup stays unready forever, which
84 # reads as an armed DJ that never speaks. dropping it lets a retry arm cleanly
85 if armed is not None:
86 async with self._dj_lock:
87 if self._dj_queues.get(queue_id) is armed:
88 del self._dj_queues[queue_id]
89 raise
90 if armed is not None:
91 if self._dj_queues.get(queue_id) is armed:
92 # only if we're still the live state: a newer switch may have replaced us.
93 # flipped after cleanup so a racing pass doesn't mark old clips' gaps served
94 armed.ready = True
95 self._schedule_replan(queue_id)
96 return await self.get_queue_dj_status()
97
98 async def get_queue_dj_status(self) -> dict[str, str]:
99 """Return the queue-to-host mapping of all active queue DJs."""
100 return {queue_id: state.host_id for queue_id, state in self._dj_queues.items()}
101
102 async def _load_queue_dj(self) -> None:
103 """Load persisted queue DJ assignments and arm their states."""
104 file_exists = await asyncio.to_thread(self._dj_file.exists)
105 if not file_exists:
106 self._dj_queues = {}
107 return
108 async with aiofiles.open(self._dj_file) as file_handle:
109 content = await file_handle.read()
110 try:
111 payload = await async_json_loads(content)
112 except ValueError as err:
113 self.logger.error("Queue DJ file is corrupt, starting without queue DJs: %s", err)
114 payload = {}
115 queues = payload.get("queues", {}) if isinstance(payload, dict) else {}
116 self._dj_queues = {}
117 if isinstance(queues, dict):
118 for queue_id, entry in queues.items():
119 host_id = str(entry.get("host_id", "")).strip() if isinstance(entry, dict) else ""
120 if host_id not in self._hosts:
121 self.logger.warning(
122 "Dropping queue DJ for %s: host %s no longer exists", queue_id, host_id
123 )
124 continue
125 # no clip cleanup precedes a boot arm, so this state may plan right away
126 self._arm_dj_state(str(queue_id), host_id).ready = True
127
128 async def _write_queue_dj(self) -> None:
129 """Persist queue DJ assignments to disk."""
130 payload = {
131 "version": 1,
132 "queues": {
133 queue_id: {"host_id": state.host_id}
134 for queue_id, state in sorted(self._dj_queues.items())
135 },
136 }
137 await self._write_json_file(self._dj_file, payload)
138
139 def _arm_dj_state(self, queue_id: str, host_id: str) -> DJQueueState:
140 """Create fresh in-memory DJ state for a queue."""
141 # a fresh session id per arm keeps clip ids from colliding with clips
142 # persisted in the queue by a previous run of this provider
143 state = DJQueueState(
144 queue_id=queue_id,
145 host_id=host_id,
146 dj_session_id=f"dj{uuid4().hex[:12]}",
147 )
148 self._dj_queues[queue_id] = state
149 return state
150
151 async def _on_dj_queue_event(self, event: MassEvent) -> None:
152 """Handle queue and player events for the queues that run a DJ."""
153 queue_id = str(event.object_id or "")
154 if queue_id not in self._dj_queues:
155 return
156 if event.event == EventType.PLAYER_REMOVED:
157 async with self._dj_lock:
158 self._dj_queues.pop(queue_id, None)
159 await self._write_queue_dj()
160 self.logger.debug("Dropped queue DJ for removed player %s", queue_id)
161 return
162 self._schedule_replan(queue_id)
163
164 def _schedule_replan(self, queue_id: str) -> None:
165 """Request a replan pass for the given queue."""
166 if self._unloading:
167 return
168 state = self._dj_queues.get(queue_id)
169 if state is None or state.replan_pending:
170 return
171 state.replan_pending = True
172 state.task = self.mass.create_task(
173 self._drain_replans(queue_id), task_id=f"ai_radio_dj_replan_{queue_id}"
174 )
175
176 async def _drain_replans(self, queue_id: str) -> None:
177 """Run replan passes until no further request landed during the last one."""
178 # the inserts of a pass re-fire QUEUE_ITEMS_UPDATED while this task still holds the
179 # replan task id, so its follow-up request is served here instead of by a new task
180 while (state := self._dj_queues.get(queue_id)) is not None and state.replan_pending:
181 if self._unloading:
182 return
183 try:
184 await self._replan_queue(queue_id)
185 except Exception:
186 # cleared (not left pending) so a later event can retry without hot-looping
187 # here; re-fetched since a re-arm mid-pass may have swapped in a new state
188 self.logger.exception("Queue DJ replan failed for %s", queue_id)
189 if (live_state := self._dj_queues.get(queue_id)) is not None:
190 live_state.replan_pending = False
191 return
192
193 async def _replan_queue(self, queue_id: str) -> None: # noqa: PLR0915
194 """Run one planning, injection and repair pass over a queue."""
195 state = self._dj_queues.get(queue_id)
196 if state is None:
197 return
198 async with state.lock:
199 # cleared up front so an event landing mid-pass requests a fresh pass
200 state.replan_pending = False
201 if not state.ready:
202 # a switch armed this state but hasn't finished clearing the old clips yet;
203 # planning now would mark their gaps served. the switch replans once ready
204 self.logger.debug("Queue %s is waiting for its DJ switch cleanup", queue_id)
205 return
206 if any(
207 session.status == "running" and session.queue_id == queue_id
208 for session in self._sessions.values()
209 ):
210 # a show plans its own breaks into this queue and its clips carry no DJ
211 # attribute, so injecting here would stack talk on top of talk
212 self.logger.debug("Queue %s is running a show, skipping replan", queue_id)
213 return
214 queue = self.mass.player_queues.get(queue_id)
215 if queue is None:
216 # usually the queue just hasn't registered yet (players appear seconds after
217 # load); state is kept, QUEUE_ADDED resumes it, PLAYER_REMOVED is what drops it
218 self.logger.debug("Queue %s is not registered (yet), skipping replan", queue_id)
219 return
220 items = self._dj_queue_items(queue_id)
221 guard_index = self._dj_guard_index(queue)
222 if self._repair_dj_clips(queue_id, state, items, guard_index):
223 items = self._dj_queue_items(queue_id)
224 # tracks that left the queue keep no decision, so the set cannot grow unbounded
225 state.decided_gap_ids &= {item.queue_item_id for item in items}
226
227 window = self._dj_window(items, guard_index)
228 if len(window) < 2:
229 self.logger.debug(
230 "Queue %s has no plannable gap ahead of the player, skipping replan", queue_id
231 )
232 return
233 # measured from the same point the planner counts from, or OPTIONAL guard
234 # positions drift between passes. recomputed every pass so state self-corrects
235 state.songs_before_window, state.minutes_before_window = self._dj_window_offsets(
236 items, window[0].queue_item_id
237 )
238 host = self._hosts.get(state.host_id)
239 if host is None:
240 self.logger.warning(
241 "Disabling queue DJ on %s: host %s no longer exists", queue_id, state.host_id
242 )
243 async with self._dj_lock:
244 self._dj_queues.pop(queue_id, None)
245 await self._write_queue_dj()
246 return
247
248 window_tracks = [
249 self._queue_item_to_track(index, item) for index, item in enumerate(window)
250 ]
251 program = self._build_program({"id": "", "name": f"AI DJ {host['name']}"}, host)
252 runtime_tokens = await self._prepare_runtime_tokens(program)
253 if self._dj_queues.get(queue_id) is not state:
254 # a switch or disable replaced this queue's state while the fetch above
255 # was in flight, so the session this pass planned for is gone
256 return
257 # every gap the planner is about to evaluate, so gaps where chance or a guard
258 # picks nothing count as decided too instead of being rolled again next pass
259 evaluated_gap_ids = {
260 str(track["item_id"])
261 for track in window_tracks[1:]
262 if track["item_id"] not in state.decided_gap_ids
263 }
264 planned, history = self._plan_sections(
265 session_id=state.dj_session_id,
266 tracks=window_tracks,
267 program=program,
268 track_index_offset=state.songs_before_window,
269 minute_offset=state.minutes_before_window,
270 history_state=state.history,
271 allowed_slot_when=["between_songs"],
272 runtime_tokens=runtime_tokens,
273 decided_next_item_ids=state.decided_gap_ids,
274 )
275 # spliced into a working copy and applied as one update; a call per clip would
276 # flood every client with queue events
277 working = self._dj_queue_items(queue_id)
278 # re-read: the planning awaits above can take seconds, letting the player buffer
279 # further ahead
280 guard_index = self._dj_guard_index(queue)
281 # insert order is not load bearing: every clip resolves its own target position
282 # in the working copy. descending walks back from the queue tail
283 injected = 0
284 skipped: dict[str, int] = {}
285 rejected: list[PlannedSection] = []
286 for section in sorted(planned, key=lambda item: item.insert_at_index, reverse=True):
287 target = window_tracks[section.insert_at_index]
288 outcome = self._splice_dj_clip(
289 queue_id=queue_id,
290 items=working,
291 guard_index=guard_index,
292 state=state,
293 program=program,
294 target=target,
295 section=section,
296 )
297 if outcome == "injected":
298 injected += 1
299 else:
300 rejected.append(section)
301 skipped[outcome] = skipped.get(outcome, 0) + 1
302 if outcome == "gap_gone":
303 # the target moved or left the queue, so nothing was decided about
304 # its gap and a later pass has to look at it again
305 evaluated_gap_ids.discard(str(target["item_id"]))
306 if injected:
307 self.mass.player_queues.update_items(queue_id, working)
308 # a rejected clip never airs, so its guard event must be removed: left in, it
309 # would block its own successor for a full guard window and double-count later
310 for section in rejected:
311 for section_id, event in section.history_events:
312 if (events := history.get(section_id)) and event in events:
313 events.remove(event)
314 # only the newest events matter to the guards (the last one for min_gap_songs,
315 # a 60 minute window for max_per_60min), so the tail is dropped
316 state.history = {
317 section_id: events[-HISTORY_EVENTS_PER_SECTION:]
318 for section_id, events in history.items()
319 }
320 state.decided_gap_ids |= evaluated_gap_ids
321 self.logger.debug(
322 "Replanned queue %s: window %s tracks, %s open gaps, planned %s, injected %s, "
323 "skipped %s, decided %s",
324 queue_id,
325 len(window_tracks),
326 len(evaluated_gap_ids),
327 len(planned),
328 injected,
329 skipped,
330 len(state.decided_gap_ids),
331 )
332
333 def _splice_dj_clip(
334 self,
335 queue_id: str,
336 items: list[QueueItem],
337 guard_index: int,
338 state: DJQueueState,
339 program: dict[str, Any],
340 target: dict[str, Any],
341 section: PlannedSection,
342 ) -> DJSpliceOutcome:
343 """Insert one planned clip in front of its target track and report the outcome."""
344 target_index = next(
345 (index for index, item in enumerate(items) if item.queue_item_id == target["item_id"]),
346 None,
347 )
348 if target_index is None:
349 return "gap_gone"
350 if target_index <= guard_index + 1:
351 return "too_close"
352 if items[target_index - 1].extra_attributes.get(ATTR_QUEUE_DJ):
353 return "occupied"
354 # the planner numbers its clips from zero every pass, so the id comes from the
355 # state counter instead to stay unique for the lifetime of the session
356 section.clip_id = f"{state.dj_session_id}_{state.clip_counter:03d}"
357 state.clip_counter += 1
358 clip = self._section_to_clip_item(queue_id, state.dj_session_id, program, section)
359 clip.extra_attributes[ATTR_QUEUE_DJ] = True
360 clip.extra_attributes[ATTR_GAP_NEXT_ID] = target["item_id"]
361 # sharing the target's sort index keeps the clip in front of the track it announces
362 # when the queue is un-shuffled, without renumbering everything behind it
363 clip.sort_index = items[target_index].sort_index
364 items.insert(target_index, clip)
365 return "injected"
366
367 def _remove_pending_dj_clips(self, queue_id: str) -> None:
368 """Remove not-yet-played DJ clips from the queue, except the armed session's own."""
369 # await-free on purpose: nothing can mutate the queue between the snapshot below and
370 # the update that applies the filtered list
371 queue = self.mass.player_queues.get(queue_id)
372 if queue is None:
373 return
374 # only the live session's clips survive, so a re-enable racing this cleanup keeps
375 # its own; clips from a session nothing remembers (ids re-roll each load) are cleared
376 live_state = self._dj_queues.get(queue_id)
377 keep_session_id = live_state.dj_session_id if live_state is not None else None
378 guard_index = self._dj_guard_index(queue)
379 items = self._dj_queue_items(queue_id)
380 # one update for the whole cleanup: a delete per clip floods every connected client
381 # with queue events. items up to the guard are what the player already owns
382 kept = [
383 item
384 for index, item in enumerate(items)
385 if index <= guard_index
386 or not item.extra_attributes.get(ATTR_QUEUE_DJ)
387 or (
388 keep_session_id is not None
389 and item.extra_attributes.get(ATTR_SESSION_ID) == keep_session_id
390 )
391 ]
392 if (removed := len(items) - len(kept)) == 0:
393 return
394 self.mass.player_queues.update_items(queue_id, kept)
395 self.logger.debug("Removed %s pending DJ clip(s) from queue %s", removed, queue_id)
396
397 def _dj_guard_index(self, queue: PlayerQueue) -> int:
398 """Return the highest queue index the player already owns."""
399 # the player owns the current and the already buffered item, and the slot right
400 # after the buffered one may be handed to the player at any moment
401 boundary_index = committed_index(queue)
402 return boundary_index if boundary_index is not None else -1
403
404 def _repair_dj_clips(
405 self, queue_id: str, state: DJQueueState, items: list[QueueItem], guard_index: int
406 ) -> bool:
407 """Delete DJ clips that no longer sit in front of the track they announce."""
408 stale_ids: set[str] = set()
409 for index in range(guard_index + 2, len(items)):
410 item = items[index]
411 if not item.extra_attributes.get(ATTR_QUEUE_DJ):
412 continue
413 successor = items[index + 1] if index + 1 < len(items) else None
414 if successor is not None and successor.queue_item_id == item.extra_attributes.get(
415 ATTR_GAP_NEXT_ID
416 ):
417 continue
418 stale_ids.add(item.queue_item_id)
419 # the gap this clip was serving is open again, so let a later pass decide it anew
420 if (gap_next_id := item.extra_attributes.get(ATTR_GAP_NEXT_ID)) is not None:
421 state.decided_gap_ids.discard(str(gap_next_id))
422 if not stale_ids:
423 return False
424 # one update for all of them, so the clients see a single queue change
425 self.mass.player_queues.update_items(
426 queue_id, [item for item in items if item.queue_item_id not in stale_ids]
427 )
428 self.logger.debug(
429 "Repaired queue %s: deleted %s stale DJ clip(s)", queue_id, len(stale_ids)
430 )
431 return True
432
433 def _dj_window(self, items: list[QueueItem], guard_index: int) -> list[QueueItem]:
434 """Return the upcoming music items that this pass may plan against."""
435 # every upcoming track, decided or not: the planner counts songs and minutes over a
436 # contiguous run, and per gap decisions are what keeps the work from being redone
437 return [
438 item
439 for item in items[guard_index + 1 :]
440 if not item.extra_attributes.get(ATTR_QUEUE_DJ)
441 ]
442
443 def _dj_window_offsets(self, items: list[QueueItem], window_start_id: str) -> tuple[int, float]:
444 """Return the songs and minutes of music playing before the first window track."""
445 behind = []
446 for item in items:
447 if item.queue_item_id == window_start_id:
448 break
449 if not item.extra_attributes.get(ATTR_QUEUE_DJ):
450 behind.append(item)
451 minutes = sum(item.duration or FALLBACK_TRACK_SECONDS for item in behind) / 60.0
452 return len(behind), minutes
453
454 def _dj_queue_items(self, queue_id: str) -> list[QueueItem]:
455 """Return all items of a queue."""
456 items: list[QueueItem] = []
457 offset = 0
458 while True:
459 page = self.mass.player_queues.items(queue_id, limit=QUEUE_PAGE_SIZE, offset=offset)
460 items.extend(page)
461 if len(page) < QUEUE_PAGE_SIZE:
462 return items
463 offset += QUEUE_PAGE_SIZE
464
465 def _queue_item_to_track(self, index: int, item: QueueItem) -> dict[str, Any]:
466 """Convert a music queue item into the track dict the planner consumes."""
467 name = ""
468 artist = ""
469 if (media_item := item.media_item) is not None:
470 name = str(media_item.name or "")
471 artists = getattr(media_item, "artists", None)
472 if artists:
473 artist = str(artists[0].name)
474 if not name:
475 raw_name = str(item.name or "")
476 artist, separator, name = raw_name.partition(" - ")
477 if not separator:
478 artist, name = "", raw_name
479 return {
480 "index": index,
481 "item_id": item.queue_item_id,
482 "name": name,
483 "artist": artist,
484 "songinfo": f"{artist} - {name}".strip(" -"),
485 "duration": item.duration,
486 "media_item": None,
487 }
488