diff --git a/controller/channel-billing-newapi.go b/controller/channel-billing-newapi.go new file mode 100644 index 000000000000..b65a745ea4df --- /dev/null +++ b/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() +} diff --git a/controller/channel-billing.go b/controller/channel-billing.go index 62982d2f5ceb..a1f60dfacbee 100644 --- a/controller/channel-billing.go +++ b/controller/channel-billing.go @@ -152,6 +152,7 @@ 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) } @@ -159,10 +160,6 @@ func GetResponseBody(method, url string, channel *model.Channel, headers http.He 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 diff --git a/controller/channel.go b/controller/channel.go index 3a1e58328923..7f61843805f7 100644 --- a/controller/channel.go +++ b/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()) diff --git a/controller/channel_authz.go b/controller/channel_authz.go index f85ffef92769..f1781ef11ce6 100644 --- a/controller/channel_authz.go +++ b/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 } diff --git a/controller/channel_authz_test.go b/controller/channel_authz_test.go index 0a57eac50dd7..dbe3ea1fe216 100644 --- a/controller/channel_authz_test.go +++ b/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() diff --git a/controller/channel_billing_balance_test.go b/controller/channel_billing_balance_test.go new file mode 100644 index 000000000000..63aac81ff7a9 --- /dev/null +++ b/controller/channel_billing_balance_test.go @@ -0,0 +1,121 @@ +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.Errorf("unexpected path %s", r.URL.Path) + http.Error(w, "unexpected path", http.StatusInternalServerError) + return + } + })) + 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: "aPerUnit, + 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: "aPerUnit, + QuotaDisplayType: "USD", + }) + assert.True(t, unlimited.Unlimited) + assert.Empty(t, unlimited.Remaining) + assert.Nil(t, legacy) +} + +func TestNormalizeNewAPIBalanceClampsNegativeRemaining(t *testing.T) { + credits, legacy := normalizeNewAPIBalance(decimal.NewFromInt(-1), false, nil) + assert.Equal(t, "0", credits.Remaining) + assert.Nil(t, legacy) + + quotaPerUnit := decimal.NewFromInt(500000) + usd, legacy := normalizeNewAPIBalance(decimal.NewFromInt(-500000), false, &newAPIStatusData{ + QuotaPerUnit: "aPerUnit, + QuotaDisplayType: "USD", + }) + assert.Equal(t, "0", usd.Remaining) + require.NotNil(t, legacy) + assert.Zero(t, *legacy) + assert.True(t, newAPIBalanceExhausted(&usd)) +} + +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) +} diff --git a/docs/channel/new-api-balance.md b/docs/channel/new-api-balance.md new file mode 100644 index 000000000000..92d809fa1aa2 --- /dev/null +++ b/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` 字段。与现有余额查询一致，多密钥渠道暂不支持。 diff --git a/model/channel.go b/model/channel.go index 2cd7c3115ff6..9d4c42f97987 100644 --- a/model/channel.go +++ b/model/channel.go @@ -21,25 +21,26 @@ import ( ) type Channel struct { - Id int `json:"id"` - Type int `json:"type" gorm:"default:0"` - Key string `json:"key" gorm:"not null"` - OpenAIOrganization *string `json:"openai_organization"` - TestModel *string `json:"test_model"` - Status int `json:"status" gorm:"default:1"` - Name string `json:"name" gorm:"index"` - Weight *uint `json:"weight" gorm:"default:0"` - CreatedTime int64 `json:"created_time" gorm:"bigint"` - TestTime int64 `json:"test_time" gorm:"bigint"` - ResponseTime int `json:"response_time"` // in milliseconds - BaseURL *string `json:"base_url" gorm:"column:base_url;default:''"` - Other string `json:"other"` - Balance float64 `json:"balance"` // in USD - BalanceUpdatedTime int64 `json:"balance_updated_time" gorm:"bigint"` - Models string `json:"models"` - Group string `json:"group" gorm:"type:varchar(64);default:'default'"` - UsedQuota int64 `json:"used_quota" gorm:"bigint;default:0"` - ModelMapping *string `json:"model_mapping" gorm:"type:text"` + Id int `json:"id"` + Type int `json:"type" gorm:"default:0"` + Key string `json:"key" gorm:"not null"` + OpenAIOrganization *string `json:"openai_organization"` + TestModel *string `json:"test_model"` + Status int `json:"status" gorm:"default:1"` + Name string `json:"name" gorm:"index"` + Weight *uint `json:"weight" gorm:"default:0"` + CreatedTime int64 `json:"created_time" gorm:"bigint"` + TestTime int64 `json:"test_time" gorm:"bigint"` + ResponseTime int `json:"response_time"` // in milliseconds + BaseURL *string `json:"base_url" gorm:"column:base_url;default:''"` + Other string `json:"other"` + Balance float64 `json:"balance"` // in USD + BalanceUpdatedTime int64 `json:"balance_updated_time" gorm:"bigint"` + BalanceInfo *ChannelBalanceInfo `json:"balance_info" gorm:"type:json"` + Models string `json:"models"` + Group string `json:"group" gorm:"type:varchar(64);default:'default'"` + UsedQuota int64 `json:"used_quota" gorm:"bigint;default:0"` + ModelMapping *string `json:"model_mapping" gorm:"type:text"` //MaxInputTokens *int `json:"max_input_tokens" gorm:"default:0"` StatusCodeMapping *string `json:"status_code_mapping" gorm:"type:varchar(1024);default:''"` Priority *int64 `json:"priority" gorm:"bigint;default:0"` diff --git a/model/channel_balance.go b/model/channel_balance.go new file mode 100644 index 000000000000..e00459010aad --- /dev/null +++ b/model/channel_balance.go @@ -0,0 +1,67 @@ +package model + +import ( + "database/sql/driver" + "fmt" + + "github.com/QuantumNous/new-api/common" +) + +const ( + ChannelBalanceUnitMoney = "money" + ChannelBalanceUnitTokens = "tokens" + ChannelBalanceUnitCredits = "credits" +) + +// ChannelBalanceInfo preserves the native unit returned by a New API upstream. +// The legacy Channel.Balance field remains USD-only. +type ChannelBalanceInfo struct { + Remaining string `json:"remaining,omitempty"` + Unit string `json:"unit,omitempty"` + Currency string `json:"currency,omitempty"` + DisplayUnit string `json:"display_unit,omitempty"` + Unlimited bool `json:"unlimited"` + UpdatedAt int64 `json:"updated_at"` +} + +func (info ChannelBalanceInfo) Value() (driver.Value, error) { + return common.Marshal(&info) +} + +func (info *ChannelBalanceInfo) Scan(value any) error { + if value == nil { + *info = ChannelBalanceInfo{} + return nil + } + var data []byte + switch typed := value.(type) { + case []byte: + data = typed + case string: + data = []byte(typed) + default: + return fmt.Errorf("unsupported channel balance info value type %T", value) + } + if len(data) == 0 { + *info = ChannelBalanceInfo{} + return nil + } + return common.Unmarshal(data, info) +} + +func (channel *Channel) UpdateBalanceInfo(info ChannelBalanceInfo, legacyBalanceUSD *float64) error { + updates := map[string]any{"balance_info": info} + if legacyBalanceUSD != nil { + updates["balance"] = *legacyBalanceUSD + updates["balance_updated_time"] = info.UpdatedAt + } + if err := DB.Model(channel).Updates(updates).Error; err != nil { + return err + } + channel.BalanceInfo = &info + if legacyBalanceUSD != nil { + channel.Balance = *legacyBalanceUSD + channel.BalanceUpdatedTime = info.UpdatedAt + } + return nil +} diff --git a/model/channel_balance_test.go b/model/channel_balance_test.go new file mode 100644 index 000000000000..bb22c9b7074c --- /dev/null +++ b/model/channel_balance_test.go @@ -0,0 +1,38 @@ +package model + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestChannelBalanceInfoPersistsNativeUnit(t *testing.T) { + require.NoError(t, DB.AutoMigrate(&Channel{}, &Ability{})) + const channelID = 970001 + require.NoError(t, DB.Unscoped().Delete(&Channel{}, "id = ?", channelID).Error) + t.Cleanup(func() { _ = DB.Unscoped().Delete(&Channel{}, "id = ?", channelID).Error }) + + baseURL := "https://upstream.example" + channel := &Channel{ + Id: channelID, Type: constant.ChannelTypeNewAPI, Key: "test-key", + Name: "balance persistence test", Status: common.ChannelStatusEnabled, BaseURL: &baseURL, + } + require.NoError(t, DB.Create(channel).Error) + + info := ChannelBalanceInfo{ + Remaining: "12.5", Unit: ChannelBalanceUnitMoney, + Currency: "USD", DisplayUnit: "$", UpdatedAt: 100, + } + legacyBalance := 12.5 + require.NoError(t, channel.UpdateBalanceInfo(info, &legacyBalance)) + + persisted, err := GetChannelById(channelID, true) + require.NoError(t, err) + require.NotNil(t, persisted.BalanceInfo) + assert.Equal(t, info, *persisted.BalanceInfo) + assert.Equal(t, legacyBalance, persisted.Balance) + assert.EqualValues(t, 100, persisted.BalanceUpdatedTime) +} diff --git a/web/src/features/channels/components/channels-columns.tsx b/web/src/features/channels/components/channels-columns.tsx index 36dc8f677625..0814cb4e95d3 100644 --- a/web/src/features/channels/components/channels-columns.tsx +++ b/web/src/features/channels/components/channels-columns.tsx @@ -56,10 +56,15 @@ import { formatTimestampToDate } from '@/lib/format' import { truncateText } from '@/lib/utils' import { getCodexUsage } from '../api' -import { CHANNEL_STATUS_CONFIG, MODEL_FETCHABLE_TYPES } from '../constants' +import { + CHANNEL_STATUS_CONFIG, + CHANNEL_TYPE_NEW_API, + MODEL_FETCHABLE_TYPES, +} from '../constants' import { formatRelativeTime, formatResponseTime, + formatNewAPIBalance, getBalanceVariant, getChannelTypeIcon, getChannelTypeLabel, @@ -333,6 +338,8 @@ function BalanceCell({ channel }: { channel: Channel }) { const isTagRow = isTagAggregateRow(channel) const balance = channel.balance || 0 const usedQuota = channel.used_quota || 0 + const newAPIBalance = + channel.type === CHANNEL_TYPE_NEW_API ? channel.balance_info : null const [isUpdating, setIsUpdating] = useState(false) const [codexUsageOpen, setCodexUsageOpen] = useState(false) const [codexUsageResponse, setCodexUsageResponse] = @@ -358,9 +365,12 @@ function BalanceCell({ channel }: { channel: Channel }) { showSymbol: layout !== 'card', }) ) - const remainingFull = withSuffix( + let remainingFull = withSuffix( formatCurrencyFromUSD(balance, balanceFormatOptions) ) + if (newAPIBalance) { + remainingFull = formatNewAPIBalance(newAPIBalance, t('Unlimited')) + } const usedDisplay = usedFull.length > MAX_INLINE_BALANCE_CHARS ? withSuffix( @@ -371,16 +381,16 @@ function BalanceCell({ channel }: { channel: Channel }) { }) ) : usedFull - const remainingDisplay = - remainingFull.length > MAX_INLINE_BALANCE_CHARS - ? withSuffix( - formatCurrencyFromUSD(balance, { - compact: true, - locale, - showSymbol: layout !== 'card', - }) - ) - : remainingFull + let remainingDisplay = remainingFull + if (!newAPIBalance && remainingFull.length > MAX_INLINE_BALANCE_CHARS) { + remainingDisplay = withSuffix( + formatCurrencyFromUSD(balance, { + compact: true, + locale, + showSymbol: layout !== 'card', + }) + ) + } const usedLabel = `${t('Used:')} ${usedFull}` const remainingLabel = `${t('Remaining:')} ${remainingFull}` const maskedUsedLabel = `${t('Used:')} ${SENSITIVE_MASK}` @@ -416,7 +426,7 @@ function BalanceCell({ channel }: { channel: Channel }) { } // Regular channel row: show used and remaining with click to update - const variant = getBalanceVariant(balance) + const variant = newAPIBalance ? 'success' : getBalanceVariant(balance) const handleClickUpdate = async () => { if (isUpdating) { @@ -1090,7 +1100,13 @@ export function useChannelsColumns( // Balance column (Used/Remaining) { - accessorKey: 'balance', + id: 'balance', + // Native-unit balances must not be sorted as if the legacy USD value + // represented the same quantity. + accessorFn: (channel) => + channel.type === CHANNEL_TYPE_NEW_API && channel.balance_info + ? null + : channel.balance, header: t('Used / Remaining'), cell: ({ row }) => , size: 180, diff --git a/web/src/features/channels/components/dialogs/balance-query-dialog.tsx b/web/src/features/channels/components/dialogs/balance-query-dialog.tsx index a9f6d11e314f..d8c054f4698e 100644 --- a/web/src/features/channels/components/dialogs/balance-query-dialog.tsx +++ b/web/src/features/channels/components/dialogs/balance-query-dialog.tsx @@ -17,7 +17,7 @@ along with this program. If not, see . For commercial licensing, please contact support@quantumnous.com */ import { useQueryClient } from '@tanstack/react-query' -import { Loader2, RefreshCw, DollarSign } from 'lucide-react' +import { DollarSign, Loader2, RefreshCw } from 'lucide-react' import { useEffect, useState } from 'react' import { useTranslation } from 'react-i18next' import { toast } from 'sonner' @@ -29,7 +29,9 @@ import { formatCurrencyFromUSD } from '@/lib/currency' import { formatTimestampToDate } from '@/lib/format' import { getCodexUsage, updateChannelBalance } from '../../api' -import { channelsQueryKeys } from '../../lib' +import { CHANNEL_TYPE_NEW_API } from '../../constants' +import { channelsQueryKeys, formatNewAPIBalance } from '../../lib' +import type { ChannelBalanceInfo } from '../../types' import { useChannels } from '../channels-provider' import { CodexUsageDialog, @@ -50,9 +52,10 @@ export function BalanceQueryDialog({ const queryClient = useQueryClient() const [isQuerying, setIsQuerying] = useState(false) const [balance, setBalance] = useState(null) - const [balanceUpdatedTime, setBalanceUpdatedTime] = useState( + const [balanceInfo, setBalanceInfo] = useState( null ) + const [balanceUpdatedAt, setBalanceUpdatedAt] = useState(null) const [codexUsageResponse, setCodexUsageResponse] = useState(null) @@ -90,25 +93,32 @@ export function BalanceQueryDialog({ setIsQuerying(true) try { const response = await updateChannelBalance(currentRow.id) - if (response.success && response.balance !== undefined) { - const newBalance = response.balance - const now = Math.floor(Date.now() / 1000) + const hasBalance = response.balance !== undefined + const hasStructuredData = response.data !== undefined + const hasPayload = hasBalance || hasStructuredData - setBalance(newBalance) - setBalanceUpdatedTime(now) - toast.success(t('Balance updated successfully')) + if (response.success && hasPayload) { + const now = Math.floor(Date.now() / 1000) + if (hasBalance) { + setBalance(response.balance ?? null) + } + if (response.data) { + setBalanceInfo(response.data) + } + setBalanceUpdatedAt(now) - // Update currentRow immediately with new balance and timestamp setCurrentRow({ ...currentRow, - balance: newBalance, - balance_updated_time: now, + balance: response.balance ?? currentRow.balance, + balance_info: response.data ?? currentRow.balance_info, + balance_updated_time: hasBalance + ? now + : currentRow.balance_updated_time, }) - - // Invalidate queries to refresh the table await queryClient.invalidateQueries({ queryKey: channelsQueryKeys.lists(), }) + toast.success(t('Balance updated successfully')) } else { toast.error(response.message || t('Failed to query balance')) } @@ -123,23 +133,12 @@ export function BalanceQueryDialog({ const handleClose = () => { setBalance(null) - setBalanceUpdatedTime(null) + setBalanceInfo(null) + setBalanceUpdatedAt(null) setCodexUsageResponse(null) onOpenChange(false) } - const formatBalance = (bal: number) => - formatCurrencyFromUSD(bal, { - digitsLarge: 2, - digitsSmall: 4, - abbreviate: false, - }) - - const formatDate = (timestamp: number) => { - if (!timestamp) return 'Never' - return formatTimestampToDate(timestamp) - } - if (isCodex) { return (
- {/* Current Balance Display */}
- + {t('Current Balance')}
-
- {balance !== null - ? formatBalance(balance) - : formatBalance(currentRow.balance)} -
+
{displayedAmount}
{t('Last updated:')}{' '} - {formatDate(balanceUpdatedTime ?? currentRow.balance_updated_time)} + {displayedUpdatedAt > 0 + ? formatTimestampToDate(displayedUpdatedAt) + : t('Never')}
diff --git a/web/src/features/channels/lib/__tests__/new-api-balance.test.ts b/web/src/features/channels/lib/__tests__/new-api-balance.test.ts new file mode 100644 index 000000000000..b46aa35a5c31 --- /dev/null +++ b/web/src/features/channels/lib/__tests__/new-api-balance.test.ts @@ -0,0 +1,71 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import assert from 'node:assert/strict' +import { describe, test } from 'node:test' + +import type { Channel, ChannelBalanceInfo } from '../../types' +import { channelNeedsAttention } from '../channel-utils' +import { formatNewAPIBalance } from '../new-api-balance' + +const balance = ( + overrides: Partial = {} +): ChannelBalanceInfo => ({ + remaining: '123.45', + unit: 'money', + currency: 'USD', + unlimited: false, + updated_at: 1_786_000_000, + ...overrides, +}) + +describe('New API balance', () => { + test('formats money, credits and unlimited quota', () => { + assert.equal(formatNewAPIBalance(balance(), 'Unlimited'), '$123.45') + assert.equal( + formatNewAPIBalance( + balance({ + unit: 'credits', + currency: undefined, + display_unit: 'credits', + }), + 'Unlimited' + ), + '123.45 credits' + ) + assert.equal( + formatNewAPIBalance(balance({ unlimited: true }), '无限制'), + '无限制' + ) + }) + + test('does not apply the legacy USD warning to native balances', () => { + const channel = { + id: 9, + type: 60, + status: 1, + balance: 0.5, + balance_info: balance({ currency: 'CNY', remaining: '0.5' }), + } as Channel + assert.equal(channelNeedsAttention(channel), false) + assert.equal( + channelNeedsAttention({ ...channel, balance_info: null }), + true + ) + }) +}) diff --git a/web/src/features/channels/lib/channel-actions.ts b/web/src/features/channels/lib/channel-actions.ts index 7efee24d3da2..e9f2de2c2da3 100644 --- a/web/src/features/channels/lib/channel-actions.ts +++ b/web/src/features/channels/lib/channel-actions.ts @@ -42,6 +42,7 @@ import { } from '../api' import { CHANNEL_STATUS, ERROR_MESSAGES, SUCCESS_MESSAGES } from '../constants' import type { ChannelTestResponse, CopyChannelParams } from '../types' +import { formatNewAPIBalance } from './new-api-balance' // ============================================================================ // Query Keys @@ -372,19 +373,31 @@ export async function handleUpdateChannelBalance( ): Promise { try { const response = await updateChannelBalance(id) - if (response.success && response.balance !== undefined) { - const balance = response.balance + const hasPayload = + response.data !== undefined || response.balance !== undefined + if (response.success && hasPayload) { + queryClient?.invalidateQueries({ queryKey: channelsQueryKeys.lists() }) + let displayBalance = '-' + if (response.data) { + displayBalance = formatNewAPIBalance( + response.data, + i18next.t('Unlimited') + ) + } else if (response.balance !== undefined) { + displayBalance = formatCurrencyFromUSD(response.balance, { + digitsLarge: 2, + digitsSmall: 4, + abbreviate: false, + }) + } toast.success( i18next.t('Balance updated: {{balance}}', { - balance: formatCurrencyFromUSD(balance, { - digitsLarge: 2, - digitsSmall: 4, - abbreviate: false, - }), + balance: displayBalance, }) ) - queryClient?.invalidateQueries({ queryKey: channelsQueryKeys.lists() }) - onSuccess?.(balance) + if (response.balance !== undefined) { + onSuccess?.(response.balance) + } } else { toast.error(response.message || i18next.t('Failed to update balance')) } @@ -462,7 +475,9 @@ export async function handleBatchEnable( toast.error(response.message || i18next.t('Failed to enable channels')) } else if (failCount > 0) { toast.error( - i18next.t('{{count}} channel(s) failed to enable', { count: failCount }) + i18next.t('{{count}} channel(s) failed to enable', { + count: failCount, + }) ) } } catch { diff --git a/web/src/features/channels/lib/channel-utils.ts b/web/src/features/channels/lib/channel-utils.ts index 9424a8521b6f..7daac7b0f7c0 100644 --- a/web/src/features/channels/lib/channel-utils.ts +++ b/web/src/features/channels/lib/channel-utils.ts @@ -21,6 +21,7 @@ import { formatTimestampToDate } from '@/lib/format' import { CHANNEL_STATUS_CONFIG, + CHANNEL_TYPE_NEW_API, CHANNEL_TYPES, MULTI_KEY_STATUS_CONFIG, RESPONSE_TIME_CONFIG, @@ -561,8 +562,11 @@ export function channelNeedsAttention(channel: Channel): boolean { return true } - // Low balance (less than $1) - if (channel.balance > 0 && channel.balance < 1) { + // Legacy balances are USD. Structured native balances must not be compared + // against the legacy one-dollar threshold. + const hasNewAPIBalance = + channel.type === CHANNEL_TYPE_NEW_API && channel.balance_info + if (!hasNewAPIBalance && channel.balance > 0 && channel.balance < 1) { return true } @@ -586,7 +590,9 @@ export function getAttentionReason(channel: Channel): string | null { if (channel.status === 3) { return 'Auto-disabled' } - if (channel.balance > 0 && channel.balance < 1) { + const hasNewAPIBalance = + channel.type === CHANNEL_TYPE_NEW_API && channel.balance_info + if (!hasNewAPIBalance && channel.balance > 0 && channel.balance < 1) { return 'Low balance' } if ( diff --git a/web/src/features/channels/lib/index.ts b/web/src/features/channels/lib/index.ts index 8c18151ceb5f..aada1dcddc8b 100644 --- a/web/src/features/channels/lib/index.ts +++ b/web/src/features/channels/lib/index.ts @@ -18,6 +18,7 @@ For commercial licensing, please contact support@quantumnous.com */ // Re-export all library functions export * from './channel-actions' +export * from './new-api-balance' export * from './channel-field-update' export * from './advanced-custom' export * from './channel-form-errors' diff --git a/web/src/features/channels/lib/new-api-balance.ts b/web/src/features/channels/lib/new-api-balance.ts new file mode 100644 index 000000000000..c8460e1046c7 --- /dev/null +++ b/web/src/features/channels/lib/new-api-balance.ts @@ -0,0 +1,43 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import type { ChannelBalanceInfo } from '../types' + +const CURRENCY_SYMBOLS: Record = { + CNY: '¥', + EUR: '€', + GBP: '£', + JPY: '¥', + KRW: '₩', + USD: '$', +} + +export function formatNewAPIBalance( + info: ChannelBalanceInfo, + unlimitedLabel: string +): string { + if (info.unlimited) return unlimitedLabel + const amount = info.remaining?.trim() + if (!amount) return '-' + if (info.unit === 'money') { + const currency = info.currency?.toUpperCase() || '' + const symbol = info.display_unit?.trim() || CURRENCY_SYMBOLS[currency] || '' + return `${symbol}${amount}`.trim() + } + return info.display_unit ? `${amount} ${info.display_unit}` : amount +} diff --git a/web/src/features/channels/types.ts b/web/src/features/channels/types.ts index f7747fa21210..df97f31702b0 100644 --- a/web/src/features/channels/types.ts +++ b/web/src/features/channels/types.ts @@ -34,6 +34,20 @@ export const channelInfoSchema = z.object({ export type ChannelInfo = z.infer +export const channelBalanceUnitSchema = z.enum(['money', 'tokens', 'credits']) + +export const channelBalanceInfoSchema = z.object({ + remaining: z.string().nullish(), + unit: channelBalanceUnitSchema.nullish(), + currency: z.string().nullish(), + display_unit: z.string().nullish(), + unlimited: z.boolean(), + updated_at: z.number(), +}) + +export type ChannelBalanceUnit = z.infer +export type ChannelBalanceInfo = z.infer + export const channelSchema = z.object({ id: z.number(), type: z.number(), @@ -50,6 +64,7 @@ export const channelSchema = z.object({ other: z.string().default(''), balance: z.number().default(0), // in USD balance_updated_time: z.number(), + balance_info: channelBalanceInfoSchema.nullish(), models: z.string().default(''), group: z.string().default('default'), used_quota: z.number().default(0), @@ -197,6 +212,7 @@ export interface ChannelBalanceResponse { message?: string balance?: number currency?: string + data?: ChannelBalanceInfo } export interface FetchModelsResponse {