Maven/internal/phraser/spawn_test.go

package phraser

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"syscall"
	"testing"
	"time"
)

// The spawn path (NewLLMPhraser, spawnLlamaServer, startLlamaProc, llamaProc.Close)
// was at 0% coverage: every test built the phraser with NewLLMPhraserAt, which
// starts no process. These tests drive the real spawn code against a fake
// llama-server script, so the startup race arms and the reaping are exercised
// without a model or a GPU.

// fakeLlama writes an executable script standing in for llama-server and returns
// its path. body runs after the script has recorded its own pid.
func fakeLlama(t *testing.T, body string) string {
	t.Helper()
	dir := t.TempDir()
	path := filepath.Join(dir, "fake-llama-server")
	script := "#!/bin/sh\n" + body + "\n"
	if err := os.WriteFile(path, []byte(script), 0o755); err != nil {
		t.Fatalf("write fake server: %v", err)
	}
	return path
}

// listensThenSleeps prints the line startLlamaProc scrapes, then stays alive
// until killed — the shape of a real llama-server that came up.
const listensThenSleeps = `echo "srv    load_model: listening on http://127.0.0.1:18081" >&2
while : ; do sleep 1 ; done`

func testCfg(bin string) Config {
	cfg := DefaultConfig("/nonexistent/model.gguf")
	cfg.BinPath = bin
	return cfg
}

func TestExtractPort(t *testing.T) {
	for _, tc := range []struct{ in, want string }{
		{"127.0.0.1:0", "0"},
		{"127.0.0.1:8080", "8080"},
		{"127.0.0.1:", "0"},
		{"", "0"},
		{"8080", "0"}, // no colon: Cut yields no port, so the caller gets the "any port" default
	} {
		if got := extractPort(tc.in); got != tc.want {
			t.Errorf("extractPort(%q) = %q, want %q", tc.in, got, tc.want)
		}
	}
}

// The prompt cache is what ate 6.8GB of the deployed server's RSS, so the cap
// has to reach the command line, and the opt-out has to leave it off.
func TestLlamaArgsCapsPromptCache(t *testing.T) {
	cfg := DefaultConfig("/m.gguf")
	if got := strings.Join(llamaArgs(cfg), " "); !strings.Contains(got, "--cache-ram 512") {
		t.Errorf("default args = %q, want --cache-ram 512", got)
	}
	cfg.CacheRAMMiB = 0
	if got := strings.Join(llamaArgs(cfg), " "); strings.Contains(got, "--cache-ram") {
		t.Errorf("args with the cap off = %q, want no --cache-ram flag", got)
	}
}

func TestStartLlamaProcScrapesPortAndReaps(t *testing.T) {
	bin := fakeLlama(t, listensThenSleeps)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	p, err := startLlamaProc(ctx, testCfg(bin))
	if err != nil {
		t.Fatalf("startLlamaProc: %v", err)
	}
	if p.BaseURL() != "http://127.0.0.1:18081" {
		t.Fatalf("BaseURL = %q, want the scraped address", p.BaseURL())
	}
	pid := p.cmd.Process.Pid

	p.cancel = cancel
	if err := p.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}
	// Close must Wait, otherwise the child lingers as a zombie.
	if p.cmd.ProcessState == nil {
		t.Fatal("Close did not reap the child: ProcessState is nil")
	}
	if err := syscall.Kill(pid, 0); err == nil {
		t.Fatalf("child %d still exists after Close", pid)
	}
}

// captureLog redirects the standard logger for the duration of a test and
// returns what was written to it.
func captureLog(t *testing.T) *bytes.Buffer {
	t.Helper()
	var buf bytes.Buffer
	old := log.Writer()
	flags := log.Flags()
	log.SetOutput(&buf)
	log.SetFlags(0)
	t.Cleanup(func() { log.SetOutput(old); log.SetFlags(flags) })
	return &buf
}

