ModelRouter/internal/scheduler/scheduler_test.go

package scheduler

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"llmsproxy/internal/types"
)

// fakeProvider is an in-memory Provider used to exercise chain scheduling
// deterministically without a Lua runtime.
type fakeProvider struct {
	name      string
	model     string
	pref      atomic.Int64
	available atomic.Bool
	busy      atomic.Bool
	fail      atomic.Bool
	chatHits  atomic.Int64
}

func fakeProv(name, model string) *fakeProvider {
	f := &fakeProvider{name: name, model: model}
	f.available.Store(true)
	return f
}

func (f *fakeProvider) Name() string { return f.name }

func (f *fakeProvider) ModelFor(reqModel string) string { return f.model }

func (f *fakeProvider) ModelAvailable(model string) bool { return f.available.Load() }

func (f *fakeProvider) Pref(model string) int64 { return f.pref.Load() }

func (f *fakeProvider) Chat(ctx context.Context, req *types.ChatRequest) (*types.UnifiedResponse, error) {
	f.chatHits.Add(1)
	if f.busy.Load() {
		return nil, types.ErrBusy
	}
	if f.fail.Load() {
		return nil, fmt.Errorf("upstream error")
	}
	return &types.UnifiedResponse{Content: f.name, FinishReason: "stop", TokenUsage: types.TokenUsage{Prompt: 1, Completion: 1, Total: 2}}, nil
}

func (f *fakeProvider) ChatStream(ctx context.Context, req *types.ChatRequest) (<-chan types.UnifiedChunk, error) {
	if f.busy.Load() {
		return nil, types.ErrBusy
	}
	if f.fail.Load() {
		return nil, fmt.Errorf("upstream error")
	}
	ch := make(chan types.UnifiedChunk, 2)
	ch <- types.UnifiedChunk{Content: f.name}
	ch <- types.UnifiedChunk{Done: true}
	close(ch)
	return ch, nil
}

func (f *fakeProvider) Image(ctx context.Context, req *types.ImageGenRequest) (*types.UnifiedResponse, error) {
	return nil, errors.New("no image")
}

// TestDirectSkipsCooledCandidate: direct paths share the AUTO-chain rule that
// cooldown is the only hard skip — a cooling candidate must never be hit, and
// a fully cooling set must fail without touching upstream.
func TestDirectSkipsCooledCandidate(t *testing.T) {
	hot := fakeProv("hot", "m")
	cold := fakeProv("cold", "m")
	cold.available.Store(false)
	s := New(2)
	req := &types.ChatRequest{
		Model:    "m",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
	}
	resp, src, _, err := s.Chat(context.Background(), []Provider{cold, hot}, req)
	if err != nil || src != "hot" {
		t.Fatalf("want hot to serve, got src=%q err=%v", src, err)
	}
	if resp == nil || resp.Content != "hot" {
		t.Fatalf("bad response: %#v", resp)
	}
	if cold.chatHits.Load() != 0 {
		t.Fatalf("cooled candidate must not be hit, got %d hits", cold.chatHits.Load())
	}
	// all candidates cooled: direct path reports it instead of hitting upstream
	hot.available.Store(false)
	hits := hot.chatHits.Load()
	_, _, _, err = s.Chat(context.Background(), []Provider{cold, hot}, req)
	if err == nil || !strings.Contains(err.Error(), "cooling down") {
		t.Fatalf("want cooling-down error, got %v", err)
	}
	if hot.chatHits.Load() != hits || cold.chatHits.Load() != 0 {
		t.Fatalf("cooled candidates must not be hit, got hot=%d cold=%d", hot.chatHits.Load(), cold.chatHits.Load())
	}
}

// lookup resolves (model, source) -> provider for chain builders in tests.
type lookup func(m, s string) Provider

func bySource(ps ...*fakeProvider) lookup {
	return func(m, s string) Provider {
		for _, p := range ps {
			if p.name == s {
				return p
			}
		}
		return nil
	}
}

func chatReq() *types.ChatRequest {
	return &types.ChatRequest{Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("hi")}}}
}

