git.iko.soy/crush
on main help sign in
cleanup app, config and root
3ad983dKujtim Hoxha
.gitignore+1−1
compare with main
unifiedsplit
@@ -41,6 +41,6 @@ debug.log
41.env
41.env
42.env.local
42.env.local
43
43
44−.termai
44+.opencode
45
45
46internal/assets/diff/index.mjs
46internal/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
2
2
3import (
3import (
4 "context"
4 "context"
5− "log/slog"
5+ "fmt"
6 "os"
6 "os"
7 "sync"
7 "sync"
8+ "time"
8
9
9 tea "github.com/charmbracelet/bubbletea"
10 tea "github.com/charmbracelet/bubbletea"
10 "github.com/kujtimiihoxha/termai/internal/app"
11 "github.com/kujtimiihoxha/termai/internal/app"
@@ -13,6 +14,7 @@ import (
13 "github.com/kujtimiihoxha/termai/internal/db"
14 "github.com/kujtimiihoxha/termai/internal/db"
14 "github.com/kujtimiihoxha/termai/internal/llm/agent"
15 "github.com/kujtimiihoxha/termai/internal/llm/agent"
15 "github.com/kujtimiihoxha/termai/internal/logging"
16 "github.com/kujtimiihoxha/termai/internal/logging"
17+ "github.com/kujtimiihoxha/termai/internal/pubsub"
16 "github.com/kujtimiihoxha/termai/internal/tui"
18 "github.com/kujtimiihoxha/termai/internal/tui"
17 zone "github.com/lrstanley/bubblezone"
19 zone "github.com/lrstanley/bubblezone"
18 "github.com/spf13/cobra"
20 "github.com/spf13/cobra"
@@ -23,111 +25,229 @@ var rootCmd = &cobra.Command{
23 Short: "A terminal ai assistant",
25 Short: "A terminal ai assistant",
24 Long: `A terminal ai assistant`,
26 Long: `A terminal ai assistant`,
25 RunE: func(cmd *cobra.Command, args []string) error {
27 RunE: func(cmd *cobra.Command, args []string) error {
28+ // If the help flag is set, show the help message
26 if cmd.Flag("help").Changed {
29 if cmd.Flag("help").Changed {
27 cmd.Help()
30 cmd.Help()
28 return nil
31 return nil
29 }
32 }
33+
34+ // Load the config
30 debug, _ := cmd.Flags().GetBool("debug")
35 debug, _ := cmd.Flags().GetBool("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+ }
31− err := config.Load(debug)
50+ _, err := config.Load(cwd, debug)
32 if err != nil {
51 if err != nil {
33 return err
52 return err
34 }
53 }
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)
44
54
45 err = assets.WriteAssets()
55 err = assets.WriteAssets()
46 if err != nil {
56 if err != nil {
47− return err
57+ logging.Error("Error writing assets: %v", err)
48 }
58 }
49
59
60+ // Connect DB, this will also run migrations
50 conn, err := db.Connect()
61 conn, err := db.Connect()
51 if err != nil {
62 if err != nil {
52 return err
63 return err
53 }
64 }
65+
66+ // Create main context for the application
54− ctx := context.Background()
67+ ctx, cancel := context.WithCancel(context.Background())
68+ defer cancel()
55
69
56 app := app.New(ctx, conn)
70 app := app.New(ctx, conn)
57− logging.Info("Starting termai...")
71+
72+ // Set up the TUI
58 zone.NewGlobal()
73 zone.NewGlobal()
59− tui := tea.NewProgram(
74+ program := tea.NewProgram(
60 tui.New(app),
75 tui.New(app),
61 tea.WithAltScreen(),
76 tea.WithAltScreen(),
62 tea.WithMouseCellMotion(),
77 tea.WithMouseCellMotion(),
63 )
78 )
64− logging.Info("Setting up subscriptions...")
65− ch, unsub := setupSubscriptions(app)
66− defer unsub()
67
79
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
68 go func() {
92 go func() {
69− // Set this up once
93+ defer tuiWg.Done()
70− agent.GetMcpTools(ctx, app.Permissions)
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
71− for msg := range ch {
106+ case msg, ok := <-ch:
107+ if !ok {
108+ logging.Info("TUI message channel closed")
109+ return
110+ }
72− tui.Send(msg)
111+ program.Send(msg)
112+ }
73 }
113 }
74 }()
114 }()
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+
75− if _, err := tui.Run(); err != nil {
137+ if err != nil {
138+ logging.Error("TUI error: %v", err)
76− return err
139+ return fmt.Errorf("TUI error: %v", err)
77 }
140 }
141+
142+ logging.Info("TUI exited with result: %v", result)
78 return nil
143 return nil
79 },
144 },
80}
145}
81
146
82−func setupSubscriptions(app *app.App) (chan tea.Msg, func()) {
147+// attemptTUIRecovery tries to recover the TUI after a panic
83− ch := make(chan tea.Msg)
148+func attemptTUIRecovery(program *tea.Program) {
84− wg := sync.WaitGroup{}
149+ logging.Info("Attempting to recover TUI after panic")
85− ctx, cancel := context.WithCancel(app.Context)
150+
86− {
151+ // We could try to restart the TUI or gracefully exit
87− sub := logging.Subscribe(ctx)
152+ // For now, we'll just quit the program to avoid further issues
88− wg.Add(1)
153+ program.Quit()
154+}
155+
156+func initMCPTools(ctx context.Context, app *app.App) {
89− go func() {
157+ go func() {
90− for ev := range sub {
158+ defer func() {
91− ch <- ev
159+ if r := recover(); r != nil {
160+ logging.Error("Panic in MCP goroutine: %v", r)
92 }
161 }
93− wg.Done()
94 }()
162 }()
95− }
163+
96− {
164+ // Create a context with timeout for the initial MCP tools fetch
97− sub := app.Sessions.Subscribe(ctx)
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+) {
98− wg.Add(1)
181+ wg.Add(1)
99− go func() {
182+ go func() {
100− for ev := range sub {
183+ defer wg.Done()
101− ch <- ev
184+ defer func() {
185+ if r := recover(); r != nil {
186+ logging.Error("Panic in %s subscription goroutine: %v", name, r)
102 }
187 }
103− wg.Done()
104 }()
188 }()
105− }
189+
106− {
190+ for {
107− sub := app.Messages.Subscribe(ctx)
191+ select {
108− wg.Add(1)
192+ case event, ok := <-subscriber(ctx):
109− go func() {
193+ if !ok {
110− for ev := range sub {
194+ logging.Info("%s subscription channel closed", name)
111− ch <- ev
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
112 }
213 }
113− wg.Done()
214+ }
114− }()
215+ }()
115− }
216+}
116− {
217+
117− sub := app.Permissions.Subscribe(ctx)
218+func setupSubscriptions(app *app.App) (chan tea.Msg, func()) {
118− wg.Add(1)
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{})
119 go func() {
236 go func() {
120− for ev := range sub {
237+ wg.Wait()
121− ch <- ev
238+ close(waitCh)
122− }
123− wg.Done()
124 }()
239 }()
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
125 }
249 }
126− return ch, func() {
250+ return ch, cleanupFunc
127− cancel()
128− wg.Wait()
129− close(ch)
130− }
131}
251}
132
252
133func Execute() {
253func Execute() {
@@ -139,5 +259,6 @@ func Execute() {
139
259
140func init() {
260func init() {
141 rootCmd.Flags().BoolP("help", "h", false, "Help")
261 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")
143}
264}
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
83// Load initializes the configuration from environment variables and config files.
83// Load initializes the configuration from environment variables and config files.
84// If debug is true, debug mode is enabled and log level is set to debug.
84// If debug is true, debug mode is enabled and log level is set to debug.
85// It returns an error if configuration loading fails.
85// 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) {
87 if cfg != nil {
87 if cfg != nil {
88− return nil
88+ return cfg, nil
89 }
89 }
90
90
91 cfg = &Config{
91 cfg = &Config{
@@ -101,7 +101,7 @@ func Load(workingDir string, debug bool) error {
101
101
102 // Read global config
102 // Read global config
103 if err := readConfig(viper.ReadInConfig()); err != nil {
103 if err := readConfig(viper.ReadInConfig()); err != nil {
104− return err
104+ return cfg, err
105 }
105 }
106
106
107 // Load and merge local config
107 // Load and merge local config
@@ -109,7 +109,7 @@ func Load(workingDir string, debug bool) error {
109
109
110 // Apply configuration to the struct
110 // Apply configuration to the struct
111 if err := viper.Unmarshal(cfg); err != nil {
111 if err := viper.Unmarshal(cfg); err != nil {
112− return err
112+ return cfg, fmt.Errorf("failed to unmarshal config: %w", err)
113 }
113 }
114
114
115 applyDefaultValues()
115 applyDefaultValues()
@@ -123,7 +123,7 @@ func Load(workingDir string, debug bool) error {
123 Level: defaultLevel,
123 Level: defaultLevel,
124 }))
124 }))
125 slog.SetDefault(logger)
125 slog.SetDefault(logger)
126− return nil
126+ return cfg, nil
127}
127}
128
128
129// configureViper sets up viper's configuration paths and environment variables.
129// configureViper sets up viper's configuration paths and environment variables.
@@ -237,7 +237,7 @@ func readConfig(err error) error {
237 return nil
237 return nil
238 }
238 }
239
239
240− return err
240+ return fmt.Errorf("failed to read config: %w", err)
241}
241}
242
242
243// mergeLocalConfig loads and merges configuration from the local directory.
243// mergeLocalConfig loads and merges configuration from the local directory.
@@ -264,14 +264,6 @@ func applyDefaultValues() {
264 }
264 }
265}
265}
266
266
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−
275// Get returns the current configuration.
267// Get returns the current configuration.
276// It's safe to call this function multiple times.
268// It's safe to call this function multiple times.
277func Get() *Config {
269func Get() *Config {
internal/history/file.go+35−38
@@ -27,45 +27,43 @@ type File struct {
27
27
28type Service interface {
28type Service interface {
29 pubsub.Suscriber[File]
29 pubsub.Suscriber[File]
30− Create(sessionID, path, content string) (File, error)
30+ Create(ctx context.Context, sessionID, path, content string) (File, error)
31− CreateVersion(sessionID, path, content string) (File, error)
31+ CreateVersion(ctx context.Context, sessionID, path, content string) (File, error)
32− Get(id string) (File, error)
32+ Get(ctx context.Context, id string) (File, error)
33− GetByPathAndSession(path, sessionID string) (File, error)
33+ GetByPathAndSession(ctx context.Context, path, sessionID string) (File, error)
34− ListBySession(sessionID string) ([]File, error)
34+ ListBySession(ctx context.Context, sessionID string) ([]File, error)
35− ListLatestSessionFiles(sessionID string) ([]File, error)
35+ ListLatestSessionFiles(ctx context.Context, sessionID string) ([]File, error)
36− Update(file File) (File, error)
36+ Update(ctx context.Context, file File) (File, error)
37− Delete(id string) error
37+ Delete(ctx context.Context, id string) error
38− DeleteSessionFiles(sessionID string) error
38+ DeleteSessionFiles(ctx context.Context, sessionID string) error
39}
39}
40
40
41type service struct {
41type service struct {
42 *pubsub.Broker[File]
42 *pubsub.Broker[File]
43− q db.Querier
43+ q db.Querier
44− ctx context.Context
45}
44}
46
45
47−func NewService(ctx context.Context, q db.Querier) Service {
46+func NewService(q db.Querier) Service {
48 return &service{
47 return &service{
49 Broker: pubsub.NewBroker[File](),
48 Broker: pubsub.NewBroker[File](),
50 q: q,
49 q: q,
51− ctx: ctx,
52 }
50 }
53}
51}
54
52
55−func (s *service) Create(sessionID, path, content string) (File, error) {
53+func (s *service) Create(ctx context.Context, sessionID, path, content string) (File, error) {
56− return s.createWithVersion(sessionID, path, content, InitialVersion)
54+ return s.createWithVersion(ctx, sessionID, path, content, InitialVersion)
57}
55}
58
56
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) {
60 // Get the latest version for this path
58 // Get the latest version for this path
61− files, err := s.q.ListFilesByPath(s.ctx, path)
59+ files, err := s.q.ListFilesByPath(ctx, path)
62 if err != nil {
60 if err != nil {
63 return File{}, err
61 return File{}, err
64 }
62 }
65
63
66 if len(files) == 0 {
64 if len(files) == 0 {
67 // No previous versions, create initial
65 // No previous versions, create initial
68− return s.Create(sessionID, path, content)
66+ return s.Create(ctx, sessionID, path, content)
69 }
67 }
70
68
71 // Get the latest version
69 // Get the latest version
@@ -89,11 +87,11 @@ func (s *service) CreateVersion(sessionID, path, content string) (File, error) {
89 nextVersion = fmt.Sprintf("v%d", latestFile.CreatedAt)
87 nextVersion = fmt.Sprintf("v%d", latestFile.CreatedAt)
90 }
88 }
91
89
92− return s.createWithVersion(sessionID, path, content, nextVersion)
90+ return s.createWithVersion(ctx, sessionID, path, content, nextVersion)
93}
91}
94
92
95−func (s *service) createWithVersion(sessionID, path, content, version string) (File, error) {
93+func (s *service) createWithVersion(ctx context.Context, sessionID, path, content, version string) (File, error) {
96− dbFile, err := s.q.CreateFile(s.ctx, db.CreateFileParams{
94+ dbFile, err := s.q.CreateFile(ctx, db.CreateFileParams{
97 ID: uuid.New().String(),
95 ID: uuid.New().String(),
98 SessionID: sessionID,
96 SessionID: sessionID,
99 Path: path,
97 Path: path,
@@ -108,16 +106,16 @@ func (s *service) createWithVersion(sessionID, path, content, version string) (F
108 return file, nil
106 return file, nil
109}
107}
110
108
111−func (s *service) Get(id string) (File, error) {
109+func (s *service) Get(ctx context.Context, id string) (File, error) {
112− dbFile, err := s.q.GetFile(s.ctx, id)
110+ dbFile, err := s.q.GetFile(ctx, id)
113 if err != nil {
111 if err != nil {
114 return File{}, err
112 return File{}, err
115 }
113 }
116 return s.fromDBItem(dbFile), nil
114 return s.fromDBItem(dbFile), nil
117}
115}
118
116
119−func (s *service) GetByPathAndSession(path, sessionID string) (File, error) {
117+func (s *service) GetByPathAndSession(ctx context.Context, path, sessionID string) (File, error) {
120− dbFile, err := s.q.GetFileByPathAndSession(s.ctx, db.GetFileByPathAndSessionParams{
118+ dbFile, err := s.q.GetFileByPathAndSession(ctx, db.GetFileByPathAndSessionParams{
121 Path: path,
119 Path: path,
122 SessionID: sessionID,
120 SessionID: sessionID,
123 })
121 })
@@ -127,8 +125,8 @@ func (s *service) GetByPathAndSession(path, sessionID string) (File, error) {
127 return s.fromDBItem(dbFile), nil
125 return s.fromDBItem(dbFile), nil
128}
126}
129
127
130−func (s *service) ListBySession(sessionID string) ([]File, error) {
128+func (s *service) ListBySession(ctx context.Context, sessionID string) ([]File, error) {
131− dbFiles, err := s.q.ListFilesBySession(s.ctx, sessionID)
129+ dbFiles, err := s.q.ListFilesBySession(ctx, sessionID)
132 if err != nil {
130 if err != nil {
133 return nil, err
131 return nil, err
134 }
132 }
@@ -139,8 +137,8 @@ func (s *service) ListBySession(sessionID string) ([]File, error) {
139 return files, nil
137 return files, nil
140}
138}
141
139
142−func (s *service) ListLatestSessionFiles(sessionID string) ([]File, error) {
140+func (s *service) ListLatestSessionFiles(ctx context.Context, sessionID string) ([]File, error) {
143− dbFiles, err := s.q.ListLatestSessionFiles(s.ctx, sessionID)
141+ dbFiles, err := s.q.ListLatestSessionFiles(ctx, sessionID)
144 if err != nil {
142 if err != nil {
145 return nil, err
143 return nil, err
146 }
144 }
@@ -151,8 +149,8 @@ func (s *service) ListLatestSessionFiles(sessionID string) ([]File, error) {
151 return files, nil
149 return files, nil
152}
150}
153
151
154−func (s *service) Update(file File) (File, error) {
152+func (s *service) Update(ctx context.Context, file File) (File, error) {
155− dbFile, err := s.q.UpdateFile(s.ctx, db.UpdateFileParams{
153+ dbFile, err := s.q.UpdateFile(ctx, db.UpdateFileParams{
156 ID: file.ID,
154 ID: file.ID,
157 Content: file.Content,
155 Content: file.Content,
158 Version: file.Version,
156 Version: file.Version,
@@ -165,12 +163,12 @@ func (s *service) Update(file File) (File, error) {
165 return updatedFile, nil
163 return updatedFile, nil
166}
164}
167
165
168−func (s *service) Delete(id string) error {
166+func (s *service) Delete(ctx context.Context, id string) error {
169− file, err := s.Get(id)
167+ file, err := s.Get(ctx, id)
170 if err != nil {
168 if err != nil {
171 return err
169 return err
172 }
170 }
173− err = s.q.DeleteFile(s.ctx, id)
171+ err = s.q.DeleteFile(ctx, id)
174 if err != nil {
172 if err != nil {
175 return err
173 return err
176 }
174 }
@@ -178,13 +176,13 @@ func (s *service) Delete(id string) error {
178 return nil
176 return nil
179}
177}
180
178
181−func (s *service) DeleteSessionFiles(sessionID string) error {
179+func (s *service) DeleteSessionFiles(ctx context.Context, sessionID string) error {
182− files, err := s.ListBySession(sessionID)
180+ files, err := s.ListBySession(ctx, sessionID)
183 if err != nil {
181 if err != nil {
184 return err
182 return err
185 }
183 }
186 for _, file := range files {
184 for _, file := range files {
187− err = s.Delete(file.ID)
185+ err = s.Delete(ctx, file.ID)
188 if err != nil {
186 if err != nil {
189 return err
187 return err
190 }
188 }
@@ -203,4 +201,3 @@ func (s *service) fromDBItem(item db.File) File {
203 UpdatedAt: item.UpdatedAt,
201 UpdatedAt: item.UpdatedAt,
204 }
202 }
205}
203}
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
51 return tools.NewTextErrorResponse(fmt.Sprintf("error creating agent: %s", err)), nil
51 return tools.NewTextErrorResponse(fmt.Sprintf("error creating agent: %s", err)), nil
52 }
52 }
53
53
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")
55 if err != nil {
55 if err != nil {
56 return tools.NewTextErrorResponse(fmt.Sprintf("error creating session: %s", err)), nil
56 return tools.NewTextErrorResponse(fmt.Sprintf("error creating session: %s", err)), nil
57 }
57 }
@@ -61,7 +61,7 @@ func (b *agentTool) Run(ctx context.Context, call tools.ToolCall) (tools.ToolRes
61 return tools.NewTextErrorResponse(fmt.Sprintf("error generating agent: %s", err)), nil
61 return tools.NewTextErrorResponse(fmt.Sprintf("error generating agent: %s", err)), nil
62 }
62 }
63
63
64− messages, err := b.app.Messages.List(session.ID)
64+ messages, err := b.app.Messages.List(ctx, session.ID)
65 if err != nil {
65 if err != nil {
66 return tools.NewTextErrorResponse(fmt.Sprintf("error listing messages: %s", err)), nil
66 return tools.NewTextErrorResponse(fmt.Sprintf("error listing messages: %s", err)), nil
67 }
67 }
@@ -74,11 +74,11 @@ func (b *agentTool) Run(ctx context.Context, call tools.ToolCall) (tools.ToolRes
74 return tools.NewTextErrorResponse("no assistant message found"), nil
74 return tools.NewTextErrorResponse("no assistant message found"), nil
75 }
75 }
76
76
77− updatedSession, err := b.app.Sessions.Get(session.ID)
77+ updatedSession, err := b.app.Sessions.Get(ctx, session.ID)
78 if err != nil {
78 if err != nil {
79 return tools.NewTextErrorResponse(fmt.Sprintf("error: %s", err)), nil
79 return tools.NewTextErrorResponse(fmt.Sprintf("error: %s", err)), nil
80 }
80 }
81− parentSession, err := b.app.Sessions.Get(b.parentSessionID)
81+ parentSession, err := b.app.Sessions.Get(ctx, b.parentSessionID)
82 if err != nil {
82 if err != nil {
83 return tools.NewTextErrorResponse(fmt.Sprintf("error: %s", err)), nil
83 return tools.NewTextErrorResponse(fmt.Sprintf("error: %s", err)), nil
84 }
84 }
@@ -87,7 +87,7 @@ func (b *agentTool) Run(ctx context.Context, call tools.ToolCall) (tools.ToolRes
87 parentSession.PromptTokens += updatedSession.PromptTokens
87 parentSession.PromptTokens += updatedSession.PromptTokens
88 parentSession.CompletionTokens += updatedSession.CompletionTokens
88 parentSession.CompletionTokens += updatedSession.CompletionTokens
89
89
90− _, err = b.app.Sessions.Save(parentSession)
90+ _, err = b.app.Sessions.Save(ctx, parentSession)
91 if err != nil {
91 if err != nil {
92 return tools.NewTextErrorResponse(fmt.Sprintf("error: %s", err)), nil
92 return tools.NewTextErrorResponse(fmt.Sprintf("error: %s", err)), nil
93 }
93 }
internal/llm/agent/agent.go+27−26
@@ -48,7 +48,7 @@ func (c *agent) handleTitleGeneration(ctx context.Context, sessionID, content st
48 return
48 return
49 }
49 }
50
50
51− session, err := c.Sessions.Get(sessionID)
51+ session, err := c.Sessions.Get(ctx, sessionID)
52 if err != nil {
52 if err != nil {
53 return
53 return
54 }
54 }
@@ -56,12 +56,12 @@ func (c *agent) handleTitleGeneration(ctx context.Context, sessionID, content st
56 session.Title = response.Content
56 session.Title = response.Content
57 session.Title = strings.TrimSpace(session.Title)
57 session.Title = strings.TrimSpace(session.Title)
58 session.Title = strings.ReplaceAll(session.Title, "\n", " ")
58 session.Title = strings.ReplaceAll(session.Title, "\n", " ")
59− c.Sessions.Save(session)
59+ c.Sessions.Save(ctx, session)
60 }
60 }
61}
61}
62
62
63−func (c *agent) TrackUsage(sessionID string, model models.Model, usage provider.TokenUsage) error {
63+func (c *agent) TrackUsage(ctx context.Context, sessionID string, model models.Model, usage provider.TokenUsage) error {
64− session, err := c.Sessions.Get(sessionID)
64+ session, err := c.Sessions.Get(ctx, sessionID)
65 if err != nil {
65 if err != nil {
66 return err
66 return err
67 }
67 }
@@ -75,11 +75,12 @@ func (c *agent) TrackUsage(sessionID string, model models.Model, usage provider.
75 session.CompletionTokens += usage.OutputTokens
75 session.CompletionTokens += usage.OutputTokens
76 session.PromptTokens += usage.InputTokens
76 session.PromptTokens += usage.InputTokens
77
77
78− _, err = c.Sessions.Save(session)
78+ _, err = c.Sessions.Save(ctx, session)
79 return err
79 return err
80}
80}
81
81
82func (c *agent) processEvent(
82func (c *agent) processEvent(
83+ ctx context.Context,
83 sessionID string,
84 sessionID string,
84 assistantMsg *message.Message,
85 assistantMsg *message.Message,
85 event provider.ProviderEvent,
86 event provider.ProviderEvent,
@@ -87,10 +88,10 @@ func (c *agent) processEvent(
87 switch event.Type {
88 switch event.Type {
88 case provider.EventThinkingDelta:
89 case provider.EventThinkingDelta:
89 assistantMsg.AppendReasoningContent(event.Content)
90 assistantMsg.AppendReasoningContent(event.Content)
90− return c.Messages.Update(*assistantMsg)
91+ return c.Messages.Update(ctx, *assistantMsg)
91 case provider.EventContentDelta:
92 case provider.EventContentDelta:
92 assistantMsg.AppendContent(event.Content)
93 assistantMsg.AppendContent(event.Content)
93− return c.Messages.Update(*assistantMsg)
94+ return c.Messages.Update(ctx, *assistantMsg)
94 case provider.EventError:
95 case provider.EventError:
95 if errors.Is(event.Error, context.Canceled) {
96 if errors.Is(event.Error, context.Canceled) {
96 return nil
97 return nil
@@ -105,11 +106,11 @@ func (c *agent) processEvent(
105 case provider.EventComplete:
106 case provider.EventComplete:
106 assistantMsg.SetToolCalls(event.Response.ToolCalls)
107 assistantMsg.SetToolCalls(event.Response.ToolCalls)
107 assistantMsg.AddFinish(event.Response.FinishReason)
108 assistantMsg.AddFinish(event.Response.FinishReason)
108− err := c.Messages.Update(*assistantMsg)
109+ err := c.Messages.Update(ctx, *assistantMsg)
109 if err != nil {
110 if err != nil {
110 return err
111 return err
111 }
112 }
112− return c.TrackUsage(sessionID, c.model, event.Response.Usage)
113+ return c.TrackUsage(ctx, sessionID, c.model, event.Response.Usage)
113 }
114 }
114
115
115 return nil
116 return nil
@@ -237,7 +238,7 @@ func (c *agent) handleToolExecution(
237 for _, toolResult := range toolResults {
238 for _, toolResult := range toolResults {
238 parts = append(parts, toolResult)
239 parts = append(parts, toolResult)
239 }
240 }
240− msg, err := c.Messages.Create(assistantMsg.SessionID, message.CreateMessageParams{
241+ msg, err := c.Messages.Create(ctx, assistantMsg.SessionID, message.CreateMessageParams{
241 Role: message.Tool,
242 Role: message.Tool,
242 Parts: parts,
243 Parts: parts,
243 })
244 })
@@ -247,7 +248,7 @@ func (c *agent) handleToolExecution(
247
248
248func (c *agent) generate(ctx context.Context, sessionID string, content string) error {
249func (c *agent) generate(ctx context.Context, sessionID string, content string) error {
249 ctx = context.WithValue(ctx, tools.SessionIDContextKey, sessionID)
250 ctx = context.WithValue(ctx, tools.SessionIDContextKey, sessionID)
250− messages, err := c.Messages.List(sessionID)
251+ messages, err := c.Messages.List(ctx, sessionID)
251 if err != nil {
252 if err != nil {
252 return err
253 return err
253 }
254 }
@@ -256,7 +257,7 @@ func (c *agent) generate(ctx context.Context, sessionID string, content string)
256 go c.handleTitleGeneration(ctx, sessionID, content)
257 go c.handleTitleGeneration(ctx, sessionID, content)
257 }
258 }
258
259
259− userMsg, err := c.Messages.Create(sessionID, message.CreateMessageParams{
260+ userMsg, err := c.Messages.Create(ctx, sessionID, message.CreateMessageParams{
260 Role: message.User,
261 Role: message.User,
261 Parts: []message.ContentPart{
262 Parts: []message.ContentPart{
262 message.TextContent{
263 message.TextContent{
@@ -272,7 +273,7 @@ func (c *agent) generate(ctx context.Context, sessionID string, content string)
272 for {
273 for {
273 select {
274 select {
274 case <-ctx.Done():
275 case <-ctx.Done():
275− assistantMsg, err := c.Messages.Create(sessionID, message.CreateMessageParams{
276+ assistantMsg, err := c.Messages.Create(ctx, sessionID, message.CreateMessageParams{
276 Role: message.Assistant,
277 Role: message.Assistant,
277 Parts: []message.ContentPart{},
278 Parts: []message.ContentPart{},
278 })
279 })
@@ -280,7 +281,7 @@ func (c *agent) generate(ctx context.Context, sessionID string, content string)
280 return err
281 return err
281 }
282 }
282 assistantMsg.AddFinish("canceled")
283 assistantMsg.AddFinish("canceled")
283− c.Messages.Update(assistantMsg)
284+ c.Messages.Update(ctx, assistantMsg)
284 return context.Canceled
285 return context.Canceled
285 default:
286 default:
286 // Continue processing
287 // Continue processing
@@ -289,7 +290,7 @@ func (c *agent) generate(ctx context.Context, sessionID string, content string)
289 eventChan, err := c.agent.StreamResponse(ctx, messages, c.tools)
290 eventChan, err := c.agent.StreamResponse(ctx, messages, c.tools)
290 if err != nil {
291 if err != nil {
291 if errors.Is(err, context.Canceled) {
292 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{
293 Role: message.Assistant,
294 Role: message.Assistant,
294 Parts: []message.ContentPart{},
295 Parts: []message.ContentPart{},
295 })
296 })
@@ -297,13 +298,13 @@ func (c *agent) generate(ctx context.Context, sessionID string, content string)
297 return err
298 return err
298 }
299 }
299 assistantMsg.AddFinish("canceled")
300 assistantMsg.AddFinish("canceled")
300− c.Messages.Update(assistantMsg)
301+ c.Messages.Update(ctx, assistantMsg)
301 return context.Canceled
302 return context.Canceled
302 }
303 }
303 return err
304 return err
304 }
305 }
305
306
306− assistantMsg, err := c.Messages.Create(sessionID, message.CreateMessageParams{
307+ assistantMsg, err := c.Messages.Create(ctx, sessionID, message.CreateMessageParams{
307 Role: message.Assistant,
308 Role: message.Assistant,
308 Parts: []message.ContentPart{},
309 Parts: []message.ContentPart{},
309 Model: c.model.ID,
310 Model: c.model.ID,
@@ -314,22 +315,22 @@ func (c *agent) generate(ctx context.Context, sessionID string, content string)
314
315
315 ctx = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)
316 ctx = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)
316 for event := range eventChan {
317 for event := range eventChan {
317− err = c.processEvent(sessionID, &assistantMsg, event)
318+ err = c.processEvent(ctx, sessionID, &assistantMsg, event)
318 if err != nil {
319 if err != nil {
319 if errors.Is(err, context.Canceled) {
320 if errors.Is(err, context.Canceled) {
320 assistantMsg.AddFinish("canceled")
321 assistantMsg.AddFinish("canceled")
321− c.Messages.Update(assistantMsg)
322+ c.Messages.Update(ctx, assistantMsg)
322 return context.Canceled
323 return context.Canceled
323 }
324 }
324 assistantMsg.AddFinish("error:" + err.Error())
325 assistantMsg.AddFinish("error:" + err.Error())
325− c.Messages.Update(assistantMsg)
326+ c.Messages.Update(ctx, assistantMsg)
326 return err
327 return err
327 }
328 }
328
329
329 select {
330 select {
330 case <-ctx.Done():
331 case <-ctx.Done():
331 assistantMsg.AddFinish("canceled")
332 assistantMsg.AddFinish("canceled")
332− c.Messages.Update(assistantMsg)
333+ c.Messages.Update(ctx, assistantMsg)
333 return context.Canceled
334 return context.Canceled
334 default:
335 default:
335 }
336 }
@@ -339,7 +340,7 @@ func (c *agent) generate(ctx context.Context, sessionID string, content string)
339 select {
340 select {
340 case <-ctx.Done():
341 case <-ctx.Done():
341 assistantMsg.AddFinish("canceled")
342 assistantMsg.AddFinish("canceled")
342− c.Messages.Update(assistantMsg)
343+ c.Messages.Update(ctx, assistantMsg)
343 return context.Canceled
344 return context.Canceled
344 default:
345 default:
345 // Continue processing
346 // Continue processing
@@ -349,13 +350,13 @@ func (c *agent) generate(ctx context.Context, sessionID string, content string)
349 if err != nil {
350 if err != nil {
350 if errors.Is(err, context.Canceled) {
351 if errors.Is(err, context.Canceled) {
351 assistantMsg.AddFinish("canceled")
352 assistantMsg.AddFinish("canceled")
352− c.Messages.Update(assistantMsg)
353+ c.Messages.Update(ctx, assistantMsg)
353 return context.Canceled
354 return context.Canceled
354 }
355 }
355 return err
356 return err
356 }
357 }
357
358
358− c.Messages.Update(assistantMsg)
359+ c.Messages.Update(ctx, assistantMsg)
359
360
360 if len(assistantMsg.ToolCalls()) == 0 {
361 if len(assistantMsg.ToolCalls()) == 0 {
361 break
362 break
@@ -370,7 +371,7 @@ func (c *agent) generate(ctx context.Context, sessionID string, content string)
370 select {
371 select {
371 case <-ctx.Done():
372 case <-ctx.Done():
372 assistantMsg.AddFinish("canceled")
373 assistantMsg.AddFinish("canceled")
373− c.Messages.Update(assistantMsg)
374+ c.Messages.Update(ctx, assistantMsg)
374 return context.Canceled
375 return context.Canceled
375 default:
376 default:
376 // Continue processing
377 // Continue processing
@@ -383,7 +384,7 @@ func getAgentProviders(ctx context.Context, model models.Model) (provider.Provid
383 maxTokens := config.Get().Model.CoderMaxTokens
384 maxTokens := config.Get().Model.CoderMaxTokens
384
385
385 providerConfig, ok := config.Get().Providers[model.Provider]
386 providerConfig, ok := config.Get().Providers[model.Provider]
386− if !ok || !providerConfig.Enabled {
387+ if !ok || providerConfig.Disabled {
387 return nil, nil, errors.New("provider is not enabled")
388 return nil, nil, errors.New("provider is not enabled")
388 }
389 }
389 var agentProvider provider.Provider
390 var agentProvider provider.Provider
internal/llm/agent/coder.go+3−2
@@ -40,12 +40,13 @@ func NewCoderAgent(app *app.App) (Agent, error) {
40 return nil, errors.New("model not supported")
40 return nil, errors.New("model not supported")
41 }
41 }
42
42
43+ ctx := context.Background()
43− agentProvider, titleGenerator, err := getAgentProviders(app.Context, model)
44+ agentProvider, titleGenerator, err := getAgentProviders(ctx, model)
44 if err != nil {
45 if err != nil {
45 return nil, err
46 return nil, err
46 }
47 }
47
48
48− otherTools := GetMcpTools(app.Context, app.Permissions)
49+ otherTools := GetMcpTools(ctx, app.Permissions)
49 if len(app.LSPClients) > 0 {
50 if len(app.LSPClients) > 0 {
50 otherTools = append(otherTools, tools.NewDiagnosticsTool(app.LSPClients))
51 otherTools = append(otherTools, tools.NewDiagnosticsTool(app.LSPClients))
51 }
52 }
internal/llm/agent/task.go+2−1
@@ -24,7 +24,8 @@ func NewTaskAgent(app *app.App) (Agent, error) {
24 return nil, errors.New("model not supported")
24 return nil, errors.New("model not supported")
25 }
25 }
26
26
27+ ctx := context.Background()
27− agentProvider, titleGenerator, err := getAgentProviders(app.Context, model)
28+ agentProvider, titleGenerator, err := getAgentProviders(ctx, model)
28 if err != nil {
29 if err != nil {
29 return nil, err
30 return nil, err
30 }
31 }
internal/message/message.go+22−24
@@ -20,34 +20,32 @@ type CreateMessageParams struct {
20
20
21type Service interface {
21type Service interface {
22 pubsub.Suscriber[Message]
22 pubsub.Suscriber[Message]
23− Create(sessionID string, params CreateMessageParams) (Message, error)
23+ Create(ctx context.Context, sessionID string, params CreateMessageParams) (Message, error)
24− Update(message Message) error
24+ Update(ctx context.Context, message Message) error
25− Get(id string) (Message, error)
25+ Get(ctx context.Context, id string) (Message, error)
26− List(sessionID string) ([]Message, error)
26+ List(ctx context.Context, sessionID string) ([]Message, error)
27− Delete(id string) error
27+ Delete(ctx context.Context, id string) error
28− DeleteSessionMessages(sessionID string) error
28+ DeleteSessionMessages(ctx context.Context, sessionID string) error
29}
29}
30
30
31type service struct {
31type service struct {
32 *pubsub.Broker[Message]
32 *pubsub.Broker[Message]
33− q db.Querier
33+ q db.Querier
34− ctx context.Context
35}
34}
36
35
37−func NewService(ctx context.Context, q db.Querier) Service {
36+func NewService(q db.Querier) Service {
38 return &service{
37 return &service{
39 Broker: pubsub.NewBroker[Message](),
38 Broker: pubsub.NewBroker[Message](),
40 q: q,
39 q: q,
41− ctx: ctx,
42 }
40 }
43}
41}
44
42
45−func (s *service) Delete(id string) error {
43+func (s *service) Delete(ctx context.Context, id string) error {
46− message, err := s.Get(id)
44+ message, err := s.Get(ctx, id)
47 if err != nil {
45 if err != nil {
48 return err
46 return err
49 }
47 }
50− err = s.q.DeleteMessage(s.ctx, message.ID)
48+ err = s.q.DeleteMessage(ctx, message.ID)
51 if err != nil {
49 if err != nil {
52 return err
50 return err
53 }
51 }
@@ -55,7 +53,7 @@ func (s *service) Delete(id string) error {
55 return nil
53 return nil
56}
54}
57
55
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) {
59 if params.Role != Assistant {
57 if params.Role != Assistant {
60 params.Parts = append(params.Parts, Finish{
58 params.Parts = append(params.Parts, Finish{
61 Reason: "stop",
59 Reason: "stop",
@@ -66,7 +64,7 @@ func (s *service) Create(sessionID string, params CreateMessageParams) (Message,
66 return Message{}, err
64 return Message{}, err
67 }
65 }
68
66
69− dbMessage, err := s.q.CreateMessage(s.ctx, db.CreateMessageParams{
67+ dbMessage, err := s.q.CreateMessage(ctx, db.CreateMessageParams{
70 ID: uuid.New().String(),
68 ID: uuid.New().String(),
71 SessionID: sessionID,
69 SessionID: sessionID,
72 Role: string(params.Role),
70 Role: string(params.Role),
@@ -84,14 +82,14 @@ func (s *service) Create(sessionID string, params CreateMessageParams) (Message,
84 return message, nil
82 return message, nil
85}
83}
86
84
87−func (s *service) DeleteSessionMessages(sessionID string) error {
85+func (s *service) DeleteSessionMessages(ctx context.Context, sessionID string) error {
88− messages, err := s.List(sessionID)
86+ messages, err := s.List(ctx, sessionID)
89 if err != nil {
87 if err != nil {
90 return err
88 return err
91 }
89 }
92 for _, message := range messages {
90 for _, message := range messages {
93 if message.SessionID == sessionID {
91 if message.SessionID == sessionID {
94− err = s.Delete(message.ID)
92+ err = s.Delete(ctx, message.ID)
95 if err != nil {
93 if err != nil {
96 return err
94 return err
97 }
95 }
@@ -100,7 +98,7 @@ func (s *service) DeleteSessionMessages(sessionID string) error {
100 return nil
98 return nil
101}
99}
102
100
103−func (s *service) Update(message Message) error {
101+func (s *service) Update(ctx context.Context, message Message) error {
104 parts, err := marshallParts(message.Parts)
102 parts, err := marshallParts(message.Parts)
105 if err != nil {
103 if err != nil {
106 return err
104 return err
@@ -110,7 +108,7 @@ func (s *service) Update(message Message) error {
110 finishedAt.Int64 = f.Time
108 finishedAt.Int64 = f.Time
111 finishedAt.Valid = true
109 finishedAt.Valid = true
112 }
110 }
113− err = s.q.UpdateMessage(s.ctx, db.UpdateMessageParams{
111+ err = s.q.UpdateMessage(ctx, db.UpdateMessageParams{
114 ID: message.ID,
112 ID: message.ID,
115 Parts: string(parts),
113 Parts: string(parts),
116 FinishedAt: finishedAt,
114 FinishedAt: finishedAt,
@@ -122,16 +120,16 @@ func (s *service) Update(message Message) error {
122 return nil
120 return nil
123}
121}
124
122
125−func (s *service) Get(id string) (Message, error) {
123+func (s *service) Get(ctx context.Context, id string) (Message, error) {
126− dbMessage, err := s.q.GetMessage(s.ctx, id)
124+ dbMessage, err := s.q.GetMessage(ctx, id)
127 if err != nil {
125 if err != nil {
128 return Message{}, err
126 return Message{}, err
129 }
127 }
130 return s.fromDBItem(dbMessage)
128 return s.fromDBItem(dbMessage)
131}
129}
132
130
133−func (s *service) List(sessionID string) ([]Message, error) {
131+func (s *service) List(ctx context.Context, sessionID string) ([]Message, error) {
134− dbMessages, err := s.q.ListMessagesBySession(s.ctx, sessionID)
132+ dbMessages, err := s.q.ListMessagesBySession(ctx, sessionID)
135 if err != nil {
133 if err != nil {
136 return nil, err
134 return nil, err
137 }
135 }
internal/session/session.go+21−23
@@ -23,22 +23,21 @@ type Session struct {
23
23
24type Service interface {
24type Service interface {
25 pubsub.Suscriber[Session]
25 pubsub.Suscriber[Session]
26− Create(title string) (Session, error)
26+ Create(ctx context.Context, title string) (Session, error)
27− CreateTaskSession(toolCallID, parentSessionID, title string) (Session, error)
27+ CreateTaskSession(ctx context.Context, toolCallID, parentSessionID, title string) (Session, error)
28− Get(id string) (Session, error)
28+ Get(ctx context.Context, id string) (Session, error)
29− List() ([]Session, error)
29+ List(ctx context.Context) ([]Session, error)
30− Save(session Session) (Session, error)
30+ Save(ctx context.Context, session Session) (Session, error)
31− Delete(id string) error
31+ Delete(ctx context.Context, id string) error
32}
32}
33
33
34type service struct {
34type service struct {
35 *pubsub.Broker[Session]
35 *pubsub.Broker[Session]
36− q db.Querier
36+ q db.Querier
37− ctx context.Context
38}
37}
39
38
40−func (s *service) Create(title string) (Session, error) {
39+func (s *service) Create(ctx context.Context, title string) (Session, error) {
41− dbSession, err := s.q.CreateSession(s.ctx, db.CreateSessionParams{
40+ dbSession, err := s.q.CreateSession(ctx, db.CreateSessionParams{
42 ID: uuid.New().String(),
41 ID: uuid.New().String(),
43 Title: title,
42 Title: title,
44 })
43 })
@@ -50,8 +49,8 @@ func (s *service) Create(title string) (Session, error) {
50 return session, nil
49 return session, nil
51}
50}
52
51
53−func (s *service) CreateTaskSession(toolCallID, parentSessionID, title string) (Session, error) {
52+func (s *service) CreateTaskSession(ctx context.Context, toolCallID, parentSessionID, title string) (Session, error) {
54− dbSession, err := s.q.CreateSession(s.ctx, db.CreateSessionParams{
53+ dbSession, err := s.q.CreateSession(ctx, db.CreateSessionParams{
55 ID: toolCallID,
54 ID: toolCallID,
56 ParentSessionID: sql.NullString{String: parentSessionID, Valid: true},
55 ParentSessionID: sql.NullString{String: parentSessionID, Valid: true},
57 Title: title,
56 Title: title,
@@ -64,12 +63,12 @@ func (s *service) CreateTaskSession(toolCallID, parentSessionID, title string) (
64 return session, nil
63 return session, nil
65}
64}
66
65
67−func (s *service) Delete(id string) error {
66+func (s *service) Delete(ctx context.Context, id string) error {
68− session, err := s.Get(id)
67+ session, err := s.Get(ctx, id)
69 if err != nil {
68 if err != nil {
70 return err
69 return err
71 }
70 }
72− err = s.q.DeleteSession(s.ctx, session.ID)
71+ err = s.q.DeleteSession(ctx, session.ID)
73 if err != nil {
72 if err != nil {
74 return err
73 return err
75 }
74 }
@@ -77,16 +76,16 @@ func (s *service) Delete(id string) error {
77 return nil
76 return nil
78}
77}
79
78
80−func (s *service) Get(id string) (Session, error) {
79+func (s *service) Get(ctx context.Context, id string) (Session, error) {
81− dbSession, err := s.q.GetSessionByID(s.ctx, id)
80+ dbSession, err := s.q.GetSessionByID(ctx, id)
82 if err != nil {
81 if err != nil {
83 return Session{}, err
82 return Session{}, err
84 }
83 }
85 return s.fromDBItem(dbSession), nil
84 return s.fromDBItem(dbSession), nil
86}
85}
87
86
88−func (s *service) Save(session Session) (Session, error) {
87+func (s *service) Save(ctx context.Context, session Session) (Session, error) {
89− dbSession, err := s.q.UpdateSession(s.ctx, db.UpdateSessionParams{
88+ dbSession, err := s.q.UpdateSession(ctx, db.UpdateSessionParams{
90 ID: session.ID,
89 ID: session.ID,
91 Title: session.Title,
90 Title: session.Title,
92 PromptTokens: session.PromptTokens,
91 PromptTokens: session.PromptTokens,
@@ -101,8 +100,8 @@ func (s *service) Save(session Session) (Session, error) {
101 return session, nil
100 return session, nil
102}
101}
103
102
104−func (s *service) List() ([]Session, error) {
103+func (s *service) List(ctx context.Context) ([]Session, error) {
105− dbSessions, err := s.q.ListSessions(s.ctx)
104+ dbSessions, err := s.q.ListSessions(ctx)
106 if err != nil {
105 if err != nil {
107 return nil, err
106 return nil, err
108 }
107 }
@@ -127,11 +126,10 @@ func (s service) fromDBItem(item db.Session) Session {
127 }
126 }
128}
127}
129
128
130−func NewService(ctx context.Context, q db.Querier) Service {
129+func NewService(q db.Querier) Service {
131 broker := pubsub.NewBroker[Session]()
130 broker := pubsub.NewBroker[Session]()
132 return &service{
131 return &service{
133 broker,
132 broker,
134 q,
133 q,
135− ctx,
136 }
134 }
137}
135}
internal/tui/components/chat/messages.go+3−2
@@ -1,6 +1,7 @@
1package chat
1package chat
2
2
3import (
3import (
4+ "context"
4 "encoding/json"
5 "encoding/json"
5 "fmt"
6 "fmt"
6 "math"
7 "math"
@@ -324,7 +325,7 @@ func (m *messagesCmp) renderToolCall(toolCall message.ToolCall, isNested bool) s
324
325
325 innerToolCalls := make([]string, 0)
326 innerToolCalls := make([]string, 0)
326 if toolCall.Name == agent.AgentToolName {
327 if toolCall.Name == agent.AgentToolName {
327− messages, _ := m.app.Messages.List(toolCall.ID)
328+ messages, _ := m.app.Messages.List(context.Background(), toolCall.ID)
328 toolCalls := make([]message.ToolCall, 0)
329 toolCalls := make([]message.ToolCall, 0)
329 for _, v := range messages {
330 for _, v := range messages {
330 toolCalls = append(toolCalls, v.ToolCalls()...)
331 toolCalls = append(toolCalls, v.ToolCalls()...)
@@ -554,7 +555,7 @@ func (m *messagesCmp) GetSize() (int, int) {
554
555
555func (m *messagesCmp) SetSession(session session.Session) tea.Cmd {
556func (m *messagesCmp) SetSession(session session.Session) tea.Cmd {
556 m.session = session
557 m.session = session
557− messages, err := m.app.Messages.List(session.ID)
558+ messages, err := m.app.Messages.List(context.Background(), session.ID)
558 if err != nil {
559 if err != nil {
559 return util.ReportError(err)
560 return util.ReportError(err)
560 }
561 }
internal/tui/components/repl/editor.go+2−2
@@ -160,7 +160,7 @@ func (m *editorCmp) Send() tea.Cmd {
160 return util.ReportWarn("Assistant is still working on the previous message")
160 return util.ReportWarn("Assistant is still working on the previous message")
161 }
161 }
162
162
163− messages, err := m.app.Messages.List(m.sessionID)
163+ messages, err := m.app.Messages.List(context.Background(), m.sessionID)
164 if err != nil {
164 if err != nil {
165 return util.ReportError(err)
165 return util.ReportError(err)
166 }
166 }
@@ -177,7 +177,7 @@ func (m *editorCmp) Send() tea.Cmd {
177 if len(content) == 0 {
177 if len(content) == 0 {
178 return util.ReportWarn("Message is empty")
178 return util.ReportWarn("Message is empty")
179 }
179 }
180− ctx, cancel := context.WithCancel(m.app.Context)
180+ ctx, cancel := context.WithCancel(context.Background())
181 m.cancelMessage = cancel
181 m.cancelMessage = cancel
182 go func() {
182 go func() {
183 defer cancel()
183 defer cancel()
internal/tui/components/repl/messages.go+4−3
@@ -1,6 +1,7 @@
1package repl
1package repl
2
2
3import (
3import (
4+ "context"
4 "encoding/json"
5 "encoding/json"
5 "fmt"
6 "fmt"
6 "sort"
7 "sort"
@@ -77,8 +78,8 @@ func (m *messagesCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
77 m.session = msg.Payload
78 m.session = msg.Payload
78 }
79 }
79 case SelectedSessionMsg:
80 case SelectedSessionMsg:
80− m.session, _ = m.app.Sessions.Get(msg.SessionID)
81+ m.session, _ = m.app.Sessions.Get(context.Background(), msg.SessionID)
81− m.messages, _ = m.app.Messages.List(m.session.ID)
82+ m.messages, _ = m.app.Messages.List(context.Background(), m.session.ID)
82 m.renderView()
83 m.renderView()
83 m.viewport.GotoBottom()
84 m.viewport.GotoBottom()
84 }
85 }
@@ -259,7 +260,7 @@ func (m *messagesCmp) renderMessageWithToolCall(content string, tools []message.
259
260
260 runningIndicator := runningStyle.Render(fmt.Sprintf("%s Running...", styles.SpinnerIcon))
261 runningIndicator := runningStyle.Render(fmt.Sprintf("%s Running...", styles.SpinnerIcon))
261 allParts = append(allParts, leftPadding.Render(runningIndicator))
262 allParts = append(allParts, leftPadding.Render(runningIndicator))
262− taskSessionMessages, _ := m.app.Messages.List(toolCall.ID)
263+ taskSessionMessages, _ := m.app.Messages.List(context.Background(), toolCall.ID)
263 for _, msg := range taskSessionMessages {
264 for _, msg := range taskSessionMessages {
264 if msg.Role == message.Assistant {
265 if msg.Role == message.Assistant {
265 for _, toolCall := range msg.ToolCalls() {
266 for _, toolCall := range msg.ToolCalls() {
internal/tui/components/repl/sessions.go+3−1
@@ -1,6 +1,7 @@
1package repl
1package repl
2
2
3import (
3import (
4+ "context"
4 "fmt"
5 "fmt"
5 "strings"
6 "strings"
6
7
@@ -57,12 +58,13 @@ var sessionKeyMapValue = sessionsKeyMap{
57}
58}
58
59
59func (i *sessionsCmp) Init() tea.Cmd {
60func (i *sessionsCmp) Init() tea.Cmd {
60− existing, err := i.app.Sessions.List()
61+ existing, err := i.app.Sessions.List(context.Background())
61 if err != nil {
62 if err != nil {
62 return util.ReportError(err)
63 return util.ReportError(err)
63 }
64 }
64 if len(existing) == 0 || existing[0].MessageCount > 0 {
65 if len(existing) == 0 || existing[0].MessageCount > 0 {
65 newSession, err := i.app.Sessions.Create(
66 newSession, err := i.app.Sessions.Create(
67+ context.Background(),
66 "New Session",
68 "New Session",
67 )
69 )
68 if err != nil {
70 if err != nil {
internal/tui/page/chat.go+5−3
@@ -1,6 +1,8 @@
1package page
1package page
2
2
3import (
3import (
4+ "context"
5+
4 "github.com/charmbracelet/bubbles/key"
6 "github.com/charmbracelet/bubbles/key"
5 tea "github.com/charmbracelet/bubbletea"
7 tea "github.com/charmbracelet/bubbletea"
6 "github.com/kujtimiihoxha/termai/internal/app"
8 "github.com/kujtimiihoxha/termai/internal/app"
@@ -36,7 +38,7 @@ func (p *chatPage) Init() tea.Cmd {
36 p.layout.Init(),
38 p.layout.Init(),
37 }
39 }
38
40
39− sessions, _ := p.app.Sessions.List()
41+ sessions, _ := p.app.Sessions.List(context.Background())
40 if len(sessions) > 0 {
42 if len(sessions) > 0 {
41 p.session = sessions[0]
43 p.session = sessions[0]
42 cmd := p.setSidebar()
44 cmd := p.setSidebar()
@@ -92,7 +94,7 @@ func (p *chatPage) clearSidebar() {
92func (p *chatPage) sendMessage(text string) tea.Cmd {
94func (p *chatPage) sendMessage(text string) tea.Cmd {
93 var cmds []tea.Cmd
95 var cmds []tea.Cmd
94 if p.session.ID == "" {
96 if p.session.ID == "" {
95− session, err := p.app.Sessions.Create("New Session")
97+ session, err := p.app.Sessions.Create(context.Background(), "New Session")
96 if err != nil {
98 if err != nil {
97 return util.ReportError(err)
99 return util.ReportError(err)
98 }
100 }
@@ -110,7 +112,7 @@ func (p *chatPage) sendMessage(text string) tea.Cmd {
110 return util.ReportError(err)
112 return util.ReportError(err)
111 }
113 }
112 go func() {
114 go func() {
113− a.Generate(p.app.Context, p.session.ID, text)
115+ a.Generate(context.Background(), p.session.ID, text)
114 }()
116 }()
115
117
116 return tea.Batch(cmds...)
118 return tea.Batch(cmds...)
internal/tui/tui.go+4−2
@@ -1,6 +1,8 @@
1package tui
1package tui
2
2
3import (
3import (
4+ "context"
5+
4 "github.com/charmbracelet/bubbles/key"
6 "github.com/charmbracelet/bubbles/key"
5 tea "github.com/charmbracelet/bubbletea"
7 tea "github.com/charmbracelet/bubbletea"
6 "github.com/charmbracelet/lipgloss"
8 "github.com/charmbracelet/lipgloss"
@@ -184,7 +186,7 @@ func (a appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
184 }
186 }
185 case key.Matches(msg, replKeyMap):
187 case key.Matches(msg, replKeyMap):
186 if a.currentPage == page.ReplPage {
188 if a.currentPage == page.ReplPage {
187− sessions, err := a.app.Sessions.List()
189+ sessions, err := a.app.Sessions.List(context.Background())
188 if err != nil {
190 if err != nil {
189 return a, util.CmdHandler(util.ReportError(err))
191 return a, util.CmdHandler(util.ReportError(err))
190 }
192 }
@@ -192,7 +194,7 @@ func (a appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
192 if lastSession.MessageCount == 0 {
194 if lastSession.MessageCount == 0 {
193 return a, util.CmdHandler(repl.SelectedSessionMsg{SessionID: lastSession.ID})
195 return a, util.CmdHandler(repl.SelectedSessionMsg{SessionID: lastSession.ID})
194 }
196 }
195− s, err := a.app.Sessions.Create("New Session")
197+ s, err := a.app.Sessions.Create(context.Background(), "New Session")
196 if err != nil {
198 if err != nil {
197 return a, util.CmdHandler(util.ReportError(err))
199 return a, util.CmdHandler(util.ReportError(err))
198 }
200 }