cleveragents-core/features/steps/bridge_coverage_boost_steps.py

"""Steps targeting remaining uncovered lines in bridge.py.

Uncovered lines addressed:
- Line 64:    early return from cleanup_tasks_async (empty _active_tasks)
- Lines 79-83: pending tasks that exceed the timeout in cleanup_tasks_async
- Lines 89-93: late tasks injected into _active_tasks during the await window
- Line 149:   MESSAGE_ROUTER node with "rules" key in create_graph_from_config
- Lines 282-290: execute_node inner coroutine in _create_node_operator
"""

from __future__ import annotations

import asyncio
import contextlib
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch

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

from cleveragents.langgraph.bridge import RxPyLangGraphBridge
from cleveragents.langgraph.nodes import NodeType
from cleveragents.langgraph.state import GraphState
from cleveragents.reactive.stream_router import (
    ReactiveStreamRouter,
    StreamMessage,
)

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


def _fresh_bridge() -> tuple[RxPyLangGraphBridge, MagicMock]:
    """Return a bridge backed by a mocked stream-router."""
    scheduler = MagicMock()
    router = ReactiveStreamRouter(scheduler=scheduler)
    router.agents = {"agent": MagicMock()}
    bridge = RxPyLangGraphBridge(router)
    return bridge, scheduler


def _make_loop() -> asyncio.AbstractEventLoop:
    """Create and install a fresh event loop."""
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    return loop


def _teardown_loop(loop: asyncio.AbstractEventLoop) -> None:
    """Cancel all pending tasks, close loop, and install a fresh one."""
    # Cancel any lingering tasks
    pending = asyncio.all_tasks(loop) if hasattr(asyncio, "all_tasks") else set()
    for t in pending:
        t.cancel()
    if pending:
        loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
    if not loop.is_closed():
        loop.close()
    asyncio.set_event_loop(asyncio.new_event_loop())


# ===================================================================
# Scenario: cleanup_tasks_async returns immediately when no tasks
# (targets line 64)
# ===================================================================


@given("a bridge with an empty active tasks set")
def step_bridge_empty_tasks(context: Context) -> None:
    bridge, _ = _fresh_bridge()
    # Ensure there are no active tasks
    bridge._active_tasks.clear()
    context.bridge = bridge
    context.async_error = None
    context._loop = _make_loop()

    def _cleanup() -> None:
        _teardown_loop(context._loop)

    context._cleanup_handlers.append(_cleanup)


@when("I call cleanup_tasks_async on the idle bridge")
def step_call_cleanup_async_idle(context: Context) -> None:
    try:
        context._loop.run_until_complete(context.bridge.cleanup_tasks_async())
    except Exception as exc:
        context.async_error = exc


@then("the async cleanup should complete without error")
def step_assert_no_async_error(context: Context) -> None:
    assert context.async_error is None, (
        f"Expected no error, got {context.async_error!r}"
    )


@then("the bridge active task set should remain empty")
def step_assert_tasks_still_empty(context: Context) -> None:
    assert len(context.bridge._active_tasks) == 0, (
        f"Expected empty _active_tasks, got {len(context.bridge._active_tasks)}"
    )


# ===================================================================
# Scenario: cleanup_tasks_async logs pending tasks that exceed timeout
# (targets lines 79-83)
# ===================================================================


@given("a bridge with a task that ignores cancellation")
def step_bridge_with_uncancellable_task(context: Context) -> None:
    bridge, _ = _fresh_bridge()
    loop = _make_loop()

    async def _stubborn() -> None:
        """Coroutine that catches CancelledError and keeps sleeping,
        so it remains in the ``pending`` set after asyncio.wait()."""
        try:
            await asyncio.sleep(99999)
        except asyncio.CancelledError:
            # Swallow the first cancellation and keep running
            with contextlib.suppress(asyncio.CancelledError):
                await asyncio.sleep(99999)

    task = loop.create_task(_stubborn(), name="stubborn_task")
    bridge._active_tasks.add(task)

    context.bridge = bridge
    context.stubborn_task = task
    context._loop = loop

    def _cleanup() -> None:
        task.cancel()
        with contextlib.suppress(Exception):
            loop.run_until_complete(asyncio.gather(task, return_exceptions=True))
        _teardown_loop(loop)

    context._cleanup_handlers.append(_cleanup)


@when("I call cleanup_tasks_async with a very short timeout")
def step_call_cleanup_async_short_timeout(context: Context) -> None:
    # Use an extremely short timeout so the stubborn task ends up in `pending`
    context._loop.run_until_complete(context.bridge.cleanup_tasks_async(timeout=0.001))


@then("pending tasks should be re-cancelled after the timeout")
def step_assert_pending_re_cancelled(context: Context) -> None:
    task = context.stubborn_task
    # After cleanup, the task should have been cancelled (a second time)
    assert task.done() or task.cancelled(), (
        f"Expected task to be done/cancelled, got state={task._state}"
    )