func TestBuildChainTierOrdering(t *testing.T) {
	a, b, c, d := fakeProv("s1", "a"), fakeProv("s2", "b"), fakeProv("s3", "c"), fakeProv("s4", "d")
	ch := BuildChain([]Rule{
		{Tier: 1, Model: "a", Source: "s1"},
		{Tier: 3, Model: "c", Source: "s3"},
		{Tier: 2, Model: "b", Source: "s2"},
		{Tier: 1, Model: "d", Source: "s4"},
		{Tier: 9, Model: "gone", Source: "missing"},
	}, bySource(a, b, c, d))
	if len(ch.Tiers) != 3 {
		t.Fatalf("tiers = %d, want 3", len(ch.Tiers))
	}
	if ch.Tiers[0].Tier != 1 || ch.Tiers[1].Tier != 2 || ch.Tiers[2].Tier != 3 {
		t.Fatalf("tier order = %d,%d,%d, want 1,2,3", ch.Tiers[0].Tier, ch.Tiers[1].Tier, ch.Tiers[2].Tier)
	}
	// same-tier slots keep configured order, unresolvable rule is dropped
	if len(ch.Tiers[0].Slots) != 2 || ch.Tiers[0].Slots[0].Model != "a" || ch.Tiers[0].Slots[1].Model != "d" {
		t.Fatalf("tier 1 slots = %+v", ch.Tiers[0].Slots)
	}
}

func TestChainRoundRobin(t *testing.T) {
	a, b := fakeProv("s1", "a"), fakeProv("s2", "b")
	ch := BuildChain([]Rule{
		{Tier: 0, Model: "a", Source: "s1"},
		{Tier: 0, Model: "b", Source: "s2"},
	}, bySource(a, b))
	s := New(0)
	var got []string
	for i := 0; i < 4; i++ {
		_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
		if err != nil {
			t.Fatalf("iter %d: %v", i, err)
		}
		got = append(got, src)
	}
	want := []string{"s1", "s2", "s1", "s2"}
	for i := range want {
		if got[i] != want[i] {
			t.Fatalf("rr order = %v, want %v", got, want)
		}
	}
}

func TestChainPreferenceSinksButStaysReachable(t *testing.T) {
	neg := fakeProv("neg", "n")
	good := fakeProv("good", "g")
	neg.pref.Store(-5)
	good.fail.Store(true)
	ch := BuildChain([]Rule{
		{Tier: 0, Model: "n", Source: "neg"},
		{Tier: 0, Model: "g", Source: "good"},
	}, bySource(neg, good))
	s := New(0)
	resp, src, model, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
	if err != nil {
		t.Fatalf("chain: %v", err)
	}
	// Round-robin starts at index 0 (neg). neg is available (pref=-5 > -20)
	// and succeeds. good is never tried because neg already won.
	if src != "neg" || model != "n" || resp.Content != "neg" {
		t.Fatalf("served src=%q model=%q content=%q", src, model, resp.Content)
	}
	if good.chatHits.Load() != 0 {
		t.Fatal("neg was tried first and succeeded; good must not be attempted")
	}
	// Second request: cursor advances. neg wins again (good hard-fails).
	resp2, src2, _, err2 := s.ChainChat(context.Background(), ch, chatReq(), nil)
	if err2 != nil {
		t.Fatalf("second chain: %v", err2)
	}
	if src2 != "neg" || resp2.Content != "neg" {
		t.Fatalf("second req src=%q content=%q", src2, resp2.Content)
	}
}

func TestChainBusySkipsWithoutPenalty(t *testing.T) {
	a, b := fakeProv("s1", "a"), fakeProv("s2", "b")
	a.busy.Store(true)
	ch := BuildChain([]Rule{
		{Tier: 0, Model: "a", Source: "s1"},
		{Tier: 0, Model: "b", Source: "s2"},
	}, bySource(a, b))
	s := New(0)
	_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
	if err != nil {
		t.Fatalf("chain: %v", err)
	}
	if src != "s2" {
		t.Fatalf("src = %q, want s2", src)
	}
	if a.chatHits.Load() == 0 {
		t.Fatal("busy slot must have been attempted")
	}
}

