cleveragents-core/robot/helper_subplan_spawn_orchestration.py

"""Helper script for subplan_spawn_orchestration.robot integration tests.

Each subcommand exercises SubplanService.spawn() in parallel mode and
verifies end-to-end orchestration: child Plan creation, identity wiring,
project link inheritance, parent status attachment, and execution mode.
"""

from __future__ import annotations

import sys
from collections.abc import Callable
from pathlib import Path
from typing import NoReturn
from unittest.mock import MagicMock

# Ensure local source tree is importable.
_ROOT: Path = Path(__file__).resolve().parents[1]
_SRC: str = str(_ROOT / "src")
_ROBOT: str = str(_ROOT / "robot")
for _p in (_SRC, _ROBOT):
    if _p not in sys.path:
        sys.path.insert(0, _p)

from cleveragents.application.services.subplan_service import (  # noqa: E402
    SpawnEntry,
    SpawnResult,
    SubplanService,
)
from cleveragents.domain.models.core.decision import (  # noqa: E402
    ContextSnapshot,
    Decision,
    DecisionType,
)
from cleveragents.domain.models.core.plan import (  # noqa: E402
    ExecutionMode,
    NamespacedName,
    Plan,
    PlanIdentity,
    PlanPhase,
    ProcessingState,
    ProjectLink,
    SubplanConfig,
)

_PARENT_PLAN_ID: str = "01HGZ6FE0AQDYTR4BXVQZ6PN00"
_ROOT_PLAN_ID: str = "01HGZ6FE0AQDYTR4BXVQZ6RF00"
_DEC_ID1: str = "01HGZ6FE0AQDYTR4BXVQZ6DA00"
_DEC_ID2: str = "01HGZ6FE0AQDYTR4BXVQZ6DB00"
_DEC_ID3: str = "01HGZ6FE0AQDYTR4BXVQZ6DC00"


def _fail(msg: str) -> NoReturn:
    """Print error message to stderr and exit with code 1."""
    print(msg, file=sys.stderr)
    sys.exit(1)


def _mock_decision_service() -> MagicMock:
    """Create a mock DecisionService."""
    svc: MagicMock = MagicMock()
    svc.list_by_type = MagicMock(return_value=[])
    return svc


def _make_parent_plan() -> Plan:
    """Create a parent plan configured for parallel subplan spawning."""
    return Plan(
        identity=PlanIdentity(
            plan_id=_PARENT_PLAN_ID,
            root_plan_id=_ROOT_PLAN_ID,
        ),
        namespaced_name=NamespacedName(namespace="local", name="parallel-parent"),
        description="Parent plan for parallel subplan orchestration",
        action_name="local/parallel-action",
        phase=PlanPhase.STRATEGIZE,
        processing_state=ProcessingState.PROCESSING,
        project_links=[
            ProjectLink(project_name="local/my-project"),
            ProjectLink(project_name="local/other-project", read_only=True),
        ],
        subplan_config=SubplanConfig(
            execution_mode=ExecutionMode.PARALLEL,
            max_parallel=5,
        ),
    )


def _make_spawn_entries() -> list[SpawnEntry]:
    """Create three spawn entries for parallel testing."""
    decisions: list[tuple[str, str, str]] = [
        (_DEC_ID1, "Spawn child A?", "local/sub-action-a"),
        (_DEC_ID2, "Spawn child B?", "local/sub-action-b"),
        (_DEC_ID3, "Spawn child C?", "local/sub-action-c"),
    ]
    entries: list[SpawnEntry] = []
    for i, (dec_id, question, action) in enumerate(decisions):
        dec: Decision = Decision(
            decision_id=dec_id,
            plan_id=_PARENT_PLAN_ID,
            decision_type=DecisionType.SUBPLAN_PARALLEL_SPAWN,
            sequence_number=i,
            question=question,
            chosen_option=action,
            context_snapshot=ContextSnapshot(),
        )
        entries.append(SpawnEntry(decision=dec, action_name=action))
    return entries


def _do_spawn() -> tuple[Plan, SubplanConfig, list[SpawnEntry], SpawnResult]:
    """Execute spawn and return all artifacts."""
    svc: SubplanService = SubplanService(
        decision_service=_mock_decision_service(),
    )
    parent: Plan = _make_parent_plan()
    config: SubplanConfig = SubplanConfig(
        execution_mode=ExecutionMode.PARALLEL,
        max_parallel=5,
    )
    entries: list[SpawnEntry] = _make_spawn_entries()
    result: SpawnResult = svc.spawn(
        parent_plan=parent,
        config=config,
        spawn_entries=entries,
    )
    return parent, config, entries, result


