llmgw/internal/cli/standbyroutes_test.go

package cli

import (
	"context"
	"encoding/json"
	"io"
	"log/slog"
	"net"
	"net/http"
	"net/http/httptest"
	"slices"
	"strings"
	"sync"
	"testing"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/accounting"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control"
)

// fakeLeaderRouting is the leader's routing API: what it was told, and the
// replicas and catalog it answers with.
type fakeLeaderRouting struct {
	mu      sync.Mutex
	got     []control.Activity
	routing control.Routing
}

func (f *fakeLeaderRouting) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost || r.URL.Path != "/api/v1/routing/activity" {
		http.NotFound(w, r)
		return
	}
	var act control.Activity
	if err := json.NewDecoder(r.Body).Decode(&act); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	f.got = append(f.got, act)
	_ = json.NewEncoder(w).Encode(f.routing)
}

func (f *fakeLeaderRouting) last() control.Activity {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.got[len(f.got)-1]
}

type fixedUsage map[string]time.Time

func (u fixedUsage) LastUsed() map[string]time.Time { return u }

// sharingUsage is a router whose limits are counted with the others'.
type sharingUsage struct {
	fixedUsage
	own    accounting.Rates
	others *accounting.Rates
}

func (u *sharingUsage) Rates() accounting.Rates          { return u.own }
func (u *sharingUsage) SetOtherRates(r accounting.Rates) { u.others = &r }

// A standby reports its counts under its name and puts what the leader says
// the rest of the cluster counted in force on its router.
func TestAStandbysLimitsAreTheClusters(t *testing.T) {
	theirs := accounting.NewRates(time.Now())
	theirs.Tokens.Current = map[string]int64{"key-1": 900}
	leader := &fakeLeaderRouting{routing: control.Routing{Rates: &theirs}}
	c := newForwardClusterWith(t, leader)
	u := &sharingUsage{own: accounting.NewRates(time.Now())}
	u.own.Requests.Current = map[string]int64{"key-1": 4}
	s := newStandbyRoutes(c.standby, u, nil, nil, slog.New(slog.DiscardHandler))
	if err := s.exchange(t.Context()); err != nil {
		t.Fatal(err)
	}
	act := leader.last()
	if act.From != "ctl-b" || act.Rates == nil || act.Rates.Requests.Current["key-1"] != 4 {
		t.Errorf("reported %+v", act)
	}
	if u.others == nil || u.others.Tokens.Current["key-1"] != 900 {
		t.Errorf("the router was told %+v", u.others)
	}
}

func TestAStandbyFollowsTheLeadersReplicas(t *testing.T) {
	leader := &fakeLeaderRouting{routing: control.Routing{Catalog: "a/1",
		Replicas: []control.RoutedReplica{{ID: "llama-r0", Model: "llama", Endpoint: "http://10.0.0.5:20001"}}}}
	c := newForwardClusterWith(t, leader)
	used := time.Date(2026, 11, 1, 12, 0, 0, 0, time.UTC)
	refreshed := 0
	s := newStandbyRoutes(c.standby, fixedUsage{"llama-r0": used}, func(_ context.Context) error {
		refreshed++
		return nil
	}, nil, slog.New(slog.DiscardHandler))

	if _, ok := s.Ready(); ok {
		t.Fatal("ready before hearing from the leader")
	}
	if err := s.exchange(t.Context()); err != nil {
		t.Fatal(err)
	}
	if reps, ok := s.Ready(); !ok || len(reps) != 1 || reps[0].ID != "llama-r0" {
		t.Fatalf("replicas %+v, %v", reps, ok)
	}
	if got := leader.last().LastUsed["llama-r0"]; !got.Equal(used) {
		t.Errorf("the leader was told the replica was last used %v", got)
	}
	if refreshed != 1 {
		t.Errorf("the catalog was read %d times on first hearing from the leader", refreshed)
	}

	// Only what is new is reported, and the catalog is read again only when
	// the leader's changed.
	s.wanted("flux")
	if err := s.exchange(t.Context()); err != nil {
		t.Fatal(err)
	}
	if act := leader.last(); len(act.LastUsed) != 0 || !slices.Equal(act.Want, []string{"flux"}) {
		t.Errorf("second report %+v", act)
	}
	if refreshed != 1 {
		t.Errorf("an unchanged catalog was read again")
	}
	leader.mu.Lock()
	leader.routing.Catalog = "a/2"
	leader.mu.Unlock()
	if err := s.exchange(t.Context()); err != nil {
		t.Fatal(err)
	}
	if act := leader.last(); len(act.Want) != 0 {
		t.Errorf("an on-demand model reported twice: %+v", act)
	}
	if refreshed != 2 {
		t.Errorf("the changed catalog was not read: %d", refreshed)
	}

	// Unheard for longer than it is trusted, it routes nothing, and requests
	// go to the leader again.
	c.standby.routes = s
	if !c.standby.answersHere("/v1/chat/completions") {
		t.Error("a standby with replicas passed the API on")
	}
	s.now = func() time.Time { return time.Now().Add(standbyTrust + time.Second) }
	if _, ok := s.Ready(); ok || c.standby.answersHere("/v1/chat/completions") {
		t.Error("a standby the leader has not answered for a minute still routes")
	}
}

// The client API stays on a standby with replicas; the portal, its admin
// pages and /mcp still go to the leader.
func TestAStandbyAnswersOnlyTheAPIItself(t *testing.T) {
	leader := &fakeLeaderRouting{routing: control.Routing{
		Replicas: []control.RoutedReplica{{ID: "llama-r0", Model: "llama", Endpoint: "http://10.0.0.5:20001"}}}}
	c := newForwardClusterWith(t, leader)
	s := newStandbyRoutes(c.standby, nil, nil, nil, slog.New(slog.DiscardHandler))
	if err := s.exchange(t.Context()); err != nil {
		t.Fatal(err)
	}
	c.standby.routes = s

	port := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte("leader " + r.URL.Path))
	}))
	defer port.Close()
	_, p, _ := net.SplitHostPort(strings.TrimPrefix(port.URL, "http://"))
	local := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = w.Write([]byte("standby " + r.URL.Path)) })
	srv := httptest.NewServer(c.standby.public(local, p))
	defer srv.Close()

	for path, want := range map[string]string{
		"/v1/chat/completions": "standby", "/v1/models": "standby",
		"/admin/replicas": "leader", "/mcp": "leader", "/": "leader",
	} {
		resp, err := http.Get(srv.URL + path) //nolint:noctx // a test
		if err != nil {
			t.Fatal(err)
		}
		b, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		if got, _, _ := strings.Cut(string(b), " "); got != want {
			t.Errorf("%s answered by %q, want the %s", path, b, want)
		}
	}
}