AirCoding/AirPlan/docs/spec/AirPlan-ParaV2/lib/air_runtime/contracts.py

from __future__ import annotations

from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from typing import Any, Dict, List

from .paths import required_project_artifacts

KNOWN_TASK_STATUSES = {"TODO", "DOING", "DONE", "BLOCKED"}
ACTIVE_TASK_STATUSES = {"TODO", "DOING"}
FINAL_WORKER_STATUSES = {"done", "blocked", "skipped"}
DOCUMENT_UPDATE_ACTIONS = {"replace_block", "append_lines", "create_file"}
DEFAULT_WORKER_RESULT_NOTE = (
    "Do not merge this result until summary, filesChanged, validations, and risks are reviewed."
)
DEFAULT_DEBUG_POLICY = {
    "enabled": True,
    "triggerOnBlocked": True,
    "triggerValidationStatuses": ["failed", "error"],
    "maxSessionsPerTask": 1,
    "preferOriginalAirDbg": True,
}
DEFAULT_XDB_POLICY = {
    "enabled": True,
    "requireForGuiTasks": True,
    "captureOnBlocked": True,
    "captureOnDone": True,
    "triggerValidationStatuses": ["failed", "error"],
    "maxSessionsPerTask": 1,
    "preferOriginalAirXDB": True,
    "remoteFirstIfConfigured": True,
    "taskKeywords": [
        "gui",
        "ui",
        "desktop",
        "browser",
        "screen",
        "screenshot",
        "visual",
        "layout",
        "popup",
        "focus",
        "canvas",
        "acceptance",
        "midscene",
        "airxdb",
        "lvgl",
        "lcd",
        "qt-app",
    ],
}
DEFAULT_REPAIR_POLICY = {
    "enabled": True,
    "triggerOnBlocked": True,
    "triggerValidationStatuses": ["failed", "error"],
    "maxAttemptsPerTask": 2,
    "requeueTodoStatus": "DOING",
    "autoPrepareWorker": True,
}
REQUIRED_PROJECT_ARTIFACTS = required_project_artifacts()


def now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _dedupe(items: List[str]) -> List[str]:
    seen = set()
    ordered: List[str] = []
    for item in items:
        cleaned = str(item).strip()
        if not cleaned or cleaned in seen:
            continue
        seen.add(cleaned)
        ordered.append(cleaned)
    return ordered


def _ensure_str_list(value: Any) -> List[str]:
    if value is None:
        return []
    if isinstance(value, list):
        return _dedupe([str(item) for item in value])
    return _dedupe([str(value)])


@dataclass
class ValidationRecord:
    kind: str
    status: str
    command: str = ""
    evidence: List[str] = field(default_factory=list)

    @classmethod
    def from_dict(cls, payload: Dict[str, Any]) -> "ValidationRecord":
        return cls(
            kind=str(payload.get("kind", "")).strip(),
            status=str(payload.get("status", "")).strip(),
            command=str(payload.get("command", "")).strip(),
            evidence=_ensure_str_list(payload.get("evidence")),
        )

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)


@dataclass
class TaskRecord:
    task_id: str
    status: str
    module: str
    task: str
    files_dirs: str
    done_when: str
    validation: str
    adr_c4_update: str
    line_number: int
    dependencies: List[str] = field(default_factory=list)
    write_paths: List[str] = field(default_factory=list)
    global_doc_paths: List[str] = field(default_factory=list)

    def normalized_write_set(self) -> List[str]:
        return _dedupe(self.write_paths + self.global_doc_paths)

    def touches_global_docs(self) -> bool:
        return bool(self.global_doc_paths)

    def to_dict(self) -> Dict[str, Any]:
        payload = asdict(self)
        payload["taskId"] = payload.pop("task_id")
        payload["filesDirs"] = payload.pop("files_dirs")
        payload["doneWhen"] = payload.pop("done_when")
        payload["adrC4Update"] = payload.pop("adr_c4_update")
        payload["lineNumber"] = payload.pop("line_number")
        payload["writeSet"] = payload.pop("write_paths")
        payload["globalDocPaths"] = payload.pop("global_doc_paths")
        return payload

    @classmethod
    def from_dict(cls, payload: Dict[str, Any]) -> "TaskRecord":
        return cls(
            task_id=str(payload.get("taskId", "")).strip(),
            status=str(payload.get("status", "")).strip().upper(),
            module=str(payload.get("module", "")).strip(),
            task=str(payload.get("task", "")).strip(),
            files_dirs=str(payload.get("filesDirs", "")).strip(),
            done_when=str(payload.get("doneWhen", "")).strip(),
            validation=str(payload.get("validation", "")).strip(),
            adr_c4_update=str(payload.get("adrC4Update", "")).strip(),
            line_number=int(payload.get("lineNumber", 0) or 0),
            dependencies=_ensure_str_list(payload.get("dependencies")),
            write_paths=_ensure_str_list(payload.get("writeSet")),
            global_doc_paths=_ensure_str_list(payload.get("globalDocPaths")),
        )


