2043151 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

chore: make agent optional

2043151Kujtim Hoxha1 year ago

cmd/root.go+3−139

compare with main

unifiedsplit

@@ -6,7 +6,6 @@ import (

6 "io"

7 "log/slog"

8 "os"

9− "sync"

10 "time"

9 "time"

12 tea "github.com/charmbracelet/bubbletea/v2"

11 tea "github.com/charmbracelet/bubbletea/v2"

@@ -16,7 +15,6 @@ import (

16 "github.com/charmbracelet/crush/internal/format"

15 "github.com/charmbracelet/crush/internal/format"

17 "github.com/charmbracelet/crush/internal/llm/agent"

16 "github.com/charmbracelet/crush/internal/llm/agent"

18 "github.com/charmbracelet/crush/internal/log"

17 "github.com/charmbracelet/crush/internal/log"

19− "github.com/charmbracelet/crush/internal/pubsub"

20 "github.com/charmbracelet/crush/internal/tui"

18 "github.com/charmbracelet/crush/internal/tui"

21 "github.com/charmbracelet/crush/internal/version"

19 "github.com/charmbracelet/crush/internal/version"

22 "github.com/charmbracelet/fang"

20 "github.com/charmbracelet/fang"

@@ -122,76 +120,17 @@ to assist developers in writing, debugging, and understanding code directly from

122 tea.WithUniformKeyLayout(),

120 tea.WithUniformKeyLayout(),

123 )

121 )

124

122

125− // Setup the subscriptions, this will send services events to the TUI

123+ go app.Subscribe(program)

126− ch, cancelSubs := setupSubscriptions(app, ctx)

127

124

128− // Create a context for the TUI message handler

129− tuiCtx, tuiCancel := context.WithCancel(ctx)

130− var tuiWg sync.WaitGroup

131− tuiWg.Add(1)

132−

133− // Set up message handling for the TUI

134− go func() {

135− defer tuiWg.Done()

136− defer log.RecoverPanic("TUI-message-handler", func() {

137− attemptTUIRecovery(program)

138− })

139−

140− for {

141− select {

142− case <-tuiCtx.Done():

143− slog.Info("TUI message handler shutting down")

144− return

145− case msg, ok := <-ch:

146− if !ok {

147− slog.Info("TUI message channel closed")

148− return

149− }

150− program.Send(msg)

151− }

152− }

153− }()

154−

155− // Cleanup function for when the program exits

156− cleanup := func() {

157− // Shutdown the app

158− app.Shutdown()

159−

160− // Cancel subscriptions first

161− cancelSubs()

162−

163− // Then cancel TUI message handler

164− tuiCancel()

165−

166− // Wait for TUI message handler to finish

167− tuiWg.Wait()

168−

169− slog.Info("All goroutines cleaned up")

170− }

171−

172− // Run the TUI

173− result, err := program.Run()

125+ if _, err := program.Run(); err != nil {

174− cleanup()

175−

176− if err != nil {

177 slog.Error(fmt.Sprintf("TUI run error: %v", err))

126 slog.Error(fmt.Sprintf("TUI run error: %v", err))

178 return fmt.Errorf("TUI error: %v", err)

127 return fmt.Errorf("TUI error: %v", err)

179 }

128 }

180−

129+ app.Shutdown()

181− slog.Info(fmt.Sprintf("TUI exited with result: %v", result))

182 return nil

130 return nil

183 },

131 },

184}

132}

185

133

186−// attemptTUIRecovery tries to recover the TUI after a panic

187−func attemptTUIRecovery(program *tea.Program) {

188− slog.Info("Attempting to recover TUI after panic")

189−

190− // We could try to restart the TUI or gracefully exit

191− // For now, we'll just quit the program to avoid further issues

192− program.Quit()

193−}

194−

