llmgw/internal/cli/judgekeeper_test.go

package cli

import (
	"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 {
		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.check()
	if last := fs.started[len(fs.started)-1]; last.Name != "qwen3guard-0.6b" || !last.CPU || last.Replicas != 2 {
		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)
	}
}