ForcePilot/backend/package/yuxi/channels/adapters/yuanbao/outbound_queue.py

from __future__ import annotations

import asyncio
from collections.abc import Callable, Awaitable
from typing import Any

from yuxi.channels.models import ChannelResponse, ChannelIdentity, DeliveryResult
from yuxi.utils.logging_config import logger


class OutboundQueue:
    def __init__(
        self,
        config: dict[str, Any],
        send_fn: Callable[[ChannelResponse], Awaitable[DeliveryResult]],
    ):
        strategy = config.get("outboundQueueStrategy", "merge-text")
        self._strategy = strategy
        self._min_chars = config.get("minChars", 2800)
        self._max_chars = config.get("maxChars", 3000)
        self._idle_ms = config.get("idleMs", 5000)
        self._send_fn = send_fn

        self._buffer: list[str] = []
        self._buffer_len = 0
        self._idle_task: asyncio.Task | None = None
        self._lock = asyncio.Lock()
        self._flush_event = asyncio.Event()

        self._first_identity: ChannelIdentity | None = None
        self._first_metadata: dict[str, Any] = {}

    async def enqueue(self, response: ChannelResponse) -> DeliveryResult | None:
        if self._strategy == "immediate":
            return await self._send_fn(response)

        return await self._enqueue_merge(response)

    async def _enqueue_merge(self, response: ChannelResponse) -> DeliveryResult | None:
        content = response.content

        async with self._lock:
            if not self._buffer:
                self._first_identity = response.identity
                self._first_metadata = response.metadata

            self._buffer.append(content)
            self._buffer_len += len(content)

            if self._buffer_len >= self._max_chars:
                return await self._flush()
            elif self._buffer_len >= self._min_chars:
                return await self._flush()

            self._cancel_idle()
            self._idle_task = asyncio.create_task(self._idle_flush())

        return None

    async def flush(self) -> DeliveryResult | None:
        async with self._lock:
            return await self._flush()

    async def _flush(self) -> DeliveryResult | None:
        if not self._buffer:
            return None

        merged = "\n\n".join(self._buffer)
        identity = self._first_identity
        metadata = self._first_metadata

        self._buffer.clear()
        self._buffer_len = 0
        self._first_identity = None
        self._first_metadata = {}
        self._cancel_idle()

        logger.debug(f"[Yuanbao] OutboundQueue flushed {len(merged)} chars")

        if identity is None:
            logger.error("[Yuanbao] OutboundQueue: _flush called without identity, dropping messages")
            return DeliveryResult(success=False, error="OutboundQueue: missing identity in flushed response")
        response = ChannelResponse(
            identity=identity,
            content=merged,
            metadata=metadata,
        )
        return await self._send_fn(response)

    async def _idle_flush(self) -> None:
        await asyncio.sleep(self._idle_ms / 1000)
        async with self._lock:
            if self._buffer:
                logger.debug("[Yuanbao] OutboundQueue idle flush triggered")
                await self._flush()

    def _cancel_idle(self) -> None:
        if self._idle_task and not self._idle_task.done():
            self._idle_task.cancel()
            self._idle_task = None