iop/packages/go/agentprovider/cli/status/quota.go

package status

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"math"
	"sort"
	"strconv"
	"strings"
	"time"
)

const (
	quotaSnapshotSchemaVersion = "1.0"
	quotaSnapshotSource        = "iop-agent-runtime quota-probe"

	quotaStateAvailable     = "available"
	quotaStateExhausted     = "exhausted"
	quotaStateUnknown       = "unknown"
	quotaStateNotApplicable = "not_applicable"

	quotaReasonCheckerError  = "checker_error"
	quotaReasonEvidence      = "cap_evidence_unknown"
	quotaReasonNotApplicable = "quota_not_applicable"
)

var durableQuotaReasonCodes = map[string]struct{}{
	quotaReasonCheckerError:  {},
	quotaReasonEvidence:      {},
	quotaReasonNotApplicable: {},
}

// QuotaCapView is the normalized, non-sensitive evidence for one required
// usage cap. It deliberately excludes UsageStatus.RawOutput.
type QuotaCapView struct {
	Name             string   `json:"name"`
	Status           string   `json:"status"`
	RemainingPercent *float64 `json:"remaining_percent"`
}

// QuotaTargetView is the selector-facing availability result for one target.
type QuotaTargetView struct {
	Adapter string `json:"adapter"`
	Target  string `json:"target"`
	Status  string `json:"status"`
}

// QuotaSnapshot is a narrow JSON bridge from the Go usage checker to the
// execution selector. It never serializes provider output or checker errors.
type QuotaSnapshot struct {
	SchemaVersion string            `json:"schema_version"`
	SnapshotID    string            `json:"snapshot_id"`
	Source        string            `json:"source"`
	CheckedAt     string            `json:"checked_at"`
	Targets       []QuotaTargetView `json:"targets"`
	RequiredCaps  []QuotaCapView    `json:"required_caps"`
	ReasonCodes   []string          `json:"reason_codes"`
}

// NormalizeQuotaSnapshot resolves the declared cap set against a parsed usage
// status. A confirmed exhausted cap wins over missing or malformed evidence so
// callers do not admit a target known to be exhausted.
func NormalizeQuotaSnapshot(adapter, target string, requiredCaps []string, checkedAt time.Time, usage *UsageStatus, checkErr error) QuotaSnapshot {
	views := make([]QuotaCapView, 0, len(requiredCaps))
	reasons := make([]string, 0, 1)
	if checkErr != nil || usage == nil {
		for _, cap := range requiredCaps {
			views = append(views, QuotaCapView{Name: cap, Status: quotaStateUnknown})
		}
		reasons = append(reasons, quotaReasonCheckerError)
	} else {
		for _, cap := range requiredCaps {
			views = append(views, resolveQuotaCap(usage, cap))
		}
	}

	status := quotaStateAvailable
	for _, view := range views {
		if view.Status == quotaStateExhausted {
			status = quotaStateExhausted
			break
		}
		if view.Status != quotaStateAvailable {
			status = quotaStateUnknown
		}
	}
	if len(views) == 0 {
		status = quotaStateNotApplicable
		reasons = []string{quotaReasonNotApplicable}
	}
	if status == quotaStateUnknown && len(reasons) == 0 {
		reasons = append(reasons, quotaReasonEvidence)
	}

	checked := checkedAt.UTC().Format(time.RFC3339Nano)
	snapshot := QuotaSnapshot{
		SchemaVersion: quotaSnapshotSchemaVersion,
		Source:        quotaSnapshotSource,
		CheckedAt:     checked,
		Targets:       []QuotaTargetView{{Adapter: adapter, Target: target, Status: status}},
		RequiredCaps:  views,
		ReasonCodes:   reasons,
	}
	snapshot.SnapshotID = quotaSnapshotID(snapshot)
	return snapshot
}

