orchestrata_ultimate.agents.task_executor

Reference › orchestrata_ultimate › agents › task_executor

task_executor is the point where a TaskDefinition meets an actual volk. TaskExecutor is built from a VolkPool (pool), a VolkRuntime (runtime), a MultiagentPrologBridge (prolog_bridge), an optional Redis state manager and an optional Postman handle. The coordinator, intent_adapter and postman_adapter all call it.

submit registers the task as pending, resolves the target alias with _resolve_alias, acquires a pool instance and marks it busy (_acquire_instance), then runs _run_and_complete. _resolve_alias uses the task's target and, when that is empty, asks Prolog for the best volk, raising TaskAssignmentError if there is none. _execute_task looks the alias up in a module-level class map filled by _init_volk_map, instantiates the volk, converts the task to a CXP packet with _task_to_packet, and awaits the volk's process method. An unknown alias or an exception produces a failed TaskResult with the error message; a response packet becomes the result output and is forwarded to the next volk through Postman by _route_response. cancel, get_task, get_task_status, get_result and list_tasks serve queries, and pending_count, executing_count and completed_count are properties.

At a glance

Class diagram

Dependency diagram

API

Constants

Classes

TaskExecutor

Line 94

Executes tasks on Volk instances.

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/agents/task_executor.py

# mypy: ignore-errors
"""Task Executor — submit and track task execution on Volk instances.

Coordinates task assignment via Prolog governance, instance acquisition
via VolkPool, execution tracking, and result collection.

Author: Mohamed
Spec: 005-multi-agent (amended)
"""

# Spec: 005 - Multi-Agent System

from __future__ import annotations

import logging
import time
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from orchestrata_ultimate.infra.redis_state import RedisStateManager

from orchestrata_ultimate.agents.exceptions import (
    TaskAssignmentError,
    TaskCancellationError,
    TaskExecutionError,
)
from orchestrata_ultimate.agents.pool import VolkPool
from orchestrata_ultimate.agents.prolog_bridge import MultiagentPrologBridge
from orchestrata_ultimate.agents.runtime import VolkRuntime
from orchestrata_ultimate.agents.types import (
    TaskDefinition,
    TaskResult,
    TaskStatus,
    VolkInstance,
)
from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.core.packet import CXPPacket

_logger = logging.getLogger(__name__)


def _utc_now() -> datetime:
    return datetime.now(UTC)


def _task_to_packet(task: TaskDefinition) -> CXPPacket:
    """Convert a TaskDefinition to a CXPPacket."""
    from orchestrata_ultimate.core.packet import CXPPacket, WorkStatus

    metadata = task.parameters or {}
    return CXPPacket(
        from_agent=metadata.get("from_agent", "SYSTEM"),
        to_agent=task.target,
        thread_id=metadata.get("thread_id", ""),
        subject=task.action,
        payload=metadata,
        work_status=WorkStatus.PENDING,
    )


_VOLK_CLASS_MAP: dict = {}


def _init_volk_map() -> None:
    """Initialize the Volk class mapping."""
    volk_aliases = [
        ("U:BA", "orchestrata_ultimate.agents.volks.analyst.agent", "AnalystAgent"),
        ("U:Ar", "orchestrata_ultimate.agents.volks.architect.volk", "ArchitectVolk"),
        (
            "U:Cd",
            "orchestrata_ultimate.agents.volks.codedirector.volk",
            "CodeDirectorVolk",
        ),
        ("U:Cp", "orchestrata_ultimate.agents.volks.compliance.volk", "ComplianceVolk"),
        ("U:Cx", "orchestrata_ultimate.agents.volks.concierge.volk", "ConciergeVolk"),
        ("J:Junior", "orchestrata_ultimate.agents.volks.junior.volk", "JuniorVolk"),
        ("U:PM", "orchestrata_ultimate.agents.volks.pm.volk", "PMVolk"),
        ("U:Qa", "orchestrata_ultimate.agents.volks.qa.volk", "QaVolk"),
        ("U:Ux", "orchestrata_ultimate.agents.volks.ux.volk", "UXVolk"),
    ]
    for alias, module_path, class_name in volk_aliases:
        try:
            module = __import__(module_path, fromlist=[class_name])
            volk_class = getattr(module, class_name)
            _VOLK_CLASS_MAP[alias] = volk_class
        except ImportError:
            pass