@then("the bridge active task set should be cleared after timeout cleanup")
def step_assert_cleared_after_timeout(context: Context) -> None:
    assert len(context.bridge._active_tasks) == 0, (
        f"Expected empty _active_tasks after timeout cleanup, "
        f"got {len(context.bridge._active_tasks)}"
    )


# ===================================================================
# Scenario: cleanup_tasks_async cancels late tasks added during await
# (targets lines 89-93)
# ===================================================================


@given("a bridge with a task whose done-callback spawns a late task")
def step_bridge_with_late_task_spawner(context: Context) -> None:
    bridge, _ = _fresh_bridge()
    loop = _make_loop()

    async def _quick() -> None:
        """Completes quickly so its done-callback fires during cleanup."""
        await asyncio.sleep(0)

    async def _late_sleeper() -> None:
        """Late task that should be found and cancelled by cleanup."""
        await asyncio.sleep(99999)

    task = loop.create_task(_quick(), name="quick_task")

    def _on_done(_t: asyncio.Task[Any]) -> None:
        """Done callback: inject a late task into bridge._active_tasks."""
        late = loop.create_task(_late_sleeper(), name="late_task")
        bridge._active_tasks.add(late)
        context.late_task = late

    task.add_done_callback(_on_done)
    bridge._active_tasks.add(task)

    context.bridge = bridge
    context.spawner_task = task
    context._loop = loop

    def _cleanup() -> None:
        for t in list(bridge._active_tasks):
            t.cancel()
        try:
            all_tasks = [t for t in asyncio.all_tasks(loop) if not t.done()]
            if all_tasks:
                loop.run_until_complete(
                    asyncio.gather(*all_tasks, return_exceptions=True)
                )
        except Exception:
            pass
        _teardown_loop(loop)

    context._cleanup_handlers.append(_cleanup)


@when("I call cleanup_tasks_async allowing callback to fire")
def step_call_cleanup_async_late(context: Context) -> None:
    context._loop.run_until_complete(context.bridge.cleanup_tasks_async(timeout=0.5))


@then("the late task should be cancelled during cleanup")
def step_assert_late_task_cancelled(context: Context) -> None:
    late = getattr(context, "late_task", None)
    assert late is not None, "Late task was never created by the done-callback"
    assert late.done() or late.cancelled(), (
        f"Expected late task to be done/cancelled, got state={late._state}"
    )


@then("the bridge active task set should be cleared after late-task cleanup")
def step_assert_cleared_after_late(context: Context) -> None:
    assert len(context.bridge._active_tasks) == 0, (
        f"Expected empty _active_tasks after late-task cleanup, "
        f"got {len(context.bridge._active_tasks)}"
    )


# ===================================================================
# Scenario: create_graph_from_config with MESSAGE_ROUTER + rules
# (targets line 149)
# ===================================================================


@given("a bridge ready to create graphs from config")
def step_bridge_for_config_creation(context: Context) -> None:
    bridge, _ = _fresh_bridge()
    context.bridge = bridge


@when("I create a graph with a MESSAGE_ROUTER node containing rules")
def step_create_graph_with_message_router_rules(context: Context) -> None:
    routing_rules = [
        {"condition": {"equals": "hello"}, "target": "greeter"},
        {"condition": {"field": "type", "value": "query"}, "target": "searcher"},
    ]
    config: dict[str, Any] = {
        "name": "router_test_graph",
        "entry_point": "start",
        "nodes": {
            "my_router": {
                "type": "message_router",
                "rules": routing_rules,
                "metadata": {"custom_key": "custom_value"},
            },
        },
        "edges": [
            {"source": "start", "target": "my_router"},
        ],
    }
    with patch("cleveragents.langgraph.bridge.LangGraph") as mock_lg_cls:
        mock_graph = MagicMock()
        mock_graph.name = "router_test_graph"
        mock_lg_cls.return_value = mock_graph

        context.created_graph = context.bridge.create_graph_from_config(config)

    # Capture the GraphConfig that was passed to the LangGraph constructor
    call_kwargs = mock_lg_cls.call_args
    context.passed_graph_config = call_kwargs.kwargs.get("config")
    context.routing_rules = routing_rules


