src/agents/planner/validator.py

"""PlannerValidator — checks a TaskList before it reaches the TaskRunner.

Runs the 8 checks from AGENT_ARCHITECTURE_CONTEXT_new.md §7.3. On failure it
raises `PlannerValidationError` with a message specific enough that the planner
can be re-prompted to self-correct (the retry loop lives in service.py).

Check #1 (Pydantic parse) is enforced at the structured-output boundary — by the
time a `TaskList` reaches here it has already parsed; this validator additionally
rejects structurally-invalid plans (duplicate ids, dangling edges, cycles).
"""

from __future__ import annotations

from pydantic import ValidationError

from src.middlewares.logging import get_logger

from ...catalog.models import Catalog
from ...query.ir.models import QueryIR
from ...query.ir.repair import IRRepairer
from ...query.ir.validator import IRValidationError, IRValidator
from ...tools.analytics.aggregation import SUPPORTED_AGGS
from .contracts import ToolRegistry
from .errors import PlannerValidationError
from .inputs import Constraints
from .schemas import PLACEHOLDER_RE, TaskList

logger = get_logger("ir_repair")

# Heuristic: a checkable success_criteria mentions a measurable signal.
_CHECKABLE_TOKENS = ("rate", "count", "match", "produced", "above", "below", "equal")

# Tool categories whose output is NOT analyzable data rows: `catalog.introspection`
# (check_data/check_knowledge → catalog metadata) and `retrieval.documents`
# (retrieve_knowledge → prose chunks). An analyze_* `data` handoff must not come from
# one of these — their output is kind="table"/documents so the structural checks pass,
# but feeding metadata rows to analyze_* makes it fail to find the requested columns.
# A denylist (not an allowlist) so any data/table-producing tool — retrieve_data AND a
# table-producing analyze_* chained into another analyze_* — stays valid.
_NON_DATA_SOURCE_CATEGORIES = frozenset({"catalog.introspection", "retrieval.documents"})

# DFS colors for cycle detection.
_WHITE, _GREY, _BLACK = 0, 1, 2


