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")
	}
}

// TestLongContextStepCompiles pins the long-context step's two accepted forms
// and its rejection of a step that would silently do nothing.
//
// Providers publish both shapes — DeepSeek says "2x past 128k", OpenAI quotes
// an absolute "$3.00/M past 200k" — so both must compile. Neither may be
// missing: a step with no rate is a rule that LOOKS like it handles long
// context and does not, which under-bills without any visible sign.
func TestLongContextStepCompiles(t *testing.T) {
	sources := []config.Source{{Name: "up", BaseURL: "https://up.example/v1"}}

	compile := func(lc config.BillingLongContext) (map[string]interface{}, error) {
		p := &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", LongContext: &lc},
			},
		}}}
		prices, err := Compile(p, sources)
		if err != nil {
			return nil, err
		}
		e := prices["models"].(map[string]interface{})["m"].(map[string]interface{})
		return e["long_context"].(map[string]interface{}), nil
	}

	// Absolute rate (OpenAI's published shape).
	lc1, err := compile(config.BillingLongContext{Above: 200000, Prompt: "3.00"})
	if err != nil {
		t.Fatalf("absolute long_context rate rejected: %v", err)
	}
	if lc1["above"] != 200000 || lc1["prompt"] != 3e-6 {
		t.Errorf("absolute step = %#v, want above=200000 prompt=3e-6", lc1)
	}

	// Multiplier (DeepSeek's shape): 2x of the 1.00/M base.
	lc2, err := compile(config.BillingLongContext{Above: 128000, Multiplier: 2})
	if err != nil {
		t.Fatalf("multiplier long_context rejected: %v", err)
	}
	if lc2["prompt"] != 2e-6 {
		t.Errorf("multiplier step = %#v, want prompt=2e-6", lc2)
	}

	// Both given: the absolute rate wins, because that is the number the
	// provider quotes. Multiplying instead would produce a figure the operator
	// never wrote.
	lc3, err := compile(config.BillingLongContext{Above: 128000, Prompt: "5.00", Multiplier: 2})
	if err != nil {
		t.Fatal(err)
	}
	if lc3["prompt"] != 5e-6 {
		t.Errorf("both given = %#v, want the quoted absolute rate 5e-6", lc3)
	}

	// Neither: a no-op step must fail, not compile into a rule that silently
	// bills long context at the base rate.
	if _, err := compile(config.BillingLongContext{Above: 128000}); err == nil {
		t.Error("a step with neither prompt nor multiplier must be rejected")
	}
	// A zero threshold would make every request "long".
	if _, err := compile(config.BillingLongContext{Above: 0, Multiplier: 2}); err == nil {
		t.Error("above=0 must be rejected")
	}
}