ModelRouter/internal/billing/wildcard_test.go

package billing

import (
	"testing"

	"llmsproxy/internal/config"
)

// The wildcard token rule must not blanket-mark every source as priced.
//
// Why this is a hazard and not a style question: priceFor() sets priced=true as
// soon as the request's SOURCE has an entry in the price table, and that flag
// is what keeps a request OUT of unpriced_reqs. A "*" rule is a model
// catalogue — it prices specific model ids wherever they come from — but
// compiling it into a {0,0,0} entry on all 21 sources made every model the
// catalogue does not list look "priced at zero". Paid providers (commandcode,
// alittokenplan, …) then stopped appearing in the unpriced warnings, which is
// exactly the silent zero-billing failure the plugin exists to surface.
func TestWildcardTokenRuleDoesNotMarkEverySourcePriced(t *testing.T) {
	sources := []config.Source{
		{Name: "free1", BaseURL: "https://free1.example.com/v1"},
		{Name: "paid1", BaseURL: "https://paid1.example.com/v1"},
		{Name: "paid2", BaseURL: "https://paid2.example.com/v1"},
	}
	profile := &config.BillingProfile{
		ID: "p",
		Rules: []config.BillingRule{
			{URL: "https://free1.example.com/v1", Mode: "free"},
			{
				URL: "*", Mode: "token",
				Models: map[string]config.BillingToken{
					"only-catalogued-model": {Prompt: "1.00", Completion: "2.00"},
				},
			},
		},
	}
	prices, err := Compile(profile, sources)
	if err != nil {
		t.Fatalf("compile: %v", err)
	}
	srcs := prices["sources"].(map[string]interface{})

	// The explicitly free source keeps its entry (free is a deliberate 0).
	if _, ok := srcs["free1"]; !ok {
		t.Error("the explicit free rule lost its source entry")
	}
	// Paid sources must NOT be marked, or their unlisted models bill $0.
	for _, name := range []string{"paid1", "paid2"} {
		if _, ok := srcs[name]; ok {
			t.Errorf("★ wildcard token rule marked %q as priced at 0 — its unlisted "+
				"models will bill $0 instead of appearing as unpriced", name)
		}
	}
	// The catalogue itself must still be present, or the rule prices nothing.
	models := prices["models"].(map[string]interface{})
	if _, ok := models["only-catalogued-model"]; !ok {
		t.Error("the wildcard rule did not publish its model price")
	}
}

// The inverse guard: a token rule for a SPECIFIC url still marks that source,
// because it genuinely declares pricing for the models served from it.
func TestSpecificURLTokenRuleStillMarksItsSource(t *testing.T) {
	sources := []config.Source{
		{Name: "mine", BaseURL: "https://mine.example.com/v1"},
		{Name: "other", BaseURL: "https://other.example.com/v1"},
	}
	profile := &config.BillingProfile{
		ID: "p",
		Rules: []config.BillingRule{{
			URL: "https://mine.example.com/v1", Mode: "token",
			Models: map[string]config.BillingToken{
				"m": {Prompt: "1.00", Completion: "2.00"},
			},
		}},
	}
	prices, err := Compile(profile, sources)
	if err != nil {
		t.Fatalf("compile: %v", err)
	}
	srcs := prices["sources"].(map[string]interface{})
	if _, ok := srcs["mine"]; !ok {
		t.Error("a specific-url token rule must still mark its own source priced")
	}
	if _, ok := srcs["other"]; ok {
		t.Error("an unrelated source was marked priced by a specific-url rule")
	}
}

// TestReasoningPriceIsOmittedUnlessSet pins the absence convention for the
// thinking-token rate.
//
// A rule that says nothing about reasoning must produce a price entry with NO
// reasoning_price key, so the billing plugin skips the repricing branch and
// bills thinking at the ordinary output rate — which is what every provider
// actually charges. Writing an explicit 0 instead would read as "thinking is
// free" and silently under-bill every request from that model.
func TestReasoningPriceIsOmittedUnlessSet(t *testing.T) {
	sources := []config.Source{{Name: "up", BaseURL: "https://up.example/v1"}}

	// Case 1: no reasoning declared → the key must not appear.
	p1 := &config.BillingProfile{ID: "p", Rules: []config.BillingRule{{
		URL: "https://up.example/v1", Mode: "token",
		Models: map[string]config.BillingToken{
			"m": {Prompt: "1.00", Completion: "2.00"},
		},
	}}}
	prices, err := Compile(p1, sources)
	if err != nil {
		t.Fatal(err)
	}
	entry := prices["models"].(map[string]interface{})["m"].(map[string]interface{})
	if _, present := entry["reasoning_price"]; present {
		t.Errorf("reasoning_price present (%v) without a rule; the plugin would treat 0 as free",
			entry["reasoning_price"])
	}

	// Case 2: declared → present, and converted to per-token like every other rate.
	p2 := &config.BillingProfile{ID: "p", Rules: []config.BillingRule{{
		URL: "https://up.example/v1", Mode: "token",
		Models: map[string]config.BillingToken{
			"m": {Prompt: "1.00", Completion: "2.00", Reasoning: "6.00"},
		},
	}}}
	prices2, err := Compile(p2, sources)
	if err != nil {
		t.Fatal(err)
	}
	e2 := prices2["models"].(map[string]interface{})["m"].(map[string]interface{})
	rp, ok := e2["reasoning_price"].(float64)
	if !ok {
		t.Fatalf("reasoning_price missing or wrong type: %#v", e2["reasoning_price"])
	}
	if rp != 6e-6 {
		t.Errorf("reasoning_price = %v, want 6e-6 (USD per TOKEN)", rp)
	}
	// A malformed rate must be rejected, not silently ignored: an operator who
	// typed "6 USD" expects either that value or an error, never a rule that
	// looks applied and bills nothing.
	p3 := &config.BillingProfile{ID: "p", Rules: []config.BillingRule{{
		URL: "https://up.example/v1", Mode: "token",
		Models: map[string]config.BillingToken{
			"m": {Prompt: "1.00", Completion: "2.00", Reasoning: "six dollars"},
		},
	}}}
	if _, err := Compile(p3, sources); err == nil {
		t.Error("a non-numeric reasoning rate must fail the compile")
	}
}