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 .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

        # 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 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}