diff --git a/.github/workflows/container-images-cd.yml b/.github/workflows/container-images-cd.yml index b9b342daf4a2..8d20efa74751 100644 --- a/.github/workflows/container-images-cd.yml +++ b/.github/workflows/container-images-cd.yml @@ -444,7 +444,7 @@ jobs: - name: Check for changes that affect Tasks Agent Temporal worker id: check_changes_tasks_agent_temporal_worker run: | - if git diff --name-only HEAD^ HEAD | grep -qE '^posthog/temporal/common|^posthog/temporal/ai/posthog_code_slack_|^posthog/temporal/ai/slack_app/|^products/slack_app/backend|^products/tasks/backend|^posthog/tasks/|^posthog/settings/temporal.py$|^posthog/management/commands/start_temporal_worker.py$|^pyproject.toml$|^bin/temporal-django-worker$|^posthog/clickhouse/migrations/'; then + if git diff --name-only HEAD^ HEAD | grep -qE '^posthog/temporal/common|^posthog/temporal/ai/posthog_code_slack_|^posthog/temporal/ai/slack_app/|^products/slack_app/backend|^products/tasks/backend|^products/wizard/backend/|^posthog/tasks/|^posthog/settings/temporal.py$|^posthog/management/commands/start_temporal_worker.py$|^pyproject.toml$|^bin/temporal-django-worker$|^posthog/clickhouse/migrations/'; then echo "changed=true" >> $GITHUB_OUTPUT else echo "changed=false" >> $GITHUB_OUTPUT diff --git a/docs/internal/wizard-registry.md b/docs/internal/wizard-registry.md index 286399c46529..fedcbec96706 100644 --- a/docs/internal/wizard-registry.md +++ b/docs/internal/wizard-registry.md @@ -1,8 +1,8 @@ # Wizard registry Status: design document for the Wizard run stack. -This layer ships only the data models and the registry parsing described under "Remote configuration". -The "API" routes and the "Cloud execution" workers describe planned behavior and do not exist yet; they land in later stack layers. +This layer ships the data models, registry parsing, and cloud execution workers. +The "API" routes describe planned behavior and land in a later stack layer. The Wizard Registry lists the programs a person can run in the Wizard. Programs are personalized by the signed-in user's distinct ID and the organization that owns the selected project. @@ -66,6 +66,26 @@ Its command is empty, which keeps the Wizard package's default command. ## Cloud execution +Wizard analytics events use the process-lifetime background capture client without a Celery hop. +Delivery is best-effort; an abrupt process exit can lose buffered events. + +Wizard execution returns only the last 64 KiB of stdout and stderr from the sandbox, preserving the command's exit code. +Log streams are drained through bounded tails without storing unbounded log files. +Repository publishing bounds sandbox reads and rejects staged contents that exceed the cumulative 35 MiB commit payload budget, including base64 expansion. + +### Dedicated worker rollout + +`WIZARD_TASK_QUEUE` controls both Wizard dispatch and worker registration. It defaults to `general-purpose-task-queue` during rollout. +`TEMPORAL_TASK_QUEUE` selects the queue a worker process polls; setting it alone does not register Wizard workflows on that queue. + +1. Deploy an image containing the Wizard workflow and activity registration before enabling the dedicated worker. +2. Set both `TEMPORAL_TASK_QUEUE` and `WIZARD_TASK_QUEUE` to `wizard-task-queue` on the dedicated worker. +3. Verify the worker is healthy and polling before setting `WIZARD_TASK_QUEUE=wizard-task-queue` on dispatching services. + +The dedicated worker uses the Tasks Agent image release. Changes to Wizard backend code must trigger that image's deployment. + +### Program execution + Cloud workers execute the program snapshot stored when the run was created. They do not evaluate the registry again. They run the pinned `wizard_version` from that snapshot, never a mutable tag such as `latest`. diff --git a/posthog/management/commands/start_temporal_worker.py b/posthog/management/commands/start_temporal_worker.py index 9479e05a7f01..1c44ccb0348e 100644 --- a/posthog/management/commands/start_temporal_worker.py +++ b/posthog/management/commands/start_temporal_worker.py @@ -275,6 +275,10 @@ ACTIVITIES as WA_DIGEST_ACTIVITIES, WORKFLOWS as WA_DIGEST_WORKFLOWS, ) +from products.wizard.backend.facade.temporal import ( + ACTIVITIES as WIZARD_ACTIVITIES, + WORKFLOWS as WIZARD_WORKFLOWS, +) # When adding modules to a queue, also update the corresponding CI trigger # in .github/workflows/container-images-cd.yml (check_changes_*_temporal_worker) @@ -414,6 +418,11 @@ TASKS_WORKFLOWS + POSTHOG_CODE_SLACK_WORKFLOWS, TASKS_ACTIVITIES + POSTHOG_CODE_SLACK_ACTIVITIES, ), + ( + settings.WIZARD_TASK_QUEUE, + WIZARD_WORKFLOWS, + WIZARD_ACTIVITIES, + ), ( settings.MAX_AI_TASK_QUEUE, AI_WORKFLOWS, diff --git a/posthog/management/commands/test/test_start_temporal_worker.py b/posthog/management/commands/test/test_start_temporal_worker.py index b0d9c9d77273..5cee1d32daa6 100644 --- a/posthog/management/commands/test/test_start_temporal_worker.py +++ b/posthog/management/commands/test/test_start_temporal_worker.py @@ -1,5 +1,7 @@ import pytest +from django.conf import settings + from posthog.management.commands.start_temporal_worker import ( DATA_SYNC_WORKFLOWS, WA_DIGEST_ACTIVITIES, @@ -9,11 +11,30 @@ workflows_include_data_import_syncs, ) +from products.wizard.backend.facade.temporal import ( + ACTIVITIES as WIZARD_ACTIVITIES, + WORKFLOWS as WIZARD_WORKFLOWS, +) + class _NotADataSyncWorkflow: pass +def test_wizard_queue_registers_workflows_and_activities() -> None: + entries = [ + (workflows, activities) + for queue, workflows, activities in _task_queue_specs + if queue == settings.WIZARD_TASK_QUEUE + ] + workflows = {workflow for registered_workflows, _ in entries for workflow in registered_workflows} + activities = {activity for _, registered_activities in entries for activity in registered_activities} + assert WIZARD_WORKFLOWS + assert WIZARD_ACTIVITIES + assert set(WIZARD_WORKFLOWS) <= workflows + assert set(WIZARD_ACTIVITIES) <= activities + + # Data-import sources import vendor SDKs (google-ads, etc.) that register protobuf descriptors into a # process-global pool exactly once. The worker eagerly loads them at boot only for queues that run # data syncs; everything else stays lazy to keep startup fast. Queue settings collapse to a single diff --git a/posthog/settings/temporal.py b/posthog/settings/temporal.py index 36e2cd0b9bf7..37a9eccbbc67 100644 --- a/posthog/settings/temporal.py +++ b/posthog/settings/temporal.py @@ -245,6 +245,11 @@ def _set_temporal_task_queue(task_queue: str) -> str: HEALTH_CHECK_TASK_QUEUE = _set_temporal_task_queue("health-check-task-queue") DUCKLAKE_TASK_QUEUE = _set_temporal_task_queue("ducklake-task-queue") TASKS_TASK_QUEUE = _set_temporal_task_queue("tasks-task-queue") +# Defaults to the general-purpose fleet so dispatch always has a live worker. Routing Wizard runs +# to a dedicated, separately-scalable worker takes two steps in order: deploy a worker fleet +# polling "wizard-task-queue", then set this env on the dispatching services. Setting it first +# strands runs on a pollerless queue until the workflow execution timeout closes them. +WIZARD_TASK_QUEUE = _set_temporal_task_queue(os.getenv("WIZARD_TASK_QUEUE", "general-purpose-task-queue")) WIZARD_RUN_ARTIFACTS_S3_BUCKET = os.getenv("WIZARD_RUN_ARTIFACTS_S3_BUCKET", "") TASKS_DISPATCHER_BATCH_SIZE = get_from_env("TASKS_DISPATCHER_BATCH_SIZE", 50, type_cast=int) TASKS_DISPATCHER_CONCURRENCY = get_from_env("TASKS_DISPATCHER_CONCURRENCY", 20, type_cast=int) diff --git a/products/tasks/backend/facade/repo_selection.py b/products/tasks/backend/facade/repo_selection.py index c1f26e6243c6..89408b003432 100644 --- a/products/tasks/backend/facade/repo_selection.py +++ b/products/tasks/backend/facade/repo_selection.py @@ -6,6 +6,7 @@ handle the typed outcomes. """ +from products.tasks.backend.github_repository_access import inaccessible_repositories_via_integration from products.tasks.backend.logic.repo_selection import ( REPO_SELECTION_DUMMY_REPOSITORY, RepoSelectionRejectedError, @@ -15,13 +16,27 @@ resolve_team_github_integration, select_repository, ) +from products.tasks.backend.temporal.process_task.utils import get_github_token + + +def repository_accessible_via_integration(team_id: int, integration_id: int, repository: str) -> bool: + return not inaccessible_repositories_via_integration(team_id, integration_id, [repository]) + + +def resolve_team_github_integration_id(team_id: int) -> int | None: + integration = resolve_team_github_integration(team_id, team_only=True) + return integration.integration.id if integration is not None else None + __all__ = [ "REPO_SELECTION_DUMMY_REPOSITORY", "RepoSelectionRejectedError", "RepoSelectionResult", "RepoSelectionUnavailableError", + "get_github_token", "list_team_connected_repositories", + "repository_accessible_via_integration", "resolve_team_github_integration", + "resolve_team_github_integration_id", "select_repository", ] diff --git a/products/tasks/backend/facade/sandbox.py b/products/tasks/backend/facade/sandbox.py index 569f51bf1dd7..c980b8d92c42 100644 --- a/products/tasks/backend/facade/sandbox.py +++ b/products/tasks/backend/facade/sandbox.py @@ -26,6 +26,7 @@ get_sandbox_class, get_sandbox_class_for_backend, is_public_sandbox_repo, + sandbox_repo_path, ) __all__ = [ @@ -45,4 +46,5 @@ "get_sandbox_class", "get_sandbox_class_for_backend", "is_public_sandbox_repo", + "sandbox_repo_path", ] diff --git a/products/tasks/backend/logic/services/sandbox.py b/products/tasks/backend/logic/services/sandbox.py index 957a0c333dbb..c34a03c0d3c7 100644 --- a/products/tasks/backend/logic/services/sandbox.py +++ b/products/tasks/backend/logic/services/sandbox.py @@ -231,12 +231,18 @@ def effective_memory_request_mb(self) -> int: """Repos the sandbox is allowed to clone unauthenticated, even when the team has no GitHub integration""" # TODO: Remove `posthog/.github` when we switch repo discovery to repo-less agent (now it works as a lightweight dummy) -SENSITIVE_AGENT_RUNTIME_ENV_NAMES: frozenset[str] = frozenset( - {"POSTHOG_TASK_RUN_EVENT_INGEST_TOKEN", "POSTHOG_TASK_RUN_SESSION_TOKEN"} +# this helps redact sensitive environment variables for logging +SENSITIVE_SANDBOX_ENV_NAMES: frozenset[str] = frozenset( + { + "GITHUB_TOKEN", + "POSTHOG_TASK_RUN_EVENT_INGEST_TOKEN", + "POSTHOG_TASK_RUN_SESSION_TOKEN", + "POSTHOG_WIZARD_API_KEY", + } ) SHELL_ARGUMENT_VALUE_PATTERN = r"'(?:[^']|'\"'\"')*'|\"(?:\\.|[^\"])*\"|\S+" -SENSITIVE_AGENT_RUNTIME_ENV_PATTERN = re.compile( - r"(?P" + "|".join(re.escape(name) for name in SENSITIVE_AGENT_RUNTIME_ENV_NAMES) + r")=" +SENSITIVE_SANDBOX_ENV_PATTERN = re.compile( + r"(?P" + "|".join(re.escape(name) for name in SENSITIVE_SANDBOX_ENV_NAMES) + r")=" rf"(?P{SHELL_ARGUMENT_VALUE_PATTERN})" ) SENSITIVE_AGENT_RUNTIME_ARGUMENT_PATTERN = re.compile( @@ -245,6 +251,7 @@ def effective_memory_request_mb(self) -> int: SENSITIVE_FILE_HEREDOC_PATTERN = re.compile( r"(?P<<'POSTHOG_FILE_EOF'\n).*?(?P\nPOSTHOG_FILE_EOF)", re.DOTALL ) +GITHUB_CLONE_TOKEN_PATTERN = re.compile(r"(?Phttps://x-access-token:)[^@\s]+(?P@github\.com/)") def is_public_sandbox_repo(repository: str | None) -> bool: @@ -258,9 +265,15 @@ def sandbox_repo_path(repository: str) -> str: def redact_sandbox_command(command: str) -> str: - redacted = SENSITIVE_AGENT_RUNTIME_ENV_PATTERN.sub(r"\g=", command) - redacted = SENSITIVE_AGENT_RUNTIME_ARGUMENT_PATTERN.sub(r"\g ", redacted) - return SENSITIVE_FILE_HEREDOC_PATTERN.sub(r"\g\g", redacted) + redacted = command + for pattern, substitution in ( + (SENSITIVE_SANDBOX_ENV_PATTERN, r"\g="), + (SENSITIVE_AGENT_RUNTIME_ARGUMENT_PATTERN, r"\g "), + (SENSITIVE_FILE_HEREDOC_PATTERN, r"\g\g"), + (GITHUB_CLONE_TOKEN_PATTERN, r"\g\g"), + ): + redacted = pattern.sub(substitution, redacted) + return redacted def build_agent_runtime_env_prefix( diff --git a/products/tasks/backend/logic/services/test_docker_sandbox.py b/products/tasks/backend/logic/services/test_docker_sandbox.py index c6ea96494581..5836879d2d54 100644 --- a/products/tasks/backend/logic/services/test_docker_sandbox.py +++ b/products/tasks/backend/logic/services/test_docker_sandbox.py @@ -126,7 +126,7 @@ def test_redact_sandbox_command_hides_event_ingest_token(self, command): assert "secret token" not in redacted assert "POSTHOG_TASK_RUN_EVENT_INGEST_TOKEN=" in redacted - def test_redact_sandbox_command_hides_mcp_credentials(self): + def test_redact_sandbox_command_hides_mcp_credentials(self) -> None: command = 'agent-server --mcpServers \'[{"headers":{"Authorization":"Bearer secret-token"}}]\' --port 8080' redacted = redact_sandbox_command(command) @@ -141,6 +141,21 @@ def test_redact_sandbox_command_hides_file_heredoc(self): assert "c2VjcmV0" not in redacted assert "" in redacted + def test_redact_sandbox_command_hides_github_clone_token(self) -> None: + command = "git clone https://x-access-token:github-secret@github.com/PostHog/posthog.git" + + redacted = redact_sandbox_command(command) + + assert "github-secret" not in redacted + assert "https://x-access-token:@github.com/PostHog/posthog.git" in redacted + + @pytest.mark.parametrize("name", ("GITHUB_TOKEN", "POSTHOG_WIZARD_API_KEY")) + def test_redact_sandbox_command_hides_worker_credentials(self, name: str) -> None: + redacted = redact_sandbox_command(f"docker run -e {name}=worker-secret sandbox") + + assert "worker-secret" not in redacted + assert f"{name}=" in redacted + @pytest.mark.parametrize( "input_url,expected_url", [ diff --git a/products/wizard/backend/facade/api.py b/products/wizard/backend/facade/api.py index 79d9783155dc..c878a910def4 100644 --- a/products/wizard/backend/facade/api.py +++ b/products/wizard/backend/facade/api.py @@ -16,19 +16,27 @@ from products.wizard.backend import metrics from products.wizard.backend.facade.contracts import ( CreatePullRequestArtifactInput, + CreateWizardRunInput, + ListWizardRunsInput, UpsertWizardSessionInput, WizardProgram, WizardRunArtifactDTO, + WizardRunCreationResult, + WizardRunDTO, WizardRunGitDiffArtifactDTO, + WizardRunPage, WizardRunPullRequestArtifactDTO, WizardSessionDTO, ) +from products.wizard.backend.facade.enums import WizardRunStage, WizardRunStatus +from products.wizard.backend.facade.validation import validate_wizard_version as validate_wizard_version_value from products.wizard.backend.logic import ( - pubsub, registry as registry_service, + runs as run_service, sessions, ) from products.wizard.backend.logic.artifacts import service as artifacts +from products.wizard.backend.logic.sessions import pubsub def upsert(params: UpsertWizardSessionInput) -> tuple[WizardSessionDTO, bool]: @@ -91,6 +99,40 @@ def get_registry(*, distinct_id: str, organization_id: str) -> tuple[WizardProgr return registry_service.get_registry(distinct_id=distinct_id, organization_id=organization_id) +def create_run(params: CreateWizardRunInput) -> WizardRunDTO: + return run_service.create_run(params) + + +def create_run_with_result(params: CreateWizardRunInput) -> WizardRunCreationResult: + return run_service.create_run_with_result(params) + + +def get_run(team_id: int, run_id: UUID) -> WizardRunDTO: + return run_service.get_run(team_id, run_id) + + +def list_runs(params: ListWizardRunsInput) -> WizardRunPage: + return run_service.list_runs(params) + + +def update_run_stage(team_id: int, run_id: UUID, stage: WizardRunStage) -> WizardRunDTO: + return run_service.update_run_stage(team_id, run_id, stage) + + +def cancel_run(team_id: int, run_id: UUID) -> WizardRunDTO: + return run_service.cancel_run(team_id, run_id) + + +def update_run_status( + team_id: int, + run_id: UUID, + status: WizardRunStatus, + *, + error_code: str | None = None, +) -> WizardRunDTO: + return run_service.transition_run(team_id, run_id, status, error_code=error_code) + + def create_git_diff_artifact(team_id: int, run_id: UUID, content: bytes) -> WizardRunGitDiffArtifactDTO | None: """ Uploads a git diff generated in a Wizard run. @@ -117,3 +159,11 @@ def get_git_diff_artifact_content(team_id: int, run_id: UUID, artifact_id: UUID) Returns the content of a git diff artifact generated in a Wizard run. """ return artifacts.get_git_diff_artifact_content(team_id, run_id, artifact_id) + + +def validate_git_repository(repository: str) -> None: + run_service.validate_git_repository_name(repository) + + +def validate_wizard_version(wizard_version: object) -> str: + return validate_wizard_version_value(wizard_version) diff --git a/products/wizard/backend/facade/temporal.py b/products/wizard/backend/facade/temporal.py new file mode 100644 index 000000000000..28d82ff3594e --- /dev/null +++ b/products/wizard/backend/facade/temporal.py @@ -0,0 +1,3 @@ +from products.wizard.backend.temporal import ACTIVITIES, WORKFLOWS + +__all__ = ["ACTIVITIES", "WORKFLOWS"] diff --git a/products/wizard/backend/logic/runs/__init__.py b/products/wizard/backend/logic/runs/__init__.py new file mode 100644 index 000000000000..a07ac0f8d41d --- /dev/null +++ b/products/wizard/backend/logic/runs/__init__.py @@ -0,0 +1,27 @@ +from products.wizard.backend.logic.runs.lifecycle import ( + cancel_run, + complete_run, + create_run, + create_run_with_result, + fail_run, + get_run, + list_runs, + start_run, + transition_run, + update_run_stage, +) +from products.wizard.backend.logic.runs.validation import validate_git_repository_name + +__all__ = [ + "cancel_run", + "complete_run", + "create_run", + "create_run_with_result", + "fail_run", + "get_run", + "list_runs", + "start_run", + "transition_run", + "update_run_stage", + "validate_git_repository_name", +] diff --git a/products/wizard/backend/logic/runs/admission.py b/products/wizard/backend/logic/runs/admission.py new file mode 100644 index 000000000000..2810b86bec57 --- /dev/null +++ b/products/wizard/backend/logic/runs/admission.py @@ -0,0 +1,36 @@ +from django.utils import timezone + +from products.wizard.backend.facade.enums import WizardRunEnvironment, WizardRunStatus +from products.wizard.backend.facade.errors import ( + ActiveWizardRunError, + WizardRunDailyLimitError, + WizardRunHourlyLimitError, +) +from products.wizard.backend.logic.runs.config import ( + CLOUD_RUN_DAILY_LIMIT, + CLOUD_RUN_DAILY_WINDOW, + CLOUD_RUN_HOURLY_LIMIT, + CLOUD_RUN_HOURLY_WINDOW, +) +from products.wizard.backend.models import WizardRun + + +def enforce_cloud_run_creation_policy(team_id: int, created_by_id: int, idempotency_key: str | None = None) -> None: + runs = WizardRun.objects.for_team(team_id).filter( + created_by_id=created_by_id, + environment=WizardRunEnvironment.CLOUD.value, + ) + + if idempotency_key is not None: + runs = runs.exclude(idempotency_key=idempotency_key) + + if runs.filter(status__in=(WizardRunStatus.CREATED.value, WizardRunStatus.RUNNING.value)).exists(): + raise ActiveWizardRunError + + now = timezone.now() + + if runs.filter(created_at__gte=now - CLOUD_RUN_HOURLY_WINDOW).count() >= CLOUD_RUN_HOURLY_LIMIT: + raise WizardRunHourlyLimitError + + if runs.filter(created_at__gte=now - CLOUD_RUN_DAILY_WINDOW).count() >= CLOUD_RUN_DAILY_LIMIT: + raise WizardRunDailyLimitError diff --git a/products/wizard/backend/logic/runs/cancellation.py b/products/wizard/backend/logic/runs/cancellation.py new file mode 100644 index 000000000000..27059ee7594c --- /dev/null +++ b/products/wizard/backend/logic/runs/cancellation.py @@ -0,0 +1,23 @@ +import logging +from uuid import UUID + +from products.wizard.backend.logic.runs import store +from products.wizard.backend.temporal import client as temporal_client +from products.wizard.backend.temporal.errors import WizardTemporalError + +logger = logging.getLogger(__name__) + + +def dispatch_cancellation(team_id: int, run_id: UUID) -> bool: + if store.get_workflow_id(team_id, run_id) is None: + return False + + try: + temporal_client.cancel_wizard_run_workflow(run_id) + except WizardTemporalError: + logger.exception("wizard_run_cancellation_failed", extra={"team_id": team_id, "run_id": str(run_id)}) + return False + + store.mark_cancellation_dispatched(team_id, run_id) + + return True diff --git a/products/wizard/backend/logic/runs/dispatch.py b/products/wizard/backend/logic/runs/dispatch.py new file mode 100644 index 000000000000..6162e945a0e8 --- /dev/null +++ b/products/wizard/backend/logic/runs/dispatch.py @@ -0,0 +1,43 @@ +from uuid import UUID + +from products.wizard.backend.facade.enums import WizardRunEnvironment, WizardRunStatus +from products.wizard.backend.logic.runs import ( + cancellation as cancellation_service, + store, +) +from products.wizard.backend.logic.runs.errors import WizardRunDispatchError +from products.wizard.backend.logic.workers.config import local_wizard_source_root +from products.wizard.backend.observability.contracts import WizardRunDispatchOutcome +from products.wizard.backend.observability.service import wizard_observability +from products.wizard.backend.temporal import client as temporal_client +from products.wizard.backend.temporal.constants import wizard_run_workflow_id +from products.wizard.backend.temporal.contracts import WizardRunActivityInput +from products.wizard.backend.temporal.errors import WizardTemporalError + + +def dispatch_created_cloud_wizard_run_to_temporal_worker(team_id: int, run_id: UUID) -> None: + run = store.get_run(team_id, run_id) + + if run.environment != WizardRunEnvironment.CLOUD or run.status != WizardRunStatus.CREATED: + return + + try: + temporal_client.start_wizard_run_workflow( + WizardRunActivityInput( + team_id=team_id, + run_id=run_id, + use_local_wizard_source=local_wizard_source_root() is not None, + ) + ) + except WizardTemporalError as error: + # marks the failed dispatch, which returns whether the run has exhausted its retry attempts + is_exhausted = store.mark_dispatch_failed(team_id, run_id) + + wizard_observability.dispatch_finished(run, WizardRunDispatchOutcome.FAILED) + + raise WizardRunDispatchError(is_exhausted=is_exhausted) from error + + store.mark_dispatch_succeeded(team_id, run_id, wizard_run_workflow_id(run_id)) + if store.cancellation_requested(team_id, run_id): + cancellation_service.dispatch_cancellation(team_id, run_id) + wizard_observability.dispatch_finished(run, WizardRunDispatchOutcome.SUCCEEDED) diff --git a/products/wizard/backend/logic/runs/errors.py b/products/wizard/backend/logic/runs/errors.py index bee1190f4a2e..c3f65c84fcac 100644 --- a/products/wizard/backend/logic/runs/errors.py +++ b/products/wizard/backend/logic/runs/errors.py @@ -1,6 +1,6 @@ class WizardRunDispatchError(Exception): - def __init__(self, *, exhausted: bool) -> None: - self.exhausted = exhausted + def __init__(self, *, is_exhausted: bool) -> None: + self.exhausted = is_exhausted super().__init__() diff --git a/products/wizard/backend/logic/runs/lifecycle.py b/products/wizard/backend/logic/runs/lifecycle.py new file mode 100644 index 000000000000..9678ac96d089 --- /dev/null +++ b/products/wizard/backend/logic/runs/lifecycle.py @@ -0,0 +1,233 @@ +import logging +from dataclasses import replace +from functools import partial +from uuid import UUID + +from django.db import transaction as database_transaction + +from posthog.models import Team, User + +from products.wizard.backend.facade.config import DEFAULT_WIZARD_VERSION +from products.wizard.backend.facade.contracts import ( + CreateWizardRunInput, + GitRepositoryWorkspace, + ListWizardRunsInput, + LocalFolderWorkspace, + WizardRunCreationResult, + WizardRunDTO, + WizardRunPage, +) +from products.wizard.backend.facade.enums import ( + WizardRunEnvironment, + WizardRunErrorCode, + WizardRunStage, + WizardRunStatus, +) +from products.wizard.backend.facade.errors import ( + IllegalStatusTransitionError, + InvalidWorkspaceEnvironmentError, + MissingWizardRunIdempotencyKeyError, + WizardProgramEnvironmentNotSupportedError, + WizardRunIdempotencyConflictError, +) +from products.wizard.backend.logic import registry as registry_service +from products.wizard.backend.logic.runs import ( + cancellation as cancellation_service, + store, +) +from products.wizard.backend.logic.runs.admission import enforce_cloud_run_creation_policy +from products.wizard.backend.logic.runs.dispatch import dispatch_created_cloud_wizard_run_to_temporal_worker +from products.wizard.backend.logic.runs.errors import WizardRunDispatchError +from products.wizard.backend.logic.runs.fingerprints import create_run_request_fingerprint +from products.wizard.backend.logic.runs.repository_access import authorize_git_repository_access +from products.wizard.backend.logic.runs.transitions import transition +from products.wizard.backend.observability.service import wizard_observability as run_observability + +logger = logging.getLogger(__name__) + + +def create_run(params: CreateWizardRunInput) -> WizardRunDTO: + return create_run_with_result(params).run + + +def create_run_with_result(params: CreateWizardRunInput) -> WizardRunCreationResult: + match params.environment, params.workspace: + case WizardRunEnvironment.LOCAL, LocalFolderWorkspace(): + pass + case WizardRunEnvironment.CLOUD, GitRepositoryWorkspace(): + pass + case _: + raise InvalidWorkspaceEnvironmentError + + request_fingerprint: str | None = None + is_cloud_run = params.environment == WizardRunEnvironment.CLOUD + + if is_cloud_run: + if params.idempotency_key is None: + raise MissingWizardRunIdempotencyKeyError + + request_fingerprint = create_run_request_fingerprint(params) + + existing = store.get_run_by_idempotency_key(params.team_id, params.idempotency_key) + + if existing is not None: + if store.get_request_fingerprint(params.team_id, existing.id) != request_fingerprint: + raise WizardRunIdempotencyConflictError + + return WizardRunCreationResult(run=existing, created=False) + + user = User.objects.only("distinct_id").get(id=params.created_by_id) + team = Team.objects.only("organization_id").get(id=params.team_id) + + program = registry_service.get_program( + program_id=params.program_id, + distinct_id=user.distinct_id, + organization_id=str(team.organization_id), + ) + + program = replace(program, wizard_version=params.wizard_version or DEFAULT_WIZARD_VERSION) + + if params.environment not in program.supported_environments: + raise WizardProgramEnvironmentNotSupportedError + + if isinstance(params.workspace, GitRepositoryWorkspace): + authorize_git_repository_access(params.team_id, params.workspace.repository) + + initial_status = ( + WizardRunStatus.RUNNING if params.environment == WizardRunEnvironment.LOCAL else WizardRunStatus.CREATED + ) + + with database_transaction.atomic(): + if is_cloud_run: + enforce_cloud_run_creation_policy(params.team_id, params.created_by_id, params.idempotency_key) + + result = store.create_run( + team_id=params.team_id, + created_by_id=params.created_by_id, + environment=params.environment, + workspace=params.workspace, + program=program, + status=initial_status, + idempotency_key=params.idempotency_key, + request_fingerprint=request_fingerprint, + ) + + if not result.created and store.get_request_fingerprint(params.team_id, result.run.id) != request_fingerprint: + raise WizardRunIdempotencyConflictError + + should_dispatch_wizard_run = is_cloud_run and result.created + + if should_dispatch_wizard_run: + database_transaction.on_commit( + partial( + _dispatch_cloud_run, + params.team_id, + result.run.id, + ), + ) + + if is_cloud_run: + result = WizardRunCreationResult(run=store.get_run(params.team_id, result.run.id), created=result.created) + + if result.created and not is_cloud_run: + run_observability.run_created(result.run) + + return result + + +def _dispatch_cloud_run(team_id: int, run_id: UUID) -> None: + run_observability.run_created(store.get_run(team_id, run_id)) + + try: + dispatch_created_cloud_wizard_run_to_temporal_worker(team_id, run_id) + + except WizardRunDispatchError as error: + logger.exception( + "wizard_run_dispatch_failed", + extra={"team_id": team_id, "run_id": str(run_id), "exhausted": error.exhausted}, + ) + if error.exhausted: + fail_run(team_id, run_id, error_code=WizardRunErrorCode.DISPATCH_FAILED.value) + + +def get_run(team_id: int, run_id: UUID) -> WizardRunDTO: + return store.get_run(team_id, run_id) + + +def list_runs(params: ListWizardRunsInput) -> WizardRunPage: + return store.list_runs(params) + + +def start_run(team_id: int, run_id: UUID) -> WizardRunDTO: + return transition_run(team_id, run_id, WizardRunStatus.RUNNING) + + +def complete_run(team_id: int, run_id: UUID) -> WizardRunDTO: + return transition_run(team_id, run_id, WizardRunStatus.COMPLETED) + + +def fail_run( + team_id: int, + run_id: UUID, + *, + error_code: str | None = None, +) -> WizardRunDTO: + return transition_run(team_id, run_id, WizardRunStatus.FAILED, error_code=error_code) + + +def cancel_run(team_id: int, run_id: UUID) -> WizardRunDTO: + with database_transaction.atomic(): + run = store.get_run_for_update(team_id, run_id) + previous = run + next_status = transition(run.status, WizardRunStatus.CANCELLED) + run = store.set_run_status(team_id, run_id, next_status, None) + + if previous.environment == WizardRunEnvironment.CLOUD: + # Persist the cancellation intent in the same transaction as the status, so a crash + # before dispatch still leaves the run in the recovery index. Send the Temporal cancel + # only after the commit, since it is an external side effect. + store.mark_cancellation_requested(team_id, run_id) + database_transaction.on_commit(partial(cancellation_service.dispatch_cancellation, team_id, run_id)) + + run_observability.run_transitioned(previous, run) + return run + + +def request_cloud_run_cancellation(team_id: int, run_id: UUID) -> None: + store.mark_cancellation_requested(team_id, run_id) + cancellation_service.dispatch_cancellation(team_id, run_id) + + +def update_run_stage(team_id: int, run_id: UUID, stage: WizardRunStage) -> WizardRunDTO: + with database_transaction.atomic(): + run = store.get_run_for_update(team_id, run_id) + previous = run + + if run.status not in (WizardRunStatus.CREATED, WizardRunStatus.RUNNING): + raise IllegalStatusTransitionError + + if run.stage == stage: + return run + + run = store.set_run_stage(team_id, run_id, stage) + + run_observability.stage_changed(previous, run) + + return run + + +def transition_run( + team_id: int, + run_id: UUID, + next_status: WizardRunStatus, + *, + error_code: str | None = None, +) -> WizardRunDTO: + with database_transaction.atomic(): + run = store.get_run_for_update(team_id, run_id) + previous = run + next_status = transition(run.status, next_status, error_code=error_code) + run = store.set_run_status(team_id, run_id, next_status, error_code) + + run_observability.run_transitioned(previous, run) + return run diff --git a/products/wizard/backend/logic/runs/repository_access.py b/products/wizard/backend/logic/runs/repository_access.py new file mode 100644 index 000000000000..f46890b56dd1 --- /dev/null +++ b/products/wizard/backend/logic/runs/repository_access.py @@ -0,0 +1,16 @@ +from products.tasks.backend.facade import repo_selection +from products.wizard.backend.facade.errors import MissingGitHubIntegrationError, RepositoryNotAccessibleError +from products.wizard.backend.logic.runs.validation import validate_git_repository_name + + +def authorize_git_repository_access(team_id: int, repository: str) -> int: + validate_git_repository_name(repository) + + integration_id = repo_selection.resolve_team_github_integration_id(team_id) + if integration_id is None: + raise MissingGitHubIntegrationError + + if not repo_selection.repository_accessible_via_integration(team_id, integration_id, repository): + raise RepositoryNotAccessibleError + + return integration_id diff --git a/products/wizard/backend/logic/runs/store.py b/products/wizard/backend/logic/runs/store.py index 027c26411786..eeb9086ef1aa 100644 --- a/products/wizard/backend/logic/runs/store.py +++ b/products/wizard/backend/logic/runs/store.py @@ -107,6 +107,10 @@ def mark_cancellation_requested(team_id: int, run_id: UUID) -> None: WizardRun.objects.for_team(team_id).filter(id=run_id).update(cancellation_requested_at=timezone.now()) +def cancellation_requested(team_id: int, run_id: UUID) -> bool: + return _get_run_record(team_id, run_id).cancellation_requested_at is not None + + def mark_cancellation_dispatched(team_id: int, run_id: UUID) -> None: WizardRun.objects.for_team(team_id).filter(id=run_id).update(cancellation_dispatched_at=timezone.now()) @@ -122,6 +126,10 @@ def mark_dispatch_succeeded(team_id: int, run_id: UUID, workflow_id: str) -> Non def mark_dispatch_failed(team_id: int, run_id: UUID) -> bool: + """ + Marks a dispatch attempt as failed and returns whether the run has exhausted its retry attempts. + """ + with database_transaction.atomic(): run = WizardRun.objects.for_team(team_id).select_for_update().filter(id=run_id).first() if run is None: diff --git a/products/wizard/backend/logic/workers/commands.py b/products/wizard/backend/logic/workers/commands.py new file mode 100644 index 000000000000..a386399e21d8 --- /dev/null +++ b/products/wizard/backend/logic/workers/commands.py @@ -0,0 +1,142 @@ +import re +import shlex +from uuid import UUID + +from django.conf import settings + +from posthog.utils import get_instance_region + +from products.wizard.backend.facade.validation import is_executable_wizard_version +from products.wizard.backend.logic.artifacts.config import MAX_GIT_DIFF_BYTES +from products.wizard.backend.logic.workers.config import ( + LOCAL_WIZARD_ARCHIVE_PATH, + LOCAL_WIZARD_INSTALL_PATH, + MAX_HANDOFF_BODY_BYTES, + WIZARD_DIFF_OUTPUT_PATH, + WIZARD_PACKAGE_INSTALL_PATH, + WIZARD_TIMEOUT_SECONDS, +) +from products.wizard.backend.logic.workers.errors import InvalidWizardProgramCommandError, InvalidWizardVersionError + +_WIZARD_PROGRAM_COMMAND_PATTERN = re.compile(r"^[a-z0-9]+(?:-[a-z0-9]+)*$") +MAX_SANDBOX_LOG_BYTES = 64 * 1024 + + +def bound_command_output(command: str, *, stdout_limit: int = MAX_SANDBOX_LOG_BYTES, tail: bool = True) -> str: + """ + Wraps a command to limit its stdout and stderr output to a given number of bytes. + Helps prevent runaway commands from filling up the worker's logs and causing OOM errors. + """ + stdout_reader = "tail" if tail else "head" + script = ( + f"( set -o pipefail; {command} ) " + f"> >({stdout_reader} -c {stdout_limit}) 2> >(tail -c {MAX_SANDBOX_LOG_BYTES} >&2); " + 'status=$?; wait; exit "$status"' + ) + return f"bash -c {shlex.quote(script)}" + + +def build_wizard_command( + repository_path: str, + team_id: int, + wizard_version: str, + program_command: tuple[str, ...], + *, + use_local_wizard_source: bool = False, +) -> str: + if not is_executable_wizard_version(wizard_version): + raise InvalidWizardVersionError + + if any(_WIZARD_PROGRAM_COMMAND_PATTERN.fullmatch(argument) is None for argument in program_command): + raise InvalidWizardProgramCommandError + + setup_commands: list[str] = [] + if use_local_wizard_source: + executable = f"node {shlex.quote(f'{LOCAL_WIZARD_INSTALL_PATH}/dist/bin.js')}" + else: + install_path = shlex.quote(WIZARD_PACKAGE_INSTALL_PATH) + package = shlex.quote(f"@posthog/wizard@{wizard_version}") + + setup_commands = [ + f"rm -rf {install_path}", + f"mkdir -p {install_path}", + # removes previous npm configs, creating an empty one + "rm -f /tmp/posthog-wizard-global.npmrc", + "touch /tmp/posthog-wizard-global.npmrc", + ( + # this helps reproducibility by avoiding global configs being inherited + "npm --userconfig=/dev/null --globalconfig=/tmp/posthog-wizard-global.npmrc install " + "--ignore-scripts --no-package-lock --no-save " + f"--prefix {install_path} --registry=https://registry.npmjs.org {package}" + ), + ] + + executable = f"node {install_path}/node_modules/@posthog/wizard/dist/bin.js" + + invocation = [ + executable, + *(shlex.quote(argument) for argument in program_command), + "--headless-DONOTUSE-EXPERIMENTAL", + "--install-dir .", + f"--region {shlex.quote(_wizard_region())}", + f"--project-id {shlex.quote(str(team_id))}", + ] + + if settings.DEBUG: + invocation.append('--base-url "$POSTHOG_API_URL"') + + commands = [*setup_commands, f"cd {shlex.quote(repository_path)}", " ".join(invocation)] + return f"timeout -k 30 {WIZARD_TIMEOUT_SECONDS} sh -c {shlex.quote(' && '.join(commands))}" + + +def build_local_wizard_preparation_command() -> str: + archive_path = shlex.quote(LOCAL_WIZARD_ARCHIVE_PATH) + install_path = shlex.quote(LOCAL_WIZARD_INSTALL_PATH) + return " && ".join( + ( + "apt-get update -qq && apt-get install -y --no-install-recommends build-essential python3", + f"rm -rf {install_path}", + f"mkdir -p {install_path}", + f"tar -xzf {archive_path} -C {install_path}", + f"cd {install_path}", + "corepack enable pnpm", + "HUSKY=0 pnpm install --frozen-lockfile", + "pnpm run prebuild", + "WIZARD_BUILD_NODE_ENV=ci pnpm exec tsdown", + "chmod +x ./dist/bin.js", + "cp -r scripts/** dist", + "rm -f dist/*.no-jest.*", + ) + ) + + +def build_git_diff_command(repository_path: str) -> str: + # Stage the diff to a file and ship at most MAX_GIT_DIFF_BYTES+1 back: the worker reads all + # stdout into memory, so the byte cap is enforced sandbox-side. The sentinel extra byte lets + # the artifact step tell "at the limit" from "oversized"; && preserves git's exit code. + return ( + f"cd {shlex.quote(repository_path)} && git add -N --all && " + f"git diff --binary --no-ext-diff HEAD > {shlex.quote(WIZARD_DIFF_OUTPUT_PATH)} && " + f"head -c {MAX_GIT_DIFF_BYTES + 1} {shlex.quote(WIZARD_DIFF_OUTPUT_PATH)}" + ) + + +def build_sanitize_repository_remote_command(repository_path: str, repository: str) -> str: + remote_url = f"https://github.com/{repository}.git" + return f"git -C {shlex.quote(repository_path)} remote set-url origin {shlex.quote(remote_url)}" + + +def wizard_handoff_output_path(run_id: UUID) -> str: + return f"/tmp/posthog-wizard-handoff-{run_id.hex}.md" + + +def build_read_handoff_command(run_id: UUID) -> str: + return f"head -c {MAX_HANDOFF_BODY_BYTES} {shlex.quote(wizard_handoff_output_path(run_id))}" + + +def pull_request_branch(run_id: UUID) -> str: + return f"posthog/wizard-{run_id.hex[:12]}" + + +def _wizard_region() -> str: + return "eu" if get_instance_region() == "EU" else "us" diff --git a/products/wizard/backend/logic/workers/config.py b/products/wizard/backend/logic/workers/config.py index 9aa28b270dc9..2d5d5b160a8b 100644 --- a/products/wizard/backend/logic/workers/config.py +++ b/products/wizard/backend/logic/workers/config.py @@ -12,6 +12,8 @@ LOCAL_WIZARD_BUILD_TIMEOUT_SECONDS = 15 * 60 LOCAL_WIZARD_ARCHIVE_PATH = "/tmp/posthog-local-wizard.tar.gz" LOCAL_WIZARD_INSTALL_PATH = "/tmp/posthog-local-wizard" +WIZARD_PACKAGE_INSTALL_PATH = "/tmp/posthog-wizard-package" +WIZARD_DIFF_OUTPUT_PATH = "/tmp/posthog-wizard-diff.patch" LOCAL_WIZARD_SOURCE_ENTRIES = ( "README.md", "bin.ts", diff --git a/products/wizard/backend/logic/workers/errors.py b/products/wizard/backend/logic/workers/errors.py new file mode 100644 index 000000000000..34d06bf1a9e5 --- /dev/null +++ b/products/wizard/backend/logic/workers/errors.py @@ -0,0 +1,8 @@ +class InvalidWizardVersionError(ValueError): + def __init__(self) -> None: + super().__init__("Invalid Wizard version") + + +class InvalidWizardProgramCommandError(ValueError): + def __init__(self) -> None: + super().__init__("Invalid Wizard program command") diff --git a/products/wizard/backend/logic/workers/lifecycle.py b/products/wizard/backend/logic/workers/lifecycle.py new file mode 100644 index 000000000000..5af144c16e09 --- /dev/null +++ b/products/wizard/backend/logic/workers/lifecycle.py @@ -0,0 +1,84 @@ +import logging +from uuid import UUID + +from django.core.exceptions import ObjectDoesNotExist +from django.db import DatabaseError + +from products.tasks.backend.facade.sandbox import SandboxCleanupError +from products.wizard.backend.facade.errors import WizardRunNotFoundError +from products.wizard.backend.logic.runs import store as run_store +from products.wizard.backend.logic.runs.errors import WizardWorkerCleanupError +from products.wizard.backend.logic.workers import ( + service as cloud_worker, + store as worker_store, +) +from products.wizard.backend.logic.workers.contracts import WizardWorkerUsageMeasurement +from products.wizard.backend.observability.contracts import WizardWorkerCleanupOutcome +from products.wizard.backend.observability.service import wizard_observability + +logger = logging.getLogger(__name__) + + +def cleanup_worker(team_id: int, run_id: UUID, sandbox_id: str) -> None: + worker_store.mark_cleanup_pending(team_id, run_id) + + try: + usage = cloud_worker.measure_worker_usage(sandbox_id) + + if usage is not None: + _record_worker_usage(team_id, run_id, sandbox_id, usage) + finally: + _destroy_worker(team_id, run_id, sandbox_id) + + _report_worker_usage(team_id, run_id) + + +def _record_worker_usage( + team_id: int, + run_id: UUID, + sandbox_id: str, + usage: WizardWorkerUsageMeasurement, +) -> None: + try: + worker_store.record_usage(team_id, run_id, usage) + except (DatabaseError, ObjectDoesNotExist, ValueError): + logger.exception( + "wizard_worker_usage_recording_failed", + extra={"team_id": team_id, "run_id": str(run_id), "sandbox_id": sandbox_id}, + ) + + +def _destroy_worker(team_id: int, run_id: UUID, sandbox_id: str) -> None: + try: + cloud_worker.destroy_worker(sandbox_id) + except SandboxCleanupError as error: + _record_cleanup_failure(team_id, run_id, sandbox_id) + wizard_observability.worker_cleanup_finished(team_id, run_id, WizardWorkerCleanupOutcome.FAILED) + raise WizardWorkerCleanupError from error + + worker_store.mark_cleaned(team_id, run_id) + wizard_observability.worker_cleanup_finished(team_id, run_id, WizardWorkerCleanupOutcome.SUCCEEDED) + + +def _record_cleanup_failure(team_id: int, run_id: UUID, sandbox_id: str) -> None: + try: + worker_store.mark_cleanup_failed(team_id, run_id) + except DatabaseError: + logger.exception( + "wizard_worker_cleanup_failure_recording_failed", + extra={"team_id": team_id, "run_id": str(run_id), "sandbox_id": sandbox_id}, + ) + + +def _report_worker_usage(team_id: int, run_id: UUID) -> None: + try: + run = run_store.get_run(team_id, run_id) + telemetry = worker_store.get_worker_telemetry(team_id, run_id) + except (DatabaseError, ObjectDoesNotExist, ValueError, WizardRunNotFoundError): + logger.exception( + "wizard_worker_usage_reporting_failed", + extra={"team_id": team_id, "run_id": str(run_id)}, + ) + return + + wizard_observability.worker_usage_recorded(run, telemetry) diff --git a/products/wizard/backend/logic/workers/repository_publisher.py b/products/wizard/backend/logic/workers/repository_publisher.py new file mode 100644 index 000000000000..b6b608da56c1 --- /dev/null +++ b/products/wizard/backend/logic/workers/repository_publisher.py @@ -0,0 +1,370 @@ +import shlex +import base64 +from typing import Any + +from posthog.dataclasses import frozen +from posthog.models.integration import GitHubIntegration, Integration + +from products.tasks.backend.facade.sandbox import SandboxBase, sandbox_repo_path +from products.wizard.backend.logic.workers.commands import bound_command_output +from products.wizard.backend.logic.workers.contracts import RepositoryPullRequest, SignedRepositoryCommit + +GIT_COMMAND_TIMEOUT_SECONDS = 60 +GITHUB_MUTATION_TIMEOUT_SECONDS = 60 +MAX_COMMIT_PAYLOAD_BYTES = 35 * 1024 * 1024 + +_CREATE_COMMIT_MUTATION = """mutation($input: CreateCommitOnBranchInput!) { + createCommitOnBranch(input: $input) { commit { oid url } } +}""" +_CREATE_BRANCH_MUTATION = """mutation($input: CreateRefInput!) { + createRef(input: $input) { ref { name } } +}""" +_BRANCH_TIP_QUERY = """query($owner: String!, $name: String!, $qualifiedName: String!) { + repository(owner: $owner, name: $name) { + id + ref(qualifiedName: $qualifiedName) { + target { ... on Commit { oid } } + } + } +}""" + + +class RepositoryPublishingError(Exception): + pass + + +@frozen +class _RepositoryName: + owner: str + name: str + + +@frozen +class _StagedChanges: + additions: list[tuple[str, str]] + deletions: list[str] + + +def create_signed_commit( + sandbox: SandboxBase, + *, + team_id: int, + integration_id: int, + repository: str, + branch: str, + message: str, + source: str, +) -> SignedRepositoryCommit: + """Create a signed commit through GitHub without requiring a local signing key.""" + repository_path = sandbox_repo_path(repository) + repository_name = _repository_parts(repository) + github = _github_integration(team_id, integration_id, source) + + _stage_all(sandbox, repository_path) + head_sha = _head_sha(sandbox, repository_path) + + repository_id, existing_tip = _repository_and_branch_tip( + github, repository_name.owner, repository_name.name, branch + ) + if existing_tip is None: + branch_existed = False + _create_branch(github, repository_id, branch, head_sha) + branch_tip = head_sha + else: + branch_existed = True + _fetch_branch(sandbox, repository_path, branch) + branch_tip = existing_tip + + changes = _staged_changes(sandbox, repository_path, branch_tip) + if not changes.additions and not changes.deletions: + if branch_existed: + return SignedRepositoryCommit(repository=repository, branch=branch, commit_shas=(branch_tip,)) + raise RepositoryPublishingError("No staged changes to commit.") + + _assert_payload_size(changes) + commit_sha = _commit_staged_changes(github, repository, branch, branch_tip, message, changes) + return SignedRepositoryCommit(repository=repository, branch=branch, commit_shas=(commit_sha,)) + + +def _stage_all(sandbox: SandboxBase, repository_path: str) -> None: + _run_git(sandbox, repository_path, "add --all", "staging") + + +def _head_sha(sandbox: SandboxBase, repository_path: str) -> str: + head_sha = _run_git(sandbox, repository_path, "rev-parse HEAD", "HEAD resolution").strip() + if not head_sha: + raise RepositoryPublishingError("Repository checkout has no HEAD commit.") + return head_sha + + +def _fetch_branch(sandbox: SandboxBase, repository_path: str, branch: str) -> None: + _run_git(sandbox, repository_path, f"fetch --no-tags origin {shlex.quote(branch)}", "branch fetch") + + +def _staged_changes(sandbox: SandboxBase, repository_path: str, branch_tip: str) -> _StagedChanges: + output = _run_git( + sandbox, + repository_path, + f"diff --cached -z --no-renames --name-status {shlex.quote(branch_tip)}", + "staged change detection", + ) + + additions: list[tuple[str, str]] = [] + deletions: list[str] = [] + tokens = output.split("\0") + remaining_bytes = MAX_COMMIT_PAYLOAD_BYTES + for index in range(0, len(tokens) - 1, 2): + status, path = tokens[index], tokens[index + 1] + if not status: + continue + if status.startswith("D"): + remaining_bytes -= len(path.encode("utf-8")) + 16 + deletions.append(path) + else: + remaining_bytes -= len(path.encode("utf-8")) + 32 + if remaining_bytes < 0: + raise RepositoryPublishingError("Staged changes exceed the GitHub commit payload limit.") + contents = _staged_file_contents(sandbox, repository_path, path, remaining_bytes) + remaining_bytes -= len(contents) + additions.append((path, contents)) + if remaining_bytes < 0: + raise RepositoryPublishingError("Staged changes exceed the GitHub commit payload limit.") + return _StagedChanges(additions=additions, deletions=deletions) + + +def _staged_file_contents(sandbox: SandboxBase, repository_path: str, path: str, remaining_bytes: int) -> str: + contents = _run_git( + sandbox, + repository_path, + f"show {shlex.quote(f':{path}')} | base64 -w 0", + f"staged file read for '{path}'", + max_output_bytes=remaining_bytes, + ).strip() + try: + base64.b64decode(contents, validate=True) + except ValueError as error: + raise RepositoryPublishingError(f"Could not read staged file '{path}'.") from error + return contents + + +def _assert_payload_size(changes: _StagedChanges) -> None: + payload_bytes = sum(len(path) + len(contents) + 32 for path, contents in changes.additions) + payload_bytes += sum(len(path) + 16 for path in changes.deletions) + if payload_bytes > MAX_COMMIT_PAYLOAD_BYTES: + raise RepositoryPublishingError("Staged changes exceed the GitHub commit payload limit.") + + +def _commit_staged_changes( + github: GitHubIntegration, + repository: str, + branch: str, + branch_tip: str, + message: str, + changes: _StagedChanges, +) -> str: + headline, _, body = message.partition("\n") + payload = _graphql_request( + github, + _CREATE_COMMIT_MUTATION, + { + "input": { + "branch": {"repositoryNameWithOwner": repository, "branchName": branch}, + "expectedHeadOid": branch_tip, + "message": {"headline": headline.strip(), "body": body.strip()}, + "fileChanges": { + "additions": [{"path": path, "contents": contents} for path, contents in changes.additions], + "deletions": [{"path": path} for path in changes.deletions], + }, + } + }, + "signed commit", + ) + + commit = _payload_dict(payload, "createCommitOnBranch", "commit") + commit_sha = commit.get("oid") if commit else None + if not isinstance(commit_sha, str): + raise RepositoryPublishingError("GitHub did not return the signed commit.") + return commit_sha + + +def _github_integration(team_id: int, integration_id: int, source: str) -> GitHubIntegration: + integration = Integration.objects.filter(team_id=team_id, id=integration_id, kind="github").first() + if integration is None: + raise RepositoryPublishingError("GitHub integration is unavailable.") + return GitHubIntegration(integration, source=source) + + +def _repository_and_branch_tip( + github: GitHubIntegration, owner: str, repository_name: str, branch: str +) -> tuple[str, str | None]: + payload = _graphql_request( + github, + _BRANCH_TIP_QUERY, + {"owner": owner, "name": repository_name, "qualifiedName": f"refs/heads/{branch}"}, + "branch tip query", + ) + + repository_data = _payload_dict(payload, "repository") + repository_id = repository_data.get("id") if repository_data else None + if not isinstance(repository_id, str): + raise RepositoryPublishingError("GitHub repository is unavailable.") + + ref_data = repository_data.get("ref") if repository_data else None + target = ref_data.get("target") if isinstance(ref_data, dict) else None + tip = target.get("oid") if isinstance(target, dict) else None + return repository_id, tip if isinstance(tip, str) else None + + +def _create_branch(github: GitHubIntegration, repository_id: str, branch: str, head_sha: str) -> None: + _graphql_request( + github, + _CREATE_BRANCH_MUTATION, + {"input": {"repositoryId": repository_id, "name": f"refs/heads/{branch}", "oid": head_sha}}, + "branch creation", + ) + + +def _graphql_request( + github: GitHubIntegration, query: str, variables: dict[str, Any], endpoint_label: str +) -> dict[str, Any]: + response = github.api_request( + "POST", + "/graphql", + endpoint=endpoint_label, + json_body={"query": query, "variables": variables}, + timeout=GITHUB_MUTATION_TIMEOUT_SECONDS, + ) + if response.status_code != 200: + raise RepositoryPublishingError(f"GitHub {endpoint_label} failed with status code {response.status_code}.") + try: + payload: Any = response.json() + except ValueError as error: + raise RepositoryPublishingError(f"GitHub {endpoint_label} returned an invalid response.") from error + if not isinstance(payload, dict): + raise RepositoryPublishingError(f"GitHub {endpoint_label} returned an invalid response.") + + errors = payload.get("errors") + if isinstance(errors, list) and errors: + messages = ", ".join( + entry["message"] for entry in errors if isinstance(entry, dict) and isinstance(entry.get("message"), str) + ) + raise RepositoryPublishingError(f"GitHub {endpoint_label} failed: {messages}") + return payload + + +def _payload_dict(payload: dict[str, Any], *keys: str) -> dict[str, Any] | None: + value: object = payload + for key in ("data", *keys): + if not isinstance(value, dict): + return None + value = value.get(key) + return value if isinstance(value, dict) else None + + +def create_pull_request( + *, + team_id: int, + integration_id: int, + repository: str, + head_branch: str, + title: str, + body: str, + source: str, +) -> RepositoryPullRequest: + repository_name = _repository_parts(repository) + integration = Integration.objects.filter(team_id=team_id, id=integration_id, kind="github").first() + if integration is None: + raise RepositoryPublishingError("GitHub integration is unavailable.") + + github = GitHubIntegration(integration, source=source) + if github.organization().casefold() != repository_name.owner.casefold(): + raise RepositoryPublishingError("GitHub integration does not own the repository.") + + base_branch = github.get_default_branch(repository_name.name) + existing = _matching_pull_request( + github.list_pull_requests(repository_name.name), + repository, + head_branch, + base_branch, + ) + if existing is not None: + return existing + + created = github.create_pull_request(repository_name.name, title, body, head_branch, base_branch) + return _created_pull_request(created, repository, head_branch, base_branch) + + +def _matching_pull_request( + payload: object, + repository: str, + head_branch: str, + base_branch: str, +) -> RepositoryPullRequest | None: + if not isinstance(payload, dict) or payload.get("success") is not True: + return None + pull_requests = payload.get("pull_requests") + if not isinstance(pull_requests, list): + return None + for value in pull_requests: + if not isinstance(value, dict) or value.get("head_branch") != head_branch: + continue + number = value.get("number") + url = value.get("url") + pull_request_base_branch = value.get("base_branch") + if isinstance(number, int) and isinstance(url, str): + return RepositoryPullRequest( + repository=repository, + number=number, + url=url, + head_branch=head_branch, + base_branch=pull_request_base_branch if isinstance(pull_request_base_branch, str) else base_branch, + ) + return None + + +def _created_pull_request( + payload: object, + repository: str, + head_branch: str, + base_branch: str, +) -> RepositoryPullRequest: + if not isinstance(payload, dict) or payload.get("success") is not True: + raise RepositoryPublishingError("GitHub did not create the pull request.") + number = payload.get("pr_number") + url = payload.get("pr_url") + if not isinstance(number, int) or not isinstance(url, str): + raise RepositoryPublishingError("GitHub returned an invalid pull request.") + return RepositoryPullRequest( + repository=repository, + number=number, + url=url, + head_branch=head_branch, + base_branch=base_branch, + ) + + +def _repository_parts(repository: str) -> _RepositoryName: + parts = repository.split("/") + if len(parts) != 2 or not all(parts): + raise RepositoryPublishingError("Repository must use owner/name format.") + return _RepositoryName(owner=parts[0], name=parts[1]) + + +def _run_git( + sandbox: SandboxBase, + repository_path: str, + arguments: str, + stage: str, + *, + max_output_bytes: int = MAX_COMMIT_PAYLOAD_BYTES, +) -> str: + command = f"git -C {shlex.quote(repository_path)} {arguments}" + result = sandbox.execute( + bound_command_output(command, stdout_limit=max_output_bytes + 1, tail=False), + timeout_seconds=GIT_COMMAND_TIMEOUT_SECONDS, + ) + if len(result.stdout.encode("utf-8")) > max_output_bytes: + raise RepositoryPublishingError("Staged changes exceed the GitHub commit payload limit.") + if result.exit_code != 0: + raise RepositoryPublishingError(f"Repository {stage} failed with exit code {result.exit_code}.") + return result.stdout diff --git a/products/wizard/backend/logic/workers/service.py b/products/wizard/backend/logic/workers/service.py new file mode 100644 index 000000000000..378ceb93abdc --- /dev/null +++ b/products/wizard/backend/logic/workers/service.py @@ -0,0 +1,372 @@ +import logging +from datetime import timedelta +from pathlib import Path +from uuid import UUID + +from django.conf import settings +from django.utils import timezone + +from posthog.dataclasses import frozen +from posthog.models.user import User +from posthog.temporal.oauth import create_wizard_oauth_access_token_for_user + +from products.tasks.backend.facade.repo_selection import get_github_token +from products.tasks.backend.facade.sandbox import ( + SandboxBase, + SandboxConfig, + SandboxExecutionError, + SandboxNotFoundError, + SandboxTimeoutError, + get_sandbox_class, + sandbox_repo_path, +) +from products.wizard.backend.logic.artifacts.config import ( + PULL_REQUEST_BODY, + PULL_REQUEST_COMMIT_MESSAGE, + PULL_REQUEST_TITLE, +) +from products.wizard.backend.logic.workers.commands import ( + bound_command_output, + build_git_diff_command, + build_local_wizard_preparation_command, + build_read_handoff_command, + build_sanitize_repository_remote_command, + build_wizard_command, + pull_request_branch, + wizard_handoff_output_path, +) +from products.wizard.backend.logic.workers.config import ( + LOCAL_WIZARD_ARCHIVE_PATH, + LOCAL_WIZARD_BUILD_TIMEOUT_SECONDS, + SANDBOX_CPU_CORES, + SANDBOX_DISK_SIZE_GB, + SANDBOX_EXECUTION_TIMEOUT_SECONDS, + SANDBOX_MEMORY_GB, + SANDBOX_TEMPLATE_BASE, + SANDBOX_TTL_SECONDS, + WIZARD_ERROR_DETAIL_LENGTH, + WIZARD_TIMEOUT_EXIT_CODE, +) +from products.wizard.backend.logic.workers.contracts import ( + RepositoryPullRequest, + WizardWorkerProvisioning, + WizardWorkerResourceUsage, + WizardWorkerUsageMeasurement, +) +from products.wizard.backend.logic.workers.local_package import build_local_wizard_source_archive +from products.wizard.backend.logic.workers.wizard_error_output import stderr_to_wizard_error_code +from products.wizard.backend.observability.service import wizard_observability + +from .repository_publisher import RepositoryPublishingError, create_pull_request, create_signed_commit + +logger = logging.getLogger(__name__) + + +@frozen +class WizardWorkerProvisionRequest: + team_id: int + created_by_id: int + run_id: UUID + + +@frozen +class GitRepositoryCloneRequest: + sandbox_id: str + github_integration_id: int + repository: str + + +@frozen +class WizardExecutionRequest: + sandbox_id: str + workspace_path: str + team_id: int + wizard_version: str + program_command: tuple[str, ...] + use_local_wizard_source: bool = False + + +@frozen +class GitRepositoryHandoffRequest: + team_id: int + run_id: UUID + sandbox_id: str + workspace_path: str + github_integration_id: int + repository: str + + +@frozen +class WizardWorkerResult: + diff: bytes + pull_request: RepositoryPullRequest | None + + +class WizardWorkerExecutionError(Exception): + def __init__( + self, + stage: str, + exit_code: int, + detail: str | None = None, + wizard_error_code: str | None = None, + ) -> None: + self.stage = stage + self.exit_code = exit_code + self.detail = detail + self.wizard_error_code = wizard_error_code + + message = f"Wizard Worker {stage} failed with exit code {exit_code}." + if detail: + message = f"{message}\n{detail}" + + super().__init__(message) + + +class WizardWorkerTimeoutError(Exception): + pass + + +def provision_wizard_worker(request: WizardWorkerProvisionRequest) -> WizardWorkerProvisioning: + user = User.objects.get(id=request.created_by_id) + wizard_token = create_wizard_oauth_access_token_for_user(user, request.team_id) + + config = _build_sandbox_config(request, wizard_token) + sandbox = get_sandbox_class().create(config) + provisioned_at = timezone.now() + + return WizardWorkerProvisioning( + sandbox_id=sandbox.id, + resource_usage=WizardWorkerResourceUsage( + cpu_cores=config.cpu_cores, + memory_gb=config.memory_gb, + disk_size_gb=config.disk_size_gb, + ttl_seconds=config.ttl_seconds, + ttl_expires_at=provisioned_at + timedelta(seconds=config.ttl_seconds), + ), + ) + + +def clone_repository(request: GitRepositoryCloneRequest) -> str: + sandbox = get_sandbox_class().get_by_id(request.sandbox_id) + github_token = get_github_token(request.github_integration_id) or "" + clone_result = sandbox.clone_repository(request.repository, github_token=github_token, shallow=True) + + _raise_for_failure( + "repository clone", + clone_result.exit_code, + stdout=clone_result.stdout, + stderr=clone_result.stderr, + sensitive_values=(github_token,), + ) + + workspace_path = sandbox_repo_path(request.repository) + sanitize_result = sandbox.execute( + build_sanitize_repository_remote_command(workspace_path, request.repository), + timeout_seconds=60, + ) + _raise_for_failure( + "repository credential cleanup", + sanitize_result.exit_code, + stdout=sanitize_result.stdout, + stderr=sanitize_result.stderr, + sensitive_values=(github_token,), + ) + + return workspace_path + + +def prepare_local_wizard(sandbox_id: str, source_root: Path) -> None: + archive = build_local_wizard_source_archive(source_root) + + try: + sandbox = get_sandbox_class().get_by_id(sandbox_id) + + upload_result = sandbox.write_file(LOCAL_WIZARD_ARCHIVE_PATH, archive) + _raise_for_failure( + "local Wizard source upload", + upload_result.exit_code, + stdout=upload_result.stdout, + stderr=upload_result.stderr, + ) + + build_result = sandbox.execute( + build_local_wizard_preparation_command(), + timeout_seconds=LOCAL_WIZARD_BUILD_TIMEOUT_SECONDS, + ) + _raise_for_failure( + "local Wizard build", + build_result.exit_code, + stdout=build_result.stdout, + stderr=build_result.stderr, + ) + except (SandboxExecutionError, SandboxNotFoundError, SandboxTimeoutError) as error: + raise WizardWorkerExecutionError("local Wizard preparation", 1, str(error)) from error + + +def execute_wizard(request: WizardExecutionRequest) -> None: + sandbox = get_sandbox_class().get_by_id(request.sandbox_id) + + wizard_result = sandbox.execute( + bound_command_output( + build_wizard_command( + request.workspace_path, + request.team_id, + request.wizard_version, + request.program_command, + use_local_wizard_source=request.use_local_wizard_source, + ) + ), + timeout_seconds=SANDBOX_EXECUTION_TIMEOUT_SECONDS, + ) + + if wizard_result.exit_code == WIZARD_TIMEOUT_EXIT_CODE: + raise WizardWorkerTimeoutError + + _raise_for_failure( + "execution", + wizard_result.exit_code, + stdout=wizard_result.stdout, + stderr=wizard_result.stderr, + ) + + +def create_git_repository_handoff(request: GitRepositoryHandoffRequest) -> WizardWorkerResult: + sandbox = get_sandbox_class().get_by_id(request.sandbox_id) + + diff_result = sandbox.execute( + build_git_diff_command(request.workspace_path), + timeout_seconds=60, + ) + + _raise_for_failure( + "diff capture", + diff_result.exit_code, + stdout=diff_result.stdout, + stderr=diff_result.stderr, + ) + + diff = diff_result.stdout.encode("utf-8") + + if not diff: + return WizardWorkerResult(diff=diff, pull_request=None) + + branch = pull_request_branch(request.run_id) + handoff_body = _read_handoff_body(sandbox, request.run_id) + + if handoff_body is None: + wizard_observability.handoff_body_fallback(request.team_id, request.run_id) + handoff_body = PULL_REQUEST_BODY + + try: + create_signed_commit( + sandbox, + team_id=request.team_id, + integration_id=request.github_integration_id, + repository=request.repository, + branch=branch, + message=PULL_REQUEST_COMMIT_MESSAGE, + source="wizard", + ) + + pull_request = create_pull_request( + team_id=request.team_id, + integration_id=request.github_integration_id, + repository=request.repository, + head_branch=branch, + title=PULL_REQUEST_TITLE, + body=handoff_body, + source="wizard", + ) + + except RepositoryPublishingError as error: + raise WizardWorkerExecutionError("publishing", 1, str(error)) from error + + return WizardWorkerResult(diff=diff, pull_request=pull_request) + + +def destroy_worker(sandbox_id: str) -> None: + try: + sandbox = get_sandbox_class().get_by_id(sandbox_id) + except SandboxNotFoundError: + return + sandbox.destroy() + + +def measure_worker_usage(sandbox_id: str) -> WizardWorkerUsageMeasurement | None: + try: + sandbox = get_sandbox_class().get_by_id(sandbox_id) + cpu_usage_usec = sandbox.read_cpu_usage_usec() + billed_cpu_usage_usec = sandbox.read_billed_cpu_usage_usec() + except (SandboxExecutionError, SandboxNotFoundError, SandboxTimeoutError): + logger.exception("wizard_worker_usage_measurement_failed", extra={"sandbox_id": sandbox_id}) + return None + + if cpu_usage_usec is None and billed_cpu_usage_usec is None: + return None + + return WizardWorkerUsageMeasurement( + cpu_usage_usec=cpu_usage_usec, + billed_cpu_usage_usec=billed_cpu_usage_usec, + measured_at=timezone.now(), + ) + + +def _raise_for_failure( + stage: str, + exit_code: int, + *, + stdout: str = "", + stderr: str = "", + sensitive_values: tuple[str, ...] = (), +) -> None: + if exit_code != 0: + raise WizardWorkerExecutionError( + stage, + exit_code, + _failure_detail(stdout, stderr, sensitive_values), + stderr_to_wizard_error_code(stderr), + ) + + +def _failure_detail(stdout: str, stderr: str, sensitive_values: tuple[str, ...]) -> str | None: + output = "\n".join(value for value in (stdout.strip(), stderr.strip()) if value) + for sensitive_value in sensitive_values: + if sensitive_value: + output = output.replace(sensitive_value, "[REDACTED]") + return output[-WIZARD_ERROR_DETAIL_LENGTH:] or None + + +def _read_handoff_body(sandbox: SandboxBase, run_id: UUID) -> str | None: + handoff_result = sandbox.execute(build_read_handoff_command(run_id), timeout_seconds=10) + if handoff_result.exit_code != 0: + return None + + handoff_body = handoff_result.stdout.strip() + return handoff_body or None + + +def _build_sandbox_config(request: WizardWorkerProvisionRequest, wizard_token: str) -> SandboxConfig: + environment_variables = { + "POSTHOG_API_URL": settings.SANDBOX_API_URL or settings.SITE_URL, + "POSTHOG_PROJECT_ID": str(request.team_id), + "POSTHOG_WIZARD_API_KEY": wizard_token, + "POSTHOG_HANDOFF_OUTPUT_PATH": wizard_handoff_output_path(request.run_id), + } + if settings.DEBUG and settings.SANDBOX_MCP_URL: + environment_variables["MCP_URL"] = settings.SANDBOX_MCP_URL + + return SandboxConfig( + name=f"wizard-{request.run_id}", + template=SANDBOX_TEMPLATE_BASE, + default_execution_timeout_seconds=SANDBOX_EXECUTION_TIMEOUT_SECONDS, + ttl_seconds=SANDBOX_TTL_SECONDS, + memory_gb=SANDBOX_MEMORY_GB, + cpu_cores=SANDBOX_CPU_CORES, + disk_size_gb=SANDBOX_DISK_SIZE_GB, + environment_variables=environment_variables, + metadata={ + "purpose": "wizard_run", + "team_id": str(request.team_id), + "wizard_run_id": str(request.run_id), + }, + ) diff --git a/products/wizard/backend/observability/config.py b/products/wizard/backend/observability/config.py index 0c6d49214557..63d3e9950584 100644 --- a/products/wizard/backend/observability/config.py +++ b/products/wizard/backend/observability/config.py @@ -1,5 +1,3 @@ -WIZARD_ANALYTICS_TASK = "wizard.capture_run_event" - WIZARD_RUN_CREATED_EVENT = "wizard run created" WIZARD_RUN_DISPATCH_FINISHED_EVENT = "wizard run dispatch finished" WIZARD_RUN_STAGE_ENTERED_EVENT = "wizard run stage entered" diff --git a/products/wizard/backend/observability/events.py b/products/wizard/backend/observability/events.py index 5ba58c73b42e..140620ff67ad 100644 --- a/products/wizard/backend/observability/events.py +++ b/products/wizard/backend/observability/events.py @@ -1,11 +1,11 @@ from uuid import NAMESPACE_URL, uuid5 -from celery import current_app as celery_app +from posthog.models import User +from posthog.ph_client import ph_background_capture from products.wizard.backend.facade.contracts import WizardRunDTO from products.wizard.backend.facade.enums import WizardRunStage from products.wizard.backend.observability.config import ( - WIZARD_ANALYTICS_TASK, WIZARD_PULL_REQUEST_CREATED_EVENT, WIZARD_RUN_CREATED_EVENT, WIZARD_RUN_DISPATCH_FINISHED_EVENT, @@ -90,7 +90,14 @@ def _enqueue_run_event( if properties is not None: event_properties.update(properties) - celery_app.signature( - WIZARD_ANALYTICS_TASK, - args=[run.team_id, run.created_by_id, str(run.id), event, str(event_uuid), event_properties], - ).apply_async() + distinct_id = str(run.id) + if run.created_by_id is not None: + user_distinct_id = User.objects.filter(id=run.created_by_id).values_list("distinct_id", flat=True).first() + if user_distinct_id is not None: + distinct_id = user_distinct_id + ph_background_capture()( + distinct_id=distinct_id, + event=event, + properties={**event_properties, "team_id": run.team_id, "wizard_run_id": str(run.id)}, + uuid=str(event_uuid), + ) diff --git a/products/wizard/backend/observability/metrics.py b/products/wizard/backend/observability/metrics.py index 965a2eb0e10b..239e32e90b3d 100644 --- a/products/wizard/backend/observability/metrics.py +++ b/products/wizard/backend/observability/metrics.py @@ -1,7 +1,12 @@ from prometheus_client import Counter, Gauge, Histogram from products.wizard.backend.facade.contracts import WizardRunArtifactDTO, WizardRunDTO -from products.wizard.backend.facade.enums import WizardRunEnvironment, WizardRunStage, WizardRunStatus +from products.wizard.backend.facade.enums import ( + WizardRunEnvironment, + WizardRunErrorCode, + WizardRunStage, + WizardRunStatus, +) from products.wizard.backend.observability.config import ( WIZARD_RUN_DURATION_BUCKETS, WIZARD_WORKER_CPU_SECONDS_BUCKETS, @@ -157,8 +162,22 @@ def report_run_status_changed(previous: WizardRunDTO, current: WizardRunDTO) -> WIZARD_CLOUD_RUNS_ACTIVE.labels(status=current.status.value, stage=current.stage.value).inc() +# Known error codes are a fixed set; the stored code can also be an open PHW_* value +# from a pinned Wizard version. Collapse anything outside the set to keep the Prometheus +# label bounded — the raw code stays in the run row and the log line. +_KNOWN_ERROR_CODES: frozenset[str] = frozenset(code.value for code in WizardRunErrorCode) + + +def _bounded_error_code(error_code: str | None) -> str: + if error_code is None: + return "none" + if error_code in _KNOWN_ERROR_CODES: + return error_code + return "other" + + def report_run_finished(run: WizardRunDTO, failure_stage: WizardRunStage | None) -> None: - error_code = run.error_code or "none" + error_code = _bounded_error_code(run.error_code) WIZARD_RUNS_FINISHED_TOTAL.labels( environment=run.environment.value, diff --git a/products/wizard/backend/observability/service.py b/products/wizard/backend/observability/service.py index 5281af706fa9..fe1d0aadb581 100644 --- a/products/wizard/backend/observability/service.py +++ b/products/wizard/backend/observability/service.py @@ -3,7 +3,7 @@ from functools import partial from uuid import UUID -from kombu.exceptions import OperationalError +from django.db import DatabaseError from products.wizard.backend.facade.contracts import WizardRunArtifactDTO, WizardRunDTO, WizardRunPullRequestArtifactDTO from products.wizard.backend.facade.enums import WizardRunStatus @@ -41,7 +41,7 @@ def _observe( if event is not None: try: event() - except OperationalError: + except (DatabaseError, ValueError): logger.exception(event_error or f"{name}_event_failed", extra=context) if level is not None: diff --git a/products/wizard/backend/tasks/__init__.py b/products/wizard/backend/tasks/__init__.py index b0ba1add2138..385a5c3da34b 100644 --- a/products/wizard/backend/tasks/__init__.py +++ b/products/wizard/backend/tasks/__init__.py @@ -1,3 +1,5 @@ from products.wizard.backend.tasks.tasks import sync_wizard_event_definitions -__all__ = ["sync_wizard_event_definitions"] +__all__ = [ + "sync_wizard_event_definitions", +] diff --git a/products/wizard/backend/tasks/config.py b/products/wizard/backend/tasks/config.py new file mode 100644 index 000000000000..9af2c712ec3d --- /dev/null +++ b/products/wizard/backend/tasks/config.py @@ -0,0 +1 @@ +RECONCILE_WIZARD_RUNS_TASK = "wizard.reconcile_runs" diff --git a/products/wizard/backend/tasks/tasks.py b/products/wizard/backend/tasks/tasks.py index e1332c8b63e5..57967d958627 100644 --- a/products/wizard/backend/tasks/tasks.py +++ b/products/wizard/backend/tasks/tasks.py @@ -1,6 +1,8 @@ import logging from typing import Any +from django.db import DatabaseError + from celery import shared_task from celery.app.task import Task @@ -10,7 +12,7 @@ PlaceholderEventDefinition, create_placeholder_event_definitions, ) -from products.wizard.backend.facade.enums import RunPhase +from products.wizard.backend.facade.enums import WizardSessionRunPhase from products.wizard.backend.models import WizardSession logger = logging.getLogger(__name__) @@ -61,14 +63,14 @@ def _planned_event_definitions(event_plan: dict[str, Any] | None) -> list[Placeh def sync_wizard_event_definitions(self: Task, team_id: int, session_id: str) -> None: try: session = WizardSession.objects.for_team(team_id).filter(session_id=session_id).first() - if session is None or session.run_phase != RunPhase.COMPLETED.value: + if session is None or session.run_phase != WizardSessionRunPhase.COMPLETED.value: return create_placeholder_event_definitions( team_id=team_id, definitions=_planned_event_definitions(session.event_plan), ) - except Exception as error: + except DatabaseError as error: countdown = min(2**self.request.retries, 60) logger.warning( "Retrying wizard event definition sync", diff --git a/products/wizard/backend/temporal/__init__.py b/products/wizard/backend/temporal/__init__.py new file mode 100644 index 000000000000..6b99c2ef09ab --- /dev/null +++ b/products/wizard/backend/temporal/__init__.py @@ -0,0 +1,18 @@ +from products.wizard.backend.temporal.activities.execution import execute_wizard +from products.wizard.backend.temporal.activities.handoff import create_run_artifacts +from products.wizard.backend.temporal.activities.lifecycle import finalize_run +from products.wizard.backend.temporal.activities.local_package import prepare_local_wizard +from products.wizard.backend.temporal.activities.workspace import clone_repository, destroy_worker, provision_worker +from products.wizard.backend.temporal.workflows.execute_run import ExecuteWizardRunWorkflow + +WORKFLOWS = [ExecuteWizardRunWorkflow] + +ACTIVITIES = [ + provision_worker, + prepare_local_wizard, + clone_repository, + execute_wizard, + create_run_artifacts, + destroy_worker, + finalize_run, +] diff --git a/products/wizard/backend/temporal/activities/errors.py b/products/wizard/backend/temporal/activities/errors.py new file mode 100644 index 000000000000..2f41f89415e6 --- /dev/null +++ b/products/wizard/backend/temporal/activities/errors.py @@ -0,0 +1,4 @@ +WIZARD_REPOSITORY_ACCESS_ERROR_TYPE = "WizardRepositoryAccessError" +WIZARD_RUN_CONFIGURATION_ERROR_TYPE = "WizardRunConfigurationError" +WIZARD_WORKER_EXECUTION_ERROR_TYPE = "WizardWorkerExecutionError" +WIZARD_WORKER_TIMEOUT_ERROR_TYPE = "WizardWorkerTimeoutError" diff --git a/products/wizard/backend/temporal/activities/execution.py b/products/wizard/backend/temporal/activities/execution.py new file mode 100644 index 000000000000..3c7981c6a2d6 --- /dev/null +++ b/products/wizard/backend/temporal/activities/execution.py @@ -0,0 +1,44 @@ +from temporalio import activity +from temporalio.exceptions import ApplicationError + +from posthog.temporal.common.utils import asyncify + +from products.wizard.backend.facade import api as wizard_facade +from products.wizard.backend.facade.enums import WizardRunStage +from products.wizard.backend.logic.workers import service as cloud_worker +from products.wizard.backend.logic.workers.service import WizardExecutionRequest +from products.wizard.backend.temporal.activities.errors import ( + WIZARD_WORKER_EXECUTION_ERROR_TYPE, + WIZARD_WORKER_TIMEOUT_ERROR_TYPE, +) +from products.wizard.backend.temporal.contracts import PreparedGitRepositoryWorkspace + + +@activity.defn(name="wizard_execute") +@asyncify +def execute_wizard(input: PreparedGitRepositoryWorkspace) -> None: + wizard_facade.update_run_stage(input.team_id, input.run_id, WizardRunStage.EXECUTING_WIZARD) + run = wizard_facade.get_run(input.team_id, input.run_id) + try: + cloud_worker.execute_wizard( + WizardExecutionRequest( + sandbox_id=input.sandbox_id, + workspace_path=input.root_path, + team_id=input.team_id, + wizard_version=run.program.wizard_version, + program_command=run.program.command, + use_local_wizard_source=input.use_local_wizard_source, + ) + ) + except cloud_worker.WizardWorkerTimeoutError as error: + raise ApplicationError( + "Wizard Worker timed out.", + type=WIZARD_WORKER_TIMEOUT_ERROR_TYPE, + non_retryable=True, + ) from error + except cloud_worker.WizardWorkerExecutionError as error: + raise ApplicationError( + str(error), + type=error.wizard_error_code or WIZARD_WORKER_EXECUTION_ERROR_TYPE, + non_retryable=True, + ) from error diff --git a/products/wizard/backend/temporal/activities/handoff.py b/products/wizard/backend/temporal/activities/handoff.py new file mode 100644 index 000000000000..bc9d185f0ddb --- /dev/null +++ b/products/wizard/backend/temporal/activities/handoff.py @@ -0,0 +1,52 @@ +from temporalio import activity +from temporalio.exceptions import ApplicationError + +from posthog.temporal.common.utils import asyncify + +from products.wizard.backend.facade import api as wizard_facade +from products.wizard.backend.facade.contracts import CreatePullRequestArtifactInput +from products.wizard.backend.facade.enums import WizardRunStage, WizardRunStatus +from products.wizard.backend.logic.workers import service as cloud_worker +from products.wizard.backend.logic.workers.service import GitRepositoryHandoffRequest +from products.wizard.backend.temporal.activities.errors import WIZARD_WORKER_EXECUTION_ERROR_TYPE +from products.wizard.backend.temporal.activities.lifecycle import transition_cloud_run +from products.wizard.backend.temporal.contracts import PreparedGitRepositoryWorkspace + + +@activity.defn(name="wizard_create_run_artifacts") +@asyncify +def create_run_artifacts(input: PreparedGitRepositoryWorkspace) -> None: + wizard_facade.update_run_stage(input.team_id, input.run_id, WizardRunStage.CREATING_ARTIFACTS) + try: + result = cloud_worker.create_git_repository_handoff( + GitRepositoryHandoffRequest( + team_id=input.team_id, + run_id=input.run_id, + sandbox_id=input.sandbox_id, + workspace_path=input.root_path, + github_integration_id=input.github_integration_id, + repository=input.repository, + ) + ) + except cloud_worker.WizardWorkerExecutionError as error: + raise ApplicationError( + str(error), + type=WIZARD_WORKER_EXECUTION_ERROR_TYPE, + non_retryable=True, + ) from error + + wizard_facade.create_git_diff_artifact(input.team_id, input.run_id, result.diff) + if result.pull_request is not None: + wizard_facade.create_pull_request_artifact( + CreatePullRequestArtifactInput( + team_id=input.team_id, + run_id=input.run_id, + url=result.pull_request.url, + number=result.pull_request.number, + repository=result.pull_request.repository, + head_branch=result.pull_request.head_branch, + base_branch=result.pull_request.base_branch, + ) + ) + + transition_cloud_run(input.team_id, input.run_id, WizardRunStatus.COMPLETED) diff --git a/products/wizard/backend/temporal/activities/lifecycle.py b/products/wizard/backend/temporal/activities/lifecycle.py new file mode 100644 index 000000000000..a2e04e9c0afd --- /dev/null +++ b/products/wizard/backend/temporal/activities/lifecycle.py @@ -0,0 +1,84 @@ +from uuid import UUID + +from temporalio import activity + +from posthog.temporal.common.utils import asyncify + +from products.wizard.backend.facade import api as wizard_facade +from products.wizard.backend.facade.contracts import WizardRunDTO +from products.wizard.backend.facade.enums import WizardRunEnvironment, WizardRunStatus +from products.wizard.backend.facade.errors import IllegalStatusTransitionError +from products.wizard.backend.temporal.contracts import WizardRunFinalizationActivityInput + + +@activity.defn(name="wizard_finalize_run") +@asyncify +def finalize_run(input: WizardRunFinalizationActivityInput) -> None: + transition_cloud_run( + input.team_id, + input.run_id, + input.status, + error_code=input.error_code, + ) + + +def transition_cloud_run( + team_id: int, + run_id: UUID, + status: WizardRunStatus, + *, + error_code: str | None = None, +) -> None: + current = _get_cloud_run(team_id, run_id) + + if _is_settled(current, status, error_code): + return + + try: + wizard_facade.update_run_status(team_id, run_id, status, error_code=error_code) + except IllegalStatusTransitionError: + current = _get_cloud_run(team_id, run_id) + + if _is_settled(current, status, error_code): + return + + raise + + +def _get_cloud_run(team_id: int, run_id: UUID) -> WizardRunDTO: + run = wizard_facade.get_run(team_id, run_id) + + if run.environment != WizardRunEnvironment.CLOUD: + raise ValueError("Wizard Run transitions require a cloud Wizard Run.") + + return run + + +def _is_settled( + run: WizardRunDTO, + status: WizardRunStatus, + error_code: str | None, +) -> bool: + """ + Determines if the Wizard Run is already in a settled state for the given status and error code. + """ + + if _matches_status_and_error_code(run, status, error_code): + return True + + return run.status == WizardRunStatus.CANCELLED and status in ( + WizardRunStatus.FAILED, + WizardRunStatus.COMPLETED, + ) + + +def _matches_status_and_error_code( + run: WizardRunDTO, + status: WizardRunStatus, + error_code: str | None, +) -> bool: + """ + Checks if the Wizard Run's status and error code match the provided values. + """ + + return run.status == status and run.error_code == error_code diff --git a/products/wizard/backend/temporal/activities/local_package.py b/products/wizard/backend/temporal/activities/local_package.py new file mode 100644 index 000000000000..4a659e442b80 --- /dev/null +++ b/products/wizard/backend/temporal/activities/local_package.py @@ -0,0 +1,40 @@ +from temporalio import activity +from temporalio.exceptions import ApplicationError + +from posthog.temporal.common.utils import asyncify + +from products.wizard.backend.logic.workers import service as cloud_worker +from products.wizard.backend.logic.workers.config import local_wizard_source_root +from products.wizard.backend.logic.workers.local_package import InvalidLocalWizardSourceError +from products.wizard.backend.temporal.activities.errors import ( + WIZARD_RUN_CONFIGURATION_ERROR_TYPE, + WIZARD_WORKER_EXECUTION_ERROR_TYPE, +) +from products.wizard.backend.temporal.contracts import ProvisionedWizardWorker + + +@activity.defn(name="wizard_prepare_local_package") +@asyncify +def prepare_local_wizard(input: ProvisionedWizardWorker) -> None: + source_root = local_wizard_source_root() + if not input.use_local_wizard_source or source_root is None: + raise ApplicationError( + "Local Wizard source is not available.", + type=WIZARD_RUN_CONFIGURATION_ERROR_TYPE, + non_retryable=True, + ) + + try: + cloud_worker.prepare_local_wizard(input.sandbox_id, source_root) + except InvalidLocalWizardSourceError as error: + raise ApplicationError( + str(error), + type=WIZARD_RUN_CONFIGURATION_ERROR_TYPE, + non_retryable=True, + ) from error + except cloud_worker.WizardWorkerExecutionError as error: + raise ApplicationError( + str(error), + type=WIZARD_WORKER_EXECUTION_ERROR_TYPE, + non_retryable=True, + ) from error diff --git a/products/wizard/backend/temporal/activities/workspace.py b/products/wizard/backend/temporal/activities/workspace.py new file mode 100644 index 000000000000..ed09c3f88300 --- /dev/null +++ b/products/wizard/backend/temporal/activities/workspace.py @@ -0,0 +1,147 @@ +from temporalio import activity +from temporalio.exceptions import ApplicationError + +from posthog.temporal.common.utils import asyncify + +from products.wizard.backend.facade import api as wizard_facade +from products.wizard.backend.facade.contracts import GitRepositoryWorkspace, WizardRunDTO +from products.wizard.backend.facade.enums import ( + WizardRunEnvironment, + WizardRunStage, + WizardRunStatus, + WizardWorkspaceType, +) +from products.wizard.backend.facade.errors import ( + InvalidRepositoryError, + MissingGitHubIntegrationError, + RepositoryNotAccessibleError, +) +from products.wizard.backend.logic.runs.repository_access import authorize_git_repository_access +from products.wizard.backend.logic.workers import ( + lifecycle as worker_lifecycle, + service as cloud_worker, + store as worker_store, +) +from products.wizard.backend.logic.workers.service import GitRepositoryCloneRequest, WizardWorkerProvisionRequest +from products.wizard.backend.temporal.activities.errors import ( + WIZARD_REPOSITORY_ACCESS_ERROR_TYPE, + WIZARD_RUN_CONFIGURATION_ERROR_TYPE, + WIZARD_WORKER_EXECUTION_ERROR_TYPE, +) +from products.wizard.backend.temporal.activities.lifecycle import transition_cloud_run +from products.wizard.backend.temporal.contracts import ( + PreparedGitRepositoryWorkspace, + ProvisionedWizardWorker, + WizardRunActivityInput, +) + + +@activity.defn(name="wizard_provision_worker") +@asyncify +def provision_worker(input: WizardRunActivityInput) -> ProvisionedWizardWorker: + run = _get_cloud_run(input) + wizard_facade.update_run_stage(input.team_id, input.run_id, WizardRunStage.PROVISIONING) + if run.created_by_id is None: + raise ApplicationError( + "Wizard run creator is no longer available.", + type=WIZARD_RUN_CONFIGURATION_ERROR_TYPE, + non_retryable=True, + ) + if not isinstance(run.workspace, GitRepositoryWorkspace): + raise ApplicationError( + "Wizard run does not have a supported cloud workspace.", + type=WIZARD_RUN_CONFIGURATION_ERROR_TYPE, + non_retryable=True, + ) + + existing_sandbox_id = worker_store.get_sandbox_id(input.team_id, input.run_id) + if existing_sandbox_id is not None: + return ProvisionedWizardWorker( + team_id=input.team_id, + run_id=input.run_id, + sandbox_id=existing_sandbox_id, + workspace_type=WizardWorkspaceType.GIT_REPOSITORY, + use_local_wizard_source=input.use_local_wizard_source, + ) + + provisioning = cloud_worker.provision_wizard_worker( + WizardWorkerProvisionRequest( + team_id=input.team_id, + created_by_id=run.created_by_id, + run_id=input.run_id, + ) + ) + worker_store.record_provisioned_worker(input.team_id, input.run_id, provisioning) + transition_cloud_run(input.team_id, input.run_id, WizardRunStatus.RUNNING) + return ProvisionedWizardWorker( + team_id=input.team_id, + run_id=input.run_id, + sandbox_id=provisioning.sandbox_id, + workspace_type=WizardWorkspaceType.GIT_REPOSITORY, + use_local_wizard_source=input.use_local_wizard_source, + ) + + +@activity.defn(name="wizard_clone_repository") +@asyncify +def clone_repository(input: ProvisionedWizardWorker) -> PreparedGitRepositoryWorkspace: + run = _get_cloud_run(WizardRunActivityInput(team_id=input.team_id, run_id=input.run_id)) + wizard_facade.update_run_stage(input.team_id, input.run_id, WizardRunStage.PREPARING_WORKSPACE) + if not isinstance(run.workspace, GitRepositoryWorkspace): + raise ApplicationError( + "Wizard run does not have a Git repository workspace.", + type=WIZARD_RUN_CONFIGURATION_ERROR_TYPE, + non_retryable=True, + ) + + try: + integration_id = authorize_git_repository_access(input.team_id, run.workspace.repository) + except (InvalidRepositoryError, MissingGitHubIntegrationError, RepositoryNotAccessibleError) as error: + raise ApplicationError( + "GitHub access to the Wizard run repository is unavailable.", + type=WIZARD_REPOSITORY_ACCESS_ERROR_TYPE, + non_retryable=True, + ) from error + + try: + root_path = cloud_worker.clone_repository( + GitRepositoryCloneRequest( + sandbox_id=input.sandbox_id, + github_integration_id=integration_id, + repository=run.workspace.repository, + ) + ) + except cloud_worker.WizardWorkerExecutionError as error: + # Clone failures are usually transient (DNS, TLS, GitHub transport) and the clone is + # idempotent, so stay retryable and let PREPARATION_RETRY_POLICY reattempt. + raise ApplicationError( + str(error), + type=WIZARD_WORKER_EXECUTION_ERROR_TYPE, + ) from error + + return PreparedGitRepositoryWorkspace( + team_id=input.team_id, + run_id=input.run_id, + sandbox_id=input.sandbox_id, + repository=run.workspace.repository, + root_path=root_path, + github_integration_id=integration_id, + use_local_wizard_source=input.use_local_wizard_source, + ) + + +@activity.defn(name="wizard_destroy_worker") +@asyncify +def destroy_worker(input: ProvisionedWizardWorker) -> None: + worker_lifecycle.cleanup_worker(input.team_id, input.run_id, input.sandbox_id) + + +def _get_cloud_run(input: WizardRunActivityInput) -> WizardRunDTO: + run = wizard_facade.get_run(input.team_id, input.run_id) + if run.environment != WizardRunEnvironment.CLOUD: + raise ApplicationError( + "Wizard Worker requires a cloud run.", + type=WIZARD_RUN_CONFIGURATION_ERROR_TYPE, + non_retryable=True, + ) + return run diff --git a/products/wizard/backend/temporal/client.py b/products/wizard/backend/temporal/client.py new file mode 100644 index 000000000000..a55cecfbf9ab --- /dev/null +++ b/products/wizard/backend/temporal/client.py @@ -0,0 +1,48 @@ +from uuid import UUID + +from django.conf import settings + +from asgiref.sync import async_to_sync +from temporalio.common import RetryPolicy, WorkflowIDReusePolicy +from temporalio.exceptions import WorkflowAlreadyStartedError +from temporalio.service import RPCError, RPCStatusCode + +from posthog.temporal.common.client import async_connect + +from products.wizard.backend.temporal.constants import EXECUTE_WIZARD_RUN_WORKFLOW, wizard_run_workflow_id +from products.wizard.backend.temporal.contracts import WizardRunActivityInput +from products.wizard.backend.temporal.errors import WizardTemporalError + +WORKFLOW_RETRY_POLICY = RetryPolicy(maximum_attempts=1) + + +@async_to_sync +async def start_wizard_run_workflow(input: WizardRunActivityInput) -> None: + try: + client = await async_connect() + + await client.start_workflow( + EXECUTE_WIZARD_RUN_WORKFLOW, + input, + id=wizard_run_workflow_id(input.run_id), + id_reuse_policy=WorkflowIDReusePolicy.REJECT_DUPLICATE, + task_queue=settings.WIZARD_TASK_QUEUE, + retry_policy=WORKFLOW_RETRY_POLICY, + ) + except WorkflowAlreadyStartedError: + return + except (RPCError, RuntimeError) as error: + raise WizardTemporalError from error + + +@async_to_sync +async def cancel_wizard_run_workflow(run_id: UUID) -> None: + try: + client = await async_connect() + await client.get_workflow_handle(wizard_run_workflow_id(run_id)).cancel() + except RPCError as error: + if error.status == RPCStatusCode.NOT_FOUND: + return + raise WizardTemporalError from error + except RuntimeError as error: + raise WizardTemporalError from error diff --git a/products/wizard/backend/temporal/config.py b/products/wizard/backend/temporal/config.py new file mode 100644 index 000000000000..d55cd72545ce --- /dev/null +++ b/products/wizard/backend/temporal/config.py @@ -0,0 +1,19 @@ +from datetime import timedelta + +from temporalio.common import RetryPolicy + +FINALIZATION_TIMEOUT = timedelta(minutes=1) +PROVISION_TIMEOUT = timedelta(minutes=5) +PREPARATION_TIMEOUT = timedelta(minutes=10) +EXECUTION_TIMEOUT = timedelta(minutes=50) +HANDOFF_TIMEOUT = timedelta(minutes=5) +LOCAL_PACKAGE_PREPARATION_TIMEOUT = timedelta(minutes=20) +CLEANUP_TIMEOUT = timedelta(minutes=5) + +FINALIZATION_RETRY_POLICY = RetryPolicy(maximum_attempts=3) +PROVISION_RETRY_POLICY = RetryPolicy(maximum_attempts=1) +PREPARATION_RETRY_POLICY = RetryPolicy(maximum_attempts=3) +EXECUTION_RETRY_POLICY = RetryPolicy(maximum_attempts=1) +HANDOFF_RETRY_POLICY = RetryPolicy(maximum_attempts=3) +LOCAL_PACKAGE_PREPARATION_RETRY_POLICY = RetryPolicy(maximum_attempts=1) +CLEANUP_RETRY_POLICY = RetryPolicy(maximum_attempts=3) diff --git a/products/wizard/backend/temporal/constants.py b/products/wizard/backend/temporal/constants.py new file mode 100644 index 000000000000..088832dcf711 --- /dev/null +++ b/products/wizard/backend/temporal/constants.py @@ -0,0 +1,7 @@ +from uuid import UUID + +EXECUTE_WIZARD_RUN_WORKFLOW = "execute-wizard-run" + + +def wizard_run_workflow_id(run_id: UUID) -> str: + return f"wizard-run-{run_id}" diff --git a/products/wizard/backend/temporal/contracts.py b/products/wizard/backend/temporal/contracts.py new file mode 100644 index 000000000000..7cbb29e726f8 --- /dev/null +++ b/products/wizard/backend/temporal/contracts.py @@ -0,0 +1,46 @@ +from uuid import UUID + +from posthog.dataclasses import frozen + +from products.wizard.backend.facade.enums import WizardRunStatus, WizardWorkspaceType + + +@frozen +class WizardRunActivityInput: + team_id: int + run_id: UUID + use_local_wizard_source: bool = False + + +@frozen +class WizardRunFinalizationActivityInput: + team_id: int + run_id: UUID + status: WizardRunStatus + error_code: str | None = None + + def __post_init__(self) -> None: + if self.status not in (WizardRunStatus.FAILED, WizardRunStatus.CANCELLED): + raise ValueError("Wizard Run finalization requires a failed or cancelled status.") + if self.status != WizardRunStatus.FAILED and self.error_code is not None: + raise ValueError("Only failed Wizard Runs can have an error code.") + + +@frozen +class ProvisionedWizardWorker: + team_id: int + run_id: UUID + sandbox_id: str + workspace_type: WizardWorkspaceType + use_local_wizard_source: bool = False + + +@frozen +class PreparedGitRepositoryWorkspace: + team_id: int + run_id: UUID + sandbox_id: str + repository: str + root_path: str + github_integration_id: int + use_local_wizard_source: bool = False diff --git a/products/wizard/backend/temporal/errors.py b/products/wizard/backend/temporal/errors.py new file mode 100644 index 000000000000..7c21bd6fbbb9 --- /dev/null +++ b/products/wizard/backend/temporal/errors.py @@ -0,0 +1,33 @@ +from temporalio.exceptions import ActivityError, ApplicationError, TimeoutError + +from products.wizard.backend.facade.enums import WizardRunErrorCode +from products.wizard.backend.facade.validation import is_wizard_error_code +from products.wizard.backend.temporal.activities.errors import ( + WIZARD_REPOSITORY_ACCESS_ERROR_TYPE, + WIZARD_WORKER_TIMEOUT_ERROR_TYPE, +) + + +class WizardTemporalError(Exception): + pass + + +def wizard_run_error_code(error: ActivityError) -> str: + cause = error.cause + if isinstance(cause, TimeoutError): + return WizardRunErrorCode.TIMEOUT.value + if isinstance(cause, ApplicationError) and cause.type == WIZARD_WORKER_TIMEOUT_ERROR_TYPE: + return WizardRunErrorCode.TIMEOUT.value + if isinstance(cause, ApplicationError) and cause.type == WIZARD_REPOSITORY_ACCESS_ERROR_TYPE: + return WizardRunErrorCode.REPOSITORY_ACCESS_FAILED.value + if isinstance(cause, ApplicationError) and is_wizard_error_code(cause.type): + return cause.type + match error.activity_type: + case "wizard_provision_worker": + return WizardRunErrorCode.PROVISIONING_FAILED.value + case "wizard_clone_repository": + return WizardRunErrorCode.WORKSPACE_PREPARATION_FAILED.value + case "wizard_create_run_artifacts": + return WizardRunErrorCode.ARTIFACT_CREATION_FAILED.value + case _: + return WizardRunErrorCode.EXECUTION_FAILED.value diff --git a/products/wizard/backend/temporal/serializers.py b/products/wizard/backend/temporal/serializers.py new file mode 100644 index 000000000000..2ff25417606d --- /dev/null +++ b/products/wizard/backend/temporal/serializers.py @@ -0,0 +1,13 @@ +import json +from uuid import UUID + +from products.wizard.backend.temporal.contracts import WizardRunActivityInput + + +def deserialize_workflow_input(inputs: list[str]) -> WizardRunActivityInput: + value = json.loads(inputs[0]) + return WizardRunActivityInput( + team_id=value["team_id"], + run_id=UUID(value["run_id"]), + use_local_wizard_source=value.get("use_local_wizard_source") is True, + ) diff --git a/products/wizard/backend/temporal/workflows/execute_run.py b/products/wizard/backend/temporal/workflows/execute_run.py new file mode 100644 index 000000000000..0bfb173daf9e --- /dev/null +++ b/products/wizard/backend/temporal/workflows/execute_run.py @@ -0,0 +1,171 @@ +import asyncio +from uuid import UUID + +from temporalio import workflow +from temporalio.exceptions import ActivityError + +from posthog.temporal.common.base import PostHogWorkflow + +from products.wizard.backend.facade.enums import WizardRunErrorCode, WizardRunStatus, WizardWorkspaceType +from products.wizard.backend.temporal.activities.execution import execute_wizard +from products.wizard.backend.temporal.activities.handoff import create_run_artifacts +from products.wizard.backend.temporal.activities.lifecycle import finalize_run +from products.wizard.backend.temporal.activities.local_package import prepare_local_wizard +from products.wizard.backend.temporal.activities.workspace import clone_repository, destroy_worker, provision_worker +from products.wizard.backend.temporal.config import ( + CLEANUP_RETRY_POLICY, + CLEANUP_TIMEOUT, + EXECUTION_RETRY_POLICY, + EXECUTION_TIMEOUT, + FINALIZATION_RETRY_POLICY, + FINALIZATION_TIMEOUT, + HANDOFF_RETRY_POLICY, + HANDOFF_TIMEOUT, + LOCAL_PACKAGE_PREPARATION_RETRY_POLICY, + LOCAL_PACKAGE_PREPARATION_TIMEOUT, + PREPARATION_RETRY_POLICY, + PREPARATION_TIMEOUT, + PROVISION_RETRY_POLICY, + PROVISION_TIMEOUT, +) +from products.wizard.backend.temporal.constants import EXECUTE_WIZARD_RUN_WORKFLOW, wizard_run_workflow_id +from products.wizard.backend.temporal.contracts import ( + PreparedGitRepositoryWorkspace, + ProvisionedWizardWorker, + WizardRunActivityInput, + WizardRunFinalizationActivityInput, +) +from products.wizard.backend.temporal.errors import wizard_run_error_code +from products.wizard.backend.temporal.serializers import deserialize_workflow_input + + +@workflow.defn(name=EXECUTE_WIZARD_RUN_WORKFLOW) +class ExecuteWizardRunWorkflow(PostHogWorkflow): + @staticmethod + def workflow_id_for(run_id: UUID) -> str: + return wizard_run_workflow_id(run_id) + + @staticmethod + def parse_inputs(inputs: list[str]) -> WizardRunActivityInput: + return deserialize_workflow_input(inputs) + + @workflow.run + async def run(self, input: WizardRunActivityInput) -> None: + worker: ProvisionedWizardWorker | None = None + try: + worker = await workflow.execute_activity( + provision_worker, + input, + start_to_close_timeout=PROVISION_TIMEOUT, + retry_policy=PROVISION_RETRY_POLICY, + cancellation_type=workflow.ActivityCancellationType.WAIT_CANCELLATION_COMPLETED, + ) + + if worker is None: + raise RuntimeError("Wizard Worker provisioning returned no worker.") + + if input.use_local_wizard_source: + await workflow.execute_activity( + prepare_local_wizard, + worker, + start_to_close_timeout=LOCAL_PACKAGE_PREPARATION_TIMEOUT, + retry_policy=LOCAL_PACKAGE_PREPARATION_RETRY_POLICY, + cancellation_type=workflow.ActivityCancellationType.WAIT_CANCELLATION_COMPLETED, + ) + + workspace = await self._prepare_workspace(worker) + + await workflow.execute_activity( + execute_wizard, + workspace, + start_to_close_timeout=EXECUTION_TIMEOUT, + retry_policy=EXECUTION_RETRY_POLICY, + cancellation_type=workflow.ActivityCancellationType.WAIT_CANCELLATION_COMPLETED, + ) + + await workflow.execute_activity( + create_run_artifacts, + workspace, + start_to_close_timeout=HANDOFF_TIMEOUT, + retry_policy=HANDOFF_RETRY_POLICY, + cancellation_type=workflow.ActivityCancellationType.WAIT_CANCELLATION_COMPLETED, + ) + except asyncio.CancelledError: + await self._destroy_worker(worker) + + await self._finalize_run( + WizardRunFinalizationActivityInput( + team_id=input.team_id, + run_id=input.run_id, + status=WizardRunStatus.CANCELLED, + ) + ) + raise + except ActivityError as error: + await self._destroy_worker(worker) + + await self._finalize_run( + WizardRunFinalizationActivityInput( + team_id=input.team_id, + run_id=input.run_id, + status=WizardRunStatus.FAILED, + error_code=wizard_run_error_code(error), + ) + ) + raise + except Exception: + workflow.logger.exception( + "wizard_run_workflow_failed", + extra={"team_id": input.team_id, "run_id": str(input.run_id)}, + ) + await self._destroy_worker(worker) + + await self._finalize_run( + WizardRunFinalizationActivityInput( + team_id=input.team_id, + run_id=input.run_id, + status=WizardRunStatus.FAILED, + error_code=WizardRunErrorCode.EXECUTION_FAILED, + ) + ) + raise + else: + await self._destroy_worker(worker) + + @staticmethod + async def _finalize_run(input: WizardRunFinalizationActivityInput) -> None: + await workflow.execute_activity( + finalize_run, + input, + start_to_close_timeout=FINALIZATION_TIMEOUT, + retry_policy=FINALIZATION_RETRY_POLICY, + ) + + @staticmethod + async def _destroy_worker(worker: ProvisionedWizardWorker | None) -> None: + if worker is None: + return + try: + await workflow.execute_activity( + destroy_worker, + worker, + start_to_close_timeout=CLEANUP_TIMEOUT, + retry_policy=CLEANUP_RETRY_POLICY, + ) + except ActivityError: + workflow.logger.exception( + "wizard_worker_cleanup_failed", + extra={"team_id": worker.team_id, "run_id": str(worker.run_id)}, + ) + + @staticmethod + async def _prepare_workspace(worker: ProvisionedWizardWorker) -> PreparedGitRepositoryWorkspace: + if worker.workspace_type == WizardWorkspaceType.GIT_REPOSITORY: + return await workflow.execute_activity( + clone_repository, + worker, + start_to_close_timeout=PREPARATION_TIMEOUT, + retry_policy=PREPARATION_RETRY_POLICY, + cancellation_type=workflow.ActivityCancellationType.WAIT_CANCELLATION_COMPLETED, + ) + raise ValueError(f"Unsupported cloud workspace type: {worker.workspace_type}") diff --git a/products/wizard/backend/tests/conftest.py b/products/wizard/backend/tests/conftest.py index ae3f00200968..3a012a95ef9d 100644 --- a/products/wizard/backend/tests/conftest.py +++ b/products/wizard/backend/tests/conftest.py @@ -1,8 +1,15 @@ import pytest +from unittest.mock import patch from posthog.models.scoping import team_scope +@pytest.fixture(autouse=True) +def _use_default_wizard_registry(): + with patch("posthoganalytics.get_feature_flag_payload", return_value=None): + yield + + @pytest.fixture(autouse=True) def _set_team_scope(request): if request.node.get_closest_marker("django_db") is None: diff --git a/products/wizard/backend/tests/runs/test_admission.py b/products/wizard/backend/tests/runs/test_admission.py new file mode 100644 index 000000000000..c67e72b8d92b --- /dev/null +++ b/products/wizard/backend/tests/runs/test_admission.py @@ -0,0 +1,74 @@ +from datetime import timedelta + +import pytest + +from django.utils import timezone + +from products.wizard.backend.facade.enums import ( + WizardRunDispatchStatus, + WizardRunEnvironment, + WizardRunStatus, + WizardWorkspaceType, +) +from products.wizard.backend.facade.errors import ( + ActiveWizardRunError, + WizardRunDailyLimitError, + WizardRunHourlyLimitError, +) +from products.wizard.backend.logic.programs import program_to_mapping +from products.wizard.backend.logic.registry.config import POSTHOG_INTEGRATION_PROGRAM +from products.wizard.backend.logic.runs.admission import enforce_cloud_run_creation_policy +from products.wizard.backend.logic.runs.config import CLOUD_RUN_DAILY_LIMIT, CLOUD_RUN_HOURLY_LIMIT +from products.wizard.backend.models import WizardRun + + +def _create_cloud_run(team_id: int, user_id: int, status: WizardRunStatus, age: timedelta) -> WizardRun: + run = WizardRun.objects.for_team(team_id).create( + team_id=team_id, + created_by_id=user_id, + environment=WizardRunEnvironment.CLOUD.value, + workspace_type=WizardWorkspaceType.GIT_REPOSITORY.value, + workspace={"repository": "posthog/posthog"}, + program=program_to_mapping(POSTHOG_INTEGRATION_PROGRAM), + status=status.value, + dispatch_status=WizardRunDispatchStatus.DISPATCHED.value, + ) + WizardRun.objects.for_team(team_id).filter(id=run.id).update(created_at=timezone.now() - age) + return run + + +@pytest.mark.django_db +def test_cloud_admission_rejects_an_active_run(team, user) -> None: + _create_cloud_run(team.id, user.id, WizardRunStatus.RUNNING, timedelta(minutes=5)) + + with pytest.raises(ActiveWizardRunError): + enforce_cloud_run_creation_policy(team.id, user.id) + + +@pytest.mark.django_db +def test_cloud_admission_limits_hourly_runs(team, user) -> None: + for _ in range(CLOUD_RUN_HOURLY_LIMIT - 1): + _create_cloud_run(team.id, user.id, WizardRunStatus.COMPLETED, timedelta(minutes=10)) + enforce_cloud_run_creation_policy(team.id, user.id) + _create_cloud_run(team.id, user.id, WizardRunStatus.FAILED, timedelta(minutes=20)) + + with pytest.raises(WizardRunHourlyLimitError): + enforce_cloud_run_creation_policy(team.id, user.id) + + +@pytest.mark.django_db +def test_cloud_admission_limits_daily_runs(team, user) -> None: + for _ in range(CLOUD_RUN_DAILY_LIMIT - 1): + _create_cloud_run(team.id, user.id, WizardRunStatus.COMPLETED, timedelta(hours=2)) + enforce_cloud_run_creation_policy(team.id, user.id) + _create_cloud_run(team.id, user.id, WizardRunStatus.COMPLETED, timedelta(hours=2)) + + with pytest.raises(WizardRunDailyLimitError): + enforce_cloud_run_creation_policy(team.id, user.id) + + +@pytest.mark.django_db +def test_cloud_admission_ignores_historical_runs(team, user) -> None: + _create_cloud_run(team.id, user.id, WizardRunStatus.COMPLETED, timedelta(days=2)) + + enforce_cloud_run_creation_policy(team.id, user.id) diff --git a/products/wizard/backend/tests/runs/test_cloud_activities.py b/products/wizard/backend/tests/runs/test_cloud_activities.py new file mode 100644 index 000000000000..68b40bedb881 --- /dev/null +++ b/products/wizard/backend/tests/runs/test_cloud_activities.py @@ -0,0 +1,355 @@ +from datetime import timedelta +from pathlib import Path +from uuid import uuid4 + +import pytest +from unittest.mock import MagicMock, patch + +from django.apps import apps +from django.test import override_settings +from django.utils import timezone + +from asgiref.sync import async_to_sync +from temporalio.exceptions import ApplicationError +from temporalio.testing import ActivityEnvironment + +from products.wizard.backend.facade import api as wizard_facade +from products.wizard.backend.facade.contracts import ( + CreateWizardRunInput, + GitRepositoryWorkspace, + WizardRunDTO, + WizardRunGitDiffArtifactDTO, + WizardRunPullRequestArtifactDTO, +) +from products.wizard.backend.facade.enums import ( + WizardRunArtifactType, + WizardRunEnvironment, + WizardRunStatus, + WizardWorkspaceType, +) +from products.wizard.backend.logic.workers.contracts import RepositoryPullRequest, WizardWorkerResourceUsage +from products.wizard.backend.logic.workers.service import ( + GitRepositoryCloneRequest, + GitRepositoryHandoffRequest, + WizardExecutionRequest, + WizardWorkerExecutionError, + WizardWorkerProvisioning, + WizardWorkerProvisionRequest, + WizardWorkerResult, + WizardWorkerTimeoutError, +) +from products.wizard.backend.temporal.activities.errors import ( + WIZARD_REPOSITORY_ACCESS_ERROR_TYPE, + WIZARD_WORKER_EXECUTION_ERROR_TYPE, + WIZARD_WORKER_TIMEOUT_ERROR_TYPE, +) +from products.wizard.backend.temporal.activities.execution import execute_wizard +from products.wizard.backend.temporal.activities.handoff import create_run_artifacts +from products.wizard.backend.temporal.activities.local_package import prepare_local_wizard +from products.wizard.backend.temporal.activities.workspace import clone_repository, provision_worker +from products.wizard.backend.temporal.contracts import ( + PreparedGitRepositoryWorkspace, + ProvisionedWizardWorker, + WizardRunActivityInput, +) + + +def _create_cloud_run( + team_id: int, + user_id: int, + *, + program_id: str = "posthog-integration", + registry_payload: object = None, +) -> WizardRunDTO: + with ( + patch("posthoganalytics.get_feature_flag_payload", return_value=registry_payload), + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.resolve_team_github_integration_id", + return_value=123, + ), + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.repository_accessible_via_integration", + return_value=True, + ), + ): + return wizard_facade.create_run( + CreateWizardRunInput( + team_id=team_id, + created_by_id=user_id, + environment=WizardRunEnvironment.CLOUD, + idempotency_key="test-cloud-activity", + workspace=GitRepositoryWorkspace(repository="posthog/posthog"), + program_id=program_id, + ) + ) + + +def _worker(run: WizardRunDTO) -> ProvisionedWizardWorker: + return ProvisionedWizardWorker( + team_id=run.team_id, + run_id=run.id, + sandbox_id="worker-id", + workspace_type=WizardWorkspaceType.GIT_REPOSITORY, + ) + + +def _provisioning() -> WizardWorkerProvisioning: + return WizardWorkerProvisioning( + sandbox_id="worker-id", + resource_usage=WizardWorkerResourceUsage( + cpu_cores=2, + memory_gb=4, + disk_size_gb=16, + ttl_seconds=4500, + ttl_expires_at=timezone.now() + timedelta(seconds=4500), + ), + ) + + +def _workspace(run: WizardRunDTO) -> PreparedGitRepositoryWorkspace: + return PreparedGitRepositoryWorkspace( + team_id=run.team_id, + run_id=run.id, + sandbox_id="worker-id", + repository="posthog/posthog", + root_path="/tmp/workspace/repos/posthog/posthog", + github_integration_id=456, + ) + + +async def _run_provision_worker(input: WizardRunActivityInput) -> ProvisionedWizardWorker: + return await ActivityEnvironment().run(provision_worker, input) + + +async def _run_clone_repository(input: ProvisionedWizardWorker) -> PreparedGitRepositoryWorkspace: + return await ActivityEnvironment().run(clone_repository, input) + + +async def _run_prepare_local_wizard(input: ProvisionedWizardWorker) -> None: + await ActivityEnvironment().run(prepare_local_wizard, input) + + +async def _run_execute_wizard(input: PreparedGitRepositoryWorkspace) -> None: + await ActivityEnvironment().run(execute_wizard, input) + + +async def _run_create_run_artifacts(input: PreparedGitRepositoryWorkspace) -> None: + await ActivityEnvironment().run(create_run_artifacts, input) + + +@pytest.mark.django_db(transaction=True) +def test_provision_worker_uses_persisted_run_identity(team, user) -> None: + run = _create_cloud_run(team.id, user.id) + + with patch( + "products.wizard.backend.temporal.activities.workspace.cloud_worker.provision_wizard_worker", + return_value=_provisioning(), + ) as provision: + result = async_to_sync(_run_provision_worker)(WizardRunActivityInput(team_id=team.id, run_id=run.id)) + + assert result == _worker(run) + persisted_run = wizard_facade.get_run(team.id, run.id) + assert persisted_run.status == WizardRunStatus.RUNNING + assert persisted_run.stage == "provisioning" + provision.assert_called_once_with( + WizardWorkerProvisionRequest(team_id=team.id, created_by_id=user.id, run_id=run.id) + ) + + worker_record = apps.get_model("wizard", "WizardWorker").objects.for_team(team.id).get(run_id=run.id) + assert worker_record.sandbox_id == "worker-id" + assert worker_record.cleanup_status == "active" + assert worker_record.resource_usage["cpu_cores"] == 2.0 + + provision.reset_mock() + repeated = async_to_sync(_run_provision_worker)(WizardRunActivityInput(team_id=team.id, run_id=run.id)) + + assert repeated == result + provision.assert_not_called() + + +@pytest.mark.django_db(transaction=True) +def test_clone_repository_rechecks_access_before_preparing_workspace(team, user) -> None: + run = _create_cloud_run(team.id, user.id) + worker = _worker(run) + + with ( + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.resolve_team_github_integration_id", + return_value=456, + ) as resolve_integration, + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.repository_accessible_via_integration", + return_value=True, + ) as repository_accessible, + patch( + "products.wizard.backend.temporal.activities.workspace.cloud_worker.clone_repository", + return_value="/tmp/workspace/repos/posthog/posthog", + ) as clone, + ): + result = async_to_sync(_run_clone_repository)(worker) + + assert result == _workspace(run) + assert wizard_facade.get_run(team.id, run.id).stage == "preparing_workspace" + resolve_integration.assert_called_once_with(team.id) + repository_accessible.assert_called_once_with(team.id, 456, "posthog/posthog") + clone.assert_called_once_with( + GitRepositoryCloneRequest( + sandbox_id=worker.sandbox_id, + github_integration_id=456, + repository="posthog/posthog", + ) + ) + + +@override_settings(DEBUG=True, LOCAL_WIZARD_ROOT="/tmp/posthog-wizard") +def test_prepare_local_wizard_uses_configured_source() -> None: + worker = ProvisionedWizardWorker( + team_id=7, + run_id=uuid4(), + sandbox_id="worker-id", + workspace_type=WizardWorkspaceType.GIT_REPOSITORY, + use_local_wizard_source=True, + ) + + with patch( + "products.wizard.backend.temporal.activities.local_package.cloud_worker.prepare_local_wizard" + ) as prepare: + async_to_sync(_run_prepare_local_wizard)(worker) + + prepare.assert_called_once_with(worker.sandbox_id, Path("/tmp/posthog-wizard").resolve()) + + +@pytest.mark.django_db(transaction=True) +def test_clone_repository_rejects_access_revoked_after_creation(team, user) -> None: + run = _create_cloud_run(team.id, user.id) + + with ( + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.resolve_team_github_integration_id", + return_value=None, + ), + patch("products.wizard.backend.temporal.activities.workspace.cloud_worker.clone_repository") as clone, + pytest.raises(ApplicationError) as error, + ): + async_to_sync(_run_clone_repository)(_worker(run)) + + assert error.value.type == WIZARD_REPOSITORY_ACCESS_ERROR_TYPE + assert error.value.non_retryable is True + clone.assert_not_called() + + +@pytest.mark.django_db(transaction=True) +def test_clone_repository_keeps_transient_clone_failure_retryable(team, user) -> None: + run = _create_cloud_run(team.id, user.id) + + with ( + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.resolve_team_github_integration_id", + return_value=456, + ), + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.repository_accessible_via_integration", + return_value=True, + ), + patch( + "products.wizard.backend.temporal.activities.workspace.cloud_worker.clone_repository", + side_effect=WizardWorkerExecutionError("repository clone", 128), + ), + pytest.raises(ApplicationError) as error, + ): + async_to_sync(_run_clone_repository)(_worker(run)) + + assert error.value.type == WIZARD_WORKER_EXECUTION_ERROR_TYPE + assert not error.value.non_retryable + + +@pytest.mark.parametrize( + "worker_error, error_type", + ( + (WizardWorkerTimeoutError(), WIZARD_WORKER_TIMEOUT_ERROR_TYPE), + (WizardWorkerExecutionError("execution", 1), WIZARD_WORKER_EXECUTION_ERROR_TYPE), + ( + WizardWorkerExecutionError("execution", 1, wizard_error_code="PHW_DETECT_NO_POSTHOG_SDK"), + "PHW_DETECT_NO_POSTHOG_SDK", + ), + ), +) +def test_execute_wizard_maps_worker_error(worker_error: Exception, error_type: str) -> None: + workspace = PreparedGitRepositoryWorkspace( + team_id=7, + run_id=uuid4(), + sandbox_id="worker-id", + repository="posthog/posthog", + root_path="/tmp/workspace/repos/posthog/posthog", + github_integration_id=456, + ) + + with ( + patch("products.wizard.backend.temporal.activities.execution.wizard_facade.update_run_stage"), + patch("products.wizard.backend.temporal.activities.execution.wizard_facade.get_run") as get_run, + patch( + "products.wizard.backend.temporal.activities.execution.cloud_worker.execute_wizard", + side_effect=worker_error, + ) as execute, + pytest.raises(ApplicationError) as error, + ): + get_run.return_value.program.command = () + get_run.return_value.program.wizard_version = "2.60.0" + async_to_sync(_run_execute_wizard)(workspace) + + execute.assert_called_once_with( + WizardExecutionRequest( + sandbox_id=workspace.sandbox_id, + workspace_path=workspace.root_path, + team_id=workspace.team_id, + wizard_version="2.60.0", + program_command=(), + ) + ) + assert error.value.type == error_type + + +@pytest.mark.django_db(transaction=True) +@patch("products.wizard.backend.logic.artifacts.service.object_storage.write") +def test_create_run_artifacts_persists_git_diff_and_pull_request(_write: MagicMock, team, user) -> None: + run = _create_cloud_run(team.id, user.id) + wizard_facade.update_run_status(team.id, run.id, WizardRunStatus.RUNNING) + workspace = _workspace(run) + diff = b"diff --git a/a b/a\n" + pull_request = RepositoryPullRequest( + repository=workspace.repository, + number=123, + url="https://github.com/posthog/posthog/pull/123", + head_branch="posthog/wizard-123", + base_branch="master", + ) + + with patch( + "products.wizard.backend.temporal.activities.handoff.cloud_worker.create_git_repository_handoff", + return_value=WizardWorkerResult(diff=diff, pull_request=pull_request), + ) as handoff: + async_to_sync(_run_create_run_artifacts)(workspace) + + handoff.assert_called_once_with( + GitRepositoryHandoffRequest( + team_id=workspace.team_id, + run_id=workspace.run_id, + sandbox_id=workspace.sandbox_id, + workspace_path=workspace.root_path, + github_integration_id=workspace.github_integration_id, + repository=workspace.repository, + ) + ) + artifacts = wizard_facade.list_run_artifacts(team.id, run.id) + assert {artifact.artifact_type for artifact in artifacts} == { + WizardRunArtifactType.GIT_DIFF, + WizardRunArtifactType.PULL_REQUEST, + } + git_diff_artifact = next(artifact for artifact in artifacts if isinstance(artifact, WizardRunGitDiffArtifactDTO)) + assert git_diff_artifact.size_bytes == len(diff) + pull_request_artifact = next( + artifact for artifact in artifacts if isinstance(artifact, WizardRunPullRequestArtifactDTO) + ) + assert pull_request_artifact.url == pull_request.url + assert pull_request_artifact.number == pull_request.number + assert wizard_facade.get_run(team.id, run.id).status == WizardRunStatus.COMPLETED diff --git a/products/wizard/backend/tests/runs/test_cloud_execution_activity.py b/products/wizard/backend/tests/runs/test_cloud_execution_activity.py new file mode 100644 index 000000000000..fb0c5778905e --- /dev/null +++ b/products/wizard/backend/tests/runs/test_cloud_execution_activity.py @@ -0,0 +1,82 @@ +import pytest +from unittest.mock import patch + +from asgiref.sync import async_to_sync +from temporalio.testing import ActivityEnvironment + +from products.wizard.backend.facade import api as wizard_facade +from products.wizard.backend.facade.contracts import CreateWizardRunInput, GitRepositoryWorkspace +from products.wizard.backend.facade.enums import WizardRunEnvironment +from products.wizard.backend.logic.workers.service import WizardExecutionRequest +from products.wizard.backend.temporal.activities.execution import execute_wizard +from products.wizard.backend.temporal.contracts import PreparedGitRepositoryWorkspace + + +async def _run_execute_wizard(input: PreparedGitRepositoryWorkspace) -> None: + await ActivityEnvironment().run(execute_wizard, input) + + +@pytest.mark.django_db(transaction=True) +def test_execute_wizard_uses_persisted_program_snapshot(team, user) -> None: + registry_payload = { + "version": 1, + "programs": [ + { + "id": "web-analytics-audit", + "name": "Web analytics audit", + "description": "Audit a project's web analytics setup", + "wizard_version": "2.60.0", + "command": ["audit", "web-analytics"], + "tags": ["audit", "web-analytics"], + "required_programs": ["posthog-integration"], + "supported_environments": ["cloud"], + } + ], + } + with ( + patch("posthoganalytics.get_feature_flag_payload", return_value=registry_payload), + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.resolve_team_github_integration_id", + return_value=123, + ), + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.repository_accessible_via_integration", + return_value=True, + ), + patch("products.wizard.backend.logic.runs.lifecycle.dispatch_created_cloud_wizard_run_to_temporal_worker"), + ): + run = wizard_facade.create_run( + CreateWizardRunInput( + team_id=team.id, + created_by_id=user.id, + environment=WizardRunEnvironment.CLOUD, + idempotency_key="test-cloud-execution", + workspace=GitRepositoryWorkspace(repository="posthog/posthog"), + program_id="web-analytics-audit", + wizard_version="2.60.0", + ) + ) + + workspace = PreparedGitRepositoryWorkspace( + team_id=team.id, + run_id=run.id, + sandbox_id="worker-id", + repository="posthog/posthog", + root_path="/tmp/workspace/repos/posthog/posthog", + github_integration_id=456, + ) + with ( + patch("posthoganalytics.get_feature_flag_payload", side_effect=AssertionError("registry was re-evaluated")), + patch("products.wizard.backend.temporal.activities.execution.cloud_worker.execute_wizard") as execute, + ): + async_to_sync(_run_execute_wizard)(workspace) + + execute.assert_called_once_with( + WizardExecutionRequest( + sandbox_id="worker-id", + workspace_path="/tmp/workspace/repos/posthog/posthog", + team_id=team.id, + wizard_version="2.60.0", + program_command=("audit", "web-analytics"), + ) + ) diff --git a/products/wizard/backend/tests/runs/test_cloud_worker.py b/products/wizard/backend/tests/runs/test_cloud_worker.py new file mode 100644 index 000000000000..e3b4deac8ccf --- /dev/null +++ b/products/wizard/backend/tests/runs/test_cloud_worker.py @@ -0,0 +1,461 @@ +import io +import tarfile +from pathlib import Path +from types import SimpleNamespace +from uuid import uuid4 + +import pytest +from unittest.mock import MagicMock, patch + +from django.test import override_settings + +from parameterized import parameterized + +from products.tasks.backend.facade.sandbox import SandboxNotFoundError +from products.wizard.backend.logic.artifacts.config import MAX_GIT_DIFF_BYTES +from products.wizard.backend.logic.workers.commands import wizard_handoff_output_path +from products.wizard.backend.logic.workers.config import ( + LOCAL_WIZARD_ARCHIVE_PATH, + LOCAL_WIZARD_INSTALL_PATH, + WIZARD_DIFF_OUTPUT_PATH, + WIZARD_PACKAGE_INSTALL_PATH, + local_wizard_source_root, +) +from products.wizard.backend.logic.workers.contracts import RepositoryPullRequest +from products.wizard.backend.logic.workers.errors import InvalidWizardProgramCommandError, InvalidWizardVersionError +from products.wizard.backend.logic.workers.service import ( + GitRepositoryCloneRequest, + GitRepositoryHandoffRequest, + WizardExecutionRequest, + WizardWorkerExecutionError, + WizardWorkerProvisionRequest, + WizardWorkerResult, + WizardWorkerTimeoutError, + clone_repository, + create_git_repository_handoff, + destroy_worker, + execute_wizard, + prepare_local_wizard, + provision_wizard_worker, +) + + +def _execution_result(*, stdout: str = "", stderr: str = "", exit_code: int = 0) -> SimpleNamespace: + return SimpleNamespace(stdout=stdout, stderr=stderr, exit_code=exit_code) + + +@override_settings(DEBUG=False, LOCAL_WIZARD_ROOT="/tmp/posthog-wizard") +def test_local_wizard_source_is_disabled_outside_debug_mode() -> None: + assert local_wizard_source_root() is None + + +@override_settings(DEBUG=True, SANDBOX_MCP_URL="http://host.docker.internal:8787/mcp") +@patch("products.wizard.backend.logic.workers.service.get_sandbox_class") +@patch("products.wizard.backend.logic.workers.service.create_wizard_oauth_access_token_for_user") +@patch("products.wizard.backend.logic.workers.service.User.objects.get") +def test_provision_worker_configures_wizard_environment( + get_user: MagicMock, + create_wizard_token: MagicMock, + get_sandbox_class: MagicMock, +) -> None: + request = WizardWorkerProvisionRequest(team_id=7, created_by_id=13, run_id=uuid4()) + create_wizard_token.return_value = "wizard-secret" + get_sandbox_class.return_value.create.return_value.id = "worker-id" + + provisioning = provision_wizard_worker(request) + + assert provisioning.sandbox_id == "worker-id" + assert provisioning.resource_usage.cpu_cores == 2 + assert provisioning.resource_usage.memory_gb == 4 + assert provisioning.resource_usage.disk_size_gb == 16 + assert provisioning.resource_usage.ttl_seconds == 75 * 60 + get_user.assert_called_once_with(id=request.created_by_id) + config = get_sandbox_class.return_value.create.call_args.args[0] + assert "GITHUB_TOKEN" not in config.environment_variables + assert config.environment_variables["POSTHOG_WIZARD_API_KEY"] == "wizard-secret" + assert config.environment_variables["MCP_URL"] == "http://host.docker.internal:8787/mcp" + assert "POSTHOG_WIZARD_RUN_ID" not in config.environment_variables + assert config.environment_variables["POSTHOG_HANDOFF_OUTPUT_PATH"] == wizard_handoff_output_path(request.run_id) + assert config.ttl_seconds == 75 * 60 + + +@override_settings(DEBUG=False, SANDBOX_MCP_URL="http://host.docker.internal:8787/mcp") +@patch("products.wizard.backend.logic.workers.service.get_sandbox_class") +@patch("products.wizard.backend.logic.workers.service.create_wizard_oauth_access_token_for_user") +@patch("products.wizard.backend.logic.workers.service.User.objects.get") +def test_provision_worker_ignores_local_mcp_endpoint_outside_debug_mode( + get_user: MagicMock, + create_wizard_token: MagicMock, + get_sandbox_class: MagicMock, +) -> None: + request = WizardWorkerProvisionRequest(team_id=7, created_by_id=13, run_id=uuid4()) + create_wizard_token.return_value = "wizard-secret" + get_sandbox_class.return_value.create.return_value.id = "worker-id" + + provision_wizard_worker(request) + + config = get_sandbox_class.return_value.create.call_args.args[0] + assert "MCP_URL" not in config.environment_variables + + +@patch("products.wizard.backend.logic.workers.service.get_sandbox_class") +@patch("products.wizard.backend.logic.workers.service.get_github_token") +def test_clone_repository_uses_integration_token( + get_github_token: MagicMock, + get_sandbox_class: MagicMock, +) -> None: + request = GitRepositoryCloneRequest( + sandbox_id="worker-id", + github_integration_id=17, + repository="PostHog/PostHog", + ) + get_github_token.return_value = "github-secret" + sandbox = get_sandbox_class.return_value.get_by_id.return_value + sandbox.clone_repository.return_value = _execution_result() + sandbox.execute.return_value = _execution_result() + + root_path = clone_repository(request) + + assert root_path == "/tmp/workspace/repos/posthog/posthog" + get_sandbox_class.return_value.get_by_id.assert_called_once_with(request.sandbox_id) + sandbox.clone_repository.assert_called_once_with(request.repository, github_token="github-secret", shallow=True) + sanitize_command = sandbox.execute.call_args.args[0] + assert "github-secret" not in sanitize_command + assert "https://github.com/PostHog/PostHog.git" in sanitize_command + + +@patch("products.wizard.backend.logic.workers.service.get_sandbox_class") +@patch("products.wizard.backend.logic.workers.service.get_github_token", return_value="github-secret") +def test_clone_repository_rejects_clone_failure( + _get_github_token: MagicMock, + get_sandbox_class: MagicMock, +) -> None: + request = GitRepositoryCloneRequest( + sandbox_id="worker-id", + github_integration_id=17, + repository="PostHog/PostHog", + ) + sandbox = get_sandbox_class.return_value.get_by_id.return_value + sandbox.clone_repository.return_value = _execution_result( + stderr="fatal: https://x-access-token:github-secret@github.com/PostHog/PostHog.git not found", + exit_code=128, + ) + + with pytest.raises(WizardWorkerExecutionError) as error: + clone_repository(request) + + assert "github-secret" not in str(error.value) + assert "[REDACTED]" in str(error.value) + + +@patch("products.wizard.backend.logic.workers.service.get_sandbox_class") +def test_prepare_local_wizard_uploads_source_and_builds_it(get_sandbox_class: MagicMock, tmp_path: Path) -> None: + source_root = tmp_path / "wizard" + source_root.mkdir() + (source_root / "package.json").write_text('{"name":"@posthog/wizard"}') + (source_root / "pnpm-lock.yaml").write_text("lockfileVersion: '9.0'") + (source_root / "pnpm-workspace.yaml").write_text("packages: []") + (source_root / "bin.ts").write_text("export {}") + (source_root / "src").mkdir() + (source_root / "src" / "error.ts").write_text("export const code = 'sdk_missing'") + (source_root / ".env").write_text("SECRET=value") + (source_root / "node_modules").mkdir() + (source_root / "node_modules" / "dependency.js").write_text("ignored") + sandbox = get_sandbox_class.return_value.get_by_id.return_value + sandbox.write_file.return_value = _execution_result() + sandbox.execute.return_value = _execution_result() + + prepare_local_wizard("worker-id", source_root) + + archive = sandbox.write_file.call_args.args[1] + with tarfile.open(fileobj=io.BytesIO(archive), mode="r:gz") as value: + names = set(value.getnames()) + + assert "src/error.ts" in names + assert ".env" not in names + assert "node_modules/dependency.js" not in names + sandbox.write_file.assert_called_once_with(LOCAL_WIZARD_ARCHIVE_PATH, archive) + command = sandbox.execute.call_args.args[0] + assert LOCAL_WIZARD_INSTALL_PATH in command + assert "WIZARD_BUILD_NODE_ENV=ci pnpm exec tsdown" in command + assert "pnpm build" not in command + + +@patch("products.wizard.backend.logic.workers.service.get_sandbox_class") +def test_prepare_local_wizard_reports_build_stderr(get_sandbox_class: MagicMock, tmp_path: Path) -> None: + source_root = tmp_path / "wizard" + source_root.mkdir() + (source_root / "package.json").write_text('{"name":"@posthog/wizard"}') + (source_root / "pnpm-lock.yaml").write_text("lockfileVersion: '9.0'") + sandbox = get_sandbox_class.return_value.get_by_id.return_value + sandbox.write_file.return_value = _execution_result() + sandbox.execute.return_value = _execution_result( + stdout="Build complete", + stderr="Local build failed", + exit_code=1, + ) + + with pytest.raises(WizardWorkerExecutionError) as error: + prepare_local_wizard("worker-id", source_root) + + assert error.value.detail == "Build complete\nLocal build failed" + + +@parameterized.expand( + ( + ( + "default", + (), + f"node {WIZARD_PACKAGE_INSTALL_PATH}/node_modules/@posthog/wizard/dist/bin.js --headless-DONOTUSE-EXPERIMENTAL", + ), + ( + "nested", + ("audit", "web-analytics"), + f"node {WIZARD_PACKAGE_INSTALL_PATH}/node_modules/@posthog/wizard/dist/bin.js audit web-analytics --headless-DONOTUSE-EXPERIMENTAL", + ), + ) +) +@patch("products.wizard.backend.logic.workers.service.get_sandbox_class") +def test_execute_wizard_uses_selected_program( + _name: str, + program_command: tuple[str, ...], + expected_invocation: str, + get_sandbox_class: MagicMock, +) -> None: + request = WizardExecutionRequest( + sandbox_id="worker-id", + workspace_path="/tmp/workspace/repos/posthog/posthog", + team_id=7, + wizard_version="2.60.0", + program_command=program_command, + use_local_wizard_source=False, + ) + sandbox = get_sandbox_class.return_value.get_by_id.return_value + sandbox.execute.return_value = _execution_result() + + execute_wizard(request) + + get_sandbox_class.return_value.get_by_id.assert_called_once_with(request.sandbox_id) + command = sandbox.execute.call_args.args[0] + assert expected_invocation in command + assert f"--prefix {WIZARD_PACKAGE_INSTALL_PATH}" in command + assert "--registry=https://registry.npmjs.org" in command + assert command.index("npm ") < command.index(f"cd {request.workspace_path}") + assert "wizard-secret" not in command + + +@patch("products.wizard.backend.logic.workers.service.get_sandbox_class") +def test_execute_wizard_uses_prepared_local_wizard(get_sandbox_class: MagicMock) -> None: + request = WizardExecutionRequest( + sandbox_id="worker-id", + workspace_path="/tmp/workspace/repos/posthog/posthog", + team_id=7, + wizard_version="2.60.0", + program_command=("audit", "web-analytics"), + use_local_wizard_source=True, + ) + sandbox = get_sandbox_class.return_value.get_by_id.return_value + sandbox.execute.return_value = _execution_result() + + execute_wizard(request) + + command = sandbox.execute.call_args.args[0] + assert f"node {LOCAL_WIZARD_INSTALL_PATH}/dist/bin.js audit web-analytics" in command + assert "npm " not in command + + +@patch("products.wizard.backend.logic.workers.service.get_sandbox_class") +def test_execute_wizard_rejects_program_options(get_sandbox_class: MagicMock) -> None: + request = WizardExecutionRequest( + sandbox_id="worker-id", + workspace_path="/tmp/workspace/repos/posthog/posthog", + team_id=7, + wizard_version="2.60.0", + program_command=("--base-url",), + ) + + with pytest.raises(InvalidWizardProgramCommandError, match="Invalid Wizard program command"): + execute_wizard(request) + + get_sandbox_class.return_value.get_by_id.return_value.execute.assert_not_called() + + +@patch("products.wizard.backend.logic.workers.service.get_sandbox_class") +def test_execute_wizard_rejects_mutable_version(get_sandbox_class: MagicMock) -> None: + request = WizardExecutionRequest( + sandbox_id="worker-id", + workspace_path="/tmp/workspace/repos/posthog/posthog", + team_id=7, + wizard_version="^2.60.0", + program_command=(), + ) + + with pytest.raises(InvalidWizardVersionError, match="Invalid Wizard version"): + execute_wizard(request) + + get_sandbox_class.return_value.get_by_id.return_value.execute.assert_not_called() + + +@patch("products.wizard.backend.logic.workers.service.get_sandbox_class") +def test_execute_wizard_maps_command_timeout(get_sandbox_class: MagicMock) -> None: + request = WizardExecutionRequest( + sandbox_id="worker-id", + workspace_path="/tmp/workspace/repos/posthog/posthog", + team_id=7, + wizard_version="2.60.0", + program_command=(), + ) + get_sandbox_class.return_value.get_by_id.return_value.execute.return_value = _execution_result(exit_code=124) + + with pytest.raises(WizardWorkerTimeoutError): + execute_wizard(request) + + +@patch("products.wizard.backend.logic.workers.service.get_sandbox_class") +def test_execute_wizard_surfaces_wizard_error_code(get_sandbox_class: MagicMock) -> None: + request = WizardExecutionRequest( + sandbox_id="worker-id", + workspace_path="/tmp/workspace/repos/posthog/posthog", + team_id=7, + wizard_version="2.60.0", + program_command=("audit",), + ) + get_sandbox_class.return_value.get_by_id.return_value.execute.return_value = _execution_result( + stderr='Audit failed\nphw-error: {"code":"PHW_DETECT_NO_POSTHOG_SDK","message":"Missing SDK"}', + exit_code=1, + ) + + with pytest.raises(WizardWorkerExecutionError) as error: + execute_wizard(request) + + assert error.value.wizard_error_code == "PHW_DETECT_NO_POSTHOG_SDK" + + +@patch("products.wizard.backend.logic.workers.service.create_pull_request") +@patch("products.wizard.backend.logic.workers.service.create_signed_commit") +@patch("products.wizard.backend.logic.workers.service.get_sandbox_class") +def test_git_repository_handoff_captures_diff_and_publishes_pull_request( + get_sandbox_class: MagicMock, + create_signed_commit: MagicMock, + create_pull_request: MagicMock, +) -> None: + request = GitRepositoryHandoffRequest( + team_id=7, + run_id=uuid4(), + sandbox_id="worker-id", + workspace_path="/tmp/workspace/repos/posthog/posthog", + github_integration_id=17, + repository="PostHog/PostHog", + ) + sandbox = get_sandbox_class.return_value.get_by_id.return_value + sandbox.execute.side_effect = ( + _execution_result(stdout="diff --git a/a b/a\n"), + _execution_result(stdout="# Setup report\n\nAll done.\n"), + ) + branch = f"posthog/wizard-{request.run_id.hex[:12]}" + pull_request = RepositoryPullRequest( + repository=request.repository, + number=123, + url="https://github.com/posthog/posthog/pull/123", + head_branch=branch, + base_branch="master", + ) + create_pull_request.return_value = pull_request + + result = create_git_repository_handoff(request) + + assert result == WizardWorkerResult(diff=b"diff --git a/a b/a\n", pull_request=pull_request) + assert "git add -N --all" in sandbox.execute.call_args_list[0].args[0] + assert WIZARD_DIFF_OUTPUT_PATH in sandbox.execute.call_args_list[0].args[0] + assert f"head -c {MAX_GIT_DIFF_BYTES + 1}" in sandbox.execute.call_args_list[0].args[0] + assert wizard_handoff_output_path(request.run_id) in sandbox.execute.call_args_list[1].args[0] + assert "head -c 60000" in sandbox.execute.call_args_list[1].args[0] + create_signed_commit.assert_called_once_with( + sandbox, + team_id=request.team_id, + integration_id=request.github_integration_id, + repository=request.repository, + branch=branch, + message="Set up PostHog", + source="wizard", + ) + create_pull_request.assert_called_once_with( + team_id=request.team_id, + integration_id=request.github_integration_id, + repository=request.repository, + head_branch=branch, + title="Instrument application with PostHog", + body="# Setup report\n\nAll done.", + source="wizard", + ) + + +@pytest.mark.parametrize( + "handoff_result", + ( + _execution_result(exit_code=1), + _execution_result(stdout=" \n"), + ), +) +@patch("products.wizard.backend.logic.workers.service.create_pull_request") +@patch("products.wizard.backend.logic.workers.service.create_signed_commit") +@patch("products.wizard.backend.logic.workers.service.get_sandbox_class") +@patch("products.wizard.backend.logic.workers.service.wizard_observability.handoff_body_fallback") +def test_git_repository_handoff_uses_generic_body_when_handoff_is_unavailable( + handoff_body_fallback: MagicMock, + get_sandbox_class: MagicMock, + create_signed_commit: MagicMock, + create_pull_request: MagicMock, + handoff_result: SimpleNamespace, +) -> None: + request = GitRepositoryHandoffRequest( + team_id=7, + run_id=uuid4(), + sandbox_id="worker-id", + workspace_path="/tmp/workspace/repos/posthog/posthog", + github_integration_id=17, + repository="PostHog/PostHog", + ) + sandbox = get_sandbox_class.return_value.get_by_id.return_value + sandbox.execute.side_effect = (_execution_result(stdout="diff --git a/a b/a\n"), handoff_result) + + create_git_repository_handoff(request) + + assert create_pull_request.call_args.kwargs["body"] == ( + "This pull request contains changes created by Wizard, PostHog's setup agent." + ) + handoff_body_fallback.assert_called_once_with(request.team_id, request.run_id) + + +@patch("products.wizard.backend.logic.workers.service.create_pull_request") +@patch("products.wizard.backend.logic.workers.service.get_sandbox_class") +def test_git_repository_handoff_skips_publish_without_changes( + get_sandbox_class: MagicMock, + create_pull_request: MagicMock, +) -> None: + request = GitRepositoryHandoffRequest( + team_id=7, + run_id=uuid4(), + sandbox_id="worker-id", + workspace_path="/tmp/workspace/repos/posthog/posthog", + github_integration_id=17, + repository="PostHog/PostHog", + ) + get_sandbox_class.return_value.get_by_id.return_value.execute.return_value = _execution_result() + + result = create_git_repository_handoff(request) + + assert result == WizardWorkerResult(diff=b"", pull_request=None) + get_sandbox_class.return_value.get_by_id.return_value.execute.assert_called_once() + create_pull_request.assert_not_called() + + +@patch("products.wizard.backend.logic.workers.service.get_sandbox_class") +def test_destroy_worker_accepts_already_destroyed_sandbox(get_sandbox_class: MagicMock) -> None: + get_sandbox_class.return_value.get_by_id.side_effect = SandboxNotFoundError( + "Worker not found.", + {"sandbox_id": "worker-id"}, + RuntimeError("not found"), + ) + + destroy_worker("worker-id") diff --git a/products/wizard/backend/tests/runs/test_cloud_workflow.py b/products/wizard/backend/tests/runs/test_cloud_workflow.py new file mode 100644 index 000000000000..11d80c3e9d3b --- /dev/null +++ b/products/wizard/backend/tests/runs/test_cloud_workflow.py @@ -0,0 +1,268 @@ +import asyncio +from dataclasses import replace +from uuid import UUID, uuid4 + +import pytest +from unittest.mock import AsyncMock, MagicMock + +from temporalio.exceptions import ActivityError, ApplicationError, RetryState, TimeoutError, TimeoutType + +from products.wizard.backend.facade.enums import WizardRunErrorCode, WizardRunStatus, WizardWorkspaceType +from products.wizard.backend.temporal.activities.execution import WIZARD_WORKER_TIMEOUT_ERROR_TYPE, execute_wizard +from products.wizard.backend.temporal.activities.handoff import create_run_artifacts +from products.wizard.backend.temporal.activities.lifecycle import finalize_run +from products.wizard.backend.temporal.activities.local_package import prepare_local_wizard +from products.wizard.backend.temporal.activities.workspace import clone_repository, destroy_worker, provision_worker +from products.wizard.backend.temporal.contracts import ( + PreparedGitRepositoryWorkspace, + ProvisionedWizardWorker, + WizardRunActivityInput, + WizardRunFinalizationActivityInput, +) +from products.wizard.backend.temporal.workflows import execute_run as execute_run_workflow_module +from products.wizard.backend.temporal.workflows.execute_run import ExecuteWizardRunWorkflow + + +def _activity_error(cause: BaseException, activity_type: str = "wizard_execute") -> ActivityError: + error = ActivityError( + "Activity failed", + scheduled_event_id=1, + started_event_id=2, + identity="worker", + activity_type=activity_type, + activity_id="activity", + retry_state=RetryState.MAXIMUM_ATTEMPTS_REACHED, + ) + error.__cause__ = cause + return error + + +@pytest.fixture +def workflow_input() -> WizardRunActivityInput: + return WizardRunActivityInput(team_id=1, run_id=uuid4()) + + +@pytest.fixture +def worker(workflow_input: WizardRunActivityInput) -> ProvisionedWizardWorker: + return ProvisionedWizardWorker( + team_id=workflow_input.team_id, + run_id=workflow_input.run_id, + sandbox_id="worker", + workspace_type=WizardWorkspaceType.GIT_REPOSITORY, + ) + + +@pytest.fixture +def workspace( + workflow_input: WizardRunActivityInput, + worker: ProvisionedWizardWorker, +) -> PreparedGitRepositoryWorkspace: + return PreparedGitRepositoryWorkspace( + team_id=workflow_input.team_id, + run_id=workflow_input.run_id, + sandbox_id=worker.sandbox_id, + repository="posthog/posthog", + root_path="/tmp/workspace/repos/posthog/posthog", + github_integration_id=123, + ) + + +@pytest.mark.asyncio +async def test_cloud_workflow_completes_after_worker_execution( + monkeypatch: pytest.MonkeyPatch, + workflow_input: WizardRunActivityInput, + worker: ProvisionedWizardWorker, + workspace: PreparedGitRepositoryWorkspace, +) -> None: + execute_activity = AsyncMock(side_effect=[worker, workspace, None, None, None]) + monkeypatch.setattr(execute_run_workflow_module.workflow, "execute_activity", execute_activity) + + await ExecuteWizardRunWorkflow().run(workflow_input) + + assert [call.args[0] for call in execute_activity.await_args_list] == [ + provision_worker, + clone_repository, + execute_wizard, + create_run_artifacts, + destroy_worker, + ] + assert execute_activity.await_args_list[0].args[1] == workflow_input + assert execute_activity.await_args_list[1].args[1] == worker + assert execute_activity.await_args_list[2].args[1] == workspace + assert execute_activity.await_args_list[3].args[1] == workspace + assert execute_activity.await_args_list[4].args[1] == worker + assert execute_activity.await_args_list[1].kwargs["retry_policy"].maximum_attempts == 3 + assert execute_activity.await_args_list[2].kwargs["retry_policy"].maximum_attempts == 1 + assert execute_activity.await_args_list[3].kwargs["retry_policy"].maximum_attempts == 3 + assert execute_activity.await_args_list[0].kwargs["retry_policy"].maximum_attempts == 1 + assert execute_activity.await_args_list[4].kwargs["retry_policy"].maximum_attempts == 3 + + +@pytest.mark.asyncio +async def test_cloud_workflow_prepares_local_wizard_before_workspace( + monkeypatch: pytest.MonkeyPatch, + workflow_input: WizardRunActivityInput, + worker: ProvisionedWizardWorker, + workspace: PreparedGitRepositoryWorkspace, +) -> None: + local_input = replace(workflow_input, use_local_wizard_source=True) + local_worker = replace(worker, use_local_wizard_source=True) + local_workspace = replace(workspace, use_local_wizard_source=True) + execute_activity = AsyncMock(side_effect=[local_worker, None, local_workspace, None, None, None]) + monkeypatch.setattr(execute_run_workflow_module.workflow, "execute_activity", execute_activity) + + await ExecuteWizardRunWorkflow().run(local_input) + + assert [call.args[0] for call in execute_activity.await_args_list] == [ + provision_worker, + prepare_local_wizard, + clone_repository, + execute_wizard, + create_run_artifacts, + destroy_worker, + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("cause", "expected_error_code"), + [ + ( + ApplicationError("Wizard Worker timed out.", type=WIZARD_WORKER_TIMEOUT_ERROR_TYPE), + WizardRunErrorCode.TIMEOUT, + ), + ( + TimeoutError("Activity timed out", type=TimeoutType.START_TO_CLOSE, last_heartbeat_details=[]), + WizardRunErrorCode.TIMEOUT, + ), + ( + ApplicationError("Wizard failed", type="PHW_DETECT_NO_POSTHOG_SDK"), + "PHW_DETECT_NO_POSTHOG_SDK", + ), + (ApplicationError("Worker failed", type="WorkerFailure"), WizardRunErrorCode.EXECUTION_FAILED), + ], +) +async def test_cloud_workflow_persists_execution_failure( + monkeypatch: pytest.MonkeyPatch, + workflow_input: WizardRunActivityInput, + worker: ProvisionedWizardWorker, + workspace: PreparedGitRepositoryWorkspace, + cause: BaseException, + expected_error_code: str, +) -> None: + activity_error = _activity_error(cause) + execute_activity = AsyncMock(side_effect=[worker, workspace, activity_error, None, None]) + monkeypatch.setattr(execute_run_workflow_module.workflow, "execute_activity", execute_activity) + + with pytest.raises(ActivityError) as raised: + await ExecuteWizardRunWorkflow().run(workflow_input) + + assert raised.value is activity_error + assert execute_activity.await_args_list[3].args == (destroy_worker, worker) + assert execute_activity.await_args_list[4].args == ( + finalize_run, + WizardRunFinalizationActivityInput( + team_id=workflow_input.team_id, + run_id=workflow_input.run_id, + status=WizardRunStatus.FAILED, + error_code=expected_error_code, + ), + ) + + +@pytest.mark.asyncio +async def test_cloud_workflow_persists_cancellation( + monkeypatch: pytest.MonkeyPatch, + workflow_input: WizardRunActivityInput, + worker: ProvisionedWizardWorker, + workspace: PreparedGitRepositoryWorkspace, +) -> None: + execute_activity = AsyncMock(side_effect=[worker, workspace, asyncio.CancelledError(), None, None]) + monkeypatch.setattr(execute_run_workflow_module.workflow, "execute_activity", execute_activity) + + with pytest.raises(asyncio.CancelledError): + await ExecuteWizardRunWorkflow().run(workflow_input) + + assert execute_activity.await_args_list[3].args == (destroy_worker, worker) + assert execute_activity.await_args_list[4].args == ( + finalize_run, + WizardRunFinalizationActivityInput( + team_id=workflow_input.team_id, + run_id=workflow_input.run_id, + status=WizardRunStatus.CANCELLED, + ), + ) + + +@pytest.mark.asyncio +async def test_cloud_workflow_persists_unexpected_failure( + monkeypatch: pytest.MonkeyPatch, + workflow_input: WizardRunActivityInput, +) -> None: + unexpected_error = RuntimeError("unexpected workflow failure") + execute_activity = AsyncMock(side_effect=[unexpected_error, None]) + workflow_logger = MagicMock() + monkeypatch.setattr(execute_run_workflow_module.workflow, "execute_activity", execute_activity) + monkeypatch.setattr(execute_run_workflow_module.workflow, "logger", workflow_logger) + + with pytest.raises(RuntimeError) as raised: + await ExecuteWizardRunWorkflow().run(workflow_input) + + assert raised.value is unexpected_error + workflow_logger.exception.assert_called_once_with( + "wizard_run_workflow_failed", + extra={"team_id": workflow_input.team_id, "run_id": str(workflow_input.run_id)}, + ) + assert execute_activity.await_args_list[1].args == ( + finalize_run, + WizardRunFinalizationActivityInput( + team_id=workflow_input.team_id, + run_id=workflow_input.run_id, + status=WizardRunStatus.FAILED, + error_code=WizardRunErrorCode.EXECUTION_FAILED, + ), + ) + + +@pytest.mark.asyncio +async def test_cloud_workflow_keeps_success_when_worker_cleanup_fails( + monkeypatch: pytest.MonkeyPatch, + workflow_input: WizardRunActivityInput, + worker: ProvisionedWizardWorker, + workspace: PreparedGitRepositoryWorkspace, +) -> None: + cleanup_error = _activity_error(ApplicationError("Worker cleanup failed", type="CleanupFailure")) + execute_activity = AsyncMock(side_effect=[worker, workspace, None, None, cleanup_error]) + workflow_logger = MagicMock() + monkeypatch.setattr(execute_run_workflow_module.workflow, "execute_activity", execute_activity) + monkeypatch.setattr(execute_run_workflow_module.workflow, "logger", workflow_logger) + + await ExecuteWizardRunWorkflow().run(workflow_input) + + assert len(execute_activity.await_args_list) == 5 + workflow_logger.exception.assert_called_once() + + +def test_cloud_workflow_identity_and_input_parsing() -> None: + run_id = uuid4() + + assert ExecuteWizardRunWorkflow.workflow_id_for(run_id) == f"wizard-run-{run_id}" + assert ExecuteWizardRunWorkflow.parse_inputs([f'{{"team_id": 1, "run_id": "{run_id}"}}']) == WizardRunActivityInput( + team_id=1, + run_id=UUID(str(run_id)), + ) + + +@pytest.mark.parametrize( + ("activity_type", "expected_error_code"), + ( + ("wizard_provision_worker", "provisioning_failed"), + ("wizard_clone_repository", "workspace_preparation_failed"), + ("wizard_execute", "execution_failed"), + ("wizard_create_run_artifacts", "artifact_creation_failed"), + ), +) +def test_activity_failures_map_to_their_stage(activity_type: str, expected_error_code: str) -> None: + error = _activity_error(ApplicationError("failed", type="ExternalFailure"), activity_type) + + assert execute_run_workflow_module.wizard_run_error_code(error) == expected_error_code diff --git a/products/wizard/backend/tests/runs/test_command_output.py b/products/wizard/backend/tests/runs/test_command_output.py new file mode 100644 index 000000000000..23e81de835c4 --- /dev/null +++ b/products/wizard/backend/tests/runs/test_command_output.py @@ -0,0 +1,63 @@ +import sys +import shlex +import subprocess +from types import SimpleNamespace + +import pytest +from unittest.mock import MagicMock, patch + +from products.wizard.backend.logic.workers.commands import MAX_SANDBOX_LOG_BYTES, bound_command_output +from products.wizard.backend.logic.workers.repository_publisher import ( + RepositoryPublishingError, + _run_git, + _staged_changes, +) + + +@pytest.mark.parametrize("exit_code", [0, 7, 124]) +def test_command_output_keeps_bounded_tails_and_exit_code(exit_code: int) -> None: + script = ( + f"import sys; sys.stdout.write('a' * {MAX_SANDBOX_LOG_BYTES * 4} + 'out'); " + f"sys.stderr.write('b' * {MAX_SANDBOX_LOG_BYTES * 4} + 'err'); sys.exit({exit_code})" + ) + command = bound_command_output(f"{shlex.quote(sys.executable)} -c {shlex.quote(script)}") + result = subprocess.run(["bash", "-c", command], capture_output=True, timeout=10) + assert result.returncode == exit_code + assert result.stdout == b"a" * (MAX_SANDBOX_LOG_BYTES - 3) + b"out" + assert result.stderr == b"b" * (MAX_SANDBOX_LOG_BYTES - 3) + b"err" + + +@pytest.mark.parametrize("size", [31, 32, 33, 100000]) +def test_git_output_rejects_overflow_before_returning_contents(size: int) -> None: + sandbox = MagicMock() + + def execute(command: str, *, timeout_seconds: int) -> SimpleNamespace: + result = subprocess.run(["bash", "-c", command], capture_output=True, text=True, timeout=timeout_seconds) + assert len(result.stdout) <= 33 + return SimpleNamespace(stdout=result.stdout, exit_code=result.returncode) + + sandbox.execute.side_effect = execute + arguments = f"--version >/dev/null; printf '%0{size}d' 0" + if size > 32: + with pytest.raises(RepositoryPublishingError, match="payload limit"): + _run_git(sandbox, ".", arguments, "test", max_output_bytes=32) + else: + assert _run_git(sandbox, ".", arguments, "test", max_output_bytes=32) == "0" * size + + +@pytest.mark.parametrize("budget", [73, 74, 75]) +def test_staged_changes_enforce_cumulative_encoded_budget(budget: int) -> None: + sandbox = MagicMock() + sandbox.execute.side_effect = [ + SimpleNamespace(stdout="A\0a\0A\0b\0", exit_code=0), + SimpleNamespace(stdout="YQ==", exit_code=0), + SimpleNamespace(stdout="Yg==", exit_code=0), + ] + with patch("products.wizard.backend.logic.workers.repository_publisher.MAX_COMMIT_PAYLOAD_BYTES", budget): + if budget < 74: + with pytest.raises(RepositoryPublishingError, match="payload limit"): + _staged_changes(sandbox, ".", "HEAD") + else: + changes = _staged_changes(sandbox, ".", "HEAD") + assert changes.additions == [("a", "YQ=="), ("b", "Yg==")] + assert changes.deletions == [] diff --git a/products/wizard/backend/tests/runs/test_creation.py b/products/wizard/backend/tests/runs/test_creation.py new file mode 100644 index 000000000000..394c7e31e569 --- /dev/null +++ b/products/wizard/backend/tests/runs/test_creation.py @@ -0,0 +1,388 @@ +import pytest +from unittest.mock import patch + +from django.db import IntegrityError, transaction +from django.db.models import NOT_PROVIDED + +from products.wizard.backend.facade import api as wizard_facade +from products.wizard.backend.facade.config import DEFAULT_WIZARD_VERSION +from products.wizard.backend.facade.contracts import ( + CreateWizardRunInput, + GitRepositoryWorkspace, + LocalFolderWorkspace, + WizardProgram, +) +from products.wizard.backend.facade.enums import WizardRunEnvironment, WizardRunStatus +from products.wizard.backend.facade.errors import ( + InvalidRepositoryError, + InvalidWorkspaceEnvironmentError, + MissingGitHubIntegrationError, + RepositoryNotAccessibleError, + WizardProgramEnvironmentNotSupportedError, +) +from products.wizard.backend.models import WizardRun +from products.wizard.backend.temporal.contracts import WizardRunActivityInput +from products.wizard.backend.temporal.errors import WizardTemporalError + +PROGRAM_DEFINITION = { + "id": "web-analytics-audit", + "name": "Web analytics audit", + "description": "Audit a project's web analytics setup", + "wizard_version": "2.60.0", + "command": ["audit", "web-analytics"], + "tags": ["audit", "web-analytics"], + "required_programs": ["posthog-integration"], + "supported_environments": ["local"], +} +PROGRAM_PAYLOAD = {"version": 1, "programs": [PROGRAM_DEFINITION]} + + +@pytest.mark.django_db +def test_create_run_persists_resolved_program(team, user) -> None: + with patch("posthoganalytics.get_feature_flag_payload", return_value=PROGRAM_PAYLOAD): + run = wizard_facade.create_run( + CreateWizardRunInput( + team_id=team.id, + created_by_id=user.id, + environment=WizardRunEnvironment.LOCAL, + workspace=LocalFolderWorkspace(project_name="example-project"), + program_id="web-analytics-audit", + ) + ) + + assert run.program == WizardProgram( + id="web-analytics-audit", + name="Web analytics audit", + description="Audit a project's web analytics setup", + wizard_version=DEFAULT_WIZARD_VERSION, + command=("audit", "web-analytics"), + tags=("audit", "web-analytics"), + required_programs=("posthog-integration",), + supported_environments=(WizardRunEnvironment.LOCAL,), + ) + assert WizardRun.objects.for_team(team.id).get(id=run.id).program == { + **PROGRAM_DEFINITION, + "wizard_version": DEFAULT_WIZARD_VERSION, + } + + +@pytest.mark.django_db +def test_create_run_rejects_unsupported_program_environment(team, user) -> None: + with ( + patch("posthoganalytics.get_feature_flag_payload", return_value=PROGRAM_PAYLOAD), + pytest.raises(WizardProgramEnvironmentNotSupportedError), + ): + wizard_facade.create_run( + CreateWizardRunInput( + team_id=team.id, + created_by_id=user.id, + environment=WizardRunEnvironment.CLOUD, + idempotency_key="unsupported-program-environment", + workspace=GitRepositoryWorkspace(repository="posthog/posthog"), + program_id="web-analytics-audit", + ) + ) + + assert not WizardRun.objects.for_team(team.id).exists() + + +@pytest.mark.django_db +def test_run_program_has_no_default_and_cannot_be_null(team, user) -> None: + program_field = WizardRun._meta.get_field("program") + + assert program_field.default is NOT_PROVIDED + assert program_field.null is False + with pytest.raises(IntegrityError), transaction.atomic(): + WizardRun.objects.create( + team=team, + created_by=user, + environment=WizardRunEnvironment.LOCAL.value, + workspace_type="local_folder", + workspace={"project_name": "example-project"}, + status=WizardRunStatus.RUNNING.value, + ) + + +@pytest.mark.django_db +def test_local_run_starts_running(team, user) -> None: + run = wizard_facade.create_run( + CreateWizardRunInput( + team_id=team.id, + created_by_id=user.id, + program_id="posthog-integration", + environment=WizardRunEnvironment.LOCAL, + workspace=LocalFolderWorkspace(project_name="example-project"), + ) + ) + + assert run.team_id == team.id + assert run.created_by_id == user.id + assert run.environment == WizardRunEnvironment.LOCAL + assert run.workspace == LocalFolderWorkspace(project_name="example-project") + assert run.status == WizardRunStatus.RUNNING + assert run.error_code is None + + +@pytest.mark.django_db +def test_create_run_rejects_unsupported_environment_workspace(team, user) -> None: + with pytest.raises(InvalidWorkspaceEnvironmentError): + wizard_facade.create_run( + CreateWizardRunInput( + team_id=team.id, + created_by_id=user.id, + program_id="posthog-integration", + environment=WizardRunEnvironment.LOCAL, + workspace=GitRepositoryWorkspace(repository="posthog/posthog"), + ) + ) + + assert not WizardRun.objects.for_team(team.id).exists() + + +@pytest.mark.django_db +def test_cloud_run_starts_created(team, user) -> None: + with ( + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.resolve_team_github_integration_id", + return_value=123, + ), + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.repository_accessible_via_integration", + return_value=True, + ), + ): + run = wizard_facade.create_run( + CreateWizardRunInput( + team_id=team.id, + created_by_id=user.id, + program_id="posthog-integration", + environment=WizardRunEnvironment.CLOUD, + idempotency_key="cloud-run-created", + workspace=GitRepositoryWorkspace(repository="posthog/posthog"), + ) + ) + + assert run.status == WizardRunStatus.CREATED + assert run.workspace == GitRepositoryWorkspace(repository="posthog/posthog") + + +@pytest.mark.django_db(transaction=True) +def test_cloud_run_dispatches_after_persistence(team, user) -> None: + with ( + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.resolve_team_github_integration_id", + return_value=123, + ), + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.repository_accessible_via_integration", + return_value=True, + ), + patch( + "products.wizard.backend.logic.runs.lifecycle.dispatch_created_cloud_wizard_run_to_temporal_worker" + ) as dispatch_wizard_run, + ): + run = wizard_facade.create_run( + CreateWizardRunInput( + team_id=team.id, + created_by_id=user.id, + program_id="posthog-integration", + environment=WizardRunEnvironment.CLOUD, + idempotency_key="cloud-run-dispatch", + workspace=GitRepositoryWorkspace(repository="posthog/posthog"), + ) + ) + + dispatch_wizard_run.assert_called_once_with(team.id, run.id) + assert run.status == WizardRunStatus.CREATED + + +@pytest.mark.django_db(transaction=True) +def test_cloud_run_cancellation_during_dispatch_reaches_temporal(team, user) -> None: + def cancel_during_dispatch(input: WizardRunActivityInput) -> None: + wizard_facade.cancel_run(input.team_id, input.run_id) + + with ( + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.resolve_team_github_integration_id", + return_value=123, + ), + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.repository_accessible_via_integration", + return_value=True, + ), + patch( + "products.wizard.backend.logic.runs.dispatch.temporal_client.start_wizard_run_workflow", + side_effect=cancel_during_dispatch, + ), + patch( + "products.wizard.backend.logic.runs.cancellation.temporal_client.cancel_wizard_run_workflow" + ) as cancel_workflow, + ): + run = wizard_facade.create_run( + CreateWizardRunInput( + team_id=team.id, + created_by_id=user.id, + program_id="posthog-integration", + environment=WizardRunEnvironment.CLOUD, + idempotency_key="cloud-run-cancel-during-dispatch", + workspace=GitRepositoryWorkspace(repository="posthog/posthog"), + ) + ) + + record = WizardRun.objects.for_team(team.id).get(id=run.id) + cancel_workflow.assert_called_once_with(run.id) + assert record.status == WizardRunStatus.CANCELLED.value + assert record.cancellation_requested_at is not None + assert record.cancellation_dispatched_at is not None + + +@pytest.mark.django_db(transaction=True) +def test_cloud_run_survives_temporal_dispatch_failure(team, user) -> None: + with ( + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.resolve_team_github_integration_id", + return_value=123, + ), + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.repository_accessible_via_integration", + return_value=True, + ), + patch( + "products.wizard.backend.logic.runs.dispatch.temporal_client.start_wizard_run_workflow", + side_effect=WizardTemporalError, + ), + ): + run = wizard_facade.create_run( + CreateWizardRunInput( + team_id=team.id, + created_by_id=user.id, + program_id="posthog-integration", + environment=WizardRunEnvironment.CLOUD, + idempotency_key="cloud-run-dispatch-failure", + workspace=GitRepositoryWorkspace(repository="posthog/posthog"), + ) + ) + + record = WizardRun.objects.for_team(team.id).get(id=run.id) + assert run.status == WizardRunStatus.CREATED + assert run.error_code is None + assert record.dispatch_status == "pending" + assert record.dispatch_attempts == 1 + assert record.dispatch_error == "Temporal dispatch failed." + + +@pytest.mark.django_db(transaction=True) +def test_cloud_run_rollback_prevents_dispatch(team, user) -> None: + with ( + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.resolve_team_github_integration_id", + return_value=123, + ), + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.repository_accessible_via_integration", + return_value=True, + ), + patch( + "products.wizard.backend.logic.runs.lifecycle.dispatch_created_cloud_wizard_run_to_temporal_worker" + ) as dispatch, + transaction.atomic(), + ): + wizard_facade.create_run( + CreateWizardRunInput( + team_id=team.id, + created_by_id=user.id, + program_id="posthog-integration", + environment=WizardRunEnvironment.CLOUD, + idempotency_key="cloud-run-rollback", + workspace=GitRepositoryWorkspace(repository="posthog/posthog"), + ) + ) + transaction.set_rollback(True) + + dispatch.assert_not_called() + + +@pytest.mark.django_db(transaction=True) +def test_local_run_does_not_dispatch(team, user) -> None: + with patch( + "products.wizard.backend.logic.runs.lifecycle.dispatch_created_cloud_wizard_run_to_temporal_worker" + ) as dispatch: + wizard_facade.create_run( + CreateWizardRunInput( + team_id=team.id, + created_by_id=user.id, + program_id="posthog-integration", + environment=WizardRunEnvironment.LOCAL, + workspace=LocalFolderWorkspace(project_name="example-project"), + ) + ) + + dispatch.assert_not_called() + + +@pytest.mark.django_db +def test_cloud_run_requires_github_integration(team, user) -> None: + with patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.resolve_team_github_integration_id", + return_value=None, + ): + with pytest.raises(MissingGitHubIntegrationError): + wizard_facade.create_run( + CreateWizardRunInput( + team_id=team.id, + created_by_id=user.id, + program_id="posthog-integration", + environment=WizardRunEnvironment.CLOUD, + idempotency_key="missing-integration", + workspace=GitRepositoryWorkspace(repository="posthog/posthog"), + ) + ) + + assert not WizardRun.objects.for_team(team.id).exists() + + +@pytest.mark.django_db +def test_cloud_run_rejects_inaccessible_repository(team, user) -> None: + with ( + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.resolve_team_github_integration_id", + return_value=123, + ), + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.repository_accessible_via_integration", + return_value=False, + ), + ): + with pytest.raises(RepositoryNotAccessibleError): + wizard_facade.create_run( + CreateWizardRunInput( + team_id=team.id, + created_by_id=user.id, + program_id="posthog-integration", + environment=WizardRunEnvironment.CLOUD, + idempotency_key="inaccessible-repository", + workspace=GitRepositoryWorkspace(repository="private/example"), + ) + ) + + +@pytest.mark.parametrize("repository", ("posthog", "/posthog", "posthog/", "posthog/posthog/extra")) +@pytest.mark.django_db +def test_cloud_run_rejects_invalid_repository_before_github_lookup(team, user, repository: str) -> None: + with patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.resolve_team_github_integration_id" + ) as resolve: + with pytest.raises(InvalidRepositoryError): + wizard_facade.create_run( + CreateWizardRunInput( + team_id=team.id, + created_by_id=user.id, + program_id="posthog-integration", + environment=WizardRunEnvironment.CLOUD, + idempotency_key=f"repository-{repository}", + workspace=GitRepositoryWorkspace(repository=repository), + ) + ) + + resolve.assert_not_called() + assert not WizardRun.objects.for_team(team.id).exists() diff --git a/products/wizard/backend/tests/runs/test_finalization_activity.py b/products/wizard/backend/tests/runs/test_finalization_activity.py new file mode 100644 index 000000000000..098b3a38ed3a --- /dev/null +++ b/products/wizard/backend/tests/runs/test_finalization_activity.py @@ -0,0 +1,133 @@ +import pytest +from unittest.mock import patch + +from asgiref.sync import async_to_sync +from temporalio.testing import ActivityEnvironment + +from products.wizard.backend.facade import api as wizard_facade +from products.wizard.backend.facade.contracts import CreateWizardRunInput, GitRepositoryWorkspace, LocalFolderWorkspace +from products.wizard.backend.facade.enums import WizardRunEnvironment, WizardRunErrorCode, WizardRunStatus +from products.wizard.backend.temporal.activities.lifecycle import finalize_run, transition_cloud_run +from products.wizard.backend.temporal.contracts import WizardRunFinalizationActivityInput + + +async def _run_activity(input: WizardRunFinalizationActivityInput) -> None: + await ActivityEnvironment().run(finalize_run, input) + + +@pytest.mark.django_db(transaction=True) +@pytest.mark.parametrize( + ("status", "error_code"), + ( + (WizardRunStatus.FAILED, WizardRunErrorCode.TIMEOUT), + (WizardRunStatus.FAILED, WizardRunErrorCode.EXECUTION_FAILED), + (WizardRunStatus.FAILED, "PHW_DETECT_NO_POSTHOG_SDK"), + (WizardRunStatus.CANCELLED, None), + ), +) +def test_finalize_run_is_retry_safe(team, user, status: WizardRunStatus, error_code: str | None) -> None: + with ( + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.resolve_team_github_integration_id", + return_value=123, + ), + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.repository_accessible_via_integration", + return_value=True, + ), + ): + run = wizard_facade.create_run( + CreateWizardRunInput( + team_id=team.id, + created_by_id=user.id, + program_id="posthog-integration", + environment=WizardRunEnvironment.CLOUD, + idempotency_key="test-finalization", + workspace=GitRepositoryWorkspace(repository="posthog/posthog"), + ) + ) + wizard_facade.update_run_status(team.id, run.id, WizardRunStatus.RUNNING) + input = WizardRunFinalizationActivityInput( + team_id=team.id, + run_id=run.id, + status=status, + error_code=error_code, + ) + + async_to_sync(_run_activity)(input) + async_to_sync(_run_activity)(input) + + persisted = wizard_facade.get_run(team.id, run.id) + assert persisted.status == status + assert persisted.error_code == error_code + + +@pytest.mark.django_db(transaction=True) +@pytest.mark.parametrize( + ("finalizer_status", "finalizer_error_code"), + ( + (WizardRunStatus.FAILED, WizardRunErrorCode.EXECUTION_FAILED), + (WizardRunStatus.COMPLETED, None), + ), +) +def test_finalize_run_preserves_user_cancellation( + team, user, finalizer_status: WizardRunStatus, finalizer_error_code: str | None +) -> None: + with ( + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.resolve_team_github_integration_id", + return_value=123, + ), + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.repository_accessible_via_integration", + return_value=True, + ), + ): + run = wizard_facade.create_run( + CreateWizardRunInput( + team_id=team.id, + created_by_id=user.id, + program_id="posthog-integration", + environment=WizardRunEnvironment.CLOUD, + idempotency_key="test-cancellation-race", + workspace=GitRepositoryWorkspace(repository="posthog/posthog"), + ) + ) + wizard_facade.update_run_status(team.id, run.id, WizardRunStatus.RUNNING) + wizard_facade.update_run_status(team.id, run.id, WizardRunStatus.CANCELLED) + + if finalizer_status == WizardRunStatus.COMPLETED: + transition_cloud_run(team.id, run.id, finalizer_status, error_code=finalizer_error_code) + else: + input = WizardRunFinalizationActivityInput( + team_id=team.id, + run_id=run.id, + status=finalizer_status, + error_code=finalizer_error_code, + ) + async_to_sync(_run_activity)(input) + + persisted = wizard_facade.get_run(team.id, run.id) + assert persisted.status == WizardRunStatus.CANCELLED + assert persisted.error_code is None + + +@pytest.mark.django_db(transaction=True) +def test_finalize_run_rejects_local_run(team, user) -> None: + run = wizard_facade.create_run( + CreateWizardRunInput( + team_id=team.id, + created_by_id=user.id, + program_id="posthog-integration", + environment=WizardRunEnvironment.LOCAL, + workspace=LocalFolderWorkspace(project_name="example-project"), + ) + ) + input = WizardRunFinalizationActivityInput( + team_id=team.id, + run_id=run.id, + status=WizardRunStatus.CANCELLED, + ) + + with pytest.raises(ValueError, match="cloud Wizard Run"): + async_to_sync(_run_activity)(input) diff --git a/products/wizard/backend/tests/runs/test_lifecycle.py b/products/wizard/backend/tests/runs/test_lifecycle.py new file mode 100644 index 000000000000..2103c137da5c --- /dev/null +++ b/products/wizard/backend/tests/runs/test_lifecycle.py @@ -0,0 +1,143 @@ +import pytest +from unittest.mock import patch + +from posthog.models import Team + +from products.wizard.backend.facade import api as wizard_facade +from products.wizard.backend.facade.contracts import ( + CreateWizardRunInput, + GitRepositoryWorkspace, + LocalFolderWorkspace, + WizardRunDTO, +) +from products.wizard.backend.facade.enums import WizardRunEnvironment, WizardRunErrorCode, WizardRunStatus +from products.wizard.backend.facade.errors import IllegalStatusTransitionError, WizardRunNotFoundError + + +def _create_local_run(team_id: int, user_id: int) -> WizardRunDTO: + return wizard_facade.create_run( + CreateWizardRunInput( + team_id=team_id, + created_by_id=user_id, + program_id="posthog-integration", + environment=WizardRunEnvironment.LOCAL, + workspace=LocalFolderWorkspace(project_name="example-project"), + ) + ) + + +def _create_cloud_run(team_id: int, user_id: int) -> WizardRunDTO: + with ( + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.resolve_team_github_integration_id", + return_value=123, + ), + patch( + "products.wizard.backend.logic.runs.repository_access.repo_selection.repository_accessible_via_integration", + return_value=True, + ), + ): + return wizard_facade.create_run( + CreateWizardRunInput( + team_id=team_id, + created_by_id=user_id, + program_id="posthog-integration", + environment=WizardRunEnvironment.CLOUD, + idempotency_key="test-lifecycle", + workspace=GitRepositoryWorkspace(repository="posthog/posthog"), + ) + ) + + +@pytest.mark.django_db +def test_get_run_is_scoped_to_team(team, user) -> None: + other_team = Team.objects.create(organization=team.organization, project=team.project, name="Other environment") + created = _create_local_run(team.id, user.id) + + assert wizard_facade.get_run(team.id, created.id) == created + + with pytest.raises(WizardRunNotFoundError): + wizard_facade.get_run(other_team.id, created.id) + + +@pytest.mark.django_db +def test_start_run_persists_running_status(team, user) -> None: + created = _create_cloud_run(team.id, user.id) + + started = wizard_facade.update_run_status(team.id, created.id, WizardRunStatus.RUNNING) + + assert started.status == WizardRunStatus.RUNNING + assert started.started_at is not None + assert started.finished_at is None + assert wizard_facade.get_run(team.id, created.id) == started + + +@pytest.mark.parametrize("expected_status", (WizardRunStatus.COMPLETED, WizardRunStatus.CANCELLED)) +@pytest.mark.django_db +def test_running_run_persists_terminal_status(team, user, expected_status: WizardRunStatus) -> None: + created = _create_local_run(team.id, user.id) + + if expected_status == WizardRunStatus.CANCELLED: + transitioned = wizard_facade.cancel_run(team.id, created.id) + else: + transitioned = wizard_facade.update_run_status(team.id, created.id, expected_status) + + assert transitioned.status == expected_status + assert transitioned.finished_at is not None + assert transitioned.stage is None + assert wizard_facade.get_run(team.id, created.id) == transitioned + + +@pytest.mark.django_db +def test_fail_run_persists_error_code(team, user) -> None: + created = _create_local_run(team.id, user.id) + + failed = wizard_facade.update_run_status( + team.id, + created.id, + WizardRunStatus.FAILED, + error_code=WizardRunErrorCode.TIMEOUT, + ) + + assert failed.status == WizardRunStatus.FAILED + assert failed.error_code == WizardRunErrorCode.TIMEOUT + assert failed.error_message == "The Wizard run timed out." + assert failed.finished_at is not None + assert wizard_facade.get_run(team.id, created.id) == failed + + +@pytest.mark.django_db +def test_fail_run_persists_wizard_error_code(team, user) -> None: + created = _create_local_run(team.id, user.id) + + failed = wizard_facade.update_run_status( + team.id, + created.id, + WizardRunStatus.FAILED, + error_code="PHW_DETECT_NO_POSTHOG_SDK", + ) + + assert failed.error_code == "PHW_DETECT_NO_POSTHOG_SDK" + assert failed.error_message == "The Wizard could not complete the selected program." + + +@pytest.mark.django_db +def test_invalid_persisted_transition_leaves_run_unchanged(team, user) -> None: + created = _create_local_run(team.id, user.id) + completed = wizard_facade.update_run_status(team.id, created.id, WizardRunStatus.COMPLETED) + + with pytest.raises(IllegalStatusTransitionError): + wizard_facade.cancel_run(team.id, created.id) + + assert wizard_facade.get_run(team.id, created.id) == completed + + +@pytest.mark.django_db +def test_transition_run_is_scoped_to_team(team, user) -> None: + other_team = Team.objects.create(organization=team.organization, project=team.project, name="Other environment") + created = _create_cloud_run(team.id, user.id) + + with pytest.raises(WizardRunNotFoundError): + wizard_facade.update_run_status(other_team.id, created.id, WizardRunStatus.RUNNING) + + assert wizard_facade.get_run(team.id, created.id) == created diff --git a/products/wizard/backend/tests/runs/test_observability.py b/products/wizard/backend/tests/runs/test_observability.py index 713c240891fb..d16e52f12dc0 100644 --- a/products/wizard/backend/tests/runs/test_observability.py +++ b/products/wizard/backend/tests/runs/test_observability.py @@ -1,19 +1,42 @@ import logging -from datetime import UTC, datetime -from uuid import UUID +from dataclasses import replace +from datetime import UTC, datetime, timedelta +from uuid import UUID, uuid4 +import pytest from unittest import TestCase from unittest.mock import patch -from kombu.exceptions import OperationalError +from django.db import DatabaseError +from django.test import override_settings + from parameterized import parameterized +from prometheus_client import REGISTRY -from products.wizard.backend.facade.contracts import LocalFolderWorkspace, WizardRunDTO -from products.wizard.backend.facade.enums import WizardRunEnvironment, WizardRunStatus +from products.wizard.backend.facade.contracts import ( + CreateWizardRunInput, + GitRepositoryWorkspace, + LocalFolderWorkspace, + WizardRunDTO, + WizardRunGitDiffArtifactDTO, +) +from products.wizard.backend.facade.enums import ( + WizardRunArtifactType, + WizardRunEnvironment, + WizardRunErrorCode, + WizardRunStage, + WizardRunStatus, +) from products.wizard.backend.logic.registry.config import POSTHOG_INTEGRATION_PROGRAM -from products.wizard.backend.observability import metrics -from products.wizard.backend.observability.contracts import WizardRunDispatchOutcome +from products.wizard.backend.logic.runs import dispatch, lifecycle +from products.wizard.backend.logic.runs.errors import WizardRunDispatchError +from products.wizard.backend.logic.workers.contracts import WizardWorkerResourceUsage, WizardWorkerTelemetry +from products.wizard.backend.observability import events, metrics, service +from products.wizard.backend.observability.contracts import WizardRunDispatchOutcome, WizardWorkerCleanupOutcome from products.wizard.backend.observability.service import WizardObservability +from products.wizard.backend.observability.worker_usage import worker_usage_observation +from products.wizard.backend.temporal.contracts import WizardRunActivityInput +from products.wizard.backend.temporal.errors import WizardTemporalError class TestWizardObservability(TestCase): @@ -46,21 +69,405 @@ def test_dispatch_observations_survive_independent_failures(self, metric_fails: with ( patch.object(metrics.WIZARD_RUN_DISPATCH_ATTEMPTS_TOTAL, "labels") as counter, - patch("products.wizard.backend.observability.events.celery_app.signature") as signature, + patch.object(events, "ph_background_capture") as background_capture, + patch.object(events.User.objects, "filter") as users, self.assertLogs("products.wizard.backend.observability.service", level=logging.INFO) as logs, ): counter.return_value.inc.side_effect = ValueError("Metric unavailable") if metric_fails else None - signature.return_value.apply_async.side_effect = ( - OperationalError("Queue unavailable") if event_fails else None - ) + users.return_value.values_list.return_value.first.return_value = "example-user" + background_capture.return_value.side_effect = ValueError("Capture unavailable") if event_fails else None WizardObservability().dispatch_finished(run, WizardRunDispatchOutcome.SUCCEEDED) counter.return_value.inc.assert_called_once_with() - signature.return_value.apply_async.assert_called_once_with() - self.assertEqual(signature.call_args.kwargs["args"][3], "wizard run dispatch finished") + background_capture.return_value.assert_called_once() + self.assertEqual(background_capture.return_value.call_args.kwargs["event"], "wizard run dispatch finished") messages = [record.getMessage() for record in logs.records] self.assertEqual("wizard_run_dispatch_metric_failed" in messages, metric_fails) self.assertEqual("wizard_run_dispatch_event_failed" in messages, event_fails) self.assertEqual(messages[-1], "wizard_run_dispatch_finished") self.assertEqual(logs.records[-1].levelno, logging.INFO) + + +def _cloud_run() -> WizardRunDTO: + created_at = datetime(2026, 8, 25, 12, tzinfo=UTC) + + return WizardRunDTO( + id=uuid4(), + team_id=1, + created_by_id=2, + environment=WizardRunEnvironment.CLOUD, + workspace=GitRepositoryWorkspace(repository="posthog/posthog"), + program=POSTHOG_INTEGRATION_PROGRAM, + status=WizardRunStatus.CREATED, + error_code=None, + error_message=None, + stage=WizardRunStage.DISPATCHING, + created_at=created_at, + updated_at=created_at, + started_at=None, + finished_at=None, + deadline_at=created_at + timedelta(hours=1), + ) + + +def _sample(name: str, labels: dict[str, str] | None = None) -> float: + value = REGISTRY.get_sample_value(name, labels) + return value if value is not None else 0 + + +@pytest.mark.django_db +def test_run_creation_emits_observability(team, user) -> None: + with patch("products.wizard.backend.logic.runs.lifecycle.run_observability") as observability: + run = lifecycle.create_run( + CreateWizardRunInput( + team_id=team.id, + created_by_id=user.id, + environment=WizardRunEnvironment.LOCAL, + workspace=LocalFolderWorkspace(project_name="example"), + program_id="posthog-integration", + ) + ) + + observability.run_created.assert_called_once_with(run) + + +@pytest.mark.django_db +def test_run_transition_emits_observability(team, user) -> None: + run = lifecycle.create_run( + CreateWizardRunInput( + team_id=team.id, + created_by_id=user.id, + environment=WizardRunEnvironment.LOCAL, + workspace=LocalFolderWorkspace(project_name="example"), + program_id="posthog-integration", + ) + ) + + with patch("products.wizard.backend.logic.runs.lifecycle.run_observability") as observability: + completed = lifecycle.complete_run(team.id, run.id) + + observability.run_transitioned.assert_called_once_with(run, completed) + + +def test_stage_update_is_idempotent() -> None: + run = replace(_cloud_run(), stage=WizardRunStage.PROVISIONING) + + with ( + patch.object(lifecycle.database_transaction, "atomic"), + patch.object(lifecycle.store, "get_run_for_update", return_value=run), + patch.object(lifecycle.store, "set_run_stage") as set_run_stage, + patch.object(lifecycle.run_observability, "stage_entered") as stage_entered, + ): + result = lifecycle.update_run_stage(run.team_id, run.id, WizardRunStage.PROVISIONING) + + assert result == run + set_run_stage.assert_not_called() + stage_entered.assert_not_called() + + +def test_failure_is_attributed_to_previous_stage() -> None: + previous = replace( + _cloud_run(), + status=WizardRunStatus.RUNNING, + stage=WizardRunStage.EXECUTING_WIZARD, + started_at=datetime(2026, 8, 25, 12, 1, tzinfo=UTC), + ) + current = replace( + previous, + status=WizardRunStatus.FAILED, + stage=None, + error_code=WizardRunErrorCode.EXECUTION_FAILED, + finished_at=datetime(2026, 8, 25, 12, 2, tzinfo=UTC), + ) + observability = WizardObservability() + + with ( + patch.object(service.metrics, "report_run_finished") as report_run_finished, + patch.object(service.events, "enqueue_run_finished") as enqueue_run_finished, + ): + observability.run_transitioned(previous, current) + + report_run_finished.assert_called_once_with(current, WizardRunStage.EXECUTING_WIZARD) + assert enqueue_run_finished.call_args.args[1] == WizardRunStage.EXECUTING_WIZARD + + +def test_observability_failures_do_not_escape() -> None: + run = replace(_cloud_run(), stage=None) + observability = WizardObservability() + + with ( + patch.object(service.metrics, "report_run_created", side_effect=ValueError("metric failed")), + patch.object(service.events, "enqueue_run_created", side_effect=DatabaseError("event failed")), + ): + observability.run_created(run) + + +def test_observability_does_not_hide_programming_errors() -> None: + run = replace(_cloud_run(), stage=None) + observability = WizardObservability() + + with ( + patch.object(service.metrics, "report_run_created", side_effect=TypeError("bug")), + patch.object(service.events, "enqueue_run_created"), + pytest.raises(TypeError, match="bug"), + ): + observability.run_created(run) + + +def test_cloud_creation_records_created_and_initial_stage_metrics() -> None: + run = _cloud_run() + created_labels = {"environment": "cloud"} + stage_labels = {"stage": "dispatching"} + created_before = _sample("posthog_wizard_runs_created_total", created_labels) + stage_before = _sample("posthog_wizard_run_stage_entered_total", stage_labels) + + with patch.object(events, "ph_background_capture"), patch.object(events.User.objects, "filter"): + WizardObservability().run_created(run) + + assert _sample("posthog_wizard_runs_created_total", created_labels) == created_before + 1 + assert _sample("posthog_wizard_run_stage_entered_total", stage_labels) == stage_before + 1 + + +def test_dispatch_and_stage_metrics_record_funnel() -> None: + dispatch_labels = {"outcome": WizardRunDispatchOutcome.SUCCEEDED.value} + stage_labels = {"stage": WizardRunStage.PROVISIONING.value} + dispatch_before = _sample("posthog_wizard_run_dispatch_attempts_total", dispatch_labels) + stage_before = _sample("posthog_wizard_run_stage_entered_total", stage_labels) + + metrics.report_dispatch_finished(WizardRunDispatchOutcome.SUCCEEDED) + metrics.report_stage_entered(WizardRunStage.PROVISIONING) + + assert _sample("posthog_wizard_run_dispatch_attempts_total", dispatch_labels) == dispatch_before + 1 + assert _sample("posthog_wizard_run_stage_entered_total", stage_labels) == stage_before + 1 + + +def test_active_cloud_run_metric_tracks_status_and_stage() -> None: + created = _cloud_run() + running = replace( + created, + status=WizardRunStatus.RUNNING, + stage=WizardRunStage.PROVISIONING, + ) + completed = replace(running, status=WizardRunStatus.COMPLETED, stage=None) + created_labels = {"status": "created", "stage": "dispatching"} + provisioning_labels = {"status": "running", "stage": "provisioning"} + created_before = _sample("posthog_wizard_cloud_runs_active", created_labels) + provisioning_before = _sample("posthog_wizard_cloud_runs_active", provisioning_labels) + + metrics.report_run_created(created) + metrics.report_run_status_changed(created, running) + metrics.report_run_status_changed(running, completed) + + assert _sample("posthog_wizard_cloud_runs_active", created_labels) == created_before + assert _sample("posthog_wizard_cloud_runs_active", provisioning_labels) == provisioning_before + + +def test_failure_metrics_record_stage_and_error_code() -> None: + run = replace( + _cloud_run(), + status=WizardRunStatus.FAILED, + stage=None, + error_code=WizardRunErrorCode.EXECUTION_FAILED, + started_at=datetime(2026, 8, 25, 12, 1, tzinfo=UTC), + finished_at=datetime(2026, 8, 25, 12, 2, tzinfo=UTC), + ) + finished_labels = { + "environment": "cloud", + "status": "failed", + "error_code": "execution_failed", + } + failure_labels = { + "environment": "cloud", + "stage": "executing_wizard", + "error_code": "execution_failed", + } + duration_labels = {"environment": "cloud", "status": "failed"} + finished_before = _sample("posthog_wizard_runs_finished_total", finished_labels) + failure_before = _sample("posthog_wizard_run_failures_total", failure_labels) + duration_before = _sample("posthog_wizard_run_duration_seconds_sum", duration_labels) + + metrics.report_run_finished(run, WizardRunStage.EXECUTING_WIZARD) + + assert _sample("posthog_wizard_runs_finished_total", finished_labels) == finished_before + 1 + assert _sample("posthog_wizard_run_failures_total", failure_labels) == failure_before + 1 + assert _sample("posthog_wizard_run_duration_seconds_sum", duration_labels) == duration_before + 60 + + +def test_failure_metrics_bound_unknown_error_code() -> None: + run = replace( + _cloud_run(), + status=WizardRunStatus.FAILED, + stage=None, + error_code="PHW_SOME_PINNED_VERSION_CODE", + started_at=datetime(2026, 8, 25, 12, 1, tzinfo=UTC), + finished_at=datetime(2026, 8, 25, 12, 2, tzinfo=UTC), + ) + bounded_labels = {"environment": "cloud", "status": "failed", "error_code": "other"} + raw_labels = {"environment": "cloud", "status": "failed", "error_code": "PHW_SOME_PINNED_VERSION_CODE"} + bounded_before = _sample("posthog_wizard_runs_finished_total", bounded_labels) + + metrics.report_run_finished(run, WizardRunStage.EXECUTING_WIZARD) + + assert _sample("posthog_wizard_runs_finished_total", bounded_labels) == bounded_before + 1 + assert REGISTRY.get_sample_value("posthog_wizard_runs_finished_total", raw_labels) is None + + +def test_dispatch_reports_success() -> None: + run = _cloud_run() + + with ( + patch.object(dispatch.store, "get_run", return_value=run), + patch.object(dispatch.temporal_client, "start_wizard_run_workflow"), + patch.object(dispatch.store, "mark_dispatch_succeeded"), + patch.object(dispatch.store, "cancellation_requested", return_value=False), + patch.object(dispatch.wizard_observability, "dispatch_finished") as dispatch_finished, + ): + dispatch.dispatch_created_cloud_wizard_run_to_temporal_worker(run.team_id, run.id) + + dispatch_finished.assert_called_once_with(run, WizardRunDispatchOutcome.SUCCEEDED) + + +@override_settings(DEBUG=True, LOCAL_WIZARD_ROOT="/tmp/posthog-wizard") +def test_dispatch_enables_local_wizard_source() -> None: + run = _cloud_run() + + with ( + patch.object(dispatch.store, "get_run", return_value=run), + patch.object(dispatch.temporal_client, "start_wizard_run_workflow") as start_workflow, + patch.object(dispatch.store, "mark_dispatch_succeeded"), + patch.object(dispatch.store, "cancellation_requested", return_value=False), + patch.object(dispatch.wizard_observability, "dispatch_finished"), + ): + dispatch.dispatch_created_cloud_wizard_run_to_temporal_worker(run.team_id, run.id) + + start_workflow.assert_called_once_with( + WizardRunActivityInput(team_id=run.team_id, run_id=run.id, use_local_wizard_source=True) + ) + + +def test_dispatch_reports_failure() -> None: + run = _cloud_run() + + with ( + patch.object(dispatch.store, "get_run", return_value=run), + patch.object( + dispatch.temporal_client, + "start_wizard_run_workflow", + side_effect=WizardTemporalError, + ), + patch.object(dispatch.store, "mark_dispatch_failed") as mark_dispatch_failed, + patch.object(dispatch.wizard_observability, "dispatch_finished") as dispatch_finished, + pytest.raises(WizardRunDispatchError), + ): + dispatch.dispatch_created_cloud_wizard_run_to_temporal_worker(run.team_id, run.id) + + mark_dispatch_failed.assert_called_once_with(run.team_id, run.id) + dispatch_finished.assert_called_once_with(run, WizardRunDispatchOutcome.FAILED) + + +def test_worker_metrics_record_usage_and_allocated_resources() -> None: + telemetry = WizardWorkerTelemetry( + resource_usage=WizardWorkerResourceUsage( + cpu_cores=2, + memory_gb=4, + disk_size_gb=16, + ttl_seconds=4500, + ttl_expires_at=datetime(2026, 8, 25, 13, 15, tzinfo=UTC), + provider_cpu_usage_usec=120_000_000, + provider_billed_cpu_usage_usec=180_000_000, + provider_usage_measured_at=datetime(2026, 8, 25, 12, 30, tzinfo=UTC), + ), + lifetime_seconds=300, + ) + lifetime_before = _sample("posthog_wizard_worker_lifetime_seconds_sum") + cpu_before = _sample("posthog_wizard_worker_cpu_usage_seconds_sum") + billed_cpu_before = _sample("posthog_wizard_worker_billed_cpu_usage_seconds_sum") + allocated_cpu_before = _sample("posthog_wizard_worker_allocated_cpu_core_seconds_sum") + allocated_memory_before = _sample("posthog_wizard_worker_allocated_memory_gb_seconds_sum") + allocated_disk_before = _sample("posthog_wizard_worker_allocated_disk_gb_seconds_sum") + + metrics.report_worker_usage(worker_usage_observation(telemetry)) + + assert _sample("posthog_wizard_worker_lifetime_seconds_sum") == lifetime_before + 300 + assert _sample("posthog_wizard_worker_cpu_usage_seconds_sum") == cpu_before + 120 + assert _sample("posthog_wizard_worker_billed_cpu_usage_seconds_sum") == billed_cpu_before + 180 + assert _sample("posthog_wizard_worker_allocated_cpu_core_seconds_sum") == allocated_cpu_before + 600 + assert _sample("posthog_wizard_worker_allocated_memory_gb_seconds_sum") == allocated_memory_before + 1200 + assert _sample("posthog_wizard_worker_allocated_disk_gb_seconds_sum") == allocated_disk_before + 4800 + + +def test_worker_metrics_skip_unavailable_provider_usage() -> None: + telemetry = WizardWorkerTelemetry( + resource_usage=WizardWorkerResourceUsage( + cpu_cores=2, + memory_gb=4, + disk_size_gb=16, + ttl_seconds=4500, + ttl_expires_at=datetime(2026, 8, 25, 13, 15, tzinfo=UTC), + ), + lifetime_seconds=300, + ) + cpu_before = _sample("posthog_wizard_worker_cpu_usage_seconds_count") + billed_cpu_before = _sample("posthog_wizard_worker_billed_cpu_usage_seconds_count") + + metrics.report_worker_usage(worker_usage_observation(telemetry)) + + assert _sample("posthog_wizard_worker_cpu_usage_seconds_count") == cpu_before + assert _sample("posthog_wizard_worker_billed_cpu_usage_seconds_count") == billed_cpu_before + + +def test_reliability_metrics_record_artifacts_cleanup_and_deadlines() -> None: + artifact = WizardRunGitDiffArtifactDTO( + id=uuid4(), + team_id=1, + run_id=uuid4(), + artifact_type=WizardRunArtifactType.GIT_DIFF, + size_bytes=10, + content_hash="hash", + additions=0, + removals=0, + created_at=datetime(2026, 8, 25, 12, tzinfo=UTC), + ) + artifact_labels = {"type": "git_diff"} + cleanup_labels = {"outcome": "succeeded"} + deadline_labels = {"environment": "cloud"} + artifact_before = _sample("posthog_wizard_artifacts_created_total", artifact_labels) + cleanup_before = _sample("posthog_wizard_worker_cleanups_total", cleanup_labels) + deadline_before = _sample("posthog_wizard_runs_past_deadline_total", deadline_labels) + + metrics.report_artifact_created(artifact) + metrics.report_worker_cleanup(WizardWorkerCleanupOutcome.SUCCEEDED) + metrics.report_run_past_deadline(_cloud_run()) + + assert _sample("posthog_wizard_artifacts_created_total", artifact_labels) == artifact_before + 1 + assert _sample("posthog_wizard_worker_cleanups_total", cleanup_labels) == cleanup_before + 1 + assert _sample("posthog_wizard_runs_past_deadline_total", deadline_labels) == deadline_before + 1 + + +@pytest.mark.parametrize("user_distinct_id", ["example-user", None]) +def test_stage_event_has_deterministic_identity_and_run_properties(user_distinct_id: str | None) -> None: + run = replace(_cloud_run(), stage=WizardRunStage.PROVISIONING) + + with ( + patch.object(events, "ph_background_capture") as background_capture, + patch.object(events.User.objects, "filter") as users, + ): + users.return_value.values_list.return_value.first.return_value = user_distinct_id + events.enqueue_stage_entered(run, WizardRunStage.PROVISIONING) + events.enqueue_stage_entered(run, WizardRunStage.PROVISIONING) + + first_args = background_capture.return_value.call_args_list[0].kwargs + second_args = background_capture.return_value.call_args_list[1].kwargs + + assert first_args["uuid"] == second_args["uuid"] + assert first_args["distinct_id"] == (user_distinct_id or str(run.id)) + assert first_args["properties"] == { + "team_id": run.team_id, + "wizard_run_id": str(run.id), + "environment": "cloud", + "workspace_type": "git_repository", + "program_id": "posthog-integration", + "wizard_version": POSTHOG_INTEGRATION_PROGRAM.wizard_version, + "stage": "provisioning", + } diff --git a/products/wizard/backend/tests/runs/test_repository_publisher.py b/products/wizard/backend/tests/runs/test_repository_publisher.py new file mode 100644 index 000000000000..8ceca06fced4 --- /dev/null +++ b/products/wizard/backend/tests/runs/test_repository_publisher.py @@ -0,0 +1,54 @@ +import base64 +from types import SimpleNamespace + +from unittest.mock import MagicMock, patch + +from products.wizard.backend.logic.workers.repository_publisher import create_signed_commit + + +def _execution_result(*, stdout: str = "", exit_code: int = 0) -> SimpleNamespace: + return SimpleNamespace(stdout=stdout, exit_code=exit_code) + + +def _graphql_response(payload: dict[str, object]) -> SimpleNamespace: + return SimpleNamespace(status_code=200, json=lambda: payload) + + +@patch("products.wizard.backend.logic.workers.repository_publisher.GitHubIntegration") +@patch("products.wizard.backend.logic.workers.repository_publisher.Integration.objects.filter") +def test_create_signed_commit_strips_head_sha_and_supports_binary_files( + filter_integrations: MagicMock, + github_integration_class: MagicMock, +) -> None: + filter_integrations.return_value.first.return_value = MagicMock() + github = github_integration_class.return_value + github.api_request.side_effect = [ + _graphql_response({"data": {"repository": {"id": "R_1", "ref": None}}}), + _graphql_response({"data": {"createRef": {"ref": {"name": "refs/heads/posthog/wizard-123"}}}}), + _graphql_response({"data": {"createCommitOnBranch": {"commit": {"oid": "newsha456"}}}}), + ] + + staged_contents = base64.b64encode(b"\x00\x01\x02").decode() + sandbox = MagicMock() + sandbox.execute.side_effect = [ + _execution_result(), + _execution_result(stdout="7a6e71985f4e0058f10517fc662813a39818f805\n"), + _execution_result(stdout="A\0src/config.py\0"), + _execution_result(stdout=staged_contents), + ] + + result = create_signed_commit( + sandbox, + team_id=7, + integration_id=13, + repository="posthog/posthog", + branch="posthog/wizard-123", + message="Set up PostHog", + source="wizard", + ) + + create_ref_input = github.api_request.call_args_list[1].kwargs["json_body"]["variables"]["input"] + commit_input = github.api_request.call_args_list[2].kwargs["json_body"]["variables"]["input"] + assert create_ref_input["oid"] == "7a6e71985f4e0058f10517fc662813a39818f805" + assert commit_input["expectedHeadOid"] == "7a6e71985f4e0058f10517fc662813a39818f805" + assert result.commit_shas == ("newsha456",) diff --git a/products/wizard/backend/tests/runs/test_temporal_client.py b/products/wizard/backend/tests/runs/test_temporal_client.py new file mode 100644 index 000000000000..c7abf4e33bb0 --- /dev/null +++ b/products/wizard/backend/tests/runs/test_temporal_client.py @@ -0,0 +1,86 @@ +from types import SimpleNamespace +from uuid import uuid4 + +import pytest +from unittest.mock import AsyncMock, MagicMock + +from django.conf import settings + +from temporalio.common import WorkflowIDReusePolicy +from temporalio.exceptions import WorkflowAlreadyStartedError +from temporalio.service import RPCError, RPCStatusCode + +from products.wizard.backend.temporal import client as temporal_client +from products.wizard.backend.temporal.contracts import WizardRunActivityInput +from products.wizard.backend.temporal.errors import WizardTemporalError +from products.wizard.backend.temporal.workflows.execute_run import ExecuteWizardRunWorkflow + + +def test_start_wizard_run_workflow_uses_stable_identity(monkeypatch: pytest.MonkeyPatch) -> None: + start_workflow = AsyncMock() + connect = AsyncMock(return_value=SimpleNamespace(start_workflow=start_workflow)) + monkeypatch.setattr(temporal_client, "async_connect", connect) + input = WizardRunActivityInput(team_id=1, run_id=uuid4()) + + temporal_client.start_wizard_run_workflow(input) + + start_workflow.assert_awaited_once_with( + ExecuteWizardRunWorkflow.get_name(), + input, + id=ExecuteWizardRunWorkflow.workflow_id_for(input.run_id), + id_reuse_policy=WorkflowIDReusePolicy.REJECT_DUPLICATE, + task_queue=settings.WIZARD_TASK_QUEUE, + retry_policy=temporal_client.WORKFLOW_RETRY_POLICY, + ) + + +def test_start_wizard_run_workflow_accepts_duplicate_dispatch(monkeypatch: pytest.MonkeyPatch) -> None: + input = WizardRunActivityInput(team_id=1, run_id=uuid4()) + start_workflow = AsyncMock( + side_effect=WorkflowAlreadyStartedError( + ExecuteWizardRunWorkflow.workflow_id_for(input.run_id), + ExecuteWizardRunWorkflow.get_name(), + ) + ) + connect = AsyncMock(return_value=SimpleNamespace(start_workflow=start_workflow)) + monkeypatch.setattr(temporal_client, "async_connect", connect) + + temporal_client.start_wizard_run_workflow(input) + + +def test_start_wizard_run_workflow_translates_rpc_failure(monkeypatch: pytest.MonkeyPatch) -> None: + input = WizardRunActivityInput(team_id=1, run_id=uuid4()) + start_workflow = AsyncMock(side_effect=RPCError("unavailable", RPCStatusCode.UNAVAILABLE, b"")) + connect = AsyncMock(return_value=SimpleNamespace(start_workflow=start_workflow)) + monkeypatch.setattr(temporal_client, "async_connect", connect) + + with pytest.raises(WizardTemporalError): + temporal_client.start_wizard_run_workflow(input) + + +def test_start_wizard_run_workflow_translates_connection_failure(monkeypatch: pytest.MonkeyPatch) -> None: + input = WizardRunActivityInput(team_id=1, run_id=uuid4()) + connect = AsyncMock(side_effect=RuntimeError("connection closed")) + monkeypatch.setattr(temporal_client, "async_connect", connect) + + with pytest.raises(WizardTemporalError): + temporal_client.start_wizard_run_workflow(input) + + +def test_cancel_wizard_run_workflow_accepts_missing_workflow(monkeypatch: pytest.MonkeyPatch) -> None: + cancel = AsyncMock(side_effect=RPCError("not found", RPCStatusCode.NOT_FOUND, b"")) + handle = MagicMock(cancel=cancel) + connect = AsyncMock(return_value=MagicMock(get_workflow_handle=MagicMock(return_value=handle))) + monkeypatch.setattr(temporal_client, "async_connect", connect) + + temporal_client.cancel_wizard_run_workflow(uuid4()) + + +def test_cancel_wizard_run_workflow_translates_rpc_failure(monkeypatch: pytest.MonkeyPatch) -> None: + cancel = AsyncMock(side_effect=RPCError("unavailable", RPCStatusCode.UNAVAILABLE, b"")) + handle = MagicMock(cancel=cancel) + connect = AsyncMock(return_value=MagicMock(get_workflow_handle=MagicMock(return_value=handle))) + monkeypatch.setattr(temporal_client, "async_connect", connect) + + with pytest.raises(WizardTemporalError): + temporal_client.cancel_wizard_run_workflow(uuid4()) diff --git a/products/wizard/backend/tests/runs/test_worker_lifecycle.py b/products/wizard/backend/tests/runs/test_worker_lifecycle.py new file mode 100644 index 000000000000..49c449299d53 --- /dev/null +++ b/products/wizard/backend/tests/runs/test_worker_lifecycle.py @@ -0,0 +1,83 @@ +from uuid import uuid4 + +import pytest +from unittest.mock import MagicMock, patch + +from django.db import DatabaseError + +from products.tasks.backend.facade.sandbox import SandboxCleanupError +from products.wizard.backend.logic.runs.errors import WizardWorkerCleanupError +from products.wizard.backend.logic.workers import lifecycle as worker_lifecycle +from products.wizard.backend.observability.contracts import WizardWorkerCleanupOutcome + + +def test_cleanup_worker_destroys_sandbox_when_usage_measurement_has_programming_error() -> None: + run_id = uuid4() + + with ( + patch.object(worker_lifecycle.worker_store, "mark_cleanup_pending"), + patch.object(worker_lifecycle.cloud_worker, "measure_worker_usage", side_effect=RuntimeError("bug")), + patch.object(worker_lifecycle.cloud_worker, "destroy_worker") as destroy_worker, + patch.object(worker_lifecycle.worker_store, "mark_cleaned"), + ): + with pytest.raises(RuntimeError, match="bug"): + worker_lifecycle.cleanup_worker(1, run_id, "sandbox-id") + + destroy_worker.assert_called_once_with("sandbox-id") + + +def test_cleanup_worker_destroys_sandbox_when_usage_persistence_fails() -> None: + run_id = uuid4() + usage = MagicMock() + + with ( + patch.object(worker_lifecycle.worker_store, "mark_cleanup_pending"), + patch.object(worker_lifecycle.cloud_worker, "measure_worker_usage", return_value=usage), + patch.object(worker_lifecycle.worker_store, "record_usage", side_effect=DatabaseError), + patch.object(worker_lifecycle.cloud_worker, "destroy_worker") as destroy_worker, + patch.object(worker_lifecycle.worker_store, "mark_cleaned") as mark_cleaned, + patch.object(worker_lifecycle, "_report_worker_usage"), + patch.object(worker_lifecycle.wizard_observability, "worker_cleanup_finished") as cleanup_finished, + ): + worker_lifecycle.cleanup_worker(1, run_id, "sandbox-id") + + destroy_worker.assert_called_once_with("sandbox-id") + mark_cleaned.assert_called_once_with(1, run_id) + cleanup_finished.assert_called_once_with(1, run_id, WizardWorkerCleanupOutcome.SUCCEEDED) + + +def test_cleanup_worker_translates_sandbox_cleanup_failure() -> None: + run_id = uuid4() + cleanup_error = SandboxCleanupError("cleanup failed", {}, RuntimeError("provider unavailable"), capture=False) + + with ( + patch.object(worker_lifecycle.worker_store, "mark_cleanup_pending"), + patch.object(worker_lifecycle.cloud_worker, "measure_worker_usage", return_value=None), + patch.object(worker_lifecycle.cloud_worker, "destroy_worker", side_effect=cleanup_error), + patch.object(worker_lifecycle.worker_store, "mark_cleanup_failed") as mark_cleanup_failed, + patch.object(worker_lifecycle.wizard_observability, "worker_cleanup_finished") as cleanup_finished, + ): + with pytest.raises(WizardWorkerCleanupError): + worker_lifecycle.cleanup_worker(1, run_id, "sandbox-id") + + mark_cleanup_failed.assert_called_once_with(1, run_id) + cleanup_finished.assert_called_once_with(1, run_id, WizardWorkerCleanupOutcome.FAILED) + + +def test_cleanup_worker_reports_persisted_usage_after_cleanup() -> None: + run_id = uuid4() + run = MagicMock() + telemetry = MagicMock() + + with ( + patch.object(worker_lifecycle.worker_store, "mark_cleanup_pending"), + patch.object(worker_lifecycle.cloud_worker, "measure_worker_usage", return_value=None), + patch.object(worker_lifecycle.cloud_worker, "destroy_worker"), + patch.object(worker_lifecycle.worker_store, "mark_cleaned"), + patch.object(worker_lifecycle.run_store, "get_run", return_value=run), + patch.object(worker_lifecycle.worker_store, "get_worker_telemetry", return_value=telemetry), + patch.object(worker_lifecycle.wizard_observability, "worker_usage_recorded") as worker_usage_recorded, + ): + worker_lifecycle.cleanup_worker(1, run_id, "sandbox-id") + + worker_usage_recorded.assert_called_once_with(run, telemetry) diff --git a/products/wizard/backend/tests/runs/test_workspace.py b/products/wizard/backend/tests/runs/test_workspace.py new file mode 100644 index 000000000000..139be83c0dad --- /dev/null +++ b/products/wizard/backend/tests/runs/test_workspace.py @@ -0,0 +1,45 @@ +import pytest + +from products.wizard.backend.facade.contracts import GitRepositoryWorkspace, LocalFolderWorkspace, WizardWorkspace +from products.wizard.backend.facade.enums import WizardWorkspaceType +from products.wizard.backend.logic.runs.mappers import record_to_workspace, workspace_to_record + + +@pytest.mark.parametrize( + ("workspace_type", "metadata", "expected_workspace"), + [ + ( + WizardWorkspaceType.LOCAL_FOLDER, + {"project_name": "example-project"}, + LocalFolderWorkspace(project_name="example-project"), + ), + ( + WizardWorkspaceType.GIT_REPOSITORY, + {"repository": "posthog/posthog"}, + GitRepositoryWorkspace(repository="posthog/posthog"), + ), + ], +) +def test_workspace_serialization_round_trip( + workspace_type: WizardWorkspaceType, + metadata: dict[str, str], + expected_workspace: WizardWorkspace, +) -> None: + workspace = record_to_workspace(workspace_type.value, metadata) + + assert workspace == expected_workspace + assert workspace_to_record(workspace) == (workspace_type, metadata) + + +@pytest.mark.parametrize( + ("workspace_type", "metadata"), + [ + (WizardWorkspaceType.LOCAL_FOLDER, {"project_name": 123}), + (WizardWorkspaceType.GIT_REPOSITORY, {"repository": None}), + ], +) +def test_workspace_rejects_invalid_serialized_value( + workspace_type: WizardWorkspaceType, metadata: dict[str, object] +) -> None: + with pytest.raises(ValueError): + record_to_workspace(workspace_type.value, metadata) diff --git a/products/wizard/turbo.json b/products/wizard/turbo.json index a765b03ee90e..71041022421c 100644 --- a/products/wizard/turbo.json +++ b/products/wizard/turbo.json @@ -7,6 +7,7 @@ "backend/presentation/**", "backend/routes.py", "backend/tasks/**", + "backend/temporal/**", "backend/models.py", "backend/migrations/**" ], diff --git a/tach.toml b/tach.toml index bf83205f27e9..80f080cb6766 100644 --- a/tach.toml +++ b/tach.toml @@ -138,6 +138,7 @@ depends_on = [ "products.mcp_analytics", "products.mcp_registry", "products.web_analytics", + "products.wizard", "products.workflows", "products.actions", "products.alerts", "products.warehouse_sources", "products.data_tools", "products.exports", "products.annotations",