// The child's buffer-size, KV-cache and offload lines are the only way to
// account for its memory on a running box, and they used to be dropped: stderr
// was scraped for the listen line and thrown away, stdout was never piped.
func TestStartLlamaProcForwardsChildOutput(t *testing.T) {
	buf := captureLog(t)
	bin := fakeLlama(t, `echo "load_tensors: Vulkan0 model buffer size = 1053.34 MiB" >&2
echo "llama_context: KV self size = 448.00 MiB"
`+listensThenSleeps)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	p, err := startLlamaProc(ctx, testCfg(bin))
	if err != nil {
		t.Fatalf("startLlamaProc: %v", err)
	}
	p.cancel = cancel
	defer p.Close()

	got := buf.String()
	for _, want := range []string{
		"llama: load_tensors: Vulkan0 model buffer size = 1053.34 MiB", // stderr
		"llama: llama_context: KV self size = 448.00 MiB",              // stdout, previously discarded
	} {
		if !strings.Contains(got, want) {
			t.Errorf("log missing %q\nlog was:\n%s", want, got)
		}
	}
}

func TestStartLlamaProcFailureArms(t *testing.T) {
	t.Run("binary missing", func(t *testing.T) {
		cfg := testCfg(filepath.Join(t.TempDir(), "does-not-exist"))
		_, err := startLlamaProc(context.Background(), cfg)
		if err == nil || !strings.Contains(err.Error(), "llm: start") {
			t.Fatalf("err = %v, want a start failure", err)
		}
	})

	t.Run("server exits without listening", func(t *testing.T) {
		// stderr closes, so the reader goroutine reports EOF on errCh. The error
		// must carry the child's last words: bare "EOF" named no cause.
		captureLog(t)
		bin := fakeLlama(t, `echo "ggml_vulkan: no device" >&2
exit 1`)
		_, err := startLlamaProc(context.Background(), testCfg(bin))
		if err == nil || !strings.Contains(err.Error(), "llm: server output") {
			t.Fatalf("err = %v, want the server-output arm", err)
		}
		if !strings.Contains(err.Error(), "ggml_vulkan: no device") {
			t.Errorf("err = %v, want the child's last output in it", err)
		}
	})

	t.Run("context cancelled during startup", func(t *testing.T) {
		// Never prints the listen line and never exits: only ctx can end this.
		bin := fakeLlama(t, `while : ; do sleep 1 ; done`)
		ctx, cancel := context.WithCancel(context.Background())
		go func() {
			time.Sleep(150 * time.Millisecond)
			cancel()
		}()
		defer cancel()
		_, err := startLlamaProc(ctx, testCfg(bin))
		if !errors.Is(err, context.Canceled) {
			t.Fatalf("err = %v, want context.Canceled", err)
		}
	})
}

func TestNewLLMPhraserSpawns(t *testing.T) {
	bin := fakeLlama(t, listensThenSleeps)
	p, err := NewLLMPhraser(context.Background(), testCfg(bin))
	if err != nil {
		t.Fatalf("NewLLMPhraser: %v", err)
	}
	if p.BaseURL() != "http://127.0.0.1:18081" {
		t.Fatalf("BaseURL = %q", p.BaseURL())
	}
	pid := p.be.(*llamaProc).cmd.Process.Pid
	if err := p.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}
	if p.BaseURL() != "" {
		t.Fatalf("BaseURL after Close = %q, want empty", p.BaseURL())
	}
	if err := syscall.Kill(pid, 0); err == nil {
		t.Fatalf("llama-server %d survived Close", pid)
	}
}

func TestNewLLMPhraserSpawnFailure(t *testing.T) {
	cfg := testCfg(filepath.Join(t.TempDir(), "does-not-exist"))
	p, err := NewLLMPhraser(context.Background(), cfg)
	if err == nil {
		p.Close()
		t.Fatal("want an error when the server cannot start")
	}
	if p != nil {
		t.Fatalf("want a nil phraser on failure, got %#v", p)
	}
}