_init_volk_map()


class TaskExecutor:
    """Executes tasks on Volk instances.

    Parameters
    ----------
    pool : VolkPool
        Pool for acquiring/releasing instances.
    runtime : VolkRuntime
        Runtime for state management.
    prolog : MultiagentPrologBridge
        Bridge for governance queries.
    redis_state : RedisStateManager, optional
        Redis state manager for task tracking.
    """

    def __init__(
        self,
        pool: VolkPool,
        runtime: VolkRuntime,
        prolog: MultiagentPrologBridge,
        redis_state: RedisStateManager | None = None,
        postman: Any | None = None,
    ) -> None:
        self._pool = pool
        self._runtime = runtime
        self._prolog = prolog
        self._redis = redis_state
        self._postman = postman
        self._tasks: dict[str, TaskDefinition] = {}
        self._task_status: dict[str, TaskStatus] = {}
        self._task_assignments: dict[str, str] = {}
        self._results: dict[str, TaskResult] = {}

    async def submit(self, task: TaskDefinition) -> TaskResult:
        """Submit a task for execution.

        Raises
        ------
        TaskAssignmentError
            If no suitable Volk is found or governance denies assignment.
        TaskExecutionError
            If execution fails.
        """
        start_time = time.monotonic()
        self._register_task(task)

        alias = self._resolve_alias(task)
        instance = await self._acquire_instance(task, alias)

        result = await self._run_and_complete(task, instance, start_time)
        return result

    def _register_task(self, task: TaskDefinition) -> None:
        """Register the task in local state and mark it PENDING."""
        self._tasks[task.task_id] = task
        self._task_status[task.task_id] = TaskStatus.PENDING

    async def _acquire_instance(self, task: TaskDefinition, alias: str) -> VolkInstance:
        """Acquire a pool instance and mark it busy for the task.

        Raises
        ------
        TaskAssignmentError
            If pool.acquire() fails.
        """
        try:
            instance = await self._pool.acquire(alias, task)
        except Exception as e:
            self._task_status[task.task_id] = TaskStatus.FAILED
            raise TaskAssignmentError(
                f"Failed to acquire instance for task '{task.task_id}': {e}",
                task_id=task.task_id,
            )

        self._runtime.mark_busy(instance.instance_id, task.task_id)
        self._task_assignments[task.task_id] = instance.instance_id
        self._task_status[task.task_id] = TaskStatus.EXECUTING
        return instance

    async def _run_and_complete(
        self, task: TaskDefinition, instance: VolkInstance, start_time: float
    ) -> TaskResult:
        """Execute the task, route response to next volk, finalise result.

        Raises
        ------
        TaskExecutionError
            If _execute_task() raises.
        """
        response_packet = None
        try:
            result, response_packet = await self._execute_task(
                task, instance.instance_id
            )
        except Exception as e:
            self._task_status[task.task_id] = TaskStatus.FAILED
            await self._pool.release(instance.instance_id)
            raise TaskExecutionError(
                f"Task '{task.task_id}' execution failed: {e}",
                task_id=task.task_id,
                instance_id=instance.instance_id,
            )

        if response_packet is not None and self._postman is not None:
            next_target = response_packet.to_agent
            if next_target and next_target != task.target:
                await self._route_response(response_packet)

        elapsed_ms = (time.monotonic() - start_time) * 1000
        result.execution_time_ms = elapsed_ms
        result.executor_instance_id = instance.instance_id

        self._task_status[task.task_id] = TaskStatus.COMPLETED
        self._results[task.task_id] = result

        inst = await self._runtime.get_instance(instance.instance_id)
        if inst:
            inst.tasks_completed += 1

        await self._pool.release(instance.instance_id)

        _logger.info(
            f"Task {task.task_id} completed on {instance.instance_id} "
            f"in {elapsed_ms:.1f}ms"
        )
        return result

    @traces(spec="001", story="US-002", reqs=["REQ-001-US-002-001"])
    async def _route_response(self, packet: CXPPacket) -> None:
        """Route response packet to next volk via Postman.

        Args:
            packet: Response CXPPacket with to_agent set.
        """
        from orchestrata_ultimate.cxp.types import DeliveryStatus

        if self._postman is None:
            return
        try:
            result = await self._postman.route(packet)
            if result.status == DeliveryStatus.DELIVERED:
                _logger.info(f"Routed response to {packet.to_agent}: {packet.subject}")
            else:
                _logger.warning(
                    f"Failed to route response to {packet.to_agent}: {result.status}"
                )
        except Exception as e:
            _logger.error(f"Error routing response: {e}")

    async def cancel(self, task_id: str) -> bool:
        """Cancel a task.

        Only PENDING tasks can be cancelled.

        Raises
        ------
        TaskCancellationError
            If the task is not in a cancellable state.
        """
        status = self._task_status.get(task_id)
        if status is None:
            raise TaskCancellationError(f"Task '{task_id}' not found", task_id=task_id)

        if status in (TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED):
            raise TaskCancellationError(
                f"Task '{task_id}' is already in terminal state '{status.value}'",
                task_id=task_id,
            )

        if status == TaskStatus.EXECUTING:
            # Release the instance if assigned
            instance_id = self._task_assignments.get(task_id)
            if instance_id:
                await self._pool.release(instance_id)

        self._task_status[task_id] = TaskStatus.CANCELLED
        return True

    async def get_task(self, task_id: str) -> TaskDefinition | None:
        return self._tasks.get(task_id)

    async def get_task_status(self, task_id: str) -> TaskStatus | None:
        return self._task_status.get(task_id)

    async def get_result(self, task_id: str) -> TaskResult | None:
        return self._results.get(task_id)

    async def list_tasks(
        self,
        status: str | None = None,
        alias: str | None = None,
    ) -> list[TaskDefinition]:
        tasks = list(self._tasks.values())
        if status is not None:
            tasks = [
                t
                for t in tasks
                if self._task_status.get(t.task_id, TaskStatus.PENDING).value == status
            ]
        if alias is not None:
            tasks = [t for t in tasks if t.target == alias]
        return tasks

    # ================================================================
    # Internal
    # ================================================================

    @traces(spec="001", story="US-002", reqs=["REQ-001-US-002-006"])
    def _resolve_alias(self, task: TaskDefinition) -> str:
        """Resolve the target Volk alias for a task.

        Uses the task's target field directly. If empty, uses
        Prolog to find the best Volk.
        """
        if task.target:
            return task.target

        # Try Prolog selection
        best = self._prolog.find_best_volk_for_task(task.task_id)
        if best:
            return best

        raise TaskAssignmentError(
            f"No target alias specified and Prolog found no candidate "
            f"for task '{task.task_id}'",
            task_id=task.task_id,
        )

    async def _execute_task(
        self, task: TaskDefinition, instance_id: str
    ) -> tuple[TaskResult, CXPPacket | None]:
        """Execute a task on a Volk instance.

        Dispatches to the actual Volk's process method:
        1. Look up the Volk class by alias
        2. Instantiate the Volk
        3. Convert TaskDefinition to CXPPacket
        4. Call volk.process(packet)
        5. Convert response to TaskResult
        """
        volk_class = _VOLK_CLASS_MAP.get(task.target)
        if volk_class is None:
            _logger.error("No volk class found for alias: %s", task.target)
            return (
                TaskResult(
                    task_id=task.task_id,
                    success=False,
                    error_message=f"Unknown volk alias: {task.target}",
                ),
                None,
            )

        volk = volk_class()
        packet = _task_to_packet(task)

        try:
            response = await volk.process(packet)
            if response is None:
                return (
                    TaskResult(
                        task_id=task.task_id,
                        success=True,
                        output={"action": task.action, "target": task.target},
                    ),
                    None,
                )
            return (
                TaskResult(
                    task_id=task.task_id,
                    success=True,
                    output={
                        "action": task.action,
                        "target": task.target,
                        "response": response.payload if response else {},
                    },
                ),
                response,
            )
        except Exception as e:
            _logger.error("Task execution failed: %s", e)
            return (
                TaskResult(
                    task_id=task.task_id,
                    success=False,
                    error_message=str(e),
                ),
                None,
            )

    @property
    def pending_count(self) -> int:
        return sum(1 for s in self._task_status.values() if s == TaskStatus.PENDING)

    @property
    def executing_count(self) -> int:
        return sum(1 for s in self._task_status.values() if s == TaskStatus.EXECUTING)

    @property
    def completed_count(self) -> int:
        return sum(1 for s in self._task_status.values() if s == TaskStatus.COMPLETED)