feat: add channel auto reset usage option

dto/channel_settings.go

@@ -43,6 +43,7 @@ type ChannelOtherSettings struct {
 	AllowSpeed                            bool                  `json:"allow_speed,omitempty"`                // 是否允许 speed 透传（仅 Claude，默认过滤以避免意外切换推理速度模式）
 	AllowSafetyIdentifier                 bool                  `json:"allow_safety_identifier,omitempty"`    // 是否允许 safety_identifier 透传（默认过滤以保护用户隐私）
 	DisableStore                          bool                  `json:"disable_store,omitempty"`              // 是否禁用 store 透传（默认允许透传，禁用后可能导致 Codex 无法使用）
+	AutoResetUsageEnabled                 bool                  `json:"auto_reset_usage_enabled,omitempty"`   // 是否在限流时自动使用一次可用重置次数
 	AllowIncludeObfuscation               bool                  `json:"allow_include_obfuscation,omitempty"`  // 是否允许 stream_options.include_obfuscation 透传（默认过滤以避免关闭流混淆保护）
 	DisableTaskPollingSleep               bool                  `json:"disable_task_polling_sleep,omitempty"` // 是否跳过异步任务轮询间隔
 	AwsKeyType                            AwsKeyType            `json:"aws_key_type,omitempty"`

relay/channel/codex/adaptor.go

@@ -1,26 +1,55 @@
 package codex
 
 import (
+	"context"
+	"crypto/sha256"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"io"
 	"net/http"
 	"strings"
+	"time"
 
 	"github.com/QuantumNous/new-api/common"
+	projectconstant "github.com/QuantumNous/new-api/constant"
 	"github.com/QuantumNous/new-api/dto"
+	"github.com/QuantumNous/new-api/logger"
 	"github.com/QuantumNous/new-api/relay/channel"
 	"github.com/QuantumNous/new-api/relay/channel/openai"
 	relaycommon "github.com/QuantumNous/new-api/relay/common"
 	relayconstant "github.com/QuantumNous/new-api/relay/constant"
+	"github.com/QuantumNous/new-api/service"
 	"github.com/QuantumNous/new-api/types"
 
 	"github.com/gin-gonic/gin"
+	"golang.org/x/sync/singleflight"
 )
 
 type Adaptor struct {
 }
 
+var (
+	codexAutoResetTimeout = 15 * time.Second
+	codexAutoResetGroup   singleflight.Group
+)
+
+const codexAutoResetLockTTL = 15 * time.Minute
+
+type codexRateLimitWindow struct {
+	UsedPercent        float64 `json:"used_percent"`
+	LimitWindowSeconds int64   `json:"limit_window_seconds"`
+}
+
+type codexUsagePayload struct {
+	PlanType  string `json:"plan_type"`
+	RateLimit struct {
+		PlanType        string                `json:"plan_type"`
+		PrimaryWindow   *codexRateLimitWindow `json:"primary_window"`
+		SecondaryWindow *codexRateLimitWindow `json:"secondary_window"`
+	} `json:"rate_limit"`
+}
+
 func (a *Adaptor) ConvertGeminiRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeminiChatRequest) (any, error) {
 	return nil, errors.New("codex channel: endpoint not supported")
 }
@@ -108,7 +137,198 @@ func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommo
 }
 
 func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
