admin/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

feat(trace): support external trace id propagation (#22623)

Browse Source 
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
qfl
2025-07-22 15:17:43 +08:00
committed by GitHub
parent c987001a19
commit 841e53dbbe
25 changed files with 236 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
api/controllers/service_api/app/completion.py
View File

@@ -1,5 +1,6 @@
import logging
from flask import request
from flask_restful import Resource, reqparse
from werkzeug.exceptions import InternalServerError, NotFound
@@ -23,6 +24,7 @@ from core.errors.error import (
ProviderTokenNotInitError,
QuotaExceededError,
)
from core.helper.trace_id_helper import get_external_trace_id
from core.model_runtime.errors.invoke import InvokeError
from libs import helper
from libs.helper import uuid_value
@@ -111,6 +113,10 @@ class ChatApi(Resource):
args = parser.parse_args()
external_trace_id = get_external_trace_id(request)
if external_trace_id:
args["external_trace_id"] = external_trace_id
streaming = args["response_mode"] == "streaming"
try:

6
api/controllers/service_api/app/workflow.py
View File

@@ -1,6 +1,7 @@
import logging
from dateutil.parser import isoparse
from flask import request
from flask_restful import Resource, fields, marshal_with, reqparse
from flask_restful.inputs import int_range
from sqlalchemy.orm import Session, sessionmaker
@@ -23,6 +24,7 @@ from core.errors.error import (
ProviderTokenNotInitError,
QuotaExceededError,
)
from core.helper.trace_id_helper import get_external_trace_id
from core.model_runtime.errors.invoke import InvokeError
from core.workflow.entities.workflow_execution import WorkflowExecutionStatus
from extensions.ext_database import db
@@ -90,7 +92,9 @@ class WorkflowRunApi(Resource):
parser.add_argument("files", type=list, required=False, location="json")
parser.add_argument("response_mode", type=str, choices=["blocking", "streaming"], location="json")
args = parser.parse_args()
external_trace_id = get_external_trace_id(request)
if external_trace_id:
args["external_trace_id"] = external_trace_id
streaming = args.get("response_mode") == "streaming"
try:

6
api/core/app/apps/advanced_chat/app_generator.py
View File

@@ -23,6 +23,7 @@ from core.app.apps.message_based_app_generator import MessageBasedAppGenerator
from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager
from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom
from core.app.entities.task_entities import ChatbotAppBlockingResponse, ChatbotAppStreamResponse
from core.helper.trace_id_helper import extract_external_trace_id_from_args
from core.model_runtime.errors.invoke import InvokeAuthorizationError
from core.ops.ops_trace_manager import TraceQueueManager
from core.prompt.utils.get_thread_messages_length import get_thread_messages_length
@@ -112,7 +113,10 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
query = query.replace("\x00", "")
inputs = args["inputs"]
extras = {"auto_generate_conversation_name": args.get("auto_generate_name", False)}
extras = {
"auto_generate_conversation_name": args.get("auto_generate_name", False),
**extract_external_trace_id_from_args(args),
}
# get conversation
conversation = None

4
api/core/app/apps/advanced_chat/generate_task_pipeline.py
View File

@@ -559,6 +559,7 @@ class AdvancedChatAppGenerateTaskPipeline:
outputs=event.outputs,
conversation_id=self._conversation_id,
trace_manager=trace_manager,
external_trace_id=self._application_generate_entity.extras.get("external_trace_id"),
)
workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response(
session=session,
@@ -590,6 +591,7 @@ class AdvancedChatAppGenerateTaskPipeline:
exceptions_count=event.exceptions_count,
conversation_id=None,
trace_manager=trace_manager,
external_trace_id=self._application_generate_entity.extras.get("external_trace_id"),
)
workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response(
session=session,
@@ -622,6 +624,7 @@ class AdvancedChatAppGenerateTaskPipeline:
conversation_id=self._conversation_id,
trace_manager=trace_manager,
exceptions_count=event.exceptions_count,
external_trace_id=self._application_generate_entity.extras.get("external_trace_id"),
)
workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response(
session=session,
@@ -653,6 +656,7 @@ class AdvancedChatAppGenerateTaskPipeline:
error_message=event.get_stop_reason(),
conversation_id=self._conversation_id,
trace_manager=trace_manager,
external_trace_id=self._application_generate_entity.extras.get("external_trace_id"),
)
workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response(
session=session,

6
api/core/app/apps/workflow/app_generator.py
View File

@@ -22,6 +22,7 @@ from core.app.apps.workflow.generate_response_converter import WorkflowAppGenera
from core.app.apps.workflow.generate_task_pipeline import WorkflowAppGenerateTaskPipeline
from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity
from core.app.entities.task_entities import WorkflowAppBlockingResponse, WorkflowAppStreamResponse
from core.helper.trace_id_helper import extract_external_trace_id_from_args
from core.model_runtime.errors.invoke import InvokeAuthorizationError
from core.ops.ops_trace_manager import TraceQueueManager
from core.repositories import DifyCoreRepositoryFactory
@@ -123,6 +124,10 @@ class WorkflowAppGenerator(BaseAppGenerator):
)
inputs: Mapping[str, Any] = args["inputs"]
extras = {
**extract_external_trace_id_from_args(args),
}
workflow_run_id = str(uuid.uuid4())
# init application generate entity
application_generate_entity = WorkflowAppGenerateEntity(
@@ -142,6 +147,7 @@ class WorkflowAppGenerator(BaseAppGenerator):
call_depth=call_depth,
trace_manager=trace_manager,
workflow_execution_id=workflow_run_id,
extras=extras,
)
contexts.plugin_tool_providers.set({})

