ModelRouter/internal/core/key_auto_test.go

package core

import (
	"path/filepath"
	"testing"

	"llmsproxy/internal/config"
	"llmsproxy/internal/scheduler"
)

// Per-key AUTO chains let an admin give one user their own scheduling chain
// while every other key keeps the global one. These tests pin the three
// properties that make that safe:
//
//  1. a key WITH its own chain gets it;
//  2. a key WITHOUT one inherits the global chain (backward compatible);
//  3. the override is actually persisted, so it survives a restart.
//
// A test that only checked (1) would pass even if AutoChainFor ignored the key
// and always returned the global chain whenever the two happened to be equal,
// so (2) uses deliberately *different* chains.

func perKeyTestConfig(t *testing.T) *config.Config {
	t.Helper()
	dir := t.TempDir()
	return &config.Config{
		Path:         filepath.Join(dir, "config.yaml"),
		AdapterDir:   filepath.Join(dir, "adapters"),
		RuntimeFile:  filepath.Join(dir, "runtime.json"),
		Listen:       "127.0.0.1:0",
		DefaultModel: "AUTO",
		GatewayKeys:  []string{"sk-gw-admin"},
		Sources: []config.Source{
			{Name: "s1", BaseURL: "http://127.0.0.1:1/v1", Adapter: "openai",
				Models: []config.Model{{ID: "gpt-4o"}, {ID: "gpt-4o-mini"}}},
		},
		// Global chain points at gpt-4o.
		Auto: []config.ModelScope{{Model: "gpt-4o", Source: "s1", Tier: 1}},
		Keys: []config.GWKey{
			{Key: "sk-gw-admin", Role: "admin"},
			{Key: "sk-gw-inherits", Role: "user", Name: "inherits"},
			// This key's own chain points at a DIFFERENT model, so returning
			// the global chain by mistake is detectable.
			{Key: "sk-gw-own", Role: "user", Name: "own",
				Auto: []config.ModelScope{{Model: "gpt-4o-mini", Source: "s1", Tier: 1}}},
		},
	}
}

func firstSlotModel(t *testing.T, c *Core, key string) string {
	t.Helper()
	ch, _ := c.AutoChainFor(key)
	if ch == nil || len(ch.Tiers) == 0 {
		t.Fatalf("key %s: no chain tiers", key)
	}
	slots := ch.Tiers[0].Slots
	if len(slots) == 0 {
		t.Fatalf("key %s: tier 1 has no slots", key)
	}
	return slots[0].Model
}

// A key that declares its own chain must be scheduled by it, not the global.
func TestAutoChainForUsesKeyOwnChain(t *testing.T) {
	c := newTestCore(t, perKeyTestConfig(t))

	if got := firstSlotModel(t, c, "sk-gw-own"); got != "gpt-4o-mini" {
		t.Fatalf("key with own chain must use it, got %q (want gpt-4o-mini)", got)
	}
	// Sanity: the global chain really is different, so the assertion above
	// cannot pass by accident.
	if got := firstSlotModel(t, c, "sk-gw-admin"); got != "gpt-4o" {
		t.Fatalf("admin must use global chain, got %q (want gpt-4o)", got)
	}
}

// A key with no own chain inherits the global one. This is what keeps every
// pre-existing deployment working with no config change.
func TestAutoChainForInheritsGlobalWhenUnset(t *testing.T) {
	c := newTestCore(t, perKeyTestConfig(t))

	if got := firstSlotModel(t, c, "sk-gw-inherits"); got != "gpt-4o" {
		t.Fatalf("key without own chain must inherit global, got %q (want gpt-4o)", got)
	}
}

// A key that is not in the key table at all (e.g. a seed key) must also
// inherit the global chain rather than erroring or returning nil.
func TestAutoChainForUnknownKeyInheritsGlobal(t *testing.T) {
	c := newTestCore(t, perKeyTestConfig(t))

	ch, ok := c.AutoChainFor("sk-gw-never-configured")
	if !ok || ch == nil || len(ch.Tiers) == 0 {
		t.Fatalf("unknown key must fall back to global chain, got ok=%v chain=%v", ok, ch)
	}
}

