ModelRouter/internal/lua/quota_opendesign_test.go

package lua

import (
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

// The OpenDesign / Vela adapter is the first quota probe that needs a credential
// the source itself does not carry. Inference runs on the source's `api_key`
// (the runtime key); the two billing endpoints on amr-api.open-design.ai only
// accept the `controlKey` the desktop app receives at sign-in.
//
// That asymmetry is why these tests exist: the obvious wrong implementation —
// "reuse the source's api_key like every other adapter" — returns 401 against
// the live service, and a test that only stubs a happy response would never
// notice.
//
// Live responses this fixture copies (both 200, ~1.3s):
//   GET https://amr-api.open-design.ai/api/v1/wallet/balance
//   GET https://amr-api.open-design.ai/api/v1/billing/summary
const (
	odWallet = `{"balanceUsd":"15.7062","expiresAt":"2026-10-08T23:01:06.662Z",` +
		`"updatedAt":"2026-10-06T08:02:20.167Z","billingScopeVersion":2,` +
		`"workspaceId":"h7cvgf86l79yy3cpmsbmjfxu",` +
		`"workspaceMemberId":"uxqzu1rrtwsmzfpwugkcuany"}`

	odSummary = `{"balanceUsd":"15.7062","creditsPerUsd":10000,` +
		`"balances":{"subscriptionCredits":"152062","rechargeCredits":"5000",` +
		`"totalAvailableCredits":"157062"},"subscriptionCreditsUsd":"15.2062",` +
		`"rechargeCreditsUsd":"0.5000","totalAvailableCreditsUsd":"15.7062",` +
		`"membershipTier":"plus","billingInterval":"monthly",` +
		`"subscriptionStatus":"active","subscriptionEntitlementStatus":"active",` +
		`"subscriptionCancelAtPeriodEnd":false,` +
		`"subscriptionCurrentPeriodStart":"2026-10-05T23:01:54.000Z",` +
		`"subscriptionCurrentPeriodEnd":"2026-11-05T23:01:54.000Z",` +
		`"pendingSubscriptionChange":null,"firstMonthIntroEligible":false,` +
		`"availableActions":["subscription_checkout","billing_portal"],` +
		`"totalRechargedUsd":"0.5000","totalConsumedUsd":"4.7938",` +
		`"todayConsumedUsd":"0.0036","rechargeCount":1,"usageCount":2982,` +
		`"updatedAt":"2026-10-06T14:55:00.407Z"}`
)

// odServer stands in for amr-api. It records the Authorization header it saw on
// the two BILLING routes so a test can assert WHICH credential the plugin sent
// — the whole point of this adapter.
//
// Only the billing paths are recorded. The other three adapters derive their
// paths from the source's base_url, not from api_base, so in these tests they
// are pointed at amr-link and refused by the outbound allowlist before they can
// reach this server. Recording every request would make the "did it probe at
// all?" assertion depend on that coincidence; recording the billing routes makes
// it mean exactly what it says.
func odServer(t *testing.T, seenAuth *[]string, wallet, summary string, wStatus, sStatus int) *httptest.Server {
	t.Helper()
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch {
		case strings.HasSuffix(r.URL.Path, "/wallet/balance"):
			if seenAuth != nil {
				*seenAuth = append(*seenAuth, r.Header.Get("Authorization"))
			}
			w.WriteHeader(wStatus)
			_, _ = w.Write([]byte(wallet))
		case strings.HasSuffix(r.URL.Path, "/billing/summary"):
			if seenAuth != nil {
				*seenAuth = append(*seenAuth, r.Header.Get("Authorization"))
			}
			w.WriteHeader(sStatus)
			_, _ = w.Write([]byte(summary))
		default:
			w.WriteHeader(http.StatusNotFound)
			_, _ = w.Write([]byte(`{"error":"not found"}`))
		}
	}))
}

// odPlugins boots a real VM with the shipped quota.lua (not a copy) and returns
// the loaded Plugins.
func odPlugins(t *testing.T) (*Plugins, string) {
	t.Helper()
	dir := t.TempDir()
	pluginDir := filepath.Join(dir, "plugins")
	if err := os.MkdirAll(pluginDir, 0o755); err != nil {
		t.Fatal(err)
	}
	vm := NewVM(filepath.Join(dir, "adapters"))
	if err := vm.Start(); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(vm.Stop)

	src := readShippedPlugin(t, "quota.lua")
	if err := os.WriteFile(filepath.Join(pluginDir, "quota.lua"), []byte(src), 0o644); err != nil {
		t.Fatal(err)
	}
	ps := NewPlugins(vm, pluginDir)
	t.Cleanup(ps.Close)
	// Set before LoadDir: the allowlist is consulted when the plugin makes a
	// request, and the empty default denies everything.
	ps.SetOutboundHosts([]string{"127.0.0.1"})
	if err := ps.LoadDir(); err != nil {
		t.Fatal(err)
	}
	p := ps.find("quota")
	if p == nil {
		t.Fatal("quota plugin not found")
	}
	if p.LoadError != "" {
		t.Fatalf("quota.lua failed to compile: %s", p.LoadError)
	}
	return ps, src
}

