CLIProxyAPI/internal/runtime/executor/helps/codex_quota.go

package helps

import (
	"net/http"
	"strings"

	"github.com/tidwall/gjson"
)

const (
	codexRateLimitsEventType      = "codex.rate_limits"
	codexQuotaAdditionalHeaderKey = "X-Codex-Additional-"
	maxCodexAdditionalRateLimits  = 8
)

// ParseCodexQuotaEventHeaders converts one Codex websocket quota event into the
// same bounded header representation used by HTTP usage observations. Codex
// sends additional_rate_limits as an object on the websocket path, while the
// /wham/usage probe exposes the equivalent data as an array; both forms are
// normalized into namespaced X-Codex headers here.
//
// The websocket payload identifies each additional limit only by limit name
// ("GPT-5.3-Codex-Spark"), while HTTP responses namespace the same limit by a
// short name (x-codex-bengalfox-*). The two forms therefore cannot produce
// identical header names; the X-Codex-Additional- prefix marks the websocket
// origin, and the snapshot-replacement semantics in QuotaState keep the two
// spellings from accumulating side by side.
func ParseCodexQuotaEventHeaders(payload []byte) http.Header {
	if len(payload) == 0 {
		return nil
	}

	eventKind := codexQuotaEventKind(payload)
	if eventKind == codexQuotaEventOther {
		return nil
	}
	root := gjson.ParseBytes(payload)
	if eventKind == codexQuotaEventError {
		return parseCodexQuotaHeadersObject(root.Get("headers"))
	}

	headers := make(http.Header)
	hasQuotaData := false

	baseRateLimits := firstCodexQuotaResult(root, "rate_limits", "rateLimit")
	if addCodexQuotaRateLimitHeaders(headers, "X-Codex-", baseRateLimits) {
		// The main rate-limit object has no user-facing name in the websocket
		// payload, so its fields retain the HTTP response header names.
		hasQuotaData = true
	}

	additional := firstCodexQuotaResult(root, "additional_rate_limits", "additionalRateLimits")
	additionalCount := 0
	if additional.Exists() && (additional.IsObject() || additional.IsArray()) {
		additional.ForEach(func(key, value gjson.Result) bool {
			if additionalCount >= maxCodexAdditionalRateLimits {
				return false
			}
			limitName := strings.TrimSpace(key.String())
			if additional.IsArray() {
				limitName = firstCodexQuotaResultString(value, "limit_name", "limitName", "name")
			}
			if limitName == "" {
				return true
			}
			identifier := normalizeCodexQuotaHeaderIdentifier(limitName)
			if identifier == "" {
				return true
			}
			rateInfo := firstCodexQuotaResult(value, "rate_limit", "rateLimit")
			if !rateInfo.Exists() {
				rateInfo = value
			}
			prefix := codexQuotaAdditionalHeaderKey + identifier + "-"
			if addCodexQuotaRateLimitHeaders(headers, prefix, rateInfo) {
				// The limit name is upstream-controlled and lands in the
				// plain-text request log, so reject control characters here the
				// same way the plan type is validated below.
				if validCodexQuotaEventText(limitName) {
					headers.Set(prefix+"Limit-Name", limitName)
				}
				hasQuotaData = true
				additionalCount++
			}
			return true
		})
	}

	codeReview := firstCodexQuotaResult(root, "code_review_rate_limits", "codeReviewRateLimits")
	if addCodexQuotaRateLimitHeaders(headers, "X-Codex-Code-Review-", codeReview) {
		hasQuotaData = true
	}

	credits := firstCodexQuotaResult(root, "credits")
	if credits.Exists() && credits.IsObject() {
		hasQuotaData = setCodexQuotaScalarHeader(headers, "X-Codex-Credits-Has-Credits", credits, "has_credits", "hasCredits") || hasQuotaData
		hasQuotaData = setCodexQuotaScalarHeader(headers, "X-Codex-Credits-Unlimited", credits, "unlimited") || hasQuotaData
		hasQuotaData = setCodexQuotaScalarHeader(headers, "X-Codex-Credits-Balance", credits, "balance") || hasQuotaData
	}

	if !hasQuotaData {
		return nil
	}
	// A malformed active-limit name only invalidates that one header; the
	// window percentages collected above stay valid and must not be discarded.
	activeLimit := firstCodexQuotaResultString(root, "metered_limit_name", "meteredLimitName", "limit_name", "limitName")
	if validCodexQuotaEventIdentifier(activeLimit) {
		headers.Set("X-Codex-Active-Limit", activeLimit)
	}
	if planType := firstCodexQuotaResultString(root, "plan_type", "planType"); validCodexQuotaEventText(planType) {
		headers.Set("X-Codex-Plan-Type", planType)
	}
	return headers
}

