fix(replay-vision): pre-launch hardening sweep

products/replay_vision/backend/admin.py

@@ -26,6 +26,8 @@ class ReplayScannerAdmin(admin.ModelAdmin):
 @admin.register(ReplayObservation)
 class ReplayObservationAdmin(admin.ModelAdmin):
     list_display = ("scanner", "session_id", "status", "triggered_by", "created_at", "completed_at")
+    # The product's largest table: without this the changelist runs a query per row for the FK columns.
+    list_select_related = ("scanner", "team")
     list_filter = ("status", "triggered_by")
     search_fields = ("session_id", "workflow_id")
     # Observations are workflow-created and immutable post-create except for status/error_reason.
@@ -80,6 +82,7 @@ class VisionActionAdmin(admin.ModelAdmin):
         "next_run_at",
         "created_at",
     )
+    list_select_related = ("team", "scanner")
     list_filter = ("trigger_type", "mode", "enabled")
     search_fields = ("name",)
     raw_id_fields = ("team", "scanner", "hog_flow", "created_by")
@@ -89,6 +92,7 @@ class VisionActionAdmin(admin.ModelAdmin):
 @admin.register(VisionActionRun)
 class VisionActionRunAdmin(admin.ModelAdmin):
     list_display = ("id", "vision_action", "team", "status", "observation_count", "scheduled_at", "created_at")
+    list_select_related = ("vision_action", "team")
     list_filter = ("status",)
     search_fields = ("idempotency_key", "temporal_workflow_id")
     raw_id_fields = ("vision_action", "team")

products/replay_vision/backend/api/errors.py

@@ -0,0 +1,7 @@
+from rest_framework import serializers
+
+
+class ReplayVisionErrorSerializer(serializers.Serializer):
+    """The shape every Replay Vision error response uses, so generated clients read one key."""
+
+    detail = serializers.CharField(help_text="Human-readable explanation of why the request was refused.")

products/replay_vision/backend/api/observations.py

@@ -5,7 +5,7 @@
 from zoneinfo import ZoneInfo
 
 from django.conf import settings
-from django.db import transaction
+from django.db import IntegrityError, transaction
 from django.db.models import Case, CharField, FloatField, Func, IntegerField, Q, QuerySet, Value, When
 from django.db.models.fields.json import KeyTextTransform, KeyTransform
 from django.db.models.functions import Cast
@@ -14,7 +14,13 @@
 import structlog
 import django_filters
 from django_filters.rest_framework import DjangoFilterBackend
-from drf_spectacular.utils import OpenApiParameter, extend_schema, extend_schema_field, extend_schema_view
+from drf_spectacular.utils import (
+    OpenApiParameter,
+    OpenApiResponse,
+    extend_schema,
+    extend_schema_field,
+    extend_schema_view,
+)
 from pydantic import ValidationError as PydanticValidationError
 from rest_framework import mixins, serializers, status, viewsets
 from rest_framework.decorators import action
@@ -30,16 +36,19 @@
 from posthog.renderers import ServerSentEventRenderer
 from posthog.utils import relative_date_parse
 
+from products.replay_vision.backend.api.errors import ReplayVisionErrorSerializer
 from products.replay_vision.backend.api.filters import MultiChoiceFilter, OrderByFilter, ordering_enum
 from products.replay_vision.backend.api.observation_progress import stream_observation_progress
 from products.replay_vision.backend.api.observation_stats import compute_observation_stats
 from products.replay_vision.backend.api.trigger import (
     WorkflowStartOutcome,
     check_observation_quota,
     check_team_in_flight_capacity,
+    claim_apply_scanner_slot,
     start_apply_scanner_workflow,
 )
 from products.replay_vision.backend.billing import observation_credits_for_model
