From fa311c46acc3f728d249ed6ee6bb1b49f84978a2 Mon Sep 17 00:00:00 2001
From: LordMathis <matus@namesny.com>
Date: Sat, 6 Dec 2025 19:52:40 +0100
Subject: [PATCH] Improve server shutdown process

---
 .vscode/launch.json      |  1 +
 cmd/server/main.go       | 15 +++++++++++++--
 pkg/database/database.go |  6 ++++++
 pkg/manager/manager.go   |  5 -----
 4 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/.vscode/launch.json b/.vscode/launch.json
index 12a8525..5422f27 100644
--- a/.vscode/launch.json
+++ b/.vscode/launch.json
@@ -14,6 +14,7 @@
                 "GO_ENV": "development",
                 "LLAMACTL_CONFIG_PATH": "${workspaceFolder}/llamactl.dev.yaml"
             },
+            "console": "integratedTerminal",
         }
     ]
 }
\ No newline at end of file
diff --git a/cmd/server/main.go b/cmd/server/main.go
index 9431b59..579b924 100644
--- a/cmd/server/main.go
+++ b/cmd/server/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"llamactl/pkg/config"
 	"llamactl/pkg/database"
@@ -11,6 +12,7 @@ import (
 	"os"
 	"os/signal"
 	"syscall"
+	"time"
 )
 
 // version is set at build time using -ldflags "-X main.version=1.0.0"
@@ -116,14 +118,23 @@ func main() {
 	<-stop
 	fmt.Println("Shutting down server...")
 
-	if err := server.Close(); err != nil {
+	// Create shutdown context with timeout
+	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer shutdownCancel()
+
+	// Shutdown HTTP server gracefully
+	if err := server.Shutdown(shutdownCtx); err != nil {
 		log.Printf("Error shutting down server: %v\n", err)
 	} else {
 		fmt.Println("Server shut down gracefully.")
 	}
 
-	// Wait for all instances to stop
+	// Stop all instances and cleanup
 	instanceManager.Shutdown()
 
+	if err := db.Close(); err != nil {
+		log.Printf("Error closing database: %v\n", err)
+	}
+
 	fmt.Println("Exiting llamactl.")
 }
diff --git a/pkg/database/database.go b/pkg/database/database.go
index 98793ff..b5ed946 100644
--- a/pkg/database/database.go
+++ b/pkg/database/database.go
@@ -107,6 +107,12 @@ func Open(config *Config) (*sqliteDB, error) {
 func (db *sqliteDB) Close() error {
 	if db.DB != nil {
 		log.Println("Closing database connection")
+
+		// Checkpoint WAL to merge changes back to main database file
+		if _, err := db.DB.Exec("PRAGMA wal_checkpoint(TRUNCATE)"); err != nil {
+			log.Printf("Warning: Failed to checkpoint WAL: %v", err)
+		}
+
 		return db.DB.Close()
 	}
 	return nil
diff --git a/pkg/manager/manager.go b/pkg/manager/manager.go
index fc23a2b..e566fe4 100644
--- a/pkg/manager/manager.go
+++ b/pkg/manager/manager.go
@@ -114,11 +114,6 @@ func (im *instanceManager) Shutdown() {
 		}
 		wg.Wait()
 		fmt.Println("All instances stopped.")
-
-		// 4. Close database connection
-		if err := im.db.Close(); err != nil {
-			log.Printf("Error closing database: %v\n", err)
-		}
 	})
 }