func TestChainAllBusyBoundedWaitThenNextTier(t *testing.T) {
	oldWait, oldPoll := busyWait, busyPoll
	busyWait, busyPoll = 60*time.Millisecond, 10*time.Millisecond
	t.Cleanup(func() { busyWait, busyPoll = oldWait, oldPoll })
	a, b, c := fakeProv("s1", "a"), fakeProv("s2", "b"), fakeProv("s3", "c")
	a.busy.Store(true)
	b.busy.Store(true)
	ch := BuildChain([]Rule{
		{Tier: 1, Model: "a", Source: "s1"},
		{Tier: 1, Model: "b", Source: "s2"},
		{Tier: 2, Model: "c", Source: "s3"},
	}, bySource(a, b, c))
	s := New(0)
	t0 := time.Now()
	_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
	el := time.Since(t0)
	if err != nil {
		t.Fatalf("chain: %v", err)
	}
	if src != "s3" {
		t.Fatalf("src = %q, want s3 (downgrade after bounded wait)", src)
	}
	if el > time.Second {
		t.Fatalf("busy wait not bounded: %v", el)
	}
}

func TestChainQuotaExhausted(t *testing.T) {
	a, b := fakeProv("s1", "a"), fakeProv("s2", "b")
	ch := BuildChain([]Rule{
		{Tier: 0, Model: "a", Source: "s1", Quota: 100, Period: "hour"},
		{Tier: 0, Model: "b", Source: "s2"},
	}, bySource(a, b))
	s := New(0)
	exhausted := func(sl *Slot) bool { return sl.Source == "s1" && sl.Quota > 0 }
	_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), exhausted)
	if err != nil {
		t.Fatalf("chain: %v", err)
	}
	if src != "s2" {
		t.Fatalf("src = %q, want s2 (quota slot dropped)", src)
	}
	if a.chatHits.Load() != 0 {
		t.Fatal("quota-exhausted slot must not be called")
	}
}

func TestChainErrSummary(t *testing.T) {
	a, b, c := fakeProv("s1", "a"), fakeProv("s2", "b"), fakeProv("s3", "c")
	a.fail.Store(true)
	b.fail.Store(true)
	c.available.Store(false) // whole tier 1 cooling
	ch := BuildChain([]Rule{
		{Tier: 0, Model: "a", Source: "s1"},
		{Tier: 0, Model: "b", Source: "s2"},
		{Tier: 1, Model: "c", Source: "s3"},
	}, bySource(a, b, c))
	s := New(0)
	_, _, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
	var ce *ChainErr
	if !errors.As(err, &ce) {
		t.Fatalf("err = %v, want *ChainErr", err)
	}
	if len(ce.Tiers) != 2 || ce.Tiers[0].Source != "s1" || ce.Tiers[1].Source != "s2" {
		t.Fatalf("tiers = %+v", ce.Tiers)
	}
	if len(ce.Skipped) != 1 {
		t.Fatalf("skipped = %+v", ce.Skipped)
	}
	msg := ce.Error()
	if !strings.Contains(msg, "s1") || !strings.Contains(msg, "s2") || !strings.Contains(msg, "tier 1: no schedulable slot") {
		t.Fatalf("summary = %q", msg)
	}
}

func TestChainStreamFallsBackBeforeFirstChunk(t *testing.T) {
	a, b := fakeProv("s1", "a"), fakeProv("s2", "b")
	a.fail.Store(true)
	ch := BuildChain([]Rule{
		{Tier: 0, Model: "a", Source: "s1"},
		{Tier: 0, Model: "b", Source: "s2"},
	}, bySource(a, b))
	s := New(0)
	chunks, src, model, err := s.ChainChatStream(context.Background(), ch, chatReq(), nil)
	if err != nil {
		t.Fatalf("chain stream: %v", err)
	}
	if src != "s2" || model != "b" {
		t.Fatalf("src=%q model=%q", src, model)
	}
	var text string
	for ck := range chunks {
		text += ck.Content
	}
	if text != "s2" {
		t.Fatalf("text = %q", text)
	}
}

func TestChainResetCooldownAfterSwap(t *testing.T) {
	a := fakeProv("s1", "a")
	ch := BuildChain([]Rule{{Tier: 0, Model: "a", Source: "s1"}}, bySource(a))
	// a freshly built chain must schedule from index 0 (cursor starts at -1)
	if base := ch.Tiers[0].NextStart(); base != 0 {
		t.Fatalf("first start = %d, want 0", base)
	}
}