fix: preserve group pricing in auth snapshots

This commit is contained in:
wucm667
2026-08-14 00:25:55 +08:00
parent fbfdcef818
commit 674570ca17
4 changed files with 58 additions and 2 deletions
@@ -84,6 +84,8 @@ type APIKeyAuthGroupSnapshot struct {
AudioRealtimePricePerMin *float64 `json:"audio_realtime_price_per_min,omitempty"`
AudioTTSPricePerMillionChars *float64 `json:"audio_tts_price_per_million_chars,omitempty"`
AudioSTTPricePerHour *float64 `json:"audio_stt_price_per_hour,omitempty"`
LongContextPricingEnabled bool `json:"long_context_pricing_enabled"`
ModelPricing []ChannelModelPricing `json:"model_pricing,omitempty"`
ClaudeCodeOnly bool `json:"claude_code_only"`
FallbackGroupID *int64 `json:"fallback_group_id,omitempty"`
FallbackGroupIDOnInvalidRequest *int64 `json:"fallback_group_id_on_invalid_request,omitempty"`
@@ -14,7 +14,7 @@ import (
"github.com/dgraph-io/ristretto"
)
const apiKeyAuthSnapshotVersion = 19 // v19: group search/audio/video_model_prices billing fields (force refresh of pre-fix snapshots)
const apiKeyAuthSnapshotVersion = 20 // v20: group long-context and model pricing fields (force refresh of pre-fix snapshots)
type apiKeyAuthCacheConfig struct {
l1Size int
@@ -406,6 +406,8 @@ func (s *APIKeyService) snapshotFromAPIKey(ctx context.Context, apiKey *APIKey)
AudioRealtimePricePerMin: apiKey.Group.AudioRealtimePricePerMin,
AudioTTSPricePerMillionChars: apiKey.Group.AudioTTSPricePerMillionChars,
AudioSTTPricePerHour: apiKey.Group.AudioSTTPricePerHour,
LongContextPricingEnabled: apiKey.Group.LongContextPricingEnabled,
ModelPricing: apiKey.Group.ModelPricing,
ClaudeCodeOnly: apiKey.Group.ClaudeCodeOnly,
FallbackGroupID: apiKey.Group.FallbackGroupID,
FallbackGroupIDOnInvalidRequest: apiKey.Group.FallbackGroupIDOnInvalidRequest,
@@ -501,6 +503,8 @@ func (s *APIKeyService) snapshotToAPIKey(key string, snapshot *APIKeyAuthSnapsho
AudioRealtimePricePerMin: snapshot.Group.AudioRealtimePricePerMin,
AudioTTSPricePerMillionChars: snapshot.Group.AudioTTSPricePerMillionChars,
AudioSTTPricePerHour: snapshot.Group.AudioSTTPricePerHour,
LongContextPricingEnabled: snapshot.Group.LongContextPricingEnabled,
ModelPricing: snapshot.Group.ModelPricing,
ClaudeCodeOnly: snapshot.Group.ClaudeCodeOnly,
FallbackGroupID: snapshot.Group.FallbackGroupID,
FallbackGroupIDOnInvalidRequest: snapshot.Group.FallbackGroupIDOnInvalidRequest,
@@ -0,0 +1,50 @@
package service
import (
"context"
"encoding/json"
"testing"
"github.com/stretchr/testify/require"
)
func TestAPIKeyAuthSnapshotGroupPricingRoundtrip(t *testing.T) {
groupID := int64(50)
inputPrice := 1e-6
outputPrice := 2e-6
apiKey := &APIKey{
ID: 82, UserID: 40, GroupID: &groupID, Key: "sk-pricing-roundtrip", Status: StatusActive,
User: &User{ID: 40, Status: StatusActive},
Group: &Group{
ID: groupID, Name: "pricing-roundtrip", Platform: PlatformAnthropic, Status: StatusActive,
LongContextPricingEnabled: true,
ModelPricing: []ChannelModelPricing{{
Models: []string{"claude-sonnet-*"}, BillingMode: BillingModeToken,
InputPrice: &inputPrice, OutputPrice: &outputPrice,
}},
},
}
svc := &APIKeyService{}
payload, err := json.Marshal(&APIKeyAuthCacheEntry{Snapshot: svc.snapshotFromAPIKey(context.Background(), apiKey)})
require.NoError(t, err)
var cached APIKeyAuthCacheEntry
require.NoError(t, json.Unmarshal(payload, &cached))
materialized, used, err := svc.applyAuthCacheEntry(apiKey.Key, &cached)
require.NoError(t, err)
require.True(t, used)
require.NotNil(t, materialized.Group)
require.True(t, materialized.Group.LongContextPricingEnabled)
require.Equal(t, apiKey.Group.ModelPricing, materialized.Group.ModelPricing)
billing := &BillingService{fallbackPrices: map[string]*ModelPricing{
"claude-sonnet-4": {InputPricePerToken: 3e-6, OutputPricePerToken: 15e-6},
}}
resolver := NewModelPricingResolver(nil, billing)
resolved := resolver.Resolve(context.Background(), PricingInput{Model: "claude-sonnet-4", Group: materialized.Group})
require.Equal(t, PricingSourceGroup, resolved.Source)
require.True(t, resolved.longContextPricingEnabled)
require.InDelta(t, inputPrice, resolved.BasePricing.InputPricePerToken, 1e-12)
require.InDelta(t, outputPrice, resolved.BasePricing.OutputPricePerToken, 1e-12)
}
@@ -53,7 +53,7 @@ func TestAPIKeyAuthSnapshotProfitControlRoundtrip(t *testing.T) {
snapshot := svc.snapshotFromAPIKey(context.Background(), apiKey)
require.NotNil(t, snapshot)
require.Equal(t, apiKeyAuthSnapshotVersion, snapshot.Version)
require.Equal(t, 19, snapshot.Version, "v19 起认证快照携带 search/audio/video_model_prices 计费字段")
require.Equal(t, 20, snapshot.Version, "v20 起认证快照携带分组长上下文与模型定价字段")
// 模拟 L2 缓存的完整 JSON 往返(与 apiKeyCache.SetAuthCache/GetAuthCache 同构)。
payload, err := json.Marshal(&APIKeyAuthCacheEntry{Snapshot: snapshot})