feat(billing-probe): optionally sync account rate from upstream declared rate

Successful upstream billing probes already persist the upstream-declared
rate as a display-only snapshot. Add a per-account opt-in that writes
that declared rate back to the account's rate_multiplier, so the account
cost basis follows upstream repricing automatically instead of drifting
until an operator notices.

- new per-account flag upstream_billing_rate_sync_enabled stored next to
  the probe flag in account extra: enabling sync force-enables the
  probe, disabling the probe cascades sync off, and eligibility follows
  IsUpstreamBillingProbeIdentity (tightened from any non-empty platform
  to an explicit whitelist of the five supported API-key platforms so
  future platforms do not silently inherit probe/sync semantics)
- only a successful probe whose declared rate survives validation
  (finite, within bounds, not rounded to zero at the rate_multiplier
  decimal(10,4) scale) writes back; failed/unsupported/invalid probes
  leave rate_multiplier unchanged
- the writeback rides the existing snapshot CAS transaction:
  UpdateUpstreamBillingProbeSnapshot takes an optional rateMultiplier
  and applies it atomically with the snapshot under the same
  identity/snapshot compare-and-swap, so a probe result observed on a
  stale account cannot clobber a concurrent admin edit
- admin edit goes through UpdateWithAccountBillingSettings, which
  applies the form without overwriting a rate that a probe synchronized
  after the edit form was loaded (nil rateMultiplier = not edited);
  once sync is enabled the edit form shows the rate as managed
- bulk update rejects a manual rate_multiplier change when any target
  account has rate sync enabled (whole batch fails with a dedicated
  error so partial writes cannot bypass the sync ownership)
- frontend: sync toggle with hints in the edit modal (probe/sync
  enable/disable coupling enforced in the form), synced-rate tooltip on
  the rate cell, bulk edit modal warns and blocks rate edits that hit
  sync-enabled accounts; en/zh copy updated
- tests: service unit tests for sync gating/validation/cascade, sqlmock
  repo tests for the extended CAS, real-PostgreSQL integration tests
  (rate written only for successful+enabled accounts, manual rate
  protected after sync disabled, admin edit preserved across concurrent
  probe sync), handler/API contract updates, frontend specs for modal
  coupling, bulk rejection and rate cell
