llmgw/internal/cli/judgekeeper_test.go

package cli

import (
	"context"
	"errors"
	"log/slog"
	"slices"
	"testing"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/modelserve"
	"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
)

type fakeServe struct {
	started []modelserve.Request
	running []modelserve.Job
	err     error
}

func (f *fakeServe) Start(r modelserve.Request) (modelserve.Job, error) {
	if f.err != nil {
		return modelserve.Job{}, f.err
	}
	f.started = append(f.started, r)
	return modelserve.Job{ID: "j", Model: r.Name, Request: r, State: modelserve.Running}, nil
}
func (f *fakeServe) Jobs() []modelserve.Job { return f.running }

func TestAJudgeAPolicyNamesIsFetchedWhenTheCatalogLacksIt(t *testing.T) {
	lib := guard.Config{Policies: map[string]guard.Policy{
		"baseline": {Judge: &guard.JudgePolicy{Models: []string{"qwen3guard-4b", "qwen3guard-0.6b"}, Input: "block"}},
		"custom":   {Judge: &guard.JudgePolicy{Model: "our-own-judge", Input: "block"}},
	}}
	inCatalog := []catalog.Model{{Name: "qwen3guard-0.6b"}}
	now := time.Date(2026, 10, 8, 12, 0, 0, 0, time.UTC)
	fs := &fakeServe{}
	leader := true
	k := &judgeKeeper{
		library: func() guard.Config { return lib },
		models:  func() []catalog.Model { return inCatalog },
		serve:   fs,
		leader:  func() bool { return leader },
		gpus:    func() []adminapi.ClassUse { return []adminapi.ClassUse{{Class: "6000-24g", Total: 1, Used: 1}} },
		log:     slog.New(slog.DiscardHandler),
		now:     func() time.Time { return now },
	}

	k.check()
	// The GPU judge alone: the CPU one is there, and one llmgw does not know
	// is not guessed at.
	if len(fs.started) != 1 || fs.started[0].Name != "qwen3guard-4b" || fs.started[0].GPUClass != "6000-24g" ||
		fs.started[0].Class != catalog.ClassStandard || fs.started[0].GuardFormat != "qwen3guard" || fs.started[0].CPU ||
		!fs.started[0].IfAbsent {
		t.Fatalf("started %+v", fs.started)
	}

	// Being fetched: not asked for again.
	fs.running = []modelserve.Job{{Model: "qwen3guard-4b", State: modelserve.Running}}
	k.check()
	if len(fs.started) != 1 {
		t.Errorf("asked for again while it is being fetched: %d", len(fs.started))
	}
	// Failed: asked for again only after a while.
	fs.running = nil
	now = now.Add(time.Minute)
	k.check()
	if len(fs.started) != 1 {
		t.Error("asked for again at once after it failed")
	}
	now = now.Add(judgeRetry)
	k.check()
	if len(fs.started) != 2 {
		t.Error("never asked for again")
	}

	// Only the leader fetches, and nothing while the guardrails are off.
	leader = false
	now = now.Add(judgeRetry)
	k.check()
	leader = true
	off := false
	lib.Enabled = &off
	k.check()
	if len(fs.started) != 2 {
		t.Errorf("a standby or the guardrails off fetched: %d", len(fs.started))
	}

	// A CPU judge on a cluster without GPUs is served on CPUs, two of it.
	lib = guard.Config{Policies: map[string]guard.Policy{"b": {Judge: &guard.JudgePolicy{Model: "qwen3guard-0.6b", Input: "block"}}}}
	inCatalog = nil
	k.gpus = func() []adminapi.ClassUse { return nil }
	k.nodeCPUs = func() []int { return []int{8, 4} }
	k.check()
	// Sized as init sizes its two judges: 2 GB, and two cores each on a
	// cluster whose smallest node has four.
	if last := fs.started[len(fs.started)-1]; last.Name != "qwen3guard-0.6b" || !last.CPU || last.Replicas != 2 ||
		last.MemoryGB != judgeMemoryGB || last.CPUCores != refitCores(2, []int{8, 4}) || last.CPUCores == 0 {
		t.Errorf("CPU judge: %+v", last)
	}
	fs.err = errors.New("the mirror is down")
	lib = guard.Config{Policies: map[string]guard.Policy{"b": {Judge: &guard.JudgePolicy{Model: "qwen3guard-8b", Input: "block"}}}}
	k.check() // no GPU, so not started, and no panic on the error path either
}

// A policy moved from Llama Guard's codes to Qwen3Guard would count nothing:
// its categories become the judge's own.
func TestCategoriesFollowTheJudges(t *testing.T) {
	lg := []string{"S1", "S2", "S10"}
	if got := categoriesFor([]string{"qwen3guard-0.6b"}, lg); !slices.Equal(got, baselineCategories["qwen3guard"]) {
		t.Errorf("Llama Guard's codes for Qwen3Guard: %v", got)
	}
	if got := categoriesFor([]string{"qwen3guard-4b", "qwen3guard-0.6b"}, []string{"Violent", "Jailbreak"}); got != nil {
		t.Errorf("its own names were replaced: %v", got)
	}
	if got := categoriesFor([]string{"llama-guard-3-1b"}, lg); got != nil {
		t.Errorf("Llama Guard's own codes were replaced: %v", got)
	}
	// Both kinds asked in turn: each must have something to count.
	got := categoriesFor([]string{"qwen3guard-4b", "llama-guard-3-1b"}, []string{"Violent"})
	if !slices.Contains(got, "Violent") || !slices.Contains(got, "S1") {
		t.Errorf("mixed: %v", got)
	}
	if got := categoriesFor([]string{"our-own-judge"}, lg); got != nil {
		t.Errorf("a judge llmgw does not know: %v", got)
	}
	if got := categoriesFor([]string{"qwen3guard-0.6b"}, nil); got != nil {
		t.Errorf("no categories (every one counts) became %v", got)
	}
}

