temp/robot/helper_plan_executor_subplan_spawning.py

"""Helper script for PlanExecutor subplan spawning integration Robot tests.

Exercises the end-to-end subplan spawning path through PlanExecutor:
  - subplan_spawn decision realisation (stub mode)
  - subplan_parallel_spawn decision realisation (stub mode)
  - no-op when no spawn decisions exist
  - parent plan status tracking for child subplan failures
  - no-op when SubplanService is not configured
"""

from __future__ import annotations

import sys
from unittest.mock import MagicMock

from cleveragents.application.services.plan_executor import PlanExecutor
from cleveragents.application.services.subplan_execution_service import (
    SubplanExecutionResult,
    SubplanExecutionService,
)
from cleveragents.application.services.subplan_service import (
    SpawnEntry,
    SpawnResult,
    SubplanService,
)
from cleveragents.domain.models.core.decision import (
    ContextSnapshot,
    Decision,
    DecisionType,
)
from cleveragents.domain.models.core.plan import (
    ExecutionMode,
    NamespacedName,
    Plan,
    PlanIdentity,
    PlanPhase,
    ProcessingState,
    SubplanConfig,
    SubplanStatus,
)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

_PLAN_ID = "01KNFMGJSG67S6RG9TVXV205TQ"
_ROOT_ID = "01KNFMGJSH67S6RG9TVXV205TR"
_DEC_ID = "01KNFMGJSH67S6RG9TVXV205TS"
_SUBPLAN_ID = "01KNFMGJSH67S6RG9TVXV205TT"


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_plan(
    *,
    subplan_config: SubplanConfig | None = None,
) -> Plan:
    plan = Plan(
        identity=PlanIdentity(
            plan_id=_PLAN_ID,
            root_plan_id=_ROOT_ID,
        ),
        namespaced_name=NamespacedName(namespace="local", name="robot-test-plan"),
        description="Robot test plan for subplan spawning",
        action_name="local/robot-test-action",
        phase=PlanPhase.EXECUTE,
        processing_state=ProcessingState.QUEUED,
        subplan_config=subplan_config,
    )
    plan.decision_root_id = _ROOT_ID
    return plan


def _make_decision(
    decision_type: DecisionType = DecisionType.SUBPLAN_SPAWN,
) -> Decision:
    return Decision(
        decision_id=_DEC_ID,
        plan_id=_PLAN_ID,
        decision_type=decision_type,
        sequence_number=0,
        question="Spawn a child plan?",
        chosen_option="local/sub-action",
        context_snapshot=ContextSnapshot(relevant_resources=[]),
    )


def _make_spawn_result(
    execution_mode: str = ExecutionMode.SEQUENTIAL,
) -> SpawnResult:
    sub_status = SubplanStatus(
        subplan_id=_SUBPLAN_ID,
        action_name="local/sub-action",
        status=ProcessingState.QUEUED,
    )
    return SpawnResult(
        spawned_statuses=[sub_status],
        metadata={},
        total_spawned=1,
        execution_mode=execution_mode,
        child_plans=[],
    )


def _make_exec_result(
    all_succeeded: bool = True,
    failed_ids: list[str] | None = None,
) -> SubplanExecutionResult:
    failed = failed_ids or []
    status = ProcessingState.COMPLETE if all_succeeded else ProcessingState.ERRORED
    sub_status = SubplanStatus(
        subplan_id=_SUBPLAN_ID,
        action_name="local/sub-action",
        status=status,
        error=None if all_succeeded else "ChildExecError: timeout",
    )
    return SubplanExecutionResult(
        all_succeeded=all_succeeded,
        statuses=[sub_status],
        merge_result=None,
        total_duration_ms=10,
        failed_subplan_ids=failed,
    )


def _make_lifecycle(plan: Plan) -> MagicMock:
    lcs = MagicMock()
    lcs.get_plan.return_value = plan
    lcs.start_execute = MagicMock()
    lcs.complete_execute = MagicMock()
    lcs.fail_execute = MagicMock()
    lcs._commit_plan = MagicMock()
    return lcs


