llmgw/internal/router/guard_test.go

package router

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus/testutil"

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

// wordGuard blocks requests saying "forbidden", rewrites "alice@example.com",
// flags "maybe", and stops answers saying "hi".
type wordGuard struct {
	err     error
	subject guard.Subject
	// answers, when set, stops answers saying "hi".
	answers bool
	noStore bool
}

func (g *wordGuard) Input(_ context.Context, s guard.Subject, body []byte) (guard.Verdict, error) {
	g.subject = s
	switch {
	case g.err != nil:
		return guard.Verdict{}, g.err
	case bytes.Contains(body, []byte("forbidden")):
		return guard.Verdict{Action: guard.Block, Policy: "words", Category: "test"}, nil
	case bytes.Contains(body, []byte("alice@example.com")):
		return guard.Verdict{Action: guard.Redact, Policy: "pii", Category: "email",
			Body: bytes.ReplaceAll(body, []byte("alice@example.com"), []byte("[EMAIL_1]"))}, nil
	case bytes.Contains(body, []byte("maybe")):
		return guard.Verdict{Action: guard.Flag, Policy: "words"}, nil
	}
	return guard.Verdict{}, nil
}

func (g *wordGuard) Output(context.Context, guard.Subject) guard.OutputFilter {
	if !g.answers {
		return nil
	}
	return &stopHi{}
}

func (g *wordGuard) Directives(guard.Subject) guard.Directives {
	return guard.Directives{NoStore: g.noStore}
}

type stopHi struct{ v guard.Verdict }

func (f *stopHi) Buffered(body []byte) []byte {
	if !bytes.Contains(body, []byte(`"hi"`)) {
		return body
	}
	f.v = guard.Verdict{Action: guard.Block, Stage: guard.Output, Policy: "answers"}
	return []byte(`{"choices":[{"index":0,"message":{"role":"assistant","content":""},"finish_reason":"content_filter"}]}`)
}
func (f *stopHi) OnEvent(d string) ([]string, bool) { return []string{d}, false }
func (f *stopHi) OnEnd() []string                   { return nil }
func (f *stopHi) Verdict() guard.Verdict            { return f.v }

func guardedRouter(t *testing.T, g guard.Guard) (*Router, *fakeEngine) {
	e := newFakeEngine(t, "m")
	rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}}))
	rt.SetGuard(g)
	return rt, e
}

func say(text string) map[string]any {
	return map[string]any{"model": "m", "messages": []any{map[string]any{"role": "user", "content": text}}}
}

func TestABlockedRequestIsRefusedBeforeTheEngine(t *testing.T) {
	rt, e := guardedRouter(t, &wordGuard{})
	rec := post(t, rt.Handler(), say("something forbidden"), nil)
	var body struct {
		Error struct{ Type, Code, Message string } `json:"error"`
	}
	_ = json.Unmarshal(rec.Body.Bytes(), &body)
	if rec.Code != http.StatusBadRequest || body.Error.Code != policyViolation ||
		rec.Header().Get(guardHeader) != "blocked" || rec.Header().Get(guardPolicyHeader) != "words" {
		t.Errorf("%d %v %s", rec.Code, rec.Header(), rec.Body.String())
	}
	if strings.Contains(rec.Body.String(), "forbidden") {
		t.Error("the refusal quotes what it refused")
	}
	if e.calls.Load() != 0 {
		t.Error("the engine was asked anyway")
	}
}

func TestARedactedRequestReachesTheEngineRewritten(t *testing.T) {
	rt, e := guardedRouter(t, &wordGuard{})
	rec := post(t, rt.Handler(), say("write to alice@example.com"), nil)
	if rec.Header().Get(guardHeader) != "redacted" {
		t.Errorf("headers %v", rec.Header())
	}
	sent, _ := e.last.Load().([]byte)
	if bytes.Contains(sent, []byte("alice@example.com")) || !bytes.Contains(sent, []byte("[EMAIL_1]")) {
		t.Errorf("the engine was sent %s", sent)
	}
}

