cleveragents-core/features/steps/subplan_spawn_service_steps.py

"""Step definitions for subplan spawn service scenarios."""

from __future__ import annotations

from unittest.mock import MagicMock

from behave import given, then, when
from behave.runner import Context

from cleveragents.application.services.subplan_service import (
    SpawnEntry,
    SubplanService,
)

from cleveragents.core.exceptions import SubplanSpawnError
from cleveragents.domain.models.core.decision import (
    ContextSnapshot,
    Decision,
    DecisionType,
    ResourceRef,
)
from cleveragents.domain.models.core.plan import (
    ExecutionMode,
    NamespacedName,
    Plan,
    PlanIdentity,
    SubplanConfig,
)

_PLAN_ID = "01HGZ6FE0AQDYTR4BXVQZ6PN00"
_ROOT_ID = "01HGZ6FE0AQDYTR4BXVQZ6RF00"
_DEC_ID1 = "01HGZ6FE0AQDYTR4BXVQZ6DA00"
_DEC_ID2 = "01HGZ6FE0AQDYTR4BXVQZ6DB00"
_DEC_ID3 = "01HGZ6FE0AQDYTR4BXVQZ6DC00"


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


def _make_decision(
    decision_id: str,
    plan_id: str = _PLAN_ID,
    decision_type: DecisionType = DecisionType.SUBPLAN_SPAWN,
    sequence: int = 0,
    chosen_option: str = "local/sub-action",
    resource_ids: list[str] | None = None,
) -> Decision:
    """Create a Decision with sensible defaults."""
    resources: list[ResourceRef] = []
    if resource_ids:
        resources = [ResourceRef(resource_id=rid) for rid in resource_ids]
    return Decision(
        decision_id=decision_id,
        plan_id=plan_id,
        decision_type=decision_type,
        sequence_number=sequence,
        question="Should we spawn a child plan?",
        chosen_option=chosen_option,
        context_snapshot=ContextSnapshot(relevant_resources=resources),
    )


def _make_plan(
    plan_id: str = _PLAN_ID,
    parent_plan_id: str | None = None,
    root_plan_id: str | None = None,
) -> Plan:
    """Create a minimal Plan for testing."""
    return Plan(
        identity=PlanIdentity(
            plan_id=plan_id,
            parent_plan_id=parent_plan_id,
            root_plan_id=root_plan_id,
        ),
        namespaced_name=NamespacedName(namespace="local", name="test-plan"),
        description="A test plan for subplan spawn",
        action_name="local/test-action",
    )


# ---- Successful spawn ----


@given("a parent plan with a subplan spawn decision")
def step_given_parent_with_spawn_decision(context: Context) -> None:
    context.decision_service = _mock_decision_service()
    context.subplan_service = SubplanService(
        decision_service=context.decision_service,
    )
    context.parent_plan = _make_plan()
    dec: Decision = _make_decision(_DEC_ID1)
    context.spawn_entries = [
        SpawnEntry(
            decision=dec,
            action_name="local/sub-action",
            description="Test subplan",
        )
    ]


@given("a subplan config with sequential execution mode")
def step_given_config_sequential(context: Context) -> None:
    context.subplan_config = SubplanConfig(
        execution_mode=ExecutionMode.SEQUENTIAL,
    )


@when("the subplan service spawns child plans")
def step_when_spawn(context: Context) -> None:
    context.spawn_error = None
    try:
        context.spawn_result = context.subplan_service.spawn(
            parent_plan=context.parent_plan,
            config=context.subplan_config,
            spawn_entries=context.spawn_entries,
            available_resources=getattr(context, "available_resources", None),
        )
    except (SubplanSpawnError, ValueError) as exc:
        context.spawn_error = exc
        context.spawn_result = None


@then("{count:d} child plan status should be created")
def step_then_count_statuses_singular(context: Context, count: int) -> None:
    assert context.spawn_result is not None, f"Spawn failed: {context.spawn_error}"
    actual: int = len(context.spawn_result.spawned_statuses)
    assert actual == count, f"Expected {count} statuses, got {actual}"


@then("{count:d} child plan statuses should be created")
def step_then_count_statuses_plural(context: Context, count: int) -> None:
    assert context.spawn_result is not None, f"Spawn failed: {context.spawn_error}"
    actual: int = len(context.spawn_result.spawned_statuses)
    assert actual == count, f"Expected {count} statuses, got {actual}"


@then("the spawn metadata should contain the parent plan id")
def step_then_metadata_parent_id(context: Context) -> None:
    assert context.spawn_result is not None
    for meta in context.spawn_result.metadata.values():
        assert meta.parent_plan_id == context.parent_plan.identity.plan_id


