From d2585097f38f8e495a0060f9b8ebf395e2139603 Mon Sep 17 00:00:00 2001 From: Tian Lee <498756723@qq.com> Date: Fri, 17 Jul 2026 14:15:05 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E5=AE=8C=E5=96=84=E4=B8=8A=E6=B8=B8=20S?= =?UTF-8?q?ub2API=20=E8=AE=A1=E8=B4=B9=E5=80=8D=E7=8E=87=E6=8E=A2=E6=B5=8B?= =?UTF-8?q?=E4=B8=8E=E8=B4=A6=E5=8F=B7=E5=B1=95=E7=A4=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../internal/handler/admin/account_handler.go | 7 +- .../account_handler_mixed_channel_test.go | 19 +++ .../handler/admin/admin_service_stub_test.go | 2 + backend/internal/repository/account_repo.go | 76 ++++++++- .../account_repo_sort_integration_test.go | 105 +++++++++++++ ...repo_upstream_billing_probe_update_test.go | 40 +++++ backend/internal/service/account_service.go | 1 + backend/internal/service/admin_account.go | 41 ++++- ...min_account_upstream_billing_probe_test.go | 83 ++++++++++ backend/internal/service/admin_service.go | 2 + .../service/upstream_billing_probe.go | 25 ++- .../service/upstream_billing_probe_test.go | 22 ++- .../account/BulkEditAccountModal.vue | 52 +++++++ .../components/account/CreateAccountModal.vue | 35 ++++- .../account/UpstreamBillingRateCell.vue | 111 +++++++++---- .../__tests__/BulkEditAccountModal.spec.ts | 67 ++++++++ .../__tests__/CreateAccountModal.spec.ts | 49 +++++- .../__tests__/UpstreamBillingRateCell.spec.ts | 146 ++++++++++++++++-- frontend/src/components/common/DataTable.vue | 60 +++---- .../common/__tests__/DataTable.spec.ts | 4 + .../src/i18n/locales/en/admin/accounts.ts | 20 ++- .../src/i18n/locales/en/admin/settings.ts | 10 ++ .../src/i18n/locales/zh/admin/accounts.ts | 20 ++- .../src/i18n/locales/zh/admin/settings.ts | 10 ++ frontend/src/types/index.ts | 1 + frontend/src/views/admin/AccountsView.vue | 141 ++++++++--------- frontend/src/views/admin/SettingsView.vue | 129 ++++++++++++++++ .../__tests__/AccountsView.bulkEdit.spec.ts | 146 +++++++++++++++++- .../AccountsView.usageWindowsHint.spec.ts | 2 + .../admin/__tests__/SettingsView.spec.ts | 68 ++++++++ 30 files changed, 1311 insertions(+), 183 deletions(-) diff --git a/backend/internal/handler/admin/account_handler.go b/backend/internal/handler/admin/account_handler.go index d9b742003f..98adcf5142 100644 --- a/backend/internal/handler/admin/account_handler.go +++ b/backend/internal/handler/admin/account_handler.go @@ -121,6 +121,7 @@ type CreateAccountRequest 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"` ConfirmMixedChannelRisk *bool `json:"confirm_mixed_channel_risk"` // 用户确认混合渠道风险 } @@ -159,6 +160,7 @@ type BulkUpdateAccountsRequest struct { GroupIDs *[]int64 `json:"group_ids"` Credentials map[string]any `json:"credentials"` Extra map[string]any `json:"extra"` + ProbeEnabled *bool `json:"upstream_billing_probe_enabled"` ConfirmMixedChannelRisk *bool `json:"confirm_mixed_channel_risk"` // 用户确认混合渠道风险 } @@ -828,6 +830,7 @@ func (h *AccountHandler) Create(c *gin.Context) { GroupIDs: req.GroupIDs, ExpiresAt: req.ExpiresAt, AutoPauseOnExpired: req.AutoPauseOnExpired, + ProbeEnabled: req.ProbeEnabled, SkipMixedChannelCheck: skipCheck, }) if execErr != nil { @@ -1907,7 +1910,8 @@ func (h *AccountHandler) BulkUpdate(c *gin.Context) { req.Schedulable != nil || req.GroupIDs != nil || len(req.Credentials) > 0 || - len(req.Extra) > 0 + len(req.Extra) > 0 || + req.ProbeEnabled != nil if !hasUpdates { response.BadRequest(c, "No updates provided") @@ -1928,6 +1932,7 @@ func (h *AccountHandler) BulkUpdate(c *gin.Context) { GroupIDs: req.GroupIDs, Credentials: req.Credentials, Extra: req.Extra, + ProbeEnabled: req.ProbeEnabled, 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 a1dc25124e..bc4a3cf49b 100644 --- a/backend/internal/handler/admin/account_handler_mixed_channel_test.go +++ b/backend/internal/handler/admin/account_handler_mixed_channel_test.go @@ -222,3 +222,22 @@ func TestBulkUpdateAcceptsFilterTargetRequest(t *testing.T) { require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp)) require.Equal(t, float64(0), resp["code"]) } + +func TestBulkUpdateAcceptsDedicatedUpstreamBillingProbeSetting(t *testing.T) { + adminSvc := newStubAdminService() + router := setupAccountMixedChannelRouter(adminSvc) + + body, _ := json.Marshal(map[string]any{ + "account_ids": []int64{1, 2}, + "upstream_billing_probe_enabled": false, + }) + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/bulk-update", 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.lastBulkUpdateAccountInput) + require.NotNil(t, adminSvc.lastBulkUpdateAccountInput.ProbeEnabled) + require.False(t, *adminSvc.lastBulkUpdateAccountInput.ProbeEnabled) +} diff --git a/backend/internal/handler/admin/admin_service_stub_test.go b/backend/internal/handler/admin/admin_service_stub_test.go index f5c238a61e..035e171c9b 100644 --- a/backend/internal/handler/admin/admin_service_stub_test.go +++ b/backend/internal/handler/admin/admin_service_stub_test.go @@ -33,6 +33,7 @@ type stubAdminService struct { createSparkShadowErr error updateAccountErr error bulkUpdateAccountErr error + lastBulkUpdateAccountInput *service.BulkUpdateAccountsInput getAccountResult *service.Account updateAccountCalls int updateAccountExtraCalls int @@ -478,6 +479,7 @@ func (s *stubAdminService) SetAccountSchedulable(ctx context.Context, id int64, } func (s *stubAdminService) BulkUpdateAccounts(ctx context.Context, input *service.BulkUpdateAccountsInput) (*service.BulkUpdateAccountsResult, error) { + s.lastBulkUpdateAccountInput = input if s.bulkUpdateAccountErr != nil { return nil, s.bulkUpdateAccountErr } diff --git a/backend/internal/repository/account_repo.go b/backend/internal/repository/account_repo.go index f014155bb2..52c19e143f 100644 --- a/backend/internal/repository/account_repo.go +++ b/backend/internal/repository/account_repo.go @@ -893,6 +893,20 @@ func (r *accountRepository) ListOpsAccountsForStats(ctx context.Context, platfor func accountListOrder(params pagination.PaginationParams) []func(*entsql.Selector) { sortBy := strings.ToLower(strings.TrimSpace(params.SortBy)) sortOrder := params.NormalizedSortOrder(pagination.SortOrderAsc) + if sortBy == "upstream_billing_rate" { + direction := "ASC" + tieOrder := entsql.Asc + if sortOrder == pagination.SortOrderDesc { + direction = "DESC" + tieOrder = entsql.Desc + } + return []func(*entsql.Selector){func(s *entsql.Selector) { + extra := s.C(dbaccount.FieldExtra) + expression := upstreamBillingRateSortExpression(extra) + s.OrderExpr(entsql.Expr(expression + " " + direction + " NULLS LAST")) + s.OrderBy(tieOrder(s.C(dbaccount.FieldID))) + }} + } field := dbaccount.FieldName defaultOrder := true @@ -934,6 +948,40 @@ func accountListOrder(params pagination.PaginationParams) []func(*entsql.Selecto return []func(*entsql.Selector){dbent.Asc(field), dbent.Asc(dbaccount.FieldID)} } +func upstreamBillingRateSortExpression(extra string) string { + status := extra + " #>> '{upstream_billing_probe,status}'" + effectiveJSON := extra + " #> '{upstream_billing_probe,data,effective_rate_multiplier}'" + effective := extra + " #>> '{upstream_billing_probe,data,effective_rate_multiplier}'" + resolvedJSON := extra + " #> '{upstream_billing_probe,data,resolved_rate_multiplier}'" + resolved := extra + " #>> '{upstream_billing_probe,data,resolved_rate_multiplier}'" + peakEnabledJSON := extra + " #> '{upstream_billing_probe,data,peak_rate_enabled}'" + peakEnabled := extra + " #>> '{upstream_billing_probe,data,peak_rate_enabled}'" + peakStart := extra + " #>> '{upstream_billing_probe,data,peak_start}'" + peakEnd := extra + " #>> '{upstream_billing_probe,data,peak_end}'" + peakMultiplierJSON := extra + " #> '{upstream_billing_probe,data,peak_rate_multiplier}'" + peakMultiplier := extra + " #>> '{upstream_billing_probe,data,peak_rate_multiplier}'" + peakMultiplierValue := "(CASE WHEN jsonb_typeof(" + peakMultiplierJSON + ") = 'number' THEN (" + peakMultiplier + ")::numeric END)" + billingScope := extra + " #>> '{upstream_billing_probe,data,billing_scope}'" + timezone := extra + " #>> '{upstream_billing_probe,data,timezone}'" + validClock := "'^([01][0-9]|2[0-3]):[0-5][0-9]$'" + startMinute := "(CASE WHEN " + peakStart + " ~ " + validClock + " THEN split_part(" + peakStart + ", ':', 1)::numeric * 60 + split_part(" + peakStart + ", ':', 2)::numeric END)" + endMinute := "(CASE WHEN " + peakEnd + " ~ " + validClock + " THEN split_part(" + peakEnd + ", ':', 1)::numeric * 60 + split_part(" + peakEnd + ", ':', 2)::numeric END)" + localMinute := "(EXTRACT(HOUR FROM (CURRENT_TIMESTAMP AT TIME ZONE (" + timezone + "))) * 60 + EXTRACT(MINUTE FROM (CURRENT_TIMESTAMP AT TIME ZONE (" + timezone + "))))" + validPeakWindow := peakStart + " ~ " + validClock + " AND " + + peakEnd + " ~ " + validClock + " AND " + + startMinute + " < " + endMinute + validPeakConfig := validPeakWindow + " AND " + peakMultiplierValue + " >= 0 AND " + + "EXISTS (SELECT 1 FROM pg_timezone_names WHERE name = " + timezone + ")" + dynamicRate := "CASE WHEN " + peakEnabled + " = 'false' THEN (" + resolved + ")::numeric WHEN " + peakEnabled + " = 'true' AND " + validPeakConfig + + " THEN (" + resolved + ")::numeric * CASE WHEN " + localMinute + " >= " + startMinute + " AND " + localMinute + " < " + endMinute + + " THEN " + peakMultiplierValue + " ELSE 1 END ELSE NULL END" + legacySnapshot := "jsonb_typeof(" + resolvedJSON + ") IS NULL AND jsonb_typeof(" + peakEnabledJSON + ") IS NULL" + + return "CASE WHEN " + status + " IN ('ok', 'failed') AND (jsonb_typeof(" + resolvedJSON + ") = 'number' OR jsonb_typeof(" + effectiveJSON + ") = 'number') THEN CASE WHEN jsonb_typeof(" + + resolvedJSON + ") = 'number' AND jsonb_typeof(" + peakEnabledJSON + ") = 'boolean' THEN CASE WHEN " + billingScope + " = 'token' THEN " + dynamicRate + " ELSE NULL END WHEN " + legacySnapshot + + " AND jsonb_typeof(" + effectiveJSON + ") = 'number' THEN (" + effective + ")::numeric END END" +} + func (r *accountRepository) ListByGroup(ctx context.Context, groupID int64) ([]service.Account, error) { accounts, err := r.queryAccountsByGroup(ctx, groupID, accountGroupQueryOptions{ status: service.StatusActive, @@ -2575,6 +2623,12 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates args = append(args, *updates.Schedulable) idx++ } + if updates.ProbeEnabled != nil { + if updates.Extra == nil { + updates.Extra = make(map[string]any) + } + updates.Extra[service.UpstreamBillingProbeEnabledExtraKey] = *updates.ProbeEnabled + } // JSONB 需要合并而非覆盖，使用 raw SQL 保持旧行为。 if len(updates.Credentials) > 0 { payload, err := json.Marshal(updates.Credentials) @@ -2605,8 +2659,14 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates setClauses = append(setClauses, "updated_at = NOW()") - query := "UPDATE accounts SET " + joinClauses(setClauses, ", ") + " WHERE id = ANY($" + itoa(idx) + ") AND deleted_at IS NULL" + whereClause := " WHERE id = ANY($" + itoa(idx) + ") AND deleted_at IS NULL" args = append(args, pq.Array(ids)) + idx++ + if updates.ProbeEnabled != nil { + whereClause += " AND platform = $" + itoa(idx) + " AND type = $" + itoa(idx+1) + args = append(args, service.PlatformOpenAI, service.AccountTypeAPIKey) + } + query := "UPDATE accounts SET " + joinClauses(setClauses, ", ") + whereClause baseCtx := ctx contextTx := dbent.TxFromContext(ctx) @@ -2635,6 +2695,20 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates if err != nil { return 0, err } + if updates.ProbeEnabled != nil { + expectedRows := int64(0) + seenIDs := make(map[int64]struct{}, len(ids)) + for _, id := range ids { + if _, seen := seenIDs[id]; seen { + continue + } + seenIDs[id] = struct{}{} + expectedRows++ + } + if rows != expectedRows { + return 0, service.ErrUpstreamBillingProbeAccountInvalid + } + } if rows > 0 { payload := map[string]any{"account_ids": ids} if err := enqueueSchedulerOutbox(ctx, exec, service.SchedulerOutboxEventAccountBulkChanged, nil, nil, payload); err != nil { diff --git a/backend/internal/repository/account_repo_sort_integration_test.go b/backend/internal/repository/account_repo_sort_integration_test.go index 098dde7bbd..19c974f37f 100644 --- a/backend/internal/repository/account_repo_sort_integration_test.go +++ b/backend/internal/repository/account_repo_sort_integration_test.go @@ -3,6 +3,9 @@ package repository import ( + "fmt" + "time" + "github.com/Wei-Shaw/sub2api/internal/pkg/pagination" "github.com/Wei-Shaw/sub2api/internal/service" ) @@ -33,3 +36,105 @@ func (s *AccountRepoSuite) TestListWithFilters_SortByPriorityDesc() { s.Require().Equal("high-priority", accounts[0].Name) s.Require().Equal("low-priority", accounts[1].Name) } + +func (s *AccountRepoSuite) TestListWithFilters_SortByUpstreamBillingRateWithMissingLast() { + makeAccount := func(name, status string, rate any) { + extra := map[string]any{} + if rate != nil { + extra[service.UpstreamBillingProbeExtraKey] = map[string]any{ + "status": status, + "data": map[string]any{"effective_rate_multiplier": rate}, + } + } + mustCreateAccount(s.T(), s.client, &service.Account{Name: name, Extra: extra}) + } + makeAccount("high-rate", service.UpstreamBillingProbeStatusOK, 0.8) + makeAccount("low-rate", service.UpstreamBillingProbeStatusOK, 0.03) + makeAccount("missing-rate", "", nil) + makeAccount("unsupported-with-retained-rate", service.UpstreamBillingProbeStatusUnsupported, 0.01) + + for _, tc := range []struct { + order string + want []string + }{ + {order: "asc", want: []string{"low-rate", "high-rate", "missing-rate", "unsupported-with-retained-rate"}}, + {order: "desc", want: []string{"high-rate", "low-rate", "unsupported-with-retained-rate", "missing-rate"}}, + } { + accounts, _, err := s.repo.ListWithFilters(s.ctx, pagination.PaginationParams{ + Page: 1, PageSize: 10, SortBy: "upstream_billing_rate", SortOrder: tc.order, + }, "", "", "", "", 0, "") + s.Require().NoError(err) + s.Require().Len(accounts, 4) + for i, name := range tc.want { + s.Require().Equal(name, accounts[i].Name) + } + } +} + +func (s *AccountRepoSuite) TestListWithFilters_SortByCurrentUpstreamBillingRateDuringPeak() { + now := time.Now() + locations := []string{"UTC", "Asia/Shanghai", "America/New_York", "Europe/London"} + var timezone string + var minute int + for _, name := range locations { + location, err := time.LoadLocation(name) + s.Require().NoError(err) + local := now.In(location) + candidate := local.Hour()*60 + local.Minute() + if candidate >= 2 && candidate <= 1436 { + timezone = name + minute = candidate + break + } + } + s.Require().NotEmpty(timezone) + + peakStart := fmt.Sprintf("%02d:%02d", (minute-2)/60, (minute-2)%60) + peakEnd := fmt.Sprintf("%02d:%02d", (minute+3)/60, (minute+3)%60) + mustCreateAccount(s.T(), s.client, &service.Account{ + Name: "current-peak-rate", + Extra: map[string]any{ + service.UpstreamBillingProbeExtraKey: map[string]any{ + "status": service.UpstreamBillingProbeStatusOK, + "data": map[string]any{ + "billing_scope": "token", + "resolved_rate_multiplier": 1.0, + "effective_rate_multiplier": 1.0, + "peak_rate_enabled": true, + "peak_start": peakStart, + "peak_end": peakEnd, + "peak_rate_multiplier": 10.0, + "timezone": timezone, + }, + }, + }, + }) + mustCreateAccount(s.T(), s.client, &service.Account{ + Name: "current-off-peak-rate", + Extra: map[string]any{ + service.UpstreamBillingProbeExtraKey: map[string]any{ + "status": service.UpstreamBillingProbeStatusOK, + "data": map[string]any{ + "effective_rate_multiplier": 5.0, + }, + }, + }, + }) + + for _, tc := range []struct { + order string + want []string + }{ + {order: "asc", want: []string{"current-off-peak-rate", "current-peak-rate"}}, + {order: "desc", want: []string{"current-peak-rate", "current-off-peak-rate"}}, + } { + accounts, _, err := s.repo.ListWithFilters(s.ctx, pagination.PaginationParams{ + Page: 1, PageSize: 10, SortBy: "upstream_billing_rate", SortOrder: tc.order, + }, "", "", "", "", 0, "") + s.Require().NoError(err) + s.Require().Len(accounts, 2) + for i, name := range tc.want { + s.Require().Equal(name, accounts[i].Name) + } + } +} 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 26fbb74dcf..2a78f38c10 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 @@ -163,6 +163,46 @@ func TestBulkUpdateNilProbeRemovesKeyInsteadOfWritingJSONNull(t *testing.T) { require.Contains(t, normalizeSQLWhitespace(exec.execQueries[0]), "- 'upstream_billing_probe'") } +func TestBulkUpdateDisablingProbeRemovesSnapshot(t *testing.T) { + exec := &recordingSQLExecutor{result: rowsAffectedResult(1)} + repo := newAccountRepositoryWithSQL(nil, exec, nil) + + _, err := repo.BulkUpdate(context.Background(), []int64{27}, service.AccountBulkUpdate{ + Extra: map[string]any{service.UpstreamBillingProbeEnabledExtraKey: false}, + }) + + require.NoError(t, err) + require.NotEmpty(t, exec.execQueries) + require.Contains(t, normalizeSQLWhitespace(exec.execQueries[0]), "- 'upstream_billing_probe'") + payload, ok := exec.execArgs[0][0].([]byte) + require.True(t, ok) + require.Equal(t, `{"upstream_billing_probe_enabled":false}`, string(payload)) +} + +func TestBulkUpdateProbeEligibilityMismatchRollsBack(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() }) + + enabled := true + mock.ExpectBegin() + mock.ExpectExec(`(?s)UPDATE accounts SET extra = .* WHERE id = ANY\(\$2\) AND deleted_at IS NULL AND platform = \$3 AND type = \$4`). + WithArgs(sqlmock.AnyArg(), `{27,28}`, service.PlatformOpenAI, service.AccountTypeAPIKey). + WillReturnResult(sqlmock.NewResult(0, 1)) + mock.ExpectRollback() + + repo := newAccountRepositoryWithSQL(client, db, nil) + rows, err := repo.BulkUpdate(context.Background(), []int64{27, 28}, service.AccountBulkUpdate{ + ProbeEnabled: &enabled, + }) + + require.ErrorIs(t, err, service.ErrUpstreamBillingProbeAccountInvalid) + require.Zero(t, rows) + require.NoError(t, mock.ExpectationsWereMet()) +} + func TestUpdateCredentialsAtomicallyClearsProbeForOpenAIAPIKeyIdentityChange(t *testing.T) { db, mock, err := sqlmock.New() require.NoError(t, err) diff --git a/backend/internal/service/account_service.go b/backend/internal/service/account_service.go index b7fe0a68cc..d1a57e6d25 100644 --- a/backend/internal/service/account_service.go +++ b/backend/internal/service/account_service.go @@ -144,6 +144,7 @@ type AccountBulkUpdate struct { Schedulable *bool Credentials map[string]any Extra map[string]any + ProbeEnabled *bool } // CreateAccountRequest 创建账号请求 diff --git a/backend/internal/service/admin_account.go b/backend/internal/service/admin_account.go index 899679221b..ae3a9da367 100644 --- a/backend/internal/service/admin_account.go +++ b/backend/internal/service/admin_account.go @@ -469,6 +469,15 @@ func buildAccountForCreate(input *CreateAccountInput, accountExtra map[string]an Status: StatusActive, Schedulable: true, } + if input.ProbeEnabled != nil && *input.ProbeEnabled { + if !isUpstreamBillingProbeAccount(account) { + return nil, ErrUpstreamBillingProbeAccountInvalid + } + if account.Extra == nil { + account.Extra = make(map[string]any) + } + account.Extra[UpstreamBillingProbeEnabledExtraKey] = true + } // 预计算固定时间重置的下次重置时间 if account.Extra != nil { if err := ValidateQuotaResetConfig(account.Extra); err != nil { @@ -842,7 +851,7 @@ func (s *adminServiceImpl) UpdateAccountExtra(ctx context.Context, id int64, upd // BulkUpdateAccounts updates multiple accounts in one request. // It merges credentials/extra keys instead of overwriting the whole object. func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUpdateAccountsInput) (*BulkUpdateAccountsResult, error) { - // Probe state is updated only through its dedicated endpoints. + // Managed probe state may only enter through the dedicated typed field below. delete(input.Extra, UpstreamBillingProbeEnabledExtraKey) delete(input.Extra, UpstreamBillingProbeExtraKey) @@ -874,13 +883,30 @@ func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUp // 预取所有目标账号，供凭据守卫/代理守卫/混合渠道检查共用，避免多次 DB 查询。 var cachedTargets []*Account - if len(input.Credentials) > 0 || input.ProxyID != nil || needMixedChannelCheck || hasLongContextBillingUpdate { + if len(input.Credentials) > 0 || input.ProxyID != nil || needMixedChannelCheck || hasLongContextBillingUpdate || input.ProbeEnabled != nil { loaded, err := s.accountRepo.GetByIDs(ctx, input.AccountIDs) if err != nil { return nil, err } cachedTargets = loaded } + if input.ProbeEnabled != nil { + targetsByID := make(map[int64]*Account, len(cachedTargets)) + for _, account := range cachedTargets { + if account != nil { + targetsByID[account.ID] = account + } + } + for _, accountID := range input.AccountIDs { + account, ok := targetsByID[accountID] + if !ok { + return nil, ErrAccountNotFound + } + if !isUpstreamBillingProbeAccount(account) { + return nil, ErrUpstreamBillingProbeAccountInvalid + } + } + } if hasLongContextBillingUpdate { for _, account := range cachedTargets { if account == nil || account.Platform != PlatformOpenAI { @@ -952,8 +978,15 @@ func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUp // Prepare bulk updates for columns and JSONB fields. repoUpdates := AccountBulkUpdate{ - Credentials: input.Credentials, - Extra: input.Extra, + Credentials: input.Credentials, + Extra: input.Extra, + ProbeEnabled: input.ProbeEnabled, + } + if input.ProbeEnabled != nil { + if repoUpdates.Extra == nil { + repoUpdates.Extra = make(map[string]any) + } + repoUpdates.Extra[UpstreamBillingProbeEnabledExtraKey] = *input.ProbeEnabled } 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 727a66cf12..88db653861 100644 --- a/backend/internal/service/admin_account_upstream_billing_probe_test.go +++ b/backend/internal/service/admin_account_upstream_billing_probe_test.go @@ -38,6 +38,32 @@ func TestCreateAccountDropsManagedUpstreamBillingProbeState(t *testing.T) { require.NotContains(t, created.Extra, UpstreamBillingProbeExtraKey) } +func TestCreateAccountAcceptsDedicatedUpstreamBillingProbeSetting(t *testing.T) { + enabled := true + repo := &upstreamBillingProbeAccountRepo{} + created, err := (&adminServiceImpl{accountRepo: repo}).CreateAccount(context.Background(), &CreateAccountInput{ + Name: "upstream", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Credentials: map[string]any{"api_key": "sk-test"}, + ProbeEnabled: &enabled, + SkipDefaultGroupBind: true, + }) + + require.NoError(t, err) + require.Equal(t, true, created.Extra[UpstreamBillingProbeEnabledExtraKey]) + + _, err = (&adminServiceImpl{accountRepo: repo}).CreateAccount(context.Background(), &CreateAccountInput{ + Name: "oauth", + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Credentials: map[string]any{"access_token": "token"}, + ProbeEnabled: &enabled, + SkipDefaultGroupBind: true, + }) + require.ErrorIs(t, err, ErrUpstreamBillingProbeAccountInvalid) +} + func TestUpdateAccountPreservesManagedUpstreamBillingProbeStateForUnrelatedEdit(t *testing.T) { accountID := int64(110) repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{ @@ -360,6 +386,63 @@ func TestBulkUpdateAccountsDropsManagedUpstreamBillingProbeState(t *testing.T) { require.NotContains(t, repo.bulkUpdates[0].Extra, UpstreamBillingProbeExtraKey) } +func TestBulkUpdateAccountsAcceptsDedicatedUpstreamBillingProbeSetting(t *testing.T) { + for _, enabled := range []bool{true, false} { + t.Run(map[bool]string{true: "enable", false: "disable"}[enabled], func(t *testing.T) { + repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{ + 1: {ID: 1, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, + 2: {ID: 2, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, + }} + + result, err := (&adminServiceImpl{accountRepo: repo}).BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{ + AccountIDs: []int64{1, 2}, + ProbeEnabled: &enabled, + }) + + require.NoError(t, err) + require.Equal(t, 2, result.Success) + require.Len(t, repo.bulkUpdates, 1) + require.Equal(t, enabled, repo.bulkUpdates[0].Extra[UpstreamBillingProbeEnabledExtraKey]) + require.NotNil(t, repo.bulkUpdates[0].ProbeEnabled) + require.Equal(t, enabled, *repo.bulkUpdates[0].ProbeEnabled) + }) + } +} + +func TestBulkUpdateAccountsRejectsProbeSettingForIneligibleTargetBeforeWrite(t *testing.T) { + for _, enabled := range []bool{true, false} { + t.Run(map[bool]string{true: "enable", false: "disable"}[enabled], func(t *testing.T) { + repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{ + 1: {ID: 1, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, + 2: {ID: 2, Platform: PlatformOpenAI, Type: AccountTypeOAuth}, + }} + + _, err := (&adminServiceImpl{accountRepo: repo}).BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{ + AccountIDs: []int64{1, 2}, + ProbeEnabled: &enabled, + }) + + require.ErrorIs(t, err, ErrUpstreamBillingProbeAccountInvalid) + require.Empty(t, repo.bulkUpdates) + }) + } +} + +func TestBulkUpdateAccountsRejectsProbeSettingWhenTargetIsMissing(t *testing.T) { + enabled := true + repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{ + 1: {ID: 1, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, + }} + + _, err := (&adminServiceImpl{accountRepo: repo}).BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{ + AccountIDs: []int64{1, 2}, + ProbeEnabled: &enabled, + }) + + require.ErrorIs(t, err, ErrAccountNotFound) + require.Empty(t, repo.bulkUpdates) +} + func TestBulkUpdateAccountsInvalidatesProbeSnapshotForIdentityCredentials(t *testing.T) { repo := &upstreamBillingProbeAccountRepo{} input := &BulkUpdateAccountsInput{ diff --git a/backend/internal/service/admin_service.go b/backend/internal/service/admin_service.go index c167b6be84..013cb04049 100644 --- a/backend/internal/service/admin_service.go +++ b/backend/internal/service/admin_service.go @@ -329,6 +329,7 @@ type CreateAccountInput struct { GroupIDs []int64 ExpiresAt *int64 AutoPauseOnExpired *bool + ProbeEnabled *bool // SkipDefaultGroupBind prevents auto-binding to platform default group when GroupIDs is empty. SkipDefaultGroupBind bool // SkipMixedChannelCheck skips the mixed channel risk check when binding groups. @@ -378,6 +379,7 @@ type BulkUpdateAccountsInput struct { GroupIDs *[]int64 Credentials map[string]any Extra map[string]any + ProbeEnabled *bool // SkipMixedChannelCheck skips the mixed channel risk check when binding groups. // This should only be set when the caller has explicitly confirmed the risk. SkipMixedChannelCheck bool diff --git a/backend/internal/service/upstream_billing_probe.go b/backend/internal/service/upstream_billing_probe.go index 70dbf8f8f2..d5c28e68b6 100644 --- a/backend/internal/service/upstream_billing_probe.go +++ b/backend/internal/service/upstream_billing_probe.go @@ -38,7 +38,7 @@ const ( upstreamBillingProbeMaxBodyBytes = 64 * 1024 upstreamBillingProbeMaxPerCycle = 20 upstreamBillingProbeConcurrency = 4 - upstreamBillingProbeMaxBackoff = 24 * time.Hour + upstreamBillingProbeMaxDelay = 24 * time.Hour upstreamBillingProbeLeaderLockKey = "upstream:billing:probe:leader" upstreamBillingProbeLeaderLockTTL = 2 * time.Minute ) @@ -620,7 +620,7 @@ func (s *UpstreamBillingProbeService) probeLoadedAccount(ctx context.Context, ac ReceivedAt: probeTimePtr(now), FreshUntil: probeTimePtr(now.Add(2 * time.Duration(intervalMinutes) * time.Minute)), LastAttemptAt: now, - NextProbeAt: now.Add(nextProbeDelay(intervalMinutes, 0, 0)), + NextProbeAt: now.Add(nextProbeDelay(intervalMinutes, 0)), HTTPStatus: resp.StatusCode, } if err := s.updateSnapshot(ctx, account, snapshot); err != nil { @@ -650,7 +650,7 @@ func (s *UpstreamBillingProbeService) persistProbeFailure( snapshot := &UpstreamBillingProbeSnapshot{ Status: status, LastAttemptAt: now, - NextProbeAt: now.Add(nextProbeDelay(intervalMinutes, failureCount, retryAfterDuration)), + NextProbeAt: now.Add(nextProbeDelay(intervalMinutes, retryAfterDuration)), FailureCount: failureCount, HTTPStatus: statusCode, LastError: reason, @@ -859,20 +859,13 @@ func (s *UpstreamBillingProbeService) currentTime() time.Time { return time.Now() } -func nextProbeDelay(intervalMinutes, failureCount int, retryAfterDuration time.Duration) time.Duration { +func nextProbeDelay(intervalMinutes int, retryAfterDuration time.Duration) time.Duration { interval := time.Duration(intervalMinutes) * time.Minute if interval < upstreamBillingProbeMinIntervalMinutes*time.Minute { interval = upstreamBillingProbeMinIntervalMinutes * time.Minute } - if failureCount > 0 { - shift := failureCount - if shift > 5 { - shift = 5 - } - interval *= time.Duration(1 << shift) - } - if interval > upstreamBillingProbeMaxBackoff { - interval = upstreamBillingProbeMaxBackoff + if interval > upstreamBillingProbeMaxDelay { + interval = upstreamBillingProbeMaxDelay } jitterRange := interval / 5 if jitterRange > 5*time.Minute { @@ -883,11 +876,11 @@ func nextProbeDelay(intervalMinutes, failureCount int, retryAfterDuration time.D } if retryAfterDuration > interval { // Retry-After is an explicit upstream instruction; do not shorten it - // with the local exponential-backoff ceiling. + // with the local maximum delay. return retryAfterDuration } - if interval > upstreamBillingProbeMaxBackoff { - return upstreamBillingProbeMaxBackoff + if interval > upstreamBillingProbeMaxDelay { + return upstreamBillingProbeMaxDelay } return interval } diff --git a/backend/internal/service/upstream_billing_probe_test.go b/backend/internal/service/upstream_billing_probe_test.go index 0ffcc6006e..eecd866a90 100644 --- a/backend/internal/service/upstream_billing_probe_test.go +++ b/backend/internal/service/upstream_billing_probe_test.go @@ -495,7 +495,7 @@ func TestUpstreamBillingProbeFailurePreservesLastSuccessAndRetryAfter(t *testing } func TestUpstreamBillingProbeRetryAfterIsNotShortened(t *testing.T) { - delay := nextProbeDelay(30, 1, 48*time.Hour) + delay := nextProbeDelay(30, 48*time.Hour) require.Equal(t, 48*time.Hour, delay) } @@ -510,11 +510,21 @@ func TestUpstreamBillingProbeEmptyResponseIsPersistedAsFailure(t *testing.T) { } repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{account.ID: account}} svc := newUpstreamBillingProbeTestService(repo, &httpUpstreamRecorder{}, &upstreamBillingProbeSettingRepo{}) + fixedNow := time.Date(2026, time.July, 13, 2, 0, 0, 0, time.UTC) + svc.now = func() time.Time { return fixedNow } snapshot, err := svc.ProbeAccount(context.Background(), account.ID) require.NoError(t, err) require.Equal(t, UpstreamBillingProbeStatusFailed, snapshot.Status) require.Equal(t, "empty_response", snapshot.LastError) + require.False(t, snapshot.NextProbeAt.Before(fixedNow.Add(24*time.Minute))) + require.False(t, snapshot.NextProbeAt.After(fixedNow.Add(36*time.Minute))) + + snapshot, err = svc.ProbeAccount(context.Background(), account.ID) + require.NoError(t, err) + require.Equal(t, 2, snapshot.FailureCount) + require.False(t, snapshot.NextProbeAt.Before(fixedNow.Add(24*time.Minute))) + require.False(t, snapshot.NextProbeAt.After(fixedNow.Add(36*time.Minute))) } func TestUpstreamBillingProbeUnsupportedAndAccountToggle(t *testing.T) { @@ -533,6 +543,8 @@ func TestUpstreamBillingProbeUnsupportedAndAccountToggle(t *testing.T) { Body: io.NopCloser(strings.NewReader("not found")), }} svc := newUpstreamBillingProbeTestService(repo, upstream, &upstreamBillingProbeSettingRepo{}) + fixedNow := time.Date(2026, time.July, 13, 2, 0, 0, 0, time.UTC) + svc.now = func() time.Time { return fixedNow } require.NoError(t, svc.SetAccountEnabled(context.Background(), account.ID, true)) require.Equal(t, true, account.Extra[UpstreamBillingProbeEnabledExtraKey]) @@ -540,6 +552,14 @@ func TestUpstreamBillingProbeUnsupportedAndAccountToggle(t *testing.T) { require.NoError(t, err) require.Equal(t, UpstreamBillingProbeStatusUnsupported, snapshot.Status) require.Equal(t, "unsupported", snapshot.LastError) + require.False(t, snapshot.NextProbeAt.Before(fixedNow.Add(24*time.Minute))) + require.False(t, snapshot.NextProbeAt.After(fixedNow.Add(36*time.Minute))) + + snapshot, err = svc.ProbeAccount(context.Background(), account.ID) + require.NoError(t, err) + require.Equal(t, 2, snapshot.FailureCount) + require.False(t, snapshot.NextProbeAt.Before(fixedNow.Add(24*time.Minute))) + require.False(t, snapshot.NextProbeAt.After(fixedNow.Add(36*time.Minute))) invalid := &Account{ID: 20, Platform: PlatformOpenAI, Type: AccountTypeOAuth} repo.accounts[invalid.ID] = invalid diff --git a/frontend/src/components/account/BulkEditAccountModal.vue b/frontend/src/components/account/BulkEditAccountModal.vue index 5f744c0fc1..a4268fb926 100644 --- a/frontend/src/components/account/BulkEditAccountModal.vue +++ b/frontend/src/components/account/BulkEditAccountModal.vue @@ -854,6 +854,45 @@ + +
+
+
+ +

+ {{ t('admin.accounts.upstreamBilling.autoProbeHint') }} +

+
+ +
+
+