// collectOpenDesign drives the plugin's own collect_source for a source that
// looks like vela. The plugin iterates ADAPTERS in order and picks the host's
// remembered vendor, so the first adapter that answers is the one under test.
func collectOpenDesign(t *testing.T, ps *Plugins, base string) map[string]interface{} {
	t.Helper()
	out, err := ps.Call("quota", "collect_source", map[string]interface{}{
		"name": "vela",
		"sources": []interface{}{
			map[string]interface{}{
				"name":     "vela",
				"base_url": base,
				"api_key":  "sk-amr-runtime-key-must-not-be-sent-to-billing",
			},
		},
	})
	if err != nil {
		t.Fatalf("collect_source failed: %v", err)
	}
	rep, _ := out.(map[string]interface{})
	return rep
}

// TestQuotaOpenDesignUsesControlKeyNotRuntimeKey is the reason this adapter
// cannot be written like the others.
//
// Mutation: replace `request(api .. "/wallet/balance", ck, ...)` with the
// function's `key` parameter (`request(api .. "/wallet/balance", key, ...)`)
// and this fails — the server records the runtime key.
func TestQuotaOpenDesignUsesControlKeyNotRuntimeKey(t *testing.T) {
	var seenAuth []string
	srv := odServer(t, &seenAuth, odWallet, odSummary, 200, 200)
	defer srv.Close()

	ps, _ := odPlugins(t)
	// The plugin's config table is `prices`; SetState routes a `prices` key
	// there rather than into accumulated state.
	if err := ps.SetState("quota", map[string]interface{}{
		"prices": map[string]interface{}{
			"control_key": "vela_ctrl_TESTKEY",
			"api_base":    srv.URL + "/api/v1",
		},
	}); err != nil {
		t.Fatalf("SetState: %v", err)
	}

	rep := collectOpenDesign(t, ps, "https://amr-link.open-design.ai/v1")
	if rep == nil {
		t.Fatal("no report produced")
	}
	if rep["vendor"] != "opendesign" {
		t.Fatalf("wrong adapter answered: vendor=%v (report=%v)", rep["vendor"], rep)
	}
	for _, got := range seenAuth {
		if strings.Contains(got, "runtime-key") {
			t.Fatalf("the source's api_key was sent to the billing endpoint (%q); "+
				"amr-api rejects it with 401 — the control key is a separate credential", got)
		}
		if got != "Bearer vela_ctrl_TESTKEY" {
			t.Fatalf("billing request carried %q, want the configured control key", got)
		}
	}
	if len(seenAuth) == 0 {
		t.Fatal("no request reached the server: the adapter never probed")
	}
}

// TestQuotaOpenDesignSkipsWithoutControlKey pins the failure mode that made the
// previous incarnation of this probe useless: with no credential the adapter
// must SKIP (not error, not report a zero balance). An error would put the
// source in the negative memo and file it under "probe failed"; a zero balance
// would tell the operator the account is empty.
//
// This test originally left api_base unset, which sent the probe to the REAL
// amr-api host. Every path then failed on the network and returned nil, so the
// test passed for the wrong reason — and a mutation that deleted the credential
// guard survived. Pointing api_base at a local stub makes the server reachable,
// so the only thing that can stop the request is the guard itself.
func TestQuotaOpenDesignSkipsWithoutControlKey(t *testing.T) {
	var seenAuth []string
	srv := odServer(t, &seenAuth, odWallet, odSummary, 200, 200)
	defer srv.Close()

	ps, _ := odPlugins(t)
	// api_base set, control_key deliberately NOT set.
	if err := ps.SetState("quota", map[string]interface{}{
		"prices": map[string]interface{}{"api_base": srv.URL + "/api/v1"},
	}); err != nil {
		t.Fatal(err)
	}

	rep := collectOpenDesign(t, ps, "https://amr-link.open-design.ai/v1")
	if rep == nil {
		t.Fatal("no report produced")
	}
	// The guard must come BEFORE the network. A reachable server that is never
	// asked anything is the only evidence that the credential check ran at all.
	if len(seenAuth) != 0 {
		t.Fatalf("the adapter probed the billing host with no control key configured "+
			"(Authorization: %q); an empty bearer is not a credential and the "+
			"response cannot be trusted", seenAuth)
	}
	if rep["skipped"] != true {
		t.Fatalf("skipped = %#v, want true", rep["skipped"])
	}
	reason, _ := rep["reason"].(string)
	if !strings.Contains(reason, "control_key") {
		t.Fatalf("reason = %q, want it to name the missing control_key so the "+
			"operator knows what to go and configure", reason)
	}
	// Whatever adapter answers, it must not be a fabricated opendesign balance.
	if rep["vendor"] == "opendesign" {
		t.Fatalf("opendesign reported a balance with no control key configured: %v", rep)
	}
}