@then("the node config metadata should contain the routing rules")
def step_assert_rules_in_metadata(context: Context) -> None:
    gc = context.passed_graph_config
    assert gc is not None, "GraphConfig was not captured"
    node_cfg = gc.nodes.get("my_router")
    assert node_cfg is not None, (
        f"Expected 'my_router' in nodes, got: {list(gc.nodes.keys())}"
    )
    assert "rules" in node_cfg.metadata, (
        f"Expected 'rules' in metadata, got keys: {list(node_cfg.metadata.keys())}"
    )
    assert node_cfg.metadata["rules"] == context.routing_rules, (
        f"Rules mismatch: {node_cfg.metadata['rules']} != {context.routing_rules}"
    )
    # Also verify the custom metadata key was preserved
    assert node_cfg.metadata.get("custom_key") == "custom_value", (
        f"Expected custom_key in metadata, got: {node_cfg.metadata}"
    )
    assert node_cfg.type == NodeType.MESSAGE_ROUTER, (
        f"Expected MESSAGE_ROUTER type, got: {node_cfg.type}"
    )


# ===================================================================
# Scenario: Node operator executes the inner coroutine and updates state
# (targets lines 282-290)
# ===================================================================


@given("a bridge with a graph containing a mock executable node")
def step_bridge_with_executable_node(context: Context) -> None:
    bridge, _ = _fresh_bridge()
    loop = _make_loop()

    # Build a mock graph with a node that has an async execute method
    mock_graph = MagicMock()
    mock_graph.name = "node_exec_graph"

    # State manager with controllable get_state / update_state
    mock_state = MagicMock(spec=GraphState)
    mock_state.messages = [{"content": "initial"}]
    mock_state.to_dict = MagicMock(return_value={"messages": mock_state.messages})

    mock_state_manager = MagicMock()
    mock_state_manager.get_state = MagicMock(return_value=mock_state)
    mock_state_manager.update_state = MagicMock()
    mock_graph.state_manager = mock_state_manager

    # Create a mock node whose execute returns state updates
    mock_node = MagicMock()
    mock_node.execute = AsyncMock(
        return_value={"messages": [{"content": "node_output"}]}
    )

    mock_graph.nodes = {"test_worker": mock_node}
    bridge.graphs["node_exec_graph"] = mock_graph

    context.bridge = bridge
    context.mock_node = mock_node
    context.mock_state_manager = mock_state_manager
    context.mock_graph = mock_graph
    context._loop = loop

    def _cleanup() -> None:
        for t in list(bridge._active_tasks):
            t.cancel()
        try:
            remaining = [t for t in bridge._active_tasks if not t.done()]
            if remaining:
                loop.run_until_complete(
                    asyncio.gather(*remaining, return_exceptions=True)
                )
        except Exception:
            pass
        _teardown_loop(loop)

    context._cleanup_handlers.append(_cleanup)


@when("I pipe a message through the node operator and await the result")
def step_pipe_message_through_node_operator(context: Context) -> None:
    import rx  # type: ignore

    bridge = context.bridge
    loop = context._loop

    results: list[Any] = []
    errors: list[Any] = []
    msg = StreamMessage(
        content="input_payload",
        metadata={"origin": "node_test"},
    )

    async def _run_in_loop() -> None:
        """Subscribe inside the running loop so get_running_loop() succeeds."""
        node_op = bridge._create_node_operator(
            {"graph": "node_exec_graph", "node": "test_worker"}
        )
        rx.just(msg).pipe(node_op).subscribe(
            on_next=lambda x: results.append(x),
            on_error=lambda e: errors.append(e),
        )
        # Give the scheduled coroutine time to complete
        await asyncio.sleep(0.05)
        # Also gather any remaining tasks
        pending = [t for t in bridge._active_tasks if not t.done()]
        if pending:
            await asyncio.gather(*pending, return_exceptions=True)

    loop.run_until_complete(_run_in_loop())

    context.node_op_results = results
    context.node_op_errors = errors


@then("the node execute method should have been called")
def step_assert_node_execute_called(context: Context) -> None:
    assert len(context.node_op_errors) == 0, (
        f"Node operator raised errors: {context.node_op_errors}"
    )
    context.mock_node.execute.assert_called_once()


@then("the state manager should have been updated with the node result")
def step_assert_state_manager_updated(context: Context) -> None:
    context.mock_state_manager.update_state.assert_called_once_with(
        {"messages": [{"content": "node_output"}]},
        node_id="test_worker",
    )


@then("the returned message metadata should contain the node and graph names")
def step_assert_node_metadata(context: Context) -> None:
    assert len(context.node_op_results) > 0, "Node operator produced no results"
    result_msg = context.node_op_results[0]
    assert isinstance(result_msg, StreamMessage), (
        f"Expected StreamMessage, got {type(result_msg).__name__}"
    )
    assert result_msg.metadata.get("node") == "test_worker", (
        f"Expected node='test_worker' in metadata, got: {result_msg.metadata}"
    )
    assert result_msg.metadata.get("graph") == "node_exec_graph", (
        f"Expected graph='node_exec_graph' in metadata, got: {result_msg.metadata}"
    )
    assert "state" in result_msg.metadata, (
        f"Expected 'state' key in metadata, got keys: {list(result_msg.metadata.keys())}"
    )