mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-11-03 00:33:13 +00:00
fix(openai-ws): bill each later turn with the current group pricing
The Responses WebSocket handler reads the API key auth snapshot once at connect and reuses it for every turn: RecordUsage takes the group pricing from that snapshot, and the profit gate reinstalled in BeforeTurn takes its selling price from the connect-time ctxkey.Group. A group price change therefore never reaches an open connection until the client reconnects, while HTTP requests pick it up once the auth cache entry is invalidated. Each later turn now re-reads the key through APIKeyService.GetByKey (the auth cache the middleware uses) in BeforeTurn. When the key ID, group ID, group platform and subscription type are unchanged, the fresh group is swapped into ctxkey.Group for the turn's profit gate and passed to that turn's RecordUsage. A failed lookup or a key moved to another group keeps the connection snapshot. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
SPHINX998andClaude Opus 5.5 committed
1 parent
42bc7f6cff
commit
73381fc9c7
3 files changed
+412
-2
No files matched your search
@@ -2125,6 +2125,101 @@ func (p *openAIWSTurnPricing) currentOr(fallback time.Time) time.Time {
|
||||
return fallback
|
||||
}
|
||||
|
||||
// openAIWSTurnAPIKeyLookup 是后续 turn 重取 API Key 认证快照的入口
|
||||
// （APIKeyService.GetByKey：经 L1/L2 认证缓存，未命中才回源）。
|
||||
type openAIWSTurnAPIKeyLookup interface {
|
||||
GetByKey(ctx context.Context, key string) (*service.APIKey, error)
|
||||
}
|
||||
|
||||
// openAIWSTurnBillingAPIKeys 按 turn 号保存每个 turn 计费用的 API Key 快照。
|
||||
//
|
||||
// API Key 认证快照只在建连时取一次。若连接内所有 turn 共用它，分组定价（倍率、
|
||||
// 高峰、图片定价、利润门参数）会停留在建连时刻：管理员调价后，已打开的长连接
|
||||
// 在客户端重连前一直按旧价计费、按旧售价过利润门，而 HTTP 请求每次都经认证缓存
|
||||
// 取快照，缓存失效后即生效。
|
||||
//
|
||||
// 因此后续 turn 在 BeforeTurn（与 HTTP 请求进入认证中间件同位）经同一认证缓存
|
||||
// 重取快照，只采用其中的分组，作为该 turn 利润门准入与用量计费共同的分组；用户、
|
||||
// Key 限额、订阅仍沿用建连快照。只在同一把 Key、同一分组且平台与订阅类型未变时
|
||||
// 采用：Key 换组或分组改平台/订阅类型不是调价，该连接按建连分组调度，继续按建连
|
||||
// 快照计费；重取失败同样保留建连快照，不断连。
|
||||
//
|
||||
// 首轮沿用刚经认证中间件取得的建连快照；没有经过 BeforeTurn 的 turn 回退建连
|
||||
// 快照。只保留当前与上一个 turn：下一 turn 的 BeforeTurn 先于上一 turn 的
|
||||
// AfterTurn 执行时，上一 turn 仍取到自己的快照。
|
||||
type openAIWSTurnBillingAPIKeys struct {
|
||||
mu sync.Mutex
|
||||
keys map[int]*service.APIKey
|
||||
}
|
||||
|
||||
// begin 在 BeforeTurn 重装利润门前调用：后续 turn 重取计费分组并按 turn 记下，
|
||||
// 返回换入该分组的上下文供本 turn 的利润门使用。
|
||||
func (k *openAIWSTurnBillingAPIKeys) begin(ctx context.Context, apiKeyService *service.APIKeyService, turn int, conn *service.APIKey) context.Context {
|
||||
turnKey := conn
|
||||
if turn > 1 && apiKeyService != nil {
|
||||
turnKey = refreshOpenAIWSTurnBillingAPIKey(ctx, apiKeyService, conn)
|
||||
}
|
||||
k.set(turn, turnKey)
|
||||
return withOpenAIWSTurnBillingGroup(ctx, turnKey)
|
||||
}
|
||||
|
||||
func (k *openAIWSTurnBillingAPIKeys) set(turn int, key *service.APIKey) {
|
||||
k.mu.Lock()
|
||||
defer k.mu.Unlock()
|
||||
if k.keys == nil {
|
||||
k.keys = make(map[int]*service.APIKey, 2)
|
||||
}
|
||||
for t := range k.keys {
|
||||
if t < turn-1 {
|
||||
delete(k.keys, t)
|
||||
}
|
||||
}
|
||||
k.keys[turn] = key
|
||||
}
|
||||
|
||||
func (k *openAIWSTurnBillingAPIKeys) forTurn(turn int, conn *service.APIKey) *service.APIKey {
|
||||
k.mu.Lock()
|
||||
defer k.mu.Unlock()
|
||||
if key := k.keys[turn]; key != nil {
|
||||
return key
|
||||
}
|
||||
return conn
|
||||
}
|
||||
|
||||
// refreshOpenAIWSTurnBillingAPIKey 返回本 turn 计费用的 API Key：分组取当前认证
|
||||
// 快照，其余字段与建连快照共享。不满足采用条件时原样返回建连快照。
|
||||
func refreshOpenAIWSTurnBillingAPIKey(ctx context.Context, lookup openAIWSTurnAPIKeyLookup, conn *service.APIKey) *service.APIKey {
|
||||
if lookup == nil || conn == nil || conn.Key == "" || conn.GroupID == nil || conn.Group == nil {
|
||||
return conn
|
||||
}
|
||||
latest, err := lookup.GetByKey(ctx, conn.Key)
|
||||
if err != nil || latest == nil || latest.ID != conn.ID || latest.GroupID == nil || *latest.GroupID != *conn.GroupID {
|
||||
return conn
|
||||
}
|
||||
group := latest.Group
|
||||
if group == nil || group.ID != conn.Group.ID ||
|
||||
group.Platform != conn.Group.Platform ||
|
||||
group.SubscriptionType != conn.Group.SubscriptionType {
|
||||
return conn
|
||||
}
|
||||
turnKey := *conn
|
||||
turnKey.Group = group
|
||||
return &turnKey
|
||||
}
|
||||
|
||||
// withOpenAIWSTurnBillingGroup 把本 turn 的计费分组换进认证分组上下文，使利润门
|
||||
// 的售价与本 turn 计费同源。上下文里没有同 ID 的认证分组时不改动。
|
||||
func withOpenAIWSTurnBillingGroup(ctx context.Context, turnKey *service.APIKey) context.Context {
|
||||
if turnKey == nil || turnKey.Group == nil {
|
||||
return ctx
|
||||
}
|
||||
current, ok := ctx.Value(ctxkey.Group).(*service.Group)
|
||||
if !ok || current == nil || current == turnKey.Group || current.ID != turnKey.Group.ID {
|
||||
return ctx
|
||||
}
|
||||
return context.WithValue(ctx, ctxkey.Group, turnKey.Group)
|
||||
}
|
||||
|
||||
// recordOpenAIProfitVeto 记录 OpenAI 侧选号循环的一次利润门终检否决：把账号
|
||||
// 加入本请求排除集并递增否决计数。返回 false 表示否决次数已达
|
||||
// maxProfitVetoAttempts，调用方必须停止重选并按「无可用账号」终止。
|
||||
@@ -2856,6 +2951,9 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
// turn 级定价：首轮回退到 TurnStarted 的所属 turn 时刻；后续 turn 由
|
||||
// BeforeTurn 重新冻结 pricingAt 并按最新门复核当前账号。
|
||||
var turnPricing openAIWSTurnPricing
|
||||
// 后续 turn 的计费分组同样在 BeforeTurn 经认证缓存重取，使分组调价对
|
||||
// 已打开的连接生效。
|
||||
var turnBillingAPIKeys openAIWSTurnBillingAPIKeys
|
||||
// Passthrough ingress does not invoke BeforeTurn for the first frame.
|
||||
if err := checkSimpleModeTurnBilling(); err != nil {
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, "billing check failed")
|
||||
@@ -2943,7 +3041,8 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
// 长连接跨峰谷/倍率刷新防护：每个 turn 按当前时刻重装门并复核
|
||||
// 当前账号，越线即要求客户端重连重选（连接绑定单一上游账号，
|
||||
// 无法中途换号）。本 turn 的准入与计费共用同一 pricingAt。
|
||||
turnCtx, turnAt := h.gatewayService.WithOpenAITurnPricingContext(ctx, apiKey.GroupID)
|
||||
turnBillingCtx := turnBillingAPIKeys.begin(ctx, h.apiKeyService, turn, apiKey)
|
||||
turnCtx, turnAt := h.gatewayService.WithOpenAITurnPricingContext(turnBillingCtx, apiKey.GroupID)
|
||||
if _, vetoed, reason := h.gatewayService.ProfitControlVetoLatest(turnCtx, account); vetoed {
|
||||
reqLog.Info("openai.websocket_turn_profit_vetoed",
|
||||
zap.Int("turn", turn),
|
||||
@@ -3059,11 +3158,12 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
sessionID := service.ExtractClientSessionID(c)
|
||||
turnRecordPricingAt := turnPricing.currentOr(turnStart)
|
||||
turnBillingAPIKey := turnBillingAPIKeys.forTurn(turn, apiKey)
|
||||
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
|
||||
h.submitOpenAIUsageRecordTask(ctx, result, func(taskCtx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(taskCtx, &service.OpenAIRecordUsageInput{
|
||||
Result: result,
|
||||
APIKey: apiKey,
|
||||
APIKey: turnBillingAPIKey,
|
||||
User: apiKey.User,
|
||||
Account: account,
|
||||
Subscription: subscription,
|
||||
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
@@ -1948,6 +1949,12 @@ type openAIResponsesWSUsageLogCase struct {
|
||||
afterFirstUpstreamRequest func(channelSvc *service.ChannelService) error
|
||||
// group 覆盖 apiKey.Group（分组级模型白名单测试用）；nil 保持原有无分组行为。
|
||||
group *service.Group
|
||||
// apiKeyService 非 nil 时模拟 API Key 认证中间件：连接认证快照经它按
|
||||
// apiKeyCredential 取得，并把其分组放入请求 ctx；handler 也使用它。
|
||||
apiKeyService *service.APIKeyService
|
||||
apiKeyCredential string
|
||||
// accountRateMultiplier 覆盖账号倍率（利润门测试用）。
|
||||
accountRateMultiplier *float64
|
||||
// firstFrameCloseExpected：首帧即被拒（连接被 1008 关闭），不期待任何响应帧。
|
||||
firstFrameCloseExpected bool
|
||||
// secondTurnCloseExpected：第二个 turn 被拒（连接被 1008 关闭）。
|
||||
@@ -2962,6 +2969,7 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
if tc.accountPlatform != "" {
|
||||
account.Platform = tc.accountPlatform
|
||||
}
|
||||
account.RateMultiplier = tc.accountRateMultiplier
|
||||
if strings.TrimSpace(tc.ingressMode) != "" {
|
||||
account.Extra["openai_apikey_responses_websockets_v2_mode"] = tc.ingressMode
|
||||
}
|
||||
@@ -3050,6 +3058,12 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
GroupID: &groupID,
|
||||
User: &service.User{ID: 1701, Status: service.StatusActive},
|
||||
}
|
||||
if tc.apiKeyService != nil {
|
||||
h.apiKeyService = tc.apiKeyService
|
||||
authKey, err := tc.apiKeyService.GetByKey(context.Background(), tc.apiKeyCredential)
|
||||
require.NoError(t, err)
|
||||
apiKey = authKey
|
||||
}
|
||||
if tc.simpleModeRejectAtRead > 0 {
|
||||
apiKey.RateLimit5h = 1
|
||||
}
|
||||
@@ -3060,6 +3074,9 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
router.Use(func(c *gin.Context) {
|
||||
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.User.ID, Concurrency: 1})
|
||||
if tc.apiKeyService != nil && apiKey.Group != nil {
|
||||
c.Request = c.Request.WithContext(context.WithValue(c.Request.Context(), ctxkey.Group, apiKey.Group))
|
||||
}
|
||||
c.Next()
|
||||
})
|
||||
router.GET("/openai/v1/responses", h.ResponsesWebSocket)
|
||||
|
||||
@@ -0,0 +1,293 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
coderws "github.com/coder/websocket"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// 分组调价对已打开的 WebSocket 连接生效：同一条连接的两个 turn 之间由管理端
|
||||
// 改分组倍率（并按分组失效认证缓存），第二个 turn 必须按新倍率计费、按新售价
|
||||
// 过利润门；HTTP 请求每次经认证缓存取快照，本来就是这个行为。
|
||||
|
||||
const (
|
||||
wsGroupPricingGroupID = int64(4201) // 与 runOpenAIResponsesWebSocketUsageLogCase 的调度分组一致
|
||||
wsGroupPricingKey = "sk-ws-group-pricing"
|
||||
wsGroupPricingFirst = `{"type":"response.create","model":"gpt-5.6-sol","stream":false}`
|
||||
wsGroupPricingSecond = `{"type":"response.create","model":"gpt-5.6-sol","stream":false}`
|
||||
)
|
||||
|
||||
// wsGroupPricingAPIKeyRepoStub 只实现认证查询与按分组列 Key（认证缓存失效用）。
|
||||
type wsGroupPricingAPIKeyRepoStub struct {
|
||||
service.APIKeyRepository
|
||||
mu sync.Mutex
|
||||
apiKey service.APIKey
|
||||
group service.Group
|
||||
err error
|
||||
lookups int
|
||||
}
|
||||
|
||||
func newWSGroupPricingAPIKeyRepoStub(rate float64) *wsGroupPricingAPIKeyRepoStub {
|
||||
groupID := wsGroupPricingGroupID
|
||||
return &wsGroupPricingAPIKeyRepoStub{
|
||||
apiKey: service.APIKey{
|
||||
ID: 1801,
|
||||
UserID: 1701,
|
||||
Key: wsGroupPricingKey,
|
||||
Status: service.StatusActive,
|
||||
GroupID: &groupID,
|
||||
User: &service.User{ID: 1701, Status: service.StatusActive, Concurrency: 1},
|
||||
},
|
||||
group: service.Group{
|
||||
ID: groupID,
|
||||
Name: "ws-group-pricing",
|
||||
Platform: service.PlatformOpenAI,
|
||||
Status: service.StatusActive,
|
||||
RateMultiplier: rate,
|
||||
Hydrated: true,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (s *wsGroupPricingAPIKeyRepoStub) GetByKeyForAuth(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.lookups++
|
||||
if s.err != nil {
|
||||
return nil, s.err
|
||||
}
|
||||
if key != s.apiKey.Key {
|
||||
return nil, service.ErrAPIKeyNotFound
|
||||
}
|
||||
apiKey := s.apiKey
|
||||
user := *s.apiKey.User
|
||||
apiKey.User = &user
|
||||
group := s.group
|
||||
apiKey.Group = &group
|
||||
return &apiKey, nil
|
||||
}
|
||||
|
||||
func (s *wsGroupPricingAPIKeyRepoStub) ListKeysByGroupID(ctx context.Context, groupID int64) ([]string, error) {
|
||||
return []string{s.apiKey.Key}, nil
|
||||
}
|
||||
|
||||
// adminUpdate 模拟管理端改 Key/分组后按分组失效认证缓存（与 admin 分组更新同一入口）。
|
||||
func (s *wsGroupPricingAPIKeyRepoStub) adminUpdate(apiKeyService *service.APIKeyService, fn func(apiKey *service.APIKey, group *service.Group)) error {
|
||||
s.mu.Lock()
|
||||
fn(&s.apiKey, &s.group)
|
||||
s.mu.Unlock()
|
||||
apiKeyService.InvalidateAuthCacheByGroupID(context.Background(), wsGroupPricingGroupID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *wsGroupPricingAPIKeyRepoStub) lookupCount() int {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.lookups
|
||||
}
|
||||
|
||||
func newWSGroupPricingAPIKeyService(repo *wsGroupPricingAPIKeyRepoStub) *service.APIKeyService {
|
||||
return service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, &config.Config{})
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_GroupRateChangeAppliesToNextTurn(t *testing.T) {
|
||||
for _, mode := range []string{service.OpenAIWSIngressModePassthrough, service.OpenAIWSIngressModeCtxPool} {
|
||||
t.Run(mode, func(t *testing.T) {
|
||||
repo := newWSGroupPricingAPIKeyRepoStub(3.0)
|
||||
apiKeyService := newWSGroupPricingAPIKeyService(repo)
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: wsGroupPricingFirst,
|
||||
secondPayload: wsGroupPricingSecond,
|
||||
ingressMode: mode,
|
||||
apiKeyService: apiKeyService,
|
||||
apiKeyCredential: wsGroupPricingKey,
|
||||
afterFirstUpstreamRequest: func(*service.ChannelService) error {
|
||||
return repo.adminUpdate(apiKeyService, func(_ *service.APIKey, group *service.Group) { group.RateMultiplier = 0.3 })
|
||||
},
|
||||
})
|
||||
|
||||
require.Len(t, got.logs, 2)
|
||||
require.InDelta(t, 3.0, got.logs[0].RateMultiplier, 1e-12, "turn 1 bills at the rate in force at connect")
|
||||
require.InDelta(t, 0.3, got.logs[1].RateMultiplier, 1e-12, "turn 2 must bill at the group rate in force when it starts")
|
||||
require.Greater(t, got.logs[0].ActualCost, 0.0)
|
||||
require.InDelta(t, got.logs[0].ActualCost/10, got.logs[1].ActualCost, 1e-15)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_GroupRateChangeReachesProfitGateOnNextTurn(t *testing.T) {
|
||||
repo := newWSGroupPricingAPIKeyRepoStub(3.0)
|
||||
repo.group.ProfitControlEnabled = true
|
||||
repo.group.ProfitMinMargin = 0.1
|
||||
apiKeyService := newWSGroupPricingAPIKeyService(repo)
|
||||
accountRate := 1.0
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: wsGroupPricingFirst,
|
||||
secondPayload: wsGroupPricingSecond,
|
||||
apiKeyService: apiKeyService,
|
||||
apiKeyCredential: wsGroupPricingKey,
|
||||
accountRateMultiplier: &accountRate,
|
||||
// 3.0 → 0.3：账号倍率 1.0 在新售价下已越线，第二个 turn 必须在转发上游前被拒。
|
||||
afterFirstUpstreamRequest: func(*service.ChannelService) error {
|
||||
return repo.adminUpdate(apiKeyService, func(_ *service.APIKey, group *service.Group) { group.RateMultiplier = 0.3 })
|
||||
},
|
||||
secondTurnCloseExpected: true,
|
||||
closeStatus: coderws.StatusTryAgainLater,
|
||||
closeReason: "no longer eligible for this connection",
|
||||
})
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_KeyMovedToAnotherGroupKeepsConnectionGroup(t *testing.T) {
|
||||
repo := newWSGroupPricingAPIKeyRepoStub(3.0)
|
||||
apiKeyService := newWSGroupPricingAPIKeyService(repo)
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: wsGroupPricingFirst,
|
||||
secondPayload: wsGroupPricingSecond,
|
||||
apiKeyService: apiKeyService,
|
||||
apiKeyCredential: wsGroupPricingKey,
|
||||
afterFirstUpstreamRequest: func(*service.ChannelService) error {
|
||||
return repo.adminUpdate(apiKeyService, func(apiKey *service.APIKey, group *service.Group) {
|
||||
other := wsGroupPricingGroupID + 1
|
||||
apiKey.GroupID = &other
|
||||
group.ID = other
|
||||
group.RateMultiplier = 0.3
|
||||
})
|
||||
},
|
||||
})
|
||||
|
||||
require.Len(t, got.logs, 2)
|
||||
require.InDelta(t, 3.0, got.logs[1].RateMultiplier, 1e-12,
|
||||
"a key moved to another group keeps the connection's group: the connection was scheduled from it")
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_FailedKeyRefreshKeepsConnectionGroup(t *testing.T) {
|
||||
repo := newWSGroupPricingAPIKeyRepoStub(3.0)
|
||||
apiKeyService := newWSGroupPricingAPIKeyService(repo)
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: wsGroupPricingFirst,
|
||||
secondPayload: wsGroupPricingSecond,
|
||||
apiKeyService: apiKeyService,
|
||||
apiKeyCredential: wsGroupPricingKey,
|
||||
afterFirstUpstreamRequest: func(*service.ChannelService) error {
|
||||
return repo.adminUpdate(apiKeyService, func(_ *service.APIKey, group *service.Group) {
|
||||
group.RateMultiplier = 0.3
|
||||
repo.err = errors.New("auth store unavailable")
|
||||
})
|
||||
},
|
||||
})
|
||||
|
||||
require.Len(t, got.logs, 2)
|
||||
require.InDelta(t, 3.0, got.logs[1].RateMultiplier, 1e-12, "a failed refresh keeps the connection snapshot")
|
||||
require.GreaterOrEqual(t, repo.lookupCount(), 2, "the second turn must have attempted a refresh")
|
||||
}
|
||||
|
||||
type wsTurnAPIKeyLookupFunc func(ctx context.Context, key string) (*service.APIKey, error)
|
||||
|
||||
func (f wsTurnAPIKeyLookupFunc) GetByKey(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return f(ctx, key)
|
||||
}
|
||||
|
||||
func TestRefreshOpenAIWSTurnBillingAPIKey(t *testing.T) {
|
||||
groupID := int64(77)
|
||||
conn := &service.APIKey{
|
||||
ID: 5,
|
||||
Key: "sk-conn",
|
||||
GroupID: &groupID,
|
||||
User: &service.User{ID: 9},
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, SubscriptionType: service.SubscriptionTypeStandard, RateMultiplier: 3},
|
||||
}
|
||||
latestWith := func(mutate func(k *service.APIKey)) *service.APIKey {
|
||||
gid := groupID
|
||||
k := &service.APIKey{
|
||||
ID: 5,
|
||||
Key: "sk-conn",
|
||||
GroupID: &gid,
|
||||
User: &service.User{ID: 9, Balance: 123},
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, SubscriptionType: service.SubscriptionTypeStandard, RateMultiplier: 0.3},
|
||||
}
|
||||
if mutate != nil {
|
||||
mutate(k)
|
||||
}
|
||||
return k
|
||||
}
|
||||
|
||||
t.Run("same group adopts the latest group and nothing else", func(t *testing.T) {
|
||||
got := refreshOpenAIWSTurnBillingAPIKey(context.Background(), wsTurnAPIKeyLookupFunc(func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
require.Equal(t, "sk-conn", key)
|
||||
return latestWith(nil), nil
|
||||
}), conn)
|
||||
require.NotSame(t, conn, got)
|
||||
require.InDelta(t, 0.3, got.Group.RateMultiplier, 1e-12)
|
||||
require.Same(t, conn.User, got.User, "only the group snapshot is refreshed")
|
||||
require.Equal(t, conn.ID, got.ID)
|
||||
require.InDelta(t, 3.0, conn.Group.RateMultiplier, 1e-12, "the connection snapshot is never mutated")
|
||||
})
|
||||
|
||||
keep := map[string]wsTurnAPIKeyLookupFunc{
|
||||
"lookup error": func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return nil, errors.New("boom")
|
||||
},
|
||||
"not found": func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return nil, service.ErrAPIKeyNotFound
|
||||
},
|
||||
"nil key": func(ctx context.Context, key string) (*service.APIKey, error) { return nil, nil },
|
||||
"different key id": func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return latestWith(func(k *service.APIKey) { k.ID = 6 }), nil
|
||||
},
|
||||
"moved group": func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return latestWith(func(k *service.APIKey) { other := int64(78); k.GroupID = &other; k.Group.ID = other }), nil
|
||||
},
|
||||
"no group": func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return latestWith(func(k *service.APIKey) { k.Group = nil }), nil
|
||||
},
|
||||
"platform changed": func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return latestWith(func(k *service.APIKey) { k.Group.Platform = service.PlatformAnthropic }), nil
|
||||
},
|
||||
"billing type changed": func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return latestWith(func(k *service.APIKey) { k.Group.SubscriptionType = service.SubscriptionTypeSubscription }), nil
|
||||
},
|
||||
}
|
||||
for name, lookup := range keep {
|
||||
t.Run(name+" keeps the connection snapshot", func(t *testing.T) {
|
||||
require.Same(t, conn, refreshOpenAIWSTurnBillingAPIKey(context.Background(), lookup, conn))
|
||||
})
|
||||
}
|
||||
|
||||
t.Run("no credential or no group skips the lookup", func(t *testing.T) {
|
||||
called := false
|
||||
lookup := wsTurnAPIKeyLookupFunc(func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
called = true
|
||||
return latestWith(nil), nil
|
||||
})
|
||||
noKey := *conn
|
||||
noKey.Key = ""
|
||||
require.Same(t, &noKey, refreshOpenAIWSTurnBillingAPIKey(context.Background(), lookup, &noKey))
|
||||
noGroup := *conn
|
||||
noGroup.Group = nil
|
||||
require.Same(t, &noGroup, refreshOpenAIWSTurnBillingAPIKey(context.Background(), lookup, &noGroup))
|
||||
require.Nil(t, refreshOpenAIWSTurnBillingAPIKey(context.Background(), lookup, nil))
|
||||
require.Same(t, conn, refreshOpenAIWSTurnBillingAPIKey(context.Background(), nil, conn))
|
||||
require.False(t, called)
|
||||
})
|
||||
}
|
||||
|
||||
func TestOpenAIWSTurnBillingAPIKeysKeepPreviousTurnUntilItIsRecorded(t *testing.T) {
|
||||
conn := &service.APIKey{ID: 1}
|
||||
turn2 := &service.APIKey{ID: 2}
|
||||
turn3 := &service.APIKey{ID: 3}
|
||||
var keys openAIWSTurnBillingAPIKeys
|
||||
require.Same(t, conn, keys.forTurn(1, conn), "a turn without BeforeTurn bills with the connection snapshot")
|
||||
keys.set(2, turn2)
|
||||
keys.set(3, turn3) // turn 3 starts before turn 2's usage is submitted
|
||||
require.Same(t, turn2, keys.forTurn(2, conn))
|
||||
require.Same(t, turn3, keys.forTurn(3, conn))
|
||||
keys.set(4, conn)
|
||||
require.Same(t, conn, keys.forTurn(2, conn), "only the current and previous turn are kept")
|
||||
require.Len(t, keys.keys, 2)
|
||||
}
|
||||
Reference in new issue
Block a user