Extra lavalink listeners/optional track logging

2 files changed, 78 insertions, 10 deletions

diff --git a/code/cogs/music.py b/code/cogs/music.py
index f3e98b4..4b3fdbb 100644
--- a/code/cogs/music.py
+++ b/code/cogs/music.py
@@ -1,14 +1,15 @@
 import discord
 from discord.ext import commands
-from discord import app_commands
 import lavalink
 from lavalink import errors
+import os
 
 from utils.config import (
     LAVALINK_HOST,
     LAVALINK_PASSWORD,
     LAVALINK_PORT,
     LOG,
+    LOG_SONGS,
 )
 from utils.command_tree import CheckPlayerError
 from utils.ai_recommendations import add_song_recommendations
@@ -95,6 +96,7 @@ class LavalinkVoiceClient(discord.VoiceProtocol):
 class Music(commands.Cog):
     def __init__(self, bot):
         self.bot = bot
+        self.log_file = "track_events.log"
 
     async def cog_load(self):
         if not hasattr(
@@ -117,11 +119,16 @@ class Music(commands.Cog):
                 )
             else:
                 await node.connect()
-                LOG.info(f"Connected to lavalink node {node.name}")
 
         self.lavalink: lavalink.Client = self.bot.lavalink
         self.lavalink.add_event_hooks(self)
 
+        if os.path.exists("/.dockerenv"):
+            self.log_file = "/config/track_events.log"
+
+        if LOG_SONGS:
+            LOG.info(f"Logging track events to {self.log_file}")
+
     def cog_unload(self):
         """Cog unload handler. This removes any event hooks that were registered."""
         self.lavalink._event_hooks.clear()
@@ -235,6 +242,54 @@ class Music(commands.Cog):
                 self.bot.openai, self.bot.user, event.player, 5, inputs
             )
 
+    @lavalink.listener(lavalink.events.NodeConnectedEvent)
+    async def node_connected(self, event: lavalink.events.NodeConnectedEvent):
+        LOG.info(f"Lavalink node {event.node.name} has connected")
+
+    @lavalink.listener(lavalink.events.NodeReadyEvent)
+    async def node_ready(self, event: lavalink.events.NodeReadyEvent):
+        LOG.info(f"Lavalink node {event.node.name} is ready")
+
+    @lavalink.listener(lavalink.events.NodeDisconnectedEvent)
+    async def node_disconnected(
+        self, event: lavalink.events.NodeDisconnectedEvent
+    ):
+        LOG.error(f"Lavalink node {event.node.name} has disconnected")
+
+    # If we get a track load failed event (like LoadError, but for some reason that
+    # wasn't the eception raised), skip the track
+    @lavalink.listener(lavalink.events.TrackLoadFailedEvent)
+    async def track_load_failed(self, event: lavalink.events.TrackEndEvent):
+        await event.player.skip()
+
+    # Only log track events if enabled
+    if LOG_SONGS:
+
+        @lavalink.listener(lavalink.events.TrackStartEvent)
+        async def track_start(self, event: lavalink.events.TrackStartEvent):
+            with open(self.log_file, "a") as f:
+                f.write(
+                    f"STARTED: {event.track.title} by {event.track.author}\n"
+                )
+
+        @lavalink.listener(lavalink.events.TrackStuckEvent)
+        async def track_stuck(self, event: lavalink.events.TrackStuckEvent):
+            with open(self.log_file, "a") as f:
+                f.write(
+                    f"STUCK: {event.track.title} by {event.track.author} -"
+                    f" {event.track.uri}\n"
+                )
+
+        @lavalink.listener(lavalink.events.TrackExceptionEvent)
+        async def track_exception(
+            self, event: lavalink.events.TrackExceptionEvent
+        ):
+            with open(self.log_file, "a") as f:
+                f.write(
+                    f"EXCEPTION{event.track.title} by {event.track.author} -"
+                    f" {event.track.uri}\n"
+                )
+
 
 async def setup(bot):
     await bot.add_cog(Music(bot))
diff --git a/code/utils/config.py b/code/utils/config.py
index 470d03c..f2e0d01 100644
--- a/code/utils/config.py
+++ b/code/utils/config.py
@@ -33,7 +33,8 @@ BOT_COLOR = None
 BOT_INVITE_LINK = None
 FEEDBACK_CHANNEL_ID = None
 BUG_CHANNEL_ID = None
-YOUTUBE_SUPPORT = None
+LOG_SONGS = False
+YOUTUBE_SUPPORT = False
 SPOTIFY_CLIENT_ID = None
 SPOTIFY_CLIENT_SECRET = None
 GENIUS_CLIENT_ID = None
@@ -54,10 +55,17 @@ schema = {
                 "bot_invite_link": {"type": "string"},
                 "feedback_channel_id": {"type": "integer"},
                 "bug_channel_id": {"type": "integer"},
-                "youtube_support": {"type": "boolean"},
+                "log_songs": {"type": "boolean"},
             },
             "required": ["token"],
         },
+        "youtube": {
+            "type": "object",
+            "properties": {
+                "enabled": {"type": "boolean"},
+            },
+            "required": ["enabled"],
+        },
         "spotify": {
             "type": "object",
             "properties": {
@@ -118,7 +126,10 @@ bot_info:
     bot_invite_link: ""
     feedback_channel_id: ""
     bug_channel_id: ""
-    youtube_support: false
+    log_songs: true
+
+youtube:
+    enabled: false
 
 spotify:
     spotify_client_id: ""
@@ -149,7 +160,7 @@ lavalink:
 
 # Thouroughly validate all of the options in the config.yaml file
 def validate_config(file_contents):
-    global TOKEN, BOT_COLOR, BOT_INVITE_LINK, FEEDBACK_CHANNEL_ID, BUG_CHANNEL_ID, YOUTUBE_SUPPORT, SPOTIFY_CLIENT_ID, SPOTIFY_CLIENT_SECRET, GENIUS_CLIENT_ID, GENIUS_CLIENT_SECRET, OPENAI_API_KEY, LAVALINK_HOST, LAVALINK_PORT, LAVALINK_PASSWORD
+    global TOKEN, BOT_COLOR, BOT_INVITE_LINK, FEEDBACK_CHANNEL_ID, BUG_CHANNEL_ID, LOG_SONGS, YOUTUBE_SUPPORT, SPOTIFY_CLIENT_ID, SPOTIFY_CLIENT_SECRET, GENIUS_CLIENT_ID, GENIUS_CLIENT_SECRET, OPENAI_API_KEY, LAVALINK_HOST, LAVALINK_PORT, LAVALINK_PASSWORD
     config = yaml.safe_load(file_contents)
 
     try:
@@ -209,10 +220,12 @@ def validate_config(file_contents):
             else:
                 BUG_CHANNEL_ID = config["bot_info"]["bug_channel_id"]
 
-    if "youtube_support" in config["bot_info"]:
-        YOUTUBE_SUPPORT = bool(config["bot_info"]["youtube_support"])
-    else:
-        YOUTUBE_SUPPORT = False
+    if "log_songs" in config["bot_info"]:
+        LOG_SONGS = bool(config["bot_info"]["log_songs"])
+
+    # Check for YouTube support
+    if "youtube" in config:
+        YOUTUBE_SUPPORT = bool(config["youtube"]["enabled"])
 
     #
     # If the SPOTIFY section is present, make sure the client ID and secret are valid