fix(billing): read account-level mappings from scheduler snapshot only when unpriced

- only consult account mappings after primary/channel candidates fail to price (or composite) - source mappings from the scheduler snapshot (same pool as account selection), no direct DB query and no per-model cache - ungrouped API keys use the scheduler's ungrouped pool - no negative caching: new mappings take effect on the next request Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-10-14 03:52:26 +08:00 · 2026-09-28 17:08:13 -04:00
parent 652a7d74f9
commit b886fb8edb
2 changed files with 199 additions and 45 deletions
@@ -2,7 +2,6 @@ package service

 import (
 	"context"
-	"fmt"
 	"math"
 	"strings"
 	"sync"
@@ -13,7 +12,6 @@ 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 余额在途预留的缓存能力（可选）。
@@ -332,53 +330,61 @@ 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 本身）。
+	// accountMappedModels 返回调度器可选账号对 model 的账号级映射结果（去重，不含 model 本身）。
 	// 准入时账号尚未选定，计费侧 billableModelWithFallback 会回退到实际转发模型（UpstreamModel，
 	// 即账号映射后的模型），因此这里取所有候选映射模型的最高估算。
-	accountMappedModels func(ctx context.Context, groupID int64, model string) []string
+	// 仅在首选/渠道候选均无法定价（或 composite 分组）时才调用；实现只读调度器快照，
+	// 不在请求路径上直接查库，也不按模型名缓存（内存不随请求模型名增长）。
+	accountMappedModels func(ctx context.Context, apiKey *APIKey, model string) []string
 }

-// inflightAccountMappingCache 缓存「分组 + 模型 → 账号级映射模型集合」，避免未定价别名每次请求都查库。
-var inflightAccountMappingCache = gocache.New(30*time.Second, time.Minute)
+// inflightSnapshotLister 调度器快照读取（与账号选择同源：Redis 快照，未命中时由快照服务自身回源并回填）。
+type inflightSnapshotLister func(ctx context.Context, groupID *int64, platform string, hasForcePlatform bool) ([]Account, error)

