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 TestKeyAPIStoresQuota(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","token_quota":50000,"req_quota":200,"period":"nhour","hours":6,"models":[{"model":"m1"}]}`)
	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)
	}
	if created.Key.TokenQuota != 50000 || created.Key.ReqQuota != 200 ||
		created.Key.Period != "nhour" || created.Key.Hours != 6 {
		t.Fatalf("created key did not carry the caps: %+v", created.Key)
	}
	// and it must survive a read-back (persisted, not just echoed)
	back := keyRecord(t, g, created.Key.Key)
	if back.TokenQuota != 50000 || back.Period != "nhour" || back.Hours != 6 {
		t.Errorf("read-back lost the caps: %+v", back)
	}
}

// Editing only the model scope must not silently clear a key's budget: the
// caps are pointers precisely so "absent" is not "zero".
func TestKeyAPIUpdateKeepsQuotaWhenOmitted(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","token_quota":50000,"period":"day-typo-free","models":[{"model":"m1"}]}`)
	rr = adminReq(t, g, "POST", "/api/keys", `{"name":"y","role":"user","token_quota":50000,"period":"hour","models":[{"model":"m1"}]}`)
	if rr.Code != 200 {
		t.Fatalf("setup create: %d %s", rr.Code, rr.Body.String())
	}
	var created struct {
		Key config.GWKey `json:"key"`
	}
	_ = json.Unmarshal(rr.Body.Bytes(), &created)

	// a scope-only edit
	rr = adminReq(t, g, "PUT", "/api/keys/"+created.Key.Key,
		`{"name":"agent-y","models":[{"model":"m1"},{"model":"m2"}]}`)
	if rr.Code != 200 {
		t.Fatalf("update: %d %s", rr.Code, rr.Body.String())
	}
	back := keyRecord(t, g, created.Key.Key)
	if back.TokenQuota != 50000 {
		t.Errorf("token_quota was cleared by a scope-only edit: %d", back.TokenQuota)
	}
	if back.Period != "hour" {
		t.Errorf("period was cleared by a scope-only edit: %q", back.Period)
	}
	if len(back.Models) != 2 {
		t.Errorf("scope edit did not apply: %+v", back.Models)
	}
}

// Sending 0 explicitly must lift the cap, not be treated as "absent".
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","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)

	rr = adminReq(t, g, "PUT", "/api/keys/"+created.Key.Key, `{"token_quota":0}`)
	if rr.Code != 200 {
		t.Fatalf("lift: %d %s", rr.Code, rr.Body.String())
	}
	if back := keyRecord(t, g, created.Key.Key); back.TokenQuota != 0 {
		t.Errorf("token_quota = %d, want 0 (cap lifted)", back.TokenQuota)
	}
}

// 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","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","token_quota":-5}`)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("want 400 for a negative quota, got %d %s", rr.Code, rr.Body.String())
	}
}

// A non-admin key must not be able to set or read another key's budget.
func TestKeyAPIQuotaIsAdminOnly(t *testing.T) {
	g := adminGateway(t,
		config.GWKey{Key: "sk-admin", Role: "admin"},
		config.GWKey{Key: "sk-u", Role: "user", TokenQuota: 10, Period: "hour"},
	)
	req, _ := http.NewRequest("POST", "/api/keys", strings.NewReader(`{"name":"x","role":"admin","token_quota":0}`))
	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 echoes the caps but never the 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 cap: %s", rr.Body.String())
	}
}

// The AUTO scope entry must honour its reset window: usage that aged out of
// the window must not count against a per-key cap.
func TestAutoScopeQuotaHonoursWindow(t *testing.T) {
	g, _ := quotaGateway(t,
		config.GWKey{Key: "sk-a", Role: "user", Models: []config.ModelScope{
			{Model: "AUTO", TokenQuota: 1000, Period: "hour"},
		}},
		config.GWKey{Key: "sk-b", Role: "user"},
	)
	ctx := quotaCtx(t, g, "sk-a")
	sc := config.ModelScope{Model: "AUTO", TokenQuota: 1000, Period: "hour"}

	// aged-out usage: 2 days old, 5M tokens — must be invisible to a 1h window
	g.stats.Record(Req{Time: nowMSOffset(-48 * 3600 * 1000), Key: keyID("sk-a"),
		Model: "m1", Source: "up", Prompt: 2500000, Compl: 2500000, OK: true, Status: 200})
	if used := g.scopeTokens(ctx, sc); used != 0 {
		t.Fatalf("AUTO scope saw %d tokens outside its 1h window; the period is being ignored", used)
	}

	// in-window usage counts
	g.stats.Record(Req{Time: nowMSOffset(0), Key: keyID("sk-a"),
		Model: "m1", Source: "up", Prompt: 400, Compl: 400, OK: true, Status: 200})
	if used := g.scopeTokens(ctx, sc); used != 800 {
		t.Fatalf("AUTO scope used = %d, want 800", used)
	}
}

var _ = fmt.Sprintf