ModelRouter/internal/lua/billing_rules_test.go

package lua

import (
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

// billing 的规则**完全在插件内**：以源名为键、由插件自己编译、自己写
// `<plugin_dir>/billing.json`。内核没有 billing 代码，也不提供 billing 端点；
// 这些判据直接打插件的导出（对应插件自己的 /api/plugins/http/billing/rules）。
func putRules(t *testing.T, ps *Plugins, rules map[string]interface{}) map[string]interface{} {
	t.Helper()
	out, err := ps.Call("billing", "rules_put", map[string]interface{}{"billing": rules})
	if err != nil {
		t.Fatalf("rules_put: %v", err)
	}
	m, _ := out.(map[string]interface{})
	return m
}

func pricesOf(t *testing.T, ps *Plugins) map[string]interface{} {
	t.Helper()
	raw := ps.Prices("billing")
	p, ok := raw.(map[string]interface{})
	if !ok {
		t.Fatalf("billing has no prices: %#v", raw)
	}
	return p
}

func oneProfile(rules []interface{}, def interface{}) map[string]interface{} {
	p := map[string]interface{}{"id": "p", "rules": rules}
	if def != nil {
		p["default"] = def
	}
	return map[string]interface{}{"active": "p", "profiles": []interface{}{p}}
}

// 插件自己把规则写到 plugin_dir/billing.json，且读取回来一致。
func TestBillingRulesPersistInPluginOwnedFile(t *testing.T) {
	ps, pdir := billingVM(t)
	rules := oneProfile([]interface{}{
		map[string]interface{}{"source": "localzen", "mode": "free"},
	}, "unpriced")
	if m := putRules(t, ps, rules); m["ok"] != true {
		t.Fatalf("rules_put 失败: %#v", m)
	}
	path := filepath.Join(pdir, "billing.json")
	raw, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("插件没有自己写规则文件 %s: %v", path, err)
	}
	var onDisk map[string]interface{}
	if err := json.Unmarshal(raw, &onDisk); err != nil {
		t.Fatalf("规则文件不是 JSON: %v\n%s", err, raw)
	}
	if !strings.Contains(string(raw), "localzen") {
		t.Fatalf("规则文件内容不对: %s", raw)
	}
	// 重新加载插件（模拟重启）：规则应从文件恢复。
	ps2 := reloadBilling(t, pdir)
	got, err := ps2.Call("billing", "rules_get", map[string]interface{}{})
	if err != nil {
		t.Fatalf("rules_get: %v", err)
	}
	b, _ := json.Marshal(got)
	if !strings.Contains(string(b), "localzen") {
		t.Fatalf("重启后没有从 billing.json 恢复规则: %s", b)
	}
	// 编译结果也恢复了：free 源条目存在。
	if p := ps2.Prices("billing").(map[string]interface{}); p != nil {
		if _, ok := p["sources"].(map[string]interface{})["localzen"]; !ok {
			t.Errorf("恢复后价格表缺 localzen: %#v", p)
		}
	}
}

func reloadBilling(t *testing.T, pdir string) *Plugins {
	t.Helper()
	vm := NewVM(filepath.Join(t.TempDir(), "adapters"))
	if err := vm.Start(); err != nil {
		t.Fatalf("vm: %v", err)
	}
	t.Cleanup(vm.Stop)
	ps := NewPlugins(vm, pdir)
	if err := ps.LoadDir(); err != nil {
		t.Fatalf("load: %v", err)
	}
	return ps
}

// 编译语义：free/subscription 写"有意零价"，token 规则**不写源条目**（未列价
// 的模型要可见地进入 unpriced），通配符只发布模型价。
func TestBillingRulesCompileSemantics(t *testing.T) {
	ps, _ := billingVM(t)
	rules := oneProfile([]interface{}{
		map[string]interface{}{"source": "free-src", "mode": "free"},
		map[string]interface{}{"source": "sub-src", "mode": "subscription", "subscription": 20},
		map[string]interface{}{"source": "tok-src", "mode": "token", "models": map[string]interface{}{
			"m1": map[string]interface{}{"prompt": "0.15", "completion": "0.60"},
		}},
		map[string]interface{}{"url": "*", "mode": "token", "models": map[string]interface{}{
			"m2": map[string]interface{}{"prompt": "1", "completion": "2"},
		}},
	}, nil)
	if m := putRules(t, ps, rules); m["ok"] != true {
		t.Fatalf("rules_put: %#v", m)
	}
	p := pricesOf(t, ps)
	srcs := p["sources"].(map[string]interface{})
	for _, n := range []string{"free-src", "sub-src"} {
		e, ok := srcs[n].(map[string]interface{})
		if !ok || e["prompt"] != 0.0 || e["completion"] != 0.0 {
			t.Errorf("%s 应写有意零价: %#v", n, srcs[n])
		}
	}
	if _, ok := srcs["tok-src"]; ok {
		t.Errorf("token 规则不得产生源条目: %#v", srcs["tok-src"])
	}
	if subs, ok := p["subscriptions"].(map[string]interface{})["sub-src"].(map[string]interface{}); !ok || subs["monthly"] != 20.0 {
		t.Errorf("subscription 固定费丢失: %#v", p["subscriptions"])
	}
	models := p["models"].(map[string]interface{})
	if models["m1"].(map[string]interface{})["prompt"] != 1.5e-7 {
		t.Errorf("m1 prompt = %#v, want 1.5e-7", models["m1"])
	}
	if _, ok := models["m2"]; !ok {
		t.Error("通配符规则的模型价丢失")
	}
}

