Diff
1diff --git a/cmd/kaiwari-server/main.go b/cmd/kaiwari-server/main.go
2index b1bdaea2e4be69982e8472b7934920f75af1890b..07805f2658301734e29475694130498787d20041 100644
3--- a/cmd/kaiwari-server/main.go
4+++ b/cmd/kaiwari-server/main.go
5@@ -1,16 +1,22 @@
6 package main
7
8 import (
9+ "context"
10 "fmt"
11 "log/slog"
12 "net/http"
13 "os"
14+ "os/signal"
15+ "syscall"
16+ "time"
17
18 "git.theedgeofrage.com/TheEdgeOfRage/kaiwari/internal/bootstrap"
19 "git.theedgeofrage.com/TheEdgeOfRage/kaiwari/internal/config"
20 "git.theedgeofrage.com/TheEdgeOfRage/kaiwari/internal/server"
21 )
22
23+const shutdownTimeout = 10 * time.Second
24+
25 func main() {
26 slog.SetDefault(slog.New(slog.NewTextHandler(os.Stderr, nil)))
27 cfg := config.ParseServerConfig()
28@@ -18,15 +24,36 @@ func main() {
29 fatalf("token is required: set --token or JP_SERVER_TOKEN")
30 }
31
32- svc, err := bootstrap.New(&cfg.Config)
33+ ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
34+ defer stop()
35+
36+ svc, err := bootstrap.New(ctx, &cfg.Config)
37 if err != nil {
38 fatalf("%v", err)
39 }
40+ defer svc.Close()
41
42 srv := server.New(svc, cfg.Token)
43+ httpSrv := &http.Server{Addr: cfg.Listen, Handler: srv.Handler()}
44+ errCh := make(chan error, 1)
45+ go func() {
46+ if err := httpSrv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
47+ errCh <- err
48+ }
49+ }()
50+
51 fmt.Printf("kaiwari-server: listening on %s\n", cfg.Listen)
52- if err := http.ListenAndServe(cfg.Listen, srv.Handler()); err != nil {
53+ select {
54+ case err := <-errCh:
55+ svc.Close()
56 fatalf("%v", err)
57+ case <-ctx.Done():
58+ shutdownCtx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
59+ defer cancel()
60+ if err := httpSrv.Shutdown(shutdownCtx); err != nil {
61+ svc.Close()
62+ fatalf("shutdown: %v", err)
63+ }
64 }
65 }
66
67diff --git a/cmd/kaiwari/main.go b/cmd/kaiwari/main.go
68index 68f1d699a43147774ce2733578f41f5eb06d27cd..9285f5e88c3445303c73110ee46726e08855383d 100644
69--- a/cmd/kaiwari/main.go
70+++ b/cmd/kaiwari/main.go
71@@ -1,8 +1,11 @@
72 package main
73
74 import (
75+ "context"
76 "fmt"
77 "os"
78+ "os/signal"
79+ "syscall"
80 "time"
81
82 "github.com/charmbracelet/bubbletea"
83@@ -24,18 +27,28 @@ const recordCap = 15 * time.Second
84 const recordCommand = "arecord"
85
86 func main() {
87+ ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
88+ defer stop()
89+
90 cfg := config.ParseConfig()
91
92- svc, err := bootstrap.New(cfg)
93+ svc, err := bootstrap.New(ctx, cfg)
94 if err != nil {
95 fatalf("%v", err)
96 }
97+ defer svc.Close()
98
99 state := game.NewState(svc.Brief)
100 orch := buildOrchestrator(svc, state)
101
102 m := ui.NewModel(state, orch, int(recordCap.Seconds()))
103- if _, err := tea.NewProgram(m).Run(); err != nil {
104+ p := tea.NewProgram(m)
105+ go func() {
106+ <-ctx.Done()
107+ p.Quit()
108+ }()
109+ if _, err := p.Run(); err != nil {
110+ svc.Close()
111 fatalf("%v", err)
112 }
113 }
114diff --git a/internal/bootstrap/bootstrap.go b/internal/bootstrap/bootstrap.go
115index 712e023508c7dca46aac3af3344141f50134b9ea..b0274203bf66ef710dd63e8479ab69f6080cf12a 100644
116--- a/internal/bootstrap/bootstrap.go
117+++ b/internal/bootstrap/bootstrap.go
118@@ -46,8 +46,9 @@ type Services struct {
119
120 // New loads the scenario, starts the self-hosted model services unless disabled,
121 // builds one LLM client per slot, and warms up the prompt slots. Any failure
122-// tears down whatever was started and is returned as an error.
123-func New(cfg *config.Config) (*Services, error) {
124+// tears down whatever was started and is returned as an error. Cancelling ctx
125+// aborts startup and cleans up any spawned children.
126+func New(ctx context.Context, cfg *config.Config) (*Services, error) {
127 sc, err := scenario.Load(cfg.Scenario)
128 if err != nil {
129 return nil, fmt.Errorf("load scenario: %w", err)
130@@ -57,7 +58,7 @@ func New(cfg *config.Config) (*Services, error) {
131 var manager *services.Manager
132 if !c.DisableModelLoading {
133 manager = services.NewManager()
134- if err := manager.Start(context.Background()); err != nil {
135+ if err := manager.Start(ctx); err != nil {
136 return nil, err
137 }
138 // Managed children bind loopback; pin the endpoints regardless of URL flags.
139@@ -81,11 +82,11 @@ func New(cfg *config.Config) (*Services, error) {
140 {compactionClient, llm.CompactionPrompt()},
141 {sheetClient, llm.SheetSystemPrompt()},
142 }
143- ctx, cancel := context.WithTimeout(context.Background(), warmupTimeout)
144+ warmCtx, cancel := context.WithTimeout(ctx, warmupTimeout)
145 defer cancel()
146 var warmErr error
147 for _, w := range warmups {
148- if warmErr = w.client.Warmup(ctx, []string{w.prompt}); warmErr != nil {
149+ if warmErr = w.client.Warmup(warmCtx, []string{w.prompt}); warmErr != nil {
150 break
151 }
152 }