Maven/internal/phraser/swap_test.go

package phraser

import (
	"context"
	"errors"
	"net/http"
	"net/http/httptest"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

// fakeModel — a stand-in llama-server. It answers /v1/models with its own name
// and /v1/chat/completions with a phrasing-contract reply that names itself, so
// a test can tell WHICH model answered a turn — the property the swap is about.
type fakeModel struct {
	srv    *httptest.Server
	name   string
	closed atomic.Bool
}

func newFakeModel(t *testing.T, name string) *fakeModel {
	t.Helper()
	f := &fakeModel{name: name}
	mux := http.NewServeMux()
	mux.HandleFunc("/v1/models", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`{"data":[{"id":"/models/` + name + `.gguf"}]}`))
	})
	mux.HandleFunc("/v1/chat/completions", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`{"choices":[{"message":{"content":"{\"response\":\"` + name + `\",\"mood\":\"neutral\"}"}}]}`))
	})
	f.srv = httptest.NewServer(mux)
	t.Cleanup(f.srv.Close)
	return f
}

func (f *fakeModel) BaseURL() string { return f.srv.URL }
func (f *fakeModel) Close() error    { f.closed.Store(true); return nil }

// fakeFleet is the injected launcher: it hands out a prepared fakeModel per
// model path, and refuses paths the test did not prepare (that is what a bad
// gguf looks like from here). It also asserts the invariant that matters on a
// laptop iGPU: never two servers alive at the same time.
type fakeFleet struct {
	mu      sync.Mutex
	models  map[string]string // model path → fake name
	live    int
	maxLive int
	launch  int
}

func (fl *fakeFleet) launcher(t *testing.T) func(context.Context, Config) (backend, error) {
	return func(ctx context.Context, cfg Config) (backend, error) {
		fl.mu.Lock()
		name, ok := fl.models[cfg.ModelPath]
		fl.launch++
		if !ok {
			fl.mu.Unlock()
			return nil, errors.New("no such model file: " + cfg.ModelPath)
		}
		fl.live++
		if fl.live > fl.maxLive {
			fl.maxLive = fl.live
		}
		fl.mu.Unlock()
		f := newFakeModel(t, name)
		return &fleetBackend{fleet: fl, model: f}, nil
	}
}

type fleetBackend struct {
	fleet *fakeFleet
	model *fakeModel
	once  sync.Once
}

func (b *fleetBackend) BaseURL() string { return b.model.BaseURL() }
func (b *fleetBackend) Close() error {
	b.once.Do(func() {
		b.fleet.mu.Lock()
		b.fleet.live--
		b.fleet.mu.Unlock()
	})
	return b.model.Close()
}

// newSwapPhraser builds an LLMPhraser with an injected launcher, so the swap
// path is exercised without a gguf or a GPU.
func newSwapPhraser(t *testing.T, fl *fakeFleet, modelPath string) *LLMPhraser {
	t.Helper()
	cfg := DefaultConfig(modelPath)
	cfg.Timeout = 5 * time.Second
	p := &LLMPhraser{
		cfg:      cfg,
		client:   &http.Client{Timeout: cfg.Timeout},
		spawnCtx: context.Background(),
		cancel:   func() {},
		launch:   fl.launcher(t),
		probe:    defaultProbe,
		live:     liveModel{ModelPath: modelPath, NGpuLayers: cfg.NGpuLayers, NCtx: cfg.NCtx},
	}
	be, err := p.launch(p.spawnCtx, cfg)
	if err != nil {
		t.Fatalf("initial launch: %v", err)
	}
	p.be = be
	t.Cleanup(func() { p.Close() })
	return p
}

