From b754ff6f1ba826a14700aadd2a841a28c8d46bfa Mon Sep 17 00:00:00 2001
From: Rain <94058511+Pluviobyte@users.noreply.github.com>
Date: Tue, 22 Sep 2026 14:51:44 +0800
Subject: [PATCH] fix(proxies): invalidate billing probe on fallback restore

---
 backend/internal/repository/account_repo.go   |  6 +-
 .../proxy_restore_probe_integration_test.go   | 66 +++++++++++++++++++
 2 files changed, 71 insertions(+), 1 deletion(-)
 create mode 100644 backend/internal/repository/proxy_restore_probe_integration_test.go

diff --git a/backend/internal/repository/account_repo.go b/backend/internal/repository/account_repo.go
index 8b8c82dd66..84e8fc156a 100644
--- a/backend/internal/repository/account_repo.go
+++ b/backend/internal/repository/account_repo.go
@@ -3887,8 +3887,12 @@ func (r *accountRepository) ResetQuotaUsedAndClearRateLimitCooldown(ctx context.
 // 仅当 proxy_fallback_origin_id IS NOT NULL 时执行更新；
 // 若影响行数为 0，则返回 ErrAccountNotInFallback（账号存在但不在 fallback 状态）。
 func (r *accountRepository) RevertProxyFallback(ctx context.Context, accountID int64) error {
+	// Probe snapshots belong to the network identity; invalidate only on a real proxy change.
 	res, err := r.sql.ExecContext(ctx, `
-		UPDATE accounts SET proxy_id=proxy_fallback_origin_id, proxy_fallback_origin_id=NULL, updated_at=NOW()
+		UPDATE accounts SET
+			extra=CASE WHEN type='apikey' AND proxy_id IS DISTINCT FROM proxy_fallback_origin_id
+				THEN extra - 'upstream_billing_probe' ELSE extra END,
+			proxy_id=proxy_fallback_origin_id, proxy_fallback_origin_id=NULL, updated_at=NOW()
 		WHERE id=$1 AND proxy_fallback_origin_id IS NOT NULL AND deleted_at IS NULL`, accountID)
 	if err != nil {
 		return err
diff --git a/backend/internal/repository/proxy_restore_probe_integration_test.go b/backend/internal/repository/proxy_restore_probe_integration_test.go
new file mode 100644
index 0000000000..ebbf77a37e
--- /dev/null
+++ b/backend/internal/repository/proxy_restore_probe_integration_test.go
@@ -0,0 +1,66 @@
+//go:build integration
+
+package repository
+
+import (
+	"encoding/json"
+	"github.com/Wei-Shaw/sub2api/internal/service"
+)
+
+func (s *ProxyExpirySuite) TestRevertClearsBillingProbe() {
+	original := s.mkProxy("original", service.FallbackModeDirect, nil, nil)
+	backup := s.mkProxy("backup", service.FallbackModeNone, nil, nil)
+	account := s.mkAccountWithProxy(backup)
+	_, err := s.tx.ExecContext(s.ctx, `UPDATE accounts SET type='apikey',proxy_fallback_origin_id=$1,extra='{"upstream_billing_probe_enabled":true,"upstream_billing_probe":{"status":"ok","data":{"balance":123}},"keep_me":true}'::jsonb WHERE id=$2`, original, account)
+	s.Require().NoError(err)
+	repo := newAccountRepositoryWithSQL(s.tx.Client(), s.tx, nil)
+	s.Require().NoError(repo.RevertProxyFallback(s.ctx, account))
+	s.Require().Equal(&original, s.accountProxyID(account))
+	var raw []byte
+	s.Require().NoError(scanSingleRow(s.ctx, s.tx, `SELECT extra FROM accounts WHERE id=$1`, []any{account}, &raw))
+	var extra map[string]any
+	s.Require().NoError(json.Unmarshal(raw, &extra))
+	s.NotContains(extra, "upstream_billing_probe", "probe from previous network identity must be invalidated")
+	s.Equal(true, extra["keep_me"])
+}
+
+func (s *ProxyExpirySuite) TestRevertProbeInvalidationScope() {
+	for _, tc := range []struct {
+		name, kind          string
+		same, direct, clear bool
+	}{
+		{"changed API key", "apikey", false, false, true},
+		{"direct API key", "apikey", false, true, true},
+		{"same proxy", "apikey", true, false, false},
+		{"oauth metadata", "oauth", false, false, false},
+	} {
+		s.Run(tc.name, func() {
+			original := s.mkProxy("origin", service.FallbackModeNone, nil, nil)
+			current := s.mkProxy("current", service.FallbackModeNone, nil, nil)
+			if tc.same {
+				current = original
+			}
+			account := s.mkAccountWithProxy(current)
+			_, err := s.tx.ExecContext(s.ctx, `UPDATE accounts SET type=$1,proxy_fallback_origin_id=$2,extra='{"upstream_billing_probe_enabled":true,"upstream_billing_probe":{"status":"ok"},"keep_me":true}'::jsonb WHERE id=$3`, tc.kind, original, account)
+			s.Require().NoError(err)
+			if tc.direct {
+				_, err = s.tx.ExecContext(s.ctx, `UPDATE accounts SET proxy_id=NULL WHERE id=$1`, account)
+				s.Require().NoError(err)
+			}
+			repo := newAccountRepositoryWithSQL(s.tx.Client(), s.tx, nil)
+			s.Require().NoError(repo.RevertProxyFallback(s.ctx, account))
+			var raw []byte
+			s.Require().NoError(scanSingleRow(s.ctx, s.tx, `SELECT extra FROM accounts WHERE id=$1`, []any{account}, &raw))
+			var extra map[string]any
+			s.Require().NoError(json.Unmarshal(raw, &extra))
+			if tc.clear {
+				s.NotContains(extra, "upstream_billing_probe")
+			} else {
+				s.Contains(extra, "upstream_billing_probe")
+			}
+			s.Equal(true, extra["upstream_billing_probe_enabled"])
+			s.Equal(true, extra["keep_me"])
+			s.Equal(&original, s.accountProxyID(account))
+		})
+	}
+}