Merge pull request #27 from kujtimiihoxha/opencode

OpenCode - Initial Implementation

1 files changed, 194 insertions, 79 deletions

diff --git a/cmd/root.go b/cmd/root.go
index bdab53e14..8777acb82 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -2,126 +2,240 @@ package cmd
 
 import (
 	"context"
-	"log/slog"
+	"fmt"
 	"os"
 	"sync"
+	"time"
 
 	tea "github.com/charmbracelet/bubbletea"
-	"github.com/kujtimiihoxha/termai/internal/app"
-	"github.com/kujtimiihoxha/termai/internal/config"
-	"github.com/kujtimiihoxha/termai/internal/db"
-	"github.com/kujtimiihoxha/termai/internal/llm/agent"
-	"github.com/kujtimiihoxha/termai/internal/logging"
-	"github.com/kujtimiihoxha/termai/internal/tui"
+	"github.com/kujtimiihoxha/opencode/internal/app"
+	"github.com/kujtimiihoxha/opencode/internal/config"
+	"github.com/kujtimiihoxha/opencode/internal/db"
+	"github.com/kujtimiihoxha/opencode/internal/llm/agent"
+	"github.com/kujtimiihoxha/opencode/internal/logging"
+	"github.com/kujtimiihoxha/opencode/internal/pubsub"
+	"github.com/kujtimiihoxha/opencode/internal/tui"
 	zone "github.com/lrstanley/bubblezone"
 	"github.com/spf13/cobra"
 )
 
 var rootCmd = &cobra.Command{
-	Use:   "termai",
-	Short: "A terminal ai assistant",
-	Long:  `A terminal ai assistant`,
+	Use:   "OpenCode",
+	Short: "A terminal AI assistant for software development",
+	Long: `OpenCode is a powerful terminal-based AI assistant that helps with software development tasks.
+It provides an interactive chat interface with AI capabilities, code analysis, and LSP integration
+to assist developers in writing, debugging, and understanding code directly from the terminal.`,
 	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)
-		cfg := config.Get()
-		defaultLevel := slog.LevelInfo
-		if cfg.Debug {
-			defaultLevel = slog.LevelDebug
+		cwd, _ := cmd.Flags().GetString("cwd")
+		if cwd != "" {
+			err := os.Chdir(cwd)
+			if err != nil {
+				return fmt.Errorf("failed to change directory: %v", err)
+			}
 		}
-		logger := slog.New(slog.NewTextHandler(logging.NewWriter(), &slog.HandlerOptions{
-			Level: defaultLevel,
-		}))
-		slog.SetDefault(logger)
-
+		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
 		}
+
+		// Connect DB, this will also run migrations
 		conn, err := db.Connect()
 		if err != nil {
 			return err
 		}
-		ctx := context.Background()
 
-		app := app.New(ctx, conn)
-		logging.Info("Starting termai...")
+		// Create main context for the application
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		app, err := app.New(ctx, conn)
+		if err != nil {
+			logging.Error("Failed to create app: %v", err)
+			return err
+		}
+
+		// 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, ctx)
+
+		// 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 logging.RecoverPanic("TUI-message-handler", func() {
+				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
-			}
-			wg.Done()
-		}()
-	}
-	{
-		sub := app.Sessions.Subscribe(ctx)
-		wg.Add(1)
-		go func() {
-			for ev := range sub {
-				ch <- ev
-			}
-			wg.Done()
-		}()
-	}
-	{
-		sub := app.Messages.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 logging.RecoverPanic("MCP-goroutine", nil)
+
+		// 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 logging.RecoverPanic(fmt.Sprintf("subscription-%s", name), nil)
+
+		subCh := subscriber(ctx)
+
+		for {
+			select {
+			case event, ok := <-subCh:
+				if !ok {
+					logging.Info("subscription channel closed", "name", name)
+					return
+				}
+
+				var msg tea.Msg = event
+
+				select {
+				case outputCh <- msg:
+				case <-time.After(2 * time.Second):
+					logging.Warn("message dropped due to slow consumer", "name", name)
+				case <-ctx.Done():
+					logging.Info("subscription cancelled", "name", name)
+					return
+				}
+			case <-ctx.Done():
+				logging.Info("subscription cancelled", "name", name)
+				return
 			}
-			wg.Done()
-		}()
-	}
-	{
-		sub := app.Permissions.Subscribe(ctx)
-		wg.Add(1)
+		}
+	}()
+}
+
+func setupSubscriptions(app *app.App, parentCtx context.Context) (chan tea.Msg, func()) {
+	ch := make(chan tea.Msg, 100)
+
+	wg := sync.WaitGroup{}
+	ctx, cancel := context.WithCancel(parentCtx) // Inherit from parent context
+
+	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)
+
+	cleanupFunc := func() {
+		logging.Info("Cancelling all subscriptions")
+		cancel() // Signal all goroutines to stop
+
+		waitCh := make(chan struct{})
 		go func() {
-			for ev := range sub {
-				ch <- ev
-			}
-			wg.Done()
+			defer logging.RecoverPanic("subscription-cleanup", nil)
+			wg.Wait()
+			close(waitCh)
 		}()
+
+		select {
+		case <-waitCh:
+			logging.Info("All subscription goroutines completed successfully")
+			close(ch) // Only close after all writers are confirmed done
+		case <-time.After(5 * time.Second):
+			logging.Warn("Timed out waiting for some subscription goroutines to complete")
+			close(ch)
+		}
 	}
-	return ch, func() {
-		cancel()
-		wg.Wait()
-		close(ch)
-	}
+	return ch, cleanupFunc
 }
 
 func Execute() {
@@ -133,5 +247,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")
 }