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",
25 Long: `A terminal ai assistant`,
26 RunE: func(cmd *cobra.Command, args []string) error {
27 // If the help flag is set, show the help message
28 if cmd.Flag("help").Changed {
29 cmd.Help()
30 return nil
31 }
32
33 // Load the config
34 debug, _ := cmd.Flags().GetBool("debug")
35 cwd, _ := cmd.Flags().GetString("cwd")
36 if cwd != "" {
37 err := os.Chdir(cwd)
38 if err != nil {
39 return fmt.Errorf("failed to change directory: %v", err)
40 }
41 }
42 if cwd == "" {
43 c, err := os.Getwd()
44 if err != nil {
45 return fmt.Errorf("failed to get current working directory: %v", err)
46 }
47 cwd = c
48 }
49 _, err := config.Load(cwd, debug)
50 if err != nil {
51 return err
52 }
53
54 // Connect DB, this will also run migrations
55 conn, err := db.Connect()
56 if err != nil {
57 return err
58 }
59
60 // Create main context for the application
61 ctx, cancel := context.WithCancel(context.Background())
62 defer cancel()
63
64 app, err := app.New(ctx, conn)
65 if err != nil {
66 logging.Error("Failed to create app: %v", err)
67 return err
68 }
69
70 // Set up the TUI
71 zone.NewGlobal()
72 program := tea.NewProgram(
73 tui.New(app),
74 tea.WithAltScreen(),
75 tea.WithMouseCellMotion(),
76 )
77
78 // Initialize MCP tools in the background
79 initMCPTools(ctx, app)
80
81 // Setup the subscriptions, this will send services events to the TUI
82 ch, cancelSubs := setupSubscriptions(app)
83
84 // Create a context for the TUI message handler
85 tuiCtx, tuiCancel := context.WithCancel(ctx)
86 var tuiWg sync.WaitGroup
87 tuiWg.Add(1)
88
89 // Set up message handling for the TUI
90 go func() {
91 defer tuiWg.Done()
92 defer logging.RecoverPanic("TUI-message-handler", func() {
93 attemptTUIRecovery(program)
94 })
95
96 for {
97 select {
98 case <-tuiCtx.Done():
99 logging.Info("TUI message handler shutting down")
100 return
101 case msg, ok := <-ch:
102 if !ok {
103 logging.Info("TUI message channel closed")
104 return
105 }
106 program.Send(msg)
107 }
108 }
109 }()
110
111 // Cleanup function for when the program exits
112 cleanup := func() {
113 // Shutdown the app
114 app.Shutdown()
115
116 // Cancel subscriptions first
117 cancelSubs()
118
119 // Then cancel TUI message handler
120 tuiCancel()
121
122 // Wait for TUI message handler to finish
123 tuiWg.Wait()
124
125 logging.Info("All goroutines cleaned up")
126 }
127
128 // Run the TUI
129 result, err := program.Run()
130 cleanup()
131
132 if err != nil {
133 logging.Error("TUI error: %v", err)
134 return fmt.Errorf("TUI error: %v", err)
135 }
136
137 logging.Info("TUI exited with result: %v", result)
138 return nil
139 },
140}
141
142// attemptTUIRecovery tries to recover the TUI after a panic
143func attemptTUIRecovery(program *tea.Program) {
144 logging.Info("Attempting to recover TUI after panic")
145
146 // We could try to restart the TUI or gracefully exit
147 // For now, we'll just quit the program to avoid further issues
148 program.Quit()
149}
150
151func initMCPTools(ctx context.Context, app *app.App) {
152 go func() {
153 defer logging.RecoverPanic("MCP-goroutine", nil)
154
155 // Create a context with timeout for the initial MCP tools fetch
156 ctxWithTimeout, cancel := context.WithTimeout(ctx, 30*time.Second)
157 defer cancel()
158
159 // Set this up once with proper error handling
160 agent.GetMcpTools(ctxWithTimeout, app.Permissions)
161 logging.Info("MCP message handling goroutine exiting")
162 }()
163}
164
165func setupSubscriber[T any](
166 ctx context.Context,
167 wg *sync.WaitGroup,
168 name string,
169 subscriber func(context.Context) <-chan pubsub.Event[T],
170 outputCh chan<- tea.Msg,
171) {
172 wg.Add(1)
173 go func() {
174 defer wg.Done()
175 defer logging.RecoverPanic(fmt.Sprintf("subscription-%s", name), nil)
176
177 for {
178 select {
179 case event, ok := <-subscriber(ctx):
180 if !ok {
181 logging.Info("%s subscription channel closed", name)
182 return
183 }
184
185 // Convert generic event to tea.Msg if needed
186 var msg tea.Msg = event
187
188 // Non-blocking send with timeout to prevent deadlocks
189 select {
190 case outputCh <- msg:
191 case <-time.After(500 * time.Millisecond):
192 logging.Warn("%s message dropped due to slow consumer", name)
193 case <-ctx.Done():
194 logging.Info("%s subscription cancelled", name)
195 return
196 }
197 case <-ctx.Done():
198 logging.Info("%s subscription cancelled", name)
199 return
200 }
201 }
202 }()
203}
204
205func setupSubscriptions(app *app.App) (chan tea.Msg, func()) {
206 ch := make(chan tea.Msg, 100)
207 // Add a buffer to prevent blocking
208 wg := sync.WaitGroup{}
209 ctx, cancel := context.WithCancel(context.Background())
210 // Setup each subscription using the helper
211 setupSubscriber(ctx, &wg, "logging", logging.Subscribe, ch)
212 setupSubscriber(ctx, &wg, "sessions", app.Sessions.Subscribe, ch)
213 setupSubscriber(ctx, &wg, "messages", app.Messages.Subscribe, ch)
214 setupSubscriber(ctx, &wg, "permissions", app.Permissions.Subscribe, ch)
215
216 // Return channel and a cleanup function
217 cleanupFunc := func() {
218 logging.Info("Cancelling all subscriptions")
219 cancel() // Signal all goroutines to stop
220
221 // Wait with a timeout for all goroutines to complete
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 case <-time.After(5 * time.Second):
233 logging.Warn("Timed out waiting for some subscription goroutines to complete")
234 }
235
236 close(ch) // Safe to close after all writers are done or timed out
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}