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

controller/channel-billing-newapi.go

@@ -0,0 +1,140 @@
+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) {
+	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

@@ -439,6 +439,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 +476,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()
+	addChannelRequest.Channel.BalanceInfo = nil
 	keys := make([]string, 0)
 	switch addChannelRequest.Mode {
 	case "multi_to_single":
@@ -1432,8 +1433,11 @@ func CopyChannel(c *gin.Context) {
 	clone.Name = origin.Name + suffix
 	clone.TestTime = 0
 	clone.ResponseTime = 0
+	// New API balance snapshots belong to one channel and must not be copied.
+	clone.BalanceInfo = nil
 	if resetBalance {
 		clone.Balance = 0
+		clone.BalanceUpdatedTime = 0
 		clone.UsedQuota = 0
 	}
 

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,6 +128,7 @@ 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)

controller/channel_billing_balance_test.go

@@ -0,0 +1,103 @@
+package controller
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strconv"
+	"testing"
+
+	"github.com/QuantumNous/new-api/common"
+	"github.com/QuantumNous/new-api/constant"
+	"github.com/QuantumNous/new-api/model"
+	"github.com/gin-gonic/gin"
+	"github.com/shopspring/decimal"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestUpdateChannelBalanceRefreshesOnlyRequestedNewAPIChannel(t *testing.T) {
+	db := openTokenControllerTestDB(t)
+	require.NoError(t, db.AutoMigrate(&model.Channel{}))
+
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "Bearer upstream-token", r.Header.Get("Authorization"))
+		switch r.URL.Path {
+		case "/api/usage/token/":
+			writeControllerBalanceJSON(t, w, map[string]any{
+				"code": true,
+				"data": map[string]any{"total_available": 3625000, "unlimited_quota": false, "expires_at": 0},
+			})
+		case "/api/status":
+			writeControllerBalanceJSON(t, w, map[string]any{
+				"success": true,
+				"data":    map[string]any{"quota_per_unit": 500000, "quota_display_type": "USD"},
+			})
+		default:
+			t.Fatalf("unexpected path %s", r.URL.Path)
+		}
+	}))
+	defer upstream.Close()
+
+	baseURL := upstream.URL
+	requested := &model.Channel{Id: 980001, Type: constant.ChannelTypeNewAPI, Key: "upstream-token", Name: "requested", Status: common.ChannelStatusEnabled, BaseURL: &baseURL}
+	untouched := &model.Channel{Id: 980002, Type: constant.ChannelTypeNewAPI, Key: "upstream-token", Name: "untouched", Status: common.ChannelStatusEnabled, BaseURL: &baseURL}
+	require.NoError(t, db.Create(requested).Error)
+	require.NoError(t, db.Create(untouched).Error)
+
+	recorder := httptest.NewRecorder()
+	ctx, _ := gin.CreateTestContext(recorder)
+	ctx.Request = httptest.NewRequest(http.MethodGet, fmt.Sprintf("/api/channel/update_balance/%d", requested.Id), nil)
+	ctx.Params = gin.Params{{Key: "id", Value: strconv.Itoa(requested.Id)}}
+	UpdateChannelBalance(ctx)
+
+	var response struct {
+		Success bool                      `json:"success"`
+		Balance float64                   `json:"balance"`
+		Data    *model.ChannelBalanceInfo `json:"data"`
+	}
+	require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
+	require.True(t, response.Success)
+	assert.Equal(t, 7.25, response.Balance)
+	require.NotNil(t, response.Data)
+	assert.Equal(t, "7.25", response.Data.Remaining)
+
+	persisted, err := model.GetChannelById(requested.Id, true)
+	require.NoError(t, err)
+	require.NotNil(t, persisted.BalanceInfo)
+	assert.Equal(t, "7.25", persisted.BalanceInfo.Remaining)
+	persisted, err = model.GetChannelById(untouched.Id, true)
+	require.NoError(t, err)
+	assert.Nil(t, persisted.BalanceInfo)
+}
+
+func TestNormalizeNewAPIBalancePreservesUpstreamDisplay(t *testing.T) {
+	quotaPerUnit := decimal.NewFromInt(500000)
+	exchangeRate := decimal.RequireFromString("7.3")
+
+	cny, legacy := normalizeNewAPIBalance(decimal.NewFromInt(2500000), false, &newAPIStatusData{
+		QuotaPerUnit:     &quotaPerUnit,
+		QuotaDisplayType: "CNY",
+		USDExchangeRate:  &exchangeRate,
+	})
+	assert.Equal(t, model.ChannelBalanceUnitMoney, cny.Unit)
+	assert.Equal(t, "CNY", cny.Currency)
+	assert.Equal(t, "36.5", cny.Remaining)
+	assert.Nil(t, legacy)
+
+	unlimited, legacy := normalizeNewAPIBalance(decimal.Zero, true, &newAPIStatusData{
+		QuotaPerUnit:     &quotaPerUnit,
+		QuotaDisplayType: "USD",
+	})
+	assert.True(t, unlimited.Unlimited)
+	assert.Empty(t, unlimited.Remaining)
+	assert.Nil(t, legacy)
+}
+
+func writeControllerBalanceJSON(t *testing.T, writer http.ResponseWriter, value any) {
+	t.Helper()
+	payload, err := common.Marshal(value)
+	require.NoError(t, err)
+	_, err = writer.Write(payload)
+	require.NoError(t, err)
+}

docs/channel/new-api-balance.md

@@ -0,0 +1,16 @@
+# New API 上游余额查询
+
+该功能只为 `New API` 渠道（类型 `60`）补充余额查询，不改变其他渠道的查询逻辑，也不新增公共 API。
+
+平台使用渠道中配置的 Bearer Token 请求上游：
+
+- `GET {base_url}/api/usage/token/`：读取 Token 剩余额度与无限额度状态。
+- `GET {base_url}/api/status`：读取额度单位和换算设置；不可用时按原始 `credits` 展示。
+
+管理员仍通过现有接口刷新：
+
+```http
+GET /api/channel/update_balance/{channel_id}
+```
+
+结果保存在渠道的 `balance_info` 中，支持金额、`credits`、`tokens` 和无限额度。只有 USD 金额会同步到原有的 `channel.balance` 字段。与现有余额查询一致，多密钥渠道暂不支持。