""" YTAudio - YouTube audio for Red-DiscordBot via yt-dlp -> mp3 -> Lavalink (local source). This cog OWNS the shared Red-Lavalink client connection: * On startup it calls ``lavalink.initialize()`` (once the bot is ready), connecting the in-process Lavalink websocket client to the standalone Lavalink node running on localhost:2333 (managed by systemd). * On unload it calls ``lavalink.close()``. Because the Red-Lavalink client is a process-global singleton, exactly ONE cog should own initialize()/close(). Other cogs (e.g. ttstoy) simply call ``lavalink.get_player(guild_id)`` and share this connection. Playback path: yt-dlp downloads bestaudio -> ffmpeg transcodes to mp3 -> the mp3 file is loaded through Lavalink's ``local`` source and played. The mp3 is deleted when the track finishes / is skipped / on stop. """ import asyncio import functools import logging import os import shutil import time as _time from typing import Dict, List, Optional import discord import lavalink import yt_dlp from redbot.core import Config, commands from redbot.core.bot import Red from redbot.core.data_manager import cog_data_path log = logging.getLogger("red.scrapyard.ytaudio") # --- Lavalink node connection (must match the standalone application.yml) --- LL_HOST = "localhost" LL_PORT = 2333 LL_PASSWORD = "youshallnotpass" LL_SECURED = False MP3_BITRATE = "192" YTDLP_BASE_OPTS = { "format": "bestaudio/best", "quiet": True, "no_warnings": True, "noplaylist": True, "ignoreerrors": False, "default_search": "ytsearch", "source_address": "0.0.0.0", "postprocessors": [ { "key": "FFmpegExtractAudio", "preferredcodec": "mp3", "preferredquality": MP3_BITRATE, } ], } class QueuedTrack: """A track queued for a guild, resolved lazily at play time.""" __slots__ = ("query", "requester", "channel", "title", "mp3_path", "status_msg") def __init__(self, query: str, requester: discord.abc.User, channel): self.query = query self.requester = requester self.channel = channel self.title: Optional[str] = None self.mp3_path: Optional[str] = None self.status_msg = None class GuildPlayer: """Per-guild queue + player-loop state. One per guild. The queue model is explicit and always inspectable: * ``pending`` is the list of tracks waiting to play (index 0 is next). * ``current`` is the track playing right now (or None). Commands read these directly, so nowplaying/queue always reflect reality. """ def __init__(self, cog: "YTAudio", guild_id: int): self.cog = cog self.guild_id = guild_id self.pending: List[QueuedTrack] = [] self.current: Optional[QueuedTrack] = None self.task: Optional[asyncio.Task] = None self.next_event = asyncio.Event() # fired when the current track ends self.wakeup = asyncio.Event() # fired when a track is added self.volume = 100 def start(self): if self.task is None or self.task.done(): self.task = asyncio.create_task(self._player_loop()) def enqueue(self, track: "QueuedTrack"): self.pending.append(track) self.wakeup.set() def clear_pending(self): cleared = self.pending self.pending = [] for t in cleared: self.cog._cleanup_file(t) async def _player_loop(self): while True: if not self.pending: self.wakeup.clear() if not self.pending: try: await self.wakeup.wait() except asyncio.CancelledError: return continue track = self.pending.pop(0) self.current = track self.next_event.clear() try: await self.cog._play_track(self.guild_id, track, self.next_event) await self.next_event.wait() except asyncio.CancelledError: self.cog._cleanup_file(track) self.current = None return except Exception: log.exception("Error playing track in guild %s", self.guild_id) finally: self.cog._cleanup_file(track) self.current = None class YTAudio(commands.Cog): """YouTube audio via yt-dlp -> mp3 -> Lavalink.""" def __init__(self, bot: Red): self.bot = bot self.config = Config.get_conf(self, identifier=0x59544155_44494F, force_registration=True) self.config.register_guild(volume=100) self.players: Dict[int, GuildPlayer] = {} self._ll_ready = False self._connect_task: Optional[asyncio.Task] = None self.mp3_dir = cog_data_path(self) / "mp3cache" self.mp3_dir.mkdir(parents=True, exist_ok=True) self._ytdlp_opts = dict(YTDLP_BASE_OPTS) self._ytdlp_opts["outtmpl"] = str(self.mp3_dir / "%(id)s.%(ext)s") # ------------------------------------------------------------------ # Lavalink client ownership # ------------------------------------------------------------------ def start_up_task(self): """Called from setup() after add_cog.""" self._connect_task = asyncio.create_task(self._connect_lavalink()) async def _connect_lavalink(self): await self.bot.wait_until_red_ready() # If a node is already connected (e.g. built-in Audio still loaded), # don't double-initialize, just share it. if lavalink.get_all_nodes(): self._ll_ready = True log.info("YTAudio: a Lavalink node is already connected; sharing it.") return max_retries = 5 for attempt in range(1, max_retries + 1): try: await lavalink.initialize( bot=self.bot, host=LL_HOST, password=LL_PASSWORD, port=LL_PORT, timeout=30, secured=LL_SECURED, ) self._ll_ready = True lavalink.register_event_listener(self._track_event_handler) log.info("YTAudio: connected Lavalink client to %s:%s", LL_HOST, LL_PORT) return except lavalink.AbortingNodeConnection: await self._safe_close() log.warning("YTAudio: Lavalink connection aborted.") return except asyncio.TimeoutError: await self._safe_close() log.warning( "YTAudio: Lavalink connect timed out (attempt %s/%s), retrying...", attempt, max_retries, ) await asyncio.sleep(2) except Exception: log.exception("YTAudio: unexpected error connecting to Lavalink.") await self._safe_close() return log.critical("YTAudio: failed to connect to Lavalink node after %s attempts.", max_retries) async def _safe_close(self): try: await lavalink.close(self.bot) except Exception: pass # ------------------------------------------------------------------ # Track end handling # ------------------------------------------------------------------ async def _track_event_handler(self, player, event, extra): if event in ( lavalink.LavalinkEvents.TRACK_END, lavalink.LavalinkEvents.TRACK_EXCEPTION, lavalink.LavalinkEvents.TRACK_STUCK, ): gp = self.players.get(player.guild.id) if gp is not None: gp.next_event.set() # ------------------------------------------------------------------ # yt-dlp resolution (blocking -> executor) with live progress # ------------------------------------------------------------------ def _make_progress_hook(self, loop, track: QueuedTrack): # Runs inside the executor thread. Schedules throttled message edits # back onto the event loop. Discord rate-limits edits, so we edit at # most once per second and only when the whole-number percent changes. state = {"last_edit": 0.0, "last_pct": -1} def hook(d): if track.status_msg is None: return status = d.get("status") if status == "downloading": total = d.get("total_bytes") or d.get("total_bytes_estimate") downloaded = d.get("downloaded_bytes") or 0 if not total: return pct = int(downloaded * 100 / total) now = _time.monotonic() if pct == state["last_pct"] or (now - state["last_edit"]) < 1.0: if pct < 100: return state["last_pct"] = pct state["last_edit"] = now content = "Downloading: {} ({}%)".format(track.query, pct) loop.call_soon_threadsafe( lambda c=content: asyncio.ensure_future(self._safe_edit(track, c)) ) elif status == "finished": content = "Converting to mp3: {}".format(track.query) loop.call_soon_threadsafe( lambda c=content: asyncio.ensure_future(self._safe_edit(track, c)) ) return hook async def _safe_edit(self, track: QueuedTrack, content: str): if track.status_msg is None: return try: await track.status_msg.edit(content=content) except Exception: pass def _blocking_extract(self, query: str, progress_hook=None) -> dict: opts = dict(self._ytdlp_opts) if progress_hook is not None: opts["progress_hooks"] = [progress_hook] with yt_dlp.YoutubeDL(opts) as ydl: info = ydl.extract_info(query, download=True) if "entries" in info: # search/playlist -> first entry info = info["entries"][0] return info async def _resolve_to_mp3(self, track: QueuedTrack): loop = asyncio.get_running_loop() hook = self._make_progress_hook(loop, track) info = await loop.run_in_executor( None, functools.partial(self._blocking_extract, track.query, hook) ) track.title = info.get("title", track.query) vid = info.get("id") mp3_path = self.mp3_dir / f"{vid}.mp3" track.mp3_path = str(mp3_path) return track def _cleanup_file(self, track: Optional[QueuedTrack]): if track and track.mp3_path and os.path.isfile(track.mp3_path): try: os.remove(track.mp3_path) except OSError: log.warning("Could not delete mp3: %s", track.mp3_path) # ------------------------------------------------------------------ # Playback # ------------------------------------------------------------------ async def _play_track(self, guild_id: int, track: QueuedTrack, done_event: asyncio.Event): status_msg = None if track.channel is not None: try: status_msg = await track.channel.send( "Downloading: {} (0%)".format(track.query) ) except Exception: status_msg = None track.status_msg = status_msg await self._resolve_to_mp3(track) player = lavalink.get_player(guild_id) gp = self.players.get(guild_id) if gp: await player.set_volume(gp.volume) loaded = await player.load_tracks(track.mp3_path) if loaded.has_error or loaded.load_type != lavalink.enums.LoadType.TRACK_LOADED: log.error("YTAudio: failed to load mp3 into Lavalink: %s", track.mp3_path) if status_msg is not None: try: await status_msg.edit(content="Failed to load: {}".format(track.query)) except Exception: pass done_event.set() return player.add(requester=track.requester, track=loaded.tracks[0]) if not player.is_playing: await player.play() if status_msg is not None: try: await status_msg.edit(content="Now playing: {}".format(track.title or track.query)) except Exception: pass def _get_gp(self, guild_id: int) -> GuildPlayer: gp = self.players.get(guild_id) if gp is None: gp = GuildPlayer(self, guild_id) self.players[guild_id] = gp gp.start() return gp async def _ensure_connected(self, ctx: commands.Context) -> bool: if not self._ll_ready and not lavalink.get_all_nodes(): await ctx.send("Lavalink is not connected yet. Try again in a moment.") return False if ctx.author.voice is None or ctx.author.voice.channel is None: await ctx.send("You need to be in a voice channel.") return False target = ctx.author.voice.channel me_voice = ctx.guild.me.voice if me_voice is None or me_voice.channel is None: await lavalink.connect(target) elif me_voice.channel.id != target.id: # Already connected to a different channel: move instead of reconnecting. try: player = lavalink.get_player(ctx.guild.id) await player.move_to(target) except Exception: await lavalink.connect(target) return True # ------------------------------------------------------------------ # Commands # ------------------------------------------------------------------ @commands.command() @commands.guild_only() async def play(self, ctx: commands.Context, *, query: str): """Play a YouTube URL or search query.""" if not await self._ensure_connected(ctx): return gp = self._get_gp(ctx.guild.id) gp.volume = await self.config.guild(ctx.guild).volume() track = QueuedTrack(query, ctx.author, ctx.channel) already_active = gp.current is not None or bool(gp.pending) gp.enqueue(track) if already_active: position = len(gp.pending) await ctx.send("Added to queue (position {}): {}".format(position, query)) @commands.command() @commands.guild_only() async def skip(self, ctx: commands.Context): """Skip the current track.""" gp = self.players.get(ctx.guild.id) if gp is None or gp.current is None: return await ctx.send("Nothing is playing.") try: player = lavalink.get_player(ctx.guild.id) except Exception: return await ctx.send("Nothing is playing.") # Stopping fires TRACK_END, which the event handler uses to advance # the queue. Do not set next_event here or it would skip two tracks. await player.stop() await ctx.send("Skipped.") @commands.command() @commands.guild_only() async def stop(self, ctx: commands.Context): """Stop playback and clear the queue.""" gp = self.players.get(ctx.guild.id) if gp: gp.clear_pending() gp.next_event.set() try: player = lavalink.get_player(ctx.guild.id) await player.stop() except Exception: pass await ctx.send("Stopped and cleared the queue.") @commands.command() @commands.guild_only() async def pause(self, ctx: commands.Context): """Pause playback.""" try: player = lavalink.get_player(ctx.guild.id) except Exception: return await ctx.send("Nothing is playing.") await player.pause(True) await ctx.send("Paused.") @commands.command() @commands.guild_only() async def resume(self, ctx: commands.Context): """Resume playback.""" try: player = lavalink.get_player(ctx.guild.id) except Exception: return await ctx.send("Nothing is playing.") await player.pause(False) await ctx.send("Resumed.") @commands.command(name="queue", aliases=["q"]) @commands.guild_only() async def queue_cmd(self, ctx: commands.Context): """Show the current track and the queue.""" gp = self.players.get(ctx.guild.id) if not gp or (gp.current is None and not gp.pending): return await ctx.send("The queue is empty.") lines = [] if gp.current: lines.append("Now playing: {}".format(gp.current.title or gp.current.query)) if gp.pending: lines.append("Up next:") for i, t in enumerate(gp.pending, 1): lines.append("{}. {}".format(i, t.title or t.query)) else: lines.append("Queue is empty.") # Keep within Discord message limits. out = "\n".join(lines) if len(out) > 1900: out = out[:1900] + "\n..." await ctx.send(out) @commands.command(name="nowplaying", aliases=["np"]) @commands.guild_only() async def nowplaying(self, ctx: commands.Context): """Show the current track.""" gp = self.players.get(ctx.guild.id) if not gp or gp.current is None: return await ctx.send("Nothing is playing.") await ctx.send("Now playing: {}".format(gp.current.title or gp.current.query)) @commands.command() @commands.guild_only() async def volume(self, ctx: commands.Context, level: int): """Set volume (0-150).""" level = max(0, min(150, level)) await self.config.guild(ctx.guild).volume.set(level) gp = self.players.get(ctx.guild.id) if gp: gp.volume = level try: player = lavalink.get_player(ctx.guild.id) await player.set_volume(level) except Exception: pass await ctx.send("Volume set to {}.".format(level)) @commands.command(name="summon", aliases=["ytjoin"]) @commands.guild_only() async def join(self, ctx: commands.Context): """Summon the bot to your voice channel.""" if await self._ensure_connected(ctx): await ctx.send("Joined your voice channel.") @commands.command(name="ytleave", aliases=["ytdc"]) @commands.guild_only() async def leave(self, ctx: commands.Context): """Leave the voice channel and clear state.""" gp = self.players.get(ctx.guild.id) if gp: gp.clear_pending() gp.next_event.set() try: player = lavalink.get_player(ctx.guild.id) await player.stop() await player.disconnect() except Exception: pass await ctx.send("Left the channel.") # ------------------------------------------------------------------ # Cleanup # ------------------------------------------------------------------ async def cog_unload(self): if self._connect_task: self._connect_task.cancel() for gp in self.players.values(): if gp.task: gp.task.cancel() try: lavalink.unregister_event_listener(self._track_event_handler) except Exception: pass # This cog owns the connection: close it on unload. await self._safe_close() # Sweep leftover mp3s. try: shutil.rmtree(self.mp3_dir, ignore_errors=True) self.mp3_dir.mkdir(parents=True, exist_ok=True) except Exception: pass