// TestPdeathsigKillsOrphan is the orphan test the task asked for. A SIGKILLed
// mavend never runs Close, so nothing but the kernel's Pdeathsig can stop its
// llama-server. Re-exec this test binary as the "daemon", let it spawn the fake
// server, SIGKILL the daemon, and assert the grandchild died with it.
func TestPdeathsigKillsOrphan(t *testing.T) {
	bin := fakeLlama(t, listensThenSleeps)

	cmd := exec.Command(os.Args[0], "-test.run=TestSpawnHelperProcess", "-test.v=false")
	cmd.Env = append(os.Environ(), "MAVEN_SPAWN_HELPER=1", "MAVEN_FAKE_LLAMA="+bin)
	out, err := cmd.StdoutPipe()
	if err != nil {
		t.Fatalf("stdout pipe: %v", err)
	}
	if err := cmd.Start(); err != nil {
		t.Fatalf("start helper: %v", err)
	}
	defer func() { _ = cmd.Process.Kill(); _ = cmd.Wait() }()

	buf := make([]byte, 256)
	n, err := out.Read(buf)
	if err != nil {
		t.Fatalf("read child pid: %v", err)
	}
	childPID, err := strconv.Atoi(strings.TrimSpace(string(buf[:n])))
	if err != nil {
		t.Fatalf("helper printed %q, want a pid: %v", string(buf[:n]), err)
	}
	if err := syscall.Kill(childPID, 0); err != nil {
		t.Fatalf("llama-server %d not running before the kill: %v", childPID, err)
	}

	// SIGKILL: the helper gets no chance to clean up, exactly like an OOM kill.
	if err := cmd.Process.Signal(syscall.SIGKILL); err != nil {
		t.Fatalf("kill helper: %v", err)
	}
	_, _ = cmd.Process.Wait()

	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		if err := syscall.Kill(childPID, 0); err != nil {
			return // gone: Pdeathsig did its job
		}
		time.Sleep(20 * time.Millisecond)
	}
	_ = syscall.Kill(childPID, syscall.SIGKILL)
	t.Fatalf("llama-server %d outlived the SIGKILLed parent", childPID)
}

// TestSpawnHelperProcess is not a test. It is the child half of
// TestPdeathsigKillsOrphan: spawn a llama-server, print its pid, then block.
func TestSpawnHelperProcess(t *testing.T) {
	if os.Getenv("MAVEN_SPAWN_HELPER") != "1" {
		t.Skip("helper for TestPdeathsigKillsOrphan")
	}
	cfg := testCfg(os.Getenv("MAVEN_FAKE_LLAMA"))
	p, err := startLlamaProc(context.Background(), cfg)
	if err != nil {
		fmt.Println("spawn failed:", err)
		os.Exit(1)
	}
	fmt.Println(p.cmd.Process.Pid)
	os.Stdout.Sync()
	select {} // wait to be killed
}

// TestKillMavenScriptMatchesRealCommandLine pins kill-maven.sh's fallback
// pattern to the command line startLlamaProc actually builds. The script leaked
// orphans twice already, both times because the pattern stopped matching: first
// `llama-server.*maven`, then a hardcoded model name after the model was swapped.
func TestKillMavenScriptMatchesRealCommandLine(t *testing.T) {
	src, err := os.ReadFile("../../kill-maven.sh")
	if err != nil {
		t.Fatalf("read kill-maven.sh: %v", err)
	}
	m := regexp.MustCompile(`(?m)^\s*LLM='([^']+)'`).FindSubmatch(src)
	if m == nil {
		t.Fatal("no default LLM='...' pattern in kill-maven.sh")
	}
	pat, err := regexp.Compile(string(m[1]))
	if err != nil {
		t.Fatalf("LLM pattern %q does not compile: %v", m[1], err)
	}

	// Rebuild the command line from the production arg list, so a change to
	// startLlamaProc that breaks the sweep fails here instead of on the box.
	cfg := DefaultConfig("/opt/maven/models/llm/Qwen3-1.7B-UD-Q4_K_XL.gguf")
	cfg.NCtx, cfg.NGpuLayers = 4096, 99
	cmdline := cfg.BinPath + " " + strings.Join(llamaArgs(cfg), " ")
	if !pat.MatchString(cmdline) {
		t.Fatalf("kill-maven.sh pattern %q does not match %q — orphans would leak", m[1], cmdline)
	}
}