cleanup app, config and root

1 files changed, 187 insertions, 66 deletions

diff --git a/cmd/root.go b/cmd/root.go
index d846a14c2..092606de7 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -2,9 +2,10 @@ package cmd
 
 import (
 	"context"
-	"log/slog"
+	"fmt"
 	"os"
 	"sync"
+	"time"
 
 	tea "github.com/charmbracelet/bubbletea"
 	"github.com/kujtimiihoxha/termai/internal/app"
@@ -13,6 +14,7 @@ import (
 	"github.com/kujtimiihoxha/termai/internal/db"
 	"github.com/kujtimiihoxha/termai/internal/llm/agent"
 	"github.com/kujtimiihoxha/termai/internal/logging"
+	"github.com/kujtimiihoxha/termai/internal/pubsub"
 	"github.com/kujtimiihoxha/termai/internal/tui"
 	zone "github.com/lrstanley/bubblezone"
 	"github.com/spf13/cobra"
@@ -23,111 +25,229 @@ var rootCmd = &cobra.Command{
 	Short: "A terminal ai assistant",
 	Long:  `A terminal ai assistant`,
 	RunE: func(cmd *cobra.Command, args []string) error {
+		// If the help flag is set, show the help message
 		if cmd.Flag("help").Changed {
 			cmd.Help()
 			return nil
 		}
+
+		// Load the config
 		debug, _ := cmd.Flags().GetBool("debug")
-		err := config.Load(debug)
+		cwd, _ := cmd.Flags().GetString("cwd")
+		if cwd != "" {
+			err := os.Chdir(cwd)
+			if err != nil {
+				return fmt.Errorf("failed to change directory: %v", err)
+			}
+		}
+		if cwd == "" {
+			c, err := os.Getwd()
+			if err != nil {
+				return fmt.Errorf("failed to get current working directory: %v", err)
+			}
+			cwd = c
+		}
+		_, err := config.Load(cwd, debug)
 		if err != nil {
 			return err
 		}
-		cfg := config.Get()
-		defaultLevel := slog.LevelInfo
-		if cfg.Debug {
-			defaultLevel = slog.LevelDebug
-		}
-		logger := slog.New(slog.NewTextHandler(logging.NewWriter(), &slog.HandlerOptions{
-			Level: defaultLevel,
-		}))
-		slog.SetDefault(logger)
 
 		err = assets.WriteAssets()
 		if err != nil {
-			return err
+			logging.Error("Error writing assets: %v", err)
 		}
 
+		// Connect DB, this will also run migrations
 		conn, err := db.Connect()
 		if err != nil {
 			return err
 		}
-		ctx := context.Background()
+
+		// Create main context for the application
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
 
 		app := app.New(ctx, conn)
-		logging.Info("Starting termai...")
+
+		// Set up the TUI
 		zone.NewGlobal()
-		tui := tea.NewProgram(
+		program := tea.NewProgram(
 			tui.New(app),
 			tea.WithAltScreen(),
 			tea.WithMouseCellMotion(),
 		)
-		logging.Info("Setting up subscriptions...")
-		ch, unsub := setupSubscriptions(app)
-		defer unsub()
 
+		// Initialize MCP tools in the background
+		initMCPTools(ctx, app)
+
+		// Setup the subscriptions, this will send services events to the TUI
+		ch, cancelSubs := setupSubscriptions(app)
+
+		// Create a context for the TUI message handler
+		tuiCtx, tuiCancel := context.WithCancel(ctx)
+		var tuiWg sync.WaitGroup
+		tuiWg.Add(1)
+
+		// Set up message handling for the TUI
 		go func() {
-			// Set this up once
-			agent.GetMcpTools(ctx, app.Permissions)
-			for msg := range ch {
-				tui.Send(msg)
+			defer tuiWg.Done()
+			defer func() {
+				if r := recover(); r != nil {
+					logging.Error("Panic in TUI message handling: %v", r)
+					attemptTUIRecovery(program)
+				}
+			}()
+
+			for {
+				select {
+				case <-tuiCtx.Done():
+					logging.Info("TUI message handler shutting down")
+					return
+				case msg, ok := <-ch:
+					if !ok {
+						logging.Info("TUI message channel closed")
+						return
+					}
+					program.Send(msg)
+				}
 			}
 		}()
-		if _, err := tui.Run(); err != nil {
-			return err
+
+		// Cleanup function for when the program exits
+		cleanup := func() {
+			// Shutdown the app
+			app.Shutdown()
+
+			// Cancel subscriptions first
+			cancelSubs()
+
+			// Then cancel TUI message handler
+			tuiCancel()
+
+			// Wait for TUI message handler to finish
+			tuiWg.Wait()
+
+			logging.Info("All goroutines cleaned up")
+		}
+
+		// Run the TUI
+		result, err := program.Run()
+		cleanup()
+
+		if err != nil {
+			logging.Error("TUI error: %v", err)
+			return fmt.Errorf("TUI error: %v", err)
 		}
+
+		logging.Info("TUI exited with result: %v", result)
 		return nil
 	},
 }
 
-func setupSubscriptions(app *app.App) (chan tea.Msg, func()) {
-	ch := make(chan tea.Msg)
-	wg := sync.WaitGroup{}
-	ctx, cancel := context.WithCancel(app.Context)
-	{
-		sub := logging.Subscribe(ctx)
-		wg.Add(1)
-		go func() {
-			for ev := range sub {
-				ch <- ev
+// attemptTUIRecovery tries to recover the TUI after a panic
+func attemptTUIRecovery(program *tea.Program) {
+	logging.Info("Attempting to recover TUI after panic")
+
+	// We could try to restart the TUI or gracefully exit
+	// For now, we'll just quit the program to avoid further issues
+	program.Quit()
+}
+
+func initMCPTools(ctx context.Context, app *app.App) {
+	go func() {
+		defer func() {
+			if r := recover(); r != nil {
+				logging.Error("Panic in MCP goroutine: %v", r)
 			}
-			wg.Done()
 		}()
-	}
-	{
-		sub := app.Sessions.Subscribe(ctx)
-		wg.Add(1)
-		go func() {
-			for ev := range sub {
-				ch <- ev
+
+		// Create a context with timeout for the initial MCP tools fetch
+		ctxWithTimeout, cancel := context.WithTimeout(ctx, 30*time.Second)
+		defer cancel()
+
+		// Set this up once with proper error handling
+		agent.GetMcpTools(ctxWithTimeout, app.Permissions)
+		logging.Info("MCP message handling goroutine exiting")
+	}()
+}
+
+func setupSubscriber[T any](
+	ctx context.Context,
+	wg *sync.WaitGroup,
+	name string,
+	subscriber func(context.Context) <-chan pubsub.Event[T],
+	outputCh chan<- tea.Msg,
+) {
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		defer func() {
+			if r := recover(); r != nil {
+				logging.Error("Panic in %s subscription goroutine: %v", name, r)
 			}
-			wg.Done()
 		}()
-	}
-	{
-		sub := app.Messages.Subscribe(ctx)
-		wg.Add(1)
-		go func() {
-			for ev := range sub {
-				ch <- ev
+
+		for {
+			select {
+			case event, ok := <-subscriber(ctx):
+				if !ok {
+					logging.Info("%s subscription channel closed", name)
+					return
+				}
+
+				// Convert generic event to tea.Msg if needed
+				var msg tea.Msg = event
+
+				// Non-blocking send with timeout to prevent deadlocks
+				select {
+				case outputCh <- msg:
+				case <-time.After(500 * time.Millisecond):
+					logging.Warn("%s message dropped due to slow consumer", name)
+				case <-ctx.Done():
+					logging.Info("%s subscription cancelled", name)
+					return
+				}
+			case <-ctx.Done():
+				logging.Info("%s subscription cancelled", name)
+				return
 			}
-			wg.Done()
-		}()
-	}
-	{
-		sub := app.Permissions.Subscribe(ctx)
-		wg.Add(1)
+		}
+	}()
+}
+
+func setupSubscriptions(app *app.App) (chan tea.Msg, func()) {
+	ch := make(chan tea.Msg, 100)
+	// Add a buffer to prevent blocking
+	wg := sync.WaitGroup{}
+	ctx, cancel := context.WithCancel(context.Background())
+	// Setup each subscription using the helper
+	setupSubscriber(ctx, &wg, "logging", logging.Subscribe, ch)
+	setupSubscriber(ctx, &wg, "sessions", app.Sessions.Subscribe, ch)
+	setupSubscriber(ctx, &wg, "messages", app.Messages.Subscribe, ch)
+	setupSubscriber(ctx, &wg, "permissions", app.Permissions.Subscribe, ch)
+
+	// Return channel and a cleanup function
+	cleanupFunc := func() {
+		logging.Info("Cancelling all subscriptions")
+		cancel() // Signal all goroutines to stop
+
+		// Wait with a timeout for all goroutines to complete
+		waitCh := make(chan struct{})
 		go func() {
-			for ev := range sub {
-				ch <- ev
-			}
-			wg.Done()
+			wg.Wait()
+			close(waitCh)
 		}()
+
+		select {
+		case <-waitCh:
+			logging.Info("All subscription goroutines completed successfully")
+		case <-time.After(5 * time.Second):
+			logging.Warn("Timed out waiting for some subscription goroutines to complete")
+		}
+
+		close(ch) // Safe to close after all writers are done or timed out
 	}
-	return ch, func() {
-		cancel()
-		wg.Wait()
-		close(ch)
-	}
+	return ch, cleanupFunc
 }
 
 func Execute() {
@@ -139,5 +259,6 @@ func Execute() {
 
 func init() {
 	rootCmd.Flags().BoolP("help", "h", false, "Help")
-	rootCmd.Flags().BoolP("debug", "d", false, "Help")
+	rootCmd.Flags().BoolP("debug", "d", false, "Debug")
+	rootCmd.Flags().StringP("cwd", "c", "", "Current working directory")
 }