cleveragents-core/features/steps/security_async_steps.py

"""Step definitions for async resource cleanup and leak prevention."""

import asyncio
import logging
import tempfile
from pathlib import Path

from behave import given, then, when

from cleveragents.a2a.events import A2aEventQueue
from cleveragents.core.async_cleanup import AsyncResourceTracker
from cleveragents.langgraph.state import GraphState, StateManager

# ---------------------------------------------------------------------------
# Helpers — lightweight mock async resource (lives in steps, not src/)
# ---------------------------------------------------------------------------


class MockAsyncResource:
    """Test double for an async-closable resource."""

    def __init__(self, name: str, close_delay: float = 0.0) -> None:
        self.name = name
        self.close_delay = close_delay
        self.closed = False
        self.close_count = 0

    async def close(self) -> None:
        if self.close_delay > 0:
            # Use the original (un-patched) asyncio.sleep so timeout-based
            # tests observe real wall-clock delays.
            _real_sleep = getattr(asyncio, "_original_sleep", asyncio.sleep)
            await _real_sleep(self.close_delay)
        self.closed = True
        self.close_count += 1


# ---------------------------------------------------------------------------
# Background
# ---------------------------------------------------------------------------


@given("I have an async resource tracker")
def step_create_tracker(context):
    context.tracker = AsyncResourceTracker()
    context.resources = {}
    context.log_handler = _CapturingHandler()
    _logger = logging.getLogger("cleveragents.core.async_cleanup")
    # Remove stale handlers from previous scenarios (sequential mode)
    for h in list(_logger.handlers):
        if isinstance(h, _CapturingHandler):
            _logger.removeHandler(h)
    _logger.addHandler(context.log_handler)
    # Ensure the logger propagates warnings even if a prior test disabled them.
    _logger.setLevel(logging.DEBUG)
    _logger.disabled = False
    context.warnings_logged = context.log_handler.records


# ---------------------------------------------------------------------------
# Givens
# ---------------------------------------------------------------------------


@given('I have a mock async resource named "{name}"')
def step_create_mock_resource(context, name):
    resource = MockAsyncResource(name)
    context.resources[name] = resource
    context.current_resource = resource
    context.current_resource_name = name


@given(
    'I have a mock async resource named "{name}" that takes {seconds:g} seconds to close'
)
def step_create_slow_resource(context, name, seconds):
    resource = MockAsyncResource(name, close_delay=seconds)
    context.resources[name] = resource
    context.current_resource = resource
    context.current_resource_name = name


@given("I have {count:d} mock async resources with varied close times")
def step_create_multiple_resources(context, count):
    context.multi_resources = []
    for i in range(count):
        name = f"resource-{i}"
        resource = MockAsyncResource(name, close_delay=0.01 * (i + 1))
        context.resources[name] = resource
        context.multi_resources.append((name, resource))
        context.tracker.register(name, resource)


@given("I register the resource with the tracker")
def step_register_resource(context):
    context.tracker.register(context.current_resource_name, context.current_resource)


@given("I have an async task tracked by the bridge")
def step_create_bridge_task(context):
    from cleveragents.langgraph.bridge import RxPyLangGraphBridge
    from cleveragents.reactive.stream_router import ReactiveStreamRouter

    router = ReactiveStreamRouter()
    context.bridge = RxPyLangGraphBridge(router)
    context.cancellation_reasons = {}

    async def _long_running():
        try:
            await asyncio.sleep(100)
        except asyncio.CancelledError:
            return

    loop = asyncio.new_event_loop()

    async def _setup():
        task = asyncio.create_task(_long_running())
        context.bridge._active_tasks.add(task)
        context.bridge_task = task

    loop.run_until_complete(_setup())
    context.bridge_loop = loop


@given("I have a state manager with a temporary checkpoint directory")
def step_create_state_manager(context):
    context.temp_dir = tempfile.mkdtemp()
    context.checkpoint_dir = Path(context.temp_dir)
    context.state_manager = StateManager(
        initial_state=GraphState(),
        checkpoint_dir=context.checkpoint_dir,
    )


@given("I have an A2A event queue with {count:d} active subscriptions")
def step_create_event_queue(context, count):
    context.event_queue = A2aEventQueue()
    context.subscription_ids = []
    for _ in range(count):
        sub_id = context.event_queue.subscribe_local(lambda _evt: None)
        context.subscription_ids.append(sub_id)


@given("I have a bridge with {count:d} slow async tasks")
def step_create_bridge_with_slow_tasks(context, count):
    from cleveragents.langgraph.bridge import RxPyLangGraphBridge
    from cleveragents.reactive.stream_router import ReactiveStreamRouter

    router = ReactiveStreamRouter()
    context.bridge = RxPyLangGraphBridge(router)

    loop = asyncio.new_event_loop()

    async def _setup():
        for _ in range(count):

            async def _slow():
                try:
                    await asyncio.sleep(100)
                except asyncio.CancelledError:
                    return

            task = asyncio.create_task(_slow())
            context.bridge._active_tasks.add(task)

    loop.run_until_complete(_setup())
    context.bridge_loop = loop


