app.go

  1package app
  2
  3import (
  4	"context"
  5	"database/sql"
  6	"errors"
  7	"fmt"
  8	"log/slog"
  9	"sync"
 10	"time"
 11
 12	"charm.land/fantasy"
 13	tea "github.com/charmbracelet/bubbletea/v2"
 14	"github.com/charmbracelet/crush/internal/agent"
 15	"github.com/charmbracelet/crush/internal/agent/tools"
 16	"github.com/charmbracelet/crush/internal/config"
 17	"github.com/charmbracelet/crush/internal/csync"
 18	"github.com/charmbracelet/crush/internal/db"
 19	"github.com/charmbracelet/crush/internal/format"
 20	"github.com/charmbracelet/crush/internal/history"
 21	"github.com/charmbracelet/crush/internal/log"
 22	"github.com/charmbracelet/crush/internal/lsp"
 23	"github.com/charmbracelet/crush/internal/message"
 24	"github.com/charmbracelet/crush/internal/permission"
 25	"github.com/charmbracelet/crush/internal/pubsub"
 26	"github.com/charmbracelet/crush/internal/session"
 27	"github.com/charmbracelet/crush/internal/shell"
 28	"github.com/charmbracelet/x/ansi"
 29)
 30
 31type App struct {
 32	Sessions    session.Service
 33	Messages    message.Service
 34	History     history.Service
 35	Permissions permission.Service
 36
 37	AgentCoordinator agent.Coordinator
 38
 39	LSPClients *csync.Map[string, *lsp.Client]
 40
 41	config *config.Config
 42
 43	serviceEventsWG *sync.WaitGroup
 44	eventsCtx       context.Context
 45	events          chan tea.Msg
 46	tuiWG           *sync.WaitGroup
 47
 48	// global context and cleanup functions
 49	globalCtx    context.Context
 50	cleanupFuncs []func() error
 51}
 52
 53// New initializes a new applcation instance.
 54func New(ctx context.Context, conn *sql.DB, cfg *config.Config) (*App, error) {
 55	q := db.New(conn)
 56	sessions := session.NewService(q)
 57	messages := message.NewService(q)
 58	files := history.NewService(q, conn)
 59	skipPermissionsRequests := cfg.Permissions != nil && cfg.Permissions.SkipRequests
 60	allowedTools := []string{}
 61	if cfg.Permissions != nil && cfg.Permissions.AllowedTools != nil {
 62		allowedTools = cfg.Permissions.AllowedTools
 63	}
 64
 65	app := &App{
 66		Sessions:    sessions,
 67		Messages:    messages,
 68		History:     files,
 69		Permissions: permission.NewPermissionService(cfg.WorkingDir(), skipPermissionsRequests, allowedTools),
 70		LSPClients:  csync.NewMap[string, *lsp.Client](),
 71
 72		globalCtx: ctx,
 73
 74		config: cfg,
 75
 76		events:          make(chan tea.Msg, 100),
 77		serviceEventsWG: &sync.WaitGroup{},
 78		tuiWG:           &sync.WaitGroup{},
 79	}
 80
 81	app.setupEvents()
 82
 83	// Initialize LSP clients in the background.
 84	app.initLSPClients(ctx)
 85
 86	// cleanup database upon app shutdown
 87	app.cleanupFuncs = append(app.cleanupFuncs, conn.Close)
 88
 89	// TODO: remove the concept of agent config, most likely.
 90	if !cfg.IsConfigured() {
 91		slog.Warn("No agent configuration found")
 92		return app, nil
 93	}
 94	if err := app.InitCoderAgent(ctx); err != nil {
 95		return nil, fmt.Errorf("failed to initialize coder agent: %w", err)
 96	}
 97	return app, nil
 98}
 99
