fix(flags): keep the replay relink alive under activity-signal muting

posthog/test/repo_invariants/setup_receivers_baseline.txt

@@ -35,7 +35,6 @@ model_activity:products.data_tools.backend.activity_logging.handle_expression_ac
 model_activity:products.experiments.backend.activity_logging.handle_experiment_change
 model_activity:products.exports.backend.models.subscription.log_subscription_activity
 model_activity:products.feature_flags.backend.activity_logging.handle_feature_flag_change
-model_activity:products.feature_flags.backend.session_recording_links.relink_teams_on_key_change
 model_activity:products.legal_documents.backend.signals.handle_legal_document_change
 model_activity:products.logs.backend.activity_logging.handle_logs_alert_activity
 model_activity:products.logs.backend.activity_logging.handle_logs_retention_rule_activity
@@ -197,6 +196,7 @@ post_save:products.feature_flags.backend.local_evaluation.evaluation_context_cha
 post_save:products.feature_flags.backend.local_evaluation.evaluation_context_name_changed
 post_save:products.feature_flags.backend.local_evaluation.experiment_changed
 post_save:products.feature_flags.backend.local_evaluation.feature_flag_changed
+post_save:products.feature_flags.backend.session_recording_links.relink_teams_on_key_change
 post_save:products.managed_warehouse.backend.models._observe_duckgres_server_save
 post_save:products.review_hog.backend.receivers.handle_task_run_saved
 post_save:products.signals.backend.receivers.capture_status_change_analytics
@@ -256,6 +256,7 @@ pre_save:products.customer_analytics.backend.models.account._enforce_account_gro
 pre_save:products.event_definitions.backend.models.event_definition._capture_event_definition_old_name
 pre_save:products.feature_flags.backend.flag_version_sync.capture_cohort_definition_before_save
 pre_save:products.feature_flags.backend.flag_version_sync.capture_flag_definition_before_save
+pre_save:products.feature_flags.backend.session_recording_links.capture_replay_link_key_before_save
 pre_save:products.signals.backend.receivers.capture_prior_state
 pre_save:products.surveys.backend.models.pre_save_survey
 user_logged_in:posthog.models.activity_logging.signal_handlers.log_user_login_activity

products/feature_flags/backend/api/feature_flag.py

@@ -120,6 +120,7 @@
     REPLAY_LINKED_FLAG_DELETE_ERROR,
     replay_linked_flag_ids,
     teams_linking_flag,
+    teams_linking_flag_in_project,
 )
 from products.feature_flags.backend.types import PropertyFilterType
 from products.feature_flags.backend.user_blast_radius import get_user_blast_radius
