From c85745b230089cedaacec636e6fc78f8ab922608 Mon Sep 17 00:00:00 2001
From: Kaden <kaidizhao7@gmail.com>
Date: Tue, 16 Jun 2026 23:26:52 +0800
Subject: [PATCH] fix: serve Ctrl-C stop

---
 cli/serve.go     | 7 ++++++-
 cmd/kage/main.go | 9 ++++++++-
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/cli/serve.go b/cli/serve.go
index 7ed3166..78fcfdd 100644
--- a/cli/serve.go
+++ b/cli/serve.go
@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"os"
 	"path/filepath"
+	"time"
 
 	"github.com/spf13/cobra"
 )
@@ -61,7 +62,11 @@ func runServe(ctx context.Context, dir, addr string) error {
 			return err
 		}
 	case <-ctx.Done():
-		_ = srv.Close()
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+		defer cancel()
+		if err := srv.Shutdown(shutdownCtx); err != nil {
+			_ = srv.Close()
+		}
 		if err := <-srvErr; err != nil && err != http.ErrServerClosed {
 			return err
 		}
diff --git a/cmd/kage/main.go b/cmd/kage/main.go
index 31fdb25..51f0cd4 100644
--- a/cmd/kage/main.go
+++ b/cmd/kage/main.go
@@ -7,6 +7,7 @@ import (
 	"context"
 	"os"
 	"os/signal"
+	"syscall"
 
 	"github.com/tamnd/kage/cli"
 	"github.com/tamnd/kage/viewer"
@@ -18,7 +19,13 @@ func main() {
 	// thread; in the default build this is a harmless no-op.
 	viewer.LockMainThread()
 
-	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
 	defer stop()
+
+	go func() {
+		<-ctx.Done()
+		stop()
+	}()
+
 	os.Exit(cli.Execute(ctx))
 }