git.iko.soy/crush
on main help sign in
cleanup app, config and root
commit3ad983db0f2c08826d56cb5de274d706c95b3353 authorKujtim Hoxha <kujtimii.h@gmail.com> dateSun 13 Apr 2025, 13:17 (2 years ago) parent5601466
21 files changed+514−281
compare with main
unifiedsplit
.gitignore+1−1 .termai.json → .opencode.json cmd/git/main.go−4 cmd/root.go+187−66 internal/app/app.go+76 internal/app/lsp.go+108 internal/app/services.go−64 internal/config/config.go+6−14 internal/history/file.go+35−38 internal/llm/agent/agent-tool.go+5−5 internal/llm/agent/agent.go+27−26 internal/llm/agent/coder.go+3−2 internal/llm/agent/task.go+2−1 internal/message/message.go+22−24 internal/session/session.go+21−23 internal/tui/components/chat/messages.go+3−2 internal/tui/components/repl/editor.go+2−2 internal/tui/components/repl/messages.go+4−3 internal/tui/components/repl/sessions.go+3−1 internal/tui/page/chat.go+5−3 internal/tui/tui.go+4−2
.gitignore+1−1
@@ -41,6 +41,6 @@ debug.log
4141.env
4242.env.local
4343
44−.termai
44+.opencode
4545
4646internal/assets/diff/index.mjs
.termai.json → .opencode.jsonrenamed

The content is the same.