// 端到端：token 规则 + 1e6/1e6 tokens @ 0.15/0.60 USD/M = 0.75，
// 且 unpriced 计数不涨（模型本身被标价）。
func TestBillingRulesTokenCostEndToEnd(t *testing.T) {
	ps, _ := billingVM(t)
	rules := oneProfile([]interface{}{
		map[string]interface{}{"source": "tok-src", "mode": "token", "models": map[string]interface{}{
			"m": map[string]interface{}{"prompt": "0.15", "completion": "0.60"},
		}},
	}, nil)
	if m := putRules(t, ps, rules); m["ok"] != true {
		t.Fatalf("rules_put: %#v", m)
	}
	ps.Fire(StageRequestEnd, map[string]interface{}{
		"model": "m", "source": "tok-src", "key": "***p", "ok": true,
		"prompt_tokens": 1000000.0, "completion_tokens": 1000000.0, "time": 1750000000000,
	})
	st := stateOf(t, ps)
	approx(t, "cost", st["total"].(map[string]interface{})["cost"].(float64), 0.75)
	if got := st["unpriced_reqs"].(float64); got != 0 {
		t.Errorf("unpriced_reqs = %v, want 0", got)
	}
	// 同源未列价的模型必须落入 unpriced（可见），而不是静默计 0。
	ps.Fire(StageRequestEnd, map[string]interface{}{
		"model": "unlisted", "source": "tok-src", "key": "***p", "ok": true,
		"prompt_tokens": 100.0, "completion_tokens": 1.0, "time": 1750000000000,
	})
	if got := stateOf(t, ps)["unpriced_reqs"].(float64); got != 1 {
		t.Errorf("未列价模型应计入 unpriced，got %v", got)
	}
}

// 非法规则必须被拒绝，且**不覆盖**已有文件（保存失败不能丢旧配置）。
func TestBillingRulesRejectInvalidWithoutClobberingFile(t *testing.T) {
	ps, pdir := billingVM(t)
	good := oneProfile([]interface{}{
		map[string]interface{}{"source": "s1", "mode": "free"},
	}, nil)
	if m := putRules(t, ps, good); m["ok"] != true {
		t.Fatalf("seed: %#v", m)
	}
	path := filepath.Join(pdir, "billing.json")
	before, _ := os.ReadFile(path)

	bad := oneProfile([]interface{}{
		map[string]interface{}{"source": "s1", "mode": "nope"},
	}, nil)
	out, err := ps.Call("billing", "rules_put", map[string]interface{}{"billing": bad})
	if err != nil {
		t.Fatalf("call: %v", err)
	}
	m := out.(map[string]interface{})
	if m["error"] == nil {
		t.Fatalf("非法 mode 应被拒绝: %#v", m)
	}
	after, _ := os.ReadFile(path)
	if string(before) != string(after) {
		t.Errorf("被拒绝的保存改动了规则文件:\nbefore=%s\nafter=%s", before, after)
	}
}

// reasoning 缺席 = 按输出价计；显式 long_context/peak 必须编译成插件读取的形状。
func TestBillingRulesOptionalPricesCompile(t *testing.T) {
	ps, _ := billingVM(t)
	rules := oneProfile([]interface{}{
		map[string]interface{}{"source": "s1", "mode": "token",
			"models": map[string]interface{}{
				"plain":     map[string]interface{}{"prompt": "1", "completion": "2"},
				"reasoning": map[string]interface{}{"prompt": "1", "completion": "2", "reasoning": "3"},
				"long": map[string]interface{}{"prompt": "1", "completion": "2",
					"long_context": map[string]interface{}{"above": 128000, "multiplier": 2}},
			},
			"peak": map[string]interface{}{"multiplier": 2, "hours": []interface{}{[]interface{}{1, 4}}},
		},
	}, nil)
	if m := putRules(t, ps, rules); m["ok"] != true {
		t.Fatalf("rules_put: %#v", m)
	}
	models := pricesOf(t, ps)["models"].(map[string]interface{})
	if _, present := models["plain"].(map[string]interface{})["reasoning_price"]; present {
		t.Error("未配置 reasoning 不应出现 reasoning_price（0 会被当成免费）")
	}
	if got := models["reasoning"].(map[string]interface{})["reasoning_price"]; got != 3e-6 {
		t.Errorf("reasoning_price = %v, want 3e-6", got)
	}
	lc := models["long"].(map[string]interface{})["long_context"].(map[string]interface{})
	if lc["above"] != 128000.0 || lc["prompt"] != 2e-6 {
		t.Errorf("long_context = %#v", lc)
	}
	pk := models["long"].(map[string]interface{})["peak"].(map[string]interface{})
	if pk["multiplier"] != 2.0 {
		t.Errorf("peak = %#v", pk)
	}
}