mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 12:02:27 +08:00
Merge pull request #4896 from Brisbanehuang/feat/upstream-billing-probe-multi-platform
feat(billing-probe): 上游计费倍率探测放宽到全部 API-key 平台账号
This commit is contained in:
@@ -625,7 +625,7 @@ func lockAndMergeAccountProbeExtra(ctx context.Context, client *dbent.Client, ac
|
||||
delete(extra, key)
|
||||
}
|
||||
probeExplicitlyDisabled := false
|
||||
probeAccount := account.Platform == service.PlatformOpenAI && account.Type == service.AccountTypeAPIKey
|
||||
probeAccount := service.IsUpstreamBillingProbeIdentity(account.Platform, account.Type)
|
||||
if probeAccount && explicitProbeEnabled != nil {
|
||||
extra[service.UpstreamBillingProbeEnabledExtraKey] = *explicitProbeEnabled
|
||||
probeExplicitlyDisabled = !*explicitProbeEnabled
|
||||
@@ -709,7 +709,7 @@ func (r *accountRepository) UpdateCredentials(ctx context.Context, id int64, cre
|
||||
credentials = $1::jsonb,
|
||||
extra = CASE
|
||||
-- 凭证整体未变化 ⇒ Ollama 组身份必然未变化;顶层 DISTINCT 守卫防止
|
||||
-- 非 Ollama 账号的无变化持久化误清 openai 探测快照或重写 NULL extra。
|
||||
-- 非 Ollama 账号的无变化持久化误清探测快照或重写 NULL extra。
|
||||
WHEN platform IN ('openai', 'anthropic')
|
||||
AND type = 'apikey'
|
||||
AND credentials IS DISTINCT FROM $1::jsonb
|
||||
@@ -720,15 +720,14 @@ func (r *accountRepository) UpdateCredentials(ctx context.Context, id int64, cre
|
||||
AND `+ollamaCloudBaseURLMatchesSQL("$1::jsonb ->> 'base_url'")+`
|
||||
)
|
||||
)
|
||||
THEN (CASE
|
||||
WHEN platform = 'openai' THEN COALESCE(extra, '{}'::jsonb) - 'upstream_billing_probe'
|
||||
ELSE COALESCE(extra, '{}'::jsonb)
|
||||
END)
|
||||
THEN COALESCE(extra, '{}'::jsonb)
|
||||
- 'upstream_billing_probe'
|
||||
- 'ollama_cloud_usage_session'
|
||||
- 'ollama_cloud_usage_auto_refresh'
|
||||
- 'ollama_cloud_usage_snapshot'
|
||||
WHEN platform = 'openai'
|
||||
AND type = 'apikey'
|
||||
-- 上游倍率探测已放宽到全部 API-key 平台:凭证变化即视为探测
|
||||
-- 身份变化,丢弃 stale 快照。
|
||||
WHEN type = 'apikey'
|
||||
AND credentials IS DISTINCT FROM $1::jsonb
|
||||
THEN COALESCE(extra, '{}'::jsonb) - 'upstream_billing_probe'
|
||||
ELSE extra
|
||||
@@ -2839,8 +2838,8 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates
|
||||
args = append(args, pq.Array(ids))
|
||||
idx++
|
||||
if updates.ProbeEnabled != nil {
|
||||
whereClause += " AND platform = $" + itoa(idx) + " AND type = $" + itoa(idx+1)
|
||||
args = append(args, service.PlatformOpenAI, service.AccountTypeAPIKey)
|
||||
whereClause += " AND type = $" + itoa(idx)
|
||||
args = append(args, service.AccountTypeAPIKey)
|
||||
}
|
||||
query := "UPDATE accounts SET " + joinClauses(setClauses, ", ") + whereClause
|
||||
|
||||
@@ -3382,7 +3381,6 @@ func (r *accountRepository) ListDueUpstreamBillingProbeAccounts(ctx context.Cont
|
||||
FROM accounts
|
||||
WHERE deleted_at IS NULL
|
||||
AND status = 'active'
|
||||
AND platform = 'openai'
|
||||
AND type = 'apikey'
|
||||
AND extra @> '{"upstream_billing_probe_enabled": true}'::jsonb
|
||||
), parsed AS MATERIALIZED (
|
||||
|
||||
+51
@@ -153,3 +153,54 @@ func TestListDueUpstreamBillingProbeAccountsSelectsEarliestDueAcrossIDs(t *testi
|
||||
}
|
||||
require.Equal(t, want, got)
|
||||
}
|
||||
|
||||
// 探测资格放宽回归:任何 API-key 平台的启用账号都进入定时探测候选并按
|
||||
// 到期时间排序;OAuth 与未启用账号仍被排除。
|
||||
func TestListDueUpstreamBillingProbeAccountsIncludesAllAPIKeyPlatforms(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
|
||||
now := time.Date(2026, time.July, 26, 3, 0, 0, 0, time.UTC)
|
||||
_, err := tx.ExecContext(ctx, `
|
||||
UPDATE accounts
|
||||
SET extra = extra - 'upstream_billing_probe_enabled' - 'upstream_billing_probe'
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
|
||||
insert := func(name, platform, accountType, nextProbeAt string) int64 {
|
||||
t.Helper()
|
||||
var id int64
|
||||
extra := fmt.Sprintf(`{
|
||||
"upstream_billing_probe_enabled": true,
|
||||
"upstream_billing_probe": {"status": "ok", "next_probe_at": %q}
|
||||
}`, nextProbeAt)
|
||||
err := scanSingleRow(ctx, tx, `
|
||||
INSERT INTO accounts (name, platform, type, status, extra)
|
||||
VALUES ($1, $2, $3, 'active', $4::jsonb)
|
||||
RETURNING id
|
||||
`, []any{name, platform, accountType, extra}, &id)
|
||||
require.NoError(t, err)
|
||||
return id
|
||||
}
|
||||
|
||||
openaiDue := insert("probe-openai-due", "openai", service.AccountTypeAPIKey, "2026-07-26T02:57:00Z")
|
||||
anthropicDue := insert("probe-anthropic-due", "anthropic", service.AccountTypeAPIKey, "2026-07-26T02:58:00Z")
|
||||
grokDue := insert("probe-grok-due", "grok", service.AccountTypeAPIKey, "2026-07-26T02:59:00Z")
|
||||
// OAuth 账号即便误持有启用标记也不得入选。
|
||||
_ = insert("probe-grok-oauth-excluded", "grok", service.AccountTypeOAuth, "2026-07-26T02:59:00Z")
|
||||
// 未启用探测的 API-key 账号不入选。
|
||||
var disabledID int64
|
||||
err = scanSingleRow(ctx, tx, `
|
||||
INSERT INTO accounts (name, platform, type, status, extra)
|
||||
VALUES ('probe-grok-disabled', 'grok', $1, 'active', '{}'::jsonb)
|
||||
RETURNING id
|
||||
`, []any{service.AccountTypeAPIKey}, &disabledID)
|
||||
require.NoError(t, err)
|
||||
|
||||
accounts, err := repo.ListDueUpstreamBillingProbeAccounts(ctx, now, 20)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, accounts, 3)
|
||||
require.Equal(t, openaiDue, accounts[0].ID)
|
||||
require.Equal(t, anthropicDue, accounts[1].ID)
|
||||
require.Equal(t, grokDue, accounts[2].ID)
|
||||
}
|
||||
|
||||
@@ -28,7 +28,8 @@ func TestAccountRepositoryListDueUpstreamBillingProbeAccountsBoundsQuery(t *test
|
||||
normalized := normalizeSQLWhitespace(capturedSQL)
|
||||
require.Contains(t, normalized, "deleted_at IS NULL")
|
||||
require.Contains(t, normalized, "status = 'active'")
|
||||
require.Contains(t, normalized, "platform = 'openai'")
|
||||
// 探测资格已放宽到全部 API-key 平台:候选 SQL 不得再按 platform 过滤。
|
||||
require.NotContains(t, normalized, "platform")
|
||||
require.Contains(t, normalized, "type = 'apikey'")
|
||||
require.Contains(t, normalized, `extra @> '{"upstream_billing_probe_enabled": true}'::jsonb`)
|
||||
require.Contains(t, normalized, "jsonb_path_query_first_tz")
|
||||
|
||||
@@ -227,8 +227,8 @@ func TestBulkUpdateProbeEligibilityMismatchRollsBack(t *testing.T) {
|
||||
|
||||
enabled := true
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec(`(?s)UPDATE accounts SET extra = .* WHERE id = ANY\(\$2\) AND deleted_at IS NULL AND platform = \$3 AND type = \$4`).
|
||||
WithArgs(sqlmock.AnyArg(), `{27,28}`, service.PlatformOpenAI, service.AccountTypeAPIKey).
|
||||
mock.ExpectExec(`(?s)UPDATE accounts SET extra = .* WHERE id = ANY\(\$2\) AND deleted_at IS NULL AND type = \$3`).
|
||||
WithArgs(sqlmock.AnyArg(), `{27,28}`, service.AccountTypeAPIKey).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectRollback()
|
||||
|
||||
|
||||
@@ -231,8 +231,7 @@ func invalidateProxyProbeSnapshots(ctx context.Context, exec sqlExecutor, proxyI
|
||||
WHERE proxy_id = $1
|
||||
AND type = 'apikey'
|
||||
AND (
|
||||
(platform = 'openai'
|
||||
AND extra ? 'upstream_billing_probe'
|
||||
(extra ? 'upstream_billing_probe'
|
||||
AND extra -> 'upstream_billing_probe' <> 'null'::jsonb)
|
||||
OR (platform IN ('openai', 'anthropic')
|
||||
AND extra ? 'ollama_cloud_usage_snapshot'
|
||||
@@ -747,7 +746,7 @@ func (r *proxyRepository) sweepOneExpiredProxyOnExec(ctx context.Context, exec s
|
||||
rows, err = exec.QueryContext(ctx, `
|
||||
UPDATE accounts SET proxy_id=NULL, proxy_fallback_origin_id=$1,
|
||||
extra=CASE
|
||||
WHEN platform='openai' AND type='apikey' AND extra ? 'upstream_billing_probe'
|
||||
WHEN type='apikey' AND extra ? 'upstream_billing_probe'
|
||||
THEN extra - 'upstream_billing_probe'
|
||||
ELSE extra
|
||||
END,
|
||||
@@ -758,7 +757,7 @@ func (r *proxyRepository) sweepOneExpiredProxyOnExec(ctx context.Context, exec s
|
||||
rows, err = exec.QueryContext(ctx, `
|
||||
UPDATE accounts SET proxy_id=$2, proxy_fallback_origin_id=$1,
|
||||
extra=CASE
|
||||
WHEN platform='openai' AND type='apikey' AND extra ? 'upstream_billing_probe'
|
||||
WHEN type='apikey' AND extra ? 'upstream_billing_probe'
|
||||
THEN extra - 'upstream_billing_probe'
|
||||
ELSE extra
|
||||
END,
|
||||
|
||||
@@ -33,7 +33,7 @@ func TestProxyUpdateInvalidatesBoundProbeSnapshotsAndEnqueuesOutboxAtomically(t
|
||||
WithArgs(int64(9)).
|
||||
WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
expectProxyUpdateReload(mock, 9, "new.example", "user", "pass")
|
||||
mock.ExpectQuery(`(?s)UPDATE accounts.*- 'upstream_billing_probe'.*- 'ollama_cloud_usage_snapshot'.*type = 'apikey'.*platform = 'openai'.*extra \? 'upstream_billing_probe'.*platform IN \('openai', 'anthropic'\).*extra \? 'ollama_cloud_usage_snapshot'.*RETURNING id`).
|
||||
mock.ExpectQuery(`(?s)UPDATE accounts.*- 'upstream_billing_probe'.*- 'ollama_cloud_usage_snapshot'.*type = 'apikey'.*extra \? 'upstream_billing_probe'.*platform IN \('openai', 'anthropic'\).*extra \? 'ollama_cloud_usage_snapshot'.*RETURNING id`).
|
||||
WithArgs(int64(9)).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(int64(17)).AddRow(int64(18)))
|
||||
mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)")).
|
||||
@@ -76,7 +76,7 @@ func TestProxyUpdateRollsBackWhenProbeInvalidationOutboxFails(t *testing.T) {
|
||||
WithArgs(int64(9)).
|
||||
WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
expectProxyUpdateReload(mock, 9, "new.example", "", "")
|
||||
mock.ExpectQuery(`(?s)UPDATE accounts.*- 'upstream_billing_probe'.*- 'ollama_cloud_usage_snapshot'.*type = 'apikey'.*platform = 'openai'.*extra \? 'upstream_billing_probe'.*platform IN \('openai', 'anthropic'\).*extra \? 'ollama_cloud_usage_snapshot'.*RETURNING id`).
|
||||
mock.ExpectQuery(`(?s)UPDATE accounts.*- 'upstream_billing_probe'.*- 'ollama_cloud_usage_snapshot'.*type = 'apikey'.*extra \? 'upstream_billing_probe'.*platform IN \('openai', 'anthropic'\).*extra \? 'ollama_cloud_usage_snapshot'.*RETURNING id`).
|
||||
WithArgs(int64(9)).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(int64(17)))
|
||||
mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)")).
|
||||
|
||||
@@ -151,16 +151,31 @@ func TestReconcileCRSUpstreamBillingProbeExtra(t *testing.T) {
|
||||
require.NotContains(t, extra, UpstreamBillingProbeExtraKey)
|
||||
})
|
||||
|
||||
// API-key 平台间切换:探测资格保留(放宽后不再限 OpenAI),开关沿用本地值,
|
||||
// 但平台属于探测身份,快照必须作废。
|
||||
for _, target := range []struct {
|
||||
name string
|
||||
platform string
|
||||
}{
|
||||
{name: "anthropic api key", platform: PlatformAnthropic},
|
||||
{name: "gemini api key", platform: PlatformGemini},
|
||||
} {
|
||||
t.Run(target.name+" keeps enabled and clears snapshot", func(t *testing.T) {
|
||||
extra := mergeMap(existing.Extra, remote)
|
||||
reconcileCRSUpstreamBillingProbeExtra(existing, target.platform, AccountTypeAPIKey, existing.Credentials, extra)
|
||||
require.Equal(t, false, extra[UpstreamBillingProbeEnabledExtraKey])
|
||||
require.NotContains(t, extra, UpstreamBillingProbeExtraKey)
|
||||
})
|
||||
}
|
||||
|
||||
for _, target := range []struct {
|
||||
name string
|
||||
platform string
|
||||
typeName string
|
||||
}{
|
||||
{name: "anthropic oauth", platform: PlatformAnthropic, typeName: AccountTypeOAuth},
|
||||
{name: "anthropic api key", platform: PlatformAnthropic, typeName: AccountTypeAPIKey},
|
||||
{name: "openai oauth", platform: PlatformOpenAI, typeName: AccountTypeOAuth},
|
||||
{name: "gemini oauth", platform: PlatformGemini, typeName: AccountTypeOAuth},
|
||||
{name: "gemini api key", platform: PlatformGemini, typeName: AccountTypeAPIKey},
|
||||
} {
|
||||
t.Run(target.name+" removes inapplicable state", func(t *testing.T) {
|
||||
extra := mergeMap(existing.Extra, remote)
|
||||
|
||||
@@ -1173,7 +1173,7 @@ func reconcileCRSUpstreamBillingProbeExtra(
|
||||
return
|
||||
}
|
||||
target := &Account{Platform: targetPlatform, Type: targetType, Credentials: targetCredentials}
|
||||
if targetPlatform == PlatformOpenAI && targetType == AccountTypeAPIKey {
|
||||
if IsUpstreamBillingProbeIdentity(targetPlatform, targetType) {
|
||||
if enabled, ok := existing.Extra[UpstreamBillingProbeEnabledExtraKey]; ok {
|
||||
extra[UpstreamBillingProbeEnabledExtraKey] = enabled
|
||||
}
|
||||
|
||||
@@ -2653,6 +2653,15 @@ func newOpenAILegacyUpstreamRateOrder(accounts []*Account, now time.Time, oauthS
|
||||
var first float64
|
||||
distinct := false
|
||||
for _, account := range accounts {
|
||||
if account == nil {
|
||||
continue
|
||||
}
|
||||
// 与 openAIUpstreamCostFactors 使用同一道平台门控:只有 OpenAI 平台账号
|
||||
// 的倍率参与 legacy 低倍率优先排序。上游自报倍率来自中转方,不能让它对
|
||||
// 其他平台的调度产生影响——否则自报低价即可吸走流量,而实际结算走本地倍率。
|
||||
if !account.IsOpenAIApiKey() && !account.IsOpenAIOAuth() {
|
||||
continue
|
||||
}
|
||||
rate, ok := openAISchedulingRate(account, now, oauthSchedulingRateMultiplier)
|
||||
if !ok {
|
||||
continue
|
||||
|
||||
@@ -491,6 +491,46 @@ func TestOpenAILegacyUpstreamRateOrderRequiresComparableRates(t *testing.T) {
|
||||
require.Negative(t, distinct.compare(&Account{ID: 2}, &Account{ID: 3}))
|
||||
}
|
||||
|
||||
// 探测资格已放宽到全部 API-key 平台,但调度侧的信任面没有跟着扩大:
|
||||
// 只有 OpenAI 平台账号的上游自报倍率参与 legacy 低倍率优先排序,
|
||||
// 否则中转方自报低价即可吸走流量,而实际结算走本地倍率。
|
||||
// 本用例钉死 newOpenAILegacyUpstreamRateOrder 与 openAIUpstreamCostFactors
|
||||
// 使用同一道平台门控。
|
||||
func TestOpenAILegacyUpstreamRateOrderIgnoresNonOpenAIPlatforms(t *testing.T) {
|
||||
now := time.Now()
|
||||
nonOpenAI := func(id int64, platform string, rate float64) *Account {
|
||||
account := upstreamCostTestAccount(id, UpstreamBillingProbeStatusOK, rate, now.Add(-time.Minute), 30*time.Minute)
|
||||
account.Platform = platform
|
||||
return account
|
||||
}
|
||||
grokCheap := nonOpenAI(1, PlatformGrok, 0.01)
|
||||
anthropicExpensive := nonOpenAI(2, PlatformAnthropic, 0.9)
|
||||
|
||||
order := newOpenAILegacyUpstreamRateOrder([]*Account{grokCheap, anthropicExpensive, nil}, now, defaultOpenAIOAuthSchedulingRateMultiplier)
|
||||
require.False(t, order.enabled)
|
||||
require.Empty(t, order.rates)
|
||||
require.Zero(t, order.compare(grokCheap, anthropicExpensive))
|
||||
|
||||
factors := openAIUpstreamCostFactors([]*Account{grokCheap, anthropicExpensive}, now, defaultOpenAIOAuthSchedulingRateMultiplier)
|
||||
require.Equal(t, openAIUpstreamCostNeutralFactor, factors[grokCheap.ID])
|
||||
require.Equal(t, openAIUpstreamCostNeutralFactor, factors[anthropicExpensive.ID])
|
||||
|
||||
// 混合候选集里,非 OpenAI 账号既不进 rates 也不影响 OpenAI 账号之间的排序。
|
||||
openAICheap := upstreamCostTestAccount(3, UpstreamBillingProbeStatusOK, 0.02, now.Add(-time.Minute), 30*time.Minute)
|
||||
openAIExpensive := upstreamCostTestAccount(4, UpstreamBillingProbeStatusOK, 0.12, now.Add(-time.Minute), 30*time.Minute)
|
||||
mixed := newOpenAILegacyUpstreamRateOrder(
|
||||
[]*Account{grokCheap, openAICheap, anthropicExpensive, openAIExpensive},
|
||||
now, defaultOpenAIOAuthSchedulingRateMultiplier,
|
||||
)
|
||||
require.True(t, mixed.enabled)
|
||||
require.Len(t, mixed.rates, 2)
|
||||
require.NotContains(t, mixed.rates, grokCheap.ID)
|
||||
require.NotContains(t, mixed.rates, anthropicExpensive.ID)
|
||||
require.Negative(t, mixed.compare(openAICheap, openAIExpensive))
|
||||
// 自报 0.01 的 grok 账号没有已知倍率,排在有倍率的 OpenAI 账号之后。
|
||||
require.Positive(t, mixed.compare(grokCheap, openAIExpensive))
|
||||
}
|
||||
|
||||
func TestOpenAISchedulingRatePlacesOAuthAtConfiguredReference(t *testing.T) {
|
||||
now := time.Now()
|
||||
cheap := upstreamCostTestAccount(1, UpstreamBillingProbeStatusOK, 0.02, now.Add(-time.Minute), 30*time.Minute)
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"math"
|
||||
"math/rand/v2"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
@@ -39,6 +40,10 @@ const (
|
||||
upstreamBillingProbeMaxPerCycle = 20
|
||||
upstreamBillingProbeConcurrency = 4
|
||||
upstreamBillingProbeMaxDelay = 24 * time.Hour
|
||||
// unsupported 账号的重探间隔倍数:上游不是 sub2api 中转就不会突然长出
|
||||
// /v1/sub2api/billing,按常规 interval 重排只会持续占满每周期
|
||||
// upstreamBillingProbeMaxPerCycle 个名额。
|
||||
upstreamBillingProbeUnsupportedDelayFactor = 8
|
||||
upstreamBillingProbeLeaderLockKey = "upstream:billing:probe:leader"
|
||||
upstreamBillingProbeLeaderLockTTL = 2 * time.Minute
|
||||
)
|
||||
@@ -51,7 +56,7 @@ var (
|
||||
"UPSTREAM_BILLING_PROBE_UNAVAILABLE", "upstream billing probe is unavailable",
|
||||
)
|
||||
ErrUpstreamBillingProbeAccountInvalid = infraerrors.BadRequest(
|
||||
"UPSTREAM_BILLING_PROBE_ACCOUNT_INVALID", "account is not an OpenAI API key account",
|
||||
"UPSTREAM_BILLING_PROBE_ACCOUNT_INVALID", "account is not an API key account",
|
||||
)
|
||||
ErrUpstreamBillingProbeIdentityChanged = infraerrors.Conflict(
|
||||
"UPSTREAM_BILLING_PROBE_IDENTITY_CHANGED", "account identity changed during upstream billing probe; retry the probe",
|
||||
@@ -551,13 +556,24 @@ func (s *UpstreamBillingProbeService) probeLoadedAccount(ctx context.Context, ac
|
||||
if s.accountTestService == nil || s.accountTestService.httpUpstream == nil {
|
||||
return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "transport_unavailable", 0)
|
||||
}
|
||||
apiKey := account.GetOpenAIApiKey()
|
||||
// 平台放宽后取数直读 credentials:所有 API-key 平台的密钥与自定义上游
|
||||
// 统一存放在 credentials.api_key / credentials.base_url。
|
||||
apiKey := account.GetCredential("api_key")
|
||||
if apiKey == "" {
|
||||
return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "missing_api_key", 0)
|
||||
}
|
||||
baseURL := account.GetOpenAIBaseURL()
|
||||
if baseURL == "" {
|
||||
baseURL = "https://api.openai.com"
|
||||
baseURL := account.GetCredential("base_url")
|
||||
if account.Platform == PlatformOpenAI {
|
||||
if baseURL == "" {
|
||||
// 保持官方语义:OpenAI 账号无自定义 base 时探官方域(404 → unsupported)。
|
||||
baseURL = "https://api.openai.com"
|
||||
}
|
||||
} else if upstreamBillingProbeTargetIsOfficialAPI(baseURL) {
|
||||
// 其他平台 base_url 为空或指向官方 API 根域(前端创建时会把空值
|
||||
// 填成官方默认域,且提供 us-east-1.api.x.ai 等官方区域预设)⇒
|
||||
// 必无 /v1/sub2api/billing;不发请求,直接记 unsupported,避免
|
||||
// 拿账号 Key 周期性请求官方域的不存在路径。
|
||||
return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "unsupported", 0)
|
||||
}
|
||||
normalizedBaseURL, err := s.accountTestService.validateUpstreamBaseURL(baseURL)
|
||||
if err != nil {
|
||||
@@ -580,7 +596,12 @@ func (s *UpstreamBillingProbeService) probeLoadedAccount(ctx context.Context, ac
|
||||
if err != nil {
|
||||
return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "request_build_failed", 0)
|
||||
}
|
||||
reqCtx := WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI)
|
||||
// OpenAI 账号保持官方 openai 传输画像;其他平台探测走默认画像。
|
||||
profile := HTTPUpstreamProfileDefault
|
||||
if account.Platform == PlatformOpenAI {
|
||||
profile = HTTPUpstreamProfileOpenAI
|
||||
}
|
||||
reqCtx := WithHTTPUpstreamProfile(req.Context(), profile)
|
||||
req = req.WithContext(WithHTTPUpstreamRedirectsDisabled(reqCtx))
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
@@ -644,13 +665,15 @@ func (s *UpstreamBillingProbeService) persistProbeFailure(
|
||||
failureCount = previous.FailureCount + 1
|
||||
}
|
||||
status := UpstreamBillingProbeStatusFailed
|
||||
delay := nextProbeDelay(intervalMinutes, retryAfterDuration)
|
||||
if reason == "unsupported" {
|
||||
status = UpstreamBillingProbeStatusUnsupported
|
||||
delay = unsupportedProbeDelay(intervalMinutes, retryAfterDuration)
|
||||
}
|
||||
snapshot := &UpstreamBillingProbeSnapshot{
|
||||
Status: status,
|
||||
LastAttemptAt: now,
|
||||
NextProbeAt: now.Add(nextProbeDelay(intervalMinutes, retryAfterDuration)),
|
||||
NextProbeAt: now.Add(delay),
|
||||
FailureCount: failureCount,
|
||||
HTTPStatus: statusCode,
|
||||
LastError: reason,
|
||||
@@ -840,8 +863,67 @@ func decodeUpstreamBillingProbeSnapshot(extra map[string]any) *UpstreamBillingPr
|
||||
return &snapshot
|
||||
}
|
||||
|
||||
// IsUpstreamBillingProbeIdentity reports whether an account identity may opt
|
||||
// in to the upstream billing probe. `/v1/sub2api/billing` is a key-scoped
|
||||
// sub2api convention: any API-key account whose base_url points at a
|
||||
// sub2api-compatible upstream answers it regardless of the account platform,
|
||||
// so eligibility is not restricted to OpenAI.
|
||||
// Non-sub2api upstreams return 404 and the snapshot records "unsupported".
|
||||
// Only AccountTypeAPIKey is in scope. OAuth/Bedrock hold no static API key to
|
||||
// present at all; AccountTypeUpstream (antigravity relay accounts) does carry
|
||||
// a base_url plus a static api_key, but it is deliberately left out of the
|
||||
// current supported set. New antigravity relay accounts are created with
|
||||
// type=apikey by the admin form, so only pre-existing type=upstream rows
|
||||
// cannot turn the probe on.
|
||||
func IsUpstreamBillingProbeIdentity(platform, accountType string) bool {
|
||||
return platform != "" && accountType == AccountTypeAPIKey
|
||||
}
|
||||
|
||||
func isUpstreamBillingProbeAccount(account *Account) bool {
|
||||
return account != nil && account.Platform == PlatformOpenAI && account.Type == AccountTypeAPIKey
|
||||
return account != nil && IsUpstreamBillingProbeIdentity(account.Platform, account.Type)
|
||||
}
|
||||
|
||||
// upstreamBillingProbeOfficialAPIDomains lists the root domains of official
|
||||
// provider APIs. The create form fills empty base_url values with official
|
||||
// defaults (and offers official regional presets like us-east-1.api.x.ai),
|
||||
// so probing them would send the account key to an official API path that
|
||||
// cannot exist. Matching is by registrable root domain — exact host or any
|
||||
// subdomain, after stripping the port and a trailing DNS dot — because no
|
||||
// third-party sub2api relay can live under these domains, while custom
|
||||
// relays (the only targets that can answer /v1/sub2api/billing) always do
|
||||
// probe. OpenAI-platform accounts never reach this check: they keep the
|
||||
// upstream-official behavior of probing api.openai.com.
|
||||
// ollama.com is a first-class configuration here (Ollama Cloud accounts are
|
||||
// platform openai/anthropic with base_url https://ollama.com/v1), and it is
|
||||
// an official provider API just like the rest, so it belongs on this list.
|
||||
var upstreamBillingProbeOfficialAPIDomains = []string{
|
||||
"anthropic.com",
|
||||
"googleapis.com",
|
||||
"x.ai",
|
||||
"grok.com",
|
||||
"openai.com",
|
||||
"ollama.com",
|
||||
}
|
||||
|
||||
func upstreamBillingProbeTargetIsOfficialAPI(baseURL string) bool {
|
||||
baseURL = strings.TrimSpace(baseURL)
|
||||
if baseURL == "" {
|
||||
return true
|
||||
}
|
||||
parsed, err := url.Parse(baseURL)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
host := strings.TrimSuffix(strings.ToLower(parsed.Hostname()), ".")
|
||||
if host == "" {
|
||||
return true
|
||||
}
|
||||
for _, domain := range upstreamBillingProbeOfficialAPIDomains {
|
||||
if host == domain || strings.HasSuffix(host, "."+domain) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func upstreamBillingProbeEnabled(account *Account) bool {
|
||||
@@ -885,6 +967,23 @@ func nextProbeDelay(intervalMinutes int, retryAfterDuration time.Duration) time.
|
||||
return interval
|
||||
}
|
||||
|
||||
// unsupportedProbeDelay 拉长 unsupported 账号的重探间隔,让无效候选自然退出
|
||||
// 热队列,不再和真正接入 sub2api 的中转账号抢每周期的探测名额。
|
||||
// 仍按 upstreamBillingProbeMaxDelay 封顶,保证上游后来接入 sub2api 时最迟一天
|
||||
// 内会被重新发现;base 本身已达上限(例如 Retry-After 明确要求更久)时原样返回,
|
||||
// 不缩短上游指令。
|
||||
func unsupportedProbeDelay(intervalMinutes int, retryAfterDuration time.Duration) time.Duration {
|
||||
base := nextProbeDelay(intervalMinutes, retryAfterDuration)
|
||||
if base >= upstreamBillingProbeMaxDelay {
|
||||
return base
|
||||
}
|
||||
stretched := base * upstreamBillingProbeUnsupportedDelayFactor
|
||||
if stretched > upstreamBillingProbeMaxDelay {
|
||||
return upstreamBillingProbeMaxDelay
|
||||
}
|
||||
return stretched
|
||||
}
|
||||
|
||||
func retryAfter(header http.Header, now time.Time) time.Duration {
|
||||
value := strings.TrimSpace(header.Get("Retry-After"))
|
||||
if value == "" {
|
||||
|
||||
@@ -0,0 +1,215 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// 探测资格放宽:/v1/sub2api/billing 是 key 级端点,任何 API-key
|
||||
// 平台账号都可开启探测;OAuth/Bedrock 无静态 Key 仍不合格。
|
||||
func TestUpstreamBillingProbeIdentityCoversAllAPIKeyPlatforms(t *testing.T) {
|
||||
for _, platform := range []string{
|
||||
PlatformOpenAI, PlatformGrok, PlatformAnthropic, PlatformGemini, PlatformAntigravity,
|
||||
} {
|
||||
require.True(t, IsUpstreamBillingProbeIdentity(platform, AccountTypeAPIKey), platform)
|
||||
require.True(t, isUpstreamBillingProbeAccount(&Account{Platform: platform, Type: AccountTypeAPIKey}), platform)
|
||||
}
|
||||
require.False(t, IsUpstreamBillingProbeIdentity(PlatformOpenAI, AccountTypeOAuth))
|
||||
require.False(t, IsUpstreamBillingProbeIdentity(PlatformGrok, AccountTypeOAuth))
|
||||
require.False(t, IsUpstreamBillingProbeIdentity(PlatformAnthropic, AccountTypeBedrock))
|
||||
require.False(t, IsUpstreamBillingProbeIdentity("", AccountTypeAPIKey))
|
||||
require.False(t, isUpstreamBillingProbeAccount(nil))
|
||||
}
|
||||
|
||||
func upstreamBillingProbeValidBody() io.ReadCloser {
|
||||
return io.NopCloser(strings.NewReader(`{
|
||||
"object":"sub2api.key_billing",
|
||||
"schema_version":1,
|
||||
"billing_scope":"token",
|
||||
"group_rate_multiplier":0.02,
|
||||
"resolved_rate_multiplier":0.02,
|
||||
"peak_rate_enabled":false,
|
||||
"effective_rate_multiplier":0.02,
|
||||
"observed_at":"2026-07-26T01:00:00Z"
|
||||
}`))
|
||||
}
|
||||
|
||||
func TestUpstreamBillingProbeGrokAccountPersistsSnapshot(t *testing.T) {
|
||||
account := &Account{
|
||||
ID: 151,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeAPIKey,
|
||||
Status: StatusActive,
|
||||
Concurrency: 2,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "sk-grok-relay",
|
||||
"base_url": "https://relay.example/v1",
|
||||
},
|
||||
}
|
||||
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{account.ID: account}}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: upstreamBillingProbeValidBody(),
|
||||
}}
|
||||
svc := newUpstreamBillingProbeTestService(repo, upstream, &upstreamBillingProbeSettingRepo{})
|
||||
fixedNow := time.Date(2026, time.July, 26, 2, 0, 0, 0, time.UTC)
|
||||
svc.now = func() time.Time { return fixedNow }
|
||||
|
||||
snapshot, err := svc.ProbeAccount(context.Background(), account.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, UpstreamBillingProbeStatusOK, snapshot.Status)
|
||||
require.Equal(t, 0.02, snapshot.Data["resolved_rate_multiplier"])
|
||||
require.Equal(t, "https://relay.example/v1/sub2api/billing", upstream.lastReq.URL.String())
|
||||
require.Equal(t, "Bearer sk-grok-relay", upstream.lastReq.Header.Get("Authorization"))
|
||||
// 非 OpenAI 平台探测使用默认传输画像。
|
||||
require.Equal(t, HTTPUpstreamProfileDefault, HTTPUpstreamProfileFromContext(upstream.lastReq.Context()))
|
||||
|
||||
persisted := decodeUpstreamBillingProbeSnapshot(account.Extra)
|
||||
require.NotNil(t, persisted)
|
||||
require.Equal(t, UpstreamBillingProbeStatusOK, persisted.Status)
|
||||
}
|
||||
|
||||
// 非 OpenAI 平台没有自定义 base_url 时,上游是各自官方 API,必无
|
||||
// /v1/sub2api/billing:不发请求,直接落 unsupported。
|
||||
func TestUpstreamBillingProbeNonOpenAIWithoutBaseURLIsUnsupportedWithoutRequest(t *testing.T) {
|
||||
account := &Account{
|
||||
ID: 152,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeAPIKey,
|
||||
Status: StatusActive,
|
||||
Credentials: map[string]any{"api_key": "sk-grok-official"},
|
||||
}
|
||||
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{account.ID: account}}
|
||||
upstream := &httpUpstreamRecorder{}
|
||||
svc := newUpstreamBillingProbeTestService(repo, upstream, &upstreamBillingProbeSettingRepo{})
|
||||
|
||||
snapshot, err := svc.ProbeAccount(context.Background(), account.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, UpstreamBillingProbeStatusUnsupported, snapshot.Status)
|
||||
require.Nil(t, upstream.lastReq)
|
||||
}
|
||||
|
||||
// 前端创建 API-key 账号时会把空 base_url 填成平台官方默认域;官方 API 同样
|
||||
// 必无 /v1/sub2api/billing,不发请求直接 unsupported。
|
||||
// 显式端口、DNS 尾点、官方区域子域(前端 Grok 预设含 us-east-1.api.x.ai 等)
|
||||
// 与跨平台官方域都不能绕过。
|
||||
func TestUpstreamBillingProbeOfficialAPIBaseURLIsUnsupportedWithoutRequest(t *testing.T) {
|
||||
cases := []struct {
|
||||
platform string
|
||||
baseURL string
|
||||
}{
|
||||
{PlatformAnthropic, "https://api.anthropic.com"},
|
||||
{PlatformAnthropic, "https://api.anthropic.com:443"},
|
||||
{PlatformAnthropic, "https://api.anthropic.com./"},
|
||||
{PlatformAnthropic, "HTTPS://API.ANTHROPIC.COM/"},
|
||||
{PlatformGemini, "https://generativelanguage.googleapis.com"},
|
||||
{PlatformAntigravity, "https://cloudcode-pa.googleapis.com"},
|
||||
{PlatformGrok, "https://api.x.ai/v1"},
|
||||
{PlatformGrok, "https://us-east-1.api.x.ai/v1"},
|
||||
{PlatformGrok, "https://eu-west-1.api.x.ai/v1"},
|
||||
{PlatformGrok, "https://cli-chat-proxy.grok.com/v1"},
|
||||
// 跨平台官方域同样必无该端点,一并拦截。
|
||||
{PlatformAnthropic, "https://api.x.ai/v1"},
|
||||
{PlatformGrok, "https://api.openai.com"},
|
||||
// Ollama Cloud 是本仓一等支持配置(platform openai/anthropic +
|
||||
// base_url https://ollama.com/v1),同为官方 API,不能拿 Key 去空探。
|
||||
{PlatformAnthropic, "https://ollama.com/v1"},
|
||||
{PlatformAnthropic, "https://ollama.com"},
|
||||
{PlatformAnthropic, "https://www.ollama.com/v1"},
|
||||
}
|
||||
for i, tc := range cases {
|
||||
account := &Account{
|
||||
ID: int64(200 + i),
|
||||
Platform: tc.platform,
|
||||
Type: AccountTypeAPIKey,
|
||||
Status: StatusActive,
|
||||
Credentials: map[string]any{"api_key": "sk-official", "base_url": tc.baseURL},
|
||||
}
|
||||
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{account.ID: account}}
|
||||
upstream := &httpUpstreamRecorder{}
|
||||
svc := newUpstreamBillingProbeTestService(repo, upstream, &upstreamBillingProbeSettingRepo{})
|
||||
|
||||
snapshot, err := svc.ProbeAccount(context.Background(), account.ID)
|
||||
require.NoError(t, err, "%s %s", tc.platform, tc.baseURL)
|
||||
require.Equal(t, UpstreamBillingProbeStatusUnsupported, snapshot.Status, "%s %s", tc.platform, tc.baseURL)
|
||||
require.Nil(t, upstream.lastReq, "%s %s", tc.platform, tc.baseURL)
|
||||
}
|
||||
}
|
||||
|
||||
// 官方域判定按规范化 host 的根域后缀匹配:端口、尾点、区域子域都拦;
|
||||
// 自定义中转与相似但不同的域照常探测。
|
||||
func TestUpstreamBillingProbeOfficialAPIHostMatchingIsNormalized(t *testing.T) {
|
||||
require.True(t, upstreamBillingProbeTargetIsOfficialAPI(""))
|
||||
require.True(t, upstreamBillingProbeTargetIsOfficialAPI("https://api.x.ai"))
|
||||
require.True(t, upstreamBillingProbeTargetIsOfficialAPI("https://api.anthropic.com:8443"))
|
||||
require.True(t, upstreamBillingProbeTargetIsOfficialAPI("https://us-west-2.api.x.ai/v1"))
|
||||
require.True(t, upstreamBillingProbeTargetIsOfficialAPI("https://generativelanguage.googleapis.com."))
|
||||
require.True(t, upstreamBillingProbeTargetIsOfficialAPI("https://x.ai"))
|
||||
// openai 官方域在全集内;openai 平台账号不经过本判定(行为级测试钉死照探)。
|
||||
require.True(t, upstreamBillingProbeTargetIsOfficialAPI("https://api.openai.com"))
|
||||
// Ollama Cloud 官方域及其子域(Ollama Cloud 账号的 base_url 允许带 www.)。
|
||||
require.True(t, upstreamBillingProbeTargetIsOfficialAPI("https://ollama.com/v1"))
|
||||
require.True(t, upstreamBillingProbeTargetIsOfficialAPI("https://ollama.com:443/v1"))
|
||||
require.True(t, upstreamBillingProbeTargetIsOfficialAPI("https://www.ollama.com/v1"))
|
||||
require.True(t, upstreamBillingProbeTargetIsOfficialAPI("HTTPS://OLLAMA.COM./v1"))
|
||||
// 相似但不同的注册域不拦:中转完全可能叫 *-x.ai 之外的任何名字。
|
||||
require.False(t, upstreamBillingProbeTargetIsOfficialAPI("https://relay.example/v1"))
|
||||
require.False(t, upstreamBillingProbeTargetIsOfficialAPI("https://notx.ai"))
|
||||
require.False(t, upstreamBillingProbeTargetIsOfficialAPI("https://anthropic.com.evil.example"))
|
||||
require.False(t, upstreamBillingProbeTargetIsOfficialAPI("https://api.relay-station.example"))
|
||||
require.False(t, upstreamBillingProbeTargetIsOfficialAPI("https://notollama.com/v1"))
|
||||
require.False(t, upstreamBillingProbeTargetIsOfficialAPI("https://ollama.com.evil.example/v1"))
|
||||
require.False(t, upstreamBillingProbeTargetIsOfficialAPI("https://ollama.example/v1"))
|
||||
}
|
||||
|
||||
// OpenAI 语义保持不变:无自定义 base 时仍探官方域,且沿用 openai 传输画像。
|
||||
func TestUpstreamBillingProbeOpenAIDefaultBaseURLPreserved(t *testing.T) {
|
||||
account := &Account{
|
||||
ID: 17,
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeAPIKey,
|
||||
Status: StatusActive,
|
||||
Credentials: map[string]any{"api_key": "sk-openai"},
|
||||
}
|
||||
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{account.ID: account}}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: upstreamBillingProbeValidBody(),
|
||||
}}
|
||||
svc := newUpstreamBillingProbeTestService(repo, upstream, &upstreamBillingProbeSettingRepo{})
|
||||
|
||||
_, err := svc.ProbeAccount(context.Background(), account.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://api.openai.com/v1/sub2api/billing", upstream.lastReq.URL.String())
|
||||
require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(upstream.lastReq.Context()))
|
||||
}
|
||||
|
||||
func TestUpstreamBillingProbeSetAccountEnabledAcceptsGrokAPIKey(t *testing.T) {
|
||||
grokAPIKey := &Account{
|
||||
ID: 151,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeAPIKey,
|
||||
Status: StatusActive,
|
||||
Credentials: map[string]any{"api_key": "sk", "base_url": "https://relay.example"},
|
||||
}
|
||||
grokOAuth := &Account{ID: 152, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive}
|
||||
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
|
||||
grokAPIKey.ID: grokAPIKey,
|
||||
grokOAuth.ID: grokOAuth,
|
||||
}}
|
||||
svc := newUpstreamBillingProbeTestService(repo, &upstreamBillingProbeHTTPStub{}, &upstreamBillingProbeSettingRepo{})
|
||||
|
||||
require.NoError(t, svc.SetAccountEnabled(context.Background(), grokAPIKey.ID, true))
|
||||
require.Equal(t, true, repo.accounts[grokAPIKey.ID].Extra[UpstreamBillingProbeEnabledExtraKey])
|
||||
|
||||
err := svc.SetAccountEnabled(context.Background(), grokOAuth.ID, true)
|
||||
require.ErrorIs(t, err, ErrUpstreamBillingProbeAccountInvalid)
|
||||
}
|
||||
@@ -499,6 +499,69 @@ func TestUpstreamBillingProbeRetryAfterIsNotShortened(t *testing.T) {
|
||||
require.Equal(t, 48*time.Hour, delay)
|
||||
}
|
||||
|
||||
// unsupported 的重探间隔明显长于普通失败,但始终有上界:上游后来接入 sub2api
|
||||
// 时最迟一天内会被重新发现,且不会缩短上游 Retry-After 指令。
|
||||
func TestUpstreamBillingProbeUnsupportedDelayIsStretchedAndBounded(t *testing.T) {
|
||||
// 默认 30 分钟 interval:普通失败 24~36 分钟,unsupported 为其 8 倍。
|
||||
stretched := unsupportedProbeDelay(30, 0)
|
||||
require.Greater(t, stretched, 36*time.Minute)
|
||||
require.GreaterOrEqual(t, stretched, 192*time.Minute)
|
||||
require.LessOrEqual(t, stretched, 288*time.Minute)
|
||||
|
||||
// 永不超过封顶值,因此 unsupported 账号不会被永久排除在重探之外。
|
||||
require.LessOrEqual(t, unsupportedProbeDelay(upstreamBillingProbeMaxIntervalMinutes, 0), upstreamBillingProbeMaxDelay)
|
||||
require.Positive(t, unsupportedProbeDelay(upstreamBillingProbeMinIntervalMinutes, 0))
|
||||
|
||||
// Retry-After 更长时原样保留,不被封顶缩短;更短时至少不早于该指令。
|
||||
require.Equal(t, 48*time.Hour, unsupportedProbeDelay(30, 48*time.Hour))
|
||||
require.GreaterOrEqual(t, unsupportedProbeDelay(30, time.Hour), time.Hour)
|
||||
}
|
||||
|
||||
// 加长退避只把 unsupported 账号移出周期性热队列,手动探测不受影响。
|
||||
func TestUpstreamBillingProbeUnsupportedBackoffDefersRunnerButNotManualProbe(t *testing.T) {
|
||||
// Ollama Cloud 形态:官方域,不发请求直接落 unsupported。
|
||||
account := &Account{
|
||||
ID: 31,
|
||||
Platform: PlatformAnthropic,
|
||||
Type: AccountTypeAPIKey,
|
||||
Status: StatusActive,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{"api_key": "sk-ollama", "base_url": "https://ollama.com/v1"},
|
||||
Extra: map[string]any{UpstreamBillingProbeEnabledExtraKey: true},
|
||||
}
|
||||
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{account.ID: account}}
|
||||
upstream := &upstreamBillingProbeHTTPStub{}
|
||||
settingsRepo := &upstreamBillingProbeSettingRepo{values: map[string]string{
|
||||
SettingKeyUpstreamBillingProbeSettings: `{"enabled":true,"interval_minutes":30}`,
|
||||
}}
|
||||
svc := newUpstreamBillingProbeTestService(repo, upstream, settingsRepo)
|
||||
start := time.Date(2026, time.July, 26, 2, 0, 0, 0, time.UTC)
|
||||
now := start
|
||||
svc.now = func() time.Time { return now }
|
||||
|
||||
require.NoError(t, svc.RunDue(context.Background()))
|
||||
first := decodeUpstreamBillingProbeSnapshot(account.Extra)
|
||||
require.NotNil(t, first)
|
||||
require.Equal(t, UpstreamBillingProbeStatusUnsupported, first.Status)
|
||||
require.Equal(t, start, first.LastAttemptAt)
|
||||
require.False(t, first.NextProbeAt.Before(start.Add(192*time.Minute)))
|
||||
require.Zero(t, upstream.calls.Load())
|
||||
|
||||
// 一个普通失败早就该重探的时间点(远超 36 分钟),runner 仍跳过该账号。
|
||||
now = start.Add(90 * time.Minute)
|
||||
require.NoError(t, svc.RunDue(context.Background()))
|
||||
deferred := decodeUpstreamBillingProbeSnapshot(account.Extra)
|
||||
require.NotNil(t, deferred)
|
||||
require.Equal(t, first.LastAttemptAt, deferred.LastAttemptAt)
|
||||
|
||||
// 手动探测无视退避窗口,管理员随时可以重试。
|
||||
manual, err := svc.ProbeAccount(context.Background(), account.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, UpstreamBillingProbeStatusUnsupported, manual.Status)
|
||||
require.Equal(t, now, manual.LastAttemptAt)
|
||||
require.Equal(t, 2, manual.FailureCount)
|
||||
}
|
||||
|
||||
func TestUpstreamBillingProbeEmptyResponseIsPersistedAsFailure(t *testing.T) {
|
||||
account := &Account{
|
||||
ID: 21,
|
||||
@@ -552,14 +615,15 @@ func TestUpstreamBillingProbeUnsupportedAndAccountToggle(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, UpstreamBillingProbeStatusUnsupported, snapshot.Status)
|
||||
require.Equal(t, "unsupported", snapshot.LastError)
|
||||
require.False(t, snapshot.NextProbeAt.Before(fixedNow.Add(24*time.Minute)))
|
||||
require.False(t, snapshot.NextProbeAt.After(fixedNow.Add(36*time.Minute)))
|
||||
// unsupported 走加长退避:默认 30 分钟 interval ⇒ (24~36) * 8 分钟。
|
||||
require.False(t, snapshot.NextProbeAt.Before(fixedNow.Add(192*time.Minute)))
|
||||
require.False(t, snapshot.NextProbeAt.After(fixedNow.Add(288*time.Minute)))
|
||||
|
||||
snapshot, err = svc.ProbeAccount(context.Background(), account.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 2, snapshot.FailureCount)
|
||||
require.False(t, snapshot.NextProbeAt.Before(fixedNow.Add(24*time.Minute)))
|
||||
require.False(t, snapshot.NextProbeAt.After(fixedNow.Add(36*time.Minute)))
|
||||
require.False(t, snapshot.NextProbeAt.Before(fixedNow.Add(192*time.Minute)))
|
||||
require.False(t, snapshot.NextProbeAt.After(fixedNow.Add(288*time.Minute)))
|
||||
|
||||
invalid := &Account{ID: 20, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
||||
repo.accounts[invalid.ID] = invalid
|
||||
|
||||
@@ -905,8 +905,8 @@
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Upstream billing auto probe (OpenAI API Key only) -->
|
||||
<div v-if="allOpenAIAPIKey" class="border-t border-gray-200 pt-4 dark:border-dark-600">
|
||||
<!-- Upstream billing auto probe (any API-key platform) -->
|
||||
<div v-if="allBillingProbeCapable" class="border-t border-gray-200 pt-4 dark:border-dark-600">
|
||||
<div class="mb-3 flex items-center justify-between">
|
||||
<div class="flex-1 pr-4">
|
||||
<label
|
||||
@@ -1401,6 +1401,15 @@ const allOpenAIAPIKey = computed(() => {
|
||||
)
|
||||
})
|
||||
|
||||
// 上游倍率自动探测已放宽到全部 API-key 平台:只要求所选类型全为 apikey,
|
||||
// 平台不限(sub2api 上游即可应答 /v1/sub2api/billing)。
|
||||
const allBillingProbeCapable = computed(() => {
|
||||
return (
|
||||
targetSelectedTypes.value.length > 0 &&
|
||||
targetSelectedTypes.value.every(t => t === 'apikey')
|
||||
)
|
||||
})
|
||||
|
||||
// 是否全部为 anthropic/openai 平台的 apikey 账号(请求头覆写仅在此条件下显示)
|
||||
// 所选平台 × 所选类型的全组合均需具备覆写资格(实际选中账号是该组合的子集,
|
||||
// 按交叉积判定偏保守但绝不放行不合资格的账号)
|
||||
|
||||
@@ -888,6 +888,20 @@
|
||||
/>
|
||||
<p class="input-hint">{{ t('admin.accounts.upstream.apiKeyHint') }}</p>
|
||||
</div>
|
||||
<!-- 上游倍率自动探测:antigravity upstream 也是 API-key 账号 -->
|
||||
<div class="flex items-center justify-between gap-4 border-t border-gray-200 pt-4 dark:border-dark-600">
|
||||
<div>
|
||||
<label class="input-label mb-0">{{ t('admin.accounts.upstreamBilling.autoProbe') }}</label>
|
||||
<p class="mt-1 text-xs text-gray-500 dark:text-gray-400">
|
||||
{{ t('admin.accounts.upstreamBilling.autoProbeHint') }}
|
||||
</p>
|
||||
</div>
|
||||
<Toggle
|
||||
v-model="upstreamBillingAutoProbeEnabled"
|
||||
data-testid="upstream-billing-auto-probe-antigravity"
|
||||
:aria-label="t('admin.accounts.upstreamBilling.autoProbe')"
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Vertex Service Account -->
|
||||
@@ -1143,8 +1157,8 @@
|
||||
<p v-if="apiKeyHint" class="input-hint">{{ apiKeyHint }}</p>
|
||||
</div>
|
||||
|
||||
<!-- 上游倍率自动探测:全部 API-key 平台可用(所在区块已限定 apikey 类型) -->
|
||||
<div
|
||||
v-if="form.platform === 'openai'"
|
||||
class="flex items-center justify-between gap-4 border-t border-gray-200 pt-4 dark:border-dark-600"
|
||||
>
|
||||
<div>
|
||||
@@ -4608,7 +4622,6 @@ const submitCreateAccount = async (payload: CreateAccountRequest) => {
|
||||
try {
|
||||
const account = await adminAPI.accounts.create(withAntigravityConfirmFlag(payload))
|
||||
if (
|
||||
payload.platform === 'openai' &&
|
||||
payload.type === 'apikey' &&
|
||||
payload.upstream_billing_probe_enabled === true
|
||||
) {
|
||||
@@ -5164,8 +5177,7 @@ const handleSubmit = async () => {
|
||||
...form,
|
||||
group_ids: form.group_ids,
|
||||
extra,
|
||||
upstream_billing_probe_enabled:
|
||||
form.platform === 'openai' ? upstreamBillingAutoProbeEnabled.value : undefined,
|
||||
upstream_billing_probe_enabled: upstreamBillingAutoProbeEnabled.value,
|
||||
auto_pause_on_expired: autoPauseOnExpired.value
|
||||
})
|
||||
}
|
||||
@@ -5294,6 +5306,9 @@ const createAccountAndFinish = async (
|
||||
rate_multiplier: form.rate_multiplier,
|
||||
group_ids: form.group_ids,
|
||||
expires_at: form.expires_at,
|
||||
// 上游倍率探测对全部 API-key 平台开放(antigravity upstream 走本 helper);
|
||||
// 非 apikey 类型(bedrock/oauth)不传,后端不动作。
|
||||
upstream_billing_probe_enabled: type === 'apikey' ? upstreamBillingAutoProbeEnabled.value : undefined,
|
||||
auto_pause_on_expired: autoPauseOnExpired.value
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1650,7 +1650,7 @@
|
||||
</div>
|
||||
|
||||
<div
|
||||
v-if="account?.platform === 'openai' && account?.type === 'apikey'"
|
||||
v-if="account?.type === 'apikey'"
|
||||
class="flex items-center justify-between gap-4 border-t border-gray-200 pt-4 dark:border-dark-600"
|
||||
>
|
||||
<div>
|
||||
@@ -4584,7 +4584,6 @@ const handleSubmit = async () => {
|
||||
} else {
|
||||
newExtra.openai_responses_mode = openAIResponsesMode.value
|
||||
}
|
||||
newExtra.upstream_billing_probe_enabled = upstreamBillingAutoProbeEnabled.value
|
||||
}
|
||||
if (autoPause5hThreshold.value != null && autoPause5hThreshold.value > 0) {
|
||||
newExtra.auto_pause_5h_threshold = autoPause5hThreshold.value / 100
|
||||
@@ -4649,6 +4648,11 @@ const handleSubmit = async () => {
|
||||
const currentExtra = (updatePayload.extra as Record<string, unknown>) ||
|
||||
(props.account.extra as Record<string, unknown>) || {}
|
||||
const newExtra: Record<string, unknown> = { ...currentExtra }
|
||||
// 上游倍率自动探测对全部 API-key 平台开放(sub2api 上游即可应答),
|
||||
// Bedrock 凭证无静态 Key 不参与。
|
||||
if (props.account.type === 'apikey') {
|
||||
newExtra.upstream_billing_probe_enabled = upstreamBillingAutoProbeEnabled.value
|
||||
}
|
||||
// Total quota
|
||||
if (editQuotaLimit.value != null && editQuotaLimit.value > 0) {
|
||||
newExtra.quota_limit = editQuotaLimit.value
|
||||
|
||||
@@ -105,7 +105,8 @@ defineEmits<{
|
||||
|
||||
const { t } = useI18n()
|
||||
const CLOCK_SKEW_TOLERANCE_MS = 5 * 60 * 1000
|
||||
const eligible = computed(() => props.account.platform === 'openai' && props.account.type === 'apikey')
|
||||
// 探测资格已放宽到全部 API-key 平台(上游是 sub2api 即可应答)。
|
||||
const eligible = computed(() => props.account.type === 'apikey')
|
||||
const snapshot = computed<UpstreamBillingProbeSnapshot | undefined>(() => props.account.extra?.upstream_billing_probe)
|
||||
const data = computed(() => snapshot.value?.data)
|
||||
const probeEnabled = computed(() => props.account.extra?.upstream_billing_probe_enabled === true)
|
||||
|
||||
@@ -387,6 +387,23 @@ describe('BulkEditAccountModal', () => {
|
||||
})
|
||||
})
|
||||
|
||||
it('非 OpenAI 平台的 API Key 批量编辑同样可开启上游倍率自动探测', async () => {
|
||||
// 探测已放宽到全部 API-key 平台,混合平台选择只要求类型全为 apikey。
|
||||
const wrapper = mountModal({
|
||||
selectedPlatforms: ['grok', 'anthropic'],
|
||||
selectedTypes: ['apikey']
|
||||
})
|
||||
|
||||
await wrapper.get('#bulk-edit-upstream-billing-auto-probe-enabled').setValue(true)
|
||||
await wrapper.get('#bulk-edit-account-form').trigger('submit.prevent')
|
||||
await flushPromises()
|
||||
|
||||
expect(adminAPI.accounts.bulkUpdate).toHaveBeenCalledTimes(1)
|
||||
expect(adminAPI.accounts.bulkUpdate).toHaveBeenCalledWith([1, 2], {
|
||||
upstream_billing_probe_enabled: true
|
||||
})
|
||||
})
|
||||
|
||||
it('OpenAI API Key 批量编辑可统一关闭上游倍率自动探测', async () => {
|
||||
const wrapper = mountModal({
|
||||
selectedPlatforms: ['openai'],
|
||||
|
||||
@@ -254,7 +254,39 @@ describe('CreateAccountModal OpenAI long-context billing', () => {
|
||||
|
||||
expect(createAccountMock).toHaveBeenCalledTimes(1)
|
||||
expect(createAccountMock.mock.calls[0]?.[0]?.extra?.openai_long_context_billing_enabled).toBeUndefined()
|
||||
expect(createAccountMock.mock.calls[0]?.[0]?.upstream_billing_probe_enabled).toBeUndefined()
|
||||
// 上游倍率探测已放宽到全部 API-key 平台:非 OpenAI 平台与 OpenAI 一致,默认开启。
|
||||
expect(createAccountMock.mock.calls[0]?.[0]?.upstream_billing_probe_enabled).toBe(true)
|
||||
})
|
||||
|
||||
it('sends an explicit disabled state when the non-OpenAI create toggle is turned off', async () => {
|
||||
await submitApiKeyAccount('anthropic', false, true)
|
||||
|
||||
expect(createAccountMock.mock.calls[0]?.[0]?.upstream_billing_probe_enabled).toBe(false)
|
||||
})
|
||||
|
||||
it('antigravity upstream 创建默认携带上游倍率探测开关', async () => {
|
||||
// antigravity upstream 走独立创建 helper,
|
||||
// 也必须与其余 API-key 平台一样默认开启探测并传递开关。
|
||||
const wrapper = mountModal()
|
||||
await selectButtonByText(wrapper, 'Antigravity')
|
||||
await selectButtonByText(wrapper, 'admin.accounts.types.antigravityApikey')
|
||||
await wrapper.get('form#create-account-form input[type="text"]').setValue('antigravity relay')
|
||||
const baseInput = wrapper
|
||||
.findAll('input')
|
||||
.find((candidate) => candidate.attributes('placeholder') === 'https://cloudcode-pa.googleapis.com')
|
||||
expect(baseInput).toBeDefined()
|
||||
await baseInput?.setValue('https://relay.example')
|
||||
await wrapper.get('form#create-account-form input[type="password"]').setValue('sk-upstream')
|
||||
await wrapper.get('form#create-account-form').trigger('submit.prevent')
|
||||
await flushPromises()
|
||||
|
||||
expect(createAccountMock).toHaveBeenCalledTimes(1)
|
||||
const payload = createAccountMock.mock.calls[0]?.[0]
|
||||
expect(payload?.platform).toBe('antigravity')
|
||||
expect(payload?.type).toBe('apikey')
|
||||
expect(payload?.upstream_billing_probe_enabled).toBe(true)
|
||||
// 创建成功后前端立即发起一次首探(与其他 apikey 平台一致)。
|
||||
expect(probeUpstreamBillingMock).toHaveBeenCalledWith(42)
|
||||
})
|
||||
|
||||
it('leaves Codex session import billing ownership to the backend', async () => {
|
||||
|
||||
@@ -657,6 +657,28 @@ describe('EditAccountModal', () => {
|
||||
expect(updateAccountMock.mock.calls[0]?.[1]?.extra?.upstream_billing_probe_enabled).toBe(true)
|
||||
})
|
||||
|
||||
it('exposes the upstream billing auto-probe toggle for non-OpenAI API-key accounts', async () => {
|
||||
// 探测已放宽到全部 API-key 平台:grok 账号同样能开启并保存。
|
||||
const account = buildAccount()
|
||||
account.platform = 'grok'
|
||||
account.name = 'grok-relay'
|
||||
account.credentials = { api_key: 'sk-grok', base_url: 'https://relay.example/v1' }
|
||||
updateAccountMock.mockReset()
|
||||
checkMixedChannelRiskMock.mockReset()
|
||||
checkMixedChannelRiskMock.mockResolvedValue({ has_risk: false })
|
||||
updateAccountMock.mockResolvedValue(account)
|
||||
|
||||
const wrapper = mountModal(account)
|
||||
const toggle = wrapper.get('[data-testid="upstream-billing-auto-probe"]')
|
||||
expect(toggle.attributes('aria-checked')).toBe('false')
|
||||
|
||||
await toggle.trigger('click')
|
||||
await wrapper.get('form#edit-account-form').trigger('submit.prevent')
|
||||
|
||||
expect(updateAccountMock).toHaveBeenCalledTimes(1)
|
||||
expect(updateAccountMock.mock.calls[0]?.[1]?.extra?.upstream_billing_probe_enabled).toBe(true)
|
||||
})
|
||||
|
||||
it('clears OpenAI APIKey Responses override when set back to auto', async () => {
|
||||
const account = buildAccount()
|
||||
account.extra = {
|
||||
|
||||
@@ -268,6 +268,11 @@ describe('UpstreamBillingRateCell', () => {
|
||||
await wrapper.get('[data-testid="upstream-billing-probe"]').trigger('click')
|
||||
expect(wrapper.emitted('probe')).toHaveLength(1)
|
||||
|
||||
// 探测已放宽到全部 API-key 平台:grok API-key 账号同样可探测。
|
||||
await wrapper.setProps({ account: makeAccount({ platform: 'grok' }) })
|
||||
await wrapper.get('[data-testid="upstream-billing-probe"]').trigger('click')
|
||||
expect(wrapper.emitted('probe')).toHaveLength(2)
|
||||
|
||||
await wrapper.setProps({ account: makeAccount({ type: 'oauth' }) })
|
||||
expect(wrapper.findAll('button')).toHaveLength(0)
|
||||
expect(wrapper.text()).toBe('-')
|
||||
|
||||
@@ -250,7 +250,7 @@ export default {
|
||||
enabled: 'On',
|
||||
disabled: 'Off',
|
||||
probeFailed: 'Failed to probe upstream rate',
|
||||
noEligibleAccounts: 'Select OpenAI API key accounts',
|
||||
noEligibleAccounts: 'Select API key accounts',
|
||||
batchLimit: 'A batch can probe at most 20 accounts',
|
||||
batchCompleted: 'Probed {count} account(s)',
|
||||
batchPartial: 'Probe partially completed: {success} succeeded, {failed} failed'
|
||||
|
||||
@@ -361,7 +361,7 @@ export default {
|
||||
},
|
||||
upstreamBillingProbe: {
|
||||
title: 'Upstream Rate Auto Detection',
|
||||
description: 'Periodically retrieve billing rates declared by upstream Sub2API sites connected to OpenAI API keys.',
|
||||
description: 'Periodically retrieve billing rates declared by upstream Sub2API sites connected to API key accounts.',
|
||||
enabled: 'Enable global auto detection',
|
||||
enabledHint: 'When enabled, scheduled detection runs only for accounts that also enable automatic detection. Disabling stops all scheduled detection; manual detection remains available.',
|
||||
intervalMinutes: 'Detection interval (minutes)',
|
||||
|
||||
@@ -203,7 +203,7 @@ export default {
|
||||
enabled: '打开',
|
||||
disabled: '关闭',
|
||||
probeFailed: '探测上游倍率失败',
|
||||
noEligibleAccounts: '请选择 OpenAI API Key 账号',
|
||||
noEligibleAccounts: '请选择 API Key 账号',
|
||||
batchLimit: '每次最多探测 20 个账号',
|
||||
batchCompleted: '已完成 {count} 个账号的倍率探测',
|
||||
batchPartial: '倍率探测部分完成:成功 {success} 个,失败 {failed} 个'
|
||||
|
||||
@@ -354,7 +354,7 @@ export default {
|
||||
},
|
||||
upstreamBillingProbe: {
|
||||
title: '上游倍率自动探测',
|
||||
description: '定期获取 OpenAI API Key 所连接上游 Sub2API 站点声明的计费倍率。',
|
||||
description: '定期获取 API Key 账号所连接上游 Sub2API 站点声明的计费倍率。',
|
||||
enabled: '启用全局自动探测',
|
||||
enabledHint: '开启后,仅对账号自身已启用自动检测的账号执行定时探测;关闭后停止所有定时探测,手动探测不受影响。',
|
||||
intervalMinutes: '探测周期(分钟)',
|
||||
|
||||
Reference in New Issue
Block a user