zhuchaowe/qwen_agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Propagate trace id through MCP metadata

Browse Source
This commit is contained in:
csh28 2026-05-08 18:59:42 +08:00
parent 951948639e
commit 1f06450402
4 changed files with 122 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
agent/deep_assistant.py
View File

@ -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]:

97
agent/mcp_trace_meta.py Normal file
View File

@ -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

16
skills/autoload/onprem/rag-retrieve/rag_retrieve_server.py
View File

@ -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",

18
skills/autoload/support/rag-retrieve/rag_retrieve_server.py
View File

@ -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",