ModelRouter/internal/gateway/discover_test.go

package gateway

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

// realCatalogBody is the actual /models response captured from the amr-link
// endpoint, not a hand-written fixture. Two properties of it are the reason
// this test exists at all:
//
//   - it nests modalities under "architecture", which the naive parser misses;
//   - it lists five image-generation ids among 34 chat ids, so a parser that
//     returns everything as "chat" silently poisons the AUTO chain.
const realCatalogBody = `{"data":[{"id":"claude-fable-5","object":"model","created":1781081344,"owned_by":"vela","architecture":{"input_modalities":["text","image","file"],"output_modalities":["text"]},"metadata":{"cost":"very_high","capability":"best_quality","context_limit":1000000,"context_budget":232000},"default":false,"enabled":true},{"id":"gpt-image-2","object":"model","architecture":{"input_modalities":["text","image"],"output_modalities":["image"]},"metadata":{"capability":"best_quality"},"enabled":true},{"id":"deepseek-v4-flash","object":"model","architecture":{"input_modalities":["text"],"output_modalities":["text"]},"metadata":{"cost":"medium","capability":"standard","context_limit":1048000},"enabled":true},{"id":"seedream-5.0-pro","object":"model","architecture":{"input_modalities":["text","image"],"output_modalities":["image"]},"metadata":{"capability":"best_quality"},"enabled":true}]}`

func TestParseModelCatalog_RealUpstreamShape(t *testing.T) {
	models, err := parseModelCatalog([]byte(realCatalogBody))
	if err != nil {
		t.Fatalf("parse: %v", err)
	}
	if len(models) != 4 {
		t.Fatalf("want 4 models, got %d", len(models))
	}

	kind := map[string]string{}
	for _, m := range models {
		kind[m.ID] = m.Kind
	}
	// The two image generators must be classified by the upstream's declared
	// output modality, not by their names.
	for _, id := range []string{"gpt-image-2", "seedream-5.0-pro"} {
		if kind[id] != "image" {
			t.Errorf("%s: got kind %q, want \"image\"", id, kind[id])
		}
	}
	for _, id := range []string{"claude-fable-5", "deepseek-v4-flash"} {
		if kind[id] != "chat" {
			t.Errorf("%s: got kind %q, want \"chat\"", id, kind[id])
		}
	}

	// Metadata the upstream sent must survive; keys it did not send must not
	// appear as zero values.
	if got := models[0].Meta["context_limit"]; got == nil {
		t.Error("claude-fable-5: context_limit dropped (it is nested under metadata in the real response)")
	} else if f, ok := got.(float64); !ok || f != 1000000 {
		t.Errorf("context_limit = %#v, want numeric 1000000 so the UI can sort on it", got)
	}
	if _, present := models[0].Meta["cost"]; !present {
		t.Error("claude-fable-5: cost dropped")
	}
	// context_budget is present on this entry and must not be dropped either.
	if _, present := models[0].Meta["context_budget"]; !present {
		t.Error("claude-fable-5: context_budget dropped")
	}
	// deepseek-v4-flash DOES send metadata; gpt-image-2 sends only capability.
	// Neither may lose it.
	if models[2].Meta["context_limit"] == nil {
		t.Error("deepseek-v4-flash: context_limit dropped")
	}
	if models[2].Meta["cost"] == nil {
		t.Error("deepseek-v4-flash: cost dropped")
	}
	if _, bogus := models[0].Meta["made_up_field"]; bogus {
		t.Error("invented a metadata key that the upstream never sent")
	}
}

