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.
//
// Two details make the test actually exercise the rule instead of passing by
// accident, and both were found by a mutation surviving:
//
//  1. The memo is pre-seeded so openDesign runs FIRST. The memo remembers which
//     vendor answered last and that vendor is tried first; once a host has ever
//     answered through openDesign it runs first forever after. With the default
//     order (openDesign last) its 401 IS the last verdict, so even a plain
//     "last error wins" rule passes.
//  2. The other three adapters must ANSWER (404), not fail. They read `base`, not
//     `api_base`, so the source base_url points at the same stub — otherwise they
//     all error out on an unreachable host and the final verdict is an error
//     regardless of the rule being tested.
func TestQuotaOpenDesignRejectedControlKeyIsAnError(t *testing.T) {
	srv := odServer(t, nil, `{"error":"unauthenticated"}`, `{"error":"unauthenticated"}`,
		http.StatusUnauthorized, http.StatusUnauthorized)
	defer srv.Close()

	ps, _ := odPlugins(t)
	// Host key must match what hostOf(base) yields for the source below.
	src := map[string]interface{}{
		"name": "vela", "base_url": srv.URL, "api_key": "k",
	}
	if err := ps.SetState("quota", map[string]interface{}{
		"prices": map[string]interface{}{
			"control_key": "vela_ctrl_WRONG",
			"api_base":    srv.URL + "/api/v1",
		},
		// Flat, NOT nested under a "state" key: SetState pulls `prices` out and
		// stores the REST of the payload as the plugin's state. Nesting it would
		// store {state: {probe_memo: ...}} and the memo would never be read.
		"probe_memo": map[string]interface{}{
			hostOnly(srv.URL): "opendesign",
		},
	}); err != nil {
		t.Fatal(err)
	}

	out, err := ps.Call("quota", "collect_source", map[string]interface{}{
		"name": "vela", "sources": []interface{}{src},
	})
	if err != nil {
		t.Fatal(err)
	}
	rep, _ := out.(map[string]interface{})
	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 an error, which the page renders
	// as a red "读取失败" row. Filing it as a skip instead shows the same blank
	// card as "this vendor has no balance endpoint" — and sends 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 other three "+
			"adapters answered 404 for this host, so the final verdict is a skip unless "+
			"a rejected credential outranks it. The operator would go looking at the "+
			"vendor instead of re-running the login.", 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 probes the inference host, which serves no
// billing route at all, and every OpenDesign source silently skips.
//
// Asserted on the exported derivation rather than on the source text or on a
// failed probe's error message. The source-text version pinned the SPELLING of
// the rule: it went red the moment the rule was tightened from an `amr-link.`
// prefix to a full hostname, while the behaviour was unchanged. The
// error-message version inferred "which host did we ask?" from an allowlist
// refusal, which stops working the moment any adapter earlier in the chain also
// fails — and one of them always does.
func TestQuotaOpenDesignHostRewritingIsUsed(t *testing.T) {
	ps, _ := odPlugins(t)
	cases := []struct{ in, want string }{
		// The real production shape: base_url carries a path.
		{"https://amr-link.open-design.ai/v1", "https://amr-api.open-design.ai/api/v1"},
		// And without one, so the derivation cannot depend on the path being there.
		{"https://amr-link.open-design.ai", "https://amr-api.open-design.ai/api/v1"},
		// A trailing slash must not produce a double slash.
		{"https://amr-link.open-design.ai/v1/", "https://amr-api.open-design.ai/api/v1"},
	}
	for _, c := range cases {
		out, err := ps.Call("quota", "resolve_api_base", map[string]interface{}{
			"base_url": c.in,
		})
		if err != nil {
			t.Fatalf("resolve_api_base(%q): %v", c.in, err)
		}
		rep, _ := out.(map[string]interface{})
		if got, _ := rep["api_base"].(string); got != c.want {
			t.Errorf("resolve_api_base(%q) = %q, want %q — the probe would go to the "+
				"inference host, which serves no billing route, and the source would "+
				"silently stay unlisted", c.in, got, c.want)
		}
	}

	// A lookalike host must NOT be rewritten into an OpenDesign address: an
	// `amr-link.` prefix rule would turn amr-link.example.com into amr-api...,
	// sending an unrelated provider's probe to a host that is not its own.
	out, err := ps.Call("quota", "resolve_api_base", map[string]interface{}{
		"base_url": "https://amr-link.example.com/v1",
	})
	if err != nil {
		t.Fatal(err)
	}
	rep, _ := out.(map[string]interface{})
	if got, _ := rep["api_base"].(string); got != "https://amr-link.example.com/api/v1" {
		t.Errorf("a lookalike host was rewritten to %q", got)
	}

	// An explicit override wins (and is trimmed), so an operator can point at a
	// moved host without waiting for a code change.
	out, err = ps.Call("quota", "resolve_api_base", map[string]interface{}{
		"base_url": "https://amr-link.open-design.ai/v1",
		"api_base": "https://amr-api.example.net/api/v1/",
	})
	if err != nil {
		t.Fatal(err)
	}
	rep, _ = out.(map[string]interface{})
	if got, _ := rep["api_base"].(string); got != "https://amr-api.example.net/api/v1" {
		t.Errorf("api_base override not honoured (or not trimmed): %q", got)
	}
}

// TestQuotaUnreadyAdapterDoesNotSuppressAConfiguredHost is the end-to-end half of
// the unready rule, and it reproduces a bug seen on the live plugin.
//
// The negative memo means "this HOST exposes no balance endpoint" and is applied
// UNCONDITIONALLY, before any adapter runs. A host judged while openDesign had no
// prices.control_key therefore got memoised as having nothing and stayed
// suppressed after the key was supplied — observed live: unset key → memo → key
// set → probe skipped without a single request, so the fix the operator had just
// applied looked like it did nothing.
//
// It is NOT fixable by "do not memoise a skip": on a real OpenDesign host the
// other three adapters answer 404 ("no dashboard here"), which is durable and
// worth remembering, so the verdict would be recorded either way. The memo has to
// carry WHICH vendors were unable to judge the host, and the read side has to
// retire that record once they can.
func TestQuotaUnreadyAdapterDoesNotSuppressAConfiguredHost(t *testing.T) {
	var seenAuth []string
	srv := odServer(t, &seenAuth, odWallet, odSummary, 200, 200)
	defer srv.Close()

	ps, _ := odPlugins(t)
	args := map[string]interface{}{
		"name": "vela",
		"sources": []interface{}{map[string]interface{}{
			"name": "vela", "base_url": srv.URL + "/v1", "api_key": "k",
		}},
	}

	// Step 1: no control_key. openDesign cannot judge the host, and no billing
	// request goes anywhere.
	out, err := ps.Call("quota", "collect_source", args)
	if err != nil {
		t.Fatal(err)
	}
	rep, _ := out.(map[string]interface{})
	if rep["skipped"] != true {
		t.Fatalf("expected a skip with no control key configured: %v", rep)
	}
	if len(seenAuth) != 0 {
		t.Fatalf("an unconfigured adapter made a request: %q", seenAuth)
	}

	// Step 2: configure the key. The very next probe must reach the upstream. It
	// is inside the backoff window (minutes of wall clock) while this runs in
	// microseconds — so a host refused now is refused for the whole window in
	// production, which is the bug, not a timing accident.
	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)
	}
	out, err = ps.Call("quota", "collect_source", args)
	if err != nil {
		t.Fatal(err)
	}
	rep, _ = out.(map[string]interface{})
	if rep["vendor"] != "opendesign" || rep["skipped"] == true {
		t.Fatalf("expected an opendesign balance once configured, got %v — the host "+
			"is still suppressed by the provisional memo it earned while the adapter "+
			"had no credential", rep)
	}
	if len(seenAuth) == 0 {
		t.Fatal("no request reached the billing endpoint after the control key was " +
			"supplied")
	}
}