alishahryar1--free-claude-code/tests/messaging/test_tree_processor.py

"""Manager-level task and cancellation ownership tests."""

import asyncio
import logging

import pytest

from free_claude_code.messaging.models import IncomingMessage, MessageScope
from free_claude_code.messaging.trees import (
    CancellationReason,
    CancellationUiOwner,
    FailureResult,
    MessageState,
    NodeClaim,
    QueueEntry,
    TreeQueueManager,
)
from free_claude_code.messaging.trees import manager as manager_module
from free_claude_code.messaging.trees import processor as processor_module
from free_claude_code.messaging.trees.node import MessageNode
from free_claude_code.messaging.trees.processor import TreeQueueProcessor
from free_claude_code.messaging.trees.runtime import MessageTree

_SCOPE = MessageScope(platform="telegram", chat_id="chat")


def _incoming(node_id: str, *, reply_to: str | None = None) -> IncomingMessage:
    return IncomingMessage(
        text=f"prompt {node_id}",
        chat_id=_SCOPE.chat_id,
        user_id="user",
        message_id=node_id,
        platform=_SCOPE.platform,
        reply_to_message_id=reply_to,
    )


async def _wait_for_no_tasks(manager: TreeQueueManager) -> None:
    """Yield deterministic ready-queue checkpoints until task cleanup completes."""
    loop = asyncio.get_running_loop()
    for _ in range(20):
        if manager.task_count() == 0:
            return
        checkpoint = asyncio.Event()
        loop.call_soon(checkpoint.set)
        await checkpoint.wait()
    assert manager.task_count() == 0


@pytest.mark.asyncio
async def test_active_cancel_returns_runner_owned_effect_and_terminal_snapshot() -> (
    None
):
    started = asyncio.Event()

    async def process(_claim: NodeClaim) -> None:
        started.set()
        await asyncio.Event().wait()

    manager = TreeQueueManager(process)
    await manager.admit(_incoming("root"), "status-root")
    await started.wait()

    result = await manager.cancel_node(
        _SCOPE,
        "root",
        reason=CancellationReason.STOP,
    )

    assert [(effect.node.node_id, effect.ui_owner) for effect in result.effects] == [
        ("root", CancellationUiOwner.RUNNER)
    ]
    assert len(result.snapshots) == 1
    assert result.snapshots[0].nodes["root"]["state"] == "error"
    view = await manager.get_node(_SCOPE, "root")
    assert view is not None and view.state is MessageState.ERROR
    assert manager.task_count() == 0


@pytest.mark.asyncio
async def test_queued_cancel_returns_workflow_effect_and_exact_queue_update() -> None:
    release_root = asyncio.Event()
    root_started = asyncio.Event()
    child_started = asyncio.Event()
    queue_updates: list[tuple[tuple[str, int], ...]] = []

    async def process(claim: NodeClaim) -> None:
        if claim.node.node_id == "root":
            root_started.set()
            await release_root.wait()
        else:
            child_started.set()

    async def capture_queue(queue: tuple[QueueEntry, ...]) -> None:
        queue_updates.append(
            tuple((entry.node.node_id, entry.position) for entry in queue)
        )

    manager = TreeQueueManager(process, queue_update_callback=capture_queue)
    await manager.admit(_incoming("root"), "status-root")
    await root_started.wait()
    decision = await manager.admit(
        _incoming("child", reply_to="root"),
        "status-child",
        parent_reference_id="root",
    )
    assert decision.position == 1

    result = await manager.cancel_node(
        _SCOPE,
        "child",
        reason=CancellationReason.STOP,
    )

    assert [(effect.node.node_id, effect.ui_owner) for effect in result.effects] == [
        ("child", CancellationUiOwner.WORKFLOW)
    ]
    assert queue_updates == [()]
    assert result.snapshots[0].nodes["child"]["state"] == "error"
    assert child_started.is_set() is False

    release_root.set()
    await _wait_for_no_tasks(manager)
    assert child_started.is_set() is False


