from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware
import uuid
from contextvars import ContextVar

# 创建一个上下文变量来存储 trace_id
trace_id_var = ContextVar('trace_id', default=None)

class TraceMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next) -> Response:
        # 生成唯一的 trace_id
        trace_id = str(uuid.uuid4())
        
        # 设置上下文变量
        token = trace_id_var.set(trace_id)
        
        try:
            # 处理请求
            response = await call_next(request)
            # 将 trace_id 添加到响应头中
            response.headers['X-Trace-ID'] = trace_id
            return response
        finally:
            # 清理上下文变量
            trace_id_var.reset(token)

def get_trace_id():
    """获取当前请求的 trace_id"""
    return trace_id_var.get()