@@ -6,6 +6,7 @@ import (
66 "io"
77 "log/slog"
88 "os"
9+ "sync"
910 "time"
1011
1112 tea "github.com/charmbracelet/bubbletea/v2"
@@ -72,9 +73,8 @@ to assist developers in writing, debugging, and understanding code directly from
7273 return err
7374 }
7475
75− // Create main context for the application
76− ctx, cancel := context.WithCancel(context.Background())
77− defer cancel()
76+ // Use the context from the command which includes signal handling
77+ ctx := cmd.Context()
7878
7979 // Connect DB, this will also run migrations
8080 conn, err := db.Connect(ctx, cfg.Options.DataDirectory)
@@ -87,8 +87,23 @@ to assist developers in writing, debugging, and understanding code directly from
8787 slog.Error(fmt.Sprintf("Failed to create app instance: %v", err))
8888 return err
8989 }
90− // Defer shutdown here so it runs for both interactive and non-interactive modes
91− defer app.Shutdown()
90+
91+ // Set up shutdown handling that works for both normal exit and signal interruption
92+ var shutdownOnce sync.Once
93+ shutdown := func() {
94+ shutdownOnce.Do(func() {
95+ slog.Info("Shutting down application")
96+ app.Shutdown()
97+ })
98+ }
99+ defer shutdown()
100+
101+ // Handle context cancellation (from signals) in a goroutine
102+ go func() {
103+ <-ctx.Done()
104+ slog.Info("Context cancelled, initiating shutdown")
105+ shutdown()
106+ }()
92107
93108 // Initialize MCP tools early for both modes
94109 initMCPTools(ctx, app, cfg)
@@ -121,7 +136,6 @@ to assist developers in writing, debugging, and understanding code directly from
121136 slog.Error(fmt.Sprintf("TUI run error: %v", err))
122137 return fmt.Errorf("TUI error: %v", err)
123138 }
124− app.Shutdown()
125139 return nil
126140 },
127141}
@@ -140,9 +154,9 @@ func initMCPTools(ctx context.Context, app *app.App, cfg *config.Config) {
140154 }()
141155}
142156
143−func Execute() {
157+func Execute(ctx context.Context) {
144158 if err := fang.Execute(
145− context.Background(),
159+ ctx,
146160 rootCmd,
147161 fang.WithVersion(version.Version),
148162 ); err != nil {