95b9a42 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

fix: faster shutdown (#1570)

95b9a42Carlos Alexandro Becker1 year ago

internal/agent/tools/diagnostics.go+1−1

compare with main

unifiedsplit

@@ -94,7 +94,7 @@ func getDiagnostics(filePath string, lsps *csync.Map[string, *lsp.Client]) strin

94 }

96 out := output.String()

97− slog.Info("Diagnostics", "output", out)

97+ slog.Debug("Diagnostics", "output", out)

98 return out

99}

100

internal/agent/tools/mcp/init.go+21−10

@@ -109,14 +109,26 @@ func GetState(name string) (ClientInfo, bool) {

109// Close closes all MCP clients. This should be called during application shutdown.

110func Close() error {

111 var errs []error

112+ var wg sync.WaitGroup

112 for name, session := range sessions.Seq2() {

113 for name, session := range sessions.Seq2() {

114+ wg.Go(func() {

115+ done := make(chan bool, 1)

116+ go func() {

113− if err := session.Close(); err != nil &&

117+ if err := session.Close(); err != nil &&

114− !errors.Is(err, io.EOF) &&

118+ !errors.Is(err, io.EOF) &&

115− !errors.Is(err, context.Canceled) &&

119+ !errors.Is(err, context.Canceled) &&

116− err.Error() != "signal: killed" {

120+ err.Error() != "signal: killed" {

117− errs = append(errs, fmt.Errorf("close mcp: %s: %w", name, err))

121+ errs = append(errs, fmt.Errorf("close mcp: %s: %w", name, err))

122+ }

123+ done <- true

124+ }()

125+ select {

126+ case <-done:

127+ case <-time.After(time.Millisecond * 250):

128+ }

118− }

129+ })

119 }

130 }

131+ wg.Wait()

120 broker.Shutdown()

132 broker.Shutdown()

121 return errors.Join(errs...)

133 return errors.Join(errs...)

122}

134}

@@ -279,9 +291,8 @@ func createSession(ctx context.Context, name string, m config.MCPConfig, resolve

279 })

291 })

280 },

292 },

281 LoggingMessageHandler: func(_ context.Context, req *mcp.LoggingMessageRequest) {

293 LoggingMessageHandler: func(_ context.Context, req *mcp.LoggingMessageRequest) {

282− slog.Info("mcp log", "name", name, "data", req.Params.Data)

294+ slog.Info("MCP log", "name", name, "data", req.Params.Data)

283 },

295 },

284− KeepAlive: time.Minute * 10,

285 },

296 },

286 )

297 )

287

298

@@ -289,14 +300,14 @@ func createSession(ctx context.Context, name string, m config.MCPConfig, resolve

289 if err != nil {

300 if err != nil {

290 err = maybeStdioErr(err, transport)

301 err = maybeStdioErr(err, transport)

291 updateState(name, StateError, maybeTimeoutErr(err, timeout), nil, Counts{})

302 updateState(name, StateError, maybeTimeoutErr(err, timeout), nil, Counts{})

292− slog.Error("error starting mcp client", "error", err, "name", name)

303+ slog.Error("MCP client failed to initialize", "error", err, "name", name)

293 cancel()

304 cancel()

294 cancelTimer.Stop()

305 cancelTimer.Stop()

295 return nil, err

306 return nil, err

296 }

307 }

297

308

298 cancelTimer.Stop()

309 cancelTimer.Stop()

299− slog.Info("Initialized mcp client", "name", name)

310+ slog.Info("MCP client initialized", "name", name)

300 return session, nil

311 return session, nil

301}

312}

302

313

internal/app/app.go+22−10