// TestQuotaOpenDesignReportsBalanceAndSpend checks the fields the operator
// actually reads. The vendor sends these as JSON STRINGS ("15.7062"), so a
// naive type assertion to number would drop every one of them.
func TestQuotaOpenDesignReportsBalanceAndSpend(t *testing.T) {
	srv := odServer(t, nil, odWallet, odSummary, 200, 200)
	defer srv.Close()

	ps, _ := odPlugins(t)
	if err := ps.SetState("quota", map[string]interface{}{
		"prices": map[string]interface{}{
			"control_key": "vela_ctrl_TESTKEY",
			"api_base":    srv.URL + "/api/v1",
		},
	}); err != nil {
		t.Fatal(err)
	}
	rep := collectOpenDesign(t, ps, "https://amr-link.open-design.ai/v1")
	if rep == nil || rep["vendor"] != "opendesign" {
		t.Fatalf("expected an opendesign report, got %v", rep)
	}

	want := map[string]float64{
		"balance":          15.7062,
		"subscription_usd": 15.2062,
		"recharge_usd":     0.5000,
		"today_used":       0.0036,
		"period_used":      4.7938,
	}
	for k, exp := range want {
		got, ok := rep[k].(float64)
		if !ok {
			t.Fatalf("%s missing or not a number: %#v (full report %v)", k, rep[k], rep)
		}
		if got != exp {
			t.Fatalf("%s = %v, want %v", k, got, exp)
		}
	}
	if rep["credits_per_usd"] != float64(10000) {
		t.Fatalf("credits_per_usd = %v, want 10000 (used for the credit split)", rep["credits_per_usd"])
	}
	if rep["tier"] != "plus" || rep["status"] != "active" {
		t.Fatalf("tier/status = %v/%v, want plus/active", rep["tier"], rep["status"])
	}
	if rep["balance_expires_at"] != "2026-10-08T23:01:06.662Z" {
		t.Fatalf("balance_expires_at = %v — a balance that expires silently is "+
			"worse than one that warns", rep["balance_expires_at"])
	}
	// usageCount arrives as a JSON number, not a string, and must survive too.
	if rep["usage_count"] != float64(2982) {
		t.Fatalf("usage_count = %#v, want 2982", rep["usage_count"])
	}
}

// TestQuotaOpenDesignHasNoFabricatedPeriodWindow pins the deliberately-absent
// field. `vela billing workspace-snapshot` answers billingState:"free" with NO
// codingPlan.windows for this account, so there is no period allowance to show.
//
// Mutation: add `rep.windows = {{...}}` (or any synthesized window) and this
// fails. The page renders windows as "X% remaining", so a fabricated one is a
// confident lie about the account.
func TestQuotaOpenDesignHasNoFabricatedPeriodWindow(t *testing.T) {
	srv := odServer(t, nil, odWallet, odSummary, 200, 200)
	defer srv.Close()

	ps, _ := odPlugins(t)
	if err := ps.SetState("quota", map[string]interface{}{
		"prices": map[string]interface{}{
			"control_key": "vela_ctrl_TESTKEY",
			"api_base":    srv.URL + "/api/v1",
		},
	}); err != nil {
		t.Fatal(err)
	}
	rep := collectOpenDesign(t, ps, "https://amr-link.open-design.ai/v1")
	if rep == nil || rep["vendor"] != "opendesign" {
		t.Fatalf("expected an opendesign report, got %v", rep)
	}
	if w, ok := rep["windows"]; ok && w != nil {
		t.Fatalf("windows = %#v, but this account has no Coding Plan "+
			"(billingState=free) — a synthesized window renders as a "+
			"percentage remaining that does not exist", w)
	}
}

