feat(channel): add balance query for New API channels

controller/channel-billing-newapi.go

@@ -0,0 +1,143 @@
+package controller
+
+import (
+	"errors"
+	"fmt"
+	"net/http"
+	"net/url"
+	"strings"
+	"time"
+
+	"github.com/QuantumNous/new-api/common"
+	"github.com/QuantumNous/new-api/model"
+	"github.com/shopspring/decimal"
+)
+
+type newAPITokenUsageResponse struct {
+	Code bool                  `json:"code"`
+	Data *newAPITokenUsageData `json:"data"`
+}
+
+type newAPITokenUsageData struct {
+	TotalAvailable *decimal.Decimal `json:"total_available"`
+	UnlimitedQuota bool             `json:"unlimited_quota"`
+	ExpiresAt      int64            `json:"expires_at"`
+}
+
+type newAPIStatusResponse struct {
+	Success bool              `json:"success"`
+	Data    *newAPIStatusData `json:"data"`
+}
+
+type newAPIStatusData struct {
+	QuotaPerUnit               *decimal.Decimal `json:"quota_per_unit"`
+	QuotaDisplayType           string           `json:"quota_display_type"`
+	USDExchangeRate            *decimal.Decimal `json:"usd_exchange_rate"`
+	CustomCurrencySymbol       string           `json:"custom_currency_symbol"`
+	CustomCurrencyExchangeRate *decimal.Decimal `json:"custom_currency_exchange_rate"`
+}
+
+func updateChannelNewAPIBalance(channel *model.Channel) (*model.ChannelBalanceInfo, *float64, error) {
+	usageBody, err := getNewAPIChannelResponse(channel, "/api/usage/token/")
+	if err != nil {
+		return nil, nil, err
+	}
+	var usage newAPITokenUsageResponse
+	if err := common.Unmarshal(usageBody, &usage); err != nil {
+		return nil, nil, fmt.Errorf("invalid New API token usage response: %w", err)
+	}
+	if !usage.Code || usage.Data == nil || usage.Data.TotalAvailable == nil {
+		return nil, nil, errors.New("New API token usage response is invalid")
+	}
+	if usage.Data.ExpiresAt > 0 && usage.Data.ExpiresAt < time.Now().Unix() {
+		return nil, nil, errors.New("New API token is expired")
+	}
+
+	var status *newAPIStatusData
+	if statusBody, statusErr := getNewAPIChannelResponse(channel, "/api/status"); statusErr == nil {
+		var parsed newAPIStatusResponse
+		if common.Unmarshal(statusBody, &parsed) == nil && parsed.Success && parsed.Data != nil {
+			status = parsed.Data
+		}
+	}
+	info, legacyBalance := normalizeNewAPIBalance(*usage.Data.TotalAvailable, usage.Data.UnlimitedQuota, status)
+	if err := channel.UpdateBalanceInfo(info, legacyBalance); err != nil {
+		return nil, nil, err
+	}
+	return &info, legacyBalance, nil
+}
+
+func getNewAPIChannelResponse(channel *model.Channel, path string) ([]byte, error) {
+	baseURL, err := url.Parse(strings.TrimSpace(channel.GetBaseURL()))
+	if err != nil || (baseURL.Scheme != "http" && baseURL.Scheme != "https") || baseURL.Host == "" || baseURL.User != nil || baseURL.RawQuery != "" || baseURL.Fragment != "" {
+		return nil, errors.New("invalid New API channel base URL")
+	}
+	baseURL.Path = strings.TrimRight(baseURL.Path, "/") + path
+	baseURL.RawPath = ""
+	return GetResponseBody(http.MethodGet, baseURL.String(), channel, GetAuthHeader(channel.Key))
+}
+
+func normalizeNewAPIBalance(remaining decimal.Decimal, unlimited bool, status *newAPIStatusData) (model.ChannelBalanceInfo, *float64) {
+	if !unlimited && remaining.IsNegative() {
+		remaining = decimal.Zero
+	}
+	info := model.ChannelBalanceInfo{
+		Remaining:   remaining.String(),
+		Unit:        model.ChannelBalanceUnitCredits,
+		DisplayUnit: "credits",
+		Unlimited:   unlimited,
+		UpdatedAt:   common.GetTimestamp(),
+	}
+	if unlimited {
+		info.Remaining = ""
+	}
+	if status == nil || status.QuotaPerUnit == nil || !status.QuotaPerUnit.IsPositive() {
+		return info, nil
+	}
+
+	convert := func(multiplier decimal.Decimal) decimal.Decimal {
+		return remaining.Div(*status.QuotaPerUnit).Mul(multiplier)
+	}
+	var legacyBalance *float64
+	switch strings.ToUpper(status.QuotaDisplayType) {
+	case "USD":
+		amount := convert(decimal.NewFromInt(1))
+		info.Unit, info.Currency, info.DisplayUnit = model.ChannelBalanceUnitMoney, "USD", "$"
+		if !unlimited {
+			info.Remaining = amount.String()
+			value := amount.InexactFloat64()
+			legacyBalance = &value
+		}
+	case "CNY":
+		if status.USDExchangeRate != nil && status.USDExchangeRate.IsPositive() {
+			amount := convert(*status.USDExchangeRate)
+			info.Unit, info.Currency, info.DisplayUnit = model.ChannelBalanceUnitMoney, "CNY", "¥"
+			if !unlimited {
+				info.Remaining = amount.String()
+			}
+		}
+	case "TOKENS":
+		info.Unit, info.DisplayUnit = model.ChannelBalanceUnitTokens, "tokens"
+	case "CUSTOM":
+		if status.CustomCurrencyExchangeRate != nil && status.CustomCurrencyExchangeRate.IsPositive() {
+			amount := convert(*status.CustomCurrencyExchangeRate)
+			info.Unit, info.Currency = model.ChannelBalanceUnitMoney, "CUSTOM"
+			info.DisplayUnit = strings.TrimSpace(status.CustomCurrencySymbol)
+			if info.DisplayUnit == "" {
+				info.DisplayUnit = "¤"
+			}
+			if !unlimited {
+				info.Remaining = amount.String()
+			}
+		}
+	}
+	return info, legacyBalance
+}
+
+func newAPIBalanceExhausted(info *model.ChannelBalanceInfo) bool {
+	if info == nil || info.Unlimited || info.Remaining == "" {
+		return false
+	}
+	remaining, err := decimal.NewFromString(info.Remaining)
+	return err == nil && !remaining.IsPositive()
+}

