ForcePilot/backend/package/yuxi/channel/message/reply_dispatcher.py

import asyncio
import logging
import random
import time
from collections.abc import Callable, Coroutine
from typing import Any

from yuxi.channel.message.models import ReplyPayload, ReplyStage, StreamingChunk, UnifiedMessage

logger = logging.getLogger(__name__)

_DEFAULT_MIN_CHARS = 80
_DEFAULT_DEBOUNCE_MS = 150.0
_DEFAULT_TIMEOUT_MS = 10_000.0

OnIdleCallback = Callable[[], None]
OnErrorCallback = Callable[[Exception], None]
BeforeDeliverCallback = Callable[[str], Coroutine[Any, Any, str | None]]


class ReplyDispatcher:
    def __init__(
        self,
        send_fn: Callable[..., Coroutine[Any, Any, str | None]],
        msg: UnifiedMessage,
        *,
        response_prefix: str = "",
        human_delay: tuple[float, float] | None = None,
        timeout_ms: float = _DEFAULT_TIMEOUT_MS,
        on_idle: OnIdleCallback | None = None,
        on_error: OnErrorCallback | None = None,
        before_deliver: BeforeDeliverCallback | None = None,
    ):
        self._send_fn = send_fn
        self._msg = msg
        self._response_prefix = response_prefix
        self._human_delay = human_delay
        self._timeout_ms = timeout_ms
        self._on_idle = on_idle
        self._on_error = on_error
        self._before_deliver = before_deliver

        self._send_chain: asyncio.Lock = asyncio.Lock()
        self._flush_lock: asyncio.Lock = asyncio.Lock()
        self._pending = 0
        self._complete = asyncio.Event()
        self._done = False

        self._blocks: list[str] = []
        self._block_deadline: float = 0.0
        self._block_timer: asyncio.Task | None = None
        self._sent_payload_keys: set[str] = set()
        self._sent_content_keys: set[str] = set()

        self._tool_count = 0
        self._block_count = 0
        self._final_count = 0
        self._failed_count = 0
        self._cancelled_count = 0
        self._started_at: float | None = None

        self._abort_event = asyncio.Event()

    async def enqueue(self, chunk: StreamingChunk) -> bool:
        if self._abort_event.is_set():
            self._cancelled_count += 1
            return False

        if chunk.stage == ReplyStage.TOOL:
            return await self._dispatch_tool(chunk)
        elif chunk.stage == ReplyStage.BLOCK:
            return await self._dispatch_block(chunk)
        elif chunk.stage == ReplyStage.FINAL:
            return await self._dispatch_final(chunk)
        return False

    async def _dispatch_tool(self, chunk: StreamingChunk) -> bool:
        self._pending += 1
        self._tool_count += 1
        try:
            prefix = self._response_prefix if self._response_prefix else ""
            content = f"{prefix}{chunk.content}" if chunk.content else ""
            if not content:
                return False

            await self._send_ordered(content)
            self._block_count += 1
            return True
        except Exception as e:
            logger.exception("Tool dispatch failed")
            self._failed_count += 1
            if self._on_error:
                self._on_error(e)
            raise
        finally:
            self._pending -= 1
            self._maybe_fire_idle()

    async def _dispatch_block(self, chunk: StreamingChunk) -> bool:
        if not chunk.content:
            return False

        async with self._flush_lock:
            self._blocks.append(chunk.content)
            now = time.monotonic()

            if self._block_deadline == 0.0:
                self._block_deadline = now + _DEFAULT_DEBOUNCE_MS / 1000.0
                self._block_timer = asyncio.create_task(self._flush_after_debounce())

            if sum(len(b) for b in self._blocks) >= _DEFAULT_MIN_CHARS:
                await self._flush_blocks()
        return True

    async def _flush_after_debounce(self) -> None:
        try:
            await asyncio.sleep(_DEFAULT_DEBOUNCE_MS / 1000.0)
        except asyncio.CancelledError:
            return
        async with self._flush_lock:
            await self._flush_blocks()

    async def _flush_blocks(self) -> None:
        if not self._blocks:
            return
        combined = "".join(self._blocks)
        self._blocks.clear()
        self._block_deadline = 0.0
        if self._block_timer and not self._block_timer.done():
            self._block_timer.cancel()
            self._block_timer = None

        payload = ReplyPayload(
            target_id=self._target_id,
            content=combined,
            reply_to_id=self._msg.msg_id if self._block_count == 0 else None,
        )
        if payload.content_key in self._sent_content_keys:
            return
        self._sent_content_keys.add(payload.content_key)

        self._pending += 1
        self._block_count += 1
        try:
            await self._send_ordered(combined)
        except Exception as e:
            logger.exception("Block dispatch failed")
            self._failed_count += 1
            if self._on_error:
                self._on_error(e)
            raise
        finally:
            self._pending -= 1
            self._maybe_fire_idle()

    async def _dispatch_final(self, chunk: StreamingChunk) -> bool:
        async with self._flush_lock:
            await self._flush_blocks()

        if chunk.content and chunk.content != "finished":
            self._pending += 1
            self._final_count += 1
            try:
                await self._send_ordered(chunk.content)
            except Exception as e:
                logger.exception("Final dispatch failed")
                self._failed_count += 1
                if self._on_error:
                    self._on_error(e)
                raise
            finally:
                self._pending -= 1
                self._maybe_fire_idle()
            self._done = True
            self._complete.set()
            return True

        self._done = True
        self._complete.set()
        self._maybe_fire_idle()
        return True

    async def _send_ordered(self, content: str) -> None:
        async with self._send_chain:
            if self._abort_event.is_set():
                return

            if self._human_delay and self._block_count > 0:
                delay = self._human_delay[0] + (self._human_delay[1] - self._human_delay[0]) * random.random()
                await asyncio.sleep(delay / 1000.0)

            if self._response_prefix and self._block_count == 0:
                prefix = self._response_prefix.format(
                    agent_name=self._msg.channel_type,
                    model="",
                )
                content = f"{prefix}{content}"

            if self._before_deliver:
                try:
                    content = await self._before_deliver(content)
                except Exception as e:
                    logger.exception("before_deliver hook failed")
                    if self._on_error:
                        self._on_error(e)
                    return
                if content is None:
                    self._cancelled_count += 1
                    return

            await self._send_fn(content)

    @property
    def _target_id(self) -> str:
        return self._msg.group.id if self._msg.group and self._msg.group.id else self._msg.sender.id

    async def wait_idle(self, timeout_ms: float | None = None) -> None:
        timeout = timeout_ms or self._timeout_ms
        try:
            await asyncio.wait_for(self._complete.wait(), timeout=timeout / 1000.0)
        except TimeoutError:
            self._abort_event.set()
            self._done = True
            self._complete.set()
            logger.warning("ReplyDispatcher idle timeout after %.0fms, aborting", timeout)

    def mark_complete(self) -> None:
        self._done = True
        self._complete.set()
        self._maybe_fire_idle()

    async def abort(self) -> None:
        async with self._flush_lock:
            if self._blocks:
                await self._flush_blocks()
        self._abort_event.set()
        self._complete.set()

    def _maybe_fire_idle(self) -> None:
        if self._done and self._pending == 0 and self._on_idle:
            self._on_idle()

    @property
    def tool_count(self) -> int:
        return self._tool_count

    @property
    def block_count(self) -> int:
        return self._block_count

    @property
    def final_count(self) -> int:
        return self._final_count

    @property
    def failed_count(self) -> int:
        return self._failed_count

    @property
    def cancelled_count(self) -> int:
        return self._cancelled_count

    @property
    def queued_counts(self) -> dict[str, int]:
        return {"tool": self._tool_count, "block": self._block_count, "final": self._final_count}

    @property
    def is_idle(self) -> bool:
        return self._done and self._pending == 0