ModelRouter/internal/gateway/key_quota_api_test.go

package gateway

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"

	"llmsproxy/internal/config"
	"llmsproxy/internal/core"
)

// adminGateway builds a gateway whose admin key can call /api/keys.
func adminGateway(t *testing.T, keys ...config.GWKey) *Gateway {
	t.Helper()
	td := t.TempDir()
	cfgPath := td + "/config.yaml"
	if err := os.WriteFile(cfgPath, []byte("listen: :0"), 0o644); err != nil {
		t.Fatal(err)
	}
	cfg := &config.Config{
		Path:        cfgPath,
		AdapterDir:  td + "/adapters",
		RuntimeFile: td + "/runtime.json",
		Keys:        keys,
	}
	if err := cfg.ApplyDefaults(); err != nil {
		t.Fatal(err)
	}
	c, err := core.NewFromConfig(cfg)
	if err != nil {
		t.Fatalf("core: %v", err)
	}
	t.Cleanup(c.Close)
	g, err := New(c, []string{"sk-admin"})
	if err != nil {
		t.Fatalf("gateway: %v", err)
	}
	return g
}

func adminReq(t *testing.T, g *Gateway, method, path, body string) *httptest.ResponseRecorder {
	t.Helper()
	req, _ := http.NewRequest(method, path, strings.NewReader(body))
	req.Header.Set("Authorization", "Bearer sk-admin")
	req.Header.Set("Content-Type", "application/json")
	rr := httptest.NewRecorder()
	g.Handler().ServeHTTP(rr, req)
	return rr
}

// keyRecord pulls one key's stored record out of the admin list.
func keyRecord(t *testing.T, g *Gateway, secret string) config.GWKey {
	t.Helper()
	rr := adminReq(t, g, "GET", "/api/keys", "")
	if rr.Code != 200 {
		t.Fatalf("GET /api/keys: %d %s", rr.Code, rr.Body.String())
	}
	var out struct {
		Keys []config.GWKey `json:"keys"`
	}
	if err := json.Unmarshal(rr.Body.Bytes(), &out); err != nil {
		t.Fatalf("decode: %v (%s)", err, rr.Body.String())
	}
	for _, k := range out.Keys {
		if k.Key == secret {
			return k
		}
	}
	t.Fatalf("key %q not found in %s", secret, rr.Body.String())
	return config.GWKey{}
}

func scopeOf(t *testing.T, k config.GWKey, model string) config.ModelScope {
	t.Helper()
	for _, m := range k.Models {
		if m.Model == model {
			return m
		}
	}
	t.Fatalf("scope %q not found in %+v", model, k.Models)
	return config.ModelScope{}
}

// Quotas live on the scope entries, not on the key: creating a key with a
// budget means creating scopes that carry it, and they must survive a
// read-back (persisted, not just echoed).
func TestKeyAPICreatesPerModelQuota(t *testing.T) {
	g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"})
	rr := adminReq(t, g, "POST", "/api/keys",
		`{"name":"agent-x","role":"user","models":[{"model":"m1","token_quota":50000,"req_quota":200,"period":"nhour","hours":6}]}`)
	if rr.Code != 200 {
		t.Fatalf("create: %d %s", rr.Code, rr.Body.String())
	}
	var created struct {
		Key config.GWKey `json:"key"`
	}
	if err := json.Unmarshal(rr.Body.Bytes(), &created); err != nil {
		t.Fatalf("decode: %v", err)
	}
	sc := scopeOf(t, created.Key, "m1")
	if sc.TokenQuota != 50000 || sc.ReqQuota != 200 || sc.Period != "nhour" || sc.Hours != 6 {
		t.Fatalf("created scope did not carry the caps: %+v", sc)
	}
	back := scopeOf(t, keyRecord(t, g, created.Key.Key), "m1")
	if back.TokenQuota != 50000 || back.ReqQuota != 200 || back.Period != "nhour" || back.Hours != 6 {
		t.Errorf("read-back lost the caps: %+v", back)
	}
}

// Two models on one key carry independent budgets.
func TestKeyAPIKeepsPerModelQuotaIndependent(t *testing.T) {
	g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"})
	rr := adminReq(t, g, "POST", "/api/keys",
		`{"name":"agent-y","role":"user","models":[{"model":"m1","token_quota":1000,"period":"hour"},{"model":"m2","token_quota":9999,"req_quota":7,"period":"week"}]}`)
	if rr.Code != 200 {
		t.Fatalf("create: %d %s", rr.Code, rr.Body.String())
	}
	var created struct {
		Key config.GWKey `json:"key"`
	}
	_ = json.Unmarshal(rr.Body.Bytes(), &created)
	back := keyRecord(t, g, created.Key.Key)
	if a := scopeOf(t, back, "m1"); a.TokenQuota != 1000 || a.ReqQuota != 0 || a.Period != "hour" {
		t.Errorf("m1 caps wrong: %+v", a)
	}
	if b := scopeOf(t, back, "m2"); b.TokenQuota != 9999 || b.ReqQuota != 7 || b.Period != "week" {
		t.Errorf("m2 caps wrong: %+v", b)
	}
}

