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