diff --git a/src/sentry/seer/autofix/autofix_agent.py b/src/sentry/seer/autofix/autofix_agent.py index e28bcb003316..e954513c397e 100644 --- a/src/sentry/seer/autofix/autofix_agent.py +++ b/src/sentry/seer/autofix/autofix_agent.py @@ -36,6 +36,7 @@ AutofixFeatureArgs, autofix_run_targets_other_issue, check_autofix_billing, + get_run_billing_feature, require_autofix_run_for_group, trigger_autofix_feature, ) @@ -438,6 +439,7 @@ def trigger_autofix_agent( group: Group, step: AutofixStep, referrer: AutofixReferrer, + *, run_id: int | None = None, stopping_point: AutofixStoppingPoint | None = None, user_context: str | None = None, @@ -458,6 +460,8 @@ def trigger_autofix_agent( group: The Sentry group (issue) to analyze step: Which autofix step to run run_id: Existing run ID to continue, or None for new run + billing_feature: Internal-only. The feature the run is billed to for its whole life. + Not exposed via the API. stopping_point: Where to stop the automated pipeline (only used for new runs) allow_free_cohort: Internal-only flag set by agentic triage to bypass quota for free cohort orgs. Not exposed via the API. @@ -498,9 +502,9 @@ def trigger_autofix_agent( args = AutofixFeatureArgs( step=step, referrer=referrer, + step_args=step_args, existing_run_id=run_id, insert_index=insert_index, - step_args=step_args, user_context=user_context, stopping_point=stopping_point, allow_free_cohort=allow_free_cohort, @@ -835,8 +839,6 @@ 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: @@ -844,6 +846,13 @@ def trigger_push_changes( else: _validate_run_belongs_to_group(state, group) + check_autofix_billing( + group.organization, + get_run_billing_feature(state), + SeerOperation.AUTOFIX_OPEN_PR, + is_new_run=False, + ) + record_autofix_event( AiAutofixPrCreatedStartedEvent( organization_id=group.organization.id, diff --git a/src/sentry/seer/autofix/feature/dispatch.py b/src/sentry/seer/autofix/feature/dispatch.py index b73537daa5e2..4c3a522e78a3 100644 --- a/src/sentry/seer/autofix/feature/dispatch.py +++ b/src/sentry/seer/autofix/feature/dispatch.py @@ -17,6 +17,7 @@ from sentry.models.organization import Organization from sentry.models.project import Project from sentry.seer.agent.client import SeerAgentClient +from sentry.seer.agent.client_models import SeerRunState from sentry.seer.agent.client_utils import ( AgentRunOptions, collect_user_org_context, @@ -43,7 +44,7 @@ get_automation_handoff, is_free_cohort_org, ) -from sentry.seer.credits import SeerFeature, SeerOperation, has_enough_seer_credits +from sentry.seer.credits import 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 @@ -88,6 +89,12 @@ def _mirror_matches_issue(group: Group) -> Q: return Q(group_id=group.id, project_id=group.project_id) +def get_run_billing_feature(state: SeerRunState) -> int: + """The billable feature the run was started under. Runs from before Seer + stored it are unattributed.""" + return state.billing_feature or billing.SEER_BILLING_FEATURE_UNSPECIFIED + + def autofix_run_targets_other_issue(group: Group, run_id: int) -> bool: """True unless ``run_id`` is mirrored to this issue's group and project. @@ -162,6 +169,7 @@ def _org_allows_coding(group: Group) -> bool: def check_autofix_billing( organization: Organization, + billing_feature: int, operation: SeerOperation, *, is_new_run: bool, @@ -169,14 +177,14 @@ def check_autofix_billing( ) -> 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. + checked for credits under the run's billing feature, 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): + ) and not has_enough_seer_credits(organization, billing_feature, operation): raise SeerOutOfCreditsException() return False @@ -198,14 +206,16 @@ def continue_single_run( group: Group, reply: str, *, + referrer: AutofixReferrer, + billing_feature: int, user: User | RpcUser, can_publish: bool, - referrer: AutofixReferrer, ) -> SeerRun: from sentry.seer.autofix.on_completion_hook import AutofixOnCompletionHook check_autofix_billing( group.organization, + billing_feature, SeerOperation.AUTOFIX_SINGLE_RUN_REPLY, is_new_run=False, ) @@ -259,10 +269,26 @@ def trigger_autofix_feature( # Avoid a circular import through the legacy Autofix dispatcher. from sentry.seer.autofix.on_completion_hook import AutofixOnCompletionHook - is_new_run = args.existing_run_id is None + existing_agent_run: SeerAgentRun | None = None + if args.existing_run_id is not None: + existing_agent_run = require_autofix_run_for_group(group, args.existing_run_id) + is_new_run = existing_agent_run is None + + billing_feature = args.billing_feature + if billing_feature is None: + # Continuations without the stored feature are unattributed rather than guessed. + billing_feature = ( + _BILLING_FEATURE_BY_REFERRER.get( + args.referrer, billing.SEER_BILLING_FEATURE_AUTOFIX_HEADLESS + ) + if is_new_run + else billing.SEER_BILLING_FEATURE_UNSPECIFIED + ) + try: uses_quota = check_autofix_billing( group.organization, + billing_feature, args.step.seer_operation, is_new_run=is_new_run, # Free cohort orgs bypass quota only when called from agentic triage @@ -382,10 +408,7 @@ def trigger_autofix_feature( if is_new_run: run = client.start_feature_run( feature_id=SINGLE_RUN_FEATURE_ID if single_run else FEATURE_ID, - billing_feature=args.billing_feature - or _BILLING_FEATURE_BY_REFERRER.get( - args.referrer, billing.SEER_BILLING_FEATURE_AUTOFIX_HEADLESS - ), + billing_feature=billing_feature, payload=payload.dict(), referrer=args.referrer.value, user_org_context=user_org_context, @@ -395,9 +418,7 @@ def trigger_autofix_feature( flush=args.flush, extras=extras, ) - elif args.existing_run_id is not None: - existing_agent_run = require_autofix_run_for_group(group, args.existing_run_id) - + elif existing_agent_run is not None: run = client.continue_feature_run( existing_agent_run=existing_agent_run, payload=payload.dict(), diff --git a/src/sentry/seer/autofix/on_completion_hook.py b/src/sentry/seer/autofix/on_completion_hook.py index 194e93b412eb..d297fa9445c5 100644 --- a/src/sentry/seer/autofix/on_completion_hook.py +++ b/src/sentry/seer/autofix/on_completion_hook.py @@ -29,6 +29,7 @@ from sentry.seer.autofix.coding_agent import IntegrationNotFound from sentry.seer.autofix.constants import AutofixReferrer from sentry.seer.autofix.exceptions import SeerOutOfCreditsException +from sentry.seer.autofix.feature.dispatch import get_run_billing_feature from sentry.seer.autofix.feature.models import FEATURE_ID as AUTOFIX_FEATURE_ID from sentry.seer.autofix.feature.models import LEGACY_FEATURE_ID as LEGACY_AUTOFIX_FEATURE_ID from sentry.seer.autofix.pr_iteration.completion import ( @@ -793,6 +794,7 @@ def _maybe_continue_pipeline( step=next_step, referrer=referrer, run_id=run_id, + billing_feature=get_run_billing_feature(state), ) except SeerOutOfCreditsException: logger.info( diff --git a/src/sentry/seer/autofix/pr_iteration/completion.py b/src/sentry/seer/autofix/pr_iteration/completion.py index ef8997de3a32..143eb7623b0b 100644 --- a/src/sentry/seer/autofix/pr_iteration/completion.py +++ b/src/sentry/seer/autofix/pr_iteration/completion.py @@ -16,7 +16,11 @@ from sentry.seer.autofix.commit_author import SeerCommitAuthor, parse_commit_author from sentry.seer.autofix.constants import AutofixReferrer from sentry.seer.autofix.exceptions import SeerOutOfCreditsException -from sentry.seer.autofix.feature.dispatch import AutofixFeatureArgs, trigger_autofix_feature +from sentry.seer.autofix.feature.dispatch import ( + AutofixFeatureArgs, + get_run_billing_feature, + trigger_autofix_feature, +) from sentry.seer.autofix.feature.models import PrIterationStepArgs from sentry.seer.autofix.pr_iteration.current_iteration import triggered_iteration_id from sentry.seer.autofix.pr_iteration.emit import ( @@ -582,6 +586,7 @@ def _retry_push_on_moved_branch( pr.repo_name: pr.pr_url for pr in state.repo_pr_states.values() if pr.pr_url }, ), + billing_feature=get_run_billing_feature(state), ), ) except SeerOutOfCreditsException: diff --git a/src/sentry/seer/autofix/pr_iteration/manual_trigger.py b/src/sentry/seer/autofix/pr_iteration/manual_trigger.py index c53e909e070e..d18990e76551 100644 --- a/src/sentry/seer/autofix/pr_iteration/manual_trigger.py +++ b/src/sentry/seer/autofix/pr_iteration/manual_trigger.py @@ -20,12 +20,13 @@ from sentry.models.group import Group from sentry.seer.autofix.autofix_agent import get_autofix_run_state from sentry.seer.autofix.constants import AutofixReferrer +from sentry.seer.autofix.feature.dispatch import get_run_billing_feature from sentry.seer.autofix.pr_iteration.emit import bootstrap_iteration from sentry.seer.autofix.pr_iteration.feedback import Feedback from sentry.seer.autofix.pr_iteration.feedback_sources.base import ConsumeTriggerSource from sentry.seer.autofix.pr_iteration.pause import PauseReason, get_pause_reason from sentry.seer.autofix.pr_iteration.queue import enqueue_autofix_feedback -from sentry.seer.credits import SeerFeature, SeerOperation, can_run_seer_operation +from sentry.seer.credits import SeerOperation, can_run_seer_operation from sentry.seer.endpoints.utils import ( SEER_AUTOFIX_CONTINUE_OUT_OF_CREDITS_DETAIL, SEER_OUT_OF_CREDITS_CODE, @@ -88,7 +89,7 @@ def handle_ui_feedback( # Check credits here, since the queued feedback is consumed after this request returns. if not can_run_seer_operation( - group.organization, SeerFeature.AUTOFIX, SeerOperation.AUTOFIX_PR_ITERATION + group.organization, get_run_billing_feature(run_state), SeerOperation.AUTOFIX_PR_ITERATION ): return Response( { diff --git a/src/sentry/seer/credits.py b/src/sentry/seer/credits.py index 6debf2e2f7f0..a52b3caec79c 100644 --- a/src/sentry/seer/credits.py +++ b/src/sentry/seer/credits.py @@ -3,6 +3,10 @@ from collections.abc import Callable from enum import StrEnum +from sentry_protos.billing.v1.common.v1.seer_billing_feature_pb2 import ( + SEER_BILLING_FEATURE_UNSPECIFIED, + SeerBillingFeature, +) from sentry_protos.billing.v1.services.seer_feature_billing.v1.endpoint_check_seer_credits_pb2 import ( CheckSeerCreditsRequest, ) @@ -21,11 +25,18 @@ class SeerOperation(StrEnum): AUTOFIX_SINGLE_RUN_REPLY = "autofix_single_run_reply" -class SeerFeature(StrEnum): - """Seer features as tagged on LLM proxy calls, which billing maps to data categories.""" - - AUTOFIX = "autofix" - EXPLORER = "explorer" +def billing_feature_label(feature: int) -> str: + """ + The label LLM proxy usage and billing use for a SeerBillingFeature value: the enum name, + lowercased, without its prefix. Missing attribution is "unattributed", like the LLM proxy. + """ + if feature == SEER_BILLING_FEATURE_UNSPECIFIED: + return "unattributed" + return ( + SeerBillingFeature.Name(SeerBillingFeature.ValueType(feature)) + .removeprefix("SEER_BILLING_FEATURE_") + .lower() + ) # TODO(CW-2314): Fake value until we have real estimates. @@ -46,7 +57,7 @@ class SeerFeature(StrEnum): def has_enough_seer_credits( - organization: Organization, feature: SeerFeature, operation: SeerOperation + organization: Organization, feature: int, operation: SeerOperation ) -> bool: """ Checks whether the organization has enough Seer credits to run the operation, @@ -55,7 +66,7 @@ def has_enough_seer_credits( response = seer_credit.backend.check_seer_credits( CheckSeerCreditsRequest( organization_id=organization.id, - feature_id=feature.value, + feature_id=billing_feature_label(feature), predicted_credit_usage=SEER_OPERATION_ESTIMATES[operation](), take_lease=False, ) @@ -64,9 +75,11 @@ def has_enough_seer_credits( def can_run_seer_operation( - organization: Organization, feature: SeerFeature, operation: SeerOperation + organization: Organization, feature: int, operation: SeerOperation ) -> bool: - """Checks whether the organization is token-billed and has enough Seer credits to run the operation.""" + """ + 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 ): diff --git a/src/sentry/seer/endpoints/admin_autofix_retry.py b/src/sentry/seer/endpoints/admin_autofix_retry.py index 7a45ad6108d9..fcd4074d8b7a 100644 --- a/src/sentry/seer/endpoints/admin_autofix_retry.py +++ b/src/sentry/seer/endpoints/admin_autofix_retry.py @@ -19,6 +19,7 @@ ) from sentry.seer.autofix.constants import AutofixReferrer from sentry.seer.autofix.exceptions import SeerOutOfCreditsException +from sentry.seer.autofix.feature.dispatch import get_run_billing_feature from sentry.seer.autofix.feature.models import FEATURE_ID, LEGACY_FEATURE_ID from sentry.seer.autofix.on_completion_hook import PIPELINE_ORDER from sentry.seer.models import SeerRun @@ -119,6 +120,7 @@ def _retry_run(organization: Organization, run_id: int) -> dict[str, Any]: step=step, referrer=step_referrer or run_referrer or AutofixReferrer.UNKNOWN, run_id=run_id, + billing_feature=get_run_billing_feature(state), insert_index=insert_index, ) except SeerOutOfCreditsException: diff --git a/src/sentry/seer/endpoints/group_ai_autofix.py b/src/sentry/seer/endpoints/group_ai_autofix.py index 21bdf55441c3..175920553276 100644 --- a/src/sentry/seer/endpoints/group_ai_autofix.py +++ b/src/sentry/seer/endpoints/group_ai_autofix.py @@ -61,7 +61,11 @@ from sentry.seer.autofix.commit_author import commit_author_for_user from sentry.seer.autofix.constants import PUBLISH_SCOPES, AutofixReferrer from sentry.seer.autofix.exceptions import NoSeerQuotaException, SeerOutOfCreditsException -from sentry.seer.autofix.feature.dispatch import project_stopping_point, should_use_single_run +from sentry.seer.autofix.feature.dispatch import ( + get_run_billing_feature, + project_stopping_point, + should_use_single_run, +) from sentry.seer.autofix.feature.models import FEATURE_ID as AUTOFIX_FEATURE_ID from sentry.seer.autofix.feature.models import SINGLE_RUN_FEATURE_ID from sentry.seer.autofix.on_completion_hook import get_run_stopping_point @@ -514,6 +518,8 @@ def post( status=status.HTTP_409_CONFLICT, ) + # The client can supply referrer. Resolve attribution from the request instead. + billing_feature: int = billing.SEER_BILLING_FEATURE_AUTOFIX_HEADLESS if resolved_run_id is not None: try: run_state = get_autofix_run_state(group, resolved_run_id) @@ -521,6 +527,7 @@ def post( if _is_unknown_run_id_error(e): return Response(status=status.HTTP_404_NOT_FOUND) raise PermissionDenied(SEER_PERMISSION_DENIED) + billing_feature = get_run_billing_feature(run_state) # A continued run keeps the stopping point it started with, and the # completion hook acts on it after this step, so gate on it too. @@ -562,8 +569,6 @@ def post( status=status.HTTP_409_CONFLICT, ) - # The client can supply referrer. Resolve attribution from the request instead. - billing_feature = billing.SEER_BILLING_FEATURE_AUTOFIX_HEADLESS if is_autofix_kickoff: actor = resolve_action_actor(request) source = resolve_action_source(request) @@ -721,6 +726,7 @@ def get(self, request: Request, group: Group) -> Response[AutofixStateResponse]: else AUTOFIX_FEATURE_ID, "status": state.status, "failure_reason": state.failure_reason, + "billing_feature": state.billing_feature, "blocks": blocks, "updated_at": state.updated_at, "pending_user_input": ( diff --git a/src/sentry/seer/endpoints/organization_seer_agent_chat.py b/src/sentry/seer/endpoints/organization_seer_agent_chat.py index 12416ecab212..e696a8799c43 100644 --- a/src/sentry/seer/endpoints/organization_seer_agent_chat.py +++ b/src/sentry/seer/endpoints/organization_seer_agent_chat.py @@ -11,6 +11,9 @@ from rest_framework.request import Request from rest_framework.response import Response from rest_framework.views import APIView +from sentry_protos.billing.v1.common.v1.seer_billing_feature_pb2 import ( + SEER_BILLING_FEATURE_SEER_AGENT, +) from sentry import features from sentry.api.api_owners import ApiOwner @@ -35,10 +38,10 @@ AutofixReferrer, ) from sentry.seer.autofix.exceptions import SeerOutOfCreditsException -from sentry.seer.autofix.feature.dispatch import continue_single_run +from sentry.seer.autofix.feature.dispatch import continue_single_run, get_run_billing_feature 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 SeerFeature, SeerOperation, can_run_seer_operation +from sentry.seer.credits import SeerOperation, can_run_seer_operation from sentry.seer.endpoints.group_ai_autofix import RUN_IN_FLIGHT_CODE from sentry.seer.endpoints.utils import ( SEER_AGENT_CHAT_OUT_OF_CREDITS_DETAIL, @@ -251,12 +254,10 @@ def _maybe_reply_to_single_run( if not request.access.has_any_project_scope(project, AUTOFIX_RUN_SCOPES): raise PermissionDenied("You do not have access to this run.") try: - if ( - SeerAgentClient(organization, request.user) - .get_run(run_id=resolved.seer_run_state_id) - .status - == "processing" - ): + run_state = SeerAgentClient(organization, request.user).get_run( + run_id=resolved.seer_run_state_id + ) + if run_state.status == "processing": return Response( { "detail": "This run is still working", @@ -270,9 +271,10 @@ def _maybe_reply_to_single_run( single_run, single_run.group, reply, + referrer=AutofixReferrer.SEER_EXPLORER, + billing_feature=get_run_billing_feature(run_state), user=replier, can_publish=request.access.has_any_project_scope(project, PUBLISH_SCOPES), - referrer=AutofixReferrer.SEER_EXPLORER, ) except SeerOutOfCreditsException: return Response( @@ -450,7 +452,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, SeerFeature.EXPLORER, SeerOperation.AGENT_CHAT_MESSAGE + organization, SEER_BILLING_FEATURE_SEER_AGENT, SeerOperation.AGENT_CHAT_MESSAGE ): return Response( {"detail": SEER_AGENT_CHAT_OUT_OF_CREDITS_DETAIL, "code": SEER_OUT_OF_CREDITS_CODE}, diff --git a/src/sentry/seer/endpoints/organization_seer_credit_check.py b/src/sentry/seer/endpoints/organization_seer_credit_check.py index 16d580ddedfc..3e36496c92ed 100644 --- a/src/sentry/seer/endpoints/organization_seer_credit_check.py +++ b/src/sentry/seer/endpoints/organization_seer_credit_check.py @@ -3,6 +3,12 @@ from rest_framework import serializers from rest_framework.request import Request from rest_framework.response import Response +from sentry_protos.billing.v1.common.v1.seer_billing_feature_pb2 import ( + SEER_BILLING_FEATURE_AUTOFIX_WEB, + SEER_BILLING_FEATURE_SEER_AGENT, + SEER_BILLING_FEATURE_UNSPECIFIED, + SeerBillingFeature, +) from sentry.api.api_owners import ApiOwner from sentry.api.api_publish_status import ApiPublishStatus @@ -10,12 +16,20 @@ from sentry.api.bases.organization import OrganizationEndpoint from sentry.models.organization import Organization from sentry.ratelimits.config import RateLimitConfig -from sentry.seer.credits import SeerFeature, SeerOperation, can_run_seer_operation +from sentry.seer.credits import SeerOperation, can_run_seer_operation from sentry.types.ratelimit import RateLimit, RateLimitCategory class SeerCreditCheckSerializer(serializers.Serializer[dict[str, str]]): operation = serializers.ChoiceField(choices=[operation.value for operation in SeerOperation]) + billing_feature = serializers.ChoiceField( + choices=[ + value + for value in SeerBillingFeature.values() + if value != SEER_BILLING_FEATURE_UNSPECIFIED + ], + required=False, + ) @cell_silo_endpoint @@ -46,12 +60,14 @@ def get(self, request: Request, organization: Organization) -> Response: return Response(serializer.errors, status=400) operation = SeerOperation(serializer.validated_data["operation"]) - # TODO: Let the UI pass the feature once it can tell agentic triage runs apart. - feature = ( - SeerFeature.EXPLORER - if operation == SeerOperation.AGENT_CHAT_MESSAGE - else SeerFeature.AUTOFIX - ) + + if "billing_feature" in serializer.validated_data: + feature = serializer.validated_data["billing_feature"] + elif operation == SeerOperation.AGENT_CHAT_MESSAGE: + feature = SEER_BILLING_FEATURE_SEER_AGENT + else: + feature = SEER_BILLING_FEATURE_AUTOFIX_WEB + return Response( {"hasEnoughCredits": can_run_seer_operation(organization, feature, operation)} ) diff --git a/src/sentry/seer/entrypoints/operator.py b/src/sentry/seer/entrypoints/operator.py index 936841268b69..3a50a3fc0329 100644 --- a/src/sentry/seer/entrypoints/operator.py +++ b/src/sentry/seer/entrypoints/operator.py @@ -3,6 +3,7 @@ from sentry_protos.billing.v1.common.v1.seer_billing_feature_pb2 import ( SEER_BILLING_FEATURE_SEER_AGENT_SLACK, + SEER_BILLING_FEATURE_UNSPECIFIED, ) from taskbroker_client.retry import Retry @@ -213,6 +214,7 @@ def trigger_autofix_agent( trigger_push_changes, ) from sentry.seer.autofix.exceptions import NoSeerQuotaException, SeerOutOfCreditsException + from sentry.seer.autofix.feature.dispatch import get_run_billing_feature from sentry.seer.autofix.feature.models import FEATURE_ID as AUTOFIX_FEATURE_ID from sentry.seer.autofix.steps import AutofixStep @@ -305,6 +307,11 @@ def trigger_autofix_agent( step=AutofixStep.from_autofix_stopping_point(stopping_point), referrer=AutofixReferrer.SLACK, run_id=run_id, + billing_feature=( + get_run_billing_feature(existing_state) + if existing_state and existing_state.run_id == run_id + else SEER_BILLING_FEATURE_UNSPECIFIED + ), user=user, ) except NoSeerQuotaException: diff --git a/src/sentry/tasks/seer/pr_iteration.py b/src/sentry/tasks/seer/pr_iteration.py index c95fce7757b6..718839c8ee17 100644 --- a/src/sentry/tasks/seer/pr_iteration.py +++ b/src/sentry/tasks/seer/pr_iteration.py @@ -53,6 +53,7 @@ from sentry.seer.autofix.commit_author import commit_author_for_feedback from sentry.seer.autofix.constants import AutofixReferrer from sentry.seer.autofix.exceptions import SeerOutOfCreditsException +from sentry.seer.autofix.feature.dispatch import get_run_billing_feature from sentry.seer.autofix.pr_iteration.bot_identity import bot_logins_for_feedback from sentry.seer.autofix.pr_iteration.cap_exhausted import assign_user_for_exhausted_cap from sentry.seer.autofix.pr_iteration.constants import PR_ITERATION_PROVIDER @@ -835,6 +836,7 @@ def _drain_queued_autofix_feedback( step=AutofixStep.PR_ITERATION, referrer=referrer, run_id=run_id, + billing_feature=get_run_billing_feature(state), user_context="\n\n".join(item.text for item in feedback_items), feedback=feedback_items, actor_user_id=actor_user_id, diff --git a/tests/sentry/seer/autofix/feature/test_dispatch.py b/tests/sentry/seer/autofix/feature/test_dispatch.py index c3838d4a4715..870dfbaba48b 100644 --- a/tests/sentry/seer/autofix/feature/test_dispatch.py +++ b/tests/sentry/seer/autofix/feature/test_dispatch.py @@ -4,6 +4,7 @@ import pytest from sentry_protos.billing.v1.common.v1.seer_billing_feature_pb2 import ( SEER_BILLING_FEATURE_AUTOFIX_AGENTIC_TRIAGE, + SEER_BILLING_FEATURE_UNSPECIFIED, ) from sentry_protos.billing.v1.services.seer_feature_billing.v1.endpoint_check_seer_credits_pb2 import ( CheckSeerCreditsResponse, @@ -375,6 +376,10 @@ def test_raises_when_out_of_credits_for_a_continued_step( self, mock_check_seer_credits: MagicMock ) -> None: """Unlike quota, credits are checked for every step, not only new runs.""" + fake_run = self.create_seer_run(organization=self.organization, seer_run_state_id=123) + self.create_seer_agent_run( + run=fake_run, source=FEATURE_ID, group=self.group, project=self.group.project + ) with ( patch("sentry.seer.autofix.feature.dispatch.SeerAgentClient") as MockClient, patch("sentry.seer.autofix.feature.dispatch.quotas"), @@ -392,6 +397,107 @@ def test_raises_when_out_of_credits_for_a_continued_step( MockClient.return_value.continue_feature_run.assert_not_called() + @with_feature(["organizations:seer-credit-checks", "organizations:token-billed-seer-enabled"]) + @patch( + "sentry.seer_credit.backend.check_seer_credits", + return_value=CheckSeerCreditsResponse(has_enough_credits=True), + ) + def test_starts_agentic_triage_run_under_its_billing_feature( + self, mock_check_seer_credits: MagicMock + ) -> None: + fake_run = self.create_seer_run(organization=self.organization, type="feature_run") + + with ( + patch("sentry.seer.autofix.feature.dispatch.SeerAgentClient") as MockClient, + patch("sentry.seer.autofix.feature.dispatch.quotas"), + ): + MockClient.return_value.start_feature_run.return_value = fake_run + + run = trigger_autofix_feature( + self.group, + AutofixFeatureArgs( + referrer=AutofixReferrer.AGENTIC_TRIAGE, + step=AutofixStep.ROOT_CAUSE, + step_args=RCAStepArgs(), + billing_feature=SEER_BILLING_FEATURE_AUTOFIX_AGENTIC_TRIAGE, + ), + ) + + assert run is fake_run + assert mock_check_seer_credits.call_args.args[0].feature_id == "autofix_agentic_triage" + start_kwargs = MockClient.return_value.start_feature_run.call_args.kwargs + assert start_kwargs["billing_feature"] == SEER_BILLING_FEATURE_AUTOFIX_AGENTIC_TRIAGE + + @with_feature(["organizations:seer-credit-checks", "organizations:token-billed-seer-enabled"]) + @patch( + "sentry.seer_credit.backend.check_seer_credits", + return_value=CheckSeerCreditsResponse(has_enough_credits=True), + ) + def test_continues_agentic_triage_run_under_its_billing_feature( + self, mock_check_seer_credits: MagicMock + ) -> None: + fake_run = self.create_seer_run(organization=self.organization, seer_run_state_id=123) + self.create_seer_agent_run( + run=fake_run, + source=FEATURE_ID, + group=self.group, + project=self.group.project, + ) + + with ( + patch("sentry.seer.autofix.feature.dispatch.SeerAgentClient") as MockClient, + patch("sentry.seer.autofix.feature.dispatch.quotas"), + ): + MockClient.return_value.continue_feature_run.return_value = fake_run + + run = trigger_autofix_feature( + self.group, + AutofixFeatureArgs( + referrer=AutofixReferrer.GROUP_AUTOFIX_ENDPOINT, + step=AutofixStep.SOLUTION, + existing_run_id=123, + step_args=SolutionStepArgs(), + billing_feature=SEER_BILLING_FEATURE_AUTOFIX_AGENTIC_TRIAGE, + ), + ) + + assert run is fake_run + assert mock_check_seer_credits.call_args.args[0].feature_id == "autofix_agentic_triage" + + @with_feature(["organizations:seer-credit-checks", "organizations:token-billed-seer-enabled"]) + @patch( + "sentry.seer_credit.backend.check_seer_credits", + return_value=CheckSeerCreditsResponse(has_enough_credits=True), + ) + def test_continues_run_without_billing_feature_as_unattributed( + self, mock_check_seer_credits: MagicMock + ) -> None: + fake_run = self.create_seer_run(organization=self.organization, seer_run_state_id=123) + self.create_seer_agent_run( + run=fake_run, + source=FEATURE_ID, + group=self.group, + project=self.group.project, + ) + + with ( + patch("sentry.seer.autofix.feature.dispatch.SeerAgentClient") as MockClient, + patch("sentry.seer.autofix.feature.dispatch.quotas"), + ): + MockClient.return_value.continue_feature_run.return_value = fake_run + + trigger_autofix_feature( + self.group, + AutofixFeatureArgs( + referrer=AutofixReferrer.SLACK, + step=AutofixStep.SOLUTION, + existing_run_id=123, + step_args=SolutionStepArgs(), + ), + ) + + assert mock_check_seer_credits.call_args.args[0].feature_id == "unattributed" + def test_allows_async_dispatch(self) -> None: fake_run = self.create_seer_run(organization=self.organization, type="feature_run") @@ -511,6 +617,10 @@ def test_pr_iteration_enables_coding_and_sends_its_step_args(self) -> None: def test_code_changes_rejected_when_coding_disabled(self) -> None: self.organization.update_option("sentry:enable_seer_coding", False) + fake_run = self.create_seer_run(organization=self.organization, seer_run_state_id=123) + self.create_seer_agent_run( + run=fake_run, source=FEATURE_ID, group=self.group, project=self.group.project + ) with pytest.raises( SeerPermissionError, match="Seer coding is not enabled for this organization" @@ -674,7 +784,12 @@ def setUp(self) -> None: self.project.update_option("sentry:seer_automated_run_stopping_point", "root_cause") self.organization.update_option("sentry:enable_seer_coding", True) - def _reply(self, *, can_publish: bool = True) -> MagicMock: + def _reply( + self, + *, + can_publish: bool = True, + billing_feature: int = SEER_BILLING_FEATURE_UNSPECIFIED, + ) -> MagicMock: with patch("sentry.seer.autofix.feature.dispatch.SeerAgentClient") as mock_client_cls: mock_client_cls.return_value.continue_feature_run.return_value = self.seer_run continue_single_run( @@ -684,6 +799,7 @@ def _reply(self, *, can_publish: bool = True) -> MagicMock: user=self.user, can_publish=can_publish, referrer=AutofixReferrer.SEER_EXPLORER, + billing_feature=billing_feature, ) return mock_client_cls @@ -739,6 +855,19 @@ def test_raises_when_out_of_credits(self, mock_check_seer_credits: MagicMock) -> user=self.user, can_publish=True, referrer=AutofixReferrer.SEER_EXPLORER, + billing_feature=SEER_BILLING_FEATURE_UNSPECIFIED, ) mock_client_cls.return_value.continue_feature_run.assert_not_called() + + @with_feature(["organizations:seer-credit-checks", "organizations:token-billed-seer-enabled"]) + @patch( + "sentry.seer_credit.backend.check_seer_credits", + return_value=CheckSeerCreditsResponse(has_enough_credits=True), + ) + def test_checks_credits_under_the_run_billing_feature( + self, mock_check_seer_credits: MagicMock + ) -> None: + self._reply(billing_feature=SEER_BILLING_FEATURE_AUTOFIX_AGENTIC_TRIAGE) + + assert mock_check_seer_credits.call_args.args[0].feature_id == "autofix_agentic_triage" diff --git a/tests/sentry/seer/autofix/test_autofix_agent.py b/tests/sentry/seer/autofix/test_autofix_agent.py index 555ef5924e20..4d301fe51661 100644 --- a/tests/sentry/seer/autofix/test_autofix_agent.py +++ b/tests/sentry/seer/autofix/test_autofix_agent.py @@ -4,6 +4,9 @@ import pytest from django.test import override_settings from rest_framework.exceptions import PermissionDenied +from sentry_protos.billing.v1.common.v1.seer_billing_feature_pb2 import ( + SEER_BILLING_FEATURE_AUTOFIX_AGENTIC_TRIAGE, +) from sentry_protos.billing.v1.services.seer_feature_billing.v1.endpoint_check_seer_credits_pb2 import ( CheckSeerCreditsResponse, ) @@ -1618,6 +1621,7 @@ def _push( self, mock_post: MagicMock, features: dict[str, bool] | None = None, + billing_feature: int | None = None, **kwargs: Any, ) -> dict[str, Any]: """Push with a minimal run state and return the payload sent to Seer.""" @@ -1630,6 +1634,7 @@ def _push( updated_at="2024-01-01T00:00:00Z", repo_pr_states={}, metadata={"group_id": self.group.id}, + billing_feature=billing_feature, ) with self.feature(features or {}): @@ -1677,14 +1682,30 @@ def test_raises_when_out_of_credits( self, mock_post: MagicMock, mock_check_seer_credits: MagicMock ) -> None: with pytest.raises(SeerOutOfCreditsException): - trigger_push_changes( - group=self.group, - run_id=123, - referrer=AutofixReferrer.UNKNOWN, - ) + self._push(mock_post) mock_post.assert_not_called() + @patch( + "sentry.seer_credit.backend.check_seer_credits", + return_value=CheckSeerCreditsResponse(has_enough_credits=True), + ) + @patch("sentry.seer.agent.client.make_agent_update_request") + def test_checks_credits_under_the_run_billing_feature( + self, mock_post: MagicMock, mock_check_seer_credits: MagicMock + ) -> None: + self._push( + mock_post, + features={ + "organizations:seer-credit-checks": True, + "organizations:token-billed-seer-enabled": True, + }, + billing_feature=SEER_BILLING_FEATURE_AUTOFIX_AGENTIC_TRIAGE, + ) + + assert mock_check_seer_credits.call_args.args[0].feature_id == "autofix_agentic_triage" + mock_post.assert_called_once() + @patch("sentry.seer.agent.client.make_agent_update_request") def test_passes_pr_description_suffix_ready_state_and_author( self, mock_post: MagicMock diff --git a/tests/sentry/seer/endpoints/test_admin_autofix_retry.py b/tests/sentry/seer/endpoints/test_admin_autofix_retry.py index e95ecb37cb35..d68d075f916d 100644 --- a/tests/sentry/seer/endpoints/test_admin_autofix_retry.py +++ b/tests/sentry/seer/endpoints/test_admin_autofix_retry.py @@ -2,6 +2,9 @@ from unittest.mock import MagicMock, patch from rest_framework.response import Response +from sentry_protos.billing.v1.common.v1.seer_billing_feature_pb2 import ( + SEER_BILLING_FEATURE_UNSPECIFIED, +) from sentry.models.organization import Organization from sentry.seer.agent.client_models import MemoryBlock, Message, RepoPRState, SeerRunState @@ -93,6 +96,7 @@ def test_retries_failed_step_as_system( step=AutofixStep.SOLUTION, referrer=AutofixReferrer.NIGHT_SHIFT, run_id=1, + billing_feature=SEER_BILLING_FEATURE_UNSPECIFIED, insert_index=2, ) assert viewer_contexts == [ diff --git a/tests/sentry/seer/endpoints/test_group_ai_autofix.py b/tests/sentry/seer/endpoints/test_group_ai_autofix.py index 91af32d5f500..0476d0137405 100644 --- a/tests/sentry/seer/endpoints/test_group_ai_autofix.py +++ b/tests/sentry/seer/endpoints/test_group_ai_autofix.py @@ -7,8 +7,13 @@ from django.test import override_settings from django.utils import timezone from sentry_protos.billing.v1.common.v1.seer_billing_feature_pb2 import ( + SEER_BILLING_FEATURE_AUTOFIX_AGENTIC_TRIAGE, SEER_BILLING_FEATURE_AUTOFIX_HEADLESS, SEER_BILLING_FEATURE_AUTOFIX_WEB, + SEER_BILLING_FEATURE_UNSPECIFIED, +) +from sentry_protos.billing.v1.services.seer_feature_billing.v1.endpoint_check_seer_credits_pb2 import ( + CheckSeerCreditsResponse, ) from sentry.integrations.services.integration import RpcIntegration @@ -42,7 +47,7 @@ from sentry.seer.autofix.pr_iteration.queue import QueuedAutofixFeedback from sentry.seer.autofix.steps import AutofixStep from sentry.seer.autofix.utils import AutofixStoppingPoint -from sentry.seer.credits import SeerFeature, SeerOperation +from sentry.seer.credits import SeerOperation from sentry.seer.models import SeerPermissionError from sentry.seer.models.run import SeerRun from sentry.testutils.cases import APITestCase, SnubaTestCase @@ -1264,7 +1269,7 @@ def test_insert_index_passed_through( group=group, step=AutofixStep.SOLUTION, referrer=AutofixReferrer.GROUP_AUTOFIX_ENDPOINT, - billing_feature=SEER_BILLING_FEATURE_AUTOFIX_HEADLESS, + billing_feature=SEER_BILLING_FEATURE_UNSPECIFIED, stopping_point=None, run_id=42, user_context=None, @@ -1706,10 +1711,50 @@ def test_pr_iteration_rejected_when_out_of_credits( "code": "seer_out_of_credits", } mock_can_run.assert_called_once_with( - group.organization, SeerFeature.AUTOFIX, SeerOperation.AUTOFIX_PR_ITERATION + group.organization, SEER_BILLING_FEATURE_UNSPECIFIED, SeerOperation.AUTOFIX_PR_ITERATION ) mock_enqueue.assert_not_called() + @with_feature(["organizations:seer-credit-checks", "organizations:token-billed-seer-enabled"]) + @patch( + "sentry.seer_credit.backend.check_seer_credits", + return_value=CheckSeerCreditsResponse(has_enough_credits=True), + ) + @patch("sentry.seer.autofix.pr_iteration.manual_trigger.trigger_consume_pr_iteration_feedback") + @patch("sentry.seer.autofix.pr_iteration.manual_trigger.enqueue_autofix_feedback") + @patch("sentry.seer.autofix.pr_iteration.manual_trigger.get_autofix_run_state") + def test_pr_iteration_checks_credits_under_the_run_billing_feature( + self, + mock_run_state: MagicMock, + mock_enqueue: MagicMock, + mock_consume: MagicMock, + mock_check_seer_credits: MagicMock, + ) -> None: + """Billing decides whether the run's feature is charged.""" + group = self.create_group() + run = self.create_seer_run(organization=self.organization, seer_run_state_id=123) + self.create_seer_agent_run(run=run, source="autofix", group=group, project=group.project) + mock_run_state.return_value = SeerRunState( + run_id=123, + blocks=[], + status="completed", + updated_at="2024-01-01T00:00:00Z", + repo_pr_states={"owner/repo": RepoPRState(repo_name="owner/repo", pr_number=7)}, + billing_feature=SEER_BILLING_FEATURE_AUTOFIX_AGENTIC_TRIAGE, + ) + + self.login_as(user=self.user) + response = self.client.post( + self._get_url(group.id), + data={"step": "pr_iteration", "run_id": 123, "user_context": "please fix this"}, + format="json", + ) + + assert response.status_code == 202, response.data + assert mock_check_seer_credits.call_args.args[0].feature_id == "autofix_agentic_triage" + mock_enqueue.assert_called_once() + mock_consume.assert_called_once() + @patch("sentry.seer.autofix.pr_iteration.manual_trigger.trigger_consume_pr_iteration_feedback") @patch("sentry.seer.autofix.pr_iteration.manual_trigger.enqueue_autofix_feedback") @patch("sentry.seer.autofix.pr_iteration.manual_trigger.get_autofix_run_state") diff --git a/tests/sentry/seer/endpoints/test_organization_seer_credit_check.py b/tests/sentry/seer/endpoints/test_organization_seer_credit_check.py index 92246f1f6a68..64c91421924a 100644 --- a/tests/sentry/seer/endpoints/test_organization_seer_credit_check.py +++ b/tests/sentry/seer/endpoints/test_organization_seer_credit_check.py @@ -33,7 +33,7 @@ def test_get_returns_has_enough_credits_when_billing_has_credits( mock_check_seer_credits.assert_called_once_with( CheckSeerCreditsRequest( organization_id=self.organization.id, - feature_id="autofix", + feature_id="autofix_web", predicted_credit_usage=SEER_OPERATION_ESTIMATES[SeerOperation.AUTOFIX_ROOT_CAUSE](), take_lease=False, ) @@ -50,7 +50,7 @@ def test_get_returns_not_enough_credits_when_billing_has_no_credits( response = self.get_success_response(self.organization.slug, operation="agent_chat_message") assert response.data == {"hasEnoughCredits": False} - assert mock_check_seer_credits.call_args.args[0].feature_id == "explorer" + assert mock_check_seer_credits.call_args.args[0].feature_id == "seer_agent" def test_get_returns_has_enough_credits_when_credit_checks_disabled(self) -> None: response = self.get_success_response(self.organization.slug, operation="autofix_solution") diff --git a/tests/sentry/seer/test_credits.py b/tests/sentry/seer/test_credits.py index ab40c1306bec..e3ec9e7ec8fb 100644 --- a/tests/sentry/seer/test_credits.py +++ b/tests/sentry/seer/test_credits.py @@ -2,17 +2,16 @@ from unittest.mock import MagicMock, patch +from sentry_protos.billing.v1.common.v1.seer_billing_feature_pb2 import ( + SEER_BILLING_FEATURE_AUTOFIX_WEB, + SEER_BILLING_FEATURE_SEER_AGENT, +) from sentry_protos.billing.v1.services.seer_feature_billing.v1.endpoint_check_seer_credits_pb2 import ( CheckSeerCreditsRequest, CheckSeerCreditsResponse, ) -from sentry.seer.credits import ( - SEER_OPERATION_ESTIMATES, - SeerFeature, - SeerOperation, - can_run_seer_operation, -) +from sentry.seer.credits import SEER_OPERATION_ESTIMATES, SeerOperation, can_run_seer_operation from sentry.testutils.cases import TestCase from sentry.testutils.helpers.features import with_feature @@ -31,7 +30,7 @@ def test_can_run_seer_operation_skips_billing_when_credit_checks_disabled( self, mock_check_seer_credits: MagicMock ) -> None: assert can_run_seer_operation( - self.organization, SeerFeature.AUTOFIX, SeerOperation.AUTOFIX_ROOT_CAUSE + self.organization, SEER_BILLING_FEATURE_AUTOFIX_WEB, SeerOperation.AUTOFIX_ROOT_CAUSE ) mock_check_seer_credits.assert_not_called() @@ -44,7 +43,7 @@ def test_can_run_seer_operation_skips_billing_when_not_token_billed( self, mock_check_seer_credits: MagicMock ) -> None: assert can_run_seer_operation( - self.organization, SeerFeature.AUTOFIX, SeerOperation.AUTOFIX_ROOT_CAUSE + self.organization, SEER_BILLING_FEATURE_AUTOFIX_WEB, SeerOperation.AUTOFIX_ROOT_CAUSE ) mock_check_seer_credits.assert_not_called() @@ -57,12 +56,12 @@ def test_can_run_seer_operation_returns_false_when_billing_has_no_credits( self, mock_check_seer_credits: MagicMock ) -> None: assert not can_run_seer_operation( - self.organization, SeerFeature.AUTOFIX, SeerOperation.AUTOFIX_ROOT_CAUSE + self.organization, SEER_BILLING_FEATURE_AUTOFIX_WEB, SeerOperation.AUTOFIX_ROOT_CAUSE ) mock_check_seer_credits.assert_called_once_with( CheckSeerCreditsRequest( organization_id=self.organization.id, - feature_id="autofix", + feature_id="autofix_web", predicted_credit_usage=SEER_OPERATION_ESTIMATES[SeerOperation.AUTOFIX_ROOT_CAUSE](), take_lease=False, ) @@ -77,5 +76,5 @@ def test_can_run_seer_operation_returns_true_when_billing_has_credits( self, mock_check_seer_credits: MagicMock ) -> None: assert can_run_seer_operation( - self.organization, SeerFeature.EXPLORER, SeerOperation.AGENT_CHAT_MESSAGE + self.organization, SEER_BILLING_FEATURE_SEER_AGENT, SeerOperation.AGENT_CHAT_MESSAGE )