@pytest.mark.asyncio
async def test_cancel_cleanup_timeout_is_bounded_and_task_remains_owned(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(manager_module, "CANCEL_TASK_DRAIN_TIMEOUT_S", 0.01)
    started = asyncio.Event()
    cancellation_seen = asyncio.Event()
    release_cleanup = asyncio.Event()

    async def process(_claim: NodeClaim) -> None:
        started.set()
        try:
            await asyncio.Event().wait()
        except asyncio.CancelledError:
            cancellation_seen.set()
            await release_cleanup.wait()
            raise

    manager = TreeQueueManager(process)
    await manager.admit(_incoming("root"), "status-root")
    await started.wait()

    try:
        result = await asyncio.wait_for(
            manager.cancel_node(
                _SCOPE,
                "root",
                reason=CancellationReason.STOP,
            ),
            timeout=0.5,
        )
        await cancellation_seen.wait()
        assert result.effects[0].node.node_id == "root"
        assert manager.task_count() == 1
    finally:
        release_cleanup.set()

    await _wait_for_no_tasks(manager)


@pytest.mark.asyncio
async def test_escaped_processor_failure_persists_effects_through_manager_owner() -> (
    None
):
    started = asyncio.Event()
    release = asyncio.Event()
    failures: list[FailureResult] = []
    queue_updates: list[tuple[QueueEntry, ...]] = []

    async def process(claim: NodeClaim) -> None:
        if claim.node.node_id == "root":
            started.set()
            await release.wait()
            raise RuntimeError("processor boundary failed")

    async def capture_queue(queue: tuple[QueueEntry, ...]) -> None:
        queue_updates.append(queue)

    manager = TreeQueueManager(
        process,
        queue_update_callback=capture_queue,
        unexpected_failure_callback=failures.append,
    )
    await manager.admit(_incoming("root"), "status-root")
    await started.wait()
    await manager.admit(
        _incoming("child", reply_to="root"),
        "status-child",
        parent_reference_id="root",
    )

    release.set()
    await _wait_for_no_tasks(manager)

    assert len(failures) == 1
    failure = failures[0]
    assert failure.snapshot is not None
    assert {target.node_id for target in failure.affected} == {"root", "child"}
    assert failure.snapshot.nodes["root"]["state"] == "error"
    assert failure.snapshot.nodes["child"]["state"] == "error"
    assert queue_updates == [()]


@pytest.mark.asyncio
async def test_wait_idle_spans_successor_publication_and_completion() -> None:
    root_started = asyncio.Event()
    release_root = asyncio.Event()
    child_started = asyncio.Event()
    release_child = asyncio.Event()

    async def process(claim: NodeClaim) -> None:
        if claim.node.node_id == "root":
            root_started.set()
            await release_root.wait()
            return
        child_started.set()
        await release_child.wait()

    manager = TreeQueueManager(process)
    await asyncio.wait_for(manager.wait_idle(), timeout=0.1)
    await manager.admit(_incoming("root"), "status-root")
    await asyncio.wait_for(root_started.wait(), timeout=1)
    await manager.admit(
        _incoming("child", reply_to="root"),
        "status-child",
        parent_reference_id="root",
    )
    idle_task = asyncio.create_task(manager.wait_idle())

    try:
        await asyncio.sleep(0)
        assert not idle_task.done()

        release_root.set()
        await asyncio.wait_for(child_started.wait(), timeout=1)
        assert not idle_task.done()

        release_child.set()
        await asyncio.wait_for(idle_task, timeout=1)
        assert manager.task_count() == 0
    finally:
        release_root.set()
        release_child.set()
        if not idle_task.done():
            idle_task.cancel()
            with pytest.raises(asyncio.CancelledError):
                await idle_task


@pytest.mark.asyncio
async def test_wait_idle_spans_pre_run_cancellation_recovery() -> None:
    finish_started = asyncio.Event()
    release_finish = asyncio.Event()

    async def process(_claim: NodeClaim) -> None:
        raise AssertionError("pre-run cancellation must not enter the processor")

    async def fail_claim(_claim: NodeClaim) -> None:
        raise AssertionError("cancellation must not fail the claim")

    async def finish_claim(_tree: MessageTree, _claim: NodeClaim) -> None:
        finish_started.set()
        await release_finish.wait()

    tree = MessageTree(
        MessageNode(
            node_id="root",
            scope=_SCOPE,
            prompt="prompt root",
            status_message_id="status-root",
        )
    )
    decision = await tree.enqueue_or_claim("root")
    assert decision.claim is not None
    processor = TreeQueueProcessor(
        process,
        claim_failure_callback=fail_claim,
        claim_finished_callback=finish_claim,
    )
    processor.launch(tree, decision.claim)
    cancelled = processor.cancel(decision.claim, CancellationReason.STOP)
    assert cancelled is not None
    assert cancelled.runner_started is False
    idle_task = asyncio.create_task(processor.wait_idle())

    try:
        await asyncio.wait_for(finish_started.wait(), timeout=1)
        assert not idle_task.done()

        release_finish.set()
        await asyncio.wait_for(cancelled.task, timeout=1)
        await asyncio.wait_for(idle_task, timeout=1)
        assert processor.task_count() == 0
    finally:
        release_finish.set()
        if not idle_task.done():
            idle_task.cancel()
            with pytest.raises(asyncio.CancelledError):
                await idle_task


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("log_messaging_error_details", "secret_is_logged"),
    [(False, False), (True, True)],
)
async def test_wait_idle_surfaces_finish_failure_without_leaking_task_owner(
    caplog: pytest.LogCaptureFixture,
    log_messaging_error_details: bool,
    secret_is_logged: bool,
) -> None:
    secret = "unique-finish-callback-secret"
    finish_error = RuntimeError(secret)
    finish_calls = 0

    async def process(_claim: NodeClaim) -> None:
        return

    async def fail_claim(_claim: NodeClaim) -> None:
        raise AssertionError("successful processing must not fail the claim")

    async def finish_claim(_tree: MessageTree, _claim: NodeClaim) -> None:
        nonlocal finish_calls
        finish_calls += 1
        raise finish_error

    tree = MessageTree(
        MessageNode(
            node_id="root",
            scope=_SCOPE,
            prompt="prompt root",
            status_message_id="status-root",
        )
    )
    decision = await tree.enqueue_or_claim("root")
    assert decision.claim is not None
    processor = TreeQueueProcessor(
        process,
        claim_failure_callback=fail_claim,
        claim_finished_callback=finish_claim,
        log_messaging_error_details=log_messaging_error_details,
    )

    with caplog.at_level(logging.ERROR):
        processor.launch(tree, decision.claim)
        with pytest.raises(RuntimeError) as raised:
            await asyncio.wait_for(processor.wait_idle(), timeout=1)

    assert raised.value is finish_error
    assert finish_calls == 1
    assert processor.task_count() == 0
    await asyncio.wait_for(processor.wait_idle(), timeout=0.1)

    messages = "\n".join(record.getMessage() for record in caplog.records)
    assert "Claim completion callback failed for node root" in messages
    assert "RuntimeError" in messages
    assert (secret in messages) is secret_is_logged