func resolveQuotaCap(usage *UsageStatus, required string) QuotaCapView {
	view := QuotaCapView{Name: required, Status: quotaStateUnknown}
	if usage == nil {
		return view
	}
	if required == "overall" {
		return quotaCapFromRemaining(required, usage.DailyLimit)
	}
	model, ok := strings.CutPrefix(required, "model:")
	if !ok || strings.TrimSpace(model) == "" {
		return view
	}
	count, err := strconv.Atoi(usage.Metadata["model_usage_count"])
	if err != nil || count < 0 {
		return view
	}
	matches := make([]string, 0, 1)
	for index := 0; index < count; index++ {
		prefix := fmt.Sprintf("model_usage_%d", index)
		if strings.EqualFold(strings.TrimSpace(usage.Metadata[prefix+"_name"]), strings.TrimSpace(model)) {
			matches = append(matches, usage.Metadata[prefix+"_used_percent"])
		}
	}
	if len(matches) != 1 {
		return view
	}
	used, ok := parsePercent(matches[0])
	if !ok {
		return view
	}
	remaining := 100 - used
	return quotaCapFromNumber(required, remaining)
}

func quotaCapFromRemaining(name, value string) QuotaCapView {
	remaining, ok := parsePercent(value)
	if !ok {
		return QuotaCapView{Name: name, Status: quotaStateUnknown}
	}
	return quotaCapFromNumber(name, remaining)
}

func quotaCapFromNumber(name string, remaining float64) QuotaCapView {
	view := QuotaCapView{Name: name, RemainingPercent: &remaining}
	if remaining == 0 {
		view.Status = quotaStateExhausted
	} else {
		view.Status = quotaStateAvailable
	}
	return view
}

func parsePercent(value string) (float64, bool) {
	value = strings.TrimSpace(strings.TrimSuffix(strings.TrimSpace(value), "%"))
	if value == "" {
		return 0, false
	}
	parsed, err := strconv.ParseFloat(value, 64)
	if err != nil || math.IsNaN(parsed) || math.IsInf(parsed, 0) || parsed < 0 || parsed > 100 {
		return 0, false
	}
	return parsed, true
}

// ValidateQuotaReasonCodes accepts only the bounded, secret-free reason
// registry emitted by NormalizeQuotaSnapshot.
func ValidateQuotaReasonCodes(reasons []string) error {
	seen := make(map[string]struct{}, len(reasons))
	for _, reason := range reasons {
		if _, ok := durableQuotaReasonCodes[reason]; !ok {
			return fmt.Errorf("quota snapshot contains an unsupported reason code")
		}
		if _, duplicate := seen[reason]; duplicate {
			return fmt.Errorf("quota snapshot repeats a reason code")
		}
		seen[reason] = struct{}{}
	}
	return nil
}

