完善消息链路trace并贯通AI与发送动作

utils/ai/unified_llm.py

@@ -15,6 +15,7 @@ from requests import HTTPError
 from loguru import logger
 
 from utils.ai.llm_registry import LLMRegistry
+from utils.trace_context import get_current_trace_id, format_trace_prefix
 
 
 class UnifiedLLMClient:
@@ -58,6 +59,7 @@ class UnifiedLLMClient:
         backend: str,
         scene: str,
         model: str,
+        trace_id: str,
         success: bool,
         latency_ms: float,
         error: str = "",
@@ -70,6 +72,7 @@ class UnifiedLLMClient:
                 "backend": str(backend or "").strip(),
                 "scene": str(scene or "").strip(),
                 "model": str(model or "").strip(),
+                "trace_id": str(trace_id or "").strip(),
                 "success": bool(success),
                 "latency_ms": round(float(latency_ms or 0.0), 2),
                 "error": str(error or "").strip()[:300],
@@ -236,6 +239,7 @@ class UnifiedLLMClient:
         started_at = time.monotonic()
         self.last_error = ""
         result: Optional[Dict[str, Any]] = None
+        current_trace_id = get_current_trace_id()
         if not self.is_available():
             self.last_error = "client_unavailable"
         elif self.provider == "dify":
@@ -272,10 +276,19 @@ class UnifiedLLMClient:
             backend=str(self.config.get("backend", "") or ""),
             scene=str(self.config.get("scene", "") or ""),
             model=self.model or str(self.mode or ""),
+            trace_id=current_trace_id,
             success=bool(result and result.get("text")),
             latency_ms=latency_ms,
             error=self.last_error,
         )
+        # 在统一出口补一条轻量 trace 日志，方便把“消息 -> AI 调用”快速串起来。
+        self.LOG.debug(
+            f"{format_trace_prefix(current_trace_id)}LLM调用结束 "
+            f"provider={self.provider} backend={self.config.get('backend', '') or '-'} "
+            f"scene={self.config.get('scene', '') or '-'} "
+            f"success={bool(result and result.get('text'))} latency_ms={round(latency_ms, 2)} "
+            f"error={self.last_error or '-'}"
+        )
         return result
 
     def _generate_openai(

utils/trace_context.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from contextvars import ContextVar, Token
+
+# 当前消息链路的 trace_id：
+# 1. 使用 ContextVar 而不是全局变量，避免并发消息之间互相覆盖；
+# 2. asyncio Task 会自动继承上下文，因此插件里再起的协程也能拿到同一个 trace_id；
+# 3. 对不在消息主链路里的后台任务，该值默认为空字符串，不会影响原有逻辑。
+_CURRENT_TRACE_ID: ContextVar[str] = ContextVar("current_trace_id", default="")
+
+
+def set_current_trace_id(trace_id: str) -> Token:
+    """写入当前上下文的 trace_id，并返回可回滚 token。"""
+    return _CURRENT_TRACE_ID.set(str(trace_id or "").strip())
+
+
+def reset_current_trace_id(token: Token) -> None:
+    """根据 token 回滚 trace_id，避免上下文泄漏到后续无关任务。"""
+    _CURRENT_TRACE_ID.reset(token)
+
+
+def get_current_trace_id() -> str:
+    """读取当前上下文中的 trace_id。"""
+    return str(_CURRENT_TRACE_ID.get("") or "").strip()
+
+
+def format_trace_prefix(trace_id: str = "") -> str:
+    """统一生成日志前缀，避免各模块手写格式不一致。"""
+    resolved_trace_id = str(trace_id or get_current_trace_id() or "").strip()
+    if not resolved_trace_id:
+        return ""
+    return f"[trace_id={resolved_trace_id}] "