tracing.py 1.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950
  1. """链路追踪:trace_id 的 ContextVar 载体。
  2. trace_id 在一次请求内贯穿「HTTP 入口 → SSE 生成器 → 工具调用 → 计费落库」,
  3. 用于把散落在各层的日志串成一条线。2.0 的做法是不依赖任何第三方 tracing
  4. SDK:一个 ContextVar 加两个绑定/解绑函数即可,跨 asyncio 任务传播由
  5. ContextVar 自身保证。
  6. """
  7. from __future__ import annotations
  8. import uuid
  9. from contextvars import ContextVar
  10. from typing import Any
  11. # 当前请求的 trace_id;未绑定时为空串(而非 None,省掉调用方的判空)
  12. _trace_id: ContextVar[str] = ContextVar("trace_id", default="")
  13. def new_trace_id() -> str:
  14. """生成一个新的 trace_id(32 位十六进制,无连字符,便于塞进日志与响应头)。"""
  15. return uuid.uuid4().hex
  16. def get_trace_id() -> str:
  17. """读取当前上下文绑定的 trace_id;未绑定返回空串。"""
  18. return _trace_id.get()
  19. def bind_trace_id(trace_id: str | None) -> Any:
  20. """绑定 trace_id,返回可用于还原的 token。
  21. 调用方应在 finally 中 ``unbind_trace_id(token)``。用 token 还原而不是
  22. 直接重置为空串,是为了支持嵌套场景(外层已有 trace 时内层不要抹掉)。
  23. """
  24. return _trace_id.set(trace_id or new_trace_id())
  25. def unbind_trace_id(token: Any) -> None:
  26. """还原到 ``bind_trace_id`` 之前的 trace_id。"""
  27. try:
  28. _trace_id.reset(token)
  29. except (ValueError, LookupError):
  30. # token 已被消费或来自别的上下文:静默降级,不影响主流程
  31. pass
  32. def trace_headers(trace_id: str | None = None) -> dict[str, str]:
  33. """构造向下游(MCP / 外部 HTTP)透传 trace_id 的请求头。"""
  34. tid = trace_id or get_trace_id()
  35. return {"X-Trace-Id": tid} if tid else {}