# ---------------------------------------------------------------------------
# Whens
# ---------------------------------------------------------------------------


@when("I register the resource with the tracker")
def step_when_register_resource(context):
    context.tracker.register(context.current_resource_name, context.current_resource)


@when("I close all tracked resources")
def step_close_all(context):
    loop = asyncio.new_event_loop()
    loop.run_until_complete(context.tracker.close_all())
    loop.close()


@when("I close all tracked resources with a {timeout:g} second timeout")
def step_close_all_with_timeout(context, timeout):
    loop = asyncio.new_event_loop()
    loop.run_until_complete(context.tracker.close_all(timeout=timeout))
    loop.close()


@when("the tracker finalizer runs without close_all")
def step_run_finalizer(context):
    context.tracker._warn_unclosed()


@when("I request graceful cleanup with a {timeout:g} second timeout")
def step_graceful_bridge_cleanup(context, timeout):
    loop = context.bridge_loop

    async def _cleanup():
        await context.bridge.cleanup_tasks_async(timeout=timeout)

    loop.run_until_complete(_cleanup())


@when('I cancel the task with reason "{reason}"')
def step_cancel_with_reason(context, reason):
    loop = context.bridge_loop

    async def _cancel():
        await context.bridge.cancel_task_with_reason(context.bridge_task, reason)

    loop.run_until_complete(_cancel())
    context.cancel_reason = reason


@when("I save a checkpoint and then close the state manager")
def step_save_and_close_state_manager(context):
    context.state_manager.update_state(
        {"messages": [{"role": "user", "content": "test"}]}
    )
    context.state_manager._save_checkpoint()
    context.state_manager.close()


@when("I close the event queue")
def step_close_event_queue(context):
    context.event_queue.close()


@when('I try to register another resource named "{name}"')
def step_try_register_duplicate(context, name):
    try:
        dup_resource = MockAsyncResource(name)
        context.tracker.register(name, dup_resource)
        context.lsp_error = None
    except ValueError as exc:
        context.lsp_error = exc


@when("I try to register a resource with an empty name")
def step_try_register_empty_name(context):
    try:
        context.tracker.register("", MockAsyncResource("empty"))
        context.lsp_error = None
    except ValueError as exc:
        context.lsp_error = exc


@when('I try to register a None resource with name "{name}"')
def step_try_register_none_resource(context, name):
    try:
        context.tracker.register(name, None)  # type: ignore[arg-type]
        context.lsp_error = None
    except ValueError as exc:
        context.lsp_error = exc


@when("I use the tracker as an async context manager and register the resource")
def step_use_context_manager(context):
    loop = asyncio.new_event_loop()

    async def _run():
        async with context.tracker as tracker:
            tracker.register(context.current_resource_name, context.current_resource)

    loop.run_until_complete(_run())
    loop.close()


@when("I close all tracked resources twice")
def step_close_all_twice(context):
    loop = asyncio.new_event_loop()
    loop.run_until_complete(context.tracker.close_all())
    loop.run_until_complete(context.tracker.close_all())
    loop.close()


@when("I run enhanced cleanup with a {timeout:g} second timeout")
def step_enhanced_bridge_cleanup(context, timeout):
    loop = context.bridge_loop

    async def _cleanup():
        await context.bridge.cleanup_tasks_async(timeout=timeout)

    loop.run_until_complete(_cleanup())


# ---------------------------------------------------------------------------
# Thens
# ---------------------------------------------------------------------------


@then('the resource "{name}" should be closed')
def step_assert_resource_closed(context, name):
    resource = context.resources[name]
    assert resource.closed, f"Resource '{name}' was not closed"


@then("the tracker should have zero open resources")
def step_assert_zero_open(context):
    assert context.tracker.open_count == 0, (
        f"Expected 0 open resources, got {context.tracker.open_count}"
    )


@then("all {count:d} resources should be closed")
def step_assert_all_closed(context, count):
    closed_count = sum(1 for _, r in context.multi_resources if r.closed)
    assert closed_count == count, f"Expected {count} closed, got {closed_count}"


@then('a warning should be logged mentioning "{text}"')
def step_assert_warning_logged(context, text):
    warnings = [r for r in context.log_handler.records if r.levelno >= logging.WARNING]
    matching = [r for r in warnings if text in r.getMessage()]
    assert matching, (
        f"No warning mentioning '{text}' found. "
        f"Warnings: {[r.getMessage() for r in warnings]}"
    )


@then("the task should be cancelled")
def step_assert_task_cancelled(context):
    task = context.bridge_task
    assert task.done(), "Task should be done after cleanup"


@then("the bridge should have no active tasks")
def step_assert_no_active_tasks(context):
    assert len(context.bridge._active_tasks) == 0, (
        f"Expected 0 active tasks, got {len(context.bridge._active_tasks)}"
    )


