vent_agent_2.0/obs/tracing.py

"""链路追踪：trace_id 的 ContextVar 载体。

trace_id 在一次请求内贯穿「HTTP 入口 → SSE 生成器 → 工具调用 → 计费落库」，
用于把散落在各层的日志串成一条线。2.0 的做法是不依赖任何第三方 tracing
SDK：一个 ContextVar 加两个绑定/解绑函数即可，跨 asyncio 任务传播由
ContextVar 自身保证。
"""

from __future__ import annotations

import uuid
import re
from contextlib import contextmanager
from contextvars import ContextVar
from typing import Any

# 当前请求的 trace_id；未绑定时为空串（而非 None，省掉调用方的判空）
_trace_id: ContextVar[str] = ContextVar("trace_id", default="")
_context: ContextVar[dict[str, Any]] = ContextVar("log_context", default={})
_TRACE_PATTERN = re.compile(r"[A-Za-z0-9._:-]{1,128}\Z")


def normalize_trace_id(value: str | None) -> str:
    return value if isinstance(value, str) and _TRACE_PATTERN.fullmatch(value) else new_trace_id()


def get_log_context() -> dict[str, Any]:
    return dict(_context.get())


def bind_log_context(**fields):
    return _context.set({**_context.get(), **fields})


def unbind_log_context(token) -> None:
    try:
        _context.reset(token)
    except (ValueError, LookupError):
        pass


@contextmanager
def log_context(**fields):
    token = bind_log_context(**fields)
    try:
        yield
    finally:
        unbind_log_context(token)


def new_trace_id() -> str:
    """生成一个新的 trace_id（32 位十六进制，无连字符，便于塞进日志与响应头）。"""
    return uuid.uuid4().hex


def get_trace_id() -> str:
    """读取当前上下文绑定的 trace_id；未绑定返回空串。"""
    return _trace_id.get()


def bind_trace_id(trace_id: str | None) -> Any:
    """绑定 trace_id，返回可用于还原的 token。

    调用方应在 finally 中 ``unbind_trace_id(token)``。用 token 还原而不是
    直接重置为空串，是为了支持嵌套场景（外层已有 trace 时内层不要抹掉）。
    """
    return _trace_id.set(normalize_trace_id(trace_id))


def unbind_trace_id(token: Any) -> None:
    """还原到 ``bind_trace_id`` 之前的 trace_id。"""
    try:
        _trace_id.reset(token)
    except (ValueError, LookupError):
        # token 已被消费或来自别的上下文：静默降级，不影响主流程
        pass


def trace_headers(trace_id: str | None = None) -> dict[str, str]:
    """构造向下游（MCP / 外部 HTTP）透传 trace_id 的请求头。"""
    tid = trace_id or get_trace_id()
    return {"X-Trace-Id": tid} if tid else {}