llmgw/test/e2e/fallback_context_test.go

package e2e

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"strings"
	"testing"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// callWithPrompt sends a prompt of a chosen size, which is what decides whether
// a fallback can hold it.
func (h *harness) callWithPrompt(t *testing.T, key, model, prompt string) (int, http.Header, string) {
	t.Helper()
	body, _ := json.Marshal(map[string]any{
		"model":      model,
		"messages":   []map[string]string{{"role": "user", "content": prompt}},
		"max_tokens": 4,
	})
	req, err := http.NewRequestWithContext(context.Background(), http.MethodPost,
		h.router.URL+"/v1/chat/completions", bytes.NewReader(body))
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+key)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	msg, _ := io.ReadAll(resp.Body)
	return resp.StatusCode, resp.Header, string(msg)
}

// timeoutAfter is the poll these tests need because spend is written off the
// request path: the budget is not spent the instant the first reply returns.
// A predicate rather than arithmetic at four call sites, so the loops read as
// loops.
func timeoutAfter() func() bool {
	deadline := time.Now().Add(5 * time.Second)
	return func() bool {
		if time.Now().After(deadline) {
			return false
		}
		time.Sleep(25 * time.Millisecond)
		return true
	}
}

const roomyModel = "roomy-model"

// contextFallbackModels is a large-context model with a small-context one
// behind it, which is the shape that made this a bug: the fit is checked
// against what the caller asked for, and something else ends up serving.
func contextFallbackModels() []catalog.Model {
	return []catalog.Model{
		{
			Name: testModel, Source: "/models/test",
			ContextLength: 32768, MaxOutput: 4096,
			Fallbacks: []string{cheapModel, roomyModel},
			Scaling:   catalog.Scaling{Min: 1, Max: 1},
			Profiles:  []catalog.Profile{{Name: "p", Engine: "llamacpp", Weight: 1}},
		},
		{
			// Too small to hold a large prompt, and first in the list, so it is
			// the one that would have been chosen.
			Name: cheapModel, Source: "/models/cheap",
			ContextLength: 512, MaxOutput: 128,
			Scaling:  catalog.Scaling{Min: 1, Max: 1},
			Profiles: []catalog.Profile{{Name: "p", Engine: "llamacpp", Weight: 1}},
		},
		{
			Name: roomyModel, Source: "/models/roomy",
			ContextLength: 32768, MaxOutput: 4096,
			Scaling:  catalog.Scaling{Min: 1, Max: 1},
			Profiles: []catalog.Profile{{Name: "p", Engine: "llamacpp", Weight: 1}},
		},
	}
}

func contextFallbackConfig(path string) func(*config.Server) {
	return func(c *config.Server) {
		paidConfig(path)(c)
		c.Models = contextFallbackModels()
		// The same engines answer for every name, so this is about model choice
		// rather than placement.
		for i := range c.Endpoints {
			c.Endpoints[i].Models = []string{testModel, cheapModel, roomyModel}
		}
	}
}

// spendTheBudget exhausts a key so the next request falls back.
func spendTheBudget(t *testing.T, path, id string) {
	t.Helper()
	withStore(t, path, func(s *store.DB) {
		budget := int64(1)
		if err := s.PutLimits(context.Background(), store.Limits{
			Scope: store.ScopeKey, ScopeID: id,
			Budget: &budget, BudgetPeriod: store.PeriodMonth,
		}); err != nil {
			t.Fatal(err)
		}
	})
}

func TestAFallbackTooSmallForThePromptIsSkipped(t *testing.T) {
	// The bug this covers: the context check runs against the model the caller
	// asked for, and the budget fallback swaps in a different one afterwards.
	// A 20k prompt sent to a 32k model would land on a 512-token one, and the
	// engine would truncate it or refuse — leaving the caller with a failure
	// naming a model they never chose.
	path, mint := provision(t)
	key := mint(t, "broke")
	spendTheBudget(t, path, "key-broke")

	h := start(t, 2, contextFallbackConfig(path))

	// Comfortably beyond the small model's 512 tokens, and well within 32768.
	big := strings.Repeat("word ", 8000)

	if code, _, msg := h.callWithPrompt(t, key, testModel, big); code != http.StatusOK {
		t.Fatalf("the first request should go through: %d %s", code, msg)
	}

	deadline := timeoutAfter()
	for deadline() {
		code, hdr, msg := h.callWithPrompt(t, key, testModel, big)
		served := hdr.Get("x-llmgw-served-model")
		if served == cheapModel {
			t.Fatalf("served by a model that cannot hold the prompt: %s", msg)
		}
		if code == http.StatusOK && served == roomyModel {
			return // skipped the small one, took the one that fits
		}
	}
	t.Fatal("the budget never ran out, or no fallback served")
}

func TestWhenNoAffordableModelFitsTheRequestIsRefused(t *testing.T) {
	// Being served by a model that cannot hold the prompt is worse than a
	// refusal: it fails later, in the engine, naming a model the caller did
	// not choose. The history is deliberately not trimmed to make it fit —
	// dropping someone's messages to stay under a limit they did not ask about
	// would be a silent loss inside a request that was already being quietly
	// moved.
	path, mint := provision(t)
	key := mint(t, "broke")
	spendTheBudget(t, path, "key-broke")

	h := start(t, 2, func(c *config.Server) {
		contextFallbackConfig(path)(c)
		// Only the small model behind it now, so nothing affordable fits.
		c.Models[0].Fallbacks = []string{cheapModel}
	})

	big := strings.Repeat("word ", 8000)
	if code, _, msg := h.callWithPrompt(t, key, testModel, big); code != http.StatusOK {
		t.Fatalf("the first request should go through: %d %s", code, msg)
	}

	deadline := timeoutAfter()
	for deadline() {
		code, hdr, _ := h.callWithPrompt(t, key, testModel, big)
		if hdr.Get("x-llmgw-served-model") == cheapModel {
			t.Fatal("a prompt was sent to a model too small to hold it")
		}
		if code == http.StatusTooManyRequests {
			return
		}
	}
	t.Fatal("the budget never ran out")
}

func TestASmallPromptStillReachesTheSmallFallback(t *testing.T) {
	// The check must not disable the feature. A prompt that fits is served by
	// the first affordable model as before.
	path, mint := provision(t)
	key := mint(t, "broke")
	spendTheBudget(t, path, "key-broke")

	h := start(t, 2, contextFallbackConfig(path))

	if code, _, msg := h.callWithPrompt(t, key, testModel, "hi"); code != http.StatusOK {
		t.Fatalf("the first request should go through: %d %s", code, msg)
	}

	deadline := timeoutAfter()
	for deadline() {
		code, hdr, _ := h.callWithPrompt(t, key, testModel, "hi")
		if code == http.StatusOK && hdr.Get("x-llmgw-served-model") == cheapModel {
			return
		}
		if code == http.StatusTooManyRequests {
			t.Fatal("a prompt that fits was refused instead of falling back")
		}
	}
	t.Fatal("the budget never ran out, or the fallback never served")
}