backend/agents/orchestrator.py

"""
Multi-Agent Orchestrator — task delegation and concurrent multi-agent
execution. Runs a planner-produced task DAG across a pool of agents,
respecting dependencies and a concurrency cap, and publishes real
status events on the message bus as work progresses.
"""
from __future__ import annotations

import asyncio
import logging
from dataclasses import dataclass
from typing import Callable, Optional

from .base import BaseAgent
from .llm_gateway import LLMGateway
from .memory import AgentMemoryStore
from .messaging import MessageBus
from .planner import PlannedTask, PlanningEngine

logger = logging.getLogger("dolor3v.agents.orchestrator")

AgentFactory = Callable[[str, str], BaseAgent]  # (agent_id, role) -> BaseAgent


@dataclass(slots=True)
class TaskResult:
    task_id: str
    agent_id: str
    output: str
    error: Optional[str] = None


class Orchestrator:
    """Runs a plan across a pool of agents, respecting the dependency DAG."""

    def __init__(
        self,
        gateway: LLMGateway,
        memory: AgentMemoryStore,
        bus: MessageBus,
        agent_factory: AgentFactory,
        max_concurrent_agents: int = 3,
    ) -> None:
        self.gateway = gateway
        self.memory = memory
        self.bus = bus
        self.agent_factory = agent_factory
        self.planner = PlanningEngine(gateway)
        self._semaphore = asyncio.Semaphore(max_concurrent_agents)

    async def run_goal(self, goal: str, context: str = "") -> dict[str, TaskResult]:
        tasks = await self.planner.plan(goal, context=context)
        results: dict[str, TaskResult] = {}

        while any(t.status in ("pending", "ready", "running") for t in tasks):
            ready = self.planner.ready_tasks(tasks)
            if not ready:
                stuck = [t.id for t in tasks if t.status == "pending"]
                if stuck:
                    raise RuntimeError(f"Plan stalled; unresolvable dependency for tasks: {stuck}")
                break

            for task in ready:
                task.status = "running"

            batch = await asyncio.gather(
                *(self._run_task(task) for task in ready),
                return_exceptions=True,
            )
            for task, outcome in zip(ready, batch):
                if isinstance(outcome, BaseException):
                    task.status = "failed"
                    results[task.id] = TaskResult(task.id, agent_id="", output="", error=str(outcome))
                    logger.error("Task %s failed: %s", task.id, outcome)
                else:
                    task.status = "done"
                    results[task.id] = outcome

        return results

    async def _run_task(self, task: PlannedTask) -> TaskResult:
        async with self._semaphore:
            agent_id = f"{task.agent_role}-{task.id}"
            agent = self.agent_factory(agent_id, task.agent_role)
            self.bus.register(agent_id)
            await self.bus.publish(
                sender="orchestrator", topic="task-status", type="started",
                payload={"task_id": task.id, "agent_id": agent_id},
            )
            try:
                output = await agent.run(task.description, task_id=task.id)
                await self.bus.publish(
                    sender="orchestrator", topic="task-status", type="completed",
                    payload={"task_id": task.id, "agent_id": agent_id},
                )
                return TaskResult(task_id=task.id, agent_id=agent_id, output=output)
            finally:
                self.bus.unregister(agent_id)