diff --git a/agent/deep_assistant.py b/agent/deep_assistant.py
index fa92c3f..7415413 100644
--- a/agent/deep_assistant.py
+++ b/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]:
diff --git a/agent/mcp_trace_meta.py b/agent/mcp_trace_meta.py
new file mode 100644
index 0000000..4e4b5c7
--- /dev/null
+++ b/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
diff --git a/skills/autoload/onprem/rag-retrieve/rag_retrieve_server.py b/skills/autoload/onprem/rag-retrieve/rag_retrieve_server.py
index 2575644..6f308e2 100644
--- a/skills/autoload/onprem/rag-retrieve/rag_retrieve_server.py
+++ b/skills/autoload/onprem/rag-retrieve/rag_retrieve_server.py
@@ -73,7 +73,7 @@ Format: ``
"""
-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",
diff --git a/skills/autoload/support/rag-retrieve/rag_retrieve_server.py b/skills/autoload/support/rag-retrieve/rag_retrieve_server.py
index 671a456..09e0924 100644
--- a/skills/autoload/support/rag-retrieve/rag_retrieve_server.py
+++ b/skills/autoload/support/rag-retrieve/rag_retrieve_server.py
@@ -73,7 +73,7 @@ Format: ``
"""
-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",