// The per-key chain must survive a save/reload round trip: an admin configuring
// a chain in the WebUI and restarting the gateway must not silently lose it.
func TestSaveKeyAutoPersistsAndApplies(t *testing.T) {
	cfg := perKeyTestConfig(t)
	c := newTestCore(t, cfg)

	if err := c.SaveKeyAuto("sk-gw-inherits", []config.ModelScope{
		{Model: "gpt-4o-mini", Source: "s1", Tier: 1},
	}); err != nil {
		t.Fatalf("SaveKeyAuto: %v", err)
	}
	if got := firstSlotModel(t, c, "sk-gw-inherits"); got != "gpt-4o-mini" {
		t.Fatalf("saved chain must apply immediately, got %q", got)
	}

	// Reload from disk exactly like a restart does.
	reloaded, err := config.Load(cfg.Path)
	if err != nil {
		t.Fatalf("reload config: %v", err)
	}
	c2 := newTestCore(t, reloaded)
	if got := firstSlotModel(t, c2, "sk-gw-inherits"); got != "gpt-4o-mini" {
		t.Fatalf("per-key chain must survive restart, got %q", got)
	}
	// The other key must be unaffected by its neighbour's override.
	if got := firstSlotModel(t, c2, "sk-gw-admin"); got != "gpt-4o" {
		t.Fatalf("override must not leak to other keys, got %q", got)
	}
}

// Clearing the override must return the key to the global chain, both in
// memory and after a restart. This is what the WebUI's "use global" toggle
// does, so a stale shadowing entry would silently pin the user to their old
// chain forever.
func TestSaveKeyAutoClearRestoresGlobal(t *testing.T) {
	cfg := perKeyTestConfig(t)
	c := newTestCore(t, cfg)

	if err := c.SaveKeyAuto("sk-gw-own", nil); err != nil {
		t.Fatalf("clear: %v", err)
	}
	if got := firstSlotModel(t, c, "sk-gw-own"); got != "gpt-4o" {
		t.Fatalf("cleared key must inherit global again, got %q", got)
	}

	reloaded, err := config.Load(cfg.Path)
	if err != nil {
		t.Fatalf("reload: %v", err)
	}
	c2 := newTestCore(t, reloaded)
	if got := firstSlotModel(t, c2, "sk-gw-own"); got != "gpt-4o" {
		t.Fatalf("clear must survive restart, got %q", got)
	}
}

// Saving an unknown key must fail loudly rather than creating a shadow entry
// that no request can ever match.
func TestSaveKeyAutoUnknownKeyFails(t *testing.T) {
	c := newTestCore(t, perKeyTestConfig(t))

	if err := c.SaveKeyAuto("sk-gw-nope", []config.ModelScope{{Model: "gpt-4o", Source: "s1"}}); err == nil {
		t.Fatal("SaveKeyAuto on unknown key must return an error")
	}
}

// Per-key quota validation must run, exactly like the model scope path: an
// override is another place an operator can write a bad budget.
func TestSaveKeyAutoRejectsBadQuota(t *testing.T) {
	c := newTestCore(t, perKeyTestConfig(t))

	err := c.SaveKeyAuto("sk-gw-inherits", []config.ModelScope{
		{Model: "gpt-4o", Source: "s1", TokenQuota: -5},
	})
	if err == nil {
		t.Fatal("negative token quota in a per-key auto chain must be rejected")
	}
}

// Source changes must recompile per-key chains. A key whose chain names a
// model that does not exist yet compiles to an empty chain (BuildChain drops
// slots it cannot resolve — sending the request into a black hole would be
// worse), and must start working as soon as the source appears.
func TestKeyAutoChainsRebuildOnSourceChange(t *testing.T) {
	cfg := perKeyTestConfig(t)
	cfg.Keys = append(cfg.Keys, config.GWKey{
		Key: "sk-gw-late", Role: "user", Name: "late",
		Auto: []config.ModelScope{{Model: "new-model", Source: "s2", Tier: 1}},
	})
	c := newTestCore(t, cfg)

	// s2 does not exist yet, so the chain has no usable slot.
	if ch, _ := c.AutoChainFor("sk-gw-late"); chainSlotsOf(ch) != 0 {
		t.Fatalf("chain naming an unknown model must compile empty, got %d slots",
			chainSlotsOf(ch))
	}
	// The other keys are unaffected by the broken one.
	if got := firstSlotModel(t, c, "sk-gw-admin"); got != "gpt-4o" {
		t.Fatalf("broken per-key chain must not affect others, got %q", got)
	}

	// Add the source; the per-key chain must pick it up without a restart.
	cfg.Sources = append(cfg.Sources, config.Source{
		Name: "s2", BaseURL: "http://127.0.0.1:2/v1", Adapter: "openai",
		Models: []config.Model{{ID: "new-model"}},
	})
	if err := c.Reload(); err != nil {
		t.Fatalf("Reload: %v", err)
	}
	if got := firstSlotModel(t, c, "sk-gw-late"); got != "new-model" {
		t.Fatalf("per-key chain must resolve after the source is added, got %q", got)
	}
}

// chainSlotsOf counts usable slots, tolerating a nil chain.
func chainSlotsOf(ch *scheduler.Chain) int {
	if ch == nil {
		return 0
	}
	n := 0
	for _, tn := range ch.Tiers {
		n += len(tn.Slots)
	}
	return n
}