100// Config returns the application configuration.
101func (app *App) Config() *config.Config {
102	return app.config
103}
104
105// RunNonInteractive handles the execution flow when a prompt is provided via
106// CLI flag.
107func (app *App) RunNonInteractive(ctx context.Context, prompt string, quiet bool) error {
108	slog.Info("Running in non-interactive mode")
109
110	ctx, cancel := context.WithCancel(ctx)
111	defer cancel()
112
113	var spinner *format.Spinner
114	if !quiet {
115		spinner = format.NewSpinner(ctx, cancel, "Generating")
116		spinner.Start()
117	}
118
119	// Helper function to stop spinner once.
120	stopSpinner := func() {
121		if !quiet && spinner != nil {
122			spinner.Stop()
123			spinner = nil
124		}
125	}
126	defer stopSpinner()
127
128	const maxPromptLengthForTitle = 100
129	titlePrefix := "Non-interactive: "
130	var titleSuffix string
131
132	if len(prompt) > maxPromptLengthForTitle {
133		titleSuffix = prompt[:maxPromptLengthForTitle] + "..."
134	} else {
135		titleSuffix = prompt
136	}
137	title := titlePrefix + titleSuffix
138
139	sess, err := app.Sessions.Create(ctx, title)
140	if err != nil {
141		return fmt.Errorf("failed to create session for non-interactive mode: %w", err)
142	}
143	slog.Info("Created session for non-interactive run", "session_id", sess.ID)
144
145	// Automatically approve all permission requests for this non-interactive session
146	app.Permissions.AutoApproveSession(sess.ID)
147
148	type response struct {
149		result *fantasy.AgentResult
150		err    error
151	}
152	done := make(chan response, 1)
153
154	go func(ctx context.Context, sessionID, prompt string) {
155		result, err := app.AgentCoordinator.Run(ctx, sess.ID, prompt)
156		if err != nil {
157			done <- response{
158				err: fmt.Errorf("failed to start agent processing stream: %w", err),
159			}
160		}
161		done <- response{
162			result: result,
163		}
164	}(ctx, sess.ID, prompt)
165
166	messageEvents := app.Messages.Subscribe(ctx)
167	messageReadBytes := make(map[string]int)
168
169	defer fmt.Printf(ansi.ResetProgressBar)
170	for {
171		// HACK: add it again on every iteration so it doesn't get hidden by
172		// the terminal due to inactivity.
173		fmt.Printf(ansi.SetIndeterminateProgressBar)
174		select {
175		case result := <-done:
176			stopSpinner()
177			if result.err != nil {
178				if errors.Is(result.err, context.Canceled) || errors.Is(result.err, agent.ErrRequestCancelled) {
179					slog.Info("Non-interactive: agent processing cancelled", "session_id", sess.ID)
180					return nil
181				}
182				return fmt.Errorf("agent processing failed: %w", result.err)
183			}
184			return nil
185
186		case event := <-messageEvents:
187			msg := event.Payload
188			if msg.SessionID == sess.ID && msg.Role == message.Assistant && len(msg.Parts) > 0 {
189				stopSpinner()
190
191				content := msg.Content().String()
192				readBytes := messageReadBytes[msg.ID]
193
194				if len(content) < readBytes {
195					slog.Error("Non-interactive: message content is shorter than read bytes", "message_length", len(content), "read_bytes", readBytes)
196					return fmt.Errorf("message content is shorter than read bytes: %d < %d", len(content), readBytes)
197				}
198
199				part := content[readBytes:]
200				fmt.Print(part)
201				messageReadBytes[msg.ID] = len(content)
202			}
203
204		case <-ctx.Done():
205			stopSpinner()
206			return ctx.Err()
207		}
208	}
209}
210
211func (app *App) UpdateAgentModel(ctx context.Context) error {
212	return app.AgentCoordinator.UpdateModels(ctx)
213}
214
215func (app *App) setupEvents() {
216	ctx, cancel := context.WithCancel(app.globalCtx)
217	app.eventsCtx = ctx
218	setupSubscriber(ctx, app.serviceEventsWG, "sessions", app.Sessions.Subscribe, app.events)
219	setupSubscriber(ctx, app.serviceEventsWG, "messages", app.Messages.Subscribe, app.events)
220	setupSubscriber(ctx, app.serviceEventsWG, "permissions", app.Permissions.Subscribe, app.events)
221	setupSubscriber(ctx, app.serviceEventsWG, "permissions-notifications", app.Permissions.SubscribeNotifications, app.events)
222	setupSubscriber(ctx, app.serviceEventsWG, "history", app.History.Subscribe, app.events)
223	setupSubscriber(ctx, app.serviceEventsWG, "mcp", tools.SubscribeMCPEvents, app.events)
224	setupSubscriber(ctx, app.serviceEventsWG, "lsp", SubscribeLSPEvents, app.events)
225	cleanupFunc := func() error {
226		cancel()
227		app.serviceEventsWG.Wait()
228		return nil
229	}
230	app.cleanupFuncs = append(app.cleanupFuncs, cleanupFunc)
231}
232
233func setupSubscriber[T any](
234	ctx context.Context,
235	wg *sync.WaitGroup,
236	name string,
237	subscriber func(context.Context) <-chan pubsub.Event[T],
238	outputCh chan<- tea.Msg,
239) {
240	wg.Go(func() {
241		subCh := subscriber(ctx)
242		for {
243			select {
244			case event, ok := <-subCh:
245				if !ok {
246					slog.Debug("subscription channel closed", "name", name)
247					return
248				}
249				var msg tea.Msg = event
250				select {
251				case outputCh <- msg:
252				case <-time.After(2 * time.Second):
253					slog.Warn("message dropped due to slow consumer", "name", name)
254				case <-ctx.Done():
255					slog.Debug("subscription cancelled", "name", name)
256					return
257				}
258			case <-ctx.Done():
259				slog.Debug("subscription cancelled", "name", name)
260				return
261			}
262		}
263	})
264}
265
266func (app *App) InitCoderAgent(ctx context.Context) error {
267	coderAgentCfg := app.config.Agents[config.AgentCoder]
268	if coderAgentCfg.ID == "" {
269		return fmt.Errorf("coder agent configuration is missing")
270	}
271	var err error
272	app.AgentCoordinator, err = agent.NewCoordinator(
273		ctx,
274		app.config,
275		app.Sessions,
276		app.Messages,
277		app.Permissions,
278		app.History,
279		app.LSPClients,
280	)
281	if err != nil {
282		slog.Error("Failed to create coder agent", "err", err)
283		return err
284	}
285
286	// Add MCP client cleanup to shutdown process
287	app.cleanupFuncs = append(app.cleanupFuncs, tools.CloseMCPClients)
288	return nil
289}
290
291// Subscribe sends events to the TUI as tea.Msgs.
292func (app *App) Subscribe(program *tea.Program) {
293	defer log.RecoverPanic("app.Subscribe", func() {
294		slog.Info("TUI subscription panic: attempting graceful shutdown")
295		program.Quit()
296	})
297
298	app.tuiWG.Add(1)
299	tuiCtx, tuiCancel := context.WithCancel(app.globalCtx)
300	app.cleanupFuncs = append(app.cleanupFuncs, func() error {
301		slog.Debug("Cancelling TUI message handler")
302		tuiCancel()
303		app.tuiWG.Wait()
304		return nil
305	})
306	defer app.tuiWG.Done()
307
308	for {
309		select {
310		case <-tuiCtx.Done():
311			slog.Debug("TUI message handler shutting down")
312			return
313		case msg, ok := <-app.events:
314			if !ok {
315				slog.Debug("TUI message channel closed")
316				return
317			}
318			program.Send(msg)
319		}
320	}
321}
322
323// Shutdown performs a graceful shutdown of the application.
324func (app *App) Shutdown() {
325	if app.AgentCoordinator != nil {
326		app.AgentCoordinator.CancelAll()
327	}
328
329	// Kill all background shells.
330	shell.GetBackgroundShellManager().KillAll()
331
332	// Shutdown all LSP clients.
333	for name, client := range app.LSPClients.Seq2() {
334		shutdownCtx, cancel := context.WithTimeout(app.globalCtx, 5*time.Second)
335		if err := client.Close(shutdownCtx); err != nil {
336			slog.Error("Failed to shutdown LSP client", "name", name, "error", err)
337		}
338		cancel()
339	}
340
341	// Call call cleanup functions.
342	for _, cleanup := range app.cleanupFuncs {
343		if cleanup != nil {
344			if err := cleanup(); err != nil {
345				slog.Error("Failed to cleanup app properly on shutdown", "error", err)
346			}
347		}
348	}
349}