195func initMCPTools(ctx context.Context, app *app.App, cfg *config.Config) {

134func initMCPTools(ctx context.Context, app *app.App, cfg *config.Config) {

196 go func() {

135 go func() {

197 defer log.RecoverPanic("MCP-goroutine", nil)

136 defer log.RecoverPanic("MCP-goroutine", nil)

@@ -206,81 +145,6 @@ func initMCPTools(ctx context.Context, app *app.App, cfg *config.Config) {

206 }()

145 }()

207}

146}

208

147

209−func setupSubscriber[T any](

210− ctx context.Context,

211− wg *sync.WaitGroup,

212− name string,

213− subscriber func(context.Context) <-chan pubsub.Event[T],

214− outputCh chan<- tea.Msg,

215−) {

216− wg.Add(1)

217− go func() {

218− defer wg.Done()

219− defer log.RecoverPanic(fmt.Sprintf("subscription-%s", name), nil)

220−

221− subCh := subscriber(ctx)

222−

223− for {

224− select {

225− case event, ok := <-subCh:

226− if !ok {

227− slog.Info("subscription channel closed", "name", name)

228− return

229− }

230−

231− var msg tea.Msg = event

232−

233− select {

234− case outputCh <- msg:

235− case <-time.After(2 * time.Second):

236− slog.Warn("message dropped due to slow consumer", "name", name)

237− case <-ctx.Done():

238− slog.Info("subscription cancelled", "name", name)

239− return

240− }

241− case <-ctx.Done():

242− slog.Info("subscription cancelled", "name", name)

243− return

244− }

245− }

246− }()

247−}

248−

249−func setupSubscriptions(app *app.App, parentCtx context.Context) (chan tea.Msg, func()) {

250− ch := make(chan tea.Msg, 100)

251−

252− wg := sync.WaitGroup{}

253− ctx, cancel := context.WithCancel(parentCtx) // Inherit from parent context

254−

255− setupSubscriber(ctx, &wg, "sessions", app.Sessions.Subscribe, ch)

256− setupSubscriber(ctx, &wg, "messages", app.Messages.Subscribe, ch)

257− setupSubscriber(ctx, &wg, "permissions", app.Permissions.Subscribe, ch)

258− setupSubscriber(ctx, &wg, "coderAgent", app.CoderAgent.Subscribe, ch)

259− setupSubscriber(ctx, &wg, "history", app.History.Subscribe, ch)

260−

261− cleanupFunc := func() {

262− slog.Info("Cancelling all subscriptions")

263− cancel() // Signal all goroutines to stop

264−

265− waitCh := make(chan struct{})

266− go func() {

267− defer log.RecoverPanic("subscription-cleanup", nil)

268− wg.Wait()

269− close(waitCh)

270− }()

271−

272− select {

273− case <-waitCh:

274− slog.Info("All subscription goroutines completed successfully")

275− close(ch) // Only close after all writers are confirmed done

276− case <-time.After(5 * time.Second):

277− slog.Warn("Timed out waiting for some subscription goroutines to complete")

278− close(ch)

279− }

280− }

281− return ch, cleanupFunc

282−}

283−

284func Execute() {

148func Execute() {

285 if err := fang.Execute(

149 if err := fang.Execute(

286 context.Background(),

150 context.Background(),

internal/app/app.go+137−28

@@ -10,11 +10,13 @@ import (

10 "sync"

11 "time"

13+ tea "github.com/charmbracelet/bubbletea/v2"

13 "github.com/charmbracelet/crush/internal/config"

14 "github.com/charmbracelet/crush/internal/config"

14 "github.com/charmbracelet/crush/internal/db"

15 "github.com/charmbracelet/crush/internal/db"

15 "github.com/charmbracelet/crush/internal/format"

16 "github.com/charmbracelet/crush/internal/format"

16 "github.com/charmbracelet/crush/internal/history"

17 "github.com/charmbracelet/crush/internal/history"

17 "github.com/charmbracelet/crush/internal/llm/agent"

18 "github.com/charmbracelet/crush/internal/llm/agent"

19+ "github.com/charmbracelet/crush/internal/pubsub"

19 "github.com/charmbracelet/crush/internal/lsp"

21 "github.com/charmbracelet/crush/internal/lsp"

20 "github.com/charmbracelet/crush/internal/message"

22 "github.com/charmbracelet/crush/internal/message"

@@ -36,9 +38,18 @@ type App struct {

37 watcherCancelFuncs []context.CancelFunc

39 watcherCancelFuncs []context.CancelFunc

38 cancelFuncsMutex sync.Mutex

40 cancelFuncsMutex sync.Mutex

39− watcherWG sync.WaitGroup

41+ lspWatcherWG sync.WaitGroup

41 config *config.Config

43 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()

42}

53}

44func New(ctx context.Context, conn *sql.DB, cfg *config.Config) (*App, error) {

55func 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) {

53 History: files,

64 History: files,

54 Permissions: permission.NewPermissionService(cfg.WorkingDir()),

65 Permissions: permission.NewPermissionService(cfg.WorkingDir()),

55 LSPClients: make(map[string]*lsp.Client),

66 LSPClients: make(map[string]*lsp.Client),

67+

68+ globalCtx: ctx,

69+

56− config: cfg,

70+ config: cfg,

71+

72+ events: make(chan tea.Msg, 100),

73+ serviceEventsWG: &sync.WaitGroup{},

74+ tuiWG: &sync.WaitGroup{},

57 }

75 }

77+ app.setupEvents()

78+

59 // Initialize LSP clients in the background

79 // Initialize LSP clients in the background

60 go app.initLSPClients(ctx)

80 go app.initLSPClients(ctx)

62 // TODO: remove the concept of agent config most likely

82 // TODO: remove the concept of agent config most likely

63− coderAgentCfg := cfg.Agents["coder"]

83+ if cfg.IsConfigured() {

64− if coderAgentCfg.ID == "" {

84+ if err := app.InitCoderAgent(); err != nil {

65− return nil, fmt.Errorf("coder agent configuration is missing")

85+ return nil, fmt.Errorf("failed to initialize coder agent: %w", err)

66− }

86+ }

67−

87+ } else {

68− var err error

88+ slog.Warn("No agent configuration found")

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

80 }

89 }

81−

82 return app, nil

90 return app, nil

83}

91}

@@ -146,32 +154,133 @@ func (a *App) RunNonInteractive(ctx context.Context, prompt string, outputFormat

146 return nil

154 return nil

147}

155}