3
api/core/app/apps/workflow/generate_task_pipeline.py
View File

@@ -490,6 +490,7 @@ class WorkflowAppGenerateTaskPipeline:
outputs=event.outputs,
conversation_id=None,
trace_manager=trace_manager,
external_trace_id=self._application_generate_entity.extras.get("external_trace_id"),
)
# save workflow app log
@@ -524,6 +525,7 @@ class WorkflowAppGenerateTaskPipeline:
exceptions_count=event.exceptions_count,
conversation_id=None,
trace_manager=trace_manager,
external_trace_id=self._application_generate_entity.extras.get("external_trace_id"),
)
# save workflow app log
@@ -561,6 +563,7 @@ class WorkflowAppGenerateTaskPipeline:
conversation_id=None,
trace_manager=trace_manager,
exceptions_count=event.exceptions_count if isinstance(event, QueueWorkflowFailedEvent) else 0,
external_trace_id=self._application_generate_entity.extras.get("external_trace_id"),
)
# save workflow app log

42
api/core/helper/trace_id_helper.py Normal file
View File

@@ -0,0 +1,42 @@
import re
from collections.abc import Mapping
from typing import Any, Optional
def is_valid_trace_id(trace_id: str) -> bool:
"""
Check if the trace_id is valid.
Requirements: 1-128 characters, only letters, numbers, '-', and '_'.
"""
return bool(re.match(r"^[a-zA-Z0-9\-_]{1,128}$", trace_id))
def get_external_trace_id(request: Any) -> Optional[str]:
"""
Retrieve the trace_id from the request.
Priority: header ('X-Trace-Id'), then parameters, then JSON body. Returns None if not provided or invalid.
"""
trace_id = request.headers.get("X-Trace-Id")
if not trace_id:
trace_id = request.args.get("trace_id")
if not trace_id and getattr(request, "is_json", False):
json_data = getattr(request, "json", None)
if json_data:
trace_id = json_data.get("trace_id")
if isinstance(trace_id, str) and is_valid_trace_id(trace_id):
return trace_id
return None
def extract_external_trace_id_from_args(args: Mapping[str, Any]) -> dict:
"""
Extract 'external_trace_id' from args.
Returns a dict suitable for use in extras. Returns an empty dict if not found.
"""
trace_id = args.get("external_trace_id")
if trace_id:
return {"external_trace_id": trace_id}
return {}

3
api/core/ops/aliyun_trace/aliyun_trace.py
View File

@@ -101,7 +101,8 @@ class AliyunDataTrace(BaseTraceInstance):
raise ValueError(f"Aliyun get run url failed: {str(e)}")
def workflow_trace(self, trace_info: WorkflowTraceInfo):
trace_id = convert_to_trace_id(trace_info.workflow_run_id)
external_trace_id = trace_info.metadata.get("external_trace_id")
trace_id = external_trace_id or convert_to_trace_id(trace_info.workflow_run_id)
workflow_span_id = convert_to_span_id(trace_info.workflow_run_id, "workflow")
self.add_workflow_span(trace_id, workflow_span_id, trace_info)

3
api/core/ops/arize_phoenix_trace/arize_phoenix_trace.py
View File

@@ -153,7 +153,8 @@ class ArizePhoenixDataTrace(BaseTraceInstance):
}
workflow_metadata.update(trace_info.metadata)
trace_id = uuid_to_trace_id(trace_info.workflow_run_id)
external_trace_id = trace_info.metadata.get("external_trace_id")
trace_id = external_trace_id or uuid_to_trace_id(trace_info.workflow_run_id)
span_id = RandomIdGenerator().generate_span_id()
context = SpanContext(
trace_id=trace_id,

5
api/core/ops/langfuse_trace/langfuse_trace.py
View File

@@ -67,13 +67,14 @@ class LangFuseDataTrace(BaseTraceInstance):
self.generate_name_trace(trace_info)
def workflow_trace(self, trace_info: WorkflowTraceInfo):
trace_id = trace_info.workflow_run_id
external_trace_id = trace_info.metadata.get("external_trace_id")
trace_id = external_trace_id or trace_info.workflow_run_id
user_id = trace_info.metadata.get("user_id")
metadata = trace_info.metadata
metadata["workflow_app_log_id"] = trace_info.workflow_app_log_id
if trace_info.message_id:
trace_id = trace_info.message_id
trace_id = external_trace_id or trace_info.message_id
name = TraceTaskName.MESSAGE_TRACE.value
trace_data = LangfuseTrace(
id=trace_id,

3
api/core/ops/langsmith_trace/langsmith_trace.py
View File

@@ -65,7 +65,8 @@ class LangSmithDataTrace(BaseTraceInstance):
self.generate_name_trace(trace_info)
def workflow_trace(self, trace_info: WorkflowTraceInfo):
trace_id = trace_info.message_id or trace_info.workflow_run_id
external_trace_id = trace_info.metadata.get("external_trace_id")
trace_id = external_trace_id or trace_info.message_id or trace_info.workflow_run_id
if trace_info.start_time is None:
trace_info.start_time = datetime.now()
message_dotted_order = (

5
api/core/ops/opik_trace/opik_trace.py
View File

@@ -96,7 +96,8 @@ class OpikDataTrace(BaseTraceInstance):
self.generate_name_trace(trace_info)
def workflow_trace(self, trace_info: WorkflowTraceInfo):
dify_trace_id = trace_info.workflow_run_id
external_trace_id = trace_info.metadata.get("external_trace_id")
dify_trace_id = external_trace_id or trace_info.workflow_run_id
opik_trace_id = prepare_opik_uuid(trace_info.start_time, dify_trace_id)
workflow_metadata = wrap_metadata(
trace_info.metadata, message_id=trace_info.message_id, workflow_app_log_id=trace_info.workflow_app_log_id
@@ -104,7 +105,7 @@ class OpikDataTrace(BaseTraceInstance):
root_span_id = None
if trace_info.message_id:
dify_trace_id = trace_info.message_id
dify_trace_id = external_trace_id or trace_info.message_id
opik_trace_id = prepare_opik_uuid(trace_info.start_time, dify_trace_id)
trace_data = {

4
api/core/ops/ops_trace_manager.py
View File

@@ -520,6 +520,10 @@ class TraceTask:
"app_id": workflow_run.app_id,
}
external_trace_id = self.kwargs.get("external_trace_id")
if external_trace_id:
metadata["external_trace_id"] = external_trace_id
workflow_trace_info = WorkflowTraceInfo(
workflow_data=workflow_run.to_dict(),
conversation_id=conversation_id,

3
api/core/ops/weave_trace/weave_trace.py
View File

@@ -87,7 +87,8 @@ class WeaveDataTrace(BaseTraceInstance):
self.generate_name_trace(trace_info)
def workflow_trace(self, trace_info: WorkflowTraceInfo):
trace_id = trace_info.message_id or trace_info.workflow_run_id
external_trace_id = trace_info.metadata.get("external_trace_id")
trace_id = external_trace_id or trace_info.message_id or trace_info.workflow_run_id
if trace_info.start_time is None:
trace_info.start_time = datetime.now()

11
api/core/workflow/workflow_cycle_manager.py
View File

@@ -85,6 +85,7 @@ class WorkflowCycleManager:
outputs: Mapping[str, Any] | None = None,
conversation_id: Optional[str] = None,
trace_manager: Optional[TraceQueueManager] = None,
external_trace_id: Optional[str] = None,
) -> WorkflowExecution:
workflow_execution = self._get_workflow_execution_or_raise_error(workflow_run_id)
@@ -96,7 +97,7 @@ class WorkflowCycleManager:
total_steps=total_steps,
)
self._add_trace_task_if_needed(trace_manager, workflow_execution, conversation_id)
self._add_trace_task_if_needed(trace_manager, workflow_execution, conversation_id, external_trace_id)
self._workflow_execution_repository.save(workflow_execution)
return workflow_execution
@@ -111,6 +112,7 @@ class WorkflowCycleManager:
exceptions_count: int = 0,
conversation_id: Optional[str] = None,
trace_manager: Optional[TraceQueueManager] = None,
external_trace_id: Optional[str] = None,
) -> WorkflowExecution:
execution = self._get_workflow_execution_or_raise_error(workflow_run_id)
@@ -123,7 +125,7 @@ class WorkflowCycleManager:
exceptions_count=exceptions_count,
)
self._add_trace_task_if_needed(trace_manager, execution, conversation_id)
self._add_trace_task_if_needed(trace_manager, execution, conversation_id, external_trace_id)
self._workflow_execution_repository.save(execution)
return execution
@@ -139,6 +141,7 @@ class WorkflowCycleManager:
conversation_id: Optional[str] = None,
trace_manager: Optional[TraceQueueManager] = None,
exceptions_count: int = 0,
external_trace_id: Optional[str] = None,
) -> WorkflowExecution:
workflow_execution = self._get_workflow_execution_or_raise_error(workflow_run_id)
now = naive_utc_now()
@@ -154,7 +157,7 @@ class WorkflowCycleManager:
)
self._fail_running_node_executions(workflow_execution.id_, error_message, now)
self._add_trace_task_if_needed(trace_manager, workflow_execution, conversation_id)
self._add_trace_task_if_needed(trace_manager, workflow_execution, conversation_id, external_trace_id)
self._workflow_execution_repository.save(workflow_execution)
return workflow_execution
@@ -312,6 +315,7 @@ class WorkflowCycleManager:
trace_manager: Optional[TraceQueueManager],
workflow_execution: WorkflowExecution,
conversation_id: Optional[str],
external_trace_id: Optional[str],
) -> None:
"""Add trace task if trace manager is provided."""
if trace_manager:
@@ -321,6 +325,7 @@ class WorkflowCycleManager:
workflow_execution=workflow_execution,
conversation_id=conversation_id,
user_id=trace_manager.user_id,
external_trace_id=external_trace_id,
)
)