// Sending 0 explicitly lifts that model's cap.
func TestKeyAPIUpdateZeroLiftsCap(t *testing.T) {
	g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"})
	rr := adminReq(t, g, "POST", "/api/keys",
		`{"name":"z","role":"user","models":[{"model":"m1","token_quota":1000,"period":"hour"}]}`)
	if rr.Code != 200 {
		t.Fatalf("create: %d %s", rr.Code, rr.Body.String())
	}
	var created struct {
		Key config.GWKey `json:"key"`
	}
	_ = json.Unmarshal(rr.Body.Bytes(), &created)
	secret := created.Key.Key

	rr = adminReq(t, g, "PUT", "/api/keys/"+secret,
		`{"models":[{"model":"m1","token_quota":0,"req_quota":0,"period":""}]}`)
	if rr.Code != 200 {
		t.Fatalf("lift: %d %s", rr.Code, rr.Body.String())
	}
	if back := scopeOf(t, keyRecord(t, g, secret), "m1"); back.TokenQuota != 0 || back.ReqQuota != 0 || back.Period != "" {
		t.Errorf("caps not lifted: %+v", back)
	}
}

// A misspelled period must be refused, not quietly turned into an all-time
// quota — which is the exact opposite of what the operator typed.
func TestKeyAPIRejectsBadPeriod(t *testing.T) {
	g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"})
	rr := adminReq(t, g, "POST", "/api/keys",
		`{"name":"bad","role":"user","models":[{"model":"m1","token_quota":1000,"period":"houre"}]}`)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("want 400 for a bad period, got %d %s", rr.Code, rr.Body.String())
	}
	if !strings.Contains(rr.Body.String(), "period") {
		t.Errorf("error should name the period field: %s", rr.Body.String())
	}
}

func TestKeyAPIRejectsNegativeQuota(t *testing.T) {
	g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"})
	rr := adminReq(t, g, "POST", "/api/keys",
		`{"name":"bad","role":"user","models":[{"model":"m1","token_quota":-5}]}`)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("want 400 for a negative quota, got %d %s", rr.Code, rr.Body.String())
	}
}

// The rejection must say WHICH model is over budget, so an operator looking at
// a key with a dozen scopes can tell which one to raise.
func TestKeyAPIRejectionNamesTheModel(t *testing.T) {
	g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"})
	rr := adminReq(t, g, "POST", "/api/keys",
		`{"name":"bad","role":"user","models":[{"model":"m1","req_quota":-1}]}`)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("want 400, got %d", rr.Code)
	}
	if !strings.Contains(rr.Body.String(), "m1") {
		t.Errorf("error should name the offending model: %s", rr.Body.String())
	}
}

// A non-admin key must not be able to mint keys.
func TestKeyAPIQuotaIsAdminOnly(t *testing.T) {
	g := adminGateway(t,
		config.GWKey{Key: "sk-admin", Role: "admin"},
		config.GWKey{Key: "sk-u", Role: "user", Models: []config.ModelScope{{Model: "m1", TokenQuota: 10, Period: "hour"}}},
	)
	req, _ := http.NewRequest("POST", "/api/keys",
		strings.NewReader(`{"name":"x","role":"admin","models":[{"model":"m1"}]}`))
	req.Header.Set("Authorization", "Bearer sk-u")
	req.Header.Set("Content-Type", "application/json")
	rr := httptest.NewRecorder()
	g.Handler().ServeHTTP(rr, req)
	if rr.Code != http.StatusForbidden {
		t.Fatalf("non-admin create: want 403, got %d %s", rr.Code, rr.Body.String())
	}
	// /api/v1/keys exposes the per-model caps but never a secret
	rr = adminReq(t, g, "GET", "/api/v1/keys", "")
	if rr.Code != 200 {
		t.Fatalf("GET /api/v1/keys: %d", rr.Code)
	}
	if strings.Contains(rr.Body.String(), "sk-u") {
		t.Error("/api/v1/keys leaked a key secret")
	}
	if !strings.Contains(rr.Body.String(), `"token_quota":10`) {
		t.Errorf("/api/v1/keys should expose the per-model caps: %s", rr.Body.String())
	}
}

// A user must be able to see their own budget: /api/keys/me is the only key
// view a non-admin gets, so a cap missing from it is invisible to the very
// client it constrains.
func TestKeyMeExposesOwnQuota(t *testing.T) {
	g := adminGateway(t,
		config.GWKey{Key: "sk-admin", Role: "admin"},
		config.GWKey{Key: "sk-u", Role: "user", Name: "agent",
			Models: []config.ModelScope{{Model: "m1", TokenQuota: 123456, ReqQuota: 42, Period: "week"}}},
	)
	req, _ := http.NewRequest("GET", "/api/keys/me", nil)
	req.Header.Set("Authorization", "Bearer sk-u")
	rr := httptest.NewRecorder()
	g.Handler().ServeHTTP(rr, req)
	if rr.Code != 200 {
		t.Fatalf("GET /api/keys/me: %d %s", rr.Code, rr.Body.String())
	}
	// the endpoint wraps the record: {"key": {...}}
	var wrap struct {
		Key config.GWKey `json:"key"`
	}
	if err := json.Unmarshal(rr.Body.Bytes(), &wrap); err != nil {
		t.Fatalf("decode: %v (%s)", err, rr.Body.String())
	}
	sc := scopeOf(t, wrap.Key, "m1")
	if sc.TokenQuota != 123456 || sc.ReqQuota != 42 || sc.Period != "week" {
		t.Errorf("own quota not visible to the key's owner: %+v", sc)
	}
}

var _ = fmt.Sprintf