ModelRouter/internal/core/source_hosts_test.go

package core

import (
	"net/url"
	"testing"

	"llmsproxy/internal/config"
)

// The plugin outbound allowlist is built from the hosts of configured sources.
// That rule is correct for every adapter whose credential and endpoint live on
// the same host — which was all of them until OpenDesign.
//
// OpenDesign splits them: inference is amr-link.open-design.ai, billing is
// amr-api.open-design.ai. The quota plugin derives the billing host from the
// inference one, so the host it asks for is NOT the host the operator
// configured. The request is then refused by the allowlist with a message that
// reads like a plugin defect ("host ... is not an allowed outbound destination")
// rather than a missing entry.
//
// This file pins the bridge: when a source looks like OpenDesign, the derived
// billing host joins the allowlist.

func TestSourceHostsIncludesConfiguredHosts(t *testing.T) {
	cfg := &config.Config{Sources: []config.Source{
		{Name: "a", BaseURL: "https://api.deepseek.com"},
		{Name: "b", BaseURL: "https://amr-link.open-design.ai/v1"},
		{Name: "empty", BaseURL: ""},
		{Name: "bad", BaseURL: "://nonsense"},
	}}
	got := map[string]bool{}
	for _, h := range sourceHosts(cfg) {
		got[h] = true
	}
	for _, want := range []string{"api.deepseek.com", "amr-link.open-design.ai"} {
		if !got[want] {
			t.Errorf("sourceHosts missing configured host %q (allowlist would refuse "+
				"every plugin request to it): %v", want, got)
		}
	}
	// An empty/unparseable base_url must not abort the whole list.
	if !got["api.deepseek.com"] {
		t.Fatal("an unparseable sibling source dropped a valid host")
	}
}

// TestSourceHostsAddsDerivedBillingHost covers the OpenDesign case at the level
// that actually matters: the host the PLUGIN asks for must be in the allowlist.
func TestSourceHostsAddsDerivedBillingHost(t *testing.T) {
	cfg := &config.Config{Sources: []config.Source{
		{Name: "vela", BaseURL: "https://amr-link.open-design.ai/v1"},
	}}
	got := map[string]bool{}
	for _, h := range sourceHosts(cfg) {
		got[h] = true
	}
	if !got["amr-link.open-design.ai"] {
		t.Fatalf("inference host missing: %v", got)
	}
	if !got["amr-api.open-design.ai"] {
		t.Fatalf("the billing host the quota plugin derives is not in the allowlist, "+
			"so every OpenDesign balance probe is refused before it leaves the "+
			"process: %v", got)
	}
}

// A deployment with no OpenDesign source must not silently gain its billing
// host: an allowlist entry is a grant, and widening it for everyone would let an
// unrelated plugin reach a host the operator never configured.
func TestSourceHostsDoesNotGrantUnconfiguredBillingHost(t *testing.T) {
	cfg := &config.Config{Sources: []config.Source{
		{Name: "ds", BaseURL: "https://api.deepseek.com"},
	}}
	for _, h := range sourceHosts(cfg) {
		if h == "amr-api.open-design.ai" {
			t.Fatal("the OpenDesign billing host was granted to a deployment with " +
				"no OpenDesign source")
		}
	}
}

// The derivation has to agree with the Lua side. Keep the two spellings of the
// rule in one place so a change to one is visible when reading the other.
func TestSourceHostsDerivationMatchesPluginRule(t *testing.T) {
	// what the Lua adapter does: origin:gsub("^amr%-link%.", "amr-api.")
	const in = "amr-link.open-design.ai"
	want := "amr-api.open-design.ai"
	got := deriveBillingHost(in)
	if got != want {
		t.Fatalf("deriveBillingHost(%q) = %q, want %q — the Go allowlist and the "+
			"Lua probe must agree on the host or the request is refused", in, got, want)
	}
	// A host that is not amr-link stays untouched, so this cannot accidentally
	// rewrite an unrelated provider into an OpenDesign address.
	for _, other := range []string{"api.deepseek.com", "amr-link.example.com", "xamr-link.a"} {
		if deriveBillingHost(other) != other {
			t.Errorf("deriveBillingHost(%q) = %q, want it unchanged", other, deriveBillingHost(other))
		}
	}
	// url.Parse round-trip keeps the helper honest about what it receives.
	u, err := url.Parse("https://amr-link.open-design.ai/v1")
	if err != nil {
		t.Fatal(err)
	}
	if deriveBillingHost(u.Hostname()) != "amr-api.open-design.ai" {
		t.Fatal("derivation failed on a real base_url hostname")
	}
}