148

156

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+

149// Shutdown performs a clean shutdown of the application

256// Shutdown performs a clean shutdown of the application

150func (app *App) Shutdown() {

257func (app *App) Shutdown() {

151− // Cancel all watcher goroutines

152 app.cancelFuncsMutex.Lock()

258 app.cancelFuncsMutex.Lock()

153 for _, cancel := range app.watcherCancelFuncs {

259 for _, cancel := range app.watcherCancelFuncs {

154 cancel()

260 cancel()

155 }

261 }

156 app.cancelFuncsMutex.Unlock()

262 app.cancelFuncsMutex.Unlock()

157− app.watcherWG.Wait()

263+ app.lspWatcherWG.Wait()

158

264

159− // Perform additional cleanup for LSP clients

160 app.clientsMutex.RLock()

265 app.clientsMutex.RLock()

161 clients := make(map[string]*lsp.Client, len(app.LSPClients))

266 clients := make(map[string]*lsp.Client, len(app.LSPClients))

162 maps.Copy(clients, app.LSPClients)

267 maps.Copy(clients, app.LSPClients)

163 app.clientsMutex.RUnlock()

268 app.clientsMutex.RUnlock()

164

269

165 for name, client := range clients {

270 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)

167 if err := client.Shutdown(shutdownCtx); err != nil {

272 if err := client.Shutdown(shutdownCtx); err != nil {

168 slog.Error("Failed to shutdown LSP client", "name", name, "error", err)

273 slog.Error("Failed to shutdown LSP client", "name", name, "error", err)

169 }

274 }

170 cancel()

275 cancel()

171 }

276 }

277+ if app.CoderAgent != nil {

172− app.CoderAgent.CancelAll()

278+ app.CoderAgent.CancelAll()

173−}

279+ }

174

280

175−func (app *App) UpdateAgentModel() error {

281+ for _, cleanup := range app.cleanupFuncs {

176− return app.CoderAgent.UpdateModel()

282+ if cleanup != nil {

283+ cleanup()

284+ }

285+ }

177}

286}

internal/app/lsp.go+2−2

@@ -71,7 +71,7 @@ func (app *App) createAndStartLSPClient(ctx context.Context, name string, comman

71 app.cancelFuncsMutex.Unlock()

73 // Add the watcher to a WaitGroup to track active goroutines

74− app.watcherWG.Add(1)

74+ app.lspWatcherWG.Add(1)

76 // Add to map with mutex protection before starting goroutine

77 app.clientsMutex.Lock()

@@ -83,7 +83,7 @@ func (app *App) createAndStartLSPClient(ctx context.Context, name string, comman

84// runWorkspaceWatcher executes the workspace watcher for an LSP client

85func (app *App) runWorkspaceWatcher(ctx context.Context, name string, workspaceWatcher *watcher.WorkspaceWatcher) {

86− defer app.watcherWG.Done()

86+ defer app.lspWatcherWG.Done()

87 defer log.RecoverPanic("LSP-"+name, func() {

88 // Try to restart the client

89 app.restartLSPClient(ctx, name)

internal/tui/page/chat/chat.go+1−1

@@ -511,7 +511,7 @@ func (p *chatPage) Bindings() []key.Binding {

511 p.keyMap.NewSession,

512 p.keyMap.AddAttachment,

513 }

514− if p.app.CoderAgent.IsBusy() {

514+ if p.app.CoderAgent != nil && p.app.CoderAgent.IsBusy() {

515 cancelBinding := p.keyMap.Cancel

516 if p.canceling {

517 cancelBinding = key.NewBinding(