root.go

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