const (
	codexQuotaEventOther = iota
	codexQuotaEventError
	codexQuotaEventRateLimits
)

// codexQuotaEventKind performs an allocation-free discriminator check for the
// common non-quota websocket event path. Codex frames carry the top-level
// type field within their opening bytes (after sequence_number at most), so
// only a bounded prefix is scanned instead of the whole frame; full JSON
// parsing is reserved for quota and error events.
func codexQuotaEventKind(payload []byte) int {
	const scanLimit = 256
	end := len(payload)
	if end > scanLimit {
		end = scanLimit
	}
	for i := 0; i+6 < end; i++ {
		if payload[i] != '"' || payload[i+1] != 't' || payload[i+2] != 'y' || payload[i+3] != 'p' || payload[i+4] != 'e' || payload[i+5] != '"' {
			continue
		}
		j := i + 6
		for j < end && (payload[j] == ' ' || payload[j] == '\t' || payload[j] == '\r' || payload[j] == '\n') {
			j++
		}
		if j >= end || payload[j] != ':' {
			continue
		}
		j++
		for j < end && (payload[j] == ' ' || payload[j] == '\t' || payload[j] == '\r' || payload[j] == '\n') {
			j++
		}
		if j+7 <= end && payload[j] == '"' && payload[j+1] == 'e' && payload[j+2] == 'r' && payload[j+3] == 'r' && payload[j+4] == 'o' && payload[j+5] == 'r' && payload[j+6] == '"' {
			return codexQuotaEventError
		}
		const rateLimits = `"codex.rate_limits"`
		if j+len(rateLimits) <= end {
			matched := true
			for offset := range rateLimits {
				if payload[j+offset] != rateLimits[offset] {
					matched = false
					break
				}
			}
			if matched {
				return codexQuotaEventRateLimits
			}
		}
	}
	return codexQuotaEventOther
}

func addCodexQuotaRateLimitHeaders(headers http.Header, prefix string, rateInfo gjson.Result) bool {
	if !rateInfo.Exists() || !rateInfo.IsObject() {
		return false
	}

	changed := false
	if setCodexQuotaScalarHeaderFromResult(headers, prefix+"Allowed", rateInfo, "allowed") {
		changed = true
	}
	if setCodexQuotaScalarHeaderFromResult(headers, prefix+"Limit-Reached", rateInfo, "limit_reached", "limitReached") {
		changed = true
	}
	for _, windowName := range []string{"primary", "secondary"} {
		window := firstCodexQuotaResult(rateInfo, windowName)
		if !window.Exists() || !window.IsObject() {
			continue
		}
		used := firstCodexQuotaResult(window, "used_percent", "usedPercent")
		minutes := firstCodexQuotaResult(window, "window_minutes", "windowMinutes")
		resetAfter := firstCodexQuotaResult(window, "reset_after_seconds", "resetAfterSeconds")
		resetAt := firstCodexQuotaResult(window, "reset_at", "resetAt")
		hasResetAfter := resetAfter.Exists() && resetAfter.Int() >= 0
		hasResetAt := resetAt.Exists() && resetAt.Int() > 0
		if !used.Exists() || !minutes.Exists() ||
			used.Float() < 0 || used.Float() > 100 || minutes.Int() <= 0 ||
			(!hasResetAfter && !hasResetAt) {
			continue
		}
		windowPrefix := prefix + strings.ToUpper(windowName[:1]) + windowName[1:] + "-"
		setCodexQuotaScalarHeaderFromResult(headers, windowPrefix+"Used-Percent", window, "used_percent", "usedPercent")
		setCodexQuotaScalarHeaderFromResult(headers, windowPrefix+"Window-Minutes", window, "window_minutes", "windowMinutes")
		if hasResetAfter {
			setCodexQuotaScalarHeaderFromResult(headers, windowPrefix+"Reset-After-Seconds", window, "reset_after_seconds", "resetAfterSeconds")
		}
		if hasResetAt {
			setCodexQuotaScalarHeaderFromResult(headers, windowPrefix+"Reset-At", window, "reset_at", "resetAt")
		}
		changed = true
	}
	return changed
}

