diff --git a/.semgrep/rules/security/idor-team-scoped-models.yaml b/.semgrep/rules/security/idor-team-scoped-models.yaml index b2c325147fef..963ff7b14c6b 100644 --- a/.semgrep/rules/security/idor-team-scoped-models.yaml +++ b/.semgrep/rules/security/idor-team-scoped-models.yaml @@ -389,7 +389,10 @@ rules: |WebAnalyticsFilterPreset |WebAnalyticsInteraction |WebAnalyticsVisit + |WizardRun + |WizardRunArtifact |WizardSession + |WizardWorker )$ - focus-metavariable: $SINK pattern-sanitizers: @@ -770,7 +773,10 @@ rules: |WebAnalyticsFilterPreset |WebAnalyticsInteraction |WebAnalyticsVisit + |WizardRun + |WizardRunArtifact |WizardSession + |WizardWorker )$ # Exclude when team filter is present - pattern-not: $M.objects. ... .$METHOD2(..., team__project_id=$T, ...) diff --git a/docs/internal/wizard-registry.md b/docs/internal/wizard-registry.md new file mode 100644 index 000000000000..286399c46529 --- /dev/null +++ b/docs/internal/wizard-registry.md @@ -0,0 +1,75 @@ +# 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. + +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. + +## API + +The registry endpoint returns a paginated list: + +```text +GET /api/projects/{project_id}/wizard/registry/?limit=100&offset=0 +``` + +Each program contains an ID, display metadata, a pinned Wizard version, Wizard command arguments, tags, prerequisite program IDs, and supported run environments. + +Creating a run requires the selected `program_id`: + +```json +{ + "program_id": "posthog-integration", + "environment": "local", + "workspace": { + "type": "local_folder", + "project_name": "example-project" + } +} +``` + +The backend resolves the ID from the personalized registry, checks that the program supports the requested environment, and stores the complete program definition on the run. +Clients cannot supply or change program metadata. +Run responses return the stored `program` snapshot. + +## Remote configuration + +The registry is stored in the `wizard-program-registry` feature flag payload: + +```json +{ + "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": ["local", "cloud"] + } + ] +} +``` + +The entire payload is invalid when its version or shape is unsupported, program IDs are duplicated, an environment is unknown, or any program field is invalid. +Each program must include a `wizard_version` pinned to an exact release, such as `2.60.0`. +The registry payload rejects `latest`, so a program that omits `wizard_version` or sets it to `latest` makes the whole payload invalid. +A valid empty program list remains empty. + +When the payload cannot be fetched or is invalid, the registry contains one `posthog-integration` program with local and cloud support. +Its command is empty, which keeps the Wizard package's default command. + +## Cloud 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`. + +Program command tokens are positional command names. +They appear after `@posthog/wizard@{program.wizard_version}` and before worker-owned flags. +Option flags and shell syntax are rejected, and every accepted token is shell-quoted before execution. diff --git a/products/wizard/backend/facade/api.py b/products/wizard/backend/facade/api.py index 6a70795d241a..0b7d8d92d430 100644 --- a/products/wizard/backend/facade/api.py +++ b/products/wizard/backend/facade/api.py @@ -13,8 +13,12 @@ from typing import Any from products.wizard.backend import metrics -from products.wizard.backend.facade.contracts import UpsertWizardSessionInput, WizardSessionDTO -from products.wizard.backend.logic import pubsub, sessions +from products.wizard.backend.facade.contracts import UpsertWizardSessionInput, WizardProgram, WizardSessionDTO +from products.wizard.backend.logic import ( + pubsub, + registry as registry_service, + sessions, +) def upsert(params: UpsertWizardSessionInput) -> tuple[WizardSessionDTO, bool]: @@ -65,3 +69,7 @@ def record_latest_session_poll(raw_source: str | None, result: str) -> None: metrics.WIZARD_LATEST_SESSION_REQUESTS_TOTAL.labels( source=metrics.poll_source_label(raw_source), result=result ).inc() + + +def get_registry(*, distinct_id: str, organization_id: str) -> tuple[WizardProgram, ...]: + return registry_service.get_registry(distinct_id=distinct_id, organization_id=organization_id) diff --git a/products/wizard/backend/facade/config.py b/products/wizard/backend/facade/config.py new file mode 100644 index 000000000000..b244f766048e --- /dev/null +++ b/products/wizard/backend/facade/config.py @@ -0,0 +1,9 @@ +from datetime import timedelta + +DEFAULT_WIZARD_VERSION = "2.67.0" +LATEST_WIZARD_VERSION = "latest" +WIZARD_REGISTRY_VERSION = 1 +WIZARD_RUN_DEADLINE = timedelta(minutes=85) +WIZARD_ERROR_CODE_MAX_LENGTH = 50 +WIZARD_ERROR_CODE_PATTERN = r"PHW_[A-Z][A-Z0-9_]*" +WIZARD_GIT_DIFF_CONTENT_TYPE = "text/x-diff; charset=utf-8" diff --git a/products/wizard/backend/facade/contracts.py b/products/wizard/backend/facade/contracts.py index f523ccb18e4e..8da93e5ac6ef 100644 --- a/products/wizard/backend/facade/contracts.py +++ b/products/wizard/backend/facade/contracts.py @@ -7,9 +7,19 @@ from dataclasses import dataclass from datetime import datetime, timedelta -from typing import Any +from typing import Any, Literal +from uuid import UUID -from .enums import RunPhase, TaskStatus +from posthog.dataclasses import frozen + +from .enums import ( + WizardRunArtifactType, + WizardRunEnvironment, + WizardRunStage, + WizardRunStatus, + WizardSessionRunPhase, + WizardSessionTaskStatus, +) STALE_AFTER = timedelta(minutes=10) @@ -18,7 +28,7 @@ class WizardTaskDTO: id: str title: str - status: TaskStatus + status: WizardSessionTaskStatus @dataclass(frozen=True) @@ -35,7 +45,7 @@ class WizardSessionDTO: workflow_id: str skill_id: str started_at: datetime - run_phase: RunPhase + run_phase: WizardSessionRunPhase tasks: tuple[WizardTaskDTO, ...] event_plan: dict[str, Any] | None error: dict[str, Any] | None @@ -47,10 +57,6 @@ class WizardSessionDTO: is_stale: bool -class WizardSessionOwnershipError(Exception): - """Raised when an upsert would overwrite a session owned by a different user.""" - - @dataclass(frozen=True) class UpsertWizardSessionRequest: """What the wizard CLI POSTs. team_id is derived from the URL, not the body.""" @@ -59,7 +65,7 @@ class UpsertWizardSessionRequest: workflow_id: str skill_id: str started_at: datetime - run_phase: RunPhase + run_phase: WizardSessionRunPhase tasks: tuple[WizardTaskDTO, ...] event_plan: dict[str, Any] | None = None error: dict[str, Any] | None = None @@ -74,7 +80,7 @@ class UpsertWizardSessionInput: workflow_id: str skill_id: str started_at: datetime - run_phase: RunPhase + run_phase: WizardSessionRunPhase tasks: tuple[WizardTaskDTO, ...] event_plan: dict[str, Any] | None error: dict[str, Any] | None @@ -82,3 +88,130 @@ class UpsertWizardSessionInput: handoff_text: str | None = None # Set on create only, never overwritten on later pushes for the same run. created_by_id: int | None = None + + +class WizardSessionOwnershipError(Exception): + pass + + +@frozen +class LocalFolderWorkspace: + project_name: str + type: Literal["local_folder"] = "local_folder" + + +@frozen +class GitRepositoryWorkspace: + repository: str + type: Literal["git_repository"] = "git_repository" + + +type WizardWorkspace = LocalFolderWorkspace | GitRepositoryWorkspace + + +@frozen +class WizardProgram: + id: str + name: str + description: str + wizard_version: str + command: tuple[str, ...] + tags: tuple[str, ...] + required_programs: tuple[str, ...] + supported_environments: tuple[WizardRunEnvironment, ...] + + +@frozen +class WizardRegistry: + programs: tuple[WizardProgram, ...] + version: Literal[1] = 1 + + +@frozen +class CreateWizardRunInput: + team_id: int + created_by_id: int + environment: WizardRunEnvironment + workspace: WizardWorkspace + program_id: str + wizard_version: str | None = None + idempotency_key: str | None = None + + +@frozen +class WizardRunCreationResult: + run: "WizardRunDTO" + created: bool + + +@frozen +class WizardRunDTO: + id: UUID + team_id: int + created_by_id: int | None + environment: WizardRunEnvironment + workspace: WizardWorkspace + program: WizardProgram + status: WizardRunStatus + error_code: str | None + error_message: str | None + stage: WizardRunStage | None + created_at: datetime + updated_at: datetime | None + started_at: datetime | None + finished_at: datetime | None + deadline_at: datetime | None + + +@frozen +class ListWizardRunsInput: + team_id: int + offset: int + limit: int + + +@frozen +class WizardRunPage: + results: tuple[WizardRunDTO, ...] + count: int + + +@frozen +class CreatePullRequestArtifactInput: + team_id: int + run_id: UUID + url: str + number: int + repository: str + head_branch: str + base_branch: str + + +@frozen +class WizardRunGitDiffArtifactDTO: + id: UUID + team_id: int + run_id: UUID + artifact_type: Literal[WizardRunArtifactType.GIT_DIFF] + size_bytes: int + content_hash: str + additions: int | None + removals: int | None + created_at: datetime + + +@frozen +class WizardRunPullRequestArtifactDTO: + id: UUID + team_id: int + run_id: UUID + artifact_type: Literal[WizardRunArtifactType.PULL_REQUEST] + url: str + number: int + repository: str + head_branch: str + base_branch: str + created_at: datetime + + +type WizardRunArtifactDTO = WizardRunGitDiffArtifactDTO | WizardRunPullRequestArtifactDTO diff --git a/products/wizard/backend/facade/enums.py b/products/wizard/backend/facade/enums.py index 1b709c165cc2..ca917ef96b93 100644 --- a/products/wizard/backend/facade/enums.py +++ b/products/wizard/backend/facade/enums.py @@ -9,18 +9,14 @@ from enum import StrEnum -class RunPhase(StrEnum): - """Phase of a run.""" - +class WizardSessionRunPhase(StrEnum): IDLE = "idle" RUNNING = "running" COMPLETED = "completed" ERROR = "error" -class TaskStatus(StrEnum): - """Status of a task.""" - +class WizardSessionTaskStatus(StrEnum): PENDING = "pending" IN_PROGRESS = "in_progress" COMPLETED = "completed" @@ -28,3 +24,60 @@ class TaskStatus(StrEnum): # These are not currently used, but we want to reserve them for future use. FAILED = "failed" CANCELED = "canceled" + + +RunPhase = WizardSessionRunPhase +TaskStatus = WizardSessionTaskStatus + + +class WizardRunStatus(StrEnum): + CREATED = "created" + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + CANCELLED = "cancelled" + + +class WizardRunDispatchStatus(StrEnum): + PENDING = "pending" + DISPATCHED = "dispatched" + + +class WizardRunStage(StrEnum): + DISPATCHING = "dispatching" + PROVISIONING = "provisioning" + PREPARING_WORKSPACE = "preparing_workspace" + EXECUTING_WIZARD = "executing_wizard" + CREATING_ARTIFACTS = "creating_artifacts" + + +class WizardWorkerCleanupStatus(StrEnum): + ACTIVE = "active" + PENDING = "pending" + CLEANED = "cleaned" + FAILED = "failed" + + +class WizardRunEnvironment(StrEnum): + LOCAL = "local" + CLOUD = "cloud" + + +class WizardWorkspaceType(StrEnum): + LOCAL_FOLDER = "local_folder" + GIT_REPOSITORY = "git_repository" + + +class WizardRunErrorCode(StrEnum): + TIMEOUT = "timeout" + PROVISIONING_FAILED = "provisioning_failed" + REPOSITORY_ACCESS_FAILED = "repository_access_failed" + WORKSPACE_PREPARATION_FAILED = "workspace_preparation_failed" + EXECUTION_FAILED = "execution_failed" + ARTIFACT_CREATION_FAILED = "artifact_creation_failed" + DISPATCH_FAILED = "dispatch_failed" + + +class WizardRunArtifactType(StrEnum): + GIT_DIFF = "git_diff" + PULL_REQUEST = "pull_request" diff --git a/products/wizard/backend/facade/errors.py b/products/wizard/backend/facade/errors.py new file mode 100644 index 000000000000..78013b27d837 --- /dev/null +++ b/products/wizard/backend/facade/errors.py @@ -0,0 +1,65 @@ +from products.wizard.backend.facade.contracts import WizardSessionOwnershipError as WizardSessionOwnershipError + + +class MissingGitHubIntegrationError(Exception): + pass + + +class RepositoryNotAccessibleError(Exception): + pass + + +class InvalidWorkspaceEnvironmentError(Exception): + pass + + +class InvalidRepositoryError(Exception): + pass + + +class WizardProgramNotAvailableError(Exception): + pass + + +class WizardProgramEnvironmentNotSupportedError(Exception): + pass + + +class IllegalStatusTransitionError(Exception): + pass + + +class InvalidTransitionMetadataError(Exception): + pass + + +class WizardRunNotFoundError(Exception): + pass + + +class WizardRunArtifactNotFoundError(Exception): + pass + + +class WizardRunArtifactTooLargeError(Exception): + pass + + +class WizardRunIdempotencyConflictError(Exception): + pass + + +class MissingWizardRunIdempotencyKeyError(Exception): + pass + + +class ActiveWizardRunError(Exception): + pass + + +class WizardRunHourlyLimitError(Exception): + pass + + +class WizardRunDailyLimitError(Exception): + pass diff --git a/products/wizard/backend/facade/validation.py b/products/wizard/backend/facade/validation.py new file mode 100644 index 000000000000..56b351eb49ab --- /dev/null +++ b/products/wizard/backend/facade/validation.py @@ -0,0 +1,88 @@ +import re +from typing import TypeGuard + +from products.wizard.backend.facade.config import ( + LATEST_WIZARD_VERSION, + WIZARD_ERROR_CODE_MAX_LENGTH, + WIZARD_ERROR_CODE_PATTERN, +) +from products.wizard.backend.facade.enums import WizardRunEnvironment + +_EXACT_WIZARD_VERSION_PATTERN = re.compile( + r"^(0|[1-9][0-9]*)\.(0|[1-9][0-9]*)\.(0|[1-9][0-9]*)" + r"(?:-(?:0|[1-9][0-9]*|[0-9]*[A-Za-z-][0-9A-Za-z-]*)(?:\.(?:0|[1-9][0-9]*|[0-9]*[A-Za-z-][0-9A-Za-z-]*))*)?" + r"(?:\+[0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*)?$" +) +_PROGRAM_ID_PATTERN = re.compile(r"^[a-z0-9]+(?:-[a-z0-9]+)*$") +_WIZARD_ERROR_CODE_PATTERN = re.compile(WIZARD_ERROR_CODE_PATTERN) + + +def validate_nonempty_string(value: object, *, error: str) -> str: + if not isinstance(value, str) or value != value.strip() or not value: + raise ValueError(error) + return value + + +def validate_program_id(value: object) -> str: + if not isinstance(value, str) or _PROGRAM_ID_PATTERN.fullmatch(value) is None: + raise ValueError("Invalid Wizard program") + return value + + +def validate_pinned_wizard_version(value: object) -> str: + if not isinstance(value, str) or _EXACT_WIZARD_VERSION_PATTERN.fullmatch(value) is None: + raise ValueError("Invalid Wizard program") + + return value + + +def validate_wizard_version(value: object) -> str: + if value == LATEST_WIZARD_VERSION: + return LATEST_WIZARD_VERSION + return validate_pinned_wizard_version(value) + + +def is_executable_wizard_version(value: object) -> bool: + try: + validate_wizard_version(value) + except ValueError: + return False + return True + + +def is_wizard_error_code(value: object) -> TypeGuard[str]: + return ( + isinstance(value, str) + and len(value) <= WIZARD_ERROR_CODE_MAX_LENGTH + and _WIZARD_ERROR_CODE_PATTERN.fullmatch(value) is not None + ) + + +def validate_program_ids(value: object) -> tuple[str, ...]: + if not isinstance(value, list): + raise ValueError("Invalid Wizard program") + values = tuple(validate_program_id(item) for item in value) + if len(values) != len(set(values)): + raise ValueError("Invalid Wizard program") + return values + + +def validate_program_environments(value: object) -> tuple[WizardRunEnvironment, ...]: + if not isinstance(value, list): + raise ValueError("Invalid Wizard program") + try: + environments = tuple(WizardRunEnvironment(item) for item in value) + except (TypeError, ValueError) as error: + raise ValueError("Invalid Wizard program") from error + if not environments or len(environments) != len(set(environments)): + raise ValueError("Invalid Wizard program") + return environments + + +def validate_workspace_metadata_value(metadata: object, key: str) -> str: + if not isinstance(metadata, dict): + raise ValueError("Wizard workspace metadata must be an object") + value = metadata.get(key) + if not isinstance(value, str): + raise ValueError(f"Wizard workspace metadata field {key!r} must be a string") + return value diff --git a/products/wizard/backend/logic/programs.py b/products/wizard/backend/logic/programs.py new file mode 100644 index 000000000000..d1db7af3fb4a --- /dev/null +++ b/products/wizard/backend/logic/programs.py @@ -0,0 +1,57 @@ +from products.wizard.backend.facade.contracts import WizardProgram +from products.wizard.backend.facade.validation import ( + validate_nonempty_string, + validate_pinned_wizard_version, + validate_program_environments, + validate_program_id, + validate_program_ids, + validate_wizard_version, +) + +_PROGRAM_FIELDS = frozenset( + { + "id", + "name", + "description", + "wizard_version", + "command", + "tags", + "required_programs", + "supported_environments", + } +) + + +def program_to_mapping(program: WizardProgram) -> dict[str, object]: + return { + "id": program.id, + "name": program.name, + "description": program.description, + "wizard_version": program.wizard_version, + "command": list(program.command), + "tags": list(program.tags), + "required_programs": list(program.required_programs), + "supported_environments": [environment.value for environment in program.supported_environments], + } + + +def program_from_mapping(value: object, *, allow_latest_version: bool = False) -> WizardProgram: + if not isinstance(value, dict): + raise ValueError("Invalid Wizard program") + fields = set(value) + if not _PROGRAM_FIELDS.issubset(fields) or (not allow_latest_version and fields != _PROGRAM_FIELDS): + raise ValueError("Invalid Wizard program") + return WizardProgram( + id=validate_program_id(value["id"]), + name=validate_nonempty_string(value["name"], error="Invalid Wizard program"), + description=validate_nonempty_string(value["description"], error="Invalid Wizard program"), + wizard_version=( + validate_wizard_version(value["wizard_version"]) + if allow_latest_version + else validate_pinned_wizard_version(value["wizard_version"]) + ), + command=validate_program_ids(value["command"]), + tags=validate_program_ids(value["tags"]), + required_programs=validate_program_ids(value["required_programs"]), + supported_environments=validate_program_environments(value["supported_environments"]), + ) diff --git a/products/wizard/backend/logic/registry/__init__.py b/products/wizard/backend/logic/registry/__init__.py new file mode 100644 index 000000000000..256d9f3efb3f --- /dev/null +++ b/products/wizard/backend/logic/registry/__init__.py @@ -0,0 +1,3 @@ +from products.wizard.backend.logic.registry.service import get_program, get_registry + +__all__ = ["get_program", "get_registry"] diff --git a/products/wizard/backend/logic/registry/config.py b/products/wizard/backend/logic/registry/config.py new file mode 100644 index 000000000000..a93944150511 --- /dev/null +++ b/products/wizard/backend/logic/registry/config.py @@ -0,0 +1,18 @@ +from products.wizard.backend.facade.config import DEFAULT_WIZARD_VERSION +from products.wizard.backend.facade.contracts import WizardProgram, WizardRegistry +from products.wizard.backend.facade.enums import WizardRunEnvironment + +REGISTRY_FEATURE_FLAG = "wizard-program-registry" + +POSTHOG_INTEGRATION_PROGRAM = WizardProgram( + id="posthog-integration", + name="PostHog integration", + description="Set up PostHog SDK integration", + wizard_version=DEFAULT_WIZARD_VERSION, + command=(), + tags=(), + required_programs=(), + supported_environments=(WizardRunEnvironment.LOCAL, WizardRunEnvironment.CLOUD), +) + +FALLBACK_REGISTRY = WizardRegistry(programs=(POSTHOG_INTEGRATION_PROGRAM,)) diff --git a/products/wizard/backend/logic/registry/parser.py b/products/wizard/backend/logic/registry/parser.py new file mode 100644 index 000000000000..5475a486274b --- /dev/null +++ b/products/wizard/backend/logic/registry/parser.py @@ -0,0 +1,30 @@ +import json + +from products.wizard.backend.facade.config import WIZARD_REGISTRY_VERSION +from products.wizard.backend.facade.contracts import WizardRegistry +from products.wizard.backend.logic.programs import program_from_mapping + +_REGISTRY_FIELDS = frozenset({"version", "programs"}) + + +def parse_registry_payload(value: object) -> WizardRegistry: + decoded = _decode_payload(value) + if not isinstance(decoded, dict) or set(decoded) != _REGISTRY_FIELDS: + raise ValueError("Invalid Wizard registry") + if type(decoded["version"]) is not int or decoded["version"] != WIZARD_REGISTRY_VERSION: + raise ValueError("Invalid Wizard registry") + if not isinstance(decoded["programs"], list): + raise ValueError("Invalid Wizard registry") + programs = tuple(program_from_mapping(program) for program in decoded["programs"]) + if len(programs) != len({program.id for program in programs}): + raise ValueError("Invalid Wizard registry") + return WizardRegistry(programs=programs) + + +def _decode_payload(value: object) -> object: + if not isinstance(value, str): + return value + try: + return json.loads(value) + except (TypeError, ValueError) as error: + raise ValueError("Invalid Wizard registry") from error diff --git a/products/wizard/backend/logic/registry/service.py b/products/wizard/backend/logic/registry/service.py new file mode 100644 index 000000000000..a5b045017f63 --- /dev/null +++ b/products/wizard/backend/logic/registry/service.py @@ -0,0 +1,44 @@ +import logging + +import posthoganalytics +from posthoganalytics.request import APIError, RequestsConnectionError, RequestsTimeout + +from products.wizard.backend.facade.contracts import WizardProgram +from products.wizard.backend.facade.errors import WizardProgramNotAvailableError +from products.wizard.backend.logic.registry.config import FALLBACK_REGISTRY, REGISTRY_FEATURE_FLAG +from products.wizard.backend.logic.registry.parser import parse_registry_payload +from products.wizard.backend.metrics import report_registry_fallback + +logger = logging.getLogger(__name__) + + +def get_registry(*, distinct_id: str, organization_id: str) -> tuple[WizardProgram, ...]: + try: + payload = posthoganalytics.get_feature_flag_payload( + REGISTRY_FEATURE_FLAG, + distinct_id=distinct_id, + groups={"organization": organization_id}, + group_properties={"organization": {"id": organization_id}}, + only_evaluate_locally=False, + send_feature_flag_events=False, + ) + except (APIError, RequestsConnectionError, RequestsTimeout): + logger.warning("wizard_registry_fallback", extra={"reason": "request_failed"}, exc_info=True) + report_registry_fallback("request_failed") + return FALLBACK_REGISTRY.programs + + try: + registry = parse_registry_payload(payload) + except ValueError: + logger.warning("wizard_registry_fallback", extra={"reason": "invalid_payload"}, exc_info=True) + report_registry_fallback("invalid_payload") + return FALLBACK_REGISTRY.programs + return registry.programs + + +def get_program(*, program_id: str, distinct_id: str, organization_id: str) -> WizardProgram: + programs = get_registry(distinct_id=distinct_id, organization_id=organization_id) + program = next((program for program in programs if program.id == program_id), None) + if program is None: + raise WizardProgramNotAvailableError + return program diff --git a/products/wizard/backend/metrics.py b/products/wizard/backend/metrics.py index fd26d5416f7a..7aaf70151d95 100644 --- a/products/wizard/backend/metrics.py +++ b/products/wizard/backend/metrics.py @@ -44,6 +44,16 @@ def poll_source_label(raw_source: str | None) -> str: labelnames=["outcome"], ) +WIZARD_REGISTRY_FALLBACK_TOTAL = Counter( + "posthog_wizard_registry_fallback_total", + "Wizard registry requests that used the built-in fallback", + labelnames=["reason"], +) + + +def report_registry_fallback(reason: str) -> None: + WIZARD_REGISTRY_FALLBACK_TOTAL.labels(reason=reason).inc() + def report_session_upserted(previous_run_phase: str | None, dto: WizardSessionDTO) -> None: if dto.run_phase not in _TERMINAL_PHASES or previous_run_phase in _TERMINAL_PHASES: diff --git a/products/wizard/backend/migrations/0006_create_wizard_run.py b/products/wizard/backend/migrations/0006_create_wizard_run.py new file mode 100644 index 000000000000..92a0dcb8dc9c --- /dev/null +++ b/products/wizard/backend/migrations/0006_create_wizard_run.py @@ -0,0 +1,152 @@ +import django.db.models.deletion +from django.conf import settings +from django.db import migrations, models + +import posthog.uuidt + + +class Migration(migrations.Migration): + dependencies = [ + ("posthog", "1319_team_posthog_team_widget_token_idx"), + ("wizard", "0005_wizardsession_handoff_text"), + migrations.swappable_dependency(settings.AUTH_USER_MODEL), + ] + + operations = [ + migrations.CreateModel( + name="WizardRun", + fields=[ + ("created_at", models.DateTimeField(auto_now_add=True)), + ("updated_at", models.DateTimeField(auto_now=True, null=True)), + ( + "id", + models.UUIDField( + default=posthog.uuidt.uuid7, + editable=False, + primary_key=True, + serialize=False, + ), + ), + ("environment", models.CharField(choices=[("local", "local"), ("cloud", "cloud")], max_length=20)), + ( + "workspace_type", + models.CharField( + choices=[("local_folder", "local_folder"), ("git_repository", "git_repository")], + max_length=30, + ), + ), + ("workspace", models.JSONField()), + ("program", models.JSONField()), + ( + "status", + models.CharField( + choices=[ + ("created", "created"), + ("running", "running"), + ("completed", "completed"), + ("failed", "failed"), + ("cancelled", "cancelled"), + ], + max_length=20, + ), + ), + ( + "error_code", + models.CharField( + blank=True, + max_length=50, + null=True, + ), + ), + ("idempotency_key", models.CharField(blank=True, max_length=255, null=True)), + ("request_fingerprint", models.CharField(blank=True, max_length=64, null=True)), + ( + "dispatch_status", + models.CharField( + blank=True, + choices=[("pending", "pending"), ("dispatched", "dispatched")], + max_length=20, + null=True, + ), + ), + ("dispatch_attempts", models.PositiveSmallIntegerField(default=0)), + ("dispatch_error", models.CharField(blank=True, max_length=255, null=True)), + ("dispatch_next_attempt_at", models.DateTimeField(blank=True, null=True)), + ("workflow_id", models.CharField(blank=True, max_length=255, null=True)), + ( + "stage", + models.CharField( + blank=True, + choices=[ + ("dispatching", "dispatching"), + ("provisioning", "provisioning"), + ("preparing_workspace", "preparing_workspace"), + ("executing_wizard", "executing_wizard"), + ("creating_artifacts", "creating_artifacts"), + ], + max_length=30, + null=True, + ), + ), + ("stage_started_at", models.DateTimeField(blank=True, null=True)), + ("error_message", models.CharField(blank=True, max_length=255, null=True)), + ("started_at", models.DateTimeField(blank=True, null=True)), + ("finished_at", models.DateTimeField(blank=True, null=True)), + ("deadline_at", models.DateTimeField(blank=True, null=True)), + ("cancellation_requested_at", models.DateTimeField(blank=True, null=True)), + ("cancellation_dispatched_at", models.DateTimeField(blank=True, null=True)), + ( + "created_by", + models.ForeignKey( + blank=True, + db_constraint=False, + null=True, + on_delete=django.db.models.deletion.SET_NULL, + related_name="+", + to=settings.AUTH_USER_MODEL, + ), + ), + ( + "team", + models.ForeignKey( + db_constraint=False, + on_delete=django.db.models.deletion.CASCADE, + related_name="+", + to="posthog.team", + ), + ), + ], + options={ + "abstract": False, + "indexes": [ + models.Index( + fields=["dispatch_next_attempt_at", "created_at"], + condition=models.Q(status="created", dispatch_status="pending"), + name="wizard_run_dispatch_idx", + ), + models.Index( + fields=["id"], + condition=models.Q( + status__in=("cancelled", "failed"), + cancellation_requested_at__isnull=False, + cancellation_dispatched_at__isnull=True, + ), + name="wizard_run_cancel_idx", + ), + models.Index( + fields=["deadline_at"], + condition=models.Q(status__in=("created", "running"), deadline_at__isnull=False), + name="wizard_run_deadline_idx", + ), + ], + }, + ), + migrations.AddConstraint( + model_name="wizardrun", + constraint=models.UniqueConstraint( + condition=models.Q(("idempotency_key__isnull", False)), + fields=("team", "idempotency_key"), + name="unique_wizard_run_idempotency_key_per_team", + ), + ), + ] diff --git a/products/wizard/backend/migrations/0007_create_wizard_run_artifact.py b/products/wizard/backend/migrations/0007_create_wizard_run_artifact.py new file mode 100644 index 000000000000..d886dc9b93a8 --- /dev/null +++ b/products/wizard/backend/migrations/0007_create_wizard_run_artifact.py @@ -0,0 +1,63 @@ +import django.db.models.deletion +from django.db import migrations, models + +import posthog.uuidt + + +class Migration(migrations.Migration): + dependencies = [("wizard", "0006_create_wizard_run")] + + operations = [ + migrations.CreateModel( + name="WizardRunArtifact", + fields=[ + ( + "id", + models.UUIDField( + default=posthog.uuidt.uuid7, + editable=False, + primary_key=True, + serialize=False, + ), + ), + ( + "type", + models.CharField( + choices=[("git_diff", "git_diff"), ("pull_request", "pull_request")], + max_length=30, + ), + ), + ("storage_path", models.CharField(blank=True, max_length=1024, null=True)), + ("external_url", models.URLField(blank=True, max_length=1024, null=True)), + ("metadata", models.JSONField(blank=True, null=True)), + ("size_bytes", models.PositiveBigIntegerField(blank=True, null=True)), + ("content_hash", models.CharField(blank=True, max_length=64, null=True)), + ("created_at", models.DateTimeField(auto_now_add=True)), + ( + "run", + models.ForeignKey( + on_delete=django.db.models.deletion.CASCADE, + related_name="artifacts", + to="wizard.wizardrun", + ), + ), + ( + "team", + models.ForeignKey( + db_constraint=False, + on_delete=django.db.models.deletion.CASCADE, + related_name="+", + to="posthog.team", + ), + ), + ], + options={"abstract": False}, + ), + migrations.AddConstraint( + model_name="wizardrunartifact", + constraint=models.UniqueConstraint( + fields=("run", "type"), + name="unique_wizard_artifact_type_per_run", + ), + ), + ] diff --git a/products/wizard/backend/migrations/0008_create_wizard_worker.py b/products/wizard/backend/migrations/0008_create_wizard_worker.py new file mode 100644 index 000000000000..72ec3c35a04a --- /dev/null +++ b/products/wizard/backend/migrations/0008_create_wizard_worker.py @@ -0,0 +1,72 @@ +import django.db.models.deletion +from django.db import migrations, models + +import posthog.uuidt + + +class Migration(migrations.Migration): + dependencies = [("wizard", "0007_create_wizard_run_artifact")] + + operations = [ + migrations.CreateModel( + name="WizardWorker", + fields=[ + ("updated_at", models.DateTimeField(auto_now=True, null=True)), + ( + "id", + models.UUIDField( + default=posthog.uuidt.uuid7, + editable=False, + primary_key=True, + serialize=False, + ), + ), + ("created_at", models.DateTimeField(auto_now_add=True)), + ("sandbox_id", models.CharField(blank=True, max_length=255, null=True, unique=True)), + ("resource_usage", models.JSONField(blank=True, null=True)), + ( + "cleanup_status", + models.CharField( + choices=[ + ("active", "active"), + ("pending", "pending"), + ("cleaned", "cleaned"), + ("failed", "failed"), + ], + default="active", + max_length=20, + ), + ), + ("cleanup_attempts", models.PositiveSmallIntegerField(default=0)), + ("cleanup_error", models.CharField(blank=True, max_length=255, null=True)), + ("cleaned_at", models.DateTimeField(blank=True, null=True)), + ( + "run", + models.OneToOneField( + on_delete=django.db.models.deletion.CASCADE, + related_name="worker", + to="wizard.wizardrun", + ), + ), + ( + "team", + models.ForeignKey( + db_constraint=False, + on_delete=django.db.models.deletion.CASCADE, + related_name="+", + to="posthog.team", + ), + ), + ], + options={ + "abstract": False, + "indexes": [ + models.Index( + fields=["cleanup_attempts"], + condition=models.Q(cleanup_status__in=("active", "pending"), sandbox_id__isnull=False), + name="wizard_worker_cleanup_idx", + ) + ], + }, + ), + ] diff --git a/products/wizard/backend/migrations/max_migration.txt b/products/wizard/backend/migrations/max_migration.txt index af339652e44a..f08bf49c32fb 100644 --- a/products/wizard/backend/migrations/max_migration.txt +++ b/products/wizard/backend/migrations/max_migration.txt @@ -1 +1 @@ -0005_wizardsession_handoff_text +0008_create_wizard_worker diff --git a/products/wizard/backend/models.py b/products/wizard/backend/models.py index 5157cbc45265..64926ab953eb 100644 --- a/products/wizard/backend/models.py +++ b/products/wizard/backend/models.py @@ -10,9 +10,18 @@ from django.db import models from posthog.models.scoping.root_mixin import TeamScopedRootMixin -from posthog.models.utils import CreatedMetaFields, UUIDModel +from posthog.models.utils import CreatedMetaFields, UpdatedMetaFields, UUIDModel -from products.wizard.backend.facade.enums import RunPhase +from products.wizard.backend.facade.enums import ( + WizardRunArtifactType, + WizardRunDispatchStatus, + WizardRunEnvironment, + WizardRunStage, + WizardRunStatus, + WizardSessionRunPhase, + WizardWorkerCleanupStatus, + WizardWorkspaceType, +) class WizardSession(UUIDModel, TeamScopedRootMixin, CreatedMetaFields): @@ -33,7 +42,7 @@ class WizardSession(UUIDModel, TeamScopedRootMixin, CreatedMetaFields): skill_id = models.CharField(max_length=255) started_at = models.DateTimeField() - run_phase = models.CharField(max_length=50, choices=[(phase.value, phase.value) for phase in RunPhase]) + run_phase = models.CharField(max_length=50, choices=[(phase.value, phase.value) for phase in WizardSessionRunPhase]) tasks = models.JSONField(default=list) event_plan = models.JSONField(null=True, blank=True) @@ -63,3 +72,144 @@ class Meta(TeamScopedRootMixin.Meta): name="wizard_sess_team_wf_start_idx", ), ] + + +class WizardRun(UUIDModel, TeamScopedRootMixin, CreatedMetaFields, UpdatedMetaFields): + team = models.ForeignKey("posthog.Team", on_delete=models.CASCADE, related_name="+", db_constraint=False) + + created_by = models.ForeignKey( + "posthog.User", + on_delete=models.SET_NULL, + null=True, + blank=True, + related_name="+", + db_constraint=False, + ) + + environment = models.CharField( + max_length=20, + choices=[(environment.value, environment.value) for environment in WizardRunEnvironment], + ) + + workspace_type = models.CharField( + max_length=30, + choices=[(workspace_type.value, workspace_type.value) for workspace_type in WizardWorkspaceType], + ) + + workspace = models.JSONField() + + program = models.JSONField() + + status = models.CharField( + max_length=20, + choices=[(status.value, status.value) for status in WizardRunStatus], + ) + + error_code = models.CharField(max_length=50, null=True, blank=True) + + idempotency_key = models.CharField(max_length=255, null=True, blank=True) + request_fingerprint = models.CharField(max_length=64, null=True, blank=True) + + dispatch_status = models.CharField( + max_length=20, + choices=[(status.value, status.value) for status in WizardRunDispatchStatus], + null=True, + blank=True, + ) + dispatch_attempts = models.PositiveSmallIntegerField(default=0) + dispatch_error = models.CharField(max_length=255, null=True, blank=True) + dispatch_next_attempt_at = models.DateTimeField(null=True, blank=True) + workflow_id = models.CharField(max_length=255, null=True, blank=True) + + stage = models.CharField( + max_length=30, + choices=[(stage.value, stage.value) for stage in WizardRunStage], + null=True, + blank=True, + ) + stage_started_at = models.DateTimeField(null=True, blank=True) + error_message = models.CharField(max_length=255, null=True, blank=True) + started_at = models.DateTimeField(null=True, blank=True) + finished_at = models.DateTimeField(null=True, blank=True) + deadline_at = models.DateTimeField(null=True, blank=True) + cancellation_requested_at = models.DateTimeField(null=True, blank=True) + cancellation_dispatched_at = models.DateTimeField(null=True, blank=True) + + class Meta(TeamScopedRootMixin.Meta): + constraints = [ + models.UniqueConstraint( + fields=["team", "idempotency_key"], + condition=models.Q(idempotency_key__isnull=False), + name="unique_wizard_run_idempotency_key_per_team", + ) + ] + indexes = [ + models.Index( + fields=["dispatch_next_attempt_at", "created_at"], + condition=models.Q(status="created", dispatch_status="pending"), + name="wizard_run_dispatch_idx", + ), + models.Index( + fields=["id"], + condition=models.Q( + status__in=("cancelled", "failed"), + cancellation_requested_at__isnull=False, + cancellation_dispatched_at__isnull=True, + ), + name="wizard_run_cancel_idx", + ), + models.Index( + fields=["deadline_at"], + condition=models.Q(status__in=("created", "running"), deadline_at__isnull=False), + name="wizard_run_deadline_idx", + ), + ] + + +class WizardRunArtifact(UUIDModel, TeamScopedRootMixin): + team = models.ForeignKey("posthog.Team", on_delete=models.CASCADE, related_name="+", db_constraint=False) + run = models.ForeignKey(WizardRun, on_delete=models.CASCADE, related_name="artifacts") + type = models.CharField( + max_length=30, + choices=[(artifact_type.value, artifact_type.value) for artifact_type in WizardRunArtifactType], + ) + storage_path = models.CharField(max_length=1024, null=True, blank=True) + external_url = models.URLField(max_length=1024, null=True, blank=True) + metadata = models.JSONField(null=True, blank=True) + size_bytes = models.PositiveBigIntegerField(null=True, blank=True) + content_hash = models.CharField(max_length=64, null=True, blank=True) + created_at = models.DateTimeField(auto_now_add=True) + + class Meta(TeamScopedRootMixin.Meta): + constraints = [models.UniqueConstraint(fields=["run", "type"], name="unique_wizard_artifact_type_per_run")] + + +# This tracks provisioned Wizard Workers, helping monitor resource +# usage, and handle provisioning-related errors. +class WizardWorker(UUIDModel, TeamScopedRootMixin, UpdatedMetaFields): + team = models.ForeignKey("posthog.Team", on_delete=models.CASCADE, related_name="+", db_constraint=False) + run = models.OneToOneField(WizardRun, on_delete=models.CASCADE, related_name="worker") + + created_at = models.DateTimeField(auto_now_add=True) + + sandbox_id = models.CharField(max_length=255, null=True, blank=True, unique=True) + + resource_usage = models.JSONField(null=True, blank=True) + + cleanup_status = models.CharField( + max_length=20, + choices=[(status.value, status.value) for status in WizardWorkerCleanupStatus], + default=WizardWorkerCleanupStatus.ACTIVE.value, + ) + cleanup_attempts = models.PositiveSmallIntegerField(default=0) + cleanup_error = models.CharField(max_length=255, null=True, blank=True) + cleaned_at = models.DateTimeField(null=True, blank=True) + + class Meta(TeamScopedRootMixin.Meta): + indexes = [ + models.Index( + fields=["cleanup_attempts"], + condition=models.Q(cleanup_status__in=("active", "pending"), sandbox_id__isnull=False), + name="wizard_worker_cleanup_idx", + ) + ] diff --git a/products/wizard/backend/tests/test_registry.py b/products/wizard/backend/tests/test_registry.py new file mode 100644 index 000000000000..c65d0f1a3213 --- /dev/null +++ b/products/wizard/backend/tests/test_registry.py @@ -0,0 +1,202 @@ +import json + +import pytest +from unittest.mock import patch + +from parameterized import parameterized +from posthoganalytics.request import APIError + +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 WizardProgram +from products.wizard.backend.facade.enums import WizardRunEnvironment +from products.wizard.backend.logic.programs import program_from_mapping, program_to_mapping +from products.wizard.backend.logic.registry.parser import parse_registry_payload + +POSTHOG_INTEGRATION_PROGRAM = WizardProgram( + id="posthog-integration", + name="PostHog integration", + description="Set up PostHog SDK integration", + wizard_version=DEFAULT_WIZARD_VERSION, + command=(), + tags=(), + required_programs=(), + supported_environments=(WizardRunEnvironment.LOCAL, WizardRunEnvironment.CLOUD), +) + +AUDIT_PROGRAM_PAYLOAD = { + "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"], +} + +REGISTRY_PAYLOAD = {"version": 1, "programs": [AUDIT_PROGRAM_PAYLOAD]} + + +def test_registry_payload_parsing() -> None: + registry = parse_registry_payload(REGISTRY_PAYLOAD) + + assert registry.programs[0].id == "web-analytics-audit" + + +@parameterized.expand( + [ + ("unexpected_field", {**REGISTRY_PAYLOAD, "unexpected": "value"}), + ("unsupported_version", {**REGISTRY_PAYLOAD, "version": 2}), + ("duplicate_program", {"version": 1, "programs": [AUDIT_PROGRAM_PAYLOAD, AUDIT_PROGRAM_PAYLOAD]}), + ] +) +def test_registry_rejects_invalid_serialized_value(_name: str, value: object) -> None: + with pytest.raises(ValueError): + parse_registry_payload(value) + + +def test_program_serialization_round_trip() -> None: + program = program_from_mapping(AUDIT_PROGRAM_PAYLOAD) + + assert program_to_mapping(program) == AUDIT_PROGRAM_PAYLOAD + assert program_from_mapping(program_to_mapping(program)) == program + + +@parameterized.expand( + [ + ("unexpected_field", {**AUDIT_PROGRAM_PAYLOAD, "unexpected": "value"}), + ("invalid_command", {**AUDIT_PROGRAM_PAYLOAD, "command": ["--override"]}), + ("duplicate_environments", {**AUDIT_PROGRAM_PAYLOAD, "supported_environments": ["local", "local"]}), + ("leading_zero_prerelease", {**AUDIT_PROGRAM_PAYLOAD, "wizard_version": "1.2.3-01"}), + ("leading_zero_dotted_prerelease", {**AUDIT_PROGRAM_PAYLOAD, "wizard_version": "1.2.3-alpha.01"}), + ] +) +def test_program_rejects_invalid_serialized_value(_name: str, value: object) -> None: + with pytest.raises(ValueError): + program_from_mapping(value) + + +def test_program_accepts_valid_prerelease_wizard_version() -> None: + program = program_from_mapping({**AUDIT_PROGRAM_PAYLOAD, "wizard_version": "1.2.3-alpha.1"}) + + assert program.wizard_version == "1.2.3-alpha.1" + + +def test_program_persisted_deserialization_accepts_legacy_wizard_version() -> None: + persisted_value = {**AUDIT_PROGRAM_PAYLOAD, "wizard_version": "latest"} + + with pytest.raises(ValueError): + program_from_mapping(persisted_value) + + assert program_from_mapping(persisted_value, allow_latest_version=True).wizard_version == "latest" + + +@parameterized.expand( + [ + ("object", REGISTRY_PAYLOAD), + ("json_string", json.dumps(REGISTRY_PAYLOAD)), + ] +) +def test_registry_returns_personalized_programs(_name: str, payload: object) -> None: + with patch("posthoganalytics.get_feature_flag_payload", return_value=payload) as get_payload: + programs = wizard_facade.get_registry(distinct_id="user-distinct-id", organization_id="organization-id") + + assert programs == ( + WizardProgram( + 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=(WizardRunEnvironment.LOCAL,), + ), + ) + get_payload.assert_called_once_with( + "wizard-program-registry", + distinct_id="user-distinct-id", + groups={"organization": "organization-id"}, + group_properties={"organization": {"id": "organization-id"}}, + only_evaluate_locally=False, + send_feature_flag_events=False, + ) + + +def test_registry_preserves_valid_empty_program_list() -> None: + with patch("posthoganalytics.get_feature_flag_payload", return_value={"version": 1, "programs": []}): + programs = wizard_facade.get_registry(distinct_id="user-distinct-id", organization_id="organization-id") + + assert programs == () + + +@parameterized.expand( + [ + ("missing_payload", None), + ("invalid_json", "{"), + ("unsupported_version", {"version": 2, "programs": []}), + ("duplicate_ids", {"version": 1, "programs": [AUDIT_PROGRAM_PAYLOAD, AUDIT_PROGRAM_PAYLOAD]}), + ( + "missing_wizard_version", + { + "version": 1, + "programs": [{key: value for key, value in AUDIT_PROGRAM_PAYLOAD.items() if key != "wizard_version"}], + }, + ), + ( + "mutable_wizard_version", + { + "version": 1, + "programs": [{**AUDIT_PROGRAM_PAYLOAD, "wizard_version": "latest"}], + }, + ), + ( + "invalid_program", + { + "version": 1, + "programs": [ + AUDIT_PROGRAM_PAYLOAD, + {**AUDIT_PROGRAM_PAYLOAD, "id": "invalid-program", "command": ["--override"]}, + ], + }, + ), + ( + "unknown_environment", + { + "version": 1, + "programs": [ + {**AUDIT_PROGRAM_PAYLOAD, "supported_environments": ["hosted"]}, + ], + }, + ), + ] +) +def test_registry_falls_back_when_payload_is_invalid(_name: str, payload: object) -> None: + with ( + patch("posthoganalytics.get_feature_flag_payload", return_value=payload), + patch("products.wizard.backend.logic.registry.service.report_registry_fallback") as report_fallback, + ): + programs = wizard_facade.get_registry(distinct_id="user-distinct-id", organization_id="organization-id") + + assert programs == (POSTHOG_INTEGRATION_PROGRAM,) + report_fallback.assert_called_once_with("invalid_payload") + + +def test_registry_does_not_hide_programming_errors() -> None: + with ( + patch("posthoganalytics.get_feature_flag_payload", side_effect=RuntimeError("bug")), + pytest.raises(RuntimeError, match="bug"), + ): + wizard_facade.get_registry(distinct_id="user-distinct-id", organization_id="organization-id") + + +def test_registry_falls_back_when_feature_flag_fetch_fails() -> None: + with ( + patch("posthoganalytics.get_feature_flag_payload", side_effect=APIError(503, "Unavailable")), + patch("products.wizard.backend.logic.registry.service.report_registry_fallback") as report_fallback, + ): + programs = wizard_facade.get_registry(distinct_id="user-distinct-id", organization_id="organization-id") + + assert programs == (POSTHOG_INTEGRATION_PROGRAM,) + report_fallback.assert_called_once_with("request_failed")