cleanup app, config and root

3 files changed, 184 insertions, 64 deletions

diff --git a/internal/app/app.go b/internal/app/app.go
new file mode 100644
index 000000000..fa4a6ee90
--- /dev/null
+++ b/internal/app/app.go
@@ -0,0 +1,76 @@
+package app
+
+import (
+	"context"
+	"database/sql"
+	"maps"
+	"sync"
+	"time"
+
+	"github.com/kujtimiihoxha/termai/internal/db"
+	"github.com/kujtimiihoxha/termai/internal/history"
+	"github.com/kujtimiihoxha/termai/internal/logging"
+	"github.com/kujtimiihoxha/termai/internal/lsp"
+	"github.com/kujtimiihoxha/termai/internal/message"
+	"github.com/kujtimiihoxha/termai/internal/permission"
+	"github.com/kujtimiihoxha/termai/internal/session"
+)
+
+type App struct {
+	Sessions    session.Service
+	Messages    message.Service
+	Files       history.Service
+	Permissions permission.Service
+
+	LSPClients map[string]*lsp.Client
+
+	clientsMutex sync.RWMutex
+
+	watcherCancelFuncs []context.CancelFunc
+	cancelFuncsMutex   sync.Mutex
+	watcherWG          sync.WaitGroup
+}
+
+func New(ctx context.Context, conn *sql.DB) *App {
+	q := db.New(conn)
+	sessions := session.NewService(q)
+	messages := message.NewService(q)
+	files := history.NewService(q)
+
+	app := &App{
+		Sessions:    sessions,
+		Messages:    messages,
+		Files:       files,
+		Permissions: permission.NewPermissionService(),
+		LSPClients:  make(map[string]*lsp.Client),
+	}
+
+	app.initLSPClients(ctx)
+
+	return app
+}
+
+// Shutdown performs a clean shutdown of the application
+func (app *App) Shutdown() {
+	// Cancel all watcher goroutines
+	app.cancelFuncsMutex.Lock()
+	for _, cancel := range app.watcherCancelFuncs {
+		cancel()
+	}
+	app.cancelFuncsMutex.Unlock()
+	app.watcherWG.Wait()
+
+	// Perform additional cleanup for LSP clients
+	app.clientsMutex.RLock()
+	clients := make(map[string]*lsp.Client, len(app.LSPClients))
+	maps.Copy(clients, app.LSPClients)
+	app.clientsMutex.RUnlock()
+
+	for name, client := range clients {
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		if err := client.Shutdown(shutdownCtx); err != nil {
+			logging.Error("Failed to shutdown LSP client", "name", name, "error", err)
+		}
+		cancel()
+	}
+}
diff --git a/internal/app/lsp.go b/internal/app/lsp.go
new file mode 100644
index 000000000..4e0568f07
--- /dev/null
+++ b/internal/app/lsp.go
@@ -0,0 +1,108 @@
+package app
+
+import (
+	"context"
+	"time"
+
+	"github.com/kujtimiihoxha/termai/internal/config"
+	"github.com/kujtimiihoxha/termai/internal/logging"
+	"github.com/kujtimiihoxha/termai/internal/lsp"
+	"github.com/kujtimiihoxha/termai/internal/lsp/watcher"
+)
+
+func (app *App) initLSPClients(ctx context.Context) {
+	cfg := config.Get()
+
+	// Initialize LSP clients
+	for name, clientConfig := range cfg.LSP {
+		app.createAndStartLSPClient(ctx, name, clientConfig.Command, clientConfig.Args...)
+	}
+}
+
+// createAndStartLSPClient creates a new LSP client, initializes it, and starts its workspace watcher
+func (app *App) createAndStartLSPClient(ctx context.Context, name string, command string, args ...string) {
+	// Create a specific context for initialization with a timeout
+	initCtx, initCancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer initCancel()
+
+	// Create the LSP client
+	lspClient, err := lsp.NewClient(initCtx, command, args...)
+	if err != nil {
+		logging.Error("Failed to create LSP client for", name, err)
+		return
+	}
+
+	// Initialize with the initialization context
+	_, err = lspClient.InitializeLSPClient(initCtx, config.WorkingDirectory())
+	if err != nil {
+		logging.Error("Initialize failed", "name", name, "error", err)
+		// Clean up the client to prevent resource leaks
+		lspClient.Close()
+		return
+	}
+
+	// Create a child context that can be canceled when the app is shutting down
+	watchCtx, cancelFunc := context.WithCancel(ctx)
+	workspaceWatcher := watcher.NewWorkspaceWatcher(lspClient)
+
+	// Store the cancel function to be called during cleanup
+	app.cancelFuncsMutex.Lock()
+	app.watcherCancelFuncs = append(app.watcherCancelFuncs, cancelFunc)
+	app.cancelFuncsMutex.Unlock()
+
+	// Add the watcher to a WaitGroup to track active goroutines
+	app.watcherWG.Add(1)
+
+	// Add to map with mutex protection before starting goroutine
+	app.clientsMutex.Lock()
+	app.LSPClients[name] = lspClient
+	app.clientsMutex.Unlock()
+
+	go app.runWorkspaceWatcher(watchCtx, name, workspaceWatcher)
+}
+
+// runWorkspaceWatcher executes the workspace watcher for an LSP client
+func (app *App) runWorkspaceWatcher(ctx context.Context, name string, workspaceWatcher *watcher.WorkspaceWatcher) {
+	defer app.watcherWG.Done()
+	defer func() {
+		if r := recover(); r != nil {
+			logging.Error("LSP client crashed", "client", name, "panic", r)
+
+			// Try to restart the client
+			app.restartLSPClient(ctx, name)
+		}
+	}()
+
+	workspaceWatcher.WatchWorkspace(ctx, config.WorkingDirectory())
+	logging.Info("Workspace watcher stopped", "client", name)
+}
+
+// restartLSPClient attempts to restart a crashed or failed LSP client
+func (app *App) restartLSPClient(ctx context.Context, name string) {
+	// Get the original configuration
+	cfg := config.Get()
+	clientConfig, exists := cfg.LSP[name]
+	if !exists {
+		logging.Error("Cannot restart client, configuration not found", "client", name)
+		return
+	}
+
+	// Clean up the old client if it exists
+	app.clientsMutex.Lock()
+	oldClient, exists := app.LSPClients[name]
+	if exists {
+		delete(app.LSPClients, name) // Remove from map before potentially slow shutdown
+	}
+	app.clientsMutex.Unlock()
+
+	if exists && oldClient != nil {
+		// Try to shut it down gracefully, but don't block on errors
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		_ = oldClient.Shutdown(shutdownCtx)
+		cancel()
+	}
+
+	// Create a new client using the shared function
+	app.createAndStartLSPClient(ctx, name, clientConfig.Command, clientConfig.Args...)
+	logging.Info("Successfully restarted LSP client", "client", name)
+}
diff --git a/internal/app/services.go b/internal/app/services.go
deleted file mode 100644
index 6ecdef03c..000000000
--- a/internal/app/services.go
+++ /dev/null
@@ -1,64 +0,0 @@
-package app
-
-import (
-	"context"
-	"database/sql"
-
-	"github.com/kujtimiihoxha/termai/internal/config"
-	"github.com/kujtimiihoxha/termai/internal/db"
-	"github.com/kujtimiihoxha/termai/internal/history"
-	"github.com/kujtimiihoxha/termai/internal/logging"
-	"github.com/kujtimiihoxha/termai/internal/lsp"
-	"github.com/kujtimiihoxha/termai/internal/lsp/watcher"
-	"github.com/kujtimiihoxha/termai/internal/message"
-	"github.com/kujtimiihoxha/termai/internal/permission"
-	"github.com/kujtimiihoxha/termai/internal/session"
-)
-
-type App struct {
-	Context context.Context
-
-	Sessions    session.Service
-	Messages    message.Service
-	Files       history.Service
-	Permissions permission.Service
-
-	LSPClients map[string]*lsp.Client
-}
-
-func New(ctx context.Context, conn *sql.DB) *App {
-	cfg := config.Get()
-	logging.Info("Debug mode enabled")
-
-	q := db.New(conn)
-	sessions := session.NewService(ctx, q)
-	messages := message.NewService(ctx, q)
-	files := history.NewService(ctx, q)
-
-	app := &App{
-		Context:     ctx,
-		Sessions:    sessions,
-		Messages:    messages,
-		Files:       files,
-		Permissions: permission.NewPermissionService(),
-		LSPClients:  make(map[string]*lsp.Client),
-	}
-
-	for name, client := range cfg.LSP {
-		lspClient, err := lsp.NewClient(ctx, client.Command, client.Args...)
-		workspaceWatcher := watcher.NewWorkspaceWatcher(lspClient)
-		if err != nil {
-			logging.Error("Failed to create LSP client for", name, err)
-			continue
-		}
-
-		_, err = lspClient.InitializeLSPClient(ctx, config.WorkingDirectory())
-		if err != nil {
-			logging.Error("Initialize failed", "error", err)
-			continue
-		}
-		go workspaceWatcher.WatchWorkspace(ctx, config.WorkingDirectory())
-		app.LSPClients[name] = lspClient
-	}
-	return app
-}