@pytest.mark.asyncio
async def test_failed_launch_rolls_idle_state_back(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    async def process(_claim: NodeClaim) -> None:
        return

    async def fail_claim(_claim: NodeClaim) -> None:
        return

    async def finish_claim(_tree: MessageTree, _claim: NodeClaim) -> None:
        return

    tree = MessageTree(
        MessageNode(
            node_id="root",
            scope=_SCOPE,
            prompt="prompt root",
            status_message_id="status-root",
        )
    )
    decision = await tree.enqueue_or_claim("root")
    assert decision.claim is not None
    processor = TreeQueueProcessor(
        process,
        claim_failure_callback=fail_claim,
        claim_finished_callback=finish_claim,
    )

    def fail_create_task(*_args: object, **_kwargs: object) -> None:
        raise RuntimeError("launch failed")

    monkeypatch.setattr(
        processor_module.asyncio,
        "create_task",
        fail_create_task,
    )

    with pytest.raises(RuntimeError, match="launch failed"):
        processor.launch(tree, decision.claim)

    assert processor.task_count() == 0
    await asyncio.wait_for(processor.wait_idle(), timeout=0.1)


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("log_messaging_error_details", "secret_is_logged"),
    [(False, False), (True, True)],
)
async def test_processor_failure_logging_respects_diagnostic_policy(
    caplog: pytest.LogCaptureFixture,
    log_messaging_error_details: bool,
    secret_is_logged: bool,
) -> None:
    secret = "unique-processor-exception-secret"

    async def process(_claim: NodeClaim) -> None:
        raise RuntimeError(secret)

    manager = TreeQueueManager(
        process,
        log_messaging_error_details=log_messaging_error_details,
    )
    with caplog.at_level(logging.ERROR):
        await manager.admit(_incoming("root"), "status-root")
        await asyncio.wait_for(manager.wait_idle(), timeout=1)

    messages = "\n".join(record.getMessage() for record in caplog.records)

    assert "Error processing node root" in messages
    assert "RuntimeError" in messages
    assert (secret in messages) is secret_is_logged