ModelRouter/internal/gateway/billing_rules_api.go

package gateway

import (
	"encoding/json"
	"net/http"

	"llmsproxy/internal/billing"
	"llmsproxy/internal/config"
)

// Billing rules API.
//
//	GET  /api/plugins/billing/rules    the active profile's rules + all profiles
//	PUT  /api/plugins/billing/rules    replace the whole DSL and re-inject prices
//	POST /api/plugins/billing/rules    { profile, rule } append one rule
//	DELETE /api/plugins/billing/rules  { profile, url } remove one rule
//
// Why a separate endpoint instead of the generic plugin state: prices are
// plugin state, but they are also a durable, reviewable CONFIGURATION. Routing
// them through /state would let an admin key PUT arbitrary plugin state and
// silently reset the accumulated accounting. Here the server owns the shape:
// it validates the DSL, writes it back to config.yaml, recompiles, and hands
// the plugin its prices table. The operator edits rules; the plugin keeps
// numbers. The two are never mixed in one payload.
//
// Editing means "what the operator typed is what config.yaml holds". A
// billing mistake found next week must be traceable to a reviewable file, not
// to a blob in the plugin's state sidecar.
func (g *Gateway) handleBillingRules(w http.ResponseWriter, r *http.Request) {
	if reqRole(r.Context()) != "admin" {
		writeError(w, http.StatusForbidden, "forbidden", "admin role required")
		return
	}
	switch r.Method {
	case http.MethodGet:
		g.getBillingRules(w)
	case http.MethodPut:
		g.putBillingRules(w, r)
	case http.MethodPost:
		g.addBillingRule(w, r)
	case http.MethodDelete:
		g.delBillingRule(w, r)
	default:
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET/PUT/POST/DELETE")
	}
}

// billingRulesPayload is the UI-facing view: the profiles as declared, plus
// which one is active and the source URLs they can match — the editor needs
// the URL list to offer suggestions, and an operator typing a URL that matches
// no source is exactly the silent-no-pricing failure this feature exists to
// remove.
type billingRulesPayload struct {
	DSL      *config.BillingDSL `json:"billing"`
	Active   string             `json:"active"`
	URLs     []string           `json:"urls"`
	Warnings []string           `json:"warnings"`
}

func (g *Gateway) billingRulesPayload() billingRulesPayload {
	out := billingRulesPayload{}
	cfg := g.core.Config()
	if cfg != nil && cfg.BillingDSL != nil {
		// Copy so the response cannot be mutated back through the pointer.
		cp := *cfg.BillingDSL
		out.DSL = &cp
		if p := cp.Resolve(""); p != nil {
			out.Active = p.ID
		}
	}
	for _, s := range g.sourceURLs() {
		out.URLs = append(out.URLs, s)
	}
	out.Warnings = g.billingRuleWarnings()
	return out
}

func (g *Gateway) sourceURLs() []string {
	cfg := g.core.Config()
	if cfg == nil {
		return nil
	}
	seen := map[string]bool{}
	var out []string
	for _, s := range cfg.Sources {
		if s.BaseURL == "" || seen[s.BaseURL] {
			continue
		}
		seen[s.BaseURL] = true
		out = append(out, s.BaseURL)
	}
	return out
}

// billingRuleWarnings reports rules that match no configured source. This is
// the silent killer of hand-written price tables: the rule parses, validates,
// and prices nothing at all, so every request on that URL lands in unpriced.
// Saying so out loud is the difference between a five-second fix and an
// afternoon wondering why the bill is zero.
func (g *Gateway) billingRuleWarnings() []string {
	cfg := g.core.Config()
	if cfg == nil || cfg.BillingDSL == nil {
		return nil
	}
	urls := map[string]bool{}
	for _, u := range g.sourceURLs() {
		urls[u] = true
	}
	var warns []string
	for _, p := range cfg.BillingDSL.Profiles {
		for _, r := range p.Rules {
			if r.URL == "*" || urls[r.URL] {
				continue
			}
			warns = append(warns, "profile "+p.ID+": rule url "+r.URL+" matches no configured source")
		}
	}
	return warns
}

func (g *Gateway) getBillingRules(w http.ResponseWriter) {
	writeJSON(w, http.StatusOK, g.billingRulesPayload())
}

func (g *Gateway) putBillingRules(w http.ResponseWriter, r *http.Request) {
	var body struct {
		Billing *config.BillingDSL `json:"billing"`
	}
	if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
		return
	}
	if body.Billing == nil {
		writeError(w, http.StatusBadRequest, "invalid_request", "billing is required")
		return
	}
	g.applyBillingDSLUpdate(w, body.Billing)
}

