@@ -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(),
@@ -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()
+ }
+ }
}