refactor: codex usage ui

controller/codex_usage.go

@@ -18,6 +18,46 @@ import (
 )
 
 func GetCodexChannelUsage(c *gin.Context) {
+	fetchCodexChannelWhamData(
+		c,
+		service.FetchCodexWhamUsage,
+		"failed to fetch codex usage",
+		"获取用量信息失败，请稍后重试",
+	)
+}
+
+func GetCodexChannelRateLimitResetCredits(c *gin.Context) {
+	fetchCodexChannelWhamData(
+		c,
+		service.FetchCodexWhamRateLimitResetCredits,
+		"failed to fetch codex reset credits",
+		"获取重置次数详情失败，请稍后重试",
+	)
+}
+
+func ResetCodexChannelUsage(c *gin.Context) {
+	fetchCodexChannelWhamData(
+		c,
+		service.ConsumeCodexWhamRateLimitResetCredit,
+		"failed to reset codex usage",
+		"重置用量失败，请稍后重试",
+	)
+}
+
+type codexWhamFetchFunc func(
+	ctx context.Context,
+	client *http.Client,
+	baseURL string,
+	accessToken string,
+	accountID string,
+) (statusCode int, body []byte, err error)
+
+func fetchCodexChannelWhamData(
+	c *gin.Context,
+	fetch codexWhamFetchFunc,
+	logPrefix string,
+	userMessage string,
+) {
 	channelId, err := strconv.Atoi(c.Param("id"))
 	if err != nil {
 		common.ApiError(c, fmt.Errorf("invalid channel id: %w", err))
@@ -68,10 +108,10 @@ func GetCodexChannelUsage(c *gin.Context) {
 	ctx, cancel := context.WithTimeout(c.Request.Context(), 15*time.Second)
 	defer cancel()
 
-	statusCode, body, err := service.FetchCodexWhamUsage(ctx, client, ch.GetBaseURL(), accessToken, accountID)
+	statusCode, body, err := fetch(ctx, client, ch.GetBaseURL(), accessToken, accountID)
 	if err != nil {
-		common.SysError("failed to fetch codex usage: " + err.Error())
-		c.JSON(http.StatusOK, gin.H{"success": false, "message": "获取用量信息失败，请稍后重试"})
+		common.SysError(logPrefix + ": " + err.Error())
+		c.JSON(http.StatusOK, gin.H{"success": false, "message": userMessage})
 		return
 	}
 
@@ -98,10 +138,10 @@ func GetCodexChannelUsage(c *gin.Context) {
 
 			ctx2, cancel2 := context.WithTimeout(c.Request.Context(), 15*time.Second)
 			defer cancel2()
-			statusCode, body, err = service.FetchCodexWhamUsage(ctx2, client, ch.GetBaseURL(), oauthKey.AccessToken, accountID)
+			statusCode, body, err = fetch(ctx2, client, ch.GetBaseURL(), oauthKey.AccessToken, accountID)
 			if err != nil {
-				common.SysError("failed to fetch codex usage after refresh: " + err.Error())
-				c.JSON(http.StatusOK, gin.H{"success": false, "message": "获取用量信息失败，请稍后重试"})
+				common.SysError(logPrefix + " after refresh: " + err.Error())
+				c.JSON(http.StatusOK, gin.H{"success": false, "message": userMessage})
 				return
 			}
 		}

router/api-router.go

@@ -251,6 +251,8 @@ func SetApiRouter(router *gin.Engine) {
 			channelRoute.POST("/fetch_models", middleware.RootAuth(), controller.FetchModels)
 			channelRoute.POST("/:id/codex/refresh", controller.RefreshCodexChannelCredential)
 			channelRoute.GET("/:id/codex/usage", controller.GetCodexChannelUsage)
+			channelRoute.GET("/:id/codex/usage/reset-credits", controller.GetCodexChannelRateLimitResetCredits)
+			channelRoute.POST("/:id/codex/usage/reset", controller.ResetCodexChannelUsage)
 			channelRoute.POST("/ollama/pull", controller.OllamaPullModel)
 			channelRoute.POST("/ollama/pull/stream", controller.OllamaPullModelStream)
 			channelRoute.DELETE("/ollama/delete", controller.OllamaDeleteModel)

service/codex_wham_usage.go

@@ -1,11 +1,15 @@
 package service
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"io"
 	"net/http"
 	"strings"
+
+	"github.com/QuantumNous/new-api/common"
+	"github.com/google/uuid"
 )
 
 func FetchCodexWhamUsage(
@@ -35,13 +39,50 @@ func FetchCodexWhamUsage(
 	if err != nil {
 		return 0, nil, err
 	}
-	req.Header.Set("Authorization", "Bearer "+at)
-	req.Header.Set("chatgpt-account-id", aid)
-	req.Header.Set("Accept", "application/json")
-	if req.Header.Get("originator") == "" {
-		req.Header.Set("originator", "codex_cli_rs")
+	setCodexWhamRequestHeaders(req, at, aid)
+
+	resp, err := client.Do(req)
+	if err != nil {
+		return 0, nil, err
+	}
+	defer resp.Body.Close()
+
+	body, err = io.ReadAll(resp.Body)
+	if err != nil {
+		return resp.StatusCode, nil, err
+	}
+	return resp.StatusCode, body, nil
+}
+
+func FetchCodexWhamRateLimitResetCredits(
+	ctx context.Context,
+	client *http.Client,
+	baseURL string,
+	accessToken string,
+	accountID string,
+) (statusCode int, body []byte, err error) {
+	if client == nil {
+		return 0, nil, fmt.Errorf("nil http client")
+	}
+	bu := strings.TrimRight(strings.TrimSpace(baseURL), "/")
+	if bu == "" {
+		return 0, nil, fmt.Errorf("empty baseURL")
+	}
+	at := strings.TrimSpace(accessToken)
+	aid := strings.TrimSpace(accountID)
+	if at == "" {
+		return 0, nil, fmt.Errorf("empty accessToken")
+	}
+	if aid == "" {
+		return 0, nil, fmt.Errorf("empty accountID")
 	}
 
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, bu+"/backend-api/wham/rate-limit-reset-credits", nil)
+	if err != nil {
+		return 0, nil, err
+	}
+	setCodexWhamRequestHeaders(req, at, aid)
+
 	resp, err := client.Do(req)
 	if err != nil {
 		return 0, nil, err
@@ -54,3 +95,67 @@ func FetchCodexWhamUsage(
 	}
 	return resp.StatusCode, body, nil
 }
+
+func ConsumeCodexWhamRateLimitResetCredit(
+	ctx context.Context,
+	client *http.Client,
+	baseURL string,
+	accessToken string,
+	accountID string,
+) (statusCode int, body []byte, err error) {
+	if client == nil {
+		return 0, nil, fmt.Errorf("nil http client")
+	}
+	bu := strings.TrimRight(strings.TrimSpace(baseURL), "/")
+	if bu == "" {
+		return 0, nil, fmt.Errorf("empty baseURL")
+	}
+	at := strings.TrimSpace(accessToken)
+	aid := strings.TrimSpace(accountID)
+	if at == "" {
+		return 0, nil, fmt.Errorf("empty accessToken")
+	}
+	if aid == "" {
+		return 0, nil, fmt.Errorf("empty accountID")
+	}
+
+	requestBody, err := common.Marshal(map[string]string{
+		"redeem_request_id": uuid.NewString(),
+	})
+	if err != nil {
+		return 0, nil, err
+	}
+
+	req, err := http.NewRequestWithContext(
+		ctx,
+		http.MethodPost,
+		bu+"/backend-api/wham/rate-limit-reset-credits/consume",
+		bytes.NewReader(requestBody),
+	)
+	if err != nil {
+		return 0, nil, err
+	}
+	setCodexWhamRequestHeaders(req, at, aid)
+	req.Header.Set("Content-Type", "application/json")
+
+	resp, err := client.Do(req)
+	if err != nil {
+		return 0, nil, err
+	}
+	defer resp.Body.Close()
+
+	body, err = io.ReadAll(resp.Body)
+	if err != nil {
+		return resp.StatusCode, nil, err
+	}
+	return resp.StatusCode, body, nil
+}
+
+func setCodexWhamRequestHeaders(req *http.Request, accessToken string, accountID string) {
+	req.Header.Set("Authorization", "Bearer "+accessToken)
+	req.Header.Set("chatgpt-account-id", accountID)
+	req.Header.Set("Accept", "application/json")
+	if req.Header.Get("originator") == "" {
+		req.Header.Set("originator", "codex_cli_rs")
+	}
+}