script-build-and-framework-tmp/agent/agent/orchestration/_task_graph.py

"""Pure task-tree operations used by :mod:`agent.orchestration.coordinator`.

The task ledger remains owned and committed by ``TaskCoordinator``.  This
module only mutates the in-memory aggregate handed to it; it never performs
I/O, acquires locks, or commits state on its own.
"""

from __future__ import annotations

from dataclasses import asdict
from typing import Any, Dict, List, Optional, Sequence, Tuple

from .errors import OrchestrationError, TaskConflict
from .models import (
    DecisionAction,
    PlannerDecision,
    TaskAttempt,
    TaskLedger,
    TaskRecord,
    TaskSpec,
    TaskStatus,
    ValidationReport,
    ValidationRunStatus,
    ValidationVerdict,
    json_values,
    new_id,
    utc_now,
)
from .state_machine import transition


TERMINAL_TASK_STATUSES = {
    TaskStatus.COMPLETED,
    TaskStatus.CANCELLED,
    TaskStatus.SUPERSEDED,
}

ACTIVE_TASK_STATUSES = {
    TaskStatus.RUNNING,
    TaskStatus.AWAITING_VALIDATION,
    TaskStatus.VALIDATING,
}


def path_key(display_path: str) -> Tuple[int, ...]:
    """Return a stable numeric sort key for a dotted task display path."""

    try:
        return tuple(int(part) for part in display_path.split("."))
    except ValueError:
        return (10**9,)


