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

"""Best-effort TaskLedger-to-GoalTree compatibility projection."""

from __future__ import annotations

import logging
from typing import Any, Awaitable, Callable, Dict, Optional

from agent.trace.goal_models import GoalTree

from .models import TaskLedger, TaskRecord, TaskStatus
from .protocols import TaskStore


logger = logging.getLogger(__name__)


GOAL_STATUS_BY_TASK_STATUS = {
    TaskStatus.PENDING: "pending",
    TaskStatus.RUNNING: "in_progress",
    TaskStatus.AWAITING_VALIDATION: "in_progress",
    TaskStatus.VALIDATING: "in_progress",
    TaskStatus.AWAITING_DECISION: "in_progress",
    TaskStatus.NEEDS_REPLAN: "in_progress",
    TaskStatus.WAITING_CHILDREN: "in_progress",
    TaskStatus.BLOCKED: "in_progress",
    TaskStatus.COMPLETED: "completed",
    TaskStatus.CANCELLED: "abandoned",
    TaskStatus.SUPERSEDED: "abandoned",
}


class GoalProjection:
    """Project authoritative task state into the legacy GoalTree view.

    The projection never commits a ledger directly. Linking a projected Goal
    back to its Task is delegated to the coordinator's mutation callback so the
    coordinator remains the single persistence boundary.
    """

    def __init__(
        self,
        *,
        task_store: TaskStore,
        get_trace_store: Callable[[], Any],
        mutate: Callable[..., Awaitable[Dict[str, Any]]],
    ) -> None:
        self._task_store = task_store
        self._get_trace_store = get_trace_store
        self._mutate = mutate

    async def project_state(self, root_trace_id: str, task_id: str) -> None:
        trace_store = self._get_trace_store()
        if not trace_store:
            return
        ledger = await self._task_store.load(root_trace_id)
        task = _task(ledger, task_id)
        if not task.goal_id:
            return
        summary = None
        if task.decision_ids:
            summary = ledger.decisions[task.decision_ids[-1]].reason
        await trace_store.update_goal(
            root_trace_id,
            task.goal_id,
            cascade_completion=False,
            status=GOAL_STATUS_BY_TASK_STATUS[task.status],
            summary=summary,
        )

    async def project_compatibility(
        self,
        root_trace_id: str,
        task_id: str,
        *,
        ensure: bool = False,
        after_task_id: Optional[str] = None,
    ) -> None:
        """Run projection without invalidating an authoritative Ledger commit."""

        if not self._get_trace_store():
            return
        try:
            if ensure:
                await self.ensure(
                    root_trace_id,
                    task_id,
                    after_task_id=after_task_id,
                )
            await self.project_state(root_trace_id, task_id)
        except Exception as exc:
            self.log_failure(root_trace_id, task_id, exc)

    @staticmethod
    def log_failure(
        root_trace_id: str,
        task_id: str,
        error: Exception,
    ) -> None:
        logger.warning(
            "Goal compatibility projection failed after Ledger commit "
            "(%s, %s): %s",
            root_trace_id,
            task_id,
            error,
        )

    async def ensure(
        self,
        root_trace_id: str,
        task_id: str,
        after_task_id: Optional[str] = None,
    ) -> None:
        trace_store = self._get_trace_store()
        if not trace_store:
            return
        ledger = await self._task_store.load(root_trace_id)
        task = _task(ledger, task_id)
        tree = await trace_store.get_goal_tree(root_trace_id)
        if tree is None:
            tree = GoalTree(mission=ledger.root_objective)
        if task.goal_id and tree.find(task.goal_id):
            return
        parent_goal_id = None
        if task.parent_task_id:
            parent = ledger.tasks[task.parent_task_id]
            if not parent.goal_id:
                await self.ensure(root_trace_id, parent.task_id)
                ledger = await self._task_store.load(root_trace_id)
                parent = ledger.tasks[task.parent_task_id]
                tree = await trace_store.get_goal_tree(root_trace_id) or tree
            parent_goal_id = parent.goal_id
        after_goal_id = None
        if after_task_id and after_task_id in ledger.tasks:
            after_goal_id = ledger.tasks[after_task_id].goal_id
        if after_goal_id and tree.find(after_goal_id):
            goal = tree.add_goals_after(
                after_goal_id,
                descriptions=[task.current_spec.objective],
                reasons=["TaskLedger compatibility projection"],
            )[0]
        else:
            goal = tree.add_goals(
                descriptions=[task.current_spec.objective],
                reasons=["TaskLedger compatibility projection"],
                parent_id=parent_goal_id,
            )[0]
        await trace_store.update_goal_tree(root_trace_id, tree)

        def link(current: TaskLedger) -> Dict[str, Any]:
            current_task = _task(current, task_id)
            if current_task.goal_id is None:
                current_task.goal_id = goal.id
            return {"task_id": task_id, "goal_id": current_task.goal_id}

        await self._mutate(root_trace_id, "goal_projection_linked", link)

    async def reconcile(self, root_trace_id: str) -> Dict[str, Any]:
        ledger = await self._task_store.load(root_trace_id)
        for task_id in ledger.tasks:
            await self.ensure(root_trace_id, task_id)
            await self.project_state(root_trace_id, task_id)
        return {
            "root_trace_id": root_trace_id,
            "reconciled_tasks": len(ledger.tasks),
        }


def _task(ledger: TaskLedger, task_id: str) -> TaskRecord:
    try:
        return ledger.tasks[task_id]
    except KeyError as exc:
        raise ValueError(f"Task not found: {task_id}") from exc


__all__ = ["GOAL_STATUS_BY_TASK_STATUS", "GoalProjection"]