ModelRouter/internal/gateway/key_auto_api_test.go

package gateway

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

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

// Per-key AUTO chain admin API. The behaviours pinned here are the ones that
// are easy to get wrong and invisible when they are:
//
//   - /api/keys/me/auto must be reachable by a NON-admin (it is their own
//     chain); routing it through the generic {key}/auto handler would 403 it.
//   - /api/keys/{key}/auto must 403 a non-admin.
//   - an empty PUT and a DELETE mean the same thing ("inherit global"), so the
//     WebUI cannot persist a chain that would then 503 with no slots.
//   - a PUT that resolves to zero slots is reported back, so an admin who
//     mistyped a model learns it now instead of from the user's next 503.

// keyAutoGateway builds a gateway with one admin key and one user key.
func keyAutoGateway(t *testing.T) *Gateway {
	t.Helper()
	td := t.TempDir()
	cfgPath := filepath.Join(td, "config.yaml")
	if err := os.WriteFile(cfgPath, []byte("listen: :0"), 0644); err != nil {
		t.Fatal(err)
	}
	cfg := &config.Config{
		Path:         cfgPath,
		AdapterDir:   filepath.Join(td, "adapters"),
		RuntimeFile:  filepath.Join(td, "runtime.json"),
		DefaultModel: "AUTO",
		GatewayKeys:  []string{"sk-admin"},
		Keys: []config.GWKey{
			{Key: "sk-admin", Role: "admin", Name: "admin"},
			{Key: "sk-user", Role: "user", Name: "user"},
		},
		Sources: []config.Source{
			{Name: "s1", BaseURL: "http://127.0.0.1:1/v1", Adapter: "openai",
				Models: []config.Model{{ID: "gpt-4o", Priority: 10}, {ID: "gpt-4o-mini", Priority: 5}}},
		},
	}
	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)
	if err != nil {
		t.Fatalf("gateway: %v", err)
	}
	return g
}

// doReqAs issues an authenticated request as a specific key.
func doReqAs(t *testing.T, g *Gateway, key, method, path, body string) *httptest.ResponseRecorder {
	t.Helper()
	req, _ := http.NewRequest(method, path, strings.NewReader(body))
	req.Header.Set("Authorization", "Bearer "+key)
	if body != "" {
		req.Header.Set("Content-Type", "application/json")
	}
	rr := httptest.NewRecorder()
	g.Handler().ServeHTTP(rr, req)
	return rr
}

func TestKeyAutoAPIPermissions(t *testing.T) {
	g := keyAutoGateway(t)

	// A user may read their own chain...
	if rr := doReqAs(t, g, "sk-user", "GET", "/api/keys/me/auto", ""); rr.Code != http.StatusOK {
		t.Fatalf("GET /api/keys/me/auto as user = %d, want 200 (body %s)", rr.Code, rr.Body.String())
	}
	// ...but not another key's.
	if rr := doReqAs(t, g, "sk-user", "GET", "/api/keys/sk-admin/auto", ""); rr.Code != http.StatusForbidden {
		t.Fatalf("GET other key's auto as user = %d, want 403", rr.Code)
	}
	// ...and not their own via the admin route either.
	if rr := doReqAs(t, g, "sk-user", "PUT", "/api/keys/sk-user/auto", `{"auto":[{"model":"gpt-4o"}]}`); rr.Code != http.StatusForbidden {
		t.Fatalf("user PUT own auto = %d, want 403 (only admin configures)", rr.Code)
	}
}

func TestKeyAutoAPICrudAndInherit(t *testing.T) {
	g := keyAutoGateway(t)

	// Initially inherits.
	var got struct {
		Inherits bool                `json:"inherits"`
		Auto     []config.ModelScope `json:"auto"`
	}
	rr := doReqAs(t, g, "sk-admin", "GET", "/api/keys/sk-user/auto", "")
	if rr.Code != http.StatusOK {
		t.Fatalf("GET = %d: %s", rr.Code, rr.Body.String())
	}
	if err := json.Unmarshal(rr.Body.Bytes(), &got); err != nil {
		t.Fatal(err)
	}
	if !got.Inherits || len(got.Auto) != 0 {
		t.Fatalf("fresh key must inherit global, got inherits=%v auto=%v", got.Inherits, got.Auto)
	}

	// Admin sets a chain; the response reports the resolved slot count.
	rr = doReqAs(t, g, "sk-admin", "PUT", "/api/keys/sk-user/auto",
		`{"auto":[{"model":"gpt-4o-mini","source":"s1","tier":1}]}`)
	if rr.Code != http.StatusOK {
		t.Fatalf("PUT = %d: %s", rr.Code, rr.Body.String())
	}
	var put struct {
		Inherits bool `json:"inherits"`
		Slots    int  `json:"slots"`
	}
	if err := json.Unmarshal(rr.Body.Bytes(), &put); err != nil {
		t.Fatal(err)
	}
	if put.Inherits {
		t.Fatal("after a non-empty PUT the key must no longer inherit")
	}
	if put.Slots != 1 {
		t.Fatalf("resolved slots = %d, want 1", put.Slots)
	}

	// And it reads back as an own chain.
	rr = doReqAs(t, g, "sk-admin", "GET", "/api/keys/sk-user/auto", "")
	if err := json.Unmarshal(rr.Body.Bytes(), &got); err != nil {
		t.Fatal(err)
	}
	if got.Inherits || len(got.Auto) != 1 || got.Auto[0].Model != "gpt-4o-mini" {
		t.Fatalf("own chain must read back, got inherits=%v auto=%v", got.Inherits, got.Auto)
	}

	// DELETE restores inheritance.
	if rr := doReqAs(t, g, "sk-admin", "DELETE", "/api/keys/sk-user/auto", ""); rr.Code != http.StatusOK {
		t.Fatalf("DELETE = %d: %s", rr.Code, rr.Body.String())
	}
	rr = doReqAs(t, g, "sk-admin", "GET", "/api/keys/sk-user/auto", "")
	if err := json.Unmarshal(rr.Body.Bytes(), &got); err != nil {
		t.Fatal(err)
	}
	if !got.Inherits {
		t.Fatal("after DELETE the key must inherit the global chain again")
	}
}

