mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 12:02:27 +08:00
HuntercodeT
e93e6368fc
filterSchedulerExtra is an allowlist and never carried openai_passthrough / openai_oauth_passthrough, while credentials.model_mapping is preserved in full. Candidate filtering (ListSchedulableAccounts) reads that trimmed projection (sched:meta:<id>), so there a passthrough account looks exactly like a whitelist-mode one: Account.IsModelSupported falls back to the model_mapping whitelist and the account is dropped as model_not_supported. Forwarding, which hydrates the full account from sched:acc:<id>, still honors passthrough - hence "testing the account directly succeeds, but the gateway reports no available accounts supporting model X". The practical impact is that a stale leftover whitelist silently shrinks the candidate pool of accounts that were explicitly configured to let the upstream decide. When no account in the group lists the requested model, the group is fully dark for it even though every account can serve it. This is not a regression of #4936: that fix corrected the predicate (Account.IsModelSupported short-circuits passthrough before the mapping check), this one restores the input the predicate is fed. The projection allowlist predates #4936 by three and a half months and has never carried the flag, so the fix simply never took effect on the snapshot path. Add both keys to the allowlist, plus a regression test that round-trips the projection through JSON (the real sched:meta write/read path) and asserts the model gate still lets a non-whitelisted model through for a passthrough account. Claude-Session: https://claude.ai/code/session_01LwvyJVV2sdqiu41Z9JDNUq
1124 lines
45 KiB
Go
1124 lines
45 KiB
Go
//go:build unit
|
|
|
|
package repository
|
|
|
|
import (
|
|
"context"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"fmt"
|
|
"strconv"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/service"
|
|
"github.com/alicebob/miniredis/v2"
|
|
"github.com/redis/go-redis/v9"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func newSchedulerCacheUnit(t *testing.T) *schedulerCache {
|
|
cache, _ := newSchedulerCacheUnitWithRedis(t)
|
|
return cache
|
|
}
|
|
|
|
func newSchedulerCacheUnitWithRedis(t *testing.T) (*schedulerCache, *miniredis.Miniredis) {
|
|
t.Helper()
|
|
mr := miniredis.RunT(t)
|
|
rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
|
|
t.Cleanup(func() { _ = rdb.Close() })
|
|
cache, ok := newSchedulerCacheWithChunkSizes(rdb, defaultSchedulerSnapshotMGetChunkSize, defaultSchedulerSnapshotWriteChunkSize).(*schedulerCache)
|
|
require.True(t, ok)
|
|
return cache, mr
|
|
}
|
|
|
|
func TestSchedulerCacheWriteAccountIDsSkipsUnencodableTimes(t *testing.T) {
|
|
ctx := context.Background()
|
|
cache := newSchedulerCacheUnit(t)
|
|
invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
|
|
|
|
accountIDs, err := cache.writeAccountIDs(ctx, []service.Account{
|
|
{ID: 111, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey},
|
|
{ID: 112, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, ExpiresAt: &invalidTime},
|
|
})
|
|
require.NoError(t, err)
|
|
require.Equal(t, []int64{111}, accountIDs)
|
|
|
|
cached, err := cache.GetAccount(ctx, 111)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, cached)
|
|
|
|
invalid, err := cache.GetAccount(ctx, 112)
|
|
require.NoError(t, err)
|
|
require.Nil(t, invalid)
|
|
}
|
|
|
|
func TestSchedulerCacheSetAccountClearsUnencodablePayload(t *testing.T) {
|
|
ctx := context.Background()
|
|
cache := newSchedulerCacheUnit(t)
|
|
|
|
account := service.Account{ID: 113, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
|
|
require.NoError(t, cache.SetAccount(ctx, &account))
|
|
|
|
invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
|
|
account.ExpiresAt = &invalidTime
|
|
require.NoError(t, cache.SetAccount(ctx, &account))
|
|
|
|
cached, err := cache.GetAccount(ctx, account.ID)
|
|
require.NoError(t, err)
|
|
require.Nil(t, cached)
|
|
}
|
|
|
|
func TestSchedulerCacheUpdateLastUsedClearsUnencodablePayload(t *testing.T) {
|
|
ctx := context.Background()
|
|
cache := newSchedulerCacheUnit(t)
|
|
account := service.Account{ID: 114, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
|
|
require.NoError(t, cache.SetAccount(ctx, &account))
|
|
|
|
invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
|
|
require.NoError(t, cache.UpdateLastUsed(ctx, map[int64]time.Time{account.ID: invalidTime}))
|
|
|
|
cached, err := cache.GetAccount(ctx, account.ID)
|
|
require.NoError(t, err)
|
|
require.Nil(t, cached)
|
|
}
|
|
|
|
func TestSchedulerCacheSnapshotAccountIDReusePreservesPayloadAndMembers(t *testing.T) {
|
|
ctx := context.Background()
|
|
cache, _ := newSchedulerCacheUnitWithRedis(t)
|
|
invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
|
|
validOne := service.Account{
|
|
ID: 701,
|
|
Name: "first",
|
|
Platform: service.PlatformOpenAI,
|
|
Type: service.AccountTypeOAuth,
|
|
Credentials: map[string]any{"model_mapping": map[string]any{"z": "last", "a": "first"}},
|
|
Extra: map[string]any{"mixed_scheduling": true},
|
|
GroupIDs: []int64{17},
|
|
}
|
|
validTwo := service.Account{ID: 702, Name: "second", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
|
|
invalid := service.Account{ID: 799, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, ExpiresAt: &invalidTime}
|
|
accounts := []service.Account{validOne, invalid, validTwo, validOne}
|
|
|
|
single := service.SchedulerBucket{GroupID: 17, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
|
|
singleToken, err := cache.CaptureBucketWriteToken(ctx, single)
|
|
require.NoError(t, err)
|
|
accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, single, singleToken, accounts)
|
|
require.NoError(t, err)
|
|
require.Equal(t, []int64{701, 702, 701}, accountIDs, "应保留可编码账号的原顺序和重复项")
|
|
|
|
wantFull, err := json.Marshal(validOne)
|
|
require.NoError(t, err)
|
|
wantMeta, err := json.Marshal(buildSchedulerMetadataAccount(validOne))
|
|
require.NoError(t, err)
|
|
fullBefore, err := cache.rdb.Get(ctx, schedulerAccountKey("701")).Bytes()
|
|
require.NoError(t, err)
|
|
metaBefore, err := cache.rdb.Get(ctx, schedulerAccountMetaKey("701")).Bytes()
|
|
require.NoError(t, err)
|
|
require.Equal(t, wantFull, fullBefore)
|
|
require.Equal(t, wantMeta, metaBefore)
|
|
|
|
forced := service.SchedulerBucket{GroupID: 17, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
|
|
forcedToken, err := cache.CaptureBucketWriteToken(ctx, forced)
|
|
require.NoError(t, err)
|
|
require.NoError(t, cache.SetSnapshotByAccountIDs(ctx, forced, forcedToken, accountIDs))
|
|
|
|
fullAfter, err := cache.rdb.Get(ctx, schedulerAccountKey("701")).Bytes()
|
|
require.NoError(t, err)
|
|
metaAfter, err := cache.rdb.Get(ctx, schedulerAccountMetaKey("701")).Bytes()
|
|
require.NoError(t, err)
|
|
require.Equal(t, fullBefore, fullAfter, "ID-only 路径不得重写完整账号键")
|
|
require.Equal(t, metaBefore, metaAfter, "ID-only 路径不得重写调度元数据键")
|
|
|
|
for _, bucket := range []service.SchedulerBucket{single, forced} {
|
|
version, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Result()
|
|
require.NoError(t, err)
|
|
members, err := cache.rdb.ZRange(ctx, schedulerSnapshotKey(bucket, version), 0, -1).Result()
|
|
require.NoError(t, err)
|
|
require.Equal(t, []string{"702", "701"}, members, bucket.String())
|
|
}
|
|
missing, err := cache.GetAccount(ctx, invalid.ID)
|
|
require.NoError(t, err)
|
|
require.Nil(t, missing)
|
|
}
|
|
|
|
func TestSchedulerCacheSetSnapshotMatchesIDPublishing(t *testing.T) {
|
|
ctx := context.Background()
|
|
cache, _ := newSchedulerCacheUnitWithRedis(t)
|
|
invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
|
|
validOne := service.Account{
|
|
ID: 721,
|
|
Name: "first",
|
|
Platform: service.PlatformOpenAI,
|
|
Type: service.AccountTypeOAuth,
|
|
Credentials: map[string]any{"model_mapping": map[string]any{"source": "target"}},
|
|
Extra: map[string]any{"mixed_scheduling": true},
|
|
GroupIDs: []int64{21},
|
|
}
|
|
validTwo := service.Account{ID: 722, Name: "second", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
|
|
invalid := service.Account{ID: 799, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, ExpiresAt: &invalidTime}
|
|
accounts := []service.Account{validOne, invalid, validTwo, validOne}
|
|
|
|
normal := service.SchedulerBucket{GroupID: 21, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
|
|
normalToken, err := cache.CaptureBucketWriteToken(ctx, normal)
|
|
require.NoError(t, err)
|
|
require.NoError(t, cache.SetSnapshot(ctx, normal, normalToken, accounts))
|
|
|
|
fullBefore, err := cache.rdb.Get(ctx, schedulerAccountKey("721")).Bytes()
|
|
require.NoError(t, err)
|
|
metaBefore, err := cache.rdb.Get(ctx, schedulerAccountMetaKey("721")).Bytes()
|
|
require.NoError(t, err)
|
|
|
|
idOnly := service.SchedulerBucket{GroupID: 21, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
|
|
idOnlyToken, err := cache.CaptureBucketWriteToken(ctx, idOnly)
|
|
require.NoError(t, err)
|
|
accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, idOnly, idOnlyToken, accounts)
|
|
require.NoError(t, err)
|
|
require.Equal(t, []int64{721, 722, 721}, accountIDs)
|
|
|
|
fullAfter, err := cache.rdb.Get(ctx, schedulerAccountKey("721")).Bytes()
|
|
require.NoError(t, err)
|
|
metaAfter, err := cache.rdb.Get(ctx, schedulerAccountMetaKey("721")).Bytes()
|
|
require.NoError(t, err)
|
|
require.Equal(t, fullBefore, fullAfter, "普通快照和 ID 发布必须写入相同完整账号 payload")
|
|
require.Equal(t, metaBefore, metaAfter, "普通快照和 ID 发布必须写入相同元数据 payload")
|
|
|
|
for _, bucket := range []service.SchedulerBucket{normal, idOnly} {
|
|
version, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Result()
|
|
require.NoError(t, err)
|
|
members, err := cache.rdb.ZRange(ctx, schedulerSnapshotKey(bucket, version), 0, -1).Result()
|
|
require.NoError(t, err)
|
|
require.Equal(t, []string{"722", "721"}, members, bucket.String())
|
|
}
|
|
}
|
|
|
|
func TestSchedulerCacheSnapshotAccountIDReuseKeepsEmptySnapshotSemantics(t *testing.T) {
|
|
ctx := context.Background()
|
|
cache := newSchedulerCacheUnit(t)
|
|
invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
|
|
accounts := []service.Account{{ID: 811, Platform: service.PlatformOpenAI, ExpiresAt: &invalidTime}}
|
|
|
|
single := service.SchedulerBucket{GroupID: 18, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
|
|
singleToken, err := cache.CaptureBucketWriteToken(ctx, single)
|
|
require.NoError(t, err)
|
|
accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, single, singleToken, accounts)
|
|
require.NoError(t, err)
|
|
require.Empty(t, accountIDs)
|
|
|
|
forced := service.SchedulerBucket{GroupID: 18, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
|
|
forcedToken, err := cache.CaptureBucketWriteToken(ctx, forced)
|
|
require.NoError(t, err)
|
|
require.NoError(t, cache.SetSnapshotByAccountIDs(ctx, forced, forcedToken, accountIDs))
|
|
|
|
for _, bucket := range []service.SchedulerBucket{single, forced} {
|
|
ready, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerReadyPrefix, bucket)).Result()
|
|
require.NoError(t, err)
|
|
require.Equal(t, "1", ready)
|
|
snapshot, hit, err := cache.GetSnapshot(ctx, bucket)
|
|
require.NoError(t, err)
|
|
require.False(t, hit, bucket.String())
|
|
require.Nil(t, snapshot)
|
|
}
|
|
}
|
|
|
|
func TestSchedulerCacheSetSnapshotByAccountIDsKeepsFencing(t *testing.T) {
|
|
ctx := context.Background()
|
|
cache := newSchedulerCacheUnit(t)
|
|
bucket := service.SchedulerBucket{GroupID: 19, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
|
|
|
|
err := cache.SetSnapshotByAccountIDs(ctx, bucket, service.SchedulerBucketWriteToken{}, []int64{901})
|
|
require.ErrorIs(t, err, service.ErrSchedulerBucketWriteFenced)
|
|
_, err = cache.rdb.Get(ctx, schedulerBucketKey(schedulerVersionPrefix, bucket)).Result()
|
|
require.ErrorIs(t, err, redis.Nil)
|
|
|
|
token, err := cache.CaptureBucketWriteToken(ctx, bucket)
|
|
require.NoError(t, err)
|
|
require.NoError(t, cache.RetireBucket(ctx, bucket))
|
|
err = cache.SetSnapshotByAccountIDs(ctx, bucket, token, []int64{901})
|
|
require.ErrorIs(t, err, service.ErrSchedulerBucketRetired)
|
|
}
|
|
|
|
func TestSchedulerCacheSetSnapshotByAccountIDsDoesNotResurrectDeletedAccount(t *testing.T) {
|
|
ctx := context.Background()
|
|
cache := newSchedulerCacheUnit(t)
|
|
account := service.Account{ID: 902, Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth}
|
|
single := service.SchedulerBucket{GroupID: 20, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
|
|
singleToken, err := cache.CaptureBucketWriteToken(ctx, single)
|
|
require.NoError(t, err)
|
|
accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, single, singleToken, []service.Account{account})
|
|
require.NoError(t, err)
|
|
require.Equal(t, []int64{account.ID}, accountIDs)
|
|
require.NoError(t, cache.DeleteAccount(ctx, account.ID))
|
|
|
|
forced := service.SchedulerBucket{GroupID: 20, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
|
|
forcedToken, err := cache.CaptureBucketWriteToken(ctx, forced)
|
|
require.NoError(t, err)
|
|
require.NoError(t, cache.SetSnapshotByAccountIDs(ctx, forced, forcedToken, accountIDs))
|
|
|
|
full, err := cache.GetAccount(ctx, account.ID)
|
|
require.NoError(t, err)
|
|
require.Nil(t, full, "ID-only 发布不得复活已删除的完整账号键")
|
|
snapshot, hit, err := cache.GetSnapshot(ctx, forced)
|
|
require.NoError(t, err)
|
|
require.False(t, hit, "元数据缺失时必须安全回源，而不是返回残缺快照")
|
|
require.Nil(t, snapshot)
|
|
}
|
|
|
|
func TestMarshalSchedulerCacheAccountKeepsEncodingJSONWireFormat(t *testing.T) {
|
|
cases := []struct {
|
|
name string
|
|
account service.Account
|
|
}{
|
|
{name: "nil collections", account: service.Account{ID: 801}},
|
|
{name: "empty collections", account: service.Account{
|
|
ID: 802,
|
|
Credentials: map[string]any{},
|
|
Extra: map[string]any{},
|
|
GroupIDs: []int64{},
|
|
Groups: []*service.Group{},
|
|
}},
|
|
{name: "nested maps and escaping", account: service.Account{
|
|
ID: 803,
|
|
Credentials: map[string]any{"model_mapping": map[string]any{"z": "<last>", "a": "&first"}},
|
|
Extra: map[string]any{"mixed_scheduling": true},
|
|
}},
|
|
}
|
|
for _, tc := range cases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
full, meta, err := marshalSchedulerCacheAccount(tc.account)
|
|
require.NoError(t, err)
|
|
wantFull, err := json.Marshal(tc.account)
|
|
require.NoError(t, err)
|
|
wantMeta, err := json.Marshal(buildSchedulerMetadataAccount(tc.account))
|
|
require.NoError(t, err)
|
|
require.Equal(t, wantFull, full)
|
|
require.Equal(t, wantMeta, meta)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestBuildSchedulerMetadataAccount_KeepsOpenAIWSFlags(t *testing.T) {
|
|
account := service.Account{
|
|
ID: 42,
|
|
Platform: service.PlatformOpenAI,
|
|
Type: service.AccountTypeOAuth,
|
|
Extra: map[string]any{
|
|
"openai_oauth_responses_websockets_v2_enabled": true,
|
|
"openai_oauth_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough,
|
|
"openai_ws_force_http": true,
|
|
"openai_responses_mode": "force_chat_completions",
|
|
"openai_responses_supported": false,
|
|
"codex_fingerprint_mode": "session",
|
|
"codex_fingerprint_seed": "11111111-1111-4111-8111-111111111111",
|
|
"mixed_scheduling": true,
|
|
"unused_large_field": "drop-me",
|
|
},
|
|
}
|
|
|
|
got := buildSchedulerMetadataAccount(account)
|
|
|
|
require.Equal(t, true, got.Extra["openai_oauth_responses_websockets_v2_enabled"])
|
|
require.Equal(t, service.OpenAIWSIngressModePassthrough, got.Extra["openai_oauth_responses_websockets_v2_mode"])
|
|
require.Equal(t, true, got.Extra["openai_ws_force_http"])
|
|
require.Equal(t, "force_chat_completions", got.Extra["openai_responses_mode"])
|
|
require.Equal(t, false, got.Extra["openai_responses_supported"])
|
|
require.Equal(t, "session", got.Extra["codex_fingerprint_mode"])
|
|
require.Equal(t, "11111111-1111-4111-8111-111111111111", got.Extra["codex_fingerprint_seed"])
|
|
require.Equal(t, true, got.Extra["mixed_scheduling"])
|
|
require.Nil(t, got.Extra["unused_large_field"])
|
|
}
|
|
|
|
func TestBuildSchedulerMetadataAccount_KeepsGrokMediaEligibility(t *testing.T) {
|
|
t.Run("explicit override", func(t *testing.T) {
|
|
account := service.Account{
|
|
ID: 43,
|
|
Platform: service.PlatformGrok,
|
|
Type: service.AccountTypeOAuth,
|
|
Extra: map[string]any{
|
|
service.GrokMediaEligibleExtraKey: false,
|
|
"unused_large_field": "drop-me",
|
|
},
|
|
}
|
|
|
|
got := buildSchedulerMetadataAccount(account)
|
|
|
|
eligible, reason := got.GrokMediaGenerationEligibility()
|
|
require.False(t, eligible)
|
|
require.Equal(t, "override_disabled", reason)
|
|
require.Equal(t, false, got.Extra[service.GrokMediaEligibleExtraKey])
|
|
require.Nil(t, got.Extra["unused_large_field"])
|
|
})
|
|
|
|
t.Run("forbidden billing observation", func(t *testing.T) {
|
|
account := service.Account{
|
|
ID: 44,
|
|
Platform: service.PlatformGrok,
|
|
Type: service.AccountTypeOAuth,
|
|
Extra: map[string]any{
|
|
"grok_billing_snapshot": map[string]any{
|
|
"status_code": 200,
|
|
"weekly_status_code": 403,
|
|
"monthly_status_code": 200,
|
|
},
|
|
},
|
|
}
|
|
|
|
got := buildSchedulerMetadataAccount(account)
|
|
|
|
eligible, reason := got.GrokMediaGenerationEligibility()
|
|
require.False(t, eligible)
|
|
require.Equal(t, "billing_forbidden", reason)
|
|
require.NotNil(t, got.Extra["grok_billing_snapshot"])
|
|
})
|
|
}
|
|
|
|
func TestBuildSchedulerMetadataAccount_KeepsSlimGroupMembership(t *testing.T) {
|
|
account := service.Account{
|
|
ID: 42,
|
|
Platform: service.PlatformAnthropic,
|
|
GroupIDs: []int64{7, 9, 7, 0},
|
|
AccountGroups: []service.AccountGroup{
|
|
{
|
|
AccountID: 42,
|
|
GroupID: 7,
|
|
Priority: 2,
|
|
Account: &service.Account{ID: 42, Name: "drop-from-metadata"},
|
|
Group: &service.Group{ID: 7, Name: "drop-from-metadata"},
|
|
},
|
|
{
|
|
AccountID: 42,
|
|
GroupID: 11,
|
|
Priority: 3,
|
|
Group: &service.Group{ID: 11, Name: "drop-from-metadata"},
|
|
},
|
|
{
|
|
AccountID: 42,
|
|
GroupID: 0,
|
|
Priority: 4,
|
|
},
|
|
},
|
|
}
|
|
|
|
got := buildSchedulerMetadataAccount(account)
|
|
|
|
require.Equal(t, []int64{7, 9, 11}, got.GroupIDs)
|
|
require.Len(t, got.AccountGroups, 2)
|
|
require.Equal(t, int64(42), got.AccountGroups[0].AccountID)
|
|
require.Equal(t, int64(7), got.AccountGroups[0].GroupID)
|
|
require.Equal(t, 2, got.AccountGroups[0].Priority)
|
|
require.Nil(t, got.AccountGroups[0].Account)
|
|
require.Nil(t, got.AccountGroups[0].Group)
|
|
require.Equal(t, int64(11), got.AccountGroups[1].GroupID)
|
|
require.Nil(t, got.Groups)
|
|
}
|
|
|
|
func TestBuildSchedulerMetadataAccount_KeepsQuotaAutoPauseFields(t *testing.T) {
|
|
account := service.Account{
|
|
ID: 88,
|
|
Extra: map[string]any{
|
|
"codex_5h_used_percent": 12.34,
|
|
"codex_7d_used_percent": 56.78,
|
|
"codex_5h_reset_at": "2026-05-29T10:00:00Z",
|
|
"codex_7d_reset_at": "2026-06-01T10:00:00Z",
|
|
"codex_5h_reset_after_seconds": 300,
|
|
"codex_7d_reset_after_seconds": 600,
|
|
"codex_usage_updated_at": "2026-05-29T09:00:00Z",
|
|
"auto_pause_5h_threshold": 0.95,
|
|
"auto_pause_7d_threshold": 0.96,
|
|
"auto_pause_5h_disabled": true,
|
|
"auto_pause_7d_disabled": false,
|
|
},
|
|
}
|
|
|
|
got := buildSchedulerMetadataAccount(account)
|
|
|
|
require.Equal(t, 12.34, got.Extra["codex_5h_used_percent"])
|
|
require.Equal(t, 56.78, got.Extra["codex_7d_used_percent"])
|
|
require.Equal(t, "2026-05-29T10:00:00Z", got.Extra["codex_5h_reset_at"])
|
|
require.Equal(t, "2026-06-01T10:00:00Z", got.Extra["codex_7d_reset_at"])
|
|
require.Equal(t, 300, got.Extra["codex_5h_reset_after_seconds"])
|
|
require.Equal(t, 600, got.Extra["codex_7d_reset_after_seconds"])
|
|
require.Equal(t, "2026-05-29T09:00:00Z", got.Extra["codex_usage_updated_at"])
|
|
require.Equal(t, 0.95, got.Extra["auto_pause_5h_threshold"])
|
|
require.Equal(t, 0.96, got.Extra["auto_pause_7d_threshold"])
|
|
require.Equal(t, true, got.Extra["auto_pause_5h_disabled"])
|
|
require.Equal(t, false, got.Extra["auto_pause_7d_disabled"])
|
|
}
|
|
|
|
func TestBuildSchedulerMetadataAccount_KeepsQuotaStateForCachedAccounts(t *testing.T) {
|
|
now := time.Now().UTC()
|
|
activeStart := now.Add(-time.Hour).Format(time.RFC3339)
|
|
expiredDailyStart := now.Add(-25 * time.Hour).Format(time.RFC3339)
|
|
expiredWeeklyStart := now.Add(-8 * 24 * time.Hour).Format(time.RFC3339)
|
|
weeklyResetDay := float64(now.AddDate(0, 0, 1).Weekday())
|
|
|
|
cases := []struct {
|
|
name string
|
|
platform string
|
|
typ string
|
|
extra map[string]any
|
|
quotaExceeded bool
|
|
}{
|
|
{
|
|
name: "anthropic api key total quota exhausted", platform: service.PlatformAnthropic, typ: service.AccountTypeAPIKey,
|
|
extra: map[string]any{"quota_limit": 10.0, "quota_used": 10.0}, quotaExceeded: true,
|
|
},
|
|
{
|
|
name: "gemini api key rolling daily quota exhausted", platform: service.PlatformGemini, typ: service.AccountTypeAPIKey,
|
|
extra: map[string]any{
|
|
"quota_daily_limit": 20.0, "quota_daily_used": 20.0,
|
|
"quota_daily_start": activeStart, "quota_daily_reset_mode": "rolling",
|
|
}, quotaExceeded: true,
|
|
},
|
|
{
|
|
name: "gemini api key expired rolling daily window", platform: service.PlatformGemini, typ: service.AccountTypeAPIKey,
|
|
extra: map[string]any{
|
|
"quota_daily_limit": 20.0, "quota_daily_used": 20.0,
|
|
"quota_daily_start": expiredDailyStart, "quota_daily_reset_mode": "rolling",
|
|
},
|
|
},
|
|
{
|
|
name: "bedrock fixed weekly quota exhausted", platform: service.PlatformAnthropic, typ: service.AccountTypeBedrock,
|
|
extra: map[string]any{
|
|
"quota_weekly_limit": 30.0, "quota_weekly_used": 30.0, "quota_weekly_start": activeStart,
|
|
"quota_weekly_reset_mode": "fixed", "quota_weekly_reset_day": weeklyResetDay,
|
|
"quota_weekly_reset_hour": 0.0, "quota_reset_timezone": "UTC",
|
|
}, quotaExceeded: true,
|
|
},
|
|
{
|
|
name: "bedrock expired fixed weekly window", platform: service.PlatformAnthropic, typ: service.AccountTypeBedrock,
|
|
extra: map[string]any{
|
|
"quota_weekly_limit": 30.0, "quota_weekly_used": 30.0, "quota_weekly_start": expiredWeeklyStart,
|
|
"quota_weekly_reset_mode": "fixed", "quota_weekly_reset_day": weeklyResetDay,
|
|
"quota_weekly_reset_hour": 0.0, "quota_reset_timezone": "UTC",
|
|
},
|
|
},
|
|
}
|
|
for i, tc := range cases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
extra := make(map[string]any, len(tc.extra)+1)
|
|
for key, value := range tc.extra {
|
|
extra[key] = value
|
|
}
|
|
extra["unrelated"] = "drop me"
|
|
account := service.Account{
|
|
ID: int64(46690 + i), Platform: tc.platform, Type: tc.typ, Extra: extra,
|
|
Status: service.StatusActive, Schedulable: true,
|
|
}
|
|
cache := newSchedulerCacheUnit(t)
|
|
ctx := context.Background()
|
|
bucket := service.SchedulerBucket{GroupID: int64(46690 + i), Platform: tc.platform, Mode: service.SchedulerModeSingle}
|
|
token, err := cache.CaptureBucketWriteToken(ctx, bucket)
|
|
require.NoError(t, err)
|
|
require.NoError(t, cache.SetSnapshot(ctx, bucket, token, []service.Account{account}))
|
|
|
|
snapshot, hit, err := cache.GetSnapshot(ctx, bucket)
|
|
require.NoError(t, err)
|
|
require.True(t, hit)
|
|
require.Len(t, snapshot, 1)
|
|
cached := snapshot[0]
|
|
require.Equal(t, tc.extra, cached.Extra)
|
|
require.NotContains(t, cached.Extra, "unrelated")
|
|
require.Equal(t, tc.quotaExceeded, cached.IsQuotaExceeded())
|
|
require.Equal(t, !tc.quotaExceeded, cached.IsSchedulable())
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestBuildSchedulerMetadataAccount_KeepsModelRateLimits(t *testing.T) {
|
|
account := service.Account{
|
|
ID: 90,
|
|
Platform: service.PlatformAntigravity,
|
|
Extra: map[string]any{
|
|
"model_rate_limits": map[string]any{
|
|
"gemini-3-flash": map[string]any{
|
|
"rate_limit_reset_at": "2026-05-30T10:10:00Z",
|
|
},
|
|
"antigravity:gemini": map[string]any{
|
|
"rate_limit_reset_at": "2026-05-30T10:10:00Z",
|
|
},
|
|
},
|
|
"unused_large_field": "drop-me",
|
|
},
|
|
}
|
|
|
|
got := buildSchedulerMetadataAccount(account)
|
|
|
|
limits, ok := got.Extra["model_rate_limits"].(map[string]any)
|
|
require.True(t, ok)
|
|
require.Contains(t, limits, "gemini-3-flash")
|
|
require.Contains(t, limits, "antigravity:gemini")
|
|
require.Nil(t, got.Extra["unused_large_field"])
|
|
}
|
|
|
|
func TestBuildSchedulerMetadataAccount_KeepsSparkShadowRoutingIdentity(t *testing.T) {
|
|
parentID := int64(100)
|
|
account := service.Account{
|
|
ID: 200,
|
|
Platform: service.PlatformOpenAI,
|
|
Type: service.AccountTypeOAuth,
|
|
ParentAccountID: &parentID,
|
|
QuotaDimension: service.QuotaDimensionSpark,
|
|
Credentials: map[string]any{
|
|
"model_mapping": map[string]any{
|
|
"gpt-5.3-codex-spark": "gpt-5.3-codex-spark",
|
|
},
|
|
"compact_model_mapping": map[string]any{
|
|
"gpt-5.4": "gpt-5.4-openai-compact",
|
|
},
|
|
"access_token": "drop-me",
|
|
},
|
|
}
|
|
|
|
got := buildSchedulerMetadataAccount(account)
|
|
|
|
require.NotNil(t, got.ParentAccountID)
|
|
require.Equal(t, parentID, *got.ParentAccountID)
|
|
require.Equal(t, service.QuotaDimensionSpark, got.QuotaDimension)
|
|
require.Equal(t, map[string]any{"gpt-5.3-codex-spark": "gpt-5.3-codex-spark"}, got.Credentials["model_mapping"])
|
|
require.Equal(t, map[string]any{"gpt-5.4": "gpt-5.4-openai-compact"}, got.Credentials["compact_model_mapping"])
|
|
require.Nil(t, got.Credentials["access_token"])
|
|
}
|
|
|
|
func TestSchedulerCacheBucketRetirementFencesWritersAndReopen(t *testing.T) {
|
|
ctx := context.Background()
|
|
cache, mr := newSchedulerCacheUnitWithRedis(t)
|
|
bucket := service.SchedulerBucket{GroupID: 41, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
|
|
otherBucket := service.SchedulerBucket{GroupID: 42, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
|
|
account := service.Account{ID: 4101, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
|
|
|
|
token, err := cache.CaptureBucketWriteToken(ctx, bucket)
|
|
require.NoError(t, err)
|
|
require.True(t, token.ValidFor(bucket))
|
|
require.NoError(t, cache.SetSnapshot(ctx, bucket, token, []service.Account{account}))
|
|
|
|
// A token is bound to the full bucket identity, not just an epoch number.
|
|
err = cache.SetSnapshot(ctx, otherBucket, token, []service.Account{account})
|
|
require.ErrorIs(t, err, service.ErrSchedulerBucketWriteFenced)
|
|
_, err = cache.rdb.Get(ctx, schedulerBucketKey(schedulerVersionPrefix, otherBucket)).Result()
|
|
require.ErrorIs(t, err, redis.Nil)
|
|
otherAccount := service.Account{ID: 4201, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
|
|
otherToken, err := cache.CaptureBucketWriteToken(ctx, otherBucket)
|
|
require.NoError(t, err)
|
|
require.NoError(t, cache.SetSnapshot(ctx, otherBucket, otherToken, []service.Account{otherAccount}))
|
|
otherEpoch := otherToken.Epoch
|
|
|
|
activeVersion, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Result()
|
|
require.NoError(t, err)
|
|
require.NoError(t, cache.RetireBucket(ctx, bucket))
|
|
retiredEpoch, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerEpochPrefix, bucket)).Int64()
|
|
require.NoError(t, err)
|
|
require.Greater(t, retiredEpoch, token.Epoch)
|
|
|
|
// Retirement is idempotent and does not advance the epoch again.
|
|
require.NoError(t, cache.RetireBucket(ctx, bucket))
|
|
retiredEpochAgain, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerEpochPrefix, bucket)).Int64()
|
|
require.NoError(t, err)
|
|
require.Equal(t, retiredEpoch, retiredEpochAgain)
|
|
|
|
// New readers miss because ready/active were removed atomically. A reader that
|
|
// captured activeVersion before retirement may still finish against that version.
|
|
_, hit, err := cache.GetSnapshot(ctx, bucket)
|
|
require.NoError(t, err)
|
|
require.False(t, hit)
|
|
ids, err := cache.rdb.ZRange(ctx, schedulerSnapshotKey(bucket, activeVersion), 0, -1).Result()
|
|
require.NoError(t, err)
|
|
require.Equal(t, []string{"4101"}, ids)
|
|
ttl, err := cache.rdb.TTL(ctx, schedulerSnapshotKey(bucket, activeVersion)).Result()
|
|
require.NoError(t, err)
|
|
require.Positive(t, ttl)
|
|
require.LessOrEqual(t, ttl, time.Duration(snapshotGraceTTLSeconds)*time.Second)
|
|
|
|
buckets, err := cache.ListBuckets(ctx)
|
|
require.NoError(t, err)
|
|
require.NotContains(t, buckets, bucket)
|
|
require.Contains(t, buckets, otherBucket)
|
|
otherSnapshot, otherHit, err := cache.GetSnapshot(ctx, otherBucket)
|
|
require.NoError(t, err)
|
|
require.True(t, otherHit)
|
|
require.Len(t, otherSnapshot, 1)
|
|
require.Equal(t, otherAccount.ID, otherSnapshot[0].ID)
|
|
otherEpochAfter, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerEpochPrefix, otherBucket)).Int64()
|
|
require.NoError(t, err)
|
|
require.Equal(t, otherEpoch, otherEpochAfter)
|
|
|
|
_, err = cache.CaptureBucketWriteToken(ctx, bucket)
|
|
require.ErrorIs(t, err, service.ErrSchedulerBucketRetired)
|
|
versionBeforeRejectedWrite, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerVersionPrefix, bucket)).Int64()
|
|
require.NoError(t, err)
|
|
err = cache.SetSnapshot(ctx, bucket, token, []service.Account{account})
|
|
require.ErrorIs(t, err, service.ErrSchedulerBucketRetired)
|
|
versionAfterRejectedWrite, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerVersionPrefix, bucket)).Int64()
|
|
require.NoError(t, err)
|
|
require.Equal(t, versionBeforeRejectedWrite, versionAfterRejectedWrite, "fenced writers must not allocate a new version")
|
|
retired, err := cache.rdb.Exists(ctx, schedulerBucketKey(schedulerRetiredPrefix, bucket)).Result()
|
|
require.NoError(t, err)
|
|
require.EqualValues(t, 1, retired, "ordinary writers must never clear the tombstone")
|
|
mr.FastForward(time.Duration(snapshotGraceTTLSeconds+1) * time.Second)
|
|
exists, err := cache.rdb.Exists(ctx, schedulerSnapshotKey(bucket, activeVersion)).Result()
|
|
require.NoError(t, err)
|
|
require.Zero(t, exists, "retired active snapshot must expire after the in-flight grace period")
|
|
|
|
newToken, err := cache.ReopenBucket(ctx, bucket)
|
|
require.NoError(t, err)
|
|
require.True(t, newToken.ValidFor(bucket))
|
|
require.Equal(t, retiredEpoch, newToken.Epoch)
|
|
reopenedAgain, err := cache.ReopenBucket(ctx, bucket)
|
|
require.NoError(t, err)
|
|
require.Equal(t, newToken, reopenedAgain, "reopen must be idempotent within one retirement generation")
|
|
err = cache.SetSnapshot(ctx, bucket, token, []service.Account{account})
|
|
require.ErrorIs(t, err, service.ErrSchedulerBucketWriteFenced)
|
|
require.NoError(t, cache.SetSnapshot(ctx, bucket, newToken, []service.Account{account}))
|
|
reopenedWhileOpen, err := cache.ReopenBucket(ctx, bucket)
|
|
require.NoError(t, err)
|
|
require.Equal(t, newToken, reopenedWhileOpen)
|
|
|
|
snapshot, hit, err := cache.GetSnapshot(ctx, bucket)
|
|
require.NoError(t, err)
|
|
require.True(t, hit)
|
|
require.Len(t, snapshot, 1)
|
|
require.Equal(t, account.ID, snapshot[0].ID)
|
|
}
|
|
|
|
func TestSchedulerCacheActivationIsFencedAfterRetire(t *testing.T) {
|
|
ctx := context.Background()
|
|
cache := newSchedulerCacheUnit(t)
|
|
bucket := service.SchedulerBucket{GroupID: 51, Platform: service.PlatformAnthropic, Mode: service.SchedulerModeMixed}
|
|
account := service.Account{ID: 5101, Platform: service.PlatformAnthropic, Type: service.AccountTypeAPIKey}
|
|
|
|
token, err := cache.CaptureBucketWriteToken(ctx, bucket)
|
|
require.NoError(t, err)
|
|
version, err := cache.allocateSnapshotVersion(ctx, bucket, token)
|
|
require.NoError(t, err)
|
|
_, err = cache.writeSnapshotVersionAndReturnAccountIDs(ctx, bucket, version, []service.Account{account})
|
|
require.NoError(t, err)
|
|
|
|
// Deterministic race C: retirement and authoritative reopen both happen after
|
|
// INCR/write but before the old writer activates.
|
|
require.NoError(t, cache.RetireBucket(ctx, bucket))
|
|
_, err = cache.ReopenBucket(ctx, bucket)
|
|
require.NoError(t, err)
|
|
err = cache.activateSnapshotVersion(ctx, bucket, token, version)
|
|
require.ErrorIs(t, err, service.ErrSchedulerBucketWriteFenced)
|
|
|
|
exists, err := cache.rdb.Exists(ctx, schedulerSnapshotKey(bucket, version)).Result()
|
|
require.NoError(t, err)
|
|
require.Zero(t, exists, "fenced activation must delete its unpublished snapshot")
|
|
exists, err = cache.rdb.Exists(
|
|
ctx,
|
|
schedulerBucketKey(schedulerReadyPrefix, bucket),
|
|
schedulerBucketKey(schedulerActivePrefix, bucket),
|
|
).Result()
|
|
require.NoError(t, err)
|
|
require.Zero(t, exists)
|
|
buckets, err := cache.ListBuckets(ctx)
|
|
require.NoError(t, err)
|
|
require.NotContains(t, buckets, bucket)
|
|
}
|
|
|
|
func TestSchedulerCacheConcurrentReopenReturnsSameToken(t *testing.T) {
|
|
ctx := context.Background()
|
|
cache := newSchedulerCacheUnit(t)
|
|
bucket := service.SchedulerBucket{GroupID: 53, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
|
|
account := service.Account{ID: 5301, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
|
|
|
|
oldToken, err := cache.CaptureBucketWriteToken(ctx, bucket)
|
|
require.NoError(t, err)
|
|
require.NoError(t, cache.RetireBucket(ctx, bucket))
|
|
|
|
type reopenResult struct {
|
|
token service.SchedulerBucketWriteToken
|
|
err error
|
|
}
|
|
start := make(chan struct{})
|
|
results := make(chan reopenResult, 2)
|
|
for range 2 {
|
|
go func() {
|
|
<-start
|
|
token, err := cache.ReopenBucket(ctx, bucket)
|
|
results <- reopenResult{token: token, err: err}
|
|
}()
|
|
}
|
|
close(start)
|
|
first := <-results
|
|
second := <-results
|
|
require.NoError(t, first.err)
|
|
require.NoError(t, second.err)
|
|
require.Equal(t, first.token, second.token)
|
|
require.Greater(t, first.token.Epoch, oldToken.Epoch)
|
|
|
|
require.ErrorIs(t, cache.SetSnapshot(ctx, bucket, oldToken, []service.Account{account}), service.ErrSchedulerBucketWriteFenced)
|
|
require.NoError(t, cache.SetSnapshot(ctx, bucket, first.token, []service.Account{account}))
|
|
}
|
|
|
|
func TestSchedulerCacheReopenExpiresPreviousActiveSnapshot(t *testing.T) {
|
|
ctx := context.Background()
|
|
cache, mr := newSchedulerCacheUnitWithRedis(t)
|
|
bucket := service.SchedulerBucket{GroupID: 52, Platform: service.PlatformGemini, Mode: service.SchedulerModeForced}
|
|
account := service.Account{ID: 5201, Platform: service.PlatformGemini, Type: service.AccountTypeAPIKey}
|
|
|
|
oldToken, err := cache.CaptureBucketWriteToken(ctx, bucket)
|
|
require.NoError(t, err)
|
|
require.NoError(t, cache.SetSnapshot(ctx, bucket, oldToken, []service.Account{account}))
|
|
oldVersion, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Result()
|
|
require.NoError(t, err)
|
|
retiredEpoch := oldToken.Epoch + 1
|
|
require.NoError(t, cache.rdb.Set(ctx, schedulerBucketKey(schedulerEpochPrefix, bucket), retiredEpoch, 0).Err())
|
|
require.NoError(t, cache.rdb.Set(ctx, schedulerBucketKey(schedulerRetiredPrefix, bucket), retiredEpoch, 0).Err())
|
|
|
|
newToken, err := cache.ReopenBucket(ctx, bucket)
|
|
require.NoError(t, err)
|
|
require.Equal(t, retiredEpoch, newToken.Epoch)
|
|
_, hit, err := cache.GetSnapshot(ctx, bucket)
|
|
require.NoError(t, err)
|
|
require.False(t, hit)
|
|
ttl, err := cache.rdb.TTL(ctx, schedulerSnapshotKey(bucket, oldVersion)).Result()
|
|
require.NoError(t, err)
|
|
require.Positive(t, ttl)
|
|
require.LessOrEqual(t, ttl, time.Duration(snapshotGraceTTLSeconds)*time.Second)
|
|
|
|
require.ErrorIs(t, cache.SetSnapshot(ctx, bucket, oldToken, []service.Account{account}), service.ErrSchedulerBucketWriteFenced)
|
|
mr.FastForward(time.Duration(snapshotGraceTTLSeconds+1) * time.Second)
|
|
exists, err := cache.rdb.Exists(ctx, schedulerSnapshotKey(bucket, oldVersion)).Result()
|
|
require.NoError(t, err)
|
|
require.Zero(t, exists)
|
|
require.NoError(t, cache.SetSnapshot(ctx, bucket, newToken, []service.Account{account}))
|
|
}
|
|
|
|
func TestSchedulerCacheGroupLifecycleLeaseConcurrentAcquireSingleOwner(t *testing.T) {
|
|
ctx := context.Background()
|
|
cache := newSchedulerCacheUnit(t)
|
|
const groupID int64 = 71
|
|
|
|
type result struct {
|
|
lease service.SchedulerGroupLifecycleLease
|
|
acquired bool
|
|
err error
|
|
}
|
|
start := make(chan struct{})
|
|
results := make(chan result, 32)
|
|
for range 32 {
|
|
go func() {
|
|
<-start
|
|
lease, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, groupID, time.Minute)
|
|
results <- result{lease: lease, acquired: acquired, err: err}
|
|
}()
|
|
}
|
|
close(start)
|
|
|
|
var owner service.SchedulerGroupLifecycleLease
|
|
acquiredCount := 0
|
|
for range 32 {
|
|
got := <-results
|
|
require.NoError(t, got.err)
|
|
if got.acquired {
|
|
acquiredCount++
|
|
owner = got.lease
|
|
require.True(t, got.lease.ValidFor(groupID))
|
|
} else {
|
|
require.Equal(t, service.SchedulerGroupLifecycleLease{}, got.lease)
|
|
}
|
|
}
|
|
require.Equal(t, 1, acquiredCount)
|
|
require.Len(t, owner.OwnerToken, schedulerGroupLifecycleOwnerTokenBytes*2)
|
|
require.Equal(t, strings.ToLower(owner.OwnerToken), owner.OwnerToken)
|
|
decodedOwner, err := hex.DecodeString(owner.OwnerToken)
|
|
require.NoError(t, err)
|
|
require.Len(t, decodedOwner, schedulerGroupLifecycleOwnerTokenBytes)
|
|
|
|
require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, owner))
|
|
next, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, groupID, time.Minute)
|
|
require.NoError(t, err)
|
|
require.True(t, acquired)
|
|
require.True(t, next.ValidFor(groupID))
|
|
require.NotEqual(t, owner.OwnerToken, next.OwnerToken)
|
|
require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, next))
|
|
}
|
|
|
|
func TestSchedulerCacheGroupLifecycleLeaseStaleReleaseCannotDeleteSuccessor(t *testing.T) {
|
|
ctx := context.Background()
|
|
cache, mr := newSchedulerCacheUnitWithRedis(t)
|
|
const groupID int64 = 72
|
|
const ttl = time.Minute
|
|
|
|
first, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, groupID, ttl)
|
|
require.NoError(t, err)
|
|
require.True(t, acquired)
|
|
|
|
mr.FastForward(ttl + time.Second)
|
|
second, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, groupID, ttl)
|
|
require.NoError(t, err)
|
|
require.True(t, acquired)
|
|
require.NotEqual(t, first.OwnerToken, second.OwnerToken)
|
|
|
|
require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(ctx, first), service.ErrSchedulerGroupLifecycleLeaseLost)
|
|
owner, err := cache.rdb.Get(ctx, schedulerGroupLifecycleLockKey(groupID)).Result()
|
|
require.NoError(t, err)
|
|
require.Equal(t, second.OwnerToken, owner)
|
|
|
|
_, acquired, err = cache.TryAcquireGroupLifecycleLease(ctx, groupID, ttl)
|
|
require.NoError(t, err)
|
|
require.False(t, acquired)
|
|
require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, second))
|
|
require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(ctx, second), service.ErrSchedulerGroupLifecycleLeaseLost)
|
|
}
|
|
|
|
func TestSchedulerCacheGroupLifecycleLeaseExpiredReleaseIsLost(t *testing.T) {
|
|
ctx := context.Background()
|
|
cache, mr := newSchedulerCacheUnitWithRedis(t)
|
|
const groupID int64 = 73
|
|
const ttl = time.Minute
|
|
|
|
lease, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, groupID, ttl)
|
|
require.NoError(t, err)
|
|
require.True(t, acquired)
|
|
mr.FastForward(ttl + time.Second)
|
|
|
|
require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(ctx, lease), service.ErrSchedulerGroupLifecycleLeaseLost)
|
|
}
|
|
|
|
func TestSchedulerCacheGroupLifecycleLeaseWrongOwnerAndCrossGroupAreLost(t *testing.T) {
|
|
ctx := context.Background()
|
|
cache := newSchedulerCacheUnit(t)
|
|
const firstGroupID int64 = 74
|
|
const secondGroupID int64 = 75
|
|
|
|
first, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, firstGroupID, time.Minute)
|
|
require.NoError(t, err)
|
|
require.True(t, acquired)
|
|
second, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, secondGroupID, time.Minute)
|
|
require.NoError(t, err)
|
|
require.True(t, acquired, "different groups must acquire independently")
|
|
require.NotEqual(t, first.OwnerToken, second.OwnerToken)
|
|
|
|
wrongOwner := first
|
|
wrongOwner.OwnerToken = strings.Repeat("0", schedulerGroupLifecycleOwnerTokenBytes*2)
|
|
if wrongOwner.OwnerToken == first.OwnerToken {
|
|
wrongOwner.OwnerToken = strings.Repeat("1", schedulerGroupLifecycleOwnerTokenBytes*2)
|
|
}
|
|
require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(ctx, wrongOwner), service.ErrSchedulerGroupLifecycleLeaseLost)
|
|
|
|
crossGroup := first
|
|
crossGroup.GroupID = secondGroupID
|
|
require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(ctx, crossGroup), service.ErrSchedulerGroupLifecycleLeaseLost)
|
|
|
|
firstOwner, err := cache.rdb.Get(ctx, schedulerGroupLifecycleLockKey(firstGroupID)).Result()
|
|
require.NoError(t, err)
|
|
require.Equal(t, first.OwnerToken, firstOwner)
|
|
secondOwner, err := cache.rdb.Get(ctx, schedulerGroupLifecycleLockKey(secondGroupID)).Result()
|
|
require.NoError(t, err)
|
|
require.Equal(t, second.OwnerToken, secondOwner)
|
|
|
|
require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, first))
|
|
require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, second))
|
|
}
|
|
|
|
func TestSchedulerCacheGroupLifecycleLeaseCanceledContextFailsClosed(t *testing.T) {
|
|
cache := newSchedulerCacheUnit(t)
|
|
canceledCtx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
|
|
lease, acquired, err := cache.TryAcquireGroupLifecycleLease(canceledCtx, 76, time.Minute)
|
|
require.ErrorIs(t, err, context.Canceled)
|
|
require.False(t, acquired)
|
|
require.Equal(t, service.SchedulerGroupLifecycleLease{}, lease)
|
|
|
|
ctx := context.Background()
|
|
lease, acquired, err = cache.TryAcquireGroupLifecycleLease(ctx, 76, time.Minute)
|
|
require.NoError(t, err)
|
|
require.True(t, acquired)
|
|
require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(canceledCtx, lease), context.Canceled)
|
|
owner, err := cache.rdb.Get(ctx, schedulerGroupLifecycleLockKey(lease.GroupID)).Result()
|
|
require.NoError(t, err)
|
|
require.Equal(t, lease.OwnerToken, owner)
|
|
require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, lease))
|
|
}
|
|
|
|
func TestSchedulerCacheGroupLifecycleLeaseRejectsInvalidInput(t *testing.T) {
|
|
ctx := context.Background()
|
|
cache := newSchedulerCacheUnit(t)
|
|
|
|
lease, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, 0, time.Minute)
|
|
require.ErrorIs(t, err, service.ErrSchedulerGroupLifecycleLeaseInvalid)
|
|
require.False(t, acquired)
|
|
require.Equal(t, service.SchedulerGroupLifecycleLease{}, lease)
|
|
|
|
lease, acquired, err = cache.TryAcquireGroupLifecycleLease(ctx, 73, 0)
|
|
require.ErrorIs(t, err, service.ErrSchedulerGroupLifecycleLeaseInvalid)
|
|
require.False(t, acquired)
|
|
require.Equal(t, service.SchedulerGroupLifecycleLease{}, lease)
|
|
|
|
canceledCtx, cancel := context.WithCancel(ctx)
|
|
cancel()
|
|
require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(canceledCtx, service.SchedulerGroupLifecycleLease{}), service.ErrSchedulerGroupLifecycleLeaseInvalid)
|
|
keys, err := cache.rdb.DBSize(ctx).Result()
|
|
require.NoError(t, err)
|
|
require.Zero(t, keys)
|
|
}
|
|
|
|
var schedulerCachePayloadBenchmarkSink int
|
|
|
|
func BenchmarkSchedulerCacheAccountPayloadReuse(b *testing.B) {
|
|
for _, size := range []int{1, 100, 10_000} {
|
|
accounts := schedulerCacheBenchmarkAccounts(size)
|
|
b.Run(fmt.Sprintf("pair_baseline_%d_accounts", size), func(b *testing.B) {
|
|
b.ReportAllocs()
|
|
for i := 0; i < b.N; i++ {
|
|
first, err := benchmarkSchedulerLegacySnapshotPayload(accounts)
|
|
if err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
second, err := benchmarkSchedulerLegacySnapshotPayload(accounts)
|
|
if err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
schedulerCachePayloadBenchmarkSink = first + second
|
|
}
|
|
})
|
|
b.Run(fmt.Sprintf("pair_reuse_%d_accounts", size), func(b *testing.B) {
|
|
b.ReportAllocs()
|
|
for i := 0; i < b.N; i++ {
|
|
ids, total, err := benchmarkSchedulerReusableSnapshotPayload(accounts)
|
|
if err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
// 第二个桶仍构造成员，只跳过账号 JSON 与全局账号键。
|
|
total += len(schedulerSnapshotMembers(ids))
|
|
schedulerCachePayloadBenchmarkSink = total
|
|
}
|
|
})
|
|
b.Run(fmt.Sprintf("first_baseline_%d_accounts", size), func(b *testing.B) {
|
|
b.ReportAllocs()
|
|
for i := 0; i < b.N; i++ {
|
|
total, err := benchmarkSchedulerLegacySnapshotPayload(accounts)
|
|
if err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
schedulerCachePayloadBenchmarkSink = total
|
|
}
|
|
})
|
|
b.Run(fmt.Sprintf("first_reuse_%d_accounts", size), func(b *testing.B) {
|
|
b.ReportAllocs()
|
|
for i := 0; i < b.N; i++ {
|
|
ids, total, err := benchmarkSchedulerReusableSnapshotPayload(accounts)
|
|
if err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
total += len(ids)
|
|
schedulerCachePayloadBenchmarkSink = total
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func benchmarkSchedulerLegacySnapshotPayload(accounts []service.Account) (int, error) {
|
|
cacheable := make([]service.Account, 0, len(accounts))
|
|
total := 0
|
|
for _, account := range accounts {
|
|
full, meta, err := marshalSchedulerCacheAccount(account)
|
|
if err != nil {
|
|
continue
|
|
}
|
|
total += len(full) + len(meta)
|
|
cacheable = append(cacheable, account)
|
|
}
|
|
members := make([]redis.Z, 0, len(cacheable))
|
|
for idx, account := range cacheable {
|
|
members = append(members, redis.Z{Score: float64(idx), Member: strconv.FormatInt(account.ID, 10)})
|
|
}
|
|
return total + len(members), nil
|
|
}
|
|
|
|
func benchmarkSchedulerReusableSnapshotPayload(accounts []service.Account) ([]int64, int, error) {
|
|
accountIDs := make([]int64, 0, len(accounts))
|
|
total := 0
|
|
for _, account := range accounts {
|
|
full, meta, err := marshalSchedulerCacheAccount(account)
|
|
if err != nil {
|
|
continue
|
|
}
|
|
total += len(full) + len(meta)
|
|
accountIDs = append(accountIDs, account.ID)
|
|
}
|
|
total += len(schedulerSnapshotMembers(accountIDs))
|
|
return accountIDs, total, nil
|
|
}
|
|
|
|
func schedulerCacheBenchmarkAccounts(size int) []service.Account {
|
|
largeValue := strings.Repeat("x", 4096)
|
|
credentials := map[string]any{
|
|
"api_key": "benchmark-key",
|
|
"model_mapping": map[string]any{"z-model": "z-target", "a-model": "a-target"},
|
|
"large_value": largeValue,
|
|
}
|
|
extra := map[string]any{
|
|
"mixed_scheduling": true,
|
|
"model_rate_limits": map[string]any{
|
|
"z-model": map[string]any{"rate_limit_reset_at": "2026-07-16T00:00:00Z"},
|
|
"a-model": map[string]any{"rate_limit_reset_at": "2026-07-16T00:00:00Z"},
|
|
},
|
|
"large_value": largeValue,
|
|
}
|
|
accounts := make([]service.Account, size)
|
|
for i := range accounts {
|
|
id := int64(i + 1)
|
|
accounts[i] = service.Account{
|
|
ID: id,
|
|
Name: "benchmark-account",
|
|
Platform: service.PlatformOpenAI,
|
|
Type: service.AccountTypeOAuth,
|
|
Credentials: credentials,
|
|
Extra: extra,
|
|
GroupIDs: []int64{7, 9},
|
|
AccountGroups: []service.AccountGroup{
|
|
{AccountID: id, GroupID: 7, Priority: 1},
|
|
{AccountID: id, GroupID: 9, Priority: 2},
|
|
},
|
|
}
|
|
}
|
|
return accounts
|
|
}
|
|
|
|
// 调度投影必须保留 OpenAI 透传开关。
|
|
//
|
|
// 候选过滤走 ListSchedulableAccounts，读的是 buildSchedulerMetadataAccount 产出的精简投影；
|
|
// Account.IsModelSupported 又靠 extra 上的透传开关短路 model_mapping 白名单（#4936）。
|
|
// 一旦投影把开关裁掉、却保留了白名单，透传账号在选号阶段就会退回白名单判定并被误判成
|
|
// model_not_supported，而转发阶段（读完整账号）仍按透传工作 —— 表现为"单独测这个账号能通、
|
|
// 走网关却报 no available accounts"。#4936 修的是判定逻辑，这里守的是喂给判定的输入。
|
|
func TestBuildSchedulerMetadataAccount_KeepsOpenAIPassthroughForModelGate(t *testing.T) {
|
|
for _, key := range []string{"openai_passthrough", "openai_oauth_passthrough"} {
|
|
t.Run(key, func(t *testing.T) {
|
|
account := service.Account{
|
|
ID: 383,
|
|
Platform: service.PlatformOpenAI,
|
|
Type: service.AccountTypeOAuth,
|
|
Credentials: map[string]any{
|
|
// 账号从白名单模式切到透传后常见的残留映射，未列出请求的模型。
|
|
"model_mapping": map[string]any{"gpt-5.5": "gpt-5.5"},
|
|
"access_token": "drop-me",
|
|
},
|
|
Extra: map[string]any{key: true},
|
|
}
|
|
require.True(t, account.IsModelSupported("gpt-5.6-sol"),
|
|
"前置条件：透传账号本应放行白名单外的模型")
|
|
|
|
meta := buildSchedulerMetadataAccount(account)
|
|
|
|
// 走一遍真实的序列化/反序列化路径（写入 sched:meta 再由 decodeCachedAccount 读回）。
|
|
payload, err := json.Marshal(meta)
|
|
require.NoError(t, err)
|
|
var restored service.Account
|
|
require.NoError(t, json.Unmarshal(payload, &restored))
|
|
|
|
require.Equal(t, true, restored.Extra[key])
|
|
require.True(t, restored.IsOpenAIPassthroughEnabled())
|
|
require.True(t, restored.IsModelSupported("gpt-5.6-sol"),
|
|
"投影裁掉透传开关会让透传账号在候选过滤阶段被误判为 model_not_supported")
|
|
// 白名单本身仍需保留：非透传账号依赖它做模型门。
|
|
require.Equal(t, map[string]any{"gpt-5.5": "gpt-5.5"}, restored.Credentials["model_mapping"])
|
|
})
|
|
}
|
|
}
|