-	return channel.DoApiRequest(a, c, info, requestBody)
+	if info == nil || !info.ChannelOtherSettings.AutoResetUsageEnabled || requestBody == nil {
+		return channel.DoApiRequest(a, c, info, requestBody)
+	}
+
+	maxBytes := int64(projectconstant.MaxRequestBodyMB)
+	if maxBytes <= 0 {
+		maxBytes = 128
+	}
+	storage, err := common.CreateBodyStorageFromReader(requestBody, info.UpstreamRequestBodySize, maxBytes<<20)
+	if err != nil {
+		return nil, err
+	}
+	defer storage.Close()
+
+	if _, err = storage.Seek(0, io.SeekStart); err != nil {
+		return nil, err
+	}
+	resp, err := channel.DoApiRequest(a, c, info, common.ReaderOnly(storage))
+	if err != nil {
+		return nil, err
+	}
+	if resp == nil || resp.StatusCode != http.StatusTooManyRequests {
+		return resp, nil
+	}
+
+	if !consumeCodexResetCredit(c, info) {
+		return resp, nil
+	}
+	_ = resp.Body.Close()
+
+	if _, err = storage.Seek(0, io.SeekStart); err != nil {
+		return nil, err
+	}
+	return channel.DoApiRequest(a, c, info, common.ReaderOnly(storage))
+}
+
+func consumeCodexResetCredit(c *gin.Context, info *relaycommon.RelayInfo) bool {
+	oauthKey, err := ParseOAuthKey(strings.TrimSpace(info.ApiKey))
+	if err != nil {
+		logger.LogWarn(c, "codex auto reset usage skipped: "+err.Error())
+		return false
+	}
+
+	client := service.GetHttpClient()
+	if info.ChannelSetting.Proxy != "" {
+		client, err = service.NewProxyHttpClient(info.ChannelSetting.Proxy)
+		if err != nil {
+			logger.LogWarn(c, "codex auto reset usage skipped: "+err.Error())
+			return false
+		}
+	}
+
+	requestContext := context.Background()
+	if c != nil && c.Request != nil {
+		requestContext = c.Request.Context()
+	}
+	resetKey := codexAutoResetKey(info.ChannelBaseUrl, oauthKey.AccountID)
+	resultChannel := codexAutoResetGroup.DoChan(resetKey, func() (any, error) {
+		ctx, cancel := context.WithTimeout(context.Background(), codexAutoResetTimeout)
+		defer cancel()
+
+		if common.RedisEnabled && common.RDB != nil {
+			acquired, lockErr := common.RDB.SetNX(
+				ctx,
+				"codex:auto-reset:lock:"+resetKey,
+				"1",
+				codexAutoResetLockTTL,
+			).Result()
+			if lockErr == nil && !acquired {
+				return false, nil
+			}
+			if lockErr != nil {
+				logger.LogWarn(c, "codex auto reset Redis lock unavailable: "+lockErr.Error())
+			}
+		}
+
+		eligible, eligibilityErr := checkCodexAutoResetEligibility(ctx, client, info, oauthKey)
+		if eligibilityErr != nil || !eligible {
+			return false, eligibilityErr
+		}
+		return performCodexAutoReset(ctx, client, info, oauthKey)
+	})
+
+	select {
+	case result := <-resultChannel:
+		if result.Err != nil {
+			if !result.Shared {
+				logger.LogWarn(c, "codex auto reset usage failed: "+result.Err.Error())
+			}
+			return false
+		}
+		reset, ok := result.Val.(bool)
+		if reset && !result.Shared {
+			logger.LogInfo(c, "codex auto reset usage ready for retry")
+		}
+		return ok && reset
+	case <-requestContext.Done():
+		return false
+	}
+}
+
+func checkCodexAutoResetEligibility(ctx context.Context, client *http.Client, info *relaycommon.RelayInfo, oauthKey *OAuthKey) (bool, error) {
+	statusCode, body, err := service.FetchCodexWhamUsage(
+		ctx,
+		client,
+		info.ChannelBaseUrl,
+		oauthKey.AccessToken,
+		oauthKey.AccountID,
+	)
+	if err != nil {
+		return false, fmt.Errorf("fetch usage: %w", err)
+	}
+	if statusCode < http.StatusOK || statusCode >= http.StatusMultipleChoices {
+		return false, fmt.Errorf("fetch usage: upstream_status=%d", statusCode)
+	}
+
+	var usage codexUsagePayload
+	if err = common.Unmarshal(body, &usage); err != nil {
+		return false, fmt.Errorf("parse usage: %w", err)
+	}
+	weeklyExhausted := false
+	for _, window := range []*codexRateLimitWindow{
+		usage.RateLimit.PrimaryWindow,
+		usage.RateLimit.SecondaryWindow,
+	} {
+		if window != nil && window.LimitWindowSeconds >= int64((24*time.Hour)/time.Second) && window.UsedPercent >= 100 {
+			weeklyExhausted = true
+			break
+		}
+	}
+	planType := usage.PlanType
+	if planType == "" {
+		planType = usage.RateLimit.PlanType
+	}
+	if !weeklyExhausted && strings.EqualFold(planType, "free") && usage.RateLimit.PrimaryWindow != nil {
+		weeklyExhausted = usage.RateLimit.PrimaryWindow.UsedPercent >= 100
+	}
+	if !weeklyExhausted {
+		return false, nil
+	}
+
+	statusCode, body, err = service.FetchCodexWhamRateLimitResetCredits(
+		ctx,
+		client,
+		info.ChannelBaseUrl,
+		oauthKey.AccessToken,
+		oauthKey.AccountID,
+	)
+	if err != nil {
+		return false, fmt.Errorf("fetch reset credits: %w", err)
+	}
+	if statusCode < http.StatusOK || statusCode >= http.StatusMultipleChoices {
+		return false, fmt.Errorf("fetch reset credits: upstream_status=%d", statusCode)
+	}
+	var credits struct {
+		AvailableCount int `json:"available_count"`
+	}
+	if err = common.Unmarshal(body, &credits); err != nil {
+		return false, fmt.Errorf("parse reset credits: %w", err)
+	}
+	if credits.AvailableCount <= 0 {
+		return false, nil
+	}
+	return true, nil
+}
+
+func performCodexAutoReset(ctx context.Context, client *http.Client, info *relaycommon.RelayInfo, oauthKey *OAuthKey) (bool, error) {
+	statusCode, body, err := service.ConsumeCodexWhamRateLimitResetCredit(
+		ctx,
+		client,
+		info.ChannelBaseUrl,
+		oauthKey.AccessToken,
+		oauthKey.AccountID,
+	)
+	if err != nil {
+		return false, fmt.Errorf("consume reset credit: %w", err)
+	}
+	if statusCode < http.StatusOK || statusCode >= http.StatusMultipleChoices {
+		return false, fmt.Errorf("consume reset credit: upstream_status=%d", statusCode)
+	}
+	var resetResult struct {
+		WindowsReset int `json:"windows_reset"`
+	}
+	if err = common.Unmarshal(body, &resetResult); err != nil {
+		return false, fmt.Errorf("parse reset result: %w", err)
+	}
+	return resetResult.WindowsReset > 0, nil
+}
+
+func codexAutoResetKey(baseURL string, accountID string) string {
+	identity := strings.TrimRight(strings.TrimSpace(baseURL), "/") + "|" + strings.TrimSpace(accountID)
+	return fmt.Sprintf("%x", sha256.Sum256([]byte(identity)))
 }
 
 func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {