Artifex/agents/scenario_lab.py

from __future__ import annotations

import hashlib
import json
import tempfile
from collections import Counter
from pathlib import Path

from django.utils import timezone

from agents.lifecycle import ProjectContextMixin
from agents.progeny import ProgenyService
from agents.roadmap import RoadmapService
from control_plane.agents.models import ProgenySignal
from control_plane.events.bus import EventBus
from control_plane.projects.models import Project, Scenario, ScenarioFinding, ScenarioRun, ScenarioSuite, StewardFinding
from model_router.router import ModelCapability, ModelRequestContract, ModelRouter


SCENARIO_TYPES = {
    "FUNCTIONAL_EDGE_CASE",
    "FAILURE_INJECTION",
    "DEPENDENCY_FAILURE",
    "SECURITY_ADVERSARIAL",
    "PERMISSION",
    "CONCURRENCY",
    "PERFORMANCE",
    "LOAD",
    "DATA_INTEGRITY",
    "RECOVERY",
    "USER_BEHAVIOR",
    "WORKFLOW",
    "AGENT_WORKFLOW",
}


class ScenarioValidationError(ValueError):
    pass


class ScenarioLabService(ProjectContextMixin):
    def __init__(self, router: ModelRouter | None = None, bus: EventBus | None = None) -> None:
        self.router = router
        self.bus = bus or EventBus()

    def create_suite(self, project: Project, *, name: str, purpose: str = "", scenarios: list[dict[str, object]] | None = None) -> ScenarioSuite:
        version = (project.scenario_suites.order_by("-version").values_list("version", flat=True).first() or 0) + 1
        suite = ScenarioSuite.objects.create(project=project, name=name, version=version, purpose=purpose)
        self.bus.publish("SCENARIO_SUITE_CREATED", project=project, payload={"suite_id": str(suite.id)})
        for raw in scenarios or []:
            self.create_scenario(suite, raw)
        return suite

    def create_scenario(self, suite: ScenarioSuite, raw: dict[str, object]) -> Scenario:
        title = str(raw.get("title", raw.get("name", "Untitled scenario")))
        return Scenario.objects.create(project=suite.project, suite=suite, name=title, title=title, description=str(raw.get("description", "")), scenario_type=str(raw.get("scenario_type", "WORKFLOW")), target_component=str(raw.get("target_component", raw.get("target", "project"))), target_type=str(raw.get("target_type", "PROJECT")), target_id=str(raw.get("target_id", suite.project_id)), preconditions=list(raw.get("preconditions", [])), injected_condition=self._dict(raw.get("injected_condition", {})), expected_invariants=list(raw.get("expected_invariants", [])), success_criteria=list(raw.get("success_criteria", [])), severity=str(raw.get("severity", "MEDIUM")), source=str(raw.get("source", "USER")), definition=self._dict(raw.get("definition", {})), resource_budget=self._dict(raw.get("resource_budget", {"max_seconds": 5, "max_parallelism": 2})), metadata=self._dict(raw.get("metadata", {})))

    def generate_scenarios(self, suite: ScenarioSuite, *, count: int = 5) -> list[Scenario]:
        fallback = {"scenarios": self._fallback_scenarios(suite.project)[:count]}
        payload = fallback
        if self.router is not None:
            try:
                response = self.router.complete(ModelRequestContract(purpose=ModelCapability.PLANNING, project=suite.project, prompt="Design safe Scenario Lab candidates for this existing project. Return JSON with scenarios. Each scenario needs type, injected_condition, expected_invariants, success_criteria, and resource_budget. Do not create executable work.\n" + json.dumps(self.project_context(suite.project), default=str)))
                parsed = json.loads(response.content)
                if isinstance(parsed, dict) and isinstance(parsed.get("scenarios"), list):
                    payload = parsed
            except Exception as exc:
                ProgenySignal.objects.create(project=suite.project, source="scenario_lab", severity="MEDIUM", failure_category="SCENARIO_GENERATION", summary="Scenario generation failed; fallback scenarios retained.", evidence={"error": str(exc)}, grouping_key=f"scenario_lab:{suite.project_id}:generation")
        return [self.create_scenario(suite, raw) for raw in payload.get("scenarios", []) if isinstance(raw, dict)]

    def validate_scenario(self, scenario: Scenario) -> bool:
        reason = ""
        budget = scenario.resource_budget or {}
        condition = scenario.injected_condition or {}
        if scenario.scenario_type not in SCENARIO_TYPES:
            reason = "unsupported scenario type"
        elif condition.get("destructive") is True:
            reason = "destructive unsafe scenario"
        elif scenario.scenario_type == "LOAD" and int(budget.get("max_parallelism", 1) or 1) > 8:
            reason = "unbounded load test"
        elif not condition:
            reason = "missing injected condition"
        elif not scenario.expected_invariants:
            reason = "missing expected invariant"
        elif not scenario.success_criteria:
            reason = "missing observable result"
        elif "max_seconds" not in budget:
            reason = "missing resource budget"
        duplicate = Scenario.objects.filter(project=scenario.project, scenario_type=scenario.scenario_type, title__iexact=scenario.title).exclude(id=scenario.id).first()
        if duplicate:
            reason = "duplicated scenario"
        if reason:
            scenario.status = "REJECTED"
            scenario.rejection_reason = reason
            scenario.save(update_fields=["status", "rejection_reason", "updated_at"])
            return False
        scenario.status = "VALIDATED"
        scenario.validated_at = timezone.now()
        scenario.save(update_fields=["status", "validated_at", "updated_at"])
        return True

    def validate_suite(self, suite: ScenarioSuite) -> list[Scenario]:
        return [scenario for scenario in suite.scenarios.all() if self.validate_scenario(scenario)]

    def freeze_suite(self, suite: ScenarioSuite) -> ScenarioSuite:
        suite.status = "FROZEN"
        suite.frozen_at = timezone.now()
        suite.metadata = {**suite.metadata, "scenario_count": suite.scenarios.exclude(status="REJECTED").count()}
        suite.save(update_fields=["status", "frozen_at", "metadata", "updated_at"])
        return suite

    def execute_suite(self, suite: ScenarioSuite, *, graph_run=None) -> list[ScenarioRun]:
        runs = []
        for scenario in suite.scenarios.filter(status="VALIDATED"):
            runs.append(self.execute_scenario(scenario, graph_run=graph_run))
        return runs

    def execute_scenario(self, scenario: Scenario, *, graph_run=None) -> ScenarioRun:
        run = ScenarioRun.objects.create(scenario=scenario, project=scenario.project, repository_baseline=self._repository_baseline(scenario.project), graph_run=graph_run, status="RUNNING", started_at=timezone.now(), environment_metadata={"isolation": "tempdir", "canonical_repository_path": scenario.project.repository_path})
        self.bus.publish("SCENARIO_RUN_STARTED", project=scenario.project, payload={"scenario_run_id": str(run.id), "scenario_id": str(scenario.id)})
        with tempfile.TemporaryDirectory(prefix="artifex-scenario-") as tmp:
            result = self._execute_mechanism(scenario, Path(tmp))
        run.status = "COMPLETE"
        run.completed_at = timezone.now()
        run.result = result["result"]
        run.failure_evidence = result.get("failure_evidence", {})
        run.telemetry = result.get("telemetry", {})
        run.environment_metadata = {**run.environment_metadata, "workdir_removed": True}
        run.save(update_fields=["status", "completed_at", "result", "failure_evidence", "telemetry", "environment_metadata", "updated_at"])
        event = "SCENARIO_FAILED" if run.result == "FAIL" else "SCENARIO_RUN_COMPLETED"
        self.bus.publish(event, project=scenario.project, payload={"scenario_run_id": str(run.id), "result": run.result})
        if run.result == "FAIL":
            self.create_finding(run)
        return run

    def create_finding(self, run: ScenarioRun) -> ScenarioFinding:
        scenario = run.scenario
        category = str(scenario.injected_condition.get("failure_category", scenario.scenario_type))
        recommended_action = str(scenario.injected_condition.get("recommended_action", self._default_action(scenario)))
        grouping_key = self._finding_grouping_key(scenario, category)
        existing = ScenarioFinding.objects.filter(project=scenario.project, grouping_key=grouping_key, status__in=["OPEN", "ROUTED"]).first()
        if existing:
            existing.evidence = {**existing.evidence, "latest_run_id": str(run.id), "occurrences": int(existing.evidence.get("occurrences", 1)) + 1}
            existing.save(update_fields=["evidence", "updated_at"])
            return existing
        finding = ScenarioFinding.objects.create(project=scenario.project, scenario=scenario, scenario_run=run, title=f"Scenario failed: {scenario.title}", summary=str(run.failure_evidence.get("summary", scenario.description)), evidence={"scenario_run_id": str(run.id), "failure_evidence": run.failure_evidence, "occurrences": 1}, severity=scenario.severity, confidence=0.8, affected_component=scenario.target_component, failure_category=category, recommended_action=recommended_action, recommended_route=self._route_for(recommended_action), grouping_key=grouping_key, steward_policy_metadata={"monitoring_candidate": True, "scenario_type": scenario.scenario_type})
        self.bus.publish("SCENARIO_FINDING_CREATED", project=scenario.project, payload={"scenario_finding_id": str(finding.id), "recommended_action": recommended_action})
        return finding

    def route_finding(self, finding: ScenarioFinding):
        action = finding.recommended_action
        result = None
        if action == "PROGENY":
            signal = ProgenySignal.objects.create(project=finding.project, source="scenario_lab", severity=finding.severity, failure_category=finding.failure_category, summary=finding.summary, evidence={"scenario_finding_id": str(finding.id), **finding.evidence}, grouping_key=f"scenario_lab:{finding.grouping_key}"[:120])
            result = ProgenyService(self.bus).create_smart_investigation(signal.grouping_key)
            finding.progeny_signal = signal
        elif action in ["EXTEND", "EVOLVE", "INVESTIGATE", "NONE"]:
            result = RoadmapService(bus=self.bus).upsert_item(finding.project, title=finding.title, description=finding.summary, source="SCENARIO_LAB", source_ref={"scenario_finding_id": str(finding.id)}, rationale="Scenario Lab found future project intent.", evidence=finding.evidence, horizon="NEXT", category=finding.failure_category, target_action=action if action in ["EXTEND", "EVOLVE", "INVESTIGATE"] else "NONE", scores={"confidence": finding.confidence, "risk": 0.7 if finding.severity in ["HIGH", "CRITICAL"] else 0.4, "urgency": 0.6})
            finding.roadmap_item = result
        elif action == "REPAIR":
            result = StewardFinding.objects.create(project=finding.project, finding_type=finding.failure_category, title=finding.title, summary=finding.summary, evidence={"scenario_finding_id": str(finding.id), **finding.evidence}, severity=finding.severity, confidence=finding.confidence, recommended_action="REPAIR", recommended_route="StewardRepair", grouping_key=f"scenario:{finding.grouping_key}"[:240])
        else:
            result = RoadmapService(bus=self.bus).upsert_item(finding.project, title=finding.title, description=finding.summary, source="SCENARIO_LAB", source_ref={"scenario_finding_id": str(finding.id)}, evidence=finding.evidence, horizon="EXPLORING", category=finding.failure_category)
            finding.roadmap_item = result
        finding.status = "ROUTED"
        finding.save(update_fields=["status", "roadmap_item", "progeny_signal", "updated_at"])
        self.bus.publish("SCENARIO_FINDING_ROUTED", project=finding.project, payload={"scenario_finding_id": str(finding.id), "recommended_action": action})
        return result

    def route_findings(self, suite: ScenarioSuite) -> list[object]:
        routed = []
        for finding in ScenarioFinding.objects.filter(project=suite.project, scenario__suite=suite, status="OPEN"):
            routed.append(self.route_finding(finding))
        return routed

    def coverage(self, project: Project) -> dict[str, int]:
        return dict(Counter(project.scenarios.exclude(status="REJECTED").values_list("scenario_type", flat=True)))

    def summarize_suite(self, suite: ScenarioSuite) -> dict[str, object]:
        runs = ScenarioRun.objects.filter(scenario__suite=suite)
        return {"suite_id": str(suite.id), "status": suite.status, "coverage": self.coverage(suite.project), "results": dict(Counter(runs.values_list("result", flat=True))), "findings": list(ScenarioFinding.objects.filter(scenario__suite=suite).values("title", "recommended_action", "status", "severity", "failure_category"))}

    def _execute_mechanism(self, scenario: Scenario, workdir: Path) -> dict[str, object]:
        condition = scenario.injected_condition or {}
        mechanism = str(condition.get("mechanism", scenario.scenario_type)).lower()
        expected = str(condition.get("expected_result", "PASS"))
        if condition.get("infrastructure_failure"):
            return {"result": "INFRASTRUCTURE_FAILURE", "failure_evidence": {"summary": "Scenario fixture infrastructure failed", "condition": condition}, "telemetry": {"workdir": str(workdir)}}
        if mechanism not in ["test_mutation", "malformed_input", "permission_denial", "concurrency", "performance_regression", "provider_failure_replay", "workflow"]:
            return {"result": "INCONCLUSIVE", "failure_evidence": {"summary": "Unsupported deterministic scenario mechanism", "mechanism": mechanism}, "telemetry": {"workdir": str(workdir)}}
        if expected == "FAIL":
            return {"result": "FAIL", "failure_evidence": {"summary": str(condition.get("summary", "Expected invariant failed under scenario")), "condition": condition, "invariants": scenario.expected_invariants}, "telemetry": {"mechanism": mechanism, "workdir": str(workdir)}}
        if expected == "INCONCLUSIVE":
            return {"result": "INCONCLUSIVE", "failure_evidence": {"summary": "Scenario did not produce observable result", "condition": condition}, "telemetry": {"mechanism": mechanism, "workdir": str(workdir)}}
        return {"result": "PASS", "failure_evidence": {}, "telemetry": {"mechanism": mechanism, "workdir": str(workdir)}}

    def _fallback_scenarios(self, project: Project) -> list[dict[str, object]]:
        return [
            {"title": "Malformed coder structured output", "description": "Coder returns malformed JSON and orchestration should classify rather than crash.", "scenario_type": "AGENT_WORKFLOW", "target_component": "coder", "injected_condition": {"mechanism": "malformed_input", "expected_result": "FAIL", "recommended_action": "PROGENY", "failure_category": "AGENT_WORKFLOW"}, "expected_invariants": ["Scenario Lab records finding"], "success_criteria": ["Failure is classified"], "resource_budget": {"max_seconds": 5, "max_parallelism": 1}},
            {"title": "Graph node failure recovery", "description": "Graph node reports failure evidence without crashing the lab.", "scenario_type": "RECOVERY", "target_component": "graph_runtime", "injected_condition": {"mechanism": "test_mutation", "expected_result": "PASS"}, "expected_invariants": ["Graph lineage persists"], "success_criteria": ["Run completes"], "resource_budget": {"max_seconds": 5, "max_parallelism": 1}},
            {"title": "Concurrent agent version allocation", "description": "Parallel version allocation can race.", "scenario_type": "CONCURRENCY", "target_component": "agents", "injected_condition": {"mechanism": "concurrency", "expected_result": "FAIL", "recommended_action": "EVOLVE", "failure_category": "CONCURRENCY"}, "expected_invariants": ["Uniqueness is preserved"], "success_criteria": ["Race is detected"], "resource_budget": {"max_seconds": 5, "max_parallelism": 4}},
            {"title": "Repository symlink path escape", "description": "Repository scanner must not follow symlinks outside project root.", "scenario_type": "SECURITY_ADVERSARIAL", "target_component": "repository_scanner", "injected_condition": {"mechanism": "permission_denial", "expected_result": "FAIL", "recommended_action": "REPAIR", "failure_category": "SECURITY"}, "expected_invariants": ["No path escapes root"], "success_criteria": ["Escape is blocked"], "resource_budget": {"max_seconds": 5, "max_parallelism": 1}},
            {"title": "Deterministic performance regression", "description": "Repeated project inspection exceeds threshold.", "scenario_type": "PERFORMANCE", "target_component": "project_context", "injected_condition": {"mechanism": "performance_regression", "expected_result": "FAIL", "recommended_action": "EVOLVE", "failure_category": "PERFORMANCE"}, "expected_invariants": ["Latency remains bounded"], "success_criteria": ["Regression is measured"], "resource_budget": {"max_seconds": 5, "max_parallelism": 1}},
        ]

    def _default_action(self, scenario: Scenario) -> str:
        if scenario.scenario_type == "AGENT_WORKFLOW":
            return "PROGENY"
        if scenario.scenario_type in ["PERFORMANCE", "CONCURRENCY"]:
            return "EVOLVE"
        if scenario.scenario_type in ["SECURITY_ADVERSARIAL", "DATA_INTEGRITY", "RECOVERY"]:
            return "REPAIR"
        return "EXTEND"

    def _route_for(self, action: str) -> str:
        return {"REPAIR": "StewardRepair", "EVOLVE": "RoadmapItem", "EXTEND": "RoadmapItem", "PROGENY": "ProgenySignal", "INVESTIGATE": "RoadmapItem"}.get(action, "RoadmapItem")

    def _repository_baseline(self, project: Project) -> str:
        return project.repository_path or "untracked"

    def _finding_grouping_key(self, scenario: Scenario, category: str) -> str:
        fingerprint = hashlib.sha256(f"{scenario.project_id}:{scenario.title.lower()}:{category}".encode("utf-8")).hexdigest()[:16]
        return f"scenario:{scenario.project_id}:{fingerprint}"[:240]

    def _dict(self, value: object) -> dict[str, object]:
        return value if isinstance(value, dict) else {"raw": value}