func TestAnAnswerAGuardStopsEndsAsContentFilter(t *testing.T) {
	rt, _ := guardedRouter(t, &wordGuard{answers: true})
	rec := post(t, rt.Handler(), say("hello"), nil)
	if rec.Code != http.StatusOK || !strings.Contains(rec.Body.String(), `"finish_reason":"content_filter"`) ||
		strings.Contains(rec.Body.String(), `"hi"`) || rec.Header().Get(guardHeader) != "blocked" {
		t.Errorf("%d %v %s", rec.Code, rec.Header(), rec.Body.String())
	}
}

func TestAGuardThatCannotJudgeRefuses(t *testing.T) {
	g := &wordGuard{err: errors.New("judge down")}
	rt, e := guardedRouter(t, g)
	rec := post(t, rt.Handler(), say("hello"), nil)
	if rec.Code != http.StatusServiceUnavailable || e.calls.Load() != 0 || strings.Contains(rec.Body.String(), "judge down") {
		t.Errorf("%d %s, engine calls %d", rec.Code, rec.Body.String(), e.calls.Load())
	}
	if g.subject.Model != "m" || g.subject.Path != "/v1/chat/completions" {
		t.Errorf("judged as %+v", g.subject)
	}
}

func TestAFlaggedRequestIsServedAndSaysSo(t *testing.T) {
	rt, e := guardedRouter(t, &wordGuard{})
	rec := post(t, rt.Handler(), say("maybe"), nil)
	if rec.Code != http.StatusOK || rec.Header().Get(guardHeader) != "flagged" || e.calls.Load() != 1 {
		t.Errorf("%d %v", rec.Code, rec.Header())
	}
}

func TestABlockedRequestIsRecordedWithItsVerdictAndCostsNothing(t *testing.T) {
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt, secret := pricedRouter(t, e.srv.URL)
	rt.SetGuard(&wordGuard{})
	rec := post(t, rt.Handler(),
		map[string]any{"model": "q14", "messages": []any{map[string]any{"role": "user", "content": "forbidden"}}},
		map[string]string{"Authorization": "Bearer " + secret, "X-Request-Id": "blocked-1"})
	if rec.Code != http.StatusBadRequest {
		t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
	}
	deadline := time.Now().Add(5 * time.Second)
	for {
		u, ok, err := rt.store.UsageByRequest(t.Context(), "blocked-1")
		if err != nil {
			t.Fatal(err)
		}
		if ok {
			if u.GuardAction != "blocked" || u.GuardPolicy != "words" || u.GuardCategory != "test" ||
				u.CostMicros != 0 || u.PromptTokens != 0 || u.Status != http.StatusBadRequest {
				t.Errorf("recorded %+v", u)
			}
			return
		}
		if time.Now().After(deadline) {
			t.Fatal("the blocked request was not recorded")
		}
		time.Sleep(20 * time.Millisecond)
	}
}

func TestAPolicyThatKeepsNothingRefusesWhatWouldBeKept(t *testing.T) {
	rt, e := guardedRouter(t, &wordGuard{noStore: true})
	rec := post(t, rt.Handler(), say("hello"), map[string]string{"x-llmgw-thread": "9b2f6c1e-4a0d"})
	if rec.Code != http.StatusBadRequest || !strings.Contains(rec.Body.String(), `"code":"no_store"`) || e.calls.Load() != 0 {
		t.Errorf("a conversation under no_store: %d %s", rec.Code, rec.Body.String())
	}
	if rec := post(t, rt.Handler(), say("hello"), nil); rec.Code != http.StatusOK {
		t.Errorf("a plain request under no_store: %d %s", rec.Code, rec.Body.String())
	}
}

func TestTheGateJudgesRealTraffic(t *testing.T) {
	g, err := guard.NewGate(guard.Config{Default: []string{"base"}, Policies: map[string]guard.Policy{"base": {
		Input:  guard.Rules{PII: map[string]string{"email": "redact"}},
		Output: guard.Rules{Deny: []guard.Pattern{{Name: "greeting", Regex: `\bhi\b`, Action: "block"}}},
	}}}, nil, nil)
	if err != nil {
		t.Fatal(err)
	}
	rt, e := guardedRouter(t, g)
	rec := post(t, rt.Handler(), say("tell carol@example.com"), nil)
	sent, _ := e.last.Load().([]byte)
	if bytes.Contains(sent, []byte("carol@")) || !bytes.Contains(sent, []byte("[EMAIL_1]")) {
		t.Errorf("the engine was sent %s", sent)
	}
	// The fake engine answers "hi", which the output policy stops.
	if !strings.Contains(rec.Body.String(), "content_filter") || rec.Header().Get(guardHeader) != "blocked" {
		t.Errorf("%v %s", rec.Header(), rec.Body.String())
	}
}