# ---------------------------------------------------------------------------
# Subcommands
# ---------------------------------------------------------------------------


def _parallel_spawn_identity() -> None:
    """Verify child Plans have correct parent/root identity."""
    _parent, _config, entries, result = _do_spawn()
    child_plans: list[Plan] | None = getattr(result, "child_plans", None)
    if child_plans is None:
        _fail("SpawnResult has no child_plans attribute")
    if len(child_plans) != len(entries):
        _fail(f"Expected {len(entries)} child plans, got {len(child_plans)}")
    for child in child_plans:
        if not isinstance(child, Plan):
            _fail(f"Expected Plan, got {type(child).__name__}")
        if child.identity.parent_plan_id != _PARENT_PLAN_ID:
            _fail(
                f"parent_plan_id={child.identity.parent_plan_id!r}, "
                f"expected {_PARENT_PLAN_ID!r}"
            )
        if child.identity.root_plan_id != _ROOT_PLAN_ID:
            _fail(
                f"root_plan_id={child.identity.root_plan_id!r}, "
                f"expected {_ROOT_PLAN_ID!r}"
            )
        if child.phase != PlanPhase.STRATEGIZE:
            _fail(f"Child phase={child.phase!r}, expected STRATEGIZE")
        if child.processing_state != ProcessingState.QUEUED:
            _fail(f"Child state={child.processing_state!r}, expected QUEUED")
    print("parallel-spawn-identity-ok")


def _parallel_spawn_project_links() -> None:
    """Verify child plans inherit project links from parent."""
    parent, _config, _entries, result = _do_spawn()
    child_plans: list[Plan] | None = getattr(result, "child_plans", None)
    if child_plans is None:
        _fail("SpawnResult has no child_plans attribute")
    parent_names: list[str] = [lnk.project_name for lnk in parent.project_links]
    for child in child_plans:
        child_names: list[str] = [lnk.project_name for lnk in child.project_links]
        if child_names != parent_names:
            _fail(f"Child project links {child_names!r} != parent {parent_names!r}")
    print("parallel-spawn-project-links-ok")


def _parallel_spawn_parent_statuses() -> None:
    """Verify spawn() attaches SubplanStatus entries to the parent."""
    parent, _config, entries, result = _do_spawn()
    if len(parent.subplan_statuses) != len(entries):
        _fail(
            f"Parent has {len(parent.subplan_statuses)} statuses, "
            f"expected {len(entries)}"
        )
    # Verify each status references a child plan
    child_ids: set[str] = set()
    child_plans: list[Plan] | None = getattr(result, "child_plans", None)
    if child_plans:
        child_ids = {c.identity.plan_id for c in child_plans}
    for status in parent.subplan_statuses:
        if status.subplan_id not in child_ids:
            _fail(f"SubplanStatus id={status.subplan_id!r} not in child plan IDs")
        if status.status != ProcessingState.QUEUED:
            _fail(f"Status state={status.status!r}, expected QUEUED")
    print("parallel-spawn-parent-statuses-ok")


def _parallel_spawn_mode() -> None:
    """Verify spawn result reports PARALLEL execution mode."""
    _parent, _config, _entries, result = _do_spawn()
    if result.execution_mode != ExecutionMode.PARALLEL:
        _fail(f"execution_mode={result.execution_mode!r}, expected PARALLEL")
    if result.total_spawned != 3:
        _fail(f"total_spawned={result.total_spawned}, expected 3")
    print("parallel-spawn-mode-ok")


# ---------------------------------------------------------------------------
# Dispatcher
# ---------------------------------------------------------------------------

_COMMANDS: dict[str, Callable[[], None]] = {
    "parallel-spawn-identity": _parallel_spawn_identity,
    "parallel-spawn-project-links": _parallel_spawn_project_links,
    "parallel-spawn-parent-statuses": _parallel_spawn_parent_statuses,
    "parallel-spawn-mode": _parallel_spawn_mode,
}

if __name__ == "__main__":
    if len(sys.argv) < 2 or sys.argv[1] not in _COMMANDS:
        print(
            f"Usage: {sys.argv[0]} <{'|'.join(_COMMANDS)}>",
            file=sys.stderr,
        )
        sys.exit(1)
    cmd: Callable[[], None] = _COMMANDS[sys.argv[1]]
    cmd()