cleveragents-core/features/steps/plan_execute_runtime_steps.py

"""Step definitions for plan execute runtime integration tests."""

from __future__ import annotations

from typing import Any
from unittest.mock import MagicMock

from behave import given, then, when
from pydantic import ValidationError as PydanticValidationError
from ulid import ULID

from cleveragents.application.services.plan_execution_context import (
    PlanExecutionContext,
    RuntimeExecuteActor,
    RuntimeExecuteResult,
)
from cleveragents.application.services.plan_executor import (
    PlanExecutor,
    StrategyDecision,
)
from cleveragents.core.exceptions import PlanError, ValidationError
from cleveragents.domain.models.core.change import (
    ChangeEntry,
    ChangeOperation,
    InMemoryChangeSetStore,
)
from cleveragents.tool.registry import ToolRegistry
from cleveragents.tool.runner import ToolRunner


def _runner() -> ToolRunner:
    return ToolRunner(registry=ToolRegistry())


def _make_decision_root_id() -> str:
    return str(ULID())


def _mock_lifecycle() -> MagicMock:
    return MagicMock()


def _entry(
    plan_id: str, tool: str = "test/tool", path: str = "src/t.py"
) -> ChangeEntry:
    return ChangeEntry(
        plan_id=plan_id,
        resource_id=str(ULID()),
        tool_name=tool,
        operation=ChangeOperation.MODIFY,
        path=path,
    )


def _make_decisions(count: int) -> list[StrategyDecision]:
    root_id = str(ULID())
    decisions = []
    for i in range(count):
        decisions.append(
            StrategyDecision(
                decision_id=root_id if i == 0 else str(ULID()),
                step_text=f"Step {i + 1}",
                sequence=i,
                parent_id=root_id if i > 0 else None,
            )
        )
    return decisions


def _make_ctx(plan_id: str, **kwargs: Any) -> PlanExecutionContext:
    """Create a PlanExecutionContext with required fields."""
    defaults: dict[str, Any] = {"plan_id": plan_id}
    defaults.update(kwargs)
    return PlanExecutionContext(**defaults)


# -- Background --


@given("a fresh in-memory changeset store")
def step_bg_store(ctx: Any) -> None:
    ctx.changeset_store = InMemoryChangeSetStore()


@given("a valid plan ID")
def step_bg_plan_id(ctx: Any) -> None:
    ctx.plan_id = str(ULID())


# -- Context creation --


@when("I create a PlanExecutionContext with a valid plan_id")
def step_ctx_valid(ctx: Any) -> None:
    ctx.exec_ctx = _make_ctx(ctx.plan_id)


@when("I create a PlanExecutionContext with all optional fields")
def step_ctx_all(ctx: Any) -> None:
    ctx.decision_root_id = str(ULID())
    ctx.sandbox_root = "/tmp/test-sandbox"
    ctx.automation_profile = "trusted"
    ctx.exec_ctx = PlanExecutionContext(
        plan_id=ctx.plan_id,
        decision_root_id=ctx.decision_root_id,
        sandbox_root=ctx.sandbox_root,
        automation_profile=ctx.automation_profile,
        project_resources={"repo": {"path": "/code"}},
        resource_bindings={},
        changeset_store=ctx.changeset_store,
    )


@when("I attempt to create a PlanExecutionContext with empty plan_id")
def step_ctx_empty(ctx: Any) -> None:
    try:
        PlanExecutionContext(plan_id="")
        ctx.error = None
    except ValidationError as e:
        ctx.error = e


@when("I attempt to create a PlanExecutionContext with None plan_id")
def step_ctx_none(ctx: Any) -> None:
    try:
        PlanExecutionContext(plan_id=None)
        ctx.error = None
    except (ValidationError, TypeError) as e:
        ctx.error = e


@when("I create a PlanExecutionContext without specifying a changeset store")
def step_ctx_no_store(ctx: Any) -> None:
    ctx.exec_ctx = _make_ctx(ctx.plan_id)


@when("I create a PlanExecutionContext with a custom changeset store")
def step_ctx_custom_store(ctx: Any) -> None:
    ctx.custom_store = InMemoryChangeSetStore()
    ctx.exec_ctx = _make_ctx(ctx.plan_id, changeset_store=ctx.custom_store)


# -- Context assertions --


@then("the context plan_id should match the provided value")
def step_a_plan_id(ctx: Any) -> None:
    assert ctx.exec_ctx.plan_id == ctx.plan_id


@then("the context should have the correct decision_root_id")
def step_a_drid(ctx: Any) -> None:
    assert ctx.exec_ctx.decision_root_id == ctx.decision_root_id


@then("the context should have the correct sandbox_root")
def step_a_sandbox(ctx: Any) -> None:
    assert ctx.exec_ctx.sandbox_root == ctx.sandbox_root


