feat(seer-billing): Check Seer credits before starting autofix steps

src/sentry/seer/autofix/autofix_agent.py

@@ -9,7 +9,7 @@
 from scm.types import GetBranchProtocol, GetRepositoryProtocol
 from sentry_sdk import traces
 
-from sentry import features, quotas
+from sentry import features
 from sentry.analytics.events.autofix_events import (
     AiAutofixAgentHandoffEvent,
     AiAutofixCodeChangesCompletedEvent,
@@ -23,18 +23,19 @@
     AiAutofixSolutionCompletedEvent,
     AiAutofixSolutionStartedEvent,
 )
-from sentry.constants import ENABLE_SEER_CODING_DEFAULT, DataCategory
+from sentry.constants import ENABLE_SEER_CODING_DEFAULT
 from sentry.integrations.services.integration import integration_service
 from sentry.models.group import Group
 from sentry.seer.agent.client import SeerAgentClient
 from sentry.seer.agent.client_models import SeerRunState
 from sentry.seer.autofix.analytics import record_autofix_event
 from sentry.seer.autofix.commit_author import SeerCommitAuthor
 from sentry.seer.autofix.constants import AutofixReferrer
-from sentry.seer.autofix.exceptions import NoSeerQuotaException
+from sentry.seer.autofix.exceptions import NoSeerQuotaException, SeerOutOfCreditsException
 from sentry.seer.autofix.feature.dispatch import (
     AutofixFeatureArgs,
     autofix_run_targets_other_issue,
+    check_autofix_billing,
     require_autofix_run_for_group,
     trigger_autofix_feature,
 )
@@ -63,6 +64,7 @@
     is_free_cohort_org,
     read_preference_from_sentry_db,
 )
+from sentry.seer.credits import SeerOperation
 from sentry.seer.entrypoints.operator import (
     SeerActivityAttribution,
     SeerAutofixOperator,
@@ -461,20 +463,6 @@ def trigger_autofix_agent(
     """
     outcome = "success"
     try:
-        # check billing quota for triggering a new autofix run
-        # Free cohort orgs bypass quota only when called from agentic triage
-        # (allow_free_cohort=True). The API endpoint never sets this flag,
-        # so manual triggers still require quota.
-        if run_id is None:
-            skip_quota_check = allow_free_cohort and is_free_cohort_org(group.organization)
-            if not skip_quota_check:
-                has_budget: bool = quotas.backend.check_seer_quota(
-                    org_id=group.organization.id,
-                    data_category=DataCategory.SEER_AUTOFIX,
-                )
-                if not has_budget:
-                    raise NoSeerQuotaException()
-
         if run_id is not None:
             _assert_existing_run_belongs_to_group(group, run_id)
 
@@ -548,7 +536,12 @@ def trigger_autofix_agent(
                 actor_user_id=actor_user_id,
             )
         return feature_run
-    except (NoSeerQuotaException, SeerPermissionError, PrIterationNoPullRequestException):
+    except (
+        NoSeerQuotaException,
+        SeerOutOfCreditsException,
+        SeerPermissionError,
+        PrIterationNoPullRequestException,
+    ):
         outcome = "declined"
         raise
     except BaseException:
@@ -839,6 +832,8 @@ def trigger_push_changes(
     ):
         raise PermissionDenied("Code generation is disabled for this organization")
 
+    check_autofix_billing(group.organization, SeerOperation.AUTOFIX_OPEN_PR, is_new_run=False)
+
     client = get_autofix_agent_client(group, user=user)
 
     if state is None:

src/sentry/seer/autofix/exceptions.py

@@ -2,6 +2,10 @@ class NoSeerQuotaException(Exception):
     """Raised when an Autofix run cannot consume Seer quota."""
 
 
+class SeerOutOfCreditsException(Exception):
+    """Raised when the organization doesn't have enough Seer credits for an Autofix operation."""
+
+
 class IssueSummaryUnavailable(Exception):
     """Raised when an issue summary cannot be generated for an expected reason."""
 

src/sentry/seer/autofix/feature/dispatch.py

@@ -13,6 +13,7 @@
     DataCategory,
 )
 from sentry.models.group import Group
+from sentry.models.organization import Organization
 from sentry.models.project import Project
 from sentry.seer.agent.client import SeerAgentClient
 from sentry.seer.agent.client_utils import (
@@ -22,7 +23,7 @@
 )
 from sentry.seer.agent.on_completion_hook import extract_hook_definition
 from sentry.seer.autofix.constants import AutofixReferrer
-from sentry.seer.autofix.exceptions import NoSeerQuotaException
+from sentry.seer.autofix.exceptions import NoSeerQuotaException, SeerOutOfCreditsException
 from sentry.seer.autofix.feature.models import (
     FEATURE_ID,
     LEGACY_FEATURE_ID,
@@ -41,6 +42,7 @@
     get_automation_handoff,
     is_free_cohort_org,
 )
+from sentry.seer.credits import SeerFeature, SeerOperation, has_enough_seer_credits
 from sentry.seer.models.run import SeerAgentRun, SeerRun
 from sentry.seer.models.seer_api_models import UNKNOWN_RUN_ID_FOR_GROUP, SeerPermissionError
 from sentry.users.models.user import User
@@ -148,6 +150,39 @@ def _org_allows_coding(group: Group) -> bool:
     )
 
 
