app.go

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