1package ollama
 2
 3import (
 4	"context"
 5	"os"
 6	"os/exec"
 7	"os/signal"
 8	"syscall"
 9	"time"
10)
11
12// setupCleanup sets up signal handlers for cleanup
13func setupCleanup() {
14	c := make(chan os.Signal, 1)
15	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
16
17	go func() {
18		<-c
19		cleanup()
20		os.Exit(0)
21	}()
22}
23
24// cleanup stops all running models and service if started by Crush
25func cleanup() {
26	processManager.mu.Lock()
27	defer processManager.mu.Unlock()
28
29	// Stop all running models using HTTP API
30	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
31	defer cancel()
32
33	if IsRunning(ctx) {
34		stopAllModels(ctx)
35	}
36
37	// Stop Ollama service if we started it
38	if processManager.crushStartedOllama && processManager.ollamaProcess != nil {
39		stopOllamaService()
40	}
41}
42
43// stopAllModels stops all running models
44func stopAllModels(ctx context.Context) {
45	runningModels, err := GetRunningModels(ctx)
46	if err != nil {
47		return
48	}
49
50	for _, model := range runningModels {
51		stopModel(ctx, model.Name)
52	}
53}
54
55// stopModel stops a specific model using CLI
56func stopModel(ctx context.Context, modelName string) error {
57	cmd := exec.CommandContext(ctx, "ollama", "stop", modelName)
58	return cmd.Run()
59}
60
61// stopOllamaService stops the Ollama service process
62func stopOllamaService() {
63	if processManager.ollamaProcess == nil {
64		return
65	}
66
67	// Try graceful shutdown first
68	if err := processManager.ollamaProcess.Process.Signal(syscall.SIGTERM); err == nil {
69		// Wait for graceful shutdown
70		done := make(chan error, 1)
71		go func() {
72			done <- processManager.ollamaProcess.Wait()
73		}()
74
75		select {
76		case <-done:
77			// Process finished gracefully
78		case <-time.After(5 * time.Second):
79			// Force kill if not shut down gracefully
80			syscall.Kill(-processManager.ollamaProcess.Process.Pid, syscall.SIGKILL)
81			processManager.ollamaProcess.Wait()
82		}
83	}
84
85	processManager.ollamaProcess = nil
86	processManager.crushStartedOllama = false
87}