git.iko.soy/crush
on 21ca2ae help sign in
crush/cmd/root.go
309 lines 8.4 kB history raw
1package cmd
3import (
4 "context"
5 "fmt"
6 "os"
7 "sync"
8 "time"
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"
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
34 # Run with debug logging
35 opencode -d
37 # Run with debug logging in a specific directory
38 opencode -d -c /path/to/project
40 # Print version
41 opencode -v
43 # Run a single non-interactive prompt
44 opencode -p "Explain the use of context in Go"
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 }
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")
67 // Validate format option
68 if !format.IsValid(outputFormat) {
69 return fmt.Errorf("invalid format option: %s\n%s", outputFormat, format.GetHelpText())
70 }
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 }
90 // Connect DB, this will also run migrations
91 conn, err := db.Connect()
92 if err != nil {
93 return err
94 }
96 // Create main context for the application
97 ctx, cancel := context.WithCancel(context.Background())
98 defer cancel()
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()
108 // Initialize MCP tools early for both modes
109 initMCPTools(ctx, app)
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 }
117 // Interactive mode
118 // Set up the TUI
119 zone.NewGlobal()
120 program := tea.NewProgram(
121 tui.New(app),
122 tea.WithAltScreen(),
123 )
125 // Setup the subscriptions, this will send services events to the TUI
126 ch, cancelSubs := setupSubscriptions(app, ctx)
128 // Create a context for the TUI message handler
129 tuiCtx, tuiCancel := context.WithCancel(ctx)
130 var tuiWg sync.WaitGroup
131 tuiWg.Add(1)
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 })
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 }()
155 // Cleanup function for when the program exits
156 cleanup := func() {
157 // Shutdown the app
158 app.Shutdown()
160 // Cancel subscriptions first
161 cancelSubs()
163 // Then cancel TUI message handler
164 tuiCancel()
166 // Wait for TUI message handler to finish
167 tuiWg.Wait()
169 logging.Info("All goroutines cleaned up")
170 }
172 // Run the TUI
173 result, err := program.Run()
174 cleanup()
176 if err != nil {
177 logging.Error("TUI error: %v", err)
178 return fmt.Errorf("TUI error: %v", err)
179 }
181 logging.Info("TUI exited with result: %v", result)
182 return nil
183 },
186// attemptTUIRecovery tries to recover the TUI after a panic
187func attemptTUIRecovery(program *tea.Program) {
188 logging.Info("Attempting to recover TUI after panic")
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()
195func initMCPTools(ctx context.Context, app *app.App) {
196 go func() {
197 defer logging.RecoverPanic("MCP-goroutine", nil)
199 // Create a context with timeout for the initial MCP tools fetch
200 ctxWithTimeout, cancel := context.WithTimeout(ctx, 30*time.Second)
201 defer cancel()
203 // Set this up once with proper error handling
204 agent.GetMcpTools(ctxWithTimeout, app.Permissions)
205 logging.Info("MCP message handling goroutine exiting")
206 }()
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)
221 subCh := subscriber(ctx)
223 for {
224 select {
225 case event, ok := <-subCh:
226 if !ok {
227 logging.Info("subscription channel closed", "name", name)
228 return
229 }
231 var msg tea.Msg = event
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 }()
249func setupSubscriptions(app *app.App, parentCtx context.Context) (chan tea.Msg, func()) {
250 ch := make(chan tea.Msg, 100)
252 wg := sync.WaitGroup{}
253 ctx, cancel := context.WithCancel(parentCtx) // Inherit from parent context
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)
261 cleanupFunc := func() {
262 logging.Info("Cancelling all subscriptions")
263 cancel() // Signal all goroutines to stop
265 waitCh := make(chan struct{})
266 go func() {
267 defer logging.RecoverPanic("subscription-cleanup", nil)
268 wg.Wait()
269 close(waitCh)
270 }()
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
284func Execute() {
285 err := rootCmd.Execute()
286 if err != nil {
287 os.Exit(1)
288 }
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")
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)")
302 // Add quiet flag to hide spinner in non-interactive mode
303 rootCmd.Flags().BoolP("quiet", "q", false, "Hide spinner in non-interactive mode")
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 })