@dataclass
class WorkerResult:
    task_id: str
    status: str
    summary: str
    files_changed: List[str] = field(default_factory=list)
    validations: List[ValidationRecord] = field(default_factory=list)
    evidence_paths: List[str] = field(default_factory=list)
    risks: List[str] = field(default_factory=list)
    blockers: List[str] = field(default_factory=list)
    recommend_global_doc_updates: List[str] = field(default_factory=list)
    global_doc_paths: List[str] = field(default_factory=list)
    document_updates: List["DocumentUpdate"] = field(default_factory=list)
    debug_sessions: List["DebugSession"] = field(default_factory=list)
    xdb_sessions: List["XdbSession"] = field(default_factory=list)
    repair_attempts: List["RepairAttempt"] = field(default_factory=list)
    notes: List[str] = field(default_factory=list)
    finalized_at: str = ""

    @classmethod
    def from_dict(cls, payload: Dict[str, Any]) -> "WorkerResult":
        return cls(
            task_id=str(payload.get("taskId", "")).strip(),
            status=str(payload.get("status", "")).strip().lower(),
            summary=str(payload.get("summary", "")).strip(),
            files_changed=_ensure_str_list(payload.get("filesChanged")),
            validations=[
                ValidationRecord.from_dict(item)
                for item in payload.get("validations", [])
            ],
            evidence_paths=_ensure_str_list(payload.get("evidencePaths")),
            risks=_ensure_str_list(payload.get("risks")),
            blockers=_ensure_str_list(payload.get("blockers")),
            recommend_global_doc_updates=_ensure_str_list(
                payload.get("recommendGlobalDocUpdates")
            ),
            global_doc_paths=_ensure_str_list(payload.get("globalDocPaths")),
            document_updates=[
                DocumentUpdate.from_dict(item)
                for item in payload.get("documentUpdates", [])
            ],
            debug_sessions=[
                DebugSession.from_dict(item)
                for item in payload.get("debugSessions", [])
            ],
            xdb_sessions=[
                XdbSession.from_dict(item)
                for item in payload.get("xdbSessions", [])
            ],
            repair_attempts=[
                RepairAttempt.from_dict(item)
                for item in payload.get("repairAttempts", [])
            ],
            notes=_ensure_str_list(payload.get("notes")),
            finalized_at=str(payload.get("finalizedAt", "")).strip(),
        )

    def validate(self) -> None:
        if not self.task_id:
            raise ValueError("worker result missing taskId")
        if self.status not in FINAL_WORKER_STATUSES:
            raise ValueError(
                f"worker result status must be one of {sorted(FINAL_WORKER_STATUSES)}"
            )
        if not self.summary:
            raise ValueError("worker result summary must not be empty")
        if self.status == "blocked" and not self.blockers:
            raise ValueError("blocked worker result must include blockers")
        for item in self.document_updates:
            item.validate()
        for item in self.debug_sessions:
            item.validate()
        for item in self.xdb_sessions:
            item.validate()
        for item in self.repair_attempts:
            item.validate()

    def has_meaningful_done_payload(self) -> bool:
        return bool(
            self.files_changed
            or self.validations
            or self.evidence_paths
            or self.document_updates
            or self.debug_sessions
            or self.xdb_sessions
        )

    def validate_for_finalize(self) -> None:
        self.validate()
        if self.status == "done" and not self.has_meaningful_done_payload():
            raise ValueError(
                "done worker result must include filesChanged, validations, evidencePaths, "
                "documentUpdates, debugSessions, or xdbSessions before finalize"
            )

    def to_dict(self) -> Dict[str, Any]:
        return {
            "taskId": self.task_id,
            "status": self.status,
            "summary": self.summary,
            "filesChanged": self.files_changed,
            "validations": [item.to_dict() for item in self.validations],
            "evidencePaths": self.evidence_paths,
            "risks": self.risks,
            "blockers": self.blockers,
            "recommendGlobalDocUpdates": self.recommend_global_doc_updates,
            "globalDocPaths": self.global_doc_paths,
            "documentUpdates": [item.to_dict() for item in self.document_updates],
            "debugSessions": [item.to_dict() for item in self.debug_sessions],
            "xdbSessions": [item.to_dict() for item in self.xdb_sessions],
            "repairAttempts": [item.to_dict() for item in self.repair_attempts],
            "notes": self.notes,
            "finalizedAt": self.finalized_at,
        }


def default_worker_result(task_id: str, summary: str = "") -> WorkerResult:
    return WorkerResult(
        task_id=task_id,
        status="done",
        summary=summary,
        files_changed=[],
        validations=[],
        evidence_paths=[],
        risks=[],
        blockers=[],
        recommend_global_doc_updates=[],
        global_doc_paths=[],
        document_updates=[],
        debug_sessions=[],
        xdb_sessions=[],
        repair_attempts=[],
        notes=[
            DEFAULT_WORKER_RESULT_NOTE
        ],
        finalized_at="",
    )


@dataclass
class DebugSession:
    session_id: str
    task_id: str
    source: str
    trigger: str
    reason: str
    request_path: str
    debug_log_path: str
    state_path: str = ""
    mode: str = "original-airdbg"
    status: str = "requested"
    created_at: str = ""

    @classmethod
    def from_dict(cls, payload: Dict[str, Any]) -> "DebugSession":
        return cls(
            session_id=str(payload.get("sessionId", "")).strip(),
            task_id=str(payload.get("taskId", "")).strip(),
            source=str(payload.get("source", "")).strip(),
            trigger=str(payload.get("trigger", "")).strip(),
            reason=str(payload.get("reason", "")).strip(),
            request_path=str(payload.get("requestPath", "")).strip(),
            debug_log_path=str(payload.get("debugLogPath", "")).strip(),
            state_path=str(payload.get("statePath", "")).strip(),
            mode=str(payload.get("mode", "original-airdbg")).strip(),
            status=str(payload.get("status", "requested")).strip(),
            created_at=str(payload.get("createdAt", "")).strip(),
        )

    def validate(self) -> None:
        if not self.session_id:
            raise ValueError("debug session missing sessionId")
        if not self.task_id:
            raise ValueError("debug session missing taskId")
        if not self.request_path:
            raise ValueError("debug session missing requestPath")
        if not self.debug_log_path:
            raise ValueError("debug session missing debugLogPath")

    def to_dict(self) -> Dict[str, Any]:
        return {
            "sessionId": self.session_id,
            "taskId": self.task_id,
            "source": self.source,
            "trigger": self.trigger,
            "reason": self.reason,
            "requestPath": self.request_path,
            "debugLogPath": self.debug_log_path,
            "statePath": self.state_path,
            "mode": self.mode,
            "status": self.status,
            "createdAt": self.created_at,
        }


