Neo-ZQYY/apps/backend/app/trace/ws_wrapper.py

# -*- coding: utf-8 -*-
"""
WebSocket 连接追踪辅助函数

提供一组轻量函数，在 WebSocket 端点内部调用，
追踪连接全生命周期（WS_CONNECT → WS_MESSAGE → WS_DISCONNECT）。

所有函数在无活跃 TraceContext 时静默跳过，不影响业务逻辑。
WS_MESSAGE span 每 5 条消息记录一次，避免 span 爆炸。
"""

from __future__ import annotations

from datetime import datetime

from app.trace.context import (
    SpanType,
    TraceContext,
    TraceSpan,
    TraceType,
    create_ws_trace,
    get_current_trace,
    set_current_trace,
)
from app.trace.config import get_trace_config
from app.trace.writer import TraceWriter


def ws_trace_connect(
    execution_id: str,
    client_info: str = "",
) -> TraceContext | None:
    """创建 WS TraceContext 并记录 WS_CONNECT span。

    返回创建的 TraceContext（供后续 span 使用），
    如果 trace 未启用则返回 None。
    """
    cfg = get_trace_config()
    if not cfg.enabled:
        return None

    ctx = create_ws_trace()
    set_current_trace(ctx)

    ctx.add_span(TraceSpan(
        span_type=SpanType.WS_CONNECT,
        module="trace.ws_wrapper",
        function="ws_trace_connect",
        description_zh=f"WebSocket 连接建立: execution_id={execution_id}",
        description_en=f"WebSocket connected: execution_id={execution_id}",
        params={"execution_id": execution_id, "client_info": client_info},
        result_summary="connected",
        duration_ms=0.0,
        timestamp=datetime.now().isoformat(),
    ))
    return ctx


def ws_trace_message(
    message_count: int,
    total_bytes: int,
) -> None:
    """记录 WS_MESSAGE span：每 5 条消息记录一次，避免 span 爆炸。

    仅当 message_count 是 5 的倍数时才记录 span。
    """
    ctx = get_current_trace()
    if ctx is None:
        return

    # 每 5 条消息记录一次
    if message_count % 5 != 0:
        return

    ctx.add_span(TraceSpan(
        span_type=SpanType.WS_MESSAGE,
        module="trace.ws_wrapper",
        function="ws_trace_message",
        description_zh=f"WS 消息推送: 累计 {message_count} 条, {total_bytes} 字节",
        description_en=f"WS message push: cumulative {message_count} msgs, {total_bytes} bytes",
        params={"message_count": message_count, "total_bytes": total_bytes},
        result_summary=f"msgs={message_count}, bytes={total_bytes}",
        duration_ms=0.0,
        timestamp=datetime.now().isoformat(),
    ))


def ws_trace_disconnect(
    reason: str,
    total_messages: int,
    total_duration_ms: float,
) -> None:
    """记录 WS_DISCONNECT span 并写入 trace 日志。"""
    ctx = get_current_trace()
    if ctx is None:
        return

    ctx.add_span(TraceSpan(
        span_type=SpanType.WS_DISCONNECT,
        module="trace.ws_wrapper",
        function="ws_trace_disconnect",
        description_zh=f"WebSocket 断开: 原因={reason}, 总消息={total_messages}, 耗时={total_duration_ms:.0f}ms",
        description_en=f"WebSocket disconnected: reason={reason}, total_msgs={total_messages}, duration={total_duration_ms:.0f}ms",
        params={
            "reason": reason,
            "total_messages": total_messages,
            "total_duration_ms": total_duration_ms,
        },
        result_summary=f"msgs={total_messages}, reason={reason}",
        duration_ms=total_duration_ms,
        timestamp=datetime.now().isoformat(),
    ))

    # 同步写入日志文件（ws_trace_disconnect 可能在同步上下文中调用）
    try:
        from app.trace.writer import serialize_trace, get_log_file
        import json
        from pathlib import Path

        writer = TraceWriter()
        data = serialize_trace(ctx)
        line = json.dumps(data, ensure_ascii=False, default=str)
        dt = ctx.start_time if isinstance(ctx.start_time, datetime) else datetime.now()
        filepath = get_log_file(dt, base_dir=writer.base_dir)
        filepath.parent.mkdir(parents=True, exist_ok=True)
        filepath = writer._rotate_if_needed(filepath)
        writer._sync_write(filepath, line)
    except Exception:
        pass  # 写入失败不影响业务