@then("the context should have the correct automation_profile")
def step_a_profile(ctx: Any) -> None:
    assert ctx.exec_ctx.automation_profile == ctx.automation_profile


@then('the plan execution context should raise a ValidationError containing "{text}"')
def step_a_valerr(ctx: Any, text: str) -> None:
    assert ctx.error is not None, "Expected error"
    assert text in str(ctx.error), f"'{text}' not in '{ctx.error}'"


@then("the plan execution context should raise a ValidationError")
def step_a_valerr2(ctx: Any) -> None:
    assert ctx.error is not None


@then("the context changeset_store should be an InMemoryChangeSetStore")
def step_a_default_store(ctx: Any) -> None:
    assert isinstance(ctx.exec_ctx.changeset_store, InMemoryChangeSetStore)


@then("the context changeset_store should be the provided store")
def step_a_custom_store(ctx: Any) -> None:
    assert ctx.exec_ctx.changeset_store is ctx.custom_store


@then("the context project_resources should be an empty dict")
def step_a_empty_res(ctx: Any) -> None:
    assert ctx.exec_ctx.project_resources == {}


@then("the context resource_bindings should be an empty dict")
def step_a_empty_bind(ctx: Any) -> None:
    assert ctx.exec_ctx.resource_bindings == {}


# -- Changeset operations --


@given("a PlanExecutionContext")
def step_g_ctx(ctx: Any) -> None:
    ctx.exec_ctx = _make_ctx(ctx.plan_id, changeset_store=ctx.changeset_store)


@given('a PlanExecutionContext with sandbox_root "{path}"')
def step_g_ctx_sandbox(ctx: Any, path: str) -> None:
    ctx.exec_ctx = _make_ctx(
        ctx.plan_id,
        sandbox_root=path,
        changeset_store=ctx.changeset_store,
    )


@given("a PlanExecutionContext with an active changeset")
def step_g_ctx_cs(ctx: Any) -> None:
    ctx.exec_ctx = _make_ctx(ctx.plan_id, changeset_store=ctx.changeset_store)
    ctx.active_cs_id = ctx.exec_ctx.start_changeset()


@given("a PlanExecutionContext without an active changeset")
def step_g_ctx_no_cs(ctx: Any) -> None:
    ctx.exec_ctx = _make_ctx(ctx.plan_id, changeset_store=ctx.changeset_store)


@given("a PlanExecutionContext with an active changeset and one recorded entry")
def step_g_ctx_entry(ctx: Any) -> None:
    ctx.exec_ctx = _make_ctx(ctx.plan_id, changeset_store=ctx.changeset_store)
    ctx.active_cs_id = ctx.exec_ctx.start_changeset()
    ctx.test_entry = _entry(ctx.plan_id)
    ctx.exec_ctx.record_change(ctx.test_entry)


@when("I call start_changeset")
def step_w_start_cs(ctx: Any) -> None:
    ctx.changeset_id_result = ctx.exec_ctx.start_changeset()


@when("I record a ChangeEntry into the context")
def step_w_record(ctx: Any) -> None:
    ctx.test_entry = _entry(ctx.plan_id, path="src/new.py")
    ctx.exec_ctx.record_change(ctx.test_entry)


@when("I attempt to record a ChangeEntry")
def step_w_record_fail(ctx: Any) -> None:
    try:
        ctx.exec_ctx.record_change(_entry(ctx.plan_id))
        ctx.error = None
    except (PlanError, ValidationError, KeyError) as e:
        ctx.error = e


@when("I record three ChangeEntry instances")
def step_w_record3(ctx: Any) -> None:
    for i in range(3):
        ctx.exec_ctx.record_change(_entry(ctx.plan_id, path=f"src/f{i}.py"))


@when("I call get_changeset with the active changeset_id")
def step_w_get_cs(ctx: Any) -> None:
    ctx.retrieved_cs = ctx.exec_ctx.get_changeset(ctx.active_cs_id)


@when("I call get_changeset with a nonexistent ID")
def step_w_get_cs_miss(ctx: Any) -> None:
    ctx.retrieved_cs = ctx.exec_ctx.get_changeset("nonexistent-id")


@then("the returned changeset_id should be a valid ULID string")
def step_a_ulid(ctx: Any) -> None:
    cid = ctx.changeset_id_result
    assert cid and len(cid) == 26


@then("the active_changeset_ids should contain the new ID")
def step_a_active_ids(ctx: Any) -> None:
    assert ctx.changeset_id_result in ctx.exec_ctx.active_changeset_ids


@then("the changeset should contain the recorded entry")
def step_a_entry_in(ctx: Any) -> None:
    cs = ctx.exec_ctx.get_changeset(ctx.active_cs_id)
    assert cs is not None
    assert any(e.entry_id == ctx.test_entry.entry_id for e in cs.entries)