// TestDiscoveredModelKind_ImageOnlyByUpstreamMetadata is the regression guard
// for the failure mode that motivated keeping kind="chat" as the default: a
// name-based heuristic marks "ideogram-v2" or "gpt-image-1" as an image model,
// but a catalog that also serves chat for that id would then have its chat
// entry removed from the AUTO chain. Only the upstream's declared OUTPUT
// modality may classify a model as image.
func TestDiscoveredModelKind_ImageOnlyByUpstreamMetadata(t *testing.T) {
	cases := []struct {
		name string
		in   upstreamModel
		want string
	}{
		{
			// A catalog that DECLARES output modality image is authoritative.
			// The "kind" string saying chat must not override it.
			name: "declared output modality image wins over a chat kind string",
			in: func() upstreamModel {
				var m upstreamModel
				m.Kind = "chat"
				m.Architecture.OutputModalities = []string{"image"}
				return m
			}(),
			want: "image",
		},
		{
			name: "bare array modality string",
			in:   upstreamModel{ID: "x", Modality: "text-to-image"},
			want: "image",
		},
		{
			name: "image INPUT is still a chat model",
			in: func() upstreamModel {
				var m upstreamModel
				m.Architecture.InputModalities = []string{"text", "image"}
				return m
			}(),
			want: "chat",
		},
		{
			name: "file input is still a chat model",
			in: func() upstreamModel {
				var m upstreamModel
				m.Architecture.InputModalities = []string{"text", "file"}
				return m
			}(),
			want: "chat",
		},
		{
			name: "vision capability flag is chat",
			in:   upstreamModel{ID: "y", Capabilities: modelCaps{Vision: true}},
			want: "chat",
		},
		{
			name: "capabilities.image_generation is image",
			in:   upstreamModel{ID: "z", Capabilities: modelCaps{ImageGeneration: true}},
			want: "image",
		},
		{
			name: "no modality info at all defaults to chat",
			in:   upstreamModel{ID: "plain"},
			want: "chat",
		},
		{
			// The default must be chat even when the name LOOKS like an image
			// model. Guessing from the name is what puts gpt-image-2 into the
			// chat AUTO chain, where every request to it fails.
			name: "image-looking name without upstream metadata stays chat",
			in:   upstreamModel{ID: "gpt-image-2"},
			want: "chat",
		},
		{
			name: "vision-looking name without upstream metadata stays chat",
			in:   upstreamModel{ID: "ideogram-v2"},
			want: "chat",
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			if got := discoveredModelKind(tc.in); got != tc.want {
				t.Errorf("got %q, want %q", got, tc.want)
			}
		})
	}
}

func TestParseModelCatalog_RejectsNonCatalog(t *testing.T) {
	// A 200 with valid JSON that is not a catalog must NOT be reported as an
	// empty result. Reporting zero models would let a login page or a
	// rate-limit body create a source with no models in it.
	for name, body := range map[string]string{
		"login page html": `<html><body>Sign in</body></html>`,
		"rate limit json": `{"error":{"message":"slow down","type":"rate_limit"}}`,
		"empty object":    `{}`,
		"null":            `null`,
		"wrong shape":     `{"objects":[{"id":"m"}]}`,
	} {
		t.Run(name, func(t *testing.T) {
			if _, err := parseModelCatalog([]byte(body)); err == nil {
				t.Errorf("expected an error for %q, got nil (would create an empty source)", name)
			}
		})
	}
}

func TestParseModelCatalog_AcceptsBareArray(t *testing.T) {
	models, err := parseModelCatalog([]byte(`[{"id":"a"},{"id":"b"}]`))
	if err != nil {
		t.Fatalf("bare array should parse: %v", err)
	}
	if len(models) != 2 {
		t.Fatalf("want 2, got %d", len(models))
	}
}

// TestDiscoverModels_NormalizesWhenOnlyBareModelsAnswers: operator pastes
// "https://host/v1" (the form every SDK example shows) but the vendor serves
// the catalog at "https://host/models" — the version segment on the path is
// implied, not doubled. Discovery must strip the segment off the base and
// report it on the endpoint, because openai's adapter endpoint is
// "/chat/completions" and the two are concatenated.
func TestDiscoverModels_NormalizesWhenOnlyBareModelsAnswers(t *testing.T) {
	var paths []string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		paths = append(paths, r.URL.Path)
		if r.URL.Path != "/models" {
			w.WriteHeader(http.StatusNotFound)
			return
		}
		_, _ = w.Write([]byte(realCatalogBody))
	}))
	defer srv.Close()

	res, err := DiscoverModels(context.Background(), srv.URL+"/v1", "sk-test")
	if err != nil {
		t.Fatalf("discovery failed: %v", err)
	}
	if res.BaseURL != srv.URL {
		t.Errorf("base_url = %q, want %q (version segment moved off the base)", res.BaseURL, srv.URL)
	}
	// "/models" carries no version segment, but the operator typed one — the
	// endpoint has to restore it or the chat call would miss /v1.
	if res.Endpoint != "/v1/chat/completions" {
		t.Errorf("endpoint = %q, want %q", res.Endpoint, "/v1/chat/completions")
	}
	if res.Note == "" {
		t.Error("silently rewriting the operator's base_url is worse than saying so")
	}
	if len(res.Models) != 2 || len(res.ImageModels) != 2 {
		t.Errorf("got %d chat / %d image, want 2 / 2", len(res.Models), len(res.ImageModels))
	}
	if len(paths) < 2 {
		t.Errorf("tried only %v; expected a fallback search", paths)
	}
}

// TestDiscoverModels_BareHostWithVersionedModelsPath: operator pastes the host
// root, vendor serves /v1/models. The base needs no rewriting; the endpoint
// does.
func TestDiscoverModels_BareHostWithVersionedModelsPath(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/v1/models" {
			w.WriteHeader(http.StatusNotFound)
			return
		}
		_, _ = w.Write([]byte(`{"data":[{"id":"m1"}]}`))
	}))
	defer srv.Close()

	res, err := DiscoverModels(context.Background(), srv.URL, "")
	if err != nil {
		t.Fatalf("discovery failed: %v", err)
	}
	if res.BaseURL != srv.URL {
		t.Errorf("base_url = %q, want %q", res.BaseURL, srv.URL)
	}
	if res.Endpoint != "/v1/chat/completions" {
		t.Errorf("endpoint = %q, want %q", res.Endpoint, "/v1/chat/completions")
	}
}

// TestDiscoverModels_KeepsWorkingBaseURLUnchanged covers the vendor that really
// does serve "/v1/models": the typed URL must survive untouched so no needless
// config rewrite is suggested.
func TestDiscoverModels_KeepsWorkingBaseURLUnchanged(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/v1/models" {
			w.WriteHeader(http.StatusNotFound)
			return
		}
		_, _ = w.Write([]byte(`{"data":[{"id":"m1"},{"id":"m2"}]}`))
	}))
	defer srv.Close()

	// Base typed WITH "/v1", catalog served at "…/v1/models". That is exactly
	// the amr-link shape: first candidate succeeds, nothing is rewritten, and
	// the adapter's default endpoint already composes.
	res, err := DiscoverModels(context.Background(), srv.URL+"/v1", "")
	if err != nil {
		t.Fatalf("discovery failed: %v", err)
	}
	if res.BaseURL != srv.URL+"/v1" {
		t.Errorf("base_url = %q, want it unchanged", res.BaseURL)
	}
	if res.Endpoint != "" {
		t.Errorf("endpoint = %q, want empty (the typed base already carries /v1)", res.Endpoint)
	}
	if res.Note != "" {
		t.Errorf("note = %q, want empty when nothing was rewritten", res.Note)
	}
	if len(res.Models) != 2 {
		t.Errorf("got %d chat models, want 2", len(res.Models))
	}
}

// TestDiscoverModels_ErrorNamesEveryTriedURL: a failed probe is useless if it
// only reports the last URL. The operator needs to see the full search.
func TestDiscoverModels_ErrorNamesEveryTriedURL(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusNotFound)
	}))
	defer srv.Close()

	_, err := DiscoverModels(context.Background(), srv.URL+"/v1", "k")
	if err == nil {
		t.Fatal("expected failure against an all-404 server")
	}
	msg := err.Error()
	for _, want := range []string{"/models", "/v1/models", "404"} {
		if !strings.Contains(msg, want) {
			t.Errorf("error %q should mention %q", msg, want)
		}
	}
}

func TestDiscoveryCandidates_DedupesAndPrefersTyped(t *testing.T) {
	cands := discoveryCandidates("https://h/v1/")
	if len(cands) == 0 {
		t.Fatal("no candidates")
	}
	if cands[0].modelsPath != "/models" || cands[0].baseURL != "https://h/v1" {
		t.Errorf("first candidate = %+v, want the typed base with a bare /models path", cands[0])
	}
	seen := map[string]bool{}
	for _, c := range cands {
		k := c.baseURL + "|" + c.modelsPath
		if seen[k] {
			t.Errorf("duplicate candidate %s", k)
		}
		seen[k] = true
	}
	if !seen["https://h|/v1/models"] {
		t.Error("host root with the canonical path was never tried; that is the shape the real amr-link endpoint uses")
	}
	// The bare "/models" form must exist too — some gateways serve it there.
	hasBare := false
	for _, c := range cands {
		if c.modelsPath == "/models" {
			hasBare = true
		}
	}
	if !hasBare {
		t.Error("no /models candidate; vendors serving the catalog at /models would be missed")
	}
	// Preference must be monotonically non-decreasing so the sort is stable.
	for i := 1; i < len(cands); i++ {
		if cands[i].prefer < cands[i-1].prefer {
			t.Errorf("candidates out of preference order at %d: %d < %d", i, cands[i].prefer, cands[i-1].prefer)
		}
	}
}