func (g *Gateway) addBillingRule(w http.ResponseWriter, r *http.Request) {
	var body struct {
		Profile string             `json:"profile"`
		Rule    config.BillingRule `json:"rule"`
	}
	if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
		return
	}
	cfg := g.core.Config()
	if cfg == nil || cfg.BillingDSL == nil || len(cfg.BillingDSL.Profiles) == 0 {
		writeError(w, http.StatusBadRequest, "no_profile", "no billing profiles configured; create one first")
		return
	}
	next := cloneBillingDSL(cfg.BillingDSL)
	idx := profileIndex(next, body.Profile)
	if idx < 0 {
		writeError(w, http.StatusBadRequest, "unknown_profile", "no such profile: "+body.Profile)
		return
	}
	p := &next.Profiles[idx]
	// Replace rather than append when the URL is already declared: two rules
	// for one URL would silently make the first unreachable (first match wins),
	// which is the exact ambiguity an editor should not be able to create.
	replaced := false
	for i := range p.Rules {
		if p.Rules[i].URL == body.Rule.URL {
			p.Rules[i] = body.Rule
			replaced = true
			break
		}
	}
	if !replaced {
		p.Rules = append(p.Rules, body.Rule)
	}
	g.applyBillingDSLUpdate(w, next)
}

func (g *Gateway) delBillingRule(w http.ResponseWriter, r *http.Request) {
	var body struct {
		Profile string `json:"profile"`
		URL     string `json:"url"`
	}
	if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
		return
	}
	cfg := g.core.Config()
	if cfg == nil || cfg.BillingDSL == nil {
		writeError(w, http.StatusBadRequest, "no_profile", "no billing profiles configured")
		return
	}
	next := cloneBillingDSL(cfg.BillingDSL)
	idx := profileIndex(next, body.Profile)
	if idx < 0 {
		writeError(w, http.StatusBadRequest, "unknown_profile", "no such profile: "+body.Profile)
		return
	}
	p := &next.Profiles[idx]
	out := p.Rules[:0]
	found := false
	for _, rule := range p.Rules {
		if rule.URL == body.URL {
			found = true
			continue
		}
		out = append(out, rule)
	}
	if !found {
		writeError(w, http.StatusNotFound, "not_found", "no rule for url "+body.URL)
		return
	}
	p.Rules = out
	g.applyBillingDSLUpdate(w, next)
}

// applyBillingDSLUpdate is the single write path: validate, persist to
// config.yaml, recompile, inject, and only then report success. If the config
// cannot be written the change is NOT applied in memory either — a rules editor
// that says "saved" and loses the edit on the next restart is worse than one
// that refuses.
func (g *Gateway) applyBillingDSLUpdate(w http.ResponseWriter, next *config.BillingDSL) {
	if err := next.Validate(); err != nil {
		writeError(w, http.StatusBadRequest, "invalid_rules", err.Error())
		return
	}
	cfg := g.core.Config()
	if cfg == nil {
		writeError(w, http.StatusInternalServerError, "no_config", "no config loaded")
		return
	}
	// Assign BEFORE saving. The previous order saved first and assigned after,
	// so every rule edit reported success while writing a config.yaml with no
	// billing section at all — the operator's edit vanished on restart and
	// nothing in the API response said so.
	prev := cfg.BillingDSL
	cfg.BillingDSL = next

	// Compile before persisting: a profile that cannot be turned into prices
	// must not reach the config file, or the next restart fails to load the
	// very rules the editor just accepted.
	prices, err := billing.CompileOpts(next.Resolve(next.Active), cfg.Sources, true)
	if err != nil {
		cfg.BillingDSL = prev // no half-applied state
		writeError(w, http.StatusBadRequest, "compile_failed", err.Error())
		return
	}
	if err := cfg.Save(); err != nil {
		cfg.BillingDSL = prev
		writeError(w, http.StatusInternalServerError, "persist_failed", err.Error())
		return
	}
	ps := g.core.Plugins()
	if ps != nil {
		if err := ps.SetState("billing", map[string]interface{}{"prices": prices}); err != nil {
			writeError(w, http.StatusBadRequest, "plugin_error", err.Error())
			return
		}
	}
	p := g.billingRulesPayload()
	writeJSON(w, http.StatusOK, map[string]interface{}{
		"ok": true, "billing": p.DSL, "active": p.Active, "warnings": p.Warnings,
	})
}

func cloneBillingDSL(d *config.BillingDSL) *config.BillingDSL {
	out := &config.BillingDSL{Active: d.Active}
	out.Profiles = make([]config.BillingProfile, len(d.Profiles))
	copy(out.Profiles, d.Profiles)
	return out
}

func profileIndex(d *config.BillingDSL, id string) int {
	if id == "" {
		return -1
	}
	for i, p := range d.Profiles {
		if p.ID == id {
			return i
		}
	}
	return -1
}