diff --git a/internal/runtime/executor/codex_websockets_execute.go b/internal/runtime/executor/codex_websockets_execute.go index 43f86ad80..5a42218d6 100644 --- a/internal/runtime/executor/codex_websockets_execute.go +++ b/internal/runtime/executor/codex_websockets_execute.go @@ -280,6 +280,7 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut payload = applyCodexIdentityConfuseResponsePayload(payload, identityState) helps.AppendAPIWebsocketResponse(ctx, e.cfg, payload) payload = helps.RestoreCodexMultiAgentV2Response(payload, optimizeMultiAgentV2) + reporter.ObserveQuotaHeaders(helps.ParseCodexQuotaEventHeaders(payload)) if wsErr, ok := parseCodexWebsocketError(payload); ok { if sess != nil { diff --git a/internal/runtime/executor/codex_websockets_stream.go b/internal/runtime/executor/codex_websockets_stream.go index 719e1a363..401614ec1 100644 --- a/internal/runtime/executor/codex_websockets_stream.go +++ b/internal/runtime/executor/codex_websockets_stream.go @@ -337,6 +337,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr payload = applyCodexIdentityConfuseResponsePayload(payload, identityState) helps.AppendAPIWebsocketResponse(ctx, e.cfg, payload) payload = helps.RestoreCodexMultiAgentV2Response(payload, optimizeMultiAgentV2) + reporter.ObserveQuotaHeaders(helps.ParseCodexQuotaEventHeaders(payload)) if wsErr, ok := parseCodexWebsocketError(payload); ok { terminateReason = "upstream_error" diff --git a/internal/runtime/executor/helps/codex_quota.go b/internal/runtime/executor/helps/codex_quota.go new file mode 100644 index 000000000..a837ceafc --- /dev/null +++ b/internal/runtime/executor/helps/codex_quota.go @@ -0,0 +1,99 @@ +package helps + +import ( + "math" + "net/http" + "strconv" + "strings" + + "github.com/tidwall/gjson" +) + +const codexRateLimitsEventType = "codex.rate_limits" + +// ParseCodexQuotaEventHeaders converts one Codex websocket quota event into the +// same bounded header representation used by HTTP usage observations. +func ParseCodexQuotaEventHeaders(payload []byte) http.Header { + if gjson.GetBytes(payload, "type").String() != codexRateLimitsEventType { + return nil + } + + activeLimit := firstCodexQuotaEventString(payload, "metered_limit_name", "limit_name") + if activeLimit == "" { + activeLimit = "codex" + } + if !validCodexQuotaEventIdentifier(activeLimit) { + return nil + } + + headers := make(http.Header) + usableWindows := 0 + for _, windowName := range []string{"primary", "secondary"} { + path := "rate_limits." + windowName + usedPercent := gjson.GetBytes(payload, path+".used_percent") + windowMinutes := gjson.GetBytes(payload, path+".window_minutes") + if !usedPercent.Exists() || !windowMinutes.Exists() { + continue + } + used := usedPercent.Float() + minutes := windowMinutes.Int() + if math.IsNaN(used) || math.IsInf(used, 0) || used < 0 || minutes <= 0 { + continue + } + + prefix := "X-Codex-" + strings.ToUpper(windowName[:1]) + windowName[1:] + "-" + resetHeader := "" + resetValue := "" + if resetAt := gjson.GetBytes(payload, path+".reset_at"); resetAt.Exists() && resetAt.Int() > 0 { + resetHeader = prefix + "Reset-At" + resetValue = strconv.FormatInt(resetAt.Int(), 10) + } else if resetAfter := gjson.GetBytes(payload, path+".reset_after_seconds"); resetAfter.Exists() && resetAfter.Int() >= 0 { + resetHeader = prefix + "Reset-After-Seconds" + resetValue = strconv.FormatInt(resetAfter.Int(), 10) + } else { + continue + } + headers.Set(prefix+"Used-Percent", strconv.FormatFloat(used, 'f', -1, 64)) + headers.Set(prefix+"Window-Minutes", strconv.FormatInt(minutes, 10)) + headers.Set(resetHeader, resetValue) + usableWindows++ + } + if usableWindows == 0 { + return nil + } + + headers.Set("X-Codex-Active-Limit", activeLimit) + if planType := firstCodexQuotaEventString(payload, "plan_type"); validCodexQuotaEventText(planType) { + headers.Set("X-Codex-Plan-Type", planType) + } + return headers +} + +func firstCodexQuotaEventString(payload []byte, paths ...string) string { + for _, path := range paths { + if value := strings.TrimSpace(gjson.GetBytes(payload, path).String()); value != "" { + return value + } + } + return "" +} + +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") +} diff --git a/internal/runtime/executor/helps/codex_quota_test.go b/internal/runtime/executor/helps/codex_quota_test.go new file mode 100644 index 000000000..02eb69436 --- /dev/null +++ b/internal/runtime/executor/helps/codex_quota_test.go @@ -0,0 +1,69 @@ +package helps + +import ( + "context" + "net/http" + "testing" + + internallogging "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" +) + +func TestParseCodexQuotaEventHeadersPreservesActiveLimit(t *testing.T) { + headers := ParseCodexQuotaEventHeaders([]byte(`{ + "type":"codex.rate_limits", + "plan_type":"pro", + "metered_limit_name":"codex_bengalfox", + "rate_limits":{ + "primary":{"used_percent":2,"window_minutes":10080,"reset_at":1782951970}, + "secondary":null + } + }`)) + if headers.Get("X-Codex-Active-Limit") != "codex_bengalfox" || + headers.Get("X-Codex-Primary-Used-Percent") != "2" || + headers.Get("X-Codex-Primary-Window-Minutes") != "10080" || + headers.Get("X-Codex-Primary-Reset-At") != "1782951970" || + headers.Get("X-Codex-Plan-Type") != "pro" { + t.Fatalf("unexpected quota headers: %#v", headers) + } +} + +func TestParseCodexQuotaEventHeadersDefaultsToCodexAndRejectsIncompleteWindows(t *testing.T) { + headers := ParseCodexQuotaEventHeaders([]byte(`{ + "type":"codex.rate_limits", + "rate_limits":{"primary":{"used_percent":42,"window_minutes":300,"reset_after_seconds":60}} + }`)) + if headers.Get("X-Codex-Active-Limit") != "codex" || headers.Get("X-Codex-Primary-Reset-After-Seconds") != "60" { + t.Fatalf("unexpected default quota headers: %#v", headers) + } + if got := ParseCodexQuotaEventHeaders([]byte(`{"type":"codex.rate_limits","rate_limits":{"primary":{"used_percent":42}}}`)); got != nil { + t.Fatalf("incomplete quota event produced headers: %#v", got) + } + + headers = ParseCodexQuotaEventHeaders([]byte(`{ + "type":"codex.rate_limits", + "rate_limits":{ + "primary":{"used_percent":42,"window_minutes":300}, + "secondary":{"used_percent":17,"window_minutes":10080,"reset_after_seconds":60} + } + }`)) + if headers.Get("X-Codex-Primary-Used-Percent") != "" || + headers.Get("X-Codex-Secondary-Used-Percent") != "17" { + t.Fatalf("partially valid quota event produced incomplete headers: %#v", headers) + } +} + +func TestUsageReporterMergesWebsocketQuotaHeaders(t *testing.T) { + ctx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(ctx, http.Header{"X-Request-Id": []string{"req-1"}}) + reporter := NewUsageReporter(ctx, "codex", "gpt-5.3-codex-spark", nil) + reporter.ObserveQuotaHeaders(http.Header{ + "X-Codex-Active-Limit": []string{"codex_bengalfox"}, + "X-Codex-Primary-Used-Percent": []string{"2"}, + "X-Codex-Primary-Window-Minutes": []string{"10080"}, + "X-Codex-Primary-Reset-At": []string{"1782951970"}, + }) + headers := reporter.responseHeaders(ctx) + if headers.Get("X-Request-Id") != "req-1" || headers.Get("X-Codex-Active-Limit") != "codex_bengalfox" { + t.Fatalf("merged response headers = %#v", headers) + } +} diff --git a/internal/runtime/executor/helps/usage_helpers.go b/internal/runtime/executor/helps/usage_helpers.go index 52e1687f3..459a3e783 100644 --- a/internal/runtime/executor/helps/usage_helpers.go +++ b/internal/runtime/executor/helps/usage_helpers.go @@ -22,24 +22,26 @@ import ( ) type UsageReporter struct { - provider string - executorType string - model string - alias string - authID string - authIndex string - authType string - apiKey string - source string - reasoning string - serviceTier string - generate bool - requestedAt time.Time - ttftMu sync.RWMutex - ttft time.Duration - ttftStart time.Time - ttftSet bool - once sync.Once + provider string + executorType string + model string + alias string + authID string + authIndex string + authType string + apiKey string + source string + reasoning string + serviceTier string + generate bool + requestedAt time.Time + quotaHeadersMu sync.RWMutex + quotaHeaders http.Header + ttftMu sync.RWMutex + ttft time.Duration + ttftStart time.Time + ttftSet bool + once sync.Once } type usageExecutor interface { @@ -236,10 +238,50 @@ func (r *UsageReporter) EnsurePublished(ctx context.Context) { } func (r *UsageReporter) publishRecord(ctx context.Context, record usage.Record) { - record.ResponseHeaders = internallogging.GetResponseHeaders(ctx) + record.ResponseHeaders = r.responseHeaders(ctx) usage.PublishRecord(ctx, record) } +// ObserveQuotaHeaders merges provider quota metadata that arrived outside the +// HTTP response header block, such as Codex websocket rate-limit events. +func (r *UsageReporter) ObserveQuotaHeaders(headers http.Header) { + if r == nil || len(headers) == 0 { + return + } + r.quotaHeadersMu.Lock() + if r.quotaHeaders == nil { + r.quotaHeaders = make(http.Header, len(headers)) + } + for key, values := range headers { + canonicalKey := http.CanonicalHeaderKey(strings.TrimSpace(key)) + if canonicalKey == "" { + continue + } + r.quotaHeaders[canonicalKey] = append([]string(nil), values...) + } + r.quotaHeadersMu.Unlock() +} + +func (r *UsageReporter) responseHeaders(ctx context.Context) http.Header { + headers := internallogging.GetResponseHeaders(ctx) + if r == nil { + return headers + } + r.quotaHeadersMu.RLock() + if len(r.quotaHeaders) == 0 { + r.quotaHeadersMu.RUnlock() + return headers + } + if headers == nil { + headers = make(http.Header, len(r.quotaHeaders)) + } + for key, values := range r.quotaHeaders { + headers[key] = append([]string(nil), values...) + } + r.quotaHeadersMu.RUnlock() + return headers +} + func (r *UsageReporter) buildRecord(detail usage.Detail, failed bool, failures ...usage.Failure) usage.Record { var fail usage.Failure if len(failures) > 0 {