app.go

  1package app
  2
  3import (
  4	"context"
  5	"database/sql"
  6	"errors"
  7	"fmt"
  8	"io"
  9	"log/slog"
 10	"sync"
 11	"time"
 12
 13	"charm.land/fantasy"
 14	tea "github.com/charmbracelet/bubbletea/v2"
 15	"github.com/charmbracelet/crush/internal/agent"
 16	"github.com/charmbracelet/crush/internal/agent/tools"
 17	"github.com/charmbracelet/crush/internal/config"
 18	"github.com/charmbracelet/crush/internal/csync"
 19	"github.com/charmbracelet/crush/internal/db"
 20	"github.com/charmbracelet/crush/internal/format"
 21	"github.com/charmbracelet/crush/internal/history"
 22	"github.com/charmbracelet/crush/internal/log"
 23	"github.com/charmbracelet/crush/internal/lsp"
 24	"github.com/charmbracelet/crush/internal/message"
 25	"github.com/charmbracelet/crush/internal/permission"
 26	"github.com/charmbracelet/crush/internal/pubsub"
 27	"github.com/charmbracelet/crush/internal/session"
 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 runs the application in non-interactive mode with the
106// given prompt, printing to stdout.
107func (app *App) RunNonInteractive(ctx context.Context, output io.Writer, 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	const 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
146	// session.
147	app.Permissions.AutoApproveSession(sess.ID)
148
149	type response struct {
150		result *fantasy.AgentResult
151		err    error
152	}
153	done := make(chan response, 1)
154
155	go func(ctx context.Context, sessionID, prompt string) {
156		result, err := app.AgentCoordinator.Run(ctx, sess.ID, prompt)
157		if err != nil {
158			done <- response{
159				err: fmt.Errorf("failed to start agent processing stream: %w", err),
160			}
161		}
162		done <- response{
163			result: result,
164		}
165	}(ctx, sess.ID, prompt)
166
167	messageEvents := app.Messages.Subscribe(ctx)
168	messageReadBytes := make(map[string]int)
169
170	defer func() {
171		_, _ = fmt.Printf(ansi.ResetProgressBar)
172
173		// Always print a newline at the end. If output is a TTY this will
174		// prevent the prompt from overwriting the last line of output.
175		_, _ = fmt.Fprintln(output)
176	}()
177
178	for {
179		// HACK: Reinitialize the terminal progress bar on every iteration so
180		// it doesn't get hidden by the terminal due to inactivity.
181		_, _ = fmt.Printf(ansi.SetIndeterminateProgressBar)
182
183		select {
184		case result := <-done:
185			stopSpinner()
186			if result.err != nil {
187				if errors.Is(result.err, context.Canceled) || errors.Is(result.err, agent.ErrRequestCancelled) {
188					slog.Info("Non-interactive: agent processing cancelled", "session_id", sess.ID)
189					return nil
190				}
191				return fmt.Errorf("agent processing failed: %w", result.err)
192			}
193			return nil
194
195		case event := <-messageEvents:
196			msg := event.Payload
197			if msg.SessionID == sess.ID && msg.Role == message.Assistant && len(msg.Parts) > 0 {
198				stopSpinner()
199
200				content := msg.Content().String()
201				readBytes := messageReadBytes[msg.ID]
202
203				if len(content) < readBytes {
204					slog.Error("Non-interactive: message content is shorter than read bytes", "message_length", len(content), "read_bytes", readBytes)
205					return fmt.Errorf("message content is shorter than read bytes: %d < %d", len(content), readBytes)
206				}
207
208				part := content[readBytes:]
209				fmt.Fprint(output, part)
210				messageReadBytes[msg.ID] = len(content)
211			}
212
213		case <-ctx.Done():
214			stopSpinner()
215			return ctx.Err()
216		}
217	}
218}
219
220func (app *App) UpdateAgentModel(ctx context.Context) error {
221	return app.AgentCoordinator.UpdateModels(ctx)
222}
223
224func (app *App) setupEvents() {
225	ctx, cancel := context.WithCancel(app.globalCtx)
226	app.eventsCtx = ctx
227	setupSubscriber(ctx, app.serviceEventsWG, "sessions", app.Sessions.Subscribe, app.events)
228	setupSubscriber(ctx, app.serviceEventsWG, "messages", app.Messages.Subscribe, app.events)
229	setupSubscriber(ctx, app.serviceEventsWG, "permissions", app.Permissions.Subscribe, app.events)
230	setupSubscriber(ctx, app.serviceEventsWG, "permissions-notifications", app.Permissions.SubscribeNotifications, app.events)
231	setupSubscriber(ctx, app.serviceEventsWG, "history", app.History.Subscribe, app.events)
232	setupSubscriber(ctx, app.serviceEventsWG, "mcp", tools.SubscribeMCPEvents, app.events)
233	setupSubscriber(ctx, app.serviceEventsWG, "lsp", SubscribeLSPEvents, app.events)
234	cleanupFunc := func() error {
235		cancel()
236		app.serviceEventsWG.Wait()
237		return nil
238	}
239	app.cleanupFuncs = append(app.cleanupFuncs, cleanupFunc)
240}
241
242func setupSubscriber[T any](
243	ctx context.Context,
244	wg *sync.WaitGroup,
245	name string,
246	subscriber func(context.Context) <-chan pubsub.Event[T],
247	outputCh chan<- tea.Msg,
248) {
249	wg.Go(func() {
250		subCh := subscriber(ctx)
251		for {
252			select {
253			case event, ok := <-subCh:
254				if !ok {
255					slog.Debug("subscription channel closed", "name", name)
256					return
257				}
258				var msg tea.Msg = event
259				select {
260				case outputCh <- msg:
261				case <-time.After(2 * time.Second):
262					slog.Warn("message dropped due to slow consumer", "name", name)
263				case <-ctx.Done():
264					slog.Debug("subscription cancelled", "name", name)
265					return
266				}
267			case <-ctx.Done():
268				slog.Debug("subscription cancelled", "name", name)
269				return
270			}
271		}
272	})
273}
274
275func (app *App) InitCoderAgent(ctx context.Context) error {
276	coderAgentCfg := app.config.Agents[config.AgentCoder]
277	if coderAgentCfg.ID == "" {
278		return fmt.Errorf("coder agent configuration is missing")
279	}
280	var err error
281	app.AgentCoordinator, err = agent.NewCoordinator(
282		ctx,
283		app.config,
284		app.Sessions,
285		app.Messages,
286		app.Permissions,
287		app.History,
288		app.LSPClients,
289	)
290	if err != nil {
291		slog.Error("Failed to create coder agent", "err", err)
292		return err
293	}
294
295	// Add MCP client cleanup to shutdown process
296	app.cleanupFuncs = append(app.cleanupFuncs, tools.CloseMCPClients)
297	return nil
298}
299
300// Subscribe sends events to the TUI as tea.Msgs.
301func (app *App) Subscribe(program *tea.Program) {
302	defer log.RecoverPanic("app.Subscribe", func() {
303		slog.Info("TUI subscription panic: attempting graceful shutdown")
304		program.Quit()
305	})
306
307	app.tuiWG.Add(1)
308	tuiCtx, tuiCancel := context.WithCancel(app.globalCtx)
309	app.cleanupFuncs = append(app.cleanupFuncs, func() error {
310		slog.Debug("Cancelling TUI message handler")
311		tuiCancel()
312		app.tuiWG.Wait()
313		return nil
314	})
315	defer app.tuiWG.Done()
316
317	for {
318		select {
319		case <-tuiCtx.Done():
320			slog.Debug("TUI message handler shutting down")
321			return
322		case msg, ok := <-app.events:
323			if !ok {
324				slog.Debug("TUI message channel closed")
325				return
326			}
327			program.Send(msg)
328		}
329	}
330}
331
332// Shutdown performs a graceful shutdown of the application.
333func (app *App) Shutdown() {
334	if app.AgentCoordinator != nil {
335		app.AgentCoordinator.CancelAll()
336	}
337
338	// Shutdown all LSP clients.
339	for name, client := range app.LSPClients.Seq2() {
340		shutdownCtx, cancel := context.WithTimeout(app.globalCtx, 5*time.Second)
341		if err := client.Close(shutdownCtx); err != nil {
342			slog.Error("Failed to shutdown LSP client", "name", name, "error", err)
343		}
344		cancel()
345	}
346
347	// Call call cleanup functions.
348	for _, cleanup := range app.cleanupFuncs {
349		if cleanup != nil {
350			if err := cleanup(); err != nil {
351				slog.Error("Failed to cleanup app properly on shutdown", "error", err)
352			}
353		}
354	}
355}