+from products.replay_vision.backend.enqueue_claims import release_enqueue_claim
 from products.replay_vision.backend.error_kinds import ERROR_REASON_HELP_TEXT
 from products.replay_vision.backend.feature_flag import ReplayVisionEnabledPermission
 from products.replay_vision.backend.models.replay_observation import (
@@ -864,7 +873,16 @@ def create_task(self, request: Request, **kwargs: Any) -> Response:
             locked.save(update_fields=["created_task_id"])
         return Response({"task_id": task_id}, status=status.HTTP_201_CREATED)
 
-    @extend_schema(request=None, responses={202: RetryResponseSerializer})
+    @extend_schema(
+        request=None,
+        responses={
+            202: RetryResponseSerializer,
+            409: OpenApiResponse(
+                response=ReplayVisionErrorSerializer, description="The previous run is still finishing."
+            ),
+            503: OpenApiResponse(response=ReplayVisionErrorSerializer, description="The retry couldn't be started."),
+        },
+    )
     @action(detail=True, methods=["post"], required_scopes=["replay_scanner:write", "session_recording:read"])
     def retry(self, request: Request, **kwargs: Any) -> Response:
         """Delete a failed observation and re-run its scanner on the same recording. Returns 202 with the workflow handle."""
@@ -873,32 +891,68 @@ def retry(self, request: Request, **kwargs: Any) -> Response:
         scanner = getattr(self, "_scanner_for_url_cache", None) or observation.scanner
         # Retry writes to the scanner; the session route's get_object only object-checks the observation row.
         self.check_object_permissions(self.request, scanner)
+        session_id = observation.session_id
+        original_pk = observation.pk
+        original_created_at = observation.created_at
         if observation.status != ObservationStatus.FAILED:
             raise ValidationError("Only failed observations can be retried.")
+        # Advisory, and deliberately outside the lock below: the atomic claim is the authoritative gate,
+        # and these two read enough to be worth keeping off a held row lock.
         check_observation_quota(self.team.organization_id, observation_credits_for_model(scanner.model))
         check_team_in_flight_capacity(self.team.id)
-        session_id = observation.session_id
-        # Free the UNIQUE(scanner, session_id) slot; the usage ledger is immutable, so the failed attempt stays counted.
-        original_pk = observation.pk
-        original_created_at = observation.created_at
-        observation.delete()
+        # Locked so two concurrent retries can't both pass the status check and both delete the row.
+        with transaction.atomic():
+            locked = ReplayObservation.objects.select_for_update().get(pk=original_pk, team_id=self.team_id)
+            if locked.status != ObservationStatus.FAILED:
+                raise ValidationError("Only failed observations can be retried.")
+            # Captured before the delete cascades it away: a run that never starts has to put the team's
+            # rating back with the row, not just the row.
+            original_label = ReplayObservationLabel.objects.filter(
+                observation_id=original_pk, team_id=locked.team_id
+            ).first()
+            label_created_at = original_label.created_at if original_label else None
+            label_updated_at = original_label.updated_at if original_label else None
+            # Claimed before the delete so a capped retry never touches the row, and so never cascades
+            # away the observation's shared label for a request that changes nothing.
+            workflow_id, claimed = claim_apply_scanner_slot(scanner, session_id)
+            if not claimed:
+                raise Throttled(detail="This team is at its in-flight observation limit. Try again in a few minutes.")
+            try:
+                # Free the UNIQUE(scanner, session_id) slot; the ledger is immutable, so the failed attempt
+                # stays counted.
+                locked.delete()
+            except Exception:
+                release_enqueue_claim(
+                    team_id=scanner.team_id, scanner_id=scanner.id, workflow_id=workflow_id, immediately=True
+                )
+                raise
         workflow_id, outcome = start_apply_scanner_workflow(
             scanner,
             session_id,
             triggered_by_user_id=cast(User, request.user).id,
             trigger=ObservationTrigger.RETRY,
+            slot_already_claimed=True,
         )
         if outcome is not WorkflowStartOutcome.STARTED:
-            # The replacement run never started, so restore the failed row (its shared label, if any, is lost
-            # to the cascade) instead of leaving the recording looking unscanned.
-            observation.pk = original_pk
-            observation.save(force_insert=True)
-            ReplayObservation.objects.filter(pk=original_pk, team_id=observation.team_id).update(
-                created_at=original_created_at
-            )
-        if outcome is WorkflowStartOutcome.CAPPED:
-            # The pre-check above passed on a snapshot; the atomic claim is the authoritative gate.
-            raise Throttled(detail="This team is at its in-flight observation limit. Try again in a few minutes.")
+            # The replacement run never started, so restore the failed row and its rating instead of leaving
+            # the recording looking unscanned and the team's feedback gone.
+            try:
+                with transaction.atomic():
+                    observation.pk = original_pk
+                    observation.save(force_insert=True)
+                    ReplayObservation.objects.filter(pk=original_pk, team_id=observation.team_id).update(
+                        created_at=original_created_at
+                    )
+                    if original_label is not None:
+                        original_label.save(force_insert=True)
+                        # auto_now_add/auto_now stamp the insert with now; the rating happened earlier.
+                        ReplayObservationLabel.objects.filter(pk=original_label.pk, team_id=observation.team_id).update(
+                            created_at=label_created_at, updated_at=label_updated_at
+                        )
+            except IntegrityError:
+                # A run we couldn't start is already persisting its own row for this (scanner, session);
+                # the recording isn't stranded, so report it as still finishing rather than 500ing.
+                outcome = WorkflowStartOutcome.ALREADY_RUNNING
         if outcome is WorkflowStartOutcome.ALREADY_RUNNING:
             # The prior run is still closing, so its deterministic id blocks the restart and no new row will appear.
             return Response(
@@ -948,16 +1002,23 @@ def label(self, request: Request, **kwargs: Any) -> Response:
             return Response(status=204)
         input_serializer = ReplayObservationLabelSerializer(data=request.data)
         input_serializer.is_valid(raise_exception=True)
-        # team_id in the lookup keeps the query team-scoped.
-        label, _ = ReplayObservationLabel.objects.update_or_create(
-            observation=observation,
-            team_id=observation.team_id,
-            defaults={
-                "is_correct": input_serializer.validated_data["is_correct"],
-                "feedback": input_serializer.validated_data.get("feedback", ""),
-                "created_by": user,
-            },
-        )
+        # Lock the parent row so two concurrent first-time ratings serialize: unlocked, both see no label,
+        # both insert, and the loser hits the OneToOne constraint as a 500.
+        with transaction.atomic():
+            # `only("pk")`: the lock is the point, and the full row drags its JSONB columns along.
+            ReplayObservation.objects.select_for_update().only("pk").filter(
+                pk=observation.pk, team_id=observation.team_id
+            ).first()
+            # team_id in the lookup keeps the query team-scoped.
+            label, _ = ReplayObservationLabel.objects.update_or_create(
+                observation=observation,
+                team_id=observation.team_id,
+                defaults={
+                    "is_correct": input_serializer.validated_data["is_correct"],
+                    "feedback": input_serializer.validated_data.get("feedback", ""),
+                    "created_by": user,
+                },
+            )
         # The core calibration signal: thumbs up/down on whether the scanner got the session right.
         report_user_action(
             user,

products/replay_vision/backend/api/prompt_suggestions.py

@@ -401,8 +401,6 @@ def evaluate(self, request: Request, **kwargs: Any) -> Response:
         input_serializer.is_valid(raise_exception=True)
         session_limit = input_serializer.validated_data["session_limit"]
         edited_config = input_serializer.validated_data.get("config")
-        if suggestion.status != SuggestionStatus.PENDING:
-            raise ValidationError("Only the current pending suggestion can be tested.")
         if not evaluation_supported(scanner):
             raise ValidationError("Testing isn't available for this scanner type.")
         # A malformed edited config must be rejected before it charges credits on runs that can't succeed,
@@ -414,29 +412,38 @@ def evaluate(self, request: Request, **kwargs: Any) -> Response:
         rated_count = self._rated_count(scanner)
         if rated_count == 0:
             raise ValidationError("Rate some results first. They are what the suggestion is tested against.")
-        # A test already in flight keeps reporting its state even if quota ran out meanwhile.
-        if evaluation_in_flight(suggestion.evaluation):
-            return Response(ReplayScannerPromptSuggestionSerializer(suggestion).data)
-        # Each re-run session charges credits like a normal observation, so refuse a test that would
-        # overspend the month. An uncapped org (no credit limit) never trips this.
-        planned = min(session_limit, rated_count)
-        planned_credits = planned * observation_credits_for_model(scanner.model)
-        quota = compute_quota_snapshot(organization_id=self.team.organization_id)
-        if quota.would_exceed(planned_credits):
-            raise QuotaLimitExceeded(
-                detail=(
-                    f"This test would use {planned_credits:,} credits but only {quota.remaining or 0:,} of the "
-                    f"monthly Replay Vision credit limit of {quota.credit_limit or 0:,} remain. Lower the test "
-                    f"session count or wait for the reset on "
-                    f"{quota.period_end.strftime('%b')} {quota.period_end.day}."
-                )
+        # Guards must run on a locked row, like apply and dismiss: on unlocked reads two concurrent tests
+        # both see "not in flight", and the second stub save moves `started_at`, which re-keys the usage
+        # receipts of the first run's still-settling sessions and charges them twice.
+        with transaction.atomic():
+            suggestion = ReplayScannerPromptSuggestion.objects.select_for_update().get(
+                team_id=self.team_id, id=suggestion.id
             )
-
-        # Stamp running first so the UI never sees a gap and the planned spend counts against quota
-        # right away. The select activity replaces this stub with the real total and fingerprint.
-        previous_evaluation = suggestion.evaluation
-        suggestion.evaluation = build_running_evaluation(total=planned, labels_fingerprint="")
-        suggestion.save(update_fields=["evaluation"])
+            if suggestion.status != SuggestionStatus.PENDING:
+                raise ValidationError("Only the current pending suggestion can be tested.")
+            # A test already in flight keeps reporting its state even if quota ran out meanwhile.
+            if evaluation_in_flight(suggestion.evaluation):
+                return Response(ReplayScannerPromptSuggestionSerializer(suggestion).data)
+            # Each re-run session charges credits like a normal observation, so refuse a test that would
+            # overspend the month. An uncapped org (no credit limit) never trips this.
+            planned = min(session_limit, rated_count)
+            planned_credits = planned * observation_credits_for_model(scanner.model)
+            quota = compute_quota_snapshot(organization_id=self.team.organization_id)
+            if quota.would_exceed(planned_credits):
+                raise QuotaLimitExceeded(
+                    detail=(
+                        f"This test would use {planned_credits:,} credits but only {quota.remaining or 0:,} of the "
+                        f"monthly Replay Vision credit limit of {quota.credit_limit or 0:,} remain. Lower the test "
+                        f"session count or wait for the reset on "
+                        f"{quota.period_end.strftime('%b')} {quota.period_end.day}."
+                    )
+                )
+            # Stamp running first so the UI never sees a gap and the planned spend counts against quota
+            # right away. The select activity replaces this stub with the real total and fingerprint.
+            previous_evaluation = suggestion.evaluation
+            suggestion.evaluation = build_running_evaluation(total=planned, labels_fingerprint="", model=scanner.model)
+            suggestion.save(update_fields=["evaluation"])
+        started_at = str(suggestion.evaluation.get("started_at") or "")
         try:
             client = sync_connect()
             async_to_sync(client.start_workflow)(  # type: ignore[misc]
@@ -446,6 +453,7 @@ def evaluate(self, request: Request, **kwargs: Any) -> Response:
                     team_id=scanner.team_id,
                     session_limit=session_limit,
                     config_override=edited_config,
+                    started_at=started_at,
                 ),
                 id=build_evaluate_prompt_suggestion_workflow_id(suggestion.id),
                 task_queue=settings.REPLAY_VISION_TASK_QUEUE,