func TestAJudgeModelIsNobodysToAskButTheGuardrails(t *testing.T) {
	e := newFakeEngine(t, "llama-guard")
	cfg := baseConfig(config.Endpoint{ID: "j", URL: e.srv.URL, Models: []string{"llama-guard"}})
	judge := catModel("llama-guard", []string{"lg"})
	judge.Role, judge.Guard = catalog.RoleGuard, &catalog.GuardModel{Format: "llamaguard3"}
	cfg.Models = []catalog.Model{judge}
	rt := newTestRouter(t, cfg)
	waitForEndpoint(t, rt)

	for _, name := range []string{"llama-guard", "lg"} {
		if rec := post(t, rt.Handler(), map[string]any{"model": name, "messages": []any{}}, nil); rec.Code != http.StatusNotFound {
			t.Errorf("a client asked %s: %d %s", name, rec.Code, rec.Body.String())
		}
	}
	req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
	rec := httptest.NewRecorder()
	rt.Handler().ServeHTTP(rec, req)
	if strings.Contains(rec.Body.String(), "llama-guard") {
		t.Errorf("listed: %s", rec.Body.String())
	}
	calls := e.calls.Load()
	out, err := rt.Judge(t.Context(), "llama-guard", "/v1/chat/completions", []byte(`{"model":"llama-guard","messages":[{"role":"user","content":"x"}]}`))
	if err != nil || !strings.Contains(string(out), `"hi"`) || e.calls.Load() != calls+1 {
		t.Errorf("the guardrails asking: %s %v", out, err)
	}
	if f, _, ok := rt.GuardFormat("lg"); !ok || f != "llamaguard3" {
		t.Errorf("format %q %v", f, ok)
	}
}

func TestWhatTheGuardrailsDidIsCounted(t *testing.T) {
	rt, _ := guardedRouter(t, &wordGuard{answers: true})
	post(t, rt.Handler(), say("something forbidden"), nil)
	post(t, rt.Handler(), say("hello"), nil)
	if n := testutil.ToFloat64(rt.m.guard.checks.WithLabelValues("input", "words", "blocked", "test", "enforce")); n != 1 {
		t.Errorf("input blocks counted %v", n)
	}
	if n := testutil.ToFloat64(rt.m.guard.checks.WithLabelValues("output", "answers", "blocked", "", "enforce")); n != 1 {
		t.Errorf("output blocks counted %v", n)
	}
}

func TestAJudgeModelIsNotNamedToClientsInErrors(t *testing.T) {
	e := newFakeEngine(t, "llama-guard")
	cfg := baseConfig(config.Endpoint{ID: "j", URL: e.srv.URL, Models: []string{"llama-guard"}})
	judge := catModel("llama-guard", nil)
	judge.Role, judge.Guard = catalog.RoleGuard, &catalog.GuardModel{Format: "llamaguard3"}
	cfg.Models = []catalog.Model{judge, catModel("absent", nil)}
	rt := newTestRouter(t, cfg)
	waitForEndpoint(t, rt)
	rec := post(t, rt.Handler(), map[string]any{"model": "absent", "messages": []any{}}, nil)
	if strings.Contains(rec.Body.String(), "llama-guard") {
		t.Errorf("named: %s", rec.Body.String())
	}
}

func TestSignInThroughAProviderNeedsNoAPIKey(t *testing.T) {
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt, _ := pricedRouter(t, e.srv.URL)
	rt.oidc = config.OIDC{Issuer: "http://127.0.0.1:1", ClientID: "llmgw"}
	rec := httptest.NewRecorder()
	rt.Handler().ServeHTTP(rec, httptest.NewRequest("GET", "/login/oidc", nil))
	// The provider is not there, so the page says so — but it is the page,
	// not the API's "no key".
	if rec.Code == http.StatusUnauthorized || !strings.Contains(rec.Body.String(), "sign-in service") {
		t.Errorf("%d %s", rec.Code, rec.Body.String())
	}
}