llmgw/internal/router/external_test.go

package router

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/notify"
	"gitea.mynodes.xyz/ilker/llmgw/internal/notify/notifytest"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/breaker"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
)

func TestALocalModelWithNoCapacityFallsBackToAnExternalProvider(t *testing.T) {
	local := newFakeEngine(t, "local")
	local.status.Store(http.StatusServiceUnavailable)

	var calls atomic.Int64
	var gotModel, gotAuth atomic.Value
	provider := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/metrics" {
			t.Error("the provider was scraped for metrics")
		}
		calls.Add(1)
		gotAuth.Store(r.Header.Get("Authorization"))
		var body map[string]any
		b, _ := io.ReadAll(r.Body)
		_ = json.Unmarshal(b, &body)
		gotModel.Store(fmt.Sprint(body["model"]))
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, `{"model":"gpt-4o-mini","choices":[{"message":{"content":"from outside"}}],"usage":{"prompt_tokens":3,"completion_tokens":2}}`)
	}))
	t.Cleanup(provider.Close)

	cfg := baseConfig(config.Endpoint{ID: "l", URL: local.srv.URL, Models: []string{"local"}})
	cloud := catalog.Model{Name: "cloud", Source: "external",
		External: &catalog.External{URL: provider.URL, Model: "gpt-4o-mini"}}
	cfg.Models = []catalog.Model{catModel("local", nil, "cloud"), cloud}
	rt := newTestRouter(t, cfg)
	rt.SetEndpoints(t.Context(), []endpoint.Info{
		{ID: "l", URL: local.srv.URL, Models: []string{"local"}},
		{ID: "external:cloud", URL: provider.URL, Models: []string{"cloud"}, APIKey: "provider-key",
			Engine: endpoint.EngineExternal, Pool: "external", Weight: 1, UpstreamModel: "gpt-4o-mini"},
	})

	rec := post(t, rt.Handler(), map[string]any{"model": "local", "messages": []any{}}, nil)
	if rec.Code != http.StatusOK {
		t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
	}
	if calls.Load() != 1 {
		t.Fatalf("the provider was called %d times", calls.Load())
	}
	if gotModel.Load() != "gpt-4o-mini" {
		t.Errorf("the provider was asked for %v, not its own name for the model", gotModel.Load())
	}
	if gotAuth.Load() != "Bearer provider-key" {
		t.Errorf("the provider got %v, not its key", gotAuth.Load())
	}
	if got := rec.Header().Get("x-llmgw-served-model"); got != "cloud" {
		t.Errorf("the client should be told it was served by cloud, got %q", got)
	}
	// Never scraped, and never stale for it.
	for _, s := range rt.reg.Snapshots("cloud") {
		if s.Stale {
			t.Error("an external endpoint reads as stale")
		}
	}
}

func TestAnExternalModelIsNotAModelWithProfiles(t *testing.T) {
	good := catalog.Model{Name: "cloud", Source: "external",
		External: &catalog.External{URL: "https://api.example.com", Model: "gpt-4o-mini"}}
	if _, err := catalog.New([]catalog.Model{good}); err != nil {
		t.Fatalf("a valid external model was refused: %v", err)
	}
	for name, m := range map[string]catalog.Model{
		"profiles": {Name: "x", Source: "external", External: &catalog.External{URL: "https://a.b", Model: "m"},
			Profiles: []catalog.Profile{{Name: "p", Engine: "llamacpp"}}},
		"no url":   {Name: "x", Source: "external", External: &catalog.External{URL: "api.example.com", Model: "m"}},
		"no model": {Name: "x", Source: "external", External: &catalog.External{URL: "https://a.b"}},
	} {
		if _, err := catalog.New([]catalog.Model{m}); err == nil {
			t.Errorf("%s: accepted", name)
		}
	}
}

// Two keys for one provider, one of them at its limit, or refused: every
// request is answered by the other key, the failing one's breaker opens, and
// a refused key is said, by its endpoint and never its value.
func TestAKeyAtItsLimitOrRefusedIsPassedOverForAnother(t *testing.T) {
	for _, status := range []int{http.StatusTooManyRequests, http.StatusUnauthorized} {
		provider := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if r.Header.Get("Authorization") == "Bearer key-a" {
				http.Error(w, "no", status)
				return
			}
			w.Header().Set("Content-Type", "application/json")
			fmt.Fprint(w, `{"model":"m","choices":[{"message":{"content":"ok"}}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
		}))
		t.Cleanup(provider.Close)
		cfg := baseConfig()
		cfg.Models = []catalog.Model{{Name: "cloud", Source: "external",
			External: &catalog.External{URL: provider.URL, Model: "m"}}}
		rt := newTestRouter(t, cfg)
		key := func(id, k string) endpoint.Info {
			return endpoint.Info{ID: id, URL: provider.URL, Models: []string{"cloud"}, APIKey: k,
				Engine: endpoint.EngineExternal, Pool: "external", Weight: 1, UpstreamModel: "m"}
		}
		rt.SetEndpoints(t.Context(), []endpoint.Info{key("external:cloud#1", "key-a"), key("external:cloud#2", "key-b")})
		n, events := notifytest.Catch(t)
		rt.SetNotifier(n)

		// Until the failing key's breaker opens: how many requests the
		// scheduler sends it before then is its own choice.
		open := func() bool {
			for _, s := range rt.reg.Snapshots("cloud") {
				if s.ID == "external:cloud#1" && s.Breaker != breaker.Closed {
					return true
				}
			}
			return false
		}
		for i := 0; i < 100 && !open(); i++ {
			if rec := post(t, rt.Handler(), map[string]any{"model": "cloud", "messages": []any{}}, nil); rec.Code != http.StatusOK {
				t.Fatalf("key-a answering %d: request %d got %d: %s", status, i, rec.Code, rec.Body.String())
			}
		}
		for _, s := range rt.reg.Snapshots("cloud") {
			if s.ID == "external:cloud#1" && s.Breaker == breaker.Closed {
				t.Errorf("%d: the failing key's breaker is still closed", status)
			}
			if s.ID == "external:cloud#2" && s.Breaker != breaker.Closed {
				t.Errorf("%d: the good key's breaker opened", status)
			}
		}
		var refused []notify.Event
		for _, e := range events() {
			if e.Kind == notify.ExternalKeyRefused {
				refused = append(refused, e)
			}
		}
		switch {
		case status == http.StatusTooManyRequests && len(refused) != 0:
			t.Errorf("a key at its limit was said to be refused: %+v", refused)
		case status == http.StatusUnauthorized && (len(refused) != 1 || refused[0].Subject != "external:cloud#1" ||
			strings.Contains(refused[0].Message, "key-a")):
			t.Errorf("a refused key was said as %+v", refused)
		}
	}
}