You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

125 lines
3.7 KiB

  1. package controller
  2. import (
  3. "context"
  4. "errors"
  5. "fmt"
  6. "net/http"
  7. "strconv"
  8. "strings"
  9. "time"
  10. "github.com/QuantumNous/new-api/common"
  11. "github.com/QuantumNous/new-api/constant"
  12. "github.com/QuantumNous/new-api/model"
  13. "github.com/QuantumNous/new-api/service"
  14. "github.com/gin-gonic/gin"
  15. )
  16. func GetCodexChannelUsage(c *gin.Context) {
  17. channelId, err := strconv.Atoi(c.Param("id"))
  18. if err != nil {
  19. common.ApiError(c, fmt.Errorf("invalid channel id: %w", err))
  20. return
  21. }
  22. ch, err := model.GetChannelById(channelId, true)
  23. if err != nil {
  24. common.ApiError(c, err)
  25. return
  26. }
  27. if ch == nil {
  28. c.JSON(http.StatusOK, gin.H{"success": false, "message": "channel not found"})
  29. return
  30. }
  31. if ch.Type != constant.ChannelTypeCodex {
  32. c.JSON(http.StatusOK, gin.H{"success": false, "message": "channel type is not Codex"})
  33. return
  34. }
  35. if ch.ChannelInfo.IsMultiKey {
  36. c.JSON(http.StatusOK, gin.H{"success": false, "message": "multi-key channel is not supported"})
  37. return
  38. }
  39. oauthKey, err := common.ParseCodexOAuthCredential(strings.TrimSpace(ch.Key))
  40. if err != nil {
  41. if errors.Is(err, common.ErrCodexOAuthCredentialRequired) {
  42. c.JSON(http.StatusOK, gin.H{"success": false, "message": "当前凭证方式不支持查看用量"})
  43. return
  44. }
  45. common.SysError("failed to parse oauth key: " + err.Error())
  46. c.JSON(http.StatusOK, gin.H{"success": false, "message": "解析凭证失败,请检查渠道配置"})
  47. return
  48. }
  49. accessToken := strings.TrimSpace(oauthKey.AccessToken)
  50. accountID := strings.TrimSpace(oauthKey.AccountID)
  51. client, err := service.NewProxyHttpClient(ch.GetSetting().Proxy)
  52. if err != nil {
  53. common.ApiError(c, err)
  54. return
  55. }
  56. ctx, cancel := context.WithTimeout(c.Request.Context(), 15*time.Second)
  57. defer cancel()
  58. statusCode, body, err := service.FetchCodexWhamUsage(ctx, client, ch.GetBaseURL(), accessToken, accountID)
  59. if err != nil {
  60. common.SysError("failed to fetch codex usage: " + err.Error())
  61. c.JSON(http.StatusOK, gin.H{"success": false, "message": "获取用量信息失败,请稍后重试"})
  62. return
  63. }
  64. if (statusCode == http.StatusUnauthorized || statusCode == http.StatusForbidden) && strings.TrimSpace(oauthKey.RefreshToken) != "" {
  65. refreshCtx, refreshCancel := context.WithTimeout(c.Request.Context(), 10*time.Second)
  66. defer refreshCancel()
  67. res, refreshErr := service.RefreshCodexOAuthTokenWithProxy(refreshCtx, oauthKey.RefreshToken, ch.GetSetting().Proxy)
  68. if refreshErr == nil {
  69. oauthKey.AccessToken = res.AccessToken
  70. oauthKey.RefreshToken = res.RefreshToken
  71. oauthKey.LastRefresh = time.Now().Format(time.RFC3339)
  72. oauthKey.Expired = res.ExpiresAt.Format(time.RFC3339)
  73. if strings.TrimSpace(oauthKey.Type) == "" {
  74. oauthKey.Type = "codex"
  75. }
  76. encoded, encErr := common.Marshal(oauthKey)
  77. if encErr == nil {
  78. _ = model.DB.Model(&model.Channel{}).Where("id = ?", ch.Id).Update("key", string(encoded)).Error
  79. model.InitChannelCache()
  80. service.ResetProxyClientCache()
  81. }
  82. ctx2, cancel2 := context.WithTimeout(c.Request.Context(), 15*time.Second)
  83. defer cancel2()
  84. statusCode, body, err = service.FetchCodexWhamUsage(ctx2, client, ch.GetBaseURL(), oauthKey.AccessToken, accountID)
  85. if err != nil {
  86. common.SysError("failed to fetch codex usage after refresh: " + err.Error())
  87. c.JSON(http.StatusOK, gin.H{"success": false, "message": "获取用量信息失败,请稍后重试"})
  88. return
  89. }
  90. }
  91. }
  92. var payload any
  93. if common.Unmarshal(body, &payload) != nil {
  94. payload = string(body)
  95. }
  96. ok := statusCode >= 200 && statusCode < 300
  97. resp := gin.H{
  98. "success": ok,
  99. "message": "",
  100. "upstream_status": statusCode,
  101. "data": payload,
  102. }
  103. if !ok {
  104. resp["message"] = fmt.Sprintf("upstream status: %d", statusCode)
  105. }
  106. c.JSON(http.StatusOK, resp)
  107. }