ForcePilot/backend/package/yuxi/channels/adapters/dingding/stream.py

from __future__ import annotations

import asyncio
import time
from typing import TYPE_CHECKING, Any

import httpx

from yuxi.channels.models import DeliveryResult
from yuxi.utils.logging_config import logger

if TYPE_CHECKING:
    from yuxi.channels.adapters.dingding.token import DingDingTokenManager

STREAM_UPDATE_MIN_INTERVAL_S = 0.16
STREAM_SIGNIFICANT_INCREMENT = 18
STREAM_NATURAL_BOUNDARIES = ("\n", "。", "！", "？", "；", "：", ",", ". ", "! ", "? ", "; ", ": ")
STREAM_PLACEHOLDER_TEXT = "⏳ Thinking..."
MAX_CONSECUTIVE_FAILURES = 5
BACKOFF_BASE_S = 0.5
BACKOFF_MAX_S = 8.0


class DingDingStreamingSession:
    def __init__(
        self,
        token_manager: DingDingTokenManager,
        open_conversation_id: str,
        robot_code: str,
        chat_id: str,
        *,
        config: dict[str, Any] | None = None,
        http_client: httpx.AsyncClient | None = None,
    ):
        self._token_manager = token_manager
        self._open_conversation_id = open_conversation_id
        self._robot_code = robot_code
        self._chat_id = chat_id
        self._config = config or {}
        self._http_client = http_client

        self._msg_id: str | None = None
        self._accumulated = ""
        self._last_update = 0.0
        self._closed = False
        self._consecutive_failures = 0
        self._started = False
        self._chat_type = "group" if chat_id.startswith("group_") else "direct"

        stream_cfg = self._config.get("streaming", {})
        if not isinstance(stream_cfg, dict):
            stream_cfg = {}
        self._stream_mode = stream_cfg.get("mode", "block")
        block_cfg = stream_cfg.get("block", {})
        if not isinstance(block_cfg, dict):
            block_cfg = {}
        self._update_interval_ms = block_cfg.get("update_interval_ms", 800)
        self._coalesce_min_chars = block_cfg.get("coalesce_min_chars", 100)

    async def start(self, initial_text: str = "") -> DeliveryResult:
        if self._started:
            return DeliveryResult(success=True, message_id=self._msg_id or "")

        self._started = True
        send_text = initial_text or STREAM_PLACEHOLDER_TEXT

        try:
            from .send import send_markdown

            result = await send_markdown(
                self._token_manager,
                self._open_conversation_id,
                self._robot_code,
                "ForcePilot",
                send_text,
                chat_type=self._chat_type,
                http_client=self._http_client,
            )
            if result.success:
                self._msg_id = result.message_id
                self._consecutive_failures = 0
            else:
                self._consecutive_failures += 1
            return result
        except Exception as e:
            logger.error(f"[DingDingStream] start error: {e}")
            self._consecutive_failures += 1
            return DeliveryResult(success=False, error=str(e))

    async def send_chunk(self, chunk: str, finished: bool = False) -> DeliveryResult:
        if self._closed:
            return DeliveryResult(success=False, error="Stream closed")

        if not self._msg_id:
            return await self.start(chunk)

        if finished:
            return await self._finish()

        self._accumulated += chunk
        return await self._maybe_update()

    async def _finish(self) -> DeliveryResult:
        self._closed = True
        try:
            from .send import send_markdown

            result = await send_markdown(
                self._token_manager,
                self._open_conversation_id,
                self._robot_code,
                "ForcePilot",
                self._accumulated,
                chat_type=self._chat_type,
                http_client=self._http_client,
            )
            if result.success:
                self._consecutive_failures = 0
            else:
                self._consecutive_failures += 1
            return result
        except Exception as e:
            logger.error(f"[DingDingStream] finish error: {e}")
            self._consecutive_failures += 1
            return DeliveryResult(success=False, error=str(e))

    async def _maybe_update(self) -> DeliveryResult:
        now = time.monotonic()
        elapsed_ms = (now - self._last_update) * 1000

        should_send = (
            len(self._accumulated) >= self._coalesce_min_chars
            and elapsed_ms >= self._update_interval_ms
        ) or any(self._accumulated.endswith(b) for b in STREAM_NATURAL_BOUNDARIES)

        if not should_send:
            return DeliveryResult(success=True, message_id=self._msg_id)

        self._last_update = now

        try:
            result = await self._edit_message(self._accumulated)
            if result.success:
                self._consecutive_failures = 0
            else:
                self._consecutive_failures += 1
                logger.debug(f"[DingDingStream] edit failed (consecutive={self._consecutive_failures}): {result.error}")
            return result
        except Exception as e:
            self._consecutive_failures += 1
            logger.debug(f"[DingDingStream] edit error: {e}")
            return DeliveryResult(success=False, error=str(e))

    async def _edit_message(self, content: str) -> DeliveryResult:
        if not self._token_manager:
            return DeliveryResult(success=False, error="Not connected")

        try:
            token = await self._token_manager.get_token()
            http_client = self._http_client or httpx.AsyncClient(timeout=httpx.Timeout(15))
            url = f"https://api.dingtalk.com/v1.0/robot/messages/{self._msg_id}"

            headers = {
                "x-acs-dingtalk-access-token": token,
                "Content-Type": "application/json",
            }

            import json

            payload = {
                "msgParam": json.dumps({"content": content}, ensure_ascii=False),
                "msgKey": "sampleMarkdown",
                "openConversationId": self._open_conversation_id,
                "robotCode": self._robot_code,
            }

            resp = await http_client.put(url, json=payload, headers=headers)
            if resp.status_code != 200:
                return DeliveryResult(success=False, error=f"DingDing API HTTP {resp.status_code}")

            data = resp.json()
            return DeliveryResult(success=True, message_id=data.get("processQueryKey", self._msg_id))
        except Exception as e:
            logger.error(f"[DingDingStream] _edit_message failed: {e}")
            return DeliveryResult(success=False, error=str(e))

    @property
    def message_id(self) -> str | None:
        return self._msg_id

    @property
    def accumulated_text(self) -> str:
        return self._accumulated

    @property
    def is_closed(self) -> bool:
        return self._closed

    @property
    def consecutive_failures(self) -> int:
        return self._consecutive_failures

    def should_degrade(self) -> bool:
        return self._consecutive_failures >= MAX_CONSECUTIVE_FAILURES

    def backoff_delay(self) -> float:
        if self._consecutive_failures == 0:
            return 0.0
        return min(BACKOFF_BASE_S * (2 ** (self._consecutive_failures - 1)), BACKOFF_MAX_S)


class DingDingStreamManager:
    def __init__(self, adapter):
        self._adapter = adapter
        self._sessions: dict[str, DingDingStreamingSession] = {}
        self._pending_block_text: dict[str, str] = {}
        self._lock = asyncio.Lock()

    async def send_chunk(self, chat_id: str, msg_id: str, chunk: str, finished: bool) -> DeliveryResult:
        config = getattr(self._adapter, "config", {})
        stream_cfg = config.get("streaming", {})
        if not isinstance(stream_cfg, dict):
            stream_cfg = {}
        stream_mode = stream_cfg.get("mode", "block")

        if stream_mode == "partial":
            return await self._send_partial(chat_id, msg_id, chunk, finished)

        return await self._send_block(chat_id, msg_id, chunk, finished)

    async def _send_partial(self, chat_id: str, msg_id: str, chunk: str, finished: bool) -> DeliveryResult:
        session = self._sessions.get(chat_id)
        if session is None or session.is_closed:
            adapter_config = getattr(self._adapter, "config", {})
            accounts = adapter_config.get("accounts", {})
            robot_code = accounts.get("default", {}).get("robot_code", "")
            open_conv_id = chat_id.replace("group_", "").replace("dm_", "")
            token_mgr = getattr(self._adapter, "_token_manager", None)
            http_client = await self._get_http_client()

            session = DingDingStreamingSession(
                token_mgr,
                open_conv_id,
                robot_code,
                chat_id,
                config=adapter_config,
                http_client=http_client,
            )
            self._sessions[chat_id] = session

        if session.should_degrade():
            logger.warning(
                f"[DingDingStream] Degrading partial stream for {chat_id}"
                f" after {session.consecutive_failures} failures"
            )
            return await self._fallback_block_send(chat_id)

        return await session.send_chunk(chunk, finished)

    async def _send_block(self, chat_id: str, msg_id: str, chunk: str, finished: bool) -> DeliveryResult:
        async with self._lock:
            if not finished:
                self._pending_block_text[chat_id] = self._pending_block_text.get(chat_id, "") + chunk
                return DeliveryResult(success=True, message_id=msg_id or "pending")

            final_text = self._pending_block_text.pop(chat_id, chunk)

        return await self._fallback_block_send(chat_id, final_text)

    async def _fallback_block_send(self, chat_id: str, text: str = "") -> DeliveryResult:
        if not text:
            return DeliveryResult(success=True)

        accounts = getattr(self._adapter, "config", {}).get("accounts", {})
        robot_code = accounts.get("default", {}).get("robot_code", "")
        open_conv_id = chat_id.replace("group_", "").replace("dm_", "")
        chat_type = "group" if chat_id.startswith("group_") else "direct"
        token_mgr = getattr(self._adapter, "_token_manager", None)
        http_client = await self._get_http_client()

        from .send import send_markdown

        async with self._lock:
            return await send_markdown(
                token_mgr,
                open_conv_id,
                robot_code,
                "ForcePilot",
                text,
                chat_type=chat_type,
                http_client=http_client,
            )

    async def _get_http_client(self) -> httpx.AsyncClient:
        if hasattr(self._adapter, "_get_http_client"):
            return await self._adapter._get_http_client()
        return httpx.AsyncClient(timeout=httpx.Timeout(15))

    def close_session(self, chat_id: str) -> None:
        self._sessions.pop(chat_id, None)
        self._pending_block_text.pop(chat_id, None)

    def cleanup(self) -> None:
        self._sessions.clear()
        self._pending_block_text.clear()