func TestVersionSegment(t *testing.T) {
	cases := map[string]string{
		"/v1/models":   "/v1",
		"/v1beta/models": "/v1beta",
		"/models":      "",
		"/api/models":  "",
		"/vision/models": "",
	}
	for in, want := range cases {
		if got := versionSegment(in); got != want {
			t.Errorf("versionSegment(%q) = %q, want %q", in, got, want)
		}
	}
}

func TestTrimVersionSuffix(t *testing.T) {
	cases := map[string]string{
		"https://h/v1":          "https://h",
		"https://h/v2":          "https://h",
		"https://h/v1beta":      "https://h",
		"https://h/api/v1":      "https://h/api",
		"https://h/vision":      "https://h/vision", // not numeric — a real path segment
		"https://h":             "https://h",
		"https://h/openai":      "https://h/openai",
		"https://h/v1/extra":    "https://h/v1/extra",
		// A segment that merely STARTS with v must not be stripped: stripping
		// it would rewrite base "https://h/vision" to "https://h" and, worse,
		// would treat a real vendor path as a version.
		"https://h/vertex":      "https://h/vertex",
		"https://h/volcengine":  "https://h/volcengine",
		"https://h/vertex-ai/v1": "https://h/vertex-ai",
	}
	for in, want := range cases {
		if got := trimVersionSuffix(in); got != want {
			t.Errorf("trimVersionSuffix(%q) = %q, want %q", in, got, want)
		}
	}
}

// TestIsVersionSegment_RequiresDigits pins the distinction the whole endpoint
// derivation rests on: "v1"/"v1beta" are versions, "vision"/"vertex" are path
// segments. Without the digit requirement, stripping turns a vendor's real
// base URL into a different one.
func TestIsVersionSegment_RequiresDigits(t *testing.T) {
	yes := []string{"v1", "v2", "v10", "v1beta", "v1alpha", "v1preview"}
	no := []string{"vision", "vertex", "volcengine", "openai", "v", "vv", "api", "beta", "v1x"}
	for _, s := range yes {
		if !isVersionSegment(s) {
			t.Errorf("isVersionSegment(%q) = false, want true", s)
		}
	}
	for _, s := range no {
		if isVersionSegment(s) {
			t.Errorf("isVersionSegment(%q) = true, want false", s)
		}
	}
}

// TestComposeEndpoint_ExactlyOneVersionSegment is the invariant test. Every row
// below is a URL form that was observed to 404 or to route wrongly when the
// version segment ended up in the wrong place.
func TestComposeEndpoint_ExactlyOneVersionSegment(t *testing.T) {
	cases := []struct {
		name          string
		modelsPath    string
		persistedBase string
		typedBase     string
		want          string
	}{
		{
			name:       "typed base has /v1 and catalog served there: base keeps it, endpoint empty",
			modelsPath: "/models", persistedBase: "https://h/v1", typedBase: "https://h/v1",
			want: "",
		},
		{
			name:       "catalog served at versioned path, base stripped: endpoint carries it",
			modelsPath: "/v1/models", persistedBase: "https://h", typedBase: "https://h/v1",
			want: "/v1/chat/completions",
		},
		{
			name:       "bare host, bare catalog, nothing typed: adapter default composes",
			modelsPath: "/models", persistedBase: "https://h", typedBase: "https://h",
			want: "",
		},
		{
			name:       "bare host, versioned catalog: endpoint owns the version",
			modelsPath: "/v1/models", persistedBase: "https://h", typedBase: "https://h",
			want: "/v1/chat/completions",
		},
		{
			name:       "typed /v1 but catalog only at bare /models: strip base, restore on endpoint",
			modelsPath: "/models", persistedBase: "https://h", typedBase: "https://h/v1",
			want: "/v1/chat/completions",
		},
		{
			name:       "non-version segment must never be promoted onto the endpoint",
			modelsPath: "/models", persistedBase: "https://h/vision", typedBase: "https://h/vision",
			want: "",
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			got := composeEndpoint(tc.modelsPath, tc.persistedBase, tc.typedBase)
			if got != tc.want {
				t.Errorf("composeEndpoint = %q, want %q", got, tc.want)
			}
			// Whichever way it resolved, base + endpoint must never carry the
			// same version segment twice.
			onBase := trimVersionSuffix(tc.persistedBase) != tc.persistedBase
			onEndpoint := versionSegment(strings.TrimSuffix(got, "/chat/completions") + "/models") != ""
			if onBase && onEndpoint {
				t.Errorf("version segment duplicated: base %q + endpoint %q", tc.persistedBase, got)
			}
		})
	}
}

