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/llm/agent"
 16	"github.com/opencode-ai/opencode/internal/logging"
 17	"github.com/opencode-ai/opencode/internal/pubsub"
 18	"github.com/opencode-ai/opencode/internal/tui"
 19	"github.com/opencode-ai/opencode/internal/version"
 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		)
 83
 84		// Initialize MCP tools in the background
 85		initMCPTools(ctx, app)
 86
 87		// Setup the subscriptions, this will send services events to the TUI
 88		ch, cancelSubs := setupSubscriptions(app, ctx)
 89
 90		// Create a context for the TUI message handler
 91		tuiCtx, tuiCancel := context.WithCancel(ctx)
 92		var tuiWg sync.WaitGroup
 93		tuiWg.Add(1)
 94
 95		// Set up message handling for the TUI
 96		go func() {
 97			defer tuiWg.Done()
 98			defer logging.RecoverPanic("TUI-message-handler", func() {
 99				attemptTUIRecovery(program)
100			})
101
102			for {
103				select {
104				case <-tuiCtx.Done():
105					logging.Info("TUI message handler shutting down")
106					return
107				case msg, ok := <-ch:
108					if !ok {
109						logging.Info("TUI message channel closed")
110						return
111					}
112					program.Send(msg)
113				}
114			}
115		}()
116
117		// Cleanup function for when the program exits
118		cleanup := func() {
119			// Shutdown the app
120			app.Shutdown()
121
122			// Cancel subscriptions first
123			cancelSubs()
124
125			// Then cancel TUI message handler
126			tuiCancel()
127
128			// Wait for TUI message handler to finish
129			tuiWg.Wait()
130
131			logging.Info("All goroutines cleaned up")
132		}
133
134		// Run the TUI
135		result, err := program.Run()
136		cleanup()
137
138		if err != nil {
139			logging.Error("TUI error: %v", err)
140			return fmt.Errorf("TUI error: %v", err)
141		}
142
143		logging.Info("TUI exited with result: %v", result)
144		return nil
145	},
146}
147
148// attemptTUIRecovery tries to recover the TUI after a panic
149func attemptTUIRecovery(program *tea.Program) {
150	logging.Info("Attempting to recover TUI after panic")
151
152	// We could try to restart the TUI or gracefully exit
153	// For now, we'll just quit the program to avoid further issues
154	program.Quit()
155}
156
157func initMCPTools(ctx context.Context, app *app.App) {
158	go func() {
159		defer logging.RecoverPanic("MCP-goroutine", nil)
160
161		// Create a context with timeout for the initial MCP tools fetch
162		ctxWithTimeout, cancel := context.WithTimeout(ctx, 30*time.Second)
163		defer cancel()
164
165		// Set this up once with proper error handling
166		agent.GetMcpTools(ctxWithTimeout, app.Permissions)
167		logging.Info("MCP message handling goroutine exiting")
168	}()
169}
170
171func setupSubscriber[T any](
172	ctx context.Context,
173	wg *sync.WaitGroup,
174	name string,
175	subscriber func(context.Context) <-chan pubsub.Event[T],
176	outputCh chan<- tea.Msg,
177) {
178	wg.Add(1)
179	go func() {
180		defer wg.Done()
181		defer logging.RecoverPanic(fmt.Sprintf("subscription-%s", name), nil)
182
183		subCh := subscriber(ctx)
184
185		for {
186			select {
187			case event, ok := <-subCh:
188				if !ok {
189					logging.Info("subscription channel closed", "name", name)
190					return
191				}
192
193				var msg tea.Msg = event
194
195				select {
196				case outputCh <- msg:
197				case <-time.After(2 * time.Second):
198					logging.Warn("message dropped due to slow consumer", "name", name)
199				case <-ctx.Done():
200					logging.Info("subscription cancelled", "name", name)
201					return
202				}
203			case <-ctx.Done():
204				logging.Info("subscription cancelled", "name", name)
205				return
206			}
207		}
208	}()
209}
210
211func setupSubscriptions(app *app.App, parentCtx context.Context) (chan tea.Msg, func()) {
212	ch := make(chan tea.Msg, 100)
213
214	wg := sync.WaitGroup{}
215	ctx, cancel := context.WithCancel(parentCtx) // Inherit from parent context
216
217	setupSubscriber(ctx, &wg, "logging", logging.Subscribe, ch)
218	setupSubscriber(ctx, &wg, "sessions", app.Sessions.Subscribe, ch)
219	setupSubscriber(ctx, &wg, "messages", app.Messages.Subscribe, ch)
220	setupSubscriber(ctx, &wg, "permissions", app.Permissions.Subscribe, ch)
221	setupSubscriber(ctx, &wg, "coderAgent", app.CoderAgent.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}