1package ollama
 2
 3import (
 4	"context"
 5	"os"
 6	"os/exec"
 7	"os/signal"
 8	"syscall"
 9	"time"
10)
11
12var processManager = &ProcessManager{
13	processes: make(map[string]*exec.Cmd),
14}
15
16// setupProcessCleanup sets up signal handlers to clean up processes on exit
17func setupProcessCleanup() {
18	c := make(chan os.Signal, 1)
19	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
20
21	go func() {
22		<-c
23		cleanupProcesses()
24		os.Exit(0)
25	}()
26}
27
28// cleanupProcesses terminates all Ollama processes started by Crush
29func cleanupProcesses() {
30	processManager.mu.Lock()
31	defer processManager.mu.Unlock()
32
33	// Use CLI approach to stop all running models
34	// This is more reliable than tracking individual processes
35	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
36	defer cancel()
37
38	if err := CLIStopAllModels(ctx); err != nil {
39		// If CLI approach fails, fall back to process tracking
40		// Clean up model processes
41		for modelName, cmd := range processManager.processes {
42			if cmd.Process != nil {
43				cmd.Process.Kill()
44				cmd.Wait() // Wait for the process to actually exit
45			}
46			delete(processManager.processes, modelName)
47		}
48	} else {
49		// CLI approach succeeded, clear our process tracking
50		processManager.processes = make(map[string]*exec.Cmd)
51	}
52
53	// Clean up Ollama server if Crush started it
54	if processManager.crushStartedOllama && processManager.ollamaServer != nil {
55		if processManager.ollamaServer.Process != nil {
56			// Kill the entire process group to ensure all children are terminated
57			syscall.Kill(-processManager.ollamaServer.Process.Pid, syscall.SIGTERM)
58
59			// Give it a moment to shut down gracefully
60			time.Sleep(2 * time.Second)
61
62			// Force kill if still running
63			if processManager.ollamaServer.ProcessState == nil {
64				syscall.Kill(-processManager.ollamaServer.Process.Pid, syscall.SIGKILL)
65			}
66
67			processManager.ollamaServer.Wait() // Wait for the process to actually exit
68		}
69		processManager.ollamaServer = nil
70		processManager.crushStartedOllama = false
71	}
72}