@@ -370,30 +370,42 @@ func (app *App) Subscribe(program *tea.Program) {

370

371// Shutdown performs a graceful shutdown of the application.

372func (app *App) Shutdown() {

373+ start := time.Now()

374+ defer func() { slog.Info("Shutdown took " + time.Since(start).String()) }()

375+ var wg sync.WaitGroup

373 if app.AgentCoordinator != nil {

376 if app.AgentCoordinator != nil {

377+ wg.Go(func() {

374− app.AgentCoordinator.CancelAll()

378+ app.AgentCoordinator.CancelAll()

379+ })

375 }

380 }

376

381

377 // Kill all background shells.

382 // Kill all background shells.

383+ wg.Go(func() {

378− shell.GetBackgroundShellManager().KillAll()

384+ shell.GetBackgroundShellManager().KillAll()

385+ })

379

386

380 // Shutdown all LSP clients.

387 // Shutdown all LSP clients.

381 for name, client := range app.LSPClients.Seq2() {

388 for name, client := range app.LSPClients.Seq2() {

389+ wg.Go(func() {

382− shutdownCtx, cancel := context.WithTimeout(app.globalCtx, 5*time.Second)

390+ shutdownCtx, cancel := context.WithTimeout(app.globalCtx, 5*time.Second)

391+ defer cancel()

383− if err := client.Close(shutdownCtx); err != nil {

392+ if err := client.Close(shutdownCtx); err != nil {

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

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

385− }

394+ }

386− cancel()

395+ })

387 }

396 }

388

397

389 // Call call cleanup functions.

398 // Call call cleanup functions.

390 for _, cleanup := range app.cleanupFuncs {

399 for _, cleanup := range app.cleanupFuncs {

391 if cleanup != nil {

400 if cleanup != nil {

401+ wg.Go(func() {

392− if err := cleanup(); err != nil {

402+ if err := cleanup(); err != nil {

393− slog.Error("Failed to cleanup app properly on shutdown", "error", err)

403+ slog.Error("Failed to cleanup app properly on shutdown", "error", err)

404+ }

394− }

405+ })

395 }

406 }

396 }

407 }

408+ wg.Wait()

397}

409}

398

410

399// checkForUpdates checks for available updates.

411// checkForUpdates checks for available updates.

internal/app/lsp.go+4−4

@@ -23,11 +23,11 @@ func (app *App) initLSPClients(ctx context.Context) {

24// createAndStartLSPClient creates a new LSP client, initializes it, and starts its workspace watcher

25func (app *App) createAndStartLSPClient(ctx context.Context, name string, config config.LSPConfig) {

26− slog.Info("Creating LSP client", "name", name, "command", config.Command, "fileTypes", config.FileTypes, "args", config.Args)

26+ slog.Debug("Creating LSP client", "name", name, "command", config.Command, "fileTypes", config.FileTypes, "args", config.Args)

28 // Check if any root markers exist in the working directory (config now has defaults)

29 if !lsp.HasRootMarkers(app.config.WorkingDir(), config.RootMarkers) {

30− slog.Info("Skipping LSP client - no root markers found", "name", name, "rootMarkers", config.RootMarkers)

30+ slog.Debug("Skipping LSP client: no root markers found", "name", name, "rootMarkers", config.RootMarkers)

31 updateLSPState(name, lsp.StateDisabled, nil, nil, 0)

32 return

33 }

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

53 // Initialize LSP client.

54 _, err = lspClient.Initialize(initCtx, app.config.WorkingDir())

55 if err != nil {

56− slog.Error("Initialize failed", "name", name, "error", err)

56+ slog.Error("LSP client initialization failed", "name", name, "error", err)

57 updateLSPState(name, lsp.StateError, err, lspClient, 0)

58 lspClient.Close(ctx)

59 return

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

68 updateLSPState(name, lsp.StateError, err, lspClient, 0)

69 } else {

70 // Server reached a ready state scuccessfully.

71− slog.Info("LSP server is ready", "name", name)

71+ slog.Debug("LSP server is ready", "name", name)

72 lspClient.SetServerState(lsp.StateReady)

73 updateLSPState(name, lsp.StateReady, nil, lspClient, 0)

74 }