import asyncio from dataclasses import FrozenInstanceError import pytest from free_claude_code.messaging.models import MessageScope from free_claude_code.messaging.trees.node import MessageNode, MessageState from free_claude_code.messaging.trees.runtime import MessageTree _SCOPE = MessageScope(platform="telegram", chat_id="chat") def _tree() -> MessageTree: return MessageTree( MessageNode( node_id="root", scope=_SCOPE, prompt="prompt root", status_message_id="status-root", ) ) async def _add( tree: MessageTree, node_id: str, status_message_id: str, parent_id: str = "root", parent_reference_id: str | None = None, ): return await tree.add_and_enqueue( node_id, _SCOPE, f"prompt {node_id}", status_message_id, parent_id, parent_reference_id or parent_id, ) @pytest.mark.asyncio async def test_enqueue_finish_and_claim_next_are_atomic_and_fifo() -> None: tree = _tree() root = await tree.enqueue_or_claim("root") child_a = await _add(tree, "child-a", "status-a") child_b = await _add(tree, "child-b", "status-b") assert root.accepted and root.claim is not None and root.position is None assert child_a.claim is None and child_a.position == 1 assert child_b.claim is None and child_b.position == 2 completion = await tree.finish_and_claim_next(root.claim.claim_id) assert completion.next_claim is not None assert completion.next_claim.node.node_id == "child-a" assert [(entry.node.node_id, entry.position) for entry in completion.queue] == [ ("child-b", 1) ] @pytest.mark.asyncio async def test_duplicate_unknown_and_terminal_admission_are_rejected_without_wedge() -> ( None ): tree = _tree() first = await tree.enqueue_or_claim("root") duplicate = await tree.enqueue_or_claim("root") unknown = await tree.enqueue_or_claim("status-root") assert first.claim is not None assert duplicate.accepted is False assert unknown.accepted is False assert duplicate.snapshot is None assert unknown.snapshot is None await tree.complete_claim(first.claim.claim_id, "session") await tree.finish_and_claim_next(first.claim.claim_id) terminal = await tree.enqueue_or_claim("root") assert terminal.accepted is False assert terminal.snapshot is None @pytest.mark.asyncio async def test_stale_claim_id_cannot_clear_a_newer_claim() -> None: tree = _tree() first = await tree.enqueue_or_claim("root") assert first.claim is not None await tree.finish_and_claim_next(first.claim.claim_id) second = await _add(tree, "child", "status-child") assert second.claim is not None stale = await tree.finish_and_claim_next(first.claim.claim_id) assert stale.next_claim is None assert await tree.record_session(second.claim.claim_id, "session") is not None await tree.finish_and_claim_next(second.claim.claim_id) third = await _add(tree, "third", "status-third") assert third.claim is not None @pytest.mark.asyncio async def test_cancellation_keeps_active_identity_until_matching_finish() -> None: tree = _tree() root = await tree.enqueue_or_claim("root") queued = await _add(tree, "queued", "status-queued") assert root.claim is not None and queued.position == 1 cancelled = await tree.cancel_node("root") assert cancelled.active_claim == root.claim assert [node.node_id for node in cancelled.nodes] == ["root"] completion = await tree.finish_and_claim_next(root.claim.claim_id) assert completion.next_claim is not None assert completion.next_claim.node.node_id == "queued" @pytest.mark.asyncio @pytest.mark.parametrize("propagate", [False, True]) async def test_cancelled_claim_rejects_late_failure_and_descendant_propagation( propagate: bool, ) -> None: tree = _tree() root = await tree.enqueue_or_claim("root") queued = await _add(tree, "queued", "status-queued") assert root.claim is not None and queued.position == 1 await tree.cancel_node("root") late_failure = await tree.fail_claim( root.claim.claim_id, propagate=propagate, ) assert late_failure.snapshot is None assert late_failure.affected == () assert late_failure.queue_update is None queued_view = await tree.node_view("queued") assert queued_view is not None assert queued_view.state is MessageState.PENDING completion = await tree.finish_and_claim_next(root.claim.claim_id) assert completion.next_claim is not None assert completion.next_claim.node.node_id == "queued" @pytest.mark.asyncio async def test_cancelled_queue_member_is_never_claimed() -> None: tree = _tree() root = await tree.enqueue_or_claim("root") await _add(tree, "a", "status-a") await _add(tree, "b", "status-b") assert root.claim is not None cancelled = await tree.cancel_node("a") assert [node.node_id for node in cancelled.nodes] == ["a"] completion = await tree.finish_and_claim_next(root.claim.claim_id) assert completion.next_claim is not None assert completion.next_claim.node.node_id == "b" @pytest.mark.asyncio async def test_prompt_subtree_removal_returns_every_platform_message() -> None: tree = _tree() await _add(tree, "child", "status-child") removed = await tree.remove_message_subtree("root") assert removed.removed_message_ids == frozenset( {"root", "status-root", "child", "status-child"} ) @pytest.mark.asyncio async def test_status_subtree_preserves_prompt_siblings_and_invalidates_resume() -> ( None ): tree = _tree() root = await tree.enqueue_or_claim("root") assert root.claim is not None await tree.record_session(root.claim.claim_id, "session-root") direct_sibling = await _add(tree, "sibling", "status-sibling") status_descendant = await _add( tree, "descendant", "status-descendant", parent_reference_id="status-root", ) assert direct_sibling.position == 1 assert status_descendant.position == 2 removed = await tree.remove_message_subtree("status-root") assert removed.removed_message_ids == frozenset( {"status-root", "descendant", "status-descendant"} ) root_view = await tree.node_view("root") sibling_view = await tree.node_view("sibling") assert root_view is not None assert root_view.state is MessageState.ERROR assert root_view.session_id is None assert sibling_view is not None assert sibling_view.state is MessageState.PENDING assert await tree.node_view("descendant") is None assert await tree.resolve_reply("status-root") is None completion = await tree.finish_and_claim_next(root.claim.claim_id) assert completion.next_claim is not None assert completion.next_claim.node.node_id == "sibling" assert completion.next_claim.parent_session_id is None @pytest.mark.asyncio async def test_concurrent_duplicate_enqueue_produces_one_claim_only() -> None: tree = _tree() gate = asyncio.Event() async def enqueue(): await gate.wait() return await tree.enqueue_or_claim("root") tasks = [asyncio.create_task(enqueue()) for _ in range(20)] gate.set() results = await asyncio.gather(*tasks) assert sum(result.claim is not None for result in results) == 1 assert sum(result.accepted for result in results) == 1 @pytest.mark.asyncio async def test_results_are_frozen_and_do_not_expose_mutable_nodes() -> None: tree = _tree() decision = await tree.enqueue_or_claim("root") assert decision.claim is not None with pytest.raises(FrozenInstanceError): decision.claim.__setattr__("claim_id", "replacement") assert not hasattr(decision.claim, "__dict__") assert not isinstance(decision.claim.node, MessageNode) def test_message_tree_has_no_lock_or_partial_mutation_escape_hatches() -> None: forbidden = { "with_lock", "enqueue", "dequeue", "put_queue_unlocked", "remove_from_queue", "set_processing_state", "clear_current_node", "set_current_task", "cancel_current_task", "set_node_error_sync", "drain_queue_and_mark_cancelled", "reset_processing_state", } assert forbidden.isdisjoint(vars(MessageTree))