func TestSwap_LoadsNewModelAndRepointsHolders(t *testing.T) {
	fl := &fakeFleet{models: map[string]string{"/m/old.gguf": "old", "/m/new.gguf": "new"}}
	p := newSwapPhraser(t, fl, "/m/old.gguf")

	// A holder of the base URL (the LLM router's client, in the daemon).
	var seen []string
	p.OnSwap(func(base string) { seen = append(seen, base) })

	before, err := p.PhraseChat(context.Background(), "привет", nil)
	if err != nil || before != "old" {
		t.Fatalf("before swap: %q, %v; want the old model to answer", before, err)
	}

	res, err := p.Swap(context.Background(), SwapSpec{ModelPath: "/m/new.gguf"})
	if err != nil {
		t.Fatalf("Swap: %v", err)
	}
	if res.Model != "new" {
		t.Errorf("res.Model = %q; want the identity the NEW server reported (%q)", res.Model, "new")
	}
	if res.RolledBack {
		t.Errorf("res.RolledBack = true on a successful swap")
	}
	after, err := p.PhraseChat(context.Background(), "привет", nil)
	if err != nil || after != "new" {
		t.Fatalf("after swap: %q, %v; want the new model to answer", after, err)
	}
	if path, _, _ := p.LiveModel(); path != "/m/new.gguf" {
		t.Errorf("LiveModel = %q; want /m/new.gguf", path)
	}
	if len(seen) != 1 || seen[0] != p.BaseURL() {
		t.Errorf("observers saw %v; want exactly one call with the new base %q", seen, p.BaseURL())
	}
	if fl.maxLive > 1 {
		t.Errorf("%d servers were alive at once; the iGPU only fits one model", fl.maxLive)
	}
}

func TestSwap_FailedLoadRollsBackToTheWorkingModel(t *testing.T) {
	fl := &fakeFleet{models: map[string]string{"/m/old.gguf": "old"}}
	p := newSwapPhraser(t, fl, "/m/old.gguf")

	res, err := p.Swap(context.Background(), SwapSpec{ModelPath: "/m/broken.gguf"})
	if err == nil {
		t.Fatal("Swap to a model that will not load returned nil error")
	}
	if !res.RolledBack {
		t.Errorf("res.RolledBack = false; a failed swap must say it rolled back")
	}
	if res.Model != "old" {
		t.Errorf("res.Model = %q; want the old model back", res.Model)
	}
	// The point of the rollback: turns keep working.
	got, err := p.PhraseChat(context.Background(), "привет", nil)
	if err != nil || got != "old" {
		t.Fatalf("after rollback: %q, %v; want the old model serving again", got, err)
	}
	if path, _, _ := p.LiveModel(); path != "/m/old.gguf" {
		t.Errorf("LiveModel = %q; want the old model", path)
	}
	if fl.maxLive > 1 {
		t.Errorf("%d servers alive at once during a rollback", fl.maxLive)
	}
}

func TestSwap_ProbeFailureIsTreatedAsAFailedLoad(t *testing.T) {
	// A server that starts but will not say what it loaded must never be
	// published — we would be serving turns from a backend we cannot confirm.
	fl := &fakeFleet{models: map[string]string{"/m/old.gguf": "old", "/m/mute.gguf": "mute"}}
	p := newSwapPhraser(t, fl, "/m/old.gguf")
	// Fail the probe once — for the newly launched server — and let the
	// rollback's probe through.
	calls := 0
	p.probe = func(ctx context.Context, base string) (string, error) {
		calls++
		if calls == 1 {
			return "", errors.New("no answer from the new server")
		}
		return defaultProbe(ctx, base)
	}

	_, err := p.Swap(context.Background(), SwapSpec{ModelPath: "/m/mute.gguf"})
	if err == nil {
		t.Fatal("Swap published a server that failed its probe")
	}
	if path, _, _ := p.LiveModel(); path != "/m/old.gguf" {
		t.Errorf("LiveModel = %q; want the old model after a failed probe", path)
	}
}