// An empty PUT must behave like DELETE. Otherwise a WebUI that submits an
// empty list (every slot deleted in the editor) would persist an empty chain
// and the next AUTO request would fail with no_provider.
func TestKeyAutoAPIEmptyPutClearsOverride(t *testing.T) {
	g := keyAutoGateway(t)

	if rr := doReqAs(t, g, "sk-admin", "PUT", "/api/keys/sk-user/auto",
		`{"auto":[{"model":"gpt-4o","source":"s1","tier":1}]}`); rr.Code != http.StatusOK {
		t.Fatalf("seed PUT = %d: %s", rr.Code, rr.Body.String())
	}
	rr := doReqAs(t, g, "sk-admin", "PUT", "/api/keys/sk-user/auto", `{"auto":[]}`)
	if rr.Code != http.StatusOK {
		t.Fatalf("empty PUT = %d: %s", rr.Code, rr.Body.String())
	}
	var got struct {
		Inherits bool `json:"inherits"`
	}
	if rr := doReqAs(t, g, "sk-admin", "GET", "/api/keys/sk-user/auto", ""); rr.Code != http.StatusOK {
		t.Fatal(rr.Body.String())
	}
	if err := json.Unmarshal(rr.Body.Bytes(), &got); err != nil {
		t.Fatal(err)
	}
	if !got.Inherits {
		t.Fatal("an empty PUT must clear the override, not persist an empty chain")
	}
}

// A chain naming a model that does not exist compiles to zero slots. The API
// must say so at write time instead of letting the user's next request 503.
func TestKeyAutoAPIReportsUnresolvableSlots(t *testing.T) {
	g := keyAutoGateway(t)

	rr := doReqAs(t, g, "sk-admin", "PUT", "/api/keys/sk-user/auto",
		`{"auto":[{"model":"does-not-exist","source":"s1","tier":1}]}`)
	if rr.Code != http.StatusOK {
		t.Fatalf("PUT = %d: %s", rr.Code, rr.Body.String())
	}
	var put struct {
		Inherits bool `json:"inherits"`
		Slots    int  `json:"slots"`
	}
	if err := json.Unmarshal(rr.Body.Bytes(), &put); err != nil {
		t.Fatal(err)
	}
	if put.Slots != 0 {
		t.Fatalf("unknown model must resolve to 0 slots, got %d", put.Slots)
	}
}

func TestKeyAutoAPIUnknownKey404(t *testing.T) {
	g := keyAutoGateway(t)

	if rr := doReqAs(t, g, "sk-admin", "GET", "/api/keys/sk-nope/auto", ""); rr.Code != http.StatusNotFound {
		t.Fatalf("GET unknown key auto = %d, want 404", rr.Code)
	}
	if rr := doReqAs(t, g, "sk-admin", "PUT", "/api/keys/sk-nope/auto", `{"auto":[{"model":"gpt-4o"}]}`); rr.Code != http.StatusNotFound {
		t.Fatalf("PUT unknown key auto = %d, want 404", rr.Code)
	}
}

// A bad quota must be rejected with 400, not persisted and not silently
// clamped — same rule as the model scope path.
func TestKeyAutoAPIRejectsBadQuota(t *testing.T) {
	g := keyAutoGateway(t)

	rr := doReqAs(t, g, "sk-admin", "PUT", "/api/keys/sk-user/auto",
		`{"auto":[{"model":"gpt-4o","source":"s1","tier":1,"token_quota":-1}]}`)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("negative quota PUT = %d, want 400 (body %s)", rr.Code, rr.Body.String())
	}
	// Nothing may have been persisted.
	var got struct {
		Inherits bool `json:"inherits"`
	}
	get := doReqAs(t, g, "sk-admin", "GET", "/api/keys/sk-user/auto", "")
	if get.Code != http.StatusOK {
		t.Fatal(get.Body.String())
	}
	if err := json.Unmarshal(get.Body.Bytes(), &got); err != nil {
		t.Fatal(err)
	}
	if !got.Inherits {
		t.Fatal("a rejected PUT must not persist a chain")
	}
}

// The other key endpoints must keep working: the {key}/auto route must not
// shadow /api/keys/{key} itself.
func TestKeyAutoAPIDoesNotShadowKeyRoutes(t *testing.T) {
	g := keyAutoGateway(t)

	if rr := doReqAs(t, g, "sk-admin", "GET", "/api/keys", ""); rr.Code != http.StatusOK {
		t.Fatalf("GET /api/keys = %d, want 200", rr.Code)
	}
	if rr := doReqAs(t, g, "sk-admin", "GET", "/api/keys/me", ""); rr.Code != http.StatusOK {
		t.Fatalf("GET /api/keys/me = %d, want 200", rr.Code)
	}
	if rr := doReqAs(t, g, "sk-user", "GET", "/api/keys/me", ""); rr.Code != http.StatusOK {
		t.Fatalf("GET /api/keys/me as user = %d, want 200", rr.Code)
	}
}