class PlannerValidator:
    def __init__(
        self,
        ir_validator: IRValidator | None = None,
        ir_repairer: IRRepairer | None = None,
    ) -> None:
        self._ir_validator = ir_validator or IRValidator()
        self._ir_repairer = ir_repairer or IRRepairer()

    def validate(
        self,
        task_list: TaskList,
        registry: ToolRegistry,
        catalog: Catalog,
        constraints: Constraints,
    ) -> None:
        tasks = task_list.tasks

        # Infeasible sentinel (planner.md "When the catalog cannot answer"): an
        # empty plan carrying `infeasible_reason` is a VALID outcome — the
        # coordinator renders it as an honest data-gap answer instead of the
        # planner force-mapping the question onto unrelated columns. A non-empty
        # plan keeps normal validation and the reason is ignored (a real plan
        # wins over a hedge).
        if task_list.infeasible_reason and not tasks:
            return

        # Check 6 — plan non-empty and within the task cap.
        if not tasks:
            raise PlannerValidationError("plan is empty: at least one task is required")
        if len(tasks) > constraints.max_tasks:
            raise PlannerValidationError(
                f"plan has {len(tasks)} tasks, exceeds max_tasks={constraints.max_tasks}"
            )

        ids = [t.id for t in tasks]
        if len(set(ids)) != len(ids):
            dupes = sorted({i for i in ids if ids.count(i) > 1})
            raise PlannerValidationError(f"duplicate task id(s): {dupes}")
        id_set = set(ids)
        tasks_by_id = {t.id: t for t in tasks}

        known_tools = registry.names()
        known_sources = {s.source_id for s in catalog.sources}

        for task in tasks:
            for call in task.tool_calls:
                # Check 2 — every tool exists in the registry.
                if call.tool not in known_tools:
                    raise PlannerValidationError(
                        f"task {task.id}: tool {call.tool!r} not in registry "
                        f"(known: {sorted(known_tools)})"
                    )
                spec = registry.get(call.tool)
                assert spec is not None  # guaranteed by the membership check above

                # Check 8a — args carry the required keys and no unknown keys.
                required = set(spec.input_schema.get("required", []))
                allowed = set(spec.input_schema.get("properties", {}).keys()) | required
                missing = required - set(call.args.keys())
                if missing:
                    raise PlannerValidationError(
                        f"task {task.id}: tool {call.tool!r} missing required arg(s): "
                        f"{sorted(missing)}"
                    )
                unknown = set(call.args.keys()) - allowed
                if unknown:
                    raise PlannerValidationError(
                        f"task {task.id}: tool {call.tool!r} has unknown arg(s): "
                        f"{sorted(unknown)} (allowed: {sorted(allowed)})"
                    )

                # Check 8c — analyze_aggregate: every aggregation FUNCTION must be one
                # the tool supports. Check 8a only validates arg *names* (`aggregations`
                # is allowed); it never looks at the function *values* inside the dict,
                # so an unsupported func like `std` otherwise passes validation and only
                # fails at execution — too late for a corrective retry, so the task
                # reaches the Assembler as a silent failure. Catch it here so the planner
                # is re-prompted to degrade to a supported function (e.g. `mean`).
                if call.tool == "analyze_aggregate":
                    aggs = call.args.get("aggregations")
                    if isinstance(aggs, dict):
                        bad = sorted(
                            {
                                f
                                for funcs in aggs.values()
                                for f in ([funcs] if isinstance(funcs, str) else funcs or [])
                                if f not in SUPPORTED_AGGS
                            }
                        )
                        if bad:
                            raise PlannerValidationError(
                                f"task {task.id}: analyze_aggregate has unsupported "
                                f"aggregation function(s) {bad} (supported: "
                                f"{sorted(SUPPORTED_AGGS)}). Use a supported function "
                                "(e.g. mean/median); for the spread of a whole column "
                                "use analyze_descriptive instead."
                            )

                # Check 3 — concrete source_id args must exist in the catalog.
                src = call.args.get("source_id")
                if isinstance(src, str) and not _is_placeholder(src):
                    if src not in known_sources:
                        raise PlannerValidationError(
                            f"task {task.id}: tool {call.tool!r} references unknown "
                            f"source_id {src!r} (known: {sorted(known_sources)})"
                        )

                # Check 8b — inline retrieve_data IR validates against the catalog.
                if call.tool == "retrieve_data":
                    self._validate_inline_ir(task.id, call.args, catalog)

                # Check 9 — a `data` handoff (Pattern A) must reference a task that
                # produces analyzable data rows, not one producing catalog metadata
                # (check_data/check_knowledge) or documents (retrieve_knowledge).
                self._validate_data_source(task.id, call, tasks_by_id, registry)

            # Check 7 — success_criteria is checkable.
            if not _is_checkable(task.success_criteria):
                raise PlannerValidationError(
                    f"task {task.id}: success_criteria is not checkable — include a "
                    f"measurable signal (one of {list(_CHECKABLE_TOKENS)}); "
                    f"got {task.success_criteria!r}"
                )

        # Check 4 — DAG: edges resolve, placeholders resolve, no cycles.
        self._validate_dag(tasks_by_id, id_set)

    def _validate_inline_ir(self, task_id: str, args: dict, catalog: Catalog) -> None:
        raw_ir = args.get("ir")
        if not isinstance(raw_ir, dict):
            raise PlannerValidationError(
                f"task {task_id}: retrieve_data.args.ir must be an inline QueryIR "
                f"object, got {type(raw_ir).__name__}"
            )
        try:
            ir = QueryIR.model_validate(raw_ir)
        except ValidationError as e:
            raise PlannerValidationError(
                f"task {task_id}: retrieve_data.args.ir is not a valid QueryIR: {e}"
            ) from e

        # Canonicalize near-miss ids (LLM dropped/mutated a char in an opaque
        # catalog id) before validating. On a successful repair, write the fixed
        # IR back into the tool call so the downstream executor runs the
        # corrected IR — not just the validator.
        ir, repairs = self._ir_repairer.repair(ir, catalog)
        if repairs:
            args["ir"] = ir.model_dump()
            for r in repairs:
                logger.info(
                    "repaired ir id",
                    task_id=task_id,
                    where=r.where,
                    from_id=r.from_id,
                    to_id=r.to_id,
                )

        try:
            self._ir_validator.validate(ir, catalog)
        except IRValidationError as e:
            raise PlannerValidationError(
                f"task {task_id}: retrieve_data IR failed catalog validation: {e}"
            ) from e

    @staticmethod
    def _validate_data_source(
        task_id: str, call, tasks_by_id: dict, registry: ToolRegistry
    ) -> None:
        """A `data` placeholder must reference a data-producing task, not a
        metadata (check_data/check_knowledge) or documents (retrieve_knowledge) one.

        Those pass the structural checks (check_* also returns kind="table"), but
        their rows are catalog schema, so a downstream analyze_* fails to find the
        requested columns. Resolving points at the referenced task's representative
        output — its last tool call (matches TaskRunner's `outputs[-1]`).
        """
        data_arg = call.args.get("data")
        if not isinstance(data_arg, str):
            return
        match = PLACEHOLDER_RE.fullmatch(data_arg.strip())
        if not match:
            return
        ref_task = tasks_by_id.get(match.group(1))
        if ref_task is None or not ref_task.tool_calls:
            return  # a dangling placeholder is reported by the DAG check
        ref_tool = ref_task.tool_calls[-1].tool
        ref_spec = registry.get(ref_tool)
        if ref_spec is not None and ref_spec.category in _NON_DATA_SOURCE_CATEGORIES:
            raise PlannerValidationError(
                f"task {task_id}: tool {call.tool!r} takes its 'data' from task "
                f"{match.group(1)} ({ref_tool!r}, category {ref_spec.category!r}), "
                "which produces metadata/documents — not analyzable data rows. Feed "
                "analyze_* from a data-producing tool (e.g. retrieve_data)."
            )

    @staticmethod
    def _validate_dag(tasks_by_id: dict, id_set: set[str]) -> None:
        for task in tasks_by_id.values():
            for dep in task.depends_on:
                if dep not in id_set:
                    raise PlannerValidationError(
                        f"task {task.id}: depends_on references unknown task {dep!r}"
                    )
                if dep == task.id:
                    raise PlannerValidationError(
                        f"task {task.id}: depends_on includes itself"
                    )

        cycle = _find_cycle(tasks_by_id)
        if cycle:
            raise PlannerValidationError(f"cycle detected in depends_on: {' -> '.join(cycle)}")

        # On an acyclic graph, a placeholder is safe iff its target is a
        # transitive ancestor — i.e. guaranteed to have completed before this
        # task runs. Requiring a *direct* depends_on would wrongly reject valid
        # plans that depend on the target through an intermediate task.
        ancestors = _all_ancestors(tasks_by_id)
        for task in tasks_by_id.values():
            for ref in _placeholder_refs(task):
                if ref not in id_set:
                    raise PlannerValidationError(
                        f"task {task.id}: placeholder '${{{ref}}}' references unknown task"
                    )
                if ref not in ancestors[task.id]:
                    raise PlannerValidationError(
                        f"task {task.id}: placeholder '${{{ref}}}' used but {ref!r} is "
                        f"not a (transitive) dependency — add it to depends_on"
                    )


