ForcePilot/backend/package/yuxi/channel/extensions/wechat-mp/passive_reply.py

import asyncio
import logging
import time
from dataclasses import dataclass, field
from datetime import datetime, UTC

from yuxi.channel.extensions.wechat_mp.message import extract_content
from yuxi.channel.message.models import MessageType, PeerInfo, UnifiedMessage
from yuxi.channel.routing.models import PeerKind

logger = logging.getLogger(__name__)

MAX_RETRY = 3
WAIT_WINDOW_SEC = 4.0
POLL_INTERVAL_SEC = 0.1

_passive_state: "PassiveReplyState | None" = None


def _get_passive_state() -> "PassiveReplyState":
    global _passive_state
    if _passive_state is None:
        _passive_state = PassiveReplyState()
    return _passive_state


def _reset_passive_state() -> None:
    global _passive_state
    _passive_state = None


@dataclass
class PassiveReplyState:
    cache: dict[str, list[tuple[str, str]]] = field(default_factory=dict)
    running: set[str] = field(default_factory=set)
    request_cnt: dict[str, int] = field(default_factory=dict)

    def is_new_request(self, from_user: str, message_id: str, content: str) -> bool:
        has_cache = from_user in self.cache and bool(self.cache[from_user])
        is_running = from_user in self.running
        is_command = content.startswith("#")
        is_new_msg = message_id not in self.request_cnt
        return (not has_cache and not is_running) or (is_command and is_new_msg)

    def mark_running(self, from_user: str):
        self.running.add(from_user)

    def mark_done(self, from_user: str):
        self.running.discard(from_user)

    def incr_request(self, message_id: str):
        self.request_cnt[message_id] = self.request_cnt.get(message_id, 0) + 1

    def is_done(self, from_user: str) -> bool:
        return from_user not in self.running

    def reap(self, from_user: str) -> tuple[str, str] | None:
        entries = self.cache.get(from_user)
        if not entries:
            return None
        result = entries.pop(0)
        if not entries:
            del self.cache[from_user]
        return result

    def add_cache(self, from_user: str, reply_type: str, content: str):
        if from_user not in self.cache:
            self.cache[from_user] = []
        self.cache[from_user].append((reply_type, content))

    def cleanup_user(self, from_user: str):
        self.cache.pop(from_user, None)
        self.running.discard(from_user)

    def reset(self):
        self.cache.clear()
        self.running.clear()
        self.request_cnt.clear()


async def handle_passive_callback(
    raw_msg,
    decrypted_xml: str,
    processor,
    outbound,
    subscribe_msg: str = "",
) -> str:
    state = _get_passive_state()

    msg_content = raw_msg.content or ""

    if not state.is_new_request(raw_msg.from_user, raw_msg.msg_id, msg_content):
        state.incr_request(raw_msg.msg_id)

    state.mark_running(raw_msg.from_user)

    asyncio.create_task(
        _process_passive_msg(raw_msg, processor, outbound),
        name=f"wechat-mp-passive-{raw_msg.from_user}",
    )

    deadline = time.time() + WAIT_WINDOW_SEC
    while time.time() < deadline:
        if state.is_done(raw_msg.from_user):
            result = state.reap(raw_msg.from_user)
            if result:
                reply_type, content = result
                return _build_passive_xml(result, raw_msg)

        await asyncio.sleep(POLL_INTERVAL_SEC)

    count = state.request_cnt.get(raw_msg.msg_id, 1)
    state.cleanup_user(raw_msg.from_user)

    if count < MAX_RETRY:
        return "success"

    return _build_text_xml(raw_msg.to_user, raw_msg.from_user, "【正在思考中，请稍后再试】")


async def _process_passive_msg(raw_msg, processor, outbound) -> None:
    state = _get_passive_state()
    from_user = raw_msg.from_user

    try:
        content = extract_content(raw_msg)
        if not content:
            state.mark_done(from_user)
            return

        msg_type_map = {
            "text": MessageType.TEXT,
            "image": MessageType.IMAGE,
            "voice": MessageType.VOICE,
            "video": MessageType.VIDEO,
            "shortvideo": MessageType.VIDEO,
            "location": MessageType.TEXT,
        }
        msg_type = msg_type_map.get(raw_msg.msg_type, MessageType.TEXT)

        unified = UnifiedMessage(
            msg_id=raw_msg.msg_id,
            channel_type="wechat-mp",
            account_id="default",
            content=content,
            message_type=msg_type,
            sender=PeerInfo(
                id=raw_msg.from_user,
                kind=PeerKind.DIRECT,
                display_name=raw_msg.from_user,
            ),
            timestamp=datetime.fromtimestamp(raw_msg.create_time, tz=UTC) if raw_msg.create_time else None,
            raw_payload=raw_msg.raw_xml,
            body_for_agent=content,
            metadata={
                "FromUserName": raw_msg.from_user,
                "ToUserName": raw_msg.to_user,
                "MsgType": raw_msg.msg_type,
                "Event": raw_msg.event,
                "EventKey": raw_msg.event_key,
                "MediaId": raw_msg.media_id,
                "PicUrl": raw_msg.pic_url,
            },
        )

        if processor is not None:
            try:
                await asyncio.wait_for(processor.process(unified), timeout=120.0)
            except TimeoutError:
                logger.error("Passive agent response timeout for wechat-mp user %s", from_user)
            except Exception:
                logger.exception("Failed to process passive WeChat MP message for user %s", from_user)
        state.mark_done(from_user)

    except Exception:
        logger.exception("Passive message processing error for user %s", from_user)
        state.mark_done(from_user)


def _build_passive_xml(result: tuple[str, str], raw_msg) -> str:
    reply_type, content = result
    to_user = raw_msg.from_user
    from_user = raw_msg.to_user

    if reply_type == "text":
        return _build_text_xml(from_user, to_user, content)

    if reply_type == "image":
        return (
            '<xml>'
            f'<ToUserName><![CDATA[{to_user}]]></ToUserName>'
            f'<FromUserName><![CDATA[{from_user}]]></FromUserName>'
            f'<CreateTime>{int(time.time())}</CreateTime>'
            '<MsgType><![CDATA[image]]></MsgType>'
            '<Image>'
            f'<MediaId><![CDATA[{content}]]></MediaId>'
            '</Image>'
            '</xml>'
        )

    if reply_type == "voice":
        return (
            '<xml>'
            f'<ToUserName><![CDATA[{to_user}]]></ToUserName>'
            f'<FromUserName><![CDATA[{from_user}]]></FromUserName>'
            f'<CreateTime>{int(time.time())}</CreateTime>'
            '<MsgType><![CDATA[voice]]></MsgType>'
            '<Voice>'
            f'<MediaId><![CDATA[{content}]]></MediaId>'
            '</Voice>'
            '</xml>'
        )

    return _build_text_xml(from_user, to_user, content)


def _build_text_xml(from_user: str, to_user: str, content: str) -> str:
    return (
        '<xml>'
        f'<ToUserName><![CDATA[{to_user}]]></ToUserName>'
        f'<FromUserName><![CDATA[{from_user}]]></FromUserName>'
        f'<CreateTime>{int(time.time())}</CreateTime>'
        '<MsgType><![CDATA[text]]></MsgType>'
        f'<Content><![CDATA[{content}]]></Content>'
        '</xml>'
    )