on main help sign in
fix: improving shutdown (#2175)

* test: use t.Context() and synctest Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com> * fix: passing down context to all shutdown funcs Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com> * perf(lsp): kill all clients on shutdown Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com> * fix: exit posthog earlier Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com> * test: fix dirs test Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com> * test: fix projects test Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com> * test: fix race Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com> * Update internal/lsp/manager.go Co-authored-by: Andrey Nering <andreynering@users.noreply.github.com> * fix: cleanup Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com> * test: race Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com> --------- Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com> Co-authored-by: Andrey Nering <andreynering@users.noreply.github.com>

commit63e009898a0fc6f2f1aa29a2f375ab4ebd579061 authorCarlos Alexandro Becker <caarlos0@users.noreply.github.com> dateMon 9 Feb 2026, 11:32 (8 months ago) parent722093b
10 files changed+68−28
go.mod+1−1 go.sum+2−2 internal/agent/tools/mcp/init.go+1−4 internal/app/app.go+16−6 internal/cmd/dirs_test.go+2 internal/cmd/root.go−3 internal/lsp/client.go+4−1 internal/lsp/manager.go+22 internal/projects/projects_test.go+6 internal/shell/background_test.go+14−11
go.mod+1−1
@@ -31,7 +31,7 @@ require (
3131 github.com/charmbracelet/x/exp/ordered v0.1.0
3232 github.com/charmbracelet/x/exp/slice v0.0.0-20251201173703-9f73bfd934ff
3333 github.com/charmbracelet/x/exp/strings v0.1.0
34− github.com/charmbracelet/x/powernap v0.0.0-20260127155452-b72a9a918687
34+ github.com/charmbracelet/x/powernap v0.0.0-20260209132835-6b065b8ba62c
3535 github.com/charmbracelet/x/term v0.2.2
3636 github.com/clipperhouse/displaywidth v0.10.0
3737 github.com/clipperhouse/uax29/v2 v2.6.0
go.sum+2−2
@@ -122,8 +122,8 @@ github.com/charmbracelet/x/exp/strings v0.1.0 h1:i69S2XI7uG1u4NLGeJPSYU++Nmjvpo9
122122github.com/charmbracelet/x/exp/strings v0.1.0/go.mod h1:/ehtMPNh9K4odGFkqYJKpIYyePhdp1hLBRvyY4bWkH8=
123123github.com/charmbracelet/x/json v0.2.0 h1:DqB+ZGx2h+Z+1s98HOuOyli+i97wsFQIxP2ZQANTPrQ=
124124github.com/charmbracelet/x/json v0.2.0/go.mod h1:opFIflx2YgXgi49xVUu8gEQ21teFAxyMwvOiZhIvWNM=
125−github.com/charmbracelet/x/powernap v0.0.0-20260127155452-b72a9a918687 h1:h1XMgTkpBt9kEJ+9DkARNBXEgaigUQ0cI2Bot7Awnt8=
126−github.com/charmbracelet/x/powernap v0.0.0-20260127155452-b72a9a918687/go.mod h1:cmdl5zlP5mR8TF2Y68UKc7hdGUDiSJ2+4hk0h04Hsx4=
125+github.com/charmbracelet/x/powernap v0.0.0-20260209132835-6b065b8ba62c h1:6E+Y7WQ6Rnw+FmeXoRBtyCBkPcXS0hSMuws6QBr+nyQ=
126+github.com/charmbracelet/x/powernap v0.0.0-20260209132835-6b065b8ba62c/go.mod h1:cmdl5zlP5mR8TF2Y68UKc7hdGUDiSJ2+4hk0h04Hsx4=
127127github.com/charmbracelet/x/term v0.2.2 h1:xVRT/S2ZcKdhhOuSP4t5cLi5o+JxklsoEObBSgfgZRk=
128128github.com/charmbracelet/x/term v0.2.2/go.mod h1:kF8CY5RddLWrsgVwpw4kAa6TESp6EB5y3uxGLeCqzAI=
129129github.com/charmbracelet/x/termios v0.1.1 h1:o3Q2bT8eqzGnGPOYheoYS8eEleT5ZVNYNy8JawjaNZY=
internal/agent/tools/mcp/init.go+1−4
@@ -123,10 +123,7 @@ func GetState(name string) (ClientInfo, bool) {
123123}
124124
125125// Close closes all MCP clients. This should be called during application shutdown.
126−func Close() error {
127− ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
128− defer cancel()
129−
126+func Close(ctx context.Context) error {
130127 var wg sync.WaitGroup
131128 for name, session := range sessions.Seq2() {
132129 wg.Go(func() {
internal/app/app.go+16−6
@@ -22,6 +22,7 @@ import (
2222 "github.com/charmbracelet/crush/internal/agent/tools/mcp"
2323 "github.com/charmbracelet/crush/internal/config"
2424 "github.com/charmbracelet/crush/internal/db"
25+ "github.com/charmbracelet/crush/internal/event"
2526 "github.com/charmbracelet/crush/internal/filetracker"
2627 "github.com/charmbracelet/crush/internal/format"
2728 "github.com/charmbracelet/crush/internal/history"
@@ -68,7 +69,7 @@ type App struct {
6869
6970 // global context and cleanup functions
7071 globalCtx context.Context
71− cleanupFuncs []func() error
72+ cleanupFuncs []func(context.Context) error
7273}
7374
7475// New initializes a new application instance.
@@ -108,7 +109,11 @@ func New(ctx context.Context, conn *sql.DB, cfg *config.Config) (*App, error) {
108109 go mcp.Initialize(ctx, app.Permissions, cfg)
109110
110111 // cleanup database upon app shutdown
111− app.cleanupFuncs = append(app.cleanupFuncs, conn.Close, mcp.Close)
112+ app.cleanupFuncs = append(
113+ app.cleanupFuncs,
114+ func(context.Context) error { return conn.Close() },
115+ mcp.Close,
116+ )
112117
113118 // TODO: remove the concept of agent config, most likely.
114119 if !cfg.IsConfigured() {
@@ -416,7 +421,7 @@ func (app *App) setupEvents() {
416421 setupSubscriber(ctx, app.serviceEventsWG, "history", app.History.Subscribe, app.events)
417422 setupSubscriber(ctx, app.serviceEventsWG, "mcp", mcp.SubscribeEvents, app.events)
418423 setupSubscriber(ctx, app.serviceEventsWG, "lsp", SubscribeLSPEvents, app.events)
419− cleanupFunc := func() error {
424+ cleanupFunc := func(context.Context) error {
420425 cancel()
421426 app.serviceEventsWG.Wait()
422427 return nil
@@ -503,7 +508,7 @@ func (app *App) Subscribe(program *tea.Program) {
503508
504509 app.tuiWG.Add(1)
505510 tuiCtx, tuiCancel := context.WithCancel(app.globalCtx)
506− app.cleanupFuncs = append(app.cleanupFuncs, func() error {
511+ app.cleanupFuncs = append(app.cleanupFuncs, func(context.Context) error {
507512 slog.Debug("Cancelling TUI message handler")
508513 tuiCancel()
509514 app.tuiWG.Wait()
@@ -544,6 +549,11 @@ func (app *App) Shutdown() {
544549 shutdownCtx, cancel := context.WithTimeout(app.globalCtx, 5*time.Second)
545550 defer cancel()
546551
552+ // Send exit event
553+ wg.Go(func() {
554+ event.AppExited()
555+ })
556+
547557 // Kill all background shells.
548558 wg.Go(func() {
549559 shell.GetBackgroundShellManager().KillAll(shutdownCtx)
@@ -551,14 +561,14 @@ func (app *App) Shutdown() {
551561
552562 // Shutdown all LSP clients.
553563 wg.Go(func() {
554− app.LSPManager.StopAll(shutdownCtx)
564+ app.LSPManager.KillAll(shutdownCtx)
555565 })
556566
557567 // Call all cleanup functions.
558568 for _, cleanup := range app.cleanupFuncs {
559569 if cleanup != nil {
560570 wg.Go(func() {
561− if err := cleanup(); err != nil {
571+ if err := cleanup(shutdownCtx); err != nil {
562572 slog.Error("Failed to cleanup app properly on shutdown", "error", err)
563573 }
564574 })
internal/cmd/dirs_test.go+2
@@ -12,6 +12,8 @@ import (
1212func init() {
1313 os.Setenv("XDG_CONFIG_HOME", "/tmp/fakeconfig")
1414 os.Setenv("XDG_DATA_HOME", "/tmp/fakedata")
15+ os.Unsetenv("CRUSH_GLOBAL_CONFIG")
16+ os.Unsetenv("CRUSH_GLOBAL_DATA")
1517}
1618
1719func TestDirs(t *testing.T) {
internal/cmd/root.go−3
@@ -106,9 +106,6 @@ crush -y
106106 }
107107 return nil
108108 },
109− PostRun: func(cmd *cobra.Command, args []string) {
110− event.AppExited()
111− },
112109}
113110
114111var heartbit = lipgloss.NewStyle().Foreground(charmtone.Dolly).SetString(`
internal/lsp/client.go+4−1
@@ -117,7 +117,10 @@ func (c *Client) Initialize(ctx context.Context, workspaceDir string) (*protocol
117117 return result, nil
118118}
119119
120−// Close closes the LSP client.
120+// Kill kills the client without doing anything else.
121+func (c *Client) Kill() { c.client.Kill() }
122+
123+// Close closes all open files in the client, then the client.
121124func (c *Client) Close(ctx context.Context) error {
122125 c.CloseAllFiles(ctx)
123126
internal/lsp/manager.go+22
@@ -255,6 +255,28 @@ func handles(server *powernapconfig.ServerConfig, filePath, workDir string) bool
255255 hasRootMarkers(workDir, server.RootMarkers)
256256}
257257
258+// KillAll force-kills all the LSP clients.
259+//
260+// This is generally faster than [Manager.StopAll] because it doesn't wait for
261+// the server to exit gracefully, but it can lead to data loss if the server is
262+// in the middle of writing something.
263+// Generally it doesn't matter when shutting down Crush, though.
264+func (s *Manager) KillAll(context.Context) {
265+ s.mu.Lock()
266+ defer s.mu.Unlock()
267+
268+ var wg sync.WaitGroup
269+ for name, client := range s.clients.Seq2() {
270+ wg.Go(func() {
271+ defer func() { s.callback(name, client) }()
272+ client.client.Kill()
273+ client.SetServerState(StateStopped)
274+ slog.Debug("Killed LSP client", "name", name)
275+ })
276+ }
277+ wg.Wait()
278+}
279+
258280// StopAll stops all running LSP clients and clears the client map.
259281func (s *Manager) StopAll(ctx context.Context) {
260282 s.mu.Lock()
internal/projects/projects_test.go+6
@@ -12,6 +12,7 @@ func TestRegisterAndList(t *testing.T) {
1212
1313 // Override the projects file path for testing
1414 t.Setenv("XDG_DATA_HOME", tmpDir)
15+ t.Setenv("CRUSH_GLOBAL_DATA", filepath.Join(tmpDir, "crush"))
1516
1617 // Test registering a project
1718 err := Register("/home/user/project1", "/home/user/project1/.crush")
@@ -61,6 +62,7 @@ func TestRegisterAndList(t *testing.T) {
6162func TestRegisterUpdatesExisting(t *testing.T) {
6263 tmpDir := t.TempDir()
6364 t.Setenv("XDG_DATA_HOME", tmpDir)
65+ t.Setenv("CRUSH_GLOBAL_DATA", filepath.Join(tmpDir, "crush"))
6466
6567 // Register a project
6668 err := Register("/home/user/project1", "/home/user/project1/.crush")
@@ -97,6 +99,7 @@ func TestRegisterUpdatesExisting(t *testing.T) {
9799func TestLoadEmptyFile(t *testing.T) {
98100 tmpDir := t.TempDir()
99101 t.Setenv("XDG_DATA_HOME", tmpDir)
102+ t.Setenv("CRUSH_GLOBAL_DATA", filepath.Join(tmpDir, "crush"))
100103
101104 // List before any projects exist
102105 projects, err := List()
@@ -112,6 +115,7 @@ func TestLoadEmptyFile(t *testing.T) {
112115func TestProjectsFilePath(t *testing.T) {
113116 tmpDir := t.TempDir()
114117 t.Setenv("XDG_DATA_HOME", tmpDir)
118+ t.Setenv("CRUSH_GLOBAL_DATA", filepath.Join(tmpDir, "crush"))
115119
116120 expected := filepath.Join(tmpDir, "crush", "projects.json")
117121 actual := projectsFilePath()
@@ -124,6 +128,7 @@ func TestProjectsFilePath(t *testing.T) {
124128func TestRegisterWithParentDataDir(t *testing.T) {
125129 tmpDir := t.TempDir()
126130 t.Setenv("XDG_DATA_HOME", tmpDir)
131+ t.Setenv("CRUSH_GLOBAL_DATA", filepath.Join(tmpDir, "crush"))
127132
128133 // Register a project where .crush is in a parent directory.
129134 // e.g., working in /home/user/monorepo/packages/app but .crush is at /home/user/monorepo/.crush
@@ -153,6 +158,7 @@ func TestRegisterWithParentDataDir(t *testing.T) {
153158func TestRegisterWithExternalDataDir(t *testing.T) {
154159 tmpDir := t.TempDir()
155160 t.Setenv("XDG_DATA_HOME", tmpDir)
161+ t.Setenv("CRUSH_GLOBAL_DATA", filepath.Join(tmpDir, "crush"))
156162
157163 // Register a project where .crush is in a completely different location.
158164 // e.g., project at /home/user/project but data stored at /var/data/crush/myproject
internal/shell/background_test.go+14−11
@@ -14,7 +14,7 @@ func TestBackgroundShellManager_Start(t *testing.T) {
1414 t.Skip("Skipping this until I figure out why its flaky")
1515 t.Parallel()
1616
17− ctx := context.Background()
17+ ctx := t.Context()
1818 workingDir := t.TempDir()
1919 manager := newBackgroundShellManager()
2020
@@ -51,7 +51,7 @@ func TestBackgroundShellManager_Start(t *testing.T) {
5151func TestBackgroundShellManager_Get(t *testing.T) {
5252 t.Parallel()
5353
54− ctx := context.Background()
54+ ctx := t.Context()
5555 workingDir := t.TempDir()
5656 manager := newBackgroundShellManager()
5757
@@ -77,7 +77,7 @@ func TestBackgroundShellManager_Get(t *testing.T) {
7777func TestBackgroundShellManager_Kill(t *testing.T) {
7878 t.Parallel()
7979
80− ctx := context.Background()
80+ ctx := t.Context()
8181 workingDir := t.TempDir()
8282 manager := newBackgroundShellManager()
8383
@@ -119,7 +119,7 @@ func TestBackgroundShellManager_KillNonExistent(t *testing.T) {
119119func TestBackgroundShell_IsDone(t *testing.T) {
120120 t.Parallel()
121121
122− ctx := context.Background()
122+ ctx := t.Context()
123123 workingDir := t.TempDir()
124124 manager := newBackgroundShellManager()
125125
@@ -142,7 +142,7 @@ func TestBackgroundShell_IsDone(t *testing.T) {
142142func TestBackgroundShell_WithBlockFuncs(t *testing.T) {
143143 t.Parallel()
144144
145− ctx := context.Background()
145+ ctx := t.Context()
146146 workingDir := t.TempDir()
147147 manager := newBackgroundShellManager()
148148
@@ -180,7 +180,7 @@ func TestBackgroundShellManager_List(t *testing.T) {
180180
181181 t.Parallel()
182182
183− ctx := context.Background()
183+ ctx := t.Context()
184184 workingDir := t.TempDir()
185185 manager := newBackgroundShellManager()
186186
@@ -224,7 +224,7 @@ func TestBackgroundShellManager_List(t *testing.T) {
224224func TestBackgroundShellManager_KillAll(t *testing.T) {
225225 t.Parallel()
226226
227− ctx := context.Background()
227+ ctx := t.Context()
228228 workingDir := t.TempDir()
229229 manager := newBackgroundShellManager()
230230
@@ -250,7 +250,7 @@ func TestBackgroundShellManager_KillAll(t *testing.T) {
250250 }
251251
252252 // Kill all shells
253− manager.KillAll(context.Background())
253+ manager.KillAll(t.Context())
254254
255255 // Verify all shells are done
256256 if !shell1.IsDone() {
@@ -286,19 +286,22 @@ func TestBackgroundShellManager_KillAll(t *testing.T) {
286286func TestBackgroundShellManager_KillAll_Timeout(t *testing.T) {
287287 t.Parallel()
288288
289+ // XXX: can't use synctest here - causes --race to trip.
290+
289291 workingDir := t.TempDir()
290292 manager := newBackgroundShellManager()
291293
292294 // Start a shell that traps signals and ignores cancellation.
293− _, err := manager.Start(context.Background(), workingDir, nil, "trap '' TERM INT; sleep 60", "")
295+ _, err := manager.Start(t.Context(), workingDir, nil, "trap '' TERM INT; sleep 60", "")
294296 require.NoError(t, err)
295297
296298 // Short timeout to test the timeout path.
297− ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
298− defer cancel()
299+ ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond)
300+ t.Cleanup(cancel)
299301
300302 start := time.Now()
301303 manager.KillAll(ctx)
304+
302305 elapsed := time.Since(start)
303306
304307 // Must return promptly after timeout, not hang for 60 seconds.