add cleanup flag for commands

discord_bot/discord_bot.go

@@ -18,6 +18,7 @@ type botImpl struct {
 	imagineQueue       imagine_queue.Queue
 	registeredCommands []*discordgo.ApplicationCommand
 	imagineCommand     string
+	removeCommands     bool
 }
 
 type Config struct {
@@ -26,6 +27,7 @@ type Config struct {
 	GuildID         string
 	ImagineQueue    imagine_queue.Queue
 	ImagineCommand  string
+	RemoveCommands  bool
 }
 
 func (b *botImpl) imagineCommandString() string {
@@ -80,6 +82,7 @@ func New(cfg Config) (Bot, error) {
 		imagineQueue:       cfg.ImagineQueue,
 		registeredCommands: make([]*discordgo.ApplicationCommand, 0),
 		imagineCommand:     cfg.ImagineCommand,
+		removeCommands:     cfg.RemoveCommands,
 	}
 
 	err = bot.addImagineCommand()
@@ -176,12 +179,18 @@ func (b *botImpl) Start() {
 
 func (b *botImpl) teardown() error {
 	// Delete all commands added by the bot
-	// for _, v := range b.registeredCommands {
-	//	err := b.botSession.ApplicationCommandDelete(b.botSession.State.User.ID, b.guildID, v.ID)
-	//	if err != nil {
-	//		log.Panicf("Cannot delete '%v' command: %v", v.Name, err)
-	//	}
-	// }
+	if b.removeCommands {
+		log.Printf("Removing all commands added by bot...")
+
+		for _, v := range b.registeredCommands {
+			log.Printf("Removing command '%v'...", v.Name)
+
+			err := b.botSession.ApplicationCommandDelete(b.botSession.State.User.ID, b.guildID, v.ID)
+			if err != nil {
+				log.Panicf("Cannot delete '%v' command: %v", v.Name, err)
+			}
+		}
+	}
 
 	return b.botSession.Close()
 }

main.go

@@ -14,11 +14,12 @@ import (
 
 // Bot parameters
 var (
-	guildID        = flag.String("guild", "", "Guild ID. If not passed - bot registers commands globally")
-	botToken       = flag.String("token", "", "Bot access token")
-	apiHost        = flag.String("host", "", "Host for the Automatic1111 API")
-	imagineCommand = flag.String("imagine", "imagine", "Imagine command name. Default is \"imagine\"")
-	devModeFlag    = flag.Bool("dev", false, "Start in development mode, using \"dev_\" prefixed commands instead")
+	guildID            = flag.String("guild", "", "Guild ID. If not passed - bot registers commands globally")
+	botToken           = flag.String("token", "", "Bot access token")
+	apiHost            = flag.String("host", "", "Host for the Automatic1111 API")
+	imagineCommand     = flag.String("imagine", "imagine", "Imagine command name. Default is \"imagine\"")
+	removeCommandsFlag = flag.Bool("remove", false, "Delete all commands when bot exits")
+	devModeFlag        = flag.Bool("dev", false, "Start in development mode, using \"dev_\" prefixed commands instead")
 )
 
 func main() {
@@ -48,6 +49,12 @@ func main() {
 		log.Printf("Starting in development mode.. all commands prefixed with \"dev_\"")
 	}
 
+	removeCommands := false
+
+	if removeCommandsFlag != nil && *removeCommandsFlag {
+		removeCommands = *removeCommandsFlag
+	}
+
 	stableDiffusionAPI, err := stable_diffusion_api.New(stable_diffusion_api.Config{
 		Host: *apiHost,
 	})
@@ -87,6 +94,7 @@ func main() {
 		GuildID:         *guildID,
 		ImagineQueue:    imagineQueue,
 		ImagineCommand:  *imagineCommand,
+		RemoveCommands:  removeCommands,
 	})
 	if err != nil {
 		log.Fatalf("Error creating Discord bot: %v", err)