root.go

  1package cmd
  2
  3import (
  4	"context"
  5	"fmt"
  6	"os"
  7	"sync"
  8	"time"
  9
 10	tea "github.com/charmbracelet/bubbletea"
 11	"github.com/kujtimiihoxha/opencode/internal/app"
 12	"github.com/kujtimiihoxha/opencode/internal/config"
 13	"github.com/kujtimiihoxha/opencode/internal/db"
 14	"github.com/kujtimiihoxha/opencode/internal/llm/agent"
 15	"github.com/kujtimiihoxha/opencode/internal/logging"
 16	"github.com/kujtimiihoxha/opencode/internal/pubsub"
 17	"github.com/kujtimiihoxha/opencode/internal/tui"
 18	zone "github.com/lrstanley/bubblezone"
 19	"github.com/spf13/cobra"
 20)
 21
 22var rootCmd = &cobra.Command{
 23	Use:   "OpenCode",
 24	Short: "A terminal AI assistant for software development",
 25	Long: `OpenCode is a powerful terminal-based AI assistant that helps with software development tasks.
 26It provides an interactive chat interface with AI capabilities, code analysis, and LSP integration
 27to assist developers in writing, debugging, and understanding code directly from the terminal.`,
 28	RunE: func(cmd *cobra.Command, args []string) error {
 29		// If the help flag is set, show the help message
 30		if cmd.Flag("help").Changed {
 31			cmd.Help()
 32			return nil
 33		}
 34
 35		// Load the config
 36		debug, _ := cmd.Flags().GetBool("debug")
 37		cwd, _ := cmd.Flags().GetString("cwd")
 38		if cwd != "" {
 39			err := os.Chdir(cwd)
 40			if err != nil {
 41				return fmt.Errorf("failed to change directory: %v", err)
 42			}
 43		}
 44		if cwd == "" {
 45			c, err := os.Getwd()
 46			if err != nil {
 47				return fmt.Errorf("failed to get current working directory: %v", err)
 48			}
 49			cwd = c
 50		}
 51		_, err := config.Load(cwd, debug)
 52		if err != nil {
 53			return err
 54		}
 55
 56		// Connect DB, this will also run migrations
 57		conn, err := db.Connect()
 58		if err != nil {
 59			return err
 60		}
 61
 62		// Create main context for the application
 63		ctx, cancel := context.WithCancel(context.Background())
 64		defer cancel()
 65
 66		app, err := app.New(ctx, conn)
 67		if err != nil {
 68			logging.Error("Failed to create app: %v", err)
 69			return err
 70		}
 71
 72		// Set up the TUI
 73		zone.NewGlobal()
 74		program := tea.NewProgram(
 75			tui.New(app),
 76			tea.WithAltScreen(),
 77			tea.WithMouseCellMotion(),
 78		)
 79
 80		// Initialize MCP tools in the background
 81		initMCPTools(ctx, app)
 82
 83		// Setup the subscriptions, this will send services events to the TUI
 84		ch, cancelSubs := setupSubscriptions(app, ctx)
 85
 86		// Create a context for the TUI message handler
 87		tuiCtx, tuiCancel := context.WithCancel(ctx)
 88		var tuiWg sync.WaitGroup
 89		tuiWg.Add(1)
 90
 91		// Set up message handling for the TUI
 92		go func() {
 93			defer tuiWg.Done()
 94			defer logging.RecoverPanic("TUI-message-handler", func() {
 95				attemptTUIRecovery(program)
 96			})
 97
 98			for {
 99				select {
100				case <-tuiCtx.Done():
101					logging.Info("TUI message handler shutting down")
102					return
103				case msg, ok := <-ch:
104					if !ok {
105						logging.Info("TUI message channel closed")
106						return
107					}
108					program.Send(msg)
109				}
110			}
111		}()
112
113		// Cleanup function for when the program exits
114		cleanup := func() {
115			// Shutdown the app
116			app.Shutdown()
117
118			// Cancel subscriptions first
119			cancelSubs()
120
121			// Then cancel TUI message handler
122			tuiCancel()
123
124			// Wait for TUI message handler to finish
125			tuiWg.Wait()
126
127			logging.Info("All goroutines cleaned up")
128		}
129
130		// Run the TUI
131		result, err := program.Run()
132		cleanup()
133
134		if err != nil {
135			logging.Error("TUI error: %v", err)
136			return fmt.Errorf("TUI error: %v", err)
137		}
138
139		logging.Info("TUI exited with result: %v", result)
140		return nil
141	},
142}
143
144// attemptTUIRecovery tries to recover the TUI after a panic
145func attemptTUIRecovery(program *tea.Program) {
146	logging.Info("Attempting to recover TUI after panic")
147
148	// We could try to restart the TUI or gracefully exit
149	// For now, we'll just quit the program to avoid further issues
150	program.Quit()
151}
152
153func initMCPTools(ctx context.Context, app *app.App) {
154	go func() {
155		defer logging.RecoverPanic("MCP-goroutine", nil)
156
157		// Create a context with timeout for the initial MCP tools fetch
158		ctxWithTimeout, cancel := context.WithTimeout(ctx, 30*time.Second)
159		defer cancel()
160
161		// Set this up once with proper error handling
162		agent.GetMcpTools(ctxWithTimeout, app.Permissions)
163		logging.Info("MCP message handling goroutine exiting")
164	}()
165}
166
167func setupSubscriber[T any](
168	ctx context.Context,
169	wg *sync.WaitGroup,
170	name string,
171	subscriber func(context.Context) <-chan pubsub.Event[T],
172	outputCh chan<- tea.Msg,
173) {
174	wg.Add(1)
175	go func() {
176		defer wg.Done()
177		defer logging.RecoverPanic(fmt.Sprintf("subscription-%s", name), nil)
178
179		subCh := subscriber(ctx)
180
181		for {
182			select {
183			case event, ok := <-subCh:
184				if !ok {
185					logging.Info("subscription channel closed", "name", name)
186					return
187				}
188
189				var msg tea.Msg = event
190
191				select {
192				case outputCh <- msg:
193				case <-time.After(2 * time.Second):
194					logging.Warn("message dropped due to slow consumer", "name", name)
195				case <-ctx.Done():
196					logging.Info("subscription cancelled", "name", name)
197					return
198				}
199			case <-ctx.Done():
200				logging.Info("subscription cancelled", "name", name)
201				return
202			}
203		}
204	}()
205}
206
207func setupSubscriptions(app *app.App, parentCtx context.Context) (chan tea.Msg, func()) {
208	ch := make(chan tea.Msg, 100)
209
210	wg := sync.WaitGroup{}
211	ctx, cancel := context.WithCancel(parentCtx) // Inherit from parent context
212
213	setupSubscriber(ctx, &wg, "logging", logging.Subscribe, ch)
214	setupSubscriber(ctx, &wg, "sessions", app.Sessions.Subscribe, ch)
215	setupSubscriber(ctx, &wg, "messages", app.Messages.Subscribe, ch)
216	setupSubscriber(ctx, &wg, "permissions", app.Permissions.Subscribe, ch)
217
218	cleanupFunc := func() {
219		logging.Info("Cancelling all subscriptions")
220		cancel() // Signal all goroutines to stop
221
222		waitCh := make(chan struct{})
223		go func() {
224			defer logging.RecoverPanic("subscription-cleanup", nil)
225			wg.Wait()
226			close(waitCh)
227		}()
228
229		select {
230		case <-waitCh:
231			logging.Info("All subscription goroutines completed successfully")
232			close(ch) // Only close after all writers are confirmed done
233		case <-time.After(5 * time.Second):
234			logging.Warn("Timed out waiting for some subscription goroutines to complete")
235			close(ch)
236		}
237	}
238	return ch, cleanupFunc
239}
240
241func Execute() {
242	err := rootCmd.Execute()
243	if err != nil {
244		os.Exit(1)
245	}
246}
247
248func init() {
249	rootCmd.Flags().BoolP("help", "h", false, "Help")
250	rootCmd.Flags().BoolP("debug", "d", false, "Debug")
251	rootCmd.Flags().StringP("cwd", "c", "", "Current working directory")
252}