type fakeProfiles struct {
	set map[string][]catalog.Profile
}

func (f *fakeProfiles) SetProfiles(_ context.Context, name string, ps []catalog.Profile, _ string) (catalog.Model, error) {
	if f.set == nil {
		f.set = map[string][]catalog.Profile{}
	}
	f.set[name] = ps
	return catalog.Model{Name: name, Profiles: ps}, nil
}

// A judge that v0.14.6 or earlier's init installed is moved to today's
// figures, once, and its replicas restarted onto them; one somebody changed
// is left alone.
func TestAnEarlierInstallsJudgeIsRefittedOnce(t *testing.T) {
	lib := guard.Config{Policies: map[string]guard.Policy{
		"baseline": {Judge: &guard.JudgePolicy{Model: "qwen3guard-0.6b", Input: "block"}},
	}}
	earlier := guardModel("http://10.0.0.5:8081/qwen3guard-0.6b/guard/g.gguf", 2, 1, nil)
	earlier.Profiles[0].MinMemoryGB, earlier.Profiles[0].Args = 4, []string{"--ctx-size", "8192"}
	inCatalog := []catalog.Model{earlier}
	fp := &fakeProfiles{}
	var restarted []string
	k := &judgeKeeper{
		library:  func() guard.Config { return lib },
		models:   func() []catalog.Model { return inCatalog },
		serve:    &fakeServe{},
		log:      slog.New(slog.DiscardHandler),
		profiles: fp,
		restart:  func(ids ...string) error { restarted = append(restarted, ids...); return nil },
		replicas: func(string) []string { return []string{"j-1", "j-2"} },
		nodeCPUs: func() []int { return []int{4} },
	}
	k.check()
	got := fp.set["qwen3guard-0.6b"]
	if len(got) != 1 || got[0].MinMemoryGB != 2 || got[0].CPUCores != 2 || !slices.Contains(got[0].Args, "--cache-ram") {
		t.Fatalf("refitted to %+v", got)
	}
	if !slices.Equal(restarted, []string{"j-1", "j-2"}) {
		t.Errorf("restarted %v", restarted)
	}
	k.check()
	if len(restarted) != 2 {
		t.Error("refitted twice")
	}

	// Changed by somebody: an extra arg. Left as it is.
	mine := guardModel("http://x/g.gguf", 2, 3, nil)
	mine.Profiles[0].Args = []string{"--ctx-size", "16384"}
	inCatalog = []catalog.Model{mine}
	fp.set, k.looked = nil, nil
	k.check()
	if fp.set != nil {
		t.Errorf("a judge set by hand was changed: %+v", fp.set)
	}
}

func TestAJudgeIsRefittedToTheCoresInitWouldGive(t *testing.T) {
	for _, c := range []struct {
		replicas int
		cpus     []int
		want     int
	}{
		{1, []int{4}, 4},
		{1, []int{2}, 2},
		{1, []int{16}, 4},
		{2, []int{4}, 2},
		{2, []int{2}, 1},
		{2, []int{4, 32}, 2},
		{2, nil, 2},
	} {
		if got := refitCores(c.replicas, c.cpus); got != c.want {
			t.Errorf("%d judges on %v: %d cores, want %d", c.replicas, c.cpus, got, c.want)
		}
	}
}

// Just after the control plane starts, the judge init is adding is left to
// init: the keeper serves one itself only after its grace.
func TestTheKeeperLeavesAJustStartedInstallsJudgeToInit(t *testing.T) {
	lib := guard.Config{Policies: map[string]guard.Policy{"b": {Judge: &guard.JudgePolicy{Model: "qwen3guard-0.6b", Input: "block"}}}}
	now := time.Date(2026, 10, 19, 12, 0, 0, 0, time.UTC)
	fs := &fakeServe{}
	k := &judgeKeeper{
		library: func() guard.Config { return lib },
		models:  func() []catalog.Model { return nil },
		serve:   fs, grace: judgeGrace,
		gpus: func() []adminapi.ClassUse { return nil },
		log:  slog.New(slog.DiscardHandler),
		now:  func() time.Time { return now },
	}
	k.check()
	now = now.Add(judgeGrace / 2)
	k.check()
	if len(fs.started) != 0 {
		t.Fatalf("served within its grace: %+v", fs.started)
	}
	now = now.Add(judgeGrace)
	k.check()
	if len(fs.started) != 1 {
		t.Errorf("never served after its grace: %d", len(fs.started))
	}
}