fix(billing): estimate unpriced aliases via billing fallback models; don't renew expired reservations

- Inflight estimate now mirrors billableModelWithFallback: when the primary
  billing model has no price, fall back to the channel-mapped model (also for
  billing_model_source=requested) and to account-level mapped models (max over
  schedulable accounts in the group, cached 30s).
- Composite groups also consider concrete account-mapped models (max) so the
  alias' family fuzzy price cannot underestimate.
- Renew Lua refuses (and purges) members whose score <= now.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
William Drake
2026-09-28 16:49:29 -04:00
co-authored by Claude Opus 5.5
parent 52c123e6f4
commit 652a7d74f9
4 changed files with 179 additions and 24 deletions
@@ -56,10 +56,17 @@ var (
return {1, tostring(sum), count}
`)
// KEYS: [1]=zset [2]=hash ARGV: [1]=member [2]=expire_at_ms [3]=key_ttl_ms
// 仅当成员仍存在时续期(XX),并把两个 key 的 TTL 至少延长到 key_ttl_ms。
// KEYS: [1]=zset [2]=hash ARGV: [1]=member [2]=expire_at_ms [3]=key_ttl_ms [4]=now_ms
// 仅当成员仍存在且尚未过期(score > now)时续期(XX),并把两个 key 的 TTL 至少延长到 key_ttl_ms。
// 已过期但尚未被惰性清理的成员不得被复活:顺手清掉并返回 0。
renewInflightBalanceScript = redis.NewScript(`
if not redis.call('ZSCORE', KEYS[1], ARGV[1]) then
local score = redis.call('ZSCORE', KEYS[1], ARGV[1])
if not score then
return 0
end
if tonumber(score) <= tonumber(ARGV[4]) then
redis.call('ZREM', KEYS[1], ARGV[1])
redis.call('HDEL', KEYS[2], ARGV[1])
return 0
end
redis.call('ZADD', KEYS[1], 'XX', tonumber(ARGV[2]), ARGV[1])
@@ -123,7 +130,8 @@ func (c *billingCache) RenewInflightBalance(ctx context.Context, userID int64, r
if ttlMs <= 0 {
ttlMs = 1
}
n, err := renewInflightBalanceScript.Run(ctx, c.rdb, []string{zkey, hkey}, requestID, time.Now().UnixMilli()+ttlMs, ttlMs).Int64()
now := time.Now().UnixMilli()
n, err := renewInflightBalanceScript.Run(ctx, c.rdb, []string{zkey, hkey}, requestID, now+ttlMs, ttlMs, now).Int64()
if err != nil {
return false, err
}
@@ -267,3 +267,21 @@ func TestInflightReservation_RenewalKeepsStreamingReservationAlive(t *testing.T)
require.NoError(t, err)
require.False(t, ok, "renew must not resurrect a released reservation")
}
func TestInflightReservation_RenewDoesNotResurrectExpiredMember(t *testing.T) {
_, cache, _ := newInflightTestEnv(t, true, 60)
ctx := context.Background()
userID := int64(79)
zkey, hkey := billingInflightKeys(userID)
// 已过期但尚未被惰性清理的成员(score 在过去)。
require.NoError(t, cache.rdb.ZAdd(ctx, zkey, redis.Z{Score: float64(time.Now().UnixMilli() - 1000), Member: "stale"}).Err())
require.NoError(t, cache.rdb.HSet(ctx, hkey, "stale", "0.5").Err())
ok, err := cache.RenewInflightBalance(ctx, userID, "stale", time.Minute)
require.NoError(t, err)
require.False(t, ok, "renew must not resurrect an expired member")
require.Equal(t, int64(0), inflightCount(t, cache, userID))
exists, err := cache.rdb.HExists(ctx, hkey, "stale").Result()
require.NoError(t, err)
require.False(t, exists)
}
@@ -2,7 +2,9 @@ package service
import (
"context"
"fmt"
"math"
"strings"
"sync"
"sync/atomic"
"time"
@@ -11,6 +13,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
"github.com/google/uuid"
gocache "github.com/patrickmn/go-cache"
)
// InflightBalanceReservationCache 余额在途预留的缓存能力(可选)。
@@ -329,6 +332,51 @@ type inflightEstimateDeps struct {
resolver *ModelPricingResolver
resolveMapping func(ctx context.Context, groupID int64, model string) ChannelMappingResult
userGroupRate func(ctx context.Context, userID, groupID int64, groupDefault float64) float64
// accountMappedModels 返回分组内可调度账号对 model 的账号级映射结果(去重,不含 model 本身)。
// 准入时账号尚未选定,计费侧 billableModelWithFallback 会回退到实际转发模型(UpstreamModel,
// 即账号映射后的模型),因此这里取所有候选映射模型的最高估算。
accountMappedModels func(ctx context.Context, groupID int64, model string) []string
}
// inflightAccountMappingCache 缓存「分组 + 模型 → 账号级映射模型集合」,避免未定价别名每次请求都查库。
var inflightAccountMappingCache = gocache.New(30*time.Second, time.Minute)
func inflightAccountMappedModelsFromRepo(repo AccountRepository) func(ctx context.Context, groupID int64, model string) []string {
if repo == nil {
return nil
}
return func(ctx context.Context, groupID int64, model string) []string {
model = strings.TrimSpace(model)
if groupID <= 0 || model == "" {
return nil
}
key := fmt.Sprintf("%d|%s", groupID, model)
if v, ok := inflightAccountMappingCache.Get(key); ok {
if out, ok := v.([]string); ok {
return out
}
}
accounts, err := repo.ListSchedulableByGroupID(ctx, groupID)
if err != nil {
return nil
}
seen := map[string]struct{}{}
var out []string
for i := range accounts {
mapped, matched := accounts[i].ResolveMappedModel(model)
mapped = strings.TrimSpace(mapped)
if !matched || mapped == "" || mapped == model {
continue
}
if _, dup := seen[mapped]; dup {
continue
}
seen[mapped] = struct{}{}
out = append(out, mapped)
}
inflightAccountMappingCache.Set(key, out, gocache.DefaultExpiration)
return out
}
}
var inflightUnpricedLastLog atomic.Int64
@@ -347,26 +395,32 @@ func logInflightUnpriced(model string, groupID *int64) {
}
// inflightBillingModelCandidates 与计费路径一致地挑选计费模型:
// channel_mapped(默认)→ 映射后模型;requested → 请求模型;
// upstream / response_model 在准入时未知 → 取请求模型与映射模型两者较高估算。
func inflightBillingModelCandidates(ctx context.Context, deps inflightEstimateDeps, apiKey *APIKey, model string) []string {
if apiKey == nil || apiKey.GroupID == nil || deps.resolveMapping == nil {
return []string{model}
// 返回 primary(计费侧首选的计费模型)与 fallbacks(计费侧 billableModelWithFallback
// 在首选模型查无价时回退到的实际转发模型:渠道映射模型 → 账号级映射模型)。
// - channel_mapped(默认)→ 映射后模型(同时估算请求模型,取较高者);
// - requested → 请求模型;
// - upstream / response_model 在准入时未知 → 取请求模型与映射模型两者较高估算。
func inflightBillingModelCandidates(ctx context.Context, deps inflightEstimateDeps, apiKey *APIKey, model string) (primary, fallbacks []string) {
if apiKey == nil || apiKey.GroupID == nil {
return []string{model}, nil
}
m := deps.resolveMapping(ctx, *apiKey.GroupID, model)
mapped := m.MappedModel
if mapped == "" || mapped == model {
return []string{model}
upstreamInput := model
primary = []string{model}
if deps.resolveMapping != nil {
m := deps.resolveMapping(ctx, *apiKey.GroupID, model)
if mapped := m.MappedModel; mapped != "" && mapped != model {
upstreamInput = mapped
if m.BillingModelSource == BillingModelSourceRequested {
fallbacks = append(fallbacks, mapped)
} else {
primary = []string{mapped, model}
}
}
}
switch m.BillingModelSource {
case BillingModelSourceRequested:
return []string{model}
case BillingModelSourceChannelMapped, "":
// 计费侧对无价映射模型会回落到具体转发模型;此处同时估算请求模型,取较高者。
return []string{mapped, model}
default:
return []string{mapped, model}
if deps.accountMappedModels != nil {
fallbacks = append(fallbacks, deps.accountMappedModels(ctx, *apiKey.GroupID, upstreamInput)...)
}
return primary, fallbacks
}
func (d inflightEstimateDeps) rates(ctx context.Context, apiKey *APIKey) (text, image float64) {
@@ -529,9 +583,23 @@ func (d inflightEstimateDeps) estimate(ctx context.Context, apiKey *APIKey, req
// 免费分组:不计费,也无需预留。
return 0, true
}
best := 0.0
for _, m := range inflightBillingModelCandidates(ctx, d, apiKey, req.Model) {
if c := d.estimateOne(ctx, apiKey, m, req, textRate, imageRate); c > best {
primary, fallbacks := inflightBillingModelCandidates(ctx, d, apiKey, req.Model)
bestOf := func(models []string) float64 {
best := 0.0
for _, m := range models {
if c := d.estimateOne(ctx, apiKey, m, req, textRate, imageRate); c > best {
best = c
}
}
return best
}
best := bestOf(primary)
// composite 分组:计费侧除非别名有显式渠道价,否则按实际转发的具体模型计费;
// 别名本身可能命中家族模糊价(低估),因此与候选具体模型一起取最高。
composite := apiKey.Group != nil && apiKey.Group.Platform == PlatformComposite
if best <= 0 || composite {
// 与 billableModelWithFallback 同口径:首选模型无价时回退到实际转发模型。
if c := bestOf(fallbacks); c > best {
best = c
}
}
@@ -578,6 +646,7 @@ func (s *GatewayService) inflightEstimateDeps() inflightEstimateDeps {
d.resolveMapping = s.channelService.ResolveChannelMapping
}
d.userGroupRate = s.getUserGroupRateMultiplier
d.accountMappedModels = inflightAccountMappedModelsFromRepo(s.accountRepo)
return d
}
@@ -596,6 +665,7 @@ func (s *OpenAIGatewayService) inflightEstimateDeps() inflightEstimateDeps {
d.resolveMapping = s.channelService.ResolveChannelMapping
}
d.userGroupRate = s.ResolveUserGroupRateMultiplier
d.accountMappedModels = inflightAccountMappedModelsFromRepo(s.accountRepo)
return d
}
@@ -4,6 +4,7 @@ package service
import (
"context"
"math"
"sync"
"sync/atomic"
"testing"
@@ -321,3 +322,61 @@ func TestSyncBalanceCacheAfterDeduction_SynchronousWhenInflightEnabled(t *testin
bal, _ := cache.GetUserBalance(context.Background(), 3)
require.InDelta(t, 0.75, bal, 1e-12)
}
// 计费侧:requested 来源下别名本身无价 → billableModelWithFallback 回退到映射模型。
// 准入估算必须同口径地 > 0,且等于按回退模型的估算。
func TestInflightEstimate_RequestedSourceUnpricedAliasFallsBackLikeBilling(t *testing.T) {
groupID := int64(30)
ch := Channel{
ID: 3,
Status: StatusActive,
GroupIDs: []int64{groupID},
BillingModelSource: BillingModelSourceRequested,
ModelMapping: map[string]map[string]string{
"anthropic": {"req-alias": "claude-sonnet-4-5"},
},
}
cs := newTestChannelService(makeStandardRepo(ch, map[int64]string{groupID: "anthropic"}))
svc := newInflightEstimateGateway(t, cs)
apiKey := &APIKey{User: &User{ID: 1}, GroupID: &groupID, Group: &Group{ID: groupID, Platform: PlatformAnthropic, RateMultiplier: 1}}
ctx := context.Background()
billed := svc.billableModelWithFallback(ctx, apiKey, "req-alias", "claude-sonnet-4-5", "req-alias")
require.Equal(t, "claude-sonnet-4-5", billed, "precondition: billing falls back to the mapped model")
est, priced := svc.EstimateInflightReservation(ctx, apiKey, InflightEstimateRequest{Model: "req-alias", BodyBytes: 4000, MaxTokens: 1000})
require.True(t, priced)
require.Greater(t, est, 0.0)
direct, _ := svc.EstimateInflightReservation(ctx, apiKey, InflightEstimateRequest{Model: billed, BodyBytes: 4000, MaxTokens: 1000})
require.InDelta(t, direct, est, 1e-12)
}
// 计费侧:别名仅在账号级映射(无渠道映射)→ UpstreamModel 为账号映射模型,计费回退到它。
// 准入时账号未选定:按分组内候选账号映射模型的最高估算。
func TestInflightEstimate_AccountLevelMappingFallsBackLikeBilling(t *testing.T) {
groupID := int64(31)
svc := newInflightEstimateGateway(t, nil)
svc.accountRepo = &modelsListAccountRepoStub{byGroup: map[int64][]Account{groupID: {
{ID: 1, Platform: PlatformAnthropic, Credentials: map[string]any{"model_mapping": map[string]any{"acct-alias-31": "claude-sonnet-4-5"}}},
{ID: 2, Platform: PlatformAnthropic, Credentials: map[string]any{"model_mapping": map[string]any{"acct-alias-31": "claude-opus-4-1"}}},
}}}
apiKey := &APIKey{User: &User{ID: 1}, GroupID: &groupID, Group: &Group{ID: groupID, Platform: PlatformAnthropic, RateMultiplier: 1}}
ctx := context.Background()
_, ok := EstimateInflightReservationCost(svc.billingService, svc.cfg.Billing.InflightReservation, "acct-alias-31", 4000, 1000, 1)
require.False(t, ok, "precondition: alias has no pricing")
est, priced := svc.EstimateInflightReservation(ctx, apiKey, InflightEstimateRequest{Model: "acct-alias-31", BodyBytes: 4000, MaxTokens: 1000})
require.True(t, priced)
req := InflightEstimateRequest{BodyBytes: 4000, MaxTokens: 1000}
best := 0.0
for _, upstream := range []string{"claude-sonnet-4-5", "claude-opus-4-1"} {
billed := svc.billableModelWithFallback(ctx, apiKey, "acct-alias-31", upstream, "acct-alias-31")
require.Equal(t, upstream, billed, "precondition: billing charges the account-mapped model")
req.Model = billed
c, _ := svc.EstimateInflightReservation(ctx, apiKey, req)
require.Greater(t, c, 0.0)
best = math.Max(best, c)
}
require.InDelta(t, best, est, 1e-12, "estimate = max over candidate account-mapped billing models")
}