XingHuo/app/utils/trace.py

"""trace_id 全链路贯通（contextvars · 架构 §5.1）。

约束（编码前必读）：
- 中间件必须在 call_next 之前 set()，禁止在 endpoint/service 内重新 set()；
- 后台任务（create_task）须显式 contextvars.copy_context()；
- 同步 def 路由跑线程池时 context 经 anyio 传播，由 B8 trace 一致性断言实测兜底。
"""

from __future__ import annotations

from contextvars import ContextVar, Token
from uuid import uuid4

_trace_id: ContextVar[str] = ContextVar("trace_id", default="")


def new_trace(trace_id: str | None = None) -> str:
    """生成并绑定新 trace_id；透传外部 X-Trace-Id 时传入该值。"""
    tid = trace_id or f"trc-{uuid4().hex[:16]}"
    _trace_id.set(tid)
    return tid


def set_trace(trace_id: str) -> Token:
    """绑定外部 trace_id（中间件入口），返回 token 供请求结束复位。"""
    return _trace_id.set(trace_id)


def reset_trace(token: Token) -> None:
    """复位到中间件之前的上下文（防异步任务间串号）。"""
    _trace_id.reset(token)


def current_trace() -> str:
    """读取当前 trace_id；未初始化时返回空串（调用方应兜底生成）。"""
    return _trace_id.get()


def ensure_trace() -> None:
    """无上下文时兜底归因（脚本/引擎入口），有值时保留（中间件场景不重新 set）。"""
    if not current_trace():
        new_trace()