app.go

  1package app
  2
  3import (
  4	"context"
  5	"database/sql"
  6	"errors"
  7	"fmt"
  8	"log/slog"
  9	"sync"
 10	"time"
 11
 12	tea "github.com/charmbracelet/bubbletea/v2"
 13	"github.com/charmbracelet/crush/internal/agent"
 14	"github.com/charmbracelet/crush/internal/agent/tools"
 15	"github.com/charmbracelet/crush/internal/config"
 16	"github.com/charmbracelet/crush/internal/csync"
 17	"github.com/charmbracelet/crush/internal/db"
 18	"github.com/charmbracelet/crush/internal/format"
 19	"github.com/charmbracelet/crush/internal/history"
 20	"github.com/charmbracelet/crush/internal/log"
 21	"github.com/charmbracelet/crush/internal/pubsub"
 22	"github.com/charmbracelet/fantasy/ai"
 23
 24	"github.com/charmbracelet/crush/internal/lsp"
 25	"github.com/charmbracelet/crush/internal/message"
 26	"github.com/charmbracelet/crush/internal/permission"
 27	"github.com/charmbracelet/crush/internal/session"
 28)
 29
 30type App struct {
 31	Sessions    session.Service
 32	Messages    message.Service
 33	History     history.Service
 34	Permissions permission.Service
 35
 36	// TODO: (kujtim) remove this when fully implemented
 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		if err := app.InitCoderAgent(); err != nil {
 92			return nil, fmt.Errorf("failed to initialize coder agent: %w", err)
 93		}
 94	} else {
 95		slog.Warn("No agent configuration found")
 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	// Start spinner if not in quiet mode.
114	var spinner *format.Spinner
115	if !quiet {
116		spinner = format.NewSpinner(ctx, cancel, "Generating")
117		spinner.Start()
118	}
119
120	// Helper function to stop spinner once.
121	stopSpinner := func() {
122		if !quiet && spinner != nil {
123			spinner.Stop()
124			spinner = nil
125		}
126	}
127	defer stopSpinner()
128
129	const maxPromptLengthForTitle = 100
130	titlePrefix := "Non-interactive: "
131	var titleSuffix string
132
133	if len(prompt) > maxPromptLengthForTitle {
134		titleSuffix = prompt[:maxPromptLengthForTitle] + "..."
135	} else {
136		titleSuffix = prompt
137	}
138	title := titlePrefix + titleSuffix
139
140	sess, err := app.Sessions.Create(ctx, title)
141	if err != nil {
142		return fmt.Errorf("failed to create session for non-interactive mode: %w", err)
143	}
144	slog.Info("Created session for non-interactive run", "session_id", sess.ID)
145
146	// Automatically approve all permission requests for this non-interactive session
147	app.Permissions.AutoApproveSession(sess.ID)
148
149	type response struct {
150		result *ai.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	for {
171		select {
172		case result := <-done:
173			stopSpinner()
174			if result.err != nil {
175				if errors.Is(result.err, context.Canceled) || errors.Is(result.err, agent.ErrRequestCancelled) {
176					slog.Info("Non-interactive: agent processing cancelled", "session_id", sess.ID)
177					return nil
178				}
179				return fmt.Errorf("agent processing failed: %w", result.err)
180			}
181			return nil
182
183		case event := <-messageEvents:
184			msg := event.Payload
185			if msg.SessionID == sess.ID && msg.Role == message.Assistant && len(msg.Parts) > 0 {
186				stopSpinner()
187
188				content := msg.Content().String()
189				readBytes := messageReadBytes[msg.ID]
190
191				if len(content) < readBytes {
192					slog.Error("Non-interactive: message content is shorter than read bytes", "message_length", len(content), "read_bytes", readBytes)
193					return fmt.Errorf("message content is shorter than read bytes: %d < %d", len(content), readBytes)
194				}
195
196				part := content[readBytes:]
197				fmt.Print(part)
198				messageReadBytes[msg.ID] = len(content)
199			}
200
201		case <-ctx.Done():
202			stopSpinner()
203			return ctx.Err()
204		}
205	}
206}
207
208func (app *App) UpdateAgentModel() error {
209	return app.AgentCoordinator.UpdateModels()
210}
211
212func (app *App) setupEvents() {
213	ctx, cancel := context.WithCancel(app.globalCtx)
214	app.eventsCtx = ctx
215	setupSubscriber(ctx, app.serviceEventsWG, "sessions", app.Sessions.Subscribe, app.events)
216	setupSubscriber(ctx, app.serviceEventsWG, "messages", app.Messages.Subscribe, app.events)
217	setupSubscriber(ctx, app.serviceEventsWG, "permissions", app.Permissions.Subscribe, app.events)
218	setupSubscriber(ctx, app.serviceEventsWG, "permissions-notifications", app.Permissions.SubscribeNotifications, app.events)
219	setupSubscriber(ctx, app.serviceEventsWG, "history", app.History.Subscribe, app.events)
220	setupSubscriber(ctx, app.serviceEventsWG, "mcp", tools.SubscribeMCPEvents, app.events)
221	setupSubscriber(ctx, app.serviceEventsWG, "lsp", SubscribeLSPEvents, app.events)
222	cleanupFunc := func() error {
223		cancel()
224		app.serviceEventsWG.Wait()
225		return nil
226	}
227	app.cleanupFuncs = append(app.cleanupFuncs, cleanupFunc)
228}
229
230func setupSubscriber[T any](
231	ctx context.Context,
232	wg *sync.WaitGroup,
233	name string,
234	subscriber func(context.Context) <-chan pubsub.Event[T],
235	outputCh chan<- tea.Msg,
236) {
237	wg.Go(func() {
238		subCh := subscriber(ctx)
239		for {
240			select {
241			case event, ok := <-subCh:
242				if !ok {
243					slog.Debug("subscription channel closed", "name", name)
244					return
245				}
246				var msg tea.Msg = event
247				select {
248				case outputCh <- msg:
249				case <-time.After(2 * time.Second):
250					slog.Warn("message dropped due to slow consumer", "name", name)
251				case <-ctx.Done():
252					slog.Debug("subscription cancelled", "name", name)
253					return
254				}
255			case <-ctx.Done():
256				slog.Debug("subscription cancelled", "name", name)
257				return
258			}
259		}
260	})
261}
262
263func (app *App) InitCoderAgent() error {
264	coderAgentCfg := app.config.Agents[config.AgentCoder]
265	if coderAgentCfg.ID == "" {
266		return fmt.Errorf("coder agent configuration is missing")
267	}
268	var err error
269	app.AgentCoordinator, err = agent.NewCoordinator(
270		app.config,
271		app.Sessions,
272		app.Messages,
273		app.Permissions,
274		app.History,
275		app.LSPClients,
276	)
277	if err != nil {
278		slog.Error("Failed to create coder agent", "err", err)
279		return err
280	}
281
282	// Add MCP client cleanup to shutdown process
283	app.cleanupFuncs = append(app.cleanupFuncs, tools.CloseMCPClients)
284	return nil
285}
286
287// Subscribe sends events to the TUI as tea.Msgs.
288func (app *App) Subscribe(program *tea.Program) {
289	defer log.RecoverPanic("app.Subscribe", func() {
290		slog.Info("TUI subscription panic: attempting graceful shutdown")
291		program.Quit()
292	})
293
294	app.tuiWG.Add(1)
295	tuiCtx, tuiCancel := context.WithCancel(app.globalCtx)
296	app.cleanupFuncs = append(app.cleanupFuncs, func() error {
297		slog.Debug("Cancelling TUI message handler")
298		tuiCancel()
299		app.tuiWG.Wait()
300		return nil
301	})
302	defer app.tuiWG.Done()
303
304	for {
305		select {
306		case <-tuiCtx.Done():
307			slog.Debug("TUI message handler shutting down")
308			return
309		case msg, ok := <-app.events:
310			if !ok {
311				slog.Debug("TUI message channel closed")
312				return
313			}
314			program.Send(msg)
315		}
316	}
317}
318
319// Shutdown performs a graceful shutdown of the application.
320func (app *App) Shutdown() {
321	if app.AgentCoordinator != nil {
322		app.AgentCoordinator.CancelAll()
323	}
324
325	// Shutdown all LSP clients.
326	for name, client := range app.LSPClients.Seq2() {
327		shutdownCtx, cancel := context.WithTimeout(app.globalCtx, 5*time.Second)
328		if err := client.Close(shutdownCtx); err != nil {
329			slog.Error("Failed to shutdown LSP client", "name", name, "error", err)
330		}
331		cancel()
332	}
333
334	// Call call cleanup functions.
335	for _, cleanup := range app.cleanupFuncs {
336		if cleanup != nil {
337			if err := cleanup(); err != nil {
338				slog.Error("Failed to cleanup app properly on shutdown", "error", err)
339			}
340		}
341	}
342}