mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-01 15:02:58 +08:00
refactor(monitor): pass loaded account through quota sources (single load per fetch)
配额 fetcher 缓存未命中时账号被加载两次:fetchUncached GetByID 一次, 下游数据源(GetUsage / CN QueryUsage / CN QueryBalance)按 ID 再各自 GetByID 一次——每次 GetByID 含 proxies/groups 联查,纯属重复劳动。 统一改为「路由前加载一次、指针直传」: - fetcher 三个数据源接口签名改收 *Account: GetUsageForAccount / QueryUsageForAccount / QueryBalanceForAccount - AccountUsageService 暴露 GetUsageForAccount 直通(getUsageForAccount 既有逻辑零改动) - CN 两侧服务抽 validateCodingPlanAccount / validatePayGAccount (加载后校验原样),新 ForAccount 入口 = NOT_CONFIGURED 守卫 → 校验 → singleflight → 探测;ID 入口 = load + 委托,语义不变 - singleflight key 保持 "cn_quota:<id>" / "cn_balance:<id>":ForAccount 与 admin handler 的 ID 入口并发探测仍按账号合并 - cn_provider_balance_check_service.checkOne 刻意留在 ID 入口:调度器 路径按 ID 取的是最新凭据,不复用可能已过期的已加载账号 校验移到 flight 之外且 ForAccount 复用同一套校验,直传无法绕过 平台/模式检查;专项测试固化(无效账号零出站请求、单次加载指针直传 require.Same)。
This commit is contained in:
@@ -503,6 +503,13 @@ func (s *AccountUsageService) GetUsage(ctx context.Context, accountID int64, for
|
||||
return s.getUsageForAccount(ctx, account, forceProbe)
|
||||
}
|
||||
|
||||
// GetUsageForAccount 已加载账号的使用量直通入口(配额监控 fetcher 复用,
|
||||
// 避免缓存未命中时账号被加载两次——每次 GetByID 含 proxies/groups 联查)。
|
||||
func (s *AccountUsageService) GetUsageForAccount(ctx context.Context, account *Account, force ...bool) (*UsageInfo, error) {
|
||||
forceProbe := len(force) > 0 && force[0]
|
||||
return s.getUsageForAccount(ctx, account, forceProbe)
|
||||
}
|
||||
|
||||
// GetUsageBatch 批量获取账号使用量。
|
||||
// Anthropic OAuth/SetupToken 统一走 passive 链路,其他账号复用现有主动查询逻辑。
|
||||
// 单个账号失败不会中断整批请求,错误会按账号返回。
|
||||
|
||||
@@ -19,10 +19,12 @@ import (
|
||||
// 渠道监控「配额模式」的配额抓取器。
|
||||
//
|
||||
// 不直接对接上游,而是把账号侧现成的用量服务归一成 domain.MonitorQuotaSnapshot:
|
||||
// - 海外 5 家(anthropic/openai/gemini/antigravity/grok)→ AccountUsageService.GetUsage
|
||||
// - 国产 coding plan(kimi/zhipu/deepseek)→ CNProviderQuotaService.QueryUsage
|
||||
// - 国产 payg(kimi/deepseek)→ CNProviderBalanceService.QueryBalance
|
||||
// (zhipu payg 无公开余额端点,QueryBalance 会返回该错误,原样透出)
|
||||
// - 海外 5 家(anthropic/openai/gemini/antigravity/grok)→ AccountUsageService.GetUsageForAccount
|
||||
// - 国产 coding plan(kimi/zhipu/deepseek)→ CNProviderQuotaService.QueryUsageForAccount
|
||||
// - 国产 payg(kimi/deepseek)→ CNProviderBalanceService.QueryBalanceForAccount
|
||||
// (zhipu payg 无公开余额端点,探测会返回该错误,原样透出)
|
||||
// 数据源统一接受已加载的 *Account:fetchUncached 路由前 GetByID 一次并传下去,
|
||||
// 下游服务不再各自重载(每次 GetByID 含 proxies/groups 联查)。
|
||||
//
|
||||
// Fetch 永不返回 error:所有失败都降级为 Success=false 的快照照常入库,
|
||||
// 由 deriveQuotaCheckResult 推导为 failed/error 状态。
|
||||
@@ -33,18 +35,19 @@ import (
|
||||
// 抓取由 singleflight 合并为一次上游查询。
|
||||
|
||||
// monitorUsageSource 海外平台账号用量查询(AccountUsageService 天然满足)。
|
||||
// 传已加载的 *Account:fetchUncached 只 GetByID 一次,下游不再重复加载。
|
||||
type monitorUsageSource interface {
|
||||
GetUsage(ctx context.Context, accountID int64, force ...bool) (*UsageInfo, error)
|
||||
GetUsageForAccount(ctx context.Context, account *Account, force ...bool) (*UsageInfo, error)
|
||||
}
|
||||
|
||||
// monitorCNQuotaSource 国产 coding plan 滚动窗口额度探测(CNProviderQuotaService 天然满足)。
|
||||
type monitorCNQuotaSource interface {
|
||||
QueryUsage(ctx context.Context, accountID int64) (*CNProviderQuotaProbeResult, error)
|
||||
QueryUsageForAccount(ctx context.Context, account *Account) (*CNProviderQuotaProbeResult, error)
|
||||
}
|
||||
|
||||
// monitorCNBalanceSource 国产 payg 余额探测(CNProviderBalanceService 天然满足)。
|
||||
type monitorCNBalanceSource interface {
|
||||
QueryBalance(ctx context.Context, accountID int64) (*CNProviderBalanceResult, error)
|
||||
QueryBalanceForAccount(ctx context.Context, account *Account) (*CNProviderBalanceResult, error)
|
||||
}
|
||||
|
||||
// monitorAccountSource 账号加载(AccountRepository 天然满足)。
|
||||
@@ -191,23 +194,26 @@ func (f *ChannelMonitorQuotaFetcher) fetchUncached(ctx context.Context, accountI
|
||||
return quotaErrorSnapshot("usage", "linked account not found", now)
|
||||
}
|
||||
|
||||
// 账号只在路由前加载这一次;已加载的 account 直接传给数据源
|
||||
// (GetUsageForAccount / QueryUsageForAccount / QueryBalanceForAccount),
|
||||
// 下游服务不再各自 GetByID(每次含 proxies/groups 联查)。
|
||||
switch account.Platform {
|
||||
case domain.PlatformKimi, domain.PlatformZhipu, domain.PlatformDeepseek:
|
||||
if account.IsCodingPlan() {
|
||||
return f.fetchCNQuota(ctx, accountID, now)
|
||||
return f.fetchCNQuota(ctx, account, now)
|
||||
}
|
||||
return f.fetchCNBalance(ctx, accountID, now)
|
||||
return f.fetchCNBalance(ctx, account, now)
|
||||
default:
|
||||
return f.fetchUsage(ctx, accountID, now)
|
||||
return f.fetchUsage(ctx, account, now)
|
||||
}
|
||||
}
|
||||
|
||||
// fetchUsage 海外平台:AccountUsageService.GetUsage → 快照。
|
||||
func (f *ChannelMonitorQuotaFetcher) fetchUsage(ctx context.Context, accountID int64, now time.Time) *domain.MonitorQuotaSnapshot {
|
||||
// fetchUsage 海外平台:AccountUsageService.GetUsageForAccount → 快照。
|
||||
func (f *ChannelMonitorQuotaFetcher) fetchUsage(ctx context.Context, account *Account, now time.Time) *domain.MonitorQuotaSnapshot {
|
||||
if f.usage == nil {
|
||||
return quotaErrorSnapshot("usage", "usage service is not configured", now)
|
||||
}
|
||||
usage, err := f.usage.GetUsage(ctx, accountID)
|
||||
usage, err := f.usage.GetUsageForAccount(ctx, account)
|
||||
if err != nil {
|
||||
msg := truncateMessage(sanitizeErrorMessage(err.Error()))
|
||||
return &domain.MonitorQuotaSnapshot{
|
||||
@@ -336,12 +342,12 @@ func sortedQuotaModelNames(quotas map[string]*AntigravityModelQuota) []string {
|
||||
return names
|
||||
}
|
||||
|
||||
// fetchCNQuota 国产 coding plan:CNProviderQuotaService.QueryUsage → 快照。
|
||||
func (f *ChannelMonitorQuotaFetcher) fetchCNQuota(ctx context.Context, accountID int64, now time.Time) *domain.MonitorQuotaSnapshot {
|
||||
// fetchCNQuota 国产 coding plan:CNProviderQuotaService.QueryUsageForAccount → 快照。
|
||||
func (f *ChannelMonitorQuotaFetcher) fetchCNQuota(ctx context.Context, account *Account, now time.Time) *domain.MonitorQuotaSnapshot {
|
||||
if f.cnQuota == nil {
|
||||
return quotaErrorSnapshot("cn_quota", "cn quota service is not configured", now)
|
||||
}
|
||||
result, err := f.cnQuota.QueryUsage(ctx, accountID)
|
||||
result, err := f.cnQuota.QueryUsageForAccount(ctx, account)
|
||||
if err != nil {
|
||||
msg := truncateMessage(sanitizeErrorMessage(err.Error()))
|
||||
return &domain.MonitorQuotaSnapshot{
|
||||
@@ -381,12 +387,12 @@ func (f *ChannelMonitorQuotaFetcher) fetchCNQuota(ctx context.Context, accountID
|
||||
return snapshot
|
||||
}
|
||||
|
||||
// fetchCNBalance 国产 payg:CNProviderBalanceService.QueryBalance → 快照。
|
||||
func (f *ChannelMonitorQuotaFetcher) fetchCNBalance(ctx context.Context, accountID int64, now time.Time) *domain.MonitorQuotaSnapshot {
|
||||
// fetchCNBalance 国产 payg:CNProviderBalanceService.QueryBalanceForAccount → 快照。
|
||||
func (f *ChannelMonitorQuotaFetcher) fetchCNBalance(ctx context.Context, account *Account, now time.Time) *domain.MonitorQuotaSnapshot {
|
||||
if f.cnBalance == nil {
|
||||
return quotaErrorSnapshot("cn_balance", "cn balance service is not configured", now)
|
||||
}
|
||||
result, err := f.cnBalance.QueryBalance(ctx, accountID)
|
||||
result, err := f.cnBalance.QueryBalanceForAccount(ctx, account)
|
||||
if err != nil {
|
||||
msg := truncateMessage(sanitizeErrorMessage(err.Error()))
|
||||
return &domain.MonitorQuotaSnapshot{
|
||||
|
||||
@@ -20,18 +20,20 @@ import (
|
||||
type stubMonitorUsageSource struct {
|
||||
usage *UsageInfo
|
||||
err error
|
||||
// block 非 nil 时 GetUsage 阻塞在该 channel 上,用于并发/singleflight 测试。
|
||||
// block 非 nil 时 GetUsageForAccount 阻塞在该 channel 上,用于并发/singleflight 测试。
|
||||
block chan struct{}
|
||||
|
||||
mu sync.Mutex
|
||||
calls int
|
||||
lastCtx context.Context
|
||||
mu sync.Mutex
|
||||
calls int
|
||||
lastCtx context.Context
|
||||
lastAccount *Account
|
||||
}
|
||||
|
||||
func (s *stubMonitorUsageSource) GetUsage(ctx context.Context, accountID int64, force ...bool) (*UsageInfo, error) {
|
||||
func (s *stubMonitorUsageSource) GetUsageForAccount(ctx context.Context, account *Account, force ...bool) (*UsageInfo, error) {
|
||||
s.mu.Lock()
|
||||
s.calls++
|
||||
s.lastCtx = ctx
|
||||
s.lastAccount = account
|
||||
s.mu.Unlock()
|
||||
if s.block != nil {
|
||||
<-s.block
|
||||
@@ -45,25 +47,35 @@ func (s *stubMonitorUsageSource) getCalls() int {
|
||||
return s.calls
|
||||
}
|
||||
|
||||
type stubMonitorCNQuotaSource struct {
|
||||
result *CNProviderQuotaProbeResult
|
||||
err error
|
||||
calls int
|
||||
func (s *stubMonitorUsageSource) getLastAccount() *Account {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.lastAccount
|
||||
}
|
||||
|
||||
func (s *stubMonitorCNQuotaSource) QueryUsage(ctx context.Context, accountID int64) (*CNProviderQuotaProbeResult, error) {
|
||||
type stubMonitorCNQuotaSource struct {
|
||||
result *CNProviderQuotaProbeResult
|
||||
err error
|
||||
calls int
|
||||
lastAccount *Account
|
||||
}
|
||||
|
||||
func (s *stubMonitorCNQuotaSource) QueryUsageForAccount(ctx context.Context, account *Account) (*CNProviderQuotaProbeResult, error) {
|
||||
s.calls++
|
||||
s.lastAccount = account
|
||||
return s.result, s.err
|
||||
}
|
||||
|
||||
type stubMonitorCNBalanceSource struct {
|
||||
result *CNProviderBalanceResult
|
||||
err error
|
||||
calls int
|
||||
result *CNProviderBalanceResult
|
||||
err error
|
||||
calls int
|
||||
lastAccount *Account
|
||||
}
|
||||
|
||||
func (s *stubMonitorCNBalanceSource) QueryBalance(ctx context.Context, accountID int64) (*CNProviderBalanceResult, error) {
|
||||
func (s *stubMonitorCNBalanceSource) QueryBalanceForAccount(ctx context.Context, account *Account) (*CNProviderBalanceResult, error) {
|
||||
s.calls++
|
||||
s.lastAccount = account
|
||||
return s.result, s.err
|
||||
}
|
||||
|
||||
@@ -192,6 +204,49 @@ func TestQuotaFetcher_PayGAccountUsesCNBalance(t *testing.T) {
|
||||
require.Empty(t, snapshot.Error)
|
||||
}
|
||||
|
||||
// P2-6:fetchUncached 只 GetByID 一次,已加载的 account 指针直传数据源,
|
||||
// 三条路由都不能让下游重载账号。
|
||||
func TestQuotaFetcher_LoadsAccountOnceAndPassesItThrough(t *testing.T) {
|
||||
t.Run("overseas usage", func(t *testing.T) {
|
||||
fetcher, usage, _, _, accounts := newQuotaFetcherTestSetup(t)
|
||||
acc := &Account{ID: 21, Platform: domain.PlatformAnthropic}
|
||||
accounts.accounts[21] = acc
|
||||
usage.usage = &UsageInfo{}
|
||||
|
||||
fetcher.Fetch(context.Background(), 21)
|
||||
|
||||
require.Equal(t, 1, accounts.calls)
|
||||
require.Same(t, acc, usage.getLastAccount())
|
||||
require.Equal(t, 1, usage.getCalls())
|
||||
})
|
||||
|
||||
t.Run("cn coding plan", func(t *testing.T) {
|
||||
fetcher, _, cnQuota, _, accounts := newQuotaFetcherTestSetup(t)
|
||||
acc := &Account{ID: 22, Platform: domain.PlatformKimi, Credentials: map[string]any{"account_mode": AccountModeCoding}}
|
||||
accounts.accounts[22] = acc
|
||||
cnQuota.result = &CNProviderQuotaProbeResult{Success: true}
|
||||
|
||||
fetcher.Fetch(context.Background(), 22)
|
||||
|
||||
require.Equal(t, 1, accounts.calls)
|
||||
require.Same(t, acc, cnQuota.lastAccount)
|
||||
require.Equal(t, 1, cnQuota.calls)
|
||||
})
|
||||
|
||||
t.Run("cn payg", func(t *testing.T) {
|
||||
fetcher, _, _, cnBalance, accounts := newQuotaFetcherTestSetup(t)
|
||||
acc := &Account{ID: 23, Platform: domain.PlatformDeepseek, Credentials: map[string]any{"account_mode": AccountModePayG}}
|
||||
accounts.accounts[23] = acc
|
||||
cnBalance.result = &CNProviderBalanceResult{Success: true, Available: true, Balance: 1, Currency: "CNY"}
|
||||
|
||||
fetcher.Fetch(context.Background(), 23)
|
||||
|
||||
require.Equal(t, 1, accounts.calls)
|
||||
require.Same(t, acc, cnBalance.lastAccount)
|
||||
require.Equal(t, 1, cnBalance.calls)
|
||||
})
|
||||
}
|
||||
|
||||
// --- 失败路径(Fetch 永不返回 error) ---
|
||||
|
||||
func TestQuotaFetcher_AccountMissingYieldsLinkedAccountSnapshot(t *testing.T) {
|
||||
|
||||
@@ -84,14 +84,27 @@ func NewCNProviderBalanceService(
|
||||
|
||||
// QueryBalance 探测指定 payg 账号的余额并落 Extra 快照。
|
||||
func (s *CNProviderBalanceService) QueryBalance(ctx context.Context, accountID int64) (*CNProviderBalanceResult, error) {
|
||||
account, err := s.loadPayGAccount(ctx, accountID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return s.QueryBalanceForAccount(ctx, account)
|
||||
}
|
||||
|
||||
// QueryBalanceForAccount 探测已加载账号(配额监控 fetcher / 周期余额检测复用,
|
||||
// 避免二次 GetByID)。singleflight key 与 QueryBalance 相同,按账号 ID 合并。
|
||||
func (s *CNProviderBalanceService) QueryBalanceForAccount(ctx context.Context, account *Account) (*CNProviderBalanceResult, error) {
|
||||
if s == nil || s.accountRepo == nil || s.httpUpstream == nil {
|
||||
return nil, infraerrors.New(http.StatusInternalServerError, "CN_BALANCE_NOT_CONFIGURED", "cn provider balance service is not configured")
|
||||
}
|
||||
key := "cn_balance:" + strconv.FormatInt(accountID, 10)
|
||||
if err := validatePayGAccount(account); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
key := "cn_balance:" + strconv.FormatInt(account.ID, 10)
|
||||
resultCh := s.flight.DoChan(key, func() (any, error) {
|
||||
probeCtx, cancel := context.WithTimeout(context.Background(), cnBalanceUpstreamTimeout+5*time.Second)
|
||||
defer cancel()
|
||||
return s.queryBalance(probeCtx, accountID)
|
||||
return s.queryBalanceForAccount(probeCtx, account)
|
||||
})
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
@@ -109,11 +122,7 @@ func (s *CNProviderBalanceService) QueryBalance(ctx context.Context, accountID i
|
||||
}
|
||||
}
|
||||
|
||||
func (s *CNProviderBalanceService) queryBalance(ctx context.Context, accountID int64) (*CNProviderBalanceResult, error) {
|
||||
account, err := s.loadPayGAccount(ctx, accountID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
func (s *CNProviderBalanceService) queryBalanceForAccount(ctx context.Context, account *Account) (*CNProviderBalanceResult, error) {
|
||||
provider := account.Platform
|
||||
if provider != PlatformKimi && provider != PlatformDeepseek {
|
||||
return nil, infraerrors.New(http.StatusBadRequest, "CN_BALANCE_NO_ENDPOINT", "account provider has no balance endpoint")
|
||||
@@ -230,17 +239,26 @@ func (s *CNProviderBalanceService) loadPayGAccount(ctx context.Context, accountI
|
||||
if err != nil {
|
||||
return nil, infraerrors.Newf(http.StatusNotFound, "CN_BALANCE_ACCOUNT_NOT_FOUND", "account not found: %v", err)
|
||||
}
|
||||
if err := validatePayGAccount(account); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return account, nil
|
||||
}
|
||||
|
||||
// validatePayGAccount 加载后的非 DB 校验(ForAccount 入口同样复用,
|
||||
// 保证直传 account 也不绕过平台/模式检查)。
|
||||
func validatePayGAccount(account *Account) error {
|
||||
if account == nil {
|
||||
return nil, infraerrors.New(http.StatusNotFound, "CN_BALANCE_ACCOUNT_NOT_FOUND", "account not found")
|
||||
return infraerrors.New(http.StatusNotFound, "CN_BALANCE_ACCOUNT_NOT_FOUND", "account not found")
|
||||
}
|
||||
if !account.IsCNProvider() {
|
||||
return nil, infraerrors.New(http.StatusBadRequest, "CN_BALANCE_INVALID_PLATFORM", "account is not a CN provider account")
|
||||
return infraerrors.New(http.StatusBadRequest, "CN_BALANCE_INVALID_PLATFORM", "account is not a CN provider account")
|
||||
}
|
||||
// coding 账号走额度探测,余额端点不适用。
|
||||
if account.IsCodingPlan() {
|
||||
return nil, infraerrors.New(http.StatusBadRequest, "CN_BALANCE_CODING_PLAN", "coding plan account has no balance endpoint; use quota probe")
|
||||
return infraerrors.New(http.StatusBadRequest, "CN_BALANCE_CODING_PLAN", "coding plan account has no balance endpoint; use quota probe")
|
||||
}
|
||||
return account, nil
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *CNProviderBalanceService) resolveProxyURL(ctx context.Context, account *Account) string {
|
||||
|
||||
@@ -0,0 +1,125 @@
|
||||
package service
|
||||
|
||||
// ForAccount 直传入口(P2-6)的校验回归测试:
|
||||
// QueryUsageForAccount / QueryBalanceForAccount 接受已加载的 *Account,
|
||||
// 但必须复用与 ID 入口相同的加载后校验——直传不能绕过平台/模式检查,
|
||||
// 且校验在 singleflight 之前完成(无效账号不得发起任何上游请求)。
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func codingAccount(platform string) *Account {
|
||||
return &Account{
|
||||
ID: 1, Platform: platform, Type: AccountTypeAPIKey, Status: StatusActive,
|
||||
Credentials: map[string]any{"account_mode": AccountModeCoding, "api_key": "sk-test"},
|
||||
}
|
||||
}
|
||||
|
||||
func paygAccount(platform string) *Account {
|
||||
return &Account{
|
||||
ID: 2, Platform: platform, Type: AccountTypeAPIKey, Status: StatusActive,
|
||||
Credentials: map[string]any{"account_mode": AccountModePayG, "api_key": "sk-test"},
|
||||
}
|
||||
}
|
||||
|
||||
func requireReason(t *testing.T, err error, reason string) {
|
||||
t.Helper()
|
||||
require.Error(t, err)
|
||||
var appErr *infraerrors.ApplicationError
|
||||
require.ErrorAs(t, err, &appErr)
|
||||
require.Equal(t, reason, appErr.Reason)
|
||||
}
|
||||
|
||||
func TestValidateCodingPlanAccount_Matrix(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
account *Account
|
||||
wantReason string
|
||||
}{
|
||||
{name: "nil", account: nil, wantReason: "CN_QUOTA_ACCOUNT_NOT_FOUND"},
|
||||
{name: "non cn provider", account: &Account{ID: 3, Platform: PlatformAnthropic}, wantReason: "CN_QUOTA_INVALID_PLATFORM"},
|
||||
{name: "payg has no quota endpoint", account: paygAccount(PlatformKimi), wantReason: "CN_QUOTA_NOT_CODING_PLAN"},
|
||||
{name: "kimi coding ok", account: codingAccount(PlatformKimi)},
|
||||
{name: "zhipu coding ok", account: codingAccount(PlatformZhipu)},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
err := validateCodingPlanAccount(tc.account)
|
||||
if tc.wantReason == "" {
|
||||
require.NoError(t, err)
|
||||
return
|
||||
}
|
||||
requireReason(t, err, tc.wantReason)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidatePayGAccount_Matrix(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
account *Account
|
||||
wantReason string
|
||||
}{
|
||||
{name: "nil", account: nil, wantReason: "CN_BALANCE_ACCOUNT_NOT_FOUND"},
|
||||
{name: "non cn provider", account: &Account{ID: 3, Platform: PlatformAnthropic}, wantReason: "CN_BALANCE_INVALID_PLATFORM"},
|
||||
{name: "coding has no balance endpoint", account: codingAccount(PlatformKimi), wantReason: "CN_BALANCE_CODING_PLAN"},
|
||||
{name: "kimi payg ok", account: paygAccount(PlatformKimi)},
|
||||
{name: "deepseek payg ok", account: paygAccount(PlatformDeepseek)},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
err := validatePayGAccount(tc.account)
|
||||
if tc.wantReason == "" {
|
||||
require.NoError(t, err)
|
||||
return
|
||||
}
|
||||
requireReason(t, err, tc.wantReason)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// 直传入口的校验在 singleflight/上游请求之前:无效账号必须零出站请求。
|
||||
func TestCNProviderQuotaService_QueryUsageForAccount_RejectsInvalidAccount(t *testing.T) {
|
||||
repo := &fakeCNProbeAccountRepo{}
|
||||
upstream := &recordingHTTPUpstream{}
|
||||
svc := NewCNProviderQuotaService(repo, nil, upstream, nil)
|
||||
|
||||
_, err := svc.QueryUsageForAccount(context.Background(), paygAccount(PlatformKimi))
|
||||
requireReason(t, err, "CN_QUOTA_NOT_CODING_PLAN")
|
||||
require.Zero(t, upstream.calls)
|
||||
|
||||
_, err = svc.QueryUsageForAccount(context.Background(), nil)
|
||||
requireReason(t, err, "CN_QUOTA_ACCOUNT_NOT_FOUND")
|
||||
require.Zero(t, upstream.calls)
|
||||
}
|
||||
|
||||
func TestCNProviderBalanceService_QueryBalanceForAccount_RejectsInvalidAccount(t *testing.T) {
|
||||
repo := &fakeCNProbeAccountRepo{}
|
||||
upstream := &recordingHTTPUpstream{}
|
||||
svc := NewCNProviderBalanceService(repo, nil, upstream, nil)
|
||||
|
||||
_, err := svc.QueryBalanceForAccount(context.Background(), codingAccount(PlatformKimi))
|
||||
requireReason(t, err, "CN_BALANCE_CODING_PLAN")
|
||||
require.Zero(t, upstream.calls)
|
||||
|
||||
_, err = svc.QueryBalanceForAccount(context.Background(), &Account{ID: 9, Platform: PlatformAnthropic})
|
||||
requireReason(t, err, "CN_BALANCE_INVALID_PLATFORM")
|
||||
require.Zero(t, upstream.calls)
|
||||
}
|
||||
|
||||
// ID 入口与 ForAccount 入口对同一账号的行为一致(loadCodingPlanAccount 的
|
||||
// 加载后校验 = validateCodingPlanAccount;余额侧对称)。
|
||||
func TestCNProviderServices_IDEntryAppliesSameValidation(t *testing.T) {
|
||||
repo := &fakeCNProbeAccountRepo{account: paygAccount(PlatformKimi)}
|
||||
upstream := &recordingHTTPUpstream{}
|
||||
svc := NewCNProviderQuotaService(repo, nil, upstream, nil)
|
||||
|
||||
_, err := svc.QueryUsage(context.Background(), 2)
|
||||
requireReason(t, err, "CN_QUOTA_NOT_CODING_PLAN")
|
||||
require.Zero(t, upstream.calls)
|
||||
}
|
||||
@@ -89,14 +89,27 @@ func NewCNProviderQuotaService(
|
||||
// QueryUsage 探测指定账号的 Coding Plan 滚动窗口用量并落 Extra 快照。
|
||||
// 同一账号的并发探测会被 singleflight 合并。
|
||||
func (s *CNProviderQuotaService) QueryUsage(ctx context.Context, accountID int64) (*CNProviderQuotaProbeResult, error) {
|
||||
account, err := s.loadCodingPlanAccount(ctx, accountID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return s.QueryUsageForAccount(ctx, account)
|
||||
}
|
||||
|
||||
// QueryUsageForAccount 探测已加载账号(配额监控 fetcher 复用,避免二次 GetByID)。
|
||||
// singleflight key 与 QueryUsage 相同,按账号 ID 与 admin 侧并发探测合并。
|
||||
func (s *CNProviderQuotaService) QueryUsageForAccount(ctx context.Context, account *Account) (*CNProviderQuotaProbeResult, error) {
|
||||
if s == nil || s.accountRepo == nil || s.httpUpstream == nil {
|
||||
return nil, infraerrors.New(http.StatusInternalServerError, "CN_QUOTA_NOT_CONFIGURED", "cn provider quota service is not configured")
|
||||
}
|
||||
key := "cn_quota:" + strconv.FormatInt(accountID, 10)
|
||||
if err := validateCodingPlanAccount(account); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
key := "cn_quota:" + strconv.FormatInt(account.ID, 10)
|
||||
resultCh := s.flight.DoChan(key, func() (any, error) {
|
||||
probeCtx, cancel := context.WithTimeout(context.Background(), cnQuotaUpstreamTimeout+5*time.Second)
|
||||
defer cancel()
|
||||
return s.queryUsage(probeCtx, accountID)
|
||||
return s.queryUsageForAccount(probeCtx, account)
|
||||
})
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
@@ -114,12 +127,7 @@ func (s *CNProviderQuotaService) QueryUsage(ctx context.Context, accountID int64
|
||||
}
|
||||
}
|
||||
|
||||
func (s *CNProviderQuotaService) queryUsage(ctx context.Context, accountID int64) (*CNProviderQuotaProbeResult, error) {
|
||||
account, err := s.loadCodingPlanAccount(ctx, accountID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
func (s *CNProviderQuotaService) queryUsageForAccount(ctx context.Context, account *Account) (*CNProviderQuotaProbeResult, error) {
|
||||
provider := account.GetCodingPlanProvider()
|
||||
if provider != PlatformKimi && provider != PlatformZhipu {
|
||||
return nil, infraerrors.New(http.StatusBadRequest, "CN_QUOTA_NOT_CODING_PLAN", "account is not a kimi/zhipu coding plan account")
|
||||
@@ -231,18 +239,27 @@ func (s *CNProviderQuotaService) loadCodingPlanAccount(ctx context.Context, acco
|
||||
if err != nil {
|
||||
return nil, infraerrors.Newf(http.StatusNotFound, "CN_QUOTA_ACCOUNT_NOT_FOUND", "account not found: %v", err)
|
||||
}
|
||||
if account == nil {
|
||||
return nil, infraerrors.New(http.StatusNotFound, "CN_QUOTA_ACCOUNT_NOT_FOUND", "account not found")
|
||||
}
|
||||
if !account.IsCNProvider() {
|
||||
return nil, infraerrors.New(http.StatusBadRequest, "CN_QUOTA_INVALID_PLATFORM", "account is not a CN provider account")
|
||||
}
|
||||
if !account.IsCodingPlan() {
|
||||
return nil, infraerrors.New(http.StatusBadRequest, "CN_QUOTA_NOT_CODING_PLAN", "account is not a coding plan account")
|
||||
if err := validateCodingPlanAccount(account); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return account, nil
|
||||
}
|
||||
|
||||
// validateCodingPlanAccount 加载后的非 DB 校验(ForAccount 入口同样复用,
|
||||
// 保证直传 account 也不绕过平台/模式检查)。
|
||||
func validateCodingPlanAccount(account *Account) error {
|
||||
if account == nil {
|
||||
return infraerrors.New(http.StatusNotFound, "CN_QUOTA_ACCOUNT_NOT_FOUND", "account not found")
|
||||
}
|
||||
if !account.IsCNProvider() {
|
||||
return infraerrors.New(http.StatusBadRequest, "CN_QUOTA_INVALID_PLATFORM", "account is not a CN provider account")
|
||||
}
|
||||
if !account.IsCodingPlan() {
|
||||
return infraerrors.New(http.StatusBadRequest, "CN_QUOTA_NOT_CODING_PLAN", "account is not a coding plan account")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *CNProviderQuotaService) resolveProxyURL(ctx context.Context, account *Account) string {
|
||||
if account == nil || account.ProxyID == nil {
|
||||
return ""
|
||||
|
||||
Reference in New Issue
Block a user