@dataclass
class XdbSession:
    session_id: str
    task_id: str
    source: str
    trigger: str
    reason: str
    request_path: str
    gui_debug_log_path: str
    report_path: str = ""
    state_path: str = ""
    mode: str = "local-screenshot"
    status: str = "captured"
    screenshots: List[str] = field(default_factory=list)
    created_at: str = ""
    remote_target: str = ""

    @classmethod
    def from_dict(cls, payload: Dict[str, Any]) -> "XdbSession":
        return cls(
            session_id=str(payload.get("sessionId", "")).strip(),
            task_id=str(payload.get("taskId", "")).strip(),
            source=str(payload.get("source", "")).strip(),
            trigger=str(payload.get("trigger", "")).strip(),
            reason=str(payload.get("reason", "")).strip(),
            request_path=str(payload.get("requestPath", "")).strip(),
            gui_debug_log_path=str(payload.get("guiDebugLogPath", "")).strip(),
            report_path=str(payload.get("reportPath", "")).strip(),
            state_path=str(payload.get("statePath", "")).strip(),
            mode=str(payload.get("mode", "local-screenshot")).strip(),
            status=str(payload.get("status", "captured")).strip(),
            screenshots=_ensure_str_list(payload.get("screenshots")),
            created_at=str(payload.get("createdAt", "")).strip(),
            remote_target=str(payload.get("remoteTarget", "")).strip(),
        )

    def validate(self) -> None:
        if not self.session_id:
            raise ValueError("xdb session missing sessionId")
        if not self.task_id:
            raise ValueError("xdb session missing taskId")
        if not self.request_path:
            raise ValueError("xdb session missing requestPath")
        if not self.gui_debug_log_path:
            raise ValueError("xdb session missing guiDebugLogPath")

    def to_dict(self) -> Dict[str, Any]:
        return {
            "sessionId": self.session_id,
            "taskId": self.task_id,
            "source": self.source,
            "trigger": self.trigger,
            "reason": self.reason,
            "requestPath": self.request_path,
            "guiDebugLogPath": self.gui_debug_log_path,
            "reportPath": self.report_path,
            "statePath": self.state_path,
            "mode": self.mode,
            "status": self.status,
            "screenshots": self.screenshots,
            "createdAt": self.created_at,
            "remoteTarget": self.remote_target,
        }


@dataclass
class RepairAttempt:
    repair_id: str
    task_id: str
    attempt_index: int
    source_status: str
    source_result_path: str
    reason: str
    repair_brief_path: str
    state_path: str = ""
    worker_brief_path: str = ""
    debug_session_ids: List[str] = field(default_factory=list)
    status: str = "queued"
    created_at: str = ""

    @classmethod
    def from_dict(cls, payload: Dict[str, Any]) -> "RepairAttempt":
        return cls(
            repair_id=str(payload.get("repairId", "")).strip(),
            task_id=str(payload.get("taskId", "")).strip(),
            attempt_index=int(payload.get("attemptIndex", 0) or 0),
            source_status=str(payload.get("sourceStatus", "")).strip(),
            source_result_path=str(payload.get("sourceResultPath", "")).strip(),
            reason=str(payload.get("reason", "")).strip(),
            repair_brief_path=str(payload.get("repairBriefPath", "")).strip(),
            state_path=str(payload.get("statePath", "")).strip(),
            worker_brief_path=str(payload.get("workerBriefPath", "")).strip(),
            debug_session_ids=_ensure_str_list(payload.get("debugSessionIds")),
            status=str(payload.get("status", "queued")).strip(),
            created_at=str(payload.get("createdAt", "")).strip(),
        )

    def validate(self) -> None:
        if not self.repair_id:
            raise ValueError("repair attempt missing repairId")
        if not self.task_id:
            raise ValueError("repair attempt missing taskId")
        if self.attempt_index <= 0:
            raise ValueError("repair attempt must have positive attemptIndex")
        if not self.repair_brief_path:
            raise ValueError("repair attempt missing repairBriefPath")

    def to_dict(self) -> Dict[str, Any]:
        return {
            "repairId": self.repair_id,
            "taskId": self.task_id,
            "attemptIndex": self.attempt_index,
            "sourceStatus": self.source_status,
            "sourceResultPath": self.source_result_path,
            "reason": self.reason,
            "repairBriefPath": self.repair_brief_path,
            "statePath": self.state_path,
            "workerBriefPath": self.worker_brief_path,
            "debugSessionIds": self.debug_session_ids,
            "status": self.status,
            "createdAt": self.created_at,
        }