@@ -1948,24 +1949,39 @@ def _free_key_held_by_soft_deleted_flags(self, key: str, exclude_pk: int | None
             soft_deleted_qs = soft_deleted_qs.exclude(pk=exclude_pk)
 
         for flag in soft_deleted_qs:
+            if teams_linking_flag_in_project(self.context["project_id"], flag.id).exists():
+                # Hard-deleting fires no save, so nothing relinks the teams gating replay on
+                # this tombstone and they keep the key the new flag is about to claim. Rename
+                # instead and `relink_teams_on_key_change` moves them onto the tombstone. The
+                # blocker check still runs first: renaming a flag an active dependent references
+                # would silently break that dependent, whether or not a team links it for replay.
+                self._raise_if_key_reuse_blocked(flag)
+                flag.key = flag.tombstoned_key()
+                flag.save(update_fields=["key"])
+                continue
             try:
                 flag.delete()
             except (deletion.RestrictedError, deletion.ProtectedError):
-                blockers = []
-                active_experiment_ids = list(flag.experiment_set.filter(deleted=False).values_list("id", flat=True))
-                if active_experiment_ids:
-                    blockers.append(f"active experiment(s) with ID(s): {', '.join(map(str, active_experiment_ids))}")
-                eaf_count = flag.features.count()
-                if eaf_count:
-                    blockers.append(f"{eaf_count} early access feature(s)")
-                if blockers:
-                    raise exceptions.ValidationError(
-                        f"Cannot reuse key '{flag.key}': a soft-deleted flag with this key is still "
-                        f"referenced by {' and '.join(blockers)}. Please contact support."
-                    )
+                self._raise_if_key_reuse_blocked(flag)
                 flag.key = flag.tombstoned_key()
                 flag.save(update_fields=["key"])
 
+    def _raise_if_key_reuse_blocked(self, flag: FeatureFlag) -> None:
+        # An active dependent (Experiment via RESTRICT, EarlyAccessFeature via PROTECT) means the
+        # tombstone can't be freed without silently breaking that dependent, so refuse the reuse.
+        blockers = []
+        active_experiment_ids = list(flag.experiment_set.filter(deleted=False).values_list("id", flat=True))
+        if active_experiment_ids:
+            blockers.append(f"active experiment(s) with ID(s): {', '.join(map(str, active_experiment_ids))}")
+        eaf_count = flag.features.count()
+        if eaf_count:
+            blockers.append(f"{eaf_count} early access feature(s)")
+        if blockers:
+            raise exceptions.ValidationError(
+                f"Cannot reuse key '{flag.key}': a soft-deleted flag with this key is still "
+                f"referenced by {' and '.join(blockers)}. Please contact support."
+            )
+
     def _reraise_duplicate_key_violation(self, exc: IntegrityError) -> NoReturn:
         # The "unique key for team" DB constraint is the authoritative guard on key
         # uniqueness. A concurrent create or rename can slip a duplicate in between the

products/feature_flags/backend/api/test/test_feature_flag.py

@@ -37,6 +37,7 @@
 from posthog.models.organization import Organization, OrganizationMembership
 from posthog.models.personal_api_key import PersonalAPIKey
 from posthog.models.project_secret_api_key import ProjectSecretAPIKey
+from posthog.models.signals import mute_selected_signals
 from posthog.models.team.extensions import get_or_create_team_extension
 from posthog.models.team.team import Team
 from posthog.models.utils import generate_random_token_personal, hash_key_value
@@ -8116,9 +8117,13 @@ def test_reuses_key_held_by_legacy_soft_deleted_flag_with_soft_deleted_experimen
     def test_reuse_key_with_inconsistent_soft_deleted_flag_referenced_by_active_experiment(self, mode: str):
         # If a tombstone is still referenced by an active experiment (invariant
         # violation), renaming it would silently break the experiment. Error out.
+        # A team gating replay on the tombstone must not change that: the replay relink
+        # takes the rename path, which would otherwise skip this check.
         legacy_flag = FeatureFlag.objects.create(team=self.team, created_by=self.user, key="inconsistent-key")
         exp = Experiment.objects.create(team=self.team, created_by=self.user, feature_flag=legacy_flag)
         FeatureFlag.objects_including_soft_deleted.filter(pk=legacy_flag.pk).update(deleted=True)
+        self.team.session_recording_linked_flag = {"id": legacy_flag.id, "key": "inconsistent-key"}
+        self.team.save()
 
         if mode == "create":
             response = self.client.post(
@@ -8149,7 +8154,8 @@ def test_reuse_key_with_inconsistent_soft_deleted_flag_referenced_by_active_expe
     def test_reuse_key_with_inconsistent_soft_deleted_flag_referenced_by_eaf(self, mode: str):
         # EarlyAccessFeature.feature_flag uses on_delete=PROTECT (sibling of
         # RESTRICT). A tombstone with an EAF must surface the same defensive
-        # error, not a 500.
+        # error, not a 500. A team gating replay on it must not skip that check
+        # by diverting to the rename path.
         legacy_flag = FeatureFlag.objects.create(team=self.team, created_by=self.user, key="eaf-key")
         EarlyAccessFeature.objects.create(
             team=self.team,
@@ -8159,6 +8165,8 @@ def test_reuse_key_with_inconsistent_soft_deleted_flag_referenced_by_eaf(self, m
             feature_flag=legacy_flag,
         )
         FeatureFlag.objects_including_soft_deleted.filter(pk=legacy_flag.pk).update(deleted=True)
+        self.team.session_recording_linked_flag = {"id": legacy_flag.id, "key": "eaf-key"}
+        self.team.save()
 
         if mode == "create":
             response = self.client.post(
@@ -14550,3 +14558,48 @@ def test_rename_does_not_save_teams_it_has_nothing_to_change(self, scope: str, m
         sibling_team.refresh_from_db()
         assert sibling_team.session_recording_linked_flag == linked_flag_before
         assert sibling_team.id not in {call.args[0] for call in mock_refresh.call_args_list}
+
+    def test_rename_still_relinks_teams_while_the_activity_signal_is_muted(self) -> None:
+        # mute_selected_signals is for mass-deletion scenarios where firing the activity-log
+        # signal thousands of times would be wasteful, but it silences every @mutable_receiver
+        # indiscriminately. Muting the audit log must not also stop teams recording sessions.
+        flag = FeatureFlag.objects.create(team=self.team, created_by=self.user, key="replay-gate")
+        self._link_flag(self.team, {"id": flag.id, "key": "replay-gate"})
+
+        with self.captureOnCommitCallbacks(execute=True):
+            with mute_selected_signals():
+                flag.key = "replay-gate-v2"
+                flag.save()
+
+        self.team.refresh_from_db()
+        assert self.team.session_recording_linked_flag == {"id": flag.id, "key": "replay-gate-v2"}
+
+    @parameterized.expand([("create",), ("rename",)])
+    def test_freeing_a_tombstoned_key_relinks_teams_to_the_tombstone(self, mode: str) -> None:
+        # Nothing here blocks the hard delete, which is what makes this the interesting case:
+        # a hard delete fires no save, so a team gating replay on this tombstone would be left
+        # on the key the new flag is about to claim. _free_key_held_by_soft_deleted_flags keeps
+        # the row and renames it instead, and the relink follows it to the tombstone. The rename
+        # path frees the key inside the update transaction, where the tombstone and the claiming
+        # flag each schedule their own relink, so it needs its own coverage.
+        old_flag = FeatureFlag.objects.create(team=self.team, created_by=self.user, key="replay-gate", deleted=True)
+        self._link_flag(self.team, {"id": old_flag.id, "key": "replay-gate"})
+
+        with self.captureOnCommitCallbacks(execute=True):
+            if mode == "create":
+                response = self.client.post(
+                    f"/api/projects/{self.team.id}/feature_flags/",
+                    {"name": "New gate", "key": "replay-gate"},
+                )
+            else:
+                claiming_flag = FeatureFlag.objects.create(team=self.team, created_by=self.user, key="replay-gate-tmp")
+                response = self.client.patch(
+                    f"/api/projects/{self.team.id}/feature_flags/{claiming_flag.id}/",
+                    {"key": "replay-gate"},
+                )
+
+        assert response.status_code in (status.HTTP_200_OK, status.HTTP_201_CREATED), response.content
+        old_flag.refresh_from_db()
+        assert old_flag.key == f"replay-gate:deleted:{old_flag.id}"
+        self.team.refresh_from_db()
+        assert self.team.session_recording_linked_flag == {"id": old_flag.id, "key": old_flag.key}

products/feature_flags/backend/field_snapshots.py

@@ -0,0 +1,50 @@
+"""Snapshots a model's persisted field values across a save, for signal receivers that act
+only when a save really changed something.
+
+A ``pre_save`` receiver stores the values the save is about to overwrite; its ``post_save``
+counterpart compares them against the instance. Each receiver passes its own attribute and
+field set, so two features watching the same model never depend on which fields the other
+happens to watch.
+"""
+
+from typing import Any
+
+from django.db import models
+
+
+def capture_fields_before_save(
+    instance: models.Model,
+    manager: models.Manager[Any],
+    fields: frozenset[str],
+    *,
+    attr: str,
+    update_fields: frozenset[str] | None,
+    raw: bool,
+) -> None:
+    """Snapshot the persisted fields this save may overwrite.
+
+    Always resets the snapshot first so a failed earlier save can never leak a stale
+    capture into a later save's comparison.
+    """
+    setattr(instance, attr, None)
+    if raw or instance.pk is None:
+        return
+    # Only fields this save will actually persist: a field changed in memory but excluded
+    # from update_fields is not written, so it must not count.
+    persisted_fields = fields if update_fields is None else fields.intersection(update_fields)
+    if not persisted_fields:
+        return
+    # sorted(): set iteration order varies between processes, and it reaches the SELECT's
+    # column list — an unstable query string defeats plan reuse and makes the query
+    # snapshots that cover flag saves fail at random.
+    setattr(instance, attr, manager.filter(pk=instance.pk).values(*sorted(persisted_fields)).first())
+
+
+def snapshot_if_changed(instance: models.Model, *, attr: str) -> dict[str, Any] | None:
+    """Pop the pre_save snapshot and return it only if a snapshotted value changed."""
+    before = instance.__dict__.pop(attr, None)
+    if before is None:
+        return None
+    if all(getattr(instance, field) == value for field, value in before.items()):
+        return None
+    return before

products/feature_flags/backend/flag_version_sync.py

@@ -30,7 +30,7 @@
 from collections import defaultdict
 from typing import Any
 
-from django.db import models, transaction
+from django.db import transaction
 from django.db.models import Value
 from django.db.models.functions import Coalesce
 from django.db.models.signals import post_save, pre_save
@@ -43,6 +43,7 @@
 from posthog.models.activity_logging.utils import activity_storage
 
 from products.cohorts.backend.models.cohort import Cohort, CohortOrEmpty
+from products.feature_flags.backend.field_snapshots import capture_fields_before_save, snapshot_if_changed
 from products.feature_flags.backend.models.feature_flag import FeatureFlag
 
 logger = structlog.get_logger(__name__)
@@ -71,42 +72,6 @@
 _DEFINITION_BEFORE_SAVE_ATTR = "_definition_before_save"
 
 
-def _capture_definition_before_save(
-    instance: models.Model,
-    manager: models.Manager[Any],
-    definition_fields: frozenset[str],
-    update_fields: frozenset[str] | None,
-    raw: bool,
-) -> None:
-    """Snapshot the persisted definition fields this save may overwrite.
-
-    Always resets the snapshot first so a failed earlier save can never leak a stale
-    capture into a later save's comparison.
-    """
-    setattr(instance, _DEFINITION_BEFORE_SAVE_ATTR, None)
-    if raw or instance.pk is None:
-        return
-    # Only fields this save will actually persist: a definition field changed in
-    # memory but excluded from update_fields is not written, so it must not count.
-    fields = definition_fields if update_fields is None else definition_fields.intersection(update_fields)
-    if not fields:
-        return
-    # sorted(): set iteration order varies between processes, and it reaches the SELECT's
-    # column list — an unstable query string defeats plan reuse and makes the query
-    # snapshots that cover flag saves fail at random.
-    setattr(instance, _DEFINITION_BEFORE_SAVE_ATTR, manager.filter(pk=instance.pk).values(*sorted(fields)).first())
-
-
-def _changed_definition_fields(instance: models.Model) -> dict[str, Any] | None:
-    """Pop the pre_save snapshot and return it only if a snapshotted value changed."""
-    before = instance.__dict__.pop(_DEFINITION_BEFORE_SAVE_ATTR, None)
-    if before is None:
-        return None
-    if all(getattr(instance, field) == value for field, value in before.items()):
-        return None
-    return before
-
-
 @receiver(pre_save, sender=Cohort)
 def capture_cohort_definition_before_save(
     sender: type[Cohort],
@@ -115,7 +80,14 @@ def capture_cohort_definition_before_save(
     update_fields: frozenset[str] | None = None,
     **kwargs: Any,
 ) -> None:
-    _capture_definition_before_save(instance, Cohort.objects, COHORT_DEFINITION_FIELDS, update_fields, raw)
+    capture_fields_before_save(
+        instance,
+        Cohort.objects,
+        COHORT_DEFINITION_FIELDS,
+        attr=_DEFINITION_BEFORE_SAVE_ATTR,
+        update_fields=update_fields,
+        raw=raw,
+    )
 
 
 @receiver(pre_save, sender=FeatureFlag)
@@ -128,8 +100,13 @@ def capture_flag_definition_before_save(
 ) -> None:
     # objects_including_soft_deleted so restoring a soft-deleted flag reads as a
     # deleted True -> False change instead of snapshotting nothing.
-    _capture_definition_before_save(
-        instance, FeatureFlag.objects_including_soft_deleted, FLAG_DEFINITION_FIELDS, update_fields, raw
+    capture_fields_before_save(
+        instance,
+        FeatureFlag.objects_including_soft_deleted,
+        FLAG_DEFINITION_FIELDS,
+        attr=_DEFINITION_BEFORE_SAVE_ATTR,
+        update_fields=update_fields,
+        raw=raw,
     )
 
 
@@ -143,7 +120,7 @@ def bump_flag_versions_on_cohort_definition_change(
 ) -> None:
     if raw or created:
         return
-    if _changed_definition_fields(instance) is None:
+    if snapshot_if_changed(instance, attr=_DEFINITION_BEFORE_SAVE_ATTR) is None:
         return
 
     project_id = instance.team.project_id
@@ -171,7 +148,7 @@ def bump_dependent_flag_versions_on_flag_definition_change(
 ) -> None:
     if raw or created:
         return
-    before = _changed_definition_fields(instance)
+    before = snapshot_if_changed(instance, attr=_DEFINITION_BEFORE_SAVE_ATTR)
     if before is None:
         return
     # An already-deleted flag is in no payload, so edits to it (e.g. the tombstone rename