diff --git a/backend/internal/handler/admin/account_handler.go b/backend/internal/handler/admin/account_handler.go index a970497226..9b80dd05d1 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 bc4a3cf49b..9e29b5893b 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 0eb661b932..4c3beb9b8a 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 1e70bf2c44..274ebb39b7 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,40 @@ func (r *accountRepository) ListCRSAccountIDs(ctx context.Context) (map[string]i } func (r *accountRepository) Update(ctx context.Context, account *service.Account) error { - return r.updateAccount(ctx, account, nil) + return r.updateAccount(ctx, account, nil, nil, account.RateMultiplier) } // UpdateWithUpstreamBillingProbeEnabled applies an explicit probe switch in the // same row-lock transaction as the rest of an admin account edit. func (r *accountRepository) UpdateWithUpstreamBillingProbeEnabled(ctx context.Context, account *service.Account, enabled bool) error { - return r.updateAccount(ctx, account, &enabled) + var rateSyncEnabled *bool + if !enabled { + disabled := false + rateSyncEnabled = &disabled + } + return r.updateAccount(ctx, account, &enabled, rateSyncEnabled, nil) } -func (r *accountRepository) updateAccount(ctx context.Context, account *service.Account, explicitProbeEnabled *bool) error { +// UpdateWithAccountBillingSettings applies an admin account edit while +// preserving a concurrently probe-synchronized rate unless the request +// explicitly includes a manual rate. +func (r *accountRepository) UpdateWithAccountBillingSettings( + ctx context.Context, + account *service.Account, + probeEnabled *bool, + rateSyncEnabled *bool, + rateMultiplier *float64, +) error { + return r.updateAccount(ctx, account, probeEnabled, rateSyncEnabled, rateMultiplier) +} + +func (r *accountRepository) updateAccount( + ctx context.Context, + account *service.Account, + explicitProbeEnabled *bool, + explicitRateSyncEnabled *bool, + explicitRateMultiplier *float64, +) error { if account == nil { return nil } @@ -430,7 +455,14 @@ func (r *accountRepository) updateAccount(ctx context.Context, account *service. } } - updated, err := r.updateLockedAccount(ctx, client, account, explicitProbeEnabled) + updated, err := r.updateLockedAccount( + ctx, + client, + account, + explicitProbeEnabled, + explicitRateSyncEnabled, + explicitRateMultiplier, + ) if err != nil { return translatePersistenceError(err, service.ErrAccountNotFound, nil) } @@ -452,8 +484,15 @@ func (r *accountRepository) updateAccount(ctx context.Context, account *service. return nil } -func (r *accountRepository) updateLockedAccount(ctx context.Context, client *dbent.Client, account *service.Account, explicitProbeEnabled *bool) (*dbent.Account, error) { - extra, err := lockAndMergeAccountProbeExtra(ctx, client, account, explicitProbeEnabled) +func (r *accountRepository) updateLockedAccount( + ctx context.Context, + client *dbent.Client, + account *service.Account, + explicitProbeEnabled *bool, + explicitRateSyncEnabled *bool, + explicitRateMultiplier *float64, +) (*dbent.Account, error) { + extra, err := lockAndMergeAccountProbeExtra(ctx, client, account, explicitProbeEnabled, explicitRateSyncEnabled) if err != nil { return nil, err } @@ -478,8 +517,8 @@ func (r *accountRepository) updateLockedAccount(ctx context.Context, client *dbe SetSchedulable(schedulable). SetAutoPauseOnExpired(account.AutoPauseOnExpired) - if account.RateMultiplier != nil { - builder.SetRateMultiplier(*account.RateMultiplier) + if explicitRateMultiplier != nil { + builder.SetRateMultiplier(*explicitRateMultiplier) } if account.LoadFactor != nil { builder.SetLoadFactor(*account.LoadFactor) @@ -542,7 +581,13 @@ func (r *accountRepository) updateLockedAccount(ctx context.Context, client *dbe return builder.Save(ctx) } -func lockAndMergeAccountProbeExtra(ctx context.Context, client *dbent.Client, account *service.Account, explicitProbeEnabled *bool) (map[string]any, error) { +func lockAndMergeAccountProbeExtra( + ctx context.Context, + client *dbent.Client, + account *service.Account, + explicitProbeEnabled *bool, + explicitRateSyncEnabled *bool, +) (map[string]any, error) { credentials, err := json.Marshal(normalizeJSONMap(account.Credentials)) if err != nil { return nil, err @@ -569,6 +614,7 @@ func lockAndMergeAccountProbeExtra(ctx context.Context, client *dbent.Client, ac ), proxy_id IS NOT DISTINCT FROM $5, extra -> 'upstream_billing_probe_enabled', + extra -> 'upstream_billing_rate_sync_enabled', extra -> 'upstream_billing_probe', extra -> 'ollama_cloud_usage_session', extra -> 'ollama_cloud_usage_auto_refresh', @@ -593,6 +639,7 @@ func lockAndMergeAccountProbeExtra(ctx context.Context, client *dbent.Client, ac ollamaGroupIdentityUnchanged bool ollamaProxyIdentityUnchanged bool currentEnabled []byte + currentRateSyncEnabled []byte currentSnapshot []byte currentOllamaSession []byte currentOllamaAutoRefresh []byte @@ -603,6 +650,7 @@ func lockAndMergeAccountProbeExtra(ctx context.Context, client *dbent.Client, ac &ollamaGroupIdentityUnchanged, &ollamaProxyIdentityUnchanged, ¤tEnabled, + ¤tRateSyncEnabled, ¤tSnapshot, ¤tOllamaSession, ¤tOllamaAutoRefresh, @@ -617,6 +665,7 @@ func lockAndMergeAccountProbeExtra(ctx context.Context, client *dbent.Client, ac extra := copyJSONMap(normalizeJSONMap(account.Extra)) for _, key := range []string{ service.UpstreamBillingProbeEnabledExtraKey, + service.UpstreamBillingRateSyncEnabledExtraKey, service.UpstreamBillingProbeExtraKey, service.OllamaCloudUsageSessionExtraKey, service.OllamaCloudUsageAutoRefreshExtraKey, @@ -624,21 +673,52 @@ func lockAndMergeAccountProbeExtra(ctx context.Context, client *dbent.Client, ac } { delete(extra, key) } - probeExplicitlyDisabled := false probeAccount := service.IsUpstreamBillingProbeIdentity(account.Platform, account.Type) - if probeAccount && explicitProbeEnabled != nil { - extra[service.UpstreamBillingProbeEnabledExtraKey] = *explicitProbeEnabled - probeExplicitlyDisabled = !*explicitProbeEnabled - } else if probeAccount { + probeEnabled := false + probeEnabledPresent := false + if probeAccount { if enabled, ok, err := decodeAccountExtraJSON(currentEnabled); err != nil { return nil, err - } else if ok { - extra[service.UpstreamBillingProbeEnabledExtraKey] = enabled - if value, isBool := enabled.(bool); isBool && !value { - probeExplicitlyDisabled = true - } + } else if value, isBool := enabled.(bool); ok && isBool { + probeEnabled = value + probeEnabledPresent = true + } + if explicitProbeEnabled != nil { + probeEnabled = *explicitProbeEnabled + probeEnabledPresent = true } } + rateSyncEnabled := false + rateSyncEnabledPresent := false + if probeAccount { + if enabled, ok, err := decodeAccountExtraJSON(currentRateSyncEnabled); err != nil { + return nil, err + } else if value, isBool := enabled.(bool); ok && isBool { + rateSyncEnabled = value + rateSyncEnabledPresent = true + } + if explicitRateSyncEnabled != nil { + rateSyncEnabled = *explicitRateSyncEnabled + rateSyncEnabledPresent = true + } + if explicitProbeEnabled != nil && !*explicitProbeEnabled { + rateSyncEnabled = false + rateSyncEnabledPresent = true + } else if rateSyncEnabled { + probeEnabled = true + probeEnabledPresent = true + } + if !probeEnabled { + rateSyncEnabled = false + } + if probeEnabledPresent { + extra[service.UpstreamBillingProbeEnabledExtraKey] = probeEnabled + } + if rateSyncEnabledPresent { + extra[service.UpstreamBillingRateSyncEnabledExtraKey] = rateSyncEnabled + } + } + probeExplicitlyDisabled := probeEnabledPresent && !probeEnabled if identityUnchanged && !probeExplicitlyDisabled { if snapshot, ok, err := decodeAccountExtraJSON(currentSnapshot); err != nil { return nil, err @@ -2527,21 +2607,25 @@ func (r *accountRepository) UpdateUpstreamBillingProbeSnapshot( ctx context.Context, account *service.Account, snapshot *service.UpstreamBillingProbeSnapshot, + rateMultiplier *float64, ) error { if account == nil || snapshot == nil { return service.ErrAccountNilInput } + if snapshot.Status != service.UpstreamBillingProbeStatusOK { + rateMultiplier = nil + } if dbent.TxFromContext(ctx) == nil { tx, err := r.client.Tx(ctx) if errors.Is(err, dbent.ErrTxStarted) { - return r.updateUpstreamBillingProbeSnapshotInTx(ctx, account, snapshot) + return r.updateUpstreamBillingProbeSnapshotInTx(ctx, account, snapshot, rateMultiplier) } if err != nil { return err } defer func() { _ = tx.Rollback() }() - if err := r.updateUpstreamBillingProbeSnapshotInTx(dbent.NewTxContext(ctx, tx), account, snapshot); err != nil { + if err := r.updateUpstreamBillingProbeSnapshotInTx(dbent.NewTxContext(ctx, tx), account, snapshot, rateMultiplier); err != nil { return err } if err := tx.Commit(); err != nil { @@ -2552,13 +2636,14 @@ func (r *accountRepository) UpdateUpstreamBillingProbeSnapshot( r.syncSchedulerAccountSnapshot(ctx, account.ID) return nil } - return r.updateUpstreamBillingProbeSnapshotInTx(ctx, account, snapshot) + return r.updateUpstreamBillingProbeSnapshotInTx(ctx, account, snapshot, rateMultiplier) } func (r *accountRepository) updateUpstreamBillingProbeSnapshotInTx( ctx context.Context, account *service.Account, snapshot *service.UpstreamBillingProbeSnapshot, + rateMultiplier *float64, ) error { payload, err := json.Marshal(map[string]any{service.UpstreamBillingProbeExtraKey: snapshot}) if err != nil { @@ -2584,6 +2669,14 @@ func (r *accountRepository) updateUpstreamBillingProbeSnapshotInTx( if err != nil { return err } + var expectedRateSyncEnabled any + if account.Extra != nil { + expectedRateSyncEnabled = account.Extra[service.UpstreamBillingRateSyncEnabledExtraKey] + } + expectedRateSyncEnabledJSON, err := json.Marshal(expectedRateSyncEnabled) + if err != nil { + return err + } client := clientFromContext(ctx, r.client) proxyMatches, err := lockAndMatchProbeProxyIdentity(ctx, client, account) if err != nil { @@ -2598,7 +2691,16 @@ func (r *accountRepository) updateUpstreamBillingProbeSnapshotInTx( } result, err := client.ExecContext(ctx, ` UPDATE accounts - SET extra = COALESCE(extra, '{}'::jsonb) || $1::jsonb, updated_at = NOW() + SET + extra = COALESCE(extra, '{}'::jsonb) || $1::jsonb, + rate_multiplier = CASE + WHEN $10::numeric IS NOT NULL + AND extra @> '{"upstream_billing_probe_enabled": true}'::jsonb + AND extra @> '{"upstream_billing_rate_sync_enabled": true}'::jsonb + THEN $10::numeric + ELSE rate_multiplier + END, + updated_at = NOW() WHERE id = $2 AND platform = $3 AND type = $4 @@ -2606,8 +2708,9 @@ func (r *accountRepository) updateUpstreamBillingProbeSnapshotInTx( AND proxy_id IS NOT DISTINCT FROM $6 AND COALESCE(extra -> 'upstream_billing_probe', 'null'::jsonb) = $7::jsonb AND COALESCE(extra -> 'upstream_billing_probe_enabled', 'null'::jsonb) = $8::jsonb + AND COALESCE(extra -> 'upstream_billing_rate_sync_enabled', 'null'::jsonb) = $9::jsonb AND deleted_at IS NULL - `, string(payload), account.ID, account.Platform, account.Type, string(credentials), proxyID, string(expectedSnapshotJSON), string(expectedEnabledJSON)) + `, string(payload), account.ID, account.Platform, account.Type, string(credentials), proxyID, string(expectedSnapshotJSON), string(expectedEnabledJSON), string(expectedRateSyncEnabledJSON), rateMultiplier) if err != nil { return err } 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 448eb42a69..868901584b 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 4de2187915..f9fe019554 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 38e786d4d0..9abeecc798 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) @@ -115,8 +115,8 @@ func TestLockAndMergeAccountProbeExtraProtectsOllamaManagedFields(t *testing.T) mock.ExpectQuery(`(?s)`+regexp.QuoteMeta("SELECT")+`.*`+regexp.QuoteMeta("FOR NO KEY UPDATE")). WithArgs(int64(29), service.PlatformAnthropic, service.AccountTypeAPIKey, `{"api_key":"key","base_url":"https://ollama.com"}`, nil). - WillReturnRows(sqlmock.NewRows([]string{"identity_unchanged", "ollama_group_unchanged", "ollama_proxy_unchanged", "enabled", "snapshot", "ollama_session", "ollama_auto", "ollama_snapshot"}). - AddRow(identityUnchanged, identityUnchanged, true, nil, nil, []byte(`"local-ciphertext"`), []byte(`true`), []byte(`{"status":"ok"}`))) + WillReturnRows(sqlmock.NewRows([]string{"identity_unchanged", "ollama_group_unchanged", "ollama_proxy_unchanged", "enabled", "rate_sync_enabled", "snapshot", "ollama_session", "ollama_auto", "ollama_snapshot"}). + AddRow(identityUnchanged, identityUnchanged, true, nil, nil, nil, []byte(`"local-ciphertext"`), []byte(`true`), []byte(`{"status":"ok"}`))) account := &service.Account{ ID: 29, Platform: service.PlatformAnthropic, Type: service.AccountTypeAPIKey, @@ -127,7 +127,7 @@ func TestLockAndMergeAccountProbeExtraProtectsOllamaManagedFields(t *testing.T) service.OllamaCloudUsageSnapshotExtraKey: map[string]any{"status": "forged"}, }, } - got, err := lockAndMergeAccountProbeExtra(context.Background(), client, account, nil) + got, err := lockAndMergeAccountProbeExtra(context.Background(), client, account, nil, nil) require.NoError(t, err) if identityUnchanged { require.Equal(t, "local-ciphertext", got[service.OllamaCloudUsageSessionExtraKey]) @@ -275,13 +275,13 @@ func TestUpdateWithUpstreamBillingProbeEnabledRollsBackWhenOutboxFails(t *testin mock.ExpectBegin() mock.ExpectQuery(`(?s)`+regexp.QuoteMeta("SELECT")+`.*`+regexp.QuoteMeta("FOR NO KEY UPDATE")). WithArgs(int64(27), service.PlatformOpenAI, service.AccountTypeAPIKey, `{"api_key":"sk-test"}`, nil). - WillReturnRows(sqlmock.NewRows([]string{"identity_unchanged", "ollama_group_unchanged", "ollama_proxy_unchanged", "enabled", "snapshot", "ollama_session", "ollama_auto", "ollama_snapshot"}). - AddRow(true, false, true, []byte(`true`), []byte(`{"status":"ok"}`), nil, nil, nil)) + WillReturnRows(sqlmock.NewRows([]string{"identity_unchanged", "ollama_group_unchanged", "ollama_proxy_unchanged", "enabled", "rate_sync_enabled", "snapshot", "ollama_session", "ollama_auto", "ollama_snapshot"}). + AddRow(true, false, true, []byte(`true`), []byte(`true`), []byte(`{"status":"ok"}`), nil, nil, nil)) mock.ExpectExec(`(?s)UPDATE .*accounts.*SET.*WHERE .*id.*`). WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectQuery(`(?s)SELECT .* FROM "accounts" WHERE "id" = \$1`). WithArgs(int64(27)). - WillReturnRows(updatedAccountRows(27, `{"upstream_billing_probe_enabled":false}`)) + WillReturnRows(updatedAccountRows(27, `{"upstream_billing_probe_enabled":false,"upstream_billing_rate_sync_enabled":false}`)) mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox")).WillReturnError(errors.New("outbox failed")) mock.ExpectRollback() @@ -305,6 +305,7 @@ func TestUpdateWithUpstreamBillingProbeEnabledRollsBackWhenOutboxFails(t *testin require.EqualError(t, err, "outbox failed") require.Equal(t, false, account.Extra[service.UpstreamBillingProbeEnabledExtraKey]) + require.Equal(t, false, account.Extra[service.UpstreamBillingRateSyncEnabledExtraKey]) require.NotContains(t, account.Extra, service.UpstreamBillingProbeExtraKey) require.NoError(t, mock.ExpectationsWereMet()) } 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 252fa47cf9..6e7631ac0e 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 6735843609..f6b1a06992 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 98263dca1c..d6905d7329 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 1238452ed1..7b1b6152c4 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 { @@ -806,23 +832,40 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U } } - probeEnabledAppliedAtomically := false - if requestedProbeEnabledUpdate != nil && isUpstreamBillingProbeAccount(account) { - if updater, ok := s.accountRepo.(accountProbeEnabledAtomicUpdater); ok { - if err := updater.UpdateWithUpstreamBillingProbeEnabled(ctx, account, *requestedProbeEnabledUpdate); err != nil { - return nil, err - } - probeEnabledAppliedAtomically = true - } + billingSettingsAppliedAtomically := false + updater := s.accountBillingRepo + if updater == nil { + // Unit tests and narrow internal callers may construct adminServiceImpl + // directly; production wiring requires this capability through + // AdminAccountRepository. + updater, _ = s.accountRepo.(AccountBillingSettingsRepository) } - if !probeEnabledAppliedAtomically { + if updater != nil { + if err := updater.UpdateWithAccountBillingSettings( + ctx, + account, + requestedProbeEnabledUpdate, + requestedRateSyncEnabledUpdate, + input.RateMultiplier, + ); err != nil { + return nil, err + } + billingSettingsAppliedAtomically = true + } + if !billingSettingsAppliedAtomically { if err := s.accountRepo.Update(ctx, account); err != nil { return nil, err } - if requestedProbeEnabledUpdate != nil && isUpstreamBillingProbeAccount(account) { - if err := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{ - UpstreamBillingProbeEnabledExtraKey: *requestedProbeEnabledUpdate, - }); err != nil { + if (requestedProbeEnabledUpdate != nil || requestedRateSyncEnabledUpdate != nil) && + isUpstreamBillingProbeAccount(account) { + settings := make(map[string]any, 2) + if requestedProbeEnabledUpdate != nil { + settings[UpstreamBillingProbeEnabledExtraKey] = *requestedProbeEnabledUpdate + } + if requestedRateSyncEnabledUpdate != nil { + settings[UpstreamBillingRateSyncEnabledExtraKey] = *requestedRateSyncEnabledUpdate + } + if err := s.accountRepo.UpdateExtra(ctx, account.ID, settings); err != nil { return nil, err } } @@ -854,6 +897,9 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U // UpdateAccountExtra 仅对 Extra JSONB 做 key 级合并,避免覆盖其它运行态键 // (如 model_rate_limits / passive_usage_* 等)。 func (s *adminServiceImpl) UpdateAccountExtra(ctx context.Context, id int64, updates map[string]any) error { + delete(updates, UpstreamBillingProbeEnabledExtraKey) + delete(updates, UpstreamBillingRateSyncEnabledExtraKey) + delete(updates, UpstreamBillingProbeExtraKey) delete(updates, OllamaCloudUsageSessionExtraKey) delete(updates, OllamaCloudUsageAutoRefreshExtraKey) delete(updates, OllamaCloudUsageSnapshotExtraKey) @@ -877,6 +923,7 @@ func (s *adminServiceImpl) UpdateAccountExtra(ctx context.Context, id int64, upd func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUpdateAccountsInput) (*BulkUpdateAccountsResult, error) { // Managed probe/session state may only enter through dedicated typed endpoints. delete(input.Extra, UpstreamBillingProbeEnabledExtraKey) + delete(input.Extra, UpstreamBillingRateSyncEnabledExtraKey) delete(input.Extra, UpstreamBillingProbeExtraKey) delete(input.Extra, OllamaCloudUsageSessionExtraKey) delete(input.Extra, OllamaCloudUsageAutoRefreshExtraKey) @@ -910,7 +957,7 @@ func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUp // 预取所有目标账号,供凭据守卫/代理守卫/混合渠道检查共用,避免多次 DB 查询。 var cachedTargets []*Account - if len(input.Credentials) > 0 || input.ProxyID != nil || needMixedChannelCheck || hasLongContextBillingUpdate || input.ProbeEnabled != nil { + if len(input.Credentials) > 0 || input.ProxyID != nil || needMixedChannelCheck || hasLongContextBillingUpdate || input.ProbeEnabled != nil || input.RateMultiplier != nil { loaded, err := s.accountRepo.GetByIDs(ctx, input.AccountIDs) if err != nil { return nil, err @@ -996,6 +1043,21 @@ func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUp if *input.RateMultiplier < 0 { return nil, errors.New("rate_multiplier must be >= 0") } + syncEnabledCount := 0 + for _, account := range cachedTargets { + if account == nil || account.Extra == nil { + continue + } + enabled, _ := account.Extra[UpstreamBillingRateSyncEnabledExtraKey].(bool) + if enabled { + syncEnabledCount++ + } + } + if syncEnabledCount > 0 { + return nil, ErrUpstreamBillingRateSyncBulkConflict.WithMetadata(map[string]string{ + "count": strconv.Itoa(syncEnabledCount), + }) + } } // 校验并规范化请求头覆写配置(批量路径为 JSONB 顶层 key 合并,直接校验增量即可) @@ -1014,6 +1076,9 @@ func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUp repoUpdates.Extra = make(map[string]any) } repoUpdates.Extra[UpstreamBillingProbeEnabledExtraKey] = *input.ProbeEnabled + if !*input.ProbeEnabled { + repoUpdates.Extra[UpstreamBillingRateSyncEnabledExtraKey] = false + } } if updatesUpstreamBillingProbeIdentity(input.Credentials) || input.ProxyID != nil { if repoUpdates.Extra == nil { 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 88db653861..6fcee5327c 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,95 @@ func (r *upstreamBillingProbeAdminRepo) ListShadowsByParent(context.Context, int return nil, nil } +type accountBillingSettingsAdminRepo struct { + *upstreamBillingProbeAccountRepo + concurrentRate *float64 + lastExplicitRate *float64 + updateCalls int +} + +func (r *accountBillingSettingsAdminRepo) UpdateWithAccountBillingSettings( + _ context.Context, + account *Account, + probeEnabled *bool, + rateSyncEnabled *bool, + rateMultiplier *float64, +) error { + r.mu.Lock() + defer r.mu.Unlock() + + current := r.accounts[account.ID] + if current == nil { + return ErrAccountNotFound + } + updated := *account + updated.Credentials = mergeMap(nil, account.Credentials) + updated.Extra = mergeMap(nil, account.Extra) + if updated.Extra == nil { + updated.Extra = make(map[string]any) + } + if probeEnabled != nil { + updated.Extra[UpstreamBillingProbeEnabledExtraKey] = *probeEnabled + } + if rateSyncEnabled != nil { + updated.Extra[UpstreamBillingRateSyncEnabledExtraKey] = *rateSyncEnabled + } + switch { + case rateMultiplier != nil: + value := *rateMultiplier + updated.RateMultiplier = &value + r.lastExplicitRate = &value + case r.concurrentRate != nil: + value := *r.concurrentRate + updated.RateMultiplier = &value + r.lastExplicitRate = nil + default: + updated.RateMultiplier = cloneAccountValuePointer(current.RateMultiplier) + r.lastExplicitRate = nil + } + r.accounts[account.ID] = &updated + r.updateCalls++ + return nil +} + +func TestUpdateAccountRoutesRateIntentThroughAtomicBillingUpdater(t *testing.T) { + accountID := int64(109) + initialRate := 0.1 + concurrentRate := 0.2 + repo := &accountBillingSettingsAdminRepo{ + upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{ + accountID: { + ID: accountID, + Name: "before", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + RateMultiplier: &initialRate, + Extra: map[string]any{ + UpstreamBillingProbeEnabledExtraKey: true, + UpstreamBillingRateSyncEnabledExtraKey: true, + }, + }, + }}, + concurrentRate: &concurrentRate, + } + svc := &adminServiceImpl{accountRepo: repo} + + updated, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{Name: "after"}) + require.NoError(t, err) + require.Equal(t, 1, repo.updateCalls) + require.Nil(t, repo.lastExplicitRate) + require.Equal(t, concurrentRate, *updated.RateMultiplier) + + zero := 0.0 + updated, err = svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{RateMultiplier: &zero}) + require.NoError(t, err) + require.Equal(t, 2, repo.updateCalls) + require.NotNil(t, repo.lastExplicitRate) + require.Zero(t, *repo.lastExplicitRate) + require.Zero(t, *updated.RateMultiplier) +} + func TestCreateAccountDropsManagedUpstreamBillingProbeState(t *testing.T) { repo := &upstreamBillingProbeAccountRepo{} svc := &adminServiceImpl{accountRepo: repo} @@ -28,13 +117,15 @@ func TestCreateAccountDropsManagedUpstreamBillingProbeState(t *testing.T) { Credentials: map[string]any{"api_key": "sk-test"}, SkipDefaultGroupBind: true, Extra: map[string]any{ - UpstreamBillingProbeEnabledExtraKey: true, - UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"}, + UpstreamBillingProbeEnabledExtraKey: true, + UpstreamBillingRateSyncEnabledExtraKey: true, + UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"}, }, }) require.NoError(t, err) require.NotContains(t, created.Extra, UpstreamBillingProbeEnabledExtraKey) + require.NotContains(t, created.Extra, UpstreamBillingRateSyncEnabledExtraKey) require.NotContains(t, created.Extra, UpstreamBillingProbeExtraKey) } @@ -73,8 +164,9 @@ func TestUpdateAccountPreservesManagedUpstreamBillingProbeStateForUnrelatedEdit( Type: AccountTypeAPIKey, Status: StatusActive, Extra: map[string]any{ - UpstreamBillingProbeEnabledExtraKey: true, - UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"}, + UpstreamBillingProbeEnabledExtraKey: true, + UpstreamBillingRateSyncEnabledExtraKey: true, + UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"}, }, }, }} @@ -86,6 +178,7 @@ func TestUpdateAccountPreservesManagedUpstreamBillingProbeStateForUnrelatedEdit( require.NoError(t, err) require.Equal(t, true, updated.Extra[UpstreamBillingProbeEnabledExtraKey]) + require.Equal(t, true, updated.Extra[UpstreamBillingRateSyncEnabledExtraKey]) require.Contains(t, updated.Extra, UpstreamBillingProbeExtraKey) require.Equal(t, "value", updated.Extra["custom"]) } @@ -196,8 +289,9 @@ func TestUpdateAccountInvalidatesProbeSnapshotWhenUpstreamIdentityChanges(t *tes "base_url": "https://old.example", }, Extra: map[string]any{ - UpstreamBillingProbeEnabledExtraKey: true, - UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"}, + UpstreamBillingProbeEnabledExtraKey: true, + UpstreamBillingRateSyncEnabledExtraKey: true, + UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"}, }, }, }} @@ -210,6 +304,7 @@ func TestUpdateAccountInvalidatesProbeSnapshotWhenUpstreamIdentityChanges(t *tes require.Equal(t, true, updated.Extra[UpstreamBillingProbeEnabledExtraKey]) } else { require.NotContains(t, updated.Extra, UpstreamBillingProbeEnabledExtraKey) + require.NotContains(t, updated.Extra, UpstreamBillingRateSyncEnabledExtraKey) } }) } @@ -289,16 +384,70 @@ func TestUpdateAccountAcceptsProbeEnabledAndRejectsInjectedSnapshot(t *testing.T svc := &adminServiceImpl{accountRepo: repo} updated, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{ Extra: map[string]any{ - UpstreamBillingProbeEnabledExtraKey: true, - UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"}, + UpstreamBillingProbeEnabledExtraKey: true, + UpstreamBillingRateSyncEnabledExtraKey: true, + UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"}, }, }) require.NoError(t, err) require.Equal(t, true, updated.Extra[UpstreamBillingProbeEnabledExtraKey]) + require.NotContains(t, updated.Extra, UpstreamBillingRateSyncEnabledExtraKey) require.NotContains(t, updated.Extra, UpstreamBillingProbeExtraKey) } +func TestUpdateAccountRateSyncControlsProbeAndManualMode(t *testing.T) { + accountID := int64(151) + repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{ + accountID: { + ID: accountID, + Platform: PlatformGemini, + Type: AccountTypeAPIKey, + Status: StatusActive, + Extra: map[string]any{}, + }, + }} + svc := &adminServiceImpl{accountRepo: repo} + + syncEnabled := true + updated, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{ + RateSyncEnabled: &syncEnabled, + }) + require.NoError(t, err) + require.Equal(t, true, updated.Extra[UpstreamBillingProbeEnabledExtraKey]) + require.Equal(t, true, updated.Extra[UpstreamBillingRateSyncEnabledExtraKey]) + + syncEnabled = false + updated, err = svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{ + RateSyncEnabled: &syncEnabled, + }) + require.NoError(t, err) + require.Equal(t, true, updated.Extra[UpstreamBillingProbeEnabledExtraKey]) + require.Equal(t, false, updated.Extra[UpstreamBillingRateSyncEnabledExtraKey]) +} + +func TestUpdateAccountRejectsSyncWithExplicitlyDisabledProbe(t *testing.T) { + accountID := int64(152) + repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{ + accountID: { + ID: accountID, + Platform: PlatformAnthropic, + Type: AccountTypeAPIKey, + Status: StatusActive, + }, + }} + probeEnabled := false + syncEnabled := true + + _, err := (&adminServiceImpl{accountRepo: repo}).UpdateAccount(context.Background(), accountID, &UpdateAccountInput{ + ProbeEnabled: &probeEnabled, + RateSyncEnabled: &syncEnabled, + }) + + require.Error(t, err) + require.Empty(t, repo.updates[accountID]) +} + func TestUpdateAccountExplicitProbeDisableUsesDedicatedExtraUpdate(t *testing.T) { accountID := int64(113) repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{ @@ -321,6 +470,7 @@ func TestUpdateAccountExplicitProbeDisableUsesDedicatedExtraUpdate(t *testing.T) require.NoError(t, err) require.Len(t, repo.updates[accountID], 1) require.Equal(t, false, repo.updates[accountID][0][UpstreamBillingProbeEnabledExtraKey]) + require.Equal(t, false, repo.updates[accountID][0][UpstreamBillingRateSyncEnabledExtraKey]) } func TestUpdateAccountExplicitUnchangedProbeEnabledStillUsesDedicatedExtraUpdate(t *testing.T) { @@ -364,15 +514,36 @@ func TestUpdateAccountRejectsInvalidProbeEnabled(t *testing.T) { require.Error(t, err) } +func TestUpdateAccountExtraDropsManagedBillingProbeFields(t *testing.T) { + accountID := int64(153) + repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{ + accountID: {ID: accountID, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, + }} + + err := (&adminServiceImpl{accountRepo: repo}).UpdateAccountExtra(context.Background(), accountID, map[string]any{ + "custom": "value", + UpstreamBillingProbeEnabledExtraKey: true, + UpstreamBillingRateSyncEnabledExtraKey: true, + UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"}, + }) + + require.NoError(t, err) + require.Equal(t, "value", repo.accounts[accountID].Extra["custom"]) + require.NotContains(t, repo.accounts[accountID].Extra, UpstreamBillingProbeEnabledExtraKey) + require.NotContains(t, repo.accounts[accountID].Extra, UpstreamBillingRateSyncEnabledExtraKey) + require.NotContains(t, repo.accounts[accountID].Extra, UpstreamBillingProbeExtraKey) +} + func TestBulkUpdateAccountsDropsManagedUpstreamBillingProbeState(t *testing.T) { repo := &upstreamBillingProbeAccountRepo{} svc := &adminServiceImpl{accountRepo: repo} input := &BulkUpdateAccountsInput{ AccountIDs: []int64{1}, Extra: map[string]any{ - "custom": "value", - UpstreamBillingProbeEnabledExtraKey: true, - UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"}, + "custom": "value", + UpstreamBillingProbeEnabledExtraKey: true, + UpstreamBillingRateSyncEnabledExtraKey: true, + UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"}, }, } @@ -383,6 +554,7 @@ func TestBulkUpdateAccountsDropsManagedUpstreamBillingProbeState(t *testing.T) { require.Len(t, repo.bulkUpdates, 1) require.Equal(t, "value", repo.bulkUpdates[0].Extra["custom"]) require.NotContains(t, repo.bulkUpdates[0].Extra, UpstreamBillingProbeEnabledExtraKey) + require.NotContains(t, repo.bulkUpdates[0].Extra, UpstreamBillingRateSyncEnabledExtraKey) require.NotContains(t, repo.bulkUpdates[0].Extra, UpstreamBillingProbeExtraKey) } @@ -403,6 +575,9 @@ func TestBulkUpdateAccountsAcceptsDedicatedUpstreamBillingProbeSetting(t *testin require.Equal(t, 2, result.Success) require.Len(t, repo.bulkUpdates, 1) require.Equal(t, enabled, repo.bulkUpdates[0].Extra[UpstreamBillingProbeEnabledExtraKey]) + if !enabled { + require.Equal(t, false, repo.bulkUpdates[0].Extra[UpstreamBillingRateSyncEnabledExtraKey]) + } require.NotNil(t, repo.bulkUpdates[0].ProbeEnabled) require.Equal(t, enabled, *repo.bulkUpdates[0].ProbeEnabled) }) diff --git a/backend/internal/service/admin_service.go b/backend/internal/service/admin_service.go index bd907125a5..7e69ef258d 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 63bf26fbd1..206efb21f1 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 7b19dbf6f5..7bd31778d7 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 344ed19d14..effeca4677 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 e1637accdb..e5a72b4c44 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 7f105445e2..35630da9ad 100644 --- a/backend/internal/service/upstream_billing_probe.go +++ b/backend/internal/service/upstream_billing_probe.go @@ -28,8 +28,9 @@ import ( const ( // These values live in accounts.extra so PR2 does not require a schema migration. - UpstreamBillingProbeExtraKey = "upstream_billing_probe" - UpstreamBillingProbeEnabledExtraKey = "upstream_billing_probe_enabled" + UpstreamBillingProbeExtraKey = "upstream_billing_probe" + UpstreamBillingProbeEnabledExtraKey = "upstream_billing_probe_enabled" + UpstreamBillingRateSyncEnabledExtraKey = "upstream_billing_rate_sync_enabled" upstreamBillingProbeDefaultIntervalMinutes = 30 upstreamBillingProbeMinIntervalMinutes = 5 @@ -44,6 +45,8 @@ const ( // /v1/sub2api/billing,按常规 interval 重排只会持续占满每周期 // upstreamBillingProbeMaxPerCycle 个名额。 upstreamBillingProbeUnsupportedDelayFactor = 8 + upstreamBillingProbeAccountRateScale = 10000.0 + upstreamBillingProbeAccountRateMax = 999999.9999 upstreamBillingProbeLeaderLockKey = "upstream:billing:probe:leader" upstreamBillingProbeLeaderLockTTL = 2 * time.Minute ) @@ -61,6 +64,10 @@ var ( ErrUpstreamBillingProbeIdentityChanged = infraerrors.Conflict( "UPSTREAM_BILLING_PROBE_IDENTITY_CHANGED", "account identity changed during upstream billing probe; retry the probe", ) + ErrUpstreamBillingRateSyncBulkConflict = infraerrors.Conflict( + "UPSTREAM_BILLING_RATE_SYNC_BULK_CONFLICT", + "account rate multiplier cannot be changed in bulk while upstream billing rate sync is enabled", + ) ) const ( @@ -197,7 +204,7 @@ type UpstreamBillingProbeService struct { } type upstreamBillingProbeSnapshotWriter interface { - UpdateUpstreamBillingProbeSnapshot(context.Context, *Account, *UpstreamBillingProbeSnapshot) error + UpdateUpstreamBillingProbeSnapshot(context.Context, *Account, *UpstreamBillingProbeSnapshot, *float64) error } type upstreamBillingProbeDueAccountLister interface { @@ -546,9 +553,11 @@ func (s *UpstreamBillingProbeService) SetAccountEnabled(ctx context.Context, acc if !isUpstreamBillingProbeAccount(account) { return ErrUpstreamBillingProbeAccountInvalid } - return s.accountRepo.UpdateExtra(ctx, accountID, map[string]any{ - UpstreamBillingProbeEnabledExtraKey: enabled, - }) + updates := map[string]any{UpstreamBillingProbeEnabledExtraKey: enabled} + if !enabled { + updates[UpstreamBillingRateSyncEnabledExtraKey] = false + } + return s.accountRepo.UpdateExtra(ctx, accountID, updates) } func (s *UpstreamBillingProbeService) probeLoadedAccount(ctx context.Context, account *Account, intervalMinutes int) (*UpstreamBillingProbeSnapshot, error) { @@ -635,6 +644,10 @@ func (s *UpstreamBillingProbeService) probeLoadedAccount(ctx context.Context, ac if err != nil { return s.persistProbeFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "invalid_response", retryAfter(resp.Header, now)) } + rateMultiplier, ok := upstreamBillingProbeAccountRate(data) + if !ok { + return s.persistProbeFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "invalid_response", retryAfter(resp.Header, now)) + } snapshot := &UpstreamBillingProbeSnapshot{ Status: UpstreamBillingProbeStatusOK, Data: data, @@ -644,7 +657,7 @@ func (s *UpstreamBillingProbeService) probeLoadedAccount(ctx context.Context, ac NextProbeAt: now.Add(nextProbeDelay(intervalMinutes, 0)), HTTPStatus: resp.StatusCode, } - if err := s.updateSnapshot(ctx, account, snapshot); err != nil { + if err := s.updateSnapshot(ctx, account, snapshot, &rateMultiplier); err != nil { return nil, err } return snapshot, nil @@ -686,18 +699,23 @@ func (s *UpstreamBillingProbeService) persistProbeFailure( snapshot.FreshUntil = probeTimePtr(previous.ReceivedAt.Add(2 * time.Duration(intervalMinutes) * time.Minute)) } } - if err := s.updateSnapshot(ctx, account, snapshot); err != nil { + if err := s.updateSnapshot(ctx, account, snapshot, nil); err != nil { return nil, err } return snapshot, nil } -func (s *UpstreamBillingProbeService) updateSnapshot(ctx context.Context, account *Account, snapshot *UpstreamBillingProbeSnapshot) error { +func (s *UpstreamBillingProbeService) updateSnapshot( + ctx context.Context, + account *Account, + snapshot *UpstreamBillingProbeSnapshot, + rateMultiplier *float64, +) error { writer, ok := s.accountRepo.(upstreamBillingProbeSnapshotWriter) if !ok { return ErrUpstreamBillingProbeUnavailable } - return writer.UpdateUpstreamBillingProbeSnapshot(ctx, account, snapshot) + return writer.UpdateUpstreamBillingProbeSnapshot(ctx, account, snapshot, rateMultiplier) } func parseUpstreamBillingProbeResponse(body []byte) (map[string]any, error) { @@ -799,6 +817,21 @@ func upstreamBillingRateAt(data map[string]any, now time.Time) (float64, bool) { return base, true } +// upstreamBillingProbeAccountRate converts the declared effective multiplier +// to the precision supported by accounts.rate_multiplier (DECIMAL(10,4)). +func upstreamBillingProbeAccountRate(data map[string]any) (float64, bool) { + value, ok := resolveAccountExtraNumber(data, "effective_rate_multiplier") + if !ok || value < 0 || value > upstreamBillingProbeAccountRateMax || + math.IsNaN(value) || math.IsInf(value, 0) { + return 0, false + } + rounded := math.Round(value*upstreamBillingProbeAccountRateScale) / upstreamBillingProbeAccountRateScale + if value > 0 && rounded == 0 { + return 0, false + } + return rounded, true +} + func upstreamBillingPeakMultiplierAt(data map[string]any, now time.Time) (float64, bool) { peakEnabled, ok := data["peak_rate_enabled"].(bool) if !ok { @@ -865,9 +898,7 @@ func decodeUpstreamBillingProbeSnapshot(extra map[string]any) *UpstreamBillingPr // IsUpstreamBillingProbeIdentity reports whether an account identity may opt // in to the upstream billing probe. `/v1/sub2api/billing` is a key-scoped -// sub2api convention: any API-key account whose base_url points at a -// sub2api-compatible upstream answers it regardless of the account platform, -// so eligibility is not restricted to OpenAI. +// sub2api convention shared by the five supported API-key platforms. // Non-sub2api upstreams return 404 and the snapshot records "unsupported". // Only AccountTypeAPIKey is in scope. OAuth/Bedrock hold no static API key to // present at all; AccountTypeUpstream (antigravity relay accounts) does carry @@ -876,7 +907,15 @@ func decodeUpstreamBillingProbeSnapshot(extra map[string]any) *UpstreamBillingPr // type=apikey by the admin form, so only pre-existing type=upstream rows // cannot turn the probe on. func IsUpstreamBillingProbeIdentity(platform, accountType string) bool { - return platform != "" && accountType == AccountTypeAPIKey + if accountType != AccountTypeAPIKey { + return false + } + switch platform { + case PlatformOpenAI, PlatformAnthropic, PlatformGemini, PlatformAntigravity, PlatformGrok: + return true + default: + return false + } } func isUpstreamBillingProbeAccount(account *Account) bool { @@ -934,6 +973,14 @@ func upstreamBillingProbeEnabled(account *Account) bool { return ok && enabled } +func upstreamBillingRateSyncEnabled(account *Account) bool { + if account == nil || account.Extra == nil { + return false + } + enabled, ok := account.Extra[UpstreamBillingRateSyncEnabledExtraKey].(bool) + return ok && enabled && upstreamBillingProbeEnabled(account) +} + func (s *UpstreamBillingProbeService) currentTime() time.Time { if s != nil && s.now != nil { return s.now() diff --git a/backend/internal/service/upstream_billing_probe_multiplatform_test.go b/backend/internal/service/upstream_billing_probe_multiplatform_test.go index a2b03a96fc..018f2b7198 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 edd297af96..fd1d7b45de 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,8 @@ func TestUpstreamBillingProbeSuccessPersistsSanitizedSnapshot(t *testing.T) { require.Equal(t, fixedNow.Add(time.Hour), *snapshot.FreshUntil) require.False(t, snapshot.NextProbeAt.Before(fixedNow.Add(24*time.Minute))) require.False(t, snapshot.NextProbeAt.After(fixedNow.Add(36*time.Minute))) + require.NotNil(t, account.RateMultiplier) + require.Equal(t, 0.9, *account.RateMultiplier) require.Equal(t, "https://upstream.example/v1/sub2api/billing", upstream.lastReq.URL.String()) require.Equal(t, http.MethodGet, upstream.lastReq.Method) require.Equal(t, "Bearer sk-sensitive", upstream.lastReq.Header.Get("Authorization")) @@ -317,6 +336,98 @@ func TestUpstreamBillingProbeSuccessPersistsSanitizedSnapshot(t *testing.T) { require.Equal(t, snapshot.Status, persisted.Status) } +func TestUpstreamBillingProbeSyncsEffectiveRateForAllAPIKeyPlatforms(t *testing.T) { + for _, platform := range []string{ + PlatformOpenAI, + PlatformAnthropic, + PlatformGemini, + PlatformAntigravity, + PlatformGrok, + } { + t.Run(platform, func(t *testing.T) { + initialRate := 0.25 + account := &Account{ + ID: 17, + Platform: platform, + Type: AccountTypeAPIKey, + Status: StatusActive, + Concurrency: 1, + RateMultiplier: &initialRate, + Credentials: map[string]any{ + "api_key": "sk-sensitive", + "base_url": "https://upstream.example", + }, + Extra: map[string]any{ + UpstreamBillingProbeEnabledExtraKey: true, + UpstreamBillingRateSyncEnabledExtraKey: true, + }, + } + repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{account.ID: account}} + svc := newUpstreamBillingProbeTestService(repo, &upstreamBillingProbeHTTPStub{}, &upstreamBillingProbeSettingRepo{}) + + snapshot, err := svc.ProbeAccount(context.Background(), account.ID) + + require.NoError(t, err) + require.Equal(t, UpstreamBillingProbeStatusOK, snapshot.Status) + require.NotNil(t, account.RateMultiplier) + require.Equal(t, 0.8, *account.RateMultiplier) + }) + } +} + +func TestUpstreamBillingProbeOnlyDoesNotChangeAccountRate(t *testing.T) { + initialRate := 0.25 + account := &Account{ + ID: 18, + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Status: StatusActive, + Concurrency: 1, + RateMultiplier: &initialRate, + Credentials: map[string]any{ + "api_key": "sk-sensitive", + "base_url": "https://upstream.example", + }, + Extra: map[string]any{UpstreamBillingProbeEnabledExtraKey: true}, + } + repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{account.ID: account}} + svc := newUpstreamBillingProbeTestService(repo, &upstreamBillingProbeHTTPStub{}, &upstreamBillingProbeSettingRepo{}) + + snapshot, err := svc.ProbeAccount(context.Background(), account.ID) + + require.NoError(t, err) + require.Equal(t, UpstreamBillingProbeStatusOK, snapshot.Status) + require.NotNil(t, account.RateMultiplier) + require.Equal(t, initialRate, *account.RateMultiplier) + require.Contains(t, account.Extra, UpstreamBillingProbeExtraKey) +} + +func TestUpstreamBillingProbeAccountRateMatchesDatabasePrecision(t *testing.T) { + tests := []struct { + name string + value float64 + want float64 + ok bool + }{ + {name: "zero", value: 0, want: 0, ok: true}, + {name: "round to four decimals", value: 0.07654, want: 0.0765, ok: true}, + {name: "maximum", value: upstreamBillingProbeAccountRateMax, want: upstreamBillingProbeAccountRateMax, ok: true}, + {name: "positive below database precision", value: 0.00001, ok: false}, + {name: "overflow before rounding", value: 999999.99991, ok: false}, + {name: "overflow after rounding", value: 999999.99996, ok: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, ok := upstreamBillingProbeAccountRate(map[string]any{"effective_rate_multiplier": tt.value}) + require.Equal(t, tt.ok, ok) + if tt.ok { + require.Equal(t, tt.want, got) + } + }) + } +} + func TestUpstreamBillingProbeRejectsMissingRequiredMultiplier(t *testing.T) { _, err := parseUpstreamBillingProbeResponse([]byte(`{ "object":"sub2api.key_billing", @@ -456,6 +567,7 @@ func TestUpstreamBillingRateAtHandlesDST(t *testing.T) { func TestUpstreamBillingProbeFailurePreservesLastSuccessAndRetryAfter(t *testing.T) { receivedAt := time.Date(2026, time.July, 12, 12, 0, 0, 0, time.UTC) + initialRate := 0.35 previous := &UpstreamBillingProbeSnapshot{ Status: UpstreamBillingProbeStatusOK, Data: map[string]any{"effective_rate_multiplier": 0.5}, @@ -463,13 +575,17 @@ func TestUpstreamBillingProbeFailurePreservesLastSuccessAndRetryAfter(t *testing FailureCount: 1, } account := &Account{ - ID: 18, - Platform: PlatformOpenAI, - Type: AccountTypeAPIKey, - Status: StatusActive, - Concurrency: 1, - Credentials: map[string]any{"api_key": "sk-test", "base_url": "https://upstream.example"}, - Extra: map[string]any{UpstreamBillingProbeExtraKey: previous}, + ID: 18, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Concurrency: 1, + RateMultiplier: &initialRate, + Credentials: map[string]any{"api_key": "sk-test", "base_url": "https://upstream.example"}, + Extra: map[string]any{ + UpstreamBillingProbeEnabledExtraKey: true, + UpstreamBillingProbeExtraKey: previous, + }, } repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{account.ID: account}} upstream := &httpUpstreamRecorder{resp: &http.Response{ @@ -492,6 +608,8 @@ func TestUpstreamBillingProbeFailurePreservesLastSuccessAndRetryAfter(t *testing require.Equal(t, "http_error", snapshot.LastError) require.False(t, snapshot.NextProbeAt.Before(fixedNow.Add(4*time.Hour))) require.NotContains(t, snapshot.LastError, "do not persist") + require.NotNil(t, account.RateMultiplier) + require.Equal(t, initialRate, *account.RateMultiplier) } func TestUpstreamBillingProbeRetryAfterIsNotShortened(t *testing.T) { @@ -611,6 +729,7 @@ func TestUpstreamBillingProbeUnsupportedAndAccountToggle(t *testing.T) { require.NoError(t, svc.SetAccountEnabled(context.Background(), account.ID, true)) require.Equal(t, true, account.Extra[UpstreamBillingProbeEnabledExtraKey]) + account.Extra[UpstreamBillingRateSyncEnabledExtraKey] = true snapshot, err := svc.ProbeAccount(context.Background(), account.ID) require.NoError(t, err) require.Equal(t, UpstreamBillingProbeStatusUnsupported, snapshot.Status) @@ -624,6 +743,9 @@ func TestUpstreamBillingProbeUnsupportedAndAccountToggle(t *testing.T) { require.Equal(t, 2, snapshot.FailureCount) require.False(t, snapshot.NextProbeAt.Before(fixedNow.Add(192*time.Minute))) require.False(t, snapshot.NextProbeAt.After(fixedNow.Add(288*time.Minute))) + require.NoError(t, svc.SetAccountEnabled(context.Background(), account.ID, false)) + require.Equal(t, false, account.Extra[UpstreamBillingProbeEnabledExtraKey]) + require.Equal(t, false, account.Extra[UpstreamBillingRateSyncEnabledExtraKey]) invalid := &Account{ID: 20, Platform: PlatformOpenAI, Type: AccountTypeOAuth} repo.accounts[invalid.ID] = invalid @@ -662,10 +784,14 @@ func TestUpstreamBillingProbeRunnerIsBoundedAndManualProbeIgnoresSwitches(t *tes require.Equal(t, int64(20), upstream.calls.Load()) accounts[25].Extra[UpstreamBillingProbeEnabledExtraKey] = false + manualRate := 0.25 + accounts[25].RateMultiplier = &manualRate snapshot, err := svc.ProbeAccount(context.Background(), 25) require.NoError(t, err) require.Equal(t, UpstreamBillingProbeStatusOK, snapshot.Status) require.Equal(t, int64(21), upstream.calls.Load()) + require.NotNil(t, accounts[25].RateMultiplier) + require.Equal(t, manualRate, *accounts[25].RateMultiplier) } func TestUpstreamBillingProbeRunnerRechecksEnabledAfterDueSelection(t *testing.T) { diff --git a/frontend/src/components/account/BulkEditAccountModal.vue b/frontend/src/components/account/BulkEditAccountModal.vue index 02ddaca92e..ad9998695d 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') }} +
+