@then("the spawn metadata should contain the root plan id")
def step_then_metadata_root_id(context: Context) -> None:
    assert context.spawn_result is not None
    plan: Plan = context.parent_plan
    expected_root: str = plan.identity.root_plan_id or plan.identity.plan_id
    for meta in context.spawn_result.metadata.values():
        assert meta.root_plan_id == expected_root


@then("the spawn metadata should contain the execution mode")
def step_then_metadata_exec_mode(context: Context) -> None:
    assert context.spawn_result is not None
    for meta in context.spawn_result.metadata.values():
        assert meta.execution_mode in {m.value for m in ExecutionMode}


# ---- Multiple entries ----


@given("a parent plan with {count:d} subplan spawn decisions")
def step_given_parent_with_n_decisions(context: Context, count: int) -> None:
    context.decision_service = _mock_decision_service()
    context.subplan_service = SubplanService(
        decision_service=context.decision_service,
    )
    context.parent_plan = _make_plan()
    context.spawn_entries = []
    for i in range(count):
        # Build unique ULID-like IDs
        suffix: str = f"{i:02d}"
        did: str = f"01HGZ6FE0AQDYTR4BXVQZ6D{suffix}0"
        dec: Decision = _make_decision(did, sequence=i)
        context.spawn_entries.append(
            SpawnEntry(
                decision=dec,
                action_name=f"local/sub-action-{i}",
                description=f"Subplan {i}",
            )
        )


@given("a subplan config with parallel execution mode and max_parallel {n:d}")
def step_given_config_parallel(context: Context, n: int) -> None:
    context.subplan_config = SubplanConfig(
        execution_mode=ExecutionMode.PARALLEL,
        max_parallel=n,
    )


@then("each status should have a unique subplan id")
def step_then_unique_ids(context: Context) -> None:
    assert context.spawn_result is not None
    ids: list[str] = [s.subplan_id for s in context.spawn_result.spawned_statuses]
    assert len(ids) == len(set(ids)), f"Duplicate IDs found: {ids}"


@then("the spawn result total_spawned should be {count:d}")
def step_then_total_spawned(context: Context, count: int) -> None:
    assert context.spawn_result is not None
    assert context.spawn_result.total_spawned == count


# ---- Validation: resource scopes ----


@given('a parent plan with a spawn entry targeting resource "{res_id}"')
def step_given_entry_with_resource(context: Context, res_id: str) -> None:
    context.decision_service = _mock_decision_service()
    context.subplan_service = SubplanService(
        decision_service=context.decision_service,
    )
    context.parent_plan = _make_plan()
    dec: Decision = _make_decision(
        _DEC_ID1,
        resource_ids=[res_id],
    )
    context.spawn_entries = [
        SpawnEntry(
            decision=dec,
            action_name="local/sub-action",
            target_resources=[res_id],
        )
    ]


@given('available resources are "{res_a}" and "{res_b}"')
def step_given_available_resources(
    context: Context,
    res_a: str,
    res_b: str,
) -> None:
    context.available_resources = frozenset({res_a, res_b})


@when("the subplan service validates the spawn request")
def step_when_validate(context: Context) -> None:
    context.validation_result = context.subplan_service.validate_spawn(
        config=context.subplan_config,
        spawn_entries=context.spawn_entries,
        available_resources=getattr(context, "available_resources", None),
    )


@then("the validation result should be invalid")
def step_then_invalid(context: Context) -> None:
    assert not context.validation_result.valid, (
        "Expected invalid, got valid with no errors"
    )


@then("the validation errors should mention unresolved resource")
def step_then_error_mentions_resource(context: Context) -> None:
    errors_str: str = " ".join(context.validation_result.errors)
    assert "resource" in errors_str.lower(), (
        f"Expected 'resource' in errors: {context.validation_result.errors}"
    )


# ---- Validation: merge strategy ----


@given("a parent plan with a valid spawn entry")
def step_given_valid_entry(context: Context) -> None:
    context.decision_service = _mock_decision_service()
    context.subplan_service = SubplanService(
        decision_service=context.decision_service,
    )
    context.parent_plan = _make_plan()
    dec: Decision = _make_decision(_DEC_ID1)
    context.spawn_entries = [
        SpawnEntry(
            decision=dec,
            action_name="local/sub-action",
        )
    ]