class TaskGraph:
    """Maintain task hierarchy and accepted child-result bindings.

    All methods operate on a caller-owned ``TaskLedger``.  Keeping persistence
    out of this component preserves ``TaskCoordinator`` as the sole writer.
    """

    @staticmethod
    def accepted_child_results(
        ledger: TaskLedger,
        task: TaskRecord,
        attempt: Optional[TaskAttempt] = None,
    ) -> List[Dict[str, Any]]:
        """Build immutable, accepted direct-child inputs for a parent Worker."""

        results: List[Dict[str, Any]] = []
        if attempt is None:
            decision_ids = TaskGraph.accepted_child_decision_ids(ledger, task)
        elif attempt.accepted_child_decision_ids is None:
            raise OrchestrationError(
                "Attempt has no frozen accepted child decision binding"
            )
        else:
            decision_ids = attempt.accepted_child_decision_ids
        if len(set(decision_ids)) != len(decision_ids):
            raise OrchestrationError("Attempt child decision bindings are not unique")
        expected_decision_ids = TaskGraph.accepted_child_decision_ids(ledger, task)
        if tuple(decision_ids) != expected_decision_ids:
            raise OrchestrationError(
                "Attempt child decision bindings do not match direct children in "
                "display_path order"
            )
        direct_children = set(task.child_task_ids)
        prior_path: Optional[Tuple[int, ...]] = None
        for decision_id in decision_ids:
            decision = ledger.decisions.get(decision_id)
            if (
                decision is None
                or decision.task_id not in direct_children
                or decision.action != DecisionAction.ACCEPT
            ):
                raise OrchestrationError(
                    "Attempt child decision binding is not a direct-child ACCEPT"
                )
            child = ledger.tasks[decision.task_id]
            child_path = path_key(child.display_path)
            if prior_path is not None and child_path < prior_path:
                raise OrchestrationError(
                    "Attempt child decision bindings are not in display_path order"
                )
            prior_path = child_path
            bound_decision, child_attempt, validation = TaskGraph.accepted_result_binding(
                ledger, child
            )
            if bound_decision.decision_id != decision_id:
                raise OrchestrationError(
                    "Attempt child decision binding is not the child's accepted result"
                )
            results.append(
                {
                    "task_id": child.task_id,
                    "task_spec": json_values(asdict(child.current_spec)),
                    "attempt_id": child_attempt.attempt_id,
                    "snapshot_id": child_attempt.snapshot_id,
                    "submission": json_values(asdict(child_attempt.submission)),
                    "validation": {
                        "validation_id": validation.validation_id,
                        "summary": validation.summary,
                        "evidence_refs": json_values(
                            asdict(validation)["evidence_refs"]
                        ),
                    },
                }
            )
        return results

    @staticmethod
    def accepted_child_decision_ids(
        ledger: TaskLedger,
        task: TaskRecord,
    ) -> Tuple[str, ...]:
        decision_ids: List[str] = []
        children = sorted(
            (ledger.tasks[child_id] for child_id in task.child_task_ids),
            key=lambda child: path_key(child.display_path),
        )
        for child in children:
            if child.status == TaskStatus.COMPLETED:
                decision, _attempt, _validation = TaskGraph.accepted_result_binding(
                    ledger, child
                )
                decision_ids.append(decision.decision_id)
        return tuple(decision_ids)

    @staticmethod
    def root_task(ledger: TaskLedger) -> TaskRecord:
        if not ledger.root_task_id:
            raise TaskConflict(
                "Rootless task ledger is incompatible with mission execution; "
                "rebuild the development trace"
            )
        root = ledger.tasks.get(ledger.root_task_id)
        top_level = [
            task.task_id
            for task in ledger.tasks.values()
            if task.parent_task_id is None
        ]
        if not root or root.parent_task_id is not None or top_level != [root.task_id]:
            raise OrchestrationError("Task ledger Root Task structure is invalid")
        return root

    @staticmethod
    def accepted_result_binding(
        ledger: TaskLedger,
        task: TaskRecord,
    ) -> Tuple[PlannerDecision, TaskAttempt, ValidationReport]:
        if task.status != TaskStatus.COMPLETED or not task.decision_ids:
            raise OrchestrationError(f"Task {task.task_id} has no accepted result")
        decision = ledger.decisions.get(task.decision_ids[-1])
        if (
            not decision
            or decision.task_id != task.task_id
            or decision.action != DecisionAction.ACCEPT
            or not decision.attempt_id
            or not decision.validation_id
        ):
            raise OrchestrationError(
                f"Task {task.task_id} has an invalid accept decision"
            )
        attempt = ledger.attempts.get(decision.attempt_id)
        validation = ledger.validations.get(decision.validation_id)
        if (
            not attempt
            or attempt.task_id != task.task_id
            or not attempt.submission
            or not attempt.snapshot_id
            or attempt.spec_version != task.current_spec_version
            or not validation
            or validation.task_id != task.task_id
            or validation.attempt_id != attempt.attempt_id
            or validation.snapshot_id != attempt.snapshot_id
            or validation.spec_version != task.current_spec_version
            or validation.status != ValidationRunStatus.COMPLETED
            or validation.verdict != ValidationVerdict.PASSED
        ):
            raise OrchestrationError(
                f"Task {task.task_id} has an invalid accepted result binding"
            )
        return decision, attempt, validation

    @staticmethod
    def nonterminal_descendants(
        ledger: TaskLedger,
        task: TaskRecord,
    ) -> List[TaskRecord]:
        descendants: List[TaskRecord] = []
        pending_ids = list(task.child_task_ids)
        while pending_ids:
            child = ledger.tasks[pending_ids.pop()]
            pending_ids.extend(child.child_task_ids)
            if child.status not in TERMINAL_TASK_STATUSES:
                descendants.append(child)
        return descendants

    def create_child_records(
        self,
        ledger: TaskLedger,
        parent: TaskRecord,
        drafts: Sequence[Dict[str, Any]],
    ) -> List[str]:
        return self.create_records(ledger, drafts, parent.task_id, parent.display_path)

    def create_sibling_records(
        self,
        ledger: TaskLedger,
        task: TaskRecord,
        drafts: Sequence[Dict[str, Any]],
    ) -> List[str]:
        parent_path = (
            ledger.tasks[task.parent_task_id].display_path
            if task.parent_task_id
            else ""
        )
        return self.create_records(ledger, drafts, task.parent_task_id, parent_path)

    @staticmethod
    def add_task_record(
        ledger: TaskLedger,
        draft: Dict[str, Any],
        parent_task_id: Optional[str],
        display_path: str,
    ) -> str:
        task_id = new_id()
        spec = TaskGraph.task_spec_from_draft(draft, version=1)
        ledger.tasks[task_id] = TaskRecord(
            task_id=task_id,
            goal_id=None,
            parent_task_id=parent_task_id,
            display_path=display_path,
            specs=[spec],
        )
        return task_id

    @staticmethod
    def task_spec_from_draft(draft: Dict[str, Any], *, version: int) -> TaskSpec:
        if not isinstance(draft, dict):
            raise ValueError("Task draft must be an object")
        return TaskSpec.from_dict(
            {
                "version": version,
                "objective": draft.get("objective", ""),
                "acceptance_criteria": draft.get("acceptance_criteria", []),
                "context_refs": draft.get("context_refs", ()),
            }
        )

    def create_records(
        self,
        ledger: TaskLedger,
        drafts: Sequence[Dict[str, Any]],
        parent_task_id: Optional[str],
        parent_path: str,
    ) -> List[str]:
        siblings = sorted(
            (
                task
                for task in ledger.tasks.values()
                if task.parent_task_id == parent_task_id
            ),
            key=lambda item: path_key(item.display_path),
        )
        created: List[str] = []
        for offset, draft in enumerate(drafts, start=1):
            index = len(siblings) + offset
            display_path = f"{parent_path}.{index}" if parent_path else str(index)
            task_id = self.add_task_record(
                ledger,
                draft,
                parent_task_id,
                display_path,
            )
            if parent_task_id:
                ledger.tasks[parent_task_id].child_task_ids.append(task_id)
            created.append(task_id)
        return created

    @staticmethod
    def mark_parent_waiting(parent: TaskRecord) -> None:
        if parent.status == TaskStatus.WAITING_CHILDREN:
            return
        transition(parent.status, TaskStatus.WAITING_CHILDREN)
        parent.status = TaskStatus.WAITING_CHILDREN
        parent.updated_at = utc_now()

    @staticmethod
    def update_parent_after_child(ledger: TaskLedger, child: TaskRecord) -> None:
        if not child.parent_task_id or child.status not in TERMINAL_TASK_STATUSES:
            return
        parent = ledger.tasks[child.parent_task_id]
        if parent.status != TaskStatus.WAITING_CHILDREN:
            return
        children = [ledger.tasks[child_id] for child_id in parent.child_task_ids]
        if children and all(
            candidate.status in TERMINAL_TASK_STATUSES for candidate in children
        ):
            transition(parent.status, TaskStatus.NEEDS_REPLAN)
            parent.status = TaskStatus.NEEDS_REPLAN
        parent.updated_at = utc_now()

    @staticmethod
    def display_path(
        ledger: TaskLedger,
        parent_task_id: Optional[str],
        index: int,
    ) -> str:
        if not parent_task_id:
            return str(index)
        return f"{ledger.tasks[parent_task_id].display_path}.{index}"

    @staticmethod
    def rebase_task_path(
        ledger: TaskLedger,
        task: TaskRecord,
        display_path: str,
    ) -> None:
        task.display_path = display_path
        for index, child_id in enumerate(task.child_task_ids, start=1):
            TaskGraph.rebase_task_path(
                ledger,
                ledger.tasks[child_id],
                f"{display_path}.{index}",
            )


__all__ = [
    "ACTIVE_TASK_STATUSES",
    "TERMINAL_TASK_STATUSES",
    "TaskGraph",
    "path_key",
]