def _is_placeholder(value: str) -> bool:
    return bool(PLACEHOLDER_RE.fullmatch(value.strip()))


def _placeholder_refs(task) -> set[str]:
    refs: set[str] = set()
    for call in task.tool_calls:
        for value in call.args.values():
            if isinstance(value, str):
                refs.update(PLACEHOLDER_RE.findall(value))
    return refs


def _is_checkable(text: str) -> bool:
    low = text.lower()
    return any(tok in low for tok in _CHECKABLE_TOKENS)


def _find_cycle(tasks_by_id: dict) -> list[str] | None:
    color = {tid: _WHITE for tid in tasks_by_id}
    stack: list[str] = []

    def dfs(node: str) -> list[str] | None:
        color[node] = _GREY
        stack.append(node)
        for dep in tasks_by_id[node].depends_on:
            if color.get(dep) == _GREY:
                idx = stack.index(dep)
                return stack[idx:] + [dep]
            if color.get(dep) == _WHITE:
                found = dfs(dep)
                if found:
                    return found
        stack.pop()
        color[node] = _BLACK
        return None

    for tid in tasks_by_id:
        if color[tid] == _WHITE:
            found = dfs(tid)
            if found:
                return found
    return None


def _all_ancestors(tasks_by_id: dict) -> dict[str, set[str]]:
    """ancestors[id] = all tasks reachable by following depends_on edges."""
    cache: dict[str, set[str]] = {}

    def visit(node: str, seen: set[str]) -> set[str]:
        if node in cache:
            return cache[node]
        acc: set[str] = set()
        for dep in tasks_by_id[node].depends_on:
            if dep in seen or dep not in tasks_by_id:
                continue
            acc.add(dep)
            acc |= visit(dep, seen | {dep})
        cache[node] = acc
        return acc

    return {tid: visit(tid, {tid}) for tid in tasks_by_id}