This commit is contained in:
Brisbanehuang
2026-08-01 22:11:09 +08:00
committed by shaw
parent 56f3d3c9b0
commit b0f5007f04
34 changed files with 1382 additions and 146 deletions
@@ -147,6 +147,8 @@ type UpdateAccountRequest struct {
GroupIDs *[]int64 `json:"group_ids"`
ExpiresAt *int64 `json:"expires_at"`
AutoPauseOnExpired *bool `json:"auto_pause_on_expired"`
ProbeEnabled *bool `json:"upstream_billing_probe_enabled"`
RateSyncEnabled *bool `json:"upstream_billing_rate_sync_enabled"`
ConfirmMixedChannelRisk *bool `json:"confirm_mixed_channel_risk"` // 用户确认混合渠道风险
}
@@ -987,6 +989,8 @@ func (h *AccountHandler) Update(c *gin.Context) {
GroupIDs: req.GroupIDs,
ExpiresAt: req.ExpiresAt,
AutoPauseOnExpired: req.AutoPauseOnExpired,
ProbeEnabled: req.ProbeEnabled,
RateSyncEnabled: req.RateSyncEnabled,
SkipMixedChannelCheck: skipCheck,
})
if err != nil {
@@ -147,6 +147,28 @@ func TestAccountHandlerUpdateMixedChannelConflictSimplifiedResponse(t *testing.T
require.False(t, hasRequireConfirmation)
}
func TestAccountHandlerUpdateMapsUpstreamBillingRateSyncSettings(t *testing.T) {
adminSvc := newStubAdminService()
router := setupAccountMixedChannelRouter(adminSvc)
body, _ := json.Marshal(map[string]any{
"name": "gemini-key",
"upstream_billing_probe_enabled": true,
"upstream_billing_rate_sync_enabled": true,
})
rec := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPut, "/api/v1/admin/accounts/42", bytes.NewReader(body))
req.Header.Set("Content-Type", "application/json")
router.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
require.NotNil(t, adminSvc.lastUpdateAccountInput)
require.NotNil(t, adminSvc.lastUpdateAccountInput.ProbeEnabled)
require.True(t, *adminSvc.lastUpdateAccountInput.ProbeEnabled)
require.NotNil(t, adminSvc.lastUpdateAccountInput.RateSyncEnabled)
require.True(t, *adminSvc.lastUpdateAccountInput.RateSyncEnabled)
}
func TestAccountHandlerBulkUpdateMixedChannelConflict(t *testing.T) {
adminSvc := newStubAdminService()
adminSvc.bulkUpdateAccountErr = &service.MixedChannelError{
@@ -32,6 +32,7 @@ type stubAdminService struct {
createAccountErr error
createSparkShadowErr error
updateAccountErr error
lastUpdateAccountInput *service.UpdateAccountInput
bulkUpdateAccountErr error
lastBulkUpdateAccountInput *service.BulkUpdateAccountsInput
getAccountResult *service.Account
@@ -501,6 +502,7 @@ func (s *stubAdminService) RecoverDuplicateAccount(ctx context.Context, id int64
func (s *stubAdminService) UpdateAccount(ctx context.Context, id int64, input *service.UpdateAccountInput) (*service.Account, error) {
s.updateAccountCalls++
s.lastUpdateAccountInput = input
if s.updateAccountErr != nil {
return nil, s.updateAccountErr
}
+127 -24
View File
@@ -58,6 +58,7 @@ var schedulerNeutralExtraKeyPrefixes = []string{
"codex_7d_",
"passive_usage_",
"upstream_billing_probe",
"upstream_billing_rate_sync",
"ollama_cloud_usage",
}
@@ -397,16 +398,40 @@ func (r *accountRepository) ListCRSAccountIDs(ctx context.Context) (map[string]i
}
func (r *accountRepository) Update(ctx context.Context, account *service.Account) error {
return r.updateAccount(ctx, account, nil)
return r.updateAccount(ctx, account, nil, nil, account.RateMultiplier)
}
// UpdateWithUpstreamBillingProbeEnabled applies an explicit probe switch in the
// same row-lock transaction as the rest of an admin account edit.
func (r *accountRepository) UpdateWithUpstreamBillingProbeEnabled(ctx context.Context, account *service.Account, enabled bool) error {
return r.updateAccount(ctx, account, &enabled)
var rateSyncEnabled *bool
if !enabled {
disabled := false
rateSyncEnabled = &disabled
}
return r.updateAccount(ctx, account, &enabled, rateSyncEnabled, nil)
}
func (r *accountRepository) updateAccount(ctx context.Context, account *service.Account, explicitProbeEnabled *bool) error {
// UpdateWithAccountBillingSettings applies an admin account edit while
// preserving a concurrently probe-synchronized rate unless the request
// explicitly includes a manual rate.
func (r *accountRepository) UpdateWithAccountBillingSettings(
ctx context.Context,
account *service.Account,
probeEnabled *bool,
rateSyncEnabled *bool,
rateMultiplier *float64,
) error {
return r.updateAccount(ctx, account, probeEnabled, rateSyncEnabled, rateMultiplier)
}
func (r *accountRepository) updateAccount(
ctx context.Context,
account *service.Account,
explicitProbeEnabled *bool,
explicitRateSyncEnabled *bool,
explicitRateMultiplier *float64,
) error {
if account == nil {
return nil
}
@@ -430,7 +455,14 @@ func (r *accountRepository) updateAccount(ctx context.Context, account *service.
}
}
updated, err := r.updateLockedAccount(ctx, client, account, explicitProbeEnabled)
updated, err := r.updateLockedAccount(
ctx,
client,
account,
explicitProbeEnabled,
explicitRateSyncEnabled,
explicitRateMultiplier,
)
if err != nil {
return translatePersistenceError(err, service.ErrAccountNotFound, nil)
}
@@ -452,8 +484,15 @@ func (r *accountRepository) updateAccount(ctx context.Context, account *service.
return nil
}
func (r *accountRepository) updateLockedAccount(ctx context.Context, client *dbent.Client, account *service.Account, explicitProbeEnabled *bool) (*dbent.Account, error) {
extra, err := lockAndMergeAccountProbeExtra(ctx, client, account, explicitProbeEnabled)
func (r *accountRepository) updateLockedAccount(
ctx context.Context,
client *dbent.Client,
account *service.Account,
explicitProbeEnabled *bool,
explicitRateSyncEnabled *bool,
explicitRateMultiplier *float64,
) (*dbent.Account, error) {
extra, err := lockAndMergeAccountProbeExtra(ctx, client, account, explicitProbeEnabled, explicitRateSyncEnabled)
if err != nil {
return nil, err
}
@@ -478,8 +517,8 @@ func (r *accountRepository) updateLockedAccount(ctx context.Context, client *dbe
SetSchedulable(schedulable).
SetAutoPauseOnExpired(account.AutoPauseOnExpired)
if account.RateMultiplier != nil {
builder.SetRateMultiplier(*account.RateMultiplier)
if explicitRateMultiplier != nil {
builder.SetRateMultiplier(*explicitRateMultiplier)
}
if account.LoadFactor != nil {
builder.SetLoadFactor(*account.LoadFactor)
@@ -542,7 +581,13 @@ func (r *accountRepository) updateLockedAccount(ctx context.Context, client *dbe
return builder.Save(ctx)
}
func lockAndMergeAccountProbeExtra(ctx context.Context, client *dbent.Client, account *service.Account, explicitProbeEnabled *bool) (map[string]any, error) {
func lockAndMergeAccountProbeExtra(
ctx context.Context,
client *dbent.Client,
account *service.Account,
explicitProbeEnabled *bool,
explicitRateSyncEnabled *bool,
) (map[string]any, error) {
credentials, err := json.Marshal(normalizeJSONMap(account.Credentials))
if err != nil {
return nil, err
@@ -569,6 +614,7 @@ func lockAndMergeAccountProbeExtra(ctx context.Context, client *dbent.Client, ac
),
proxy_id IS NOT DISTINCT FROM $5,
extra -> 'upstream_billing_probe_enabled',
extra -> 'upstream_billing_rate_sync_enabled',
extra -> 'upstream_billing_probe',
extra -> 'ollama_cloud_usage_session',
extra -> 'ollama_cloud_usage_auto_refresh',
@@ -593,6 +639,7 @@ func lockAndMergeAccountProbeExtra(ctx context.Context, client *dbent.Client, ac
ollamaGroupIdentityUnchanged bool
ollamaProxyIdentityUnchanged bool
currentEnabled []byte
currentRateSyncEnabled []byte
currentSnapshot []byte
currentOllamaSession []byte
currentOllamaAutoRefresh []byte
@@ -603,6 +650,7 @@ func lockAndMergeAccountProbeExtra(ctx context.Context, client *dbent.Client, ac
&ollamaGroupIdentityUnchanged,
&ollamaProxyIdentityUnchanged,
&currentEnabled,
&currentRateSyncEnabled,
&currentSnapshot,
&currentOllamaSession,
&currentOllamaAutoRefresh,
@@ -617,6 +665,7 @@ func lockAndMergeAccountProbeExtra(ctx context.Context, client *dbent.Client, ac
extra := copyJSONMap(normalizeJSONMap(account.Extra))
for _, key := range []string{
service.UpstreamBillingProbeEnabledExtraKey,
service.UpstreamBillingRateSyncEnabledExtraKey,
service.UpstreamBillingProbeExtraKey,
service.OllamaCloudUsageSessionExtraKey,
service.OllamaCloudUsageAutoRefreshExtraKey,
@@ -624,21 +673,52 @@ func lockAndMergeAccountProbeExtra(ctx context.Context, client *dbent.Client, ac
} {
delete(extra, key)
}
probeExplicitlyDisabled := false
probeAccount := service.IsUpstreamBillingProbeIdentity(account.Platform, account.Type)
if probeAccount && explicitProbeEnabled != nil {
extra[service.UpstreamBillingProbeEnabledExtraKey] = *explicitProbeEnabled
probeExplicitlyDisabled = !*explicitProbeEnabled
} else if probeAccount {
probeEnabled := false
probeEnabledPresent := false
if probeAccount {
if enabled, ok, err := decodeAccountExtraJSON(currentEnabled); err != nil {
return nil, err
} else if ok {
extra[service.UpstreamBillingProbeEnabledExtraKey] = enabled
if value, isBool := enabled.(bool); isBool && !value {
probeExplicitlyDisabled = true
}
} else if value, isBool := enabled.(bool); ok && isBool {
probeEnabled = value
probeEnabledPresent = true
}
if explicitProbeEnabled != nil {
probeEnabled = *explicitProbeEnabled
probeEnabledPresent = true
}
}
rateSyncEnabled := false
rateSyncEnabledPresent := false
if probeAccount {
if enabled, ok, err := decodeAccountExtraJSON(currentRateSyncEnabled); err != nil {
return nil, err
} else if value, isBool := enabled.(bool); ok && isBool {
rateSyncEnabled = value
rateSyncEnabledPresent = true
}
if explicitRateSyncEnabled != nil {
rateSyncEnabled = *explicitRateSyncEnabled
rateSyncEnabledPresent = true
}
if explicitProbeEnabled != nil && !*explicitProbeEnabled {
rateSyncEnabled = false
rateSyncEnabledPresent = true
} else if rateSyncEnabled {
probeEnabled = true
probeEnabledPresent = true
}
if !probeEnabled {
rateSyncEnabled = false
}
if probeEnabledPresent {
extra[service.UpstreamBillingProbeEnabledExtraKey] = probeEnabled
}
if rateSyncEnabledPresent {
extra[service.UpstreamBillingRateSyncEnabledExtraKey] = rateSyncEnabled
}
}
probeExplicitlyDisabled := probeEnabledPresent && !probeEnabled
if identityUnchanged && !probeExplicitlyDisabled {
if snapshot, ok, err := decodeAccountExtraJSON(currentSnapshot); err != nil {
return nil, err
@@ -2527,21 +2607,25 @@ func (r *accountRepository) UpdateUpstreamBillingProbeSnapshot(
ctx context.Context,
account *service.Account,
snapshot *service.UpstreamBillingProbeSnapshot,
rateMultiplier *float64,
) error {
if account == nil || snapshot == nil {
return service.ErrAccountNilInput
}
if snapshot.Status != service.UpstreamBillingProbeStatusOK {
rateMultiplier = nil
}
if dbent.TxFromContext(ctx) == nil {
tx, err := r.client.Tx(ctx)
if errors.Is(err, dbent.ErrTxStarted) {
return r.updateUpstreamBillingProbeSnapshotInTx(ctx, account, snapshot)
return r.updateUpstreamBillingProbeSnapshotInTx(ctx, account, snapshot, rateMultiplier)
}
if err != nil {
return err
}
defer func() { _ = tx.Rollback() }()
if err := r.updateUpstreamBillingProbeSnapshotInTx(dbent.NewTxContext(ctx, tx), account, snapshot); err != nil {
if err := r.updateUpstreamBillingProbeSnapshotInTx(dbent.NewTxContext(ctx, tx), account, snapshot, rateMultiplier); err != nil {
return err
}
if err := tx.Commit(); err != nil {
@@ -2552,13 +2636,14 @@ func (r *accountRepository) UpdateUpstreamBillingProbeSnapshot(
r.syncSchedulerAccountSnapshot(ctx, account.ID)
return nil
}
return r.updateUpstreamBillingProbeSnapshotInTx(ctx, account, snapshot)
return r.updateUpstreamBillingProbeSnapshotInTx(ctx, account, snapshot, rateMultiplier)
}
func (r *accountRepository) updateUpstreamBillingProbeSnapshotInTx(
ctx context.Context,
account *service.Account,
snapshot *service.UpstreamBillingProbeSnapshot,
rateMultiplier *float64,
) error {
payload, err := json.Marshal(map[string]any{service.UpstreamBillingProbeExtraKey: snapshot})
if err != nil {
@@ -2584,6 +2669,14 @@ func (r *accountRepository) updateUpstreamBillingProbeSnapshotInTx(
if err != nil {
return err
}
var expectedRateSyncEnabled any
if account.Extra != nil {
expectedRateSyncEnabled = account.Extra[service.UpstreamBillingRateSyncEnabledExtraKey]
}
expectedRateSyncEnabledJSON, err := json.Marshal(expectedRateSyncEnabled)
if err != nil {
return err
}
client := clientFromContext(ctx, r.client)
proxyMatches, err := lockAndMatchProbeProxyIdentity(ctx, client, account)
if err != nil {
@@ -2598,7 +2691,16 @@ func (r *accountRepository) updateUpstreamBillingProbeSnapshotInTx(
}
result, err := client.ExecContext(ctx, `
UPDATE accounts
SET extra = COALESCE(extra, '{}'::jsonb) || $1::jsonb, updated_at = NOW()
SET
extra = COALESCE(extra, '{}'::jsonb) || $1::jsonb,
rate_multiplier = CASE
WHEN $10::numeric IS NOT NULL
AND extra @> '{"upstream_billing_probe_enabled": true}'::jsonb
AND extra @> '{"upstream_billing_rate_sync_enabled": true}'::jsonb
THEN $10::numeric
ELSE rate_multiplier
END,
updated_at = NOW()
WHERE id = $2
AND platform = $3
AND type = $4
@@ -2606,8 +2708,9 @@ func (r *accountRepository) updateUpstreamBillingProbeSnapshotInTx(
AND proxy_id IS NOT DISTINCT FROM $6
AND COALESCE(extra -> 'upstream_billing_probe', 'null'::jsonb) = $7::jsonb
AND COALESCE(extra -> 'upstream_billing_probe_enabled', 'null'::jsonb) = $8::jsonb
AND COALESCE(extra -> 'upstream_billing_rate_sync_enabled', 'null'::jsonb) = $9::jsonb
AND deleted_at IS NULL
`, string(payload), account.ID, account.Platform, account.Type, string(credentials), proxyID, string(expectedSnapshotJSON), string(expectedEnabledJSON))
`, string(payload), account.ID, account.Platform, account.Type, string(credentials), proxyID, string(expectedSnapshotJSON), string(expectedEnabledJSON), string(expectedRateSyncEnabledJSON), rateMultiplier)
if err != nil {
return err
}
@@ -229,7 +229,7 @@ func TestLockAndMergeAccountProbeExtraCoalescesNullableOllamaGroupIdentity(t *te
loaded, err := newAccountRepositoryWithSQL(tx.Client(), tx, nil).GetByID(ctx, account.ID)
require.NoError(t, err)
merged, err := lockAndMergeAccountProbeExtra(ctx, tx.Client(), loaded, nil)
merged, err := lockAndMergeAccountProbeExtra(ctx, tx.Client(), loaded, nil, nil)
require.NoError(t, err, "a NULL Ollama eligibility expression must scan as false")
require.NotContains(t, merged, service.OllamaCloudUsageSessionExtraKey)
@@ -41,8 +41,8 @@ func TestUpdateUpstreamBillingProbeSnapshotRequiresSameIdentityAndSnapshot(t *te
WithArgs(int64(9)).
WillReturnRows(sqlmock.NewRows([]string{"protocol", "host", "port", "username", "password", "status"}).
AddRow("http", "127.0.0.1", 3128, "user", "pass", service.StatusActive))
mock.ExpectExec(`(?s)`+regexp.QuoteMeta("UPDATE accounts")+`.*`+regexp.QuoteMeta("WHERE id = $2")+`.*`+regexp.QuoteMeta("AND platform = $3")+`.*`+regexp.QuoteMeta("AND type = $4")+`.*`+regexp.QuoteMeta("AND credentials = $5::jsonb")+`.*`+regexp.QuoteMeta("AND proxy_id IS NOT DISTINCT FROM $6")+`.*`+regexp.QuoteMeta("COALESCE(extra -> 'upstream_billing_probe', 'null'::jsonb) = $7::jsonb")+`.*`+regexp.QuoteMeta("COALESCE(extra -> 'upstream_billing_probe_enabled', 'null'::jsonb) = $8::jsonb")).
WithArgs(sqlmock.AnyArg(), int64(17), service.PlatformOpenAI, service.AccountTypeAPIKey, `{"api_key":"sk-test","base_url":"http://127.0.0.1:8080"}`, int64(9), `{"status":"stale"}`, "null").
mock.ExpectExec(`(?s)`+regexp.QuoteMeta("UPDATE accounts")+`.*`+regexp.QuoteMeta("WHERE id = $2")+`.*`+regexp.QuoteMeta("AND platform = $3")+`.*`+regexp.QuoteMeta("AND type = $4")+`.*`+regexp.QuoteMeta("AND credentials = $5::jsonb")+`.*`+regexp.QuoteMeta("AND proxy_id IS NOT DISTINCT FROM $6")+`.*`+regexp.QuoteMeta("COALESCE(extra -> 'upstream_billing_probe', 'null'::jsonb) = $7::jsonb")+`.*`+regexp.QuoteMeta("COALESCE(extra -> 'upstream_billing_probe_enabled', 'null'::jsonb) = $8::jsonb")+`.*`+regexp.QuoteMeta("COALESCE(extra -> 'upstream_billing_rate_sync_enabled', 'null'::jsonb) = $9::jsonb")).
WithArgs(sqlmock.AnyArg(), int64(17), service.PlatformOpenAI, service.AccountTypeAPIKey, `{"api_key":"sk-test","base_url":"http://127.0.0.1:8080"}`, int64(9), `{"status":"stale"}`, "null", "null", nil).
WillReturnResult(sqlmock.NewResult(0, tt.affected))
if tt.affected > 0 {
mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox")).
@@ -75,7 +75,7 @@ func TestUpdateUpstreamBillingProbeSnapshotRequiresSameIdentityAndSnapshot(t *te
}
txCtx := dbent.NewTxContext(context.Background(), tx)
err = repo.UpdateUpstreamBillingProbeSnapshot(txCtx, account, &service.UpstreamBillingProbeSnapshot{Status: service.UpstreamBillingProbeStatusOK})
err = repo.UpdateUpstreamBillingProbeSnapshot(txCtx, account, &service.UpstreamBillingProbeSnapshot{Status: service.UpstreamBillingProbeStatusOK}, nil)
if tt.wantErr != nil {
require.ErrorIs(t, err, tt.wantErr)
@@ -98,8 +98,8 @@ func TestUpdateUpstreamBillingProbeSnapshotCommitsSnapshotAndOutboxAtomically(t
t.Cleanup(func() { _ = client.Close() })
mock.ExpectBegin()
mock.ExpectExec(`(?s)`+regexp.QuoteMeta("UPDATE accounts")+`.*`+regexp.QuoteMeta("AND credentials = $5::jsonb")+`.*`+regexp.QuoteMeta("AND proxy_id IS NOT DISTINCT FROM $6")+`.*`+regexp.QuoteMeta("COALESCE(extra -> 'upstream_billing_probe', 'null'::jsonb) = $7::jsonb")).
WithArgs(sqlmock.AnyArg(), int64(17), service.PlatformOpenAI, service.AccountTypeAPIKey, `{"api_key":"sk-test"}`, nil, "null", "null").
mock.ExpectExec(`(?s)`+regexp.QuoteMeta("UPDATE accounts")+`.*`+regexp.QuoteMeta("rate_multiplier = CASE")+`.*`+regexp.QuoteMeta("THEN $10::numeric")+`.*`+regexp.QuoteMeta("AND credentials = $5::jsonb")+`.*`+regexp.QuoteMeta("AND proxy_id IS NOT DISTINCT FROM $6")+`.*`+regexp.QuoteMeta("COALESCE(extra -> 'upstream_billing_probe', 'null'::jsonb) = $7::jsonb")).
WithArgs(sqlmock.AnyArg(), int64(17), service.PlatformOpenAI, service.AccountTypeAPIKey, `{"api_key":"sk-test"}`, nil, "null", "true", "true", 0.065).
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox")).
WithArgs(service.SchedulerOutboxEventAccountChanged, int64(17), nil, nil, sqlmock.AnyArg()).
@@ -112,9 +112,19 @@ func TestUpdateUpstreamBillingProbeSnapshotCommitsSnapshotAndOutboxAtomically(t
Platform: service.PlatformOpenAI,
Type: service.AccountTypeAPIKey,
Credentials: map[string]any{"api_key": "sk-test"},
Extra: map[string]any{
service.UpstreamBillingProbeEnabledExtraKey: true,
service.UpstreamBillingRateSyncEnabledExtraKey: true,
},
}
rateMultiplier := 0.065
err = repo.UpdateUpstreamBillingProbeSnapshot(context.Background(), account, &service.UpstreamBillingProbeSnapshot{Status: service.UpstreamBillingProbeStatusOK})
err = repo.UpdateUpstreamBillingProbeSnapshot(
context.Background(),
account,
&service.UpstreamBillingProbeSnapshot{Status: service.UpstreamBillingProbeStatusOK},
&rateMultiplier,
)
require.NoError(t, err)
require.NoError(t, mock.ExpectationsWereMet())
@@ -148,7 +158,7 @@ func TestUpdateUpstreamBillingProbeSnapshotRejectsChangedProxyIdentity(t *testin
},
}
repo := newAccountRepositoryWithSQL(client, db, nil)
err = repo.UpdateUpstreamBillingProbeSnapshot(dbent.NewTxContext(context.Background(), tx), account, &service.UpstreamBillingProbeSnapshot{Status: service.UpstreamBillingProbeStatusOK})
err = repo.UpdateUpstreamBillingProbeSnapshot(dbent.NewTxContext(context.Background(), tx), account, &service.UpstreamBillingProbeSnapshot{Status: service.UpstreamBillingProbeStatusOK}, nil)
require.ErrorIs(t, err, service.ErrUpstreamBillingProbeIdentityChanged)
mock.ExpectRollback()
@@ -166,7 +176,7 @@ func TestUpdateUpstreamBillingProbeSnapshotRollsBackWhenOutboxFails(t *testing.T
mock.ExpectBegin()
mock.ExpectExec(`(?s)`+regexp.QuoteMeta("UPDATE accounts")+`.*`+regexp.QuoteMeta("AND proxy_id IS NOT DISTINCT FROM $6")+`.*`+regexp.QuoteMeta("COALESCE(extra -> 'upstream_billing_probe', 'null'::jsonb) = $7::jsonb")).
WithArgs(sqlmock.AnyArg(), int64(18), service.PlatformOpenAI, service.AccountTypeAPIKey, `{"api_key":"sk-test"}`, nil, "null", "null").
WithArgs(sqlmock.AnyArg(), int64(18), service.PlatformOpenAI, service.AccountTypeAPIKey, `{"api_key":"sk-test"}`, nil, "null", "true", "true", 0.7).
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox")).WillReturnError(errors.New("outbox failed"))
mock.ExpectRollback()
@@ -177,9 +187,19 @@ func TestUpdateUpstreamBillingProbeSnapshotRollsBackWhenOutboxFails(t *testing.T
Platform: service.PlatformOpenAI,
Type: service.AccountTypeAPIKey,
Credentials: map[string]any{"api_key": "sk-test"},
Extra: map[string]any{
service.UpstreamBillingProbeEnabledExtraKey: true,
service.UpstreamBillingRateSyncEnabledExtraKey: true,
},
}
rateMultiplier := 0.7
err = repo.UpdateUpstreamBillingProbeSnapshot(context.Background(), account, &service.UpstreamBillingProbeSnapshot{Status: service.UpstreamBillingProbeStatusOK})
err = repo.UpdateUpstreamBillingProbeSnapshot(
context.Background(),
account,
&service.UpstreamBillingProbeSnapshot{Status: service.UpstreamBillingProbeStatusOK},
&rateMultiplier,
)
require.EqualError(t, err, "outbox failed")
require.NoError(t, mock.ExpectationsWereMet())
@@ -81,8 +81,8 @@ func TestLockAndMergeAccountProbeExtraUsesCurrentDatabaseSnapshot(t *testing.T)
mock.ExpectQuery(`(?s)`+regexp.QuoteMeta("SELECT")+`.*`+regexp.QuoteMeta("FOR NO KEY UPDATE")).
WithArgs(int64(27), service.PlatformOpenAI, service.AccountTypeAPIKey, `{"api_key":"sk-test"}`, nil).
WillReturnRows(sqlmock.NewRows([]string{"identity_unchanged", "ollama_group_unchanged", "ollama_proxy_unchanged", "enabled", "snapshot", "ollama_session", "ollama_auto", "ollama_snapshot"}).
AddRow(tt.identityUnchanged, false, true, tt.databaseEnabled, tt.databaseSnapshot, nil, nil, nil))
WillReturnRows(sqlmock.NewRows([]string{"identity_unchanged", "ollama_group_unchanged", "ollama_proxy_unchanged", "enabled", "rate_sync_enabled", "snapshot", "ollama_session", "ollama_auto", "ollama_snapshot"}).
AddRow(tt.identityUnchanged, false, true, tt.databaseEnabled, nil, tt.databaseSnapshot, nil, nil, nil))
account := &service.Account{
ID: 27,
@@ -91,7 +91,7 @@ func TestLockAndMergeAccountProbeExtraUsesCurrentDatabaseSnapshot(t *testing.T)
Credentials: map[string]any{"api_key": "sk-test"},
Extra: tt.inputExtra,
}
got, err := lockAndMergeAccountProbeExtra(context.Background(), client, account, nil)
got, err := lockAndMergeAccountProbeExtra(context.Background(), client, account, nil, nil)
require.NoError(t, err)
if tt.wantSnapshot == nil {
require.NotContains(t, got, service.UpstreamBillingProbeExtraKey)
@@ -115,8 +115,8 @@ func TestLockAndMergeAccountProbeExtraProtectsOllamaManagedFields(t *testing.T)
mock.ExpectQuery(`(?s)`+regexp.QuoteMeta("SELECT")+`.*`+regexp.QuoteMeta("FOR NO KEY UPDATE")).
WithArgs(int64(29), service.PlatformAnthropic, service.AccountTypeAPIKey, `{"api_key":"key","base_url":"https://ollama.com"}`, nil).
WillReturnRows(sqlmock.NewRows([]string{"identity_unchanged", "ollama_group_unchanged", "ollama_proxy_unchanged", "enabled", "snapshot", "ollama_session", "ollama_auto", "ollama_snapshot"}).
AddRow(identityUnchanged, identityUnchanged, true, nil, nil, []byte(`"local-ciphertext"`), []byte(`true`), []byte(`{"status":"ok"}`)))
WillReturnRows(sqlmock.NewRows([]string{"identity_unchanged", "ollama_group_unchanged", "ollama_proxy_unchanged", "enabled", "rate_sync_enabled", "snapshot", "ollama_session", "ollama_auto", "ollama_snapshot"}).
AddRow(identityUnchanged, identityUnchanged, true, nil, nil, nil, []byte(`"local-ciphertext"`), []byte(`true`), []byte(`{"status":"ok"}`)))
account := &service.Account{
ID: 29, Platform: service.PlatformAnthropic, Type: service.AccountTypeAPIKey,
@@ -127,7 +127,7 @@ func TestLockAndMergeAccountProbeExtraProtectsOllamaManagedFields(t *testing.T)
service.OllamaCloudUsageSnapshotExtraKey: map[string]any{"status": "forged"},
},
}
got, err := lockAndMergeAccountProbeExtra(context.Background(), client, account, nil)
got, err := lockAndMergeAccountProbeExtra(context.Background(), client, account, nil, nil)
require.NoError(t, err)
if identityUnchanged {
require.Equal(t, "local-ciphertext", got[service.OllamaCloudUsageSessionExtraKey])
@@ -275,13 +275,13 @@ func TestUpdateWithUpstreamBillingProbeEnabledRollsBackWhenOutboxFails(t *testin
mock.ExpectBegin()
mock.ExpectQuery(`(?s)`+regexp.QuoteMeta("SELECT")+`.*`+regexp.QuoteMeta("FOR NO KEY UPDATE")).
WithArgs(int64(27), service.PlatformOpenAI, service.AccountTypeAPIKey, `{"api_key":"sk-test"}`, nil).
WillReturnRows(sqlmock.NewRows([]string{"identity_unchanged", "ollama_group_unchanged", "ollama_proxy_unchanged", "enabled", "snapshot", "ollama_session", "ollama_auto", "ollama_snapshot"}).
AddRow(true, false, true, []byte(`true`), []byte(`{"status":"ok"}`), nil, nil, nil))
WillReturnRows(sqlmock.NewRows([]string{"identity_unchanged", "ollama_group_unchanged", "ollama_proxy_unchanged", "enabled", "rate_sync_enabled", "snapshot", "ollama_session", "ollama_auto", "ollama_snapshot"}).
AddRow(true, false, true, []byte(`true`), []byte(`true`), []byte(`{"status":"ok"}`), nil, nil, nil))
mock.ExpectExec(`(?s)UPDATE .*accounts.*SET.*WHERE .*id.*`).
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectQuery(`(?s)SELECT .* FROM "accounts" WHERE "id" = \$1`).
WithArgs(int64(27)).
WillReturnRows(updatedAccountRows(27, `{"upstream_billing_probe_enabled":false}`))
WillReturnRows(updatedAccountRows(27, `{"upstream_billing_probe_enabled":false,"upstream_billing_rate_sync_enabled":false}`))
mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox")).WillReturnError(errors.New("outbox failed"))
mock.ExpectRollback()
@@ -305,6 +305,7 @@ func TestUpdateWithUpstreamBillingProbeEnabledRollsBackWhenOutboxFails(t *testin
require.EqualError(t, err, "outbox failed")
require.Equal(t, false, account.Extra[service.UpstreamBillingProbeEnabledExtraKey])
require.Equal(t, false, account.Extra[service.UpstreamBillingRateSyncEnabledExtraKey])
require.NotContains(t, account.Extra, service.UpstreamBillingProbeExtraKey)
require.NoError(t, mock.ExpectationsWereMet())
}
@@ -30,7 +30,7 @@ func TestAccountUpdatePreservesConcurrentProbeSnapshot(t *testing.T) {
require.NoError(t, repo.UpdateUpstreamBillingProbeSnapshot(ctx, stale, &service.UpstreamBillingProbeSnapshot{
Status: service.UpstreamBillingProbeStatusOK,
LastAttemptAt: time.Now().UTC(),
}))
}, nil))
stale.Name = "ordinary-edit"
require.NoError(t, repo.Update(ctx, stale))
@@ -46,6 +46,115 @@ func TestAccountUpdatePreservesConcurrentProbeSnapshot(t *testing.T) {
require.NotContains(t, disabled.Extra, service.UpstreamBillingProbeExtraKey)
}
func TestAdminAccountEditPreservesRateSynchronizedAfterLoad(t *testing.T) {
ctx := context.Background()
tx := testEntTx(t)
repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
initialRate := 0.1
account := mustCreateAccount(t, tx.Client(), &service.Account{
Name: "probe-rate-concurrent-edit",
Platform: service.PlatformOpenAI,
Type: service.AccountTypeAPIKey,
RateMultiplier: &initialRate,
Credentials: map[string]any{"api_key": "sk-test"},
Extra: map[string]any{
service.UpstreamBillingProbeEnabledExtraKey: true,
service.UpstreamBillingRateSyncEnabledExtraKey: true,
},
})
staleAdminEdit, err := repo.GetByID(ctx, account.ID)
require.NoError(t, err)
probeAccount, err := repo.GetByID(ctx, account.ID)
require.NoError(t, err)
synchronizedRate := 0.2
require.NoError(t, repo.UpdateUpstreamBillingProbeSnapshot(ctx, probeAccount, &service.UpstreamBillingProbeSnapshot{
Status: service.UpstreamBillingProbeStatusOK,
LastAttemptAt: time.Now().UTC(),
}, &synchronizedRate))
staleAdminEdit.Name = "name-only-edit"
require.NoError(t, repo.UpdateWithAccountBillingSettings(ctx, staleAdminEdit, nil, nil, nil))
got, err := repo.GetByID(ctx, account.ID)
require.NoError(t, err)
require.Equal(t, "name-only-edit", got.Name)
require.NotNil(t, got.RateMultiplier)
require.Equal(t, synchronizedRate, *got.RateMultiplier)
}
func TestProbeSnapshotSyncsRateOnlyForSuccessfulEnabledAccount(t *testing.T) {
ctx := context.Background()
tx := testEntTx(t)
repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
initialRate := 0.25
account := mustCreateAccount(t, tx.Client(), &service.Account{
Name: "probe-rate-sync",
Platform: service.PlatformGemini,
Type: service.AccountTypeAPIKey,
RateMultiplier: &initialRate,
Credentials: map[string]any{"api_key": "sk-test"},
Extra: map[string]any{
service.UpstreamBillingProbeEnabledExtraKey: true,
service.UpstreamBillingRateSyncEnabledExtraKey: true,
},
})
loaded, err := repo.GetByID(ctx, account.ID)
require.NoError(t, err)
syncedRate := 0.065
require.NoError(t, repo.UpdateUpstreamBillingProbeSnapshot(ctx, loaded, &service.UpstreamBillingProbeSnapshot{
Status: service.UpstreamBillingProbeStatusOK,
LastAttemptAt: time.Now().UTC(),
}, &syncedRate))
got, err := repo.GetByID(ctx, account.ID)
require.NoError(t, err)
require.NotNil(t, got.RateMultiplier)
require.Equal(t, syncedRate, *got.RateMultiplier)
failedRate := 0.9
require.NoError(t, repo.UpdateUpstreamBillingProbeSnapshot(ctx, got, &service.UpstreamBillingProbeSnapshot{
Status: service.UpstreamBillingProbeStatusFailed,
LastAttemptAt: time.Now().UTC(),
}, &failedRate))
got, err = repo.GetByID(ctx, account.ID)
require.NoError(t, err)
require.NotNil(t, got.RateMultiplier)
require.Equal(t, syncedRate, *got.RateMultiplier)
require.NoError(t, repo.UpdateExtra(ctx, account.ID, map[string]any{
service.UpstreamBillingRateSyncEnabledExtraKey: false,
}))
manual, err := repo.GetByID(ctx, account.ID)
require.NoError(t, err)
manualProbeRate := 0.4
require.NoError(t, repo.UpdateUpstreamBillingProbeSnapshot(ctx, manual, &service.UpstreamBillingProbeSnapshot{
Status: service.UpstreamBillingProbeStatusOK,
LastAttemptAt: time.Now().UTC(),
}, &manualProbeRate))
got, err = repo.GetByID(ctx, account.ID)
require.NoError(t, err)
require.NotNil(t, got.RateMultiplier)
require.Equal(t, syncedRate, *got.RateMultiplier)
require.NoError(t, repo.UpdateExtra(ctx, account.ID, map[string]any{
service.UpstreamBillingProbeEnabledExtraKey: false,
service.UpstreamBillingRateSyncEnabledExtraKey: false,
}))
disabled, err := repo.GetByID(ctx, account.ID)
require.NoError(t, err)
require.NoError(t, repo.UpdateUpstreamBillingProbeSnapshot(ctx, disabled, &service.UpstreamBillingProbeSnapshot{
Status: service.UpstreamBillingProbeStatusOK,
LastAttemptAt: time.Now().UTC(),
}, &manualProbeRate))
got, err = repo.GetByID(ctx, account.ID)
require.NoError(t, err)
require.NotNil(t, got.RateMultiplier)
require.Equal(t, syncedRate, *got.RateMultiplier)
}
func TestAccountUpdatePreservesConcurrentProbeEnableFlag(t *testing.T) {
ctx := context.Background()
tx := testEntTx(t)
@@ -166,7 +275,7 @@ func TestProbeSnapshotCASIncludesLoadedEnabledState(t *testing.T) {
err = repo.UpdateUpstreamBillingProbeSnapshot(ctx, inFlight, &service.UpstreamBillingProbeSnapshot{
Status: service.UpstreamBillingProbeStatusOK,
LastAttemptAt: time.Now().UTC(),
})
}, nil)
if tt.wantConflict {
require.ErrorIs(t, err, service.ErrUpstreamBillingProbeIdentityChanged)
} else {
@@ -183,6 +292,46 @@ func TestProbeSnapshotCASIncludesLoadedEnabledState(t *testing.T) {
}
}
func TestProbeSnapshotCASProtectsManualRateAfterSyncDisabled(t *testing.T) {
ctx := context.Background()
tx := testEntTx(t)
repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
initialRate := 0.25
account := mustCreateAccount(t, tx.Client(), &service.Account{
Name: "probe-sync-cas",
Platform: service.PlatformAnthropic,
Type: service.AccountTypeAPIKey,
RateMultiplier: &initialRate,
Credentials: map[string]any{"api_key": "sk-test"},
Extra: map[string]any{
service.UpstreamBillingProbeEnabledExtraKey: true,
service.UpstreamBillingRateSyncEnabledExtraKey: true,
},
})
inFlight, err := repo.GetByID(ctx, account.ID)
require.NoError(t, err)
manual, err := repo.GetByID(ctx, account.ID)
require.NoError(t, err)
manualRate := 0.8
syncDisabled := false
require.NoError(t, repo.UpdateWithAccountBillingSettings(ctx, manual, nil, &syncDisabled, &manualRate))
probedRate := 0.1
err = repo.UpdateUpstreamBillingProbeSnapshot(ctx, inFlight, &service.UpstreamBillingProbeSnapshot{
Status: service.UpstreamBillingProbeStatusOK,
LastAttemptAt: time.Now().UTC(),
}, &probedRate)
require.ErrorIs(t, err, service.ErrUpstreamBillingProbeIdentityChanged)
got, err := repo.GetByID(ctx, account.ID)
require.NoError(t, err)
require.NotNil(t, got.RateMultiplier)
require.Equal(t, manualRate, *got.RateMultiplier)
require.Equal(t, false, got.Extra[service.UpstreamBillingRateSyncEnabledExtraKey])
require.NotContains(t, got.Extra, service.UpstreamBillingProbeExtraKey)
}
func boolPtr(value bool) *bool {
return &value
}
@@ -249,7 +398,7 @@ func TestProxyIdentityUpdateInvalidatesProbeAndRejectsInFlightSnapshot(t *testin
err = accountRepo.UpdateUpstreamBillingProbeSnapshot(ctx, inFlight, &service.UpstreamBillingProbeSnapshot{
Status: service.UpstreamBillingProbeStatusOK,
LastAttemptAt: time.Now().UTC(),
})
}, nil)
require.ErrorIs(t, err, service.ErrUpstreamBillingProbeIdentityChanged)
rows, err := tx.QueryContext(ctx, `
@@ -1819,6 +1819,16 @@ func (s *stubAccountRepo) Update(ctx context.Context, account *service.Account)
return errors.New("not implemented")
}
func (s *stubAccountRepo) UpdateWithAccountBillingSettings(
ctx context.Context,
account *service.Account,
probeEnabled *bool,
rateSyncEnabled *bool,
rateMultiplier *float64,
) error {
return errors.New("not implemented")
}
func (s *stubAccountRepo) Delete(ctx context.Context, id int64) error {
return errors.New("not implemented")
}
@@ -131,11 +131,25 @@ type AccountDuplicateRepository interface {
CreateWithAccountGroups(ctx context.Context, account *Account, groups []AccountGroup) error
}
// AccountBillingSettingsRepository applies an admin edit without overwriting a
// rate_multiplier that a successful upstream probe synchronized after the edit
// form was loaded. A nil rateMultiplier means the request did not edit it.
type AccountBillingSettingsRepository interface {
UpdateWithAccountBillingSettings(
ctx context.Context,
account *Account,
probeEnabled *bool,
rateSyncEnabled *bool,
rateMultiplier *float64,
) error
}
// AdminAccountRepository makes the account-duplication write capability an explicit
// construction dependency without forcing read-only gateway test doubles to implement it.
type AdminAccountRepository interface {
AccountRepository
AccountDuplicateRepository
AccountBillingSettingsRepository
}
// AccountBulkUpdate describes the fields that can be updated in a bulk operation.
+91 -26
View File
@@ -455,6 +455,7 @@ func normalizeGrokMediaEligibilityUpdateExtra(account *Account, input *UpdateAcc
func buildAccountForCreate(input *CreateAccountInput, accountExtra map[string]any) (*Account, error) {
// Probe/session state is system-managed. New accounts always start with automatic refresh disabled.
delete(accountExtra, UpstreamBillingProbeEnabledExtraKey)
delete(accountExtra, UpstreamBillingRateSyncEnabledExtraKey)
delete(accountExtra, UpstreamBillingProbeExtraKey)
delete(accountExtra, OllamaCloudUsageSessionExtraKey)
delete(accountExtra, OllamaCloudUsageAutoRefreshExtraKey)
@@ -594,10 +595,6 @@ func (s *adminServiceImpl) CreateAccount(ctx context.Context, input *CreateAccou
return account, nil
}
type accountProbeEnabledAtomicUpdater interface {
UpdateWithUpstreamBillingProbeEnabled(context.Context, *Account, bool) error
}
func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *UpdateAccountInput) (*Account, error) {
account, err := s.accountRepo.GetByID(ctx, id)
if err != nil {
@@ -667,7 +664,8 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U
}
// Extra 使用 map:需要区分“未提供(nil)”与“显式清空({})”。
// 关闭配额限制时前端会删除 quota_* 键并提交 extra:{},此时也必须落库。
var requestedProbeEnabledUpdate *bool
requestedProbeEnabledUpdate := input.ProbeEnabled
requestedRateSyncEnabledUpdate := input.RateSyncEnabled
if input.Extra != nil {
requestedProbeEnabled, hasRequestedProbeEnabled := normalizedExtra[UpstreamBillingProbeEnabledExtraKey]
if hasRequestedProbeEnabled {
@@ -675,9 +673,13 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U
if !ok {
return nil, infraerrors.BadRequest("INVALID_UPSTREAM_BILLING_PROBE_ENABLED", "upstream_billing_probe_enabled must be a boolean")
}
if requestedProbeEnabledUpdate != nil && *requestedProbeEnabledUpdate != enabled {
return nil, infraerrors.BadRequest("CONFLICTING_UPSTREAM_BILLING_PROBE_ENABLED", "conflicting upstream_billing_probe_enabled values")
}
requestedProbeEnabledUpdate = &enabled
}
delete(normalizedExtra, UpstreamBillingProbeEnabledExtraKey)
delete(normalizedExtra, UpstreamBillingRateSyncEnabledExtraKey)
delete(normalizedExtra, UpstreamBillingProbeExtraKey)
delete(normalizedExtra, OllamaCloudUsageSessionExtraKey)
delete(normalizedExtra, OllamaCloudUsageAutoRefreshExtraKey)
@@ -691,6 +693,7 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U
"quota_weekly_start",
grokBillingExtraKey,
UpstreamBillingProbeEnabledExtraKey,
UpstreamBillingRateSyncEnabledExtraKey,
UpstreamBillingProbeExtraKey,
OllamaCloudUsageSessionExtraKey,
OllamaCloudUsageAutoRefreshExtraKey,
@@ -700,13 +703,6 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U
normalizedExtra[key] = v
}
}
if hasRequestedProbeEnabled {
if isUpstreamBillingProbeAccount(account) {
normalizedExtra[UpstreamBillingProbeEnabledExtraKey] = requestedProbeEnabled
} else {
delete(normalizedExtra, UpstreamBillingProbeEnabledExtraKey)
}
}
account.Extra = normalizedExtra
if account.Platform == PlatformAntigravity && wasOveragesEnabled && !account.IsOveragesEnabled() {
delete(account.Extra, "antigravity_credits_overages") // 清理旧版 overages 运行态
@@ -726,6 +722,35 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U
ComputeQuotaResetAt(account.Extra)
NormalizeFixedQuotaWindows(account.Extra)
}
if requestedRateSyncEnabledUpdate != nil && *requestedRateSyncEnabledUpdate {
if requestedProbeEnabledUpdate != nil && !*requestedProbeEnabledUpdate {
return nil, infraerrors.BadRequest(
"UPSTREAM_BILLING_RATE_SYNC_REQUIRES_PROBE",
"upstream billing rate sync requires upstream billing probe",
)
}
enabled := true
requestedProbeEnabledUpdate = &enabled
}
if requestedProbeEnabledUpdate != nil && !*requestedProbeEnabledUpdate {
disabled := false
requestedRateSyncEnabledUpdate = &disabled
}
if (requestedProbeEnabledUpdate != nil && *requestedProbeEnabledUpdate) ||
(requestedRateSyncEnabledUpdate != nil && *requestedRateSyncEnabledUpdate) {
if !isUpstreamBillingProbeAccount(account) {
return nil, ErrUpstreamBillingProbeAccountInvalid
}
}
if account.Extra == nil && (requestedProbeEnabledUpdate != nil || requestedRateSyncEnabledUpdate != nil) {
account.Extra = make(map[string]any)
}
if requestedProbeEnabledUpdate != nil {
account.Extra[UpstreamBillingProbeEnabledExtraKey] = *requestedProbeEnabledUpdate
}
if requestedRateSyncEnabledUpdate != nil {
account.Extra[UpstreamBillingRateSyncEnabledExtraKey] = *requestedRateSyncEnabledUpdate
}
// 影子代理恒继承母账号(由 propagateProxyToShadows 同步),不接受独立编辑——外审 B/P1;
// 否则要等母账号下次改 proxy 才被覆盖,期间影子会出现"有时继承、有时独立"的漂移。
if input.ProxyID != nil && !account.IsCredentialShadow() {
@@ -741,6 +766,7 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U
delete(account.Extra, UpstreamBillingProbeExtraKey)
if !isUpstreamBillingProbeAccount(account) {
delete(account.Extra, UpstreamBillingProbeEnabledExtraKey)
delete(account.Extra, UpstreamBillingRateSyncEnabledExtraKey)
}
}
if account.Extra != nil {
@@ -806,23 +832,40 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U
}
}
probeEnabledAppliedAtomically := false
if requestedProbeEnabledUpdate != nil && isUpstreamBillingProbeAccount(account) {
if updater, ok := s.accountRepo.(accountProbeEnabledAtomicUpdater); ok {
if err := updater.UpdateWithUpstreamBillingProbeEnabled(ctx, account, *requestedProbeEnabledUpdate); err != nil {
return nil, err
}
probeEnabledAppliedAtomically = true
}
billingSettingsAppliedAtomically := false
updater := s.accountBillingRepo
if updater == nil {
// Unit tests and narrow internal callers may construct adminServiceImpl
// directly; production wiring requires this capability through
// AdminAccountRepository.
updater, _ = s.accountRepo.(AccountBillingSettingsRepository)
}
if !probeEnabledAppliedAtomically {
if updater != nil {
if err := updater.UpdateWithAccountBillingSettings(
ctx,
account,
requestedProbeEnabledUpdate,
requestedRateSyncEnabledUpdate,
input.RateMultiplier,
); err != nil {
return nil, err
}
billingSettingsAppliedAtomically = true
}
if !billingSettingsAppliedAtomically {
if err := s.accountRepo.Update(ctx, account); err != nil {
return nil, err
}
if requestedProbeEnabledUpdate != nil && isUpstreamBillingProbeAccount(account) {
if err := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{
UpstreamBillingProbeEnabledExtraKey: *requestedProbeEnabledUpdate,
}); err != nil {
if (requestedProbeEnabledUpdate != nil || requestedRateSyncEnabledUpdate != nil) &&
isUpstreamBillingProbeAccount(account) {
settings := make(map[string]any, 2)
if requestedProbeEnabledUpdate != nil {
settings[UpstreamBillingProbeEnabledExtraKey] = *requestedProbeEnabledUpdate
}
if requestedRateSyncEnabledUpdate != nil {
settings[UpstreamBillingRateSyncEnabledExtraKey] = *requestedRateSyncEnabledUpdate
}
if err := s.accountRepo.UpdateExtra(ctx, account.ID, settings); err != nil {
return nil, err
}
}
@@ -854,6 +897,9 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U
// UpdateAccountExtra 仅对 Extra JSONB 做 key 级合并,避免覆盖其它运行态键
// (如 model_rate_limits / passive_usage_* 等)。
func (s *adminServiceImpl) UpdateAccountExtra(ctx context.Context, id int64, updates map[string]any) error {
delete(updates, UpstreamBillingProbeEnabledExtraKey)
delete(updates, UpstreamBillingRateSyncEnabledExtraKey)
delete(updates, UpstreamBillingProbeExtraKey)
delete(updates, OllamaCloudUsageSessionExtraKey)
delete(updates, OllamaCloudUsageAutoRefreshExtraKey)
delete(updates, OllamaCloudUsageSnapshotExtraKey)
@@ -877,6 +923,7 @@ func (s *adminServiceImpl) UpdateAccountExtra(ctx context.Context, id int64, upd
func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUpdateAccountsInput) (*BulkUpdateAccountsResult, error) {
// Managed probe/session state may only enter through dedicated typed endpoints.
delete(input.Extra, UpstreamBillingProbeEnabledExtraKey)
delete(input.Extra, UpstreamBillingRateSyncEnabledExtraKey)
delete(input.Extra, UpstreamBillingProbeExtraKey)
delete(input.Extra, OllamaCloudUsageSessionExtraKey)
delete(input.Extra, OllamaCloudUsageAutoRefreshExtraKey)
@@ -910,7 +957,7 @@ func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUp
// 预取所有目标账号,供凭据守卫/代理守卫/混合渠道检查共用,避免多次 DB 查询。
var cachedTargets []*Account
if len(input.Credentials) > 0 || input.ProxyID != nil || needMixedChannelCheck || hasLongContextBillingUpdate || input.ProbeEnabled != nil {
if len(input.Credentials) > 0 || input.ProxyID != nil || needMixedChannelCheck || hasLongContextBillingUpdate || input.ProbeEnabled != nil || input.RateMultiplier != nil {
loaded, err := s.accountRepo.GetByIDs(ctx, input.AccountIDs)
if err != nil {
return nil, err
@@ -996,6 +1043,21 @@ func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUp
if *input.RateMultiplier < 0 {
return nil, errors.New("rate_multiplier must be >= 0")
}
syncEnabledCount := 0
for _, account := range cachedTargets {
if account == nil || account.Extra == nil {
continue
}
enabled, _ := account.Extra[UpstreamBillingRateSyncEnabledExtraKey].(bool)
if enabled {
syncEnabledCount++
}
}
if syncEnabledCount > 0 {
return nil, ErrUpstreamBillingRateSyncBulkConflict.WithMetadata(map[string]string{
"count": strconv.Itoa(syncEnabledCount),
})
}
}
// 校验并规范化请求头覆写配置(批量路径为 JSONB 顶层 key 合并,直接校验增量即可)
@@ -1014,6 +1076,9 @@ func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUp
repoUpdates.Extra = make(map[string]any)
}
repoUpdates.Extra[UpstreamBillingProbeEnabledExtraKey] = *input.ProbeEnabled
if !*input.ProbeEnabled {
repoUpdates.Extra[UpstreamBillingRateSyncEnabledExtraKey] = false
}
}
if updatesUpstreamBillingProbeIdentity(input.Credentials) || input.ProxyID != nil {
if repoUpdates.Extra == nil {
@@ -17,6 +17,95 @@ func (r *upstreamBillingProbeAdminRepo) ListShadowsByParent(context.Context, int
return nil, nil
}
type accountBillingSettingsAdminRepo struct {
*upstreamBillingProbeAccountRepo
concurrentRate *float64
lastExplicitRate *float64
updateCalls int
}
func (r *accountBillingSettingsAdminRepo) UpdateWithAccountBillingSettings(
_ context.Context,
account *Account,
probeEnabled *bool,
rateSyncEnabled *bool,
rateMultiplier *float64,
) error {
r.mu.Lock()
defer r.mu.Unlock()
current := r.accounts[account.ID]
if current == nil {
return ErrAccountNotFound
}
updated := *account
updated.Credentials = mergeMap(nil, account.Credentials)
updated.Extra = mergeMap(nil, account.Extra)
if updated.Extra == nil {
updated.Extra = make(map[string]any)
}
if probeEnabled != nil {
updated.Extra[UpstreamBillingProbeEnabledExtraKey] = *probeEnabled
}
if rateSyncEnabled != nil {
updated.Extra[UpstreamBillingRateSyncEnabledExtraKey] = *rateSyncEnabled
}
switch {
case rateMultiplier != nil:
value := *rateMultiplier
updated.RateMultiplier = &value
r.lastExplicitRate = &value
case r.concurrentRate != nil:
value := *r.concurrentRate
updated.RateMultiplier = &value
r.lastExplicitRate = nil
default:
updated.RateMultiplier = cloneAccountValuePointer(current.RateMultiplier)
r.lastExplicitRate = nil
}
r.accounts[account.ID] = &updated
r.updateCalls++
return nil
}
func TestUpdateAccountRoutesRateIntentThroughAtomicBillingUpdater(t *testing.T) {
accountID := int64(109)
initialRate := 0.1
concurrentRate := 0.2
repo := &accountBillingSettingsAdminRepo{
upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {
ID: accountID,
Name: "before",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
RateMultiplier: &initialRate,
Extra: map[string]any{
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingRateSyncEnabledExtraKey: true,
},
},
}},
concurrentRate: &concurrentRate,
}
svc := &adminServiceImpl{accountRepo: repo}
updated, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{Name: "after"})
require.NoError(t, err)
require.Equal(t, 1, repo.updateCalls)
require.Nil(t, repo.lastExplicitRate)
require.Equal(t, concurrentRate, *updated.RateMultiplier)
zero := 0.0
updated, err = svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{RateMultiplier: &zero})
require.NoError(t, err)
require.Equal(t, 2, repo.updateCalls)
require.NotNil(t, repo.lastExplicitRate)
require.Zero(t, *repo.lastExplicitRate)
require.Zero(t, *updated.RateMultiplier)
}
func TestCreateAccountDropsManagedUpstreamBillingProbeState(t *testing.T) {
repo := &upstreamBillingProbeAccountRepo{}
svc := &adminServiceImpl{accountRepo: repo}
@@ -28,13 +117,15 @@ func TestCreateAccountDropsManagedUpstreamBillingProbeState(t *testing.T) {
Credentials: map[string]any{"api_key": "sk-test"},
SkipDefaultGroupBind: true,
Extra: map[string]any{
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"},
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingRateSyncEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"},
},
})
require.NoError(t, err)
require.NotContains(t, created.Extra, UpstreamBillingProbeEnabledExtraKey)
require.NotContains(t, created.Extra, UpstreamBillingRateSyncEnabledExtraKey)
require.NotContains(t, created.Extra, UpstreamBillingProbeExtraKey)
}
@@ -73,8 +164,9 @@ func TestUpdateAccountPreservesManagedUpstreamBillingProbeStateForUnrelatedEdit(
Type: AccountTypeAPIKey,
Status: StatusActive,
Extra: map[string]any{
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"},
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingRateSyncEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"},
},
},
}}
@@ -86,6 +178,7 @@ func TestUpdateAccountPreservesManagedUpstreamBillingProbeStateForUnrelatedEdit(
require.NoError(t, err)
require.Equal(t, true, updated.Extra[UpstreamBillingProbeEnabledExtraKey])
require.Equal(t, true, updated.Extra[UpstreamBillingRateSyncEnabledExtraKey])
require.Contains(t, updated.Extra, UpstreamBillingProbeExtraKey)
require.Equal(t, "value", updated.Extra["custom"])
}
@@ -196,8 +289,9 @@ func TestUpdateAccountInvalidatesProbeSnapshotWhenUpstreamIdentityChanges(t *tes
"base_url": "https://old.example",
},
Extra: map[string]any{
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"},
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingRateSyncEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"},
},
},
}}
@@ -210,6 +304,7 @@ func TestUpdateAccountInvalidatesProbeSnapshotWhenUpstreamIdentityChanges(t *tes
require.Equal(t, true, updated.Extra[UpstreamBillingProbeEnabledExtraKey])
} else {
require.NotContains(t, updated.Extra, UpstreamBillingProbeEnabledExtraKey)
require.NotContains(t, updated.Extra, UpstreamBillingRateSyncEnabledExtraKey)
}
})
}
@@ -289,16 +384,70 @@ func TestUpdateAccountAcceptsProbeEnabledAndRejectsInjectedSnapshot(t *testing.T
svc := &adminServiceImpl{accountRepo: repo}
updated, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
Extra: map[string]any{
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"},
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingRateSyncEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"},
},
})
require.NoError(t, err)
require.Equal(t, true, updated.Extra[UpstreamBillingProbeEnabledExtraKey])
require.NotContains(t, updated.Extra, UpstreamBillingRateSyncEnabledExtraKey)
require.NotContains(t, updated.Extra, UpstreamBillingProbeExtraKey)
}
func TestUpdateAccountRateSyncControlsProbeAndManualMode(t *testing.T) {
accountID := int64(151)
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {
ID: accountID,
Platform: PlatformGemini,
Type: AccountTypeAPIKey,
Status: StatusActive,
Extra: map[string]any{},
},
}}
svc := &adminServiceImpl{accountRepo: repo}
syncEnabled := true
updated, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
RateSyncEnabled: &syncEnabled,
})
require.NoError(t, err)
require.Equal(t, true, updated.Extra[UpstreamBillingProbeEnabledExtraKey])
require.Equal(t, true, updated.Extra[UpstreamBillingRateSyncEnabledExtraKey])
syncEnabled = false
updated, err = svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
RateSyncEnabled: &syncEnabled,
})
require.NoError(t, err)
require.Equal(t, true, updated.Extra[UpstreamBillingProbeEnabledExtraKey])
require.Equal(t, false, updated.Extra[UpstreamBillingRateSyncEnabledExtraKey])
}
func TestUpdateAccountRejectsSyncWithExplicitlyDisabledProbe(t *testing.T) {
accountID := int64(152)
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {
ID: accountID,
Platform: PlatformAnthropic,
Type: AccountTypeAPIKey,
Status: StatusActive,
},
}}
probeEnabled := false
syncEnabled := true
_, err := (&adminServiceImpl{accountRepo: repo}).UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
ProbeEnabled: &probeEnabled,
RateSyncEnabled: &syncEnabled,
})
require.Error(t, err)
require.Empty(t, repo.updates[accountID])
}
func TestUpdateAccountExplicitProbeDisableUsesDedicatedExtraUpdate(t *testing.T) {
accountID := int64(113)
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
@@ -321,6 +470,7 @@ func TestUpdateAccountExplicitProbeDisableUsesDedicatedExtraUpdate(t *testing.T)
require.NoError(t, err)
require.Len(t, repo.updates[accountID], 1)
require.Equal(t, false, repo.updates[accountID][0][UpstreamBillingProbeEnabledExtraKey])
require.Equal(t, false, repo.updates[accountID][0][UpstreamBillingRateSyncEnabledExtraKey])
}
func TestUpdateAccountExplicitUnchangedProbeEnabledStillUsesDedicatedExtraUpdate(t *testing.T) {
@@ -364,15 +514,36 @@ func TestUpdateAccountRejectsInvalidProbeEnabled(t *testing.T) {
require.Error(t, err)
}
func TestUpdateAccountExtraDropsManagedBillingProbeFields(t *testing.T) {
accountID := int64(153)
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
accountID: {ID: accountID, Platform: PlatformOpenAI, Type: AccountTypeAPIKey},
}}
err := (&adminServiceImpl{accountRepo: repo}).UpdateAccountExtra(context.Background(), accountID, map[string]any{
"custom": "value",
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingRateSyncEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"},
})
require.NoError(t, err)
require.Equal(t, "value", repo.accounts[accountID].Extra["custom"])
require.NotContains(t, repo.accounts[accountID].Extra, UpstreamBillingProbeEnabledExtraKey)
require.NotContains(t, repo.accounts[accountID].Extra, UpstreamBillingRateSyncEnabledExtraKey)
require.NotContains(t, repo.accounts[accountID].Extra, UpstreamBillingProbeExtraKey)
}
func TestBulkUpdateAccountsDropsManagedUpstreamBillingProbeState(t *testing.T) {
repo := &upstreamBillingProbeAccountRepo{}
svc := &adminServiceImpl{accountRepo: repo}
input := &BulkUpdateAccountsInput{
AccountIDs: []int64{1},
Extra: map[string]any{
"custom": "value",
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"},
"custom": "value",
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingRateSyncEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"},
},
}
@@ -383,6 +554,7 @@ func TestBulkUpdateAccountsDropsManagedUpstreamBillingProbeState(t *testing.T) {
require.Len(t, repo.bulkUpdates, 1)
require.Equal(t, "value", repo.bulkUpdates[0].Extra["custom"])
require.NotContains(t, repo.bulkUpdates[0].Extra, UpstreamBillingProbeEnabledExtraKey)
require.NotContains(t, repo.bulkUpdates[0].Extra, UpstreamBillingRateSyncEnabledExtraKey)
require.NotContains(t, repo.bulkUpdates[0].Extra, UpstreamBillingProbeExtraKey)
}
@@ -403,6 +575,9 @@ func TestBulkUpdateAccountsAcceptsDedicatedUpstreamBillingProbeSetting(t *testin
require.Equal(t, 2, result.Success)
require.Len(t, repo.bulkUpdates, 1)
require.Equal(t, enabled, repo.bulkUpdates[0].Extra[UpstreamBillingProbeEnabledExtraKey])
if !enabled {
require.Equal(t, false, repo.bulkUpdates[0].Extra[UpstreamBillingRateSyncEnabledExtraKey])
}
require.NotNil(t, repo.bulkUpdates[0].ProbeEnabled)
require.Equal(t, enabled, *repo.bulkUpdates[0].ProbeEnabled)
})
@@ -376,6 +376,8 @@ type UpdateAccountInput struct {
GroupIDs *[]int64
ExpiresAt *int64
AutoPauseOnExpired *bool
ProbeEnabled *bool
RateSyncEnabled *bool
SkipMixedChannelCheck bool // 跳过混合渠道检查(用户已确认风险)
}
@@ -617,6 +619,7 @@ type adminServiceImpl struct {
groupDuplicateRepo GroupDuplicateRepository
accountRepo AccountRepository
accountDuplicateRepo AccountDuplicateRepository
accountBillingRepo AccountBillingSettingsRepository
proxyRepo ProxyRepository
apiKeyRepo APIKeyRepository
redeemCodeRepo RedeemCodeRepository
@@ -675,6 +678,7 @@ func NewAdminService(
groupDuplicateRepo: groupRepo,
accountRepo: accountRepo,
accountDuplicateRepo: accountRepo,
accountBillingRepo: accountRepo,
proxyRepo: proxyRepo,
apiKeyRepo: apiKeyRepo,
redeemCodeRepo: redeemCodeRepo,
@@ -5,9 +5,11 @@ package service
import (
"context"
"errors"
"net/http"
"reflect"
"testing"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/stretchr/testify/require"
)
@@ -153,6 +155,38 @@ func TestAdminService_BulkUpdateAccounts_AllSuccessIDs(t *testing.T) {
require.Len(t, result.Results, 3)
}
func TestAdminService_BulkUpdateAccounts_RejectsRateChangeForSyncedAccounts(t *testing.T) {
repo := &accountRepoStubForBulkUpdate{
getByIDsAccounts: []*Account{
{
ID: 1,
Extra: map[string]any{
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingRateSyncEnabledExtraKey: true,
},
},
{ID: 2, Extra: map[string]any{}},
},
}
svc := &adminServiceImpl{accountRepo: repo}
rateMultiplier := 0.5
result, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{
AccountIDs: []int64{1, 2},
RateMultiplier: &rateMultiplier,
})
require.Nil(t, result)
require.Error(t, err)
var appErr *infraerrors.ApplicationError
require.ErrorAs(t, err, &appErr)
require.Equal(t, int32(http.StatusConflict), appErr.Code)
require.Equal(t, "UPSTREAM_BILLING_RATE_SYNC_BULK_CONFLICT", appErr.Reason)
require.Equal(t, "1", appErr.Metadata["count"])
require.True(t, repo.getByIDsCalled)
require.Empty(t, repo.bulkUpdateIDs, "rate conflict must be rejected before any write")
}
// TestAdminService_BulkUpdateAccounts_PartialFailureIDs 验证部分失败时 success_ids/failed_ids 正确。
func TestAdminService_BulkUpdateAccounts_PartialFailureIDs(t *testing.T) {
repo := &accountRepoStubForBulkUpdate{
@@ -141,6 +141,7 @@ func TestDuplicateAccountCopiesConfigurationAndResetsRuntimeState(t *testing.T)
SessionWindowStatus: "active",
}
source.Extra[UpstreamBillingProbeEnabledExtraKey] = true
source.Extra[UpstreamBillingRateSyncEnabledExtraKey] = true
source.Extra[UpstreamBillingProbeExtraKey] = map[string]any{"status": "ok"}
require.NoError(t, repo.Create(ctx, source))
@@ -162,6 +163,7 @@ func TestDuplicateAccountCopiesConfigurationAndResetsRuntimeState(t *testing.T)
"quota_limit": float64(1000),
"codex_cli_only": true,
}, duplicate.Extra)
require.NotContains(t, duplicate.Extra, UpstreamBillingRateSyncEnabledExtraKey)
require.NotNil(t, duplicate.ExpiresAt)
require.True(t, source.ExpiresAt.Equal(*duplicate.ExpiresAt))
require.Equal(t, source.Notes, duplicate.Notes)
@@ -115,15 +115,17 @@ func TestShouldCreateAccount(t *testing.T) {
func TestReconcileCRSUpstreamBillingProbeExtra(t *testing.T) {
remote := map[string]any{
"crs_account_id": "remote-1",
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "remote"},
"crs_account_id": "remote-1",
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingRateSyncEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: map[string]any{"status": "remote"},
}
t.Run("create drops remote managed fields", func(t *testing.T) {
extra := mergeMap(nil, remote)
reconcileCRSUpstreamBillingProbeExtra(nil, PlatformOpenAI, AccountTypeAPIKey, map[string]any{"api_key": "new"}, extra)
require.NotContains(t, extra, UpstreamBillingProbeEnabledExtraKey)
require.NotContains(t, extra, UpstreamBillingRateSyncEnabledExtraKey)
require.NotContains(t, extra, UpstreamBillingProbeExtraKey)
})
@@ -132,8 +134,9 @@ func TestReconcileCRSUpstreamBillingProbeExtra(t *testing.T) {
Type: AccountTypeAPIKey,
Credentials: map[string]any{"api_key": "local", "base_url": "http://127.0.0.1:8080"},
Extra: map[string]any{
UpstreamBillingProbeEnabledExtraKey: false,
UpstreamBillingProbeExtraKey: map[string]any{"status": "local"},
UpstreamBillingProbeEnabledExtraKey: false,
UpstreamBillingRateSyncEnabledExtraKey: false,
UpstreamBillingProbeExtraKey: map[string]any{"status": "local"},
},
}
@@ -141,13 +144,27 @@ func TestReconcileCRSUpstreamBillingProbeExtra(t *testing.T) {
extra := mergeMap(existing.Extra, remote)
reconcileCRSUpstreamBillingProbeExtra(existing, existing.Platform, existing.Type, mergeMap(existing.Credentials, nil), extra)
require.Equal(t, false, extra[UpstreamBillingProbeEnabledExtraKey])
require.Equal(t, false, extra[UpstreamBillingRateSyncEnabledExtraKey])
require.Equal(t, map[string]any{"status": "local"}, extra[UpstreamBillingProbeExtraKey])
})
t.Run("same identity preserves enabled rate sync", func(t *testing.T) {
enabled := *existing
enabled.Extra = mergeMap(existing.Extra, map[string]any{
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingRateSyncEnabledExtraKey: true,
})
extra := mergeMap(enabled.Extra, remote)
reconcileCRSUpstreamBillingProbeExtra(&enabled, enabled.Platform, enabled.Type, mergeMap(enabled.Credentials, nil), extra)
require.Equal(t, true, extra[UpstreamBillingProbeEnabledExtraKey])
require.Equal(t, true, extra[UpstreamBillingRateSyncEnabledExtraKey])
})
t.Run("identity change keeps enabled and clears snapshot", func(t *testing.T) {
extra := mergeMap(existing.Extra, remote)
reconcileCRSUpstreamBillingProbeExtra(existing, PlatformOpenAI, AccountTypeAPIKey, map[string]any{"api_key": "changed"}, extra)
require.Equal(t, false, extra[UpstreamBillingProbeEnabledExtraKey])
require.Equal(t, false, extra[UpstreamBillingRateSyncEnabledExtraKey])
require.NotContains(t, extra, UpstreamBillingProbeExtraKey)
})
@@ -164,6 +181,7 @@ func TestReconcileCRSUpstreamBillingProbeExtra(t *testing.T) {
extra := mergeMap(existing.Extra, remote)
reconcileCRSUpstreamBillingProbeExtra(existing, target.platform, AccountTypeAPIKey, existing.Credentials, extra)
require.Equal(t, false, extra[UpstreamBillingProbeEnabledExtraKey])
require.Equal(t, false, extra[UpstreamBillingRateSyncEnabledExtraKey])
require.NotContains(t, extra, UpstreamBillingProbeExtraKey)
})
}
@@ -181,6 +199,7 @@ func TestReconcileCRSUpstreamBillingProbeExtra(t *testing.T) {
extra := mergeMap(existing.Extra, remote)
reconcileCRSUpstreamBillingProbeExtra(existing, target.platform, target.typeName, existing.Credentials, extra)
require.NotContains(t, extra, UpstreamBillingProbeEnabledExtraKey)
require.NotContains(t, extra, UpstreamBillingRateSyncEnabledExtraKey)
require.NotContains(t, extra, UpstreamBillingProbeExtraKey)
})
}
@@ -1162,6 +1162,7 @@ func reconcileCRSUpstreamBillingProbeExtra(
) {
for _, key := range []string{
UpstreamBillingProbeEnabledExtraKey,
UpstreamBillingRateSyncEnabledExtraKey,
UpstreamBillingProbeExtraKey,
OllamaCloudUsageSessionExtraKey,
OllamaCloudUsageAutoRefreshExtraKey,
@@ -1174,8 +1175,13 @@ func reconcileCRSUpstreamBillingProbeExtra(
}
target := &Account{Platform: targetPlatform, Type: targetType, Credentials: targetCredentials}
if IsUpstreamBillingProbeIdentity(targetPlatform, targetType) {
probeEnabled := false
if enabled, ok := existing.Extra[UpstreamBillingProbeEnabledExtraKey]; ok {
extra[UpstreamBillingProbeEnabledExtraKey] = enabled
probeEnabled, _ = enabled.(bool)
}
if enabled, ok := existing.Extra[UpstreamBillingRateSyncEnabledExtraKey].(bool); ok {
extra[UpstreamBillingRateSyncEnabledExtraKey] = enabled && probeEnabled
}
if reflect.DeepEqual(upstreamBillingProbeIdentity(existing), upstreamBillingProbeIdentity(target)) {
if snapshot, ok := existing.Extra[UpstreamBillingProbeExtraKey]; ok {
@@ -28,8 +28,9 @@ import (
const (
// These values live in accounts.extra so PR2 does not require a schema migration.
UpstreamBillingProbeExtraKey = "upstream_billing_probe"
UpstreamBillingProbeEnabledExtraKey = "upstream_billing_probe_enabled"
UpstreamBillingProbeExtraKey = "upstream_billing_probe"
UpstreamBillingProbeEnabledExtraKey = "upstream_billing_probe_enabled"
UpstreamBillingRateSyncEnabledExtraKey = "upstream_billing_rate_sync_enabled"
upstreamBillingProbeDefaultIntervalMinutes = 30
upstreamBillingProbeMinIntervalMinutes = 5
@@ -44,6 +45,8 @@ const (
// /v1/sub2api/billing,按常规 interval 重排只会持续占满每周期
// upstreamBillingProbeMaxPerCycle 个名额。
upstreamBillingProbeUnsupportedDelayFactor = 8
upstreamBillingProbeAccountRateScale = 10000.0
upstreamBillingProbeAccountRateMax = 999999.9999
upstreamBillingProbeLeaderLockKey = "upstream:billing:probe:leader"
upstreamBillingProbeLeaderLockTTL = 2 * time.Minute
)
@@ -61,6 +64,10 @@ var (
ErrUpstreamBillingProbeIdentityChanged = infraerrors.Conflict(
"UPSTREAM_BILLING_PROBE_IDENTITY_CHANGED", "account identity changed during upstream billing probe; retry the probe",
)
ErrUpstreamBillingRateSyncBulkConflict = infraerrors.Conflict(
"UPSTREAM_BILLING_RATE_SYNC_BULK_CONFLICT",
"account rate multiplier cannot be changed in bulk while upstream billing rate sync is enabled",
)
)
const (
@@ -197,7 +204,7 @@ type UpstreamBillingProbeService struct {
}
type upstreamBillingProbeSnapshotWriter interface {
UpdateUpstreamBillingProbeSnapshot(context.Context, *Account, *UpstreamBillingProbeSnapshot) error
UpdateUpstreamBillingProbeSnapshot(context.Context, *Account, *UpstreamBillingProbeSnapshot, *float64) error
}
type upstreamBillingProbeDueAccountLister interface {
@@ -546,9 +553,11 @@ func (s *UpstreamBillingProbeService) SetAccountEnabled(ctx context.Context, acc
if !isUpstreamBillingProbeAccount(account) {
return ErrUpstreamBillingProbeAccountInvalid
}
return s.accountRepo.UpdateExtra(ctx, accountID, map[string]any{
UpstreamBillingProbeEnabledExtraKey: enabled,
})
updates := map[string]any{UpstreamBillingProbeEnabledExtraKey: enabled}
if !enabled {
updates[UpstreamBillingRateSyncEnabledExtraKey] = false
}
return s.accountRepo.UpdateExtra(ctx, accountID, updates)
}
func (s *UpstreamBillingProbeService) probeLoadedAccount(ctx context.Context, account *Account, intervalMinutes int) (*UpstreamBillingProbeSnapshot, error) {
@@ -635,6 +644,10 @@ func (s *UpstreamBillingProbeService) probeLoadedAccount(ctx context.Context, ac
if err != nil {
return s.persistProbeFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "invalid_response", retryAfter(resp.Header, now))
}
rateMultiplier, ok := upstreamBillingProbeAccountRate(data)
if !ok {
return s.persistProbeFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "invalid_response", retryAfter(resp.Header, now))
}
snapshot := &UpstreamBillingProbeSnapshot{
Status: UpstreamBillingProbeStatusOK,
Data: data,
@@ -644,7 +657,7 @@ func (s *UpstreamBillingProbeService) probeLoadedAccount(ctx context.Context, ac
NextProbeAt: now.Add(nextProbeDelay(intervalMinutes, 0)),
HTTPStatus: resp.StatusCode,
}
if err := s.updateSnapshot(ctx, account, snapshot); err != nil {
if err := s.updateSnapshot(ctx, account, snapshot, &rateMultiplier); err != nil {
return nil, err
}
return snapshot, nil
@@ -686,18 +699,23 @@ func (s *UpstreamBillingProbeService) persistProbeFailure(
snapshot.FreshUntil = probeTimePtr(previous.ReceivedAt.Add(2 * time.Duration(intervalMinutes) * time.Minute))
}
}
if err := s.updateSnapshot(ctx, account, snapshot); err != nil {
if err := s.updateSnapshot(ctx, account, snapshot, nil); err != nil {
return nil, err
}
return snapshot, nil
}
func (s *UpstreamBillingProbeService) updateSnapshot(ctx context.Context, account *Account, snapshot *UpstreamBillingProbeSnapshot) error {
func (s *UpstreamBillingProbeService) updateSnapshot(
ctx context.Context,
account *Account,
snapshot *UpstreamBillingProbeSnapshot,
rateMultiplier *float64,
) error {
writer, ok := s.accountRepo.(upstreamBillingProbeSnapshotWriter)
if !ok {
return ErrUpstreamBillingProbeUnavailable
}
return writer.UpdateUpstreamBillingProbeSnapshot(ctx, account, snapshot)
return writer.UpdateUpstreamBillingProbeSnapshot(ctx, account, snapshot, rateMultiplier)
}
func parseUpstreamBillingProbeResponse(body []byte) (map[string]any, error) {
@@ -799,6 +817,21 @@ func upstreamBillingRateAt(data map[string]any, now time.Time) (float64, bool) {
return base, true
}
// upstreamBillingProbeAccountRate converts the declared effective multiplier
// to the precision supported by accounts.rate_multiplier (DECIMAL(10,4)).
func upstreamBillingProbeAccountRate(data map[string]any) (float64, bool) {
value, ok := resolveAccountExtraNumber(data, "effective_rate_multiplier")
if !ok || value < 0 || value > upstreamBillingProbeAccountRateMax ||
math.IsNaN(value) || math.IsInf(value, 0) {
return 0, false
}
rounded := math.Round(value*upstreamBillingProbeAccountRateScale) / upstreamBillingProbeAccountRateScale
if value > 0 && rounded == 0 {
return 0, false
}
return rounded, true
}
func upstreamBillingPeakMultiplierAt(data map[string]any, now time.Time) (float64, bool) {
peakEnabled, ok := data["peak_rate_enabled"].(bool)
if !ok {
@@ -865,9 +898,7 @@ func decodeUpstreamBillingProbeSnapshot(extra map[string]any) *UpstreamBillingPr
// 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.
// sub2api convention shared by the five supported API-key platforms.
// 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
@@ -876,7 +907,15 @@ func decodeUpstreamBillingProbeSnapshot(extra map[string]any) *UpstreamBillingPr
// 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
if accountType != AccountTypeAPIKey {
return false
}
switch platform {
case PlatformOpenAI, PlatformAnthropic, PlatformGemini, PlatformAntigravity, PlatformGrok:
return true
default:
return false
}
}
func isUpstreamBillingProbeAccount(account *Account) bool {
@@ -934,6 +973,14 @@ func upstreamBillingProbeEnabled(account *Account) bool {
return ok && enabled
}
func upstreamBillingRateSyncEnabled(account *Account) bool {
if account == nil || account.Extra == nil {
return false
}
enabled, ok := account.Extra[UpstreamBillingRateSyncEnabledExtraKey].(bool)
return ok && enabled && upstreamBillingProbeEnabled(account)
}
func (s *UpstreamBillingProbeService) currentTime() time.Time {
if s != nil && s.now != nil {
return s.now()
@@ -11,8 +11,8 @@ import (
"github.com/stretchr/testify/require"
)
// 探测资格放宽:/v1/sub2api/billing 是 key 级端点,任何 API-key
// 平台账号都可开启探测;OAuth/Bedrock 无静态 Key 仍不合格。
// 探测资格:/v1/sub2api/billing 是 key 级端点,五个
// 受支持平台的 API-key 账号都可开启探测;OAuth/Bedrock 无静态 Key 仍不合格。
func TestUpstreamBillingProbeIdentityCoversAllAPIKeyPlatforms(t *testing.T) {
for _, platform := range []string{
PlatformOpenAI, PlatformGrok, PlatformAnthropic, PlatformGemini, PlatformAntigravity,
@@ -24,6 +24,7 @@ func TestUpstreamBillingProbeIdentityCoversAllAPIKeyPlatforms(t *testing.T) {
require.False(t, IsUpstreamBillingProbeIdentity(PlatformGrok, AccountTypeOAuth))
require.False(t, IsUpstreamBillingProbeIdentity(PlatformAnthropic, AccountTypeBedrock))
require.False(t, IsUpstreamBillingProbeIdentity("", AccountTypeAPIKey))
require.False(t, IsUpstreamBillingProbeIdentity("future-platform", AccountTypeAPIKey))
require.False(t, isUpstreamBillingProbeAccount(nil))
}
@@ -110,7 +110,12 @@ func (r *upstreamBillingProbeAccountRepo) UpdateExtra(_ context.Context, id int6
return nil
}
func (r *upstreamBillingProbeAccountRepo) UpdateUpstreamBillingProbeSnapshot(_ context.Context, expected *Account, snapshot *UpstreamBillingProbeSnapshot) error {
func (r *upstreamBillingProbeAccountRepo) UpdateUpstreamBillingProbeSnapshot(
_ context.Context,
expected *Account,
snapshot *UpstreamBillingProbeSnapshot,
rateMultiplier *float64,
) error {
r.mu.Lock()
defer r.mu.Unlock()
account := r.accounts[expected.ID]
@@ -121,6 +126,12 @@ func (r *upstreamBillingProbeAccountRepo) UpdateUpstreamBillingProbeSnapshot(_ c
account.Extra = make(map[string]any)
}
account.Extra[UpstreamBillingProbeExtraKey] = snapshot
if snapshot.Status == UpstreamBillingProbeStatusOK &&
rateMultiplier != nil &&
upstreamBillingRateSyncEnabled(account) {
value := *rateMultiplier
account.RateMultiplier = &value
}
return nil
}
@@ -259,6 +270,7 @@ func TestUpstreamBillingProbeSettingsDefaultsAndValidation(t *testing.T) {
}
func TestUpstreamBillingProbeSuccessPersistsSanitizedSnapshot(t *testing.T) {
initialRate := 0.25
account := &Account{
ID: 17,
Platform: PlatformOpenAI,
@@ -269,6 +281,11 @@ func TestUpstreamBillingProbeSuccessPersistsSanitizedSnapshot(t *testing.T) {
"api_key": "sk-sensitive",
"base_url": "https://upstream.example/v1",
},
Extra: map[string]any{
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingRateSyncEnabledExtraKey: true,
},
RateMultiplier: &initialRate,
}
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{account.ID: account}}
upstream := &httpUpstreamRecorder{resp: &http.Response{
@@ -307,6 +324,8 @@ func TestUpstreamBillingProbeSuccessPersistsSanitizedSnapshot(t *testing.T) {
require.Equal(t, fixedNow.Add(time.Hour), *snapshot.FreshUntil)
require.False(t, snapshot.NextProbeAt.Before(fixedNow.Add(24*time.Minute)))
require.False(t, snapshot.NextProbeAt.After(fixedNow.Add(36*time.Minute)))
require.NotNil(t, account.RateMultiplier)
require.Equal(t, 0.9, *account.RateMultiplier)
require.Equal(t, "https://upstream.example/v1/sub2api/billing", upstream.lastReq.URL.String())
require.Equal(t, http.MethodGet, upstream.lastReq.Method)
require.Equal(t, "Bearer sk-sensitive", upstream.lastReq.Header.Get("Authorization"))
@@ -317,6 +336,98 @@ func TestUpstreamBillingProbeSuccessPersistsSanitizedSnapshot(t *testing.T) {
require.Equal(t, snapshot.Status, persisted.Status)
}
func TestUpstreamBillingProbeSyncsEffectiveRateForAllAPIKeyPlatforms(t *testing.T) {
for _, platform := range []string{
PlatformOpenAI,
PlatformAnthropic,
PlatformGemini,
PlatformAntigravity,
PlatformGrok,
} {
t.Run(platform, func(t *testing.T) {
initialRate := 0.25
account := &Account{
ID: 17,
Platform: platform,
Type: AccountTypeAPIKey,
Status: StatusActive,
Concurrency: 1,
RateMultiplier: &initialRate,
Credentials: map[string]any{
"api_key": "sk-sensitive",
"base_url": "https://upstream.example",
},
Extra: map[string]any{
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingRateSyncEnabledExtraKey: true,
},
}
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{account.ID: account}}
svc := newUpstreamBillingProbeTestService(repo, &upstreamBillingProbeHTTPStub{}, &upstreamBillingProbeSettingRepo{})
snapshot, err := svc.ProbeAccount(context.Background(), account.ID)
require.NoError(t, err)
require.Equal(t, UpstreamBillingProbeStatusOK, snapshot.Status)
require.NotNil(t, account.RateMultiplier)
require.Equal(t, 0.8, *account.RateMultiplier)
})
}
}
func TestUpstreamBillingProbeOnlyDoesNotChangeAccountRate(t *testing.T) {
initialRate := 0.25
account := &Account{
ID: 18,
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Status: StatusActive,
Concurrency: 1,
RateMultiplier: &initialRate,
Credentials: map[string]any{
"api_key": "sk-sensitive",
"base_url": "https://upstream.example",
},
Extra: map[string]any{UpstreamBillingProbeEnabledExtraKey: true},
}
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{account.ID: account}}
svc := newUpstreamBillingProbeTestService(repo, &upstreamBillingProbeHTTPStub{}, &upstreamBillingProbeSettingRepo{})
snapshot, err := svc.ProbeAccount(context.Background(), account.ID)
require.NoError(t, err)
require.Equal(t, UpstreamBillingProbeStatusOK, snapshot.Status)
require.NotNil(t, account.RateMultiplier)
require.Equal(t, initialRate, *account.RateMultiplier)
require.Contains(t, account.Extra, UpstreamBillingProbeExtraKey)
}
func TestUpstreamBillingProbeAccountRateMatchesDatabasePrecision(t *testing.T) {
tests := []struct {
name string
value float64
want float64
ok bool
}{
{name: "zero", value: 0, want: 0, ok: true},
{name: "round to four decimals", value: 0.07654, want: 0.0765, ok: true},
{name: "maximum", value: upstreamBillingProbeAccountRateMax, want: upstreamBillingProbeAccountRateMax, ok: true},
{name: "positive below database precision", value: 0.00001, ok: false},
{name: "overflow before rounding", value: 999999.99991, ok: false},
{name: "overflow after rounding", value: 999999.99996, ok: false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, ok := upstreamBillingProbeAccountRate(map[string]any{"effective_rate_multiplier": tt.value})
require.Equal(t, tt.ok, ok)
if tt.ok {
require.Equal(t, tt.want, got)
}
})
}
}
func TestUpstreamBillingProbeRejectsMissingRequiredMultiplier(t *testing.T) {
_, err := parseUpstreamBillingProbeResponse([]byte(`{
"object":"sub2api.key_billing",
@@ -456,6 +567,7 @@ func TestUpstreamBillingRateAtHandlesDST(t *testing.T) {
func TestUpstreamBillingProbeFailurePreservesLastSuccessAndRetryAfter(t *testing.T) {
receivedAt := time.Date(2026, time.July, 12, 12, 0, 0, 0, time.UTC)
initialRate := 0.35
previous := &UpstreamBillingProbeSnapshot{
Status: UpstreamBillingProbeStatusOK,
Data: map[string]any{"effective_rate_multiplier": 0.5},
@@ -463,13 +575,17 @@ func TestUpstreamBillingProbeFailurePreservesLastSuccessAndRetryAfter(t *testing
FailureCount: 1,
}
account := &Account{
ID: 18,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
Concurrency: 1,
Credentials: map[string]any{"api_key": "sk-test", "base_url": "https://upstream.example"},
Extra: map[string]any{UpstreamBillingProbeExtraKey: previous},
ID: 18,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
Concurrency: 1,
RateMultiplier: &initialRate,
Credentials: map[string]any{"api_key": "sk-test", "base_url": "https://upstream.example"},
Extra: map[string]any{
UpstreamBillingProbeEnabledExtraKey: true,
UpstreamBillingProbeExtraKey: previous,
},
}
repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{account.ID: account}}
upstream := &httpUpstreamRecorder{resp: &http.Response{
@@ -492,6 +608,8 @@ func TestUpstreamBillingProbeFailurePreservesLastSuccessAndRetryAfter(t *testing
require.Equal(t, "http_error", snapshot.LastError)
require.False(t, snapshot.NextProbeAt.Before(fixedNow.Add(4*time.Hour)))
require.NotContains(t, snapshot.LastError, "do not persist")
require.NotNil(t, account.RateMultiplier)
require.Equal(t, initialRate, *account.RateMultiplier)
}
func TestUpstreamBillingProbeRetryAfterIsNotShortened(t *testing.T) {
@@ -611,6 +729,7 @@ func TestUpstreamBillingProbeUnsupportedAndAccountToggle(t *testing.T) {
require.NoError(t, svc.SetAccountEnabled(context.Background(), account.ID, true))
require.Equal(t, true, account.Extra[UpstreamBillingProbeEnabledExtraKey])
account.Extra[UpstreamBillingRateSyncEnabledExtraKey] = true
snapshot, err := svc.ProbeAccount(context.Background(), account.ID)
require.NoError(t, err)
require.Equal(t, UpstreamBillingProbeStatusUnsupported, snapshot.Status)
@@ -624,6 +743,9 @@ func TestUpstreamBillingProbeUnsupportedAndAccountToggle(t *testing.T) {
require.Equal(t, 2, snapshot.FailureCount)
require.False(t, snapshot.NextProbeAt.Before(fixedNow.Add(192*time.Minute)))
require.False(t, snapshot.NextProbeAt.After(fixedNow.Add(288*time.Minute)))
require.NoError(t, svc.SetAccountEnabled(context.Background(), account.ID, false))
require.Equal(t, false, account.Extra[UpstreamBillingProbeEnabledExtraKey])
require.Equal(t, false, account.Extra[UpstreamBillingRateSyncEnabledExtraKey])
invalid := &Account{ID: 20, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
repo.accounts[invalid.ID] = invalid
@@ -662,10 +784,14 @@ func TestUpstreamBillingProbeRunnerIsBoundedAndManualProbeIgnoresSwitches(t *tes
require.Equal(t, int64(20), upstream.calls.Load())
accounts[25].Extra[UpstreamBillingProbeEnabledExtraKey] = false
manualRate := 0.25
accounts[25].RateMultiplier = &manualRate
snapshot, err := svc.ProbeAccount(context.Background(), 25)
require.NoError(t, err)
require.Equal(t, UpstreamBillingProbeStatusOK, snapshot.Status)
require.Equal(t, int64(21), upstream.calls.Load())
require.NotNil(t, accounts[25].RateMultiplier)
require.Equal(t, manualRate, *accounts[25].RateMultiplier)
}
func TestUpstreamBillingProbeRunnerRechecksEnabledAfterDueSelection(t *testing.T) {
@@ -750,6 +750,14 @@
aria-labelledby="bulk-edit-rate-multiplier-label"
/>
<p class="input-hint">{{ t('admin.accounts.billingRateMultiplierHint') }}</p>
<p
v-if="enableRateMultiplier"
class="mt-2 flex items-start gap-1 text-xs text-amber-700 dark:text-amber-300"
data-testid="bulk-rate-sync-warning"
>
<Icon name="exclamationTriangle" size="xs" class="mt-0.5 flex-shrink-0" />
<span>{{ t('admin.accounts.bulkEdit.rateSyncWarning') }}</span>
</p>
</div>
</div>
@@ -1985,6 +1993,10 @@ const submitBulkUpdate = async (baseUpdates: Record<string, unknown>) => {
pendingUpdatesForConfirm.value = baseUpdates
mixedChannelWarningMessage.value = error.message
showMixedChannelWarning.value = true
} else if (error.reason === 'UPSTREAM_BILLING_RATE_SYNC_BULK_CONFLICT') {
appStore.showError(t('admin.accounts.bulkEdit.rateSyncConflict', {
count: error.metadata?.count ?? 1
}))
} else {
appStore.showError(error.message || t('admin.accounts.bulkEdit.failed'))
console.error('Error bulk updating accounts:', error)
@@ -1441,8 +1441,43 @@
</div>
<div>
<label class="input-label">{{ t('admin.accounts.billingRateMultiplier') }}</label>
<input v-model.number="form.rate_multiplier" type="number" min="0" step="0.001" class="input" />
<p class="input-hint">{{ t('admin.accounts.billingRateMultiplierHint') }}</p>
<input
v-model.number="form.rate_multiplier"
type="number"
min="0"
step="0.001"
class="input disabled:cursor-not-allowed disabled:opacity-60"
data-testid="account-rate-multiplier"
:disabled="upstreamBillingRateSyncEnabled"
/>
<p class="input-hint">
{{
t(
upstreamBillingRateSyncEnabled
? 'admin.accounts.upstreamBilling.syncRateManagedHint'
: 'admin.accounts.billingRateMultiplierHint'
)
}}
</p>
<div
v-if="account?.type === 'apikey'"
class="mt-3 flex items-center justify-between gap-3"
>
<div class="min-w-0">
<p class="text-xs font-medium text-gray-700 dark:text-gray-200">
{{ t('admin.accounts.upstreamBilling.syncRate') }}
</p>
<p class="mt-1 text-xs text-gray-500 dark:text-gray-400">
{{ t('admin.accounts.upstreamBilling.syncRateHint') }}
</p>
</div>
<Toggle
:model-value="upstreamBillingRateSyncEnabled"
data-testid="upstream-billing-rate-sync"
:aria-label="t('admin.accounts.upstreamBilling.syncRate')"
@update:model-value="handleUpstreamBillingRateSyncChange"
/>
</div>
</div>
</div>
<div class="border-t border-gray-200 pt-4 dark:border-dark-600">
@@ -1660,9 +1695,10 @@
</p>
</div>
<Toggle
v-model="upstreamBillingAutoProbeEnabled"
:model-value="upstreamBillingAutoProbeEnabled"
data-testid="upstream-billing-auto-probe"
:aria-label="t('admin.accounts.upstreamBilling.autoProbe')"
@update:model-value="handleUpstreamBillingAutoProbeChange"
/>
</div>
@@ -2822,6 +2858,7 @@ const autoPause7dThreshold = ref<number | null>(null)
const autoPause5hDisabled = ref(false)
const autoPause7dDisabled = ref(false)
const upstreamBillingAutoProbeEnabled = ref(false)
const upstreamBillingRateSyncEnabled = ref(false)
const mixedScheduling = ref(false) // For antigravity accounts: enable mixed scheduling
const allowOverages = ref(false) // For antigravity accounts: enable AI Credits overages
const antigravityProjectId = ref('')
@@ -3191,6 +3228,20 @@ const form = reactive({
expires_at: null as number | null
})
const handleUpstreamBillingRateSyncChange = (enabled: boolean) => {
upstreamBillingRateSyncEnabled.value = enabled
if (enabled) {
upstreamBillingAutoProbeEnabled.value = true
}
}
const handleUpstreamBillingAutoProbeChange = (enabled: boolean) => {
upstreamBillingAutoProbeEnabled.value = enabled
if (!enabled) {
upstreamBillingRateSyncEnabled.value = false
}
}
const statusOptions = computed(() => {
const options = [
{ value: 'active', label: t('common.active') },
@@ -3306,6 +3357,8 @@ const syncFormFromAccount = (newAccount: Account | null) => {
autoPause5hDisabled.value = extra?.auto_pause_5h_disabled === true
autoPause7dDisabled.value = extra?.auto_pause_7d_disabled === true
upstreamBillingAutoProbeEnabled.value = extra?.upstream_billing_probe_enabled === true
upstreamBillingRateSyncEnabled.value =
upstreamBillingAutoProbeEnabled.value && extra?.upstream_billing_rate_sync_enabled === true
// Load OpenAI passthrough toggle (OpenAI OAuth/SetupToken/API Key)
openaiPassthroughEnabled.value = false
@@ -4086,6 +4139,13 @@ const handleSubmit = async () => {
updatePayload.load_factor = 0
}
updatePayload.auto_pause_on_expired = autoPauseOnExpired.value
if (props.account.type === 'apikey') {
updatePayload.upstream_billing_probe_enabled = upstreamBillingAutoProbeEnabled.value
updatePayload.upstream_billing_rate_sync_enabled = upstreamBillingRateSyncEnabled.value
if (upstreamBillingRateSyncEnabled.value) {
delete updatePayload.rate_multiplier
}
}
// For apikey type, handle credentials update
if (props.account.type === 'apikey') {
@@ -4651,7 +4711,8 @@ const handleSubmit = async () => {
// 上游倍率自动探测对全部 API-key 平台开放(sub2api 上游即可应答),
// Bedrock 凭证无静态 Key 不参与。
if (props.account.type === 'apikey') {
newExtra.upstream_billing_probe_enabled = upstreamBillingAutoProbeEnabled.value
delete newExtra.upstream_billing_probe_enabled
delete newExtra.upstream_billing_rate_sync_enabled
}
// Total quota
if (editQuotaLimit.value != null && editQuotaLimit.value > 0) {
@@ -88,6 +88,7 @@ import { computed } from 'vue'
import { useI18n } from 'vue-i18n'
import HelpTooltip from '@/components/common/HelpTooltip.vue'
import Icon from '@/components/icons/Icon.vue'
import { formatMultiplier } from '@/utils/formatters'
import type { Account, UpstreamBillingProbeSnapshot } from '@/types'
const props = withDefaults(defineProps<{
@@ -191,7 +192,7 @@ const elapsedSinceLastSuccess = computed(() => {
const effectiveRate = computed(() => {
if (!validTimestamps.value || stale.value || !['ok', 'failed'].includes(snapshot.value?.status ?? '')) return '-'
const value = currentEffectiveRate.value
return value == null ? '-' : `${Number(value.toPrecision(12))}x`
return value == null ? '-' : `${formatMultiplier(value)}x`
})
const statusLabel = computed(() => {
if (!snapshot.value) return t('admin.accounts.upstreamBilling.notProbed')
@@ -4,9 +4,13 @@ import BulkEditAccountModal from '../BulkEditAccountModal.vue'
import ModelWhitelistSelector from '../ModelWhitelistSelector.vue'
import { adminAPI } from '@/api/admin'
const { showError } = vi.hoisted(() => ({
showError: vi.fn()
}))
vi.mock('@/stores/app', () => ({
useAppStore: () => ({
showError: vi.fn(),
showError,
showSuccess: vi.fn(),
showInfo: vi.fn()
})
@@ -77,6 +81,7 @@ describe('BulkEditAccountModal', () => {
beforeEach(() => {
vi.mocked(adminAPI.accounts.bulkUpdate).mockReset()
vi.mocked(adminAPI.accounts.checkMixedChannelRisk).mockReset()
showError.mockReset()
vi.mocked(adminAPI.accounts.bulkUpdate).mockResolvedValue({
success: 2,
@@ -88,6 +93,33 @@ describe('BulkEditAccountModal', () => {
} as any)
})
it('批量修改倍率时提示自动同步账号需要先关闭同步', async () => {
const wrapper = mountModal()
expect(wrapper.find('[data-testid="bulk-rate-sync-warning"]').exists()).toBe(false)
await wrapper.get('#bulk-edit-rate-multiplier-enabled').setValue(true)
expect(wrapper.get('[data-testid="bulk-rate-sync-warning"]').text()).toContain(
'admin.accounts.bulkEdit.rateSyncWarning'
)
})
it('后端拒绝修改同步账号倍率时展示专用错误', async () => {
vi.mocked(adminAPI.accounts.bulkUpdate).mockRejectedValueOnce({
status: 409,
reason: 'UPSTREAM_BILLING_RATE_SYNC_BULK_CONFLICT',
metadata: { count: '2' },
message: 'conflict'
})
const wrapper = mountModal()
await wrapper.get('#bulk-edit-rate-multiplier-enabled').setValue(true)
await wrapper.get('#bulk-edit-account-form').trigger('submit.prevent')
await flushPromises()
expect(showError).toHaveBeenCalledWith('admin.accounts.bulkEdit.rateSyncConflict')
})
it('antigravity 白名单包含 Gemini 图片模型且过滤掉普通 GPT 模型', async () => {
const wrapper = mountModal()
const selector = wrapper.findComponent(ModelWhitelistSelector)
@@ -654,7 +654,10 @@ describe('EditAccountModal', () => {
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)
expect(updateAccountMock.mock.calls[0]?.[1]?.upstream_billing_probe_enabled).toBe(true)
expect(updateAccountMock.mock.calls[0]?.[1]?.extra).not.toHaveProperty(
'upstream_billing_probe_enabled'
)
})
it('exposes the upstream billing auto-probe toggle for non-OpenAI API-key accounts', async () => {
@@ -676,7 +679,92 @@ describe('EditAccountModal', () => {
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)
expect(updateAccountMock.mock.calls[0]?.[1]?.upstream_billing_probe_enabled).toBe(true)
})
it('enabling rate sync also enables probing and stops submitting a manual rate', async () => {
const account = buildAccount()
updateAccountMock.mockReset()
checkMixedChannelRiskMock.mockReset()
checkMixedChannelRiskMock.mockResolvedValue({ has_risk: false })
updateAccountMock.mockResolvedValue(account)
const wrapper = mountModal(account)
const syncToggle = wrapper.get('[data-testid="upstream-billing-rate-sync"]')
const probeToggle = wrapper.get('[data-testid="upstream-billing-auto-probe"]')
const rateInput = wrapper.get<HTMLInputElement>('[data-testid="account-rate-multiplier"]')
expect(syncToggle.attributes('aria-checked')).toBe('false')
expect(probeToggle.attributes('aria-checked')).toBe('false')
expect(rateInput.element.disabled).toBe(false)
expect(wrapper.text()).toContain('admin.accounts.billingRateMultiplierHint')
expect(wrapper.text()).not.toContain('admin.accounts.upstreamBilling.syncRateManagedHint')
await syncToggle.trigger('click')
expect(syncToggle.attributes('aria-checked')).toBe('true')
expect(probeToggle.attributes('aria-checked')).toBe('true')
expect(rateInput.element.disabled).toBe(true)
expect(wrapper.text()).toContain('admin.accounts.upstreamBilling.syncRateManagedHint')
expect(wrapper.text()).not.toContain('admin.accounts.billingRateMultiplierHint')
await wrapper.get('form#edit-account-form').trigger('submit.prevent')
const payload = updateAccountMock.mock.calls[0]?.[1]
expect(payload?.upstream_billing_probe_enabled).toBe(true)
expect(payload?.upstream_billing_rate_sync_enabled).toBe(true)
expect(payload).not.toHaveProperty('rate_multiplier')
})
it('disabling probing also disables rate sync and restores manual rate editing', async () => {
const account = buildAccount()
account.extra = {
upstream_billing_probe_enabled: true,
upstream_billing_rate_sync_enabled: true
}
updateAccountMock.mockReset()
checkMixedChannelRiskMock.mockReset()
checkMixedChannelRiskMock.mockResolvedValue({ has_risk: false })
updateAccountMock.mockResolvedValue(account)
const wrapper = mountModal(account)
const syncToggle = wrapper.get('[data-testid="upstream-billing-rate-sync"]')
const probeToggle = wrapper.get('[data-testid="upstream-billing-auto-probe"]')
const rateInput = wrapper.get<HTMLInputElement>('[data-testid="account-rate-multiplier"]')
expect(syncToggle.attributes('aria-checked')).toBe('true')
expect(rateInput.element.disabled).toBe(true)
await probeToggle.trigger('click')
expect(probeToggle.attributes('aria-checked')).toBe('false')
expect(syncToggle.attributes('aria-checked')).toBe('false')
expect(rateInput.element.disabled).toBe(false)
await wrapper.get('form#edit-account-form').trigger('submit.prevent')
const payload = updateAccountMock.mock.calls[0]?.[1]
expect(payload?.upstream_billing_probe_enabled).toBe(false)
expect(payload?.upstream_billing_rate_sync_enabled).toBe(false)
expect(payload?.rate_multiplier).toBe(1)
})
it('disabling only rate sync keeps automatic probing enabled', async () => {
const account = buildAccount()
account.extra = {
upstream_billing_probe_enabled: true,
upstream_billing_rate_sync_enabled: true
}
updateAccountMock.mockReset()
checkMixedChannelRiskMock.mockReset()
checkMixedChannelRiskMock.mockResolvedValue({ has_risk: false })
updateAccountMock.mockResolvedValue(account)
const wrapper = mountModal(account)
await wrapper.get('[data-testid="upstream-billing-rate-sync"]').trigger('click')
expect(wrapper.get('[data-testid="upstream-billing-auto-probe"]').attributes('aria-checked')).toBe(
'true'
)
await wrapper.get('form#edit-account-form').trigger('submit.prevent')
const payload = updateAccountMock.mock.calls[0]?.[1]
expect(payload?.upstream_billing_probe_enabled).toBe(true)
expect(payload?.upstream_billing_rate_sync_enabled).toBe(false)
expect(payload?.rate_multiplier).toBe(1)
})
it('clears OpenAI APIKey Responses override when set back to auto', async () => {
@@ -88,11 +88,11 @@ describe('UpstreamBillingRateCell', () => {
}
})
expect(wrapper.text()).toContain('0.6x')
expect(wrapper.text()).toContain('0.60x')
await wrapper.setProps({ now: Date.parse('2026-07-13T01:00:00Z') })
expect(wrapper.text()).toContain('0.9x')
expect(wrapper.text()).toContain('0.90x')
await wrapper.setProps({ now: Date.parse('2026-07-13T10:00:00Z') })
expect(wrapper.text()).toContain('0.6x')
expect(wrapper.text()).toContain('0.60x')
expect(wrapper.text()).not.toContain('admin.accounts.upstreamBilling.latest')
expect(wrapper.get('[data-testid="upstream-billing-probe"]').text()).toBe('')
expect(wrapper.get('[data-testid="upstream-billing-probe"]').attributes('aria-label')).toBe(
@@ -132,11 +132,11 @@ describe('UpstreamBillingRateCell', () => {
}
})
})
expect(wrapper.text()).toContain('0.6x')
expect(wrapper.text()).toContain('0.60x')
expect(wrapper.text()).toContain('admin.accounts.upstreamBilling.failed')
await wrapper.setProps({ now: Date.parse('2026-07-13T01:00:00Z') })
expect(wrapper.text()).toContain('0.9x')
expect(wrapper.text()).toContain('0.90x')
expect(wrapper.text()).not.toContain('admin.accounts.upstreamBilling.stale')
await wrapper.setProps({ now: Date.parse('2026-07-13T01:00:00.001Z') })
@@ -312,7 +312,7 @@ describe('UpstreamBillingRateCell', () => {
await wrapper.setProps({ account: malformedAccount({}, { received_at: 'not-a-time' }) })
expect(wrapper.get('[data-testid="upstream-billing-rate"]').text()).toBe('admin.accounts.upstreamBilling.stale')
await wrapper.setProps({ account: malformedAccount({}, { received_at: '2026-07-13T00:31:00Z' }) })
expect(wrapper.get('[data-testid="upstream-billing-rate"]').text()).toBe('0.6x')
expect(wrapper.get('[data-testid="upstream-billing-rate"]').text()).toBe('0.60x')
await wrapper.setProps({ account: malformedAccount({}, { received_at: '2026-07-13T00:36:00Z' }) })
expect(wrapper.get('[data-testid="upstream-billing-rate"]').text()).toBe('admin.accounts.upstreamBilling.stale')
await wrapper.setProps({ account: malformedAccount({}, { fresh_until: '2026-07-12T23:59:00Z' }) })
@@ -225,7 +225,11 @@ export default {
upstreamBilling: {
trustWarning: 'This rate is declared by the upstream site for the current API key. Sub2API cannot verify that it matches actual charges. The upstream site or an intermediary may return forged, stale, or modified data. Verify it against bills, balance changes, and actual usage.',
autoProbe: 'Automatically probe upstream declared rate',
autoProbeHint: 'Probe this account\'s upstream declared rate on the global interval when global probing is enabled.',
autoProbeHint: 'Refresh the upstream declared rate on the global interval. This switch alone does not change the account rate.',
syncRate: 'Sync upstream declared rate',
syncRateHint: 'Update the account rate after each successful probe. Failed or invalid declarations leave it unchanged.',
syncRateManagedHint: 'The current rate is maintained automatically from the upstream declared rate.',
syncedRateTooltip: 'This account rate is synchronized from the upstream declared rate',
manualProbe: 'Probe upstream rate now',
stale: 'Stale',
unsupported: 'Unsupported',
@@ -420,6 +424,8 @@ export default {
failed: 'Bulk update failed',
noSelection: 'Please select accounts to edit',
noFieldsSelected: 'Select at least one field to update',
rateSyncWarning: 'Accounts with upstream rate sync enabled cannot be changed in bulk. Disable sync in the account editor first.',
rateSyncConflict: 'Cannot change account rates: {count} target account(s) have upstream rate sync enabled.',
mixedPlatformWarning: 'Selected accounts span multiple platforms ({platforms}). Model mapping presets shown are combined — ensure mappings are appropriate for each platform.'
},
bulkDeleteTitle: 'Bulk Delete Accounts',
@@ -361,7 +361,7 @@ export default {
},
upstreamBillingProbe: {
title: 'Upstream Rate Auto Detection',
description: 'Periodically retrieve billing rates declared by upstream Sub2API sites connected to API key accounts.',
description: 'Periodically retrieve rates declared by upstream Sub2API sites. Account rates change only when the separate sync switch is enabled.',
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)',
@@ -178,7 +178,11 @@ export default {
upstreamBilling: {
trustWarning: '此倍率由上游站点针对当前 API Key 自行声明。Sub2API 无法验证该值是否与实际扣费一致;上游站点或中间代理可能返回伪造、过期或被篡改的数据。请结合账单、余额变化和实际用量自行核验。',
autoProbe: '自动探测上游声明倍率',
autoProbeHint: '启用后按全局探测周期查询此账号的上游声明倍率;全局探测关闭时不会执行。',
autoProbeHint: '启用后按全局周期刷新上游声明倍率;此开关本身不会修改账号倍率。',
syncRate: '同步上游声明倍率',
syncRateHint: '成功探测后自动更新账号倍率;探测失败或声明无效时保持不变。',
syncRateManagedHint: '当前倍率由上游声明倍率自动维护。',
syncedRateTooltip: '该账号倍率由上游声明倍率自动同步',
manualProbe: '立即探测上游倍率',
stale: '已过期',
unsupported: '不支持',
@@ -488,6 +492,8 @@ export default {
failed: '批量更新失败',
noSelection: '请选择要编辑的账号',
noFieldsSelected: '请至少选择一个要更新的字段',
rateSyncWarning: '已开启上游倍率同步的账号不能批量手工修改倍率,请先在账号编辑页关闭同步。',
rateSyncConflict: '无法修改账号倍率:{count} 个目标账号已开启上游倍率同步。',
mixedPlatformWarning: '所选账号跨越多个平台({platforms})。显示的模型映射预设为合并结果——请确保映射对每个平台都适用。'
},
bulkDeleteTitle: '批量删除账号',
@@ -354,7 +354,7 @@ export default {
},
upstreamBillingProbe: {
title: '上游倍率自动探测',
description: '定期获取 API Key 账号所连接上游 Sub2API 站点声明的计费倍率。',
description: '定期获取 API Key 账号所连接上游 Sub2API 站点声明的计费倍率;只有另行开启“同步上游声明倍率”的账号才会更新账号倍率。',
enabled: '启用全局自动探测',
enabledHint: '开启后,仅对账号自身已启用自动检测的账号执行定时探测;关闭后停止所有定时探测,手动探测不受影响。',
intervalMinutes: '探测周期(分钟)',
+3
View File
@@ -1051,6 +1051,7 @@ export interface Account {
model_rate_limits?: Record<string, { rate_limited_at: string; rate_limit_reset_at: string }>
antigravity_credits_overages?: Record<string, { activated_at: string; active_until: string }>
upstream_billing_probe_enabled?: boolean
upstream_billing_rate_sync_enabled?: boolean
upstream_billing_probe?: UpstreamBillingProbeSnapshot
} & Record<string, unknown>)
proxy_id: number | null
@@ -1354,6 +1355,8 @@ export interface UpdateAccountRequest {
group_ids?: number[]
expires_at?: number | null
auto_pause_on_expired?: boolean
upstream_billing_probe_enabled?: boolean
upstream_billing_rate_sync_enabled?: boolean
confirm_mixed_channel_risk?: boolean
}
+21 -4
View File
@@ -342,8 +342,17 @@
</div>
</template>
<template #cell-rate_multiplier="{ row }">
<span class="text-sm font-mono text-gray-700 dark:text-gray-300">
{{ (row.rate_multiplier ?? 1).toFixed(2) }}x
<span class="inline-flex items-center gap-1 text-sm font-mono text-gray-700 dark:text-gray-300">
<span>{{ formatMultiplier(row.rate_multiplier ?? 1) }}x</span>
<span
v-if="row.extra?.upstream_billing_rate_sync_enabled === true"
class="inline-flex cursor-help text-emerald-600 dark:text-emerald-400"
:aria-label="t('admin.accounts.upstreamBilling.syncedRateTooltip')"
:title="t('admin.accounts.upstreamBilling.syncedRateTooltip')"
data-testid="account-rate-sync-indicator"
>
<Icon name="sync" size="xs" />
</span>
</span>
</template>
<template #header-upstream_billing_rate="{ column }">
@@ -516,6 +525,7 @@ import { proxyExpiryBadgeClass, proxyExpiryLabelKey } from '@/utils/proxyExpiry'
import { extractApiErrorMessage } from '@/utils/apiError'
import { sanitizeUrl } from '@/utils/url'
import { getFloatingPanelPosition } from '@/utils/floatingPanel'
import { formatMultiplier } from '@/utils/formatters'
import type { Account, AccountPlatform, AccountSchedulerGroupScore, AccountType, Proxy as AccountProxy, AdminGroup, WindowStats, ClaudeModel, UpstreamBillingProbeSnapshot } from '@/types'
const { t } = useI18n()
@@ -1591,7 +1601,7 @@ const handleBulkProbeUpstreamBilling = async () => {
patched = true
}
})
if (patched) await refreshUpstreamBillingSortedList(true)
if (patched) await refreshAccountsAfterUpstreamBillingProbe()
const failed = results.filter(result => result.error).length
if (failed > 0) {
appStore.showError(t('admin.accounts.upstreamBilling.batchPartial', { success: results.length - failed, failed }))
@@ -1890,6 +1900,13 @@ const patchUpstreamBillingSnapshot = (accountID: number, snapshot: UpstreamBilli
extra: { ...account.extra, upstream_billing_probe: snapshot }
})
}
const refreshAccountsAfterUpstreamBillingProbe = async () => {
try {
await load()
} catch (error) {
console.error('Failed to refresh accounts after upstream billing probe:', error)
}
}
const handleProbeUpstreamBilling = async (account: Account) => {
if (probingUpstreamBilling.has(account.id)) return
probingUpstreamBilling.add(account.id)
@@ -1897,7 +1914,7 @@ const handleProbeUpstreamBilling = async (account: Account) => {
const result = await adminAPI.accounts.probeUpstreamBilling(account.id)
if (result.snapshot) {
patchUpstreamBillingSnapshot(account.id, result.snapshot)
await refreshUpstreamBillingSortedList(true)
await refreshAccountsAfterUpstreamBillingProbe()
}
} catch (error) {
console.error('Failed to probe upstream billing:', error)
@@ -10,7 +10,10 @@ const {
getUpstreamBillingProbeSettings,
getAllProxies,
getAllGroups,
probeUpstreamBillingBatch
probeUpstreamBilling,
probeUpstreamBillingBatch,
showError,
showSuccess
} = vi.hoisted(() => ({
listAccounts: vi.fn(),
listWithEtag: vi.fn(),
@@ -18,7 +21,10 @@ const {
getUpstreamBillingProbeSettings: vi.fn(),
getAllProxies: vi.fn(),
getAllGroups: vi.fn(),
probeUpstreamBillingBatch: vi.fn()
probeUpstreamBilling: vi.fn(),
probeUpstreamBillingBatch: vi.fn(),
showError: vi.fn(),
showSuccess: vi.fn()
}))
vi.mock('@/api/admin', () => ({
@@ -31,6 +37,7 @@ vi.mock('@/api/admin', () => ({
delete: vi.fn(),
batchClearError: vi.fn(),
batchRefresh: vi.fn(),
probeUpstreamBilling,
probeUpstreamBillingBatch,
toggleSchedulable: vi.fn()
},
@@ -45,8 +52,8 @@ vi.mock('@/api/admin', () => ({
vi.mock('@/stores/app', () => ({
useAppStore: () => ({
showError: vi.fn(),
showSuccess: vi.fn(),
showError,
showSuccess,
showInfo: vi.fn()
})
}))
@@ -75,6 +82,19 @@ const DataTableStub = {
<div v-for="row in data" :key="row.id">
<div data-test="select-row"><slot name="cell-select" :row="row" /></div>
<slot name="cell-created_at" :value="row.created_at" :row="row" />
<div data-test="account-rate"><slot name="cell-rate_multiplier" :row="row" /></div>
</div>
</div>
`
}
const ProbeDataTableStub = {
props: ['data'],
template: `
<div>
<div v-for="row in data" :key="row.id">
<div data-test="account-rate"><slot name="cell-rate_multiplier" :row="row" /></div>
<slot name="cell-upstream_billing_rate" :row="row" />
</div>
</div>
`
@@ -111,7 +131,10 @@ describe('admin AccountsView bulk edit scope', () => {
getUpstreamBillingProbeSettings.mockReset()
getAllProxies.mockReset()
getAllGroups.mockReset()
probeUpstreamBilling.mockReset()
probeUpstreamBillingBatch.mockReset()
showError.mockReset()
showSuccess.mockReset()
listAccounts.mockResolvedValue({
items: [],
@@ -129,6 +152,7 @@ describe('admin AccountsView bulk edit scope', () => {
getUpstreamBillingProbeSettings.mockResolvedValue({ enabled: true, interval_minutes: 30 })
getAllProxies.mockResolvedValue([])
getAllGroups.mockResolvedValue([])
probeUpstreamBilling.mockResolvedValue({})
probeUpstreamBillingBatch.mockResolvedValue([])
})
@@ -376,33 +400,110 @@ describe('admin AccountsView bulk edit scope', () => {
expect(probeUpstreamBillingBatch).toHaveBeenCalledWith([7, 11])
})
it('reloads the server-sorted list after a batch probe changes a snapshot', async () => {
localStorage.setItem('account-table-sort', JSON.stringify({ key: 'upstream_billing_rate', order: 'asc' }))
const account = (id: number) => ({
it('refreshes the current page after a batch probe and displays the synced rate', async () => {
const account = (id: number, rateMultiplier: number) => ({
id,
name: `account-${id}`,
platform: 'openai',
type: 'apikey',
status: 'active',
schedulable: true,
rate_multiplier: rateMultiplier,
created_at: '2026-07-13T00:00:00Z',
updated_at: '2026-07-13T00:00:00Z'
})
listAccounts
.mockResolvedValueOnce({ items: [account(7)], total: 1, page: 1, page_size: 20, pages: 1 })
.mockResolvedValueOnce({ items: [account(7)], total: 1, page: 1, page_size: 20, pages: 1 })
.mockResolvedValueOnce({ items: [account(7, 0.25)], total: 2, page: 1, page_size: 1, pages: 2 })
.mockResolvedValueOnce({ items: [account(11, 0.25)], total: 2, page: 2, page_size: 1, pages: 2 })
.mockResolvedValueOnce({ items: [account(11, 0.065)], total: 2, page: 2, page_size: 1, pages: 2 })
probeUpstreamBillingBatch.mockResolvedValue([
{
account_id: 7,
account_id: 11,
snapshot: {
status: 'ok',
data: { effective_rate_multiplier: 0.5 },
data: { effective_rate_multiplier: 0.065 },
last_attempt_at: '2026-07-13T00:00:00Z',
next_probe_at: '2026-07-13T00:30:00Z'
}
}
])
const wrapper = mount(AccountsView, {
global: {
stubs: {
AppLayout: { template: '<div><slot /></div>' },
TablePageLayout: { template: '<div><slot name="table" /><slot name="pagination" /></div>' },
DataTable: DataTableStub,
AccountBulkActionsBar: AccountBulkActionsBarStub,
AccountTableActions: true,
AccountTableFilters: true,
AccountActionMenu: true,
Pagination: PaginationStub,
ConfirmDialog: true,
ImportDataModal: true,
ReAuthAccountModal: true,
AccountTestModal: true,
AccountStatsModal: true,
ScheduledTestsPanel: true,
SyncFromCrsModal: true,
TempUnschedStatusModal: true,
ErrorPassthroughRulesModal: true,
TLSFingerprintProfilesModal: true,
CreateAccountModal: true,
EditAccountModal: true,
BulkEditAccountModal: BulkEditAccountModalStub,
PlatformTypeBadge: true,
AccountCapacityCell: true,
AccountStatusIndicator: true,
AccountTodayStatsCell: true,
AccountGroupsCell: true,
AccountUsageCell: true,
Icon: true
}
}
})
await flushPromises()
await wrapper.get('[data-test="next-page"]').trigger('click')
await flushPromises()
await wrapper.get('[data-test="select-row"] input').trigger('change')
await wrapper.get('[data-test="probe-upstream-billing"]').trigger('click')
await flushPromises()
expect(probeUpstreamBillingBatch).toHaveBeenCalledWith([11])
expect(listAccounts).toHaveBeenCalledTimes(3)
expect(listAccounts.mock.calls[2]?.[0]).toBe(2)
expect(wrapper.get('[data-test="account-rate"]').text()).toBe('0.065x')
})
it('does not report a successful batch probe as failed when the list refresh fails', async () => {
const account = {
id: 7,
name: 'account-7',
platform: 'openai',
type: 'apikey',
status: 'active',
schedulable: true,
rate_multiplier: 0.25,
created_at: '2026-07-13T00:00:00Z',
updated_at: '2026-07-13T00:00:00Z'
}
listAccounts
.mockResolvedValueOnce({ items: [account], total: 1, page: 1, page_size: 20, pages: 1 })
.mockRejectedValueOnce(new Error('refresh failed'))
probeUpstreamBillingBatch.mockResolvedValue([
{
account_id: 7,
snapshot: {
status: 'ok',
data: { effective_rate_multiplier: 0.065 },
last_attempt_at: '2026-07-13T00:00:00Z',
next_probe_at: '2026-07-13T00:30:00Z'
}
}
])
const consoleError = vi.spyOn(console, 'error').mockImplementation(() => {})
const wrapper = mount(AccountsView, {
global: {
stubs: {
@@ -443,7 +544,78 @@ describe('admin AccountsView bulk edit scope', () => {
await wrapper.get('[data-test="probe-upstream-billing"]').trigger('click')
await flushPromises()
expect(probeUpstreamBillingBatch).toHaveBeenCalledWith([7])
expect(showError).not.toHaveBeenCalled()
expect(showSuccess).toHaveBeenCalledWith('admin.accounts.upstreamBilling.batchCompleted')
consoleError.mockRestore()
})
it('refreshes the account row after a successful single-account probe', async () => {
const account = (rateMultiplier: number) => ({
id: 7,
name: 'account-7',
platform: 'openai',
type: 'apikey',
status: 'active',
schedulable: true,
rate_multiplier: rateMultiplier,
extra: { upstream_billing_probe_enabled: true },
created_at: '2026-07-13T00:00:00Z',
updated_at: '2026-07-13T00:00:00Z'
})
listAccounts
.mockResolvedValueOnce({ items: [account(0.25)], total: 1, page: 1, page_size: 20, pages: 1 })
.mockResolvedValueOnce({ items: [account(0.065)], total: 1, page: 1, page_size: 20, pages: 1 })
probeUpstreamBilling.mockResolvedValue({
account_id: 7,
snapshot: {
status: 'ok',
data: { effective_rate_multiplier: 0.065 },
last_attempt_at: '2026-07-13T00:00:00Z',
next_probe_at: '2026-07-13T00:30:00Z'
}
})
const wrapper = mount(AccountsView, {
global: {
stubs: {
AppLayout: { template: '<div><slot /></div>' },
TablePageLayout: { template: '<div><slot name="table" /></div>' },
DataTable: ProbeDataTableStub,
AccountBulkActionsBar: true,
AccountTableActions: true,
AccountTableFilters: true,
AccountActionMenu: true,
Pagination: true,
ConfirmDialog: true,
ImportDataModal: true,
ReAuthAccountModal: true,
AccountTestModal: true,
AccountStatsModal: true,
ScheduledTestsPanel: true,
SyncFromCrsModal: true,
TempUnschedStatusModal: true,
ErrorPassthroughRulesModal: true,
TLSFingerprintProfilesModal: true,
CreateAccountModal: true,
EditAccountModal: true,
BulkEditAccountModal: true,
PlatformTypeBadge: true,
AccountCapacityCell: true,
AccountStatusIndicator: true,
AccountTodayStatsCell: true,
AccountGroupsCell: true,
AccountUsageCell: true,
Icon: true
}
}
})
await flushPromises()
await wrapper.get('[data-testid="upstream-billing-probe"]').trigger('click')
await flushPromises()
expect(probeUpstreamBilling).toHaveBeenCalledWith(7)
expect(listAccounts).toHaveBeenCalledTimes(2)
expect(wrapper.get('[data-test="account-rate"]').text()).toBe('0.065x')
})
})
@@ -75,6 +75,9 @@ const DataTableStub = {
<slot :name="'header-' + column.key" :column="column" />
</div>
</template>
<div v-for="row in data" :key="row.id" data-test="account-rate">
<slot name="cell-rate_multiplier" :row="row" />
</div>
</div>
`
}
@@ -189,4 +192,35 @@ describe('admin AccountsView usage windows hint', () => {
const columns = wrapper.getComponent(DataTableStub).props('columns') as Array<{ key: string; sortable: boolean }>
expect(columns.find(column => column.key === 'upstream_billing_rate')?.sortable).toBe(true)
})
it('shows account multipliers with enough precision to match declared rates', async () => {
listAccounts.mockResolvedValueOnce({
items: [{
id: 7,
name: 'precision-account',
platform: 'gemini',
type: 'apikey',
status: 'active',
schedulable: true,
rate_multiplier: 0.065,
extra: {
upstream_billing_probe_enabled: true,
upstream_billing_rate_sync_enabled: true
},
created_at: '2026-07-13T00:00:00Z',
updated_at: '2026-07-13T00:00:00Z'
}],
total: 1,
page: 1,
page_size: 20,
pages: 1
})
const wrapper = mountView()
await flushPromises()
expect(wrapper.get('[data-test="account-rate"]').text()).toBe('0.065x')
const indicator = wrapper.get('[data-testid="account-rate-sync-indicator"]')
expect(indicator.attributes('title')).toBe('admin.accounts.upstreamBilling.syncedRateTooltip')
})
})