1package app
2
3import (
4 "context"
5 "database/sql"
6 "errors"
7 "fmt"
8 "log/slog"
9 "maps"
10 "sync"
11 "time"
12
13 tea "github.com/charmbracelet/bubbletea/v2"
14 "github.com/charmbracelet/crush/internal/config"
15 "github.com/charmbracelet/crush/internal/db"
16 "github.com/charmbracelet/crush/internal/format"
17 "github.com/charmbracelet/crush/internal/history"
18 "github.com/charmbracelet/crush/internal/llm/agent"
19 "github.com/charmbracelet/crush/internal/pubsub"
20
21 "github.com/charmbracelet/crush/internal/lsp"
22 "github.com/charmbracelet/crush/internal/message"
23 "github.com/charmbracelet/crush/internal/permission"
24 "github.com/charmbracelet/crush/internal/session"
25)
26
27type App struct {
28 Sessions session.Service
29 Messages message.Service
30 History history.Service
31 Permissions permission.Service
32
33 CoderAgent agent.Service
34
35 LSPClients map[string]*lsp.Client
36
37 clientsMutex sync.RWMutex
38
39 watcherCancelFuncs []context.CancelFunc
40 cancelFuncsMutex sync.Mutex
41 lspWatcherWG sync.WaitGroup
42
43 config *config.Config
44
45 serviceEventsWG *sync.WaitGroup
46 eventsCtx context.Context
47 events chan tea.Msg
48 tuiWG *sync.WaitGroup
49
50 // global context and cleanup functions
51 globalCtx context.Context
52 cleanupFuncs []func()
53}
54
55func New(ctx context.Context, conn *sql.DB, cfg *config.Config) (*App, error) {
56 q := db.New(conn)
57 sessions := session.NewService(q)
58 messages := message.NewService(q)
59 files := history.NewService(q, conn)
60
61 app := &App{
62 Sessions: sessions,
63 Messages: messages,
64 History: files,
65 Permissions: permission.NewPermissionService(cfg.WorkingDir()),
66 LSPClients: make(map[string]*lsp.Client),
67
68 globalCtx: ctx,
69
70 config: cfg,
71
72 events: make(chan tea.Msg, 100),
73 serviceEventsWG: &sync.WaitGroup{},
74 tuiWG: &sync.WaitGroup{},
75 }
76
77 app.setupEvents()
78
79 // Initialize LSP clients in the background
80 go app.initLSPClients(ctx)
81
82 // TODO: remove the concept of agent config most likely
83 if cfg.IsConfigured() {
84 if err := app.InitCoderAgent(); err != nil {
85 return nil, fmt.Errorf("failed to initialize coder agent: %w", err)
86 }
87 } else {
88 slog.Warn("No agent configuration found")
89 }
90 return app, nil
91}
92
93// RunNonInteractive handles the execution flow when a prompt is provided via CLI flag.
94func (a *App) RunNonInteractive(ctx context.Context, prompt string, outputFormat string, quiet bool) error {
95 slog.Info("Running in non-interactive mode")
96
97 // Start spinner if not in quiet mode
98 var spinner *format.Spinner
99 if !quiet {
100 spinner = format.NewSpinner("Thinking...")
101 spinner.Start()
102 defer spinner.Stop()
103 }
104
105 const maxPromptLengthForTitle = 100
106 titlePrefix := "Non-interactive: "
107 var titleSuffix string
108
109 if len(prompt) > maxPromptLengthForTitle {
110 titleSuffix = prompt[:maxPromptLengthForTitle] + "..."
111 } else {
112 titleSuffix = prompt
113 }
114 title := titlePrefix + titleSuffix
115
116 sess, err := a.Sessions.Create(ctx, title)
117 if err != nil {
118 return fmt.Errorf("failed to create session for non-interactive mode: %w", err)
119 }
120 slog.Info("Created session for non-interactive run", "session_id", sess.ID)
121
122 // Automatically approve all permission requests for this non-interactive session
123 a.Permissions.AutoApproveSession(sess.ID)
124
125 done, err := a.CoderAgent.Run(ctx, sess.ID, prompt)
126 if err != nil {
127 return fmt.Errorf("failed to start agent processing stream: %w", err)
128 }
129
130 result := <-done
131 if result.Error != nil {
132 if errors.Is(result.Error, context.Canceled) || errors.Is(result.Error, agent.ErrRequestCancelled) {
133 slog.Info("Agent processing cancelled", "session_id", sess.ID)
134 return nil
135 }
136 return fmt.Errorf("agent processing failed: %w", result.Error)
137 }
138
139 // Stop spinner before printing output
140 if !quiet && spinner != nil {
141 spinner.Stop()
142 }
143
144 // Get the text content from the response
145 content := "No content available"
146 if result.Message.Content().String() != "" {
147 content = result.Message.Content().String()
148 }
149
150 fmt.Println(format.FormatOutput(content, outputFormat))
151
152 slog.Info("Non-interactive run completed", "session_id", sess.ID)
153
154 return nil
155}
156
157func (app *App) UpdateAgentModel() error {
158 return app.CoderAgent.UpdateModel()
159}
160
161func (app *App) setupEvents() {
162 ctx, cancel := context.WithCancel(app.globalCtx)
163 app.eventsCtx = ctx
164 setupSubscriber(ctx, app.serviceEventsWG, "sessions", app.Sessions.Subscribe, app.events)
165 setupSubscriber(ctx, app.serviceEventsWG, "messages", app.Messages.Subscribe, app.events)
166 setupSubscriber(ctx, app.serviceEventsWG, "permissions", app.Permissions.Subscribe, app.events)
167 setupSubscriber(ctx, app.serviceEventsWG, "history", app.History.Subscribe, app.events)
168 cleanupFunc := func() {
169 cancel()
170 app.serviceEventsWG.Wait()
171 }
172 app.cleanupFuncs = append(app.cleanupFuncs, cleanupFunc)
173}
174
175func setupSubscriber[T any](
176 ctx context.Context,
177 wg *sync.WaitGroup,
178 name string,
179 subscriber func(context.Context) <-chan pubsub.Event[T],
180 outputCh chan<- tea.Msg,
181) {
182 wg.Add(1)
183 go func() {
184 defer wg.Done()
185 subCh := subscriber(ctx)
186 for {
187 select {
188 case event, ok := <-subCh:
189 if !ok {
190 slog.Debug("subscription channel closed", "name", name)
191 return
192 }
193 var msg tea.Msg = event
194 select {
195 case outputCh <- msg:
196 case <-time.After(2 * time.Second):
197 slog.Warn("message dropped due to slow consumer", "name", name)
198 case <-ctx.Done():
199 slog.Debug("subscription cancelled", "name", name)
200 return
201 }
202 case <-ctx.Done():
203 slog.Debug("subscription cancelled", "name", name)
204 return
205 }
206 }
207 }()
208}
209
210func (app *App) InitCoderAgent() error {
211 coderAgentCfg := app.config.Agents["coder"]
212 if coderAgentCfg.ID == "" {
213 return fmt.Errorf("coder agent configuration is missing")
214 }
215 var err error
216 app.CoderAgent, err = agent.NewAgent(
217 coderAgentCfg,
218 app.Permissions,
219 app.Sessions,
220 app.Messages,
221 app.History,
222 app.LSPClients,
223 )
224 if err != nil {
225 slog.Error("Failed to create coder agent", "err", err)
226 return err
227 }
228 setupSubscriber(app.eventsCtx, app.serviceEventsWG, "coderAgent", app.CoderAgent.Subscribe, app.events)
229 return nil
230}
231
232func (app *App) Subscribe(program *tea.Program) {
233 app.tuiWG.Add(1)
234 tuiCtx, tuiCancel := context.WithCancel(app.globalCtx)
235 app.cleanupFuncs = append(app.cleanupFuncs, func() {
236 slog.Debug("Cancelling TUI message handler")
237 tuiCancel()
238 app.tuiWG.Wait()
239 })
240 defer app.tuiWG.Done()
241 for {
242 select {
243 case <-tuiCtx.Done():
244 slog.Debug("TUI message handler shutting down")
245 return
246 case msg, ok := <-app.events:
247 if !ok {
248 slog.Debug("TUI message channel closed")
249 return
250 }
251 program.Send(msg)
252 }
253 }
254}
255
256// Shutdown performs a clean shutdown of the application
257func (app *App) Shutdown() {
258 app.cancelFuncsMutex.Lock()
259 for _, cancel := range app.watcherCancelFuncs {
260 cancel()
261 }
262 app.cancelFuncsMutex.Unlock()
263 app.lspWatcherWG.Wait()
264
265 app.clientsMutex.RLock()
266 clients := make(map[string]*lsp.Client, len(app.LSPClients))
267 maps.Copy(clients, app.LSPClients)
268 app.clientsMutex.RUnlock()
269
270 for name, client := range clients {
271 shutdownCtx, cancel := context.WithTimeout(app.globalCtx, 5*time.Second)
272 if err := client.Shutdown(shutdownCtx); err != nil {
273 slog.Error("Failed to shutdown LSP client", "name", name, "error", err)
274 }
275 cancel()
276 }
277 if app.CoderAgent != nil {
278 app.CoderAgent.CancelAll()
279 }
280
281 for _, cleanup := range app.cleanupFuncs {
282 if cleanup != nil {
283 cleanup()
284 }
285 }
286}