backend/agents/planner.py

"""
Planning Engine — decomposes a high-level goal into a dependency-checked
task DAG using the LLM. Real cycle detection and dependency validation,
not a placeholder linear list.
"""
from __future__ import annotations

import json
import uuid
from dataclasses import dataclass, field
from typing import Any

from .llm_gateway import LLMGateway

PLANNER_SYSTEM_PROMPT = (
    "You are a planning engine for a multi-agent coding system. Given a goal, "
    "break it into a minimal set of concrete subtasks. Each subtask must have "
    "a short id, a one-sentence description, a suggested agent role "
    "(e.g. 'coder', 'reviewer', 'tester', 'researcher'), and a list of subtask "
    "ids it depends on (may be empty). Respond with ONLY a JSON array, no prose."
)


@dataclass(slots=True)
class PlannedTask:
    id: str
    description: str
    agent_role: str
    depends_on: list[str] = field(default_factory=list)
    status: str = "pending"  # pending | ready | running | done | failed


class PlanningEngine:
    def __init__(self, gateway: LLMGateway) -> None:
        self.gateway = gateway

    async def plan(self, goal: str, context: str = "") -> list[PlannedTask]:
        messages = [
            {"role": "system", "content": PLANNER_SYSTEM_PROMPT},
            {"role": "user", "content": f"Context:\n{context}\n\nGoal: {goal}"},
        ]
        raw = await self.gateway.complete(messages)
        items = self._parse_tasks(raw)
        tasks = [
            PlannedTask(
                id=str(item.get("id") or uuid.uuid4()),
                description=item["description"],
                agent_role=item.get("agent_role", "coder"),
                depends_on=list(item.get("depends_on", [])),
            )
            for item in items
        ]
        self._validate_dag(tasks)
        return tasks

    @staticmethod
    def _parse_tasks(raw: str) -> list[dict[str, Any]]:
        text = raw.strip()
        start, end = text.find("["), text.rfind("]")
        if start == -1 or end == -1:
            raise ValueError(f"Planner did not return a JSON array:\n{raw[:500]}")
        return json.loads(text[start : end + 1])

    @staticmethod
    def _validate_dag(tasks: list[PlannedTask]) -> None:
        by_id = {t.id: t for t in tasks}
        for task in tasks:
            unknown = [d for d in task.depends_on if d not in by_id]
            if unknown:
                raise ValueError(f"Task '{task.id}' depends on unknown task(s): {unknown}")

        visited: dict[str, int] = {}  # 0=unvisited, 1=in-progress, 2=done

        def visit(task_id: str) -> None:
            state = visited.get(task_id, 0)
            if state == 1:
                raise ValueError(f"Cycle detected involving task '{task_id}'")
            if state == 2:
                return
            visited[task_id] = 1
            for dep in by_id[task_id].depends_on:
                visit(dep)
            visited[task_id] = 2

        for task in tasks:
            visit(task.id)

    @staticmethod
    def ready_tasks(tasks: list[PlannedTask]) -> list[PlannedTask]:
        done_ids = {t.id for t in tasks if t.status == "done"}
        return [
            t for t in tasks
            if t.status == "pending" and all(dep in done_ids for dep in t.depends_on)
        ]