func parseCodexQuotaHeadersObject(headersNode gjson.Result) http.Header {
	if !headersNode.Exists() || !headersNode.IsObject() {
		return nil
	}
	headers := make(http.Header)
	headersNode.ForEach(func(key, value gjson.Result) bool {
		name := http.CanonicalHeaderKey(strings.TrimSpace(key.String()))
		if !isCodexQuotaHeaderName(name) {
			return true
		}
		if raw := codexQuotaScalarValue(value); raw != "" {
			headers.Set(name, raw)
		}
		return true
	})
	if len(headers) == 0 {
		return nil
	}
	return headers
}

func isCodexQuotaHeaderName(name string) bool {
	lower := strings.ToLower(strings.TrimSpace(name))
	if lower == "retry-after" || strings.HasPrefix(lower, "x-ratelimit-") {
		return true
	}
	if lower == "x-codex-active-limit" || lower == "x-codex-plan-type" ||
		strings.HasPrefix(lower, "x-codex-credits-") {
		return true
	}
	if !strings.HasPrefix(lower, "x-codex-") {
		return false
	}
	for _, marker := range []string{
		"-allowed",
		"-limit-reached",
		"-limit-name",
		"-used-percent",
		"-window-minutes",
		"-reset-after-seconds",
		"-reset-at",
		"-over-secondary-limit-percent",
	} {
		if strings.Contains(lower, marker) {
			return true
		}
	}
	return false
}

func setCodexQuotaScalarHeader(headers http.Header, name string, object gjson.Result, paths ...string) bool {
	return setCodexQuotaScalarHeaderFromResult(headers, name, object, paths...)
}

func setCodexQuotaScalarHeaderFromResult(headers http.Header, name string, object gjson.Result, paths ...string) bool {
	value := firstCodexQuotaResult(object, paths...)
	if !value.Exists() {
		return false
	}
	raw := codexQuotaScalarValue(value)
	if raw == "" {
		return false
	}
	headers.Set(name, raw)
	return true
}

func codexQuotaScalarValue(value gjson.Result) string {
	switch value.Type {
	case gjson.String:
		return strings.TrimSpace(value.String())
	case gjson.Number, gjson.True, gjson.False:
		return strings.TrimSpace(value.Raw)
	default:
		return ""
	}
}

func firstCodexQuotaResult(object gjson.Result, paths ...string) gjson.Result {
	for _, path := range paths {
		value := object.Get(path)
		if value.Exists() && value.Type != gjson.Null {
			return value
		}
	}
	return gjson.Result{}
}

func firstCodexQuotaResultString(object gjson.Result, paths ...string) string {
	for _, path := range paths {
		value := firstCodexQuotaResult(object, path)
		if raw := codexQuotaScalarValue(value); raw != "" {
			return raw
		}
	}
	return ""
}

func normalizeCodexQuotaHeaderIdentifier(value string) string {
	value = strings.TrimSpace(value)
	if value == "" || len(value) > 128 {
		return ""
	}
	var builder strings.Builder
	builder.Grow(len(value))
	lastDash := false
	for _, char := range value {
		switch {
		case (char >= 'a' && char <= 'z') || (char >= 'A' && char <= 'Z') || (char >= '0' && char <= '9') || char == '.' || char == '_':
			builder.WriteRune(char)
			lastDash = false
		case char == '-':
			builder.WriteRune(char)
			lastDash = false
		default:
			if !lastDash {
				builder.WriteByte('-')
				lastDash = true
			}
		}
	}
	identifier := strings.Trim(builder.String(), "-_.")
	if identifier == "" || !validCodexQuotaEventIdentifier(identifier) {
		return ""
	}
	return identifier
}

func validCodexQuotaEventIdentifier(value string) bool {
	value = strings.TrimSpace(value)
	if value == "" || len(value) > 256 {
		return false
	}
	for _, char := range value {
		if (char >= 'a' && char <= 'z') || (char >= 'A' && char <= 'Z') ||
			(char >= '0' && char <= '9') || char == '_' || char == '-' || char == '.' {
			continue
		}
		return false
	}
	return true
}

func validCodexQuotaEventText(value string) bool {
	value = strings.TrimSpace(value)
	return value != "" && len(value) <= 256 && !strings.ContainsAny(value, "\r\n")
}