+def check_autofix_billing(
+    organization: Organization,
+    operation: SeerOperation,
+    *,
+    is_new_run: bool,
+    skip_quota_for_free_cohort: bool = False,
+) -> bool:
+    """
+    Raise if the organization can't pay for the operation. Token-billed orgs are
+    checked for credits, and everyone else is checked for quota once per run.
+    Returns whether the operation uses quota.
+    """
+    # Check credits for token-billed orgs.
+    if features.has("organizations:token-billed-seer-enabled", organization):
+        if features.has(
+            "organizations:seer-credit-checks", organization
+        ) and not has_enough_seer_credits(organization, SeerFeature.AUTOFIX, operation):
+            raise SeerOutOfCreditsException()
+        return False
+
+    # Check quota for non-token-billed orgs once per run.
+    # Free cohort orgs skip quota when the caller passes skip_quota_for_free_cohort=True.
+    # Only agentic triage does, and it is not exposed via API.
+    skips_quota = skip_quota_for_free_cohort and is_free_cohort_org(organization)
+    uses_quota = is_new_run and not skips_quota
+    if uses_quota and not quotas.backend.check_seer_quota(
+        org_id=organization.id,
+        data_category=DataCategory.SEER_AUTOFIX,
+    ):
+        raise NoSeerQuotaException()
+    return uses_quota
+
+
 def continue_single_run(
     agent_run: SeerAgentRun,
     group: Group,
@@ -159,6 +194,12 @@ def continue_single_run(
 ) -> SeerRun:
     from sentry.seer.autofix.on_completion_hook import AutofixOnCompletionHook
 
+    check_autofix_billing(
+        group.organization,
+        SeerOperation.AUTOFIX_SINGLE_RUN_REPLY,
+        is_new_run=False,
+    )
+
     enable_coding = _org_allows_coding(group)
     open_pr = enable_coding and can_publish
     payload = AutofixSingleRunFeaturePayload(
@@ -209,24 +250,36 @@ def trigger_autofix_feature(
     from sentry.seer.autofix.on_completion_hook import AutofixOnCompletionHook
 
     is_new_run = args.existing_run_id is None
-    # Free cohort orgs bypass quota only when called from agentic triage
-    # (allow_free_cohort=True). Not exposed via the API.
-    skip_quota = is_new_run and args.allow_free_cohort and is_free_cohort_org(group.organization)
-    if is_new_run and not skip_quota:
-        has_budget: bool = quotas.backend.check_seer_quota(
-            org_id=group.organization.id,
-            data_category=DataCategory.SEER_AUTOFIX,
+    try:
+        uses_quota = check_autofix_billing(
+            group.organization,
+            args.step.seer_operation,
+            is_new_run=is_new_run,
+            # Free cohort orgs bypass quota only when called from agentic triage
+            # (allow_free_cohort=True). Not exposed via the API.
+            skip_quota_for_free_cohort=args.allow_free_cohort,
+        )
+    except NoSeerQuotaException:
+        logger.warning(
+            "autofix_feature.dispatch.quota_denied",
+            extra={
+                "group_id": group.id,
+                "organization_id": group.organization.id,
+                "referrer": args.referrer.value,
+            },
+        )
+        raise
+    except SeerOutOfCreditsException:
+        logger.info(
+            "autofix_feature.dispatch.out_of_credits",
+            extra={
+                "group_id": group.id,
+                "organization_id": group.organization.id,
+                "operation": args.step.seer_operation.value,
+                "referrer": args.referrer.value,
+            },
         )
-        if not has_budget:
-            logger.warning(
-                "autofix_feature.dispatch.quota_denied",
-                extra={
-                    "group_id": group.id,
-                    "organization_id": group.organization.id,
-                    "referrer": args.referrer.value,
-                },
-            )
-            raise NoSeerQuotaException()
+        raise
 
     single_run = should_use_single_run(
         group,
@@ -342,7 +395,7 @@ def trigger_autofix_feature(
     else:
         raise Exception("Unhandled run_id branch, this should never happen")
 
-    if is_new_run and not skip_quota:
+    if uses_quota:
         quotas.backend.record_seer_run(
             group.organization.id, group.project.id, DataCategory.SEER_AUTOFIX
         )

src/sentry/seer/autofix/steps.py

@@ -3,6 +3,7 @@
 from enum import StrEnum
 
 from sentry.seer.autofix.utils import AutofixStoppingPoint
+from sentry.seer.credits import SeerOperation
 
 
 class AutofixStep(StrEnum):
@@ -13,6 +14,19 @@ class AutofixStep(StrEnum):
     CODE_CHANGES = "code_changes"
     PR_ITERATION = "pr_iteration"
 
+    @property
+    def seer_operation(self) -> SeerOperation:
+        """The Seer credit check operation corresponding to this step."""
+        match self:
+            case AutofixStep.ROOT_CAUSE:
+                return SeerOperation.AUTOFIX_ROOT_CAUSE
+            case AutofixStep.SOLUTION:
+                return SeerOperation.AUTOFIX_SOLUTION
+            case AutofixStep.CODE_CHANGES:
+                return SeerOperation.AUTOFIX_CODE_CHANGES
+            case AutofixStep.PR_ITERATION:
+                return SeerOperation.AUTOFIX_PR_ITERATION
+
     @staticmethod
     def from_autofix_stopping_point(
         autofix_stopping_point: AutofixStoppingPoint,

src/sentry/seer/credits.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from collections.abc import Callable
-from dataclasses import dataclass
 from enum import StrEnum
 
 from sentry_protos.billing.v1.services.seer_feature_billing.v1.endpoint_check_seer_credits_pb2 import (
@@ -19,6 +18,7 @@ class SeerOperation(StrEnum):
     AUTOFIX_CODE_CHANGES = "autofix_code_changes"
     AUTOFIX_OPEN_PR = "autofix_open_pr"
     AUTOFIX_PR_ITERATION = "autofix_pr_iteration"
+    AUTOFIX_SINGLE_RUN_REPLY = "autofix_single_run_reply"
 
 
 class SeerFeature(StrEnum):
@@ -28,52 +28,47 @@ class SeerFeature(StrEnum):
     EXPLORER = "explorer"
 
 
-@dataclass(frozen=True)
-class SeerOperationConfig:
-    feature: SeerFeature
-    estimate: Callable[[], int]
-
-
 # TODO(CW-2314): Fake value until we have real estimates.
 PLACEHOLDER_CREDIT_ESTIMATE = 100
 
 
-SEER_OPERATIONS: dict[SeerOperation, SeerOperationConfig] = {
-    SeerOperation.AGENT_CHAT_MESSAGE: SeerOperationConfig(
-        SeerFeature.EXPLORER, lambda: PLACEHOLDER_CREDIT_ESTIMATE
-    ),
-    SeerOperation.AUTOFIX_ROOT_CAUSE: SeerOperationConfig(
-        SeerFeature.AUTOFIX, lambda: PLACEHOLDER_CREDIT_ESTIMATE
-    ),
-    SeerOperation.AUTOFIX_SOLUTION: SeerOperationConfig(
-        SeerFeature.AUTOFIX, lambda: PLACEHOLDER_CREDIT_ESTIMATE
-    ),
-    SeerOperation.AUTOFIX_CODE_CHANGES: SeerOperationConfig(
-        SeerFeature.AUTOFIX, lambda: PLACEHOLDER_CREDIT_ESTIMATE
-    ),
-    SeerOperation.AUTOFIX_OPEN_PR: SeerOperationConfig(
-        SeerFeature.AUTOFIX, lambda: PLACEHOLDER_CREDIT_ESTIMATE
-    ),
-    SeerOperation.AUTOFIX_PR_ITERATION: SeerOperationConfig(
-        SeerFeature.AUTOFIX, lambda: PLACEHOLDER_CREDIT_ESTIMATE
-    ),
+# The operation sets the estimate. Callers pass the feature, since an operation can run under
+# more than one feature.
+SEER_OPERATION_ESTIMATES: dict[SeerOperation, Callable[[], int]] = {
+    SeerOperation.AGENT_CHAT_MESSAGE: lambda: PLACEHOLDER_CREDIT_ESTIMATE,
+    SeerOperation.AUTOFIX_ROOT_CAUSE: lambda: PLACEHOLDER_CREDIT_ESTIMATE,
+    SeerOperation.AUTOFIX_SOLUTION: lambda: PLACEHOLDER_CREDIT_ESTIMATE,
+    SeerOperation.AUTOFIX_CODE_CHANGES: lambda: PLACEHOLDER_CREDIT_ESTIMATE,
+    SeerOperation.AUTOFIX_OPEN_PR: lambda: PLACEHOLDER_CREDIT_ESTIMATE,
+    SeerOperation.AUTOFIX_PR_ITERATION: lambda: PLACEHOLDER_CREDIT_ESTIMATE,
+    SeerOperation.AUTOFIX_SINGLE_RUN_REPLY: lambda: PLACEHOLDER_CREDIT_ESTIMATE,
 }
 
 
-def can_run_seer_operation(organization: Organization, operation: SeerOperation) -> bool:
-    """Whether the organization is token-billed and has enough Seer credits to run the operation."""
-    if not features.has("organizations:seer-credit-checks", organization) or not features.has(
-        "organizations:token-billed-seer-enabled", organization
-    ):
-        return True
-
-    config = SEER_OPERATIONS[operation]
+def has_enough_seer_credits(
+    organization: Organization, feature: SeerFeature, operation: SeerOperation
+) -> bool:
+    """
+    Checks whether the organization has enough Seer credits to run the operation,
+    for callers that already checked the credit check and token-billed flags.
+    """
     response = seer_credit.backend.check_seer_credits(
         CheckSeerCreditsRequest(
             organization_id=organization.id,
-            feature_id=config.feature.value,
-            predicted_credit_usage=config.estimate(),
+            feature_id=feature.value,
+            predicted_credit_usage=SEER_OPERATION_ESTIMATES[operation](),
             take_lease=False,
         )
     )
     return response.has_enough_credits
+
+
+def can_run_seer_operation(
+    organization: Organization, feature: SeerFeature, operation: SeerOperation
+) -> bool:
+    """Checks whether the organization is token-billed and has enough Seer credits to run the operation."""
+    if not features.has("organizations:seer-credit-checks", organization) or not features.has(
+        "organizations:token-billed-seer-enabled", organization
+    ):
+        return True
+    return has_enough_seer_credits(organization, feature, operation)

src/sentry/seer/endpoints/organization_seer_agent_chat.py

@@ -37,7 +37,7 @@
 from sentry.seer.autofix.feature.dispatch import continue_single_run
 from sentry.seer.autofix.feature.models import SINGLE_RUN_FEATURE_ID
 from sentry.seer.autofix.single_run import is_single_run
-from sentry.seer.credits import SeerOperation, can_run_seer_operation
+from sentry.seer.credits import SeerFeature, SeerOperation, can_run_seer_operation
 from sentry.seer.endpoints.group_ai_autofix import RUN_IN_FLIGHT_CODE
 from sentry.seer.endpoints.utils import ResolvedSeerRun, resolve_seer_run
 from sentry.seer.models import SeerApiError, SeerPermissionError
@@ -439,7 +439,7 @@ def post(
             raise PermissionDenied(error)
         # Do not bill for dashboard generation.
         if not is_dashboard_generate_run and not can_run_seer_operation(
-            organization, SeerOperation.AGENT_CHAT_MESSAGE
+            organization, SeerFeature.EXPLORER, SeerOperation.AGENT_CHAT_MESSAGE
         ):
             return Response(
                 {