@given("a subplan config with no merge strategy")
def step_given_config_no_merge(context: Context) -> None:
    # Create a config and then forcibly set merge_strategy to None
    config: SubplanConfig = SubplanConfig(
        execution_mode=ExecutionMode.SEQUENTIAL,
    )
    # Pydantic models allow assignment when validate_assignment is True
    # We use model_copy to create a version with None merge_strategy
    # Since SubplanMergeStrategy is required, we use object.__setattr__
    object.__setattr__(config, "merge_strategy", None)
    context.subplan_config = config


@then("the validation errors should mention merge strategy")
def step_then_error_mentions_merge(context: Context) -> None:
    errors_str: str = " ".join(context.validation_result.errors)
    assert "merge strategy" in errors_str.lower(), (
        f"Expected 'merge strategy' in errors: {context.validation_result.errors}"
    )


# ---- Validation: max_parallel ----


@then("the validation errors should mention max_parallel")
def step_then_error_mentions_max_parallel(context: Context) -> None:
    errors_str: str = " ".join(context.validation_result.errors)
    assert "max_parallel" in errors_str.lower(), (
        f"Expected 'max_parallel' in errors: {context.validation_result.errors}"
    )


# ---- Metadata persistence ----


@given('a parent plan that is a subplan with root plan id "{root_id}"')
def step_given_subplan_parent(context: Context, root_id: str) -> None:
    context.decision_service = _mock_decision_service()
    context.subplan_service = SubplanService(
        decision_service=context.decision_service,
    )
    context.parent_plan = _make_plan(
        parent_plan_id="01HGZ6FE0AQDYTR4BXVQZ6PP00",
        root_plan_id=root_id,
    )


@given("a valid spawn decision")
def step_given_valid_spawn_decision(context: Context) -> None:
    dec: Decision = _make_decision(_DEC_ID1)
    context.spawn_entries = [
        SpawnEntry(
            decision=dec,
            action_name="local/sub-action",
        )
    ]


@then('the spawn metadata root_plan_id should be "{expected}"')
def step_then_metadata_root_equals(context: Context, expected: str) -> None:
    assert context.spawn_result is not None
    for meta in context.spawn_result.metadata.values():
        assert meta.root_plan_id == expected, (
            f"Expected root '{expected}', got '{meta.root_plan_id}'"
        )


@then('the spawn metadata execution_mode should be "{expected}"')
def step_then_metadata_mode_equals(context: Context, expected: str) -> None:
    assert context.spawn_result is not None
    for meta in context.spawn_result.metadata.values():
        assert meta.execution_mode == expected, (
            f"Expected mode '{expected}', got '{meta.execution_mode}'"
        )


# ---- Decision type validation ----


@given("a parent plan with a non-spawn decision type")
def step_given_non_spawn_decision(context: Context) -> None:
    context.decision_service = _mock_decision_service()
    context.subplan_service = SubplanService(
        decision_service=context.decision_service,
    )
    context.parent_plan = _make_plan()
    dec: Decision = _make_decision(
        _DEC_ID1,
        decision_type=DecisionType.STRATEGY_CHOICE,
    )
    context.spawn_entries = [
        SpawnEntry(
            decision=dec,
            action_name="local/sub-action",
        )
    ]


@then("the validation errors should mention decision type")
def step_then_error_mentions_decision_type(context: Context) -> None:
    errors_str: str = " ".join(context.validation_result.errors)
    assert "type" in errors_str.lower(), (
        f"Expected 'type' in errors: {context.validation_result.errors}"
    )


# ---- Get spawn decisions ----


@given('a decision service with spawn decisions for plan "{plan_id}"')
def step_given_decision_service_with_spawns(
    context: Context,
    plan_id: str,
) -> None:
    spawn_dec: Decision = _make_decision(
        _DEC_ID1,
        plan_id=plan_id,
        sequence=0,
    )
    parallel_dec: Decision = _make_decision(
        _DEC_ID2,
        plan_id=plan_id,
        decision_type=DecisionType.SUBPLAN_PARALLEL_SPAWN,
        sequence=1,
    )

    mock_svc: MagicMock = _mock_decision_service()
    mock_svc.list_by_type = MagicMock(
        side_effect=lambda pid, dtype: (
            [spawn_dec] if dtype == DecisionType.SUBPLAN_SPAWN.value else [parallel_dec]
        )
    )
    context.decision_service = mock_svc
    context.subplan_service = SubplanService(
        decision_service=mock_svc,
    )
    context.plan_id = plan_id


@when("the subplan service retrieves spawn decisions")
def step_when_get_spawn_decisions(context: Context) -> None:
    context.spawn_decisions = context.subplan_service.get_spawn_decisions(
        context.plan_id,
    )