@then("a warning should be logged about forced termination")
def step_assert_forced_termination_warning(context):
    warnings = [r for r in context.log_handler.records if r.levelno >= logging.WARNING]
    matching = [
        r
        for r in warnings
        if "timeout" in r.getMessage().lower() or "forced" in r.getMessage().lower()
    ]
    assert matching, (
        f"No forced termination warning found. "
        f"Warnings: {[r.getMessage() for r in warnings]}"
    )


@then("the tracker should report the timed-out resource")
def step_assert_timeout_reported(context):
    assert len(context.tracker.timed_out_resources) > 0, (
        "Expected at least one timed-out resource"
    )


@then('the cancellation reason should be "{reason}"')
def step_assert_cancellation_reason(context, reason):
    assert context.bridge.cancellation_reasons.get(context.bridge_task) == reason, (
        f"Expected reason '{reason}', "
        f"got '{context.bridge.cancellation_reasons.get(context.bridge_task)}'"
    )


@then("the checkpoint file should exist and be readable")
def step_assert_checkpoint_exists(context):
    checkpoints = list(context.checkpoint_dir.glob("checkpoint_*.json"))
    assert len(checkpoints) >= 1, "No checkpoint files found"
    content = checkpoints[0].read_text(encoding="utf-8")
    assert len(content) > 0, "Checkpoint file is empty"


@then("the state manager should be marked as closed")
def step_assert_state_manager_closed(context):
    assert context.state_manager.is_closed, "StateManager should be closed"


@then("all subscriptions should be removed")
def step_assert_subscriptions_removed(context):
    assert len(context.event_queue._subscriptions) == 0, (
        f"Expected 0 subscriptions, got {len(context.event_queue._subscriptions)}"
    )


@then("the subscription count should be zero")
def step_assert_subscription_count_zero(context):
    assert len(context.event_queue._subscriptions) == 0


@then('the resource "{name}" should be closed after exiting the context')
def step_assert_closed_after_context(context, name):
    resource = context.resources[name]
    assert resource.closed, f"Resource '{name}' should be closed after context exit"


@then('the resource "{name}" should be closed exactly once')
def step_assert_closed_once(context, name):
    resource = context.resources[name]
    assert resource.closed, f"Resource '{name}' was not closed"
    assert resource.close_count == 1, (
        f"Resource '{name}' was closed {resource.close_count} times, expected 1"
    )


@then("completed tasks should be logged")
def step_assert_tasks_logged(context):
    infos = [r for r in context.log_handler.records if r.levelno >= logging.INFO]
    has_cleanup_log = any(
        "cleanup" in r.getMessage().lower() or "task" in r.getMessage().lower()
        for r in infos
    )
    # The cleanup operation should have been logged, and tasks cleared
    assert has_cleanup_log or len(context.bridge._active_tasks) == 0


# ---------------------------------------------------------------------------
# Logging capture helper
# ---------------------------------------------------------------------------


class _CapturingHandler(logging.Handler):
    """Captures log records for assertion in tests."""

    def __init__(self) -> None:
        super().__init__()
        self.records: list[logging.LogRecord] = []

    def emit(self, record: logging.LogRecord) -> None:
        self.records.append(record)


# ---------------------------------------------------------------------------
# Steps for review-feedback scenarios (T1-T4)
# ---------------------------------------------------------------------------


@when('I try to register a resource named "{name}" after close_all')
def step_try_register_after_close(context, name):
    try:
        context.tracker.register(name, MockAsyncResource(name))
        context.runtime_error = None
    except RuntimeError as exc:
        context.runtime_error = exc


@then('a RuntimeError should be raised mentioning "{text}"')
def step_assert_runtime_error(context, text):
    assert context.runtime_error is not None, (
        "Expected RuntimeError but none was raised"
    )
    assert text in str(context.runtime_error), (
        f"Expected '{text}' in error message, got: {context.runtime_error}"
    )


@when("I close the state manager")
def step_close_state_manager(context):
    context.state_manager.close()


@when("I try to update state after close")
def step_try_update_after_close(context):
    try:
        context.state_manager.update_state(
            {"messages": [{"role": "user", "content": "x"}]}
        )
        context.runtime_error = None
    except RuntimeError as exc:
        context.runtime_error = exc


@when("I try to reset state after close")
def step_try_reset_after_close(context):
    try:
        context.state_manager.reset()
        context.runtime_error = None
    except RuntimeError as exc:
        context.runtime_error = exc


@when("I try to publish an event after close")
def step_try_publish_after_close(context):
    from cleveragents.a2a.models import A2aEvent

    try:
        context.event_queue.publish(A2aEvent(event_type="test", data={}))
        context.runtime_error = None
    except RuntimeError as exc:
        context.runtime_error = exc


@then("the event queue is_closed should be False")
def step_assert_queue_not_closed(context):
    assert not context.event_queue.is_closed, "Expected is_closed=False"


@then("the event queue is_closed should be True")
def step_assert_queue_closed(context):
    assert context.event_queue.is_closed, "Expected is_closed=True"