// TestQuotaOpenDesignRejectedControlKeyIsAnError guards the diagnosis. A 401
// means the credential is wrong — a fact the operator must fix — so it must not
// be filed under "skip" alongside "this vendor has no balance endpoint". The two
// produce the same blank card but wildly different next actions.
func TestQuotaOpenDesignRejectedControlKeyIsAnError(t *testing.T) {
	srv := odServer(t, nil, `{"error":"unauthenticated"}`, `{"error":"unauthenticated"}`,
		http.StatusUnauthorized, http.StatusUnauthorized)
	defer srv.Close()

	ps, _ := odPlugins(t)
	if err := ps.SetState("quota", map[string]interface{}{
		"prices": map[string]interface{}{
			"control_key": "vela_ctrl_WRONG",
			"api_base":    srv.URL + "/api/v1",
		},
	}); err != nil {
		t.Fatal(err)
	}
	rep := collectOpenDesign(t, ps, "https://amr-link.open-design.ai/v1")
	if rep == nil {
		t.Fatal("no report produced")
	}
	if rep["vendor"] == "opendesign" {
		t.Fatalf("a rejected control key must not produce a balance: %v", rep)
	}
	// The plugin reports a credential failure as kind="error", which the page
	// renders as a red "读取失败" row. Filing it as a skip instead would show the
	// same blank card as "this vendor has no balance endpoint" — and send the
	// operator to investigate the vendor rather than re-run the login.
	if rep["skipped"] == true {
		t.Fatalf("a rejected control key was reported as a SKIP: %v — the two "+
			"require different fixes and must not look alike", rep)
	}
	msg, _ := rep["error"].(string)
	if !strings.Contains(msg, "rejected") {
		t.Fatalf("error = %q, want it to name the rejected credential so the "+
			"operator knows to re-run the login rather than to investigate the vendor", msg)
	}
}

// TestQuotaOpenDesignDoesNotDefaultMissingRate pins the "do not invent an
// exchange rate" rule. creditsPerUsd is the only thing that converts the credit
// breakdown into dollars; when it is absent the vendor's own *Usd fields still
// give a correct balance, so the card must render with the split missing and
// FLAG it, not divide by a remembered 10000.
func TestQuotaOpenDesignDoesNotDefaultMissingRate(t *testing.T) {
	// Same summary with creditsPerUsd removed.
	noRate := strings.Replace(odSummary, `"creditsPerUsd":10000,`, "", 1)
	if strings.Contains(noRate, "creditsPerUsd") {
		t.Fatal("fixture still carries creditsPerUsd")
	}
	srv := odServer(t, nil, odWallet, noRate, 200, 200)
	defer srv.Close()

	ps, _ := odPlugins(t)
	if err := ps.SetState("quota", map[string]interface{}{
		"prices": map[string]interface{}{
			"control_key": "vela_ctrl_TESTKEY",
			"api_base":    srv.URL + "/api/v1",
		},
	}); err != nil {
		t.Fatal(err)
	}
	rep := collectOpenDesign(t, ps, "https://amr-link.open-design.ai/v1")
	if rep == nil || rep["vendor"] != "opendesign" {
		t.Fatalf("expected an opendesign report, got %v", rep)
	}
	// The balance is authoritative and comes straight from the vendor.
	if rep["balance"] != 15.7062 {
		t.Fatalf("balance = %#v, want 15.7062 (it must not depend on the rate)", rep["balance"])
	}
	if rep["credits_rate_missing"] != true {
		t.Fatalf("credits_rate_missing = %#v, want true: the page must show that "+
			"the credit split is unavailable rather than a silent zero", rep["credits_rate_missing"])
	}
	if v, ok := rep["credits_per_usd"]; ok && v != nil {
		t.Fatalf("credits_per_usd = %#v with no rate in the response — a defaulted "+
			"rate produces confidently wrong dollar amounts", v)
	}
	// subscription_usd/recharge_usd come from the vendor's own USD fields, so
	// they are still correct and must still be present.
	if rep["subscription_usd"] != 15.2062 {
		t.Fatalf("subscription_usd = %#v, want 15.2062 from the vendor's own field",
			rep["subscription_usd"])
	}
}

// TestQuotaOpenDesignHostRewritingIsUsed keeps the amr-link -> amr-api default
// working. Without it the adapter would probe the inference host, which answers
// no billing route at all, and every OpenDesign source would silently skip.
func TestQuotaOpenDesignHostRewritingIsUsed(t *testing.T) {
	var paths []string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		paths = append(paths, r.URL.Path)
		_, _ = w.Write([]byte(odWallet))
	}))
	defer srv.Close()

	ps, src := odPlugins(t)
	// api_base unset, so the adapter derives the host from base_url. The
	// derived host (amr-api.open-design.ai) is not reachable from a test, so
	// assert on the SOURCE instead — this is the guard that survives without a
	// network.
	_ = srv
	if !strings.Contains(src, `origin:gsub("^amr%-link%.", "amr-api.")`) {
		t.Fatal("quota.lua no longer rewrites amr-link.* to amr-api.* — every " +
			"OpenDesign source would probe the inference host, which serves no " +
			"billing route, and skip")
	}
	_ = ps
}