func TestSwap_RollbackFailureLeavesNoBackendAndDegrades(t *testing.T) {
	fl := &fakeFleet{models: map[string]string{"/m/old.gguf": "old"}}
	p := newSwapPhraser(t, fl, "/m/old.gguf")
	// Make the rollback fail too: the old file "disappears" mid-swap.
	fl.mu.Lock()
	delete(fl.models, "/m/old.gguf")
	fl.mu.Unlock()

	_, err := p.Swap(context.Background(), SwapSpec{ModelPath: "/m/broken.gguf"})
	if err == nil {
		t.Fatal("Swap returned nil when both the load and the rollback failed")
	}
	// Nothing is loaded, and the request path says so rather than panicking.
	if _, _, aerr := p.acquire(); !errors.Is(aerr, ErrNoBackend) {
		t.Errorf("acquire error = %v; want ErrNoBackend", aerr)
	}
	// Phrasing degrades to its fallback instead of failing the turn.
	got, err := p.PhraseChat(context.Background(), "привет", nil)
	if err != nil {
		t.Fatalf("PhraseChat after a total failure returned an error: %v", err)
	}
	if got == "" {
		t.Error("PhraseChat returned empty; the fallback must still say something")
	}
}

func TestSwap_WaitsForInFlightTurnAndRefusesNewOnes(t *testing.T) {
	fl := &fakeFleet{models: map[string]string{"/m/old.gguf": "old", "/m/new.gguf": "new"}}
	p := newSwapPhraser(t, fl, "/m/old.gguf")

	// Hold one turn open by taking a slot directly — the same slot every
	// request path takes.
	base, release, err := p.acquire()
	if err != nil {
		t.Fatalf("acquire: %v", err)
	}
	if base == "" {
		t.Fatal("acquire returned an empty base URL")
	}

	swapped := make(chan error, 1)
	go func() { _, e := p.Swap(context.Background(), SwapSpec{ModelPath: "/m/new.gguf"}); swapped <- e }()

	// While the swap waits to drain, a NEW turn is refused immediately rather
	// than blocked for the length of a model load.
	deadline := time.Now().Add(2 * time.Second)
	for {
		_, rel, aerr := p.acquire()
		if rel != nil {
			rel()
		}
		if errors.Is(aerr, ErrSwapping) {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("new turns were never refused during a swap (last error: %v)", aerr)
		}
		time.Sleep(10 * time.Millisecond)
	}

	// The swap cannot have completed while our turn was still in flight.
	select {
	case e := <-swapped:
		t.Fatalf("Swap finished before the in-flight turn released: %v", e)
	case <-time.After(50 * time.Millisecond):
	}

	release()
	if e := <-swapped; e != nil {
		t.Fatalf("Swap after drain: %v", e)
	}
	got, err := p.PhraseChat(context.Background(), "привет", nil)
	if err != nil || got != "new" {
		t.Fatalf("after swap: %q, %v; want the new model", got, err)
	}
}

func TestSwap_RefusedWhenWeDoNotOwnTheServer(t *testing.T) {
	// NewLLMPhraserAt points at a shared server the eval harness owns. Swapping
	// there would kill a server another process depends on.
	p := NewLLMPhraserAt("http://127.0.0.1:1/", DefaultConfig("/m/old.gguf"))
	if _, err := p.Swap(context.Background(), SwapSpec{ModelPath: "/m/new.gguf"}); !errors.Is(err, ErrSwapNotOwned) {
		t.Fatalf("Swap on a borrowed server = %v; want ErrSwapNotOwned", err)
	}
}

func TestSwap_SameModelIsANoOp(t *testing.T) {
	fl := &fakeFleet{models: map[string]string{"/m/old.gguf": "old"}}
	p := newSwapPhraser(t, fl, "/m/old.gguf")
	launchesBefore := fl.launch

	res, err := p.Swap(context.Background(), SwapSpec{ModelPath: "/m/old.gguf"})
	if err != nil {
		t.Fatalf("Swap to the live model: %v", err)
	}
	if res.Model != "old" {
		t.Errorf("res.Model = %q; want old", res.Model)
	}
	if fl.launch != launchesBefore {
		t.Errorf("%d extra launches; swapping to the live model must not reload weights", fl.launch-launchesBefore)
	}
}

func TestSwap_EmptyModelPathRefused(t *testing.T) {
	fl := &fakeFleet{models: map[string]string{"/m/old.gguf": "old"}}
	p := newSwapPhraser(t, fl, "/m/old.gguf")
	if _, err := p.Swap(context.Background(), SwapSpec{}); err == nil {
		t.Fatal("Swap with no model path returned nil error")
	}
}