-func inflightAccountMappedModelsFromRepo(repo AccountRepository) func(ctx context.Context, groupID int64, model string) []string {
-	if repo == nil {
+// inflightAccountMappedModelsFromSnapshot 基于调度器快照在内存中匹配账号级映射（支持通配规则）。
+// resolvePlatform 返回与调度器一致的平台；无分组 API Key 使用调度器的未分组账号池（groupID=nil）。
+func inflightAccountMappedModelsFromSnapshot(list inflightSnapshotLister, resolvePlatform func(ctx context.Context, apiKey *APIKey, model string) (string, bool, bool)) func(ctx context.Context, apiKey *APIKey, model string) []string {
+	if list == nil || resolvePlatform == nil {
 		return nil
 	}
-	return func(ctx context.Context, groupID int64, model string) []string {
+	return func(ctx context.Context, apiKey *APIKey, model string) []string {
 		model = strings.TrimSpace(model)
-		if groupID <= 0 || model == "" {
+		if apiKey == nil || 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
-			}
+		platform, forced, ok := resolvePlatform(ctx, apiKey, model)
+		if !ok || platform == "" {
+			return nil
 		}
-		accounts, err := repo.ListSchedulableByGroupID(ctx, groupID)
+		accounts, err := list(ctx, apiKey.GroupID, platform, forced)
 		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
+		return accountMappedModelsFrom(accounts, model)
 	}
 }

+func accountMappedModelsFrom(accounts []Account, model string) []string {
+	var 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 seen == nil {
+			seen = map[string]struct{}{}
+		}
+		if _, dup := seen[mapped]; dup {
+			continue
+		}
+		seen[mapped] = struct{}{}
+		out = append(out, mapped)
+	}
+	return out
+}
+
 var inflightUnpricedLastLog atomic.Int64

 func logInflightUnpriced(model string, groupID *int64) {
@@ -400,12 +406,12 @@ func logInflightUnpriced(model string, groupID *int64) {
 //   - 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
-	}
-	upstreamInput := model
+func inflightBillingModelCandidates(ctx context.Context, deps inflightEstimateDeps, apiKey *APIKey, model string) (primary, fallbacks []string, upstreamInput string) {
+	upstreamInput = model
 	primary = []string{model}
+	if apiKey == nil || apiKey.GroupID == nil {
+		return primary, nil, upstreamInput
+	}
 	if deps.resolveMapping != nil {
 		m := deps.resolveMapping(ctx, *apiKey.GroupID, model)
 		if mapped := m.MappedModel; mapped != "" && mapped != model {
@@ -417,10 +423,7 @@ func inflightBillingModelCandidates(ctx context.Context, deps inflightEstimateDe
 			}
 		}
 	}
-	if deps.accountMappedModels != nil {
-		fallbacks = append(fallbacks, deps.accountMappedModels(ctx, *apiKey.GroupID, upstreamInput)...)
-	}
-	return primary, fallbacks
+	return primary, fallbacks, upstreamInput
 }

 func (d inflightEstimateDeps) rates(ctx context.Context, apiKey *APIKey) (text, image float64) {
@@ -583,7 +586,7 @@ func (d inflightEstimateDeps) estimate(ctx context.Context, apiKey *APIKey, req
 		// 免费分组：不计费，也无需预留。
 		return 0, true
 	}
-	primary, fallbacks := inflightBillingModelCandidates(ctx, d, apiKey, req.Model)
+	primary, fallbacks, upstreamInput := inflightBillingModelCandidates(ctx, d, apiKey, req.Model)
 	bestOf := func(models []string) float64 {
 		best := 0.0
 		for _, m := range models {
@@ -602,6 +605,12 @@ func (d inflightEstimateDeps) estimate(ctx context.Context, apiKey *APIKey, req
 		if c := bestOf(fallbacks); c > best {
 			best = c
 		}
+		// 账号级映射候选（读调度器快照）仅在仍无法定价或 composite 时才查，已定价模型不触发。
+		if (best <= 0 || composite) && d.accountMappedModels != nil {
+			if c := bestOf(d.accountMappedModels(ctx, apiKey, upstreamInput)); c > best {
+				best = c
+			}
+		}
 	}
 	if best <= 0 {
 		logInflightUnpriced(req.Model, apiKey.GroupID)
@@ -646,7 +655,19 @@ func (s *GatewayService) inflightEstimateDeps() inflightEstimateDeps {
 		d.resolveMapping = s.channelService.ResolveChannelMapping
 	}
 	d.userGroupRate = s.getUserGroupRateMultiplier
-	d.accountMappedModels = inflightAccountMappedModelsFromRepo(s.accountRepo)
+	if s.schedulerSnapshot != nil {
+		snap := s.schedulerSnapshot
+		d.accountMappedModels = inflightAccountMappedModelsFromSnapshot(
+			func(ctx context.Context, groupID *int64, platform string, forced bool) ([]Account, error) {
+				accounts, _, err := snap.ListSchedulableAccounts(ctx, groupID, platform, forced)
+				return accounts, err
+			},
+			func(ctx context.Context, apiKey *APIKey, model string) (string, bool, bool) {
+				platform, forced, err := s.resolvePlatform(ctx, apiKey.GroupID, apiKey.Group, model)
+				return platform, forced, err == nil
+			},
+		)
+	}
 	return d
 }

@@ -665,7 +686,22 @@ func (s *OpenAIGatewayService) inflightEstimateDeps() inflightEstimateDeps {
 		d.resolveMapping = s.channelService.ResolveChannelMapping
 	}
 	d.userGroupRate = s.ResolveUserGroupRateMultiplier
-	d.accountMappedModels = inflightAccountMappedModelsFromRepo(s.accountRepo)
+	if s.schedulerSnapshot != nil {
+		snap := s.schedulerSnapshot
+		d.accountMappedModels = inflightAccountMappedModelsFromSnapshot(
+			func(ctx context.Context, groupID *int64, platform string, forced bool) ([]Account, error) {
+				accounts, _, err := snap.ListSchedulableAccounts(ctx, groupID, platform, forced)
+				return accounts, err
+			},
+			func(ctx context.Context, apiKey *APIKey, model string) (string, bool, bool) {
+				platform := PlatformOpenAI
+				if apiKey.Group != nil && apiKey.Group.Platform != "" && apiKey.Group.Platform != PlatformComposite {
+					platform = apiKey.Group.Platform
+				}
+				return NormalizeOpenAICompatiblePlatform(platform), false, true
+			},
+		)
+	}
 	return d
 }

@@ -4,6 +4,8 @@ package service

 import (
 	"context"
+	"fmt"
+	"runtime"
 	"math"
 	"sync"
 	"sync/atomic"
@@ -356,10 +358,11 @@ func TestInflightEstimate_RequestedSourceUnpricedAliasFallsBackLikeBilling(t *te
 func TestInflightEstimate_AccountLevelMappingFallsBackLikeBilling(t *testing.T) {
 	groupID := int64(31)
 	svc := newInflightEstimateGateway(t, nil)
-	svc.accountRepo = &modelsListAccountRepoStub{byGroup: map[int64][]Account{groupID: {
+	snap := &inflightSnapshotCacheStub{byBucket: map[string][]Account{inflightBucketKey(groupID, PlatformAnthropic): {
 		{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"}}},
 	}}}
+	attachInflightSnapshot(svc, snap)
 	apiKey := &APIKey{User: &User{ID: 1}, GroupID: &groupID, Group: &Group{ID: groupID, Platform: PlatformAnthropic, RateMultiplier: 1}}
 	ctx := context.Background()

@@ -380,3 +383,118 @@ func TestInflightEstimate_AccountLevelMappingFallsBackLikeBilling(t *testing.T)
 	}
 	require.InDelta(t, best, est, 1e-12, "estimate = max over candidate account-mapped billing models")
 }
+
+type inflightSnapshotCacheStub struct {
+	SchedulerCache
+	mu       sync.Mutex
+	byBucket map[string][]Account
+	reads    atomic.Int64
+}
+
+func inflightBucketKey(groupID int64, platform string) string {
+	return fmt.Sprintf("%d|%s", groupID, platform)
+}
+
+func (c *inflightSnapshotCacheStub) set(groupID int64, platform string, accounts []Account) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.byBucket[inflightBucketKey(groupID, platform)] = accounts
+}
+
+func (c *inflightSnapshotCacheStub) GetSnapshot(ctx context.Context, bucket SchedulerBucket) ([]*Account, bool, error) {
+	c.reads.Add(1)
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	src := c.byBucket[inflightBucketKey(bucket.GroupID, bucket.Platform)]
+	out := make([]*Account, 0, len(src))
+	for i := range src {
+		a := src[i]
+		out = append(out, &a)
+	}
+	return out, true, nil
+}
+
+// inflightCountingAccountRepo 统计请求路径上的任何直接查库。
+type inflightCountingAccountRepo struct {
+	AccountRepository
+	dbCalls atomic.Int64
+}
+
+func (r *inflightCountingAccountRepo) ListSchedulableByGroupID(ctx context.Context, groupID int64) ([]Account, error) {
+	r.dbCalls.Add(1)
+	return nil, nil
+}
+
+func attachInflightSnapshot(svc *GatewayService, snap *inflightSnapshotCacheStub) *inflightCountingAccountRepo {
+	repo := &inflightCountingAccountRepo{}
+	svc.accountRepo = repo
+	svc.schedulerSnapshot = NewSchedulerSnapshotService(snap, nil, repo, nil, svc.cfg)
+	return repo
+}
+
+// 已定价模型永不查账号映射；随机未定价模型名不直接查库、不产生按模型名的缓存（内存有界）。
+func TestInflightEstimate_AccountMappingNoDBAndBoundedMemory(t *testing.T) {
+	groupID := int64(40)
+	svc := newInflightEstimateGateway(t, nil)
+	snap := &inflightSnapshotCacheStub{byBucket: map[string][]Account{inflightBucketKey(groupID, PlatformAnthropic): {
+		{ID: 1, Platform: PlatformAnthropic, Credentials: map[string]any{"model_mapping": map[string]any{"acct-alias-40": "claude-sonnet-4-5"}}},
+	}}}
+	repo := attachInflightSnapshot(svc, snap)
+	apiKey := &APIKey{User: &User{ID: 1}, GroupID: &groupID, Group: &Group{ID: groupID, Platform: PlatformAnthropic, RateMultiplier: 1}}
+	ctx := context.Background()
+
+	for i := 0; i < 100; i++ {
+		_, priced := svc.EstimateInflightReservation(ctx, apiKey, InflightEstimateRequest{Model: "claude-sonnet-4-5", BodyBytes: 4000, MaxTokens: 1000})
+		require.True(t, priced)
+	}
+	require.Zero(t, snap.reads.Load(), "priced models must never trigger account-mapping lookup")
+
+	var before, after runtime.MemStats
+	runtime.GC()
+	runtime.ReadMemStats(&before)
+	const n = 20000
+	for i := 0; i < n; i++ {
+		_, priced := svc.EstimateInflightReservation(ctx, apiKey, InflightEstimateRequest{Model: fmt.Sprintf("rand-%d-%d", i, time.Now().UnixNano()), BodyBytes: 4000, MaxTokens: 1000})
+		require.False(t, priced)
+	}
+	runtime.GC()
+	runtime.ReadMemStats(&after)
+	require.Zero(t, repo.dbCalls.Load(), "no direct DB query on the request path")
+	require.Equal(t, int64(n), snap.reads.Load(), "unpriced lookups read the scheduler snapshot only")
+	require.Less(t, int64(after.HeapAlloc)-int64(before.HeapAlloc), int64(8<<20), "no per-model cache growth")
+}
+
+// 无分组 API Key：使用调度器的未分组账号池，估算与计费回退的账号映射模型同口径。
+func TestInflightEstimate_NoGroupKeyUsesUngroupedAccountMapping(t *testing.T) {
+	svc := newInflightEstimateGateway(t, nil)
+	snap := &inflightSnapshotCacheStub{byBucket: map[string][]Account{inflightBucketKey(0, PlatformAnthropic): {
+		{ID: 1, Platform: PlatformAnthropic, Credentials: map[string]any{"model_mapping": map[string]any{"acct-alias-nog": "claude-opus-4-1"}}},
+	}}}
+	attachInflightSnapshot(svc, snap)
+	apiKey := &APIKey{User: &User{ID: 1}}
+	ctx := context.Background()
+
+	est, priced := svc.EstimateInflightReservation(ctx, apiKey, InflightEstimateRequest{Model: "acct-alias-nog", BodyBytes: 4000, MaxTokens: 1000})
+	require.True(t, priced)
+	direct, ok := svc.EstimateInflightReservation(ctx, apiKey, InflightEstimateRequest{Model: "claude-opus-4-1", BodyBytes: 4000, MaxTokens: 1000})
+	require.True(t, ok)
+	require.InDelta(t, direct, est, 1e-12)
+}
+
+// 管理员新增账号映射后立即生效（不缓存负结果）。
+func TestInflightEstimate_AccountMappingNoNegativeCaching(t *testing.T) {
+	groupID := int64(41)
+	svc := newInflightEstimateGateway(t, nil)
+	snap := &inflightSnapshotCacheStub{byBucket: map[string][]Account{}}
+	attachInflightSnapshot(svc, snap)
+	apiKey := &APIKey{User: &User{ID: 1}, GroupID: &groupID, Group: &Group{ID: groupID, Platform: PlatformAnthropic, RateMultiplier: 1}}
+	ctx := context.Background()
+	req := InflightEstimateRequest{Model: "acct-alias-41", BodyBytes: 4000, MaxTokens: 1000}
+
+	_, priced := svc.EstimateInflightReservation(ctx, apiKey, req)
+	require.False(t, priced)
+	snap.set(groupID, PlatformAnthropic, []Account{{ID: 9, Platform: PlatformAnthropic, Credentials: map[string]any{"model_mapping": map[string]any{"acct-*": "claude-sonnet-4-5"}}}})
+	est, priced := svc.EstimateInflightReservation(ctx, apiKey, req)
+	require.True(t, priced, "new mapping (incl. wildcard) visible on next request")
+	require.Greater(t, est, 0.0)
+}