2 files changed, 12 insertions, 25 deletions

diff --git a/pkg/bot/bot.go b/pkg/bot/bot.go
index 18001f9..abee7bc 100644
--- a/pkg/bot/bot.go
+++ b/pkg/bot/bot.go
@@ -1,8 +1,6 @@
 package bot
 
 import (
-	"context"
-	"errors"
 	"jinx/pkg/discord"
 	"jinx/pkg/discord/events"
 
@@ -10,30 +8,18 @@ import (
 )
 
 type Bot struct {
-	client        *discord.Discord
-	logger        *zerolog.Logger
-	cancelContext context.CancelFunc
+	client *discord.Discord
+	logger *zerolog.Logger
 }
 
 func NewBot(token string, logger *zerolog.Logger) *Bot {
 	return &Bot{
-		client:        discord.New(token, logger),
-		logger:        logger,
-		cancelContext: nil,
+		client: discord.New(token, logger),
+		logger: logger,
 	}
 }
 
 func (b *Bot) Start() error {
-	var err error
-	ctx, cancel := context.WithCancel(context.Background())
-	b.cancelContext = cancel
-
-	defer func() {
-		if err != nil {
-			cancel()
-		}
-	}()
-
 	b.client.On(events.READY, func(_ any) {
 		b.logger.Info().Msg("ready!")
 	})
@@ -49,7 +35,7 @@ func (b *Bot) Start() error {
 		}
 	})
 
-	if err := b.client.Connect(ctx); err != nil {
+	if err := b.client.Connect(); err != nil {
 		return err
 	}
 
@@ -57,10 +43,5 @@ func (b *Bot) Start() error {
 }
 
 func (b *Bot) Stop() error {
-	if b.cancelContext == nil {
-		return errors.New("bot has no context")
-	}
-
-	b.cancelContext()
 	return b.client.Disconnect()
 }
diff --git a/pkg/discord/discord.go b/pkg/discord/discord.go
index b0b6f7d..49ddc3b 100644
--- a/pkg/discord/discord.go
+++ b/pkg/discord/discord.go
@@ -16,6 +16,8 @@ type Discord struct {
 	gateway      gateway.Gateway
 	eventHandler events.EventHandler
 	rest         rest.REST
+
+	cancel context.CancelFunc
 }
 
 func New(token string, logger *zerolog.Logger) *Discord {
@@ -34,7 +36,10 @@ func New(token string, logger *zerolog.Logger) *Discord {
 	}
 }
 
-func (d *Discord) Connect(ctx context.Context) error {
+func (d *Discord) Connect() error {
+	ctx, cancel := context.WithCancel(context.Background())
+	d.cancel = cancel
+
 	gatewayURL, err := d.rest.Gateway()
 	if err != nil {
 		return err
@@ -52,6 +57,7 @@ func (d *Discord) Connect(ctx context.Context) error {
 }
 
 func (d *Discord) Disconnect() error {
+	d.cancel()
 	return d.gateway.Close()
 }