src/free_claude_code/messaging/trees/processor.py

"""Task execution for claims returned by messaging tree aggregates."""

import asyncio
import contextlib
from collections.abc import Awaitable, Callable
from dataclasses import dataclass

from loguru import logger

from ..safe_diagnostics import format_exception_for_log
from .runtime import MessageTree
from .transitions import CancellationReason, NodeClaim, QueueEntry

NodeProcessor = Callable[[NodeClaim], Awaitable[None]]
QueueUpdateCallback = Callable[[tuple[QueueEntry, ...]], Awaitable[None]]
NodeStartedCallback = Callable[[NodeClaim], Awaitable[None]]
ClaimFailureCallback = Callable[[NodeClaim], Awaitable[None]]
ClaimFinishedCallback = Callable[[MessageTree, NodeClaim], Awaitable[None]]


@dataclass(slots=True)
class _TaskSlot:
    tree: MessageTree
    claim: NodeClaim
    task: asyncio.Task[None] | None = None
    runner_started: bool = False
    transitioned: bool = False
    recovery_task: asyncio.Task[None] | None = None
    cancellation_requested: bool = False
    cancellation_reason: CancellationReason | None = None


@dataclass(frozen=True, slots=True)
class CancelledTask:
    """Task handle plus whether the node runner owns cancellation UI."""

    task: asyncio.Task[None]
    runner_started: bool


