wecom_it_smart_desk/backend/app/services/automation/progress_publisher.py

# =============================================================================
# 企微IT智能服务台 — 阶段5 自动化 进度推送
# =============================================================================
# 说明：负责把自动化处置过程实时推送给前端：
#   1. 专用 WS 通道 /ws/automation/{session_id}（坐席工作台 + 员工 H5 均可连）
#   2. 兜底推送：同时向坐席(/ws/{agent_id})和员工(/ws/h5/{employee_id})推送，
#      保证未连专用 WS 时也能收到事件。
#   3. 静默关单调度：处置成功后 N 分钟员工无异议则自动关单。
#
# 事件名统一以 automation. 前缀（见 app.constants）。
# =============================================================================

from __future__ import annotations

import asyncio
import logging
from typing import Any, Dict, Optional, Set

from app.constants import (
    AUTOMATION_SILENT_CLOSE_TTL,
    AUTOMATION_WS_ACTION_REQUIRED,
    AUTOMATION_WS_ERROR,
    AUTOMATION_WS_PROGRESS,
    AUTOMATION_WS_RESOLVED,
    AUTOMATION_WS_TAKEOVER,
)
from app.services.ws_manager import manager as ws_manager

logger = logging.getLogger(__name__)

# 会话级专用 WS 连接注册表：session_id -> {websocket, ...}
_automation_ws: Dict[str, Set[Any]] = {}

# 会话参与方（用于兜底推送）：session_id -> {"agent_id":..,"employee_id":..}
_session_parties: Dict[str, Dict[str, Optional[str]]] = {}

# 静默关单任务：session_id -> asyncio.Task
_silent_close_tasks: Dict[str, asyncio.Task] = {}


def register_ws(session_id: str, websocket: Any) -> None:
    """注册专用 WS 连接。"""
    _automation_ws.setdefault(session_id, set()).add(websocket)


def unregister_ws(session_id: str, websocket: Any) -> None:
    """注销专用 WS 连接。"""
    conns = _automation_ws.get(session_id)
    if conns:
        conns.discard(websocket)
        if not conns:
            _automation_ws.pop(session_id, None)


def set_parties(
    session_id: str,
    agent_id: Optional[str] = None,
    employee_id: Optional[str] = None,
) -> None:
    """记录会话参与方，用于兜底推送。"""
    parties = _session_parties.setdefault(
        session_id, {"agent_id": None, "employee_id": None}
    )
    if agent_id is not None:
        parties["agent_id"] = agent_id
    if employee_id is not None:
        parties["employee_id"] = employee_id


def _build_message(event_type: str, session_id: str, data: Any) -> Dict[str, Any]:
    """构造统一 WS 消息信封。"""
    return {"type": event_type, "session_id": session_id, "data": data or {}}


async def _send_to_automation_ws(session_id: str, message: Dict[str, Any]) -> None:
    """向专用 WS 连接推送（并清理失效连接）。"""
    conns = list(_automation_ws.get(session_id, set()))
    for ws in conns:
        try:
            await ws.send_json(message)
        except Exception:  # noqa: BLE001
            unregister_ws(session_id, ws)


async def _publish(event_type: str, session_id: str, data: Any) -> None:
    """统一推送：专用 WS + 坐席/员工兜底 WS。"""
    message = _build_message(event_type, session_id, data)
    await _send_to_automation_ws(session_id, message)

    parties = _session_parties.get(session_id, {})
    # 兜底推送给坐席
    if parties.get("agent_id"):
        await ws_manager.send_to_agent(parties["agent_id"], message)
    # 兜底推送给员工
    if parties.get("employee_id"):
        await ws_manager.send_to_employee(parties["employee_id"], message)


async def publish_progress(
    session_id: str, step: str, message: str, action_id: Optional[str] = None
) -> None:
    """推送进度事件。"""
    await _publish(
        AUTOMATION_WS_PROGRESS,
        session_id,
        {"step": step, "message": message, "action_id": action_id},
    )


async def publish_action_required(
    session_id: str,
    action: Any,
    ticket: Any,
) -> None:
    """推送需要审批/确认事件（坐席审批或员工 H5 确认）。"""
    await _publish(
        AUTOMATION_WS_ACTION_REQUIRED,
        session_id,
        {
            "action": {
                "id": action.id,
                "action_type": action.action_type,
                "title": action.title,
                "description": action.description,
                "risk_level": action.risk_level,
                "payload": action.payload,
            },
            "ticket": {
                "id": ticket.id,
                "channel": ticket.channel,
                "status": ticket.status,
                "reason": ticket.reason,
            },
        },
    )


async def publish_resolved(session_id: str, summary: str) -> None:
    """推送处置成功事件。"""
    await _publish(AUTOMATION_WS_RESOLVED, session_id, {"summary": summary})


async def publish_takeover(session_id: str, reason: str) -> None:
    """推送转人工事件。"""
    await _publish(AUTOMATION_WS_TAKEOVER, session_id, {"reason": reason})


async def publish_error(session_id: str, code: int, message: str) -> None:
    """推送异常事件。"""
    await _publish(AUTOMATION_WS_ERROR, session_id, {"code": code, "message": message})


def schedule_silent_close(
    session_id: str, ttl: int = AUTOMATION_SILENT_CLOSE_TTL, on_expire=None
) -> None:
    """调度静默关单：ttl 秒后若会话仍为 resolved，则自动关单。

    Args:
        session_id: 会话ID
        ttl: 静默期秒数（默认 600 = 10 分钟）
        on_expire: 到期回调 coroutine（通常 = AutoSessionService.auto_close）
    """
    # 取消已有的同名任务，避免重复调度
    existing = _silent_close_tasks.get(session_id)
    if existing is not None and not existing.done():
        existing.cancel()

    if on_expire is None:
        return

    async def _wait_and_close() -> None:
        try:
            await asyncio.sleep(ttl)
            await on_expire(session_id)
        except asyncio.CancelledError:  # 被新的调度取消
            pass
        except Exception as e:  # noqa: BLE001
            logger.warning(f"静默关单回调异常 session={session_id}: {e}")
        finally:
            _silent_close_tasks.pop(session_id, None)

    _silent_close_tasks[session_id] = asyncio.create_task(_wait_and_close())


def cancel_silent_close(session_id: str) -> None:
    """取消静默关单调度（如会话已被接管/关单）。"""
    task = _silent_close_tasks.pop(session_id, None)
    if task is not None and not task.done():
        task.cancel()