cmd/git/main.godeleted−4
@@ -1,4 +0,0 @@
1−package main
2−
3−func main() {
4−}
cmd/root.go+187−66
@@ -2,9 +2,10 @@ package cmd
22
33import (
44 "context"
5− "log/slog"
5+ "fmt"
66 "os"
77 "sync"
8+ "time"
89
910 tea "github.com/charmbracelet/bubbletea"
1011 "github.com/kujtimiihoxha/termai/internal/app"
@@ -13,6 +14,7 @@ import (
1314 "github.com/kujtimiihoxha/termai/internal/db"
1415 "github.com/kujtimiihoxha/termai/internal/llm/agent"
1516 "github.com/kujtimiihoxha/termai/internal/logging"
17+ "github.com/kujtimiihoxha/termai/internal/pubsub"
1618 "github.com/kujtimiihoxha/termai/internal/tui"
1719 zone "github.com/lrstanley/bubblezone"
1820 "github.com/spf13/cobra"
@@ -23,111 +25,229 @@ var rootCmd = &cobra.Command{
2325 Short: "A terminal ai assistant",
2426 Long: `A terminal ai assistant`,
2527 RunE: func(cmd *cobra.Command, args []string) error {
28+ // If the help flag is set, show the help message
2629 if cmd.Flag("help").Changed {
2730 cmd.Help()
2831 return nil
2932 }
33+
34+ // Load the config
3035 debug, _ := cmd.Flags().GetBool("debug")
31− err := config.Load(debug)
36+ cwd, _ := cmd.Flags().GetString("cwd")
37+ if cwd != "" {
38+ err := os.Chdir(cwd)
39+ if err != nil {
40+ return fmt.Errorf("failed to change directory: %v", err)
41+ }
42+ }
43+ if cwd == "" {
44+ c, err := os.Getwd()
45+ if err != nil {
46+ return fmt.Errorf("failed to get current working directory: %v", err)
47+ }
48+ cwd = c
49+ }
50+ _, err := config.Load(cwd, debug)
3251 if err != nil {
3352 return err
3453 }
35− cfg := config.Get()
36− defaultLevel := slog.LevelInfo
37− if cfg.Debug {
38− defaultLevel = slog.LevelDebug
39− }
40− logger := slog.New(slog.NewTextHandler(logging.NewWriter(), &slog.HandlerOptions{
41− Level: defaultLevel,
42− }))
43− slog.SetDefault(logger)
4454
4555 err = assets.WriteAssets()
4656 if err != nil {
47− return err
57+ logging.Error("Error writing assets: %v", err)
4858 }
4959
60+ // Connect DB, this will also run migrations
5061 conn, err := db.Connect()
5162 if err != nil {
5263 return err
5364 }
54− ctx := context.Background()
65+
66+ // Create main context for the application
67+ ctx, cancel := context.WithCancel(context.Background())
68+ defer cancel()
5569
5670 app := app.New(ctx, conn)
57− logging.Info("Starting termai...")
71+
72+ // Set up the TUI
5873 zone.NewGlobal()
59− tui := tea.NewProgram(
74+ program := tea.NewProgram(
6075 tui.New(app),
6176 tea.WithAltScreen(),
6277 tea.WithMouseCellMotion(),
6378 )
64− logging.Info("Setting up subscriptions...")
65− ch, unsub := setupSubscriptions(app)
66− defer unsub()
6779
80+ // Initialize MCP tools in the background
81+ initMCPTools(ctx, app)
82+
83+ // Setup the subscriptions, this will send services events to the TUI
84+ ch, cancelSubs := setupSubscriptions(app)
85+
86+ // Create a context for the TUI message handler
87+ tuiCtx, tuiCancel := context.WithCancel(ctx)
88+ var tuiWg sync.WaitGroup
89+ tuiWg.Add(1)
90+
91+ // Set up message handling for the TUI
6892 go func() {
69− // Set this up once
70− agent.GetMcpTools(ctx, app.Permissions)
71− for msg := range ch {
72− tui.Send(msg)
93+ defer tuiWg.Done()
94+ defer func() {
95+ if r := recover(); r != nil {
96+ logging.Error("Panic in TUI message handling: %v", r)
97+ attemptTUIRecovery(program)
98+ }
99+ }()
100+
101+ for {
102+ select {
103+ case <-tuiCtx.Done():
104+ logging.Info("TUI message handler shutting down")
105+ return
106+ case msg, ok := <-ch:
107+ if !ok {
108+ logging.Info("TUI message channel closed")
109+ return
110+ }
111+ program.Send(msg)
112+ }
73113 }
74114 }()
75− if _, err := tui.Run(); err != nil {
76− return err
115+
116+ // Cleanup function for when the program exits
117+ cleanup := func() {
118+ // Shutdown the app
119+ app.Shutdown()
120+
121+ // Cancel subscriptions first
122+ cancelSubs()
123+
124+ // Then cancel TUI message handler
125+ tuiCancel()
126+
127+ // Wait for TUI message handler to finish
128+ tuiWg.Wait()
129+
130+ logging.Info("All goroutines cleaned up")
131+ }
132+
133+ // Run the TUI
134+ result, err := program.Run()
135+ cleanup()
136+
137+ if err != nil {
138+ logging.Error("TUI error: %v", err)
139+ return fmt.Errorf("TUI error: %v", err)
77140 }
141+
142+ logging.Info("TUI exited with result: %v", result)
78143 return nil
79144 },
80145}
81146
82−func setupSubscriptions(app *app.App) (chan tea.Msg, func()) {
83− ch := make(chan tea.Msg)
84− wg := sync.WaitGroup{}
85− ctx, cancel := context.WithCancel(app.Context)
86− {
87− sub := logging.Subscribe(ctx)
88− wg.Add(1)
89− go func() {
90− for ev := range sub {
91− ch <- ev
147+// attemptTUIRecovery tries to recover the TUI after a panic
148+func attemptTUIRecovery(program *tea.Program) {
149+ logging.Info("Attempting to recover TUI after panic")
150+
151+ // We could try to restart the TUI or gracefully exit
152+ // For now, we'll just quit the program to avoid further issues
153+ program.Quit()
154+}
155+
156+func initMCPTools(ctx context.Context, app *app.App) {
157+ go func() {
158+ defer func() {
159+ if r := recover(); r != nil {
160+ logging.Error("Panic in MCP goroutine: %v", r)
92161 }
93− wg.Done()
94162 }()
95− }
96− {
97− sub := app.Sessions.Subscribe(ctx)
98− wg.Add(1)
99− go func() {
100− for ev := range sub {
101− ch <- ev
163+
164+ // Create a context with timeout for the initial MCP tools fetch
165+ ctxWithTimeout, cancel := context.WithTimeout(ctx, 30*time.Second)
166+ defer cancel()
167+
168+ // Set this up once with proper error handling
169+ agent.GetMcpTools(ctxWithTimeout, app.Permissions)
170+ logging.Info("MCP message handling goroutine exiting")
171+ }()
172+}
173+
174+func setupSubscriber[T any](
175+ ctx context.Context,
176+ wg *sync.WaitGroup,
177+ name string,
178+ subscriber func(context.Context) <-chan pubsub.Event[T],
179+ outputCh chan<- tea.Msg,
180+) {
181+ wg.Add(1)
182+ go func() {
183+ defer wg.Done()
184+ defer func() {
185+ if r := recover(); r != nil {
186+ logging.Error("Panic in %s subscription goroutine: %v", name, r)
102187 }
103− wg.Done()
104188 }()
105− }
106− {
107− sub := app.Messages.Subscribe(ctx)
108− wg.Add(1)
109− go func() {
110− for ev := range sub {
111− ch <- ev
189+
190+ for {
191+ select {
192+ case event, ok := <-subscriber(ctx):
193+ if !ok {
194+ logging.Info("%s subscription channel closed", name)
195+ return
196+ }
197+
198+ // Convert generic event to tea.Msg if needed
199+ var msg tea.Msg = event
200+
201+ // Non-blocking send with timeout to prevent deadlocks
202+ select {
203+ case outputCh <- msg:
204+ case <-time.After(500 * time.Millisecond):
205+ logging.Warn("%s message dropped due to slow consumer", name)
206+ case <-ctx.Done():
207+ logging.Info("%s subscription cancelled", name)
208+ return
209+ }
210+ case <-ctx.Done():
211+ logging.Info("%s subscription cancelled", name)
212+ return
112213 }
113− wg.Done()
114− }()
115− }
116− {
117− sub := app.Permissions.Subscribe(ctx)
118− wg.Add(1)
214+ }
215+ }()
216+}
217+
218+func setupSubscriptions(app *app.App) (chan tea.Msg, func()) {
219+ ch := make(chan tea.Msg, 100)
220+ // Add a buffer to prevent blocking
221+ wg := sync.WaitGroup{}
222+ ctx, cancel := context.WithCancel(context.Background())
223+ // Setup each subscription using the helper
224+ setupSubscriber(ctx, &wg, "logging", logging.Subscribe, ch)
225+ setupSubscriber(ctx, &wg, "sessions", app.Sessions.Subscribe, ch)
226+ setupSubscriber(ctx, &wg, "messages", app.Messages.Subscribe, ch)
227+ setupSubscriber(ctx, &wg, "permissions", app.Permissions.Subscribe, ch)
228+
229+ // Return channel and a cleanup function
230+ cleanupFunc := func() {
231+ logging.Info("Cancelling all subscriptions")
232+ cancel() // Signal all goroutines to stop
233+
234+ // Wait with a timeout for all goroutines to complete
235+ waitCh := make(chan struct{})
119236 go func() {
120− for ev := range sub {
121− ch <- ev
122− }
123− wg.Done()
237+ wg.Wait()
238+ close(waitCh)
124239 }()
240+
241+ select {
242+ case <-waitCh:
243+ logging.Info("All subscription goroutines completed successfully")
244+ case <-time.After(5 * time.Second):
245+ logging.Warn("Timed out waiting for some subscription goroutines to complete")
246+ }
247+
248+ close(ch) // Safe to close after all writers are done or timed out
125249 }
126− return ch, func() {
127− cancel()
128− wg.Wait()
129− close(ch)
130− }
250+ return ch, cleanupFunc
131251}
132252
133253func Execute() {
@@ -139,5 +259,6 @@ func Execute() {
139259
140260func init() {
141261 rootCmd.Flags().BoolP("help", "h", false, "Help")
142− rootCmd.Flags().BoolP("debug", "d", false, "Help")
262+ rootCmd.Flags().BoolP("debug", "d", false, "Debug")
263+ rootCmd.Flags().StringP("cwd", "c", "", "Current working directory")
143264}
internal/app/app.gonew file+76
@@ -0,0 +1,76 @@
1+package app
2+
3+import (
4+ "context"
5+ "database/sql"
6+ "maps"
7+ "sync"
8+ "time"
9+
10+ "github.com/kujtimiihoxha/termai/internal/db"
11+ "github.com/kujtimiihoxha/termai/internal/history"
12+ "github.com/kujtimiihoxha/termai/internal/logging"
13+ "github.com/kujtimiihoxha/termai/internal/lsp"
14+ "github.com/kujtimiihoxha/termai/internal/message"
15+ "github.com/kujtimiihoxha/termai/internal/permission"
16+ "github.com/kujtimiihoxha/termai/internal/session"
17+)
18+
19+type App struct {
20+ Sessions session.Service
21+ Messages message.Service
22+ Files history.Service
23+ Permissions permission.Service
24+
25+ LSPClients map[string]*lsp.Client
26+
27+ clientsMutex sync.RWMutex
28+
29+ watcherCancelFuncs []context.CancelFunc
30+ cancelFuncsMutex sync.Mutex
31+ watcherWG sync.WaitGroup
32+}
33+
34+func New(ctx context.Context, conn *sql.DB) *App {
35+ q := db.New(conn)
36+ sessions := session.NewService(q)
37+ messages := message.NewService(q)
38+ files := history.NewService(q)
39+
40+ app := &App{
41+ Sessions: sessions,
42+ Messages: messages,
43+ Files: files,
44+ Permissions: permission.NewPermissionService(),
45+ LSPClients: make(map[string]*lsp.Client),
46+ }
47+
48+ app.initLSPClients(ctx)
49+
50+ return app
51+}
52+
53+// Shutdown performs a clean shutdown of the application
54+func (app *App) Shutdown() {
55+ // Cancel all watcher goroutines
56+ app.cancelFuncsMutex.Lock()
57+ for _, cancel := range app.watcherCancelFuncs {
58+ cancel()
59+ }
60+ app.cancelFuncsMutex.Unlock()
61+ app.watcherWG.Wait()
62+
63+ // Perform additional cleanup for LSP clients
64+ app.clientsMutex.RLock()
65+ clients := make(map[string]*lsp.Client, len(app.LSPClients))
66+ maps.Copy(clients, app.LSPClients)
67+ app.clientsMutex.RUnlock()
68+
69+ for name, client := range clients {
70+ shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
71+ if err := client.Shutdown(shutdownCtx); err != nil {
72+ logging.Error("Failed to shutdown LSP client", "name", name, "error", err)
73+ }
74+ cancel()
75+ }
76+}
internal/app/lsp.gonew file+108
@@ -0,0 +1,108 @@
1+package app
2+
3+import (
4+ "context"
5+ "time"
6+
7+ "github.com/kujtimiihoxha/termai/internal/config"
8+ "github.com/kujtimiihoxha/termai/internal/logging"
9+ "github.com/kujtimiihoxha/termai/internal/lsp"
10+ "github.com/kujtimiihoxha/termai/internal/lsp/watcher"
11+)
12+
13+func (app *App) initLSPClients(ctx context.Context) {
14+ cfg := config.Get()
15+
16+ // Initialize LSP clients
17+ for name, clientConfig := range cfg.LSP {
18+ app.createAndStartLSPClient(ctx, name, clientConfig.Command, clientConfig.Args...)
19+ }
20+}
21+
22+// createAndStartLSPClient creates a new LSP client, initializes it, and starts its workspace watcher
23+func (app *App) createAndStartLSPClient(ctx context.Context, name string, command string, args ...string) {
24+ // Create a specific context for initialization with a timeout
25+ initCtx, initCancel := context.WithTimeout(context.Background(), 30*time.Second)
26+ defer initCancel()
27+
28+ // Create the LSP client
29+ lspClient, err := lsp.NewClient(initCtx, command, args...)
30+ if err != nil {
31+ logging.Error("Failed to create LSP client for", name, err)
32+ return
33+ }
34+
35+ // Initialize with the initialization context
36+ _, err = lspClient.InitializeLSPClient(initCtx, config.WorkingDirectory())
37+ if err != nil {
38+ logging.Error("Initialize failed", "name", name, "error", err)
39+ // Clean up the client to prevent resource leaks
40+ lspClient.Close()
41+ return
42+ }
43+
44+ // Create a child context that can be canceled when the app is shutting down
45+ watchCtx, cancelFunc := context.WithCancel(ctx)
46+ workspaceWatcher := watcher.NewWorkspaceWatcher(lspClient)
47+
48+ // Store the cancel function to be called during cleanup
49+ app.cancelFuncsMutex.Lock()
50+ app.watcherCancelFuncs = append(app.watcherCancelFuncs, cancelFunc)
51+ app.cancelFuncsMutex.Unlock()
52+
53+ // Add the watcher to a WaitGroup to track active goroutines
54+ app.watcherWG.Add(1)
55+
56+ // Add to map with mutex protection before starting goroutine
57+ app.clientsMutex.Lock()
58+ app.LSPClients[name] = lspClient
59+ app.clientsMutex.Unlock()
60+
61+ go app.runWorkspaceWatcher(watchCtx, name, workspaceWatcher)
62+}
63+
64+// runWorkspaceWatcher executes the workspace watcher for an LSP client
65+func (app *App) runWorkspaceWatcher(ctx context.Context, name string, workspaceWatcher *watcher.WorkspaceWatcher) {
66+ defer app.watcherWG.Done()
67+ defer func() {
68+ if r := recover(); r != nil {
69+ logging.Error("LSP client crashed", "client", name, "panic", r)
70+
71+ // Try to restart the client
72+ app.restartLSPClient(ctx, name)
73+ }
74+ }()
75+
76+ workspaceWatcher.WatchWorkspace(ctx, config.WorkingDirectory())
77+ logging.Info("Workspace watcher stopped", "client", name)
78+}
79+
80+// restartLSPClient attempts to restart a crashed or failed LSP client
81+func (app *App) restartLSPClient(ctx context.Context, name string) {
82+ // Get the original configuration
83+ cfg := config.Get()
84+ clientConfig, exists := cfg.LSP[name]
85+ if !exists {
86+ logging.Error("Cannot restart client, configuration not found", "client", name)
87+ return
88+ }
89+
90+ // Clean up the old client if it exists
91+ app.clientsMutex.Lock()
92+ oldClient, exists := app.LSPClients[name]
93+ if exists {
94+ delete(app.LSPClients, name) // Remove from map before potentially slow shutdown
95+ }
96+ app.clientsMutex.Unlock()
97+
98+ if exists && oldClient != nil {
99+ // Try to shut it down gracefully, but don't block on errors
100+ shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
101+ _ = oldClient.Shutdown(shutdownCtx)
102+ cancel()
103+ }
104+
105+ // Create a new client using the shared function
106+ app.createAndStartLSPClient(ctx, name, clientConfig.Command, clientConfig.Args...)
107+ logging.Info("Successfully restarted LSP client", "client", name)
108+}
internal/app/services.godeleted−64
@@ -1,64 +0,0 @@
1−package app
2−
3−import (
4− "context"
5− "database/sql"
6−
7− "github.com/kujtimiihoxha/termai/internal/config"
8− "github.com/kujtimiihoxha/termai/internal/db"
9− "github.com/kujtimiihoxha/termai/internal/history"
10− "github.com/kujtimiihoxha/termai/internal/logging"
11− "github.com/kujtimiihoxha/termai/internal/lsp"
12− "github.com/kujtimiihoxha/termai/internal/lsp/watcher"
13− "github.com/kujtimiihoxha/termai/internal/message"
14− "github.com/kujtimiihoxha/termai/internal/permission"
15− "github.com/kujtimiihoxha/termai/internal/session"
16−)
17−
18−type App struct {
19− Context context.Context
20−
21− Sessions session.Service
22− Messages message.Service
23− Files history.Service
24− Permissions permission.Service
25−
26− LSPClients map[string]*lsp.Client
27−}
28−
29−func New(ctx context.Context, conn *sql.DB) *App {
30− cfg := config.Get()
31− logging.Info("Debug mode enabled")
32−
33− q := db.New(conn)
34− sessions := session.NewService(ctx, q)
35− messages := message.NewService(ctx, q)
36− files := history.NewService(ctx, q)
37−
38− app := &App{
39− Context: ctx,
40− Sessions: sessions,
41− Messages: messages,
42− Files: files,
43− Permissions: permission.NewPermissionService(),
44− LSPClients: make(map[string]*lsp.Client),
45− }
46−
47− for name, client := range cfg.LSP {
48− lspClient, err := lsp.NewClient(ctx, client.Command, client.Args...)
49− workspaceWatcher := watcher.NewWorkspaceWatcher(lspClient)
50− if err != nil {
51− logging.Error("Failed to create LSP client for", name, err)
52− continue
53− }
54−
55− _, err = lspClient.InitializeLSPClient(ctx, config.WorkingDirectory())
56− if err != nil {
57− logging.Error("Initialize failed", "error", err)
58− continue
59− }
60− go workspaceWatcher.WatchWorkspace(ctx, config.WorkingDirectory())
61− app.LSPClients[name] = lspClient
62− }
63− return app
64−}
internal/config/config.go+6−14
@@ -83,9 +83,9 @@ var cfg *Config
8383// Load initializes the configuration from environment variables and config files.
8484// If debug is true, debug mode is enabled and log level is set to debug.
8585// It returns an error if configuration loading fails.
86−func Load(workingDir string, debug bool) error {
86+func Load(workingDir string, debug bool) (*Config, error) {
8787 if cfg != nil {
88− return nil
88+ return cfg, nil
8989 }
9090
9191 cfg = &Config{
@@ -101,7 +101,7 @@ func Load(workingDir string, debug bool) error {
101101
102102 // Read global config
103103 if err := readConfig(viper.ReadInConfig()); err != nil {
104− return err
104+ return cfg, err
105105 }
106106
107107 // Load and merge local config
@@ -109,7 +109,7 @@ func Load(workingDir string, debug bool) error {
109109
110110 // Apply configuration to the struct
111111 if err := viper.Unmarshal(cfg); err != nil {
112− return err
112+ return cfg, fmt.Errorf("failed to unmarshal config: %w", err)
113113 }
114114
115115 applyDefaultValues()
@@ -123,7 +123,7 @@ func Load(workingDir string, debug bool) error {
123123 Level: defaultLevel,
124124 }))
125125 slog.SetDefault(logger)
126− return nil
126+ return cfg, nil
127127}
128128
129129// configureViper sets up viper's configuration paths and environment variables.
@@ -237,7 +237,7 @@ func readConfig(err error) error {
237237 return nil
238238 }
239239
240− return err
240+ return fmt.Errorf("failed to read config: %w", err)
241241}
242242
243243// mergeLocalConfig loads and merges configuration from the local directory.
@@ -264,14 +264,6 @@ func applyDefaultValues() {
264264 }
265265}
266266
267−// setWorkingDirectory stores the current working directory in the configuration.
268−func setWorkingDirectory() {
269− workdir, err := os.Getwd()
270− if err == nil {
271− viper.Set("wd", workdir)
272− }
273−}
274−
275267// Get returns the current configuration.
276268// It's safe to call this function multiple times.
277269func Get() *Config {
internal/history/file.go+35−38
@@ -27,45 +27,43 @@ type File struct {
2727
2828type Service interface {
2929 pubsub.Suscriber[File]
30− Create(sessionID, path, content string) (File, error)
31− CreateVersion(sessionID, path, content string) (File, error)
32− Get(id string) (File, error)
33− GetByPathAndSession(path, sessionID string) (File, error)
34− ListBySession(sessionID string) ([]File, error)
35− ListLatestSessionFiles(sessionID string) ([]File, error)
36− Update(file File) (File, error)
37− Delete(id string) error
38− DeleteSessionFiles(sessionID string) error
30+ Create(ctx context.Context, sessionID, path, content string) (File, error)
31+ CreateVersion(ctx context.Context, sessionID, path, content string) (File, error)
32+ Get(ctx context.Context, id string) (File, error)
33+ GetByPathAndSession(ctx context.Context, path, sessionID string) (File, error)
34+ ListBySession(ctx context.Context, sessionID string) ([]File, error)
35+ ListLatestSessionFiles(ctx context.Context, sessionID string) ([]File, error)
36+ Update(ctx context.Context, file File) (File, error)
37+ Delete(ctx context.Context, id string) error
38+ DeleteSessionFiles(ctx context.Context, sessionID string) error
3939}
4040
4141type service struct {
4242 *pubsub.Broker[File]
43− q db.Querier
44− ctx context.Context
43+ q db.Querier
4544}
4645
47−func NewService(ctx context.Context, q db.Querier) Service {
46+func NewService(q db.Querier) Service {
4847 return &service{
4948 Broker: pubsub.NewBroker[File](),
5049 q: q,
51− ctx: ctx,
5250 }
5351}
5452
55−func (s *service) Create(sessionID, path, content string) (File, error) {
56− return s.createWithVersion(sessionID, path, content, InitialVersion)
53+func (s *service) Create(ctx context.Context, sessionID, path, content string) (File, error) {
54+ return s.createWithVersion(ctx, sessionID, path, content, InitialVersion)
5755}
5856
59−func (s *service) CreateVersion(sessionID, path, content string) (File, error) {
57+func (s *service) CreateVersion(ctx context.Context, sessionID, path, content string) (File, error) {
6058 // Get the latest version for this path
61− files, err := s.q.ListFilesByPath(s.ctx, path)
59+ files, err := s.q.ListFilesByPath(ctx, path)
6260 if err != nil {
6361 return File{}, err
6462 }
6563
6664 if len(files) == 0 {
6765 // No previous versions, create initial
68− return s.Create(sessionID, path, content)
66+ return s.Create(ctx, sessionID, path, content)
6967 }
7068
7169 // Get the latest version
@@ -89,11 +87,11 @@ func (s *service) CreateVersion(sessionID, path, content string) (File, error) {
8987 nextVersion = fmt.Sprintf("v%d", latestFile.CreatedAt)
9088 }
9189
92− return s.createWithVersion(sessionID, path, content, nextVersion)
90+ return s.createWithVersion(ctx, sessionID, path, content, nextVersion)
9391}
9492
95−func (s *service) createWithVersion(sessionID, path, content, version string) (File, error) {
96− dbFile, err := s.q.CreateFile(s.ctx, db.CreateFileParams{
93+func (s *service) createWithVersion(ctx context.Context, sessionID, path, content, version string) (File, error) {
94+ dbFile, err := s.q.CreateFile(ctx, db.CreateFileParams{
9795 ID: uuid.New().String(),
9896 SessionID: sessionID,
9997 Path: path,
@@ -108,16 +106,16 @@ func (s *service) createWithVersion(sessionID, path, content, version string) (F
108106 return file, nil
109107}
110108
111−func (s *service) Get(id string) (File, error) {
112− dbFile, err := s.q.GetFile(s.ctx, id)
109+func (s *service) Get(ctx context.Context, id string) (File, error) {
110+ dbFile, err := s.q.GetFile(ctx, id)
113111 if err != nil {
114112 return File{}, err
115113 }
116114 return s.fromDBItem(dbFile), nil
117115}
118116
119−func (s *service) GetByPathAndSession(path, sessionID string) (File, error) {
120− dbFile, err := s.q.GetFileByPathAndSession(s.ctx, db.GetFileByPathAndSessionParams{
117+func (s *service) GetByPathAndSession(ctx context.Context, path, sessionID string) (File, error) {
118+ dbFile, err := s.q.GetFileByPathAndSession(ctx, db.GetFileByPathAndSessionParams{
121119 Path: path,
122120 SessionID: sessionID,
123121 })
@@ -127,8 +125,8 @@ func (s *service) GetByPathAndSession(path, sessionID string) (File, error) {
127125 return s.fromDBItem(dbFile), nil
128126}
129127
130−func (s *service) ListBySession(sessionID string) ([]File, error) {
131− dbFiles, err := s.q.ListFilesBySession(s.ctx, sessionID)
128+func (s *service) ListBySession(ctx context.Context, sessionID string) ([]File, error) {
129+ dbFiles, err := s.q.ListFilesBySession(ctx, sessionID)
132130 if err != nil {
133131 return nil, err
134132 }
@@ -139,8 +137,8 @@ func (s *service) ListBySession(sessionID string) ([]File, error) {
139137 return files, nil
140138}
141139
142−func (s *service) ListLatestSessionFiles(sessionID string) ([]File, error) {
143− dbFiles, err := s.q.ListLatestSessionFiles(s.ctx, sessionID)
140+func (s *service) ListLatestSessionFiles(ctx context.Context, sessionID string) ([]File, error) {
141+ dbFiles, err := s.q.ListLatestSessionFiles(ctx, sessionID)
144142 if err != nil {
145143 return nil, err
146144 }
@@ -151,8 +149,8 @@ func (s *service) ListLatestSessionFiles(sessionID string) ([]File, error) {
151149 return files, nil
152150}
153151
154−func (s *service) Update(file File) (File, error) {
155− dbFile, err := s.q.UpdateFile(s.ctx, db.UpdateFileParams{
152+func (s *service) Update(ctx context.Context, file File) (File, error) {
153+ dbFile, err := s.q.UpdateFile(ctx, db.UpdateFileParams{
156154 ID: file.ID,
157155 Content: file.Content,
158156 Version: file.Version,
@@ -165,12 +163,12 @@ func (s *service) Update(file File) (File, error) {
165163 return updatedFile, nil
166164}
167165
168−func (s *service) Delete(id string) error {
169− file, err := s.Get(id)
166+func (s *service) Delete(ctx context.Context, id string) error {
167+ file, err := s.Get(ctx, id)
170168 if err != nil {
171169 return err
172170 }
173− err = s.q.DeleteFile(s.ctx, id)
171+ err = s.q.DeleteFile(ctx, id)
174172 if err != nil {
175173 return err
176174 }
@@ -178,13 +176,13 @@ func (s *service) Delete(id string) error {
178176 return nil
179177}
180178
181−func (s *service) DeleteSessionFiles(sessionID string) error {
182− files, err := s.ListBySession(sessionID)
179+func (s *service) DeleteSessionFiles(ctx context.Context, sessionID string) error {
180+ files, err := s.ListBySession(ctx, sessionID)
183181 if err != nil {
184182 return err
185183 }
186184 for _, file := range files {
187− err = s.Delete(file.ID)
185+ err = s.Delete(ctx, file.ID)
188186 if err != nil {
189187 return err
190188 }
@@ -203,4 +201,3 @@ func (s *service) fromDBItem(item db.File) File {
203201 UpdatedAt: item.UpdatedAt,
204202 }
205203}
206−
internal/llm/agent/agent-tool.go+5−5
@@ -51,7 +51,7 @@ func (b *agentTool) Run(ctx context.Context, call tools.ToolCall) (tools.ToolRes
5151 return tools.NewTextErrorResponse(fmt.Sprintf("error creating agent: %s", err)), nil
5252 }
5353
54− session, err := b.app.Sessions.CreateTaskSession(call.ID, b.parentSessionID, "New Agent Session")
54+ session, err := b.app.Sessions.CreateTaskSession(ctx, call.ID, b.parentSessionID, "New Agent Session")
5555 if err != nil {
5656 return tools.NewTextErrorResponse(fmt.Sprintf("error creating session: %s", err)), nil
5757 }
@@ -61,7 +61,7 @@ func (b *agentTool) Run(ctx context.Context, call tools.ToolCall) (tools.ToolRes
6161 return tools.NewTextErrorResponse(fmt.Sprintf("error generating agent: %s", err)), nil
6262 }
6363
64− messages, err := b.app.Messages.List(session.ID)
64+ messages, err := b.app.Messages.List(ctx, session.ID)
6565 if err != nil {
6666 return tools.NewTextErrorResponse(fmt.Sprintf("error listing messages: %s", err)), nil
6767 }
@@ -74,11 +74,11 @@ func (b *agentTool) Run(ctx context.Context, call tools.ToolCall) (tools.ToolRes
7474 return tools.NewTextErrorResponse("no assistant message found"), nil
7575 }
7676
77− updatedSession, err := b.app.Sessions.Get(session.ID)
77+ updatedSession, err := b.app.Sessions.Get(ctx, session.ID)
7878 if err != nil {
7979 return tools.NewTextErrorResponse(fmt.Sprintf("error: %s", err)), nil
8080 }
81− parentSession, err := b.app.Sessions.Get(b.parentSessionID)
81+ parentSession, err := b.app.Sessions.Get(ctx, b.parentSessionID)
8282 if err != nil {
8383 return tools.NewTextErrorResponse(fmt.Sprintf("error: %s", err)), nil
8484 }
@@ -87,7 +87,7 @@ func (b *agentTool) Run(ctx context.Context, call tools.ToolCall) (tools.ToolRes
8787 parentSession.PromptTokens += updatedSession.PromptTokens
8888 parentSession.CompletionTokens += updatedSession.CompletionTokens
8989
90− _, err = b.app.Sessions.Save(parentSession)
90+ _, err = b.app.Sessions.Save(ctx, parentSession)
9191 if err != nil {
9292 return tools.NewTextErrorResponse(fmt.Sprintf("error: %s", err)), nil
9393 }
internal/llm/agent/agent.go+27−26
@@ -48,7 +48,7 @@ func (c *agent) handleTitleGeneration(ctx context.Context, sessionID, content st
4848 return
4949 }
5050
51− session, err := c.Sessions.Get(sessionID)
51+ session, err := c.Sessions.Get(ctx, sessionID)
5252 if err != nil {
5353 return
5454 }
@@ -56,12 +56,12 @@ func (c *agent) handleTitleGeneration(ctx context.Context, sessionID, content st
5656 session.Title = response.Content
5757 session.Title = strings.TrimSpace(session.Title)
5858 session.Title = strings.ReplaceAll(session.Title, "\n", " ")
59− c.Sessions.Save(session)
59+ c.Sessions.Save(ctx, session)
6060 }
6161}
6262
63−func (c *agent) TrackUsage(sessionID string, model models.Model, usage provider.TokenUsage) error {
64− session, err := c.Sessions.Get(sessionID)
63+func (c *agent) TrackUsage(ctx context.Context, sessionID string, model models.Model, usage provider.TokenUsage) error {
64+ session, err := c.Sessions.Get(ctx, sessionID)
6565 if err != nil {
6666 return err
6767 }
@@ -75,11 +75,12 @@ func (c *agent) TrackUsage(sessionID string, model models.Model, usage provider.
7575 session.CompletionTokens += usage.OutputTokens
7676 session.PromptTokens += usage.InputTokens
7777
78− _, err = c.Sessions.Save(session)
78+ _, err = c.Sessions.Save(ctx, session)
7979 return err
8080}
8181
8282func (c *agent) processEvent(
83+ ctx context.Context,
8384 sessionID string,
8485 assistantMsg *message.Message,
8586 event provider.ProviderEvent,
@@ -87,10 +88,10 @@ func (c *agent) processEvent(
8788 switch event.Type {
8889 case provider.EventThinkingDelta:
8990 assistantMsg.AppendReasoningContent(event.Content)
90− return c.Messages.Update(*assistantMsg)
91+ return c.Messages.Update(ctx, *assistantMsg)
9192 case provider.EventContentDelta:
9293 assistantMsg.AppendContent(event.Content)
93− return c.Messages.Update(*assistantMsg)
94+ return c.Messages.Update(ctx, *assistantMsg)
9495 case provider.EventError:
9596 if errors.Is(event.Error, context.Canceled) {
9697 return nil
@@ -105,11 +106,11 @@ func (c *agent) processEvent(
105106 case provider.EventComplete:
106107 assistantMsg.SetToolCalls(event.Response.ToolCalls)
107108 assistantMsg.AddFinish(event.Response.FinishReason)
108− err := c.Messages.Update(*assistantMsg)
109+ err := c.Messages.Update(ctx, *assistantMsg)
109110 if err != nil {
110111 return err
111112 }
112− return c.TrackUsage(sessionID, c.model, event.Response.Usage)
113+ return c.TrackUsage(ctx, sessionID, c.model, event.Response.Usage)
113114 }
114115
115116 return nil
@@ -237,7 +238,7 @@ func (c *agent) handleToolExecution(
237238 for _, toolResult := range toolResults {
238239 parts = append(parts, toolResult)
239240 }
240− msg, err := c.Messages.Create(assistantMsg.SessionID, message.CreateMessageParams{
241+ msg, err := c.Messages.Create(ctx, assistantMsg.SessionID, message.CreateMessageParams{
241242 Role: message.Tool,
242243 Parts: parts,
243244 })
@@ -247,7 +248,7 @@ func (c *agent) handleToolExecution(
247248
248249func (c *agent) generate(ctx context.Context, sessionID string, content string) error {
249250 ctx = context.WithValue(ctx, tools.SessionIDContextKey, sessionID)
250− messages, err := c.Messages.List(sessionID)
251+ messages, err := c.Messages.List(ctx, sessionID)
251252 if err != nil {
252253 return err
253254 }
@@ -256,7 +257,7 @@ func (c *agent) generate(ctx context.Context, sessionID string, content string)
256257 go c.handleTitleGeneration(ctx, sessionID, content)
257258 }
258259
259− userMsg, err := c.Messages.Create(sessionID, message.CreateMessageParams{
260+ userMsg, err := c.Messages.Create(ctx, sessionID, message.CreateMessageParams{
260261 Role: message.User,
261262 Parts: []message.ContentPart{
262263 message.TextContent{
@@ -272,7 +273,7 @@ func (c *agent) generate(ctx context.Context, sessionID string, content string)
272273 for {
273274 select {
274275 case <-ctx.Done():
275− assistantMsg, err := c.Messages.Create(sessionID, message.CreateMessageParams{
276+ assistantMsg, err := c.Messages.Create(ctx, sessionID, message.CreateMessageParams{
276277 Role: message.Assistant,
277278 Parts: []message.ContentPart{},
278279 })
@@ -280,7 +281,7 @@ func (c *agent) generate(ctx context.Context, sessionID string, content string)
280281 return err
281282 }
282283 assistantMsg.AddFinish("canceled")
283− c.Messages.Update(assistantMsg)
284+ c.Messages.Update(ctx, assistantMsg)
284285 return context.Canceled
285286 default:
286287 // Continue processing
@@ -289,7 +290,7 @@ func (c *agent) generate(ctx context.Context, sessionID string, content string)
289290 eventChan, err := c.agent.StreamResponse(ctx, messages, c.tools)
290291 if err != nil {
291292 if errors.Is(err, context.Canceled) {
292− assistantMsg, err := c.Messages.Create(sessionID, message.CreateMessageParams{
293+ assistantMsg, err := c.Messages.Create(ctx, sessionID, message.CreateMessageParams{
293294 Role: message.Assistant,
294295 Parts: []message.ContentPart{},
295296 })
@@ -297,13 +298,13 @@ func (c *agent) generate(ctx context.Context, sessionID string, content string)
297298 return err
298299 }
299300 assistantMsg.AddFinish("canceled")
300− c.Messages.Update(assistantMsg)
301+ c.Messages.Update(ctx, assistantMsg)
301302 return context.Canceled
302303 }
303304 return err
304305 }
305306
306− assistantMsg, err := c.Messages.Create(sessionID, message.CreateMessageParams{
307+ assistantMsg, err := c.Messages.Create(ctx, sessionID, message.CreateMessageParams{
307308 Role: message.Assistant,
308309 Parts: []message.ContentPart{},
309310 Model: c.model.ID,
@@ -314,22 +315,22 @@ func (c *agent) generate(ctx context.Context, sessionID string, content string)
314315
315316 ctx = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)
316317 for event := range eventChan {
317− err = c.processEvent(sessionID, &assistantMsg, event)
318+ err = c.processEvent(ctx, sessionID, &assistantMsg, event)
318319 if err != nil {
319320 if errors.Is(err, context.Canceled) {
320321 assistantMsg.AddFinish("canceled")
321− c.Messages.Update(assistantMsg)
322+ c.Messages.Update(ctx, assistantMsg)
322323 return context.Canceled
323324 }
324325 assistantMsg.AddFinish("error:" + err.Error())
325− c.Messages.Update(assistantMsg)
326+ c.Messages.Update(ctx, assistantMsg)
326327 return err
327328 }
328329
329330 select {
330331 case <-ctx.Done():
331332 assistantMsg.AddFinish("canceled")
332− c.Messages.Update(assistantMsg)
333+ c.Messages.Update(ctx, assistantMsg)
333334 return context.Canceled
334335 default:
335336 }
@@ -339,7 +340,7 @@ func (c *agent) generate(ctx context.Context, sessionID string, content string)
339340 select {
340341 case <-ctx.Done():
341342 assistantMsg.AddFinish("canceled")
342− c.Messages.Update(assistantMsg)
343+ c.Messages.Update(ctx, assistantMsg)
343344 return context.Canceled
344345 default:
345346 // Continue processing
@@ -349,13 +350,13 @@ func (c *agent) generate(ctx context.Context, sessionID string, content string)
349350 if err != nil {
350351 if errors.Is(err, context.Canceled) {
351352 assistantMsg.AddFinish("canceled")
352− c.Messages.Update(assistantMsg)
353+ c.Messages.Update(ctx, assistantMsg)
353354 return context.Canceled
354355 }
355356 return err
356357 }
357358
358− c.Messages.Update(assistantMsg)
359+ c.Messages.Update(ctx, assistantMsg)
359360
360361 if len(assistantMsg.ToolCalls()) == 0 {
361362 break
@@ -370,7 +371,7 @@ func (c *agent) generate(ctx context.Context, sessionID string, content string)
370371 select {
371372 case <-ctx.Done():
372373 assistantMsg.AddFinish("canceled")
373− c.Messages.Update(assistantMsg)
374+ c.Messages.Update(ctx, assistantMsg)
374375 return context.Canceled
375376 default:
376377 // Continue processing
@@ -383,7 +384,7 @@ func getAgentProviders(ctx context.Context, model models.Model) (provider.Provid
383384 maxTokens := config.Get().Model.CoderMaxTokens
384385
385386 providerConfig, ok := config.Get().Providers[model.Provider]
386− if !ok || !providerConfig.Enabled {
387+ if !ok || providerConfig.Disabled {
387388 return nil, nil, errors.New("provider is not enabled")
388389 }
389390 var agentProvider provider.Provider
internal/llm/agent/coder.go+3−2
@@ -40,12 +40,13 @@ func NewCoderAgent(app *app.App) (Agent, error) {
4040 return nil, errors.New("model not supported")
4141 }
4242
43− agentProvider, titleGenerator, err := getAgentProviders(app.Context, model)
43+ ctx := context.Background()
44+ agentProvider, titleGenerator, err := getAgentProviders(ctx, model)
4445 if err != nil {
4546 return nil, err
4647 }
4748
48− otherTools := GetMcpTools(app.Context, app.Permissions)
49+ otherTools := GetMcpTools(ctx, app.Permissions)
4950 if len(app.LSPClients) > 0 {
5051 otherTools = append(otherTools, tools.NewDiagnosticsTool(app.LSPClients))
5152 }
internal/llm/agent/task.go+2−1
@@ -24,7 +24,8 @@ func NewTaskAgent(app *app.App) (Agent, error) {
2424 return nil, errors.New("model not supported")
2525 }
2626
27− agentProvider, titleGenerator, err := getAgentProviders(app.Context, model)
27+ ctx := context.Background()
28+ agentProvider, titleGenerator, err := getAgentProviders(ctx, model)
2829 if err != nil {
2930 return nil, err
3031 }
internal/message/message.go+22−24
@@ -20,34 +20,32 @@ type CreateMessageParams struct {
2020
2121type Service interface {
2222 pubsub.Suscriber[Message]
23− Create(sessionID string, params CreateMessageParams) (Message, error)
24− Update(message Message) error
25− Get(id string) (Message, error)
26− List(sessionID string) ([]Message, error)
27− Delete(id string) error
28− DeleteSessionMessages(sessionID string) error
23+ Create(ctx context.Context, sessionID string, params CreateMessageParams) (Message, error)
24+ Update(ctx context.Context, message Message) error
25+ Get(ctx context.Context, id string) (Message, error)
26+ List(ctx context.Context, sessionID string) ([]Message, error)
27+ Delete(ctx context.Context, id string) error
28+ DeleteSessionMessages(ctx context.Context, sessionID string) error
2929}
3030
3131type service struct {
3232 *pubsub.Broker[Message]
33− q db.Querier
34− ctx context.Context
33+ q db.Querier
3534}
3635
37−func NewService(ctx context.Context, q db.Querier) Service {
36+func NewService(q db.Querier) Service {
3837 return &service{
3938 Broker: pubsub.NewBroker[Message](),
4039 q: q,
41− ctx: ctx,
4240 }
4341}
4442
45−func (s *service) Delete(id string) error {
46− message, err := s.Get(id)
43+func (s *service) Delete(ctx context.Context, id string) error {
44+ message, err := s.Get(ctx, id)
4745 if err != nil {
4846 return err
4947 }
50− err = s.q.DeleteMessage(s.ctx, message.ID)
48+ err = s.q.DeleteMessage(ctx, message.ID)
5149 if err != nil {
5250 return err
5351 }
@@ -55,7 +53,7 @@ func (s *service) Delete(id string) error {
5553 return nil
5654}
5755
58−func (s *service) Create(sessionID string, params CreateMessageParams) (Message, error) {
56+func (s *service) Create(ctx context.Context, sessionID string, params CreateMessageParams) (Message, error) {
5957 if params.Role != Assistant {
6058 params.Parts = append(params.Parts, Finish{
6159 Reason: "stop",
@@ -66,7 +64,7 @@ func (s *service) Create(sessionID string, params CreateMessageParams) (Message,
6664 return Message{}, err
6765 }
6866
69− dbMessage, err := s.q.CreateMessage(s.ctx, db.CreateMessageParams{
67+ dbMessage, err := s.q.CreateMessage(ctx, db.CreateMessageParams{
7068 ID: uuid.New().String(),
7169 SessionID: sessionID,
7270 Role: string(params.Role),
@@ -84,14 +82,14 @@ func (s *service) Create(sessionID string, params CreateMessageParams) (Message,
8482 return message, nil
8583}
8684
87−func (s *service) DeleteSessionMessages(sessionID string) error {
88− messages, err := s.List(sessionID)
85+func (s *service) DeleteSessionMessages(ctx context.Context, sessionID string) error {
86+ messages, err := s.List(ctx, sessionID)
8987 if err != nil {
9088 return err
9189 }
9290 for _, message := range messages {
9391 if message.SessionID == sessionID {
94− err = s.Delete(message.ID)
92+ err = s.Delete(ctx, message.ID)
9593 if err != nil {
9694 return err
9795 }
@@ -100,7 +98,7 @@ func (s *service) DeleteSessionMessages(sessionID string) error {
10098 return nil
10199}
102100
103−func (s *service) Update(message Message) error {
101+func (s *service) Update(ctx context.Context, message Message) error {
104102 parts, err := marshallParts(message.Parts)
105103 if err != nil {
106104 return err
@@ -110,7 +108,7 @@ func (s *service) Update(message Message) error {
110108 finishedAt.Int64 = f.Time
111109 finishedAt.Valid = true
112110 }
113− err = s.q.UpdateMessage(s.ctx, db.UpdateMessageParams{
111+ err = s.q.UpdateMessage(ctx, db.UpdateMessageParams{
114112 ID: message.ID,
115113 Parts: string(parts),
116114 FinishedAt: finishedAt,
@@ -122,16 +120,16 @@ func (s *service) Update(message Message) error {
122120 return nil
123121}
124122
125−func (s *service) Get(id string) (Message, error) {
126− dbMessage, err := s.q.GetMessage(s.ctx, id)
123+func (s *service) Get(ctx context.Context, id string) (Message, error) {
124+ dbMessage, err := s.q.GetMessage(ctx, id)
127125 if err != nil {
128126 return Message{}, err
129127 }
130128 return s.fromDBItem(dbMessage)
131129}
132130
133−func (s *service) List(sessionID string) ([]Message, error) {
134− dbMessages, err := s.q.ListMessagesBySession(s.ctx, sessionID)
131+func (s *service) List(ctx context.Context, sessionID string) ([]Message, error) {
132+ dbMessages, err := s.q.ListMessagesBySession(ctx, sessionID)
135133 if err != nil {
136134 return nil, err
137135 }
internal/session/session.go+21−23
@@ -23,22 +23,21 @@ type Session struct {
2323
2424type Service interface {
2525 pubsub.Suscriber[Session]
26− Create(title string) (Session, error)
27− CreateTaskSession(toolCallID, parentSessionID, title string) (Session, error)
28− Get(id string) (Session, error)
29− List() ([]Session, error)
30− Save(session Session) (Session, error)
31− Delete(id string) error
26+ Create(ctx context.Context, title string) (Session, error)
27+ CreateTaskSession(ctx context.Context, toolCallID, parentSessionID, title string) (Session, error)
28+ Get(ctx context.Context, id string) (Session, error)
29+ List(ctx context.Context) ([]Session, error)
30+ Save(ctx context.Context, session Session) (Session, error)
31+ Delete(ctx context.Context, id string) error
3232}
3333
3434type service struct {
3535 *pubsub.Broker[Session]
36− q db.Querier
37− ctx context.Context
36+ q db.Querier
3837}
3938
40−func (s *service) Create(title string) (Session, error) {
41− dbSession, err := s.q.CreateSession(s.ctx, db.CreateSessionParams{
39+func (s *service) Create(ctx context.Context, title string) (Session, error) {
40+ dbSession, err := s.q.CreateSession(ctx, db.CreateSessionParams{
4241 ID: uuid.New().String(),
4342 Title: title,
4443 })
@@ -50,8 +49,8 @@ func (s *service) Create(title string) (Session, error) {
5049 return session, nil
5150}
5251
53−func (s *service) CreateTaskSession(toolCallID, parentSessionID, title string) (Session, error) {
54− dbSession, err := s.q.CreateSession(s.ctx, db.CreateSessionParams{
52+func (s *service) CreateTaskSession(ctx context.Context, toolCallID, parentSessionID, title string) (Session, error) {
53+ dbSession, err := s.q.CreateSession(ctx, db.CreateSessionParams{
5554 ID: toolCallID,
5655 ParentSessionID: sql.NullString{String: parentSessionID, Valid: true},
5756 Title: title,
@@ -64,12 +63,12 @@ func (s *service) CreateTaskSession(toolCallID, parentSessionID, title string) (
6463 return session, nil
6564}
6665
67−func (s *service) Delete(id string) error {
68− session, err := s.Get(id)
66+func (s *service) Delete(ctx context.Context, id string) error {
67+ session, err := s.Get(ctx, id)
6968 if err != nil {
7069 return err
7170 }
72− err = s.q.DeleteSession(s.ctx, session.ID)
71+ err = s.q.DeleteSession(ctx, session.ID)
7372 if err != nil {
7473 return err
7574 }
@@ -77,16 +76,16 @@ func (s *service) Delete(id string) error {
7776 return nil
7877}
7978
80−func (s *service) Get(id string) (Session, error) {
81− dbSession, err := s.q.GetSessionByID(s.ctx, id)
79+func (s *service) Get(ctx context.Context, id string) (Session, error) {
80+ dbSession, err := s.q.GetSessionByID(ctx, id)
8281 if err != nil {
8382 return Session{}, err
8483 }
8584 return s.fromDBItem(dbSession), nil
8685}
8786
88−func (s *service) Save(session Session) (Session, error) {
89− dbSession, err := s.q.UpdateSession(s.ctx, db.UpdateSessionParams{
87+func (s *service) Save(ctx context.Context, session Session) (Session, error) {
88+ dbSession, err := s.q.UpdateSession(ctx, db.UpdateSessionParams{
9089 ID: session.ID,
9190 Title: session.Title,
9291 PromptTokens: session.PromptTokens,
@@ -101,8 +100,8 @@ func (s *service) Save(session Session) (Session, error) {
101100 return session, nil
102101}
103102
104−func (s *service) List() ([]Session, error) {
105− dbSessions, err := s.q.ListSessions(s.ctx)
103+func (s *service) List(ctx context.Context) ([]Session, error) {
104+ dbSessions, err := s.q.ListSessions(ctx)
106105 if err != nil {
107106 return nil, err
108107 }
@@ -127,11 +126,10 @@ func (s service) fromDBItem(item db.Session) Session {
127126 }
128127}
129128
130−func NewService(ctx context.Context, q db.Querier) Service {
129+func NewService(q db.Querier) Service {
131130 broker := pubsub.NewBroker[Session]()
132131 return &service{
133132 broker,
134133 q,
135− ctx,
136134 }
137135}
internal/tui/components/chat/messages.go+3−2
@@ -1,6 +1,7 @@
11package chat
22
33import (
4+ "context"
45 "encoding/json"
56 "fmt"
67 "math"
@@ -324,7 +325,7 @@ func (m *messagesCmp) renderToolCall(toolCall message.ToolCall, isNested bool) s
324325
325326 innerToolCalls := make([]string, 0)
326327 if toolCall.Name == agent.AgentToolName {
327− messages, _ := m.app.Messages.List(toolCall.ID)
328+ messages, _ := m.app.Messages.List(context.Background(), toolCall.ID)
328329 toolCalls := make([]message.ToolCall, 0)
329330 for _, v := range messages {
330331 toolCalls = append(toolCalls, v.ToolCalls()...)
@@ -554,7 +555,7 @@ func (m *messagesCmp) GetSize() (int, int) {
554555
555556func (m *messagesCmp) SetSession(session session.Session) tea.Cmd {
556557 m.session = session
557− messages, err := m.app.Messages.List(session.ID)
558+ messages, err := m.app.Messages.List(context.Background(), session.ID)
558559 if err != nil {
559560 return util.ReportError(err)
560561 }
internal/tui/components/repl/editor.go+2−2
@@ -160,7 +160,7 @@ func (m *editorCmp) Send() tea.Cmd {
160160 return util.ReportWarn("Assistant is still working on the previous message")
161161 }
162162
163− messages, err := m.app.Messages.List(m.sessionID)
163+ messages, err := m.app.Messages.List(context.Background(), m.sessionID)
164164 if err != nil {
165165 return util.ReportError(err)
166166 }
@@ -177,7 +177,7 @@ func (m *editorCmp) Send() tea.Cmd {
177177 if len(content) == 0 {
178178 return util.ReportWarn("Message is empty")
179179 }
180− ctx, cancel := context.WithCancel(m.app.Context)
180+ ctx, cancel := context.WithCancel(context.Background())
181181 m.cancelMessage = cancel
182182 go func() {
183183 defer cancel()
internal/tui/components/repl/messages.go+4−3
@@ -1,6 +1,7 @@
11package repl
22
33import (
4+ "context"
45 "encoding/json"
56 "fmt"
67 "sort"
@@ -77,8 +78,8 @@ func (m *messagesCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
7778 m.session = msg.Payload
7879 }
7980 case SelectedSessionMsg:
80− m.session, _ = m.app.Sessions.Get(msg.SessionID)
81− m.messages, _ = m.app.Messages.List(m.session.ID)
81+ m.session, _ = m.app.Sessions.Get(context.Background(), msg.SessionID)
82+ m.messages, _ = m.app.Messages.List(context.Background(), m.session.ID)
8283 m.renderView()
8384 m.viewport.GotoBottom()
8485 }
@@ -259,7 +260,7 @@ func (m *messagesCmp) renderMessageWithToolCall(content string, tools []message.
259260
260261 runningIndicator := runningStyle.Render(fmt.Sprintf("%s Running...", styles.SpinnerIcon))
261262 allParts = append(allParts, leftPadding.Render(runningIndicator))
262− taskSessionMessages, _ := m.app.Messages.List(toolCall.ID)
263+ taskSessionMessages, _ := m.app.Messages.List(context.Background(), toolCall.ID)
263264 for _, msg := range taskSessionMessages {
264265 if msg.Role == message.Assistant {
265266 for _, toolCall := range msg.ToolCalls() {
internal/tui/components/repl/sessions.go+3−1
@@ -1,6 +1,7 @@
11package repl
22
33import (
4+ "context"
45 "fmt"
56 "strings"
67
@@ -57,12 +58,13 @@ var sessionKeyMapValue = sessionsKeyMap{
5758}
5859
5960func (i *sessionsCmp) Init() tea.Cmd {
60− existing, err := i.app.Sessions.List()
61+ existing, err := i.app.Sessions.List(context.Background())
6162 if err != nil {
6263 return util.ReportError(err)
6364 }
6465 if len(existing) == 0 || existing[0].MessageCount > 0 {
6566 newSession, err := i.app.Sessions.Create(
67+ context.Background(),
6668 "New Session",
6769 )
6870 if err != nil {
internal/tui/page/chat.go+5−3
@@ -1,6 +1,8 @@
11package page
22
33import (
4+ "context"
5+
46 "github.com/charmbracelet/bubbles/key"
57 tea "github.com/charmbracelet/bubbletea"
68 "github.com/kujtimiihoxha/termai/internal/app"
@@ -36,7 +38,7 @@ func (p *chatPage) Init() tea.Cmd {
3638 p.layout.Init(),
3739 }
3840
39− sessions, _ := p.app.Sessions.List()
41+ sessions, _ := p.app.Sessions.List(context.Background())
4042 if len(sessions) > 0 {
4143 p.session = sessions[0]
4244 cmd := p.setSidebar()
@@ -92,7 +94,7 @@ func (p *chatPage) clearSidebar() {
9294func (p *chatPage) sendMessage(text string) tea.Cmd {
9395 var cmds []tea.Cmd
9496 if p.session.ID == "" {
95− session, err := p.app.Sessions.Create("New Session")
97+ session, err := p.app.Sessions.Create(context.Background(), "New Session")
9698 if err != nil {
9799 return util.ReportError(err)
98100 }
@@ -110,7 +112,7 @@ func (p *chatPage) sendMessage(text string) tea.Cmd {
110112 return util.ReportError(err)
111113 }
112114 go func() {
113− a.Generate(p.app.Context, p.session.ID, text)
115+ a.Generate(context.Background(), p.session.ID, text)
114116 }()
115117
116118 return tea.Batch(cmds...)
internal/tui/tui.go+4−2
@@ -1,6 +1,8 @@
11package tui
22
33import (
4+ "context"
5+
46 "github.com/charmbracelet/bubbles/key"
57 tea "github.com/charmbracelet/bubbletea"
68 "github.com/charmbracelet/lipgloss"
@@ -184,7 +186,7 @@ func (a appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
184186 }
185187 case key.Matches(msg, replKeyMap):
186188 if a.currentPage == page.ReplPage {
187− sessions, err := a.app.Sessions.List()
189+ sessions, err := a.app.Sessions.List(context.Background())
188190 if err != nil {
189191 return a, util.CmdHandler(util.ReportError(err))
190192 }
@@ -192,7 +194,7 @@ func (a appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
192194 if lastSession.MessageCount == 0 {
193195 return a, util.CmdHandler(repl.SelectedSessionMsg{SessionID: lastSession.ID})
194196 }
195− s, err := a.app.Sessions.Create("New Session")
197+ s, err := a.app.Sessions.Create(context.Background(), "New Session")
196198 if err != nil {
197199 return a, util.CmdHandler(util.ReportError(err))
198200 }