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

from __future__ import annotations

import asyncio
from collections.abc import Awaitable, Callable

from yuxi.utils.logging_config import logger


SENTINEL = object()


class EventQueue:
    def __init__(self, max_size: int = 10000):
        self._queue: asyncio.Queue = asyncio.Queue(maxsize=max_size)
        self._handler: Callable[[dict], Awaitable[None]] | None = None
        self._task: asyncio.Task | None = None
        self._running = False
        self._dropped_count = 0

    async def start(self, handler: Callable[[dict], Awaitable[None]]) -> None:
        self._handler = handler
        self._running = True
        self._task = asyncio.create_task(self._process_loop())
        logger.info("[Yuanbao] Event queue started")

    async def stop(self) -> None:
        self._running = False
        self._queue.put_nowait(SENTINEL)
        if self._task and not self._task.done():
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            self._task = None
        logger.info("[Yuanbao] Event queue stopped")

    async def enqueue(self, event: dict) -> None:
        try:
            self._queue.put_nowait(event)
        except asyncio.QueueFull:
            event_type = event.get("type", "unknown")
            msg_id = event.get("msg_id", "N/A")
            logger.warning(
                f"[Yuanbao] Event queue full, dropping event: "
                f"type={event_type} msg_id={msg_id}"
            )
            self._dropped_count += 1

    async def _process_loop(self) -> None:
        while self._running:
            try:
                event = await self._queue.get()
                if event is SENTINEL:
                    break
                if self._handler:
                    await self._handler(event)
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"[Yuanbao] Event queue process error: {e}", exc_info=True)

    @property
    def size(self) -> int:
        return self._queue.qsize()