Propagate trace id through MCP metadata

agent/deep_assistant.py

@@ -24,6 +24,7 @@ from .guideline_middleware import GuidelineMiddleware
 from .tool_output_length_middleware import ToolOutputLengthMiddleware
 from .tool_use_cleanup_middleware import ToolUseCleanupMiddleware
 from .filepath_fix_middleware import FilePathFixMiddleware
+from .mcp_trace_meta import patch_mcp_client_session_trace_meta
 from utils.settings import (
     SUMMARIZATION_MAX_TOKENS,
     SUMMARIZATION_TOKENS_TO_KEEP,
@@ -123,6 +124,7 @@ def read_system_prompt():
 
 async def get_tools_from_mcp(mcp):
     """Extract tools from MCP configuration with caching."""
+    patch_mcp_client_session_trace_meta()
     start_time = time.time()
     # Defensive handling: ensure mcp is a non-empty list containing mcpServers
     if not isinstance(mcp, list) or len(mcp) == 0 or "mcpServers" not in mcp[0]:

agent/mcp_trace_meta.py

@@ -0,0 +1,97 @@
+import logging
+from functools import wraps
+from typing import Any
+
+try:
+    from mcp import ClientSession, types
+except ImportError:
+    from mcp.client.session import ClientSession
+    from mcp import types
+
+from utils.log_util.context import g
+
+logger = logging.getLogger("app")
+
+_PATCHED_ATTR = "_catalog_trace_meta_patched"
+_TRACE_META_TOOL_NAMES = {"rag_retrieve", "table_rag_retrieve"}
+
+
+def _get_trace_id() -> str:
+    try:
+        trace_id = getattr(g, "trace_id", "")
+    except (LookupError, KeyError):
+        return ""
+    return str(trace_id) if trace_id else ""
+
+
+def _get_tool_name(args: tuple[Any, ...], kwargs: dict[str, Any]) -> str:
+    name = args[0] if args else kwargs.get("name")
+    return str(name) if name else ""
+
+
+def patch_mcp_client_session_trace_meta() -> None:
+    """Attach catalog trace id to MCP tools/call params._meta."""
+    if getattr(ClientSession.call_tool, _PATCHED_ATTR, False):
+        return
+
+    original_call_tool = ClientSession.call_tool
+
+    @wraps(original_call_tool)
+    async def call_tool_with_trace_meta(self: ClientSession, *args: Any, **kwargs: Any) -> Any:
+        tool_name = _get_tool_name(args, kwargs)
+        trace_id = _get_trace_id() if tool_name in _TRACE_META_TOOL_NAMES else ""
+        if trace_id:
+            meta = kwargs.get("meta")
+            if isinstance(meta, dict):
+                meta = {**meta, "trace_id": meta.get("trace_id") or trace_id}
+            else:
+                meta = {"trace_id": trace_id}
+            kwargs["meta"] = meta
+
+        try:
+            return await original_call_tool(self, *args, **kwargs)
+        except TypeError as exc:
+            if trace_id and "meta" in kwargs and "unexpected keyword argument" in str(exc):
+                return await _call_tool_with_meta_compat(self, *args, **kwargs)
+            raise
+
+    setattr(call_tool_with_trace_meta, _PATCHED_ATTR, True)
+    ClientSession.call_tool = call_tool_with_trace_meta
+
+
+async def _call_tool_with_meta_compat(self: ClientSession, *args: Any, **kwargs: Any) -> Any:
+    """Call tools/call with _meta for MCP SDK versions before call_tool(meta=...)."""
+    name = _get_tool_name(args, kwargs)
+    if not name:
+        raise TypeError("call_tool() missing required argument: 'name'")
+
+    arguments = args[1] if len(args) > 1 else kwargs.get("arguments", kwargs.get("args"))
+    read_timeout_seconds = (
+        args[2] if len(args) > 2 else kwargs.get("read_timeout_seconds")
+    )
+    progress_callback = (
+        args[3] if len(args) > 3 else kwargs.get("progress_callback")
+    )
+    meta = kwargs.get("meta")
+
+    request_meta = meta if isinstance(meta, dict) else None
+    result = await self.send_request(
+        types.ClientRequest(
+            types.CallToolRequest(
+                params=types.CallToolRequestParams(
+                    name=name,
+                    arguments=arguments,
+                    _meta=request_meta,
+                ),
+            )
+        ),
+        types.CallToolResult,
+        request_read_timeout_seconds=read_timeout_seconds,
+        progress_callback=progress_callback,
+    )
+
+    validate_tool_result = getattr(self, "_validate_tool_result", None)
+    if validate_tool_result and not result.isError:
+        await validate_tool_result(name, result)
+
+    return result

skills/autoload/onprem/rag-retrieve/rag_retrieve_server.py

@@ -73,7 +73,7 @@ Format: `<CITATION file="file_id" filename="name.xlsx" sheet=1 rows=[2, 4] />`
 
 """
 
-def rag_retrieve(query: str, top_k: int = 100) -> Dict[str, Any]:
+def rag_retrieve(query: str, top_k: int = 100, trace_id: str = "") -> Dict[str, Any]:
     """Call the RAG retrieval API."""
     try:
         bot_id = ""
@@ -100,6 +100,8 @@ def rag_retrieve(query: str, top_k: int = 100) -> Dict[str, Any]:
             "content-type": "application/json",
             "authorization": f"Bearer {auth_token}"
         }
+        if trace_id:
+            headers["X-Request-ID"] = trace_id
         data = {
             "query": query,
             "top_k": top_k
@@ -172,7 +174,7 @@ def rag_retrieve(query: str, top_k: int = 100) -> Dict[str, Any]:
         }
 
 
-def table_rag_retrieve(query: str) -> Dict[str, Any]:
+def table_rag_retrieve(query: str, trace_id: str = "") -> Dict[str, Any]:
     """Call the Table RAG retrieval API."""
     try:
         bot_id = ""
@@ -189,6 +191,8 @@ def table_rag_retrieve(query: str) -> Dict[str, Any]:
             "content-type": "application/json",
             "authorization": f"Bearer {auth_token}"
         }
+        if trace_id:
+            headers["X-Request-ID"] = trace_id
         data = {
             "query": query,
         }
@@ -220,7 +224,7 @@ def table_rag_retrieve(query: str) -> Dict[str, Any]:
         if "markdown" in response_data:
             markdown_content = response_data["markdown"]
             if re.search(r"^no excel files found", markdown_content, re.IGNORECASE):
-                rag_result = rag_retrieve(query)
+                rag_result = rag_retrieve(query, trace_id=trace_id)
                 content = rag_result.get("content", [])
                 if content and content[0].get("type") == "text":
                     content[0]["text"] = "No table_rag_retrieve results were found. The content below is the fallback result from rag_retrieve:\n\n" + content[0]["text"]
@@ -302,6 +306,8 @@ async def handle_request(request: Dict[str, Any]) -> Dict[str, Any]:
         elif method == "tools/call":
             tool_name = params.get("name")
             arguments = params.get("arguments", {})
+            meta = params.get("_meta") or params.get("meta") or {}
+            trace_id = meta.get("trace_id", "") if isinstance(meta, dict) else ""
 
             if tool_name == "rag_retrieve":
                 query = arguments.get("query", "")
@@ -310,7 +316,7 @@ async def handle_request(request: Dict[str, Any]) -> Dict[str, Any]:
                 if not query:
                     return create_error_response(request_id, -32602, "Missing required parameter: query")
 
-                result = rag_retrieve(query, top_k)
+                result = rag_retrieve(query, top_k, trace_id)
 
                 return {
                     "jsonrpc": "2.0",
@@ -324,7 +330,7 @@ async def handle_request(request: Dict[str, Any]) -> Dict[str, Any]:
                 if not query:
                     return create_error_response(request_id, -32602, "Missing required parameter: query")
 
-                result = table_rag_retrieve(query)
+                result = table_rag_retrieve(query, trace_id)
 
                 return {
                     "jsonrpc": "2.0",

skills/autoload/support/rag-retrieve/rag_retrieve_server.py

@@ -73,7 +73,7 @@ Format: `<CITATION file="file_id" filename="name.xlsx" sheet=1 rows=[2, 4] />`
 
 """
 
-def rag_retrieve(query: str, top_k: int = 100) -> Dict[str, Any]:
+def rag_retrieve(query: str, top_k: int = 100, trace_id: str = "") -> Dict[str, Any]:
     """Call the RAG retrieval API."""
     try:
         bot_id = ""
@@ -100,6 +100,8 @@ def rag_retrieve(query: str, top_k: int = 100) -> Dict[str, Any]:
             "content-type": "application/json",
             "authorization": f"Bearer {auth_token}"
         }
