/
/
/
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 decided_before = state.decided_gap_ids
226 state.decided_gap_ids = decided_before & {item.queue_item_id for item in items}
227 if decided_before and not state.decided_gap_ids:
228 # nothing this history was recorded against is left, so its queue is gone
229 self._drop_unaired_dj_history(state, items, guard_index)
230
231 window = self._dj_window(items, guard_index)
232 if len(window) < 2:
233 self.logger.debug(
234 "Queue %s has no plannable gap ahead of the player, skipping replan", queue_id
235 )
236 return
237 # measured from the same point the planner counts from, or OPTIONAL guard
238 # positions drift between passes. recomputed every pass so state self-corrects
239 offsets = self._dj_window_offsets(items, window[0].queue_item_id)
240 # a lower song count means the queue rewound under the history, e.g. a clear or a
241 # jump back to the top; minutes dip on their own when a probed duration lands
242 if offsets[0] < state.songs_before_window:
243 self._rebase_dj_history(state, *offsets)
244 state.songs_before_window, state.minutes_before_window = offsets
245 host = self._hosts.get(state.host_id)
246 if host is None:
247 self.logger.warning(
248 "Disabling queue DJ on %s: host %s no longer exists", queue_id, state.host_id
249 )
250 async with self._dj_lock:
251 self._dj_queues.pop(queue_id, None)
252 await self._write_queue_dj()
253 return
254
255 window_tracks = [
256 self._queue_item_to_track(index, item) for index, item in enumerate(window)
257 ]
258 program = self._build_program({"id": "", "name": f"AI DJ {host['name']}"}, host)
259 runtime_tokens = await self._prepare_runtime_tokens(program)
260 if self._dj_queues.get(queue_id) is not state:
261 # a switch or disable replaced this queue's state while the fetch above
262 # was in flight, so the session this pass planned for is gone
263 return
264 # every gap the planner is about to evaluate, so gaps where chance or a guard
265 # picks nothing count as decided too instead of being rolled again next pass
266 evaluated_gap_ids = {
267 str(track["item_id"])
268 for track in window_tracks[1:]
269 if track["item_id"] not in state.decided_gap_ids
270 }
271 planned, history = self._plan_sections(
272 session_id=state.dj_session_id,
273 tracks=window_tracks,
274 program=program,
275 track_index_offset=state.songs_before_window,
276 minute_offset=state.minutes_before_window,
277 history_state=state.history,
278 allowed_slot_when=["between_songs"],
279 runtime_tokens=runtime_tokens,
280 decided_next_item_ids=state.decided_gap_ids,
281 )
282 # spliced into a working copy and applied as one update; a call per clip would
283 # flood every client with queue events
284 working = self._dj_queue_items(queue_id)
285 # re-read: the planning awaits above can take seconds, letting the player buffer
286 # further ahead
287 guard_index = self._dj_guard_index(queue)
288 # insert order is not load bearing: every clip resolves its own target position
289 # in the working copy. descending walks back from the queue tail
290 injected = 0
291 skipped: dict[str, int] = {}
292 rejected: list[PlannedSection] = []
293 for section in sorted(planned, key=lambda item: item.insert_at_index, reverse=True):
294 target = window_tracks[section.insert_at_index]
295 outcome = self._splice_dj_clip(
296 queue_id=queue_id,
297 items=working,
298 guard_index=guard_index,
299 state=state,
300 program=program,
301 target=target,
302 section=section,
303 )
304 if outcome == "injected":
305 injected += 1
306 else:
307 rejected.append(section)
308 skipped[outcome] = skipped.get(outcome, 0) + 1
309 if outcome == "gap_gone":
310 # the target moved or left the queue, so nothing was decided about
311 # its gap and a later pass has to look at it again
312 evaluated_gap_ids.discard(str(target["item_id"]))
313 if injected:
314 self.mass.player_queues.update_items(queue_id, working)
315 # a rejected clip never airs, so its guard event must be removed: left in, it
316 # would block its own successor for a full guard window and double-count later
317 for section in rejected:
318 for section_id, event in section.history_events:
319 if (events := history.get(section_id)) and event in events:
320 events.remove(event)
321 # only the newest events matter to the guards (the last one for min_gap_songs,
322 # a 60 minute window for max_per_60min), so the tail is dropped
323 state.history = {
324 section_id: events[-HISTORY_EVENTS_PER_SECTION:]
325 for section_id, events in history.items()
326 }
327 state.decided_gap_ids |= evaluated_gap_ids
328 self.logger.debug(
329 "Replanned queue %s: window %s tracks, %s open gaps, planned %s, injected %s, "
330 "skipped %s, decided %s",
331 queue_id,
332 len(window_tracks),
333 len(evaluated_gap_ids),
334 len(planned),
335 injected,
336 skipped,
337 len(state.decided_gap_ids),
338 )
339
340 def _splice_dj_clip(
341 self,
342 queue_id: str,
343 items: list[QueueItem],
344 guard_index: int,
345 state: DJQueueState,
346 program: dict[str, Any],
347 target: dict[str, Any],
348 section: PlannedSection,
349 ) -> DJSpliceOutcome:
350 """Insert one planned clip in front of its target track and report the outcome."""
351 target_index = next(
352 (index for index, item in enumerate(items) if item.queue_item_id == target["item_id"]),
353 None,
354 )
355 if target_index is None:
356 return "gap_gone"
357 if target_index <= guard_index + 1:
358 return "too_close"
359 if items[target_index - 1].extra_attributes.get(ATTR_QUEUE_DJ):
360 return "occupied"
361 # the planner numbers its clips from zero every pass, so the id comes from the
362 # state counter instead to stay unique for the lifetime of the session
363 section.clip_id = f"{state.dj_session_id}_{state.clip_counter:03d}"
364 state.clip_counter += 1
365 clip = self._section_to_clip_item(queue_id, state.dj_session_id, program, section)
366 clip.extra_attributes[ATTR_QUEUE_DJ] = True
367 clip.extra_attributes[ATTR_GAP_NEXT_ID] = target["item_id"]
368 # sharing the target's sort index keeps the clip in front of the track it announces
369 # when the queue is un-shuffled, without renumbering everything behind it
370 clip.sort_index = items[target_index].sort_index
371 items.insert(target_index, clip)
372 return "injected"
373
374 def _remove_pending_dj_clips(self, queue_id: str) -> None:
375 """Remove not-yet-played DJ clips from the queue, except the armed session's own."""
376 # await-free on purpose: nothing can mutate the queue between the snapshot below and
377 # the update that applies the filtered list
378 queue = self.mass.player_queues.get(queue_id)
379 if queue is None:
380 return
381 # only the live session's clips survive, so a re-enable racing this cleanup keeps
382 # its own; clips from a session nothing remembers (ids re-roll each load) are cleared
383 live_state = self._dj_queues.get(queue_id)
384 keep_session_id = live_state.dj_session_id if live_state is not None else None
385 guard_index = self._dj_guard_index(queue)
386 items = self._dj_queue_items(queue_id)
387 # one update for the whole cleanup: a delete per clip floods every connected client
388 # with queue events. items up to the guard are what the player already owns
389 kept = [
390 item
391 for index, item in enumerate(items)
392 if index <= guard_index
393 or not item.extra_attributes.get(ATTR_QUEUE_DJ)
394 or (
395 keep_session_id is not None
396 and item.extra_attributes.get(ATTR_SESSION_ID) == keep_session_id
397 )
398 ]
399 if (removed := len(items) - len(kept)) == 0:
400 return
401 self.mass.player_queues.update_items(queue_id, kept)
402 self.logger.debug("Removed %s pending DJ clip(s) from queue %s", removed, queue_id)
403
404 def _dj_guard_index(self, queue: PlayerQueue) -> int:
405 """Return the highest queue index the player already owns."""
406 # the player owns the current and the already buffered item, and the slot right
407 # after the buffered one may be handed to the player at any moment
408 boundary_index = committed_index(queue)
409 return boundary_index if boundary_index is not None else -1
410
411 def _repair_dj_clips(
412 self, queue_id: str, state: DJQueueState, items: list[QueueItem], guard_index: int
413 ) -> bool:
414 """Delete DJ clips that no longer sit in front of the track they announce."""
415 stale_ids: set[str] = set()
416 for index in range(guard_index + 2, len(items)):
417 item = items[index]
418 if not item.extra_attributes.get(ATTR_QUEUE_DJ):
419 continue
420 successor = items[index + 1] if index + 1 < len(items) else None
421 if successor is not None and successor.queue_item_id == item.extra_attributes.get(
422 ATTR_GAP_NEXT_ID
423 ):
424 continue
425 stale_ids.add(item.queue_item_id)
426 # the gap this clip was serving is open again, so let a later pass decide it anew
427 if (gap_next_id := item.extra_attributes.get(ATTR_GAP_NEXT_ID)) is not None:
428 state.decided_gap_ids.discard(str(gap_next_id))
429 if not stale_ids:
430 return False
431 # one update for all of them, so the clients see a single queue change
432 self.mass.player_queues.update_items(
433 queue_id, [item for item in items if item.queue_item_id not in stale_ids]
434 )
435 self.logger.debug(
436 "Repaired queue %s: deleted %s stale DJ clip(s)", queue_id, len(stale_ids)
437 )
438 return True
439
440 def _dj_window(self, items: list[QueueItem], guard_index: int) -> list[QueueItem]:
441 """Return the upcoming music items that this pass may plan against."""
442 # every upcoming track, decided or not: the planner counts songs and minutes over a
443 # contiguous run, and per gap decisions are what keeps the work from being redone
444 return [
445 item
446 for item in items[guard_index + 1 :]
447 if not item.extra_attributes.get(ATTR_QUEUE_DJ)
448 ]
449
450 def _dj_window_offsets(self, items: list[QueueItem], window_start_id: str) -> tuple[int, float]:
451 """Return the songs and minutes of music playing before the first window track."""
452 behind = []
453 for item in items:
454 if item.queue_item_id == window_start_id:
455 break
456 if not item.extra_attributes.get(ATTR_QUEUE_DJ):
457 behind.append(item)
458 minutes = sum(item.duration or FALLBACK_TRACK_SECONDS for item in behind) / 60.0
459 return len(behind), minutes
460
461 def _drop_unaired_dj_history(
462 self, state: DJQueueState, items: list[QueueItem], guard_index: int
463 ) -> None:
464 """Forget the guard history of breaks that were planned but never reached the player."""
465 # events strictly behind the window start have aired; one exactly on it is ambiguous,
466 # since its clip may be owned by the player (airing or buffered) or still one slot
467 # beyond the guard. a clip surviving in the owned head is what tells the two apart
468 owns_clip = any(
469 item.extra_attributes.get(ATTR_QUEUE_DJ) for item in items[: guard_index + 1]
470 )
471 boundary = state.songs_before_window if owns_clip else state.songs_before_window - 1
472 state.history = {
473 section_id: [(song, minute) for song, minute in events if song <= boundary]
474 for section_id, events in state.history.items()
475 }
476
477 def _rebase_dj_history(
478 self, state: DJQueueState, songs_before_window: int, minutes_before_window: float
479 ) -> None:
480 """Re-anchor the guard history onto the given start of the planning window."""
481 # events behind the window move with it so they keep their distance, while the ones
482 # ahead are capped at the window start: their old position no longer means anything
483 song_delta = min(songs_before_window - state.songs_before_window, 0)
484 minute_delta = min(minutes_before_window - state.minutes_before_window, 0.0)
485 state.history = {
486 section_id: [
487 (
488 min(song + song_delta, songs_before_window),
489 min(minute + minute_delta, minutes_before_window),
490 )
491 for song, minute in events
492 ]
493 for section_id, events in state.history.items()
494 }
495
496 def _dj_queue_items(self, queue_id: str) -> list[QueueItem]:
497 """Return all items of a queue."""
498 items: list[QueueItem] = []
499 offset = 0
500 while True:
501 page = self.mass.player_queues.items(queue_id, limit=QUEUE_PAGE_SIZE, offset=offset)
502 items.extend(page)
503 if len(page) < QUEUE_PAGE_SIZE:
504 return items
505 offset += QUEUE_PAGE_SIZE
506
507 def _queue_item_to_track(self, index: int, item: QueueItem) -> dict[str, Any]:
508 """Convert a music queue item into the track dict the planner consumes."""
509 name = ""
510 artist = ""
511 if (media_item := item.media_item) is not None:
512 name = str(media_item.name or "")
513 artists = getattr(media_item, "artists", None)
514 if artists:
515 artist = str(artists[0].name)
516 if not name:
517 raw_name = str(item.name or "")
518 artist, separator, name = raw_name.partition(" - ")
519 if not separator:
520 artist, name = "", raw_name
521 return {
522 "index": index,
523 "item_id": item.queue_item_id,
524 "name": name,
525 "artist": artist,
526 "songinfo": f"{artist} - {name}".strip(" -"),
527 "duration": item.duration,
528 "media_item": None,
529 }
530