Maven/internal/llm/remote_test.go

package llm

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

// completionServer stands in for a llama-server. It counts what reached it, so
// a test can say which of the two models answered.
func completionServer(t *testing.T, reply string, hits *atomic.Int64) *httptest.Server {
	t.Helper()
	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"choices":[{"message":{"content":"` + reply + `"}}]}`))
	}))
	t.Cleanup(s.Close)
	return s
}

func healthServer(t *testing.T, ok *atomic.Bool) *httptest.Server {
	t.Helper()
	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !ok.Load() {
			w.WriteHeader(http.StatusServiceUnavailable)
			return
		}
		w.WriteHeader(http.StatusOK)
	}))
	t.Cleanup(s.Close)
	return s
}

// waitFor polls until cond holds or the deadline passes. The prober runs on its
// own goroutine, so a test has to wait for it rather than assume it has run.
func waitFor(t *testing.T, cond func() bool) bool {
	t.Helper()
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		if cond() {
			return true
		}
		time.Sleep(5 * time.Millisecond)
	}
	return false
}

// The unconfigured deploy. No remote, no probing, every call to the floor —
// exactly what the box does today.
func TestNoRemoteGoesToTheFloor(t *testing.T) {
	var floorHits atomic.Int64
	floor := completionServer(t, "floor", &floorHits)

	p := NewPair(nil, New(floor.URL, time.Second), "", time.Second)
	p.Start(context.Background())
	defer p.Stop()

	if p.Available() {
		t.Fatal("a Pair with no remote reports available")
	}
	out, err := p.Complete(context.Background(), Req{User: "привет"})
	if err != nil {
		t.Fatalf("complete: %v", err)
	}
	if out != "floor" || floorHits.Load() != 1 {
		t.Fatalf("out = %q, floor hits = %d", out, floorHits.Load())
	}
}

// The workstation is up, so it answers and the resident model is not touched.
func TestAvailableRemoteAnswers(t *testing.T) {
	var remoteHits, floorHits atomic.Int64
	remote := completionServer(t, "remote", &remoteHits)
	floor := completionServer(t, "floor", &floorHits)
	up := &atomic.Bool{}
	up.Store(true)
	health := healthServer(t, up)

	p := NewPair(New(remote.URL, time.Second), New(floor.URL, time.Second), health.URL, 20*time.Millisecond)
	p.Start(context.Background())
	defer p.Stop()
	if !waitFor(t, p.Available) {
		t.Fatal("prober never saw the remote come up")
	}

	out, err := p.Complete(context.Background(), Req{User: "привет"})
	if err != nil {
		t.Fatalf("complete: %v", err)
	}
	if out != "remote" || floorHits.Load() != 0 {
		t.Fatalf("out = %q, floor hits = %d", out, floorHits.Load())
	}
}

// The card is busy, so /health refuses and Complete degrades silently. This is
// the constraint from 483: the workstation being down is indistinguishable from
// today's behaviour.
func TestBusyCardFallsBackSilently(t *testing.T) {
	var remoteHits, floorHits atomic.Int64
	remote := completionServer(t, "remote", &remoteHits)
	floor := completionServer(t, "floor", &floorHits)
	health := healthServer(t, &atomic.Bool{}) // never ok

	p := NewPair(New(remote.URL, time.Second), New(floor.URL, time.Second), health.URL, 20*time.Millisecond)
	p.Start(context.Background())
	defer p.Stop()
	time.Sleep(60 * time.Millisecond)

	out, err := p.Complete(context.Background(), Req{User: "привет"})
	if err != nil {
		t.Fatalf("complete: %v", err)
	}
	if out != "floor" || remoteHits.Load() != 0 {
		t.Fatalf("out = %q, remote hits = %d", out, remoteHits.Load())
	}
}

// The cached admission answer can be one interval out of date, so a remote that
// dies between probes must still not break the turn.
func TestRemoteErrorMidRequestFallsBack(t *testing.T) {
	var floorHits atomic.Int64
	dead := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
	}))
	defer dead.Close()
	floor := completionServer(t, "floor", &floorHits)
	up := &atomic.Bool{}
	up.Store(true)
	health := healthServer(t, up)

	p := NewPair(New(dead.URL, time.Second), New(floor.URL, time.Second), health.URL, time.Hour)
	p.Start(context.Background())
	defer p.Stop()
	if !waitFor(t, p.Available) {
		t.Fatal("prober never saw the remote come up")
	}

	out, err := p.Complete(context.Background(), Req{User: "привет"})
	if err != nil {
		t.Fatalf("complete: %v", err)
	}
	if out != "floor" || floorHits.Load() != 1 {
		t.Fatalf("out = %q, floor hits = %d", out, floorHits.Load())
	}
	// The failed request must have corrected the cached answer, so the next
	// one does not walk into the same hole.
	if p.Available() {
		t.Fatal("a failed remote request left the admission answer up")
	}
}

// A workstation that accepts the connection and then hangs must not spend the
// whole turn budget. It used to: the remote got the caller's context unchanged,
// so the fallback ran on an expired one and the floor returned the deadline
// error instead of an answer. The turn broke on the workstation being slow,
// which is the one outcome docs/offload.md rules out.
func TestHangingRemoteLeavesTheFloorABudget(t *testing.T) {
	var floorHits atomic.Int64
	// released, not r.Context().Done(): httptest.Server.Close waits for the
	// handler, and a handler that only watches the request context can outlive
	// the test when the client hangs up without the server noticing.
	released := make(chan struct{})
	hang := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		select {
		case <-released:
		case <-r.Context().Done():
		}
	}))
	defer hang.Close()
	defer close(released)
	floor := completionServer(t, "floor", &floorHits)
	up := &atomic.Bool{}
	up.Store(true)
	health := healthServer(t, up)

	p := NewPair(New(hang.URL, time.Minute), New(floor.URL, time.Minute), health.URL, time.Hour)
	p.Start(context.Background())
	defer p.Stop()
	if !waitFor(t, p.Available) {
		t.Fatal("prober never saw the remote come up")
	}

	ctx, cancel := context.WithTimeout(context.Background(), 400*time.Millisecond)
	defer cancel()
	out, err := p.Complete(ctx, Req{User: "привет"})
	if err != nil {
		t.Fatalf("complete: %v", err)
	}
	if out != "floor" || floorHits.Load() != 1 {
		t.Fatalf("out = %q, floor hits = %d", out, floorHits.Load())
	}
}

// The naming half of the degradation rule. A world question must not be handed
// to the resident model, because it answers by inventing.
func TestCompleteRemoteNamesTheGap(t *testing.T) {
	var floorHits atomic.Int64
	floor := completionServer(t, "floor", &floorHits)
	health := healthServer(t, &atomic.Bool{}) // never ok

	p := NewPair(New("http://127.0.0.1:1", time.Second), New(floor.URL, time.Second), health.URL, 20*time.Millisecond)
	p.Start(context.Background())
	defer p.Stop()
	time.Sleep(60 * time.Millisecond)

	if _, err := p.CompleteRemote(context.Background(), Req{User: "почему небо голубое"}); !errors.Is(err, ErrRemoteUnavailable) {
		t.Fatalf("err = %v, want ErrRemoteUnavailable", err)
	}
	if floorHits.Load() != 0 {
		t.Fatalf("CompleteRemote fell back to the floor %d times", floorHits.Load())
	}
}

// Routing sits on the hot path and must never pay for a health check. Available
// reads a cached flag, so it costs no network at all.
func TestAvailableDoesNotProbe(t *testing.T) {
	var probes atomic.Int64
	health := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		probes.Add(1)
		w.WriteHeader(http.StatusOK)
	}))
	defer health.Close()

	p := NewPair(New("http://127.0.0.1:1", time.Second), New("http://127.0.0.1:1", time.Second), health.URL, time.Hour)
	p.Start(context.Background())
	defer p.Stop()
	if !waitFor(t, p.Available) {
		t.Fatal("prober never ran")
	}

	before := probes.Load()
	for range 1000 {
		p.Available()
	}
	if got := probes.Load(); got != before {
		t.Fatalf("1000 Available calls made %d probes", got-before)
	}
}

// A Pair with no floor is a configuration mistake, and it must say so rather
// than silently having nowhere to degrade to.
func TestNoFloorIsAnError(t *testing.T) {
	p := NewPair(nil, nil, "", time.Second)
	if _, err := p.Complete(context.Background(), Req{User: "привет"}); !errors.Is(err, ErrNoFloor) {
		t.Fatalf("err = %v, want ErrNoFloor", err)
	}
}

// The workstation is behind mavgpud, which requires a bearer token on the
// completion and on /health alike. A probe that did not carry it would answer
// 401, Pair would read that as a card that is busy, and every turn would fall
// back to the resident model with nothing in the log naming why.
func TestPairSignsTheProbeAndTheCompletion(t *testing.T) {
	seen := make(chan string, 2)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		seen <- r.Header.Get("Authorization")
		if r.URL.Path == "/health" {
			return
		}
		_, _ = w.Write([]byte(`{"choices":[{"message":{"content":"ok"}}]}`))
	}))
	defer srv.Close()

	remote := New(srv.URL, time.Second)
	remote.SetToken("s3cret")
	p := NewPair(remote, New(srv.URL, time.Second), srv.URL+"/health", time.Hour)
	p.probe(context.Background())
	if !p.Available() {
		t.Fatal("the probe did not admit an answering workstation")
	}
	if _, err := p.Complete(context.Background(), Req{User: "привет"}); err != nil {
		t.Fatalf("complete: %v", err)
	}
	for i := 0; i < 2; i++ {
		if got := <-seen; got != "Bearer s3cret" {
			t.Errorf("request %d carried %q, want the bearer token", i, got)
		}
	}
}