// ValidateQuotaSnapshot verifies the complete normalized snapshot shape and
// recomputes its content-bound identity. Callers must validate before
// projecting any snapshot into durable policy state.
func ValidateQuotaSnapshot(snapshot QuotaSnapshot) error {
	if snapshot.SchemaVersion != quotaSnapshotSchemaVersion {
		return fmt.Errorf("quota snapshot has an unsupported schema")
	}
	if snapshot.Source != quotaSnapshotSource {
		return fmt.Errorf("quota snapshot has an unsupported source")
	}
	checkedAt, err := time.Parse(time.RFC3339Nano, snapshot.CheckedAt)
	if err != nil || checkedAt.IsZero() ||
		checkedAt.UTC().Format(time.RFC3339Nano) != snapshot.CheckedAt {
		return fmt.Errorf("quota snapshot has a malformed checked time")
	}
	if len(snapshot.Targets) != 1 {
		return fmt.Errorf("quota snapshot must contain exactly one target")
	}
	target := snapshot.Targets[0]
	if !safeQuotaText(target.Adapter) || !safeQuotaText(target.Target) {
		return fmt.Errorf("quota snapshot has a malformed target identity")
	}
	if err := ValidateQuotaReasonCodes(snapshot.ReasonCodes); err != nil {
		return err
	}

	capNames := make(map[string]struct{}, len(snapshot.RequiredCaps))
	derivedStatus := quotaStateAvailable
	for _, cap := range snapshot.RequiredCaps {
		if !safeQuotaText(cap.Name) {
			return fmt.Errorf("quota snapshot has a malformed cap identity")
		}
		if _, duplicate := capNames[cap.Name]; duplicate {
			return fmt.Errorf("quota snapshot repeats a required cap")
		}
		capNames[cap.Name] = struct{}{}
		switch cap.Status {
		case quotaStateAvailable:
			if !validRemaining(cap.RemainingPercent) || *cap.RemainingPercent <= 0 {
				return fmt.Errorf("quota snapshot has invalid available cap evidence")
			}
		case quotaStateExhausted:
			if !validRemaining(cap.RemainingPercent) || *cap.RemainingPercent != 0 {
				return fmt.Errorf("quota snapshot has invalid exhausted cap evidence")
			}
			derivedStatus = quotaStateExhausted
		case quotaStateUnknown:
			if cap.RemainingPercent != nil {
				return fmt.Errorf("quota snapshot has invalid unknown cap evidence")
			}
			if derivedStatus != quotaStateExhausted {
				derivedStatus = quotaStateUnknown
			}
		default:
			return fmt.Errorf("quota snapshot has an unsupported cap state")
		}
	}

	switch {
	case len(snapshot.RequiredCaps) == 0:
		if target.Status != quotaStateNotApplicable ||
			!sameReasons(snapshot.ReasonCodes, []string{quotaReasonNotApplicable}) {
			return fmt.Errorf("quota snapshot has invalid not-applicable evidence")
		}
	case target.Status != derivedStatus:
		return fmt.Errorf("quota snapshot target state conflicts with cap evidence")
	case target.Status == quotaStateUnknown:
		if len(snapshot.ReasonCodes) != 1 ||
			(snapshot.ReasonCodes[0] != quotaReasonCheckerError &&
				snapshot.ReasonCodes[0] != quotaReasonEvidence) {
			return fmt.Errorf("quota snapshot has invalid unknown evidence")
		}
	case target.Status == quotaStateAvailable || target.Status == quotaStateExhausted:
		if len(snapshot.ReasonCodes) != 0 {
			return fmt.Errorf("quota snapshot has unexpected reason evidence")
		}
	default:
		return fmt.Errorf("quota snapshot has an unsupported target state")
	}

	if !safeQuotaText(snapshot.SnapshotID) ||
		snapshot.SnapshotID != quotaSnapshotID(snapshot) {
		return fmt.Errorf("quota snapshot identity does not match normalized content")
	}
	return nil
}

func quotaSnapshotID(snapshot QuotaSnapshot) string {
	caps := append([]QuotaCapView(nil), snapshot.RequiredCaps...)
	sort.Slice(caps, func(i, j int) bool {
		if caps[i].Name != caps[j].Name {
			return caps[i].Name < caps[j].Name
		}
		return caps[i].Status < caps[j].Status
	})
	reasons := append([]string(nil), snapshot.ReasonCodes...)
	sort.Strings(reasons)
	payload := struct {
		SchemaVersion string            `json:"schema_version"`
		Source        string            `json:"source"`
		CheckedAt     string            `json:"checked_at"`
		Targets       []QuotaTargetView `json:"targets"`
		Caps          []QuotaCapView    `json:"caps"`
		Reasons       []string          `json:"reasons"`
	}{
		SchemaVersion: snapshot.SchemaVersion,
		Source:        snapshot.Source,
		CheckedAt:     snapshot.CheckedAt,
		Targets:       append([]QuotaTargetView(nil), snapshot.Targets...),
		Caps:          caps,
		Reasons:       reasons,
	}
	encoded, _ := json.Marshal(payload)
	digest := sha256.Sum256(encoded)
	return "quota-" + hex.EncodeToString(digest[:])
}

func safeQuotaText(value string) bool {
	return value != "" && strings.TrimSpace(value) == value &&
		!strings.ContainsAny(value, "\x00\r\n")
}

func validRemaining(value *float64) bool {
	return value != nil && !math.IsNaN(*value) && !math.IsInf(*value, 0) &&
		*value >= 0 && *value <= 100
}

func sameReasons(left, right []string) bool {
	if len(left) != len(right) {
		return false
	}
	for index := range left {
		if left[index] != right[index] {
			return false
		}
	}
	return true
}