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, ctx)
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 subCh := subscriber(ctx)
178
179 for {
180 select {
181 case event, ok := <-subCh:
182 if !ok {
183 logging.Info("%s subscription channel closed", name)
184 return
185 }
186
187 var msg tea.Msg = event
188
189 select {
190 case outputCh <- msg:
191 case <-time.After(2 * time.Second):
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, parentCtx context.Context) (chan tea.Msg, func()) {
206 ch := make(chan tea.Msg, 100)
207
208 wg := sync.WaitGroup{}
209 ctx, cancel := context.WithCancel(parentCtx) // Inherit from parent context
210
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 cleanupFunc := func() {
217 logging.Info("Cancelling all subscriptions")
218 cancel() // Signal all goroutines to stop
219
220 waitCh := make(chan struct{})
221 go func() {
222 defer logging.RecoverPanic("subscription-cleanup", nil)
223 wg.Wait()
224 close(waitCh)
225 }()
226
227 select {
228 case <-waitCh:
229 logging.Info("All subscription goroutines completed successfully")
230 close(ch) // Only close after all writers are confirmed done
231 case <-time.After(5 * time.Second):
232 logging.Warn("Timed out waiting for some subscription goroutines to complete")
233 close(ch)
234 }
235 }
236 return ch, cleanupFunc
237}
238
239func Execute() {
240 err := rootCmd.Execute()
241 if err != nil {
242 os.Exit(1)
243 }
244}
245
246func init() {
247 rootCmd.Flags().BoolP("help", "h", false, "Help")
248 rootCmd.Flags().BoolP("debug", "d", false, "Debug")
249 rootCmd.Flags().StringP("cwd", "c", "", "Current working directory")
250}