// TestDiscoverModels_DoesNotProbeEveryModel guards the decision not to spend
// money verifying each catalog entry. Only the catalog request is made.
func TestDiscoverModels_DoesNotProbeEveryModel(t *testing.T) {
	var calls int
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls++
		_, _ = w.Write([]byte(`{"data":[{"id":"a"},{"id":"b"},{"id":"c"}]}`))
	}))
	defer srv.Close()

	if _, err := DiscoverModels(context.Background(), srv.URL+"/v1", "k"); err != nil {
		t.Fatalf("discovery failed: %v", err)
	}
	if calls != 1 {
		t.Errorf("made %d requests, want exactly 1 (the catalog)", calls)
	}
}

func TestModelsToConfig_PrioritiesAndKinds(t *testing.T) {
	got := modelsToConfig([]string{"a", "b"}, []string{"img1"})
	if len(got) != 3 {
		t.Fatalf("want 3 models, got %d", len(got))
	}
	if got[2].ID != "img1" || got[2].Kind != "image" {
		t.Errorf("image model mis-sorted: %+v", got[2])
	}
	if got[2].Priority <= got[0].Priority {
		t.Errorf("image model priority %d should sort after chat %d", got[2].Priority, got[0].Priority)
	}
}

// TestDiscoverModels_ImageModelsNeverNull guards the UI contract end to end,
// not just on the struct: the WebUI does out.concat(image_models), and a JSON
// null there renders "undefined" next to every image model row. This asserts on
// the served bytes rather than the Go field, because that is what the bug was.
func TestDiscoverModels_ImageModelsNeverNull(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Only chat models — the case where the slice is never appended to.
		_, _ = w.Write([]byte(`{"data":[{"id":"a"},{"id":"b"}]}`))
	}))
	defer srv.Close()

	res, err := DiscoverModels(context.Background(), srv.URL+"/v1", "k")
	if err != nil {
		t.Fatalf("discovery failed: %v", err)
	}
	if len(res.ImageModels) != 0 {
		t.Fatalf("precondition: want zero image models, got %d", len(res.ImageModels))
	}
	b, err := json.Marshal(res)
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(string(b), `"image_models":[]`) {
		t.Errorf("served %s\nwant \"image_models\":[] — null breaks concat in the UI", b)
	}
}

// TestDiscoveryResult_NilSlicesSerializeAsArrays covers the zero value of
// discoveryResult, which is what a caller gets if it builds one directly
// instead of going through DiscoverModels. DiscoverModels happens to
// initialize the slices, so testing only that path would let the MarshalJSON
// nil guard look redundant — and removing it would ship a latent UI bug the
// first time some other caller builds the struct.
func TestDiscoveryResult_NilSlicesSerializeAsArrays(t *testing.T) {
	b, err := json.Marshal(&discoveryResult{BaseURL: "https://h"})
	if err != nil {
		t.Fatal(err)
	}
	got := string(b)
	for _, want := range []string{`"models":[]`, `"image_models":[]`, `"tried":[]`} {
		if !strings.Contains(got, want) {
			t.Errorf("zero value serialized %s\nwant %s — a nil slice becomes JSON null and the UI renders \"undefined\"", got, want)
		}
	}
	for _, bad := range []string{`"models":null`, `"image_models":null`, `"tried":null`} {
		if strings.Contains(got, bad) {
			t.Errorf("serialized %s, which the UI cannot concatenate", bad)
		}
	}
}

// TestDiscoveryResultJSONShape pins the field names the WebUI will bind to.
func TestDiscoveryResultJSONShape(t *testing.T) {
	b, err := json.Marshal(&discoveryResult{
		BaseURL: "https://h", Models: []string{"a"}, ImageModels: []string{},
		ModelsPath: "/models", Adapter: "openai",
	})
	if err != nil {
		t.Fatal(err)
	}
	var got map[string]interface{}
	if err := json.Unmarshal(b, &got); err != nil {
		t.Fatal(err)
	}
	// image_models must serialize as [] and never null: the UI does
	// out.concat(image_models) and null would render "undefined".
	if v, ok := got["image_models"]; !ok || v == nil {
		t.Errorf("image_models = %v, want an empty array (null breaks concat in the UI)", v)
	}
	for _, k := range []string{"base_url", "endpoint", "models", "image_models", "models_path", "tried"} {
		if _, ok := got[k]; !ok {
			t.Errorf("missing field %q in discovery response", k)
		}
	}
}