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