@dataclass
class DocumentUpdate:
    path: str
    action: str
    content: str = ""
    marker: str = ""
    append_lines: List[str] = field(default_factory=list)

    @classmethod
    def from_dict(cls, payload: Dict[str, Any]) -> "DocumentUpdate":
        return cls(
            path=str(payload.get("path", "")).strip(),
            action=str(payload.get("action", "")).strip(),
            content=str(payload.get("content", "")).rstrip(),
            marker=str(payload.get("marker", "")).strip(),
            append_lines=_ensure_str_list(payload.get("appendLines")),
        )

    def validate(self) -> None:
        if not self.path:
            raise ValueError("document update missing path")
        if self.action not in DOCUMENT_UPDATE_ACTIONS:
            raise ValueError(
                f"document update action must be one of {sorted(DOCUMENT_UPDATE_ACTIONS)}"
            )
        if self.action == "replace_block" and not self.marker:
            raise ValueError("replace_block document update requires marker")
        if self.action == "append_lines" and not self.append_lines:
            raise ValueError("append_lines document update requires appendLines")
        if self.action in {"replace_block", "create_file"} and not self.content:
            raise ValueError(f"{self.action} document update requires content")
        if self.action in {"replace_block", "create_file"} and "TODO" in self.content:
            raise ValueError(
                f"{self.action} document update for {self.path} still contains TODO placeholders"
            )
        if self.action == "append_lines" and any("TODO" in line for line in self.append_lines):
            raise ValueError(
                f"append_lines document update for {self.path} still contains TODO placeholders"
            )

    def to_dict(self) -> Dict[str, Any]:
        payload = {
            "path": self.path,
            "action": self.action,
        }
        if self.content:
            payload["content"] = self.content
        if self.marker:
            payload["marker"] = self.marker
        if self.append_lines:
            payload["appendLines"] = self.append_lines
        return payload


@dataclass
class ReviewConflict:
    task_ids: List[str]
    reason: str
    overlap_paths: List[str] = field(default_factory=list)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "taskIds": self.task_ids,
            "reason": self.reason,
            "overlapPaths": self.overlap_paths,
        }

    @classmethod
    def from_dict(cls, payload: Dict[str, Any]) -> "ReviewConflict":
        return cls(
            task_ids=_ensure_str_list(payload.get("taskIds")),
            reason=str(payload.get("reason", "")).strip(),
            overlap_paths=_ensure_str_list(payload.get("overlapPaths")),
        )


@dataclass
class ParallelGroup:
    name: str
    task_ids: List[str]
    reason: str

    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "taskIds": self.task_ids,
            "reason": self.reason,
        }

    @classmethod
    def from_dict(cls, payload: Dict[str, Any]) -> "ParallelGroup":
        return cls(
            name=str(payload.get("name", "")).strip(),
            task_ids=_ensure_str_list(payload.get("taskIds")),
            reason=str(payload.get("reason", "")).strip(),
        )


@dataclass
class ParallelReview:
    source_todo: str
    generated_at: str
    tasks_considered: List[TaskRecord]
    dependency_edges: List[Dict[str, str]]
    parallel_groups: List[ParallelGroup]
    conflicts: List[ReviewConflict]
    serialization_points: List[Dict[str, Any]]
    notes: List[str]

    def to_dict(self) -> Dict[str, Any]:
        return {
            "sourceTodo": self.source_todo,
            "generatedAt": self.generated_at,
            "tasksConsidered": [task.to_dict() for task in self.tasks_considered],
            "dependencyEdges": self.dependency_edges,
            "parallelGroups": [group.to_dict() for group in self.parallel_groups],
            "conflicts": [conflict.to_dict() for conflict in self.conflicts],
            "serializationPoints": self.serialization_points,
            "notes": self.notes,
        }

    @classmethod
    def from_dict(cls, payload: Dict[str, Any]) -> "ParallelReview":
        return cls(
            source_todo=str(payload.get("sourceTodo", "")).strip(),
            generated_at=str(payload.get("generatedAt", "")).strip(),
            tasks_considered=[
                TaskRecord.from_dict(item)
                for item in payload.get("tasksConsidered", [])
            ],
            dependency_edges=list(payload.get("dependencyEdges", [])),
            parallel_groups=[
                ParallelGroup.from_dict(item)
                for item in payload.get("parallelGroups", [])
            ],
            conflicts=[
                ReviewConflict.from_dict(item)
                for item in payload.get("conflicts", [])
            ],
            serialization_points=list(payload.get("serializationPoints", [])),
            notes=_ensure_str_list(payload.get("notes")),
        )