+        if trace_id:
+            headers["X-Request-ID"] = trace_id
         data = {
             "query": query,
             "top_k": top_k
@@ -172,7 +174,7 @@ def rag_retrieve(query: str, top_k: int = 100) -> Dict[str, Any]:
         }
 
 
-def table_rag_retrieve(query: str) -> Dict[str, Any]:
+def table_rag_retrieve(query: str, trace_id: str = "") -> Dict[str, Any]:
     """Call the Table RAG retrieval API."""
     try:
         bot_id = ""
@@ -189,6 +191,8 @@ def table_rag_retrieve(query: str) -> Dict[str, Any]:
             "content-type": "application/json",
             "authorization": f"Bearer {auth_token}"
         }
+        if trace_id:
+            headers["X-Request-ID"] = trace_id
         data = {
             "query": query,
         }
@@ -220,7 +224,7 @@ def table_rag_retrieve(query: str) -> Dict[str, Any]:
         if "markdown" in response_data:
             markdown_content = response_data["markdown"]
             if re.search(r"^no excel files found", markdown_content, re.IGNORECASE):
-                rag_result = rag_retrieve(query)
+                rag_result = rag_retrieve(query, trace_id=trace_id)
                 content = rag_result.get("content", [])
                 if content and content[0].get("type") == "text":
                     content[0]["text"] = "No table_rag_retrieve results were found. The content below is the fallback result from rag_retrieve：\n\n" + content[0]["text"]
@@ -302,7 +306,9 @@ async def handle_request(request: Dict[str, Any]) -> Dict[str, Any]:
         elif method == "tools/call":
             tool_name = params.get("name")
             arguments = params.get("arguments", {})
-            
+            meta = params.get("_meta") or params.get("meta") or {}
+            trace_id = meta.get("trace_id", "") if isinstance(meta, dict) else ""
+
             if tool_name == "rag_retrieve":
                 query = arguments.get("query", "")
                 top_k = arguments.get("top_k", 100)
@@ -310,7 +316,7 @@ async def handle_request(request: Dict[str, Any]) -> Dict[str, Any]:
                 if not query:
                     return create_error_response(request_id, -32602, "Missing required parameter: query")
 
-                result = rag_retrieve(query, top_k)
+                result = rag_retrieve(query, top_k, trace_id)
 
                 return {
                     "jsonrpc": "2.0",
@@ -324,7 +330,7 @@ async def handle_request(request: Dict[str, Any]) -> Dict[str, Any]:
                 if not query:
                     return create_error_response(request_id, -32602, "Missing required parameter: query")
 
-                result = table_rag_retrieve(query)
+                result = table_rag_retrieve(query, trace_id)
 
                 return {
                     "jsonrpc": "2.0",