@@ -10,11 +10,13 @@ import (
1010 "sync"
1111 "time"
1212
13+ tea "github.com/charmbracelet/bubbletea/v2"
1314 "github.com/charmbracelet/crush/internal/config"
1415 "github.com/charmbracelet/crush/internal/db"
1516 "github.com/charmbracelet/crush/internal/format"
1617 "github.com/charmbracelet/crush/internal/history"
1718 "github.com/charmbracelet/crush/internal/llm/agent"
19+ "github.com/charmbracelet/crush/internal/pubsub"
1820
1921 "github.com/charmbracelet/crush/internal/lsp"
2022 "github.com/charmbracelet/crush/internal/message"
@@ -36,9 +38,18 @@ type App struct {
3638
3739 watcherCancelFuncs []context.CancelFunc
3840 cancelFuncsMutex sync.Mutex
39− watcherWG sync.WaitGroup
41+ lspWatcherWG sync.WaitGroup
4042
4143 config *config.Config
44+
45+ serviceEventsWG *sync.WaitGroup
46+ eventsCtx context.Context
47+ events chan tea.Msg
48+ tuiWG *sync.WaitGroup
49+
50+ // global context and cleanup functions
51+ globalCtx context.Context
52+ cleanupFuncs []func()
4253}
4354
4455func 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) {
5364 History: files,
5465 Permissions: permission.NewPermissionService(cfg.WorkingDir()),
5566 LSPClients: make(map[string]*lsp.Client),
56− config: cfg,
67+
68+ globalCtx: ctx,
69+
70+ config: cfg,
71+
72+ events: make(chan tea.Msg, 100),
73+ serviceEventsWG: &sync.WaitGroup{},
74+ tuiWG: &sync.WaitGroup{},
5775 }
5876
77+ app.setupEvents()
78+
5979 // Initialize LSP clients in the background
6080 go app.initLSPClients(ctx)
6181
6282 // TODO: remove the concept of agent config most likely
63− coderAgentCfg := cfg.Agents["coder"]
64− if coderAgentCfg.ID == "" {
65− return nil, fmt.Errorf("coder agent configuration is missing")
66− }
67−
68− var err error
69− app.CoderAgent, err = agent.NewAgent(
70− coderAgentCfg,
71− app.Permissions,
72− app.Sessions,
73− app.Messages,
74− app.History,
75− app.LSPClients,
76− )
77− if err != nil {
78− slog.Error("Failed to create coder agent", "err", err)
79− return nil, err
83+ if cfg.IsConfigured() {
84+ if err := app.InitCoderAgent(); err != nil {
85+ return nil, fmt.Errorf("failed to initialize coder agent: %w", err)
86+ }
87+ } else {
88+ slog.Warn("No agent configuration found")
8089 }
81−
8290 return app, nil
8391}
8492
@@ -146,32 +154,133 @@ func (a *App) RunNonInteractive(ctx context.Context, prompt string, outputFormat
146154 return nil
147155}
148156
157+func (app *App) UpdateAgentModel() error {
158+ return app.CoderAgent.UpdateModel()
159+}
160+
161+func (app *App) setupEvents() {
162+ ctx, cancel := context.WithCancel(app.globalCtx)
163+ app.eventsCtx = ctx
164+ setupSubscriber(ctx, app.serviceEventsWG, "sessions", app.Sessions.Subscribe, app.events)
165+ setupSubscriber(ctx, app.serviceEventsWG, "messages", app.Messages.Subscribe, app.events)
166+ setupSubscriber(ctx, app.serviceEventsWG, "permissions", app.Permissions.Subscribe, app.events)
167+ setupSubscriber(ctx, app.serviceEventsWG, "history", app.History.Subscribe, app.events)
168+ cleanupFunc := func() {
169+ cancel()
170+ app.serviceEventsWG.Wait()
171+ }
172+ app.cleanupFuncs = append(app.cleanupFuncs, cleanupFunc)
173+}
174+
175+func setupSubscriber[T any](
176+ ctx context.Context,
177+ wg *sync.WaitGroup,
178+ name string,
179+ subscriber func(context.Context) <-chan pubsub.Event[T],
180+ outputCh chan<- tea.Msg,
181+) {
182+ wg.Add(1)
183+ go func() {
184+ defer wg.Done()
185+ subCh := subscriber(ctx)
186+ for {
187+ select {
188+ case event, ok := <-subCh:
189+ if !ok {
190+ slog.Debug("subscription channel closed", "name", name)
191+ return
192+ }
193+ var msg tea.Msg = event
194+ select {
195+ case outputCh <- msg:
196+ case <-time.After(2 * time.Second):
197+ slog.Warn("message dropped due to slow consumer", "name", name)
198+ case <-ctx.Done():
199+ slog.Debug("subscription cancelled", "name", name)
200+ return
201+ }
202+ case <-ctx.Done():
203+ slog.Debug("subscription cancelled", "name", name)
204+ return
205+ }
206+ }
207+ }()
208+}
209+
210+func (app *App) InitCoderAgent() error {
211+ coderAgentCfg := app.config.Agents["coder"]
212+ if coderAgentCfg.ID == "" {
213+ return fmt.Errorf("coder agent configuration is missing")
214+ }
215+ var err error
216+ app.CoderAgent, err = agent.NewAgent(
217+ coderAgentCfg,
218+ app.Permissions,
219+ app.Sessions,
220+ app.Messages,
221+ app.History,
222+ app.LSPClients,
223+ )
224+ if err != nil {
225+ slog.Error("Failed to create coder agent", "err", err)
226+ return err
227+ }
228+ setupSubscriber(app.eventsCtx, app.serviceEventsWG, "coderAgent", app.CoderAgent.Subscribe, app.events)
229+ return nil
230+}
231+
232+func (app *App) Subscribe(program *tea.Program) {
233+ app.tuiWG.Add(1)
234+ tuiCtx, tuiCancel := context.WithCancel(app.globalCtx)
235+ app.cleanupFuncs = append(app.cleanupFuncs, func() {
236+ slog.Debug("Cancelling TUI message handler")
237+ tuiCancel()
238+ app.tuiWG.Wait()
239+ })
240+ defer app.tuiWG.Done()
241+ for {
242+ select {
243+ case <-tuiCtx.Done():
244+ slog.Debug("TUI message handler shutting down")
245+ return
246+ case msg, ok := <-app.events:
247+ if !ok {
248+ slog.Debug("TUI message channel closed")
249+ return
250+ }
251+ program.Send(msg)
252+ }
253+ }
254+}
255+
149256// Shutdown performs a clean shutdown of the application
150257func (app *App) Shutdown() {
151− // Cancel all watcher goroutines
152258 app.cancelFuncsMutex.Lock()
153259 for _, cancel := range app.watcherCancelFuncs {
154260 cancel()
155261 }
156262 app.cancelFuncsMutex.Unlock()
157− app.watcherWG.Wait()
263+ app.lspWatcherWG.Wait()
158264
159− // Perform additional cleanup for LSP clients
160265 app.clientsMutex.RLock()
161266 clients := make(map[string]*lsp.Client, len(app.LSPClients))
162267 maps.Copy(clients, app.LSPClients)
163268 app.clientsMutex.RUnlock()
164269
165270 for name, client := range clients {
166− shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
271+ shutdownCtx, cancel := context.WithTimeout(app.globalCtx, 5*time.Second)
167272 if err := client.Shutdown(shutdownCtx); err != nil {
168273 slog.Error("Failed to shutdown LSP client", "name", name, "error", err)
169274 }
170275 cancel()
171276 }
172− app.CoderAgent.CancelAll()
173−}
277+ if app.CoderAgent != nil {
278+ app.CoderAgent.CancelAll()
279+ }
174280
175−func (app *App) UpdateAgentModel() error {
176− return app.CoderAgent.UpdateModel()
281+ for _, cleanup := range app.cleanupFuncs {
282+ if cleanup != nil {
283+ cleanup()
284+ }
285+ }
177286}