From 35a45b3c4b222f5e8fd756d649f85f9e54c11d7c Mon Sep 17 00:00:00 2001 From: Zach Collins Date: Mon, 4 Mar 2024 14:25:31 -0800 Subject: [PATCH] feat(autofix): timeout with autofix, query state (#66241) --- src/sentry/api/endpoints/group_ai_autofix.py | 12 ++- src/sentry/api/endpoints/seer_rpc.py | 38 +++++++++- .../api/endpoints/test_group_ai_autofix.py | 3 +- tests/sentry/api/endpoints/test_seer_rpc.py | 73 +++++++++++++++++++ 4 files changed, 122 insertions(+), 4 deletions(-) create mode 100644 tests/sentry/api/endpoints/test_seer_rpc.py diff --git a/src/sentry/api/endpoints/group_ai_autofix.py b/src/sentry/api/endpoints/group_ai_autofix.py index e53edea874e4b5..a9db0ac70e500f 100644 --- a/src/sentry/api/endpoints/group_ai_autofix.py +++ b/src/sentry/api/endpoints/group_ai_autofix.py @@ -112,6 +112,7 @@ def _call_autofix( repos: list[dict], event_entries: list[dict], additional_context: str, + timeout_secs: int, ): response = requests.post( f"{settings.SEER_AUTOFIX_URL}/v0/automation/autofix", @@ -123,10 +124,12 @@ def _call_autofix( "issue": { "id": group.id, "title": group.title, - "short_id": group.short_id, + "short_id": group.qualified_short_id, "events": [{"entries": event_entries}], }, "additional_context": additional_context, + "timeout_secs": timeout_secs, + "last_updated": datetime.now().isoformat(), "invoking_user": ( { "id": user.id, @@ -192,7 +195,12 @@ def post(self, request: Request, group: Group) -> Response: try: self._call_autofix( - request.user, group, repos, event_entries, data.get("additional_context", "") + request.user, + group, + repos, + event_entries, + data.get("additional_context", ""), + TIMEOUT_SECONDS, ) # Mark the task as completed after TIMEOUT_SECONDS diff --git a/src/sentry/api/endpoints/seer_rpc.py b/src/sentry/api/endpoints/seer_rpc.py index c2c2537721b15b..cde7f37b58a68e 100644 --- a/src/sentry/api/endpoints/seer_rpc.py +++ b/src/sentry/api/endpoints/seer_rpc.py @@ -5,6 +5,7 @@ from django.conf import settings from django.contrib.auth.models import AnonymousUser +from django.core.exceptions import ObjectDoesNotExist from rest_framework.exceptions import ( AuthenticationFailed, NotFound, @@ -25,6 +26,7 @@ from sentry.services.hybrid_cloud.sig import SerializableFunctionValueException from sentry.silo.base import SiloMode from sentry.utils import json +from sentry.utils.env import in_test_environment def compare_signature(url: str, body: bytes, signature: str) -> bool: @@ -131,8 +133,13 @@ def post(self, request: Request, method_name: str) -> Response: except SerializableFunctionValueException as e: capture_exception() raise ParseError from e + except ObjectDoesNotExist as e: + # Let this fall through, this is normal. + capture_exception() + raise NotFound from e except Exception as e: - # Produce more detailed log + if in_test_environment(): + raise if settings.DEBUG: raise Exception(f"Problem processing seer rpc endpoint {method_name}") from e capture_exception() @@ -174,7 +181,36 @@ def on_autofix_complete(*, issue_id: int, status: str, steps: list[dict], fix: d group.save() +def get_autofix_state(*, issue_id: int) -> dict: + group: Group = Group.objects.get(id=issue_id) + + metadata = group.data.get("metadata", {}) + autofix_data = metadata.get("autofix", {}) + + return autofix_data + + seer_method_registry = { "on_autofix_step_update": on_autofix_step_update, "on_autofix_complete": on_autofix_complete, + "get_autofix_state": get_autofix_state, } + + +def generate_request_signature(url_path: str, body: bytes) -> str: + """ + Generate a signature for the request body + with the first shared secret. If there are other + shared secrets in the list they are only to be used + for verfication during key rotation. + """ + if not settings.SEER_RPC_SHARED_SECRET: + raise RpcAuthenticationSetupException("Cannot sign RPC requests without RPC_SHARED_SECRET") + + signature_input = b"%s:%s" % ( + url_path.encode("utf8"), + body, + ) + secret = settings.SEER_RPC_SHARED_SECRET[0] + signature = hmac.new(secret.encode("utf-8"), signature_input, hashlib.sha256).hexdigest() + return f"rpc0:{signature}" diff --git a/tests/sentry/api/endpoints/test_group_ai_autofix.py b/tests/sentry/api/endpoints/test_group_ai_autofix.py index 6cbdf151940c7d..ca19273fab54fb 100644 --- a/tests/sentry/api/endpoints/test_group_ai_autofix.py +++ b/tests/sentry/api/endpoints/test_group_ai_autofix.py @@ -1,6 +1,6 @@ from unittest.mock import ANY, patch -from sentry.api.endpoints.group_ai_autofix import GroupAiAutofixEndpoint +from sentry.api.endpoints.group_ai_autofix import TIMEOUT_SECONDS, GroupAiAutofixEndpoint from sentry.models.group import Group from sentry.testutils.cases import APITestCase, SnubaTestCase from sentry.testutils.helpers.datetime import before_now @@ -93,6 +93,7 @@ def test_ai_autofix_post_endpoint(self): ], ANY, "Yes", + TIMEOUT_SECONDS, ) actual_group_arg = mock_call.call_args[0][1] diff --git a/tests/sentry/api/endpoints/test_seer_rpc.py b/tests/sentry/api/endpoints/test_seer_rpc.py new file mode 100644 index 00000000000000..a1bf808fa4c62e --- /dev/null +++ b/tests/sentry/api/endpoints/test_seer_rpc.py @@ -0,0 +1,73 @@ +from typing import Any + +from django.test import override_settings +from django.urls import reverse + +from sentry.api.endpoints.seer_rpc import generate_request_signature +from sentry.testutils.cases import APITestCase +from sentry.utils import json + + +@override_settings(SEER_RPC_SHARED_SECRET=["a-long-value-that-is-hard-to-guess"]) +class TestSeerRpc(APITestCase): + @staticmethod + def _get_path(method_name: str) -> str: + return reverse( + "sentry-api-0-seer-rpc-service", + kwargs={"method_name": method_name}, + ) + + def auth_header(self, path: str, data: dict | str) -> str: + if isinstance(data, dict): + data = json.dumps(data) + signature = generate_request_signature(path, data.encode("utf8")) + + return f"rpcsignature {signature}" + + def test_invalid_endpoint(self): + path = self._get_path("not_a_method") + response = self.client.post(path) + assert response.status_code == 403 + + def test_invalid_authentication(self): + path = self._get_path("on_autofix_step_update") + data: dict[str, Any] = {"args": {"issued_id": 1, "status": "", "steps": []}, "meta": {}} + response = self.client.post(path, data=data, HTTP_AUTHORIZATION="rpcsignature trash") + assert response.status_code == 401 + + def test_404(self): + path = self._get_path("get_autofix_state") + data: dict[str, Any] = {"args": {"issue_id": 1}, "meta": {}} + response = self.client.post( + path, data=data, HTTP_AUTHORIZATION=self.auth_header(path, data) + ) + assert response.status_code == 404 + + def test_step_state_management(self): + group = self.create_group() + + path = self._get_path("get_autofix_state") + data: dict[str, Any] = {"args": {"issue_id": group.id}, "meta": {}} + response = self.client.post( + path, data=data, HTTP_AUTHORIZATION=self.auth_header(path, data) + ) + assert response.status_code == 200 + assert response.json() == {} + + path = self._get_path("on_autofix_step_update") + data = { + "args": {"issue_id": group.id, "status": "thing", "steps": [1, 2, 3]}, + "meta": {}, + } + response = self.client.post( + path, data=data, HTTP_AUTHORIZATION=self.auth_header(path, data) + ) + assert response.status_code == 200 + + path = self._get_path("get_autofix_state") + data = {"args": {"issue_id": group.id}, "meta": {}} + response = self.client.post( + path, data=data, HTTP_AUTHORIZATION=self.auth_header(path, data) + ) + assert response.status_code == 200 + assert response.json() == {"status": "thing", "steps": [1, 2, 3]}