from __future__ import annotations import asyncio import logging import time from typing import Any from yuxi.channel.extensions.qqbot.api_client import QQBotApiClient logger = logging.getLogger(__name__) STREAM_THROTTLE_MS = 500 STREAM_MIN_THROTTLE_MS = 300 STREAM_LONG_INTERVAL_MS = 2000 STREAM_LONG_BATCH_WINDOW_MS = 300 class FlushController: def __init__(self, throttle_ms: int = STREAM_THROTTLE_MS): self._ready = False self._flush_lock = asyncio.Lock() self._needs_reflush = False self._is_completed = False self._last_update_time = 0.0 self._throttle_ms = throttle_ms @property def ready(self) -> bool: return self._ready @ready.setter def ready(self, value: bool) -> None: self._ready = value @property def is_completed(self) -> bool: return self._is_completed @is_completed.setter def is_completed(self, value: bool) -> None: self._is_completed = value @property def needs_reflush(self) -> bool: return self._needs_reflush @needs_reflush.setter def needs_reflush(self, value: bool) -> None: self._needs_reflush = value @property def last_update_time(self) -> float: return self._last_update_time @last_update_time.setter def last_update_time(self, value: float) -> None: self._last_update_time = value class C2CStreamingController: def __init__(self, api_client: QQBotApiClient): self._api_client = api_client self._sessions: dict[str, FlushController] = {} self._stream_ids: dict[str, str] = {} self._pending_text: dict[str, str] = {} self._flush_tasks: dict[str, asyncio.Task] = {} self._msg_seqs: dict[str, int] = {} async def start_streaming_session(self, openid: str) -> str: if openid in self._stream_ids: return self._stream_ids[openid] data = await self._api_client.start_c2c_stream(openid) stream_id = data.get("stream_id", data.get("id", "")) self._stream_ids[openid] = stream_id self._pending_text[openid] = "" self._sessions[openid] = FlushController() logger.debug("Streaming session started: openid=%s, stream_id=%s", openid, stream_id) return stream_id async def feed(self, openid: str, text: str) -> None: if openid not in self._pending_text: return self._pending_text[openid] += text controller = self._sessions.get(openid) if controller: controller.last_update_time = time.time() controller.needs_reflush = True await self._schedule_flush(openid) async def flush(self, openid: str) -> None: text = self._pending_text.get(openid, "") stream_id = self._stream_ids.get(openid, "") controller = self._sessions.get(openid) if not text or not stream_id: return try: if not controller or not controller.ready: await self._api_client.send_c2c_stream_chunk(openid, stream_id, text) if controller: controller.ready = True else: await self._api_client.send_c2c_stream_chunk(openid, stream_id, text) self._pending_text[openid] = "" if controller: controller.last_update_time = time.time() controller.needs_reflush = False except Exception as e: logger.warning("Stream flush failed for openid=%s: %s", openid, e) raise async def complete(self, openid: str) -> None: stream_id = self._stream_ids.pop(openid, "") if not stream_id: return controller = self._sessions.pop(openid, None) if controller: controller.is_completed = True pending = self._pending_text.pop(openid, "") if pending: try: await self._api_client.send_c2c_stream_chunk(openid, stream_id, pending) except Exception: logger.exception("Final stream flush failed for openid=%s", openid) await self._cancel_flush_task(openid) try: await self._api_client.complete_c2c_stream(openid, stream_id) except Exception: logger.exception("Stream completion failed for openid=%s", openid) async def abort(self, openid: str, reason: str | None = None) -> None: stream_id = self._stream_ids.pop(openid, "") if not stream_id: return self._sessions.pop(openid, None) self._pending_text.pop(openid, "") await self._cancel_flush_task(openid) try: await self._api_client.abort_c2c_stream(openid, stream_id) except Exception: logger.exception("Stream abort failed for openid=%s", openid) async def _schedule_flush(self, openid: str) -> None: controller = self._sessions.get(openid) if not controller or controller.is_completed: return await self._cancel_flush_task(openid) delay_ms = controller._throttle_ms now = time.time() elapsed = (now - controller.last_update_time) * 1000 if elapsed > STREAM_LONG_INTERVAL_MS: delay_ms = STREAM_LONG_BATCH_WINDOW_MS else: delay_ms = max(STREAM_MIN_THROTTLE_MS, delay_ms) async def _delayed_flush(): await asyncio.sleep(delay_ms / 1000.0) if openid in self._pending_text and controller and not controller.is_completed: try: async with controller._flush_lock: await self.flush(openid) except Exception: pass self._flush_tasks[openid] = asyncio.create_task(_delayed_flush()) async def _cancel_flush_task(self, openid: str) -> None: task = self._flush_tasks.pop(openid, None) if task and not task.done(): task.cancel() try: await task except asyncio.CancelledError: pass async def fallback_to_static(self, openid: str, text: str) -> bool: await self.abort(openid) if text: await self._api_client.send_c2c_message(openid, text, msg_type=0) return True