def _make_subplan_service(
    decisions: list[Decision] | None = None,
    spawn_result: SpawnResult | None = None,
) -> MagicMock:
    svc = MagicMock(spec=SubplanService)
    svc.get_spawn_decisions.return_value = decisions or []
    svc.build_spawn_entries.return_value = [
        SpawnEntry(decision=d, action_name="local/sub-action")
        for d in (decisions or [])
    ]
    svc.spawn.return_value = spawn_result or _make_spawn_result()
    return svc


def _make_subplan_execution_service(
    exec_result: SubplanExecutionResult | None = None,
) -> MagicMock:
    svc = MagicMock(spec=SubplanExecutionService)
    svc.execute_all.return_value = exec_result or _make_exec_result()
    return svc


def _make_execute_actor(plan: Plan) -> MagicMock:
    """Build a mock execute actor that returns a minimal result."""
    actor = MagicMock()
    result = MagicMock()
    result.changeset_id = "01JROBOT000000000000CS00001"
    result.sandbox_refs = []
    result.tool_calls_count = 0
    actor.execute.return_value = result
    return actor


# ---------------------------------------------------------------------------
# Test functions
# ---------------------------------------------------------------------------


def _test_spawn_decision_realisation() -> None:
    """subplan_spawn decision is realised as child plan execution."""
    decision = _make_decision(DecisionType.SUBPLAN_SPAWN)
    spawn_result = _make_spawn_result()
    exec_result = _make_exec_result(all_succeeded=True)

    mock_subplan_svc = _make_subplan_service(
        decisions=[decision],
        spawn_result=spawn_result,
    )
    mock_exec_svc = _make_subplan_execution_service(exec_result=exec_result)

    plan = _make_plan()
    lcs = _make_lifecycle(plan)
    execute_actor = _make_execute_actor(plan)

    executor = PlanExecutor(
        lifecycle_service=lcs,
        execute_actor=execute_actor,
        subplan_service=mock_subplan_svc,
        subplan_execution_service=mock_exec_svc,
    )

    executor.run_execute(plan_id=_PLAN_ID)

    # Verify the full spawning pipeline was invoked
    mock_subplan_svc.get_spawn_decisions.assert_called_once_with(_PLAN_ID)
    mock_subplan_svc.build_spawn_entries.assert_called_once()
    mock_subplan_svc.spawn.assert_called_once()
    mock_exec_svc.execute_all.assert_called_once()

    # Verify _commit_plan was called (plan state persisted)
    lcs._commit_plan.assert_called()

    print("spawn-decision-realisation-ok")


def _test_parallel_spawn_decision_realisation() -> None:
    """subplan_parallel_spawn decision triggers parallel execution."""
    decision = _make_decision(DecisionType.SUBPLAN_PARALLEL_SPAWN)
    spawn_result = _make_spawn_result(execution_mode=ExecutionMode.PARALLEL)
    exec_result = _make_exec_result(all_succeeded=True)

    mock_subplan_svc = _make_subplan_service(
        decisions=[decision],
        spawn_result=spawn_result,
    )
    mock_exec_svc = _make_subplan_execution_service(exec_result=exec_result)

    plan = _make_plan()
    lcs = _make_lifecycle(plan)
    execute_actor = _make_execute_actor(plan)

    executor = PlanExecutor(
        lifecycle_service=lcs,
        execute_actor=execute_actor,
        subplan_service=mock_subplan_svc,
        subplan_execution_service=mock_exec_svc,
    )

    executor.run_execute(plan_id=_PLAN_ID)

    mock_subplan_svc.get_spawn_decisions.assert_called_once_with(_PLAN_ID)
    mock_subplan_svc.spawn.assert_called_once()
    mock_exec_svc.execute_all.assert_called_once()

    print("parallel-spawn-decision-realisation-ok")


def _test_noop_when_no_spawn_decisions() -> None:
    """No-op when no spawn decisions are recorded."""
    mock_subplan_svc = _make_subplan_service(decisions=[])
    mock_exec_svc = _make_subplan_execution_service()

    plan = _make_plan()
    lcs = _make_lifecycle(plan)
    execute_actor = _make_execute_actor(plan)

    executor = PlanExecutor(
        lifecycle_service=lcs,
        execute_actor=execute_actor,
        subplan_service=mock_subplan_svc,
        subplan_execution_service=mock_exec_svc,
    )

    executor.run_execute(plan_id=_PLAN_ID)

    # get_spawn_decisions called but spawn and execute_all NOT called
    mock_subplan_svc.get_spawn_decisions.assert_called_once_with(_PLAN_ID)
    mock_subplan_svc.spawn.assert_not_called()
    mock_exec_svc.execute_all.assert_not_called()

    print("noop-no-spawn-decisions-ok")