@then("the returned decisions should only contain spawn types")
def step_then_only_spawn_types(context: Context) -> None:
    for dec in context.spawn_decisions:
        assert dec.decision_type in (
            DecisionType.SUBPLAN_SPAWN,
            DecisionType.SUBPLAN_PARALLEL_SPAWN,
        ), f"Unexpected decision type: {dec.decision_type}"


# ---- Build spawn entries ----


@given("a list of spawn decisions with resource refs")
def step_given_decisions_with_resources(context: Context) -> None:
    context.decision_service = _mock_decision_service()
    context.subplan_service = SubplanService(
        decision_service=context.decision_service,
    )
    context.input_decisions = [
        _make_decision(
            _DEC_ID1,
            resource_ids=["01HGZ6FE0AQDYTR4BXVQZ6RA00"],
            chosen_option="local/test-sub",
        ),
        _make_decision(
            _DEC_ID2,
            sequence=1,
            resource_ids=[
                "01HGZ6FE0AQDYTR4BXVQZ6RB00",
                "01HGZ6FE0AQDYTR4BXVQZ6RC00",
            ],
            chosen_option="local/test-sub-2",
        ),
    ]


@when("the subplan service builds spawn entries")
def step_when_build_entries(context: Context) -> None:
    context.built_entries = context.subplan_service.build_spawn_entries(
        context.input_decisions,
    )


@then("each entry should have action_name and target_resources")
def step_then_entries_have_fields(context: Context) -> None:
    for entry in context.built_entries:
        assert entry.action_name, "Entry missing action_name"
        # First entry should have 1 resource, second should have 2
    assert len(context.built_entries) == 2
    assert len(context.built_entries[0].target_resources) == 1
    assert len(context.built_entries[1].target_resources) == 2


# ---- Service construction ----


@when("the subplan service is constructed with None decision_service")
def step_when_construct_none(context: Context) -> None:
    context.construct_error = None
    try:
        SubplanService(decision_service=None)  # type: ignore[arg-type]
    except ValueError as exc:
        context.construct_error = exc


@then("a ValueError should be raised for subplan service")
def step_then_value_error(context: Context) -> None:
    error: Exception | None = getattr(
        context,
        "construct_error",
        None,
    ) or getattr(context, "spawn_error", None)
    assert error is not None, "Expected ValueError but none was raised"
    assert isinstance(error, ValueError), (
        f"Expected ValueError, got {type(error).__name__}"
    )


# ---- Null guard scenarios ----


@given("a valid subplan service")
def step_given_valid_service(context: Context) -> None:
    context.decision_service = _mock_decision_service()
    context.subplan_service = SubplanService(
        decision_service=context.decision_service,
    )
    context.subplan_config = SubplanConfig()


@given("a parent plan with identity")
def step_given_plan_with_identity(context: Context) -> None:
    context.parent_plan = _make_plan()


@when("spawn is called with None parent_plan")
def step_when_spawn_none_plan(context: Context) -> None:
    context.spawn_error = None
    try:
        context.subplan_service.spawn(
            parent_plan=None,  # type: ignore[arg-type]
            config=context.subplan_config,
            spawn_entries=[
                SpawnEntry(
                    decision=_make_decision(_DEC_ID1),
                    action_name="local/test",
                )
            ],
        )
    except ValueError as exc:
        context.spawn_error = exc


@when("spawn is called with empty spawn_entries")
def step_when_spawn_empty_entries(context: Context) -> None:
    context.spawn_error = None
    try:
        context.subplan_service.spawn(
            parent_plan=context.parent_plan,
            config=context.subplan_config,
            spawn_entries=[],
        )
    except ValueError as exc:
        context.spawn_error = exc


@then('a spawn ValueError should be raised with message "{msg}"')
def step_then_value_error_msg(context: Context, msg: str) -> None:
    error: Exception | None = getattr(
        context,
        "construct_error",
        None,
    ) or getattr(context, "spawn_error", None)
    assert error is not None, f"Expected ValueError with '{msg}'"
    assert isinstance(error, ValueError)
    assert msg in str(error), f"Expected '{msg}' in '{error}'"


# ---- Parallel spawn decision ----


@given("a parent plan with a subplan_parallel_spawn decision")
def step_given_parallel_spawn_decision(context: Context) -> None:
    context.decision_service = _mock_decision_service()
    context.subplan_service = SubplanService(
        decision_service=context.decision_service,
    )
    context.parent_plan = _make_plan()
    dec: Decision = _make_decision(
        _DEC_ID1,
        decision_type=DecisionType.SUBPLAN_PARALLEL_SPAWN,
    )
    context.spawn_entries = [
        SpawnEntry(
            decision=dec,
            action_name="local/parallel-sub",
        )
    ]