class TreeQueueProcessor:
    """Own asyncio tasks while MessageTree owns scheduling state."""

    def __init__(
        self,
        node_processor: NodeProcessor,
        *,
        claim_failure_callback: ClaimFailureCallback,
        claim_finished_callback: ClaimFinishedCallback,
        queue_update_callback: QueueUpdateCallback | None = None,
        node_started_callback: NodeStartedCallback | None = None,
        log_messaging_error_details: bool = False,
    ) -> None:
        self._node_processor = node_processor
        self._claim_failure_callback = claim_failure_callback
        self._claim_finished_callback = claim_finished_callback
        self._queue_update_callback = queue_update_callback
        self._node_started_callback = node_started_callback
        self._log_messaging_error_details = log_messaging_error_details
        self._tasks: dict[str, _TaskSlot] = {}
        self._completion_failures: list[Exception] = []
        self._idle = asyncio.Event()
        self._idle.set()

    @staticmethod
    def _key(claim: NodeClaim) -> str:
        return claim.claim_id

    def launch(
        self,
        tree: MessageTree,
        claim: NodeClaim,
        *,
        announce_started: bool = False,
        queue: tuple[QueueEntry, ...] = (),
    ) -> None:
        """Attach a task synchronously before another coroutine can cancel it."""
        key = self._key(claim)
        if key in self._tasks:
            raise RuntimeError(f"Claim {key} already has a task")
        slot = _TaskSlot(tree=tree, claim=claim)
        self._tasks[key] = slot
        self._idle.clear()
        ownership_ready = asyncio.Event()
        claim_runner = self._run_claim(
            slot,
            ownership_ready=ownership_ready,
            announce_started=announce_started,
            queue=queue,
        )
        try:
            task = asyncio.create_task(
                claim_runner,
                name=(f"messaging-claim-{claim.identity.root_id}-{claim.claim_id[:8]}"),
            )
        except BaseException:
            claim_runner.close()
            if self._tasks.get(key) is slot:
                self._tasks.pop(key)
            if not self._tasks:
                self._idle.set()
            raise
        slot.task = task
        task.add_done_callback(lambda _task, claim_key=key: self._task_done(claim_key))
        ownership_ready.set()

    def _task_done(self, key: str) -> None:
        """Recover a claim if its task was cancelled before entering its body."""
        slot = self._tasks.get(key)
        if slot is None or slot.transitioned or slot.recovery_task is not None:
            return
        slot.recovery_task = asyncio.create_task(
            self._recover_unentered_task(slot),
            name=f"messaging-claim-recovery-{key[:8]}",
        )

    async def _recover_unentered_task(self, slot: _TaskSlot) -> None:
        task = slot.task
        if task is not None:
            with contextlib.suppress(asyncio.CancelledError, Exception):
                await task
        if not slot.transitioned:
            await self._finish_and_continue(slot)

    async def _notify_queue_updated(self, queue: tuple[QueueEntry, ...]) -> None:
        if self._queue_update_callback is None:
            return
        try:
            await self._queue_update_callback(queue)
        except Exception as exc:
            logger.warning(
                "Queue update callback failed: {}",
                format_exception_for_log(
                    exc,
                    log_full_message=self._log_messaging_error_details,
                ),
            )

    async def notify_queue_updated(self, queue: tuple[QueueEntry, ...]) -> None:
        """Publish a transition-owned queue snapshot."""
        await self._notify_queue_updated(queue)

    async def _notify_node_started(self, claim: NodeClaim) -> None:
        if self._node_started_callback is None:
            return
        try:
            await self._node_started_callback(claim)
        except Exception as exc:
            logger.warning(
                "Node started callback failed: {}",
                format_exception_for_log(
                    exc,
                    log_full_message=self._log_messaging_error_details,
                ),
            )

    async def _run_claim(
        self,
        slot: _TaskSlot,
        *,
        ownership_ready: asyncio.Event,
        announce_started: bool,
        queue: tuple[QueueEntry, ...],
    ) -> None:
        await ownership_ready.wait()
        claim = slot.claim
        try:
            if announce_started:
                await self._notify_node_started(claim)
                await self._notify_queue_updated(queue)
            if slot.cancellation_requested:
                if slot.cancellation_reason is None:
                    raise asyncio.CancelledError
                raise asyncio.CancelledError(slot.cancellation_reason)
            slot.runner_started = True
            await self._node_processor(claim)
        except asyncio.CancelledError:
            logger.info("Task for node {} was cancelled", claim.node.node_id)
            raise
        except Exception as exc:
            logger.error(
                "Error processing node {}: {}",
                claim.node.node_id,
                format_exception_for_log(
                    exc,
                    log_full_message=self._log_messaging_error_details,
                ),
            )
            await self._claim_failure_callback(claim)
        finally:
            if not slot.transitioned:
                await self._finish_and_continue(slot)

    async def _finish_and_continue(self, slot: _TaskSlot) -> None:
        current = asyncio.current_task()
        if current is not None:
            while current.cancelling():
                current.uncancel()
        try:
            while True:
                try:
                    await self._claim_finished_callback(slot.tree, slot.claim)
                    slot.transitioned = True
                    break
                except asyncio.CancelledError:
                    if current is not None:
                        while current.cancelling():
                            current.uncancel()
                    continue
        except Exception as exc:
            self._completion_failures.append(exc)
            logger.error(
                "Claim completion callback failed for node {}: {}",
                slot.claim.node.node_id,
                format_exception_for_log(
                    exc,
                    log_full_message=self._log_messaging_error_details,
                ),
            )
        finally:
            key = self._key(slot.claim)
            if self._tasks.get(key) is slot:
                self._tasks.pop(key)
            if not self._tasks:
                self._idle.set()

    def cancel(
        self,
        claim: NodeClaim,
        reason: CancellationReason | None,
    ) -> CancelledTask | None:
        """Cancel exactly the task bound to one aggregate claim."""
        slot = self._tasks.get(self._key(claim))
        if slot is None:
            return None
        slot.cancellation_requested = True
        slot.cancellation_reason = reason
        task = slot.task
        if task is None or task.done():
            return None
        if reason is None:
            task.cancel()
        else:
            task.cancel(reason)

        if slot.runner_started:
            return CancelledTask(task=task, runner_started=True)

        if slot.recovery_task is None:
            slot.recovery_task = asyncio.create_task(
                self._recover_unentered_task(slot),
                name=f"messaging-claim-recovery-{claim.claim_id[:8]}",
            )
        return CancelledTask(task=slot.recovery_task, runner_started=False)

    def task_count(self) -> int:
        """Return the number of attached claims for observability."""
        return len(self._tasks)

    async def wait_idle(self) -> None:
        """Wait for every task and hand completion failures to the caller once."""
        await self._idle.wait()
        if not self._completion_failures:
            return
        failures = self._completion_failures
        self._completion_failures = []
        if len(failures) == 1:
            raise failures[0]
        raise ExceptionGroup("Messaging claim completion failures", failures)


__all__ = ["CancelledTask", "TreeQueueProcessor"]