@then("a PlanError should be raised about no active changeset")
def step_a_plan_err(ctx: Any) -> None:
    assert ctx.error is not None, "Expected error about no changeset"


@then("the changeset should contain exactly three entries")
def step_a_three(ctx: Any) -> None:
    cs = ctx.exec_ctx.get_changeset(ctx.active_cs_id)
    assert cs and len(cs.entries) == 3


@then("the returned changeset should not be None")
def step_a_not_none(ctx: Any) -> None:
    assert ctx.retrieved_cs is not None


@then("the changeset plan_id should match")
def step_a_cs_pid(ctx: Any) -> None:
    assert ctx.retrieved_cs.plan_id == ctx.plan_id


@then("the returned changeset should be None")
def step_a_none(ctx: Any) -> None:
    assert ctx.retrieved_cs is None


# -- Summarize --


@when("I call summarize on the context")
def step_w_summarize(ctx: Any) -> None:
    ctx.summary = ctx.exec_ctx.summarize()


@then("the summary should contain plan_id")
def step_a_sum_pid(ctx: Any) -> None:
    assert ctx.summary["plan_id"] == ctx.plan_id


@then("the summary should contain resource_binding_count of {n:d}")
def step_a_sum_count(ctx: Any, n: int) -> None:
    assert ctx.summary["resource_binding_count"] == n


@then("the summary should contain decision_root_id")
def step_a_sum_drid(ctx: Any) -> None:
    assert "decision_root_id" in ctx.summary


@then("the summary resource_binding_count should be {n:d}")
def step_a_sum_count2(ctx: Any, n: int) -> None:
    assert ctx.summary["resource_binding_count"] == n


# -- RuntimeExecuteResult model --


@when("I create a RuntimeExecuteResult with valid fields")
def step_w_result_ok(ctx: Any) -> None:
    ctx.result = RuntimeExecuteResult(
        changeset_id=str(ULID()),
        tool_call_count=5,
        sandbox_refs=["/tmp/sb"],
        decision_ids_processed=["d1", "d2"],
        execution_duration_ms=42.5,
    )


@when("I attempt to create a RuntimeExecuteResult with negative tool_call_count")
def step_w_result_neg(ctx: Any) -> None:
    try:
        RuntimeExecuteResult(
            changeset_id=str(ULID()),
            tool_call_count=-1,
        )
        ctx.error = None
    except (PydanticValidationError, ValidationError, ValueError) as e:
        ctx.error = e


@when("I attempt to create a RuntimeExecuteResult with empty changeset_id")
def step_w_result_empty_csid(ctx: Any) -> None:
    try:
        RuntimeExecuteResult(
            changeset_id="",
            tool_call_count=0,
        )
        ctx.error = None
    except (PydanticValidationError, ValidationError, ValueError) as e:
        ctx.error = e


@then("the runtime execute result should have a correct changeset_id")
def step_a_res_csid(ctx: Any) -> None:
    assert ctx.result.changeset_id
    assert len(ctx.result.changeset_id) == 26


@then("the runtime execute result tool_call_count should be non-negative")
def step_a_res_count(ctx: Any) -> None:
    assert ctx.result.tool_call_count >= 0


@then("the runtime result should fail validation")
def step_a_pydantic_err(ctx: Any) -> None:
    assert ctx.error is not None


# -- RuntimeExecuteActor --


@given("a ToolRunner with an empty registry")
def step_g_runner(ctx: Any) -> None:
    ctx.tool_runner = _runner()


@given("a RuntimeExecuteActor")
def step_g_actor(ctx: Any) -> None:
    ctx.runtime_actor = RuntimeExecuteActor(
        tool_runner=ctx.tool_runner,
        execution_context=ctx.exec_ctx,
    )


@given("a stream callback collector")
def step_g_collector(ctx: Any) -> None:
    ctx.stream_events = []
    ctx.stream_callback = lambda t, d: ctx.stream_events.append((t, d))


@when("I execute with a list of two stub decisions")
def step_w_exec2(ctx: Any) -> None:
    ctx.result = ctx.runtime_actor.execute(
        decisions=_make_decisions(2),
    )


@when("I execute with one stub decision and the stream callback")
def step_w_exec1_cb(ctx: Any) -> None:
    ctx.result = ctx.runtime_actor.execute(
        decisions=_make_decisions(1),
        stream_callback=ctx.stream_callback,
    )


@when("I execute with one stub decision")
def step_w_exec1(ctx: Any) -> None:
    ctx.result = ctx.runtime_actor.execute(
        decisions=_make_decisions(1),
    )


@when("I attempt to create a RuntimeExecuteActor with None tool_runner")
def step_w_actor_no_run(ctx: Any) -> None:
    try:
        RuntimeExecuteActor(
            tool_runner=None,
            execution_context=ctx.exec_ctx,
        )
        ctx.error = None
    except (ValidationError, TypeError) as e:
        ctx.error = e