controller/channel-billing.go

@@ -152,17 +152,14 @@ func GetResponseBody(method, url string, channel *model.Channel, headers http.He
 	if err != nil {
 		return nil, err
 	}
+	defer res.Body.Close()
 	if res.StatusCode != http.StatusOK {
 		return nil, fmt.Errorf("status code: %d", res.StatusCode)
 	}
 	body, err := io.ReadAll(res.Body)
 	if err != nil {
 		return nil, err
 	}
-	err = res.Body.Close()
-	if err != nil {
-		return nil, err
-	}
 	return body, nil
 }
 
@@ -439,6 +436,19 @@ func UpdateChannelBalance(c *gin.Context) {
 		})
 		return
 	}
+	if channel.Type == constant.ChannelTypeNewAPI {
+		info, legacyBalance, refreshErr := updateChannelNewAPIBalance(channel)
+		if refreshErr != nil {
+			common.ApiError(c, refreshErr)
+			return
+		}
+		response := gin.H{"success": true, "message": "", "data": info}
+		if legacyBalance != nil {
+			response["balance"] = *legacyBalance
+		}
+		c.JSON(http.StatusOK, response)
+		return
+	}
 	balance, err := updateChannelBalance(channel)
 	if err != nil {
 		common.ApiError(c, err)
@@ -463,6 +473,14 @@ func updateAllChannelsBalance() error {
 		if channel.ChannelInfo.IsMultiKey {
 			continue // skip multi-key channels
 		}
+		if channel.Type == constant.ChannelTypeNewAPI {
+			info, _, refreshErr := updateChannelNewAPIBalance(channel)
+			if refreshErr == nil && newAPIBalanceExhausted(info) {
+				service.DisableChannel(*types.NewChannelError(channel.Id, channel.Type, channel.Name, false, "", channel.GetAutoBan()), "余额不足")
+			}
+			time.Sleep(common.RequestInterval)
+			continue
+		}
 		// TODO: support Azure
 		//if channel.Type != common.ChannelTypeOpenAI && channel.Type != common.ChannelTypeCustom {
 		//	continue

controller/channel.go

@@ -627,6 +627,7 @@ func AddChannel(c *gin.Context) {
 	}
 
 	addChannelRequest.Channel.CreatedTime = common.GetTimestamp()
+	applyChannelBalanceReset(addChannelRequest.Channel, true)
 	keys := make([]string, 0)
 	switch addChannelRequest.Mode {
 	case "multi_to_single":
@@ -713,6 +714,16 @@ func AddChannel(c *gin.Context) {
 	return
 }
 
+func applyChannelBalanceReset(channel *model.Channel, reset bool) {
+	if channel == nil || !reset {
+		return
+	}
+	channel.Balance = 0
+	channel.BalanceUpdatedTime = 0
+	channel.BalanceInfo = nil
+	channel.UsedQuota = 0
+}
+
 func DeleteChannel(c *gin.Context) {
 	id, _ := strconv.Atoi(c.Param("id"))
 	channelName := ""
@@ -1432,10 +1443,7 @@ func CopyChannel(c *gin.Context) {
 	clone.Name = origin.Name + suffix
 	clone.TestTime = 0
 	clone.ResponseTime = 0
-	if resetBalance {
-		clone.Balance = 0
-		clone.UsedQuota = 0
-	}
+	applyChannelBalanceReset(&clone, resetBalance)
 
 	if err := clone.ValidateSettings(); err != nil {
 		common.SysError("failed to validate cloned channel: " + err.Error())

controller/channel_authz.go

@@ -87,6 +87,7 @@ var channelReadOnlyFields = map[string]struct{}{
 	"response_time":        {},
 	"balance":              {},
 	"balance_updated_time": {},
+	"balance_info":         {},
 	"used_quota":           {},
 }
 
@@ -106,6 +107,9 @@ func clearChannelReadOnlyFields(channel *PatchChannel, requestData map[string]an
 	if _, ok := requestData["balance_updated_time"]; ok {
 		channel.BalanceUpdatedTime = 0
 	}
+	if _, ok := requestData["balance_info"]; ok {
+		channel.BalanceInfo = nil
+	}
 	if _, ok := requestData["used_quota"]; ok {
 		channel.UsedQuota = 0
 	}

controller/channel_authz_test.go

@@ -85,11 +85,13 @@ func TestChannelHasSensitiveChanges(t *testing.T) {
 	t.Run("read-only fields are ignored by sensitivity check", func(t *testing.T) {
 		updated := PatchChannel{Channel: *origin}
 		updated.Balance = 99
+		updated.BalanceInfo = &model.ChannelBalanceInfo{Unit: model.ChannelBalanceUnitCredits}
 		updated.UsedQuota = 100
 		updated.ResponseTime = 200
 
 		assert.False(t, channelHasSensitiveChanges(&updated, origin, map[string]any{
 			"balance":       updated.Balance,
+			"balance_info":  updated.BalanceInfo,
 			"used_quota":    updated.UsedQuota,
 			"response_time": updated.ResponseTime,
 		}))
@@ -103,6 +105,7 @@ func TestClearChannelReadOnlyFields(t *testing.T) {
 		ResponseTime:       33,
 		Balance:            44.5,
 		BalanceUpdatedTime: 55,
+		BalanceInfo:        &model.ChannelBalanceInfo{Unit: model.ChannelBalanceUnitCredits},
 		UsedQuota:          66,
 		Models:             "gpt-4o",
 		Group:              "default",
@@ -114,6 +117,7 @@ func TestClearChannelReadOnlyFields(t *testing.T) {
 		"response_time":        channel.ResponseTime,
 		"balance":              channel.Balance,
 		"balance_updated_time": channel.BalanceUpdatedTime,
+		"balance_info":         channel.BalanceInfo,
 		"used_quota":           channel.UsedQuota,
 		"models":               channel.Models,
 		"group":                channel.Group,
@@ -124,11 +128,37 @@ func TestClearChannelReadOnlyFields(t *testing.T) {
 	assert.Zero(t, channel.ResponseTime)
 	assert.Zero(t, channel.Balance)
 	assert.Zero(t, channel.BalanceUpdatedTime)
+	assert.Nil(t, channel.BalanceInfo)
 	assert.Zero(t, channel.UsedQuota)
 	assert.Equal(t, "gpt-4o", channel.Models)
 	assert.Equal(t, "default", channel.Group)
 }
 
+func TestApplyChannelBalanceReset(t *testing.T) {
+	original := model.Channel{
+		Balance:            12.5,
+		BalanceUpdatedTime: 123,
+		BalanceInfo:        &model.ChannelBalanceInfo{Remaining: "12.5"},
+		UsedQuota:          456,
+	}
+
+	preserved := original
+	preservedInfo := *original.BalanceInfo
+	preserved.BalanceInfo = &preservedInfo
+	applyChannelBalanceReset(&preserved, false)
+	assert.Equal(t, original.Balance, preserved.Balance)
+	assert.Equal(t, original.BalanceUpdatedTime, preserved.BalanceUpdatedTime)
+	assert.Equal(t, original.BalanceInfo, preserved.BalanceInfo)
+	assert.Equal(t, original.UsedQuota, preserved.UsedQuota)
+
+	reset := original
+	applyChannelBalanceReset(&reset, true)
+	assert.Zero(t, reset.Balance)
+	assert.Zero(t, reset.BalanceUpdatedTime)
+	assert.Nil(t, reset.BalanceInfo)
+	assert.Zero(t, reset.UsedQuota)
+}
+
 func TestUpdateChannelRejectsStatusField(t *testing.T) {
 	gin.SetMode(gin.TestMode)
 	recorder := httptest.NewRecorder()