From 20431519f48af6548bd08327145f3c064ec7fc1d Mon Sep 17 00:00:00 2001 From: Kujtim Hoxha Date: Tue, 8 Jul 2025 13:15:53 +0200 Subject: [PATCH] chore: make agent optional --- cmd/root.go | 142 +--------------------------- internal/app/app.go | 165 +++++++++++++++++++++++++++------ internal/app/lsp.go | 4 +- internal/tui/page/chat/chat.go | 2 +- 4 files changed, 143 insertions(+), 170 deletions(-) diff --git a/cmd/root.go b/cmd/root.go index 02ebbd0104e88a59b659c322a0605189b01afaaf..8ce81d4346d77e3bf239442d695c2ea329fdcf0c 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -6,7 +6,6 @@ import ( "io" "log/slog" "os" - "sync" "time" tea "github.com/charmbracelet/bubbletea/v2" @@ -16,7 +15,6 @@ import ( "github.com/charmbracelet/crush/internal/format" "github.com/charmbracelet/crush/internal/llm/agent" "github.com/charmbracelet/crush/internal/log" - "github.com/charmbracelet/crush/internal/pubsub" "github.com/charmbracelet/crush/internal/tui" "github.com/charmbracelet/crush/internal/version" "github.com/charmbracelet/fang" @@ -122,76 +120,17 @@ to assist developers in writing, debugging, and understanding code directly from tea.WithUniformKeyLayout(), ) - // Setup the subscriptions, this will send services events to the TUI - ch, cancelSubs := setupSubscriptions(app, ctx) + go app.Subscribe(program) - // Create a context for the TUI message handler - tuiCtx, tuiCancel := context.WithCancel(ctx) - var tuiWg sync.WaitGroup - tuiWg.Add(1) - - // Set up message handling for the TUI - go func() { - defer tuiWg.Done() - defer log.RecoverPanic("TUI-message-handler", func() { - attemptTUIRecovery(program) - }) - - for { - select { - case <-tuiCtx.Done(): - slog.Info("TUI message handler shutting down") - return - case msg, ok := <-ch: - if !ok { - slog.Info("TUI message channel closed") - return - } - program.Send(msg) - } - } - }() - - // Cleanup function for when the program exits - cleanup := func() { - // Shutdown the app - app.Shutdown() - - // Cancel subscriptions first - cancelSubs() - - // Then cancel TUI message handler - tuiCancel() - - // Wait for TUI message handler to finish - tuiWg.Wait() - - slog.Info("All goroutines cleaned up") - } - - // Run the TUI - result, err := program.Run() - cleanup() - - if err != nil { + if _, err := program.Run(); err != nil { slog.Error(fmt.Sprintf("TUI run error: %v", err)) return fmt.Errorf("TUI error: %v", err) } - - slog.Info(fmt.Sprintf("TUI exited with result: %v", result)) + app.Shutdown() return nil }, } -// attemptTUIRecovery tries to recover the TUI after a panic -func attemptTUIRecovery(program *tea.Program) { - slog.Info("Attempting to recover TUI after panic") - - // We could try to restart the TUI or gracefully exit - // For now, we'll just quit the program to avoid further issues - program.Quit() -} - func initMCPTools(ctx context.Context, app *app.App, cfg *config.Config) { go func() { defer log.RecoverPanic("MCP-goroutine", nil) @@ -206,81 +145,6 @@ func initMCPTools(ctx context.Context, app *app.App, cfg *config.Config) { }() } -func setupSubscriber[T any]( - ctx context.Context, - wg *sync.WaitGroup, - name string, - subscriber func(context.Context) <-chan pubsub.Event[T], - outputCh chan<- tea.Msg, -) { - wg.Add(1) - go func() { - defer wg.Done() - defer log.RecoverPanic(fmt.Sprintf("subscription-%s", name), nil) - - subCh := subscriber(ctx) - - for { - select { - case event, ok := <-subCh: - if !ok { - slog.Info("subscription channel closed", "name", name) - return - } - - var msg tea.Msg = event - - select { - case outputCh <- msg: - case <-time.After(2 * time.Second): - slog.Warn("message dropped due to slow consumer", "name", name) - case <-ctx.Done(): - slog.Info("subscription cancelled", "name", name) - return - } - case <-ctx.Done(): - slog.Info("subscription cancelled", "name", name) - return - } - } - }() -} - -func setupSubscriptions(app *app.App, parentCtx context.Context) (chan tea.Msg, func()) { - ch := make(chan tea.Msg, 100) - - wg := sync.WaitGroup{} - ctx, cancel := context.WithCancel(parentCtx) // Inherit from parent context - - setupSubscriber(ctx, &wg, "sessions", app.Sessions.Subscribe, ch) - setupSubscriber(ctx, &wg, "messages", app.Messages.Subscribe, ch) - setupSubscriber(ctx, &wg, "permissions", app.Permissions.Subscribe, ch) - setupSubscriber(ctx, &wg, "coderAgent", app.CoderAgent.Subscribe, ch) - setupSubscriber(ctx, &wg, "history", app.History.Subscribe, ch) - - cleanupFunc := func() { - slog.Info("Cancelling all subscriptions") - cancel() // Signal all goroutines to stop - - waitCh := make(chan struct{}) - go func() { - defer log.RecoverPanic("subscription-cleanup", nil) - wg.Wait() - close(waitCh) - }() - - select { - case <-waitCh: - slog.Info("All subscription goroutines completed successfully") - close(ch) // Only close after all writers are confirmed done - case <-time.After(5 * time.Second): - slog.Warn("Timed out waiting for some subscription goroutines to complete") - close(ch) - } - } - return ch, cleanupFunc -} - func Execute() { if err := fang.Execute( context.Background(), diff --git a/internal/app/app.go b/internal/app/app.go index da014df81367665caf1df793760e3d832c223648..c06f859e5fa083bde55bd2c4e5d07036190be8bf 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -10,11 +10,13 @@ import ( "sync" "time" + tea "github.com/charmbracelet/bubbletea/v2" "github.com/charmbracelet/crush/internal/config" "github.com/charmbracelet/crush/internal/db" "github.com/charmbracelet/crush/internal/format" "github.com/charmbracelet/crush/internal/history" "github.com/charmbracelet/crush/internal/llm/agent" + "github.com/charmbracelet/crush/internal/pubsub" "github.com/charmbracelet/crush/internal/lsp" "github.com/charmbracelet/crush/internal/message" @@ -36,9 +38,18 @@ type App struct { watcherCancelFuncs []context.CancelFunc cancelFuncsMutex sync.Mutex - watcherWG sync.WaitGroup + lspWatcherWG sync.WaitGroup config *config.Config + + serviceEventsWG *sync.WaitGroup + eventsCtx context.Context + events chan tea.Msg + tuiWG *sync.WaitGroup + + // global context and cleanup functions + globalCtx context.Context + cleanupFuncs []func() } func New(ctx context.Context, conn *sql.DB, cfg *config.Config) (*App, error) { @@ -53,32 +64,29 @@ func New(ctx context.Context, conn *sql.DB, cfg *config.Config) (*App, error) { History: files, Permissions: permission.NewPermissionService(cfg.WorkingDir()), LSPClients: make(map[string]*lsp.Client), - config: cfg, + + globalCtx: ctx, + + config: cfg, + + events: make(chan tea.Msg, 100), + serviceEventsWG: &sync.WaitGroup{}, + tuiWG: &sync.WaitGroup{}, } + app.setupEvents() + // Initialize LSP clients in the background go app.initLSPClients(ctx) // TODO: remove the concept of agent config most likely - coderAgentCfg := cfg.Agents["coder"] - if coderAgentCfg.ID == "" { - return nil, fmt.Errorf("coder agent configuration is missing") - } - - var err error - app.CoderAgent, err = agent.NewAgent( - coderAgentCfg, - app.Permissions, - app.Sessions, - app.Messages, - app.History, - app.LSPClients, - ) - if err != nil { - slog.Error("Failed to create coder agent", "err", err) - return nil, err + if cfg.IsConfigured() { + if err := app.InitCoderAgent(); err != nil { + return nil, fmt.Errorf("failed to initialize coder agent: %w", err) + } + } else { + slog.Warn("No agent configuration found") } - return app, nil } @@ -146,32 +154,133 @@ func (a *App) RunNonInteractive(ctx context.Context, prompt string, outputFormat return nil } +func (app *App) UpdateAgentModel() error { + return app.CoderAgent.UpdateModel() +} + +func (app *App) setupEvents() { + ctx, cancel := context.WithCancel(app.globalCtx) + app.eventsCtx = ctx + setupSubscriber(ctx, app.serviceEventsWG, "sessions", app.Sessions.Subscribe, app.events) + setupSubscriber(ctx, app.serviceEventsWG, "messages", app.Messages.Subscribe, app.events) + setupSubscriber(ctx, app.serviceEventsWG, "permissions", app.Permissions.Subscribe, app.events) + setupSubscriber(ctx, app.serviceEventsWG, "history", app.History.Subscribe, app.events) + cleanupFunc := func() { + cancel() + app.serviceEventsWG.Wait() + } + app.cleanupFuncs = append(app.cleanupFuncs, cleanupFunc) +} + +func setupSubscriber[T any]( + ctx context.Context, + wg *sync.WaitGroup, + name string, + subscriber func(context.Context) <-chan pubsub.Event[T], + outputCh chan<- tea.Msg, +) { + wg.Add(1) + go func() { + defer wg.Done() + subCh := subscriber(ctx) + for { + select { + case event, ok := <-subCh: + if !ok { + slog.Debug("subscription channel closed", "name", name) + return + } + var msg tea.Msg = event + select { + case outputCh <- msg: + case <-time.After(2 * time.Second): + slog.Warn("message dropped due to slow consumer", "name", name) + case <-ctx.Done(): + slog.Debug("subscription cancelled", "name", name) + return + } + case <-ctx.Done(): + slog.Debug("subscription cancelled", "name", name) + return + } + } + }() +} + +func (app *App) InitCoderAgent() error { + coderAgentCfg := app.config.Agents["coder"] + if coderAgentCfg.ID == "" { + return fmt.Errorf("coder agent configuration is missing") + } + var err error + app.CoderAgent, err = agent.NewAgent( + coderAgentCfg, + app.Permissions, + app.Sessions, + app.Messages, + app.History, + app.LSPClients, + ) + if err != nil { + slog.Error("Failed to create coder agent", "err", err) + return err + } + setupSubscriber(app.eventsCtx, app.serviceEventsWG, "coderAgent", app.CoderAgent.Subscribe, app.events) + return nil +} + +func (app *App) Subscribe(program *tea.Program) { + app.tuiWG.Add(1) + tuiCtx, tuiCancel := context.WithCancel(app.globalCtx) + app.cleanupFuncs = append(app.cleanupFuncs, func() { + slog.Debug("Cancelling TUI message handler") + tuiCancel() + app.tuiWG.Wait() + }) + defer app.tuiWG.Done() + for { + select { + case <-tuiCtx.Done(): + slog.Debug("TUI message handler shutting down") + return + case msg, ok := <-app.events: + if !ok { + slog.Debug("TUI message channel closed") + return + } + program.Send(msg) + } + } +} + // Shutdown performs a clean shutdown of the application func (app *App) Shutdown() { - // Cancel all watcher goroutines app.cancelFuncsMutex.Lock() for _, cancel := range app.watcherCancelFuncs { cancel() } app.cancelFuncsMutex.Unlock() - app.watcherWG.Wait() + app.lspWatcherWG.Wait() - // Perform additional cleanup for LSP clients app.clientsMutex.RLock() clients := make(map[string]*lsp.Client, len(app.LSPClients)) maps.Copy(clients, app.LSPClients) app.clientsMutex.RUnlock() for name, client := range clients { - shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + shutdownCtx, cancel := context.WithTimeout(app.globalCtx, 5*time.Second) if err := client.Shutdown(shutdownCtx); err != nil { slog.Error("Failed to shutdown LSP client", "name", name, "error", err) } cancel() } - app.CoderAgent.CancelAll() -} + if app.CoderAgent != nil { + app.CoderAgent.CancelAll() + } -func (app *App) UpdateAgentModel() error { - return app.CoderAgent.UpdateModel() + for _, cleanup := range app.cleanupFuncs { + if cleanup != nil { + cleanup() + } + } } diff --git a/internal/app/lsp.go b/internal/app/lsp.go index 1777a653a4153dc42cc87444f6122df01e82cedd..ba98d4b3a074c2e9abcef87eb3030a21be669eab 100644 --- a/internal/app/lsp.go +++ b/internal/app/lsp.go @@ -71,7 +71,7 @@ func (app *App) createAndStartLSPClient(ctx context.Context, name string, comman app.cancelFuncsMutex.Unlock() // Add the watcher to a WaitGroup to track active goroutines - app.watcherWG.Add(1) + app.lspWatcherWG.Add(1) // Add to map with mutex protection before starting goroutine app.clientsMutex.Lock() @@ -83,7 +83,7 @@ func (app *App) createAndStartLSPClient(ctx context.Context, name string, comman // runWorkspaceWatcher executes the workspace watcher for an LSP client func (app *App) runWorkspaceWatcher(ctx context.Context, name string, workspaceWatcher *watcher.WorkspaceWatcher) { - defer app.watcherWG.Done() + defer app.lspWatcherWG.Done() defer log.RecoverPanic("LSP-"+name, func() { // Try to restart the client app.restartLSPClient(ctx, name) diff --git a/internal/tui/page/chat/chat.go b/internal/tui/page/chat/chat.go index 73eeec6895e0591a37978c6cbc8c41fd34e74164..fc529bba8b2be2bbbab1e5c7dd5668bcf7923ab8 100644 --- a/internal/tui/page/chat/chat.go +++ b/internal/tui/page/chat/chat.go @@ -511,7 +511,7 @@ func (p *chatPage) Bindings() []key.Binding { p.keyMap.NewSession, p.keyMap.AddAttachment, } - if p.app.CoderAgent.IsBusy() { + if p.app.CoderAgent != nil && p.app.CoderAgent.IsBusy() { cancelBinding := p.keyMap.Cancel if p.canceling { cancelBinding = key.NewBinding(