diff --git a/src/sentry/onboarding/agentic_progress/__init__.py b/src/sentry/onboarding/agentic_progress/__init__.py index c3b83455a4f2..cb4d150e3e7b 100644 --- a/src/sentry/onboarding/agentic_progress/__init__.py +++ b/src/sentry/onboarding/agentic_progress/__init__.py @@ -1,10 +1,13 @@ from .model import OnboardingRun, ProgressUpdate, RunStatus, Stage, StageStatus, apply_update +from .service import OnboardingProgressService, get_onboarding_progress_service __all__ = [ + "OnboardingProgressService", "OnboardingRun", "ProgressUpdate", "RunStatus", "Stage", "StageStatus", "apply_update", + "get_onboarding_progress_service", ] diff --git a/src/sentry/onboarding/agentic_progress/service.py b/src/sentry/onboarding/agentic_progress/service.py new file mode 100644 index 000000000000..0246a356b90b --- /dev/null +++ b/src/sentry/onboarding/agentic_progress/service.py @@ -0,0 +1,327 @@ +from __future__ import annotations + +import hashlib +import hmac +import time +import uuid +from collections.abc import Callable +from dataclasses import replace +from datetime import datetime, timedelta, timezone +from typing import NamedTuple, cast + +from django.conf import settings +from redis.exceptions import WatchError +from sentry_redis_tools.clients import RedisCluster, StrictRedis + +from sentry.utils import json, redis + +from .model import ( + InvalidOnboardingRun, + OnboardingRun, + OnboardingRunTerminal, + ProgressUpdate, + RunStatus, + apply_update, + initial_stages, +) + +RUN_LIFETIME = timedelta(days=7) +"""Maximum lifetime of a registered onboarding run.""" + +CLIENT_CLAIM_RETRIES = 50 +"""Maximum attempts to claim a browser run when concurrent registrations race.""" + +CLIENT_CLAIM_RETRY_DELAY = 0.01 +"""Delay between client-run claim attempts while another registration initializes.""" + +ATOMIC_UPDATE_RETRIES = 50 +"""Maximum attempts to update a run when concurrent Redis transactions race.""" + +TOKEN_LENGTH = 10 +"""Length of the short handoff code included in the onboarding prompt.""" + + +class RegisteredRun(NamedTuple): + run: OnboardingRun + onboarding_code: str + + +class UpdatedRun(NamedTuple): + run: OnboardingRun + changed: bool + + +class ClientRunClaim(NamedTuple): + existing: OnboardingRun | None = None + claimed_run_id: str | None = None + + +def get_redis_client() -> RedisCluster[bytes] | StrictRedis[bytes]: + return redis.redis_clusters.get_binary("default") + + +class RunNotFound(Exception): + pass + + +class RunOwnershipMismatch(Exception): + pass + + +class OnboardingProgressService: + def __init__(self) -> None: + self.redis = get_redis_client() + self.signing_secret = settings.SECRET_KEY.encode() + + def create_or_resume( + self, + *, + user_id: int, + organization_id: int, + client_run_id: str, + onboarding_code: str, + ) -> RegisteredRun: + index_key = self._client_index_key(user_id, organization_id, client_run_id) + claim = self._claim_client_run(index_key) + if claim.existing is not None: + if hmac.compare_digest(claim.existing.token_hash, self._hash_token(onboarding_code)): + return RegisteredRun(claim.existing, onboarding_code) + raise RunOwnershipMismatch + + assert claim.claimed_run_id is not None + now = datetime.now(timezone.utc) + ttl = int(RUN_LIFETIME.total_seconds()) + token_hash = self._hash_token(onboarding_code) + token_key = self._token_index_key(token_hash) + try: + self._claim_token_index(token_key, claim.claimed_run_id, ttl) + run = OnboardingRun( + run_id=claim.claimed_run_id, + channel_id=uuid.uuid4().hex, + token_hash=token_hash, + client_run_id=client_run_id, + user_id=user_id, + organization_id=organization_id, + created_at=now, + updated_at=now, + expires_at=now + RUN_LIFETIME, + sequence=0, + stages=initial_stages(), + ) + if not self.redis.set( + self._state_key(claim.claimed_run_id), self._serialize(run), ex=ttl + ): + raise RuntimeError("Unable to store onboarding run") + except Exception: + self._release_index(token_key, claim.claimed_run_id) + self._release_index(index_key, claim.claimed_run_id) + raise + + return RegisteredRun(run, onboarding_code) + + def get(self, *, run_id: str, user_id: int, organization_id: int) -> OnboardingRun | None: + run = self._load(self._state_key(run_id)) + if run is None or datetime.now(timezone.utc) >= run.expires_at: + return None + if run.user_id != user_id or run.organization_id != organization_id: + return None + return run + + def update( + self, + *, + token: str, + user_id: int, + organization_id: int, + update: ProgressUpdate, + ) -> UpdatedRun: + if len(token) != TOKEN_LENGTH or not token.isalnum(): + raise RunNotFound + token_hash = self._hash_token(token) + run_id = self._decode(self.redis.get(self._token_index_key(token_hash))) + if run_id is None: + raise RunNotFound + + def mutate(run: OnboardingRun) -> OnboardingRun: + if run.token_hash != token_hash: + raise RunNotFound + if run.user_id != user_id or run.organization_id != organization_id: + raise RunOwnershipMismatch + return apply_update(run, update, datetime.now(timezone.utc)) + + return self._atomic_update(run_id, mutate) + + def cancel(self, *, run_id: str, user_id: int, organization_id: int) -> OnboardingRun: + def mutate(run: OnboardingRun) -> OnboardingRun: + if run.user_id != user_id or run.organization_id != organization_id: + raise RunNotFound + if run.run_status is RunStatus.CANCELLED: + return run + if run.run_status is not RunStatus.ACTIVE: + raise OnboardingRunTerminal("Onboarding run is terminal") + now = datetime.now(timezone.utc) + return replace( + run, + run_status=RunStatus.CANCELLED, + updated_at=now, + sequence=run.sequence + 1, + ) + + result = self._atomic_update(run_id, mutate) + return result.run + + def _claim_client_run(self, index_key: str) -> ClientRunClaim: + """Return the active indexed run or atomically claim the client id for a new run.""" + for attempt in range(CLIENT_CLAIM_RETRIES): + try: + with self.redis.pipeline() as pipeline: + pipeline.watch(index_key) + indexed_run_id = self._decode(cast(bytes | str | None, pipeline.get(index_key))) + if indexed_run_id is not None: + existing = self._load(self._state_key(indexed_run_id)) + if existing is None and attempt < CLIENT_CLAIM_RETRIES - 1: + time.sleep(CLIENT_CLAIM_RETRY_DELAY) + continue + if ( + existing is not None + and existing.run_status is RunStatus.ACTIVE + and datetime.now(timezone.utc) < existing.expires_at + ): + return ClientRunClaim(existing=existing) + claimed_run_id = uuid.uuid4().hex + pipeline.multi() + pipeline.set( + index_key, + claimed_run_id, + ex=int(RUN_LIFETIME.total_seconds()), + ) + pipeline.execute() + return ClientRunClaim(claimed_run_id=claimed_run_id) + except WatchError: + continue + + raise RuntimeError("Unable to claim an onboarding client run") + + def _claim_token_index(self, key: str, run_id: str, ttl: int) -> None: + for attempt in range(CLIENT_CLAIM_RETRIES): + try: + with self.redis.pipeline() as pipeline: + pipeline.watch(key) + indexed_run_id = self._decode(cast(bytes | str | None, pipeline.get(key))) + if indexed_run_id is not None: + existing = self._load(self._state_key(indexed_run_id)) + if existing is None and attempt < CLIENT_CLAIM_RETRIES - 1: + time.sleep(CLIENT_CLAIM_RETRY_DELAY) + continue + if ( + existing is not None + and existing.run_status is RunStatus.ACTIVE + and datetime.now(timezone.utc) < existing.expires_at + ): + raise ValueError("Onboarding code is already in use") + + pipeline.multi() + pipeline.set(key, run_id, ex=ttl) + pipeline.execute() + return + except WatchError: + continue + + raise RuntimeError("Unable to claim onboarding code") + + def _release_index(self, key: str, run_id: str) -> None: + for _ in range(CLIENT_CLAIM_RETRIES): + try: + with self.redis.pipeline() as pipeline: + pipeline.watch(key) + indexed_run_id = self._decode(cast(bytes | str | None, pipeline.get(key))) + if indexed_run_id != run_id: + return + + pipeline.multi() + pipeline.delete(key) + pipeline.execute() + return + except WatchError: + continue + + raise RuntimeError("Unable to release onboarding index") + + def _atomic_update( + self, run_id: str, mutate: Callable[[OnboardingRun], OnboardingRun] + ) -> UpdatedRun: + key = self._state_key(run_id) + for _ in range(ATOMIC_UPDATE_RETRIES): + try: + with self.redis.pipeline() as pipeline: + pipeline.watch(key) + raw = cast(bytes | str | None, pipeline.get(key)) + if raw is None: + raise RunNotFound + current = self._deserialize(raw) + updated = mutate(current) + if updated == current: + return UpdatedRun(current, False) + + pipeline.multi() + pipeline.set(key, self._serialize(updated), ex=self._remaining_ttl(updated)) + pipeline.execute() + return UpdatedRun(updated, True) + except WatchError: + continue + + raise RuntimeError("Unable to update onboarding progress") + + def _remaining_ttl(self, run: OnboardingRun) -> int: + remaining = run.expires_at - datetime.now(timezone.utc) + if remaining.total_seconds() <= 0: + raise RunNotFound + return max(1, int(remaining.total_seconds())) + + def _hash_token(self, token: str) -> str: + # The handoff code is intentionally short enough to paste into a prompt, which + # also makes it practical to enumerate. A keyed digest avoids storing the raw + # code and prevents a Redis snapshot alone from validating offline guesses. + return hmac.new(self.signing_secret, token.encode(), hashlib.sha256).hexdigest() + + def _load(self, key: str) -> OnboardingRun | None: + raw = self.redis.get(key) + if raw is None: + return None + + try: + return self._deserialize(raw) + except RunNotFound: + return None + + @staticmethod + def _state_key(run_id: str) -> str: + return f"agentic-onboarding:run:{{{run_id}}}:state" + + @staticmethod + def _token_index_key(token_hash: str) -> str: + return f"agentic-onboarding:token:{token_hash}" + + @staticmethod + def _client_index_key(user_id: int, organization_id: int, client_run_id: str) -> str: + digest = hashlib.sha256(client_run_id.encode()).hexdigest() + return f"agentic-onboarding:client:{organization_id}:{user_id}:{digest}" + + @staticmethod + def _serialize(run: OnboardingRun) -> str: + return json.dumps(run.to_dict()) + + @staticmethod + def _deserialize(raw: bytes | str) -> OnboardingRun: + try: + return OnboardingRun.from_dict(json.loads(raw)) + except (json.JSONDecodeError, InvalidOnboardingRun, TypeError) as error: + raise RunNotFound from error + + @staticmethod + def _decode(value: bytes | str | None) -> str | None: + return value.decode() if isinstance(value, bytes) else value + + +def get_onboarding_progress_service() -> OnboardingProgressService: + return OnboardingProgressService() diff --git a/tests/sentry/onboarding/agentic_progress/test_service.py b/tests/sentry/onboarding/agentic_progress/test_service.py new file mode 100644 index 000000000000..f46bf7aa2920 --- /dev/null +++ b/tests/sentry/onboarding/agentic_progress/test_service.py @@ -0,0 +1,326 @@ +from collections.abc import Iterator +from dataclasses import replace +from datetime import datetime, timedelta +from typing import Any +from unittest.mock import patch + +import pytest +from django.utils import timezone as django_timezone +from redis.exceptions import WatchError + +from sentry.onboarding.agentic_progress.model import ( + OnboardingRunTerminal, + ProgressUpdate, + RunStatus, + Stage, + StageStatus, +) +from sentry.onboarding.agentic_progress.service import ( + ATOMIC_UPDATE_RETRIES, + CLIENT_CLAIM_RETRIES, + RUN_LIFETIME, + OnboardingProgressService, + RegisteredRun, + RunNotFound, + RunOwnershipMismatch, +) +from sentry.testutils.helpers.datetime import freeze_time + + +@pytest.fixture +def frozen_time() -> Iterator[Any]: + with freeze_time("2020-01-01T00:00:00Z") as frozen: + yield frozen + + +@pytest.fixture +def service(frozen_time: Any) -> Iterator[OnboardingProgressService]: + yield OnboardingProgressService() + + +def create_run( + service: OnboardingProgressService, + *, + user_id: int, + organization_id: int, + client_run_id: str, + onboarding_code: str = "abcdefghij", +) -> RegisteredRun: + return service.create_or_resume( + user_id=user_id, + organization_id=organization_id, + client_run_id=client_run_id, + onboarding_code=onboarding_code, + ) + + +def test_create_get_and_resume(service: OnboardingProgressService) -> None: + created, token = create_run( + service, user_id=1, organization_id=2, client_run_id="browser-session" + ) + fetched = service.get(run_id=created.run_id, user_id=1, organization_id=2) + resumed, resumed_token = create_run( + service, user_id=1, organization_id=2, client_run_id="browser-session" + ) + + assert token == "abcdefghij" + assert len(created.channel_id) == 32 + assert fetched == created + assert resumed.run_id == created.run_id + assert resumed_token == token + assert service.get(run_id=created.run_id, user_id=9, organization_id=2) is None + + +def test_resume_preserves_supplied_onboarding_code(service: OnboardingProgressService) -> None: + created, token = create_run( + service, user_id=1, organization_id=2, client_run_id="browser-session" + ) + + resumed, resumed_token = create_run( + service, + user_id=1, + organization_id=2, + client_run_id="browser-session", + onboarding_code=token, + ) + + assert resumed == created + assert resumed_token == token + + +def test_resume_rejects_a_different_onboarding_code(service: OnboardingProgressService) -> None: + create_run(service, user_id=1, organization_id=2, client_run_id="browser-session") + + with pytest.raises(RunOwnershipMismatch): + create_run( + service, + user_id=1, + organization_id=2, + client_run_id="browser-session", + onboarding_code="klmnopqrst", + ) + + +def test_update_uses_canonical_state_key(service: OnboardingProgressService) -> None: + created, token = create_run( + service, user_id=1, organization_id=2, client_run_id="browser-session" + ) + + updated, changed = service.update( + token=token, + user_id=1, + organization_id=2, + update=ProgressUpdate(stage=Stage.CREATE_PROJECT, status=StageStatus.COMPLETED), + ) + + stored = service.redis.get(f"agentic-onboarding:run:{{{created.run_id}}}:state") + assert stored is not None + assert updated.sequence == 1 + assert changed is True + + +def test_duplicate_update_is_idempotent(service: OnboardingProgressService) -> None: + _, token = create_run(service, user_id=1, organization_id=2, client_run_id="browser-session") + first, first_changed = service.update( + token=token, + user_id=1, + organization_id=2, + update=ProgressUpdate(stage=Stage.CONNECT_MCP, status=StageStatus.COMPLETED), + ) + duplicate, duplicate_changed = service.update( + token=token, + user_id=1, + organization_id=2, + update=ProgressUpdate(stage=Stage.CONNECT_MCP, status=StageStatus.COMPLETED), + ) + + assert duplicate == first + assert duplicate.sequence == 1 + assert first_changed is True + assert duplicate_changed is False + + +def test_token_and_ownership_are_validated(service: OnboardingProgressService) -> None: + _, token = create_run(service, user_id=1, organization_id=2, client_run_id="browser-session") + + with pytest.raises(RunNotFound): + service.update( + token="invalid000", + user_id=1, + organization_id=2, + update=ProgressUpdate(stage=Stage.CONNECT_MCP, status=StageStatus.COMPLETED), + ) + with pytest.raises(RunOwnershipMismatch): + service.update( + token=token, + user_id=9, + organization_id=2, + update=ProgressUpdate(stage=Stage.CONNECT_MCP, status=StageStatus.COMPLETED), + ) + + +def test_cancel_is_terminal_and_idempotent(service: OnboardingProgressService) -> None: + created, _ = create_run(service, user_id=1, organization_id=2, client_run_id="browser-session") + + cancelled = service.cancel(run_id=created.run_id, user_id=1, organization_id=2) + replay = service.cancel(run_id=created.run_id, user_id=1, organization_id=2) + + assert cancelled.run_status is RunStatus.CANCELLED + assert replay == cancelled + + +def test_cancel_rejects_completed_run(service: OnboardingProgressService) -> None: + created, token = create_run( + service, user_id=1, organization_id=2, client_run_id="browser-session" + ) + service.update( + token=token, + user_id=1, + organization_id=2, + update=ProgressUpdate( + stage=Stage.CHECK_STACK_TRACE_QUALITY, + status=StageStatus.SKIPPED, + run_status=RunStatus.COMPLETED, + ), + ) + + with pytest.raises(OnboardingRunTerminal): + service.cancel(run_id=created.run_id, user_id=1, organization_id=2) + + +def test_expiration_is_absolute(service: OnboardingProgressService, frozen_time: Any) -> None: + created, token = create_run( + service, user_id=1, organization_id=2, client_run_id="browser-session" + ) + frozen_time.shift(RUN_LIFETIME - timedelta(hours=1)) + service.update( + token=token, + user_id=1, + organization_id=2, + update=ProgressUpdate(stage=Stage.CONNECT_MCP, status=StageStatus.COMPLETED), + ) + + state_key = f"agentic-onboarding:run:{{{created.run_id}}}:state" + assert service.redis.ttl(state_key) == int(timedelta(hours=1).total_seconds()) + + frozen_time.shift(timedelta(hours=1)) + assert service.get(run_id=created.run_id, user_id=1, organization_id=2) is None + with pytest.raises(ValueError, match="expired"): + service.update( + token=token, + user_id=1, + organization_id=2, + update=ProgressUpdate(stage=Stage.ANALYZE_PROJECT, status=StageStatus.COMPLETED), + ) + + +def test_new_run_uses_one_week_ttl(service: OnboardingProgressService) -> None: + created, _ = create_run(service, user_id=1, organization_id=2, client_run_id="browser-session") + assert service.redis.ttl(f"agentic-onboarding:run:{{{created.run_id}}}:state") == int( + RUN_LIFETIME.total_seconds() + ) + + +def test_create_replaces_application_expired_run_retained_in_redis( + service: OnboardingProgressService, frozen_time: Any +) -> None: + expired, _ = create_run(service, user_id=1, organization_id=2, client_run_id="browser-session") + expired_state_key = f"agentic-onboarding:run:{{{expired.run_id}}}:state" + retained = replace(expired, expires_at=expired.expires_at - RUN_LIFETIME) + service.redis.set( + expired_state_key, service._serialize(retained), ex=int(RUN_LIFETIME.total_seconds()) + ) + + assert service.redis.get(expired_state_key) is not None + + created, token = create_run( + service, user_id=1, organization_id=2, client_run_id="browser-session" + ) + + assert created.run_id != expired.run_id + assert created.created_at == django_timezone.now() + assert token == "abcdefghij" + + +def test_registration_failure_releases_claimed_indexes( + service: OnboardingProgressService, +) -> None: + original_set = service.redis.set + + def fail_state_write(key: str, *args: Any, **kwargs: Any) -> Any: + if key.endswith(":state"): + raise RuntimeError("write failed") + return original_set(key, *args, **kwargs) + + with patch.object(service.redis, "set", side_effect=fail_state_write): + with pytest.raises(RuntimeError, match="write failed"): + create_run(service, user_id=1, organization_id=2, client_run_id="browser-session") + + token_key = service._token_index_key(service._hash_token("abcdefghij")) + client_key = service._client_index_key(1, 2, "browser-session") + assert service.redis.get(token_key) is None + assert service.redis.get(client_key) is None + + +def test_index_cleanup_does_not_delete_a_new_owner(service: OnboardingProgressService) -> None: + key = service._token_index_key(service._hash_token("abcdefghij")) + service.redis.set(key, "winning-run") + + service._release_index(key, "losing-run") + + assert service.redis.get(key) == b"winning-run" + + +def test_token_claim_waits_for_registration_state(service: OnboardingProgressService) -> None: + key = service._token_index_key(service._hash_token("abcdefghij")) + service.redis.set(key, "registering-run") + + with ( + patch.object(service, "_load", return_value=None) as load, + patch("sentry.onboarding.agentic_progress.service.time.sleep") as sleep, + ): + service._claim_token_index(key, "competing-run", 60) + + assert load.call_count == CLIENT_CLAIM_RETRIES + assert sleep.call_count == CLIENT_CLAIM_RETRIES - 1 + assert service.redis.get(key) == b"competing-run" + + +def test_redis_round_trip_preserves_datetime_fields(service: OnboardingProgressService) -> None: + created, _ = create_run(service, user_id=1, organization_id=2, client_run_id="browser-session") + + restored = service.get(run_id=created.run_id, user_id=1, organization_id=2) + + assert restored is not None + assert isinstance(restored.created_at, datetime) + assert isinstance(restored.updated_at, datetime) + assert isinstance(restored.expires_at, datetime) + + +def test_corrupted_state_is_treated_as_missing(service: OnboardingProgressService) -> None: + created, token = create_run( + service, user_id=1, organization_id=2, client_run_id="browser-session" + ) + service.redis.set(f"agentic-onboarding:run:{{{created.run_id}}}:state", b"not-json") + + assert service.get(run_id=created.run_id, user_id=1, organization_id=2) is None + with pytest.raises(RunNotFound): + service.update( + token=token, + user_id=1, + organization_id=2, + update=ProgressUpdate(stage=Stage.CONNECT_MCP, status=StageStatus.COMPLETED), + ) + + +def test_atomic_update_bounds_contention_retries(service: OnboardingProgressService) -> None: + created, _ = create_run(service, user_id=1, organization_id=2, client_run_id="browser-session") + + with patch.object(service.redis, "pipeline") as pipeline: + active_pipeline = pipeline.return_value.__enter__.return_value + active_pipeline.get.return_value = service._serialize(created) + active_pipeline.execute.side_effect = WatchError + with pytest.raises(RuntimeError, match="Unable to update onboarding progress"): + service.cancel(run_id=created.run_id, user_id=1, organization_id=2) + + assert pipeline.call_count == ATOMIC_UPDATE_RETRIES