diff --git a/backend/internal/handler/admin/account_handler.go b/backend/internal/handler/admin/account_handler.go index a97049722..9b80dd05d 100644 --- a/backend/internal/handler/admin/account_handler.go +++ b/backend/internal/handler/admin/account_handler.go @@ -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 { diff --git a/backend/internal/handler/admin/account_handler_mixed_channel_test.go b/backend/internal/handler/admin/account_handler_mixed_channel_test.go index bc4a3cf49..9e29b5893 100644 --- a/backend/internal/handler/admin/account_handler_mixed_channel_test.go +++ b/backend/internal/handler/admin/account_handler_mixed_channel_test.go @@ -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{ diff --git a/backend/internal/handler/admin/admin_service_stub_test.go b/backend/internal/handler/admin/admin_service_stub_test.go index 0eb661b93..4c3beb9b8 100644 --- a/backend/internal/handler/admin/admin_service_stub_test.go +++ b/backend/internal/handler/admin/admin_service_stub_test.go @@ -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 } diff --git a/backend/internal/repository/account_repo.go b/backend/internal/repository/account_repo.go index 1e70bf2c4..92f098768 100644 --- a/backend/internal/repository/account_repo.go +++ b/backend/internal/repository/account_repo.go @@ -58,6 +58,7 @@ var schedulerNeutralExtraKeyPrefixes = []string{ "codex_7d_", "passive_usage_", "upstream_billing_probe", + "upstream_billing_rate_sync", "ollama_cloud_usage", } @@ -397,16 +398,29 @@ 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) +// 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) error { +func (r *accountRepository) updateAccount( + ctx context.Context, + account *service.Account, + explicitProbeEnabled *bool, + explicitRateSyncEnabled *bool, + explicitRateMultiplier *float64, +) error { if account == nil { return nil } @@ -430,7 +444,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 +473,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 +506,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 +570,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 +603,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 +628,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 +639,7 @@ func lockAndMergeAccountProbeExtra(ctx context.Context, client *dbent.Client, ac &ollamaGroupIdentityUnchanged, &ollamaProxyIdentityUnchanged, ¤tEnabled, + ¤tRateSyncEnabled, ¤tSnapshot, ¤tOllamaSession, ¤tOllamaAutoRefresh, @@ -617,6 +654,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 +662,53 @@ 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 + } + // 同步依赖探测,方向是单向的:探测关闭(或探测键缺失)一律把同步归零。 + // 不做反向推导——由 rate_sync=true 推出 probe=true 会让一条"同步开、探测键 + // 缺失"的僵尸记录在任意一次无关编辑时静默打开周期性外呼。需要同时打开两个 + // 开关的调用方(管理端编辑)自己显式传 explicitProbeEnabled=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 +2597,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 +2626,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 +2659,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 +2681,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 +2698,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 } diff --git a/backend/internal/repository/account_repo_ollama_cloud_usage_integration_test.go b/backend/internal/repository/account_repo_ollama_cloud_usage_integration_test.go index 448eb42a6..868901584 100644 --- a/backend/internal/repository/account_repo_ollama_cloud_usage_integration_test.go +++ b/backend/internal/repository/account_repo_ollama_cloud_usage_integration_test.go @@ -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) diff --git a/backend/internal/repository/account_repo_upstream_billing_probe_cas_test.go b/backend/internal/repository/account_repo_upstream_billing_probe_cas_test.go index 4de218791..f9fe01955 100644 --- a/backend/internal/repository/account_repo_upstream_billing_probe_cas_test.go +++ b/backend/internal/repository/account_repo_upstream_billing_probe_cas_test.go @@ -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()) diff --git a/backend/internal/repository/account_repo_upstream_billing_probe_update_test.go b/backend/internal/repository/account_repo_upstream_billing_probe_update_test.go index 38e786d4d..46d78ac5b 100644 --- a/backend/internal/repository/account_repo_upstream_billing_probe_update_test.go +++ b/backend/internal/repository/account_repo_upstream_billing_probe_update_test.go @@ -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) @@ -104,6 +104,102 @@ func TestLockAndMergeAccountProbeExtraUsesCurrentDatabaseSnapshot(t *testing.T) } } +func probeBoolPtr(value bool) *bool { + return &value +} + +// The probe switch drives the rate-sync switch and never the other way round: +// syncing depends on probing, so a row where sync is on but the probe key is +// missing must lose the sync flag instead of silently gaining periodic +// outbound calls on the next unrelated edit. +func TestLockAndMergeAccountProbeExtraNeverInfersProbeFromRateSync(t *testing.T) { + tests := []struct { + name string + databaseEnabled any + databaseRateSync any + explicitProbeEnabled *bool + explicitRateSync *bool + wantEnabled any + wantRateSync any + }{ + { + name: "sync on with missing probe key zeroes sync and keeps probing off", + databaseEnabled: nil, + databaseRateSync: []byte(`true`), + wantEnabled: nil, + wantRateSync: false, + }, + { + name: "sync on with probe off zeroes sync", + databaseEnabled: []byte(`false`), + databaseRateSync: []byte(`true`), + wantEnabled: false, + wantRateSync: false, + }, + { + name: "both on in database stay on", + databaseEnabled: []byte(`true`), + databaseRateSync: []byte(`true`), + wantEnabled: true, + wantRateSync: true, + }, + { + name: "admin enabling both explicitly still turns probing on", + databaseEnabled: nil, + databaseRateSync: nil, + explicitProbeEnabled: probeBoolPtr(true), + explicitRateSync: probeBoolPtr(true), + wantEnabled: true, + wantRateSync: true, + }, + { + name: "explicit probe disable clears sync", + databaseEnabled: []byte(`true`), + databaseRateSync: []byte(`true`), + explicitProbeEnabled: probeBoolPtr(false), + wantEnabled: false, + wantRateSync: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + client := dbent.NewClient(dbent.Driver(entsql.OpenDB(dialect.Postgres, db))) + t.Cleanup(func() { _ = client.Close() }) + + mock.ExpectQuery(`(?s)`+regexp.QuoteMeta("SELECT")+`.*`+regexp.QuoteMeta("FOR NO KEY UPDATE")). + WithArgs(int64(31), service.PlatformOpenAI, service.AccountTypeAPIKey, `{"api_key":"sk-test"}`, 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, tt.databaseEnabled, tt.databaseRateSync, nil, nil, nil, nil)) + + account := &service.Account{ + ID: 31, + Platform: service.PlatformOpenAI, + Type: service.AccountTypeAPIKey, + Credentials: map[string]any{"api_key": "sk-test"}, + } + got, err := lockAndMergeAccountProbeExtra( + context.Background(), client, account, tt.explicitProbeEnabled, tt.explicitRateSync, + ) + require.NoError(t, err) + if tt.wantEnabled == nil { + require.NotContains(t, got, service.UpstreamBillingProbeEnabledExtraKey) + } else { + require.Equal(t, tt.wantEnabled, got[service.UpstreamBillingProbeEnabledExtraKey]) + } + if tt.wantRateSync == nil { + require.NotContains(t, got, service.UpstreamBillingRateSyncEnabledExtraKey) + } else { + require.Equal(t, tt.wantRateSync, got[service.UpstreamBillingRateSyncEnabledExtraKey]) + } + require.NoError(t, mock.ExpectationsWereMet()) + }) + } +} + func TestLockAndMergeAccountProbeExtraProtectsOllamaManagedFields(t *testing.T) { for _, identityUnchanged := range []bool{true, false} { t.Run(map[bool]string{true: "same identity keeps snapshot", false: "changed identity clears snapshot"}[identityUnchanged], func(t *testing.T) { @@ -115,8 +211,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 +223,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]) @@ -265,7 +361,7 @@ func TestUpdateCredentialsAtomicallyClearsProbeForOpenAIAPIKeyIdentityChange(t * require.NoError(t, mock.ExpectationsWereMet()) } -func TestUpdateWithUpstreamBillingProbeEnabledRollsBackWhenOutboxFails(t *testing.T) { +func TestUpdateWithAccountBillingSettingsRollsBackWhenOutboxFails(t *testing.T) { db, mock, err := sqlmock.New() require.NoError(t, err) t.Cleanup(func() { _ = db.Close() }) @@ -275,13 +371,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() @@ -301,10 +397,12 @@ func TestUpdateWithUpstreamBillingProbeEnabledRollsBackWhenOutboxFails(t *testin Schedulable: true, } - err = repo.UpdateWithUpstreamBillingProbeEnabled(context.Background(), account, false) + probeDisabled := false + err = repo.UpdateWithAccountBillingSettings(context.Background(), account, &probeDisabled, nil, nil) 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()) } diff --git a/backend/internal/repository/upstream_billing_probe_persistence_integration_test.go b/backend/internal/repository/upstream_billing_probe_persistence_integration_test.go index 252fa47cf..6e7631ac0 100644 --- a/backend/internal/repository/upstream_billing_probe_persistence_integration_test.go +++ b/backend/internal/repository/upstream_billing_probe_persistence_integration_test.go @@ -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, ` diff --git a/backend/internal/server/api_contract_test.go b/backend/internal/server/api_contract_test.go index 673584360..f6b1a0699 100644 --- a/backend/internal/server/api_contract_test.go +++ b/backend/internal/server/api_contract_test.go @@ -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") } diff --git a/backend/internal/service/account_service.go b/backend/internal/service/account_service.go index 98263dca1..d6905d732 100644 --- a/backend/internal/service/account_service.go +++ b/backend/internal/service/account_service.go @@ -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. diff --git a/backend/internal/service/admin_account.go b/backend/internal/service/admin_account.go index 1238452ed..044a7ae87 100644 --- a/backend/internal/service/admin_account.go +++ b/backend/internal/service/admin_account.go @@ -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 { @@ -766,6 +792,13 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U if *input.RateMultiplier < 0 { return nil, errors.New("rate_multiplier must be >= 0") } + // 同步开启时倍率归上游所有,手工值活不过下一次成功探测(表现为"改了又自己 + // 变回去"),与批量路径一样直接拒绝。判断的是本次请求生效后的状态:上面 + // 已把请求携带的两个开关落进 account.Extra,所以"同一请求关闭同步 + 改倍率" + // (用户显式收回所有权)会走到这里时读到 false,正常放行。 + if upstreamBillingRateSyncEnabled(account) { + return nil, ErrUpstreamBillingRateSyncConflict + } account.RateMultiplier = input.RateMultiplier } if input.LoadFactor != nil { @@ -806,23 +839,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 +904,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 +930,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 +964,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 +1050,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 +1083,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 { diff --git a/backend/internal/service/admin_account_upstream_billing_probe_test.go b/backend/internal/service/admin_account_upstream_billing_probe_test.go index 88db65386..f76dde0b2 100644 --- a/backend/internal/service/admin_account_upstream_billing_probe_test.go +++ b/backend/internal/service/admin_account_upstream_billing_probe_test.go @@ -17,6 +17,101 @@ 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) + + // 手工倍率只有在同步不再开启时才被接受,所以同一请求先关闭同步再设值 + // (同步仍开启时的手工倍率由 TestUpdateAccountRejectsManualRateWhileRateSyncEnabled 覆盖)。 + zero := 0.0 + syncDisabled := false + updated, err = svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{ + RateSyncEnabled: &syncDisabled, + 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 +123,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 +170,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 +184,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 +295,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 +310,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 +390,149 @@ 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 TestUpdateAccountRejectsManualRateWhileRateSyncEnabled(t *testing.T) { + newRepo := func(accountID int64, extra map[string]any) *upstreamBillingProbeAccountRepo { + initialRate := 0.25 + return &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{ + accountID: { + ID: accountID, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + RateMultiplier: &initialRate, + Extra: extra, + }, + }} + } + manualRate := 3.5 + syncEnabled := map[string]any{ + UpstreamBillingProbeEnabledExtraKey: true, + UpstreamBillingRateSyncEnabledExtraKey: true, + } + + t.Run("sync enabled rejects manual rate", func(t *testing.T) { + accountID := int64(153) + repo := newRepo(accountID, mergeMap(nil, syncEnabled)) + + _, err := (&adminServiceImpl{accountRepo: repo}).UpdateAccount(context.Background(), accountID, &UpdateAccountInput{ + RateMultiplier: &manualRate, + }) + + require.ErrorIs(t, err, ErrUpstreamBillingRateSyncConflict) + require.Equal(t, 0.25, *repo.accounts[accountID].RateMultiplier) + }) + + t.Run("enabling sync in the same request rejects manual rate", func(t *testing.T) { + accountID := int64(154) + repo := newRepo(accountID, map[string]any{}) + enable := true + + _, err := (&adminServiceImpl{accountRepo: repo}).UpdateAccount(context.Background(), accountID, &UpdateAccountInput{ + RateSyncEnabled: &enable, + RateMultiplier: &manualRate, + }) + + require.ErrorIs(t, err, ErrUpstreamBillingRateSyncConflict) + require.Equal(t, 0.25, *repo.accounts[accountID].RateMultiplier) + }) + + // 用户显式收回所有权:同一请求关闭同步并改倍率必须放行。 + t.Run("disabling sync in the same request allows manual rate", func(t *testing.T) { + accountID := int64(155) + repo := newRepo(accountID, mergeMap(nil, syncEnabled)) + disable := false + + updated, err := (&adminServiceImpl{accountRepo: repo}).UpdateAccount(context.Background(), accountID, &UpdateAccountInput{ + RateSyncEnabled: &disable, + RateMultiplier: &manualRate, + }) + + require.NoError(t, err) + require.Equal(t, false, updated.Extra[UpstreamBillingRateSyncEnabledExtraKey]) + require.NotNil(t, updated.RateMultiplier) + require.Equal(t, manualRate, *updated.RateMultiplier) + }) + + t.Run("sync disabled allows manual rate", func(t *testing.T) { + accountID := int64(156) + repo := newRepo(accountID, map[string]any{UpstreamBillingProbeEnabledExtraKey: true}) + + updated, err := (&adminServiceImpl{accountRepo: repo}).UpdateAccount(context.Background(), accountID, &UpdateAccountInput{ + RateMultiplier: &manualRate, + }) + + require.NoError(t, err) + require.NotNil(t, updated.RateMultiplier) + require.Equal(t, manualRate, *updated.RateMultiplier) + }) +} + +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 +555,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 +599,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 +639,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 +660,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) }) diff --git a/backend/internal/service/admin_service.go b/backend/internal/service/admin_service.go index bd907125a..7e69ef258 100644 --- a/backend/internal/service/admin_service.go +++ b/backend/internal/service/admin_service.go @@ -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, diff --git a/backend/internal/service/admin_service_bulk_update_test.go b/backend/internal/service/admin_service_bulk_update_test.go index 63bf26fbd..206efb21f 100644 --- a/backend/internal/service/admin_service_bulk_update_test.go +++ b/backend/internal/service/admin_service_bulk_update_test.go @@ -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{ diff --git a/backend/internal/service/admin_service_duplicate_account_test.go b/backend/internal/service/admin_service_duplicate_account_test.go index 7b19dbf6f..7bd31778d 100644 --- a/backend/internal/service/admin_service_duplicate_account_test.go +++ b/backend/internal/service/admin_service_duplicate_account_test.go @@ -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) diff --git a/backend/internal/service/crs_sync_helpers_test.go b/backend/internal/service/crs_sync_helpers_test.go index 344ed19d1..effeca467 100644 --- a/backend/internal/service/crs_sync_helpers_test.go +++ b/backend/internal/service/crs_sync_helpers_test.go @@ -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) }) } diff --git a/backend/internal/service/crs_sync_service.go b/backend/internal/service/crs_sync_service.go index e1637accd..e5a72b4c4 100644 --- a/backend/internal/service/crs_sync_service.go +++ b/backend/internal/service/crs_sync_service.go @@ -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 { diff --git a/backend/internal/service/upstream_billing_probe.go b/backend/internal/service/upstream_billing_probe.go index 7f105445e..d66cd905d 100644 --- a/backend/internal/service/upstream_billing_probe.go +++ b/backend/internal/service/upstream_billing_probe.go @@ -8,6 +8,7 @@ import ( "errors" "fmt" "io" + "log/slog" "math" "math/rand/v2" "net/http" @@ -28,8 +29,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 +46,7 @@ const ( // /v1/sub2api/billing,按常规 interval 重排只会持续占满每周期 // upstreamBillingProbeMaxPerCycle 个名额。 upstreamBillingProbeUnsupportedDelayFactor = 8 + upstreamBillingProbeAccountRateScale = 10000.0 upstreamBillingProbeLeaderLockKey = "upstream:billing:probe:leader" upstreamBillingProbeLeaderLockTTL = 2 * time.Minute ) @@ -51,6 +54,21 @@ const ( // UpstreamBillingProbeMaxBatchSize limits one manual batch and one runner cycle. const UpstreamBillingProbeMaxBatchSize = upstreamBillingProbeMaxPerCycle +// upstreamBillingRateSyncMaxMultiplier bounds the value the automatic +// write-back may push into accounts.rate_multiplier. +// +// No other code path bounds that column from above — admins may type any +// non-negative number and the only ceiling is the DECIMAL(10,4) column itself +// (999999.9999). That ceiling is meaningless as a guard: rate_multiplier +// scales the per-request account cost that feeds quota_used, so a single +// declared 999999 would exhaust any account quota on the first request and +// poison cost reporting. 100 is picked as a deliberately generous bound: it is +// two orders of magnitude above the 1.0 default and far above any plausible +// upstream resale markup, so no legitimate declaration is rejected while an +// absurd or hostile one cannot reach the quota control plane unattended. +// It only constrains the automatic path; manual edits keep their old range. +const upstreamBillingRateSyncMaxMultiplier = 100.0 + var ( ErrUpstreamBillingProbeUnavailable = infraerrors.ServiceUnavailable( "UPSTREAM_BILLING_PROBE_UNAVAILABLE", "upstream billing probe is unavailable", @@ -61,6 +79,14 @@ 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", + ) + ErrUpstreamBillingRateSyncConflict = infraerrors.Conflict( + "UPSTREAM_BILLING_RATE_SYNC_CONFLICT", + "account rate multiplier cannot be changed while upstream billing rate sync is enabled", + ) ) const ( @@ -87,6 +113,12 @@ type UpstreamBillingProbeSnapshot struct { FailureCount int `json:"failure_count,omitempty"` HTTPStatus int `json:"http_status,omitempty"` LastError string `json:"last_error,omitempty"` + // SyncedRateMultiplier records the value this probe wrote into + // accounts.rate_multiplier. It is only set when the account opted into rate + // sync and the declared value passed the write-back range check, so the + // stored snapshot always answers "did this probe move the account rate, and + // to what" without a separate history table. + SyncedRateMultiplier *float64 `json:"synced_rate_multiplier,omitempty"` } // UpstreamBillingProbeResult is returned by manual probe endpoints. @@ -197,7 +229,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 +578,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) { @@ -644,9 +678,39 @@ 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 { + // 账号级值域与精度只在真要写回时才有影响:只观察上游声明、未开启同步的 + // 账号不因声明值不适配 accounts.rate_multiplier 而被记成探测失败并进入 + // 指数退避——探测本身成功了,原始声明照常存进快照供展示。 + var syncRate *float64 + previousRate := account.BillingRateMultiplier() + if upstreamBillingRateSyncEnabled(account) { + if value, valid := upstreamBillingProbeSyncRate(data); valid { + syncRate = &value + snapshot.SyncedRateMultiplier = &value + } else { + declared, _ := resolveAccountExtraNumber(data, "resolved_rate_multiplier") + slog.Warn("upstream_billing_rate_sync_rejected", + "source", "upstream_billing_probe", + "account_id", account.ID, + "declared_resolved_rate_multiplier", declared, + "max_rate_multiplier", upstreamBillingRateSyncMaxMultiplier, + "current_rate_multiplier", previousRate, + ) + } + } + if err := s.updateSnapshot(ctx, account, snapshot, syncRate); err != nil { return nil, err } + if syncRate != nil { + // 写回是后台任务的裸 SQL,不经过管理端路由,因此不会产生 audit_logs 行。 + // old_rate_multiplier 是本次探测开始时读到的值(写回的 CAS 不比对该列)。 + slog.Info("upstream_billing_rate_sync_applied", + "source", "upstream_billing_probe", + "account_id", account.ID, + "old_rate_multiplier", previousRate, + "new_rate_multiplier", *syncRate, + ) + } return snapshot, nil } @@ -686,18 +750,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 +868,37 @@ func upstreamBillingRateAt(data map[string]any, now time.Time) (float64, bool) { return base, true } +// upstreamBillingProbeSyncRate converts the declared multiplier into the value +// the automatic write-back may store in accounts.rate_multiplier, at the +// precision that column supports (DECIMAL(10,4)). +// +// It reads resolved_rate_multiplier, not effective_rate_multiplier: the +// effective value folds in the peak coefficient that happened to apply at the +// instant of the probe, so writing it would freeze one probe cycle's peak (or +// off-peak) factor into a static column, while display and scheduling +// recompute the peak factor for the current time through upstreamBillingRateAt. +// +// The accepted range is deliberately narrower than the column: +// - 0 is rejected. accountCost multiplies the request cost by this value, so +// an upstream-declared 0 would stop quota_used from ever growing and every +// admin-configured account quota and cost alert would silently stop +// working. Admins may still set 0 by hand; only the automatic path refuses. +// - anything above upstreamBillingRateSyncMaxMultiplier is rejected. +// +// A rejected declaration leaves the current multiplier untouched; the probe +// still records an OK snapshot carrying the raw declaration for display. +func upstreamBillingProbeSyncRate(data map[string]any) (float64, bool) { + value, ok := resolveAccountExtraNumber(data, "resolved_rate_multiplier") + if !ok || math.IsNaN(value) || math.IsInf(value, 0) { + return 0, false + } + rounded := math.Round(value*upstreamBillingProbeAccountRateScale) / upstreamBillingProbeAccountRateScale + if rounded <= 0 || rounded > upstreamBillingRateSyncMaxMultiplier { + 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 +965,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 +974,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 +1040,18 @@ func upstreamBillingProbeEnabled(account *Account) bool { return ok && enabled } +// upstreamBillingRateSyncEnabled is the probe-side pre-filter deciding whether +// a rate is even proposed for write-back. It is a necessary condition, not the +// authority: the repository CAS re-checks both switches against the row it +// updates, so a switch flipped between load and write can never sneak a rate in. +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() diff --git a/backend/internal/service/upstream_billing_probe_multiplatform_test.go b/backend/internal/service/upstream_billing_probe_multiplatform_test.go index a2b03a96f..018f2b719 100644 --- a/backend/internal/service/upstream_billing_probe_multiplatform_test.go +++ b/backend/internal/service/upstream_billing_probe_multiplatform_test.go @@ -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)) } diff --git a/backend/internal/service/upstream_billing_probe_test.go b/backend/internal/service/upstream_billing_probe_test.go index edd297af9..e032e9aa1 100644 --- a/backend/internal/service/upstream_billing_probe_test.go +++ b/backend/internal/service/upstream_billing_probe_test.go @@ -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,12 @@ 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))) + // 写回的是不含高峰因子的 resolved 倍率(0.6),不是探测那一刻含高峰的 + // effective 倍率(0.9)——否则一个探测周期的峰值会被冻结进静态列。 + require.NotNil(t, account.RateMultiplier) + require.Equal(t, 0.6, *account.RateMultiplier) + require.NotNil(t, snapshot.SyncedRateMultiplier) + require.Equal(t, 0.6, *snapshot.SyncedRateMultiplier) 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 +340,217 @@ func TestUpstreamBillingProbeSuccessPersistsSanitizedSnapshot(t *testing.T) { require.Equal(t, snapshot.Status, persisted.Status) } +func TestUpstreamBillingProbeSyncsResolvedRateForAllAPIKeyPlatforms(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 TestUpstreamBillingProbeSyncRateRangeAndPrecision(t *testing.T) { + tests := []struct { + name string + value float64 + want float64 + ok bool + }{ + {name: "round to four decimals", value: 0.07654, want: 0.0765, ok: true}, + {name: "maximum", value: upstreamBillingRateSyncMaxMultiplier, want: upstreamBillingRateSyncMaxMultiplier, ok: true}, + // 0 会让 accountCost 恒为 0,账号配额与成本告警全部静默失效, + // 自动写回一律拒绝(管理员手工设 0 仍然允许)。 + {name: "zero is rejected", value: 0, ok: false}, + {name: "positive below database precision rounds to zero", value: 0.00001, ok: false}, + {name: "just above the write-back ceiling", value: 100.0001, ok: false}, + {name: "column ceiling is far above the write-back ceiling", value: 999999.9999, ok: false}, + {name: "negative", value: -1, ok: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, ok := upstreamBillingProbeSyncRate(map[string]any{"resolved_rate_multiplier": tt.value}) + require.Equal(t, tt.ok, ok) + if tt.ok { + require.Equal(t, tt.want, got) + } + }) + } +} + +// 只读取 resolved(时间无关的基准倍率):effective 含探测那一刻的高峰系数, +// 写回它会把一个探测周期的峰值/谷值冻结进静态列。 +func TestUpstreamBillingProbeSyncRateIgnoresEffectiveRate(t *testing.T) { + got, ok := upstreamBillingProbeSyncRate(map[string]any{ + "resolved_rate_multiplier": 0.6, + "effective_rate_multiplier": 0.9, + }) + require.True(t, ok) + require.Equal(t, 0.6, got) + + _, ok = upstreamBillingProbeSyncRate(map[string]any{"effective_rate_multiplier": 0.9}) + require.False(t, ok) +} + +// 上游声明超出自动写回值域时保持原倍率,但探测本身是成功的: +// 快照照常记 ok,不累计 failure_count、不进入退避。 +func TestUpstreamBillingProbeKeepsRateWhenDeclarationOutOfSyncRange(t *testing.T) { + for _, tt := range []struct { + name string + declared string + }{ + {name: "zero", declared: "0"}, + {name: "above ceiling", declared: "1000"}, + } { + t.Run(tt.name, func(t *testing.T) { + initialRate := 0.25 + account := &Account{ + ID: 21, + Platform: PlatformOpenAI, + 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}} + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(fmt.Sprintf(`{ + "object":"sub2api.key_billing", + "schema_version":1, + "billing_scope":"token", + "group_rate_multiplier":%[1]s, + "resolved_rate_multiplier":%[1]s, + "peak_rate_enabled":false, + "effective_rate_multiplier":%[1]s, + "observed_at":"2026-07-13T01:00:00Z" + }`, tt.declared))), + }} + svc := newUpstreamBillingProbeTestService(repo, upstream, &upstreamBillingProbeSettingRepo{}) + + snapshot, err := svc.ProbeAccount(context.Background(), account.ID) + + require.NoError(t, err) + require.Equal(t, UpstreamBillingProbeStatusOK, snapshot.Status) + require.Zero(t, snapshot.FailureCount) + require.Nil(t, snapshot.SyncedRateMultiplier) + require.NotNil(t, account.RateMultiplier) + require.Equal(t, initialRate, *account.RateMultiplier) + // 原始声明仍进快照供展示。 + require.Equal(t, snapshot.Data["resolved_rate_multiplier"], snapshot.Data["effective_rate_multiplier"]) + }) + } +} + +// 未开启同步的账号只观察上游声明:声明值不适配 accounts.rate_multiplier +// 不得被记成探测失败(否则会累计 failure_count 并进入指数退避)。 +func TestUpstreamBillingProbeWithoutSyncIgnoresUnusableDeclaredRate(t *testing.T) { + account := &Account{ + ID: 22, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Concurrency: 1, + 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}} + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{ + "object":"sub2api.key_billing", + "schema_version":1, + "billing_scope":"token", + "group_rate_multiplier":0, + "resolved_rate_multiplier":0, + "peak_rate_enabled":false, + "effective_rate_multiplier":0, + "observed_at":"2026-07-13T01:00:00Z" + }`)), + }} + svc := newUpstreamBillingProbeTestService(repo, upstream, &upstreamBillingProbeSettingRepo{}) + + snapshot, err := svc.ProbeAccount(context.Background(), account.ID) + + require.NoError(t, err) + require.Equal(t, UpstreamBillingProbeStatusOK, snapshot.Status) + require.Zero(t, snapshot.FailureCount) + require.Empty(t, snapshot.LastError) + require.Nil(t, snapshot.SyncedRateMultiplier) + require.Equal(t, float64(0), snapshot.Data["resolved_rate_multiplier"]) + require.Nil(t, account.RateMultiplier) +} + func TestUpstreamBillingProbeRejectsMissingRequiredMultiplier(t *testing.T) { _, err := parseUpstreamBillingProbeResponse([]byte(`{ "object":"sub2api.key_billing", @@ -456,6 +690,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 +698,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 +731,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 +852,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 +866,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 +907,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) { diff --git a/frontend/src/components/account/BulkEditAccountModal.vue b/frontend/src/components/account/BulkEditAccountModal.vue index 02ddaca92..ad9998695 100644 --- a/frontend/src/components/account/BulkEditAccountModal.vue +++ b/frontend/src/components/account/BulkEditAccountModal.vue @@ -750,6 +750,14 @@ aria-labelledby="bulk-edit-rate-multiplier-label" />
{{ t('admin.accounts.billingRateMultiplierHint') }}
+
+
{{ t('admin.accounts.billingRateMultiplierHint') }}
+ ++ {{ + t( + upstreamBillingRateSyncEnabled + ? 'admin.accounts.upstreamBilling.syncRateManagedHint' + : 'admin.accounts.billingRateMultiplierHint' + ) + }} +
++ {{ t('admin.accounts.upstreamBilling.syncRate') }} +
++ {{ t('admin.accounts.upstreamBilling.syncRateHint') }} +
+