@when("I attempt to create a RuntimeExecuteActor with None execution_context")
def step_w_actor_no_ctx(ctx: Any) -> None:
    try:
        RuntimeExecuteActor(
            tool_runner=ctx.tool_runner,
            execution_context=None,
        )
        ctx.error = None
    except (ValidationError, TypeError) as e:
        ctx.error = e


@then("the runtime execute result should be a RuntimeExecuteResult")
def step_a_type(ctx: Any) -> None:
    assert isinstance(ctx.result, RuntimeExecuteResult)


@then("the runtime execute result changeset_id should be a valid ULID")
def step_a_csid_ulid(ctx: Any) -> None:
    assert len(ctx.result.changeset_id) == 26


@then("the runtime execute result decision_ids_processed should have two entries")
def step_a_2ids(ctx: Any) -> None:
    assert len(ctx.result.decision_ids_processed) == 2


@then("the callback should have received runtime_execute_started event")
def step_a_evt_start(ctx: Any) -> None:
    types = [e[0] for e in ctx.stream_events]
    assert "runtime_execute_started" in types


@then("the callback should have received runtime_execute_step event")
def step_a_evt_step(ctx: Any) -> None:
    types = [e[0] for e in ctx.stream_events]
    assert "runtime_execute_step" in types


@then("the callback should have received runtime_execute_complete event")
def step_a_evt_done(ctx: Any) -> None:
    types = [e[0] for e in ctx.stream_events]
    assert "runtime_execute_complete" in types


@then('the runtime execute result sandbox_refs should contain "{path}"')
def step_a_srefs(ctx: Any, path: str) -> None:
    assert ctx.result.sandbox_refs and path in ctx.result.sandbox_refs


# -- PlanExecutor integration --


@given("a mock lifecycle service")
def step_g_mock_lc(ctx: Any) -> None:
    ctx.lifecycle = _mock_lifecycle()


@given("a PlanExecutionContext with a known changeset store")
def step_g_ctx_known(ctx: Any) -> None:
    ctx.known_store = InMemoryChangeSetStore()
    ctx.exec_ctx = _make_ctx(ctx.plan_id, changeset_store=ctx.known_store)


@when("I create a PlanExecutor with a mock tool_calling_runtime")
def step_w_pe_with_runtime(ctx: Any) -> None:
    ctx.executor = PlanExecutor(
        lifecycle_service=ctx.lifecycle,
        tool_runner=ctx.tool_runner,
        execution_context=ctx.exec_ctx,
    )


@when("I create a PlanExecutor with the execution_context")
def step_w_pe_ctx(ctx: Any) -> None:
    if not hasattr(ctx, "exec_ctx"):
        ctx.exec_ctx = _make_ctx(ctx.plan_id, changeset_store=ctx.changeset_store)
    ctx.executor = PlanExecutor(
        lifecycle_service=ctx.lifecycle,
        execution_context=ctx.exec_ctx,
    )


@when("I create a PlanExecutor without execution_context")
def step_w_pe_no_ctx(ctx: Any) -> None:
    ctx.executor = PlanExecutor(lifecycle_service=ctx.lifecycle)


@then("has_runtime should be True")
def step_a_hr_true(ctx: Any) -> None:
    assert ctx.executor.has_runtime is True


@then("has_runtime should be False")
def step_a_hr_false(ctx: Any) -> None:
    assert ctx.executor.has_runtime is False


@then("changeset_store should be the same object as the context store")
def step_a_cs_same(ctx: Any) -> None:
    assert ctx.executor.changeset_store is ctx.known_store


@then("changeset_store should be None")
def step_a_cs_none(ctx: Any) -> None:
    assert ctx.executor.changeset_store is None


@then("execution_context property should return the context")
def step_a_ec_prop(ctx: Any) -> None:
    assert ctx.executor.execution_context is ctx.exec_ctx


# -- Error cases --


@when("I attempt to create a PlanExecutor with None lifecycle_service")
def step_w_pe_none_lc(ctx: Any) -> None:
    try:
        PlanExecutor(lifecycle_service=None)
        ctx.error = None
    except ValidationError as e:
        ctx.error = e


@given("a PlanExecutor without execution_context")
def step_g_pe_no_ctx(ctx: Any) -> None:
    if not hasattr(ctx, "lifecycle"):
        ctx.lifecycle = _mock_lifecycle()
    ctx.executor = PlanExecutor(lifecycle_service=ctx.lifecycle)


@when("I attempt to run execute with an empty plan_id")
def step_w_exec_empty(ctx: Any) -> None:
    try:
        ctx.executor.run_execute("")
        ctx.error = None
    except (ValidationError, PlanError) as e:
        ctx.error = e