ForcePilot/backend/package/yuxi/utils/trace_context.py

"""请求级 trace_id 上下文（共享模块）。

通过 ContextVar 存储每个请求的 trace_id，供异常处理器、日志适配器
等跨层读取。中间件入口生成（或从 X-Request-Id header 读取）。

本模块为 server 层与 channels 层共享的 trace_id 机制，确保两套
ContextVar（server 层 ``_trace_id_var`` 与 channels 层 ``_current_span``）
保持同步。
"""

from __future__ import annotations

import uuid
from contextvars import ContextVar

_trace_id_var: ContextVar[str | None] = ContextVar("yuxi_request_trace_id", default=None)


def set_trace_id(trace_id: str) -> None:
    """设置当前请求的 trace_id。"""
    _trace_id_var.set(trace_id)


def get_trace_id() -> str | None:
    """获取当前请求的 trace_id，未设置时返回 None。"""
    return _trace_id_var.get()


def generate_trace_id() -> str:
    """生成新的 trace_id（uuid4 十六进制）。"""
    return uuid.uuid4().hex


def reset_trace_id() -> None:
    """请求结束时重置 trace_id。

    ContextVar 在请求结束后自动回收，显式重置作为保险。
    """
    _trace_id_var.set(None)