86
api/tests/unit_tests/core/helper/test_trace_id_helper.py Normal file
View File

@@ -0,0 +1,86 @@
import pytest
from core.helper.trace_id_helper import extract_external_trace_id_from_args, get_external_trace_id, is_valid_trace_id
class DummyRequest:
def __init__(self, headers=None, args=None, json=None, is_json=False):
self.headers = headers or {}
self.args = args or {}
self.json = json
self.is_json = is_json
class TestTraceIdHelper:
"""Test cases for trace_id_helper.py"""
@pytest.mark.parametrize(
("trace_id", "expected"),
[
("abc123", True),
("A-B_C-123", True),
("a" * 128, True),
("", False),
("a" * 129, False),
("abc!@#", False),
("空格", False),
("with space", False),
],
)
def test_is_valid_trace_id(self, trace_id, expected):
"""Test trace_id validation for various cases"""
assert is_valid_trace_id(trace_id) is expected
def test_get_external_trace_id_from_header(self):
"""Should extract valid trace_id from header"""
req = DummyRequest(headers={"X-Trace-Id": "abc123"})
assert get_external_trace_id(req) == "abc123"
def test_get_external_trace_id_from_args(self):
"""Should extract valid trace_id from args if header missing"""
req = DummyRequest(args={"trace_id": "abc123"})
assert get_external_trace_id(req) == "abc123"
def test_get_external_trace_id_from_json(self):
"""Should extract valid trace_id from JSON body if header and args missing"""
req = DummyRequest(is_json=True, json={"trace_id": "abc123"})
assert get_external_trace_id(req) == "abc123"
def test_get_external_trace_id_priority(self):
"""Header > args > json priority"""
req = DummyRequest(
headers={"X-Trace-Id": "header_id"},
args={"trace_id": "args_id"},
is_json=True,
json={"trace_id": "json_id"},
)
assert get_external_trace_id(req) == "header_id"
req2 = DummyRequest(args={"trace_id": "args_id"}, is_json=True, json={"trace_id": "json_id"})
assert get_external_trace_id(req2) == "args_id"
req3 = DummyRequest(is_json=True, json={"trace_id": "json_id"})
assert get_external_trace_id(req3) == "json_id"
@pytest.mark.parametrize(
"req",
[
DummyRequest(headers={"X-Trace-Id": "!!!"}),
DummyRequest(args={"trace_id": "!!!"}),
DummyRequest(is_json=True, json={"trace_id": "!!!"}),
DummyRequest(),
],
)
def test_get_external_trace_id_invalid(self, req):
"""Should return None for invalid or missing trace_id"""
assert get_external_trace_id(req) is None
@pytest.mark.parametrize(
("args", "expected"),
[
({"external_trace_id": "abc123"}, {"external_trace_id": "abc123"}),
({"other": "value"}, {}),
({}, {}),
],
)
def test_extract_external_trace_id_from_args(self, args, expected):
"""Test extraction of external_trace_id from args mapping"""
assert extract_external_trace_id_from_args(args) == expected