diff --git a/api/.env.example b/api/.env.example index f98f2ec4a25..0896ee87c89 100644 --- a/api/.env.example +++ b/api/.env.example @@ -98,9 +98,11 @@ REDIS_KEEPALIVE=true CELERY_BROKER_URL=redis://:difyai123456@localhost:${REDIS_PORT}/1 CELERY_BACKEND=redis -# Ops trace retry configuration -OPS_TRACE_RETRYABLE_DISPATCH_MAX_RETRIES=60 +# Ops trace configuration +OPS_TRACE_UNIFIED_ENABLED=false +OPS_TRACE_RETRYABLE_DISPATCH_MAX_RETRIES=300 OPS_TRACE_RETRYABLE_DISPATCH_DELAY_SECONDS=5 +OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS=1800 # Database configuration DB_TYPE=postgresql diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py index 06ae6c2718e..c3f1be5d045 100644 --- a/api/configs/feature/__init__.py +++ b/api/configs/feature/__init__.py @@ -12,6 +12,7 @@ from pydantic import ( PositiveInt, computed_field, field_validator, + model_validator, ) from pydantic_settings import BaseSettings @@ -1336,9 +1337,16 @@ class NewAgentBetaConfig(BaseSettings): class OpsTraceConfig(BaseSettings): + OPS_TRACE_UNIFIED_ENABLED: bool = Field( + description="Enable unified ops tracing for providers registered in the unified registry.", + default=False, + ) + + # Include scheduling and export grace after the parent workflow's maximum execution time. + # Recommended: max_retries >= ceil((WORKFLOW_MAX_EXECUTION_TIME + grace_seconds) / delay_seconds). OPS_TRACE_RETRYABLE_DISPATCH_MAX_RETRIES: PositiveInt = Field( description="Maximum retry attempts for transient ops trace provider dispatch failures.", - default=60, + default=300, ) OPS_TRACE_RETRYABLE_DISPATCH_DELAY_SECONDS: PositiveInt = Field( @@ -1346,6 +1354,20 @@ class OpsTraceConfig(BaseSettings): default=5, ) + OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS: PositiveInt = Field( + description="Retention in seconds for unified tracing parent contexts.", + default=1800, + ) + + @model_validator(mode="after") + def validate_parent_context_retention(self) -> "OpsTraceConfig": + if not self.OPS_TRACE_UNIFIED_ENABLED: + return self + retry_window = self.OPS_TRACE_RETRYABLE_DISPATCH_MAX_RETRIES * self.OPS_TRACE_RETRYABLE_DISPATCH_DELAY_SECONDS + if retry_window > self.OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS: + raise ValueError("OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS must cover the retry window") + return self + class CeleryBeatConfig(BaseSettings): CELERY_BEAT_SCHEDULER_TIME: int = Field( diff --git a/api/core/app/apps/advanced_chat/generate_task_pipeline.py b/api/core/app/apps/advanced_chat/generate_task_pipeline.py index 77539b103f1..a8c333f22c2 100644 --- a/api/core/app/apps/advanced_chat/generate_task_pipeline.py +++ b/api/core/app/apps/advanced_chat/generate_task_pipeline.py @@ -229,7 +229,9 @@ class AdvancedChatAppGenerateTaskPipeline(GraphRuntimeStateSupport): :return: """ self._conversation_name_generate_thread = self._message_cycle_manager.generate_conversation_name( - conversation_id=self._conversation_id, query=self._application_generate_entity.query + conversation_id=self._conversation_id, + query=self._application_generate_entity.query, + message_id=self._message_id, ) generator = self._wrapper_process_stream_response(trace_manager=self._application_generate_entity.trace_manager) diff --git a/api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py b/api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py index 42b42a7f273..dba0087d3e2 100644 --- a/api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py +++ b/api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py @@ -125,7 +125,9 @@ class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline[EasyUIAppGenerat if self._application_generate_entity.app_config.app_mode != AppMode.COMPLETION: # start generate conversation name thread self._conversation_name_generate_thread = self._message_cycle_manager.generate_conversation_name( - conversation_id=self._conversation_id, query=self._application_generate_entity.query + conversation_id=self._conversation_id, + query=self._application_generate_entity.query, + message_id=self._message_id, ) generator = self._wrapper_process_stream_response(trace_manager=self._application_generate_entity.trace_manager) diff --git a/api/core/app/task_pipeline/message_cycle_manager.py b/api/core/app/task_pipeline/message_cycle_manager.py index 91b67a03e47..5c8fc6d7531 100644 --- a/api/core/app/task_pipeline/message_cycle_manager.py +++ b/api/core/app/task_pipeline/message_cycle_manager.py @@ -75,7 +75,9 @@ class MessageCycleManager: return StreamEvent.MESSAGE - def generate_conversation_name(self, *, conversation_id: str, query: str) -> Thread | None: + def generate_conversation_name( + self, *, conversation_id: str, query: str, message_id: str | None = None + ) -> Thread | None: """ Generate conversation name. :param conversation_id: conversation id @@ -100,6 +102,7 @@ class MessageCycleManager: "flask_app": current_app._get_current_object(), # type: ignore "conversation_id": conversation_id, "query": query, + "message_id": message_id, }, ) thread.daemon = True @@ -110,7 +113,13 @@ class MessageCycleManager: return thread - def _generate_conversation_name_worker(self, flask_app: Flask, conversation_id: str, query: str): + def _generate_conversation_name_worker( + self, + flask_app: Flask, + conversation_id: str, + query: str, + message_id: str | None = None, + ): with flask_app.app_context(): with session_factory.create_session() as session: # get conversation and message @@ -135,7 +144,11 @@ class MessageCycleManager: else: try: name = LLMGenerator.generate_conversation_name( - app_model.tenant_id, query, conversation_id, conversation.app_id + app_model.tenant_id, + query, + conversation_id, + conversation.app_id, + message_id=message_id, ) redis_client.setex(cache_key, 3600, name) except Exception: diff --git a/api/core/llm_generator/llm_generator.py b/api/core/llm_generator/llm_generator.py index caead37e3f5..32e16830f4f 100644 --- a/api/core/llm_generator/llm_generator.py +++ b/api/core/llm_generator/llm_generator.py @@ -188,7 +188,12 @@ class LLMGenerator: @classmethod def generate_conversation_name( - cls, tenant_id: str, query, conversation_id: str | None = None, app_id: str | None = None + cls, + tenant_id: str, + query, + conversation_id: str | None = None, + app_id: str | None = None, + message_id: str | None = None, ): prompt = CONVERSATION_TITLE_PROMPT @@ -239,6 +244,7 @@ class LLMGenerator: TraceTask( TraceTaskName.GENERATE_NAME_TRACE, conversation_id=conversation_id, + message_id=message_id, generate_conversation_name=name, inputs=prompt, timer=timer, diff --git a/api/core/ops/entities/config_entity.py b/api/core/ops/entities/config_entity.py index d78ce90aa1a..b3df7993eed 100644 --- a/api/core/ops/entities/config_entity.py +++ b/api/core/ops/entities/config_entity.py @@ -54,3 +54,11 @@ class BaseTracingConfig(BaseModel): OPS_FILE_PATH = "ops_trace/" OPS_TRACE_FAILED_KEY = "FAILED_OPS_TRACE" + + +def ops_trace_payload_path(app_id: str, file_id: str) -> str: + return f"{OPS_FILE_PATH}{app_id}/{file_id}.json" + + +def workflow_final_trace_file_id(workflow_run_id: str) -> str: + return f"workflow-final-{workflow_run_id}" diff --git a/api/core/ops/entities/trace_entity.py b/api/core/ops/entities/trace_entity.py index e183ad7f340..bfbc5a8cc2e 100644 --- a/api/core/ops/entities/trace_entity.py +++ b/api/core/ops/entities/trace_entity.py @@ -9,6 +9,7 @@ from core.helper.trace_id_helper import ParentTraceContext class BaseTraceInfo(BaseModel): + operation_id: str | None = None message_id: str | None = None message_data: Any | None = None inputs: Union[str, dict[str, Any], list[Any]] | None = None diff --git a/api/core/ops/exceptions.py b/api/core/ops/exceptions.py index 45517046871..5e695dcdbd8 100644 --- a/api/core/ops/exceptions.py +++ b/api/core/ops/exceptions.py @@ -9,6 +9,14 @@ class RetryableTraceDispatchError(RuntimeError): """Base class for transient trace dispatch failures that Celery may retry.""" +class TraceParentContextAccessError(RetryableTraceDispatchError): + """Raised when unified parent context storage is temporarily unavailable.""" + + +class InvalidTraceParentContextError(RuntimeError): + """Raised when stored unified parent context cannot be safely restored.""" + + class PendingTraceParentContextError(RetryableTraceDispatchError): """Raised when a nested trace arrives before its parent span context is available.""" diff --git a/api/core/ops/ops_trace_manager.py b/api/core/ops/ops_trace_manager.py index 0066d035e54..8227b42221b 100644 --- a/api/core/ops/ops_trace_manager.py +++ b/api/core/ops/ops_trace_manager.py @@ -16,14 +16,16 @@ from pydantic import TypeAdapter from sqlalchemy import select from sqlalchemy.orm import Session, sessionmaker +from configs import dify_config from core.helper.encrypter import batch_decrypt_token, encrypt_token, obfuscated_token from core.helper.trace_id_helper import ParentTraceContext from core.ops.entities.config_entity import ( - OPS_FILE_PATH, BaseTracingConfig, TracingProviderEnum, + ops_trace_payload_path, ) from core.ops.entities.trace_entity import ( + BaseTraceInfo, DatasetRetrievalTraceInfo, DraftNodeExecutionTrace, GenerateNameTraceInfo, @@ -37,6 +39,7 @@ from core.ops.entities.trace_entity import ( WorkflowNodeTraceInfo, WorkflowTraceInfo, ) +from core.ops.unified_trace.registry import UnifiedProviderConfigEntry, unified_provider_config_map from core.ops.utils import JSON_DICT_ADAPTER, get_message_data from extensions.ext_database import db from extensions.ext_storage import storage @@ -344,6 +347,22 @@ class OpsTraceManager: decrypted_configs_cache: LRUCache = LRUCache(maxsize=128) _decryption_cache_lock = threading.RLock() + @classmethod + def _get_dispatch_entry( + cls, tracing_provider: str + ) -> tuple[str, TracingProviderConfigEntry | UnifiedProviderConfigEntry]: + """Select one tracing mode before constructing an instance. + + Registered unified providers never fall back after construction or dispatch starts. + Unregistered providers continue through the legacy registry. + """ + if dify_config.OPS_TRACE_UNIFIED_ENABLED: + try: + return "unified", unified_provider_config_map[tracing_provider] + except KeyError: + pass + return "legacy", provider_config_map[tracing_provider] + @classmethod def encrypt_tracing_config( cls, tenant_id: str, tracing_provider: str, tracing_config: dict[str, Any], current_trace_config=None @@ -526,17 +545,16 @@ class OpsTraceManager: if not decrypt_trace_config: return None - trace_instance, config_class = ( - provider_config_map[tracing_provider]["trace_instance"], - provider_config_map[tracing_provider]["config_class"], - ) - decrypt_trace_config_key = json.dumps(decrypt_trace_config, sort_keys=True) - tracing_instance = cls.ops_trace_instances_cache.get(decrypt_trace_config_key) + mode, dispatch_entry = cls._get_dispatch_entry(tracing_provider) + trace_instance = dispatch_entry["trace_instance"] + config_class = dispatch_entry["config_class"] + cache_key = (mode, tracing_provider, json.dumps(decrypt_trace_config, sort_keys=True)) + tracing_instance = cls.ops_trace_instances_cache.get(cache_key) if tracing_instance is None: - # create new tracing_instance and update the cache if it absent + # The mode is part of the key so unified and legacy clients never share mutable SDK state. tracing_instance = trace_instance(config_class(**decrypt_trace_config)) - cls.ops_trace_instances_cache[decrypt_trace_config_key] = tracing_instance - logger.info("new tracing_instance for app_id: %s", app_id) + cls.ops_trace_instances_cache[cache_key] = tracing_instance + logger.info("new %s tracing_instance for app_id: %s", mode, app_id) return tracing_instance @classmethod @@ -739,6 +757,8 @@ class TraceTask: trace_type: Any, message_id: str | None = None, workflow_execution: "WorkflowExecution | None" = None, + workflow_run_id: str | None = None, + workflow_total_tokens: int | None = None, conversation_id: str | None = None, user_id: str | None = None, timer: Any | None = None, @@ -746,8 +766,8 @@ class TraceTask: ): self.trace_type = trace_type self.message_id = message_id - self.workflow_run_id = workflow_execution.id_ if workflow_execution else None - self.workflow_total_tokens: int | None = workflow_execution.total_tokens if workflow_execution else None + self.workflow_run_id = workflow_execution.id_ if workflow_execution else workflow_run_id + self.workflow_total_tokens = workflow_execution.total_tokens if workflow_execution else workflow_total_tokens self.conversation_id = conversation_id self.user_id = user_id self.timer = timer @@ -1275,6 +1295,7 @@ class TraceTask: generate_name_trace_info = GenerateNameTraceInfo( trace_id=self.trace_id, + message_id=self.message_id, conversation_id=conversation_id, inputs=inputs, outputs=generate_conversation_name, @@ -1539,30 +1560,57 @@ class TraceQueueManager: trace_manager_timer.daemon = False trace_manager_timer.start() + def _resolve_storage_id(self, task: TraceTask) -> str | None: + storage_id = task.app_id + if storage_id is not None: + return storage_id + + tenant_id = task.kwargs.get("tenant_id") + if tenant_id: + return f"tenant-{tenant_id}" + + logger.warning("Skipping trace without app_id or tenant_id, trace_type: %s", task.trace_type) + return None + + def persist_trace_task(self, task: TraceTask, *, file_id: str | None = None) -> dict[str, str] | None: + if not (self._enterprise_telemetry_enabled or self.trace_instance): + return None + + task.app_id = self.app_id + storage_id = self._resolve_storage_id(task) + if storage_id is None: + return None + + resolved_file_id = file_id or uuid4().hex + trace_info = task.execute() + if isinstance(trace_info, BaseTraceInfo) and trace_info.operation_id is None: + trace_info = trace_info.model_copy(update={"operation_id": str(uuid4())}) + task_data = TaskData( + app_id=storage_id, + trace_info_type=type(trace_info).__name__, + trace_info=trace_info.model_dump() if trace_info else None, + ) + storage.save( + ops_trace_payload_path(storage_id, resolved_file_id), + task_data.model_dump_json().encode("utf-8"), + ) + return {"file_id": resolved_file_id, "app_id": storage_id} + + def enqueue_persisted_trace(self, file_info: dict[str, str]) -> None: + process_trace_tasks.apply_async( + args=(file_info,), + retry=True, + retry_policy={ + "max_retries": 3, + "interval_start": 0, + "interval_step": 1, + "interval_max": 2, + }, + ) + def send_to_celery(self, tasks: list[TraceTask]): with self.flask_app.app_context(): for task in tasks: - storage_id = task.app_id - if storage_id is None: - tenant_id = task.kwargs.get("tenant_id") - if tenant_id: - storage_id = f"tenant-{tenant_id}" - else: - logger.warning("Skipping trace without app_id or tenant_id, trace_type: %s", task.trace_type) - continue - - file_id = uuid4().hex - trace_info = task.execute() - - task_data = TaskData( - app_id=storage_id, - trace_info_type=type(trace_info).__name__, - trace_info=trace_info.model_dump() if trace_info else None, - ) - file_path = f"{OPS_FILE_PATH}{storage_id}/{file_id}.json" - storage.save(file_path, task_data.model_dump_json().encode("utf-8")) - file_info = { - "file_id": file_id, - "app_id": storage_id, - } - process_trace_tasks.delay(file_info) # type: ignore + file_info = self.persist_trace_task(task) + if file_info is not None: + self.enqueue_persisted_trace(file_info) diff --git a/api/core/ops/unified_trace/__init__.py b/api/core/ops/unified_trace/__init__.py new file mode 100644 index 00000000000..bf131bef68c --- /dev/null +++ b/api/core/ops/unified_trace/__init__.py @@ -0,0 +1 @@ +"""Provider-neutral trace construction and dispatch for opt-in unified tracing.""" diff --git a/api/core/ops/unified_trace/entities.py b/api/core/ops/unified_trace/entities.py new file mode 100644 index 00000000000..e8d2a52f089 --- /dev/null +++ b/api/core/ops/unified_trace/entities.py @@ -0,0 +1,79 @@ +"""Provider-independent trace entities emitted by the unified trace builder.""" + +from datetime import datetime +from enum import StrEnum +from typing import Any, Self + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from core.helper.trace_id_helper import ParentTraceContext + + +class CanonicalSpanKind(StrEnum): + CHAIN = "chain" + LLM = "llm" + RETRIEVER = "retriever" + TOOL = "tool" + AGENT = "agent" + + +class CanonicalSpanStatus(StrEnum): + OK = "ok" + ERROR = "error" + + +class CanonicalSpan(BaseModel): + """One provider-neutral operation with an explicit parent.""" + + id: str + parent_id: str | None + name: str + kind: CanonicalSpanKind + start_time: datetime + end_time: datetime | None + inputs: Any = None + outputs: Any = None + status: CanonicalSpanStatus + error: str | None = None + metadata: dict[str, Any] = Field(default_factory=dict) + synthetic: bool = False + can_parent_workflow: bool = False + publishes_parent_context: bool = False + links: tuple[str, ...] = () + + model_config = ConfigDict(extra="forbid", frozen=True) + + +class CanonicalTrace(BaseModel): + """A parent-before-child span tree ready for a provider adapter.""" + + trace_id: str + session_id: str + root_span_id: str + spans: tuple[CanonicalSpan, ...] + external_parent: ParentTraceContext | None = None + required_parent_context_id: str | None = None + + model_config = ConfigDict(extra="forbid", frozen=True) + + @model_validator(mode="after") + def validate_fragment(self) -> Self: + if self.external_parent is not None and self.required_parent_context_id is not None: + raise ValueError("canonical trace cannot require two external parent modes") + + seen: set[str] = set() + root_seen = False + for span in self.spans: + if span.id in seen: + raise ValueError(f"duplicate canonical span id: {span.id}") + if span.id == self.root_span_id: + root_seen = True + if span.parent_id is not None: + raise ValueError("canonical trace root cannot have a local parent") + elif span.parent_id is None or span.parent_id not in seen: + raise ValueError(f"canonical span parent must appear first: {span.id}") + seen.add(span.id) + + if not root_seen: + raise ValueError("canonical trace root is missing") + return self diff --git a/api/core/ops/unified_trace/hierarchy.py b/api/core/ops/unified_trace/hierarchy.py new file mode 100644 index 00000000000..0f08ceda597 --- /dev/null +++ b/api/core/ops/unified_trace/hierarchy.py @@ -0,0 +1,194 @@ +"""Deterministically reconstruct provider-neutral workflow execution hierarchy.""" + +import re +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime, timedelta +from typing import Any, Literal + +from graphon.entities import WorkflowNodeExecution + +_WRAPPER_INDEX_PATTERN = re.compile(r"^[A-Za-z0-9_.:-]+$") +_WRAPPER_FIELDS: tuple[tuple[Literal["iteration", "loop"], str, str], ...] = ( + ("iteration", "iteration_id", "iteration_index"), + ("loop", "loop_id", "loop_index"), +) + + +WorkflowExecutionLike = WorkflowNodeExecution + + +def _read_attribute(value: object, name: str, default: Any = None) -> Any: + """Read fields shared by persisted executions and legacy trace objects.""" + return getattr(value, name, default) # guard-ignore: no-new-getattr -- supports legacy trace objects + + +@dataclass(frozen=True) +class WrapperKey: + kind: Literal["iteration", "loop"] + container_execution_id: str + index: str + + +@dataclass(frozen=True) +class WrapperSpec: + id: str + key: WrapperKey + parent_execution_id: str + child_execution_ids: frozenset[str] + start_time: datetime + end_time: datetime + has_error: bool + + +@dataclass(frozen=True) +class WorkflowHierarchy: + parent_by_execution_id: Mapping[str, str] + wrapper_by_child_execution_id: Mapping[str, WrapperSpec] + wrappers: tuple[WrapperSpec, ...] + + +def execution_id(execution: WorkflowExecutionLike) -> str: + """Return the persisted execution identifier used as a canonical span ID.""" + return str(_read_attribute(execution, "id") or execution.node_execution_id) + + +def execution_metadata(execution: object) -> Mapping[str, Any]: + """Read metadata from repository models and older trace test doubles.""" + value = _read_attribute(execution, "execution_metadata_dict") + if not isinstance(value, Mapping): + value = _read_attribute(execution, "metadata") + if not isinstance(value, Mapping): + return {} + return {str(key): item for key, item in value.items()} + + +def _unique_execution_by_node_id(executions: Sequence[WorkflowExecutionLike]) -> dict[str, str]: + result: dict[str, str] = {} + ambiguous: set[str] = set() + for item in executions: + node_id = item.node_id + if not isinstance(node_id, str) or node_id in ambiguous: + continue + item_execution_id = execution_id(item) + previous = result.get(node_id) + if previous is None: + result[node_id] = item_execution_id + elif previous != item_execution_id: + result.pop(node_id, None) + ambiguous.add(node_id) + return result + + +def _metadata_or_attr(execution: object, metadata: Mapping[str, Any], key: str) -> Any: + value = metadata.get(key) + return value if value is not None else _read_attribute(execution, key) + + +def _normalize_index(value: Any) -> str | None: + if isinstance(value, bool): + return None + if isinstance(value, int): + return str(value) if value >= 0 else None + if isinstance(value, str) and _WRAPPER_INDEX_PATTERN.fullmatch(value): + return value + return None + + +def _finished_at(execution: WorkflowExecutionLike) -> datetime: + explicit_end = _read_attribute(execution, "end_time") + if isinstance(explicit_end, datetime): + return explicit_end + started_at = execution.created_at or _read_attribute(execution, "start_time") or datetime.now() + return started_at + timedelta(seconds=execution.elapsed_time or 0) + + +def _failed(execution: WorkflowExecutionLike) -> bool: + status = _read_attribute(execution.status, "value", execution.status) + return status == "failed" + + +def _remove_cycles(parent_by_execution_id: dict[str, str]) -> None: + cyclic: set[str] = set() + for start in sorted(parent_by_execution_id): + path: list[str] = [] + position: dict[str, int] = {} + current = start + while current in parent_by_execution_id: + if current in position: + cyclic.update(path[position[current] :]) + break + position[current] = len(path) + path.append(current) + current = parent_by_execution_id[current] + for execution in cyclic: + parent_by_execution_id.pop(execution, None) + + +def build_workflow_hierarchy(executions: Sequence[WorkflowExecutionLike]) -> WorkflowHierarchy: + """Build stable parents and synthetic loop/iteration wrappers. + + Ambiguous repeated graph node IDs and cyclic predecessor data are dropped + rather than guessed. Their spans consequently fall back to the workflow root. + """ + execution_by_node_id = _unique_execution_by_node_id(executions) + parent_by_execution_id: dict[str, str] = {} + + for item in executions: + item_execution_id = execution_id(item) + predecessor_node_id = item.predecessor_node_id + parent_execution_id = ( + execution_by_node_id.get(predecessor_node_id) if isinstance(predecessor_node_id, str) else None + ) + if parent_execution_id is None: + metadata = execution_metadata(item) + for structured_key in ("iteration_id", "loop_id"): + container_node_id = _metadata_or_attr(item, metadata, structured_key) + if isinstance(container_node_id, str): + parent_execution_id = execution_by_node_id.get(container_node_id) + if parent_execution_id is not None: + break + if parent_execution_id and parent_execution_id != item_execution_id: + parent_by_execution_id[item_execution_id] = parent_execution_id + + _remove_cycles(parent_by_execution_id) + + grouped: dict[WrapperKey, list[WorkflowExecutionLike]] = {} + for item in executions: + metadata = execution_metadata(item) + for kind, container_key, index_key in _WRAPPER_FIELDS: + container_node_id = _metadata_or_attr(item, metadata, container_key) + index = _normalize_index(_metadata_or_attr(item, metadata, index_key)) + if not isinstance(container_node_id, str) or index is None: + continue + container_execution_id = execution_by_node_id.get(container_node_id) + if container_execution_id is None or container_execution_id == execution_id(item): + continue + wrapper_key = WrapperKey(kind=kind, container_execution_id=container_execution_id, index=index) + grouped.setdefault(wrapper_key, []).append(item) + break + + wrappers: list[WrapperSpec] = [] + wrapper_by_child_execution_id: dict[str, WrapperSpec] = {} + for wrapper_key in sorted(grouped, key=lambda item: (item.kind, item.container_execution_id, item.index)): + children = grouped[wrapper_key] + child_ids = frozenset(execution_id(item) for item in children) + wrapper = WrapperSpec( + id=f"{wrapper_key.kind}:{wrapper_key.container_execution_id}:{wrapper_key.index}", + key=wrapper_key, + parent_execution_id=wrapper_key.container_execution_id, + child_execution_ids=child_ids, + start_time=min(item.created_at or datetime.now() for item in children), + end_time=max(_finished_at(item) for item in children), + has_error=any(_failed(item) for item in children), + ) + wrappers.append(wrapper) + for child_id in child_ids: + wrapper_by_child_execution_id[child_id] = wrapper + parent_by_execution_id[child_id] = wrapper.id + + return WorkflowHierarchy( + parent_by_execution_id=parent_by_execution_id, + wrapper_by_child_execution_id=wrapper_by_child_execution_id, + wrappers=tuple(wrappers), + ) diff --git a/api/core/ops/unified_trace/parent_context.py b/api/core/ops/unified_trace/parent_context.py new file mode 100644 index 00000000000..25645ffe7c3 --- /dev/null +++ b/api/core/ops/unified_trace/parent_context.py @@ -0,0 +1,249 @@ +"""Coordinate nested-workflow provider parent contexts through Redis. + +The coordinator owns storage, compatibility decisions, validation, and retry +signals. Provider adapters only create and consume their opaque context fields. +""" + +import hashlib +import json +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from enum import StrEnum +from typing import Literal, Protocol + +from pydantic import BaseModel, ConfigDict, ValidationError +from sqlalchemy import select +from sqlalchemy.orm import Session + +from configs import dify_config +from core.helper.trace_id_helper import ParentTraceContext +from core.ops.exceptions import ( + InvalidTraceParentContextError, + PendingTraceParentContextError, + TraceParentContextAccessError, +) +from extensions.ext_database import db +from models.model import App, TraceAppConfig +from models.workflow import WorkflowRun + +_PARENT_CONTEXT_KEY_PREFIX = "trace:unified:parent:" + + +class RedisParentContextStore(Protocol): + def setex(self, name: str, time: int, value: str) -> object: + raise NotImplementedError + + def get(self, name: str) -> bytes | str | None: + raise NotImplementedError + + +class ProviderParentContext(BaseModel): + """Versioned envelope containing the minimum provider restoration state.""" + + version: Literal[1] = 1 + provider: str + scope: str + trace_id: str + parent_id: str + provider_context: dict[str, str] + + model_config = ConfigDict(extra="forbid", frozen=True) + + +@dataclass(frozen=True) +class ParentDestination: + provider: str + scope: str + unified: bool + + +class ParentResolutionKind(StrEnum): + RESTORED = "restored" + LINKED_ROOT = "linked_root" + + +@dataclass(frozen=True) +class ParentResolution: + kind: ParentResolutionKind + context: ProviderParentContext | None = None + linked_parent: ParentTraceContext | None = None + + @classmethod + def restored(cls, context: ProviderParentContext) -> "ParentResolution": + return cls(kind=ParentResolutionKind.RESTORED, context=context) + + @classmethod + def linked_root(cls, parent: ParentTraceContext) -> "ParentResolution": + return cls(kind=ParentResolutionKind.LINKED_ROOT, linked_parent=parent) + + +ParentDestinationResolver = Callable[[str], ParentDestination | None] + + +def destination_scope(provider: str, endpoint: str, project: str) -> str: + """Return a stable non-secret fingerprint for a provider destination.""" + value = f"{provider}\0{endpoint.rstrip('/')}\0{project}" + return hashlib.sha256(value.encode()).hexdigest() + + +def parent_destination_from_config( + provider: str, + tracing_config: Mapping[str, object], + *, + unified: bool, +) -> ParentDestination: + """Build destination compatibility from non-secret persisted fields.""" + endpoint = tracing_config.get("endpoint") + project = tracing_config.get("project") + return ParentDestination( + provider=provider, + scope=destination_scope( + provider, + endpoint if isinstance(endpoint, str) else "", + project if isinstance(project, str) else "", + ), + unified=unified, + ) + + +def resolve_parent_destination(parent_workflow_run_id: str) -> ParentDestination | None: + """Resolve whether a parent workflow can publish compatible unified context.""" + with Session(db.engine) as session: + workflow_run = session.get(WorkflowRun, parent_workflow_run_id) + if workflow_run is None: + return None + app = session.get(App, workflow_run.app_id) + if app is None or not app.tracing: + return None + try: + app_tracing = json.loads(app.tracing) + except (TypeError, json.JSONDecodeError): + return None + if not isinstance(app_tracing, Mapping) or not app_tracing.get("enabled"): + return None + provider = app_tracing.get("tracing_provider") + if not isinstance(provider, str): + return None + trace_config = session.scalar( + select(TraceAppConfig) + .where(TraceAppConfig.app_id == app.id, TraceAppConfig.tracing_provider == provider) + .limit(1) + ) + if trace_config is None or not isinstance(trace_config.tracing_config, Mapping): + return None + + unified = False + if dify_config.OPS_TRACE_UNIFIED_ENABLED: + from core.ops.unified_trace.registry import unified_provider_config_map + + try: + unified_provider_config_map[provider] + unified = True + except KeyError: + pass + return parent_destination_from_config(provider, trace_config.tracing_config, unified=unified) + + +class ParentContextCoordinator: + """Publish and resolve cross-task parent contexts for unified providers.""" + + def __init__( + self, + store: RedisParentContextStore, + resolve_parent_destination: ParentDestinationResolver, + ) -> None: + self._store = store + self._resolve_parent_destination = resolve_parent_destination + + @staticmethod + def _key(parent_node_execution_id: str) -> str: + return f"{_PARENT_CONTEXT_KEY_PREFIX}{parent_node_execution_id}" + + def publish(self, parent_node_execution_id: str, context: ProviderParentContext) -> None: + """Persist an accepted provider parent so a nested task can restore it.""" + try: + self._store.setex( + self._key(parent_node_execution_id), + dify_config.OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS, + context.model_dump_json(), + ) + except Exception as error: + raise TraceParentContextAccessError( + f"Could not publish unified parent context for node_execution_id={parent_node_execution_id}" + ) from error + + def resolve( + self, + parent: ParentTraceContext, + *, + expected_provider: str, + expected_scope: str, + ) -> ParentResolution: + """Restore compatible context or explicitly select a linked new root.""" + destination = self._resolve_parent_destination(parent.parent_workflow_run_id) + if ( + destination is None + or not destination.unified + or destination.provider != expected_provider + or destination.scope != expected_scope + ): + return ParentResolution.linked_root(parent) + + parent_node_execution_id = parent.parent_node_execution_id + if not parent_node_execution_id: + raise InvalidTraceParentContextError("Nested workflow parent context has no node execution ID") + + return ParentResolution.restored( + self._restore( + parent_node_execution_id, + expected_provider=expected_provider, + expected_scope=expected_scope, + ) + ) + + def resolve_required( + self, + parent_context_id: str, + *, + expected_provider: str, + expected_scope: str, + ) -> ParentResolution: + """Restore a parent context that must exist for an asynchronous child.""" + return ParentResolution.restored( + self._restore( + parent_context_id, + expected_provider=expected_provider, + expected_scope=expected_scope, + ) + ) + + def _restore( + self, + parent_context_id: str, + *, + expected_provider: str, + expected_scope: str, + ) -> ProviderParentContext: + try: + raw_context = self._store.get(self._key(parent_context_id)) + except Exception as error: + raise TraceParentContextAccessError( + f"Could not read unified parent context for parent_context_id={parent_context_id}" + ) from error + + if raw_context is None: + raise PendingTraceParentContextError(parent_context_id) + + try: + context = ProviderParentContext.model_validate_json(raw_context) + except (ValidationError, ValueError, TypeError) as error: + raise InvalidTraceParentContextError( + f"Invalid unified parent context for parent_context_id={parent_context_id}" + ) from error + + if context.provider != expected_provider or context.scope != expected_scope: + raise InvalidTraceParentContextError( + "Stored unified parent context does not match the expected provider destination: " + f"parent_context_id={parent_context_id}" + ) + return context diff --git a/api/core/ops/unified_trace/provider.py b/api/core/ops/unified_trace/provider.py new file mode 100644 index 00000000000..a05d36cfc84 --- /dev/null +++ b/api/core/ops/unified_trace/provider.py @@ -0,0 +1,82 @@ +"""Unified trace runtime and provider adapter contract.""" + +from collections.abc import Callable +from typing import Protocol, override + +from core.ops.base_trace_instance import BaseTraceInstance +from core.ops.entities.config_entity import BaseTracingConfig +from core.ops.entities.trace_entity import BaseTraceInfo +from core.ops.unified_trace.entities import CanonicalTrace +from core.ops.unified_trace.parent_context import ( + ParentContextCoordinator, + ParentResolution, + ProviderParentContext, +) +from core.ops.unified_trace.trace_builder import CanonicalTraceBuilder + +ParentContextPublisher = Callable[[str, ProviderParentContext], None] + + +class UnifiedTraceAdapter(Protocol): + @property + def provider_name(self) -> str: + raise NotImplementedError + + @property + def scope(self) -> str: + raise NotImplementedError + + def emit( + self, + trace: CanonicalTrace, + parent: ParentResolution | None, + publish_parent_context: ParentContextPublisher, + ) -> None: + """Emit one Core-resolved, parent-first canonical fragment. + + Return only after the provider-specific synchronous acceptance step succeeds. + Raise RetryableTraceDispatchError when acceptance is unconfirmed because of a + recoverable provider or transport failure. Raise another exception for a + terminal failure. Publish parent context only after the corresponding provider + parent span has been accepted. + """ + raise NotImplementedError + + +class UnifiedTraceInstance(BaseTraceInstance): + """Build, coordinate, and emit a trace without a legacy fallback path.""" + + def __init__( + self, + trace_config: BaseTracingConfig, + *, + builder: CanonicalTraceBuilder, + adapter: UnifiedTraceAdapter, + coordinator: ParentContextCoordinator, + ) -> None: + super().__init__(trace_config) + self._builder = builder + self._adapter = adapter + self._coordinator = coordinator + + @override + def trace(self, trace_info: BaseTraceInfo) -> None: + canonical_trace = self._builder.build(trace_info) + if canonical_trace is None: + return + + parent_resolution = None + if canonical_trace.external_parent is not None: + parent_resolution = self._coordinator.resolve( + canonical_trace.external_parent, + expected_provider=self._adapter.provider_name, + expected_scope=self._adapter.scope, + ) + elif canonical_trace.required_parent_context_id is not None: + parent_resolution = self._coordinator.resolve_required( + canonical_trace.required_parent_context_id, + expected_provider=self._adapter.provider_name, + expected_scope=self._adapter.scope, + ) + + self._adapter.emit(canonical_trace, parent_resolution, self._coordinator.publish) diff --git a/api/core/ops/unified_trace/registry.py b/api/core/ops/unified_trace/registry.py new file mode 100644 index 00000000000..7619b28923f --- /dev/null +++ b/api/core/ops/unified_trace/registry.py @@ -0,0 +1,35 @@ +"""Lazy registry for providers implemented by the unified tracing path.""" + +import collections +from typing import TypedDict, override + +from core.ops.base_trace_instance import BaseTraceInstance +from core.ops.entities.config_entity import BaseTracingConfig, TracingProviderEnum + + +class UnifiedProviderConfigEntry(TypedDict): + config_class: type[BaseTracingConfig] + trace_instance: type[BaseTraceInstance] + + +class UnifiedTraceProviderConfigMap(collections.UserDict[str, UnifiedProviderConfigEntry]): + """Resolve unified providers without importing their SDKs until selected.""" + + @override + def __getitem__(self, key: str) -> UnifiedProviderConfigEntry: + match key: + case TracingProviderEnum.PHOENIX: + from dify_trace_arize_phoenix.config import PhoenixConfig + from dify_trace_arize_phoenix.unified_trace import UnifiedPhoenixTrace + + return {"config_class": PhoenixConfig, "trace_instance": UnifiedPhoenixTrace} + case TracingProviderEnum.LANGSMITH: + from dify_trace_langsmith.config import LangSmithConfig + from dify_trace_langsmith.unified_trace import UnifiedLangSmithTrace + + return {"config_class": LangSmithConfig, "trace_instance": UnifiedLangSmithTrace} + case _: + raise KeyError(f"Unified tracing provider is not registered: {key}") + + +unified_provider_config_map = UnifiedTraceProviderConfigMap() diff --git a/api/core/ops/unified_trace/trace_builder.py b/api/core/ops/unified_trace/trace_builder.py new file mode 100644 index 00000000000..8390c89020a --- /dev/null +++ b/api/core/ops/unified_trace/trace_builder.py @@ -0,0 +1,452 @@ +"""Build provider-neutral traces from core ops trace entities. + +Workflow persistence is accessed through an injected loader. Provider adapters +therefore receive a complete parent-first tree and never query Dify models. +""" + +from collections.abc import Callable, Mapping, Sequence +from datetime import datetime, timedelta +from typing import Any +from uuid import uuid4 + +from pydantic import ValidationError +from sqlalchemy.orm import sessionmaker + +from core.app.workflow.retry_history import RETRY_HISTORY_PROCESS_DATA_KEY +from core.helper.trace_id_helper import ParentTraceContext +from core.ops.entities.trace_entity import ( + BaseTraceInfo, + DatasetRetrievalTraceInfo, + GenerateNameTraceInfo, + MessageTraceInfo, + ModerationTraceInfo, + SuggestedQuestionTraceInfo, + ToolTraceInfo, + WorkflowTraceInfo, +) +from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace +from core.ops.unified_trace.hierarchy import ( + WorkflowExecutionLike, + build_workflow_hierarchy, + execution_id, + execution_metadata, +) +from core.repositories import DifyCoreRepositoryFactory +from extensions.ext_database import db +from models import Account +from models.workflow import WorkflowNodeExecutionTriggeredFrom + +WorkflowExecutionLoader = Callable[[WorkflowTraceInfo], Sequence[WorkflowExecutionLike]] +ServiceAccountResolver = Callable[[str], Account] + + +class RepositoryWorkflowExecutionLoader: + """Load one workflow's executions through the tenant-scoped core repository.""" + + def __init__(self, get_service_account: ServiceAccountResolver) -> None: + self._get_service_account = get_service_account + + def __call__(self, trace_info: WorkflowTraceInfo) -> Sequence[WorkflowExecutionLike]: + app_id = trace_info.metadata.get("app_id") + if not isinstance(app_id, str) or not app_id: + raise ValueError("No app_id found in workflow trace metadata") + repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( + session_factory=sessionmaker(bind=db.engine), + tenant_id=trace_info.tenant_id, + user=self._get_service_account(app_id), + app_id=app_id, + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + return repository.get_by_workflow_execution(workflow_execution_id=trace_info.workflow_run_id) + + +_NODE_KIND: dict[str, CanonicalSpanKind] = { + "llm": CanonicalSpanKind.LLM, + "knowledge-retrieval": CanonicalSpanKind.RETRIEVER, + "tool": CanonicalSpanKind.TOOL, + "agent": CanonicalSpanKind.AGENT, +} +_RETRY_SUMMARY_FIELDS = ("retry_index", "error", "elapsed_time", "created_at", "finished_at") + + +def _read_attribute(value: object, name: str, default: Any = None) -> Any: + """Read fields shared by persisted trace models and legacy trace objects.""" + return getattr(value, name, default) # guard-ignore: no-new-getattr -- supports legacy trace objects + + +def _retry_metadata(process_data: Mapping[str, Any]) -> dict[str, Any]: + raw_history = process_data.get(RETRY_HISTORY_PROCESS_DATA_KEY) + if not isinstance(raw_history, list): + return {} + + attempts: list[dict[str, Any]] = [] + for raw_attempt in raw_history: + if not isinstance(raw_attempt, Mapping): + continue + retry_index = raw_attempt.get("retry_index") + if isinstance(retry_index, bool) or not isinstance(retry_index, int) or retry_index <= 0: + continue + attempts.append({field: raw_attempt.get(field) for field in _RETRY_SUMMARY_FIELDS}) + + return {"retry_count": len(attempts), "retry_attempts": attempts} if attempts else {} + + +def resolve_session_id(trace_info: WorkflowTraceInfo | MessageTraceInfo) -> str: + """Resolve an explicit trace session before stable Dify fallbacks.""" + custom_session_id = trace_info.metadata.get("trace_session_id") + if isinstance(custom_session_id, str) and custom_session_id: + return custom_session_id + + if isinstance(trace_info, WorkflowTraceInfo): + if trace_info.conversation_id: + return trace_info.conversation_id + parent_workflow_run_id, _ = trace_info.resolved_parent_context + return parent_workflow_run_id or trace_info.workflow_run_id + + if trace_info.message_data is None: + return "" + conversation_id = _read_attribute(trace_info.message_data, "conversation_id") + return conversation_id if isinstance(conversation_id, str) else "" + + +def _status(error: str | None, status: Any = None) -> CanonicalSpanStatus: + status_value = _read_attribute(status, "value", status) + return CanonicalSpanStatus.ERROR if error or status_value in {"failed", "exception"} else CanonicalSpanStatus.OK + + +def _external_parent(trace_info: WorkflowTraceInfo) -> ParentTraceContext | None: + value = trace_info.metadata.get("parent_trace_context") + if isinstance(value, ParentTraceContext): + return value + if isinstance(value, Mapping): + try: + return ParentTraceContext.model_validate(value) + except ValidationError: + return None + return None + + +def _started_at(value: datetime | None) -> datetime: + return value or datetime.now() + + +def _single_session_id(trace_info: BaseTraceInfo) -> str: + value = trace_info.metadata.get("trace_session_id") + return value if isinstance(value, str) else "" + + +class CanonicalTraceBuilder: + """Convert supported ops trace entities into canonical parent-first trees.""" + + def __init__(self, load_workflow_executions: WorkflowExecutionLoader) -> None: + self._load_workflow_executions = load_workflow_executions + + def build(self, trace_info: BaseTraceInfo) -> CanonicalTrace | None: + match trace_info: + case WorkflowTraceInfo(): + return self._build_workflow(trace_info) + case MessageTraceInfo(): + return self._build_message(trace_info) + case ModerationTraceInfo(): + return self._build_moderation(trace_info) + case SuggestedQuestionTraceInfo(): + return self._build_suggested_question(trace_info) + case DatasetRetrievalTraceInfo(): + return self._build_dataset_retrieval(trace_info) + case ToolTraceInfo(): + return self._build_tool(trace_info) + case GenerateNameTraceInfo(): + return self._build_generate_name(trace_info) + case _: + return None + + def _build_workflow(self, trace_info: WorkflowTraceInfo) -> CanonicalTrace: + executions = self._load_workflow_executions(trace_info) + hierarchy = build_workflow_hierarchy(executions) + workflow_data = trace_info.workflow_data + workflow_start = _started_at(_read_attribute(workflow_data, "created_at") or trace_info.start_time) + workflow_end = _read_attribute(workflow_data, "finished_at") or trace_info.end_time + message_span_id = trace_info.message_id + workflow_span_id = trace_info.workflow_run_id + spans: dict[str, CanonicalSpan] = {} + + if message_span_id: + spans[message_span_id] = CanonicalSpan( + id=message_span_id, + parent_id=None, + name=f"chatflow_{trace_info.workflow_run_id}", + kind=CanonicalSpanKind.CHAIN, + start_time=_started_at(trace_info.start_time or workflow_start), + end_time=trace_info.end_time or workflow_end, + inputs=trace_info.query or dict(trace_info.workflow_run_inputs), + outputs=dict(trace_info.workflow_run_outputs), + status=_status(trace_info.error), + error=trace_info.error or None, + metadata={**trace_info.metadata, "trace_entity_type": "message"}, + publishes_parent_context=True, + ) + root_span_id = message_span_id + workflow_parent_id: str | None = message_span_id + else: + root_span_id = workflow_span_id + workflow_parent_id = None + + spans[workflow_span_id] = CanonicalSpan( + id=workflow_span_id, + parent_id=workflow_parent_id, + name=f"workflow_{workflow_span_id}", + kind=CanonicalSpanKind.CHAIN, + start_time=workflow_start, + end_time=workflow_end, + inputs=dict(trace_info.workflow_run_inputs), + outputs=dict(trace_info.workflow_run_outputs), + status=_status(trace_info.error, trace_info.workflow_run_status), + error=trace_info.error or None, + metadata={ + **trace_info.metadata, + "workflow_id": trace_info.workflow_id, + "workflow_run_id": workflow_span_id, + "workflow_app_log_id": trace_info.workflow_app_log_id, + "total_tokens": trace_info.total_tokens, + }, + ) + + execution_by_id = {execution_id(item): item for item in executions} + for wrapper in hierarchy.wrappers: + spans[wrapper.id] = CanonicalSpan( + id=wrapper.id, + parent_id=wrapper.parent_execution_id, + name=f"{wrapper.key.kind}[{wrapper.key.index}]", + kind=CanonicalSpanKind.CHAIN, + start_time=wrapper.start_time, + end_time=wrapper.end_time, + status=CanonicalSpanStatus.ERROR if wrapper.has_error else CanonicalSpanStatus.OK, + error="wrapper child failed" if wrapper.has_error else None, + metadata={ + "wrapper_type": wrapper.key.kind, + "wrapper_index": wrapper.key.index, + "container_execution_id": wrapper.parent_execution_id, + }, + synthetic=True, + ) + + for item_execution_id, item in execution_by_id.items(): + process_data = _read_attribute(item, "process_data") or {} + outputs = _read_attribute(item, "outputs") or {} + node_type = str(_read_attribute(item, "node_type", "")) + error = _read_attribute(item, "error") + started_at = _started_at(_read_attribute(item, "created_at")) + elapsed_time = _read_attribute(item, "elapsed_time") or 0 + metadata = dict(execution_metadata(item)) + metadata.update( + { + "node_id": _read_attribute(item, "node_id", ""), + "node_execution_id": item_execution_id, + "node_type": node_type, + "status": _read_attribute(item, "status", ""), + "model_provider": process_data.get("model_provider"), + "model_name": process_data.get("model_name"), + } + ) + usage = process_data.get("usage") or (outputs.get("usage") if isinstance(outputs, Mapping) else None) or {} + if isinstance(usage, Mapping): + metadata.update( + { + "prompt_tokens": usage.get("prompt_tokens", 0), + "completion_tokens": usage.get("completion_tokens", 0), + "total_tokens": usage.get("total_tokens", 0), + } + ) + metadata.update(_retry_metadata(process_data)) + title = _read_attribute(item, "title") + name = f"{node_type}_{title}" if isinstance(title, str) and title else node_type + spans[item_execution_id] = CanonicalSpan( + id=item_execution_id, + parent_id=hierarchy.parent_by_execution_id.get(item_execution_id, workflow_span_id), + name=name, + kind=_NODE_KIND.get(node_type, CanonicalSpanKind.CHAIN), + start_time=started_at, + end_time=started_at + timedelta(seconds=elapsed_time), + inputs=process_data.get("prompts", []) if node_type == "llm" else _read_attribute(item, "inputs") or {}, + outputs=outputs, + status=_status(error, _read_attribute(item, "status")), + error=error, + metadata=metadata, + can_parent_workflow=node_type == "tool", + ) + + ordered: list[CanonicalSpan] = [] + emitted: set[str] = set() + + def emit(span_id: str) -> None: + if span_id in emitted: + return + span = spans[span_id] + if span.parent_id in spans: + emit(span.parent_id) + emitted.add(span_id) + ordered.append(span) + + for span_id in sorted(spans): + emit(span_id) + + return CanonicalTrace( + trace_id=trace_info.resolved_trace_id or root_span_id, + session_id=resolve_session_id(trace_info), + root_span_id=root_span_id, + spans=tuple(ordered), + external_parent=_external_parent(trace_info), + ) + + def _single_trace( + self, + trace_info: BaseTraceInfo, + *, + name: str, + kind: CanonicalSpanKind, + inputs: Any, + outputs: Any, + error: str | None = None, + start_time: datetime | None = None, + end_time: datetime | None = None, + parent_context_id: str | None = None, + span_id: str | None = None, + session_id: str | None = None, + ) -> CanonicalTrace: + operation_id = span_id or trace_info.operation_id or str(uuid4()) + trace_id = trace_info.resolved_trace_id or parent_context_id or operation_id + span = CanonicalSpan( + id=operation_id, + parent_id=None, + name=name, + kind=kind, + start_time=_started_at(start_time or trace_info.start_time), + end_time=end_time or trace_info.end_time, + inputs=inputs, + outputs=outputs, + status=_status(error), + error=error, + metadata=dict(trace_info.metadata), + ) + return CanonicalTrace( + trace_id=trace_id, + session_id=session_id if session_id is not None else _single_session_id(trace_info), + root_span_id=operation_id, + spans=(span,), + required_parent_context_id=parent_context_id, + ) + + def _build_message(self, trace_info: MessageTraceInfo) -> CanonicalTrace | None: + message = trace_info.message_data + if message is None: + return None + message_id = trace_info.message_id or str(_read_attribute(message, "id", "")) or str(uuid4()) + started_at = _started_at(trace_info.start_time or _read_attribute(message, "created_at")) + ended_at = trace_info.end_time or _read_attribute(message, "updated_at") + answer = _read_attribute(message, "answer", trace_info.outputs) + message_error = trace_info.error or _read_attribute(message, "error") + metadata = { + **trace_info.metadata, + "trace_entity_type": "message", + "model_provider": _read_attribute(message, "model_provider"), + "model_name": _read_attribute(message, "model_id"), + "prompt_tokens": trace_info.message_tokens, + "completion_tokens": trace_info.answer_tokens, + "total_tokens": trace_info.total_tokens, + } + spans: list[CanonicalSpan] = [ + CanonicalSpan( + id=message_id, + parent_id=None, + name="message", + kind=CanonicalSpanKind.CHAIN, + start_time=started_at, + end_time=ended_at, + inputs=trace_info.inputs, + outputs=answer, + status=_status(message_error), + error=message_error, + metadata=metadata, + publishes_parent_context=True, + ), + CanonicalSpan( + id=f"{message_id}:llm", + parent_id=message_id, + name="llm", + kind=CanonicalSpanKind.LLM, + start_time=started_at, + end_time=ended_at, + inputs=trace_info.inputs, + outputs=trace_info.outputs if trace_info.outputs is not None else answer, + status=_status(message_error), + error=message_error, + metadata=metadata, + synthetic=True, + ), + ] + return CanonicalTrace( + trace_id=trace_info.resolved_trace_id or message_id, + session_id=resolve_session_id(trace_info), + root_span_id=message_id, + spans=tuple(spans), + ) + + def _build_moderation(self, trace_info: ModerationTraceInfo) -> CanonicalTrace | None: + if trace_info.message_data is None: + return None + return self._single_trace( + trace_info, + name="moderation", + kind=CanonicalSpanKind.TOOL, + inputs=trace_info.inputs, + outputs={"action": trace_info.action, "flagged": trace_info.flagged}, + parent_context_id=trace_info.message_id, + ) + + def _build_suggested_question(self, trace_info: SuggestedQuestionTraceInfo) -> CanonicalTrace | None: + if trace_info.message_data is None: + return None + return self._single_trace( + trace_info, + name="suggested_question", + kind=CanonicalSpanKind.TOOL, + inputs=trace_info.inputs, + outputs=trace_info.suggested_question, + error=trace_info.error, + parent_context_id=trace_info.message_id, + ) + + def _build_dataset_retrieval(self, trace_info: DatasetRetrievalTraceInfo) -> CanonicalTrace | None: + if trace_info.message_data is None: + return None + return self._single_trace( + trace_info, + name="dataset_retrieval", + kind=CanonicalSpanKind.RETRIEVER, + inputs=trace_info.inputs, + outputs={"documents": trace_info.documents}, + error=trace_info.error, + parent_context_id=trace_info.message_id, + ) + + def _build_tool(self, trace_info: ToolTraceInfo) -> CanonicalTrace: + return self._single_trace( + trace_info, + name=trace_info.tool_name, + kind=CanonicalSpanKind.TOOL, + inputs=trace_info.tool_inputs, + outputs=trace_info.tool_outputs, + error=trace_info.error, + parent_context_id=trace_info.message_id, + ) + + def _build_generate_name(self, trace_info: GenerateNameTraceInfo) -> CanonicalTrace: + return self._single_trace( + trace_info, + name="generate_name", + kind=CanonicalSpanKind.TOOL, + inputs=trace_info.inputs, + outputs=trace_info.outputs, + parent_context_id=trace_info.message_id, + session_id=_single_session_id(trace_info) or trace_info.conversation_id or "", + ) diff --git a/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/unified_trace.py b/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/unified_trace.py new file mode 100644 index 00000000000..f82505c5f94 --- /dev/null +++ b/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/unified_trace.py @@ -0,0 +1,195 @@ +"""Phoenix adapter for the provider-neutral unified tracing runtime.""" + +import json +from datetime import datetime +from typing import cast +from urllib.parse import urlparse + +from openinference.semconv.trace import OpenInferenceMimeTypeValues, OpenInferenceSpanKindValues, SpanAttributes +from opentelemetry.context import _SUPPRESS_INSTRUMENTATION_KEY, Context, attach, detach, set_value +from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter +from opentelemetry.sdk import trace as trace_sdk +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace.export import SpanExportResult +from opentelemetry.trace import Span, Status, StatusCode, get_current_span, set_span_in_context +from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator +from opentelemetry.util.types import AttributeValue + +from core.ops.exceptions import InvalidTraceParentContextError, RetryableTraceDispatchError +from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace +from core.ops.unified_trace.parent_context import ( + ParentContextCoordinator, + ParentResolution, + ParentResolutionKind, + ProviderParentContext, + destination_scope, + resolve_parent_destination, +) +from core.ops.unified_trace.provider import ParentContextPublisher, UnifiedTraceInstance +from core.ops.unified_trace.trace_builder import CanonicalTraceBuilder, RepositoryWorkflowExecutionLoader +from dify_trace_arize_phoenix.config import PhoenixConfig +from extensions.ext_redis import redis_client + +_KIND_MAP: dict[CanonicalSpanKind, OpenInferenceSpanKindValues] = { + CanonicalSpanKind.CHAIN: OpenInferenceSpanKindValues.CHAIN, + CanonicalSpanKind.LLM: OpenInferenceSpanKindValues.LLM, + CanonicalSpanKind.RETRIEVER: OpenInferenceSpanKindValues.RETRIEVER, + CanonicalSpanKind.TOOL: OpenInferenceSpanKindValues.TOOL, + CanonicalSpanKind.AGENT: OpenInferenceSpanKindValues.AGENT, +} + + +def _nanos(value: datetime | None) -> int | None: + return int(value.timestamp() * 1_000_000_000) if value is not None else None + + +def _json(value: object) -> str: + return json.dumps(value, default=str, ensure_ascii=False) + + +def setup_unified_tracer(config: PhoenixConfig) -> tuple[trace_sdk.Tracer, OTLPSpanExporter]: + """Create an isolated Phoenix tracer without touching the legacy provider.""" + parsed = urlparse(config.endpoint) + endpoint = f"{parsed.scheme}://{parsed.netloc}{parsed.path.rstrip('/')}/v1/traces" + exporter = OTLPSpanExporter( + endpoint=endpoint, + headers={ + "api_key": config.api_key or "", + "authorization": f"Bearer {config.api_key or ''}", + }, + timeout=30, + ) + provider = trace_sdk.TracerProvider( + resource=Resource( + attributes={ + "openinference.project.name": config.project or "", + "model_id": config.project or "", + } + ) + ) + return cast(trace_sdk.Tracer, provider.get_tracer(f"unified_phoenix_{config.project}")), exporter + + +class UnifiedPhoenixAdapter: + """Translate canonical spans to isolated OpenTelemetry/OpenInference spans.""" + + provider_name = "phoenix" + + def __init__(self, config: PhoenixConfig) -> None: + self._config = config + self._tracer, self._exporter = setup_unified_tracer(config) + self._propagator = TraceContextTextMapPropagator() + self._scope = destination_scope(self.provider_name, config.endpoint, config.project or "") + + @property + def scope(self) -> str: + return self._scope + + def _root_context(self, parent: ParentResolution | None) -> Context | None: + if parent is None or parent.kind is ParentResolutionKind.LINKED_ROOT: + return None + if parent.context is None: + return None + traceparent = parent.context.provider_context.get("traceparent") + if not traceparent: + raise InvalidTraceParentContextError("Phoenix parent context is missing traceparent") + context = self._propagator.extract(carrier={"traceparent": traceparent}) + span_context = get_current_span(context).get_span_context() + if not span_context.is_valid or not span_context.is_remote: + raise InvalidTraceParentContextError("Phoenix parent context contains an invalid traceparent") + return context + + def _attributes( + self, + canonical_span: CanonicalSpan, + trace: CanonicalTrace, + parent: ParentResolution | None, + ) -> dict[str, AttributeValue]: + metadata = dict(canonical_span.metadata) + if ( + canonical_span.id == trace.root_span_id + and parent is not None + and parent.kind is ParentResolutionKind.LINKED_ROOT + and parent.linked_parent is not None + ): + metadata["linked_parent_workflow_run_id"] = parent.linked_parent.parent_workflow_run_id + metadata["linked_parent_node_execution_id"] = parent.linked_parent.parent_node_execution_id + metadata["dify.span.kind"] = canonical_span.kind.value + metadata.pop("dify.span.links", None) + if canonical_span.links: + metadata["dify.span.links"] = list(canonical_span.links) + return { + SpanAttributes.OPENINFERENCE_SPAN_KIND: _KIND_MAP[canonical_span.kind].value, + SpanAttributes.INPUT_VALUE: _json(canonical_span.inputs), + SpanAttributes.INPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON.value, + SpanAttributes.OUTPUT_VALUE: _json(canonical_span.outputs), + SpanAttributes.OUTPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON.value, + SpanAttributes.METADATA: _json(metadata), + SpanAttributes.SESSION_ID: trace.session_id, + "dify.span.id": canonical_span.id, + "dify.span.synthetic": canonical_span.synthetic, + } + + def emit( + self, + trace: CanonicalTrace, + parent: ParentResolution | None, + publish_parent_context: ParentContextPublisher, + ) -> None: + span_by_id: dict[str, Span] = {} + root_context = self._root_context(parent) + + for canonical_span in trace.spans: + local_parent = span_by_id.get(canonical_span.parent_id or "") + context = set_span_in_context(local_parent) if local_parent is not None else root_context + span = self._tracer.start_span( + name=canonical_span.name, + context=context, + attributes=self._attributes(canonical_span, trace, parent), + start_time=_nanos(canonical_span.start_time), + ) + span_by_id[canonical_span.id] = span + provider_parent_context: ProviderParentContext | None = None + try: + if canonical_span.can_parent_workflow or canonical_span.publishes_parent_context: + carrier: dict[str, str] = {} + self._propagator.inject(carrier, context=set_span_in_context(span)) + provider_parent_context = ProviderParentContext( + provider=self.provider_name, + scope=self.scope, + trace_id=trace.trace_id, + parent_id=canonical_span.id, + provider_context=carrier, + ) + if canonical_span.status is CanonicalSpanStatus.ERROR: + error = canonical_span.error or "trace operation failed" + span.set_status(Status(StatusCode.ERROR, error)) + span.record_exception(RuntimeError(error)) + else: + span.set_status(Status(StatusCode.OK)) + finally: + span.end(end_time=_nanos(canonical_span.end_time)) + token = attach(set_value(_SUPPRESS_INSTRUMENTATION_KEY, True)) + try: + try: + export_result = self._exporter.export((cast(trace_sdk.ReadableSpan, span),)) + except Exception as error: + raise RetryableTraceDispatchError("Phoenix span export failed") from error + finally: + detach(token) + if export_result is not SpanExportResult.SUCCESS: + raise RetryableTraceDispatchError(f"Phoenix span export failed: canonical_span_id={canonical_span.id}") + if provider_parent_context is not None: + publish_parent_context(canonical_span.id, provider_parent_context) + + +class UnifiedPhoenixTrace(UnifiedTraceInstance): + """Fully isolated unified Phoenix trace instance selected by the new registry.""" + + def __init__(self, config: PhoenixConfig) -> None: + super().__init__( + config, + builder=CanonicalTraceBuilder(RepositoryWorkflowExecutionLoader(self.get_service_account_with_tenant)), + adapter=UnifiedPhoenixAdapter(config), + coordinator=ParentContextCoordinator(redis_client, resolve_parent_destination), + ) diff --git a/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/unified_trace.py b/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/unified_trace.py new file mode 100644 index 00000000000..a0d3bf6eef5 --- /dev/null +++ b/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/unified_trace.py @@ -0,0 +1,183 @@ +"""LangSmith adapter for the provider-neutral unified tracing runtime.""" + +from collections.abc import Mapping +from typing import Any, Literal +from uuid import NAMESPACE_URL, UUID, uuid5 + +from langsmith import Client +from langsmith.utils import ( + LangSmithAPIError, + LangSmithAuthError, + LangSmithConnectionError, + LangSmithError, + LangSmithRateLimitError, + LangSmithRequestTimeout, + LangSmithUserError, +) + +from core.ops.exceptions import InvalidTraceParentContextError, RetryableTraceDispatchError +from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace +from core.ops.unified_trace.parent_context import ( + ParentContextCoordinator, + ParentResolution, + ParentResolutionKind, + ProviderParentContext, + destination_scope, + resolve_parent_destination, +) +from core.ops.unified_trace.provider import ParentContextPublisher, UnifiedTraceInstance +from core.ops.unified_trace.trace_builder import CanonicalTraceBuilder, RepositoryWorkflowExecutionLoader +from core.ops.utils import generate_dotted_order +from dify_trace_langsmith.config import LangSmithConfig +from extensions.ext_redis import redis_client + +type LangSmithRunType = Literal["chain", "llm", "retriever", "tool"] + +_RUN_TYPE: dict[CanonicalSpanKind, LangSmithRunType] = { + CanonicalSpanKind.CHAIN: "chain", + CanonicalSpanKind.LLM: "llm", + CanonicalSpanKind.RETRIEVER: "retriever", + CanonicalSpanKind.TOOL: "tool", + CanonicalSpanKind.AGENT: "chain", +} + + +def _provider_run_id(canonical_id: str) -> str: + """Keep UUID execution IDs and deterministically map synthetic wrapper IDs.""" + try: + return str(UUID(canonical_id)) + except ValueError: + return str(uuid5(NAMESPACE_URL, f"dify-unified-trace:{canonical_id}")) + + +def _langsmith_value(value: Any, key: str) -> dict[str, Any]: + if isinstance(value, Mapping): + return dict(value) + return {key: value} + + +def _langsmith_inputs(span: CanonicalSpan) -> dict[str, Any]: + if span.metadata.get("trace_entity_type") == "message" and isinstance(span.inputs, str): + return {"messages": [{"role": "user", "content": span.inputs}]} + return _langsmith_value(span.inputs, "input") + + +class UnifiedLangSmithAdapter: + """Translate canonical spans to LangSmith Runs with explicit hierarchy.""" + + provider_name = "langsmith" + + def __init__(self, config: LangSmithConfig) -> None: + # Parent context is published only after create_run returns, so unified + # ordering requires synchronous writes rather than the SDK's default queue. + self._client = Client(api_key=config.api_key, api_url=config.endpoint, auto_batch_tracing=False) + self._project_name = config.project + self._scope = destination_scope(self.provider_name, config.endpoint, config.project) + + @property + def scope(self) -> str: + return self._scope + + def emit( + self, + trace: CanonicalTrace, + parent: ParentResolution | None, + publish_parent_context: ParentContextPublisher, + ) -> None: + provider_id_by_canonical_id = {span.id: _provider_run_id(span.id) for span in trace.spans} + root_provider_id = provider_id_by_canonical_id[trace.root_span_id] + restored_context = parent.context if parent and parent.kind is ParentResolutionKind.RESTORED else None + external_parent_id: str | None + external_parent_order: str | None + + if restored_context is not None: + trace_id = restored_context.trace_id + external_parent_id = restored_context.parent_id + external_parent_order = restored_context.provider_context.get("dotted_order") + if not external_parent_order: + raise InvalidTraceParentContextError("LangSmith parent context is missing dotted_order") + else: + trace_id = root_provider_id + external_parent_id = None + external_parent_order = None + + dotted_order_by_canonical_id: dict[str, str] = {} + for canonical_span in trace.spans: + provider_id = provider_id_by_canonical_id[canonical_span.id] + local_parent_id = ( + provider_id_by_canonical_id.get(canonical_span.parent_id or "") if canonical_span.parent_id else None + ) + parent_run_id = local_parent_id + parent_order = dotted_order_by_canonical_id.get(canonical_span.parent_id or "") + if canonical_span.id == trace.root_span_id: + parent_run_id = external_parent_id + parent_order = external_parent_order + + dotted_order = generate_dotted_order(provider_id, canonical_span.start_time, parent_order) + metadata = dict(canonical_span.metadata) + if canonical_span.id == trace.root_span_id: + if trace.session_id: + metadata["session_id"] = trace.session_id + if trace.trace_id != trace_id: + metadata.setdefault("external_trace_id", trace.trace_id) + if parent and parent.kind is ParentResolutionKind.LINKED_ROOT and parent.linked_parent: + metadata["linked_parent_workflow_run_id"] = parent.linked_parent.parent_workflow_run_id + metadata["linked_parent_node_execution_id"] = parent.linked_parent.parent_node_execution_id + metadata["dify.span.kind"] = canonical_span.kind.value + metadata.pop("dify.span.links", None) + if canonical_span.links: + metadata["dify.span.links"] = list(canonical_span.links) + + try: + self._client.create_run( + id=provider_id, + name=canonical_span.name, + inputs=_langsmith_inputs(canonical_span), + outputs=_langsmith_value(canonical_span.outputs, "output"), + run_type=_RUN_TYPE[canonical_span.kind], + start_time=canonical_span.start_time, + end_time=canonical_span.end_time, + error=canonical_span.error if canonical_span.status is CanonicalSpanStatus.ERROR else None, + extra={"metadata": metadata}, + tags=["dify", "synthetic" if canonical_span.synthetic else "execution"], + parent_run_id=parent_run_id, + trace_id=trace_id, + dotted_order=dotted_order, + session_name=self._project_name, + ) + except ( + LangSmithAPIError, + LangSmithConnectionError, + LangSmithRateLimitError, + LangSmithRequestTimeout, + ) as error: + raise RetryableTraceDispatchError("LangSmith run export failed") from error + except (LangSmithAuthError, LangSmithUserError) as error: + raise RuntimeError("LangSmith run export rejected") from error + except LangSmithError as error: + raise RuntimeError("LangSmith run export failed") from error + dotted_order_by_canonical_id[canonical_span.id] = dotted_order + + if canonical_span.can_parent_workflow or canonical_span.publishes_parent_context: + publish_parent_context( + canonical_span.id, + ProviderParentContext( + provider=self.provider_name, + scope=self.scope, + trace_id=trace_id, + parent_id=provider_id, + provider_context={"dotted_order": dotted_order}, + ), + ) + + +class UnifiedLangSmithTrace(UnifiedTraceInstance): + """Fully isolated unified LangSmith trace instance selected by the new registry.""" + + def __init__(self, config: LangSmithConfig) -> None: + super().__init__( + config, + builder=CanonicalTraceBuilder(RepositoryWorkflowExecutionLoader(self.get_service_account_with_tenant)), + adapter=UnifiedLangSmithAdapter(config), + coordinator=ParentContextCoordinator(redis_client, resolve_parent_destination), + ) diff --git a/api/tests/unit_tests/configs/test_dify_config.py b/api/tests/unit_tests/configs/test_dify_config.py index 85ae1054677..c4aa23ced0f 100644 --- a/api/tests/unit_tests/configs/test_dify_config.py +++ b/api/tests/unit_tests/configs/test_dify_config.py @@ -3,14 +3,43 @@ from typing import override import pytest from flask import Flask from packaging.version import Version -from pydantic import SecretStr +from pydantic import SecretStr, ValidationError from pydantic_settings import BaseSettings, PydanticBaseSettingsSource from yarl import URL from configs.app_config import DifyConfig +from configs.feature import OpsTraceConfig from enums import DeploymentEdition +def test_ops_trace_config_rejects_parent_context_ttl_shorter_than_retry_window() -> None: + with pytest.raises(ValidationError, match="must cover the retry window"): + OpsTraceConfig( + OPS_TRACE_UNIFIED_ENABLED=True, + OPS_TRACE_RETRYABLE_DISPATCH_MAX_RETRIES=4, + OPS_TRACE_RETRYABLE_DISPATCH_DELAY_SECONDS=5, + OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS=19, + ) + + +def test_ops_trace_config_skips_parent_context_validation_when_unified_tracing_is_disabled() -> None: + OpsTraceConfig( + OPS_TRACE_UNIFIED_ENABLED=False, + OPS_TRACE_RETRYABLE_DISPATCH_MAX_RETRIES=4, + OPS_TRACE_RETRYABLE_DISPATCH_DELAY_SECONDS=5, + OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS=19, + ) + + +def test_ops_trace_config_accepts_parent_context_ttl_covering_retry_window() -> None: + OpsTraceConfig( + OPS_TRACE_UNIFIED_ENABLED=True, + OPS_TRACE_RETRYABLE_DISPATCH_MAX_RETRIES=4, + OPS_TRACE_RETRYABLE_DISPATCH_DELAY_SECONDS=5, + OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS=20, + ) + + class _IsolatedDifyConfig(DifyConfig): """Load explicit test values and packaging metadata without consulting process state.""" diff --git a/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline.py b/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline.py index 5ebbd2112a9..71ea52be14b 100644 --- a/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline.py +++ b/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline.py @@ -34,6 +34,23 @@ def _build_pipeline() -> pipeline_module.AdvancedChatAppGenerateTaskPipeline: return pipeline +def test_process_passes_message_id_to_conversation_name_generation() -> None: + pipeline = _build_pipeline() + pipeline._conversation_id = "conversation-1" + pipeline._application_generate_entity = SimpleNamespace(query="hello", trace_manager=None) + pipeline._message_cycle_manager = mock.Mock() + pipeline._base_task_pipeline = SimpleNamespace(stream=True) + pipeline._wrapper_process_stream_response = mock.Mock(return_value=iter(())) + pipeline._to_stream_response = mock.Mock(return_value="stream") + + result = pipeline.process() + + assert result == "stream" + pipeline._message_cycle_manager.generate_conversation_name.assert_called_once_with( + conversation_id="conversation-1", query="hello", message_id="message-1" + ) + + def test_persist_human_input_extra_content_adds_record(monkeypatch: pytest.MonkeyPatch) -> None: pipeline = _build_pipeline() monkeypatch.setattr(pipeline, "_load_human_input_form_id", lambda **kwargs: "form-1") diff --git a/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline_core.py b/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline_core.py index 776d14aab9b..bdb4fc42317 100644 --- a/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline_core.py +++ b/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline_core.py @@ -858,7 +858,7 @@ class TestEasyUiBasedGenerateTaskPipeline: assert result == "streamed" pipeline._message_cycle_manager.generate_conversation_name.assert_called_once_with( - conversation_id="conv", query="hello" + conversation_id="conv", query="hello", message_id="msg" ) def test_process_routes_to_blocking_for_completion_mode(self): diff --git a/api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager_optimization.py b/api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager_optimization.py index c74dcda3c3e..535bd04bcfa 100644 --- a/api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager_optimization.py +++ b/api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager_optimization.py @@ -291,7 +291,9 @@ class TestMessageCycleManagerOptimization: ), patch("core.app.task_pipeline.message_cycle_manager.Timer", DummyTimer), ): - thread = message_cycle_manager.generate_conversation_name(conversation_id="conv-1", query="hello") + thread = message_cycle_manager.generate_conversation_name( + conversation_id="conv-1", query="hello", message_id="message-1" + ) assert isinstance(thread, DummyTimer) assert thread.interval == 1 @@ -301,6 +303,7 @@ class TestMessageCycleManagerOptimization: assert thread.kwargs["flask_app"] is flask_app assert thread.kwargs["conversation_id"] == "conv-1" assert thread.kwargs["query"] == "hello" + assert thread.kwargs["message_id"] == "message-1" assert message_cycle_manager._application_generate_entity.is_new_conversation is False def test_generate_conversation_name_skips_thread_when_auto_generate_disabled(self, message_cycle_manager): @@ -377,13 +380,18 @@ class TestMessageCycleManagerOptimization: mock_redis.get.return_value = None mock_llm_generator.generate_conversation_name.return_value = "generated-title" - message_cycle_manager._generate_conversation_name_worker(flask_app, "conv-1", "hello") + message_cycle_manager._generate_conversation_name_worker( + flask_app, "conv-1", "hello", message_id="message-1" + ) assert cycle_db.in_transaction() is False with Session(sqlite_engine) as verification_session: conversation = verification_session.get(Conversation, "conv-1") assert conversation is not None assert conversation.name == "generated-title" + mock_llm_generator.generate_conversation_name.assert_called_once_with( + "tenant-1", "hello", "conv-1", "app-id", message_id="message-1" + ) mock_redis.setex.assert_called_once() def test_generate_conversation_name_worker_falls_back_when_generation_fails( diff --git a/api/tests/unit_tests/core/llm_generator/test_llm_generator.py b/api/tests/unit_tests/core/llm_generator/test_llm_generator.py index af4aa42719a..2efef0b8fff 100644 --- a/api/tests/unit_tests/core/llm_generator/test_llm_generator.py +++ b/api/tests/unit_tests/core/llm_generator/test_llm_generator.py @@ -186,9 +186,13 @@ class TestLLMGenerator: mock_model_instance.invoke_llm.return_value = mock_response with patch("core.llm_generator.llm_generator.TraceQueueManager") as mock_trace: - name = LLMGenerator.generate_conversation_name("tenant_id", "test query") + name = LLMGenerator.generate_conversation_name( + "tenant_id", "test query", "conversation-1", "app-1", message_id="message-1" + ) assert name == "Test Conversation Name" - mock_trace.assert_called_once() + mock_trace.assert_called_once_with(app_id="app-1") + trace_task = mock_trace.return_value.add_trace_task.call_args.args[0] + assert trace_task.message_id == "message-1" def test_generate_conversation_name_truncated(self, mock_model_instance): long_query = "a" * 2100 diff --git a/api/tests/unit_tests/core/ops/test_ops_trace_manager.py b/api/tests/unit_tests/core/ops/test_ops_trace_manager.py index 3c3011597f0..687837c6caf 100644 --- a/api/tests/unit_tests/core/ops/test_ops_trace_manager.py +++ b/api/tests/unit_tests/core/ops/test_ops_trace_manager.py @@ -17,6 +17,7 @@ from sqlalchemy import Engine from sqlalchemy.orm import Session, sessionmaker import core.ops.ops_trace_manager as module +from configs import dify_config from core.ops.ops_trace_manager import OpsTraceManager, TraceQueueManager, TraceTask, TraceTaskName from core.rag.models.document import Document as RetrievalDocument from graphon.enums import WorkflowExecutionStatus @@ -49,6 +50,15 @@ class DummyTraceInstance: return "https://project.fake" +class DummyUnifiedTraceInstance(DummyTraceInstance): + pass + + +class FailingUnifiedTraceInstance(DummyTraceInstance): + def __init__(self, config): + raise RuntimeError("unified constructor failed") + + class FakeProviderMap: def __init__(self, data): self._data = data @@ -66,6 +76,11 @@ PROVIDER_ENTRY = { "trace_instance": DummyTraceInstance, } +UNIFIED_PROVIDER_ENTRY = { + "config_class": DummyConfig, + "trace_instance": DummyUnifiedTraceInstance, +} + class DummyTimer: def __init__(self, interval, function): @@ -120,9 +135,11 @@ class RecordingStorage: class RecordingDispatcher: def __init__(self) -> None: self.payloads: list[dict[str, str]] = [] + self.options: list[dict[str, object]] = [] - def delay(self, payload: dict[str, str]) -> None: - self.payloads.append(payload) + def apply_async(self, *, args: list[dict[str, str]], **kwargs: object) -> None: + self.payloads.append(args[0]) + self.options.append(kwargs) @pytest.fixture @@ -137,6 +154,8 @@ def database(sqlite_engine: Engine, sqlite_session: Session) -> Iterator[Session @pytest.fixture def trace_environment(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: monkeypatch.setattr(module, "provider_config_map", FakeProviderMap({"dummy": PROVIDER_ENTRY})) + monkeypatch.setattr(module, "unified_provider_config_map", FakeProviderMap({})) + monkeypatch.setattr(dify_config, "OPS_TRACE_UNIFIED_ENABLED", False) OpsTraceManager.ops_trace_instances_cache.clear() OpsTraceManager.decrypted_configs_cache.clear() monkeypatch.setattr(module.threading, "Timer", DummyTimer) @@ -346,6 +365,84 @@ def test_ops_trace_instance_uses_persisted_enabled_state_and_cache( assert OpsTraceManager.get_ops_trace_instance("missing") is None +@pytest.mark.parametrize( + ("enabled", "registered", "expected_type"), + [ + (False, False, DummyTraceInstance), + (False, True, DummyTraceInstance), + (True, False, DummyTraceInstance), + (True, True, DummyUnifiedTraceInstance), + ], +) +def test_ops_trace_instance_routes_by_unified_switch( + enabled: bool, + registered: bool, + expected_type: type[DummyTraceInstance], + monkeypatch: pytest.MonkeyPatch, + trace_environment: None, + encryption_functions, + database: Session, +) -> None: + app = _app(database, tracing=json.dumps({"enabled": True, "tracing_provider": "dummy"})) + database.add(TraceAppConfig(app_id=app.id, tracing_provider="dummy", tracing_config={})) + database.commit() + monkeypatch.setattr(dify_config, "OPS_TRACE_UNIFIED_ENABLED", enabled) + entries = {"dummy": UNIFIED_PROVIDER_ENTRY} if registered else {} + monkeypatch.setattr(module, "unified_provider_config_map", FakeProviderMap(entries)) + + instance = OpsTraceManager.get_ops_trace_instance(app.id) + + assert type(instance) is expected_type + + +def test_registered_unified_provider_does_not_fallback_when_construction_fails( + monkeypatch: pytest.MonkeyPatch, + trace_environment: None, + encryption_functions, + database: Session, +) -> None: + app = _app(database, tracing=json.dumps({"enabled": True, "tracing_provider": "dummy"})) + database.add(TraceAppConfig(app_id=app.id, tracing_provider="dummy", tracing_config={})) + database.commit() + monkeypatch.setattr(dify_config, "OPS_TRACE_UNIFIED_ENABLED", True) + monkeypatch.setattr( + module, + "unified_provider_config_map", + FakeProviderMap( + { + "dummy": { + "config_class": DummyConfig, + "trace_instance": FailingUnifiedTraceInstance, + } + } + ), + ) + + with pytest.raises(RuntimeError, match="unified constructor failed"): + OpsTraceManager.get_ops_trace_instance(app.id) + + +def test_unified_and_legacy_instances_have_separate_cache_entries( + monkeypatch: pytest.MonkeyPatch, + trace_environment: None, + encryption_functions, + database: Session, +) -> None: + app = _app(database, tracing=json.dumps({"enabled": True, "tracing_provider": "dummy"})) + database.add(TraceAppConfig(app_id=app.id, tracing_provider="dummy", tracing_config={})) + database.commit() + monkeypatch.setattr(module, "unified_provider_config_map", FakeProviderMap({"dummy": UNIFIED_PROVIDER_ENTRY})) + + monkeypatch.setattr(dify_config, "OPS_TRACE_UNIFIED_ENABLED", False) + legacy = OpsTraceManager.get_ops_trace_instance(app.id) + monkeypatch.setattr(dify_config, "OPS_TRACE_UNIFIED_ENABLED", True) + unified = OpsTraceManager.get_ops_trace_instance(app.id) + + assert type(legacy) is DummyTraceInstance + assert type(unified) is DummyUnifiedTraceInstance + assert legacy is not unified + + def test_message_config_lookup_uses_real_conversation_and_model_config(database: Session) -> None: app = _app(database) config = AppModelConfig(app_id=app.id, model='{"provider":"openai"}') @@ -402,7 +499,10 @@ def test_message_trace_reads_real_conversation_app_and_message_file( database.add(file) database.commit() monkeypatch.setattr(module, "get_message_data", lambda _message_id: _message_data()) - result = TraceTask(trace_type=TraceTaskName.MESSAGE_TRACE, message_id=message.id).message_trace(message.id) + result = TraceTask( + trace_type=TraceTaskName.MESSAGE_TRACE, + message_id=message.id, + ).preprocess() assert result.message_id == message.id assert result.conversation_mode == AppMode.CHAT assert result.file_list[0].endswith("path/to/file") @@ -565,7 +665,7 @@ def test_trace_helpers_and_streaming_metrics(trace_environment: None) -> None: assert OpsTraceManager.check_trace_config_is_effective({}, "dummy") assert OpsTraceManager.get_trace_config_project_key({}, "dummy") == "fake-key" assert OpsTraceManager.get_trace_config_project_url({}, "dummy") == "https://project.fake" - task = TraceTask(trace_type=TraceTaskName.MESSAGE_TRACE) + task = TraceTask(trace_type=TraceTaskName.MESSAGE_TRACE, message_id="message-1") assert task.conversation_trace(foo="bar") == {"foo": "bar"} assert task._extract_streaming_metrics(_message_data(message_metadata="invalid")) == {} assert task.generate_name_trace("conversation", {"start": 1, "end": 2}, tenant_id=None) == {} @@ -578,6 +678,7 @@ def test_trace_helpers_and_streaming_metrics(trace_environment: None) -> None: ) assert generated.outputs == "name" assert generated.tenant_id == "tenant-1" + assert generated.message_id == "message-1" def test_trace_queue_collect_run_and_storage_boundary(monkeypatch: pytest.MonkeyPatch, trace_environment: None) -> None: @@ -597,7 +698,7 @@ def test_trace_queue_collect_run_and_storage_boundary(monkeypatch: pytest.Monkey recording_storage = RecordingStorage() dispatcher = RecordingDispatcher() monkeypatch.setattr(module.storage, "save", recording_storage.save) - monkeypatch.setattr(module.process_trace_tasks, "delay", dispatcher.delay) + monkeypatch.setattr(module.process_trace_tasks, "apply_async", dispatcher.apply_async) file_id = UUID("00000000-0000-0000-0000-000000000123") monkeypatch.setattr(module, "uuid4", lambda: file_id) manager.add_trace_task(task) @@ -608,3 +709,55 @@ def test_trace_queue_collect_run_and_storage_boundary(monkeypatch: pytest.Monkey assert path.endswith(f"app-id/{file_id.hex}.json") assert json.loads(data)["app_id"] == "app-id" assert dispatcher.payloads == [{"file_id": file_id.hex, "app_id": "app-id"}] + + +def test_trace_queue_persists_with_caller_supplied_file_id( + monkeypatch: pytest.MonkeyPatch, trace_environment: None +) -> None: + monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", classmethod(lambda cls, _app_id: True)) + manager = TraceQueueManager(app_id="app-id", user_id="user-1") + task = TraceTask( + trace_type=TraceTaskName.GENERATE_NAME_TRACE, + conversation_id="conversation-1", + timer={"start": 1, "end": 2}, + tenant_id="tenant-1", + generate_conversation_name="name", + inputs="query", + ) + recording_storage = RecordingStorage() + monkeypatch.setattr(module.storage, "save", recording_storage.save) + + file_info = manager.persist_trace_task(task, file_id="workflow-final-run-1") + + assert file_info == {"file_id": "workflow-final-run-1", "app_id": "app-id"} + path, data = recording_storage.writes[0] + payload = json.loads(data) + assert path == "ops_trace/app-id/workflow-final-run-1.json" + assert UUID(payload["trace_info"]["operation_id"]) + + +def test_trace_queue_persistence_error_propagates(monkeypatch: pytest.MonkeyPatch, trace_environment: None) -> None: + monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", classmethod(lambda cls, _app_id: True)) + manager = TraceQueueManager(app_id="app-id", user_id="user-1") + task = TraceTask(trace_type=TraceTaskName.GENERATE_NAME_TRACE) + + def fail_save(_path: str, _data: bytes) -> None: + raise OSError("storage unavailable") + + monkeypatch.setattr(module.storage, "save", fail_save) + + with pytest.raises(OSError, match="storage unavailable"): + manager.persist_trace_task(task, file_id="workflow-final-run-1") + + +def test_trace_queue_enqueue_error_propagates(monkeypatch: pytest.MonkeyPatch, trace_environment: None) -> None: + monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", classmethod(lambda cls, _app_id: True)) + manager = TraceQueueManager(app_id="app-id", user_id="user-1") + + def fail_enqueue(*_args: object, **_kwargs: object) -> None: + raise ConnectionError("broker unavailable") + + monkeypatch.setattr(module.process_trace_tasks, "apply_async", fail_enqueue) + + with pytest.raises(ConnectionError, match="broker unavailable"): + manager.enqueue_persisted_trace({"file_id": "workflow-final-run-1", "app_id": "app-id"}) diff --git a/api/tests/unit_tests/core/ops/unified_trace/__init__.py b/api/tests/unit_tests/core/ops/unified_trace/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/api/tests/unit_tests/core/ops/unified_trace/test_entities.py b/api/tests/unit_tests/core/ops/unified_trace/test_entities.py new file mode 100644 index 00000000000..d4fb795f3b5 --- /dev/null +++ b/api/tests/unit_tests/core/ops/unified_trace/test_entities.py @@ -0,0 +1,108 @@ +from datetime import datetime + +import pytest +from pydantic import ValidationError + +from core.helper.trace_id_helper import ParentTraceContext +from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace + + +def span(span_id: str, parent_id: str | None = None) -> CanonicalSpan: + return CanonicalSpan( + id=span_id, + parent_id=parent_id, + name=span_id, + kind=CanonicalSpanKind.CHAIN, + start_time=datetime(2025, 1, 1), + end_time=None, + status=CanonicalSpanStatus.OK, + ) + + +def test_canonical_trace_rejects_unknown_fields() -> None: + with pytest.raises(ValidationError): + CanonicalTrace( + trace_id="trace-1", + session_id="session-1", + root_span_id="root-1", + spans=(), + unknown=True, # pyrefly: ignore[unexpected-keyword] + ) + + +def test_canonical_span_supports_links() -> None: + linked_span = CanonicalSpan( + id="linked-1", + parent_id="message-2", + name="linked", + kind=CanonicalSpanKind.CHAIN, + start_time=datetime(2025, 1, 1), + end_time=None, + status=CanonicalSpanStatus.OK, + links=("message-1",), + ) + + assert linked_span.links == ("message-1",) + + +@pytest.mark.parametrize( + ("root_span_id", "spans"), + [ + ("missing", (span("root"),)), + ("root", (span("root"), span("root", "root"))), + ("root", (span("root", "outside"),)), + ("root", (span("root"), span("child"))), + ("root", (span("root"), span("child", "later"), span("later", "root"))), + ], +) +def test_canonical_trace_rejects_invalid_fragment( + root_span_id: str, + spans: tuple[CanonicalSpan, ...], +) -> None: + with pytest.raises(ValidationError): + CanonicalTrace( + trace_id="trace-1", + session_id="session-1", + root_span_id=root_span_id, + spans=spans, + ) + + +def test_canonical_trace_rejects_conflicting_external_parent_modes() -> None: + with pytest.raises(ValidationError): + CanonicalTrace( + trace_id="trace-1", + session_id="session-1", + root_span_id="root", + spans=(span("root"),), + external_parent=ParentTraceContext( + parent_workflow_run_id="outer-run", + parent_node_execution_id="outer-node", + ), + required_parent_context_id="message-1", + ) + + +def test_canonical_span_is_immutable() -> None: + span = CanonicalSpan( + id="span-1", + parent_id=None, + name="root", + kind=CanonicalSpanKind.CHAIN, + start_time=datetime(2025, 1, 1), + end_time=None, + status=CanonicalSpanStatus.OK, + ) + + assert span.publishes_parent_context is False + + trace = CanonicalTrace( + trace_id="trace-1", + session_id="session-1", + root_span_id=span.id, + spans=(span,), + ) + assert trace.required_parent_context_id is None + + with pytest.raises(ValidationError): + span.name = "changed" # pyrefly: ignore[read-only] diff --git a/api/tests/unit_tests/core/ops/unified_trace/test_hierarchy.py b/api/tests/unit_tests/core/ops/unified_trace/test_hierarchy.py new file mode 100644 index 00000000000..e591e139343 --- /dev/null +++ b/api/tests/unit_tests/core/ops/unified_trace/test_hierarchy.py @@ -0,0 +1,107 @@ +from datetime import datetime, timedelta +from types import SimpleNamespace +from typing import cast + +from core.ops.unified_trace.hierarchy import WorkflowExecutionLike, build_workflow_hierarchy + + +def execution(**overrides: object) -> WorkflowExecutionLike: + values: dict[str, object] = { + "id": "exec-1", + "node_execution_id": None, + "node_id": "node-1", + "node_type": "tool", + "predecessor_node_id": None, + "iteration_id": None, + "iteration_index": None, + "loop_id": None, + "loop_index": None, + "created_at": datetime(2025, 1, 1), + "elapsed_time": 1.0, + "status": "succeeded", + "metadata": {}, + } + values.update(overrides) + return cast(WorkflowExecutionLike, SimpleNamespace(**values)) + + +def test_predecessor_becomes_parent_independent_of_repository_order() -> None: + start = execution(id="exec-start", node_id="start") + llm = execution(id="exec-llm", node_id="llm", predecessor_node_id="start") + + forward = build_workflow_hierarchy([start, llm]) + reverse = build_workflow_hierarchy([llm, start]) + + assert forward.parent_by_execution_id == {"exec-llm": "exec-start"} + assert reverse.parent_by_execution_id == forward.parent_by_execution_id + + +def test_repeated_graph_node_id_is_not_guessed_as_parent() -> None: + first = execution(id="exec-a1", node_id="a") + second = execution(id="exec-a2", node_id="a") + child = execution(id="exec-b", node_id="b", predecessor_node_id="a") + + result = build_workflow_hierarchy([first, second, child]) + + assert "exec-b" not in result.parent_by_execution_id + + +def test_iteration_child_is_parented_to_stable_wrapper() -> None: + container = execution(id="iteration-exec", node_id="iteration", node_type="iteration") + child = execution(id="child-exec", node_id="child", iteration_id="iteration", iteration_index=0) + + result = build_workflow_hierarchy([child, container]) + wrapper = result.wrapper_by_child_execution_id["child-exec"] + + assert wrapper.id == "iteration:iteration-exec:0" + assert wrapper.parent_execution_id == "iteration-exec" + assert result.parent_by_execution_id["child-exec"] == wrapper.id + + +def test_loop_wrapper_covers_child_times_and_failure() -> None: + container = execution(id="loop-exec", node_id="loop", node_type="loop") + first = execution( + id="first", + node_id="first-node", + loop_id="loop", + loop_index=2, + created_at=datetime(2025, 1, 1, 0, 0, 1), + elapsed_time=2, + ) + second = execution( + id="second", + node_id="second-node", + loop_id="loop", + loop_index=2, + created_at=datetime(2025, 1, 1, 0, 0, 2), + elapsed_time=4, + status="failed", + ) + + result = build_workflow_hierarchy([second, container, first]) + wrapper = result.wrappers[0] + + assert wrapper.id == "loop:loop-exec:2" + assert wrapper.start_time == first.created_at + assert wrapper.end_time == second.created_at + timedelta(seconds=4) + assert wrapper.has_error is True + assert wrapper.child_execution_ids == frozenset({"first", "second"}) + + +def test_invalid_wrapper_indexes_do_not_create_wrappers() -> None: + container = execution(id="loop-exec", node_id="loop", node_type="loop") + negative = execution(id="negative", node_id="negative-node", loop_id="loop", loop_index=-1) + boolean = execution(id="boolean", node_id="boolean-node", loop_id="loop", loop_index=True) + + result = build_workflow_hierarchy([container, negative, boolean]) + + assert result.wrappers == () + + +def test_cycle_edges_are_removed_deterministically() -> None: + first = execution(id="a-exec", node_id="a", predecessor_node_id="b") + second = execution(id="b-exec", node_id="b", predecessor_node_id="a") + + result = build_workflow_hierarchy([first, second]) + + assert result.parent_by_execution_id == {} diff --git a/api/tests/unit_tests/core/ops/unified_trace/test_parent_context.py b/api/tests/unit_tests/core/ops/unified_trace/test_parent_context.py new file mode 100644 index 00000000000..bd95b3636d5 --- /dev/null +++ b/api/tests/unit_tests/core/ops/unified_trace/test_parent_context.py @@ -0,0 +1,164 @@ +import json +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from core.helper.trace_id_helper import ParentTraceContext +from core.ops.exceptions import ( + InvalidTraceParentContextError, + PendingTraceParentContextError, + TraceParentContextAccessError, +) +from core.ops.unified_trace.parent_context import ( + ParentContextCoordinator, + ParentDestination, + ParentResolutionKind, + ProviderParentContext, + destination_scope, + parent_destination_from_config, +) + + +def parent() -> ParentTraceContext: + return ParentTraceContext(parent_workflow_run_id="outer-run", parent_node_execution_id="outer-tool") + + +def context(**overrides: object) -> ProviderParentContext: + values: dict[str, object] = { + "provider": "langsmith", + "scope": "scope-a", + "trace_id": "root-run", + "parent_id": "outer-tool", + "provider_context": {"dotted_order": "root.tool"}, + } + values.update(overrides) + return ProviderParentContext.model_validate(values) + + +def coordinator(redis: MagicMock, destination: ParentDestination | None) -> ParentContextCoordinator: + return ParentContextCoordinator(redis, lambda _workflow_run_id: destination) + + +def test_parent_destination_uses_non_secret_provider_scope() -> None: + destination = parent_destination_from_config( + "langsmith", + {"api_key": "secret", "endpoint": "https://smith.example", "project": "project-a"}, + unified=True, + ) + + assert destination == ParentDestination( + provider="langsmith", + scope=destination_scope("langsmith", "https://smith.example", "project-a"), + unified=True, + ) + assert "secret" not in destination.scope + + +def test_publish_uses_unified_namespace_and_configured_ttl(monkeypatch: pytest.MonkeyPatch) -> None: + redis = MagicMock() + value = context() + monkeypatch.setattr( + "core.ops.unified_trace.parent_context.dify_config", + SimpleNamespace(OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS=1_800), + ) + + coordinator(redis, None).publish("outer-tool", value) + + key, ttl, payload = redis.setex.call_args.args + assert key == "trace:unified:parent:outer-tool" + assert ttl == 1_800 + assert json.loads(payload)["provider"] == "langsmith" + + +def test_resolve_returns_compatible_context() -> None: + redis = MagicMock() + redis.get.return_value = context().model_dump_json().encode() + subject = coordinator(redis, ParentDestination(provider="langsmith", scope="scope-a", unified=True)) + + result = subject.resolve(parent(), expected_provider="langsmith", expected_scope="scope-a") + + assert result.kind is ParentResolutionKind.RESTORED + assert result.context == context() + assert result.linked_parent is None + + +def test_resolve_required_restores_message_context_without_destination_lookup() -> None: + redis = MagicMock() + redis.get.return_value = context().model_dump_json().encode() + resolve_destination = MagicMock() + subject = ParentContextCoordinator(redis, resolve_destination) + + result = subject.resolve_required( + "message-1", + expected_provider="langsmith", + expected_scope="scope-a", + ) + + assert result.kind is ParentResolutionKind.RESTORED + assert result.context == context() + redis.get.assert_called_once_with("trace:unified:parent:message-1") + resolve_destination.assert_not_called() + + +def test_missing_required_message_context_is_retryable() -> None: + redis = MagicMock() + redis.get.return_value = None + subject = coordinator(redis, None) + + with pytest.raises(PendingTraceParentContextError): + subject.resolve_required( + "message-1", + expected_provider="langsmith", + expected_scope="scope-a", + ) + + +def test_missing_compatible_context_is_retryable() -> None: + redis = MagicMock() + redis.get.return_value = None + subject = coordinator(redis, ParentDestination(provider="langsmith", scope="scope-a", unified=True)) + + with pytest.raises(PendingTraceParentContextError): + subject.resolve(parent(), expected_provider="langsmith", expected_scope="scope-a") + + +def test_non_unified_or_incompatible_parent_becomes_linked_root() -> None: + redis = MagicMock() + destinations = [ + None, + ParentDestination(provider="langsmith", scope="scope-a", unified=False), + ParentDestination(provider="phoenix", scope="scope-a", unified=True), + ParentDestination(provider="langsmith", scope="scope-b", unified=True), + ] + + for destination in destinations: + result = coordinator(redis, destination).resolve( + parent(), expected_provider="langsmith", expected_scope="scope-a" + ) + assert result.kind is ParentResolutionKind.LINKED_ROOT + assert result.linked_parent == parent() + + redis.get.assert_not_called() + + +def test_malformed_or_stale_context_is_terminal() -> None: + redis = MagicMock() + subject = coordinator(redis, ParentDestination(provider="langsmith", scope="scope-a", unified=True)) + + for payload in (b"not-json", b'{"version": 2}', context(scope="scope-b").model_dump_json().encode()): + redis.get.return_value = payload + with pytest.raises(InvalidTraceParentContextError): + subject.resolve(parent(), expected_provider="langsmith", expected_scope="scope-a") + + +def test_redis_read_and_write_failures_are_retryable() -> None: + redis = MagicMock() + redis.get.side_effect = ConnectionError("down") + redis.setex.side_effect = ConnectionError("down") + subject = coordinator(redis, ParentDestination(provider="langsmith", scope="scope-a", unified=True)) + + with pytest.raises(TraceParentContextAccessError): + subject.resolve(parent(), expected_provider="langsmith", expected_scope="scope-a") + with pytest.raises(TraceParentContextAccessError): + subject.publish("outer-tool", context()) diff --git a/api/tests/unit_tests/core/ops/unified_trace/test_provider.py b/api/tests/unit_tests/core/ops/unified_trace/test_provider.py new file mode 100644 index 00000000000..c0c4b2b7310 --- /dev/null +++ b/api/tests/unit_tests/core/ops/unified_trace/test_provider.py @@ -0,0 +1,120 @@ +from collections.abc import Callable +from datetime import datetime +from unittest.mock import MagicMock + +import pytest + +from core.helper.trace_id_helper import ParentTraceContext +from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace +from core.ops.unified_trace.parent_context import ParentResolution +from core.ops.unified_trace.provider import UnifiedTraceInstance + + +def canonical_trace(*, nested: bool = False, required_parent_context_id: str | None = None) -> CanonicalTrace: + parent = ( + ParentTraceContext(parent_workflow_run_id="outer-run", parent_node_execution_id="outer-tool") + if nested + else None + ) + return CanonicalTrace( + trace_id="trace-1", + session_id="session-1", + root_span_id="root-1", + external_parent=parent, + required_parent_context_id=required_parent_context_id, + spans=( + CanonicalSpan( + id="root-1", + parent_id=None, + name="root", + kind=CanonicalSpanKind.CHAIN, + start_time=datetime(2025, 1, 1), + end_time=datetime(2025, 1, 1, 0, 0, 1), + status=CanonicalSpanStatus.OK, + ), + ), + ) + + +def make_runtime(trace: CanonicalTrace) -> tuple[UnifiedTraceInstance, MagicMock, MagicMock, MagicMock]: + builder = MagicMock() + builder.build.return_value = trace + adapter = MagicMock(provider_name="langsmith", scope="scope-a") + adapter.provider_name = "langsmith" + adapter.scope = "scope-a" + coordinator = MagicMock() + runtime = UnifiedTraceInstance(MagicMock(), builder=builder, adapter=adapter, coordinator=coordinator) + return runtime, builder, adapter, coordinator + + +def test_runtime_passes_core_publisher_to_adapter() -> None: + runtime, _, adapter, coordinator = make_runtime(canonical_trace()) + provider_context = MagicMock() + + def emit( + _trace: CanonicalTrace, + _parent: ParentResolution | None, + publish_parent_context: Callable[[str, object], None], + ) -> None: + publish_parent_context("tool-exec", provider_context) + + adapter.emit.side_effect = emit + + runtime.trace(MagicMock()) + + coordinator.publish.assert_called_once_with("tool-exec", provider_context) + + +def test_runtime_resolves_nested_parent_before_emission() -> None: + runtime, _, adapter, coordinator = make_runtime(canonical_trace(nested=True)) + resolution = ParentResolution.restored(MagicMock()) + coordinator.resolve.return_value = resolution + + runtime.trace(MagicMock()) + + coordinator.resolve.assert_called_once() + assert coordinator.resolve.call_args.kwargs == { + "expected_provider": "langsmith", + "expected_scope": "scope-a", + } + adapter.emit.assert_called_once() + assert adapter.emit.call_args.args[1] is resolution + + +def test_runtime_resolves_required_message_parent_before_emission() -> None: + runtime, _, adapter, coordinator = make_runtime(canonical_trace(required_parent_context_id="message-1")) + resolution = ParentResolution.restored(MagicMock()) + events: list[str] = [] + coordinator.resolve_required.side_effect = lambda *_args, **_kwargs: events.append("resolve") or resolution + adapter.emit.side_effect = lambda *_args: events.append("emit") + + runtime.trace(MagicMock()) + + assert events == ["resolve", "emit"] + coordinator.resolve_required.assert_called_once_with( + "message-1", + expected_provider="langsmith", + expected_scope="scope-a", + ) + coordinator.resolve.assert_not_called() + assert adapter.emit.call_args.args[1] is resolution + + +def test_runtime_does_not_publish_when_adapter_fails_before_callback() -> None: + runtime, _, adapter, coordinator = make_runtime(canonical_trace()) + adapter.emit.side_effect = RuntimeError("provider rejected run") + + with pytest.raises(RuntimeError, match="provider rejected run"): + runtime.trace(MagicMock()) + + coordinator.publish.assert_not_called() + + +def test_runtime_has_no_legacy_fallback() -> None: + runtime, _, adapter, _ = make_runtime(canonical_trace()) + adapter.emit.side_effect = RuntimeError("terminal") + + with pytest.raises(RuntimeError, match="terminal"): + runtime.trace(MagicMock()) + + assert not hasattr(runtime, "legacy_provider") diff --git a/api/tests/unit_tests/core/ops/unified_trace/test_registry.py b/api/tests/unit_tests/core/ops/unified_trace/test_registry.py new file mode 100644 index 00000000000..0dd38b94bfc --- /dev/null +++ b/api/tests/unit_tests/core/ops/unified_trace/test_registry.py @@ -0,0 +1,14 @@ +import pytest + +from core.ops.entities.config_entity import TracingProviderEnum +from core.ops.unified_trace.registry import unified_provider_config_map + + +def test_registry_exposes_only_implemented_providers() -> None: + phoenix = unified_provider_config_map[TracingProviderEnum.PHOENIX] + langsmith = unified_provider_config_map[TracingProviderEnum.LANGSMITH] + + assert phoenix["trace_instance"].__name__ == "UnifiedPhoenixTrace" + assert langsmith["trace_instance"].__name__ == "UnifiedLangSmithTrace" + with pytest.raises(KeyError): + unified_provider_config_map[TracingProviderEnum.LANGFUSE] diff --git a/api/tests/unit_tests/core/ops/unified_trace/test_trace_builder.py b/api/tests/unit_tests/core/ops/unified_trace/test_trace_builder.py new file mode 100644 index 00000000000..3d1d0ee6325 --- /dev/null +++ b/api/tests/unit_tests/core/ops/unified_trace/test_trace_builder.py @@ -0,0 +1,536 @@ +from copy import deepcopy +from datetime import datetime, timedelta +from types import SimpleNamespace +from typing import cast +from unittest.mock import MagicMock +from uuid import UUID + +import pytest + +from core.ops.entities.trace_entity import ( + BaseTraceInfo, + DatasetRetrievalTraceInfo, + GenerateNameTraceInfo, + MessageTraceInfo, + ModerationTraceInfo, + SuggestedQuestionTraceInfo, + ToolTraceInfo, + WorkflowTraceInfo, +) +from core.ops.unified_trace.entities import CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace +from core.ops.unified_trace.trace_builder import ( + CanonicalTraceBuilder, + RepositoryWorkflowExecutionLoader, + WorkflowExecutionLike, + resolve_session_id, +) + + +def make_workflow_trace_info(**overrides: object) -> WorkflowTraceInfo: + values: dict[str, object] = { + "workflow_data": SimpleNamespace(), + "conversation_id": None, + "workflow_id": "workflow-1", + "tenant_id": "tenant-1", + "workflow_run_id": "run-1", + "workflow_run_elapsed_time": 1.0, + "workflow_run_status": "succeeded", + "workflow_run_inputs": {}, + "workflow_run_outputs": {}, + "workflow_run_version": "1", + "total_tokens": 0, + "file_list": [], + "query": "", + "metadata": {}, + } + values.update(overrides) + return WorkflowTraceInfo.model_validate(values) + + +def test_custom_session_id_wins_over_conversation_id() -> None: + info = make_workflow_trace_info( + conversation_id="conversation-1", + metadata={"trace_session_id": "customer-session"}, + ) + + assert resolve_session_id(info) == "customer-session" + + +def test_workflow_session_falls_back_to_conversation_then_run() -> None: + assert resolve_session_id(make_workflow_trace_info(conversation_id="conversation-1")) == "conversation-1" + assert resolve_session_id(make_workflow_trace_info()) == "run-1" + + +def test_nested_workflow_session_falls_back_to_parent_workflow() -> None: + info = make_workflow_trace_info( + metadata={ + "parent_trace_context": { + "parent_workflow_run_id": "parent-run", + "parent_node_execution_id": "parent-node-execution", + } + } + ) + + assert resolve_session_id(info) == "parent-run" + + +def test_message_session_falls_back_to_message_conversation() -> None: + info = MessageTraceInfo( + conversation_model="chat", + message_tokens=0, + answer_tokens=0, + total_tokens=0, + conversation_mode="chat", + message_data=SimpleNamespace(conversation_id="conversation-1", created_at=datetime(2025, 1, 1)), + metadata={}, + ) + + assert resolve_session_id(info) == "conversation-1" + + +def node_execution(**overrides: object) -> WorkflowExecutionLike: + values: dict[str, object] = { + "id": "node-exec", + "node_execution_id": None, + "node_id": "node", + "title": "Node", + "node_type": "tool", + "predecessor_node_id": None, + "iteration_id": None, + "iteration_index": None, + "loop_id": None, + "loop_index": None, + "created_at": datetime(2025, 1, 1), + "elapsed_time": 1.0, + "status": "succeeded", + "error": None, + "inputs": {"input": "value"}, + "outputs": {"output": "value"}, + "process_data": {}, + "metadata": {}, + } + values.update(overrides) + return cast(WorkflowExecutionLike, SimpleNamespace(**values)) + + +def workflow_info(**overrides: object) -> WorkflowTraceInfo: + workflow_data = SimpleNamespace( + created_at=datetime(2025, 1, 1), + finished_at=datetime(2025, 1, 1) + timedelta(seconds=5), + graph_dict={"nodes": []}, + ) + values: dict[str, object] = { + "workflow_data": workflow_data, + "start_time": workflow_data.created_at, + "end_time": workflow_data.finished_at, + "metadata": {"app_id": "app-1"}, + } + values.update(overrides) + return make_workflow_trace_info(**values) + + +CHATFLOW_INPUTS: dict[str, object] = { + "sys.app_id": "19a6d372-b8bc-4ad4-9b83-7e6e7138de31", + "sys.dialogue_count": 1, + "sys.files": [], + "sys.query": "hi", + "sys.user_id": "ca877a63-4d75-4ba3-a417-3edffe5e545c", + "sys.workflow_id": "8b81be9e-d7c1-4fa7-b90f-03791fa015ba", + "sys.workflow_run_id": "4eec02ea-4ed5-47cc-87fd-7c3821dd935d", +} + + +def test_chatflow_message_uses_query_while_workflow_keeps_complete_inputs() -> None: + builder = CanonicalTraceBuilder(lambda _info: []) + + trace = builder.build( + workflow_info( + message_id="message-1", + query="hi", + workflow_run_inputs=CHATFLOW_INPUTS, + ) + ) + + assert trace is not None + spans = {span.id: span for span in trace.spans} + assert trace.root_span_id == "message-1" + assert [span.id for span in trace.spans[:2]] == ["message-1", "run-1"] + assert spans["message-1"].name == "chatflow_run-1" + assert spans["message-1"].inputs == "hi" + assert spans["message-1"].metadata["trace_entity_type"] == "message" + assert spans["message-1"].publishes_parent_context is True + assert spans["run-1"].name == "workflow_run-1" + assert spans["run-1"].parent_id == "message-1" + assert spans["run-1"].inputs == CHATFLOW_INPUTS + + +def test_chatflow_message_falls_back_to_complete_inputs_for_empty_query() -> None: + builder = CanonicalTraceBuilder(lambda _info: []) + + trace = builder.build( + workflow_info( + message_id="message-1", + query="", + workflow_run_inputs=CHATFLOW_INPUTS, + ) + ) + + assert trace is not None + assert trace.spans[0].inputs == CHATFLOW_INPUTS + + +def test_workflow_without_message_keeps_complete_inputs_on_root() -> None: + builder = CanonicalTraceBuilder(lambda _info: []) + + trace = builder.build(workflow_info(query="hi", workflow_run_inputs=CHATFLOW_INPUTS)) + + assert trace is not None + assert trace.root_span_id == "run-1" + assert trace.spans[0].inputs == CHATFLOW_INPUTS + + +def test_build_workflow_trace_is_parent_first_and_uses_wrappers() -> None: + container = node_execution(id="iteration-exec", node_id="iteration", node_type="iteration", title="Items") + child = node_execution( + id="llm-exec", + node_id="llm", + node_type="llm", + title="Summarize", + iteration_id="iteration", + iteration_index=0, + process_data={ + "prompts": [{"role": "user", "text": "hello"}], + "model_provider": "openai", + "model_name": "gpt-4", + "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + }, + ) + loader = MagicMock(return_value=[child, container]) + builder = CanonicalTraceBuilder(loader) + + trace = builder.build(workflow_info()) + + assert trace is not None + assert [span.id for span in trace.spans] == [ + "run-1", + "iteration-exec", + "iteration:iteration-exec:0", + "llm-exec", + ] + spans = {span.id: span for span in trace.spans} + assert spans["llm-exec"].parent_id == "iteration:iteration-exec:0" + assert spans["llm-exec"].kind is CanonicalSpanKind.LLM + assert spans["llm-exec"].metadata["total_tokens"] == 30 + loader.assert_called_once() + + +def test_workflow_tool_is_marked_as_nested_workflow_parent() -> None: + builder = CanonicalTraceBuilder(lambda _info: [node_execution(id="tool-exec")]) + + trace = builder.build(workflow_info()) + + assert trace is not None + assert trace.spans[-1].can_parent_workflow is True + + +def retry_attempt(retry_index: object, **overrides: object) -> dict[str, object]: + values: dict[str, object] = { + "retry_index": retry_index, + "inputs": {"attempt": retry_index}, + "process_data": {"request": f"attempt-{retry_index}"}, + "outputs": {"status_code": 500}, + "error": f"attempt {retry_index} failed", + "elapsed_time": float(retry_index) if isinstance(retry_index, int) else 0.0, + "execution_metadata": {"internal": True}, + "created_at": 1_700_000_000, + "finished_at": 1_700_000_001, + } + values.update(overrides) + return values + + +def test_node_metadata_contains_compact_retry_summary_and_skips_malformed_entries() -> None: + history = [ + retry_attempt(1), + "malformed", + retry_attempt(True), + retry_attempt(2, error="attempt 2 timed out", elapsed_time=2.5), + retry_attempt(3), + ] + builder = CanonicalTraceBuilder(lambda _info: [node_execution(process_data={"__dify_retry_history": history})]) + + trace = builder.build(workflow_info()) + + assert trace is not None + metadata = trace.spans[-1].metadata + assert metadata["retry_count"] == 3 + assert metadata["retry_attempts"] == [ + { + "retry_index": 1, + "error": "attempt 1 failed", + "elapsed_time": 1.0, + "created_at": 1_700_000_000, + "finished_at": 1_700_000_001, + }, + { + "retry_index": 2, + "error": "attempt 2 timed out", + "elapsed_time": 2.5, + "created_at": 1_700_000_000, + "finished_at": 1_700_000_001, + }, + { + "retry_index": 3, + "error": "attempt 3 failed", + "elapsed_time": 3.0, + "created_at": 1_700_000_000, + "finished_at": 1_700_000_001, + }, + ] + assert all("inputs" not in attempt for attempt in metadata["retry_attempts"]) + assert all("process_data" not in attempt for attempt in metadata["retry_attempts"]) + assert all("outputs" not in attempt for attempt in metadata["retry_attempts"]) + assert all("execution_metadata" not in attempt for attempt in metadata["retry_attempts"]) + + +def test_node_without_retry_history_has_no_retry_metadata() -> None: + builder = CanonicalTraceBuilder(lambda _info: [node_execution()]) + + trace = builder.build(workflow_info()) + + assert trace is not None + assert "retry_count" not in trace.spans[-1].metadata + assert "retry_attempts" not in trace.spans[-1].metadata + + +def test_failed_node_preserves_error_without_mutating_trace_metadata() -> None: + metadata = {"app_id": "app-1", "custom": {"nested": True}} + original = deepcopy(metadata) + builder = CanonicalTraceBuilder( + lambda _info: [node_execution(status="failed", error="boom", metadata={"attempt": 1})] + ) + + trace = builder.build(workflow_info(metadata=metadata)) + + assert trace is not None + assert trace.spans[-1].status is CanonicalSpanStatus.ERROR + assert trace.spans[-1].error == "boom" + assert metadata == original + + +def test_nested_workflow_exposes_typed_external_parent() -> None: + builder = CanonicalTraceBuilder(lambda _info: []) + info = workflow_info( + metadata={ + "app_id": "app-1", + "parent_trace_context": { + "parent_workflow_run_id": "outer-run", + "parent_node_execution_id": "outer-tool", + }, + } + ) + + trace = builder.build(info) + + assert trace is not None + assert trace.external_parent is not None + assert trace.external_parent.parent_node_execution_id == "outer-tool" + + +def test_repository_loader_scopes_repository_to_trace(monkeypatch: pytest.MonkeyPatch) -> None: + repository = MagicMock() + repository.get_by_workflow_execution.return_value = [node_execution()] + factory = MagicMock() + factory.create_workflow_node_execution_repository.return_value = repository + monkeypatch.setattr("core.ops.unified_trace.trace_builder.DifyCoreRepositoryFactory", factory) + monkeypatch.setattr("core.ops.unified_trace.trace_builder.db", MagicMock(engine="engine")) + account = MagicMock() + get_account = MagicMock(return_value=account) + loader = RepositoryWorkflowExecutionLoader(get_account) + info = workflow_info() + + result = loader(info) + + assert len(result) == 1 + get_account.assert_called_once_with("app-1") + factory.create_workflow_node_execution_repository.assert_called_once() + call = factory.create_workflow_node_execution_repository.call_args.kwargs + assert call["tenant_id"] == "tenant-1" + assert call["app_id"] == "app-1" + repository.get_by_workflow_execution.assert_called_once_with(workflow_execution_id="run-1") + + +def test_message_trace_does_not_load_workflow_executions() -> None: + loader = MagicMock() + builder = CanonicalTraceBuilder(loader) + info = MessageTraceInfo( + conversation_model="chat", + message_tokens=2, + answer_tokens=3, + total_tokens=5, + conversation_mode="chat", + message_id="message-1", + message_data=SimpleNamespace( + id="message-1", + conversation_id="conversation-1", + answer="hello", + created_at=datetime(2025, 1, 1), + updated_at=datetime(2025, 1, 1, 0, 0, 1), + ), + inputs="hi", + metadata={}, + ) + + trace = builder.build(info) + + assert trace is not None + assert trace.root_span_id == "message-1" + assert [span.name for span in trace.spans] == ["message", "llm"] + assert trace.spans[0].outputs == "hello" + assert trace.spans[0].metadata["trace_entity_type"] == "message" + assert trace.spans[0].publishes_parent_context is True + assert trace.spans[1].parent_id == "message-1" + assert trace.spans[1].kind is CanonicalSpanKind.LLM + assert trace.spans[1].metadata["total_tokens"] == 5 + loader.assert_not_called() + + +def test_generate_name_uses_message_parent_and_conversation_session() -> None: + builder = CanonicalTraceBuilder(lambda _info: []) + info = GenerateNameTraceInfo( + tenant_id="tenant-1", + conversation_id="conversation-1", + message_id="message-1", + inputs="title prompt", + outputs="title", + metadata={}, + ) + + trace = builder.build(info) + + assert trace is not None + assert trace.trace_id == "message-1" + assert trace.session_id == "conversation-1" + assert trace.required_parent_context_id == "message-1" + assert trace.spans[0].parent_id is None + + +@pytest.mark.parametrize( + "info", + [ + ModerationTraceInfo( + message_id="message-1", + message_data=SimpleNamespace(id="message-1"), + flagged=False, + action="direct_output", + preset_response="", + query="hello", + metadata={}, + ), + SuggestedQuestionTraceInfo( + message_id="message-1", + message_data=SimpleNamespace(id="message-1"), + total_tokens=1, + suggested_question=["next"], + level="info", + metadata={}, + ), + DatasetRetrievalTraceInfo( + message_id="message-1", + message_data=SimpleNamespace(id="message-1"), + documents=[], + metadata={}, + ), + ToolTraceInfo( + message_id="message-1", + tool_name="search", + tool_inputs={}, + tool_outputs="done", + tool_config={}, + time_cost=0.1, + tool_parameters={}, + metadata={}, + ), + GenerateNameTraceInfo( + tenant_id="tenant-1", + conversation_id="conversation-1", + message_id="message-1", + metadata={}, + ), + ], +) +def test_standalone_message_child_uses_explicit_required_parent(info: BaseTraceInfo) -> None: + trace = CanonicalTraceBuilder(lambda _info: []).build(info) + + assert trace is not None + assert trace.spans[0].parent_id is None + assert trace.required_parent_context_id == "message-1" + assert_fragment_is_parent_first(trace) + + +def assert_fragment_is_parent_first(trace: CanonicalTrace) -> None: + seen: set[str] = set() + ids = [span.id for span in trace.spans] + assert len(ids) == len(set(ids)) + assert trace.root_span_id in ids + for span in trace.spans: + if span.id == trace.root_span_id: + assert span.parent_id is None + elif span.parent_id is not None: + assert span.parent_id in seen + seen.add(span.id) + + +def test_generate_name_without_message_remains_root_in_conversation_session() -> None: + builder = CanonicalTraceBuilder(lambda _info: []) + info = GenerateNameTraceInfo( + tenant_id="tenant-1", + conversation_id="conversation-1", + inputs="title prompt", + outputs="title", + metadata={}, + ) + + trace = builder.build(info) + + assert trace is not None + assert trace.session_id == "conversation-1" + assert trace.required_parent_context_id is None + assert trace.spans[0].parent_id is None + + +def test_standalone_trace_reuses_persisted_operation_id() -> None: + builder = CanonicalTraceBuilder(lambda _info: []) + info = GenerateNameTraceInfo( + tenant_id="tenant-1", + conversation_id="conversation-1", + inputs="title prompt", + outputs="title", + operation_id="00000000-0000-0000-0000-000000000099", + metadata={}, + ) + + first = builder.build(info) + second = builder.build(info) + + assert first is not None + assert second is not None + assert first.root_span_id == "00000000-0000-0000-0000-000000000099" + assert second.root_span_id == first.root_span_id + + +def test_standalone_trace_accepts_legacy_payload_without_operation_id() -> None: + info = GenerateNameTraceInfo.model_validate( + { + "tenant_id": "tenant-1", + "conversation_id": "conversation-1", + "inputs": "title prompt", + "outputs": "title", + "metadata": {}, + } + ) + + trace = CanonicalTraceBuilder(lambda _info: []).build(info) + + assert trace is not None + assert UUID(trace.root_span_id) diff --git a/api/tests/unit_tests/providers/trace/unified_trace/test_arize_phoenix.py b/api/tests/unit_tests/providers/trace/unified_trace/test_arize_phoenix.py new file mode 100644 index 00000000000..6f45ec277dd --- /dev/null +++ b/api/tests/unit_tests/providers/trace/unified_trace/test_arize_phoenix.py @@ -0,0 +1,258 @@ +import json +from datetime import datetime +from typing import cast +from unittest.mock import MagicMock + +import pytest +from dify_trace_arize_phoenix.config import PhoenixConfig +from dify_trace_arize_phoenix.unified_trace import UnifiedPhoenixAdapter +from openinference.semconv.trace import SpanAttributes +from opentelemetry.sdk.trace.export import SpanExportResult +from opentelemetry.trace import StatusCode +from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator + +from core.ops.exceptions import InvalidTraceParentContextError, RetryableTraceDispatchError +from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace +from core.ops.unified_trace.parent_context import ParentResolution, ProviderParentContext, destination_scope + +VALID_TRACEPARENT = "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01" +PhoenixAdapterFixture = tuple[UnifiedPhoenixAdapter, MagicMock, list[MagicMock]] + + +def _inject_traceparent(carrier: dict[str, str], context: object) -> None: + assert context is not None + carrier["traceparent"] = VALID_TRACEPARENT + + +def span(**overrides: object) -> CanonicalSpan: + values: dict[str, object] = { + "id": "root", + "parent_id": None, + "name": "root", + "kind": CanonicalSpanKind.CHAIN, + "start_time": datetime(2025, 1, 1), + "end_time": datetime(2025, 1, 1, 0, 0, 1), + "inputs": {"input": "value"}, + "outputs": {"output": "value"}, + "status": CanonicalSpanStatus.OK, + "metadata": {}, + } + values.update(overrides) + return CanonicalSpan.model_validate(values) + + +def trace(*spans: CanonicalSpan, session_id: str = "session-1") -> CanonicalTrace: + values = spans or (span(),) + return CanonicalTrace( + trace_id="trace-1", + session_id=session_id, + root_span_id=values[0].id, + spans=values, + ) + + +@pytest.fixture +def adapter(monkeypatch: pytest.MonkeyPatch) -> PhoenixAdapterFixture: + tracer = MagicMock() + otel_spans = [MagicMock() for _ in range(8)] + tracer.start_span.side_effect = otel_spans + exporter = MagicMock() + exporter.export.return_value = SpanExportResult.SUCCESS + monkeypatch.setattr( + "dify_trace_arize_phoenix.unified_trace.setup_unified_tracer", + lambda _config: (tracer, exporter), + ) + value = UnifiedPhoenixAdapter( + PhoenixConfig(api_key="secret", project="project-a", endpoint="https://phoenix.example") + ) + value._propagator = MagicMock() + return value, tracer, otel_spans + + +def test_emit_creates_parent_before_child_and_maps_session(adapter: PhoenixAdapterFixture) -> None: + subject, tracer, _ = adapter + root = span() + child = span(id="child", parent_id="root", name="llm", kind=CanonicalSpanKind.LLM) + + subject.emit(trace(root, child, session_id="customer-session"), None, MagicMock()) + + assert [call.kwargs["name"] for call in tracer.start_span.call_args_list] == ["root", "llm"] + root_attributes = tracer.start_span.call_args_list[0].kwargs["attributes"] + assert root_attributes[SpanAttributes.SESSION_ID] == "customer-session" + assert root_attributes[SpanAttributes.OPENINFERENCE_SPAN_KIND] == "CHAIN" + assert tracer.start_span.call_args_list[1].kwargs["context"] is not None + + +@pytest.mark.parametrize("kind", list(CanonicalSpanKind)) +def test_emit_maps_every_canonical_kind(adapter: PhoenixAdapterFixture, kind: CanonicalSpanKind) -> None: + expected = { + CanonicalSpanKind.CHAIN: "CHAIN", + CanonicalSpanKind.LLM: "LLM", + CanonicalSpanKind.RETRIEVER: "RETRIEVER", + CanonicalSpanKind.TOOL: "TOOL", + CanonicalSpanKind.AGENT: "AGENT", + }[kind] + subject, tracer, _ = adapter + + subject.emit(trace(span(kind=kind)), None, MagicMock()) + + attributes = tracer.start_span.call_args.kwargs["attributes"] + assert attributes[SpanAttributes.OPENINFERENCE_SPAN_KIND] == expected + metadata = json.loads(attributes[SpanAttributes.METADATA]) + assert metadata["dify.span.kind"] == kind.value + + +def test_emit_preserves_logical_links_and_overrides_reserved_metadata(adapter: PhoenixAdapterFixture) -> None: + subject, tracer, _ = adapter + linked_span = span( + metadata={"dify.span.kind": "forged", "dify.span.links": ["forged"]}, + links=("message-a",), + ) + + subject.emit(trace(linked_span), None, MagicMock()) + + attributes = tracer.start_span.call_args.kwargs["attributes"] + metadata = json.loads(attributes[SpanAttributes.METADATA]) + assert metadata["dify.span.kind"] == "chain" + assert metadata["dify.span.links"] == ["message-a"] + + +def test_emit_restores_w3c_parent_context(adapter: PhoenixAdapterFixture) -> None: + subject, _, _ = adapter + subject._propagator = MagicMock(wraps=TraceContextTextMapPropagator()) + provider_context = ProviderParentContext( + provider="phoenix", + scope=subject.scope, + trace_id="outer-trace", + parent_id="outer-tool", + provider_context={"traceparent": VALID_TRACEPARENT}, + ) + + subject.emit(trace(), ParentResolution.restored(provider_context), MagicMock()) + + subject._propagator.extract.assert_called_once_with(carrier={"traceparent": VALID_TRACEPARENT}) + + +def test_emit_rejects_restored_context_without_traceparent(adapter: PhoenixAdapterFixture) -> None: + subject, _, _ = adapter + provider_context = ProviderParentContext( + provider="phoenix", + scope=subject.scope, + trace_id="outer-trace", + parent_id="outer-tool", + provider_context={}, + ) + + with pytest.raises(InvalidTraceParentContextError): + subject.emit(trace(), ParentResolution.restored(provider_context), MagicMock()) + + +def test_emit_rejects_malformed_traceparent(adapter: PhoenixAdapterFixture) -> None: + subject, _, _ = adapter + subject._propagator = TraceContextTextMapPropagator() + provider_context = ProviderParentContext( + provider="phoenix", + scope=subject.scope, + trace_id="outer-trace", + parent_id="outer-tool", + provider_context={"traceparent": "malformed"}, + ) + + with pytest.raises(InvalidTraceParentContextError): + subject.emit(trace(), ParentResolution.restored(provider_context), MagicMock()) + + +def test_emit_publishes_tool_context_after_span_export(adapter: PhoenixAdapterFixture) -> None: + subject, tracer, otel_spans = adapter + cast(MagicMock, subject._propagator.inject).side_effect = _inject_traceparent + events: list[str] = [] + otel_spans[0].end.side_effect = lambda **_kwargs: events.append("end") + publish = MagicMock(side_effect=lambda *_args: events.append("publish")) + tool = span(id="tool-exec", kind=CanonicalSpanKind.TOOL, can_parent_workflow=True) + + subject.emit(trace(tool), None, publish) + + assert tracer.start_span.called + node_execution_id, context = publish.call_args.args + assert node_execution_id == "tool-exec" + assert context.provider == "phoenix" + assert context.provider_context == {"traceparent": VALID_TRACEPARENT} + assert events == ["end", "publish"] + + +def test_emit_does_not_publish_parent_context_when_export_fails(adapter: PhoenixAdapterFixture) -> None: + subject, _, _ = adapter + cast(MagicMock, subject._exporter.export).return_value = SpanExportResult.FAILURE + publish = MagicMock() + tool = span(id="tool-exec", kind=CanonicalSpanKind.TOOL, can_parent_workflow=True) + + with pytest.raises(RetryableTraceDispatchError, match="Phoenix span export failed"): + subject.emit(trace(tool), None, publish) + + publish.assert_not_called() + + +def test_emit_maps_exporter_exception_to_retryable_failure(adapter: PhoenixAdapterFixture) -> None: + subject, _, _ = adapter + cast(MagicMock, subject._exporter.export).side_effect = ConnectionError("network unavailable") + publish = MagicMock() + + with pytest.raises(RetryableTraceDispatchError, match="Phoenix span export failed"): + subject.emit(trace(), None, publish) + + publish.assert_not_called() + + +def test_emit_publishes_message_context(adapter: PhoenixAdapterFixture) -> None: + subject, _, _ = adapter + cast(MagicMock, subject._propagator.inject).side_effect = _inject_traceparent + publish = MagicMock() + message = span(id="message-1", name="message", publishes_parent_context=True) + + subject.emit(trace(message), None, publish) + + parent_id, context = publish.call_args.args + assert parent_id == "message-1" + assert context.provider_context == {"traceparent": VALID_TRACEPARENT} + + +def test_emit_records_error_status(adapter: PhoenixAdapterFixture) -> None: + subject, _, otel_spans = adapter + failed = span(status=CanonicalSpanStatus.ERROR, error="boom") + + subject.emit(trace(failed), None, MagicMock()) + + status = otel_spans[0].set_status.call_args.args[0] + assert status.status_code is StatusCode.ERROR + otel_spans[0].record_exception.assert_called_once() + + +def test_retry_metadata_is_serialized_for_phoenix(adapter: PhoenixAdapterFixture) -> None: + subject, tracer, _ = adapter + retry_metadata = { + "retry_count": 1, + "retry_attempts": [ + { + "retry_index": 1, + "error": "HTTP 500", + "elapsed_time": 1.2, + "created_at": 1_700_000_000, + "finished_at": 1_700_000_001, + } + ], + } + + subject.emit(trace(span(metadata=retry_metadata)), None, MagicMock()) + + attributes = tracer.start_span.call_args.kwargs["attributes"] + assert json.loads(attributes[SpanAttributes.METADATA]) == { + **retry_metadata, + "dify.span.kind": "chain", + } + + +def test_scope_does_not_include_api_key(adapter: PhoenixAdapterFixture) -> None: + subject, _, _ = adapter + + assert subject.scope == destination_scope("phoenix", "https://phoenix.example", "project-a") + assert "secret" not in subject.scope diff --git a/api/tests/unit_tests/providers/trace/unified_trace/test_langsmith.py b/api/tests/unit_tests/providers/trace/unified_trace/test_langsmith.py new file mode 100644 index 00000000000..da3532d316c --- /dev/null +++ b/api/tests/unit_tests/providers/trace/unified_trace/test_langsmith.py @@ -0,0 +1,328 @@ +from datetime import datetime +from unittest.mock import MagicMock + +import pytest +from dify_trace_langsmith.config import LangSmithConfig +from dify_trace_langsmith.unified_trace import UnifiedLangSmithAdapter +from langsmith.utils import ( + LangSmithAPIError, + LangSmithAuthError, + LangSmithConnectionError, + LangSmithError, + LangSmithRateLimitError, + LangSmithRequestTimeout, + LangSmithUserError, +) + +from core.ops.exceptions import InvalidTraceParentContextError, RetryableTraceDispatchError +from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace +from core.ops.unified_trace.parent_context import ParentResolution, ProviderParentContext, destination_scope + +ROOT_ID = "00000000-0000-0000-0000-000000000001" +CHILD_ID = "00000000-0000-0000-0000-000000000002" +TOOL_ID = "00000000-0000-0000-0000-000000000003" +LangSmithAdapterFixture = tuple[UnifiedLangSmithAdapter, MagicMock] + + +def span(**overrides: object) -> CanonicalSpan: + values: dict[str, object] = { + "id": ROOT_ID, + "parent_id": None, + "name": "root", + "kind": CanonicalSpanKind.CHAIN, + "start_time": datetime(2025, 1, 1), + "end_time": datetime(2025, 1, 1, 0, 0, 1), + "inputs": {"input": "value"}, + "outputs": {"output": "value"}, + "status": CanonicalSpanStatus.OK, + "metadata": {"external_trace_id": "customer-trace"}, + } + values.update(overrides) + return CanonicalSpan.model_validate(values) + + +def trace(*spans: CanonicalSpan, session_id: str = "session-1") -> CanonicalTrace: + values = spans or (span(),) + return CanonicalTrace( + trace_id="customer-trace", + session_id=session_id, + root_span_id=values[0].id, + spans=values, + ) + + +def test_client_disables_async_batching_before_parent_coordination(monkeypatch: pytest.MonkeyPatch) -> None: + client_class = MagicMock() + monkeypatch.setattr("dify_trace_langsmith.unified_trace.Client", client_class) + config = LangSmithConfig(api_key="secret", project="project-a", endpoint="https://smith.example") + + UnifiedLangSmithAdapter(config) + + client_class.assert_called_once_with( + api_key="secret", + api_url="https://smith.example", + auto_batch_tracing=False, + ) + + +@pytest.fixture +def adapter(monkeypatch: pytest.MonkeyPatch) -> LangSmithAdapterFixture: + client = MagicMock() + monkeypatch.setattr("dify_trace_langsmith.unified_trace.Client", lambda **_kwargs: client) + subject = UnifiedLangSmithAdapter( + LangSmithConfig(api_key="secret", project="project-a", endpoint="https://smith.example") + ) + return subject, client + + +def test_root_trace_id_equals_root_run_id_and_sets_thread_session(adapter: LangSmithAdapterFixture) -> None: + subject, client = adapter + + subject.emit(trace(session_id="customer-session"), None, MagicMock()) + + root = client.create_run.call_args.kwargs + assert root["id"] == ROOT_ID + assert root["trace_id"] == ROOT_ID + assert root["extra"]["metadata"]["session_id"] == "customer-session" + assert root["extra"]["metadata"]["external_trace_id"] == "customer-trace" + + +def test_message_span_uses_explicit_langsmith_human_message_schema(adapter: LangSmithAdapterFixture) -> None: + subject, client = adapter + message = span( + name="message", + inputs="hi", + metadata={"trace_entity_type": "message"}, + ) + + subject.emit(trace(message), None, MagicMock()) + + assert client.create_run.call_args.kwargs["inputs"] == {"messages": [{"role": "user", "content": "hi"}]} + + +def test_mapping_inputs_remain_unchanged(adapter: LangSmithAdapterFixture) -> None: + subject, client = adapter + raw_inputs: dict[str, object] = {"sys.app_id": "app-1", "sys.files": []} + message = span( + name="message", + inputs=raw_inputs, + metadata={"trace_entity_type": "message"}, + ) + + subject.emit(trace(message), None, MagicMock()) + + assert client.create_run.call_args.kwargs["inputs"] == raw_inputs + + +def test_empty_session_is_not_written_to_root_metadata(adapter: LangSmithAdapterFixture) -> None: + subject, client = adapter + + subject.emit(trace(session_id=""), None, MagicMock()) + + metadata = client.create_run.call_args.kwargs["extra"]["metadata"] + assert "session_id" not in metadata + + +def test_child_uses_actual_parent_run_and_dotted_order(adapter: LangSmithAdapterFixture) -> None: + subject, client = adapter + root = span() + child = span(id=CHILD_ID, parent_id=ROOT_ID, name="llm", kind=CanonicalSpanKind.LLM) + + subject.emit(trace(root, child), None, MagicMock()) + + root_run = client.create_run.call_args_list[0].kwargs + child_run = client.create_run.call_args_list[1].kwargs + assert child_run["parent_run_id"] == ROOT_ID + assert child_run["trace_id"] == ROOT_ID + assert child_run["dotted_order"].startswith(f"{root_run['dotted_order']}.") + assert child_run["run_type"] == "llm" + + +@pytest.mark.parametrize("kind", list(CanonicalSpanKind)) +def test_emit_maps_every_canonical_kind(adapter: LangSmithAdapterFixture, kind: CanonicalSpanKind) -> None: + expected = { + CanonicalSpanKind.CHAIN: "chain", + CanonicalSpanKind.LLM: "llm", + CanonicalSpanKind.RETRIEVER: "retriever", + CanonicalSpanKind.TOOL: "tool", + CanonicalSpanKind.AGENT: "chain", + }[kind] + subject, client = adapter + + subject.emit(trace(span(kind=kind)), None, MagicMock()) + + run = client.create_run.call_args.kwargs + assert run["run_type"] == expected + assert run["extra"]["metadata"]["dify.span.kind"] == kind.value + + +def test_emit_preserves_logical_links_and_overrides_reserved_metadata(adapter: LangSmithAdapterFixture) -> None: + subject, client = adapter + linked_span = span( + metadata={"dify.span.kind": "forged", "dify.span.links": ["forged"]}, + links=("message-a",), + ) + + subject.emit(trace(linked_span), None, MagicMock()) + + metadata = client.create_run.call_args.kwargs["extra"]["metadata"] + assert metadata["dify.span.kind"] == "chain" + assert metadata["dify.span.links"] == ["message-a"] + + +def test_synthetic_ids_are_mapped_consistently(adapter: LangSmithAdapterFixture) -> None: + subject, client = adapter + root = span() + wrapper = span(id="iteration:container:0", parent_id=ROOT_ID, name="iteration[0]") + child = span(id=CHILD_ID, parent_id=wrapper.id) + + subject.emit(trace(root, wrapper, child), None, MagicMock()) + + wrapper_run = client.create_run.call_args_list[1].kwargs + child_run = client.create_run.call_args_list[2].kwargs + assert wrapper_run["id"] != wrapper.id + assert child_run["parent_run_id"] == wrapper_run["id"] + + +def test_nested_workflow_restores_parent_trace_and_order(adapter: LangSmithAdapterFixture) -> None: + subject, client = adapter + parent = ProviderParentContext( + provider="langsmith", + scope=subject.scope, + trace_id="00000000-0000-0000-0000-000000000010", + parent_id="00000000-0000-0000-0000-000000000011", + provider_context={"dotted_order": "parent.order"}, + ) + + subject.emit(trace(), ParentResolution.restored(parent), MagicMock()) + + root = client.create_run.call_args.kwargs + assert root["trace_id"] == parent.trace_id + assert root["parent_run_id"] == parent.parent_id + assert root["dotted_order"].startswith("parent.order.") + + +def test_nested_workflow_rejects_parent_without_dotted_order(adapter: LangSmithAdapterFixture) -> None: + subject, _ = adapter + parent = ProviderParentContext( + provider="langsmith", + scope=subject.scope, + trace_id="00000000-0000-0000-0000-000000000010", + parent_id="00000000-0000-0000-0000-000000000011", + provider_context={}, + ) + + with pytest.raises(InvalidTraceParentContextError): + subject.emit(trace(), ParentResolution.restored(parent), MagicMock()) + + +def test_tool_context_is_published_only_after_create_run_succeeds(adapter: LangSmithAdapterFixture) -> None: + subject, client = adapter + publish = MagicMock() + tool = span(id=TOOL_ID, kind=CanonicalSpanKind.TOOL, can_parent_workflow=True) + + subject.emit(trace(tool), None, publish) + + assert client.create_run.called + node_execution_id, context = publish.call_args.args + assert node_execution_id == TOOL_ID + assert context.parent_id == TOOL_ID + assert context.provider_context["dotted_order"] + + client.reset_mock() + publish.reset_mock() + client.create_run.side_effect = RuntimeError("rejected") + with pytest.raises(RuntimeError, match="rejected"): + subject.emit(trace(tool), None, publish) + publish.assert_not_called() + + +@pytest.mark.parametrize( + "error", + [ + LangSmithConnectionError("down"), + LangSmithRequestTimeout("slow"), + LangSmithRateLimitError("limited"), + LangSmithAPIError("server"), + ], +) +def test_emit_maps_recoverable_sdk_errors_to_retryable_failure( + adapter: LangSmithAdapterFixture, error: LangSmithError +) -> None: + subject, client = adapter + client.create_run.side_effect = error + + with pytest.raises(RetryableTraceDispatchError) as exc_info: + subject.emit(trace(), None, MagicMock()) + assert exc_info.value.__cause__ is error + + +@pytest.mark.parametrize("error", [LangSmithAuthError("unauthorized"), LangSmithUserError("invalid")]) +def test_emit_maps_known_terminal_sdk_errors_to_runtime_failure( + adapter: LangSmithAdapterFixture, error: LangSmithError +) -> None: + subject, client = adapter + client.create_run.side_effect = error + + with pytest.raises(RuntimeError, match="LangSmith run export rejected") as exc_info: + subject.emit(trace(), None, MagicMock()) + assert not isinstance(exc_info.value, RetryableTraceDispatchError) + assert exc_info.value.__cause__ is error + + +def test_emit_maps_unknown_sdk_error_to_terminal_runtime_failure(adapter: LangSmithAdapterFixture) -> None: + subject, client = adapter + error = LangSmithError("unknown") + client.create_run.side_effect = error + + with pytest.raises(RuntimeError, match="LangSmith run export failed") as exc_info: + subject.emit(trace(), None, MagicMock()) + assert not isinstance(exc_info.value, RetryableTraceDispatchError) + assert exc_info.value.__cause__ is error + + +def test_retry_metadata_is_forwarded_to_langsmith(adapter: LangSmithAdapterFixture) -> None: + subject, client = adapter + retry_metadata = { + "retry_count": 1, + "retry_attempts": [ + { + "retry_index": 1, + "error": "HTTP 500", + "elapsed_time": 1.2, + "created_at": 1_700_000_000, + "finished_at": 1_700_000_001, + } + ], + } + node = span(metadata=retry_metadata) + + subject.emit(trace(node), None, MagicMock()) + + assert client.create_run.call_args.kwargs["extra"]["metadata"] == { + **retry_metadata, + "session_id": "session-1", + "external_trace_id": "customer-trace", + "dify.span.kind": "chain", + } + + +def test_message_context_is_published_after_create_run(adapter: LangSmithAdapterFixture) -> None: + subject, client = adapter + publish = MagicMock() + message = span(name="message", publishes_parent_context=True) + + subject.emit(trace(message), None, publish) + + assert client.create_run.called + parent_id, context = publish.call_args.args + assert parent_id == ROOT_ID + assert context.parent_id == ROOT_ID + assert context.provider_context["dotted_order"] + + +def test_scope_does_not_include_api_key(adapter: LangSmithAdapterFixture) -> None: + subject, _ = adapter + + assert subject.scope == destination_scope("langsmith", "https://smith.example", "project-a") + assert "secret" not in subject.scope diff --git a/api/tests/unit_tests/tasks/test_ops_trace_task.py b/api/tests/unit_tests/tasks/test_ops_trace_task.py index 5844c55c047..c23372250ae 100644 --- a/api/tests/unit_tests/tasks/test_ops_trace_task.py +++ b/api/tests/unit_tests/tasks/test_ops_trace_task.py @@ -7,6 +7,7 @@ from unittest.mock import MagicMock, patch import pytest from celery.exceptions import Retry +from configs import dify_config from core.ops.entities.config_entity import OPS_TRACE_FAILED_KEY from core.ops.exceptions import RetryableTraceDispatchError from tasks.ops_trace_task import process_trace_tasks @@ -215,8 +216,11 @@ def test_process_trace_tasks_skips_enterprise_trace_when_retry_payload_was_alrea mock_incr.assert_not_called() -def test_process_trace_tasks_default_retry_window_covers_parent_span_context_ttl(): - assert process_trace_tasks.max_retries * process_trace_tasks.default_retry_delay >= 300 +def test_process_trace_tasks_default_retry_window_covers_workflow_and_export_grace_period(): + assert ( + process_trace_tasks.max_retries * process_trace_tasks.default_retry_delay + >= dify_config.WORKFLOW_MAX_EXECUTION_TIME + 300 + ) def test_process_trace_tasks_deletes_payload_on_success(): diff --git a/dify-agent/.example.env b/dify-agent/.example.env index 5aeb0032157..1466a89d947 100644 --- a/dify-agent/.example.env +++ b/dify-agent/.example.env @@ -55,6 +55,8 @@ DIFY_AGENT_SHELL_REDACT_PATTERNS= # Use an HTTP(S) service root or an explicit /agent-stub API root. # Leave empty to avoid injecting DIFY_AGENT_STUB_* into shell.run jobs. DIFY_AGENT_STUB_API_BASE_URL=http://localhost:5050/agent-stub +# Optional bind override used only when DIFY_AGENT_STUB_API_BASE_URL uses grpc://. +DIFY_AGENT_STUB_GRPC_BIND_ADDRESS= # Dify API base URL reachable from the Sandbox for the signed /files/* data plane, # including Config file and skill pulls. DIFY_AGENT_SANDBOX_FILES_BASE_URL=http://localhost:5001 diff --git a/docker/envs/core-services/shared.env.example b/docker/envs/core-services/shared.env.example index 3cd595d9ea2..dead019db72 100644 --- a/docker/envs/core-services/shared.env.example +++ b/docker/envs/core-services/shared.env.example @@ -97,8 +97,12 @@ LOG_TZ=UTC DEBUG=false FLASK_DEBUG=false ENABLE_REQUEST_LOGGING=False +OPS_TRACE_UNIFIED_ENABLED=false OPS_TRACE_RETRYABLE_DISPATCH_MAX_RETRIES=60 OPS_TRACE_RETRYABLE_DISPATCH_DELAY_SECONDS=5 +# OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS should cover the retry window +# (OPS_TRACE_RETRYABLE_DISPATCH_MAX_RETRIES * OPS_TRACE_RETRYABLE_DISPATCH_DELAY_SECONDS). +OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS=3600 WORKFLOW_LOG_CLEANUP_ENABLED=false WORKFLOW_LOG_RETENTION_DAYS=30 WORKFLOW_LOG_CLEANUP_BATCH_SIZE=100