def _test_noop_when_no_subplan_service() -> None:
    """No-op when SubplanService is not configured."""
    plan = _make_plan()
    lcs = _make_lifecycle(plan)
    execute_actor = _make_execute_actor(plan)

    executor = PlanExecutor(
        lifecycle_service=lcs,
        execute_actor=execute_actor,
        subplan_service=None,
        subplan_execution_service=None,
    )

    # Should complete without error
    executor.run_execute(plan_id=_PLAN_ID)

    # Verify lifecycle completed normally
    lcs.complete_execute.assert_called_once_with(_PLAN_ID)

    print("noop-no-subplan-service-ok")


def _test_parent_status_tracking_on_failure() -> None:
    """Parent plan error_details records failed subplan IDs."""
    decision = _make_decision(DecisionType.SUBPLAN_SPAWN)
    spawn_result = _make_spawn_result()
    exec_result = _make_exec_result(
        all_succeeded=False,
        failed_ids=[_SUBPLAN_ID],
    )

    mock_subplan_svc = _make_subplan_service(
        decisions=[decision],
        spawn_result=spawn_result,
    )
    mock_exec_svc = _make_subplan_execution_service(exec_result=exec_result)

    plan = _make_plan()
    lcs = _make_lifecycle(plan)
    execute_actor = _make_execute_actor(plan)

    executor = PlanExecutor(
        lifecycle_service=lcs,
        execute_actor=execute_actor,
        subplan_service=mock_subplan_svc,
        subplan_execution_service=mock_exec_svc,
    )

    executor.run_execute(plan_id=_PLAN_ID)

    # Verify _commit_plan was called with a plan that has failure info
    commit_calls = lcs._commit_plan.call_args_list
    assert commit_calls, "Expected _commit_plan to be called"

    found_failure_info = False
    for c in commit_calls:
        committed_plan = c[0][0]
        details = getattr(committed_plan, "error_details", None) or {}
        if (
            "failed_subplan_ids" in details
            and details.get("subplan_execution_failed") == "true"
        ):
            found_failure_info = True
            break

    assert found_failure_info, (
        f"Expected error_details with failure info in _commit_plan calls. "
        f"Calls: {commit_calls}"
    )

    print("parent-status-tracking-failure-ok")


def _test_subplan_service_property() -> None:
    """PlanExecutor exposes subplan_service and subplan_execution_service properties."""
    mock_subplan_svc = _make_subplan_service()
    mock_exec_svc = _make_subplan_execution_service()

    plan = _make_plan()
    lcs = _make_lifecycle(plan)

    executor = PlanExecutor(
        lifecycle_service=lcs,
        subplan_service=mock_subplan_svc,
        subplan_execution_service=mock_exec_svc,
    )

    assert executor.subplan_service is mock_subplan_svc
    assert executor.subplan_execution_service is mock_exec_svc

    print("subplan-service-properties-ok")


# ---------------------------------------------------------------------------
# Dispatch
# ---------------------------------------------------------------------------

COMMANDS: dict[str, object] = {
    "spawn-decision-realisation": _test_spawn_decision_realisation,
    "parallel-spawn-decision-realisation": _test_parallel_spawn_decision_realisation,
    "noop-no-spawn-decisions": _test_noop_when_no_spawn_decisions,
    "noop-no-subplan-service": _test_noop_when_no_subplan_service,
    "parent-status-tracking-failure": _test_parent_status_tracking_on_failure,
    "subplan-service-properties": _test_subplan_service_property,
}


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: helper_plan_executor_subplan_spawning.py <command>")
        sys.exit(1)

    cmd = sys.argv[1]
    fn = COMMANDS.get(cmd)
    if fn is None:
        print(f"Unknown command: {cmd}")
        sys.exit(1)

    fn()  # type: ignore[operator]