diff --git a/.devcontainer/post_create_command.sh b/.devcontainer/post_create_command.sh index 988eda2d6bed93..89a5e9c85f61b8 100755 --- a/.devcontainer/post_create_command.sh +++ b/.devcontainer/post_create_command.sh @@ -7,7 +7,7 @@ cd web && pnpm install pipx install uv echo "alias start-api=\"cd $WORKSPACE_ROOT/api && uv run python -m flask run --host 0.0.0.0 --port=5001 --debug\"" >> ~/.bashrc -echo "alias start-worker=\"cd $WORKSPACE_ROOT/api && uv run python -m celery -A app.celery worker -P threads -c 1 --loglevel INFO -Q dataset,dataset_summary,priority_dataset,priority_pipeline,pipeline,mail,ops_trace,app_deletion,plugin,workflow_storage,conversation,workflow,schedule_poller,schedule_executor,triggered_workflow_dispatcher,trigger_refresh_publisher,trigger_refresh_executor,retention\"" >> ~/.bashrc +echo "alias start-worker=\"$WORKSPACE_ROOT/dev/start-worker --pool threads\"" >> ~/.bashrc echo "alias start-web=\"cd $WORKSPACE_ROOT/web && pnpm dev\"" >> ~/.bashrc echo "alias start-web-prod=\"cd $WORKSPACE_ROOT/web && pnpm build && pnpm start\"" >> ~/.bashrc echo "alias start-containers=\"cd $WORKSPACE_ROOT/docker && docker-compose -f docker-compose.middleware.yaml -p dify --env-file middleware.env up -d\"" >> ~/.bashrc diff --git a/AGENTS.md b/AGENTS.md index 144d9a45269dad..f8a6a409c53d22 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -8,4 +8,4 @@ Before starting backend development, read the [architecture](docs/BACKEND_ARCHIT - Run backend commands through `uv run --project api `. - Backend integration tests are CI-only and are not expected to run locally. -- Keep `docker/.env.example` limited to variables required for a default Docker Compose deployment to start. Put optional and provider-specific settings in the matching `docker/envs/*.env.example` file; `docker/.env` overrides those service-specific env files. +- Keep `docker/.env.example` limited to settings that require manual user configuration. Omit settings whose built-in defaults work unchanged. Put optional and provider-specific settings in the matching `docker/envs/*.env.example` file; `docker/.env` overrides those service-specific env files. diff --git a/api/.env.example b/api/.env.example index e051a8ca8f5b2d..21ef0068c9ac62 100644 --- a/api/.env.example +++ b/api/.env.example @@ -118,10 +118,11 @@ CELERY_BROKER_URL=redis://:difyai123456@localhost:${REDIS_PORT}/1 CELERY_BACKEND=redis # Ops trace configuration -OPS_TRACE_UNIFIED_ENABLED=false -OPS_TRACE_RETRYABLE_DISPATCH_MAX_RETRIES=780 -OPS_TRACE_RETRYABLE_DISPATCH_DELAY_SECONDS=5 -OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS=3900 +# Keep new workers separate from legacy ops_trace workers during upgrades. +OPS_TRACE_QUEUE=ops_trace_v2 +OPS_TRACE_EXPORT_MAX_ATTEMPTS=6 +# Daily cleanup of trace payloads an interrupted delivery left in storage. +ENABLE_TRACE_PAYLOAD_SWEEP_TASK=true # Database configuration DB_TYPE=postgresql diff --git a/api/.importlinter b/api/.importlinter index 789f8fec427519..3502fc9237c030 100644 --- a/api/.importlinter +++ b/api/.importlinter @@ -6,6 +6,7 @@ root_packages = context configs controllers + domain extensions factories fields @@ -17,6 +18,26 @@ root_packages = services include_external_packages = True +[importlinter:contract:domain-boundary] +name = Domain code depends on no application, persistence, framework, configuration, queue, transport, or legacy kernel code +type = forbidden +source_modules = + domain +forbidden_modules = + celery + configs + controllers + core + enterprise + extensions + flask + models + repositories + services + sqlalchemy + tasks + werkzeug + [importlinter:contract:backend-layers] name = Backend modules follow the transport, application, execution-kernel, and shared-library layers type = layers @@ -1328,3 +1349,81 @@ forbidden_modules = sqlalchemy tasks werkzeug + +[importlinter:contract:tracing-application-boundary] +name = Trace pipeline and Agent trace collection services and contracts are framework and persistence neutral +type = forbidden +source_modules = + services.tracing.canonical + services.tracing.collection + services.tracing.collection_ports + services.tracing.delivery + services.tracing.errors + services.tracing.hierarchy + services.tracing.provider_configs + services.tracing.providers + services.tracing.records + services.tracing.recovery + services.tracing.scopes + services.tracing.service +forbidden_modules = + celery + configs + controllers + core + enterprise + extensions + flask + models + repositories + sqlalchemy + tasks + werkzeug + services.tracing.adapters + services.tracing.egress + services.tracing.registry + services.tracing.request_log_protection + services.tracing.runtime_bridge + +[importlinter:contract:tracing-controller-boundary] +name = Tracing handlers reach persistence and tracing adapters only through the composition root +type = forbidden +source_modules = + controllers.console.agent.tracing + controllers.console.agent.tracing_runs + controllers.console.app.ops_trace + controllers.inner_api.agent.tracing +forbidden_modules = + extensions.ext_database + models.application_trace + models.model + repositories + sqlalchemy + services.tracing.adapters + services.tracing.registry + services.tracing.runtime_bridge + services.tracing.service + domain +# The composition root (extensions.ext_application_services) wires repositories and +# adapters, so only direct imports are checked here. +allow_indirect_imports = True + +[importlinter:contract:tracing-repository-boundary] +name = Tracing persistence adapters only depend on the tracing contracts they implement +type = forbidden +source_modules = + repositories.app.tracing_config_repository + repositories.application_trace_repository + repositories.trace_record_repository + repositories.tracing_store_errors +forbidden_modules = + controllers + extensions + flask + tasks + werkzeug + services.tracing.adapters + services.tracing.registry + services.tracing.runtime_bridge + services.tracing.service +allow_indirect_imports = True diff --git a/api/app_factory.py b/api/app_factory.py index 51e38ed857452a..bbd4dc60364acc 100644 --- a/api/app_factory.py +++ b/api/app_factory.py @@ -186,7 +186,6 @@ def initialize_extensions(app: DifyApp): ext_database, ext_enterprise_telemetry, ext_fastopenapi, - ext_forward_refs, ext_hosting_provider, ext_import_modules, ext_key_provider, @@ -218,7 +217,6 @@ def initialize_extensions(app: DifyApp): ext_warnings, ext_import_modules, ext_orjson, - ext_forward_refs, ext_compress, ext_code_based_extension, ext_database, diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py index 774605b2cd3cb4..d627c7ca37f5bd 100644 --- a/api/configs/feature/__init__.py +++ b/api/configs/feature/__init__.py @@ -13,7 +13,6 @@ PositiveInt, computed_field, field_validator, - model_validator, ) from pydantic_settings import BaseSettings @@ -1438,6 +1437,11 @@ class MultiModalTransferConfig(BaseSettings): class OpsTraceConfig(BaseSettings): + OPS_TRACE_QUEUE: str = Field( + default="ops_trace_v2", + pattern=r"^[A-Za-z0-9_.-]+$", + description="Celery queue for tracing and sandbox metering; isolate new workers during rolling upgrades.", + ) OPS_AGENT_TRACE_ENABLED: bool = Field( default=False, description="Enable application-scoped Agent tracing collection and delivery." ) @@ -1463,37 +1467,11 @@ class OpsTraceConfig(BaseSettings): default=7, description="Retention of trace payloads and encrypted destination snapshots." ) - 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=780, - ) - - OPS_TRACE_RETRYABLE_DISPATCH_DELAY_SECONDS: PositiveInt = Field( - description="Delay in seconds between transient ops trace provider dispatch retry attempts.", - default=5, - ) - - OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS: PositiveInt = Field( - description="Retention in seconds for unified tracing parent contexts.", - default=3900, + OPS_TRACE_EXPORT_MAX_ATTEMPTS: PositiveInt = Field( + default=6, + description="Attempts to send one trace to its provider when the provider is temporarily unavailable.", ) - @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( @@ -1503,6 +1481,10 @@ class CeleryBeatConfig(BaseSettings): class CeleryScheduleTasksConfig(BaseSettings): + ENABLE_TRACE_PAYLOAD_SWEEP_TASK: bool = Field( + description="Enable the daily cleanup of trace payloads that an interrupted delivery left in storage", + default=True, + ) ENABLE_APPLICATION_TRACE_RECONCILE_TASK: bool = Field( description="Enable periodic application trace recovery and retention cleanup when Agent tracing is enabled", default=True, diff --git a/api/controllers/console/agent/tracing.py b/api/controllers/console/agent/tracing.py index 32c775991c540b..26eeec6ec5809c 100644 --- a/api/controllers/console/agent/tracing.py +++ b/api/controllers/console/agent/tracing.py @@ -1,6 +1,7 @@ """Agent tracing settings address Agent resources, never hidden inline Apps.""" from collections.abc import Callable +from dataclasses import asdict from typing import Literal from uuid import UUID @@ -9,7 +10,6 @@ from pydantic import Field from werkzeug.exceptions import NotFound -from configs import dify_config from controllers.common.rbac import AgentId, PlainApp, RBACCheck from controllers.common.schema import query_params_from_model, register_response_schema_models from controllers.console import console_ns @@ -206,22 +206,22 @@ def delete(self, context: RequestContext, agent_id: UUID): def _provider_capabilities(): - from services.tracing.provider_dispatch import provider_capabilities - query = TraceProviderQuery.model_validate(request.args.to_dict(flat=True)) try: - capabilities = provider_capabilities(query.tracing_provider) - except ValueError as error: + capabilities = application_services().tracing.provider_capabilities.get(query.tracing_provider) + except AppTracingConfigInvalidProviderError as error: raise UnsupportedTracingProviderError() from error + except AppTracingConfigProviderUnavailableError as error: + raise TracingProviderUnavailableError() from error return dump_response( TraceProviderCapabilitiesResponse, - {**capabilities, "api_collection_enabled": dify_config.OPS_AGENT_TRACE_ENABLED}, + {**asdict(capabilities.support), "api_collection_enabled": capabilities.api_collection_enabled}, ) @console_ns.route("/agent//trace-capabilities") class AgentTraceCapabilitiesApi(Resource): - @error_responses(UnsupportedTracingProviderError) + @error_responses(TracingProviderUnavailableError, UnsupportedTracingProviderError) @console_ns.doc("get_agent_trace_capabilities", params=query_params_from_model(TraceProviderQuery)) @console_ns.response( 200, "Provider support declarations", console_ns.models[TraceProviderCapabilitiesResponse.__name__] @@ -234,7 +234,7 @@ def get(self, context: RequestContext, agent_id: UUID): @console_ns.route("/apps//trace-capabilities") class AppTraceCapabilitiesApi(Resource): - @error_responses(UnsupportedTracingProviderError) + @error_responses(TracingProviderUnavailableError, UnsupportedTracingProviderError) @console_ns.doc("get_app_trace_capabilities", params=query_params_from_model(TraceProviderQuery)) @console_ns.response( 200, "Provider support declarations", console_ns.models[TraceProviderCapabilitiesResponse.__name__] diff --git a/api/controllers/console/agent/tracing_runs.py b/api/controllers/console/agent/tracing_runs.py index 60149273af5cf0..0b16860a3f43f2 100644 --- a/api/controllers/console/agent/tracing_runs.py +++ b/api/controllers/console/agent/tracing_runs.py @@ -23,9 +23,8 @@ from libs.helper import dump_response from machinery.context import RequestContext from models.account import TenantAccountRole -from repositories.application_trace_repository import TraceCollectionError from services.app.agent_app_contracts import AgentAppNotFoundError -from services.tracing.service import application_tracing_service +from services.tracing.errors import TraceCollectionError type TraceSourceKind = Literal["message", "workflow", "node"] _source_kind_adapter = TypeAdapter(TraceSourceKind) @@ -155,7 +154,7 @@ def _agent_app_id(context: RequestContext, agent_id: UUID) -> str: def _status(context: RequestContext, app_id: str, kind: TraceSourceKind, source_id: str): node_id = _query_node_id() result = _run_call( - lambda: application_tracing_service().status( + lambda: application_services().tracing.agent_traces.status( tenant_id=context.active_workspace_id, app_id=app_id, source_kind=kind, @@ -170,7 +169,7 @@ def _retry(context: RequestContext, app_id: str, kind: TraceSourceKind, source_i node_id = _query_node_id() payload = TraceRunRetryPayload.model_validate(console_ns.payload or {}) result = _run_call( - lambda: application_tracing_service().retry( + lambda: application_services().tracing.agent_traces.retry( tenant_id=context.active_workspace_id, app_id=app_id, source_kind=kind, diff --git a/api/controllers/console/app/ops_trace.py b/api/controllers/console/app/ops_trace.py index ce1463d8d28ca1..4fb0080bf1497e 100644 --- a/api/controllers/console/app/ops_trace.py +++ b/api/controllers/console/app/ops_trace.py @@ -73,7 +73,12 @@ class TraceAppConfigResponse(ResponseModel): id: str | None = None app_id: str | None = None tracing_provider: str | None = None - tracing_config: dict[str, str | None] | None = Field(default=None) + tracing_config: dict[str, str | None] | None = Field( + default=None, description="Non-secret fields; secrets are write-only and never returned" + ) + configured_secrets: list[str] | None = Field( + default=None, description="Secret fields that hold a value; submit them empty to keep that value" + ) is_active: bool | None = None created_at: str | None = None updated_at: str | None = None @@ -85,6 +90,7 @@ def from_record(cls, record: AppTracingConfigRecord) -> "TraceAppConfigResponse" app_id=record.app_id, tracing_provider=record.tracing_provider, tracing_config=record.tracing_config, + configured_secrets=list(record.configured_secrets), is_active=record.is_active, created_at=str(record.created_at), updated_at=str(record.updated_at), diff --git a/api/controllers/inner_api/agent/tracing.py b/api/controllers/inner_api/agent/tracing.py index 945f42eb09f891..b171debfd42773 100644 --- a/api/controllers/inner_api/agent/tracing.py +++ b/api/controllers/inner_api/agent/tracing.py @@ -13,9 +13,14 @@ from controllers.common.schema import register_response_schema_models, register_schema_models from controllers.inner_api import inner_api_ns from controllers.inner_api.wraps import agent_inner_api_only +from extensions.ext_application_services import application_services from fields.base import ResponseModel -from repositories.application_trace_repository import TraceCollectionError -from services.tracing.service import application_tracing_service +from services.tracing.errors import ( + TraceCollectionConflictError, + TraceCollectionError, + TraceCollectionLimitError, + TraceCollectionNotFoundError, +) class AgentTraceBatchResponse(ResponseModel): @@ -38,12 +43,17 @@ def post(self): request.max_content_length = dify_config.OPS_AGENT_TRACE_BATCH_MAX_BYTES try: payload = AgentTraceBatch.model_validate(inner_api_ns.payload or {}) - accepted = application_tracing_service().ingest(payload) + accepted = application_services().tracing.agent_traces.ingest(payload) except ValidationError as error: raise BadRequest("Invalid Agent trace batch") from error + except TraceCollectionNotFoundError as error: + raise NotFound(error.code) from error + except TraceCollectionConflictError as error: + raise Conflict(error.code) from error + except TraceCollectionLimitError as error: + raise RequestEntityTooLarge(error.code) from error except TraceCollectionError as error: - error_type = {404: NotFound, 409: Conflict, 413: RequestEntityTooLarge}.get(error.status, BadRequest) - raise error_type(error.code) from error + raise BadRequest(error.code) from error return AgentTraceBatchResponse(result="success", batch_id=payload.batch_id, accepted_spans=accepted).model_dump( mode="json" ) diff --git a/api/core/agent/cot_agent_runner.py b/api/core/agent/cot_agent_runner.py index b508cc70497382..0fab683ed629af 100644 --- a/api/core/agent/cot_agent_runner.py +++ b/api/core/agent/cot_agent_runner.py @@ -5,6 +5,7 @@ from collections.abc import Generator, Mapping, Sequence from typing import Any, TypedDict +from dify_trace_sdk.emitter import TraceScope from sqlalchemy.orm import Session from core.agent.base_agent_runner import BaseAgentRunner @@ -15,7 +16,6 @@ from core.app.entities.queue_entities import QueueAgentThoughtEvent, QueueMessageEndEvent, QueueMessageFileEvent from core.credit_usage import CreditUsageAppType, CreditUsageCreatedBy from core.model_context import use_credit_usage_metadata -from core.ops.ops_trace_manager import TraceQueueManager from core.prompt.agent_history_prompt_transform import AgentHistoryPromptTransform from core.tools.__base.tool import Tool from core.tools.entities.tool_entities import ToolInvokeMeta @@ -64,7 +64,7 @@ def run( self._repack_app_generate_entity(app_generate_entity) self._init_react_state(query) - trace_manager = app_generate_entity.trace_manager + trace_scope = app_generate_entity.trace_scope # check model mode if "Observation" not in app_generate_entity.model_conf.stop: @@ -247,7 +247,7 @@ def increase_usage(final_llm_usage_dict: dict[str, LLMUsage | None], usage: LLMU action=scratchpad.action, tool_instances=tool_instances, message_file_ids=message_file_ids, - trace_manager=trace_manager, + trace_scope=trace_scope, ) scratchpad.observation = tool_invoke_response scratchpad.agent_response = tool_invoke_response @@ -314,14 +314,14 @@ def _handle_invoke_action( action: AgentScratchpadUnit.Action, tool_instances: Mapping[str, Tool], message_file_ids: list[str], - trace_manager: TraceQueueManager | None = None, + trace_scope: TraceScope | None, ) -> tuple[str, ToolInvokeMeta]: """ handle invoke action :param action: action :param tool_instances: tool instances :param message_file_ids: message file ids - :param trace_manager: trace manager + :param trace_scope: where the tool call is reported, if anywhere :return: observation, meta """ # action is tool call, invoke tool @@ -348,7 +348,7 @@ def _handle_invoke_action( message=self.message, invoke_from=self.application_generate_entity.invoke_from, agent_tool_callback=self.agent_callback, - trace_manager=trace_manager, + trace_scope=trace_scope, ) session.commit() session.close() diff --git a/api/core/agent/fc_agent_runner.py b/api/core/agent/fc_agent_runner.py index 1b204608632436..64a38b34ce649a 100644 --- a/api/core/agent/fc_agent_runner.py +++ b/api/core/agent/fc_agent_runner.py @@ -125,8 +125,7 @@ def run(self, session: Session, message: Message, query: str, **kwargs: Any) -> final_answer = "" prompt_messages: list = [] # Initialize prompt_messages - # get tracing instance - trace_manager = app_generate_entity.trace_manager + trace_scope = app_generate_entity.trace_scope def increase_usage(final_llm_usage_dict: dict[str, LLMUsage | None], usage: LLMUsage): if not final_llm_usage_dict["usage"]: @@ -331,7 +330,7 @@ def increase_usage(final_llm_usage_dict: dict[str, LLMUsage | None], usage: LLMU message=self.message, invoke_from=self.application_generate_entity.invoke_from, agent_tool_callback=self.agent_callback, - trace_manager=trace_manager, + trace_scope=trace_scope, app_id=self.application_generate_entity.app_config.app_id, message_id=self.message.id, conversation_id=self.conversation.id, diff --git a/api/core/app/apps/advanced_chat/app_generator.py b/api/core/app/apps/advanced_chat/app_generator.py index 1b7f217058525c..ba669dc77ffaed 100644 --- a/api/core/app/apps/advanced_chat/app_generator.py +++ b/api/core/app/apps/advanced_chat/app_generator.py @@ -43,11 +43,11 @@ ) from core.app.layers.pause_state_persist_layer import PauseStateLayerConfig from core.helper.trace_id_helper import extract_external_trace_id_from_args, extract_trace_session_id_from_args -from core.ops.ops_trace_manager import TraceQueueManager from core.prompt.utils.get_thread_messages_length import get_thread_messages_length from core.repositories import DifyCoreRepositoryFactory from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository from extensions.ext_database import db +from extensions.tracing import create_trace_scope from factories import file_factory from graphon.engine.filter import ResponseStreamFilter from graphon.model_runtime.errors.invoke import InvokeAuthorizationError @@ -200,9 +200,12 @@ def generate( # convert to app config app_config = AdvancedChatAppConfigManager.get_app_config(app_model=app_model, workflow=workflow) - # get tracing instance - trace_manager = TraceQueueManager( - app_id=app_model.id, user_id=user.id if isinstance(user, Account) else user.session_id + trace_scope = create_trace_scope( + tenant_id=app_model.tenant_id, + app_id=app_model.id, + app_tracing=app_model.tracing, + user_id=user.id if isinstance(user, Account) else user.session_id, + extras=extras, ) if invoke_from == InvokeFrom.DEBUGGER: @@ -229,7 +232,7 @@ def generate( stream=streaming, invoke_from=invoke_from, extras=extras, - trace_manager=trace_manager, + trace_scope=trace_scope, workflow_run_id=str(workflow_run_id), ) contexts.plugin_tool_providers.set({}) @@ -292,15 +295,18 @@ def resume( """ Resume a paused advanced chat execution. - ``trace_manager`` is transient and excluded from generate-entity serialization, - so resumed executions rebuild it here before persistence layers receive the entity. + ``trace_scope`` is transient and excluded from generate-entity serialization, + so resumed executions create it again before persistence layers receive the entity. """ - if application_generate_entity.trace_manager is None: + if application_generate_entity.trace_scope is None: application_generate_entity = application_generate_entity.model_copy( update={ - "trace_manager": TraceQueueManager( + "trace_scope": create_trace_scope( + tenant_id=app_model.tenant_id, app_id=app_model.id, + app_tracing=app_model.tracing, user_id=user.id if isinstance(user, Account) else user.session_id, + extras=application_generate_entity.extras, ) } ) diff --git a/api/core/app/apps/advanced_chat/app_runner.py b/api/core/app/apps/advanced_chat/app_runner.py index f350345c764bd2..12687c89ee037b 100644 --- a/api/core/app/apps/advanced_chat/app_runner.py +++ b/api/core/app/apps/advanced_chat/app_runner.py @@ -360,7 +360,7 @@ def moderation_for_inputs( inputs=dict(inputs), query=query or "", message_id=message_id, - trace_manager=app_generate_entity.trace_manager, + trace_scope=app_generate_entity.trace_scope, ) def _initialize_conversation_variables(self) -> list[Variable]: 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 3d9581f06e705d..6a25ebc046e3e8 100644 --- a/api/core/app/apps/advanced_chat/generate_task_pipeline.py +++ b/api/core/app/apps/advanced_chat/generate_task_pipeline.py @@ -8,6 +8,7 @@ from threading import Thread from typing import Any +from dify_trace_sdk.events import MessageCompleted from sqlalchemy import select, update from sqlalchemy.orm import Session @@ -72,8 +73,6 @@ from core.app.task_pipeline.message_cycle_manager import MessageCycleManager from core.base.tts import AppGeneratorTTSPublisher from core.db.session_factory import session_factory -from core.ops.entities.trace_entity import TraceTaskName -from core.ops.ops_trace_manager import TraceQueueManager, TraceTask from core.workflow.file_reference import resolve_file_record_id from core.workflow.nodes.human_input.pause_reason import HumanInputRequired from core.workflow.system_variables import build_system_variables @@ -229,7 +228,7 @@ def process( message_id=self._message_id, ) - generator = self._wrapper_process_stream_response(trace_manager=self._application_generate_entity.trace_manager) + generator = self._wrapper_process_stream_response() if self._base_task_pipeline.stream: return self._to_stream_response(generator) @@ -350,9 +349,7 @@ def _listen_audio_msg(self, publisher: AppGeneratorTTSPublisher | None, task_id: return None raise RuntimeError(f"TTS publisher returned an unknown status: {audio_msg.status}") - def _wrapper_process_stream_response( - self, trace_manager: TraceQueueManager | None = None - ) -> Generator[StreamResponse]: + def _wrapper_process_stream_response(self) -> Generator[StreamResponse]: tts_publisher = None task_id = self._application_generate_entity.task_id tenant_id = self._application_generate_entity.app_config.tenant_id @@ -372,7 +369,7 @@ def _wrapper_process_stream_response( ) try: - for response in self._process_stream_response(tts_publisher=tts_publisher, trace_manager=trace_manager): + for response in self._process_stream_response(tts_publisher=tts_publisher): while audio_response := self._listen_audio_msg(publisher=tts_publisher, task_id=task_id): yield audio_response if tts_publisher and isinstance(response, ErrorStreamResponse): @@ -651,12 +648,9 @@ def _handle_loop_completed_event(self, event: QueueLoopCompletedEvent, **kwargs) def _handle_workflow_succeeded_event( self, event: QueueWorkflowSucceededEvent, - *, - trace_manager: TraceQueueManager | None = None, **kwargs, ) -> Generator[StreamResponse]: """Handle workflow succeeded events.""" - _ = trace_manager self._ensure_workflow_initialized() validated_state = self._ensure_graph_runtime_initialized() workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response( @@ -674,12 +668,9 @@ def _handle_workflow_succeeded_event( def _handle_workflow_partial_success_event( self, event: QueueWorkflowPartialSuccessEvent, - *, - trace_manager: TraceQueueManager | None = None, **kwargs, ) -> Generator[StreamResponse]: """Handle workflow partial success events.""" - _ = trace_manager self._ensure_workflow_initialized() validated_state = self._ensure_graph_runtime_initialized() workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response( @@ -738,12 +729,9 @@ def _handle_workflow_paused_event( def _handle_workflow_failed_event( self, event: QueueWorkflowFailedEvent, - *, - trace_manager: TraceQueueManager | None = None, **kwargs, ) -> Generator[StreamResponse]: """Handle workflow failed events.""" - _ = trace_manager self._ensure_workflow_initialized() validated_state = self._ensure_graph_runtime_initialized() @@ -768,11 +756,9 @@ def _handle_stop_event( event: QueueStopEvent, *, graph_runtime_state: RuntimeState | None = None, - trace_manager: TraceQueueManager | None = None, **kwargs, ) -> Generator[StreamResponse]: """Handle stop events.""" - _ = trace_manager resolved_state = None if self._workflow_run_id: resolved_state = self._resolve_graph_runtime_state(graph_runtime_state) @@ -935,7 +921,6 @@ def _dispatch_event( event: Any, *, tts_publisher: AppGeneratorTTSPublisher | None = None, - trace_manager: TraceQueueManager | None = None, queue_message: WorkflowQueueMessage | MessageQueueMessage | None = None, ) -> Generator[StreamResponse]: """Dispatch events using elegant pattern matching.""" @@ -947,7 +932,6 @@ def _dispatch_event( yield from handler( event, tts_publisher=tts_publisher, - trace_manager=trace_manager, queue_message=queue_message, ) return @@ -963,7 +947,6 @@ def _dispatch_event( yield from self._handle_node_failed_events( event, tts_publisher=tts_publisher, - trace_manager=trace_manager, queue_message=queue_message, ) return @@ -974,7 +957,6 @@ def _dispatch_event( def _process_stream_response( self, tts_publisher: AppGeneratorTTSPublisher | None = None, - trace_manager: TraceQueueManager | None = None, ) -> Generator[StreamResponse]: """ Process stream response using elegant Fluent Python patterns. @@ -993,22 +975,22 @@ def _process_stream_response( break case QueueWorkflowFailedEvent(): - yield from self._handle_workflow_failed_event(event, trace_manager=trace_manager) + yield from self._handle_workflow_failed_event(event) break case QueueWorkflowPausedEvent(): yield from self._handle_workflow_paused_event(event) break case QueueWorkflowSucceededEvent(): - yield from self._handle_workflow_succeeded_event(event, trace_manager=trace_manager) + yield from self._handle_workflow_succeeded_event(event) break case QueueWorkflowPartialSuccessEvent(): - yield from self._handle_workflow_partial_success_event(event, trace_manager=trace_manager) + yield from self._handle_workflow_partial_success_event(event) break case QueueStopEvent(): - yield from self._handle_stop_event(event, graph_runtime_state=None, trace_manager=trace_manager) + yield from self._handle_stop_event(event, graph_runtime_state=None) break # Handle all other events through elegant dispatch @@ -1017,7 +999,6 @@ def _process_stream_response( self._dispatch_event( event, tts_publisher=tts_publisher, - trace_manager=trace_manager, queue_message=queue_message, ) ): @@ -1086,16 +1067,10 @@ def _save_message(self, *, session: Session, graph_runtime_state: RuntimeState | session.add_all(message_files) def _emit_message_trace(self) -> None: - trace_manager = self._application_generate_entity.trace_manager - if trace_manager: - trace_manager.add_trace_task( - TraceTask( - TraceTaskName.MESSAGE_TRACE, - conversation_id=self._conversation_id, - message_id=self._message_id, - trace_session_id=self._application_generate_entity.extras.get("trace_session_id"), - ) - ) + """Report the finished message once its records are committed.""" + trace_scope = self._application_generate_entity.trace_scope + if trace_scope is not None: + trace_scope.emit(MessageCompleted(message_id=self._message_id, conversation_id=self._conversation_id)) def _seed_graph_runtime_state_from_queue_manager(self) -> None: """Bootstrap the cached runtime state from the queue manager when present.""" diff --git a/api/core/app/apps/agent_app/app_generator.py b/api/core/app/apps/agent_app/app_generator.py index 0c402bccf5fbe8..2e4aa90185bd2e 100644 --- a/api/core/app/apps/agent_app/app_generator.py +++ b/api/core/app/apps/agent_app/app_generator.py @@ -49,8 +49,8 @@ ) from core.credit_usage import CreditUsageAppType from core.db.session_factory import session_factory -from core.ops.ops_trace_manager import TraceQueueManager from extensions.ext_database import db +from extensions.tracing import application_tracing_hooks, create_trace_scope from factories import file_factory from models import Account, App, AppModelConfig, Conversation, EndUser, Message, MessageAnnotation from models.agent import ( @@ -142,7 +142,14 @@ def generate( else [] ) - trace_manager = TraceQueueManager(app_model.id, user.id if isinstance(user, Account) else user.session_id) + extras = {"auto_generate_conversation_name": args.get("auto_generate_name", True)} + trace_scope = create_trace_scope( + tenant_id=app_model.tenant_id, + app_id=app_model.id, + app_tracing=app_model.tracing, + user_id=user.id if isinstance(user, Account) else user.session_id, + extras=extras, + ) application_generate_entity = AgentAppGenerateEntity( task_id=str(uuid.uuid4()), app_config=app_config, @@ -162,11 +169,9 @@ def generate( user_id=user.id, stream=streaming, invoke_from=invoke_from, - extras={ - "auto_generate_conversation_name": args.get("auto_generate_name", True), - }, + extras=extras, call_depth=0, - trace_manager=trace_manager, + trace_scope=trace_scope, agent_id=agent.id, agent_config_snapshot_id=agent_config_id, agent_config_version_kind=agent_config_version_kind, @@ -314,8 +319,6 @@ def _generate_worker( # A local stop or failed cancellation wait is still just intent. pass except Exception as e: - from core.ops.base_trace_instance import application_tracing_hooks - tracing_hooks = application_tracing_hooks() if tracing_hooks is not None: tracing_hooks.capture_agent_app_terminal( @@ -387,7 +390,7 @@ def _run_input_guards( inputs=dict(application_generate_entity.inputs), query=query or "", message_id=message.id, - trace_manager=application_generate_entity.trace_manager, + trace_scope=application_generate_entity.trace_scope, ) except ModerationError as e: publish_text_answer(queue_manager=queue_manager, model_name=model_name, answer=str(e), user_query=query) diff --git a/api/core/app/apps/agent_app/app_runner.py b/api/core/app/apps/agent_app/app_runner.py index b2e7f04ae62c05..51375b395aa949 100644 --- a/api/core/app/apps/agent_app/app_runner.py +++ b/api/core/app/apps/agent_app/app_runner.py @@ -50,6 +50,7 @@ QueueRetrieverResourcesEvent, ) from extensions.ext_database import db +from extensions.tracing import application_tracing_hooks from graphon.file import File from graphon.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage from graphon.model_runtime.entities.message_entities import ( @@ -765,8 +766,6 @@ def run( binding_id=runtime.binding_id, snapshot=terminal.session_snapshot, ) - from core.ops.base_trace_instance import application_tracing_hooks - tracing_hooks = application_tracing_hooks() if tracing_hooks is not None: tracing_hooks.capture_agent_app_terminal( @@ -1059,8 +1058,6 @@ def _finish_stopped_terminal( # snapshot/usage, but let the original business failure own the root. # Check after cleanup so a Stop observed during its wait still wins. return - from core.ops.base_trace_instance import application_tracing_hooks - tracing_hooks = application_tracing_hooks() if tracing_hooks is not None: outcome: Literal["succeeded", "failed", "cancelled"] diff --git a/api/core/app/apps/agent_app/runtime_request_builder.py b/api/core/app/apps/agent_app/runtime_request_builder.py index c2d8accae1f1e1..65235768807601 100644 --- a/api/core/app/apps/agent_app/runtime_request_builder.py +++ b/api/core/app/apps/agent_app/runtime_request_builder.py @@ -37,7 +37,6 @@ from configs import dify_config from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom from core.app.llm.model_access import resolve_model_context_window, resolve_model_supports_vision -from core.ops.base_trace_instance import application_tracing_hooks from core.plugin.provider_identity import normalize_plugin_daemon_provider_identity from core.workflow.file_reference import build_file_reference, is_canonical_file_reference from core.workflow.nodes.agent_v2.dify_tools_builder import ( @@ -54,6 +53,7 @@ build_shell_layer_config, load_runtime_agent_skill_configs, ) +from extensions.tracing import application_tracing_hooks from graphon.file import File, FileTransferMethod, FileType, file_manager from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent from models.agent_config_entities import AgentKnowledgeSpaceConfig, AgentSoulConfig, AgentSoulToolsConfig diff --git a/api/core/app/apps/agent_chat/app_generator.py b/api/core/app/apps/agent_chat/app_generator.py index 2f661f066452a5..10aac614d4c9e6 100644 --- a/api/core/app/apps/agent_chat/app_generator.py +++ b/api/core/app/apps/agent_chat/app_generator.py @@ -23,7 +23,7 @@ from core.app.entities.app_invoke_entities import AgentChatAppGenerateEntity, InvokeFrom from core.db.session_factory import session_factory from core.helper.trace_id_helper import extract_trace_session_id_from_args -from core.ops.ops_trace_manager import TraceQueueManager +from extensions.tracing import create_trace_scope from factories import file_factory from graphon.model_runtime.errors.invoke import InvokeAuthorizationError from libs.flask_utils import preserve_flask_contexts @@ -179,8 +179,13 @@ def generate( annotation_reply=annotation_reply, ) - # get tracing instance - trace_manager = TraceQueueManager(app_model.id, user.id if isinstance(user, Account) else user.session_id) + trace_scope = create_trace_scope( + tenant_id=app_model.tenant_id, + app_id=app_model.id, + app_tracing=app_model.tracing, + user_id=user.id if isinstance(user, Account) else user.session_id, + extras=extras, + ) # init application generate entity application_generate_entity = AgentChatAppGenerateEntity( @@ -204,7 +209,7 @@ def generate( invoke_from=invoke_from, extras=extras, call_depth=0, - trace_manager=trace_manager, + trace_scope=trace_scope, ) # init generate records diff --git a/api/core/app/apps/base_app_runner.py b/api/core/app/apps/base_app_runner.py index 774cf255968bc7..899b1a2e4fc133 100644 --- a/api/core/app/apps/base_app_runner.py +++ b/api/core/app/apps/base_app_runner.py @@ -482,7 +482,7 @@ def moderation_for_inputs( inputs=dict(inputs), query=query or "", message_id=message_id, - trace_manager=app_generate_entity.trace_manager, + trace_scope=app_generate_entity.trace_scope, ) def check_hosting_moderation( diff --git a/api/core/app/apps/chat/app_generator.py b/api/core/app/apps/chat/app_generator.py index d0be650e4871a2..badb527a366e16 100644 --- a/api/core/app/apps/chat/app_generator.py +++ b/api/core/app/apps/chat/app_generator.py @@ -23,8 +23,8 @@ from core.app.entities.app_invoke_entities import ChatAppGenerateEntity, InvokeFrom from core.db.session_factory import session_factory from core.helper.trace_id_helper import extract_trace_session_id_from_args -from core.ops.ops_trace_manager import TraceQueueManager from extensions.ext_database import db +from extensions.tracing import create_trace_scope from factories import file_factory from graphon.model_runtime.errors.invoke import InvokeAuthorizationError from models import Account @@ -168,9 +168,12 @@ def generate( annotation_reply=annotation_reply, ) - # get tracing instance - trace_manager = TraceQueueManager( - app_id=app_model.id, user_id=user.id if isinstance(user, Account) else user.session_id + trace_scope = create_trace_scope( + tenant_id=app_model.tenant_id, + app_id=app_model.id, + app_tracing=app_model.tracing, + user_id=user.id if isinstance(user, Account) else user.session_id, + extras=extras, ) # init application generate entity @@ -193,7 +196,7 @@ def generate( user_id=user.id, invoke_from=invoke_from, extras=extras, - trace_manager=trace_manager, + trace_scope=trace_scope, stream=streaming, ) diff --git a/api/core/app/apps/completion/app_generator.py b/api/core/app/apps/completion/app_generator.py index 7312b079fe8e75..fde123a4dfc292 100644 --- a/api/core/app/apps/completion/app_generator.py +++ b/api/core/app/apps/completion/app_generator.py @@ -23,8 +23,8 @@ from core.app.entities.task_entities import AppBlockingResponse, AppStreamResponse from core.db.session_factory import session_factory from core.helper.trace_id_helper import extract_trace_session_id_from_args -from core.ops.ops_trace_manager import TraceQueueManager from extensions.ext_database import db +from extensions.tracing import create_trace_scope from factories import file_factory from graphon.model_runtime.errors.invoke import InvokeAuthorizationError from models import Account, App, Conversation, EndUser, Message @@ -151,9 +151,13 @@ def generate( annotation_reply=annotation_reply, ) - # get tracing instance - trace_manager = TraceQueueManager( - app_id=app_model.id, user_id=user.id if isinstance(user, Account) else user.session_id + extras = extract_trace_session_id_from_args(args) + trace_scope = create_trace_scope( + tenant_id=app_model.tenant_id, + app_id=app_model.id, + app_tracing=app_model.tracing, + user_id=user.id if isinstance(user, Account) else user.session_id, + extras=extras, ) # init application generate entity @@ -170,10 +174,8 @@ def generate( user_id=user.id, stream=streaming, invoke_from=invoke_from, - extras={ - **extract_trace_session_id_from_args(args), - }, - trace_manager=trace_manager, + extras=extras, + trace_scope=trace_scope, ) # init generate records diff --git a/api/core/app/apps/workflow/app_generator.py b/api/core/app/apps/workflow/app_generator.py index 4a34825bde6a6a..e5223b9de359f4 100644 --- a/api/core/app/apps/workflow/app_generator.py +++ b/api/core/app/apps/workflow/app_generator.py @@ -38,12 +38,12 @@ extract_parent_trace_context_from_args, extract_trace_session_id_from_args, ) -from core.ops.ops_trace_manager import TraceQueueManager from core.repositories import DifyCoreRepositoryFactory from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository from core.trigger.constants import is_trigger_node_type from core.workflow.node_factory import get_default_root_node_id from extensions.ext_database import db +from extensions.tracing import create_trace_scope from factories import file_factory from graphon.engine.filter import ResponseStreamFilter from graphon.engine.layer import Layer @@ -197,12 +197,6 @@ def generate( workflow=workflow, ) - # get tracing instance - trace_manager = TraceQueueManager( - app_id=app_model.id, - user_id=user.id if isinstance(user, Account) else user.session_id, - ) - inputs: Mapping[str, Any] = args["inputs"] extras = { @@ -210,6 +204,13 @@ def generate( **extract_parent_trace_context_from_args(args), **extract_trace_session_id_from_args(args), } + trace_scope = create_trace_scope( + tenant_id=app_model.tenant_id, + app_id=app_model.id, + app_tracing=app_model.tracing, + user_id=user.id if isinstance(user, Account) else user.session_id, + extras=extras, + ) workflow_run_id = str(workflow_run_id or uuid.uuid4()) root_node_id = root_node_id or get_default_root_node_id(workflow.graph_dict) root_node_config = workflow.get_node_config_by_id(root_node_id) @@ -235,7 +236,7 @@ def generate( invoke_from=invoke_from, allow_human_input=allow_human_input, call_depth=call_depth, - trace_manager=trace_manager, + trace_scope=trace_scope, workflow_execution_id=workflow_run_id, extras=extras, ) @@ -303,15 +304,18 @@ def resume( """ Resume a paused workflow execution using the persisted runtime state. - ``trace_manager`` is transient and excluded from generate-entity serialization, - so resumed executions rebuild it here before persistence layers receive the entity. + ``trace_scope`` is transient and excluded from generate-entity serialization, + so resumed executions create it again before persistence layers receive the entity. """ - if application_generate_entity.trace_manager is None: + if application_generate_entity.trace_scope is None: application_generate_entity = application_generate_entity.model_copy( update={ - "trace_manager": TraceQueueManager( + "trace_scope": create_trace_scope( + tenant_id=app_model.tenant_id, app_id=app_model.id, + app_tracing=app_model.tracing, user_id=user.id if isinstance(user, Account) else user.session_id, + extras=application_generate_entity.extras, ) } ) diff --git a/api/core/app/apps/workflow/generate_task_pipeline.py b/api/core/app/apps/workflow/generate_task_pipeline.py index 1ea0fb1127700e..ef8df9d4283ae5 100644 --- a/api/core/app/apps/workflow/generate_task_pipeline.py +++ b/api/core/app/apps/workflow/generate_task_pipeline.py @@ -57,7 +57,6 @@ ) from core.app.task_pipeline.based_generate_task_pipeline import BasedGenerateTaskPipeline from core.base.tts import AppGeneratorTTSPublisher -from core.ops.ops_trace_manager import TraceQueueManager from core.workflow.system_variables import build_system_variables from extensions.ext_database import db from graphon.entities import WorkflowStartReason @@ -127,7 +126,7 @@ def process( Process generate task pipeline. :return: """ - generator = self._wrapper_process_stream_response(trace_manager=self._application_generate_entity.trace_manager) + generator = self._wrapper_process_stream_response() if self._base_task_pipeline.stream: return self._to_stream_response(generator) else: @@ -247,9 +246,7 @@ def _listen_audio_msg(self, publisher: AppGeneratorTTSPublisher | None, task_id: return None raise RuntimeError(f"TTS publisher returned an unknown status: {audio_msg.status}") - def _wrapper_process_stream_response( - self, trace_manager: TraceQueueManager | None = None - ) -> Generator[StreamResponse]: + def _wrapper_process_stream_response(self) -> Generator[StreamResponse]: tts_publisher = None task_id = self._application_generate_entity.task_id tenant_id = self._application_generate_entity.app_config.tenant_id @@ -269,7 +266,7 @@ def _wrapper_process_stream_response( ) try: - for response in self._process_stream_response(tts_publisher=tts_publisher, trace_manager=trace_manager): + for response in self._process_stream_response(tts_publisher=tts_publisher): while audio_response := self._listen_audio_msg(publisher=tts_publisher, task_id=task_id): yield audio_response if tts_publisher and isinstance(response, ErrorStreamResponse): @@ -468,12 +465,9 @@ def _handle_loop_completed_event(self, event: QueueLoopCompletedEvent, **kwargs) def _handle_workflow_succeeded_event( self, event: QueueWorkflowSucceededEvent, - *, - trace_manager: TraceQueueManager | None = None, **kwargs, ) -> Generator[StreamResponse]: """Handle workflow succeeded events.""" - _ = trace_manager self._ensure_workflow_initialized() validated_state = self._ensure_graph_runtime_initialized() workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response( @@ -488,12 +482,9 @@ def _handle_workflow_succeeded_event( def _handle_workflow_partial_success_event( self, event: QueueWorkflowPartialSuccessEvent, - *, - trace_manager: TraceQueueManager | None = None, **kwargs, ) -> Generator[StreamResponse]: """Handle workflow partial success events.""" - _ = trace_manager self._ensure_workflow_initialized() validated_state = self._ensure_graph_runtime_initialized() workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response( @@ -523,12 +514,9 @@ def _handle_workflow_paused_event( def _handle_workflow_failed_and_stop_events( self, event: QueueWorkflowFailedEvent | QueueStopEvent, - *, - trace_manager: TraceQueueManager | None = None, **kwargs, ) -> Generator[StreamResponse]: """Handle workflow failed and stop events.""" - _ = trace_manager self._ensure_workflow_initialized() validated_state = self._ensure_graph_runtime_initialized() @@ -641,7 +629,6 @@ def _dispatch_event( event: AppQueueEvent, *, tts_publisher: AppGeneratorTTSPublisher | None = None, - trace_manager: TraceQueueManager | None = None, queue_message: WorkflowQueueMessage | MessageQueueMessage | None = None, ) -> Generator[StreamResponse]: """Dispatch events using elegant pattern matching.""" @@ -653,7 +640,6 @@ def _dispatch_event( yield from handler( event, tts_publisher=tts_publisher, - trace_manager=trace_manager, queue_message=queue_message, ) return @@ -669,7 +655,6 @@ def _dispatch_event( yield from self._handle_node_failed_events( event, tts_publisher=tts_publisher, - trace_manager=trace_manager, queue_message=queue_message, ) return @@ -679,7 +664,6 @@ def _dispatch_event( yield from self._handle_workflow_failed_and_stop_events( event, tts_publisher=tts_publisher, - trace_manager=trace_manager, queue_message=queue_message, ) return @@ -690,7 +674,6 @@ def _dispatch_event( def _process_stream_response( self, tts_publisher: AppGeneratorTTSPublisher | None = None, - trace_manager: TraceQueueManager | None = None, ) -> Generator[StreamResponse]: """ Process stream response using elegant Fluent Python patterns. @@ -730,7 +713,6 @@ def _process_stream_response( self._dispatch_event( event, tts_publisher=tts_publisher, - trace_manager=trace_manager, queue_message=queue_message, ) ): diff --git a/api/core/app/entities/app_invoke_entities.py b/api/core/app/entities/app_invoke_entities.py index 6f0baa4a586401..c2214eb088aabf 100644 --- a/api/core/app/entities/app_invoke_entities.py +++ b/api/core/app/entities/app_invoke_entities.py @@ -1,7 +1,8 @@ from collections.abc import Mapping, Sequence from enum import StrEnum -from typing import TYPE_CHECKING, Any, Literal +from typing import Any, Literal +from dify_trace_sdk.emitter import TraceScope from pydantic import BaseModel, ConfigDict, Field, ValidationInfo, field_validator from constants import UUID_NIL @@ -21,10 +22,6 @@ from models.model import AppMode from repositories.entities.workflow_execution import WorkflowExecutionTiming -if TYPE_CHECKING: - from core.ops.ops_trace_manager import TraceQueueManager - - DIFY_RUN_CONTEXT_KEY = "_dify" @@ -204,8 +201,8 @@ class AppGenerateEntity(BaseModel): # extra parameters, like: auto_generate_conversation_name extras: dict[str, Any] = Field(default_factory=dict) - # tracing instance - trace_manager: "TraceQueueManager | None" = Field(default=None, exclude=True, repr=False) + # Process-local: a resumed execution creates its scope again. + trace_scope: TraceScope | None = Field(default=None, exclude=True, repr=False) class EasyUIBasedAppGenerateEntity(AppGenerateEntity): @@ -387,17 +384,3 @@ class RagPipelineGenerateEntity(WorkflowAppGenerateEntity): document_id: str | None = None original_document_id: str | None = None start_node_id: str | None = None - - -# Resolve circular model references after model definitions. -from core.ops.ops_trace_manager import TraceQueueManager # noqa: E402 - -AppGenerateEntity.model_rebuild() -EasyUIBasedAppGenerateEntity.model_rebuild() -ConversationAppGenerateEntity.model_rebuild() -ChatAppGenerateEntity.model_rebuild() -CompletionAppGenerateEntity.model_rebuild() -AgentChatAppGenerateEntity.model_rebuild() -AdvancedChatAppGenerateEntity.model_rebuild() -WorkflowAppGenerateEntity.model_rebuild() -RagPipelineGenerateEntity.model_rebuild() 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 aa490fbf5462fa..ad27ea51eb8313 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 @@ -5,6 +5,7 @@ from threading import Thread from typing import Any, cast +from dify_trace_sdk.events import MessageCompleted from sqlalchemy import select from sqlalchemy.orm import Session @@ -49,8 +50,6 @@ from core.base.tts import AppGeneratorTTSPublisher from core.db.session_factory import session_factory from core.model_manager import ModelInstance -from core.ops.entities.trace_entity import TraceTaskName -from core.ops.ops_trace_manager import TraceQueueManager, TraceTask from core.prompt.utils.prompt_message_util import PromptMessageUtil from core.prompt.utils.prompt_template_parser import PromptTemplateParser from events.message_event import message_was_created @@ -130,7 +129,7 @@ def process( message_id=self._message_id, ) - generator = self._wrapper_process_stream_response(trace_manager=self._application_generate_entity.trace_manager) + generator = self._wrapper_process_stream_response() if self.stream: return self._to_stream_response(generator) else: @@ -218,9 +217,7 @@ def _listen_audio_msg(self, publisher: AppGeneratorTTSPublisher | None, task_id: return None raise RuntimeError(f"TTS publisher returned an unknown status: {audio_msg.status}") - def _wrapper_process_stream_response( - self, trace_manager: TraceQueueManager | None = None - ) -> Generator[StreamResponse]: + def _wrapper_process_stream_response(self) -> Generator[StreamResponse]: tenant_id = self._application_generate_entity.app_config.tenant_id task_id = self._application_generate_entity.task_id publisher = None @@ -238,7 +235,7 @@ def _wrapper_process_stream_response( get_credit_usage_app_type(self._app_config.app_mode), ) try: - for response in self._process_stream_response(publisher=publisher, trace_manager=trace_manager): + for response in self._process_stream_response(publisher=publisher): while audio_response := self._listen_audio_msg(publisher, task_id): yield audio_response if publisher and isinstance(response, ErrorStreamResponse): @@ -271,9 +268,7 @@ def _wrapper_process_stream_response( if publisher: publisher.cancel() - def _process_stream_response( - self, publisher: AppGeneratorTTSPublisher | None, trace_manager: TraceQueueManager | None = None - ) -> Generator[StreamResponse]: + def _process_stream_response(self, publisher: AppGeneratorTTSPublisher | None) -> Generator[StreamResponse]: """ Process stream response. :return: @@ -288,16 +283,7 @@ def _process_stream_response( with session_factory.create_session() as session: err = self.handle_error(event=event, session=session, message_id=self._message_id) session.commit() - - if trace_manager: - trace_manager.add_trace_task( - TraceTask( - TraceTaskName.MESSAGE_TRACE, - conversation_id=self._conversation_id, - message_id=self._message_id, - trace_session_id=self._application_generate_entity.extras.get("trace_session_id"), - ) - ) + self._emit_message_trace() yield self.error_to_stream_response(err) break @@ -323,15 +309,11 @@ def _process_stream_response( # A stopped Agent run may persist provider-reported usage after # cancellation completes. Do not replace it with local token estimates. if isinstance(event, QueueStopEvent): - saved_message = self._save_message( - session=session, - trace_manager=trace_manager, - preserve_existing_usage=True, - ) + saved_message = self._save_message(session=session, preserve_existing_usage=True) if self._app_config.app_mode == AppMode.AGENT: stopped_agent_message = saved_message else: - self._save_message(session=session, trace_manager=trace_manager) + self._save_message(session=session) session.commit() if stopped_agent_message is not None: # Provider/billing subscribers may perform external I/O. @@ -340,6 +322,9 @@ def _process_stream_response( stopped_agent_message, application_generate_entity=self._application_generate_entity, ) + else: + # The Agent worker reports a stopped Agent run's terminal itself. + self._emit_message_trace() message_end_resp = self._message_end_to_stream_response() yield message_end_resp case QueueRetrieverResourcesEvent(): @@ -436,7 +421,6 @@ def _save_message( self, *, session: Session, - trace_manager: TraceQueueManager | None = None, preserve_existing_usage: bool = False, ) -> Message: """ @@ -500,16 +484,6 @@ def _save_message( metadata["agent_stop_requested"] = True message.message_metadata = json.dumps(metadata, ensure_ascii=False) - if trace_manager and not agent_stop_requested: - trace_manager.add_trace_task( - TraceTask( - TraceTaskName.MESSAGE_TRACE, - conversation_id=self._conversation_id, - message_id=self._message_id, - trace_session_id=self._application_generate_entity.extras.get("trace_session_id"), - ) - ) - if not agent_stop_requested: message_was_created.send( message, @@ -517,6 +491,12 @@ def _save_message( ) return message + def _emit_message_trace(self) -> None: + """Report the finished message once its records are committed.""" + trace_scope = self._application_generate_entity.trace_scope + if trace_scope is not None: + trace_scope.emit(MessageCompleted(message_id=self._message_id, conversation_id=self._conversation_id)) + def _handle_stop(self, event: QueueStopEvent): """ Handle stop. diff --git a/api/core/app/task_pipeline/message_cycle_manager.py b/api/core/app/task_pipeline/message_cycle_manager.py index b41563a7de4b4e..cd8a496410f26d 100644 --- a/api/core/app/task_pipeline/message_cycle_manager.py +++ b/api/core/app/task_pipeline/message_cycle_manager.py @@ -145,6 +145,8 @@ def _generate_conversation_name_worker( conversation_id, conversation.app_id, message_id=message_id, + # The request's scope keeps the caller's trace ID with the message. + trace_scope=self._application_generate_entity.trace_scope, ) redis_client.setex(cache_key, 3600, name) except Exception: diff --git a/api/core/callback_handler/agent_tool_callback_handler.py b/api/core/callback_handler/agent_tool_callback_handler.py index 79b3b86980305f..9dc705b32534f3 100644 --- a/api/core/callback_handler/agent_tool_callback_handler.py +++ b/api/core/callback_handler/agent_tool_callback_handler.py @@ -4,8 +4,6 @@ from pydantic import BaseModel from configs import dify_config -from core.ops.entities.trace_entity import TraceTaskName -from core.ops.ops_trace_manager import TraceQueueManager, TraceTask from core.tools.entities.tool_entities import ToolInvokeMessage _TEXT_COLOR_MAPPING = { @@ -58,9 +56,6 @@ def on_tool_end( tool_name: str, tool_inputs: Mapping[str, Any], tool_outputs: Iterable[ToolInvokeMessage] | str, - message_id: str | None = None, - timer: Any | None = None, - trace_manager: TraceQueueManager | None = None, ): """If not the final action, print out observation.""" if dify_config.DEBUG: @@ -70,18 +65,6 @@ def on_tool_end( print_text("Outputs: " + str(tool_outputs)[:1000] + "\n", color=self.color) print_text("\n") - if trace_manager: - trace_manager.add_trace_task( - TraceTask( - TraceTaskName.TOOL_TRACE, - message_id=message_id, - tool_name=tool_name, - tool_inputs=tool_inputs, - tool_outputs=tool_outputs, - timer=timer, - ) - ) - def on_tool_error(self, error: Exception | KeyboardInterrupt, **kwargs: Any): """Do nothing.""" if dify_config.DEBUG: diff --git a/api/core/callback_handler/workflow_tool_callback_handler.py b/api/core/callback_handler/workflow_tool_callback_handler.py index f3c33706e01c7e..4edcf908ca678f 100644 --- a/api/core/callback_handler/workflow_tool_callback_handler.py +++ b/api/core/callback_handler/workflow_tool_callback_handler.py @@ -3,7 +3,6 @@ from configs import dify_config from core.callback_handler.agent_tool_callback_handler import DifyAgentCallbackHandler, print_text -from core.ops.ops_trace_manager import TraceQueueManager from core.tools.entities.tool_entities import ToolInvokeMessage @@ -15,9 +14,6 @@ def on_tool_execution( tool_name: str, tool_inputs: Mapping[str, Any], tool_outputs: Iterable[ToolInvokeMessage], - message_id: str | None = None, - timer: Any | None = None, - trace_manager: TraceQueueManager | None = None, ) -> Generator[ToolInvokeMessage]: for tool_output in tool_outputs: if dify_config.DEBUG: diff --git a/api/core/llm_generator/llm_generator.py b/api/core/llm_generator/llm_generator.py index dcec1b65b541b0..790cd9f5189391 100644 --- a/api/core/llm_generator/llm_generator.py +++ b/api/core/llm_generator/llm_generator.py @@ -5,6 +5,8 @@ from typing import Any, Literal, Protocol, TypedDict, cast import json_repair +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import ConversationNamed from sqlalchemy import select from sqlalchemy.orm import Session @@ -24,9 +26,6 @@ ) from core.model_context import with_credit_usage_created_by from core.model_manager import ModelInstance, ModelManager -from core.ops.entities.trace_entity import TraceTaskName -from core.ops.ops_trace_manager import TraceQueueManager, TraceTask -from core.ops.utils import measure_time from core.prompt.utils.prompt_template_parser import PromptTemplateParser from core.telemetry import PromptGenerationEvent, TelemetryContext from core.telemetry import emit as telemetry_emit @@ -37,6 +36,7 @@ from graphon.model_runtime.entities.message_entities import PromptMessage, SystemPromptMessage, UserPromptMessage from graphon.model_runtime.entities.model_entities import ModelType from graphon.model_runtime.errors.invoke import InvokeError +from libs.datetime_utils import Timing, measure_time from models import App, Message, WorkflowNodeExecutionModel from models.workflow import Workflow @@ -141,7 +141,7 @@ def _emit_prompt_generation( generated_output: str, model_provider: str, model_name: str, - timer: dict, + timer: Timing, result: LLMResult | None = None, error: str | None = None, ) -> None: @@ -165,11 +165,12 @@ def _emit_prompt_generation( "prompt_tokens": usage.prompt_tokens if usage else 0, "completion_tokens": usage.completion_tokens if usage else 0, "total_tokens": usage.total_tokens if usage else 0, - "latency": (timer["end"] - timer["start"]).total_seconds() if timer.get("end") else 0.0, + "latency": (timer.ended_at - timer.started_at).total_seconds(), "total_price": float(usage.total_price) if usage and usage.total_price else None, "currency": usage.currency if usage else None, - "timer": timer, "error": error, + "started_at": timer.started_at, + "ended_at": timer.ended_at, }, ) ) @@ -187,6 +188,7 @@ def generate_conversation_name( message_id: str | None = None, *, model_instance: ModelInstance | None = None, + trace_scope: TraceScope | None, ) -> str: """Generate with a resolved model, or resolve it for legacy background callers. @@ -234,19 +236,17 @@ def generate_conversation_name( if len(name) > 75: name = name[:75] + "..." - # get tracing instance - trace_manager = TraceQueueManager(app_id=app_id) - trace_manager.add_trace_task( - TraceTask( - TraceTaskName.GENERATE_NAME_TRACE, - conversation_id=conversation_id, - message_id=message_id, - generate_conversation_name=name, - inputs=prompt, - timer=timer, - tenant_id=tenant_id, + if trace_scope is not None: + trace_scope.emit( + ConversationNamed( + conversation_id=conversation_id, + message_id=message_id, + inputs=prompt, + name=name, + started_at=timer.started_at, + ended_at=timer.ended_at, + ) ) - ) return name diff --git a/api/core/moderation/input_moderation.py b/api/core/moderation/input_moderation.py index 9edcef3be38d1b..d85dc86bd34b14 100644 --- a/api/core/moderation/input_moderation.py +++ b/api/core/moderation/input_moderation.py @@ -2,14 +2,15 @@ from collections.abc import Mapping from typing import Any +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import ModerationChecked, json_safe + from core.app.app_config.entities import AppConfig from core.app.entities.app_invoke_entities import get_credit_usage_app_type from core.model_context import use_credit_usage_metadata from core.moderation.base import ModerationAction, ModerationError from core.moderation.factory import ModerationFactory -from core.ops.entities.trace_entity import TraceTaskName -from core.ops.ops_trace_manager import TraceQueueManager, TraceTask -from core.ops.utils import measure_time +from libs.datetime_utils import measure_time logger = logging.getLogger(__name__) @@ -23,7 +24,7 @@ def check( inputs: Mapping[str, Any], query: str, message_id: str, - trace_manager: TraceQueueManager | None = None, + trace_scope: TraceScope | None, ) -> tuple[bool, Mapping[str, Any], str]: """ Process sensitive_word_avoidance. @@ -33,7 +34,7 @@ def check( :param inputs: inputs :param query: query :param message_id: message id - :param trace_manager: trace manager + :param trace_scope: where the moderation verdict is reported, if anywhere :return: """ inputs = dict(inputs) @@ -51,14 +52,17 @@ def check( with measure_time() as timer: moderation_result = moderation_factory.moderation_for_inputs(inputs, query) - if trace_manager: - trace_manager.add_trace_task( - TraceTask( - TraceTaskName.MODERATION_TRACE, + if trace_scope is not None: + trace_scope.emit( + ModerationChecked( message_id=message_id, - moderation_result=moderation_result, - inputs=inputs, - timer=timer, + inputs={key: json_safe(value) for key, value in inputs.items()}, + query=moderation_result.query, + flagged=moderation_result.flagged, + action=moderation_result.action.value, + preset_response=moderation_result.preset_response, + started_at=timer.started_at, + ended_at=timer.ended_at, ) ) diff --git a/api/core/ops/base_trace_instance.py b/api/core/ops/base_trace_instance.py deleted file mode 100644 index bcd2a663584132..00000000000000 --- a/api/core/ops/base_trace_instance.py +++ /dev/null @@ -1,144 +0,0 @@ -from __future__ import annotations - -from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Literal, Protocol, cast - -from flask import Flask, current_app, has_app_context -from sqlalchemy import select -from sqlalchemy.orm import Session - -from configs import dify_config -from core.ops.entities.config_entity import BaseTracingConfig -from core.ops.entities.trace_entity import BaseTraceInfo -from core.ops.exceptions import RetryableTraceDispatchError -from extensions.ext_database import db -from models import Account, App - -if TYPE_CHECKING: - from dify_agent.layers.execution_context import ApplicationTracingContext - - from core.app.apps.agent_app.runtime_request_builder import AgentAppRuntimeBuildContext - from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, WorkflowAppGenerateEntity - from core.workflow.nodes.agent_v2.runtime_request_builder import WorkflowAgentRuntimeBuildContext - - -class ApplicationTracingHooks(Protocol): - """Execution-owned port for the application tracing lifecycle. - - The composition root supplies an app-scoped implementation. Core callers - carry trusted execution context and do not construct services or vendor SDKs. - """ - - def prepare_agent_app(self, context: AgentAppRuntimeBuildContext) -> ApplicationTracingContext | None: ... - - def prepare_workflow_agent( - self, context: WorkflowAgentRuntimeBuildContext, conversation_id: str | None - ) -> ApplicationTracingContext | None: ... - - def freeze_workflow_run( - self, - entity: AdvancedChatAppGenerateEntity | WorkflowAppGenerateEntity, - *, - workflow_id: str, - workflow_run_id: str, - conversation_id: object, - graph_data: object, - ) -> None: ... - - def queue_enabled_at_start(self, app_id: object) -> bool: ... - - def capture_terminal(self, info: BaseTraceInfo) -> bool: ... - - def capture_agent_app_terminal( - self, - *, - tenant_id: str, - app_id: str, - message_id: str, - outcome: Literal["succeeded", "failed", "cancelled"], - output: str | None, - ) -> None: ... - - -class ApplicationTracingHooksNotRegisteredError(RetryableTraceDispatchError): - """Enabled tracing must not silently fall back to a different export path.""" - - def __init__(self) -> None: - super().__init__("application_tracing_hooks_not_registered") - - -_APPLICATION_TRACING_HOOKS_KEY = "application_tracing_hooks" - - -def bind_application_tracing_hooks(app: Flask, hooks: ApplicationTracingHooks) -> None: - """Bind the execution port to one Flask application, never a process singleton.""" - app.extensions[_APPLICATION_TRACING_HOOKS_KEY] = hooks - - -def application_tracing_hooks() -> ApplicationTracingHooks | None: - """Resolve the current app's port; a disabled deployment needs no Flask context.""" - if not dify_config.OPS_AGENT_TRACE_ENABLED: - return None - if not has_app_context(): - raise ApplicationTracingHooksNotRegisteredError() - hooks = current_app.extensions.get(_APPLICATION_TRACING_HOOKS_KEY) - if hooks is None: - raise ApplicationTracingHooksNotRegisteredError() - return cast(ApplicationTracingHooks, hooks) - - -class BaseTraceInstance(ABC): - """ - Base trace instance for ops trace services - """ - - def __init__(self, trace_config: BaseTracingConfig): - """ - Initializer for the trace instance. - Distribute trace tasks by matching entities - """ - self.trace_config = trace_config - - @abstractmethod - def trace(self, trace_info: BaseTraceInfo): - """ - Abstract method to trace activities. - Subclasses must implement specific tracing logic for activities. - """ - ... - - def get_service_account_with_tenant(self, app_id: str) -> Account: - """ - Get the app creator with the app's tenant and membership role. - - Trace tasks may run after the creator switches workspaces. Bind only the - returned account's context without changing the creator's current workspace. - - Args: - app_id: The ID of the app - - Returns: - Account: The service account with the app's tenant set up - - Raises: - ValueError: If the app, creator account, app tenant or membership cannot be found - """ - with Session(db.engine, expire_on_commit=False) as session: - # Get the app to find its creator - app_stmt = select(App).where(App.id == app_id) - app = session.scalar(app_stmt) - if not app: - raise ValueError(f"App with id {app_id} not found") - - if not app.created_by: - raise ValueError(f"App with id {app_id} has no creator (created_by is None)") - account_stmt = select(Account).where(Account.id == app.created_by) - service_account = session.scalar(account_stmt) - if not service_account: - raise ValueError(f"Creator account with id {app.created_by} not found for app {app_id}") - - service_account.set_tenant_id_with_session(app.tenant_id, session=session) - if service_account.current_tenant_id != app.tenant_id: - raise ValueError(f"App tenant {app.tenant_id} not found for account {service_account.id}") - - return service_account diff --git a/api/core/ops/entities/config_entity.py b/api/core/ops/entities/config_entity.py deleted file mode 100644 index b3df7993eedbab..00000000000000 --- a/api/core/ops/entities/config_entity.py +++ /dev/null @@ -1,64 +0,0 @@ -from enum import StrEnum - -from pydantic import BaseModel - -from core.ops.utils import validate_project_name, validate_url - - -class TracingProviderEnum(StrEnum): - ARIZE = "arize" - PHOENIX = "phoenix" - LANGFUSE = "langfuse" - LANGSMITH = "langsmith" - OPIK = "opik" - WEAVE = "weave" - ALIYUN = "aliyun" - MLFLOW = "mlflow" - DATABRICKS = "databricks" - TENCENT = "tencent" - - -class BaseTracingConfig(BaseModel): - """ - Base model class for tracing configurations - """ - - @classmethod - def validate_endpoint_url(cls, v: str, default_url: str) -> str: - """ - Common endpoint URL validation logic - - Args: - v: URL value to validate - default_url: Default URL to use if input is None or empty - - Returns: - Validated and normalized URL - """ - return validate_url(v, default_url) - - @classmethod - def validate_project_field(cls, v: str, default_name: str) -> str: - """ - Common project name validation logic - - Args: - v: Project name to validate - default_name: Default name to use if input is None or empty - - Returns: - Validated project name - """ - return validate_project_name(v, default_name) - - -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 deleted file mode 100644 index e9685a208eec26..00000000000000 --- a/api/core/ops/entities/trace_entity.py +++ /dev/null @@ -1,292 +0,0 @@ -from collections.abc import Mapping -from datetime import datetime -from enum import StrEnum -from typing import Any - -from pydantic import BaseModel, ConfigDict, field_serializer, field_validator - -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: str | dict[str, Any] | list[Any] | None = None - outputs: str | dict[str, Any] | list[Any] | None = None - start_time: datetime | None = None - end_time: datetime | None = None - metadata: dict[str, Any] - trace_id: str | None = None - - @field_validator("inputs", "outputs") - @classmethod - def ensure_type(cls, v: str | dict[str, Any] | list[Any] | None) -> str | dict[str, Any] | list[Any] | None: - if v is None: - return None - if isinstance(v, str | dict | list): - return v - return "" - - model_config = ConfigDict(protected_namespaces=()) - - @property - def resolved_trace_id(self) -> str | None: - """Get trace_id with intelligent fallback. - - Priority: - 1. External trace_id (from X-Trace-Id header) - 2. workflow_run_id (if this trace type has it) - 3. message_id (as final fallback) - """ - if self.trace_id: - return self.trace_id - - # Try workflow_run_id (only exists on workflow-related traces) - workflow_run_id = getattr(self, "workflow_run_id", None) - if workflow_run_id: - return workflow_run_id - - # Final fallback to message_id - return str(self.message_id) if self.message_id else None - - @property - def resolved_parent_context(self) -> tuple[str | None, str | None]: - """Resolve cross-workflow parent linking from metadata. - - Extracts typed parent IDs from the ``parent_trace_context`` metadata - payload (set by tool_node when invoking nested workflows). - - Returns: - (trace_correlation_override, parent_span_id_source) where - trace_correlation_override is the outer workflow_run_id and - parent_span_id_source is the outer node_execution_id. - """ - parent_ctx = self.metadata.get("parent_trace_context") - match parent_ctx: - case ParentTraceContext(): - context = parent_ctx - case Mapping(): - try: - context = ParentTraceContext.model_validate(parent_ctx) - except ValueError: - return None, None - case _: - return None, None - return ( - context.parent_workflow_run_id, - context.parent_node_execution_id, - ) - - @field_serializer("start_time", "end_time") - def serialize_datetime(self, dt: datetime | None) -> str | None: - if dt is None: - return None - return dt.isoformat() - - -class WorkflowTraceInfo(BaseTraceInfo): - workflow_data: Any = None - conversation_id: str | None = None - workflow_app_log_id: str | None = None - workflow_id: str - tenant_id: str - workflow_run_id: str - workflow_run_elapsed_time: int | float - workflow_run_status: str - workflow_run_inputs: Mapping[str, Any] - workflow_run_outputs: Mapping[str, Any] - workflow_run_version: str - error: str | None = None - total_tokens: int - prompt_tokens: int | None = None - completion_tokens: int | None = None - file_list: list[str] - invoked_by: str | None = None - query: str - metadata: dict[str, Any] - - -class MessageTraceInfo(BaseTraceInfo): - conversation_model: str - message_tokens: int - answer_tokens: int - total_tokens: int - error: str | None = None - file_list: str | dict[str, Any] | list[Any] | None = None - message_file_data: Any | None = None - conversation_mode: str - gen_ai_server_time_to_first_token: float | None = None - llm_streaming_time_to_generate: float | None = None - is_streaming_request: bool = False - - -class ModerationTraceInfo(BaseTraceInfo): - flagged: bool - action: str - preset_response: str - query: str - - -class SuggestedQuestionTraceInfo(BaseTraceInfo): - total_tokens: int - status: str | None = None - error: str | None = None - from_account_id: str | None = None - agent_based: bool | None = None - from_source: str | None = None - model_provider: str | None = None - model_id: str | None = None - suggested_question: list[str] - level: str - status_message: str | None = None - workflow_run_id: str | None = None - - model_config = ConfigDict(protected_namespaces=()) - - -class DatasetRetrievalTraceInfo(BaseTraceInfo): - documents: Any = None - error: str | None = None - - -class ToolTraceInfo(BaseTraceInfo): - tool_name: str - tool_inputs: dict[str, Any] - tool_outputs: str - metadata: dict[str, Any] - message_file_data: Any = None - error: str | None = None - tool_config: dict[str, Any] - time_cost: int | float - tool_parameters: dict[str, Any] - file_url: str | None | list[str] = None - - -class GenerateNameTraceInfo(BaseTraceInfo): - conversation_id: str | None = None - tenant_id: str - - -class PromptGenerationTraceInfo(BaseTraceInfo): - """Trace information for prompt generation operations (rule-generate, code-generate, etc.).""" - - tenant_id: str - user_id: str - app_id: str | None = None - - operation_type: str - instruction: str - - prompt_tokens: int - completion_tokens: int - total_tokens: int - - model_provider: str - model_name: str - - latency: float - - total_price: float | None = None - currency: str | None = None - - error: str | None = None - - model_config = ConfigDict(protected_namespaces=()) - - -class WorkflowNodeTraceInfo(BaseTraceInfo): - workflow_id: str - workflow_run_id: str - tenant_id: str - node_execution_id: str - node_id: str - node_type: str - title: str - - status: str - error: str | None = None - elapsed_time: float - - index: int - predecessor_node_id: str | None = None - - total_tokens: int = 0 - total_price: float = 0.0 - currency: str | None = None - - model_provider: str | None = None - model_name: str | None = None - prompt_tokens: int | None = None - completion_tokens: int | None = None - - tool_name: str | None = None - - iteration_id: str | None = None - iteration_index: int | None = None - loop_id: str | None = None - loop_index: int | None = None - parallel_id: str | None = None - - node_inputs: Mapping[str, Any] | None = None - node_outputs: Mapping[str, Any] | None = None - process_data: Mapping[str, Any] | None = None - - invoked_by: str | None = None - - model_config = ConfigDict(protected_namespaces=()) - - -class DraftNodeExecutionTrace(WorkflowNodeTraceInfo): - pass - - -class TaskData(BaseModel): - app_id: str - trace_info_type: str - trace_info: Any = None - - -trace_info_info_map = { - "WorkflowTraceInfo": WorkflowTraceInfo, - "MessageTraceInfo": MessageTraceInfo, - "ModerationTraceInfo": ModerationTraceInfo, - "SuggestedQuestionTraceInfo": SuggestedQuestionTraceInfo, - "DatasetRetrievalTraceInfo": DatasetRetrievalTraceInfo, - "ToolTraceInfo": ToolTraceInfo, - "GenerateNameTraceInfo": GenerateNameTraceInfo, - "PromptGenerationTraceInfo": PromptGenerationTraceInfo, - "WorkflowNodeTraceInfo": WorkflowNodeTraceInfo, - "DraftNodeExecutionTrace": DraftNodeExecutionTrace, -} - - -class OperationType(StrEnum): - """Operation type for token metric labels. - - Used as a metric attribute on ``dify.tokens.input`` / ``dify.tokens.output`` - counters so consumers can break down token usage by operation. - """ - - WORKFLOW = "workflow" - NODE_EXECUTION = "node_execution" - MESSAGE = "message" - RULE_GENERATE = "rule_generate" - CODE_GENERATE = "code_generate" - STRUCTURED_OUTPUT = "structured_output" - INSTRUCTION_MODIFY = "instruction_modify" - - -class TraceTaskName(StrEnum): - CONVERSATION_TRACE = "conversation" - DRAFT_NODE_EXECUTION_TRACE = "draft_node_execution" - WORKFLOW_TRACE = "workflow" - MESSAGE_TRACE = "message" - MODERATION_TRACE = "moderation" - SUGGESTED_QUESTION_TRACE = "suggested_question" - DATASET_RETRIEVAL_TRACE = "dataset_retrieval" - TOOL_TRACE = "tool" - GENERATE_NAME_TRACE = "generate_conversation_name" - PROMPT_GENERATION_TRACE = "prompt_generation" - NODE_EXECUTION_TRACE = "node_execution" - DATASOURCE_TRACE = "datasource" diff --git a/api/core/ops/exceptions.py b/api/core/ops/exceptions.py deleted file mode 100644 index e80eb3b9649841..00000000000000 --- a/api/core/ops/exceptions.py +++ /dev/null @@ -1,37 +0,0 @@ -"""Core exceptions shared by ops trace dispatchers and trace providers. - -Provider packages may raise these types to request generic task behavior, but -generic Celery tasks should not import provider-specific exception classes. -""" - - -class TraceProviderNotInstalledError(ImportError): - def __init__(self, tracing_provider: str, module_name: str) -> None: - super().__init__(f"Tracing provider {tracing_provider} requires the missing module {module_name}.") - - -class RetryableTraceDispatchError(RuntimeError): - """Base class for transient trace dispatch failures that Celery may retry.""" - - acceptance_unknown: bool = False - - -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.""" - - parent_node_execution_id: str - - def __init__(self, parent_node_execution_id: str) -> None: - self.parent_node_execution_id = parent_node_execution_id - super().__init__( - "Pending trace parent context for parent_node_execution_id=" - f"{parent_node_execution_id}. Retry after the parent span context is published." - ) diff --git a/api/core/ops/ops_trace_manager.py b/api/core/ops/ops_trace_manager.py deleted file mode 100644 index 2814b11e3572cf..00000000000000 --- a/api/core/ops/ops_trace_manager.py +++ /dev/null @@ -1,1589 +0,0 @@ -import collections -import importlib.util -import json -import logging -import os -import queue -import threading -import time -from collections.abc import Mapping -from datetime import timedelta -from typing import TYPE_CHECKING, Any, TypedDict, override -from uuid import UUID, uuid4 - -from cachetools import LRUCache -from flask import current_app -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.base_trace_instance import application_tracing_hooks -from core.ops.entities.config_entity import ( - BaseTracingConfig, - TracingProviderEnum, - ops_trace_payload_path, -) -from core.ops.entities.trace_entity import ( - BaseTraceInfo, - DatasetRetrievalTraceInfo, - DraftNodeExecutionTrace, - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - PromptGenerationTraceInfo, - SuggestedQuestionTraceInfo, - TaskData, - ToolTraceInfo, - TraceTaskName, - WorkflowNodeTraceInfo, - WorkflowTraceInfo, -) -from core.ops.exceptions import TraceProviderNotInstalledError -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 -from models.account import Tenant -from models.dataset import Dataset -from models.model import App, AppModelConfig, Conversation, Message, MessageFile, TraceAppConfig -from models.provider import Provider, ProviderCredential, ProviderModel, ProviderModelCredential, ProviderType -from models.tools import ApiToolProvider, BuiltinToolProvider, MCPToolProvider, WorkflowToolProvider -from models.workflow import WorkflowAppLog -from tasks.ops_trace_task import process_trace_tasks - -if TYPE_CHECKING: - from graphon.entities import WorkflowExecution - -logger = logging.getLogger(__name__) - - -def _dump_parent_trace_context(parent_trace_context: Any) -> dict[str, str] | None: - if isinstance(parent_trace_context, ParentTraceContext): - return parent_trace_context.model_dump(exclude_none=True) - if isinstance(parent_trace_context, dict): - try: - return ParentTraceContext.model_validate(parent_trace_context).model_dump(exclude_none=True) - except ValueError: - return None - return None - - -def _get_trace_session_id(kwargs: Mapping[str, Any]) -> str | None: - value = kwargs.get("trace_session_id") - return value if isinstance(value, str) and value else None - - -class _AppTracingConfig(TypedDict, total=False): - enabled: bool - tracing_provider: str | None - - -_app_tracing_config_adapter: TypeAdapter[_AppTracingConfig] = TypeAdapter(_AppTracingConfig) - - -def _lookup_app_and_workspace_names(app_id: str | None, tenant_id: str | None) -> tuple[str, str]: - """Return (app_name, workspace_name) for the given IDs. Falls back to empty strings.""" - app_name = "" - workspace_name = "" - if not app_id and not tenant_id: - return app_name, workspace_name - with Session(db.engine) as session: - if app_id: - name = session.scalar(select(App.name).where(App.id == app_id)) - if name: - app_name = name - if tenant_id: - name = session.scalar(select(Tenant.name).where(Tenant.id == tenant_id)) - if name: - workspace_name = name - return app_name, workspace_name - - -_PROVIDER_TYPE_TO_MODEL: dict[str, type] = { - "builtin": BuiltinToolProvider, - "plugin": BuiltinToolProvider, - "api": ApiToolProvider, - "workflow": WorkflowToolProvider, - "mcp": MCPToolProvider, -} - - -def _lookup_credential_name(credential_id: str | None, provider_type: str | None) -> str: - if not credential_id: - return "" - model_cls = _PROVIDER_TYPE_TO_MODEL.get(provider_type or "") - if not model_cls: - return "" - with Session(db.engine) as session: - name = session.scalar(select(model_cls.name).where(model_cls.id == credential_id)) # type: ignore[attr-defined] - return str(name) if name else "" - - -def _lookup_llm_credential_info( - tenant_id: str | None, provider: str | None, model: str | None, model_type: str | None = "llm" -) -> tuple[str | None, str]: - """ - Lookup LLM credential ID and name for the given provider and model. - Returns (credential_id, credential_name). - - Handles async timing issues gracefully - if credential is deleted between lookups, - returns the ID but empty name rather than failing. - """ - if not tenant_id or not provider: - return None, "" - - try: - with Session(db.engine) as session: - # Try to find provider-level or model-level configuration - provider_record = session.scalar( - select(Provider).where( - Provider.tenant_id == tenant_id, - Provider.provider_name == provider, - Provider.provider_type == ProviderType.CUSTOM, - ) - ) - - if not provider_record: - return None, "" - - # Check if there's a model-specific config - credential_id = None - credential_name = "" - is_model_level = False - - if model: - # Try model-level first - model_record = session.scalar( - select(ProviderModel).where( - ProviderModel.tenant_id == tenant_id, - ProviderModel.provider_name == provider, - ProviderModel.model_name == model, - ProviderModel.model_type == model_type, - ) - ) - - if model_record and model_record.credential_id: - credential_id = model_record.credential_id - is_model_level = True - - if not credential_id and provider_record.credential_id: - # Fall back to provider-level credential - credential_id = provider_record.credential_id - is_model_level = False - - # Lookup credential_name if we have credential_id - if credential_id: - try: - if is_model_level: - # Query ProviderModelCredential - cred_name = session.scalar( - select(ProviderModelCredential.credential_name).where( - ProviderModelCredential.id == credential_id - ) - ) - else: - # Query ProviderCredential - cred_name = session.scalar( - select(ProviderCredential.credential_name).where(ProviderCredential.id == credential_id) - ) - - if cred_name: - credential_name = str(cred_name) - except Exception as e: - # Credential might have been deleted between lookups (async timing) - # Return ID but empty name rather than failing - logger.warning( - "Failed to lookup credential name for credential_id=%s (provider=%s, model=%s): %s", - credential_id, - provider, - model, - str(e), - exc_info=True, - ) - - return credential_id, credential_name - except Exception as e: - # Database query failed or other unexpected error - # Return empty rather than propagating error to telemetry emission - logger.warning( - "Failed to lookup LLM credential info for tenant_id=%s, provider=%s, model=%s: %s", - tenant_id, - provider, - model, - str(e), - exc_info=True, - ) - return None, "" - - -class TracingProviderConfigEntry(TypedDict): - config_class: type[BaseTracingConfig] - secret_keys: list[str] - other_keys: list[str] - trace_instance: type[Any] - - -class OpsTraceProviderConfigMap(collections.UserDict[str, TracingProviderConfigEntry]): - @override - def __getitem__(self, key: str) -> TracingProviderConfigEntry: - try: - match key: - case TracingProviderEnum.LANGFUSE: - from dify_trace_langfuse.config import LangfuseConfig - from dify_trace_langfuse.langfuse_trace import LangFuseDataTrace - - return { - "config_class": LangfuseConfig, - "secret_keys": ["public_key", "secret_key"], - "other_keys": ["host", "project_key"], - "trace_instance": LangFuseDataTrace, - } - - case TracingProviderEnum.LANGSMITH: - from dify_trace_langsmith.config import LangSmithConfig - from dify_trace_langsmith.langsmith_trace import LangSmithDataTrace - - return { - "config_class": LangSmithConfig, - "secret_keys": ["api_key"], - "other_keys": ["project", "endpoint"], - "trace_instance": LangSmithDataTrace, - } - - case TracingProviderEnum.OPIK: - from dify_trace_opik.config import OpikConfig - from dify_trace_opik.opik_trace import OpikDataTrace - - return { - "config_class": OpikConfig, - "secret_keys": ["api_key"], - "other_keys": ["project", "url", "workspace"], - "trace_instance": OpikDataTrace, - } - - case TracingProviderEnum.WEAVE: - from dify_trace_weave.config import WeaveConfig - from dify_trace_weave.weave_trace import WeaveDataTrace - - return { - "config_class": WeaveConfig, - "secret_keys": ["api_key"], - "other_keys": ["project", "entity", "endpoint", "host"], - "trace_instance": WeaveDataTrace, - } - case TracingProviderEnum.ARIZE: - from dify_trace_arize_phoenix.arize_phoenix_trace import ArizePhoenixDataTrace - from dify_trace_arize_phoenix.config import ArizeConfig - - return { - "config_class": ArizeConfig, - "secret_keys": ["api_key", "space_id"], - "other_keys": ["project", "endpoint"], - "trace_instance": ArizePhoenixDataTrace, - } - case TracingProviderEnum.PHOENIX: - from dify_trace_arize_phoenix.arize_phoenix_trace import ArizePhoenixDataTrace - from dify_trace_arize_phoenix.config import PhoenixConfig - - return { - "config_class": PhoenixConfig, - "secret_keys": ["api_key"], - "other_keys": ["project", "endpoint"], - "trace_instance": ArizePhoenixDataTrace, - } - case TracingProviderEnum.ALIYUN: - from dify_trace_aliyun.aliyun_trace import AliyunDataTrace - from dify_trace_aliyun.config import AliyunConfig - - return { - "config_class": AliyunConfig, - "secret_keys": ["license_key"], - "other_keys": ["endpoint", "app_name"], - "trace_instance": AliyunDataTrace, - } - case TracingProviderEnum.MLFLOW: - from dify_trace_mlflow.config import MLflowConfig - from dify_trace_mlflow.mlflow_trace import MLflowDataTrace - - return { - "config_class": MLflowConfig, - "secret_keys": ["password"], - "other_keys": ["tracking_uri", "experiment_id", "username"], - "trace_instance": MLflowDataTrace, - } - case TracingProviderEnum.DATABRICKS: - from dify_trace_mlflow.config import DatabricksConfig - from dify_trace_mlflow.mlflow_trace import MLflowDataTrace - - return { - "config_class": DatabricksConfig, - "secret_keys": ["personal_access_token", "client_secret"], - "other_keys": ["host", "client_id", "experiment_id"], - "trace_instance": MLflowDataTrace, - } - - case TracingProviderEnum.TENCENT: - from dify_trace_tencent.config import TencentConfig - from dify_trace_tencent.tencent_trace import TencentDataTrace - - return { - "config_class": TencentConfig, - "secret_keys": ["token"], - "other_keys": ["endpoint", "service_name"], - "trace_instance": TencentDataTrace, - } - - case _: - raise KeyError(f"Unsupported tracing provider: {key}") - except ModuleNotFoundError as error: - if error.name is None or importlib.util.find_spec(error.name.partition(".")[0]) is not None: - raise - raise TraceProviderNotInstalledError(key, error.name) from error - - -provider_config_map = OpsTraceProviderConfigMap() - - -class OpsTraceManager: - ops_trace_instances_cache: LRUCache = LRUCache(maxsize=128) - 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 - ): - """ - Encrypt tracing config. - :param tenant_id: tenant id - :param tracing_provider: tracing provider - :param tracing_config: tracing config dictionary to be encrypted - :param current_trace_config: current tracing configuration for keeping existing values - :return: encrypted tracing configuration - """ - # Get the configuration class and the keys that require encryption - config_class, secret_keys, other_keys = ( - provider_config_map[tracing_provider]["config_class"], - provider_config_map[tracing_provider]["secret_keys"], - provider_config_map[tracing_provider]["other_keys"], - ) - - new_config: dict[str, Any] = {} - # Encrypt necessary keys - for key in secret_keys: - if key in tracing_config: - if tracing_config[key] is None: - # Preserve explicit clearing for optional credentials; the - # provider model still rejects null on a required secret. - new_config[key] = None - elif "*" in tracing_config[key]: - # If the key contains '*', retain the original value from the current config - if current_trace_config: - new_config[key] = current_trace_config.get(key, tracing_config[key]) - else: - new_config[key] = tracing_config[key] - else: - # Otherwise, encrypt the key - new_config[key] = encrypt_token(tenant_id, tracing_config[key]) - - for key in other_keys: - new_config[key] = tracing_config.get(key, "") - - # Create a new instance of the config class with the new configuration - encrypted_config = config_class(**new_config) - return encrypted_config.model_dump() - - @classmethod - def decrypt_tracing_config(cls, tenant_id: str, tracing_provider: str, tracing_config: dict[str, Any]): - """ - Decrypt tracing config - :param tenant_id: tenant id - :param tracing_provider: tracing provider - :param tracing_config: tracing config - :return: - """ - config_json = json.dumps(tracing_config, sort_keys=True) - decrypted_config_key = ( - tenant_id, - tracing_provider, - config_json, - ) - - # First check without lock for performance - cached_config = cls.decrypted_configs_cache.get(decrypted_config_key) - if cached_config is not None: - return dict(cached_config) - - with cls._decryption_cache_lock: - # Second check (double-checked locking) to prevent race conditions - cached_config = cls.decrypted_configs_cache.get(decrypted_config_key) - if cached_config is not None: - return dict(cached_config) - - config_class, secret_keys, other_keys = ( - provider_config_map[tracing_provider]["config_class"], - provider_config_map[tracing_provider]["secret_keys"], - provider_config_map[tracing_provider]["other_keys"], - ) - new_config: dict[str, Any] = { - key: None for key in secret_keys if key in tracing_config and tracing_config[key] is None - } - keys_to_decrypt = [key for key in secret_keys if tracing_config.get(key) is not None] - if keys_to_decrypt: - decrypted_values = batch_decrypt_token(tenant_id, [tracing_config[key] for key in keys_to_decrypt]) - new_config.update(zip(keys_to_decrypt, decrypted_values, strict=False)) - - for key in other_keys: - new_config[key] = tracing_config.get(key, "") - - decrypted_config = config_class(**new_config).model_dump() - cls.decrypted_configs_cache[decrypted_config_key] = decrypted_config - return dict(decrypted_config) - - @classmethod - def obfuscated_decrypt_token(cls, tracing_provider: str, decrypt_tracing_config: dict[str, Any]): - """ - Decrypt tracing config - :param tracing_provider: tracing provider - :param decrypt_tracing_config: tracing config - :return: - """ - config_class, secret_keys, other_keys = ( - provider_config_map[tracing_provider]["config_class"], - provider_config_map[tracing_provider]["secret_keys"], - provider_config_map[tracing_provider]["other_keys"], - ) - new_config: dict[str, Any] = {} - for key in secret_keys: - if key in decrypt_tracing_config: - new_config[key] = obfuscated_token(decrypt_tracing_config[key]) - - for key in other_keys: - new_config[key] = decrypt_tracing_config.get(key, "") - return config_class(**new_config).model_dump() - - @classmethod - def get_decrypted_tracing_config(cls, app_id: str, tracing_provider: str): - """ - Get decrypted tracing config - :param app_id: app id - :param tracing_provider: tracing provider - :return: - """ - trace_config_data: TraceAppConfig | None = db.session.scalar( - select(TraceAppConfig) - .where(TraceAppConfig.app_id == app_id, TraceAppConfig.tracing_provider == tracing_provider) - .limit(1) - ) - - if not trace_config_data: - return None - # decrypt_token - stmt = select(App).where(App.id == app_id) - app = db.session.scalar(stmt) - if not app: - raise ValueError("App not found") - - tenant_id = app.tenant_id - if trace_config_data.tracing_config is None: - raise ValueError("Tracing config cannot be None.") - decrypt_tracing_config = cls.decrypt_tracing_config( - tenant_id, tracing_provider, trace_config_data.tracing_config - ) - - return decrypt_tracing_config - - @classmethod - def get_ops_trace_instance( - cls, - app_id: UUID | str | None = None, - ): - """ - Get ops trace through model config - :param app_id: app_id - :return: - """ - if isinstance(app_id, UUID): - app_id = str(app_id) - - if app_id is None: - return None - - # Handle storage_id format (tenant-{uuid}) - not a real app_id - if isinstance(app_id, str) and app_id.startswith("tenant-"): - return None - - app = db.session.get(App, app_id) - - if app is None: - return None - - app_ops_trace_config = _app_tracing_config_adapter.validate_json(app.tracing) if app.tracing else None - if app_ops_trace_config is None: - return None - if not app_ops_trace_config.get("enabled"): - return None - - tracing_provider = app_ops_trace_config.get("tracing_provider") - if tracing_provider is None: - return None - try: - provider_config_map[tracing_provider] - except KeyError: - return None - - # decrypt_token - decrypt_trace_config = cls.get_decrypted_tracing_config(app_id, tracing_provider) - if not decrypt_trace_config: - return None - - 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: - # 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[cache_key] = tracing_instance - logger.info("new %s tracing_instance for app_id: %s", mode, app_id) - return tracing_instance - - @classmethod - def get_app_config_through_message_id(cls, message_id: str) -> AppModelConfig | str | None: - app_model_config: AppModelConfig | str | None = None - message_stmt = select(Message).where(Message.id == message_id) - message_data = db.session.scalar(message_stmt) - if not message_data: - return None - conversation_id = message_data.conversation_id - conversation_stmt = select(Conversation).where(Conversation.id == conversation_id) - conversation_data = db.session.scalar(conversation_stmt) - if not conversation_data: - return None - - if conversation_data.app_model_config_id: - config_stmt = select(AppModelConfig).where(AppModelConfig.id == conversation_data.app_model_config_id) - app_model_config = db.session.scalar(config_stmt) - elif conversation_data.app_model_config_id is None and conversation_data.override_model_configs: - app_model_config = conversation_data.override_model_configs - - return app_model_config - - @staticmethod - def check_trace_config_is_effective(tracing_config: dict[str, Any], tracing_provider: str): - """ - Check trace config is effective - :param tracing_config: tracing config - :param tracing_provider: tracing provider - :return: - """ - config_type, trace_instance = ( - provider_config_map[tracing_provider]["config_class"], - provider_config_map[tracing_provider]["trace_instance"], - ) - config = config_type(**tracing_config) - return trace_instance(config).api_check() - - @staticmethod - def get_trace_config_project_key(tracing_config: dict[str, Any], tracing_provider: str): - """ - get trace config is project key - :param tracing_config: tracing config - :param tracing_provider: tracing provider - :return: - """ - config_type, trace_instance = ( - provider_config_map[tracing_provider]["config_class"], - provider_config_map[tracing_provider]["trace_instance"], - ) - config = config_type(**tracing_config) - return trace_instance(config).get_project_key() - - @staticmethod - def get_trace_config_project_url(tracing_config: dict[str, Any], tracing_provider: str): - """ - get trace config is project key - :param tracing_config: tracing config - :param tracing_provider: tracing provider - :return: - """ - config_type, trace_instance = ( - provider_config_map[tracing_provider]["config_class"], - provider_config_map[tracing_provider]["trace_instance"], - ) - config = config_type(**tracing_config) - return trace_instance(config).get_project_url() - - -class TraceTask: - _workflow_run_repo = None - _repo_lock = threading.Lock() - - @classmethod - def _get_workflow_run_repo(cls): - if cls._workflow_run_repo is None: - with cls._repo_lock: - if cls._workflow_run_repo is None: - # Lazy import to avoid circular import during module initialization - from repositories.factory import DifyAPIRepositoryFactory - - session_maker = sessionmaker(bind=db.engine, expire_on_commit=False) - cls._workflow_run_repo = DifyAPIRepositoryFactory.create_api_workflow_run_repository(session_maker) - return cls._workflow_run_repo - - @classmethod - def _calculate_workflow_token_split( - cls, session: "Session", workflow_run_id: str, tenant_id: str - ) -> tuple[int, int]: - """Sum prompt/completion tokens across all node executions for a workflow run. - - Reads from the ``outputs`` column (where LLM nodes store ``usage.prompt_tokens`` - and ``usage.completion_tokens``) rather than ``execution_metadata``, which only - carries ``total_tokens``. Projects only the ``outputs`` column to avoid loading - large JSON blobs unnecessarily. - """ - - from models.workflow import WorkflowNodeExecutionModel - - rows = ( - session.execute( - select(WorkflowNodeExecutionModel.outputs).where( - WorkflowNodeExecutionModel.tenant_id == tenant_id, - WorkflowNodeExecutionModel.workflow_run_id == workflow_run_id, - ) - ) - .scalars() - .all() - ) - - total_prompt = 0 - total_completion = 0 - - for raw in rows: - if not raw: - continue - try: - outputs = JSON_DICT_ADAPTER.validate_json(raw) if isinstance(raw, str) else raw - except (ValueError, TypeError): - continue - if not isinstance(outputs, dict): - continue - usage = outputs.get("usage") - if not isinstance(usage, dict): - continue - prompt = usage.get("prompt_tokens") - if isinstance(prompt, (int, float)): - total_prompt += int(prompt) - completion = usage.get("completion_tokens") - if isinstance(completion, (int, float)): - total_completion += int(completion) - - return (total_prompt, total_completion) - - @classmethod - def _get_user_id_from_metadata(cls, metadata: dict[str, Any]) -> str: - """Extract user ID from metadata, prioritizing end_user over account. - - Returns the actual user ID (end_user or account) who invoked the workflow, - regardless of invoke_from context. - """ - # Priority 1: End user (external users via API/WebApp) - if user_id := metadata.get("from_end_user_id"): - return f"end_user:{user_id}" - - # Priority 2: Account user (internal users via console/debugger) - if user_id := metadata.get("from_account_id"): - return f"account:{user_id}" - - # Priority 3: User (internal users via console/debugger) - if user_id := metadata.get("user_id"): - return f"user:{user_id}" - - return "anonymous" - - def __init__( - self, - 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, - **kwargs, - ): - self.trace_type = trace_type - self.message_id = message_id - 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 - self.file_base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001") - self.app_id = None - self.trace_id = None - self.kwargs = kwargs - if user_id is not None and "user_id" not in self.kwargs: - self.kwargs["user_id"] = user_id - external_trace_id = kwargs.get("external_trace_id") - if external_trace_id: - self.trace_id = external_trace_id - - def execute(self): - return self.preprocess() - - def preprocess(self): - preprocess_map = { - TraceTaskName.CONVERSATION_TRACE: lambda: self.conversation_trace(**self.kwargs), - TraceTaskName.WORKFLOW_TRACE: lambda: self.workflow_trace( - workflow_run_id=self.workflow_run_id, - conversation_id=self.conversation_id, - user_id=self.user_id, - total_tokens_override=self.workflow_total_tokens, - ), - TraceTaskName.MESSAGE_TRACE: lambda: self.message_trace(message_id=self.message_id, **self.kwargs), - TraceTaskName.MODERATION_TRACE: lambda: self.moderation_trace( - message_id=self.message_id, timer=self.timer, **self.kwargs - ), - TraceTaskName.SUGGESTED_QUESTION_TRACE: lambda: self.suggested_question_trace( - message_id=self.message_id, timer=self.timer, **self.kwargs - ), - TraceTaskName.DATASET_RETRIEVAL_TRACE: lambda: self.dataset_retrieval_trace( - message_id=self.message_id, timer=self.timer, **self.kwargs - ), - TraceTaskName.TOOL_TRACE: lambda: self.tool_trace( - message_id=self.message_id, timer=self.timer, **self.kwargs - ), - TraceTaskName.GENERATE_NAME_TRACE: lambda: self.generate_name_trace( - conversation_id=self.conversation_id, timer=self.timer, **self.kwargs - ), - TraceTaskName.PROMPT_GENERATION_TRACE: lambda: self.prompt_generation_trace(**self.kwargs), - TraceTaskName.NODE_EXECUTION_TRACE: lambda: self.node_execution_trace(**self.kwargs), - TraceTaskName.DRAFT_NODE_EXECUTION_TRACE: lambda: self.draft_node_execution_trace(**self.kwargs), - } - - return preprocess_map.get(self.trace_type, lambda: None)() - - # process methods for different trace types - def conversation_trace(self, **kwargs): - return kwargs - - def workflow_trace( - self, - *, - workflow_run_id: str | None, - conversation_id: str | None, - user_id: str | None, - total_tokens_override: int | None = None, - ): - if not workflow_run_id: - return {} - - workflow_run_repo = self._get_workflow_run_repo() - workflow_run = workflow_run_repo.get_workflow_run_by_id_without_tenant(run_id=workflow_run_id) - if not workflow_run: - raise ValueError("Workflow run not found") - - workflow_id = workflow_run.workflow_id - tenant_id = workflow_run.tenant_id - workflow_run_id = workflow_run.id - workflow_run_elapsed_time = workflow_run.elapsed_time - workflow_run_status = workflow_run.status - workflow_run_inputs = workflow_run.inputs_dict - workflow_run_outputs = workflow_run.outputs_dict - workflow_run_version = workflow_run.version - error = workflow_run.error or "" - - total_tokens = total_tokens_override if total_tokens_override is not None else workflow_run.total_tokens - - file_list = workflow_run_inputs.get("sys.file") or [] - query = workflow_run_inputs.get("query") or workflow_run_inputs.get("sys.query") or "" - - with Session(db.engine) as session: - # get workflow_app_log_id - workflow_app_log_data_stmt = select(WorkflowAppLog.id).where( - WorkflowAppLog.tenant_id == tenant_id, - WorkflowAppLog.app_id == workflow_run.app_id, - WorkflowAppLog.workflow_run_id == workflow_run.id, - ) - workflow_app_log_id = session.scalar(workflow_app_log_data_stmt) - # get message_id - message_id = None - if conversation_id: - message_data_stmt = select(Message.id).where( - Message.conversation_id == conversation_id, - Message.workflow_run_id == workflow_run_id, - ) - message_id = session.scalar(message_data_stmt) - prompt_tokens, completion_tokens = self._calculate_workflow_token_split( - session, workflow_run_id=workflow_run_id, tenant_id=tenant_id - ) - - from core.telemetry.gateway import is_enterprise_telemetry_enabled - - if is_enterprise_telemetry_enabled(): - app_name, workspace_name = _lookup_app_and_workspace_names(workflow_run.app_id, tenant_id) - else: - app_name, workspace_name = "", "" - - metadata: dict[str, Any] = { - "workflow_id": workflow_id, - "conversation_id": conversation_id, - "workflow_run_id": workflow_run_id, - "tenant_id": tenant_id, - "elapsed_time": workflow_run_elapsed_time, - "status": workflow_run_status, - "version": workflow_run_version, - "total_tokens": total_tokens, - "file_list": file_list, - "triggered_from": workflow_run.triggered_from, - "user_id": user_id, - "app_id": workflow_run.app_id, - "app_name": app_name, - "workspace_name": workspace_name, - } - - parent_trace_context = self.kwargs.get("parent_trace_context") - dumped_parent_trace_context = _dump_parent_trace_context(parent_trace_context) - if dumped_parent_trace_context: - metadata["parent_trace_context"] = dumped_parent_trace_context - - trace_session_id = _get_trace_session_id(self.kwargs) - if trace_session_id: - metadata["trace_session_id"] = trace_session_id - - workflow_trace_info = WorkflowTraceInfo( - trace_id=self.trace_id, - workflow_data=workflow_run.to_dict(), - conversation_id=conversation_id, - workflow_id=workflow_id, - tenant_id=tenant_id, - workflow_run_id=workflow_run_id, - workflow_run_elapsed_time=workflow_run_elapsed_time, - workflow_run_status=workflow_run_status, - workflow_run_inputs=workflow_run_inputs, - workflow_run_outputs=workflow_run_outputs, - workflow_run_version=workflow_run_version, - error=error, - total_tokens=total_tokens, - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - file_list=file_list, - query=query, - metadata=metadata, - workflow_app_log_id=workflow_app_log_id, - message_id=message_id, - start_time=workflow_run.created_at, - end_time=workflow_run.finished_at, - invoked_by=self._get_user_id_from_metadata(metadata), - ) - return workflow_trace_info - - def message_trace(self, message_id: str | None, **kwargs): - if not message_id: - return {} - message_data = get_message_data(message_id, db.session) - if not message_data: - return {} - conversation_mode_stmt = select(Conversation.mode).where(Conversation.id == message_data.conversation_id) - conversation_modes = db.session.scalars(conversation_mode_stmt).all() - if not conversation_modes or len(conversation_modes) == 0: - return {} - conversation_mode = conversation_modes[0] - created_at = message_data.created_at - inputs = message_data.message - - # get message file data - message_file_data = db.session.scalar(select(MessageFile).where(MessageFile.message_id == message_id).limit(1)) - file_list = [] - if message_file_data and message_file_data.url is not None: - file_url = f"{self.file_base_url}/{message_file_data.url}" if message_file_data else "" - file_list.append(file_url) - - streaming_metrics = self._extract_streaming_metrics(message_data) - - tenant_id = "" - with Session(db.engine) as session: - tid = session.scalar(select(App.tenant_id).where(App.id == message_data.app_id)) - if tid: - tenant_id = str(tid) - - from core.telemetry.gateway import is_enterprise_telemetry_enabled - - if is_enterprise_telemetry_enabled(): - app_name, workspace_name = _lookup_app_and_workspace_names(message_data.app_id, tenant_id) - else: - app_name, workspace_name = "", "" - - metadata = { - "conversation_id": message_data.conversation_id, - "ls_provider": message_data.model_provider, - "ls_model_name": message_data.model_id, - "status": message_data.status, - "from_end_user_id": message_data.from_end_user_id, - "from_account_id": message_data.from_account_id, - "agent_based": message_data.agent_based, - "workflow_run_id": message_data.workflow_run_id, - "from_source": message_data.from_source, - "message_id": message_id, - "tenant_id": tenant_id, - "app_id": message_data.app_id, - "user_id": message_data.from_end_user_id or message_data.from_account_id, - "app_name": app_name, - "workspace_name": workspace_name, - } - if node_execution_id := kwargs.get("node_execution_id"): - metadata["node_execution_id"] = node_execution_id - - trace_session_id = _get_trace_session_id(kwargs) - if trace_session_id: - metadata["trace_session_id"] = trace_session_id - - message_tokens = message_data.message_tokens - - message_trace_info = MessageTraceInfo( - trace_id=self.trace_id, - message_id=message_id, - message_data=message_data.to_dict(session=db.session()), - conversation_model=conversation_mode, - message_tokens=message_tokens, - answer_tokens=message_data.answer_tokens, - total_tokens=message_tokens + message_data.answer_tokens, - error=message_data.error or "", - inputs=inputs, - outputs=message_data.answer, - file_list=file_list, - start_time=created_at, - end_time=message_data.updated_at - if message_data.updated_at and message_data.updated_at > created_at - else created_at + timedelta(seconds=message_data.provider_response_latency), - metadata=metadata, - message_file_data=message_file_data, - conversation_mode=conversation_mode, - gen_ai_server_time_to_first_token=streaming_metrics.get("gen_ai_server_time_to_first_token"), - llm_streaming_time_to_generate=streaming_metrics.get("llm_streaming_time_to_generate"), - is_streaming_request=streaming_metrics.get("is_streaming_request", False), - ) - - return message_trace_info - - def moderation_trace(self, message_id, timer, **kwargs): - moderation_result = kwargs.get("moderation_result") - if not moderation_result: - return {} - inputs = kwargs.get("inputs") - message_data = get_message_data(message_id, db.session) - if not message_data: - return {} - metadata = { - "message_id": message_id, - "action": moderation_result.action, - "preset_response": moderation_result.preset_response, - "query": moderation_result.query, - } - if node_execution_id := kwargs.get("node_execution_id"): - metadata["node_execution_id"] = node_execution_id - - # get workflow_app_log_id - workflow_app_log_id = None - if message_data.workflow_run_id: - workflow_app_log_data = db.session.scalar( - select(WorkflowAppLog).where(WorkflowAppLog.workflow_run_id == message_data.workflow_run_id).limit(1) - ) - workflow_app_log_id = str(workflow_app_log_data.id) if workflow_app_log_data else None - - moderation_trace_info = ModerationTraceInfo( - trace_id=self.trace_id, - message_id=workflow_app_log_id or message_id, - inputs=inputs, - message_data=message_data.to_dict(session=db.session()), - flagged=moderation_result.flagged, - action=moderation_result.action, - preset_response=moderation_result.preset_response, - query=moderation_result.query, - start_time=timer.get("start"), - end_time=timer.get("end"), - metadata=metadata, - ) - - return moderation_trace_info - - def suggested_question_trace(self, message_id, timer, **kwargs): - suggested_question = kwargs.get("suggested_question", []) - message_data = get_message_data(message_id, db.session) - if not message_data: - return {} - metadata = { - "message_id": message_id, - "ls_provider": message_data.model_provider, - "ls_model_name": message_data.model_id, - "status": message_data.status, - "from_end_user_id": message_data.from_end_user_id, - "from_account_id": message_data.from_account_id, - "agent_based": message_data.agent_based, - "workflow_run_id": message_data.workflow_run_id, - "from_source": message_data.from_source, - } - if node_execution_id := kwargs.get("node_execution_id"): - metadata["node_execution_id"] = node_execution_id - - # get workflow_app_log_id - workflow_app_log_id = None - if message_data.workflow_run_id: - workflow_app_log_data = db.session.scalar( - select(WorkflowAppLog).where(WorkflowAppLog.workflow_run_id == message_data.workflow_run_id).limit(1) - ) - workflow_app_log_id = str(workflow_app_log_data.id) if workflow_app_log_data else None - - suggested_question_trace_info = SuggestedQuestionTraceInfo( - trace_id=self.trace_id, - message_id=workflow_app_log_id or message_id, - message_data=message_data.to_dict(session=db.session()), - inputs=message_data.message, - outputs=message_data.answer, - start_time=timer.get("start"), - end_time=timer.get("end"), - metadata=metadata, - total_tokens=message_data.message_tokens + message_data.answer_tokens, - status=message_data.status, - error=message_data.error, - from_account_id=message_data.from_account_id, - agent_based=message_data.agent_based, - from_source=message_data.from_source, - model_provider=message_data.model_provider, - model_id=message_data.model_id, - suggested_question=suggested_question, - level=message_data.status, - status_message=message_data.error, - ) - - return suggested_question_trace_info - - def dataset_retrieval_trace(self, message_id, timer, **kwargs): - documents = kwargs.get("documents") - message_data = get_message_data(message_id, db.session) - if not message_data: - return {} - - tenant_id = "" - with Session(db.engine) as session: - tid = session.scalar(select(App.tenant_id).where(App.id == message_data.app_id)) - if tid: - tenant_id = str(tid) - - from core.telemetry.gateway import is_enterprise_telemetry_enabled - - if is_enterprise_telemetry_enabled(): - app_name, workspace_name = _lookup_app_and_workspace_names(message_data.app_id, tenant_id) - else: - app_name, workspace_name = "", "" - - doc_list = [doc.model_dump() for doc in documents] if documents else [] - dataset_ids: set[str] = set() - for doc in doc_list: - doc_meta = doc.get("metadata") or {} - did = doc_meta.get("dataset_id") - if did: - dataset_ids.add(did) - - embedding_models: dict[str, dict[str, str]] = {} - if dataset_ids: - with Session(db.engine) as session: - rows = session.execute( - select(Dataset.id, Dataset.embedding_model, Dataset.embedding_model_provider).where( - Dataset.id.in_(list(dataset_ids)) - ) - ).all() - for row in rows: - embedding_models[str(row[0])] = { - "embedding_model": row[1] or "", - "embedding_model_provider": row[2] or "", - } - - # Extract rerank model info from retrieval_model kwargs - rerank_model_provider = "" - rerank_model_name = "" - if "retrieval_model" in kwargs: - retrieval_model = kwargs["retrieval_model"] - if isinstance(retrieval_model, dict): - reranking_model = retrieval_model.get("reranking_model") - if isinstance(reranking_model, dict): - rerank_model_provider = reranking_model.get("reranking_provider_name", "") - rerank_model_name = reranking_model.get("reranking_model_name", "") - - metadata = { - "message_id": message_id, - "ls_provider": message_data.model_provider, - "ls_model_name": message_data.model_id, - "status": message_data.status, - "from_end_user_id": message_data.from_end_user_id, - "from_account_id": message_data.from_account_id, - "agent_based": message_data.agent_based, - "workflow_run_id": message_data.workflow_run_id, - "from_source": message_data.from_source, - "tenant_id": tenant_id, - "app_id": message_data.app_id, - "user_id": message_data.from_end_user_id or message_data.from_account_id, - "app_name": app_name, - "workspace_name": workspace_name, - "embedding_models": embedding_models, - "rerank_model_provider": rerank_model_provider, - "rerank_model_name": rerank_model_name, - } - if node_execution_id := kwargs.get("node_execution_id"): - metadata["node_execution_id"] = node_execution_id - - dataset_retrieval_trace_info = DatasetRetrievalTraceInfo( - trace_id=self.trace_id, - message_id=message_id, - inputs=message_data.query or message_data.inputs_with_session(session=db.session()), - documents=doc_list, - start_time=timer.get("start"), - end_time=timer.get("end"), - metadata=metadata, - message_data=message_data.to_dict(session=db.session()), - error=kwargs.get("error"), - ) - - return dataset_retrieval_trace_info - - def tool_trace(self, message_id, timer, **kwargs): - tool_name = kwargs.get("tool_name", "") - tool_inputs = kwargs.get("tool_inputs", {}) - tool_outputs = kwargs.get("tool_outputs", {}) - message_data = get_message_data(message_id, db.session) - if not message_data: - return {} - tool_config = {} - time_cost = 0 - error = None - tool_parameters = {} - created_time = message_data.created_at - end_time = message_data.updated_at - agent_thoughts = message_data.agent_thoughts_with_session(session=db.session()) - for agent_thought in agent_thoughts: - if tool_name in agent_thought.tools: - created_time = agent_thought.created_at - tool_meta_data = agent_thought.tool_meta.get(tool_name, {}) - tool_config = tool_meta_data.get("tool_config", {}) - time_cost = tool_meta_data.get("time_cost", 0) - end_time = created_time + timedelta(seconds=time_cost) - error = tool_meta_data.get("error", "") - tool_parameters = tool_meta_data.get("tool_parameters", {}) - metadata = { - "message_id": message_id, - "tool_name": tool_name, - "tool_inputs": tool_inputs, - "tool_outputs": tool_outputs, - "tool_config": tool_config, - "time_cost": time_cost, - "error": error, - "tool_parameters": tool_parameters, - } - if message_data.workflow_run_id: - metadata["workflow_run_id"] = message_data.workflow_run_id - if node_execution_id := kwargs.get("node_execution_id"): - metadata["node_execution_id"] = node_execution_id - - file_url = "" - message_file_data = db.session.scalar(select(MessageFile).where(MessageFile.message_id == message_id).limit(1)) - if message_file_data: - message_file_id = message_file_data.id if message_file_data else None - type = message_file_data.type - created_by_role = message_file_data.created_by_role - created_user_id = message_file_data.created_by - file_url = f"{self.file_base_url}/{message_file_data.url}" - - metadata.update( - { - "message_file_id": message_file_id, - "created_by_role": created_by_role, - "created_user_id": created_user_id, - "type": type, - } - ) - - tool_trace_info = ToolTraceInfo( - trace_id=self.trace_id, - message_id=message_id, - message_data=message_data.to_dict(session=db.session()), - tool_name=tool_name, - start_time=timer.get("start") if timer else created_time, - end_time=timer.get("end") if timer else end_time, - tool_inputs=tool_inputs, - tool_outputs=tool_outputs, - metadata=metadata, - message_file_data=message_file_data, - error=error, - inputs=message_data.message, - outputs=message_data.answer, - tool_config=tool_config, - time_cost=time_cost, - tool_parameters=tool_parameters, - file_url=file_url, - ) - - return tool_trace_info - - def generate_name_trace(self, conversation_id, timer, **kwargs): - generate_conversation_name = kwargs.get("generate_conversation_name") - inputs = kwargs.get("inputs") - tenant_id = kwargs.get("tenant_id") - if not tenant_id: - return {} - start_time = timer.get("start") - end_time = timer.get("end") - - metadata = { - "conversation_id": conversation_id, - "tenant_id": tenant_id, - } - if node_execution_id := kwargs.get("node_execution_id"): - metadata["node_execution_id"] = node_execution_id - - 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, - start_time=start_time, - end_time=end_time, - metadata=metadata, - tenant_id=tenant_id, - ) - - return generate_name_trace_info - - def prompt_generation_trace(self, **kwargs) -> PromptGenerationTraceInfo | dict: - tenant_id = kwargs.get("tenant_id", "") - user_id = kwargs.get("user_id", "") - app_id = kwargs.get("app_id") - operation_type = kwargs.get("operation_type", "") - instruction = kwargs.get("instruction", "") - generated_output = kwargs.get("generated_output", "") - - prompt_tokens = kwargs.get("prompt_tokens", 0) - completion_tokens = kwargs.get("completion_tokens", 0) - total_tokens = kwargs.get("total_tokens", 0) - - model_provider = kwargs.get("model_provider", "") - model_name = kwargs.get("model_name", "") - - latency = kwargs.get("latency", 0.0) - - timer = kwargs.get("timer") - start_time = timer.get("start") if timer else None - end_time = timer.get("end") if timer else None - - total_price = kwargs.get("total_price") - currency = kwargs.get("currency") - - error = kwargs.get("error") - - app_name = None - workspace_name = None - if app_id: - app_name, workspace_name = _lookup_app_and_workspace_names(app_id, tenant_id) - - metadata = { - "tenant_id": tenant_id, - "user_id": user_id, - "app_id": app_id or "", - "app_name": app_name, - "workspace_name": workspace_name, - "operation_type": operation_type, - "model_provider": model_provider, - "model_name": model_name, - } - if node_execution_id := kwargs.get("node_execution_id"): - metadata["node_execution_id"] = node_execution_id - - return PromptGenerationTraceInfo( - trace_id=self.trace_id, - inputs=instruction, - outputs=generated_output, - start_time=start_time, - end_time=end_time, - metadata=metadata, - tenant_id=tenant_id, - user_id=user_id, - app_id=app_id, - operation_type=operation_type, - instruction=instruction, - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=total_tokens, - model_provider=model_provider, - model_name=model_name, - latency=latency, - total_price=total_price, - currency=currency, - error=error, - ) - - def node_execution_trace(self, **kwargs) -> WorkflowNodeTraceInfo | dict[str, Any]: - node_data: dict[str, Any] = kwargs.get("node_execution_data", {}) - if not node_data: - return {} - - from core.telemetry.gateway import is_enterprise_telemetry_enabled - - if is_enterprise_telemetry_enabled(): - app_name, workspace_name = _lookup_app_and_workspace_names( - node_data.get("app_id"), node_data.get("tenant_id") - ) - else: - app_name, workspace_name = "", "" - - # Try tool credential lookup first - credential_id = node_data.get("credential_id") - if is_enterprise_telemetry_enabled(): - credential_name = _lookup_credential_name(credential_id, node_data.get("credential_provider_type")) - # If no credential_id found (e.g., LLM nodes), try LLM credential lookup - if not credential_id: - llm_cred_id, llm_cred_name = _lookup_llm_credential_info( - tenant_id=node_data.get("tenant_id"), - provider=node_data.get("model_provider"), - model=node_data.get("model_name"), - model_type="llm", - ) - if llm_cred_id: - credential_id = llm_cred_id - credential_name = llm_cred_name - else: - credential_name = "" - metadata: dict[str, Any] = { - "tenant_id": node_data.get("tenant_id"), - "app_id": node_data.get("app_id"), - "app_name": app_name, - "workspace_name": workspace_name, - "user_id": node_data.get("user_id"), - "invoke_from": node_data.get("invoke_from"), - "credential_id": credential_id, - "credential_name": credential_name, - "dataset_ids": node_data.get("dataset_ids"), - "dataset_names": node_data.get("dataset_names"), - "plugin_name": node_data.get("plugin_name"), - } - - parent_trace_context = node_data.get("parent_trace_context") - dumped_parent_trace_context = _dump_parent_trace_context(parent_trace_context) - if dumped_parent_trace_context: - metadata["parent_trace_context"] = dumped_parent_trace_context - - message_id: str | None = None - conversation_id = node_data.get("conversation_id") - workflow_execution_id = node_data.get("workflow_execution_id") - if conversation_id and workflow_execution_id and not dumped_parent_trace_context: - with Session(db.engine) as session: - msg_id = session.scalar( - select(Message.id).where( - Message.conversation_id == conversation_id, - Message.workflow_run_id == workflow_execution_id, - ) - ) - if msg_id: - message_id = str(msg_id) - metadata["message_id"] = message_id - if conversation_id: - metadata["conversation_id"] = conversation_id - - return WorkflowNodeTraceInfo( - trace_id=self.trace_id, - message_id=message_id, - start_time=node_data.get("created_at"), - end_time=node_data.get("finished_at"), - metadata=metadata, - workflow_id=node_data.get("workflow_id", ""), - workflow_run_id=node_data.get("workflow_execution_id", ""), - tenant_id=node_data.get("tenant_id", ""), - node_execution_id=node_data.get("node_execution_id", ""), - node_id=node_data.get("node_id", ""), - node_type=node_data.get("node_type", ""), - title=node_data.get("title", ""), - status=node_data.get("status", ""), - error=node_data.get("error"), - elapsed_time=node_data.get("elapsed_time", 0.0), - index=node_data.get("index", 0), - predecessor_node_id=node_data.get("predecessor_node_id"), - total_tokens=node_data.get("total_tokens", 0), - total_price=node_data.get("total_price", 0.0), - currency=node_data.get("currency"), - model_provider=node_data.get("model_provider"), - model_name=node_data.get("model_name"), - prompt_tokens=node_data.get("prompt_tokens"), - completion_tokens=node_data.get("completion_tokens"), - tool_name=node_data.get("tool_name"), - iteration_id=node_data.get("iteration_id"), - iteration_index=node_data.get("iteration_index"), - loop_id=node_data.get("loop_id"), - loop_index=node_data.get("loop_index"), - parallel_id=node_data.get("parallel_id"), - node_inputs=node_data.get("node_inputs"), - node_outputs=node_data.get("node_outputs"), - process_data=node_data.get("process_data"), - invoked_by=self._get_user_id_from_metadata(metadata), - ) - - def draft_node_execution_trace(self, **kwargs) -> DraftNodeExecutionTrace | dict: - node_trace = self.node_execution_trace(**kwargs) - if not isinstance(node_trace, WorkflowNodeTraceInfo): - return node_trace - return DraftNodeExecutionTrace(**node_trace.model_dump()) - - def _extract_streaming_metrics(self, message_data) -> dict[str, Any]: - if not message_data.message_metadata: - return {} - - try: - metadata = JSON_DICT_ADAPTER.validate_json(message_data.message_metadata) - usage = metadata.get("usage", {}) - time_to_first_token = usage.get("time_to_first_token") - time_to_generate = usage.get("time_to_generate") - - return { - "gen_ai_server_time_to_first_token": time_to_first_token, - "llm_streaming_time_to_generate": time_to_generate, - "is_streaming_request": time_to_first_token is not None, - } - except (ValueError, AttributeError): - return {} - - -trace_manager_timer: threading.Timer | None = None -trace_manager_queue: queue.Queue = queue.Queue() -trace_manager_interval = int(os.getenv("TRACE_QUEUE_MANAGER_INTERVAL", 5)) -trace_manager_batch_size = int(os.getenv("TRACE_QUEUE_MANAGER_BATCH_SIZE", 100)) - - -class TraceQueueManager: - def __init__(self, app_id=None, user_id=None): - global trace_manager_timer - - self.app_id = app_id - self.user_id = user_id - self.trace_instance = None - self._application_tracing_enabled = False - if dify_config.OPS_AGENT_TRACE_ENABLED: - tracing_hooks = application_tracing_hooks() - self._application_tracing_enabled = ( - tracing_hooks.queue_enabled_at_start(app_id) if tracing_hooks is not None else False - ) - else: - self.trace_instance = OpsTraceManager.get_ops_trace_instance(app_id) - self.flask_app = current_app._get_current_object() # type: ignore - - from core.telemetry.gateway import is_enterprise_telemetry_enabled - - self._enterprise_telemetry_enabled = is_enterprise_telemetry_enabled() - if trace_manager_timer is None: - self.start_timer() - - def add_trace_task(self, trace_task: TraceTask): - global trace_manager_timer, trace_manager_queue - try: - if self._enterprise_telemetry_enabled or self._application_tracing_enabled or self.trace_instance: - trace_task.app_id = self.app_id - trace_manager_queue.put(trace_task) - except Exception: - logger.exception("Error adding trace task, trace_type %s", trace_task.trace_type) - finally: - self.start_timer() - - def collect_tasks(self): - global trace_manager_queue - tasks: list[TraceTask] = [] - while len(tasks) < trace_manager_batch_size and not trace_manager_queue.empty(): - task = trace_manager_queue.get_nowait() - tasks.append(task) - trace_manager_queue.task_done() - return tasks - - def run(self): - try: - tasks = self.collect_tasks() - if tasks: - self.send_to_celery(tasks) - except Exception: - logger.exception("Error processing trace tasks") - - def start_timer(self): - global trace_manager_timer - if trace_manager_timer is None or not trace_manager_timer.is_alive(): - trace_manager_timer = threading.Timer(trace_manager_interval, self.run) - trace_manager_timer.name = f"trace_manager_timer_{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}" - 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: - 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: - 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 deleted file mode 100644 index bf131bef68cfa3..00000000000000 --- a/api/core/ops/unified_trace/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""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 deleted file mode 100644 index e8d2a52f089686..00000000000000 --- a/api/core/ops/unified_trace/entities.py +++ /dev/null @@ -1,79 +0,0 @@ -"""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/parent_context.py b/api/core/ops/unified_trace/parent_context.py deleted file mode 100644 index 25645ffe7c355f..00000000000000 --- a/api/core/ops/unified_trace/parent_context.py +++ /dev/null @@ -1,249 +0,0 @@ -"""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 deleted file mode 100644 index a05d36cfc84d55..00000000000000 --- a/api/core/ops/unified_trace/provider.py +++ /dev/null @@ -1,82 +0,0 @@ -"""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 deleted file mode 100644 index 7619b28923fcb5..00000000000000 --- a/api/core/ops/unified_trace/registry.py +++ /dev/null @@ -1,35 +0,0 @@ -"""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 deleted file mode 100644 index 2847bda22f78cb..00000000000000 --- a/api/core/ops/unified_trace/trace_builder.py +++ /dev/null @@ -1,452 +0,0 @@ -"""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: - """Export caller executions; Workflow Tool source internals stay in authorized console traces.""" - - 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/core/ops/utils.py b/api/core/ops/utils.py deleted file mode 100644 index 0526c6ab5be832..00000000000000 --- a/api/core/ops/utils.py +++ /dev/null @@ -1,175 +0,0 @@ -from contextlib import contextmanager -from datetime import datetime -from typing import Any -from urllib.parse import urlparse - -from pydantic import TypeAdapter -from sqlalchemy import select -from sqlalchemy.orm import Session, scoped_session - -from models.model import Message - -JSON_DICT_ADAPTER: TypeAdapter[dict[str, Any]] = TypeAdapter(dict[str, Any]) - - -def filter_none_values(data: dict[str, Any]) -> dict[str, Any]: - new_data = {} - for key, value in data.items(): - if value is None: - continue - if isinstance(value, datetime): - new_data[key] = value.isoformat() - else: - new_data[key] = value - return new_data - - -def get_message_data(message_id: str, session: Session | scoped_session): - return session.scalar(select(Message).where(Message.id == message_id)) - - -@contextmanager -def measure_time(): - timing_info = {"start": datetime.now(), "end": None} - try: - yield timing_info - finally: - # pyrefly: ignore [bad-assignment] - timing_info["end"] = datetime.now() - - -def replace_text_with_content(data): - match data: - case dict(): - new_data = {} - for key, value in data.items(): - if key == "text": - new_data["content"] = value - else: - new_data[key] = replace_text_with_content(value) - return new_data - case list(): - return [replace_text_with_content(item) for item in data] - case _: - return data - - -def generate_dotted_order(run_id: str, start_time: str | datetime, parent_dotted_order: str | None = None) -> str: - """ - generate dotted_order for langsmith - """ - start_time = datetime.fromisoformat(start_time) if isinstance(start_time, str) else start_time - timestamp = start_time.strftime("%Y%m%dT%H%M%S%f") + "Z" - current_segment = f"{timestamp}{run_id}" - - if parent_dotted_order is None: - return current_segment - - return f"{parent_dotted_order}.{current_segment}" - - -def validate_url(url: str, default_url: str, allowed_schemes: tuple = ("https", "http")) -> str: - """ - Validate and normalize URL with proper error handling. - - NOTE: This function does not retain the `path` component of the provided URL. - In most cases, it is recommended to use `validate_url_with_path` instead. - - This function is deprecated and retained only for compatibility purposes. - New implementations should use `validate_url_with_path`. - - Args: - url: The URL to validate - default_url: Default URL to use if input is None or empty - allowed_schemes: Tuple of allowed URL schemes (default: https, http) - - Returns: - Normalized URL string - - Raises: - ValueError: If URL format is invalid or scheme not allowed - """ - if not url or url.strip() == "": - return default_url - - # Parse URL to validate format - parsed = urlparse(url) - - # Check if scheme is allowed - if parsed.scheme not in allowed_schemes: - raise ValueError(f"URL scheme must be one of: {', '.join(allowed_schemes)}") - - # Reconstruct URL with only scheme, netloc (removing path, query, fragment) - normalized_url = f"{parsed.scheme}://{parsed.netloc}" - - return normalized_url - - -def validate_url_with_path( - url: str, - default_url: str, - required_suffix: str | None = None, - *, - allowed_schemes: tuple[str, ...] = ("https", "http"), -) -> str: - """ - Validate URL that may include path components - - Args: - url: The URL to validate - default_url: Default URL to use if input is None or empty - required_suffix: Optional suffix that URL must end with - allowed_schemes: Tuple of allowed URL schemes (default: https, http) - - Returns: - Validated URL string, returned verbatim so path, query and trailing - separators survive — `required_suffix` consumers depend on that - - Raises: - ValueError: If URL format is invalid or doesn't match required suffix - """ - if not url or url.strip() == "": - return default_url - url = url.strip() - - # Parse URL to validate format - parsed = urlparse(url) - - # Check if scheme is allowed - if parsed.scheme not in allowed_schemes: - expected = " or ".join(f"{scheme}://" for scheme in allowed_schemes) - raise ValueError(f"URL must start with {expected}") - - # Check required suffix if specified - if required_suffix and not url.endswith(required_suffix): - raise ValueError(f"URL should end with {required_suffix}") - - return url - - -def validate_project_name(project: str, default_name: str) -> str: - """ - Validate and normalize project name - - Args: - project: Project name to validate - default_name: Default name to use if input is None or empty - - Returns: - Normalized project name - """ - if not project or project.strip() == "": - return default_name - - return project.strip() - - -def validate_integer_id(id_str: str) -> str: - """ - Validate and normalize integer ID - """ - id_str = id_str.strip() - if not id_str.isdigit(): - raise ValueError("ID must be a valid integer") - - return id_str diff --git a/api/core/rag/retrieval/dataset_retrieval.py b/api/core/rag/retrieval/dataset_retrieval.py index 6b266800c16762..eee5b887691721 100644 --- a/api/core/rag/retrieval/dataset_retrieval.py +++ b/api/core/rag/retrieval/dataset_retrieval.py @@ -8,6 +8,7 @@ from collections.abc import Generator, Mapping from typing import Any, cast +from dify_trace_sdk.events import DatasetRetrieved from flask import Flask, current_app from opentelemetry.trace import get_current_span from sqlalchemy import and_, func, literal, or_, select, update @@ -35,9 +36,6 @@ from core.memory.token_buffer_memory import TokenBufferMemory from core.model_context import with_credit_usage_created_by, with_credit_usage_metadata from core.model_manager import ModelInstance, ModelManager -from core.ops.entities.trace_entity import TraceTaskName -from core.ops.ops_trace_manager import TraceQueueManager, TraceTask -from core.ops.utils import measure_time from core.prompt.advanced_prompt_transform import AdvancedPromptTransform from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate from core.prompt.simple_prompt_transform import ModelMode @@ -80,6 +78,7 @@ from graphon.model_runtime.entities.message_entities import PromptMessage, PromptMessageRole, PromptMessageTool from graphon.model_runtime.entities.model_entities import ModelFeature, ModelType from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel +from libs.datetime_utils import Timing, measure_time from libs.helper import parse_uuid_str_or_none from libs.json_in_md_parser import parse_and_check_json_markdown from models.dataset import ( @@ -949,7 +948,7 @@ def _on_retrieval_end( flask_app: Flask, documents: list[Document], message_id: str | None = None, - timer: dict[str, Any] | None = None, + timer: Timing | None = None, ): """Handle retrieval end.""" with flask_app.app_context(): @@ -1062,17 +1061,29 @@ def _on_retrieval_end( self._send_trace_task(message_id, documents, timer) - def _send_trace_task(self, message_id: str | None, documents: list[Document], timer: dict[str, Any] | None): - """Send trace task if trace manager is available.""" - trace_manager: TraceQueueManager | None = ( - self.application_generate_entity.trace_manager if self.application_generate_entity else None - ) - if trace_manager: - trace_manager.add_trace_task( - TraceTask( - TraceTaskName.DATASET_RETRIEVAL_TRACE, message_id=message_id, documents=documents, timer=timer - ) + def _send_trace_task(self, message_id: str | None, documents: list[Document], timer: Timing | None): + trace_scope = self.application_generate_entity.trace_scope if self.application_generate_entity else None + if trace_scope is None or timer is None: + return + trace_scope.emit( + DatasetRetrieved( + message_id=message_id, + # Embeddings only bloat traces. + documents=[ + document.model_dump( + mode="json", + exclude={ + "vector": True, + "children": {"__all__": {"vector"}}, + "attachments": {"__all__": {"vector"}}, + }, + ) + for document in documents + ], + started_at=timer.started_at, + ended_at=timer.ended_at, ) + ) @staticmethod def _resolve_creator_user_role(user_from: str) -> CreatorUserRole | None: diff --git a/api/core/telemetry/__init__.py b/api/core/telemetry/__init__.py index 13385f514a65d6..db126091dab1aa 100644 --- a/api/core/telemetry/__init__.py +++ b/api/core/telemetry/__init__.py @@ -6,7 +6,6 @@ from __future__ import annotations -from core.ops.entities.trace_entity import TraceTaskName from core.telemetry.events import ( AppCreatedEvent, AppDeletedEvent, @@ -30,6 +29,5 @@ "PromptGenerationPayload", "TelemetryContext", "TelemetryEvent", - "TraceTaskName", "emit", ] diff --git a/api/core/telemetry/events.py b/api/core/telemetry/events.py index c1872d5b1d1616..dbf3a425d4aec8 100644 --- a/api/core/telemetry/events.py +++ b/api/core/telemetry/events.py @@ -2,9 +2,11 @@ from collections.abc import Mapping from dataclasses import dataclass +from datetime import datetime from typing import Any, NotRequired, Protocol, TypedDict, runtime_checkable -from core.ops.entities.trace_entity import TraceTaskName +from dify_trace_sdk.events import DraftNodeExecuted, PromptGenerated, TracePayload, json_safe + from enterprise.telemetry.contracts import SignalType, TelemetryCase # --------------------------------------------------------------------------- @@ -70,10 +72,11 @@ class PromptGenerationPayload(TypedDict): completion_tokens: int total_tokens: int latency: float - total_price: NotRequired[float | None] - currency: NotRequired[str | None] - timer: NotRequired[dict] - error: NotRequired[str | None] + total_price: float | None + currency: str | None + error: str | None + started_at: datetime + ended_at: datetime class FeedbackCreatedPayload(TypedDict): @@ -126,8 +129,13 @@ def case(self) -> TelemetryCase: ... def signal_type(self) -> SignalType: ... @property def ce_eligible(self) -> bool: ... - @property - def trace_task_name(self) -> TraceTaskName | None: ... + + +@runtime_checkable +class TraceSignalEvent(TelemetryEvent, Protocol): + """A ``SignalType.TRACE`` event, delivered through the trace pipeline.""" + + def trace_payload(self) -> TracePayload: ... # --------------------------------------------------------------------------- @@ -142,7 +150,10 @@ class DraftNodeExecutionTraceEvent: case: TelemetryCase = TelemetryCase.DRAFT_NODE_EXECUTION signal_type: SignalType = SignalType.TRACE ce_eligible: bool = False - trace_task_name: TraceTaskName | None = TraceTaskName.DRAFT_NODE_EXECUTION_TRACE + + def trace_payload(self) -> TracePayload: + node_execution = self.payload["node_execution_data"] + return DraftNodeExecuted(node_execution={key: json_safe(value) for key, value in node_execution.items()}) @dataclass(frozen=True) @@ -152,7 +163,25 @@ class PromptGenerationEvent: case: TelemetryCase = TelemetryCase.PROMPT_GENERATION signal_type: SignalType = SignalType.TRACE ce_eligible: bool = False - trace_task_name: TraceTaskName | None = TraceTaskName.PROMPT_GENERATION_TRACE + + def trace_payload(self) -> TracePayload: + payload = self.payload + return PromptGenerated( + operation_type=payload["operation_type"], + instruction=payload["instruction"], + generated_output=payload["generated_output"], + model_provider=payload["model_provider"], + model_name=payload["model_name"], + prompt_tokens=payload["prompt_tokens"], + completion_tokens=payload["completion_tokens"], + total_tokens=payload["total_tokens"], + latency=payload["latency"], + total_price=payload["total_price"], + currency=payload["currency"], + error=payload["error"], + started_at=payload["started_at"], + ended_at=payload["ended_at"], + ) @dataclass(frozen=True) @@ -162,7 +191,6 @@ class AppCreatedEvent: case: TelemetryCase = TelemetryCase.APP_CREATED signal_type: SignalType = SignalType.METRIC_LOG ce_eligible: bool = False - trace_task_name: TraceTaskName | None = None @dataclass(frozen=True) @@ -172,7 +200,6 @@ class AppUpdatedEvent: case: TelemetryCase = TelemetryCase.APP_UPDATED signal_type: SignalType = SignalType.METRIC_LOG ce_eligible: bool = False - trace_task_name: TraceTaskName | None = None @dataclass(frozen=True) @@ -182,7 +209,6 @@ class AppDeletedEvent: case: TelemetryCase = TelemetryCase.APP_DELETED signal_type: SignalType = SignalType.METRIC_LOG ce_eligible: bool = False - trace_task_name: TraceTaskName | None = None @dataclass(frozen=True) @@ -192,4 +218,3 @@ class FeedbackCreatedEvent: case: TelemetryCase = TelemetryCase.FEEDBACK_CREATED signal_type: SignalType = SignalType.METRIC_LOG ce_eligible: bool = False - trace_task_name: TraceTaskName | None = None diff --git a/api/core/telemetry/gateway.py b/api/core/telemetry/gateway.py index 92cde4e79f8a68..75874baf95aeea 100644 --- a/api/core/telemetry/gateway.py +++ b/api/core/telemetry/gateway.py @@ -1,10 +1,10 @@ """Telemetry gateway — single routing layer for all editions. -Dispatches ``TelemetryEvent`` instances to either the CE/EE trace -pipeline (``TraceQueueManager``) or the enterprise-only metric/log -Celery queue. Each event class carries its own routing metadata -(``signal_type``, ``ce_eligible``, ``trace_task_name``), so this -module contains no per-case mapping tables. +Dispatches ``TelemetryEvent`` instances to either the trace pipeline +or the enterprise-only metric/log Celery queue. Each event class +carries its own routing metadata (``signal_type``, ``ce_eligible``, and +``trace_payload`` for trace signals), so this module contains no +per-case mapping tables. This module lives in ``core/`` so both CE and EE share one ``emit()`` entry point. Enterprise-specific dispatch (Celery task, @@ -17,14 +17,12 @@ import json import logging import uuid -from typing import TYPE_CHECKING, Any, cast +from typing import Any, cast -from core.telemetry.events import TelemetryEvent +from core.telemetry.events import TelemetryEvent, TraceSignalEvent from enterprise.telemetry.contracts import SignalType from extensions.ext_storage import storage - -if TYPE_CHECKING: - from core.ops.ops_trace_manager import TraceQueueManager +from extensions.tracing import create_trace_scope_for_app logger = logging.getLogger(__name__) @@ -82,14 +80,14 @@ def _handle_payload_sizing( # --------------------------------------------------------------------------- -def emit(event: TelemetryEvent, trace_manager: TraceQueueManager | None = None) -> None: +def emit(event: TelemetryEvent) -> None: """Emit a telemetry event.""" if not event.ce_eligible and not is_enterprise_telemetry_enabled(): logger.debug("Dropping EE-only event: case=%s (EE disabled)", event.case) return if event.signal_type == SignalType.TRACE: - _emit_trace(event, trace_manager) + _emit_trace(event) else: _emit_metric_log(event) @@ -99,21 +97,17 @@ def emit(event: TelemetryEvent, trace_manager: TraceQueueManager | None = None) # --------------------------------------------------------------------------- -def _emit_trace(event: TelemetryEvent, trace_manager: TraceQueueManager | None) -> None: - from core.ops.ops_trace_manager import TraceQueueManager as LocalTraceQueueManager - from core.ops.ops_trace_manager import TraceTask - - if event.trace_task_name is None: - logger.warning("No trace_task_name on event: case=%s", event.case) +def _emit_trace(event: TelemetryEvent) -> None: + if not isinstance(event, TraceSignalEvent): + logger.warning("Trace signal without a trace payload: case=%s", event.case) return - ctx = event.context - queue_manager: TraceQueueManager = trace_manager or LocalTraceQueueManager( - app_id=ctx.app_id, - user_id=ctx.user_id, - ) - queue_manager.add_trace_task(TraceTask(event.trace_task_name, user_id=ctx.user_id, **event.payload)) - logger.debug("Enqueued trace task: case=%s, app_id=%s", event.case, ctx.app_id) + if ctx.tenant_id is None: + logger.warning("Trace signal without a workspace: case=%s", event.case) + return + scope = create_trace_scope_for_app(tenant_id=ctx.tenant_id, app_id=ctx.app_id, user_id=ctx.user_id, extras={}) + if scope is not None: + scope.emit(event.trace_payload()) def _emit_metric_log(event: TelemetryEvent) -> None: diff --git a/api/core/tools/tool_engine.py b/api/core/tools/tool_engine.py index 47210354367e15..1d22a41fc3b876 100644 --- a/api/core/tools/tool_engine.py +++ b/api/core/tools/tool_engine.py @@ -7,13 +7,14 @@ from mimetypes import guess_type from typing import Any, cast +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import ToolInvoked, json_safe from sqlalchemy.orm import Session, sessionmaker from yarl import URL from core.app.entities.app_invoke_entities import InvokeFrom from core.callback_handler.agent_tool_callback_handler import DifyAgentCallbackHandler from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler -from core.ops.ops_trace_manager import TraceQueueManager from core.tools.__base.tool import Tool from core.tools.entities.tool_entities import ( ToolInvokeMessage, @@ -34,6 +35,7 @@ from core.tools.workflow_as_tool.tool import WorkflowTool from extensions.ext_database import db from graphon.file import FileTransferMethod, FileType +from libs.datetime_utils import utc_now from models.enums import CreatorUserRole, MessageFileBelongsTo from models.model import Message, MessageFile @@ -55,7 +57,7 @@ def agent_invoke( message: Message, invoke_from: InvokeFrom, agent_tool_callback: DifyAgentCallbackHandler, - trace_manager: TraceQueueManager | None = None, + trace_scope: TraceScope | None, conversation_id: str | None = None, app_id: str | None = None, message_id: str | None = None, @@ -83,6 +85,7 @@ def agent_invoke( # hit the callback handler agent_tool_callback.on_tool_start(tool_name=tool.entity.identity.name, tool_inputs=tool_parameters) + started_at = utc_now() messages = ToolEngine._invoke(session, tool, tool_parameters, user_id, conversation_id, app_id, message_id) invocation_meta_dict: dict[str, ToolInvokeMeta] = {} @@ -121,9 +124,21 @@ def message_callback( tool_name=tool.entity.identity.name, tool_inputs=tool_parameters, tool_outputs=plain_text, - message_id=message.id, - trace_manager=trace_manager, ) + if trace_scope is not None: + trace_scope.emit( + ToolInvoked( + message_id=message.id, + tool_name=tool.entity.identity.name, + tool_inputs={key: json_safe(value) for key, value in tool_parameters.items()}, + tool_outputs=plain_text, + tool_config={key: json_safe(value) for key, value in (meta.tool_config or {}).items()}, + error=meta.error, + time_cost=meta.time_cost, + started_at=started_at, + ended_at=utc_now(), + ) + ) # transform tool invoke message to get LLM friendly message return plain_text, message_files, meta diff --git a/api/core/workflow/node_factory.py b/api/core/workflow/node_factory.py index 15d838b4a6e740..e36845875347ac 100644 --- a/api/core/workflow/node_factory.py +++ b/api/core/workflow/node_factory.py @@ -50,7 +50,6 @@ from core.workflow.nodes.agent.runtime_support import AgentRuntimeSupport from core.workflow.nodes.agent_v2 import DifyAgentNode from core.workflow.nodes.agent_v2.binding_resolver import WorkflowAgentBindingResolver -from core.workflow.nodes.agent_v2.discriminator import is_dify_agent_node_data from core.workflow.nodes.agent_v2.output_adapter import WorkflowAgentOutputAdapter from core.workflow.nodes.agent_v2.runtime_request_builder import WorkflowAgentRuntimeRequestBuilder from core.workflow.nodes.human_input.callback import DifyHITLCallback @@ -59,6 +58,7 @@ from core.workflow.system_variables import SystemVariableKey, get_system_text, system_variable_selector from core.workflow.template_rendering import CodeExecutorJinja2TemplateRenderer from core.workflow.workflow_tool_node import DifyWorkflowToolNode +from domain.workflow.agent_node import is_dify_agent_node_data from extensions.application_services.data_sources import build_data_source_credentials from graphon.entities.base_node_data import BaseNodeData from graphon.entities.graph_config import NodeConfigDict, NodeConfigDictAdapter diff --git a/api/core/workflow/nodes/agent_v2/runtime_request_builder.py b/api/core/workflow/nodes/agent_v2/runtime_request_builder.py index b94300f3cbb0b3..5cb9f220333a7c 100644 --- a/api/core/workflow/nodes/agent_v2/runtime_request_builder.py +++ b/api/core/workflow/nodes/agent_v2/runtime_request_builder.py @@ -49,9 +49,9 @@ from configs import dify_config from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom from core.app.llm.model_access import resolve_model_context_window -from core.ops.base_trace_instance import application_tracing_hooks from core.plugin.provider_identity import normalize_plugin_daemon_provider_identity from core.workflow.system_variables import SystemVariableKey, get_system_text, get_system_value +from extensions.tracing import application_tracing_hooks from graphon.file import File, FileTransferMethod from graphon.variables.segments import Segment from graphon.variables.template_resolution import convert_template diff --git a/api/core/workflow/nodes/agent_v2/validators.py b/api/core/workflow/nodes/agent_v2/validators.py index c92a3be9bb362a..82fa82cdaf4d40 100644 --- a/api/core/workflow/nodes/agent_v2/validators.py +++ b/api/core/workflow/nodes/agent_v2/validators.py @@ -12,6 +12,7 @@ workflow_previous_node_output_refs_from_selectors, ) from core.workflow.graph_topology import WorkflowGraphTopology +from domain.workflow.agent_node import is_dify_agent_node_data from graphon.enums import BuiltinNodeTypes, ErrorStrategy from models.agent import ( Agent, @@ -32,7 +33,6 @@ from models.workflow import Workflow from services.agent.knowledge_datasets import list_missing_tenant_knowledge_dataset_ids -from .discriminator import is_dify_agent_node_data from .entities import DifyAgentNodeData diff --git a/api/docker/entrypoint.sh b/api/docker/entrypoint.sh index 51dbe8332ac64a..f99aaf73796086 100755 --- a/api/docker/entrypoint.sh +++ b/api/docker/entrypoint.sh @@ -35,10 +35,10 @@ if [[ "${MODE}" == "worker" ]]; then if [[ -z "${CELERY_QUEUES}" ]]; then if [[ "${DEPLOYMENT_EDITION:-COMMUNITY}" == "CLOUD" ]]; then # Cloud edition: separate queues for dataset and trigger tasks - DEFAULT_QUEUES="api_token,dataset,dataset_summary,priority_dataset,priority_pipeline,pipeline,mail,human_input_delivery,human_input_contact_sync,ops_trace,app_deletion,app_rbac,plugin,workflow_storage,conversation,workflow_professional,workflow_team,workflow_sandbox,schedule_poller,schedule_executor,triggered_workflow_dispatcher,trigger_refresh_publisher,trigger_refresh_executor,retention,workflow_based_app_execution,knowledge_fs_lifecycle" + DEFAULT_QUEUES="api_token,dataset,dataset_summary,priority_dataset,priority_pipeline,pipeline,mail,human_input_delivery,human_input_contact_sync,${OPS_TRACE_QUEUE:-ops_trace_v2},app_deletion,app_rbac,plugin,workflow_storage,conversation,workflow_professional,workflow_team,workflow_sandbox,schedule_poller,schedule_executor,triggered_workflow_dispatcher,trigger_refresh_publisher,trigger_refresh_executor,retention,workflow_based_app_execution,knowledge_fs_lifecycle" else # Self-hosted editions: dataset, pipeline and workflow have separate queues - DEFAULT_QUEUES="api_token,dataset,dataset_summary,priority_dataset,priority_pipeline,pipeline,mail,human_input_delivery,human_input_contact_sync,ops_trace,app_deletion,app_rbac,plugin,workflow_storage,conversation,workflow,schedule_poller,schedule_executor,triggered_workflow_dispatcher,trigger_refresh_publisher,trigger_refresh_executor,retention,workflow_based_app_execution,knowledge_fs_lifecycle" + DEFAULT_QUEUES="api_token,dataset,dataset_summary,priority_dataset,priority_pipeline,pipeline,mail,human_input_delivery,human_input_contact_sync,${OPS_TRACE_QUEUE:-ops_trace_v2},app_deletion,app_rbac,plugin,workflow_storage,conversation,workflow,schedule_poller,schedule_executor,triggered_workflow_dispatcher,trigger_refresh_publisher,trigger_refresh_executor,retention,workflow_based_app_execution,knowledge_fs_lifecycle" fi else DEFAULT_QUEUES="${CELERY_QUEUES}" diff --git a/api/core/ops/__init__.py b/api/domain/__init__.py similarity index 100% rename from api/core/ops/__init__.py rename to api/domain/__init__.py diff --git a/api/core/ops/entities/__init__.py b/api/domain/tracing/__init__.py similarity index 100% rename from api/core/ops/entities/__init__.py rename to api/domain/tracing/__init__.py diff --git a/api/domain/tracing/errors.py b/api/domain/tracing/errors.py new file mode 100644 index 00000000000000..72ce44beb2e3b6 --- /dev/null +++ b/api/domain/tracing/errors.py @@ -0,0 +1,19 @@ +"""Business failures of trace runs; ``code`` names the reason clients and run records show.""" + + +class TraceRunError(Exception): + def __init__(self, code: str) -> None: + self.code = code + super().__init__(code) + + +class TraceRunNotFoundError(TraceRunError): + """The run collected nothing, or a batch names an invocation the run does not have.""" + + +class TraceRunConflictError(TraceRunError): + """The operation conflicts with the run's current state.""" + + +class TraceRunLimitError(TraceRunError): + """A runtime batch exceeds the run's collection limits.""" diff --git a/api/services/tracing/normalization.py b/api/domain/tracing/normalization.py similarity index 52% rename from api/services/tracing/normalization.py rename to api/domain/tracing/normalization.py index 0216e12431e29b..aba79110d456ff 100644 --- a/api/services/tracing/normalization.py +++ b/api/domain/tracing/normalization.py @@ -1,12 +1,12 @@ -"""Normalize Agent SDK observations without conflating business and provider IDs.""" +"""Normalize Agent runtime observations and outer traces for the Agent trace collector.""" -import hashlib import json +import math from collections.abc import Mapping, Sequence from datetime import UTC, datetime from typing import Any -from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace +from dify_trace_sdk.model import Cost, Span, SpanKind, SpanStatus, TokenUsage, Trace MAX_VALUE_CHARS = 16000 _SECRET_KEYS = frozenset( @@ -41,12 +41,6 @@ ) -def span_id(logical_id: str) -> str: - """Stable nonzero OTel span identity for Dify's persisted logical ID.""" - value = hashlib.sha256(f"dify-application-span:{logical_id}".encode()).hexdigest()[:16] - return value if int(value, 16) else "0000000000000001" - - def sanitize(value: Any, *, depth: int = 0) -> Any: if depth > 12: return "[truncated]" @@ -67,6 +61,8 @@ def sanitize(value: Any, *, depth: int = 0) -> Any: if isinstance(decoded, (dict, list)): return sanitize(decoded, depth=depth + 1) return value[:MAX_VALUE_CHARS] + ("[truncated]" if len(value) > MAX_VALUE_CHARS else "") + if isinstance(value, float) and not math.isfinite(value): + return str(value) if value is None or isinstance(value, (int, float, bool)): return value return str(value)[:MAX_VALUE_CHARS] @@ -143,48 +139,68 @@ def _first(attributes: Mapping[str, Any], keys: Sequence[str]) -> Any: return None -def normalize_agent_span(data: Mapping[str, Any], *, capture_content: bool) -> CanonicalSpan: - attributes = safe_attributes(data.get("attributes") or {}, capture_content=capture_content) +def _count(value: object) -> int | None: + return int(value) if isinstance(value, int | float) and not isinstance(value, bool) and value >= 0 else None + + +def _amount(value: object) -> float | None: + """A finite, non-negative price; the Agent runtime reports its decimal price as a string.""" + if isinstance(value, bool) or not isinstance(value, int | float | str): + return None + try: + amount = float(value) + except ValueError: + return None + return amount if math.isfinite(amount) and amount >= 0 else None + + +def _agent_kind(attributes: Mapping[str, Any]) -> SpanKind: operation = str(attributes.get("gen_ai.operation.name", "")) declared = str(attributes.get("openinference.span.kind", "")).lower() if operation in {"invoke_agent", "agent"} or declared == "agent": - kind = CanonicalSpanKind.AGENT - elif operation in {"execute_tool", "tool"} or "gen_ai.tool.name" in attributes or declared == "tool": - kind = CanonicalSpanKind.TOOL - elif ( + return SpanKind.AGENT + if operation in {"execute_tool", "tool"} or "gen_ai.tool.name" in attributes or declared == "tool": + return SpanKind.TOOL + if ( operation in {"chat", "text_completion", "generate_content", "embeddings"} or "gen_ai.request.model" in attributes or declared == "llm" ): - kind = CanonicalSpanKind.LLM - elif declared == "retriever": - kind = CanonicalSpanKind.RETRIEVER - else: - kind = CanonicalSpanKind.CHAIN - metadata = {key: value for key, value in attributes.items() if key not in _CONTENT_ATTRIBUTES} - mappings = { - "model_name": ("gen_ai.response.model", "gen_ai.request.model", "llm.model_name"), - "model_provider": ("gen_ai.provider.name", "gen_ai.system"), - "prompt_tokens": ("gen_ai.usage.input_tokens", "llm.token_count.prompt"), - "completion_tokens": ("gen_ai.usage.output_tokens", "llm.token_count.completion"), - "cache_read_tokens": ("gen_ai.usage.cache_read.input_tokens", "gen_ai.usage.input_tokens.cached"), - "total_price": ("dify.llm.total_price",), - "currency": ("dify.llm.currency",), - "time_to_first_token": ("gen_ai.response.time_to_first_token",), - } - for key, candidates in mappings.items(): - value = _first(attributes, candidates) - if value is not None: - metadata[key] = value - if kind == CanonicalSpanKind.LLM: - input_tokens, output_tokens = metadata.get("prompt_tokens"), metadata.get("completion_tokens") - if isinstance(input_tokens, (int, float)) and isinstance(output_tokens, (int, float)): - metadata["total_tokens"] = input_tokens + output_tokens + return SpanKind.LLM + if declared == "retriever": + return SpanKind.RETRIEVER + return SpanKind.CHAIN + + +def normalize_agent_span(data: Mapping[str, Any], *, capture_content: bool) -> Span: + attributes = safe_attributes(data.get("attributes") or {}, capture_content=capture_content) + kind = _agent_kind(attributes) + metadata: dict[str, Any] = {key: value for key, value in attributes.items() if key not in _CONTENT_ATTRIBUTES} + model_name = _first(attributes, ("gen_ai.response.model", "gen_ai.request.model", "llm.model_name")) + model_provider = _first(attributes, ("gen_ai.provider.name", "gen_ai.system")) + prompt_tokens = _count(_first(attributes, ("gen_ai.usage.input_tokens", "llm.token_count.prompt"))) + completion_tokens = _count(_first(attributes, ("gen_ai.usage.output_tokens", "llm.token_count.completion"))) + cache_read_tokens = _first(attributes, ("gen_ai.usage.cache_read.input_tokens", "gen_ai.usage.input_tokens.cached")) + if cache_read_tokens is not None: + metadata["cache_read_tokens"] = cache_read_tokens + price = _amount(_first(attributes, ("dify.llm.total_price",))) + currency = _first(attributes, ("dify.llm.currency",)) + time_to_first_token = _first(attributes, ("gen_ai.response.time_to_first_token",)) + usage = cost = None + if kind is SpanKind.LLM: + if prompt_tokens is not None or completion_tokens is not None: + usage = TokenUsage( + prompt_tokens=prompt_tokens or 0, + completion_tokens=completion_tokens or 0, + total_tokens=(prompt_tokens or 0) + (completion_tokens or 0), + ) + if price is not None: + cost = Cost(amount=price, currency=str(currency or "USD")) else: - # Agent/chain aggregate usage is useful context, never another billed generation. - for key in ("prompt_tokens", "completion_tokens", "total_tokens", "total_price"): - if key in metadata: - metadata[f"aggregate_{key}"] = metadata.pop(key) + # Agent and chain aggregates are context, never another billed generation. + for key, value in (("prompt_tokens", prompt_tokens), ("completion_tokens", completion_tokens)): + if value is not None: + metadata[f"aggregate_{key}"] = value error = str(data.get("status_message") or "Agent operation failed") if data.get("status") == "error" else None if error and not capture_content: error = "Agent operation failed" @@ -218,26 +234,30 @@ def normalize_agent_span(data: Mapping[str, Any], *, capture_content: bool) -> C if capture_content else None ) - return CanonicalSpan( - id=str(data["span_id"]), - parent_id=data.get("parent_span_id"), + return Span( + span_id=str(data["span_id"]), + parent_span_id=data.get("parent_span_id"), name=str(data["name"]), kind=kind, start_time=datetime.fromtimestamp(int(data["start_time_ns"]) / 1e9, UTC), end_time=datetime.fromtimestamp(int(data["end_time_ns"]) / 1e9, UTC), + status=SpanStatus.ERROR if error else SpanStatus.OK, + error=error, inputs=sanitize(_decoded(inputs)), outputs=sanitize(_decoded(outputs)), - status=CanonicalSpanStatus.ERROR if error else CanonicalSpanStatus.OK, - error=error, - metadata=metadata, + model_name=str(model_name) if kind is SpanKind.LLM and model_name else None, + model_provider=str(model_provider) if kind is SpanKind.LLM and model_provider else None, + usage=usage, + cost=cost, + time_to_first_token=float(time_to_first_token) + if kind is SpanKind.LLM and isinstance(time_to_first_token, int | float) and time_to_first_token >= 0 + else None, + attributes=sanitize(metadata), ) -def remap_outer_trace( - trace: CanonicalTrace, *, trace_id: str, capture_content: bool, agent_message: bool -) -> CanonicalTrace: - """Use the same IDs allocated before runtime and suppress synthetic Agent usage.""" - safe_metadata = { +_SAFE_OUTER_ATTRIBUTES = frozenset( + { "workflow_id", "workflow_run_id", "workflow_app_log_id", @@ -245,13 +265,6 @@ def remap_outer_trace( "node_execution_id", "node_type", "status", - "model_provider", - "model_name", - "prompt_tokens", - "completion_tokens", - "total_tokens", - "total_price", - "currency", "trace_entity_type", "wrapper_type", "wrapper_index", @@ -263,64 +276,107 @@ def remap_outer_trace( "conversation_id", "external_trace_id", "message_execution_state", + "parent_workflow_run_id", + "parent_node_execution_id", "linked_parent_workflow_run_id", "linked_parent_node_execution_id", "linked_parent_application_trace_id", "linked_parent_trace_id", } - spans = [] +) + + +def remap_outer_trace( + trace: Trace, *, trace_id: str, capture_content: bool, synthetic_span_ids: frozenset[str] +) -> Trace: + """Adopt the run's trace ID and capture policy; drop spans the Agent runtime reports itself.""" + spans: list[Span] = [] for original in trace.spans: - if agent_message and original.synthetic and original.kind == CanonicalSpanKind.LLM: + if original.span_id in synthetic_span_ids: continue - metadata = ( - sanitize(original.metadata) - if capture_content - else {k: sanitize(v) for k, v in original.metadata.items() if k in safe_metadata} - ) - if trace.trace_id != trace_id: - metadata.setdefault("external_trace_id", trace.trace_id) - if trace.external_parent is not None: - # This pipeline exports a complete local tree. Cross-workflow - # correlation is explicit until a durable cross-run merge exists. - metadata["linked_parent_workflow_run_id"] = trace.external_parent.parent_workflow_run_id - if trace.external_parent.parent_node_execution_id: - metadata["linked_parent_node_execution_id"] = trace.external_parent.parent_node_execution_id - metadata.pop("parent_trace_context", None) - if original.kind != CanonicalSpanKind.LLM: - for key in ("prompt_tokens", "completion_tokens", "total_tokens", "total_price"): - if key in metadata: - metadata[f"aggregate_{key}"] = metadata.pop(key) + attributes = dict(original.attributes) + if trace.external_trace_id: + attributes.setdefault("external_trace_id", trace.external_trace_id) + if not capture_content: + attributes = {key: value for key, value in attributes.items() if key in _SAFE_OUTER_ATTRIBUTES} spans.append( original.model_copy( update={ - "id": span_id(original.id), - "parent_id": span_id(original.parent_id) if original.parent_id else None, "inputs": sanitize(original.inputs) if capture_content else None, "outputs": sanitize(original.outputs) if capture_content else None, - "metadata": metadata, + "attributes": sanitize(attributes), "error": original.error if capture_content or not original.error else "Application execution failed", - "can_parent_workflow": False, - "publishes_parent_context": False, } ) ) - return CanonicalTrace( - trace_id=trace_id, session_id=trace.session_id, root_span_id=span_id(trace.root_span_id), spans=tuple(spans) + return Trace( + trace_id=trace_id, + external_trace_id=None, + session_id=trace.session_id, + user_id=trace.user_id, + parent=None, + spans=tuple(spans), + ) + + +def legacy_outer_trace(data: Mapping[str, Any]) -> Trace: + """Read an outer trace in the shape the previous release stored, so runs captured before an + upgrade still export. Its spans already carry this ledger's identifiers and capture policy.""" + spans: list[Span] = [] + for raw in data["spans"]: + metadata = raw.get("metadata") or {} + llm = raw["kind"] == SpanKind.LLM.value + prompt, completion = _count(metadata.get("prompt_tokens")), _count(metadata.get("completion_tokens")) + amount, currency = _amount(metadata.get("total_price")), metadata.get("currency") + spans.append( + Span( + span_id=raw["id"], + parent_span_id=raw["parent_id"], + name=raw["name"], + kind=SpanKind(raw["kind"]), + start_time=raw["start_time"], + end_time=raw["end_time"] or raw["start_time"], + status=SpanStatus(raw["status"]), + error=raw.get("error"), + inputs=sanitize(raw.get("inputs")), + outputs=sanitize(raw.get("outputs")), + model_name=metadata.get("model_name") if llm else None, + model_provider=metadata.get("model_provider") if llm else None, + usage=TokenUsage( + prompt_tokens=prompt or 0, + completion_tokens=completion or 0, + total_tokens=_count(metadata.get("total_tokens")) or (prompt or 0) + (completion or 0), + ) + if llm and (prompt is not None or completion is not None) + else None, + cost=Cost(amount=amount, currency=currency) + if llm and amount is not None and isinstance(currency, str) and currency + else None, + attributes=sanitize(metadata), + ) + ) + return Trace( + trace_id=data["trace_id"], + external_trace_id=None, + session_id=data.get("session_id") or None, + user_id=None, + parent=None, + spans=tuple(spans), ) -def order_spans(root_id: str, spans: Sequence[CanonicalSpan]) -> tuple[CanonicalSpan, ...]: - by_id: dict[str, CanonicalSpan] = {} +def order_spans(root_id: str, spans: Sequence[Span]) -> tuple[Span, ...]: + by_id: dict[str, Span] = {} for span in spans: - prior = by_id.get(span.id) + prior = by_id.get(span.span_id) if prior is not None and prior != span: raise ValueError("Conflicting trace span identity") - by_id[span.id] = span + by_id[span.span_id] = span if root_id not in by_id: raise ValueError("Trace root is missing") - ordered: list[CanonicalSpan] = [] + ordered: list[Span] = [] complete: set[str] = set() visiting: set[str] = set() @@ -333,14 +389,15 @@ def visit(identifier: str) -> None: if span is None: raise ValueError("Trace parent is missing") visiting.add(identifier) - if span.parent_id is not None: - visit(span.parent_id) - elif span.id != root_id: + if span.parent_span_id is not None: + visit(span.parent_span_id) + elif span.span_id != root_id: raise ValueError("Unexpected trace root") visiting.remove(identifier) complete.add(identifier) ordered.append(span) + visit(root_id) for span in spans: - visit(span.id) + visit(span.span_id) return tuple(ordered) diff --git a/api/domain/tracing/sources.py b/api/domain/tracing/sources.py new file mode 100644 index 00000000000000..61ffc41f758f51 --- /dev/null +++ b/api/domain/tracing/sources.py @@ -0,0 +1,49 @@ +"""Committed facts about the records a trace run is collected for.""" + +from dataclasses import dataclass +from datetime import datetime + + +@dataclass(frozen=True, slots=True) +class SourceWorkflowRun: + id: str + workflow_id: str + graph: str | None + """The graph snapshot of the run; ``None`` when the backend does not store it.""" + status: str + finished_at: datetime | None + """``None`` while the run has not finished.""" + created_by: str + + +@dataclass(frozen=True, slots=True) +class SourceNodeExecution: + id: str + node_id: str + workflow_id: str | None + """``None`` when the backend does not store the workflow.""" + workflow_run_id: str | None + """``None`` for a single-node preview.""" + title: str + status: str + error: str | None + inputs: str | None + outputs: str | None + created_at: datetime + elapsed_time: float + + +@dataclass(frozen=True, slots=True) +class MessageProgress: + """What recovery may infer about a message from its committed row.""" + + message_id: str + conversation_id: str + agent_app: bool + paused: bool + explicit_error: bool + """The message failed and recorded why.""" + answer_saved: bool + """The final answer transaction committed, not just an early usage flush.""" + agent_stop_requested: bool + updated_at: datetime diff --git a/api/domain/tracing/terminal_recovery.py b/api/domain/tracing/terminal_recovery.py new file mode 100644 index 00000000000000..64dd46b6dcd281 --- /dev/null +++ b/api/domain/tracing/terminal_recovery.py @@ -0,0 +1,65 @@ +"""Decide when a run's dropped terminal event can be recovered from its source's committed state. + +The decision combines the run's invocations with what its source committed. +``now`` is naive UTC, like the run's timestamps. +""" + +from datetime import datetime, timedelta + +from domain.tracing.errors import TraceRunError +from domain.tracing.sources import MessageProgress, SourceWorkflowRun +from domain.tracing.trace_run_agg import TraceRunAgg +from graphon.enums import WorkflowExecutionStatus + +_SETTLE_SECONDS = 5 +_WORKFLOW_TERMINAL = { + WorkflowExecutionStatus.SUCCEEDED, + WorkflowExecutionStatus.PARTIAL_SUCCEEDED, + WorkflowExecutionStatus.FAILED, + WorkflowExecutionStatus.STOPPED, +} + + +def message_terminal_due( + run: TraceRunAgg, progress: MessageProgress, *, now: datetime, collection_timeout: int +) -> bool: + """Whether the message's committed row is the run's business terminal now. + + Raises: + TraceRunError: The terminal can no longer arrive. + """ + invocations = list(run.invocations.values()) + pending_agent_stop = progress.agent_app and progress.agent_stop_requested + if progress.paused and not pending_agent_stop: + return False + if any(not invocation.finalized or invocation.finalized_at is None for invocation in invocations): + if (progress.explicit_error or pending_agent_stop) and progress.updated_at + timedelta( + seconds=collection_timeout + ) < now: + raise TraceRunError("agent_trace_collection_incomplete") + return False + finalized = [invocation.finalized_at for invocation in invocations if invocation.finalized_at is not None] + anchor = max([progress.updated_at, *finalized]) + if pending_agent_stop: + # Even a fully saved partial answer/usage cannot establish which + # runtime terminal won the stop race. Wait for the existing worker + # handoff, or converge to incomplete delivery after its deadline. + if anchor + timedelta(seconds=collection_timeout) < now: + raise TraceRunError("business_terminal_not_persisted") + return False + if run.capture_agent and not invocations: + return False + if anchor + timedelta(seconds=_SETTLE_SECONDS) > now: + return False + if not (progress.explicit_error or progress.answer_saved): + if finalized and max(finalized) + timedelta(seconds=collection_timeout) < now: + raise TraceRunError("business_terminal_not_persisted") + return False + return True + + +def workflow_terminal_due(workflow: SourceWorkflowRun, *, now: datetime) -> bool: + """Whether the workflow run finished long enough ago for its node executions to be saved.""" + if workflow.status not in _WORKFLOW_TERMINAL or workflow.finished_at is None: + return False + return workflow.finished_at + timedelta(seconds=_SETTLE_SECONDS) <= now diff --git a/api/domain/tracing/trace_run_agg.py b/api/domain/tracing/trace_run_agg.py new file mode 100644 index 00000000000000..edacb58cf979bb --- /dev/null +++ b/api/domain/tracing/trace_run_agg.py @@ -0,0 +1,642 @@ +"""The Agent trace run aggregate. + +A run is the trace of one Agent-managed execution: the destination and capture +policy frozen when the execution started, the runtime invocations that report +spans into it, its business terminal, and its export lifecycle. Callers pass +the time, generated identifiers, and the facts a rule needs. +""" + +import hashlib +import json +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field, replace +from datetime import UTC, datetime, timedelta +from enum import StrEnum +from typing import Any, Literal +from uuid import UUID, uuid5 + +from dify_agent.protocol.tracing import AgentTraceBatch, AgentTraceSpan +from dify_trace_sdk.ids import span_id_for +from dify_trace_sdk.model import Span, SpanKind, SpanStatus, Trace + +from domain.tracing.errors import TraceRunConflictError, TraceRunError, TraceRunLimitError, TraceRunNotFoundError +from domain.tracing.normalization import ( + legacy_outer_trace, + normalize_agent_span, + order_spans, + remap_outer_trace, + safe_attributes, + sanitize, +) +from domain.tracing.sources import SourceNodeExecution + +type TraceSourceKind = Literal["message", "workflow", "node"] + + +class TraceRunState(StrEnum): + NOT_COLLECTED = "not_collected" + """The app did not trace when the execution started; nothing is collected or exported.""" + COLLECTING = "collecting" + READY = "ready" + """Released by a manual retry; exported once its collection is complete.""" + EXPORTING = "exporting" + RETRYING = "retrying" + ACCEPTED = "accepted" + FAILED = "failed" + UNCERTAIN = "uncertain" + """The destination may have accepted the trace; only an acknowledged retry sends it again.""" + + @property + def accepts_terminal(self) -> bool: + """A run that collects nothing or is closed keeps no business terminal.""" + return self in {TraceRunState.COLLECTING, TraceRunState.READY, TraceRunState.RETRYING} + + +# A closed run accepts no runtime batches and no other business terminal. +_CLOSED_STATES = frozenset( + {TraceRunState.ACCEPTED, TraceRunState.EXPORTING, TraceRunState.UNCERTAIN, TraceRunState.FAILED} +) +_COLLECTING_STATES = frozenset({TraceRunState.COLLECTING, TraceRunState.READY}) +_SETTLED_STATES = frozenset( + {TraceRunState.NOT_COLLECTED, TraceRunState.ACCEPTED, TraceRunState.FAILED, TraceRunState.UNCERTAIN} +) +_UNFINISHED_NODE_STATUSES = frozenset({"running", "pending", "paused"}) +_SESSION_ID_CHARS = 255 + + +@dataclass(frozen=True, slots=True) +class TraceInvocation: + """One Agent runtime invocation that reports spans into a run.""" + + id: str + run_id: str + agent_id: str + parent_span_id: str + span_id: str + node_execution_id: str | None + """``None`` for the invocation of an Agent app.""" + attempt: int + finalized: bool + finalized_at: datetime | None + """``None`` until the runtime reports the invocation's final batch.""" + + +@dataclass(frozen=True, slots=True) +class TraceBatchReceipt: + """What a run remembers of a stored batch to recognize its replays.""" + + id: str + invocation_id: str + payload_hash: str + """The SHA-256 of the batch's original payload.""" + + +@dataclass(frozen=True, slots=True) +class TraceBatch: + """A runtime batch as a run stores it; its spans already follow the run's capture policy.""" + + id: str + invocation_id: str + payload_hash: str + spans: Sequence[Mapping[str, Any]] + final: bool + + +@dataclass(kw_only=True, eq=False, slots=True) +class TraceRunAgg: + """An Agent trace run, its invocations, and the receipts of its stored batches. + + State changes only through the methods below. ``invocations`` keeps them by + attempt, then creation. ``new_batches`` holds the batches ingested since the + run was loaded, for the repository to store. + """ + + id: str + tenant_id: str + app_id: str + source_kind: TraceSourceKind + source_id: str + provider: str + """``""`` when the run froze no destination.""" + config_revision: str + """``""`` when the run froze no destination.""" + provider_config: Mapping[str, Any] + """The frozen stored configuration, secrets still encrypted.""" + trace_id: str + root_span_id: str + session_id: str + capture_content: bool + capture_agent: bool + state: TraceRunState + outer_trace: Mapping[str, Any] | None + """The captured business terminal; ``None`` until it is captured.""" + span_count: int + payload_bytes: int + attempts: int + error_code: str | None + """``None`` unless the run failed, waits for a retry, or its export outcome is uncertain.""" + trace_url: str | None + provider_trace_id: str | None + span_urls: Mapping[str, str] + export_token: str | None + """The token fencing the current export attempt; ``None`` when no attempt holds the run.""" + export_started_at: datetime | None + export_heartbeat_at: datetime | None + """``None`` for attempts started before heartbeats were recorded.""" + next_attempt_at: datetime | None + """``None`` unless a retryable export failure waits for its backoff.""" + accepted_at: datetime | None + terminal_at: datetime | None + """When the business terminal was captured; ``None`` until then.""" + updated_at: datetime + invocations: dict[str, TraceInvocation] + receipts: dict[str, TraceBatchReceipt] + new_batches: list[TraceBatch] = field(default_factory=list, init=False) + + @classmethod + def freeze( + cls, + *, + run_id: str, + trace_id: str, + tenant_id: str, + app_id: str, + source_kind: TraceSourceKind, + source_id: str, + session_id: str, + policy: Mapping[str, Any], + stored_config: Mapping[str, Any] | None, + now: datetime, + ) -> "TraceRunAgg": + """Freeze the app's destination and capture policy for a source. + + A run collects only when the policy enables the provider it names and + that provider has a stored configuration; otherwise it is frozen as not + collected. The revision identifies the frozen policy and configuration. + """ + named = policy.get("tracing_provider") + destination = ( + (named, dict(stored_config)) + if isinstance(named, str) and policy.get("enabled") is True and stored_config is not None + else None + ) + provider, frozen = destination if destination is not None else ("", {}) + collecting = destination is not None + revision = hashlib.sha256( + json.dumps({"config": frozen, "policy": dict(policy)}, sort_keys=True, default=str).encode() + ).hexdigest() + return cls._new( + run_id=run_id, + trace_id=trace_id, + tenant_id=tenant_id, + app_id=app_id, + source_kind=source_kind, + source_id=source_id, + provider=provider, + config_revision=revision, + provider_config=frozen, + session_id=session_id[:_SESSION_ID_CHARS] or source_id, + capture_content=collecting and policy.get("capture_content", False) is True, + capture_agent=collecting and policy.get("capture_agent", True) is True, + state=TraceRunState.COLLECTING if collecting else TraceRunState.NOT_COLLECTED, + error_code=None, + terminal_at=None, + now=now, + ) + + @classmethod + def initialization_failed( + cls, + *, + run_id: str, + trace_id: str, + tenant_id: str, + app_id: str, + source_kind: TraceSourceKind, + source_id: str, + now: datetime, + ) -> "TraceRunAgg": + """Close a source whose start snapshot is missing. + + It must never fall back to an exporter without its capture policy. + """ + return cls._new( + run_id=run_id, + trace_id=trace_id, + tenant_id=tenant_id, + app_id=app_id, + source_kind=source_kind, + source_id=source_id, + provider="", + config_revision="", + provider_config={}, + session_id=source_id, + capture_content=False, + capture_agent=False, + state=TraceRunState.FAILED, + error_code="tracing_initialization_failed", + terminal_at=now, + now=now, + ) + + @classmethod + def _new( + cls, + *, + run_id: str, + trace_id: str, + tenant_id: str, + app_id: str, + source_kind: TraceSourceKind, + source_id: str, + provider: str, + config_revision: str, + provider_config: Mapping[str, Any], + session_id: str, + capture_content: bool, + capture_agent: bool, + state: TraceRunState, + error_code: str | None, + terminal_at: datetime | None, + now: datetime, + ) -> "TraceRunAgg": + return cls( + id=run_id, + tenant_id=tenant_id, + app_id=app_id, + source_kind=source_kind, + source_id=source_id, + provider=provider, + config_revision=config_revision, + provider_config=provider_config, + trace_id=trace_id, + root_span_id=span_id_for(source_id), + session_id=session_id, + capture_content=capture_content, + capture_agent=capture_agent, + state=state, + outer_trace=None, + span_count=0, + payload_bytes=0, + attempts=0, + error_code=error_code, + trace_url=None, + provider_trace_id=None, + span_urls={}, + export_token=None, + export_started_at=None, + export_heartbeat_at=None, + next_attempt_at=None, + accepted_at=None, + terminal_at=terminal_at, + updated_at=now, + invocations={}, + receipts={}, + ) + + @property + def awaits_terminal(self) -> bool: + """The run collects and has not captured its business terminal yet.""" + return self.outer_trace is None and self.state in _COLLECTING_STATES + + def open_invocation( + self, *, agent_id: str, node_execution_id: str | None, attempt: int, call_key: str + ) -> TraceInvocation: + """Create the invocation of a call key once; a repeated call key must name the same Agent binding. + + Raises: + TraceRunConflictError: The run stopped collecting, or the call key names another invocation. + """ + if self.state not in _COLLECTING_STATES: + raise TraceRunConflictError("tracing_run_not_collecting") + identifier = str(uuid5(UUID(self.id), call_key)) + current = self.invocations.get(identifier) + if current is not None: + if (current.agent_id, current.node_execution_id, current.attempt) != (agent_id, node_execution_id, attempt): + raise TraceRunConflictError("tracing_invocation_conflict") + return current + invocation = TraceInvocation( + id=identifier, + run_id=self.id, + agent_id=agent_id, + parent_span_id=span_id_for(node_execution_id) if node_execution_id else self.root_span_id, + span_id=span_id_for(f"agent-attempt:{identifier}"), + node_execution_id=node_execution_id, + attempt=attempt, + finalized=False, + finalized_at=None, + ) + self.invocations[identifier] = invocation + return invocation + + def ingest(self, batch: AgentTraceBatch, *, max_spans: int, max_bytes: int, now: datetime) -> int: + """Append a runtime batch and return its span count; a final batch finalizes its invocation. + + A batch ID replays only with the same validated payload, even after the run closed. + + Raises: + TraceRunNotFoundError: The batch's invocation is not one of this run's. + TraceRunConflictError: The batch conflicts with a stored one, or the run or invocation stopped collecting. + TraceRunLimitError: The batch exceeds the run's limits. + """ + # Hash the original validated payload. Sanitization cannot turn two + # conflicting producer payloads into a falsely successful replay. + canonical = json.dumps(batch.model_dump(mode="json"), sort_keys=True, separators=(",", ":")).encode() + digest = hashlib.sha256(canonical).hexdigest() + invocation = self.invocations.get(str(batch.invocation_id)) + if invocation is None: + raise TraceRunNotFoundError("tracing_invocation_not_found") + batch_id = str(batch.batch_id) + spans = batch.spans + previous = self.receipts.get(batch_id) + if previous is not None: + if previous.invocation_id != invocation.id or previous.payload_hash != digest: + raise TraceRunConflictError("tracing_batch_conflict") + return len(spans) + if self.state in _CLOSED_STATES or invocation.finalized: + raise TraceRunConflictError("tracing_collection_closed") + if any(span.trace_id != self.trace_id for span in spans): + raise TraceRunConflictError("tracing_trace_id_mismatch") + if self.span_count + len(spans) > max_spans or self.payload_bytes + len(canonical) > max_bytes: + raise TraceRunLimitError("tracing_payload_limit") + final = batch.final + self.new_batches.append( + TraceBatch( + id=batch_id, + invocation_id=invocation.id, + payload_hash=digest, + spans=[self._stored_span(span) for span in spans], + final=final, + ) + ) + self.receipts[batch_id] = TraceBatchReceipt(id=batch_id, invocation_id=invocation.id, payload_hash=digest) + self.span_count += len(spans) + self.payload_bytes += len(canonical) + self.updated_at = now + if final: + self.invocations[invocation.id] = replace(invocation, finalized=True, finalized_at=now) + return len(spans) + + def _stored_span(self, span: AgentTraceSpan) -> dict[str, Any]: + stored = span.model_dump(mode="json") + stored["attributes"] = safe_attributes(span.attributes, capture_content=self.capture_content) + if not self.capture_content: + stored["events"] = [] + stored["status_message"] = "Agent operation failed" if span.status == "error" else None + return stored + + def capture_terminal(self, trace: Trace, *, now: datetime) -> None: + """Keep the execution's business terminal under the run's trace ID, session, and capture policy. + + The Agent runtime reports a message's LLM work itself, so the + terminal's own LLM span is dropped. + """ + synthetic = frozenset({span_id_for(f"{self.source_id}:llm")}) if self.source_kind == "message" else frozenset() + adopted = remap_outer_trace( + trace, trace_id=self.trace_id, capture_content=self.capture_content, synthetic_span_ids=synthetic + ) + self._keep_terminal(adopted.model_copy(update={"session_id": self.session_id}), now=now) + + def settle_preview(self, node: SourceNodeExecution, *, now: datetime) -> None: + """Make a finished single-node preview its own terminal; a preview reports no business terminal.""" + if self.source_kind != "node" or self.outer_trace is not None or node.status in _UNFINISHED_NODE_STATUSES: + return + error = node.error + started_at = node.created_at.replace(tzinfo=UTC) + span = Span( + span_id=self.root_span_id, + parent_span_id=None, + name=node.title or "Agent preview", + kind=SpanKind.AGENT, + start_time=started_at, + end_time=started_at + timedelta(seconds=node.elapsed_time or 0), + status=SpanStatus.ERROR if error else SpanStatus.OK, + error=error if self.capture_content else "Node execution failed" if error else None, + inputs=sanitize(node.inputs) if self.capture_content else None, + outputs=sanitize(node.outputs) if self.capture_content else None, + attributes={"node_execution_id": node.id, "node_id": node.node_id, "source": "single_step"}, + ) + trace = Trace( + trace_id=self.trace_id, + external_trace_id=None, + session_id=self.session_id, + user_id=None, + parent=None, + spans=(span,), + ) + self._keep_terminal(trace, now=now) + + def _keep_terminal(self, trace: Trace, *, now: datetime) -> None: + if not self.state.accepts_terminal: + return + stored = trace.model_dump(mode="json") + if self.outer_trace is not None and self.outer_trace != stored: + # A duplicate terminal task must not rewrite a completed immutable + # span. The first persisted terminal envelope owns the export. + return + self.outer_trace = stored + self.root_span_id = trace.root.span_id + self.terminal_at = self.terminal_at or now + self.updated_at = now + + def recovery_checked(self, error_code: str | None, *, now: datetime) -> None: + """Rotate incomplete work behind other pending runs, or fail it with ``error_code``. + + A run that captured its terminal or left collection is not changed. + """ + if not self.awaits_terminal: + return + self.updated_at = now + if error_code is not None: + self.state, self.error_code = TraceRunState.FAILED, error_code + + def claim_export( + self, *, token: str, owner_exists: bool, collection_timeout: int, lease_seconds: int, now: datetime + ) -> bool: + """Lease a complete run for one export attempt that ``token`` fences; ``False`` when it is not exportable now. + + A run is complete once its terminal is captured and every invocation is + finalized; ``collection_timeout`` seconds after the terminal an + incomplete run fails. An export whose last heartbeat is older than + ``lease_seconds`` becomes uncertain, and a run whose app is gone fails. + """ + if self.state in _SETTLED_STATES: + return False + if self.state is TraceRunState.EXPORTING: + last_heartbeat = self.export_heartbeat_at or self.export_started_at + if last_heartbeat and last_heartbeat + timedelta(seconds=lease_seconds) < now: + self.state, self.error_code = TraceRunState.UNCERTAIN, "export_lease_expired" + self.updated_at = now + return False + if self.outer_trace is None: + return False + if any(not invocation.finalized for invocation in self.invocations.values()): + if self.terminal_at and self.terminal_at + timedelta(seconds=collection_timeout) < now: + self.state, self.error_code = TraceRunState.FAILED, "agent_trace_collection_incomplete" + self.updated_at = now + return False + if self.next_attempt_at and self.next_attempt_at > now: + return False + if not owner_exists: + self.state, self.error_code = TraceRunState.FAILED, "tracing_owner_not_found" + return False + self.state, self.export_token = TraceRunState.EXPORTING, token + self.export_started_at = self.export_heartbeat_at = self.updated_at = now + self.attempts += 1 + return True + + def renew_export_lease(self, token: str, *, now: datetime) -> bool: + """Renew only the current export attempt; ``False`` once the attempt no longer holds the run.""" + if self.state is not TraceRunState.EXPORTING or self.export_token != token: + return False + self.export_heartbeat_at = self.updated_at = now + return True + + def export_accepted( + self, + token: str, + *, + provider_trace_id: str, + trace_url: str | None, + span_urls: Mapping[str, str], + now: datetime, + ) -> None: + """Record the destination's acceptance. + + An acceptance that arrives after its lease expired still records the + destination's result, unless a manual retry replaced its token. + """ + late = self.state is TraceRunState.UNCERTAIN and self.error_code == "export_lease_expired" + if self.export_token != token or (self.state is not TraceRunState.EXPORTING and not late): + return + self._finish_export(TraceRunState.ACCEPTED, None, now=now) + self.provider_trace_id, self.trace_url, self.span_urls = provider_trace_id, trace_url, dict(span_urls) + self.accepted_at = now + + def export_deferred(self, token: str, *, max_attempts: int, now: datetime) -> None: + """Wait with exponential backoff after a retryable failure until the attempts run out.""" + if self.attempts >= max_attempts: + self.export_failed(token, "provider_retry_exhausted", now=now) + return + if not self._holds_export(token): + return + self._finish_export(TraceRunState.RETRYING, "provider_temporarily_unavailable", now=now) + self.next_attempt_at = now + timedelta(seconds=min(300, 5 * 2 ** min(self.attempts - 1, 6))) + + def export_uncertain(self, token: str, *, now: datetime) -> None: + """The destination may have accepted the trace; only an acknowledged retry sends it again.""" + if self._holds_export(token): + self._finish_export(TraceRunState.UNCERTAIN, "provider_acceptance_unknown", now=now) + + def export_failed(self, token: str, error_code: str, *, now: datetime) -> None: + if self._holds_export(token): + self._finish_export(TraceRunState.FAILED, error_code, now=now) + + def _holds_export(self, token: str) -> bool: + return self.state is TraceRunState.EXPORTING and self.export_token == token + + def _finish_export(self, state: TraceRunState, error_code: str | None, *, now: datetime) -> None: + self.state, self.error_code = state, error_code + self.export_token = None + self.next_attempt_at = None + self.updated_at = now + + def release_for_retry(self, *, acknowledge_possible_duplicates: bool, now: datetime) -> None: + """Release a failed or uncertain run for another export. + + Raises: + TraceRunNotFoundError: Nothing was collected. + TraceRunConflictError: The run is accepted, not retryable, or needs the duplicate acknowledgement. + """ + if self.state is TraceRunState.NOT_COLLECTED: + raise TraceRunNotFoundError("tracing_run_not_found") + if self.state is TraceRunState.ACCEPTED: + raise TraceRunConflictError("tracing_already_accepted") + if self.state not in {TraceRunState.FAILED, TraceRunState.UNCERTAIN} or self.outer_trace is None: + raise TraceRunConflictError("tracing_retry_not_allowed") + if any(not invocation.finalized for invocation in self.invocations.values()) or ( + self.error_code or "" + ).startswith("tracing_"): + raise TraceRunConflictError("tracing_retry_not_allowed") + if self.state is TraceRunState.UNCERTAIN and not acknowledge_possible_duplicates: + raise TraceRunConflictError("tracing_duplicate_acknowledgement_required") + self.state, self.error_code, self.export_token = TraceRunState.READY, None, None + self.next_attempt_at = None + self.updated_at = now + + def exported_trace(self, batches: Sequence[TraceBatch]) -> Trace: + """The trace this run exports: its business terminal with each invocation's spans under an Agent span. + + ``batches`` are the run's stored batches by arrival. + + Raises: + TraceRunError: The stored tree is not valid, so no export can repair it. + ValueError: A stored span is malformed. + """ + if self.outer_trace is None: + raise TraceRunError("tracing_terminal_missing") + # Runs captured before the upgrade stored the previous trace shape, named by its root span. + stored = self.outer_trace + outer = legacy_outer_trace(stored) if "root_span_id" in stored else Trace.model_validate(stored) + spans = list(outer.spans) + internal_by_invocation: dict[str, dict[str, Span]] = {invocation_id: {} for invocation_id in self.invocations} + for batch in batches: + if batch.invocation_id not in self.invocations: + raise TraceRunError("tracing_invocation_not_found") + for data in batch.spans: + span = normalize_agent_span(data, capture_content=self.capture_content) + current = internal_by_invocation[batch.invocation_id].get(span.span_id) + if current is not None and current != span: + raise TraceRunError("tracing_span_conflict") + internal_by_invocation[batch.invocation_id][span.span_id] = span + for invocation in self.invocations.values(): + children = list(internal_by_invocation[invocation.id].values()) + if not children: + raise TraceRunError("tracing_invocation_empty") + valid_ids = {child.span_id for child in children} + if any( + child.parent_span_id not in valid_ids and child.parent_span_id != invocation.span_id + for child in children + ): + raise TraceRunError("tracing_invocation_parent_invalid") + root_children = [child for child in children if child.parent_span_id == invocation.span_id] + if not root_children: + raise TraceRunError("tracing_invocation_root_missing") + errors = [child.error for child in root_children if child.status is SpanStatus.ERROR] + spans.append( + Span( + span_id=invocation.span_id, + parent_span_id=invocation.parent_span_id, + name=f"Agent attempt {invocation.attempt + 1}", + kind=SpanKind.AGENT, + start_time=min(child.start_time for child in children), + end_time=max(child.end_time for child in children), + status=SpanStatus.ERROR if errors else SpanStatus.OK, + error=errors[0] if errors else None, + attributes={ + "agent_id": invocation.agent_id, + "attempt": invocation.attempt, + "node_execution_id": invocation.node_execution_id, + "invocation_id": invocation.id, + }, + ) + ) + spans.extend(children) + trusted: dict[str, Any] = { + "tenant_id": self.tenant_id, + "app_id": self.app_id, + "session_id": self.session_id, + "dify.tenant_id": self.tenant_id, + "dify.app_id": self.app_id, + "dify.tracing.run_id": self.id, + "dify.source.id": self.source_id, + } + spans = [span.model_copy(update={"attributes": {**span.attributes, **trusted}}) for span in spans] + return Trace( + trace_id=self.trace_id, + external_trace_id=None, + session_id=self.session_id, + user_id=outer.user_id, + parent=None, + spans=order_spans(outer.root.span_id, spans), + ) diff --git a/api/domain/workflow/__init__.py b/api/domain/workflow/__init__.py new file mode 100644 index 00000000000000..b0ad6df941c98a --- /dev/null +++ b/api/domain/workflow/__init__.py @@ -0,0 +1 @@ +"""Pure workflow contracts shared by application services and execution adapters.""" diff --git a/api/core/workflow/nodes/agent_v2/discriminator.py b/api/domain/workflow/agent_node.py similarity index 92% rename from api/core/workflow/nodes/agent_v2/discriminator.py rename to api/domain/workflow/agent_node.py index 8df480ca5e7668..54dfd9462f01ee 100644 --- a/api/core/workflow/nodes/agent_v2/discriminator.py +++ b/api/domain/workflow/agent_node.py @@ -1,3 +1,5 @@ +"""Identify Dify Agent nodes without importing their execution runtime.""" + from collections.abc import Mapping from typing import Any diff --git a/api/core/app/workflow/retry_history.py b/api/domain/workflow/retry_history.py similarity index 88% rename from api/core/app/workflow/retry_history.py rename to api/domain/workflow/retry_history.py index 75060c55fd1555..8147a5515d6dd3 100644 --- a/api/core/app/workflow/retry_history.py +++ b/api/domain/workflow/retry_history.py @@ -1,4 +1,4 @@ -"""Validated internal payloads for persisted workflow node retry history.""" +"""Validated retry history shared by workflow persistence and tracing.""" from collections.abc import Mapping from typing import Any diff --git a/api/enterprise/telemetry/enterprise_trace.py b/api/enterprise/telemetry/enterprise_trace.py index 42874e1121fae7..020feaf1ceb52e 100644 --- a/api/enterprise/telemetry/enterprise_trace.py +++ b/api/enterprise/telemetry/enterprise_trace.py @@ -35,7 +35,15 @@ from opentelemetry.util.types import AttributeValue -from core.ops.entities.trace_entity import ( +from enterprise.telemetry.entities import ( + EnterpriseTelemetryCounter, + EnterpriseTelemetryEvent, + EnterpriseTelemetryHistogram, + EnterpriseTelemetrySpan, + TokenMetricLabels, +) +from enterprise.telemetry.telemetry_log import emit_metric_only_event, emit_telemetry_log +from services.tracing.records import ( BaseTraceInfo, DatasetRetrievalTraceInfo, DraftNodeExecutionTrace, @@ -49,14 +57,6 @@ WorkflowNodeTraceInfo, WorkflowTraceInfo, ) -from enterprise.telemetry.entities import ( - EnterpriseTelemetryCounter, - EnterpriseTelemetryEvent, - EnterpriseTelemetryHistogram, - EnterpriseTelemetrySpan, - TokenMetricLabels, -) -from enterprise.telemetry.telemetry_log import emit_metric_only_event, emit_telemetry_log logger = logging.getLogger(__name__) @@ -261,15 +261,9 @@ def _workflow_trace(self, info: WorkflowTraceInfo) -> None: invoke_from=invoke_from, ), ) - # Prefer wall-clock timestamps over the elapsed_time field: elapsed_time defaults - # to 0 in the DB and can be stale if the Celery write races with the trace task. - # start_time = workflow_run.created_at, end_time = workflow_run.finished_at. - if info.start_time and info.end_time: - workflow_duration = (info.end_time - info.start_time).total_seconds() - elif info.workflow_run_elapsed_time: - workflow_duration = float(info.workflow_run_elapsed_time) - else: - workflow_duration = 0.0 + # Wall-clock timestamps, not the elapsed_time field: elapsed_time defaults to 0 + # in the DB and can be stale if the Celery write races with the trace task. + workflow_duration = (info.end_time - info.start_time).total_seconds() self._exporter.record_histogram( EnterpriseTelemetryHistogram.WORKFLOW_DURATION, workflow_duration, diff --git a/api/events/event_handlers/sync_knowledge_fs_bindings_when_app_published_workflow_updated.py b/api/events/event_handlers/sync_knowledge_fs_bindings_when_app_published_workflow_updated.py index f6ca2271c03d82..c3c10ceb6e6589 100644 --- a/api/events/event_handlers/sync_knowledge_fs_bindings_when_app_published_workflow_updated.py +++ b/api/events/event_handlers/sync_knowledge_fs_bindings_when_app_published_workflow_updated.py @@ -6,7 +6,7 @@ from sqlalchemy.orm import Session from core.db.session_factory import session_factory -from core.workflow.nodes.agent_v2.discriminator import is_dify_agent_node_data +from domain.workflow.agent_node import is_dify_agent_node_data from events.app_event import app_published_workflow_was_updated from extensions.ext_database import db from models.knowledge_fs import AppKnowledgeFSSpaceJoin, KnowledgeFSAppSpaceJoinType diff --git a/api/extensions/application_services/tracing.py b/api/extensions/application_services/tracing.py new file mode 100644 index 00000000000000..b4ef5cd0e4189e --- /dev/null +++ b/api/extensions/application_services/tracing.py @@ -0,0 +1,109 @@ +"""Composition of the trace pipeline, Agent trace collection, and tracing provider configurations.""" + +from dataclasses import dataclass + +from sqlalchemy.orm import Session, sessionmaker + +from configs import dify_config +from enterprise.telemetry.exporter import is_enterprise_telemetry_enabled +from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository +from repositories.application_trace_repository import ApplicationTraceRepository +from repositories.factory import DifyAPIRepositoryFactory +from repositories.trace_record_repository import TraceRecordRepository +from services.tracing.adapters import ( + CeleryExportRetryQueue, + CeleryTraceEmitter, + WorkspaceTracingSecretCipher, + schedule_agent_trace_export, + send_to_enterprise_telemetry, +) +from services.tracing.canonical import CanonicalTraceBuilder +from services.tracing.collection import AgentTraceSettings +from services.tracing.delivery import TraceDeliveryService +from services.tracing.egress import SSRFEgress +from services.tracing.provider_configs import ProviderConfigGateway, TraceProviderCapabilitiesQuery +from services.tracing.providers import ProviderConfigCodec +from services.tracing.registry import EntryPointTraceProviderCatalog +from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks +from services.tracing.scopes import TraceScopes +from services.tracing.service import ApplicationTracingService + + +@dataclass(frozen=True, slots=True) +class TracingServices: + scopes: TraceScopes + delivery: TraceDeliveryService + agent_traces: ApplicationTracingService + hooks: RuntimeApplicationTracingHooks | None + """``None`` when the deployment does not collect Agent traces.""" + provider_configs: ProviderConfigGateway + provider_capabilities: TraceProviderCapabilitiesQuery + + +def build_tracing_services(*, database_client: sessionmaker[Session]) -> TracingServices: + egress = SSRFEgress() + providers = EntryPointTraceProviderCatalog() + codec = ProviderConfigCodec(WorkspaceTracingSecretCipher()) + builder = CanonicalTraceBuilder() + configs = SQLAlchemyAppTracingConfigRepository(session_factory=database_client) + enterprise_telemetry = is_enterprise_telemetry_enabled() + workflow_runs = DifyAPIRepositoryFactory.create_api_workflow_run_repository(database_client) + node_executions = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(database_client) + agent_settings = AgentTraceSettings( + enabled=dify_config.OPS_AGENT_TRACE_ENABLED, + max_spans=dify_config.OPS_AGENT_TRACE_MAX_SPANS, + max_bytes=dify_config.OPS_AGENT_TRACE_MAX_BYTES, + collection_timeout_seconds=dify_config.OPS_AGENT_TRACE_COLLECTION_TIMEOUT, + export_lease_seconds=dify_config.OPS_AGENT_TRACE_EXPORT_LEASE_SECONDS, + max_export_attempts=dify_config.OPS_AGENT_TRACE_MAX_ATTEMPTS, + retention_days=dify_config.OPS_AGENT_TRACE_RETENTION_DAYS, + ) + agent_trace_runs = ApplicationTraceRepository(database_client) + agent_traces = ApplicationTracingService( + runs=agent_trace_runs, + run_query=agent_trace_runs, + sources=agent_trace_runs, + # App, workspace, and credential names are enterprise telemetry dimensions only. + records=TraceRecordRepository( + database_client, + workflow_runs=workflow_runs, + node_executions=node_executions, + files_url=dify_config.FILES_URL, + enterprise_details=False, + ), + builder=builder, + providers=providers, + codec=codec, + egress=egress, + schedule_export=schedule_agent_trace_export, + settings=agent_settings, + ) + return TracingServices( + scopes=TraceScopes( + emitter=CeleryTraceEmitter(), policies=configs, enterprise_enabled=is_enterprise_telemetry_enabled + ), + delivery=TraceDeliveryService( + records=TraceRecordRepository( + database_client, + workflow_runs=workflow_runs, + node_executions=node_executions, + files_url=dify_config.FILES_URL, + enterprise_details=enterprise_telemetry, + ), + destinations=configs, + providers=providers, + codec=codec, + builder=builder, + egress=egress, + retries=CeleryExportRetryQueue(), + enterprise_sink=send_to_enterprise_telemetry if enterprise_telemetry else None, + collector=agent_traces if agent_settings.enabled else None, + max_export_attempts=dify_config.OPS_TRACE_EXPORT_MAX_ATTEMPTS, + ), + agent_traces=agent_traces, + hooks=RuntimeApplicationTracingHooks(agent_traces) if agent_settings.enabled else None, + provider_configs=ProviderConfigGateway(providers=providers, codec=codec, egress=egress), + provider_capabilities=TraceProviderCapabilitiesQuery( + providers=providers, api_collection_enabled=agent_settings.enabled + ), + ) diff --git a/api/extensions/ext_application_services.py b/api/extensions/ext_application_services.py index 66d46a01180e35..5ff3b4dfa940a9 100644 --- a/api/extensions/ext_application_services.py +++ b/api/extensions/ext_application_services.py @@ -17,7 +17,6 @@ from configs import dify_config from constants.dsl_version import CURRENT_APP_DSL_VERSION from core.db.session_factory import get_session_maker -from core.ops.base_trace_instance import bind_application_tracing_hooks from core.schemas.schema_manager import SchemaManager from core.tools.tool_file_manager import ToolFileManager from core.tools.tool_manager import ToolManager @@ -39,6 +38,7 @@ ) from extensions.application_services.rbac import RBACServices, build_rbac_services from extensions.application_services.resource_access_token import build_resource_access_token_service +from extensions.application_services.tracing import TracingServices, build_tracing_services from extensions.application_services.trial_app import TrialAppServices, build_trial_app_services from extensions.application_services.workspace import ( WorkspaceServices, @@ -47,6 +47,7 @@ ) from extensions.ext_redis import RedisClientWrapper, redis_client from extensions.ext_storage import storage +from extensions.tracing import bind_tracing from libs.helper import RateLimiter from libs.passport import PassportService from models.model import EndUser @@ -132,13 +133,12 @@ from services.app_site_service import AppSiteService from services.app_statistic_query import AppStatisticQuery from services.app_task_service import AppTaskControlService -from services.app_tracing_config_gateway import OpsTraceManagerGateway from services.app_tracing_config_service import AppTracingConfigService from services.billing_portal_service import BillingPortalService from services.billing_service import BillingService from services.compliance_download_service import ComplianceDownloadService from services.conversation_command_service import ConversationCommandService -from services.conversation_name_generator import generate_conversation_name +from services.conversation_name_generator import DefaultModelConversationNameGenerator from services.conversation_name_service import ConversationNameService from services.conversation_query_service import ConversationQueryService from services.conversation_service import ConversationService @@ -217,7 +217,7 @@ from services.system_feature_service import SystemFeatureService from services.tag_application_service import TagApplicationService from services.tool_file_download_service import ToolFileDownloadService -from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks +from services.tracing.adapters import StoredPolicyTraceScopes from services.upload_file_delivery_service import UploadFileDeliveryService from services.web_app_runtime_query_service import WebAppRuntimeQueryService from services.web_authentication_adapters import ( @@ -366,6 +366,7 @@ class ApplicationServices: web_passport: WebPassportService tags: TagApplicationService workflow_statistics: WorkflowStatisticQueryService + tracing: TracingServices def _build_oauth_server_service( @@ -516,9 +517,10 @@ def build_application_services( apps=console_apps, datasets=dataset_dependencies.datasets, ) + tracing = build_tracing_services(database_client=database_client) tracing_configs = AppTracingConfigService( configs=SQLAlchemyAppTracingConfigRepository(session_factory=database_client), - provider=OpsTraceManagerGateway(), + provider=tracing.provider_configs, ) apps = build_app_services( database_client=database_client, @@ -584,7 +586,10 @@ def build_application_services( invitation_tokens=invitation_tokens, ) conversations = ConversationRepository(session_factory=database_client) - conversation_names = ConversationNameService(conversations=conversations, generate_name=generate_conversation_name) + conversation_names = ConversationNameService( + conversations=conversations, + generate_name=DefaultModelConversationNameGenerator(trace_scopes=tracing.scopes), + ) conversation_commands = ConversationCommandService( conversations=conversations, enqueue_delete_cleanup=ConversationService.enqueue_delete_cleanup ) @@ -769,7 +774,7 @@ def build_application_services( ), message_suggested_questions=MessageSuggestedQuestionsService( queries=suggested_questions, - generator=SuggestedQuestionsGenerator(), + generator=SuggestedQuestionsGenerator(trace_scopes=tracing.scopes), ), plugin_file_uploads=PluginFileUploadService( owners=SQLAlchemyPluginFileUploadOwnerRepository(session_factory=database_client), @@ -862,6 +867,7 @@ def build_application_services( session_maker=database_client, ), ), + tracing=tracing, ) @@ -879,7 +885,7 @@ def init_app(app: Flask) -> None: redis=redis_client, ) app.extensions[_EXTENSION_KEY] = services - bind_application_tracing_hooks(app, RuntimeApplicationTracingHooks()) + bind_tracing(app, scopes=StoredPolicyTraceScopes(services.tracing.scopes), hooks=services.tracing.hooks) bind_account_loader(app, services.accounts.identity.load_user) diff --git a/api/extensions/ext_celery.py b/api/extensions/ext_celery.py index f9019d65224909..66c9ca5b5f405f 100644 --- a/api/extensions/ext_celery.py +++ b/api/extensions/ext_celery.py @@ -159,6 +159,11 @@ def __call__(self, *args: object, **kwargs: object) -> object: timezone=pytz.timezone(dify_config.LOG_TZ or "UTC"), task_ignore_result=True, task_annotations=dify_config.CELERY_TASK_ANNOTATIONS, + task_routes={ + "tasks.trace_delivery_task.*": {"queue": dify_config.OPS_TRACE_QUEUE}, + "tasks.application_trace_task.*": {"queue": dify_config.OPS_TRACE_QUEUE}, + "schedule.collect_agent_sandbox_usage.collect_agent_sandbox_usage": {"queue": dify_config.OPS_TRACE_QUEUE}, + }, ) if dify_config.CELERY_BACKEND == "redis": @@ -209,6 +214,8 @@ def __call__(self, *args: object, **kwargs: object) -> object: "tasks.knowledge_fs_upgrade_tasks", # legacy Dataset upgrades use a dedicated queue and worker "tasks.knowledge_fs_background_tasks", # local-engine tasks on dedicated prefork queues "tasks.workflow_run_archive_deletion_tasks", # permanent monthly archive deletion + "tasks.trace_delivery_task", # trace delivery and provider export retries + "tasks.ops_trace_task", # retired endpoint; old workers must drain pre-upgrade tasks ] day = dify_config.CELERY_BEAT_SCHEDULER_TIME @@ -254,6 +261,11 @@ def __call__(self, *args: object, **kwargs: object) -> object: "task": "tasks.knowledge_fs_upgrade_tasks.cleanup_deferred_knowledge_fs_upgrade_files", "schedule": timedelta(seconds=dify_config.KNOWLEDGE_FS_LIFECYCLE_POLL_INTERVAL_SECONDS), } + if dify_config.ENABLE_TRACE_PAYLOAD_SWEEP_TASK: + beat_schedule["sweep_trace_payloads"] = { + "task": "tasks.trace_delivery_task.sweep_trace_payloads", + "schedule": crontab(minute="30", hour="3"), + } if dify_config.OPS_AGENT_TRACE_ENABLED: imports.append("tasks.application_trace_task") if dify_config.OPS_AGENT_TRACE_ENABLED and dify_config.ENABLE_APPLICATION_TRACE_RECONCILE_TASK: diff --git a/api/extensions/ext_forward_refs.py b/api/extensions/ext_forward_refs.py deleted file mode 100644 index d3b13295325bc0..00000000000000 --- a/api/extensions/ext_forward_refs.py +++ /dev/null @@ -1,49 +0,0 @@ -import logging - -from dify_app import DifyApp - - -def is_enabled() -> bool: - return True - - -def init_app(app: DifyApp): - """Resolve Pydantic forward refs that would otherwise cause circular imports. - - Rebuilds models in core.app.entities.app_invoke_entities with the real TraceQueueManager type. - Safe to run multiple times. - """ - logger = logging.getLogger(__name__) - try: - from core.app.entities.app_invoke_entities import ( - AdvancedChatAppGenerateEntity, - AgentChatAppGenerateEntity, - AppGenerateEntity, - ChatAppGenerateEntity, - CompletionAppGenerateEntity, - ConversationAppGenerateEntity, - EasyUIBasedAppGenerateEntity, - RagPipelineGenerateEntity, - WorkflowAppGenerateEntity, - ) - from core.ops.ops_trace_manager import TraceQueueManager # heavy import, do it at startup only - - ns = {"TraceQueueManager": TraceQueueManager} - for model in ( - AppGenerateEntity, - EasyUIBasedAppGenerateEntity, - ConversationAppGenerateEntity, - ChatAppGenerateEntity, - CompletionAppGenerateEntity, - AgentChatAppGenerateEntity, - AdvancedChatAppGenerateEntity, - WorkflowAppGenerateEntity, - RagPipelineGenerateEntity, - ): - try: - model.model_rebuild(_types_namespace=ns) - except Exception as e: - logger.debug("model_rebuild skipped for %s: %s", model.__name__, e) - except Exception as e: - # Don't block app startup; just log at debug level. - logger.debug("ext_forward_refs init skipped: %s", e) diff --git a/api/extensions/ext_sentry.py b/api/extensions/ext_sentry.py index 19a153c6f786f2..05a5cdd168604b 100644 --- a/api/extensions/ext_sentry.py +++ b/api/extensions/ext_sentry.py @@ -13,24 +13,6 @@ def init_app(app: DifyApp): from graphon.model_runtime.errors.invoke import InvokeRateLimitError - try: - from langfuse._utils import parse_error - - _langfuse_error_response = parse_error.defaultErrorResponse - except (ImportError, AttributeError): - _langfuse_error_response = ( - "Unexpected error occurred. Please check your request" - " and contact support: https://langfuse.com/support." - ) - - def before_send(event, hint): - if "exc_info" in hint: - _, exc_value, _ = hint["exc_info"] - if _langfuse_error_response in str(exc_value): - return None - - return event - sentry_sdk.init( dsn=dify_config.SENTRY_DSN, integrations=[FlaskIntegration(), CeleryIntegration()], @@ -39,13 +21,11 @@ def before_send(event, hint): ValueError, FileNotFoundError, InvokeRateLimitError, - _langfuse_error_response, ], traces_sample_rate=dify_config.SENTRY_TRACES_SAMPLE_RATE, profiles_sample_rate=dify_config.SENTRY_PROFILES_SAMPLE_RATE, environment=dify_config.DEPLOY_ENV, release=f"dify-{dify_config.project.version}-{dify_config.COMMIT_SHA}", - before_send=before_send, ) if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD: diff --git a/api/extensions/logstore/repositories/logstore_api_workflow_node_execution_repository.py b/api/extensions/logstore/repositories/logstore_api_workflow_node_execution_repository.py index 9b4953fac1e38b..9a549fbd59a273 100644 --- a/api/extensions/logstore/repositories/logstore_api_workflow_node_execution_repository.py +++ b/api/extensions/logstore/repositories/logstore_api_workflow_node_execution_repository.py @@ -22,6 +22,7 @@ from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom from repositories.api_workflow_node_execution_repository import ( DifyAPIWorkflowNodeExecutionRepository, + WorkflowNodeExecutionPayloads, WorkflowNodeExecutionSnapshot, ) from repositories.sqlalchemy_api_workflow_node_execution_repository import ( @@ -161,11 +162,32 @@ def delete_executions_by_app(self, tenant_id: str, app_id: str, batch_size: int tenant_id=tenant_id, app_id=app_id, batch_size=batch_size ) + @override + def load_full_inputs(self, execution: WorkflowNodeExecutionModel) -> Mapping[str, Any] | None: + return execution.inputs_dict + + @override + def load_full_outputs(self, execution: WorkflowNodeExecutionModel) -> Mapping[str, Any] | None: + return execution.outputs_dict + @override def load_full_process_data(self, execution: WorkflowNodeExecutionModel) -> Mapping[str, Any] | None: """Return complete Process Data already stored inline by LogStore.""" return execution.process_data_dict + @override + def load_full_payloads( + self, executions: Sequence[WorkflowNodeExecutionModel] + ) -> Mapping[str, WorkflowNodeExecutionPayloads]: + return { + execution.id: WorkflowNodeExecutionPayloads( + inputs=execution.inputs_dict, + process_data=execution.process_data_dict, + outputs=execution.outputs_dict, + ) + for execution in executions + } + @override def get_node_last_execution( self, diff --git a/api/extensions/logstore/repositories/logstore_workflow_node_execution_repository.py b/api/extensions/logstore/repositories/logstore_workflow_node_execution_repository.py index 714cc694a750c2..b9b2b08732a3fc 100644 --- a/api/extensions/logstore/repositories/logstore_workflow_node_execution_repository.py +++ b/api/extensions/logstore/repositories/logstore_workflow_node_execution_repository.py @@ -12,11 +12,11 @@ from datetime import datetime from typing import Any, override +from pydantic import TypeAdapter from sqlalchemy.engine import Engine from sqlalchemy.orm import sessionmaker from configs import dify_config -from core.ops.utils import JSON_DICT_ADAPTER from core.repositories import SQLAlchemyWorkflowNodeExecutionRepository from core.repositories.factory import OrderConfig, WorkflowNodeExecutionRepository from core.workflow.node_execution import WorkflowNodeExecution as DifyWorkflowNodeExecution @@ -37,6 +37,8 @@ logger = logging.getLogger(__name__) +_JSON_DICT_ADAPTER: TypeAdapter[dict[str, Any]] = TypeAdapter(dict[str, Any]) + def _dict_to_workflow_node_execution(data: dict[str, Any]) -> WorkflowNodeExecution: """ @@ -50,10 +52,10 @@ def _dict_to_workflow_node_execution(data: dict[str, Any]) -> WorkflowNodeExecut """ logger.debug("_dict_to_workflow_node_execution: data keys=%s", list(data.keys())[:5]) # Parse JSON fields - inputs = JSON_DICT_ADAPTER.validate_json(data.get("inputs") or "{}") - process_data = JSON_DICT_ADAPTER.validate_json(data.get("process_data") or "{}") - outputs = JSON_DICT_ADAPTER.validate_json(data.get("outputs") or "{}") - metadata = JSON_DICT_ADAPTER.validate_json(data.get("execution_metadata") or "{}") + inputs = _JSON_DICT_ADAPTER.validate_json(data.get("inputs") or "{}") + process_data = _JSON_DICT_ADAPTER.validate_json(data.get("process_data") or "{}") + outputs = _JSON_DICT_ADAPTER.validate_json(data.get("outputs") or "{}") + metadata = _JSON_DICT_ADAPTER.validate_json(data.get("execution_metadata") or "{}") # Convert metadata to domain enum keys domain_metadata = {} diff --git a/api/extensions/tracing.py b/api/extensions/tracing.py new file mode 100644 index 00000000000000..d5273e28f8e6ee --- /dev/null +++ b/api/extensions/tracing.py @@ -0,0 +1,155 @@ +"""Execution-side access to the trace pipeline, bound to one Flask application. + +Legacy producers in the execution kernel import only these accessors; the +composition root binds the services they reach. Every accessor fails open: +tracing never breaks the operation that is being traced. +""" + +import logging +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Literal, Protocol, runtime_checkable + +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import TraceCorrelation +from flask import Flask, current_app, has_app_context + +from core.helper.trace_id_helper import ParentTraceContext + +if TYPE_CHECKING: + from dify_agent.layers.execution_context import ApplicationTracingContext + + from core.app.apps.agent_app.runtime_request_builder import AgentAppRuntimeBuildContext + from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, WorkflowAppGenerateEntity + from core.workflow.nodes.agent_v2.runtime_request_builder import WorkflowAgentRuntimeBuildContext + +logger = logging.getLogger(__name__) + +_SCOPES_KEY = "trace_scopes" +_HOOKS_KEY = "application_tracing_hooks" + + +@runtime_checkable +class ExecutionTraceScopes(Protocol): + """Trace scopes for legacy producers, which hold the app's stored tracing policy as its raw column value.""" + + def for_execution( + self, + *, + tenant_id: str, + app_id: str | None, + app_tracing: str | None, + user_id: str | None, + correlation: TraceCorrelation, + ) -> TraceScope | None: ... + + def for_app( + self, *, tenant_id: str, app_id: str | None, user_id: str | None, correlation: TraceCorrelation + ) -> TraceScope | None: ... + + +@runtime_checkable +class ApplicationTracingHooks(Protocol): + """Agent trace collection lifecycle; it never replaces the trace pipeline.""" + + def prepare_agent_app(self, context: "AgentAppRuntimeBuildContext") -> "ApplicationTracingContext | None": ... + + def prepare_workflow_agent( + self, context: "WorkflowAgentRuntimeBuildContext", conversation_id: str | None + ) -> "ApplicationTracingContext | None": ... + + def freeze_workflow_run( + self, + entity: "AdvancedChatAppGenerateEntity | WorkflowAppGenerateEntity", + *, + workflow_id: str, + workflow_run_id: str, + conversation_id: object, + graph_data: object, + ) -> None: ... + + def capture_agent_app_terminal( + self, + *, + tenant_id: str, + app_id: str, + message_id: str, + outcome: Literal["succeeded", "failed", "cancelled"], + output: str | None, + ) -> None: ... + + +def bind_tracing(app: Flask, *, scopes: ExecutionTraceScopes, hooks: ApplicationTracingHooks | None) -> None: + """``hooks`` is ``None`` when this deployment does not collect Agent traces.""" + app.extensions[_SCOPES_KEY] = scopes + app.extensions[_HOOKS_KEY] = hooks + + +def _scopes() -> ExecutionTraceScopes | None: + if not has_app_context(): + return None + scopes = current_app.extensions.get(_SCOPES_KEY) + return scopes if isinstance(scopes, ExecutionTraceScopes) else None + + +def correlation_from_extras(extras: Mapping[str, Any]) -> TraceCorrelation: + """The caller-supplied correlation an execution carries in its ``extras``.""" + parent = extras.get("parent_trace_context") + if isinstance(parent, Mapping): + try: + parent = ParentTraceContext.model_validate(parent) + except ValueError: + parent = None + external_trace_id, session_id = extras.get("external_trace_id"), extras.get("trace_session_id") + return TraceCorrelation( + external_trace_id=external_trace_id if isinstance(external_trace_id, str) and external_trace_id else None, + session_id=session_id if isinstance(session_id, str) and session_id else None, + parent_workflow_run_id=parent.parent_workflow_run_id if isinstance(parent, ParentTraceContext) else None, + parent_node_execution_id=parent.parent_node_execution_id if isinstance(parent, ParentTraceContext) else None, + ) + + +def create_trace_scope( + *, tenant_id: str, app_id: str | None, app_tracing: str | None, user_id: str | None, extras: Mapping[str, Any] +) -> TraceScope | None: + """The scope for an execution whose app row is loaded; ``None`` when nothing consumes its events.""" + scopes = _scopes() + if scopes is None: + return None + try: + return scopes.for_execution( + tenant_id=tenant_id, + app_id=app_id, + app_tracing=app_tracing, + user_id=user_id, + correlation=correlation_from_extras(extras), + ) + except Exception: + logger.warning("Tracing scope unavailable: app_id=%s", app_id, exc_info=True) + return None + + +def create_trace_scope_for_app( + *, tenant_id: str, app_id: str | None, user_id: str | None, extras: Mapping[str, Any] +) -> TraceScope | None: + """The scope for a caller that does not hold the app row.""" + scopes = _scopes() + if scopes is None: + return None + try: + return scopes.for_app( + tenant_id=tenant_id, app_id=app_id, user_id=user_id, correlation=correlation_from_extras(extras) + ) + except Exception: + logger.warning("Tracing scope unavailable: app_id=%s", app_id, exc_info=True) + return None + + +def application_tracing_hooks() -> ApplicationTracingHooks | None: + """The Agent collection hooks; ``None`` when Agent tracing is off or no application is bound here. + + Agent runs started without hooks are reported as not collected instead of failing. + """ + if not has_app_context(): + return None + hooks = current_app.extensions.get(_HOOKS_KEY) + return hooks if isinstance(hooks, ApplicationTracingHooks) else None diff --git a/api/libs/datetime_utils.py b/api/libs/datetime_utils.py index 95697a352c8190..e25fb2d913cf61 100644 --- a/api/libs/datetime_utils.py +++ b/api/libs/datetime_utils.py @@ -1,5 +1,8 @@ import abc import datetime +from collections.abc import Generator +from contextlib import contextmanager +from dataclasses import dataclass from typing import Protocol import pytz # type: ignore[import-untyped] @@ -22,6 +25,23 @@ def utc_now() -> datetime.datetime: return _now_func(datetime.UTC) +@dataclass +class Timing: + started_at: datetime.datetime + ended_at: datetime.datetime + + +@contextmanager +def measure_time() -> Generator[Timing]: + """Time a block in aware UTC; ``ended_at`` is final once the block exits.""" + started_at = utc_now() + timing = Timing(started_at=started_at, ended_at=started_at) + try: + yield timing + finally: + timing.ended_at = utc_now() + + def naive_utc_now() -> datetime.datetime: """Return a naive datetime object (without timezone information) representing current UTC time. diff --git a/api/migrations/versions/2026_10_23_1200-53c10126d551_unique_trace_app_config_provider.py b/api/migrations/versions/2026_10_23_1200-53c10126d551_unique_trace_app_config_provider.py new file mode 100644 index 00000000000000..78cd9b7c8efe14 --- /dev/null +++ b/api/migrations/versions/2026_10_23_1200-53c10126d551_unique_trace_app_config_provider.py @@ -0,0 +1,50 @@ +"""keep one tracing configuration per app and provider + +Revision ID: 53c10126d551 +Revises: 0cc3158b50b0 +Create Date: 2026-10-23 12:00:00.000000 + +Configuration writes lock the app row, but rows written before that lock existed can +duplicate an (app, provider) pair, and reads then pick one arbitrarily. The most +recently updated row wins; the unique index keeps the pair unique from now on. +""" + +import sqlalchemy as sa +from alembic import op + +revision = "53c10126d551" +down_revision = "0cc3158b50b0" +branch_labels = None +depends_on = None + +_TABLE = "trace_app_config" +_INDEX = "trace_app_config_app_provider_uq" + +# MySQL cannot read the table a DELETE targets except through a derived table. +_DELETE_SUPERSEDED = sa.text( + """ + DELETE FROM trace_app_config + WHERE id IN ( + SELECT id FROM ( + SELECT older.id + FROM trace_app_config older + JOIN trace_app_config newer + ON newer.app_id = older.app_id + AND newer.tracing_provider = older.tracing_provider + AND ( + newer.updated_at > older.updated_at + OR (newer.updated_at = older.updated_at AND newer.id > older.id) + ) + ) superseded + ) + """ +) + + +def upgrade() -> None: + op.execute(_DELETE_SUPERSEDED) + op.create_index(_INDEX, _TABLE, ["app_id", "tracing_provider"], unique=True) + + +def downgrade() -> None: + op.drop_index(_INDEX, table_name=_TABLE) diff --git a/api/models/model.py b/api/models/model.py index 88b1909b03ec4f..08071d3174f73e 100644 --- a/api/models/model.py +++ b/api/models/model.py @@ -2640,6 +2640,7 @@ class TraceAppConfig(TypeBase): __table_args__ = ( sa.PrimaryKeyConstraint("id", name="tracing_app_config_pkey"), sa.Index("trace_app_config_app_id_idx", "app_id"), + sa.Index("trace_app_config_app_provider_uq", "app_id", "tracing_provider", unique=True), ) id: Mapped[str] = mapped_column( diff --git a/api/openapi/markdown/console-openapi.md b/api/openapi/markdown/console-openapi.md index f7a1a956b08f2c..41d1fd96830b0e 100644 --- a/api/openapi/markdown/console-openapi.md +++ b/api/openapi/markdown/console-openapi.md @@ -1457,7 +1457,7 @@ Get available TTS voices for an Agent and language | Code | Description | Schema | | ---- | ----------- | ------ | | 200 | Provider support declarations | **application/json**: [TraceProviderCapabilitiesResponse](#traceprovidercapabilitiesresponse)
| -| 400 | Bad Request | **application/json**: [UnsupportedTracingProviderError](#unsupportedtracingprovidererror)
[ErrorResponse](#errorresponse)
| +| 400 | Bad Request | **application/json**: [TracingProviderUnavailableError](#tracingproviderunavailableerror)
[UnsupportedTracingProviderError](#unsupportedtracingprovidererror)
[ErrorResponse](#errorresponse)
| | default | Error | **application/json**: [ErrorResponse](#errorresponse)
| ### [DELETE] /agent/{agent_id}/trace-config @@ -3829,7 +3829,7 @@ Update app tracing configuration | Code | Description | Schema | | ---- | ----------- | ------ | | 200 | Provider support declarations | **application/json**: [TraceProviderCapabilitiesResponse](#traceprovidercapabilitiesresponse)
| -| 400 | Bad Request | **application/json**: [UnsupportedTracingProviderError](#unsupportedtracingprovidererror)
[ErrorResponse](#errorresponse)
| +| 400 | Bad Request | **application/json**: [TracingProviderUnavailableError](#tracingproviderunavailableerror)
[UnsupportedTracingProviderError](#unsupportedtracingprovidererror)
[ErrorResponse](#errorresponse)
| | default | Error | **application/json**: [ErrorResponse](#errorresponse)
| ### [DELETE] /apps/{app_id}/trace-config @@ -35957,13 +35957,14 @@ Enum class for tool provider | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | app_id | string | | No | +| configured_secrets | [ string ] | Secret fields that hold a value; submit them empty to keep that value | No | | created_at | string | | No | | error | string | | No | | has_not_configured | boolean | | No | | id | string | | No | | is_active | boolean | | No | | result | string | | No | -| tracing_config | object | | No | +| tracing_config | object | Non-secret fields; secrets are write-only and never returned | No | | tracing_provider | string | | No | | updated_at | string | | No | diff --git a/api/providers/trace/README.md b/api/providers/trace/README.md index 0f280c6ee84499..4e445daceb0f1b 100644 --- a/api/providers/trace/README.md +++ b/api/providers/trace/README.md @@ -1,77 +1,132 @@ # Trace providers -This directory holds **optional workspace packages** that send Dify **ops tracing** data (workflows, messages, tools, moderation, etc.) to an external observability backend (Langfuse, LangSmith, OpenTelemetry-style exporters, and others). +Each directory here is an optional uv workspace package that sends Dify traces (messages, workflow runs, tool calls, retrieval, moderation, and so on) to one external observability backend. `trace-sdk/` is the contract every provider implements; it is not a provider itself. -Unlike VDB providers, trace plugins are **not** discovered via entry points. The API core imports your package **explicitly** from `core/ops/ops_trace_manager.py` after you register the provider id and mapping. +## The contract (`trace-sdk`, imported as `dify_trace_sdk`) -## Architecture +A provider is a subclass of `dify_trace_sdk.provider.TraceProvider[ConfigT]`: -| Layer | Location | Role | -|--------|----------|------| -| Contracts | `api/core/ops/base_trace_instance.py`, `api/core/ops/entities/trace_entity.py`, `api/core/ops/entities/config_entity.py` | `BaseTraceInstance`, `BaseTracingConfig`, and typed `*TraceInfo` payloads | -| Registry | `api/core/ops/ops_trace_manager.py` | `TracingProviderEnum`, `OpsTraceProviderConfigMap` — maps provider **string** → config class, encrypted keys, and trace class | -| Your package | `api/providers/trace/trace-/` | Pydantic config + subclass of `BaseTraceInstance` | +| Member | Purpose | +|--------|---------| +| `id` | The stable provider ID stored in app tracing settings, such as `"langfuse"`. | +| `config_model` | A `ProviderConfig` subclass. Mark credentials with `secret()` and values that verification resolves with `derived()`. | +| `capabilities` | A `ProviderCapabilities` value: whether the backend prices usage natively and links to traces and spans. | +| `verify(config)` | Checks a new or edited configuration against the backend and returns it with derived fields filled. Raises `VerificationError(code)` with a stable code that contains no secret. | +| `project_url(config)` | The console link for a configuration. It must not perform I/O. | +| `send(config, trace)` | Delivers one trace unit and classifies every failure (see below). | -At runtime, `OpsTraceManager` decrypts stored credentials, builds your config model, caches a trace instance, and calls `trace(trace_info)` with a concrete `BaseTraceInfo` subtype. +What `send` receives (`dify_trace_sdk.model`): -## What you implement +- A `Trace` is a parent-first tree of completed `Span`s that uses OpenTelemetry ID formats. + - Each span carries timing, inputs and outputs, status, and typed LLM fields (`model_name`, `model_provider`, `usage`, `cost`, `time_to_first_token`). + - All values are strict JSON. +- A message trace has a `message` root and an `llm` child. +- A workflow run has a `workflow` span with node spans and loop or iteration wrapper spans. For chatflows these hang under the `message` root. +- Auxiliary operations join their message's trace as separate units. `trace.parent` names the message root span, and the first span of the unit points at it. +- A workflow nested in another one is its own trace. Its root carries `parent_workflow_run_id` and `parent_node_execution_id`. +- A unit can be sent again after a retryable failure. Derive vendor IDs deterministically from the trace and span IDs (`dify_trace_sdk.ids`), so a resend updates the same records instead of duplicating them. + +Failure classes (`dify_trace_sdk.errors`): + +| Raise | When | Pipeline behaviour | +|-------|------|--------------------| +| `ExportRetryableError` | Nothing was accepted and the cause is transient (unreachable, 408/425/429/5xx). | Sends the same unit again later, up to `OPS_TRACE_EXPORT_MAX_ATTEMPTS` attempts in total. | +| `ExportRejectedError` | The backend refused the unit, or the configuration cannot work. | Drops the unit and logs a warning. | +| `ExportOutcomeUnknownError` | Part or all of the unit may have been accepted. | Drops the unit; sending again could duplicate it. | + +Error messages are logged and must never contain credentials, request bodies, or vendor response bodies. Raise them `from None`. + +## Network access + +A provider never opens sockets, starts threads, or reads environment variables or files. All traffic goes through `self.egress`: + +- `egress.request(...)` sends one HTTP request through the SSRF proxy, with no retries or redirects, a bounded response, and no ambient cookies. +- `egress.grpc_channel(endpoint)` opens a channel pinned to a vetted address. +- The helpers in `dify_trace_sdk.exporting` and `dify_trace_sdk.otlp` turn egress failures and HTTP or gRPC statuses into the failure classes above. + +If you must use a vendor SDK, route its HTTP through `EgressHttpxTransport` or `EgressRequestsSession` (`dify_trace_sdk.transports`) and disable its batching, background threads, and retries. Importing the SDK must not change process-wide state; some SDKs install a Sentry client or patch other libraries on import. + +## Configuration and secrets + +Stored configurations keep the provider's fields at the top level, with `secret()` fields encrypted per workspace. Secrets are write-only: + +- The console never receives them back; it gets the names of the secrets that are set. +- When an edit leaves a secret empty or omits it, the stored value is kept. +- Clients cannot set `derived()` fields. + +`ProviderConfig` hides inputs from validation errors; keep validator messages fixed texts too. + +## Adding a provider + +1. Create `api/providers/trace/trace-/` with: + - a `pyproject.toml` for the project `dify-trace-`, depending on `dify-trace-sdk` (plus any vendor SDK that is safe to import); + - `src/dify_trace_/{config.py,provider.py,py.typed}`; + - tests that mirror the sources, e.g. `src/dify_trace_/provider.py` maps to `tests/unit_tests/src/dify_trace_/test_provider.py`. + + List `py.typed` under `[tool.setuptools.package-data]`. +2. Register the provider class as an entry point: + + ```toml + [project.entry-points."dify.trace_providers"] + = "dify_trace_.provider:Provider" + ``` + +3. Add the ID to `SUPPORTED_PROVIDERS` in `api/services/tracing/providers.py` and the provider's form fields to the web console's tracing field registry. +4. In `api/pyproject.toml`, add `dify-trace- = { workspace = true }` under `[tool.uv.sources]` and add the package to the `trace-` and `trace-all` dependency groups. Then run `uv lock` and `uv sync` from `api/`. +5. Test with `dify_trace_sdk.testing.RecordingEgress` and the sample traces there. Tests must not reach the network. + +`trace-langfuse/` is the smallest reference implementation; `trace-arize-phoenix/` and `trace-tencent/` show OTLP over HTTP and gRPC. + +## How traces reach a provider + +1. Producers emit typed events (`dify_trace_sdk.events`) through a `TraceScope` after their records are committed: + - generators, task pipelines and the workflow persistence layer; + - moderation, tools, retrieval, conversation naming and suggested questions. + + A scope exists only when the app has tracing enabled or enterprise telemetry is on. + The producer snapshots the event and hands it to a background publisher, so storage and + broker I/O do not delay the execution. Each process buffers at most 256 pending events; + a full buffer drops new events with a warning. Graceful shutdown drains accepted events, + but events not yet published can be lost if the process crashes. +2. `tasks.trace_delivery_task` hands each event on `OPS_TRACE_QUEUE` (default `ops_trace_v2`) to `services.tracing.delivery.TraceDeliveryService`: + - `repositories.trace_record_repository` reloads the records through the event's workspace. + - Enterprise telemetry receives every record. + - The Agent trace collector (`services.tracing.service`) takes the executions it manages. + - `services.tracing.canonical` builds the trace for the app's current provider configuration, and the provider's `send` runs. +3. A retryable failure stores only the built trace and schedules that send again. Payloads that an interrupted delivery leaves in storage are removed daily by `sweep_trace_payloads`. + +The tracing services depend only on the ports they declare: records, destinations, the provider catalog, secret encryption, the retry queue and the Agent trace ledger. `extensions/application_services/tracing.py` is the one place that chooses their adapters (`services/tracing/adapters.py`, `services/tracing/egress.py`, `services/tracing/registry.py` and the tracing repositories) and reads tracing configuration. + +## Rolling upgrade from the previous pipeline + +The new task names and payloads are incompatible with old workers. Use separate queues while +both versions run. `OPS_TRACE_QUEUE` defaults to `ops_trace_v2` for event delivery, export retries, +Agent trace exports/reconciliation, payload sweeping, and the sandbox metering task that previously +shared `ops_trace`. Celery routing is configured in `extensions/ext_celery.py`. +This isolates tracing tasks; other queues and database migrations still need their normal +compatibility checks when running both versions. + +1. Apply the database migrations and start new workers with `OPS_TRACE_QUEUE=ops_trace_v2`. + The Docker and development worker launchers include this queue by default and exclude + `ops_trace`. Explicit `CELERY_QUEUES`, `CELERY_WORKER_QUEUES`, or `--queues` lists must select + `ops_trace_v2` for the new pool. Keep the old tracing worker pool subscribed only to `ops_trace`. +2. Roll out the new API and all other producers with the same setting. Replace the old Beat + scheduler with the new one; do not run duplicate schedules. Allow old in-flight requests, + workflow jobs, and their buffered trace producers to finish before stopping them. +3. Let the old workers finish exporting their tasks. Check broker messages **and** worker + active, reserved/prefetched, scheduled/ETA, and retry deliveries. An empty broker queue alone + does not prove it has drained. Then stop the old workers gracefully and confirm no old + producers or workers remain. The new sweeper never deletes legacy `ops_trace/` payloads. +4. Keeping `ops_trace_v2` is supported. To restore the old queue name, first have **new-version** + workers consume both `ops_trace_v2,ops_trace`. Set `OPS_TRACE_QUEUE=ops_trace` on every new API, + worker, and Beat process and roll them out. Keep consumers for both queues until the new + queue, active tasks, reservations, ETA tasks, and retries drain; a redelivery can retain its + original queue. Then remove `ops_trace_v2` from worker subscriptions. + +Do not use the new version's retired `process_trace_tasks` handler to drain old messages: it can +only discard them, not export the legacy format. Do not purge queues or legacy payloads during +the transition. Returning to old binaries after publishing new-format tasks also requires +keeping new-version consumers available until those tasks finish. -### 1. Config model (`BaseTracingConfig`) - -Subclass `BaseTracingConfig` from `core.ops.entities.config_entity`. Use Pydantic validators; reuse helpers from `core.ops.utils` (for example `validate_url`, `validate_url_with_path`, `validate_project_name`) where appropriate. - -Fields fall into two groups used by the manager: - -- **`secret_keys`** — names of fields that are **encrypted at rest** (API keys, tokens, passwords). -- **`other_keys`** — non-secret connection settings (hosts, project names, endpoints). - -List these key names in your `OpsTraceProviderConfigMap` entry so encrypt/decrypt and merge logic stay correct. - -### 2. Trace instance (`BaseTraceInstance`) - -Subclass `BaseTraceInstance` and implement: - -```python -def trace(self, trace_info: BaseTraceInfo) -> None: ... -``` - -Dispatch on the concrete type with `isinstance` (see `trace_langfuse` or `trace_langsmith` for full patterns). Payload types are defined in `core/ops/entities/trace_entity.py`, including: - -- `WorkflowTraceInfo`, `WorkflowNodeTraceInfo`, `DraftNodeExecutionTrace` -- `MessageTraceInfo`, `ToolTraceInfo`, `ModerationTraceInfo`, `SuggestedQuestionTraceInfo` -- `DatasetRetrievalTraceInfo`, `GenerateNameTraceInfo`, `PromptGenerationTraceInfo` - -You may ignore categories your backend does not support; existing providers often no-op unhandled types. - -Optional: use `get_service_account_with_tenant(app_id)` from the base class when you need tenant-scoped account context. - -### 3. Register in the API core - -Upstream changes are required so Dify knows your provider exists: - -1. **`TracingProviderEnum`** (`api/core/ops/entities/config_entity.py`) — add a new member whose **value** is the stable string stored in app tracing config (e.g. `"mybackend"`). -2. **`OpsTraceProviderConfigMap.__getitem__`** (`api/core/ops/ops_trace_manager.py`) — add a `match` case for that enum member returning: - - `config_class`: your Pydantic config type - - `secret_keys` / `other_keys`: lists of field names as above - - `trace_instance`: your `BaseTraceInstance` subclass - Lazy-import your package inside the case so missing optional installs raise a clear `ImportError`. - -If the `match` case is missing, the provider string will not resolve and tracing will be disabled for that app. - -## Package layout - -Each provider is a normal uv workspace member, for example: - -- `api/providers/trace/trace-/pyproject.toml` — project name `dify-trace-`, dependencies on vendor SDKs -- `api/providers/trace/trace-/src/dify_trace_/` — `config.py`, `_trace.py`, optional `entities/`, and an empty **`py.typed`** file (PEP 561) so the API type checker treats the package as typed; list `py.typed` under `[tool.setuptools.package-data]` for that import name in `pyproject.toml`. - -Reference implementations: `trace-langfuse/`, `trace-langsmith/`, `trace-opik/`. - -## Wiring into the `api` workspace - -In `api/pyproject.toml`: - -1. **`[tool.uv.sources]`** — `dify-trace- = { workspace = true }` -2. **`[dependency-groups]`** — add `trace- = ["dify-trace-"]` and include `dify-trace-` in `trace-all` if it should ship with the default bundle - -After changing metadata, run **`uv sync`** from `api/`. +Saved provider configuration and credentials need no changes. Queue migration is deployment +configuration, shared by the API, workers, and Beat; it does not change app tracing settings. diff --git a/api/providers/trace/trace-aliyun/pyproject.toml b/api/providers/trace/trace-aliyun/pyproject.toml index bcef7e9fb15a1a..6e2732b7f4801c 100644 --- a/api/providers/trace/trace-aliyun/pyproject.toml +++ b/api/providers/trace/trace-aliyun/pyproject.toml @@ -2,13 +2,12 @@ name = "dify-trace-aliyun" version = "0.0.1" dependencies = [ - # versions inherited from parent - "opentelemetry-api", - "opentelemetry-exporter-otlp-proto-grpc", - "opentelemetry-sdk", - "opentelemetry-semantic-conventions", + "dify-trace-sdk", ] description = "Dify ops tracing provider (Aliyun)." +[project.entry-points."dify.trace_providers"] +aliyun = "dify_trace_aliyun.provider:AliyunProvider" + [tool.setuptools.packages.find] where = ["src"] diff --git a/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/agent_trace.py b/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/agent_trace.py deleted file mode 100644 index ec1b1a39fb3d5f..00000000000000 --- a/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/agent_trace.py +++ /dev/null @@ -1,27 +0,0 @@ -"""Aliyun LLM tracing export with synchronous OTLP/HTTP acknowledgment.""" - -from collections.abc import Mapping -from typing import Any - -from core.ops.unified_trace.entities import CanonicalTrace -from dify_trace_aliyun.configuration import AliyunConfigurationError, ingestion_url -from services.tracing.genai_attributes import attributes -from services.tracing.otlp_export import build_request, export_http -from services.tracing.provider_dispatch import TraceExportRejectedError, TraceExportResult - - -def export_trace(provider: str, config: Mapping[str, Any], trace: CanonicalTrace) -> TraceExportResult: - try: - url = ingestion_url(config) - except AliyunConfigurationError: - raise TraceExportRejectedError("Aliyun tracing configuration is invalid") from None - payload = build_request( - trace, - { - "service.name": str(config.get("app_name") or "dify_app"), - "acs.arms.service.feature": "genai_app", - }, - lambda span: attributes(span, trace, provider), - ) - export_http(provider, url, {}, payload) - return TraceExportResult(trace_id=trace.trace_id) diff --git a/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/aliyun_trace.py b/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/aliyun_trace.py deleted file mode 100644 index 5b7efaca491ec3..00000000000000 --- a/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/aliyun_trace.py +++ /dev/null @@ -1,850 +0,0 @@ -import logging -from collections.abc import Mapping, Sequence -from typing import Any, override - -from opentelemetry.trace import SpanKind -from sqlalchemy.orm import sessionmaker - -from core.ops.base_trace_instance import BaseTraceInstance -from core.ops.entities.trace_entity import ( - BaseTraceInfo, - DatasetRetrievalTraceInfo, - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - SuggestedQuestionTraceInfo, - ToolTraceInfo, - WorkflowTraceInfo, -) -from core.repositories import DifyCoreRepositoryFactory -from dify_trace_aliyun.config import AliyunConfig -from dify_trace_aliyun.configuration import console_url, ingestion_url, verify_configuration -from dify_trace_aliyun.data_exporter.traceclient import ( - TraceClient, - convert_datetime_to_nanoseconds, - convert_to_span_id, - convert_to_trace_id, - generate_span_id, -) -from dify_trace_aliyun.entities.aliyun_trace_entity import SpanData, TraceMetadata -from dify_trace_aliyun.entities.semconv import ( - DIFY_APP_ID, - DIFY_NODE_TYPE, - GEN_AI_AGENT_NAME, - GEN_AI_COMPLETION, - GEN_AI_INPUT_MESSAGE, - GEN_AI_OPERATION_NAME, - GEN_AI_OUTPUT_MESSAGE, - GEN_AI_PROMPT, - GEN_AI_PROVIDER_NAME, - GEN_AI_REACT_FINISH_REASON, - GEN_AI_REACT_ROUND, - GEN_AI_REQUEST_MODEL, - GEN_AI_RESPONSE_FINISH_REASON, - GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN, - GEN_AI_USAGE_INPUT_TOKENS, - GEN_AI_USAGE_OUTPUT_TOKENS, - GEN_AI_USAGE_TOTAL_TOKENS, - OPERATION_NAME_CHAT, - OPERATION_NAME_INVOKE_AGENT, - OPERATION_NAME_REACT, - RETRIEVAL_DOCUMENT, - RETRIEVAL_QUERY, - GenAISpanKind, -) -from dify_trace_aliyun.utils import ( - AgentLogEntry, - convert_seconds_to_nanoseconds, - create_common_span_attributes, - create_gen_ai_tool_attributes, - create_links_from_trace_id, - create_status_from_agent_log_entry, - create_status_from_error, - extract_model_name_from_thought_label, - extract_react_round_number, - extract_retrieval_documents, - extract_tool_description, - extract_tool_name_from_call_label, - format_input_messages, - format_output_messages, - format_retrieval_documents, - get_user_id_from_message_data, - get_workflow_node_status, - is_llm_thought_entry, - is_tool_call_entry, - map_gen_ai_tool_type, - parse_agent_log_entries, - serialize_json_data, -) -from extensions.ext_database import db -from graphon.entities import WorkflowNodeExecution -from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey -from models import WorkflowNodeExecutionTriggeredFrom - -logger = logging.getLogger(__name__) - - -class AliyunDataTrace(BaseTraceInstance): - def __init__( - self, - aliyun_config: AliyunConfig, - ): - super().__init__(aliyun_config) - self._config = aliyun_config.model_copy(deep=True) - self.trace_client = TraceClient( - service_name=self._config.app_name, endpoint=ingestion_url(self._config.model_dump()) - ) - - @override - def trace(self, trace_info: BaseTraceInfo): - with self.trace_client.batch(): - match trace_info: - case WorkflowTraceInfo(): - self.workflow_trace(trace_info) - case MessageTraceInfo(): - self.message_trace(trace_info) - case ModerationTraceInfo(): - pass - case SuggestedQuestionTraceInfo(): - self.suggested_question_trace(trace_info) - case DatasetRetrievalTraceInfo(): - self.dataset_retrieval_trace(trace_info) - case ToolTraceInfo(): - self.tool_trace(trace_info) - case GenerateNameTraceInfo(): - pass - case _: - pass - - def api_check(self) -> bool: - verify_configuration(self._config) - return True - - def get_project_url(self): - return console_url() - - def _extract_app_id(self, trace_info: BaseTraceInfo) -> str: - """Extract app_id from trace_info, trying metadata first then message_data.""" - app_id = trace_info.metadata.get("app_id") - if app_id: - return str(app_id) - message_data = getattr(trace_info, "message_data", None) - if message_data is not None: - return str(getattr(message_data, "app_id", "")) - return "" - - def workflow_trace(self, trace_info: WorkflowTraceInfo): - trace_metadata = TraceMetadata( - trace_id=convert_to_trace_id(trace_info.workflow_run_id), - workflow_span_id=convert_to_span_id(trace_info.workflow_run_id, "workflow"), - session_id=trace_info.metadata.get("conversation_id") or "", - user_id=str(trace_info.metadata.get("user_id") or ""), - links=create_links_from_trace_id(trace_info.trace_id), - ) - - self.add_workflow_span(trace_info, trace_metadata) - - workflow_node_executions = self.get_workflow_node_executions(trace_info) - for node_execution in workflow_node_executions: - node_span = self.build_workflow_node_span(node_execution, trace_info, trace_metadata) - self.trace_client.add_span(node_span) - if node_span is not None and node_execution.node_type == BuiltinNodeTypes.AGENT: - for react_span in self.build_agent_react_spans(node_execution, trace_metadata): - self.trace_client.add_span(react_span) - - def message_trace(self, trace_info: MessageTraceInfo): - message_data = trace_info.message_data - if message_data is None: - return - - message_id = trace_info.message_id - user_id = get_user_id_from_message_data(message_data) - status = create_status_from_error(trace_info.error) - - trace_metadata = TraceMetadata( - trace_id=convert_to_trace_id(message_id), - workflow_span_id=0, - session_id=trace_info.metadata.get("conversation_id") or "", - user_id=user_id, - links=create_links_from_trace_id(trace_info.trace_id), - ) - - inputs_json = serialize_json_data(trace_info.inputs) - outputs_str = str(trace_info.outputs) - - message_span_id = convert_to_span_id(message_id, "message") - message_span = SpanData( - trace_id=trace_metadata.trace_id, - parent_span_id=None, - span_id=message_span_id, - name="message", - start_time=convert_datetime_to_nanoseconds(trace_info.start_time), - end_time=convert_datetime_to_nanoseconds(trace_info.end_time), - attributes={ - **create_common_span_attributes( - session_id=trace_metadata.session_id, - user_id=trace_metadata.user_id, - span_kind=GenAISpanKind.CHAIN, - inputs=inputs_json, - outputs=outputs_str, - ), - DIFY_APP_ID: self._extract_app_id(trace_info), - }, - status=status, - links=trace_metadata.links, - span_kind=SpanKind.SERVER, - ) - self.trace_client.add_span(message_span) - - llm_attributes: dict[str, Any] = { - **create_common_span_attributes( - session_id=trace_metadata.session_id, - user_id=trace_metadata.user_id, - span_kind=GenAISpanKind.LLM, - inputs=inputs_json, - outputs=outputs_str, - ), - GEN_AI_OPERATION_NAME: OPERATION_NAME_CHAT, - GEN_AI_REQUEST_MODEL: trace_info.metadata.get("ls_model_name") or "", - GEN_AI_PROVIDER_NAME: trace_info.metadata.get("ls_provider") or "", - GEN_AI_USAGE_INPUT_TOKENS: str(trace_info.message_tokens), - GEN_AI_USAGE_OUTPUT_TOKENS: str(trace_info.answer_tokens), - GEN_AI_USAGE_TOTAL_TOKENS: str(trace_info.total_tokens), - GEN_AI_PROMPT: inputs_json, - GEN_AI_COMPLETION: outputs_str, - } - if trace_info.gen_ai_server_time_to_first_token is not None: - llm_attributes[GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN] = convert_seconds_to_nanoseconds( - trace_info.gen_ai_server_time_to_first_token - ) - - llm_span = SpanData( - trace_id=trace_metadata.trace_id, - parent_span_id=message_span_id, - span_id=convert_to_span_id(message_id, "llm"), - name="llm", - start_time=convert_datetime_to_nanoseconds(trace_info.start_time), - end_time=convert_datetime_to_nanoseconds(trace_info.end_time), - attributes=llm_attributes, - status=status, - links=trace_metadata.links, - ) - self.trace_client.add_span(llm_span) - - def dataset_retrieval_trace(self, trace_info: DatasetRetrievalTraceInfo): - if trace_info.message_data is None: - return - - message_id = trace_info.message_id - - trace_metadata = TraceMetadata( - trace_id=convert_to_trace_id(message_id), - workflow_span_id=0, - session_id=trace_info.metadata.get("conversation_id") or "", - user_id=str(trace_info.metadata.get("user_id") or ""), - links=create_links_from_trace_id(trace_info.trace_id), - ) - - documents_data = extract_retrieval_documents(trace_info.documents) - documents_json = serialize_json_data(documents_data) - inputs_str = str(trace_info.inputs) - - dataset_retrieval_span = SpanData( - trace_id=trace_metadata.trace_id, - parent_span_id=convert_to_span_id(message_id, "message"), - span_id=generate_span_id(), - name="dataset_retrieval", - start_time=convert_datetime_to_nanoseconds(trace_info.start_time), - end_time=convert_datetime_to_nanoseconds(trace_info.end_time), - attributes={ - **create_common_span_attributes( - session_id=trace_metadata.session_id, - user_id=trace_metadata.user_id, - span_kind=GenAISpanKind.RETRIEVER, - inputs=inputs_str, - outputs=documents_json, - ), - RETRIEVAL_QUERY: inputs_str, - RETRIEVAL_DOCUMENT: documents_json, - }, - links=trace_metadata.links, - ) - self.trace_client.add_span(dataset_retrieval_span) - - def tool_trace(self, trace_info: ToolTraceInfo): - if trace_info.message_data is None: - return - - message_id = trace_info.message_id - status = create_status_from_error(trace_info.error) - - trace_metadata = TraceMetadata( - trace_id=convert_to_trace_id(message_id), - workflow_span_id=0, - session_id=trace_info.metadata.get("conversation_id") or "", - user_id=str(trace_info.metadata.get("user_id") or ""), - links=create_links_from_trace_id(trace_info.trace_id), - ) - - tool_config = trace_info.tool_config if isinstance(trace_info.tool_config, Mapping) else {} - tool_inputs_json = serialize_json_data(trace_info.tool_inputs) - tool_result = str(trace_info.tool_outputs) - inputs_json = serialize_json_data(trace_info.inputs) - provider_type = tool_config.get("tool_provider_type") or tool_config.get("provider_type") - - tool_span = SpanData( - trace_id=trace_metadata.trace_id, - parent_span_id=convert_to_span_id(message_id, "message"), - span_id=generate_span_id(), - name=trace_info.tool_name, - start_time=convert_datetime_to_nanoseconds(trace_info.start_time), - end_time=convert_datetime_to_nanoseconds(trace_info.end_time), - attributes={ - **create_common_span_attributes( - session_id=trace_metadata.session_id, - user_id=trace_metadata.user_id, - span_kind=GenAISpanKind.TOOL, - inputs=inputs_json, - outputs=tool_result, - ), - **create_gen_ai_tool_attributes( - tool_name=trace_info.tool_name, - tool_type=map_gen_ai_tool_type(str(provider_type) if provider_type else None), - tool_description=extract_tool_description(tool_config), - tool_call_id=str(trace_info.metadata.get("node_execution_id") or ""), - tool_call_arguments=tool_inputs_json, - tool_call_result=tool_result, - ), - }, - status=status, - links=trace_metadata.links, - ) - self.trace_client.add_span(tool_span) - - def get_workflow_node_executions(self, trace_info: WorkflowTraceInfo) -> Sequence[WorkflowNodeExecution]: - app_id = trace_info.metadata.get("app_id") - if not app_id: - raise ValueError("No app_id found in trace_info metadata") - - service_account = self.get_service_account_with_tenant(app_id) - - session_factory = sessionmaker(bind=db.engine) - workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( - session_factory=session_factory, - tenant_id=trace_info.tenant_id, - user=service_account, - app_id=app_id, - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - - return workflow_node_execution_repository.get_by_workflow_execution( - workflow_execution_id=trace_info.workflow_run_id - ) - - def build_workflow_node_span( - self, node_execution: WorkflowNodeExecution, trace_info: WorkflowTraceInfo, trace_metadata: TraceMetadata - ): - try: - if node_execution.node_type == BuiltinNodeTypes.LLM: - node_span = self.build_workflow_llm_span(trace_info, node_execution, trace_metadata) - elif node_execution.node_type == BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL: - node_span = self.build_workflow_retrieval_span(trace_info, node_execution, trace_metadata) - elif node_execution.node_type == BuiltinNodeTypes.TOOL: - node_span = self.build_workflow_tool_span(trace_info, node_execution, trace_metadata) - elif node_execution.node_type == BuiltinNodeTypes.AGENT: - node_span = self.build_workflow_agent_span(trace_info, node_execution, trace_metadata) - else: - node_span = self.build_workflow_task_span(trace_info, node_execution, trace_metadata) - return node_span - except Exception as e: - logger.warning("Error occurred in build_workflow_node_span: %s", e, exc_info=True) - return None - - def build_workflow_task_span( - self, trace_info: WorkflowTraceInfo, node_execution: WorkflowNodeExecution, trace_metadata: TraceMetadata - ) -> SpanData: - inputs = node_execution.inputs - if node_execution.node_type == BuiltinNodeTypes.HTTP_REQUEST and isinstance( - node_execution.process_data, Mapping - ): - request = node_execution.process_data.get("request") - if isinstance(request, str) and request.strip(): - # Reuse the node's masked request log; do not reconstruct the outbound request. - inputs = {"request": request} - inputs_json = serialize_json_data(inputs) - outputs_json = serialize_json_data(node_execution.outputs) - return SpanData( - trace_id=trace_metadata.trace_id, - parent_span_id=trace_metadata.workflow_span_id, - span_id=convert_to_span_id(node_execution.id, "node"), - name=node_execution.title, - start_time=convert_datetime_to_nanoseconds(node_execution.created_at), - end_time=convert_datetime_to_nanoseconds(node_execution.finished_at), - attributes={ - **create_common_span_attributes( - session_id=trace_metadata.session_id, - user_id=trace_metadata.user_id, - span_kind=GenAISpanKind.TASK, - inputs=inputs_json, - outputs=outputs_json, - ), - DIFY_NODE_TYPE: node_execution.node_type, - }, - status=get_workflow_node_status(node_execution), - links=trace_metadata.links, - ) - - def build_workflow_tool_span( - self, trace_info: WorkflowTraceInfo, node_execution: WorkflowNodeExecution, trace_metadata: TraceMetadata - ) -> SpanData: - tool_info: Mapping[str, Any] = {} - if isinstance(node_execution.metadata, Mapping): - raw_tool_info = node_execution.metadata.get(WorkflowNodeExecutionMetadataKey.TOOL_INFO, {}) - if isinstance(raw_tool_info, Mapping): - tool_info = raw_tool_info - - inputs_json = serialize_json_data(node_execution.inputs or {}) - outputs_json = serialize_json_data(node_execution.outputs) - provider_type = tool_info.get("provider_type") - - return SpanData( - trace_id=trace_metadata.trace_id, - parent_span_id=trace_metadata.workflow_span_id, - span_id=convert_to_span_id(node_execution.id, "node"), - name=node_execution.title, - start_time=convert_datetime_to_nanoseconds(node_execution.created_at), - end_time=convert_datetime_to_nanoseconds(node_execution.finished_at), - attributes={ - **create_common_span_attributes( - session_id=trace_metadata.session_id, - user_id=trace_metadata.user_id, - span_kind=GenAISpanKind.TOOL, - inputs=inputs_json, - outputs=outputs_json, - ), - **create_gen_ai_tool_attributes( - tool_name=node_execution.title, - tool_type=map_gen_ai_tool_type(str(provider_type) if provider_type else None), - tool_description=extract_tool_description(tool_info), - tool_call_id=str(node_execution.id or ""), - tool_call_arguments=inputs_json, - tool_call_result=outputs_json, - ), - }, - status=get_workflow_node_status(node_execution), - links=trace_metadata.links, - ) - - def build_workflow_retrieval_span( - self, trace_info: WorkflowTraceInfo, node_execution: WorkflowNodeExecution, trace_metadata: TraceMetadata - ) -> SpanData: - input_value = str(node_execution.inputs.get("query", "")) if node_execution.inputs else "" - output_value = serialize_json_data(node_execution.outputs.get("result", [])) if node_execution.outputs else "" - - retrieval_documents = node_execution.outputs.get("result", []) if node_execution.outputs else [] - semantic_retrieval_documents = format_retrieval_documents(retrieval_documents) - semantic_retrieval_documents_json = serialize_json_data(semantic_retrieval_documents) - - return SpanData( - trace_id=trace_metadata.trace_id, - parent_span_id=trace_metadata.workflow_span_id, - span_id=convert_to_span_id(node_execution.id, "node"), - name=node_execution.title, - start_time=convert_datetime_to_nanoseconds(node_execution.created_at), - end_time=convert_datetime_to_nanoseconds(node_execution.finished_at), - attributes={ - **create_common_span_attributes( - session_id=trace_metadata.session_id, - user_id=trace_metadata.user_id, - span_kind=GenAISpanKind.RETRIEVER, - inputs=input_value, - outputs=output_value, - ), - RETRIEVAL_QUERY: input_value, - RETRIEVAL_DOCUMENT: semantic_retrieval_documents_json, - }, - status=get_workflow_node_status(node_execution), - links=trace_metadata.links, - ) - - def build_workflow_llm_span( - self, trace_info: WorkflowTraceInfo, node_execution: WorkflowNodeExecution, trace_metadata: TraceMetadata - ) -> SpanData: - process_data = node_execution.process_data if isinstance(node_execution.process_data, Mapping) else {} - inputs = node_execution.inputs if isinstance(node_execution.inputs, Mapping) else {} - outputs = node_execution.outputs if isinstance(node_execution.outputs, Mapping) else {} - usage_data = process_data.get("usage", {}) if "usage" in process_data else outputs.get("usage", {}) - if not isinstance(usage_data, Mapping): - usage_data = {} - - # On invoke failure graphon leaves process_data empty, but prep already wrote - # model identity (and template variables / context) into node inputs. - prompts = process_data.get("prompts") or [] - if prompts: - prompts_json = serialize_json_data(prompts) - elif inputs: - prompts_json = serialize_json_data(inputs) - else: - prompts_json = serialize_json_data([]) - - text_output = str(outputs.get("text") or "") - if not text_output: - text_output = str(outputs.get("error_message") or node_execution.error or "") - - finish_reason = outputs.get("finish_reason") or outputs.get("error_type") or "" - model_name = process_data.get("model_name") or inputs.get("model_name") or "" - model_provider = process_data.get("model_provider") or inputs.get("model_provider") or "" - - gen_ai_input_message = format_input_messages(process_data) - gen_ai_output_message = format_output_messages(outputs) - - attributes: dict[str, Any] = { - **create_common_span_attributes( - session_id=trace_metadata.session_id, - user_id=trace_metadata.user_id, - span_kind=GenAISpanKind.LLM, - inputs=prompts_json, - outputs=text_output, - ), - GEN_AI_OPERATION_NAME: OPERATION_NAME_CHAT, - GEN_AI_REQUEST_MODEL: str(model_name), - GEN_AI_PROVIDER_NAME: str(model_provider), - GEN_AI_USAGE_INPUT_TOKENS: str(usage_data.get("prompt_tokens", 0)), - GEN_AI_USAGE_OUTPUT_TOKENS: str(usage_data.get("completion_tokens", 0)), - GEN_AI_USAGE_TOTAL_TOKENS: str(usage_data.get("total_tokens", 0)), - GEN_AI_PROMPT: prompts_json, - GEN_AI_COMPLETION: text_output, - GEN_AI_RESPONSE_FINISH_REASON: str(finish_reason), - GEN_AI_INPUT_MESSAGE: gen_ai_input_message, - GEN_AI_OUTPUT_MESSAGE: gen_ai_output_message, - } - time_to_first_token = usage_data.get("time_to_first_token") - if isinstance(time_to_first_token, (int, float)): - attributes[GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN] = convert_seconds_to_nanoseconds(float(time_to_first_token)) - - return SpanData( - trace_id=trace_metadata.trace_id, - parent_span_id=trace_metadata.workflow_span_id, - span_id=convert_to_span_id(node_execution.id, "node"), - name=node_execution.title, - start_time=convert_datetime_to_nanoseconds(node_execution.created_at), - end_time=convert_datetime_to_nanoseconds(node_execution.finished_at), - attributes=attributes, - status=get_workflow_node_status(node_execution), - links=trace_metadata.links, - ) - - def build_workflow_agent_span( - self, trace_info: WorkflowTraceInfo, node_execution: WorkflowNodeExecution, trace_metadata: TraceMetadata - ) -> SpanData: - """Build an AGENT-kind span for an agent-strategy node (instead of a generic TASK span).""" - inputs_json = serialize_json_data(node_execution.inputs) - outputs = node_execution.outputs if isinstance(node_execution.outputs, Mapping) else {} - usage_data = outputs.get("usage", {}) - if not isinstance(usage_data, Mapping): - usage_data = {} - text_output = str(outputs.get("text", "")) - - attributes: dict[str, Any] = { - **create_common_span_attributes( - session_id=trace_metadata.session_id, - user_id=trace_metadata.user_id, - span_kind=GenAISpanKind.AGENT, - inputs=inputs_json, - outputs=text_output, - ), - GEN_AI_OPERATION_NAME: OPERATION_NAME_INVOKE_AGENT, - GEN_AI_AGENT_NAME: node_execution.title, - GEN_AI_USAGE_INPUT_TOKENS: str(usage_data.get("prompt_tokens", 0)), - GEN_AI_USAGE_OUTPUT_TOKENS: str(usage_data.get("completion_tokens", 0)), - GEN_AI_USAGE_TOTAL_TOKENS: str(usage_data.get("total_tokens", 0)), - } - time_to_first_token = usage_data.get("time_to_first_token") - if isinstance(time_to_first_token, (int, float)): - attributes[GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN] = convert_seconds_to_nanoseconds(float(time_to_first_token)) - - return SpanData( - trace_id=trace_metadata.trace_id, - parent_span_id=trace_metadata.workflow_span_id, - span_id=convert_to_span_id(node_execution.id, "node"), - name=node_execution.title, - start_time=convert_datetime_to_nanoseconds(node_execution.created_at), - end_time=convert_datetime_to_nanoseconds(node_execution.finished_at), - attributes=attributes, - status=get_workflow_node_status(node_execution), - links=trace_metadata.links, - ) - - def build_agent_react_spans( - self, node_execution: WorkflowNodeExecution, trace_metadata: TraceMetadata - ) -> list[SpanData]: - """Build ReAct STEP spans and child LLM / TOOL spans from the agent execution log. - - The agent log lives in ``outputs["json"]``; ``started_at``/``finished_at`` there are - monotonic-clock seconds, so they are mapped onto wall-clock time by anchoring the - earliest ``started_at`` to the node's start time. Entries without timing fall back - to the node's start/end times. Returns an empty list when no log is available. - """ - try: - outputs = node_execution.outputs or {} - round_entries = parse_agent_log_entries(outputs) - if not round_entries: - return [] - - agent_span_id = convert_to_span_id(node_execution.id, "node") - node_start_ns = convert_datetime_to_nanoseconds(node_execution.created_at) - node_end_ns = convert_datetime_to_nanoseconds(node_execution.finished_at) - - monotonic_starts = [ - entry.metadata["started_at"] - for round_entry in round_entries - for entry in [round_entry, *round_entry.children] - if isinstance(entry.metadata.get("started_at"), (int, float)) - ] - base_monotonic = min(monotonic_starts) if monotonic_starts else None - - def to_wall_clock_ns(monotonic_seconds: Any, fallback: int | None) -> int | None: - if ( - isinstance(monotonic_seconds, (int, float)) - and base_monotonic is not None - and node_start_ns is not None - ): - return node_start_ns + convert_seconds_to_nanoseconds(float(monotonic_seconds) - base_monotonic) - return fallback - - spans: list[SpanData] = [] - for index, round_entry in enumerate(round_entries, start=1): - round_number = extract_react_round_number(round_entry.label, index) - step_span_id = generate_span_id() - step_attributes: dict[str, Any] = { - **create_common_span_attributes( - session_id=trace_metadata.session_id, - user_id=trace_metadata.user_id, - span_kind=GenAISpanKind.STEP, - inputs="", - outputs=serialize_json_data(round_entry.data), - ), - GEN_AI_OPERATION_NAME: OPERATION_NAME_REACT, - GEN_AI_REACT_ROUND: round_number, - } - if round_entry.error: - step_attributes[GEN_AI_REACT_FINISH_REASON] = "error" - spans.append( - SpanData( - trace_id=trace_metadata.trace_id, - parent_span_id=agent_span_id, - span_id=step_span_id, - name=round_entry.label or f"react step {round_number}", - start_time=to_wall_clock_ns(round_entry.metadata.get("started_at"), node_start_ns), - end_time=to_wall_clock_ns(round_entry.metadata.get("finished_at"), node_end_ns), - attributes=step_attributes, - status=create_status_from_agent_log_entry(round_entry), - links=trace_metadata.links, - ) - ) - - for child in round_entry.children: - child_start = to_wall_clock_ns(child.metadata.get("started_at"), node_start_ns) - child_end = to_wall_clock_ns(child.metadata.get("finished_at"), node_end_ns) - if is_tool_call_entry(child): - spans.append( - self._build_agent_tool_call_span( - entry=child, - step_span_id=step_span_id, - trace_metadata=trace_metadata, - start_time=child_start, - end_time=child_end, - ) - ) - elif is_llm_thought_entry(child): - spans.append( - self._build_agent_llm_call_span( - entry=child, - step_span_id=step_span_id, - trace_metadata=trace_metadata, - start_time=child_start, - end_time=child_end, - ) - ) - return spans - except Exception as e: - logger.warning("Error occurred in build_agent_react_spans: %s", e, exc_info=True) - return [] - - def _build_agent_llm_call_span( - self, - entry: AgentLogEntry, - step_span_id: int, - trace_metadata: TraceMetadata, - start_time: int | None, - end_time: int | None, - ) -> SpanData: - completion = str(entry.data.get("thought") or entry.data.get("action") or "") - return SpanData( - trace_id=trace_metadata.trace_id, - parent_span_id=step_span_id, - span_id=generate_span_id(), - name=entry.label or "llm", - start_time=start_time, - end_time=end_time, - attributes={ - **create_common_span_attributes( - session_id=trace_metadata.session_id, - user_id=trace_metadata.user_id, - span_kind=GenAISpanKind.LLM, - inputs="", - outputs=serialize_json_data(entry.data), - ), - GEN_AI_OPERATION_NAME: OPERATION_NAME_CHAT, - GEN_AI_REQUEST_MODEL: extract_model_name_from_thought_label(entry.label), - GEN_AI_PROVIDER_NAME: str(entry.metadata.get("provider") or ""), - GEN_AI_USAGE_TOTAL_TOKENS: str(entry.metadata.get("total_tokens", 0)), - GEN_AI_COMPLETION: completion, - }, - status=create_status_from_agent_log_entry(entry), - links=trace_metadata.links, - ) - - def _build_agent_tool_call_span( - self, - entry: AgentLogEntry, - step_span_id: int, - trace_metadata: TraceMetadata, - start_time: int | None, - end_time: int | None, - ) -> SpanData: - tool_name = str(entry.data.get("tool_name") or extract_tool_name_from_call_label(entry.label) or "tool") - tool_parameters = entry.data.get("tool_call_args") - if tool_parameters is None: - tool_parameters = entry.data.get("tool_call_input") - if tool_parameters is None: - tool_parameters = entry.data - tool_arguments_json = serialize_json_data(tool_parameters) - tool_result = entry.data.get("output", entry.data) - tool_result_json = tool_result if isinstance(tool_result, str) else serialize_json_data(tool_result) - provider_type = entry.metadata.get("provider_type") or entry.data.get("provider_type") - return SpanData( - trace_id=trace_metadata.trace_id, - parent_span_id=step_span_id, - span_id=generate_span_id(), - name=entry.label or f"CALL {tool_name}", - start_time=start_time, - end_time=end_time, - attributes={ - **create_common_span_attributes( - session_id=trace_metadata.session_id, - user_id=trace_metadata.user_id, - span_kind=GenAISpanKind.TOOL, - inputs=tool_arguments_json, - outputs=tool_result_json, - ), - **create_gen_ai_tool_attributes( - tool_name=tool_name, - tool_type=map_gen_ai_tool_type(str(provider_type) if provider_type else None), - tool_description=extract_tool_description(entry.data) or extract_tool_description(entry.metadata), - tool_call_id=entry.id, - tool_call_arguments=tool_arguments_json, - tool_call_result=tool_result_json, - ), - }, - status=create_status_from_agent_log_entry(entry), - links=trace_metadata.links, - ) - - def add_workflow_span(self, trace_info: WorkflowTraceInfo, trace_metadata: TraceMetadata): - message_span_id = None - if trace_info.message_id: - message_span_id = convert_to_span_id(trace_info.message_id, "message") - status = create_status_from_error(trace_info.error) - - inputs_json = serialize_json_data(trace_info.workflow_run_inputs) - outputs_json = serialize_json_data(trace_info.workflow_run_outputs) - - app_id = self._extract_app_id(trace_info) - - if message_span_id: - message_span = SpanData( - trace_id=trace_metadata.trace_id, - parent_span_id=None, - span_id=message_span_id, - name="message", - start_time=convert_datetime_to_nanoseconds(trace_info.start_time), - end_time=convert_datetime_to_nanoseconds(trace_info.end_time), - attributes={ - **create_common_span_attributes( - session_id=trace_metadata.session_id, - user_id=trace_metadata.user_id, - span_kind=GenAISpanKind.CHAIN, - inputs=trace_info.workflow_run_inputs.get("sys.query") or "", - outputs=outputs_json, - ), - DIFY_APP_ID: app_id, - }, - status=status, - links=trace_metadata.links, - span_kind=SpanKind.SERVER, - ) - self.trace_client.add_span(message_span) - - workflow_span = SpanData( - trace_id=trace_metadata.trace_id, - parent_span_id=message_span_id, - span_id=trace_metadata.workflow_span_id, - name="workflow", - start_time=convert_datetime_to_nanoseconds(trace_info.start_time), - end_time=convert_datetime_to_nanoseconds(trace_info.end_time), - attributes={ - **create_common_span_attributes( - session_id=trace_metadata.session_id, - user_id=trace_metadata.user_id, - span_kind=GenAISpanKind.CHAIN, - inputs=inputs_json, - outputs=outputs_json, - ), - **({DIFY_APP_ID: app_id} if message_span_id is None else {}), - }, - status=status, - links=trace_metadata.links, - span_kind=SpanKind.SERVER if message_span_id is None else SpanKind.INTERNAL, - ) - self.trace_client.add_span(workflow_span) - - def suggested_question_trace(self, trace_info: SuggestedQuestionTraceInfo): - message_id = trace_info.message_id - status = create_status_from_error(trace_info.error) - - trace_metadata = TraceMetadata( - trace_id=convert_to_trace_id(message_id), - workflow_span_id=0, - session_id=trace_info.metadata.get("conversation_id") or "", - user_id=str(trace_info.metadata.get("user_id") or ""), - links=create_links_from_trace_id(trace_info.trace_id), - ) - - inputs_json = serialize_json_data(trace_info.inputs) - suggested_question_json = serialize_json_data(trace_info.suggested_question) - - suggested_question_span = SpanData( - trace_id=trace_metadata.trace_id, - parent_span_id=convert_to_span_id(message_id, "message"), - span_id=convert_to_span_id(message_id, "suggested_question"), - name="suggested_question", - start_time=convert_datetime_to_nanoseconds(trace_info.start_time), - end_time=convert_datetime_to_nanoseconds(trace_info.end_time), - attributes={ - **create_common_span_attributes( - session_id=trace_metadata.session_id, - user_id=trace_metadata.user_id, - span_kind=GenAISpanKind.LLM, - inputs=inputs_json, - outputs=suggested_question_json, - ), - GEN_AI_REQUEST_MODEL: trace_info.metadata.get("ls_model_name") or "", - GEN_AI_PROVIDER_NAME: trace_info.metadata.get("ls_provider") or "", - GEN_AI_PROMPT: inputs_json, - GEN_AI_COMPLETION: suggested_question_json, - }, - status=status, - links=trace_metadata.links, - ) - self.trace_client.add_span(suggested_question_span) diff --git a/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/config.py b/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/config.py index bae1c06d16d9a6..3f9df0ecdbf5df 100644 --- a/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/config.py +++ b/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/config.py @@ -1,36 +1,94 @@ -from pydantic import ConfigDict, ValidationInfo, field_validator +"""Aliyun configuration. -from core.ops.entities.config_entity import BaseTracingConfig -from core.ops.utils import validate_url_with_path +Users paste collector URLs that already carry the license key in their path +(``/adapt_/api/otlp/traces``). Such endpoints are rejected rather than +parsed, and validation messages are fixed texts that never include input. +""" +import re +from urllib.parse import unquote, urlsplit, urlunsplit -class AliyunConfig(BaseTracingConfig): - """ - Model class for Aliyun tracing config. - """ +from dify_trace_sdk.config import ProviderConfig, destination, name_or_default, secret +from pydantic import field_validator - # Endpoints may be pasted with an embedded LicenseKey. Even rejected input - # must not appear in formatted validation exceptions or chained API errors. - model_config = ConfigDict(hide_input_in_errors=True) +DEFAULT_APP_NAME = "dify_app" +DEFAULT_ENDPOINT = "https://tracing-analysis-dc-hz.aliyuncs.com" - app_name: str = "dify_app" - license_key: str - endpoint: str +_INVALID_APP_NAME = "Aliyun requires an application name without control characters." +_INVALID_LICENSE_KEY = "Aliyun requires a nonempty license key without whitespace or control characters." +_INVALID_ENDPOINT = "Aliyun requires an HTTP(S) base URL without credentials or a collector suffix." + + +def _is_control(char: str) -> bool: + return ord(char) < 32 or ord(char) == 127 + + +def _normalized_endpoint(value: str) -> str | None: + """The base URL with a lowercase ASCII host and no trailing slash; ``None`` when Aliyun cannot use it.""" + value = value.strip(" ") or DEFAULT_ENDPOINT + try: + parsed = urlsplit(value) + host, port = parsed.hostname, parsed.port + path = parsed.path.rstrip("/") + # Decode only the public path to detect hidden collector routes. License + # keys are opaque values and must never be unquoted or otherwise altered. + decoded_path = unquote(path, errors="strict") + valid = ( + parsed.scheme in {"http", "https"} + and host + and parsed.username is None + and parsed.password is None + and "?" not in value + and "#" not in value + and "\\" not in value + and not any(char == " " or _is_control(char) for char in value) + and not parsed.netloc.endswith(":") + and (port is None or 1 <= port <= 65535) + and "%" not in host + and not re.search(r"%(?![0-9a-fA-F]{2})", path) + and not re.search(r"%(?:2f|5c)", path, re.IGNORECASE) + and not any(char in decoded_path for char in "\\?#%") + and not any(_is_control(char) for char in decoded_path) + and "//" not in parsed.path + and not any(part in {".", ".."} or part.lower().startswith("adapt_") for part in decoded_path.split("/")) + and not decoded_path.lower().endswith(("/v1/traces", "/api/otlp/traces")) + ) + if not valid or host is None: + return None + host = host.encode("idna").decode("ascii").lower() + except (UnicodeError, ValueError): + # Return instead of raising here so no exception carrying the input is chained. + return None + authority = f"[{host}]" if ":" in host else host + if port is not None: + authority = f"{authority}:{port}" + return urlunsplit((parsed.scheme, authority, path, "", "")) + + +class AliyunConfig(ProviderConfig): + app_name: str = DEFAULT_APP_NAME + license_key: str = secret() + endpoint: str = destination() @field_validator("app_name") @classmethod - def app_name_validator(cls, v, info: ValidationInfo): - return cls.validate_project_field(v, "dify_app") + def _app_name(cls, value: str) -> str: + if any(_is_control(char) for char in value): + raise ValueError(_INVALID_APP_NAME) + return name_or_default(value, DEFAULT_APP_NAME) @field_validator("license_key") @classmethod - def license_key_validator(cls, v, info: ValidationInfo): - if not v or v.strip() == "": - raise ValueError("License key cannot be empty") - return v + def _license_key(cls, value: str) -> str: + # The key becomes a URL path segment: lone surrogates cannot be percent-encoded. + if not value or any(char.isspace() or _is_control(char) or "\ud800" <= char <= "\udfff" for char in value): + raise ValueError(_INVALID_LICENSE_KEY) + return value @field_validator("endpoint") @classmethod - def endpoint_validator(cls, v, info: ValidationInfo): - # aliyun uses two URL formats, which may include a URL path - return validate_url_with_path(v, "https://tracing-analysis-dc-hz.aliyuncs.com") + def _endpoint(cls, value: str) -> str: + endpoint = _normalized_endpoint(value) + if endpoint is None: + raise ValueError(_INVALID_ENDPOINT) + return endpoint diff --git a/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/configuration.py b/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/configuration.py deleted file mode 100644 index ca7c9bf689c78a..00000000000000 --- a/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/configuration.py +++ /dev/null @@ -1,158 +0,0 @@ -"""Aliyun destinations and a resource-only, zero-span collector access check. - -GET and URL helpers are pure. Verification makes one explicit collector request; -it never constructs a tracer, creates span IDs, or exports business telemetry. -""" - -import re -from collections.abc import Mapping -from typing import Literal -from urllib.parse import quote, unquote, urlsplit, urlunsplit - -from dify_trace_aliyun.config import AliyunConfig - -type ConfigurationCode = Literal["invalid_endpoint", "invalid_app_name", "invalid_license_key", "verification_failed"] -_MESSAGES: dict[ConfigurationCode, str] = { - "invalid_endpoint": "Aliyun requires an HTTP(S) base URL without credentials or a collector suffix.", - "invalid_app_name": "Aliyun requires an application name without control characters.", - "invalid_license_key": "Aliyun requires a nonempty license key without whitespace or control characters.", - "verification_failed": "Aliyun tracing configuration could not be verified.", -} -_DEFAULT_ENDPOINT = "https://tracing-analysis-dc-hz.aliyuncs.com" - - -class AliyunConfigurationError(ValueError): - """Diagnostics never include submitted URLs, credentials, or response bodies.""" - - def __init__(self, code: ConfigurationCode) -> None: - self.code = code - super().__init__(_MESSAGES[code]) - - -def _endpoint(value: object) -> str: - if not isinstance(value, str): - raise AliyunConfigurationError("invalid_endpoint") - value = value.strip(" ") or _DEFAULT_ENDPOINT - try: - parsed = urlsplit(value) - host, port = parsed.hostname, parsed.port - path = parsed.path.rstrip("/") - # Decode only the public path to detect hidden collector routes. License - # keys are opaque values and must never be unquoted or otherwise altered. - decoded_path = unquote(path, errors="strict") - valid = ( - parsed.scheme in {"http", "https"} - and host - and parsed.username is None - and parsed.password is None - and "?" not in value - and "#" not in value - and "\\" not in value - and not any(ord(char) <= 32 or ord(char) == 127 for char in value) - and not parsed.netloc.endswith(":") - and (port is None or 1 <= port <= 65535) - and "%" not in host - and not re.search(r"%(?![0-9a-fA-F]{2})", path) - and not re.search(r"%(?:2f|5c)", path, re.IGNORECASE) - and not any(char in decoded_path for char in "\\?#%") - and not any(ord(char) < 32 or ord(char) == 127 for char in decoded_path) - and "//" not in parsed.path - and not any(part in {".", ".."} or part.lower().startswith("adapt_") for part in decoded_path.split("/")) - and not decoded_path.lower().endswith(("/v1/traces", "/api/otlp/traces")) - ) - if not valid or host is None: - raise AliyunConfigurationError("invalid_endpoint") - host = host.encode("idna").decode("ascii").lower() - authority = f"[{host}]" if ":" in host else host - if port is not None: - authority = f"{authority}:{port}" - return urlunsplit((parsed.scheme, authority, path, "", "")) - except (UnicodeError, ValueError): - raise AliyunConfigurationError("invalid_endpoint") from None - - -def _license_key(value: object) -> str: - if ( - not isinstance(value, str) - or not value - or any(char.isspace() or ord(char) < 32 or ord(char) == 127 for char in value) - ): - raise AliyunConfigurationError("invalid_license_key") - try: - value.encode("utf-8") - except UnicodeError: - raise AliyunConfigurationError("invalid_license_key") from None - return value - - -def normalize_configuration(values: Mapping[str, object]) -> AliyunConfig: - """Validate new writes without tightening the model used for historical GET. - - Masked keys survive this phase so the gateway can restore the original - ciphertext on PATCH. Exporters reject masks after that restoration. - """ - name = values.get("app_name", "dify_app") - if not isinstance(name, str) or any(ord(char) < 32 or ord(char) == 127 for char in name): - raise AliyunConfigurationError("invalid_app_name") - return AliyunConfig( - app_name=name.strip() or "dify_app", - license_key=_license_key(values.get("license_key")), - endpoint=_endpoint(values.get("endpoint")), - ) - - -def ingestion_url(config: Mapping[str, object]) -> str: - """Build one destination, encoding the raw license exactly once as a path segment.""" - normalized = normalize_configuration(config) - if "*" in normalized.license_key: - raise AliyunConfigurationError("invalid_license_key") - hostname = (urlsplit(normalized.endpoint).hostname or "").removesuffix(".") - is_sls = hostname == "log.aliyuncs.com" or hostname.endswith(".log.aliyuncs.com") - suffix = "api/v1/traces" if is_sls else "api/otlp/traces" - # Aliyun's documented path-form token contains literal '@'. RFC 3986 path - # segments permit it, and this collector does not treat '%40' equivalently. - # Preserve only that required reserved character; never unquote raw keys. - return f"{normalized.endpoint}/adapt_{quote(normalized.license_key, safe='@')}/{suffix}" - - -def console_url() -> str: - """Return the platform console, not a destination-specific project or trace link.""" - return "https://arms.console.aliyun.com/#/llm" - - -def verify_configuration(config: AliyunConfig) -> None: - """Check access with one resource envelope and strictly zero scopes/spans. - - This collector distinguishes valid from invalid licenses for this request; - HEAD 405 and a completely empty request are not verification. A successful - acknowledgment does not establish business-span retention or queryability. - """ - from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest - - from services.tracing.otlp_export import is_aliyun_http_success_ack - from services.tracing.provider_common import request - - normalized = normalize_configuration(config.model_dump()) - url = ingestion_url(normalized.model_dump()) - payload = ExportTraceServiceRequest() - resource = payload.resource_spans.add().resource - resource.attributes.add(key="service.name").value.string_value = normalized.app_name - try: - response = request( - "POST", - url, - "aliyun", - headers={"Content-Type": "application/x-protobuf", "Accept": "application/x-protobuf"}, - content=payload.SerializeToString(), - ) - try: - # Only the account-tested resource-only ACK establishes this check. - # A generic empty OTLP ACK remains valid for ordinary export, but - # cannot be substituted as evidence that this license was checked. - if not is_aliyun_http_success_ack(response.status_code, response.content): - raise AliyunConfigurationError("verification_failed") - finally: - response.close() - except Exception: - # No URL, LicenseKey, vendor body, or nested transport exception escapes. - raise AliyunConfigurationError("verification_failed") from None diff --git a/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/data_exporter/traceclient.py b/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/data_exporter/traceclient.py deleted file mode 100644 index adbb9aefee6c0c..00000000000000 --- a/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/data_exporter/traceclient.py +++ /dev/null @@ -1,187 +0,0 @@ -import hashlib -import random -import socket -import uuid -from collections.abc import Generator, Sequence -from contextlib import contextmanager -from contextvars import ContextVar -from datetime import datetime -from typing import Final - -from opentelemetry import trace as trace_api -from opentelemetry.exporter.otlp.proto.common.trace_encoder import encode_spans -from opentelemetry.sdk.resources import Resource -from opentelemetry.sdk.trace import ReadableSpan -from opentelemetry.sdk.util.instrumentation import InstrumentationScope -from opentelemetry.semconv._incubating.attributes.deployment_attributes import ( # type: ignore[import-untyped] - DEPLOYMENT_ENVIRONMENT, -) -from opentelemetry.semconv._incubating.attributes.host_attributes import ( # type: ignore[import-untyped] - HOST_NAME, -) -from opentelemetry.semconv.attributes import service_attributes -from opentelemetry.trace import Link, SpanContext, TraceFlags - -from configs import dify_config -from dify_trace_aliyun.configuration import ingestion_url -from dify_trace_aliyun.entities.aliyun_trace_entity import SpanData -from dify_trace_aliyun.entities.semconv import ACS_ARMS_SERVICE_FEATURE -from services.tracing.otlp_export import export_http - -INVALID_SPAN_ID: Final[int] = 0x0000000000000000 -INVALID_TRACE_ID: Final[int] = 0x00000000000000000000000000000000 - - -class TraceClient: - """Collect one dispatch locally and synchronously acknowledge its exact batch. - - Instances are cached, so pending spans belong to the current context rather - than the client. A mapper failure discards only that dispatch; an uncertain - response propagates to the durable task instead of silently retrying a batch. - No SDK exporter, process-global configuration or background queue is created. - """ - - def __init__(self, service_name: str, endpoint: str) -> None: - self.endpoint = endpoint - self.resource = Resource( - attributes={ - service_attributes.SERVICE_NAME: service_name, - service_attributes.SERVICE_VERSION: f"dify-{dify_config.project.version}-{dify_config.COMMIT_SHA}", - DEPLOYMENT_ENVIRONMENT: f"{dify_config.DEPLOY_ENV}-{dify_config.DEPLOYMENT_EDITION.value}", - HOST_NAME: socket.gethostname(), - ACS_ARMS_SERVICE_FEATURE: "genai_app", - } - ) - self.span_builder = SpanBuilder(self.resource) - self._dispatch_spans: ContextVar[list[ReadableSpan] | None] = ContextVar("aliyun_dispatch_spans", default=None) - - @contextmanager - def batch(self) -> Generator[None]: - """Export only after the whole mapper succeeds, and always restore context.""" - spans: list[ReadableSpan] = [] - token = self._dispatch_spans.set(spans) - try: - yield - if spans: - self.export(spans) - finally: - self._dispatch_spans.reset(token) - - def export(self, spans: Sequence[ReadableSpan]) -> None: - if spans: - export_http("aliyun", self.endpoint, {}, encode_spans(spans)) - - def add_span(self, span_data: SpanData | None) -> None: - if span_data is None: - return - spans = self._dispatch_spans.get() - if spans is None: - raise RuntimeError("Aliyun spans require an active trace dispatch") - spans.append(self.span_builder.build_span(span_data)) - - -class SpanBuilder: - def __init__(self, resource: Resource) -> None: - self.resource = resource - self.instrumentation_scope = InstrumentationScope( - __name__, - "", - None, - None, - ) - - def build_span(self, span_data: SpanData) -> ReadableSpan: - span_context = trace_api.SpanContext( - trace_id=span_data.trace_id, - span_id=span_data.span_id, - is_remote=False, - trace_flags=trace_api.TraceFlags(trace_api.TraceFlags.SAMPLED), - trace_state=None, - ) - - parent_span_context = None - if span_data.parent_span_id is not None: - parent_span_context = trace_api.SpanContext( - trace_id=span_data.trace_id, - span_id=span_data.parent_span_id, - is_remote=False, - trace_flags=trace_api.TraceFlags(trace_api.TraceFlags.SAMPLED), - trace_state=None, - ) - - span = ReadableSpan( - name=span_data.name, - context=span_context, - parent=parent_span_context, - resource=self.resource, - attributes=span_data.attributes, - events=span_data.events, - links=span_data.links, - kind=span_data.span_kind, - status=span_data.status, - start_time=span_data.start_time, - end_time=span_data.end_time, - instrumentation_scope=self.instrumentation_scope, - ) - return span - - -def create_link(trace_id_str: str) -> Link: - placeholder_span_id = INVALID_SPAN_ID - try: - trace_id = int(trace_id_str, 16) - except ValueError as e: - raise ValueError(f"Invalid trace ID format: {trace_id_str}") from e - - span_context = SpanContext( - trace_id=trace_id, span_id=placeholder_span_id, is_remote=False, trace_flags=TraceFlags(TraceFlags.SAMPLED) - ) - - return Link(span_context) - - -def generate_span_id() -> int: - span_id = random.getrandbits(64) - while span_id == INVALID_SPAN_ID: - span_id = random.getrandbits(64) - return span_id - - -def convert_to_trace_id(uuid_v4: str | None) -> int: - if uuid_v4 is None: - raise ValueError("UUID cannot be None") - try: - uuid_obj = uuid.UUID(uuid_v4) - return uuid_obj.int - except ValueError as e: - raise ValueError(f"Invalid UUID input: {uuid_v4}") from e - - -def convert_string_to_id(string: str | None) -> int: - if not string: - return generate_span_id() - hash_bytes = hashlib.sha256(string.encode("utf-8")).digest() - return int.from_bytes(hash_bytes[:8], byteorder="big", signed=False) - - -def convert_to_span_id(uuid_v4: str | None, span_type: str) -> int: - if uuid_v4 is None: - raise ValueError("UUID cannot be None") - try: - uuid_obj = uuid.UUID(uuid_v4) - except ValueError as e: - raise ValueError(f"Invalid UUID input: {uuid_v4}") from e - combined_key = f"{uuid_obj.hex}-{span_type}" - return convert_string_to_id(combined_key) - - -def convert_datetime_to_nanoseconds(start_time_a: datetime | None) -> int | None: - if start_time_a is None: - return None - timestamp_in_seconds = start_time_a.timestamp() - return int(timestamp_in_seconds * 1e9) - - -def build_endpoint(base_url: str, license_key: str) -> str: - """Compatibility wrapper around the shared, credential-safe URL builder.""" - return ingestion_url({"endpoint": base_url, "license_key": license_key}) diff --git a/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/entities/__init__.py b/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/entities/__init__.py deleted file mode 100644 index e69de29bb2d1d6..00000000000000 diff --git a/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/entities/aliyun_trace_entity.py b/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/entities/aliyun_trace_entity.py deleted file mode 100644 index 9078031490d4d6..00000000000000 --- a/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/entities/aliyun_trace_entity.py +++ /dev/null @@ -1,37 +0,0 @@ -from collections.abc import Sequence -from dataclasses import dataclass -from typing import Any - -from opentelemetry import trace as trace_api -from opentelemetry.sdk.trace import Event -from opentelemetry.trace import SpanKind, Status, StatusCode -from pydantic import BaseModel, Field - - -@dataclass -class TraceMetadata: - """Metadata for trace operations, containing common attributes for all spans in a trace.""" - - trace_id: int - workflow_span_id: int - session_id: str - user_id: str - links: list[trace_api.Link] - - -class SpanData(BaseModel): - """Data model for span information in Aliyun trace system.""" - - model_config = {"arbitrary_types_allowed": True} - - trace_id: int = Field(..., description="The unique identifier for the trace.") - parent_span_id: int | None = Field(None, description="The ID of the parent span, if any.") - span_id: int = Field(..., description="The unique identifier for this span.") - name: str = Field(..., description="The name of the span.") - attributes: dict[str, Any] = Field(default_factory=dict, description="Attributes associated with the span.") - events: Sequence[Event] = Field(default_factory=list, description="Events recorded in the span.") - links: Sequence[trace_api.Link] = Field(default_factory=list, description="Links to other spans.") - status: Status = Field(default=Status(StatusCode.UNSET), description="The status of the span.") - start_time: int | None = Field(..., description="The start time of the span in nanoseconds.") - end_time: int | None = Field(..., description="The end time of the span in nanoseconds.") - span_kind: SpanKind = Field(default=SpanKind.INTERNAL, description="The OpenTelemetry SpanKind for this span.") diff --git a/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/entities/semconv.py b/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/entities/semconv.py deleted file mode 100644 index e9ad2504728ff2..00000000000000 --- a/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/entities/semconv.py +++ /dev/null @@ -1,84 +0,0 @@ -from enum import StrEnum -from typing import Final - -ACS_ARMS_SERVICE_FEATURE: Final[str] = "acs.arms.service.feature" - -# Dify-specific attributes -DIFY_APP_ID: Final[str] = "dify.app_id" -DIFY_NODE_TYPE: Final[str] = "dify.node.type" - -# Public attributes -GEN_AI_SESSION_ID: Final[str] = "gen_ai.session.id" -GEN_AI_USER_ID: Final[str] = "gen_ai.user.id" -GEN_AI_USER_NAME: Final[str] = "gen_ai.user.name" -GEN_AI_SPAN_KIND: Final[str] = "gen_ai.span.kind" -GEN_AI_OPERATION_NAME: Final[str] = "gen_ai.operation.name" -GEN_AI_FRAMEWORK: Final[str] = "gen_ai.framework" - -# Chain attributes -INPUT_VALUE: Final[str] = "input.value" -OUTPUT_VALUE: Final[str] = "output.value" - -# Retriever attributes -RETRIEVAL_QUERY: Final[str] = "retrieval.query" -RETRIEVAL_DOCUMENT: Final[str] = "retrieval.document" - -# LLM attributes -GEN_AI_REQUEST_MODEL: Final[str] = "gen_ai.request.model" -GEN_AI_PROVIDER_NAME: Final[str] = "gen_ai.provider.name" -GEN_AI_USAGE_INPUT_TOKENS: Final[str] = "gen_ai.usage.input_tokens" -GEN_AI_USAGE_OUTPUT_TOKENS: Final[str] = "gen_ai.usage.output_tokens" -GEN_AI_USAGE_TOTAL_TOKENS: Final[str] = "gen_ai.usage.total_tokens" -GEN_AI_PROMPT: Final[str] = "gen_ai.prompt" -GEN_AI_COMPLETION: Final[str] = "gen_ai.completion" -GEN_AI_RESPONSE_FINISH_REASON: Final[str] = "gen_ai.response.finish_reason" -# Time to first token of the model response in a streaming scenario, in nanoseconds. -GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN: Final[str] = "gen_ai.response.time_to_first_token" - -GEN_AI_INPUT_MESSAGE: Final[str] = "gen_ai.input.messages" -GEN_AI_OUTPUT_MESSAGE: Final[str] = "gen_ai.output.messages" - -# Tool attributes (GenAI semantic conventions) -GEN_AI_TOOL_CALL_ID: Final[str] = "gen_ai.tool.call.id" -GEN_AI_TOOL_DESCRIPTION: Final[str] = "gen_ai.tool.description" -GEN_AI_TOOL_NAME: Final[str] = "gen_ai.tool.name" -GEN_AI_TOOL_TYPE: Final[str] = "gen_ai.tool.type" -GEN_AI_TOOL_CALL_ARGUMENTS: Final[str] = "gen_ai.tool.call.arguments" -GEN_AI_TOOL_CALL_RESULT: Final[str] = "gen_ai.tool.call.result" - -# Skill attributes (conditionally required when loading a Skill) -GEN_AI_SKILL_ID: Final[str] = "gen_ai.skill.id" -GEN_AI_SKILL_NAME: Final[str] = "gen_ai.skill.name" -GEN_AI_SKILL_DESCRIPTION: Final[str] = "gen_ai.skill.description" -GEN_AI_SKILL_VERSION: Final[str] = "gen_ai.skill.version" - -# Agent attributes -GEN_AI_AGENT_NAME: Final[str] = "gen_ai.agent.name" - -# ReAct step attributes -GEN_AI_REACT_ROUND: Final[str] = "gen_ai.react.round" -GEN_AI_REACT_FINISH_REASON: Final[str] = "gen_ai.react.finish_reason" - -# gen_ai.operation.name values (see Aliyun LLM Trace field definitions) -OPERATION_NAME_CHAT: Final[str] = "chat" -OPERATION_NAME_EXECUTE_TOOL: Final[str] = "execute_tool" -OPERATION_NAME_INVOKE_AGENT: Final[str] = "invoke_agent" -OPERATION_NAME_REACT: Final[str] = "react" - -# gen_ai.tool.type values -TOOL_TYPE_FUNCTION: Final[str] = "function" -TOOL_TYPE_EXTENSION: Final[str] = "extension" -TOOL_TYPE_DATASTORE: Final[str] = "datastore" - - -class GenAISpanKind(StrEnum): - CHAIN = "CHAIN" - RETRIEVER = "RETRIEVER" - RERANKER = "RERANKER" - LLM = "LLM" - EMBEDDING = "EMBEDDING" - TOOL = "TOOL" - AGENT = "AGENT" - TASK = "TASK" - # Marks one Reasoning-Acting iteration of an agent (ReAct step). - STEP = "STEP" diff --git a/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/provider.py b/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/provider.py new file mode 100644 index 00000000000000..354a6a43ae0194 --- /dev/null +++ b/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/provider.py @@ -0,0 +1,104 @@ +"""Aliyun LLM application monitoring over OTLP/HTTP. + +The license key travels in the collector URL path, so every request marks its +URL as sensitive. The collector acknowledges an export with the literal body +``"success"`` rather than an OTLP response; standard acknowledgments are +accepted too. +""" + +from typing import override +from urllib.parse import quote, urlsplit + +from dify_trace_sdk.egress import EgressResponse +from dify_trace_sdk.errors import EgressError, ExportRejectedError, VerificationError +from dify_trace_sdk.genai import GenAIDialect, genai_attributes +from dify_trace_sdk.model import SpanKind, Trace +from dify_trace_sdk.otlp import encode_trace, post_otlp, read_otlp_ack +from dify_trace_sdk.provider import ProviderCapabilities, SendResult, TraceProvider +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest + +from dify_trace_aliyun.config import AliyunConfig + +VERIFY_TIMEOUT_SECONDS = 10.0 +_DESTINATION = "Aliyun" +_CONSOLE_URL = "https://arms.console.aliyun.com/#/llm" +_DIALECT = GenAIDialect( + kind_labels={kind: kind.value.upper() for kind in SpanKind}, + input_key="input.value", + output_key="output.value", + mark_entry_span=False, +) + + +def _ingestion_url(config: AliyunConfig) -> str: + """The collector route, with the raw license key encoded exactly once as a path segment.""" + hostname = (urlsplit(config.endpoint).hostname or "").removesuffix(".") + is_sls = hostname == "log.aliyuncs.com" or hostname.endswith(".log.aliyuncs.com") + suffix = "api/v1/traces" if is_sls else "api/otlp/traces" + # Aliyun's documented path-form key contains a literal '@'. RFC 3986 path + # segments permit it, and the collector does not treat '%40' equivalently. + return f"{config.endpoint}/adapt_{quote(config.license_key, safe='@')}/{suffix}" + + +def _is_masked(config: AliyunConfig) -> bool: + """A displayed mask of a key passes validation but is not a credential to send.""" + return "*" in config.license_key + + +def _is_success_ack(response: EgressResponse) -> bool: + """The observed acknowledgment of Aliyun's HTTP collector; never normalize other bodies into it.""" + return response.status_code == 200 and response.content == b'"success"' + + +class AliyunProvider(TraceProvider[AliyunConfig]): + id = "aliyun" + config_model = AliyunConfig + capabilities = ProviderCapabilities(native_cost=False, trace_url=False, span_url=False) + + @override + def verify(self, config: AliyunConfig) -> AliyunConfig: + """Post one resource with zero spans and require the collector's literal acknowledgment. + + The collector distinguishes valid from invalid license keys for this + request, whereas an empty OTLP acknowledgment proves nothing. Success + does not establish that business spans are retained. + """ + if _is_masked(config): + raise VerificationError("invalid_license_key") + probe = ExportTraceServiceRequest() + resource = probe.resource_spans.add().resource + resource.attributes.add(key="service.name").value.string_value = config.app_name + try: + response = self.egress.request( + "POST", + _ingestion_url(config), + headers={"Content-Type": "application/x-protobuf", "Accept": "application/x-protobuf"}, + content=probe.SerializeToString(), + timeout=VERIFY_TIMEOUT_SECONDS, + sensitive_url=True, + ) + except EgressError: + raise VerificationError("verification_failed") from None + if not _is_success_ack(response): + raise VerificationError("verification_failed") + return config + + @override + def project_url(self, config: AliyunConfig) -> str | None: + return _CONSOLE_URL + + @override + def send(self, config: AliyunConfig, trace: Trace) -> SendResult: + if _is_masked(config): + raise ExportRejectedError("Aliyun tracing configuration is invalid") + request = encode_trace( + trace, + resource={"service.name": config.app_name, "acs.arms.service.feature": "genai_app"}, + attributes=lambda span: genai_attributes(span, trace, _DIALECT), + ) + response = post_otlp( + self.egress, _ingestion_url(config), request, destination=_DESTINATION, headers={}, sensitive_url=True + ) + if not _is_success_ack(response): + read_otlp_ack(response, _DESTINATION) + return SendResult(trace_id=trace.trace_id, trace_url=None, span_urls={}) diff --git a/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/utils.py b/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/utils.py deleted file mode 100644 index c9145a95b546b8..00000000000000 --- a/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/utils.py +++ /dev/null @@ -1,369 +0,0 @@ -import json -import re -from collections.abc import Mapping -from dataclasses import dataclass, field -from typing import Any, TypedDict - -from opentelemetry.trace import Link, Status, StatusCode - -from core.rag.models.document import Document -from dify_trace_aliyun.entities.semconv import ( - GEN_AI_FRAMEWORK, - GEN_AI_OPERATION_NAME, - GEN_AI_SESSION_ID, - GEN_AI_SPAN_KIND, - GEN_AI_TOOL_CALL_ARGUMENTS, - GEN_AI_TOOL_CALL_ID, - GEN_AI_TOOL_CALL_RESULT, - GEN_AI_TOOL_DESCRIPTION, - GEN_AI_TOOL_NAME, - GEN_AI_TOOL_TYPE, - GEN_AI_USER_ID, - INPUT_VALUE, - OPERATION_NAME_EXECUTE_TOOL, - OUTPUT_VALUE, - TOOL_TYPE_DATASTORE, - TOOL_TYPE_EXTENSION, - TOOL_TYPE_FUNCTION, - GenAISpanKind, -) -from extensions.ext_database import db -from graphon.entities import WorkflowNodeExecution -from graphon.enums import WorkflowNodeExecutionStatus -from models import EndUser - -# Constants -DEFAULT_JSON_ENSURE_ASCII = False -DEFAULT_FRAMEWORK_NAME = "dify" - - -def get_user_id_from_message_data(message_data) -> str: - user_id = message_data.from_account_id - if message_data.from_end_user_id: - end_user_data: EndUser | None = db.session.get(EndUser, message_data.from_end_user_id) - if end_user_data is not None: - user_id = end_user_data.session_id - return user_id - - -def create_status_from_error(error: str | None) -> Status: - if error: - return Status(StatusCode.ERROR, error) - return Status(StatusCode.OK) - - -def get_workflow_node_status(node_execution: WorkflowNodeExecution) -> Status: - if node_execution.status == WorkflowNodeExecutionStatus.SUCCEEDED: - return Status(StatusCode.OK) - if node_execution.status in [WorkflowNodeExecutionStatus.FAILED, WorkflowNodeExecutionStatus.EXCEPTION]: - return Status(StatusCode.ERROR, str(node_execution.error)) - return Status(StatusCode.UNSET) - - -def create_links_from_trace_id(trace_id: str | None) -> list[Link]: - from dify_trace_aliyun.data_exporter.traceclient import create_link - - links = [] - if trace_id: - links.append(create_link(trace_id_str=trace_id)) - return links - - -class RetrievalDocumentMetadataDict(TypedDict): - dataset_id: Any - doc_id: Any - document_id: Any - - -class RetrievalDocumentDict(TypedDict): - content: str - metadata: RetrievalDocumentMetadataDict - score: Any - - -def extract_retrieval_documents(documents: list[Document]) -> list[RetrievalDocumentDict]: - documents_data: list[RetrievalDocumentDict] = [] - for document in documents: - document_data: RetrievalDocumentDict = { - "content": document.page_content, - "metadata": { - "dataset_id": document.metadata.get("dataset_id"), - "doc_id": document.metadata.get("doc_id"), - "document_id": document.metadata.get("document_id"), - }, - "score": document.metadata.get("score"), - } - documents_data.append(document_data) - return documents_data - - -def serialize_json_data(data: Any, ensure_ascii: bool = DEFAULT_JSON_ENSURE_ASCII) -> str: - return json.dumps(data, ensure_ascii=ensure_ascii) - - -def create_common_span_attributes( - session_id: str = "", - user_id: str = "", - span_kind: str = GenAISpanKind.CHAIN, - framework: str = DEFAULT_FRAMEWORK_NAME, - inputs: str = "", - outputs: str = "", -) -> dict[str, str]: - return { - GEN_AI_SESSION_ID: session_id, - GEN_AI_USER_ID: user_id, - GEN_AI_SPAN_KIND: span_kind, - GEN_AI_FRAMEWORK: framework, - INPUT_VALUE: inputs, - OUTPUT_VALUE: outputs, - } - - -def map_gen_ai_tool_type(provider_type: str | None) -> str: - """Map Dify tool provider type to GenAI ``gen_ai.tool.type`` values.""" - normalized = (provider_type or "").strip().lower() - if normalized in {"dataset-retrieval", "datastore"}: - return TOOL_TYPE_DATASTORE - if normalized == "extension": - return TOOL_TYPE_EXTENSION - return TOOL_TYPE_FUNCTION - - -def extract_tool_description(tool_meta: Mapping[str, Any] | None) -> str: - if not isinstance(tool_meta, Mapping): - return "" - for key in ("description", "tool_description"): - value = tool_meta.get(key) - if value: - return str(value) - return "" - - -def create_gen_ai_tool_attributes( - *, - tool_name: str, - tool_type: str = TOOL_TYPE_FUNCTION, - tool_description: str = "", - tool_call_id: str = "", - tool_call_arguments: str = "", - tool_call_result: str = "", -) -> dict[str, str]: - """Build GenAI-compliant attributes for a TOOL span.""" - return { - GEN_AI_OPERATION_NAME: OPERATION_NAME_EXECUTE_TOOL, - GEN_AI_TOOL_NAME: tool_name, - GEN_AI_TOOL_TYPE: tool_type or TOOL_TYPE_FUNCTION, - GEN_AI_TOOL_DESCRIPTION: tool_description, - GEN_AI_TOOL_CALL_ID: tool_call_id, - GEN_AI_TOOL_CALL_ARGUMENTS: tool_call_arguments, - GEN_AI_TOOL_CALL_RESULT: tool_call_result, - } - - -def format_retrieval_documents(retrieval_documents: list) -> list: - try: - if not isinstance(retrieval_documents, list): - return [] - - semantic_documents = [] - for doc in retrieval_documents: - if not isinstance(doc, dict): - continue - - metadata = doc.get("metadata", {}) - content = doc.get("content", "") - title = doc.get("title", "") - score = metadata.get("score", 0.0) - document_id = metadata.get("document_id", "") - - semantic_metadata = {} - if title: - semantic_metadata["title"] = title - if metadata.get("source"): - semantic_metadata["source"] = metadata["source"] - elif metadata.get("_source"): - semantic_metadata["source"] = metadata["_source"] - if metadata.get("doc_metadata"): - doc_metadata = metadata["doc_metadata"] - if isinstance(doc_metadata, dict): - semantic_metadata.update(doc_metadata) - - semantic_doc = { - "document": {"content": content, "metadata": semantic_metadata, "score": score, "id": document_id} - } - semantic_documents.append(semantic_doc) - - return semantic_documents - except Exception: - return [] - - -def format_input_messages(process_data: Mapping[str, Any]) -> str: - try: - if not isinstance(process_data, dict): - return serialize_json_data([]) - - prompts = process_data.get("prompts", []) - if not prompts: - return serialize_json_data([]) - - valid_roles = {"system", "user", "assistant", "tool"} - input_messages = [] - for prompt in prompts: - if not isinstance(prompt, dict): - continue - - role = prompt.get("role", "") - text = prompt.get("text", "") - - if not role or role not in valid_roles: - continue - - if text: - message = {"role": role, "parts": [{"type": "text", "content": text}]} - input_messages.append(message) - - return serialize_json_data(input_messages) - except Exception: - return serialize_json_data([]) - - -def convert_seconds_to_nanoseconds(seconds: float) -> int: - return int(seconds * 1e9) - - -_REACT_ROUND_LABEL_PATTERN = re.compile(r"ROUND\s+(\d+)", re.IGNORECASE) -_LLM_THOUGHT_LABEL_SUFFIX = " Thought" -_TOOL_CALL_LABEL_PREFIX = "CALL " - - -@dataclass -class AgentLogEntry: - """One entry of the agent-strategy execution log (``outputs["json"]`` of an agent node). - - Entries form a tree via ``parent_id``: top-level entries are ReAct rounds - (label like ``ROUND 1``) and their children are LLM thoughts / tool calls. - ``started_at``/``finished_at`` in ``metadata`` are monotonic-clock seconds - (``time.perf_counter``), not epoch timestamps. - """ - - id: str - parent_id: str | None - label: str - status: str - error: str | None - data: dict[str, Any] - metadata: dict[str, Any] - children: list["AgentLogEntry"] = field(default_factory=list) - - -def parse_agent_log_entries(outputs: Mapping[str, Any]) -> list[AgentLogEntry]: - """Parse agent node outputs into a tree of log entries, returning top-level rounds in order. - - Entries without an ``id`` (e.g. the trailing ``{"data": []}`` element) are skipped. - Children whose parent is missing are dropped. - """ - raw_entries = outputs.get("json") - if not isinstance(raw_entries, list): - return [] - - entries: list[AgentLogEntry] = [] - entries_by_id: dict[str, AgentLogEntry] = {} - for raw_entry in raw_entries: - if not isinstance(raw_entry, dict): - continue - entry_id = raw_entry.get("id") - if not entry_id: - continue - data = raw_entry.get("data") - metadata = raw_entry.get("metadata") - entry = AgentLogEntry( - id=str(entry_id), - parent_id=raw_entry.get("parent_id"), - label=str(raw_entry.get("label") or ""), - status=str(raw_entry.get("status") or ""), - error=raw_entry.get("error"), - data=data if isinstance(data, dict) else {}, - metadata=metadata if isinstance(metadata, dict) else {}, - ) - entries.append(entry) - entries_by_id[entry.id] = entry - - roots: list[AgentLogEntry] = [] - for entry in entries: - if entry.parent_id is None: - roots.append(entry) - else: - parent = entries_by_id.get(entry.parent_id) - if parent is not None: - parent.children.append(entry) - return roots - - -def extract_react_round_number(label: str, fallback: int) -> int: - match = _REACT_ROUND_LABEL_PATTERN.search(label) - if match: - return int(match.group(1)) - return fallback - - -def is_tool_call_entry(entry: AgentLogEntry) -> bool: - """Tool invocations use labels like ``CALL {tool_name}`` (also carry a provider in metadata).""" - return entry.label.startswith(_TOOL_CALL_LABEL_PREFIX) - - -def is_llm_thought_entry(entry: AgentLogEntry) -> bool: - """LLM thought entries use labels like ``{model} Thought``. - - Do not key off ``metadata.provider`` alone: tool CALL entries also set provider - (to the tool provider), which previously misclassified them as LLM spans. - """ - if is_tool_call_entry(entry): - return False - return entry.label.endswith(_LLM_THOUGHT_LABEL_SUFFIX) or bool(entry.metadata.get("provider")) - - -def extract_model_name_from_thought_label(label: str) -> str: - if label.endswith(_LLM_THOUGHT_LABEL_SUFFIX): - return label.removesuffix(_LLM_THOUGHT_LABEL_SUFFIX) - return "" - - -def extract_tool_name_from_call_label(label: str) -> str: - if label.startswith(_TOOL_CALL_LABEL_PREFIX): - return label.removeprefix(_TOOL_CALL_LABEL_PREFIX).strip() - return "" - - -def create_status_from_agent_log_entry(entry: AgentLogEntry) -> Status: - if entry.error: - return Status(StatusCode.ERROR, str(entry.error)) - if entry.status == "success": - return Status(StatusCode.OK) - return Status(StatusCode.UNSET) - - -def format_output_messages(outputs: Mapping[str, Any]) -> str: - try: - if not isinstance(outputs, dict): - return serialize_json_data([]) - - text = outputs.get("text", "") - finish_reason = outputs.get("finish_reason", "") - - if not text: - return serialize_json_data([]) - - valid_finish_reasons = {"stop", "length", "content_filter", "tool_call", "error"} - if finish_reason not in valid_finish_reasons: - finish_reason = "stop" - - output_message = { - "role": "assistant", - "parts": [{"type": "text", "content": text}], - "finish_reason": finish_reason, - } - - return serialize_json_data([output_message]) - except Exception: - return serialize_json_data([]) diff --git a/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/data_exporter/test_traceclient.py b/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/data_exporter/test_traceclient.py deleted file mode 100644 index 2d520a644bd470..00000000000000 --- a/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/data_exporter/test_traceclient.py +++ /dev/null @@ -1,160 +0,0 @@ -import uuid -from datetime import datetime -from unittest.mock import patch - -import pytest -from dify_trace_aliyun.data_exporter.traceclient import ( - INVALID_SPAN_ID, - SpanBuilder, - TraceClient, - build_endpoint, - convert_datetime_to_nanoseconds, - convert_string_to_id, - convert_to_span_id, - convert_to_trace_id, - create_link, - generate_span_id, -) -from dify_trace_aliyun.entities.aliyun_trace_entity import SpanData -from opentelemetry.sdk.resources import Resource -from opentelemetry.sdk.trace import ReadableSpan -from opentelemetry.trace import SpanKind, Status, StatusCode - - -class TestTraceClient: - def test_constructor_has_no_background_exporter(self) -> None: - client = TraceClient(service_name="test-service", endpoint="https://collector.test") - assert client.endpoint == "https://collector.test" - assert not hasattr(client, "exporter") - assert not hasattr(client, "worker_thread") - assert not hasattr(client, "queue") - - def test_add_requires_dispatch(self) -> None: - client = TraceClient(service_name="test-service", endpoint="https://collector.test") - client.add_span(None) - with pytest.raises(RuntimeError, match="active trace dispatch"): - client.add_span(SpanData(name="span", trace_id=123, span_id=456, start_time=1000, end_time=2000)) - - -class TestSpanBuilder: - def test_build_span(self): - resource = Resource({"service.name": "test"}) - builder = SpanBuilder(resource) - - span_data = SpanData( - name="test-span", - trace_id=123, - span_id=456, - parent_span_id=789, - start_time=1000, - end_time=2000, - status=Status(StatusCode.OK), - span_kind=SpanKind.INTERNAL, - attributes={"attr1": "val1"}, - events=[], - links=[], - ) - - span = builder.build_span(span_data) - assert isinstance(span, ReadableSpan) - assert span.name == "test-span" - assert span.context is not None - assert span.context.trace_id == 123 - assert span.context.span_id == 456 - assert span.parent is not None - assert span.parent.span_id == 789 - assert span.resource == resource - assert span.attributes == {"attr1": "val1"} - - def test_build_span_no_parent(self): - resource = Resource({"service.name": "test"}) - builder = SpanBuilder(resource) - - span_data = SpanData( - name="test-span", - trace_id=123, - span_id=456, - parent_span_id=None, - start_time=1000, - end_time=2000, - status=Status(StatusCode.OK), - span_kind=SpanKind.INTERNAL, - ) - - span = builder.build_span(span_data) - assert span.parent is None - - -def test_create_link(): - trace_id_str = "0123456789abcdef0123456789abcdef" - link = create_link(trace_id_str) - assert link.context.trace_id == int(trace_id_str, 16) - assert link.context.span_id == INVALID_SPAN_ID - - with pytest.raises(ValueError, match="Invalid trace ID format"): - create_link("invalid-hex") - - -def test_generate_span_id(): - # Test normal generation - span_id = generate_span_id() - assert isinstance(span_id, int) - assert span_id != INVALID_SPAN_ID - - # Test retry loop - with patch("dify_trace_aliyun.data_exporter.traceclient.random.getrandbits") as mock_rand: - mock_rand.side_effect = [INVALID_SPAN_ID, 999] - span_id = generate_span_id() - assert span_id == 999 - assert mock_rand.call_count == 2 - - -def test_convert_to_trace_id(): - uid = str(uuid.uuid4()) - trace_id = convert_to_trace_id(uid) - assert trace_id == uuid.UUID(uid).int - - with pytest.raises(ValueError, match="UUID cannot be None"): - convert_to_trace_id(None) - - with pytest.raises(ValueError, match="Invalid UUID input"): - convert_to_trace_id("not-a-uuid") - - -def test_convert_string_to_id(): - assert convert_string_to_id("test") > 0 - # Test with None string - with patch("dify_trace_aliyun.data_exporter.traceclient.generate_span_id") as mock_gen: - mock_gen.return_value = 12345 - assert convert_string_to_id(None) == 12345 - - -def test_convert_to_span_id(): - uid = str(uuid.uuid4()) - span_id = convert_to_span_id(uid, "test-type") - assert isinstance(span_id, int) - - with pytest.raises(ValueError, match="UUID cannot be None"): - convert_to_span_id(None, "test") - - with pytest.raises(ValueError, match="Invalid UUID input"): - convert_to_span_id("not-a-uuid", "test") - - -def test_convert_datetime_to_nanoseconds(): - dt = datetime(2023, 1, 1, 12, 0, 0) - ns = convert_datetime_to_nanoseconds(dt) - assert ns == int(dt.timestamp() * 1e9) - assert convert_datetime_to_nanoseconds(None) is None - - -def test_build_endpoint(): - license_key = "abc" - - # CMS 2.0 endpoint - url1 = "https://log.aliyuncs.com" - assert build_endpoint(url1, license_key) == "https://log.aliyuncs.com/adapt_abc/api/v1/traces" - - # XTrace endpoint - url2 = "https://example.com" - assert build_endpoint(url2, license_key) == "https://example.com/adapt_abc/api/otlp/traces" diff --git a/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/entities/test_aliyun_trace_entity.py b/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/entities/test_aliyun_trace_entity.py deleted file mode 100644 index a6808fec0a75b7..00000000000000 --- a/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/entities/test_aliyun_trace_entity.py +++ /dev/null @@ -1,88 +0,0 @@ -import pytest -from dify_trace_aliyun.entities.aliyun_trace_entity import SpanData, TraceMetadata -from opentelemetry import trace as trace_api -from opentelemetry.sdk.trace import Event -from opentelemetry.trace import SpanKind, Status, StatusCode -from pydantic import ValidationError - - -class TestTraceMetadata: - def test_trace_metadata_init(self): - links = [trace_api.Link(context=trace_api.SpanContext(0, 0, False))] - metadata = TraceMetadata( - trace_id=123, workflow_span_id=456, session_id="session_1", user_id="user_1", links=links - ) - assert metadata.trace_id == 123 - assert metadata.workflow_span_id == 456 - assert metadata.session_id == "session_1" - assert metadata.user_id == "user_1" - assert metadata.links == links - - -class TestSpanData: - def test_span_data_init_required_fields(self): - span_data = SpanData(trace_id=123, span_id=456, name="test_span", start_time=1000, end_time=2000) - assert span_data.trace_id == 123 - assert span_data.span_id == 456 - assert span_data.name == "test_span" - assert span_data.start_time == 1000 - assert span_data.end_time == 2000 - - # Check defaults - assert span_data.parent_span_id is None - assert span_data.attributes == {} - assert span_data.events == [] - assert span_data.links == [] - assert span_data.status.status_code == StatusCode.UNSET - assert span_data.span_kind == SpanKind.INTERNAL - - def test_span_data_with_optional_fields(self): - event = Event(name="event_1", timestamp=1500) - link = trace_api.Link(context=trace_api.SpanContext(0, 0, False)) - status = Status(StatusCode.OK) - - span_data = SpanData( - trace_id=123, - parent_span_id=111, - span_id=456, - name="test_span", - attributes={"key": "value"}, - events=[event], - links=[link], - status=status, - start_time=1000, - end_time=2000, - span_kind=SpanKind.SERVER, - ) - - assert span_data.parent_span_id == 111 - assert span_data.attributes == {"key": "value"} - assert span_data.events == [event] - assert span_data.links == [link] - assert span_data.status.status_code == status.status_code - assert span_data.span_kind == SpanKind.SERVER - - def test_span_data_missing_required_fields(self): - with pytest.raises(ValidationError): - SpanData.model_validate( - { - "trace_id": 123, - "name": "test_span", - "start_time": 1000, - "end_time": 2000, - } - ) - - def test_span_data_arbitrary_types_allowed(self): - # opentelemetry.trace.Status and Event are "arbitrary types" for Pydantic - # This test ensures they are accepted thanks to model_config - status = Status(StatusCode.ERROR, description="error occurred") - event = Event(name="exception", timestamp=1234, attributes={"exception.type": "ValueError"}) - - span_data = SpanData( - trace_id=123, span_id=456, name="test_span", status=status, events=[event], start_time=1000, end_time=2000 - ) - - assert span_data.status.status_code == status.status_code - assert span_data.status.description == status.description - assert span_data.events == [event] diff --git a/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/entities/test_semconv.py b/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/entities/test_semconv.py deleted file mode 100644 index e07aa718f29bfa..00000000000000 --- a/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/entities/test_semconv.py +++ /dev/null @@ -1,101 +0,0 @@ -from dify_trace_aliyun.entities.semconv import ( - ACS_ARMS_SERVICE_FEATURE, - GEN_AI_AGENT_NAME, - GEN_AI_COMPLETION, - GEN_AI_FRAMEWORK, - GEN_AI_INPUT_MESSAGE, - GEN_AI_OPERATION_NAME, - GEN_AI_OUTPUT_MESSAGE, - GEN_AI_PROMPT, - GEN_AI_PROVIDER_NAME, - GEN_AI_REACT_FINISH_REASON, - GEN_AI_REACT_ROUND, - GEN_AI_REQUEST_MODEL, - GEN_AI_RESPONSE_FINISH_REASON, - GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN, - GEN_AI_SESSION_ID, - GEN_AI_SKILL_DESCRIPTION, - GEN_AI_SKILL_ID, - GEN_AI_SKILL_NAME, - GEN_AI_SKILL_VERSION, - GEN_AI_SPAN_KIND, - GEN_AI_TOOL_CALL_ARGUMENTS, - GEN_AI_TOOL_CALL_ID, - GEN_AI_TOOL_CALL_RESULT, - GEN_AI_TOOL_DESCRIPTION, - GEN_AI_TOOL_NAME, - GEN_AI_TOOL_TYPE, - GEN_AI_USAGE_INPUT_TOKENS, - GEN_AI_USAGE_OUTPUT_TOKENS, - GEN_AI_USAGE_TOTAL_TOKENS, - GEN_AI_USER_ID, - GEN_AI_USER_NAME, - INPUT_VALUE, - OPERATION_NAME_EXECUTE_TOOL, - OUTPUT_VALUE, - RETRIEVAL_DOCUMENT, - RETRIEVAL_QUERY, - TOOL_TYPE_DATASTORE, - TOOL_TYPE_EXTENSION, - TOOL_TYPE_FUNCTION, - GenAISpanKind, -) - - -def test_constants(): - assert ACS_ARMS_SERVICE_FEATURE == "acs.arms.service.feature" - assert GEN_AI_SESSION_ID == "gen_ai.session.id" - assert GEN_AI_USER_ID == "gen_ai.user.id" - assert GEN_AI_USER_NAME == "gen_ai.user.name" - assert GEN_AI_SPAN_KIND == "gen_ai.span.kind" - assert GEN_AI_FRAMEWORK == "gen_ai.framework" - assert INPUT_VALUE == "input.value" - assert OUTPUT_VALUE == "output.value" - assert RETRIEVAL_QUERY == "retrieval.query" - assert RETRIEVAL_DOCUMENT == "retrieval.document" - assert GEN_AI_REQUEST_MODEL == "gen_ai.request.model" - assert GEN_AI_PROVIDER_NAME == "gen_ai.provider.name" - assert GEN_AI_USAGE_INPUT_TOKENS == "gen_ai.usage.input_tokens" - assert GEN_AI_USAGE_OUTPUT_TOKENS == "gen_ai.usage.output_tokens" - assert GEN_AI_USAGE_TOTAL_TOKENS == "gen_ai.usage.total_tokens" - assert GEN_AI_PROMPT == "gen_ai.prompt" - assert GEN_AI_COMPLETION == "gen_ai.completion" - assert GEN_AI_RESPONSE_FINISH_REASON == "gen_ai.response.finish_reason" - assert GEN_AI_INPUT_MESSAGE == "gen_ai.input.messages" - assert GEN_AI_OUTPUT_MESSAGE == "gen_ai.output.messages" - assert GEN_AI_TOOL_CALL_ID == "gen_ai.tool.call.id" - assert GEN_AI_TOOL_DESCRIPTION == "gen_ai.tool.description" - assert GEN_AI_TOOL_NAME == "gen_ai.tool.name" - assert GEN_AI_TOOL_TYPE == "gen_ai.tool.type" - assert GEN_AI_TOOL_CALL_ARGUMENTS == "gen_ai.tool.call.arguments" - assert GEN_AI_TOOL_CALL_RESULT == "gen_ai.tool.call.result" - assert GEN_AI_SKILL_ID == "gen_ai.skill.id" - assert GEN_AI_SKILL_NAME == "gen_ai.skill.name" - assert GEN_AI_SKILL_DESCRIPTION == "gen_ai.skill.description" - assert GEN_AI_SKILL_VERSION == "gen_ai.skill.version" - assert GEN_AI_OPERATION_NAME == "gen_ai.operation.name" - assert OPERATION_NAME_EXECUTE_TOOL == "execute_tool" - assert TOOL_TYPE_FUNCTION == "function" - assert TOOL_TYPE_EXTENSION == "extension" - assert TOOL_TYPE_DATASTORE == "datastore" - assert GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN == "gen_ai.response.time_to_first_token" - assert GEN_AI_AGENT_NAME == "gen_ai.agent.name" - assert GEN_AI_REACT_ROUND == "gen_ai.react.round" - assert GEN_AI_REACT_FINISH_REASON == "gen_ai.react.finish_reason" - - -def test_gen_ai_span_kind_enum(): - assert GenAISpanKind.CHAIN == "CHAIN" - assert GenAISpanKind.RETRIEVER == "RETRIEVER" - assert GenAISpanKind.RERANKER == "RERANKER" - assert GenAISpanKind.LLM == "LLM" - assert GenAISpanKind.EMBEDDING == "EMBEDDING" - assert GenAISpanKind.TOOL == "TOOL" - assert GenAISpanKind.AGENT == "AGENT" - assert GenAISpanKind.TASK == "TASK" - assert GenAISpanKind.STEP == "STEP" - - # Verify iteration works (covers the class definition) - kinds = list(GenAISpanKind) - assert len(kinds) == 9 - assert "LLM" in kinds diff --git a/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_aliyun_trace.py b/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_aliyun_trace.py deleted file mode 100644 index 0225221faf5155..00000000000000 --- a/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_aliyun_trace.py +++ /dev/null @@ -1,1145 +0,0 @@ -from __future__ import annotations - -import json -from contextlib import nullcontext -from datetime import UTC, datetime -from types import SimpleNamespace -from typing import Any, cast -from unittest.mock import MagicMock - -import dify_trace_aliyun.aliyun_trace as aliyun_trace_module -import httpx -import pytest -from dify_trace_aliyun.aliyun_trace import AliyunDataTrace -from dify_trace_aliyun.config import AliyunConfig -from dify_trace_aliyun.configuration import AliyunConfigurationError -from dify_trace_aliyun.entities.aliyun_trace_entity import SpanData, TraceMetadata -from dify_trace_aliyun.entities.semconv import ( - DIFY_NODE_TYPE, - GEN_AI_AGENT_NAME, - GEN_AI_COMPLETION, - GEN_AI_INPUT_MESSAGE, - GEN_AI_OPERATION_NAME, - GEN_AI_OUTPUT_MESSAGE, - GEN_AI_PROMPT, - GEN_AI_PROVIDER_NAME, - GEN_AI_REACT_ROUND, - GEN_AI_REQUEST_MODEL, - GEN_AI_RESPONSE_FINISH_REASON, - GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN, - GEN_AI_TOOL_CALL_ARGUMENTS, - GEN_AI_TOOL_CALL_ID, - GEN_AI_TOOL_CALL_RESULT, - GEN_AI_TOOL_DESCRIPTION, - GEN_AI_TOOL_NAME, - GEN_AI_TOOL_TYPE, - GEN_AI_USAGE_TOTAL_TOKENS, - INPUT_VALUE, - OPERATION_NAME_CHAT, - OPERATION_NAME_EXECUTE_TOOL, - OPERATION_NAME_INVOKE_AGENT, - OPERATION_NAME_REACT, - OUTPUT_VALUE, - RETRIEVAL_DOCUMENT, - RETRIEVAL_QUERY, - TOOL_TYPE_DATASTORE, - TOOL_TYPE_FUNCTION, - GenAISpanKind, -) -from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest -from opentelemetry.trace import Link, SpanContext, SpanKind, Status, StatusCode, TraceFlags - -from core.helper import ssrf_proxy -from core.ops.entities.trace_entity import ( - DatasetRetrievalTraceInfo, - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - SuggestedQuestionTraceInfo, - ToolTraceInfo, - WorkflowTraceInfo, -) -from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus -from tests.helpers.unit.core.ops.trace_fixtures import workflow_node_execution - - -class RecordingTraceClient: - def __init__(self, service_name: str = "service", endpoint: str = "endpoint"): - self.service_name = service_name - self.endpoint = endpoint - self.added_spans: list[object] = [] - - def batch(self) -> nullcontext[None]: - return nullcontext() - - def add_span(self, span: object) -> None: - self.added_spans.append(span) - - def api_check(self) -> bool: - return True - - def get_project_url(self) -> str: - return "project-url" - - -def _dt() -> datetime: - return datetime(2024, 1, 1, 0, 0, 0, tzinfo=UTC) - - -def _make_link(trace_id: int = 1, span_id: int = 2) -> Link: - context = SpanContext( - trace_id=trace_id, - span_id=span_id, - is_remote=False, - trace_flags=TraceFlags(TraceFlags.SAMPLED), - ) - return Link(context) - - -def _make_trace_metadata( - trace_id: int = 1, - workflow_span_id: int = 2, - session_id: str = "s", - user_id: str = "u", - links: list[Link] | None = None, -) -> TraceMetadata: - return TraceMetadata( - trace_id=trace_id, - workflow_span_id=workflow_span_id, - session_id=session_id, - user_id=user_id, - links=[] if links is None else links, - ) - - -def _recording_trace_client(trace_instance: AliyunDataTrace) -> RecordingTraceClient: - return cast(RecordingTraceClient, trace_instance.trace_client) - - -def _recorded_span_data(trace_instance: AliyunDataTrace) -> list[SpanData]: - return cast(list[SpanData], _recording_trace_client(trace_instance).added_spans) - - -def _make_workflow_trace_info(**overrides) -> WorkflowTraceInfo: - defaults = { - "workflow_id": "workflow-id", - "tenant_id": "tenant-id", - "workflow_run_id": "00000000-0000-0000-0000-000000000001", - "workflow_run_elapsed_time": 1.0, - "workflow_run_status": "succeeded", - "workflow_run_inputs": {"sys.query": "hello"}, - "workflow_run_outputs": {"answer": "world"}, - "workflow_run_version": "v1", - "total_tokens": 1, - "file_list": [], - "query": "hello", - "metadata": {"conversation_id": "conv", "user_id": "u", "app_id": "app"}, - "message_id": None, - "start_time": _dt(), - "end_time": _dt(), - "trace_id": "550e8400-e29b-41d4-a716-446655440000", - } - defaults.update(overrides) - return WorkflowTraceInfo(**defaults) - - -def _make_message_trace_info(**overrides) -> MessageTraceInfo: - defaults = { - "conversation_model": "chat", - "message_tokens": 1, - "answer_tokens": 2, - "total_tokens": 3, - "conversation_mode": "chat", - "metadata": {"conversation_id": "conv", "ls_model_name": "m", "ls_provider": "p"}, - "message_id": "00000000-0000-0000-0000-000000000002", - "message_data": SimpleNamespace(from_account_id="acc", from_end_user_id=None), - "inputs": {"prompt": "hi"}, - "outputs": "ok", - "start_time": _dt(), - "end_time": _dt(), - "error": None, - "trace_id": "550e8400-e29b-41d4-a716-446655440000", - } - defaults.update(overrides) - return MessageTraceInfo(**defaults) - - -def _make_dataset_retrieval_trace_info(**overrides) -> DatasetRetrievalTraceInfo: - defaults = { - "metadata": {"conversation_id": "conv", "user_id": "u"}, - "message_id": "00000000-0000-0000-0000-000000000003", - "message_data": SimpleNamespace(), - "inputs": "q", - "documents": [SimpleNamespace()], - "start_time": _dt(), - "end_time": _dt(), - "trace_id": "550e8400-e29b-41d4-a716-446655440000", - } - defaults.update(overrides) - return DatasetRetrievalTraceInfo(**defaults) - - -def _make_tool_trace_info(**overrides) -> ToolTraceInfo: - defaults = { - "tool_name": "tool", - "tool_inputs": {"x": 1}, - "tool_outputs": "out", - "tool_config": {"desc": "d"}, - "tool_parameters": {}, - "time_cost": 0.1, - "metadata": {"conversation_id": "conv", "user_id": "u"}, - "message_id": "00000000-0000-0000-0000-000000000004", - "message_data": SimpleNamespace(), - "inputs": {"i": "v"}, - "outputs": {"o": "v"}, - "start_time": _dt(), - "end_time": _dt(), - "error": None, - "trace_id": "550e8400-e29b-41d4-a716-446655440000", - } - defaults.update(overrides) - return ToolTraceInfo(**defaults) - - -def _make_suggested_question_trace_info(**overrides) -> SuggestedQuestionTraceInfo: - defaults = { - "suggested_question": ["q1", "q2"], - "level": "info", - "total_tokens": 1, - "metadata": {"conversation_id": "conv", "user_id": "u", "ls_model_name": "m", "ls_provider": "p"}, - "message_id": "00000000-0000-0000-0000-000000000005", - "inputs": {"i": 1}, - "start_time": _dt(), - "end_time": _dt(), - "error": None, - "trace_id": "550e8400-e29b-41d4-a716-446655440000", - } - defaults.update(overrides) - return SuggestedQuestionTraceInfo(**defaults) - - -@pytest.fixture -def trace_instance(monkeypatch: pytest.MonkeyPatch) -> AliyunDataTrace: - monkeypatch.setattr(aliyun_trace_module, "TraceClient", RecordingTraceClient) - # Mock get_service_account_with_tenant to avoid DB errors - monkeypatch.setattr(AliyunDataTrace, "get_service_account_with_tenant", lambda self, app_id: MagicMock()) - - config = AliyunConfig(app_name="app", license_key="k", endpoint="https://example.com") - trace = AliyunDataTrace(config) - return trace - - -def test_init_builds_endpoint_and_client() -> None: - config = AliyunConfig(app_name="my-app", license_key="license", endpoint="https://example.com/prefix") - trace = AliyunDataTrace(config) - config.license_key = "mutated" - assert trace.trace_client.endpoint == "https://example.com/prefix/adapt_license/api/otlp/traces" - assert trace.trace_client.resource.attributes["service.name"] == "my-app" - assert trace.trace_config == config - - -def test_trace_dispatches_to_correct_methods(trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch): - workflow_trace = MagicMock() - message_trace = MagicMock() - suggested_question_trace = MagicMock() - dataset_retrieval_trace = MagicMock() - tool_trace = MagicMock() - monkeypatch.setattr(trace_instance, "workflow_trace", workflow_trace) - monkeypatch.setattr(trace_instance, "message_trace", message_trace) - monkeypatch.setattr(trace_instance, "suggested_question_trace", suggested_question_trace) - monkeypatch.setattr(trace_instance, "dataset_retrieval_trace", dataset_retrieval_trace) - monkeypatch.setattr(trace_instance, "tool_trace", tool_trace) - - trace_instance.trace(_make_workflow_trace_info()) - workflow_trace.assert_called_once() - - trace_instance.trace(_make_message_trace_info()) - message_trace.assert_called_once() - - trace_instance.trace(_make_suggested_question_trace_info()) - suggested_question_trace.assert_called_once() - - trace_instance.trace(_make_dataset_retrieval_trace_info()) - dataset_retrieval_trace.assert_called_once() - - trace_instance.trace(_make_tool_trace_info()) - tool_trace.assert_called_once() - - # Branches that do nothing but should be covered - trace_instance.trace(ModerationTraceInfo(flagged=False, action="allow", preset_response="", query="", metadata={})) - trace_instance.trace(GenerateNameTraceInfo(tenant_id="t", metadata={})) - - -@pytest.mark.parametrize( - ("status", "body", "verified"), - [ - (200, b'"success"', True), - (200, b"", False), - (200, b"{}", False), - (403, b"DO_NOT_EXPOSE_BODY", False), - (405, b"", False), - ], -) -def test_api_check_verifies_credentials_with_zero_spans( - trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch, status: int, body: bytes, verified: bool -) -> None: - sent: list[ExportTraceServiceRequest] = [] - - def send(method: str, url: str, **options: Any) -> httpx.Response: - assert method == "POST" - assert url == "https://example.com/adapt_k/api/otlp/traces" - assert options["max_retries"] == 0 - assert options["follow_redirects"] is False - payload = ExportTraceServiceRequest() - payload.ParseFromString(options["content"]) - sent.append(payload) - return httpx.Response(status, content=body) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - if verified: - assert trace_instance.api_check() is True - else: - with pytest.raises(AliyunConfigurationError) as error: - trace_instance.api_check() - assert "DO_NOT_EXPOSE_BODY" not in str(error.value) - assert "adapt_" not in str(error.value) - assert len(sent) == 1 - assert len(sent[0].resource_spans) == 1 - resource = sent[0].resource_spans[0] - assert len(resource.scope_spans) == 0 - assert {item.key: item.value.string_value for item in resource.resource.attributes} == {"service.name": "app"} - - -def test_get_project_url_success(trace_instance: AliyunDataTrace): - assert trace_instance.get_project_url() == "https://arms.console.aliyun.com/#/llm" - - -def test_workflow_trace_adds_workflow_and_node_spans(trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(aliyun_trace_module, "convert_to_trace_id", lambda _: 111) - monkeypatch.setattr( - aliyun_trace_module, "convert_to_span_id", lambda _, span_type: {"workflow": 222}.get(span_type, 0) - ) - monkeypatch.setattr(aliyun_trace_module, "create_links_from_trace_id", lambda _: []) - - add_workflow_span = MagicMock() - get_workflow_node_executions = MagicMock(return_value=[MagicMock(), MagicMock()]) - build_workflow_node_span = MagicMock(side_effect=["span-1", "span-2"]) - monkeypatch.setattr(trace_instance, "add_workflow_span", add_workflow_span) - monkeypatch.setattr(trace_instance, "get_workflow_node_executions", get_workflow_node_executions) - monkeypatch.setattr(trace_instance, "build_workflow_node_span", build_workflow_node_span) - - trace_info = _make_workflow_trace_info( - trace_id="abcd", metadata={"conversation_id": "c", "user_id": "u", "app_id": "app"} - ) - trace_instance.workflow_trace(trace_info) - - add_workflow_span.assert_called_once() - passed_trace_metadata = cast(TraceMetadata, add_workflow_span.call_args.args[1]) - assert passed_trace_metadata.trace_id == 111 - assert passed_trace_metadata.workflow_span_id == 222 - assert passed_trace_metadata.session_id == "c" - assert passed_trace_metadata.user_id == "u" - assert passed_trace_metadata.links == [] - - assert _recording_trace_client(trace_instance).added_spans == ["span-1", "span-2"] - - -def test_message_trace_returns_early_if_no_message_data(trace_instance: AliyunDataTrace): - trace_info = _make_message_trace_info(message_data=None) - trace_instance.message_trace(trace_info) - assert _recording_trace_client(trace_instance).added_spans == [] - - -def test_message_trace_creates_message_and_llm_spans(trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(aliyun_trace_module, "convert_to_trace_id", lambda _: 10) - monkeypatch.setattr( - aliyun_trace_module, - "convert_to_span_id", - lambda _, span_type: {"message": 20, "llm": 30}.get(span_type, 0), - ) - monkeypatch.setattr(aliyun_trace_module, "convert_datetime_to_nanoseconds", lambda _: 123) - monkeypatch.setattr(aliyun_trace_module, "get_user_id_from_message_data", lambda _: "user") - monkeypatch.setattr(aliyun_trace_module, "create_links_from_trace_id", lambda _: []) - - status = Status(StatusCode.OK) - monkeypatch.setattr(aliyun_trace_module, "create_status_from_error", lambda _: status) - - trace_info = _make_message_trace_info( - metadata={"conversation_id": "conv", "ls_model_name": "model", "ls_provider": "provider"}, - message_tokens=7, - answer_tokens=11, - total_tokens=18, - outputs="completion", - ) - trace_instance.message_trace(trace_info) - - spans = _recorded_span_data(trace_instance) - assert len(spans) == 2 - message_span, llm_span = spans - - assert message_span.name == "message" - assert message_span.trace_id == 10 - assert message_span.parent_span_id is None - assert message_span.span_id == 20 - assert message_span.span_kind == SpanKind.SERVER - assert message_span.status == status - assert message_span.attributes["gen_ai.span.kind"] == GenAISpanKind.CHAIN - - assert llm_span.name == "llm" - assert llm_span.parent_span_id == 20 - assert llm_span.span_id == 30 - assert llm_span.status == status - assert llm_span.attributes[GEN_AI_REQUEST_MODEL] == "model" - assert llm_span.attributes[GEN_AI_USAGE_TOTAL_TOKENS] == "18" - - -def test_dataset_retrieval_trace_returns_early_if_no_message_data(trace_instance: AliyunDataTrace): - trace_info = _make_dataset_retrieval_trace_info(message_data=None) - trace_instance.dataset_retrieval_trace(trace_info) - assert _recording_trace_client(trace_instance).added_spans == [] - - -def test_dataset_retrieval_trace_creates_span(trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(aliyun_trace_module, "convert_to_trace_id", lambda _: 1) - monkeypatch.setattr( - aliyun_trace_module, "convert_to_span_id", lambda _, span_type: {"message": 2}.get(span_type, 0) - ) - monkeypatch.setattr(aliyun_trace_module, "generate_span_id", lambda: 3) - monkeypatch.setattr(aliyun_trace_module, "convert_datetime_to_nanoseconds", lambda _: 123) - monkeypatch.setattr(aliyun_trace_module, "create_links_from_trace_id", lambda _: []) - monkeypatch.setattr(aliyun_trace_module, "extract_retrieval_documents", lambda _: [{"doc": "d"}]) - - trace_instance.dataset_retrieval_trace(_make_dataset_retrieval_trace_info(inputs="query")) - spans = _recorded_span_data(trace_instance) - assert len(spans) == 1 - span = spans[0] - assert span.name == "dataset_retrieval" - assert span.attributes[RETRIEVAL_QUERY] == "query" - assert span.attributes[RETRIEVAL_DOCUMENT] == '[{"doc": "d"}]' - - -def test_tool_trace_returns_early_if_no_message_data(trace_instance: AliyunDataTrace): - trace_info = _make_tool_trace_info(message_data=None) - trace_instance.tool_trace(trace_info) - assert _recording_trace_client(trace_instance).added_spans == [] - - -def test_tool_trace_creates_span(trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(aliyun_trace_module, "convert_to_trace_id", lambda _: 10) - monkeypatch.setattr( - aliyun_trace_module, "convert_to_span_id", lambda _, span_type: {"message": 20}.get(span_type, 0) - ) - monkeypatch.setattr(aliyun_trace_module, "generate_span_id", lambda: 30) - monkeypatch.setattr(aliyun_trace_module, "convert_datetime_to_nanoseconds", lambda _: 123) - monkeypatch.setattr(aliyun_trace_module, "create_links_from_trace_id", lambda _: []) - status = Status(StatusCode.OK) - monkeypatch.setattr(aliyun_trace_module, "create_status_from_error", lambda _: status) - - trace_instance.tool_trace( - _make_tool_trace_info( - tool_name="my-tool", - tool_inputs={"a": 1}, - tool_outputs="tool-out", - tool_config={"description": "x", "tool_provider_type": "builtin"}, - inputs={"i": 1}, - metadata={"conversation_id": "conv", "user_id": "u", "node_execution_id": "exec-1"}, - ) - ) - - spans = _recorded_span_data(trace_instance) - assert len(spans) == 1 - span = spans[0] - assert span.name == "my-tool" - assert span.status == status - assert span.attributes[GEN_AI_OPERATION_NAME] == OPERATION_NAME_EXECUTE_TOOL - assert span.attributes[GEN_AI_TOOL_NAME] == "my-tool" - assert span.attributes[GEN_AI_TOOL_DESCRIPTION] == "x" - assert span.attributes[GEN_AI_TOOL_TYPE] == TOOL_TYPE_FUNCTION - assert span.attributes[GEN_AI_TOOL_CALL_ID] == "exec-1" - assert span.attributes[GEN_AI_TOOL_CALL_ARGUMENTS] == '{"a": 1}' - assert span.attributes[GEN_AI_TOOL_CALL_RESULT] == "tool-out" - - -def test_get_workflow_node_executions_requires_app_id(trace_instance: AliyunDataTrace): - trace_info = _make_workflow_trace_info(metadata={"conversation_id": "c"}) - with pytest.raises(ValueError, match="No app_id found in trace_info metadata"): - trace_instance.get_workflow_node_executions(trace_info) - - -def test_get_workflow_node_executions_builds_repo_and_fetches( - trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch -): - trace_info = _make_workflow_trace_info(metadata={"app_id": "app", "conversation_id": "c", "user_id": "u"}) - - account = object() - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", MagicMock(return_value=account)) - monkeypatch.setattr(aliyun_trace_module, "sessionmaker", MagicMock()) - monkeypatch.setattr(aliyun_trace_module, "db", SimpleNamespace(engine="engine")) - - repo = MagicMock() - repo.get_by_workflow_execution.return_value = ["node1"] - mock_factory = MagicMock() - mock_factory.create_workflow_node_execution_repository.return_value = repo - monkeypatch.setattr(aliyun_trace_module, "DifyCoreRepositoryFactory", mock_factory) - - result = trace_instance.get_workflow_node_executions(trace_info) - assert result == ["node1"] - repo.get_by_workflow_execution.assert_called_once_with(workflow_execution_id=trace_info.workflow_run_id) - - -def test_build_workflow_node_span_routes_llm_type(trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch): - node_execution = workflow_node_execution() - trace_info = _make_workflow_trace_info() - trace_metadata = _make_trace_metadata() - - monkeypatch.setattr(trace_instance, "build_workflow_llm_span", MagicMock(return_value="llm")) - - node_execution.node_type = BuiltinNodeTypes.LLM - assert trace_instance.build_workflow_node_span(node_execution, trace_info, trace_metadata) == "llm" - - -def test_build_workflow_node_span_routes_knowledge_retrieval_type( - trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch -): - node_execution = workflow_node_execution() - trace_info = _make_workflow_trace_info() - trace_metadata = _make_trace_metadata() - - monkeypatch.setattr(trace_instance, "build_workflow_retrieval_span", MagicMock(return_value="retrieval")) - - node_execution.node_type = BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL - assert trace_instance.build_workflow_node_span(node_execution, trace_info, trace_metadata) == "retrieval" - - -def test_build_workflow_node_span_routes_tool_type(trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch): - node_execution = workflow_node_execution() - trace_info = _make_workflow_trace_info() - trace_metadata = _make_trace_metadata() - - monkeypatch.setattr(trace_instance, "build_workflow_tool_span", MagicMock(return_value="tool")) - - node_execution.node_type = BuiltinNodeTypes.TOOL - assert trace_instance.build_workflow_node_span(node_execution, trace_info, trace_metadata) == "tool" - - -def test_build_workflow_node_span_routes_code_type(trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch): - node_execution = workflow_node_execution() - trace_info = _make_workflow_trace_info() - trace_metadata = _make_trace_metadata() - - monkeypatch.setattr(trace_instance, "build_workflow_task_span", MagicMock(return_value="task")) - - node_execution.node_type = BuiltinNodeTypes.CODE - assert trace_instance.build_workflow_node_span(node_execution, trace_info, trace_metadata) == "task" - - -def test_build_workflow_node_span_handles_errors( - trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture -): - node_execution = workflow_node_execution() - trace_info = _make_workflow_trace_info() - trace_metadata = _make_trace_metadata() - - monkeypatch.setattr(trace_instance, "build_workflow_task_span", MagicMock(side_effect=RuntimeError("boom"))) - node_execution.node_type = BuiltinNodeTypes.CODE - - assert trace_instance.build_workflow_node_span(node_execution, trace_info, trace_metadata) is None - assert "Error occurred in build_workflow_node_span" in caplog.text - - -@pytest.mark.parametrize( - "node_type", - [ - BuiltinNodeTypes.QUESTION_CLASSIFIER, - BuiltinNodeTypes.HTTP_REQUEST, - BuiltinNodeTypes.CODE, - BuiltinNodeTypes.AGENT, - BuiltinNodeTypes.IF_ELSE, - "custom-node", - ], -) -def test_build_workflow_task_span(trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch, node_type: str): - monkeypatch.setattr(aliyun_trace_module, "convert_to_span_id", lambda _, __: 9) - monkeypatch.setattr(aliyun_trace_module, "convert_datetime_to_nanoseconds", lambda _: 123) - status = Status(StatusCode.OK) - monkeypatch.setattr(aliyun_trace_module, "get_workflow_node_status", lambda _: status) - - trace_metadata = _make_trace_metadata(links=[_make_link()]) - node_execution = workflow_node_execution() - node_execution.id = "node-id" - node_execution.title = "title" - node_execution.node_type = node_type - node_execution.process_data = None - node_execution.inputs = {"a": 1} - node_execution.outputs = {"b": 2} - node_execution.created_at = _dt() - node_execution.finished_at = _dt() - - span = trace_instance.build_workflow_task_span(_make_workflow_trace_info(), node_execution, trace_metadata) - assert span.trace_id == 1 - assert span.span_id == 9 - assert span.parent_span_id == trace_metadata.workflow_span_id - assert span.name == "title" - assert span.start_time == span.end_time == 123 - assert span.links == trace_metadata.links - assert span.span_kind == SpanKind.INTERNAL - assert span.status.status_code == StatusCode.OK - assert span.attributes["gen_ai.span.kind"] == GenAISpanKind.TASK - assert span.attributes[DIFY_NODE_TYPE] == node_type - assert json.loads(span.attributes[INPUT_VALUE]) == {"a": 1} - assert json.loads(span.attributes[OUTPUT_VALUE]) == {"b": 2} - - -@pytest.mark.parametrize("error", [None, "Connection refused"]) -def test_http_task_span_uses_request_snapshot(trace_instance: AliyunDataTrace, error: str | None): - request = ( - "POST /collect HTTP/1.1\r\nHost: example.com\r\n" - "Authorization: Bearer ******\r\nContent-Type: application/json\r\n\r\n" - '{"message":"synthetic 测试"}' - ) - node_execution = workflow_node_execution() - node_execution.id = "00000000-0000-0000-0000-000000000006" - node_execution.node_type = BuiltinNodeTypes.HTTP_REQUEST - node_execution.title = "HTTP request" - node_execution.inputs = {"fallback": True} - node_execution.outputs = {"status_code": 200} if error is None else {"error_message": error} - node_execution.process_data = {"request": request, "unrelated": "must not be exported"} - node_execution.error = error - node_execution.status = ( - WorkflowNodeExecutionStatus.SUCCEEDED if error is None else WorkflowNodeExecutionStatus.FAILED - ) - node_execution.created_at = _dt() - node_execution.finished_at = _dt() - metadata = _make_trace_metadata(links=[_make_link()]) - - span = trace_instance.build_workflow_node_span(node_execution, _make_workflow_trace_info(), metadata) - - assert span is not None - assert json.loads(span.attributes[INPUT_VALUE]) == {"request": request} - assert json.loads(span.attributes[OUTPUT_VALUE]) == node_execution.outputs - assert span.attributes[DIFY_NODE_TYPE] == "http-request" - assert span.attributes["gen_ai.span.kind"] == GenAISpanKind.TASK - assert span.name == node_execution.title - assert span.parent_span_id == metadata.workflow_span_id - assert span.trace_id == metadata.trace_id - assert span.links == metadata.links - assert span.span_kind == SpanKind.INTERNAL - assert span.status.status_code == (StatusCode.OK if error is None else StatusCode.ERROR) - assert span.status.description == error - assert node_execution.process_data == {"request": request, "unrelated": "must not be exported"} - - -@pytest.mark.parametrize( - "process_data", - [ - None, - {}, - [], - "invalid", - 42, - {"other": "data"}, - {"request": None}, - {"request": ""}, - {"request": " \r\n"}, - {"request": 42}, - {"request": True}, - {"request": {}}, - {"request": []}, - ], -) -@pytest.mark.parametrize("inputs", [None, {}, {"fallback": "original input"}]) -def test_http_task_span_falls_back_to_inputs(trace_instance: AliyunDataTrace, process_data, inputs): - node_execution = workflow_node_execution() - node_execution.id = "00000000-0000-0000-0000-000000000006" - node_execution.node_type = BuiltinNodeTypes.HTTP_REQUEST - node_execution.title = "HTTP request" - node_execution.inputs = inputs - node_execution.outputs = {} - node_execution.process_data = process_data - node_execution.error = None - node_execution.status = WorkflowNodeExecutionStatus.SUCCEEDED - node_execution.created_at = _dt() - node_execution.finished_at = _dt() - - span = trace_instance.build_workflow_node_span(node_execution, _make_workflow_trace_info(), _make_trace_metadata()) - - assert span is not None - assert span.attributes[INPUT_VALUE] == aliyun_trace_module.serialize_json_data(inputs) - assert span.attributes[DIFY_NODE_TYPE] == "http-request" - - -def test_non_http_task_span_ignores_request_snapshot(trace_instance: AliyunDataTrace): - node_execution = workflow_node_execution() - node_execution.id = "00000000-0000-0000-0000-000000000007" - node_execution.node_type = BuiltinNodeTypes.CODE - node_execution.title = "Code" - node_execution.inputs = {"original": "input"} - node_execution.outputs = {"result": "output"} - node_execution.process_data = {"request": "POST /collect HTTP/1.1\r\nHost: example.com"} - node_execution.error = None - node_execution.status = WorkflowNodeExecutionStatus.SUCCEEDED - node_execution.created_at = _dt() - node_execution.finished_at = _dt() - - span = trace_instance.build_workflow_node_span(node_execution, _make_workflow_trace_info(), _make_trace_metadata()) - - assert span is not None - assert json.loads(span.attributes[INPUT_VALUE]) == node_execution.inputs - assert span.attributes[DIFY_NODE_TYPE] == "code" - - -def test_build_workflow_tool_span(trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(aliyun_trace_module, "convert_to_span_id", lambda _, __: 9) - monkeypatch.setattr(aliyun_trace_module, "convert_datetime_to_nanoseconds", lambda _: 123) - status = Status(StatusCode.OK) - monkeypatch.setattr(aliyun_trace_module, "get_workflow_node_status", lambda _: status) - - trace_metadata = _make_trace_metadata(links=[_make_link()]) - node_execution = workflow_node_execution() - node_execution.id = "node-id" - node_execution.title = "my-tool" - node_execution.inputs = {"a": 1} - node_execution.outputs = {"b": 2} - node_execution.created_at = _dt() - node_execution.finished_at = _dt() - node_execution.metadata = { - WorkflowNodeExecutionMetadataKey.TOOL_INFO: { - "provider_type": "dataset-retrieval", - "description": "search docs", - } - } - - span = trace_instance.build_workflow_tool_span(_make_workflow_trace_info(), node_execution, trace_metadata) - assert span.attributes[GEN_AI_OPERATION_NAME] == OPERATION_NAME_EXECUTE_TOOL - assert span.attributes[GEN_AI_TOOL_NAME] == "my-tool" - assert span.attributes[GEN_AI_TOOL_DESCRIPTION] == "search docs" - assert span.attributes[GEN_AI_TOOL_TYPE] == TOOL_TYPE_DATASTORE - assert span.attributes[GEN_AI_TOOL_CALL_ID] == "node-id" - assert span.attributes[GEN_AI_TOOL_CALL_ARGUMENTS] == '{"a": 1}' - assert span.attributes[GEN_AI_TOOL_CALL_RESULT] == '{"b": 2}' - assert span.status.status_code == StatusCode.OK - - # Cover metadata is None and inputs is None - node_execution.metadata = None - node_execution.inputs = None - span2 = trace_instance.build_workflow_tool_span(_make_workflow_trace_info(), node_execution, trace_metadata) - assert span2.attributes[GEN_AI_TOOL_DESCRIPTION] == "" - assert span2.attributes[GEN_AI_TOOL_TYPE] == TOOL_TYPE_FUNCTION - assert span2.attributes[GEN_AI_TOOL_CALL_ARGUMENTS] == "{}" - - -def test_build_workflow_retrieval_span(trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(aliyun_trace_module, "convert_to_span_id", lambda _, __: 9) - monkeypatch.setattr(aliyun_trace_module, "convert_datetime_to_nanoseconds", lambda _: 123) - status = Status(StatusCode.OK) - monkeypatch.setattr(aliyun_trace_module, "get_workflow_node_status", lambda _: status) - monkeypatch.setattr( - aliyun_trace_module, "format_retrieval_documents", lambda docs: [{"formatted": True}] if docs else [] - ) - - trace_metadata = _make_trace_metadata() - node_execution = workflow_node_execution() - node_execution.id = "node-id" - node_execution.title = "retrieval" - node_execution.inputs = {"query": "q"} - node_execution.outputs = {"result": [{"doc": "d"}]} - node_execution.created_at = _dt() - node_execution.finished_at = _dt() - - span = trace_instance.build_workflow_retrieval_span(_make_workflow_trace_info(), node_execution, trace_metadata) - assert span.attributes[RETRIEVAL_QUERY] == "q" - assert span.attributes[RETRIEVAL_DOCUMENT] == '[{"formatted": true}]' - - # Cover empty inputs/outputs - node_execution.inputs = None - node_execution.outputs = None - span2 = trace_instance.build_workflow_retrieval_span(_make_workflow_trace_info(), node_execution, trace_metadata) - assert span2.attributes[RETRIEVAL_QUERY] == "" - assert span2.attributes[RETRIEVAL_DOCUMENT] == "[]" - - -def test_build_workflow_llm_span(trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(aliyun_trace_module, "convert_to_span_id", lambda _, __: 9) - monkeypatch.setattr(aliyun_trace_module, "convert_datetime_to_nanoseconds", lambda _: 123) - status = Status(StatusCode.OK) - monkeypatch.setattr(aliyun_trace_module, "get_workflow_node_status", lambda _: status) - monkeypatch.setattr(aliyun_trace_module, "format_input_messages", lambda _: "in") - monkeypatch.setattr(aliyun_trace_module, "format_output_messages", lambda _: "out") - - trace_metadata = _make_trace_metadata() - node_execution = workflow_node_execution() - node_execution.id = "node-id" - node_execution.title = "llm" - node_execution.inputs = {} - node_execution.error = None - node_execution.process_data = { - "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, - "prompts": ["p"], - "model_name": "m", - "model_provider": "p1", - } - node_execution.outputs = {"text": "t", "finish_reason": "stop"} - node_execution.created_at = _dt() - node_execution.finished_at = _dt() - - span = trace_instance.build_workflow_llm_span(_make_workflow_trace_info(), node_execution, trace_metadata) - assert span.attributes[GEN_AI_USAGE_TOTAL_TOKENS] == "3" - assert span.attributes[GEN_AI_REQUEST_MODEL] == "m" - assert span.attributes[GEN_AI_PROMPT] == '["p"]' - assert span.attributes[GEN_AI_COMPLETION] == "t" - assert span.attributes[GEN_AI_RESPONSE_FINISH_REASON] == "stop" - assert span.attributes[GEN_AI_INPUT_MESSAGE] == "in" - assert span.attributes[GEN_AI_OUTPUT_MESSAGE] == "out" - - # Cover usage from outputs if not in process_data - node_execution.process_data = {"prompts": []} - node_execution.outputs = {"usage": {"total_tokens": 10}, "text": ""} - span2 = trace_instance.build_workflow_llm_span(_make_workflow_trace_info(), node_execution, trace_metadata) - assert span2.attributes[GEN_AI_USAGE_TOTAL_TOKENS] == "10" - - -def test_build_workflow_llm_span_falls_back_to_inputs_on_invoke_failure( - trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch -): - """Model invoke failures leave process_data empty; prep still populates inputs.""" - monkeypatch.setattr(aliyun_trace_module, "convert_to_span_id", lambda _, __: 9) - monkeypatch.setattr(aliyun_trace_module, "convert_datetime_to_nanoseconds", lambda _: 123) - monkeypatch.setattr(aliyun_trace_module, "get_workflow_node_status", lambda _: Status(StatusCode.ERROR, "boom")) - - trace_metadata = _make_trace_metadata() - node_execution = workflow_node_execution() - node_execution.id = "node-id" - node_execution.title = "llm" - node_execution.process_data = {} - node_execution.inputs = { - "model_name": "qwen-plus", - "model_provider": "langgenius/tongyi/tongyi", - "#context#": "ctx", - "query": "hello", - } - node_execution.outputs = {"error_message": "API rate limit", "error_type": "InvokeError"} - node_execution.error = "API rate limit" - node_execution.created_at = _dt() - node_execution.finished_at = _dt() - - span = trace_instance.build_workflow_llm_span(_make_workflow_trace_info(), node_execution, trace_metadata) - - expected_prompt = json.dumps(node_execution.inputs, ensure_ascii=False) - assert span.attributes[GEN_AI_REQUEST_MODEL] == "qwen-plus" - assert span.attributes[GEN_AI_PROVIDER_NAME] == "langgenius/tongyi/tongyi" - assert span.attributes[GEN_AI_PROMPT] == expected_prompt - assert span.attributes[INPUT_VALUE] == expected_prompt - assert span.attributes[GEN_AI_COMPLETION] == "API rate limit" - assert span.attributes[OUTPUT_VALUE] == "API rate limit" - assert span.attributes[GEN_AI_RESPONSE_FINISH_REASON] == "InvokeError" - # prompts were never persisted, so structured input.messages stays empty - assert span.attributes[GEN_AI_INPUT_MESSAGE] == "[]" - - -def _make_agent_outputs() -> dict: - return { - "text": " pong!", - "usage": { - "prompt_tokens": 100, - "completion_tokens": 225, - "total_tokens": 325, - "time_to_first_token": 0.5, - }, - "files": [], - "json": [ - { - "id": "round-1", - "parent_id": None, - "error": None, - "status": "success", - "data": {"action_input": "", "action_name": "", "observation": "", "thought": ""}, - "label": "ROUND 1", - "metadata": {"started_at": 6055.211092814, "finished_at": 6056.990671049, "total_tokens": 325}, - "node_id": "n1", - }, - { - "id": "thought-1", - "parent_id": "round-1", - "error": None, - "status": "success", - "data": {"action": " pong!", "thought": ""}, - "label": "deepseek-v4-flash Thought", - "metadata": { - "started_at": 6055.211450345, - "finished_at": 6056.990473571, - "provider": "langgenius/deepseek/deepseek", - "total_tokens": 325, - }, - "node_id": "n1", - }, - { - "id": "tool-1", - "parent_id": "round-1", - "error": None, - "status": "success", - "data": { - "tool_name": "current_time", - "tool_call_args": {"timezone": "Asia/Shanghai"}, - "output": "2026-07-28 18:00:00", - }, - "label": "CALL current_time", - # Tool CALL logs also set provider (tool provider); must not become LLM spans. - "metadata": { - "started_at": 6056.990500000, - "finished_at": 6056.990600000, - "provider": "langgenius/time/time", - }, - "node_id": "n1", - }, - {"data": []}, - ], - } - - -def test_build_workflow_node_span_routes_agent_type(trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch): - node_execution = workflow_node_execution() - trace_info = _make_workflow_trace_info() - trace_metadata = _make_trace_metadata() - - monkeypatch.setattr(trace_instance, "build_workflow_agent_span", MagicMock(return_value="agent")) - - node_execution.node_type = BuiltinNodeTypes.AGENT - assert trace_instance.build_workflow_node_span(node_execution, trace_info, trace_metadata) == "agent" - - -def test_build_workflow_agent_span(trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(aliyun_trace_module, "convert_to_span_id", lambda _, __: 9) - monkeypatch.setattr(aliyun_trace_module, "convert_datetime_to_nanoseconds", lambda _: 123) - status = Status(StatusCode.OK) - monkeypatch.setattr(aliyun_trace_module, "get_workflow_node_status", lambda _: status) - - trace_metadata = _make_trace_metadata() - node_execution = workflow_node_execution() - node_execution.id = "node-id" - node_execution.title = "my-agent" - node_execution.inputs = {"query": "ping"} - node_execution.outputs = _make_agent_outputs() - node_execution.created_at = _dt() - node_execution.finished_at = _dt() - - span = trace_instance.build_workflow_agent_span(_make_workflow_trace_info(), node_execution, trace_metadata) - assert span.attributes["gen_ai.span.kind"] == GenAISpanKind.AGENT - assert span.attributes[GEN_AI_OPERATION_NAME] == OPERATION_NAME_INVOKE_AGENT - assert span.attributes[GEN_AI_AGENT_NAME] == "my-agent" - assert span.attributes[GEN_AI_USAGE_TOTAL_TOKENS] == "325" - assert span.attributes[GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN] == 500_000_000 - assert span.attributes["output.value"] == " pong!" - - # TTFT attribute must be absent when usage does not carry it (e.g. blocking mode) - node_execution.outputs = {"text": "t", "usage": {"total_tokens": 1, "time_to_first_token": None}} - span2 = trace_instance.build_workflow_agent_span(_make_workflow_trace_info(), node_execution, trace_metadata) - assert GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN not in span2.attributes - - # Malformed usage payloads must not break agent span building - node_execution.outputs = {"text": "t", "usage": "not-a-mapping"} - span3 = trace_instance.build_workflow_agent_span(_make_workflow_trace_info(), node_execution, trace_metadata) - assert span3.attributes[GEN_AI_USAGE_TOTAL_TOKENS] == "0" - assert GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN not in span3.attributes - - -def test_build_agent_react_spans(trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch): - node_start_ns = 1_000_000_000 - monkeypatch.setattr(aliyun_trace_module, "convert_to_span_id", lambda _, __: 9) - monkeypatch.setattr(aliyun_trace_module, "convert_datetime_to_nanoseconds", lambda _: node_start_ns) - span_ids = iter([100, 200, 300]) - monkeypatch.setattr(aliyun_trace_module, "generate_span_id", lambda: next(span_ids)) - - trace_metadata = _make_trace_metadata() - node_execution = workflow_node_execution() - node_execution.id = "node-id" - node_execution.outputs = _make_agent_outputs() - node_execution.created_at = _dt() - node_execution.finished_at = _dt() - - spans = trace_instance.build_agent_react_spans(node_execution, trace_metadata) - assert len(spans) == 3 - step_span, llm_span, tool_span = spans - - assert step_span.parent_span_id == 9 - assert step_span.span_id == 100 - assert step_span.name == "ROUND 1" - assert step_span.attributes["gen_ai.span.kind"] == GenAISpanKind.STEP - assert step_span.attributes[GEN_AI_OPERATION_NAME] == OPERATION_NAME_REACT - assert step_span.attributes[GEN_AI_REACT_ROUND] == 1 - # Round starts at the earliest monotonic timestamp, anchored to node start time - assert step_span.start_time == node_start_ns - assert step_span.end_time == node_start_ns + int((6056.990671049 - 6055.211092814) * 1e9) - assert step_span.status.status_code == StatusCode.OK - - assert llm_span.parent_span_id == 100 - assert llm_span.span_id == 200 - assert llm_span.name == "deepseek-v4-flash Thought" - assert llm_span.attributes["gen_ai.span.kind"] == GenAISpanKind.LLM - assert llm_span.attributes[GEN_AI_OPERATION_NAME] == OPERATION_NAME_CHAT - assert llm_span.attributes[GEN_AI_REQUEST_MODEL] == "deepseek-v4-flash" - assert llm_span.attributes[GEN_AI_PROVIDER_NAME] == "langgenius/deepseek/deepseek" - assert llm_span.attributes[GEN_AI_USAGE_TOTAL_TOKENS] == "325" - assert llm_span.attributes[GEN_AI_COMPLETION] == " pong!" - assert llm_span.start_time == node_start_ns + int((6055.211450345 - 6055.211092814) * 1e9) - - assert tool_span.parent_span_id == 100 - assert tool_span.span_id == 300 - assert tool_span.name == "CALL current_time" - assert tool_span.attributes["gen_ai.span.kind"] == GenAISpanKind.TOOL - assert tool_span.attributes[GEN_AI_OPERATION_NAME] == OPERATION_NAME_EXECUTE_TOOL - assert tool_span.attributes[GEN_AI_TOOL_NAME] == "current_time" - assert tool_span.attributes[GEN_AI_TOOL_TYPE] == TOOL_TYPE_FUNCTION - assert tool_span.attributes[GEN_AI_TOOL_CALL_ID] == "tool-1" - assert tool_span.attributes[GEN_AI_TOOL_CALL_ARGUMENTS] == '{"timezone": "Asia/Shanghai"}' - assert tool_span.attributes[GEN_AI_TOOL_CALL_RESULT] == "2026-07-28 18:00:00" - assert tool_span.start_time == node_start_ns + int((6056.990500000 - 6055.211092814) * 1e9) - - -def test_build_agent_react_spans_returns_empty_without_log(trace_instance: AliyunDataTrace): - node_execution = workflow_node_execution() - node_execution.id = "node-id" - node_execution.outputs = {"text": "t"} - node_execution.created_at = _dt() - node_execution.finished_at = _dt() - - assert trace_instance.build_agent_react_spans(node_execution, _make_trace_metadata()) == [] - - -def test_workflow_trace_adds_react_spans_for_agent_nodes( - trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch -): - monkeypatch.setattr(aliyun_trace_module, "convert_to_trace_id", lambda _: 111) - monkeypatch.setattr(aliyun_trace_module, "convert_to_span_id", lambda _, __: 222) - monkeypatch.setattr(aliyun_trace_module, "create_links_from_trace_id", lambda _: []) - - agent_node = workflow_node_execution() - agent_node.node_type = BuiltinNodeTypes.AGENT - code_node = workflow_node_execution() - code_node.node_type = BuiltinNodeTypes.CODE - - monkeypatch.setattr(trace_instance, "add_workflow_span", MagicMock()) - monkeypatch.setattr(trace_instance, "get_workflow_node_executions", MagicMock(return_value=[agent_node, code_node])) - monkeypatch.setattr(trace_instance, "build_workflow_node_span", MagicMock(side_effect=["agent-span", "task-span"])) - build_agent_react_spans = MagicMock(return_value=["step-span", "llm-span"]) - monkeypatch.setattr(trace_instance, "build_agent_react_spans", build_agent_react_spans) - - trace_instance.workflow_trace(_make_workflow_trace_info()) - - build_agent_react_spans.assert_called_once() - assert _recording_trace_client(trace_instance).added_spans == [ - "agent-span", - "step-span", - "llm-span", - "task-span", - ] - - -def test_build_workflow_llm_span_records_time_to_first_token( - trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch -): - monkeypatch.setattr(aliyun_trace_module, "convert_to_span_id", lambda _, __: 9) - monkeypatch.setattr(aliyun_trace_module, "convert_datetime_to_nanoseconds", lambda _: 123) - monkeypatch.setattr(aliyun_trace_module, "get_workflow_node_status", lambda _: Status(StatusCode.OK)) - - trace_metadata = _make_trace_metadata() - node_execution = workflow_node_execution() - node_execution.id = "node-id" - node_execution.title = "llm" - node_execution.inputs = {} - node_execution.error = None - node_execution.process_data = {"prompts": []} - node_execution.outputs = {"text": "t", "usage": {"total_tokens": 1, "time_to_first_token": 0.123}} - node_execution.created_at = _dt() - node_execution.finished_at = _dt() - - span = trace_instance.build_workflow_llm_span(_make_workflow_trace_info(), node_execution, trace_metadata) - assert span.attributes[GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN] == 123_000_000 - assert span.attributes[GEN_AI_OPERATION_NAME] == OPERATION_NAME_CHAT - - # Absent when usage does not carry TTFT (blocking mode) - node_execution.outputs = {"text": "t", "usage": {"total_tokens": 1, "time_to_first_token": None}} - span2 = trace_instance.build_workflow_llm_span(_make_workflow_trace_info(), node_execution, trace_metadata) - assert GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN not in span2.attributes - - # Non-numeric TTFT must not raise and must be omitted, so the span is still built - node_execution.outputs = {"text": "t", "usage": {"total_tokens": 1, "time_to_first_token": "n/a"}} - span3 = trace_instance.build_workflow_llm_span(_make_workflow_trace_info(), node_execution, trace_metadata) - assert GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN not in span3.attributes - - -def test_message_trace_records_time_to_first_token(trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(aliyun_trace_module, "convert_to_trace_id", lambda _: 10) - monkeypatch.setattr(aliyun_trace_module, "convert_to_span_id", lambda _, span_type: 0) - monkeypatch.setattr(aliyun_trace_module, "convert_datetime_to_nanoseconds", lambda _: 123) - monkeypatch.setattr(aliyun_trace_module, "get_user_id_from_message_data", lambda _: "user") - monkeypatch.setattr(aliyun_trace_module, "create_links_from_trace_id", lambda _: []) - - trace_instance.message_trace(_make_message_trace_info(gen_ai_server_time_to_first_token=0.25)) - - llm_span = _recorded_span_data(trace_instance)[1] - assert llm_span.attributes[GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN] == 250_000_000 - - -def test_add_workflow_span(trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr( - aliyun_trace_module, "convert_to_span_id", lambda _, span_type: {"message": 20}.get(span_type, 0) - ) - monkeypatch.setattr(aliyun_trace_module, "convert_datetime_to_nanoseconds", lambda _: 123) - status = Status(StatusCode.OK) - monkeypatch.setattr(aliyun_trace_module, "create_status_from_error", lambda _: status) - - trace_metadata = _make_trace_metadata() - - # CASE 1: With message_id - trace_info = _make_workflow_trace_info( - message_id="msg-1", workflow_run_inputs={"sys.query": "hi"}, workflow_run_outputs={"ans": "ok"} - ) - trace_instance.add_workflow_span(trace_info, trace_metadata) - - client = _recording_trace_client(trace_instance) - spans = _recorded_span_data(trace_instance) - assert len(spans) == 2 - message_span = spans[0] - workflow_span = spans[1] - - assert message_span.name == "message" - assert message_span.span_kind == SpanKind.SERVER - assert message_span.parent_span_id is None - - assert workflow_span.name == "workflow" - assert workflow_span.span_kind == SpanKind.INTERNAL - assert workflow_span.parent_span_id == 20 - - client.added_spans.clear() - - # CASE 2: Without message_id - trace_info_no_msg = _make_workflow_trace_info(message_id=None) - trace_instance.add_workflow_span(trace_info_no_msg, trace_metadata) - spans = _recorded_span_data(trace_instance) - assert len(spans) == 1 - span = spans[0] - assert span.name == "workflow" - assert span.span_kind == SpanKind.SERVER - assert span.parent_span_id is None - - -def test_suggested_question_trace(trace_instance: AliyunDataTrace, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(aliyun_trace_module, "convert_to_trace_id", lambda _: 10) - monkeypatch.setattr( - aliyun_trace_module, - "convert_to_span_id", - lambda _, span_type: {"message": 20, "suggested_question": 21}.get(span_type, 0), - ) - monkeypatch.setattr(aliyun_trace_module, "convert_datetime_to_nanoseconds", lambda _: 123) - monkeypatch.setattr(aliyun_trace_module, "create_links_from_trace_id", lambda _: []) - status = Status(StatusCode.OK) - monkeypatch.setattr(aliyun_trace_module, "create_status_from_error", lambda _: status) - - trace_info = _make_suggested_question_trace_info(suggested_question=["how?"]) - trace_instance.suggested_question_trace(trace_info) - - spans = _recorded_span_data(trace_instance) - assert len(spans) == 1 - span = spans[0] - assert span.name == "suggested_question" - assert span.attributes[GEN_AI_COMPLETION] == '["how?"]' diff --git a/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_config.py b/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_config.py index 8068ee1328dd90..a4c41566197b16 100644 --- a/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_config.py +++ b/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_config.py @@ -1,85 +1,165 @@ +import traceback + import pytest from dify_trace_aliyun.config import AliyunConfig +from dify_trace_sdk.config import FormField from pydantic import ValidationError +ENDPOINT = "http://tracing-analysis-dc-qd.aliyuncs.com" +EMBEDDED_KEY = "DO-NOT-PRINT" + + +def _config(**overrides: object) -> AliyunConfig: + return AliyunConfig.model_validate({"endpoint": ENDPOINT, "license_key": "synthetic-key", **overrides}) + + +def _rendered(error: ValidationError) -> str: + return str(error) + repr(error) + "".join(traceback.format_exception(error)) + + +def test_form_declares_the_license_key_as_the_only_secret() -> None: + assert AliyunConfig.secret_fields() == {"license_key"} + assert AliyunConfig.derived_fields() == set() + assert AliyunConfig.form_fields() == ( + FormField(name="app_name", secret=False, required=False), + FormField(name="license_key", secret=True, required=True), + FormField(name="endpoint", secret=False, required=True), + ) + + +def test_stored_configuration_validates_unchanged() -> None: + stored = { + "app_name": "test_app", + "license_key": "test_license_key", + "endpoint": "https://custom.tracing-analysis-dc-hz.aliyuncs.com", + } + config = AliyunConfig.model_validate({**stored, "retired_field": "ignored"}) + assert config.model_dump() == stored + assert AliyunConfig.model_validate(config.model_dump()) == config + + +def test_normalization_keeps_the_license_key_and_application_name_opaque() -> None: + config = _config( + endpoint=" HTTPS://Project.CN-QINGDAO.LOG.ALIYUNCS.COM:443/apm/trace/opentelemetry/ ", + app_name=" 中文 agent_Name 01 ", + license_key="literal%2F@a+b/c", + ) + assert config.endpoint == "https://project.cn-qingdao.log.aliyuncs.com:443/apm/trace/opentelemetry" + assert config.app_name == "中文 agent_Name 01" + assert config.license_key == "literal%2F@a+b/c" + + +@pytest.mark.parametrize("name", ["", " ", "dify_app", "应用 & Space"]) +def test_empty_application_names_fall_back_to_the_default(name: str) -> None: + assert _config(app_name=name).app_name == (name.strip() or "dify_app") + + +def test_omitted_application_name_and_empty_endpoint_use_the_defaults() -> None: + config = _config(endpoint="") + assert config.app_name == "dify_app" + assert config.endpoint == "https://tracing-analysis-dc-hz.aliyuncs.com" + + +@pytest.mark.parametrize( + "endpoint", + [ + "https://proj-xtrace-123456-cn-heyuan.cn-heyuan.log.aliyuncs.com/apm/trace/opentelemetry", + "https://proj-xtrace-123456-cn-heyuan.cn-heyuan-intranet.log.aliyuncs.com/apm/trace/opentelemetry", + "http://tracing-cn-heyuan.arms.aliyuncs.com", + "http://tracing-cn-heyuan-internal.arms.aliyuncs.com", + "http://[::1]:4318/reverse", + ], +) +def test_documented_endpoint_formats_are_kept(endpoint: str) -> None: + assert _config(endpoint=endpoint).endpoint == endpoint + + +def test_internationalized_hosts_are_ascii_encoded() -> None: + assert _config(endpoint="https://BÜCHER.example/collector/").endpoint == "https://xn--bcher-kva.example/collector" + + +@pytest.mark.parametrize( + "endpoint", + [ + "ftp://host.test", + "host.test", + "https:///missing", + "https://user:secret@host.test", + "https://host.test?", + "https://host.test#", + "https://host.test/path?key=DO-NOT-PRINT", + "https://host.test:bad", + "https://host.test:0", + "https://host.test:65536", + "https://host.test:", + "https://[invalid", + "\nhttps://host.test", + "https://ho\nst.test", + "https://host.test\\@elsewhere", + "https://host.test/a/../b", + "https://host.test/a/%2e%2e/b", + "https://host.test/a%2fb", + "https://host.test/%252e%252e", + "https://host.test/%zz", + "https://host.test/%ff", + "https://host.test/double//path", + "https://host.test/api/otlp/traces", + "https://host.test/v1/traces", + ], +) +def test_unsafe_or_completed_routes_are_rejected_without_echo(endpoint: str) -> None: + with pytest.raises(ValidationError) as caught: + _config(endpoint=endpoint) + rendered = _rendered(caught.value) + assert "Aliyun requires an HTTP(S) base URL" in rendered + assert "host.test" not in rendered + assert "DO-NOT-PRINT" not in rendered + + +@pytest.mark.parametrize( + "endpoint", + [ + f"https://tracing-analysis-dc-hz.aliyuncs.com/adapt_{EMBEDDED_KEY}@abc_{EMBEDDED_KEY}@def/api/otlp/traces", + f"https://host.test/reverse/adapt_{EMBEDDED_KEY}/api/v1/traces/", + f"https://host.test/%61dapt_{EMBEDDED_KEY}", + ], +) +def test_pasted_collector_urls_embedding_the_license_key_are_rejected_without_echo(endpoint: str) -> None: + with pytest.raises(ValidationError) as caught: + _config(license_key=EMBEDDED_KEY, endpoint=endpoint) + assert EMBEDDED_KEY not in _rendered(caught.value) + + +@pytest.mark.parametrize("endpoint", [None, 1]) +def test_non_string_endpoints_are_rejected(endpoint: object) -> None: + with pytest.raises(ValidationError): + _config(endpoint=endpoint) + + +@pytest.mark.parametrize("field", ["endpoint", "license_key"]) +def test_required_fields_are_not_invented(field: str) -> None: + stored = {"endpoint": ENDPOINT, "license_key": "synthetic-key"} + del stored[field] + with pytest.raises(ValidationError): + AliyunConfig.model_validate(stored) + + +@pytest.mark.parametrize( + "key", ["", " ", " DO-NOT-PRINT", "DO-NOT-PRINT ", "DO-NOT\nPRINT", "DO-NOT\x00PRINT", "DO-NOT\ud800PRINT"] +) +def test_unusable_license_keys_are_rejected_without_echo(key: str) -> None: + with pytest.raises(ValidationError) as caught: + _config(license_key=key) + rendered = _rendered(caught.value) + assert "Aliyun requires a nonempty license key" in rendered + assert "DO-NOT" not in rendered + + +def test_masked_license_keys_pass_validation_for_verify_and_send_to_refuse() -> None: + assert _config(license_key="***masked").license_key == "***masked" + -class TestAliyunConfig: - """Test cases for AliyunConfig""" - - def test_valid_config(self): - """Test valid Aliyun configuration""" - config = AliyunConfig( - app_name="test_app", - license_key="test_license_key", - endpoint="https://custom.tracing-analysis-dc-hz.aliyuncs.com", - ) - assert config.app_name == "test_app" - assert config.license_key == "test_license_key" - assert config.endpoint == "https://custom.tracing-analysis-dc-hz.aliyuncs.com" - - def test_default_values(self): - """Test default values are set correctly""" - config = AliyunConfig(license_key="test_license", endpoint="https://tracing-analysis-dc-hz.aliyuncs.com") - assert config.app_name == "dify_app" - - def test_missing_required_fields(self): - """Test that required fields are enforced""" - with pytest.raises(ValidationError): - AliyunConfig.model_validate({}) - - with pytest.raises(ValidationError): - AliyunConfig.model_validate({"license_key": "test_license"}) - - with pytest.raises(ValidationError): - AliyunConfig.model_validate({"endpoint": "https://tracing-analysis-dc-hz.aliyuncs.com"}) - - def test_app_name_validation_empty(self): - """Test app_name validation with empty value""" - config = AliyunConfig( - license_key="test_license", endpoint="https://tracing-analysis-dc-hz.aliyuncs.com", app_name="" - ) - assert config.app_name == "dify_app" - - def test_endpoint_validation_empty(self): - """Test endpoint validation with empty value""" - config = AliyunConfig(license_key="test_license", endpoint="") - assert config.endpoint == "https://tracing-analysis-dc-hz.aliyuncs.com" - - def test_endpoint_validation_with_path(self): - """Test endpoint validation preserves path for Aliyun endpoints""" - config = AliyunConfig( - license_key="test_license", endpoint="https://tracing-analysis-dc-hz.aliyuncs.com/api/v1/traces" - ) - assert config.endpoint == "https://tracing-analysis-dc-hz.aliyuncs.com/api/v1/traces" - - def test_endpoint_validation_invalid_scheme(self): - """Test endpoint validation rejects invalid schemes""" - with pytest.raises(ValidationError, match="URL must start with https:// or http://"): - AliyunConfig(license_key="test_license", endpoint="ftp://invalid.tracing-analysis-dc-hz.aliyuncs.com") - - def test_endpoint_validation_no_scheme(self): - """Test endpoint validation rejects URLs without scheme""" - with pytest.raises(ValidationError, match="URL must start with https:// or http://"): - AliyunConfig(license_key="test_license", endpoint="invalid.tracing-analysis-dc-hz.aliyuncs.com") - - def test_license_key_required(self): - """Test that license_key is required and cannot be empty""" - with pytest.raises(ValidationError): - AliyunConfig(license_key="", endpoint="https://tracing-analysis-dc-hz.aliyuncs.com") - - def test_valid_endpoint_format_examples(self): - """Test valid endpoint format examples from comments""" - valid_endpoints = [ - # cms2.0 public endpoint - "https://proj-xtrace-123456-cn-heyuan.cn-heyuan.log.aliyuncs.com/apm/trace/opentelemetry", - # cms2.0 intranet endpoint - "https://proj-xtrace-123456-cn-heyuan.cn-heyuan-intranet.log.aliyuncs.com/apm/trace/opentelemetry", - # xtrace public endpoint - "http://tracing-cn-heyuan.arms.aliyuncs.com", - # xtrace intranet endpoint - "http://tracing-cn-heyuan-internal.arms.aliyuncs.com", - ] - - for endpoint in valid_endpoints: - config = AliyunConfig(license_key="test_license", endpoint=endpoint) - assert config.endpoint == endpoint +@pytest.mark.parametrize("name", ["app\nname", "app\x00name", "app\x7fname"]) +def test_application_names_with_control_characters_are_rejected(name: str) -> None: + with pytest.raises(ValidationError, match="application name without control characters"): + _config(app_name=name) diff --git a/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_provider.py b/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_provider.py new file mode 100644 index 00000000000000..41cd072b8670b6 --- /dev/null +++ b/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_provider.py @@ -0,0 +1,288 @@ +import traceback +from collections.abc import Iterable +from urllib.parse import unquote, urlsplit + +import pytest +from dify_trace_aliyun.config import AliyunConfig +from dify_trace_aliyun.provider import AliyunProvider +from dify_trace_sdk.egress import EgressResponse +from dify_trace_sdk.errors import ( + EgressBlockedError, + EgressError, + ExportOutcomeUnknownError, + ExportRejectedError, + ExportRetryableError, + VerificationError, +) +from dify_trace_sdk.provider import SendResult +from dify_trace_sdk.testing import ( + MESSAGE_ID, + RecordedRequest, + RecordingEgress, + message_trace, + response, + tool_fragment, + workflow_trace, +) +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( + ExportTraceServiceRequest, + ExportTraceServiceResponse, +) +from opentelemetry.proto.common.v1.common_pb2 import KeyValue +from opentelemetry.proto.resource.v1.resource_pb2 import Resource +from opentelemetry.proto.trace.v1.trace_pb2 import Span as OtlpSpan +from opentelemetry.proto.trace.v1.trace_pb2 import Status + +LICENSE_KEY = "SYNTHETIC-LICENSE@key_1" +CONFIG = AliyunConfig( + app_name="test_app", license_key=LICENSE_KEY, endpoint="https://tracing-analysis-dc-qd.aliyuncs.com/prefix" +) +INGESTION_URL = f"https://tracing-analysis-dc-qd.aliyuncs.com/prefix/adapt_{LICENSE_KEY}/api/otlp/traces" +SUCCESS_ACK = response(200, b'"success"') + + +def _provider(*responses: EgressResponse | Exception) -> tuple[AliyunProvider, RecordingEgress]: + egress = RecordingEgress(responses, grpc_target=None) + return AliyunProvider(egress), egress + + +def _partial_success() -> bytes: + partial = ExportTraceServiceResponse() + partial.partial_success.rejected_spans = 1 + return partial.SerializeToString() + + +def _attributes(values: Iterable[KeyValue]) -> dict[str, object]: + decoded: dict[str, object] = {} + for item in values: + match item.value.WhichOneof("value"): + case "string_value": + decoded[item.key] = item.value.string_value + case "int_value": + decoded[item.key] = item.value.int_value + case "bool_value": + decoded[item.key] = item.value.bool_value + case "double_value": + decoded[item.key] = item.value.double_value + case other: + raise AssertionError(f"unexpected attribute type: {other}") + return decoded + + +def _payload(sent: RecordedRequest) -> ExportTraceServiceRequest: + assert sent.method == "POST" + assert sent.sensitive_url is True + assert sent.headers["Content-Type"] == "application/x-protobuf" + assert sent.content is not None + return ExportTraceServiceRequest.FromString(sent.content) + + +def _exported(egress: RecordingEgress) -> tuple[Resource, list[OtlpSpan]]: + (sent,) = egress.requests + assert sent.url == INGESTION_URL + (resource_spans,) = _payload(sent).resource_spans + (scope_spans,) = resource_spans.scope_spans + return resource_spans.resource, list(scope_spans.spans) + + +def _rendered(error: Exception) -> str: + return str(error) + repr(error) + "".join(traceback.format_exception(error)) + + +@pytest.mark.parametrize("ack", [SUCCESS_ACK, response(200, "success")]) +def test_verify_posts_one_resource_without_spans_and_accepts_the_literal_ack(ack: EgressResponse) -> None: + provider, egress = _provider(ack) + assert provider.verify(CONFIG) == CONFIG + + (sent,) = egress.requests + assert sent.url == INGESTION_URL + assert sent.headers == {"Content-Type": "application/x-protobuf", "Accept": "application/x-protobuf"} + (resource_spans,) = _payload(sent).resource_spans + assert list(resource_spans.scope_spans) == [] + assert _attributes(resource_spans.resource.attributes) == {"service.name": "test_app"} + + +@pytest.mark.parametrize( + "outcome", + [ + response(200, b""), + response(200, b"{}"), + response(200, b"success"), + response(200, b'"success"\n'), + response(200, b'"failure"'), + response(201, b'"success"'), + response(204, b""), + response(401, b"private-response"), + response(403, b"private-response"), + response(405, b"private-response"), + response(503, b"private-response"), + EgressError("refused", request_sent=False), + EgressError("read timeout", request_sent=True), + EgressBlockedError("private address"), + ], +) +def test_verify_failures_have_one_code_and_expose_neither_key_nor_body(outcome: EgressResponse | Exception) -> None: + provider, egress = _provider(outcome) + with pytest.raises(VerificationError) as caught: + provider.verify(CONFIG) + assert caught.value.code == "verification_failed" + rendered = _rendered(caught.value) + assert LICENSE_KEY not in rendered + assert "private-response" not in rendered + assert len(egress.requests) == 1 + + +def test_masked_license_keys_are_never_sent() -> None: + masked = AliyunConfig.model_validate({**CONFIG.model_dump(), "license_key": "SYNTHETIC***key"}) + provider, egress = _provider() + with pytest.raises(VerificationError) as caught: + provider.verify(masked) + assert caught.value.code == "invalid_license_key" + with pytest.raises(ExportRejectedError): + provider.send(masked, message_trace()) + assert egress.requests == [] + + +def test_project_url_is_the_llm_console() -> None: + provider, _ = _provider() + assert provider.project_url(CONFIG) == "https://arms.console.aliyun.com/#/llm" + + +@pytest.mark.parametrize( + ("endpoint", "expected"), + [ + ("https://log.aliyuncs.com", "https://log.aliyuncs.com/adapt_k/api/v1/traces"), + ( + "https://CN-QINGDAO.LOG.ALIYUNCS.COM.:443/prefix/", + "https://cn-qingdao.log.aliyuncs.com.:443/prefix/adapt_k/api/v1/traces", + ), + ("https://log.aliyuncs.com.evil.test", "https://log.aliyuncs.com.evil.test/adapt_k/api/otlp/traces"), + ("https://evil-log.aliyuncs.com", "https://evil-log.aliyuncs.com/adapt_k/api/otlp/traces"), + ("https://collector.test/log.aliyuncs.com", "https://collector.test/log.aliyuncs.com/adapt_k/api/otlp/traces"), + ("http://[::1]:4318/reverse/", "http://[::1]:4318/reverse/adapt_k/api/otlp/traces"), + ], +) +def test_collector_route_depends_only_on_the_hostname(endpoint: str, expected: str) -> None: + provider, egress = _provider(SUCCESS_ACK) + provider.send(AliyunConfig(license_key="k", endpoint=endpoint), message_trace()) + (sent,) = egress.requests + assert sent.url == expected + + +@pytest.mark.parametrize("key", ["raw%2F", "raw%252F", "a/b?c#d", "left@right_one+two", "密钥"]) +def test_license_key_is_encoded_exactly_once_as_one_path_segment(key: str) -> None: + provider, egress = _provider(SUCCESS_ACK) + provider.verify(AliyunConfig(license_key=key, endpoint="https://collector.test")) + (sent,) = egress.requests + path = urlsplit(sent.url).path + assert path.startswith("/adapt_") + assert path.endswith("/api/otlp/traces") + encoded = path.removeprefix("/adapt_").removesuffix("/api/otlp/traces") + assert "/" not in encoded + assert unquote(encoded) == key + + +def test_message_trace_is_exported_with_aliyun_genai_attributes() -> None: + provider, egress = _provider(SUCCESS_ACK) + trace = message_trace() + assert provider.send(CONFIG, trace) == SendResult(trace_id=trace.trace_id, trace_url=None, span_urls={}) + + resource, (root, llm) = _exported(egress) + assert _attributes(resource.attributes) == {"service.name": "test_app", "acs.arms.service.feature": "genai_app"} + assert root.trace_id.hex() == llm.trace_id.hex() == trace.trace_id == MESSAGE_ID.replace("-", "") + assert (root.span_id.hex(), root.parent_span_id) == (trace.spans[0].span_id, b"") + assert (llm.span_id.hex(), llm.parent_span_id.hex()) == (trace.spans[1].span_id, trace.spans[0].span_id) + assert (root.name, llm.name) == ("message", "llm") + assert root.status.code == llm.status.code == Status.STATUS_CODE_OK + root_attributes = _attributes(root.attributes) + assert ( + root_attributes.items() + >= { + "gen_ai.span.kind": "CHAIN", + "gen_ai.framework": "dify", + "gen_ai.session.id": "conversation-1", + "gen_ai.user.id": "end-user-1", + "input.value": '{"query": "What is Dify?"}', + "output.value": '"An LLM application platform."', + "dify.app_id": "app-1", + "dify.tenant_id": "tenant-1", + "dify.message_id": MESSAGE_ID, + }.items() + ) + assert "gen_ai.is_entry" not in root_attributes + assert ( + _attributes(llm.attributes).items() + >= { + "gen_ai.span.kind": "LLM", + "gen_ai.operation.name": "chat", + "gen_ai.request.model": "gpt-4o", + "gen_ai.provider.name": "openai", + "gen_ai.usage.input_tokens": 12, + "gen_ai.usage.output_tokens": 8, + "gen_ai.usage.total_tokens": 20, + "gen_ai.response.time_to_first_token": 250_000_000, + "dify.app_id": "app-1", + }.items() + ) + + +def test_workflow_trace_keeps_hierarchy_kinds_and_errors() -> None: + provider, egress = _provider(SUCCESS_ACK) + trace = workflow_trace() + provider.send(CONFIG, trace) + + _, (run, llm, tool) = _exported(egress) + assert [span.span_id.hex() for span in (run, llm, tool)] == [span.span_id for span in trace.spans] + assert (llm.parent_span_id.hex(), tool.parent_span_id.hex()) == (run.span_id.hex(), llm.span_id.hex()) + assert [_attributes(span.attributes)["gen_ai.span.kind"] for span in (run, llm, tool)] == ["CHAIN", "LLM", "TOOL"] + assert (run.status.code, run.status.message) == (Status.STATUS_CODE_ERROR, "Tool failed") + assert (tool.status.code, tool.status.message) == (Status.STATUS_CODE_ERROR, "Tool failed") + assert llm.status.code == Status.STATUS_CODE_OK + assert _attributes(tool.attributes)["gen_ai.tool.name"] == "tool_Search" + + +def test_fragment_root_points_at_the_separately_exported_message() -> None: + provider, egress = _provider(SUCCESS_ACK) + fragment = tool_fragment() + provider.send(CONFIG, fragment) + + _, (tool,) = _exported(egress) + assert fragment.parent is not None + assert tool.trace_id.hex() == message_trace().trace_id + assert tool.parent_span_id.hex() == fragment.parent.span_id == message_trace().root.span_id + assert _attributes(tool.attributes)["gen_ai.span.kind"] == "TOOL" + + +@pytest.mark.parametrize( + "ack", + [response(200, b""), response(200, ExportTraceServiceResponse().SerializeToString()), response(200, {})], +) +def test_standard_otlp_acknowledgments_are_accepted(ack: EgressResponse) -> None: + provider, _ = _provider(ack) + provider.send(CONFIG, message_trace()) + + +@pytest.mark.parametrize( + ("outcome", "expected"), + [ + (response(200, _partial_success()), ExportOutcomeUnknownError), + (response(200, b"not-protobuf-\xff"), ExportOutcomeUnknownError), + (response(200, "failure"), ExportOutcomeUnknownError), + (response(201, b'"success"'), ExportOutcomeUnknownError), + (response(503, b"private-response"), ExportRetryableError), + (response(429, b"private-response"), ExportRetryableError), + (response(403, b"private-response"), ExportRejectedError), + (EgressError("refused", request_sent=False), ExportRetryableError), + (EgressError("read timeout", request_sent=True), ExportOutcomeUnknownError), + (EgressBlockedError("private address"), ExportRejectedError), + ], +) +def test_send_failures_are_classified_without_exposing_the_key( + outcome: EgressResponse | Exception, expected: type[Exception] +) -> None: + provider, _ = _provider(outcome) + with pytest.raises(expected) as caught: + provider.send(CONFIG, message_trace()) + rendered = _rendered(caught.value) + assert LICENSE_KEY not in rendered + assert "private-response" not in rendered diff --git a/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_utils.py b/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_utils.py deleted file mode 100644 index f38034100ec6da..00000000000000 --- a/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_utils.py +++ /dev/null @@ -1,316 +0,0 @@ -"""Unit tests for Aliyun trace utility transformations and database lookups.""" - -import json -from collections.abc import Mapping -from typing import Any, cast -from unittest.mock import MagicMock - -import pytest -from dify_trace_aliyun.entities.semconv import ( - GEN_AI_FRAMEWORK, - GEN_AI_OPERATION_NAME, - GEN_AI_SESSION_ID, - GEN_AI_SPAN_KIND, - GEN_AI_TOOL_CALL_ARGUMENTS, - GEN_AI_TOOL_NAME, - GEN_AI_TOOL_TYPE, - GEN_AI_USER_ID, - INPUT_VALUE, - OPERATION_NAME_EXECUTE_TOOL, - OUTPUT_VALUE, - TOOL_TYPE_DATASTORE, - TOOL_TYPE_EXTENSION, - TOOL_TYPE_FUNCTION, -) -from dify_trace_aliyun.utils import ( - create_common_span_attributes, - create_gen_ai_tool_attributes, - create_links_from_trace_id, - create_status_from_error, - extract_retrieval_documents, - extract_tool_description, - format_input_messages, - format_output_messages, - format_retrieval_documents, - get_user_id_from_message_data, - get_workflow_node_status, - map_gen_ai_tool_type, - serialize_json_data, -) -from opentelemetry.trace import Link, SpanContext, StatusCode -from sqlalchemy.orm import Session - -from core.rag.models.document import Document -from graphon.enums import WorkflowNodeExecutionStatus -from models import EndUser -from models.enums import EndUserType -from tests.helpers.unit.core.ops.trace_fixtures import workflow_node_execution - - -def test_get_user_id_from_message_data_no_end_user(monkeypatch: pytest.MonkeyPatch): - message_data = MagicMock() - message_data.from_account_id = "account_id" - message_data.from_end_user_id = None - - assert get_user_id_from_message_data(message_data) == "account_id" - - -@pytest.mark.parametrize("sqlite3_session", [(EndUser,)], indirect=True) -def test_get_user_id_from_message_data_with_end_user(monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session) -> None: - message_data = MagicMock() - message_data.from_account_id = "account_id" - message_data.from_end_user_id = "end_user_id" - - end_user_data = EndUser( - id="end_user_id", - tenant_id="tenant_id", - app_id="app_id", - type=EndUserType.BROWSER, - session_id="session_id", - ) - sqlite3_session.add(end_user_data) - sqlite3_session.commit() - - from dify_trace_aliyun.utils import db - - monkeypatch.setattr(db, "session", sqlite3_session) - - assert get_user_id_from_message_data(message_data) == "session_id" - - -@pytest.mark.parametrize("sqlite3_session", [(EndUser,)], indirect=True) -def test_get_user_id_from_message_data_end_user_not_found( - monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session -) -> None: - message_data = MagicMock() - message_data.from_account_id = "account_id" - message_data.from_end_user_id = "end_user_id" - - from dify_trace_aliyun.utils import db - - monkeypatch.setattr(db, "session", sqlite3_session) - - assert get_user_id_from_message_data(message_data) == "account_id" - - -def test_create_status_from_error(): - # Case OK - status_ok = create_status_from_error(None) - assert status_ok.status_code == StatusCode.OK - - # Case Error - status_err = create_status_from_error("some error") - assert status_err.status_code == StatusCode.ERROR - assert status_err.description == "some error" - - -def test_get_workflow_node_status(): - node_execution = workflow_node_execution() - - # SUCCEEDED - node_execution.status = WorkflowNodeExecutionStatus.SUCCEEDED - status = get_workflow_node_status(node_execution) - assert status.status_code == StatusCode.OK - - # FAILED - node_execution.status = WorkflowNodeExecutionStatus.FAILED - node_execution.error = "node fail" - status = get_workflow_node_status(node_execution) - assert status.status_code == StatusCode.ERROR - assert status.description == "node fail" - - # EXCEPTION - node_execution.status = WorkflowNodeExecutionStatus.EXCEPTION - node_execution.error = "node exception" - status = get_workflow_node_status(node_execution) - assert status.status_code == StatusCode.ERROR - assert status.description == "node exception" - - # UNSET/OTHER - node_execution.status = WorkflowNodeExecutionStatus.RUNNING - status = get_workflow_node_status(node_execution) - assert status.status_code == StatusCode.UNSET - - -def test_create_links_from_trace_id(monkeypatch: pytest.MonkeyPatch): - # Mock create_link - link = Link(SpanContext(trace_id=1, span_id=2, is_remote=False)) - import dify_trace_aliyun.data_exporter.traceclient - - monkeypatch.setattr(dify_trace_aliyun.data_exporter.traceclient, "create_link", lambda trace_id_str: link) - - # Trace ID None - assert create_links_from_trace_id(None) == [] - - # Trace ID Present - links = create_links_from_trace_id("trace_id") - assert len(links) == 1 - assert links[0] == link - - -def test_extract_retrieval_documents(): - doc1 = Document( - page_content="content1", - metadata={"dataset_id": "ds1", "doc_id": "di1", "document_id": "dd1", "score": 0.9}, - ) - doc2 = Document(page_content="content2", metadata={"dataset_id": "ds2"}) # Missing some keys - - documents = [doc1, doc2] - extracted = extract_retrieval_documents(documents) - - assert len(extracted) == 2 - assert extracted[0]["content"] == "content1" - assert extracted[0]["metadata"]["dataset_id"] == "ds1" - assert extracted[0]["score"] == 0.9 - - assert extracted[1]["content"] == "content2" - assert extracted[1]["metadata"]["dataset_id"] == "ds2" - assert extracted[1]["metadata"]["doc_id"] is None - assert extracted[1]["score"] is None - - -def test_serialize_json_data(): - data = {"a": 1} - # Test ensure_ascii default (False) - assert serialize_json_data(data) == json.dumps(data, ensure_ascii=False) - # Test ensure_ascii True - assert serialize_json_data(data, ensure_ascii=True) == json.dumps(data, ensure_ascii=True) - - -def test_create_common_span_attributes(): - attrs = create_common_span_attributes( - session_id="s1", user_id="u1", span_kind="kind1", framework="fw1", inputs="in1", outputs="out1" - ) - assert attrs[GEN_AI_SESSION_ID] == "s1" - assert attrs[GEN_AI_USER_ID] == "u1" - assert attrs[GEN_AI_SPAN_KIND] == "kind1" - assert attrs[GEN_AI_FRAMEWORK] == "fw1" - assert attrs[INPUT_VALUE] == "in1" - assert attrs[OUTPUT_VALUE] == "out1" - - -def test_format_retrieval_documents(): - # Not a list - assert format_retrieval_documents(cast(list[object], "not a list")) == [] - - # Valid list - docs = [ - {"metadata": {"score": 0.8, "document_id": "doc1", "source": "src1"}, "content": "c1", "title": "t1"}, - { - "metadata": {"_source": "src2", "doc_metadata": {"extra": "val"}}, - "content": "c2", - # Missing title - }, - "not a dict", # Should be skipped - ] - formatted = format_retrieval_documents(docs) - - assert len(formatted) == 2 - assert formatted[0]["document"]["content"] == "c1" - assert formatted[0]["document"]["metadata"]["title"] == "t1" - assert formatted[0]["document"]["metadata"]["source"] == "src1" - assert formatted[0]["document"]["score"] == 0.8 - assert formatted[0]["document"]["id"] == "doc1" - - assert formatted[1]["document"]["content"] == "c2" - assert formatted[1]["document"]["metadata"]["source"] == "src2" - assert formatted[1]["document"]["metadata"]["extra"] == "val" - assert "title" not in formatted[1]["document"]["metadata"] - assert formatted[1]["document"]["score"] == 0.0 # Default - - # Exception handling - # We can trigger an exception by passing something that causes an error in the loop logic, - # but the try/except covers the whole function. - # Passing a list that contains something that throws when calling .get() - though dicts won't. - # Let's mock a dict that raises on get. - class BadDict: - def get(self, *args, **kwargs): - raise Exception("boom") - - assert format_retrieval_documents([BadDict()]) == [] - - -def test_format_input_messages(): - # Not a dict - assert format_input_messages(cast(Mapping[str, Any], None)) == serialize_json_data([]) - - # No prompts - assert format_input_messages({}) == serialize_json_data([]) - - # Valid prompts - process_data = { - "prompts": [ - {"role": "user", "text": "hello"}, - {"role": "assistant", "text": "hi"}, - {"role": "system", "text": "be helpful"}, - {"role": "tool", "text": "result"}, - {"role": "invalid", "text": "skip me"}, - "not a dict", - {"role": "user", "text": ""}, # Empty text, should be skipped? Code says `if text: message = ...` - ] - } - result = format_input_messages(process_data) - result_list = json.loads(result) - - assert len(result_list) == 4 - assert result_list[0]["role"] == "user" - assert result_list[0]["parts"][0]["content"] == "hello" - assert result_list[1]["role"] == "assistant" - assert result_list[2]["role"] == "system" - assert result_list[3]["role"] == "tool" - - # Exception path - assert format_input_messages({"prompts": [None]}) == serialize_json_data([]) - - -def test_format_output_messages(): - # Not a dict - assert format_output_messages(cast(Mapping[str, Any], None)) == serialize_json_data([]) - - # No text - assert format_output_messages({"finish_reason": "stop"}) == serialize_json_data([]) - - # Valid - outputs = {"text": "done", "finish_reason": "length"} - result = format_output_messages(outputs) - result_list = json.loads(result) - assert len(result_list) == 1 - assert result_list[0]["role"] == "assistant" - assert result_list[0]["parts"][0]["content"] == "done" - assert result_list[0]["finish_reason"] == "length" - - # Invalid finish reason - outputs2 = {"text": "done", "finish_reason": "unknown"} - result2 = format_output_messages(outputs2) - result_list2 = json.loads(result2) - assert result_list2[0]["finish_reason"] == "stop" - - # Exception path - # Trigger exception in serialize_json_data by passing non-serializable - assert format_output_messages({"text": MagicMock()}) == serialize_json_data([]) - - -def test_map_gen_ai_tool_type(): - assert map_gen_ai_tool_type("dataset-retrieval") == TOOL_TYPE_DATASTORE - assert map_gen_ai_tool_type("extension") == TOOL_TYPE_EXTENSION - assert map_gen_ai_tool_type("builtin") == TOOL_TYPE_FUNCTION - assert map_gen_ai_tool_type(None) == TOOL_TYPE_FUNCTION - - -def test_extract_tool_description(): - assert extract_tool_description({"description": "d"}) == "d" - assert extract_tool_description({"tool_description": "td"}) == "td" - assert extract_tool_description({"other": 1}) == "" - assert extract_tool_description(None) == "" - - -def test_create_gen_ai_tool_attributes(): - attrs = create_gen_ai_tool_attributes( - tool_name="search", - tool_type=TOOL_TYPE_DATASTORE, - tool_call_arguments='{"q": 1}', - ) - assert attrs[GEN_AI_OPERATION_NAME] == OPERATION_NAME_EXECUTE_TOOL - assert attrs[GEN_AI_TOOL_NAME] == "search" - assert attrs[GEN_AI_TOOL_TYPE] == TOOL_TYPE_DATASTORE - assert attrs[GEN_AI_TOOL_CALL_ARGUMENTS] == '{"q": 1}' diff --git a/api/providers/trace/trace-arize-phoenix/pyproject.toml b/api/providers/trace/trace-arize-phoenix/pyproject.toml index b9a10e83888e77..4e0b66863d7394 100644 --- a/api/providers/trace/trace-arize-phoenix/pyproject.toml +++ b/api/providers/trace/trace-arize-phoenix/pyproject.toml @@ -2,9 +2,13 @@ name = "dify-trace-arize-phoenix" version = "0.0.1" dependencies = [ - "arize-phoenix-otel==0.15.0", + "dify-trace-sdk", ] description = "Dify ops tracing provider (Arize / Phoenix)." +[project.entry-points."dify.trace_providers"] +arize = "dify_trace_arize_phoenix.provider:ArizeProvider" +phoenix = "dify_trace_arize_phoenix.provider:PhoenixProvider" + [tool.setuptools.packages.find] where = ["src"] diff --git a/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/_arize_project.py b/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/_arize_project.py new file mode 100644 index 00000000000000..f82a016dcc6e59 --- /dev/null +++ b/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/_arize_project.py @@ -0,0 +1,147 @@ +"""The link to an Arize Cloud project: looked up by ``verify``, then built without I/O. + +Arize Cloud addresses a project by organization, space, and project IDs, while +a configuration names it. The lookup runs only for the default Cloud collector, +so credentials for a custom or regional collector never reach an assumed API +host. It is optional: a key may ingest without read access, so every failure +means no link rather than a failed verification. +""" + +import re +import time +from collections.abc import Mapping +from typing import TypeIs +from urllib.parse import quote, urlencode, urlsplit + +from dify_trace_sdk.egress import Egress +from dify_trace_sdk.errors import EgressError + +from dify_trace_arize_phoenix.config import ArizeConfig + +API_BASE = "https://api.arize.com/v2" +UI_BASE = "https://app.arize.com" +CLOUD_COLLECTOR_HOST = "otlp.arize.com" +LOOKUP_BUDGET_SECONDS = 5.0 +"""Shared by every lookup request, because a user waits for verification.""" +MAX_RESPONSE_BYTES = 1024 * 1024 +PAGE_LIMIT = 100 +MAX_PAGES = 3 +MAX_CURSOR_LENGTH = 4096 +_OPAQUE_ID = re.compile(r"[A-Za-z0-9+/_=-]{1,1024}") + + +def _opaque_id(value: object) -> TypeIs[str]: + # Never decode IDs (e.g. to derive an organization from a space) or accept URL metacharacters. + return isinstance(value, str) and _OPAQUE_ID.fullmatch(value) is not None + + +def _cloud_collector(endpoint: str) -> bool: + parsed = urlsplit(endpoint) + try: + port = parsed.port + except ValueError: + return False + return parsed.scheme == "https" and parsed.hostname == CLOUD_COLLECTOR_HOST and port in {None, 443} + + +def _identity(config: ArizeConfig) -> dict[str, str] | None: + """What a destination is bound to, or ``None`` when the collector is not the default Cloud one.""" + if not _cloud_collector(config.endpoint) or not _opaque_id(config.space_id): + return None + return {"endpoint": config.endpoint, "space_id": config.space_id, "project": config.project} + + +def project_url(config: ArizeConfig) -> str | None: + """Built from the stored destination alone: presentation reads the space ID still encrypted. + + Verification resolves the destination again on every save, so a destination + that names another endpoint or project only remains from a manual edit. + """ + destination = config.arize_project_destination + if destination is None or not _cloud_collector(config.endpoint): + return None + if destination.get("endpoint") != config.endpoint or destination.get("project") != config.project: + return None + organization_id, space_id, project_id = ( + destination.get("organization_id"), + destination.get("space_id"), + destination.get("project_id"), + ) + if not _opaque_id(organization_id) or not _opaque_id(space_id) or not _opaque_id(project_id): + return None + return ( + f"{UI_BASE}/organizations/{quote(organization_id, safe='')}" + f"/spaces/{quote(space_id, safe='')}/projects/{quote(project_id, safe='')}" + ) + + +def resolve_destination(egress: Egress, config: ArizeConfig) -> dict[str, str] | None: + """Find the one project of the space named exactly ``config.project``. + + The API's name filter matches substrings, so up to ``MAX_PAGES`` pages are + read for an exact match. Denied reads, malformed or ambiguous answers, and + an exhausted budget all mean no destination. + """ + identity = _identity(config) + if identity is None: + return None + deadline = time.monotonic() + LOOKUP_BUDGET_SECONDS + headers = {"Authorization": f"Bearer {config.api_key}", "Accept": "application/json"} + + def read(path: str) -> Mapping[str, object] | None: + remaining = deadline - time.monotonic() + if remaining <= 0: + return None + try: + # Every lookup URL embeds the space ID, which is a secret field. + response = egress.request( + "GET", API_BASE + path, headers=headers, content=None, timeout=remaining, sensitive_url=True + ) + except EgressError: + return None + if response.status_code != 200 or len(response.content) > MAX_RESPONSE_BYTES or time.monotonic() > deadline: + return None + try: + body = response.json() + except (ValueError, RecursionError): + return None + return body if isinstance(body, dict) else None + + space = read(f"/spaces/{quote(config.space_id, safe='')}") + organization_id = space.get("organization_id") if space is not None else None + if space is None or space.get("id") != config.space_id or not _opaque_id(organization_id): + return None + project_ids: set[str] = set() + seen_cursors: set[str] = set() + query: dict[str, str | int] = {"space_id": config.space_id, "name": config.project, "limit": PAGE_LIMIT} + for _ in range(MAX_PAGES): + page = read("/projects?" + urlencode(query, quote_via=quote)) + projects = page.get("projects") if page is not None else None + pagination = page.get("pagination") if page is not None else None + if not isinstance(projects, list) or len(projects) > PAGE_LIMIT or not isinstance(pagination, dict): + return None + for project in projects: + if not isinstance(project, dict) or project.get("space_id") != config.space_id: + return None + project_id, name = project.get("id"), project.get("name") + if not _opaque_id(project_id) or not isinstance(name, str): + return None + if name == config.project: + project_ids.add(project_id) + if len(project_ids) > 1: + return None + has_more, cursor = pagination.get("has_more"), pagination.get("next_cursor") + if has_more is False: + if len(project_ids) != 1: + return None + return {**identity, "organization_id": organization_id, "project_id": project_ids.pop()} + if ( + has_more is not True + or not isinstance(cursor, str) + or not 0 < len(cursor) <= MAX_CURSOR_LENGTH + or cursor in seen_cursors + ): + return None + seen_cursors.add(cursor) + query["cursor"] = cursor + return None diff --git a/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/agent_trace.py b/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/agent_trace.py deleted file mode 100644 index 7100d76615ec46..00000000000000 --- a/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/agent_trace.py +++ /dev/null @@ -1,54 +0,0 @@ -"""Stable-ID OpenInference mapping for Phoenix HTTP and Arize gRPC ingestion.""" - -from collections.abc import Mapping -from typing import Any - -from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalTrace -from services.tracing.otlp_export import build_request, export_grpc, export_http -from services.tracing.provider_common import endpoint, json_value, token_usage -from services.tracing.provider_dispatch import TraceExportResult - - -def export_trace(provider: str, config: Mapping[str, Any], trace: CanonicalTrace) -> TraceExportResult: - url = endpoint( - config.get("endpoint", "https://otlp.arize.com" if provider == "arize" else "https://app.phoenix.arize.com") - ) - api_key = str(config.get("api_key") or "") - project = str(config.get("project") or "default") - - def attributes(span: CanonicalSpan) -> dict[str, object]: - result: dict[str, object] = { - "openinference.span.kind": span.kind.value.upper(), - "input.value": json_value(span.inputs), - "input.mime_type": "application/json", - "output.value": json_value(span.outputs), - "output.mime_type": "application/json", - "metadata": json_value(span.metadata), - "session.id": trace.session_id, - "dify.span.id": span.id, - } - if span.kind is CanonicalSpanKind.LLM: - model = span.metadata.get("model") or span.metadata.get("model_name") - if model: - result["llm.model_name"] = str(model) - if span.metadata.get("model_provider"): - result["llm.provider"] = str(span.metadata["model_provider"]) - names = {"input": "prompt", "output": "completion", "total": "total"} - result.update({f"llm.token_count.{names[key]}": value for key, value in token_usage(span).items()}) - elif span.kind is CanonicalSpanKind.TOOL: - result["tool.name"] = span.name - result["tool.parameters"] = json_value(span.inputs) - return result - - payload = build_request( - trace, - {"service.name": "dify-application", "openinference.project.name": project, "model_id": project}, - attributes, - ) - headers = {"api_key": api_key, "authorization": f"Bearer {api_key}"} if api_key or provider == "arize" else {} - if provider == "arize": - headers["space_id"] = str(config.get("space_id") or "") - export_grpc(provider, url, headers, payload) - else: - export_http(provider, f"{url}/v1/traces", headers, payload) - return TraceExportResult(trace_id=trace.trace_id) diff --git a/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/arize_configuration.py b/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/arize_configuration.py deleted file mode 100644 index 32053dc40f428c..00000000000000 --- a/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/arize_configuration.py +++ /dev/null @@ -1,304 +0,0 @@ -"""Optional Arize Cloud project metadata, resolved only on explicit config writes. - -Ingestion verification is separate from metadata reads. A key without read access -may still ingest successfully, so an unavailable project link is never a reason -to reject that configuration. Readers use only the frozen destination cache. -""" - -import hashlib -import json -import re -import time -from collections.abc import Mapping -from typing import Literal -from urllib.parse import quote, urlsplit - -from pydantic import JsonValue, TypeAdapter - -from core.helper import ssrf_proxy -from dify_trace_arize_phoenix.config import ArizeConfig - -DESTINATION_KEY = "arize_project_destination" -_API_BASE = "https://api.arize.com/v2" -_UI_BASE = "https://app.arize.com" -_BUDGET_SECONDS = 5.0 -_MAX_RESPONSE_BYTES = 1024 * 1024 -_PAGE_LIMIT = 100 -_MAX_PAGES = 3 -_JSON_OBJECT = TypeAdapter(dict[str, JsonValue]) -_VERIFY_TIMEOUT_SECONDS = 5.0 - -type VerificationCode = Literal[ - "invalid_configuration", "authentication_failed", "permission_denied", "connection_failed", "unexpected_response" -] - - -class ArizeConfigurationError(ValueError): - """Fixed diagnostics never contain credentials, vendor details or endpoints.""" - - def __init__(self, code: VerificationCode) -> None: - self.code = code - super().__init__(f"Arize ingestion verification failed: {code}") - - -def verify_configuration(config: ArizeConfig) -> None: - """Require one synchronous OTLP ACK, without creating any span or project. - - The default Arize collector was verified to authenticate this resource-only - request for both valid and invalid credentials. It establishes this ingest - protocol's access, not REST read permission or a project's eventual creation. - No SDK exporter, span processor or global/environment credentials are used. - """ - import grpc # type: ignore[import-untyped] - from opentelemetry.context import _SUPPRESS_INSTRUMENTATION_KEY, attach, detach, set_value - from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest - from opentelemetry.proto.collector.trace.v1.trace_service_pb2_grpc import TraceServiceStub - - try: - parsed = urlsplit(config.endpoint) - host, port = parsed.hostname, parsed.port - if ( - parsed.scheme not in {"http", "https"} - or not host - or parsed.username is not None - or parsed.password is not None - or (port is not None and not 1 <= port <= 65535) - ): - raise ArizeConfigurationError("invalid_configuration") - for value in (config.api_key, config.space_id): - if ( - not isinstance(value, str) - or not value - or "*" in value - or any(char.isspace() or ord(char) < 32 or ord(char) == 127 for char in value) - ): - raise ArizeConfigurationError("invalid_configuration") - except ValueError: - raise ArizeConfigurationError("invalid_configuration") from None - - payload = ExportTraceServiceRequest() - resource = payload.resource_spans.add().resource - for name in ("model_id", "openinference.project.name"): - resource.attributes.add(key=name).value.string_value = config.project or "default" - headers = ( - ("api_key", config.api_key), - ("space_id", config.space_id), - ("authorization", f"Bearer {config.api_key}"), - ) - authority = f"[{host}]" if ":" in host else host - target = f"{authority}:{port or 443}" - options = (("grpc.enable_retries", 0), ("grpc.max_receive_message_length", _MAX_RESPONSE_BYTES)) - channel: grpc.Channel | None = None - context = attach(set_value(_SUPPRESS_INSTRUMENTATION_KEY, True)) - try: - try: - channel = ( - grpc.insecure_channel(target, options=options) - if parsed.scheme == "http" - else grpc.secure_channel(target, grpc.ssl_channel_credentials(), options=options) - ) - result = TraceServiceStub(channel).Export( - payload, metadata=headers, timeout=_VERIFY_TIMEOUT_SECONDS, wait_for_ready=False - ) - if result.partial_success.rejected_spans or result.partial_success.error_message: - raise ArizeConfigurationError("unexpected_response") - finally: - if channel is not None: - channel.close() - except ArizeConfigurationError: - raise - except grpc.RpcError as error: - code = error.code() - if code is grpc.StatusCode.UNAUTHENTICATED: - raise ArizeConfigurationError("authentication_failed") from None - if code is grpc.StatusCode.PERMISSION_DENIED: - raise ArizeConfigurationError("permission_denied") from None - raise ArizeConfigurationError("connection_failed") from None - except Exception: - raise ArizeConfigurationError("connection_failed") from None - finally: - detach(context) - - -def _opaque_id(value: object) -> bool: - # IDs remain opaque. In particular, never derive organization IDs by - # decoding the project or space IDs, or accept path/query metacharacters. - return isinstance(value, str) and re.fullmatch(r"[A-Za-z0-9+/_=-]{1,1024}", value) is not None - - -def _identity(values: Mapping[str, object]) -> dict[str, str] | None: - try: - config = ArizeConfig.model_validate( - {key: values[key] for key in ("endpoint", "space_id", "project") if key in values} - ) - endpoint = urlsplit(config.endpoint) - # These fixed API/UI hosts belong to the documented default collector. - # Custom or other regional collectors keep ingestion support, but must - # not send their credentials to an assumed Cloud metadata destination. - if ( - endpoint.scheme != "https" - or endpoint.hostname != "otlp.arize.com" - or endpoint.port not in (None, 443) - or endpoint.username is not None - or endpoint.password is not None - or not _opaque_id(config.space_id) - ): - return None - except ValueError: - return None - assert config.space_id is not None - return {"endpoint": config.endpoint, "space_id": config.space_id, "project": config.project or "default"} - - -def project_url(values: Mapping[str, object]) -> str | None: - """Compute only a verified cached destination; never perform vendor I/O.""" - destination = values.get(DESTINATION_KEY) - if not isinstance(destination, Mapping): - return None - identity = _identity(destination) - # Legacy Arize encrypts space_id as well as api_key. Bind the public cache - # to this exact persisted record without decrypting it on status reads. - digest = _source_digest(values) - if identity is None or digest is None or destination.get("source_digest") != digest: - return None - organization_id, project_id = destination.get("organization_id"), destination.get("project_id") - if not _opaque_id(organization_id) or not _opaque_id(project_id): - return None - assert isinstance(organization_id, str) - assert isinstance(project_id, str) - return ( - f"{_UI_BASE}/organizations/{quote(organization_id, safe='')}" - f"/spaces/{quote(identity['space_id'], safe='')}/projects/{quote(project_id, safe='')}" - ) - - -def _source_digest(values: Mapping[str, object]) -> str | None: - parts = [values.get(key) for key in ("endpoint", "space_id", "project")] - if not all(isinstance(part, str) for part in parts): - return None - return hashlib.sha256(json.dumps(parts, ensure_ascii=True, separators=(",", ":")).encode()).hexdigest() - - -def bind_destination(values: Mapping[str, object], destination: dict[str, str]) -> dict[str, str]: - """Bind server-resolved metadata after legacy secret encryption has finished.""" - digest = _source_digest(values) - if digest is None: - raise ValueError("Arize persisted destination is incomplete") - return {**destination, "source_digest": digest} - - -def _get_json( - path: str, headers: dict[str, str], deadline: float, *, params: dict[str, str | int] | None = None -) -> dict[str, JsonValue] | None: - remaining = deadline - time.monotonic() - if remaining <= 0: - return None - try: - response = ssrf_proxy.get( - _API_BASE + path, - headers=headers, - params=params, - auth=None, - timeout=remaining, - max_retries=0, - follow_redirects=False, - stream_response=True, - ) - try: - if response.status_code != 200 or time.monotonic() >= deadline: - return None - encoding = response.headers.get("content-encoding", "identity").strip().lower() - if encoding not in {"", "identity"}: - return None - content = bytearray() - for chunk in response.iter_bytes(): - if time.monotonic() >= deadline or len(content) + len(chunk) > _MAX_RESPONSE_BYTES: - return None - content.extend(chunk) - if time.monotonic() >= deadline: - return None - return _JSON_OBJECT.validate_json(content, strict=True) - finally: - response.close() - except Exception: - # Metadata is optional. Do not expose request headers, response bodies, - # credential-bearing exceptions, or turn read permission into an ingest - # permission requirement. - return None - - -def resolve_project_destination(values: Mapping[str, object]) -> dict[str, str] | None: - """Resolve exact project/space ownership within one shared five-second budget. - - The documented name filter is a substring match. Read at most three scoped - pages and require one exact name/space match. Missing organization IDs on - older servers, denied reads, malformed responses and budget exhaustion all - mean no optional link, without modifying ingestion configuration or history. - """ - identity = _identity(values) - key = values.get("api_key") - if ( - identity is None - or not isinstance(key, str) - or not key - or "*" in key - or any(char.isspace() or ord(char) < 32 or ord(char) == 127 for char in key) - ): - return None - deadline = time.monotonic() + _BUDGET_SECONDS - headers = { - "Authorization": f"Bearer {key}", - "Accept": "application/json", - "Accept-Encoding": "identity", - "Cookie": "", - } - space = _get_json("/spaces/" + quote(identity["space_id"], safe=""), headers, deadline) - if space is None or space.get("id") != identity["space_id"] or not _opaque_id(space.get("organization_id")): - return None - organization_id = space["organization_id"] - assert isinstance(organization_id, str) - project_ids: set[str] = set() - seen_cursors: set[str] = set() - params: dict[str, str | int] = { - "space_id": identity["space_id"], - "name": identity["project"], - "limit": _PAGE_LIMIT, - } - for _ in range(_MAX_PAGES): - page = _get_json("/projects", headers, deadline, params=params) - if page is None: - return None - projects, pagination = page.get("projects"), page.get("pagination") - if not isinstance(projects, list) or len(projects) > _PAGE_LIMIT or not isinstance(pagination, dict): - return None - for project in projects: - if ( - not isinstance(project, dict) - or not _opaque_id(project.get("id")) - or not isinstance(project.get("name"), str) - or project.get("space_id") != identity["space_id"] - ): - return None - if project["name"] == identity["project"]: - project_id = project["id"] - assert isinstance(project_id, str) - project_ids.add(project_id) - if len(project_ids) > 1: - return None - has_more = pagination.get("has_more") - if has_more is False: - if len(project_ids) != 1: - return None - return {**identity, "organization_id": organization_id, "project_id": next(iter(project_ids))} - cursor = pagination.get("next_cursor") - if ( - has_more is not True - or not isinstance(cursor, str) - or not cursor - or len(cursor) > 4096 - or cursor in seen_cursors - ): - return None - seen_cursors.add(cursor) - params = {**params, "cursor": cursor} - return None diff --git a/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/arize_phoenix_trace.py b/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/arize_phoenix_trace.py deleted file mode 100644 index d5ab4a570a3206..00000000000000 --- a/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/arize_phoenix_trace.py +++ /dev/null @@ -1,1593 +0,0 @@ -import json -import logging -import os -import re -import traceback -from collections.abc import Mapping, Sequence -from dataclasses import dataclass, field -from datetime import datetime, timedelta -from typing import Any, Protocol, cast, override -from urllib.parse import urlparse - -from openinference.semconv.trace import ( - MessageAttributes, - OpenInferenceMimeTypeValues, - OpenInferenceSpanKindValues, - SpanAttributes, - ToolCallAttributes, -) -from opentelemetry.context import Context -from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter as GrpcOTLPSpanExporter -from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter as HttpOTLPSpanExporter -from opentelemetry.sdk import trace as trace_sdk -from opentelemetry.sdk.resources import Resource -from opentelemetry.sdk.trace.export import SimpleSpanProcessor -from opentelemetry.semconv.attributes import exception_attributes -from opentelemetry.trace import Span, Status, StatusCode, get_current_span, set_span_in_context, use_span -from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator -from opentelemetry.util.types import AttributeValue -from sqlalchemy.orm import sessionmaker - -from core.ops.base_trace_instance import BaseTraceInstance -from core.ops.entities.trace_entity import ( - BaseTraceInfo, - DatasetRetrievalTraceInfo, - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - SuggestedQuestionTraceInfo, - ToolTraceInfo, - TraceTaskName, - WorkflowTraceInfo, -) -from core.ops.exceptions import PendingTraceParentContextError -from core.ops.utils import JSON_DICT_ADAPTER -from core.repositories import DifyCoreRepositoryFactory -from dify_trace_arize_phoenix.config import ArizeConfig, PhoenixConfig -from extensions.ext_database import db -from extensions.ext_redis import redis_client -from graphon.enums import WorkflowNodeExecutionStatus -from models.model import App, EndUser, MessageFile -from models.workflow import WorkflowNodeExecutionTriggeredFrom, WorkflowRun - -logger = logging.getLogger(__name__) - -# This parent-span carrier store is intentionally Phoenix-local for the current -# nested workflow tracing feature. If other trace providers need the same -# cross-task parent restoration behavior, move the storage and retry signaling -# behind a core trace coordination interface instead of duplicating it here. -_PHOENIX_PARENT_SPAN_CONTEXT_TTL_SECONDS = 300 -_TRACEPARENT_PATTERN = re.compile( - r"^(?P[0-9a-f]{2})-(?P[0-9a-f]{32})-(?P[0-9a-f]{16})-(?P[0-9a-f]{2})$" -) -_WRAPPER_INDEX_PATTERN = re.compile(r"^[A-Za-z0-9_.:-]+$") - - -def _phoenix_parent_span_redis_key(parent_node_execution_id: str) -> str: - """Build the Redis key that stores a restorable Phoenix parent span carrier.""" - return f"trace:phoenix:parent_span:{parent_node_execution_id}" - - -def _publish_parent_span_context(parent_node_execution_id: str, carrier: Mapping[str, str]) -> None: - """Persist a tracecontext carrier so nested workflow spans can restore the tool span parent.""" - redis_client.setex( - _phoenix_parent_span_redis_key(parent_node_execution_id), - _PHOENIX_PARENT_SPAN_CONTEXT_TTL_SECONDS, - safe_json_dumps(dict(carrier)), - ) - - -def _resolve_published_parent_span_context(parent_node_execution_id: str) -> dict[str, str]: - """Load a previously published tool-span carrier for nested workflow parenting.""" - raw_carrier = redis_client.get(_phoenix_parent_span_redis_key(parent_node_execution_id)) - if raw_carrier is None: - raise PendingTraceParentContextError(parent_node_execution_id) - - if isinstance(raw_carrier, bytes): - raw_carrier = raw_carrier.decode("utf-8") - - carrier = json.loads(raw_carrier) - if not isinstance(carrier, dict): - raise ValueError( - "Phoenix parent span context must be stored as a JSON object: " - f"parent_node_execution_id={parent_node_execution_id}" - ) - - normalized_carrier = {str(key): str(value) for key, value in carrier.items()} - if not normalized_carrier: - raise ValueError( - f"Phoenix parent span context payload is empty: parent_node_execution_id={parent_node_execution_id}" - ) - - traceparent = normalized_carrier.get("traceparent") - if not isinstance(traceparent, str): - raise ValueError( - "Phoenix parent span context payload is missing traceparent: " - f"parent_node_execution_id={parent_node_execution_id}" - ) - - traceparent_match = _TRACEPARENT_PATTERN.fullmatch(traceparent) - if traceparent_match is None: - raise ValueError( - "Phoenix parent span context payload has invalid traceparent format: " - f"parent_node_execution_id={parent_node_execution_id}" - ) - - if traceparent_match.group("version") == "ff": - raise ValueError( - "Phoenix parent span context payload has unsupported traceparent version: " - f"parent_node_execution_id={parent_node_execution_id}" - ) - - if traceparent_match.group("trace_id") == "0" * 32: - raise ValueError( - "Phoenix parent span context payload has zero trace_id in traceparent: " - f"parent_node_execution_id={parent_node_execution_id}" - ) - - if traceparent_match.group("span_id") == "0" * 16: - raise ValueError( - "Phoenix parent span context payload has zero span_id in traceparent: " - f"parent_node_execution_id={parent_node_execution_id}" - ) - - extracted_context = TraceContextTextMapPropagator().extract(carrier=normalized_carrier) - extracted_span_context = get_current_span(extracted_context).get_span_context() - if not extracted_span_context.is_valid or not extracted_span_context.is_remote: - raise ValueError( - "Phoenix parent span context payload could not be restored into a valid parent span: " - f"parent_node_execution_id={parent_node_execution_id}" - ) - - return normalized_carrier - - -def _app_uses_phoenix_provider(app_tracing_config: Mapping[str, Any] | None) -> bool: - if not app_tracing_config or not app_tracing_config.get("enabled"): - return False - return app_tracing_config.get("tracing_provider") in {"arize", "phoenix"} - - -def _parent_workflow_can_publish_span_context(parent_workflow_run_id: str) -> bool: - parent_run = db.session.query(WorkflowRun).where(WorkflowRun.id == parent_workflow_run_id).first() - if parent_run is None: - return True - - parent_app = db.session.query(App).where(App.id == parent_run.app_id).first() - if parent_app is None or not parent_app.tracing: - return False - - try: - app_tracing_config = json.loads(parent_app.tracing) - except (TypeError, json.JSONDecodeError): - return False - if not isinstance(app_tracing_config, Mapping): - return False - - return _app_uses_phoenix_provider(app_tracing_config) - - -def _resolve_workflow_parent_carrier( - parent_node_execution_id: str, parent_workflow_run_id: str | None -) -> dict[str, str] | None: - try: - return _resolve_published_parent_span_context(parent_node_execution_id) - except PendingTraceParentContextError: - if parent_workflow_run_id and not _parent_workflow_can_publish_span_context(parent_workflow_run_id): - logger.info( - "[Arize/Phoenix] Parent workflow cannot publish Phoenix span context; falling back to root span: " - "parent_workflow_run_id=%s parent_node_execution_id=%s", - parent_workflow_run_id, - parent_node_execution_id, - ) - return None - raise - - -def setup_tracer(arize_phoenix_config: ArizeConfig | PhoenixConfig) -> tuple[trace_sdk.Tracer, SimpleSpanProcessor]: - """Configure OpenTelemetry tracer with OTLP exporter for Arize/Phoenix.""" - try: - # Choose the appropriate exporter based on config type - exporter: GrpcOTLPSpanExporter | HttpOTLPSpanExporter - - # Inspect the provided endpoint to determine its structure - parsed = urlparse(arize_phoenix_config.endpoint) - base_endpoint = f"{parsed.scheme}://{parsed.netloc}" - path = parsed.path.rstrip("/") - - if isinstance(arize_phoenix_config, ArizeConfig): - arize_endpoint = f"{base_endpoint}/v1" - arize_headers = { - "api_key": arize_phoenix_config.api_key or "", - "space_id": arize_phoenix_config.space_id or "", - "authorization": f"Bearer {arize_phoenix_config.api_key or ''}", - } - exporter = GrpcOTLPSpanExporter( - endpoint=arize_endpoint, - headers=arize_headers, - timeout=30, - ) - else: - phoenix_endpoint = f"{base_endpoint}{path}/v1/traces" - # Keep headers nonempty so the SDK does not fall back to ambient - # OTEL exporter credentials for an anonymous Phoenix destination. - phoenix_headers = {"Content-Type": "application/x-protobuf"} - if arize_phoenix_config.api_key: - phoenix_headers.update( - { - "api_key": arize_phoenix_config.api_key, - "authorization": f"Bearer {arize_phoenix_config.api_key}", - } - ) - exporter = HttpOTLPSpanExporter( - endpoint=phoenix_endpoint, - headers=phoenix_headers, - timeout=30, - ) - - attributes = { - "openinference.project.name": arize_phoenix_config.project or "", - "model_id": arize_phoenix_config.project or "", - } - resource = Resource(attributes=attributes) - provider = trace_sdk.TracerProvider(resource=resource) - processor = SimpleSpanProcessor( - exporter, - ) - provider.add_span_processor(processor) - - # Create a named tracer instead of setting the global provider - tracer_name = f"arize_phoenix_tracer_{arize_phoenix_config.project}" - logger.info("[Arize/Phoenix] Created tracer with name: %s", tracer_name) - return cast(trace_sdk.Tracer, provider.get_tracer(tracer_name)), processor - except Exception as e: - logger.error("[Arize/Phoenix] Failed to setup the tracer: %s", str(e), exc_info=True) - raise - - -def datetime_to_nanos(dt: datetime | None) -> int: - """Convert datetime to nanoseconds since epoch for Arize/Phoenix.""" - if dt is None: - dt = datetime.now() - return int(dt.timestamp() * 1_000_000_000) - - -def error_to_string(error: Exception | str | None) -> str: - """Convert an error to a string with traceback information for Arize/Phoenix.""" - error_message = "Empty Stack Trace" - if error: - if isinstance(error, Exception): - string_stacktrace = "".join(traceback.format_exception(error)) - error_message = f"{error.__class__.__name__}: {error}\n\n{string_stacktrace}" - else: - error_message = str(error) - return error_message - - -def set_span_status(current_span: Span, error: Exception | str | None = None): - """Set the status of the current span based on the presence of an error for Arize/Phoenix.""" - if error: - error_string = error_to_string(error) - current_span.set_status(Status(StatusCode.ERROR, error_string)) - _record_exception_event(current_span, error) - else: - current_span.set_status(Status(StatusCode.OK)) - - -def safe_json_dumps(obj: Any) -> str: - """A convenience wrapper to ensure that any object can be safely encoded for Arize/Phoenix.""" - return json.dumps(obj, default=str, ensure_ascii=False) - - -def wrap_span_metadata(metadata, **kwargs): - """Add common metatada to all trace entity types for Arize/Phoenix.""" - metadata["created_from"] = "Dify" - metadata.update(kwargs) - return metadata - - -# Mapping from built-in node type strings to OpenInference span kinds. -# Node types not listed here default to CHAIN. -_NODE_TYPE_TO_SPAN_KIND: dict[str, OpenInferenceSpanKindValues] = { - "llm": OpenInferenceSpanKindValues.LLM, - "knowledge-retrieval": OpenInferenceSpanKindValues.RETRIEVER, - "tool": OpenInferenceSpanKindValues.TOOL, - "agent": OpenInferenceSpanKindValues.AGENT, -} - - -def _get_node_span_kind(node_type: str) -> OpenInferenceSpanKindValues: - """Return the OpenInference span kind for a given workflow node type. - - Covers every built-in node type string. Nodes that do not have a - specialised span kind (e.g. ``start``, ``end``, ``if-else``, - ``code``, ``loop``, ``iteration``, etc.) are mapped to ``CHAIN``. - """ - return _NODE_TYPE_TO_SPAN_KIND.get(node_type, OpenInferenceSpanKindValues.CHAIN) - - -def _metadata_trace_session_id(trace_info: BaseTraceInfo) -> str | None: - value = trace_info.metadata.get("trace_session_id") - return value if isinstance(value, str) and value else None - - -def _resolve_workflow_session_fallback(trace_info: WorkflowTraceInfo) -> str: - if trace_info.conversation_id: - return trace_info.conversation_id - - parent_workflow_run_id, _ = _resolve_workflow_parent_context(trace_info) - if parent_workflow_run_id: - return parent_workflow_run_id - - return trace_info.workflow_run_id - - -def _resolve_message_session_fallback(trace_info: MessageTraceInfo) -> str: - if trace_info.message_data is not None: - conversation_id = getattr(trace_info.message_data, "conversation_id", None) - if conversation_id: - return conversation_id - return "" - - -def _resolve_trace_session_id(trace_info: WorkflowTraceInfo | MessageTraceInfo) -> str: - trace_session_id = _metadata_trace_session_id(trace_info) - if trace_session_id: - return trace_session_id - if isinstance(trace_info, WorkflowTraceInfo): - return _resolve_workflow_session_fallback(trace_info) - return _resolve_message_session_fallback(trace_info) - - -def _resolve_workflow_session_id(trace_info: WorkflowTraceInfo) -> str: - """Resolve the workflow session ID for Phoenix workflow spans.""" - return _resolve_trace_session_id(trace_info) - - -def _resolve_workflow_parent_context(trace_info: BaseTraceInfo) -> tuple[str | None, str | None]: - """Expose the typed parent context already resolved on the trace info.""" - return trace_info.resolved_parent_context - - -def _resolve_workflow_root_trace_id(trace_info: WorkflowTraceInfo) -> str: - """Resolve the canonical root trace ID for Phoenix workflow spans.""" - trace_correlation_override, _ = _resolve_workflow_parent_context(trace_info) - return trace_correlation_override or trace_info.resolved_trace_id or trace_info.workflow_run_id - - -class _NodeExecutionIdentityLike(Protocol): - @property - def node_execution_id(self) -> str | None: ... - - @property - def node_id(self) -> str: ... - - @property - def predecessor_node_id(self) -> str | None: ... - - -class _NodeExecutionLike(_NodeExecutionIdentityLike, Protocol): - @property - def id(self) -> str: ... - - @property - def node_type(self) -> str: ... - - @property - def title(self) -> str | None: ... - - @property - def inputs(self) -> Mapping[str, Any] | None: ... - - @property - def process_data(self) -> Mapping[str, Any] | None: ... - - @property - def outputs(self) -> Mapping[str, Any] | None: ... - - @property - def status(self) -> WorkflowNodeExecutionStatus: ... - - @property - def error(self) -> str | None: ... - - @property - def elapsed_time(self) -> float | None: ... - - @property - def metadata(self) -> Mapping[Any, Any] | None: ... - - @property - def created_at(self) -> datetime | None: ... - - -_PHOENIX_STRUCTURED_NODE_TYPES = frozenset({"start", "end", "loop", "iteration"}) - - -def _resolve_workflow_span_name(trace_info: WorkflowTraceInfo) -> str: - """Resolve the Phoenix workflow span display name.""" - workflow_run_id = trace_info.workflow_run_id.strip() if trace_info.workflow_run_id else "" - if workflow_run_id: - return f"{TraceTaskName.WORKFLOW_TRACE.value}_{workflow_run_id}" - return TraceTaskName.WORKFLOW_TRACE.value - - -def _build_node_title_by_id(trace_info: WorkflowTraceInfo) -> dict[str, str]: - """Build an authoritative node-title index from the persisted workflow graph.""" - workflow_data = trace_info.workflow_data - workflow_graph = getattr(workflow_data, "graph_dict", None) - if not isinstance(workflow_graph, Mapping): - workflow_graph = workflow_data.get("graph") if isinstance(workflow_data, Mapping) else None - if not isinstance(workflow_graph, Mapping): - return {} - - graph_nodes = workflow_graph.get("nodes") - if not isinstance(graph_nodes, Sequence): - return {} - - node_title_by_id: dict[str, str] = {} - for graph_node in graph_nodes: - if not isinstance(graph_node, Mapping): - continue - node_id = graph_node.get("id") - node_data = graph_node.get("data") - if not isinstance(node_id, str) or not isinstance(node_data, Mapping): - continue - node_title = node_data.get("title") - if isinstance(node_title, str) and node_title.strip(): - node_title_by_id[node_id] = node_title.strip() - - return node_title_by_id - - -def _resolve_workflow_node_span_name( - node_execution: _NodeExecutionLike, - node_title_by_id: Mapping[str, str] | None = None, -) -> str: - """Resolve the Phoenix workflow node span display name.""" - node_type = str(node_execution.node_type or "") - graph_node_title = None - if node_title_by_id is not None and isinstance(node_execution.node_id, str): - graph_node_title = node_title_by_id.get(node_execution.node_id) - - node_title = graph_node_title or (node_execution.title.strip() if isinstance(node_execution.title, str) else "") - if node_title: - return f"{node_type}_{node_title}" - return node_type - - -def _get_node_execution_id(node_execution: _NodeExecutionIdentityLike) -> str: - """Return the stable execution identifier for a workflow node execution.""" - return str(getattr(node_execution, "id", None) or node_execution.node_execution_id) - - -def _build_execution_id_by_node_id(node_executions: Sequence[_NodeExecutionIdentityLike]) -> dict[str, str]: - """Index unique workflow graph node ids by execution id. - - This Phoenix-local hierarchy reconstruction intentionally drops ambiguous - node ids instead of guessing based on repository order. That keeps parent - selection deterministic until upstream tracing exposes explicit parent span - data for repeated executions. - """ - execution_id_by_node_id: dict[str, str] = {} - ambiguous_node_ids: set[str] = set() - - for node_execution in node_executions: - node_id = node_execution.node_id - if not isinstance(node_id, str): - continue - execution_id = _get_node_execution_id(node_execution) - - if node_id in ambiguous_node_ids: - continue - - existing_execution_id = execution_id_by_node_id.get(node_id) - if existing_execution_id is None: - execution_id_by_node_id[node_id] = execution_id - continue - - if existing_execution_id != execution_id: - ambiguous_node_ids.add(node_id) - execution_id_by_node_id.pop(node_id, None) - - return execution_id_by_node_id - - -@dataclass(frozen=True) -class _WrapperGroupKey: - wrapper_type: str - container_execution_id: str - index: str - - -@dataclass -class _WrapperGroup: - key: _WrapperGroupKey - container_execution_id: str - child_execution_ids: set[str] = field(default_factory=set) - start_time: datetime | None = None - end_time: datetime | None = None - has_error: bool = False - - -def _normalize_wrapper_index(value: Any) -> str | None: - """Normalize stable loop/iteration indexes used for synthetic wrapper spans.""" - 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 _execution_metadata_mapping(node_execution: object) -> Mapping[Any, Any]: - """Return execution metadata from repository models and test doubles.""" - execution_metadata = getattr(node_execution, "execution_metadata_dict", None) - if isinstance(execution_metadata, Mapping): - return execution_metadata - - execution_metadata = getattr(node_execution, "metadata", None) - if isinstance(execution_metadata, Mapping): - return execution_metadata - - return {} - - -def _node_finished_at(node_execution: object) -> datetime: - end_time = getattr(node_execution, "end_time", None) - if isinstance(end_time, datetime): - return end_time - created_at = getattr(node_execution, "created_at", None) or getattr(node_execution, "start_time", None) - if not isinstance(created_at, datetime): - created_at = datetime.now() - elapsed_time = getattr(node_execution, "elapsed_time", None) or 0.0 - return created_at + timedelta(seconds=elapsed_time) - - -def _metadata_or_attr(node_execution: object, node_metadata: Mapping[str, Any], key: str) -> Any: - value = node_metadata.get(key) - if value is not None: - return value - return getattr(node_execution, key, None) - - -def _node_execution_failed(node_execution: object) -> bool: - status = getattr(node_execution, "status", None) - return status in {WorkflowNodeExecutionStatus.FAILED, "failed"} - - -def _node_execution_handled_exception(node_execution: object) -> bool: - status = getattr(node_execution, "status", None) - return status in {WorkflowNodeExecutionStatus.EXCEPTION, "exception"} - - -def _record_exception_event(current_span: Span, error: Exception | str | None = None) -> None: - if not error: - return - - error_string = error_to_string(error) - if isinstance(error, Exception): - current_span.record_exception(error) - return - - exception_message = str(error) or repr(error) - attributes: dict[str, AttributeValue] = { - exception_attributes.EXCEPTION_TYPE: error.__class__.__name__, - exception_attributes.EXCEPTION_MESSAGE: exception_message, - exception_attributes.EXCEPTION_ESCAPED: False, - exception_attributes.EXCEPTION_STACKTRACE: error_string, - } - current_span.add_event(name="exception", attributes=attributes) - - -def _resolve_wrapper_group_key( - node_execution: _NodeExecutionIdentityLike, - node_metadata: Mapping[Any, Any], - execution_id_by_node_id: Mapping[str, str], -) -> _WrapperGroupKey | None: - for wrapper_type, container_key, index_key in ( - ("iteration", "iteration_id", "iteration_index"), - ("loop", "loop_id", "loop_index"), - ): - container_id = _metadata_or_attr(node_execution, node_metadata, container_key) - if not isinstance(container_id, str) or not container_id: - continue - - container_execution_id = execution_id_by_node_id.get(container_id) - if container_execution_id is None or container_execution_id == _get_node_execution_id(node_execution): - continue - - index = _normalize_wrapper_index(_metadata_or_attr(node_execution, node_metadata, index_key)) - if index is None: - continue - - return _WrapperGroupKey( - wrapper_type=wrapper_type, - container_execution_id=container_execution_id, - index=index, - ) - - return None - - -def _build_wrapper_groups( - node_executions: Sequence[_NodeExecutionIdentityLike], -) -> dict[_WrapperGroupKey, _WrapperGroup]: - """Group repeated loop/iteration body executions behind synthetic Phoenix spans.""" - execution_id_by_node_id = _build_execution_id_by_node_id(node_executions) - groups: dict[_WrapperGroupKey, _WrapperGroup] = {} - - for node_execution in node_executions: - node_metadata = _execution_metadata_mapping(node_execution) - group_key = _resolve_wrapper_group_key(node_execution, node_metadata, execution_id_by_node_id) - if group_key is None: - continue - - group = groups.setdefault( - group_key, - _WrapperGroup(key=group_key, container_execution_id=group_key.container_execution_id), - ) - execution_id = _get_node_execution_id(node_execution) - group.child_execution_ids.add(execution_id) - - created_at = getattr(node_execution, "created_at", None) or getattr(node_execution, "start_time", None) - if not isinstance(created_at, datetime): - created_at = datetime.now() - finished_at = _node_finished_at(node_execution) - group.start_time = created_at if group.start_time is None else min(group.start_time, created_at) - group.end_time = finished_at if group.end_time is None else max(group.end_time, finished_at) - group.has_error = group.has_error or _node_execution_failed(node_execution) - - return groups - - -def _build_graph_parent_index(node_executions: Sequence[_NodeExecutionIdentityLike]) -> dict[str, str]: - """Build an execution-id parent index from predecessor node ids.""" - execution_id_by_node_id = _build_execution_id_by_node_id(node_executions) - graph_parent_index: dict[str, str] = {} - - for node_execution in node_executions: - predecessor_node_id = node_execution.predecessor_node_id - if not isinstance(predecessor_node_id, str): - continue - - predecessor_execution_id = execution_id_by_node_id.get(predecessor_node_id) - if predecessor_execution_id is not None: - execution_id = _get_node_execution_id(node_execution) - graph_parent_index[execution_id] = predecessor_execution_id - - return graph_parent_index - - -def _resolve_structured_parent_execution_id( - node_execution: object, execution_id_by_node_id: Mapping[str, str] -) -> str | None: - """Resolve Phoenix-local structured parents from loop/iteration node ids. - - Any execution carrying ``iteration_id`` or ``loop_id`` belongs to an - enclosing structured node. When predecessor node ids are ambiguous because - the graph node repeats inside that structure, Phoenix can still keep the - child span under the enclosing loop/iteration span without relying on - execution-order heuristics. - """ - execution_metadata = getattr(node_execution, "execution_metadata_dict", None) - if not isinstance(execution_metadata, Mapping): - execution_metadata = getattr(node_execution, "metadata", None) - if not isinstance(execution_metadata, Mapping): - execution_metadata = {} - - for enclosing_node_id in ( - getattr(node_execution, "iteration_id", None), - getattr(node_execution, "loop_id", None), - execution_metadata.get("iteration_id"), - execution_metadata.get("loop_id"), - ): - if not isinstance(enclosing_node_id, str): - continue - - enclosing_execution_id = execution_id_by_node_id.get(enclosing_node_id) - if enclosing_execution_id is not None: - return enclosing_execution_id - - return None - - -def _resolve_node_parent( - execution_id: str, - predecessor_execution_id: str | None, - structured_parent_execution_id: str | None, - span_by_execution_id: Mapping[str, Span], - graph_parent_index: Mapping[str, str], - workflow_span: Span, -) -> Span: - """Resolve the parent span for a workflow node execution.""" - if predecessor_execution_id is not None: - predecessor_span = span_by_execution_id.get(predecessor_execution_id) - if predecessor_span is not None: - return predecessor_span - - graph_parent_execution_id = graph_parent_index.get(execution_id) - if graph_parent_execution_id is not None: - graph_parent_span = span_by_execution_id.get(graph_parent_execution_id) - if graph_parent_span is not None: - return graph_parent_span - - if structured_parent_execution_id is not None: - structured_parent_span = span_by_execution_id.get(structured_parent_execution_id) - if structured_parent_span is not None: - return structured_parent_span - - return workflow_span - - -class ArizePhoenixDataTrace(BaseTraceInstance): - def __init__( - self, - arize_phoenix_config: ArizeConfig | PhoenixConfig, - ): - super().__init__(arize_phoenix_config) - self.arize_phoenix_config = arize_phoenix_config - self.tracer, self.processor = setup_tracer(arize_phoenix_config) - self.project = arize_phoenix_config.project - self.file_base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001") - self.propagator = TraceContextTextMapPropagator() - self.dify_trace_ids: set[str] = set() - self.root_span_carriers: dict[str, dict[str, str]] = {} - self.carrier: dict[str, str] = {} - - @override - def trace(self, trace_info: BaseTraceInfo): - logger.info("[Arize/Phoenix] Trace Entity Info: %s", trace_info) - logger.info("[Arize/Phoenix] Trace Entity Type: %s", type(trace_info)) - try: - match trace_info: - case WorkflowTraceInfo(): - self.workflow_trace(trace_info) - case MessageTraceInfo(): - self.message_trace(trace_info) - case ModerationTraceInfo(): - self.moderation_trace(trace_info) - case SuggestedQuestionTraceInfo(): - self.suggested_question_trace(trace_info) - case DatasetRetrievalTraceInfo(): - self.dataset_retrieval_trace(trace_info) - case ToolTraceInfo(): - self.tool_trace(trace_info) - case GenerateNameTraceInfo(): - self.generate_name_trace(trace_info) - case _: - pass - - except Exception as e: - logger.error("[Arize/Phoenix] Trace Entity Error: %s", str(e), exc_info=True) - raise - - def workflow_trace(self, trace_info: WorkflowTraceInfo): - file_list = trace_info.file_list if isinstance(trace_info.file_list, list) else [] - - metadata = wrap_span_metadata( - trace_info.metadata, - trace_id=trace_info.trace_id or "", - message_id=trace_info.message_id or "", - status=trace_info.workflow_run_status or "", - status_message=trace_info.error or "", - level="ERROR" if trace_info.error else "DEFAULT", - trace_entity_type="workflow", - conversation_id=trace_info.conversation_id or "", - workflow_app_log_id=trace_info.workflow_app_log_id or "", - workflow_id=trace_info.workflow_id or "", - tenant_id=trace_info.tenant_id or "", - workflow_run_id=trace_info.workflow_run_id or "", - workflow_run_elapsed_time=trace_info.workflow_run_elapsed_time or 0, - workflow_run_version=trace_info.workflow_run_version or "", - total_tokens=trace_info.total_tokens or 0, - file_list=safe_json_dumps(file_list), - query=trace_info.query or "", - ) - workflow_session_id = _resolve_trace_session_id(trace_info) - parent_workflow_run_id, parent_node_execution_id = _resolve_workflow_parent_context(trace_info) - logger.info( - "[Arize/Phoenix] Workflow session resolution: workflow_run_id=%s conversation_id=%s " - "parent_workflow_run_id=%s parent_node_execution_id=%s resolved_session_id=%s", - trace_info.workflow_run_id, - trace_info.conversation_id, - parent_workflow_run_id, - parent_node_execution_id, - workflow_session_id, - ) - - workflow_parent_carrier: dict[str, str] | None = None - if parent_node_execution_id: - workflow_parent_carrier = _resolve_workflow_parent_carrier(parent_node_execution_id, parent_workflow_run_id) - - if workflow_parent_carrier is None: - root_trace_id = _resolve_workflow_root_trace_id(trace_info) - workflow_root_span_name: str | None = trace_info.workflow_run_id - if not isinstance(workflow_root_span_name, str) or not workflow_root_span_name.strip(): - workflow_root_span_name = None - - root_span_kwargs: dict[str, Any] = { - "root_span_name": workflow_root_span_name, - "root_span_attributes": { - SpanAttributes.INPUT_VALUE: safe_json_dumps(trace_info.workflow_run_inputs), - SpanAttributes.INPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON.value, - SpanAttributes.OUTPUT_VALUE: safe_json_dumps(trace_info.workflow_run_outputs), - SpanAttributes.OUTPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON.value, - SpanAttributes.SESSION_ID: workflow_session_id or "", - }, - } - if trace_info.error: - root_span_kwargs["root_span_error"] = trace_info.error - - workflow_parent_carrier = self.ensure_root_span(root_trace_id, **root_span_kwargs) - - workflow_span_context = self.propagator.extract(carrier=workflow_parent_carrier) - - workflow_span = self.tracer.start_span( - name=_resolve_workflow_span_name(trace_info), - attributes={ - SpanAttributes.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.CHAIN.value, - SpanAttributes.INPUT_VALUE: safe_json_dumps(trace_info.workflow_run_inputs), - SpanAttributes.INPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON.value, - SpanAttributes.OUTPUT_VALUE: safe_json_dumps(trace_info.workflow_run_outputs), - SpanAttributes.OUTPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON.value, - SpanAttributes.METADATA: safe_json_dumps(metadata), - SpanAttributes.SESSION_ID: workflow_session_id or "", - }, - start_time=datetime_to_nanos(trace_info.start_time), - context=workflow_span_context, - ) - - # Through workflow_run_id, get all_nodes_execution using repository - session_factory = sessionmaker(bind=db.engine) - - # Find the app's creator account - app_id = trace_info.metadata.get("app_id") - if not app_id: - raise ValueError("No app_id found in trace_info metadata") - - service_account = self.get_service_account_with_tenant(app_id) - - workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( - session_factory=session_factory, - tenant_id=trace_info.tenant_id, - user=service_account, - app_id=app_id, - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - - # Get all executions for this workflow run - workflow_node_executions = workflow_node_execution_repository.get_by_workflow_execution( - workflow_execution_id=trace_info.workflow_run_id - ) - node_title_by_id = _build_node_title_by_id(trace_info) - execution_id_by_node_id = _build_execution_id_by_node_id(workflow_node_executions) - graph_parent_index = _build_graph_parent_index(workflow_node_executions) - node_execution_by_execution_id = { - _get_node_execution_id(node_execution): node_execution for node_execution in workflow_node_executions - } - span_by_execution_id: dict[str, Span] = {} - emitting_execution_ids: set[str] = set() - wrapper_groups = _build_wrapper_groups(workflow_node_executions) - wrapper_span_by_key: dict[_WrapperGroupKey, Span] = {} - finalized_wrapper_keys: set[_WrapperGroupKey] = set() - wrapper_key_by_child_execution_id: dict[str, _WrapperGroupKey] = {} - for group_key, group in wrapper_groups.items(): - for child_execution_id in group.child_execution_ids: - wrapper_key_by_child_execution_id[child_execution_id] = group_key - - workflow_span_error: Exception | str | None = trace_info.error - try: - - def emit_wrapper_span(group_key: _WrapperGroupKey) -> Span | None: - existing_span = wrapper_span_by_key.get(group_key) - if existing_span is not None: - return existing_span - - group = wrapper_groups.get(group_key) - if group is None: - return None - - if group.container_execution_id not in span_by_execution_id: - parent_node_execution = node_execution_by_execution_id.get(group.container_execution_id) - if parent_node_execution is not None: - emit_node_span(parent_node_execution) - - container_span = span_by_execution_id.get(group.container_execution_id) - if container_span is None: - return None - - metadata = { - "synthetic": True, - "wrapper_type": group_key.wrapper_type, - "wrapper_index": group_key.index, - "container_execution_id": group.container_execution_id, - } - wrapper_span = self.tracer.start_span( - name=f"{group_key.wrapper_type}[{group_key.index}]", - attributes={ - SpanAttributes.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.CHAIN.value, - SpanAttributes.METADATA: safe_json_dumps(metadata), - SpanAttributes.SESSION_ID: workflow_session_id or "", - "dify.wrapper.synthetic": True, - "dify.wrapper.type": group_key.wrapper_type, - "dify.wrapper.index": group_key.index, - "dify.wrapper.container_execution_id": group.container_execution_id, - }, - start_time=datetime_to_nanos(group.start_time), - context=set_span_in_context(container_span), - ) - wrapper_span_by_key[group_key] = wrapper_span - return wrapper_span - - def finalize_wrapper_spans() -> None: - for group_key, wrapper_span in list(wrapper_span_by_key.items()): - if group_key in finalized_wrapper_keys: - continue - group = wrapper_groups[group_key] - set_span_status(wrapper_span, "wrapper child failed" if group.has_error else None) - wrapper_span.end(end_time=datetime_to_nanos(group.end_time)) - finalized_wrapper_keys.add(group_key) - - def emit_node_span(node_execution: _NodeExecutionLike) -> Span: - execution_id = _get_node_execution_id(node_execution) - existing_span = span_by_execution_id.get(execution_id) - if existing_span is not None: - return existing_span - - graph_parent_execution_id = graph_parent_index.get(execution_id) - structured_parent_execution_id = _resolve_structured_parent_execution_id( - node_execution, execution_id_by_node_id - ) - - if execution_id not in emitting_execution_ids: - emitting_execution_ids.add(execution_id) - try: - for parent_execution_id in (graph_parent_execution_id, structured_parent_execution_id): - if parent_execution_id is None or parent_execution_id == execution_id: - continue - if parent_execution_id in span_by_execution_id: - continue - parent_node_execution = node_execution_by_execution_id.get(parent_execution_id) - if parent_node_execution is not None: - emit_node_span(parent_node_execution) - finally: - emitting_execution_ids.discard(execution_id) - - tenant_id = trace_info.tenant_id # Use from trace_info instead - app_id = trace_info.metadata.get("app_id") # Use from trace_info instead - inputs_value = node_execution.inputs or {} - outputs_value = node_execution.outputs or {} - - created_at = node_execution.created_at or datetime.now() - finished_at = _node_finished_at(node_execution) - - process_data = node_execution.process_data or {} - execution_metadata = _execution_metadata_mapping(node_execution) - node_metadata = {str(k): v for k, v in execution_metadata.items()} - - node_metadata.update( - { - "node_id": node_execution.id, - "node_type": node_execution.node_type, - "node_status": node_execution.status, - "tenant_id": tenant_id, - "app_id": app_id, - "app_name": node_execution.title, - "status": node_execution.status, - "status_message": node_execution.error or "", - "level": "ERROR" if node_execution.status == WorkflowNodeExecutionStatus.FAILED else "DEFAULT", - "loop_id": _metadata_or_attr(node_execution, node_metadata, "loop_id"), - "loop_index": _metadata_or_attr(node_execution, node_metadata, "loop_index"), - "iteration_id": _metadata_or_attr(node_execution, node_metadata, "iteration_id"), - "iteration_index": _metadata_or_attr(node_execution, node_metadata, "iteration_index"), - } - ) - - # Determine the correct span kind based on node type - span_kind = _get_node_span_kind(node_execution.node_type) - if node_execution.node_type == "llm": - provider = process_data.get("model_provider") - model = process_data.get("model_name") - if provider: - node_metadata["ls_provider"] = provider - if model: - node_metadata["ls_model_name"] = model - - usage_data = ( - process_data.get("usage", {}) if "usage" in process_data else outputs_value.get("usage", {}) - ) - if usage_data: - node_metadata["total_tokens"] = usage_data.get("total_tokens", 0) - node_metadata["prompt_tokens"] = usage_data.get("prompt_tokens", 0) - node_metadata["completion_tokens"] = usage_data.get("completion_tokens", 0) - - wrapper_group_key = wrapper_key_by_child_execution_id.get(execution_id) - wrapper_parent_span = emit_wrapper_span(wrapper_group_key) if wrapper_group_key else None - parent_span = wrapper_parent_span or _resolve_node_parent( - execution_id=execution_id, - predecessor_execution_id=None, - structured_parent_execution_id=structured_parent_execution_id, - span_by_execution_id=span_by_execution_id, - graph_parent_index=graph_parent_index, - workflow_span=workflow_span, - ) - workflow_span_context = set_span_in_context(parent_span) - loop_index = node_metadata.get("loop_index") - iteration_index = node_metadata.get("iteration_index") - node_span = self.tracer.start_span( - name=_resolve_workflow_node_span_name(node_execution, node_title_by_id), - attributes={ - SpanAttributes.OPENINFERENCE_SPAN_KIND: span_kind.value, - SpanAttributes.INPUT_VALUE: safe_json_dumps(inputs_value), - SpanAttributes.INPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON.value, - SpanAttributes.OUTPUT_VALUE: safe_json_dumps(outputs_value), - SpanAttributes.OUTPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON.value, - SpanAttributes.METADATA: safe_json_dumps(node_metadata), - SpanAttributes.SESSION_ID: workflow_session_id or "", - "dify.node.execution_id": execution_id, - "dify.node.loop_id": cast(AttributeValue, node_metadata.get("loop_id") or ""), - "dify.node.loop_index": cast(AttributeValue, loop_index if loop_index is not None else ""), - "dify.node.iteration_id": cast(AttributeValue, node_metadata.get("iteration_id") or ""), - "dify.node.iteration_index": cast( - AttributeValue, iteration_index if iteration_index is not None else "" - ), - }, - start_time=datetime_to_nanos(created_at), - context=workflow_span_context, - ) - span_by_execution_id[execution_id] = node_span - node_span_error: Exception | str | None = None - try: - if node_execution.node_type == "tool": - parent_span_carrier: dict[str, str] = {} - with use_span(node_span, end_on_exit=False): - self.propagator.inject(carrier=parent_span_carrier) - _publish_parent_span_context(execution_id, parent_span_carrier) - - if node_execution.node_type == "llm": - llm_attributes: dict[str, Any] = { - SpanAttributes.INPUT_VALUE: json.dumps(process_data.get("prompts", []), ensure_ascii=False), - } - provider = process_data.get("model_provider") - model = process_data.get("model_name") - if provider: - llm_attributes[SpanAttributes.LLM_PROVIDER] = provider - if model: - llm_attributes[SpanAttributes.LLM_MODEL_NAME] = model - usage_data = ( - process_data.get("usage", {}) if "usage" in process_data else outputs_value.get("usage", {}) - ) - if usage_data: - llm_attributes[SpanAttributes.LLM_TOKEN_COUNT_TOTAL] = usage_data.get("total_tokens", 0) - llm_attributes[SpanAttributes.LLM_TOKEN_COUNT_PROMPT] = usage_data.get("prompt_tokens", 0) - llm_attributes[SpanAttributes.LLM_TOKEN_COUNT_COMPLETION] = usage_data.get( - "completion_tokens", 0 - ) - llm_attributes.update(self._construct_llm_attributes(process_data.get("prompts", []))) - node_span.set_attributes(llm_attributes) - except Exception as e: - node_span_error = e - raise - finally: - if node_span_error is not None: - set_span_status(node_span, node_span_error) - elif _node_execution_failed(node_execution) or _node_execution_handled_exception(node_execution): - set_span_status(node_span, node_execution.error) - else: - set_span_status(node_span) - node_span.end(end_time=datetime_to_nanos(finished_at)) - return node_span - - try: - for node_execution in workflow_node_executions: - emit_node_span(node_execution) - finally: - finalize_wrapper_spans() - except Exception as e: - workflow_span_error = e - raise - finally: - set_span_status(workflow_span, workflow_span_error) - workflow_span.end(end_time=datetime_to_nanos(trace_info.end_time)) - - def message_trace(self, trace_info: MessageTraceInfo): - if trace_info.message_data is None: - logger.warning("[Arize/Phoenix] Message data is None, skipping message trace.") - return - - file_list = trace_info.file_list if isinstance(trace_info.file_list, list) else [] - message_file_data: MessageFile | None = trace_info.message_file_data - - if message_file_data is not None: - file_url = f"{self.file_base_url}/{message_file_data.url}" if message_file_data else "" - file_list.append(file_url) - - metadata = wrap_span_metadata( - trace_info.metadata, - trace_id=trace_info.trace_id or "", - message_id=trace_info.message_id or "", - status=trace_info.message_data.status or "", - status_message=trace_info.error or "", - level="ERROR" if trace_info.error else "DEFAULT", - trace_entity_type="message", - conversation_model=trace_info.conversation_model or "", - message_tokens=trace_info.message_tokens or 0, - answer_tokens=trace_info.answer_tokens or 0, - total_tokens=trace_info.total_tokens or 0, - conversation_mode=trace_info.conversation_mode or "", - gen_ai_server_time_to_first_token=trace_info.gen_ai_server_time_to_first_token or 0, - llm_streaming_time_to_generate=trace_info.llm_streaming_time_to_generate or 0, - is_streaming_request=trace_info.is_streaming_request or False, - user_id=trace_info.message_data.from_account_id or "", - file_list=safe_json_dumps(file_list), - model_provider=trace_info.message_data.model_provider or "", - model_id=trace_info.message_data.model_id or "", - ) - message_session_id = _resolve_trace_session_id(trace_info) - - # Add end user data if available - if trace_info.message_data.from_end_user_id: - end_user_data: EndUser | None = db.session.get(EndUser, trace_info.message_data.from_end_user_id) - if end_user_data is not None: - metadata["end_user_id"] = end_user_data.session_id - - attributes = { - SpanAttributes.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.CHAIN.value, - SpanAttributes.INPUT_VALUE: trace_info.message_data.query, - SpanAttributes.INPUT_MIME_TYPE: OpenInferenceMimeTypeValues.TEXT.value, - SpanAttributes.OUTPUT_VALUE: trace_info.message_data.answer, - SpanAttributes.OUTPUT_MIME_TYPE: OpenInferenceMimeTypeValues.TEXT.value, - SpanAttributes.METADATA: safe_json_dumps(metadata), - SpanAttributes.SESSION_ID: message_session_id or "", - } - - dify_trace_id = trace_info.trace_id or trace_info.message_id - self.ensure_root_span(dify_trace_id) - root_span_context = self.propagator.extract(carrier=self.carrier) - - message_span = self.tracer.start_span( - name=TraceTaskName.MESSAGE_TRACE.value, - attributes=attributes, - start_time=datetime_to_nanos(trace_info.start_time), - context=root_span_context, - ) - - try: - # Convert outputs to string based on type - outputs_mime_type = OpenInferenceMimeTypeValues.TEXT.value - match trace_info.outputs: - case dict() | list(): - outputs_str = safe_json_dumps(trace_info.outputs) - outputs_mime_type = OpenInferenceMimeTypeValues.JSON.value - case str(): - outputs_str = trace_info.outputs - case _: - outputs_str = str(trace_info.outputs) - - llm_attributes: dict[str, Any] = { - SpanAttributes.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.LLM.value, - SpanAttributes.INPUT_VALUE: safe_json_dumps(trace_info.inputs), - SpanAttributes.INPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON.value, - SpanAttributes.OUTPUT_VALUE: outputs_str, - SpanAttributes.OUTPUT_MIME_TYPE: outputs_mime_type, - SpanAttributes.METADATA: safe_json_dumps(metadata), - SpanAttributes.SESSION_ID: message_session_id or "", - } - llm_attributes.update(self._construct_llm_attributes(trace_info.inputs)) - if trace_info.total_tokens is not None and trace_info.total_tokens > 0: - llm_attributes[SpanAttributes.LLM_TOKEN_COUNT_TOTAL] = trace_info.total_tokens - if trace_info.message_tokens is not None and trace_info.message_tokens > 0: - llm_attributes[SpanAttributes.LLM_TOKEN_COUNT_PROMPT] = trace_info.message_tokens - if trace_info.answer_tokens is not None and trace_info.answer_tokens > 0: - llm_attributes[SpanAttributes.LLM_TOKEN_COUNT_COMPLETION] = trace_info.answer_tokens - - if trace_info.message_data.model_id is not None: - llm_attributes[SpanAttributes.LLM_MODEL_NAME] = trace_info.message_data.model_id - if trace_info.message_data.model_provider is not None: - llm_attributes[SpanAttributes.LLM_PROVIDER] = trace_info.message_data.model_provider - - if trace_info.message_data and trace_info.message_data.message_metadata: - metadata_dict = JSON_DICT_ADAPTER.validate_json(trace_info.message_data.message_metadata) - if model_params := metadata_dict.get("model_parameters"): - llm_attributes[SpanAttributes.LLM_INVOCATION_PARAMETERS] = json.dumps(model_params) - - message_span_context = set_span_in_context(message_span) - llm_span = self.tracer.start_span( - name="llm", - attributes=llm_attributes, - start_time=datetime_to_nanos(trace_info.start_time), - context=message_span_context, - ) - - try: - if trace_info.message_data.error: - set_span_status(llm_span, trace_info.message_data.error) - else: - set_span_status(llm_span) - finally: - llm_span.end(end_time=datetime_to_nanos(trace_info.end_time)) - finally: - if trace_info.error: - set_span_status(message_span, trace_info.error) - else: - set_span_status(message_span) - message_span.end(end_time=datetime_to_nanos(trace_info.end_time)) - - def moderation_trace(self, trace_info: ModerationTraceInfo): - if trace_info.message_data is None: - logger.warning("[Arize/Phoenix] Message data is None, skipping moderation trace.") - return - - metadata = wrap_span_metadata( - trace_info.metadata, - trace_id=trace_info.trace_id or "", - message_id=trace_info.message_id or "", - status=trace_info.message_data.status or "", - status_message=trace_info.message_data.error or "", - level="ERROR" if trace_info.message_data.error else "DEFAULT", - trace_entity_type="moderation", - model_provider=trace_info.message_data.model_provider or "", - model_id=trace_info.message_data.model_id or "", - ) - - dify_trace_id = trace_info.trace_id or trace_info.message_id - self.ensure_root_span(dify_trace_id) - root_span_context = self.propagator.extract(carrier=self.carrier) - - span = self.tracer.start_span( - name=TraceTaskName.MODERATION_TRACE.value, - attributes={ - SpanAttributes.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.TOOL.value, - SpanAttributes.INPUT_VALUE: safe_json_dumps(trace_info.inputs), - SpanAttributes.INPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON.value, - SpanAttributes.OUTPUT_VALUE: safe_json_dumps( - { - "flagged": trace_info.flagged, - "action": trace_info.action, - "preset_response": trace_info.preset_response, - "query": trace_info.query, - } - ), - SpanAttributes.OUTPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON.value, - SpanAttributes.METADATA: safe_json_dumps(metadata), - }, - start_time=datetime_to_nanos(trace_info.start_time), - context=root_span_context, - ) - - try: - if trace_info.message_data.error: - set_span_status(span, trace_info.message_data.error) - else: - set_span_status(span) - finally: - span.end(end_time=datetime_to_nanos(trace_info.end_time)) - - def suggested_question_trace(self, trace_info: SuggestedQuestionTraceInfo): - if trace_info.message_data is None: - logger.warning("[Arize/Phoenix] Message data is None, skipping suggested question trace.") - return - - start_time = trace_info.start_time or trace_info.message_data.created_at - end_time = trace_info.end_time or trace_info.message_data.updated_at - - metadata = wrap_span_metadata( - trace_info.metadata, - trace_id=trace_info.trace_id or "", - message_id=trace_info.message_id or "", - status=trace_info.status or "", - status_message=trace_info.status_message or "", - level=trace_info.level or "", - trace_entity_type="suggested_question", - total_tokens=trace_info.total_tokens or 0, - from_account_id=trace_info.from_account_id or "", - agent_based=trace_info.agent_based or False, - from_source=trace_info.from_source or "", - model_provider=trace_info.model_provider or "", - model_id=trace_info.model_id or "", - workflow_run_id=trace_info.workflow_run_id or "", - ) - - dify_trace_id = trace_info.trace_id or trace_info.message_id - self.ensure_root_span(dify_trace_id) - root_span_context = self.propagator.extract(carrier=self.carrier) - - span = self.tracer.start_span( - name=TraceTaskName.SUGGESTED_QUESTION_TRACE.value, - attributes={ - SpanAttributes.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.TOOL.value, - SpanAttributes.INPUT_VALUE: safe_json_dumps(trace_info.inputs), - SpanAttributes.INPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON.value, - SpanAttributes.OUTPUT_VALUE: safe_json_dumps(trace_info.suggested_question), - SpanAttributes.OUTPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON.value, - SpanAttributes.METADATA: safe_json_dumps(metadata), - }, - start_time=datetime_to_nanos(start_time), - context=root_span_context, - ) - - try: - if trace_info.error: - set_span_status(span, trace_info.error) - else: - set_span_status(span) - finally: - span.end(end_time=datetime_to_nanos(end_time)) - - def dataset_retrieval_trace(self, trace_info: DatasetRetrievalTraceInfo): - if trace_info.message_data is None: - logger.warning("[Arize/Phoenix] Message data is None, skipping dataset retrieval trace.") - return - - start_time = trace_info.start_time or trace_info.message_data.created_at - end_time = trace_info.end_time or trace_info.message_data.updated_at - - metadata = wrap_span_metadata( - trace_info.metadata, - trace_id=trace_info.trace_id or "", - message_id=trace_info.message_id or "", - status=trace_info.message_data.status or "", - status_message=trace_info.error or "", - level="ERROR" if trace_info.error else "DEFAULT", - trace_entity_type="dataset_retrieval", - model_provider=trace_info.message_data.model_provider or "", - model_id=trace_info.message_data.model_id or "", - ) - - dify_trace_id = trace_info.trace_id or trace_info.message_id - self.ensure_root_span(dify_trace_id) - root_span_context = self.propagator.extract(carrier=self.carrier) - - span = self.tracer.start_span( - name=TraceTaskName.DATASET_RETRIEVAL_TRACE.value, - attributes={ - SpanAttributes.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.RETRIEVER.value, - SpanAttributes.INPUT_VALUE: safe_json_dumps(trace_info.inputs), - SpanAttributes.INPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON.value, - SpanAttributes.OUTPUT_VALUE: safe_json_dumps({"documents": trace_info.documents}), - SpanAttributes.OUTPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON.value, - SpanAttributes.METADATA: safe_json_dumps(metadata), - }, - start_time=datetime_to_nanos(start_time), - context=root_span_context, - ) - - try: - if trace_info.error: - set_span_status(span, trace_info.error) - else: - set_span_status(span) - finally: - span.end(end_time=datetime_to_nanos(end_time)) - - def tool_trace(self, trace_info: ToolTraceInfo): - if trace_info.message_data is None: - logger.warning("[Arize/Phoenix] Message data is None, skipping tool trace.") - return - - metadata = wrap_span_metadata( - trace_info.metadata, - trace_id=trace_info.trace_id or "", - message_id=trace_info.message_id or "", - status=trace_info.message_data.status or "", - status_message=trace_info.error or "", - level="ERROR" if trace_info.error else "DEFAULT", - trace_entity_type="tool", - tool_config=safe_json_dumps(trace_info.tool_config), - time_cost=trace_info.time_cost or 0, - file_url=trace_info.file_url or "", - ) - - dify_trace_id = trace_info.trace_id or trace_info.message_id - self.ensure_root_span(dify_trace_id) - root_span_context = self.propagator.extract(carrier=self.carrier) - - span = self.tracer.start_span( - name=trace_info.tool_name, - attributes={ - SpanAttributes.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.TOOL.value, - SpanAttributes.INPUT_VALUE: safe_json_dumps(trace_info.tool_inputs), - SpanAttributes.INPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON.value, - SpanAttributes.OUTPUT_VALUE: trace_info.tool_outputs, - SpanAttributes.OUTPUT_MIME_TYPE: OpenInferenceMimeTypeValues.TEXT.value, - SpanAttributes.METADATA: safe_json_dumps(metadata), - SpanAttributes.TOOL_NAME: trace_info.tool_name, - SpanAttributes.TOOL_PARAMETERS: safe_json_dumps(trace_info.tool_parameters), - }, - start_time=datetime_to_nanos(trace_info.start_time), - context=root_span_context, - ) - - try: - if trace_info.error: - set_span_status(span, trace_info.error) - else: - set_span_status(span) - finally: - span.end(end_time=datetime_to_nanos(trace_info.end_time)) - - def generate_name_trace(self, trace_info: GenerateNameTraceInfo): - if trace_info.message_data is None: - logger.warning("[Arize/Phoenix] Message data is None, skipping generate name trace.") - return - - metadata = wrap_span_metadata( - trace_info.metadata, - trace_id=trace_info.trace_id or "", - message_id=trace_info.message_id or "", - status=trace_info.message_data.status or "", - status_message=trace_info.message_data.error or "", - level="ERROR" if trace_info.message_data.error else "DEFAULT", - trace_entity_type="generate_name", - model_provider=trace_info.message_data.model_provider or "", - model_id=trace_info.message_data.model_id or "", - conversation_id=trace_info.conversation_id or "", - tenant_id=trace_info.tenant_id, - ) - - dify_trace_id = trace_info.trace_id or trace_info.message_id or trace_info.conversation_id - self.ensure_root_span(dify_trace_id) - root_span_context = self.propagator.extract(carrier=self.carrier) - - span = self.tracer.start_span( - name=TraceTaskName.GENERATE_NAME_TRACE.value, - attributes={ - SpanAttributes.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.CHAIN.value, - SpanAttributes.INPUT_VALUE: safe_json_dumps(trace_info.inputs), - SpanAttributes.INPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON.value, - SpanAttributes.OUTPUT_VALUE: safe_json_dumps(trace_info.outputs), - SpanAttributes.OUTPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON.value, - SpanAttributes.METADATA: safe_json_dumps(metadata), - SpanAttributes.SESSION_ID: trace_info.conversation_id or "", - }, - start_time=datetime_to_nanos(trace_info.start_time), - context=root_span_context, - ) - - try: - if trace_info.message_data.error: - set_span_status(span, trace_info.message_data.error) - else: - set_span_status(span) - finally: - span.end(end_time=datetime_to_nanos(trace_info.end_time)) - - def ensure_root_span( - self, - dify_trace_id: str | None, - *, - root_span_name: str | None = None, - root_span_error: Exception | str | None = None, - root_span_attributes: Mapping[str, AttributeValue] | None = None, - ): - """Ensure a unique root span exists for the given Dify trace ID.""" - trace_key = str(dify_trace_id) - if trace_key not in self.dify_trace_ids: - carrier: dict[str, str] = {} - - span_name = root_span_name.strip() if isinstance(root_span_name, str) and root_span_name.strip() else "Dify" - root_span_attributes_dict: dict[str, AttributeValue] = { - SpanAttributes.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.CHAIN.value, - "dify_project_name": str(self.project), - "dify_trace_id": trace_key, - } - if root_span_attributes: - root_span_attributes_dict.update(root_span_attributes) - - root_span = self.tracer.start_span(name=span_name, attributes=root_span_attributes_dict, context=Context()) - - with use_span(root_span, end_on_exit=False): - self.propagator.inject(carrier=carrier) - - set_span_status(root_span, root_span_error) - root_span.end() - self.dify_trace_ids.add(trace_key) - self.root_span_carriers[trace_key] = carrier - - self.carrier = self.root_span_carriers[trace_key] - return self.carrier - - def api_check(self): - if isinstance(self.arize_phoenix_config, ArizeConfig): - from dify_trace_arize_phoenix.arize_configuration import verify_configuration - - verify_configuration(self.arize_phoenix_config) - return True - try: - with self.tracer.start_span("api_check") as span: - span.set_attribute("test", "true") - return True - except Exception as e: - logger.info("[Arize/Phoenix] API check failed: %s", str(e), exc_info=True) - raise ValueError(f"[Arize/Phoenix] API check failed: {str(e)}") from e - - def get_project_url(self): - """Get a Phoenix redirect; Arize Cloud requires verified destination metadata.""" - try: - project_name = self.arize_phoenix_config.project - endpoint = self.arize_phoenix_config.endpoint.rstrip("/") - - # Arize Cloud links require verified organization/space/project IDs. - # Only explicit config writes resolve and cache them; this legacy - # instance has none of that metadata and cannot promise a project. - if isinstance(self.arize_phoenix_config, ArizeConfig): - return None - - # Phoenix - return f"{endpoint}/projects/?redirect_project_name={project_name}" - - except Exception as e: - logger.info("[Arize/Phoenix] Failed to construct project URL: %s", str(e), exc_info=True) - raise ValueError(f"[Arize/Phoenix] Failed to construct project URL: {str(e)}") from e - - def _construct_llm_attributes(self, prompts: dict[str, Any] | list[Any] | str | None) -> dict[str, str]: - """Construct LLM attributes with passed prompts for Arize/Phoenix.""" - attributes: dict[str, str] = {} - - def set_attribute(path: str, value: object) -> None: - """Store an attribute safely as a string.""" - if value is None: - return - try: - if isinstance(value, (dict, list)): - value = safe_json_dumps(value) - attributes[path] = str(value) - except Exception: - attributes[path] = str(value) - - def set_message_attribute(message_index: int, key: str, value: object) -> None: - path = f"{SpanAttributes.LLM_INPUT_MESSAGES}.{message_index}.{key}" - set_attribute(path, value) - - def set_tool_call_attributes( - message_index: int, tool_index: int, tool_call: dict[str, Any] | object | None - ) -> None: - """Extract and assign tool call details safely.""" - if not tool_call: - return - - def safe_get(obj, key, default=None): - if isinstance(obj, dict): - return obj.get(key, default) - return getattr(obj, key, default) - - function_obj = safe_get(tool_call, "function", {}) - function_name = safe_get(function_obj, "name", "") - function_args = safe_get(function_obj, "arguments", {}) - call_id = safe_get(tool_call, "id", "") - - base_path = ( - f"{SpanAttributes.LLM_INPUT_MESSAGES}." - f"{message_index}.{MessageAttributes.MESSAGE_TOOL_CALLS}.{tool_index}" - ) - - set_attribute(f"{base_path}.{ToolCallAttributes.TOOL_CALL_FUNCTION_NAME}", function_name) - set_attribute(f"{base_path}.{ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON}", function_args) - set_attribute(f"{base_path}.{ToolCallAttributes.TOOL_CALL_ID}", call_id) - - # Handle list of messages - match prompts: - case list(): - for message_index, message in enumerate(prompts): - if not isinstance(message, dict): - continue - - role = message.get("role", "user") - content = message.get("text") or message.get("content") or "" - - set_message_attribute(message_index, MessageAttributes.MESSAGE_ROLE, role) - set_message_attribute(message_index, MessageAttributes.MESSAGE_CONTENT, content) - - tool_calls = message.get("tool_calls") or [] - if isinstance(tool_calls, list): - for tool_index, tool_call in enumerate(tool_calls): - set_tool_call_attributes(message_index, tool_index, tool_call) - - # Handle single dict or plain string prompt - case dict() | str(): - set_message_attribute(0, MessageAttributes.MESSAGE_CONTENT, prompts) - set_message_attribute(0, MessageAttributes.MESSAGE_ROLE, "user") - - return attributes diff --git a/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/config.py b/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/config.py index 6eac5b30d222b0..6dd0d5fa0e1364 100644 --- a/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/config.py +++ b/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/config.py @@ -1,45 +1,67 @@ -from pydantic import ValidationInfo, field_validator +"""Phoenix and Arize configurations. -from core.ops.entities.config_entity import BaseTracingConfig -from core.ops.utils import validate_url_with_path +Field names match the configurations Dify stored before this package used the +provider SDK, so stored records keep loading; keys these models no longer +declare, such as a stored ``project_url``, are ignored. +""" +from urllib.parse import urlsplit, urlunsplit -class ArizeConfig(BaseTracingConfig): - """ - Model class for Arize tracing config. - """ +from dify_trace_sdk.config import ProviderConfig, derived, destination, http_url, name_or_default, secret +from pydantic import field_validator + +DEFAULT_PROJECT = "default" +PHOENIX_DEFAULT_ENDPOINT = "https://app.phoenix.arize.com" +ARIZE_DEFAULT_ENDPOINT = "https://otlp.arize.com" +_SCHEMES = ("https", "http") + + +def _project_or_default(value: object) -> object: + # Stored records may hold ``None`` for the project. + return name_or_default(value, DEFAULT_PROJECT) if value is None or isinstance(value, str) else value - api_key: str | None = None - space_id: str | None = None - project: str | None = None - endpoint: str = "https://otlp.arize.com" - @field_validator("project") +class PhoenixConfig(ProviderConfig): + api_key: str | None = secret(default=None) + """Optional: self-hosted Phoenix may run without authentication.""" + project: str = DEFAULT_PROJECT + endpoint: str = destination(default=PHOENIX_DEFAULT_ENDPOINT) + """The Phoenix base URL. A path prefix, such as a Phoenix Cloud space, is kept.""" + + @field_validator("project", mode="before") @classmethod - def project_validator(cls, v, info: ValidationInfo): - return cls.validate_project_field(v, "default") + def _project(cls, value: object) -> object: + return _project_or_default(value) @field_validator("endpoint") @classmethod - def endpoint_validator(cls, v, info: ValidationInfo): - return cls.validate_endpoint_url(v, "https://otlp.arize.com") + def _endpoint(cls, value: str) -> str: + return http_url(value, allowed_schemes=_SCHEMES) if value.strip() else PHOENIX_DEFAULT_ENDPOINT -class PhoenixConfig(BaseTracingConfig): - """ - Model class for Phoenix tracing config. - """ +class ArizeConfig(ProviderConfig): + api_key: str = secret() + space_id: str = secret() + project: str = DEFAULT_PROJECT + endpoint: str = destination(default=ARIZE_DEFAULT_ENDPOINT) + """The OTLP/gRPC collector. gRPC addresses a host, so a path is dropped as it always was.""" + arize_project_destination: dict[str, str] | None = derived() + """The Arize Cloud project verification found, with the endpoint, space, and project it was found for. - api_key: str | None = None - project: str | None = None - endpoint: str = "https://app.phoenix.arize.com" + Records saved before this package used the provider SDK also carry a + ``source_digest`` key; ``project_url`` ignores it and compares the + destination with the configuration instead. + """ - @field_validator("project") + @field_validator("project", mode="before") @classmethod - def project_validator(cls, v, info: ValidationInfo): - return cls.validate_project_field(v, "default") + def _project(cls, value: object) -> object: + return _project_or_default(value) @field_validator("endpoint") @classmethod - def endpoint_validator(cls, v, info: ValidationInfo): - return validate_url_with_path(v, "https://app.phoenix.arize.com") + def _endpoint(cls, value: str) -> str: + if not value.strip(): + return ARIZE_DEFAULT_ENDPOINT + parsed = urlsplit(http_url(value, allowed_schemes=_SCHEMES)) + return urlunsplit((parsed.scheme, parsed.netloc, "", "", "")) diff --git a/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/configuration.py b/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/configuration.py deleted file mode 100644 index 8085c7462ea6ee..00000000000000 --- a/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/configuration.py +++ /dev/null @@ -1,182 +0,0 @@ -"""Phoenix configuration checks never create a tracer, project, or test span. - -A bounded project-list GET establishes read access only, including on a keyless -self-hosted server. It does not establish ingestion or project write permission. -Legacy Arize/Phoenix runtime adapters remain separate owners. -""" - -import re -from collections.abc import Mapping -from dataclasses import dataclass -from typing import Literal -from urllib.parse import quote, unquote, urlsplit, urlunsplit - -from pydantic import TypeAdapter - -from core.helper import ssrf_proxy -from dify_trace_arize_phoenix.config import PhoenixConfig - -type VerificationCode = Literal[ - "invalid_endpoint", - "invalid_project", - "invalid_api_key", - "authentication_failed", - "permission_denied", - "endpoint_unavailable", - "malformed_response", - "connection_failed", -] -_MESSAGES: dict[VerificationCode, str] = { - "invalid_endpoint": "Phoenix requires an HTTP(S) base URL, optionally with a path prefix, without /v1/traces.", - "invalid_project": "Phoenix requires a valid project name.", - "invalid_api_key": "Phoenix requires an API key without whitespace or mask placeholders, or no key.", - "authentication_failed": "Phoenix authentication failed.", - "permission_denied": "Phoenix project-list access was denied.", - "endpoint_unavailable": "Phoenix endpoint could not be verified.", - "malformed_response": "Phoenix verification returned an invalid response.", - "connection_failed": "Phoenix verification request failed.", -} -_JSON_OBJECT = TypeAdapter(dict[str, object]) -_MAX_RESPONSE_BYTES = 1024 * 1024 -_DEFAULT_ENDPOINT = "https://app.phoenix.arize.com" - - -class PhoenixConfigurationError(ValueError): - """Safe diagnostics exclude credentials, request URLs, and vendor response bodies.""" - - def __init__(self, code: VerificationCode) -> None: - self.code = code - super().__init__(_MESSAGES[code]) - - -@dataclass(frozen=True, slots=True) -class ReadablePhoenixEndpoint: - endpoint: str - project: str - - -def _endpoint(value: object) -> str: - if not isinstance(value, str): - raise PhoenixConfigurationError("invalid_endpoint") - value = value.strip(" ") or _DEFAULT_ENDPOINT - try: - parsed = urlsplit(value) - host, port = parsed.hostname, parsed.port - path = parsed.path.rstrip("/") - decoded_path = unquote(path, errors="strict") - valid = ( - parsed.scheme in {"http", "https"} - and host - and parsed.username is None - and parsed.password is None - and "?" not in value - and "#" not in value - and "\\" not in value - and not any(ord(char) <= 32 or ord(char) == 127 for char in value) - and not parsed.netloc.endswith(":") - and (port is None or 1 <= port <= 65535) - and "%" not in host - and not re.search(r"%(?![0-9a-fA-F]{2})", path) - and not re.search(r"%(?:2f|5c)", path, re.IGNORECASE) - and not any(char in decoded_path for char in "\\?#%") - and not any(ord(char) < 32 or ord(char) == 127 for char in decoded_path) - and "//" not in parsed.path - and not any(part in {".", ".."} for part in decoded_path.split("/")) - and not decoded_path.endswith("/v1/traces") - ) - if not valid or host is None: - raise PhoenixConfigurationError("invalid_endpoint") - host = host.encode("idna").decode("ascii").lower() - authority = f"[{host}]" if ":" in host else host - if port is not None: - authority = f"{authority}:{port}" - return urlunsplit((parsed.scheme, authority, path, "", "")) - except (UnicodeError, ValueError): - raise PhoenixConfigurationError("invalid_endpoint") from None - - -def _project(value: object) -> str: - if value is None: - return "default" - if not isinstance(value, str) or any(ord(char) < 32 or ord(char) == 127 for char in value): - raise PhoenixConfigurationError("invalid_project") - result = value.strip() or "default" - if result in {".", ".."}: - raise PhoenixConfigurationError("invalid_project") - return result - - -def normalize_configuration(values: Mapping[str, object]) -> PhoenixConfig: - """Normalize new writes without tightening the model used to read old records. - - Secret contents are checked only after PATCH restores masked credentials. - The project redirect loader uses a GraphQL name variable, not REST's more - restrictive project-identifier path. Do not impose that REST limit here. - """ - key = values.get("api_key") - if key is not None and not isinstance(key, str): - raise PhoenixConfigurationError("invalid_api_key") - return PhoenixConfig( - endpoint=_endpoint(values.get("endpoint", _DEFAULT_ENDPOINT)), - project=_project(values.get("project")), - api_key=key, - ) - - -def project_url(config: PhoenixConfig) -> str: - """Compute the configured UI redirect without SDK initialization or network I/O.""" - return f"{_endpoint(config.endpoint)}/redirects/projects/{quote(_project(config.project), safe='')}" - - -def verify_configuration(config: PhoenixConfig) -> ReadablePhoenixEndpoint: - """Read at most one page; an empty project list is valid and creates nothing.""" - config = normalize_configuration(config.model_dump()) - key = config.api_key - if key and ("*" in key or any(char.isspace() or ord(char) < 32 or ord(char) == 127 for char in key)): - raise PhoenixConfigurationError("invalid_api_key") - headers = {"Accept": "application/json", "Accept-Encoding": "identity", "Cookie": ""} - if key: - headers["Authorization"] = f"Bearer {key}" - try: - response = ssrf_proxy.get( - config.endpoint + "/v1/projects", - params={"limit": 1}, - headers=headers, - auth=None, - timeout=30.0, - max_retries=0, - follow_redirects=False, - stream_response=True, - ) - except Exception: - raise PhoenixConfigurationError("connection_failed") from None - try: - if response.status_code == 401: - raise PhoenixConfigurationError("authentication_failed") - if response.status_code == 403: - raise PhoenixConfigurationError("permission_denied") - if response.status_code != 200: - raise PhoenixConfigurationError("endpoint_unavailable") - try: - buffered = ssrf_proxy.buffer_response(response, max_response_bytes=_MAX_RESPONSE_BYTES) - data = _JSON_OBJECT.validate_python(buffered.json(), strict=True) - except (ValueError, TypeError): - raise PhoenixConfigurationError("malformed_response") from None - except Exception: - raise PhoenixConfigurationError("connection_failed") from None - finally: - response.close() - rows, cursor = data.get("data"), data.get("next_cursor") - if not isinstance(rows, list) or "next_cursor" not in data or (cursor is not None and not isinstance(cursor, str)): - raise PhoenixConfigurationError("malformed_response") - for row in rows: - if ( - not isinstance(row, dict) - or not isinstance(row.get("id"), str) - or not row["id"] - or not isinstance(row.get("name"), str) - or not row["name"] - or (row.get("description") is not None and not isinstance(row["description"], str)) - ): - raise PhoenixConfigurationError("malformed_response") - return ReadablePhoenixEndpoint(config.endpoint, _project(config.project)) diff --git a/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/provider.py b/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/provider.py new file mode 100644 index 00000000000000..dbd5949ec8e45a --- /dev/null +++ b/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/provider.py @@ -0,0 +1,285 @@ +"""Phoenix and Arize destinations: OTLP spans with OpenInference attributes. + +Phoenix ingests OTLP/HTTP at ``{endpoint}/v1/traces``; Arize ingests OTLP/gRPC. +Spans keep the trace's OpenTelemetry IDs, so a unit whose root has an external +parent joins the tree that was exported separately. +""" + +import re +from typing import override +from urllib.parse import quote, unquote, urlsplit, urlunsplit + +import grpc # type: ignore[import-untyped] +from dify_trace_sdk.errors import EgressError, ExportRejectedError, VerificationError +from dify_trace_sdk.exporting import export_endpoint, json_value +from dify_trace_sdk.model import Span, SpanKind, Trace +from dify_trace_sdk.otlp import encode_trace, export_otlp_grpc, post_otlp, read_otlp_ack +from dify_trace_sdk.provider import ProviderCapabilities, SendResult, TraceProvider +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest +from opentelemetry.proto.collector.trace.v1.trace_service_pb2_grpc import TraceServiceStub + +from dify_trace_arize_phoenix import _arize_project +from dify_trace_arize_phoenix.config import DEFAULT_PROJECT, ArizeConfig, PhoenixConfig + +SERVICE_NAME = "dify-application" +JSON_MIME_TYPE = "application/json" +PHOENIX_VERIFY_TIMEOUT_SECONDS = 30.0 +ARIZE_VERIFY_TIMEOUT_SECONDS = 5.0 +MAX_VERIFY_RESPONSE_BYTES = 1024 * 1024 +_CAPABILITIES = ProviderCapabilities(native_cost=False, trace_url=False, span_url=False) +_MALFORMED_ESCAPE = re.compile(r"%(?![0-9a-f]{2})|%(?:2f|5c)", re.IGNORECASE) +"""A ``%`` that starts no escape, or an escaped ``/`` or ``\\``.""" +_PROBE_FAILURES = { + grpc.StatusCode.UNAUTHENTICATED: "authentication_failed", + grpc.StatusCode.PERMISSION_DENIED: "permission_denied", +} + + +def _span_attributes(trace: Trace, span: Span) -> dict[str, object]: + attributes: dict[str, object] = { + "openinference.span.kind": span.kind.value.upper(), + "input.value": json_value(span.inputs), + "input.mime_type": JSON_MIME_TYPE, + "output.value": json_value(span.outputs), + "output.mime_type": JSON_MIME_TYPE, + "metadata": json_value(span.attributes), + "session.id": trace.session_id, + "user.id": trace.user_id, + "dify.span.id": span.span_id, + } + # Only model calls report usage: a wrapper's usage repeats its children's. + if span.kind is SpanKind.LLM: + attributes["llm.model_name"] = span.model_name + attributes["llm.provider"] = span.model_provider + if span.usage is not None: + attributes["llm.token_count.prompt"] = span.usage.prompt_tokens + attributes["llm.token_count.completion"] = span.usage.completion_tokens + attributes["llm.token_count.total"] = span.usage.total_tokens + elif span.kind is SpanKind.TOOL: + attributes["tool.name"] = span.name + attributes["tool.parameters"] = json_value(span.inputs) + return attributes + + +def _export_request(project: str, trace: Trace) -> ExportTraceServiceRequest: + return encode_trace( + trace, + resource={"service.name": SERVICE_NAME, "openinference.project.name": project, "model_id": project}, + attributes=lambda span: _span_attributes(trace, span), + ) + + +def _is_credential(value: str) -> bool: + """Non-empty printable ASCII without whitespace or the ``*`` of a masked secret. + + Real keys never contain these. gRPC refuses non-ASCII metadata with an + error that echoes the value, so it must never see one. + """ + return bool(value) and "*" not in value and all("!" <= char <= "~" for char in value) + + +def _phoenix_base_url(value: str) -> str | None: + """Normalize a Phoenix base URL, or return ``None`` when it cannot be one. + + The path prefix stays (Phoenix Cloud spaces and reverse proxies need it), + but nothing that changes what an appended path names: no query, fragment, + backslash, escaped separator, empty or dot segment, or the ``/v1/traces`` + suffix of the ingestion URL. + """ + try: + parsed = urlsplit(value) + host, port = parsed.hostname, parsed.port + path = parsed.path.rstrip("/") + decoded_path = unquote(path, errors="strict") + if ( + parsed.scheme not in {"http", "https"} + or not host + or parsed.username is not None + or parsed.password is not None + or any(char in value for char in "?#\\") + or any(ord(char) <= 32 or ord(char) == 127 for char in value) + or parsed.netloc.endswith(":") + or port == 0 + or "%" in host + or _MALFORMED_ESCAPE.search(path) + or any(char in decoded_path for char in "\\?#%") + or any(ord(char) < 32 or ord(char) == 127 for char in decoded_path) + or "//" in parsed.path + or any(segment in {".", ".."} for segment in decoded_path.split("/")) + or decoded_path.endswith("/v1/traces") + ): + return None + host = host.encode("idna").decode("ascii").lower() + except (UnicodeError, ValueError): + return None + authority = f"[{host}]" if ":" in host else host + return urlunsplit((parsed.scheme, authority if port is None else f"{authority}:{port}", path, "", "")) + + +def _phoenix_project(value: str) -> str | None: + # The project redirect looks names up through GraphQL, so REST's stricter identifier rules do not apply. + project = value.strip() or DEFAULT_PROJECT + if project in {".", ".."} or any(ord(char) < 32 or ord(char) == 127 for char in project): + return None + return project + + +def _is_project_page(body: object) -> bool: + """Whether ``body`` is a page of Phoenix's ``GET /v1/projects``; a proxy or login page is not.""" + if not isinstance(body, dict) or not isinstance(body.get("data"), list) or "next_cursor" not in body: + return False + if body["next_cursor"] is not None and not isinstance(body["next_cursor"], str): + return False + return all( + isinstance(row, dict) + and isinstance(row.get("id"), str) + and bool(row["id"]) + and isinstance(row.get("name"), str) + and bool(row["name"]) + and (row.get("description") is None or isinstance(row["description"], str)) + for row in body["data"] + ) + + +class PhoenixProvider(TraceProvider[PhoenixConfig]): + id = "phoenix" + config_model = PhoenixConfig + capabilities = _CAPABILITIES + + @override + def verify(self, config: PhoenixConfig) -> PhoenixConfig: + """Read one page of the project list and return the normalized endpoint and project. + + This proves read access only, which a keyless self-hosted server grants + to anyone. It creates nothing and does not prove ingestion access. + """ + endpoint = _phoenix_base_url(config.endpoint) + if endpoint is None: + raise VerificationError("invalid_endpoint") + project = _phoenix_project(config.project) + if project is None: + raise VerificationError("invalid_project") + headers = {"Accept": JSON_MIME_TYPE} + if config.api_key: + if not _is_credential(config.api_key): + raise VerificationError("invalid_api_key") + headers["Authorization"] = f"Bearer {config.api_key}" + try: + response = self.egress.request( + "GET", + f"{endpoint}/v1/projects?limit=1", + headers=headers, + content=None, + timeout=PHOENIX_VERIFY_TIMEOUT_SECONDS, + sensitive_url=False, + ) + except EgressError: + raise VerificationError("connection_failed") from None + if response.status_code == 401: + raise VerificationError("authentication_failed") + if response.status_code == 403: + raise VerificationError("permission_denied") + if response.status_code != 200: + raise VerificationError("endpoint_unavailable") + try: + body = response.json() if len(response.content) <= MAX_VERIFY_RESPONSE_BYTES else None + except (ValueError, RecursionError): + body = None + if not _is_project_page(body): + raise VerificationError("malformed_response") + return config.model_copy(update={"endpoint": endpoint, "project": project}) + + @override + def project_url(self, config: PhoenixConfig) -> str | None: + endpoint, project = _phoenix_base_url(config.endpoint), _phoenix_project(config.project) + if endpoint is None or project is None: + return None + return f"{endpoint}/redirects/projects/{quote(project, safe='')}" + + @override + def send(self, config: PhoenixConfig, trace: Trace) -> SendResult: + endpoint = export_endpoint(config.endpoint) + headers: dict[str, str] = {} + if config.api_key: + if not _is_credential(config.api_key): + raise ExportRejectedError("Phoenix API key is not a valid header value") + headers = {"api_key": config.api_key, "authorization": f"Bearer {config.api_key}"} + response = post_otlp( + self.egress, + f"{endpoint}/v1/traces", + _export_request(config.project, trace), + destination="Phoenix", + headers=headers, + sensitive_url=False, + ) + read_otlp_ack(response, "Phoenix") + return SendResult(trace_id=trace.trace_id, trace_url=None, span_urls={}) + + +def _arize_metadata(config: ArizeConfig) -> dict[str, str]: + return {"api_key": config.api_key, "space_id": config.space_id, "authorization": f"Bearer {config.api_key}"} + + +def _has_valid_credentials(config: ArizeConfig) -> bool: + return _is_credential(config.api_key) and _is_credential(config.space_id) + + +def _has_valid_port(url: str) -> bool: + try: + return urlsplit(url).port != 0 + except ValueError: + return False + + +class ArizeProvider(TraceProvider[ArizeConfig]): + id = "arize" + config_model = ArizeConfig + capabilities = _CAPABILITIES + + @override + def verify(self, config: ArizeConfig) -> ArizeConfig: + """Require the collector to acknowledge an export, then look up the project link. + + The export carries a resource and no spans, so it creates nothing. It + proves ingestion access only; the link lookup may fail without failing + verification, because a key may ingest without read access. + """ + if not _has_valid_credentials(config) or not _has_valid_port(config.endpoint): + raise VerificationError("invalid_configuration") + probe = ExportTraceServiceRequest() + resource = probe.resource_spans.add().resource + for key in ("model_id", "openinference.project.name"): + resource.attributes.add(key=key).value.string_value = config.project + try: + channel = self.egress.grpc_channel(config.endpoint) + except EgressError: + raise VerificationError("connection_failed") from None + try: + result = TraceServiceStub(channel).Export( + probe, metadata=tuple(_arize_metadata(config).items()), timeout=ARIZE_VERIFY_TIMEOUT_SECONDS + ) + except grpc.RpcError as error: + raise VerificationError(_PROBE_FAILURES.get(error.code(), "connection_failed")) from None + finally: + channel.close() + if result.partial_success.rejected_spans or result.partial_success.error_message: + raise VerificationError("unexpected_response") + destination = _arize_project.resolve_destination(self.egress, config) + return config.model_copy(update={"arize_project_destination": destination}) + + @override + def project_url(self, config: ArizeConfig) -> str | None: + return _arize_project.project_url(config) + + @override + def send(self, config: ArizeConfig, trace: Trace) -> SendResult: + if not _has_valid_credentials(config): + raise ExportRejectedError("Arize credentials are not valid gRPC metadata") + export_otlp_grpc( + self.egress, + config.endpoint, + _export_request(config.project, trace), + destination="Arize", + metadata=_arize_metadata(config), + ) + return SendResult(trace_id=trace.trace_id, trace_url=None, span_urls={}) 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 deleted file mode 100644 index f82505c5f9428a..00000000000000 --- a/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/unified_trace.py +++ /dev/null @@ -1,195 +0,0 @@ -"""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-arize-phoenix/tests/unit_tests/src/dify_trace_arize_phoenix/test_arize_phoenix_trace.py b/api/providers/trace/trace-arize-phoenix/tests/unit_tests/src/dify_trace_arize_phoenix/test_arize_phoenix_trace.py deleted file mode 100644 index 6394aefafe5abb..00000000000000 --- a/api/providers/trace/trace-arize-phoenix/tests/unit_tests/src/dify_trace_arize_phoenix/test_arize_phoenix_trace.py +++ /dev/null @@ -1,2130 +0,0 @@ -"""Unit tests for Arize/Phoenix tracing with a real SQLite session factory.""" - -import json -from collections.abc import Sequence -from datetime import UTC, datetime, timedelta -from types import SimpleNamespace -from typing import Any, cast, override -from unittest.mock import MagicMock, patch - -import dify_trace_arize_phoenix.arize_phoenix_trace as arize_phoenix_trace_module -import pytest -from dify_trace_arize_phoenix.arize_phoenix_trace import ( - _NODE_TYPE_TO_SPAN_KIND, - ArizePhoenixDataTrace, - _app_uses_phoenix_provider, - _build_graph_parent_index, - _build_wrapper_groups, - _get_node_span_kind, - _normalize_wrapper_index, - _parent_workflow_can_publish_span_context, - _phoenix_parent_span_redis_key, - _resolve_node_parent, - _resolve_published_parent_span_context, - _resolve_structured_parent_execution_id, - _resolve_trace_session_id, - _resolve_workflow_parent_context, - _resolve_workflow_session_id, - datetime_to_nanos, - error_to_string, - safe_json_dumps, - set_span_status, - setup_tracer, - wrap_span_metadata, -) -from dify_trace_arize_phoenix.config import ArizeConfig, PhoenixConfig -from openinference.semconv.trace import OpenInferenceSpanKindValues, SpanAttributes -from opentelemetry.context import Context -from opentelemetry.sdk import trace as trace_sdk -from opentelemetry.sdk.trace import ReadableSpan, Tracer -from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExporter, SpanExportResult - -# pyrefly: ignore [deprecated] -from opentelemetry.semconv.trace import SpanAttributes as OTELSpanAttributes -from opentelemetry.trace import NonRecordingSpan, SpanContext, StatusCode, TraceFlags, TraceState, use_span -from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator -from sqlalchemy.orm import Session - -from core.ops.entities.trace_entity import ( - DatasetRetrievalTraceInfo, - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - SuggestedQuestionTraceInfo, - ToolTraceInfo, - TraceTaskName, - WorkflowNodeTraceInfo, - WorkflowTraceInfo, -) -from core.ops.exceptions import PendingTraceParentContextError -from graphon.enums import BUILT_IN_NODE_TYPES, BuiltinNodeTypes, WorkflowNodeExecutionStatus -from models import WorkflowNodeExecutionModel - -pytestmark = pytest.mark.parametrize("sqlite3_session", [(WorkflowNodeExecutionModel,)], indirect=True) - -# --- Helpers --- - - -@pytest.fixture(autouse=True) -def _sqlite_trace_database(monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session) -> None: - """Route trace-created session factories to an isolated SQLite engine.""" - - monkeypatch.setattr( - arize_phoenix_trace_module, - "db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - - -def _dt(): - return datetime(2024, 1, 1, 0, 0, 0, tzinfo=UTC) - - -def _make_workflow_info(**kwargs): - defaults = { - "workflow_id": "w1", - "tenant_id": "t1", - "workflow_run_id": "r1", - "workflow_run_elapsed_time": 1.0, - "workflow_run_status": "succeeded", - "workflow_run_inputs": {"in": "val"}, - "workflow_run_outputs": {"out": "val"}, - "workflow_run_version": "1.0", - "total_tokens": 10, - "file_list": ["f1"], - "query": "hi", - "metadata": {"app_id": "app1"}, - "start_time": _dt(), - "end_time": _dt() + timedelta(seconds=1), - } - defaults.update(kwargs) - return WorkflowTraceInfo(**defaults) - - -def _make_message_info(**kwargs): - defaults = { - "conversation_model": "chat", - "message_tokens": 5, - "answer_tokens": 5, - "total_tokens": 10, - "conversation_mode": "chat", - "metadata": {"app_id": "app1"}, - "inputs": {"in": "val"}, - "outputs": "val", - "start_time": _dt(), - "end_time": _dt(), - "message_id": "m1", - } - defaults.update(kwargs) - return MessageTraceInfo(**defaults) - - -def _get_start_span_call(start_span_mock, *, span_name: str): - for call in start_span_mock.call_args_list: - if call.kwargs.get("name") == span_name: - return call - raise AssertionError(f"Could not find start_span call with name={span_name!r}") - - -def _get_start_span_call_by_kind(start_span_mock, *, span_kind: str): - for call in start_span_mock.call_args_list: - attributes = call.kwargs.get("attributes", {}) - if attributes.get(SpanAttributes.OPENINFERENCE_SPAN_KIND) == span_kind: - return call - raise AssertionError(f"Could not find start_span call with span kind={span_kind!r}") - - -class _FakeQuery: - def __init__(self, result): - self._result = result - - def filter(self, *args, **kwargs): - return self - - def where(self, *args, **kwargs): - return self - - def first(self): - return self._result - - -class _CollectingSpanExporter(SpanExporter): - def __init__(self): - self.spans: list[ReadableSpan] = [] - - @override - def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: - self.spans.extend(spans) - return SpanExportResult.SUCCESS - - @override - def shutdown(self) -> None: - return None - - -def _make_node_execution(**kwargs): - defaults = { - "node_type": "tool", - "status": "succeeded", - "inputs": {}, - "outputs": {}, - "created_at": _dt(), - "elapsed_time": 1.0, - "process_data": {}, - "metadata": {}, - "title": "Node", - "id": "node-execution-1", - "node_execution_id": "node-execution-1", - "node_id": "node-1", - "predecessor_node_id": None, - "iteration_id": None, - "loop_id": None, - "error": None, - } - defaults.update(kwargs) - node_execution = MagicMock() - for key, value in defaults.items(): - setattr(node_execution, key, value) - return node_execution - - -def _make_workflow_trace_info(**kwargs) -> WorkflowTraceInfo: - defaults = { - "workflow_id": "workflow-1", - "tenant_id": "tenant-1", - "workflow_run_id": "workflow-run-1", - "workflow_run_elapsed_time": 1.0, - "workflow_run_status": "succeeded", - "workflow_run_inputs": {"input": "value"}, - "workflow_run_outputs": {"output": "value"}, - "workflow_run_version": "1.0", - "total_tokens": 10, - "file_list": ["file-1"], - "query": "hello", - "metadata": {"app_id": "app-1"}, - "start_time": _dt(), - "end_time": _dt() + timedelta(seconds=1), - } - defaults.update(kwargs) - return WorkflowTraceInfo(**defaults) - - -def _make_workflow_node_trace_info(**kwargs) -> WorkflowNodeTraceInfo: - defaults = { - "workflow_id": "workflow-1", - "workflow_run_id": "workflow-run-1", - "tenant_id": "tenant-1", - "node_execution_id": "node-execution-1", - "node_id": "node-1", - "node_type": "tool", - "title": "Node 1", - "status": "succeeded", - "elapsed_time": 1.0, - "index": 1, - "metadata": {"app_id": "app-1"}, - "start_time": _dt(), - "end_time": _dt() + timedelta(seconds=1), - } - defaults.update(kwargs) - return WorkflowNodeTraceInfo(**defaults) - - -# --- Utility Function Tests --- - - -def test_datetime_to_nanos(): - dt = _dt() - expected = int(dt.timestamp() * 1_000_000_000) - assert datetime_to_nanos(dt) == expected - - with patch("dify_trace_arize_phoenix.arize_phoenix_trace.datetime") as mock_dt: - mock_now = MagicMock() - mock_now.timestamp.return_value = 1704110400.0 - mock_dt.now.return_value = mock_now - assert datetime_to_nanos(None) == 1704110400000000000 - - -def test_error_to_string(): - try: - raise ValueError("boom") - except ValueError as e: - err = e - - res = error_to_string(err) - assert "ValueError: boom" in res - assert "traceback" in res.lower() or "line" in res.lower() - - assert error_to_string("str error") == "str error" - assert error_to_string(None) == "Empty Stack Trace" - - -def test_set_span_status(): - span = MagicMock() - # OK - set_span_status(span, None) - span.set_status.assert_called() - assert span.set_status.call_args[0][0].status_code == StatusCode.OK - - # Error Exception - span.reset_mock() - set_span_status(span, ValueError("fail")) - assert span.set_status.call_args[0][0].status_code == StatusCode.ERROR - span.record_exception.assert_called() - - # Error String - span.reset_mock() - set_span_status(span, "fail-str") - assert span.set_status.call_args[0][0].status_code == StatusCode.ERROR - span.add_event.assert_called() - - # repr branch - class SilentError: - @override - def __str__(self): - return "" - - @override - def __repr__(self): - return "SilentErrorRepr" - - span.reset_mock() - set_span_status(span, cast(Exception | str | None, SilentError())) - assert span.add_event.call_args[1]["attributes"][OTELSpanAttributes.EXCEPTION_MESSAGE] == "SilentErrorRepr" - - -def test_safe_json_dumps(): - assert safe_json_dumps({"a": _dt()}) == '{"a": "2024-01-01 00:00:00+00:00"}' - - -def test_wrap_span_metadata(): - res = wrap_span_metadata({"a": 1}, b=2) - assert res == {"a": 1, "b": 2, "created_from": "Dify"} - - -def test_app_uses_phoenix_provider_only_for_enabled_arize_or_phoenix(): - assert _app_uses_phoenix_provider({"enabled": True, "tracing_provider": "phoenix"}) is True - assert _app_uses_phoenix_provider({"enabled": True, "tracing_provider": "arize"}) is True - assert _app_uses_phoenix_provider({"enabled": False, "tracing_provider": "phoenix"}) is False - assert _app_uses_phoenix_provider({"enabled": True, "tracing_provider": "langfuse"}) is False - assert _app_uses_phoenix_provider(None) is False - - -def test_normalize_wrapper_index_accepts_stable_values(): - assert _normalize_wrapper_index(0) == "0" - assert _normalize_wrapper_index(12) == "12" - assert _normalize_wrapper_index("01") == "01" - assert _normalize_wrapper_index("branch-1_A.2:3") == "branch-1_A.2:3" - - -@pytest.mark.parametrize( - "value", - [ - True, - False, - -1, - 1.0, - "", - " 1", - "1 ", - "group/1", - "group]1", - None, - ], -) -def test_normalize_wrapper_index_rejects_unstable_values(value): - assert _normalize_wrapper_index(value) is None - - -def test_parent_workflow_can_publish_span_context_keeps_unknown_parent_retryable(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr( - "dify_trace_arize_phoenix.arize_phoenix_trace.db.session.query", - lambda model: _FakeQuery(None), - ) - - assert _parent_workflow_can_publish_span_context("missing-run") is True - - -def test_parent_workflow_can_publish_span_context_checks_parent_app_tracing(monkeypatch: pytest.MonkeyPatch): - parent_run = SimpleNamespace(app_id="parent-app") - parent_app = SimpleNamespace(tracing=json.dumps({"enabled": True, "tracing_provider": "phoenix"})) - - def fake_query(model): - if getattr(model, "__tablename__", None) == "workflow_runs": - return _FakeQuery(parent_run) - if getattr(model, "__tablename__", None) == "apps": - return _FakeQuery(parent_app) - raise AssertionError(f"Unexpected model query: {model}") - - monkeypatch.setattr("dify_trace_arize_phoenix.arize_phoenix_trace.db.session.query", fake_query) - - assert _parent_workflow_can_publish_span_context("parent-run") is True - - parent_app.tracing = json.dumps({"enabled": False, "tracing_provider": "phoenix"}) - assert _parent_workflow_can_publish_span_context("parent-run") is False - - parent_app.tracing = json.dumps({"enabled": True, "tracing_provider": "langfuse"}) - assert _parent_workflow_can_publish_span_context("parent-run") is False - - -class TestGetNodeSpanKind: - def test_all_node_types_are_mapped_correctly(self): - special_mappings = { - BuiltinNodeTypes.LLM: OpenInferenceSpanKindValues.LLM, - BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL: OpenInferenceSpanKindValues.RETRIEVER, - BuiltinNodeTypes.TOOL: OpenInferenceSpanKindValues.TOOL, - BuiltinNodeTypes.AGENT: OpenInferenceSpanKindValues.AGENT, - } - - for node_type in BUILT_IN_NODE_TYPES: - expected_span_kind = special_mappings.get(node_type, OpenInferenceSpanKindValues.CHAIN) - actual_span_kind = _get_node_span_kind(node_type) - assert actual_span_kind == expected_span_kind, ( - f"Node type {node_type!r} was mapped to {actual_span_kind}, but {expected_span_kind} was expected." - ) - - def test_unknown_string_defaults_to_chain(self): - assert _get_node_span_kind("some-future-node-type") == OpenInferenceSpanKindValues.CHAIN - - def test_stale_dataset_retrieval_not_in_mapping(self): - assert "dataset_retrieval" not in _NODE_TYPE_TO_SPAN_KIND - - -class TestWorkflowSessionResolution: - def test_resolve_workflow_session_id_prefers_trace_session_id_metadata(self): - trace_info = _make_workflow_info( - conversation_id="conversation-1", - workflow_run_id="workflow-run-1", - metadata={"app_id": "app-1", "trace_session_id": "session-1"}, - ) - - assert _resolve_trace_session_id(trace_info) == "session-1" - assert _resolve_workflow_session_id(trace_info) == "session-1" - - def test_resolve_workflow_session_id_falls_back_to_existing_workflow_behavior(self): - trace_info = _make_workflow_info( - conversation_id="conversation-1", - workflow_run_id="workflow-run-1", - metadata={"app_id": "app-1"}, - ) - - assert _resolve_trace_session_id(trace_info) == "conversation-1" - - def test_resolve_message_session_id_prefers_trace_session_id_metadata(self): - message_data = SimpleNamespace(conversation_id="conversation-1") - trace_info = _make_message_info( - message_data=message_data, - metadata={"app_id": "app-1", "trace_session_id": "session-1"}, - ) - - assert _resolve_trace_session_id(trace_info) == "session-1" - - def test_prefers_conversation_id(self): - info = _make_workflow_trace_info(conversation_id="conversation-1") - - assert _resolve_workflow_session_id(info) == "conversation-1" - - def test_nested_workflow_keeps_own_conversation_id_when_parent_context_exists(self): - info = _make_workflow_trace_info( - conversation_id="conversation-1", - metadata={ - "app_id": "app-1", - "parent_trace_context": { - "parent_workflow_run_id": "outer-workflow-run-1", - "parent_node_execution_id": "outer-node-execution-1", - }, - }, - ) - - assert _resolve_workflow_session_id(info) == "conversation-1" - - def test_uses_parent_workflow_run_id_for_nested_parent_trace_context(self): - info = _make_workflow_trace_info( - conversation_id=None, - metadata={ - "app_id": "app-1", - "parent_trace_context": { - "parent_workflow_run_id": "outer-workflow-run-1", - "parent_node_execution_id": "outer-node-execution-1", - }, - }, - ) - - assert _resolve_workflow_session_id(info) == "outer-workflow-run-1" - - def test_falls_back_to_workflow_run_id(self): - info = _make_workflow_trace_info(conversation_id=None) - - assert _resolve_workflow_session_id(info) == "workflow-run-1" - - def test_parent_context_helper_delegates_to_resolved_parent_context(self): - info = MagicMock() - info.resolved_parent_context = ("outer-workflow-run-1", "outer-node-execution-1") - - assert _resolve_workflow_parent_context(info) == info.resolved_parent_context - - -class TestPhoenixParentSpanBridgeHelpers: - def test_parent_span_redis_key_is_stable(self): - assert _phoenix_parent_span_redis_key("outer-node-execution-1") == ( - "trace:phoenix:parent_span:outer-node-execution-1" - ) - - def test_pending_parent_exception_exposes_execution_id(self): - error = PendingTraceParentContextError("outer-node-execution-1") - - assert error.parent_node_execution_id == "outer-node-execution-1" - assert "outer-node-execution-1" in str(error) - - def test_resolve_parent_span_context_rejects_payload_without_traceparent(self, monkeypatch: pytest.MonkeyPatch): - mock_redis = MagicMock() - mock_redis.get.return_value = '{"tracestate": "vendor=value"}' - monkeypatch.setattr(arize_phoenix_trace_module, "redis_client", mock_redis) - - with pytest.raises(ValueError, match="traceparent"): - _resolve_published_parent_span_context("outer-node-execution-1") - - @pytest.mark.parametrize( - "stored_payload", - [ - '{"traceparent": ""}', - '{"traceparent": "not-a-traceparent"}', - '{"traceparent": "00-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-bbbbbbbbbbbbbbbb"}', - ], - ) - def test_resolve_parent_span_context_rejects_malformed_traceparent(self, monkeypatch, stored_payload): - mock_redis = MagicMock() - mock_redis.get.return_value = stored_payload - monkeypatch.setattr(arize_phoenix_trace_module, "redis_client", mock_redis) - - with pytest.raises(ValueError, match="traceparent"): - _resolve_published_parent_span_context("outer-node-execution-1") - - -class TestWorkflowHierarchyHelpers: - def test_build_graph_parent_index_uses_predecessor_nodes_without_order_heuristics(self): - later_node = _make_workflow_node_trace_info( - node_execution_id="node-execution-3", - node_id="node-3", - predecessor_node_id="node-2", - index=3, - ) - root_node = _make_workflow_node_trace_info( - node_execution_id="node-execution-1", - node_id="node-1", - predecessor_node_id=None, - index=1, - ) - middle_node = _make_workflow_node_trace_info( - node_execution_id="node-execution-2", - node_id="node-2", - predecessor_node_id="node-1", - index=2, - ) - - graph_parent_index = _build_graph_parent_index([later_node, root_node, middle_node]) - - assert graph_parent_index == { - "node-execution-2": "node-execution-1", - "node-execution-3": "node-execution-2", - } - - def test_build_graph_parent_index_drops_ambiguous_parallel_like_predecessors(self): - first_parallel_node = _make_workflow_node_trace_info( - node_execution_id="parallel-node-execution-1", - node_id="parallel-node-1", - predecessor_node_id=None, - index=1, - parallel_id="parallel-1", - ) - second_parallel_node = _make_workflow_node_trace_info( - node_execution_id="parallel-node-execution-2", - node_id="parallel-node-1", - predecessor_node_id=None, - index=2, - parallel_id="parallel-2", - ) - child_node = _make_workflow_node_trace_info( - node_execution_id="child-node-execution-1", - node_id="child-node-1", - predecessor_node_id="parallel-node-1", - index=3, - ) - - graph_parent_index = _build_graph_parent_index([child_node, first_parallel_node, second_parallel_node]) - - assert graph_parent_index == {} - - def test_resolve_node_parent_prefers_predecessor_span(self): - workflow_span = MagicMock(name="workflow-span") - predecessor_span = MagicMock(name="predecessor-span") - graph_parent_span = MagicMock(name="graph-parent-span") - - parent = _resolve_node_parent( - execution_id="node-execution-2", - predecessor_execution_id="node-execution-1", - structured_parent_execution_id=None, - span_by_execution_id={ - "node-execution-1": predecessor_span, - "node-execution-0": graph_parent_span, - }, - graph_parent_index={ - "node-execution-2": "node-execution-0", - }, - workflow_span=workflow_span, - ) - - assert parent is predecessor_span - - def test_resolve_node_parent_falls_back_to_graph_parent_span(self): - workflow_span = MagicMock(name="workflow-span") - graph_parent_span = MagicMock(name="graph-parent-span") - - parent = _resolve_node_parent( - execution_id="node-execution-2", - predecessor_execution_id="missing-predecessor", - structured_parent_execution_id=None, - span_by_execution_id={ - "node-execution-0": graph_parent_span, - }, - graph_parent_index={ - "node-execution-2": "node-execution-0", - }, - workflow_span=workflow_span, - ) - - assert parent is graph_parent_span - - def test_resolve_node_parent_falls_back_to_workflow_span(self): - workflow_span = MagicMock(name="workflow-span") - - parent = _resolve_node_parent( - execution_id="node-execution-2", - predecessor_execution_id=None, - structured_parent_execution_id=None, - span_by_execution_id={}, - graph_parent_index={}, - workflow_span=workflow_span, - ) - - assert parent is workflow_span - - def test_resolve_structured_parent_execution_id_allows_body_nodes_to_use_enclosing_structure(self): - body_node = _make_workflow_node_trace_info( - node_execution_id="body-execution-1", - node_id="body-node-1", - node_type="tool", - loop_id="loop-node-1", - ) - - structured_parent_execution_id = _resolve_structured_parent_execution_id( - body_node, - execution_id_by_node_id={ - "loop-node-1": "loop-execution-1", - }, - ) - - assert structured_parent_execution_id == "loop-execution-1" - - def test_resolve_structured_parent_execution_id_reads_execution_metadata_dict_for_models(self): - body_node = SimpleNamespace( - node_execution_id="body-execution-1", - node_id="body-node-1", - execution_metadata_dict={ - "iteration_id": "iteration-node-1", - "loop_id": "loop-node-1", - }, - ) - - structured_parent_execution_id = _resolve_structured_parent_execution_id( - body_node, - execution_id_by_node_id={ - "iteration-node-1": "iteration-execution-1", - "loop-node-1": "loop-execution-1", - }, - ) - - assert structured_parent_execution_id == "iteration-execution-1" - - def test_build_wrapper_groups_groups_loop_children_by_index(self): - loop = _make_workflow_node_trace_info( - node_execution_id="loop-execution-1", - node_id="loop-node-1", - node_type="loop", - ) - first = _make_workflow_node_trace_info( - node_execution_id="first-execution-1", - node_id="body-node-1", - node_type="tool", - loop_id="loop-node-1", - loop_index=0, - ) - second = _make_workflow_node_trace_info( - node_execution_id="second-execution-1", - node_id="body-node-2", - node_type="tool", - loop_id="loop-node-1", - loop_index=1, - ) - - groups = _build_wrapper_groups([loop, first, second]) - - assert [(group.key.wrapper_type, group.key.index) for group in groups.values()] == [ - ("loop", "0"), - ("loop", "1"), - ] - assert [group.container_execution_id for group in groups.values()] == [ - "loop-execution-1", - "loop-execution-1", - ] - - def test_build_wrapper_groups_skips_ambiguous_container_graph_ids(self): - first_loop = _make_workflow_node_trace_info( - node_execution_id="loop-execution-1", - node_id="loop-node-1", - node_type="loop", - ) - second_loop = _make_workflow_node_trace_info( - node_execution_id="loop-execution-2", - node_id="loop-node-1", - node_type="loop", - ) - child = _make_workflow_node_trace_info( - node_execution_id="child-execution-1", - node_id="body-node-1", - node_type="tool", - loop_id="loop-node-1", - loop_index=0, - ) - - assert _build_wrapper_groups([first_loop, second_loop, child]) == {} - - -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.GrpcOTLPSpanExporter") -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.trace_sdk.TracerProvider") -def test_setup_tracer_arize(mock_provider, mock_exporter): - config = ArizeConfig(endpoint="http://a.com", api_key="k", space_id="s", project="p") - setup_tracer(config) - mock_exporter.assert_called_once() - assert mock_exporter.call_args[1]["endpoint"] == "http://a.com/v1" - - -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.HttpOTLPSpanExporter") -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.trace_sdk.TracerProvider") -def test_setup_tracer_phoenix(mock_provider, mock_exporter): - config = PhoenixConfig(endpoint="http://p.com", project="p") - setup_tracer(config) - mock_exporter.assert_called_once() - assert mock_exporter.call_args[1]["endpoint"] == "http://p.com/v1/traces" - - -def test_setup_tracer_exception(): - config = ArizeConfig(endpoint="http://a.com", project="p") - with patch("dify_trace_arize_phoenix.arize_phoenix_trace.urlparse", side_effect=Exception("boom")): - with pytest.raises(Exception, match="boom"): - setup_tracer(config) - - -# --- ArizePhoenixDataTrace Class Tests --- - - -@pytest.fixture -def span_exporter(): - return _CollectingSpanExporter() - - -@pytest.fixture -def trace_instance(span_exporter): - provider = trace_sdk.TracerProvider() - processor = SimpleSpanProcessor(span_exporter) - provider.add_span_processor(processor) - tracer = provider.get_tracer("test-phoenix") - try: - with ( - patch.object(tracer, "start_span", wraps=tracer.start_span), - patch("dify_trace_arize_phoenix.arize_phoenix_trace.setup_tracer", return_value=(tracer, processor)), - patch("dify_trace_arize_phoenix.arize_phoenix_trace.redis_client", new=MagicMock()) as mock_redis, - ): - config = ArizeConfig(endpoint="http://a.com", api_key="k", space_id="s", project="p") - instance = ArizePhoenixDataTrace(config) - cast(Any, instance)._mock_redis_client = mock_redis - yield instance - finally: - provider.shutdown() - - -def test_trace_dispatch(trace_instance): - with ( - patch.object(trace_instance, "workflow_trace") as m1, - patch.object(trace_instance, "message_trace") as m2, - patch.object(trace_instance, "moderation_trace") as m3, - patch.object(trace_instance, "suggested_question_trace") as m4, - patch.object(trace_instance, "dataset_retrieval_trace") as m5, - patch.object(trace_instance, "tool_trace") as m6, - patch.object(trace_instance, "generate_name_trace") as m7, - ): - trace_instance.trace(_make_workflow_info()) - m1.assert_called() - - trace_instance.trace(_make_message_info()) - m2.assert_called() - - trace_instance.trace(ModerationTraceInfo(flagged=True, action="a", preset_response="p", query="q", metadata={})) - m3.assert_called() - - trace_instance.trace(SuggestedQuestionTraceInfo(suggested_question=[], total_tokens=0, level="i", metadata={})) - m4.assert_called() - - trace_instance.trace(DatasetRetrievalTraceInfo(metadata={})) - m5.assert_called() - - trace_instance.trace( - ToolTraceInfo( - tool_name="t", - tool_inputs={}, - tool_outputs="o", - metadata={}, - tool_config={}, - time_cost=1, - tool_parameters={}, - ) - ) - m6.assert_called() - - trace_instance.trace(GenerateNameTraceInfo(tenant_id="t", metadata={})) - m7.assert_called() - - -def test_trace_exception(trace_instance): - with patch.object(trace_instance, "workflow_trace", side_effect=RuntimeError("fail")): - with pytest.raises(RuntimeError): - trace_instance.trace(_make_workflow_info()) - - -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_full(mock_repo_factory, trace_instance): - info = _make_workflow_info() - repo = MagicMock() - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - - node1 = MagicMock() - node1.node_type = "llm" - node1.status = "succeeded" - node1.inputs = {"q": "hi"} - node1.outputs = {"a": "bye", "usage": {"total_tokens": 5}} - node1.created_at = _dt() - node1.elapsed_time = 1.0 - node1.process_data = { - "prompts": [{"role": "user", "content": "hi"}], - "model_provider": "openai", - "model_name": "gpt-4", - } - node1.metadata = {"k": "v"} - node1.title = "title" - node1.id = "n1" - node1.error = None - - repo.get_by_workflow_execution.return_value = [node1] - - with patch.object(trace_instance, "get_service_account_with_tenant"): - trace_instance.workflow_trace(info) - - assert trace_instance.tracer.start_span.call_count >= 2 - - -def test_workflow_trace_no_app_id(trace_instance): - info = _make_workflow_info() - info.metadata = {} - with pytest.raises(ValueError, match="No app_id found in trace_info metadata"): - trace_instance.workflow_trace(info) - - -def test_workflow_trace_queries_real_repository_with_sqlite_session_factory( - trace_instance, sqlite3_session: Session -) -> None: - """Exercise the repository query through the session factory created by the trace.""" - - info = _make_workflow_info( - tenant_id="00000000-0000-0000-0000-000000000001", - workflow_id="00000000-0000-0000-0000-000000000002", - workflow_run_id="00000000-0000-0000-0000-000000000003", - metadata={"app_id": "00000000-0000-0000-0000-000000000004"}, - ) - service_account = MagicMock() - service_account.id = "00000000-0000-0000-0000-000000000005" - - with patch.object(trace_instance, "get_service_account_with_tenant", return_value=service_account): - trace_instance.workflow_trace(info) - - workflow_span_call = _get_start_span_call( - trace_instance.tracer.start_span, - span_name="workflow_00000000-0000-0000-0000-000000000003", - ) - assert workflow_span_call.kwargs["attributes"][SpanAttributes.SESSION_ID] == info.workflow_run_id - - -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_uses_canonical_root_context_for_top_level_workflow(mock_repo_factory, trace_instance): - info = _make_workflow_info( - message_id="message-1", - workflow_run_id="workflow-run-1", - metadata={"app_id": "app1", "trace_session_id": "trace-session-1"}, - ) - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [] - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - - root_carrier = {} - root_context = Context() - - with ( - patch.object(trace_instance, "get_service_account_with_tenant", return_value=MagicMock()), - patch.object(trace_instance, "ensure_root_span", return_value=root_carrier) as mock_ensure_root_span, - patch.object(trace_instance.propagator, "extract", return_value=root_context) as mock_extract, - ): - trace_instance.workflow_trace(info) - - mock_ensure_root_span.assert_called_once_with( - info.resolved_trace_id, - root_span_name="workflow-run-1", - root_span_attributes={ - SpanAttributes.INPUT_VALUE: safe_json_dumps(info.workflow_run_inputs), - SpanAttributes.INPUT_MIME_TYPE: "application/json", - SpanAttributes.OUTPUT_VALUE: safe_json_dumps(info.workflow_run_outputs), - SpanAttributes.OUTPUT_MIME_TYPE: "application/json", - SpanAttributes.SESSION_ID: "trace-session-1", - }, - ) - mock_extract.assert_called_once_with(carrier=root_carrier) - workflow_span_call = _get_start_span_call(trace_instance.tracer.start_span, span_name="workflow_workflow-run-1") - assert workflow_span_call.kwargs["context"] is root_context - - -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_uses_workflow_run_id_for_root_span_and_populates_root_inputs_outputs( - mock_repo_factory, - trace_instance, -): - info = _make_workflow_info( - workflow_run_inputs={"prompt": "hello"}, - workflow_run_outputs={"result": "world"}, - metadata={ - "app_id": "app1", - "app_name": "Workflow Name", - }, - workflow_run_id="workflow-run-xyz", - ) - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [] - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - - with patch.object(trace_instance, "get_service_account_with_tenant", return_value=MagicMock()): - trace_instance.workflow_trace(info) - - root_span_call = _get_start_span_call(trace_instance.tracer.start_span, span_name="workflow-run-xyz") - assert root_span_call.kwargs["attributes"][SpanAttributes.INPUT_VALUE] == safe_json_dumps(info.workflow_run_inputs) - assert root_span_call.kwargs["attributes"][SpanAttributes.OUTPUT_VALUE] == safe_json_dumps( - info.workflow_run_outputs - ) - assert root_span_call.kwargs["attributes"][SpanAttributes.INPUT_MIME_TYPE] == "application/json" - assert root_span_call.kwargs["attributes"][SpanAttributes.OUTPUT_MIME_TYPE] == "application/json" - - -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_propagates_workflow_error_to_root_span( - mock_repo_factory, - trace_instance, -): - info = _make_workflow_info( - workflow_run_status="failed", - error="Traceback (most recent call last): RuntimeError: workflow failed", - ) - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [] - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - - with ( - patch.object(trace_instance, "get_service_account_with_tenant", return_value=MagicMock()), - patch.object(trace_instance, "ensure_root_span", return_value={}) as mock_ensure_root_span, - ): - trace_instance.workflow_trace(info) - - mock_ensure_root_span.assert_called_once() - assert mock_ensure_root_span.call_args.kwargs["root_span_error"] == info.error - - -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_falls_back_to_dify_name_when_workflow_run_id_is_blank( - mock_repo_factory, - trace_instance, -): - info = _make_workflow_info( - metadata={ - "app_id": "app1", - "app_name": "", - }, - workflow_run_id="", - ) - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [] - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - - with patch.object(trace_instance, "get_service_account_with_tenant", return_value=MagicMock()): - trace_instance.workflow_trace(info) - - root_span_call = _get_start_span_call(trace_instance.tracer.start_span, span_name="Dify") - assert root_span_call.kwargs["attributes"]["dify_trace_id"] == "" - - -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_reuses_upstream_parent_workflow_context_when_no_parent_node_execution_id_is_available( - mock_repo_factory, trace_instance -): - info = _make_workflow_info( - message_id="message-1", - workflow_run_id="workflow-run-1", - metadata={ - "app_id": "app1", - "parent_trace_context": { - "parent_workflow_run_id": "outer-workflow-run-1", - }, - }, - ) - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [] - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - - parent_carrier = {} - parent_context = Context() - - with ( - patch.object(trace_instance, "get_service_account_with_tenant", return_value=MagicMock()), - patch.object(trace_instance, "ensure_root_span", return_value=parent_carrier) as mock_ensure_root_span, - patch.object(trace_instance.propagator, "extract", return_value=parent_context) as mock_extract, - ): - trace_instance.workflow_trace(info) - - mock_ensure_root_span.assert_called_once_with( - "outer-workflow-run-1", - root_span_name="workflow-run-1", - root_span_attributes={ - SpanAttributes.INPUT_VALUE: safe_json_dumps(info.workflow_run_inputs), - SpanAttributes.INPUT_MIME_TYPE: "application/json", - SpanAttributes.OUTPUT_VALUE: safe_json_dumps(info.workflow_run_outputs), - SpanAttributes.OUTPUT_MIME_TYPE: "application/json", - SpanAttributes.SESSION_ID: "outer-workflow-run-1", - }, - ) - mock_extract.assert_called_once_with(carrier=parent_carrier) - workflow_span_call = _get_start_span_call(trace_instance.tracer.start_span, span_name="workflow_workflow-run-1") - assert workflow_span_call.kwargs["context"] is parent_context - - -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_uses_published_parent_node_context_for_nested_workflow( - mock_repo_factory, - trace_instance, -): - info = _make_workflow_info( - message_id="message-1", - workflow_run_id="workflow-run-1", - metadata={ - "app_id": "app1", - "parent_trace_context": { - "parent_workflow_run_id": "outer-workflow-run-1", - "parent_node_execution_id": "outer-node-execution-1", - }, - }, - ) - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [] - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - stored_carrier = '{"traceparent":"00-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-bbbbbbbbbbbbbbbb-01"}' - trace_instance._mock_redis_client.get.return_value = stored_carrier - parent_context = Context() - - with ( - patch.object(trace_instance, "get_service_account_with_tenant", return_value=MagicMock()), - patch.object(trace_instance, "ensure_root_span") as mock_ensure_root_span, - patch.object(trace_instance.propagator, "extract", return_value=parent_context) as mock_extract, - ): - trace_instance.workflow_trace(info) - - trace_instance._mock_redis_client.get.assert_called_once_with( - _phoenix_parent_span_redis_key("outer-node-execution-1") - ) - mock_ensure_root_span.assert_not_called() - mock_extract.assert_called_once_with( - carrier={"traceparent": "00-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-bbbbbbbbbbbbbbbb-01"} - ) - workflow_span_call = _get_start_span_call(trace_instance.tracer.start_span, span_name="workflow_workflow-run-1") - assert workflow_span_call.kwargs["context"] is parent_context - - -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_raises_pending_parent_error_when_parent_node_context_is_missing( - mock_repo_factory, - trace_instance, -): - info = _make_workflow_info( - message_id="message-1", - workflow_run_id="workflow-run-1", - metadata={ - "app_id": "app1", - "parent_trace_context": { - "parent_workflow_run_id": "outer-workflow-run-1", - "parent_node_execution_id": "outer-node-execution-1", - }, - }, - ) - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [] - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - trace_instance._mock_redis_client.get.return_value = None - - with ( - patch.object(trace_instance, "get_service_account_with_tenant", return_value=MagicMock()), - patch( - "dify_trace_arize_phoenix.arize_phoenix_trace._parent_workflow_can_publish_span_context", - return_value=True, - ), - patch.object(trace_instance, "ensure_root_span") as mock_ensure_root_span, - pytest.raises(PendingTraceParentContextError) as exc_info, - ): - trace_instance.workflow_trace(info) - - assert exc_info.value.parent_node_execution_id == "outer-node-execution-1" - trace_instance._mock_redis_client.get.assert_called_once_with( - _phoenix_parent_span_redis_key("outer-node-execution-1") - ) - mock_ensure_root_span.assert_not_called() - - -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_falls_back_when_parent_app_tracing_cannot_publish_parent_context( - mock_repo_factory, - trace_instance, -): - info = _make_workflow_info( - message_id="message-1", - workflow_run_id="workflow-run-1", - metadata={ - "app_id": "app1", - "parent_trace_context": { - "parent_workflow_run_id": "outer-workflow-run-1", - "parent_node_execution_id": "outer-node-execution-1", - }, - }, - ) - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [] - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - trace_instance._mock_redis_client.get.return_value = None - - parent_carrier = {} - parent_context = Context() - - with ( - patch.object(trace_instance, "get_service_account_with_tenant", return_value=MagicMock()), - patch( - "dify_trace_arize_phoenix.arize_phoenix_trace._parent_workflow_can_publish_span_context", - return_value=False, - ), - patch.object(trace_instance, "ensure_root_span", return_value=parent_carrier) as mock_ensure_root_span, - patch.object(trace_instance.propagator, "extract", return_value=parent_context) as mock_extract, - ): - trace_instance.workflow_trace(info) - - mock_ensure_root_span.assert_called_once_with( - "outer-workflow-run-1", - root_span_name="workflow-run-1", - root_span_attributes={ - SpanAttributes.INPUT_VALUE: safe_json_dumps(info.workflow_run_inputs), - SpanAttributes.INPUT_MIME_TYPE: "application/json", - SpanAttributes.OUTPUT_VALUE: safe_json_dumps(info.workflow_run_outputs), - SpanAttributes.OUTPUT_MIME_TYPE: "application/json", - SpanAttributes.SESSION_ID: "outer-workflow-run-1", - }, - ) - mock_extract.assert_called_once_with(carrier=parent_carrier) - workflow_span_call = _get_start_span_call(trace_instance.tracer.start_span, span_name="workflow_workflow-run-1") - assert workflow_span_call.kwargs["context"] is parent_context - - -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_still_retries_when_parent_app_can_publish_parent_context( - mock_repo_factory, - trace_instance, -): - info = _make_workflow_info( - message_id="message-1", - workflow_run_id="workflow-run-1", - metadata={ - "app_id": "app1", - "parent_trace_context": { - "parent_workflow_run_id": "outer-workflow-run-1", - "parent_node_execution_id": "outer-node-execution-1", - }, - }, - ) - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [] - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - trace_instance._mock_redis_client.get.return_value = None - - with ( - patch.object(trace_instance, "get_service_account_with_tenant", return_value=MagicMock()), - patch( - "dify_trace_arize_phoenix.arize_phoenix_trace._parent_workflow_can_publish_span_context", - return_value=True, - ), - patch.object(trace_instance, "ensure_root_span") as mock_ensure_root_span, - pytest.raises(PendingTraceParentContextError) as exc_info, - ): - trace_instance.workflow_trace(info) - - assert exc_info.value.parent_node_execution_id == "outer-node-execution-1" - mock_ensure_root_span.assert_not_called() - - -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_uses_parent_workflow_run_id_for_workflow_and_nodes_when_nested_context_is_present( - mock_repo_factory, trace_instance -): - info = _make_workflow_info( - conversation_id=None, - metadata={ - "app_id": "app1", - "parent_trace_context": { - "parent_workflow_run_id": "outer-workflow-run-1", - }, - }, - ) - repo = MagicMock() - node_execution = MagicMock() - node_execution.node_type = "tool" - node_execution.status = "succeeded" - node_execution.inputs = {"tool_input": "value"} - node_execution.outputs = {"tool_output": "value"} - node_execution.created_at = _dt() - node_execution.elapsed_time = 1.0 - node_execution.process_data = {} - node_execution.metadata = {} - node_execution.title = "Tool node" - node_execution.id = "node-1" - node_execution.error = None - repo.get_by_workflow_execution.return_value = [node_execution] - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - - with patch.object(trace_instance, "get_service_account_with_tenant", return_value=MagicMock()): - trace_instance.workflow_trace(info) - - workflow_span_call = _get_start_span_call(trace_instance.tracer.start_span, span_name="workflow_r1") - node_span_call = _get_start_span_call(trace_instance.tracer.start_span, span_name="tool_Tool node") - - assert workflow_span_call.kwargs["attributes"][SpanAttributes.SESSION_ID] == "outer-workflow-run-1" - assert node_span_call.kwargs["attributes"][SpanAttributes.SESSION_ID] == "outer-workflow-run-1" - - -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_falls_back_to_node_type_when_node_title_is_blank(mock_repo_factory, trace_instance): - info = _make_workflow_info() - repo = MagicMock() - node_execution = _make_node_execution( - id="node-execution-1", - node_execution_id="node-execution-1", - node_id="node-1", - node_type="tool", - title=" ", - ) - repo.get_by_workflow_execution.return_value = [node_execution] - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - - with patch.object(trace_instance, "get_service_account_with_tenant", return_value=MagicMock()): - trace_instance.workflow_trace(info) - - node_span_call = _get_start_span_call(trace_instance.tracer.start_span, span_name="tool") - assert node_span_call.kwargs["attributes"][SpanAttributes.SESSION_ID] == "r1" - - -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_prefers_workflow_graph_node_title_over_execution_title(mock_repo_factory, trace_instance): - info = _make_workflow_info( - workflow_data={ - "graph": { - "nodes": [ - { - "id": "nested-tool-node", - "data": { - "type": "tool", - "title": "nested workflow tool", - }, - } - ] - } - } - ) - repo = MagicMock() - node_execution = _make_node_execution( - id="node-execution-1", - node_execution_id="node-execution-1", - node_id="nested-tool-node", - node_type="tool", - title="2", - ) - repo.get_by_workflow_execution.return_value = [node_execution] - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - - with patch.object(trace_instance, "get_service_account_with_tenant", return_value=MagicMock()): - trace_instance.workflow_trace(info) - - node_span_call = _get_start_span_call(trace_instance.tracer.start_span, span_name="tool_nested workflow tool") - assert node_span_call.kwargs["attributes"][SpanAttributes.SESSION_ID] == "r1" - - -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_keeps_nested_conversation_session_while_reusing_parent_root_context( - mock_repo_factory, trace_instance -): - info = _make_workflow_info( - conversation_id="conversation-1", - message_id="message-1", - workflow_run_id="workflow-run-1", - metadata={ - "app_id": "app1", - "parent_trace_context": { - "parent_workflow_run_id": "outer-workflow-run-1", - }, - }, - ) - repo = MagicMock() - node_execution = _make_node_execution( - id="node-execution-1", - node_execution_id="node-execution-1", - node_id="node-1", - node_type="tool", - ) - repo.get_by_workflow_execution.return_value = [node_execution] - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - - parent_carrier = {} - parent_context = Context() - - with ( - patch.object(trace_instance, "get_service_account_with_tenant", return_value=MagicMock()), - patch.object(trace_instance, "ensure_root_span", return_value=parent_carrier) as mock_ensure_root_span, - patch.object(trace_instance.propagator, "extract", return_value=parent_context) as mock_extract, - ): - trace_instance.workflow_trace(info) - - mock_ensure_root_span.assert_called_once_with( - "outer-workflow-run-1", - root_span_name="workflow-run-1", - root_span_attributes={ - SpanAttributes.INPUT_VALUE: safe_json_dumps(info.workflow_run_inputs), - SpanAttributes.INPUT_MIME_TYPE: "application/json", - SpanAttributes.OUTPUT_VALUE: safe_json_dumps(info.workflow_run_outputs), - SpanAttributes.OUTPUT_MIME_TYPE: "application/json", - SpanAttributes.SESSION_ID: "conversation-1", - }, - ) - mock_extract.assert_called_once_with(carrier=parent_carrier) - workflow_span_call = _get_start_span_call(trace_instance.tracer.start_span, span_name="workflow_workflow-run-1") - node_span_call = _get_start_span_call(trace_instance.tracer.start_span, span_name="tool_Node") - assert workflow_span_call.kwargs["context"] is parent_context - assert workflow_span_call.kwargs["attributes"][SpanAttributes.SESSION_ID] == "conversation-1" - assert node_span_call.kwargs["attributes"][SpanAttributes.SESSION_ID] == "conversation-1" - - -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_publishes_tool_node_parent_span_context_to_redis( - mock_repo_factory, - trace_instance, -): - info = _make_workflow_info() - repo = MagicMock() - node_execution = _make_node_execution( - id="tool-execution-1", - node_execution_id="tool-execution-1", - node_id="tool-node-1", - node_type="tool", - ) - repo.get_by_workflow_execution.return_value = [node_execution] - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - - workflow_span = MagicMock(name="workflow-span") - workflow_span._context_label = "workflow" - tool_span = MagicMock(name="tool-span") - tool_span._context_label = "tool" - trace_instance.tracer.start_span.side_effect = [workflow_span, tool_span] - - def inject_side_effect(carrier): - carrier["traceparent"] = "00-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-bbbbbbbbbbbbbbbb-01" - - with ( - patch.object(trace_instance, "get_service_account_with_tenant", return_value=MagicMock()), - patch.object(trace_instance, "ensure_root_span", return_value={}), - patch.object(trace_instance.propagator, "extract", return_value="root-context"), - patch.object(trace_instance.propagator, "inject", side_effect=inject_side_effect) as mock_inject, - patch( - "dify_trace_arize_phoenix.arize_phoenix_trace.set_span_in_context", - side_effect=lambda span: f"context:{span._context_label}", - ), - ): - trace_instance.workflow_trace(info) - - mock_inject.assert_called_once() - trace_instance._mock_redis_client.setex.assert_called_once_with( - _phoenix_parent_span_redis_key("tool-execution-1"), - 300, - '{"traceparent": "00-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-bbbbbbbbbbbbbbbb-01"}', - ) - - -@pytest.mark.parametrize( - ("failing_step", "expected_message"), - [ - ("inject", "inject failed"), - ("publish", "publish failed"), - ], -) -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_cleans_up_tool_span_when_parent_context_publish_fails( - mock_repo_factory, - trace_instance, - failing_step, - expected_message, -): - info = _make_workflow_info() - repo = MagicMock() - node_execution = _make_node_execution( - id="tool-execution-1", - node_execution_id="tool-execution-1", - node_id="tool-node-1", - node_type="tool", - ) - repo.get_by_workflow_execution.return_value = [node_execution] - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - - workflow_span = MagicMock(name="workflow-span") - workflow_span._context_label = "workflow" - tool_span = MagicMock(name="tool-span") - tool_span._context_label = "tool" - trace_instance.tracer.start_span.side_effect = [workflow_span, tool_span] - - inject_side_effect = None - if failing_step == "inject": - inject_side_effect = RuntimeError(expected_message) - else: - trace_instance._mock_redis_client.setex.side_effect = RuntimeError(expected_message) - - def inject_side_effect(carrier): - carrier["traceparent"] = "00-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-bbbbbbbbbbbbbbbb-01" - - with ( - patch.object(trace_instance, "get_service_account_with_tenant", return_value=MagicMock()), - patch.object(trace_instance, "ensure_root_span", return_value={}), - patch.object(trace_instance.propagator, "extract", return_value="root-context"), - patch.object(trace_instance.propagator, "inject", side_effect=inject_side_effect), - patch( - "dify_trace_arize_phoenix.arize_phoenix_trace.set_span_in_context", - side_effect=lambda span: f"context:{span._context_label}", - ), - pytest.raises(RuntimeError, match=expected_message), - ): - trace_instance.workflow_trace(info) - - tool_span.end.assert_called_once() - workflow_span.end.assert_called_once() - - -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_parents_serial_nodes_to_resolved_predecessor_span(mock_repo_factory, trace_instance): - info = _make_workflow_info() - repo = MagicMock() - second_node = _make_node_execution( - id="node-execution-2", - node_execution_id="node-execution-2", - node_id="node-2", - node_type="llm", - predecessor_node_id="node-1", - process_data={ - "prompts": [{"role": "user", "content": "hi"}], - "model_provider": "openai", - "model_name": "gpt-4", - }, - ) - first_node = _make_node_execution( - id="node-execution-1", - node_execution_id="node-execution-1", - node_id="node-1", - node_type="tool", - ) - repo.get_by_workflow_execution.return_value = [second_node, first_node] - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - - workflow_span = MagicMock(name="workflow-span") - workflow_span._context_label = "workflow" - first_node_span = MagicMock(name="first-node-span") - first_node_span._context_label = "node-1" - second_node_span = MagicMock(name="second-node-span") - second_node_span._context_label = "node-2" - trace_instance.tracer.start_span.side_effect = [workflow_span, first_node_span, second_node_span] - - with ( - patch.object(trace_instance, "get_service_account_with_tenant", return_value=MagicMock()), - patch.object(trace_instance, "ensure_root_span", return_value={}), - patch.object(trace_instance.propagator, "extract", return_value="root-context"), - patch( - "dify_trace_arize_phoenix.arize_phoenix_trace.set_span_in_context", - side_effect=lambda span: f"context:{span._context_label}", - ), - ): - trace_instance.workflow_trace(info) - - first_node_call = _get_start_span_call(trace_instance.tracer.start_span, span_name="tool_Node") - second_node_call = _get_start_span_call(trace_instance.tracer.start_span, span_name="llm_Node") - assert first_node_call.kwargs["context"] == "context:workflow" - assert second_node_call.kwargs["context"] == "context:node-1" - - -@pytest.mark.parametrize( - ("enclosing_node_type", "structured_field"), - [ - ("loop", "loop_id"), - ("iteration", "iteration_id"), - ], -) -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_parents_structured_start_nodes_to_enclosing_structure_span( - mock_repo_factory, - trace_instance, - enclosing_node_type, - structured_field, -): - info = _make_workflow_info() - repo = MagicMock() - enclosing_node = _make_node_execution( - id=f"{enclosing_node_type}-execution-1", - node_execution_id=f"{enclosing_node_type}-execution-1", - node_id=f"{enclosing_node_type}-node-1", - node_type=enclosing_node_type, - ) - structured_kwargs = {structured_field: f"{enclosing_node_type}-node-1"} - start_node = _make_node_execution( - id="start-execution-1", - node_execution_id="start-execution-1", - node_id="start-node-1", - node_type="start", - **structured_kwargs, - ) - repo.get_by_workflow_execution.return_value = [start_node, enclosing_node] - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - - workflow_span = MagicMock(name="workflow-span") - workflow_span._context_label = "workflow" - enclosing_node_span = MagicMock(name="enclosing-node-span") - enclosing_node_span._context_label = enclosing_node_type - start_node_span = MagicMock(name="start-node-span") - start_node_span._context_label = "start" - trace_instance.tracer.start_span.side_effect = [workflow_span, enclosing_node_span, start_node_span] - - with ( - patch.object(trace_instance, "get_service_account_with_tenant", return_value=MagicMock()), - patch.object(trace_instance, "ensure_root_span", return_value={}), - patch.object(trace_instance.propagator, "extract", return_value="root-context"), - patch( - "dify_trace_arize_phoenix.arize_phoenix_trace.set_span_in_context", - side_effect=lambda span: f"context:{span._context_label}", - ), - ): - trace_instance.workflow_trace(info) - - start_node_call = _get_start_span_call(trace_instance.tracer.start_span, span_name="start_Node") - assert start_node_call.kwargs["context"] == f"context:{enclosing_node_type}" - - -@pytest.mark.parametrize( - ("enclosing_node_type", "structured_field"), - [ - ("loop", "loop_id"), - ("iteration", "iteration_id"), - ], -) -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_keeps_duplicate_body_node_children_under_enclosing_structure( - mock_repo_factory, - trace_instance, - enclosing_node_type, - structured_field, -): - info = _make_workflow_info() - repo = MagicMock() - enclosing_node = _make_node_execution( - id=f"{enclosing_node_type}-execution-1", - node_execution_id=f"{enclosing_node_type}-execution-1", - node_id=f"{enclosing_node_type}-node-1", - node_type=enclosing_node_type, - ) - structured_kwargs = {structured_field: f"{enclosing_node_type}-node-1"} - repeated_body_node_1 = _make_node_execution( - id="body-execution-1", - node_execution_id="body-execution-1", - node_id="body-node-1", - node_type="tool", - **structured_kwargs, - ) - repeated_body_node_2 = _make_node_execution( - id="body-execution-2", - node_execution_id="body-execution-2", - node_id="body-node-1", - node_type="tool", - **structured_kwargs, - ) - child_node = _make_node_execution( - id="child-execution-1", - node_execution_id="child-execution-1", - node_id="child-node-1", - node_type="llm", - predecessor_node_id="body-node-1", - process_data={ - "prompts": [{"role": "user", "content": "hi"}], - "model_provider": "openai", - "model_name": "gpt-4", - }, - **structured_kwargs, - ) - repo.get_by_workflow_execution.return_value = [ - child_node, - repeated_body_node_1, - repeated_body_node_2, - enclosing_node, - ] - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - - workflow_span = MagicMock(name="workflow-span") - workflow_span._context_label = "workflow" - enclosing_node_span = MagicMock(name="enclosing-node-span") - enclosing_node_span._context_label = enclosing_node_type - child_node_span = MagicMock(name="child-node-span") - child_node_span._context_label = "child" - repeated_body_node_1_span = MagicMock(name="repeated-body-node-1-span") - repeated_body_node_1_span._context_label = "body-1" - repeated_body_node_2_span = MagicMock(name="repeated-body-node-2-span") - repeated_body_node_2_span._context_label = "body-2" - trace_instance.tracer.start_span.side_effect = [ - workflow_span, - enclosing_node_span, - child_node_span, - repeated_body_node_1_span, - repeated_body_node_2_span, - ] - - with ( - patch.object(trace_instance, "get_service_account_with_tenant", return_value=MagicMock()), - patch.object(trace_instance, "ensure_root_span", return_value={}), - patch.object(trace_instance.propagator, "extract", return_value="root-context"), - patch( - "dify_trace_arize_phoenix.arize_phoenix_trace.set_span_in_context", - side_effect=lambda span: f"context:{span._context_label}", - ), - ): - trace_instance.workflow_trace(info) - - child_node_call = _get_start_span_call(trace_instance.tracer.start_span, span_name="llm_Node") - assert child_node_call.kwargs["context"] == f"context:{enclosing_node_type}" - - -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_records_exception_node_event_without_failing_root_span(mock_repo_factory, trace_instance): - info = _make_workflow_info(workflow_run_status="succeeded", error=None) - repo = MagicMock() - handled_error_node = _make_node_execution( - id="code-execution-1", - node_execution_id="code-execution-1", - node_id="code-node-1", - node_type="code", - title="Code with error handling", - status=WorkflowNodeExecutionStatus.EXCEPTION, - error="Traceback (most recent call last): RuntimeError: handled failure", - metadata={"error_strategy": "fail-branch"}, - ) - repo.get_by_workflow_execution.return_value = [handled_error_node] - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - - root_span = MagicMock(name="root-span") - workflow_span = MagicMock(name="workflow-span") - node_span = MagicMock(name="node-span") - trace_instance.tracer.start_span.side_effect = [root_span, workflow_span, node_span] - - with ( - patch.object(trace_instance, "get_service_account_with_tenant", return_value=MagicMock()), - patch.object(trace_instance.propagator, "extract", return_value="root-context"), - ): - trace_instance.workflow_trace(info) - - root_span.set_status.assert_called_once() - assert root_span.set_status.call_args.args[0].status_code == StatusCode.OK - workflow_span.set_status.assert_called_once() - assert workflow_span.set_status.call_args.args[0].status_code == StatusCode.OK - node_span.set_status.assert_called_once() - assert node_span.set_status.call_args.args[0].status_code == StatusCode.ERROR - node_span.add_event.assert_called_once() - assert node_span.add_event.call_args.kwargs["name"] == "exception" - assert node_span.add_event.call_args.kwargs["attributes"][OTELSpanAttributes.EXCEPTION_MESSAGE] == ( - "Traceback (most recent call last): RuntimeError: handled failure" - ) - - -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_groups_loop_iteration_children_under_wrapper_spans(mock_repo_factory, trace_instance): - info = _make_workflow_info(conversation_id="conversation-1") - repo = MagicMock() - loop_node = _make_node_execution( - id="loop-execution-1", - node_execution_id="loop-execution-1", - node_id="loop-node-1", - node_type="loop", - title="Main Loop", - ) - first_body_node = _make_node_execution( - id="body-execution-1", - node_execution_id="body-execution-1", - node_id="body-node-1", - node_type="tool", - loop_id="loop-node-1", - metadata={"loop_id": "loop-node-1", "loop_index": 0}, - ) - second_body_node = _make_node_execution( - id="body-execution-2", - node_execution_id="body-execution-2", - node_id="body-node-1", - node_type="tool", - loop_id="loop-node-1", - metadata={"loop_id": "loop-node-1", "loop_index": 1}, - ) - repo.get_by_workflow_execution.return_value = [loop_node, first_body_node, second_body_node] - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - - workflow_span = MagicMock(name="workflow-span") - workflow_span._context_label = "workflow" - loop_span = MagicMock(name="loop-span") - loop_span._context_label = "loop" - first_wrapper_span = MagicMock(name="loop-wrapper-0") - first_wrapper_span._context_label = "loop-wrapper-0" - first_body_span = MagicMock(name="body-span-1") - first_body_span._context_label = "body-1" - second_wrapper_span = MagicMock(name="loop-wrapper-1") - second_wrapper_span._context_label = "loop-wrapper-1" - second_body_span = MagicMock(name="body-span-2") - second_body_span._context_label = "body-2" - trace_instance.tracer.start_span.side_effect = [ - workflow_span, - loop_span, - first_wrapper_span, - first_body_span, - second_wrapper_span, - second_body_span, - ] - - with ( - patch.object(trace_instance, "get_service_account_with_tenant", return_value=MagicMock()), - patch.object(trace_instance, "ensure_root_span", return_value={}), - patch.object(trace_instance.propagator, "extract", return_value="root-context"), - patch( - "dify_trace_arize_phoenix.arize_phoenix_trace.set_span_in_context", - side_effect=lambda span: f"context:{span._context_label}", - ), - ): - trace_instance.workflow_trace(info) - - first_wrapper_call = _get_start_span_call(trace_instance.tracer.start_span, span_name="loop[0]") - second_wrapper_call = _get_start_span_call(trace_instance.tracer.start_span, span_name="loop[1]") - first_body_call = trace_instance.tracer.start_span.call_args_list[3] - second_body_call = trace_instance.tracer.start_span.call_args_list[5] - - assert first_wrapper_call.kwargs["context"] == "context:loop" - assert second_wrapper_call.kwargs["context"] == "context:loop" - assert first_wrapper_call.kwargs["attributes"]["dify.wrapper.synthetic"] is True - assert first_wrapper_call.kwargs["attributes"]["dify.wrapper.index"] == "0" - assert first_body_call.kwargs["context"] == "context:loop-wrapper-0" - assert second_body_call.kwargs["context"] == "context:loop-wrapper-1" - assert first_body_call.kwargs["attributes"][SpanAttributes.SESSION_ID] == "conversation-1" - assert json.loads(first_body_call.kwargs["attributes"][SpanAttributes.METADATA])["loop_index"] == 0 - assert first_body_call.kwargs["attributes"]["dify.node.loop_index"] == 0 - - -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_finalizes_loop_wrapper_with_child_time_bounds_and_error_status( - mock_repo_factory, trace_instance -): - info = _make_workflow_info() - repo = MagicMock() - loop_node = _make_node_execution( - id="loop-execution-1", - node_execution_id="loop-execution-1", - node_id="loop-node-1", - node_type="loop", - ) - first_body_node = _make_node_execution( - id="body-execution-1", - node_execution_id="body-execution-1", - node_id="body-node-1", - node_type="tool", - created_at=_dt() + timedelta(seconds=3), - elapsed_time=2.0, - loop_id="loop-node-1", - metadata={"loop_id": "loop-node-1", "loop_index": 0}, - ) - failed_body_node = _make_node_execution( - id="body-execution-2", - node_execution_id="body-execution-2", - node_id="body-node-2", - node_type="tool", - status=WorkflowNodeExecutionStatus.FAILED, - error="body failed", - created_at=_dt() + timedelta(seconds=4), - elapsed_time=5.0, - loop_id="loop-node-1", - metadata={"loop_id": "loop-node-1", "loop_index": 0}, - ) - repo.get_by_workflow_execution.return_value = [loop_node, first_body_node, failed_body_node] - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - - workflow_span = MagicMock(name="workflow-span") - workflow_span._context_label = "workflow" - loop_span = MagicMock(name="loop-span") - loop_span._context_label = "loop" - wrapper_span = MagicMock(name="loop-wrapper-0") - wrapper_span._context_label = "loop-wrapper-0" - first_body_span = MagicMock(name="body-span-1") - first_body_span._context_label = "body-1" - failed_body_span = MagicMock(name="body-span-2") - failed_body_span._context_label = "body-2" - trace_instance.tracer.start_span.side_effect = [ - workflow_span, - loop_span, - wrapper_span, - first_body_span, - failed_body_span, - ] - - with ( - patch.object(trace_instance, "get_service_account_with_tenant", return_value=MagicMock()), - patch.object(trace_instance, "ensure_root_span", return_value={}), - patch.object(trace_instance.propagator, "extract", return_value="root-context"), - patch( - "dify_trace_arize_phoenix.arize_phoenix_trace.set_span_in_context", - side_effect=lambda span: f"context:{span._context_label}", - ), - ): - trace_instance.workflow_trace(info) - - wrapper_call = _get_start_span_call(trace_instance.tracer.start_span, span_name="loop[0]") - assert wrapper_call.kwargs["start_time"] == datetime_to_nanos(first_body_node.created_at) - wrapper_span.end.assert_called_once_with( - end_time=datetime_to_nanos(failed_body_node.created_at + timedelta(seconds=failed_body_node.elapsed_time)) - ) - wrapper_span.set_status.assert_called_once() - assert wrapper_span.set_status.call_args.args[0].status_code == StatusCode.ERROR - - -@patch("dify_trace_arize_phoenix.arize_phoenix_trace.DifyCoreRepositoryFactory") -def test_workflow_trace_falls_back_to_workflow_span_for_parallel_like_ambiguous_predecessors( - mock_repo_factory, trace_instance -): - info = _make_workflow_info() - repo = MagicMock() - child_node = _make_node_execution( - id="child-execution-1", - node_execution_id="child-execution-1", - node_id="child-node-1", - node_type="llm", - predecessor_node_id="parallel-node-1", - process_data={ - "prompts": [{"role": "user", "content": "hi"}], - "model_provider": "openai", - "model_name": "gpt-4", - }, - ) - first_parallel_node = _make_node_execution( - id="parallel-execution-1", - node_execution_id="parallel-execution-1", - node_id="parallel-node-1", - node_type="tool", - parallel_id="parallel-1", - ) - second_parallel_node = _make_node_execution( - id="parallel-execution-2", - node_execution_id="parallel-execution-2", - node_id="parallel-node-1", - node_type="tool", - parallel_id="parallel-2", - ) - repo.get_by_workflow_execution.return_value = [child_node, first_parallel_node, second_parallel_node] - mock_repo_factory.create_workflow_node_execution_repository.return_value = repo - - workflow_span = MagicMock(name="workflow-span") - workflow_span._context_label = "workflow" - child_node_span = MagicMock(name="child-node-span") - child_node_span._context_label = "child" - first_parallel_node_span = MagicMock(name="first-parallel-node-span") - first_parallel_node_span._context_label = "parallel-1" - second_parallel_node_span = MagicMock(name="second-parallel-node-span") - second_parallel_node_span._context_label = "parallel-2" - trace_instance.tracer.start_span.side_effect = [ - workflow_span, - child_node_span, - first_parallel_node_span, - second_parallel_node_span, - ] - - with ( - patch.object(trace_instance, "get_service_account_with_tenant", return_value=MagicMock()), - patch.object(trace_instance, "ensure_root_span", return_value={}), - patch.object(trace_instance.propagator, "extract", return_value="root-context"), - patch( - "dify_trace_arize_phoenix.arize_phoenix_trace.set_span_in_context", - side_effect=lambda span: f"context:{span._context_label}", - ), - ): - trace_instance.workflow_trace(info) - - child_node_call = _get_start_span_call(trace_instance.tracer.start_span, span_name="llm_Node") - assert child_node_call.kwargs["context"] == "context:workflow" - - -def test_message_trace_keeps_conversation_id_as_session(trace_instance): - info = _make_message_info() - info.message_data = MagicMock() - info.message_data.conversation_id = "conversation-2" - info.message_data.from_account_id = "acc2" - info.message_data.from_end_user_id = None - info.message_data.query = "q2" - info.message_data.answer = "a2" - info.message_data.status = "s2" - info.message_data.model_id = "m2" - info.message_data.model_provider = "p2" - info.message_data.message_metadata = "{}" - info.message_data.error = None - info.error = None - - root_span = MagicMock() - message_span = MagicMock() - llm_span = MagicMock() - trace_instance.tracer.start_span.side_effect = [root_span, message_span, llm_span] - - trace_instance.message_trace(info) - - message_span_call = _get_start_span_call( - trace_instance.tracer.start_span, span_name=TraceTaskName.MESSAGE_TRACE.value - ) - assert message_span_call.kwargs["attributes"][SpanAttributes.SESSION_ID] == "conversation-2" - - -def test_message_trace_uses_trace_session_id_metadata_as_session(trace_instance): - info = _make_message_info(metadata={"app_id": "app-1", "trace_session_id": "session-1"}) - info.message_data = MagicMock() - info.message_data.conversation_id = "conversation-2" - info.message_data.from_account_id = "acc2" - info.message_data.from_end_user_id = None - info.message_data.query = "q2" - info.message_data.answer = "a2" - info.message_data.status = "s2" - info.message_data.model_id = "m2" - info.message_data.model_provider = "p2" - info.message_data.message_metadata = "{}" - info.message_data.error = None - info.error = None - - root_span = MagicMock() - message_span = MagicMock() - llm_span = MagicMock() - trace_instance.tracer.start_span.side_effect = [root_span, message_span, llm_span] - - trace_instance.message_trace(info) - - message_span_call = _get_start_span_call( - trace_instance.tracer.start_span, span_name=TraceTaskName.MESSAGE_TRACE.value - ) - llm_span_call = _get_start_span_call_by_kind( - trace_instance.tracer.start_span, span_kind=OpenInferenceSpanKindValues.LLM.value - ) - assert message_span_call.kwargs["attributes"][SpanAttributes.SESSION_ID] == "session-1" - assert llm_span_call.kwargs["attributes"][SpanAttributes.SESSION_ID] == "session-1" - - -def test_message_trace_with_error(trace_instance): - info = _make_message_info() - info.message_data = MagicMock() - info.message_data.from_account_id = "acc1" - info.message_data.from_end_user_id = None - info.message_data.query = "q" - info.message_data.answer = "a" - info.message_data.status = "s" - info.message_data.model_id = "m" - info.message_data.model_provider = "p" - info.message_data.message_metadata = "{}" - info.message_data.error = "processing failed" - info.error = "message error" - - trace_instance.message_trace(info) - assert trace_instance.tracer.start_span.call_count >= 1 - - -def test_trace_methods_return_early_with_no_message_data(trace_instance): - info = MagicMock() - info.message_data = None - - trace_instance.moderation_trace(info) - trace_instance.suggested_question_trace(info) - trace_instance.dataset_retrieval_trace(info) - trace_instance.tool_trace(info) - trace_instance.generate_name_trace(info) - - assert trace_instance.tracer.start_span.call_count == 0 - - -def test_moderation_trace_ok(trace_instance): - info = ModerationTraceInfo(flagged=True, action="a", preset_response="p", query="q", metadata={}) - info.message_data = MagicMock() - info.message_data.error = None - trace_instance.moderation_trace(info) - # root span (1) + moderation span (1) = 2 - assert trace_instance.tracer.start_span.call_count >= 1 - - -def test_suggested_question_trace_ok(trace_instance): - info = SuggestedQuestionTraceInfo(suggested_question=["?"], total_tokens=1, level="i", metadata={}) - info.message_data = MagicMock() - info.error = None - trace_instance.suggested_question_trace(info) - assert trace_instance.tracer.start_span.call_count >= 1 - - -def test_dataset_retrieval_trace_ok(trace_instance): - info = DatasetRetrievalTraceInfo(documents=[], metadata={}) - info.message_data = MagicMock() - info.error = None - trace_instance.dataset_retrieval_trace(info) - assert trace_instance.tracer.start_span.call_count >= 1 - - -def test_tool_trace_ok(trace_instance): - info = ToolTraceInfo( - tool_name="t", tool_inputs={}, tool_outputs="o", metadata={}, tool_config={}, time_cost=1, tool_parameters={} - ) - info.message_data = MagicMock() - info.error = None - trace_instance.tool_trace(info) - assert trace_instance.tracer.start_span.call_count >= 1 - - -def test_generate_name_trace_ok(trace_instance): - info = GenerateNameTraceInfo(tenant_id="t", metadata={}) - info.message_data = MagicMock() - info.message_data.error = None - trace_instance.generate_name_trace(info) - assert trace_instance.tracer.start_span.call_count >= 1 - - -def test_get_project_url_phoenix(trace_instance): - trace_instance.arize_phoenix_config = PhoenixConfig(endpoint="http://p.com", project="p") - assert "p.com/projects/" in trace_instance.get_project_url() - - -def test_set_attribute_none_logic(trace_instance): - # Test role can be None - attrs = trace_instance._construct_llm_attributes([{"role": None, "content": "hi"}]) - assert "llm.input_messages.0.message.role" not in attrs - - # Test tool call id can be None - tool_call_none_id = {"id": None, "function": {"name": "f1"}} - attrs = trace_instance._construct_llm_attributes([{"role": "assistant", "tool_calls": [tool_call_none_id]}]) - assert "llm.input_messages.0.message.tool_calls.0.tool_call.id" not in attrs - - -def test_construct_llm_attributes_dict_branch(trace_instance): - attrs = trace_instance._construct_llm_attributes({"prompt": "hi"}) - assert '"prompt": "hi"' in attrs["llm.input_messages.0.message.content"] - assert attrs["llm.input_messages.0.message.role"] == "user" - - -def test_phoenix_api_check_success(trace_instance): - trace_instance.arize_phoenix_config = PhoenixConfig(project="test", endpoint="https://phoenix.test") - assert trace_instance.api_check() is True - - -def test_ensure_root_span_basic(trace_instance): - trace_instance.ensure_root_span("tid") - assert "tid" in trace_instance.dify_trace_ids - - -def test_ensure_root_span_ignores_unsampled_ambient_otel_parent(): - exporter = _CollectingSpanExporter() - provider = trace_sdk.TracerProvider() - provider.add_span_processor(SimpleSpanProcessor(exporter)) - trace_instance = ArizePhoenixDataTrace.__new__(ArizePhoenixDataTrace) - trace_instance.tracer = cast(Tracer, provider.get_tracer("test-phoenix-root-span")) - trace_instance.propagator = TraceContextTextMapPropagator() - trace_instance.project = "p" - trace_instance.dify_trace_ids = set() - trace_instance.root_span_carriers = {} - trace_instance.carrier = {} - - ambient_span_context = SpanContext( - trace_id=0x11111111111111111111111111111111, - span_id=0x2222222222222222, - is_remote=True, - trace_flags=TraceFlags(0), - trace_state=TraceState(), - ) - - with use_span(NonRecordingSpan(ambient_span_context), end_on_exit=False): - carrier = trace_instance.ensure_root_span("tid") - - assert len(exporter.spans) == 1 - root_span = exporter.spans[0] - root_span_context = root_span.get_span_context() - assert root_span_context is not None - assert root_span.parent is None - assert root_span_context.trace_id != ambient_span_context.trace_id - assert carrier["traceparent"].split("-")[1] == f"{root_span_context.trace_id:032x}" - - -def test_ensure_root_span_uses_custom_name_and_attributes(trace_instance): - root_attributes = { - SpanAttributes.INPUT_VALUE: '{"input":"value"}', - SpanAttributes.OUTPUT_VALUE: '{"output":"value"}', - } - - trace_instance.ensure_root_span("tid", root_span_name="Workflow Name", root_span_attributes=root_attributes) - - trace_instance.tracer.start_span.assert_called_once_with( - name="Workflow Name", - attributes={ - SpanAttributes.OPENINFERENCE_SPAN_KIND: "CHAIN", - "dify_project_name": "p", - "dify_trace_id": "tid", - SpanAttributes.INPUT_VALUE: '{"input":"value"}', - SpanAttributes.OUTPUT_VALUE: '{"output":"value"}', - }, - context=Context(), - ) - - -def test_ensure_root_span_records_error_status_and_exception_event(trace_instance, span_exporter): - trace_instance.ensure_root_span("tid", root_span_error="workflow failed") - - assert len(span_exporter.spans) == 1 - root_span = span_exporter.spans[0] - assert root_span.status.status_code == StatusCode.ERROR - assert len(root_span.events) == 1 - assert root_span.events[0].name == "exception" - assert root_span.events[0].attributes[OTELSpanAttributes.EXCEPTION_MESSAGE] == "workflow failed" - - -def test_ensure_root_span_falls_back_to_dify_name_when_custom_name_is_blank(trace_instance): - trace_instance.ensure_root_span("tid", root_span_name=" ") - - trace_instance.tracer.start_span.assert_called_once() - assert trace_instance.tracer.start_span.call_args.kwargs["name"] == "Dify" diff --git a/api/providers/trace/trace-arize-phoenix/tests/unit_tests/src/dify_trace_arize_phoenix/test_config.py b/api/providers/trace/trace-arize-phoenix/tests/unit_tests/src/dify_trace_arize_phoenix/test_config.py index 11e951c3b1e281..9bcee89f493623 100644 --- a/api/providers/trace/trace-arize-phoenix/tests/unit_tests/src/dify_trace_arize_phoenix/test_config.py +++ b/api/providers/trace/trace-arize-phoenix/tests/unit_tests/src/dify_trace_arize_phoenix/test_config.py @@ -2,87 +2,95 @@ from dify_trace_arize_phoenix.config import ArizeConfig, PhoenixConfig from pydantic import ValidationError - -class TestArizeConfig: - """Test cases for ArizeConfig""" - - def test_valid_config(self): - """Test valid Arize configuration""" - config = ArizeConfig( - api_key="test_key", space_id="test_space", project="test_project", endpoint="https://custom.arize.com" - ) - assert config.api_key == "test_key" - assert config.space_id == "test_space" - assert config.project == "test_project" - assert config.endpoint == "https://custom.arize.com" - - def test_default_values(self): - """Test default values are set correctly""" - config = ArizeConfig() - assert config.api_key is None - assert config.space_id is None - assert config.project is None - assert config.endpoint == "https://otlp.arize.com" - - def test_project_validation_empty(self): - """Test project validation with empty value""" - config = ArizeConfig(project="") - assert config.project == "default" - - def test_project_validation_none(self): - """Test project validation with None value""" - config = ArizeConfig(project=None) - assert config.project == "default" - - def test_endpoint_validation_empty(self): - """Test endpoint validation with empty value""" - config = ArizeConfig(endpoint="") - assert config.endpoint == "https://otlp.arize.com" - - def test_endpoint_validation_with_path(self): - """Test endpoint validation normalizes URL by removing path""" - config = ArizeConfig(endpoint="https://custom.arize.com/api/v1") - assert config.endpoint == "https://custom.arize.com" - - def test_endpoint_validation_invalid_scheme(self): - """Test endpoint validation rejects invalid schemes""" - with pytest.raises(ValidationError, match="URL scheme must be one of"): - ArizeConfig(endpoint="ftp://invalid.com") - - def test_endpoint_validation_no_scheme(self): - """Test endpoint validation rejects URLs without scheme""" - with pytest.raises(ValidationError, match="URL scheme must be one of"): - ArizeConfig(endpoint="invalid.com") - - -class TestPhoenixConfig: - """Test cases for PhoenixConfig""" - - def test_valid_config(self): - """Test valid Phoenix configuration""" - config = PhoenixConfig(api_key="test_key", project="test_project", endpoint="https://custom.phoenix.com") - assert config.api_key == "test_key" - assert config.project == "test_project" - assert config.endpoint == "https://custom.phoenix.com" - - def test_default_values(self): - """Test default values are set correctly""" - config = PhoenixConfig() - assert config.api_key is None - assert config.project is None - assert config.endpoint == "https://app.phoenix.arize.com" - - def test_project_validation_empty(self): - """Test project validation with empty value""" - config = PhoenixConfig(project="") - assert config.project == "default" - - def test_endpoint_validation_with_path(self): - """Test endpoint validation with path""" - config = PhoenixConfig(endpoint="https://app.phoenix.arize.com/s/dify-integration") - assert config.endpoint == "https://app.phoenix.arize.com/s/dify-integration" - - def test_endpoint_validation_without_path(self): - """Test endpoint validation without path""" - config = PhoenixConfig(endpoint="https://app.phoenix.arize.com") - assert config.endpoint == "https://app.phoenix.arize.com" +SPACE = "U3BhY2U6dGVzdA==" + + +def test_phoenix_key_is_an_optional_secret() -> None: + assert PhoenixConfig.secret_fields() == {"api_key"} + assert PhoenixConfig.derived_fields() == set() + assert [(field.name, field.secret, field.required) for field in PhoenixConfig.form_fields()] == [ + ("api_key", True, False), + ("project", False, False), + ("endpoint", False, False), + ] + config = PhoenixConfig() + assert (config.api_key, config.project, config.endpoint) == (None, "default", "https://app.phoenix.arize.com") + + +def test_arize_credentials_are_required_secrets_and_the_destination_is_derived() -> None: + assert ArizeConfig.secret_fields() == {"api_key", "space_id"} + assert ArizeConfig.derived_fields() == {"arize_project_destination"} + assert [(field.name, field.secret, field.required) for field in ArizeConfig.form_fields()] == [ + ("api_key", True, True), + ("space_id", True, True), + ("project", False, False), + ("endpoint", False, False), + ] + config = ArizeConfig(api_key="key", space_id=SPACE) + assert (config.project, config.endpoint, config.arize_project_destination) == ( + "default", + "https://otlp.arize.com", + None, + ) + with pytest.raises(ValidationError): + ArizeConfig.model_validate({"api_key": "key"}) + + +@pytest.mark.parametrize(("project", "expected"), [(None, "default"), ("", "default"), (" ", "default"), (" p ", "p")]) +def test_empty_projects_fall_back_to_default(project: str | None, expected: str) -> None: + assert PhoenixConfig.model_validate({"project": project}).project == expected + assert ArizeConfig.model_validate({"api_key": "key", "space_id": SPACE, "project": project}).project == expected + + +def test_phoenix_endpoint_keeps_its_path() -> None: + assert PhoenixConfig(endpoint=" https://phoenix.example.com/s/space/ ").endpoint == ( + "https://phoenix.example.com/s/space/" + ) + assert PhoenixConfig(endpoint=" ").endpoint == "https://app.phoenix.arize.com" + + +def test_arize_endpoint_is_reduced_to_the_collector_origin() -> None: + def endpoint(value: str) -> str: + return ArizeConfig(api_key="key", space_id=SPACE, endpoint=value).endpoint + + assert endpoint("https://otlp.eu-west-1a.arize.com/v1/") == "https://otlp.eu-west-1a.arize.com" + assert endpoint("http://collector.internal:4317") == "http://collector.internal:4317" + assert endpoint("") == "https://otlp.arize.com" + + +@pytest.mark.parametrize( + "endpoint", ["ftp://collector.example.com", "collector.example.com", "https://user:hunter2@collector.example.com"] +) +def test_invalid_endpoints_are_rejected_without_echoing_them(endpoint: str) -> None: + for model, values in ((PhoenixConfig, {}), (ArizeConfig, {"api_key": "key", "space_id": SPACE})): + with pytest.raises(ValidationError) as caught: + model.model_validate({**values, "endpoint": endpoint}) + assert "hunter2" not in str(caught.value) + + +def test_stored_rows_still_load() -> None: + phoenix = PhoenixConfig.model_validate( + {"api_key": None, "project": None, "endpoint": "https://app.phoenix.arize.com", "project_url": "https://x"} + ) + assert (phoenix.project, phoenix.endpoint) == ("default", "https://app.phoenix.arize.com") + + destination = { + "endpoint": "https://otlp.arize.com", + "space_id": SPACE, + "project": "exact", + "organization_id": "T3JnOnRlc3Q=", + "project_id": "UHJvamVjdDovKz0=", + "source_digest": "0" * 64, + } + arize = ArizeConfig.model_validate( + { + "api_key": "key", + "space_id": SPACE, + "project": "exact", + "endpoint": "https://otlp.arize.com", + "arize_project_destination": destination, + "project_url": "https://app.arize.com/organizations/forged", + } + ) + assert arize.arize_project_destination == destination + assert "project_url" not in arize.model_dump() diff --git a/api/providers/trace/trace-arize-phoenix/tests/unit_tests/src/dify_trace_arize_phoenix/test_provider.py b/api/providers/trace/trace-arize-phoenix/tests/unit_tests/src/dify_trace_arize_phoenix/test_provider.py new file mode 100644 index 00000000000000..628a0d0b1f6614 --- /dev/null +++ b/api/providers/trace/trace-arize-phoenix/tests/unit_tests/src/dify_trace_arize_phoenix/test_provider.py @@ -0,0 +1,637 @@ +from collections.abc import Iterable, Iterator +from concurrent import futures + +import grpc # type: ignore[import-untyped] +import pytest +from dify_trace_arize_phoenix.config import ArizeConfig, PhoenixConfig +from dify_trace_arize_phoenix.provider import ArizeProvider, PhoenixProvider +from dify_trace_sdk.egress import EgressResponse +from dify_trace_sdk.errors import ( + EgressBlockedError, + EgressError, + ExportOutcomeUnknownError, + ExportRejectedError, + ExportRetryableError, + VerificationError, +) +from dify_trace_sdk.exporting import json_value +from dify_trace_sdk.model import SpanStatus, Trace +from dify_trace_sdk.testing import ( + WORKFLOW_RUN_ID, + RecordingEgress, + message_trace, + response, + tool_fragment, + workflow_trace, +) +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( + ExportTraceServiceRequest, + ExportTraceServiceResponse, +) +from opentelemetry.proto.collector.trace.v1.trace_service_pb2_grpc import add_TraceServiceServicer_to_server +from opentelemetry.proto.common.v1.common_pb2 import AnyValue, KeyValue +from opentelemetry.proto.trace.v1.trace_pb2 import Status + +SPACE = "U3BhY2U6dGVzdA==" +ORGANIZATION = "T3JnOnRlc3Q=" +PROJECT_ID = "UHJvamVjdDovKz0=" +ARIZE_CONFIG = ArizeConfig(api_key="ak-secret-key", space_id=SPACE, project="exact") +DESTINATION = { + "endpoint": "https://otlp.arize.com", + "space_id": SPACE, + "project": "exact", + "organization_id": ORGANIZATION, + "project_id": PROJECT_ID, +} +URL = "https://app.arize.com/organizations/T3JnOnRlc3Q%3D/spaces/U3BhY2U6dGVzdA%3D%3D/projects/UHJvamVjdDovKz0%3D" + + +class _Collector: + """An in-process OTLP collector; gRPC dispatches to ``Export`` by name.""" + + def __init__(self) -> None: + self.target = "" + self.status = grpc.StatusCode.OK + self.ack = ExportTraceServiceResponse() + self.received: list[tuple[ExportTraceServiceRequest, dict[str, str | bytes]]] = [] + + def Export( # noqa: N802 - gRPC method name + self, request: ExportTraceServiceRequest, context: grpc.ServicerContext + ) -> ExportTraceServiceResponse: + self.received.append((request, dict(context.invocation_metadata()))) + if self.status != grpc.StatusCode.OK: + context.abort(self.status, "vendor detail about ak-secret-key") + return self.ack + + +@pytest.fixture +def collector() -> Iterator[_Collector]: + collector = _Collector() + server = grpc.server(futures.ThreadPoolExecutor(max_workers=1)) + add_TraceServiceServicer_to_server(collector, server) + collector.target = f"127.0.0.1:{server.add_insecure_port('127.0.0.1:0')}" + server.start() + try: + yield collector + finally: + server.stop(None) + + +def _arize_provider( + grpc_target: str | None, *responses: EgressResponse | Exception +) -> tuple[ArizeProvider, RecordingEgress]: + egress = RecordingEgress(responses, grpc_target=grpc_target) + return ArizeProvider(egress), egress + + +def _arize_config(**changes: object) -> ArizeConfig: + return ArizeConfig.model_validate({**ARIZE_CONFIG.model_dump(), **changes}) + + +def _ack(rejected_spans: int, error_message: str) -> ExportTraceServiceResponse: + ack = ExportTraceServiceResponse() + ack.partial_success.rejected_spans = rejected_spans + ack.partial_success.error_message = error_message + return ack + + +def _space() -> EgressResponse: + return response(200, {"id": SPACE, "organization_id": ORGANIZATION}) + + +def _page(projects: list[dict[str, str]], cursor: str | None) -> EgressResponse: + return response(200, {"projects": projects, "pagination": {"has_more": cursor is not None, "next_cursor": cursor}}) + + +def _project(name: str, project_id: str) -> dict[str, str]: + return {"id": project_id, "name": name, "space_id": SPACE} + + +class TestArizeProvider: + @staticmethod + def test_verify_requires_an_empty_export_ack_and_resolves_the_project_link(collector: _Collector) -> None: + provider, egress = _arize_provider( + collector.target, + _space(), + _page([_project("exact-substring", "UHJvamVjdDoy")], "next"), + _page([_project("exact", PROJECT_ID)], None), + ) + verified = provider.verify(ARIZE_CONFIG) + + assert verified.arize_project_destination == DESTINATION + assert provider.project_url(verified) == URL + assert egress.grpc_endpoints == ["https://otlp.arize.com"] + ((probe, metadata),) = collector.received + (resource_spans,) = probe.resource_spans + assert {item.key: item.value.string_value for item in resource_spans.resource.attributes} == { + "model_id": "exact", + "openinference.project.name": "exact", + } + assert len(resource_spans.scope_spans) == 0 + assert (metadata["api_key"], metadata["space_id"], metadata["authorization"]) == ( + "ak-secret-key", + SPACE, + "Bearer ak-secret-key", + ) + lookup = "https://api.arize.com/v2/projects?space_id=U3BhY2U6dGVzdA%3D%3D&name=exact&limit=100" + assert [request.url for request in egress.requests] == [ + "https://api.arize.com/v2/spaces/U3BhY2U6dGVzdA%3D%3D", + lookup, + lookup + "&cursor=next", + ] + for request in egress.requests: + assert request.method == "GET" + assert request.headers == {"Authorization": "Bearer ak-secret-key", "Accept": "application/json"} + assert request.sensitive_url is True + + @staticmethod + @pytest.mark.parametrize( + ("status", "code"), + [ + (grpc.StatusCode.UNAUTHENTICATED, "authentication_failed"), + (grpc.StatusCode.PERMISSION_DENIED, "permission_denied"), + (grpc.StatusCode.UNAVAILABLE, "connection_failed"), + (grpc.StatusCode.RESOURCE_EXHAUSTED, "connection_failed"), + ], + ) + def test_rejected_probe_has_a_stable_code(collector: _Collector, status: grpc.StatusCode, code: str) -> None: + collector.status = status + provider, egress = _arize_provider(collector.target) + with pytest.raises(VerificationError) as caught: + provider.verify(ARIZE_CONFIG) + assert caught.value.code == code + assert caught.value.__suppress_context__ + assert "ak-secret-key" not in str(caught.value) + assert len(collector.received) == 1 + assert egress.requests == [] + + @staticmethod + @pytest.mark.parametrize("ack", [_ack(1, ""), _ack(0, "vendor detail")]) + def test_partial_probe_ack_is_unexpected(collector: _Collector, ack: ExportTraceServiceResponse) -> None: + collector.ack = ack + provider, egress = _arize_provider(collector.target) + with pytest.raises(VerificationError) as caught: + provider.verify(ARIZE_CONFIG) + assert caught.value.code == "unexpected_response" + assert egress.requests == [] + + @staticmethod + @pytest.mark.parametrize( + "changes", + [ + {"api_key": ""}, + {"api_key": "******"}, + {"api_key": "key\nheader"}, + {"api_key": "ключ"}, + {"space_id": " "}, + {"space_id": "******"}, + {"endpoint": "https://otlp.arize.com:0"}, + {"endpoint": "https://otlp.arize.com:99999"}, + ], + ) + def test_invalid_values_never_open_a_channel(changes: dict[str, str]) -> None: + provider, egress = _arize_provider(None) + with pytest.raises(VerificationError) as caught: + provider.verify(_arize_config(**changes)) + assert caught.value.code == "invalid_configuration" + assert egress.grpc_endpoints == [] + + @staticmethod + def test_disallowed_collector_is_a_connection_failure() -> None: + provider, egress = _arize_provider(None) + with pytest.raises(VerificationError) as caught: + provider.verify(ARIZE_CONFIG) + assert caught.value.code == "connection_failed" + assert egress.grpc_endpoints == ["https://otlp.arize.com"] + + @staticmethod + @pytest.mark.parametrize( + "lookups", + [ + [response(403, {"message": "denied"})], + [EgressError("read timeout", request_sent=True)], + [response(200, b"not json")], + [response(200, {"id": "U3BhY2U6b3RoZXI=", "organization_id": ORGANIZATION})], + [response(200, {"id": SPACE, "organization_id": "../escape"})], + [_space(), _page([_project("exact", PROJECT_ID), _project("exact", "UHJvamVjdDoy")], None)], + [_space(), _page([_project("EXACT", PROJECT_ID)], None)], + [_space(), _page([{**_project("exact", PROJECT_ID), "space_id": "U3BhY2U6b3RoZXI="}], None)], + [_space(), response(200, {"projects": [_project("exact", PROJECT_ID)], "pagination": {"has_more": True}})], + [_space(), *(_page([_project("other", PROJECT_ID)], str(index)) for index in range(3))], + [_space(), _page([], "same"), _page([], "same")], + ], + ) + def test_project_link_is_optional(collector: _Collector, lookups: list[EgressResponse | Exception]) -> None: + provider, egress = _arize_provider(collector.target, *lookups) + verified = provider.verify(ARIZE_CONFIG) + assert verified.arize_project_destination is None + assert provider.project_url(verified) is None + assert len(egress.requests) == len(lookups) + + @staticmethod + @pytest.mark.parametrize( + "endpoint", ["https://otlp.eu-west-1a.arize.com", "http://otlp.arize.com", "https://otlp.arize.com:4317"] + ) + def test_other_collectors_never_send_credentials_to_arize_cloud(collector: _Collector, endpoint: str) -> None: + provider, egress = _arize_provider(collector.target) + verified = provider.verify(_arize_config(endpoint=endpoint)) + assert verified.arize_project_destination is None + assert egress.grpc_endpoints == [endpoint] + assert egress.requests == [] + + @staticmethod + def test_project_url_reads_a_stored_destination_without_io() -> None: + provider, egress = _arize_provider(None) + stored = _arize_config(arize_project_destination={**DESTINATION, "source_digest": "0" * 64}) + assert provider.project_url(stored) == URL + assert (egress.requests, egress.grpc_endpoints) == ([], []) + + @staticmethod + def test_project_url_needs_no_decrypted_space_id() -> None: + provider, _ = _arize_provider(None) + stored = _arize_config( + arize_project_destination=DESTINATION, space_id="ZW5jcnlwdGVkLXNwYWNlLWlkLWNpcGhlcnRleHQ=" + ) + assert provider.project_url(stored) == URL + + @staticmethod + @pytest.mark.parametrize( + "changes", + [ + {"project": "renamed"}, + {"endpoint": "https://otlp.arize.com:443"}, + {"arize_project_destination": {**DESTINATION, "organization_id": "bad?query"}}, + {"arize_project_destination": {**DESTINATION, "project_id": "../escape"}}, + {"arize_project_destination": None}, + ], + ) + def test_project_url_refuses_stale_or_forged_destinations(changes: dict[str, object]) -> None: + provider, _ = _arize_provider(None) + assert provider.project_url(_arize_config(**{"arize_project_destination": DESTINATION, **changes})) is None + + @staticmethod + def test_send_exports_openinference_spans_with_arize_metadata(collector: _Collector) -> None: + provider, egress = _arize_provider(collector.target) + trace = message_trace() + result = provider.send(ARIZE_CONFIG, trace) + + assert egress.grpc_endpoints == ["https://otlp.arize.com"] + ((request, metadata),) = collector.received + assert (metadata["api_key"], metadata["space_id"], metadata["authorization"]) == ( + "ak-secret-key", + SPACE, + "Bearer ak-secret-key", + ) + (resource_spans,) = request.resource_spans + assert {item.key: item.value.string_value for item in resource_spans.resource.attributes} == { + "service.name": "dify-application", + "openinference.project.name": "exact", + "model_id": "exact", + } + _, llm = resource_spans.scope_spans[0].spans + llm_attributes = {item.key: item.value for item in llm.attributes} + assert llm_attributes["openinference.span.kind"].string_value == "LLM" + assert llm_attributes["llm.model_name"].string_value == "gpt-4o" + assert llm_attributes["llm.token_count.total"].int_value == 20 + assert llm_attributes["session.id"].string_value == "conversation-1" + assert (result.trace_id, result.trace_url, result.span_urls) == (trace.trace_id, None, {}) + + @staticmethod + @pytest.mark.parametrize("trace", [message_trace(), workflow_trace(), tool_fragment()]) + def test_send_keeps_ids_parents_and_status(collector: _Collector, trace: Trace) -> None: + provider, _ = _arize_provider(collector.target) + provider.send(ARIZE_CONFIG, trace) + + ((request, _),) = collector.received + spans = request.resource_spans[0].scope_spans[0].spans + sent = [ + (span.trace_id.hex(), span.span_id.hex(), span.parent_span_id.hex(), span.status.code) for span in spans + ] + assert sent == [ + ( + trace.trace_id, + span.span_id, + span.parent_span_id or "", + Status.STATUS_CODE_ERROR if span.status is SpanStatus.ERROR else Status.STATUS_CODE_OK, + ) + for span in trace.spans + ] + + @staticmethod + @pytest.mark.parametrize( + ("status", "expected"), + [ + (grpc.StatusCode.UNAVAILABLE, ExportRetryableError), + (grpc.StatusCode.RESOURCE_EXHAUSTED, ExportRetryableError), + (grpc.StatusCode.DEADLINE_EXCEEDED, ExportOutcomeUnknownError), + (grpc.StatusCode.UNAUTHENTICATED, ExportRejectedError), + (grpc.StatusCode.INVALID_ARGUMENT, ExportRejectedError), + ], + ) + def test_send_failures_are_classified( + collector: _Collector, status: grpc.StatusCode, expected: type[Exception] + ) -> None: + collector.status = status + provider, _ = _arize_provider(collector.target) + with pytest.raises(expected) as caught: + provider.send(ARIZE_CONFIG, message_trace()) + assert "ak-secret-key" not in str(caught.value) + + @staticmethod + def test_partial_acceptance_is_outcome_unknown(collector: _Collector) -> None: + collector.ack = _ack(1, "") + provider, _ = _arize_provider(collector.target) + with pytest.raises(ExportOutcomeUnknownError): + provider.send(ARIZE_CONFIG, message_trace()) + + @staticmethod + def test_disallowed_collector_rejects_the_export() -> None: + provider, _ = _arize_provider(None) + with pytest.raises(ExportRejectedError): + provider.send(ARIZE_CONFIG, message_trace()) + + @staticmethod + @pytest.mark.parametrize("changes", [{"api_key": "ключ"}, {"space_id": "line\nbreak"}]) + def test_credentials_grpc_cannot_carry_are_rejected_before_export( + collector: _Collector, changes: dict[str, str] + ) -> None: + provider, egress = _arize_provider(collector.target) + with pytest.raises(ExportRejectedError) as caught: + provider.send(_arize_config(**changes), message_trace()) + assert "ключ" not in str(caught.value) + assert (egress.grpc_endpoints, collector.received) == ([], []) + + +PHOENIX_CONFIG = PhoenixConfig( + api_key="phx-secret-key", project="dify-app", endpoint="https://phoenix.example.com/s/space" +) +EMPTY_PAGE = {"data": [], "next_cursor": None} + + +def _phoenix_provider(*responses: EgressResponse | Exception) -> tuple[PhoenixProvider, RecordingEgress]: + egress = RecordingEgress(responses, grpc_target=None) + return PhoenixProvider(egress), egress + + +def _phoenix_config(**changes: str | None) -> PhoenixConfig: + return PhoenixConfig.model_validate({**PHOENIX_CONFIG.model_dump(), **changes}) + + +def _value(value: AnyValue) -> object: + match value.WhichOneof("value"): + case "string_value": + return value.string_value + case "int_value": + return value.int_value + case "double_value": + return value.double_value + case "bool_value": + return value.bool_value + case other: + raise AssertionError(f"unexpected attribute value type {other}") + + +def _attributes(attributes: Iterable[KeyValue]) -> dict[str, object]: + return {attribute.key: _value(attribute.value) for attribute in attributes} + + +def _sent(egress: RecordingEgress) -> ExportTraceServiceRequest: + (sent,) = egress.requests + return ExportTraceServiceRequest.FromString(sent.content or b"") + + +def _partial_ack() -> bytes: + ack = ExportTraceServiceResponse() + ack.partial_success.rejected_spans = 1 + return ack.SerializeToString() + + +class TestPhoenixProvider: + @staticmethod + def test_verify_normalizes_the_destination_and_reads_one_project_page() -> None: + provider, egress = _phoenix_provider(response(200, EMPTY_PAGE)) + verified = provider.verify( + _phoenix_config(endpoint=" HTTP://Phoenix.EXAMPLE.com:6006/proxy/s/space/ ", project=" project name ") + ) + assert (verified.endpoint, verified.project) == ( + "http://phoenix.example.com:6006/proxy/s/space", + "project name", + ) + (sent,) = egress.requests + assert (sent.method, sent.url, sent.content) == ( + "GET", + "http://phoenix.example.com:6006/proxy/s/space/v1/projects?limit=1", + None, + ) + assert sent.headers == {"Accept": "application/json", "Authorization": "Bearer phx-secret-key"} + assert sent.sensitive_url is False + + @staticmethod + @pytest.mark.parametrize("api_key", [None, ""]) + def test_verify_without_a_key_sends_no_authorization(api_key: str | None) -> None: + provider, egress = _phoenix_provider(response(200, EMPTY_PAGE)) + provider.verify(PhoenixConfig(api_key=api_key)) + (sent,) = egress.requests + assert sent.url == "https://app.phoenix.arize.com/v1/projects?limit=1" + assert sent.headers == {"Accept": "application/json"} + + @staticmethod + def test_verify_does_not_require_the_configured_project_or_follow_the_cursor() -> None: + page = {"data": [{"id": "UHJvamVjdDox", "name": "other", "description": None}], "next_cursor": "opaque"} + provider, egress = _phoenix_provider(response(200, page)) + assert provider.verify(PHOENIX_CONFIG).project == "dify-app" + assert len(egress.requests) == 1 + + @staticmethod + @pytest.mark.parametrize( + ("changes", "code"), + [ + ({"endpoint": "https://phoenix.example.com/v1/traces"}, "invalid_endpoint"), + ({"endpoint": "https://phoenix.example.com/s/../space"}, "invalid_endpoint"), + ({"endpoint": "https://phoenix.example.com/s%2fspace"}, "invalid_endpoint"), + ({"endpoint": "https://phoenix.example.com//space"}, "invalid_endpoint"), + ({"endpoint": "https://phoenix.example.com/s?token=secret"}, "invalid_endpoint"), + ({"endpoint": "https://phoenix.example.com:0"}, "invalid_endpoint"), + ({"project": ".."}, "invalid_project"), + ({"project": "line\nbreak"}, "invalid_project"), + ({"api_key": "******"}, "invalid_api_key"), + ({"api_key": "two words"}, "invalid_api_key"), + ({"api_key": "ключ"}, "invalid_api_key"), + ], + ) + def test_invalid_values_fail_before_any_request(changes: dict[str, str], code: str) -> None: + provider, egress = _phoenix_provider() + with pytest.raises(VerificationError) as caught: + provider.verify(_phoenix_config(**changes)) + assert caught.value.code == code + assert egress.requests == [] + + @staticmethod + @pytest.mark.parametrize( + ("outcome", "code"), + [ + (response(401, {"detail": "phx-secret-key is invalid"}), "authentication_failed"), + (response(403, {}), "permission_denied"), + (response(404, {}), "endpoint_unavailable"), + (response(302, {}), "endpoint_unavailable"), + (response(500, {}), "endpoint_unavailable"), + (response(204, b""), "endpoint_unavailable"), + (response(200, b"login"), "malformed_response"), + (response(200, []), "malformed_response"), + (response(200, {"data": []}), "malformed_response"), + (response(200, {"data": [], "next_cursor": 3}), "malformed_response"), + (response(200, {"data": [{"id": "1"}], "next_cursor": None}), "malformed_response"), + (response(200, {"data": [{"id": "1", "name": ""}], "next_cursor": None}), "malformed_response"), + ( + response(200, {"data": [{"id": "1", "name": "p", "description": {}}], "next_cursor": None}), + "malformed_response", + ), + (response(200, {**EMPTY_PAGE, "padding": "x" * 1024 * 1024}), "malformed_response"), + (EgressError("connect timeout", request_sent=False), "connection_failed"), + (EgressBlockedError("private address"), "connection_failed"), + ], + ) + def test_verify_failures_have_stable_codes(outcome: EgressResponse | Exception, code: str) -> None: + provider, egress = _phoenix_provider(outcome) + with pytest.raises(VerificationError) as caught: + provider.verify(PHOENIX_CONFIG) + assert caught.value.code == code + assert "phx-secret-key" not in str(caught.value) + assert caught.value.__context__ is None or caught.value.__suppress_context__ + assert len(egress.requests) == 1 + + @staticmethod + def test_project_url_redirects_by_name_without_io() -> None: + provider, egress = _phoenix_provider() + config = PhoenixConfig(endpoint="https://Phoenix.example.com:8443/s/space/", project="A /? & #中文") + assert provider.project_url(config) == ( + "https://phoenix.example.com:8443/s/space/redirects/projects/A%20%2F%3F%20%26%20%23%E4%B8%AD%E6%96%87" + ) + assert provider.project_url(PhoenixConfig()) == "https://app.phoenix.arize.com/redirects/projects/default" + assert provider.project_url(PhoenixConfig(endpoint="https://phoenix.example.com/v1/traces")) is None + assert egress.requests == [] + + @staticmethod + def test_message_trace_becomes_openinference_spans() -> None: + provider, egress = _phoenix_provider(response(200, ExportTraceServiceResponse().SerializeToString())) + trace = message_trace() + result = provider.send(PHOENIX_CONFIG, trace) + + (sent,) = egress.requests + assert sent.url == "https://phoenix.example.com/s/space/v1/traces" + assert sent.headers["Content-Type"] == "application/x-protobuf" + assert (sent.headers["api_key"], sent.headers["authorization"]) == ("phx-secret-key", "Bearer phx-secret-key") + (resource_spans,) = _sent(egress).resource_spans + assert _attributes(resource_spans.resource.attributes) == { + "service.name": "dify-application", + "openinference.project.name": "dify-app", + "model_id": "dify-app", + } + message, llm = resource_spans.scope_spans[0].spans + assert [(span.trace_id.hex(), span.span_id.hex(), span.parent_span_id.hex()) for span in (message, llm)] == [ + (trace.trace_id, trace.spans[0].span_id, ""), + (trace.trace_id, trace.spans[1].span_id, trace.spans[0].span_id), + ] + assert (message.name, message.status.code) == ("message", Status.STATUS_CODE_OK) + assert _attributes(message.attributes) == { + "openinference.span.kind": "CHAIN", + "input.value": '{"query": "What is Dify?"}', + "input.mime_type": "application/json", + "output.value": '"An LLM application platform."', + "output.mime_type": "application/json", + "metadata": json_value(trace.spans[0].attributes), + "session.id": "conversation-1", + "user.id": "end-user-1", + "dify.span.id": trace.spans[0].span_id, + } + llm_attributes = _attributes(llm.attributes) + assert llm_attributes["openinference.span.kind"] == "LLM" + assert llm_attributes["input.value"] == '[{"role": "user", "content": "What is Dify?"}]' + assert {key: value for key, value in llm_attributes.items() if key.startswith("llm.")} == { + "llm.model_name": "gpt-4o", + "llm.provider": "openai", + "llm.token_count.prompt": 12, + "llm.token_count.completion": 8, + "llm.token_count.total": 20, + } + assert (result.trace_id, result.trace_url, result.span_urls) == (trace.trace_id, None, {}) + + @staticmethod + def test_workflow_spans_keep_hierarchy_errors_and_tool_fields() -> None: + provider, egress = _phoenix_provider(response(200, b"")) + trace = workflow_trace() + provider.send(PHOENIX_CONFIG, trace) + + run, llm, tool = _sent(egress).resource_spans[0].scope_spans[0].spans + assert [span.parent_span_id.hex() for span in (run, llm, tool)] == [ + "", + trace.spans[0].span_id, + trace.spans[1].span_id, + ] + assert (run.status.code, run.status.message) == (Status.STATUS_CODE_ERROR, "Tool failed") + assert (llm.status.code, llm.status.message) == (Status.STATUS_CODE_OK, "") + assert (tool.status.code, tool.status.message) == (Status.STATUS_CODE_ERROR, "Tool failed") + tool_attributes = _attributes(tool.attributes) + assert tool_attributes["openinference.span.kind"] == "TOOL" + assert (tool_attributes["tool.name"], tool_attributes["tool.parameters"]) == ( + "tool_Search", + '{"query": "tracing"}', + ) + assert (tool_attributes["session.id"], tool_attributes["user.id"]) == (WORKFLOW_RUN_ID, "account-1") + assert not any(key.startswith("llm.") for key in tool_attributes) + assert _attributes(llm.attributes)["llm.token_count.total"] == 40 + + @staticmethod + def test_fragment_root_points_at_the_separately_sent_parent() -> None: + provider, egress = _phoenix_provider(response(200, b"")) + fragment = tool_fragment() + provider.send(PHOENIX_CONFIG, fragment) + + (span,) = _sent(egress).resource_spans[0].scope_spans[0].spans + assert fragment.parent is not None + assert (span.trace_id.hex(), span.span_id.hex(), span.parent_span_id.hex()) == ( + fragment.trace_id, + fragment.root.span_id, + fragment.parent.span_id, + ) + attributes = _attributes(span.attributes) + assert (attributes["openinference.span.kind"], attributes["tool.name"]) == ("TOOL", "weather") + assert attributes["output.value"] == '"Sunny"' + + @staticmethod + def test_send_without_a_key_sends_no_credentials() -> None: + provider, egress = _phoenix_provider(response(200, b"")) + provider.send(PhoenixConfig(endpoint="http://phoenix.internal:6006/"), message_trace()) + (sent,) = egress.requests + assert sent.url == "http://phoenix.internal:6006/v1/traces" + assert sent.headers == {"Content-Type": "application/x-protobuf", "Accept": "application/x-protobuf"} + + @staticmethod + @pytest.mark.parametrize( + ("outcome", "expected"), + [ + (response(503, {}), ExportRetryableError), + (response(429, {}), ExportRetryableError), + (response(401, {"detail": "phx-secret-key is invalid"}), ExportRejectedError), + (EgressError("connection refused", request_sent=False), ExportRetryableError), + (EgressError("read timeout", request_sent=True), ExportOutcomeUnknownError), + (EgressBlockedError("private address"), ExportRejectedError), + (response(200, _partial_ack()), ExportOutcomeUnknownError), + (response(200, b"not-protobuf-\xff"), ExportOutcomeUnknownError), + ], + ) + def test_send_failures_are_classified(outcome: EgressResponse | Exception, expected: type[Exception]) -> None: + provider, _ = _phoenix_provider(outcome) + with pytest.raises(expected) as caught: + provider.send(PHOENIX_CONFIG, message_trace()) + assert "phx-secret-key" not in str(caught.value) + + @staticmethod + @pytest.mark.parametrize( + "changes", [{"api_key": "line\nbreak"}, {"endpoint": "https://phoenix.example.com/?token=1"}] + ) + def test_unusable_configuration_is_rejected_without_a_request(changes: dict[str, str]) -> None: + provider, egress = _phoenix_provider() + with pytest.raises(ExportRejectedError): + provider.send(_phoenix_config(**changes), message_trace()) + assert egress.requests == [] diff --git a/api/providers/trace/trace-langfuse/pyproject.toml b/api/providers/trace/trace-langfuse/pyproject.toml index 27d2273a6943b6..6b23f9e64de8b2 100644 --- a/api/providers/trace/trace-langfuse/pyproject.toml +++ b/api/providers/trace/trace-langfuse/pyproject.toml @@ -2,9 +2,12 @@ name = "dify-trace-langfuse" version = "0.0.1" dependencies = [ - "langfuse>=4.2.0,<5.0.0", + "dify-trace-sdk", ] description = "Dify ops tracing provider (Langfuse)." +[project.entry-points."dify.trace_providers"] +langfuse = "dify_trace_langfuse.provider:LangfuseProvider" + [tool.setuptools.packages.find] where = ["src"] diff --git a/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/agent_trace.py b/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/agent_trace.py deleted file mode 100644 index 841bd1227c919b..00000000000000 --- a/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/agent_trace.py +++ /dev/null @@ -1,75 +0,0 @@ -"""Export complete Agent/Workflow trees through Langfuse v4 OTLP ingestion.""" - -import base64 -from collections.abc import Mapping -from datetime import timedelta -from typing import Any -from urllib.parse import quote - -from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace -from services.tracing.otlp_export import build_request, export_http -from services.tracing.provider_common import endpoint, json_value, required, token_usage, usd_cost, utc -from services.tracing.provider_dispatch import TraceExportResult - - -def export_trace(provider: str, config: Mapping[str, Any], trace: CanonicalTrace) -> TraceExportResult: - host = endpoint(config.get("host", "https://cloud.langfuse.com")) - public_key, secret_key = required(config, "public_key"), required(config, "secret_key") - auth = base64.b64encode(f"{public_key}:{secret_key}".encode()).decode("ascii") - root = trace.spans[0] - project_url = config.get("project_url") - trace_url = f"{endpoint(project_url)}/traces/{trace.trace_id}" if project_url else None - - def attributes(span: CanonicalSpan) -> dict[str, object]: - values: dict[str, object] = { - "langfuse.trace.name": root.name, - "langfuse.session.id": trace.session_id, - "langfuse.observation.type": "generation" if span.kind is CanonicalSpanKind.LLM else span.kind.value, - "langfuse.observation.input": json_value(span.inputs), - "langfuse.observation.output": json_value(span.outputs), - "langfuse.observation.level": "ERROR" if span.status is CanonicalSpanStatus.ERROR else "DEFAULT", - "dify.span.id": span.id, - } - for key, value in root.metadata.items(): - if value is not None: - values[f"langfuse.trace.metadata.{key}"] = value if isinstance(value, str) else json_value(value) - for key, value in span.metadata.items(): - if value is not None: - values[f"langfuse.observation.metadata.{key}"] = value if isinstance(value, str) else json_value(value) - user_id = span.metadata.get("user_id") or root.metadata.get("user_id") - if user_id: - values["langfuse.user.id"] = str(user_id) - if span.error: - values["langfuse.observation.status_message"] = span.error - if span.kind is CanonicalSpanKind.LLM: - model = span.metadata.get("model") or span.metadata.get("model_name") - if model: - values["langfuse.observation.model.name"] = str(model) - usage = token_usage(span) - if usage: - values["langfuse.observation.usage_details"] = json_value(usage) - cost = usd_cost(span) - if cost is not None: - values["langfuse.observation.cost_details"] = json_value({"total": cost}) - ttft = span.metadata.get("time_to_first_token", span.metadata.get("ttft")) - if isinstance(ttft, (int, float)) and ttft >= 0: - values["langfuse.observation.completion_start_time"] = ( - utc(span.start_time) + timedelta(seconds=ttft) - ).isoformat() - return values - - payload = build_request(trace, {"service.name": "dify-application"}, attributes) - export_http( - provider, - f"{host}/api/public/otel/v1/traces", - { - "Authorization": f"Basic {auth}", - "x-langfuse-ingestion-version": "4", - }, - payload, - ) - return TraceExportResult( - trace_id=trace.trace_id, - trace_url=trace_url, - span_urls={span.id: f"{trace_url}?observation={quote(span.id)}" for span in trace.spans} if trace_url else {}, - ) diff --git a/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/config.py b/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/config.py index 90d1a2846bfeb3..0564891d6f01e5 100644 --- a/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/config.py +++ b/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/config.py @@ -1,19 +1,23 @@ -from pydantic import ValidationInfo, field_validator +from dify_trace_sdk.config import ProviderConfig, derived, destination, http_url, secret +from pydantic import field_validator -from core.ops.entities.config_entity import BaseTracingConfig -from core.ops.utils import validate_url_with_path +DEFAULT_HOST = "https://api.langfuse.com" -class LangfuseConfig(BaseTracingConfig): - """ - Model class for Langfuse tracing config. - """ - - public_key: str - secret_key: str - host: str = "https://api.langfuse.com" +class LangfuseConfig(ProviderConfig): + public_key: str = secret() + secret_key: str = secret() + host: str = destination(default=DEFAULT_HOST) + project_key: str | None = derived() + # Older configurations stored the verified URL instead of the project ID. + project_url: str | None = derived() @field_validator("host") @classmethod - def host_validator(cls, v, info: ValidationInfo): - return validate_url_with_path(v, "https://api.langfuse.com") + def _host(cls, value: str) -> str: + return http_url(value, allowed_schemes=("https", "http")) if value.strip() else DEFAULT_HOST + + @field_validator("project_url") + @classmethod + def _project_url(cls, value: str | None) -> str | None: + return http_url(value, allowed_schemes=("https", "http")) if value and value.strip() else None diff --git a/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/entities/__init__.py b/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/entities/__init__.py deleted file mode 100644 index e69de29bb2d1d6..00000000000000 diff --git a/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/entities/langfuse_trace_entity.py b/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/entities/langfuse_trace_entity.py deleted file mode 100644 index 5b1a01e4babe72..00000000000000 --- a/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/entities/langfuse_trace_entity.py +++ /dev/null @@ -1,284 +0,0 @@ -from collections.abc import Mapping -from datetime import datetime -from enum import StrEnum -from typing import Any - -from pydantic import BaseModel, ConfigDict, Field, field_validator -from pydantic_core.core_schema import ValidationInfo - -from core.ops.utils import replace_text_with_content - - -def validate_input_output(v, field_name): - """ - Validate input output - :param v: - :param field_name: - :return: - """ - if v == {} or v is None: - return v - match v: - case str(): - return [ - { - "role": "assistant" if field_name == "output" else "user", - "content": v, - } - ] - case list(): - if len(v) > 0 and isinstance(v[0], dict): - v = replace_text_with_content(data=v) - return v - else: - return [ - { - "role": "assistant" if field_name == "output" else "user", - "content": str(v), - } - ] - - return v - - -class LevelEnum(StrEnum): - DEBUG = "DEBUG" - WARNING = "WARNING" - ERROR = "ERROR" - DEFAULT = "DEFAULT" - - -class LangfuseTrace(BaseModel): - """ - Langfuse trace model - """ - - id: str | None = Field( - default=None, - description="The id of the trace can be set, defaults to a random id. Used to link traces to external systems " - "or when creating a distributed trace. Traces are upserted on id.", - ) - name: str | None = Field( - default=None, - description="Identifier of the trace. Useful for sorting/filtering in the UI.", - ) - input: str | dict[str, Any] | list | None | None = Field( - default=None, description="The input of the trace. Can be any JSON object." - ) - output: str | dict[str, Any] | list | None | None = Field( - default=None, description="The output of the trace. Can be any JSON object." - ) - metadata: dict[str, Any] | None = Field( - default=None, - description="Additional metadata of the trace. Can be any JSON object. Metadata is merged when being updated " - "via the API.", - ) - user_id: str | None = Field( - default=None, - description="The id of the user that triggered the execution. Used to provide user-level analytics.", - ) - session_id: str | None = Field( - default=None, - description="Used to group multiple traces into a session in Langfuse. Use your own session/thread identifier.", - ) - version: str | None = Field( - default=None, - description="The version of the trace type. Used to understand how changes to the trace type affect metrics. " - "Useful in debugging.", - ) - release: str | None = Field( - default=None, - description="The release identifier of the current deployment. Used to understand how changes of different " - "deployments affect metrics. Useful in debugging.", - ) - tags: list[str] | None = Field( - default=None, - description="Tags are used to categorize or label traces. Traces can be filtered by tags in the UI and GET " - "API. Tags can also be changed in the UI. Tags are merged and never deleted via the API.", - ) - public: bool | None = Field( - default=None, - description="You can make a trace public to share it via a public link. This allows others to view the trace " - "without needing to log in or be members of your Langfuse project.", - ) - - @field_validator("input", "output") - @classmethod - def ensure_dict(cls, v, info: ValidationInfo): - field_name = info.field_name - return validate_input_output(v, field_name) - - -class LangfuseSpan(BaseModel): - """ - Langfuse span model - """ - - id: str | None = Field( - default=None, - description="The id of the span can be set, otherwise a random id is generated. Spans are upserted on id.", - ) - session_id: str | None = Field( - default=None, - description="Used to group multiple spans into a session in Langfuse. Use your own session/thread identifier.", - ) - trace_id: str | None = Field( - default=None, - description="The id of the trace the span belongs to. Used to link spans to traces.", - ) - user_id: str | None = Field( - default=None, - description="The id of the user that triggered the execution. Used to provide user-level analytics.", - ) - start_time: datetime | None = Field( - default_factory=datetime.now, - description="The time at which the span started, defaults to the current time.", - ) - end_time: datetime | None = Field( - default=None, - description="The time at which the span ended. Automatically set by span.end().", - ) - name: str | None = Field( - default=None, - description="Identifier of the span. Useful for sorting/filtering in the UI.", - ) - metadata: dict[str, Any] | None = Field( - default=None, - description="Additional metadata of the span. Can be any JSON object. Metadata is merged when being updated " - "via the API.", - ) - level: LevelEnum | None = Field( - default=None, - description="The level of the span. Can be DEBUG, DEFAULT, WARNING or ERROR. Used for sorting/filtering of " - "traces with elevated error levels and for highlighting in the UI.", - ) - status_message: str | None = Field( - default=None, - description="The status message of the span. Additional field for context of the event. E.g. the error " - "message of an error event.", - ) - input: str | Mapping[str, Any] | list | None | None = Field( - default=None, description="The input of the span. Can be any JSON object." - ) - output: str | Mapping[str, Any] | list | None | None = Field( - default=None, description="The output of the span. Can be any JSON object." - ) - version: str | None = Field( - default=None, - description="The version of the span type. Used to understand how changes to the span type affect metrics. " - "Useful in debugging.", - ) - parent_observation_id: str | None = Field( - default=None, - description="The id of the observation the span belongs to. Used to link spans to observations.", - ) - - @field_validator("input", "output") - @classmethod - def ensure_dict(cls, v, info: ValidationInfo): - field_name = info.field_name - return validate_input_output(v, field_name) - - -class UnitEnum(StrEnum): - CHARACTERS = "CHARACTERS" - TOKENS = "TOKENS" - SECONDS = "SECONDS" - MILLISECONDS = "MILLISECONDS" - IMAGES = "IMAGES" - - -class GenerationUsage(BaseModel): - promptTokens: int | None = None - completionTokens: int | None = None - total: int | None = None - input: int | None = None - output: int | None = None - unit: UnitEnum | None = None - inputCost: float | None = None - outputCost: float | None = None - totalCost: float | None = None - - @field_validator("input", "output") - @classmethod - def ensure_dict(cls, v, info: ValidationInfo): - field_name = info.field_name - return validate_input_output(v, field_name) - - -class LangfuseGeneration(BaseModel): - id: str | None = Field( - default=None, - description="The id of the generation can be set, defaults to random id.", - ) - trace_id: str | None = Field( - default=None, - description="The id of the trace the generation belongs to. Used to link generations to traces.", - ) - parent_observation_id: str | None = Field( - default=None, - description="The id of the observation the generation belongs to. Used to link generations to observations.", - ) - name: str | None = Field( - default=None, - description="Identifier of the generation. Useful for sorting/filtering in the UI.", - ) - start_time: datetime | None = Field( - default_factory=datetime.now, - description="The time at which the generation started, defaults to the current time.", - ) - completion_start_time: datetime | None = Field( - default=None, - description="The time at which the completion started (streaming). Set it to get latency analytics broken " - "down into time until completion started and completion duration.", - ) - end_time: datetime | None = Field( - default=None, - description="The time at which the generation ended. Automatically set by generation.end().", - ) - model: str | None = Field(default=None, description="The name of the model used for the generation.") - model_parameters: dict[str, Any] | None = Field( - default=None, - description="The parameters of the model used for the generation; can be any key-value pairs.", - ) - input: Any | None = Field( - default=None, - description="The prompt used for the generation. Can be any string or JSON object.", - ) - output: Any | None = Field( - default=None, - description="The completion generated by the model. Can be any string or JSON object.", - ) - usage: GenerationUsage | None = Field( - default=None, - description="The usage object supports the OpenAi structure with tokens and a more generic version with " - "detailed costs and units.", - ) - metadata: dict[str, Any] | None = Field( - default=None, - description="Additional metadata of the generation. Can be any JSON object. Metadata is merged when being " - "updated via the API.", - ) - level: LevelEnum | None = Field( - default=None, - description="The level of the generation. Can be DEBUG, DEFAULT, WARNING or ERROR. Used for sorting/filtering " - "of traces with elevated error levels and for highlighting in the UI.", - ) - status_message: str | None = Field( - default=None, - description="The status message of the generation. Additional field for context of the event. E.g. the error " - "message of an error event.", - ) - version: str | None = Field( - default=None, - description="The version of the generation type. Used to understand how changes to the span type affect " - "metrics. Useful in debugging.", - ) - - model_config = ConfigDict(protected_namespaces=()) - - @field_validator("input", "output") - @classmethod - def ensure_dict(cls, v, info: ValidationInfo): - field_name = info.field_name - return validate_input_output(v, field_name) diff --git a/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/langfuse_trace.py b/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/langfuse_trace.py deleted file mode 100644 index 728a63a82c84c6..00000000000000 --- a/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/langfuse_trace.py +++ /dev/null @@ -1,600 +0,0 @@ -import logging -import os -import uuid -from datetime import UTC, datetime, timedelta -from typing import override - -import httpx -from langfuse import __version__ as langfuse_version -from langfuse.api import ( - CreateGenerationBody, - CreateSpanBody, - IngestionEvent_GenerationCreate, - IngestionEvent_SpanCreate, - IngestionEvent_TraceCreate, - LangfuseAPI, - TraceBody, -) -from langfuse.api.commons.types.usage import Usage -from sqlalchemy.orm import sessionmaker - -from core.ops.base_trace_instance import BaseTraceInstance -from core.ops.entities.trace_entity import ( - BaseTraceInfo, - DatasetRetrievalTraceInfo, - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - SuggestedQuestionTraceInfo, - ToolTraceInfo, - TraceTaskName, - WorkflowTraceInfo, -) -from core.ops.utils import filter_none_values -from core.repositories import DifyCoreRepositoryFactory -from dify_trace_langfuse.config import LangfuseConfig -from dify_trace_langfuse.entities.langfuse_trace_entity import ( - GenerationUsage, - LangfuseGeneration, - LangfuseSpan, - LangfuseTrace, - LevelEnum, - UnitEnum, -) -from extensions.ext_database import db -from graphon.enums import BuiltinNodeTypes -from models import EndUser, WorkflowNodeExecutionTriggeredFrom -from models.enums import MessageStatus - -logger = logging.getLogger(__name__) - - -class LangFuseDataTrace(BaseTraceInstance): - def __init__( - self, - langfuse_config: LangfuseConfig, - ): - super().__init__(langfuse_config) - timeout = int(os.environ.get("LANGFUSE_TIMEOUT", 5)) - self._http_client: httpx.Client | None = httpx.Client(timeout=timeout) - self.langfuse_client = LangfuseAPI( - base_url=langfuse_config.host, - username=langfuse_config.public_key, - password=langfuse_config.secret_key, - x_langfuse_sdk_name="python", - x_langfuse_sdk_version=langfuse_version, - x_langfuse_public_key=langfuse_config.public_key, - timeout=timeout, - httpx_client=self._http_client, - ) - self.file_base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001") - - def close(self) -> None: - client = getattr(self, "_http_client", None) - if client is None: - return - self._http_client = None - try: - client.close() - except Exception: - logger.debug("Failed to close Langfuse HTTP client", exc_info=True) - - def __del__(self) -> None: - self.close() - - @staticmethod - def _get_completion_start_time( - start_time: datetime | None, time_to_first_token: float | int | None - ) -> datetime | None: - """Convert a relative TTFT value in seconds into Langfuse's absolute completion start time.""" - if start_time is None or time_to_first_token is None: - return None - - try: - ttft_seconds = float(time_to_first_token) - except (TypeError, ValueError): - return None - - if ttft_seconds < 0: - return None - - return start_time + timedelta(seconds=ttft_seconds) - - @override - def trace(self, trace_info: BaseTraceInfo): - match trace_info: - case WorkflowTraceInfo(): - self.workflow_trace(trace_info) - case MessageTraceInfo(): - self.message_trace(trace_info) - case ModerationTraceInfo(): - self.moderation_trace(trace_info) - case SuggestedQuestionTraceInfo(): - self.suggested_question_trace(trace_info) - case DatasetRetrievalTraceInfo(): - self.dataset_retrieval_trace(trace_info) - case ToolTraceInfo(): - self.tool_trace(trace_info) - case GenerateNameTraceInfo(): - self.generate_name_trace(trace_info) - case _: - pass - - def workflow_trace(self, trace_info: WorkflowTraceInfo): - trace_id = trace_info.trace_id or trace_info.workflow_run_id - user_id = trace_info.metadata.get("user_id") - metadata = trace_info.metadata - metadata["workflow_app_log_id"] = trace_info.workflow_app_log_id - - if trace_info.message_id: - trace_id = trace_info.trace_id or trace_info.message_id - name = TraceTaskName.MESSAGE_TRACE - trace_data = LangfuseTrace( - id=trace_id, - user_id=user_id, - name=name, - input=dict(trace_info.workflow_run_inputs), - output=dict(trace_info.workflow_run_outputs), - metadata=metadata, - session_id=trace_info.conversation_id, - tags=["message", "workflow"], - version=trace_info.workflow_run_version, - ) - self.add_trace(langfuse_trace_data=trace_data) - workflow_span_data = LangfuseSpan( - id=trace_info.workflow_run_id, - name=TraceTaskName.WORKFLOW_TRACE, - input=dict(trace_info.workflow_run_inputs), - output=dict(trace_info.workflow_run_outputs), - trace_id=trace_id, - start_time=trace_info.start_time, - end_time=trace_info.end_time, - metadata=metadata, - level=LevelEnum.DEFAULT if trace_info.error == "" else LevelEnum.ERROR, - status_message=trace_info.error or "", - ) - self.add_span(langfuse_span_data=workflow_span_data) - else: - trace_data = LangfuseTrace( - id=trace_id, - user_id=user_id, - name=TraceTaskName.WORKFLOW_TRACE, - input=dict(trace_info.workflow_run_inputs), - output=dict(trace_info.workflow_run_outputs), - metadata=metadata, - session_id=trace_info.conversation_id, - tags=["workflow"], - version=trace_info.workflow_run_version, - ) - self.add_trace(langfuse_trace_data=trace_data) - - # through workflow_run_id get all_nodes_execution using repository - session_factory = sessionmaker(bind=db.engine) - # Find the app's creator account - app_id = trace_info.metadata.get("app_id") - if not app_id: - raise ValueError("No app_id found in trace_info metadata") - - service_account = self.get_service_account_with_tenant(app_id) - - workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( - session_factory=session_factory, - tenant_id=trace_info.tenant_id, - user=service_account, - app_id=app_id, - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - - # Get all executions for this workflow run - workflow_node_executions = workflow_node_execution_repository.get_by_workflow_execution( - workflow_execution_id=trace_info.workflow_run_id - ) - - for node_execution in workflow_node_executions: - node_execution_id = node_execution.id - tenant_id = trace_info.tenant_id # Use from trace_info instead - app_id = trace_info.metadata.get("app_id") # Use from trace_info instead - node_name = node_execution.title - node_type = node_execution.node_type - status = node_execution.status - if node_type == BuiltinNodeTypes.LLM: - inputs = node_execution.process_data.get("prompts", {}) if node_execution.process_data else {} - else: - inputs = node_execution.inputs or {} - outputs = node_execution.outputs or {} - created_at = node_execution.created_at or datetime.now() - elapsed_time = node_execution.elapsed_time - finished_at = created_at + timedelta(seconds=elapsed_time) - - execution_metadata = node_execution.metadata or {} - metadata = {str(k): v for k, v in execution_metadata.items()} - metadata.update( - { - "workflow_run_id": trace_info.workflow_run_id, - "node_execution_id": node_execution_id, - "tenant_id": tenant_id, - "app_id": app_id, - "node_name": node_name, - "node_type": node_type, - "status": status, - } - ) - process_data = node_execution.process_data or {} - model_provider = process_data.get("model_provider", None) - model_name = process_data.get("model_name", None) - if model_provider is not None and model_name is not None: - metadata.update( - { - "model_provider": model_provider, - "model_name": model_name, - } - ) - - # add generation span - if process_data and process_data.get("model_mode") == "chat": - total_token = metadata.get("total_tokens", 0) - prompt_tokens = 0 - completion_tokens = 0 - completion_start_time = None - try: - usage_data = process_data.get("usage") - if not isinstance(usage_data, dict): - usage_data = outputs.get("usage") - if not isinstance(usage_data, dict): - usage_data = {} - prompt_tokens = usage_data.get("prompt_tokens", 0) - completion_tokens = usage_data.get("completion_tokens", 0) - completion_start_time = self._get_completion_start_time( - created_at, usage_data.get("time_to_first_token") - ) - except Exception: - logger.error("Failed to extract usage", exc_info=True) - - # add generation - generation_usage = GenerationUsage( - input=prompt_tokens, - output=completion_tokens, - total=total_token, - unit=UnitEnum.TOKENS, - ) - - node_generation_data = LangfuseGeneration( - id=node_execution_id, - name=node_name, - trace_id=trace_id, - model=process_data.get("model_name"), - start_time=created_at, - completion_start_time=completion_start_time, - end_time=finished_at, - input=inputs, - output=outputs, - metadata=metadata, - level=(LevelEnum.DEFAULT if status == "succeeded" else LevelEnum.ERROR), - status_message=trace_info.error or "", - parent_observation_id=trace_info.workflow_run_id if trace_info.message_id else None, - usage=generation_usage, - ) - - self.add_generation(langfuse_generation_data=node_generation_data) - - # add normal span - else: - span_data = LangfuseSpan( - id=node_execution_id, - name=node_name, - input=inputs, - output=outputs, - trace_id=trace_id, - start_time=created_at, - end_time=finished_at, - metadata=metadata, - level=(LevelEnum.DEFAULT if status == "succeeded" else LevelEnum.ERROR), - status_message=trace_info.error or "", - parent_observation_id=trace_info.workflow_run_id if trace_info.message_id else None, - ) - - self.add_span(langfuse_span_data=span_data) - - def message_trace(self, trace_info: MessageTraceInfo, **kwargs): - # get message file data - file_list = trace_info.file_list - metadata = trace_info.metadata - message_data = trace_info.message_data - if message_data is None: - return - message_id = message_data.id - - user_id = message_data.from_account_id - if message_data.from_end_user_id: - end_user_data: EndUser | None = db.session.get(EndUser, message_data.from_end_user_id) - if end_user_data is not None: - user_id = end_user_data.session_id - metadata["user_id"] = user_id - - trace_id = trace_info.trace_id or message_id - - trace_data = LangfuseTrace( - id=trace_id, - user_id=user_id, - name=TraceTaskName.MESSAGE_TRACE, - input={ - "message": trace_info.inputs, - "files": file_list, - "message_tokens": trace_info.message_tokens, - "answer_tokens": trace_info.answer_tokens, - "total_tokens": trace_info.total_tokens, - "error": trace_info.error, - "provider_response_latency": message_data.provider_response_latency, - "created_at": trace_info.start_time, - }, - output=trace_info.outputs, - metadata=metadata, - session_id=message_data.conversation_id, - tags=["message", str(trace_info.conversation_mode)], - version=None, - release=None, - public=None, - ) - self.add_trace(langfuse_trace_data=trace_data) - - # add generation - generation_usage = GenerationUsage( - input=trace_info.message_tokens, - output=trace_info.answer_tokens, - total=trace_info.total_tokens, - unit=UnitEnum.TOKENS, - totalCost=message_data.total_price, - ) - completion_start_time = self._get_completion_start_time( - trace_info.start_time, - trace_info.gen_ai_server_time_to_first_token, - ) - - langfuse_generation_data = LangfuseGeneration( - name="llm", - trace_id=trace_id, - start_time=trace_info.start_time, - completion_start_time=completion_start_time, - end_time=trace_info.end_time, - model=message_data.model_id, - input=trace_info.inputs, - output=message_data.answer, - metadata=metadata, - level=(LevelEnum.DEFAULT if message_data.status != MessageStatus.ERROR else LevelEnum.ERROR), - status_message=message_data.error or "", - usage=generation_usage, - ) - - self.add_generation(langfuse_generation_data) - - def moderation_trace(self, trace_info: ModerationTraceInfo): - if trace_info.message_data is None: - return - span_data = LangfuseSpan( - name=TraceTaskName.MODERATION_TRACE, - input=trace_info.inputs, - output={ - "action": trace_info.action, - "flagged": trace_info.flagged, - "preset_response": trace_info.preset_response, - "inputs": trace_info.inputs, - }, - trace_id=trace_info.trace_id or trace_info.message_id, - start_time=trace_info.start_time or trace_info.message_data.created_at, - end_time=trace_info.end_time or trace_info.message_data.created_at, - metadata=trace_info.metadata, - ) - - self.add_span(langfuse_span_data=span_data) - - def suggested_question_trace(self, trace_info: SuggestedQuestionTraceInfo): - message_data = trace_info.message_data - if message_data is None: - return - generation_usage = GenerationUsage( - total=len(str(trace_info.suggested_question)), - input=len(trace_info.inputs) if trace_info.inputs else 0, - output=len(trace_info.suggested_question), - unit=UnitEnum.CHARACTERS, - ) - - generation_data = LangfuseGeneration( - name=TraceTaskName.SUGGESTED_QUESTION_TRACE, - input=trace_info.inputs, - output=str(trace_info.suggested_question), - trace_id=trace_info.trace_id or trace_info.message_id, - start_time=trace_info.start_time, - end_time=trace_info.end_time, - metadata=trace_info.metadata, - level=(LevelEnum.DEFAULT if message_data.status != MessageStatus.ERROR else LevelEnum.ERROR), - status_message=message_data.error or "", - usage=generation_usage, - ) - - self.add_generation(langfuse_generation_data=generation_data) - - def dataset_retrieval_trace(self, trace_info: DatasetRetrievalTraceInfo): - if trace_info.message_data is None: - return - dataset_retrieval_span_data = LangfuseSpan( - name=TraceTaskName.DATASET_RETRIEVAL_TRACE, - input=trace_info.inputs, - output={"documents": trace_info.documents}, - trace_id=trace_info.trace_id or trace_info.message_id, - start_time=trace_info.start_time or trace_info.message_data.created_at, - end_time=trace_info.end_time or trace_info.message_data.updated_at, - metadata=trace_info.metadata, - ) - - self.add_span(langfuse_span_data=dataset_retrieval_span_data) - - def tool_trace(self, trace_info: ToolTraceInfo): - tool_span_data = LangfuseSpan( - name=trace_info.tool_name, - input=trace_info.tool_inputs, - output=trace_info.tool_outputs, - trace_id=trace_info.trace_id or trace_info.message_id, - start_time=trace_info.start_time, - end_time=trace_info.end_time, - metadata=trace_info.metadata, - level=(LevelEnum.DEFAULT if trace_info.error == "" or trace_info.error is None else LevelEnum.ERROR), - status_message=trace_info.error, - ) - - self.add_span(langfuse_span_data=tool_span_data) - - def generate_name_trace(self, trace_info: GenerateNameTraceInfo): - name_generation_trace_data = LangfuseTrace( - name=TraceTaskName.GENERATE_NAME_TRACE, - input=trace_info.inputs, - output=trace_info.outputs, - user_id=trace_info.tenant_id, - metadata=trace_info.metadata, - session_id=trace_info.conversation_id, - ) - - self.add_trace(langfuse_trace_data=name_generation_trace_data) - - name_generation_span_data = LangfuseSpan( - name=TraceTaskName.GENERATE_NAME_TRACE, - input=trace_info.inputs, - output=trace_info.outputs, - trace_id=trace_info.conversation_id, - start_time=trace_info.start_time, - end_time=trace_info.end_time, - metadata=trace_info.metadata, - ) - self.add_span(langfuse_span_data=name_generation_span_data) - - def _make_event_id(self) -> str: - return str(uuid.uuid4()) - - def _now_iso(self) -> str: - return datetime.now(UTC).isoformat() - - def add_trace(self, langfuse_trace_data: LangfuseTrace | None = None): - data = filter_none_values(langfuse_trace_data.model_dump()) if langfuse_trace_data else {} - try: - body = TraceBody( - id=data.get("id"), - name=data.get("name"), - user_id=data.get("user_id"), - input=data.get("input"), - output=data.get("output"), - metadata=data.get("metadata"), - session_id=data.get("session_id"), - version=data.get("version"), - release=data.get("release"), - tags=data.get("tags"), - public=data.get("public"), - ) - event = IngestionEvent_TraceCreate( - body=body, - id=self._make_event_id(), - timestamp=self._now_iso(), - ) - self.langfuse_client.ingestion.batch(batch=[event]) - logger.debug("LangFuse Trace created successfully") - except Exception as e: - raise ValueError(f"LangFuse Failed to create trace: {str(e)}") from e - - def add_span(self, langfuse_span_data: LangfuseSpan | None = None): - data = filter_none_values(langfuse_span_data.model_dump()) if langfuse_span_data else {} - try: - body = CreateSpanBody( - id=data.get("id"), - trace_id=data.get("trace_id"), - name=data.get("name"), - start_time=data.get("start_time"), - end_time=data.get("end_time"), - input=data.get("input"), - output=data.get("output"), - metadata=data.get("metadata"), - level=data.get("level"), - status_message=data.get("status_message"), - parent_observation_id=data.get("parent_observation_id"), - version=data.get("version"), - ) - event = IngestionEvent_SpanCreate( - body=body, - id=self._make_event_id(), - timestamp=self._now_iso(), - ) - self.langfuse_client.ingestion.batch(batch=[event]) - logger.debug("LangFuse Span created successfully") - except Exception as e: - raise ValueError(f"LangFuse Failed to create span: {str(e)}") from e - - def update_span(self, span, langfuse_span_data: LangfuseSpan | None = None): - format_span_data = filter_none_values(langfuse_span_data.model_dump()) if langfuse_span_data else {} - - span.end(**format_span_data) - - def add_generation(self, langfuse_generation_data: LangfuseGeneration | None = None): - data = filter_none_values(langfuse_generation_data.model_dump()) if langfuse_generation_data else {} - try: - usage_data = data.pop("usage", None) - usage = None - if usage_data: - usage = Usage( - input=usage_data.get("input", 0) or 0, - output=usage_data.get("output", 0) or 0, - total=usage_data.get("total", 0) or 0, - unit=usage_data.get("unit"), - input_cost=usage_data.get("inputCost"), - output_cost=usage_data.get("outputCost"), - total_cost=usage_data.get("totalCost"), - ) - - body = CreateGenerationBody( - id=data.get("id"), - trace_id=data.get("trace_id"), - name=data.get("name"), - start_time=data.get("start_time"), - end_time=data.get("end_time"), - model=data.get("model"), - model_parameters=data.get("model_parameters"), - input=data.get("input"), - output=data.get("output"), - usage=usage, - metadata=data.get("metadata"), - level=data.get("level"), - status_message=data.get("status_message"), - parent_observation_id=data.get("parent_observation_id"), - version=data.get("version"), - completion_start_time=data.get("completion_start_time"), - ) - event = IngestionEvent_GenerationCreate( - body=body, - id=self._make_event_id(), - timestamp=self._now_iso(), - ) - self.langfuse_client.ingestion.batch(batch=[event]) - logger.debug("LangFuse Generation created successfully") - except Exception as e: - raise ValueError(f"LangFuse Failed to create generation: {str(e)}") from e - - def update_generation(self, generation, langfuse_generation_data: LangfuseGeneration | None = None): - format_generation_data = ( - filter_none_values(langfuse_generation_data.model_dump()) if langfuse_generation_data else {} - ) - - generation.end(**format_generation_data) - - def api_check(self): - try: - projects = self.langfuse_client.projects.get() - except Exception as e: - logger.debug("LangFuse API check failed", exc_info=True) - raise ValueError(f"LangFuse API check failed: {str(e)}") from e - if not projects.data: - raise ValueError("LangFuse API check failed: no project found for the provided credentials") - return True - - def get_project_key(self): - try: - projects = self.langfuse_client.projects.get() - return projects.data[0].id - except Exception as e: - logger.debug("LangFuse get project key failed", exc_info=True) - raise ValueError(f"LangFuse get project key failed: {str(e)}") from e diff --git a/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/provider.py b/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/provider.py new file mode 100644 index 00000000000000..620ecd12388c87 --- /dev/null +++ b/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/provider.py @@ -0,0 +1,240 @@ +"""Langfuse destination over the public ingestion API. + +Each trace unit becomes ingestion events: the root span is the Langfuse trace +and every other span an observation (a generation for LLM calls). Langfuse +upserts by ID, so re-sending a unit after a partial failure is safe. +""" + +import base64 +from collections.abc import Mapping, Sequence +from datetime import datetime, timedelta +from typing import override +from urllib.parse import quote +from uuid import uuid5 + +from dify_trace_sdk.errors import ( + EgressError, + ExportOutcomeUnknownError, + ExportRejectedError, + ExportRetryableError, + VerificationError, +) +from dify_trace_sdk.exporting import export_endpoint, json_value, send_export_request +from dify_trace_sdk.ids import span_uuid, trace_uuid +from dify_trace_sdk.model import Span, SpanKind, SpanStatus, Trace +from dify_trace_sdk.provider import ProviderCapabilities, SendResult, TraceProvider +from pydantic import JsonValue + +from dify_trace_langfuse.config import LangfuseConfig + +VERIFY_TIMEOUT_SECONDS = 10.0 +MAX_BATCH_BYTES = 3_000_000 +"""Stay below Langfuse Cloud's ingestion request limit.""" + + +def _iso(value: datetime) -> str: + return value.isoformat() + + +def _messages(value: JsonValue, *, role: str) -> JsonValue: + """Preserve Langfuse's chat preview without rewriting multimodal content blocks.""" + if isinstance(value, str): + return [{"role": role, "content": value}] + if isinstance(value, list): + return [ + {("content" if key == "text" else key): part for key, part in item.items()} + if isinstance(item, dict) and "role" in item and "content" not in item + else item + for item in value + ] + return value + + +class LangfuseProvider(TraceProvider[LangfuseConfig]): + id = "langfuse" + config_model = LangfuseConfig + capabilities = ProviderCapabilities(native_cost=True, trace_url=True, span_url=True) + + @staticmethod + def _headers(config: LangfuseConfig) -> dict[str, str]: + token = base64.b64encode(f"{config.public_key}:{config.secret_key}".encode()).decode() + return {"Authorization": f"Basic {token}", "Accept": "application/json"} + + @override + def verify(self, config: LangfuseConfig) -> LangfuseConfig: + try: + response = self.egress.request( + "GET", + f"{config.host.rstrip('/')}/api/public/projects", + headers=self._headers(config), + content=None, + timeout=VERIFY_TIMEOUT_SECONDS, + sensitive_url=False, + ) + except EgressError: + raise VerificationError("connection_failed") from None + if response.status_code in {401, 403}: + raise VerificationError("authentication_failed") + if not 200 <= response.status_code < 300: + raise VerificationError("endpoint_unavailable") + try: + body = response.json() + projects = body["data"] if isinstance(body, dict) else None + project_id = projects[0]["id"] if isinstance(projects, list) and projects else None + except (ValueError, KeyError, TypeError): + raise VerificationError("malformed_response") from None + if not isinstance(project_id, str) or not project_id: + raise VerificationError("project_not_found") + return config.model_copy(update={"project_key": project_id}) + + @override + def project_url(self, config: LangfuseConfig) -> str: + host = config.host.rstrip("/") + return f"{host}/project/{config.project_key}" if config.project_key else config.project_url or host + + @override + def send(self, config: LangfuseConfig, trace: Trace) -> SendResult: + host = export_endpoint(config.host) + trace_key = trace.external_trace_id or str(trace_uuid(trace.trace_id)) + events = self._events(trace, trace_key) + for batch in self._batches(events): + try: + response = send_export_request( + self.egress, + "POST", + f"{host}/api/public/ingestion", + destination="Langfuse", + headers={**self._headers(config), "Content-Type": "application/json"}, + content=batch, + sensitive_url=False, + ) + except ExportOutcomeUnknownError: + raise ExportRetryableError("Langfuse ingestion outcome is unknown; re-sending is idempotent") from None + try: + body = response.json() if response.content else {} + except ValueError: + # Something between Dify and Langfuse answered; re-sending is idempotent. + raise ExportRetryableError("Langfuse ingestion answered with an unreadable body") from None + self._check_results(body) + trace_url = ( + f"{self.project_url(config).rstrip('/')}/traces/{quote(trace_key, safe='')}" + if config.project_key or config.project_url + else None + ) + span_urls = ( + { + span.span_id: f"{trace_url}?observation={self._observation_id(trace, span.span_id)}" + for span in trace.spans + } + if trace_url + else {} + ) + return SendResult(trace_id=trace_key, trace_url=trace_url, span_urls=span_urls) + + @staticmethod + def _observation_id(trace: Trace, span_id: str) -> str: + return str(span_uuid(trace.trace_id, span_id)) + + def _events(self, trace: Trace, trace_key: str) -> list[dict[str, object]]: + events: list[dict[str, object]] = [] + if trace.parent is None: + root = trace.root + events.append( + self._event( + trace, + "trace-create", + { + "id": trace_key, + "name": root.name, + "timestamp": _iso(root.start_time), + "userId": trace.user_id, + "sessionId": trace.session_id, + "input": _messages(root.inputs, role="user"), + "output": _messages(root.outputs, role="assistant"), + "metadata": {**root.attributes, "dify_trace_id": trace.trace_id}, + }, + ) + ) + # The trace itself represents the root; its children are top-level observations. + trace_level_span_id = root.span_id + observations: Sequence[Span] = trace.spans[1:] + else: + trace_level_span_id = trace.parent.span_id + observations = trace.spans + for span in observations: + parent = ( + None + if span.parent_span_id == trace_level_span_id or span.parent_span_id is None + else self._observation_id(trace, span.parent_span_id) + ) + body: dict[str, object] = { + "id": self._observation_id(trace, span.span_id), + "traceId": trace_key, + "parentObservationId": parent, + "name": span.name, + "startTime": _iso(span.start_time), + "endTime": _iso(span.end_time), + "input": _messages(span.inputs, role="user"), + "output": _messages(span.outputs, role="assistant"), + "metadata": {**span.attributes, "dify_span_kind": span.kind.value}, + "level": "ERROR" if span.status is SpanStatus.ERROR else "DEFAULT", + "statusMessage": span.error, + } + if span.kind is SpanKind.LLM: + body.update(self._generation_fields(span)) + events.append(self._event(trace, "generation-create", body)) + else: + events.append(self._event(trace, "span-create", body)) + return events + + @staticmethod + def _generation_fields(span: Span) -> dict[str, object]: + fields: dict[str, object] = {"model": span.model_name} + if span.usage is not None: + usage: dict[str, object] = { + "input": span.usage.prompt_tokens, + "output": span.usage.completion_tokens, + "total": span.usage.total_tokens, + "unit": "TOKENS", + } + if span.cost is not None and span.cost.currency.upper() == "USD": + usage["totalCost"] = span.cost.amount + fields["usage"] = usage + if span.time_to_first_token is not None: + fields["completionStartTime"] = _iso(span.start_time + timedelta(seconds=span.time_to_first_token)) + return fields + + @staticmethod + def _event(trace: Trace, event_type: str, body: Mapping[str, object]) -> dict[str, object]: + # Deterministic event IDs let Langfuse deduplicate a redelivered unit. + return { + "id": str(uuid5(trace_uuid(trace.trace_id), f"{event_type}:{body['id']}")), + "type": event_type, + "timestamp": body.get("startTime") or body.get("timestamp"), + "body": {key: value for key, value in body.items() if value is not None}, + } + + @staticmethod + def _batches(events: list[dict[str, object]]) -> list[bytes]: + batches: list[bytes] = [] + current: list[bytes] = [] + size = 0 + for event in events: + encoded = json_value(event).encode() + if current and size + len(encoded) > MAX_BATCH_BYTES: + batches.append(b'{"batch":[' + b",".join(current) + b"]}") + current, size = [], 0 + current.append(encoded) + size += len(encoded) + 1 + batches.append(b'{"batch":[' + b",".join(current) + b"]}") + return batches + + @staticmethod + def _check_results(body: object) -> None: + errors = body.get("errors") if isinstance(body, dict) else None + if not isinstance(errors, list) or not errors: + return + statuses = [error.get("status") for error in errors if isinstance(error, dict)] + if any(isinstance(status, int) and (status == 429 or status >= 500) for status in statuses): + raise ExportRetryableError("Langfuse could not ingest every event; re-sending is idempotent") + raise ExportRejectedError(f"Langfuse rejected {len(errors)} ingestion events") diff --git a/api/providers/trace/trace-langfuse/tests/unit_tests/src/dify_trace_langfuse/test_config.py b/api/providers/trace/trace-langfuse/tests/unit_tests/src/dify_trace_langfuse/test_config.py index 0c3c3fc81e79b8..9518995a62e1cc 100644 --- a/api/providers/trace/trace-langfuse/tests/unit_tests/src/dify_trace_langfuse/test_config.py +++ b/api/providers/trace/trace-langfuse/tests/unit_tests/src/dify_trace_langfuse/test_config.py @@ -3,40 +3,15 @@ from pydantic import ValidationError -class TestLangfuseConfig: - """Test cases for LangfuseConfig""" - - def test_valid_config(self): - """Test valid Langfuse configuration""" - config = LangfuseConfig(public_key="public_key", secret_key="secret_key", host="https://custom.langfuse.com") - assert config.public_key == "public_key" - assert config.secret_key == "secret_key" - assert config.host == "https://custom.langfuse.com" - - def test_valid_config_with_path(self): - host = "https://custom.langfuse.com/api/v1" - config = LangfuseConfig(public_key="public_key", secret_key="secret_key", host=host) - assert config.public_key == "public_key" - assert config.secret_key == "secret_key" - assert config.host == host - - def test_default_values(self): - """Test default values are set correctly""" - config = LangfuseConfig(public_key="public", secret_key="secret") - assert config.host == "https://api.langfuse.com" - - def test_missing_required_fields(self): - """Test that required fields are enforced""" - with pytest.raises(ValidationError): - LangfuseConfig.model_validate({}) - - with pytest.raises(ValidationError): - LangfuseConfig.model_validate({"public_key": "public"}) - - with pytest.raises(ValidationError): - LangfuseConfig.model_validate({"secret_key": "secret"}) - - def test_host_validation_empty(self): - """Test host validation with empty value""" - config = LangfuseConfig(public_key="public", secret_key="secret", host="") - assert config.host == "https://api.langfuse.com" +def test_config_secrets_and_defaults() -> None: + assert LangfuseConfig.secret_fields() == {"public_key", "secret_key"} + assert LangfuseConfig.derived_fields() == {"project_key", "project_url"} + assert LangfuseConfig(public_key="p", secret_key="s", host=" ").host == "https://api.langfuse.com" + with pytest.raises(ValidationError): + LangfuseConfig(public_key="p", secret_key="s", host="ftp://cloud.langfuse.com") + + +@pytest.mark.parametrize("project_url", ["javascript:alert(1)", "https://user:secret@langfuse.example/project/1"]) +def test_stored_project_urls_must_be_safe_links(project_url: str) -> None: + with pytest.raises(ValidationError): + LangfuseConfig(public_key="p", secret_key="s", project_url=project_url) diff --git a/api/providers/trace/trace-langfuse/tests/unit_tests/src/dify_trace_langfuse/test_langfuse_trace.py b/api/providers/trace/trace-langfuse/tests/unit_tests/src/dify_trace_langfuse/test_langfuse_trace.py deleted file mode 100644 index bf13dad521eee9..00000000000000 --- a/api/providers/trace/trace-langfuse/tests/unit_tests/src/dify_trace_langfuse/test_langfuse_trace.py +++ /dev/null @@ -1,911 +0,0 @@ -import collections -import logging -from datetime import UTC, datetime, timedelta -from types import SimpleNamespace -from typing import cast, override -from unittest.mock import MagicMock, patch - -import pytest -from dify_trace_langfuse.config import LangfuseConfig -from dify_trace_langfuse.entities.langfuse_trace_entity import ( - LangfuseGeneration, - LangfuseSpan, - LangfuseTrace, - LevelEnum, - UnitEnum, -) -from dify_trace_langfuse.langfuse_trace import LangFuseDataTrace -from sqlalchemy.orm import Session - -from core.ops.entities.trace_entity import ( - DatasetRetrievalTraceInfo, - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - SuggestedQuestionTraceInfo, - ToolTraceInfo, - TraceTaskName, - WorkflowTraceInfo, -) -from graphon.enums import BuiltinNodeTypes -from models import EndUser -from models.enums import EndUserType, MessageStatus -from tests.helpers.unit.core.ops.trace_fixtures import message_trace_info, tool_trace_info, workflow_trace_info - -"""Unit tests for Langfuse trace translation with real SQLite-backed lookups.""" - - -def _dt() -> datetime: - return datetime(2024, 1, 1, 0, 0, 0, tzinfo=UTC) - - -class TestLangfuseTraceCases: - @staticmethod - @pytest.fixture - def langfuse_config(): - return LangfuseConfig(public_key="pk-123", secret_key="sk-123", host="https://cloud.langfuse.com") - - @staticmethod - @pytest.fixture - def trace_instance(langfuse_config, monkeypatch: pytest.MonkeyPatch): - # Mock Langfuse client to avoid network calls - mock_client = MagicMock() - monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.LangfuseAPI", lambda **kwargs: mock_client) - - instance = LangFuseDataTrace(langfuse_config) - return instance - - @staticmethod - def test_init(langfuse_config, monkeypatch: pytest.MonkeyPatch): - mock_langfuse = MagicMock() - monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.LangfuseAPI", mock_langfuse) - monkeypatch.setenv("FILES_URL", "http://test.url") - - instance = LangFuseDataTrace(langfuse_config) - - mock_langfuse.assert_called_once() - kwargs = mock_langfuse.call_args.kwargs - assert kwargs["username"] == langfuse_config.public_key - assert kwargs["password"] == langfuse_config.secret_key - assert kwargs["base_url"] == langfuse_config.host - assert instance.file_base_url == "http://test.url" - - @staticmethod - def test_same_public_key_uses_independent_api_clients(monkeypatch: pytest.MonkeyPatch): - clients = [MagicMock(), MagicMock()] - create_client = MagicMock(side_effect=clients) - monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.LangfuseAPI", create_client) - - first = LangFuseDataTrace( - LangfuseConfig(public_key="shared", secret_key="secret-a", host="https://tenant-a.example") - ) - second = LangFuseDataTrace( - LangfuseConfig(public_key="shared", secret_key="secret-b", host="https://tenant-b.example") - ) - - assert first.langfuse_client is clients[0] - assert second.langfuse_client is clients[1] - assert [ - (call.kwargs["base_url"], call.kwargs["username"], call.kwargs["password"]) - for call in create_client.call_args_list - ] == [ - ("https://tenant-a.example", "shared", "secret-a"), - ("https://tenant-b.example", "shared", "secret-b"), - ] - - @staticmethod - def test_close_is_idempotent(langfuse_config, monkeypatch: pytest.MonkeyPatch): - http_client = MagicMock() - monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.httpx.Client", lambda **kwargs: http_client) - monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.LangfuseAPI", MagicMock()) - - instance = LangFuseDataTrace(langfuse_config) - instance.close() - instance.close() - - http_client.close.assert_called_once() - - @staticmethod - def test_trace_dispatch(trace_instance, monkeypatch: pytest.MonkeyPatch): - methods = [ - "workflow_trace", - "message_trace", - "moderation_trace", - "suggested_question_trace", - "dataset_retrieval_trace", - "tool_trace", - "generate_name_trace", - ] - mocks = {method: MagicMock() for method in methods} - for method, m in mocks.items(): - monkeypatch.setattr(trace_instance, method, m) - - # WorkflowTraceInfo - info = workflow_trace_info() - trace_instance.trace(info) - mocks["workflow_trace"].assert_called_once_with(info) - - # MessageTraceInfo - info = message_trace_info() - trace_instance.trace(info) - mocks["message_trace"].assert_called_once_with(info) - - # ModerationTraceInfo - info = ModerationTraceInfo(metadata={}, flagged=False, action="allow", preset_response="", query="hello") - trace_instance.trace(info) - mocks["moderation_trace"].assert_called_once_with(info) - - # SuggestedQuestionTraceInfo - info = SuggestedQuestionTraceInfo(metadata={}, total_tokens=0, suggested_question=[], level="info") - trace_instance.trace(info) - mocks["suggested_question_trace"].assert_called_once_with(info) - - # DatasetRetrievalTraceInfo - info = DatasetRetrievalTraceInfo(metadata={}) - trace_instance.trace(info) - mocks["dataset_retrieval_trace"].assert_called_once_with(info) - - # ToolTraceInfo - info = tool_trace_info() - trace_instance.trace(info) - mocks["tool_trace"].assert_called_once_with(info) - - # GenerateNameTraceInfo - info = GenerateNameTraceInfo(tenant_id="tenant-1", metadata={}) - trace_instance.trace(info) - mocks["generate_name_trace"].assert_called_once_with(info) - - @staticmethod - @pytest.mark.parametrize("sqlite3_session", [()], indirect=True) - def test_workflow_trace_with_message_id( - trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session - ) -> None: - # Setup trace info - trace_info = WorkflowTraceInfo( - workflow_id="wf-1", - tenant_id="tenant-1", - workflow_run_id="run-1", - workflow_run_elapsed_time=1.0, - workflow_run_status="succeeded", - workflow_run_inputs={"input": "hi"}, - workflow_run_outputs={"output": "hello"}, - workflow_run_version="1.0", - message_id="msg-1", - conversation_id="conv-1", - total_tokens=100, - file_list=[], - query="hi", - start_time=_dt(), - end_time=_dt() + timedelta(seconds=1), - trace_id="trace-1", - metadata={"app_id": "app-1", "user_id": "user-1"}, - workflow_app_log_id="log-1", - error="", - ) - - monkeypatch.setattr( - "dify_trace_langfuse.langfuse_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - - # Mock node executions - node_llm = MagicMock() - node_llm.id = "node-llm" - node_llm.title = "LLM Node" - node_llm.node_type = BuiltinNodeTypes.LLM - node_llm.status = "succeeded" - node_llm.process_data = { - "model_mode": "chat", - "model_name": "gpt-4", - "model_provider": "openai", - "usage": {"prompt_tokens": 10, "completion_tokens": 20}, - } - node_llm.inputs = {"prompts": "p"} - node_llm.outputs = {"text": "t"} - node_llm.created_at = _dt() - node_llm.elapsed_time = 0.5 - node_llm.metadata = {"foo": "bar"} - - node_other = MagicMock() - node_other.id = "node-other" - node_other.title = "Other Node" - node_other.node_type = BuiltinNodeTypes.CODE - node_other.status = "failed" - node_other.process_data = None - node_other.inputs = {"code": "print"} - node_other.outputs = {"result": "ok"} - node_other.created_at = None # Trigger datetime.now() branch - node_other.elapsed_time = 0.2 - node_other.metadata = None - - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [node_llm, node_other] - - mock_factory = MagicMock() - mock_factory.create_workflow_node_execution_repository.return_value = repo - monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.DifyCoreRepositoryFactory", mock_factory) - - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - # Track calls to add_trace, add_span, add_generation - trace_instance.add_trace = MagicMock() - trace_instance.add_span = MagicMock() - trace_instance.add_generation = MagicMock() - - trace_instance.workflow_trace(trace_info) - - # Verify add_trace (Workflow Level) - trace_instance.add_trace.assert_called_once() - trace_data = trace_instance.add_trace.call_args[1]["langfuse_trace_data"] - assert trace_data.id == "trace-1" - assert trace_data.name == TraceTaskName.MESSAGE_TRACE - assert "message" in trace_data.tags - assert "workflow" in trace_data.tags - - # Verify add_span (Workflow Run Span) - assert trace_instance.add_span.call_count >= 1 - # First span should be workflow run span because message_id is present - workflow_span = trace_instance.add_span.call_args_list[0][1]["langfuse_span_data"] - assert workflow_span.id == "run-1" - assert workflow_span.name == TraceTaskName.WORKFLOW_TRACE - - # Verify Generation for LLM node - trace_instance.add_generation.assert_called_once() - gen_data = trace_instance.add_generation.call_args[1]["langfuse_generation_data"] - assert gen_data.id == "node-llm" - assert gen_data.usage.input == 10 - assert gen_data.usage.output == 20 - - # Verify normal span for Other node - # Second add_span call - other_span = trace_instance.add_span.call_args_list[1][1]["langfuse_span_data"] - assert other_span.id == "node-other" - assert other_span.level == LevelEnum.ERROR - - @staticmethod - @pytest.mark.parametrize("sqlite3_session", [()], indirect=True) - def test_workflow_trace_no_message_id( - trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session - ) -> None: - trace_info = WorkflowTraceInfo( - workflow_id="wf-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.0", - total_tokens=0, - file_list=[], - query="", - message_id=None, - conversation_id="conv-1", - start_time=_dt(), - end_time=_dt(), - trace_id=None, # Should fallback to workflow_run_id - metadata={"app_id": "app-1"}, - workflow_app_log_id="log-1", - error="", - ) - - monkeypatch.setattr( - "dify_trace_langfuse.langfuse_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [] - mock_factory = MagicMock() - mock_factory.create_workflow_node_execution_repository.return_value = repo - monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.DifyCoreRepositoryFactory", mock_factory) - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - trace_instance.add_trace = MagicMock() - trace_instance.workflow_trace(trace_info) - - trace_instance.add_trace.assert_called_once() - trace_data = trace_instance.add_trace.call_args[1]["langfuse_trace_data"] - assert trace_data.id == "run-1" - assert trace_data.name == TraceTaskName.WORKFLOW_TRACE - - @staticmethod - @pytest.mark.parametrize("sqlite3_session", [()], indirect=True) - def test_workflow_trace_missing_app_id( - trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session - ) -> None: - trace_info = WorkflowTraceInfo( - workflow_id="wf-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.0", - total_tokens=0, - file_list=[], - query="", - message_id=None, - conversation_id="conv-1", - start_time=_dt(), - end_time=_dt(), - metadata={}, # Missing app_id - workflow_app_log_id="log-1", - error="", - ) - monkeypatch.setattr( - "dify_trace_langfuse.langfuse_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - - with pytest.raises(ValueError, match="No app_id found in trace_info metadata"): - trace_instance.workflow_trace(trace_info) - - @staticmethod - def test_message_trace_basic(trace_instance, monkeypatch: pytest.MonkeyPatch): - message_data = MagicMock() - message_data.id = "msg-1" - message_data.from_account_id = "acc-1" - message_data.from_end_user_id = None - message_data.provider_response_latency = 0.5 - message_data.conversation_id = "conv-1" - message_data.total_price = 0.01 - message_data.model_id = "gpt-4" - message_data.answer = "hello" - message_data.status = MessageStatus.NORMAL - message_data.error = None - - trace_info = MessageTraceInfo( - message_id="msg-1", - message_data=message_data, - inputs={"query": "hi"}, - outputs={"answer": "hello"}, - message_tokens=10, - answer_tokens=20, - total_tokens=30, - start_time=_dt(), - end_time=_dt() + timedelta(seconds=1), - trace_id="trace-1", - metadata={"foo": "bar"}, - conversation_mode="chat", - conversation_model="gpt-4", - file_list=[], - error=None, - ) - - trace_instance.add_trace = MagicMock() - trace_instance.add_generation = MagicMock() - - trace_instance.message_trace(trace_info) - - trace_instance.add_trace.assert_called_once() - trace_instance.add_generation.assert_called_once() - - gen_data = trace_instance.add_generation.call_args[0][0] - assert gen_data.name == "llm" - assert gen_data.usage.total == 30 - - @staticmethod - @pytest.mark.parametrize("sqlite3_session", [(EndUser,)], indirect=True) - def test_message_trace_with_end_user( - trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session - ) -> None: - message_data = MagicMock() - message_data.id = "msg-1" - message_data.from_account_id = "acc-1" - message_data.from_end_user_id = "end-user-1" - message_data.conversation_id = "conv-1" - message_data.status = MessageStatus.NORMAL - message_data.model_id = "gpt-4" - message_data.error = "" - message_data.answer = "hello" - message_data.total_price = 0.0 - message_data.provider_response_latency = 0.1 - - trace_info = MessageTraceInfo( - message_id="msg-1", - message_data=message_data, - inputs={}, - outputs={}, - message_tokens=0, - answer_tokens=0, - total_tokens=0, - start_time=_dt(), - end_time=_dt(), - metadata={}, - conversation_mode="chat", - conversation_model="gpt-4", - file_list=[], - error=None, - ) - - end_user = EndUser( - id="end-user-1", - tenant_id="tenant-1", - app_id="app-1", - type=EndUserType.BROWSER, - session_id="session-id-123", - ) - engine = sqlite3_session.get_bind() - with Session(engine) as write_session: - write_session.add(end_user) - write_session.commit() - - with Session(engine) as read_session: - monkeypatch.setattr( - "dify_trace_langfuse.langfuse_trace.db", - SimpleNamespace(engine=engine, session=read_session), - ) - - trace_instance.add_trace = MagicMock() - trace_instance.add_generation = MagicMock() - - trace_instance.message_trace(trace_info) - - trace_data = trace_instance.add_trace.call_args[1]["langfuse_trace_data"] - assert trace_data.user_id == "session-id-123" - assert trace_data.metadata["user_id"] == "session-id-123" - - @staticmethod - def test_message_trace_none_data(trace_instance): - trace_info = SimpleNamespace(message_data=None, file_list=[], metadata={}) - trace_instance.add_trace = MagicMock() - trace_instance.message_trace(trace_info) - trace_instance.add_trace.assert_not_called() - - @staticmethod - def test_moderation_trace(trace_instance): - message_data = MagicMock() - message_data.created_at = _dt() - - trace_info = ModerationTraceInfo( - message_id="msg-1", - message_data=message_data, - inputs={"q": "hi"}, - action="stop", - flagged=True, - preset_response="blocked", - start_time=None, - end_time=None, - metadata={"foo": "bar"}, - trace_id="trace-1", - query="hi", - ) - - trace_instance.add_span = MagicMock() - trace_instance.moderation_trace(trace_info) - - trace_instance.add_span.assert_called_once() - span_data = trace_instance.add_span.call_args[1]["langfuse_span_data"] - assert span_data.name == TraceTaskName.MODERATION_TRACE - assert span_data.output["flagged"] is True - - @staticmethod - def test_suggested_question_trace(trace_instance): - message_data = MagicMock() - message_data.status = MessageStatus.NORMAL - message_data.error = None - - trace_info = SuggestedQuestionTraceInfo( - message_id="msg-1", - message_data=message_data, - inputs="hi", - suggested_question=["q1"], - total_tokens=10, - level="info", - start_time=_dt(), - end_time=_dt(), - metadata={}, - trace_id="trace-1", - ) - - trace_instance.add_generation = MagicMock() - trace_instance.suggested_question_trace(trace_info) - - trace_instance.add_generation.assert_called_once() - gen_data = trace_instance.add_generation.call_args[1]["langfuse_generation_data"] - assert gen_data.name == TraceTaskName.SUGGESTED_QUESTION_TRACE - assert gen_data.usage.unit == UnitEnum.CHARACTERS - - @staticmethod - def test_dataset_retrieval_trace(trace_instance): - message_data = MagicMock() - message_data.created_at = _dt() - message_data.updated_at = _dt() - - trace_info = DatasetRetrievalTraceInfo( - message_id="msg-1", - message_data=message_data, - inputs="query", - documents=[{"id": "doc1"}], - start_time=None, - end_time=None, - metadata={}, - trace_id="trace-1", - ) - - trace_instance.add_span = MagicMock() - trace_instance.dataset_retrieval_trace(trace_info) - - trace_instance.add_span.assert_called_once() - span_data = trace_instance.add_span.call_args[1]["langfuse_span_data"] - assert span_data.name == TraceTaskName.DATASET_RETRIEVAL_TRACE - assert span_data.output["documents"] == [{"id": "doc1"}] - - @staticmethod - def test_tool_trace(trace_instance): - trace_info = ToolTraceInfo( - message_id="msg-1", - message_data=MagicMock(), - inputs={}, - outputs={}, - tool_name="my_tool", - tool_inputs={"a": 1}, - tool_outputs="result_string", - time_cost=0.1, - start_time=_dt(), - end_time=_dt(), - metadata={}, - trace_id="trace-1", - tool_config={}, - tool_parameters={}, - error="some error", - ) - - trace_instance.add_span = MagicMock() - trace_instance.tool_trace(trace_info) - - trace_instance.add_span.assert_called_once() - span_data = trace_instance.add_span.call_args[1]["langfuse_span_data"] - assert span_data.name == "my_tool" - assert span_data.level == LevelEnum.ERROR - - @staticmethod - def test_generate_name_trace(trace_instance): - trace_info = GenerateNameTraceInfo( - inputs={"q": "hi"}, - outputs={"name": "new"}, - tenant_id="tenant-1", - conversation_id="conv-1", - start_time=_dt(), - end_time=_dt(), - metadata={"m": 1}, - ) - - trace_instance.add_trace = MagicMock() - trace_instance.add_span = MagicMock() - - trace_instance.generate_name_trace(trace_info) - - trace_instance.add_trace.assert_called_once() - trace_instance.add_span.assert_called_once() - - trace_data = trace_instance.add_trace.call_args[1]["langfuse_trace_data"] - assert trace_data.name == TraceTaskName.GENERATE_NAME_TRACE - assert trace_data.user_id == "tenant-1" - - span_data = trace_instance.add_span.call_args[1]["langfuse_span_data"] - assert span_data.trace_id == "conv-1" - - @staticmethod - def test_add_trace_success(trace_instance): - data = LangfuseTrace(id="t1", name="trace") - trace_instance.add_trace(data) - trace_instance.langfuse_client.ingestion.batch.assert_called_once() - - @staticmethod - def test_add_trace_error(trace_instance): - trace_instance.langfuse_client.ingestion.batch.side_effect = Exception("error") - data = LangfuseTrace(id="t1", name="trace") - with pytest.raises(ValueError, match="LangFuse Failed to create trace: error"): - trace_instance.add_trace(data) - - @staticmethod - def test_add_span_success(trace_instance): - data = LangfuseSpan(id="s1", name="span", trace_id="t1") - trace_instance.add_span(data) - trace_instance.langfuse_client.ingestion.batch.assert_called_once() - - @staticmethod - def test_add_span_error(trace_instance): - trace_instance.langfuse_client.ingestion.batch.side_effect = Exception("error") - data = LangfuseSpan(id="s1", name="span", trace_id="t1") - with pytest.raises(ValueError, match="LangFuse Failed to create span: error"): - trace_instance.add_span(data) - - @staticmethod - def test_update_span(trace_instance): - span = MagicMock() - data = LangfuseSpan(id="s1", name="span", trace_id="t1") - trace_instance.update_span(span, data) - span.end.assert_called_once() - - @staticmethod - def test_add_generation_success(trace_instance): - data = LangfuseGeneration(id="g1", name="gen", trace_id="t1") - trace_instance.add_generation(data) - trace_instance.langfuse_client.ingestion.batch.assert_called_once() - - @staticmethod - def test_add_generation_error(trace_instance): - trace_instance.langfuse_client.ingestion.batch.side_effect = Exception("error") - data = LangfuseGeneration(id="g1", name="gen", trace_id="t1") - with pytest.raises(ValueError, match="LangFuse Failed to create generation: error"): - trace_instance.add_generation(data) - - @staticmethod - def test_update_generation(trace_instance): - gen = MagicMock() - data = LangfuseGeneration(id="g1", name="gen", trace_id="t1") - trace_instance.update_generation(gen, data) - gen.end.assert_called_once() - - @staticmethod - def test_api_check_success(trace_instance): - trace_instance.langfuse_client.projects.get.return_value = MagicMock(data=[MagicMock()]) - assert trace_instance.api_check() is True - - @staticmethod - def test_api_check_rejects_empty_project_list(trace_instance): - trace_instance.langfuse_client.projects.get.return_value = MagicMock(data=[]) - with pytest.raises(ValueError, match="no project found for the provided credentials"): - trace_instance.api_check() - - @staticmethod - def test_api_check_error(trace_instance): - trace_instance.langfuse_client.projects.get.side_effect = Exception("fail") - with pytest.raises(ValueError, match="LangFuse API check failed: fail"): - trace_instance.api_check() - - @staticmethod - def test_get_project_key_success(trace_instance): - mock_data = MagicMock() - mock_data.id = "proj-1" - trace_instance.langfuse_client.projects.get.return_value = MagicMock(data=[mock_data]) - assert trace_instance.get_project_key() == "proj-1" - - @staticmethod - def test_get_project_key_error(trace_instance): - trace_instance.langfuse_client.projects.get.side_effect = Exception("fail") - with pytest.raises(ValueError, match="LangFuse get project key failed: fail"): - trace_instance.get_project_key() - - @staticmethod - def test_moderation_trace_none(trace_instance): - trace_info = ModerationTraceInfo( - message_id="m", - message_data=None, - inputs={}, - action="s", - flagged=False, - preset_response="", - query="", - metadata={}, - ) - trace_instance.add_span = MagicMock() - trace_instance.moderation_trace(trace_info) - trace_instance.add_span.assert_not_called() - - @staticmethod - def test_suggested_question_trace_none(trace_instance): - trace_info = SuggestedQuestionTraceInfo( - message_id="m", message_data=None, inputs={}, suggested_question=[], total_tokens=0, level="i", metadata={} - ) - trace_instance.add_generation = MagicMock() - trace_instance.suggested_question_trace(trace_info) - trace_instance.add_generation.assert_not_called() - - @staticmethod - def test_dataset_retrieval_trace_none(trace_instance): - trace_info = DatasetRetrievalTraceInfo(message_id="m", message_data=None, inputs={}, documents=[], metadata={}) - trace_instance.add_span = MagicMock() - trace_instance.dataset_retrieval_trace(trace_info) - trace_instance.add_span.assert_not_called() - - @staticmethod - def test_langfuse_trace_entity_with_list_dict_input(): - # To cover lines 29-31 in langfuse_trace_entity.py - # We need to mock replace_text_with_content or just check if it works - # Actually replace_text_with_content is imported from core.ops.utils - data = LangfuseTrace(id="t1", name="n", input=[{"text": "hello"}]) - assert isinstance(data.input, list) - assert data.input[0]["content"] == "hello" - - @staticmethod - @pytest.mark.parametrize("sqlite3_session", [()], indirect=True) - def test_workflow_trace_handles_usage_extraction_error( - trace_instance, - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - sqlite3_session: Session, - ): - # Setup trace info to trigger LLM node usage extraction - trace_info = WorkflowTraceInfo( - workflow_id="wf-1", - tenant_id="t", - workflow_run_id="r", - workflow_run_elapsed_time=1.0, - workflow_run_status="s", - workflow_run_inputs={}, - workflow_run_outputs={}, - workflow_run_version="1", - total_tokens=0, - file_list=[], - query="", - message_id=None, - conversation_id="c", - start_time=_dt(), - end_time=_dt(), - metadata={"app_id": "app-1"}, - workflow_app_log_id="l", - error="", - ) - - node = MagicMock() - node.id = "n1" - node.title = "LLM Node" - node.node_type = BuiltinNodeTypes.LLM - node.status = "succeeded" - - class BadDict(collections.UserDict): - @override - def get(self, key, default=None): - if key == "usage": - raise Exception("Usage extraction failed") - return super().get(key, default) - - node.process_data = BadDict({"model_mode": "chat", "model_name": "gpt-4", "usage": True, "prompts": ["p"]}) - node.created_at = _dt() - node.elapsed_time = 0.1 - node.metadata = {} - node.outputs = {} - - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [node] - mock_factory = MagicMock() - mock_factory.create_workflow_node_execution_repository.return_value = repo - monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.DifyCoreRepositoryFactory", mock_factory) - monkeypatch.setattr( - "dify_trace_langfuse.langfuse_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - trace_instance.add_trace = MagicMock() - trace_instance.add_generation = MagicMock() - - with caplog.at_level(logging.ERROR): - trace_instance.workflow_trace(trace_info) - - assert "Failed to extract usage" in caplog.text - trace_instance.add_generation.assert_called_once() - - -"""Tests for Langfuse TTFT reporting support.""" - - -def _create_trace_instance() -> LangFuseDataTrace: - with patch("dify_trace_langfuse.langfuse_trace.LangfuseAPI", autospec=True): - return LangFuseDataTrace( - LangfuseConfig( - public_key="public-key", - secret_key="secret-key", - host="https://cloud.langfuse.com", - ) - ) - - -class TestLangfuseTraceCasesUnitTests: - class TestLangFuseDataTraceCompletionStartTime: - def test_message_trace_reports_completion_start_time(self): - trace = _create_trace_instance() - start_time = datetime(2026, 3, 11, 13, 0, 0) - trace_info = MessageTraceInfo( - trace_id="trace-123", - message_id="message-123", - message_data=SimpleNamespace( - id="message-123", - from_account_id="account-1", - from_end_user_id=None, - conversation_id="conversation-1", - model_id="gpt-4o-mini", - answer="hi there", - status="normal", - error="", - total_price=0.12, - provider_response_latency=3.5, - ), - conversation_model="chat", - message_tokens=10, - answer_tokens=20, - total_tokens=30, - error="", - inputs="hello", - outputs="hi there", - file_list=[], - start_time=start_time, - end_time=start_time + timedelta(seconds=3.5), - metadata={}, - message_file_data=None, - conversation_mode="chat", - gen_ai_server_time_to_first_token=1.2, - llm_streaming_time_to_generate=2.3, - is_streaming_request=True, - ) - - with patch.object(trace, "add_trace"), patch.object(trace, "add_generation") as add_generation: - trace.message_trace(trace_info) - - generation = add_generation.call_args.args[0] - assert generation.completion_start_time == start_time + timedelta(seconds=1.2) - - def test_workflow_trace_reports_completion_start_time_from_llm_usage(self): - trace = _create_trace_instance() - start_time = datetime(2026, 3, 11, 13, 0, 0) - node_execution = SimpleNamespace( - id="node-exec-1", - title="Chat LLM", - node_type=BuiltinNodeTypes.LLM, - status="succeeded", - process_data={ - "model_mode": "chat", - "model_name": "gpt-4o-mini", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 20, - "time_to_first_token": 1.2, - }, - }, - inputs={"question": "hello"}, - outputs={"text": "hi there"}, - created_at=start_time, - elapsed_time=3.5, - metadata={}, - ) - trace_info = WorkflowTraceInfo( - trace_id="trace-123", - workflow_data={}, - conversation_id=None, - workflow_app_log_id=None, - workflow_id="workflow-1", - tenant_id="tenant-1", - workflow_run_id="workflow-run-1", - workflow_run_elapsed_time=3.5, - workflow_run_status="succeeded", - workflow_run_inputs={"question": "hello"}, - workflow_run_outputs={"answer": "hi there"}, - workflow_run_version="1", - error="", - total_tokens=30, - file_list=[], - query="hello", - metadata={"app_id": "app-1", "user_id": "user-1"}, - start_time=start_time, - end_time=start_time + timedelta(seconds=3.5), - ) - repository = MagicMock() - repository.get_by_workflow_execution.return_value = [node_execution] - - with ( - patch.object(trace, "add_trace"), - patch.object(trace, "add_span"), - patch.object(trace, "add_generation") as add_generation, - patch.object(trace, "get_service_account_with_tenant", return_value=MagicMock()), - patch("dify_trace_langfuse.langfuse_trace.db", MagicMock()), - patch( - "dify_trace_langfuse.langfuse_trace.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - return_value=repository, - ), - ): - trace.workflow_trace(trace_info) - - generation = add_generation.call_args.kwargs["langfuse_generation_data"] - assert generation.completion_start_time == start_time + timedelta(seconds=1.2) - - def test_ignores_invalid_ttft_values(self): - trace = _create_trace_instance() - start_time = datetime(2026, 3, 11, 13, 0, 0) - - assert trace._get_completion_start_time(start_time, None) is None - assert trace._get_completion_start_time(start_time, -1) is None - assert trace._get_completion_start_time(start_time, cast(float | int | None, "invalid")) is None diff --git a/api/providers/trace/trace-langfuse/tests/unit_tests/src/dify_trace_langfuse/test_provider.py b/api/providers/trace/trace-langfuse/tests/unit_tests/src/dify_trace_langfuse/test_provider.py new file mode 100644 index 00000000000000..8227f931cca99e --- /dev/null +++ b/api/providers/trace/trace-langfuse/tests/unit_tests/src/dify_trace_langfuse/test_provider.py @@ -0,0 +1,236 @@ +import base64 + +import pytest +from dify_trace_langfuse.config import LangfuseConfig +from dify_trace_langfuse.provider import MAX_BATCH_BYTES, LangfuseProvider +from dify_trace_sdk.egress import EgressResponse +from dify_trace_sdk.errors import ( + EgressError, + ExportRejectedError, + ExportRetryableError, + VerificationError, +) +from dify_trace_sdk.ids import span_uuid +from dify_trace_sdk.testing import MESSAGE_ID, RecordingEgress, message_trace, response, tool_fragment, workflow_trace +from pydantic import JsonValue + +CONFIG = LangfuseConfig(public_key="pk-lf-1", secret_key="sk-lf-1", host="https://cloud.langfuse.com") + + +def _provider(*responses: EgressResponse | Exception) -> tuple[LangfuseProvider, RecordingEgress]: + egress = RecordingEgress(responses, grpc_target=None) + return LangfuseProvider(egress), egress + + +def _sent_events(egress: RecordingEgress) -> list[dict]: + (sent,) = egress.requests + body = sent.json() + assert isinstance(body, dict) + return body["batch"] + + +def test_verify_resolves_the_project() -> None: + provider, egress = _provider(response(200, {"data": [{"id": "project-1", "name": "Dify"}]})) + verified = provider.verify(CONFIG) + assert verified.project_key == "project-1" + assert provider.project_url(verified) == "https://cloud.langfuse.com/project/project-1" + (sent,) = egress.requests + assert sent.method == "GET" + assert sent.url == "https://cloud.langfuse.com/api/public/projects" + assert sent.headers["Authorization"] == "Basic " + base64.b64encode(b"pk-lf-1:sk-lf-1").decode() + + +def test_legacy_config_keeps_project_and_trace_links_without_verification() -> None: + provider, egress = _provider(response(200, {})) + config = LangfuseConfig.model_validate( + { + "public_key": "pk-lf-1", + "secret_key": "sk-lf-1", + "host": "https://collector.example/langfuse", + "project_url": "https://dashboard.example/langfuse/project/legacy-project/", + } + ) + + result = provider.send(config, message_trace()) + + assert provider.project_url(config) == config.project_url + trace_url = f"https://dashboard.example/langfuse/project/legacy-project/traces/{MESSAGE_ID}" + assert result.trace_url == trace_url + assert result.span_urls + assert all(url.startswith(trace_url + "?observation=") for url in result.span_urls.values()) + assert len(egress.requests) == 1 + assert egress.requests[0].url == "https://collector.example/langfuse/api/public/ingestion" + verified = config.model_copy(update={"project_key": "new-project"}) + assert provider.project_url(verified) == "https://collector.example/langfuse/project/new-project" + + +@pytest.mark.parametrize( + ("outcome", "code"), + [ + (response(401, {}), "authentication_failed"), + (response(500, {}), "endpoint_unavailable"), + (response(200, {"data": []}), "project_not_found"), + (response(200, b"not json"), "malformed_response"), + (EgressError("refused", request_sent=False), "connection_failed"), + ], +) +def test_verify_failures_have_stable_codes(outcome: EgressResponse | Exception, code: str) -> None: + provider, _ = _provider(outcome) + with pytest.raises(VerificationError) as caught: + provider.verify(CONFIG) + assert caught.value.code == code + + +def test_message_trace_becomes_trace_and_generation() -> None: + provider, egress = _provider(response(207, {"successes": [], "errors": []})) + trace = message_trace() + result = provider.send(CONFIG.model_copy(update={"project_key": "project-1"}), trace) + + trace_event, generation = _sent_events(egress) + assert trace_event["type"] == "trace-create" + assert trace_event["body"]["id"] == MESSAGE_ID + assert trace_event["body"]["sessionId"] == "conversation-1" + assert trace_event["body"]["userId"] == "end-user-1" + assert generation["type"] == "generation-create" + body = generation["body"] + assert body["traceId"] == MESSAGE_ID + assert "parentObservationId" not in body + assert body["model"] == "gpt-4o" + assert body["usage"] == {"input": 12, "output": 8, "total": 20, "unit": "TOKENS", "totalCost": 0.0004} + assert body["completionStartTime"] == "2026-10-01T08:00:00.250000+00:00" + assert body["input"] == [{"role": "user", "content": "What is Dify?"}] + assert body["output"] == [{"role": "assistant", "content": "An LLM application platform."}] + assert result.trace_id == MESSAGE_ID + assert result.trace_url == f"https://cloud.langfuse.com/project/project-1/traces/{MESSAGE_ID}" + + +@pytest.mark.parametrize( + ("inputs", "expected"), + [ + ("Hello", [{"role": "user", "content": "Hello"}]), + ( + [{"role": "system", "text": "Be brief", "files": []}, {"role": "user", "text": "Hello"}], + [{"role": "system", "content": "Be brief", "files": []}, {"role": "user", "content": "Hello"}], + ), + ( + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this"}, + {"type": "image_url", "image_url": {"url": "https://example.com/image.png"}}, + ], + } + ], + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this"}, + {"type": "image_url", "image_url": {"url": "https://example.com/image.png"}}, + ], + } + ], + ), + ({"text": "workflow data"}, {"text": "workflow data"}), + ([{"text": "retrieval data"}], [{"text": "retrieval data"}]), + ], +) +def test_chat_inputs_keep_the_langfuse_preview_format(inputs: JsonValue, expected: JsonValue) -> None: + provider, egress = _provider(response(200, {})) + trace = message_trace() + trace = trace.model_copy(update={"spans": (trace.root, trace.spans[1].model_copy(update={"inputs": inputs}))}) + source = trace.model_dump() + + provider.send(CONFIG, trace) + + _, generation = _sent_events(egress) + assert generation["body"]["input"] == expected + assert trace.model_dump() == source + + +def test_workflow_spans_keep_hierarchy_and_errors() -> None: + provider, egress = _provider(response(200, {})) + trace = workflow_trace() + provider.send(CONFIG, trace) + + _, llm, tool = _sent_events(egress) + assert "parentObservationId" not in llm["body"] + assert tool["body"]["parentObservationId"] == str(span_uuid(trace.trace_id, trace.spans[1].span_id)) + assert tool["type"] == "span-create" + assert tool["body"]["level"] == "ERROR" + assert tool["body"]["statusMessage"] == "Tool failed" + + +def test_fragment_joins_the_message_trace_without_recreating_it() -> None: + provider, egress = _provider(response(207, {"successes": [{"id": "x", "status": 201}], "errors": []})) + provider.send(CONFIG, tool_fragment()) + + (tool,) = _sent_events(egress) + assert tool["type"] == "span-create" + assert tool["body"]["traceId"] == MESSAGE_ID + assert "parentObservationId" not in tool["body"] + + +def test_external_trace_id_is_kept_verbatim() -> None: + provider, egress = _provider(response(200, {})) + trace = message_trace().model_copy(update={"external_trace_id": "caller-trace-1"}) + assert provider.send(CONFIG, trace).trace_id == "caller-trace-1" + trace_event, generation = _sent_events(egress) + assert trace_event["body"]["id"] == "caller-trace-1" + assert generation["body"]["traceId"] == "caller-trace-1" + + +def test_event_ids_are_stable_across_redelivery() -> None: + first, first_egress = _provider(response(200, {})) + second, second_egress = _provider(response(200, {})) + first.send(CONFIG, workflow_trace()) + second.send(CONFIG, workflow_trace()) + assert [event["id"] for event in _sent_events(first_egress)] == [ + event["id"] for event in _sent_events(second_egress) + ] + + +@pytest.mark.parametrize( + ("outcome", "expected"), + [ + ( + response(207, {"successes": [], "errors": [{"id": "e", "status": 400, "message": "bad"}]}), + ExportRejectedError, + ), + ( + response(207, {"successes": [], "errors": [{"id": "e", "status": 500, "message": "busy"}]}), + ExportRetryableError, + ), + (response(401, {}), ExportRejectedError), + (EgressError("read timeout", request_sent=True), ExportRetryableError), + ], +) +def test_send_failures_are_classified(outcome: EgressResponse | Exception, expected: type[Exception]) -> None: + provider, _ = _provider(outcome) + with pytest.raises(expected): + provider.send(CONFIG, message_trace()) + + +def test_unreadable_ingestion_bodies_are_resent() -> None: + provider, _ = _provider(EgressResponse(status_code=200, headers={}, content=b"proxy")) + with pytest.raises(ExportRetryableError): + provider.send(CONFIG, message_trace()) + + +def test_batches_are_sized_in_bytes() -> None: + wide = message_trace().model_copy( + update={ + "spans": tuple( + span.model_copy(update={"inputs": "中" * (MAX_BATCH_BYTES // 4)}) for span in message_trace().spans + ) + } + ) + provider, egress = _provider( + response(207, {"successes": [], "errors": []}), response(207, {"successes": [], "errors": []}) + ) + + provider.send(CONFIG, wide) + + assert len(egress.requests) >= 2 + assert all(len(request.content or b"") <= MAX_BATCH_BYTES * 1.1 for request in egress.requests) diff --git a/api/providers/trace/trace-langsmith/pyproject.toml b/api/providers/trace/trace-langsmith/pyproject.toml index 618bec79d0ee56..a961929309b157 100644 --- a/api/providers/trace/trace-langsmith/pyproject.toml +++ b/api/providers/trace/trace-langsmith/pyproject.toml @@ -2,9 +2,13 @@ name = "dify-trace-langsmith" version = "0.0.1" dependencies = [ + "dify-trace-sdk", "langsmith==0.8.18", ] description = "Dify ops tracing provider (LangSmith)." +[project.entry-points."dify.trace_providers"] +langsmith = "dify_trace_langsmith.provider:LangSmithProvider" + [tool.setuptools.packages.find] where = ["src"] diff --git a/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/agent_trace.py b/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/agent_trace.py deleted file mode 100644 index 0f0be6c5e64318..00000000000000 --- a/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/agent_trace.py +++ /dev/null @@ -1,105 +0,0 @@ -"""Canonical Agent trees mapped to synchronous, isolated LangSmith runs.""" - -from collections.abc import Mapping -from typing import Any -from urllib.parse import quote -from uuid import NAMESPACE_URL, uuid5 - -from core.ops.exceptions import RetryableTraceDispatchError -from core.ops.unified_trace.entities import CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace -from core.ops.utils import generate_dotted_order -from services.tracing.provider_common import endpoint, finished_at, required, token_usage, usd_cost, utc -from services.tracing.provider_dispatch import TraceExportRejectedError, TraceExportResult, TraceExportUncertainError -from services.tracing.sdk_transport import SSRFSession - - -def export_trace(provider: str, config: Mapping[str, Any], trace: CanonicalTrace) -> TraceExportResult: - from langsmith import Client - from urllib3.util.retry import Retry - - api_url = endpoint(config.get("endpoint", "https://api.smith.langchain.com")) - project = required(config, "project") - ids = {span.id: str(uuid5(NAMESPACE_URL, f"dify:{trace.trace_id}:{span.id}")) for span in trace.spans} - root_id = ids[trace.root_span_id] - orders: dict[str, str] = {} - accepted: set[str] = set() - client = Client( - api_key=required(config, "api_key"), - api_url=api_url, - session=SSRFSession(provider), - retry_config=Retry(total=0), - auto_batch_tracing=False, - tracing_sampling_rate=1.0, - tracing_mode="langsmith", - omit_traced_runtime_info=True, - hide_inputs=False, - hide_outputs=False, - info={}, - ) - try: - for span in trace.spans: - orders[span.id] = generate_dotted_order( - ids[span.id], utc(span.start_time), orders.get(span.parent_id or "") - ) - metadata = { - **span.metadata, - "session_id": trace.session_id, - "external_trace_id": trace.trace_id, - "dify.span.id": span.id, - "dify.span.kind": span.kind.value, - } - model = span.metadata.get("model") or span.metadata.get("model_name") - if span.kind is CanonicalSpanKind.LLM and model: - metadata["ls_model_name"] = model - metadata["ls_provider"] = span.metadata.get("model_provider") - values: dict[str, Any] = {} - usage = token_usage(span) - values.update( - {f"{key}_tokens" if key != "total" else "total_tokens": value for key, value in usage.items()} - ) - cost = usd_cost(span) - if cost is not None: - values["total_cost"] = cost - outputs = dict(span.outputs) if isinstance(span.outputs, Mapping) else {"output": span.outputs} - if usage: - outputs["usage_metadata"] = {f"{key}_tokens": value for key, value in usage.items()} - client.create_run( - id=ids[span.id], - name=span.name, - inputs=dict(span.inputs) if isinstance(span.inputs, Mapping) else {"input": span.inputs}, - outputs=outputs, - run_type="chain" if span.kind is CanonicalSpanKind.AGENT else span.kind.value, - project_name=project, - start_time=utc(span.start_time), - end_time=finished_at(span), - parent_run_id=ids.get(span.parent_id or ""), - trace_id=root_id, - dotted_order=orders[span.id], - extra={"metadata": metadata}, - error=(span.error or "Trace operation failed") if span.status is CanonicalSpanStatus.ERROR else None, - **values, - ) - accepted.add(span.id) - except Exception as error: - if accepted: - raise TraceExportUncertainError( - "LangSmith accepted part of the trace; reconciliation is required" - ) from None - # LangSmith wraps injected transport exceptions; recover our sanitized outcome. - cause: BaseException | None = error - while cause is not None: - if isinstance(cause, (RetryableTraceDispatchError, TraceExportRejectedError)): - raise cause from None - cause = cause.__cause__ - raise TraceExportRejectedError("LangSmith rejected the trace") from None - finally: - client.close() - project_url = config.get("project_url") - trace_url = f"{str(project_url).rstrip('/')}/r/{quote(root_id)}" if project_url else None - return TraceExportResult( - trace_id=root_id, - trace_url=trace_url, - span_urls={span.id: f"{str(project_url).rstrip('/')}/r/{ids[span.id]}" for span in trace.spans} - if project_url - else {}, - ) diff --git a/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/config.py b/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/config.py index a4d75a3c2b5b51..81f22d35e11f10 100644 --- a/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/config.py +++ b/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/config.py @@ -1,20 +1,34 @@ -from pydantic import ValidationInfo, field_validator +from dify_trace_sdk.config import ProviderConfig, derived, destination, http_url, secret +from pydantic import field_validator -from core.ops.entities.config_entity import BaseTracingConfig -from core.ops.utils import validate_url_with_path +DEFAULT_ENDPOINT = "https://api.smith.langchain.com" -class LangSmithConfig(BaseTracingConfig): - """ - Model class for Langsmith tracing config. - """ - - api_key: str +class LangSmithConfig(ProviderConfig): + api_key: str = secret() project: str - endpoint: str = "https://api.smith.langchain.com" + endpoint: str = destination(default=DEFAULT_ENDPOINT) + project_url: str | None = derived() + + @field_validator("api_key") + @classmethod + def _api_key(cls, value: str) -> str: + # The SDK strips the key before exporting; verification must send the same value. + key = value.strip() + if not key: + raise ValueError("API key is required") + return key + + @field_validator("project") + @classmethod + def _project(cls, value: str) -> str: + # The SDK replaces an empty project name with the host's LANGSMITH_PROJECT. + if not value.strip(): + raise ValueError("Project name is required") + return value @field_validator("endpoint") @classmethod - def endpoint_validator(cls, v, info: ValidationInfo): - # LangSmith only allows HTTPS; self-hosted deployments may sit under a path prefix - return validate_url_with_path(v, "https://api.smith.langchain.com", allowed_schemes=("https",)) + def _endpoint(cls, value: str) -> str: + # LangSmith is HTTPS-only; self-hosted deployments may serve the API under a path prefix. + return http_url(value, allowed_schemes=("https",)) if value.strip() else DEFAULT_ENDPOINT diff --git a/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/entities/__init__.py b/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/entities/__init__.py deleted file mode 100644 index e69de29bb2d1d6..00000000000000 diff --git a/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/entities/langsmith_trace_entity.py b/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/entities/langsmith_trace_entity.py deleted file mode 100644 index cdbedaf739955c..00000000000000 --- a/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/entities/langsmith_trace_entity.py +++ /dev/null @@ -1,149 +0,0 @@ -from collections.abc import Mapping -from datetime import datetime -from enum import StrEnum -from typing import Any - -from pydantic import BaseModel, Field, field_validator -from pydantic_core.core_schema import ValidationInfo - -from core.ops.utils import replace_text_with_content - - -class LangSmithRunType(StrEnum): - tool = "tool" - chain = "chain" - llm = "llm" - retriever = "retriever" - embedding = "embedding" - prompt = "prompt" - parser = "parser" - - -class LangSmithTokenUsage(BaseModel): - input_tokens: int | None = None - output_tokens: int | None = None - total_tokens: int | None = None - - -class LangSmithMultiModel(BaseModel): - file_list: list[str] | None = Field(None, description="List of files") - - -class LangSmithRunModel(LangSmithTokenUsage, LangSmithMultiModel): - name: str | None = Field(..., description="Name of the run") - inputs: str | Mapping[str, Any] | list | None | None = Field(None, description="Inputs of the run") - outputs: str | Mapping[str, Any] | list | None | None = Field(None, description="Outputs of the run") - run_type: LangSmithRunType = Field(..., description="Type of the run") - start_time: datetime | str | None = Field(None, description="Start time of the run") - end_time: datetime | str | None = Field(None, description="End time of the run") - extra: dict[str, Any] | None = Field(None, description="Extra information of the run") - error: str | None = Field(None, description="Error message of the run") - serialized: dict[str, Any] | None = Field(None, description="Serialized data of the run") - parent_run_id: str | None = Field(None, description="Parent run ID") - events: list[dict[str, Any]] | None = Field(None, description="Events associated with the run") - tags: list[str] | None = Field(None, description="Tags associated with the run") - trace_id: str | None = Field(None, description="Trace ID associated with the run") - dotted_order: str | None = Field(None, description="Dotted order of the run") - id: str | None = Field(None, description="ID of the run") - session_id: str | None = Field(None, description="Session ID associated with the run") - session_name: str | None = Field(None, description="Session name associated with the run") - reference_example_id: str | None = Field(None, description="Reference example ID associated with the run") - input_attachments: dict[str, Any] | None = Field(None, description="Input attachments of the run") - output_attachments: dict[str, Any] | None = Field(None, description="Output attachments of the run") - - @field_validator("inputs", "outputs") - @classmethod - def ensure_dict(cls, v, info: ValidationInfo): - field_name = info.field_name - values = info.data - if v == {} or v is None: - return v - usage_metadata = { - "input_tokens": values.get("input_tokens", 0), - "output_tokens": values.get("output_tokens", 0), - "total_tokens": values.get("total_tokens", 0), - } - file_list = values.get("file_list", []) - match v: - case str(): - match field_name: - case "inputs": - return { - "messages": { - "role": "user", - "content": v, - "usage_metadata": usage_metadata, - "file_list": file_list, - }, - } - case "outputs": - return { - "choices": { - "role": "ai", - "content": v, - "usage_metadata": usage_metadata, - "file_list": file_list, - }, - } - case _: - pass - case list(): - data = {} - if len(v) > 0 and isinstance(v[0], dict): - # rename text to content - v = replace_text_with_content(data=v) - match field_name: - case "inputs": - data = { - "messages": v, - } - case "outputs": - data = { - "choices": { - "role": "ai", - "content": v, - "usage_metadata": usage_metadata, - "file_list": file_list, - }, - } - case _: - pass - return data - else: - return { - "choices": { - "role": "ai" if field_name == "outputs" else "user", - "content": str(v), - "usage_metadata": usage_metadata, - "file_list": file_list, - }, - } - if isinstance(v, dict): - v["usage_metadata"] = usage_metadata - v["file_list"] = file_list - return v - return v - - @classmethod - @field_validator("start_time", "end_time") - def format_time(cls, v, info: ValidationInfo): - if not isinstance(v, datetime): - raise ValueError(f"{info.field_name} must be a datetime object") - else: - return v.strftime("%Y-%m-%dT%H:%M:%S.%fZ") - - -class LangSmithRunUpdateModel(BaseModel): - run_id: str = Field(..., description="ID of the run") - trace_id: str | None = Field(None, description="Trace ID associated with the run") - dotted_order: str | None = Field(None, description="Dotted order of the run") - parent_run_id: str | None = Field(None, description="Parent run ID") - end_time: datetime | str | None = Field(None, description="End time of the run") - error: str | None = Field(None, description="Error message of the run") - inputs: dict[str, Any] | None = Field(None, description="Inputs of the run") - outputs: dict[str, Any] | None = Field(None, description="Outputs of the run") - events: list[dict[str, Any]] | None = Field(None, description="Events associated with the run") - tags: list[str] | None = Field(None, description="Tags associated with the run") - extra: dict[str, Any] | None = Field(None, description="Extra information of the run") - input_attachments: dict[str, Any] | None = Field(None, description="Input attachments of the run") - output_attachments: dict[str, Any] | None = Field(None, description="Output attachments of the run") diff --git a/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/langsmith_trace.py b/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/langsmith_trace.py deleted file mode 100644 index bdd576b1d294ad..00000000000000 --- a/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/langsmith_trace.py +++ /dev/null @@ -1,532 +0,0 @@ -import logging -import os -import uuid -from datetime import datetime, timedelta -from typing import cast, override - -from langsmith import Client -from langsmith.schemas import RunBase -from sqlalchemy.orm import sessionmaker - -from core.ops.base_trace_instance import BaseTraceInstance -from core.ops.entities.trace_entity import ( - BaseTraceInfo, - DatasetRetrievalTraceInfo, - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - SuggestedQuestionTraceInfo, - ToolTraceInfo, - TraceTaskName, - WorkflowTraceInfo, -) -from core.ops.utils import filter_none_values, generate_dotted_order -from core.repositories import DifyCoreRepositoryFactory -from dify_trace_langsmith.config import LangSmithConfig -from dify_trace_langsmith.entities.langsmith_trace_entity import ( - LangSmithRunModel, - LangSmithRunType, - LangSmithRunUpdateModel, -) -from extensions.ext_database import db -from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey -from models import EndUser, MessageFile, WorkflowNodeExecutionTriggeredFrom - -logger = logging.getLogger(__name__) - - -class LangSmithDataTrace(BaseTraceInstance): - def __init__( - self, - langsmith_config: LangSmithConfig, - ): - super().__init__(langsmith_config) - self.langsmith_key = langsmith_config.api_key - self.project_name = langsmith_config.project - self.project_id = None - self.langsmith_client = Client(api_key=langsmith_config.api_key, api_url=langsmith_config.endpoint) - self.file_base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001") - - @override - def trace(self, trace_info: BaseTraceInfo): - match trace_info: - case WorkflowTraceInfo(): - self.workflow_trace(trace_info) - case MessageTraceInfo(): - self.message_trace(trace_info) - case ModerationTraceInfo(): - self.moderation_trace(trace_info) - case SuggestedQuestionTraceInfo(): - self.suggested_question_trace(trace_info) - case DatasetRetrievalTraceInfo(): - self.dataset_retrieval_trace(trace_info) - case ToolTraceInfo(): - self.tool_trace(trace_info) - case GenerateNameTraceInfo(): - self.generate_name_trace(trace_info) - case _: - pass - - def workflow_trace(self, trace_info: WorkflowTraceInfo): - # trace_id must equal the root run's run_id (LangSmith protocol); external trace_id - # cannot be used here as it would cause HTTP 400. - trace_id = trace_info.message_id or trace_info.workflow_run_id - if trace_info.start_time is None: - trace_info.start_time = datetime.now() - message_dotted_order = ( - generate_dotted_order(trace_info.message_id, trace_info.start_time) if trace_info.message_id else None - ) - workflow_dotted_order = generate_dotted_order( - trace_info.workflow_run_id, - trace_info.workflow_data.created_at, - message_dotted_order, - ) - metadata = trace_info.metadata - metadata["workflow_app_log_id"] = trace_info.workflow_app_log_id - if trace_info.trace_id: - metadata["external_trace_id"] = trace_info.trace_id - - if trace_info.message_id: - message_run = LangSmithRunModel( - id=trace_info.message_id, - name=TraceTaskName.MESSAGE_TRACE, - inputs=dict(trace_info.workflow_run_inputs), - outputs=dict(trace_info.workflow_run_outputs), - run_type=LangSmithRunType.chain, - start_time=trace_info.start_time, - end_time=trace_info.end_time, - extra={ - "metadata": metadata, - }, - tags=["message", "workflow"], - error=trace_info.error, - trace_id=trace_id, - dotted_order=message_dotted_order, - file_list=[], - serialized=None, - parent_run_id=None, - events=[], - session_id=None, - session_name=None, - reference_example_id=None, - input_attachments={}, - output_attachments={}, - ) - self.add_run(message_run) - - langsmith_run = LangSmithRunModel( - file_list=trace_info.file_list, - total_tokens=trace_info.total_tokens, - id=trace_info.workflow_run_id, - name=TraceTaskName.WORKFLOW_TRACE, - inputs=dict(trace_info.workflow_run_inputs), - run_type=LangSmithRunType.tool, - start_time=trace_info.workflow_data.created_at, - end_time=trace_info.workflow_data.finished_at, - outputs=dict(trace_info.workflow_run_outputs), - extra={ - "metadata": metadata, - }, - error=trace_info.error, - tags=["workflow"], - parent_run_id=trace_info.message_id or None, - trace_id=trace_id, - dotted_order=workflow_dotted_order, - serialized=None, - events=[], - session_id=None, - session_name=None, - reference_example_id=None, - input_attachments={}, - output_attachments={}, - ) - - self.add_run(langsmith_run) - - # through workflow_run_id get all_nodes_execution using repository - session_factory = sessionmaker(bind=db.engine) - # Find the app's creator account - app_id = trace_info.metadata.get("app_id") - if not app_id: - raise ValueError("No app_id found in trace_info metadata") - - service_account = self.get_service_account_with_tenant(app_id) - - workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( - session_factory=session_factory, - tenant_id=trace_info.tenant_id, - user=service_account, - app_id=app_id, - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - - # Get all executions for this workflow run - workflow_node_executions = workflow_node_execution_repository.get_by_workflow_execution( - workflow_execution_id=trace_info.workflow_run_id - ) - - for node_execution in workflow_node_executions: - node_execution_id = node_execution.id - tenant_id = trace_info.tenant_id # Use from trace_info instead - app_id = trace_info.metadata.get("app_id") # Use from trace_info instead - node_name = node_execution.title - node_type = node_execution.node_type - status = node_execution.status - if node_type == BuiltinNodeTypes.LLM: - inputs = node_execution.process_data.get("prompts", {}) if node_execution.process_data else {} - else: - inputs = node_execution.inputs or {} - outputs = node_execution.outputs or {} - created_at = node_execution.created_at or datetime.now() - elapsed_time = node_execution.elapsed_time - finished_at = created_at + timedelta(seconds=elapsed_time) - - execution_metadata = node_execution.metadata or {} - node_total_tokens = execution_metadata.get(WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS) or 0 - metadata = {str(key): value for key, value in execution_metadata.items()} - metadata.update( - { - "workflow_run_id": trace_info.workflow_run_id, - "node_execution_id": node_execution_id, - "tenant_id": tenant_id, - "app_id": app_id, - "app_name": node_name, - "node_type": node_type, - "status": status, - } - ) - - process_data = node_execution.process_data or {} - - if process_data and process_data.get("model_mode") == "chat": - run_type = LangSmithRunType.llm - metadata.update( - { - "ls_provider": process_data.get("model_provider", ""), - "ls_model_name": process_data.get("model_name", ""), - } - ) - elif node_type == BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL: - run_type = LangSmithRunType.retriever - else: - run_type = LangSmithRunType.tool - - prompt_tokens = 0 - completion_tokens = 0 - try: - usage_data = process_data.get("usage", {}) if "usage" in process_data else outputs.get("usage", {}) - prompt_tokens = usage_data.get("prompt_tokens", 0) - completion_tokens = usage_data.get("completion_tokens", 0) - except Exception: - logger.error("Failed to extract usage", exc_info=True) - - node_dotted_order = generate_dotted_order(node_execution_id, created_at, workflow_dotted_order) - langsmith_run = LangSmithRunModel( - total_tokens=node_total_tokens, - input_tokens=prompt_tokens, - output_tokens=completion_tokens, - name=node_type, - inputs=inputs, - run_type=run_type, - start_time=created_at, - end_time=finished_at, - outputs=outputs, - file_list=trace_info.file_list, - extra={ - "metadata": metadata, - }, - parent_run_id=trace_info.workflow_run_id, - tags=["node_execution"], - id=node_execution_id, - trace_id=trace_id, - dotted_order=node_dotted_order, - error="", - serialized=None, - events=[], - session_id=None, - session_name=None, - reference_example_id=None, - input_attachments={}, - output_attachments={}, - ) - - self.add_run(langsmith_run) - - def message_trace(self, trace_info: MessageTraceInfo): - # get message file data - file_list = cast(list[str], trace_info.file_list) or [] - message_file_data: MessageFile | None = trace_info.message_file_data - file_url = f"{self.file_base_url}/{message_file_data.url}" if message_file_data else "" - file_list.append(file_url) - metadata = trace_info.metadata - message_data = trace_info.message_data - if message_data is None: - return - message_id = message_data.id - - user_id = message_data.from_account_id - metadata["user_id"] = user_id - - if message_data.from_end_user_id: - end_user_data: EndUser | None = db.session.get(EndUser, message_data.from_end_user_id) - if end_user_data is not None: - end_user_id = end_user_data.session_id - metadata["end_user_id"] = end_user_id - - message_run = LangSmithRunModel( - input_tokens=trace_info.message_tokens, - output_tokens=trace_info.answer_tokens, - total_tokens=trace_info.total_tokens, - id=message_id, - name=TraceTaskName.MESSAGE_TRACE, - inputs=trace_info.inputs, - run_type=LangSmithRunType.chain, - start_time=trace_info.start_time, - end_time=trace_info.end_time, - outputs=message_data.answer, - extra={"metadata": metadata}, - tags=["message", str(trace_info.conversation_mode)], - error=trace_info.error, - file_list=file_list, - serialized=None, - events=[], - session_id=None, - session_name=None, - reference_example_id=None, - input_attachments={}, - output_attachments={}, - trace_id=trace_info.trace_id, - dotted_order=None, - parent_run_id=None, - ) - self.add_run(message_run) - - # create llm run parented to message run - llm_run = LangSmithRunModel( - input_tokens=trace_info.message_tokens, - output_tokens=trace_info.answer_tokens, - total_tokens=trace_info.total_tokens, - name="llm", - inputs=trace_info.inputs, - run_type=LangSmithRunType.llm, - start_time=trace_info.start_time, - end_time=trace_info.end_time, - outputs=message_data.answer, - extra={"metadata": metadata}, - parent_run_id=message_id, - tags=["llm", str(trace_info.conversation_mode)], - error=trace_info.error, - file_list=file_list, - serialized=None, - events=[], - session_id=None, - session_name=None, - reference_example_id=None, - input_attachments={}, - output_attachments={}, - trace_id=trace_info.trace_id, - dotted_order=None, - id=str(uuid.uuid4()), - ) - self.add_run(llm_run) - - def moderation_trace(self, trace_info: ModerationTraceInfo): - if trace_info.message_data is None: - return - langsmith_run = LangSmithRunModel( - name=TraceTaskName.MODERATION_TRACE, - inputs=trace_info.inputs, - outputs={ - "action": trace_info.action, - "flagged": trace_info.flagged, - "preset_response": trace_info.preset_response, - "inputs": trace_info.inputs, - }, - run_type=LangSmithRunType.tool, - extra={"metadata": trace_info.metadata}, - tags=["moderation"], - parent_run_id=trace_info.message_id, - start_time=trace_info.start_time or trace_info.message_data.created_at, - end_time=trace_info.end_time or trace_info.message_data.updated_at, - id=str(uuid.uuid4()), - serialized=None, - events=[], - session_id=None, - session_name=None, - reference_example_id=None, - input_attachments={}, - output_attachments={}, - trace_id=trace_info.trace_id, - dotted_order=None, - error="", - file_list=[], - ) - - self.add_run(langsmith_run) - - def suggested_question_trace(self, trace_info: SuggestedQuestionTraceInfo): - message_data = trace_info.message_data - if message_data is None: - return - suggested_question_run = LangSmithRunModel( - name=TraceTaskName.SUGGESTED_QUESTION_TRACE, - inputs=trace_info.inputs, - outputs=trace_info.suggested_question, - run_type=LangSmithRunType.tool, - extra={"metadata": trace_info.metadata}, - tags=["suggested_question"], - parent_run_id=trace_info.message_id, - start_time=trace_info.start_time or message_data.created_at, - end_time=trace_info.end_time or message_data.updated_at, - id=str(uuid.uuid4()), - serialized=None, - events=[], - session_id=None, - session_name=None, - reference_example_id=None, - input_attachments={}, - output_attachments={}, - trace_id=trace_info.trace_id, - dotted_order=None, - error="", - file_list=[], - ) - - self.add_run(suggested_question_run) - - def dataset_retrieval_trace(self, trace_info: DatasetRetrievalTraceInfo): - if trace_info.message_data is None: - return - dataset_retrieval_run = LangSmithRunModel( - name=TraceTaskName.DATASET_RETRIEVAL_TRACE, - inputs=trace_info.inputs, - outputs={"documents": trace_info.documents}, - run_type=LangSmithRunType.retriever, - extra={"metadata": trace_info.metadata}, - tags=["dataset_retrieval"], - parent_run_id=trace_info.message_id, - start_time=trace_info.start_time or trace_info.message_data.created_at, - end_time=trace_info.end_time or trace_info.message_data.updated_at, - id=str(uuid.uuid4()), - serialized=None, - events=[], - session_id=None, - session_name=None, - reference_example_id=None, - input_attachments={}, - output_attachments={}, - trace_id=trace_info.trace_id, - dotted_order=None, - error="", - file_list=[], - ) - - self.add_run(dataset_retrieval_run) - - def tool_trace(self, trace_info: ToolTraceInfo): - tool_run = LangSmithRunModel( - name=trace_info.tool_name, - inputs=trace_info.tool_inputs, - outputs=trace_info.tool_outputs, - run_type=LangSmithRunType.tool, - extra={ - "metadata": trace_info.metadata, - }, - tags=["tool", trace_info.tool_name], - parent_run_id=trace_info.message_id, - start_time=trace_info.start_time, - end_time=trace_info.end_time, - file_list=[cast(str, trace_info.file_url)], - id=str(uuid.uuid4()), - serialized=None, - events=[], - session_id=None, - session_name=None, - reference_example_id=None, - input_attachments={}, - output_attachments={}, - trace_id=trace_info.trace_id, - dotted_order=None, - error=trace_info.error or "", - ) - - self.add_run(tool_run) - - def generate_name_trace(self, trace_info: GenerateNameTraceInfo): - name_run = LangSmithRunModel( - name=TraceTaskName.GENERATE_NAME_TRACE, - inputs=trace_info.inputs, - outputs=trace_info.outputs, - run_type=LangSmithRunType.tool, - extra={"metadata": trace_info.metadata}, - tags=["generate_name"], - start_time=trace_info.start_time or datetime.now(), - end_time=trace_info.end_time or datetime.now(), - id=str(uuid.uuid4()), - serialized=None, - events=[], - session_id=None, - session_name=None, - reference_example_id=None, - input_attachments={}, - output_attachments={}, - trace_id=trace_info.trace_id, - dotted_order=None, - error="", - file_list=[], - parent_run_id=None, - ) - - self.add_run(name_run) - - def add_run(self, run_data: LangSmithRunModel): - data = run_data.model_dump() - if self.project_id: - data["session_id"] = self.project_id - elif self.project_name: - data["session_name"] = self.project_name - - data = filter_none_values(data) - try: - self.langsmith_client.create_run(**data) - logger.debug("LangSmith Run created successfully.") - except Exception as e: - raise ValueError(f"LangSmith Failed to create run: {str(e)}") from e - - def update_run(self, update_run_data: LangSmithRunUpdateModel): - data = update_run_data.model_dump() - data = filter_none_values(data) - try: - self.langsmith_client.update_run(**data) - logger.debug("LangSmith Run updated successfully.") - except Exception as e: - raise ValueError(f"LangSmith Failed to update run: {str(e)}") from e - - def api_check(self): - try: - random_project_name = f"test_project_{datetime.now().strftime('%Y%m%d%H%M%S')}" - self.langsmith_client.create_project(project_name=random_project_name) - self.langsmith_client.delete_project(project_name=random_project_name) - return True - except Exception as e: - logger.debug("LangSmith API check failed", exc_info=True) - raise ValueError(f"LangSmith API check failed: {str(e)}") from e - - def get_project_url(self): - try: - run_data = RunBase( - id=uuid.uuid4(), - name="tool", - inputs={"input": "test"}, - outputs={"output": "test"}, - run_type=LangSmithRunType.tool, - start_time=datetime.now(), - ) - - project_url = self.langsmith_client.get_run_url( - run=run_data, project_id=self.project_id, project_name=self.project_name - ) - return project_url.split("/r/")[0] - except Exception as e: - logger.debug("LangSmith get run url failed", exc_info=True) - raise ValueError(f"LangSmith get run url failed: {str(e)}") from e diff --git a/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/provider.py b/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/provider.py new file mode 100644 index 00000000000000..aa48e92110fd80 --- /dev/null +++ b/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/provider.py @@ -0,0 +1,329 @@ +"""LangSmith destination through the LangSmith SDK. + +Each span becomes one complete run (outputs and end time included) with an +explicit trace ID and dotted order, so the SDK posts it synchronously and +``send`` returns only once LangSmith accepted every run. Run IDs derive from +the trace and span IDs. A root run takes the trace UUID, which is the message +or workflow run ID unless the caller supplied a trace ID, so a separately sent +fragment can attach under it without coordination. +""" + +from collections.abc import Mapping +from datetime import datetime +from typing import Literal, NoReturn, override +from urllib.parse import urlencode, urlsplit, urlunsplit +from uuid import UUID + +from dify_trace_sdk.egress import Egress, EgressResponse +from dify_trace_sdk.errors import ( + EgressBlockedError, + EgressError, + ExportOutcomeUnknownError, + ExportRejectedError, + ExportRetryableError, + TraceProviderError, + VerificationError, +) +from dify_trace_sdk.exporting import classify_egress_error, export_endpoint, raise_for_export_status +from dify_trace_sdk.ids import span_uuid, trace_uuid +from dify_trace_sdk.model import Span, SpanKind, SpanStatus, Trace +from dify_trace_sdk.provider import ProviderCapabilities, SendResult, TraceProvider +from dify_trace_sdk.transports import EgressRequestsSession +from langsmith import Client +from pydantic import JsonValue +from urllib3.util.retry import Retry + +from dify_trace_langsmith.config import LangSmithConfig + +VERIFY_TIMEOUT_SECONDS = 10.0 +DEFAULT_WEB_URL = "https://smith.langchain.com" +_LOCAL_HOSTS = frozenset({"localhost", "127.0.0.1", "0.0.0.0", "::1", "0:0:0:0:0:0:0:1"}) +_REGIONAL_WEB_URLS = ( + ("eu.", "https://eu.smith.langchain.com"), + ("aws.", "https://aws.smith.langchain.com"), + ("apac.", "https://apac.smith.langchain.com"), + ("dev.", "https://dev.smith.langchain.com"), + ("beta.", "https://beta.smith.langchain.com"), +) + +type _RunType = Literal["chain", "llm", "retriever", "tool"] + +_RUN_TYPES: Mapping[SpanKind, _RunType] = { + SpanKind.CHAIN: "chain", + SpanKind.LLM: "llm", + SpanKind.RETRIEVER: "retriever", + SpanKind.TOOL: "tool", + SpanKind.AGENT: "chain", +} + + +def _web_url(endpoint: str) -> str: + """The LangSmith web app serving an API endpoint, mapped like the SDK's ``get_host_url``. + + The SDK also counts hosts that resolve to loopback as local; resolving one + here would reach the network outside the egress. + """ + parsed = urlsplit(endpoint.rstrip("/")) + host = parsed.hostname or "" + if host in _LOCAL_HOSTS: + return "http://localhost" + for suffix in ("/api", "/api/v1"): + if parsed.path.endswith(suffix): + return urlunsplit(parsed._replace(path=parsed.path.removesuffix(suffix), query="", fragment="")) + for prefix, web_url in _REGIONAL_WEB_URLS: + if host.startswith(prefix): + return web_url + return DEFAULT_WEB_URL + + +def _first_session_field(sessions: list[object], field: str) -> str | None: + """A UUID field of the first listed tracing project, or ``None`` when none is listed.""" + if not sessions: + return None + session = sessions[0] + value = session.get(field) if isinstance(session, dict) else None + if not isinstance(value, str): + raise VerificationError("malformed_response") + try: + return str(UUID(value)) + except ValueError: + raise VerificationError("malformed_response") from None + + +def _dotted_order(run_id: str, start_time: datetime, parent: str | None) -> str: + """LangSmith's ordering key: one ``Z`` segment per run from the trace root down.""" + segment = f"{start_time:%Y%m%dT%H%M%S%f}Z{run_id}" + return segment if parent is None else f"{parent}.{segment}" + + +def _payload(value: JsonValue, key: str) -> dict[str, JsonValue]: + """Run inputs and outputs are objects; any other value is wrapped under ``key``.""" + if isinstance(value, dict): + return dict(value) + return {} if value is None else {key: value} + + +def _usage(span: Span) -> dict[str, JsonValue]: + """An LLM span's token counts and USD cost, keyed like LangSmith's ``usage_metadata``.""" + usage: dict[str, JsonValue] = {} + if span.kind is not SpanKind.LLM: + return usage + if span.usage is not None: + usage["input_tokens"] = span.usage.prompt_tokens + usage["output_tokens"] = span.usage.completion_tokens + usage["total_tokens"] = span.usage.total_tokens + if span.cost is not None and span.cost.currency.upper() == "USD": + usage["total_cost"] = span.cost.amount + return usage + + +def _metadata(trace: Trace, span: Span) -> dict[str, JsonValue]: + metadata: dict[str, JsonValue] = {**span.attributes, "dify_trace_id": trace.trace_id} + if trace.session_id is not None: + # LangSmith groups traces into threads by this key. + metadata["session_id"] = trace.session_id + if trace.external_trace_id is not None: + metadata["external_trace_id"] = trace.external_trace_id + if span.kind is SpanKind.LLM and span.model_name: + # LangSmith labels model calls, and prices those sent without a cost, by these keys. + metadata["ls_model_name"] = span.model_name + if span.model_provider: + metadata["ls_provider"] = span.model_provider + return metadata + + +def _trace_run_id(trace: Trace) -> str: + """The run the unit hangs from: the message or workflow run itself, or its root span when a + caller's trace ID groups several of them, since LangSmith keeps one root run per trace.""" + if trace.external_trace_id is None: + return str(trace_uuid(trace.trace_id)) + anchor = trace.parent.span_id if trace.parent is not None else trace.root.span_id + return str(span_uuid(trace.trace_id, anchor)) + + +def _run_link_base(project_url: str | None) -> str | None: + """Run links need a project page; until the project exists, ``verify`` stores the project list.""" + if project_url is None or "/projects/p/" not in project_url: + return None + return project_url.rstrip("/") + + +def _export_failure(error: BaseException) -> TraceProviderError: + """Recover the failure ``_ClassifyingEgress`` raised from the SDK's wrapping exceptions.""" + cause: BaseException | None = error + while cause is not None: + if isinstance(cause, TraceProviderError): + return cause + cause = cause.__cause__ + return ExportRejectedError("LangSmith rejected the trace") + + +class _ClassifyingEgress: + """The egress the SDK exports through, classifying each failure while its HTTP status is known. + + The SDK folds most statuses into an untyped ``LangSmithError`` chained to + whatever its session raised, so ``send`` recovers these errors from that chain. + """ + + def __init__(self, egress: Egress) -> None: + self._egress = egress + + def request( + self, + method: str, + url: str, + *, + headers: Mapping[str, str], + content: bytes | None, + timeout: float, + sensitive_url: bool, + ) -> EgressResponse: + try: + response = self._egress.request( + method, url, headers=headers, content=content, timeout=timeout, sensitive_url=sensitive_url + ) + except EgressError as error: + raise classify_egress_error(error, "LangSmith") from None + # LangSmith answers 409 for a run it already has, which the SDK counts as accepted. + if response.status_code != 409: + raise_for_export_status(response.status_code, "LangSmith") + return response + + def grpc_channel(self, endpoint: str) -> NoReturn: + raise EgressBlockedError("LangSmith export does not use gRPC") + + +class LangSmithProvider(TraceProvider[LangSmithConfig]): + id = "langsmith" + config_model = LangSmithConfig + capabilities = ProviderCapabilities(native_cost=True, trace_url=True, span_url=True) + + @override + def verify(self, config: LangSmithConfig) -> LangSmithConfig: + # Like the SDK's ``_get_tenant_id``: any listed project names the key's workspace. + tenant_id = _first_session_field(self._sessions(config, {"limit": "1"}), "tenant_id") + # Like the SDK's ``read_project``. + project_id = _first_session_field( + self._sessions(config, {"limit": "1", "name": config.project, "include_stats": "false"}), "id" + ) + web_url = _web_url(config.endpoint) + if tenant_id is None: + project_url = web_url + elif project_id is None: + # LangSmith creates the project with its first run. + project_url = f"{web_url}/o/{tenant_id}/projects" + else: + project_url = f"{web_url}/o/{tenant_id}/projects/p/{project_id}" + return config.model_copy(update={"project_url": project_url}) + + def _sessions(self, config: LangSmithConfig, query: Mapping[str, str]) -> list[object]: + """List tracing projects, which the LangSmith API calls sessions.""" + try: + response = self.egress.request( + "GET", + f"{config.endpoint.rstrip('/')}/sessions?{urlencode(query)}", + headers={"x-api-key": config.api_key, "Accept": "application/json"}, + content=None, + timeout=VERIFY_TIMEOUT_SECONDS, + sensitive_url=False, + ) + except EgressError: + raise VerificationError("connection_failed") from None + if response.status_code in {401, 403}: + raise VerificationError("authentication_failed") + if not 200 <= response.status_code < 300: + raise VerificationError("endpoint_unavailable") + try: + body = response.json() + except ValueError: + raise VerificationError("malformed_response") from None + if not isinstance(body, list): + raise VerificationError("malformed_response") + return body + + @override + def project_url(self, config: LangSmithConfig) -> str | None: + return config.project_url + + @override + def send(self, config: LangSmithConfig, trace: Trace) -> SendResult: + api_url = export_endpoint(config.endpoint) + root_run_id = _trace_run_id(trace) + run_ids = {span.span_id: str(span_uuid(trace.trace_id, span.span_id)) for span in trace.spans} + dotted_orders: dict[str, str] = {} + if trace.parent is None: + run_ids[trace.root.span_id] = root_run_id + else: + # The trace root was sent in another unit; its run ID and dotted order are derivable. + run_ids[trace.parent.span_id] = root_run_id + dotted_orders[trace.parent.span_id] = _dotted_order(root_run_id, trace.parent.start_time, None) + client = self._client(config, api_url) + accepted: list[str] = [] + try: + for span in trace.spans: + run_id = run_ids[span.span_id] + parent_order = dotted_orders[span.parent_span_id] if span.parent_span_id else None + dotted_orders[span.span_id] = _dotted_order(run_id, span.start_time, parent_order) + outputs = _payload(span.outputs, "output") + if usage := _usage(span): + # Run ingestion takes token counts and costs from here, not from top-level run fields. + outputs["usage_metadata"] = usage + client.create_run( + id=run_id, + name=span.name, + inputs=_payload(span.inputs, "input"), + outputs=outputs, + run_type=_RUN_TYPES[span.kind], + project_name=config.project, + start_time=span.start_time, + end_time=span.end_time, + parent_run_id=run_ids[span.parent_span_id] if span.parent_span_id else None, + trace_id=root_run_id, + dotted_order=dotted_orders[span.span_id], + extra={"metadata": _metadata(trace, span)}, + error=(span.error or "Trace operation failed") if span.status is SpanStatus.ERROR else None, + ) + accepted.append(run_id) + except Exception as error: + failure = _export_failure(error) + if isinstance(failure, ExportRetryableError | ExportOutcomeUnknownError) and accepted: + # Run IDs are deterministic and LangSmith answers 409 for runs it has, so resending is safe. + raise ExportRetryableError("LangSmith accepted part of the trace") from None + raise failure from None + finally: + client.close() + link_base = _run_link_base(config.project_url) + if link_base is None: + return SendResult(trace_id=root_run_id, trace_url=None, span_urls={}) + return SendResult( + trace_id=root_run_id, + trace_url=f"{link_base}/r/{root_run_id}", + span_urls={span.span_id: f"{link_base}/r/{run_ids[span.span_id]}" for span in trace.spans}, + ) + + def _client(self, config: LangSmithConfig, api_url: str) -> Client: + session = EgressRequestsSession(_ClassifyingEgress(self.egress)) + try: + # Synchronous posts without a background thread or retries, no host runtime details, + # no GET /info discovery, and every argument that would otherwise fall back to a + # LANGSMITH_* environment variable pinned. + return Client( + api_url=api_url, + api_key=config.api_key, + session=session, + retry_config=Retry(total=0), + auto_batch_tracing=False, + tracing_sampling_rate=1.0, + tracing_mode="langsmith", + omit_traced_runtime_info=True, + hide_inputs=False, + hide_outputs=False, + hide_metadata=False, + workspace_id="", + info={}, + # Closing a client that uses the prompt cache shuts down the process-wide cache. + disable_prompt_cache=True, + ) + except Exception: + # The SDK rejects some host environments, e.g. conflicting LANGSMITH_* endpoints. + raise ExportRejectedError("LangSmith client could not be created") from None 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 deleted file mode 100644 index a0d3bf6eef57a5..00000000000000 --- a/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/unified_trace.py +++ /dev/null @@ -1,183 +0,0 @@ -"""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/providers/trace/trace-langsmith/tests/unit_tests/src/dify_trace_langsmith/test_config.py b/api/providers/trace/trace-langsmith/tests/unit_tests/src/dify_trace_langsmith/test_config.py index d941dca987c5f2..2edb9ca10cdb05 100644 --- a/api/providers/trace/trace-langsmith/tests/unit_tests/src/dify_trace_langsmith/test_config.py +++ b/api/providers/trace/trace-langsmith/tests/unit_tests/src/dify_trace_langsmith/test_config.py @@ -2,44 +2,47 @@ from dify_trace_langsmith.config import LangSmithConfig from pydantic import ValidationError - -class TestLangSmithConfig: - """Test cases for LangSmithConfig""" - - def test_valid_config(self): - """Test valid LangSmith configuration""" - config = LangSmithConfig(api_key="test_key", project="test_project", endpoint="https://custom.smith.com") - assert config.api_key == "test_key" - assert config.project == "test_project" - assert config.endpoint == "https://custom.smith.com" - - def test_default_values(self): - """Test default values are set correctly""" - config = LangSmithConfig(api_key="key", project="project") - assert config.endpoint == "https://api.smith.langchain.com" - - def test_missing_required_fields(self): - """Test that required fields are enforced""" - with pytest.raises(ValidationError): - LangSmithConfig.model_validate({}) - - with pytest.raises(ValidationError): - LangSmithConfig.model_validate({"api_key": "key"}) - - with pytest.raises(ValidationError): - LangSmithConfig.model_validate({"project": "project"}) - - def test_endpoint_validation_https_only(self): - """Test endpoint validation only allows HTTPS""" - with pytest.raises(ValidationError, match="URL must start with https://"): - LangSmithConfig(api_key="key", project="project", endpoint="http://insecure.com") - - def test_endpoint_preserves_path(self): - """Self-hosted LangSmith endpoints keep their API path prefix""" - config = LangSmithConfig(api_key="key", project="project", endpoint="https://langsmith.internal/api") - assert config.endpoint == "https://langsmith.internal/api" - - def test_endpoint_preserves_versioned_path(self): - """Self-hosted LangSmith endpoints keep multi-segment API paths""" - config = LangSmithConfig(api_key="key", project="project", endpoint="https://langsmith.internal/api/v1") - assert config.endpoint == "https://langsmith.internal/api/v1" +API_KEY = "lsv2_pt_test-key" +TENANT_ID = "4f1f0b4e-6a51-4c5a-9a8e-2d7f0c3b9e11" +PROJECT_ID = "0d9b8f61-2c3e-4b7a-8f5d-6e1a9c4b7d22" +PROJECT_URL = f"https://smith.langchain.com/o/{TENANT_ID}/projects/p/{PROJECT_ID}" + + +def test_config_declares_secret_and_derived_fields() -> None: + assert LangSmithConfig.secret_fields() == {"api_key"} + assert LangSmithConfig.derived_fields() == {"project_url"} + assert [(field.name, field.secret, field.required) for field in LangSmithConfig.form_fields()] == [ + ("api_key", True, True), + ("project", False, True), + ("endpoint", False, False), + ] + + +def test_config_reads_stored_values() -> None: + stored = {"api_key": " lsv2_pt_stored ", "project": "dify-app", "endpoint": " ", "project_url": PROJECT_URL} + config = LangSmithConfig.model_validate(stored) + assert config.api_key == "lsv2_pt_stored" + assert config.endpoint == "https://api.smith.langchain.com" + assert config.project_url == PROJECT_URL + + +@pytest.mark.parametrize("endpoint", ["https://langsmith.internal/api", "https://langsmith.internal/api/v1"]) +def test_config_keeps_the_endpoint_path(endpoint: str) -> None: + assert LangSmithConfig(api_key=API_KEY, project="dify-app", endpoint=endpoint).endpoint == endpoint + + +@pytest.mark.parametrize( + "fields", + [ + {"project": "dify-app"}, + {"api_key": API_KEY}, + {"api_key": " ", "project": "dify-app"}, + {"api_key": API_KEY, "project": " "}, + {"api_key": API_KEY, "project": "dify-app", "endpoint": "http://langsmith.internal"}, + {"api_key": API_KEY, "project": "dify-app", "endpoint": "https://user:hunter2@langsmith.internal"}, + ], +) +def test_config_rejects_invalid_values_without_echoing_them(fields: dict[str, str]) -> None: + with pytest.raises(ValidationError) as caught: + LangSmithConfig.model_validate(fields) + assert "hunter2" not in str(caught.value) diff --git a/api/providers/trace/trace-langsmith/tests/unit_tests/src/dify_trace_langsmith/test_langsmith_trace.py b/api/providers/trace/trace-langsmith/tests/unit_tests/src/dify_trace_langsmith/test_langsmith_trace.py deleted file mode 100644 index 264ab4d2862274..00000000000000 --- a/api/providers/trace/trace-langsmith/tests/unit_tests/src/dify_trace_langsmith/test_langsmith_trace.py +++ /dev/null @@ -1,743 +0,0 @@ -"""Unit tests for LangSmith trace translation with SQLite-backed lookups.""" - -import collections -from datetime import datetime, timedelta -from types import SimpleNamespace -from typing import override -from unittest.mock import MagicMock - -import pytest -from dify_trace_langsmith.config import LangSmithConfig -from dify_trace_langsmith.entities.langsmith_trace_entity import ( - LangSmithRunModel, - LangSmithRunType, - LangSmithRunUpdateModel, -) -from dify_trace_langsmith.langsmith_trace import LangSmithDataTrace -from sqlalchemy.orm import Session - -from core.ops.entities.trace_entity import ( - DatasetRetrievalTraceInfo, - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - SuggestedQuestionTraceInfo, - ToolTraceInfo, - TraceTaskName, - WorkflowTraceInfo, -) -from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey -from models import EndUser -from models.enums import EndUserType -from tests.helpers.unit.core.ops.trace_fixtures import message_trace_info, tool_trace_info, workflow_trace_info - - -def _dt() -> datetime: - return datetime(2024, 1, 1, 0, 0, 0) - - -@pytest.fixture -def langsmith_config(): - return LangSmithConfig(api_key="ls-123", project="default", endpoint="https://api.smith.langchain.com") - - -@pytest.fixture -def trace_instance(langsmith_config, monkeypatch: pytest.MonkeyPatch): - # Mock LangSmith client - mock_client = MagicMock() - monkeypatch.setattr("dify_trace_langsmith.langsmith_trace.Client", lambda **kwargs: mock_client) - - instance = LangSmithDataTrace(langsmith_config) - return instance - - -def test_init(langsmith_config, monkeypatch: pytest.MonkeyPatch): - mock_client_class = MagicMock() - monkeypatch.setattr("dify_trace_langsmith.langsmith_trace.Client", mock_client_class) - monkeypatch.setenv("FILES_URL", "http://test.url") - - instance = LangSmithDataTrace(langsmith_config) - - mock_client_class.assert_called_once_with(api_key=langsmith_config.api_key, api_url=langsmith_config.endpoint) - assert instance.langsmith_key == langsmith_config.api_key - assert instance.project_name == langsmith_config.project - assert instance.file_base_url == "http://test.url" - - -def test_init_passes_self_hosted_path_to_client(monkeypatch: pytest.MonkeyPatch): - config = LangSmithConfig(api_key="ls-123", project="default", endpoint="https://langsmith.internal/api") - mock_client_class = MagicMock() - monkeypatch.setattr("dify_trace_langsmith.langsmith_trace.Client", mock_client_class) - - LangSmithDataTrace(config) - - mock_client_class.assert_called_once_with(api_key="ls-123", api_url="https://langsmith.internal/api") - - -def test_trace_dispatch(trace_instance, monkeypatch: pytest.MonkeyPatch): - methods = [ - "workflow_trace", - "message_trace", - "moderation_trace", - "suggested_question_trace", - "dataset_retrieval_trace", - "tool_trace", - "generate_name_trace", - ] - mocks = {method: MagicMock() for method in methods} - for method, m in mocks.items(): - monkeypatch.setattr(trace_instance, method, m) - - # WorkflowTraceInfo - info = workflow_trace_info() - trace_instance.trace(info) - mocks["workflow_trace"].assert_called_once_with(info) - - # MessageTraceInfo - info = message_trace_info() - trace_instance.trace(info) - mocks["message_trace"].assert_called_once_with(info) - - # ModerationTraceInfo - info = ModerationTraceInfo(metadata={}, flagged=False, action="allow", preset_response="", query="hello") - trace_instance.trace(info) - mocks["moderation_trace"].assert_called_once_with(info) - - # SuggestedQuestionTraceInfo - info = SuggestedQuestionTraceInfo(metadata={}, total_tokens=0, suggested_question=[], level="info") - trace_instance.trace(info) - mocks["suggested_question_trace"].assert_called_once_with(info) - - # DatasetRetrievalTraceInfo - info = DatasetRetrievalTraceInfo(metadata={}) - trace_instance.trace(info) - mocks["dataset_retrieval_trace"].assert_called_once_with(info) - - # ToolTraceInfo - info = tool_trace_info() - trace_instance.trace(info) - mocks["tool_trace"].assert_called_once_with(info) - - # GenerateNameTraceInfo - info = GenerateNameTraceInfo(tenant_id="tenant-1", metadata={}) - trace_instance.trace(info) - mocks["generate_name_trace"].assert_called_once_with(info) - - -@pytest.mark.parametrize("sqlite3_session", [()], indirect=True) -def test_workflow_trace(trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session) -> None: - # Setup trace info - workflow_data = MagicMock() - workflow_data.created_at = _dt() - workflow_data.finished_at = _dt() + timedelta(seconds=1) - - trace_info = WorkflowTraceInfo( - tenant_id="tenant-1", - workflow_id="wf-1", - workflow_run_id="run-1", - workflow_run_inputs={"input": "hi"}, - workflow_run_outputs={"output": "hello"}, - workflow_run_status="succeeded", - workflow_run_version="1.0", - workflow_run_elapsed_time=1.0, - total_tokens=100, - file_list=[], - query="hi", - message_id="msg-1", - conversation_id="conv-1", - start_time=_dt(), - end_time=_dt() + timedelta(seconds=1), - trace_id="trace-1", - metadata={"app_id": "app-1"}, - workflow_app_log_id="log-1", - error="", - workflow_data=workflow_data, - ) - - monkeypatch.setattr( - "dify_trace_langsmith.langsmith_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - - # Mock node executions - node_llm = MagicMock() - node_llm.id = "node-llm" - node_llm.title = "LLM Node" - node_llm.node_type = BuiltinNodeTypes.LLM - node_llm.status = "succeeded" - node_llm.process_data = { - "model_mode": "chat", - "model_name": "gpt-4", - "model_provider": "openai", - "usage": {"prompt_tokens": 10, "completion_tokens": 20}, - } - node_llm.inputs = {"prompts": "p"} - node_llm.outputs = {"text": "t"} - node_llm.created_at = _dt() - node_llm.elapsed_time = 0.5 - node_llm.metadata = {WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: 30} - - node_other = MagicMock() - node_other.id = "node-other" - node_other.title = "Tool Node" - node_other.node_type = BuiltinNodeTypes.TOOL - node_other.status = "succeeded" - node_other.process_data = None - node_other.inputs = {"tool_input": "val"} - node_other.outputs = {"tool_output": "val"} - node_other.created_at = None # Trigger datetime.now() - node_other.elapsed_time = 0.2 - node_other.metadata = {} - - node_retrieval = MagicMock() - node_retrieval.id = "node-retrieval" - node_retrieval.title = "Retrieval Node" - node_retrieval.node_type = BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL - node_retrieval.status = "succeeded" - node_retrieval.process_data = None - node_retrieval.inputs = {"query": "val"} - node_retrieval.outputs = {"results": "val"} - node_retrieval.created_at = _dt() - node_retrieval.elapsed_time = 0.2 - node_retrieval.metadata = {} - - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [node_llm, node_other, node_retrieval] - - mock_factory = MagicMock() - mock_factory.create_workflow_node_execution_repository.return_value = repo - monkeypatch.setattr("dify_trace_langsmith.langsmith_trace.DifyCoreRepositoryFactory", mock_factory) - - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - trace_instance.add_run = MagicMock() - - trace_instance.workflow_trace(trace_info) - - # Verify add_run calls - # 1. message run (id="msg-1") - # 2. workflow run (id="run-1") - # 3. node llm run (id="node-llm") - # 4. node other run (id="node-other") - # 5. node retrieval run (id="node-retrieval") - assert trace_instance.add_run.call_count == 5 - - call_args = [call[0][0] for call in trace_instance.add_run.call_args_list] - - assert call_args[0].id == "msg-1" - assert call_args[0].name == TraceTaskName.MESSAGE_TRACE - # trace_id must equal root run's id (message_id), not the external trace_id "trace-1" - assert call_args[0].trace_id == "msg-1" - - assert call_args[1].id == "run-1" - assert call_args[1].name == TraceTaskName.WORKFLOW_TRACE - assert call_args[1].parent_run_id == "msg-1" - assert call_args[1].trace_id == "msg-1" - - assert call_args[2].id == "node-llm" - assert call_args[2].run_type == LangSmithRunType.llm - assert call_args[2].trace_id == "msg-1" - - assert call_args[3].id == "node-other" - assert call_args[3].run_type == LangSmithRunType.tool - - assert call_args[4].id == "node-retrieval" - assert call_args[4].run_type == LangSmithRunType.retriever - - -@pytest.mark.parametrize("sqlite3_session", [()], indirect=True) -def test_workflow_trace_no_start_time( - trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session -) -> None: - workflow_data = MagicMock() - workflow_data.created_at = _dt() - workflow_data.finished_at = _dt() + timedelta(seconds=1) - - trace_info = WorkflowTraceInfo( - tenant_id="tenant-1", - workflow_id="wf-1", - workflow_run_id="run-1", - workflow_run_inputs={}, - workflow_run_outputs={}, - workflow_run_status="succeeded", - workflow_run_version="1.0", - workflow_run_elapsed_time=1.0, - total_tokens=10, - file_list=[], - query="hi", - message_id="msg-1", - conversation_id="conv-1", - start_time=None, - end_time=None, - trace_id="trace-1", - metadata={"app_id": "app-1"}, - workflow_app_log_id="log-1", - error="", - workflow_data=workflow_data, - ) - - monkeypatch.setattr( - "dify_trace_langsmith.langsmith_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [] - mock_factory = MagicMock() - mock_factory.create_workflow_node_execution_repository.return_value = repo - monkeypatch.setattr("dify_trace_langsmith.langsmith_trace.DifyCoreRepositoryFactory", mock_factory) - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - trace_instance.add_run = MagicMock() - trace_instance.workflow_trace(trace_info) - assert trace_instance.add_run.called - - -@pytest.mark.parametrize("sqlite3_session", [()], indirect=True) -def test_workflow_trace_missing_app_id( - trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session -) -> None: - trace_info = workflow_trace_info() - trace_info.trace_id = "trace-1" - trace_info.message_id = None - trace_info.workflow_run_id = "run-1" - trace_info.start_time = None - trace_info.workflow_data = MagicMock() - trace_info.workflow_data.created_at = _dt() - trace_info.metadata = {} # Empty metadata - trace_info.workflow_app_log_id = "log-1" - trace_info.file_list = [] - trace_info.total_tokens = 0 - trace_info.workflow_run_inputs = {} - trace_info.workflow_run_outputs = {} - trace_info.error = "" - - monkeypatch.setattr( - "dify_trace_langsmith.langsmith_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - - with pytest.raises(ValueError, match="No app_id found in trace_info metadata"): - trace_instance.workflow_trace(trace_info) - - -@pytest.mark.parametrize("sqlite3_session", [(EndUser,)], indirect=True) -def test_message_trace(trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session) -> None: - message_data = MagicMock() - message_data.id = "msg-1" - message_data.from_account_id = "acc-1" - message_data.from_end_user_id = "end-user-1" - message_data.answer = "hello answer" - - trace_info = MessageTraceInfo( - message_id="msg-1", - message_data=message_data, - inputs={"input": "hi"}, - outputs={"answer": "hello"}, - message_tokens=10, - answer_tokens=20, - total_tokens=30, - start_time=_dt(), - end_time=_dt() + timedelta(seconds=1), - trace_id="trace-1", - metadata={"foo": "bar"}, - conversation_mode="chat", - conversation_model="gpt-4", - file_list=[], - error=None, - message_file_data=MagicMock(url="file-url"), - ) - - end_user = EndUser( - id="end-user-1", - tenant_id="tenant-1", - app_id="app-1", - type=EndUserType.BROWSER, - session_id="session-id-123", - ) - sqlite3_session.add(end_user) - sqlite3_session.commit() - monkeypatch.setattr( - "dify_trace_langsmith.langsmith_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - - trace_instance.add_run = MagicMock() - - trace_instance.message_trace(trace_info) - - # 1. message run - # 2. llm run - assert trace_instance.add_run.call_count == 2 - - call_args = [call[0][0] for call in trace_instance.add_run.call_args_list] - assert call_args[0].id == "msg-1" - assert call_args[0].extra["metadata"]["end_user_id"] == "session-id-123" - assert call_args[1].parent_run_id == "msg-1" - assert call_args[1].name == "llm" - - -def test_message_trace_no_data(trace_instance): - trace_info = message_trace_info() - trace_info.message_data = None - trace_info.file_list = [] - trace_info.message_file_data = None - trace_info.metadata = {} - trace_instance.add_run = MagicMock() - trace_instance.message_trace(trace_info) - trace_instance.add_run.assert_not_called() - - -def test_moderation_trace_no_data(trace_instance): - trace_info = ModerationTraceInfo(metadata={}, flagged=False, action="allow", preset_response="", query="hello") - trace_info.message_data = None - trace_instance.add_run = MagicMock() - trace_instance.moderation_trace(trace_info) - trace_instance.add_run.assert_not_called() - - -def test_suggested_question_trace_no_data(trace_instance): - trace_info = SuggestedQuestionTraceInfo(metadata={}, total_tokens=0, suggested_question=[], level="info") - trace_info.message_data = None - trace_instance.add_run = MagicMock() - trace_instance.suggested_question_trace(trace_info) - trace_instance.add_run.assert_not_called() - - -def test_dataset_retrieval_trace_no_data(trace_instance): - trace_info = DatasetRetrievalTraceInfo(metadata={}) - trace_info.message_data = None - trace_instance.add_run = MagicMock() - trace_instance.dataset_retrieval_trace(trace_info) - trace_instance.add_run.assert_not_called() - - -def test_moderation_trace(trace_instance): - message_data = MagicMock() - message_data.created_at = _dt() - message_data.updated_at = _dt() - - trace_info = ModerationTraceInfo( - message_id="msg-1", - message_data=message_data, - inputs={"q": "hi"}, - action="stop", - flagged=True, - preset_response="blocked", - start_time=None, - end_time=None, - metadata={}, - trace_id="trace-1", - query="hi", - ) - - trace_instance.add_run = MagicMock() - trace_instance.moderation_trace(trace_info) - trace_instance.add_run.assert_called_once() - assert trace_instance.add_run.call_args[0][0].name == TraceTaskName.MODERATION_TRACE - - -def test_suggested_question_trace(trace_instance): - message_data = MagicMock() - message_data.created_at = _dt() - message_data.updated_at = _dt() - - trace_info = SuggestedQuestionTraceInfo( - message_id="msg-1", - message_data=message_data, - inputs="hi", - suggested_question=["q1"], - total_tokens=10, - level="info", - start_time=None, - end_time=None, - metadata={}, - trace_id="trace-1", - ) - - trace_instance.add_run = MagicMock() - trace_instance.suggested_question_trace(trace_info) - trace_instance.add_run.assert_called_once() - assert trace_instance.add_run.call_args[0][0].name == TraceTaskName.SUGGESTED_QUESTION_TRACE - - -def test_dataset_retrieval_trace(trace_instance): - message_data = MagicMock() - message_data.created_at = _dt() - message_data.updated_at = _dt() - - trace_info = DatasetRetrievalTraceInfo( - message_id="msg-1", - message_data=message_data, - inputs="query", - documents=[{"id": "doc1"}], - start_time=None, - end_time=None, - metadata={}, - trace_id="trace-1", - ) - - trace_instance.add_run = MagicMock() - trace_instance.dataset_retrieval_trace(trace_info) - trace_instance.add_run.assert_called_once() - assert trace_instance.add_run.call_args[0][0].name == TraceTaskName.DATASET_RETRIEVAL_TRACE - - -def test_tool_trace(trace_instance): - trace_info = ToolTraceInfo( - message_id="msg-1", - message_data=MagicMock(), - inputs={}, - outputs={}, - tool_name="my_tool", - tool_inputs={"a": 1}, - tool_outputs="result", - time_cost=0.1, - start_time=_dt(), - end_time=_dt(), - metadata={}, - trace_id="trace-1", - tool_config={}, - tool_parameters={}, - file_url="http://file", - ) - - trace_instance.add_run = MagicMock() - trace_instance.tool_trace(trace_info) - trace_instance.add_run.assert_called_once() - assert trace_instance.add_run.call_args[0][0].name == "my_tool" - - -def test_generate_name_trace(trace_instance): - trace_info = GenerateNameTraceInfo( - inputs={"q": "hi"}, - outputs={"name": "new"}, - tenant_id="tenant-1", - conversation_id="conv-1", - start_time=None, - end_time=None, - metadata={}, - trace_id="trace-1", - ) - - trace_instance.add_run = MagicMock() - trace_instance.generate_name_trace(trace_info) - trace_instance.add_run.assert_called_once() - assert trace_instance.add_run.call_args[0][0].name == TraceTaskName.GENERATE_NAME_TRACE - - -def test_add_run_success(trace_instance): - run_data = LangSmithRunModel( - id="run-1", name="test", inputs={}, outputs={}, run_type=LangSmithRunType.tool, start_time=_dt() - ) - trace_instance.project_id = "proj-1" - trace_instance.add_run(run_data) - trace_instance.langsmith_client.create_run.assert_called_once() - args, kwargs = trace_instance.langsmith_client.create_run.call_args - assert kwargs["session_id"] == "proj-1" - - -def test_add_run_error(trace_instance): - run_data = LangSmithRunModel(id="run-1", name="test", run_type=LangSmithRunType.tool, start_time=_dt()) - trace_instance.langsmith_client.create_run.side_effect = Exception("failed") - with pytest.raises(ValueError, match="LangSmith Failed to create run: failed"): - trace_instance.add_run(run_data) - - -def test_update_run_success(trace_instance): - update_data = LangSmithRunUpdateModel(run_id="run-1", outputs={"out": "val"}) - trace_instance.update_run(update_data) - trace_instance.langsmith_client.update_run.assert_called_once() - - -def test_update_run_error(trace_instance): - update_data = LangSmithRunUpdateModel(run_id="run-1") - trace_instance.langsmith_client.update_run.side_effect = Exception("failed") - with pytest.raises(ValueError, match="LangSmith Failed to update run: failed"): - trace_instance.update_run(update_data) - - -@pytest.mark.parametrize("sqlite3_session", [()], indirect=True) -def test_workflow_trace_usage_extraction_error( - trace_instance, - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - sqlite3_session: Session, -) -> None: - workflow_data = MagicMock() - workflow_data.created_at = _dt() - workflow_data.finished_at = _dt() + timedelta(seconds=1) - - trace_info = WorkflowTraceInfo( - tenant_id="tenant-1", - workflow_id="wf-1", - workflow_run_id="run-1", - workflow_run_inputs={}, - workflow_run_outputs={}, - workflow_run_status="succeeded", - workflow_run_version="1.0", - workflow_run_elapsed_time=1.0, - total_tokens=100, - file_list=[], - query="hi", - message_id="msg-1", - conversation_id="conv-1", - start_time=_dt(), - end_time=_dt(), - trace_id="trace-1", - metadata={"app_id": "app-1"}, - workflow_app_log_id="log-1", - error="", - workflow_data=workflow_data, - ) - - class BadDict(collections.UserDict): - @override - def get(self, key, default=None): - if key == "usage": - raise Exception("Usage extraction failed") - return super().get(key, default) - - node_llm = MagicMock() - node_llm.id = "node-llm" - node_llm.title = "LLM Node" - node_llm.node_type = BuiltinNodeTypes.LLM - node_llm.status = "succeeded" - node_llm.process_data = BadDict({"model_mode": "chat", "model_name": "gpt-4", "usage": True, "prompts": ["p"]}) - node_llm.inputs = {} - node_llm.outputs = {} - node_llm.created_at = _dt() - node_llm.elapsed_time = 0.5 - node_llm.metadata = {} - - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [node_llm] - - mock_factory = MagicMock() - mock_factory.create_workflow_node_execution_repository.return_value = repo - monkeypatch.setattr("dify_trace_langsmith.langsmith_trace.DifyCoreRepositoryFactory", mock_factory) - monkeypatch.setattr( - "dify_trace_langsmith.langsmith_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - trace_instance.add_run = MagicMock() - - import logging - - with caplog.at_level(logging.ERROR): - trace_instance.workflow_trace(trace_info) - - assert "Failed to extract usage" in caplog.text - - -def test_api_check_success(trace_instance): - assert trace_instance.api_check() is True - assert trace_instance.langsmith_client.create_project.called - assert trace_instance.langsmith_client.delete_project.called - - -def test_api_check_error(trace_instance): - trace_instance.langsmith_client.create_project.side_effect = Exception("error") - with pytest.raises(ValueError, match="LangSmith API check failed: error"): - trace_instance.api_check() - - -def test_get_project_url_success(trace_instance): - trace_instance.langsmith_client.get_run_url.return_value = "https://smith.langchain.com/o/org/p/proj/r/run" - url = trace_instance.get_project_url() - assert url == "https://smith.langchain.com/o/org/p/proj" - - -def test_get_project_url_error(trace_instance): - trace_instance.langsmith_client.get_run_url.side_effect = Exception("error") - with pytest.raises(ValueError, match="LangSmith get run url failed: error"): - trace_instance.get_project_url() - - -def _make_workflow_trace_info( - *, message_id: str | None, workflow_run_id: str, trace_id: str | None -) -> WorkflowTraceInfo: - workflow_data = MagicMock() - workflow_data.created_at = _dt() - workflow_data.finished_at = _dt() + timedelta(seconds=1) - return WorkflowTraceInfo( - tenant_id="tenant-1", - workflow_id="wf-1", - workflow_run_id=workflow_run_id, - workflow_run_inputs={}, - workflow_run_outputs={}, - workflow_run_status="succeeded", - workflow_run_version="1.0", - workflow_run_elapsed_time=1.0, - total_tokens=0, - file_list=[], - query="q", - message_id=message_id, - conversation_id="conv-1" if message_id else None, - start_time=_dt(), - end_time=_dt() + timedelta(seconds=1), - trace_id=trace_id, - metadata={"app_id": "app-1"}, - workflow_app_log_id=None, - error=None, - workflow_data=workflow_data, - ) - - -def _patch_workflow_trace_deps(monkeypatch, trace_instance, sqlite3_session: Session) -> None: - monkeypatch.setattr( - "dify_trace_langsmith.langsmith_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [] - factory = MagicMock() - factory.create_workflow_node_execution_repository.return_value = repo - monkeypatch.setattr("dify_trace_langsmith.langsmith_trace.DifyCoreRepositoryFactory", factory) - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - trace_instance.add_run = MagicMock() - - -@pytest.mark.parametrize("sqlite3_session", [()], indirect=True) -def test_workflow_trace_id_uses_message_id_not_external( - trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session -) -> None: - """Chatflow with external trace_id: LangSmith trace_id must be message_id, not external.""" - trace_info = _make_workflow_trace_info( - message_id="msg-abc", - workflow_run_id="run-xyz", - trace_id="external-999", - ) - _patch_workflow_trace_deps(monkeypatch, trace_instance, sqlite3_session) - - trace_instance.workflow_trace(trace_info) - - calls = [c[0][0] for c in trace_instance.add_run.call_args_list] - # message run (root) and workflow run (child) must both use message_id as trace_id - assert calls[0].id == "msg-abc" - assert calls[0].trace_id == "msg-abc" - assert calls[1].id == "run-xyz" - assert calls[1].trace_id == "msg-abc" - # external_trace_id preserved in metadata - assert trace_info.metadata.get("external_trace_id") == "external-999" - - -@pytest.mark.parametrize("sqlite3_session", [()], indirect=True) -def test_workflow_trace_id_pure_workflow_uses_run_id( - trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session -) -> None: - """Pure workflow (no message_id) with external trace_id: trace_id must be workflow_run_id.""" - trace_info = _make_workflow_trace_info( - message_id=None, - workflow_run_id="run-xyz", - trace_id="external-999", - ) - _patch_workflow_trace_deps(monkeypatch, trace_instance, sqlite3_session) - - trace_instance.workflow_trace(trace_info) - - calls = [c[0][0] for c in trace_instance.add_run.call_args_list] - # workflow run is the root; trace_id must equal its run_id - assert calls[0].id == "run-xyz" - assert calls[0].trace_id == "run-xyz" diff --git a/api/providers/trace/trace-langsmith/tests/unit_tests/src/dify_trace_langsmith/test_provider.py b/api/providers/trace/trace-langsmith/tests/unit_tests/src/dify_trace_langsmith/test_provider.py new file mode 100644 index 00000000000000..10ed1402d3b1f7 --- /dev/null +++ b/api/providers/trace/trace-langsmith/tests/unit_tests/src/dify_trace_langsmith/test_provider.py @@ -0,0 +1,376 @@ +import threading + +import pytest +from dify_trace_langsmith.config import LangSmithConfig +from dify_trace_langsmith.provider import LangSmithProvider +from dify_trace_sdk.egress import EgressResponse +from dify_trace_sdk.errors import ( + EgressBlockedError, + EgressError, + ExportOutcomeUnknownError, + ExportRejectedError, + ExportRetryableError, + VerificationError, +) +from dify_trace_sdk.ids import span_uuid +from dify_trace_sdk.model import Cost +from dify_trace_sdk.provider import SendResult +from dify_trace_sdk.testing import ( + MESSAGE_ID, + WORKFLOW_RUN_ID, + RecordingEgress, + message_trace, + response, + tool_fragment, + workflow_trace, +) + +API_KEY = "lsv2_pt_test-key" +TENANT_ID = "4f1f0b4e-6a51-4c5a-9a8e-2d7f0c3b9e11" +PROJECT_ID = "0d9b8f61-2c3e-4b7a-8f5d-6e1a9c4b7d22" +OTHER_PROJECT_ID = "a7c3e9d1-5b2f-4e8a-9c6d-3f1b8e2a4c33" +CONFIG = LangSmithConfig(api_key=API_KEY, project="dify-app") +PROJECT_URL = f"https://smith.langchain.com/o/{TENANT_ID}/projects/p/{PROJECT_ID}" +VERIFIED = CONFIG.model_copy(update={"project_url": PROJECT_URL}) + + +def _provider(*responses: EgressResponse | Exception) -> tuple[LangSmithProvider, RecordingEgress]: + egress = RecordingEgress(responses, grpc_target=None) + return LangSmithProvider(egress), egress + + +def _accepted(count: int) -> list[EgressResponse]: + return [response(202, {}) for _ in range(count)] + + +def _session(session_id: str, name: str) -> dict[str, str]: + return {"id": session_id, "tenant_id": TENANT_ID, "name": name} + + +def _runs(egress: RecordingEgress) -> list[dict]: + """The runs the SDK posted, decoded from the recorded request bodies.""" + runs = [] + for sent in egress.requests: + assert (sent.method, sent.url) == ("POST", "https://api.smith.langchain.com/runs") + assert sent.headers["x-api-key"] == API_KEY + assert sent.headers["Content-Type"] == "application/json" + body = sent.json() + assert isinstance(body, dict) + runs.append(body) + return runs + + +def test_verify_resolves_the_project_page_with_read_only_requests() -> None: + provider, egress = _provider( + response(200, [_session(OTHER_PROJECT_ID, "another-app")]), + response(200, [_session(PROJECT_ID, "dify-app")]), + ) + + verified = provider.verify(CONFIG) + + assert verified.project_url == PROJECT_URL + tenant_lookup, project_lookup = egress.requests + assert tenant_lookup.url == "https://api.smith.langchain.com/sessions?limit=1" + assert project_lookup.url == "https://api.smith.langchain.com/sessions?limit=1&name=dify-app&include_stats=false" + for sent in egress.requests: + assert sent.method == "GET" + assert sent.content is None + assert sent.headers["x-api-key"] == API_KEY + assert not sent.sensitive_url + + +def test_verify_links_the_project_list_until_the_first_run_creates_the_project() -> None: + provider, _ = _provider(response(200, [_session(OTHER_PROJECT_ID, "another-app")]), response(200, [])) + assert provider.verify(CONFIG).project_url == f"https://smith.langchain.com/o/{TENANT_ID}/projects" + + +def test_verify_of_a_workspace_without_projects_links_the_web_app() -> None: + provider, _ = _provider(response(200, []), response(200, [])) + assert provider.verify(CONFIG).project_url == "https://smith.langchain.com" + + +@pytest.mark.parametrize( + ("endpoint", "web_url"), + [ + ("https://langsmith.example.com/api/v1", "https://langsmith.example.com"), + ("https://langsmith.example.com/team/api/", "https://langsmith.example.com/team"), + ("https://eu.api.smith.langchain.com", "https://eu.smith.langchain.com"), + ("https://localhost:1984", "http://localhost"), + ], +) +def test_verify_derives_the_web_app_like_the_sdk(endpoint: str, web_url: str) -> None: + provider, egress = _provider( + response(200, [_session(PROJECT_ID, "dify-app")]), response(200, [_session(PROJECT_ID, "dify-app")]) + ) + + verified = provider.verify(LangSmithConfig(api_key=API_KEY, project="dify-app", endpoint=endpoint)) + + assert verified.project_url == f"{web_url}/o/{TENANT_ID}/projects/p/{PROJECT_ID}" + assert egress.requests[0].url == f"{endpoint.rstrip('/')}/sessions?limit=1" + + +@pytest.mark.parametrize( + ("outcomes", "code"), + [ + ([response(401, {"detail": "Invalid token"})], "authentication_failed"), + ([response(403, {"error": "org_scoped_key_requires_workspace"})], "authentication_failed"), + ([response(404, {"detail": "Not Found"})], "endpoint_unavailable"), + ([response(503, {})], "endpoint_unavailable"), + ([response(200, b"")], "malformed_response"), + ([response(200, {"detail": "not a list"})], "malformed_response"), + ([response(200, [{"id": PROJECT_ID, "tenant_id": "tenant-1"}])], "malformed_response"), + ([response(200, [_session(OTHER_PROJECT_ID, "another-app")]), response(500, {})], "endpoint_unavailable"), + ( + [response(200, [_session(OTHER_PROJECT_ID, "another-app")]), response(200, [{"name": "dify-app"}])], + "malformed_response", + ), + ([EgressError("connection refused", request_sent=False)], "connection_failed"), + ([EgressBlockedError("private address")], "connection_failed"), + ], +) +def test_verify_failures_have_stable_codes(outcomes: list[EgressResponse | Exception], code: str) -> None: + provider, _ = _provider(*outcomes) + with pytest.raises(VerificationError) as caught: + provider.verify(CONFIG) + assert caught.value.code == code + + +def test_project_url_returns_the_stored_link_without_io() -> None: + provider, egress = _provider() + assert provider.project_url(VERIFIED) == PROJECT_URL + assert provider.project_url(CONFIG) is None + assert egress.requests == [] + + +def test_message_trace_is_posted_as_a_root_run_and_an_llm_run() -> None: + provider, egress = _provider(*_accepted(2)) + trace = message_trace() + message, llm = trace.spans + llm_run_id = str(span_uuid(trace.trace_id, llm.span_id)) + root_order = f"20261001T080000000000Z{MESSAGE_ID}" + + result = provider.send(VERIFIED, trace) + + assert _runs(egress) == [ + { + "id": MESSAGE_ID, + "trace_id": MESSAGE_ID, + "parent_run_id": None, + "dotted_order": root_order, + "session_name": "dify-app", + "name": "message", + "run_type": "chain", + "start_time": "2026-10-01T08:00:00+00:00", + "end_time": "2026-10-01T08:00:02+00:00", + "inputs": {"query": "What is Dify?"}, + "outputs": {"output": "An LLM application platform."}, + "error": None, + "extra": { + "metadata": { + "message_id": MESSAGE_ID, + "app_id": "app-1", + "tenant_id": "tenant-1", + "dify_trace_id": trace.trace_id, + "session_id": "conversation-1", + } + }, + }, + { + "id": llm_run_id, + "trace_id": MESSAGE_ID, + "parent_run_id": MESSAGE_ID, + "dotted_order": f"{root_order}.20261001T080000000000Z{llm_run_id}", + "session_name": "dify-app", + "name": "llm", + "run_type": "llm", + "start_time": "2026-10-01T08:00:00+00:00", + "end_time": "2026-10-01T08:00:02+00:00", + "inputs": {"input": [{"role": "user", "content": "What is Dify?"}]}, + "outputs": { + "output": "An LLM application platform.", + "usage_metadata": {"input_tokens": 12, "output_tokens": 8, "total_tokens": 20, "total_cost": 0.0004}, + }, + "error": None, + "extra": { + "metadata": { + "message_id": MESSAGE_ID, + "dify_trace_id": trace.trace_id, + "session_id": "conversation-1", + "ls_model_name": "gpt-4o", + "ls_provider": "openai", + } + }, + }, + ] + assert result == SendResult( + trace_id=MESSAGE_ID, + trace_url=f"{PROJECT_URL}/r/{MESSAGE_ID}", + span_urls={message.span_id: f"{PROJECT_URL}/r/{MESSAGE_ID}", llm.span_id: f"{PROJECT_URL}/r/{llm_run_id}"}, + ) + + +def test_workflow_runs_keep_the_hierarchy_and_failures() -> None: + provider, egress = _provider(*_accepted(3)) + trace = workflow_trace() + _, llm, tool = trace.spans + llm_run_id = str(span_uuid(trace.trace_id, llm.span_id)) + tool_run_id = str(span_uuid(trace.trace_id, tool.span_id)) + + provider.send(VERIFIED, trace) + + runs = _runs(egress) + workflow_run, llm_run, tool_run = runs + assert [run["id"] for run in runs] == [WORKFLOW_RUN_ID, llm_run_id, tool_run_id] + assert [run["parent_run_id"] for run in runs] == [None, WORKFLOW_RUN_ID, llm_run_id] + assert {run["trace_id"] for run in runs} == {WORKFLOW_RUN_ID} + assert tool_run["dotted_order"] == ( + f"20261001T080000000000Z{WORKFLOW_RUN_ID}" + f".20261001T080001000000Z{llm_run_id}" + f".20261001T080003000000Z{tool_run_id}" + ) + assert [run["run_type"] for run in runs] == ["chain", "llm", "tool"] + assert [run["error"] for run in runs] == ["Tool failed", None, "Tool failed"] + assert llm_run["outputs"] == { + "text": "Tracing records execution.", + "usage_metadata": {"input_tokens": 30, "output_tokens": 10, "total_tokens": 40}, + } + assert tool_run["outputs"] == {} + assert workflow_run["extra"]["metadata"]["session_id"] == WORKFLOW_RUN_ID + + +def test_llm_cost_is_sent_only_in_usd() -> None: + provider, egress = _provider(*_accepted(2)) + trace = message_trace() + message, llm = trace.spans + priced_in_rmb = llm.model_copy(update={"cost": Cost(amount=0.003, currency="RMB")}) + + provider.send(VERIFIED, trace.model_copy(update={"spans": (message, priced_in_rmb)})) + + assert _runs(egress)[1]["outputs"]["usage_metadata"] == {"input_tokens": 12, "output_tokens": 8, "total_tokens": 20} + + +def test_external_trace_id_is_kept_in_metadata() -> None: + provider, egress = _provider(*_accepted(2)) + provider.send(VERIFIED, message_trace().model_copy(update={"external_trace_id": "caller-trace-1"})) + assert [run["extra"]["metadata"]["external_trace_id"] for run in _runs(egress)] == ["caller-trace-1"] * 2 + + +def test_fragment_attaches_under_the_message_run() -> None: + message_provider, message_egress = _provider(*_accepted(2)) + message_provider.send(VERIFIED, message_trace()) + message_order = _runs(message_egress)[0]["dotted_order"] + provider, egress = _provider(*_accepted(1)) + fragment = tool_fragment() + tool_run_id = str(span_uuid(fragment.trace_id, fragment.root.span_id)) + + result = provider.send(VERIFIED, fragment) + + (tool_run,) = _runs(egress) + assert tool_run["id"] == tool_run_id + assert tool_run["parent_run_id"] == MESSAGE_ID + assert tool_run["trace_id"] == MESSAGE_ID + assert tool_run["dotted_order"] == f"{message_order}.20261001T080000500000Z{tool_run_id}" + assert result == SendResult( + trace_id=MESSAGE_ID, + trace_url=f"{PROJECT_URL}/r/{MESSAGE_ID}", + span_urls={fragment.root.span_id: f"{PROJECT_URL}/r/{tool_run_id}"}, + ) + + +def test_runs_are_posted_under_the_endpoint_path() -> None: + provider, egress = _provider(*_accepted(2)) + config = LangSmithConfig(api_key=API_KEY, project="dify-app", endpoint="https://langsmith.internal/api/v1/") + provider.send(config, message_trace()) + assert [sent.url for sent in egress.requests] == ["https://langsmith.internal/api/v1/runs"] * 2 + + +def test_redelivery_posts_identical_runs() -> None: + first, first_egress = _provider(*_accepted(3)) + second, second_egress = _provider(*_accepted(3)) + first.send(VERIFIED, workflow_trace()) + second.send(VERIFIED, workflow_trace()) + assert [sent.content for sent in first_egress.requests] == [sent.content for sent in second_egress.requests] + + +def test_runs_langsmith_already_has_count_as_accepted() -> None: + provider, egress = _provider(response(409, {"detail": "Run already exists"}), *_accepted(1)) + assert provider.send(VERIFIED, message_trace()).trace_id == MESSAGE_ID + assert len(egress.requests) == 2 + + +@pytest.mark.parametrize( + "project_url", + [None, f"https://smith.langchain.com/o/{TENANT_ID}/projects", "https://smith.langchain.com"], +) +def test_runs_are_linked_only_once_the_project_page_is_known(project_url: str | None) -> None: + provider, _ = _provider(*_accepted(2)) + result = provider.send(CONFIG.model_copy(update={"project_url": project_url}), message_trace()) + assert result == SendResult(trace_id=MESSAGE_ID, trace_url=None, span_urls={}) + + +def test_send_leaves_no_background_threads() -> None: + before = set(threading.enumerate()) + provider, _ = _provider(*_accepted(2)) + provider.send(VERIFIED, message_trace()) + assert set(threading.enumerate()) == before + + +@pytest.mark.parametrize( + ("outcome", "expected"), + [ + (response(429, {}), ExportRetryableError), + (response(500, {}), ExportRetryableError), + (response(503, {}), ExportRetryableError), + (response(400, {"detail": "invalid dotted_order"}), ExportRejectedError), + (response(401, {"detail": "Invalid token"}), ExportRejectedError), + (EgressError("connection refused", request_sent=False), ExportRetryableError), + (EgressError("read timeout", request_sent=True), ExportOutcomeUnknownError), + (EgressBlockedError("private address"), ExportRejectedError), + ], +) +def test_send_failures_are_classified(outcome: EgressResponse | Exception, expected: type[Exception]) -> None: + provider, egress = _provider(outcome) + with pytest.raises(expected) as caught: + provider.send(VERIFIED, message_trace()) + assert len(egress.requests) == 1 + assert caught.value.__cause__ is None + assert caught.value.__suppress_context__ + + +@pytest.mark.parametrize( + "failure", + [response(500, {}), EgressError("connection refused", request_sent=False), EgressError("reset", request_sent=True)], +) +def test_partial_acceptance_is_resent_because_runs_are_idempotent(failure: EgressResponse | Exception) -> None: + provider, _ = _provider(response(202, {}), failure) + with pytest.raises(ExportRetryableError): + provider.send(VERIFIED, message_trace()) + + +def test_partial_acceptance_of_a_rejected_trace_is_not_resent() -> None: + provider, _ = _provider(response(202, {}), response(422, {})) + with pytest.raises(ExportRejectedError): + provider.send(VERIFIED, message_trace()) + + +def test_runs_that_share_a_caller_trace_id_keep_their_own_root_runs() -> None: + first, second = ( + message_trace().model_copy(update={"external_trace_id": "caller-trace"}), + workflow_trace().model_copy(update={"external_trace_id": "caller-trace"}), + ) + provider, egress = _provider(*_accepted(len(first.spans) + len(second.spans))) + + provider.send(VERIFIED, first) + provider.send(VERIFIED, second) + + roots = [run for run in _runs(egress) if run["parent_run_id"] is None] + assert len(roots) == 2 + assert len({run["id"] for run in roots}) == 2 + assert all(run["trace_id"] == run["id"] for run in roots) + + +def test_failures_never_echo_credentials_urls_or_bodies() -> None: + provider, _ = _provider(response(422, {"detail": "rejected-body-detail"})) + with pytest.raises(ExportRejectedError) as caught: + provider.send(VERIFIED, message_trace()) + assert str(caught.value) == "LangSmith ingestion rejected (HTTP 422)" diff --git a/api/providers/trace/trace-mlflow/pyproject.toml b/api/providers/trace/trace-mlflow/pyproject.toml index a72507d87716ab..7ed65107a83d75 100644 --- a/api/providers/trace/trace-mlflow/pyproject.toml +++ b/api/providers/trace/trace-mlflow/pyproject.toml @@ -2,9 +2,14 @@ name = "dify-trace-mlflow" version = "0.0.1" dependencies = [ + "dify-trace-sdk", "mlflow-skinny>=3.11.1,<4.0.0", ] description = "Dify ops tracing provider (MLflow / Databricks)." +[project.entry-points."dify.trace_providers"] +mlflow = "dify_trace_mlflow.provider:MLflowProvider" +databricks = "dify_trace_mlflow.provider:DatabricksProvider" + [tool.setuptools.packages.find] where = ["src"] diff --git a/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/_databricks.py b/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/_databricks.py new file mode 100644 index 00000000000000..cb1db11002e044 --- /dev/null +++ b/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/_databricks.py @@ -0,0 +1,178 @@ +"""Databricks experiment traces: the SDK's V3 trace API and trace artifact, sent through the egress.""" + +from collections.abc import Mapping +from typing import Any, NoReturn, override +from urllib.parse import quote, urlsplit + +from dify_trace_sdk.egress import Egress +from dify_trace_sdk.errors import ExportOutcomeUnknownError, ExportRejectedError +from dify_trace_sdk.exporting import json_value, nanos, send_export_request +from dify_trace_sdk.model import SpanStatus, Trace +from google.protobuf.json_format import ParseDict, ParseError +from google.protobuf.message import Message +from mlflow.entities import TraceInfo, TraceLocation, TraceState +from mlflow.protos.databricks_artifacts_pb2 import ( + ArtifactCredentialType, + DatabricksMlflowArtifactsService, + GetCredentialsForTraceDataUpload, +) +from mlflow.store.tracking.rest_store import RestStore +from mlflow.tracing.constant import ( + TRACE_REQUEST_RESPONSE_PREVIEW_MAX_LENGTH_DBX, + CostKey, + TokenUsageKey, + TraceMetadataKey, + TraceTagKey, +) +from mlflow.utils.rest_utils import extract_api_info_for_service + +from dify_trace_mlflow._spans import token_usage, usd_cost + +_DESTINATION = "Databricks" +_SIGNED_URL_CREDENTIALS = { + ArtifactCredentialType.AWS_PRESIGNED_URL, + ArtifactCredentialType.GCP_SIGNED_URL, + ArtifactCredentialType.AZURE_SAS_URI, +} + + +def _no_sdk_transport() -> NoReturn: + raise ExportRejectedError("The MLflow SDK must not open its own Databricks connection") + + +class DatabricksStore(RestStore): + """The SDK's V3 ``start_trace`` with its request sent through the egress. + + Only ``start_trace``, a POST, is used. The SDK never receives host + credentials, so any request that bypassed ``_call_endpoint`` would fail + instead of resolving credentials from the process. The SDK's retry timeout + is ignored: each call is one request, and the trace pipeline owns retries. + """ + + def __init__(self, egress: Egress, host: str, headers: Mapping[str, str]) -> None: + super().__init__(_no_sdk_transport) + self._egress = egress + self._host = host + self._headers = headers + + @override + def _call_endpoint( + self, + api: Any, + json_body: str | None = None, + endpoint: str | None = None, + retry_timeout_seconds: int | None = None, + response_proto: Message | None = None, + ) -> Message: + methods = self._V3_METHOD_TO_INFO if api in self._V3_APIS else self._METHOD_TO_INFO + default_endpoint, method = methods[api] + response = send_export_request( + self._egress, + method, + self._host + (endpoint or default_endpoint), + destination=_DESTINATION, + headers={"Content-Type": "application/json", "Accept": "application/json", **self._headers}, + content=None if json_body is None else json_body.encode(), + sensitive_url=False, + ) + try: + body = response.json() + if not isinstance(body, dict): + raise ValueError("acknowledgment is not an object") + return ParseDict( + body, response_proto if response_proto is not None else api.Response(), ignore_unknown_fields=True + ) + except (ParseError, TypeError, ValueError): + raise ExportOutcomeUnknownError("Databricks returned an unreadable trace acknowledgment") from None + + +def _preview(value: object) -> str | None: + return None if value is None else json_value(value)[:TRACE_REQUEST_RESPONSE_PREVIEW_MAX_LENGTH_DBX] + + +def trace_info(trace: Trace, experiment_id: str, trace_id: str) -> TraceInfo: + """The trace-level record of a standalone trace: session, user, totals, previews, and name. + + Raises: + ValueError: The root's inputs or outputs cannot be encoded as JSON. + """ + root = trace.root + metadata: dict[str, str] = {} + if trace.session_id is not None: + metadata[TraceMetadataKey.TRACE_SESSION] = trace.session_id + if trace.user_id is not None: + metadata[TraceMetadataKey.TRACE_USER] = trace.user_id + usages = [usage for span in trace.spans if (usage := token_usage(span)) is not None] + if usages: + keys = (TokenUsageKey.INPUT_TOKENS, TokenUsageKey.OUTPUT_TOKENS, TokenUsageKey.TOTAL_TOKENS) + metadata[TraceMetadataKey.TOKEN_USAGE] = json_value({key: sum(usage[key] for usage in usages) for key in keys}) + costs = [cost for span in trace.spans if (cost := usd_cost(span)) is not None] + if costs: + metadata[TraceMetadataKey.COST] = json_value({CostKey.TOTAL_COST: sum(costs)}) + return TraceInfo( + trace_id=trace_id, + trace_location=TraceLocation.from_experiment_id(experiment_id), + request_time=nanos(root.start_time) // 1_000_000, + execution_duration=(nanos(root.end_time) - nanos(root.start_time)) // 1_000_000, + state=TraceState.ERROR if root.status is SpanStatus.ERROR else TraceState.OK, + request_preview=_preview(root.inputs), + response_preview=_preview(root.outputs), + client_request_id=trace.trace_id, + trace_metadata=metadata, + tags={TraceTagKey.TRACE_NAME: root.name}, + ) + + +def _upload(egress: Egress, method: str, url: str, headers: Mapping[str, str], content: bytes | None) -> None: + # Signed URLs embed a storage credential and receive no Databricks token. + send_export_request( + egress, method, url, destination=_DESTINATION, headers=headers, content=content, sensitive_url=True + ) + + +def upload_trace_data(egress: Egress, host: str, headers: Mapping[str, str], trace_id: str, payload: bytes) -> None: + """Store the spans as the trace's artifact through a signed cloud storage URL. + + Raises: + TraceProviderError: Any failure; the caller decides the outcome. + """ + path, method = extract_api_info_for_service(DatabricksMlflowArtifactsService, "/api/2.0")[ + GetCredentialsForTraceDataUpload + ] + response = send_export_request( + egress, + method, + host + path.replace("{request_id}", quote(trace_id, safe="")), + destination=_DESTINATION, + headers={"Accept": "application/json", **headers}, + content=None, + sensitive_url=False, + ) + try: + body = response.json() + if not isinstance(body, dict): + raise ValueError("credentials are not an object") + credential = ParseDict( + body, GetCredentialsForTraceDataUpload.Response(), ignore_unknown_fields=True + ).credential_info + except (ParseError, TypeError, ValueError): + raise ExportOutcomeUnknownError("Databricks returned unreadable artifact credentials") from None + signed_url = credential.signed_uri + parsed = urlsplit(signed_url) + if parsed.scheme not in {"http", "https"} or not parsed.hostname or parsed.username or parsed.password: + raise ExportRejectedError("Invalid Databricks artifact upload destination") + upload_headers = {header.name: header.value for header in credential.headers} + # An unknown type is dropped while parsing and would read as the enum's first value. + credential_type = credential.type if credential.HasField("type") else None + if credential_type == ArtifactCredentialType.AZURE_ADLS_GEN2_SAS_URI: + # The whole artifact fits one append; flushing commits it. + separator = "&" if parsed.query else "?" + adls_headers = {key: value for key, value in upload_headers.items() if key.lower() != "x-ms-blob-type"} + _upload(egress, "PUT", f"{signed_url}{separator}resource=file", adls_headers, None) + _upload(egress, "PATCH", f"{signed_url}{separator}action=append&position=0&flush=true", adls_headers, payload) + elif credential_type in _SIGNED_URL_CREDENTIALS: + if credential_type == ArtifactCredentialType.AZURE_SAS_URI: + upload_headers["x-ms-blob-type"] = "BlockBlob" + _upload(egress, "PUT", signed_url, upload_headers, payload) + else: + raise ExportRejectedError("Unsupported Databricks artifact credential type") diff --git a/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/_spans.py b/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/_spans.py new file mode 100644 index 00000000000000..aee92de2426076 --- /dev/null +++ b/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/_spans.py @@ -0,0 +1,203 @@ +"""Dify spans as MLflow span entities, serialized by the MLflow SDK itself. + +Building immutable ``mlflow.entities.Span`` objects reuses MLflow's attribute +encoding and its OTLP and trace-artifact serialization without a tracer, trace +manager, or exporter, so no MLflow or OpenTelemetry pipeline state is involved. +""" + +import json + +from dify_trace_sdk.exporting import json_value, nanos +from dify_trace_sdk.ids import trace_id_for +from dify_trace_sdk.model import Span, SpanKind, SpanStatus, Trace +from mlflow.entities import Span as MlflowSpan +from mlflow.entities import TraceData +from mlflow.tracing.constant import CostKey, GenAiSemconvKey, SpanAttributeKey, TokenUsageKey +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import Event, ReadableSpan +from opentelemetry.trace import SpanContext, Status, StatusCode + +PARENT_TRACE_ID_ATTRIBUTE = "dify.parent_trace_id" +PARENT_SPAN_ID_ATTRIBUTE = "dify.parent_span_id" +EXTERNAL_TRACE_ID_ATTRIBUTE = "dify.external_trace_id" + + +def standalone(trace: Trace) -> Trace: + """Deliver a fragment as a trace of its own. + + MLflow cannot safely append spans to a trace it already stored, so a + fragment gets a deterministic trace ID and its root records where it belongs. + A caller's trace ID may group several runs, and MLflow keeps one root per + trace, so such a run gets its own trace ID and keeps the caller's as an attribute. + """ + root = trace.root + if trace.parent is None: + if trace.external_trace_id is None: + return trace + tagged = root.model_copy( + update={"attributes": {**root.attributes, EXTERNAL_TRACE_ID_ATTRIBUTE: trace.external_trace_id}} + ) + return trace.model_copy( + update={"trace_id": trace_id_for(f"{trace.trace_id}:{root.span_id}"), "spans": (tagged, *trace.spans[1:])} + ) + detached = root.model_copy( + update={ + "parent_span_id": None, + "attributes": { + **root.attributes, + PARENT_TRACE_ID_ATTRIBUTE: trace.trace_id, + PARENT_SPAN_ID_ATTRIBUTE: trace.parent.span_id, + }, + } + ) + return trace.model_copy( + update={ + "trace_id": trace_id_for(f"{trace.trace_id}:{root.span_id}"), + "parent": None, + "spans": (detached, *trace.spans[1:]), + } + ) + + +def token_usage(span: Span) -> dict[str, int] | None: + """Usage of a model call. MLflow sums span usage into the trace total, so wrappers report none.""" + if span.kind is not SpanKind.LLM or span.usage is None: + return None + return { + TokenUsageKey.INPUT_TOKENS: span.usage.prompt_tokens, + TokenUsageKey.OUTPUT_TOKENS: span.usage.completion_tokens, + TokenUsageKey.TOTAL_TOKENS: span.usage.total_tokens, + } + + +def usd_cost(span: Span) -> float | None: + """MLflow costs are USD totals of model calls.""" + if span.kind is not SpanKind.LLM or span.cost is None or span.cost.currency.upper() != "USD": + return None + return span.cost.amount + + +def _chat_tools(value: object) -> list[dict[str, object]] | None: + """Map OpenTelemetry's flat function definitions to MLflow's chat tool schema. + + The source attribute is kept as is. Setting the native attribute stops + MLflow's OTLP translator from copying the incompatible flat shape. Entries + that are not named functions are left out instead of becoming fake tools, + and JSON Schemas are kept whole although MLflow's ChatTool model accepts + only a subset of them. + """ + if isinstance(value, str): + try: + value = json.loads(value) + except ValueError: + return None + if not isinstance(value, list): + return None + tools: list[dict[str, object]] = [] + for item in value: + if not isinstance(item, dict) or item.get("type") != "function": + continue + function = item.get("function") + if not isinstance(function, dict): + function = {key: item[key] for key in ("name", "description", "parameters", "strict") if key in item} + name = function.get("name") + parameters = function.get("parameters") + if not isinstance(name, str) or not name.strip(): + continue + if parameters is not None and not isinstance(parameters, dict): + continue + tools.append({"type": "function", "function": function}) + return tools + + +def _attributes(trace: Trace, span: Span, experiment_id: str, request_id: str) -> dict[str, object]: + attributes: dict[str, object] = { + **span.attributes, + SpanAttributeKey.REQUEST_ID: request_id, + SpanAttributeKey.EXPERIMENT_ID: experiment_id, + SpanAttributeKey.SPAN_TYPE: span.kind.value.upper(), + } + if span.parent_span_id is None: + # MLflow's own exporters carry trace-level user and session on the root span. + if trace.session_id is not None: + attributes[SpanAttributeKey.SESSION_ID] = trace.session_id + if trace.user_id is not None: + attributes[SpanAttributeKey.USER_ID] = trace.user_id + if span.inputs is not None: + attributes[SpanAttributeKey.INPUTS] = span.inputs + if span.outputs is not None: + attributes[SpanAttributeKey.OUTPUTS] = span.outputs + if SpanAttributeKey.CHAT_TOOLS not in attributes: + tools = _chat_tools(span.attributes.get(GenAiSemconvKey.TOOL_DEFINITIONS)) + if tools is not None: + attributes[SpanAttributeKey.CHAT_TOOLS] = tools + if (usage := token_usage(span)) is not None: + attributes[SpanAttributeKey.CHAT_USAGE] = usage + if (cost := usd_cost(span)) is not None: + attributes[SpanAttributeKey.LLM_COST] = {CostKey.TOTAL_COST: cost} + if span.model_name: + attributes[SpanAttributeKey.MODEL] = span.model_name + if span.model_provider: + attributes[SpanAttributeKey.MODEL_PROVIDER] = span.model_provider + return attributes + + +def mlflow_spans(trace: Trace, experiment_id: str, request_id: str) -> list[MlflowSpan]: + """Spans of a standalone trace; ``request_id`` is MLflow's identifier of the trace. + + Raises: + ValueError: A value cannot be encoded as JSON (e.g. NaN). + """ + otel_trace_id = int(trace.trace_id, 16) + spans: list[MlflowSpan] = [] + for span in trace.spans: + end_time = nanos(span.end_time) + parent = ( + None + if span.parent_span_id is None + else SpanContext(trace_id=otel_trace_id, span_id=int(span.parent_span_id, 16), is_remote=False) + ) + status = Status(StatusCode.OK) + events: list[Event] = [] + if span.status is SpanStatus.ERROR: + status = Status(StatusCode.ERROR, span.error) + if span.error: + # MLflow's tracer records errors as exception events, which its trace UI displays. + attributes = {"exception.type": "Error", "exception.message": span.error} + events.append(Event(name="exception", attributes=attributes, timestamp=end_time)) + spans.append( + MlflowSpan( + ReadableSpan( + name=span.name, + context=SpanContext(trace_id=otel_trace_id, span_id=int(span.span_id, 16), is_remote=False), + parent=parent, + start_time=nanos(span.start_time), + end_time=end_time, + attributes={ + key: json_value(value) + for key, value in _attributes(trace, span, experiment_id, request_id).items() + }, + status=status, + events=events, + resource=Resource.get_empty(), + ) + ) + ) + return spans + + +def otlp_request(spans: list[MlflowSpan]) -> ExportTraceServiceRequest: + """The OTLP request MLflow's REST store sends for these spans. + + Raises: + ValueError: An attribute does not fit its OTLP type (e.g. an integer beyond 64 bits). + """ + request = ExportTraceServiceRequest() + request.resource_spans.add().scope_spans.add().spans.extend(span.to_otel_proto() for span in spans) + return request + + +def trace_data(spans: list[MlflowSpan]) -> bytes: + """The trace artifact MLflow's tracing client uploads for these spans.""" + return json_value(TraceData(spans=spans).to_dict()).encode() diff --git a/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/agent_trace.py b/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/agent_trace.py deleted file mode 100644 index 3bb26942e9b503..00000000000000 --- a/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/agent_trace.py +++ /dev/null @@ -1,259 +0,0 @@ -"""MLflow SDK span serialization with per-export credentials and bounded HTTP. - -OSS uses its OTLP experiment endpoint (MLflow >= 3.4). Databricks retains the -existing experiment destination through the SDK's V3 trace and artifact APIs; -its UC-only OTLP endpoint is not interchangeable with an experiment ID. -""" - -import base64 -import json -from collections.abc import Mapping -from typing import Any, override -from urllib.parse import quote, urlencode, urlsplit - -from google.protobuf.json_format import ParseDict, ParseError -from mlflow.entities import Span, TraceData, TraceInfo, TraceLocation, TraceState -from mlflow.protos.databricks_artifacts_pb2 import ( - ArtifactCredentialType, - DatabricksMlflowArtifactsService, - GetCredentialsForTraceDataUpload, -) -from mlflow.store.tracking.rest_store import RestStore -from mlflow.tracing.constant import SpanAttributeKey, TraceMetadataKey, TraceTagKey -from mlflow.tracing.utils.otlp import MLFLOW_EXPERIMENT_ID_HEADER, OTLP_TRACES_PATH -from mlflow.utils.rest_utils import MlflowHostCreds, extract_api_info_for_service -from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest -from opentelemetry.sdk.resources import Resource -from opentelemetry.sdk.trace import ReadableSpan -from opentelemetry.trace import SpanContext, Status, StatusCode - -from core.ops.unified_trace.entities import CanonicalSpanStatus, CanonicalTrace -from services.tracing.otlp_export import export_http -from services.tracing.provider_common import endpoint, json_value, nanos, request, required, token_usage, usd_cost -from services.tracing.provider_dispatch import TraceExportRejectedError, TraceExportResult, TraceExportUncertainError - - -def _chat_tools(value: object) -> list[dict[str, Any]] | None: - """Map OTel's flat function definitions to MLflow's chat tool UI schema. - - Preserve the original gen_ai attribute separately. An explicit native field - prevents MLflow's OTLP translator from copying the incompatible flat shape. - Invalid/unsupported entries stay in the source metadata, not in chat tools. - """ - if isinstance(value, str): - try: - value = json.loads(value) - except ValueError: - return None - if not isinstance(value, list): - return None - tools = [] - for item in value: - if not isinstance(item, dict) or item.get("type") != "function": - continue - function = item.get("function") - if not isinstance(function, dict): - function = {key: item[key] for key in ("name", "description", "parameters", "strict") if key in item} - name = function.get("name") - if not isinstance(name, str) or not name.strip(): - continue - tool = {"type": "function", "function": function} - parameters = function.get("parameters") - if parameters is not None and not isinstance(parameters, dict): - continue - # MLflow's Python ChatTool model accepts only a subset of JSON Schema - # (e.g. it requires properties and restricts enum values). Preserve the - # full schema for the UI instead of silently dropping valid definitions. - tools.append(tool) - return tools - - -def _spans(trace: CanonicalTrace, experiment_id: str, request_id: str) -> list[Span]: - result = [] - for span in trace.spans: - attributes: dict[str, Any] = { - **span.metadata, - SpanAttributeKey.REQUEST_ID: request_id, - SpanAttributeKey.EXPERIMENT_ID: experiment_id, - SpanAttributeKey.SPAN_TYPE: span.kind.value.upper(), - SpanAttributeKey.SESSION_ID: trace.session_id, - } - if span.inputs is not None: - attributes[SpanAttributeKey.INPUTS] = span.inputs - if span.outputs is not None: - attributes[SpanAttributeKey.OUTPUTS] = span.outputs - if SpanAttributeKey.CHAT_TOOLS not in attributes: - tools = _chat_tools(span.metadata.get("gen_ai.tool.definitions")) - if tools is not None: - attributes[SpanAttributeKey.CHAT_TOOLS] = tools - if usage := token_usage(span): - attributes[SpanAttributeKey.CHAT_USAGE] = {f"{key}_tokens": value for key, value in usage.items()} - if (cost := usd_cost(span)) is not None: - attributes[SpanAttributeKey.LLM_COST] = {"total_cost": cost} - if model := span.metadata.get("model_name") or span.metadata.get("model"): - attributes[SpanAttributeKey.MODEL] = model - context = SpanContext(trace_id=int(trace.trace_id, 16), span_id=int(span.id, 16), is_remote=False) - parent = ( - SpanContext(trace_id=context.trace_id, span_id=int(span.parent_id, 16), is_remote=False) - if span.parent_id - else None - ) - result.append( - Span( - ReadableSpan( - name=span.name, - context=context, - parent=parent, - start_time=nanos(span.start_time), - end_time=nanos(span.end_time) if span.end_time else None, - attributes={key: json_value(value) for key, value in attributes.items()}, - status=Status(StatusCode.ERROR, span.error) - if span.status == CanonicalSpanStatus.ERROR - else Status(StatusCode.OK), - resource=Resource.get_empty(), - ) - ) - ) - return result - - -def _auth(provider: str, config: Mapping[str, Any], host: str) -> dict[str, str]: - if provider == "mlflow": - if config.get("username") and config.get("password"): - token = base64.b64encode(f"{config['username']}:{config['password']}".encode()).decode() - return {"Authorization": f"Basic {token}"} - return {} - if config.get("client_id") and config.get("client_secret"): - response = request( - "POST", - host + "/oidc/v1/token", - provider, - auth=(config["client_id"], config["client_secret"]), - data={"grant_type": "client_credentials", "scope": "all-apis"}, - ) - token = response.json().get("access_token") - if not isinstance(token, str) or not token: - raise TraceExportRejectedError("Databricks OAuth token was not returned") - else: - token = required(config, "personal_access_token") - return {"Authorization": f"Bearer {token}"} - - -class _Store(RestStore): - """Reuse native V3 request/response schemas without process-wide SDK credentials.""" - - def __init__(self, host: str, headers: dict[str, str]) -> None: - super().__init__(lambda: MlflowHostCreds(host)) - self._host = host - self._headers = headers - - @override - def _call_endpoint(self, api, json_body=None, endpoint=None, retry_timeout_seconds=None, response_proto=None): - methods = self._V3_METHOD_TO_INFO if api in self._V3_APIS else self._METHOD_TO_INFO - default_endpoint, method = methods[api] - response = request( - method, - self._host + (endpoint or default_endpoint), - "databricks", - headers={"Content-Type": "application/json", **self._headers}, - content=json_body, - ) - try: - return ParseDict(response.json(), response_proto if response_proto is not None else api.Response()) - except (ValueError, ParseError): - raise TraceExportUncertainError("Databricks returned an unreadable trace acknowledgment") from None - - -def _upload_databricks(host: str, headers: dict[str, str], trace_id: str, payload: bytes) -> None: - path, method = extract_api_info_for_service(DatabricksMlflowArtifactsService, "/api/2.0")[ - GetCredentialsForTraceDataUpload - ] - response = request( - method, host + path.replace("{request_id}", quote(trace_id, safe="")), "databricks", headers=headers - ) - credential = ParseDict(response.json(), GetCredentialsForTraceDataUpload.Response()).credential_info - signed_url = credential.signed_uri - parsed = urlsplit(signed_url) - if parsed.scheme not in {"http", "https"} or not parsed.hostname or parsed.username or parsed.password: - raise TraceExportRejectedError("Invalid Databricks artifact upload destination") - upload_headers = {header.name: header.value for header in credential.headers} - if credential.type == ArtifactCredentialType.AZURE_ADLS_GEN2_SAS_URI: - # The bounded trace payload fits a single append. Flush commits that append. - separator = "&" if parsed.query else "?" - create_headers = {key: value for key, value in upload_headers.items() if key.lower() != "x-ms-blob-type"} - request("PUT", signed_url + separator + "resource=file", "databricks", headers=create_headers) - request( - "PATCH", - signed_url + separator + "action=append&position=0&flush=true", - "databricks", - headers=create_headers, - content=payload, - ) - elif credential.type in { - ArtifactCredentialType.AWS_PRESIGNED_URL, - ArtifactCredentialType.GCP_SIGNED_URL, - ArtifactCredentialType.AZURE_SAS_URI, - }: - if credential.type == ArtifactCredentialType.AZURE_SAS_URI: - upload_headers["x-ms-blob-type"] = "BlockBlob" - request("PUT", signed_url, "databricks", headers=upload_headers, content=payload) - else: - raise TraceExportRejectedError("Unsupported Databricks artifact credential type") - - -def export_trace(provider: str, config: Mapping[str, Any], trace: CanonicalTrace) -> TraceExportResult: - host = endpoint(required(config, "host" if provider == "databricks" else "tracking_uri")) - experiment = str(config.get("experiment_id") or "0") - if not experiment.isdigit(): - raise TraceExportRejectedError("MLflow experiment ID must be numeric") - headers = _auth(provider, config, host) - trace_id = "tr-" + trace.trace_id - spans = _spans(trace, experiment, trace_id) - if provider == "mlflow": - payload = ExportTraceServiceRequest() - scope = payload.resource_spans.add().scope_spans.add() - scope.spans.extend(span.to_otel_proto() for span in spans) - export_http(provider, host + OTLP_TRACES_PATH, {MLFLOW_EXPERIMENT_ID_HEADER: experiment, **headers}, payload) - # MLflow 3.4 has classic (selectedTraceId) and GenAI-v3 - # (selectedEvaluationId) views; 3.17's v4 reads traceId. On close/change, - # v4 removes its selectedEvaluationId fallback before updating traceId. - # Each older view reads only its own key and ignores the other aliases. - query = urlencode({"traceId": trace_id, "selectedTraceId": trace_id, "selectedEvaluationId": trace_id}) - url = f"{host}/#/experiments/{experiment}/traces?{query}" - else: - root = next(span for span in trace.spans if span.id == trace.root_span_id) - metadata = {TraceMetadataKey.TRACE_SESSION: trace.session_id} - usages = [token_usage(span) for span in trace.spans] - if any(usages): - metadata[TraceMetadataKey.TOKEN_USAGE] = json_value( - {f"{key}_tokens": sum(usage.get(key, 0) for usage in usages) for key in ("input", "output", "total")} - ) - costs = [cost for span in trace.spans if (cost := usd_cost(span)) is not None] - if costs: - metadata[TraceMetadataKey.COST] = json_value({"total_cost": sum(costs)}) - info = TraceInfo( - trace_id=trace_id, - trace_location=TraceLocation.from_experiment_id(experiment), - request_time=nanos(root.start_time) // 1_000_000, - execution_duration=(nanos(root.end_time) - nanos(root.start_time)) // 1_000_000 if root.end_time else None, - state=TraceState.ERROR if root.status == CanonicalSpanStatus.ERROR else TraceState.OK, - request_preview=json_value(root.inputs)[:10000] if root.inputs is not None else None, - response_preview=json_value(root.outputs)[:10000] if root.outputs is not None else None, - trace_metadata=metadata, - tags={TraceTagKey.TRACE_NAME: root.name}, - client_request_id=trace.trace_id, - ) - stored = _Store(host, headers).start_trace(info) - try: - trace_id = stored.trace_id - if not trace_id: - raise TraceExportUncertainError("Databricks did not acknowledge a trace identifier") - data = TraceData(spans=_spans(trace, experiment, trace_id)) - _upload_databricks(host, headers, trace_id, json_value(data.to_dict()).encode()) - except Exception: - raise TraceExportUncertainError( - "Databricks accepted trace metadata but artifact delivery is incomplete" - ) from None - # Databricks opens its trace detail panel with a different key from OSS MLflow. - url = f"{host}/ml/experiments/{experiment}/traces?selectedEvaluationId={quote(trace_id, safe='')}" - return TraceExportResult(trace_id=trace_id, trace_url=url) diff --git a/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/config.py b/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/config.py index 84914165e3b247..6815d2efa3eb48 100644 --- a/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/config.py +++ b/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/config.py @@ -1,46 +1,38 @@ -from pydantic import ValidationInfo, field_validator +from dify_trace_sdk.config import ProviderConfig, destination, http_url, integer_id, secret +from pydantic import field_validator -from core.ops.entities.config_entity import BaseTracingConfig -from core.ops.utils import validate_integer_id, validate_url_with_path +DEFAULT_TRACKING_URI = "http://localhost:5000" -class MLflowConfig(BaseTracingConfig): - """ - Model class for MLflow tracing config. - """ - - tracking_uri: str = "http://localhost:5000" - experiment_id: str = "0" # Default experiment id in MLflow is 0 +class MLflowConfig(ProviderConfig): + tracking_uri: str = destination(default=DEFAULT_TRACKING_URI) + experiment_id: str = "0" username: str | None = None - password: str | None = None + password: str | None = secret(default=None) @field_validator("tracking_uri") @classmethod - def tracking_uri_validator(cls, v, info: ValidationInfo): - if isinstance(v, str) and v.startswith("databricks"): - raise ValueError( - "Please use Databricks tracing config below to record traces to Databricks-managed MLflow instances." - ) - return validate_url_with_path(v, "http://localhost:5000") + def _tracking_uri(cls, value: str) -> str: + if not value.strip(): + return DEFAULT_TRACKING_URI + if value.strip().startswith("databricks"): + raise ValueError("Use the Databricks provider for Databricks-managed MLflow") + return http_url(value, allowed_schemes=("https", "http")) @field_validator("experiment_id") @classmethod - def experiment_id_validator(cls, v, info: ValidationInfo): - return validate_integer_id(v) - + def _experiment_id(cls, value: str) -> str: + return integer_id(value) -class DatabricksConfig(BaseTracingConfig): - """ - Model class for Databricks (Databricks-managed MLflow) tracing config. - """ +class DatabricksConfig(ProviderConfig): experiment_id: str - host: str + host: str = destination() client_id: str | None = None - client_secret: str | None = None - personal_access_token: str | None = None + client_secret: str | None = secret(default=None) + personal_access_token: str | None = secret(default=None) @field_validator("experiment_id") @classmethod - def experiment_id_validator(cls, v, info: ValidationInfo): - return validate_integer_id(v) + def _experiment_id(cls, value: str) -> str: + return integer_id(value) diff --git a/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/configuration.py b/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/configuration.py deleted file mode 100644 index 326930318d2ddb..00000000000000 --- a/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/configuration.py +++ /dev/null @@ -1,163 +0,0 @@ -"""Read Databricks experiment metadata without changing MLflow's global state. - -This verifies authentication and read access to an existing active experiment. -It does not prove write permission, tracing API availability, or artifact delivery. -""" - -from dataclasses import dataclass -from typing import Literal -from urllib.parse import urlsplit, urlunsplit - -from google.protobuf.json_format import ParseDict, ParseError -from mlflow.protos.service_pb2 import GetExperiment -from pydantic import TypeAdapter - -from core.helper import ssrf_proxy -from dify_trace_mlflow.config import DatabricksConfig - -type VerificationCode = Literal[ - "invalid_workspace_url", - "invalid_experiment_id", - "authentication_failed", - "permission_denied", - "experiment_not_found_or_inaccessible", - "experiment_not_active", - "experiment_id_mismatch", - "malformed_response", - "endpoint_unavailable", - "connection_failed", -] - -_MESSAGES: dict[VerificationCode, str] = { - "invalid_workspace_url": "Databricks requires an HTTPS workspace origin.", - "invalid_experiment_id": "Databricks requires a numeric experiment ID.", - "authentication_failed": "Databricks authentication failed.", - "permission_denied": "Databricks experiment access was denied.", - "experiment_not_found_or_inaccessible": "The selected Databricks experiment was not found or is inaccessible.", - "experiment_not_active": "The selected Databricks experiment is not active.", - "experiment_id_mismatch": "Databricks returned a different experiment.", - "malformed_response": "Databricks verification returned an invalid response.", - "endpoint_unavailable": "Databricks endpoint could not be verified.", - "connection_failed": "Databricks verification request failed.", -} -_JSON_OBJECT = TypeAdapter(dict[str, object]) - - -class DatabricksConfigurationError(ValueError): - """Fixed safe diagnostics; never include credentials or a vendor response.""" - - def __init__(self, code: VerificationCode) -> None: - self.code = code - super().__init__(_MESSAGES[code]) - - -@dataclass(frozen=True, slots=True) -class ReadableDatabricksExperiment: - host: str - experiment_id: str - - -def _target(config: DatabricksConfig) -> ReadableDatabricksExperiment: - host = config.host.strip() - try: - parsed = urlsplit(host) - valid = ( - parsed.scheme == "https" - and parsed.hostname - and parsed.username is None - and parsed.password is None - and parsed.path in {"", "/"} - and not parsed.query - and not parsed.fragment - and "\\" not in host - and not any(ord(char) <= 32 or ord(char) == 127 for char in host) - ) - _ = parsed.port - except ValueError: - raise DatabricksConfigurationError("invalid_workspace_url") from None - if not valid: - raise DatabricksConfigurationError("invalid_workspace_url") - experiment_id = config.experiment_id.strip() - if not experiment_id or not experiment_id.isascii() or not experiment_id.isdigit(): - raise DatabricksConfigurationError("invalid_experiment_id") - return ReadableDatabricksExperiment( - host=urlunsplit((parsed.scheme, parsed.netloc, "", "", "")), - experiment_id=experiment_id.lstrip("0") or "0", - ) - - -def project_url(config: DatabricksConfig) -> str: - """Pure URL presentation, including for historical configurations.""" - target = _target(config) - return f"{target.host}/ml/experiments/{target.experiment_id}/traces" - - -def _request(method: Literal["GET", "POST"], url: str, **options: object) -> dict[str, object]: - request = ssrf_proxy.get if method == "GET" else ssrf_proxy.post - try: - response = request(url, max_retries=0, timeout=30.0, follow_redirects=False, **options) - except Exception: - raise DatabricksConfigurationError("connection_failed") from None - try: - if response.status_code == 401: - raise DatabricksConfigurationError("authentication_failed") - if response.status_code == 403: - raise DatabricksConfigurationError("permission_denied") - if response.status_code == 404 and method == "GET": - raise DatabricksConfigurationError("experiment_not_found_or_inaccessible") - if not 200 <= response.status_code < 300: - raise DatabricksConfigurationError("endpoint_unavailable") - try: - return _JSON_OBJECT.validate_python(response.json(), strict=True) - except ValueError: - raise DatabricksConfigurationError("malformed_response") from None - finally: - response.close() - - -def verify_configuration(config: DatabricksConfig) -> ReadableDatabricksExperiment: - """Read the exact experiment with the exporter's existing PAT/M2M precedence. - - GetExperiment uses the official MLflow request/response schema. No permissions, - SCIM or account API is needed, and no experiment, run or trace is created. - Read access alone is intentionally not presented as write authorization. - """ - target = _target(config) - if config.client_id and config.client_secret: - data = _request( - "POST", - target.host + "/oidc/v1/token", - auth=(config.client_id, config.client_secret), - data={"grant_type": "client_credentials", "scope": "all-apis"}, - ) - access_token = data.get("access_token") - token_type = data.get("token_type") - if ( - not isinstance(access_token, str) - or not access_token.strip() - or not isinstance(token_type, str) - or token_type.lower() != "bearer" - ): - raise DatabricksConfigurationError("authentication_failed") - else: - access_token = config.personal_access_token - if not access_token or not access_token.strip(): - raise DatabricksConfigurationError("authentication_failed") - request = GetExperiment(experiment_id=target.experiment_id) - data = _request( - "GET", - target.host + "/api/2.0/mlflow/experiments/get", - params={"experiment_id": request.experiment_id}, - headers={"Authorization": f"Bearer {access_token}", "Accept": "application/json"}, - ) - try: - result = ParseDict(data, GetExperiment.Response(), ignore_unknown_fields=True) - except (ValueError, TypeError, ParseError): - raise DatabricksConfigurationError("malformed_response") from None - if not result.HasField("experiment") or not result.experiment.experiment_id: - raise DatabricksConfigurationError("malformed_response") - if result.experiment.experiment_id != target.experiment_id: - raise DatabricksConfigurationError("experiment_id_mismatch") - if result.experiment.lifecycle_stage != "active": - raise DatabricksConfigurationError("experiment_not_active") - return target diff --git a/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/mlflow_trace.py b/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/mlflow_trace.py deleted file mode 100644 index 0f7ffb08de190e..00000000000000 --- a/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/mlflow_trace.py +++ /dev/null @@ -1,694 +0,0 @@ -import json -import logging -import os -from collections.abc import Callable -from contextvars import ContextVar -from datetime import datetime, timedelta -from typing import Any, cast, override - -import mlflow -from mlflow.entities import Document, LiveSpan, Span, SpanEvent, SpanStatusCode, SpanType, TraceData, TraceInfo -from mlflow.entities.span import NoOpSpan -from mlflow.exceptions import RestException -from mlflow.tracing.client import TracingClient -from mlflow.tracing.constant import SpanAttributeKey, TokenUsageKey, TraceMetadataKey -from mlflow.tracing.export.mlflow_v3 import MlflowV3SpanExporter -from mlflow.tracing.fluent import start_span_no_context, update_current_trace -from mlflow.tracing.processor.mlflow_v3 import MlflowV3SpanProcessor -from mlflow.tracing.provider import detach_span_from_context, set_span_in_context -from mlflow.tracing.trace_manager import InMemoryTraceManager -from mlflow.tracing.utils import encode_span_id, get_otel_attribute -from opentelemetry import trace as otel_trace -from opentelemetry.context import Context -from opentelemetry.sdk.resources import Resource -from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.sdk.trace.sampling import ALWAYS_ON -from sqlalchemy import select - -from core.ops.base_trace_instance import BaseTraceInstance -from core.ops.entities.trace_entity import ( - BaseTraceInfo, - DatasetRetrievalTraceInfo, - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - SuggestedQuestionTraceInfo, - ToolTraceInfo, - TraceTaskName, - WorkflowTraceInfo, -) -from core.ops.utils import JSON_DICT_ADAPTER -from dify_trace_mlflow.config import DatabricksConfig, MLflowConfig -from extensions.ext_database import db -from graphon.enums import BuiltinNodeTypes -from models import EndUser -from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom - -logger = logging.getLogger(__name__) - -type SpanAttributes = dict[str, object] - - -def datetime_to_nanoseconds(dt: datetime | None) -> int | None: - """Convert datetime to nanosecond timestamp for MLflow API""" - if dt is None: - return None - return int(dt.timestamp() * 1_000_000_000) - - -class _CheckedTracingClient(TracingClient): - """Remember errors hidden by SDK exporters without changing SDK serialization or I/O.""" - - def __init__(self, tracking_uri: str, failed: ContextVar[bool]): - super().__init__(tracking_uri) - self._failed = failed - - def _checked[T](self, operation: Callable[[], T]) -> T: - try: - return operation() - except Exception: - self._failed.set(True) - raise - - @override - def start_trace(self, trace_info: TraceInfo) -> TraceInfo: - return self._checked(lambda: super(_CheckedTracingClient, self).start_trace(trace_info)) - - @override - def _upload_trace_data(self, trace_info: TraceInfo, trace_data: TraceData) -> None: - self._checked(lambda: super(_CheckedTracingClient, self)._upload_trace_data(trace_info, trace_data)) - - @override - def log_spans(self, location: str, spans: list[Span]) -> list[Span]: - try: - return super().log_spans(location, spans) - except NotImplementedError: - # SDK falls back to the full trace artifact on pre-3.4/file-store servers. - raise - except RestException as error: - if "REST OTLP span logging is not supported" not in error.message: - self._failed.set(True) - raise - except Exception: - self._failed.set(True) - raise - - -class _SynchronousMlflowExporter(MlflowV3SpanExporter): - """Celery owns asynchronous dispatch; no SDK queue may survive adapter LRU eviction.""" - - def __init__(self, tracking_uri: str): - self.failed: ContextVar[bool] = ContextVar("mlflow_legacy_export_failed", default=False) - super().__init__(tracking_uri) - self._client = _CheckedTracingClient(tracking_uri, self.failed) - - @override - def _should_enable_async_logging(self) -> bool: - return False - - -class MLflowDataTrace(BaseTraceInstance): - def __init__(self, config: MLflowConfig | DatabricksConfig): - super().__init__(config) - self._tracer: otel_trace.Tracer | None = None - self._exporter: _SynchronousMlflowExporter | None = None - self._experiment_id = config.experiment_id - if isinstance(config, DatabricksConfig): - # Preserve the SDK's existing Databricks/UC destination selection. - self._setup_databricks(config) - os.environ["MLFLOW_ENABLE_ASYNC_TRACE_LOGGING"] = "true" - return - self._setup_mlflow(config) - - # The public MLflow destination API still replaces the SDK's shared provider - # and can mutate OTEL_RESOURCE_ATTRIBUTES. Keep this legacy SDK pipeline private. - # No SDK background queue or atexit callback may retain an evicted adapter. - self._provider = TracerProvider(resource=Resource({}), sampler=ALWAYS_ON, shutdown_on_exit=False) - self._exporter = _SynchronousMlflowExporter(tracking_uri=self._tracking_uri) - self._provider.add_span_processor(MlflowV3SpanProcessor(self._exporter, export_metrics=False)) - self._tracer = self._provider.get_tracer(__name__) - - def _setup_databricks(self, config: DatabricksConfig): - """Setup connection to Databricks-managed MLflow instances""" - os.environ["DATABRICKS_HOST"] = config.host - - if config.client_id and config.client_secret: - # OAuth: https://docs.databricks.com/aws/en/dev-tools/auth/oauth-m2m?language=Environment - os.environ["DATABRICKS_CLIENT_ID"] = config.client_id - os.environ["DATABRICKS_CLIENT_SECRET"] = config.client_secret - elif config.personal_access_token: - # PAT: https://docs.databricks.com/aws/en/dev-tools/auth/pat - os.environ["DATABRICKS_TOKEN"] = config.personal_access_token - else: - raise ValueError( - "Either Databricks token (PAT) or client id and secret (OAuth) must be provided" - "See https://docs.databricks.com/aws/en/dev-tools/auth/#what-authorization-option-should-i-choose " - "for more information about the authorization options." - ) - self._tracking_uri = "databricks" - mlflow.set_tracking_uri("databricks") - mlflow.set_experiment(experiment_id=config.experiment_id) - - # Remove trailing slash from host - config.host = config.host.rstrip("/") - self._project_url = f"{config.host}/ml/experiments/{config.experiment_id}/traces" - - def _setup_mlflow(self, config: MLflowConfig): - """Setup connection to MLflow instances""" - self._tracking_uri = config.tracking_uri - - # Simple auth if provided - if config.username and config.password: - os.environ["MLFLOW_TRACKING_USERNAME"] = config.username - os.environ["MLFLOW_TRACKING_PASSWORD"] = config.password - - self._project_url = f"{config.tracking_uri}/#/experiments/{config.experiment_id}/traces" - - def _start_span_no_context( - self, - *, - name: str, - span_type: str, - parent_span: LiveSpan | None = None, - inputs: object | None = None, - attributes: SpanAttributes | None = None, - start_time_ns: int | None = None, - ) -> LiveSpan: - """Use the real SDK processor without inheriting a process-wide destination. - - Every span needs the experiment attribute: the SDK's incremental exporter - resolves each child independently, rather than inheriting the root's location. - A missing/nonrecording span is an export failure, never successful dispatch. - """ - if self._tracer is None: - span = start_span_no_context( - name=name, - span_type=span_type, - parent_span=parent_span, - inputs=inputs, - attributes=cast(dict[str, str] | None, attributes), - start_time_ns=start_time_ns, - ) - if isinstance(span, NoOpSpan): - raise RuntimeError("mlflow_legacy_span_not_recording") - return span - context = Context() - if parent_span is not None: - parent = parent_span._span - if not isinstance(parent, otel_trace.Span): - raise RuntimeError("mlflow_legacy_parent_span_invalid") - context = otel_trace.set_span_in_context(parent, context) - try: - otel_span = self._tracer.start_span( - name, - context=context, - attributes={SpanAttributeKey.EXPERIMENT_ID: json.dumps(self._experiment_id)}, - start_time=start_time_ns, - ) - except Exception: - raise RuntimeError("mlflow_legacy_span_start_failed") from None - if not otel_span.is_recording(): - raise RuntimeError("mlflow_legacy_span_not_recording") - trace_id = get_otel_attribute(otel_span, SpanAttributeKey.REQUEST_ID) - manager = InMemoryTraceManager.get_instance() - span = ( - manager.get_span_from_id(trace_id, encode_span_id(otel_span.get_span_context().span_id)) - if trace_id - else None - ) - if not isinstance(span, LiveSpan): - manager.pop_trace(otel_span.get_span_context().trace_id) - otel_span.end() - raise RuntimeError("mlflow_legacy_span_not_registered") - span.set_span_type(span_type) - if inputs is not None: - span.set_inputs(inputs) - span.set_attributes(attributes or {}) - span.set_attribute(SpanAttributeKey.EXPERIMENT_ID, self._experiment_id) - return span - - @override - def trace(self, trace_info: BaseTraceInfo): - """Dispatch legacy data; synchronous OSS export failures must reach the task owner.""" - failure_token = self._exporter.failed.set(False) if self._exporter is not None else None - try: - match trace_info: - case WorkflowTraceInfo(): - self.workflow_trace(trace_info) - case MessageTraceInfo(): - self.message_trace(trace_info) - case ToolTraceInfo(): - self.tool_trace(trace_info) - case ModerationTraceInfo(): - self.moderation_trace(trace_info) - case DatasetRetrievalTraceInfo(): - self.dataset_retrieval_trace(trace_info) - case SuggestedQuestionTraceInfo(): - self.suggested_question_trace(trace_info) - case GenerateNameTraceInfo(): - self.generate_name_trace(trace_info) - if self._exporter is not None and self._exporter.failed.get(): - raise RuntimeError("mlflow_legacy_export_failed") - except Exception: - logger.exception("[MLflow] Trace error") - raise - finally: - if self._exporter is not None and failure_token is not None: - self._exporter.failed.reset(failure_token) - - def workflow_trace(self, trace_info: WorkflowTraceInfo): - """Create workflow span as root, with node spans as children""" - # fields with sys.xyz is added by Dify, they are duplicate to trace_info.metadata - raw_inputs = trace_info.workflow_run_inputs or {} - workflow_inputs = {k: v for k, v in raw_inputs.items() if not k.startswith("sys.")} - - # Special inputs propagated by system - if trace_info.query: - workflow_inputs["query"] = trace_info.query - - workflow_span = self._start_span_no_context( - name=TraceTaskName.WORKFLOW_TRACE.value, - span_type=SpanType.CHAIN, - inputs=workflow_inputs, - attributes=trace_info.metadata, - start_time_ns=datetime_to_nanoseconds(trace_info.start_time), - ) - - # Set reserved fields in trace-level metadata - trace_metadata = {} - if user_id := trace_info.metadata.get("user_id"): - trace_metadata[TraceMetadataKey.TRACE_USER] = user_id - if session_id := trace_info.conversation_id: - trace_metadata[TraceMetadataKey.TRACE_SESSION] = session_id - self._set_trace_metadata(workflow_span, trace_metadata) - - try: - # Create child spans for workflow nodes - for node in self._get_workflow_nodes(trace_info): - inputs = None - attributes: SpanAttributes = { - "node_id": node.id, - "node_type": node.node_type, - "status": node.status, - "tenant_id": node.tenant_id, - "app_id": node.app_id, - "app_name": node.title, - } - - if node.node_type in (BuiltinNodeTypes.LLM, BuiltinNodeTypes.QUESTION_CLASSIFIER): - inputs, llm_attributes = self._parse_llm_inputs_and_attributes(node) - attributes.update(llm_attributes) - elif node.node_type == BuiltinNodeTypes.HTTP_REQUEST: - inputs = node.process_data # contains request URL - - if not inputs: - inputs = JSON_DICT_ADAPTER.validate_json(node.inputs) if node.inputs else {} - - node_span = self._start_span_no_context( - name=node.title, - span_type=self._get_node_span_type(node.node_type), - parent_span=workflow_span, - inputs=inputs, - attributes=attributes, - start_time_ns=datetime_to_nanoseconds(node.created_at), - ) - - # Handle node errors - if node.status != "succeeded": - node_span.set_status(SpanStatusCode.ERROR) - node_span.add_event( - SpanEvent( # type: ignore[abstract] - name="exception", - attributes={ - "exception.message": f"Node failed with status: {node.status}", - "exception.type": "Error", - "exception.stacktrace": f"Node failed with status: {node.status}", - }, - ) - ) - - # End node span - finished_at = node.created_at + timedelta(seconds=node.elapsed_time) - outputs = JSON_DICT_ADAPTER.validate_json(node.outputs) if node.outputs else {} - if node.node_type == BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL: - outputs = self._parse_knowledge_retrieval_outputs(outputs) - elif node.node_type == BuiltinNodeTypes.LLM: - outputs = outputs.get("text", outputs) - node_span.end( - outputs=outputs, - end_time_ns=datetime_to_nanoseconds(finished_at), - ) - - # Handle workflow-level errors - if trace_info.error: - workflow_span.set_status(SpanStatusCode.ERROR) - workflow_span.add_event( - SpanEvent( # type: ignore[abstract] - name="exception", - attributes={ - "exception.message": trace_info.error, - "exception.type": "Error", - "exception.stacktrace": trace_info.error, - }, - ) - ) - - finally: - workflow_span.end( - outputs=trace_info.workflow_run_outputs, - end_time_ns=datetime_to_nanoseconds(trace_info.end_time), - ) - - def _parse_llm_inputs_and_attributes(self, node: WorkflowNodeExecutionModel) -> tuple[object, SpanAttributes]: - """Parse LLM inputs and attributes from LLM workflow node""" - if node.process_data is None: - return {}, {} - - try: - data = JSON_DICT_ADAPTER.validate_json(node.process_data) - except (ValueError, TypeError): - return {}, {} - - inputs = self._parse_prompts(data.get("prompts")) - attributes = { - "model_name": data.get("model_name"), - "model_provider": data.get("model_provider"), - "finish_reason": data.get("finish_reason"), - } - - if hasattr(SpanAttributeKey, "MESSAGE_FORMAT"): - attributes[SpanAttributeKey.MESSAGE_FORMAT] = "dify" - - if usage := data.get("usage"): - # Set reserved token usage attributes - attributes[SpanAttributeKey.CHAT_USAGE] = { - TokenUsageKey.INPUT_TOKENS: usage.get("prompt_tokens", 0), - TokenUsageKey.OUTPUT_TOKENS: usage.get("completion_tokens", 0), - TokenUsageKey.TOTAL_TOKENS: usage.get("total_tokens", 0), - } - # Store raw usage data as well as it includes more data like price - attributes["usage"] = usage - - return inputs, attributes - - def _parse_knowledge_retrieval_outputs(self, outputs: dict[str, Any]): - """Parse KR outputs and attributes from KR workflow node""" - retrieved = outputs.get("result", []) - - if not retrieved or not isinstance(retrieved, list): - return outputs - - documents = [] - for item in retrieved: - documents.append(Document(page_content=item.get("content", ""), metadata=item.get("metadata", {}))) - return documents - - def message_trace(self, trace_info: MessageTraceInfo): - """Create span for CHATBOT message processing""" - if not trace_info.message_data: - return - - file_list = cast(list[str], trace_info.file_list) or [] - if message_file_data := trace_info.message_file_data: - base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001") - file_list.append(f"{base_url}/{message_file_data.url}") - - span = self._start_span_no_context( - name=TraceTaskName.MESSAGE_TRACE.value, - span_type=SpanType.LLM, - inputs=self._parse_prompts(trace_info.inputs), - attributes={ - "message_id": trace_info.message_id, - "model_provider": trace_info.message_data.model_provider, - "model_id": trace_info.message_data.model_id, - "conversation_mode": trace_info.conversation_mode, - "file_list": file_list, - "total_price": trace_info.message_data.total_price, - **trace_info.metadata, - }, - start_time_ns=datetime_to_nanoseconds(trace_info.start_time), - ) - - if hasattr(SpanAttributeKey, "MESSAGE_FORMAT"): - span.set_attribute(SpanAttributeKey.MESSAGE_FORMAT, "dify") - - # Set token usage - span.set_attribute( - SpanAttributeKey.CHAT_USAGE, - { - TokenUsageKey.INPUT_TOKENS: trace_info.message_tokens or 0, - TokenUsageKey.OUTPUT_TOKENS: trace_info.answer_tokens or 0, - TokenUsageKey.TOTAL_TOKENS: trace_info.total_tokens or 0, - }, - ) - - # Set reserved fields in trace-level metadata - trace_metadata = {} - if user_id := self._get_message_user_id(trace_info.metadata): - trace_metadata[TraceMetadataKey.TRACE_USER] = user_id - if session_id := trace_info.metadata.get("conversation_id"): - trace_metadata[TraceMetadataKey.TRACE_SESSION] = session_id - self._set_trace_metadata(span, trace_metadata) - - if trace_info.error: - span.set_status(SpanStatusCode.ERROR) - span.add_event( - SpanEvent( # type: ignore[abstract] - name="error", - attributes={ - "exception.message": trace_info.error, - "exception.type": "Error", - "exception.stacktrace": trace_info.error, - }, - ) - ) - - span.end( - outputs=trace_info.message_data.answer, - end_time_ns=datetime_to_nanoseconds(trace_info.end_time), - ) - - def _get_message_user_id(self, metadata: dict[str, Any]) -> str | None: - if (end_user_id := metadata.get("from_end_user_id")) and ( - end_user_data := db.session.get(EndUser, end_user_id) - ): - return end_user_data.session_id - - return metadata.get("from_account_id") # type: ignore[return-value] - - def tool_trace(self, trace_info: ToolTraceInfo): - span = self._start_span_no_context( - name=trace_info.tool_name, - span_type=SpanType.TOOL, - inputs=trace_info.tool_inputs, - attributes={ - "message_id": trace_info.message_id, - "metadata": trace_info.metadata, - "tool_config": trace_info.tool_config, - "tool_parameters": trace_info.tool_parameters, - }, - start_time_ns=datetime_to_nanoseconds(trace_info.start_time), - ) - - # Handle tool errors - if trace_info.error: - span.set_status(SpanStatusCode.ERROR) - span.add_event( - SpanEvent( # type: ignore[abstract] - name="error", - attributes={ - "exception.message": trace_info.error, - "exception.type": "Error", - "exception.stacktrace": trace_info.error, - }, - ) - ) - - span.end( - outputs=trace_info.tool_outputs, - end_time_ns=datetime_to_nanoseconds(trace_info.end_time), - ) - - def moderation_trace(self, trace_info: ModerationTraceInfo): - if trace_info.message_data is None: - return - - start_time = trace_info.start_time or trace_info.message_data.created_at - span = self._start_span_no_context( - name=TraceTaskName.MODERATION_TRACE.value, - span_type=SpanType.TOOL, - inputs=trace_info.inputs or {}, - attributes={ - "message_id": trace_info.message_id, - "metadata": trace_info.metadata, - }, - start_time_ns=datetime_to_nanoseconds(start_time), - ) - - span.end( - outputs={ - "action": trace_info.action, - "flagged": trace_info.flagged, - "preset_response": trace_info.preset_response, - }, - end_time_ns=datetime_to_nanoseconds(trace_info.end_time), - ) - - def dataset_retrieval_trace(self, trace_info: DatasetRetrievalTraceInfo): - if trace_info.message_data is None: - return - - span = self._start_span_no_context( - name=TraceTaskName.DATASET_RETRIEVAL_TRACE.value, - span_type=SpanType.RETRIEVER, - inputs=trace_info.inputs, - attributes={ - "message_id": trace_info.message_id, - "metadata": trace_info.metadata, - }, - start_time_ns=datetime_to_nanoseconds(trace_info.start_time), - ) - span.end(outputs={"documents": trace_info.documents}, end_time_ns=datetime_to_nanoseconds(trace_info.end_time)) - - def suggested_question_trace(self, trace_info: SuggestedQuestionTraceInfo): - if trace_info.message_data is None: - return - - start_time = trace_info.start_time or trace_info.message_data.created_at - end_time = trace_info.end_time or trace_info.message_data.updated_at - - span = self._start_span_no_context( - name=TraceTaskName.SUGGESTED_QUESTION_TRACE.value, - span_type=SpanType.TOOL, - inputs=trace_info.inputs, - attributes={ - "message_id": trace_info.message_id, - "model_provider": trace_info.model_provider, - "model_id": trace_info.model_id, - "total_tokens": trace_info.total_tokens or 0, - }, - start_time_ns=datetime_to_nanoseconds(start_time), - ) - - if trace_info.error: - span.set_status(SpanStatusCode.ERROR) - span.add_event( - SpanEvent( # type: ignore[abstract] - name="error", - attributes={ - "exception.message": trace_info.error, - "exception.type": "Error", - "exception.stacktrace": trace_info.error, - }, - ) - ) - - span.end(outputs=trace_info.suggested_question, end_time_ns=datetime_to_nanoseconds(end_time)) - - def generate_name_trace(self, trace_info: GenerateNameTraceInfo): - span = self._start_span_no_context( - name=TraceTaskName.GENERATE_NAME_TRACE.value, - span_type=SpanType.CHAIN, - inputs=trace_info.inputs, - attributes={"message_id": trace_info.message_id}, - start_time_ns=datetime_to_nanoseconds(trace_info.start_time), - ) - span.end(outputs=trace_info.outputs, end_time_ns=datetime_to_nanoseconds(trace_info.end_time)) - - def _get_workflow_nodes(self, trace_info: WorkflowTraceInfo): - """Export caller nodes without disclosing Workflow Tool source internals.""" - 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") - workflow_nodes = db.session.scalars( - select(WorkflowNodeExecutionModel) - .where( - WorkflowNodeExecutionModel.workflow_run_id == trace_info.workflow_run_id, - WorkflowNodeExecutionModel.tenant_id == trace_info.tenant_id, - WorkflowNodeExecutionModel.app_id == app_id, - WorkflowNodeExecutionModel.triggered_from == WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - .order_by(WorkflowNodeExecutionModel.created_at) - ).all() - return workflow_nodes - - def _get_node_span_type(self, node_type: str) -> str: - """Map Dify node types to MLflow span types""" - node_type_mapping = { - BuiltinNodeTypes.LLM: SpanType.LLM, - BuiltinNodeTypes.QUESTION_CLASSIFIER: SpanType.LLM, - BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL: SpanType.RETRIEVER, - BuiltinNodeTypes.TOOL: SpanType.TOOL, - BuiltinNodeTypes.CODE: SpanType.TOOL, - BuiltinNodeTypes.HTTP_REQUEST: SpanType.TOOL, - BuiltinNodeTypes.AGENT: SpanType.AGENT, - } - return node_type_mapping.get(node_type, "CHAIN") # type: ignore[arg-type,call-overload] - - def _set_trace_metadata(self, span: Span, metadata: dict[str, Any]): - token = None - try: - # NB: Set span in context such that we can use update_current_trace() API - token = set_span_in_context(span) - update_current_trace(metadata=metadata) - finally: - if token: - detach_span_from_context(token) - - def _parse_prompts(self, prompts): - """Postprocess prompts format to be standard chat messages""" - match prompts: - case str(): - return prompts - case dict(): - return self._parse_single_message(prompts) - case list(): - messages = [self._parse_single_message(item) for item in prompts] - messages = self._resolve_tool_call_ids(messages) - return messages - return prompts # Fallback to original format - - def _parse_single_message(self, item: dict[str, Any]): - """Postprocess single message format to be standard chat message""" - role = item.get("role", "user") - msg = {"role": role, "content": item.get("text", "")} - - if ( - (tool_calls := item.get("tool_calls")) - # Tool message does not contain tool calls normally - and role != "tool" - ): - msg["tool_calls"] = tool_calls - - if files := item.get("files"): - msg["files"] = files - - return msg - - def _resolve_tool_call_ids(self, messages: list[dict]): - """ - The tool call message from Dify does not contain tool call ids, which is not - ideal for debugging. This method resolves the tool call ids by matching the - tool call name and parameters with the tool instruction messages. - """ - tool_call_ids = [] - for msg in messages: - if tool_calls := msg.get("tool_calls"): - tool_call_ids = [t["id"] for t in tool_calls] - if msg["role"] == "tool": - # Get the tool call id in the order of the tool call messages - # assuming Dify runs tools sequentially - if tool_call_ids: - msg["tool_call_id"] = tool_call_ids.pop(0) - return messages - - def api_check(self): - """Simple connection test""" - try: - mlflow.MlflowClient(tracking_uri=self._tracking_uri).search_experiments(max_results=1) - return True - except Exception as e: - raise ValueError(f"MLflow connection failed: {str(e)}") from e - - def get_project_url(self): - return self._project_url diff --git a/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/oss_configuration.py b/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/oss_configuration.py deleted file mode 100644 index 20c201e00f0126..00000000000000 --- a/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/oss_configuration.py +++ /dev/null @@ -1,201 +0,0 @@ -"""Read one OSS MLflow experiment without changing SDK or authentication state. - -Configuration verification proves only read access to an existing active experiment. -It does not establish trace ingestion, artifact delivery, or write permission. The -legacy exporter remains a separate owner; configuration never constructs it. -""" - -import base64 -import re -from collections.abc import Mapping -from dataclasses import dataclass -from typing import Literal -from urllib.parse import unquote, urlsplit, urlunsplit - -from google.protobuf.json_format import ParseDict, ParseError -from mlflow.protos.service_pb2 import GetExperiment -from pydantic import TypeAdapter - -from core.helper import ssrf_proxy -from dify_trace_mlflow.config import MLflowConfig - -type VerificationCode = Literal[ - "invalid_tracking_uri", - "invalid_experiment_id", - "incomplete_basic_auth", - "authentication_failed", - "permission_denied", - "experiment_not_found_or_inaccessible", - "experiment_not_active", - "experiment_id_mismatch", - "malformed_response", - "endpoint_unavailable", - "connection_failed", -] - -_MESSAGES: dict[VerificationCode, str] = { - "invalid_tracking_uri": "MLflow requires an HTTP(S) tracking server URL with an optional path prefix.", - "invalid_experiment_id": "MLflow requires a numeric experiment ID.", - "incomplete_basic_auth": "MLflow username and password must both be supplied or both be empty.", - "authentication_failed": "MLflow authentication failed.", - "permission_denied": "MLflow experiment access was denied.", - "experiment_not_found_or_inaccessible": "The selected MLflow experiment was not found or is inaccessible.", - "experiment_not_active": "The selected MLflow experiment is not active.", - "experiment_id_mismatch": "MLflow returned a different experiment.", - "malformed_response": "MLflow verification returned an invalid response.", - "endpoint_unavailable": "MLflow endpoint could not be verified.", - "connection_failed": "MLflow verification request failed.", -} -_JSON_OBJECT = TypeAdapter(dict[str, object]) -_MAX_RESPONSE_BYTES = 1024 * 1024 - - -class MLflowConfigurationError(ValueError): - """Fixed diagnostics that never include request values or vendor response text.""" - - def __init__(self, code: VerificationCode) -> None: - self.code = code - super().__init__(_MESSAGES[code]) - - -@dataclass(frozen=True, slots=True) -class ReadableMLflowExperiment: - tracking_uri: str - experiment_id: str - - -def _tracking_uri(value: object) -> str: - if not isinstance(value, str): - raise MLflowConfigurationError("invalid_tracking_uri") - # Only surrounding ordinary spaces are normalization; URL parsers silently - # discard some control characters, so inspect the original input first. - value = value.strip(" ") - try: - parsed = urlsplit(value) - host, port = parsed.hostname, parsed.port - path = parsed.path.rstrip("/") - decoded_path = unquote(path, errors="strict") - valid = ( - parsed.scheme in {"http", "https"} - and host - and parsed.username is None - and parsed.password is None - and "?" not in value - and "#" not in value - and "\\" not in value - and not any(ord(char) <= 32 or ord(char) == 127 for char in value) - and not parsed.netloc.endswith(":") - and (port is None or 1 <= port <= 65535) - and "%" not in host - and not re.search(r"%(?![0-9a-fA-F]{2})", path) - and not re.search(r"%(?:2f|5c)", path, re.IGNORECASE) - and not any(char in decoded_path for char in "\\?#%") - and not any(ord(char) < 32 or ord(char) == 127 for char in decoded_path) - and "//" not in parsed.path - and not any(part in {".", ".."} for part in decoded_path.split("/")) - ) - if not valid or host is None: - raise MLflowConfigurationError("invalid_tracking_uri") - host = host.encode("idna").decode("ascii").lower() - authority = f"[{host}]" if ":" in host else host - if port is not None: - authority = f"{authority}:{port}" - return urlunsplit((parsed.scheme, authority, path, "", "")) - except (UnicodeError, ValueError): - raise MLflowConfigurationError("invalid_tracking_uri") from None - - -def _experiment_id(value: str) -> str: - value = value.strip() - if not value or not value.isascii() or not value.isdigit(): - raise MLflowConfigurationError("invalid_experiment_id") - return value.lstrip("0") or "0" - - -def normalize_configuration(values: Mapping[str, object]) -> MLflowConfig: - """Validate new destinations while retaining masked/nullable password values. - - Pair validation belongs after masked credentials are restored and decrypted. - Keeping it out of MLflowConfig also leaves historical records readable and - avoids interpreting RSA ciphertext as a user-supplied password. - """ - config = MLflowConfig.model_validate(values) - return config.model_copy( - update={ - "tracking_uri": _tracking_uri(values.get("tracking_uri", config.tracking_uri)), - "experiment_id": _experiment_id(config.experiment_id), - } - ) - - -def project_url(config: MLflowConfig) -> str: - """Pure link calculation, independent of authentication and server access.""" - host = _tracking_uri(config.tracking_uri) - experiment_id = _experiment_id(config.experiment_id) - return f"{host}/#/experiments/{experiment_id}/traces" - - -def verify_configuration(config: MLflowConfig) -> ReadableMLflowExperiment: - """GET the exact experiment with only this request's optional Basic credentials.""" - config = normalize_configuration(config.model_dump()) - username, password = config.username, config.password - has_username, has_password = bool(username), bool(password) - if has_username != has_password or ( - has_username - and ( - username is None - or password is None - or not username.strip() - or not password.strip() - or ":" in username - or any(ord(char) < 32 or ord(char) == 127 for char in username + password) - ) - ): - raise MLflowConfigurationError("incomplete_basic_auth") - - # Explicit Cookie suppresses the pooled HTTP client's ambient cookie jar. - # No MLflow environment variables, SDK client, or process-wide auth is read. - headers = {"Accept": "application/json", "Accept-Encoding": "identity", "Cookie": ""} - if username and password: - encoded = base64.b64encode(f"{username}:{password}".encode()).decode() - headers["Authorization"] = f"Basic {encoded}" - request = GetExperiment(experiment_id=config.experiment_id) - try: - response = ssrf_proxy.get( - config.tracking_uri + "/api/2.0/mlflow/experiments/get", - params={"experiment_id": request.experiment_id}, - headers=headers, - auth=None, - max_retries=0, - timeout=30.0, - follow_redirects=False, - stream_response=True, - ) - except Exception: - raise MLflowConfigurationError("connection_failed") from None - try: - if response.status_code == 401: - raise MLflowConfigurationError("authentication_failed") - if response.status_code == 403: - raise MLflowConfigurationError("permission_denied") - if response.status_code == 404: - raise MLflowConfigurationError("experiment_not_found_or_inaccessible") - if not 200 <= response.status_code < 300: - raise MLflowConfigurationError("endpoint_unavailable") - try: - buffered = ssrf_proxy.buffer_response(response, max_response_bytes=_MAX_RESPONSE_BYTES) - data = _JSON_OBJECT.validate_python(buffered.json(), strict=True) - result = ParseDict(data, GetExperiment.Response(), ignore_unknown_fields=True) - except (ValueError, TypeError, ParseError): - raise MLflowConfigurationError("malformed_response") from None - except Exception: - raise MLflowConfigurationError("connection_failed") from None - finally: - response.close() - if not result.HasField("experiment") or not result.experiment.experiment_id: - raise MLflowConfigurationError("malformed_response") - if result.experiment.experiment_id != config.experiment_id: - raise MLflowConfigurationError("experiment_id_mismatch") - if result.experiment.lifecycle_stage != "active": - raise MLflowConfigurationError("experiment_not_active") - return ReadableMLflowExperiment(config.tracking_uri, config.experiment_id) diff --git a/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/provider.py b/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/provider.py new file mode 100644 index 00000000000000..59fd12d3561a74 --- /dev/null +++ b/api/providers/trace/trace-mlflow/src/dify_trace_mlflow/provider.py @@ -0,0 +1,349 @@ +"""MLflow tracking servers and Databricks-managed MLflow experiments. + +Spans are serialized with MLflow's own entities, but every request carries this +call's credentials through the egress. Nothing is read from or written to the +process environment, and no MLflow tracer, fluent API, or SDK client runs. + +OSS MLflow ingests OTLP/HTTP on its experiment endpoint (MLflow >= 3.4). +Databricks keeps the experiment destination through the V3 trace API and the +trace artifact; its OTLP endpoint writes Unity Catalog tables instead. +""" + +import base64 +import re +from collections.abc import Mapping +from typing import override +from urllib.parse import quote, unquote, urlencode, urlsplit, urlunsplit + +from dify_trace_sdk.egress import Egress +from dify_trace_sdk.errors import ( + EgressError, + ExportOutcomeUnknownError, + ExportRejectedError, + ExportRetryableError, + TraceProviderError, + VerificationError, +) +from dify_trace_sdk.exporting import export_endpoint, send_export_request +from dify_trace_sdk.model import Trace +from dify_trace_sdk.otlp import post_otlp, read_otlp_ack +from dify_trace_sdk.provider import ProviderCapabilities, SendResult, TraceProvider +from google.protobuf.json_format import ParseDict, ParseError +from mlflow.protos.service_pb2 import GetExperiment +from mlflow.tracing.constant import TRACE_REQUEST_ID_PREFIX +from mlflow.tracing.utils.otlp import MLFLOW_EXPERIMENT_ID_HEADER, OTLP_TRACES_PATH + +from dify_trace_mlflow._databricks import DatabricksStore, trace_info, upload_trace_data +from dify_trace_mlflow._spans import mlflow_spans, otlp_request, standalone, trace_data +from dify_trace_mlflow.config import DatabricksConfig, MLflowConfig + +VERIFY_TIMEOUT_SECONDS = 10.0 +_GET_EXPERIMENT_PATH = "/api/2.0/mlflow/experiments/get" +_OAUTH_TOKEN_PATH = "/oidc/v1/token" +_OAUTH_TOKEN_REQUEST = urlencode({"grant_type": "client_credentials", "scope": "all-apis"}).encode() + + +def _has_control_character(value: str) -> bool: + return any(ord(char) < 32 or ord(char) == 127 for char in value) + + +def _basic_authorization(username: str, password: str) -> str: + return "Basic " + base64.b64encode(f"{username}:{password}".encode()).decode() + + +def _basic_credentials(config: MLflowConfig) -> tuple[str, str] | None: + """Both or neither; the username cannot contain the Basic separator.""" + username, password = config.username or "", config.password or "" + if not username and not password: + return None + if not username.strip() or not password.strip() or ":" in username or _has_control_character(username + password): + raise VerificationError("incomplete_basic_auth") + return username, password + + +def _experiment_id(value: str) -> str: + stripped = value.strip() + if not stripped.isascii() or not stripped.isdigit(): + raise VerificationError("invalid_experiment_id") + return stripped.lstrip("0") or "0" + + +def _tracking_uri(value: str) -> str: + """Normalize an OSS tracking server URL with an optional path prefix. + + Anything a URL parser or proxy could read in two ways is rejected, and the + raw input is inspected first because parsers silently drop some characters. + """ + value = value.strip(" ") + try: + parsed = urlsplit(value) + host, port = parsed.hostname, parsed.port + path = parsed.path.rstrip("/") + decoded_path = unquote(path, errors="strict") + valid = ( + parsed.scheme in {"http", "https"} + and host + and parsed.username is None + and parsed.password is None + and "?" not in value + and "#" not in value + and "\\" not in value + and not any(ord(char) <= 32 or ord(char) == 127 for char in value) + and not parsed.netloc.endswith(":") + and (port is None or 1 <= port <= 65535) + and "%" not in host + and not re.search(r"%(?![0-9a-fA-F]{2})", path) + and not re.search(r"%(?:2f|5c)", path, re.IGNORECASE) + and not any(char in decoded_path for char in "\\?#%") + and not _has_control_character(decoded_path) + and "//" not in parsed.path + and not any(part in {".", ".."} for part in decoded_path.split("/")) + ) + if not valid or host is None: + raise VerificationError("invalid_tracking_uri") + host = host.encode("idna").decode("ascii").lower() + except (UnicodeError, ValueError): + raise VerificationError("invalid_tracking_uri") from None + authority = f"[{host}]" if ":" in host else host + if port is not None: + authority = f"{authority}:{port}" + return urlunsplit((parsed.scheme, authority, path, "", "")) + + +def _workspace(config: DatabricksConfig) -> tuple[str, str]: + """The HTTPS workspace origin and the experiment ID, normalized.""" + host = config.host.strip() + try: + parsed = urlsplit(host) + port = parsed.port + except ValueError: + raise VerificationError("invalid_workspace_url") from None + valid = ( + parsed.scheme == "https" + and parsed.hostname + and parsed.username is None + and parsed.password is None + and parsed.path in {"", "/"} + and not parsed.query + and not parsed.fragment + and port != 0 + and "\\" not in host + and not any(ord(char) <= 32 or ord(char) == 127 for char in host) + ) + if not valid: + raise VerificationError("invalid_workspace_url") + return urlunsplit((parsed.scheme, parsed.netloc, "", "", "")), _experiment_id(config.experiment_id) + + +def _verification_json( + egress: Egress, + method: str, + url: str, + *, + headers: Mapping[str, str], + content: bytes | None, + not_found: str, +) -> dict[str, object]: + """Send one verification request; ``not_found`` is the code for HTTP 404.""" + try: + response = egress.request( + method, url, headers=headers, content=content, timeout=VERIFY_TIMEOUT_SECONDS, sensitive_url=False + ) + except EgressError: + raise VerificationError("connection_failed") from None + if response.status_code == 401: + raise VerificationError("authentication_failed") + if response.status_code == 403: + raise VerificationError("permission_denied") + if response.status_code == 404: + raise VerificationError(not_found) + if not 200 <= response.status_code < 300: + raise VerificationError("endpoint_unavailable") + try: + body = response.json() + except ValueError: + raise VerificationError("malformed_response") from None + if not isinstance(body, dict): + raise VerificationError("malformed_response") + return body + + +def _check_experiment(egress: Egress, base_url: str, experiment_id: str, headers: Mapping[str, str]) -> None: + """Read the exact experiment. Read access does not prove trace write permission.""" + body = _verification_json( + egress, + "GET", + f"{base_url}{_GET_EXPERIMENT_PATH}?{urlencode({'experiment_id': experiment_id})}", + headers={"Accept": "application/json", **headers}, + content=None, + not_found="experiment_not_found_or_inaccessible", + ) + try: + result = ParseDict(body, GetExperiment.Response(), ignore_unknown_fields=True) + except (ParseError, TypeError, ValueError): + raise VerificationError("malformed_response") from None + if not result.HasField("experiment") or not result.experiment.experiment_id: + raise VerificationError("malformed_response") + if result.experiment.experiment_id != experiment_id: + raise VerificationError("experiment_id_mismatch") + if result.experiment.lifecycle_stage != "active": + raise VerificationError("experiment_not_active") + + +def _oauth_headers(client_id: str, client_secret: str) -> dict[str, str]: + return { + "Authorization": _basic_authorization(client_id, client_secret), + "Content-Type": "application/x-www-form-urlencoded", + "Accept": "application/json", + } + + +class MLflowProvider(TraceProvider[MLflowConfig]): + id = "mlflow" + config_model = MLflowConfig + capabilities = ProviderCapabilities(native_cost=True, trace_url=True, span_url=False) + + @override + def verify(self, config: MLflowConfig) -> MLflowConfig: + tracking_uri = _tracking_uri(config.tracking_uri) + experiment_id = _experiment_id(config.experiment_id) + credentials = _basic_credentials(config) + headers = {} if credentials is None else {"Authorization": _basic_authorization(*credentials)} + _check_experiment(self.egress, tracking_uri, experiment_id, headers) + return config.model_copy(update={"tracking_uri": tracking_uri, "experiment_id": experiment_id}) + + @override + def project_url(self, config: MLflowConfig) -> str | None: + try: + tracking_uri = _tracking_uri(config.tracking_uri) + experiment_id = _experiment_id(config.experiment_id) + except VerificationError: + return None + return f"{tracking_uri}/#/experiments/{experiment_id}/traces" + + @override + def send(self, config: MLflowConfig, trace: Trace) -> SendResult: + tracking_uri = export_endpoint(config.tracking_uri) + unit = standalone(trace) + trace_id = TRACE_REQUEST_ID_PREFIX + unit.trace_id + try: + request = otlp_request(mlflow_spans(unit, config.experiment_id, trace_id)) + except ValueError: + raise ExportRejectedError("MLflow cannot encode the trace") from None + headers = {MLFLOW_EXPERIMENT_ID_HEADER: config.experiment_id} + if config.username and config.password: + headers["Authorization"] = _basic_authorization(config.username, config.password) + response = post_otlp( + self.egress, + tracking_uri + OTLP_TRACES_PATH, + request, + destination="MLflow", + headers=headers, + sensitive_url=False, + ) + read_otlp_ack(response, "MLflow") + # MLflow 3.4's classic and GenAI trace views read selectedTraceId and + # selectedEvaluationId; later views read traceId. Each ignores the others. + query = urlencode({"traceId": trace_id, "selectedTraceId": trace_id, "selectedEvaluationId": trace_id}) + trace_url = f"{tracking_uri}/#/experiments/{config.experiment_id}/traces?{query}" + return SendResult(trace_id=trace_id, trace_url=trace_url, span_urls={}) + + +class DatabricksProvider(TraceProvider[DatabricksConfig]): + id = "databricks" + config_model = DatabricksConfig + capabilities = ProviderCapabilities(native_cost=True, trace_url=True, span_url=False) + + def _verification_token(self, host: str, config: DatabricksConfig) -> str: + """OAuth machine-to-machine credentials take precedence over a personal access token.""" + if config.client_id and config.client_secret: + body = _verification_json( + self.egress, + "POST", + host + _OAUTH_TOKEN_PATH, + headers=_oauth_headers(config.client_id, config.client_secret), + content=_OAUTH_TOKEN_REQUEST, + not_found="endpoint_unavailable", + ) + token, token_type = body.get("access_token"), body.get("token_type") + if ( + not isinstance(token, str) + or not token.strip() + or not isinstance(token_type, str) + or token_type.lower() != "bearer" + ): + raise VerificationError("authentication_failed") + return token + if not config.personal_access_token or not config.personal_access_token.strip(): + raise VerificationError("authentication_failed") + return config.personal_access_token + + def _export_token(self, host: str, config: DatabricksConfig) -> str: + if config.client_id and config.client_secret: + try: + response = send_export_request( + self.egress, + "POST", + host + _OAUTH_TOKEN_PATH, + destination="Databricks", + headers=_oauth_headers(config.client_id, config.client_secret), + content=_OAUTH_TOKEN_REQUEST, + sensitive_url=False, + ) + except ExportOutcomeUnknownError: + # Nothing is exported before the token arrives, so asking again is safe. + raise ExportRetryableError("Databricks OAuth token request failed in transport") from None + try: + body = response.json() + except ValueError: + body = None + token = body.get("access_token") if isinstance(body, dict) else None + if not isinstance(token, str) or not token: + raise ExportRejectedError("Databricks OAuth token was not returned") + return token + if not config.personal_access_token: + raise ExportRejectedError("Databricks credentials are not configured") + return config.personal_access_token + + @override + def verify(self, config: DatabricksConfig) -> DatabricksConfig: + host, experiment_id = _workspace(config) + token = self._verification_token(host, config) + _check_experiment(self.egress, host, experiment_id, {"Authorization": f"Bearer {token}"}) + return config.model_copy(update={"host": host, "experiment_id": experiment_id}) + + @override + def project_url(self, config: DatabricksConfig) -> str | None: + try: + host, experiment_id = _workspace(config) + except VerificationError: + return None + return f"{host}/ml/experiments/{experiment_id}/traces" + + @override + def send(self, config: DatabricksConfig, trace: Trace) -> SendResult: + host = export_endpoint(config.host) + unit = standalone(trace) + requested_id = TRACE_REQUEST_ID_PREFIX + unit.trace_id + try: + info = trace_info(unit, config.experiment_id, requested_id) + artifact = trace_data(mlflow_spans(unit, config.experiment_id, requested_id)) + except ValueError: + raise ExportRejectedError("Databricks cannot encode the trace") from None + headers = {"Authorization": f"Bearer {self._export_token(host, config)}"} + trace_id = DatabricksStore(self.egress, host, headers).start_trace(info).trace_id + if not trace_id: + raise ExportOutcomeUnknownError("Databricks did not acknowledge a trace identifier") + try: + if trace_id != requested_id: + artifact = trace_data(mlflow_spans(unit, config.experiment_id, trace_id)) + upload_trace_data(self.egress, host, headers, trace_id, artifact) + except TraceProviderError: + raise ExportOutcomeUnknownError( + "Databricks accepted trace metadata but artifact delivery is incomplete" + ) from None + # Databricks opens its trace panel with a different query key from OSS MLflow. + trace_url = ( + f"{host}/ml/experiments/{config.experiment_id}/traces?selectedEvaluationId={quote(trace_id, safe='')}" + ) + return SendResult(trace_id=trace_id, trace_url=trace_url, span_urls={}) diff --git a/api/providers/trace/trace-mlflow/tests/unit_tests/src/dify_trace_mlflow/test_config.py b/api/providers/trace/trace-mlflow/tests/unit_tests/src/dify_trace_mlflow/test_config.py new file mode 100644 index 00000000000000..d0f53619a744c0 --- /dev/null +++ b/api/providers/trace/trace-mlflow/tests/unit_tests/src/dify_trace_mlflow/test_config.py @@ -0,0 +1,60 @@ +import pytest +from dify_trace_mlflow.config import DatabricksConfig, MLflowConfig +from dify_trace_sdk.config import FormField +from pydantic import ValidationError + + +def test_mlflow_config_defaults_and_secrets() -> None: + config = MLflowConfig() + assert config.tracking_uri == "http://localhost:5000" + assert config.experiment_id == "0" + assert config.username is None + assert config.password is None + assert MLflowConfig(tracking_uri=" ").tracking_uri == "http://localhost:5000" + assert MLflowConfig.secret_fields() == {"password"} + assert MLflowConfig.derived_fields() == set() + + +def test_mlflow_config_reads_stored_configurations() -> None: + config = MLflowConfig.model_validate( + { + "tracking_uri": " https://mlflow.example/proxy/ ", + "experiment_id": " 42 ", + "username": "admin", + "password": "s3cret", + "retired_field": "ignored", + } + ) + assert config.tracking_uri == "https://mlflow.example/proxy/" + assert config.experiment_id == "42" + assert (config.username, config.password) == ("admin", "s3cret") + + +@pytest.mark.parametrize( + "tracking_uri", ["databricks", "databricks://profile", "ftp://mlflow.example", "https://admin:pw@mlflow.example"] +) +def test_mlflow_config_rejects_other_destinations(tracking_uri: str) -> None: + with pytest.raises(ValidationError) as caught: + MLflowConfig(tracking_uri=tracking_uri) + assert "admin:pw" not in str(caught.value) + + +def test_experiment_ids_must_be_integers() -> None: + with pytest.raises(ValidationError): + MLflowConfig(experiment_id="experiment-1") + with pytest.raises(ValidationError): + DatabricksConfig(experiment_id="experiment-1", host="https://dbc.example") + + +def test_databricks_config_fields() -> None: + config = DatabricksConfig.model_validate({"experiment_id": " 42 ", "host": "https://dbc.example"}) + assert config.experiment_id == "42" + assert (config.client_id, config.client_secret, config.personal_access_token) == (None, None, None) + assert DatabricksConfig.secret_fields() == {"client_secret", "personal_access_token"} + assert DatabricksConfig.form_fields() == ( + FormField(name="experiment_id", secret=False, required=True), + FormField(name="host", secret=False, required=True), + FormField(name="client_id", secret=False, required=False), + FormField(name="client_secret", secret=True, required=False), + FormField(name="personal_access_token", secret=True, required=False), + ) diff --git a/api/providers/trace/trace-mlflow/tests/unit_tests/src/dify_trace_mlflow/test_mlflow_trace.py b/api/providers/trace/trace-mlflow/tests/unit_tests/src/dify_trace_mlflow/test_mlflow_trace.py deleted file mode 100644 index e4ecf4db806f49..00000000000000 --- a/api/providers/trace/trace-mlflow/tests/unit_tests/src/dify_trace_mlflow/test_mlflow_trace.py +++ /dev/null @@ -1,1052 +0,0 @@ -"""Comprehensive tests for dify_trace_mlflow.mlflow_trace module.""" - -from __future__ import annotations - -import json -import os -from datetime import UTC, datetime -from types import SimpleNamespace -from unittest.mock import MagicMock, patch - -import pytest -from dify_trace_mlflow.config import DatabricksConfig, MLflowConfig -from dify_trace_mlflow.mlflow_trace import MLflowDataTrace, datetime_to_nanoseconds -from mlflow.tracing.constant import SpanAttributeKey, TokenUsageKey - -from core.ops.entities.trace_entity import ( - DatasetRetrievalTraceInfo, - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - SuggestedQuestionTraceInfo, - ToolTraceInfo, - WorkflowTraceInfo, -) -from graphon.enums import BuiltinNodeTypes - -# ── Helpers ────────────────────────────────────────────────────────────────── - - -def _dt() -> datetime: - return datetime(2024, 1, 1, 0, 0, 0, tzinfo=UTC) - - -def _make_workflow_trace_info(**overrides) -> WorkflowTraceInfo: - defaults = { - "workflow_id": "wf-id", - "tenant_id": "tenant", - "workflow_run_id": "run-1", - "workflow_run_elapsed_time": 1.0, - "workflow_run_status": "succeeded", - "workflow_run_inputs": {"key": "val"}, - "workflow_run_outputs": {"answer": "42"}, - "workflow_run_version": "v1", - "total_tokens": 10, - "file_list": [], - "query": "hello", - "metadata": {"app_id": "app-1", "user_id": "u1", "conversation_id": "c1"}, - "start_time": _dt(), - "end_time": _dt(), - } - defaults.update(overrides) - return WorkflowTraceInfo(**defaults) - - -def _make_message_trace_info(**overrides) -> MessageTraceInfo: - defaults = { - "conversation_model": "chat", - "message_tokens": 5, - "answer_tokens": 10, - "total_tokens": 15, - "conversation_mode": "chat", - "metadata": {"conversation_id": "c1", "from_account_id": "a1"}, - "message_id": "msg-1", - "message_data": SimpleNamespace( - model_provider="openai", - model_id="gpt-4", - total_price=0.01, - answer="response text", - ), - "inputs": {"prompt": "hi"}, - "outputs": "ok", - "start_time": _dt(), - "end_time": _dt(), - "error": None, - } - defaults.update(overrides) - return MessageTraceInfo(**defaults) - - -def _make_tool_trace_info(**overrides) -> ToolTraceInfo: - defaults = { - "tool_name": "my_tool", - "tool_inputs": {"x": 1}, - "tool_outputs": "output", - "tool_config": {"desc": "d"}, - "tool_parameters": {"p": "v"}, - "time_cost": 0.5, - "metadata": {"user_id": "u1"}, - "message_id": "msg-1", - "inputs": {"i": "v"}, - "outputs": {"o": "v"}, - "start_time": _dt(), - "end_time": _dt(), - "error": None, - } - defaults.update(overrides) - return ToolTraceInfo(**defaults) - - -def _make_moderation_trace_info(**overrides) -> ModerationTraceInfo: - defaults = { - "flagged": False, - "action": "allow", - "preset_response": "", - "query": "test", - "metadata": {"user_id": "u1"}, - "message_id": "msg-1", - } - defaults.update(overrides) - return ModerationTraceInfo(**defaults) - - -def _make_dataset_retrieval_trace_info(**overrides) -> DatasetRetrievalTraceInfo: - defaults = { - "metadata": {"user_id": "u1"}, - "message_id": "msg-1", - "message_data": SimpleNamespace(), - "inputs": "query", - "documents": [{"content": "doc"}], - "start_time": _dt(), - "end_time": _dt(), - } - defaults.update(overrides) - return DatasetRetrievalTraceInfo(**defaults) - - -def _make_suggested_question_trace_info(**overrides) -> SuggestedQuestionTraceInfo: - defaults = { - "suggested_question": ["q1", "q2"], - "level": "info", - "total_tokens": 5, - "metadata": {"user_id": "u1"}, - "message_id": "msg-1", - "message_data": SimpleNamespace(created_at=_dt(), updated_at=_dt()), - "inputs": {"i": 1}, - "start_time": _dt(), - "end_time": _dt(), - "error": None, - } - defaults.update(overrides) - return SuggestedQuestionTraceInfo(**defaults) - - -def _make_generate_name_trace_info(**overrides) -> GenerateNameTraceInfo: - defaults = { - "tenant_id": "t1", - "metadata": {"user_id": "u1"}, - "message_id": "msg-1", - "inputs": {"i": 1}, - "outputs": {"name": "test"}, - "start_time": _dt(), - "end_time": _dt(), - } - defaults.update(overrides) - return GenerateNameTraceInfo(**defaults) - - -def _make_node(**overrides): - """Create a mock workflow node execution row.""" - defaults = { - "id": "node-1", - "tenant_id": "t1", - "app_id": "app-1", - "title": "Node Title", - "node_type": BuiltinNodeTypes.CODE, - "status": "succeeded", - "inputs": '{"key": "value"}', - "outputs": '{"result": "ok"}', - "created_at": _dt(), - "elapsed_time": 1.0, - "process_data": None, - "execution_metadata": None, - } - defaults.update(overrides) - return SimpleNamespace(**defaults) - - -# ── Fixtures ───────────────────────────────────────────────────────────────── - - -@pytest.fixture -def mock_mlflow(): - with patch("dify_trace_mlflow.mlflow_trace.mlflow") as mock: - yield mock - - -@pytest.fixture -def mock_tracing(): - """Patch all MLflow tracing functions used by the module.""" - with ( - patch("dify_trace_mlflow.mlflow_trace.MLflowDataTrace._start_span_no_context") as mock_start, - patch("dify_trace_mlflow.mlflow_trace.update_current_trace") as mock_update, - patch("dify_trace_mlflow.mlflow_trace.set_span_in_context") as mock_set, - patch("dify_trace_mlflow.mlflow_trace.detach_span_from_context") as mock_detach, - ): - yield { - "start": mock_start, - "update": mock_update, - "set": mock_set, - "detach": mock_detach, - } - - -@pytest.fixture -def mock_db(): - with patch("dify_trace_mlflow.mlflow_trace.db") as mock: - yield mock - - -@pytest.fixture -def trace_instance(mock_mlflow): - """Create an MLflowDataTrace using a basic MLflowConfig (no auth).""" - config = MLflowConfig(tracking_uri="http://localhost:5000", experiment_id="0") - return MLflowDataTrace(config) - - -# ── datetime_to_nanoseconds ───────────────────────────────────────────────── - - -class TestDatetimeToNanoseconds: - def test_none_returns_none(self): - assert datetime_to_nanoseconds(None) is None - - def test_converts_datetime(self): - dt = datetime(2024, 1, 1, 0, 0, 0, tzinfo=UTC) - expected = int(dt.timestamp() * 1_000_000_000) - assert datetime_to_nanoseconds(dt) == expected - - -# ── __init__ / setup ───────────────────────────────────────────────────────── - - -class TestInit: - def test_mlflow_config_no_auth(self, mock_mlflow): - config = MLflowConfig(tracking_uri="http://localhost:5000", experiment_id="0") - trace = MLflowDataTrace(config) - mock_mlflow.set_tracking_uri.assert_not_called() - mock_mlflow.set_experiment.assert_not_called() - assert trace.get_project_url() == "http://localhost:5000/#/experiments/0/traces" - assert trace._tracking_uri == "http://localhost:5000" - - def test_mlflow_config_with_auth(self, mock_mlflow): - config = MLflowConfig( - tracking_uri="http://localhost:5000", - experiment_id="1", - username="user", - password="pass", - ) - MLflowDataTrace(config) - assert os.environ["MLFLOW_TRACKING_USERNAME"] == "user" - assert os.environ["MLFLOW_TRACKING_PASSWORD"] == "pass" - - def test_databricks_oauth(self, mock_mlflow): - config = DatabricksConfig( - host="https://db.com/", - experiment_id="42", - client_id="cid", - client_secret="csec", - ) - trace = MLflowDataTrace(config) - assert os.environ["DATABRICKS_HOST"] == "https://db.com/" - assert os.environ["DATABRICKS_CLIENT_ID"] == "cid" - assert os.environ["DATABRICKS_CLIENT_SECRET"] == "csec" - mock_mlflow.set_tracking_uri.assert_called_with("databricks") - assert trace._tracking_uri == "databricks" - # Trailing slash stripped - assert trace.get_project_url() == "https://db.com/ml/experiments/42/traces" - - def test_databricks_pat(self, mock_mlflow): - config = DatabricksConfig( - host="https://db.com", - experiment_id="1", - personal_access_token="pat", - ) - trace = MLflowDataTrace(config) - assert os.environ["DATABRICKS_TOKEN"] == "pat" - assert "db.com/ml/experiments/1/traces" in trace.get_project_url() - - def test_databricks_no_creds_raises(self, mock_mlflow): - config = DatabricksConfig(host="https://db.com", experiment_id="1") - with pytest.raises(ValueError, match="Either Databricks token"): - MLflowDataTrace(config) - - -# ── trace dispatcher ──────────────────────────────────────────────────────── - - -class TestTraceDispatcher: - def test_dispatches_workflow(self, trace_instance, mock_tracing, mock_db): - with patch.object(trace_instance, "workflow_trace") as mock_wt: - trace_instance.trace(_make_workflow_trace_info()) - mock_wt.assert_called_once() - - def test_dispatches_message(self, trace_instance, mock_tracing, mock_db): - with patch.object(trace_instance, "message_trace") as mock_mt: - trace_instance.trace(_make_message_trace_info()) - mock_mt.assert_called_once() - - def test_dispatches_tool(self, trace_instance, mock_tracing, mock_db): - with patch.object(trace_instance, "tool_trace") as mock_tt: - trace_instance.trace(_make_tool_trace_info()) - mock_tt.assert_called_once() - - def test_dispatches_moderation(self, trace_instance, mock_tracing, mock_db): - with patch.object(trace_instance, "moderation_trace") as mock_mod: - trace_instance.trace(_make_moderation_trace_info(message_data=SimpleNamespace(created_at=_dt()))) - mock_mod.assert_called_once() - - def test_dispatches_dataset_retrieval(self, trace_instance, mock_tracing, mock_db): - with patch.object(trace_instance, "dataset_retrieval_trace") as mock_dr: - trace_instance.trace(_make_dataset_retrieval_trace_info()) - mock_dr.assert_called_once() - - def test_dispatches_suggested_question(self, trace_instance, mock_tracing, mock_db): - with patch.object(trace_instance, "suggested_question_trace") as mock_sq: - trace_instance.trace(_make_suggested_question_trace_info()) - mock_sq.assert_called_once() - - def test_dispatches_generate_name(self, trace_instance, mock_tracing, mock_db): - with patch.object(trace_instance, "generate_name_trace") as mock_gn: - trace_instance.trace(_make_generate_name_trace_info()) - mock_gn.assert_called_once() - - def test_reraises_exception(self, trace_instance, mock_tracing, mock_db): - with patch.object(trace_instance, "workflow_trace", side_effect=RuntimeError("boom")): - with pytest.raises(RuntimeError, match="boom"): - trace_instance.trace(_make_workflow_trace_info()) - - -# ── workflow_trace ─────────────────────────────────────────────────────────── - - -class TestWorkflowTrace: - def test_basic_workflow_no_nodes(self, trace_instance, mock_tracing, mock_db): - mock_db.session.scalars.return_value.all.return_value = [] - span = MagicMock() - mock_tracing["start"].return_value = span - mock_tracing["set"].return_value = "token" - - trace_info = _make_workflow_trace_info(conversation_id="sess-1") - trace_instance.workflow_trace(trace_info) - - # Workflow span started and ended - mock_tracing["start"].assert_called_once() - span.end.assert_called_once() - - def test_workflow_filters_sys_inputs_and_adds_query(self, trace_instance, mock_tracing, mock_db): - mock_db.session.scalars.return_value.all.return_value = [] - span = MagicMock() - mock_tracing["start"].return_value = span - mock_tracing["set"].return_value = "token" - - trace_info = _make_workflow_trace_info( - workflow_run_inputs={"sys.app_id": "x", "user_input": "hi"}, - query="hello", - ) - trace_instance.workflow_trace(trace_info) - - call_kwargs = mock_tracing["start"].call_args - inputs = call_kwargs.kwargs["inputs"] - assert "sys.app_id" not in inputs - assert inputs["user_input"] == "hi" - assert inputs["query"] == "hello" - - def test_workflow_with_llm_node(self, trace_instance, mock_tracing, mock_db): - usage = {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15} - llm_node = _make_node( - node_type=BuiltinNodeTypes.LLM, - process_data=json.dumps( - { - "prompts": [{"role": "user", "text": "hi"}], - "model_name": "gpt-4", - "model_provider": "openai", - "finish_reason": "stop", - "usage": usage, - } - ), - outputs='{"text": "hello world"}', - ) - mock_db.session.scalars.return_value.all.return_value = [llm_node] - - workflow_span = MagicMock() - node_span = MagicMock() - mock_tracing["start"].side_effect = [workflow_span, node_span] - mock_tracing["set"].return_value = "token" - - trace_instance.workflow_trace(_make_workflow_trace_info()) - assert mock_tracing["start"].call_count == 2 - node_start_call = mock_tracing["start"].call_args_list[1] - attrs = node_start_call.kwargs["attributes"] - assert attrs[SpanAttributeKey.CHAT_USAGE] == { - TokenUsageKey.INPUT_TOKENS: 5, - TokenUsageKey.OUTPUT_TOKENS: 10, - TokenUsageKey.TOTAL_TOKENS: 15, - } - assert attrs["usage"] == usage - node_span.end.assert_called_once() - workflow_span.end.assert_called_once() - - def test_workflow_with_question_classifier_node(self, trace_instance, mock_tracing, mock_db): - qc_node = _make_node( - node_type=BuiltinNodeTypes.QUESTION_CLASSIFIER, - process_data=json.dumps( - { - "prompts": "classify this", - "model_name": "gpt-4", - "model_provider": "openai", - } - ), - ) - mock_db.session.scalars.return_value.all.return_value = [qc_node] - workflow_span = MagicMock() - node_span = MagicMock() - mock_tracing["start"].side_effect = [workflow_span, node_span] - mock_tracing["set"].return_value = "token" - - trace_instance.workflow_trace(_make_workflow_trace_info()) - assert mock_tracing["start"].call_count == 2 - - def test_workflow_with_http_request_node(self, trace_instance, mock_tracing, mock_db): - http_node = _make_node( - node_type=BuiltinNodeTypes.HTTP_REQUEST, - process_data='{"url": "https://api.com"}', - ) - mock_db.session.scalars.return_value.all.return_value = [http_node] - workflow_span = MagicMock() - node_span = MagicMock() - mock_tracing["start"].side_effect = [workflow_span, node_span] - mock_tracing["set"].return_value = "token" - - trace_instance.workflow_trace(_make_workflow_trace_info()) - # HTTP_REQUEST uses process_data as inputs - node_start_call = mock_tracing["start"].call_args_list[1] - assert node_start_call.kwargs["inputs"] == '{"url": "https://api.com"}' - - def test_workflow_with_knowledge_retrieval_node(self, trace_instance, mock_tracing, mock_db): - kr_node = _make_node( - node_type=BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL, - outputs=json.dumps( - { - "result": [ - {"content": "doc1", "metadata": {"source": "s1"}}, - {"content": "doc2", "metadata": {}}, - ] - } - ), - ) - mock_db.session.scalars.return_value.all.return_value = [kr_node] - workflow_span = MagicMock() - node_span = MagicMock() - mock_tracing["start"].side_effect = [workflow_span, node_span] - mock_tracing["set"].return_value = "token" - - trace_instance.workflow_trace(_make_workflow_trace_info()) - # outputs should be parsed to Document objects - end_call = node_span.end.call_args - outputs = end_call.kwargs["outputs"] - assert len(outputs) == 2 - - def test_workflow_with_failed_node(self, trace_instance, mock_tracing, mock_db): - failed_node = _make_node(status="failed") - mock_db.session.scalars.return_value.all.return_value = [failed_node] - workflow_span = MagicMock() - node_span = MagicMock() - mock_tracing["start"].side_effect = [workflow_span, node_span] - mock_tracing["set"].return_value = "token" - - trace_instance.workflow_trace(_make_workflow_trace_info()) - node_span.set_status.assert_called_once() - node_span.add_event.assert_called_once() - - def test_workflow_with_workflow_error(self, trace_instance, mock_tracing, mock_db): - mock_db.session.scalars.return_value.all.return_value = [] - workflow_span = MagicMock() - mock_tracing["start"].return_value = workflow_span - mock_tracing["set"].return_value = "token" - - trace_info = _make_workflow_trace_info(error="workflow failed") - trace_instance.workflow_trace(trace_info) - workflow_span.set_status.assert_called_once() - workflow_span.add_event.assert_called_once() - # Still ends the span via finally - workflow_span.end.assert_called_once() - - def test_workflow_node_no_inputs_no_outputs(self, trace_instance, mock_tracing, mock_db): - node = _make_node(inputs=None, outputs=None) - mock_db.session.scalars.return_value.all.return_value = [node] - workflow_span = MagicMock() - node_span = MagicMock() - mock_tracing["start"].side_effect = [workflow_span, node_span] - mock_tracing["set"].return_value = "token" - - trace_instance.workflow_trace(_make_workflow_trace_info()) - node_call = mock_tracing["start"].call_args_list[1] - assert node_call.kwargs["inputs"] == {} - end_call = node_span.end.call_args - assert end_call.kwargs["outputs"] == {} - - def test_workflow_no_user_id_no_conversation_id(self, trace_instance, mock_tracing, mock_db): - mock_db.session.scalars.return_value.all.return_value = [] - span = MagicMock() - mock_tracing["start"].return_value = span - mock_tracing["set"].return_value = "token" - - trace_info = _make_workflow_trace_info( - metadata={"app_id": "app-1"}, - conversation_id=None, - ) - trace_instance.workflow_trace(trace_info) - # _set_trace_metadata still called with empty metadata - mock_tracing["update"].assert_called_once() - - def test_workflow_empty_query(self, trace_instance, mock_tracing, mock_db): - """When query is empty string, it's falsy so no query key added.""" - mock_db.session.scalars.return_value.all.return_value = [] - span = MagicMock() - mock_tracing["start"].return_value = span - mock_tracing["set"].return_value = "token" - - trace_info = _make_workflow_trace_info(query="") - trace_instance.workflow_trace(trace_info) - call_kwargs = mock_tracing["start"].call_args - inputs = call_kwargs.kwargs["inputs"] - assert "query" not in inputs - - -# ── _parse_llm_inputs_and_attributes ───────────────────────────────────────── - - -class TestParseLlmInputsAndAttributes: - def test_none_process_data(self, trace_instance): - node = _make_node(process_data=None) - inputs, attrs = trace_instance._parse_llm_inputs_and_attributes(node) - assert inputs == {} - assert attrs == {} - - def test_invalid_json(self, trace_instance): - node = _make_node(process_data="not json") - inputs, attrs = trace_instance._parse_llm_inputs_and_attributes(node) - assert inputs == {} - assert attrs == {} - - def test_valid_process_data_with_usage(self, trace_instance): - node = _make_node( - process_data=json.dumps( - { - "prompts": [{"role": "user", "text": "hi"}], - "model_name": "gpt-4", - "model_provider": "openai", - "finish_reason": "stop", - "usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15}, - } - ) - ) - inputs, attrs = trace_instance._parse_llm_inputs_and_attributes(node) - assert isinstance(inputs, list) - assert attrs["model_name"] == "gpt-4" - assert "usage" in attrs - - def test_valid_process_data_without_usage(self, trace_instance): - node = _make_node( - process_data=json.dumps( - { - "prompts": "simple prompt", - "model_name": "gpt-3.5", - } - ) - ) - inputs, attrs = trace_instance._parse_llm_inputs_and_attributes(node) - assert inputs == "simple prompt" - assert attrs["model_name"] == "gpt-3.5" - - -# ── _parse_knowledge_retrieval_outputs ─────────────────────────────────────── - - -class TestParseKnowledgeRetrievalOutputs: - def test_with_results(self, trace_instance): - outputs = {"result": [{"content": "c1", "metadata": {"s": "1"}}]} - docs = trace_instance._parse_knowledge_retrieval_outputs(outputs) - assert len(docs) == 1 - assert docs[0].page_content == "c1" - - def test_empty_result(self, trace_instance): - outputs = {"result": []} - result = trace_instance._parse_knowledge_retrieval_outputs(outputs) - assert result == outputs - - def test_no_result_key(self, trace_instance): - outputs = {"other": "data"} - result = trace_instance._parse_knowledge_retrieval_outputs(outputs) - assert result == outputs - - def test_result_not_list(self, trace_instance): - outputs = {"result": "not a list"} - result = trace_instance._parse_knowledge_retrieval_outputs(outputs) - assert result == outputs - - -# ── message_trace ──────────────────────────────────────────────────────────── - - -class TestMessageTrace: - def test_returns_early_if_no_message_data(self, trace_instance, mock_tracing, mock_db): - trace_info = _make_message_trace_info(message_data=None) - trace_instance.message_trace(trace_info) - mock_tracing["start"].assert_not_called() - - def test_basic_message_trace(self, trace_instance, mock_tracing, mock_db): - span = MagicMock() - mock_tracing["start"].return_value = span - mock_tracing["set"].return_value = "token" - - trace_instance.message_trace(_make_message_trace_info()) - mock_tracing["start"].assert_called_once() - span.end.assert_called_once() - - def test_message_trace_with_error(self, trace_instance, mock_tracing, mock_db): - span = MagicMock() - mock_tracing["start"].return_value = span - mock_tracing["set"].return_value = "token" - - trace_info = _make_message_trace_info(error="something broke") - trace_instance.message_trace(trace_info) - span.set_status.assert_called_once() - span.add_event.assert_called_once() - - def test_message_trace_with_file_data(self, trace_instance, mock_tracing, mock_db, monkeypatch: pytest.MonkeyPatch): - span = MagicMock() - mock_tracing["start"].return_value = span - mock_tracing["set"].return_value = "token" - monkeypatch.setenv("FILES_URL", "http://files.test") - - file_data = SimpleNamespace(url="path/to/file.png") - trace_info = _make_message_trace_info( - message_file_data=file_data, - file_list=["existing_file.txt"], - ) - trace_instance.message_trace(trace_info) - call_kwargs = mock_tracing["start"].call_args - attrs = call_kwargs.kwargs["attributes"] - assert "http://files.test/path/to/file.png" in attrs["file_list"] - assert "existing_file.txt" in attrs["file_list"] - - def test_message_trace_preserves_structured_span_attributes(self, trace_instance, mock_tracing, mock_db): - span = MagicMock() - mock_tracing["start"].return_value = span - mock_tracing["set"].return_value = "token" - - trace_info = _make_message_trace_info( - metadata={ - "conversation_id": "c1", - "from_account_id": "a1", - "routing": {"node": "answer", "score": 0.7}, - }, - file_list=["existing_file.txt"], - ) - trace_instance.message_trace(trace_info) - - attrs = mock_tracing["start"].call_args.kwargs["attributes"] - assert attrs["message_id"] == "msg-1" - assert attrs["total_price"] == 0.01 - assert attrs["routing"] == {"node": "answer", "score": 0.7} - assert attrs["file_list"] == ["existing_file.txt"] - - def test_message_trace_file_list_none(self, trace_instance, mock_tracing, mock_db): - span = MagicMock() - mock_tracing["start"].return_value = span - mock_tracing["set"].return_value = "token" - - trace_info = _make_message_trace_info(file_list=None, message_file_data=None) - trace_instance.message_trace(trace_info) - mock_tracing["start"].assert_called_once() - - def test_message_trace_with_end_user(self, trace_instance, mock_tracing, mock_db): - span = MagicMock() - mock_tracing["start"].return_value = span - mock_tracing["set"].return_value = "token" - - end_user = MagicMock() - end_user.session_id = "session-xyz" - - trace_info = _make_message_trace_info( - metadata={"from_end_user_id": "eu-1", "conversation_id": "c1"}, - ) - trace_instance.message_trace(trace_info) - # update_current_trace called with user id from EndUser - mock_tracing["update"].assert_called_once() - - def test_message_trace_with_no_conversation_id(self, trace_instance, mock_tracing, mock_db): - span = MagicMock() - mock_tracing["start"].return_value = span - mock_tracing["set"].return_value = "token" - - trace_info = _make_message_trace_info( - metadata={"from_account_id": "acc-1"}, - ) - trace_instance.message_trace(trace_info) - mock_tracing["update"].assert_called_once() - - -# ── _get_message_user_id ───────────────────────────────────────────────────── - - -class TestGetMessageUserId: - def test_returns_end_user_session_id(self, trace_instance, mock_db): - end_user = MagicMock() - end_user.session_id = "session-1" - mock_db.session.get.return_value = end_user - result = trace_instance._get_message_user_id({"from_end_user_id": "eu-1"}) - assert result == "session-1" - - def test_returns_account_id_when_no_end_user(self, trace_instance, mock_db): - mock_db.session.get.return_value = None - result = trace_instance._get_message_user_id({"from_end_user_id": "eu-1", "from_account_id": "acc-1"}) - assert result == "acc-1" - - def test_returns_account_id_when_no_end_user_id(self, trace_instance, mock_db): - result = trace_instance._get_message_user_id({"from_account_id": "acc-1"}) - assert result == "acc-1" - - def test_returns_none_when_nothing(self, trace_instance, mock_db): - result = trace_instance._get_message_user_id({}) - assert result is None - - -# ── tool_trace ─────────────────────────────────────────────────────────────── - - -class TestToolTrace: - def test_basic_tool_trace(self, trace_instance, mock_tracing): - span = MagicMock() - mock_tracing["start"].return_value = span - - trace_instance.tool_trace(_make_tool_trace_info()) - mock_tracing["start"].assert_called_once() - span.end.assert_called_once() - span.set_status.assert_not_called() - - def test_tool_trace_with_error(self, trace_instance, mock_tracing): - span = MagicMock() - mock_tracing["start"].return_value = span - - trace_instance.tool_trace(_make_tool_trace_info(error="tool failed")) - span.set_status.assert_called_once() - span.add_event.assert_called_once() - span.end.assert_called_once() - - -# ── moderation_trace ───────────────────────────────────────────────────────── - - -class TestModerationTrace: - def test_returns_early_if_no_message_data(self, trace_instance, mock_tracing): - trace_info = _make_moderation_trace_info(message_data=None) - trace_instance.moderation_trace(trace_info) - mock_tracing["start"].assert_not_called() - - def test_basic_moderation_trace(self, trace_instance, mock_tracing): - span = MagicMock() - mock_tracing["start"].return_value = span - - trace_info = _make_moderation_trace_info( - message_data=SimpleNamespace(created_at=_dt()), - start_time=_dt(), - end_time=_dt(), - ) - trace_instance.moderation_trace(trace_info) - mock_tracing["start"].assert_called_once() - span.end.assert_called_once() - end_kwargs = span.end.call_args.kwargs["outputs"] - assert end_kwargs["action"] == "allow" - assert end_kwargs["flagged"] is False - - def test_moderation_uses_message_data_created_at_if_no_start_time(self, trace_instance, mock_tracing): - span = MagicMock() - mock_tracing["start"].return_value = span - - trace_info = _make_moderation_trace_info( - message_data=SimpleNamespace(created_at=_dt()), - start_time=None, - end_time=_dt(), - ) - trace_instance.moderation_trace(trace_info) - mock_tracing["start"].assert_called_once() - - -# ── dataset_retrieval_trace ────────────────────────────────────────────────── - - -class TestDatasetRetrievalTrace: - def test_returns_early_if_no_message_data(self, trace_instance, mock_tracing): - trace_info = _make_dataset_retrieval_trace_info(message_data=None) - trace_instance.dataset_retrieval_trace(trace_info) - mock_tracing["start"].assert_not_called() - - def test_basic_dataset_retrieval_trace(self, trace_instance, mock_tracing): - span = MagicMock() - mock_tracing["start"].return_value = span - - trace_instance.dataset_retrieval_trace(_make_dataset_retrieval_trace_info()) - mock_tracing["start"].assert_called_once() - span.end.assert_called_once() - - -# ── suggested_question_trace ───────────────────────────────────────────────── - - -class TestSuggestedQuestionTrace: - def test_returns_early_if_no_message_data(self, trace_instance, mock_tracing): - trace_info = _make_suggested_question_trace_info(message_data=None) - trace_instance.suggested_question_trace(trace_info) - mock_tracing["start"].assert_not_called() - - def test_basic_suggested_question_trace(self, trace_instance, mock_tracing): - span = MagicMock() - mock_tracing["start"].return_value = span - - trace_instance.suggested_question_trace(_make_suggested_question_trace_info()) - mock_tracing["start"].assert_called_once() - span.end.assert_called_once() - - def test_suggested_question_with_error(self, trace_instance, mock_tracing): - span = MagicMock() - mock_tracing["start"].return_value = span - - trace_info = _make_suggested_question_trace_info(error="failed") - trace_instance.suggested_question_trace(trace_info) - span.set_status.assert_called_once() - span.add_event.assert_called_once() - - def test_uses_message_data_times_when_no_start_end(self, trace_instance, mock_tracing): - span = MagicMock() - mock_tracing["start"].return_value = span - - trace_info = _make_suggested_question_trace_info( - start_time=None, - end_time=None, - ) - trace_instance.suggested_question_trace(trace_info) - mock_tracing["start"].assert_called_once() - span.end.assert_called_once() - - -# ── generate_name_trace ────────────────────────────────────────────────────── - - -class TestGenerateNameTrace: - def test_basic_generate_name_trace(self, trace_instance, mock_tracing): - span = MagicMock() - mock_tracing["start"].return_value = span - - trace_instance.generate_name_trace(_make_generate_name_trace_info()) - mock_tracing["start"].assert_called_once() - span.end.assert_called_once() - - -# ── _get_workflow_nodes ────────────────────────────────────────────────────── - - -class TestGetWorkflowNodes: - def test_queries_db(self, trace_instance, mock_db): - mock_db.session.scalars.return_value.all.return_value = ["n1", "n2"] - result = trace_instance._get_workflow_nodes(_make_workflow_trace_info()) - assert result == ["n1", "n2"] - - @pytest.mark.parametrize("metadata", [{}, {"app_id": ""}, {"app_id": 123}], ids=["missing", "empty", "non-string"]) - def test_rejects_invalid_app_id_before_querying_nodes(self, metadata, trace_instance, mock_db): - with pytest.raises(ValueError, match="No app_id found in workflow trace metadata"): - trace_instance._get_workflow_nodes(_make_workflow_trace_info(metadata=metadata)) - - mock_db.session.scalars.assert_not_called() - - -# ── _get_node_span_type ───────────────────────────────────────────────────── - - -class TestGetNodeSpanType: - @pytest.mark.parametrize( - ("node_type", "expected_contains"), - [ - (BuiltinNodeTypes.LLM, "LLM"), - (BuiltinNodeTypes.QUESTION_CLASSIFIER, "LLM"), - (BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL, "RETRIEVER"), - (BuiltinNodeTypes.TOOL, "TOOL"), - (BuiltinNodeTypes.CODE, "TOOL"), - (BuiltinNodeTypes.HTTP_REQUEST, "TOOL"), - (BuiltinNodeTypes.AGENT, "AGENT"), - ], - ) - def test_mapped_types(self, trace_instance, node_type, expected_contains): - result = trace_instance._get_node_span_type(node_type) - assert expected_contains in str(result) - - def test_unknown_type_returns_chain(self, trace_instance): - result = trace_instance._get_node_span_type("unknown_node") - assert result == "CHAIN" - - -# ── _set_trace_metadata ───────────────────────────────────────────────────── - - -class TestSetTraceMetadata: - def test_sets_and_detaches(self, trace_instance, mock_tracing): - span = MagicMock() - mock_tracing["set"].return_value = "token" - - trace_instance._set_trace_metadata(span, {"key": "val"}) - mock_tracing["set"].assert_called_once_with(span) - mock_tracing["update"].assert_called_once_with(metadata={"key": "val"}) - mock_tracing["detach"].assert_called_once_with("token") - - def test_detaches_even_on_error(self, trace_instance, mock_tracing): - span = MagicMock() - mock_tracing["set"].return_value = "token" - mock_tracing["update"].side_effect = RuntimeError("fail") - - with pytest.raises(RuntimeError): - trace_instance._set_trace_metadata(span, {}) - mock_tracing["detach"].assert_called_once_with("token") - - def test_no_detach_when_token_is_none(self, trace_instance, mock_tracing): - span = MagicMock() - mock_tracing["set"].return_value = None - - trace_instance._set_trace_metadata(span, {}) - mock_tracing["detach"].assert_not_called() - - -# ── _parse_prompts ─────────────────────────────────────────────────────────── - - -class TestParsePrompts: - def test_string_input(self, trace_instance): - assert trace_instance._parse_prompts("hello") == "hello" - - def test_dict_input(self, trace_instance): - result = trace_instance._parse_prompts({"role": "user", "text": "hi"}) - assert result == {"role": "user", "content": "hi"} - - def test_list_input(self, trace_instance): - prompts = [ - {"role": "user", "text": "hi"}, - {"role": "assistant", "text": "hello"}, - ] - result = trace_instance._parse_prompts(prompts) - assert len(result) == 2 - assert result[0]["role"] == "user" - - def test_none_input(self, trace_instance): - assert trace_instance._parse_prompts(None) is None - - def test_int_passthrough(self, trace_instance): - assert trace_instance._parse_prompts(42) == 42 - - -# ── _parse_single_message ─────────────────────────────────────────────────── - - -class TestParseSingleMessage: - def test_basic_message(self, trace_instance): - result = trace_instance._parse_single_message({"role": "user", "text": "hello"}) - assert result == {"role": "user", "content": "hello"} - - def test_default_role(self, trace_instance): - result = trace_instance._parse_single_message({"text": "hello"}) - assert result["role"] == "user" - - def test_with_tool_calls(self, trace_instance): - item = { - "role": "assistant", - "text": "", - "tool_calls": [{"id": "tc1", "function": {"name": "fn"}}], - } - result = trace_instance._parse_single_message(item) - assert "tool_calls" in result - - def test_tool_role_ignores_tool_calls(self, trace_instance): - item = { - "role": "tool", - "text": "result", - "tool_calls": [{"id": "tc1"}], - } - result = trace_instance._parse_single_message(item) - assert "tool_calls" not in result - - def test_with_files(self, trace_instance): - item = {"role": "user", "text": "look", "files": ["f1.png"]} - result = trace_instance._parse_single_message(item) - assert result["files"] == ["f1.png"] - - def test_no_files(self, trace_instance): - result = trace_instance._parse_single_message({"role": "user", "text": "hi"}) - assert "files" not in result - - -# ── _resolve_tool_call_ids ─────────────────────────────────────────────────── - - -class TestResolveToolCallIds: - def test_resolves_tool_call_ids(self, trace_instance): - messages = [ - { - "role": "assistant", - "content": "", - "tool_calls": [{"id": "tc1"}, {"id": "tc2"}], - }, - {"role": "tool", "content": "result1"}, - {"role": "tool", "content": "result2"}, - ] - result = trace_instance._resolve_tool_call_ids(messages) - assert result[1]["tool_call_id"] == "tc1" - assert result[2]["tool_call_id"] == "tc2" - - def test_no_tool_calls(self, trace_instance): - messages = [ - {"role": "user", "content": "hi"}, - {"role": "assistant", "content": "hello"}, - ] - result = trace_instance._resolve_tool_call_ids(messages) - assert "tool_call_id" not in result[0] - assert "tool_call_id" not in result[1] - - def test_tool_message_no_ids_available(self, trace_instance): - """Tool message with no preceding tool_calls should not crash.""" - messages = [ - {"role": "tool", "content": "result"}, - ] - result = trace_instance._resolve_tool_call_ids(messages) - assert "tool_call_id" not in result[0] - - -# ── api_check ──────────────────────────────────────────────────────────────── - - -class TestApiCheck: - def test_success(self, trace_instance, mock_mlflow): - mock_mlflow.MlflowClient.return_value.search_experiments.return_value = [] - assert trace_instance.api_check() is True - - def test_failure(self, trace_instance, mock_mlflow): - mock_mlflow.MlflowClient.return_value.search_experiments.side_effect = ConnectionError("refused") - with pytest.raises(ValueError, match="MLflow connection failed"): - trace_instance.api_check() - - -# ── get_project_url ────────────────────────────────────────────────────────── - - -class TestGetProjectUrl: - def test_returns_url(self, trace_instance): - assert "experiments" in trace_instance.get_project_url() diff --git a/api/providers/trace/trace-mlflow/tests/unit_tests/src/dify_trace_mlflow/test_provider.py b/api/providers/trace/trace-mlflow/tests/unit_tests/src/dify_trace_mlflow/test_provider.py new file mode 100644 index 00000000000000..9624c813690de2 --- /dev/null +++ b/api/providers/trace/trace-mlflow/tests/unit_tests/src/dify_trace_mlflow/test_provider.py @@ -0,0 +1,865 @@ +import ast +import base64 +import json +import os +from pathlib import Path +from typing import Any +from urllib.parse import parse_qs, urlsplit + +import dify_trace_mlflow +import pytest +from dify_trace_mlflow.config import DatabricksConfig, MLflowConfig +from dify_trace_mlflow.provider import DatabricksProvider, MLflowProvider +from dify_trace_sdk.egress import EgressResponse +from dify_trace_sdk.errors import ( + EgressBlockedError, + EgressError, + ExportOutcomeUnknownError, + ExportRejectedError, + ExportRetryableError, + VerificationError, +) +from dify_trace_sdk.ids import trace_id_for +from dify_trace_sdk.model import Cost, Span, TokenUsage, Trace +from dify_trace_sdk.provider import SendResult +from dify_trace_sdk.testing import ( + MESSAGE_ID, + WORKFLOW_RUN_ID, + RecordedRequest, + RecordingEgress, + message_trace, + response, + tool_fragment, + workflow_trace, +) +from mlflow.entities import Span as MlflowSpan +from mlflow.entities import SpanStatusCode, TraceData +from mlflow.tracing.constant import SpanAttributeKey +from mlflow.tracing.otel.translation import translate_span_when_storing +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( + ExportTraceServiceRequest, + ExportTraceServiceResponse, +) +from pydantic import JsonValue + +MLFLOW_CONFIG = MLflowConfig( + tracking_uri="https://mlflow.example/proxy", experiment_id="42", username="admin", password="s3cret" +) +BASIC_AUTH = "Basic " + base64.b64encode(b"admin:s3cret").decode() +ACK = EgressResponse(status_code=200, headers={"content-type": "application/x-protobuf"}, content=b"") +MLFLOW_EXPERIMENT = {"experiment": {"experiment_id": "42", "name": "Dify", "lifecycle_stage": "active"}} + + +def _mlflow_provider(*responses: EgressResponse | Exception) -> tuple[MLflowProvider, RecordingEgress]: + egress = RecordingEgress(responses, grpc_target=None) + return MLflowProvider(egress), egress + + +def _otlp(egress: RecordingEgress) -> ExportTraceServiceRequest: + (sent,) = egress.requests + assert sent.content is not None + return ExportTraceServiceRequest.FromString(sent.content) + + +def _exported(egress: RecordingEgress) -> list[MlflowSpan]: + """Read the export the way the MLflow server's OTLP endpoint does.""" + return [MlflowSpan.from_otel_proto(span) for span in _otlp(egress).resource_spans[0].scope_spans[0].spans] + + +def _stored_attribute(span: MlflowSpan, key: str) -> JsonValue: + """An attribute as the MLflow server's tracking store keeps it.""" + return json.loads(translate_span_when_storing(span)["attributes"][key]) + + +def _with_span(trace: Trace, index: int, span: Span) -> Trace: + return trace.model_copy(update={"spans": (*trace.spans[:index], span, *trace.spans[index + 1 :])}) + + +def _ack(rejected_spans: int) -> EgressResponse: + ack = ExportTraceServiceResponse() + ack.partial_success.rejected_spans = rejected_spans + return EgressResponse( + status_code=200, headers={"content-type": "application/x-protobuf"}, content=ack.SerializeToString() + ) + + +FLAT_TOOL = { + "type": "function", + "name": "calculate", + "description": "Add numbers", + "parameters": {"type": "object", "properties": {"value": {"type": "integer"}}, "$comment": "kept"}, + "strict": True, +} +NESTED_TOOL = {"type": "function", "function": {"name": "already_nested", "description": "Nested definition"}} + + +class TestMLflowProvider: + @staticmethod + def test_verify_reads_the_experiment_and_normalizes_the_destination() -> None: + provider, egress = _mlflow_provider(response(200, MLFLOW_EXPERIMENT)) + config = MLflowConfig( + tracking_uri="https://MLflow.Example/proxy/", experiment_id="042", username="admin", password="s3cret" + ) + + verified = provider.verify(config) + + assert verified.tracking_uri == "https://mlflow.example/proxy" + assert verified.experiment_id == "42" + assert (verified.username, verified.password) == ("admin", "s3cret") + (sent,) = egress.requests + assert sent.method == "GET" + assert sent.url == "https://mlflow.example/proxy/api/2.0/mlflow/experiments/get?experiment_id=42" + assert sent.headers == {"Accept": "application/json", "Authorization": BASIC_AUTH} + assert sent.content is None + + @staticmethod + def test_verify_without_credentials_sends_no_authorization() -> None: + provider, egress = _mlflow_provider(response(200, MLFLOW_EXPERIMENT)) + provider.verify(MLflowConfig(tracking_uri="http://localhost:5000", experiment_id="42")) + (sent,) = egress.requests + assert sent.headers == {"Accept": "application/json"} + + @staticmethod + @pytest.mark.parametrize( + ("outcome", "code"), + [ + (response(401, {}), "authentication_failed"), + (response(403, {}), "permission_denied"), + (response(404, {"error_code": "RESOURCE_DOES_NOT_EXIST"}), "experiment_not_found_or_inaccessible"), + (response(503, {}), "endpoint_unavailable"), + (response(302, {}), "endpoint_unavailable"), + (response(200, b""), "malformed_response"), + (response(200, ["not", "an", "object"]), "malformed_response"), + (response(200, {}), "malformed_response"), + (response(200, {"experiment": {"experiment_id": ["42"]}}), "malformed_response"), + ( + response(200, {"experiment": {"experiment_id": "7", "lifecycle_stage": "active"}}), + "experiment_id_mismatch", + ), + ( + response(200, {"experiment": {"experiment_id": "42", "lifecycle_stage": "deleted"}}), + "experiment_not_active", + ), + (EgressError("connection refused", request_sent=False), "connection_failed"), + (EgressBlockedError("private address"), "connection_failed"), + ], + ) + def test_verify_failures_have_stable_codes(outcome: EgressResponse | Exception, code: str) -> None: + provider, _ = _mlflow_provider(outcome) + with pytest.raises(VerificationError) as caught: + provider.verify(MLFLOW_CONFIG) + assert caught.value.code == code + + @staticmethod + @pytest.mark.parametrize( + ("values", "code"), + [ + ({"tracking_uri": "https://mlflow.example/a/../b"}, "invalid_tracking_uri"), + ({"tracking_uri": "https://mlflow.example/a%2Fb"}, "invalid_tracking_uri"), + ({"tracking_uri": "https://mlflow.example//proxy"}, "invalid_tracking_uri"), + ({"tracking_uri": "https://mlflow.example/proxy?token=1"}, "invalid_tracking_uri"), + ({"tracking_uri": "https://mlflow.example/my proxy"}, "invalid_tracking_uri"), + ({"tracking_uri": "https://mlflow.example:0"}, "invalid_tracking_uri"), + ({"password": None}, "incomplete_basic_auth"), + ({"username": None}, "incomplete_basic_auth"), + ({"username": " "}, "incomplete_basic_auth"), + ({"username": "ad:min"}, "incomplete_basic_auth"), + ({"password": "s3cret\n"}, "incomplete_basic_auth"), + ], + ) + def test_verify_rejects_unusable_configurations_without_requests(values: dict[str, str | None], code: str) -> None: + provider, egress = _mlflow_provider() + config = MLflowConfig.model_validate({**MLFLOW_CONFIG.model_dump(), **values}) + with pytest.raises(VerificationError) as caught: + provider.verify(config) + assert caught.value.code == code + assert egress.requests == [] + + @staticmethod + def test_project_url_is_normalized_without_requests() -> None: + provider, egress = _mlflow_provider() + config = MLflowConfig(tracking_uri="https://MLflow.Example/proxy/", experiment_id="042") + assert provider.project_url(config) == "https://mlflow.example/proxy/#/experiments/42/traces" + assert provider.project_url(MLflowConfig(tracking_uri="https://mlflow.example/a/../b")) is None + assert egress.requests == [] + + @staticmethod + def test_message_trace_is_exported_as_mlflow_spans_over_otlp() -> None: + provider, egress = _mlflow_provider(ACK) + trace = message_trace() + + result = provider.send(MLFLOW_CONFIG, trace) + + (sent,) = egress.requests + assert sent.method == "POST" + assert sent.url == "https://mlflow.example/proxy/v1/traces" + assert sent.headers["Content-Type"] == "application/x-protobuf" + assert sent.headers["x-mlflow-experiment-id"] == "42" + assert sent.headers["Authorization"] == BASIC_AUTH + assert not sent.sensitive_url + (otlp_span, _) = _otlp(egress).resource_spans[0].scope_spans[0].spans + assert otlp_span.trace_id.hex() == trace.trace_id + root, llm = _exported(egress) + request_id = f"tr-{trace.trace_id}" + assert root.trace_id == request_id + assert llm.trace_id == request_id + assert (root.span_id, root.parent_id) == (trace.spans[0].span_id, None) + assert (llm.span_id, llm.parent_id) == (trace.spans[1].span_id, trace.spans[0].span_id) + assert root.span_type == "CHAIN" + assert llm.span_type == "LLM" + assert (root.start_time_ns, root.end_time_ns) == (1790841600000000000, 1790841602000000000) + assert root.status.status_code == SpanStatusCode.OK + assert root.inputs == {"query": "What is Dify?"} + assert root.outputs == "An LLM application platform." + assert root.get_attribute("message_id") == MESSAGE_ID + assert root.get_attribute("tenant_id") == "tenant-1" + assert root.get_attribute(SpanAttributeKey.EXPERIMENT_ID) == "42" + assert root.get_attribute(SpanAttributeKey.SESSION_ID) == "conversation-1" + assert root.get_attribute(SpanAttributeKey.USER_ID) == "end-user-1" + assert root.get_attribute(SpanAttributeKey.CHAT_USAGE) is None + assert llm.get_attribute(SpanAttributeKey.SESSION_ID) is None + assert llm.inputs == [{"role": "user", "content": "What is Dify?"}] + assert llm.model_name == "gpt-4o" + assert llm.get_attribute(SpanAttributeKey.MODEL_PROVIDER) == "openai" + usage = {"input_tokens": 12, "output_tokens": 8, "total_tokens": 20} + assert _stored_attribute(llm, SpanAttributeKey.CHAT_USAGE) == usage + assert _stored_attribute(llm, SpanAttributeKey.LLM_COST) == {"total_cost": 0.0004} + assert result.trace_id == request_id + assert result.span_urls == {} + assert result.trace_url is not None + link = urlsplit(result.trace_url) + path, _, query = link.fragment.partition("?") + assert f"{link.scheme}://{link.netloc}{link.path}" == "https://mlflow.example/proxy/" + assert path == "/experiments/42/traces" + assert parse_qs(query) == { + "traceId": [request_id], + "selectedTraceId": [request_id], + "selectedEvaluationId": [request_id], + } + + @staticmethod + def test_send_without_credentials_sends_no_authorization() -> None: + provider, egress = _mlflow_provider(ACK) + provider.send(MLflowConfig(tracking_uri="http://localhost:5000/", username="admin"), message_trace()) + (sent,) = egress.requests + assert sent.url == "http://localhost:5000/v1/traces" + assert sent.headers["x-mlflow-experiment-id"] == "0" + assert "Authorization" not in sent.headers + + @staticmethod + def test_workflow_spans_keep_hierarchy_and_errors() -> None: + provider, egress = _mlflow_provider(ACK) + trace = workflow_trace() + + provider.send(MLFLOW_CONFIG, trace) + + run, llm, tool = _exported(egress) + assert llm.parent_id == run.span_id + assert tool.parent_id == llm.span_id + assert tool.span_type == "TOOL" + assert run.status.status_code == SpanStatusCode.ERROR + assert run.status.description == "Tool failed" + assert [(event.name, event.attributes) for event in tool.events] == [ + ("exception", {"exception.type": "Error", "exception.message": "Tool failed"}) + ] + assert llm.status.status_code == SpanStatusCode.OK + assert llm.events == [] + assert run.get_attribute(SpanAttributeKey.SESSION_ID) == WORKFLOW_RUN_ID + assert run.get_attribute(SpanAttributeKey.USER_ID) == "account-1" + assert run.get_attribute("workflow_run_id") == WORKFLOW_RUN_ID + assert llm.get_attribute("node_type") == "llm" + assert llm.get_attribute(SpanAttributeKey.LLM_COST) is None + + @staticmethod + def test_usage_and_cost_are_reported_by_model_calls_only_and_in_usd() -> None: + trace = message_trace() + usage = TokenUsage(prompt_tokens=12, completion_tokens=8, total_tokens=20) + trace = _with_span(trace, 0, trace.spans[0].model_copy(update={"usage": usage})) + trace = _with_span(trace, 1, trace.spans[1].model_copy(update={"cost": Cost(amount=0.03, currency="CNY")})) + provider, egress = _mlflow_provider(ACK) + + provider.send(MLFLOW_CONFIG, trace) + + root, llm = _exported(egress) + assert root.get_attribute(SpanAttributeKey.CHAT_USAGE) is None + assert llm.get_attribute(SpanAttributeKey.CHAT_USAGE) == { + "input_tokens": 12, + "output_tokens": 8, + "total_tokens": 20, + } + assert llm.get_attribute(SpanAttributeKey.LLM_COST) is None + + @staticmethod + @pytest.mark.parametrize( + ("definitions", "expected"), + [ + ( + [FLAT_TOOL, NESTED_TOOL], + [{"type": "function", "function": {k: v for k, v in FLAT_TOOL.items() if k != "type"}}, NESTED_TOOL], + ), + (json.dumps([NESTED_TOOL]), [NESTED_TOOL]), + ([None, 12, {"type": "other"}, {"type": "function"}, {**FLAT_TOOL, "parameters": "schema"}], []), + ("not json", None), + ], + ) + def test_tool_definitions_become_native_chat_tools(definitions: JsonValue, expected: JsonValue) -> None: + trace = message_trace() + llm = trace.spans[1] + attributes = {**llm.attributes, "gen_ai.tool.definitions": definitions} + provider, egress = _mlflow_provider(ACK) + + provider.send(MLFLOW_CONFIG, _with_span(trace, 1, llm.model_copy(update={"attributes": attributes}))) + + _, exported = _exported(egress) + assert exported.get_attribute(SpanAttributeKey.CHAT_TOOLS) == expected + assert exported.get_attribute("gen_ai.tool.definitions") == definitions + + @staticmethod + def test_native_chat_tools_take_precedence() -> None: + native = [{"type": "function", "function": {"name": "native"}}] + trace = message_trace() + llm = trace.spans[1] + attributes = {**llm.attributes, SpanAttributeKey.CHAT_TOOLS: native, "gen_ai.tool.definitions": [FLAT_TOOL]} + provider, egress = _mlflow_provider(ACK) + + provider.send(MLFLOW_CONFIG, _with_span(trace, 1, llm.model_copy(update={"attributes": attributes}))) + + _, exported = _exported(egress) + assert exported.get_attribute(SpanAttributeKey.CHAT_TOOLS) == native + + @staticmethod + def test_fragment_is_sent_as_its_own_trace_linked_to_its_parent() -> None: + provider, egress = _mlflow_provider(ACK) + fragment = tool_fragment() + assert fragment.parent is not None + own_trace_id = trace_id_for(f"{fragment.trace_id}:{fragment.root.span_id}") + + result = provider.send(MLFLOW_CONFIG, fragment) + + (otlp_span,) = _otlp(egress).resource_spans[0].scope_spans[0].spans + assert otlp_span.trace_id.hex() == own_trace_id + assert otlp_span.parent_span_id == b"" + (tool,) = _exported(egress) + assert tool.trace_id == f"tr-{own_trace_id}" + assert tool.span_id == fragment.root.span_id + assert tool.parent_id is None + assert tool.get_attribute("dify.parent_trace_id") == fragment.trace_id + assert tool.get_attribute("dify.parent_span_id") == fragment.parent.span_id + assert tool.get_attribute(SpanAttributeKey.SESSION_ID) == "conversation-1" + assert result.trace_id == f"tr-{own_trace_id}" + assert own_trace_id != fragment.trace_id + + @staticmethod + def test_runs_that_share_a_caller_trace_id_become_separate_traces() -> None: + first = message_trace().model_copy(update={"external_trace_id": "caller-trace"}) + second = workflow_trace().model_copy(update={"external_trace_id": "caller-trace"}) + provider, _ = _mlflow_provider(ACK, ACK) + + first_result = provider.send(MLFLOW_CONFIG, first) + second_result = provider.send(MLFLOW_CONFIG, second) + + assert first_result.trace_id != second_result.trace_id + assert first_result.trace_id == f"tr-{trace_id_for(f'{first.trace_id}:{first.root.span_id}')}" + + @staticmethod + def test_runs_without_a_caller_trace_id_keep_their_own_trace_id() -> None: + provider, _ = _mlflow_provider(ACK) + trace = message_trace() + + assert provider.send(MLFLOW_CONFIG, trace).trace_id == f"tr-{trace.trace_id}" + + @staticmethod + def test_redelivery_keeps_identifiers() -> None: + first, first_egress = _mlflow_provider(ACK) + second, second_egress = _mlflow_provider(ACK) + assert first.send(MLFLOW_CONFIG, tool_fragment()) == second.send(MLFLOW_CONFIG, tool_fragment()) + assert first_egress.requests[0].content == second_egress.requests[0].content + + @staticmethod + @pytest.mark.parametrize( + ("outcome", "expected"), + [ + (_ack(1), ExportOutcomeUnknownError), + ( + EgressResponse(status_code=200, headers={"content-type": "application/json"}, content=b"[]"), + ExportOutcomeUnknownError, + ), + (response(400, {"error_code": "INVALID_PARAMETER_VALUE"}), ExportRejectedError), + (response(401, {}), ExportRejectedError), + (response(503, {}), ExportRetryableError), + (EgressError("connection reset", request_sent=True), ExportOutcomeUnknownError), + (EgressError("connection refused", request_sent=False), ExportRetryableError), + (EgressBlockedError("private address"), ExportRejectedError), + ], + ) + def test_send_failures_are_classified(outcome: EgressResponse | Exception, expected: type[Exception]) -> None: + provider, _ = _mlflow_provider(outcome) + with pytest.raises(expected) as caught: + provider.send(MLFLOW_CONFIG, message_trace()) + assert "s3cret" not in str(caught.value) + + @staticmethod + @pytest.mark.parametrize( + "update", + [{"inputs": {"score": float("nan")}}, {"attributes": {"count": 2**70}}], + ) + def test_unencodable_traces_are_rejected_before_sending(update: dict[str, JsonValue]) -> None: + trace = message_trace() + provider, egress = _mlflow_provider() + with pytest.raises(ExportRejectedError): + provider.send(MLFLOW_CONFIG, _with_span(trace, 0, trace.spans[0].model_copy(update=update))) + assert egress.requests == [] + + +DATABRICKS_CONFIG = DatabricksConfig(experiment_id="42", host="https://dbc.example", personal_access_token="dapi-pat") +OAUTH_CONFIG = DatabricksConfig( + experiment_id="42", host="https://dbc.example", client_id="client", client_secret="client-secret" +) +CLIENT_AUTH = "Basic " + base64.b64encode(b"client:client-secret").decode() +TOKEN = response(200, {"access_token": "oauth-token", "token_type": "Bearer", "expires_in": 3600}) +DATABRICKS_EXPERIMENT = response(200, {"experiment": {"experiment_id": "42", "lifecycle_stage": "active"}}) +SERVER_TRACE_ID = "tr-abcdef0123456789abcdef0123456789" +STARTED = response(200, {"trace": {"trace_info": {"trace_id": SERVER_TRACE_ID}}}) +SIGNED_URL = "https://storage.example/traces/data?sig=signature" +STORED = response(201, b"") + + +def _databricks_provider(*responses: EgressResponse | Exception) -> tuple[DatabricksProvider, RecordingEgress]: + egress = RecordingEgress(responses, grpc_target=None) + return DatabricksProvider(egress), egress + + +def _credentials(credential_type: str, headers: dict[str, str]) -> EgressResponse: + credential_headers = [{"name": name, "value": value} for name, value in headers.items()] + return response( + 200, {"credential_info": {"signed_uri": SIGNED_URL, "type": credential_type, "headers": credential_headers}} + ) + + +def _json(sent: RecordedRequest) -> Any: + return json.loads(sent.content or b"null") + + +class TestDatabricksProvider: + @staticmethod + def test_verify_reads_the_experiment_with_a_personal_access_token() -> None: + provider, egress = _databricks_provider(DATABRICKS_EXPERIMENT) + config = DatabricksConfig(experiment_id="042", host=" https://dbc.example/ ", personal_access_token="dapi-pat") + + verified = provider.verify(config) + + assert verified.host == "https://dbc.example" + assert verified.experiment_id == "42" + (sent,) = egress.requests + assert sent.method == "GET" + assert sent.url == "https://dbc.example/api/2.0/mlflow/experiments/get?experiment_id=42" + assert sent.headers == {"Accept": "application/json", "Authorization": "Bearer dapi-pat"} + + @staticmethod + def test_verify_prefers_oauth_credentials() -> None: + provider, egress = _databricks_provider(TOKEN, DATABRICKS_EXPERIMENT) + + provider.verify(OAUTH_CONFIG.model_copy(update={"personal_access_token": "dapi-pat"})) + + token, experiment = egress.requests + assert (token.method, token.url) == ("POST", "https://dbc.example/oidc/v1/token") + assert token.headers["Authorization"] == CLIENT_AUTH + assert token.headers["Content-Type"] == "application/x-www-form-urlencoded" + assert token.content == b"grant_type=client_credentials&scope=all-apis" + assert experiment.headers["Authorization"] == "Bearer oauth-token" + + @staticmethod + @pytest.mark.parametrize( + ("values", "code"), + [ + ({"host": "http://dbc.example"}, "invalid_workspace_url"), + ({"host": "dbc.example"}, "invalid_workspace_url"), + ({"host": "https://dbc.example/ml"}, "invalid_workspace_url"), + ({"host": "https://dbc.example?o=1"}, "invalid_workspace_url"), + ({"host": "https://token@dbc.example"}, "invalid_workspace_url"), + ({"host": "https://dbc.example:99999"}, "invalid_workspace_url"), + ({"personal_access_token": None}, "authentication_failed"), + ({"personal_access_token": " "}, "authentication_failed"), + ({"personal_access_token": None, "client_id": "client"}, "authentication_failed"), + ], + ) + def test_verify_rejects_unusable_configurations_without_requests(values: dict[str, str | None], code: str) -> None: + provider, egress = _databricks_provider() + config = DatabricksConfig.model_validate({**DATABRICKS_CONFIG.model_dump(), **values}) + with pytest.raises(VerificationError) as caught: + provider.verify(config) + assert caught.value.code == code + assert egress.requests == [] + + @staticmethod + @pytest.mark.parametrize( + ("outcome", "code"), + [ + (response(401, {}), "authentication_failed"), + (response(403, {}), "permission_denied"), + (response(404, {}), "experiment_not_found_or_inaccessible"), + (response(500, {}), "endpoint_unavailable"), + (response(200, b"not json"), "malformed_response"), + ( + response(200, {"experiment": {"experiment_id": "7", "lifecycle_stage": "active"}}), + "experiment_id_mismatch", + ), + ( + response(200, {"experiment": {"experiment_id": "42", "lifecycle_stage": "deleted"}}), + "experiment_not_active", + ), + (EgressError("read timeout", request_sent=True), "connection_failed"), + ], + ) + def test_verify_failures_have_stable_codes(outcome: EgressResponse | Exception, code: str) -> None: + provider, _ = _databricks_provider(outcome) + with pytest.raises(VerificationError) as caught: + provider.verify(DATABRICKS_CONFIG) + assert caught.value.code == code + + @staticmethod + @pytest.mark.parametrize( + ("outcome", "code"), + [ + (response(401, {"error": "invalid_client"}), "authentication_failed"), + (response(404, {}), "endpoint_unavailable"), + (response(200, {"token_type": "Bearer"}), "authentication_failed"), + (response(200, {"access_token": "oauth-token", "token_type": "mac"}), "authentication_failed"), + (response(200, b"not json"), "malformed_response"), + (EgressError("connection refused", request_sent=False), "connection_failed"), + ], + ) + def test_oauth_verification_failures_have_stable_codes(outcome: EgressResponse | Exception, code: str) -> None: + provider, egress = _databricks_provider(outcome) + with pytest.raises(VerificationError) as caught: + provider.verify(OAUTH_CONFIG) + assert caught.value.code == code + assert len(egress.requests) == 1 + + @staticmethod + def test_project_url_is_normalized_without_requests() -> None: + provider, egress = _databricks_provider() + config = DatabricksConfig(experiment_id="042", host="https://dbc.example/") + assert provider.project_url(config) == "https://dbc.example/ml/experiments/42/traces" + assert provider.project_url(DatabricksConfig(experiment_id="42", host="dbc.example")) is None + assert egress.requests == [] + + @staticmethod + def test_trace_metadata_is_started_before_the_artifact_upload() -> None: + provider, egress = _databricks_provider(STARTED, _credentials("AWS_PRESIGNED_URL", {}), STORED) + trace = message_trace() + + result = provider.send(DATABRICKS_CONFIG, trace) + + start, credentials, upload = egress.requests + assert (start.method, start.url) == ("POST", "https://dbc.example/api/3.0/mlflow/traces") + assert start.headers["Authorization"] == "Bearer dapi-pat" + assert start.headers["Content-Type"] == "application/json" + info = _json(start)["trace"]["trace_info"] + assert info["trace_id"] == f"tr-{trace.trace_id}" + assert info["client_request_id"] == trace.trace_id + assert info["trace_location"] == {"type": "MLFLOW_EXPERIMENT", "mlflow_experiment": {"experiment_id": "42"}} + assert info["state"] == "OK" + assert info["request_time"] == "2026-10-01T08:00:00Z" + assert info["execution_duration"] == "2s" + assert info["request_preview"] == '{"query": "What is Dify?"}' + assert info["response_preview"] == '"An LLM application platform."' + assert info["tags"] == {"mlflow.traceName": "message"} + metadata = info["trace_metadata"] + assert metadata["mlflow.trace.session"] == "conversation-1" + assert metadata["mlflow.trace.user"] == "end-user-1" + assert json.loads(metadata["mlflow.trace.tokenUsage"]) == { + "input_tokens": 12, + "output_tokens": 8, + "total_tokens": 20, + } + assert json.loads(metadata["mlflow.trace.cost"]) == {"total_cost": 0.0004} + assert credentials.method == "GET" + assert ( + credentials.url + == f"https://dbc.example/api/2.0/mlflow/traces/{SERVER_TRACE_ID}/credentials-for-data-upload" + ) + assert credentials.headers["Authorization"] == "Bearer dapi-pat" + assert "Authorization" not in upload.headers + root, llm = TraceData.from_dict(_json(upload)).spans + assert root.trace_id == SERVER_TRACE_ID + assert llm.trace_id == SERVER_TRACE_ID + assert (root.span_id, root.parent_id) == (trace.spans[0].span_id, None) + assert llm.parent_id == root.span_id + assert root.get_attribute(SpanAttributeKey.SESSION_ID) == "conversation-1" + assert llm.get_attribute(SpanAttributeKey.CHAT_USAGE) == { + "input_tokens": 12, + "output_tokens": 8, + "total_tokens": 20, + } + assert llm.get_attribute(SpanAttributeKey.LLM_COST) == {"total_cost": 0.0004} + assert result == SendResult( + trace_id=SERVER_TRACE_ID, + trace_url=f"https://dbc.example/ml/experiments/42/traces?selectedEvaluationId={SERVER_TRACE_ID}", + span_urls={}, + ) + + @staticmethod + @pytest.mark.parametrize( + ("credential_type", "credential_headers", "uploads"), + [ + ( + "AWS_PRESIGNED_URL", + {"x-amz-server-side-encryption": "AES256"}, + [("PUT", SIGNED_URL, {"x-amz-server-side-encryption": "AES256"}, True)], + ), + ("GCP_SIGNED_URL", {}, [("PUT", SIGNED_URL, {}, True)]), + ( + "AZURE_SAS_URI", + {"x-ms-version": "2021-08-06"}, + [("PUT", SIGNED_URL, {"x-ms-version": "2021-08-06", "x-ms-blob-type": "BlockBlob"}, True)], + ), + ( + "AZURE_ADLS_GEN2_SAS_URI", + {"x-ms-version": "2021-08-06", "x-ms-blob-type": "BlockBlob"}, + [ + ("PUT", f"{SIGNED_URL}&resource=file", {"x-ms-version": "2021-08-06"}, False), + ( + "PATCH", + f"{SIGNED_URL}&action=append&position=0&flush=true", + {"x-ms-version": "2021-08-06"}, + True, + ), + ], + ), + ], + ) + def test_artifact_upload_follows_the_credential_type( + credential_type: str, + credential_headers: dict[str, str], + uploads: list[tuple[str, str, dict[str, str], bool]], + ) -> None: + provider, egress = _databricks_provider( + STARTED, _credentials(credential_type, credential_headers), *[STORED for _ in uploads] + ) + + provider.send(DATABRICKS_CONFIG, message_trace()) + + sent = egress.requests[2:] + assert [ + (request.method, request.url, request.headers, request.content is not None) for request in sent + ] == uploads + assert all(request.sensitive_url for request in sent) + assert len(TraceData.from_dict(_json(sent[-1])).spans) == 2 + + @staticmethod + def test_oauth_token_is_exchanged_before_export() -> None: + provider, egress = _databricks_provider(TOKEN, STARTED, _credentials("GCP_SIGNED_URL", {}), STORED) + + provider.send(OAUTH_CONFIG, message_trace()) + + token, start, credentials, upload = egress.requests + assert (token.method, token.url) == ("POST", "https://dbc.example/oidc/v1/token") + assert token.headers["Authorization"] == CLIENT_AUTH + assert token.content == b"grant_type=client_credentials&scope=all-apis" + assert start.headers["Authorization"] == "Bearer oauth-token" + assert credentials.headers["Authorization"] == "Bearer oauth-token" + assert "Authorization" not in upload.headers + + @staticmethod + def test_fragment_is_its_own_databricks_trace() -> None: + fragment = tool_fragment() + assert fragment.parent is not None + own_trace_id = f"tr-{trace_id_for(f'{fragment.trace_id}:{fragment.root.span_id}')}" + provider, egress = _databricks_provider( + response(200, {"trace": {"trace_info": {"trace_id": own_trace_id}}}), + _credentials("GCP_SIGNED_URL", {}), + STORED, + ) + + result = provider.send(DATABRICKS_CONFIG, fragment) + + start, _, upload = egress.requests + assert _json(start)["trace"]["trace_info"]["trace_id"] == own_trace_id + assert _json(start)["trace"]["trace_info"]["tags"] == {"mlflow.traceName": "weather"} + (tool,) = TraceData.from_dict(_json(upload)).spans + assert tool.trace_id == own_trace_id + assert tool.parent_id is None + assert tool.get_attribute("dify.parent_trace_id") == fragment.trace_id + assert tool.get_attribute("dify.parent_span_id") == fragment.parent.span_id + assert result.trace_id == own_trace_id + + @staticmethod + @pytest.mark.parametrize( + ("outcome", "expected"), + [ + (response(401, {"error": "invalid_client"}), ExportRejectedError), + (response(503, {}), ExportRetryableError), + (EgressError("connection reset", request_sent=True), ExportRetryableError), + (response(200, {"token_type": "Bearer"}), ExportRejectedError), + (response(200, b"not json"), ExportRejectedError), + ], + ) + def test_oauth_failures_export_nothing(outcome: EgressResponse | Exception, expected: type[Exception]) -> None: + provider, egress = _databricks_provider(outcome) + with pytest.raises(expected) as caught: + provider.send(OAUTH_CONFIG, message_trace()) + assert "client-secret" not in str(caught.value) + assert len(egress.requests) == 1 + + @staticmethod + @pytest.mark.parametrize( + ("responses", "expected"), + [ + ([response(503, {})], ExportRetryableError), + ([response(400, {"error_code": "INVALID_PARAMETER_VALUE"})], ExportRejectedError), + ([EgressError("connection reset", request_sent=True)], ExportOutcomeUnknownError), + ([response(200, b"not json")], ExportOutcomeUnknownError), + ([response(200, {"trace": {}})], ExportOutcomeUnknownError), + ], + ) + def test_trace_start_failures_are_classified( + responses: list[EgressResponse | Exception], expected: type[Exception] + ) -> None: + provider, egress = _databricks_provider(*responses) + with pytest.raises(expected) as caught: + provider.send(DATABRICKS_CONFIG, message_trace()) + assert "dapi-pat" not in str(caught.value) + assert len(egress.requests) == 1 + + @staticmethod + @pytest.mark.parametrize( + "responses", + [ + [response(503, {})], + [response(200, b"not json")], + [_credentials("NOT_A_CREDENTIAL_TYPE", {})], + [ + response( + 200, {"credential_info": {"signed_uri": "ftp://storage.example/data", "type": "GCP_SIGNED_URL"}} + ) + ], + [_credentials("AWS_PRESIGNED_URL", {}), response(403, {})], + [_credentials("AWS_PRESIGNED_URL", {}), EgressError("connection refused", request_sent=False)], + [_credentials("AZURE_ADLS_GEN2_SAS_URI", {}), STORED, response(500, {})], + ], + ) + def test_incomplete_artifact_after_accepted_metadata_is_an_unknown_outcome( + responses: list[EgressResponse | Exception], + ) -> None: + provider, egress = _databricks_provider(STARTED, *responses) + with pytest.raises(ExportOutcomeUnknownError, match="metadata"): + provider.send(DATABRICKS_CONFIG, message_trace()) + assert len(egress.requests) == 1 + len(responses) + + @staticmethod + def test_missing_credentials_are_rejected_without_requests() -> None: + provider, egress = _databricks_provider() + with pytest.raises(ExportRejectedError): + provider.send(DatabricksConfig(experiment_id="42", host="https://dbc.example"), message_trace()) + assert egress.requests == [] + + @staticmethod + def test_unencodable_traces_are_rejected_before_sending() -> None: + trace = message_trace() + child = trace.spans[1].model_copy(update={"outputs": {"score": float("inf")}}) + provider, egress = _databricks_provider() + with pytest.raises(ExportRejectedError): + provider.send(DATABRICKS_CONFIG, trace.model_copy(update={"spans": (trace.spans[0], child)})) + assert egress.requests == [] + + +CREDENTIALS = response( + 200, {"credential_info": {"signed_uri": "https://storage.example/data?sig=signature", "type": "GCP_SIGNED_URL"}} +) +UPLOADED = response(200, b"") +ALLOWED_IMPORTS = { + "base64", + "collections", + "dify_trace_mlflow", + "dify_trace_sdk", + "google", + "json", + "mlflow", + "opentelemetry", + "pydantic", + "re", + "typing", + "urllib", +} +ALLOWED_MLFLOW_IMPORTS = ( + "mlflow.entities", + "mlflow.protos.", + "mlflow.store.tracking.rest_store", + "mlflow.tracing.constant", + "mlflow.tracing.utils.otlp", + "mlflow.utils.rest_utils", +) + + +def _basic(username: str, password: str) -> str: + return "Basic " + base64.b64encode(f"{username}:{password}".encode()).decode() + + +def _experiment(experiment_id: str) -> EgressResponse: + return response(200, {"experiment": {"experiment_id": experiment_id, "lifecycle_stage": "active"}}) + + +def _started(trace_id: str) -> EgressResponse: + return response(200, {"trace": {"trace_info": {"trace_id": trace_id}}}) + + +class TestIsolation: + """One worker process serves many tenants, so nothing may outlive or leak across a call.""" + + @staticmethod + def test_tenants_in_one_process_keep_their_own_credentials_and_leave_the_environment_alone() -> None: + environment = dict(os.environ) + token = response(200, {"access_token": "oauth-token-c", "token_type": "Bearer"}) + egress = RecordingEgress( + [ + *(_experiment("1"), ACK, _experiment("2"), ACK), + *(token, _experiment("3"), token, _started("tr-c"), CREDENTIALS, UPLOADED), + *(_experiment("4"), _started("tr-d"), CREDENTIALS, UPLOADED), + ], + grpc_target=None, + ) + mlflow_provider = MLflowProvider(egress) + databricks_provider = DatabricksProvider(egress) + + for mlflow_config in ( + MLflowConfig( + tracking_uri="https://a.example", experiment_id="1", username="tenant-a", password="password-a" + ), + MLflowConfig(tracking_uri="https://b.example", experiment_id="2"), + ): + mlflow_provider.send(mlflow_provider.verify(mlflow_config), message_trace()) + for databricks_config in ( + DatabricksConfig( + experiment_id="3", host="https://c.example", client_id="client-c", client_secret="secret-c" + ), + DatabricksConfig(experiment_id="4", host="https://d.example", personal_access_token="pat-d"), + ): + databricks_provider.send(databricks_provider.verify(databricks_config), message_trace()) + + assert dict(os.environ) == environment + authorization: dict[str, set[str | None]] = {} + for sent in egress.requests: + authorization.setdefault(urlsplit(sent.url).netloc, set()).add(sent.headers.get("Authorization")) + assert authorization == { + "a.example": {_basic("tenant-a", "password-a")}, + "b.example": {None}, + "c.example": {_basic("client-c", "secret-c"), "Bearer oauth-token-c"}, + "d.example": {"Bearer pat-d"}, + "storage.example": {None}, + } + + @staticmethod + def test_the_package_imports_nothing_from_the_application_or_the_process() -> None: + for path in Path(dify_trace_mlflow.__file__).parent.glob("*.py"): + for node in ast.walk(ast.parse(path.read_text())): + if isinstance(node, ast.Import): + modules = [alias.name for alias in node.names] + elif isinstance(node, ast.ImportFrom): + modules = [node.module or ""] + else: + continue + for module in modules: + assert module.split(".")[0] in ALLOWED_IMPORTS, f"{path.name} imports {module}" + if module.startswith("mlflow"): + assert module.startswith(ALLOWED_MLFLOW_IMPORTS), f"{path.name} imports {module}" diff --git a/api/providers/trace/trace-opik/pyproject.toml b/api/providers/trace/trace-opik/pyproject.toml index 874997168ecd1b..6ccf1ede364310 100644 --- a/api/providers/trace/trace-opik/pyproject.toml +++ b/api/providers/trace/trace-opik/pyproject.toml @@ -2,9 +2,12 @@ name = "dify-trace-opik" version = "0.0.1" dependencies = [ - "opik~=1.11.2", + "dify-trace-sdk", ] description = "Dify ops tracing provider (Opik)." +[project.entry-points."dify.trace_providers"] +opik = "dify_trace_opik.provider:OpikProvider" + [tool.setuptools.packages.find] where = ["src"] diff --git a/api/providers/trace/trace-opik/src/dify_trace_opik/agent_trace.py b/api/providers/trace/trace-opik/src/dify_trace_opik/agent_trace.py deleted file mode 100644 index db03c1a8d36d9f..00000000000000 --- a/api/providers/trace/trace-opik/src/dify_trace_opik/agent_trace.py +++ /dev/null @@ -1,129 +0,0 @@ -"""Synchronous Opik SDK ingestion with native UUIDv7, thread, and usage fields.""" - -from collections.abc import Mapping -from datetime import datetime -from typing import Any -from uuid import NAMESPACE_URL, UUID, uuid5 - -from core.ops.exceptions import RetryableTraceDispatchError -from core.ops.unified_trace.entities import CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace -from services.tracing.provider_common import endpoint, finished_at, json_value, token_usage, usd_cost, utc -from services.tracing.provider_dispatch import TraceExportRejectedError, TraceExportResult, TraceExportUncertainError -from services.tracing.sdk_transport import SSRFTransport - - -def _provider_id(trace_id: str, span_id: str, started_at: datetime) -> str: - from opik.id_helpers import uuid4_to_uuid7 - - seeded = UUID(str(uuid5(NAMESPACE_URL, f"dify:{trace_id}:{span_id}")), version=4) - return str(uuid4_to_uuid7(utc(started_at), str(seeded))) - - -def _content(value: Any) -> dict[str, Any] | list[dict[str, Any]] | str | None: - """Opik accepts objects, object lists and strings, but not every JSON value.""" - if value is None or isinstance(value, (dict, str)): - return value - if isinstance(value, list) and all(isinstance(item, dict) for item in value): - return value - return json_value(value) - - -def export_trace(provider: str, config: Mapping[str, Any], trace: CanonicalTrace) -> TraceExportResult: - import httpx - from opik.rest_api.client import OpikApi - from opik.rest_api.types.error_info_write import ErrorInfoWrite - from opik.rest_api.types.span_write import SpanWrite - - url = endpoint(config.get("url", "https://www.comet.com/opik/api")) - project = str(config.get("project") or "Default Project") - root = trace.spans[0] - trace_id = _provider_id(trace.trace_id, "trace", root.start_time) - ids = {span.id: _provider_id(trace.trace_id, span.id, span.start_time) for span in trace.spans} - accepted = False - with httpx.Client(transport=SSRFTransport(provider), timeout=30) as http_client: - client = OpikApi( - base_url=url, - api_key=config.get("api_key"), - workspace_name=str(config.get("workspace") or "default"), - httpx_client=http_client, - ) - spans = [] - for span in trace.spans: - usage = token_usage(span) - native_usage = { - "prompt_tokens" if key == "input" else "completion_tokens" if key == "output" else "total_tokens": value - for key, value in usage.items() - } - ttft = span.metadata.get("time_to_first_token", span.metadata.get("ttft")) - spans.append( - SpanWrite( - id=ids[span.id], - trace_id=trace_id, - parent_span_id=ids.get(span.parent_id or ""), - project_name=project, - name=span.name, - type="llm" - if span.kind is CanonicalSpanKind.LLM - else "tool" - if span.kind is CanonicalSpanKind.TOOL - else "general", - start_time=utc(span.start_time), - end_time=finished_at(span), - input=_content(span.inputs), - output=_content(span.outputs), - metadata={ - **span.metadata, - "dify.span.kind": span.kind.value, - "dify.span.id": span.id, - "dify.trace_id": trace.trace_id, - "session_id": trace.session_id, - }, - model=str(span.metadata.get("model") or span.metadata.get("model_name") or "") - if span.kind is CanonicalSpanKind.LLM - else None, - provider=str(span.metadata.get("model_provider") or "") - if span.kind is CanonicalSpanKind.LLM - else None, - usage=native_usage or None, - total_estimated_cost=usd_cost(span), - ttft=float(ttft) * 1000 - if span.kind is CanonicalSpanKind.LLM and isinstance(ttft, (int, float)) and ttft >= 0 - else None, - error_info=ErrorInfoWrite( - exception_type="DifyExecutionError", - message=span.error or "Trace operation failed", - traceback="", - ) - if span.status is CanonicalSpanStatus.ERROR - else None, - ) - ) - try: - client.traces.create_trace( - id=trace_id, - project_name=project, - name=root.name, - start_time=utc(root.start_time), - end_time=finished_at(root), - input=_content(root.inputs), - output=_content(root.outputs), - metadata={**root.metadata, "dify.trace_id": trace.trace_id}, - error_info=ErrorInfoWrite( - exception_type="DifyExecutionError", - message=root.error or "Trace operation failed", - traceback="", - ) - if root.status is CanonicalSpanStatus.ERROR - else None, - thread_id=trace.session_id, - request_options={"max_retries": 0}, - ) - accepted = True - client.spans.create_spans(spans=spans, request_options={"max_retries": 0}) - except Exception as error: - if accepted: - raise TraceExportUncertainError("Opik accepted part of the trace; reconciliation is required") from None - if isinstance(error, (RetryableTraceDispatchError, TraceExportRejectedError)): - raise - raise TraceExportRejectedError("Opik rejected the trace") from None - return TraceExportResult(trace_id=trace_id) diff --git a/api/providers/trace/trace-opik/src/dify_trace_opik/config.py b/api/providers/trace/trace-opik/src/dify_trace_opik/config.py index c16ff1d9035ad4..1f6b213bd8968e 100644 --- a/api/providers/trace/trace-opik/src/dify_trace_opik/config.py +++ b/api/providers/trace/trace-opik/src/dify_trace_opik/config.py @@ -1,25 +1,35 @@ -from pydantic import ValidationInfo, field_validator +from dify_trace_sdk.config import ProviderConfig, destination, http_url, name_or_default, secret +from pydantic import field_validator -from core.ops.entities.config_entity import BaseTracingConfig -from core.ops.utils import validate_url_with_path +DEFAULT_URL = "https://www.comet.com/opik/api/" +DEFAULT_PROJECT = "Default Project" -class OpikConfig(BaseTracingConfig): - """ - Model class for Opik tracing config. - """ +class OpikConfig(ProviderConfig): + """Self-hosted Opik needs no API key; without a workspace, Opik uses the account's default one.""" - api_key: str | None = None - project: str | None = None + api_key: str | None = secret(default=None) + project: str = DEFAULT_PROJECT workspace: str | None = None - url: str = "https://www.comet.com/opik/api/" + url: str = destination(default=DEFAULT_URL) - @field_validator("project") + @field_validator("project", mode="before") @classmethod - def project_validator(cls, v, info: ValidationInfo): - return cls.validate_project_field(v, "Default Project") + def _project(cls, value: object) -> object: + # Stored configurations hold null or blank for the default project. + return name_or_default(value, DEFAULT_PROJECT) if value is None or isinstance(value, str) else value + + @field_validator("workspace") + @classmethod + def _workspace(cls, value: str | None) -> str | None: + return (value.strip() or None) if value else None @field_validator("url") @classmethod - def url_validator(cls, v, info: ValidationInfo): - return validate_url_with_path(v, "https://www.comet.com/opik/api/", required_suffix="/api/") + def _url(cls, value: str) -> str: + if not value.strip(): + return DEFAULT_URL + url = http_url(value, allowed_schemes=("https", "http")) + if not url.endswith("/api/"): + raise ValueError("URL should end with /api/") + return url diff --git a/api/providers/trace/trace-opik/src/dify_trace_opik/configuration.py b/api/providers/trace/trace-opik/src/dify_trace_opik/configuration.py deleted file mode 100644 index e2152b906fb261..00000000000000 --- a/api/providers/trace/trace-opik/src/dify_trace_opik/configuration.py +++ /dev/null @@ -1,156 +0,0 @@ -"""Check Opik workspace access without SDK streamers or ambient credentials. - -The auth endpoint is read-only despite using POST. Verification does not create -projects or traces and does not establish project write or ingestion permission. -""" - -from dataclasses import dataclass -from typing import Literal -from urllib.parse import quote, urlencode, urlsplit - -from pydantic import TypeAdapter - -from core.helper import ssrf_proxy -from dify_trace_opik.config import OpikConfig - -type VerificationCode = Literal[ - "invalid_endpoint", - "invalid_workspace", - "invalid_project", - "authentication_failed", - "permission_denied", - "endpoint_unavailable", - "malformed_response", - "connection_failed", -] - -_MESSAGES: dict[VerificationCode, str] = { - "invalid_endpoint": "Opik requires an HTTP(S) API URL ending in /api/.", - "invalid_workspace": "Opik requires a valid workspace name.", - "invalid_project": "Opik requires a valid project name.", - "authentication_failed": "Opik authentication failed.", - "permission_denied": "Opik workspace access was denied.", - "endpoint_unavailable": "Opik endpoint could not be verified.", - "malformed_response": "Opik verification returned an invalid response.", - "connection_failed": "Opik verification request failed.", -} -_JSON_OBJECT = TypeAdapter(dict[str, object]) - - -class OpikConfigurationError(ValueError): - """Fixed safe codes; no credentials, vendor bodies or request details.""" - - def __init__(self, code: VerificationCode) -> None: - self.code = code - super().__init__(_MESSAGES[code]) - - -@dataclass(frozen=True, slots=True) -class VerifiedOpikWorkspace: - url: str - workspace: str - project: str - - -def _api_url(config: OpikConfig) -> str: - value = config.url.strip() - try: - parsed = urlsplit(value) - valid = ( - parsed.scheme in {"http", "https"} - and parsed.hostname - and parsed.username is None - and parsed.password is None - and not parsed.query - and not parsed.fragment - and parsed.path.endswith("/api/") - and "\\" not in value - and not any(ord(char) <= 32 or ord(char) == 127 for char in value) - ) - _ = parsed.port - except ValueError: - raise OpikConfigurationError("invalid_endpoint") from None - if not valid: - raise OpikConfigurationError("invalid_endpoint") - return value - - -def _name(value: str | None, default: str, code: Literal["invalid_workspace", "invalid_project"]) -> str: - result = value.strip() if value else default - result = result or default - if any(ord(char) < 32 or ord(char) == 127 for char in result) or ( - code == "invalid_workspace" and result in {".", ".."} - ): - raise OpikConfigurationError(code) - return result - - -def console_url(config: OpikConfig) -> str: - """Use the configured host and prefix, never the SDK's global OpikConfig.""" - return _api_url(config).removesuffix("api/") - - -def project_url(config: OpikConfig) -> str | None: - """Pure fallback; historical default workspaces cannot be resolved on GET.""" - base = console_url(config) - workspace = _name(config.workspace, "default", "invalid_workspace") - if workspace == "default": - return None - project = _name(config.project, "Default Project", "invalid_project") - return f"{base}{quote(workspace, safe='')}/redirect/projects?{urlencode({'name': project})}" - - -def _request(method: Literal["GET", "POST"], url: str, headers: dict[str, str]) -> dict[str, object] | None: - try: - response = ssrf_proxy.make_request( - method, - url, - headers=headers, - json={} if method == "POST" else None, - max_retries=0, - timeout=30.0, - follow_redirects=False, - ) - except Exception: - raise OpikConfigurationError("connection_failed") from None - try: - if response.status_code == 401: - raise OpikConfigurationError("authentication_failed") - if response.status_code == 403: - raise OpikConfigurationError("permission_denied") - if not 200 <= response.status_code < 300: - raise OpikConfigurationError("endpoint_unavailable") - if method == "POST": - # The official auth contract is No Content, not an arbitrary 2xx page. - if response.content.strip(): - raise OpikConfigurationError("malformed_response") - return None - try: - return _JSON_OBJECT.validate_python(response.json(), strict=True) - except ValueError: - raise OpikConfigurationError("malformed_response") from None - finally: - response.close() - - -def verify_configuration(config: OpikConfig) -> VerifiedOpikWorkspace: - """Check auth/workspace via the same endpoints as Opik 1.11.2's auth_check. - - /v1/private/auth checks workspace access and returns no content. Only the - sentinel default workspace needs /v1/private/auth/workspace resolution. - Project creation remains an ingestion concern, as in the legacy behavior. - """ - url = _api_url(config) - workspace = _name(config.workspace, "default", "invalid_workspace") - project = _name(config.project, "Default Project", "invalid_project") - headers = {"Comet-Workspace": workspace, "Accept": "application/json"} - if config.api_key: - headers["Authorization"] = config.api_key - _request("POST", url + "v1/private/auth", headers) - if workspace == "default": - data = _request("GET", url + "v1/private/auth/workspace", headers) - resolved = data.get("workspace_name") if data else None - if not isinstance(resolved, str) or not resolved.strip(): - raise OpikConfigurationError("malformed_response") - workspace = _name(resolved, "default", "invalid_workspace") - return VerifiedOpikWorkspace(url=url, workspace=workspace, project=project) diff --git a/api/providers/trace/trace-opik/src/dify_trace_opik/opik_trace.py b/api/providers/trace/trace-opik/src/dify_trace_opik/opik_trace.py deleted file mode 100644 index 7323542c0f6264..00000000000000 --- a/api/providers/trace/trace-opik/src/dify_trace_opik/opik_trace.py +++ /dev/null @@ -1,479 +0,0 @@ -import hashlib -import logging -import os -import uuid -from datetime import datetime, timedelta -from typing import Any, cast, override - -from opik import Opik, Trace -from opik.id_helpers import uuid4_to_uuid7 -from sqlalchemy.orm import sessionmaker - -from core.ops.base_trace_instance import BaseTraceInstance -from core.ops.entities.trace_entity import ( - BaseTraceInfo, - DatasetRetrievalTraceInfo, - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - SuggestedQuestionTraceInfo, - ToolTraceInfo, - TraceTaskName, - WorkflowTraceInfo, -) -from core.repositories import DifyCoreRepositoryFactory -from dify_trace_opik.config import OpikConfig -from extensions.ext_database import db -from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey -from models import EndUser, MessageFile, WorkflowNodeExecutionTriggeredFrom - -logger = logging.getLogger(__name__) - - -def wrap_dict(key_name, data): - """Make sure that the input data is a dict""" - if not isinstance(data, dict): - return {key_name: data} - - return data - - -def wrap_metadata(metadata, **kwargs): - """Add common metatada to all Traces and Spans""" - metadata["created_from"] = "dify" - - metadata.update(kwargs) - - return metadata - - -def _seed_to_uuid4(seed: str) -> str: - """Derive a deterministic UUID4-formatted string from an arbitrary seed. - - uuid4_to_uuid7 requires a valid UUID v4 string, but some Dify identifiers - are not UUIDs (e.g. a workflow_run_id with a "-root" suffix appended to - distinguish the root span from the trace). This helper hashes the seed - with MD5 and patches the version/variant bits so the result satisfies the - UUID v4 contract. - """ - raw = hashlib.md5(seed.encode()).digest() - ba = bytearray(raw) - ba[6] = (ba[6] & 0x0F) | 0x40 # version 4 - ba[8] = (ba[8] & 0x3F) | 0x80 # variant 1 - return str(uuid.UUID(bytes=bytes(ba))) - - -def prepare_opik_uuid(user_datetime: datetime | None, user_uuid: str | None) -> uuid.UUID: - """Preserve UUID4 mappings and deterministically adapt external trace IDs. - - OpenTelemetry trace IDs and other external identifiers need not be UUID4, - which the Opik conversion helper requires. Reuse the existing stable seed - mapping for those identifiers. Missing identifiers retain random-ID behavior. - """ - - if user_datetime is None: - user_datetime = datetime.now() - - if user_uuid is None: - user_uuid = str(uuid.uuid4()) - - try: - is_uuid4 = uuid.UUID(user_uuid).version == 4 - except ValueError: - is_uuid4 = False - if not is_uuid4: - user_uuid = _seed_to_uuid4(user_uuid) - - return uuid4_to_uuid7(user_datetime, user_uuid) - - -class OpikDataTrace(BaseTraceInstance): - def __init__( - self, - opik_config: OpikConfig, - ): - super().__init__(opik_config) - self.opik_client = Opik( - project_name=opik_config.project, - workspace=opik_config.workspace, - host=opik_config.url, - api_key=opik_config.api_key, - ) - self.project = opik_config.project - self.file_base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001") - - @override - def trace(self, trace_info: BaseTraceInfo): - match trace_info: - case WorkflowTraceInfo(): - self.workflow_trace(trace_info) - case MessageTraceInfo(): - self.message_trace(trace_info) - case ModerationTraceInfo(): - self.moderation_trace(trace_info) - case SuggestedQuestionTraceInfo(): - self.suggested_question_trace(trace_info) - case DatasetRetrievalTraceInfo(): - self.dataset_retrieval_trace(trace_info) - case ToolTraceInfo(): - self.tool_trace(trace_info) - case GenerateNameTraceInfo(): - self.generate_name_trace(trace_info) - case _: - pass - - def workflow_trace(self, trace_info: WorkflowTraceInfo): - workflow_metadata = wrap_metadata( - trace_info.metadata, message_id=trace_info.message_id, workflow_app_log_id=trace_info.workflow_app_log_id - ) - - if trace_info.message_id: - dify_trace_id = trace_info.trace_id or trace_info.message_id - trace_name = TraceTaskName.MESSAGE_TRACE - trace_tags = ["message", "workflow"] - root_span_seed = trace_info.workflow_run_id - else: - dify_trace_id = trace_info.trace_id or trace_info.workflow_run_id - trace_name = TraceTaskName.WORKFLOW_TRACE - trace_tags = ["workflow"] - root_span_seed = _seed_to_uuid4(trace_info.workflow_run_id + "-root") - - opik_trace_id = prepare_opik_uuid(trace_info.start_time, dify_trace_id) - - trace_data = { - "id": opik_trace_id, - "name": trace_name, - "start_time": trace_info.start_time, - "end_time": trace_info.end_time, - "metadata": workflow_metadata, - "input": wrap_dict("input", trace_info.workflow_run_inputs), - "output": wrap_dict("output", trace_info.workflow_run_outputs), - "thread_id": trace_info.conversation_id, - "tags": trace_tags, - "project_name": self.project, - } - self.add_trace(trace_data) - - root_span_id = prepare_opik_uuid(trace_info.start_time, root_span_seed) - span_data = { - "id": root_span_id, - "parent_span_id": None, - "trace_id": opik_trace_id, - "name": TraceTaskName.WORKFLOW_TRACE, - "input": wrap_dict("input", trace_info.workflow_run_inputs), - "output": wrap_dict("output", trace_info.workflow_run_outputs), - "start_time": trace_info.start_time, - "end_time": trace_info.end_time, - "metadata": workflow_metadata, - "tags": ["workflow"], - "project_name": self.project, - } - self.add_span(span_data) - - # through workflow_run_id get all_nodes_execution using repository - session_factory = sessionmaker(bind=db.engine) - # Find the app's creator account - app_id = trace_info.metadata.get("app_id") - if not app_id: - raise ValueError("No app_id found in trace_info metadata") - - service_account = self.get_service_account_with_tenant(app_id) - - workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( - session_factory=session_factory, - tenant_id=trace_info.tenant_id, - user=service_account, - app_id=app_id, - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - - # Get all executions for this workflow run - workflow_node_executions = workflow_node_execution_repository.get_by_workflow_execution( - workflow_execution_id=trace_info.workflow_run_id - ) - - for node_execution in workflow_node_executions: - node_execution_id = node_execution.id - tenant_id = trace_info.tenant_id # Use from trace_info instead - app_id = trace_info.metadata.get("app_id") # Use from trace_info instead - node_name = node_execution.title - node_type = node_execution.node_type - status = node_execution.status - if node_type == BuiltinNodeTypes.LLM: - inputs = node_execution.process_data.get("prompts", {}) if node_execution.process_data else {} - else: - inputs = node_execution.inputs or {} - outputs = node_execution.outputs or {} - created_at = node_execution.created_at or datetime.now() - elapsed_time = node_execution.elapsed_time - finished_at = created_at + timedelta(seconds=elapsed_time) - - execution_metadata = node_execution.metadata or {} - metadata = {str(k): v for k, v in execution_metadata.items()} - metadata.update( - { - "workflow_run_id": trace_info.workflow_run_id, - "node_execution_id": node_execution_id, - "tenant_id": tenant_id, - "app_id": app_id, - "app_name": node_name, - "node_type": node_type, - "status": status, - } - ) - - process_data = node_execution.process_data or {} - - provider = None - model = None - total_tokens = 0 - completion_tokens = 0 - prompt_tokens = 0 - - if process_data and process_data.get("model_mode") == "chat": - run_type = "llm" - provider = process_data.get("model_provider", None) - model = process_data.get("model_name", "") - metadata.update( - { - "ls_provider": provider, - "ls_model_name": model, - } - ) - - try: - usage_data = process_data.get("usage", {}) if "usage" in process_data else outputs.get("usage", {}) - total_tokens = usage_data.get("total_tokens", 0) - prompt_tokens = usage_data.get("prompt_tokens", 0) - completion_tokens = usage_data.get("completion_tokens", 0) - except Exception: - logger.error("Failed to extract usage", exc_info=True) - - else: - run_type = "tool" - - if not total_tokens: - total_tokens = execution_metadata.get(WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS) or 0 - - span_data = { - "trace_id": opik_trace_id, - "id": prepare_opik_uuid(created_at, node_execution_id), - "parent_span_id": root_span_id, - "name": node_name, - "type": run_type, - "start_time": created_at, - "end_time": finished_at, - "metadata": wrap_metadata(metadata), - "input": wrap_dict("input", inputs), - "output": wrap_dict("output", outputs), - "tags": ["node_execution"], - "project_name": self.project, - "usage": { - "total_tokens": total_tokens, - "completion_tokens": completion_tokens, - "prompt_tokens": prompt_tokens, - }, - "model": model, - "provider": provider, - } - - self.add_span(span_data) - - def message_trace(self, trace_info: MessageTraceInfo): - # get message file data - file_list = cast(list[str], trace_info.file_list) or [] - message_file_data: MessageFile | None = trace_info.message_file_data - - if message_file_data is not None: - file_url = f"{self.file_base_url}/{message_file_data.url}" if message_file_data else "" - file_list.append(file_url) - - message_data = trace_info.message_data - if message_data is None: - return - - metadata = trace_info.metadata - dify_trace_id = trace_info.trace_id or trace_info.message_id - - user_id = message_data.from_account_id - metadata["user_id"] = user_id - metadata["file_list"] = file_list - - if message_data.from_end_user_id: - end_user_data: EndUser | None = db.session.get(EndUser, message_data.from_end_user_id) - if end_user_data is not None: - end_user_id = end_user_data.session_id - metadata["end_user_id"] = end_user_id - - trace_data = { - "id": prepare_opik_uuid(trace_info.start_time, dify_trace_id), - "name": TraceTaskName.MESSAGE_TRACE, - "start_time": trace_info.start_time, - "end_time": trace_info.end_time, - "metadata": wrap_metadata(metadata), - "input": trace_info.inputs, - "output": message_data.answer, - "thread_id": message_data.conversation_id, - "tags": ["message", str(trace_info.conversation_mode)], - "project_name": self.project, - } - trace = self.add_trace(trace_data) - - span_data = { - "trace_id": trace.id, - "name": "llm", - "type": "llm", - "start_time": trace_info.start_time, - "end_time": trace_info.end_time, - "metadata": wrap_metadata(metadata), - "input": {"input": trace_info.inputs}, - "output": {"output": message_data.answer}, - "tags": ["llm", str(trace_info.conversation_mode)], - "usage": { - "completion_tokens": trace_info.answer_tokens, - "prompt_tokens": trace_info.message_tokens, - "total_tokens": trace_info.total_tokens, - }, - "project_name": self.project, - } - self.add_span(span_data) - - def moderation_trace(self, trace_info: ModerationTraceInfo): - if trace_info.message_data is None: - return - - start_time = trace_info.start_time or trace_info.message_data.created_at - - span_data = { - "trace_id": prepare_opik_uuid(start_time, trace_info.trace_id or trace_info.message_id), - "name": TraceTaskName.MODERATION_TRACE, - "type": "tool", - "start_time": start_time, - "end_time": trace_info.end_time or trace_info.message_data.updated_at, - "metadata": wrap_metadata(trace_info.metadata), - "input": wrap_dict("input", trace_info.inputs), - "output": { - "action": trace_info.action, - "flagged": trace_info.flagged, - "preset_response": trace_info.preset_response, - "inputs": trace_info.inputs, - }, - "tags": ["moderation"], - } - - self.add_span(span_data) - - def suggested_question_trace(self, trace_info: SuggestedQuestionTraceInfo): - message_data = trace_info.message_data - if message_data is None: - return - - start_time = trace_info.start_time or message_data.created_at - - span_data = { - "trace_id": prepare_opik_uuid(start_time, trace_info.trace_id or trace_info.message_id), - "name": TraceTaskName.SUGGESTED_QUESTION_TRACE, - "type": "tool", - "start_time": start_time, - "end_time": trace_info.end_time or message_data.updated_at, - "metadata": wrap_metadata(trace_info.metadata), - "input": wrap_dict("input", trace_info.inputs), - "output": wrap_dict("output", trace_info.suggested_question), - "tags": ["suggested_question"], - } - - self.add_span(span_data) - - def dataset_retrieval_trace(self, trace_info: DatasetRetrievalTraceInfo): - if trace_info.message_data is None: - return - - start_time = trace_info.start_time or trace_info.message_data.created_at - - span_data = { - "trace_id": prepare_opik_uuid(start_time, trace_info.trace_id or trace_info.message_id), - "name": TraceTaskName.DATASET_RETRIEVAL_TRACE, - "type": "tool", - "start_time": start_time, - "end_time": trace_info.end_time or trace_info.message_data.updated_at, - "metadata": wrap_metadata(trace_info.metadata), - "input": wrap_dict("input", trace_info.inputs), - "output": {"documents": trace_info.documents}, - "tags": ["dataset_retrieval"], - } - - self.add_span(span_data) - - def tool_trace(self, trace_info: ToolTraceInfo): - span_data = { - "trace_id": prepare_opik_uuid(trace_info.start_time, trace_info.trace_id or trace_info.message_id), - "name": trace_info.tool_name, - "type": "tool", - "start_time": trace_info.start_time, - "end_time": trace_info.end_time, - "metadata": wrap_metadata(trace_info.metadata), - "input": wrap_dict("input", trace_info.tool_inputs), - "output": wrap_dict("output", trace_info.tool_outputs), - "tags": ["tool", trace_info.tool_name], - } - - self.add_span(span_data) - - def generate_name_trace(self, trace_info: GenerateNameTraceInfo): - trace_data = { - "id": prepare_opik_uuid(trace_info.start_time, trace_info.trace_id or trace_info.message_id), - "name": TraceTaskName.GENERATE_NAME_TRACE, - "start_time": trace_info.start_time, - "end_time": trace_info.end_time, - "metadata": wrap_metadata(trace_info.metadata), - "input": trace_info.inputs, - "output": trace_info.outputs, - "thread_id": trace_info.conversation_id, - "tags": ["generate_name"], - "project_name": self.project, - } - - trace = self.add_trace(trace_data) - - span_data = { - "trace_id": trace.id, - "name": TraceTaskName.GENERATE_NAME_TRACE, - "start_time": trace_info.start_time, - "end_time": trace_info.end_time, - "metadata": wrap_metadata(trace_info.metadata), - "input": wrap_dict("input", trace_info.inputs), - "output": wrap_dict("output", trace_info.outputs), - "tags": ["generate_name"], - } - - self.add_span(span_data) - - def add_trace(self, opik_trace_data: dict[str, Any]) -> Trace: - try: - trace = self.opik_client.trace(**opik_trace_data) - logger.debug("Opik Trace created successfully") - return trace - except Exception as e: - raise ValueError(f"Opik Failed to create trace: {str(e)}") from e - - def add_span(self, opik_span_data: dict[str, Any]): - try: - self.opik_client.span(**opik_span_data) - logger.debug("Opik Span created successfully") - except Exception as e: - raise ValueError(f"Opik Failed to create span: {str(e)}") from e - - def api_check(self): - try: - self.opik_client.auth_check() - return True - except Exception as e: - logger.info("Opik API check failed: %s", str(e), exc_info=True) - raise ValueError(f"Opik API check failed: {str(e)}") from e - - def get_project_url(self): - try: - return self.opik_client.get_project_url(project_name=self.project) - except Exception as e: - logger.info("Opik get run url failed: %s", str(e), exc_info=True) - raise ValueError(f"Opik get run url failed: {str(e)}") from e diff --git a/api/providers/trace/trace-opik/src/dify_trace_opik/provider.py b/api/providers/trace/trace-opik/src/dify_trace_opik/provider.py new file mode 100644 index 00000000000000..a1f6de4a5dd229 --- /dev/null +++ b/api/providers/trace/trace-opik/src/dify_trace_opik/provider.py @@ -0,0 +1,334 @@ +"""Opik destination over its REST API. + +Each trace unit becomes an Opik trace, created first, and its spans, created in +batches; the root span is both the trace and its top-level span. Opik requires +UUIDv7 identifiers, so they are derived from the Dify trace and span start +times: a redelivered unit keeps its identifiers, and a fragment joins the trace +of its message. The Opik SDK is never imported: importing it replaces the +process's Sentry client with Comet's. +""" + +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import datetime +from typing import override +from urllib.parse import quote, urlencode, urlsplit +from uuid import NAMESPACE_URL, UUID, uuid5 + +from dify_trace_sdk.egress import EgressResponse +from dify_trace_sdk.errors import ( + EgressError, + ExportOutcomeUnknownError, + ExportRejectedError, + TraceProviderError, + VerificationError, +) +from dify_trace_sdk.exporting import json_value, send_export_request +from dify_trace_sdk.model import Span, SpanKind, SpanStatus, Trace +from dify_trace_sdk.provider import ProviderCapabilities, SendResult, TraceProvider +from pydantic import JsonValue + +from dify_trace_opik.config import DEFAULT_PROJECT, OpikConfig + +VERIFY_TIMEOUT_SECONDS = 30.0 +MAX_SPANS_PER_BATCH = 1000 +"""The Opik SDK's span batch size; larger units are sent in several batches.""" +DEFAULT_WORKSPACE = "default" +"""Opik's name for the account's default workspace; only ``verify`` resolves it.""" +_SPAN_TYPES: Mapping[SpanKind, str] = {SpanKind.LLM: "llm", SpanKind.TOOL: "tool"} +"""Opik has no chain, agent, or retriever spans; those are ``general``.""" + + +@dataclass(frozen=True) +class _Settings: + """The configuration as Opik receives it.""" + + url: str + """The API base URL, ending in ``/api/``.""" + workspace: str + project: str + + +def _settings(config: OpikConfig) -> _Settings: + """Check every value that reaches a URL or header before a credential is sent. + + Raises: + VerificationError: A value cannot be sent safely; ``code`` names it. + """ + return _Settings( + url=_api_url(config.url), + workspace=_workspace_name(config.workspace), + project=_project_name(config.project), + ) + + +def _api_url(value: str) -> str: + url = value.strip() + try: + parsed = urlsplit(url) + _ = parsed.port # raises ValueError for an invalid port + except ValueError: + raise VerificationError("invalid_endpoint") from None + if ( + parsed.scheme not in {"http", "https"} + or not parsed.hostname + or parsed.username is not None + or parsed.password is not None + or parsed.query + or parsed.fragment + or not url.endswith("/api/") + or "\\" in url + or any(ord(char) <= 32 or ord(char) == 127 for char in url) + ): + raise VerificationError("invalid_endpoint") + return url + + +def _has_control_character(value: str) -> bool: + return any(ord(char) < 32 or ord(char) == 127 for char in value) + + +def _workspace_name(value: str | None) -> str: + """The workspace travels as a header value and as a console path segment.""" + name = (value.strip() if value else "") or DEFAULT_WORKSPACE + if _has_control_character(name) or name in {".", ".."}: + raise VerificationError("invalid_workspace") + return name + + +def _project_name(value: str) -> str: + name = value.strip() or DEFAULT_PROJECT + if _has_control_character(name): + raise VerificationError("invalid_project") + return name + + +def _headers(api_key: str | None, workspace: str) -> dict[str, str]: + headers = {"Comet-Workspace": workspace, "Accept": "application/json"} + if api_key: + headers["Authorization"] = api_key + return headers + + +def _resolved_workspace(response: EgressResponse) -> str: + try: + body = response.json() + except ValueError: + raise VerificationError("malformed_response") from None + name = body.get("workspace_name") if isinstance(body, dict) else None + if not isinstance(name, str) or not name.strip(): + raise VerificationError("malformed_response") + return _workspace_name(name) + + +def _opik_id(trace_id: str, key: str, started_at: datetime) -> str: + """A UUIDv7 that is stable for a Dify trace, a key within it, and a start time. + + Identical to the legacy exporter's IDs: a uuid5 seed whose leading bits are + replaced as ``opik.id_helpers.uuid4_to_uuid7`` does. + """ + seed = uuid5(NAMESPACE_URL, f"dify:{trace_id}:{key}").bytes + value = bytearray(int(started_at.timestamp() * 1000).to_bytes(6, "big") + seed[6:]) + value[6] = 0x70 | (seed[6] & 0x0F) + value[8] = 0x80 | (seed[8] & 0x3F) + return str(UUID(bytes=bytes(value))) + + +def _iso(value: datetime) -> str: + """UTC with a ``Z`` suffix, as the Opik SDK sends it.""" + return value.isoformat().replace("+00:00", "Z") + + +def _content(value: JsonValue) -> JsonValue: + """Opik takes an object, a list of objects, or a string as input and output.""" + if value is None or isinstance(value, dict | str): + return value + if isinstance(value, list) and all(isinstance(item, dict) for item in value): + return value + return json_value(value) + + +def _error_info(span: Span) -> dict[str, str] | None: + if span.status is not SpanStatus.ERROR: + return None + return {"exception_type": "DifyExecutionError", "message": span.error or "Trace operation failed", "traceback": ""} + + +def _model_call_fields(span: Span) -> dict[str, object]: + """Only model calls carry usage, cost, and latency; wrapper spans would count them twice.""" + fields: dict[str, object] = {"model": span.model_name, "provider": span.model_provider} + if span.usage is not None: + fields["usage"] = { + "prompt_tokens": span.usage.prompt_tokens, + "completion_tokens": span.usage.completion_tokens, + "total_tokens": span.usage.total_tokens, + } + # Opik costs are in USD; without one, Opik may estimate it from the model and usage. + if span.cost is not None and span.cost.currency.upper() == "USD": + fields["total_estimated_cost"] = span.cost.amount + if span.time_to_first_token is not None: + fields["ttft"] = span.time_to_first_token * 1000 + return fields + + +def _without_none(body: Mapping[str, object]) -> dict[str, object]: + return {key: value for key, value in body.items() if value is not None} + + +def _encode(body: Mapping[str, object]) -> bytes: + return json_value(body).encode() + + +class OpikProvider(TraceProvider[OpikConfig]): + id = "opik" + config_model = OpikConfig + capabilities = ProviderCapabilities(native_cost=True, trace_url=False, span_url=False) + + @override + def verify(self, config: OpikConfig) -> OpikConfig: + """Check workspace access with the endpoints of the Opik SDK's ``auth_check``. + + ``POST v1/private/auth`` answers No Content when the key may use the + workspace; the default workspace is then resolved to its name. Neither + request creates the project: Opik creates it on first ingestion. + """ + settings = _settings(config) + headers = _headers(config.api_key, settings.workspace) + access = self._check( + "POST", f"{settings.url}v1/private/auth", {**headers, "Content-Type": "application/json"}, b"{}" + ) + if access.content.strip(): + # A 2xx page with content, such as a login form, is not Opik's auth endpoint. + raise VerificationError("malformed_response") + workspace = settings.workspace + if workspace == DEFAULT_WORKSPACE: + workspace = _resolved_workspace( + self._check("GET", f"{settings.url}v1/private/auth/workspace", headers, None) + ) + return config.model_copy(update={"url": settings.url, "workspace": workspace, "project": settings.project}) + + @override + def project_url(self, config: OpikConfig) -> str | None: + try: + settings = _settings(config) + except VerificationError: + return None + console = settings.url.removesuffix("api/") + if settings.workspace == DEFAULT_WORKSPACE: + # Only verify can resolve the default workspace's name; link the console instead. + return console + query = urlencode({"name": settings.project}) + return f"{console}{quote(settings.workspace, safe='')}/redirect/projects?{query}" + + @override + def send(self, config: OpikConfig, trace: Trace) -> SendResult: + try: + settings = _settings(config) + except VerificationError as error: + raise ExportRejectedError(f"Opik configuration is invalid ({error.code})") from None + # A fragment's parent is the root of its trace, so both units derive the same trace ID. + trace_start = trace.parent.start_time if trace.parent is not None else trace.root.start_time + trace_id = _opik_id(trace.trace_id, "trace", trace_start) + try: + requests = self._requests(trace, trace_id, settings) + except ValueError: + raise ExportRejectedError("Opik trace contains values that cannot be encoded as JSON") from None + headers = {**_headers(config.api_key, settings.workspace), "Content-Type": "application/json"} + accepted = False + for url, content in requests: + try: + send_export_request( + self.egress, + "POST", + url, + destination="Opik", + headers=headers, + content=content, + sensitive_url=False, + ) + except TraceProviderError as error: + if accepted: + raise ExportOutcomeUnknownError( + f"Opik accepted part of the trace before a failure: {error}" + ) from None + raise + accepted = True + return SendResult(trace_id=trace_id, trace_url=None, span_urls={}) + + def _check(self, method: str, url: str, headers: Mapping[str, str], content: bytes | None) -> EgressResponse: + try: + response = self.egress.request( + method, + url, + headers=headers, + content=content, + timeout=VERIFY_TIMEOUT_SECONDS, + sensitive_url=False, + ) + except EgressError: + raise VerificationError("connection_failed") from None + if response.status_code == 401: + raise VerificationError("authentication_failed") + if response.status_code == 403: + raise VerificationError("permission_denied") + if not 200 <= response.status_code < 300: + raise VerificationError("endpoint_unavailable") + return response + + def _requests(self, trace: Trace, trace_id: str, settings: _Settings) -> list[tuple[str, bytes]]: + """Encode every request before the first is sent, so a bad value cannot leave a partial trace.""" + span_ids = {span.span_id: _opik_id(trace.trace_id, span.span_id, span.start_time) for span in trace.spans} + spans = [self._span(trace, span, trace_id, span_ids, settings.project) for span in trace.spans] + requests: list[tuple[str, bytes]] = [] + if trace.parent is None: + created = self._trace(trace, trace_id, settings.project) + requests.append((f"{settings.url}v1/private/traces", _encode(created))) + for start in range(0, len(spans), MAX_SPANS_PER_BATCH): + batch = {"spans": spans[start : start + MAX_SPANS_PER_BATCH]} + requests.append((f"{settings.url}v1/private/spans/batch", _encode(batch))) + return requests + + @staticmethod + def _trace(trace: Trace, trace_id: str, project: str) -> dict[str, object]: + root = trace.root + return _without_none( + { + "id": trace_id, + "project_name": project, + "name": root.name, + "start_time": _iso(root.start_time), + "end_time": _iso(root.end_time), + "input": _content(root.inputs), + "output": _content(root.outputs), + "metadata": {**root.attributes, "dify_trace_id": trace.trace_id, "user_id": trace.user_id}, + "error_info": _error_info(root), + "thread_id": trace.session_id, + } + ) + + @staticmethod + def _span(trace: Trace, span: Span, trace_id: str, span_ids: Mapping[str, str], project: str) -> dict[str, object]: + body: dict[str, object] = { + "id": span_ids[span.span_id], + "project_name": project, + "trace_id": trace_id, + # The unit's root is top-level, and so is a fragment's root: its parent belongs to another unit. + "parent_span_id": span_ids.get(span.parent_span_id) if span.parent_span_id is not None else None, + "name": span.name, + "type": _SPAN_TYPES.get(span.kind, "general"), + "start_time": _iso(span.start_time), + "end_time": _iso(span.end_time), + "input": _content(span.inputs), + "output": _content(span.outputs), + "metadata": { + **span.attributes, + "dify_span_kind": span.kind.value, + "dify_span_id": span.span_id, + "dify_trace_id": trace.trace_id, + "session_id": trace.session_id, + }, + "error_info": _error_info(span), + } + if span.kind is SpanKind.LLM: + body.update(_model_call_fields(span)) + return _without_none(body) diff --git a/api/providers/trace/trace-opik/tests/unit_tests/src/dify_trace_opik/test_config.py b/api/providers/trace/trace-opik/tests/unit_tests/src/dify_trace_opik/test_config.py index 5a54b70bba77c4..8f5a26ec31fa9e 100644 --- a/api/providers/trace/trace-opik/tests/unit_tests/src/dify_trace_opik/test_config.py +++ b/api/providers/trace/trace-opik/tests/unit_tests/src/dify_trace_opik/test_config.py @@ -1,48 +1,35 @@ import pytest from dify_trace_opik.config import OpikConfig +from dify_trace_sdk.config import FormField from pydantic import ValidationError - -class TestOpikConfig: - """Test cases for OpikConfig""" - - def test_valid_config(self): - """Test valid Opik configuration""" - config = OpikConfig( - api_key="test_key", - project="test_project", - workspace="test_workspace", - url="https://custom.comet.com/opik/api/", - ) - assert config.api_key == "test_key" - assert config.project == "test_project" - assert config.workspace == "test_workspace" - assert config.url == "https://custom.comet.com/opik/api/" - - def test_default_values(self): - """Test default values are set correctly""" - config = OpikConfig() - assert config.api_key is None - assert config.project is None - assert config.workspace is None - assert config.url == "https://www.comet.com/opik/api/" - - def test_project_validation_empty(self): - """Test project validation with empty value""" - config = OpikConfig(project="") - assert config.project == "Default Project" - - def test_url_validation_empty(self): - """Test URL validation with empty value""" - config = OpikConfig(url="") - assert config.url == "https://www.comet.com/opik/api/" - - def test_url_validation_missing_suffix(self): - """Test URL validation requires /api/ suffix""" - with pytest.raises(ValidationError, match="URL should end with /api/"): - OpikConfig(url="https://custom.comet.com/opik/") - - def test_url_validation_invalid_scheme(self): - """Test URL validation rejects invalid schemes""" - with pytest.raises(ValidationError, match="URL must start with https:// or http://"): - OpikConfig(url="ftp://custom.comet.com/opik/api/") +API_URL = "https://www.comet.com/opik/api/" + + +def test_config_fields_and_defaults() -> None: + assert OpikConfig.secret_fields() == {"api_key"} + assert OpikConfig.derived_fields() == frozenset() + assert OpikConfig.form_fields() == ( + FormField(name="api_key", secret=True, required=False), + FormField(name="project", secret=False, required=False), + FormField(name="workspace", secret=False, required=False), + FormField(name="url", secret=False, required=False), + ) + stored = OpikConfig.model_validate({"api_key": None, "project": None, "workspace": " ", "url": ""}) + assert (stored.api_key, stored.project, stored.workspace, stored.url) == (None, "Default Project", None, API_URL) + named = OpikConfig(project=" Dify ", workspace=" team-1 ", url=" http://localhost:5173/api/ ") + assert (named.project, named.workspace, named.url) == ("Dify", "team-1", "http://localhost:5173/api/") + + +@pytest.mark.parametrize( + ("url", "message"), + [ + ("ftp://comet.example.com/api/", "URL must start with https:// or http://"), + ("https://comet.example.com/opik/", "URL should end with /api/"), + ("https://user:secret@comet.example.com/api/", "URL must name a host and must not contain credentials"), + ], +) +def test_config_rejects_unusable_urls_without_echoing_them(url: str, message: str) -> None: + with pytest.raises(ValidationError, match=message) as caught: + OpikConfig(url=url) + assert "secret" not in str(caught.value) diff --git a/api/providers/trace/trace-opik/tests/unit_tests/src/dify_trace_opik/test_opik_trace.py b/api/providers/trace/trace-opik/tests/unit_tests/src/dify_trace_opik/test_opik_trace.py deleted file mode 100644 index c228c00ebefc04..00000000000000 --- a/api/providers/trace/trace-opik/tests/unit_tests/src/dify_trace_opik/test_opik_trace.py +++ /dev/null @@ -1,1052 +0,0 @@ -from __future__ import annotations - -import collections -import logging -import uuid -from datetime import UTC, datetime, timedelta -from types import SimpleNamespace -from typing import cast, override -from unittest.mock import MagicMock, patch - -import pytest -from dify_trace_opik.config import OpikConfig -from dify_trace_opik.opik_trace import OpikDataTrace, _seed_to_uuid4, prepare_opik_uuid, wrap_dict, wrap_metadata -from sqlalchemy.orm import Session - -from core.ops.entities.trace_entity import ( - DatasetRetrievalTraceInfo, - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - SuggestedQuestionTraceInfo, - ToolTraceInfo, - TraceTaskName, - WorkflowTraceInfo, -) -from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey -from models import EndUser -from models.enums import EndUserType, MessageStatus -from tests.helpers.unit.core.ops.trace_fixtures import message_trace_info, tool_trace_info, workflow_trace_info - -"""Unit tests for Opik trace translation with SQLite-backed lookups.""" - - -def _dt() -> datetime: - return datetime(2024, 1, 1, 0, 0, 0, tzinfo=UTC) - - -class TestOpikTraceCases: - @staticmethod - @pytest.fixture - def opik_config(): - return OpikConfig( - project="test-project", workspace="test-workspace", url="https://cloud.opik.com/api/", api_key="api-key-123" - ) - - @staticmethod - @pytest.fixture - def trace_instance(opik_config, monkeypatch: pytest.MonkeyPatch): - mock_client = MagicMock() - monkeypatch.setattr("dify_trace_opik.opik_trace.Opik", lambda **kwargs: mock_client) - - instance = OpikDataTrace(opik_config) - return instance - - @staticmethod - def test_wrap_dict(): - assert wrap_dict("input", {"a": 1}) == {"a": 1} - assert wrap_dict("input", "hello") == {"input": "hello"} - - @staticmethod - def test_wrap_metadata(): - assert wrap_metadata({"a": 1}, b=2) == {"a": 1, "b": 2, "created_from": "dify"} - - @staticmethod - def test_prepare_opik_uuid(): - # Test with valid datetime and uuid string - dt = datetime(2024, 1, 1) - uuid_str = "b3e8e918-472e-4b69-8051-12502c34fc07" - result = prepare_opik_uuid(dt, uuid_str) - assert result is not None - # We won't test the exact uuid7 value but just that it returns a string id - - # Test with None dt and uuid_str - result = prepare_opik_uuid(None, None) - assert result is not None - - @staticmethod - def test_init(opik_config, monkeypatch: pytest.MonkeyPatch): - mock_opik = MagicMock() - monkeypatch.setattr("dify_trace_opik.opik_trace.Opik", mock_opik) - monkeypatch.setenv("FILES_URL", "http://test.url") - - instance = OpikDataTrace(opik_config) - - mock_opik.assert_called_once_with( - project_name=opik_config.project, - workspace=opik_config.workspace, - host=opik_config.url, - api_key=opik_config.api_key, - ) - assert instance.file_base_url == "http://test.url" - assert instance.project == opik_config.project - - @staticmethod - def test_trace_dispatch(trace_instance, monkeypatch: pytest.MonkeyPatch): - methods = [ - "workflow_trace", - "message_trace", - "moderation_trace", - "suggested_question_trace", - "dataset_retrieval_trace", - "tool_trace", - "generate_name_trace", - ] - mocks = {method: MagicMock() for method in methods} - for method, m in mocks.items(): - monkeypatch.setattr(trace_instance, method, m) - - # WorkflowTraceInfo - info = workflow_trace_info() - trace_instance.trace(info) - mocks["workflow_trace"].assert_called_once_with(info) - - # MessageTraceInfo - info = message_trace_info() - trace_instance.trace(info) - mocks["message_trace"].assert_called_once_with(info) - - # ModerationTraceInfo - info = ModerationTraceInfo(metadata={}, flagged=False, action="allow", preset_response="", query="hello") - trace_instance.trace(info) - mocks["moderation_trace"].assert_called_once_with(info) - - # SuggestedQuestionTraceInfo - info = SuggestedQuestionTraceInfo(metadata={}, total_tokens=0, suggested_question=[], level="info") - trace_instance.trace(info) - mocks["suggested_question_trace"].assert_called_once_with(info) - - # DatasetRetrievalTraceInfo - info = DatasetRetrievalTraceInfo(metadata={}) - trace_instance.trace(info) - mocks["dataset_retrieval_trace"].assert_called_once_with(info) - - # ToolTraceInfo - info = tool_trace_info() - trace_instance.trace(info) - mocks["tool_trace"].assert_called_once_with(info) - - # GenerateNameTraceInfo - info = GenerateNameTraceInfo(tenant_id="tenant-1", metadata={}) - trace_instance.trace(info) - mocks["generate_name_trace"].assert_called_once_with(info) - - @staticmethod - @pytest.mark.parametrize("sqlite3_session", [()], indirect=True) - def test_workflow_trace_with_message_id( - trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session - ) -> None: - # Define constants for better readability - workflow_id = "fb05c7cd-6cec-4add-8a84-df03a408b4ce" - workflow_run_id = "33c67568-7a8a-450e-8916-a5f135baeaef" - message_id = "04ec3956-85f3-488a-8539-1017251dc8c6" - conversation_id = "d3d01066-23ae-4830-9ce4-eb5640b42a7e" - trace_id = "bf26d929-6f15-4c2f-9abc-761c217056f3" - workflow_app_log_id = "ca0e018e-edd4-43fb-a05a-ea001ca8ef4b" - llm_node_id = "80d7dfa8-08f4-4ab7-aa37-0ca7d27207e3" - code_node_id = "b9cd9a7b-c534-4aa9-b5da-efd454140900" - - trace_info = WorkflowTraceInfo( - workflow_id=workflow_id, - tenant_id="tenant-1", - workflow_run_id=workflow_run_id, - workflow_run_elapsed_time=1.0, - workflow_run_status="succeeded", - workflow_run_inputs={"input": "hi"}, - workflow_run_outputs={"output": "hello"}, - workflow_run_version="1.0", - message_id=message_id, - conversation_id=conversation_id, - total_tokens=100, - file_list=[], - query="hi", - start_time=_dt(), - end_time=_dt() + timedelta(seconds=1), - trace_id=trace_id, - metadata={"app_id": "app-1", "user_id": "user-1"}, - workflow_app_log_id=workflow_app_log_id, - error="", - ) - - monkeypatch.setattr( - "dify_trace_opik.opik_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - - node_llm = MagicMock() - node_llm.id = llm_node_id - node_llm.title = "LLM Node" - node_llm.node_type = BuiltinNodeTypes.LLM - node_llm.status = "succeeded" - node_llm.process_data = { - "model_mode": "chat", - "model_name": "gpt-4", - "model_provider": "openai", - "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, - } - node_llm.inputs = {"prompts": "p"} - node_llm.outputs = {"text": "t"} - node_llm.created_at = _dt() - node_llm.elapsed_time = 0.5 - node_llm.metadata = {"foo": "bar"} - - node_other = MagicMock() - node_other.id = code_node_id - node_other.title = "Other Node" - node_other.node_type = BuiltinNodeTypes.CODE - node_other.status = "failed" - node_other.process_data = None - node_other.inputs = {"code": "print"} - node_other.outputs = {"result": "ok"} - node_other.created_at = None - node_other.elapsed_time = 0.2 - node_other.metadata = {WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS.value: 10} - - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [node_llm, node_other] - - mock_factory = MagicMock() - mock_factory.create_workflow_node_execution_repository.return_value = repo - monkeypatch.setattr("dify_trace_opik.opik_trace.DifyCoreRepositoryFactory", mock_factory) - - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - trace_instance.add_trace = MagicMock() - trace_instance.add_span = MagicMock() - - trace_instance.workflow_trace(trace_info) - - trace_instance.add_trace.assert_called_once() - trace_data = trace_instance.add_trace.call_args[1].get( - "opik_trace_data", trace_instance.add_trace.call_args[0][0] - ) - assert trace_data["name"] == TraceTaskName.MESSAGE_TRACE - assert "message" in trace_data["tags"] - assert "workflow" in trace_data["tags"] - - assert trace_instance.add_span.call_count >= 1 - - @staticmethod - @pytest.mark.parametrize("sqlite3_session", [()], indirect=True) - def test_workflow_trace_no_message_id( - trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session - ) -> None: - # Define constants for better readability - workflow_id = "f0708b36-b1d7-42b3-a876-1d01b7d8f1a3" - workflow_run_id = "d42ec285-c2fd-4248-8866-5c9386b101ac" - conversation_id = "88a17f2e-9436-4472-bab9-4b1601d5af3c" - workflow_app_log_id = "41780d0d-ffba-4220-bc0c-401e4c89cdfb" - - trace_info = WorkflowTraceInfo( - workflow_id=workflow_id, - tenant_id="tenant-1", - workflow_run_id=workflow_run_id, - workflow_run_elapsed_time=1.0, - workflow_run_status="succeeded", - workflow_run_inputs={}, - workflow_run_outputs={}, - workflow_run_version="1.0", - total_tokens=0, - file_list=[], - query="", - message_id=None, - conversation_id=conversation_id, - start_time=_dt(), - end_time=_dt(), - trace_id=None, - metadata={"app_id": "app-1"}, - workflow_app_log_id=workflow_app_log_id, - error="", - ) - - monkeypatch.setattr( - "dify_trace_opik.opik_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [] - mock_factory = MagicMock() - mock_factory.create_workflow_node_execution_repository.return_value = repo - monkeypatch.setattr("dify_trace_opik.opik_trace.DifyCoreRepositoryFactory", mock_factory) - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - trace_instance.add_trace = MagicMock() - trace_instance.workflow_trace(trace_info) - - trace_instance.add_trace.assert_called_once() - - @staticmethod - @pytest.mark.parametrize("sqlite3_session", [()], indirect=True) - def test_workflow_trace_missing_app_id( - trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session - ) -> None: - trace_info = WorkflowTraceInfo( - workflow_id="5745f1b8-f8e6-4859-8110-996acb6c8d6a", - tenant_id="tenant-1", - workflow_run_id="46f53304-1659-464b-bee5-116585f0bec8", - workflow_run_elapsed_time=1.0, - workflow_run_status="succeeded", - workflow_run_inputs={}, - workflow_run_outputs={}, - workflow_run_version="1.0", - total_tokens=0, - file_list=[], - query="", - message_id=None, - conversation_id="83f86b89-caef-4de8-a0f9-f164eddae1ea", - start_time=_dt(), - end_time=_dt(), - metadata={}, - workflow_app_log_id="339760b2-4b94-4532-8c81-133a97e4680e", - error="", - ) - monkeypatch.setattr( - "dify_trace_opik.opik_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - - with pytest.raises(ValueError, match="No app_id found in trace_info metadata"): - trace_instance.workflow_trace(trace_info) - - @staticmethod - def test_message_trace_basic(trace_instance, monkeypatch: pytest.MonkeyPatch): - # Define constants for better readability - message_data_id = "e3a26712-8cac-4a25-94a4-a3bff21ee3ab" - conversation_id = "9d3f3751-7521-4c19-9307-20e3cf6789a3" - message_trace_id = "710ace2f-bca8-41be-858c-54da42742a77" - opit_trace_id = "f7dfd978-0d10-4549-8abf-00f2cbc49d2c" - - message_data = MagicMock() - message_data.id = message_data_id - message_data.from_account_id = "acc-1" - message_data.from_end_user_id = None - message_data.provider_response_latency = 0.5 - message_data.conversation_id = conversation_id - message_data.total_price = 0.01 - message_data.model_id = "gpt-4" - message_data.answer = "hello" - message_data.status = MessageStatus.NORMAL - message_data.error = None - - trace_info = MessageTraceInfo( - message_id=message_trace_id, - message_data=message_data, - inputs={"query": "hi"}, - outputs={"answer": "hello"}, - message_tokens=10, - answer_tokens=20, - total_tokens=30, - start_time=_dt(), - end_time=_dt() + timedelta(seconds=1), - trace_id=opit_trace_id, - metadata={"foo": "bar"}, - conversation_mode="chat", - conversation_model="gpt-4", - file_list=[], - error=None, - message_file_data=MagicMock(url="test.png"), - ) - - trace_instance.add_trace = MagicMock(return_value=MagicMock(id="trace_id_1")) - trace_instance.add_span = MagicMock() - - trace_instance.message_trace(trace_info) - - trace_instance.add_trace.assert_called_once() - trace_instance.add_span.assert_called_once() - - @staticmethod - @pytest.mark.parametrize("sqlite3_session", [(EndUser,)], indirect=True) - def test_message_trace_with_end_user( - trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session - ) -> None: - write_session = sqlite3_session - message_data = MagicMock() - message_data.id = "85411059-79fb-4deb-a76c-c2e215f1b97e" - message_data.from_account_id = "acc-1" - message_data.from_end_user_id = "end-user-1" - message_data.conversation_id = "7d9f96d8-3be2-4e93-9c0e-922ff98dccc6" - message_data.status = MessageStatus.NORMAL - message_data.model_id = "gpt-4" - message_data.error = "" - message_data.answer = "hello" - message_data.total_price = 0.0 - message_data.provider_response_latency = 0.1 - - trace_info = MessageTraceInfo( - message_id="6bff35c7-33b7-4acb-ba21-44569a0327d0", - message_data=message_data, - inputs={}, - outputs={}, - message_tokens=0, - answer_tokens=0, - total_tokens=0, - start_time=_dt(), - end_time=_dt(), - metadata={}, - conversation_mode="chat", - conversation_model="gpt-4", - file_list=["url1"], - error=None, - ) - - end_user = EndUser( - id="end-user-1", - tenant_id="tenant-1", - app_id="app-1", - type=EndUserType.BROWSER, - session_id="session-id-123", - ) - write_session.add(end_user) - write_session.commit() - - with Session(write_session.get_bind()) as read_session: - monkeypatch.setattr( - "dify_trace_opik.opik_trace.db", - SimpleNamespace(engine=read_session.get_bind(), session=read_session), - ) - - trace_instance.add_trace = MagicMock(return_value=MagicMock(id="trace_id_2")) - trace_instance.add_span = MagicMock() - - trace_instance.message_trace(trace_info) - - trace_data = trace_instance.add_trace.call_args[0][0] - assert trace_data["metadata"]["user_id"] == "acc-1" - assert trace_data["metadata"]["end_user_id"] == "session-id-123" - - @staticmethod - def test_message_trace_none_data(trace_instance): - trace_info = SimpleNamespace(message_data=None, file_list=[], message_file_data=None, metadata={}) - trace_instance.add_trace = MagicMock() - trace_instance.message_trace(trace_info) - trace_instance.add_trace.assert_not_called() - - @staticmethod - def test_moderation_trace(trace_instance): - message_data = MagicMock() - message_data.created_at = _dt() - message_data.updated_at = _dt() - - trace_info = ModerationTraceInfo( - message_id="489d0dfd-065c-4106-8f9c-daded296c92d", - message_data=message_data, - inputs={"q": "hi"}, - action="stop", - flagged=True, - preset_response="blocked", - start_time=None, - end_time=None, - metadata={"foo": "bar"}, - trace_id="6f16cf18-9f4b-4955-8b6b-43cfa10978fc", - query="hi", - ) - - trace_instance.add_span = MagicMock() - trace_instance.moderation_trace(trace_info) - - trace_instance.add_span.assert_called_once() - span_data = trace_instance.add_span.call_args[0][0] - assert span_data["name"] == TraceTaskName.MODERATION_TRACE - assert span_data["output"]["flagged"] is True - - @staticmethod - def test_moderation_trace_none(trace_instance): - trace_info = ModerationTraceInfo( - message_id="cd732e4e-37f1-4c7e-8c64-820308bedcbf", - message_data=None, - inputs={}, - action="s", - flagged=False, - preset_response="", - query="", - metadata={}, - ) - trace_instance.add_span = MagicMock() - trace_instance.moderation_trace(trace_info) - trace_instance.add_span.assert_not_called() - - @staticmethod - def test_suggested_question_trace(trace_instance): - message_data = MagicMock() - message_data.created_at = _dt() - message_data.updated_at = _dt() - - trace_info = SuggestedQuestionTraceInfo( - message_id="7de55bda-a91d-477e-98ab-85c53c438469", - message_data=message_data, - inputs="hi", - suggested_question=["q1"], - total_tokens=10, - level="info", - start_time=_dt(), - end_time=_dt(), - metadata={}, - trace_id="a6687292-68c7-42ba-ae51-285579944d7b", - ) - - trace_instance.add_span = MagicMock() - trace_instance.suggested_question_trace(trace_info) - - trace_instance.add_span.assert_called_once() - span_data = trace_instance.add_span.call_args[0][0] - assert span_data["name"] == TraceTaskName.SUGGESTED_QUESTION_TRACE - - @staticmethod - def test_suggested_question_trace_none(trace_instance): - trace_info = SuggestedQuestionTraceInfo( - message_id="23696fc5-7e7f-46ec-bce8-1adc3c7f297d", - message_data=None, - inputs={}, - suggested_question=[], - total_tokens=0, - level="i", - metadata={}, - ) - trace_instance.add_span = MagicMock() - trace_instance.suggested_question_trace(trace_info) - trace_instance.add_span.assert_not_called() - - @staticmethod - def test_dataset_retrieval_trace(trace_instance): - message_data = MagicMock() - message_data.created_at = _dt() - message_data.updated_at = _dt() - - trace_info = DatasetRetrievalTraceInfo( - message_id="3e1a819f-c391-4950-adfd-96f82e5419a1", - message_data=message_data, - inputs="query", - documents=[{"id": "doc1"}], - start_time=None, - end_time=None, - metadata={}, - trace_id="41361000-e9be-4d11-b5e4-ab27ce0817d6", - ) - - trace_instance.add_span = MagicMock() - trace_instance.dataset_retrieval_trace(trace_info) - - trace_instance.add_span.assert_called_once() - span_data = trace_instance.add_span.call_args[0][0] - assert span_data["name"] == TraceTaskName.DATASET_RETRIEVAL_TRACE - - @staticmethod - def test_dataset_retrieval_trace_none(trace_instance): - trace_info = DatasetRetrievalTraceInfo( - message_id="35d6d44c-bccb-4e6e-8bd8-859257723ea8", message_data=None, inputs={}, documents=[], metadata={} - ) - trace_instance.add_span = MagicMock() - trace_instance.dataset_retrieval_trace(trace_info) - trace_instance.add_span.assert_not_called() - - @staticmethod - def test_tool_trace(trace_instance): - trace_info = ToolTraceInfo( - message_id="99db92c4-2254-496a-b5cc-18153315ce35", - message_data=MagicMock(), - inputs={}, - outputs={}, - tool_name="my_tool", - tool_inputs={"a": 1}, - tool_outputs="result_string", - time_cost=0.1, - start_time=_dt(), - end_time=_dt(), - metadata={}, - trace_id="a15a5fcb-7ffd-4458-8330-208f4cb1f796", - tool_config={}, - tool_parameters={}, - error="some error", - ) - - trace_instance.add_span = MagicMock() - trace_instance.tool_trace(trace_info) - - trace_instance.add_span.assert_called_once() - span_data = trace_instance.add_span.call_args[0][0] - assert span_data["name"] == "my_tool" - - @staticmethod - def test_generate_name_trace(trace_instance): - trace_info = GenerateNameTraceInfo( - inputs={"q": "hi"}, - outputs={"name": "new"}, - tenant_id="tenant-1", - conversation_id="271fe28f-6b86-416b-8d6b-bbbbfa9db791", - start_time=_dt(), - end_time=_dt(), - metadata={"921f010e-6878-4831-ae6b-271bf68c56fb": 1}, - ) - - trace_instance.add_trace = MagicMock(return_value=MagicMock(id="trace_id_3")) - trace_instance.add_span = MagicMock() - - trace_instance.generate_name_trace(trace_info) - - trace_instance.add_trace.assert_called_once() - trace_instance.add_span.assert_called_once() - - trace_data = trace_instance.add_trace.call_args[0][0] - assert trace_data["name"] == TraceTaskName.GENERATE_NAME_TRACE - - span_data = trace_instance.add_span.call_args[0][0] - assert span_data["trace_id"] == "trace_id_3" - - @staticmethod - def test_add_trace_success(trace_instance): - trace_data = {"id": "t1", "name": "trace"} - trace_instance.opik_client.trace.return_value = MagicMock(id="t1") - trace = trace_instance.add_trace(trace_data) - trace_instance.opik_client.trace.assert_called_once() - assert trace.id == "t1" - - @staticmethod - def test_add_trace_error(trace_instance): - trace_instance.opik_client.trace.side_effect = Exception("error") - trace_data = {"id": "t1", "name": "trace"} - with pytest.raises(ValueError, match="Opik Failed to create trace: error"): - trace_instance.add_trace(trace_data) - - @staticmethod - def test_add_span_success(trace_instance): - span_data = {"id": "s1", "name": "span", "trace_id": "t1"} - trace_instance.add_span(span_data) - trace_instance.opik_client.span.assert_called_once() - - @staticmethod - def test_add_span_error(trace_instance): - trace_instance.opik_client.span.side_effect = Exception("error") - span_data = {"id": "s1", "name": "span", "trace_id": "t1"} - with pytest.raises(ValueError, match="Opik Failed to create span: error"): - trace_instance.add_span(span_data) - - @staticmethod - def test_api_check_success(trace_instance): - trace_instance.opik_client.auth_check.return_value = True - assert trace_instance.api_check() is True - - @staticmethod - def test_api_check_error(trace_instance): - trace_instance.opik_client.auth_check.side_effect = Exception("fail") - with pytest.raises(ValueError, match="Opik API check failed: fail"): - trace_instance.api_check() - - @staticmethod - def test_get_project_url_success(trace_instance): - trace_instance.opik_client.get_project_url.return_value = "http://project.url" - assert trace_instance.get_project_url() == "http://project.url" - trace_instance.opik_client.get_project_url.assert_called_once_with(project_name=trace_instance.project) - - @staticmethod - def test_get_project_url_error(trace_instance): - trace_instance.opik_client.get_project_url.side_effect = Exception("fail") - with pytest.raises(ValueError, match="Opik get run url failed: fail"): - trace_instance.get_project_url() - - @staticmethod - @pytest.mark.parametrize("sqlite3_session", [()], indirect=True) - def test_workflow_trace_usage_extraction_error_fixed( - trace_instance, - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - sqlite3_session: Session, - ) -> None: - trace_info = WorkflowTraceInfo( - workflow_id="86a52565-4a6b-4a1b-9bfd-98e4595e70de", - tenant_id="66e8e918-472e-4b69-8051-12502c34fc07", - workflow_run_id="8403965c-3344-4d22-a8fe-d8d55cee64d9", - workflow_run_elapsed_time=1.0, - workflow_run_status="s", - workflow_run_inputs={}, - workflow_run_outputs={}, - workflow_run_version="1", - total_tokens=0, - file_list=[], - query="", - message_id=None, - conversation_id="7a02cb9d-6949-4c59-a89d-f25bbc881e0e", - start_time=_dt(), - end_time=_dt(), - metadata={"app_id": "77e8e918-472e-4b69-8051-12502c34fc07"}, - workflow_app_log_id="82268424-e193-476c-a6db-f473388ee5fe", - error="", - ) - - node = MagicMock() - node.id = "88e8e918-472e-4b69-8051-12502c34fc07" - node.title = "LLM Node" - node.node_type = BuiltinNodeTypes.LLM - node.status = "succeeded" - - class BadDict(collections.UserDict): - @override - def get(self, key, default=None): - if key == "usage": - raise Exception("Usage extraction failed") - return super().get(key, default) - - node.process_data = BadDict({"model_mode": "chat", "model_name": "gpt-4", "usage": True, "prompts": ["p"]}) - node.created_at = _dt() - node.elapsed_time = 0.1 - node.metadata = {} - node.outputs = {} - - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [node] - mock_factory = MagicMock() - mock_factory.create_workflow_node_execution_repository.return_value = repo - monkeypatch.setattr("dify_trace_opik.opik_trace.DifyCoreRepositoryFactory", mock_factory) - monkeypatch.setattr( - "dify_trace_opik.opik_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - trace_instance.add_trace = MagicMock() - trace_instance.add_span = MagicMock() - - with caplog.at_level(logging.ERROR): - trace_instance.workflow_trace(trace_info) - - assert "Failed to extract usage" in caplog.text - assert trace_instance.add_span.call_count >= 1 - # Verify that at least one of the spans is for the LLM Node - span_names = [call.args[0]["name"] for call in trace_instance.add_span.call_args_list] - assert "LLM Node" in span_names - - -"""Tests for OpikDataTrace workflow_trace changes. - -Covers: -- _seed_to_uuid4 helper: produces valid UUID4 strings deterministically -- prepare_opik_uuid helper: basic contract -- workflow_trace without message_id now creates a root span parented to None -- workflow_trace without message_id: node spans parent to root_span_id (not workflow_app_log_id) -- workflow_trace with message_id still creates root span keyed on workflow_run_id (unchanged path) -""" - - -# A stable UUID4 used as the workflow_run_id throughout all tests. -_WORKFLOW_RUN_ID = "a3f1b2c4-d5e6-4f78-9a0b-c1d2e3f4a5b6" - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _make_workflow_trace_info( - *, - message_id: str | None = None, - workflow_app_log_id: str | None = None, - workflow_run_id: str = _WORKFLOW_RUN_ID, -) -> WorkflowTraceInfo: - """Return a minimal WorkflowTraceInfo suitable for unit testing.""" - return WorkflowTraceInfo( - message_id=message_id, - workflow_id="wf-id", - tenant_id="tenant-id", - workflow_run_id=workflow_run_id, - workflow_app_log_id=workflow_app_log_id, - workflow_run_elapsed_time=1.5, - workflow_run_status="succeeded", - workflow_run_inputs={"query": "hello"}, - workflow_run_outputs={"result": "world"}, - workflow_run_version="1", - total_tokens=42, - file_list=[], - query="hello", - start_time=datetime(2025, 1, 1, 12, 0, 0), - end_time=datetime(2025, 1, 1, 12, 0, 1), - metadata={"app_id": "app-abc"}, - conversation_id=None, - ) - - -def _make_opik_trace_instance() -> OpikDataTrace: - """Construct an OpikDataTrace with the Opik SDK client mocked out.""" - with patch("dify_trace_opik.opik_trace.Opik"): - from dify_trace_opik.config import OpikConfig - - config = OpikConfig(api_key="key", project="test-project", url="https://www.comet.com/opik/api/") - instance = OpikDataTrace(config) - - instance.add_trace = MagicMock(return_value=MagicMock(id="mock-trace-id")) - instance.add_span = MagicMock() - instance.get_service_account_with_tenant = MagicMock(return_value=MagicMock()) - return instance - - -def _add_trace_mock(instance: OpikDataTrace) -> MagicMock: - return cast(MagicMock, instance.add_trace) - - -def _add_span_mock(instance: OpikDataTrace) -> MagicMock: - return cast(MagicMock, instance.add_span) - - -class TestOpikTraceCasesUnitTests: - # --------------------------------------------------------------------------- - # _seed_to_uuid4 - # --------------------------------------------------------------------------- - - class TestSeedToUuid4: - def test_returns_valid_uuid4_string(self): - result = _seed_to_uuid4("some-arbitrary-seed") - parsed = uuid.UUID(result) - assert parsed.version == 4 - - def test_is_deterministic(self): - assert _seed_to_uuid4("seed-abc") == _seed_to_uuid4("seed-abc") - - def test_different_seeds_give_different_results(self): - assert _seed_to_uuid4("seed-1") != _seed_to_uuid4("seed-2") - - def test_workflow_run_id_with_root_suffix_is_valid_uuid4(self): - """The primary use-case: deriving a root-span UUID from workflow_run_id + '-root'.""" - seed = _WORKFLOW_RUN_ID + "-root" - result = _seed_to_uuid4(seed) - parsed = uuid.UUID(result) - assert parsed.version == 4 - - def test_seed_and_seed_root_produce_different_uuids(self): - """Root span UUID must differ from the base workflow UUID to avoid ID collisions.""" - base = _seed_to_uuid4(_WORKFLOW_RUN_ID) - with_root = _seed_to_uuid4(_WORKFLOW_RUN_ID + "-root") - assert base != with_root - - # --------------------------------------------------------------------------- - # prepare_opik_uuid - # --------------------------------------------------------------------------- - - class TestPrepareOpikUuid: - def test_is_deterministic(self): - dt = datetime(2025, 6, 15, 10, 30, 0) - uid = str(uuid.uuid4()) - assert prepare_opik_uuid(dt, uid) == prepare_opik_uuid(dt, uid) - - def test_different_uuids_give_different_results(self): - dt = datetime(2025, 6, 15, 10, 30, 0) - assert prepare_opik_uuid(dt, str(uuid.uuid4())) != prepare_opik_uuid(dt, str(uuid.uuid4())) - - def test_none_datetime_does_not_raise(self): - assert prepare_opik_uuid(None, str(uuid.uuid4())) is not None - - def test_none_uuid_does_not_raise(self): - assert prepare_opik_uuid(datetime(2025, 1, 1), None) is not None - - # --------------------------------------------------------------------------- - # workflow_trace — no message_id (new code path) - # --------------------------------------------------------------------------- - - class TestWorkflowTraceWithoutMessageId: - def _run(self, trace_info: WorkflowTraceInfo, node_executions: list | None = None): - instance = _make_opik_trace_instance() - fake_repo = MagicMock() - fake_repo.get_by_workflow_execution.return_value = node_executions or [] - - with ( - patch("dify_trace_opik.opik_trace.db") as mock_db, - patch("dify_trace_opik.opik_trace.sessionmaker"), - patch( - "dify_trace_opik.opik_trace.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - return_value=fake_repo, - ), - ): - mock_db.engine = MagicMock() - instance.workflow_trace(trace_info) - - return instance - - def _expected_root_span_id(self, trace_info: WorkflowTraceInfo): - return prepare_opik_uuid( - trace_info.start_time, - _seed_to_uuid4(trace_info.workflow_run_id + "-root"), - ) - - def test_root_span_is_created(self): - trace_info = _make_workflow_trace_info(message_id=None) - instance = self._run(trace_info) - assert _add_span_mock(instance).called - - def test_root_span_id_matches_expected(self): - trace_info = _make_workflow_trace_info(message_id=None) - instance = self._run(trace_info) - - expected = self._expected_root_span_id(trace_info) - root_span_kwargs = _add_span_mock(instance).call_args_list[0][0][0] - assert root_span_kwargs["id"] == expected - - def test_root_span_has_no_parent(self): - trace_info = _make_workflow_trace_info(message_id=None) - instance = self._run(trace_info) - - root_span_kwargs = _add_span_mock(instance).call_args_list[0][0][0] - assert root_span_kwargs["parent_span_id"] is None - - def test_trace_name_is_workflow_trace(self): - """Without message_id, the Opik trace itself should be named WORKFLOW_TRACE.""" - trace_info = _make_workflow_trace_info(message_id=None) - instance = self._run(trace_info) - - trace_kwargs = _add_trace_mock(instance).call_args_list[0][0][0] - assert trace_kwargs["name"] == TraceTaskName.WORKFLOW_TRACE - - def test_root_span_name_is_workflow_trace(self): - trace_info = _make_workflow_trace_info(message_id=None) - instance = self._run(trace_info) - - root_span_kwargs = _add_span_mock(instance).call_args_list[0][0][0] - assert root_span_kwargs["name"] == TraceTaskName.WORKFLOW_TRACE - - def test_root_span_has_workflow_tag(self): - trace_info = _make_workflow_trace_info(message_id=None) - instance = self._run(trace_info) - - root_span_kwargs = _add_span_mock(instance).call_args_list[0][0][0] - assert "workflow" in root_span_kwargs["tags"] - - def test_node_execution_spans_are_parented_to_root(self): - """Node spans must use root_span_id as parent, not any other ID.""" - trace_info = _make_workflow_trace_info(message_id=None) - expected_root_span_id = self._expected_root_span_id(trace_info) - - node_exec = MagicMock() - node_exec.id = str(uuid.uuid4()) - node_exec.title = "LLM Node" - node_exec.node_type = "llm" - node_exec.status = "succeeded" - node_exec.process_data = {} - node_exec.inputs = {"prompt": "hi"} - node_exec.outputs = {"text": "hello"} - node_exec.created_at = datetime(2025, 1, 1, 12, 0, 0) - node_exec.elapsed_time = 0.5 - node_exec.metadata = {} - - instance = self._run(trace_info, node_executions=[node_exec]) - - # call_args_list[0] = root span, [1] = node execution span - add_span = _add_span_mock(instance) - assert add_span.call_count == 2 - node_span_kwargs = add_span.call_args_list[1][0][0] - assert node_span_kwargs["parent_span_id"] == expected_root_span_id - - def test_node_span_not_parented_to_workflow_app_log_id(self): - """Old behaviour derived parent from workflow_app_log_id; that must no longer apply.""" - trace_info = _make_workflow_trace_info( - message_id=None, - workflow_app_log_id=str(uuid.uuid4()), - ) - - node_exec = MagicMock() - node_exec.id = str(uuid.uuid4()) - node_exec.title = "Tool Node" - node_exec.node_type = "tool" - node_exec.status = "succeeded" - node_exec.process_data = {} - node_exec.inputs = {} - node_exec.outputs = {} - node_exec.created_at = datetime(2025, 1, 1, 12, 0, 0) - node_exec.elapsed_time = 0.2 - node_exec.metadata = {} - - instance = self._run(trace_info, node_executions=[node_exec]) - - old_parent_id = prepare_opik_uuid(trace_info.start_time, trace_info.workflow_app_log_id) - node_span_kwargs = _add_span_mock(instance).call_args_list[1][0][0] - assert node_span_kwargs["parent_span_id"] != old_parent_id - - def test_root_span_id_differs_from_trace_id(self): - """The root span must have a different ID from the Opik trace to maintain correct hierarchy.""" - trace_info = _make_workflow_trace_info(message_id=None) - dify_trace_id = trace_info.trace_id or trace_info.workflow_run_id - opik_trace_id = prepare_opik_uuid(trace_info.start_time, dify_trace_id) - root_span_id = self._expected_root_span_id(trace_info) - assert root_span_id != opik_trace_id - - # --------------------------------------------------------------------------- - # workflow_trace — with message_id (unchanged path, guard against regression) - # --------------------------------------------------------------------------- - - class TestWorkflowTraceWithMessageId: - _MESSAGE_ID = str(uuid.uuid4()) - - def _run(self, trace_info: WorkflowTraceInfo, node_executions: list | None = None): - instance = _make_opik_trace_instance() - fake_repo = MagicMock() - fake_repo.get_by_workflow_execution.return_value = node_executions or [] - - with ( - patch("dify_trace_opik.opik_trace.db") as mock_db, - patch("dify_trace_opik.opik_trace.sessionmaker"), - patch( - "dify_trace_opik.opik_trace.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - return_value=fake_repo, - ), - ): - mock_db.engine = MagicMock() - instance.workflow_trace(trace_info) - - return instance - - def test_trace_name_is_message_trace(self): - """With message_id, the Opik trace should be named MESSAGE_TRACE.""" - trace_info = _make_workflow_trace_info(message_id=self._MESSAGE_ID) - instance = self._run(trace_info) - - trace_kwargs = _add_trace_mock(instance).call_args_list[0][0][0] - assert trace_kwargs["name"] == TraceTaskName.MESSAGE_TRACE - - def test_root_span_uses_workflow_run_id_directly(self): - """When message_id is set, root_span_id = prepare_opik_uuid(start_time, workflow_run_id).""" - trace_info = _make_workflow_trace_info(message_id=self._MESSAGE_ID) - instance = self._run(trace_info) - - expected_root_span_id = prepare_opik_uuid(trace_info.start_time, trace_info.workflow_run_id) - root_span_kwargs = _add_span_mock(instance).call_args_list[0][0][0] - assert root_span_kwargs["id"] == expected_root_span_id - - def test_root_span_id_differs_from_no_message_id_case(self): - """The two branches must produce different root span IDs for the same workflow_run_id.""" - id_with_message = prepare_opik_uuid( - datetime(2025, 1, 1, 12, 0, 0), - _WORKFLOW_RUN_ID, - ) - id_without_message = prepare_opik_uuid( - datetime(2025, 1, 1, 12, 0, 0), - _seed_to_uuid4(_WORKFLOW_RUN_ID + "-root"), - ) - assert id_with_message != id_without_message - - def test_node_spans_parented_to_workflow_run_root_span(self): - """Node spans must still parent to root_span_id derived from workflow_run_id.""" - trace_info = _make_workflow_trace_info(message_id=self._MESSAGE_ID) - expected_root_span_id = prepare_opik_uuid(trace_info.start_time, trace_info.workflow_run_id) - - node_exec = MagicMock() - node_exec.id = str(uuid.uuid4()) - node_exec.title = "LLM" - node_exec.node_type = "llm" - node_exec.status = "succeeded" - node_exec.process_data = {} - node_exec.inputs = {} - node_exec.outputs = {} - node_exec.created_at = datetime(2025, 1, 1, 12, 0, 0) - node_exec.elapsed_time = 0.3 - node_exec.metadata = {} - - instance = self._run(trace_info, node_executions=[node_exec]) - - node_span_kwargs = _add_span_mock(instance).call_args_list[1][0][0] - assert node_span_kwargs["parent_span_id"] == expected_root_span_id diff --git a/api/providers/trace/trace-opik/tests/unit_tests/src/dify_trace_opik/test_provider.py b/api/providers/trace/trace-opik/tests/unit_tests/src/dify_trace_opik/test_provider.py new file mode 100644 index 00000000000000..00a9890e8f6022 --- /dev/null +++ b/api/providers/trace/trace-opik/tests/unit_tests/src/dify_trace_opik/test_provider.py @@ -0,0 +1,526 @@ +import subprocess +import sys +from typing import Any +from urllib.parse import parse_qs, urlsplit + +import pytest +from dify_trace_opik.config import OpikConfig +from dify_trace_opik.provider import OpikProvider +from dify_trace_sdk.egress import EgressResponse +from dify_trace_sdk.errors import ( + EgressBlockedError, + EgressError, + ExportOutcomeUnknownError, + ExportRejectedError, + ExportRetryableError, + VerificationError, +) +from dify_trace_sdk.ids import span_id_for +from dify_trace_sdk.model import Cost, Span, SpanKind, SpanStatus, TokenUsage, Trace +from dify_trace_sdk.provider import SendResult +from dify_trace_sdk.testing import ( + MESSAGE_ID, + WORKFLOW_RUN_ID, + RecordedRequest, + RecordingEgress, + message_trace, + response, + tool_fragment, + workflow_trace, +) + +API_URL = "https://www.comet.com/opik/api/" +API_KEY = "opik-key-1" +CONFIG = OpikConfig(api_key=API_KEY, project="Dify", workspace="team-1", url=API_URL) +DEFAULT_WORKSPACE_CONFIG = OpikConfig(api_key=API_KEY) +HEADERS = { + "Comet-Workspace": "team-1", + "Accept": "application/json", + "Authorization": API_KEY, + "Content-Type": "application/json", +} +CREATED = response(201, b"") +NO_CONTENT = response(204, b"") +FAILURE = {"exception_type": "DifyExecutionError", "message": "Tool failed", "traceback": ""} + +# The legacy agent_trace exporter assigned these IDs to the sample traces; keeping +# them lets units sent before and after the migration join the same Opik trace. +MESSAGE_TRACE_ID = "01a0f67a-3800-73b3-8caf-ef474b73d890" +MESSAGE_SPAN_ID = "01a0f67a-3800-7312-8296-b10f658726ef" +MESSAGE_LLM_SPAN_ID = "01a0f67a-3800-712a-9d04-7e6b84a30a0c" +WORKFLOW_TRACE_ID = "01a0f67a-3800-7875-8be5-d677d1b69abc" +WORKFLOW_SPAN_ID = "01a0f67a-3800-7c27-98a2-cb7f6dafe47e" +WORKFLOW_LLM_SPAN_ID = "01a0f67a-3be8-7fb9-b3b4-887648ce8273" +WORKFLOW_TOOL_SPAN_ID = "01a0f67a-43b8-7ed6-a7e7-e44532c25e28" +TOOL_SPAN_ID = "01a0f67a-39f4-777f-b4fd-452b9ab6b95f" + + +def _provider(*responses: EgressResponse | Exception) -> tuple[OpikProvider, RecordingEgress]: + egress = RecordingEgress(responses, grpc_target=None) + return OpikProvider(egress), egress + + +def _json(request: RecordedRequest) -> dict[str, Any]: + body = request.json() + assert isinstance(body, dict) + return body + + +def test_importing_the_provider_leaves_the_opik_sdk_unloaded() -> None: + # Importing any opik module replaces the process's Sentry client with Comet's. + script = "import sys, dify_trace_opik.provider; print(sorted(n for n in sys.modules if n.split('.')[0] == 'opik'))" + result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=True) + assert result.stdout.strip() == "[]" + + +def test_verify_checks_access_to_a_named_workspace() -> None: + provider, egress = _provider(NO_CONTENT) + config = CONFIG.model_copy(update={"url": "https://opik.example.com/prefix/api/"}) + + assert provider.verify(config) == config + + (sent,) = egress.requests + assert (sent.method, sent.url) == ("POST", "https://opik.example.com/prefix/api/v1/private/auth") + assert sent.headers == HEADERS + assert sent.content == b"{}" + assert sent.sensitive_url is False + + +@pytest.mark.parametrize("workspace", [None, "", "default", " default "]) +def test_verify_resolves_the_default_workspace(workspace: str | None) -> None: + provider, egress = _provider(NO_CONTENT, response(200, {"workspace_name": "resolved-team"})) + + verified = provider.verify(OpikConfig(api_key=API_KEY, workspace=workspace)) + + assert verified == OpikConfig(api_key=API_KEY, project="Default Project", workspace="resolved-team", url=API_URL) + access, lookup = egress.requests + assert (access.method, access.headers["Comet-Workspace"]) == ("POST", "default") + assert (lookup.method, lookup.url, lookup.content) == ("GET", f"{API_URL}v1/private/auth/workspace", None) + assert lookup.headers == {"Comet-Workspace": "default", "Accept": "application/json", "Authorization": API_KEY} + + +def test_verify_without_an_api_key_sends_no_authorization() -> None: + provider, egress = _provider(NO_CONTENT) + + provider.verify(OpikConfig(workspace="local", url="http://selfhost.test:5173/api/")) + + (sent,) = egress.requests + assert sent.headers == { + "Comet-Workspace": "local", + "Accept": "application/json", + "Content-Type": "application/json", + } + + +@pytest.mark.parametrize( + ("config", "outcomes", "code"), + [ + (CONFIG, [response(401, {"message": "private-vendor-error"})], "authentication_failed"), + (CONFIG, [response(403, {})], "permission_denied"), + (CONFIG, [response(302, b"")], "endpoint_unavailable"), + (CONFIG, [response(500, {})], "endpoint_unavailable"), + (CONFIG, [response(200, b"Login")], "malformed_response"), + (CONFIG, [EgressError("refused", request_sent=False)], "connection_failed"), + (CONFIG, [EgressBlockedError("private address")], "connection_failed"), + (DEFAULT_WORKSPACE_CONFIG, [NO_CONTENT, response(401, {})], "authentication_failed"), + (DEFAULT_WORKSPACE_CONFIG, [NO_CONTENT, response(200, b"not json")], "malformed_response"), + (DEFAULT_WORKSPACE_CONFIG, [NO_CONTENT, response(200, ["team"])], "malformed_response"), + (DEFAULT_WORKSPACE_CONFIG, [NO_CONTENT, response(200, {"workspace_name": " "})], "malformed_response"), + (DEFAULT_WORKSPACE_CONFIG, [NO_CONTENT, response(200, {"workspace_name": 123})], "malformed_response"), + (DEFAULT_WORKSPACE_CONFIG, [NO_CONTENT, response(200, {"workspace_name": ".."})], "invalid_workspace"), + ], +) +def test_verify_failures_have_stable_codes( + config: OpikConfig, outcomes: list[EgressResponse | Exception], code: str +) -> None: + provider, egress = _provider(*outcomes) + with pytest.raises(VerificationError) as caught: + provider.verify(config) + assert caught.value.code == code + assert API_KEY not in str(caught.value) + assert len(egress.requests) == len(outcomes) + + +@pytest.mark.parametrize( + ("settings", "code"), + [ + ({"url": "https://private.test/api/?next=/api/"}, "invalid_endpoint"), + ({"url": "https://private.test/api/#/api/"}, "invalid_endpoint"), + ({"url": "https://private.test:invalid/api/"}, "invalid_endpoint"), + ({"url": "https://private.test/with space/api/"}, "invalid_endpoint"), + ({"url": "https://private.test/back\\slash/api/"}, "invalid_endpoint"), + ({"workspace": ".."}, "invalid_workspace"), + ({"workspace": "bad\nworkspace"}, "invalid_workspace"), + ({"project": "bad\nproject"}, "invalid_project"), + ], +) +def test_unsafe_settings_never_reach_the_network(settings: dict[str, str], code: str) -> None: + config = OpikConfig.model_validate({"api_key": API_KEY, **settings}) + provider, egress = _provider() + + with pytest.raises(VerificationError) as caught: + provider.verify(config) + assert caught.value.code == code + with pytest.raises(ExportRejectedError, match=code): + provider.send(config, message_trace()) + assert provider.project_url(config) is None + assert egress.requests == [] + + +def test_project_url_links_the_project_in_a_named_workspace() -> None: + provider, egress = _provider() + assert provider.project_url(CONFIG) == "https://www.comet.com/opik/team-1/redirect/projects?name=Dify" + + url = provider.project_url( + OpikConfig(url="https://private.test:7443/prefix/api/", workspace="team/name", project="A & B/? 中文") + ) + + assert url is not None + parsed = urlsplit(url) + assert (parsed.netloc, parsed.path) == ("private.test:7443", "/prefix/team%2Fname/redirect/projects") + assert parse_qs(parsed.query) == {"name": ["A & B/? 中文"]} + assert egress.requests == [] + + +@pytest.mark.parametrize("workspace", [None, "", "default"]) +def test_project_url_falls_back_to_the_console_for_the_default_workspace(workspace: str | None) -> None: + provider, _ = _provider() + config = OpikConfig(url="http://selfhost.test:5173/prefix/api/", workspace=workspace) + assert provider.project_url(config) == "http://selfhost.test:5173/prefix/" + + +def test_message_trace_creates_the_trace_then_its_spans() -> None: + provider, egress = _provider(CREATED, NO_CONTENT) + trace = message_trace() + message, llm = trace.spans + + result = provider.send(CONFIG, trace) + + assert result == SendResult(trace_id=MESSAGE_TRACE_ID, trace_url=None, span_urls={}) + create_trace, create_spans = egress.requests + assert (create_trace.method, create_trace.url) == ("POST", f"{API_URL}v1/private/traces") + assert (create_spans.method, create_spans.url) == ("POST", f"{API_URL}v1/private/spans/batch") + assert create_trace.headers == create_spans.headers == HEADERS + assert create_trace.sensitive_url is create_spans.sensitive_url is False + assert create_trace.json() == { + "id": MESSAGE_TRACE_ID, + "project_name": "Dify", + "name": "message", + "start_time": "2026-10-01T08:00:00Z", + "end_time": "2026-10-01T08:00:02Z", + "input": {"query": "What is Dify?"}, + "output": "An LLM application platform.", + "metadata": { + "message_id": MESSAGE_ID, + "app_id": "app-1", + "tenant_id": "tenant-1", + "dify_trace_id": trace.trace_id, + "user_id": "end-user-1", + }, + "thread_id": "conversation-1", + } + assert create_spans.json() == { + "spans": [ + { + "id": MESSAGE_SPAN_ID, + "project_name": "Dify", + "trace_id": MESSAGE_TRACE_ID, + "name": "message", + "type": "general", + "start_time": "2026-10-01T08:00:00Z", + "end_time": "2026-10-01T08:00:02Z", + "input": {"query": "What is Dify?"}, + "output": "An LLM application platform.", + "metadata": { + "message_id": MESSAGE_ID, + "app_id": "app-1", + "tenant_id": "tenant-1", + "dify_span_kind": "chain", + "dify_span_id": message.span_id, + "dify_trace_id": trace.trace_id, + "session_id": "conversation-1", + }, + }, + { + "id": MESSAGE_LLM_SPAN_ID, + "project_name": "Dify", + "trace_id": MESSAGE_TRACE_ID, + "parent_span_id": MESSAGE_SPAN_ID, + "name": "llm", + "type": "llm", + "start_time": "2026-10-01T08:00:00Z", + "end_time": "2026-10-01T08:00:02Z", + "input": [{"role": "user", "content": "What is Dify?"}], + "output": "An LLM application platform.", + "metadata": { + "message_id": MESSAGE_ID, + "dify_span_kind": "llm", + "dify_span_id": llm.span_id, + "dify_trace_id": trace.trace_id, + "session_id": "conversation-1", + }, + "model": "gpt-4o", + "provider": "openai", + "usage": {"prompt_tokens": 12, "completion_tokens": 8, "total_tokens": 20}, + "total_estimated_cost": 0.0004, + "ttft": 250.0, + }, + ] + } + + +def test_workflow_trace_keeps_hierarchy_errors_and_usage() -> None: + provider, egress = _provider(CREATED, NO_CONTENT) + trace = workflow_trace() + run, llm, tool = trace.spans + + assert provider.send(CONFIG, trace).trace_id == WORKFLOW_TRACE_ID + + create_trace, create_spans = egress.requests + assert create_trace.json() == { + "id": WORKFLOW_TRACE_ID, + "project_name": "Dify", + "name": "workflow", + "start_time": "2026-10-01T08:00:00Z", + "end_time": "2026-10-01T08:00:05Z", + "input": {"topic": "tracing"}, + "output": {}, + "metadata": { + "workflow_run_id": WORKFLOW_RUN_ID, + "app_id": "app-1", + "dify_trace_id": trace.trace_id, + "user_id": "account-1", + }, + "error_info": FAILURE, + "thread_id": WORKFLOW_RUN_ID, + } + assert create_spans.json() == { + "spans": [ + { + "id": WORKFLOW_SPAN_ID, + "project_name": "Dify", + "trace_id": WORKFLOW_TRACE_ID, + "name": "workflow", + "type": "general", + "start_time": "2026-10-01T08:00:00Z", + "end_time": "2026-10-01T08:00:05Z", + "input": {"topic": "tracing"}, + "output": {}, + "metadata": { + "workflow_run_id": WORKFLOW_RUN_ID, + "app_id": "app-1", + "dify_span_kind": "chain", + "dify_span_id": run.span_id, + "dify_trace_id": trace.trace_id, + "session_id": WORKFLOW_RUN_ID, + }, + "error_info": FAILURE, + }, + { + "id": WORKFLOW_LLM_SPAN_ID, + "project_name": "Dify", + "trace_id": WORKFLOW_TRACE_ID, + "parent_span_id": WORKFLOW_SPAN_ID, + "name": "llm_Summarize", + "type": "llm", + "start_time": "2026-10-01T08:00:01Z", + "end_time": "2026-10-01T08:00:03Z", + "input": [{"role": "user", "content": "Summarize tracing"}], + "output": {"text": "Tracing records execution."}, + "metadata": { + "node_execution_id": "node-execution-llm", + "node_type": "llm", + "dify_span_kind": "llm", + "dify_span_id": llm.span_id, + "dify_trace_id": trace.trace_id, + "session_id": WORKFLOW_RUN_ID, + }, + "model": "gpt-4o", + "provider": "openai", + "usage": {"prompt_tokens": 30, "completion_tokens": 10, "total_tokens": 40}, + }, + { + "id": WORKFLOW_TOOL_SPAN_ID, + "project_name": "Dify", + "trace_id": WORKFLOW_TRACE_ID, + "parent_span_id": WORKFLOW_LLM_SPAN_ID, + "name": "tool_Search", + "type": "tool", + "start_time": "2026-10-01T08:00:03Z", + "end_time": "2026-10-01T08:00:05Z", + "input": {"query": "tracing"}, + "output": {}, + "metadata": { + "node_execution_id": "node-execution-tool", + "node_type": "tool", + "dify_span_kind": "tool", + "dify_span_id": tool.span_id, + "dify_trace_id": trace.trace_id, + "session_id": WORKFLOW_RUN_ID, + }, + "error_info": FAILURE, + }, + ] + } + + +def test_fragment_joins_the_message_trace_without_creating_it() -> None: + provider, egress = _provider(NO_CONTENT) + fragment = tool_fragment() + + assert provider.send(CONFIG, fragment).trace_id == MESSAGE_TRACE_ID + + (create_spans,) = egress.requests + assert create_spans.url == f"{API_URL}v1/private/spans/batch" + assert create_spans.json() == { + "spans": [ + { + "id": TOOL_SPAN_ID, + "project_name": "Dify", + "trace_id": MESSAGE_TRACE_ID, + "name": "weather", + "type": "tool", + "start_time": "2026-10-01T08:00:00.500000Z", + "end_time": "2026-10-01T08:00:01Z", + "input": {"city": "Berlin"}, + "output": "Sunny", + "metadata": { + "message_id": MESSAGE_ID, + "dify_span_kind": "tool", + "dify_span_id": fragment.root.span_id, + "dify_trace_id": fragment.trace_id, + "session_id": "conversation-1", + }, + } + ] + } + + +def test_redelivery_sends_identical_requests() -> None: + first, first_egress = _provider(CREATED, NO_CONTENT) + second, second_egress = _provider(CREATED, NO_CONTENT) + + first.send(CONFIG, workflow_trace()) + second.send(CONFIG, workflow_trace()) + + assert [request.content for request in first_egress.requests] == [ + request.content for request in second_egress.requests + ] + + +def test_opik_specific_fields_are_sent_only_where_opik_accepts_them() -> None: + trace = message_trace() + message, llm = trace.spans + message = message.model_copy( + update={ + "inputs": [1, "two"], + "outputs": 42, + "model_name": "gpt-4o", + "usage": TokenUsage(prompt_tokens=1, completion_tokens=1, total_tokens=2), + "cost": Cost(amount=1.0, currency="USD"), + "time_to_first_token": 1.0, + } + ) + llm = llm.model_copy(update={"cost": Cost(amount=0.1, currency="EUR"), "time_to_first_token": None}) + provider, egress = _provider(CREATED, NO_CONTENT) + + provider.send(CONFIG, trace.model_copy(update={"spans": (message, llm)})) + + create_trace, create_spans = egress.requests + message_span, llm_span = _json(create_spans)["spans"] + assert (_json(create_trace)["input"], _json(create_trace)["output"]) == ('[1, "two"]', "42") + assert (message_span["input"], message_span["output"]) == ('[1, "two"]', "42") + assert not {"model", "provider", "usage", "total_estimated_cost", "ttft"} & message_span.keys() + assert "total_estimated_cost" not in llm_span + assert "ttft" not in llm_span + assert llm_span["usage"] == {"prompt_tokens": 12, "completion_tokens": 8, "total_tokens": 20} + + +def test_large_units_send_spans_in_batches() -> None: + run = workflow_trace().root.model_copy(update={"status": SpanStatus.OK, "error": None}) + nodes = tuple( + Span( + span_id=span_id_for(f"node-execution-{index}"), + parent_span_id=run.span_id, + name=f"node_{index}", + kind=SpanKind.TOOL, + start_time=run.start_time, + end_time=run.end_time, + status=SpanStatus.OK, + ) + for index in range(1000) + ) + trace = Trace( + trace_id=workflow_trace().trace_id, + external_trace_id=None, + session_id=None, + user_id=None, + parent=None, + spans=(run, *nodes), + ) + provider, egress = _provider(CREATED, NO_CONTENT, NO_CONTENT) + + provider.send(CONFIG, trace) + + _, first_batch, second_batch = egress.requests + assert [len(_json(batch)["spans"]) for batch in (first_batch, second_batch)] == [1000, 1] + + +def test_values_that_are_not_json_fail_before_anything_is_sent() -> None: + trace = message_trace() + message, llm = trace.spans + llm = llm.model_copy(update={"attributes": {"score": float("nan")}}) + provider, egress = _provider() + + with pytest.raises(ExportRejectedError): + provider.send(CONFIG, trace.model_copy(update={"spans": (message, llm)})) + assert egress.requests == [] + + +@pytest.mark.parametrize( + ("outcome", "expected"), + [ + (response(503, {}), ExportRetryableError), + (response(429, {}), ExportRetryableError), + (response(401, {}), ExportRejectedError), + (response(422, {"errors": ["private-vendor-error"]}), ExportRejectedError), + (EgressError("refused", request_sent=False), ExportRetryableError), + (EgressError("read timeout", request_sent=True), ExportOutcomeUnknownError), + (EgressBlockedError("private address"), ExportRejectedError), + ], +) +def test_trace_creation_failures_are_classified(outcome: EgressResponse | Exception, expected: type[Exception]) -> None: + provider, egress = _provider(outcome) + with pytest.raises(expected) as caught: + provider.send(CONFIG, message_trace()) + assert len(egress.requests) == 1 + assert API_KEY not in str(caught.value) + assert "private" not in str(caught.value) + + +@pytest.mark.parametrize( + "outcome", + [ + response(503, {}), + response(400, {}), + EgressError("refused", request_sent=False), + EgressBlockedError("private address"), + ], +) +def test_span_failures_after_the_trace_was_created_are_outcome_unknown(outcome: EgressResponse | Exception) -> None: + provider, egress = _provider(CREATED, outcome) + with pytest.raises(ExportOutcomeUnknownError): + provider.send(CONFIG, message_trace()) + assert len(egress.requests) == 2 + + +@pytest.mark.parametrize( + ("outcome", "expected"), + [(response(503, {}), ExportRetryableError), (response(400, {}), ExportRejectedError)], +) +def test_fragment_failures_are_classified_as_nothing_was_accepted( + outcome: EgressResponse, expected: type[Exception] +) -> None: + provider, _ = _provider(outcome) + with pytest.raises(expected): + provider.send(CONFIG, tool_fragment()) diff --git a/api/providers/trace/trace-sdk/pyproject.toml b/api/providers/trace/trace-sdk/pyproject.toml new file mode 100644 index 00000000000000..c6ea723773222c --- /dev/null +++ b/api/providers/trace/trace-sdk/pyproject.toml @@ -0,0 +1,18 @@ +[project] +name = "dify-trace-sdk" +version = "0.0.1" +dependencies = [ + "grpcio", + "httpx", + "opentelemetry-proto", + "protobuf", + "pydantic>=2", + "requests", +] +description = "Contracts shared by Dify trace producers, the trace pipeline, and trace providers." + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.setuptools.package-data] +dify_trace_sdk = ["py.typed"] diff --git a/api/providers/trace/trace-sdk/src/dify_trace_sdk/__init__.py b/api/providers/trace/trace-sdk/src/dify_trace_sdk/__init__.py new file mode 100644 index 00000000000000..dc15482aa672aa --- /dev/null +++ b/api/providers/trace/trace-sdk/src/dify_trace_sdk/__init__.py @@ -0,0 +1,5 @@ +"""Contracts shared by Dify trace producers, the trace pipeline, and trace providers. + +Providers depend on this package and their vendor SDK only. They never import +the Dify application (database, Redis, configuration, services). +""" diff --git a/api/providers/trace/trace-sdk/src/dify_trace_sdk/config.py b/api/providers/trace/trace-sdk/src/dify_trace_sdk/config.py new file mode 100644 index 00000000000000..228332494ceaa1 --- /dev/null +++ b/api/providers/trace/trace-sdk/src/dify_trace_sdk/config.py @@ -0,0 +1,99 @@ +"""Provider configuration models. + +A provider declares its configuration once, as a pydantic model. Field +metadata marks which values are secrets (encrypted at rest, never returned), +which are derived (resolved by ``verify``, never accepted from users), and +which name the destination secrets are sent to. +""" + +from dataclasses import dataclass +from typing import Any +from urllib.parse import urlsplit + +from pydantic import BaseModel, ConfigDict, Field +from pydantic.fields import FieldInfo +from pydantic_core import PydanticUndefined + +_SECRET = "x-dify-secret" +_DERIVED = "x-dify-derived" +_DESTINATION = "x-dify-destination" + + +def secret(*, default: Any = PydanticUndefined) -> Any: + """Declare a credential field: encrypted at rest and write-only through the API.""" + return Field(default=default, json_schema_extra={_SECRET: True}) + + +def derived() -> Any: + """Declare a field that ``verify`` resolves from the destination; users cannot set it.""" + return Field(default=None, json_schema_extra={_DERIVED: True}) + + +def destination(*, default: Any = PydanticUndefined) -> Any: + """Declare a field naming where secrets are sent; changing it requires entering the secrets again.""" + return Field(default=default, json_schema_extra={_DESTINATION: True}) + + +def _flag(field: FieldInfo, name: str) -> bool: + extra = field.json_schema_extra + return isinstance(extra, dict) and extra.get(name) is True + + +@dataclass(frozen=True) +class FormField: + name: str + secret: bool + required: bool + + +class ProviderConfig(BaseModel): + """Base class for provider configuration. + + Validation errors never echo input values: users paste credentials into + endpoint fields, and these errors reach logs and API responses. + """ + + model_config = ConfigDict(extra="ignore", frozen=True, hide_input_in_errors=True) + + @classmethod + def secret_fields(cls) -> frozenset[str]: + return frozenset(name for name, field in cls.model_fields.items() if _flag(field, _SECRET)) + + @classmethod + def derived_fields(cls) -> frozenset[str]: + return frozenset(name for name, field in cls.model_fields.items() if _flag(field, _DERIVED)) + + @classmethod + def destination_fields(cls) -> frozenset[str]: + return frozenset(name for name, field in cls.model_fields.items() if _flag(field, _DESTINATION)) + + @classmethod + def form_fields(cls) -> tuple[FormField, ...]: + """The user-editable fields, in declaration order.""" + return tuple( + FormField(name=name, secret=_flag(field, _SECRET), required=field.is_required()) + for name, field in cls.model_fields.items() + if not _flag(field, _DERIVED) + ) + + +def http_url(value: str, *, allowed_schemes: tuple[str, ...]) -> str: + """Validate a destination URL, keeping its path; reject credentials in the authority.""" + url = value.strip() + parsed = urlsplit(url) + if parsed.scheme not in allowed_schemes: + raise ValueError("URL must start with " + " or ".join(f"{scheme}://" for scheme in allowed_schemes)) + if not parsed.hostname or parsed.username is not None or parsed.password is not None: + raise ValueError("URL must name a host and must not contain credentials") + return url + + +def name_or_default(value: str | None, default: str) -> str: + return value.strip() if value and value.strip() else default + + +def integer_id(value: str) -> str: + stripped = value.strip() + if not (stripped.isascii() and stripped.isdigit()): + raise ValueError("ID must be a valid integer") + return stripped diff --git a/api/providers/trace/trace-sdk/src/dify_trace_sdk/egress.py b/api/providers/trace/trace-sdk/src/dify_trace_sdk/egress.py new file mode 100644 index 00000000000000..ad303761b204cd --- /dev/null +++ b/api/providers/trace/trace-sdk/src/dify_trace_sdk/egress.py @@ -0,0 +1,54 @@ +"""The only network access a provider has. + +The Dify API supplies an ``Egress`` that applies its outbound policy (SSRF +proxy, timeouts, response size limits, log redaction). Providers never open +sockets, HTTP clients, or gRPC channels of their own. +""" + +import json +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Protocol + +import grpc # type: ignore[import-untyped] + + +@dataclass(frozen=True) +class EgressResponse: + status_code: int + headers: Mapping[str, str] + """Header names are lowercase.""" + content: bytes + + def json(self) -> object: + return json.loads(self.content) + + +class Egress(Protocol): + def request( + self, + method: str, + url: str, + *, + headers: Mapping[str, str], + content: bytes | None, + timeout: float, + sensitive_url: bool, + ) -> EgressResponse: + """Send one request without retries or redirects. + + ``sensitive_url`` marks a URL that embeds a credential, so transport + diagnostics must not log it. + + Raises: + EgressError: The request failed in transport or was blocked by policy. + """ + ... + + def grpc_channel(self, endpoint: str) -> grpc.Channel: + """Open a gRPC channel to an ``http(s)://host[:port]`` endpoint; the caller closes it. + + Raises: + EgressError: The endpoint is not allowed. + """ + ... diff --git a/api/providers/trace/trace-sdk/src/dify_trace_sdk/emitter.py b/api/providers/trace/trace-sdk/src/dify_trace_sdk/emitter.py new file mode 100644 index 00000000000000..7892a166f8fa0d --- /dev/null +++ b/api/providers/trace/trace-sdk/src/dify_trace_sdk/emitter.py @@ -0,0 +1,50 @@ +"""The producer-side entry point of the trace pipeline.""" + +import logging +from dataclasses import dataclass +from datetime import UTC, datetime +from typing import Protocol, runtime_checkable +from uuid import uuid4 + +from dify_trace_sdk.events import TraceCorrelation, TraceEvent, TracePayload + +logger = logging.getLogger(__name__) + + +@runtime_checkable +class TraceEmitter(Protocol): + def emit(self, event: TraceEvent) -> None: + """Hand one event to the trace pipeline without blocking on delivery.""" + ... + + +@dataclass(frozen=True) +class TraceScope: + """Who emitted events belong to, plus the caller's correlation. + + A scope exists only when someone consumes the events, so producers check + for ``None`` instead of asking whether tracing is enabled. + """ + + emitter: TraceEmitter + tenant_id: str + app_id: str | None + user_id: str | None + correlation: TraceCorrelation + + def emit(self, payload: TracePayload) -> None: + """Report a finished operation. Call after its records are committed; never raises.""" + try: + self.emitter.emit( + TraceEvent( + event_id=uuid4().hex, + tenant_id=self.tenant_id, + app_id=self.app_id, + user_id=self.user_id, + occurred_at=datetime.now(UTC), + correlation=self.correlation, + payload=payload, + ) + ) + except Exception: + logger.warning("Trace event dropped: kind=%s app_id=%s", payload.kind, self.app_id, exc_info=True) diff --git a/api/providers/trace/trace-sdk/src/dify_trace_sdk/errors.py b/api/providers/trace/trace-sdk/src/dify_trace_sdk/errors.py new file mode 100644 index 00000000000000..e94ea813c8caf6 --- /dev/null +++ b/api/providers/trace/trace-sdk/src/dify_trace_sdk/errors.py @@ -0,0 +1,47 @@ +"""Failure classes a provider reports to the trace pipeline. + +Messages must never contain credentials, request bodies, or vendor response +bodies: they are logged and persisted as delivery diagnostics. +""" + + +class TraceProviderError(Exception): + """Base class for provider failures.""" + + +class ExportRetryableError(TraceProviderError): + """Nothing was accepted and the cause is transient, so sending again is safe.""" + + +class ExportRejectedError(TraceProviderError): + """The destination refused the trace, or the configuration cannot work. Do not retry.""" + + +class ExportOutcomeUnknownError(TraceProviderError): + """Part or all of the trace may have been accepted, so a retry could duplicate it.""" + + +class VerificationError(TraceProviderError): + """A configuration could not be verified. ``code`` is a stable, non-secret reason.""" + + def __init__(self, code: str) -> None: + self.code = code + super().__init__(code) + + +class EgressError(TraceProviderError): + """The request failed in transport or was blocked by the egress policy. + + ``request_sent`` tells whether the destination may have received the request. + """ + + def __init__(self, message: str, *, request_sent: bool) -> None: + self.request_sent = request_sent + super().__init__(message) + + +class EgressBlockedError(EgressError): + """The egress policy does not allow this destination. Retrying cannot help.""" + + def __init__(self, message: str) -> None: + super().__init__(message, request_sent=False) diff --git a/api/providers/trace/trace-sdk/src/dify_trace_sdk/events.py b/api/providers/trace/trace-sdk/src/dify_trace_sdk/events.py new file mode 100644 index 00000000000000..3f1b8dd3ac3714 --- /dev/null +++ b/api/providers/trace/trace-sdk/src/dify_trace_sdk/events.py @@ -0,0 +1,164 @@ +"""Trace events: what producers report when a traced operation finishes. + +An event carries Dify record identifiers plus the facts that are persisted +nowhere else (moderation verdicts, tool calls, generated names, timings). The +trace pipeline reloads everything else from the database after validating +the owner chain, so events stay small and never carry credentials. +""" + +import json +from datetime import datetime +from typing import Annotated, Literal + +from pydantic import BaseModel, ConfigDict, Field, JsonValue + +_EVENT_CONFIG = ConfigDict(extra="forbid", frozen=True) + + +def json_safe(value: object) -> JsonValue: + """Convert a producer value to strict JSON: other objects become strings, and so do NaN and Infinity.""" + return json.loads(json.dumps(value, ensure_ascii=False, default=str), parse_constant=str) + + +class TraceCorrelation(BaseModel): + """Caller-supplied correlation, propagated unchanged to the trace.""" + + external_trace_id: str | None + session_id: str | None + parent_workflow_run_id: str | None + parent_node_execution_id: str | None + + model_config = _EVENT_CONFIG + + +class MessageCompleted(BaseModel): + kind: Literal["message"] = "message" + message_id: str + conversation_id: str | None + + model_config = _EVENT_CONFIG + + +class WorkflowRunCompleted(BaseModel): + kind: Literal["workflow_run"] = "workflow_run" + workflow_run_id: str + conversation_id: str | None + total_tokens: int | None + """The in-memory total when the persisted run may lag behind.""" + + model_config = _EVENT_CONFIG + + +class ModerationChecked(BaseModel): + kind: Literal["moderation"] = "moderation" + message_id: str | None + inputs: dict[str, JsonValue] + query: str + flagged: bool + action: str + preset_response: str + started_at: datetime + ended_at: datetime + + model_config = _EVENT_CONFIG + + +class SuggestedQuestionsGenerated(BaseModel): + kind: Literal["suggested_questions"] = "suggested_questions" + message_id: str + questions: list[str] + started_at: datetime + ended_at: datetime + + model_config = _EVENT_CONFIG + + +class DatasetRetrieved(BaseModel): + kind: Literal["dataset_retrieval"] = "dataset_retrieval" + message_id: str | None + documents: list[dict[str, JsonValue]] + started_at: datetime + ended_at: datetime + + model_config = _EVENT_CONFIG + + +class ToolInvoked(BaseModel): + kind: Literal["tool"] = "tool" + message_id: str + tool_name: str + tool_inputs: dict[str, JsonValue] + tool_outputs: str + tool_config: dict[str, JsonValue] + error: str | None + time_cost: float + started_at: datetime + ended_at: datetime + + model_config = _EVENT_CONFIG + + +class ConversationNamed(BaseModel): + kind: Literal["conversation_name"] = "conversation_name" + conversation_id: str | None + message_id: str | None + inputs: str + name: str + started_at: datetime + ended_at: datetime + + model_config = _EVENT_CONFIG + + +class PromptGenerated(BaseModel): + kind: Literal["prompt_generation"] = "prompt_generation" + operation_type: str + instruction: str + generated_output: str + model_provider: str + model_name: str + prompt_tokens: int + completion_tokens: int + total_tokens: int + latency: float + total_price: float | None + currency: str | None + error: str | None + started_at: datetime + ended_at: datetime + + model_config = _EVENT_CONFIG + + +class DraftNodeExecuted(BaseModel): + kind: Literal["draft_node_execution"] = "draft_node_execution" + node_execution: dict[str, JsonValue] + + model_config = _EVENT_CONFIG + + +type TracePayload = Annotated[ + MessageCompleted + | WorkflowRunCompleted + | ModerationChecked + | SuggestedQuestionsGenerated + | DatasetRetrieved + | ToolInvoked + | ConversationNamed + | PromptGenerated + | DraftNodeExecuted, + Field(discriminator="kind"), +] + + +class TraceEvent(BaseModel): + event_id: str + tenant_id: str + app_id: str | None + """``None`` only for workspace-level operations such as prompt generation.""" + user_id: str | None + occurred_at: datetime + correlation: TraceCorrelation + payload: TracePayload + + model_config = _EVENT_CONFIG diff --git a/api/providers/trace/trace-sdk/src/dify_trace_sdk/exporting.py b/api/providers/trace/trace-sdk/src/dify_trace_sdk/exporting.py new file mode 100644 index 00000000000000..98547e76528c54 --- /dev/null +++ b/api/providers/trace/trace-sdk/src/dify_trace_sdk/exporting.py @@ -0,0 +1,81 @@ +"""Helpers for implementing ``TraceProvider.send``.""" + +import json +from collections.abc import Mapping +from datetime import UTC, datetime +from urllib.parse import urlsplit, urlunsplit + +from dify_trace_sdk.egress import Egress, EgressResponse +from dify_trace_sdk.errors import ( + EgressBlockedError, + EgressError, + ExportOutcomeUnknownError, + ExportRejectedError, + ExportRetryableError, + TraceProviderError, +) + +EXPORT_TIMEOUT_SECONDS = 30.0 + + +def json_value(value: object) -> str: + return json.dumps(value, ensure_ascii=False, default=str, allow_nan=False) + + +def nanos(value: datetime) -> int: + aware = value.replace(tzinfo=UTC) if value.tzinfo is None else value + return int(aware.timestamp() * 1_000_000_000) + + +def export_endpoint(value: str) -> str: + """Validate a configured base URL for export and drop its trailing slash.""" + parsed = urlsplit(value.strip()) + if parsed.scheme not in {"http", "https"} or not parsed.hostname or parsed.username or parsed.password: + raise ExportRejectedError("Invalid tracing endpoint") + if parsed.query or parsed.fragment: + raise ExportRejectedError("Tracing endpoint must not contain query parameters or a fragment") + return urlunsplit(parsed).rstrip("/") + + +def classify_egress_error(error: EgressError, destination: str) -> TraceProviderError: + """The pipeline's reading of a transport failure; raise it ``from None`` so no request detail is chained.""" + if isinstance(error, EgressBlockedError): + return ExportRejectedError(f"{destination} endpoint is not allowed by the egress policy") + if error.request_sent: + return ExportOutcomeUnknownError(f"{destination} ingestion outcome is unknown after transport failure") + return ExportRetryableError(f"{destination} is unreachable") + + +def raise_for_export_status(status_code: int, destination: str) -> None: + if 200 <= status_code < 300: + return + # 501: the destination does not implement this ingestion API, so sending again cannot help. + if status_code in {408, 425, 429} or (status_code >= 500 and status_code != 501): + raise ExportRetryableError(f"{destination} ingestion temporarily unavailable (HTTP {status_code})") + raise ExportRejectedError(f"{destination} ingestion rejected (HTTP {status_code})") + + +def send_export_request( + egress: Egress, + method: str, + url: str, + *, + destination: str, + headers: Mapping[str, str], + content: bytes | None, + sensitive_url: bool, +) -> EgressResponse: + """Send one export request and classify every failure for the pipeline.""" + try: + response = egress.request( + method, + url, + headers=headers, + content=content, + timeout=EXPORT_TIMEOUT_SECONDS, + sensitive_url=sensitive_url, + ) + except EgressError as error: + raise classify_egress_error(error, destination) from None + raise_for_export_status(response.status_code, destination) + return response diff --git a/api/providers/trace/trace-sdk/src/dify_trace_sdk/genai.py b/api/providers/trace/trace-sdk/src/dify_trace_sdk/genai.py new file mode 100644 index 00000000000000..6cdc1069662739 --- /dev/null +++ b/api/providers/trace/trace-sdk/src/dify_trace_sdk/genai.py @@ -0,0 +1,71 @@ +"""OpenTelemetry GenAI semantic-convention attributes shared by GenAI-dialect providers.""" + +from collections.abc import Mapping +from dataclasses import dataclass + +from dify_trace_sdk.exporting import json_value +from dify_trace_sdk.model import Span, SpanKind, Trace + +_DIFY_ID_ATTRIBUTES = ("tenant_id", "app_id", "agent_id", "workflow_run_id", "node_execution_id", "message_id") + + +@dataclass(frozen=True) +class GenAIDialect: + """The few places where GenAI-compatible destinations disagree.""" + + kind_labels: Mapping[SpanKind, str] + input_key: str + output_key: str + mark_entry_span: bool + + +def genai_attributes(span: Span, trace: Trace, dialect: GenAIDialect) -> dict[str, object]: + root = trace.root + result: dict[str, object] = { + "gen_ai.span.kind": dialect.kind_labels.get(span.kind, span.kind.value.upper()), + "gen_ai.session.id": trace.session_id, + "gen_ai.framework": "dify", + "dify.span.id": span.span_id, + "dify.metadata": json_value(span.attributes), + dialect.input_key: json_value(span.inputs), + dialect.output_key: json_value(span.outputs), + } + if dialect.mark_entry_span and span is root and trace.parent is None: + result["gen_ai.is_entry"] = "true" + for key in _DIFY_ID_ATTRIBUTES: + value = span.attributes.get(key) or root.attributes.get(key) + if value: + result[f"dify.{key}"] = str(value) + if trace.user_id: + result["gen_ai.user.id"] = trace.user_id + if span.kind is SpanKind.LLM: + result["gen_ai.operation.name"] = "chat" + result["gen_ai.prompt"] = json_value(span.inputs) + result["gen_ai.completion"] = json_value(span.outputs) + result["gen_ai.input.messages"] = json_value(span.inputs) + result["gen_ai.output.messages"] = json_value(span.outputs) + if span.model_name: + result["gen_ai.request.model"] = span.model_name + result["gen_ai.response.model"] = span.model_name + if span.model_provider: + result["gen_ai.provider.name"] = span.model_provider + if span.usage is not None: + result["gen_ai.usage.input_tokens"] = span.usage.prompt_tokens + result["gen_ai.usage.output_tokens"] = span.usage.completion_tokens + result["gen_ai.usage.total_tokens"] = span.usage.total_tokens + if span.time_to_first_token is not None: + result["gen_ai.response.time_to_first_token"] = int(span.time_to_first_token * 1_000_000_000) + elif span.kind is SpanKind.AGENT: + result["gen_ai.operation.name"] = "invoke_agent" + result["gen_ai.agent.name"] = span.name + elif span.kind is SpanKind.TOOL: + result["gen_ai.operation.name"] = "execute_tool" + result["gen_ai.tool.name"] = span.name + result["gen_ai.tool.call.arguments"] = json_value(span.inputs) + result["gen_ai.tool.call.result"] = json_value(span.outputs) + result["tool.name"] = span.name + result["tool.parameters"] = json_value(span.inputs) + elif span.kind is SpanKind.RETRIEVER: + result["retrieval.query"] = json_value(span.inputs) + result["retrieval.document"] = json_value(span.outputs) + return result diff --git a/api/providers/trace/trace-sdk/src/dify_trace_sdk/ids.py b/api/providers/trace/trace-sdk/src/dify_trace_sdk/ids.py new file mode 100644 index 00000000000000..6b77a24789adeb --- /dev/null +++ b/api/providers/trace/trace-sdk/src/dify_trace_sdk/ids.py @@ -0,0 +1,41 @@ +"""Deterministic identifiers. + +The same Dify record always maps to the same trace and span identifiers, so a +redelivered trace carries the identifiers of its first delivery and separately +sent units (a tool call, its message) join the same tree without coordination. +""" + +import hashlib +import re +from uuid import UUID, uuid5 + +_HEX32 = re.compile(r"[0-9a-f]{32}") + + +def trace_id_for(logical_id: str) -> str: + """Return the trace ID for a Dify record ID or caller-supplied trace ID. + + UUIDs and 32-hex values keep their value, so a message ID or a caller's + ``X-Trace-Id`` stays recognizable in the destination. Other values are hashed. + """ + candidate = logical_id.replace("-", "").lower() + if _HEX32.fullmatch(candidate) and int(candidate, 16): + return candidate + digest = hashlib.sha256(f"dify-trace:{logical_id}".encode()).hexdigest()[:32] + return digest if int(digest, 16) else "0" * 31 + "1" + + +def span_id_for(logical_id: str) -> str: + """Return the span ID for a Dify record ID (message, workflow run, node execution, ...).""" + digest = hashlib.sha256(f"dify-application-span:{logical_id}".encode()).hexdigest()[:16] + return digest if int(digest, 16) else "0" * 15 + "1" + + +def trace_uuid(trace_id: str) -> UUID: + """The trace ID as a UUID, for destinations whose identifiers are UUIDs.""" + return UUID(hex=trace_id) + + +def span_uuid(trace_id: str, span_id: str) -> UUID: + """A span ID as a UUID that is unique within its trace.""" + return uuid5(trace_uuid(trace_id), span_id) diff --git a/api/providers/trace/trace-sdk/src/dify_trace_sdk/model.py b/api/providers/trace/trace-sdk/src/dify_trace_sdk/model.py new file mode 100644 index 00000000000000..3203233f084595 --- /dev/null +++ b/api/providers/trace/trace-sdk/src/dify_trace_sdk/model.py @@ -0,0 +1,194 @@ +"""The provider-neutral trace tree. + +This model is the lowest common denominator of every supported destination: a +tree of completed spans with timing, input and output, status, and LLM usage. +Identifiers use the OpenTelemetry formats (32 and 16 lowercase hex characters) +so OTLP destinations accept them directly and other destinations can derive +their native identifiers deterministically (see ``dify_trace_sdk.ids``). +""" + +import math +import re +from datetime import UTC, datetime +from enum import StrEnum +from typing import Self + +from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator, model_validator + +_TRACE_ID = re.compile(r"[0-9a-f]{32}") +_SPAN_ID = re.compile(r"[0-9a-f]{16}") + + +def _aware_utc(value: datetime) -> datetime: + """Dify persists naive UTC datetimes; never interpret them in the host timezone.""" + return value.replace(tzinfo=UTC) if value.tzinfo is None else value.astimezone(UTC) + + +def _check_id(value: str, pattern: re.Pattern[str], label: str) -> str: + if not pattern.fullmatch(value) or int(value, 16) == 0: + raise ValueError(f"{label} must be nonzero lowercase hex") + return value + + +def _require_finite[T](value: T) -> T: + """Strict JSON has no NaN or Infinity, so a span holding one could never be sent.""" + pending: list[object] = [value] + while pending: + item = pending.pop() + if isinstance(item, float) and not math.isfinite(item): + raise ValueError("JSON numbers must be finite") + if isinstance(item, dict): + pending.extend(item.values()) + elif isinstance(item, list): + pending.extend(item) + return value + + +class SpanKind(StrEnum): + CHAIN = "chain" + LLM = "llm" + RETRIEVER = "retriever" + TOOL = "tool" + AGENT = "agent" + + +class SpanStatus(StrEnum): + OK = "ok" + ERROR = "error" + + +class TokenUsage(BaseModel): + prompt_tokens: int = Field(ge=0) + completion_tokens: int = Field(ge=0) + total_tokens: int = Field(ge=0) + + model_config = ConfigDict(extra="forbid", frozen=True) + + +class Cost(BaseModel): + amount: float = Field(ge=0) + currency: str = Field(min_length=1) + + model_config = ConfigDict(extra="forbid", frozen=True) + + @field_validator("amount") + @classmethod + def _finite(cls, value: float) -> float: + if not math.isfinite(value): + raise ValueError("cost amount must be finite") + return value + + +class Span(BaseModel): + """One completed operation. A span without a parent is the root of its trace.""" + + span_id: str + parent_span_id: str | None + name: str = Field(min_length=1) + kind: SpanKind + start_time: datetime + end_time: datetime + status: SpanStatus + error: str | None = None + inputs: JsonValue = None + outputs: JsonValue = None + model_name: str | None = None + model_provider: str | None = None + usage: TokenUsage | None = None + cost: Cost | None = None + time_to_first_token: float | None = Field(default=None, ge=0, allow_inf_nan=False) + attributes: dict[str, JsonValue] = Field(default_factory=dict) + + model_config = ConfigDict(extra="forbid", frozen=True) + + @field_validator("span_id") + @classmethod + def _span_id(cls, value: str) -> str: + return _check_id(value, _SPAN_ID, "span_id") + + @field_validator("inputs", "outputs", "attributes") + @classmethod + def _finite(cls, value: JsonValue) -> JsonValue: + return _require_finite(value) + + @field_validator("parent_span_id") + @classmethod + def _parent_span_id(cls, value: str | None) -> str | None: + return None if value is None else _check_id(value, _SPAN_ID, "parent_span_id") + + @field_validator("start_time", "end_time") + @classmethod + def _utc(cls, value: datetime) -> datetime: + return _aware_utc(value) + + @model_validator(mode="after") + def _ordered(self) -> Self: + if self.end_time < self.start_time: + raise ValueError(f"span {self.span_id} ends before it starts") + if self.parent_span_id == self.span_id: + raise ValueError(f"span {self.span_id} cannot be its own parent") + return self + + +class SpanRef(BaseModel): + """A span of the same trace that was or will be sent separately.""" + + span_id: str + start_time: datetime + + model_config = ConfigDict(extra="forbid", frozen=True) + + @field_validator("span_id") + @classmethod + def _span_id(cls, value: str) -> str: + return _check_id(value, _SPAN_ID, "span_id") + + @field_validator("start_time") + @classmethod + def _utc(cls, value: datetime) -> datetime: + return _aware_utc(value) + + +class Trace(BaseModel): + """A parent-first span tree, delivered as one unit. + + ``parent`` is set when this tree joins a trace whose root was sent in + another unit, e.g. a tool call under its message. ``parent`` is always the + root span of that trace, and the first span of this unit points at it. + """ + + trace_id: str + external_trace_id: str | None + """The caller-supplied trace ID, kept verbatim for destinations that accept free-form IDs.""" + session_id: str | None + user_id: str | None + parent: SpanRef | None + spans: tuple[Span, ...] = Field(min_length=1) + + model_config = ConfigDict(extra="forbid", frozen=True) + + @field_validator("trace_id") + @classmethod + def _trace_id(cls, value: str) -> str: + return _check_id(value, _TRACE_ID, "trace_id") + + @model_validator(mode="after") + def _parent_first(self) -> Self: + root = self.spans[0] + expected_root_parent = self.parent.span_id if self.parent is not None else None + if root.parent_span_id != expected_root_parent: + raise ValueError("the first span must be the root of this unit") + seen = {root.span_id} + for span in self.spans[1:]: + if span.span_id in seen: + raise ValueError(f"duplicate span_id: {span.span_id}") + if span.parent_span_id is None or span.parent_span_id not in seen: + raise ValueError(f"span {span.span_id} must follow its parent") + seen.add(span.span_id) + if expected_root_parent in seen: + raise ValueError("the external parent cannot be a span of this unit") + return self + + @property + def root(self) -> Span: + return self.spans[0] diff --git a/api/providers/trace/trace-sdk/src/dify_trace_sdk/otlp.py b/api/providers/trace/trace-sdk/src/dify_trace_sdk/otlp.py new file mode 100644 index 00000000000000..3b35fa83d10be0 --- /dev/null +++ b/api/providers/trace/trace-sdk/src/dify_trace_sdk/otlp.py @@ -0,0 +1,138 @@ +"""OTLP encoding and delivery for providers that ingest OpenTelemetry traces.""" + +from collections.abc import Callable, Mapping + +import grpc # type: ignore[import-untyped] +from google.protobuf.json_format import ParseDict, ParseError +from google.protobuf.message import DecodeError +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( + ExportTraceServiceRequest, + ExportTraceServiceResponse, +) +from opentelemetry.proto.collector.trace.v1.trace_service_pb2_grpc import TraceServiceStub +from opentelemetry.proto.common.v1.common_pb2 import AnyValue, KeyValue +from opentelemetry.proto.trace.v1.trace_pb2 import Span as OtlpSpan +from opentelemetry.proto.trace.v1.trace_pb2 import Status + +from dify_trace_sdk.egress import Egress, EgressResponse +from dify_trace_sdk.errors import EgressError, ExportOutcomeUnknownError, ExportRejectedError, ExportRetryableError +from dify_trace_sdk.exporting import classify_egress_error, json_value, nanos, send_export_request +from dify_trace_sdk.model import Span, SpanStatus, Trace + +type AttributeMapper = Callable[[Span], Mapping[str, object]] + +OTLP_GRPC_TIMEOUT_SECONDS = 30.0 + + +def _attribute(key: str, value: object) -> KeyValue: + if isinstance(value, bool): + encoded = AnyValue(bool_value=value) + elif isinstance(value, int): + encoded = AnyValue(int_value=value) + elif isinstance(value, float): + encoded = AnyValue(double_value=value) + elif isinstance(value, str): + encoded = AnyValue(string_value=value) + else: + encoded = AnyValue(string_value=json_value(value)) + return KeyValue(key=key, value=encoded) + + +def encode_trace( + trace: Trace, *, resource: Mapping[str, object], attributes: AttributeMapper +) -> ExportTraceServiceRequest: + """Encode one trace unit; ``None`` attribute values are omitted.""" + payload = ExportTraceServiceRequest() + resource_spans = payload.resource_spans.add() + resource_spans.resource.attributes.extend(_attribute(k, v) for k, v in resource.items() if v is not None) + scope = resource_spans.scope_spans.add() + scope.scope.name = "dify" + scope.scope.version = "1" + trace_id = bytes.fromhex(trace.trace_id) + for span in trace.spans: + target = scope.spans.add( + trace_id=trace_id, + span_id=bytes.fromhex(span.span_id), + parent_span_id=bytes.fromhex(span.parent_span_id) if span.parent_span_id else b"", + name=span.name, + kind=OtlpSpan.SPAN_KIND_INTERNAL, + start_time_unix_nano=nanos(span.start_time), + end_time_unix_nano=nanos(span.end_time), + flags=1, + ) + target.status.code = Status.STATUS_CODE_ERROR if span.status is SpanStatus.ERROR else Status.STATUS_CODE_OK + if span.error: + target.status.message = span.error + target.attributes.extend(_attribute(k, v) for k, v in attributes(span).items() if v is not None) + return payload + + +def read_otlp_ack(response: EgressResponse, destination: str) -> None: + """Accept a standard OTLP acknowledgment; partial acceptance is never retried.""" + result = ExportTraceServiceResponse() + if response.content: + try: + if "json" in response.headers.get("content-type", "").lower(): + body = response.json() + if not isinstance(body, dict): + raise ValueError("acknowledgment is not an object") + ParseDict(body, result) + else: + result.ParseFromString(response.content) + except (DecodeError, ParseError, ValueError): + raise ExportOutcomeUnknownError(f"{destination} returned an unreadable acknowledgment") from None + if result.partial_success.rejected_spans: + raise ExportOutcomeUnknownError(f"{destination} accepted part of the trace") + + +def post_otlp( + egress: Egress, + url: str, + request: ExportTraceServiceRequest, + *, + destination: str, + headers: Mapping[str, str], + sensitive_url: bool, +) -> EgressResponse: + """POST protobuf OTLP; the caller reads the acknowledgment.""" + return send_export_request( + egress, + "POST", + url, + destination=destination, + headers={"Content-Type": "application/x-protobuf", "Accept": "application/x-protobuf", **headers}, + content=request.SerializeToString(), + sensitive_url=sensitive_url, + ) + + +def export_otlp_grpc( + egress: Egress, + endpoint: str, + request: ExportTraceServiceRequest, + *, + destination: str, + metadata: Mapping[str, str], +) -> None: + try: + channel = egress.grpc_channel(endpoint) + except EgressError as error: + raise classify_egress_error(error, destination) from None + try: + result = TraceServiceStub(channel).Export( + request, metadata=tuple(metadata.items()), timeout=OTLP_GRPC_TIMEOUT_SECONDS + ) + except grpc.RpcError as error: + code = error.code() + if code in {grpc.StatusCode.UNAVAILABLE, grpc.StatusCode.RESOURCE_EXHAUSTED, grpc.StatusCode.ABORTED}: + raise ExportRetryableError(f"{destination} ingestion temporarily unavailable") from None + if code in {grpc.StatusCode.DEADLINE_EXCEEDED, grpc.StatusCode.CANCELLED, grpc.StatusCode.UNKNOWN}: + raise ExportOutcomeUnknownError(f"{destination} ingestion outcome is unknown") from None + raise ExportRejectedError(f"{destination} ingestion rejected") from None + except ValueError: + # gRPC rejects invalid metadata before sending and echoes it, credentials included, in the message. + raise ExportRejectedError(f"{destination} export metadata is invalid") from None + finally: + channel.close() + if result.partial_success.rejected_spans: + raise ExportOutcomeUnknownError(f"{destination} accepted part of the trace") diff --git a/api/providers/trace/trace-sdk/src/dify_trace_sdk/provider.py b/api/providers/trace/trace-sdk/src/dify_trace_sdk/provider.py new file mode 100644 index 00000000000000..54928f66c3ccd6 --- /dev/null +++ b/api/providers/trace/trace-sdk/src/dify_trace_sdk/provider.py @@ -0,0 +1,77 @@ +"""The interface every trace provider implements. + +It is the lowest common denominator of the supported destinations: validate a +configuration, verify it against the destination, link to the destination's +project, and send one trace. Everything else (collecting data, retries, +credentials storage, tenancy) belongs to the Dify trace pipeline. +""" + +from abc import ABC, abstractmethod +from collections.abc import Mapping +from dataclasses import dataclass +from typing import ClassVar + +from dify_trace_sdk.config import ProviderConfig +from dify_trace_sdk.egress import Egress +from dify_trace_sdk.model import Trace + + +@dataclass(frozen=True) +class ProviderCapabilities: + """What a destination's mapping supports. These describe the adapter, not delivery health.""" + + native_cost: bool + trace_url: bool + span_url: bool + + +@dataclass(frozen=True) +class SendResult: + trace_id: str + """The destination's identifier for the trace.""" + trace_url: str | None + span_urls: Mapping[str, str] + """Destination links keyed by ``Span.span_id``.""" + + +class TraceProvider[ConfigT: ProviderConfig](ABC): + """One trace destination. + + Instances are cheap and hold no connections between calls: the pipeline + creates one per operation with the egress it must use. + """ + + id: ClassVar[str] + config_model: ClassVar[type[ProviderConfig]] + capabilities: ClassVar[ProviderCapabilities] + + def __init__(self, egress: Egress) -> None: + self.egress = egress + + @abstractmethod + def verify(self, config: ConfigT) -> ConfigT: + """Check that the destination accepts this configuration. + + Return the configuration with derived fields resolved. Make read-only + requests only; never create resources in the destination. + + Raises: + VerificationError: The destination rejected or could not confirm the configuration. + """ + + @abstractmethod + def project_url(self, config: ConfigT) -> str | None: + """A link to the destination's project for this configuration. Performs no I/O.""" + + @abstractmethod + def send(self, config: ConfigT, trace: Trace) -> SendResult: + """Deliver one trace unit and return once the destination acknowledged it. + + When ``trace.parent`` is set, attach the root under that span if the + destination supports it; otherwise deliver the unit on its own. + + Raises: + ExportRetryableError: Nothing was accepted; sending again is safe. + ExportRejectedError: The destination refused the trace. + ExportOutcomeUnknownError: The destination may have accepted part of the trace. + """ diff --git a/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/data_exporter/__init__.py b/api/providers/trace/trace-sdk/src/dify_trace_sdk/py.typed similarity index 100% rename from api/providers/trace/trace-aliyun/src/dify_trace_aliyun/data_exporter/__init__.py rename to api/providers/trace/trace-sdk/src/dify_trace_sdk/py.typed diff --git a/api/providers/trace/trace-sdk/src/dify_trace_sdk/testing.py b/api/providers/trace/trace-sdk/src/dify_trace_sdk/testing.py new file mode 100644 index 00000000000000..846d38716062a4 --- /dev/null +++ b/api/providers/trace/trace-sdk/src/dify_trace_sdk/testing.py @@ -0,0 +1,204 @@ +"""Test support for provider packages: a recording egress and representative traces.""" + +import json +from collections import deque +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta + +import grpc # type: ignore[import-untyped] + +from dify_trace_sdk.egress import EgressResponse +from dify_trace_sdk.errors import EgressBlockedError +from dify_trace_sdk.ids import span_id_for, trace_id_for +from dify_trace_sdk.model import Cost, Span, SpanKind, SpanRef, SpanStatus, TokenUsage, Trace + +MESSAGE_ID = "7b3c6c0e-38a1-4b52-9a7e-3ad8f3a1c001" +WORKFLOW_RUN_ID = "7b3c6c0e-38a1-4b52-9a7e-3ad8f3a1c002" +STARTED_AT = datetime(2026, 10, 1, 8, 0, 0, tzinfo=UTC) + + +@dataclass(frozen=True) +class RecordedRequest: + method: str + url: str + headers: Mapping[str, str] + content: bytes | None + timeout: float + sensitive_url: bool + + def json(self) -> object: + return json.loads(self.content or b"null") + + +class RecordingEgress: + """Answers requests from a script and records what a provider sent.""" + + def __init__(self, responses: Iterable[EgressResponse | Exception], *, grpc_target: str | None) -> None: + self._responses = deque(responses) + self._grpc_target = grpc_target + self.requests: list[RecordedRequest] = [] + self.grpc_endpoints: list[str] = [] + + def request( + self, + method: str, + url: str, + *, + headers: Mapping[str, str], + content: bytes | None, + timeout: float, + sensitive_url: bool, + ) -> EgressResponse: + self.requests.append( + RecordedRequest( + method=method, + url=url, + headers=dict(headers), + content=content, + timeout=timeout, + sensitive_url=sensitive_url, + ) + ) + if not self._responses: + raise AssertionError(f"unexpected request: {method} {url}") + response = self._responses.popleft() + if isinstance(response, Exception): + raise response + return response + + def grpc_channel(self, endpoint: str) -> grpc.Channel: + self.grpc_endpoints.append(endpoint) + if self._grpc_target is None: + raise EgressBlockedError("gRPC is not available in this test") + return grpc.insecure_channel(self._grpc_target) + + +def response(status_code: int, body: object) -> EgressResponse: + content = body if isinstance(body, bytes) else json.dumps(body).encode() + content_type = "application/octet-stream" if isinstance(body, bytes) else "application/json" + return EgressResponse(status_code=status_code, headers={"content-type": content_type}, content=content) + + +def message_trace() -> Trace: + """A chat message: the message root and the LLM call that answered it.""" + message_span_id = span_id_for(MESSAGE_ID) + return Trace( + trace_id=trace_id_for(MESSAGE_ID), + external_trace_id=None, + session_id="conversation-1", + user_id="end-user-1", + parent=None, + spans=( + Span( + span_id=message_span_id, + parent_span_id=None, + name="message", + kind=SpanKind.CHAIN, + start_time=STARTED_AT, + end_time=STARTED_AT + timedelta(seconds=2), + status=SpanStatus.OK, + inputs={"query": "What is Dify?"}, + outputs="An LLM application platform.", + attributes={"message_id": MESSAGE_ID, "app_id": "app-1", "tenant_id": "tenant-1"}, + ), + Span( + span_id=span_id_for(f"{MESSAGE_ID}:llm"), + parent_span_id=message_span_id, + name="llm", + kind=SpanKind.LLM, + start_time=STARTED_AT, + end_time=STARTED_AT + timedelta(seconds=2), + status=SpanStatus.OK, + inputs=[{"role": "user", "content": "What is Dify?"}], + outputs="An LLM application platform.", + model_name="gpt-4o", + model_provider="openai", + usage=TokenUsage(prompt_tokens=12, completion_tokens=8, total_tokens=20), + cost=Cost(amount=0.0004, currency="USD"), + time_to_first_token=0.25, + attributes={"message_id": MESSAGE_ID}, + ), + ), + ) + + +def workflow_trace() -> Trace: + """A workflow run with a failed tool node under a successful LLM node.""" + run_span_id = span_id_for(WORKFLOW_RUN_ID) + llm_span_id = span_id_for("node-execution-llm") + return Trace( + trace_id=trace_id_for(WORKFLOW_RUN_ID), + external_trace_id=None, + session_id=WORKFLOW_RUN_ID, + user_id="account-1", + parent=None, + spans=( + Span( + span_id=run_span_id, + parent_span_id=None, + name="workflow", + kind=SpanKind.CHAIN, + start_time=STARTED_AT, + end_time=STARTED_AT + timedelta(seconds=5), + status=SpanStatus.ERROR, + error="Tool failed", + inputs={"topic": "tracing"}, + outputs={}, + attributes={"workflow_run_id": WORKFLOW_RUN_ID, "app_id": "app-1"}, + ), + Span( + span_id=llm_span_id, + parent_span_id=run_span_id, + name="llm_Summarize", + kind=SpanKind.LLM, + start_time=STARTED_AT + timedelta(seconds=1), + end_time=STARTED_AT + timedelta(seconds=3), + status=SpanStatus.OK, + inputs=[{"role": "user", "content": "Summarize tracing"}], + outputs={"text": "Tracing records execution."}, + model_name="gpt-4o", + model_provider="openai", + usage=TokenUsage(prompt_tokens=30, completion_tokens=10, total_tokens=40), + attributes={"node_execution_id": "node-execution-llm", "node_type": "llm"}, + ), + Span( + span_id=span_id_for("node-execution-tool"), + parent_span_id=llm_span_id, + name="tool_Search", + kind=SpanKind.TOOL, + start_time=STARTED_AT + timedelta(seconds=3), + end_time=STARTED_AT + timedelta(seconds=5), + status=SpanStatus.ERROR, + error="Tool failed", + inputs={"query": "tracing"}, + outputs={}, + attributes={"node_execution_id": "node-execution-tool", "node_type": "tool"}, + ), + ), + ) + + +def tool_fragment() -> Trace: + """A tool call sent separately and attached under the message of ``message_trace``.""" + return Trace( + trace_id=trace_id_for(MESSAGE_ID), + external_trace_id=None, + session_id="conversation-1", + user_id="end-user-1", + parent=SpanRef(span_id=span_id_for(MESSAGE_ID), start_time=STARTED_AT), + spans=( + Span( + span_id=span_id_for("tool-event-1"), + parent_span_id=span_id_for(MESSAGE_ID), + name="weather", + kind=SpanKind.TOOL, + start_time=STARTED_AT + timedelta(milliseconds=500), + end_time=STARTED_AT + timedelta(seconds=1), + status=SpanStatus.OK, + inputs={"city": "Berlin"}, + outputs="Sunny", + attributes={"message_id": MESSAGE_ID}, + ), + ), + ) diff --git a/api/providers/trace/trace-sdk/src/dify_trace_sdk/transports.py b/api/providers/trace/trace-sdk/src/dify_trace_sdk/transports.py new file mode 100644 index 00000000000000..7bdaafa56dcce3 --- /dev/null +++ b/api/providers/trace/trace-sdk/src/dify_trace_sdk/transports.py @@ -0,0 +1,67 @@ +"""Adapters that let a vendor SDK send its HTTP requests through ``Egress``. + +The SDK keeps its request and response validation; the egress keeps the +network policy. Both adapters forbid SDK-side retries by construction: each +SDK request is exactly one egress request. +""" + +from typing import Any, override + +import httpx +import requests + +from dify_trace_sdk.egress import Egress + +SDK_REQUEST_TIMEOUT_SECONDS = 30.0 + + +class EgressHttpxTransport(httpx.BaseTransport): + def __init__(self, egress: Egress) -> None: + self._egress = egress + + @override + def handle_request(self, request: httpx.Request) -> httpx.Response: + response = self._egress.request( + request.method, + str(request.url), + headers=dict(request.headers), + content=request.read(), + timeout=SDK_REQUEST_TIMEOUT_SECONDS, + sensitive_url=False, + ) + return httpx.Response( + response.status_code, headers=dict(response.headers), content=response.content, request=request + ) + + +class EgressRequestsSession(requests.Session): + def __init__(self, egress: Egress) -> None: + super().__init__() + self._egress = egress + # Host proxies, CA bundles, and .netrc credentials must not reach a tenant's destination. + self.trust_env = False + + @override + def send(self, request: requests.PreparedRequest, **kwargs: Any) -> requests.Response: + body = request.body + if not request.method or not request.url: + raise ValueError("The SDK prepared an incomplete request") + if isinstance(body, str): + body = body.encode() + if body is not None and not isinstance(body, bytes): + raise ValueError("The SDK prepared a streaming request body") + received = self._egress.request( + request.method, + request.url, + headers=dict(request.headers), + content=body, + timeout=SDK_REQUEST_TIMEOUT_SECONDS, + sensitive_url=False, + ) + response = requests.Response() + response.status_code = received.status_code + response.headers.update(received.headers) + response._content = received.content + response.url = request.url + response.request = request + return response diff --git a/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_config.py b/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_config.py new file mode 100644 index 00000000000000..0eeba1ac73dedd --- /dev/null +++ b/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_config.py @@ -0,0 +1,51 @@ +import pytest +from dify_trace_sdk.config import FormField, ProviderConfig, derived, http_url, integer_id, name_or_default, secret +from pydantic import ValidationError, field_validator + + +class _Config(ProviderConfig): + api_key: str = secret() + password: str | None = secret(default=None) + endpoint: str = "https://example.com" + project_url: str | None = derived() + + @field_validator("endpoint") + @classmethod + def _endpoint(cls, value: str) -> str: + return http_url(value, allowed_schemes=("https",)) + + +def test_field_metadata_drives_secrets_derived_fields_and_form() -> None: + assert _Config.secret_fields() == {"api_key", "password"} + assert _Config.derived_fields() == {"project_url"} + assert _Config.form_fields() == ( + FormField(name="api_key", secret=True, required=True), + FormField(name="password", secret=True, required=False), + FormField(name="endpoint", secret=False, required=False), + ) + + +def test_validation_errors_never_echo_input() -> None: + with pytest.raises(ValidationError) as caught: + _Config.model_validate({"api_key": "sk-123", "endpoint": "http://user:pw@example.com"}) + assert "user:pw" not in str(caught.value) + + +def test_unknown_stored_keys_are_ignored() -> None: + config = _Config.model_validate({"api_key": "k", "legacy_field": "x"}) + assert not hasattr(config, "legacy_field") + + +def test_url_and_name_helpers() -> None: + assert http_url(" https://example.com/base/ ", allowed_schemes=("https",)) == "https://example.com/base/" + with pytest.raises(ValueError, match="URL must start with"): + http_url("ftp://example.com", allowed_schemes=("https", "http")) + with pytest.raises(ValueError, match="must not contain credentials"): + http_url("https://user:secret@example.com", allowed_schemes=("https",)) + assert name_or_default(" ", "default") == "default" + assert name_or_default(" project ", "default") == "project" + assert integer_id(" 42 ") == "42" + with pytest.raises(ValueError, match="valid integer"): + integer_id("4x") + with pytest.raises(ValueError, match="valid integer"): + integer_id("\u0664\u0662") diff --git a/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_emitter.py b/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_emitter.py new file mode 100644 index 00000000000000..67b5b7b56afc2d --- /dev/null +++ b/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_emitter.py @@ -0,0 +1,40 @@ +import logging + +import pytest +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import MessageCompleted, TraceCorrelation, TraceEvent + +CORRELATION = TraceCorrelation( + external_trace_id="caller-1", session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None +) + + +class _CollectingEmitter: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +class _BrokenEmitter: + def emit(self, event: TraceEvent) -> None: + raise RuntimeError("broker down") + + +def test_scope_stamps_owner_and_correlation() -> None: + emitter = _CollectingEmitter() + scope = TraceScope(emitter=emitter, tenant_id="t", app_id="a", user_id="u", correlation=CORRELATION) + scope.emit(MessageCompleted(message_id="m", conversation_id="c")) + (event,) = emitter.events + assert (event.tenant_id, event.app_id, event.user_id) == ("t", "a", "u") + assert event.correlation == CORRELATION + assert event.payload == MessageCompleted(message_id="m", conversation_id="c") + assert TraceEvent.model_validate_json(event.model_dump_json()) == event + + +def test_scope_never_raises_into_the_producer(caplog: pytest.LogCaptureFixture) -> None: + scope = TraceScope(emitter=_BrokenEmitter(), tenant_id="t", app_id="a", user_id=None, correlation=CORRELATION) + with caplog.at_level(logging.WARNING): + scope.emit(MessageCompleted(message_id="m", conversation_id=None)) + assert "Trace event dropped" in caplog.text diff --git a/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_exporting.py b/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_exporting.py new file mode 100644 index 00000000000000..08fb18ca2d5b67 --- /dev/null +++ b/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_exporting.py @@ -0,0 +1,49 @@ +import pytest +from dify_trace_sdk.egress import EgressResponse +from dify_trace_sdk.errors import ( + EgressBlockedError, + EgressError, + ExportOutcomeUnknownError, + ExportRejectedError, + ExportRetryableError, +) +from dify_trace_sdk.exporting import EXPORT_TIMEOUT_SECONDS, export_endpoint, send_export_request +from dify_trace_sdk.testing import RecordingEgress, response + + +def _send(egress: RecordingEgress) -> None: + send_export_request( + egress, + "POST", + "https://collector.example.com/v1/traces", + destination="Example", + headers={}, + content=b"{}", + sensitive_url=False, + ) + + +@pytest.mark.parametrize( + ("outcome", "expected"), + [ + (EgressError("refused", request_sent=False), ExportRetryableError), + (EgressError("read timeout", request_sent=True), ExportOutcomeUnknownError), + (EgressBlockedError("private address"), ExportRejectedError), + (response(503, {}), ExportRetryableError), + (response(429, {}), ExportRetryableError), + (response(401, {}), ExportRejectedError), + (response(501, {}), ExportRejectedError), + ], +) +def test_export_failures_are_classified(outcome: EgressResponse | Exception, expected: type[Exception]) -> None: + egress = RecordingEgress([outcome], grpc_target=None) + with pytest.raises(expected): + _send(egress) + assert egress.requests[0].timeout == EXPORT_TIMEOUT_SECONDS + + +def test_export_endpoint_validation() -> None: + assert export_endpoint("https://collector.example.com/") == "https://collector.example.com" + for value in ("ftp://x", "https://u:p@x.example.com", "https://x.example.com/?token=1"): + with pytest.raises(ExportRejectedError): + export_endpoint(value) diff --git a/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_ids.py b/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_ids.py new file mode 100644 index 00000000000000..78b6f5454dbdcb --- /dev/null +++ b/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_ids.py @@ -0,0 +1,16 @@ +from dify_trace_sdk.ids import span_id_for, span_uuid, trace_id_for, trace_uuid +from dify_trace_sdk.testing import MESSAGE_ID + + +def test_trace_ids_keep_uuid_identity() -> None: + assert trace_id_for(MESSAGE_ID) == MESSAGE_ID.replace("-", "") + assert str(trace_uuid(trace_id_for(MESSAGE_ID))) == MESSAGE_ID + assert trace_id_for("caller-trace-1") == trace_id_for("caller-trace-1") + assert len(trace_id_for("caller-trace-1")) == 32 + + +def test_span_ids_are_stable_and_distinct() -> None: + assert span_id_for("node-1") == span_id_for("node-1") + assert span_id_for("node-1") != span_id_for("node-2") + trace_id = trace_id_for(MESSAGE_ID) + assert span_uuid(trace_id, span_id_for("node-1")) == span_uuid(trace_id, span_id_for("node-1")) diff --git a/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_model.py b/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_model.py new file mode 100644 index 00000000000000..a66539abd6a8d6 --- /dev/null +++ b/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_model.py @@ -0,0 +1,115 @@ +from datetime import UTC, datetime, timedelta + +import pytest +from dify_trace_sdk.ids import span_id_for, trace_id_for +from dify_trace_sdk.model import Cost, Span, SpanKind, SpanRef, SpanStatus, Trace +from pydantic import ValidationError + +STARTED_AT = datetime(2026, 1, 1, tzinfo=UTC) + + +def _span(span_id: str, parent_span_id: str | None) -> Span: + return Span( + span_id=span_id, + parent_span_id=parent_span_id, + name="step", + kind=SpanKind.CHAIN, + start_time=STARTED_AT, + end_time=STARTED_AT + timedelta(seconds=1), + status=SpanStatus.OK, + ) + + +def test_naive_datetimes_are_read_as_utc() -> None: + span = Span( + span_id=span_id_for("a"), + parent_span_id=None, + name="step", + kind=SpanKind.CHAIN, + start_time=datetime(2026, 1, 1, 8, 0), + end_time=datetime(2026, 1, 1, 8, 1), + status=SpanStatus.OK, + ) + assert span.start_time == datetime(2026, 1, 1, 8, 0, tzinfo=UTC) + + +def test_span_rejects_end_before_start_and_invalid_ids() -> None: + with pytest.raises(ValidationError): + Span( + span_id=span_id_for("a"), + parent_span_id=None, + name="step", + kind=SpanKind.CHAIN, + start_time=STARTED_AT, + end_time=STARTED_AT - timedelta(seconds=1), + status=SpanStatus.OK, + ) + with pytest.raises(ValidationError): + _span("0" * 16, None) + with pytest.raises(ValidationError): + _span("ABCDEF0123456789", None) + + +def test_trace_requires_parent_first_order() -> None: + root, child = span_id_for("root"), span_id_for("child") + with pytest.raises(ValidationError, match="must follow its parent"): + Trace( + trace_id=trace_id_for("t"), + external_trace_id=None, + session_id=None, + user_id=None, + parent=None, + spans=(_span(root, None), _span(span_id_for("orphan"), child)), + ) + with pytest.raises(ValidationError, match="root of this unit"): + Trace( + trace_id=trace_id_for("t"), + external_trace_id=None, + session_id=None, + user_id=None, + parent=None, + spans=(_span(child, root),), + ) + + +def test_fragment_root_must_point_at_its_external_parent() -> None: + parent = SpanRef(span_id=span_id_for("message"), start_time=STARTED_AT) + with pytest.raises(ValidationError, match="root of this unit"): + Trace( + trace_id=trace_id_for("t"), + external_trace_id=None, + session_id=None, + user_id=None, + parent=parent, + spans=(_span(span_id_for("tool"), None),), + ) + + +def test_cost_must_be_finite() -> None: + with pytest.raises(ValidationError): + Cost(amount=float("inf"), currency="USD") + + +@pytest.mark.parametrize( + "fields", + [ + {"inputs": {"score": float("nan")}}, + {"outputs": [1.0, float("inf")]}, + {"attributes": {"ratio": float("-inf")}}, + {"time_to_first_token": float("inf")}, + ], +) +def test_spans_reject_numbers_strict_json_cannot_carry(fields: dict[str, object]) -> None: + with pytest.raises(ValidationError): + Span.model_validate( + { + "span_id": span_id_for("a"), + "parent_span_id": None, + "name": "step", + "kind": SpanKind.CHAIN, + "start_time": STARTED_AT, + "end_time": STARTED_AT, + "status": SpanStatus.OK, + **fields, + } + ) diff --git a/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_otlp.py b/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_otlp.py new file mode 100644 index 00000000000000..c72e06a793e7d8 --- /dev/null +++ b/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_otlp.py @@ -0,0 +1,111 @@ +from concurrent import futures +from typing import override + +import grpc # type: ignore[import-untyped] +import pytest +from dify_trace_sdk.errors import ( + EgressError, + ExportOutcomeUnknownError, + ExportRejectedError, + ExportRetryableError, +) +from dify_trace_sdk.otlp import encode_trace, export_otlp_grpc, post_otlp, read_otlp_ack +from dify_trace_sdk.testing import RecordingEgress, message_trace, response, tool_fragment +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceResponse +from opentelemetry.proto.collector.trace.v1.trace_service_pb2_grpc import add_TraceServiceServicer_to_server + + +def test_otlp_encoding_keeps_ids_and_external_parent() -> None: + fragment = tool_fragment() + request = encode_trace(fragment, resource={"service.name": "dify"}, attributes=lambda span: {"a": 1, "b": None}) + (span,) = request.resource_spans[0].scope_spans[0].spans + assert span.trace_id.hex() == fragment.trace_id + assert span.span_id.hex() == fragment.root.span_id + assert fragment.parent is not None + assert span.parent_span_id.hex() == fragment.parent.span_id + assert [attribute.key for attribute in span.attributes] == ["a"] + + +def test_otlp_ack_partial_success_is_outcome_unknown() -> None: + partial = ExportTraceServiceResponse() + partial.partial_success.rejected_spans = 1 + with pytest.raises(ExportOutcomeUnknownError): + read_otlp_ack(response(200, partial.SerializeToString()), "Example") + with pytest.raises(ExportOutcomeUnknownError): + read_otlp_ack(response(200, b"not-protobuf-\xff"), "Example") + read_otlp_ack(response(200, ExportTraceServiceResponse().SerializeToString()), "Example") + + +def test_post_otlp_sends_protobuf() -> None: + egress = RecordingEgress([response(200, b"")], grpc_target=None) + request = encode_trace(message_trace(), resource={}, attributes=lambda span: {}) + post_otlp(egress, "https://c.example.com/v1/traces", request, destination="Example", headers={}, sensitive_url=True) + (sent,) = egress.requests + assert sent.headers["Content-Type"] == "application/x-protobuf" + assert sent.content == request.SerializeToString() + assert sent.sensitive_url is True + + +class _Collector: + """An in-process OTLP collector; gRPC dispatches to ``Export`` by name.""" + + def __init__(self) -> None: + self.received = [] + + def Export(self, request, context): # noqa: N802 - gRPC method name + self.received.append((request, dict(context.invocation_metadata()))) + return ExportTraceServiceResponse() + + +def test_otlp_grpc_export_uses_the_egress_channel() -> None: + collector = _Collector() + server = grpc.server(futures.ThreadPoolExecutor(max_workers=1)) + add_TraceServiceServicer_to_server(collector, server) + port = server.add_insecure_port("127.0.0.1:0") + server.start() + try: + egress = RecordingEgress([], grpc_target=f"127.0.0.1:{port}") + request = encode_trace(message_trace(), resource={}, attributes=lambda span: {}) + export_otlp_grpc(egress, "https://collector.example.com", request, destination="Example", metadata={"k": "v"}) + finally: + server.stop(None) + assert egress.grpc_endpoints == ["https://collector.example.com"] + ((received, metadata),) = collector.received + assert received == request + assert metadata["k"] == "v" + + +def test_otlp_grpc_invalid_metadata_is_rejected_without_echoing_it() -> None: + egress = RecordingEgress([], grpc_target="127.0.0.1:9") + request = encode_trace(message_trace(), resource={}, attributes=lambda span: {}) + with pytest.raises(ExportRejectedError) as raised: + export_otlp_grpc( + egress, "https://collector.example.com", request, destination="Example", metadata={"api_key": "s3cret\n"} + ) + assert "s3cret" not in str(raised.value) + assert raised.value.__cause__ is None + + +class _UnresolvableEgress(RecordingEgress): + @override + def grpc_channel(self, endpoint: str) -> grpc.Channel: + raise EgressError("name resolution failed", request_sent=False) + + +def test_otlp_grpc_unresolvable_endpoint_is_retryable() -> None: + request = encode_trace(message_trace(), resource={}, attributes=lambda span: {}) + with pytest.raises(ExportRetryableError): + export_otlp_grpc( + _UnresolvableEgress([], grpc_target=None), + "https://collector.example.com", + request, + destination="Example", + metadata={}, + ) + + +def test_otlp_grpc_blocked_endpoint_is_rejected() -> None: + egress = RecordingEgress([], grpc_target=None) + request = encode_trace(message_trace(), resource={}, attributes=lambda span: {}) + with pytest.raises(ExportRejectedError): + export_otlp_grpc(egress, "https://169.254.169.254", request, destination="Example", metadata={}) diff --git a/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_testing.py b/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_testing.py new file mode 100644 index 00000000000000..7a97ac9c91c363 --- /dev/null +++ b/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_testing.py @@ -0,0 +1,9 @@ +from dify_trace_sdk.testing import message_trace, tool_fragment, workflow_trace + + +def test_sample_traces_are_valid_parent_first_trees() -> None: + assert message_trace().root.name == "message" + assert len(workflow_trace().spans) == 3 + fragment = tool_fragment() + assert fragment.parent is not None + assert fragment.root.parent_span_id == fragment.parent.span_id diff --git a/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_transports.py b/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_transports.py new file mode 100644 index 00000000000000..b5204fdea8231f --- /dev/null +++ b/api/providers/trace/trace-sdk/tests/unit_tests/src/dify_trace_sdk/test_transports.py @@ -0,0 +1,23 @@ +import httpx +from dify_trace_sdk.testing import RecordingEgress, response +from dify_trace_sdk.transports import EgressHttpxTransport, EgressRequestsSession + + +def test_httpx_transport_routes_sdk_requests_through_egress() -> None: + egress = RecordingEgress([response(201, {"ok": True})], grpc_target=None) + with httpx.Client(transport=EgressHttpxTransport(egress)) as client: + result = client.post("https://api.example.com/v1/items", json={"a": 1}) + assert result.status_code == 201 + assert result.json() == {"ok": True} + (sent,) = egress.requests + assert sent.method == "POST" + assert sent.json() == {"a": 1} + + +def test_requests_session_routes_sdk_requests_through_egress() -> None: + egress = RecordingEgress([response(200, {"ok": True})], grpc_target=None) + session = EgressRequestsSession(egress) + result = session.post("https://api.example.com/v1/items", json={"a": 1}) + assert result.json() == {"ok": True} + assert egress.requests[0].url == "https://api.example.com/v1/items" + assert session.trust_env is False diff --git a/api/providers/trace/trace-tencent/pyproject.toml b/api/providers/trace/trace-tencent/pyproject.toml index eab06fc70805a9..f5e2fb6c1d4a39 100644 --- a/api/providers/trace/trace-tencent/pyproject.toml +++ b/api/providers/trace/trace-tencent/pyproject.toml @@ -2,13 +2,12 @@ name = "dify-trace-tencent" version = "0.0.1" dependencies = [ - # versions inherited from parent - "opentelemetry-api", - "opentelemetry-exporter-otlp-proto-grpc", - "opentelemetry-sdk", - "opentelemetry-semantic-conventions", + "dify-trace-sdk", ] description = "Dify ops tracing provider (Tencent APM)." +[project.entry-points."dify.trace_providers"] +tencent = "dify_trace_tencent.provider:TencentProvider" + [tool.setuptools.packages.find] where = ["src"] diff --git a/api/providers/trace/trace-tencent/src/dify_trace_tencent/agent_trace.py b/api/providers/trace/trace-tencent/src/dify_trace_tencent/agent_trace.py deleted file mode 100644 index 78c99f94e61152..00000000000000 --- a/api/providers/trace/trace-tencent/src/dify_trace_tencent/agent_trace.py +++ /dev/null @@ -1,19 +0,0 @@ -"""Tencent APM Agent spans over its existing authenticated OTLP/gRPC protocol.""" - -from collections.abc import Mapping -from typing import Any - -from core.ops.unified_trace.entities import CanonicalTrace -from services.tracing.genai_attributes import attributes -from services.tracing.otlp_export import build_request, export_grpc -from services.tracing.provider_common import endpoint, required -from services.tracing.provider_dispatch import TraceExportResult - - -def export_trace(provider: str, config: Mapping[str, Any], trace: CanonicalTrace) -> TraceExportResult: - url = endpoint(required(config, "endpoint")) - payload = build_request( - trace, {"service.name": required(config, "service_name")}, lambda span: attributes(span, trace, provider) - ) - export_grpc(provider, url, {"authorization": f"Bearer {required(config, 'token')}"}, payload) - return TraceExportResult(trace_id=trace.trace_id) diff --git a/api/providers/trace/trace-tencent/src/dify_trace_tencent/client.py b/api/providers/trace/trace-tencent/src/dify_trace_tencent/client.py deleted file mode 100644 index a628ce4bb51fc8..00000000000000 --- a/api/providers/trace/trace-tencent/src/dify_trace_tencent/client.py +++ /dev/null @@ -1,569 +0,0 @@ -""" -Tencent APM Trace Client - handles network operations, metrics, and API communication -""" - -from __future__ import annotations - -import importlib -import json -import logging -import os -import socket -from typing import TYPE_CHECKING -from urllib.parse import urlparse - -try: - from importlib.metadata import version -except ImportError: - from importlib_metadata import version # type: ignore[import-not-found] - -if TYPE_CHECKING: - from opentelemetry.metrics import Histogram, Meter - from opentelemetry.sdk.metrics.export import MetricReader - -from opentelemetry import trace as trace_api -from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter -from opentelemetry.sdk.resources import Resource -from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.sdk.trace.export import BatchSpanProcessor -from opentelemetry.semconv._incubating.attributes.deployment_attributes import ( # type: ignore[import-untyped] - DEPLOYMENT_ENVIRONMENT, -) -from opentelemetry.semconv._incubating.attributes.host_attributes import ( # type: ignore[import-untyped] - HOST_NAME, -) -from opentelemetry.semconv.attributes import service_attributes -from opentelemetry.trace import SpanKind -from opentelemetry.util.types import AttributeValue - -from configs import dify_config - -from .entities.semconv import ( - GEN_AI_SERVER_TIME_TO_FIRST_TOKEN, - GEN_AI_STREAMING_TIME_TO_GENERATE, - GEN_AI_TOKEN_USAGE, - GEN_AI_TRACE_DURATION, - LLM_OPERATION_DURATION, -) -from .entities.tencent_trace_entity import SpanData - -logger = logging.getLogger(__name__) - - -def _get_opentelemetry_sdk_version() -> str: - """Get OpenTelemetry SDK version dynamically.""" - try: - return version("opentelemetry-sdk") - except Exception: - logger.debug("Failed to get opentelemetry-sdk version, using default") - return "1.27.0" # fallback version - - -class TencentTraceClient: - """Tencent APM trace client using OpenTelemetry OTLP exporter""" - - def __init__( - self, - service_name: str, - endpoint: str, - token: str, - max_queue_size: int = 1000, - schedule_delay_sec: int = 5, - max_export_batch_size: int = 50, - metrics_export_interval_sec: int = 10, - ): - self.endpoint = endpoint - self.token = token - self.service_name = service_name - self.metrics_export_interval_sec = metrics_export_interval_sec - - self.resource = Resource( - attributes={ - service_attributes.SERVICE_NAME: service_name, - service_attributes.SERVICE_VERSION: f"dify-{dify_config.project.version}-{dify_config.COMMIT_SHA}", - DEPLOYMENT_ENVIRONMENT: f"{dify_config.DEPLOY_ENV}-{dify_config.DEPLOYMENT_EDITION.value}", - HOST_NAME: socket.gethostname(), - "telemetry.sdk.language": "python", - "telemetry.sdk.name": "opentelemetry", - "telemetry.sdk.version": _get_opentelemetry_sdk_version(), - } - ) - # Prepare gRPC endpoint/metadata - grpc_endpoint, insecure, _, _ = self._resolve_grpc_target(endpoint) - - headers = (("authorization", f"Bearer {token}"),) - - self.exporter = OTLPSpanExporter( - endpoint=grpc_endpoint, - headers=headers, - insecure=insecure, - timeout=30, - ) - - self.tracer_provider = TracerProvider(resource=self.resource) - self.span_processor = BatchSpanProcessor( - span_exporter=self.exporter, - max_queue_size=max_queue_size, - schedule_delay_millis=schedule_delay_sec * 1000, - max_export_batch_size=max_export_batch_size, - ) - self.tracer_provider.add_span_processor(self.span_processor) - - # use dify api version as tracer version - self.tracer = self.tracer_provider.get_tracer("dify-sdk", dify_config.project.version) - - # Store span contexts for parent-child relationships - self.span_contexts: dict[int, trace_api.SpanContext] = {} - - self.meter: Meter | None = None - self.meter_provider: MeterProvider | None = None - self.hist_llm_duration: Histogram | None = None - self.hist_token_usage: Histogram | None = None - self.hist_time_to_first_token: Histogram | None = None - self.hist_time_to_generate: Histogram | None = None - self.hist_trace_duration: Histogram | None = None - self.metric_reader: MetricReader | None = None - - # Metrics exporter and instruments - try: - from opentelemetry.sdk.metrics import Histogram as SdkHistogram - from opentelemetry.sdk.metrics import MeterProvider - from opentelemetry.sdk.metrics.export import AggregationTemporality, PeriodicExportingMetricReader - - protocol = os.getenv("OTEL_EXPORTER_OTLP_PROTOCOL", "").strip().lower() - use_http_protobuf = protocol in {"http/protobuf", "http-protobuf"} - use_http_json = protocol in {"http/json", "http-json"} - - # Tencent APM works best with delta aggregation temporality - preferred_temporality: dict[type, AggregationTemporality] = {SdkHistogram: AggregationTemporality.DELTA} - - def _create_metric_exporter(exporter_cls, **kwargs): - """Create metric exporter with preferred_temporality support""" - try: - return exporter_cls(**kwargs, preferred_temporality=preferred_temporality) - except Exception: - return exporter_cls(**kwargs) - - metric_reader = None - if use_http_json: - exporter_cls = None - for mod_path in ( - "opentelemetry.exporter.otlp.http.json.metric_exporter", - "opentelemetry.exporter.otlp.json.metric_exporter", - ): - try: - mod = importlib.import_module(mod_path) - exporter_cls = getattr(mod, "OTLPMetricExporter", None) - if exporter_cls: - break - except Exception: - continue - if exporter_cls is not None: - metric_exporter = _create_metric_exporter( - exporter_cls, - endpoint=endpoint, - headers={"authorization": f"Bearer {token}"}, - ) - else: - from opentelemetry.exporter.otlp.proto.http.metric_exporter import ( - OTLPMetricExporter as HttpMetricExporter, - ) - - metric_exporter = _create_metric_exporter( - HttpMetricExporter, - endpoint=endpoint, - headers={"authorization": f"Bearer {token}"}, - ) - metric_reader = PeriodicExportingMetricReader( - metric_exporter, export_interval_millis=self.metrics_export_interval_sec * 1000 - ) - - elif use_http_protobuf: - from opentelemetry.exporter.otlp.proto.http.metric_exporter import ( - OTLPMetricExporter as HttpMetricExporter, - ) - - metric_exporter = _create_metric_exporter( - HttpMetricExporter, - endpoint=endpoint, - headers={"authorization": f"Bearer {token}"}, - ) - metric_reader = PeriodicExportingMetricReader( - metric_exporter, export_interval_millis=self.metrics_export_interval_sec * 1000 - ) - else: - from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import ( - OTLPMetricExporter as GrpcMetricExporter, - ) - - m_grpc_endpoint, m_insecure, _, _ = self._resolve_grpc_target(endpoint) - - metric_exporter = _create_metric_exporter( - GrpcMetricExporter, - endpoint=m_grpc_endpoint, - headers={"authorization": f"Bearer {token}"}, - insecure=m_insecure, - ) - metric_reader = PeriodicExportingMetricReader( - metric_exporter, export_interval_millis=self.metrics_export_interval_sec * 1000 - ) - - if metric_reader is not None: - # Use instance-level MeterProvider instead of global to support config changes - # without worker restart. Each TencentTraceClient manages its own MeterProvider. - provider = MeterProvider(resource=self.resource, metric_readers=[metric_reader]) - self.meter_provider = provider - self.meter = provider.get_meter("dify-sdk", dify_config.project.version) - - # LLM operation duration histogram - self.hist_llm_duration = self.meter.create_histogram( - name=LLM_OPERATION_DURATION, - unit="s", - description="LLM operation duration (seconds)", - ) - - # Token usage histogram with exponential buckets - self.hist_token_usage = self.meter.create_histogram( - name=GEN_AI_TOKEN_USAGE, - unit="token", - description="Number of tokens used in prompt and completions", - ) - - # Time to first token histogram - self.hist_time_to_first_token = self.meter.create_histogram( - name=GEN_AI_SERVER_TIME_TO_FIRST_TOKEN, - unit="s", - description="Time to first token for streaming LLM responses (seconds)", - ) - - # Time to generate histogram - self.hist_time_to_generate = self.meter.create_histogram( - name=GEN_AI_STREAMING_TIME_TO_GENERATE, - unit="s", - description="Total time to generate streaming LLM responses (seconds)", - ) - - # Trace duration histogram - self.hist_trace_duration = self.meter.create_histogram( - name=GEN_AI_TRACE_DURATION, - unit="s", - description="End-to-end GenAI trace duration (seconds)", - ) - - self.metric_reader = metric_reader - else: - self.meter = None - self.meter_provider = None - self.hist_llm_duration = None - self.hist_token_usage = None - self.hist_time_to_first_token = None - self.hist_time_to_generate = None - self.hist_trace_duration = None - self.metric_reader = None - except Exception: - logger.exception("[Tencent APM] Metrics initialization failed; metrics disabled") - self.meter = None - self.meter_provider = None - self.hist_llm_duration = None - self.hist_token_usage = None - self.hist_time_to_first_token = None - self.hist_time_to_generate = None - self.hist_trace_duration = None - self.metric_reader = None - - def add_span(self, span_data: SpanData) -> None: - """Create and export span using OpenTelemetry Tracer API""" - try: - self._create_and_export_span(span_data) - logger.debug("[Tencent APM] Created span: %s", span_data.name) - - except Exception: - logger.exception("[Tencent APM] Failed to create span: %s", span_data.name) - - # Metrics recording API - def record_llm_duration(self, latency_seconds: float, attributes: dict[str, str] | None = None) -> None: - """Record LLM operation duration histogram in seconds.""" - try: - if not hasattr(self, "hist_llm_duration") or self.hist_llm_duration is None: - return - attrs: dict[str, str] = {} - if attributes: - for k, v in attributes.items(): - attrs[k] = str(v) if not isinstance(v, (str, int, float, bool)) else v # type: ignore[assignment] - - logger.info( - "[Tencent Metrics] Metric: %s | Value: %.4f | Attributes: %s", - LLM_OPERATION_DURATION, - latency_seconds, - json.dumps(attrs, ensure_ascii=False), - ) - - self.hist_llm_duration.record(latency_seconds, attrs) # type: ignore[attr-defined] - except Exception: - logger.debug("[Tencent APM] Failed to record LLM duration", exc_info=True) - - def record_token_usage( - self, - token_count: int, - token_type: str, - operation_name: str, - request_model: str, - response_model: str, - server_address: str, - provider: str, - ) -> None: - """Record token usage histogram. - - Args: - token_count: Number of tokens used - token_type: "input" or "output" - operation_name: Operation name (e.g., "chat") - request_model: Model used in request - response_model: Model used in response - server_address: Server address - provider: Model provider name - """ - try: - if not hasattr(self, "hist_token_usage") or self.hist_token_usage is None: - return - - attributes = { - "gen_ai.operation.name": operation_name, - "gen_ai.request.model": request_model, - "gen_ai.response.model": response_model, - "gen_ai.system": provider, - "gen_ai.token.type": token_type, - "server.address": server_address, - } - - logger.info( - "[Tencent Metrics] Metric: %s | Value: %d | Attributes: %s", - GEN_AI_TOKEN_USAGE, - token_count, - json.dumps(attributes, ensure_ascii=False), - ) - - self.hist_token_usage.record(token_count, attributes) # type: ignore[attr-defined] - except Exception: - logger.debug("[Tencent APM] Failed to record token usage", exc_info=True) - - def record_time_to_first_token( - self, ttft_seconds: float, provider: str, model: str, operation_name: str = "chat" - ) -> None: - """Record time to first token histogram. - - Args: - ttft_seconds: Time to first token in seconds - provider: Model provider name - model: Model name - operation_name: Operation name (default: "chat") - """ - try: - if not hasattr(self, "hist_time_to_first_token") or self.hist_time_to_first_token is None: - return - - attributes = { - "gen_ai.operation.name": operation_name, - "gen_ai.system": provider, - "gen_ai.request.model": model, - "gen_ai.response.model": model, - "stream": "true", - } - - logger.info( - "[Tencent Metrics] Metric: %s | Value: %.4f | Attributes: %s", - GEN_AI_SERVER_TIME_TO_FIRST_TOKEN, - ttft_seconds, - json.dumps(attributes, ensure_ascii=False), - ) - - self.hist_time_to_first_token.record(ttft_seconds, attributes) # type: ignore[attr-defined] - except Exception: - logger.debug("[Tencent APM] Failed to record time to first token", exc_info=True) - - def record_time_to_generate( - self, ttg_seconds: float, provider: str, model: str, operation_name: str = "chat" - ) -> None: - """Record time to generate histogram. - - Args: - ttg_seconds: Time to generate in seconds - provider: Model provider name - model: Model name - operation_name: Operation name (default: "chat") - """ - try: - if not hasattr(self, "hist_time_to_generate") or self.hist_time_to_generate is None: - return - - attributes = { - "gen_ai.operation.name": operation_name, - "gen_ai.system": provider, - "gen_ai.request.model": model, - "gen_ai.response.model": model, - "stream": "true", - } - - logger.info( - "[Tencent Metrics] Metric: %s | Value: %.4f | Attributes: %s", - GEN_AI_STREAMING_TIME_TO_GENERATE, - ttg_seconds, - json.dumps(attributes, ensure_ascii=False), - ) - - self.hist_time_to_generate.record(ttg_seconds, attributes) # type: ignore[attr-defined] - except Exception: - logger.debug("[Tencent APM] Failed to record time to generate", exc_info=True) - - def record_trace_duration(self, duration_seconds: float, attributes: dict[str, str] | None = None) -> None: - """Record end-to-end trace duration histogram in seconds. - - Args: - duration_seconds: Trace duration in seconds - attributes: Optional attributes (e.g., conversation_mode, app_id) - """ - try: - if not hasattr(self, "hist_trace_duration") or self.hist_trace_duration is None: - return - - attrs: dict[str, str] = {} - if attributes: - for k, v in attributes.items(): - attrs[k] = str(v) if not isinstance(v, (str, int, float, bool)) else v # type: ignore[assignment] - - logger.info( - "[Tencent Metrics] Metric: %s | Value: %.4f | Attributes: %s", - GEN_AI_TRACE_DURATION, - duration_seconds, - json.dumps(attrs, ensure_ascii=False), - ) - - self.hist_trace_duration.record(duration_seconds, attrs) # type: ignore[attr-defined] - except Exception: - logger.debug("[Tencent APM] Failed to record trace duration", exc_info=True) - - def _create_and_export_span(self, span_data: SpanData) -> None: - """Create span using OpenTelemetry Tracer API""" - try: - parent_context = None - if span_data.parent_span_id and span_data.parent_span_id in self.span_contexts: - parent_context = trace_api.set_span_in_context( - trace_api.NonRecordingSpan(self.span_contexts[span_data.parent_span_id]) - ) - - span = self.tracer.start_span( - name=span_data.name, - context=parent_context, - kind=SpanKind.INTERNAL, - start_time=span_data.start_time, - ) - self.span_contexts[span_data.span_id] = span.get_span_context() - - if span_data.attributes: - attributes: dict[str, AttributeValue] = {} - for key, value in span_data.attributes.items(): - if isinstance(value, (int, float, bool)): - attributes[key] = value - else: - attributes[key] = str(value) - span.set_attributes(attributes) - - if span_data.events: - for event in span_data.events: - span.add_event(event.name, event.attributes, event.timestamp) - - if span_data.status: - span.set_status(span_data.status) - - # Manually end span; do not use context manager to avoid double-end warnings - span.end(end_time=span_data.end_time) - - except Exception: - logger.exception("[Tencent APM] Error creating span: %s", span_data.name) - - def api_check(self) -> bool: - """Check API connectivity using socket connection test for gRPC endpoints""" - try: - # Resolve gRPC target consistently with exporters - _, _, host, port = self._resolve_grpc_target(self.endpoint) - - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.settimeout(5) - result = sock.connect_ex((host, port)) - sock.close() - - if result == 0: - logger.info("[Tencent APM] Endpoint %s:%s is accessible", host, port) - return True - else: - logger.warning("[Tencent APM] Endpoint %s:%s is not accessible", host, port) - if host in ["127.0.0.1", "localhost"]: - logger.info("[Tencent APM] Development environment detected, allowing config save") - return True - return False - - except Exception: - logger.exception("[Tencent APM] API check failed") - return bool("127.0.0.1" in self.endpoint or "localhost" in self.endpoint) - - def get_project_url(self) -> str: - """Get project console URL""" - return "https://console.cloud.tencent.com/apm" - - def shutdown(self) -> None: - """Shutdown the client and export remaining spans""" - try: - if self.span_processor: - logger.info("[Tencent APM] Flushing remaining spans before shutdown") - _ = self.span_processor.force_flush() - self.span_processor.shutdown() - - if self.tracer_provider: - self.tracer_provider.shutdown() - - # Shutdown instance-level meter provider - if self.meter_provider is not None: - try: - self.meter_provider.shutdown() # type: ignore[attr-defined] - except Exception: - logger.debug("[Tencent APM] Error shutting down meter provider", exc_info=True) - - if self.metric_reader is not None: - try: - self.metric_reader.shutdown() # type: ignore[attr-defined] - except Exception: - logger.debug("[Tencent APM] Error shutting down metric reader", exc_info=True) - - except Exception: - logger.exception("[Tencent APM] Error during client shutdown") - - @staticmethod - def _resolve_grpc_target(endpoint: str, default_port: int = 4317) -> tuple[str, bool, str, int]: - """Normalize endpoint to gRPC target and security flag. - - Returns: - (grpc_endpoint, insecure, host, port) - """ - try: - if endpoint.startswith(("http://", "https://")): - parsed = urlparse(endpoint) - host = parsed.hostname or "localhost" - port = parsed.port or default_port - insecure = parsed.scheme == "http" - return f"{host}:{port}", insecure, host, port - - host = endpoint - port = default_port - if ":" in endpoint: - parts = endpoint.rsplit(":", 1) - host = parts[0] or "localhost" - try: - port = int(parts[1]) - except Exception: - port = default_port - - insecure = ("localhost" in host) or ("127.0.0.1" in host) - return f"{host}:{port}", insecure, host, port - except Exception: - host, port = "localhost", default_port - return f"{host}:{port}", True, host, port diff --git a/api/providers/trace/trace-tencent/src/dify_trace_tencent/config.py b/api/providers/trace/trace-tencent/src/dify_trace_tencent/config.py index 398e6c55a86d48..bb4cd3b9262ba5 100644 --- a/api/providers/trace/trace-tencent/src/dify_trace_tencent/config.py +++ b/api/providers/trace/trace-tencent/src/dify_trace_tencent/config.py @@ -1,30 +1,75 @@ -from pydantic import ValidationInfo, field_validator +"""Tencent APM configuration. Validation messages are fixed texts that never include input.""" -from core.ops.entities.config_entity import BaseTracingConfig +import re +from urllib.parse import urlsplit, urlunsplit +from dify_trace_sdk.config import ProviderConfig, destination, secret +from pydantic import field_validator -class TencentConfig(BaseTracingConfig): - """ - Tencent APM tracing config - """ +SERVICE_NAME = re.compile(r"[a-z](?:[a-z0-9-]{0,61}[a-z0-9])?") +_LEGACY_DEFAULT_SERVICE_NAME = "dify_app" +_INVALID_TOKEN = "Tencent APM requires a nonempty token of printable ASCII characters." +_INVALID_ENDPOINT = "Tencent APM requires an HTTP(S) collector origin with a valid host and port." - token: str - endpoint: str + +def _normalized_endpoint(value: str) -> str | None: + """The collector origin with a lowercase ASCII host; ``None`` when the value is not an origin.""" + value = value.strip() + try: + parsed = urlsplit(value) + host, port = parsed.hostname, parsed.port + valid = ( + parsed.scheme in {"http", "https"} + and host + and parsed.username is None + and parsed.password is None + and parsed.path in {"", "/"} + and not parsed.query + and not parsed.fragment + and "?" not in value + and "#" not in value + and "\\" not in value + and not any(ord(char) <= 32 or ord(char) == 127 for char in value) + and not parsed.netloc.endswith(":") + and (port is None or 1 <= port <= 65535) + ) + if not valid or host is None: + return None + host = host.encode("idna").decode("ascii").lower() + except (UnicodeError, ValueError): + # Return instead of raising here so no exception carrying the input is chained. + return None + authority = f"[{host}]" if ":" in host else host + if port is not None: + authority = f"{authority}:{port}" + return urlunsplit((parsed.scheme, authority, "", "", "")) + + +class TencentConfig(ProviderConfig): + token: str = secret() + endpoint: str = destination() service_name: str @field_validator("token") @classmethod - def token_validator(cls, v, info: ValidationInfo): - if not v or v.strip() == "": - raise ValueError("Token cannot be empty") - return v + def _token(cls, value: str) -> str: + # The token is opaque and kept verbatim. gRPC metadata carries only printable + # ASCII, and gRPC reports any other value in an exception that quotes it. + if not value.strip() or not (value.isascii() and value.isprintable()): + raise ValueError(_INVALID_TOKEN) + return value @field_validator("endpoint") @classmethod - def endpoint_validator(cls, v, info: ValidationInfo): - return cls.validate_endpoint_url(v, "https://apm.tencentcloudapi.com") + def _endpoint(cls, value: str) -> str: + endpoint = _normalized_endpoint(value) + if endpoint is None: + raise ValueError(_INVALID_ENDPOINT) + return endpoint @field_validator("service_name") @classmethod - def service_name_validator(cls, v, info: ValidationInfo): - return cls.validate_project_field(v, "dify_app") + def _service_name(cls, value: str) -> str: + # Names saved before verification checked ``SERVICE_NAME`` (such as the old + # ``dify_app`` default) were exported as entered and must keep working. + return value.strip() or _LEGACY_DEFAULT_SERVICE_NAME diff --git a/api/providers/trace/trace-tencent/src/dify_trace_tencent/configuration.py b/api/providers/trace/trace-tencent/src/dify_trace_tencent/configuration.py deleted file mode 100644 index 04146a6e45022b..00000000000000 --- a/api/providers/trace/trace-tencent/src/dify_trace_tencent/configuration.py +++ /dev/null @@ -1,142 +0,0 @@ -"""Tencent configuration syntax and console navigation without exporter state. - -These checks do not establish token validity or ingestion permission. -Collectors may acknowledge zero-span OTLP requests without authenticating -them, so an empty Export acknowledgment must not be presented as authentication. -""" - -import re -from collections.abc import Mapping -from typing import Literal -from urllib.parse import urlsplit, urlunsplit - -from dify_trace_tencent.config import TencentConfig - -type ConfigurationCode = Literal[ - "invalid_endpoint", "invalid_service_name", "invalid_token", "connection_failed", "unexpected_response" -] - -_MESSAGES: dict[ConfigurationCode, str] = { - "invalid_endpoint": "Tencent APM requires an HTTP(S) collector origin with a valid host and port.", - "invalid_service_name": ( - "Tencent APM service names must contain 1 to 63 lowercase letters, digits or hyphens, " - "start with a letter and end with a letter or digit." - ), - "invalid_token": "Tencent APM requires a nonempty token.", - "connection_failed": "Tencent APM collector protocol connection failed.", - "unexpected_response": "Tencent APM collector returned an unexpected protocol acknowledgment.", -} -_SERVICE_NAME = re.compile(r"[a-z](?:[a-z0-9-]{0,61}[a-z0-9])?") - - -class TencentConfigurationError(ValueError): - """Fixed safe errors, without request values or credentials.""" - - def __init__(self, code: ConfigurationCode) -> None: - self.code = code - super().__init__(_MESSAGES[code]) - - -def _endpoint(value: object) -> str: - if not isinstance(value, str): - raise TencentConfigurationError("invalid_endpoint") - value = value.strip() - try: - parsed = urlsplit(value) - host, port = parsed.hostname, parsed.port - valid = ( - parsed.scheme in {"http", "https"} - and host - and parsed.username is None - and parsed.password is None - and parsed.path in {"", "/"} - and not parsed.query - and not parsed.fragment - and "?" not in value - and "#" not in value - and "\\" not in value - and not any(ord(char) <= 32 or ord(char) == 127 for char in value) - and not parsed.netloc.endswith(":") - and (port is None or 1 <= port <= 65535) - ) - if not valid or host is None: - raise TencentConfigurationError("invalid_endpoint") - host = host.encode("idna").decode("ascii").lower() - authority = f"[{host}]" if ":" in host else host - if port is not None: - authority = f"{authority}:{port}" - return urlunsplit((parsed.scheme, authority, "", "", "")) - except (UnicodeError, ValueError): - raise TencentConfigurationError("invalid_endpoint") from None - - -def normalize_configuration(values: Mapping[str, object]) -> TencentConfig: - """Validate new writes before legacy model normalization drops URL parts. - - The permissive legacy model remains available to display existing stored - configurations. Masked tokens are retained here so PATCH can restore the - encrypted original through the shared configuration owner. - """ - endpoint = _endpoint(values.get("endpoint")) - service_name = values.get("service_name") - if not isinstance(service_name, str) or not _SERVICE_NAME.fullmatch(service_name.strip()): - raise TencentConfigurationError("invalid_service_name") - token = values.get("token") - if not isinstance(token, str) or not token.strip(): - raise TencentConfigurationError("invalid_token") - return TencentConfig(endpoint=endpoint, service_name=service_name.strip(), token=token) - - -def console_url() -> str: - """Return the console entry, not a destination-specific project link.""" - return "https://console.cloud.tencent.com/apm" - - -def verify_configuration(config: TencentConfig) -> TencentConfig: - """Check syntax and one zero-span OTLP/gRPC exchange, not authentication. - - A collector may acknowledge invalid-token or unauthenticated empty - requests. Actual credential validity and ingestion must therefore be - assessed from a business trace, never inferred from this connection check. - No tracer, span processor, metrics reader or global SDK state is created. - """ - import grpc # type: ignore[import-untyped] - from opentelemetry.context import _SUPPRESS_INSTRUMENTATION_KEY, attach, detach, set_value - from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest - from opentelemetry.proto.collector.trace.v1.trace_service_pb2_grpc import TraceServiceStub - - config = normalize_configuration(config.model_dump()) - parsed = urlsplit(config.endpoint) - host = parsed.hostname - assert host is not None # Validated above, before constructing a channel. - authority = f"[{host}]" if ":" in host else host - target = f"{authority}:{parsed.port or 4317}" - options = (("grpc.enable_retries", 0), ("grpc.max_receive_message_length", 4 * 1024 * 1024)) - channel: grpc.Channel | None = None - context = attach(set_value(_SUPPRESS_INSTRUMENTATION_KEY, True)) - try: - try: - channel = ( - grpc.insecure_channel(target, options=options) - if parsed.scheme == "http" - else grpc.secure_channel(target, grpc.ssl_channel_credentials(), options=options) - ) - response = TraceServiceStub(channel).Export( - ExportTraceServiceRequest(), - metadata=(("authorization", f"Bearer {config.token}"),), - timeout=10.0, - wait_for_ready=False, - ) - if response.partial_success.rejected_spans or response.partial_success.error_message: - raise TencentConfigurationError("unexpected_response") - finally: - if channel is not None: - channel.close() - except TencentConfigurationError: - raise - except Exception: - # gRPC errors may contain metadata or vendor details. Expose neither. - raise TencentConfigurationError("connection_failed") from None - finally: - detach(context) - return config diff --git a/api/providers/trace/trace-tencent/src/dify_trace_tencent/entities/__init__.py b/api/providers/trace/trace-tencent/src/dify_trace_tencent/entities/__init__.py deleted file mode 100644 index b1602628ed7756..00000000000000 --- a/api/providers/trace/trace-tencent/src/dify_trace_tencent/entities/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Tencent trace entities module diff --git a/api/providers/trace/trace-tencent/src/dify_trace_tencent/entities/semconv.py b/api/providers/trace/trace-tencent/src/dify_trace_tencent/entities/semconv.py deleted file mode 100644 index cd2dbade8bcc33..00000000000000 --- a/api/providers/trace/trace-tencent/src/dify_trace_tencent/entities/semconv.py +++ /dev/null @@ -1,89 +0,0 @@ -from enum import Enum - -# public -GEN_AI_SESSION_ID = "gen_ai.session.id" - -GEN_AI_USER_ID = "gen_ai.user.id" - -GEN_AI_USER_NAME = "gen_ai.user.name" - -GEN_AI_SPAN_KIND = "gen_ai.span.kind" - -GEN_AI_FRAMEWORK = "gen_ai.framework" - -GEN_AI_IS_ENTRY = "gen_ai.is_entry" # mark to count the LLM-related traces - -# Chain -INPUT_VALUE = "gen_ai.entity.input" - -OUTPUT_VALUE = "gen_ai.entity.output" - - -# Retriever -RETRIEVAL_QUERY = "retrieval.query" - -RETRIEVAL_DOCUMENT = "retrieval.document" - - -# GENERATION -GEN_AI_MODEL_NAME = "gen_ai.response.model" - -GEN_AI_PROVIDER = "gen_ai.provider.name" - - -GEN_AI_USAGE_INPUT_TOKENS = "gen_ai.usage.input_tokens" - -GEN_AI_USAGE_OUTPUT_TOKENS = "gen_ai.usage.output_tokens" - -GEN_AI_USAGE_TOTAL_TOKENS = "gen_ai.usage.total_tokens" - -GEN_AI_PROMPT_TEMPLATE_TEMPLATE = "gen_ai.prompt_template.template" - -GEN_AI_PROMPT_TEMPLATE_VARIABLE = "gen_ai.prompt_template.variable" - -GEN_AI_PROMPT = "gen_ai.prompt" - -GEN_AI_COMPLETION = "gen_ai.completion" - -GEN_AI_RESPONSE_FINISH_REASON = "gen_ai.response.finish_reason" - -# Streaming Span Attributes -GEN_AI_IS_STREAMING_REQUEST = "llm.is_streaming" # Same as OpenLLMetry semconv - -# Tool -TOOL_NAME = "tool.name" - -TOOL_DESCRIPTION = "tool.description" - -TOOL_PARAMETERS = "tool.parameters" - -# Instrumentation Library -INSTRUMENTATION_NAME = "dify-sdk" -INSTRUMENTATION_VERSION = "0.1.0" -INSTRUMENTATION_LANGUAGE = "python" - - -# Metrics -LLM_OPERATION_DURATION = "gen_ai.client.operation.duration" -GEN_AI_TOKEN_USAGE = "gen_ai.client.token.usage" -GEN_AI_SERVER_TIME_TO_FIRST_TOKEN = "gen_ai.server.time_to_first_token" -GEN_AI_STREAMING_TIME_TO_GENERATE = "gen_ai.streaming.time_to_generate" -# The LLM trace duration which is exclusive to tencent apm -GEN_AI_TRACE_DURATION = "gen_ai.trace.duration" - -# Token Usage Attributes -GEN_AI_OPERATION_NAME = "gen_ai.operation.name" -GEN_AI_REQUEST_MODEL = "gen_ai.request.model" -GEN_AI_RESPONSE_MODEL = "gen_ai.response.model" -GEN_AI_SYSTEM = "gen_ai.system" -GEN_AI_TOKEN_TYPE = "gen_ai.token.type" -SERVER_ADDRESS = "server.address" - - -class GenAISpanKind(Enum): - WORKFLOW = "WORKFLOW" # OpenLLMetry - RETRIEVER = "RETRIEVER" # RAG - GENERATION = "GENERATION" # Langfuse - TOOL = "TOOL" # OpenLLMetry - AGENT = "AGENT" # OpenLLMetry - TASK = "TASK" # OpenLLMetry diff --git a/api/providers/trace/trace-tencent/src/dify_trace_tencent/entities/tencent_trace_entity.py b/api/providers/trace/trace-tencent/src/dify_trace_tencent/entities/tencent_trace_entity.py deleted file mode 100644 index 428850f1098a64..00000000000000 --- a/api/providers/trace/trace-tencent/src/dify_trace_tencent/entities/tencent_trace_entity.py +++ /dev/null @@ -1,21 +0,0 @@ -from collections.abc import Sequence - -from opentelemetry import trace as trace_api -from opentelemetry.sdk.trace import Event -from opentelemetry.trace import Status, StatusCode -from pydantic import BaseModel, Field - - -class SpanData(BaseModel): - model_config = {"arbitrary_types_allowed": True} - - trace_id: int = Field(..., description="The unique identifier for the trace.") - parent_span_id: int | None = Field(None, description="The ID of the parent span, if any.") - span_id: int = Field(..., description="The unique identifier for this span.") - name: str = Field(..., description="The name of the span.") - attributes: dict[str, str] = Field(default_factory=dict, description="Attributes associated with the span.") - events: Sequence[Event] = Field(default_factory=list, description="Events recorded in the span.") - links: Sequence[trace_api.Link] = Field(default_factory=list, description="Links to other spans.") - status: Status = Field(default=Status(StatusCode.UNSET), description="The status of the span.") - start_time: int = Field(..., description="The start time of the span in nanoseconds.") - end_time: int = Field(..., description="The end time of the span in nanoseconds.") diff --git a/api/providers/trace/trace-tencent/src/dify_trace_tencent/provider.py b/api/providers/trace/trace-tencent/src/dify_trace_tencent/provider.py new file mode 100644 index 00000000000000..bee6adf4ac6565 --- /dev/null +++ b/api/providers/trace/trace-tencent/src/dify_trace_tencent/provider.py @@ -0,0 +1,98 @@ +"""Tencent APM over its OTLP/gRPC collector. + +The collector authenticates a bearer token sent as gRPC metadata. Tencent APM +counts LLM traces by their entry span, so the root of every unit that starts a +trace carries ``gen_ai.is_entry``; a fragment joining an existing trace does not. +""" + +from typing import override +from urllib.parse import urlsplit + +import grpc # type: ignore[import-untyped] +from dify_trace_sdk.errors import EgressError, VerificationError +from dify_trace_sdk.genai import GenAIDialect, genai_attributes +from dify_trace_sdk.model import SpanKind, Trace +from dify_trace_sdk.otlp import encode_trace, export_otlp_grpc +from dify_trace_sdk.provider import ProviderCapabilities, SendResult, TraceProvider +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest +from opentelemetry.proto.collector.trace.v1.trace_service_pb2_grpc import TraceServiceStub + +from dify_trace_tencent.config import SERVICE_NAME, TencentConfig + +DEFAULT_GRPC_PORT = 4317 +VERIFY_TIMEOUT_SECONDS = 10.0 +_DESTINATION = "Tencent APM" +_CONSOLE_URL = "https://console.cloud.tencent.com/apm" +_DIALECT = GenAIDialect( + kind_labels={ + **{kind: kind.value.upper() for kind in SpanKind}, + SpanKind.LLM: "GENERATION", + SpanKind.CHAIN: "WORKFLOW", + }, + input_key="gen_ai.entity.input", + output_key="gen_ai.entity.output", + mark_entry_span=True, +) + + +def _grpc_endpoint(config: TencentConfig) -> str: + """The collector listens on the OTLP/gRPC port unless the endpoint names another one.""" + if urlsplit(config.endpoint).port is None: + return f"{config.endpoint}:{DEFAULT_GRPC_PORT}" + return config.endpoint + + +def _metadata(config: TencentConfig) -> dict[str, str]: + return {"authorization": f"Bearer {config.token}"} + + +class TencentProvider(TraceProvider[TencentConfig]): + id = "tencent" + config_model = TencentConfig + capabilities = ProviderCapabilities(native_cost=False, trace_url=False, span_url=False) + + @override + def verify(self, config: TencentConfig) -> TencentConfig: + """Exchange one zero-span Export with the collector. + + This shows the endpoint is an OTLP/gRPC collector, not that the token + is valid: collectors may acknowledge empty requests without + authenticating them, so only business traces establish ingestion. + """ + if not SERVICE_NAME.fullmatch(config.service_name): + raise VerificationError("invalid_service_name") + try: + channel = self.egress.grpc_channel(_grpc_endpoint(config)) + except EgressError: + raise VerificationError("connection_failed") from None + try: + response = TraceServiceStub(channel).Export( + ExportTraceServiceRequest(), + metadata=tuple(_metadata(config).items()), + timeout=VERIFY_TIMEOUT_SECONDS, + wait_for_ready=False, + ) + except (grpc.RpcError, ValueError): + # gRPC raises ValueError for metadata it cannot send, quoting the token. + raise VerificationError("connection_failed") from None + finally: + channel.close() + if response.partial_success.rejected_spans or response.partial_success.error_message: + raise VerificationError("unexpected_response") + return config + + @override + def project_url(self, config: TencentConfig) -> str | None: + return _CONSOLE_URL + + @override + def send(self, config: TencentConfig, trace: Trace) -> SendResult: + request = encode_trace( + trace, + resource={"service.name": config.service_name}, + attributes=lambda span: genai_attributes(span, trace, _DIALECT), + ) + export_otlp_grpc( + self.egress, _grpc_endpoint(config), request, destination=_DESTINATION, metadata=_metadata(config) + ) + return SendResult(trace_id=trace.trace_id, trace_url=None, span_urls={}) diff --git a/api/providers/trace/trace-tencent/src/dify_trace_tencent/span_builder.py b/api/providers/trace/trace-tencent/src/dify_trace_tencent/span_builder.py deleted file mode 100644 index a473e9141afb08..00000000000000 --- a/api/providers/trace/trace-tencent/src/dify_trace_tencent/span_builder.py +++ /dev/null @@ -1,383 +0,0 @@ -""" -Tencent APM Span Builder - handles all span construction logic -""" - -import json -import logging -from datetime import datetime - -from opentelemetry.trace import Status, StatusCode - -from core.ops.entities.trace_entity import ( - DatasetRetrievalTraceInfo, - MessageTraceInfo, - ToolTraceInfo, - WorkflowTraceInfo, -) -from core.rag.models.document import Document -from dify_trace_tencent.entities.semconv import ( - GEN_AI_COMPLETION, - GEN_AI_FRAMEWORK, - GEN_AI_IS_ENTRY, - GEN_AI_IS_STREAMING_REQUEST, - GEN_AI_MODEL_NAME, - GEN_AI_PROMPT, - GEN_AI_PROVIDER, - GEN_AI_RESPONSE_FINISH_REASON, - GEN_AI_SESSION_ID, - GEN_AI_SPAN_KIND, - GEN_AI_USAGE_INPUT_TOKENS, - GEN_AI_USAGE_OUTPUT_TOKENS, - GEN_AI_USAGE_TOTAL_TOKENS, - GEN_AI_USER_ID, - INPUT_VALUE, - OUTPUT_VALUE, - RETRIEVAL_DOCUMENT, - RETRIEVAL_QUERY, - TOOL_DESCRIPTION, - TOOL_NAME, - TOOL_PARAMETERS, - GenAISpanKind, -) -from dify_trace_tencent.entities.tencent_trace_entity import SpanData -from dify_trace_tencent.utils import TencentTraceUtils -from graphon.entities import WorkflowNodeExecution -from graphon.enums import WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus - -logger = logging.getLogger(__name__) - - -class TencentSpanBuilder: - """Builder class for constructing different types of spans""" - - @staticmethod - def _get_time_nanoseconds(time_value: datetime | None) -> int: - """Convert datetime to nanoseconds for span creation.""" - return TencentTraceUtils.convert_datetime_to_nanoseconds(time_value) - - @staticmethod - def build_workflow_spans( - trace_info: WorkflowTraceInfo, trace_id: int, user_id: str, links: list | None = None - ) -> list[SpanData]: - """Build workflow-related spans""" - spans = [] - links = links or [] - - message_span_id = None - workflow_span_id = TencentTraceUtils.convert_to_span_id(trace_info.workflow_run_id, "workflow") - - if hasattr(trace_info, "metadata") and trace_info.metadata.get("conversation_id"): - message_span_id = TencentTraceUtils.convert_to_span_id(trace_info.workflow_run_id, "message") - - status = Status(StatusCode.OK) - if trace_info.error: - status = Status(StatusCode.ERROR, trace_info.error) - - if message_span_id: - message_span = TencentSpanBuilder._build_message_span( - trace_info, trace_id, message_span_id, user_id, status, links - ) - spans.append(message_span) - - workflow_span = TencentSpanBuilder._build_workflow_span( - trace_info, trace_id, workflow_span_id, message_span_id, user_id, status, links - ) - spans.append(workflow_span) - - return spans - - @staticmethod - def _build_message_span( - trace_info: WorkflowTraceInfo, trace_id: int, message_span_id: int, user_id: str, status: Status, links: list - ) -> SpanData: - """Build message span for chatflow""" - return SpanData( - trace_id=trace_id, - parent_span_id=None, - span_id=message_span_id, - name="message", - start_time=TencentSpanBuilder._get_time_nanoseconds(trace_info.start_time), - end_time=TencentSpanBuilder._get_time_nanoseconds(trace_info.end_time), - attributes={ - GEN_AI_SESSION_ID: trace_info.metadata.get("conversation_id", ""), - GEN_AI_USER_ID: str(user_id), - GEN_AI_SPAN_KIND: GenAISpanKind.WORKFLOW.value, - GEN_AI_FRAMEWORK: "dify", - GEN_AI_IS_ENTRY: "true", - INPUT_VALUE: trace_info.workflow_run_inputs.get("sys.query", ""), - OUTPUT_VALUE: json.dumps(trace_info.workflow_run_outputs, ensure_ascii=False), - }, - status=status, - links=links, - ) - - @staticmethod - def _build_workflow_span( - trace_info: WorkflowTraceInfo, - trace_id: int, - workflow_span_id: int, - message_span_id: int | None, - user_id: str, - status: Status, - links: list, - ) -> SpanData: - """Build workflow span""" - attributes = { - GEN_AI_USER_ID: str(user_id), - GEN_AI_SPAN_KIND: GenAISpanKind.WORKFLOW.value, - GEN_AI_FRAMEWORK: "dify", - INPUT_VALUE: json.dumps(trace_info.workflow_run_inputs, ensure_ascii=False), - OUTPUT_VALUE: json.dumps(trace_info.workflow_run_outputs, ensure_ascii=False), - } - - if message_span_id is None: - attributes[GEN_AI_IS_ENTRY] = "true" - - return SpanData( - trace_id=trace_id, - parent_span_id=message_span_id, - span_id=workflow_span_id, - name="workflow", - start_time=TencentSpanBuilder._get_time_nanoseconds(trace_info.start_time), - end_time=TencentSpanBuilder._get_time_nanoseconds(trace_info.end_time), - attributes=attributes, - status=status, - links=links, - ) - - @staticmethod - def build_workflow_llm_span( - trace_id: int, workflow_span_id: int, trace_info: WorkflowTraceInfo, node_execution: WorkflowNodeExecution - ) -> SpanData: - """Build LLM span for workflow nodes.""" - process_data = node_execution.process_data or {} - outputs = node_execution.outputs or {} - usage_data = process_data.get("usage", {}) if "usage" in process_data else outputs.get("usage", {}) - - session_id = trace_info.metadata.get("conversation_id") - attributes = { - GEN_AI_SESSION_ID: "" if session_id is None else session_id, - GEN_AI_SPAN_KIND: GenAISpanKind.GENERATION.value, - GEN_AI_FRAMEWORK: "dify", - GEN_AI_MODEL_NAME: process_data.get("model_name", ""), - GEN_AI_PROVIDER: process_data.get("model_provider", ""), - GEN_AI_USAGE_INPUT_TOKENS: str(usage_data.get("prompt_tokens", 0)), - GEN_AI_USAGE_OUTPUT_TOKENS: str(usage_data.get("completion_tokens", 0)), - GEN_AI_USAGE_TOTAL_TOKENS: str(usage_data.get("total_tokens", 0)), - GEN_AI_PROMPT: json.dumps(process_data.get("prompts", []), ensure_ascii=False), - GEN_AI_COMPLETION: str(outputs.get("text", "")), - GEN_AI_RESPONSE_FINISH_REASON: outputs.get("finish_reason", ""), - INPUT_VALUE: json.dumps(process_data.get("prompts", []), ensure_ascii=False), - OUTPUT_VALUE: str(outputs.get("text", "")), - } - - if usage_data.get("time_to_first_token") is not None: - attributes[GEN_AI_IS_STREAMING_REQUEST] = "true" - - return SpanData( - trace_id=trace_id, - parent_span_id=workflow_span_id, - span_id=TencentTraceUtils.convert_to_span_id(node_execution.id, "node"), - name="GENERATION", - start_time=TencentSpanBuilder._get_time_nanoseconds(node_execution.created_at), - end_time=TencentSpanBuilder._get_time_nanoseconds(node_execution.finished_at), - attributes=attributes, - status=TencentSpanBuilder._get_workflow_node_status(node_execution), - ) - - @staticmethod - def build_message_span( - trace_info: MessageTraceInfo, trace_id: int, user_id: str, links: list | None = None - ) -> SpanData: - """Build message span.""" - links = links or [] - status = Status(StatusCode.OK) - if trace_info.error: - status = Status(StatusCode.ERROR, trace_info.error) - - attributes = { - GEN_AI_SESSION_ID: trace_info.metadata.get("conversation_id", ""), - GEN_AI_USER_ID: str(user_id), - GEN_AI_SPAN_KIND: GenAISpanKind.WORKFLOW.value, - GEN_AI_FRAMEWORK: "dify", - GEN_AI_IS_ENTRY: "true", - INPUT_VALUE: str(trace_info.inputs or ""), - OUTPUT_VALUE: str(trace_info.outputs or ""), - } - - if trace_info.is_streaming_request: - attributes[GEN_AI_IS_STREAMING_REQUEST] = "true" - - return SpanData( - trace_id=trace_id, - parent_span_id=None, - span_id=TencentTraceUtils.convert_to_span_id(trace_info.message_id, "message"), - name="message", - start_time=TencentSpanBuilder._get_time_nanoseconds(trace_info.start_time), - end_time=TencentSpanBuilder._get_time_nanoseconds(trace_info.end_time), - attributes=attributes, - status=status, - links=links, - ) - - @staticmethod - def build_tool_span(trace_info: ToolTraceInfo, trace_id: int, parent_span_id: int) -> SpanData: - """Build tool span.""" - status = Status(StatusCode.OK) - if trace_info.error: - status = Status(StatusCode.ERROR, trace_info.error) - - return SpanData( - trace_id=trace_id, - parent_span_id=parent_span_id, - span_id=TencentTraceUtils.convert_to_span_id(trace_info.message_id, "tool"), - name=trace_info.tool_name, - start_time=TencentSpanBuilder._get_time_nanoseconds(trace_info.start_time), - end_time=TencentSpanBuilder._get_time_nanoseconds(trace_info.end_time), - attributes={ - GEN_AI_SPAN_KIND: GenAISpanKind.TOOL.value, - GEN_AI_FRAMEWORK: "dify", - TOOL_NAME: trace_info.tool_name, - TOOL_DESCRIPTION: "", - TOOL_PARAMETERS: json.dumps(trace_info.tool_parameters, ensure_ascii=False), - INPUT_VALUE: json.dumps(trace_info.tool_inputs, ensure_ascii=False), - OUTPUT_VALUE: str(trace_info.tool_outputs), - }, - status=status, - ) - - @staticmethod - def build_retrieval_span(trace_info: DatasetRetrievalTraceInfo, trace_id: int, parent_span_id: int) -> SpanData: - """Build dataset retrieval span.""" - status = Status(StatusCode.OK) - if getattr(trace_info, "error", None): - status = Status(StatusCode.ERROR, trace_info.error) # type: ignore[arg-type] - - documents_data = TencentSpanBuilder._extract_retrieval_documents(trace_info.documents) - - return SpanData( - trace_id=trace_id, - parent_span_id=parent_span_id, - span_id=TencentTraceUtils.convert_to_span_id(trace_info.message_id, "retrieval"), - name="retrieval", - start_time=TencentSpanBuilder._get_time_nanoseconds(trace_info.start_time), - end_time=TencentSpanBuilder._get_time_nanoseconds(trace_info.end_time), - attributes={ - GEN_AI_SPAN_KIND: GenAISpanKind.RETRIEVER.value, - GEN_AI_FRAMEWORK: "dify", - RETRIEVAL_QUERY: str(trace_info.inputs or ""), - RETRIEVAL_DOCUMENT: json.dumps(documents_data, ensure_ascii=False), - INPUT_VALUE: str(trace_info.inputs or ""), - OUTPUT_VALUE: json.dumps(documents_data, ensure_ascii=False), - }, - status=status, - ) - - @staticmethod - def _get_workflow_node_status(node_execution: WorkflowNodeExecution) -> Status: - """Get workflow node execution status.""" - if node_execution.status == WorkflowNodeExecutionStatus.SUCCEEDED: - return Status(StatusCode.OK) - elif node_execution.status in [WorkflowNodeExecutionStatus.FAILED, WorkflowNodeExecutionStatus.EXCEPTION]: - return Status(StatusCode.ERROR, str(node_execution.error)) - return Status(StatusCode.UNSET) - - @staticmethod - def build_workflow_retrieval_span( - trace_id: int, workflow_span_id: int, trace_info: WorkflowTraceInfo, node_execution: WorkflowNodeExecution - ) -> SpanData: - """Build knowledge retrieval span for workflow nodes.""" - input_value = "" - if node_execution.inputs: - input_value = str(node_execution.inputs.get("query", "")) - output_value = "" - if node_execution.outputs: - output_value = json.dumps(node_execution.outputs.get("result", []), ensure_ascii=False) - - session_id = trace_info.metadata.get("conversation_id") - return SpanData( - trace_id=trace_id, - parent_span_id=workflow_span_id, - span_id=TencentTraceUtils.convert_to_span_id(node_execution.id, "node"), - name=node_execution.title, - start_time=TencentSpanBuilder._get_time_nanoseconds(node_execution.created_at), - end_time=TencentSpanBuilder._get_time_nanoseconds(node_execution.finished_at), - attributes={ - GEN_AI_SESSION_ID: "" if session_id is None else session_id, - GEN_AI_SPAN_KIND: GenAISpanKind.RETRIEVER.value, - GEN_AI_FRAMEWORK: "dify", - RETRIEVAL_QUERY: input_value, - RETRIEVAL_DOCUMENT: output_value, - INPUT_VALUE: input_value, - OUTPUT_VALUE: output_value, - }, - status=TencentSpanBuilder._get_workflow_node_status(node_execution), - ) - - @staticmethod - def build_workflow_tool_span( - trace_id: int, workflow_span_id: int, trace_info: WorkflowTraceInfo, node_execution: WorkflowNodeExecution - ) -> SpanData: - """Build tool span for workflow nodes.""" - tool_des = {} - if node_execution.metadata: - tool_des = node_execution.metadata.get(WorkflowNodeExecutionMetadataKey.TOOL_INFO, {}) - - return SpanData( - trace_id=trace_id, - parent_span_id=workflow_span_id, - span_id=TencentTraceUtils.convert_to_span_id(node_execution.id, "node"), - name=node_execution.title, - start_time=TencentSpanBuilder._get_time_nanoseconds(node_execution.created_at), - end_time=TencentSpanBuilder._get_time_nanoseconds(node_execution.finished_at), - attributes={ - GEN_AI_SPAN_KIND: GenAISpanKind.TOOL.value, - GEN_AI_FRAMEWORK: "dify", - TOOL_NAME: node_execution.title, - TOOL_DESCRIPTION: json.dumps(tool_des, ensure_ascii=False), - TOOL_PARAMETERS: json.dumps(node_execution.inputs or {}, ensure_ascii=False), - INPUT_VALUE: json.dumps(node_execution.inputs or {}, ensure_ascii=False), - OUTPUT_VALUE: json.dumps(node_execution.outputs, ensure_ascii=False), - }, - status=TencentSpanBuilder._get_workflow_node_status(node_execution), - ) - - @staticmethod - def build_workflow_task_span( - trace_id: int, workflow_span_id: int, trace_info: WorkflowTraceInfo, node_execution: WorkflowNodeExecution - ) -> SpanData: - """Build generic task span for workflow nodes.""" - session_id = trace_info.metadata.get("conversation_id") - return SpanData( - trace_id=trace_id, - parent_span_id=workflow_span_id, - span_id=TencentTraceUtils.convert_to_span_id(node_execution.id, "node"), - name=node_execution.title, - start_time=TencentSpanBuilder._get_time_nanoseconds(node_execution.created_at), - end_time=TencentSpanBuilder._get_time_nanoseconds(node_execution.finished_at), - attributes={ - GEN_AI_SESSION_ID: "" if session_id is None else session_id, - GEN_AI_SPAN_KIND: GenAISpanKind.TASK.value, - GEN_AI_FRAMEWORK: "dify", - INPUT_VALUE: json.dumps(node_execution.inputs, ensure_ascii=False), - OUTPUT_VALUE: json.dumps(node_execution.outputs, ensure_ascii=False), - }, - status=TencentSpanBuilder._get_workflow_node_status(node_execution), - ) - - @staticmethod - def _extract_retrieval_documents(documents: list[Document]): - """Extract documents data for retrieval tracing.""" - documents_data = [] - for document in documents: - document_data = { - "content": document.page_content, - "metadata": { - "dataset_id": document.metadata.get("dataset_id"), - "doc_id": document.metadata.get("doc_id"), - "document_id": document.metadata.get("document_id"), - }, - "score": document.metadata.get("score"), - } - documents_data.append(document_data) - return documents_data diff --git a/api/providers/trace/trace-tencent/src/dify_trace_tencent/tencent_trace.py b/api/providers/trace/trace-tencent/src/dify_trace_tencent/tencent_trace.py deleted file mode 100644 index e123042f4cf5eb..00000000000000 --- a/api/providers/trace/trace-tencent/src/dify_trace_tencent/tencent_trace.py +++ /dev/null @@ -1,538 +0,0 @@ -"""Tencent APM tracing with idempotent client cleanup.""" - -import inspect -import logging -from typing import override - -from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker - -from core.ops.base_trace_instance import BaseTraceInstance -from core.ops.entities.trace_entity import ( - BaseTraceInfo, - DatasetRetrievalTraceInfo, - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - SuggestedQuestionTraceInfo, - ToolTraceInfo, - WorkflowTraceInfo, -) -from core.repositories import SQLAlchemyWorkflowNodeExecutionRepository -from dify_trace_tencent.client import TencentTraceClient -from dify_trace_tencent.config import TencentConfig -from dify_trace_tencent.entities.tencent_trace_entity import SpanData -from dify_trace_tencent.span_builder import TencentSpanBuilder -from dify_trace_tencent.utils import TencentTraceUtils -from extensions.ext_database import db -from graphon.entities.workflow_node_execution import ( - WorkflowNodeExecution, -) -from graphon.nodes import BuiltinNodeTypes -from models import Account, App, TenantAccountJoin, WorkflowNodeExecutionTriggeredFrom - -logger = logging.getLogger(__name__) - - -class TencentDataTrace(BaseTraceInstance): - """ - Tencent APM trace implementation with single responsibility principle. - Acts as a coordinator that delegates specific tasks to specialized classes. - - The instance owns a long-lived ``TencentTraceClient``. Cleanup may happen - explicitly in tests or implicitly during garbage collection, so shutdown - must be safe to call multiple times. - """ - - trace_client: TencentTraceClient - _closed: bool - - def __init__(self, tencent_config: TencentConfig): - super().__init__(tencent_config) - self._closed = False - self.trace_client = TencentTraceClient( - service_name=tencent_config.service_name, - endpoint=tencent_config.endpoint, - token=tencent_config.token, - metrics_export_interval_sec=5, - ) - - @override - def trace(self, trace_info: BaseTraceInfo) -> None: - """Main tracing entry point - coordinates different trace types.""" - match trace_info: - case WorkflowTraceInfo(): - self.workflow_trace(trace_info) - case MessageTraceInfo(): - self.message_trace(trace_info) - case ModerationTraceInfo(): - pass - case SuggestedQuestionTraceInfo(): - self.suggested_question_trace(trace_info) - case DatasetRetrievalTraceInfo(): - self.dataset_retrieval_trace(trace_info) - case ToolTraceInfo(): - self.tool_trace(trace_info) - case GenerateNameTraceInfo(): - pass - - def api_check(self) -> bool: - return self.trace_client.api_check() - - def get_project_url(self) -> str: - return self.trace_client.get_project_url() - - def workflow_trace(self, trace_info: WorkflowTraceInfo) -> None: - """Handle workflow tracing by coordinating data retrieval and span construction.""" - try: - trace_id = TencentTraceUtils.convert_to_trace_id(trace_info.workflow_run_id) - - links = [] - if trace_info.trace_id: - links.append(TencentTraceUtils.create_link(trace_info.trace_id)) - - user_id = self._get_user_id(trace_info) - - workflow_spans = TencentSpanBuilder.build_workflow_spans(trace_info, trace_id, str(user_id), links) - - for span in workflow_spans: - self.trace_client.add_span(span) - - self._process_workflow_nodes(trace_info, trace_id) - - # Record trace duration for entry span - self._record_workflow_trace_duration(trace_info) - - except Exception: - logger.exception("[Tencent APM] Failed to process workflow trace") - - def message_trace(self, trace_info: MessageTraceInfo) -> None: - """Handle message tracing.""" - try: - trace_id = TencentTraceUtils.convert_to_trace_id(trace_info.message_id) - user_id = self._get_user_id(trace_info) - - links = [] - if trace_info.trace_id: - links.append(TencentTraceUtils.create_link(trace_info.trace_id)) - - message_span = TencentSpanBuilder.build_message_span(trace_info, trace_id, str(user_id), links) - - self.trace_client.add_span(message_span) - - self._record_message_llm_metrics(trace_info) - - # Record trace duration for entry span - self._record_message_trace_duration(trace_info) - - except Exception: - logger.exception("[Tencent APM] Failed to process message trace") - - def tool_trace(self, trace_info: ToolTraceInfo) -> None: - """Handle tool tracing.""" - try: - parent_span_id = None - trace_root_id = None - - if trace_info.message_id: - parent_span_id = TencentTraceUtils.convert_to_span_id(trace_info.message_id, "message") - trace_root_id = trace_info.message_id - - if parent_span_id and trace_root_id: - trace_id = TencentTraceUtils.convert_to_trace_id(trace_root_id) - - tool_span = TencentSpanBuilder.build_tool_span(trace_info, trace_id, parent_span_id) - - self.trace_client.add_span(tool_span) - - except Exception: - logger.exception("[Tencent APM] Failed to process tool trace") - - def dataset_retrieval_trace(self, trace_info: DatasetRetrievalTraceInfo) -> None: - """Handle dataset retrieval tracing.""" - try: - parent_span_id = None - trace_root_id = None - - if trace_info.message_id: - parent_span_id = TencentTraceUtils.convert_to_span_id(trace_info.message_id, "message") - trace_root_id = trace_info.message_id - - if parent_span_id and trace_root_id: - trace_id = TencentTraceUtils.convert_to_trace_id(trace_root_id) - - retrieval_span = TencentSpanBuilder.build_retrieval_span(trace_info, trace_id, parent_span_id) - - self.trace_client.add_span(retrieval_span) - - except Exception: - logger.exception("[Tencent APM] Failed to process dataset retrieval trace") - - def suggested_question_trace(self, trace_info: SuggestedQuestionTraceInfo) -> None: - """Handle suggested question tracing""" - try: - logger.info("[Tencent APM] Processing suggested question trace") - - except Exception: - logger.exception("[Tencent APM] Failed to process suggested question trace") - - def _process_workflow_nodes(self, trace_info: WorkflowTraceInfo, trace_id: int) -> None: - """Process workflow node executions.""" - try: - workflow_span_id = TencentTraceUtils.convert_to_span_id(trace_info.workflow_run_id, "workflow") - - node_executions = self._get_workflow_node_executions(trace_info) - - for node_execution in node_executions: - try: - node_span = self._build_workflow_node_span(node_execution, trace_id, trace_info, workflow_span_id) - if node_span: - self.trace_client.add_span(node_span) - - if node_execution.node_type == BuiltinNodeTypes.LLM: - self._record_llm_metrics(node_execution) - except Exception: - logger.exception("[Tencent APM] Failed to process node execution: %s", node_execution.id) - - except Exception: - logger.exception("[Tencent APM] Failed to process workflow nodes") - - def _build_workflow_node_span( - self, node_execution: WorkflowNodeExecution, trace_id: int, trace_info: WorkflowTraceInfo, workflow_span_id: int - ) -> SpanData | None: - """Build span for different node types""" - try: - if node_execution.node_type == BuiltinNodeTypes.LLM: - return TencentSpanBuilder.build_workflow_llm_span( - trace_id, workflow_span_id, trace_info, node_execution - ) - elif node_execution.node_type == BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL: - return TencentSpanBuilder.build_workflow_retrieval_span( - trace_id, workflow_span_id, trace_info, node_execution - ) - elif node_execution.node_type == BuiltinNodeTypes.TOOL: - return TencentSpanBuilder.build_workflow_tool_span( - trace_id, workflow_span_id, trace_info, node_execution - ) - else: - # Handle all other node types as generic tasks - return TencentSpanBuilder.build_workflow_task_span( - trace_id, workflow_span_id, trace_info, node_execution - ) - except Exception: - logger.debug( - "[Tencent APM] Error building span for node %s: %s", - node_execution.id, - node_execution.node_type, - exc_info=True, - ) - return None - - def _get_workflow_node_executions(self, trace_info: WorkflowTraceInfo) -> list[WorkflowNodeExecution]: - """Retrieve workflow node executions from database.""" - try: - session_maker = sessionmaker(bind=db.engine) - - with Session(db.engine, expire_on_commit=False) as session: - app_id = trace_info.metadata.get("app_id") - if not app_id: - raise ValueError("No app_id found in trace_info metadata") - - app_stmt = select(App).where(App.id == app_id) - app = session.scalar(app_stmt) - if not app: - raise ValueError(f"App with id {app_id} not found") - - if not app.created_by: - raise ValueError(f"App with id {app_id} has no creator") - - account_stmt = select(Account).where(Account.id == app.created_by) - service_account = session.scalar(account_stmt) - if not service_account: - raise ValueError(f"Creator account not found for app {app_id}") - - repository = SQLAlchemyWorkflowNodeExecutionRepository( - session_factory=session_maker, - tenant_id=app.tenant_id, - user=service_account, - app_id=trace_info.metadata.get("app_id"), - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - - executions = repository.get_by_workflow_execution(workflow_execution_id=trace_info.workflow_run_id) - return list(executions) - - except Exception: - logger.exception("[Tencent APM] Failed to get workflow node executions") - return [] - - def _get_user_id(self, trace_info: BaseTraceInfo) -> str: - """Get user ID from trace info.""" - try: - tenant_id = None - user_id = None - - if isinstance(trace_info, (WorkflowTraceInfo, GenerateNameTraceInfo)): - tenant_id = trace_info.tenant_id - - if hasattr(trace_info, "metadata") and trace_info.metadata: - user_id = trace_info.metadata.get("user_id") - - if user_id and tenant_id: - stmt = ( - select(Account.name) - .join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id) - .where(Account.id == user_id, TenantAccountJoin.tenant_id == tenant_id) - ) - - session_maker = sessionmaker(bind=db.engine) - with session_maker() as session: - account_name = session.scalar(stmt) - return account_name or str(user_id) - elif user_id: - return str(user_id) - - return "anonymous" - - except Exception: - logger.exception("[Tencent APM] Failed to get user ID") - return "unknown" - - def _record_llm_metrics(self, node_execution: WorkflowNodeExecution) -> None: - """Record LLM performance metrics""" - try: - process_data = node_execution.process_data or {} - outputs = node_execution.outputs or {} - usage = process_data.get("usage", {}) if "usage" in process_data else outputs.get("usage", {}) - - model_provider = process_data.get("model_provider", "unknown") - model_name = process_data.get("model_name", "unknown") - model_mode = process_data.get("model_mode", "chat") - - # Record LLM duration - if hasattr(self.trace_client, "record_llm_duration"): - latency_s = float(usage.get("latency", 0.0)) - - if latency_s > 0: - # Determine if streaming from usage metrics - is_streaming = usage.get("time_to_first_token") is not None - - attributes = { - "gen_ai.system": model_provider, - "gen_ai.response.model": model_name, - "gen_ai.operation.name": model_mode, - "stream": "true" if is_streaming else "false", - } - self.trace_client.record_llm_duration(latency_s, attributes) - - # Record streaming metrics from usage - time_to_first_token = usage.get("time_to_first_token") - if time_to_first_token is not None and hasattr(self.trace_client, "record_time_to_first_token"): - ttft_seconds = float(time_to_first_token) - if ttft_seconds > 0: - self.trace_client.record_time_to_first_token( - ttft_seconds=ttft_seconds, provider=model_provider, model=model_name, operation_name=model_mode - ) - - time_to_generate = usage.get("time_to_generate") - if time_to_generate is not None and hasattr(self.trace_client, "record_time_to_generate"): - ttg_seconds = float(time_to_generate) - if ttg_seconds > 0: - self.trace_client.record_time_to_generate( - ttg_seconds=ttg_seconds, provider=model_provider, model=model_name, operation_name=model_mode - ) - - # Record token usage - if hasattr(self.trace_client, "record_token_usage"): - # Extract token counts - input_tokens = int(usage.get("prompt_tokens", 0)) - output_tokens = int(usage.get("completion_tokens", 0)) - - if input_tokens > 0 or output_tokens > 0: - server_address = f"{model_provider}" - - # Record input tokens - if input_tokens > 0: - self.trace_client.record_token_usage( - token_count=input_tokens, - token_type="input", - operation_name=model_mode, - request_model=model_name, - response_model=model_name, - server_address=server_address, - provider=model_provider, - ) - - # Record output tokens - if output_tokens > 0: - self.trace_client.record_token_usage( - token_count=output_tokens, - token_type="output", - operation_name=model_mode, - request_model=model_name, - response_model=model_name, - server_address=server_address, - provider=model_provider, - ) - - except Exception: - logger.debug("[Tencent APM] Failed to record LLM metrics") - - def _record_message_llm_metrics(self, trace_info: MessageTraceInfo) -> None: - """Record LLM metrics for message traces""" - try: - trace_metadata = trace_info.metadata or {} - message_data = trace_info.message_data or {} - provider_latency = 0.0 - if isinstance(message_data, dict): - provider_latency = float(message_data.get("provider_response_latency", 0.0) or 0.0) - else: - provider_latency = float(getattr(message_data, "provider_response_latency", 0.0) or 0.0) - - model_provider = trace_metadata.get("ls_provider") or ( - message_data.get("model_provider", "") if isinstance(message_data, dict) else "" - ) - model_name = trace_metadata.get("ls_model_name") or ( - message_data.get("model_id", "") if isinstance(message_data, dict) else "" - ) - - # Record LLM duration - if provider_latency > 0 and hasattr(self.trace_client, "record_llm_duration"): - is_streaming = trace_info.is_streaming_request - - duration_attributes = { - "gen_ai.system": model_provider, - "gen_ai.response.model": model_name, - "gen_ai.operation.name": "chat", # Message traces are always chat - "stream": "true" if is_streaming else "false", - } - self.trace_client.record_llm_duration(provider_latency, duration_attributes) - - # Record streaming metrics for message traces - if trace_info.is_streaming_request: - # Record time to first token - if trace_info.gen_ai_server_time_to_first_token is not None and hasattr( - self.trace_client, "record_time_to_first_token" - ): - ttft_seconds = float(trace_info.gen_ai_server_time_to_first_token) - if ttft_seconds > 0: - self.trace_client.record_time_to_first_token( - ttft_seconds=ttft_seconds, provider=str(model_provider or ""), model=str(model_name or "") - ) - - # Record time to generate - if trace_info.llm_streaming_time_to_generate is not None and hasattr( - self.trace_client, "record_time_to_generate" - ): - ttg_seconds = float(trace_info.llm_streaming_time_to_generate) - if ttg_seconds > 0: - self.trace_client.record_time_to_generate( - ttg_seconds=ttg_seconds, provider=str(model_provider or ""), model=str(model_name or "") - ) - - # Record token usage - if hasattr(self.trace_client, "record_token_usage"): - input_tokens = int(trace_info.message_tokens or 0) - output_tokens = int(trace_info.answer_tokens or 0) - - if input_tokens > 0: - self.trace_client.record_token_usage( - token_count=input_tokens, - token_type="input", - operation_name="chat", - request_model=str(model_name or ""), - response_model=str(model_name or ""), - server_address=str(model_provider or ""), - provider=str(model_provider or ""), - ) - - if output_tokens > 0: - self.trace_client.record_token_usage( - token_count=output_tokens, - token_type="output", - operation_name="chat", - request_model=str(model_name or ""), - response_model=str(model_name or ""), - server_address=str(model_provider or ""), - provider=str(model_provider or ""), - ) - - except Exception: - logger.debug("[Tencent APM] Failed to record message LLM metrics") - - def _record_workflow_trace_duration(self, trace_info: WorkflowTraceInfo) -> None: - """Record end-to-end workflow trace duration.""" - try: - if not hasattr(self.trace_client, "record_trace_duration"): - return - - # Calculate duration from start_time and end_time to match span duration - if trace_info.start_time and trace_info.end_time: - duration_s = (trace_info.end_time - trace_info.start_time).total_seconds() - else: - # Fallback to workflow_run_elapsed_time if timestamps not available - duration_s = float(trace_info.workflow_run_elapsed_time) - - if duration_s > 0: - attributes = { - "conversation_mode": "workflow", - "workflow_status": trace_info.workflow_run_status, - } - - # Add conversation_id if available - if trace_info.conversation_id: - attributes["has_conversation"] = "true" - else: - attributes["has_conversation"] = "false" - - self.trace_client.record_trace_duration(duration_s, attributes) - - except Exception: - logger.debug("[Tencent APM] Failed to record workflow trace duration") - - def _record_message_trace_duration(self, trace_info: MessageTraceInfo) -> None: - """Record end-to-end message trace duration.""" - try: - if not hasattr(self.trace_client, "record_trace_duration"): - return - - # Calculate duration from start_time and end_time - if trace_info.start_time and trace_info.end_time: - duration = (trace_info.end_time - trace_info.start_time).total_seconds() - - if duration > 0: - attributes = { - "conversation_mode": trace_info.conversation_mode, - } - - # Add streaming flag if available - if hasattr(trace_info, "is_streaming_request"): - attributes["stream"] = "true" if trace_info.is_streaming_request else "false" - - self.trace_client.record_trace_duration(duration, attributes) - - except Exception: - logger.debug("[Tencent APM] Failed to record message trace duration") - - def close(self) -> None: - """Synchronously and idempotently shutdown the underlying trace client.""" - if getattr(self, "_closed", False): - return - - self._closed = True - trace_client = getattr(self, "trace_client", None) - if trace_client is None: - return - - try: - shutdown_result = trace_client.shutdown() - if inspect.isawaitable(shutdown_result): - close_awaitable = getattr(shutdown_result, "close", None) - if callable(close_awaitable): - close_awaitable() - except Exception: - logger.exception("[Tencent APM] Failed to shutdown trace client during cleanup") - - def __del__(self): - """Ensure best-effort cleanup on garbage collection without retrying shutdown.""" - self.close() diff --git a/api/providers/trace/trace-tencent/src/dify_trace_tencent/utils.py b/api/providers/trace/trace-tencent/src/dify_trace_tencent/utils.py deleted file mode 100644 index 216efd3d8445bf..00000000000000 --- a/api/providers/trace/trace-tencent/src/dify_trace_tencent/utils.py +++ /dev/null @@ -1,64 +0,0 @@ -""" -Utility functions for Tencent APM tracing -""" - -import hashlib -import random -import uuid -from datetime import datetime - -from opentelemetry.trace import Link, SpanContext, TraceFlags - - -class TencentTraceUtils: - """Utility class for common tracing operations.""" - - INVALID_SPAN_ID = 0x0000000000000000 - INVALID_TRACE_ID = 0x00000000000000000000000000000000 - - @staticmethod - def convert_to_trace_id(uuid_v4: str | None) -> int: - try: - uuid_obj = uuid.UUID(uuid_v4) if uuid_v4 else uuid.uuid4() - except Exception as e: - raise ValueError(f"Invalid UUID input: {e}") from e - return uuid_obj.int - - @staticmethod - def convert_to_span_id(uuid_v4: str | None, span_type: str) -> int: - try: - uuid_obj = uuid.UUID(uuid_v4) if uuid_v4 else uuid.uuid4() - except Exception as e: - raise ValueError(f"Invalid UUID input: {e}") from e - combined_key = f"{uuid_obj.hex}-{span_type}" - hash_bytes = hashlib.sha256(combined_key.encode("utf-8")).digest() - return int.from_bytes(hash_bytes[:8], byteorder="big", signed=False) - - @staticmethod - def generate_span_id() -> int: - span_id = random.getrandbits(64) - while span_id == TencentTraceUtils.INVALID_SPAN_ID: - span_id = random.getrandbits(64) - return span_id - - @staticmethod - def convert_datetime_to_nanoseconds(start_time: datetime | None) -> int: - if start_time is None: - start_time = datetime.now() - timestamp_in_seconds = start_time.timestamp() - return int(timestamp_in_seconds * 1e9) - - @staticmethod - def create_link(trace_id_str: str) -> Link: - try: - trace_id = int(trace_id_str, 16) if len(trace_id_str) == 32 else uuid.UUID(trace_id_str).int - except (ValueError, TypeError): - trace_id = uuid.uuid4().int - - span_context = SpanContext( - trace_id=trace_id, - span_id=TencentTraceUtils.INVALID_SPAN_ID, - is_remote=False, - trace_flags=TraceFlags(TraceFlags.SAMPLED), - ) - return Link(span_context) diff --git a/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_client.py b/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_client.py deleted file mode 100644 index cf2cf4a1be01e3..00000000000000 --- a/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_client.py +++ /dev/null @@ -1,616 +0,0 @@ -"""Tests for the TencentTraceClient helpers that drive tracing and metrics.""" - -from __future__ import annotations - -import logging -import sys -import types -from types import SimpleNamespace -from typing import Any, TypedDict, cast -from unittest.mock import MagicMock - -import pytest -from dify_trace_tencent import client as client_module -from dify_trace_tencent.client import TencentTraceClient, _get_opentelemetry_sdk_version -from dify_trace_tencent.entities.tencent_trace_entity import SpanData -from opentelemetry.sdk.trace import Event -from opentelemetry.trace import SpanContext, Status, StatusCode, TraceFlags - -from enums import DeploymentEdition - -metric_reader_instances: list[DummyMetricReader] = [] -meter_provider_instances: list[DummyMeterProvider] = [] - - -class DummyHistogram: - """Placeholder histogram type used by the stubbed metric stack.""" - - -class AggregationTemporality: - DELTA = "delta" - - -class DummyMeter: - def __init__(self) -> None: - self.created: list[tuple[dict[str, object], MagicMock]] = [] - - def create_histogram(self, **kwargs: object) -> MagicMock: - hist = MagicMock(name=f"hist-{kwargs.get('name')}") - self.created.append((kwargs, hist)) - return hist - - -class DummyMeterProvider: - def __init__(self, resource: object, metric_readers: list[object]) -> None: - self.resource = resource - self.metric_readers = metric_readers - self.meter = DummyMeter() - self.shutdown = MagicMock(name="meter_provider_shutdown") - meter_provider_instances.append(self) - - def get_meter(self, name: str, version: str) -> DummyMeter: - return self.meter - - -class DummyMetricReader: - def __init__(self, exporter: object, export_interval_millis: int) -> None: - self.exporter = exporter - self.export_interval_millis = export_interval_millis - self.shutdown = MagicMock(name="metric_reader_shutdown") - metric_reader_instances.append(self) - - -class DummyGrpcMetricExporter: - def __init__(self, **kwargs: object) -> None: - self.kwargs = kwargs - - -class DummyHttpMetricExporter: - def __init__(self, **kwargs: object) -> None: - self.kwargs = kwargs - - -class DummyJsonMetricExporter: - def __init__(self, **kwargs: object) -> None: - self.kwargs = kwargs - - -class DummyJsonMetricExporterNoTemporality: - """Exporter that rejects preferred_temporality to exercise fallback.""" - - def __init__(self, **kwargs: object) -> None: - if "preferred_temporality" in kwargs: - raise RuntimeError("unsupported preferred_temporality") - self.kwargs = kwargs - - -class PatchedCoreComponents(TypedDict): - span_exporter: MagicMock - span_processor: MagicMock - tracer: MagicMock - span: MagicMock - tracer_provider: MagicMock - trace_api: Any - - -def _add_stub_modules(monkeypatch: pytest.MonkeyPatch) -> None: - """Drop fake metric modules into sys.modules so the client imports resolve.""" - - metrics_module = cast(Any, types.ModuleType("opentelemetry.sdk.metrics")) - metrics_module.Histogram = DummyHistogram - metrics_module.MeterProvider = DummyMeterProvider - monkeypatch.setitem(sys.modules, "opentelemetry.sdk.metrics", metrics_module) - - metrics_export_module = cast(Any, types.ModuleType("opentelemetry.sdk.metrics.export")) - metrics_export_module.AggregationTemporality = AggregationTemporality - metrics_export_module.PeriodicExportingMetricReader = DummyMetricReader - monkeypatch.setitem(sys.modules, "opentelemetry.sdk.metrics.export", metrics_export_module) - - grpc_module = cast(Any, types.ModuleType("opentelemetry.exporter.otlp.proto.grpc.metric_exporter")) - grpc_module.OTLPMetricExporter = DummyGrpcMetricExporter - monkeypatch.setitem(sys.modules, "opentelemetry.exporter.otlp.proto.grpc.metric_exporter", grpc_module) - - http_module = cast(Any, types.ModuleType("opentelemetry.exporter.otlp.proto.http.metric_exporter")) - http_module.OTLPMetricExporter = DummyHttpMetricExporter - monkeypatch.setitem(sys.modules, "opentelemetry.exporter.otlp.proto.http.metric_exporter", http_module) - - http_json_module = cast(Any, types.ModuleType("opentelemetry.exporter.otlp.http.json.metric_exporter")) - http_json_module.OTLPMetricExporter = DummyJsonMetricExporter - monkeypatch.setitem(sys.modules, "opentelemetry.exporter.otlp.http.json.metric_exporter", http_json_module) - - legacy_json_module = cast(Any, types.ModuleType("opentelemetry.exporter.otlp.json.metric_exporter")) - legacy_json_module.OTLPMetricExporter = DummyJsonMetricExporter - monkeypatch.setitem(sys.modules, "opentelemetry.exporter.otlp.json.metric_exporter", legacy_json_module) - - -@pytest.fixture(autouse=True) -def stub_metric_modules(monkeypatch: pytest.MonkeyPatch) -> None: - metric_reader_instances.clear() - meter_provider_instances.clear() - _add_stub_modules(monkeypatch) - - -@pytest.fixture(autouse=True) -def patch_core_components(monkeypatch: pytest.MonkeyPatch) -> PatchedCoreComponents: - span_exporter = MagicMock(name="span_exporter") - monkeypatch.setattr(client_module, "OTLPSpanExporter", MagicMock(return_value=span_exporter)) - - span_processor = MagicMock(name="span_processor") - monkeypatch.setattr(client_module, "BatchSpanProcessor", MagicMock(return_value=span_processor)) - - tracer = MagicMock(name="tracer") - span = MagicMock(name="span") - tracer.start_span.return_value = span - - tracer_provider = MagicMock(name="tracer_provider") - tracer_provider.get_tracer.return_value = tracer - tracer_provider.shutdown = MagicMock(name="tracer_provider_shutdown") - monkeypatch.setattr(client_module, "TracerProvider", MagicMock(return_value=tracer_provider)) - - resource = MagicMock(name="resource") - monkeypatch.setattr(client_module, "Resource", MagicMock(return_value=resource)) - - trace_api_stub = SimpleNamespace( - set_span_in_context=MagicMock(name="set_span_in_context", return_value="trace-context"), - NonRecordingSpan=MagicMock(name="non_recording_span", side_effect=lambda ctx: f"non-{ctx}"), - ) - monkeypatch.setattr(client_module, "trace_api", trace_api_stub) - - fake_config = SimpleNamespace( - project=SimpleNamespace(version="test"), - COMMIT_SHA="sha", - DEPLOY_ENV="dev", - DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, - ) - monkeypatch.setattr(client_module, "dify_config", fake_config) - - monkeypatch.setattr(client_module.socket, "gethostname", lambda: "fake-host") - monkeypatch.setenv("OTEL_EXPORTER_OTLP_PROTOCOL", "") - - return { - "span_exporter": span_exporter, - "span_processor": span_processor, - "tracer": tracer, - "span": span, - "tracer_provider": tracer_provider, - "trace_api": trace_api_stub, - } - - -def _make_span_context(trace_id: int = 1, span_id: int = 2) -> SpanContext: - return SpanContext( - trace_id=trace_id, - span_id=span_id, - is_remote=False, - trace_flags=TraceFlags(TraceFlags.SAMPLED), - ) - - -def _build_client() -> TencentTraceClient: - return TencentTraceClient( - service_name="service", - endpoint="https://trace.example.com:4317", - token="token", - ) - - -def test_get_opentelemetry_sdk_version_reads_install(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(client_module, "version", lambda pkg: "2.0.0") - assert _get_opentelemetry_sdk_version() == "2.0.0" - - -def test_get_opentelemetry_sdk_version_falls_back(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(client_module, "version", MagicMock(side_effect=RuntimeError("boom"))) - assert _get_opentelemetry_sdk_version() == "1.27.0" - - -@pytest.mark.parametrize( - ("endpoint", "expected"), - [ - ( - "https://example.com:9090", - ("example.com:9090", False, "example.com", 9090), - ), - ( - "http://localhost", - ("localhost:4317", True, "localhost", 4317), - ), - ( - "example.com:bad", - ("example.com:4317", False, "example.com", 4317), - ), - ], -) -def test_resolve_grpc_target_parsable_variants(endpoint: str, expected: tuple[str, bool, str, int]) -> None: - assert TencentTraceClient._resolve_grpc_target(endpoint) == expected - - -def test_resolve_grpc_target_handles_errors() -> None: - assert TencentTraceClient._resolve_grpc_target(cast(str, 123)) == ("localhost:4317", True, "localhost", 4317) - - -@pytest.mark.parametrize( - ("method", "attr_name", "args"), - [ - ("record_llm_duration", "hist_llm_duration", (0.3, {"foo": object()})), - ("record_token_usage", "hist_token_usage", (5, "input", "chat", "gpt", "gpt", "addr", "dify")), - ("record_time_to_first_token", "hist_time_to_first_token", (0.4, "dify", "gpt")), - ("record_time_to_generate", "hist_time_to_generate", (0.6, "dify", "gpt")), - ("record_trace_duration", "hist_trace_duration", (1.0, {"meta": object()})), - ], -) -def test_record_methods_call_histograms(method: str, attr_name: str, args: tuple[object, ...]) -> None: - client = _build_client() - hist_mock = MagicMock(name=attr_name) - setattr(client, attr_name, hist_mock) - - getattr(client, method)(*args) - hist_mock.record.assert_called_once() - - -def test_record_methods_skip_when_histogram_missing() -> None: - client = _build_client() - client.hist_llm_duration = None - client.record_llm_duration(0.1) - - client.hist_token_usage = None - client.record_token_usage(1, "go", "chat", "model", "model", "addr", "provider") - - client.hist_time_to_first_token = None - client.record_time_to_first_token(0.2, "prov", "model") - - client.hist_time_to_generate = None - client.record_time_to_generate(0.3, "prov", "model") - - client.hist_trace_duration = None - client.record_trace_duration(0.5) - - -def test_record_llm_duration_handles_exceptions(caplog: pytest.LogCaptureFixture) -> None: - client = _build_client() - client.hist_llm_duration = MagicMock(name="hist_llm_duration") - client.hist_llm_duration.record.side_effect = RuntimeError("boom") - - caplog.set_level(logging.DEBUG, logger=client_module.logger.name) - client.record_llm_duration(0.2) - - assert "[Tencent APM] Failed to record LLM duration" in caplog.text - - -def test_create_and_export_span_sets_attributes(patch_core_components: PatchedCoreComponents) -> None: - client = _build_client() - span = patch_core_components["span"] - ctx = _make_span_context(span_id=2) - span.get_span_context.return_value = ctx - - data = SpanData( - trace_id=1, - parent_span_id=None, - span_id=2, - name="span", - attributes={"key": "value"}, - events=[Event(name="evt", attributes={"k": "v"}, timestamp=123)], - status=Status(StatusCode.OK), - start_time=10, - end_time=20, - ) - - client._create_and_export_span(data) - span.set_attributes.assert_called_once() - span.add_event.assert_called_once() - span.set_status.assert_called_once() - span.end.assert_called_once_with(end_time=20) - assert client.span_contexts[2] == ctx - - -def test_create_and_export_span_uses_parent_context(patch_core_components: PatchedCoreComponents) -> None: - client = _build_client() - existing_context = _make_span_context(span_id=10) - client.span_contexts[10] = existing_context - span = patch_core_components["span"] - span.get_span_context.return_value = _make_span_context(span_id=11) - - data = SpanData( - trace_id=1, - parent_span_id=10, - span_id=11, - name="span", - attributes={}, - events=[], - start_time=0, - end_time=1, - ) - - client._create_and_export_span(data) - trace_api = patch_core_components["trace_api"] - trace_api.NonRecordingSpan.assert_called_once_with(existing_context) - trace_api.set_span_in_context.assert_called_once() - - -def test_create_and_export_span_exception_logs_error( - patch_core_components: PatchedCoreComponents, caplog: pytest.LogCaptureFixture -) -> None: - client = _build_client() - span = patch_core_components["span"] - span.get_span_context.return_value = _make_span_context(span_id=2) - # pyrefly: ignore [missing-attribute] - client.tracer.start_span.side_effect = RuntimeError("boom") - - caplog.set_level(logging.DEBUG, logger=client_module.logger.name) - client._create_and_export_span( - SpanData( - trace_id=1, - parent_span_id=None, - span_id=2, - name="span", - attributes={}, - events=[], - start_time=0, - end_time=1, - ) - ) - - error_records = [record for record in caplog.records if record.levelno == logging.ERROR] - assert len(error_records) == 1 - assert error_records[0].getMessage() == "[Tencent APM] Error creating span: span" - - -def test_api_check_connects_successfully(monkeypatch: pytest.MonkeyPatch) -> None: - client = _build_client() - - monkeypatch.setattr( - TencentTraceClient, - "_resolve_grpc_target", - MagicMock(return_value=("host:123", False, "host", 123)), - ) - - socket_mock = MagicMock() - socket_instance = MagicMock() - socket_instance.connect_ex.return_value = 0 - socket_mock.return_value = socket_instance - monkeypatch.setattr(client_module.socket, "socket", socket_mock) - - assert client.api_check() - socket_instance.connect_ex.assert_called_once() - - -def test_api_check_returns_false_and_handles_local(monkeypatch: pytest.MonkeyPatch) -> None: - client = _build_client() - - monkeypatch.setattr( - TencentTraceClient, - "_resolve_grpc_target", - MagicMock(return_value=("host:123", False, "host", 123)), - ) - - socket_mock = MagicMock() - socket_instance = MagicMock() - socket_instance.connect_ex.return_value = 1 - socket_mock.return_value = socket_instance - monkeypatch.setattr(client_module.socket, "socket", socket_mock) - - assert not client.api_check() - - monkeypatch.setattr( - TencentTraceClient, - "_resolve_grpc_target", - MagicMock(return_value=("localhost:4317", True, "localhost", 4317)), - ) - socket_instance.connect_ex.return_value = 1 - assert client.api_check() - - -def test_api_check_handles_exceptions(monkeypatch: pytest.MonkeyPatch) -> None: - client = TencentTraceClient("svc", "https://localhost", "token") - - monkeypatch.setattr(client_module.socket, "socket", MagicMock(side_effect=RuntimeError("boom"))) - assert client.api_check() - - -def test_get_project_url() -> None: - client = _build_client() - assert client.get_project_url() == "https://console.cloud.tencent.com/apm" - - -def test_shutdown_flushes_all_components(patch_core_components: PatchedCoreComponents) -> None: - client = _build_client() - span_processor = patch_core_components["span_processor"] - tracer_provider = patch_core_components["tracer_provider"] - - client.shutdown() - span_processor.force_flush.assert_called_once() - span_processor.shutdown.assert_called_once() - tracer_provider.shutdown.assert_called_once() - - meter_provider = meter_provider_instances[-1] - metric_reader = metric_reader_instances[-1] - meter_provider.shutdown.assert_called_once() - metric_reader.shutdown.assert_called_once() - - -def test_shutdown_logs_when_meter_provider_fails(caplog: pytest.LogCaptureFixture) -> None: - client = _build_client() - meter_provider = meter_provider_instances[-1] - meter_provider.shutdown.side_effect = RuntimeError("boom") - assert client.metric_reader is not None - # pyrefly: ignore [missing-attribute] - client.metric_reader.shutdown.side_effect = RuntimeError("boom") - - caplog.set_level(logging.DEBUG, logger=client_module.logger.name) - client.shutdown() - - assert "[Tencent APM] Error shutting down meter provider" in caplog.text - assert "[Tencent APM] Error shutting down metric reader" in caplog.text - - -def test_metrics_initialization_failure_sets_histogram_attributes(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(DummyMeterProvider, "__init__", MagicMock(side_effect=RuntimeError("err"))) - client = _build_client() - - assert client.meter is None - assert client.meter_provider is None - assert client.hist_llm_duration is None - assert client.hist_token_usage is None - assert client.hist_time_to_first_token is None - assert client.hist_time_to_generate is None - assert client.hist_trace_duration is None - assert client.metric_reader is None - - -def test_add_span_logs_exception(monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture) -> None: - client = _build_client() - monkeypatch.setattr(client, "_create_and_export_span", MagicMock(side_effect=RuntimeError("boom"))) - - caplog.set_level(logging.DEBUG, logger=client_module.logger.name) - client.add_span( - SpanData( - trace_id=1, - parent_span_id=None, - span_id=2, - name="span", - attributes={}, - events=[], - start_time=0, - end_time=1, - ) - ) - - error_records = [record for record in caplog.records if record.levelno == logging.ERROR] - assert len(error_records) == 1 - assert error_records[0].getMessage() == "[Tencent APM] Failed to create span: span" - - -def test_create_and_export_span_converts_attribute_types(patch_core_components: PatchedCoreComponents) -> None: - client = _build_client() - span = patch_core_components["span"] - span.get_span_context.return_value = _make_span_context(span_id=2) - - data = SpanData.model_construct( - trace_id=1, - parent_span_id=None, - span_id=2, - name="span", - attributes={"num": 5, "flag": True, "pi": float("3.14"), "text": "value"}, - events=[], - links=[], - status=Status(StatusCode.OK), - start_time=0, - end_time=1, - ) - - client._create_and_export_span(data) - (attrs,) = span.set_attributes.call_args.args - assert attrs["num"] == 5 - assert attrs["flag"] is True - assert attrs["pi"] == float("3.14") - assert attrs["text"] == "value" - - -def test_record_llm_duration_converts_attributes() -> None: - client = _build_client() - hist_mock = MagicMock(name="hist_llm_duration") - client.hist_llm_duration = hist_mock - - client.record_llm_duration(0.3, cast(dict[str, str], {"foo": object(), "bar": 2})) - _, attrs = hist_mock.record.call_args.args - assert isinstance(attrs["foo"], str) - assert attrs["bar"] == 2 - - -def test_record_trace_duration_converts_attributes() -> None: - client = _build_client() - hist_mock = MagicMock(name="hist_trace_duration") - client.hist_trace_duration = hist_mock - - client.record_trace_duration(1.0, cast(dict[str, str], {"meta": object(), "ok": True})) - _, attrs = hist_mock.record.call_args.args - assert isinstance(attrs["meta"], str) - assert attrs["ok"] is True - - -@pytest.mark.parametrize( - ("method", "attr_name", "args"), - [ - ("record_token_usage", "hist_token_usage", (5, "input", "chat", "gpt", "gpt", "addr", "dify")), - ("record_time_to_first_token", "hist_time_to_first_token", (0.4, "dify", "gpt")), - ("record_time_to_generate", "hist_time_to_generate", (0.6, "dify", "gpt")), - ("record_trace_duration", "hist_trace_duration", (1.0, {"meta": object()})), - ], -) -def test_record_methods_handle_exceptions( - method: str, attr_name: str, args: tuple[object, ...], caplog: pytest.LogCaptureFixture -) -> None: - client = _build_client() - hist_mock = MagicMock(name=attr_name) - hist_mock.record.side_effect = RuntimeError("boom") - setattr(client, attr_name, hist_mock) - - caplog.set_level(logging.DEBUG, logger=client_module.logger.name) - getattr(client, method)(*args) - - assert any( - record.levelno == logging.DEBUG and record.getMessage().startswith("[Tencent APM] Failed to record") - for record in caplog.records - ) - - -def test_metrics_initializes_grpc_metric_exporter() -> None: - client = _build_client() - metric_reader = metric_reader_instances[-1] - exporter = metric_reader.exporter - - assert isinstance(exporter, DummyGrpcMetricExporter) - assert metric_reader.export_interval_millis == client.metrics_export_interval_sec * 1000 - assert exporter.kwargs["endpoint"] == "trace.example.com:4317" - assert exporter.kwargs["insecure"] is False - assert cast(dict[str, dict[str, str]], exporter.kwargs)["headers"]["authorization"] == "Bearer token" - - -def test_metrics_initializes_http_protobuf_metric_exporter(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("OTEL_EXPORTER_OTLP_PROTOCOL", "http/protobuf") - client = _build_client() - metric_reader = metric_reader_instances[-1] - exporter = metric_reader.exporter - - assert isinstance(exporter, DummyHttpMetricExporter) - assert metric_reader.export_interval_millis == client.metrics_export_interval_sec * 1000 - assert exporter.kwargs["endpoint"] == client.endpoint - assert cast(dict[str, dict[str, str]], exporter.kwargs)["headers"]["authorization"] == "Bearer token" - - -def test_metrics_initializes_http_json_metric_exporter(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("OTEL_EXPORTER_OTLP_PROTOCOL", "http/json") - client = _build_client() - metric_reader = metric_reader_instances[-1] - exporter = metric_reader.exporter - - assert isinstance(exporter, DummyJsonMetricExporter) - assert metric_reader.export_interval_millis == client.metrics_export_interval_sec * 1000 - assert exporter.kwargs["endpoint"] == client.endpoint - assert cast(dict[str, dict[str, str]], exporter.kwargs)["headers"]["authorization"] == "Bearer token" - assert "preferred_temporality" in exporter.kwargs - - -def test_metrics_http_json_metric_exporter_falls_back_without_temporality(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("OTEL_EXPORTER_OTLP_PROTOCOL", "http/json") - exporter_module = sys.modules["opentelemetry.exporter.otlp.http.json.metric_exporter"] - monkeypatch.setattr(exporter_module, "OTLPMetricExporter", DummyJsonMetricExporterNoTemporality) - _ = _build_client() - metric_reader = metric_reader_instances[-1] - exporter = metric_reader.exporter - - assert isinstance(exporter, DummyJsonMetricExporterNoTemporality) - assert "preferred_temporality" not in exporter.kwargs - - -def test_metrics_http_json_uses_http_fallback_when_no_json_exporter(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("OTEL_EXPORTER_OTLP_PROTOCOL", "http/json") - - def _fail_import(mod_path: str) -> types.ModuleType: - raise ModuleNotFoundError(mod_path) - - monkeypatch.setattr(client_module.importlib, "import_module", _fail_import) - - _ = _build_client() - metric_reader = metric_reader_instances[-1] - assert isinstance(metric_reader.exporter, DummyHttpMetricExporter) diff --git a/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_config.py b/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_config.py new file mode 100644 index 00000000000000..50344886c874e5 --- /dev/null +++ b/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_config.py @@ -0,0 +1,142 @@ +import traceback + +import pytest +from dify_trace_sdk.config import FormField +from dify_trace_tencent.config import TencentConfig +from pydantic import ValidationError + +TOKEN = "synthetic-token" + + +def _config(**overrides: object) -> TencentConfig: + return TencentConfig.model_validate( + {"endpoint": "https://collector.example:4320", "service_name": "dify-agent", "token": TOKEN, **overrides} + ) + + +def _rendered(error: ValidationError) -> str: + return str(error) + repr(error) + "".join(traceback.format_exception(error)) + + +def test_form_declares_the_token_as_the_only_secret() -> None: + assert TencentConfig.secret_fields() == {"token"} + assert TencentConfig.derived_fields() == set() + assert TencentConfig.form_fields() == ( + FormField(name="token", secret=True, required=True), + FormField(name="endpoint", secret=False, required=True), + FormField(name="service_name", secret=False, required=True), + ) + + +def test_normalization_keeps_the_token_opaque_and_revalidates_unchanged() -> None: + config = _config( + endpoint=" https://COLLECTOR.Example:4320/ ", + service_name=" dify-agent ", + token=" opaque token/+== ", + retired_field="ignored", + ) + assert config.model_dump() == { + "token": " opaque token/+== ", + "endpoint": "https://collector.example:4320", + "service_name": "dify-agent", + } + assert TencentConfig.model_validate(config.model_dump()) == config + + +@pytest.mark.parametrize("service_name", ["a", "a" * 63, "app-1", "a--b"]) +def test_documented_service_name_boundaries_are_accepted(service_name: str) -> None: + assert _config(service_name=service_name).service_name == service_name + + +@pytest.mark.parametrize( + ("service_name", "stored"), [("dify_app", "dify_app"), ("Dify App ", "Dify App"), (" ", "dify_app")] +) +def test_stored_service_names_from_before_verification_stay_usable(service_name: str, stored: str) -> None: + assert _config(service_name=service_name).service_name == stored + + +@pytest.mark.parametrize( + ("endpoint", "expected"), + [ + ("https://COLLECTOR.Example:4320/", "https://collector.example:4320"), + ("http://collector.example:4319", "http://collector.example:4319"), + ("https://collector.example/", "https://collector.example"), + ("http://collector.example", "http://collector.example"), + ("https://collector.example:443", "https://collector.example:443"), + ("http://127.0.0.1:1/", "http://127.0.0.1:1"), + ("https://[2001:DB8::1]:65535/", "https://[2001:db8::1]:65535"), + ("https://BÜCHER.example:4320/", "https://xn--bcher-kva.example:4320"), + ], +) +def test_origins_are_normalized_without_changing_an_explicit_port(endpoint: str, expected: str) -> None: + assert _config(endpoint=endpoint).endpoint == expected + + +@pytest.mark.parametrize( + "endpoint", + [ + "", + " ", + "collector.example:4320", + "ftp://collector.example:4320", + "https://", + "https:///", + "https://user:synthetic-secret@collector.example:4320", + "https://@collector.example:4320", + "https://collector.example:4320/v1/traces", + "https://collector.example:4320//", + "https://collector.example:4320?token=synthetic-secret", + "https://collector.example:4320?", + "https://collector.example:4320#project", + "https://collector.example:4320#", + "https://collector.example\\@other.example:4320", + "https://col lector.example:4320", + "https://collector.example:\n4320", + "https://collector.example:\t4320", + "https://collector\x00.example:4320", + "https://collector.example:0", + "https://collector.example:65536", + "https://collector.example:-1", + "https://collector.example:invalid", + "https://collector.example:", + "https://[2001:db8::1", + "https://[not-an-ipv6-address]:4320", + ], +) +def test_invalid_endpoints_are_rejected_without_echo(endpoint: str) -> None: + with pytest.raises(ValidationError) as caught: + _config(endpoint=endpoint) + rendered = _rendered(caught.value) + assert "Tencent APM requires an HTTP(S) collector origin" in rendered + assert "synthetic-secret" not in rendered + assert TOKEN not in rendered + + +@pytest.mark.parametrize( + "token", ["", " ", "\t\n", "synthetic\ntoken", "synthetic\x00token", "synthetic\x7ftoken", "synthetic-tökén"] +) +def test_tokens_grpc_metadata_cannot_carry_are_rejected_without_echo(token: str) -> None: + with pytest.raises(ValidationError) as caught: + _config(token=token) + rendered = _rendered(caught.value) + assert "Tencent APM requires a nonempty token" in rendered + assert "synthetic" not in rendered + + +def test_masked_tokens_pass_validation() -> None: + assert _config(token="***").token == "***" + + +@pytest.mark.parametrize("field", ["endpoint", "service_name", "token"]) +@pytest.mark.parametrize("value", [None, 42]) +def test_non_string_values_are_rejected(field: str, value: object) -> None: + with pytest.raises(ValidationError): + _config(**{field: value}) + + +@pytest.mark.parametrize("field", ["endpoint", "service_name", "token"]) +def test_required_fields_have_no_defaults(field: str) -> None: + stored = {"endpoint": "https://collector.example:4320", "service_name": "dify-agent", "token": TOKEN} + del stored[field] + with pytest.raises(ValidationError): + TencentConfig.model_validate(stored) diff --git a/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_provider.py b/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_provider.py new file mode 100644 index 00000000000000..c5031d76f5d9b8 --- /dev/null +++ b/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_provider.py @@ -0,0 +1,302 @@ +import traceback +from collections.abc import Iterable, Iterator +from concurrent import futures + +import grpc # type: ignore[import-untyped] +import pytest +from dify_trace_sdk.errors import ( + ExportOutcomeUnknownError, + ExportRejectedError, + ExportRetryableError, + VerificationError, +) +from dify_trace_sdk.provider import SendResult +from dify_trace_sdk.testing import MESSAGE_ID, RecordingEgress, message_trace, tool_fragment, workflow_trace +from dify_trace_tencent.config import TencentConfig +from dify_trace_tencent.provider import TencentProvider +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( + ExportTraceServiceRequest, + ExportTraceServiceResponse, +) +from opentelemetry.proto.collector.trace.v1.trace_service_pb2_grpc import add_TraceServiceServicer_to_server +from opentelemetry.proto.common.v1.common_pb2 import KeyValue +from opentelemetry.proto.resource.v1.resource_pb2 import Resource +from opentelemetry.proto.trace.v1.trace_pb2 import Span as OtlpSpan +from opentelemetry.proto.trace.v1.trace_pb2 import Status + +TOKEN = "synthetic-token" +UNSENDABLE_TOKEN = "synthetic\ntoken" +CONFIG = TencentConfig(token=TOKEN, endpoint="https://apm.example.com", service_name="dify-app") + + +class _Collector: + """An in-process OTLP/gRPC collector; gRPC dispatches to ``Export`` by name.""" + + def __init__(self, target: str) -> None: + self.target = target + self.received: list[tuple[ExportTraceServiceRequest, dict[str, str | bytes]]] = [] + self.response = ExportTraceServiceResponse() + self.abort_with: grpc.StatusCode | None = None + + def Export( # noqa: N802 - gRPC method name + self, request: ExportTraceServiceRequest, context: grpc.ServicerContext + ) -> ExportTraceServiceResponse: + self.received.append((request, dict(context.invocation_metadata()))) + if self.abort_with is not None: + context.abort(self.abort_with, "scripted failure") + return self.response + + +@pytest.fixture +def collector() -> Iterator[_Collector]: + server = grpc.server(futures.ThreadPoolExecutor(max_workers=1)) + collector = _Collector(target=f"127.0.0.1:{server.add_insecure_port('127.0.0.1:0')}") + add_TraceServiceServicer_to_server(collector, server) + server.start() + yield collector + server.stop(None) + + +def _provider(grpc_target: str | None) -> tuple[TencentProvider, RecordingEgress]: + egress = RecordingEgress([], grpc_target=grpc_target) + return TencentProvider(egress), egress + + +def _config(**overrides: object) -> TencentConfig: + return TencentConfig.model_validate( + {"endpoint": "https://collector.example:4320", "service_name": "dify-agent", "token": TOKEN, **overrides} + ) + + +def _attributes(values: Iterable[KeyValue]) -> dict[str, object]: + decoded: dict[str, object] = {} + for item in values: + match item.value.WhichOneof("value"): + case "string_value": + decoded[item.key] = item.value.string_value + case "int_value": + decoded[item.key] = item.value.int_value + case "bool_value": + decoded[item.key] = item.value.bool_value + case "double_value": + decoded[item.key] = item.value.double_value + case other: + raise AssertionError(f"unexpected attribute type: {other}") + return decoded + + +def _exported(collector: _Collector) -> tuple[Resource, list[OtlpSpan]]: + ((request, metadata),) = collector.received + assert metadata["authorization"] == f"Bearer {TOKEN}" + (resource_spans,) = request.resource_spans + (scope_spans,) = resource_spans.scope_spans + return resource_spans.resource, list(scope_spans.spans) + + +def _rendered(error: Exception) -> str: + return str(error) + repr(error) + "".join(traceback.format_exception(error)) + + +def test_verify_exchanges_one_empty_export_with_the_bearer_token(collector: _Collector) -> None: + provider, egress = _provider(collector.target) + assert provider.verify(CONFIG) == CONFIG + assert egress.grpc_endpoints == ["https://apm.example.com:4317"] + assert egress.requests == [] + ((request, metadata),) = collector.received + assert request == ExportTraceServiceRequest() + assert metadata["authorization"] == f"Bearer {TOKEN}" + + +@pytest.mark.parametrize( + ("endpoint", "dialed"), + [ + ("https://apm.example.com", "https://apm.example.com:4317"), + ("http://apm.example.com", "http://apm.example.com:4317"), + ("http://[2001:db8::1]", "http://[2001:db8::1]:4317"), + ("https://apm.example.com:4320", "https://apm.example.com:4320"), + ("https://apm.example.com:443", "https://apm.example.com:443"), + ], +) +def test_the_otlp_grpc_port_is_added_only_when_the_endpoint_names_none( + collector: _Collector, endpoint: str, dialed: str +) -> None: + provider, egress = _provider(collector.target) + config = TencentConfig(token=TOKEN, endpoint=endpoint, service_name="dify-app") + provider.verify(config) + provider.send(config, message_trace()) + assert egress.grpc_endpoints == [dialed, dialed] + + +@pytest.mark.parametrize( + "status", [grpc.StatusCode.UNAUTHENTICATED, grpc.StatusCode.UNAVAILABLE, grpc.StatusCode.UNIMPLEMENTED] +) +def test_verify_rpc_failures_are_connection_failures(collector: _Collector, status: grpc.StatusCode) -> None: + collector.abort_with = status + provider, _ = _provider(collector.target) + with pytest.raises(VerificationError) as caught: + provider.verify(CONFIG) + assert caught.value.code == "connection_failed" + assert TOKEN not in _rendered(caught.value) + + +@pytest.mark.parametrize(("rejected_spans", "error_message"), [(1, ""), (0, f"private-collector-detail {TOKEN}")]) +def test_verify_partial_acknowledgments_are_unexpected( + collector: _Collector, rejected_spans: int, error_message: str +) -> None: + collector.response.partial_success.rejected_spans = rejected_spans + collector.response.partial_success.error_message = error_message + provider, _ = _provider(collector.target) + with pytest.raises(VerificationError) as caught: + provider.verify(CONFIG) + assert caught.value.code == "unexpected_response" + rendered = _rendered(caught.value) + assert "private-collector-detail" not in rendered + assert TOKEN not in rendered + + +def test_metadata_grpc_cannot_send_never_leaks_the_token(collector: _Collector) -> None: + # Validation rejects such tokens; this covers a configuration that bypassed it. + unsendable = TencentConfig.model_construct( + token=UNSENDABLE_TOKEN, endpoint="https://apm.example.com", service_name="dify-app" + ) + provider, _ = _provider(collector.target) + with pytest.raises(VerificationError) as verify_error: + provider.verify(unsendable) + assert verify_error.value.code == "connection_failed" + with pytest.raises(ExportRejectedError) as send_error: + provider.send(unsendable, message_trace()) + for error in (verify_error.value, send_error.value): + assert "synthetic" not in _rendered(error) + assert collector.received == [] + + +def test_verify_endpoint_refused_by_the_egress_is_a_connection_failure() -> None: + provider, egress = _provider(None) + with pytest.raises(VerificationError) as caught: + provider.verify(CONFIG) + assert caught.value.code == "connection_failed" + assert egress.grpc_endpoints == ["https://apm.example.com:4317"] + + +@pytest.mark.parametrize( + "service_name", + ["", " ", "dify_app", "Dify", "1app", "-app", "app-", "a" * 64, "应用", "a.b", "a/b", "a b", "a\nb"], +) +def test_verification_rejects_service_names_outside_the_documented_pattern(service_name: str) -> None: + egress = RecordingEgress([], grpc_target=None) + with pytest.raises(VerificationError) as raised: + TencentProvider(egress).verify(_config(service_name=service_name)) + assert raised.value.code == "invalid_service_name" + assert egress.grpc_endpoints == [] + + +def test_project_url_is_the_apm_console() -> None: + provider, _ = _provider(None) + assert provider.project_url(CONFIG) == "https://console.cloud.tencent.com/apm" + + +def test_message_trace_is_exported_with_tencent_genai_attributes(collector: _Collector) -> None: + provider, _ = _provider(collector.target) + trace = message_trace() + assert provider.send(CONFIG, trace) == SendResult(trace_id=trace.trace_id, trace_url=None, span_urls={}) + + resource, (root, llm) = _exported(collector) + assert _attributes(resource.attributes) == {"service.name": "dify-app"} + assert root.trace_id.hex() == llm.trace_id.hex() == trace.trace_id == MESSAGE_ID.replace("-", "") + assert (root.span_id.hex(), root.parent_span_id) == (trace.spans[0].span_id, b"") + assert (llm.span_id.hex(), llm.parent_span_id.hex()) == (trace.spans[1].span_id, trace.spans[0].span_id) + assert (root.name, llm.name) == ("message", "llm") + assert root.status.code == llm.status.code == Status.STATUS_CODE_OK + root_attributes = _attributes(root.attributes) + assert ( + root_attributes.items() + >= { + "gen_ai.span.kind": "WORKFLOW", + "gen_ai.is_entry": "true", + "gen_ai.framework": "dify", + "gen_ai.session.id": "conversation-1", + "gen_ai.user.id": "end-user-1", + "gen_ai.entity.input": '{"query": "What is Dify?"}', + "gen_ai.entity.output": '"An LLM application platform."', + "dify.app_id": "app-1", + "dify.message_id": MESSAGE_ID, + }.items() + ) + assert "input.value" not in root_attributes + llm_attributes = _attributes(llm.attributes) + assert ( + llm_attributes.items() + >= { + "gen_ai.span.kind": "GENERATION", + "gen_ai.operation.name": "chat", + "gen_ai.request.model": "gpt-4o", + "gen_ai.provider.name": "openai", + "gen_ai.usage.input_tokens": 12, + "gen_ai.usage.output_tokens": 8, + "gen_ai.usage.total_tokens": 20, + "gen_ai.response.time_to_first_token": 250_000_000, + }.items() + ) + assert "gen_ai.is_entry" not in llm_attributes + + +def test_workflow_trace_marks_only_its_root_as_the_entry(collector: _Collector) -> None: + provider, _ = _provider(collector.target) + trace = workflow_trace() + provider.send(CONFIG, trace) + + _, (run, llm, tool) = _exported(collector) + assert (llm.parent_span_id.hex(), tool.parent_span_id.hex()) == (run.span_id.hex(), llm.span_id.hex()) + attributes = [_attributes(span.attributes) for span in (run, llm, tool)] + assert [values["gen_ai.span.kind"] for values in attributes] == ["WORKFLOW", "GENERATION", "TOOL"] + assert [values.get("gen_ai.is_entry") for values in attributes] == ["true", None, None] + assert (run.status.code, run.status.message) == (Status.STATUS_CODE_ERROR, "Tool failed") + assert (tool.status.code, tool.status.message) == (Status.STATUS_CODE_ERROR, "Tool failed") + + +def test_fragment_joins_the_message_without_claiming_the_entry(collector: _Collector) -> None: + provider, _ = _provider(collector.target) + fragment = tool_fragment() + provider.send(CONFIG, fragment) + + _, (tool,) = _exported(collector) + assert fragment.parent is not None + assert tool.trace_id.hex() == message_trace().trace_id + assert tool.parent_span_id.hex() == fragment.parent.span_id == message_trace().root.span_id + tool_attributes = _attributes(tool.attributes) + assert tool_attributes["gen_ai.span.kind"] == "TOOL" + assert "gen_ai.is_entry" not in tool_attributes + + +@pytest.mark.parametrize( + ("status", "expected"), + [ + (grpc.StatusCode.UNAVAILABLE, ExportRetryableError), + (grpc.StatusCode.RESOURCE_EXHAUSTED, ExportRetryableError), + (grpc.StatusCode.DEADLINE_EXCEEDED, ExportOutcomeUnknownError), + (grpc.StatusCode.UNAUTHENTICATED, ExportRejectedError), + (grpc.StatusCode.INVALID_ARGUMENT, ExportRejectedError), + ], +) +def test_send_rpc_failures_are_classified_without_exposing_the_token( + collector: _Collector, status: grpc.StatusCode, expected: type[Exception] +) -> None: + collector.abort_with = status + provider, _ = _provider(collector.target) + with pytest.raises(expected) as caught: + provider.send(CONFIG, message_trace()) + assert TOKEN not in _rendered(caught.value) + + +def test_partial_acceptance_is_an_unknown_outcome(collector: _Collector) -> None: + collector.response.partial_success.rejected_spans = 1 + provider, _ = _provider(collector.target) + with pytest.raises(ExportOutcomeUnknownError): + provider.send(CONFIG, message_trace()) + + +def test_endpoint_refused_by_the_egress_is_rejected() -> None: + provider, egress = _provider(None) + with pytest.raises(ExportRejectedError): + provider.send(CONFIG, message_trace()) + assert egress.grpc_endpoints == ["https://apm.example.com:4317"] diff --git a/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_span_builder.py b/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_span_builder.py deleted file mode 100644 index e1be98258c3706..00000000000000 --- a/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_span_builder.py +++ /dev/null @@ -1,361 +0,0 @@ -from datetime import datetime -from unittest.mock import patch - -from dify_trace_tencent.entities.semconv import ( - GEN_AI_IS_ENTRY, - GEN_AI_IS_STREAMING_REQUEST, - GEN_AI_MODEL_NAME, - GEN_AI_SPAN_KIND, - GEN_AI_USAGE_INPUT_TOKENS, - INPUT_VALUE, - RETRIEVAL_DOCUMENT, - RETRIEVAL_QUERY, - TOOL_DESCRIPTION, - TOOL_NAME, - TOOL_PARAMETERS, - GenAISpanKind, -) -from dify_trace_tencent.span_builder import TencentSpanBuilder -from opentelemetry.trace import StatusCode - -from core.ops.entities.trace_entity import ( - DatasetRetrievalTraceInfo, -) -from core.rag.models.document import Document -from graphon.enums import WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus -from tests.helpers.unit.core.ops.trace_fixtures import ( - message_trace_info, - tool_trace_info, - workflow_node_execution, - workflow_trace_info, -) - - -class TestTencentSpanBuilder: - def test_get_time_nanoseconds(self): - with patch("dify_trace_tencent.utils.TencentTraceUtils.convert_datetime_to_nanoseconds") as mock_convert: - mock_convert.return_value = 123456789 - dt = datetime.now() - result = TencentSpanBuilder._get_time_nanoseconds(dt) - assert result == 123456789 - mock_convert.assert_called_once_with(dt) - - def test_build_workflow_spans(self): - trace_info = workflow_trace_info() - trace_info.workflow_run_id = "run_id" - trace_info.error = None - trace_info.start_time = datetime.now() - trace_info.end_time = datetime.now() - trace_info.workflow_run_inputs = {"sys.query": "hello"} - trace_info.workflow_run_outputs = {"answer": "world"} - trace_info.metadata = {"conversation_id": "conv_id"} - - with patch("dify_trace_tencent.utils.TencentTraceUtils.convert_to_span_id") as mock_convert_id: - mock_convert_id.side_effect = [1, 2] # workflow_span_id, message_span_id - with patch.object(TencentSpanBuilder, "_get_time_nanoseconds", return_value=100): - spans = TencentSpanBuilder.build_workflow_spans(trace_info, 123, "user_1") - - assert len(spans) == 2 - assert spans[0].name == "message" - assert spans[0].span_id == 2 - assert spans[1].name == "workflow" - assert spans[1].span_id == 1 - assert spans[1].parent_span_id == 2 - - def test_build_workflow_spans_no_message(self): - trace_info = workflow_trace_info() - trace_info.workflow_run_id = "run_id" - trace_info.error = "some error" - trace_info.start_time = datetime.now() - trace_info.end_time = datetime.now() - trace_info.workflow_run_inputs = {} - trace_info.workflow_run_outputs = {} - trace_info.metadata = {} # No conversation_id - - with patch("dify_trace_tencent.utils.TencentTraceUtils.convert_to_span_id") as mock_convert_id: - mock_convert_id.return_value = 1 - with patch.object(TencentSpanBuilder, "_get_time_nanoseconds", return_value=100): - spans = TencentSpanBuilder.build_workflow_spans(trace_info, 123, "user_1") - - assert len(spans) == 1 - assert spans[0].name == "workflow" - assert spans[0].status.status_code == StatusCode.ERROR - assert spans[0].status.description == "some error" - assert spans[0].attributes[GEN_AI_IS_ENTRY] == "true" - - def test_build_workflow_llm_span(self): - trace_info = workflow_trace_info() - trace_info.metadata = {"conversation_id": "conv_id"} - - node_execution = workflow_node_execution() - node_execution.id = "node_id" - node_execution.created_at = datetime.now() - node_execution.finished_at = datetime.now() - node_execution.status = WorkflowNodeExecutionStatus.SUCCEEDED - node_execution.process_data = { - "model_name": "gpt-4", - "model_provider": "openai", - "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30, "time_to_first_token": 0.5}, - "prompts": ["hello"], - } - node_execution.outputs = {"text": "world"} - - with patch("dify_trace_tencent.utils.TencentTraceUtils.convert_to_span_id") as mock_convert_id: - mock_convert_id.return_value = 456 - with patch.object(TencentSpanBuilder, "_get_time_nanoseconds", return_value=100): - span = TencentSpanBuilder.build_workflow_llm_span(123, 1, trace_info, node_execution) - - assert span.name == "GENERATION" - assert span.attributes[GEN_AI_MODEL_NAME] == "gpt-4" - assert span.attributes[GEN_AI_IS_STREAMING_REQUEST] == "true" - assert span.attributes[GEN_AI_USAGE_INPUT_TOKENS] == "10" - - def test_build_workflow_llm_span_usage_in_outputs(self): - trace_info = workflow_trace_info() - trace_info.metadata = {} - - node_execution = workflow_node_execution() - node_execution.id = "node_id" - node_execution.created_at = datetime.now() - node_execution.finished_at = datetime.now() - node_execution.status = WorkflowNodeExecutionStatus.SUCCEEDED - node_execution.process_data = {} - node_execution.outputs = { - "text": "world", - "usage": {"prompt_tokens": 15, "completion_tokens": 25, "total_tokens": 40}, - } - - with patch("dify_trace_tencent.utils.TencentTraceUtils.convert_to_span_id") as mock_convert_id: - mock_convert_id.return_value = 456 - with patch.object(TencentSpanBuilder, "_get_time_nanoseconds", return_value=100): - span = TencentSpanBuilder.build_workflow_llm_span(123, 1, trace_info, node_execution) - - assert span.attributes[GEN_AI_USAGE_INPUT_TOKENS] == "15" - assert GEN_AI_IS_STREAMING_REQUEST not in span.attributes - - def test_build_message_span_standalone(self): - trace_info = message_trace_info() - trace_info.message_id = "msg_id" - trace_info.error = None - trace_info.start_time = datetime.now() - trace_info.end_time = datetime.now() - trace_info.inputs = {"q": "hi"} - trace_info.outputs = "hello" - trace_info.metadata = {"conversation_id": "conv_id"} - trace_info.is_streaming_request = True - - with patch("dify_trace_tencent.utils.TencentTraceUtils.convert_to_span_id") as mock_convert_id: - mock_convert_id.return_value = 789 - with patch.object(TencentSpanBuilder, "_get_time_nanoseconds", return_value=100): - span = TencentSpanBuilder.build_message_span(trace_info, 123, "user_1") - - assert span.name == "message" - assert span.attributes[GEN_AI_IS_STREAMING_REQUEST] == "true" - assert span.attributes[INPUT_VALUE] == str(trace_info.inputs) - - def test_build_message_span_standalone_with_error(self): - trace_info = message_trace_info() - trace_info.message_id = "msg_id" - trace_info.error = "some error" - trace_info.start_time = datetime.now() - trace_info.end_time = datetime.now() - trace_info.inputs = None - trace_info.outputs = None - trace_info.metadata = {} - trace_info.is_streaming_request = False - - with patch("dify_trace_tencent.utils.TencentTraceUtils.convert_to_span_id") as mock_convert_id: - mock_convert_id.return_value = 789 - with patch.object(TencentSpanBuilder, "_get_time_nanoseconds", return_value=100): - span = TencentSpanBuilder.build_message_span(trace_info, 123, "user_1") - - assert span.status.status_code == StatusCode.ERROR - assert span.status.description == "some error" - assert span.attributes[INPUT_VALUE] == "" - - def test_build_tool_span(self): - trace_info = tool_trace_info() - trace_info.message_id = "msg_id" - trace_info.tool_name = "search" - trace_info.error = "tool error" - trace_info.start_time = datetime.now() - trace_info.end_time = datetime.now() - trace_info.tool_parameters = {"p": 1} - trace_info.tool_inputs = {"i": 2} - trace_info.tool_outputs = "result" - - with patch("dify_trace_tencent.utils.TencentTraceUtils.convert_to_span_id") as mock_convert_id: - mock_convert_id.return_value = 101 - with patch.object(TencentSpanBuilder, "_get_time_nanoseconds", return_value=100): - span = TencentSpanBuilder.build_tool_span(trace_info, 123, 1) - - assert span.name == "search" - assert span.status.status_code == StatusCode.ERROR - assert span.attributes[TOOL_NAME] == "search" - - def test_build_retrieval_span(self): - trace_info = DatasetRetrievalTraceInfo(metadata={}) - trace_info.message_id = "msg_id" - trace_info.inputs = "query" - trace_info.error = None - trace_info.start_time = datetime.now() - trace_info.end_time = datetime.now() - - doc = Document( - page_content="content", metadata={"dataset_id": "d1", "doc_id": "di1", "document_id": "du1", "score": 0.9} - ) - trace_info.documents = [doc] - - with patch("dify_trace_tencent.utils.TencentTraceUtils.convert_to_span_id") as mock_convert_id: - mock_convert_id.return_value = 202 - with patch.object(TencentSpanBuilder, "_get_time_nanoseconds", return_value=100): - span = TencentSpanBuilder.build_retrieval_span(trace_info, 123, 1) - - assert span.name == "retrieval" - assert span.attributes[RETRIEVAL_QUERY] == "query" - assert "content" in span.attributes[RETRIEVAL_DOCUMENT] - - def test_build_retrieval_span_with_error(self): - trace_info = DatasetRetrievalTraceInfo(metadata={}) - trace_info.message_id = "msg_id" - trace_info.inputs = "" - trace_info.error = "retrieval failed" - trace_info.start_time = datetime.now() - trace_info.end_time = datetime.now() - trace_info.documents = [] - - with patch("dify_trace_tencent.utils.TencentTraceUtils.convert_to_span_id") as mock_convert_id: - mock_convert_id.return_value = 202 - with patch.object(TencentSpanBuilder, "_get_time_nanoseconds", return_value=100): - span = TencentSpanBuilder.build_retrieval_span(trace_info, 123, 1) - - assert span.status.status_code == StatusCode.ERROR - assert span.status.description == "retrieval failed" - - def test_get_workflow_node_status(self): - node = workflow_node_execution() - - node.status = WorkflowNodeExecutionStatus.SUCCEEDED - assert TencentSpanBuilder._get_workflow_node_status(node).status_code == StatusCode.OK - - node.status = WorkflowNodeExecutionStatus.FAILED - node.error = "fail" - status = TencentSpanBuilder._get_workflow_node_status(node) - assert status.status_code == StatusCode.ERROR - assert status.description == "fail" - - node.status = WorkflowNodeExecutionStatus.EXCEPTION - node.error = "exc" - status = TencentSpanBuilder._get_workflow_node_status(node) - assert status.status_code == StatusCode.ERROR - assert status.description == "exc" - - node.status = WorkflowNodeExecutionStatus.RUNNING - assert TencentSpanBuilder._get_workflow_node_status(node).status_code == StatusCode.UNSET - - def test_build_workflow_retrieval_span(self): - trace_info = workflow_trace_info() - trace_info.metadata = {"conversation_id": "conv_id"} - - node_execution = workflow_node_execution() - node_execution.id = "node_id" - node_execution.title = "my retrieval" - node_execution.status = WorkflowNodeExecutionStatus.SUCCEEDED - node_execution.inputs = {"query": "q1"} - node_execution.outputs = {"result": [{"content": "c1"}]} - node_execution.created_at = datetime.now() - node_execution.finished_at = datetime.now() - - with patch("dify_trace_tencent.utils.TencentTraceUtils.convert_to_span_id") as mock_convert_id: - mock_convert_id.return_value = 303 - with patch.object(TencentSpanBuilder, "_get_time_nanoseconds", return_value=100): - span = TencentSpanBuilder.build_workflow_retrieval_span(123, 1, trace_info, node_execution) - - assert span.name == "my retrieval" - assert span.attributes[RETRIEVAL_QUERY] == "q1" - assert "c1" in span.attributes[RETRIEVAL_DOCUMENT] - - def test_build_workflow_retrieval_span_empty(self): - trace_info = workflow_trace_info() - trace_info.metadata = {} - - node_execution = workflow_node_execution() - node_execution.id = "node_id" - node_execution.title = "my retrieval" - node_execution.status = WorkflowNodeExecutionStatus.SUCCEEDED - node_execution.inputs = {} - node_execution.outputs = {} - node_execution.created_at = datetime.now() - node_execution.finished_at = datetime.now() - - with patch("dify_trace_tencent.utils.TencentTraceUtils.convert_to_span_id") as mock_convert_id: - mock_convert_id.return_value = 303 - with patch.object(TencentSpanBuilder, "_get_time_nanoseconds", return_value=100): - span = TencentSpanBuilder.build_workflow_retrieval_span(123, 1, trace_info, node_execution) - - assert span.attributes[RETRIEVAL_QUERY] == "" - assert span.attributes[RETRIEVAL_DOCUMENT] == "" - - def test_build_workflow_tool_span(self): - trace_info = workflow_trace_info() - - node_execution = workflow_node_execution() - node_execution.id = "node_id" - node_execution.title = "my tool" - node_execution.status = WorkflowNodeExecutionStatus.SUCCEEDED - node_execution.metadata = {WorkflowNodeExecutionMetadataKey.TOOL_INFO: {"info": "some"}} - node_execution.inputs = {"param": "val"} - node_execution.outputs = {"res": "ok"} - node_execution.created_at = datetime.now() - node_execution.finished_at = datetime.now() - - with patch("dify_trace_tencent.utils.TencentTraceUtils.convert_to_span_id") as mock_convert_id: - mock_convert_id.return_value = 404 - with patch.object(TencentSpanBuilder, "_get_time_nanoseconds", return_value=100): - span = TencentSpanBuilder.build_workflow_tool_span(123, 1, trace_info, node_execution) - - assert span.name == "my tool" - assert span.attributes[TOOL_NAME] == "my tool" - assert "some" in span.attributes[TOOL_DESCRIPTION] - - def test_build_workflow_tool_span_no_metadata(self): - trace_info = workflow_trace_info() - - node_execution = workflow_node_execution() - node_execution.id = "node_id" - node_execution.title = "my tool" - node_execution.status = WorkflowNodeExecutionStatus.SUCCEEDED - node_execution.metadata = None - node_execution.inputs = None - node_execution.outputs = {"res": "ok"} - node_execution.created_at = datetime.now() - node_execution.finished_at = datetime.now() - - with patch("dify_trace_tencent.utils.TencentTraceUtils.convert_to_span_id") as mock_convert_id: - mock_convert_id.return_value = 404 - with patch.object(TencentSpanBuilder, "_get_time_nanoseconds", return_value=100): - span = TencentSpanBuilder.build_workflow_tool_span(123, 1, trace_info, node_execution) - - assert span.attributes[TOOL_DESCRIPTION] == "{}" - assert span.attributes[TOOL_PARAMETERS] == "{}" - - def test_build_workflow_task_span(self): - trace_info = workflow_trace_info() - trace_info.metadata = {"conversation_id": "conv_id"} - - node_execution = workflow_node_execution() - node_execution.id = "node_id" - node_execution.title = "my task" - node_execution.status = WorkflowNodeExecutionStatus.SUCCEEDED - node_execution.inputs = {"in": 1} - node_execution.outputs = {"out": 2} - node_execution.created_at = datetime.now() - node_execution.finished_at = datetime.now() - - with patch("dify_trace_tencent.utils.TencentTraceUtils.convert_to_span_id") as mock_convert_id: - mock_convert_id.return_value = 505 - with patch.object(TencentSpanBuilder, "_get_time_nanoseconds", return_value=100): - span = TencentSpanBuilder.build_workflow_task_span(123, 1, trace_info, node_execution) - - assert span.name == "my task" - assert span.attributes[GEN_AI_SPAN_KIND] == GenAISpanKind.TASK.value diff --git a/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_tencent_trace.py b/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_tencent_trace.py deleted file mode 100644 index 5f53ecba4c8206..00000000000000 --- a/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_tencent_trace.py +++ /dev/null @@ -1,754 +0,0 @@ -"""Unit tests for Tencent tracing, including SQLite-backed account resolution.""" - -import gc -import logging -import warnings -from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest -from dify_trace_tencent.config import TencentConfig -from dify_trace_tencent.tencent_trace import TencentDataTrace -from sqlalchemy.orm import Session - -from core.ops.entities.trace_entity import ( - DatasetRetrievalTraceInfo, - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - SuggestedQuestionTraceInfo, - ToolTraceInfo, - WorkflowTraceInfo, -) -from graphon.enums import BuiltinNodeTypes -from models import Account, App, Tenant, TenantAccountJoin -from models.account import TenantAccountRole -from models.model import AppMode, IconType -from tests.helpers.unit.core.ops.trace_fixtures import ( - message_trace_info, - tool_trace_info, - workflow_node_execution, - workflow_trace_info, -) - -logger = logging.getLogger(__name__) - - -@pytest.fixture -def tencent_config(): - return TencentConfig(service_name="test-service", endpoint="https://test-endpoint", token="test-token") - - -@pytest.fixture -def mock_trace_client(): - with patch("dify_trace_tencent.tencent_trace.TencentTraceClient") as mock: - yield mock - - -@pytest.fixture -def mock_span_builder(): - with patch("dify_trace_tencent.tencent_trace.TencentSpanBuilder") as mock: - yield mock - - -@pytest.fixture -def mock_trace_utils(): - with patch("dify_trace_tencent.tencent_trace.TencentTraceUtils") as mock: - yield mock - - -@pytest.fixture -def tencent_data_trace(tencent_config, mock_trace_client): - return TencentDataTrace(tencent_config) - - -class TestTencentDataTrace: - def test_init(self, tencent_config, mock_trace_client): - trace = TencentDataTrace(tencent_config) - mock_trace_client.assert_called_once_with( - service_name=tencent_config.service_name, - endpoint=tencent_config.endpoint, - token=tencent_config.token, - metrics_export_interval_sec=5, - ) - assert trace.trace_client == mock_trace_client.return_value - - def test_trace_dispatch(self, tencent_data_trace): - methods = [ - ( - WorkflowTraceInfo( - workflow_id="wf", - tenant_id="t", - workflow_run_id="run", - workflow_run_elapsed_time=1.0, - workflow_run_status="s", - workflow_run_inputs={}, - workflow_run_outputs={}, - workflow_run_version="v", - total_tokens=0, - file_list=[], - query="", - metadata={}, - ), - "workflow_trace", - ), - ( - MessageTraceInfo( - message_id="msg", - message_data={}, - inputs={}, - outputs={}, - start_time=None, - end_time=None, - conversation_mode="chat", - conversation_model="gpt-3.5-turbo", - message_tokens=0, - answer_tokens=0, - total_tokens=0, - metadata={}, - ), - "message_trace", - ), - ( - ModerationTraceInfo( - flagged=False, action="a", preset_response="p", query="q", metadata={}, message_id="m" - ), - None, - ), # Pass - ( - SuggestedQuestionTraceInfo( - suggested_question=[], - level="l", - total_tokens=0, - metadata={}, - message_id="m", - message_data={}, - inputs={}, - start_time=None, - end_time=None, - ), - "suggested_question_trace", - ), - ( - DatasetRetrievalTraceInfo( - metadata={}, - message_id="m", - message_data={}, - inputs={}, - documents=[], - start_time=None, - end_time=None, - ), - "dataset_retrieval_trace", - ), - ( - ToolTraceInfo( - tool_name="t", - tool_inputs={}, - tool_outputs="", - tool_config={}, - tool_parameters={}, - time_cost=0, - metadata={}, - message_id="m", - inputs={}, - outputs={}, - start_time=None, - end_time=None, - ), - "tool_trace", - ), - ( - GenerateNameTraceInfo( - tenant_id="t", metadata={}, message_id="m", inputs={}, outputs={}, start_time=None, end_time=None - ), - None, - ), # Pass - ] - - for trace_info, method_name in methods: - if method_name: - with patch.object(tencent_data_trace, method_name) as mock_method: - tencent_data_trace.trace(trace_info) - mock_method.assert_called_once_with(trace_info) - else: - tencent_data_trace.trace(trace_info) - - def test_api_check(self, tencent_data_trace): - tencent_data_trace.trace_client.api_check.return_value = True - assert tencent_data_trace.api_check() is True - tencent_data_trace.trace_client.api_check.assert_called_once() - - def test_get_project_url(self, tencent_data_trace): - tencent_data_trace.trace_client.get_project_url.return_value = "http://url" - assert tencent_data_trace.get_project_url() == "http://url" - tencent_data_trace.trace_client.get_project_url.assert_called_once() - - def test_workflow_trace(self, tencent_data_trace, mock_trace_utils, mock_span_builder): - trace_info = workflow_trace_info() - trace_info.workflow_run_id = "run-id" - trace_info.trace_id = "parent-trace-id" - - mock_trace_utils.convert_to_trace_id.return_value = 123 - mock_trace_utils.create_link.return_value = "link" - - with ( - patch.object(tencent_data_trace, "_get_user_id", return_value="user-1"), - patch.object(tencent_data_trace, "_process_workflow_nodes") as mock_proc, - patch.object(tencent_data_trace, "_record_workflow_trace_duration") as mock_dur, - ): - mock_span_builder.build_workflow_spans.return_value = [MagicMock(), MagicMock()] - - tencent_data_trace.workflow_trace(trace_info) - - mock_trace_utils.convert_to_trace_id.assert_called_once_with("run-id") - mock_trace_utils.create_link.assert_called_once_with("parent-trace-id") - mock_span_builder.build_workflow_spans.assert_called_once() - assert tencent_data_trace.trace_client.add_span.call_count == 2 - mock_proc.assert_called_once_with(trace_info, 123) - mock_dur.assert_called_once_with(trace_info) - - def test_workflow_trace_exception(self, tencent_data_trace, caplog: pytest.LogCaptureFixture): - trace_info = workflow_trace_info() - trace_info.workflow_run_id = "run-id" - - with ( - patch( - "dify_trace_tencent.tencent_trace.TencentTraceUtils.convert_to_trace_id", side_effect=Exception("error") - ), - caplog.at_level(logging.ERROR), - ): - tencent_data_trace.workflow_trace(trace_info) - assert "[Tencent APM] Failed to process workflow trace" in caplog.text - - def test_message_trace(self, tencent_data_trace, mock_trace_utils, mock_span_builder): - trace_info = message_trace_info() - trace_info.message_id = "msg-id" - trace_info.trace_id = "parent-trace-id" - - mock_trace_utils.convert_to_trace_id.return_value = 123 - mock_trace_utils.create_link.return_value = "link" - - with ( - patch.object(tencent_data_trace, "_get_user_id", return_value="user-1"), - patch.object(tencent_data_trace, "_record_message_llm_metrics") as mock_metrics, - patch.object(tencent_data_trace, "_record_message_trace_duration") as mock_dur, - ): - mock_span_builder.build_message_span.return_value = MagicMock() - - tencent_data_trace.message_trace(trace_info) - - mock_trace_utils.convert_to_trace_id.assert_called_once_with("msg-id") - mock_trace_utils.create_link.assert_called_once_with("parent-trace-id") - mock_span_builder.build_message_span.assert_called_once() - tencent_data_trace.trace_client.add_span.assert_called_once() - mock_metrics.assert_called_once_with(trace_info) - mock_dur.assert_called_once_with(trace_info) - - def test_message_trace_exception(self, tencent_data_trace, caplog: pytest.LogCaptureFixture): - trace_info = message_trace_info() - - with ( - patch( - "dify_trace_tencent.tencent_trace.TencentTraceUtils.convert_to_trace_id", side_effect=Exception("error") - ), - caplog.at_level(logging.ERROR), - ): - tencent_data_trace.message_trace(trace_info) - assert "[Tencent APM] Failed to process message trace" in caplog.text - - def test_tool_trace(self, tencent_data_trace, mock_trace_utils, mock_span_builder): - trace_info = tool_trace_info() - trace_info.message_id = "msg-id" - - mock_trace_utils.convert_to_span_id.return_value = 456 - mock_trace_utils.convert_to_trace_id.return_value = 123 - - tencent_data_trace.tool_trace(trace_info) - - mock_trace_utils.convert_to_span_id.assert_called_once_with("msg-id", "message") - mock_trace_utils.convert_to_trace_id.assert_called_once_with("msg-id") - mock_span_builder.build_tool_span.assert_called_once_with(trace_info, 123, 456) - tencent_data_trace.trace_client.add_span.assert_called_once() - - def test_tool_trace_no_msg_id(self, tencent_data_trace): - trace_info = tool_trace_info() - trace_info.message_id = None - - tencent_data_trace.tool_trace(trace_info) - tencent_data_trace.trace_client.add_span.assert_not_called() - - def test_tool_trace_exception(self, tencent_data_trace, caplog: pytest.LogCaptureFixture): - trace_info = tool_trace_info() - trace_info.message_id = "msg-id" - - with ( - patch( - "dify_trace_tencent.tencent_trace.TencentTraceUtils.convert_to_span_id", side_effect=Exception("error") - ), - caplog.at_level(logging.ERROR), - ): - tencent_data_trace.tool_trace(trace_info) - assert "[Tencent APM] Failed to process tool trace" in caplog.text - - def test_dataset_retrieval_trace(self, tencent_data_trace, mock_trace_utils, mock_span_builder): - trace_info = DatasetRetrievalTraceInfo(metadata={}) - trace_info.message_id = "msg-id" - - mock_trace_utils.convert_to_span_id.return_value = 456 - mock_trace_utils.convert_to_trace_id.return_value = 123 - - tencent_data_trace.dataset_retrieval_trace(trace_info) - - mock_trace_utils.convert_to_span_id.assert_called_once_with("msg-id", "message") - mock_trace_utils.convert_to_trace_id.assert_called_once_with("msg-id") - mock_span_builder.build_retrieval_span.assert_called_once_with(trace_info, 123, 456) - tencent_data_trace.trace_client.add_span.assert_called_once() - - def test_dataset_retrieval_trace_no_msg_id(self, tencent_data_trace): - trace_info = DatasetRetrievalTraceInfo(metadata={}) - trace_info.message_id = None - - tencent_data_trace.dataset_retrieval_trace(trace_info) - tencent_data_trace.trace_client.add_span.assert_not_called() - - def test_dataset_retrieval_trace_exception(self, tencent_data_trace, caplog: pytest.LogCaptureFixture): - trace_info = DatasetRetrievalTraceInfo(metadata={}) - trace_info.message_id = "msg-id" - - with ( - patch( - "dify_trace_tencent.tencent_trace.TencentTraceUtils.convert_to_span_id", side_effect=Exception("error") - ), - caplog.at_level(logging.ERROR), - ): - tencent_data_trace.dataset_retrieval_trace(trace_info) - assert "[Tencent APM] Failed to process dataset retrieval trace" in caplog.text - - def test_suggested_question_trace(self, tencent_data_trace, caplog: pytest.LogCaptureFixture): - trace_info = SuggestedQuestionTraceInfo(metadata={}, total_tokens=0, suggested_question=[], level="info") - with caplog.at_level(logging.INFO): - tencent_data_trace.suggested_question_trace(trace_info) - assert "[Tencent APM] Processing suggested question trace" in caplog.text - - def test_suggested_question_trace_exception( - self, tencent_data_trace, monkeypatch, caplog: pytest.LogCaptureFixture - ): - trace_info = SuggestedQuestionTraceInfo(metadata={}, total_tokens=0, suggested_question=[], level="info") - target_logger = logging.getLogger("dify_trace_tencent.tencent_trace") - monkeypatch.setattr(target_logger, "info", MagicMock(side_effect=Exception("error"))) - with caplog.at_level(logging.ERROR): - tencent_data_trace.suggested_question_trace(trace_info) - assert "[Tencent APM] Failed to process suggested question trace" in caplog.text - - def test_process_workflow_nodes(self, tencent_data_trace, mock_trace_utils): - trace_info = workflow_trace_info() - trace_info.workflow_run_id = "run-id" - mock_trace_utils.convert_to_span_id.return_value = 111 - - node1 = workflow_node_execution() - node1.id = "n1" - node1.node_type = BuiltinNodeTypes.LLM - node2 = workflow_node_execution() - node2.id = "n2" - node2.node_type = BuiltinNodeTypes.TOOL - - with ( - patch.object(tencent_data_trace, "_get_workflow_node_executions", return_value=[node1, node2]), - patch.object(tencent_data_trace, "_build_workflow_node_span", side_effect=["span1", "span2"]), - patch.object(tencent_data_trace, "_record_llm_metrics") as mock_metrics, - ): - tencent_data_trace._process_workflow_nodes(trace_info, 123) - - assert tencent_data_trace.trace_client.add_span.call_count == 2 - mock_metrics.assert_called_once_with(node1) - - def test_process_workflow_nodes_node_exception( - self, tencent_data_trace, mock_trace_utils, caplog: pytest.LogCaptureFixture - ): - trace_info = workflow_trace_info() - mock_trace_utils.convert_to_span_id.return_value = 111 - - node = workflow_node_execution() - node.id = "n1" - next_node = workflow_node_execution() - next_node.id = "n2" - - with ( - patch.object(tencent_data_trace, "_get_workflow_node_executions", return_value=[node, next_node]), - patch.object( - tencent_data_trace, "_build_workflow_node_span", side_effect=[Exception("node error"), "next-span"] - ), - caplog.at_level(logging.ERROR), - ): - tencent_data_trace._process_workflow_nodes(trace_info, 123) - assert "[Tencent APM] Failed to process node execution: n1" in caplog.text - tencent_data_trace.trace_client.add_span.assert_called_once_with("next-span") - - def test_process_workflow_nodes_exception( - self, tencent_data_trace, mock_trace_utils, caplog: pytest.LogCaptureFixture - ): - trace_info = workflow_trace_info() - mock_trace_utils.convert_to_span_id.side_effect = Exception("outer error") - - with caplog.at_level(logging.ERROR): - tencent_data_trace._process_workflow_nodes(trace_info, 123) - assert "[Tencent APM] Failed to process workflow nodes" in caplog.text - - def test_build_workflow_node_span(self, tencent_data_trace, mock_span_builder): - trace_info = workflow_trace_info() - - nodes = [ - (BuiltinNodeTypes.LLM, mock_span_builder.build_workflow_llm_span), - (BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL, mock_span_builder.build_workflow_retrieval_span), - (BuiltinNodeTypes.TOOL, mock_span_builder.build_workflow_tool_span), - (BuiltinNodeTypes.CODE, mock_span_builder.build_workflow_task_span), - ] - - for node_type, builder_method in nodes: - node = workflow_node_execution() - node.node_type = node_type - builder_method.return_value = "span" - - result = tencent_data_trace._build_workflow_node_span(node, 123, trace_info, 456) - - assert result == "span" - builder_method.assert_called_once_with(123, 456, trace_info, node) - - def test_build_workflow_node_span_exception( - self, tencent_data_trace, mock_span_builder, caplog: pytest.LogCaptureFixture - ): - node = workflow_node_execution() - node.node_type = BuiltinNodeTypes.LLM - node.id = "n1" - mock_span_builder.build_workflow_llm_span.side_effect = Exception("error") - - with caplog.at_level(logging.DEBUG): - result = tencent_data_trace._build_workflow_node_span(node, 123, MagicMock(), 456) - assert result is None - assert len([r for r in caplog.records if r.levelno == logging.DEBUG]) >= 1 - - @pytest.mark.parametrize( - "sqlite3_session", - [(Account, App, Tenant, TenantAccountJoin)], - indirect=True, - ) - def test_get_workflow_node_executions( - self, - tencent_data_trace, - monkeypatch: pytest.MonkeyPatch, - sqlite3_session: Session, - ) -> None: - account = Account(name="Trace User", email="trace-user@example.com") - tenant = Tenant(name="Trace Tenant") - sqlite3_session.add_all([account, tenant]) - sqlite3_session.flush() - app = App( - id="app-1", - tenant_id=tenant.id, - name="Trace App", - description="", - mode=AppMode.WORKFLOW, - icon_type=IconType.EMOJI, - icon="robot", - icon_background="#FFFFFF", - enable_site=False, - enable_api=False, - created_by=account.id, - max_active_requests=0, - ) - tenant_join = TenantAccountJoin( - tenant_id=tenant.id, - account_id=account.id, - current=True, - role=TenantAccountRole.OWNER, - ) - sqlite3_session.add_all([app, tenant_join]) - sqlite3_session.commit() - - trace_info = workflow_trace_info() - trace_info.metadata = {"app_id": app.id} - trace_info.workflow_run_id = "run-1" - database = SimpleNamespace(engine=sqlite3_session.get_bind()) - monkeypatch.setattr("dify_trace_tencent.tencent_trace.db", database) - - with patch("dify_trace_tencent.tencent_trace.SQLAlchemyWorkflowNodeExecutionRepository") as mock_repo: - mock_repo.return_value.get_by_workflow_execution.return_value = [] - results = tencent_data_trace._get_workflow_node_executions(trace_info) - - assert results == [] - service_account = mock_repo.call_args.kwargs["user"] - assert isinstance(service_account, Account) - assert service_account.id == account.id - assert mock_repo.call_args.kwargs["tenant_id"] == tenant.id - - @pytest.mark.parametrize("sqlite3_session", [()], indirect=True) - def test_get_workflow_node_executions_no_app_id( - self, - tencent_data_trace, - caplog: pytest.LogCaptureFixture, - monkeypatch: pytest.MonkeyPatch, - sqlite3_session: Session, - ) -> None: - trace_info = workflow_trace_info() - trace_info.metadata = {} - monkeypatch.setattr( - "dify_trace_tencent.tencent_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind()), - ) - - with caplog.at_level(logging.ERROR): - results = tencent_data_trace._get_workflow_node_executions(trace_info) - assert results == [] - assert len([r for r in caplog.records if r.levelno == logging.ERROR]) >= 1 - - @pytest.mark.parametrize("sqlite3_session", [(App,)], indirect=True) - def test_get_workflow_node_executions_app_not_found( - self, - tencent_data_trace, - caplog: pytest.LogCaptureFixture, - monkeypatch: pytest.MonkeyPatch, - sqlite3_session: Session, - ) -> None: - trace_info = workflow_trace_info() - trace_info.metadata = {"app_id": "app-1"} - monkeypatch.setattr( - "dify_trace_tencent.tencent_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind()), - ) - - with caplog.at_level(logging.ERROR): - results = tencent_data_trace._get_workflow_node_executions(trace_info) - assert results == [] - assert len([r for r in caplog.records if r.levelno == logging.ERROR]) >= 1 - - def test_get_user_id_workflow(self, tencent_data_trace): - trace_info = workflow_trace_info() - trace_info.tenant_id = "tenant-1" - trace_info.metadata = {"user_id": "user-1"} - - with ( - patch("dify_trace_tencent.tencent_trace.sessionmaker", side_effect=Exception("Database error")), - patch("dify_trace_tencent.tencent_trace.db") as mock_db, - ): - mock_db.init_app = MagicMock() - mock_db.engine = MagicMock() - - user_id = tencent_data_trace._get_user_id(trace_info) - assert user_id == "unknown" - - def test_get_user_id_only_user_id(self, tencent_data_trace): - trace_info = message_trace_info() - trace_info.metadata = {"user_id": "user-1"} - - user_id = tencent_data_trace._get_user_id(trace_info) - assert user_id == "user-1" - - def test_get_user_id_anonymous(self, tencent_data_trace): - trace_info = message_trace_info() - trace_info.metadata = {} - - user_id = tencent_data_trace._get_user_id(trace_info) - assert user_id == "anonymous" - - def test_get_user_id_exception(self, tencent_data_trace, caplog: pytest.LogCaptureFixture): - trace_info = workflow_trace_info() - trace_info.tenant_id = "t" - trace_info.metadata = {"user_id": "u"} - - with ( - patch("dify_trace_tencent.tencent_trace.sessionmaker", side_effect=Exception("error")), - caplog.at_level(logging.ERROR), - ): - user_id = tencent_data_trace._get_user_id(trace_info) - assert user_id == "unknown" - assert "[Tencent APM] Failed to get user ID" in caplog.text - - def test_record_llm_metrics_usage_in_process_data(self, tencent_data_trace): - node = workflow_node_execution() - node.process_data = { - "usage": { - "latency": 2.5, - "time_to_first_token": 0.5, - "time_to_generate": 2.0, - "prompt_tokens": 10, - "completion_tokens": 20, - }, - "model_provider": "openai", - "model_name": "gpt-4", - "model_mode": "chat", - } - node.outputs = {} - - tencent_data_trace._record_llm_metrics(node) - - tencent_data_trace.trace_client.record_llm_duration.assert_called_once() - tencent_data_trace.trace_client.record_time_to_first_token.assert_called_once() - tencent_data_trace.trace_client.record_time_to_generate.assert_called_once() - assert tencent_data_trace.trace_client.record_token_usage.call_count == 2 - - def test_record_llm_metrics_usage_in_outputs(self, tencent_data_trace): - node = workflow_node_execution() - node.process_data = {} - node.outputs = {"usage": {"latency": 1.0, "prompt_tokens": 5}} - - tencent_data_trace._record_llm_metrics(node) - tencent_data_trace.trace_client.record_llm_duration.assert_called_once() - tencent_data_trace.trace_client.record_token_usage.assert_called_once() - - def test_record_llm_metrics_exception(self, tencent_data_trace, caplog: pytest.LogCaptureFixture): - node = workflow_node_execution() - node.process_data = None - node.outputs = None - - with caplog.at_level(logging.DEBUG): - tencent_data_trace._record_llm_metrics(node) - # Should not crash - - def test_record_message_llm_metrics(self, tencent_data_trace): - trace_info = message_trace_info() - trace_info.metadata = {"ls_provider": "openai", "ls_model_name": "gpt-4"} - trace_info.message_data = {"provider_response_latency": 1.1} - trace_info.is_streaming_request = True - trace_info.gen_ai_server_time_to_first_token = 0.2 - trace_info.llm_streaming_time_to_generate = 0.9 - trace_info.message_tokens = 15 - trace_info.answer_tokens = 25 - - tencent_data_trace._record_message_llm_metrics(trace_info) - - tencent_data_trace.trace_client.record_llm_duration.assert_called_once() - tencent_data_trace.trace_client.record_time_to_first_token.assert_called_once() - tencent_data_trace.trace_client.record_time_to_generate.assert_called_once() - assert tencent_data_trace.trace_client.record_token_usage.call_count == 2 - - def test_record_message_llm_metrics_object_data(self, tencent_data_trace): - trace_info = message_trace_info() - trace_info.metadata = {} - msg_data = MagicMock() - msg_data.provider_response_latency = 1.1 - msg_data.model_provider = "anthropic" - msg_data.model_id = "claude" - trace_info.message_data = msg_data - trace_info.is_streaming_request = False - - tencent_data_trace._record_message_llm_metrics(trace_info) - tencent_data_trace.trace_client.record_llm_duration.assert_called_once() - - def test_record_message_llm_metrics_exception(self, tencent_data_trace, caplog: pytest.LogCaptureFixture): - trace_info = message_trace_info() - # Inject malformed metadata to exercise the metrics error handler. - trace_info = trace_info.model_copy(update={"metadata": None}) - - with caplog.at_level(logging.DEBUG): - tencent_data_trace._record_message_llm_metrics(trace_info) - # Should not crash - - def test_record_workflow_trace_duration(self, tencent_data_trace): - trace_info = workflow_trace_info() - from datetime import datetime, timedelta - - now = datetime.now() - trace_info.start_time = now - trace_info.end_time = now + timedelta(seconds=3) - trace_info.workflow_run_status = "succeeded" - trace_info.conversation_id = "conv-1" - - # Mock the record_trace_duration method to capture arguments - with patch.object(tencent_data_trace.trace_client, "record_trace_duration") as mock_record: - tencent_data_trace._record_workflow_trace_duration(trace_info) - - # Assert the method was called once - mock_record.assert_called_once() - - # Extract arguments passed to the method - args, kwargs = mock_record.call_args - - # Validate the duration argument - assert args[0] == 3.0 - - # Validate the attributes dict in kwargs - attributes = kwargs["attributes"] if "attributes" in kwargs else args[1] if len(args) > 1 else {} - assert attributes["conversation_mode"] == "workflow" - assert attributes["has_conversation"] == "true" - - def test_record_workflow_trace_duration_fallback(self, tencent_data_trace): - trace_info = workflow_trace_info() - trace_info.start_time = None - trace_info.workflow_run_elapsed_time = 4.5 - trace_info.workflow_run_status = "failed" - trace_info.conversation_id = None - - with patch.object(tencent_data_trace.trace_client, "record_trace_duration") as mock_record: - tencent_data_trace._record_workflow_trace_duration(trace_info) - mock_record.assert_called_once() - args, kwargs = mock_record.call_args - assert args[0] == 4.5 - # Check attributes dict (either in kwargs or as second positional arg) - attributes = kwargs["attributes"] if "attributes" in kwargs else args[1] if len(args) > 1 else {} - assert attributes["has_conversation"] == "false" - - def test_record_workflow_trace_duration_exception(self, tencent_data_trace, caplog: pytest.LogCaptureFixture): - trace_info = workflow_trace_info() - trace_info.start_time = MagicMock() # This might cause total_seconds() to fail if not mocked right - - with caplog.at_level(logging.DEBUG): - tencent_data_trace._record_workflow_trace_duration(trace_info) - - def test_record_message_trace_duration(self, tencent_data_trace): - trace_info = message_trace_info() - from datetime import datetime, timedelta - - now = datetime.now() - trace_info.start_time = now - trace_info.end_time = now + timedelta(seconds=2) - trace_info.conversation_mode = "chat" - trace_info.is_streaming_request = True - - tencent_data_trace._record_message_trace_duration(trace_info) - tencent_data_trace.trace_client.record_trace_duration.assert_called_once_with( - 2.0, {"conversation_mode": "chat", "stream": "true"} - ) - - def test_record_message_trace_duration_exception(self, tencent_data_trace, caplog: pytest.LogCaptureFixture): - trace_info = message_trace_info() - trace_info.start_time = None - - with caplog.at_level(logging.DEBUG): - tencent_data_trace._record_message_trace_duration(trace_info) - - def test_close(self, tencent_data_trace): - client = tencent_data_trace.trace_client - tencent_data_trace.close() - client.shutdown.assert_called_once() - - def test_close_is_idempotent(self, tencent_data_trace): - client = tencent_data_trace.trace_client - - tencent_data_trace.close() - tencent_data_trace.close() - - client.shutdown.assert_called_once() - - def test_close_exception(self, tencent_data_trace, caplog: pytest.LogCaptureFixture): - tencent_data_trace.trace_client.shutdown.side_effect = Exception("error") - with caplog.at_level(logging.ERROR): - tencent_data_trace.close() - assert "[Tencent APM] Failed to shutdown trace client during cleanup" in caplog.text - - def test_close_handles_async_shutdown_mock(self, tencent_data_trace): - shutdown = AsyncMock() - tencent_data_trace.trace_client.shutdown = shutdown - - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - tencent_data_trace.close() - gc.collect() - - shutdown.assert_called_once() - assert not [ - warning - for warning in caught - if issubclass(warning.category, RuntimeWarning) - and "AsyncMockMixin._execute_mock_call" in str(warning.message) - ] diff --git a/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_utils.py b/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_utils.py deleted file mode 100644 index 63c6d680d786da..00000000000000 --- a/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_utils.py +++ /dev/null @@ -1,105 +0,0 @@ -"""Unit tests for Tencent APM tracing utilities.""" - -from __future__ import annotations - -import hashlib -import uuid -from datetime import UTC, datetime -from unittest.mock import patch - -import pytest -from dify_trace_tencent.utils import TencentTraceUtils -from opentelemetry.trace import Link, TraceFlags - - -def test_convert_to_trace_id_with_valid_uuid() -> None: - uuid_str = "12345678-1234-5678-1234-567812345678" - assert TencentTraceUtils.convert_to_trace_id(uuid_str) == uuid.UUID(uuid_str).int - - -def test_convert_to_trace_id_uses_uuid4_when_none() -> None: - expected_uuid = uuid.UUID("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa") - with patch("dify_trace_tencent.utils.uuid.uuid4", return_value=expected_uuid) as uuid4_mock: - assert TencentTraceUtils.convert_to_trace_id(None) == expected_uuid.int - uuid4_mock.assert_called_once() - - -def test_convert_to_trace_id_raises_value_error_for_invalid_uuid() -> None: - with pytest.raises(ValueError, match=r"^Invalid UUID input:"): - TencentTraceUtils.convert_to_trace_id("not-a-uuid") - - -def test_convert_to_span_id_is_deterministic_and_sensitive_to_type() -> None: - uuid_str = "12345678-1234-5678-1234-567812345678" - span_type = "llm" - - uuid_obj = uuid.UUID(uuid_str) - combined_key = f"{uuid_obj.hex}-{span_type}" - hash_bytes = hashlib.sha256(combined_key.encode("utf-8")).digest() - expected = int.from_bytes(hash_bytes[:8], byteorder="big", signed=False) - - assert TencentTraceUtils.convert_to_span_id(uuid_str, span_type) == expected - assert TencentTraceUtils.convert_to_span_id(uuid_str, "other") != expected - - -def test_convert_to_span_id_uses_uuid4_when_none() -> None: - expected_uuid = uuid.UUID("bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb") - with patch("dify_trace_tencent.utils.uuid.uuid4", return_value=expected_uuid) as uuid4_mock: - span_id = TencentTraceUtils.convert_to_span_id(None, "workflow") - assert isinstance(span_id, int) - uuid4_mock.assert_called_once() - - -def test_convert_to_span_id_raises_value_error_for_invalid_uuid() -> None: - with pytest.raises(ValueError, match=r"^Invalid UUID input:"): - TencentTraceUtils.convert_to_span_id("bad-uuid", "span") - - -def test_generate_span_id_skips_invalid_span_id() -> None: - with patch( - "dify_trace_tencent.utils.random.getrandbits", - side_effect=[TencentTraceUtils.INVALID_SPAN_ID, 42], - ) as bits_mock: - assert TencentTraceUtils.generate_span_id() == 42 - assert bits_mock.call_count == 2 - - -def test_convert_datetime_to_nanoseconds_accepts_datetime() -> None: - start_time = datetime(2024, 1, 1, 0, 0, 0, tzinfo=UTC) - expected = int(start_time.timestamp() * 1e9) - assert TencentTraceUtils.convert_datetime_to_nanoseconds(start_time) == expected - - -def test_convert_datetime_to_nanoseconds_uses_now_when_none() -> None: - fixed = datetime(2024, 1, 2, 3, 4, 5, tzinfo=UTC) - expected = int(fixed.timestamp() * 1e9) - - with patch("dify_trace_tencent.utils.datetime") as datetime_mock: - datetime_mock.now.return_value = fixed - assert TencentTraceUtils.convert_datetime_to_nanoseconds(None) == expected - datetime_mock.now.assert_called_once() - - -@pytest.mark.parametrize( - ("trace_id_str", "expected_trace_id"), - [ - ("0" * 31 + "1", int("0" * 31 + "1", 16)), - (str(uuid.UUID("cccccccc-cccc-cccc-cccc-cccccccccccc")), uuid.UUID("cccccccc-cccc-cccc-cccc-cccccccccccc").int), - ], -) -def test_create_link_accepts_hex_or_uuid(trace_id_str: str, expected_trace_id: int) -> None: - link = TencentTraceUtils.create_link(trace_id_str) - assert isinstance(link, Link) - assert link.context.trace_id == expected_trace_id - assert link.context.span_id == TencentTraceUtils.INVALID_SPAN_ID - assert link.context.is_remote is False - assert link.context.trace_flags == TraceFlags(TraceFlags.SAMPLED) - - -@pytest.mark.parametrize("trace_id_str", ["g" * 32, "not-a-uuid", None]) -def test_create_link_falls_back_to_uuid4(trace_id_str: object) -> None: - fallback_uuid = uuid.UUID("dddddddd-dddd-dddd-dddd-dddddddddddd") - with patch("dify_trace_tencent.utils.uuid.uuid4", return_value=fallback_uuid) as uuid4_mock: - link = TencentTraceUtils.create_link(trace_id_str) # type: ignore[arg-type] - assert link.context.trace_id == fallback_uuid.int - uuid4_mock.assert_called_once() diff --git a/api/providers/trace/trace-weave/pyproject.toml b/api/providers/trace/trace-weave/pyproject.toml index 8225cdbf563dd0..9cdd488d095d62 100644 --- a/api/providers/trace/trace-weave/pyproject.toml +++ b/api/providers/trace/trace-weave/pyproject.toml @@ -2,9 +2,12 @@ name = "dify-trace-weave" version = "0.0.1" dependencies = [ - "weave==0.52.36", + "dify-trace-sdk", ] description = "Dify ops tracing provider (Weave)." +[project.entry-points."dify.trace_providers"] +weave = "dify_trace_weave.provider:WeaveProvider" + [tool.setuptools.packages.find] where = ["src"] diff --git a/api/providers/trace/trace-weave/src/dify_trace_weave/agent_trace.py b/api/providers/trace/trace-weave/src/dify_trace_weave/agent_trace.py deleted file mode 100644 index 8799e78f0b8bfa..00000000000000 --- a/api/providers/trace/trace-weave/src/dify_trace_weave/agent_trace.py +++ /dev/null @@ -1,89 +0,0 @@ -"""Completed Agent trees through Weave's typed call API and an isolated transport.""" - -from collections.abc import Mapping -from typing import Any, cast -from urllib.parse import quote -from uuid import UUID, uuid5 - -from weave.trace_server.trace_server_interface import ( - CallEndReq, - CallStartReq, - EndedCallSchemaForInsert, - StartedCallSchemaForInsert, - SummaryInsertMap, -) - -from core.ops.unified_trace.entities import CanonicalTrace -from dify_trace_weave.transport import TraceServer, resolve_project -from services.tracing.provider_common import endpoint, required, token_usage, utc -from services.tracing.provider_dispatch import TraceExportRejectedError, TraceExportResult, TraceExportUncertainError - - -def export_trace(provider: str, config: Mapping[str, Any], trace: CanonicalTrace) -> TraceExportResult: - api_key = required(config, "api_key") - entity, project, ui_host = resolve_project(config, api_key) - # Weave's concrete HTTP client leaves unrelated optional API methods as protocol stubs. - server = TraceServer( # pyrefly: ignore[bad-instantiation] - endpoint(config.get("endpoint") or "https://trace.wandb.ai"), auth=("api", api_key) - ) - project_id = f"{entity}/{project}" - trace_uuid = UUID(hex=trace.trace_id) - span_ids = {span.id: str(uuid5(trace_uuid, span.id)) for span in trace.spans} - # Weave counts calls whose id equals turn_id as conversation turns. - turn_id = span_ids[trace.root_span_id] if trace.session_id else None - wrote = False - try: - for span in trace.spans: - started = server.call_start( - CallStartReq( - start=StartedCallSchemaForInsert( - project_id=project_id, - id=span_ids[span.id], - op_name=span.name, - display_name=span.name, - trace_id=str(trace_uuid), - parent_id=span_ids.get(span.parent_id) if span.parent_id else None, - thread_id=trace.session_id or None, - turn_id=turn_id, - started_at=utc(span.start_time), - attributes={**span.metadata, "dify.span_kind": span.kind.value, "dify.span_id": span.id}, - inputs=span.inputs if isinstance(span.inputs, dict) else {"input": span.inputs}, - ) - ) - ) - wrote = True - if started.id != span_ids[span.id] or started.trace_id != str(trace_uuid): - raise TraceExportUncertainError("Weave did not acknowledge the submitted identifiers") - for span in reversed(trace.spans): - usage = token_usage(span) - summary: dict[str, Any] = {} - if usage: - model = str(span.metadata.get("model_name") or span.metadata.get("model") or "unknown") - summary["usage"] = { - model: { - {"input": "prompt_tokens", "output": "completion_tokens", "total": "total_tokens"}[key]: value - for key, value in usage.items() - } - | {"requests": 1} - } - if span.end_time is None: - raise TraceExportRejectedError("Weave requires completed spans") - server.call_end( - CallEndReq( - end=EndedCallSchemaForInsert( - project_id=project_id, - id=span_ids[span.id], - ended_at=utc(span.end_time), - output=span.outputs, - exception=span.error or ("Agent span failed" if span.status.value == "error" else None), - summary=cast(SummaryInsertMap, summary), - ) - ) - ) - except Exception: - if wrote: - raise TraceExportUncertainError("Weave trace was partially accepted") from None - raise - base = f"{ui_host}/{quote(entity, safe='')}/{quote(project, safe='')}/r/call/" - urls = {span_id: base + call_id for span_id, call_id in span_ids.items()} - return TraceExportResult(trace_id=str(trace_uuid), trace_url=urls[trace.root_span_id], span_urls=urls) diff --git a/api/providers/trace/trace-weave/src/dify_trace_weave/config.py b/api/providers/trace/trace-weave/src/dify_trace_weave/config.py index b54b7fd1be9010..a3d366743503b5 100644 --- a/api/providers/trace/trace-weave/src/dify_trace_weave/config.py +++ b/api/providers/trace/trace-weave/src/dify_trace_weave/config.py @@ -1,29 +1,41 @@ -from pydantic import ValidationInfo, field_validator +from urllib.parse import urlsplit -from core.ops.entities.config_entity import BaseTracingConfig -from core.ops.utils import validate_url_with_path +from dify_trace_sdk.config import ProviderConfig, destination, http_url, secret +from pydantic import field_validator +DEFAULT_ENDPOINT = "https://trace.wandb.ai" +DEFAULT_HOST = "https://api.wandb.ai" -class WeaveConfig(BaseTracingConfig): - """ - Model class for Weave tracing config. - """ - api_key: str +def _base_url(value: str, allowed_schemes: tuple[str, ...]) -> str: + """Validate a URL that request paths are appended to, keeping its path prefix.""" + url = http_url(value, allowed_schemes=allowed_schemes) + if any(char in "?#\\" or ord(char) <= 32 or ord(char) == 127 for char in url): + raise ValueError("URL must not contain a query, a fragment, whitespace, or backslashes") + try: + _ = urlsplit(url).port + except ValueError: + raise ValueError("URL port is invalid") from None + return url + + +class WeaveConfig(ProviderConfig): + api_key: str = secret() entity: str | None = None + """The W&B team or user; ``verify`` resolves the API key's default entity when it is unset.""" project: str - endpoint: str = "https://trace.wandb.ai" - host: str | None = None + """A project name, or ``entity/project`` when ``entity`` is unset.""" + endpoint: str = destination(default=DEFAULT_ENDPOINT) + """The Weave trace server.""" + host: str | None = destination(default=None) + """The W&B API host, ``DEFAULT_HOST`` when unset.""" @field_validator("endpoint") @classmethod - def endpoint_validator(cls, v, info: ValidationInfo): - # Weave only allows HTTPS for endpoint - return validate_url_with_path(v, "https://trace.wandb.ai", allowed_schemes=("https",)) + def _endpoint(cls, value: str) -> str: + return _base_url(value, ("https",)) if value.strip() else DEFAULT_ENDPOINT @field_validator("host") @classmethod - def host_validator(cls, v, info: ValidationInfo): - if v is not None and v.strip() != "": - return validate_url_with_path(v, v, allowed_schemes=("https", "http")) - return v + def _host(cls, value: str | None) -> str | None: + return _base_url(value, ("https", "http")) if value and value.strip() else None diff --git a/api/providers/trace/trace-weave/src/dify_trace_weave/configuration.py b/api/providers/trace/trace-weave/src/dify_trace_weave/configuration.py deleted file mode 100644 index e5f92f7055f168..00000000000000 --- a/api/providers/trace/trace-weave/src/dify_trace_weave/configuration.py +++ /dev/null @@ -1,250 +0,0 @@ -"""Read-only Weave configuration checks without SDK login or global client state. - -Projects must already exist. Configuration verification never creates a project, -trace, or local credential file; actual legacy exports retain their SDK lifecycle. -""" - -from dataclasses import dataclass -from typing import Literal -from urllib.parse import quote, urlsplit, urlunsplit - -from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError -from weave.trace_server.trace_server_interface import CallsQueryStatsReq, CallsQueryStatsRes - -from core.helper import ssrf_proxy -from dify_trace_weave.config import WeaveConfig - -type VerificationCode = Literal[ - "invalid_project", - "invalid_endpoint", - "authentication_failed", - "permission_denied", - "project_not_found_or_inaccessible", - "project_read_only", - "project_permission_unknown", - "graphql_error", - "malformed_response", - "endpoint_unavailable", - "connection_failed", -] - -_MESSAGES: dict[VerificationCode, str] = { - "invalid_project": "Weave requires a valid entity and project name.", - "invalid_endpoint": "Weave endpoint is invalid.", - "authentication_failed": "Weave authentication failed.", - "permission_denied": "Weave permission was denied.", - "project_not_found_or_inaccessible": "Create the Weave project first or check project access.", - "project_read_only": "Weave project is read-only for these credentials.", - "project_permission_unknown": "Weave project write permission could not be verified.", - "graphql_error": "Weave project verification returned an error.", - "malformed_response": "Weave verification returned an invalid response.", - "endpoint_unavailable": "Weave endpoint could not be verified.", - "connection_failed": "Weave verification request failed.", -} - -_VIEWER_QUERY = "query DifyWeaveViewer { viewer { id entity } }" -_PROJECT_QUERY = """query DifyWeaveProject($entity: String!, $project: String!) { - viewer { id } - project(name: $project, entityName: $entity) { id name entityName readOnly } -}""" -_JSON_OBJECT = TypeAdapter(dict[str, object]) - - -class WeaveConfigurationError(ValueError): - """Fixed diagnostic codes: never include keys, request data, or server errors.""" - - def __init__(self, code: VerificationCode) -> None: - self.code = code - super().__init__(_MESSAGES[code]) - - -@dataclass(frozen=True, slots=True) -class VerifiedWeaveProject: - entity: str - project: str - project_url: str - - -class _Viewer(BaseModel): - id: str = Field(min_length=1) - entity: str | None = None - model_config = ConfigDict(strict=True) - - -class _Project(BaseModel): - id: str = Field(min_length=1) - name: str = Field(min_length=1) - entity_name: str = Field(alias="entityName", min_length=1) - read_only: bool | None = Field(alias="readOnly") - model_config = ConfigDict(strict=True) - - -def _base_url(value: str, *, allow_http: bool) -> str: - try: - parsed = urlsplit(value) - valid = ( - parsed.scheme in ({"https", "http"} if allow_http else {"https"}) - and parsed.hostname - and parsed.username is None - and parsed.password is None - and not parsed.query - and not parsed.fragment - and "\\" not in value - and not any(ord(char) <= 32 or ord(char) == 127 for char in value) - ) - _ = parsed.port - except ValueError: - raise WeaveConfigurationError("invalid_endpoint") from None - if not valid: - raise WeaveConfigurationError("invalid_endpoint") - return urlunsplit(parsed).rstrip("/") - - -def _name(value: str) -> str: - value = value.strip() - if ( - not value - or value in {".", ".."} - or "/" in value - or "\\" in value - or any(ord(char) < 32 or ord(char) == 127 for char in value) - ): - raise WeaveConfigurationError("invalid_project") - return value - - -def _target(config: WeaveConfig) -> tuple[str | None, str]: - entity = config.entity.strip() if config.entity else None - project = config.project.strip() - if "/" in project: - parts = project.split("/") - if entity or len(parts) != 2: - raise WeaveConfigurationError("invalid_project") - entity, project = _name(parts[0]), parts[1] - return (_name(entity) if entity else None), _name(project) - - -def console_url(config: WeaveConfig) -> str: - """Map the public API origin to its UI; preserve private origins and prefixes.""" - host = _base_url(config.host or "https://api.wandb.ai", allow_http=True) - parsed = urlsplit(host) - if ( - parsed.scheme == "https" - and parsed.hostname == "api.wandb.ai" - and parsed.port in {None, 443} - and not parsed.path - ): - return "https://wandb.ai" - return host - - -def project_url(config: WeaveConfig) -> str | None: - """Pure presentation fallback; never resolve a default entity over the network.""" - entity, project = _target(config) - if entity is None: - return None - return f"{console_url(config)}/{quote(entity, safe='')}/{quote(project, safe='')}" - - -def _post(url: str, *, api_key: str, body: dict[str, object]) -> dict[str, object]: - try: - response = ssrf_proxy.post( - url, - auth=("api", api_key), - json=body, - max_retries=0, - timeout=30.0, - follow_redirects=False, - ) - except Exception: - raise WeaveConfigurationError("connection_failed") from None - try: - if response.status_code == 401: - raise WeaveConfigurationError("authentication_failed") - if response.status_code == 403: - raise WeaveConfigurationError("permission_denied") - if not 200 <= response.status_code < 300: - raise WeaveConfigurationError("endpoint_unavailable") - try: - return _JSON_OBJECT.validate_python(response.json(), strict=True) - except ValueError: - raise WeaveConfigurationError("malformed_response") from None - finally: - response.close() - - -def _graphql(host: str, *, api_key: str, body: dict[str, object]) -> dict[str, object]: - response = _post(host + "/graphql", api_key=api_key, body=body) - errors = response.get("errors") - if errors is not None: - if not isinstance(errors, list): - raise WeaveConfigurationError("malformed_response") - if errors: - raise WeaveConfigurationError("graphql_error") - try: - return _JSON_OBJECT.validate_python(response.get("data"), strict=True) - except ValidationError: - raise WeaveConfigurationError("malformed_response") from None - - -def verify_configuration(config: WeaveConfig) -> VerifiedWeaveProject: - """Verify credentials, an existing writable project, and the trace read endpoint. - - The GraphQL fields match W&B's Viewer, GetProject and IsProjectReadOnly - queries. CallsQueryStats uses the same typed read API as the Weave SDK. - These checks do not write a test trace or replace ingestion authorization. - """ - if not config.api_key.strip(): - raise WeaveConfigurationError("authentication_failed") - entity, project = _target(config) - host = _base_url(config.host or "https://api.wandb.ai", allow_http=True) - trace_endpoint = _base_url(config.endpoint, allow_http=False) - try: - if entity is None: - data = _graphql(host, api_key=config.api_key, body={"query": _VIEWER_QUERY}) - if "viewer" not in data: - raise WeaveConfigurationError("malformed_response") - if data.get("viewer") is None: - raise WeaveConfigurationError("authentication_failed") - viewer = _Viewer.model_validate(data["viewer"]) - if not viewer.entity: - raise WeaveConfigurationError("invalid_project") - entity = _name(viewer.entity) - - data = _graphql( - host, - api_key=config.api_key, - body={"query": _PROJECT_QUERY, "variables": {"entity": entity, "project": project}}, - ) - if "viewer" not in data or "project" not in data: - raise WeaveConfigurationError("malformed_response") - if data.get("viewer") is None: - raise WeaveConfigurationError("authentication_failed") - _Viewer.model_validate(data["viewer"]) - if data.get("project") is None: - raise WeaveConfigurationError("project_not_found_or_inaccessible") - resolved = _Project.model_validate(data["project"]) - if resolved.name != project or resolved.entity_name != entity: - raise WeaveConfigurationError("invalid_project") - if resolved.read_only is True: - raise WeaveConfigurationError("project_read_only") - if resolved.read_only is None: - raise WeaveConfigurationError("project_permission_unknown") - - stats_request = CallsQueryStatsReq(project_id=f"{entity}/{project}", limit=1) - stats_response = _post( - trace_endpoint + "/calls/query_stats", - api_key=config.api_key, - body=stats_request.model_dump(mode="json", exclude_none=True), - ) - stats = CallsQueryStatsRes.model_validate(stats_response, strict=True) - if stats.count < 0: - raise WeaveConfigurationError("malformed_response") - except ValidationError: - raise WeaveConfigurationError("malformed_response") from None - - return VerifiedWeaveProject( - entity=entity, - project=project, - project_url=f"{console_url(config)}/{quote(entity, safe='')}/{quote(project, safe='')}", - ) diff --git a/api/providers/trace/trace-weave/src/dify_trace_weave/entities/__init__.py b/api/providers/trace/trace-weave/src/dify_trace_weave/entities/__init__.py deleted file mode 100644 index e69de29bb2d1d6..00000000000000 diff --git a/api/providers/trace/trace-weave/src/dify_trace_weave/entities/weave_trace_entity.py b/api/providers/trace/trace-weave/src/dify_trace_weave/entities/weave_trace_entity.py deleted file mode 100644 index f8a3e3723cebc2..00000000000000 --- a/api/providers/trace/trace-weave/src/dify_trace_weave/entities/weave_trace_entity.py +++ /dev/null @@ -1,100 +0,0 @@ -from collections.abc import Mapping -from typing import Any - -from pydantic import BaseModel, Field, field_validator -from pydantic_core.core_schema import ValidationInfo - -from core.ops.utils import replace_text_with_content - - -class WeaveTokenUsage(BaseModel): - input_tokens: int | None = None - output_tokens: int | None = None - total_tokens: int | None = None - - -class WeaveMultiModel(BaseModel): - file_list: list[str] | None = Field(None, description="List of files") - - -class WeaveTraceModel(WeaveTokenUsage, WeaveMultiModel): - id: str = Field(..., description="ID of the trace") - op: str = Field(..., description="Name of the operation") - inputs: str | Mapping[str, Any] | list | None | None = Field(None, description="Inputs of the trace") - outputs: str | Mapping[str, Any] | list | None | None = Field(None, description="Outputs of the trace") - attributes: str | dict[str, Any] | list | None | None = Field( - None, description="Metadata and attributes associated with trace" - ) - exception: str | None = Field(None, description="Exception message of the trace") - - @field_validator("inputs", "outputs") - @classmethod - def ensure_dict(cls, v, info: ValidationInfo): - field_name = info.field_name - values = info.data - if v == {} or v is None: - return v - usage_metadata = { - "input_tokens": values.get("input_tokens", 0), - "output_tokens": values.get("output_tokens", 0), - "total_tokens": values.get("total_tokens", 0), - } - file_list = values.get("file_list", []) - match v: - case str(): - if field_name == "inputs": - return { - "messages": { - "role": "user", - "content": v, - "usage_metadata": usage_metadata, - "file_list": file_list, - }, - } - elif field_name == "outputs": - return { - "choices": { - "role": "ai", - "content": v, - "usage_metadata": usage_metadata, - "file_list": file_list, - }, - } - case list(): - data = {} - if len(v) > 0 and isinstance(v[0], dict): - # rename text to content - v = replace_text_with_content(data=v) - if field_name == "inputs": - data = { - "messages": [ - dict(msg, **{"usage_metadata": usage_metadata, "file_list": file_list}) # type: ignore - for msg in v - ] - if isinstance(v, list) - else v, - } - elif field_name == "outputs": - data = { - "choices": { - "role": "ai", - "content": v, - "usage_metadata": usage_metadata, - "file_list": file_list, - }, - } - return data - else: - return { - "choices": { - "role": "ai" if field_name == "outputs" else "user", - "content": str(v), - "usage_metadata": usage_metadata, - "file_list": file_list, - }, - } - if isinstance(v, dict): - v["usage_metadata"] = usage_metadata - v["file_list"] = file_list - return v - return v diff --git a/api/providers/trace/trace-weave/src/dify_trace_weave/provider.py b/api/providers/trace/trace-weave/src/dify_trace_weave/provider.py new file mode 100644 index 00000000000000..9554d3bf201a20 --- /dev/null +++ b/api/providers/trace/trace-weave/src/dify_trace_weave/provider.py @@ -0,0 +1,374 @@ +"""Weave destination over the trace server's call API. + +Each span becomes one Weave call, started parent-first and ended children-first +as the Weave SDK does. Call IDs derive from span IDs, so a redelivered unit +reuses the IDs of its first delivery. Starts and ends are separate requests, so +a failure after the first accepted call leaves the outcome unknown. The W&B +GraphQL API is only read, to verify a configuration and to resolve an unset +entity. The ``weave`` package is not imported because importing it patches PIL +for the whole process. +""" + +import base64 +from collections.abc import Mapping +from typing import override +from urllib.parse import quote, urlsplit +from uuid import UUID + +from dify_trace_sdk.egress import EgressResponse +from dify_trace_sdk.errors import ( + EgressError, + ExportOutcomeUnknownError, + ExportRejectedError, + ExportRetryableError, + VerificationError, +) +from dify_trace_sdk.exporting import json_value, send_export_request +from dify_trace_sdk.ids import span_uuid, trace_uuid +from dify_trace_sdk.model import Span, SpanKind, SpanStatus, Trace +from dify_trace_sdk.provider import ProviderCapabilities, SendResult, TraceProvider +from pydantic import BaseModel, ConfigDict, Field, JsonValue, ValidationError + +from dify_trace_weave.config import DEFAULT_HOST, WeaveConfig + +VERIFY_TIMEOUT_SECONDS = 10.0 +PUBLIC_CONSOLE_URL = "https://wandb.ai" +# The trace server rejects calls whose names are longer than these. +MAX_OP_NAME_LENGTH = 128 +MAX_DISPLAY_NAME_LENGTH = 1024 + +_VIEWER_QUERY = "query DifyWeaveViewer { viewer { id entity } }" +_PROJECT_QUERY = """query DifyWeaveProject($entity: String!, $project: String!) { + viewer { id } + project(name: $project, entityName: $entity) { id name entityName readOnly } +}""" + + +class _InvalidNameError(ValueError): + """An entity or project is not a usable W&B name.""" + + +class _Viewer(BaseModel): + id: str = Field(min_length=1) + entity: str | None = None + + model_config = ConfigDict(strict=True) + + +class _Project(BaseModel): + id: str = Field(min_length=1) + name: str = Field(min_length=1) + entity_name: str = Field(alias="entityName", min_length=1) + read_only: bool | None = Field(alias="readOnly") + + model_config = ConfigDict(strict=True) + + +class _CallStats(BaseModel): + count: int = Field(ge=0) + + model_config = ConfigDict(strict=True) + + +def _name(value: str) -> str: + name = value.strip() + if not name or name in {".", ".."} or any(char in "/\\" or ord(char) < 32 or ord(char) == 127 for char in name): + raise _InvalidNameError + return name + + +def _target(config: WeaveConfig) -> tuple[str | None, str]: + """The entity, None when the API key's default entity applies, and the project. + + Raises: + _InvalidNameError: A name is unusable, or an ``entity/project`` path conflicts with ``entity``. + """ + entity = config.entity.strip() if config.entity else None + project = config.project.strip() + if "/" in project: + if entity: + raise _InvalidNameError + entity, _, project = project.partition("/") + return _name(entity), _name(project) + return (_name(entity) if entity else None), _name(project) + + +def _api_host(config: WeaveConfig) -> str: + return (config.host or DEFAULT_HOST).rstrip("/") + + +def _console_url(config: WeaveConfig) -> str: + """The W&B UI: the public API host maps to wandb.ai, other hosts serve the UI themselves.""" + host = _api_host(config) + parsed = urlsplit(host) + if ( + parsed.scheme == "https" + and parsed.hostname == "api.wandb.ai" + and parsed.port in {None, 443} + and not parsed.path + ): + return PUBLIC_CONSOLE_URL + return host + + +def _project_url(config: WeaveConfig, entity: str, project: str) -> str: + return f"{_console_url(config)}/{quote(entity, safe='')}/{quote(project, safe='')}" + + +def _headers(config: WeaveConfig) -> dict[str, str]: + token = base64.b64encode(f"api:{config.api_key}".encode()).decode() + return {"Authorization": f"Basic {token}", "Content-Type": "application/json", "Accept": "application/json"} + + +def _json(response: EgressResponse) -> object: + try: + return response.json() + except ValueError: + return None + + +def _lookup(document: object, *path: str) -> object: + """The value at ``path`` in a JSON document, or None when the document has another shape.""" + for key in path: + document = document.get(key) if isinstance(document, dict) else None + return document + + +def _same_id(value: object, expected: str) -> bool: + """Whether ``value`` names the ``expected`` UUID, in any spelling.""" + if not isinstance(value, str): + return False + try: + return UUID(value) == UUID(expected) + except ValueError: + return False + + +def _call_id(trace: Trace, span_id: str) -> str: + return str(span_uuid(trace.trace_id, span_id)) + + +def _inputs(value: JsonValue) -> dict[str, JsonValue]: + """Weave call inputs must be an object.""" + if isinstance(value, dict): + return value + return {} if value is None else {"input": value} + + +def _summary(span: Span) -> dict[str, object]: + """Weave reads token usage per model name from the call summary.""" + if span.kind is not SpanKind.LLM or span.usage is None: + return {} + return { + "usage": { + span.model_name or "unknown": { + "prompt_tokens": span.usage.prompt_tokens, + "completion_tokens": span.usage.completion_tokens, + "total_tokens": span.usage.total_tokens, + "requests": 1, + } + } + } + + +def _call_bodies(trace: Trace, project_id: str, trace_key: str) -> tuple[list[bytes], list[bytes]]: + """Encoded call starts in parent-first order and call ends in children-first order.""" + thread_id = trace.session_id or None + # Weave counts a call as a conversation turn when its id is its turn_id. Every + # call of a root unit carries the root's id; fragments carry none, so they + # never count as turns of their own. + turn_id = _call_id(trace, trace.root.span_id) if thread_id and trace.parent is None else None + starts = [ + { + "start": { + "project_id": project_id, + "id": _call_id(trace, span.span_id), + "op_name": span.name[:MAX_OP_NAME_LENGTH], + "display_name": span.name[:MAX_DISPLAY_NAME_LENGTH], + "trace_id": trace_key, + "parent_id": _call_id(trace, span.parent_span_id) if span.parent_span_id else None, + "thread_id": thread_id, + "turn_id": turn_id, + "started_at": span.start_time.isoformat(), + "attributes": {**span.attributes, "dify_span_kind": span.kind.value, "dify_span_id": span.span_id}, + "inputs": _inputs(span.inputs), + } + } + for span in trace.spans + ] + ends = [ + { + "end": { + "project_id": project_id, + "id": _call_id(trace, span.span_id), + "ended_at": span.end_time.isoformat(), + "exception": (span.error or "Span failed") if span.status is SpanStatus.ERROR else None, + "output": span.outputs, + "summary": _summary(span), + } + } + for span in reversed(trace.spans) + ] + try: + return [json_value(body).encode() for body in starts], [json_value(body).encode() for body in ends] + except ValueError: + raise ExportRejectedError("Weave trace contains numbers that JSON cannot represent") from None + + +def _acknowledged(response: EgressResponse, call_id: str, trace_key: str) -> bool: + """Whether a call start response echoes the submitted call and trace IDs.""" + body = _json(response) + return _same_id(_lookup(body, "id"), call_id) and _same_id(_lookup(body, "trace_id"), trace_key) + + +class WeaveProvider(TraceProvider[WeaveConfig]): + id = "weave" + config_model = WeaveConfig + capabilities = ProviderCapabilities(native_cost=False, trace_url=True, span_url=True) + + @override + def verify(self, config: WeaveConfig) -> WeaveConfig: + """Check the credentials, an existing writable project, and the trace server's read API.""" + if not config.api_key.strip(): + raise VerificationError("authentication_failed") + try: + entity, project = _target(config) + if entity is None: + entity = self._viewer_entity(config) + self._verify_project(config, entity, project) + stats = self._verify_request( + config, + f"{config.endpoint.rstrip('/')}/calls/query_stats", + {"project_id": f"{entity}/{project}", "limit": 1}, + ) + _CallStats.model_validate(stats) + except _InvalidNameError: + raise VerificationError("invalid_project") from None + except ValidationError: + raise VerificationError("malformed_response") from None + return config.model_copy(update={"entity": entity, "project": project}) + + @override + def project_url(self, config: WeaveConfig) -> str | None: + try: + entity, project = _target(config) + except _InvalidNameError: + return _console_url(config) + return _project_url(config, entity, project) if entity else _console_url(config) + + @override + def send(self, config: WeaveConfig, trace: Trace) -> SendResult: + try: + entity, project = _target(config) + except _InvalidNameError: + raise ExportRejectedError("Weave entity or project name is invalid") from None + if entity is None: + entity = self._default_entity(config) + trace_key = str(trace_uuid(trace.trace_id)) + starts, ends = _call_bodies(trace, f"{entity}/{project}", trace_key) + endpoint = config.endpoint.rstrip("/") + written = False + try: + for span, start in zip(trace.spans, starts, strict=True): + response = self._post(config, f"{endpoint}/call/start", start) + written = True + if not _acknowledged(response, _call_id(trace, span.span_id), trace_key): + raise ExportOutcomeUnknownError("Weave did not acknowledge the submitted call identifiers") + for end in ends: + self._post(config, f"{endpoint}/call/end", end) + except (ExportRetryableError, ExportRejectedError): + if written: + raise ExportOutcomeUnknownError("Weave accepted only part of the trace") from None + raise + calls_url = f"{_project_url(config, entity, project)}/r/call" + root_span_id = trace.parent.span_id if trace.parent is not None else trace.root.span_id + return SendResult( + trace_id=trace_key, + trace_url=f"{calls_url}/{_call_id(trace, root_span_id)}", + span_urls={span.span_id: f"{calls_url}/{_call_id(trace, span.span_id)}" for span in trace.spans}, + ) + + def _post(self, config: WeaveConfig, url: str, content: bytes) -> EgressResponse: + return send_export_request( + self.egress, + "POST", + url, + destination="Weave", + headers=_headers(config), + content=content, + sensitive_url=False, + ) + + def _default_entity(self, config: WeaveConfig) -> str: + """Resolve the API key's default entity, which Weave uses when none is configured.""" + try: + response = self._post(config, f"{_api_host(config)}/graphql", json_value({"query": _VIEWER_QUERY}).encode()) + except ExportOutcomeUnknownError: + # The lookup only reads, so a lost answer cannot have written anything. + raise ExportRetryableError("Weave entity lookup failed in transport") from None + entity = _lookup(_json(response), "data", "viewer", "entity") + try: + return _name(entity if isinstance(entity, str) else "") + except _InvalidNameError: + raise ExportRejectedError("Weave could not resolve the default entity of the API key") from None + + def _viewer_entity(self, config: WeaveConfig) -> str: + data = self._graphql(config, {"query": _VIEWER_QUERY}) + if "viewer" not in data: + raise VerificationError("malformed_response") + if data["viewer"] is None: + raise VerificationError("authentication_failed") + return _name(_Viewer.model_validate(data["viewer"]).entity or "") + + def _verify_project(self, config: WeaveConfig, entity: str, project: str) -> None: + data = self._graphql(config, {"query": _PROJECT_QUERY, "variables": {"entity": entity, "project": project}}) + if "viewer" not in data or "project" not in data: + raise VerificationError("malformed_response") + if data["viewer"] is None: + raise VerificationError("authentication_failed") + _Viewer.model_validate(data["viewer"]) + if data["project"] is None: + raise VerificationError("project_not_found_or_inaccessible") + found = _Project.model_validate(data["project"]) + if found.name != project or found.entity_name != entity: + raise VerificationError("invalid_project") + if found.read_only is None: + raise VerificationError("project_permission_unknown") + if found.read_only: + raise VerificationError("project_read_only") + + def _graphql(self, config: WeaveConfig, body: Mapping[str, object]) -> dict[str, object]: + """Run a W&B GraphQL query and return its ``data``.""" + answer = self._verify_request(config, f"{_api_host(config)}/graphql", body) + errors = answer.get("errors") + if errors is not None and not isinstance(errors, list): + raise VerificationError("malformed_response") + if errors: + raise VerificationError("graphql_error") + data = answer.get("data") + if not isinstance(data, dict): + raise VerificationError("malformed_response") + return data + + def _verify_request(self, config: WeaveConfig, url: str, body: Mapping[str, object]) -> dict[str, object]: + try: + response = self.egress.request( + "POST", + url, + headers=_headers(config), + content=json_value(body).encode(), + timeout=VERIFY_TIMEOUT_SECONDS, + sensitive_url=False, + ) + except EgressError: + raise VerificationError("connection_failed") from None + if response.status_code == 401: + raise VerificationError("authentication_failed") + if response.status_code == 403: + raise VerificationError("permission_denied") + if not 200 <= response.status_code < 300: + raise VerificationError("endpoint_unavailable") + answer = _json(response) + if not isinstance(answer, dict): + raise VerificationError("malformed_response") + return answer diff --git a/api/providers/trace/trace-weave/src/dify_trace_weave/transport.py b/api/providers/trace/trace-weave/src/dify_trace_weave/transport.py deleted file mode 100644 index ed578b5d768124..00000000000000 --- a/api/providers/trace/trace-weave/src/dify_trace_weave/transport.py +++ /dev/null @@ -1,56 +0,0 @@ -"""Instance-authenticated Weave SDK transport shared by managed and legacy tracing.""" - -from collections.abc import Mapping -from typing import Any, override - -from pydantic import BaseModel -from weave.trace_server_bindings.remote_http_trace_server import RemoteHTTPTraceServer - -from dify_trace_weave.config import WeaveConfig -from dify_trace_weave.configuration import console_url -from services.tracing.provider_common import endpoint, request, required -from services.tracing.provider_dispatch import TraceExportRejectedError, TraceExportUncertainError - - -class TraceServer(RemoteHTTPTraceServer): - """Use SDK request/response validation without SDK retries or global W&B login.""" - - @override - def _post_request_executor(self, url: str, req: BaseModel, stream: bool = False): - return request( - "POST", - self.trace_server_url + url, - "weave", - auth=self._auth, - content=req.model_dump_json(by_alias=True).encode(), - headers={"Content-Type": "application/json"}, - ) - - @override - def _generic_request(self, url, req, req_model, res_model, method="POST", params=None): - response = self._post_request_executor(url, req) - try: - return res_model.model_validate(response.json()) - except ValueError: - raise TraceExportUncertainError("Weave returned an unreadable ingestion acknowledgment") from None - - -def resolve_project(config: Mapping[str, Any], api_key: str) -> tuple[str, str, str]: - project = required(config, "project") - host = endpoint(config.get("host") or "https://api.wandb.ai") - entity = config.get("entity") - if not entity and "/" in project: - entity, project = project.split("/", 1) - if not entity: - response = request( - "POST", - host + "/graphql", - "weave", - auth=("api", api_key), - json={"query": "query DifyTraceEntity { viewer { entity } }"}, - ) - entity = (response.json().get("data") or {}).get("viewer", {}).get("entity") - if not isinstance(entity, str) or not entity or "/" in entity or not project or "/" in project: - raise TraceExportRejectedError("Weave project or entity could not be resolved") - ui_host = console_url(WeaveConfig(api_key=api_key, entity=entity, project=project, host=host)) - return entity, project, ui_host diff --git a/api/providers/trace/trace-weave/src/dify_trace_weave/weave_trace.py b/api/providers/trace/trace-weave/src/dify_trace_weave/weave_trace.py deleted file mode 100644 index 6766536203e626..00000000000000 --- a/api/providers/trace/trace-weave/src/dify_trace_weave/weave_trace.py +++ /dev/null @@ -1,544 +0,0 @@ -import logging -import os -import uuid -from contextvars import ContextVar -from datetime import UTC, datetime, timedelta -from typing import Any, cast, override - -from sqlalchemy.orm import sessionmaker -from weave.trace_server.trace_server_interface import ( - CallEndReq, - CallStartReq, - EndedCallSchemaForInsert, - StartedCallSchemaForInsert, - SummaryInsertMap, - TraceStatus, -) - -from core.ops.base_trace_instance import BaseTraceInstance -from core.ops.entities.trace_entity import ( - BaseTraceInfo, - DatasetRetrievalTraceInfo, - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - SuggestedQuestionTraceInfo, - ToolTraceInfo, - TraceTaskName, - WorkflowTraceInfo, -) -from core.ops.exceptions import RetryableTraceDispatchError -from core.repositories import DifyCoreRepositoryFactory -from dify_trace_weave.config import WeaveConfig -from dify_trace_weave.configuration import console_url, project_url, verify_configuration -from dify_trace_weave.entities.weave_trace_entity import WeaveTraceModel -from dify_trace_weave.transport import TraceServer, resolve_project -from extensions.ext_database import db -from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey -from models import EndUser, MessageFile, WorkflowNodeExecutionTriggeredFrom -from services.tracing.provider_common import endpoint -from services.tracing.provider_dispatch import TraceExportUncertainError - -logger = logging.getLogger(__name__) -_dispatch_started_calls: ContextVar[dict[str, dict[str, Any]] | None] = ContextVar("weave_dispatch_calls", default=None) - - -def _same_identifier(actual: str, expected: str) -> bool: - """Accept UUID formatting differences without rewriting legacy wire identifiers.""" - if actual == expected: - return True - try: - return uuid.UUID(actual) == uuid.UUID(expected) - except ValueError: - return False - - -class WeaveDataTrace(BaseTraceInstance): - def __init__( - self, - weave_config: WeaveConfig, - ): - super().__init__(weave_config) - # Construction is local: no global SDK client, login, netrc, or background queue. - self._config = weave_config.model_copy(deep=True) - self.weave_api_key = self._config.api_key - self.project_name = self._config.project - self.entity = self._config.entity - self.host = self._config.host - self.server = TraceServer( # pyrefly: ignore[bad-instantiation] - endpoint(self._config.endpoint), should_batch=False, auth=("api", self.weave_api_key) - ) - self._project_id: str | None = None - self.file_base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001") - self.calls: dict[str, Any] = {} - - @property - def project_id(self) -> str: - # Existing saved configurations may omit entity. Resolve it lazily with - # this instance's credentials, never W&B's process-global defaults. - if self._project_id is None: - entity, project, _ = resolve_project(self._config.model_dump(), self.weave_api_key) - self._project_id = f"{entity}/{project}" - return self._project_id - - def get_project_url(self): - return project_url(self._config) or console_url(self._config) - - @override - def trace(self, trace_info: BaseTraceInfo): - logger.debug("Trace info: %s", trace_info) - accepted_calls: dict[str, dict[str, Any]] = {} - token = _dispatch_started_calls.set(accepted_calls) - try: - match trace_info: - case WorkflowTraceInfo(): - self.workflow_trace(trace_info) - case MessageTraceInfo(): - self.message_trace(trace_info) - case ModerationTraceInfo(): - self.moderation_trace(trace_info) - case SuggestedQuestionTraceInfo(): - self.suggested_question_trace(trace_info) - case DatasetRetrievalTraceInfo(): - self.dataset_retrieval_trace(trace_info) - case ToolTraceInfo(): - self.tool_trace(trace_info) - case GenerateNameTraceInfo(): - self.generate_name_trace(trace_info) - case _: - pass - except RetryableTraceDispatchError as error: - if accepted_calls and not error.acceptance_unknown: - raise TraceExportUncertainError( - "Weave trace was partially accepted; automatic retry is unsafe" - ) from None - raise - finally: - for call_id, call_meta in accepted_calls.items(): - if self.calls.get(call_id) is call_meta: - self.calls.pop(call_id, None) - _dispatch_started_calls.reset(token) - - def workflow_trace(self, trace_info: WorkflowTraceInfo): - trace_id = trace_info.trace_id or trace_info.message_id or trace_info.workflow_run_id - if trace_info.start_time is None: - trace_info.start_time = datetime.now() - - if trace_info.message_id: - message_attributes = trace_info.metadata - message_attributes["workflow_app_log_id"] = trace_info.workflow_app_log_id - - message_attributes["message_id"] = trace_info.message_id - message_attributes["workflow_run_id"] = trace_info.workflow_run_id - message_attributes["trace_id"] = trace_id - message_attributes["start_time"] = trace_info.start_time - message_attributes["end_time"] = trace_info.end_time - message_attributes["tags"] = ["message", "workflow"] - - message_run = WeaveTraceModel( - id=trace_info.message_id, - op=str(TraceTaskName.MESSAGE_TRACE), - inputs=dict(trace_info.workflow_run_inputs), - outputs=dict(trace_info.workflow_run_outputs), - total_tokens=trace_info.total_tokens, - attributes=message_attributes, - exception=trace_info.error, - file_list=[], - ) - self.start_call(message_run, parent_run_id=trace_info.workflow_run_id) - self.finish_call(message_run) - - workflow_attributes = trace_info.metadata - workflow_attributes["workflow_run_id"] = trace_info.workflow_run_id - workflow_attributes["trace_id"] = trace_id - workflow_attributes["start_time"] = trace_info.start_time - workflow_attributes["end_time"] = trace_info.end_time - workflow_attributes["tags"] = ["dify_workflow"] - - workflow_run = WeaveTraceModel( - file_list=trace_info.file_list, - total_tokens=trace_info.total_tokens, - id=trace_info.workflow_run_id, - op=str(TraceTaskName.WORKFLOW_TRACE), - inputs=dict(trace_info.workflow_run_inputs), - outputs=dict(trace_info.workflow_run_outputs), - attributes=workflow_attributes, - exception=trace_info.error, - ) - - self.start_call(workflow_run, parent_run_id=trace_info.message_id) - - # through workflow_run_id get all_nodes_execution using repository - session_factory = sessionmaker(bind=db.engine) - # Find the app's creator account - app_id = trace_info.metadata.get("app_id") - if not app_id: - raise ValueError("No app_id found in trace_info metadata") - - service_account = self.get_service_account_with_tenant(app_id) - - workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( - session_factory=session_factory, - tenant_id=trace_info.tenant_id, - user=service_account, - app_id=app_id, - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - - # Get all executions for this workflow run - workflow_node_executions = workflow_node_execution_repository.get_by_workflow_execution( - workflow_execution_id=trace_info.workflow_run_id - ) - - # rearrange workflow_node_executions by starting time - workflow_node_executions = sorted(workflow_node_executions, key=lambda x: x.created_at) - - for node_execution in workflow_node_executions: - node_execution_id = node_execution.id - tenant_id = trace_info.tenant_id # Use from trace_info instead - app_id = trace_info.metadata.get("app_id") # Use from trace_info instead - node_name = node_execution.title - node_type = node_execution.node_type - status = node_execution.status - if node_type == BuiltinNodeTypes.LLM: - inputs = node_execution.process_data.get("prompts", {}) if node_execution.process_data else {} - else: - inputs = node_execution.inputs or {} - outputs = node_execution.outputs or {} - created_at = node_execution.created_at or datetime.now() - elapsed_time = node_execution.elapsed_time - finished_at = created_at + timedelta(seconds=elapsed_time) - - execution_metadata = node_execution.metadata or {} - node_total_tokens = execution_metadata.get(WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS) or 0 - attributes = {str(k): v for k, v in execution_metadata.items()} - attributes.update( - { - "workflow_run_id": trace_info.workflow_run_id, - "node_execution_id": node_execution_id, - "tenant_id": tenant_id, - "app_id": app_id, - "app_name": node_name, - "node_type": node_type, - "status": status, - } - ) - - process_data = node_execution.process_data or {} - if process_data and process_data.get("model_mode") == "chat": - attributes.update( - { - "ls_provider": process_data.get("model_provider", ""), - "ls_model_name": process_data.get("model_name", ""), - } - ) - attributes["tags"] = ["node_execution"] - attributes["start_time"] = created_at - attributes["end_time"] = finished_at - attributes["elapsed_time"] = elapsed_time - attributes["workflow_run_id"] = trace_info.workflow_run_id - attributes["trace_id"] = trace_id - node_run = WeaveTraceModel( - total_tokens=node_total_tokens, - op=node_type, - inputs=inputs, - outputs=outputs, - file_list=trace_info.file_list, - attributes=attributes, - id=node_execution_id, - exception=None, - ) - - self.start_call(node_run, parent_run_id=trace_info.workflow_run_id) - self.finish_call(node_run) - - self.finish_call(workflow_run) - - def message_trace(self, trace_info: MessageTraceInfo): - # get message file data - file_list = cast(list[str], trace_info.file_list) or [] - message_file_data: MessageFile | None = trace_info.message_file_data - file_url = f"{self.file_base_url}/{message_file_data.url}" if message_file_data else "" - file_list.append(file_url) - attributes = trace_info.metadata - message_data = trace_info.message_data - if message_data is None: - return - message_id = message_data.id - - user_id = message_data.from_account_id - attributes["user_id"] = user_id - - if message_data.from_end_user_id: - end_user_data: EndUser | None = db.session.get(EndUser, message_data.from_end_user_id) - if end_user_data is not None: - end_user_id = end_user_data.session_id - attributes["end_user_id"] = end_user_id - - attributes["message_id"] = message_id - attributes["start_time"] = trace_info.start_time - attributes["end_time"] = trace_info.end_time - attributes["tags"] = ["message", str(trace_info.conversation_mode)] - - trace_id = trace_info.trace_id or message_id - attributes["trace_id"] = trace_id - - message_run = WeaveTraceModel( - id=trace_id, - op=str(TraceTaskName.MESSAGE_TRACE), - input_tokens=trace_info.message_tokens, - output_tokens=trace_info.answer_tokens, - total_tokens=trace_info.total_tokens, - inputs=trace_info.inputs, - outputs=trace_info.outputs, - exception=trace_info.error, - file_list=file_list, - attributes=attributes, - ) - self.start_call(message_run) - - # create llm run parented to message run - llm_run = WeaveTraceModel( - id=str(uuid.uuid4()), - input_tokens=trace_info.message_tokens, - output_tokens=trace_info.answer_tokens, - total_tokens=trace_info.total_tokens, - op="llm", - inputs=trace_info.inputs, - outputs=trace_info.outputs, - attributes=attributes, - file_list=[], - exception=None, - ) - self.start_call( - llm_run, - parent_run_id=trace_id, - ) - self.finish_call(llm_run) - self.finish_call(message_run) - - def moderation_trace(self, trace_info: ModerationTraceInfo): - if trace_info.message_data is None: - return - - attributes = trace_info.metadata - attributes["tags"] = ["moderation"] - attributes["message_id"] = trace_info.message_id - attributes["start_time"] = trace_info.start_time or trace_info.message_data.created_at - attributes["end_time"] = trace_info.end_time or trace_info.message_data.updated_at - - trace_id = trace_info.trace_id or trace_info.message_id - attributes["trace_id"] = trace_id - - moderation_run = WeaveTraceModel( - id=str(uuid.uuid4()), - op=str(TraceTaskName.MODERATION_TRACE), - inputs=trace_info.inputs, - outputs={ - "action": trace_info.action, - "flagged": trace_info.flagged, - "preset_response": trace_info.preset_response, - "inputs": trace_info.inputs, - }, - attributes=attributes, - exception=getattr(trace_info, "error", None), - file_list=[], - ) - self.start_call(moderation_run, parent_run_id=trace_id) - self.finish_call(moderation_run) - - def suggested_question_trace(self, trace_info: SuggestedQuestionTraceInfo): - message_data = trace_info.message_data - if message_data is None: - return - attributes = trace_info.metadata - attributes["message_id"] = trace_info.message_id - attributes["tags"] = ["suggested_question"] - attributes["start_time"] = (trace_info.start_time or message_data.created_at,) - attributes["end_time"] = (trace_info.end_time or message_data.updated_at,) - - trace_id = trace_info.trace_id or trace_info.message_id - attributes["trace_id"] = trace_id - - suggested_question_run = WeaveTraceModel( - id=str(uuid.uuid4()), - op=str(TraceTaskName.SUGGESTED_QUESTION_TRACE), - inputs=trace_info.inputs, - outputs=trace_info.suggested_question, - attributes=attributes, - exception=trace_info.error, - file_list=[], - ) - - self.start_call(suggested_question_run, parent_run_id=trace_id) - self.finish_call(suggested_question_run) - - def dataset_retrieval_trace(self, trace_info: DatasetRetrievalTraceInfo): - if trace_info.message_data is None: - return - attributes = trace_info.metadata - attributes["message_id"] = trace_info.message_id - attributes["tags"] = ["dataset_retrieval"] - attributes["start_time"] = (trace_info.start_time or trace_info.message_data.created_at,) - attributes["end_time"] = (trace_info.end_time or trace_info.message_data.updated_at,) - - trace_id = trace_info.trace_id or trace_info.message_id - attributes["trace_id"] = trace_id - - dataset_retrieval_run = WeaveTraceModel( - id=str(uuid.uuid4()), - op=str(TraceTaskName.DATASET_RETRIEVAL_TRACE), - inputs=trace_info.inputs, - outputs={"documents": trace_info.documents}, - attributes=attributes, - exception=getattr(trace_info, "error", None), - file_list=[], - ) - - self.start_call(dataset_retrieval_run, parent_run_id=trace_id) - self.finish_call(dataset_retrieval_run) - - def tool_trace(self, trace_info: ToolTraceInfo): - attributes = trace_info.metadata - attributes["tags"] = ["tool", trace_info.tool_name] - attributes["start_time"] = trace_info.start_time - attributes["end_time"] = trace_info.end_time - - message_id = trace_info.message_id or getattr(trace_info, "conversation_id", None) - message_id = message_id or None - trace_id = trace_info.trace_id or message_id - attributes["trace_id"] = trace_id - - tool_run = WeaveTraceModel( - id=str(uuid.uuid4()), - op=trace_info.tool_name, - inputs=trace_info.tool_inputs, - outputs=trace_info.tool_outputs, - file_list=[cast(str, trace_info.file_url)] if trace_info.file_url else [], - attributes=attributes, - exception=trace_info.error, - ) - self.start_call(tool_run, parent_run_id=trace_id) - self.finish_call(tool_run) - - def generate_name_trace(self, trace_info: GenerateNameTraceInfo): - attributes = trace_info.metadata - attributes["tags"] = ["generate_name"] - attributes["start_time"] = trace_info.start_time - attributes["end_time"] = trace_info.end_time - - name_run = WeaveTraceModel( - id=str(uuid.uuid4()), - op=str(TraceTaskName.GENERATE_NAME_TRACE), - inputs=trace_info.inputs, - outputs=trace_info.outputs, - attributes=attributes, - exception=getattr(trace_info, "error", None), - file_list=[], - ) - - self.start_call(name_run) - self.finish_call(name_run) - - def api_check(self): - # Read-only credential/project/endpoint verification; it creates no calls. - verify_configuration(self._config) - return True - - def _normalize_time(self, dt: datetime | None) -> datetime: - if dt is None: - return datetime.now(UTC) - if dt.tzinfo is None: - return dt.replace(tzinfo=UTC) - return dt - - def start_call(self, run_data: WeaveTraceModel, parent_run_id: str | None = None): - inputs = run_data.inputs - if inputs is None: - inputs = {} - elif not isinstance(inputs, dict): - inputs = {"inputs": str(inputs)} - - attributes = run_data.attributes - if attributes is None: - attributes = {} - elif not isinstance(attributes, dict): - attributes = {"attributes": str(attributes)} - - start_time = attributes.get("start_time") if isinstance(attributes, dict) else None - started_at = self._normalize_time(start_time if isinstance(start_time, datetime) else None) - trace_id = attributes.get("trace_id") if isinstance(attributes, dict) else None - if trace_id is None: - trace_id = run_data.id - - call_start_req = CallStartReq( - start=StartedCallSchemaForInsert( - project_id=self.project_id, - id=run_data.id, - op_name=str(run_data.op), - trace_id=trace_id, - parent_id=parent_run_id, - started_at=started_at, - attributes=attributes, - inputs=inputs, - wb_user_id=None, - ) - ) - started = self.server.call_start(call_start_req) - if not _same_identifier(started.id, run_data.id) or not _same_identifier(started.trace_id, trace_id): - raise TraceExportUncertainError("Weave did not acknowledge the submitted identifiers") - call_meta = {"trace_id": trace_id, "parent_id": parent_run_id} - self.calls[run_data.id] = call_meta - dispatch_calls = _dispatch_started_calls.get() - if dispatch_calls is not None: - dispatch_calls[run_data.id] = call_meta - - def finish_call(self, run_data: WeaveTraceModel): - call_meta = self.calls.get(run_data.id) - if not call_meta: - raise ValueError(f"Call with id {run_data.id} not found") - - attributes = run_data.attributes - if attributes is None: - attributes = {} - elif not isinstance(attributes, dict): - attributes = {"attributes": str(attributes)} - - start_time = attributes.get("start_time") if isinstance(attributes, dict) else None - end_time = attributes.get("end_time") if isinstance(attributes, dict) else None - started_at = self._normalize_time(start_time if isinstance(start_time, datetime) else None) - ended_at = self._normalize_time(end_time if isinstance(end_time, datetime) else None) - elapsed_ms = int((ended_at - started_at).total_seconds() * 1000) - if elapsed_ms < 0: - elapsed_ms = 0 - - status_counts = { - TraceStatus.SUCCESS: 0, - TraceStatus.ERROR: 0, - } - if run_data.exception: - status_counts[TraceStatus.ERROR] = 1 - else: - status_counts[TraceStatus.SUCCESS] = 1 - - summary: dict[str, Any] = { - "status_counts": status_counts, - "weave": {"latency_ms": elapsed_ms}, - } - - exception_str = str(run_data.exception) if run_data.exception else None - - call_end_req = CallEndReq( - end=EndedCallSchemaForInsert( - project_id=self.project_id, - id=run_data.id, - ended_at=ended_at, - exception=exception_str, - output=run_data.outputs, - summary=cast(SummaryInsertMap, summary), - ) - ) - try: - self.server.call_end(call_end_req) - finally: - self.calls.pop(run_data.id, None) diff --git a/api/providers/trace/trace-weave/tests/unit_tests/src/dify_trace_weave/test_config.py b/api/providers/trace/trace-weave/tests/unit_tests/src/dify_trace_weave/test_config.py index c6c4f7a09ceac6..ed2b46719f49ed 100644 --- a/api/providers/trace/trace-weave/tests/unit_tests/src/dify_trace_weave/test_config.py +++ b/api/providers/trace/trace-weave/tests/unit_tests/src/dify_trace_weave/test_config.py @@ -3,74 +3,37 @@ from pydantic import ValidationError -class TestWeaveConfig: - """Test cases for WeaveConfig""" - - def test_valid_config(self): - """Test valid Weave configuration""" - config = WeaveConfig( - api_key="test_key", - entity="test_entity", - project="test_project", - endpoint="https://custom.wandb.ai", - host="https://custom.host.com", - ) - assert config.api_key == "test_key" - assert config.entity == "test_entity" - assert config.project == "test_project" - assert config.endpoint == "https://custom.wandb.ai" - assert config.host == "https://custom.host.com" - - def test_default_values(self): - """Test default values are set correctly""" - config = WeaveConfig(api_key="key", project="project") - assert config.entity is None - assert config.endpoint == "https://trace.wandb.ai" - assert config.host is None - - def test_missing_required_fields(self): - """Test that required fields are enforced""" - with pytest.raises(ValidationError): - WeaveConfig.model_validate({}) - - with pytest.raises(ValidationError): - WeaveConfig.model_validate({"api_key": "key"}) - +def test_config_fields_and_defaults() -> None: + assert WeaveConfig.secret_fields() == {"api_key"} + assert WeaveConfig.derived_fields() == set() + assert [field.name for field in WeaveConfig.form_fields()] == ["api_key", "entity", "project", "endpoint", "host"] + stored = WeaveConfig.model_validate( + {"api_key": "k", "entity": None, "project": "demo", "endpoint": " ", "host": "", "project_url": "ignored"} + ) + assert (stored.entity, stored.endpoint, stored.host) == (None, "https://trace.wandb.ai", None) + private = WeaveConfig( + api_key="k", project="demo", endpoint="https://wandb.internal/traces/", host="http://wandb.internal/wandb" + ) + assert (private.endpoint, private.host) == ("https://wandb.internal/traces/", "http://wandb.internal/wandb") + for incomplete in ({"api_key": "k"}, {"project": "demo"}): with pytest.raises(ValidationError): - WeaveConfig.model_validate({"project": "project"}) - - def test_endpoint_validation_https_only(self): - """Test endpoint validation only allows HTTPS""" - with pytest.raises(ValidationError, match="URL must start with https://"): - WeaveConfig(api_key="key", project="project", endpoint="http://insecure.wandb.ai") - - def test_host_validation_optional(self): - """Test host validation is optional but validates when provided""" - config = WeaveConfig(api_key="key", project="project", host=None) - assert config.host is None - - config = WeaveConfig(api_key="key", project="project", host="") - assert config.host == "" - - config = WeaveConfig(api_key="key", project="project", host="https://valid.host.com") - assert config.host == "https://valid.host.com" - - def test_host_validation_invalid_scheme(self): - """Test host validation rejects invalid schemes when provided""" - with pytest.raises(ValidationError, match="URL must start with https:// or http://"): - WeaveConfig(api_key="key", project="project", host="ftp://invalid.host.com") - - def test_endpoint_preserves_path(self): - """Self-hosted Weave endpoints keep their path prefix""" - config = WeaveConfig(api_key="key", project="project", endpoint="https://wandb.internal/api") - assert config.endpoint == "https://wandb.internal/api" - - def test_host_preserves_path(self): - """Self-hosted W&B hosts keep their path prefix""" - config = WeaveConfig(api_key="key", project="project", host="https://wandb.internal/wandb") - assert config.host == "https://wandb.internal/wandb" - - def test_host_preserves_http_path(self): - """Self-hosted W&B hosts may be plain http and keep their path""" - config = WeaveConfig(api_key="key", project="project", host="http://wandb.internal/wandb") - assert config.host == "http://wandb.internal/wandb" + WeaveConfig.model_validate(incomplete) + + +@pytest.mark.parametrize( + ("field", "url"), + [ + ("endpoint", "http://trace.internal"), + ("host", "ftp://wandb.internal"), + ("host", "https://user:secret@wandb.internal"), + ("endpoint", "https://trace.internal/?token=secret"), + ("host", "https://wandb.internal/#secret"), + ("host", "https://wandb.internal/se cret"), + ("endpoint", "https://trace.internal\\secret"), + ("host", "https://wandb.internal:secret"), + ], +) +def test_config_rejects_unusable_urls_without_echoing_them(field: str, url: str) -> None: + with pytest.raises(ValidationError) as caught: + WeaveConfig.model_validate({"api_key": "k", "project": "demo", field: url}) + assert "secret" not in str(caught.value) diff --git a/api/providers/trace/trace-weave/tests/unit_tests/src/dify_trace_weave/test_provider.py b/api/providers/trace/trace-weave/tests/unit_tests/src/dify_trace_weave/test_provider.py new file mode 100644 index 00000000000000..6253d07e071a19 --- /dev/null +++ b/api/providers/trace/trace-weave/tests/unit_tests/src/dify_trace_weave/test_provider.py @@ -0,0 +1,634 @@ +import base64 +import subprocess +import sys +from typing import Any +from uuid import UUID, uuid5 + +import pytest +from dify_trace_sdk.egress import EgressResponse +from dify_trace_sdk.errors import ( + EgressBlockedError, + EgressError, + ExportOutcomeUnknownError, + ExportRejectedError, + ExportRetryableError, + VerificationError, +) +from dify_trace_sdk.model import Trace +from dify_trace_sdk.provider import SendResult +from dify_trace_sdk.testing import ( + MESSAGE_ID, + WORKFLOW_RUN_ID, + RecordedRequest, + RecordingEgress, + message_trace, + response, + tool_fragment, + workflow_trace, +) +from dify_trace_weave.config import WeaveConfig +from dify_trace_weave.provider import WeaveProvider +from pydantic import JsonValue + +CONFIG = WeaveConfig(api_key="weave-key", entity="team", project="demo") +AUTHORIZATION = "Basic " + base64.b64encode(b"api:weave-key").decode() +TRACE_SERVER = "https://trace.wandb.ai" +CALLS_URL = "https://wandb.ai/team/demo/r/call" +PROJECT = {"id": "project-1", "name": "demo", "entityName": "team", "readOnly": False} +PROJECT_ANSWER = response(200, {"data": {"viewer": {"id": "viewer-1"}, "project": PROJECT}}) +STATS = response(200, {"count": 0}) + + +def _provider(*responses: EgressResponse | Exception) -> tuple[WeaveProvider, RecordingEgress]: + egress = RecordingEgress(responses, grpc_target=None) + return WeaveProvider(egress), egress + + +def _body(request: RecordedRequest) -> dict[str, Any]: + body = request.json() + assert isinstance(body, dict) + return body + + +def _call_id(trace: Trace, span_id: str) -> str: + return str(uuid5(UUID(hex=trace.trace_id), span_id)) + + +def _acknowledgments(trace: Trace) -> list[EgressResponse | Exception]: + """What the trace server answers to every call start and end of ``trace``.""" + trace_key = str(UUID(hex=trace.trace_id)) + starts: list[EgressResponse | Exception] = [ + response(200, {"id": _call_id(trace, span.span_id), "trace_id": trace_key}) for span in trace.spans + ] + return starts + [response(200, {}) for _ in trace.spans] + + +def _with_root(trace: Trace, **fields: JsonValue) -> Trace: + return trace.model_copy(update={"spans": (trace.root.model_copy(update=fields), *trace.spans[1:])}) + + +def test_verify_checks_the_project_and_trace_server_with_read_only_requests() -> None: + provider, egress = _provider(PROJECT_ANSWER, STATS) + assert provider.verify(CONFIG) == CONFIG + + graphql, stats = egress.requests + assert graphql.url == "https://api.wandb.ai/graphql" + query = _body(graphql) + assert "mutation" not in query["query"] + assert query["variables"] == {"entity": "team", "project": "demo"} + assert stats.url == "https://trace.wandb.ai/calls/query_stats" + assert _body(stats) == {"project_id": "team/demo", "limit": 1} + assert all(request.method == "POST" for request in egress.requests) + assert all(request.headers["Authorization"] == AUTHORIZATION for request in egress.requests) + + +@pytest.mark.parametrize( + ("entity", "project", "lookups"), + [("team", "demo", 0), (None, "team/demo", 0), (" team ", " demo ", 0), (None, "demo", 1)], +) +def test_verify_resolves_the_entity_and_project(entity: str | None, project: str, lookups: int) -> None: + viewer = response(200, {"data": {"viewer": {"id": "viewer-1", "entity": "team"}}}) + provider, egress = _provider(*([viewer] * lookups), PROJECT_ANSWER, STATS) + verified = provider.verify(WeaveConfig(api_key="weave-key", entity=entity, project=project)) + assert (verified.entity, verified.project) == ("team", "demo") + assert len(egress.requests) == lookups + 2 + + +def test_verify_keeps_private_path_prefixes() -> None: + config = WeaveConfig( + api_key="weave-key", + entity="team", + project="demo", + host="http://wandb.internal/wandb/", + endpoint="https://trace.internal/proxy/", + ) + provider, egress = _provider(PROJECT_ANSWER, STATS) + provider.verify(config) + assert [request.url for request in egress.requests] == [ + "http://wandb.internal/wandb/graphql", + "https://trace.internal/proxy/calls/query_stats", + ] + + +@pytest.mark.parametrize( + ("entity", "project"), + [ + ("team", ""), + ("team", ".."), + ("/", "demo"), + (None, "/demo"), + (None, "team/"), + (None, "team/sub/demo"), + ("other", "team/demo"), + ("team", "de\\mo"), + ("team", "de\nmo"), + ], +) +def test_verify_rejects_unusable_names_before_any_request(entity: str | None, project: str) -> None: + provider, egress = _provider() + with pytest.raises(VerificationError) as caught: + provider.verify(WeaveConfig(api_key="weave-key", entity=entity, project=project)) + assert caught.value.code == "invalid_project" + assert egress.requests == [] + + +def test_verify_requires_an_api_key() -> None: + provider, egress = _provider() + with pytest.raises(VerificationError) as caught: + provider.verify(WeaveConfig(api_key=" ", entity="team", project="demo")) + assert caught.value.code == "authentication_failed" + assert egress.requests == [] + + +@pytest.mark.parametrize( + ("outcome", "code"), + [ + (response(401, {"error": "DO_NOT_EXPOSE"}), "authentication_failed"), + (response(403, {}), "permission_denied"), + (response(302, {}), "endpoint_unavailable"), + (response(404, {}), "endpoint_unavailable"), + (response(429, {}), "endpoint_unavailable"), + (response(503, {}), "endpoint_unavailable"), + (response(200, b"not json"), "malformed_response"), + (response(200, []), "malformed_response"), + (response(200, {"errors": [{"message": "DO_NOT_EXPOSE"}]}), "graphql_error"), + (response(200, {"errors": "DO_NOT_EXPOSE"}), "malformed_response"), + (response(200, {"data": []}), "malformed_response"), + (response(200, {"data": {}}), "malformed_response"), + (EgressError("DO_NOT_EXPOSE", request_sent=True), "connection_failed"), + (EgressBlockedError("DO_NOT_EXPOSE"), "connection_failed"), + ], +) +def test_verify_failures_have_stable_codes(outcome: EgressResponse | Exception, code: str) -> None: + provider, egress = _provider(outcome) + with pytest.raises(VerificationError) as caught: + provider.verify(CONFIG) + assert caught.value.code == code + assert str(caught.value) == code + assert len(egress.requests) == 1 + + +@pytest.mark.parametrize( + ("data", "code"), + [ + ({"viewer": {"id": "viewer-1"}, "project": None}, "project_not_found_or_inaccessible"), + ({"viewer": None, "project": PROJECT}, "authentication_failed"), + ({"viewer": {"id": ""}, "project": PROJECT}, "malformed_response"), + ({"viewer": {"id": "viewer-1"}}, "malformed_response"), + ({"viewer": {"id": "viewer-1"}, "project": {**PROJECT, "readOnly": True}}, "project_read_only"), + ({"viewer": {"id": "viewer-1"}, "project": {**PROJECT, "readOnly": None}}, "project_permission_unknown"), + ({"viewer": {"id": "viewer-1"}, "project": {**PROJECT, "readOnly": "false"}}, "malformed_response"), + ({"viewer": {"id": "viewer-1"}, "project": {**PROJECT, "id": ""}}, "malformed_response"), + ({"viewer": {"id": "viewer-1"}, "project": {**PROJECT, "name": "other"}}, "invalid_project"), + ({"viewer": {"id": "viewer-1"}, "project": {**PROJECT, "entityName": "other"}}, "invalid_project"), + ], +) +def test_verify_requires_an_existing_writable_project(data: dict[str, object], code: str) -> None: + provider, egress = _provider(response(200, {"data": data})) + with pytest.raises(VerificationError) as caught: + provider.verify(CONFIG) + assert caught.value.code == code + assert len(egress.requests) == 1 + + +@pytest.mark.parametrize( + ("data", "code"), + [ + ({}, "malformed_response"), + ({"viewer": None}, "authentication_failed"), + ({"viewer": {"id": "viewer-1", "entity": None}}, "invalid_project"), + ({"viewer": {"id": "viewer-1", "entity": "team/other"}}, "invalid_project"), + ({"viewer": {"id": "viewer-1", "entity": 7}}, "malformed_response"), + ], +) +def test_verify_needs_a_default_entity_when_none_is_configured(data: dict[str, object], code: str) -> None: + provider, egress = _provider(response(200, {"data": data})) + with pytest.raises(VerificationError) as caught: + provider.verify(WeaveConfig(api_key="weave-key", project="demo")) + assert caught.value.code == code + assert len(egress.requests) == 1 + + +@pytest.mark.parametrize("count", [-1, True, "0", None]) +def test_verify_rejects_unusable_trace_server_answers(count: object) -> None: + provider, egress = _provider(PROJECT_ANSWER, response(200, {"count": count})) + with pytest.raises(VerificationError) as caught: + provider.verify(CONFIG) + assert caught.value.code == "malformed_response" + assert len(egress.requests) == 2 + + +@pytest.mark.parametrize( + ("config", "url"), + [ + (CONFIG, "https://wandb.ai/team/demo"), + (WeaveConfig(api_key="k", project="team/demo"), "https://wandb.ai/team/demo"), + (WeaveConfig(api_key="k", project="demo"), "https://wandb.ai"), + (WeaveConfig(api_key="k", project="team/sub/demo"), "https://wandb.ai"), + ( + WeaveConfig(api_key="k", entity="team", project="demo", host="https://API.WANDB.AI:443/"), + "https://wandb.ai/team/demo", + ), + ( + WeaveConfig(api_key="k", entity="team one", project="demo & notes#", host="http://wandb.internal/wandb/"), + "http://wandb.internal/wandb/team%20one/demo%20%26%20notes%23", + ), + ], +) +def test_project_url_makes_no_requests(config: WeaveConfig, url: str) -> None: + provider, egress = _provider() + assert provider.project_url(config) == url + assert egress.requests == [] + + +def test_message_trace_wire_format() -> None: + trace = message_trace() + provider, egress = _provider(*_acknowledgments(trace)) + result = provider.send(CONFIG, trace) + + message, llm = trace.spans + message_call, llm_call = _call_id(trace, message.span_id), _call_id(trace, llm.span_id) + assert [request.url for request in egress.requests] == [ + f"{TRACE_SERVER}/call/start", + f"{TRACE_SERVER}/call/start", + f"{TRACE_SERVER}/call/end", + f"{TRACE_SERVER}/call/end", + ] + assert all(request.method == "POST" for request in egress.requests) + assert all(request.headers["Authorization"] == AUTHORIZATION for request in egress.requests) + assert all(request.headers["Content-Type"] == "application/json" for request in egress.requests) + assert [_body(request) for request in egress.requests] == [ + { + "start": { + "project_id": "team/demo", + "id": message_call, + "op_name": "message", + "display_name": "message", + "trace_id": MESSAGE_ID, + "parent_id": None, + "thread_id": "conversation-1", + "turn_id": message_call, + "started_at": "2026-10-01T08:00:00+00:00", + "attributes": { + "message_id": MESSAGE_ID, + "app_id": "app-1", + "tenant_id": "tenant-1", + "dify_span_kind": "chain", + "dify_span_id": message.span_id, + }, + "inputs": {"query": "What is Dify?"}, + } + }, + { + "start": { + "project_id": "team/demo", + "id": llm_call, + "op_name": "llm", + "display_name": "llm", + "trace_id": MESSAGE_ID, + "parent_id": message_call, + "thread_id": "conversation-1", + "turn_id": message_call, + "started_at": "2026-10-01T08:00:00+00:00", + "attributes": {"message_id": MESSAGE_ID, "dify_span_kind": "llm", "dify_span_id": llm.span_id}, + "inputs": {"input": [{"role": "user", "content": "What is Dify?"}]}, + } + }, + { + "end": { + "project_id": "team/demo", + "id": llm_call, + "ended_at": "2026-10-01T08:00:02+00:00", + "exception": None, + "output": "An LLM application platform.", + "summary": { + "usage": { + "gpt-4o": {"prompt_tokens": 12, "completion_tokens": 8, "total_tokens": 20, "requests": 1} + } + }, + } + }, + { + "end": { + "project_id": "team/demo", + "id": message_call, + "ended_at": "2026-10-01T08:00:02+00:00", + "exception": None, + "output": "An LLM application platform.", + "summary": {}, + } + }, + ] + assert result == SendResult( + trace_id=MESSAGE_ID, + trace_url=f"{CALLS_URL}/{message_call}", + span_urls={message.span_id: f"{CALLS_URL}/{message_call}", llm.span_id: f"{CALLS_URL}/{llm_call}"}, + ) + + +def test_workflow_trace_wire_format() -> None: + trace = workflow_trace() + provider, egress = _provider(*_acknowledgments(trace)) + provider.send(CONFIG, trace) + + run, llm, tool = trace.spans + run_call, llm_call, tool_call = (_call_id(trace, span.span_id) for span in trace.spans) + assert [request.url.rsplit("/", 1)[1] for request in egress.requests] == ["start"] * 3 + ["end"] * 3 + assert [_body(request) for request in egress.requests] == [ + { + "start": { + "project_id": "team/demo", + "id": run_call, + "op_name": "workflow", + "display_name": "workflow", + "trace_id": WORKFLOW_RUN_ID, + "parent_id": None, + "thread_id": WORKFLOW_RUN_ID, + "turn_id": run_call, + "started_at": "2026-10-01T08:00:00+00:00", + "attributes": { + "workflow_run_id": WORKFLOW_RUN_ID, + "app_id": "app-1", + "dify_span_kind": "chain", + "dify_span_id": run.span_id, + }, + "inputs": {"topic": "tracing"}, + } + }, + { + "start": { + "project_id": "team/demo", + "id": llm_call, + "op_name": "llm_Summarize", + "display_name": "llm_Summarize", + "trace_id": WORKFLOW_RUN_ID, + "parent_id": run_call, + "thread_id": WORKFLOW_RUN_ID, + "turn_id": run_call, + "started_at": "2026-10-01T08:00:01+00:00", + "attributes": { + "node_execution_id": "node-execution-llm", + "node_type": "llm", + "dify_span_kind": "llm", + "dify_span_id": llm.span_id, + }, + "inputs": {"input": [{"role": "user", "content": "Summarize tracing"}]}, + } + }, + { + "start": { + "project_id": "team/demo", + "id": tool_call, + "op_name": "tool_Search", + "display_name": "tool_Search", + "trace_id": WORKFLOW_RUN_ID, + "parent_id": llm_call, + "thread_id": WORKFLOW_RUN_ID, + "turn_id": run_call, + "started_at": "2026-10-01T08:00:03+00:00", + "attributes": { + "node_execution_id": "node-execution-tool", + "node_type": "tool", + "dify_span_kind": "tool", + "dify_span_id": tool.span_id, + }, + "inputs": {"query": "tracing"}, + } + }, + { + "end": { + "project_id": "team/demo", + "id": tool_call, + "ended_at": "2026-10-01T08:00:05+00:00", + "exception": "Tool failed", + "output": {}, + "summary": {}, + } + }, + { + "end": { + "project_id": "team/demo", + "id": llm_call, + "ended_at": "2026-10-01T08:00:03+00:00", + "exception": None, + "output": {"text": "Tracing records execution."}, + "summary": { + "usage": { + "gpt-4o": {"prompt_tokens": 30, "completion_tokens": 10, "total_tokens": 40, "requests": 1} + } + }, + } + }, + { + "end": { + "project_id": "team/demo", + "id": run_call, + "ended_at": "2026-10-01T08:00:05+00:00", + "exception": "Tool failed", + "output": {}, + "summary": {}, + } + }, + ] + + +def test_fragment_joins_the_message_call_without_naming_a_turn() -> None: + message = message_trace() + message_provider, message_egress = _provider(*_acknowledgments(message)) + message_provider.send(CONFIG, message) + message_start = _body(message_egress.requests[0])["start"] + + fragment = tool_fragment() + provider, egress = _provider(*_acknowledgments(fragment)) + result = provider.send(CONFIG, fragment) + + tool = fragment.root + tool_call = _call_id(fragment, tool.span_id) + start, end = (_body(request) for request in egress.requests) + assert start == { + "start": { + "project_id": "team/demo", + "id": tool_call, + "op_name": "weather", + "display_name": "weather", + "trace_id": message_start["trace_id"], + "parent_id": message_start["id"], + "thread_id": "conversation-1", + "turn_id": None, + "started_at": "2026-10-01T08:00:00.500000+00:00", + "attributes": {"message_id": MESSAGE_ID, "dify_span_kind": "tool", "dify_span_id": tool.span_id}, + "inputs": {"city": "Berlin"}, + } + } + assert end == { + "end": { + "project_id": "team/demo", + "id": tool_call, + "ended_at": "2026-10-01T08:00:01+00:00", + "exception": None, + "output": "Sunny", + "summary": {}, + } + } + assert result == SendResult( + trace_id=MESSAGE_ID, + trace_url=f"{CALLS_URL}/{message_start['id']}", + span_urls={tool.span_id: f"{CALLS_URL}/{tool_call}"}, + ) + + +def test_unit_without_session_has_no_thread_or_turn() -> None: + trace = message_trace().model_copy(update={"session_id": None}) + provider, egress = _provider(*_acknowledgments(trace)) + provider.send(CONFIG, trace) + starts = [_body(request)["start"] for request in egress.requests[:2]] + assert [(start["thread_id"], start["turn_id"]) for start in starts] == [(None, None), (None, None)] + + +@pytest.mark.parametrize( + ("inputs", "expected"), + [(None, {}), ("What is Dify?", {"input": "What is Dify?"}), ({"query": "q"}, {"query": "q"})], +) +def test_call_inputs_are_objects(inputs: JsonValue, expected: dict[str, JsonValue]) -> None: + trace = _with_root(message_trace(), inputs=inputs) + provider, egress = _provider(*_acknowledgments(trace)) + provider.send(CONFIG, trace) + assert _body(egress.requests[0])["start"]["inputs"] == expected + + +def test_names_are_cut_to_the_trace_server_limits() -> None: + trace = _with_root(message_trace(), name="n" * 2000) + provider, egress = _provider(*_acknowledgments(trace)) + provider.send(CONFIG, trace) + start = _body(egress.requests[0])["start"] + assert (start["op_name"], start["display_name"]) == ("n" * 128, "n" * 1024) + + +def test_redelivery_sends_identical_requests() -> None: + first, first_egress = _provider(*_acknowledgments(workflow_trace())) + second, second_egress = _provider(*_acknowledgments(workflow_trace())) + assert first.send(CONFIG, workflow_trace()) == second.send(CONFIG, workflow_trace()) + assert [request.content for request in first_egress.requests] == [ + request.content for request in second_egress.requests + ] + + +def test_send_resolves_an_unset_entity_with_the_configured_key() -> None: + trace = message_trace() + viewer = response(200, {"data": {"viewer": {"id": "viewer-1", "entity": "default-team"}}}) + provider, egress = _provider(viewer, *_acknowledgments(trace)) + result = provider.send(WeaveConfig(api_key="weave-key", project="demo", host="https://wandb.internal/w/"), trace) + + lookup, *calls = egress.requests + assert lookup.url == "https://wandb.internal/w/graphql" + assert lookup.headers["Authorization"] == AUTHORIZATION + assert "mutation" not in _body(lookup)["query"] + assert {_body(request)["start"]["project_id"] for request in calls[:2]} == {"default-team/demo"} + root_call = _call_id(trace, trace.root.span_id) + assert result.trace_url == f"https://wandb.internal/w/default-team/demo/r/call/{root_call}" + + +def test_send_takes_the_entity_from_a_project_path() -> None: + provider, egress = _provider(*_acknowledgments(message_trace())) + provider.send(WeaveConfig(api_key="weave-key", project="team/demo"), message_trace()) + assert _body(egress.requests[0])["start"]["project_id"] == "team/demo" + + +@pytest.mark.parametrize( + ("outcome", "expected"), + [ + (response(200, {"data": {"viewer": None}}), ExportRejectedError), + (response(200, {"data": {"viewer": {"id": "viewer-1", "entity": "team/other"}}}), ExportRejectedError), + (response(200, {"errors": [{"message": "DO_NOT_EXPOSE"}]}), ExportRejectedError), + (response(200, b"not json"), ExportRejectedError), + (response(401, {"error": "DO_NOT_EXPOSE"}), ExportRejectedError), + (response(503, {}), ExportRetryableError), + (EgressError("DO_NOT_EXPOSE", request_sent=True), ExportRetryableError), + ], +) +def test_entity_lookup_failures_write_no_calls(outcome: EgressResponse | Exception, expected: type[Exception]) -> None: + provider, egress = _provider(outcome) + with pytest.raises(expected) as caught: + provider.send(WeaveConfig(api_key="weave-key", project="demo"), message_trace()) + assert "DO_NOT_EXPOSE" not in str(caught.value) + assert len(egress.requests) == 1 + + +@pytest.mark.parametrize(("entity", "project"), [("team", "other/demo"), (None, "team/sub/demo"), ("team", "..")]) +def test_send_rejects_unusable_names_before_any_request(entity: str | None, project: str) -> None: + provider, egress = _provider() + with pytest.raises(ExportRejectedError): + provider.send(WeaveConfig(api_key="weave-key", entity=entity, project=project), message_trace()) + assert egress.requests == [] + + +def test_numbers_json_cannot_represent_are_rejected_before_any_request() -> None: + provider, egress = _provider() + with pytest.raises(ExportRejectedError): + provider.send(CONFIG, _with_root(message_trace(), outputs=float("nan"))) + assert egress.requests == [] + + +@pytest.mark.parametrize( + ("outcome", "expected"), + [ + (response(503, {"detail": "DO_NOT_EXPOSE"}), ExportRetryableError), + (response(429, {}), ExportRetryableError), + (response(401, {"detail": "DO_NOT_EXPOSE"}), ExportRejectedError), + (EgressError("DO_NOT_EXPOSE", request_sent=False), ExportRetryableError), + (EgressError("DO_NOT_EXPOSE", request_sent=True), ExportOutcomeUnknownError), + (EgressBlockedError("DO_NOT_EXPOSE"), ExportRejectedError), + (response(200, b"not json"), ExportOutcomeUnknownError), + ], +) +def test_first_call_failures_are_classified(outcome: EgressResponse | Exception, expected: type[Exception]) -> None: + provider, egress = _provider(outcome) + with pytest.raises(expected) as caught: + provider.send(CONFIG, message_trace()) + assert "DO_NOT_EXPOSE" not in str(caught.value) + assert "weave-key" not in str(caught.value) + assert len(egress.requests) == 1 + + +@pytest.mark.parametrize("failing_request", [1, 2, 3]) +@pytest.mark.parametrize( + "outcome", + [response(503, {}), response(400, {}), EgressError("connection refused", request_sent=False)], +) +def test_failure_after_an_accepted_call_is_outcome_unknown( + failing_request: int, outcome: EgressResponse | Exception +) -> None: + trace = message_trace() + provider, egress = _provider(*_acknowledgments(trace)[:failing_request], outcome) + with pytest.raises(ExportOutcomeUnknownError): + provider.send(CONFIG, trace) + assert len(egress.requests) == failing_request + 1 + + +@pytest.mark.parametrize("field", ["id", "trace_id"]) +@pytest.mark.parametrize("failing_start", [0, 1]) +def test_call_start_must_echo_the_submitted_ids(field: str, failing_start: int) -> None: + trace = message_trace() + acknowledgment = {"id": _call_id(trace, trace.spans[failing_start].span_id), "trace_id": MESSAGE_ID} + acknowledgment[field] = "00000000-0000-4000-8000-000000000000" + provider, egress = _provider(*_acknowledgments(trace)[:failing_start], response(200, acknowledgment)) + with pytest.raises(ExportOutcomeUnknownError): + provider.send(CONFIG, trace) + assert len(egress.requests) == failing_start + 1 + + +def test_call_start_may_spell_the_echoed_ids_differently() -> None: + trace = message_trace() + starts = [ + response(200, {"id": UUID(_call_id(trace, span.span_id)).hex.upper(), "trace_id": f"{{{MESSAGE_ID}}}"}) + for span in trace.spans + ] + provider, egress = _provider(*starts, *(response(200, {}) for _ in trace.spans)) + provider.send(CONFIG, trace) + assert len(egress.requests) == 4 + + +def test_the_weave_package_is_never_imported() -> None: + # Importing ``weave`` patches PIL for the whole process. + check = "import sys, dify_trace_weave.provider; sys.exit('weave' in sys.modules)" + assert subprocess.run([sys.executable, "-c", check], check=False).returncode == 0 diff --git a/api/providers/trace/trace-weave/tests/unit_tests/src/dify_trace_weave/test_weave_trace.py b/api/providers/trace/trace-weave/tests/unit_tests/src/dify_trace_weave/test_weave_trace.py deleted file mode 100644 index b836b06a861f31..00000000000000 --- a/api/providers/trace/trace-weave/tests/unit_tests/src/dify_trace_weave/test_weave_trace.py +++ /dev/null @@ -1,1109 +0,0 @@ -"""Comprehensive tests for dify_trace_weave.weave_trace module.""" - -from __future__ import annotations - -from datetime import UTC, datetime, timedelta -from types import SimpleNamespace -from unittest.mock import MagicMock, patch - -import pytest -from dify_trace_weave.config import WeaveConfig -from dify_trace_weave.configuration import WeaveConfigurationError -from dify_trace_weave.entities.weave_trace_entity import WeaveTraceModel -from dify_trace_weave.weave_trace import WeaveDataTrace -from weave.trace_server.trace_server_interface import CallStartRes, TraceStatus - -from core.ops.entities.trace_entity import ( - DatasetRetrievalTraceInfo, - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - SuggestedQuestionTraceInfo, - ToolTraceInfo, - TraceTaskName, - WorkflowTraceInfo, -) -from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey - -# ── Helpers ────────────────────────────────────────────────────────────────── - - -def _dt() -> datetime: - return datetime(2024, 1, 1, 0, 0, 0, tzinfo=UTC) - - -def _make_weave_config(**overrides) -> WeaveConfig: - defaults = { - "api_key": "wv-api-key", - "project": "my-project", - "entity": "my-entity", - "host": None, - } - defaults.update(overrides) - return WeaveConfig(**defaults) - - -def _make_workflow_trace_info(**overrides) -> WorkflowTraceInfo: - defaults = { - "workflow_id": "wf-id", - "tenant_id": "tenant-1", - "workflow_run_id": "run-1", - "workflow_run_elapsed_time": 1.0, - "workflow_run_status": "succeeded", - "workflow_run_inputs": {"key": "val"}, - "workflow_run_outputs": {"answer": "42"}, - "workflow_run_version": "v1", - "total_tokens": 10, - "file_list": [], - "query": "hello", - "metadata": {"user_id": "u1", "app_id": "app-1"}, - "start_time": _dt(), - "end_time": _dt() + timedelta(seconds=1), - } - defaults.update(overrides) - return WorkflowTraceInfo(**defaults) - - -def _make_message_trace_info(**overrides) -> MessageTraceInfo: - msg_data = MagicMock() - msg_data.id = "msg-1" - msg_data.from_account_id = "acc-1" - msg_data.from_end_user_id = None - defaults = { - "conversation_model": "chat", - "message_tokens": 5, - "answer_tokens": 10, - "total_tokens": 15, - "conversation_mode": "chat", - "metadata": {"conversation_id": "c1"}, - "message_id": "msg-1", - "message_data": msg_data, - "inputs": {"prompt": "hi"}, - "outputs": "ok", - "start_time": _dt(), - "end_time": _dt() + timedelta(seconds=1), - "error": None, - } - defaults.update(overrides) - return MessageTraceInfo(**defaults) - - -def _make_moderation_trace_info(**overrides) -> ModerationTraceInfo: - defaults = { - "flagged": False, - "action": "allow", - "preset_response": "", - "query": "test", - "metadata": {"user_id": "u1"}, - "message_id": "msg-1", - } - defaults.update(overrides) - return ModerationTraceInfo(**defaults) - - -def _make_suggested_question_trace_info(**overrides) -> SuggestedQuestionTraceInfo: - defaults = { - "suggested_question": ["q1", "q2"], - "level": "info", - "total_tokens": 5, - "metadata": {"user_id": "u1"}, - "message_id": "msg-1", - "message_data": SimpleNamespace(created_at=_dt(), updated_at=_dt()), - "inputs": {"i": 1}, - "start_time": _dt(), - "end_time": _dt() + timedelta(seconds=1), - "error": None, - } - defaults.update(overrides) - return SuggestedQuestionTraceInfo(**defaults) - - -def _make_dataset_retrieval_trace_info(**overrides) -> DatasetRetrievalTraceInfo: - msg_data = MagicMock() - msg_data.created_at = _dt() - msg_data.updated_at = _dt() - defaults = { - "metadata": {"user_id": "u1"}, - "message_id": "msg-1", - "message_data": msg_data, - "inputs": "query", - "documents": [{"content": "doc"}], - "start_time": _dt(), - "end_time": _dt() + timedelta(seconds=1), - } - defaults.update(overrides) - return DatasetRetrievalTraceInfo(**defaults) - - -def _make_tool_trace_info(**overrides) -> ToolTraceInfo: - defaults = { - "tool_name": "my_tool", - "tool_inputs": {"x": 1}, - "tool_outputs": "output", - "tool_config": {"desc": "d"}, - "tool_parameters": {"p": "v"}, - "time_cost": 0.5, - "metadata": {"user_id": "u1"}, - "message_id": "msg-1", - "inputs": {"i": "v"}, - "outputs": {"o": "v"}, - "start_time": _dt(), - "end_time": _dt() + timedelta(seconds=1), - "error": None, - } - defaults.update(overrides) - return ToolTraceInfo(**defaults) - - -def _make_generate_name_trace_info(**overrides) -> GenerateNameTraceInfo: - defaults = { - "tenant_id": "t1", - "metadata": {"user_id": "u1"}, - "message_id": "msg-1", - "inputs": {"i": 1}, - "outputs": {"name": "test"}, - "start_time": _dt(), - "end_time": _dt() + timedelta(seconds=1), - } - defaults.update(overrides) - return GenerateNameTraceInfo(**defaults) - - -def _make_node(**overrides): - """Create a mock workflow node execution object.""" - defaults = { - "id": "node-1", - "title": "Node Title", - "node_type": BuiltinNodeTypes.CODE, - "status": "succeeded", - "inputs": {"key": "value"}, - "outputs": {"result": "ok"}, - "created_at": _dt(), - "elapsed_time": 1.0, - "process_data": None, - "metadata": {}, - } - defaults.update(overrides) - return SimpleNamespace(**defaults) - - -# ── Fixtures ───────────────────────────────────────────────────────────────── - - -@pytest.fixture -def mock_server(): - with patch("dify_trace_weave.weave_trace.TraceServer") as factory: - server = factory.return_value - server.call_start.side_effect = lambda req: CallStartRes(id=req.start.id, trace_id=req.start.trace_id) - yield factory - - -@pytest.fixture -def trace_instance(mock_server): - return WeaveDataTrace(_make_weave_config()) - - -@pytest.fixture -def trace_instance_with_host(mock_server): - return WeaveDataTrace(_make_weave_config(host="https://my.wandb.host")) - - -class TestInit: - def test_init_has_explicit_endpoint_and_auth_without_batching(self, mock_server): - instance = WeaveDataTrace(_make_weave_config(endpoint="https://trace.test/prefix")) - mock_server.assert_called_once_with("https://trace.test/prefix", should_batch=False, auth=("api", "wv-api-key")) - assert instance.weave_api_key == "wv-api-key" - assert instance.project_id == "my-entity/my-project" - assert instance.calls == {} - mock_server.return_value.call_start.assert_not_called() - - @pytest.mark.parametrize("host", [None, "https://my.wandb.host", "https://wandb.internal/api"]) - def test_init_preserves_host(self, mock_server, host): - instance = WeaveDataTrace(_make_weave_config(host=host)) - assert instance.host == host - - def test_init_without_entity_does_not_resolve_over_network(self, mock_server): - with patch("dify_trace_weave.weave_trace.resolve_project") as resolve: - instance = WeaveDataTrace(_make_weave_config(entity=None)) - resolve.assert_not_called() - assert instance.get_project_url() == "https://wandb.ai" - resolve.assert_not_called() - - def test_qualified_project_does_not_require_default_entity(self, mock_server): - instance = WeaveDataTrace(_make_weave_config(entity=None, project="team/project")) - assert instance.project_id == "team/project" - assert instance.get_project_url() == "https://wandb.ai/team/project" - - def test_init_files_url_from_env(self, mock_server, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("FILES_URL", "http://files.example.com") - assert WeaveDataTrace(_make_weave_config()).file_base_url == "http://files.example.com" - - def test_init_files_url_default(self, mock_server, monkeypatch: pytest.MonkeyPatch): - monkeypatch.delenv("FILES_URL", raising=False) - assert WeaveDataTrace(_make_weave_config()).file_base_url == "http://127.0.0.1:5001" - - -class TestGetProjectUrl: - def test_get_project_url_with_entity(self, trace_instance): - assert trace_instance.get_project_url() == "https://wandb.ai/my-entity/my-project" - - def test_get_project_url_preserves_private_prefix(self, mock_server): - instance = WeaveDataTrace(_make_weave_config(host="https://wandb.internal/prefix")) - assert instance.get_project_url() == "https://wandb.internal/prefix/my-entity/my-project" - - -# ── TestTraceDispatcher ───────────────────────────────────────────────────── - - -class TestTraceDispatcher: - def test_dispatches_workflow_trace(self, trace_instance): - with patch.object(trace_instance, "workflow_trace") as mock_wt: - trace_instance.trace(_make_workflow_trace_info()) - mock_wt.assert_called_once() - - def test_dispatches_message_trace(self, trace_instance): - with patch.object(trace_instance, "message_trace") as mock_mt: - trace_instance.trace(_make_message_trace_info()) - mock_mt.assert_called_once() - - def test_dispatches_moderation_trace(self, trace_instance): - with patch.object(trace_instance, "moderation_trace") as mock_mod: - msg_data = MagicMock() - msg_data.created_at = _dt() - trace_instance.trace(_make_moderation_trace_info(message_data=msg_data)) - mock_mod.assert_called_once() - - def test_dispatches_suggested_question_trace(self, trace_instance): - with patch.object(trace_instance, "suggested_question_trace") as mock_sq: - trace_instance.trace(_make_suggested_question_trace_info()) - mock_sq.assert_called_once() - - def test_dispatches_dataset_retrieval_trace(self, trace_instance): - with patch.object(trace_instance, "dataset_retrieval_trace") as mock_dr: - trace_instance.trace(_make_dataset_retrieval_trace_info()) - mock_dr.assert_called_once() - - def test_dispatches_tool_trace(self, trace_instance): - with patch.object(trace_instance, "tool_trace") as mock_tool: - trace_instance.trace(_make_tool_trace_info()) - mock_tool.assert_called_once() - - def test_dispatches_generate_name_trace(self, trace_instance): - with patch.object(trace_instance, "generate_name_trace") as mock_gn: - trace_instance.trace(_make_generate_name_trace_info()) - mock_gn.assert_called_once() - - -# ── TestNormalizeTime ───────────────────────────────────────────────────────── - - -class TestNormalizeTime: - def test_none_returns_utc_now(self, trace_instance): - now_before = datetime.now(UTC) - result = trace_instance._normalize_time(None) - now_after = datetime.now(UTC) - assert result.tzinfo is not None - assert now_before <= result <= now_after - - def test_naive_datetime_gets_utc(self, trace_instance): - naive = datetime(2024, 6, 15, 12, 0, 0) - result = trace_instance._normalize_time(naive) - assert result.tzinfo == UTC - assert result.year == 2024 - assert result.month == 6 - - def test_aware_datetime_unchanged(self, trace_instance): - aware = datetime(2024, 6, 15, 12, 0, 0, tzinfo=UTC) - result = trace_instance._normalize_time(aware) - assert result == aware - assert result.tzinfo == UTC - - -# ── TestStartCall ───────────────────────────────────────────────────────────── - - -class TestStartCall: - def test_start_call_basic(self, trace_instance): - """Test basic start_call stores call metadata.""" - run = WeaveTraceModel( - id="run-1", - op="test-op", - inputs={"key": "val"}, - attributes={"trace_id": "t-1", "start_time": _dt()}, - ) - trace_instance.start_call(run) - - assert "run-1" in trace_instance.calls - assert trace_instance.calls["run-1"]["trace_id"] == "t-1" - assert trace_instance.calls["run-1"]["parent_id"] is None - trace_instance.server.call_start.assert_called_once() - - def test_start_call_with_parent(self, trace_instance): - """Test start_call records parent_run_id.""" - run = WeaveTraceModel( - id="child-1", - op="child-op", - inputs={}, - attributes={"trace_id": "t-1", "start_time": _dt()}, - ) - trace_instance.start_call(run, parent_run_id="parent-1") - - assert trace_instance.calls["child-1"]["parent_id"] == "parent-1" - - def test_start_call_none_inputs_becomes_empty_dict(self, trace_instance): - """Test that None inputs is normalized to {}.""" - run = WeaveTraceModel( - id="run-2", - op="op", - inputs=None, - attributes={"trace_id": "t-2", "start_time": _dt()}, - ) - trace_instance.start_call(run) - call_args = trace_instance.server.call_start.call_args - req = call_args[0][0] - assert req.start.inputs == {} - - def test_start_call_non_dict_inputs_becomes_str_dict(self, trace_instance): - """Test that non-dict inputs is wrapped as string.""" - run = WeaveTraceModel( - id="run-3", - op="op", - inputs="some string input", - attributes={"trace_id": "t-3", "start_time": _dt()}, - ) - trace_instance.start_call(run) - call_args = trace_instance.server.call_start.call_args - req = call_args[0][0] - # String inputs gets converted by validator to a dict - assert isinstance(req.start.inputs, dict) - - def test_start_call_none_attributes_becomes_empty_dict(self, trace_instance): - """Test that None attributes is handled properly.""" - run = WeaveTraceModel( - id="run-4", - op="op", - inputs={}, - attributes=None, - ) - trace_instance.start_call(run) - # trace_id should fall back to run_data.id - assert trace_instance.calls["run-4"]["trace_id"] == "run-4" - - def test_start_call_non_dict_attributes_becomes_dict(self, trace_instance): - """Test that non-dict attributes is wrapped.""" - run = WeaveTraceModel( - id="run-5", - op="op", - inputs={}, - attributes=None, - ) - # Manually override after construction - run.attributes = "some-attr-string" - trace_instance.start_call(run) - call_args = trace_instance.server.call_start.call_args - req = call_args[0][0] - assert isinstance(req.start.attributes, dict) - assert req.start.attributes == {"attributes": "some-attr-string"} - - def test_start_call_trace_id_falls_back_to_run_id(self, trace_instance): - """When trace_id not in attributes, falls back to run_data.id.""" - run = WeaveTraceModel( - id="run-6", - op="op", - inputs={}, - attributes={"start_time": _dt()}, - ) - trace_instance.start_call(run) - assert trace_instance.calls["run-6"]["trace_id"] == "run-6" - - -# ── TestFinishCall ────────────────────────────────────────────────────────── - - -class TestFinishCall: - def _setup_call(self, trace_instance, run_id="run-1", trace_id="t-1"): - """Helper: register a call so finish_call can find it.""" - trace_instance.calls[run_id] = {"trace_id": trace_id, "parent_id": None} - - def test_finish_call_success(self, trace_instance): - """Test finish_call sends call_end with SUCCESS status.""" - self._setup_call(trace_instance) - run = WeaveTraceModel( - id="run-1", - op="op", - inputs={}, - outputs={"result": "ok"}, - attributes={"start_time": _dt(), "end_time": _dt() + timedelta(seconds=1)}, - exception=None, - ) - trace_instance.finish_call(run) - trace_instance.server.call_end.assert_called_once() - call_args = trace_instance.server.call_end.call_args - req = call_args[0][0] - assert req.end.summary["status_counts"][TraceStatus.SUCCESS] == 1 - assert req.end.summary["status_counts"][TraceStatus.ERROR] == 0 - assert req.end.exception is None - - def test_finish_call_with_error(self, trace_instance): - """Test finish_call sends call_end with ERROR status when exception is set.""" - self._setup_call(trace_instance) - run = WeaveTraceModel( - id="run-1", - op="op", - inputs={}, - outputs={}, - attributes={"start_time": _dt(), "end_time": _dt() + timedelta(seconds=1)}, - exception="Something broke", - ) - trace_instance.finish_call(run) - call_args = trace_instance.server.call_end.call_args - req = call_args[0][0] - assert req.end.summary["status_counts"][TraceStatus.ERROR] == 1 - assert req.end.summary["status_counts"][TraceStatus.SUCCESS] == 0 - assert req.end.exception == "Something broke" - - def test_finish_call_missing_id_raises(self, trace_instance): - """Test finish_call raises ValueError when call id not found.""" - run = WeaveTraceModel( - id="nonexistent", - op="op", - inputs={}, - ) - with pytest.raises(ValueError, match="Call with id nonexistent not found"): - trace_instance.finish_call(run) - - def test_finish_call_elapsed_negative_clamped_to_zero(self, trace_instance): - """Test that negative elapsed time is clamped to 0.""" - self._setup_call(trace_instance) - run = WeaveTraceModel( - id="run-1", - op="op", - inputs={}, - attributes={ - "start_time": _dt() + timedelta(seconds=5), - "end_time": _dt(), # end before start - }, - ) - trace_instance.finish_call(run) - call_args = trace_instance.server.call_end.call_args - req = call_args[0][0] - assert req.end.summary["weave"]["latency_ms"] == 0 - - def test_finish_call_none_attributes(self, trace_instance): - """Test finish_call handles None attributes.""" - self._setup_call(trace_instance) - run = WeaveTraceModel( - id="run-1", - op="op", - inputs={}, - attributes=None, - ) - trace_instance.finish_call(run) - trace_instance.server.call_end.assert_called_once() - - def test_finish_call_non_dict_attributes(self, trace_instance): - """Test finish_call handles non-dict attributes.""" - self._setup_call(trace_instance) - run = WeaveTraceModel( - id="run-1", - op="op", - inputs={}, - attributes=None, - ) - run.attributes = "some string attr" - trace_instance.finish_call(run) - trace_instance.server.call_end.assert_called_once() - - -# ── TestWorkflowTrace ───────────────────────────────────────────────────────── - - -class TestWorkflowTrace: - def _setup_repo(self, monkeypatch: pytest.MonkeyPatch, nodes=None): - """Helper to patch session/repo dependencies.""" - if nodes is None: - nodes = [] - - repo = MagicMock() - repo.get_by_workflow_execution.return_value = nodes - - mock_factory = MagicMock() - mock_factory.create_workflow_node_execution_repository.return_value = repo - - monkeypatch.setattr("dify_trace_weave.weave_trace.DifyCoreRepositoryFactory", mock_factory) - monkeypatch.setattr("dify_trace_weave.weave_trace.sessionmaker", lambda bind: MagicMock()) - monkeypatch.setattr("dify_trace_weave.weave_trace.db", MagicMock(engine="engine")) - return repo - - def test_workflow_trace_no_nodes_no_message_id(self, trace_instance, monkeypatch: pytest.MonkeyPatch): - """Workflow trace with no nodes and no message_id.""" - self._setup_repo(monkeypatch, nodes=[]) - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_workflow_trace_info(message_id=None) - trace_instance.workflow_trace(trace_info) - - # Only workflow run: start_call and finish_call each called once - assert trace_instance.start_call.call_count == 1 - assert trace_instance.finish_call.call_count == 1 - - def test_workflow_trace_with_message_id(self, trace_instance, monkeypatch: pytest.MonkeyPatch): - """Workflow trace with message_id creates both message and workflow runs.""" - self._setup_repo(monkeypatch, nodes=[]) - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_workflow_trace_info(message_id="msg-1") - trace_instance.workflow_trace(trace_info) - - # message run + workflow run = 2 start_call / finish_call - assert trace_instance.start_call.call_count == 2 - assert trace_instance.finish_call.call_count == 2 - - def test_workflow_trace_with_node_execution(self, trace_instance, monkeypatch: pytest.MonkeyPatch): - """Workflow trace iterates node executions and creates node runs.""" - node = _make_node( - id="node-1", - node_type=BuiltinNodeTypes.CODE, - inputs={"k": "v"}, - outputs={"r": "ok"}, - elapsed_time=0.5, - created_at=_dt(), - metadata={WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: 5}, - ) - self._setup_repo(monkeypatch, nodes=[node]) - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_workflow_trace_info(message_id=None) - trace_instance.workflow_trace(trace_info) - - # workflow run + node run = 2 calls - assert trace_instance.start_call.call_count == 2 - - def test_workflow_trace_with_llm_node(self, trace_instance, monkeypatch: pytest.MonkeyPatch): - """LLM node uses process_data prompts as inputs.""" - node = _make_node( - node_type=BuiltinNodeTypes.LLM, - process_data={ - "prompts": [{"role": "user", "content": "hi"}], - "model_mode": "chat", - "model_provider": "openai", - "model_name": "gpt-4", - }, - inputs={"key": "val"}, - ) - self._setup_repo(monkeypatch, nodes=[node]) - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_workflow_trace_info(message_id=None) - trace_instance.workflow_trace(trace_info) - - # Check node start_call was called with prompts input - node_call_args = trace_instance.start_call.call_args_list[-1] - node_run = node_call_args[0][0] - # WeaveTraceModel validator wraps list prompts into {"messages": [...]} - # The key "messages" should be present (validator transforms the list) - assert "messages" in node_run.inputs - - def test_workflow_trace_with_non_llm_node_uses_inputs(self, trace_instance, monkeypatch: pytest.MonkeyPatch): - """Non-LLM node uses node_execution.inputs directly.""" - node = _make_node( - node_type=BuiltinNodeTypes.TOOL, - inputs={"tool_input": "val"}, - process_data=None, - ) - self._setup_repo(monkeypatch, nodes=[node]) - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_workflow_trace_info(message_id=None) - trace_instance.workflow_trace(trace_info) - - # node run inputs should be from node.inputs; validator adds usage_metadata + file_list - node_call_args = trace_instance.start_call.call_args_list[-1] - node_run = node_call_args[0][0] - assert node_run.inputs.get("tool_input") == "val" - - def test_workflow_trace_missing_app_id_raises(self, trace_instance, monkeypatch: pytest.MonkeyPatch): - """Raises ValueError when app_id is missing from metadata.""" - monkeypatch.setattr("dify_trace_weave.weave_trace.sessionmaker", lambda bind: MagicMock()) - monkeypatch.setattr("dify_trace_weave.weave_trace.db", MagicMock(engine="engine")) - - trace_info = _make_workflow_trace_info( - message_id=None, - metadata={"user_id": "u1"}, # no app_id - ) - - with pytest.raises(ValueError, match="No app_id found in trace_info metadata"): - trace_instance.workflow_trace(trace_info) - - def test_workflow_trace_start_time_none_defaults_to_now(self, trace_instance, monkeypatch: pytest.MonkeyPatch): - """start_time defaults to datetime.now() when None.""" - self._setup_repo(monkeypatch, nodes=[]) - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_workflow_trace_info(message_id=None, start_time=None) - trace_instance.workflow_trace(trace_info) - - assert trace_instance.start_call.call_count == 1 - - def test_workflow_trace_node_created_at_none(self, trace_instance, monkeypatch: pytest.MonkeyPatch): - """Node with created_at=None uses datetime.now().""" - node = _make_node(created_at=None, elapsed_time=0.5) - self._setup_repo(monkeypatch, nodes=[node]) - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_workflow_trace_info(message_id=None) - trace_instance.workflow_trace(trace_info) - assert trace_instance.start_call.call_count == 2 - - def test_workflow_trace_chat_mode_llm_node_adds_provider(self, trace_instance, monkeypatch: pytest.MonkeyPatch): - """Chat mode LLM node adds ls_provider and ls_model_name to attributes.""" - node = _make_node( - node_type=BuiltinNodeTypes.LLM, - process_data={"model_mode": "chat", "model_provider": "openai", "model_name": "gpt-4", "prompts": []}, - ) - self._setup_repo(monkeypatch, nodes=[node]) - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - start_calls = [] - - def capture_start(run, parent_run_id=None): - start_calls.append((run, parent_run_id)) - - trace_instance.start_call = capture_start - trace_instance.finish_call = MagicMock() - - trace_info = _make_workflow_trace_info(message_id=None) - trace_instance.workflow_trace(trace_info) - - # Last start call is the node run - node_run, _ = start_calls[-1] - assert node_run.attributes.get("ls_provider") == "openai" - assert node_run.attributes.get("ls_model_name") == "gpt-4" - - def test_workflow_trace_nodes_sorted_by_created_at(self, trace_instance, monkeypatch: pytest.MonkeyPatch): - """Nodes are sorted by created_at before processing.""" - node1 = _make_node(id="node-b", created_at=_dt() + timedelta(seconds=2)) - node2 = _make_node(id="node-a", created_at=_dt()) - self._setup_repo(monkeypatch, nodes=[node1, node2]) - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - processed_ids = [] - - def capture_start(run, parent_run_id=None): - processed_ids.append(run.id) - - trace_instance.start_call = capture_start - trace_instance.finish_call = MagicMock() - - trace_info = _make_workflow_trace_info(message_id=None) - trace_instance.workflow_trace(trace_info) - - # First call = workflow run, then node-a, then node-b - assert processed_ids[1] == "node-a" - assert processed_ids[2] == "node-b" - - -# ── TestMessageTrace ────────────────────────────────────────────────────────── - - -class TestMessageTrace: - def test_returns_early_when_no_message_data(self, trace_instance): - """message_trace returns early when message_data is None.""" - trace_info = _make_message_trace_info(message_data=None) - trace_instance.start_call = MagicMock() - trace_instance.message_trace(trace_info) - trace_instance.start_call.assert_not_called() - - def test_basic_message_trace(self, trace_instance, monkeypatch: pytest.MonkeyPatch): - """message_trace creates message run and llm child run.""" - monkeypatch.setattr( - "dify_trace_weave.weave_trace.db.session.get", - lambda model, pk: None, - ) - - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_message_trace_info() - trace_instance.message_trace(trace_info) - - # message run + llm child run - assert trace_instance.start_call.call_count == 2 - assert trace_instance.finish_call.call_count == 2 - - def test_message_trace_with_file_data(self, trace_instance, monkeypatch: pytest.MonkeyPatch): - """message_trace appends file URL to file_list.""" - file_data = MagicMock() - file_data.url = "path/to/file.png" - trace_instance.file_base_url = "http://files.test" - - mock_db = MagicMock() - mock_db.session.get.return_value = None - monkeypatch.setattr("dify_trace_weave.weave_trace.db", mock_db) - - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_message_trace_info( - message_file_data=file_data, - file_list=["existing.txt"], - ) - trace_instance.message_trace(trace_info) - - # The first start_call arg (the message run) should have file in outputs or inputs - message_run = trace_instance.start_call.call_args_list[0][0][0] - assert "http://files.test/path/to/file.png" in message_run.file_list - - def test_message_trace_with_end_user(self, trace_instance, monkeypatch: pytest.MonkeyPatch): - """message_trace looks up end user and sets end_user_id attribute.""" - end_user = MagicMock() - end_user.session_id = "session-xyz" - - mock_db = MagicMock() - mock_db.session.get.return_value = end_user - monkeypatch.setattr("dify_trace_weave.weave_trace.db", mock_db) - - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - msg_data = MagicMock() - msg_data.id = "msg-1" - msg_data.from_account_id = "acc-1" - msg_data.from_end_user_id = "eu-1" - - trace_info = _make_message_trace_info(message_data=msg_data) - trace_instance.message_trace(trace_info) - - message_run = trace_instance.start_call.call_args_list[0][0][0] - assert message_run.attributes.get("end_user_id") == "session-xyz" - - def test_message_trace_no_end_user(self, trace_instance, monkeypatch: pytest.MonkeyPatch): - """message_trace handles when from_end_user_id is None.""" - mock_db = MagicMock() - mock_db.session.get.return_value = None - monkeypatch.setattr("dify_trace_weave.weave_trace.db", mock_db) - - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - msg_data = MagicMock() - msg_data.id = "msg-1" - msg_data.from_account_id = "acc-1" - msg_data.from_end_user_id = None - - trace_info = _make_message_trace_info(message_data=msg_data) - trace_instance.message_trace(trace_info) - assert trace_instance.start_call.call_count == 2 - - def test_message_trace_trace_id_fallback_to_message_id(self, trace_instance, monkeypatch: pytest.MonkeyPatch): - """trace_id falls back to message_id when trace_id is None.""" - mock_db = MagicMock() - mock_db.session.get.return_value = None - monkeypatch.setattr("dify_trace_weave.weave_trace.db", mock_db) - - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_message_trace_info(trace_id=None) - trace_instance.message_trace(trace_info) - - message_run = trace_instance.start_call.call_args_list[0][0][0] - assert message_run.id == "msg-1" - - def test_message_trace_file_list_none(self, trace_instance, monkeypatch: pytest.MonkeyPatch): - """message_trace handles file_list=None gracefully.""" - mock_db = MagicMock() - mock_db.session.get.return_value = None - monkeypatch.setattr("dify_trace_weave.weave_trace.db", mock_db) - - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_message_trace_info(file_list=None, message_file_data=None) - trace_instance.message_trace(trace_info) - assert trace_instance.start_call.call_count == 2 - - -# ── TestModerationTrace ─────────────────────────────────────────────────────── - - -class TestModerationTrace: - def test_returns_early_when_no_message_data(self, trace_instance): - """moderation_trace returns early when message_data is None.""" - trace_info = _make_moderation_trace_info(message_data=None) - trace_instance.start_call = MagicMock() - trace_instance.moderation_trace(trace_info) - trace_instance.start_call.assert_not_called() - - def test_basic_moderation_trace(self, trace_instance): - """moderation_trace creates a run with correct outputs.""" - msg_data = MagicMock() - msg_data.created_at = _dt() - msg_data.updated_at = _dt() - - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_moderation_trace_info( - message_data=msg_data, - start_time=_dt(), - end_time=_dt() + timedelta(seconds=1), - action="block", - flagged=True, - preset_response="blocked", - ) - trace_instance.moderation_trace(trace_info) - - trace_instance.start_call.assert_called_once() - trace_instance.finish_call.assert_called_once() - - run = trace_instance.start_call.call_args[0][0] - assert run.outputs["action"] == "block" - assert run.outputs["flagged"] is True - - def test_moderation_trace_with_no_times_uses_message_data_times(self, trace_instance): - """When start/end times are None, uses message_data created_at/updated_at.""" - msg_data = MagicMock() - msg_data.created_at = _dt() - msg_data.updated_at = _dt() + timedelta(seconds=1) - - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_moderation_trace_info( - message_data=msg_data, - start_time=None, - end_time=None, - ) - trace_instance.moderation_trace(trace_info) - trace_instance.start_call.assert_called_once() - - def test_moderation_trace_trace_id_fallback(self, trace_instance): - """trace_id falls back to message_id when trace_id is None.""" - msg_data = MagicMock() - msg_data.created_at = _dt() - - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_moderation_trace_info( - message_data=msg_data, - trace_id=None, - ) - trace_instance.moderation_trace(trace_info) - - _, kwargs = trace_instance.start_call.call_args - assert kwargs.get("parent_run_id") == "msg-1" - - -# ── TestSuggestedQuestionTrace ──────────────────────────────────────────────── - - -class TestSuggestedQuestionTrace: - def test_returns_early_when_no_message_data(self, trace_instance): - """suggested_question_trace returns early when message_data is None.""" - trace_info = _make_suggested_question_trace_info(message_data=None) - trace_instance.start_call = MagicMock() - trace_instance.suggested_question_trace(trace_info) - trace_instance.start_call.assert_not_called() - - def test_basic_suggested_question_trace(self, trace_instance): - """suggested_question_trace creates a run parented to trace_id.""" - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_suggested_question_trace_info(trace_id="t-1") - trace_instance.suggested_question_trace(trace_info) - - trace_instance.start_call.assert_called_once() - trace_instance.finish_call.assert_called_once() - - _, kwargs = trace_instance.start_call.call_args - assert kwargs.get("parent_run_id") == "t-1" - - def test_suggested_question_trace_trace_id_fallback(self, trace_instance): - """trace_id falls back to message_id when trace_id is None.""" - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_suggested_question_trace_info(trace_id=None) - trace_instance.suggested_question_trace(trace_info) - - _, kwargs = trace_instance.start_call.call_args - assert kwargs.get("parent_run_id") == "msg-1" - - -# ── TestDatasetRetrievalTrace ───────────────────────────────────────────────── - - -class TestDatasetRetrievalTrace: - def test_returns_early_when_no_message_data(self, trace_instance): - """dataset_retrieval_trace returns early when message_data is None.""" - trace_info = _make_dataset_retrieval_trace_info(message_data=None) - trace_instance.start_call = MagicMock() - trace_instance.dataset_retrieval_trace(trace_info) - trace_instance.start_call.assert_not_called() - - def test_basic_dataset_retrieval_trace(self, trace_instance): - """dataset_retrieval_trace creates a run with documents as outputs.""" - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_dataset_retrieval_trace_info( - documents=[{"id": "d1"}, {"id": "d2"}], - trace_id="t-1", - ) - trace_instance.dataset_retrieval_trace(trace_info) - - run = trace_instance.start_call.call_args[0][0] - # WeaveTraceModel validator injects usage_metadata/file_list into dict outputs - assert run.outputs.get("documents") == [{"id": "d1"}, {"id": "d2"}] - _, kwargs = trace_instance.start_call.call_args - assert kwargs.get("parent_run_id") == "t-1" - - def test_dataset_retrieval_trace_trace_id_fallback(self, trace_instance): - """trace_id falls back to message_id when trace_id is None.""" - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_dataset_retrieval_trace_info(trace_id=None) - trace_instance.dataset_retrieval_trace(trace_info) - - _, kwargs = trace_instance.start_call.call_args - assert kwargs.get("parent_run_id") == "msg-1" - - -# ── TestToolTrace ───────────────────────────────────────────────────────────── - - -class TestToolTrace: - def test_basic_tool_trace(self, trace_instance): - """tool_trace creates a run with correct op as tool_name.""" - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_tool_trace_info(trace_id="t-1") - trace_instance.tool_trace(trace_info) - - run = trace_instance.start_call.call_args[0][0] - assert run.op == "my_tool" - # WeaveTraceModel validator injects usage_metadata/file_list into dict inputs - assert run.inputs.get("x") == 1 - - def test_tool_trace_with_file_url(self, trace_instance): - """tool_trace adds file_url to file_list when provided.""" - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_tool_trace_info(file_url="http://files/file.pdf") - trace_instance.tool_trace(trace_info) - - run = trace_instance.start_call.call_args[0][0] - assert "http://files/file.pdf" in run.file_list - - def test_tool_trace_without_file_url(self, trace_instance): - """tool_trace uses empty file_list when file_url is None.""" - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_tool_trace_info(file_url=None) - trace_instance.tool_trace(trace_info) - - run = trace_instance.start_call.call_args[0][0] - assert run.file_list == [] - - def test_tool_trace_trace_id_from_message_id(self, trace_instance): - """trace_id uses message_id fallback.""" - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_tool_trace_info(trace_id=None) - trace_instance.tool_trace(trace_info) - - _, kwargs = trace_instance.start_call.call_args - assert kwargs.get("parent_run_id") == "msg-1" - - def test_tool_trace_message_id_none_uses_conversation_id(self, trace_instance): - """When message_id is None, tries conversation_id attribute.""" - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_tool_trace_info(trace_id=None, message_id=None) - trace_instance.tool_trace(trace_info) - - # No crash; parent_run_id is None since no fallback - _, kwargs = trace_instance.start_call.call_args - # parent_run_id should be None when no message_id and no trace_id - assert kwargs.get("parent_run_id") is None - - -# ── TestGenerateNameTrace ───────────────────────────────────────────────────── - - -class TestGenerateNameTrace: - def test_basic_generate_name_trace(self, trace_instance): - """generate_name_trace creates a run with correct op.""" - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_generate_name_trace_info() - trace_instance.generate_name_trace(trace_info) - - trace_instance.start_call.assert_called_once() - trace_instance.finish_call.assert_called_once() - - run = trace_instance.start_call.call_args[0][0] - assert run.op == str(TraceTaskName.GENERATE_NAME_TRACE) - - def test_generate_name_trace_no_parent(self, trace_instance): - """generate_name_trace has no parent run (no parent_run_id).""" - trace_instance.start_call = MagicMock() - trace_instance.finish_call = MagicMock() - - trace_info = _make_generate_name_trace_info() - trace_instance.generate_name_trace(trace_info) - - _, kwargs = trace_instance.start_call.call_args - # No parent_run_id passed to generate_name start_call - assert kwargs == {} or kwargs.get("parent_run_id") is None - - -# ── TestApiCheck ────────────────────────────────────────────────────────────── - - -class TestApiCheck: - def test_api_check_uses_read_only_verification(self, trace_instance): - with patch("dify_trace_weave.weave_trace.verify_configuration") as verify: - assert trace_instance.api_check() is True - verify.assert_called_once_with(trace_instance._config) - trace_instance.server.call_start.assert_not_called() - trace_instance.server.call_end.assert_not_called() - - def test_api_check_propagates_safe_failure(self, trace_instance): - with patch( - "dify_trace_weave.weave_trace.verify_configuration", - side_effect=WeaveConfigurationError("authentication_failed"), - ): - with pytest.raises(WeaveConfigurationError, match="Weave authentication failed"): - trace_instance.api_check() diff --git a/api/pyproject.toml b/api/pyproject.toml index e226e804b6705a..b0988e728430ee 100644 --- a/api/pyproject.toml +++ b/api/pyproject.toml @@ -10,6 +10,7 @@ dependencies = [ "celery>=5.6.3,<6.0.0", "croniter>=6.2.2,<7.0.0", "dify-agent", + "dify-trace-sdk", "flask>=3.1.3,<4.0.0", "flask-cors>=6.0.5,<7.0.0", "gevent>=26.4.0,<26.5.0", @@ -36,14 +37,18 @@ dependencies = [ "httpx[socks]==0.28.1", "markdown-it-py>=4.0.0,<5.0.0", "opentelemetry-distro==0.65b0", + "opentelemetry-exporter-otlp-proto-grpc>=1.44.0,<2.0.0", + "opentelemetry-exporter-otlp-proto-http>=1.44.0,<2.0.0", "opentelemetry-instrumentation-celery==0.65b0", "opentelemetry-instrumentation-flask==0.65b0", "opentelemetry-instrumentation-httpx==0.65b0", "opentelemetry-instrumentation-redis==0.65b0", "opentelemetry-instrumentation-sqlalchemy==0.65b0", "opentelemetry-propagator-b3>=1.45.0,<2.0.0", + "pydantic-settings>=2.14.2,<3.0.0", "readabilipy==0.3.0", "resend>=2.27.0,<3.0.0", + "sentry-sdk>=2.57.0,<3.0.0", "werkzeug>=3.1.7,<4.0.0", "zstandard==0.25.0", # Emerging: newer and fast-moving, use compatible pins @@ -110,6 +115,7 @@ dify-trace-arize-phoenix = { workspace = true } dify-trace-langfuse = { workspace = true } dify-trace-langsmith = { workspace = true } dify-trace-mlflow = { workspace = true } +dify-trace-sdk = { workspace = true } dify-trace-opik = { workspace = true } dify-trace-tencent = { workspace = true } dify-trace-weave = { workspace = true } diff --git a/api/pytest.ini b/api/pytest.ini index 9958e8e1447d85..5c24e6632d1d0e 100644 --- a/api/pytest.ini +++ b/api/pytest.ini @@ -1,8 +1,7 @@ [pytest] pythonpath = . -# These optional SDK plugins are not used by the backend suite. In particular, -# Opik eagerly imports LiteLLM, which also loads the developer's .env into os.environ. -addopts = -p no:opik -p no:langsmith_plugin --cov=./api --cov-report=json --import-mode=importlib --cov-branch --cov-report=xml +# The LangSmith SDK's pytest plugin is not used by the backend suite. +addopts = -p no:langsmith_plugin --cov=./api --cov-report=json --import-mode=importlib --cov-branch --cov-report=xml env = ANTHROPIC_API_KEY = sk-ant-api11-IamNotARealKeyJustForMockTestKawaiiiiiiiiii-NotBaka-ASkksz AZURE_OPENAI_API_BASE = https://difyai-openai.openai.azure.com diff --git a/api/repositories/api_workflow_node_execution_repository.py b/api/repositories/api_workflow_node_execution_repository.py index 540cf689109797..cbf2834bc2642c 100644 --- a/api/repositories/api_workflow_node_execution_repository.py +++ b/api/repositories/api_workflow_node_execution_repository.py @@ -37,6 +37,15 @@ class WorkflowNodeExecutionSnapshotRow(Protocol): execution_metadata: str | None +@dataclass(frozen=True) +class WorkflowNodeExecutionPayloads: + """Full node payloads; ``None`` means the execution has no value for that field.""" + + inputs: Mapping[str, Any] | None + process_data: Mapping[str, Any] | None + outputs: Mapping[str, Any] | None + + @dataclass(frozen=True) class WorkflowNodeExecutionSnapshot: """ @@ -184,6 +193,24 @@ def get_execution_snapshots_by_workflow_run( """ ... + def load_full_inputs(self, execution: WorkflowNodeExecutionModel) -> Mapping[str, Any] | None: + """Return complete inputs, resolving offloaded files within the admitted execution's tenant.""" + ... + + def load_full_outputs(self, execution: WorkflowNodeExecutionModel) -> Mapping[str, Any] | None: + """Return complete outputs, resolving offloaded files within the admitted execution's tenant.""" + ... + + def load_full_payloads( + self, executions: Sequence[WorkflowNodeExecutionModel] + ) -> Mapping[str, WorkflowNodeExecutionPayloads]: + """Batch-load payloads by execution ID for already admitted executions. + + Resolve file references within each execution's tenant in bounded queries, + then read storage after the database session closes. Shared files load once. + """ + ... + def get_workflow_tool_executions( self, tenant_id: str, diff --git a/api/repositories/app/tracing_config_repository.py b/api/repositories/app/tracing_config_repository.py index afa4968b3b8dbb..7b4ce9e0c59fdf 100644 --- a/api/repositories/app/tracing_config_repository.py +++ b/api/repositories/app/tracing_config_repository.py @@ -1,4 +1,4 @@ -"""SQLAlchemy persistence adapter for app tracing provider configurations.""" +"""SQLAlchemy persistence adapter for app tracing provider configurations and policies.""" import json from dataclasses import asdict, replace @@ -19,11 +19,32 @@ AppTracingConfigRecord, AppTracingConfigStore, ) +from services.tracing.delivery import StoredTraceDestination, TraceDestinationQuery +from services.tracing.scopes import TracingPolicyQuery _trace_settings = TypeAdapter(AppTraceSettings) -class SQLAlchemyAppTracingConfigRepository(AppTracingConfigStore): +def trace_settings_from_column(raw: str | None) -> AppTraceSettings: + """Read ``App.tracing`` for the trace pipeline: a value it cannot read disables tracing.""" + if not raw: + return AppTraceSettings() + try: + stored = json.loads(raw) + except ValueError: + return AppTraceSettings() + if not isinstance(stored, dict): + return AppTraceSettings() + provider = stored.get("tracing_provider") + return AppTraceSettings( + enabled=stored.get("enabled") is True, + tracing_provider=provider if isinstance(provider, str) else None, + capture_content=stored.get("capture_content") is True, + capture_agent=stored.get("capture_agent") is not False, + ) + + +class SQLAlchemyAppTracingConfigRepository(AppTracingConfigStore, TracingPolicyQuery, TraceDestinationQuery): def __init__(self, *, session_factory: sessionmaker[Session]) -> None: self._session_factory = session_factory @@ -127,6 +148,34 @@ def set_trace( raise AppTracingConfigNotFoundError app.tracing = json.dumps(asdict(settings)) + @override + def policy(self, *, tenant_id: str, app_id: str) -> AppTraceSettings | None: + with self._session_factory() as session: + row = session.execute(select(App.tracing).where(App.id == app_id, App.tenant_id == tenant_id)).first() + return trace_settings_from_column(row.tracing) if row is not None else None + + @override + def active_destination(self, *, tenant_id: str, app_id: str) -> StoredTraceDestination | None: + with self._session_factory() as session: + settings = trace_settings_from_column( + session.scalar(select(App.tracing).where(App.id == app_id, App.tenant_id == tenant_id)) + ) + if not settings.enabled or settings.tracing_provider is None: + return None + stored = session.scalar( + select(TraceAppConfig.tracing_config) + .join(App, App.id == TraceAppConfig.app_id) + .where( + TraceAppConfig.app_id == app_id, + App.tenant_id == tenant_id, + TraceAppConfig.tracing_provider == settings.tracing_provider, + TraceAppConfig.is_active.is_(True), + ) + ) + if not isinstance(stored, dict): + return None + return StoredTraceDestination(provider_id=settings.tracing_provider, config=dict(stored)) + @staticmethod def _settings(app: App) -> AppTraceSettings: return _trace_settings.validate_json(app.tracing) if app.tracing else AppTraceSettings() diff --git a/api/repositories/application_trace_repository.py b/api/repositories/application_trace_repository.py index 653e66ba6e5c48..26dae5de587f38 100644 --- a/api/repositories/application_trace_repository.py +++ b/api/repositories/application_trace_repository.py @@ -1,164 +1,511 @@ -"""Bounded transactions for application trace collection and delivery leases.""" +"""Persistence of Agent trace runs and reads of the sources they are collected for. + +Runs load and save as whole aggregates under their row lock; reads return +records. Rows become aggregates or records before their session closes. Every +database failure is reported as TracingStoreUnavailableError. +""" -import hashlib import json +import math +from collections.abc import Generator, Mapping +from contextlib import contextmanager from datetime import timedelta -from typing import Any, Literal -from uuid import UUID, uuid5 +from typing import Any, override -from dify_agent.protocol.tracing import AgentTraceBatch from sqlalchemy import delete, or_, select +from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session, sessionmaker +from domain.tracing.sources import MessageProgress, SourceNodeExecution, SourceWorkflowRun +from domain.tracing.trace_run_agg import ( + TraceBatch, + TraceBatchReceipt, + TraceInvocation, + TraceRunAgg, + TraceRunState, + TraceSourceKind, +) from libs.datetime_utils import naive_utc_now -from libs.uuid_utils import uuidv7 from models.application_trace import ApplicationTraceBatch, ApplicationTraceInvocation, ApplicationTraceRun -from models.model import App, Message, TraceAppConfig +from models.enums import MessageStatus +from models.model import App, AppMode, Conversation, Message, TraceAppConfig from models.workflow import Workflow, WorkflowNodeExecutionModel, WorkflowRun from repositories.factory import DifyAPIRepositoryFactory -from services.tracing.normalization import safe_attributes, span_id +from repositories.tracing_store_errors import tracing_store_errors +from services.tracing.collection import AppTracingSnapshot, SourceMessage, TraceApp, TraceAppMode, TraceRun +from services.tracing.collection_ports import LockedTraceRun, TraceRunQuery, TraceRunRepository, TraceSourceQuery +from services.tracing.errors import TraceCollectionConflictError + +_SOURCE_KINDS: dict[str, TraceSourceKind] = {"message": "message", "workflow": "workflow", "node": "node"} +_APP_MODES = { + AppMode.AGENT: TraceAppMode.AGENT, + AppMode.ADVANCED_CHAT: TraceAppMode.ADVANCED_CHAT, + AppMode.WORKFLOW: TraceAppMode.WORKFLOW, +} + + +def _message_metadata(message: Message) -> Mapping[str, Any]: + try: + metadata = json.loads(message.message_metadata) if message.message_metadata else {} + except (ValueError, TypeError): + return {} + return metadata if isinstance(metadata, dict) else {} + + +def _agent_stop_requested(message: Message) -> bool: + """Read the consumer's stop intent without inferring cancellation success.""" + return _message_metadata(message).get("agent_stop_requested") is True + + +def _answer_saved(message: Message) -> bool: + """Recognize the EasyUI _save_message transaction, not an early usage flush.""" + if message.status is not MessageStatus.NORMAL or message.updated_at <= message.created_at: + return False + metadata = _message_metadata(message) + if not {"annotation_reply", "retriever_resources", "reasoning"} <= metadata.keys(): + return False + usage = metadata.get("usage") + if not isinstance(usage, dict): + return False + for name in ("prompt_tokens", "completion_tokens"): + value = usage.get(name) + if not isinstance(value, (int, float)) or isinstance(value, bool) or not math.isfinite(value) or value < 0: + return False + latency = message.provider_response_latency + return isinstance(latency, (int, float)) and math.isfinite(latency) and latency >= 0 + + +def _stored_policy(raw: str | None) -> Mapping[str, Any]: + """``App.tracing`` as written; a policy that cannot be read is empty.""" + try: + policy = json.loads(raw) if raw else {} + except (TypeError, ValueError): + return {} + return policy if isinstance(policy, dict) else {} + + +def _run(row: ApplicationTraceRun) -> TraceRun: + return TraceRun( + id=row.id, + tenant_id=row.tenant_id, + app_id=row.app_id, + source_kind=_SOURCE_KINDS[row.source_kind], + source_id=row.source_id, + provider=row.provider, + config_revision=row.config_revision, + provider_config=dict(row.provider_config), + trace_id=row.trace_id, + root_span_id=row.root_span_id, + session_id=row.session_id, + capture_content=row.capture_content, + capture_agent=row.capture_agent, + state=TraceRunState(row.state), + outer_trace=dict(row.outer_trace) if row.outer_trace is not None else None, + span_count=row.span_count, + attempts=row.attempts, + error_code=row.error_code, + trace_url=row.trace_url, + provider_trace_id=row.provider_trace_id, + span_urls=dict(row.span_urls), + accepted_at=row.accepted_at, + ) + + +def _aggregate( + row: ApplicationTraceRun, + invocations: list[ApplicationTraceInvocation], + receipts: list[tuple[str, str, str]], +) -> TraceRunAgg: + return TraceRunAgg( + id=row.id, + tenant_id=row.tenant_id, + app_id=row.app_id, + source_kind=_SOURCE_KINDS[row.source_kind], + source_id=row.source_id, + provider=row.provider, + config_revision=row.config_revision, + provider_config=dict(row.provider_config), + trace_id=row.trace_id, + root_span_id=row.root_span_id, + session_id=row.session_id, + capture_content=row.capture_content, + capture_agent=row.capture_agent, + state=TraceRunState(row.state), + outer_trace=dict(row.outer_trace) if row.outer_trace is not None else None, + span_count=row.span_count, + payload_bytes=row.payload_bytes, + attempts=row.attempts, + error_code=row.error_code, + trace_url=row.trace_url, + provider_trace_id=row.provider_trace_id, + span_urls=dict(row.span_urls), + export_token=row.export_token, + export_started_at=row.export_started_at, + export_heartbeat_at=row.export_heartbeat_at, + next_attempt_at=row.next_attempt_at, + accepted_at=row.accepted_at, + terminal_at=row.terminal_at, + updated_at=row.updated_at, + invocations={invocation.id: _invocation(invocation) for invocation in invocations}, + receipts={ + batch_id: TraceBatchReceipt(id=batch_id, invocation_id=invocation_id, payload_hash=payload_hash) + for batch_id, invocation_id, payload_hash in receipts + }, + ) + + +def _new_row(run: TraceRunAgg) -> ApplicationTraceRun: + return ApplicationTraceRun( + id=run.id, + tenant_id=run.tenant_id, + app_id=run.app_id, + source_kind=run.source_kind, + source_id=run.source_id, + provider=run.provider, + config_revision=run.config_revision, + provider_config=dict(run.provider_config), + trace_id=run.trace_id, + root_span_id=run.root_span_id, + session_id=run.session_id, + capture_content=run.capture_content, + capture_agent=run.capture_agent, + state=run.state, + error_code=run.error_code, + terminal_at=run.terminal_at, + updated_at=run.updated_at, + ) + +def _save(row: ApplicationTraceRun, run: TraceRunAgg) -> None: + """Copy the fields a run's transitions change; its source, destination, and policy stay frozen.""" + row.state = run.state + row.outer_trace = dict(run.outer_trace) if run.outer_trace is not None else None + row.root_span_id = run.root_span_id + row.span_count = run.span_count + row.payload_bytes = run.payload_bytes + row.attempts = run.attempts + row.error_code = run.error_code + row.trace_url = run.trace_url + row.provider_trace_id = run.provider_trace_id + row.span_urls = dict(run.span_urls) + row.export_token = run.export_token + row.export_started_at = run.export_started_at + row.export_heartbeat_at = run.export_heartbeat_at + row.next_attempt_at = run.next_attempt_at + row.accepted_at = run.accepted_at + row.terminal_at = run.terminal_at + row.updated_at = run.updated_at -class TraceCollectionError(Exception): - def __init__(self, code: str, status: int = 400) -> None: - self.code = code - self.status = status - super().__init__(code) +def _invocation(row: ApplicationTraceInvocation) -> TraceInvocation: + return TraceInvocation( + id=row.id, + run_id=row.run_id, + agent_id=row.agent_id, + parent_span_id=row.parent_span_id, + span_id=row.span_id, + node_execution_id=row.node_execution_id, + attempt=row.attempt, + finalized=row.finalized, + finalized_at=row.finalized_at, + ) -class ApplicationTraceRepository: + +class ApplicationTraceRepository(TraceRunRepository, TraceRunQuery, TraceSourceQuery): def __init__(self, sessions: sessionmaker[Session]) -> None: - self.sessions = sessions + self._sessions = sessions + + @contextmanager + def _read(self) -> Generator[Session]: + with tracing_store_errors(), self._sessions() as session: + yield session + + @contextmanager + def _write(self) -> Generator[Session]: + with tracing_store_errors(), self._sessions.begin() as session: + yield session - def _source( - self, *, tenant_id: str, app_id: str, kind: str, source_id: str + def _owned_app_id(self, session: Session, *, tenant_id: str, app_id: str) -> str | None: + return session.scalar(select(App.id).where(App.id == app_id, App.tenant_id == tenant_id)) + + def _source_row( + self, *, tenant_id: str, app_id: str, kind: TraceSourceKind, source_id: str ) -> Message | WorkflowRun | WorkflowNodeExecutionModel | None: # Close the SQL owner read before consulting a configured backend: LogStore # can perform network I/O and need not have a corresponding SQL run/node row. - with self.sessions() as session: - if session.scalar(select(App.id).where(App.id == app_id, App.tenant_id == tenant_id)) is None: + with self._read() as session: + if self._owned_app_id(session, tenant_id=tenant_id, app_id=app_id) is None: return None if kind == "message": return session.scalar(select(Message).where(Message.id == source_id, Message.app_id == app_id)) source: WorkflowRun | WorkflowNodeExecutionModel | None - if kind == "workflow": - repository = DifyAPIRepositoryFactory.create_api_workflow_run_repository(self.sessions) - source = repository.get_workflow_run_by_id(tenant_id=tenant_id, app_id=app_id, run_id=source_id) - elif kind == "node": - node_repository = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(self.sessions) - source = node_repository.get_execution_by_id(execution_id=source_id, tenant_id=tenant_id) - else: - return None + with tracing_store_errors(): + if kind == "workflow": + repository = DifyAPIRepositoryFactory.create_api_workflow_run_repository(self._sessions) + source = repository.get_workflow_run_by_id(tenant_id=tenant_id, app_id=app_id, run_id=source_id) + else: + node_repository = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(self._sessions) + source = node_repository.get_execution_by_id(execution_id=source_id, tenant_id=tenant_id) # The node reader has no app filter. Check the entire chain on either # backend's result before exposing a destination or authorizing a retry. if source is None or (source.id, source.tenant_id, source.app_id) != (source_id, tenant_id, app_id): return None return source - def start( - self, - *, - tenant_id: str, - app_id: str, - source_kind: Literal["message", "workflow", "node"], - source_id: str, - session_id: str, - workflow_id: str | None = None, - ) -> ApplicationTraceRun | None: - # Re-entry uses the already validated, frozen source. It must not depend - # on LogStore availability or indexing latency after collection began. - with self.sessions() as session: - if session.scalar(select(App.id).where(App.id == app_id, App.tenant_id == tenant_id)) is None: + @override + def app(self, *, tenant_id: str, app_id: str) -> TraceApp | None: + with self._read() as session: + app = session.scalar(select(App).where(App.id == app_id, App.tenant_id == tenant_id)) + if app is None: return None - existing = session.scalar( - select(ApplicationTraceRun).where( - ApplicationTraceRun.tenant_id == tenant_id, - ApplicationTraceRun.app_id == app_id, - ApplicationTraceRun.source_kind == source_kind, - ApplicationTraceRun.source_id == source_id, + return TraceApp( + id=app.id, + tenant_id=app.tenant_id, + mode=_APP_MODES.get(app.mode, TraceAppMode.OTHER), + workflow_id=app.workflow_id, + ) + + @override + def app_workflow_graph(self, *, tenant_id: str, app_id: str, workflow_id: str) -> str | None: + with self._read() as session: + return session.scalar( + select(Workflow.graph).where( + Workflow.id == workflow_id, Workflow.app_id == app_id, Workflow.tenant_id == tenant_id ) ) - if existing is not None: - return existing - source = self._source(tenant_id=tenant_id, app_id=app_id, kind=source_kind, source_id=source_id) - # The app row serializes parallel node starts, so one workflow freezes one - # destination even when its first two Agent nodes start concurrently. - with self.sessions.begin() as session: - app = session.scalar(select(App).where(App.id == app_id, App.tenant_id == tenant_id).with_for_update()) - if app is None: + + @override + def message(self, *, tenant_id: str, app_id: str, message_id: str) -> SourceMessage | None: + with self._read() as session: + message = session.scalar( + select(Message) + .join(App, App.id == Message.app_id) + .where(Message.id == message_id, Message.app_id == app_id, App.tenant_id == tenant_id) + ) + if message is None: return None - existing = session.scalar( - select(ApplicationTraceRun).where( - ApplicationTraceRun.tenant_id == tenant_id, - ApplicationTraceRun.app_id == app_id, - ApplicationTraceRun.source_kind == source_kind, - ApplicationTraceRun.source_id == source_id, + return SourceMessage( + id=message.id, + conversation_id=message.conversation_id, + workflow_run_id=message.workflow_run_id, + agent_stop_requested=_agent_stop_requested(message), + ) + + @override + def workflow_run(self, *, tenant_id: str, app_id: str, workflow_run_id: str) -> SourceWorkflowRun | None: + run = self._source_row(tenant_id=tenant_id, app_id=app_id, kind="workflow", source_id=workflow_run_id) + if not isinstance(run, WorkflowRun): + return None + return SourceWorkflowRun( + id=run.id, + workflow_id=run.workflow_id, + graph=run.graph, + status=run.status, + finished_at=run.finished_at, + created_by=run.created_by, + ) + + @override + def node_execution(self, *, tenant_id: str, app_id: str, node_execution_id: str) -> SourceNodeExecution | None: + node = self._source_row(tenant_id=tenant_id, app_id=app_id, kind="node", source_id=node_execution_id) + if not isinstance(node, WorkflowNodeExecutionModel): + return None + return SourceNodeExecution( + id=node.id, + node_id=node.node_id, + workflow_id=node.workflow_id, + workflow_run_id=node.workflow_run_id, + title=node.title, + status=str(node.status), + error=node.error, + inputs=node.inputs, + outputs=node.outputs, + created_at=node.created_at, + elapsed_time=node.elapsed_time, + ) + + @override + def workflow_run_conversation_id(self, *, tenant_id: str, app_id: str, workflow_run_id: str) -> str | None: + with self._read() as session: + return session.scalar( + select(Message.conversation_id) + .join(App, App.id == Message.app_id) + .join(Conversation, Conversation.id == Message.conversation_id) + .where( + Message.app_id == app_id, + App.tenant_id == tenant_id, + Message.workflow_run_id == workflow_run_id, + Conversation.app_id == app_id, ) + .limit(1) ) - if existing is not None: - return existing - if source is None: - # A single-node preview allocates its execution ID before the - # node row is saved. Only the trusted workflow builder may reserve - # it, after checking its persisted owner workflow here. - if source_kind not in {"node", "workflow"} or workflow_id is None: - return None - workflow = session.scalar( - select(Workflow).where( - Workflow.id == workflow_id, Workflow.app_id == app_id, Workflow.tenant_id == tenant_id - ) + + @override + def message_progress(self, *, tenant_id: str, app_id: str, message_id: str) -> MessageProgress | None: + with self._read() as session: + source = session.execute( + select(Message, App.mode) + .join(App, App.id == Message.app_id) + .join(Conversation, Conversation.id == Message.conversation_id) + .where( + Message.id == message_id, + Message.app_id == app_id, + App.tenant_id == tenant_id, + Conversation.app_id == app_id, ) - if workflow is None: - return None - try: - policy = json.loads(app.tracing) if app.tracing else {} - except (TypeError, ValueError): - policy = {} - if not isinstance(policy, dict): - policy = {} + ).one_or_none() + if source is None: + return None + message, app_mode = source + return MessageProgress( + message_id=message.id, + conversation_id=message.conversation_id, + agent_app=app_mode == AppMode.AGENT, + paused=message.status is MessageStatus.PAUSED, + explicit_error=message.status is MessageStatus.ERROR and bool(message.error), + answer_saved=_answer_saved(message), + agent_stop_requested=_agent_stop_requested(message), + updated_at=message.updated_at, + ) + + @override + def app_tracing(self, *, tenant_id: str, app_id: str) -> AppTracingSnapshot | None: + with self._read() as session: + app = session.execute(select(App.tracing).where(App.id == app_id, App.tenant_id == tenant_id)).first() + if app is None: + return None + policy = _stored_policy(app.tracing) provider = policy.get("tracing_provider") config = ( session.scalar( - select(TraceAppConfig).where( - TraceAppConfig.app_id == app_id, - TraceAppConfig.tracing_provider == provider, + select(TraceAppConfig.tracing_config).where( + TraceAppConfig.app_id == app_id, TraceAppConfig.tracing_provider == provider ) ) - if isinstance(provider, str) and policy.get("enabled") is True + if isinstance(provider, str) else None ) - enabled = config is not None and isinstance(config.tracing_config, dict) - frozen = ( - dict(config.tracing_config) if config is not None and isinstance(config.tracing_config, dict) else {} + return AppTracingSnapshot(policy=policy, config=dict(config) if isinstance(config, dict) else None) + + @override + def add(self, run: TraceRunAgg) -> bool: + with self._write() as session: + try: + with session.begin_nested(): + session.add(_new_row(run)) + except IntegrityError: + # A locking read sees the committed winner even under MySQL's REPEATABLE READ snapshot. + winner = session.scalar( + select(ApplicationTraceRun.id) + .where( + ApplicationTraceRun.tenant_id == run.tenant_id, + ApplicationTraceRun.app_id == run.app_id, + ApplicationTraceRun.source_kind == run.source_kind, + ApplicationTraceRun.source_id == run.source_id, + ) + .with_for_update() + ) + if winner is None: + raise + return False + return True + + @staticmethod + def _invocation_models(session: Session, run_id: str) -> list[ApplicationTraceInvocation]: + return list( + session.scalars( + select(ApplicationTraceInvocation) + .where(ApplicationTraceInvocation.run_id == run_id) + .order_by( + ApplicationTraceInvocation.attempt, + ApplicationTraceInvocation.created_at, + ApplicationTraceInvocation.id, + ) ) - revision = hashlib.sha256( - json.dumps({"config": frozen, "policy": policy}, sort_keys=True, default=str).encode() - ).hexdigest() - record = ApplicationTraceRun( - id=str(uuidv7()), - tenant_id=tenant_id, - app_id=app_id, - source_kind=source_kind, - source_id=source_id, - provider=provider if enabled and isinstance(provider, str) else "", - config_revision=revision, - provider_config=frozen, - trace_id=uuidv7().hex, - root_span_id=span_id(source_id), - session_id=session_id[:255] or source_id, - capture_content=enabled and policy.get("capture_content", False) is True, - capture_agent=enabled and policy.get("capture_agent", True) is True, - state="collecting" if enabled else "not_collected", + ) + + @staticmethod + def _load(session: Session, row: ApplicationTraceRun, invocations: list[ApplicationTraceInvocation]) -> TraceRunAgg: + receipts = [ + (batch_id, invocation_id, payload_hash) + for batch_id, invocation_id, payload_hash in session.execute( + select( + ApplicationTraceBatch.id, + ApplicationTraceBatch.invocation_id, + ApplicationTraceBatch.payload_hash, + ).where(ApplicationTraceBatch.run_id == row.id) ) - session.add(record) + ] + return _aggregate(row, invocations, receipts) + + @override + def load(self, run_id: str) -> TraceRunAgg | None: + with self._read() as session: + row = session.get(ApplicationTraceRun, run_id) + if row is None: + return None + return self._load(session, row, self._invocation_models(session, run_id)) + + @override + @contextmanager + def change(self, run_id: str, *, tenant_id: str) -> Generator[LockedTraceRun | None]: + with self._write() as session: + row = session.scalar( + select(ApplicationTraceRun) + .where(ApplicationTraceRun.id == run_id, ApplicationTraceRun.tenant_id == tenant_id) + .with_for_update() + ) + if row is None: + yield None + return + invocation_models = self._invocation_models(session, run_id) + run = self._load(session, row, invocation_models) + invocation_rows = {invocation.id: invocation for invocation in invocation_models} + owner_exists = self._owned_app_id(session, tenant_id=row.tenant_id, app_id=row.app_id) is not None + yield LockedTraceRun(run=run, owner_exists=owner_exists) + _save(row, run) + for invocation in run.invocations.values(): + stored = invocation_rows.get(invocation.id) + if stored is None: + session.add( + ApplicationTraceInvocation( + id=invocation.id, + run_id=invocation.run_id, + agent_id=invocation.agent_id, + parent_span_id=invocation.parent_span_id, + span_id=invocation.span_id, + node_execution_id=invocation.node_execution_id, + attempt=invocation.attempt, + finalized=invocation.finalized, + finalized_at=invocation.finalized_at, + ) + ) + elif stored.finalized != invocation.finalized: + stored.finalized, stored.finalized_at = invocation.finalized, invocation.finalized_at session.flush() - return record + try: + with session.begin_nested(): + session.add_all( + ApplicationTraceBatch( + id=batch.id, + run_id=run.id, + invocation_id=batch.invocation_id, + payload_hash=batch.payload_hash, + spans=[dict(span) for span in batch.spans], + final=batch.final, + ) + for batch in run.new_batches + ) + except IntegrityError: + # Batch IDs are unique across runs; another run already stores this one. + raise TraceCollectionConflictError("tracing_batch_conflict") from None + @override def source_run( - self, *, tenant_id: str, app_id: str, source_kind: str, source_id: str - ) -> ApplicationTraceRun | None: - with self.sessions() as session: - return session.scalar( + self, *, tenant_id: str, app_id: str, source_kind: TraceSourceKind, source_id: str + ) -> TraceRun | None: + with self._read() as session: + run = session.scalar( select(ApplicationTraceRun).where( ApplicationTraceRun.tenant_id == tenant_id, ApplicationTraceRun.app_id == app_id, @@ -166,291 +513,67 @@ def source_run( ApplicationTraceRun.source_id == source_id, ) ) + return _run(run) if run is not None else None - def record_initialization_failure( - self, - *, - tenant_id: str, - app_id: str, - source_kind: str, - source_id: str, - ) -> ApplicationTraceRun | None: - """A missing start snapshot must never fall back to an unmasked exporter.""" - source = self._source(tenant_id=tenant_id, app_id=app_id, kind=source_kind, source_id=source_id) - if source is None: - return None - with self.sessions.begin() as session: - app = session.scalar(select(App).where(App.id == app_id, App.tenant_id == tenant_id).with_for_update()) - if app is None: - return None - current = session.scalar( - select(ApplicationTraceRun).where( - ApplicationTraceRun.tenant_id == tenant_id, - ApplicationTraceRun.app_id == app_id, - ApplicationTraceRun.source_kind == source_kind, - ApplicationTraceRun.source_id == source_id, - ) - ) - if current is not None: - return current - current = ApplicationTraceRun( - id=str(uuidv7()), - tenant_id=tenant_id, - app_id=app_id, - source_kind=source_kind, - source_id=source_id, - provider="", - config_revision="", - provider_config={}, - trace_id=uuidv7().hex, - root_span_id=span_id(source_id), - session_id=source_id, - capture_content=False, - capture_agent=False, - state="failed", - error_code="tracing_initialization_failed", - terminal_at=naive_utc_now(), - ) - session.add(current) - return current - - def invocation( - self, *, run_id: str, agent_id: str, node_execution_id: str | None, attempt: int, call_key: str - ) -> ApplicationTraceInvocation: - identifier = str(uuid5(UUID(run_id), call_key)) - with self.sessions.begin() as session: - run = session.scalar(select(ApplicationTraceRun).where(ApplicationTraceRun.id == run_id).with_for_update()) - if run is None or run.state not in {"collecting", "ready"}: - raise TraceCollectionError("tracing_run_not_collecting", 409) - invocation = session.get(ApplicationTraceInvocation, identifier) - if invocation is not None and ( - invocation.agent_id != agent_id - or invocation.node_execution_id != node_execution_id - or invocation.attempt != attempt - ): - raise TraceCollectionError("tracing_invocation_conflict", 409) - if invocation is None: - invocation = ApplicationTraceInvocation( - id=identifier, - run_id=run_id, - agent_id=agent_id, - parent_span_id=span_id(node_execution_id) if node_execution_id else run.root_span_id, - span_id=span_id(f"agent-attempt:{identifier}"), - node_execution_id=node_execution_id, - attempt=attempt, - ) - session.add(invocation) - return invocation - - def ingest(self, batch: AgentTraceBatch, *, max_spans: int, max_bytes: int) -> tuple[int, str]: - # Hash the original validated payload. Sanitization cannot turn two - # conflicting producer payloads into a falsely successful replay. - data = batch.model_dump(mode="json") - canonical = json.dumps(data, sort_keys=True, separators=(",", ":")).encode() - digest = hashlib.sha256(canonical).hexdigest() - with self.sessions.begin() as session: + @override + def workflow_run_trace(self, *, tenant_id: str, workflow_run_id: str) -> TraceRun | None: + with self._read() as session: run = session.scalar( - select(ApplicationTraceRun).where(ApplicationTraceRun.id == str(batch.run_id)).with_for_update() - ) - if run is None: - raise TraceCollectionError("tracing_run_not_found", 404) - app = session.scalar(select(App.id).where(App.id == run.app_id, App.tenant_id == run.tenant_id)) - if app is None: - raise TraceCollectionError("tracing_owner_not_found", 404) - invocation = session.scalar( - select(ApplicationTraceInvocation).where( - ApplicationTraceInvocation.id == str(batch.invocation_id), - ApplicationTraceInvocation.run_id == run.id, - ) - ) - if invocation is None: - raise TraceCollectionError("tracing_invocation_not_found", 404) - previous = session.get(ApplicationTraceBatch, str(batch.batch_id)) - if previous is not None: - if ( - previous.run_id != run.id - or previous.invocation_id != invocation.id - or previous.payload_hash != digest - ): - raise TraceCollectionError("tracing_batch_conflict", 409) - return len(previous.spans), run.id - if run.state in {"accepted", "exporting", "uncertain", "failed"} or invocation.finalized: - raise TraceCollectionError("tracing_collection_closed", 409) - if any(span.trace_id != run.trace_id for span in batch.spans): - raise TraceCollectionError("tracing_trace_id_mismatch", 409) - if run.span_count + len(batch.spans) > max_spans or run.payload_bytes + len(canonical) > max_bytes: - raise TraceCollectionError("tracing_payload_limit", 413) - spans: list[dict[str, Any]] = [] - for value in batch.spans: - span = value.model_dump(mode="json") - span["attributes"] = safe_attributes(span["attributes"], capture_content=run.capture_content) - if not run.capture_content: - span["events"] = [] - span["status_message"] = "Agent operation failed" if span["status"] == "error" else None - spans.append(span) - session.add( - ApplicationTraceBatch( - id=str(batch.batch_id), - run_id=run.id, - invocation_id=invocation.id, - payload_hash=digest, - spans=spans, - final=batch.final, + select(ApplicationTraceRun) + .join(App, App.id == ApplicationTraceRun.app_id) + .where( + ApplicationTraceRun.source_kind == "workflow", + ApplicationTraceRun.source_id == workflow_run_id, + ApplicationTraceRun.tenant_id == tenant_id, + App.tenant_id == tenant_id, ) ) - run.span_count += len(spans) - run.payload_bytes += len(canonical) - run.updated_at = naive_utc_now() - if batch.final: - invocation.finalized = True - invocation.finalized_at = naive_utc_now() - return len(spans), run.id - - def terminal(self, run_id: str, trace: dict[str, Any], root_span_id: str) -> None: - with self.sessions.begin() as session: - run = session.scalar(select(ApplicationTraceRun).where(ApplicationTraceRun.id == run_id).with_for_update()) - if run is None or run.state in {"accepted", "exporting", "uncertain", "failed"}: - return - if run.outer_trace is not None and run.outer_trace != trace: - # A duplicate terminal task must not rewrite a completed immutable - # span. The first persisted terminal envelope owns the export. - return - run.outer_trace = trace - run.root_span_id = root_span_id - run.terminal_at = run.terminal_at or naive_utc_now() - run.updated_at = naive_utc_now() - - def get(self, run_id: str) -> ApplicationTraceRun | None: - with self.sessions() as session: - return session.get(ApplicationTraceRun, run_id) - - def collected(self, run_id: str) -> tuple[list[ApplicationTraceInvocation], list[ApplicationTraceBatch]]: - with self.sessions() as session: - invocations = list( - session.scalars( - select(ApplicationTraceInvocation) - .where( - ApplicationTraceInvocation.run_id == run_id, - ) - .order_by( - ApplicationTraceInvocation.attempt, - ApplicationTraceInvocation.created_at, - ApplicationTraceInvocation.id, - ) + return _run(run) if run is not None else None + + @override + def get(self, run_id: str) -> TraceRun | None: + with self._read() as session: + run = session.get(ApplicationTraceRun, run_id) + return _run(run) if run is not None else None + + @override + def invocations(self, run_id: str) -> list[TraceInvocation]: + with self._read() as session: + return [_invocation(row) for row in self._invocation_models(session, run_id)] + + @override + def batches(self, run_id: str) -> list[TraceBatch]: + with self._read() as session: + return [ + TraceBatch( + id=row.id, + invocation_id=row.invocation_id, + payload_hash=row.payload_hash, + spans=list(row.spans), + final=row.final, ) - ) - batches = list( - session.scalars( + for row in session.scalars( select(ApplicationTraceBatch) .where(ApplicationTraceBatch.run_id == run_id) .order_by(ApplicationTraceBatch.created_at) ) - ) - return invocations, batches - - def claim( - self, run_id: str, *, collection_timeout: int, lease_seconds: int - ) -> tuple[ApplicationTraceRun, str] | None: - now = naive_utc_now() - with self.sessions.begin() as session: - run = session.scalar(select(ApplicationTraceRun).where(ApplicationTraceRun.id == run_id).with_for_update()) - if run is None or run.state in {"not_collected", "accepted", "failed", "uncertain"}: - return None - if run.state == "exporting": - last_heartbeat = run.export_heartbeat_at or run.export_started_at - if last_heartbeat and last_heartbeat + timedelta(seconds=lease_seconds) < now: - run.state, run.error_code = "uncertain", "export_lease_expired" - run.updated_at = now - return None - if run.outer_trace is None: - return None - incomplete = session.scalar( - select(ApplicationTraceInvocation.id) - .where(ApplicationTraceInvocation.run_id == run_id, ApplicationTraceInvocation.finalized.is_(False)) - .limit(1) - ) - if incomplete: - if run.terminal_at and run.terminal_at + timedelta(seconds=collection_timeout) < now: - run.state, run.error_code = "failed", "agent_trace_collection_incomplete" - run.updated_at = now - return None - if run.next_attempt_at and run.next_attempt_at > now: - return None - app = session.scalar(select(App.id).where(App.id == run.app_id, App.tenant_id == run.tenant_id)) - if app is None: - run.state, run.error_code = "failed", "tracing_owner_not_found" - return None - token = str(uuidv7()) - run.state = "exporting" - run.export_token = token - run.export_started_at = now - run.export_heartbeat_at = now - run.updated_at = now - run.attempts += 1 - return run, token - - def renew_export_lease(self, run_id: str, token: str) -> bool: - """Renew only the current export attempt, in a short independent transaction.""" - with self.sessions.begin() as session: - run = session.scalar(select(ApplicationTraceRun).where(ApplicationTraceRun.id == run_id).with_for_update()) - if run is None or run.state != "exporting" or run.export_token != token: - return False - heartbeat_at = naive_utc_now() - run.export_heartbeat_at = heartbeat_at - run.updated_at = heartbeat_at - return True - - def complete( - self, - run_id: str, - token: str, - *, - state: str, - error_code: str | None = None, - result: Any = None, - retry_delay: int | None = None, - ) -> None: - with self.sessions.begin() as session: - run = session.scalar(select(ApplicationTraceRun).where(ApplicationTraceRun.id == run_id).with_for_update()) - if run is None or run.export_token != token: - return - late_acceptance = ( - state == "accepted" - and result is not None - and run.state == "uncertain" - and run.error_code == "export_lease_expired" - ) - if run.state != "exporting" and not late_acceptance: - return - run.state, run.error_code = state, error_code - run.updated_at = naive_utc_now() - run.export_token = None - if result is not None: - run.provider_trace_id, run.trace_url, run.span_urls = ( - result.trace_id, - result.trace_url, - result.span_urls, - ) - run.accepted_at = naive_utc_now() - run.next_attempt_at = naive_utc_now() + timedelta(seconds=retry_delay) if retry_delay is not None else None - - def recovery_checked(self, run_id: str, *, error_code: str | None = None) -> None: - """Rotate checked incomplete work without changing a concurrent export.""" - with self.sessions.begin() as session: - run = session.scalar(select(ApplicationTraceRun).where(ApplicationTraceRun.id == run_id).with_for_update()) - if run is None or run.outer_trace is not None or run.state not in {"collecting", "ready"}: - return - run.updated_at = naive_utc_now() - if error_code is not None: - run.state, run.error_code = "failed", error_code + ] - def pending_ids(self, limit: int = 100) -> list[str]: - with self.sessions() as session: + @override + def pending_ids(self, limit: int) -> list[str]: + with self._read() as session: return list( session.scalars( select(ApplicationTraceRun.id) .where( - ApplicationTraceRun.state.in_(["collecting", "ready", "retrying", "exporting"]), + ApplicationTraceRun.state.in_( + [ + TraceRunState.COLLECTING, + TraceRunState.READY, + TraceRunState.RETRYING, + TraceRunState.EXPORTING, + ] + ), or_( ApplicationTraceRun.next_attempt_at.is_(None), ApplicationTraceRun.next_attempt_at <= naive_utc_now(), @@ -461,34 +584,10 @@ def pending_ids(self, limit: int = 100) -> list[str]: ) ) - def retry(self, *, run_id: str, acknowledge_possible_duplicates: bool) -> None: - with self.sessions.begin() as session: - run = session.scalar(select(ApplicationTraceRun).where(ApplicationTraceRun.id == run_id).with_for_update()) - if run is None or run.state == "not_collected": - raise TraceCollectionError("tracing_run_not_found", 404) - if run.state == "accepted": - raise TraceCollectionError("tracing_already_accepted", 409) - if run.state not in {"failed", "uncertain"} or run.outer_trace is None: - raise TraceCollectionError("tracing_retry_not_allowed", 409) - incomplete = session.scalar( - select(ApplicationTraceInvocation.id) - .where( - ApplicationTraceInvocation.run_id == run.id, - ApplicationTraceInvocation.finalized.is_(False), - ) - .limit(1) - ) - if incomplete or (run.error_code or "").startswith("tracing_"): - raise TraceCollectionError("tracing_retry_not_allowed", 409) - if run.state == "uncertain" and not acknowledge_possible_duplicates: - raise TraceCollectionError("tracing_duplicate_acknowledgement_required", 409) - run.state, run.error_code, run.export_token = "ready", None, None - run.next_attempt_at = None - run.updated_at = naive_utc_now() - - def purge(self, *, retention_days: int, limit: int = 100) -> int: + @override + def purge(self, *, retention_days: int, limit: int) -> int: cutoff = naive_utc_now() - timedelta(days=retention_days) - with self.sessions.begin() as session: + with self._write() as session: ids = list( session.scalars( select(ApplicationTraceRun.id) diff --git a/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py b/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py index 219dbd32721cb6..632cc2e55e3cea 100644 --- a/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py +++ b/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py @@ -5,19 +5,28 @@ using SQLAlchemy 2.0 style queries for WorkflowNodeExecutionModel operations. """ +import json from collections.abc import Mapping, Sequence from datetime import datetime +from itertools import batched from typing import Any, cast, override -from sqlalchemy import String, asc, delete, desc, func, or_, select +from sqlalchemy import String, asc, delete, desc, func, or_, select, tuple_ from sqlalchemy.engine import CursorResult from sqlalchemy.orm import Session, sessionmaker from extensions.ext_storage import storage from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionStatus -from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionOffload, WorkflowNodeExecutionTriggeredFrom +from models.model import UploadFile +from models.workflow import ( + ExecutionOffLoadType, + WorkflowNodeExecutionModel, + WorkflowNodeExecutionOffload, + WorkflowNodeExecutionTriggeredFrom, +) from repositories.api_workflow_node_execution_repository import ( DifyAPIWorkflowNodeExecutionRepository, + WorkflowNodeExecutionPayloads, WorkflowNodeExecutionSnapshot, WorkflowNodeExecutionSnapshotRow, ) @@ -52,11 +61,65 @@ def __init__(self, session_maker: sessionmaker[Session]): """ self._session_maker = session_maker + @override + def load_full_inputs(self, execution: WorkflowNodeExecutionModel) -> Mapping[str, Any] | None: + return self._load_full_payload(execution, ExecutionOffLoadType.INPUTS, execution.inputs_dict) + + @override + def load_full_outputs(self, execution: WorkflowNodeExecutionModel) -> Mapping[str, Any] | None: + return self._load_full_payload(execution, ExecutionOffLoadType.OUTPUTS, execution.outputs_dict) + @override def load_full_process_data(self, execution: WorkflowNodeExecutionModel) -> Mapping[str, Any] | None: """Load Process Data from external storage when the execution was offloaded.""" + return self._load_full_payload(execution, ExecutionOffLoadType.PROCESS_DATA, execution.process_data_dict) + + def _load_full_payload( + self, + execution: WorkflowNodeExecutionModel, + type_: ExecutionOffLoadType, + inline: Mapping[str, Any] | None, + ) -> Mapping[str, Any] | None: + offload = next((item for item in execution.offload_data if item.type_ == type_), None) + if offload is None: + return inline + reference = (execution.tenant_id, offload.file_id) + return self._load_offloaded_payloads({reference})[reference] + + def _load_offloaded_payloads( + self, references: set[tuple[str, str]] + ) -> dict[tuple[str, str], Mapping[str, Any] | None]: + if not references: + return {} + keys: dict[tuple[str, str], str] = {} with self._session_maker() as session: - return execution.load_full_process_data(session, storage) + for batch in batched(references, 1000, strict=False): + rows = session.execute( + select(UploadFile.tenant_id, UploadFile.id, UploadFile.key).where( + tuple_(UploadFile.tenant_id, UploadFile.id).in_(batch) + ) + ) + keys.update(((tenant_id, file_id), key) for tenant_id, file_id, key in rows) + for tenant_id, file_id in references: + assert (tenant_id, file_id) in keys, f"UploadFile with id {file_id} should exist in the execution's tenant" + return {reference: json.loads(storage.load(keys[reference])) for reference in references} + + @override + def load_full_payloads( + self, executions: Sequence[WorkflowNodeExecutionModel] + ) -> Mapping[str, WorkflowNodeExecutionPayloads]: + loaded = self._load_offloaded_payloads( + {(execution.tenant_id, offload.file_id) for execution in executions for offload in execution.offload_data} + ) + result: dict[str, WorkflowNodeExecutionPayloads] = {} + for execution in executions: + full = {offload.type_: loaded[execution.tenant_id, offload.file_id] for offload in execution.offload_data} + result[execution.id] = WorkflowNodeExecutionPayloads( + inputs=full.get(ExecutionOffLoadType.INPUTS, execution.inputs_dict), + process_data=full.get(ExecutionOffLoadType.PROCESS_DATA, execution.process_data_dict), + outputs=full.get(ExecutionOffLoadType.OUTPUTS, execution.outputs_dict), + ) + return result @override def get_node_last_execution( diff --git a/api/repositories/trace_record_repository.py b/api/repositories/trace_record_repository.py new file mode 100644 index 00000000000000..821505601888b5 --- /dev/null +++ b/api/repositories/trace_record_repository.py @@ -0,0 +1,823 @@ +"""Assemble trace records from events and trusted database state. + +Every record is read through the event's tenant and app: a reference that does +not belong to them is missing, never trusted. Events carry only identifiers +and facts persisted nowhere else, so records reflect committed state. Workflow +tool descendants are admitted only through their owning caller in the same +tenant and run. Persistence models never cross the application port. +""" + +import json +from collections import deque +from collections.abc import Mapping, Sequence +from datetime import datetime, timedelta +from typing import Any, override + +from dify_trace_sdk.events import ( + ConversationNamed, + DatasetRetrieved, + DraftNodeExecuted, + MessageCompleted, + ModerationChecked, + PromptGenerated, + SuggestedQuestionsGenerated, + ToolInvoked, + TraceEvent, + WorkflowRunCompleted, +) +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from graphon.entities import WorkflowNodeExecution +from graphon.enums import ( + BuiltinNodeTypes, + WorkflowExecutionStatus, + WorkflowNodeExecutionMetadataKey, + WorkflowNodeExecutionStatus, +) +from models import App, Tenant +from models.dataset import Dataset +from models.model import Conversation, EndUser, Message, MessageFile +from models.provider import Provider, ProviderCredential, ProviderModel, ProviderModelCredential, ProviderType +from models.tools import ApiToolProvider, BuiltinToolProvider, MCPToolProvider, WorkflowToolProvider +from models.workflow import WorkflowAppLog, WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom, WorkflowRun +from repositories.api_workflow_node_execution_repository import ( + DifyAPIWorkflowNodeExecutionRepository, + WorkflowNodeExecutionPayloads, +) +from repositories.api_workflow_run_repository import APIWorkflowRunRepository +from repositories.tracing_store_errors import tracing_store_errors +from services.tracing.delivery import TraceRecordSource +from services.tracing.errors import SourceNotFoundError, SourceNotReadyError +from services.tracing.records import ( + BaseTraceInfo, + DatasetRetrievalTraceInfo, + DraftNodeExecutionTrace, + GenerateNameTraceInfo, + MessageTraceInfo, + ModerationTraceInfo, + PromptGenerationTraceInfo, + SuggestedQuestionTraceInfo, + ToolTraceInfo, + TraceMessage, + WorkflowTraceInfo, +) + +_TERMINAL_WORKFLOW_STATUSES = frozenset( + { + WorkflowExecutionStatus.SUCCEEDED, + WorkflowExecutionStatus.PARTIAL_SUCCEEDED, + WorkflowExecutionStatus.FAILED, + WorkflowExecutionStatus.STOPPED, + } +) +_TOOL_PROVIDER_MODELS: dict[ + str, type[BuiltinToolProvider | ApiToolProvider | WorkflowToolProvider | MCPToolProvider] +] = { + "builtin": BuiltinToolProvider, + "plugin": BuiltinToolProvider, + "api": ApiToolProvider, + "workflow": WorkflowToolProvider, + "mcp": MCPToolProvider, +} + + +def _user_of(metadata: Mapping[str, Any]) -> str: + """Who invoked the operation: end user before account, as the enterprise data dictionary defines.""" + if user_id := metadata.get("from_end_user_id"): + return f"end_user:{user_id}" + if user_id := metadata.get("from_account_id"): + return f"account:{user_id}" + if user_id := metadata.get("user_id"): + return f"user:{user_id}" + return "anonymous" + + +def _text_input(inputs: Mapping[str, Any], *names: str) -> str: + """The first text value among ``names``; apps may define inputs of these names with any type.""" + return next((value for name in names if isinstance(value := inputs.get(name), str) and value), "") + + +def _trace_message(message: Message) -> TraceMessage: + return TraceMessage( + id=message.id, + app_id=message.app_id, + conversation_id=message.conversation_id, + workflow_run_id=message.workflow_run_id, + query=message.query, + answer=message.answer, + error=message.error, + created_at=message.created_at, + model_provider=message.model_provider, + model_id=message.model_id, + total_price=message.total_price, + currency=message.currency, + ) + + +def _streaming_metrics(message: Message) -> dict[str, Any]: + if not message.message_metadata: + return {} + try: + metadata = json.loads(message.message_metadata) + except ValueError: + return {} + usage = metadata.get("usage") if isinstance(metadata, dict) else None + if not isinstance(usage, dict): + return {} + time_to_first_token = usage.get("time_to_first_token") + return { + "gen_ai_server_time_to_first_token": time_to_first_token, + "llm_streaming_time_to_generate": usage.get("time_to_generate"), + "is_streaming_request": time_to_first_token is not None, + } + + +class TraceRecordRepository(TraceRecordSource): + def __init__( + self, + session_factory: sessionmaker[Session], + *, + workflow_runs: APIWorkflowRunRepository, + node_executions: DifyAPIWorkflowNodeExecutionRepository, + files_url: str, + enterprise_details: bool, + ) -> None: + self._sessions = session_factory + self._workflow_runs = workflow_runs + self._node_executions_repository = node_executions + self._files_url = files_url + # App, workspace, and credential names are enterprise telemetry dimensions only. + self._enterprise_details = enterprise_details + + @override + def assemble(self, event: TraceEvent) -> BaseTraceInfo: + match event.payload: + case MessageCompleted() as payload: + return self._message(event, payload) + case WorkflowRunCompleted() as payload: + return self._workflow(event, payload) + case ModerationChecked() as payload: + return self._moderation(event, payload) + case SuggestedQuestionsGenerated() as payload: + return self._suggested_questions(event, payload) + case DatasetRetrieved() as payload: + return self._dataset_retrieval(event, payload) + case ToolInvoked() as payload: + return self._tool(event, payload) + case ConversationNamed() as payload: + return self._conversation_name(event, payload) + case PromptGenerated() as payload: + return self._prompt_generation(event, payload) + case DraftNodeExecuted() as payload: + return self._draft_node(event, payload) + + # Shared lookups + + def _app(self, session: Session, event: TraceEvent) -> App: + app = session.scalar(select(App).where(App.id == event.app_id, App.tenant_id == event.tenant_id)) + if app is None: + raise SourceNotFoundError + return app + + def _message_in_app(self, session: Session, event: TraceEvent, message_id: str | None) -> Message: + if message_id is None: + raise SourceNotFoundError + message = session.scalar( + select(Message) + .join(App, App.id == Message.app_id) + .where(Message.id == message_id, Message.app_id == event.app_id, App.tenant_id == event.tenant_id) + ) + if message is None: + # Producers inside a request transaction may report before their message commits. + raise SourceNotReadyError + return message + + @staticmethod + def _end_user_session_id(session: Session, message: Message, tenant_id: str) -> str | None: + if not message.from_end_user_id: + return None + return session.scalar( + select(EndUser.session_id).where(EndUser.id == message.from_end_user_id, EndUser.tenant_id == tenant_id) + ) + + def _names(self, session: Session, app_id: str | None, tenant_id: str) -> tuple[str, str]: + if not self._enterprise_details: + return "", "" + app_name = ( + session.scalar(select(App.name).where(App.id == app_id, App.tenant_id == tenant_id)) if app_id else None + ) + workspace_name = session.scalar(select(Tenant.name).where(Tenant.id == tenant_id)) + return app_name or "", workspace_name or "" + + @staticmethod + def _workflow_app_log_id(session: Session, message: Message) -> str | None: + if not message.workflow_run_id: + return None + log_id = session.scalar( + select(WorkflowAppLog.id) + .where(WorkflowAppLog.app_id == message.app_id, WorkflowAppLog.workflow_run_id == message.workflow_run_id) + .limit(1) + ) + return str(log_id) if log_id else None + + @staticmethod + def _message_metadata(message: Message) -> dict[str, Any]: + return { + "message_id": message.id, + "ls_provider": message.model_provider, + "ls_model_name": message.model_id, + "status": message.status, + "from_end_user_id": message.from_end_user_id, + "from_account_id": message.from_account_id, + "agent_based": message.agent_based, + "workflow_run_id": message.workflow_run_id, + "from_source": message.from_source, + } + + @staticmethod + def _correlation_metadata(event: TraceEvent) -> dict[str, Any]: + metadata: dict[str, Any] = {"trace_event_id": event.event_id} + if event.correlation.session_id: + metadata["trace_session_id"] = event.correlation.session_id + return metadata + + # Records + + def _message(self, event: TraceEvent, payload: MessageCompleted) -> MessageTraceInfo: + with tracing_store_errors(), self._sessions() as session: + message = self._message_in_app(session, event, payload.message_id) + conversation_mode = session.scalar( + select(Conversation.mode).where( + Conversation.id == message.conversation_id, Conversation.app_id == message.app_id + ) + ) + if conversation_mode is None: + raise SourceNotFoundError + message_file = session.scalar(select(MessageFile).where(MessageFile.message_id == message.id).limit(1)) + end_user_session_id = self._end_user_session_id(session, message, event.tenant_id) + app_name, workspace_name = self._names(session, message.app_id, event.tenant_id) + + file_list = [f"{self._files_url}/{message_file.url}"] if message_file and message_file.url else [] + metadata = { + **self._correlation_metadata(event), + **self._message_metadata(message), + "conversation_id": message.conversation_id, + "tenant_id": event.tenant_id, + "app_id": message.app_id, + "user_id": message.from_end_user_id or message.from_account_id, + "end_user_session_id": end_user_session_id, + "app_name": app_name, + "workspace_name": workspace_name, + } + created_at = message.created_at + streaming = _streaming_metrics(message) + return MessageTraceInfo( + trace_id=event.correlation.external_trace_id, + message_id=message.id, + message_data=_trace_message(message), + conversation_model=conversation_mode, + message_tokens=message.message_tokens, + answer_tokens=message.answer_tokens, + total_tokens=message.message_tokens + message.answer_tokens, + error=message.error or "", + inputs=message.message, + outputs=message.answer, + file_list=file_list, + start_time=created_at, + end_time=message.updated_at + if message.updated_at and message.updated_at > created_at + else created_at + timedelta(seconds=message.provider_response_latency), + metadata=metadata, + conversation_mode=conversation_mode, + gen_ai_server_time_to_first_token=streaming.get("gen_ai_server_time_to_first_token"), + llm_streaming_time_to_generate=streaming.get("llm_streaming_time_to_generate"), + is_streaming_request=streaming.get("is_streaming_request", False), + ) + + def _workflow(self, event: TraceEvent, payload: WorkflowRunCompleted) -> WorkflowTraceInfo: + if event.app_id is None: + raise SourceNotFoundError + with tracing_store_errors(): + run = self._workflow_runs.get_workflow_run_by_id( + tenant_id=event.tenant_id, app_id=event.app_id, run_id=payload.workflow_run_id + ) + # Workflow runs may be persisted asynchronously (including LogStore indexing). + if run is None or run.status not in _TERMINAL_WORKFLOW_STATUSES or run.finished_at is None: + raise SourceNotReadyError + with tracing_store_errors(), self._sessions() as session: + self._app(session, event) + workflow_app_log_id = session.scalar( + select(WorkflowAppLog.id).where( + WorkflowAppLog.tenant_id == event.tenant_id, + WorkflowAppLog.app_id == event.app_id, + WorkflowAppLog.workflow_run_id == run.id, + ) + ) + message_row = ( + session.execute( + select(Message.id, Message.created_at).where( + Message.app_id == event.app_id, + Message.conversation_id == payload.conversation_id, + Message.workflow_run_id == run.id, + ) + ).first() + if payload.conversation_id + else None + ) + app_name, workspace_name = self._names(session, event.app_id, event.tenant_id) + executions, workflow_tool_parents = self._node_executions(run) + prompt_tokens, completion_tokens = self._token_split(executions) + inputs = run.inputs_dict + file_list = inputs.get("sys.file") or [] + total_tokens = payload.total_tokens if payload.total_tokens is not None else run.total_tokens + metadata: dict[str, Any] = { + **self._correlation_metadata(event), + "workflow_id": run.workflow_id, + "conversation_id": payload.conversation_id, + "workflow_run_id": run.id, + "tenant_id": event.tenant_id, + "elapsed_time": run.elapsed_time, + "status": run.status, + "version": run.version, + "total_tokens": total_tokens, + "file_list": file_list, + "triggered_from": run.triggered_from, + "user_id": event.user_id, + "app_id": run.app_id, + "app_name": app_name, + "workspace_name": workspace_name, + } + correlation = event.correlation + if correlation.parent_workflow_run_id: + metadata["parent_trace_context"] = { + "parent_workflow_run_id": correlation.parent_workflow_run_id, + "parent_node_execution_id": correlation.parent_node_execution_id, + } + return WorkflowTraceInfo( + trace_id=correlation.external_trace_id, + node_executions=executions, + workflow_tool_parents=workflow_tool_parents, + conversation_id=payload.conversation_id, + message_started_at=message_row[1] if message_row else None, + workflow_id=run.workflow_id, + tenant_id=event.tenant_id, + workflow_run_id=run.id, + workflow_run_elapsed_time=run.elapsed_time, + workflow_run_status=run.status, + workflow_run_inputs=inputs, + workflow_run_outputs=run.outputs_dict, + workflow_run_version=run.version, + error=run.error or "", + total_tokens=total_tokens, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + file_list=file_list, + query=_text_input(inputs, "query", "sys.query"), + metadata=metadata, + workflow_app_log_id=workflow_app_log_id, + message_id=message_row[0] if message_row else None, + start_time=run.created_at, + end_time=run.finished_at, + invoked_by=_user_of(metadata), + ) + + def _node_executions(self, run: WorkflowRun) -> tuple[Sequence[WorkflowNodeExecution], Mapping[str, str]]: + repository = self._node_executions_repository + rows: list[WorkflowNodeExecutionModel] = [] + tool_parents: dict[str, str] = {} + seen: set[str] = set() + with tracing_store_errors(): + pending: deque[tuple[WorkflowNodeExecutionModel, str | None]] = deque( + (row, None) + for row in repository.get_executions_by_workflow_run(run.tenant_id, run.app_id, run.id) + if row.triggered_from == WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN + ) + while pending: + row, parent_id = pending.popleft() + if row.id in seen or row.status == WorkflowNodeExecutionStatus.PAUSED: + continue + seen.add(row.id) + rows.append(row) + if parent_id is not None: + tool_parents[row.id] = parent_id + if row.node_type == BuiltinNodeTypes.TOOL: + # Descendants are admitted through an already owned Tool, never by source app alone. + pending.extend( + (child, row.id) + for child in repository.get_workflow_tool_executions( + run.tenant_id, run.id, row.node_execution_id or row.id + ) + ) + payloads = repository.load_full_payloads(rows) + return [self._node_execution(row, payloads[row.id]) for row in rows], tool_parents + + @staticmethod + def _node_execution( + row: WorkflowNodeExecutionModel, payloads: WorkflowNodeExecutionPayloads + ) -> WorkflowNodeExecution: + return WorkflowNodeExecution( + id=row.id, + node_execution_id=row.node_execution_id, + workflow_id=row.workflow_id, + workflow_execution_id=row.workflow_run_id, + index=row.index, + predecessor_node_id=row.predecessor_node_id, + node_id=row.node_id, + node_type=row.node_type, + title=row.title, + inputs=payloads.inputs, + process_data=payloads.process_data, + outputs=payloads.outputs, + status=WorkflowNodeExecutionStatus(row.status), + error=row.error, + elapsed_time=row.elapsed_time, + metadata={ + WorkflowNodeExecutionMetadataKey(key): value + for key, value in row.execution_metadata_dict.items() + if key in WorkflowNodeExecutionMetadataKey + }, + created_at=row.created_at, + finished_at=row.finished_at, + ) + + @staticmethod + def _token_split(executions: Sequence[WorkflowNodeExecution]) -> tuple[int, int]: + """Prompt and completion totals; the run row only stores the overall total.""" + prompt_total = completion_total = 0 + for execution in executions: + usage = (execution.outputs or {}).get("usage") + if not isinstance(usage, Mapping): + continue + prompt, completion = usage.get("prompt_tokens"), usage.get("completion_tokens") + if isinstance(prompt, int | float): + prompt_total += int(prompt) + if isinstance(completion, int | float): + completion_total += int(completion) + return prompt_total, completion_total + + def _moderation(self, event: TraceEvent, payload: ModerationChecked) -> ModerationTraceInfo: + with tracing_store_errors(), self._sessions() as session: + message = self._message_in_app(session, event, payload.message_id) + workflow_app_log_id = self._workflow_app_log_id(session, message) + end_user_session_id = self._end_user_session_id(session, message, event.tenant_id) + metadata = { + **self._correlation_metadata(event), + "message_id": message.id, + "action": payload.action, + "preset_response": payload.preset_response, + "query": payload.query, + "user_id": message.from_end_user_id or message.from_account_id, + "end_user_session_id": end_user_session_id, + } + return ModerationTraceInfo( + trace_id=event.correlation.external_trace_id, + # Enterprise telemetry keys chatflow moderation by its workflow log. + message_id=workflow_app_log_id or message.id, + inputs=payload.inputs, + message_data=_trace_message(message), + flagged=payload.flagged, + action=payload.action, + preset_response=payload.preset_response, + query=payload.query, + start_time=payload.started_at, + end_time=payload.ended_at, + metadata=metadata, + ) + + def _suggested_questions( + self, event: TraceEvent, payload: SuggestedQuestionsGenerated + ) -> SuggestedQuestionTraceInfo: + with tracing_store_errors(), self._sessions() as session: + message = self._message_in_app(session, event, payload.message_id) + workflow_app_log_id = self._workflow_app_log_id(session, message) + end_user_session_id = self._end_user_session_id(session, message, event.tenant_id) + metadata = { + **self._correlation_metadata(event), + **self._message_metadata(message), + "user_id": message.from_end_user_id or message.from_account_id, + "end_user_session_id": end_user_session_id, + } + return SuggestedQuestionTraceInfo( + trace_id=event.correlation.external_trace_id, + message_id=workflow_app_log_id or message.id, + message_data=_trace_message(message), + inputs=message.message, + outputs=message.answer, + start_time=payload.started_at, + end_time=payload.ended_at, + metadata=metadata, + total_tokens=message.message_tokens + message.answer_tokens, + status=message.status, + error=message.error, + from_account_id=message.from_account_id, + agent_based=message.agent_based, + from_source=message.from_source, + model_provider=message.model_provider, + model_id=message.model_id, + suggested_question=payload.questions, + level=message.status, + status_message=message.error, + workflow_run_id=None, + ) + + def _dataset_retrieval(self, event: TraceEvent, payload: DatasetRetrieved) -> DatasetRetrievalTraceInfo: + dataset_ids: set[str] = set() + for document in payload.documents: + document_metadata = document.get("metadata") + if isinstance(document_metadata, dict) and isinstance(document_metadata.get("dataset_id"), str): + dataset_ids.add(str(document_metadata["dataset_id"])) + with tracing_store_errors(), self._sessions() as session: + message = self._message_in_app(session, event, payload.message_id) + end_user_session_id = self._end_user_session_id(session, message, event.tenant_id) + app_name, workspace_name = self._names(session, message.app_id, event.tenant_id) + embedding_models = { + str(row.id): { + "embedding_model": row.embedding_model or "", + "embedding_model_provider": row.embedding_model_provider or "", + } + for row in session.execute( + select(Dataset.id, Dataset.embedding_model, Dataset.embedding_model_provider).where( + Dataset.id.in_(sorted(dataset_ids)), Dataset.tenant_id == event.tenant_id + ) + ) + } + inputs = message.query or message.inputs_with_session(session=session) + metadata = { + **self._correlation_metadata(event), + **self._message_metadata(message), + "tenant_id": event.tenant_id, + "app_id": message.app_id, + "user_id": message.from_end_user_id or message.from_account_id, + "end_user_session_id": end_user_session_id, + "app_name": app_name, + "workspace_name": workspace_name, + "embedding_models": embedding_models, + "rerank_model_provider": "", + "rerank_model_name": "", + } + return DatasetRetrievalTraceInfo( + trace_id=event.correlation.external_trace_id, + message_id=message.id, + inputs=inputs, + documents=payload.documents, + start_time=payload.started_at, + end_time=payload.ended_at, + metadata=metadata, + message_data=_trace_message(message), + error=None, + ) + + def _tool(self, event: TraceEvent, payload: ToolInvoked) -> ToolTraceInfo: + with tracing_store_errors(), self._sessions() as session: + message = self._message_in_app(session, event, payload.message_id) + message_file = session.scalar(select(MessageFile).where(MessageFile.message_id == message.id).limit(1)) + end_user_session_id = self._end_user_session_id(session, message, event.tenant_id) + error = payload.error or None + metadata: dict[str, Any] = { + **self._correlation_metadata(event), + "message_id": message.id, + "tool_name": payload.tool_name, + "tool_inputs": payload.tool_inputs, + "tool_outputs": payload.tool_outputs, + "tool_config": payload.tool_config, + "time_cost": payload.time_cost, + "error": error, + "tool_parameters": payload.tool_inputs, + "user_id": message.from_end_user_id or message.from_account_id, + "end_user_session_id": end_user_session_id, + } + if message.workflow_run_id: + metadata["workflow_run_id"] = message.workflow_run_id + file_url = None + if message_file is not None: + file_url = f"{self._files_url}/{message_file.url}" + metadata.update( + { + "message_file_id": message_file.id, + "created_by_role": message_file.created_by_role, + "created_user_id": message_file.created_by, + "type": message_file.type, + } + ) + return ToolTraceInfo( + trace_id=event.correlation.external_trace_id, + message_id=message.id, + message_data=_trace_message(message), + tool_name=payload.tool_name, + start_time=payload.started_at, + end_time=payload.ended_at, + tool_inputs=payload.tool_inputs, + tool_outputs=payload.tool_outputs, + metadata=metadata, + error=error, + inputs=message.message, + outputs=message.answer, + tool_config=payload.tool_config, + time_cost=payload.time_cost, + tool_parameters=payload.tool_inputs, + file_url=file_url, + ) + + def _conversation_name(self, event: TraceEvent, payload: ConversationNamed) -> GenerateNameTraceInfo: + with tracing_store_errors(), self._sessions() as session: + if event.app_id is not None: + self._app(session, event) + message = ( + self._message_in_app(session, event, payload.message_id) if payload.message_id is not None else None + ) + metadata = { + **self._correlation_metadata(event), + "conversation_id": payload.conversation_id, + "tenant_id": event.tenant_id, + "app_id": event.app_id, + "user_id": event.user_id, + } + return GenerateNameTraceInfo( + trace_id=event.correlation.external_trace_id, + message_id=payload.message_id, + message_data=_trace_message(message) if message is not None else None, + conversation_id=payload.conversation_id, + inputs=payload.inputs, + outputs=payload.name, + start_time=payload.started_at, + end_time=payload.ended_at, + metadata=metadata, + tenant_id=event.tenant_id, + ) + + def _prompt_generation(self, event: TraceEvent, payload: PromptGenerated) -> PromptGenerationTraceInfo: + with tracing_store_errors(), self._sessions() as session: + if event.app_id is not None: + self._app(session, event) + app_name, workspace_name = self._names(session, event.app_id, event.tenant_id) + user_id = event.user_id or "" + metadata = { + **self._correlation_metadata(event), + "tenant_id": event.tenant_id, + "user_id": user_id, + "app_id": event.app_id or "", + "app_name": app_name or None, + "workspace_name": workspace_name or None, + "operation_type": payload.operation_type, + "model_provider": payload.model_provider, + "model_name": payload.model_name, + } + return PromptGenerationTraceInfo( + trace_id=event.correlation.external_trace_id, + inputs=payload.instruction, + outputs=payload.generated_output, + start_time=payload.started_at, + end_time=payload.ended_at, + metadata=metadata, + tenant_id=event.tenant_id, + user_id=user_id, + app_id=event.app_id, + operation_type=payload.operation_type, + instruction=payload.instruction, + prompt_tokens=payload.prompt_tokens, + completion_tokens=payload.completion_tokens, + total_tokens=payload.total_tokens, + model_provider=payload.model_provider, + model_name=payload.model_name, + latency=payload.latency, + total_price=payload.total_price, + currency=payload.currency, + error=payload.error, + ) + + def _draft_node(self, event: TraceEvent, payload: DraftNodeExecuted) -> DraftNodeExecutionTrace: + node = payload.node_execution + if node.get("tenant_id") != event.tenant_id or node.get("app_id") != event.app_id: + raise SourceNotFoundError + with tracing_store_errors(), self._sessions() as session: + app_name, workspace_name = self._names(session, event.app_id, event.tenant_id) + credential_id, credential_name = self._credential(session, node, event.tenant_id) + metadata: dict[str, Any] = { + **self._correlation_metadata(event), + "tenant_id": event.tenant_id, + "app_id": event.app_id, + "app_name": app_name, + "workspace_name": workspace_name, + "user_id": node.get("user_id"), + "invoke_from": node.get("invoke_from"), + "credential_id": credential_id, + "credential_name": credential_name, + "dataset_ids": node.get("dataset_ids"), + "dataset_names": node.get("dataset_names"), + "plugin_name": node.get("plugin_name"), + } + return DraftNodeExecutionTrace( + trace_id=event.correlation.external_trace_id, + start_time=_datetime(node.get("created_at")), + end_time=_datetime(node.get("finished_at")), + metadata=metadata, + workflow_id=str(node.get("workflow_id") or ""), + workflow_run_id=str(node.get("workflow_execution_id") or ""), + tenant_id=event.tenant_id, + node_execution_id=str(node.get("node_execution_id") or ""), + node_id=str(node.get("node_id") or ""), + node_type=str(node.get("node_type") or ""), + title=str(node.get("title") or ""), + status=str(node.get("status") or ""), + error=_text(node.get("error")), + elapsed_time=_number(node.get("elapsed_time")), + index=int(_number(node.get("index"))), + predecessor_node_id=_text(node.get("predecessor_node_id")), + total_tokens=int(_number(node.get("total_tokens"))), + total_price=_number(node.get("total_price")), + currency=_text(node.get("currency")), + model_provider=_text(node.get("model_provider")), + model_name=_text(node.get("model_name")), + prompt_tokens=_optional_int(node.get("prompt_tokens")), + completion_tokens=_optional_int(node.get("completion_tokens")), + tool_name=_text(node.get("tool_name")), + iteration_id=_text(node.get("iteration_id")), + iteration_index=_optional_int(node.get("iteration_index")), + loop_id=_text(node.get("loop_id")), + loop_index=_optional_int(node.get("loop_index")), + parallel_id=_text(node.get("parallel_id")), + node_inputs=_mapping(node.get("node_inputs")), + node_outputs=_mapping(node.get("node_outputs")), + process_data=_mapping(node.get("process_data")), + invoked_by=_user_of(metadata), + ) + + def _credential(self, session: Session, node: Mapping[str, Any], tenant_id: str) -> tuple[str | None, str]: + """Resolve the tool or model credential a node ran with, for enterprise telemetry only.""" + if not self._enterprise_details: + return _text(node.get("credential_id")), "" + credential_id = _text(node.get("credential_id")) + if credential_id: + model_class = _TOOL_PROVIDER_MODELS.get(str(node.get("credential_provider_type") or "")) + if model_class is None: + return credential_id, "" + name = session.scalar( + select(model_class.name).where(model_class.id == credential_id, model_class.tenant_id == tenant_id) + ) + return credential_id, str(name or "") + provider_name = _text(node.get("model_provider")) + if not provider_name: + return None, "" + provider = session.scalar( + select(Provider).where( + Provider.tenant_id == tenant_id, + Provider.provider_name == provider_name, + Provider.provider_type == ProviderType.CUSTOM, + ) + ) + if provider is None: + return None, "" + model_name = _text(node.get("model_name")) + model = ( + session.scalar( + select(ProviderModel).where( + ProviderModel.tenant_id == tenant_id, + ProviderModel.provider_name == provider_name, + ProviderModel.model_name == model_name, + ProviderModel.model_type == "llm", + ) + ) + if model_name + else None + ) + if model is not None and model.credential_id: + name = session.scalar( + select(ProviderModelCredential.credential_name).where( + ProviderModelCredential.id == model.credential_id, ProviderModelCredential.tenant_id == tenant_id + ) + ) + return model.credential_id, str(name or "") + if provider.credential_id: + name = session.scalar( + select(ProviderCredential.credential_name).where( + ProviderCredential.id == provider.credential_id, ProviderCredential.tenant_id == tenant_id + ) + ) + return provider.credential_id, str(name or "") + return None, "" + + +def _text(value: object) -> str | None: + return str(value) if value is not None and value != "" else None + + +def _number(value: object) -> float: + return float(value) if isinstance(value, int | float) and not isinstance(value, bool) else 0.0 + + +def _optional_int(value: object) -> int | None: + return int(value) if isinstance(value, int | float) and not isinstance(value, bool) else None + + +def _mapping(value: object) -> Mapping[str, Any] | None: + return value if isinstance(value, dict) else None + + +def _datetime(value: object) -> datetime | None: + if isinstance(value, str): + try: + return datetime.fromisoformat(value) + except ValueError: + return None + return None diff --git a/api/repositories/tracing_store_errors.py b/api/repositories/tracing_store_errors.py new file mode 100644 index 00000000000000..c57344e6c2402a --- /dev/null +++ b/api/repositories/tracing_store_errors.py @@ -0,0 +1,16 @@ +"""Database failures of the tracing repositories, reported as the tracing slice's error.""" + +from collections.abc import Generator +from contextlib import contextmanager + +from sqlalchemy.exc import SQLAlchemyError + +from services.tracing.errors import TracingStoreUnavailableError + + +@contextmanager +def tracing_store_errors() -> Generator[None]: + try: + yield + except SQLAlchemyError as error: + raise TracingStoreUnavailableError from error diff --git a/api/schedule/collect_agent_sandbox_usage.py b/api/schedule/collect_agent_sandbox_usage.py index 27a7223ead330d..66e5bdd2ac91e0 100644 --- a/api/schedule/collect_agent_sandbox_usage.py +++ b/api/schedule/collect_agent_sandbox_usage.py @@ -19,7 +19,6 @@ class _CollectionResponse(BaseModel): @shared_task( name="schedule.collect_agent_sandbox_usage.collect_agent_sandbox_usage", - queue="ops_trace", soft_time_limit=270, time_limit=300, ignore_result=True, diff --git a/api/services/agent/composer_candidates.py b/api/services/agent/composer_candidates.py index 72fbd1528081f3..4feea48111c9fb 100644 --- a/api/services/agent/composer_candidates.py +++ b/api/services/agent/composer_candidates.py @@ -21,7 +21,7 @@ from collections.abc import Callable, Mapping from typing import Any -from core.workflow.nodes.agent_v2.discriminator import is_dify_agent_node_data +from domain.workflow.agent_node import is_dify_agent_node_data from models.agent_config_entities import ( AgentSoulConfig, DeclaredOutputConfig, diff --git a/api/services/agent/dsl_service.py b/api/services/agent/dsl_service.py index 5012f54e0eafcf..4defd6da3148f8 100644 --- a/api/services/agent/dsl_service.py +++ b/api/services/agent/dsl_service.py @@ -18,8 +18,8 @@ from sqlalchemy import func, select, update from sqlalchemy.orm import Session -from core.workflow.nodes.agent_v2.discriminator import is_dify_agent_node_data from core.workflow.nodes.agent_v2.validators import WorkflowAgentNodeValidator +from domain.workflow.agent_node import is_dify_agent_node_data from libs.datetime_utils import naive_utc_now from models import Account from models.agent import ( diff --git a/api/services/app_tracing_config_gateway.py b/api/services/app_tracing_config_gateway.py deleted file mode 100644 index c513cb215f943e..00000000000000 --- a/api/services/app_tracing_config_gateway.py +++ /dev/null @@ -1,499 +0,0 @@ -"""Compatibility gateway for the legacy tracing provider implementations.""" - -import logging -from typing import Any, override - -from pydantic import ValidationError - -from core.ops.entities.config_entity import BaseTracingConfig, TracingProviderEnum -from core.ops.exceptions import TraceProviderNotInstalledError -from core.ops.ops_trace_manager import OpsTraceManager, TracingProviderConfigEntry, provider_config_map -from services.app_tracing_config_service import ( - AppTracingConfigInvalidConfigurationError, - AppTracingConfigInvalidProviderError, - AppTracingConfigProcessingError, - AppTracingConfigProviderUnavailableError, - AppTracingConfigVerificationFailedError, - TracingConfigProviderGateway, -) - -logger = logging.getLogger(__name__) - -_PROJECT_URL_FALLBACKS = { - "phoenix": "https://app.phoenix.arize.com/projects/", - "langsmith": "https://smith.langchain.com/", - "opik": "https://www.comet.com/opik/", - "weave": "https://wandb.ai/", - "aliyun": "https://arms.console.aliyun.com/", - "tencent": "https://console.cloud.tencent.com/apm", - "mlflow": "http://localhost:5000/", - "databricks": "https://www.databricks.com/", -} - - -class OpsTraceManagerGateway(TracingConfigProviderGateway): - @override - def validate_provider(self, tracing_provider: str) -> None: - if tracing_provider not in TracingProviderEnum: - raise AppTracingConfigInvalidProviderError(tracing_provider) - - def require_provider_available(self, tracing_provider: str) -> None: - self._provider_config(tracing_provider) - - @override - def validate_saved_config( - self, *, workspace_id: str, tracing_provider: str, tracing_config: dict[str, Any] - ) -> None: - provider_config = self._provider_config(tracing_provider) - try: - decrypted = OpsTraceManager.decrypt_tracing_config(workspace_id, tracing_provider, tracing_config) - except Exception as error: - raise AppTracingConfigProcessingError from error - self._validate_config(provider_config, decrypted) - - @override - def prepare_new_config( - self, - *, - workspace_id: str, - tracing_provider: str, - tracing_config: dict[str, Any], - ) -> dict[str, Any]: - provider_config = self._provider_config(tracing_provider) - if tracing_provider == "tencent": - normalized_config = self._normalize_tencent_config(tracing_config) - elif tracing_provider == "mlflow": - normalized_config = self._normalize_mlflow_config(tracing_config) - elif tracing_provider == "phoenix": - normalized_config = self._normalize_phoenix_config(tracing_config) - elif tracing_provider == "aliyun": - normalized_config = self._normalize_aliyun_config(tracing_config) - else: - normalized_config = self._normalize_config(provider_config, tracing_config) - self._verify_config(normalized_config, tracing_provider) - - project_url = ( - None if tracing_provider == "arize" else self._get_project_url(normalized_config, tracing_provider) - ) - try: - encrypted_config = OpsTraceManager.encrypt_tracing_config( - workspace_id, - tracing_provider, - normalized_config, - ) - except Exception as error: - raise AppTracingConfigProcessingError from error - if tracing_provider == "arize": - self._resolve_arize_destination(encrypted_config, normalized_config) - elif project_url: - encrypted_config["project_url"] = project_url - return encrypted_config - - @override - def prepare_updated_config( - self, - *, - workspace_id: str, - tracing_provider: str, - tracing_config: dict[str, Any], - current_tracing_config: dict[str, Any] | None, - ) -> dict[str, Any]: - provider_config = self._provider_config(tracing_provider) - if tracing_provider == "tencent": - # Check the original URL before the legacy model discards its path - # and query. Masked tokens are restored by encryption below. - tracing_config = self._normalize_tencent_config(tracing_config) - elif tracing_provider == "mlflow": - tracing_config = self._normalize_mlflow_config(tracing_config) - elif tracing_provider == "phoenix": - tracing_config = self._normalize_phoenix_config(tracing_config) - elif tracing_provider == "aliyun": - tracing_config = self._normalize_aliyun_config(tracing_config) - self._validate_config(provider_config, tracing_config) - try: - encrypted_config = OpsTraceManager.encrypt_tracing_config( - workspace_id, - tracing_provider, - dict(tracing_config), - current_tracing_config, - ) - decrypted_config = OpsTraceManager.decrypt_tracing_config( - workspace_id, - tracing_provider, - encrypted_config, - ) - except Exception as error: - raise AppTracingConfigProcessingError from error - - self._verify_config(decrypted_config, tracing_provider) - if tracing_provider == "weave": - # Verification resolves an optional entity and qualified project name. - # Encryption happened first to restore masked credentials; persist the - # verified non-secret destination alongside that unchanged ciphertext. - encrypted_config["entity"] = decrypted_config["entity"] - encrypted_config["project"] = decrypted_config["project"] - elif tracing_provider == "databricks": - encrypted_config["host"] = decrypted_config["host"] - encrypted_config["experiment_id"] = decrypted_config["experiment_id"] - elif tracing_provider == "mlflow": - encrypted_config["tracking_uri"] = decrypted_config["tracking_uri"] - encrypted_config["experiment_id"] = decrypted_config["experiment_id"] - elif tracing_provider == "phoenix": - encrypted_config["endpoint"] = decrypted_config["endpoint"] - encrypted_config["project"] = decrypted_config["project"] - elif tracing_provider == "opik": - for key in ("url", "workspace", "project"): - encrypted_config[key] = decrypted_config[key] - # Re-resolve after destination changes; never retain a previous project's link. - project_url = None if tracing_provider == "arize" else self._get_project_url(decrypted_config, tracing_provider) - if tracing_provider == "arize": - self._resolve_arize_destination(encrypted_config, decrypted_config) - elif project_url: - encrypted_config["project_url"] = project_url - return encrypted_config - - @override - def present_config( - self, - *, - workspace_id: str, - tracing_provider: str, - tracing_config: dict[str, Any] | None, - ) -> dict[str, Any]: - if tracing_config is None: - raise AppTracingConfigProcessingError - - try: - decrypted_config = OpsTraceManager.decrypt_tracing_config( - workspace_id, - tracing_provider, - tracing_config, - ) - presented_config = OpsTraceManager.obfuscated_decrypt_token(tracing_provider, decrypted_config) - except TraceProviderNotInstalledError as error: - raise AppTracingConfigProviderUnavailableError from error - except Exception as error: - raise AppTracingConfigProcessingError from error - - if tracing_provider == "arize": - from dify_trace_arize_phoenix.arize_configuration import project_url - - # The encrypted record carries server-owned non-secret metadata. - # Legacy decryption intentionally retains only provider input fields. - presented_config["project_url"] = project_url(tracing_config) - elif tracing_provider == "aliyun": - from dify_trace_aliyun.configuration import console_url - - # GET never constructs the legacy queue/exporter or probes credentials. - presented_config["project_url"] = console_url() - elif tracing_provider == "tencent": - from dify_trace_tencent.configuration import console_url - - # Historical names remain readable. GET performs no connection - # check and never constructs legacy tracing or metrics exporters. - presented_config["project_url"] = console_url() - elif tracing_provider == "opik": - from dify_trace_opik.config import OpikConfig - from dify_trace_opik.configuration import console_url, project_url - - try: - config = OpikConfig.model_validate(decrypted_config) - presented_config["project_url"] = project_url(config) or console_url(config) - except ValueError: - presented_config["project_url"] = _PROJECT_URL_FALLBACKS["opik"] - elif tracing_provider == "phoenix": - from dify_trace_arize_phoenix.config import PhoenixConfig - from dify_trace_arize_phoenix.configuration import project_url - - try: - presented_config["project_url"] = project_url(PhoenixConfig.model_validate(decrypted_config)) - except ValueError: - # Historical destinations stay readable; never initialize a - # legacy exporter or invent a link for an unsafe base URL. - presented_config["project_url"] = None - elif tracing_provider == "mlflow": - from dify_trace_mlflow.config import MLflowConfig - from dify_trace_mlflow.oss_configuration import project_url - - try: - presented_config["project_url"] = project_url(MLflowConfig.model_validate(decrypted_config)) - except ValueError: - # OSS deployments have no universal project destination. Never - # initialize an SDK or substitute localhost for an invalid URL. - presented_config["project_url"] = None - elif tracing_provider == "databricks": - from dify_trace_mlflow.config import DatabricksConfig - from dify_trace_mlflow.configuration import project_url - - try: - presented_config["project_url"] = project_url(DatabricksConfig.model_validate(decrypted_config)) - except ValueError: - presented_config["project_url"] = _PROJECT_URL_FALLBACKS["databricks"] - elif tracing_provider == "weave": - from dify_trace_weave.config import WeaveConfig - from dify_trace_weave.configuration import console_url, project_url - - try: - config = WeaveConfig.model_validate(decrypted_config) - presented_config["project_url"] = project_url(config) or console_url(config) - except ValueError: - presented_config["project_url"] = _PROJECT_URL_FALLBACKS["weave"] - elif tracing_provider == "langfuse" and not decrypted_config.get("project_key"): - try: - project_key = OpsTraceManager.get_trace_config_project_key(decrypted_config, tracing_provider) - presented_config["project_url"] = f"{decrypted_config.get('host')}/project/{project_key}" - except Exception: - presented_config["project_url"] = f"{decrypted_config.get('host')}/" - elif tracing_provider in _PROJECT_URL_FALLBACKS and not decrypted_config.get("project_url"): - try: - presented_config["project_url"] = OpsTraceManager.get_trace_config_project_url( - decrypted_config, - tracing_provider, - ) - except Exception: - presented_config["project_url"] = _PROJECT_URL_FALLBACKS[tracing_provider] - - return presented_config - - @staticmethod - def _resolve_arize_destination(encrypted_config: dict[str, Any], verified_config: dict[str, Any]) -> None: - from dify_trace_arize_phoenix.arize_configuration import ( - DESTINATION_KEY, - bind_destination, - project_url, - resolve_project_destination, - ) - - # Only the explicit write path may replace this cache. Never trust IDs - # or URLs submitted by a client, or retain a previous target on failure. - encrypted_config.pop(DESTINATION_KEY, None) - encrypted_config.pop("project_url", None) - destination = resolve_project_destination(verified_config) - if destination is not None: - encrypted_config[DESTINATION_KEY] = bind_destination(encrypted_config, destination) - encrypted_config["project_url"] = project_url(encrypted_config) - - @staticmethod - def _provider_config(tracing_provider: str) -> TracingProviderConfigEntry: - try: - return provider_config_map[tracing_provider] - except KeyError as error: - raise AppTracingConfigInvalidProviderError(tracing_provider) from error - except TraceProviderNotInstalledError as error: - raise AppTracingConfigProviderUnavailableError from error - - @staticmethod - def _normalize_config( - provider_config: TracingProviderConfigEntry, - tracing_config: dict[str, Any], - ) -> dict[str, Any]: - normalized_config = dict(tracing_config) - default_config = OpsTraceManagerGateway._validate_config(provider_config, normalized_config) - - default_values = default_config.model_dump() - for key in provider_config["other_keys"]: - if normalized_config.get(key) == "": - normalized_config[key] = default_values.get(key) - return normalized_config - - @staticmethod - def _normalize_aliyun_config(tracing_config: dict[str, Any]) -> dict[str, Any]: - from dify_trace_aliyun.configuration import normalize_configuration - - try: - return normalize_configuration(tracing_config).model_dump() - except ValueError: - raise AppTracingConfigInvalidConfigurationError from None - - @staticmethod - def _normalize_tencent_config(tracing_config: dict[str, Any]) -> dict[str, Any]: - from dify_trace_tencent.configuration import TencentConfigurationError, normalize_configuration - - try: - return normalize_configuration(tracing_config).model_dump() - except TencentConfigurationError as error: - raise AppTracingConfigInvalidConfigurationError from error - - @staticmethod - def _normalize_mlflow_config(tracing_config: dict[str, Any]) -> dict[str, Any]: - from dify_trace_mlflow.oss_configuration import normalize_configuration - - try: - return normalize_configuration(tracing_config).model_dump() - except ValueError as error: - raise AppTracingConfigInvalidConfigurationError from error - - @staticmethod - def _normalize_phoenix_config(tracing_config: dict[str, Any]) -> dict[str, Any]: - from dify_trace_arize_phoenix.configuration import normalize_configuration - - try: - return normalize_configuration(tracing_config).model_dump() - except ValueError as error: - raise AppTracingConfigInvalidConfigurationError from error - - @staticmethod - def _validate_config( - provider_config: TracingProviderConfigEntry, - tracing_config: dict[str, Any], - ) -> BaseTracingConfig: - config_class: type[BaseTracingConfig] = provider_config["config_class"] - try: - return config_class.model_validate(tracing_config) - except ValidationError as error: - # Aliyun credentials can be embedded in a pasted endpoint. Do not - # propagate a Pydantic error object whose errors() includes inputs. - if config_class.__module__ == "dify_trace_aliyun.config": - raise AppTracingConfigInvalidConfigurationError from None - raise AppTracingConfigInvalidConfigurationError from error - - @staticmethod - def _verify_config(tracing_config: dict[str, Any], tracing_provider: str) -> None: - if tracing_provider == "arize": - from dify_trace_arize_phoenix.arize_configuration import ArizeConfigurationError, verify_configuration - from dify_trace_arize_phoenix.config import ArizeConfig - - try: - verify_configuration(ArizeConfig.model_validate(tracing_config)) - except ArizeConfigurationError as error: - logger.warning("Arize ingestion verification failed: error_code=%s", error.code) - raise AppTracingConfigVerificationFailedError from None - return - if tracing_provider == "aliyun": - from dify_trace_aliyun.config import AliyunConfig - from dify_trace_aliyun.configuration import AliyunConfigurationError, verify_configuration - - try: - verify_configuration(AliyunConfig.model_validate(tracing_config)) - except AliyunConfigurationError as error: - logger.warning("Aliyun configuration verification failed: error_code=%s", error.code) - raise AppTracingConfigVerificationFailedError from None - return - if tracing_provider == "phoenix": - from dify_trace_arize_phoenix.config import PhoenixConfig - from dify_trace_arize_phoenix.configuration import PhoenixConfigurationError, verify_configuration - - try: - resolved = verify_configuration(PhoenixConfig.model_validate(tracing_config)) - except PhoenixConfigurationError as error: - logger.warning("Phoenix configuration verification failed: error_code=%s", error.code) - raise AppTracingConfigVerificationFailedError from error - tracing_config["endpoint"] = resolved.endpoint - tracing_config["project"] = resolved.project - return - if tracing_provider == "mlflow": - from dify_trace_mlflow.config import MLflowConfig - from dify_trace_mlflow.oss_configuration import MLflowConfigurationError, verify_configuration - - try: - resolved = verify_configuration(MLflowConfig.model_validate(tracing_config)) - except MLflowConfigurationError as error: - logger.warning("MLflow configuration verification failed: error_code=%s", error.code) - raise AppTracingConfigVerificationFailedError from error - tracing_config["tracking_uri"] = resolved.tracking_uri - tracing_config["experiment_id"] = resolved.experiment_id - return - if tracing_provider == "tencent": - from dify_trace_tencent.config import TencentConfig - from dify_trace_tencent.configuration import TencentConfigurationError, verify_configuration - - try: - verify_configuration(TencentConfig.model_validate(tracing_config)) - except TencentConfigurationError as error: - logger.warning("Tencent collector connection check failed: error_code=%s", error.code) - raise AppTracingConfigVerificationFailedError from error - return - if tracing_provider == "opik": - from dify_trace_opik.config import OpikConfig - from dify_trace_opik.configuration import OpikConfigurationError, verify_configuration - - try: - resolved = verify_configuration(OpikConfig.model_validate(tracing_config)) - except OpikConfigurationError as error: - logger.warning("Opik configuration verification failed: error_code=%s", error.code) - raise AppTracingConfigVerificationFailedError from error - tracing_config["url"] = resolved.url - tracing_config["workspace"] = resolved.workspace - tracing_config["project"] = resolved.project - return - if tracing_provider == "databricks": - from dify_trace_mlflow.config import DatabricksConfig - from dify_trace_mlflow.configuration import DatabricksConfigurationError, verify_configuration - - try: - resolved = verify_configuration(DatabricksConfig.model_validate(tracing_config)) - except DatabricksConfigurationError as error: - logger.warning("Databricks configuration verification failed: error_code=%s", error.code) - raise AppTracingConfigVerificationFailedError from error - tracing_config["host"] = resolved.host - tracing_config["experiment_id"] = resolved.experiment_id - return - if tracing_provider == "weave": - from dify_trace_weave.config import WeaveConfig - from dify_trace_weave.configuration import WeaveConfigurationError, verify_configuration - - try: - resolved = verify_configuration(WeaveConfig.model_validate(tracing_config)) - except WeaveConfigurationError as error: - logger.warning("Weave configuration verification failed: error_code=%s", error.code) - raise AppTracingConfigVerificationFailedError from error - tracing_config["entity"] = resolved.entity - tracing_config["project"] = resolved.project - return - try: - is_effective = OpsTraceManager.check_trace_config_is_effective(tracing_config, tracing_provider) - except ValueError as error: - logger.warning("Tracing configuration verification failed for provider %s", tracing_provider, exc_info=True) - raise AppTracingConfigVerificationFailedError from error - if not is_effective: - raise AppTracingConfigVerificationFailedError - - @staticmethod - def _get_project_url(tracing_config: dict[str, Any], tracing_provider: str) -> str | None: - try: - if tracing_provider == "aliyun": - from dify_trace_aliyun.configuration import console_url - - return console_url() - if tracing_provider == "tencent": - from dify_trace_tencent.configuration import console_url - - return console_url() - if tracing_provider == "opik": - from dify_trace_opik.config import OpikConfig - from dify_trace_opik.configuration import console_url, project_url - - config = OpikConfig.model_validate(tracing_config) - return project_url(config) or console_url(config) - if tracing_provider == "databricks": - from dify_trace_mlflow.config import DatabricksConfig - from dify_trace_mlflow.configuration import project_url - - return project_url(DatabricksConfig.model_validate(tracing_config)) - if tracing_provider == "mlflow": - from dify_trace_mlflow.config import MLflowConfig - from dify_trace_mlflow.oss_configuration import project_url - - return project_url(MLflowConfig.model_validate(tracing_config)) - if tracing_provider == "weave": - from dify_trace_weave.config import WeaveConfig - from dify_trace_weave.configuration import project_url - - return project_url(WeaveConfig.model_validate(tracing_config)) - if tracing_provider == "phoenix": - from dify_trace_arize_phoenix.config import PhoenixConfig - from dify_trace_arize_phoenix.configuration import project_url - - return project_url(PhoenixConfig.model_validate(tracing_config)) - if tracing_provider == "arize": - from dify_trace_arize_phoenix.arize_configuration import project_url - - return project_url(tracing_config) - if tracing_provider == "langfuse": - project_key = OpsTraceManager.get_trace_config_project_key(tracing_config, tracing_provider) - return f"{tracing_config.get('host')}/project/{project_key}" - if tracing_provider in ("langsmith", "opik"): - return OpsTraceManager.get_trace_config_project_url(tracing_config, tracing_provider) - except Exception: - return None - return None diff --git a/api/services/app_tracing_config_service.py b/api/services/app_tracing_config_service.py index 11d3b1028ff413..366fd58b21dd21 100644 --- a/api/services/app_tracing_config_service.py +++ b/api/services/app_tracing_config_service.py @@ -24,6 +24,14 @@ class AppTracingConfigRecord: is_active: bool created_at: datetime updated_at: datetime + configured_secrets: tuple[str, ...] = () + """Secret fields that hold a value. Secrets themselves are never presented.""" + + +@dataclass(frozen=True, slots=True) +class PresentedTracingConfig: + values: dict[str, Any] + configured_secrets: tuple[str, ...] class AppTracingConfigStore(Protocol): @@ -99,7 +107,7 @@ def present_config( workspace_id: str, tracing_provider: str, tracing_config: dict[str, Any] | None, - ) -> dict[str, Any]: ... + ) -> PresentedTracingConfig: ... class AppTracingConfigError(Exception): @@ -178,12 +186,12 @@ def get( if record is None: return None - tracing_config = self._provider.present_config( + presented = self._provider.present_config( workspace_id=workspace_id, tracing_provider=tracing_provider, tracing_config=record.tracing_config, ) - return replace(record, tracing_config=tracing_config) + return replace(record, tracing_config=presented.values, configured_secrets=presented.configured_secrets) def create( self, diff --git a/api/services/conversation_name_generator.py b/api/services/conversation_name_generator.py index 441ceae8eaf157..ffc3a395c19ea8 100644 --- a/api/services/conversation_name_generator.py +++ b/api/services/conversation_name_generator.py @@ -1,5 +1,7 @@ """Prepare a conversation-title model separately from its invocation.""" +from typing import override + from flask import current_app from core.app.entities.app_invoke_entities import get_credit_usage_app_type @@ -9,25 +11,37 @@ from core.model_manager import ModelManager from extensions.ext_database import db from graphon.model_runtime.entities.model_entities import ModelType +from services.conversation_name_service import ConversationNameGenerator +from services.tracing.scopes import UNCORRELATED, TraceScopes + + +class DefaultModelConversationNameGenerator(ConversationNameGenerator): + """Name conversations with the workspace's default model and report each naming to its app's trace.""" + def __init__(self, *, trace_scopes: TraceScopes) -> None: + self._trace_scopes = trace_scopes -def generate_conversation_name(*, tenant_id: str, app_id: str, conversation_id: str, query: str, app_mode: str) -> str: - """Release model-preparation sessions before LLM I/O, preserving the caller's session.""" - # Provider resolution may commit; own its scope so it cannot affect the caller. - with ( - current_app.app_context(), - use_credit_usage_metadata( - {"app_type": get_credit_usage_app_type(app_mode), "created_by": CreditUsageCreatedBy.CONVERSATION_NAME} - ), - ): - model = ModelManager.for_tenant(tenant_id=tenant_id).get_default_model_instance( - tenant_id=tenant_id, model_type=ModelType.LLM - ) - db.session.remove() - return LLMGenerator.generate_conversation_name( - tenant_id=tenant_id, - app_id=app_id, - conversation_id=conversation_id, - query=query, - model_instance=model, - ) + @override + def __call__(self, *, tenant_id: str, app_id: str, conversation_id: str, query: str, app_mode: str) -> str: + """Release model-preparation sessions before LLM I/O, preserving the caller's session.""" + # Provider resolution may commit; own its scope so it cannot affect the caller. + with ( + current_app.app_context(), + use_credit_usage_metadata( + {"app_type": get_credit_usage_app_type(app_mode), "created_by": CreditUsageCreatedBy.CONVERSATION_NAME} + ), + ): + model = ModelManager.for_tenant(tenant_id=tenant_id).get_default_model_instance( + tenant_id=tenant_id, model_type=ModelType.LLM + ) + db.session.remove() + return LLMGenerator.generate_conversation_name( + tenant_id=tenant_id, + app_id=app_id, + conversation_id=conversation_id, + query=query, + model_instance=model, + trace_scope=self._trace_scopes.for_app( + tenant_id=tenant_id, app_id=app_id, user_id=None, correlation=UNCORRELATED + ), + ) diff --git a/api/services/message_suggested_questions_generator.py b/api/services/message_suggested_questions_generator.py index 696e0aa8d435a1..ca2036097b5873 100644 --- a/api/services/message_suggested_questions_generator.py +++ b/api/services/message_suggested_questions_generator.py @@ -7,22 +7,22 @@ from functools import partial from typing import TYPE_CHECKING, cast +from dify_trace_sdk.events import SuggestedQuestionsGenerated from flask import current_app from core.app.entities.app_invoke_entities import CreditUsageCreatedBy, get_credit_usage_app_type from core.llm_generator.llm_generator import _normalize_completion_params from core.model_context import use_credit_usage_metadata, with_credit_usage_created_by from core.model_manager import ModelInstance, ModelManager -from core.ops.entities.trace_entity import TraceTaskName -from core.ops.ops_trace_manager import TraceQueueManager, TraceTask -from core.ops.utils import measure_time from core.plugin.impl.base import use_plugin_daemon_request_timeout from core.prompt.utils.prompt_template_parser import PromptTemplateParser from extensions.ext_database import db from graphon.model_runtime.entities.llm_entities import LLMResult from graphon.model_runtime.entities.message_entities import PromptMessage, UserPromptMessage from graphon.model_runtime.entities.model_entities import ModelType, ParameterType +from libs.datetime_utils import measure_time from services.message_suggested_questions_parser import SuggestedQuestionsAfterAnswerOutputParser +from services.tracing.scopes import UNCORRELATED, TraceScopes if TYPE_CHECKING: from core.memory.token_buffer_memory import PreparedHistory @@ -79,6 +79,9 @@ def _default_suggested_questions_model_parameters(model_instance: ModelInstance) class SuggestedQuestionsGenerator: """Own model resolution, detached-history rendering, generation and tracing.""" + def __init__(self, *, trace_scopes: TraceScopes) -> None: + self._trace_scopes = trace_scopes + @contextmanager def prepare( self, @@ -122,7 +125,7 @@ def generate( ): model = self._prepare_model(tenant_id=context.tenant_id, model_config=model_config) db.session.remove() - questions = ( + questions: list[str] = ( list(self._invoke(prepared_model=model, histories=histories, instruction_prompt=instruction_prompt)) if model is not None else [] @@ -131,14 +134,18 @@ def generate( # Invocation catches provider failures, including failed database # transactions. Tracing always starts with a fresh scoped session. db.session.remove() - TraceQueueManager(app_id=context.app_id).add_trace_task( - TraceTask( - TraceTaskName.SUGGESTED_QUESTION_TRACE, - message_id=context.message_id, - suggested_question=questions, - timer=timer, - ) + trace_scope = self._trace_scopes.for_app( + tenant_id=context.tenant_id, app_id=context.app_id, user_id=None, correlation=UNCORRELATED ) + if trace_scope is not None: + trace_scope.emit( + SuggestedQuestionsGenerated( + message_id=context.message_id, + questions=questions, + started_at=timer.started_at, + ended_at=timer.ended_at, + ) + ) return questions @staticmethod diff --git a/api/services/tracing/adapters.py b/api/services/tracing/adapters.py new file mode 100644 index 00000000000000..b4d4a024b5f906 --- /dev/null +++ b/api/services/tracing/adapters.py @@ -0,0 +1,173 @@ +"""Runtime adapters of the trace pipeline over Celery, storage, workspace keys, and enterprise telemetry. + +Events are handed to a bounded background publisher; large ones travel through +storage instead of the broker. Dated payload keys let the sweeper expire files +that a crashed delivery left behind. +""" + +import atexit +import logging +import os +from collections.abc import Callable, Sequence +from concurrent.futures import ThreadPoolExecutor +from datetime import UTC, datetime +from threading import BoundedSemaphore +from typing import override +from uuid import uuid4 + +from celery import current_app as celery_app +from celery.signals import worker_process_shutdown +from dify_trace_sdk.emitter import TraceEmitter, TraceScope +from dify_trace_sdk.events import TraceCorrelation, TraceEvent + +from core.helper.encrypter import batch_decrypt_token, encrypt_token, obfuscated_token +from enterprise.telemetry.enterprise_trace import EnterpriseOtelTrace +from extensions.ext_storage import storage +from extensions.tracing import ExecutionTraceScopes +from repositories.app.tracing_config_repository import trace_settings_from_column +from services.tracing.delivery import ExportJob, ExportRetryQueue +from services.tracing.providers import TracingSecretCipher +from services.tracing.records import BaseTraceInfo +from services.tracing.scopes import TraceScopes + +DELIVER_TASK = "tasks.trace_delivery_task.deliver_trace_event" +RETRY_TASK = "tasks.trace_delivery_task.retry_trace_export" +AGENT_TRACE_EXPORT_TASK = "tasks.application_trace_task.export_application_trace" +PAYLOAD_PREFIX = "trace_payloads" +INLINE_EVENT_BYTES = 256 * 1024 +PENDING_EVENT_LIMIT = 256 +# Node executions may be saved by another worker after their run row; give them time to land. +WORKFLOW_SETTLE_SECONDS = 5 + +logger = logging.getLogger(__name__) + + +def payload_key(day: datetime, kind: str, identifier: str) -> str: + """Storage key for a trace payload; the date prefix lets the sweeper expire it.""" + return f"{PAYLOAD_PREFIX}/{day:%Y-%m-%d}/{kind}/{identifier}.json" + + +class CeleryTraceEmitter(TraceEmitter): + """Best-effort publication without storage or broker I/O on producer threads. + + A full buffer drops the new event with a warning. Normal shutdown drains + accepted events; a process crash can lose events not yet published. + """ + + def __init__(self) -> None: + self._closed = False + self._reset_publisher() + os.register_at_fork(after_in_child=self._after_fork) + atexit.register(self.shutdown) + worker_process_shutdown.connect(self.shutdown, dispatch_uid=f"trace-publisher-{id(self)}") + + def _reset_publisher(self) -> None: + self._slots = BoundedSemaphore(PENDING_EVENT_LIMIT) + self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="trace-publisher") + + def _after_fork(self) -> None: + # Parent threads and pending events belong to the parent process. + if not self._closed: + self._reset_publisher() + + @override + def emit(self, event: TraceEvent) -> None: + if not self._slots.acquire(blocking=False): + logger.warning( + "Trace publisher buffer full; event dropped: event_id=%s app_id=%s", event.event_id, event.app_id + ) + return + try: + # Snapshot mutable payload fields and the caller's Celery app before handoff. + encoded = event.model_dump_json() + self._executor.submit(self._publish, event, encoded, celery_app.send_task) + except Exception: + self._slots.release() + raise + + def _publish(self, event: TraceEvent, encoded: str, send_task: Callable[..., object]) -> None: + try: + data = encoded.encode() + if len(data) > INLINE_EVENT_BYTES: + key = payload_key(event.occurred_at, "events", event.event_id) + storage.save(key, data) + envelope = {"payload_key": key} + else: + envelope = {"event": encoded} + countdown = WORKFLOW_SETTLE_SECONDS if event.payload.kind == "workflow_run" else 0 + send_task(DELIVER_TASK, args=(envelope,), countdown=countdown) + except Exception as error: + logger.warning( + "Trace publication failed: event_id=%s app_id=%s error_type=%s", + event.event_id, + event.app_id, + type(error).__name__, + ) + finally: + self._slots.release() + + def shutdown(self, **_: object) -> None: + """Drain accepted events at interpreter or Celery child-process shutdown.""" + self._closed = True + self._executor.shutdown(wait=True) + + +class StoredPolicyTraceScopes(ExecutionTraceScopes): + """Read the raw stored policy that legacy producers hold, as the policy repository reads it.""" + + def __init__(self, scopes: TraceScopes) -> None: + self._scopes = scopes + + @override + def for_execution( + self, + *, + tenant_id: str, + app_id: str | None, + app_tracing: str | None, + user_id: str | None, + correlation: TraceCorrelation, + ) -> TraceScope | None: + return self._scopes.for_execution( + tenant_id=tenant_id, + app_id=app_id, + policy=trace_settings_from_column(app_tracing), + user_id=user_id, + correlation=correlation, + ) + + @override + def for_app( + self, *, tenant_id: str, app_id: str | None, user_id: str | None, correlation: TraceCorrelation + ) -> TraceScope | None: + return self._scopes.for_app(tenant_id=tenant_id, app_id=app_id, user_id=user_id, correlation=correlation) + + +class CeleryExportRetryQueue(ExportRetryQueue): + @override + def schedule(self, job: ExportJob, *, delay_seconds: int) -> None: + key = payload_key(datetime.now(UTC), "exports", uuid4().hex) + storage.save(key, job.model_dump_json().encode()) + celery_app.send_task(RETRY_TASK, args=(key,), countdown=delay_seconds) + + +class WorkspaceTracingSecretCipher(TracingSecretCipher): + @override + def encrypt(self, tenant_id: str, value: str) -> str: + return encrypt_token(tenant_id, value) + + @override + def decrypt(self, tenant_id: str, tokens: Sequence[str]) -> list[str]: + return batch_decrypt_token(tenant_id, list(tokens)) + + @override + def mask(self, value: str) -> str: + return obfuscated_token(value) + + +def schedule_agent_trace_export(run_id: str) -> None: + celery_app.send_task(AGENT_TRACE_EXPORT_TASK, args=(run_id,)) + + +def send_to_enterprise_telemetry(record: BaseTraceInfo) -> None: + EnterpriseOtelTrace().trace(record) diff --git a/api/services/tracing/canonical.py b/api/services/tracing/canonical.py new file mode 100644 index 00000000000000..0f5bcb2ffd4253 --- /dev/null +++ b/api/services/tracing/canonical.py @@ -0,0 +1,443 @@ +"""Build provider-neutral traces from assembled records. + +Identity is derived, never coordinated. A trace is named by the caller's trace +ID, else by its root record: the message for chats and chatflows, the workflow +run for workflow apps. Each span is named by the Dify record it represents. +Auxiliary operations (tool calls, retrieval, moderation, suggested questions, +conversation naming) join their message's trace as separate units under its +root span; a workflow run nested in another one is its own trace, linked +through attributes. +""" + +import math +from collections.abc import Mapping +from decimal import Decimal +from typing import Any + +from dify_trace_sdk.events import json_safe +from dify_trace_sdk.ids import span_id_for, trace_id_for +from dify_trace_sdk.model import Cost, Span, SpanKind, SpanRef, SpanStatus, TokenUsage, Trace +from pydantic import JsonValue + +from domain.workflow.retry_history import RETRY_HISTORY_PROCESS_DATA_KEY +from graphon.entities import WorkflowNodeExecution +from graphon.enums import WorkflowNodeExecutionStatus +from services.tracing.hierarchy import build_workflow_hierarchy, execution_metadata, finished_at +from services.tracing.records import ( + BaseTraceInfo, + DatasetRetrievalTraceInfo, + GenerateNameTraceInfo, + MessageTraceInfo, + ModerationTraceInfo, + SuggestedQuestionTraceInfo, + ToolTraceInfo, + WorkflowTraceInfo, +) + +_NODE_KIND: dict[str, SpanKind] = { + "llm": SpanKind.LLM, + "knowledge-retrieval": SpanKind.RETRIEVER, + "tool": SpanKind.TOOL, + "agent": SpanKind.AGENT, +} +_RETRY_SUMMARY_FIELDS = ("retry_index", "error", "elapsed_time", "created_at", "finished_at") + + +def message_trace_id(external_trace_id: str | None, message_id: str) -> str: + return trace_id_for(external_trace_id or message_id) + + +def _status(error: str | None, status: object) -> SpanStatus: + failed = str(status) in {WorkflowNodeExecutionStatus.FAILED.value, WorkflowNodeExecutionStatus.EXCEPTION.value} + return SpanStatus.ERROR if error or failed else SpanStatus.OK + + +def _cost(amount: Decimal | float | None, currency: str | None) -> Cost | None: + if amount is None or not currency: + return None + value = float(amount) + if not math.isfinite(value) or value < 0: + return None + return Cost(amount=value, currency=currency) + + +def _usage(prompt: object, completion: object, total: object) -> TokenUsage | None: + def count(value: object) -> int | None: + return int(value) if isinstance(value, int | float) and not isinstance(value, bool) and value >= 0 else None + + prompt_tokens, completion_tokens = count(prompt), count(completion) + if prompt_tokens is None and completion_tokens is None: + return None + total_tokens = count(total) + return TokenUsage( + prompt_tokens=prompt_tokens or 0, + completion_tokens=completion_tokens or 0, + total_tokens=total_tokens if total_tokens is not None else (prompt_tokens or 0) + (completion_tokens or 0), + ) + + +def _attributes(metadata: Mapping[str, Any]) -> dict[str, JsonValue]: + result = json_safe({key: value for key, value in metadata.items() if value is not None}) + return result if isinstance(result, dict) else {} + + +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 _session_id(record: BaseTraceInfo, fallback: str | None) -> str | None: + explicit = record.metadata.get("trace_session_id") + return explicit if isinstance(explicit, str) and explicit else fallback + + +def _user_id(record: BaseTraceInfo) -> str | None: + """The caller's own end-user identifier when there is one, else Dify's user ID.""" + value = record.metadata.get("end_user_session_id") or record.metadata.get("user_id") + return str(value) if value else None + + +def _seconds(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, int | float): + return None + return float(value) if math.isfinite(value) and value >= 0 else None + + +class CanonicalTraceBuilder: + """Convert records into canonical trace units; ``None`` means nothing to send.""" + + def build(self, record: BaseTraceInfo) -> Trace | None: + match record: + case WorkflowTraceInfo(): + return self._workflow(record) + case MessageTraceInfo(): + return self._message(record) + case ModerationTraceInfo(): + return self._auxiliary( + record, + name="moderation", + kind=SpanKind.TOOL, + inputs=record.inputs, + outputs={"action": record.action, "flagged": record.flagged}, + error=None, + ) + case SuggestedQuestionTraceInfo(): + return self._auxiliary( + record, + name="suggested_question", + kind=SpanKind.TOOL, + inputs=record.inputs, + outputs=record.suggested_question, + error=record.error, + ) + case DatasetRetrievalTraceInfo(): + return self._auxiliary( + record, + name="dataset_retrieval", + kind=SpanKind.RETRIEVER, + inputs=record.inputs, + outputs={"documents": record.documents}, + error=record.error, + ) + case ToolTraceInfo(): + return self._auxiliary( + record, + name=record.tool_name, + kind=SpanKind.TOOL, + inputs=record.tool_inputs, + outputs=record.tool_outputs, + error=record.error or None, + ) + case GenerateNameTraceInfo(): + return self._auxiliary( + record, + name="generate_name", + kind=SpanKind.TOOL, + inputs=record.inputs, + outputs=record.outputs, + error=None, + ) + case _: + return None + + def _message(self, record: MessageTraceInfo) -> Trace: + message = record.message_data + root_span_id = span_id_for(message.id) + started_at = record.start_time or message.created_at + ended_at = max(record.end_time or started_at, started_at) + error = record.error or message.error or None + status = _status(error, None) + attributes = _attributes({**record.metadata, "trace_entity_type": "message"}) + return Trace( + trace_id=message_trace_id(record.trace_id, message.id), + external_trace_id=record.trace_id, + session_id=_session_id(record, message.conversation_id), + user_id=_user_id(record), + parent=None, + spans=( + Span( + span_id=root_span_id, + parent_span_id=None, + name="message", + kind=SpanKind.CHAIN, + start_time=started_at, + end_time=ended_at, + status=status, + error=error, + inputs=json_safe(record.inputs), + outputs=json_safe(message.answer), + attributes=attributes, + ), + Span( + span_id=span_id_for(f"{message.id}:llm"), + parent_span_id=root_span_id, + name="llm", + kind=SpanKind.LLM, + start_time=started_at, + end_time=ended_at, + status=status, + error=error, + inputs=json_safe(record.inputs), + outputs=json_safe(record.outputs if record.outputs is not None else message.answer), + model_name=message.model_id, + model_provider=message.model_provider, + usage=_usage(record.message_tokens, record.answer_tokens, record.total_tokens), + cost=_cost(message.total_price, message.currency), + time_to_first_token=_seconds(record.gen_ai_server_time_to_first_token), + attributes=attributes, + ), + ), + ) + + def _workflow(self, record: WorkflowTraceInfo) -> Trace: + executions = record.node_executions + hierarchy = build_workflow_hierarchy(executions, record.workflow_tool_parents) + workflow_start = record.start_time + workflow_end = max(record.end_time, workflow_start) + workflow_span_id = span_id_for(record.workflow_run_id) + workflow_status = _status(record.error, record.workflow_run_status) + spans: list[Span] = [] + span_id_by_execution: dict[str, str] = {} + + if record.message_id: + # A chatflow run is shown under its message, which is the trace root. + root_logical_id = record.message_id + root_span_id = span_id_for(record.message_id) + # Units attached to this message (tool calls, retrieval, ...) name its creation as the trace start. + message_start = record.message_started_at or workflow_start + spans.append( + Span( + span_id=root_span_id, + parent_span_id=None, + name="message", + kind=SpanKind.CHAIN, + start_time=message_start, + end_time=max(record.end_time, message_start), + status=workflow_status, + error=record.error or None, + inputs=json_safe(record.query or dict(record.workflow_run_inputs)), + outputs=json_safe(dict(record.workflow_run_outputs)), + attributes=_attributes({**record.metadata, "trace_entity_type": "message"}), + ) + ) + workflow_parent: str | None = root_span_id + else: + root_logical_id = record.workflow_run_id + workflow_parent = None + + spans.append( + Span( + span_id=workflow_span_id, + parent_span_id=workflow_parent, + name="workflow", + kind=SpanKind.CHAIN, + start_time=workflow_start, + end_time=workflow_end, + status=workflow_status, + error=record.error or None, + inputs=json_safe(dict(record.workflow_run_inputs)), + outputs=json_safe(dict(record.workflow_run_outputs)), + attributes=_attributes( + { + **record.metadata, + "workflow_id": record.workflow_id, + "workflow_run_id": record.workflow_run_id, + "workflow_app_log_id": record.workflow_app_log_id, + "total_tokens": record.total_tokens, + } + ), + ) + ) + + for execution in executions: + span_id_by_execution[execution.id] = span_id_for(execution.id) + for wrapper in hierarchy.wrappers: + span_id_by_execution[wrapper.id] = span_id_for(wrapper.id) + + unordered: dict[str, Span] = {} + for wrapper in hierarchy.wrappers: + unordered[wrapper.id] = Span( + span_id=span_id_by_execution[wrapper.id], + parent_span_id=span_id_by_execution.get(wrapper.parent_execution_id, workflow_span_id), + name=f"{wrapper.key.kind}[{wrapper.key.index}]", + kind=SpanKind.CHAIN, + start_time=wrapper.start_time, + end_time=max(wrapper.end_time, wrapper.start_time), + status=SpanStatus.ERROR if wrapper.has_error else SpanStatus.OK, + error="wrapper child failed" if wrapper.has_error else None, + attributes={ + "wrapper_type": wrapper.key.kind, + "wrapper_index": wrapper.key.index, + "container_execution_id": wrapper.parent_execution_id, + }, + ) + for execution in executions: + parent_logical = hierarchy.parent_by_execution_id.get(execution.id) + unordered[execution.id] = self._node_span( + execution, + parent_span_id=span_id_by_execution.get(parent_logical, workflow_span_id) + if parent_logical + else workflow_span_id, + ) + + spans.extend(self._parent_first(unordered, {span.span_id for span in spans})) + parent_run_id, parent_node_id = record.resolved_parent_context + if parent_run_id: + spans[0] = spans[0].model_copy( + update={ + "attributes": { + **spans[0].attributes, + "parent_workflow_run_id": parent_run_id, + "parent_node_execution_id": parent_node_id, + } + } + ) + return Trace( + trace_id=trace_id_for(record.trace_id or root_logical_id), + external_trace_id=record.trace_id, + session_id=_session_id(record, record.conversation_id or parent_run_id or record.workflow_run_id), + user_id=_user_id(record), + parent=None, + spans=tuple(spans), + ) + + @staticmethod + def _node_span(execution: WorkflowNodeExecution, *, parent_span_id: str) -> Span: + process_data = dict(execution.process_data or {}) + outputs = dict(execution.outputs or {}) + node_type = str(execution.node_type) + metadata = execution_metadata(execution) + usage_source = process_data.get("usage") or outputs.get("usage") or {} + usage = ( + _usage( + usage_source.get("prompt_tokens"), + usage_source.get("completion_tokens"), + usage_source.get("total_tokens"), + ) + if isinstance(usage_source, Mapping) + else None + ) + kind = _NODE_KIND.get(node_type, SpanKind.CHAIN) + time_to_first_token = usage_source.get("time_to_first_token") if isinstance(usage_source, Mapping) else None + return Span( + span_id=span_id_for(execution.id), + parent_span_id=parent_span_id, + name=f"{node_type}_{execution.title}" if execution.title else node_type, + kind=kind, + start_time=execution.created_at, + end_time=max(finished_at(execution), execution.created_at), + status=_status(execution.error, execution.status), + error=execution.error or None, + inputs=json_safe(process_data.get("prompts", []) if kind is SpanKind.LLM else dict(execution.inputs or {})), + outputs=json_safe(outputs), + model_name=process_data.get("model_name") if kind is SpanKind.LLM else None, + model_provider=process_data.get("model_provider") if kind is SpanKind.LLM else None, + usage=usage if kind is SpanKind.LLM else None, + cost=_cost(metadata.get("total_price"), metadata.get("currency")) if kind is SpanKind.LLM else None, + time_to_first_token=_seconds(time_to_first_token) if kind is SpanKind.LLM else None, + attributes=_attributes( + { + **metadata, + "node_id": execution.node_id, + "node_execution_id": execution.id, + "node_type": node_type, + "status": str(execution.status), + **_retry_metadata(process_data), + } + ), + ) + + @staticmethod + def _parent_first(spans: Mapping[str, Span], emitted_ids: set[str]) -> list[Span]: + ordered: list[Span] = [] + pending = sorted(spans.values(), key=lambda span: (span.start_time, span.span_id)) + while pending: + ready = [span for span in pending if span.parent_span_id in emitted_ids] + if not ready: + # The hierarchy only names existing executions or wrappers as parents. + raise ValueError("workflow span hierarchy is not a tree") + for span in ready: + ordered.append(span) + emitted_ids.add(span.span_id) + pending = [span for span in pending if span.span_id not in emitted_ids] + return ordered + + def _auxiliary( + self, + record: BaseTraceInfo, + *, + name: str, + kind: SpanKind, + inputs: object, + outputs: object, + error: str | None, + ) -> Trace | None: + message = record.message_data + operation_id = record.metadata.get("trace_event_id") + if not isinstance(operation_id, str): + return None + started_at = record.start_time or (message.created_at if message is not None else None) + if started_at is None: + return None + ended_at = max(record.end_time or started_at, started_at) + span_id = span_id_for(f"event:{operation_id}") + if message is None: + trace_id = trace_id_for(record.trace_id or f"event:{operation_id}") + parent = None + parent_span_id = None + else: + trace_id = message_trace_id(record.trace_id, message.id) + parent = SpanRef(span_id=span_id_for(message.id), start_time=message.created_at) + parent_span_id = parent.span_id + return Trace( + trace_id=trace_id, + external_trace_id=record.trace_id, + session_id=_session_id(record, message.conversation_id if message is not None else None), + user_id=_user_id(record), + parent=parent, + spans=( + Span( + span_id=span_id, + parent_span_id=parent_span_id, + name=name or kind.value, + kind=kind, + start_time=started_at, + end_time=ended_at, + status=SpanStatus.ERROR if error else SpanStatus.OK, + error=error, + inputs=json_safe(inputs), + outputs=json_safe(outputs), + attributes=_attributes(record.metadata), + ), + ), + ) diff --git a/api/services/tracing/collection.py b/api/services/tracing/collection.py new file mode 100644 index 00000000000000..f0280fb3fe8280 --- /dev/null +++ b/api/services/tracing/collection.py @@ -0,0 +1,152 @@ +"""Records of Agent trace collection and its settings. + +Agent apps and workflows with Agent nodes produce spans that the Dify Agent +runtime reports separately. A run freezes the app's destination and capture +policy when its execution starts, collects runtime batches in a ledger, and +is exported once its business terminal and every invocation are complete. +""" + +import json +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import datetime +from enum import StrEnum +from typing import Any + +from domain.tracing.sources import SourceNodeExecution, SourceWorkflowRun +from domain.tracing.trace_run_agg import TraceRunState, TraceSourceKind +from domain.workflow.agent_node import is_dify_agent_node_data +from services.tracing.providers import TraceProviderSupport + + +def workflow_has_agent_v2(graph: object) -> bool: + """Scope Agent trace collection to graphs containing the Dify Agent v2 node.""" + if isinstance(graph, str): + try: + graph = json.loads(graph) + except ValueError: + return False + if not isinstance(graph, Mapping) or not isinstance(graph.get("nodes"), list): + return False + return any( + isinstance(node, Mapping) and isinstance(node.get("data"), Mapping) and is_dify_agent_node_data(node["data"]) + for node in graph["nodes"] + ) + + +class TraceAppMode(StrEnum): + """The app modes collection distinguishes; every other mode is ``OTHER``.""" + + AGENT = "agent" + ADVANCED_CHAT = "advanced-chat" + WORKFLOW = "workflow" + OTHER = "other" + + +@dataclass(frozen=True, slots=True) +class TraceRun: + id: str + tenant_id: str + app_id: str + source_kind: TraceSourceKind + source_id: str + provider: str + """``""`` when the run froze no destination.""" + config_revision: str + """``""`` when the run froze no destination.""" + provider_config: Mapping[str, Any] + """The frozen stored configuration, secrets still encrypted.""" + trace_id: str + root_span_id: str + session_id: str + capture_content: bool + capture_agent: bool + state: TraceRunState + outer_trace: Mapping[str, Any] | None + """The captured business terminal; ``None`` until it is captured.""" + span_count: int + attempts: int + error_code: str | None + """``None`` unless the run failed or its export outcome is uncertain.""" + trace_url: str | None + """``None`` until the destination accepts the trace with a link.""" + provider_trace_id: str | None + """``None`` until the destination accepts the trace.""" + span_urls: Mapping[str, str] + accepted_at: datetime | None + """``None`` until the destination accepts the trace.""" + + +@dataclass(frozen=True, slots=True) +class TraceApp: + id: str + tenant_id: str + mode: TraceAppMode + workflow_id: str | None + """``None`` for apps without a published workflow.""" + + +@dataclass(frozen=True, slots=True) +class SourceMessage: + id: str + conversation_id: str + workflow_run_id: str | None + """``None`` until a chatflow run is linked to the message.""" + agent_stop_requested: bool + """Whether the consumer asked to stop the Agent; it does not say whether the stop won.""" + + +type TraceSource = SourceMessage | SourceWorkflowRun | SourceNodeExecution + + +@dataclass(frozen=True, slots=True) +class AppTracingSnapshot: + """What a run freezes from its app: the stored tracing policy and the provider configuration it names.""" + + policy: Mapping[str, Any] + """The stored policy as written; empty when it is missing or unreadable.""" + config: Mapping[str, Any] | None + """The stored configuration of the provider the policy names, secrets still encrypted. + + ``None`` when the policy names no provider or that provider has no stored configuration. + """ + + +@dataclass(frozen=True, slots=True) +class TraceRunStatus: + """The delivery state of a source's frozen run. + + Every field after ``source_id`` is ``None`` when nothing was collected. + """ + + status: TraceRunState + source_kind: TraceSourceKind + source_id: str + provider: str | None + trace_id: str | None + """The destination's trace ID once accepted, otherwise the run's own.""" + trace_url: str | None + span_url: str | None + """The link of the selected node's Agent span; ``None`` when no node is selected.""" + project_url: str | None + """The frozen destination's project link; never resolved from the app's current configuration.""" + config_revision: str | None + capture_content: bool | None + capture_agent: bool | None + span_count: int | None + attempts: int | None + error_code: str | None + accepted_at: datetime | None + capabilities: TraceProviderSupport | None + """``None`` when the frozen provider is not installed here.""" + + +@dataclass(frozen=True, slots=True) +class AgentTraceSettings: + enabled: bool + max_spans: int + max_bytes: int + collection_timeout_seconds: int + export_lease_seconds: int + max_export_attempts: int + retention_days: int diff --git a/api/services/tracing/collection_ports.py b/api/services/tracing/collection_ports.py new file mode 100644 index 00000000000000..473edcd678cc9c --- /dev/null +++ b/api/services/tracing/collection_ports.py @@ -0,0 +1,107 @@ +"""Ports of Agent trace collection: the run repository, run and source queries, and export scheduling. + +Reads and writes are scoped by the tenant and app that own the run. Every +operation raises TracingStoreUnavailableError when the database cannot be +reached. Each ledger operation runs in its own transaction. Concurrent writers +are coordinated by the unique source key of a run, by row locks taken for every +state transition, and by the export token that fences one export attempt. +""" + +from collections.abc import Callable +from contextlib import AbstractContextManager +from dataclasses import dataclass +from typing import Protocol + +from domain.tracing.sources import MessageProgress, SourceNodeExecution, SourceWorkflowRun +from domain.tracing.trace_run_agg import TraceBatch, TraceInvocation, TraceRunAgg, TraceSourceKind +from services.tracing.collection import AppTracingSnapshot, SourceMessage, TraceApp, TraceRun + +type ScheduleTraceExport = Callable[[str], object] + + +@dataclass(frozen=True, slots=True) +class LockedTraceRun: + run: TraceRunAgg + owner_exists: bool + """Whether the app that owns the run still exists in its workspace.""" + + +class TraceRunRepository(Protocol): + """Whole Agent trace runs. + + The unique source key keeps the first run frozen for a source. A run's row + lock is held from loading it in ``change`` until its changes are saved. + """ + + def add(self, run: TraceRunAgg) -> bool: + """Insert a newly frozen run; ``False`` when another run already holds its source key.""" + ... + + def load(self, run_id: str) -> TraceRunAgg | None: + """Read a run without its lock, to decide what to change; ``None`` when it does not exist.""" + ... + + def change(self, run_id: str, *, tenant_id: str) -> AbstractContextManager[LockedTraceRun | None]: + """Load a run under its row lock and save it when the block exits without an error. + + ``None`` when the run does not exist in the supplied tenant. The lock is held until the block + exits, so the block only changes the run. A new batch whose ID another + run already stores raises TraceCollectionConflictError. + """ + ... + + def purge(self, *, retention_days: int, limit: int) -> int: + """Delete up to ``limit`` whole runs created more than ``retention_days`` ago; returns how many.""" + ... + + +class TraceRunQuery(Protocol): + def get(self, run_id: str) -> TraceRun | None: ... + + def source_run( + self, *, tenant_id: str, app_id: str, source_kind: TraceSourceKind, source_id: str + ) -> TraceRun | None: ... + + def workflow_run_trace(self, *, tenant_id: str, workflow_run_id: str) -> TraceRun | None: + """The run collected for a workflow run of any app in the workspace.""" + ... + + def invocations(self, run_id: str) -> list[TraceInvocation]: + """The run's invocations by attempt, then creation.""" + ... + + def batches(self, run_id: str) -> list[TraceBatch]: + """The run's stored batches by arrival.""" + ... + + def pending_ids(self, limit: int) -> list[str]: + """Runs that may still export and wait for no backoff, those without a terminal last.""" + ... + + +class TraceSourceQuery(Protocol): + """The records a run is collected for.""" + + def app(self, *, tenant_id: str, app_id: str) -> TraceApp | None: ... + + def app_tracing(self, *, tenant_id: str, app_id: str) -> AppTracingSnapshot | None: + """``None`` when the app is not in the workspace.""" + ... + + def app_workflow_graph(self, *, tenant_id: str, app_id: str, workflow_id: str) -> str | None: + """The graph of the app's workflow definition; ``None`` when the definition does not exist.""" + ... + + def message(self, *, tenant_id: str, app_id: str, message_id: str) -> SourceMessage | None: ... + + def workflow_run(self, *, tenant_id: str, app_id: str, workflow_run_id: str) -> SourceWorkflowRun | None: + """Read through the configured workflow run backend, outside any open transaction.""" + ... + + def node_execution(self, *, tenant_id: str, app_id: str, node_execution_id: str) -> SourceNodeExecution | None: ... + + def workflow_run_conversation_id(self, *, tenant_id: str, app_id: str, workflow_run_id: str) -> str | None: + """The conversation of the chatflow message a run answers; ``None`` outside chatflows.""" + ... + + def message_progress(self, *, tenant_id: str, app_id: str, message_id: str) -> MessageProgress | None: ... diff --git a/api/services/tracing/delivery.py b/api/services/tracing/delivery.py new file mode 100644 index 00000000000000..8727861a3093d0 --- /dev/null +++ b/api/services/tracing/delivery.py @@ -0,0 +1,249 @@ +"""Deliver trace events to every sink that consumes them. + +Sinks are independent: enterprise telemetry receives every record, the Agent +trace collector owns records of Agent-managed executions, and the app's +tracing provider receives the rest. Only a provider send is retried: a +retryable failure persists the canonical trace (never the record) and +schedules that send alone. The configuration is read when sending, so a +rotated or removed credential takes effect immediately. +""" + +import logging +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Any, Protocol + +from dify_trace_sdk.config import ProviderConfig +from dify_trace_sdk.egress import Egress +from dify_trace_sdk.errors import ExportOutcomeUnknownError, ExportRejectedError, ExportRetryableError +from dify_trace_sdk.events import TraceEvent +from dify_trace_sdk.model import Trace +from dify_trace_sdk.provider import TraceProvider +from opentelemetry import metrics +from pydantic import BaseModel, ConfigDict, ValidationError + +from services.tracing.canonical import CanonicalTraceBuilder +from services.tracing.errors import ( + ProviderConfigUnreadableError, + TraceCollectorUnavailableError, + TraceProviderUnavailableError, + UnknownTraceProviderError, +) +from services.tracing.providers import ProviderConfigCodec, TraceProviderCatalog +from services.tracing.records import BaseTraceInfo, MessageTraceInfo, WorkflowTraceInfo + +logger = logging.getLogger(__name__) + +_meter = metrics.get_meter(__name__) +_event_counter = _meter.create_counter("dify.tracing.events", description="Trace events by kind and outcome") +_export_counter = _meter.create_counter("dify.tracing.exports", description="Provider sends by provider and outcome") + + +class TraceRecordSource(Protocol): + def assemble(self, event: TraceEvent) -> BaseTraceInfo: + """Read the event's record through its tenant and app. + + Workflow records include the admitted root's workflow-tool descendants, + scoped to the same tenant and run and linked by their persisted callers. + + Raises: + SourceNotFoundError: The record does not exist for the event's tenant and app. + SourceNotReadyError: The record is not persisted or not finished yet. + """ + ... + + +@dataclass(frozen=True, slots=True) +class StoredTraceDestination: + provider_id: str + config: Mapping[str, Any] + """The stored configuration, secrets still encrypted.""" + + +class TraceDestinationQuery(Protocol): + def active_destination(self, *, tenant_id: str, app_id: str) -> StoredTraceDestination | None: + """The configuration of the provider the app traces to, read through the app's workspace. + + ``None`` when the app is not in the workspace, tracing is disabled, or the + enabled provider has no active configuration. + """ + ... + + +class TraceCollector(Protocol): + def capture_terminal(self, *, tenant_id: str, info: BaseTraceInfo) -> bool: + """Take ownership of a record of an Agent-managed execution; ``False`` leaves it to the provider.""" + ... + + def link_managed_parent(self, *, tenant_id: str, info: WorkflowTraceInfo) -> WorkflowTraceInfo: + """Name the Agent-collected parent that a provider-exported nested workflow cannot attach to.""" + ... + + +type EnterpriseTraceSink = Callable[[BaseTraceInfo], None] + + +class ExportJob(BaseModel): + """A provider send waiting for its next attempt.""" + + tenant_id: str + app_id: str + provider_id: str + attempt: int + trace: Trace + + model_config = ConfigDict(extra="forbid", frozen=True) + + +class ExportRetryQueue(Protocol): + def schedule(self, job: ExportJob, *, delay_seconds: int) -> None: + """Persist the job and run it once the delay passes; ``TraceDeliveryService.retry`` receives it.""" + ... + + +@dataclass(frozen=True) +class Destination: + provider: TraceProvider[Any] + config: ProviderConfig + + +def retry_delay_seconds(attempt: int) -> int: + return min(600, 10 * 2 ** (attempt - 1)) + + +class TraceDeliveryService: + def __init__( + self, + *, + records: TraceRecordSource, + destinations: TraceDestinationQuery, + providers: TraceProviderCatalog, + codec: ProviderConfigCodec, + builder: CanonicalTraceBuilder, + egress: Egress, + retries: ExportRetryQueue, + enterprise_sink: EnterpriseTraceSink | None, + collector: TraceCollector | None, + max_export_attempts: int, + ) -> None: + """``enterprise_sink`` and ``collector`` are ``None`` when the deployment does not run that sink.""" + self._records = records + self._destinations = destinations + self._providers = providers + self._codec = codec + self._builder = builder + self._egress = egress + self._retries = retries + self._enterprise_sink = enterprise_sink + self._collector = collector + self._max_export_attempts = max_export_attempts + + def deliver(self, event: TraceEvent, *, enterprise_delivered: bool) -> None: + """Raises SourceNotReadyError (redeliver later) or TraceCollectorUnavailableError.""" + kind = event.payload.kind + try: + record = self._records.assemble(event) + except ValidationError: + logger.warning("Trace record not assemblable: kind=%s app_id=%s", kind, event.app_id) + _event_counter.add(1, {"kind": kind, "outcome": "invalid"}) + return + if self._enterprise_sink is not None and not enterprise_delivered: + try: + self._enterprise_sink(record) + except Exception: + logger.exception("Enterprise trace failed: kind=%s app_id=%s", kind, event.app_id) + if self._collector is not None: + try: + if self._collector.capture_terminal(tenant_id=event.tenant_id, info=record): + _event_counter.add(1, {"kind": kind, "outcome": "collected"}) + return + if isinstance(record, WorkflowTraceInfo): + record = self._collector.link_managed_parent(tenant_id=event.tenant_id, info=record) + except Exception: + raise TraceCollectorUnavailableError from None + if isinstance(record, MessageTraceInfo) and record.message_data.workflow_run_id: + # A chatflow message is the root of its workflow trace; sending it again would duplicate it. + _event_counter.add(1, {"kind": kind, "outcome": "covered"}) + return + if event.app_id is None: + return + destination = self.destination(event.tenant_id, event.app_id) + if destination is None: + _event_counter.add(1, {"kind": kind, "outcome": "no_destination"}) + return + try: + trace = self._builder.build(record) + except ValueError: + # A record the builder cannot represent would fail the same way on every redelivery. + logger.warning("Trace record not buildable: kind=%s app_id=%s", kind, event.app_id, exc_info=True) + _event_counter.add(1, {"kind": kind, "outcome": "invalid"}) + return + if trace is None: + _event_counter.add(1, {"kind": kind, "outcome": "not_exported"}) + return + _event_counter.add(1, {"kind": kind, "outcome": "exported"}) + self.export( + ExportJob( + tenant_id=event.tenant_id, + app_id=event.app_id, + provider_id=destination.provider.id, + attempt=1, + trace=trace, + ), + destination, + ) + + def retry(self, job: ExportJob) -> None: + destination = self.destination(job.tenant_id, job.app_id) + if destination is None or destination.provider.id != job.provider_id: + # The app now traces elsewhere or not at all; the old destination must not receive data. + _export_counter.add(1, {"provider": job.provider_id, "outcome": "destination_changed"}) + return + self.export(job, destination) + + def export(self, job: ExportJob, destination: Destination) -> None: + provider_id = job.provider_id + try: + destination.provider.send(destination.config, job.trace) + except ExportRetryableError as error: + if job.attempt >= self._max_export_attempts: + logger.warning("Trace export retries exhausted: provider=%s app_id=%s", provider_id, job.app_id) + _export_counter.add(1, {"provider": provider_id, "outcome": "retries_exhausted"}) + return + logger.info("Trace export will retry: provider=%s attempt=%s reason=%s", provider_id, job.attempt, error) + _export_counter.add(1, {"provider": provider_id, "outcome": "retrying"}) + self._retries.schedule( + job.model_copy(update={"attempt": job.attempt + 1}), delay_seconds=retry_delay_seconds(job.attempt) + ) + except ExportOutcomeUnknownError as error: + logger.warning( + "Trace export outcome unknown: provider=%s app_id=%s reason=%s", provider_id, job.app_id, error + ) + _export_counter.add(1, {"provider": provider_id, "outcome": "outcome_unknown"}) + except ExportRejectedError as error: + logger.warning("Trace export rejected: provider=%s app_id=%s reason=%s", provider_id, job.app_id, error) + _export_counter.add(1, {"provider": provider_id, "outcome": "rejected"}) + except Exception as error: + # Unclassified failures may carry request details, so only their type is logged. + logger.warning( + "Trace export failed: provider=%s app_id=%s error=%s", provider_id, job.app_id, type(error).__name__ + ) + _export_counter.add(1, {"provider": provider_id, "outcome": "error"}) + else: + _export_counter.add(1, {"provider": provider_id, "outcome": "accepted"}) + + def destination(self, tenant_id: str, app_id: str) -> Destination | None: + """The app's current provider destination, read through the app's workspace.""" + stored = self._destinations.active_destination(tenant_id=tenant_id, app_id=app_id) + if stored is None: + return None + try: + provider = self._providers.resolve(stored.provider_id) + config = self._codec.decode(tenant_id, provider, stored.config) + except (UnknownTraceProviderError, TraceProviderUnavailableError): + logger.warning("Tracing provider unavailable: provider=%s app_id=%s", stored.provider_id, app_id) + return None + except ProviderConfigUnreadableError: + logger.warning("Tracing configuration unreadable: provider=%s app_id=%s", stored.provider_id, app_id) + return None + return Destination(provider=provider(self._egress), config=config) diff --git a/api/services/tracing/egress.py b/api/services/tracing/egress.py new file mode 100644 index 00000000000000..2fd08f35c9f8b2 --- /dev/null +++ b/api/services/tracing/egress.py @@ -0,0 +1,143 @@ +"""The network policy trace providers run under. + +HTTP goes through the SSRF proxy owner without retries or redirects, with a +bounded response. gRPC cannot traverse the SSRF proxy (it only tunnels to +port 443), so channels are opened directly to a resolved and pinned address +after rejecting loopback, link-local, and other non-routable destinations. +Private networks stay reachable for VPC collector endpoints. +""" + +import ipaddress +import socket +from collections.abc import Mapping +from typing import override +from urllib.parse import urlsplit + +import grpc # type: ignore[import-untyped] +import httpx +from dify_trace_sdk.egress import Egress, EgressResponse +from dify_trace_sdk.errors import EgressBlockedError, EgressError +from opentelemetry.context import _SUPPRESS_INSTRUMENTATION_KEY, attach, detach, set_value + +from core.helper import ssrf_proxy +from core.tools.errors import ToolSSRFError +from services.tracing.request_log_protection import protect_http_request_logs + +MAX_RESPONSE_BYTES = 8 * 1024 * 1024 +_GRPC_MAX_RECEIVE_BYTES = 4 * 1024 * 1024 + + +def _blocked_address(address: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool: + if isinstance(address, ipaddress.IPv6Address) and address.ipv4_mapped is not None: + address = address.ipv4_mapped + return ( + address.is_loopback + or address.is_link_local + or address.is_unspecified + or address.is_multicast + or address.is_reserved + ) + + +class SSRFEgress(Egress): + @override + def request( + self, + method: str, + url: str, + *, + headers: Mapping[str, str], + content: bytes | None, + timeout: float, + sensitive_url: bool, + ) -> EgressResponse: + # Bounded buffering needs identity bodies, whatever encodings an SDK offers. + request_headers = {name: value for name, value in headers.items() if name.lower() != "accept-encoding"} + request_headers["Accept-Encoding"] = "identity" + if not any(name.lower() == "cookie" for name in request_headers): + # The shared SSRF client keeps a cookie jar; an explicit header keeps other tenants' cookies out. + request_headers["Cookie"] = "" + # Export calls must not be traced by Dify's own instrumentation. + token = attach(set_value(_SUPPRESS_INSTRUMENTATION_KEY, True)) + try: + with protect_http_request_logs(sensitive_url): + response = self._send(method, url, request_headers, content, timeout) + finally: + detach(token) + return EgressResponse( + status_code=response.status_code, + headers={name.lower(): value for name, value in response.headers.items()}, + content=response.content, + ) + + @staticmethod + def _send(method: str, url: str, headers: dict[str, str], content: bytes | None, timeout: float) -> httpx.Response: + # Raise outside the original handler so no credential-bearing exception is chained. + failure: EgressError | None = None + response: httpx.Response | None = None + try: + streamed = ssrf_proxy.make_request( + method, + url, + max_retries=0, + stream_response=True, + timeout=timeout, + follow_redirects=False, + headers=headers, + content=content, + ) + response = ssrf_proxy.buffer_response(streamed, max_response_bytes=MAX_RESPONSE_BYTES) + except ToolSSRFError: + failure = EgressBlockedError("Destination blocked by the SSRF policy") + except httpx.ProxyError as error: + # The SSRF proxy refuses a disallowed HTTPS tunnel with 403 before anything is sent. + if "403" in str(error): + failure = EgressBlockedError("Destination blocked by the SSRF policy") + else: + failure = EgressError("SSRF proxy unavailable", request_sent=False) + except ( + httpx.ConnectError, + httpx.ConnectTimeout, + httpx.PoolTimeout, + httpx.UnsupportedProtocol, + httpx.LocalProtocolError, + ): + failure = EgressError("Destination unreachable", request_sent=False) + except (httpx.TransportError, ssrf_proxy.ResponseLimitError): + failure = EgressError("Transport failed after the request was sent", request_sent=True) + except ValueError: + # Request construction rejects malformed input before sending; its message may quote the URL. + failure = EgressError("Request could not be built", request_sent=False) + if failure is not None: + raise failure + if response is None: + raise EgressError("Transport returned no response", request_sent=True) + return response + + @override + def grpc_channel(self, endpoint: str) -> grpc.Channel: + parsed = urlsplit(endpoint) + host = parsed.hostname + if parsed.scheme not in {"http", "https"} or not host or parsed.username or parsed.password: + raise EgressBlockedError("Invalid gRPC endpoint") + try: + port = parsed.port or (443 if parsed.scheme == "https" else 80) + resolved = socket.getaddrinfo(host, port, type=socket.SOCK_STREAM) + except (OSError, ValueError): + raise EgressError("gRPC endpoint could not be resolved", request_sent=False) from None + addresses = [ipaddress.ip_address(item[4][0]) for item in resolved] + if not addresses or any(_blocked_address(address) for address in addresses): + raise EgressBlockedError("gRPC endpoint resolves to a disallowed address") + # Pin the checked address so a second DNS answer cannot redirect the channel. + pinned = addresses[0] + target = f"[{pinned}]:{port}" if pinned.version == 6 else f"{pinned}:{port}" + authority = f"[{host}]:{port}" if ":" in host else f"{host}:{port}" + options = [ + ("grpc.enable_retries", 0), + ("grpc.max_receive_message_length", _GRPC_MAX_RECEIVE_BYTES), + ("grpc.default_authority", authority), + ] + if parsed.scheme == "http": + return grpc.insecure_channel(target, options=options) + options.append(("grpc.ssl_target_name_override", host)) + return grpc.secure_channel(target, grpc.ssl_channel_credentials(), options=options) diff --git a/api/services/tracing/errors.py b/api/services/tracing/errors.py new file mode 100644 index 00000000000000..dd2c1a916b7b3c --- /dev/null +++ b/api/services/tracing/errors.py @@ -0,0 +1,59 @@ +"""Failures of the trace pipeline and of Agent trace collection. + +Errors name what went wrong inside the tracing slice. They carry no protocol +status: handlers choose the response for each type. +""" + + +class TracingError(Exception): + """Base error of the tracing slice.""" + + +class SourceNotFoundError(TracingError): + """The referenced record does not exist for the event's tenant and app.""" + + +class SourceNotReadyError(TracingError): + """The referenced record is not persisted or not finished yet.""" + + +class TraceCollectorUnavailableError(TracingError): + """The Agent trace collector could not decide ownership; the event must be redelivered.""" + + +class UnknownTraceProviderError(TracingError): + def __init__(self, provider_id: str) -> None: + super().__init__(f"Unknown tracing provider: {provider_id}") + + +class TraceProviderUnavailableError(TracingError): + def __init__(self, provider_id: str) -> None: + super().__init__(f"Tracing provider {provider_id} is not installed in this deployment") + + +class ProviderConfigUnreadableError(TracingError): + """A stored configuration cannot be decrypted or no longer validates.""" + + +class TracingStoreUnavailableError(TracingError): + """The trace ledger or a business record could not be read or written; a later attempt may succeed.""" + + +class TraceCollectionError(TracingError): + """Agent trace collection cannot proceed; ``code`` names the reason for clients and run records.""" + + def __init__(self, code: str) -> None: + self.code = code + super().__init__(code) + + +class TraceCollectionNotFoundError(TraceCollectionError): + """The source, run, invocation, or owner does not exist in the caller's scope.""" + + +class TraceCollectionConflictError(TraceCollectionError): + """The request conflicts with the current state of the collected run.""" + + +class TraceCollectionLimitError(TraceCollectionError): + """A runtime batch exceeds the run's collection limits.""" diff --git a/api/services/tracing/genai_attributes.py b/api/services/tracing/genai_attributes.py deleted file mode 100644 index 1b1abcfe8140cc..00000000000000 --- a/api/services/tracing/genai_attributes.py +++ /dev/null @@ -1,58 +0,0 @@ -"""GenAI semantic attributes shared by the Aliyun and Tencent adapters.""" - -from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalTrace -from services.tracing.provider_common import json_value, token_usage - - -def attributes(span: CanonicalSpan, trace: CanonicalTrace, provider: str) -> dict[str, object]: - kind = span.kind.value.upper() - if provider == "tencent": - kind = {"LLM": "GENERATION", "CHAIN": "WORKFLOW"}.get(kind, kind) - result: dict[str, object] = { - "gen_ai.span.kind": kind, - "gen_ai.session.id": trace.session_id, - "gen_ai.framework": "dify", - "dify.span.id": span.id, - "dify.metadata": json_value(span.metadata), - "gen_ai.entity.input" if provider == "tencent" else "input.value": json_value(span.inputs), - "gen_ai.entity.output" if provider == "tencent" else "output.value": json_value(span.outputs), - } - if provider == "tencent" and span.id == trace.root_span_id: - result["gen_ai.is_entry"] = "true" - for key in ("tenant_id", "app_id", "agent_id", "workflow_run_id", "node_execution_id", "message_id"): - value = span.metadata.get(key) or trace.spans[0].metadata.get(key) - if value: - result[f"dify.{key}"] = str(value) - user_id = span.metadata.get("user_id") or trace.spans[0].metadata.get("user_id") - if user_id: - result["gen_ai.user.id"] = str(user_id) - if span.kind is CanonicalSpanKind.LLM: - result["gen_ai.operation.name"] = "chat" - result["gen_ai.prompt"] = json_value(span.inputs) - result["gen_ai.completion"] = json_value(span.outputs) - result["gen_ai.input.messages"] = json_value(span.inputs) - result["gen_ai.output.messages"] = json_value(span.outputs) - model = span.metadata.get("model") or span.metadata.get("model_name") - if model: - result["gen_ai.request.model"] = str(model) - result["gen_ai.response.model"] = str(model) - if span.metadata.get("model_provider"): - result["gen_ai.provider.name"] = str(span.metadata["model_provider"]) - result.update({f"gen_ai.usage.{key}_tokens": value for key, value in token_usage(span).items()}) - ttft = span.metadata.get("time_to_first_token", span.metadata.get("ttft")) - if isinstance(ttft, (int, float)) and ttft >= 0: - result["gen_ai.response.time_to_first_token"] = int(ttft * 1_000_000_000) - elif span.kind is CanonicalSpanKind.AGENT: - result["gen_ai.operation.name"] = "invoke_agent" - result["gen_ai.agent.name"] = span.name - elif span.kind is CanonicalSpanKind.TOOL: - result["gen_ai.operation.name"] = "execute_tool" - result["gen_ai.tool.name"] = span.name - result["gen_ai.tool.call.arguments"] = json_value(span.inputs) - result["gen_ai.tool.call.result"] = json_value(span.outputs) - result["tool.name"] = span.name - result["tool.parameters"] = json_value(span.inputs) - elif span.kind is CanonicalSpanKind.RETRIEVER: - result["retrieval.query"] = json_value(span.inputs) - result["retrieval.document"] = json_value(span.outputs) - return result diff --git a/api/core/ops/unified_trace/hierarchy.py b/api/services/tracing/hierarchy.py similarity index 50% rename from api/core/ops/unified_trace/hierarchy.py rename to api/services/tracing/hierarchy.py index 0f08ceda597d24..77f0af85a272ea 100644 --- a/api/core/ops/unified_trace/hierarchy.py +++ b/api/services/tracing/hierarchy.py @@ -1,4 +1,9 @@ -"""Deterministically reconstruct provider-neutral workflow execution hierarchy.""" +"""Deterministically reconstruct the span hierarchy of one workflow run. + +Parents come from graph predecessors, workflow tools, and iteration/loop containers; +each iteration or loop round gets a synthetic wrapper. Ambiguous repeated node +IDs within an invocation and cyclic predecessor data are dropped rather than guessed. +""" import re from collections.abc import Mapping, Sequence @@ -7,6 +12,7 @@ from typing import Any, Literal from graphon.entities import WorkflowNodeExecution +from graphon.enums import WorkflowNodeExecutionStatus _WRAPPER_INDEX_PATTERN = re.compile(r"^[A-Za-z0-9_.:-]+$") _WRAPPER_FIELDS: tuple[tuple[Literal["iteration", "loop"], str, str], ...] = ( @@ -15,14 +21,6 @@ ) -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"] @@ -48,43 +46,40 @@ class WorkflowHierarchy: 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: WorkflowNodeExecution) -> dict[str, Any]: + return {str(key): value for key, value in (execution.metadata or {}).items()} + + +def finished_at(execution: WorkflowNodeExecution) -> datetime: + return execution.finished_at or execution.created_at + timedelta(seconds=execution.elapsed_time or 0) + + +type _NodeKey = tuple[str, str | None, str | None, str] -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 _node_key(item: WorkflowNodeExecution, node_id: str, tool_parents: Mapping[str, str]) -> _NodeKey: + invocation = (item.process_data or {}).get("workflow_tool_invocation_id") + return item.workflow_id, tool_parents.get(item.id), invocation if isinstance(invocation, str) else None, node_id -def _unique_execution_by_node_id(executions: Sequence[WorkflowExecutionLike]) -> dict[str, str]: - result: dict[str, str] = {} - ambiguous: set[str] = set() +def _unique_execution_by_node_id( + executions: Sequence[WorkflowNodeExecution], tool_parents: Mapping[str, str] +) -> dict[_NodeKey, str]: + result: dict[_NodeKey, str] = {} + ambiguous: set[_NodeKey] = set() for item in executions: - node_id = item.node_id - if not isinstance(node_id, str) or node_id in ambiguous: + key = _node_key(item, item.node_id, tool_parents) + if key in ambiguous: continue - item_execution_id = execution_id(item) - previous = result.get(node_id) + previous = result.get(key) if previous is None: - result[node_id] = item_execution_id - elif previous != item_execution_id: - result.pop(node_id, None) - ambiguous.add(node_id) + result[key] = item.id + elif previous != item.id: + result.pop(key, None) + ambiguous.add(key) 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 @@ -95,19 +90,6 @@ def _normalize_index(value: Any) -> str | None: 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): @@ -125,44 +107,43 @@ def _remove_cycles(parent_by_execution_id: dict[str, str]) -> None: 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) +def build_workflow_hierarchy( + executions: Sequence[WorkflowNodeExecution], tool_parents: Mapping[str, str] +) -> WorkflowHierarchy: + execution_by_node_id = _unique_execution_by_node_id(executions, tool_parents) 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 + execution_by_node_id.get(_node_key(item, item.predecessor_node_id, tool_parents)) + if item.predecessor_node_id is not None + 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) + container_node_id = metadata.get(structured_key) if isinstance(container_node_id, str): - parent_execution_id = execution_by_node_id.get(container_node_id) + parent_execution_id = execution_by_node_id.get(_node_key(item, container_node_id, tool_parents)) 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 + if parent_execution_id is None: + parent_execution_id = tool_parents.get(item.id) + if parent_execution_id and parent_execution_id != item.id: + parent_by_execution_id[item.id] = parent_execution_id _remove_cycles(parent_by_execution_id) - grouped: dict[WrapperKey, list[WorkflowExecutionLike]] = {} + grouped: dict[WrapperKey, list[WorkflowNodeExecution]] = {} 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)) + container_node_id = metadata.get(container_key) + index = _normalize_index(metadata.get(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): + container_execution_id = execution_by_node_id.get(_node_key(item, container_node_id, tool_parents)) + if container_execution_id is None or container_execution_id == item.id: continue wrapper_key = WrapperKey(kind=kind, container_execution_id=container_execution_id, index=index) grouped.setdefault(wrapper_key, []).append(item) @@ -170,17 +151,17 @@ def build_workflow_hierarchy(executions: Sequence[WorkflowExecutionLike]) -> Wor 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)): + for wrapper_key in sorted(grouped, key=lambda key: (key.kind, key.container_execution_id, key.index)): children = grouped[wrapper_key] - child_ids = frozenset(execution_id(item) for item in children) + child_ids = frozenset(item.id 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), + start_time=min(item.created_at for item in children), + end_time=max(finished_at(item) for item in children), + has_error=any(item.status == WorkflowNodeExecutionStatus.FAILED for item in children), ) wrappers.append(wrapper) for child_id in child_ids: diff --git a/api/services/tracing/otlp_export.py b/api/services/tracing/otlp_export.py deleted file mode 100644 index 9507e82f495834..00000000000000 --- a/api/services/tracing/otlp_export.py +++ /dev/null @@ -1,202 +0,0 @@ -"""Encode completed spans with stable IDs and explicit provider ingestion ACKs. - -No tracer or global SDK configuration is created: a replay must preserve the -original IDs and timestamps, and destination credentials belong to one export. -""" - -import re -from collections.abc import Callable, Mapping -from datetime import datetime -from typing import Any -from urllib.parse import urlsplit -from uuid import UUID - -from core.ops.exceptions import RetryableTraceDispatchError -from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanStatus, CanonicalTrace -from services.tracing.provider_common import json_value, nanos, request -from services.tracing.provider_dispatch import TraceExportRejectedError, TraceExportUncertainError - - -def _id(value: str, length: int) -> bytes: - if not re.fullmatch(rf"[0-9a-f]{{{length}}}", value) or int(value, 16) == 0: - raise TraceExportRejectedError("Application trace contains an invalid OpenTelemetry identifier") - return bytes.fromhex(value) - - -def build_request( - trace: CanonicalTrace, - resource: Mapping[str, object], - attributes: Callable[[CanonicalSpan], Mapping[str, object]], -) -> Any: - from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest - from opentelemetry.proto.common.v1.common_pb2 import AnyValue, KeyValue - from opentelemetry.proto.trace.v1.trace_pb2 import Status - - def attribute(key: str, value: object) -> Any: - if isinstance(value, bool): - encoded = AnyValue(bool_value=value) - elif isinstance(value, int): - encoded = AnyValue(int_value=value) - elif isinstance(value, float): - encoded = AnyValue(double_value=value) - elif isinstance(value, str): - encoded = AnyValue(string_value=value) - else: - encoded = AnyValue(string_value=json_value(value)) - return KeyValue(key=key, value=encoded) - - trace_id = _id(trace.trace_id, 32) - payload = ExportTraceServiceRequest() - resource_spans = payload.resource_spans.add() - resource_spans.resource.attributes.extend(attribute(k, v) for k, v in resource.items() if v is not None) - scope = resource_spans.scope_spans.add() - scope.scope.name = "dify.application.tracing" - scope.scope.version = "1" - for span in trace.spans: - if span.end_time is None or nanos(span.end_time) < nanos(span.start_time): - raise TraceExportRejectedError("Application trace contains incomplete or invalid timestamps") - target = scope.spans.add( - trace_id=trace_id, - span_id=_id(span.id, 16), - parent_span_id=_id(span.parent_id, 16) if span.parent_id else b"", - name=span.name, - kind=1, - start_time_unix_nano=nanos(span.start_time), - end_time_unix_nano=nanos(span.end_time), - flags=1, - ) - target.status.code = ( - Status.STATUS_CODE_ERROR if span.status is CanonicalSpanStatus.ERROR else Status.STATUS_CODE_OK - ) - if span.error: - target.status.message = span.error - target.attributes.extend(attribute(k, v) for k, v in attributes(span).items() if v is not None) - return payload - - -def _check_response(response: Any, provider: str) -> None: - if response.partial_success.rejected_spans: - # Retrying the entire batch would duplicate the accepted subset. - raise TraceExportUncertainError(f"{provider} partially accepted the trace; reconciliation is required") - - -def is_aliyun_http_success_ack(status_code: int, content: bytes) -> bool: - """Recognize the exact ACK observed from Aliyun's HTTP Trace collector. - - This literal response is an observed compatibility behavior, not the - documented OTLP ExportTraceServiceResponse. Do not normalize other bodies - or accept it for another status or provider. - """ - return status_code == 200 and content == b'"success"' - - -def _langfuse_enqueue_ack(body: object) -> bool: - """Recognize Langfuse's queue acknowledgement without exposing its auth data. - - The v4 HTTP route returns the job after blob upload and awaited queue.add, - rather than an OTLP ExportTraceServiceResponse. This confirms enqueue only. - Upstream: langfuse/langfuse@bc1ade0aacca75104bad934e92d70a454ec16e67, - web/src/server/otel/processOtelIngestion.ts:277 and - packages/shared/src/server/otel/OtelIngestionProcessor.ts:350. - """ - if not isinstance(body, Mapping) or body.get("name") != "otel-ingestion-job": - return False - if any(key in body for key in ("partialSuccess", "partial_success", "error", "errors")) or body.get("failedReason"): - return False - data = body.get("data") - if not isinstance(data, Mapping) or data.get("name") != "otel-ingestion-job": - return False - identifier, timestamp = data.get("id"), data.get("timestamp") - if not isinstance(identifier, str) or not isinstance(timestamp, str): - return False - try: - if UUID(identifier).int == 0 or datetime.fromisoformat(timestamp).tzinfo is None: - return False - except ValueError: - return False - payload = data.get("payload") - if not isinstance(payload, Mapping): - return False - stored_data, auth = payload.get("data"), payload.get("authCheck") - if not isinstance(stored_data, Mapping) or not isinstance(auth, Mapping) or auth.get("validKey") is not True: - return False - scope = auth.get("scope") - if not isinstance(scope, Mapping) or scope.get("accessLevel") != "project": - return False - return all( - isinstance(value, str) and bool(value.strip()) - for value in (stored_data.get("fileKey"), stored_data.get("publicKey"), scope.get("projectId")) - ) - - -def export_http(provider: str, url: str, headers: Mapping[str, str], payload: Any) -> None: - from google.protobuf.json_format import ParseDict, ParseError - from google.protobuf.message import DecodeError - from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceResponse - - response = request( - "POST", - url, - provider, - headers={"Content-Type": "application/x-protobuf", "Accept": "application/x-protobuf", **headers}, - content=payload.SerializeToString(), - ) - if provider == "aliyun" and is_aliyun_http_success_ack(response.status_code, response.content): - return - result = ExportTraceServiceResponse() - if response.content: - try: - if "json" in response.headers.get("content-type", "").lower(): - body = response.json() - if not isinstance(body, dict): - raise ValueError("Invalid ingestion acknowledgment object") - if ( - provider == "langfuse" - and response.status_code == 200 - and any( - key.lower() == "x-langfuse-ingestion-version" and value == "4" for key, value in headers.items() - ) - and _langfuse_enqueue_ack(body) - ): - return - ParseDict(body, result) - else: - result.ParseFromString(response.content) - except (DecodeError, ParseError, ValueError): - raise TraceExportUncertainError(f"{provider} returned an unreadable ingestion acknowledgment") from None - _check_response(result, provider) - - -def export_grpc(provider: str, url: str, headers: Mapping[str, str], payload: Any) -> None: - """Use the provider's gRPC protocol; HTTP exports use the SSRF HTTP owner.""" - import grpc # type: ignore[import-untyped] - from opentelemetry.proto.collector.trace.v1.trace_service_pb2_grpc import TraceServiceStub - - parsed = urlsplit(url) - if not parsed.hostname or parsed.username or parsed.password or parsed.scheme not in {"http", "https"}: - raise TraceExportRejectedError("Invalid OTLP gRPC endpoint") - port = parsed.port or (4317 if provider == "tencent" else 443) - host = f"[{parsed.hostname}]" if ":" in parsed.hostname else parsed.hostname - target = f"{host}:{port}" - channel = ( - grpc.insecure_channel(target) - if parsed.scheme == "http" - else grpc.secure_channel(target, grpc.ssl_channel_credentials()) - ) - try: - result = TraceServiceStub(channel).Export(payload, metadata=tuple(headers.items()), timeout=30) - except grpc.RpcError as error: - code = error.code() - if code in { - grpc.StatusCode.DEADLINE_EXCEEDED, - grpc.StatusCode.CANCELLED, - grpc.StatusCode.UNKNOWN, - grpc.StatusCode.UNAVAILABLE, - }: - raise TraceExportUncertainError(f"{provider} ingestion acknowledgment is unknown") from None - if code in {grpc.StatusCode.RESOURCE_EXHAUSTED, grpc.StatusCode.ABORTED}: - raise RetryableTraceDispatchError(f"{provider} ingestion temporarily unavailable") from None - raise TraceExportRejectedError(f"{provider} ingestion rejected") from None - finally: - channel.close() - _check_response(result, provider) diff --git a/api/services/tracing/provider_common.py b/api/services/tracing/provider_common.py deleted file mode 100644 index 88fdef0636a46d..00000000000000 --- a/api/services/tracing/provider_common.py +++ /dev/null @@ -1,127 +0,0 @@ -"""Shared serialization and HTTP boundaries for application trace adapters.""" - -import json -import math -from collections.abc import Mapping -from datetime import UTC, datetime -from typing import Any -from urllib.parse import urlsplit, urlunsplit - -from core.ops.exceptions import RetryableTraceDispatchError -from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind -from services.tracing.provider_dispatch import TraceExportRejectedError, TraceExportUncertainError -from services.tracing.request_log_protection import protect_http_request_logs - - -def json_value(value: object) -> str: - return json.dumps(value, ensure_ascii=False, default=str, allow_nan=False) - - -def utc(value: datetime) -> datetime: - return value.replace(tzinfo=UTC) if value.tzinfo is None else value.astimezone(UTC) - - -def nanos(value: datetime) -> int: - return int(utc(value).timestamp() * 1_000_000_000) - - -def finished_at(span: CanonicalSpan) -> datetime: - if span.end_time is None: - raise TraceExportRejectedError("Application trace requires completed spans") - return utc(span.end_time) - - -def endpoint(value: object) -> str: - """Validate a destination without allowing credentials in HTTP authority.""" - if not isinstance(value, str): - raise TraceExportRejectedError("Tracing endpoint is required") - parsed = urlsplit(value) - if parsed.scheme not in {"http", "https"} or not parsed.hostname or parsed.username or parsed.password: - raise TraceExportRejectedError("Invalid tracing endpoint") - if parsed.query or parsed.fragment: - raise TraceExportRejectedError("Tracing endpoint must not contain query parameters or a fragment") - return urlunsplit(parsed).rstrip("/") - - -def required(config: Mapping[str, Any], name: str) -> str: - value = config.get(name) - if not isinstance(value, str) or not value.strip(): - raise TraceExportRejectedError(f"Tracing configuration requires {name}") - return value - - -def token_usage(span: CanonicalSpan) -> dict[str, int]: - """Only model-call spans carry billable usage; wrappers retain metadata.""" - if span.kind is not CanonicalSpanKind.LLM: - return {} - result = {} - for source, target in (("prompt_tokens", "input"), ("completion_tokens", "output"), ("total_tokens", "total")): - value = span.metadata.get(source) - if isinstance(value, int) and not isinstance(value, bool) and value >= 0: - result[target] = value - return result - - -def usd_cost(span: CanonicalSpan) -> float | None: - if span.kind is not CanonicalSpanKind.LLM or str(span.metadata.get("currency", "USD")).upper() != "USD": - return None - try: - value = float(span.metadata["total_price"]) - except (KeyError, TypeError, ValueError): - return None - return value if math.isfinite(value) and value >= 0 else None - - -def check_http_status(status: int, provider: str) -> None: - if 200 <= status < 300: - return - if status in {408, 425, 429} or status >= 500: - raise RetryableTraceDispatchError(f"{provider} ingestion temporarily unavailable (HTTP {status})") - raise TraceExportRejectedError(f"{provider} ingestion rejected (HTTP {status})") - - -def request(method: str, url: str, provider: str, **kwargs: Any): - """Use the configured SSRF boundary, with retries owned by the durable task.""" - import httpx - from opentelemetry.context import _SUPPRESS_INSTRUMENTATION_KEY, attach, detach, set_value - - from core.helper import ssrf_proxy - - token = attach(set_value(_SUPPRESS_INSTRUMENTATION_KEY, True)) - try: - with protect_http_request_logs(provider == "aliyun"): - if provider == "aliyun": - response = _request_sensitive_aliyun(method, url, **kwargs) - else: - response = ssrf_proxy.make_request( - method, url, max_retries=0, timeout=30.0, follow_redirects=False, **kwargs - ) - except httpx.TransportError: - raise TraceExportUncertainError(f"{provider} ingestion outcome is unknown after transport failure") from None - finally: - detach(token) - check_http_status(response.status_code, provider) - return response - - -def _request_sensitive_aliyun(method: str, url: str, **kwargs: Any): - """Discard private HTTP exception objects before raising a public failure.""" - import httpx - - from core.helper import ssrf_proxy - - response: httpx.Response | None = None - failure: RuntimeError | None = None - try: - response = ssrf_proxy.make_request(method, url, max_retries=0, timeout=30.0, follow_redirects=False, **kwargs) - except httpx.RequestError: - failure = TraceExportUncertainError("aliyun ingestion outcome is unknown after transport failure") - except Exception: - failure = TraceExportRejectedError("aliyun HTTP request was rejected by the transport boundary") - # Raise after leaving the original exception handler so neither its cause - # nor implicit context retains a credential-bearing request or message. - if failure is not None: - raise failure from None - if response is None: - raise TraceExportUncertainError("aliyun HTTP transport returned no acknowledgment") - return response diff --git a/api/services/tracing/provider_configs.py b/api/services/tracing/provider_configs.py new file mode 100644 index 00000000000000..6623dd4a442130 --- /dev/null +++ b/api/services/tracing/provider_configs.py @@ -0,0 +1,197 @@ +"""Verify and present tracing provider configurations for the configuration service. + +Stored rows keep the shape they already have: the provider's fields at the +top level, secrets as workspace-encrypted tokens, and derived fields resolved +by verification. Secrets are write-only: they are decrypted to verify or to +deliver traces and are never presented back. +""" + +import logging +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any, override + +from dify_trace_sdk.config import ProviderConfig +from dify_trace_sdk.egress import Egress +from dify_trace_sdk.errors import VerificationError +from dify_trace_sdk.provider import TraceProvider +from pydantic import ValidationError + +from services.app_tracing_config_service import ( + AppTracingConfigInvalidConfigurationError, + AppTracingConfigInvalidProviderError, + AppTracingConfigProcessingError, + AppTracingConfigProviderUnavailableError, + AppTracingConfigVerificationFailedError, + PresentedTracingConfig, + TracingConfigProviderGateway, +) +from services.tracing.errors import ( + ProviderConfigUnreadableError, + TraceProviderUnavailableError, + UnknownTraceProviderError, +) +from services.tracing.providers import ( + SUPPORTED_PROVIDERS, + ProviderConfigCodec, + TraceProviderCatalog, + TraceProviderSupport, + configured_secrets, + support_of, +) + +logger = logging.getLogger(__name__) + + +def resolve_provider(providers: TraceProviderCatalog, provider_id: str) -> type[TraceProvider[Any]]: + try: + return providers.resolve(provider_id) + except UnknownTraceProviderError: + raise AppTracingConfigInvalidProviderError(provider_id) from None + except TraceProviderUnavailableError: + raise AppTracingConfigProviderUnavailableError from None + + +def _comparable(value: object) -> object: + return value.strip().rstrip("/") if isinstance(value, str) else value + + +class ProviderConfigGateway(TracingConfigProviderGateway): + def __init__(self, *, providers: TraceProviderCatalog, codec: ProviderConfigCodec, egress: Egress) -> None: + self._providers = providers + self._codec = codec + self._egress = egress + + @override + def validate_provider(self, tracing_provider: str) -> None: + # Reading, deleting, or disabling a configuration must work while its package is not installed. + if tracing_provider not in SUPPORTED_PROVIDERS: + raise AppTracingConfigInvalidProviderError(tracing_provider) + + @override + def validate_saved_config( + self, *, workspace_id: str, tracing_provider: str, tracing_config: dict[str, Any] + ) -> None: + provider = resolve_provider(self._providers, tracing_provider) + try: + self._codec.decode(workspace_id, provider, tracing_config) + except ProviderConfigUnreadableError: + raise AppTracingConfigProcessingError from None + + @override + def prepare_new_config( + self, *, workspace_id: str, tracing_provider: str, tracing_config: dict[str, Any] + ) -> dict[str, Any]: + provider = resolve_provider(self._providers, tracing_provider) + return self._verified(workspace_id, provider, self._submitted(provider, tracing_config, None)) + + @override + def prepare_updated_config( + self, + *, + workspace_id: str, + tracing_provider: str, + tracing_config: dict[str, Any], + current_tracing_config: dict[str, Any] | None, + ) -> dict[str, Any]: + provider = resolve_provider(self._providers, tracing_provider) + current = None + if current_tracing_config is not None: + try: + current = self._codec.decode(workspace_id, provider, current_tracing_config) + except ProviderConfigUnreadableError: + # An unreadable row can still be replaced by re-entering its secrets. + current = None + return self._verified(workspace_id, provider, self._submitted(provider, tracing_config, current)) + + @override + def present_config( + self, *, workspace_id: str, tracing_provider: str, tracing_config: dict[str, Any] | None + ) -> PresentedTracingConfig: + if tracing_config is None: + raise AppTracingConfigProcessingError + provider = resolve_provider(self._providers, tracing_provider) + model = provider.config_model + hidden = model.secret_fields() | model.derived_fields() + try: + # Secrets stay encrypted here: presentation never needs them. + config = model.model_validate(tracing_config) + except ValidationError: + # Historical rows stay readable even when they no longer validate. + values = { + key: value for key, value in tracing_config.items() if key not in hidden and isinstance(value, str) + } + return PresentedTracingConfig(values={**values, "project_url": None}, configured_secrets=()) + values = {key: value for key, value in config.model_dump(mode="json").items() if key not in hidden} + values["project_url"] = provider(self._egress).project_url(config) + return PresentedTracingConfig(values=values, configured_secrets=configured_secrets(provider, tracing_config)) + + def _submitted( + self, provider: type[TraceProvider[Any]], submitted: Mapping[str, Any], current: ProviderConfig | None + ) -> dict[str, Any]: + """Merge a form submission over the current configuration. + + A secret left empty, omitted, or echoed in its masked form keeps its + current value; ``None`` clears an optional secret. Derived fields are + never accepted from the client. Stored secrets are kept only while the + destination stays the same, so they never follow an edited host. + """ + model = provider.config_model + current_values = current.model_dump() if current is not None else {} + moved = any( + name in submitted and _comparable(submitted[name]) != _comparable(current_values.get(name)) + for name in model.destination_fields() + ) + values: dict[str, Any] = {} + for field in model.form_fields(): + current_value = current_values.get(field.name) + if field.secret and self._left_unchanged(submitted, field.name, current_value): + if not moved and current_value is not None: + values[field.name] = current_value + continue + if field.name in submitted: + values[field.name] = submitted[field.name] + return values + + def _left_unchanged(self, submitted: Mapping[str, Any], name: str, current_value: object) -> bool: + """Whether the form leaves a secret as it is: omitted, blank, or echoed in its masked form.""" + if name not in submitted: + return True + value = submitted[name] + return value == "" or ( + isinstance(current_value, str) and bool(current_value) and value == self._codec.mask(current_value) + ) + + def _verified( + self, workspace_id: str, provider: type[TraceProvider[Any]], values: dict[str, Any] + ) -> dict[str, Any]: + try: + config = provider.config_model.model_validate(values) + except ValidationError: + raise AppTracingConfigInvalidConfigurationError from None + try: + verified = provider(self._egress).verify(config) + except VerificationError as error: + logger.warning("Tracing configuration verification failed: provider=%s code=%s", provider.id, error.code) + raise AppTracingConfigVerificationFailedError from None + return self._codec.encode(workspace_id, verified) + + +@dataclass(frozen=True, slots=True) +class TraceProviderCapabilities: + support: TraceProviderSupport + api_collection_enabled: bool + """Whether this deployment collects Agent runtime spans; it does not probe the Agent runtime.""" + + +class TraceProviderCapabilitiesQuery: + def __init__(self, *, providers: TraceProviderCatalog, api_collection_enabled: bool) -> None: + self._providers = providers + self._api_collection_enabled = api_collection_enabled + + def get(self, provider_id: str) -> TraceProviderCapabilities: + """Raises AppTracingConfigInvalidProviderError or AppTracingConfigProviderUnavailableError.""" + return TraceProviderCapabilities( + support=support_of(resolve_provider(self._providers, provider_id)), + api_collection_enabled=self._api_collection_enabled, + ) diff --git a/api/services/tracing/provider_dispatch.py b/api/services/tracing/provider_dispatch.py deleted file mode 100644 index 06317c582e8e68..00000000000000 --- a/api/services/tracing/provider_dispatch.py +++ /dev/null @@ -1,122 +0,0 @@ -"""Dispatch a completed application trace without exposing destination credentials. - -An accepted export acknowledges ingestion, not query availability. Completed spans -must not be exported again after acceptance; some providers do not deduplicate IDs. -The durable task owner decides whether an uncertain transport outcome is retried. -""" - -from collections.abc import Mapping -from dataclasses import dataclass, field -from importlib import import_module -from typing import Any, Literal -from urllib.parse import urlsplit - -from core.ops.exceptions import RetryableTraceDispatchError, TraceProviderNotInstalledError -from core.ops.unified_trace.entities import CanonicalTrace - - -@dataclass(frozen=True) -class TraceExportResult: - trace_id: str - trace_url: str | None = None - span_urls: dict[str, str] = field(default_factory=dict) - status: Literal["accepted"] = "accepted" - - -class TraceExportRejectedError(RuntimeError): - """A terminal export rejection, containing no credential or response body.""" - - -class TraceExportUncertainError(RetryableTraceDispatchError): - """The transport failed without proving whether ingestion accepted the data.""" - - acceptance_unknown = True - - -def frozen_project_url(provider: str, config: Mapping[str, Any]) -> str | None: - """Read only persisted destination metadata, without credentials or vendor I/O.""" - if provider == "arize": - try: - from dify_trace_arize_phoenix.arize_configuration import project_url - except ModuleNotFoundError: - return None - - # Never substitute the current app destination for a historical run. - # Unverified legacy redirects remain stored, but are not presented as - # project links. A future config save does not backfill that snapshot. - return project_url(config) - # These adapters expose a platform console, not the run's destination project. - if provider in {"aliyun", "tencent"}: - return None - value = config.get("project_url") - if ( - not isinstance(value, str) - or not value - or any(ord(char) <= 32 or ord(char) == 127 for char in value) - or "\\" in value - ): - return None - try: - parsed = urlsplit(value) - if ( - parsed.scheme not in {"http", "https"} - or not parsed.hostname - or parsed.username is not None - or parsed.password is not None - ): - return None - _ = parsed.port - except ValueError: - return None - return value - - -_MODULES = { - "langfuse": "dify_trace_langfuse.agent_trace", - "langsmith": "dify_trace_langsmith.agent_trace", - "phoenix": "dify_trace_arize_phoenix.agent_trace", - "arize": "dify_trace_arize_phoenix.agent_trace", - "opik": "dify_trace_opik.agent_trace", - "weave": "dify_trace_weave.agent_trace", - "mlflow": "dify_trace_mlflow.agent_trace", - "databricks": "dify_trace_mlflow.agent_trace", - "aliyun": "dify_trace_aliyun.agent_trace", - "tencent": "dify_trace_tencent.agent_trace", -} - - -def provider_capabilities(provider: str) -> dict[str, bool | str]: - """Describe implemented mappings, independently of installation/account tests.""" - if provider not in _MODULES: - raise ValueError("Unsupported tracing provider") - return { - "provider": provider, - "agent_spans": True, - "session": True, - "native_usage": True, - "native_cost": provider in {"langfuse", "langsmith", "opik", "mlflow", "databricks"}, - "trace_url": provider in {"langfuse", "langsmith", "weave", "mlflow", "databricks"}, - "span_url": provider in {"langfuse", "langsmith", "weave"}, - "export_ack": "accepted", - } - - -def export_trace(provider: str, config: Mapping[str, Any], trace: CanonicalTrace) -> TraceExportResult: - """Export one complete tree to the caller's already-authorized destination.""" - module_name = _MODULES.get(provider) - if module_name is None: - raise TraceExportRejectedError("Unsupported tracing provider") - if trace.external_parent is not None or trace.required_parent_context_id is not None: - raise TraceExportRejectedError("Application trace export requires a complete resolved tree") - if any(span.end_time is None for span in trace.spans): - raise TraceExportRejectedError("Application trace export requires completed spans") - try: - module = import_module(module_name) - return module.export_trace(provider, config, trace) - except (TraceExportRejectedError, RetryableTraceDispatchError): - raise - except ImportError as error: - raise TraceProviderNotInstalledError(provider, error.name or module_name) from None - except Exception: - # Vendor exceptions can contain full requests, headers, and endpoint secrets. - raise TraceExportRejectedError(f"{provider} trace export failed") from None diff --git a/api/services/tracing/providers.py b/api/services/tracing/providers.py new file mode 100644 index 00000000000000..d7d58385afa58a --- /dev/null +++ b/api/services/tracing/providers.py @@ -0,0 +1,129 @@ +"""Trace providers as the pipeline reaches them, and their stored configurations. + +Dify supports a fixed set of provider IDs; a deployment may install any subset +of their packages. An unknown ID is a client error, a known but uninstalled one +is a deployment fact. Stored configurations keep the provider's fields at the +top level and its secrets as workspace-encrypted tokens. +""" + +import logging +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any, Literal, Protocol + +from dify_trace_sdk.config import ProviderConfig +from dify_trace_sdk.provider import TraceProvider +from pydantic import ValidationError + +from services.tracing.errors import ProviderConfigUnreadableError + +logger = logging.getLogger(__name__) + +SUPPORTED_PROVIDERS = ( + "langfuse", + "langsmith", + "opik", + "weave", + "arize", + "phoenix", + "aliyun", + "mlflow", + "databricks", + "tencent", +) + + +class TraceProviderCatalog(Protocol): + def resolve(self, provider_id: str) -> type[TraceProvider[Any]]: + """The installed provider class of a provider ID. + + Raises: + UnknownTraceProviderError: Dify does not support the ID. + TraceProviderUnavailableError: This deployment does not install the provider. + """ + ... + + +class TracingSecretCipher(Protocol): + """Workspace-scoped encryption of provider secrets.""" + + def encrypt(self, tenant_id: str, value: str) -> str: ... + + def decrypt(self, tenant_id: str, tokens: Sequence[str]) -> list[str]: + """Decrypt tokens in order; raises when the workspace key or a token is unreadable.""" + ... + + def mask(self, value: str) -> str: + """The masked form clients were shown before secrets became write-only.""" + ... + + +class ProviderConfigCodec: + """Convert provider configurations between their validated and stored forms.""" + + def __init__(self, cipher: TracingSecretCipher) -> None: + self._cipher = cipher + + def encode(self, tenant_id: str, config: ProviderConfig) -> dict[str, Any]: + stored = config.model_dump(mode="json") + for name in config.secret_fields(): + value = stored.get(name) + # An empty secret stays empty so it is not reported as configured. + if isinstance(value, str) and value: + stored[name] = self._cipher.encrypt(tenant_id, value) + return stored + + def decode(self, tenant_id: str, provider: type[TraceProvider[Any]], stored: Mapping[str, Any]) -> ProviderConfig: + """Decrypt and validate a stored configuration for verification or delivery. + + Raises: + ProviderConfigUnreadableError: A secret cannot be decrypted or the configuration no longer validates. + """ + model = provider.config_model + values = dict(stored) + encrypted = [ + name for name in sorted(model.secret_fields()) if isinstance(values.get(name), str) and values[name] + ] + if encrypted: + try: + decrypted = self._cipher.decrypt(tenant_id, [values[name] for name in encrypted]) + except Exception as error: + # A missing workspace key or a corrupt token; the type alone is safe to log. + logger.warning("Tracing configuration secrets unreadable: %s", type(error).__name__) + raise ProviderConfigUnreadableError from None + values.update(zip(encrypted, decrypted, strict=True)) + try: + return model.model_validate(values) + except ValidationError: + raise ProviderConfigUnreadableError from None + + def mask(self, value: str) -> str: + return self._cipher.mask(value) + + +def configured_secrets(provider: type[TraceProvider[Any]], stored: Mapping[str, Any]) -> tuple[str, ...]: + return tuple(name for name in sorted(provider.config_model.secret_fields()) if stored.get(name)) + + +@dataclass(frozen=True, slots=True) +class TraceProviderSupport: + """What a provider's mapping supports. Describes the adapter, not delivery health.""" + + provider: str + native_cost: bool + trace_url: bool + span_url: bool + agent_spans: Literal[True] = True + session: Literal[True] = True + native_usage: Literal[True] = True + export_ack: Literal["accepted"] = "accepted" + """An exporter acknowledgement does not confirm that the destination can query the trace.""" + + +def support_of(provider: type[TraceProvider[Any]]) -> TraceProviderSupport: + return TraceProviderSupport( + provider=provider.id, + native_cost=provider.capabilities.native_cost, + trace_url=provider.capabilities.trace_url, + span_url=provider.capabilities.span_url, + ) diff --git a/api/services/tracing/records.py b/api/services/tracing/records.py new file mode 100644 index 00000000000000..4737be9a441e0d --- /dev/null +++ b/api/services/tracing/records.py @@ -0,0 +1,242 @@ +"""Assembled trace records: one Dify operation with the data its sinks need. + +Records are plain values copied from committed rows before their session +closes. They live only inside the delivery worker; the canonical trace built +from them is what gets persisted for retries. +""" + +from collections.abc import Mapping, Sequence +from datetime import datetime +from decimal import Decimal +from enum import StrEnum +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field, StrictStr, ValidationError + +from graphon.entities import WorkflowNodeExecution + + +class TraceParentContext(BaseModel): + """The outer workflow run and node execution that started a nested workflow run.""" + + parent_workflow_run_id: StrictStr + parent_node_execution_id: StrictStr | None = None + """``None`` when the outer node execution is unknown.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + +class TraceMessage(BaseModel): + """The facts of one message that trace sinks read.""" + + id: str + app_id: str + conversation_id: str + workflow_run_id: str | None + """The chatflow run that produced the message; ``None`` for other app modes.""" + query: str + answer: str + error: str | None + """``None`` when the message recorded no error.""" + created_at: datetime + model_provider: str | None + """``None`` when no model answered, such as an annotation reply.""" + model_id: str | None + """``None`` when no model answered, such as an annotation reply.""" + total_price: Decimal | None + """``None`` when the price was not recorded.""" + currency: str + + model_config = ConfigDict(frozen=True) + + +class BaseTraceInfo(BaseModel): + message_id: str | None = None + message_data: TraceMessage | None = None + """The message the operation belongs to; ``None`` for operations outside a conversation.""" + inputs: str | dict[str, Any] | list[Any] | None = None + outputs: str | dict[str, Any] | list[Any] | None = None + start_time: datetime | None = None + end_time: datetime | None = None + metadata: dict[str, Any] + trace_id: str | None = None + """The caller-supplied trace ID, if any.""" + + model_config = ConfigDict(protected_namespaces=(), frozen=True) + + @property + def resolved_trace_id(self) -> str | None: + """The caller's trace ID, then the workflow run, then the message.""" + if self.trace_id: + return self.trace_id + if ( + isinstance(self, WorkflowTraceInfo | WorkflowNodeTraceInfo | SuggestedQuestionTraceInfo) + and self.workflow_run_id + ): + return self.workflow_run_id + return self.message_id + + @property + def resolved_parent_context(self) -> tuple[str | None, str | None]: + """(outer workflow_run_id, outer node_execution_id) of a nested workflow.""" + raw = self.metadata.get("parent_trace_context") + if not isinstance(raw, Mapping): + return None, None + try: + context = TraceParentContext.model_validate(raw) + except ValidationError: + return None, None + return context.parent_workflow_run_id, context.parent_node_execution_id + + +class WorkflowTraceInfo(BaseTraceInfo): + start_time: datetime + """When the run was created.""" + end_time: datetime + """When the run finished; only finished runs are traced.""" + node_executions: Sequence[WorkflowNodeExecution] + workflow_tool_parents: Mapping[str, str] = Field(default_factory=dict) + """Child execution row ID to its admitted caller row ID; empty outside workflow tools.""" + conversation_id: str | None + message_started_at: datetime | None + """Creation time of the chatflow message this run answers.""" + workflow_app_log_id: str | None + workflow_id: str + tenant_id: str + workflow_run_id: str + workflow_run_elapsed_time: int | float + workflow_run_status: str + workflow_run_inputs: Mapping[str, Any] + workflow_run_outputs: Mapping[str, Any] + workflow_run_version: str + error: str | None + total_tokens: int + prompt_tokens: int | None + completion_tokens: int | None + file_list: list[str] + invoked_by: str | None + query: str + + +class MessageTraceInfo(BaseTraceInfo): + message_data: TraceMessage + conversation_model: str + message_tokens: int + answer_tokens: int + total_tokens: int + error: str | None + file_list: list[str] + conversation_mode: str + gen_ai_server_time_to_first_token: float | None + llm_streaming_time_to_generate: float | None + is_streaming_request: bool + + +class ModerationTraceInfo(BaseTraceInfo): + flagged: bool + action: str + preset_response: str + query: str + + +class SuggestedQuestionTraceInfo(BaseTraceInfo): + total_tokens: int + status: str | None + error: str | None + from_account_id: str | None + agent_based: bool | None + from_source: str | None + model_provider: str | None + model_id: str | None + suggested_question: list[str] + level: str + status_message: str | None + workflow_run_id: str | None + + +class DatasetRetrievalTraceInfo(BaseTraceInfo): + documents: list[dict[str, Any]] + error: str | None + + +class ToolTraceInfo(BaseTraceInfo): + tool_name: str + tool_inputs: dict[str, Any] + tool_outputs: str + error: str | None + tool_config: dict[str, Any] + time_cost: int | float + tool_parameters: dict[str, Any] + file_url: str | None + + +class GenerateNameTraceInfo(BaseTraceInfo): + conversation_id: str | None + tenant_id: str + + +class PromptGenerationTraceInfo(BaseTraceInfo): + """Prompt generation operations (rule-generate, code-generate, etc.).""" + + tenant_id: str + user_id: str + app_id: str | None + operation_type: str + instruction: str + prompt_tokens: int + completion_tokens: int + total_tokens: int + model_provider: str + model_name: str + latency: float + total_price: float | None + currency: str | None + error: str | None + + +class WorkflowNodeTraceInfo(BaseTraceInfo): + workflow_id: str + workflow_run_id: str + tenant_id: str + node_execution_id: str + node_id: str + node_type: str + title: str + status: str + error: str | None = None + elapsed_time: float + index: int + predecessor_node_id: str | None = None + total_tokens: int = 0 + total_price: float = 0.0 + currency: str | None = None + model_provider: str | None = None + model_name: str | None = None + prompt_tokens: int | None = None + completion_tokens: int | None = None + tool_name: str | None = None + iteration_id: str | None = None + iteration_index: int | None = None + loop_id: str | None = None + loop_index: int | None = None + parallel_id: str | None = None + node_inputs: Mapping[str, Any] | None = None + node_outputs: Mapping[str, Any] | None = None + process_data: Mapping[str, Any] | None = None + invoked_by: str | None = None + + +class DraftNodeExecutionTrace(WorkflowNodeTraceInfo): + pass + + +class OperationType(StrEnum): + """Operation labels on enterprise token metrics, so totals are not double-counted.""" + + WORKFLOW = "workflow" + NODE_EXECUTION = "node_execution" + MESSAGE = "message" + RULE_GENERATE = "rule_generate" + CODE_GENERATE = "code_generate" + STRUCTURED_OUTPUT = "structured_output" + INSTRUCTION_MODIFY = "instruction_modify" diff --git a/api/services/tracing/recovery.py b/api/services/tracing/recovery.py index 4e9c8729edbbb3..562207b463a3a1 100644 --- a/api/services/tracing/recovery.py +++ b/api/services/tracing/recovery.py @@ -5,190 +5,113 @@ Message is complete. Recovery never invents answers, usage, or success states. """ -import json import logging -import math -from datetime import datetime, timedelta -from typing import Protocol, cast +from datetime import UTC, datetime +from uuid import uuid4 -from sqlalchemy import select -from sqlalchemy.exc import SQLAlchemyError +from dify_trace_sdk.events import MessageCompleted, TraceCorrelation, TraceEvent, TracePayload, WorkflowRunCompleted -from configs import dify_config -from core.ops.entities.trace_entity import BaseTraceInfo, MessageTraceInfo, TraceTaskName, WorkflowTraceInfo -from core.ops.ops_trace_manager import TraceTask -from graphon.enums import WorkflowExecutionStatus +from domain.tracing.errors import TraceRunError +from domain.tracing.terminal_recovery import message_terminal_due, workflow_terminal_due +from domain.tracing.trace_run_agg import TraceRunAgg from libs.datetime_utils import naive_utc_now -from models.application_trace import ApplicationTraceInvocation, ApplicationTraceRun -from models.enums import MessageStatus -from models.model import App, AppMode, Conversation, Message, MessageDict -from models.workflow import WorkflowRun -from repositories.application_trace_repository import ApplicationTraceRepository, TraceCollectionError +from services.tracing.collection_ports import TraceRunRepository, TraceSourceQuery +from services.tracing.delivery import TraceCollector, TraceRecordSource +from services.tracing.errors import ( + SourceNotFoundError, + SourceNotReadyError, + TraceCollectionError, + TraceCollectionNotFoundError, + TracingStoreUnavailableError, +) +from services.tracing.records import MessageTraceInfo, WorkflowTraceInfo logger = logging.getLogger(__name__) -_SETTLE_SECONDS = 5 -_WORKFLOW_TERMINAL = { - WorkflowExecutionStatus.SUCCEEDED, - WorkflowExecutionStatus.PARTIAL_SUCCEEDED, - WorkflowExecutionStatus.FAILED, - WorkflowExecutionStatus.STOPPED, -} -class RecoveryTarget(Protocol): - repository: ApplicationTraceRepository - - def capture_terminal(self, info: BaseTraceInfo) -> bool: ... - - -def agent_stop_requested(message: Message) -> bool: - """Read the consumer's stop intent without inferring cancellation success.""" - try: - metadata = json.loads(message.message_metadata) if message.message_metadata else {} - except (ValueError, TypeError): - return False - return isinstance(metadata, dict) and metadata.get("agent_stop_requested") is True - - -def _complete_message_metadata(message: Message) -> bool: - """Recognize the EasyUI _save_message transaction, not an early usage flush.""" - try: - metadata = json.loads(message.message_metadata) if message.message_metadata else {} - except (ValueError, TypeError): - return False - if ( - not isinstance(metadata, dict) - or not {"annotation_reply", "retriever_resources", "reasoning"} <= metadata.keys() - ): - return False - usage = metadata.get("usage") - if not isinstance(usage, dict): - return False - for name in ("prompt_tokens", "completion_tokens"): - value = usage.get(name) - if not isinstance(value, (int, float)) or isinstance(value, bool) or not math.isfinite(value) or value < 0: - return False - latency = message.provider_response_latency - return isinstance(latency, (int, float)) and math.isfinite(latency) and latency >= 0 - - -def _message_task(repository: ApplicationTraceRepository, run: ApplicationTraceRun, now: datetime) -> TraceTask | None: - with repository.sessions() as session: - source = session.execute( - select(Message, App.mode) - .join(App, App.id == Message.app_id) - .join(Conversation, Conversation.id == Message.conversation_id) - .where( - Message.id == run.source_id, - Message.app_id == run.app_id, - App.tenant_id == run.tenant_id, - Conversation.app_id == run.app_id, - ) - ).one_or_none() - if source is None: - raise TraceCollectionError("tracing_source_not_found", 404) - message, app_mode = source - pending_agent_stop = app_mode == AppMode.AGENT and agent_stop_requested(message) - if message.status is MessageStatus.PAUSED and not pending_agent_stop: - return None - invocations = list( - session.scalars(select(ApplicationTraceInvocation).where(ApplicationTraceInvocation.run_id == run.id)) - ) - explicit_error = message.status is MessageStatus.ERROR and bool(message.error) - if any(not invocation.finalized or invocation.finalized_at is None for invocation in invocations): - if (explicit_error or pending_agent_stop) and message.updated_at + timedelta( - seconds=dify_config.OPS_AGENT_TRACE_COLLECTION_TIMEOUT - ) < now: - raise TraceCollectionError("agent_trace_collection_incomplete") - return None - finalized = [invocation.finalized_at for invocation in invocations if invocation.finalized_at is not None] - anchor = max([message.updated_at, *finalized]) - if pending_agent_stop: - # Even a fully saved partial answer/usage cannot establish which - # runtime terminal won the stop race. Wait for the existing worker - # handoff, or converge to incomplete delivery after its deadline. - if anchor + timedelta(seconds=dify_config.OPS_AGENT_TRACE_COLLECTION_TIMEOUT) < now: - raise TraceCollectionError("business_terminal_not_persisted") - return None - if run.capture_agent and not invocations: - return None - if anchor + timedelta(seconds=_SETTLE_SECONDS) > now: - return None - complete = explicit_error or ( - message.status is MessageStatus.NORMAL - and message.updated_at > message.created_at - and _complete_message_metadata(message) - ) - if not complete: - if finalized and max(finalized) + timedelta(seconds=dify_config.OPS_AGENT_TRACE_COLLECTION_TIMEOUT) < now: - raise TraceCollectionError("business_terminal_not_persisted") - return None - return TraceTask( - TraceTaskName.MESSAGE_TRACE, - message_id=message.id, - conversation_id=message.conversation_id, - trace_session_id=run.session_id, - ) +def _event(run: TraceRunAgg, payload: TracePayload, user_id: str | None) -> TraceEvent: + return TraceEvent( + event_id=uuid4().hex, + tenant_id=run.tenant_id, + app_id=run.app_id, + user_id=user_id, + occurred_at=datetime.now(UTC), + correlation=TraceCorrelation( + external_trace_id=None, + session_id=run.session_id, + parent_workflow_run_id=None, + parent_node_execution_id=None, + ), + payload=payload, + ) + + +def _message_event( + sources: TraceSourceQuery, run: TraceRunAgg, now: datetime, collection_timeout: int +) -> TraceEvent | None: + progress = sources.message_progress(tenant_id=run.tenant_id, app_id=run.app_id, message_id=run.source_id) + if progress is None: + raise TraceCollectionNotFoundError("tracing_source_not_found") + if not message_terminal_due(run, progress, now=now, collection_timeout=collection_timeout): + return None + return _event(run, MessageCompleted(message_id=progress.message_id, conversation_id=progress.conversation_id), None) -def _workflow_task(repository: ApplicationTraceRepository, run: ApplicationTraceRun, now: datetime) -> TraceTask | None: - workflow = repository._source(tenant_id=run.tenant_id, app_id=run.app_id, kind="workflow", source_id=run.source_id) +def _workflow_event(sources: TraceSourceQuery, run: TraceRunAgg, now: datetime) -> TraceEvent | None: + workflow = sources.workflow_run(tenant_id=run.tenant_id, app_id=run.app_id, workflow_run_id=run.source_id) # An asynchronous workflow repository may not have saved the reserved run yet. - if ( - not isinstance(workflow, WorkflowRun) - or workflow.status not in _WORKFLOW_TERMINAL - or workflow.finished_at is None - ): + if workflow is None or not workflow_terminal_due(workflow, now=now): return None - if workflow.finished_at + timedelta(seconds=_SETTLE_SECONDS) > now: - return None - with repository.sessions() as session: - conversation_id = session.scalar( - select(Message.conversation_id) - .join(Conversation, Conversation.id == Message.conversation_id) - .where( - Message.app_id == run.app_id, Message.workflow_run_id == workflow.id, Conversation.app_id == run.app_id - ) - .limit(1) - ) - return TraceTask( - TraceTaskName.WORKFLOW_TRACE, - workflow_run_id=workflow.id, - conversation_id=conversation_id, - trace_session_id=run.session_id, - user_id=workflow.created_by, - ) - - -def recover_terminal(target: RecoveryTarget, run: ApplicationTraceRun) -> bool: - """Attempt one bounded recovery and rotate unresolved work behind other runs.""" - if ( - run.outer_trace is not None - or run.state not in {"collecting", "ready"} - or run.source_kind not in {"message", "workflow"} - ): + conversation_id = sources.workflow_run_conversation_id( + tenant_id=run.tenant_id, app_id=run.app_id, workflow_run_id=workflow.id + ) + return _event( + run, + WorkflowRunCompleted(workflow_run_id=workflow.id, conversation_id=conversation_id, total_tokens=None), + workflow.created_by, + ) + + +def recover_terminal( + run: TraceRunAgg, + *, + runs: TraceRunRepository, + sources: TraceSourceQuery, + records: TraceRecordSource, + collector: TraceCollector, + collection_timeout: int, + now: datetime, +) -> bool: + """Attempt one bounded recovery and rotate unresolved work behind other runs. + + ``now`` is naive UTC, like the ledger's timestamps. + """ + if not run.awaits_terminal or run.source_kind not in {"message", "workflow"}: return False error_code: str | None = None try: - task = (_message_task if run.source_kind == "message" else _workflow_task)( - target.repository, run, naive_utc_now() + event = ( + _message_event(sources, run, now, collection_timeout) + if run.source_kind == "message" + else _workflow_event(sources, run, now) ) - if task is None: + if event is None: + return False + try: + info = records.assemble(event) + except SourceNotReadyError: return False - info = task.execute() - if isinstance(info, MessageTraceInfo) and isinstance(info.message_data, dict): - info = info.model_copy(update={"message_data": Message.from_dict(cast(MessageDict, info.message_data))}) - if isinstance(info, WorkflowTraceInfo) and isinstance(info.workflow_data, dict): - info = info.model_copy(update={"workflow_data": WorkflowRun.from_dict(info.workflow_data)}) + except SourceNotFoundError: + raise TraceCollectionError("business_terminal_recovery_incomplete") from None if not isinstance(info, (MessageTraceInfo, WorkflowTraceInfo)): raise TraceCollectionError("business_terminal_recovery_incomplete") - if not target.capture_terminal(info): + if not collector.capture_terminal(tenant_id=run.tenant_id, info=info): raise TraceCollectionError("business_terminal_recovery_unclaimed") return True - except TraceCollectionError as error: + except (TraceCollectionError, TraceRunError) as error: error_code = error.code return False - except SQLAlchemyError: + except TracingStoreUnavailableError: # Leave the durable source intact. The next reconciler pass can retry. logger.warning("Application trace recovery database unavailable: run_id=%s", run.id) return False @@ -197,4 +120,6 @@ def recover_terminal(target: RecoveryTarget, run: ApplicationTraceRun) -> bool: logger.warning("Application trace recovery failed: run_id=%s error_type=%s", run.id, type(error).__name__) return False finally: - target.repository.recovery_checked(run.id, error_code=error_code) + with runs.change(run.id, tenant_id=run.tenant_id) as locked: + if locked is not None: + locked.run.recovery_checked(error_code, now=naive_utc_now()) diff --git a/api/services/tracing/registry.py b/api/services/tracing/registry.py new file mode 100644 index 00000000000000..02f847ae13941e --- /dev/null +++ b/api/services/tracing/registry.py @@ -0,0 +1,45 @@ +"""Discover installed trace providers through ``dify.trace_providers`` entry points. + +Neither an unknown nor an uninstalled provider ever raises into a business +request: callers translate both errors. +""" + +import functools +from importlib.metadata import EntryPoint, entry_points +from typing import Any, override + +from dify_trace_sdk.provider import TraceProvider + +from services.tracing.errors import TraceProviderUnavailableError, UnknownTraceProviderError +from services.tracing.providers import SUPPORTED_PROVIDERS, TraceProviderCatalog + +ENTRY_POINT_GROUP = "dify.trace_providers" + + +@functools.cache +def _entry_points() -> dict[str, EntryPoint]: + return {entry_point.name: entry_point for entry_point in entry_points(group=ENTRY_POINT_GROUP)} + + +@functools.cache +def _load(provider_id: str) -> type[TraceProvider[Any]]: + entry_point = _entry_points().get(provider_id) + if entry_point is None: + raise TraceProviderUnavailableError(provider_id) + try: + loaded = entry_point.load() + except ImportError: + raise TraceProviderUnavailableError(provider_id) from None + if not isinstance(loaded, type) or not issubclass(loaded, TraceProvider) or loaded.id != provider_id: + raise TraceProviderUnavailableError(provider_id) + return loaded + + +class EntryPointTraceProviderCatalog(TraceProviderCatalog): + """Installed packages are fixed for the process, so each provider is loaded once.""" + + @override + def resolve(self, provider_id: str) -> type[TraceProvider[Any]]: + if provider_id not in SUPPORTED_PROVIDERS: + raise UnknownTraceProviderError(provider_id) + return _load(provider_id) diff --git a/api/services/tracing/runtime_bridge.py b/api/services/tracing/runtime_bridge.py index f94aae3b192b8d..f8adfe2e4ce496 100644 --- a/api/services/tracing/runtime_bridge.py +++ b/api/services/tracing/runtime_bridge.py @@ -1,205 +1,133 @@ -"""Fail-open business hooks and durable terminal handoff for application tracing. +"""Fail-open execution hooks that start Agent trace collection. -The deployment gate is checked before constructing any database-backed service. -Execution context carries opaque correlation and policy only, never provider keys. +The composition root binds these hooks only when the deployment collects Agent +traces. Execution context carries opaque correlation and policy only, never +provider keys. """ from __future__ import annotations -import json import logging -from collections.abc import Mapping -from typing import TYPE_CHECKING, Literal +from typing import TYPE_CHECKING, Literal, override from dify_agent.layers.execution_context import ApplicationTracingContext -from configs import dify_config +from extensions.tracing import ApplicationTracingHooks +from services.tracing.collection import workflow_has_agent_v2 +from services.tracing.errors import TraceCollectionError +from services.tracing.service import ApplicationTracingService if TYPE_CHECKING: from core.app.apps.agent_app.runtime_request_builder import AgentAppRuntimeBuildContext from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, WorkflowAppGenerateEntity - from core.ops.entities.trace_entity import BaseTraceInfo from core.workflow.nodes.agent_v2.runtime_request_builder import WorkflowAgentRuntimeBuildContext - from services.tracing.service import ApplicationTracingService logger = logging.getLogger(__name__) -def workflow_has_agent_v2(graph: object) -> bool: - """Scope the new pipeline to graphs containing the Dify Agent v2 node.""" - from core.workflow.nodes.agent_v2.discriminator import is_dify_agent_node_data - - if isinstance(graph, str): - try: - graph = json.loads(graph) - except ValueError: - return False - if not isinstance(graph, Mapping) or not isinstance(graph.get("nodes"), list): - return False - return any( - isinstance(node, Mapping) and isinstance(node.get("data"), Mapping) and is_dify_agent_node_data(node["data"]) - for node in graph["nodes"] - ) - - -def _service() -> ApplicationTracingService: - from services.tracing.service import application_tracing_service - - return application_tracing_service() - - def _log_failure(stage: str, error: Exception) -> None: - from repositories.application_trace_repository import TraceCollectionError - code = error.code if isinstance(error, TraceCollectionError) else type(error).__name__ logger.warning("Application tracing hook skipped: stage=%s error_code=%s", stage, code) -def prepare_agent_app(context: AgentAppRuntimeBuildContext) -> ApplicationTracingContext | None: - if not dify_config.OPS_AGENT_TRACE_ENABLED: - return None - try: - return _service().prepare_invocation( - tenant_id=context.dify_context.tenant_id, - app_id=context.dify_context.app_id, - agent_id=context.agent_id, - source_kind="message", - source_id=context.idempotency_key, - session_id=context.dify_context.trace_session_id or context.conversation_id, - node_execution_id=None, - attempt=0, - call_key=context.idempotency_key, - ) - except Exception as error: - _log_failure("agent_app", error) - return None - - -def capture_agent_app_terminal( - *, - tenant_id: str, - app_id: str, - message_id: str, - outcome: Literal["succeeded", "failed", "cancelled"], - output: str | None, -) -> None: - if not dify_config.OPS_AGENT_TRACE_ENABLED: - return - try: - _service().capture_agent_app_terminal( - tenant_id=tenant_id, app_id=app_id, message_id=message_id, outcome=outcome, output=output - ) - except Exception as error: - _log_failure("agent_app_terminal", error) - +class RuntimeApplicationTracingHooks(ApplicationTracingHooks): + def __init__(self, traces: ApplicationTracingService) -> None: + self._traces = traces -def prepare_workflow_agent( - context: WorkflowAgentRuntimeBuildContext, conversation_id: str | None -) -> ApplicationTracingContext | None: - if not dify_config.OPS_AGENT_TRACE_ENABLED: - return None - try: - # Hidden Tool nodes own a different app but share the caller's run and are - # absent from its outer trace. Neither a source-app snapshot nor an orphan - # parent is safe. Keep this new nested scope uncollected until it has an - # explicit owner/parent contract; ordinary nodes and previews are unchanged. - if context.dify_context.workflow_tool_invocation_id is not None: + @override + def prepare_agent_app(self, context: AgentAppRuntimeBuildContext) -> ApplicationTracingContext | None: + try: + return self._traces.prepare_invocation( + tenant_id=context.dify_context.tenant_id, + app_id=context.dify_context.app_id, + agent_id=context.agent_id, + source_kind="message", + source_id=context.idempotency_key, + session_id=context.dify_context.trace_session_id or context.conversation_id, + node_execution_id=None, + attempt=0, + call_key=context.idempotency_key, + workflow_id=None, + ) + except Exception as error: + _log_failure("agent_app", error) return None - source_id = context.workflow_run_id or context.node_execution_id - return _service().prepare_invocation( - tenant_id=context.dify_context.tenant_id, - app_id=context.dify_context.app_id, - agent_id=context.agent.id, - source_kind="workflow" if context.workflow_run_id else "node", - source_id=source_id, - session_id=context.dify_context.trace_session_id or conversation_id or source_id, - node_execution_id=context.node_execution_id, - attempt=context.attempt, - call_key=f"{context.node_execution_id}:{context.attempt}", - workflow_id=context.workflow_id, - ) - except Exception as error: - _log_failure("workflow_agent", error) - return None - - -def freeze_workflow_run( - entity: AdvancedChatAppGenerateEntity | WorkflowAppGenerateEntity, - *, - workflow_id: str, - workflow_run_id: str, - conversation_id: object, - graph_data: object, -) -> None: - """Freeze policy immediately after the business workflow row is saved.""" - if not dify_config.OPS_AGENT_TRACE_ENABLED: - return - if not workflow_has_agent_v2(graph_data): - return - try: - explicit_session = entity.extras.get("trace_session_id") - session_id = next( - ( - value - for value in (explicit_session, conversation_id, workflow_run_id) - if isinstance(value, str) and value - ), - workflow_run_id, - ) - _service().start( - tenant_id=entity.app_config.tenant_id, - app_id=entity.app_config.app_id, - source_kind="workflow", - source_id=workflow_run_id, - session_id=session_id, - workflow_id=workflow_id, - ) - except Exception as error: - _log_failure("workflow_start", error) + @override + def capture_agent_app_terminal( + self, + *, + tenant_id: str, + app_id: str, + message_id: str, + outcome: Literal["succeeded", "failed", "cancelled"], + output: str | None, + ) -> None: + try: + self._traces.capture_agent_app_terminal( + tenant_id=tenant_id, app_id=app_id, message_id=message_id, outcome=outcome, output=output + ) + except Exception as error: + _log_failure("agent_app_terminal", error) + + @override + def prepare_workflow_agent( + self, context: WorkflowAgentRuntimeBuildContext, conversation_id: str | None + ) -> ApplicationTracingContext | None: + try: + # Hidden Tool nodes own a different app but share the caller's run and are + # absent from its outer trace. Neither a source-app snapshot nor an orphan + # parent is safe. Keep this new nested scope uncollected until it has an + # explicit owner/parent contract; ordinary nodes and previews are unchanged. + if context.dify_context.workflow_tool_invocation_id is not None: + return None + source_id = context.workflow_run_id or context.node_execution_id + return self._traces.prepare_invocation( + tenant_id=context.dify_context.tenant_id, + app_id=context.dify_context.app_id, + agent_id=context.agent.id, + source_kind="workflow" if context.workflow_run_id else "node", + source_id=source_id, + session_id=context.dify_context.trace_session_id or conversation_id or source_id, + node_execution_id=context.node_execution_id, + attempt=context.attempt, + call_key=f"{context.node_execution_id}:{context.attempt}", + workflow_id=context.workflow_id, + ) + except Exception as error: + _log_failure("workflow_agent", error) + return None -def queue_enabled_at_start(app_id: object) -> bool: - """Read the existing application switch without constructing a vendor SDK.""" - if not dify_config.OPS_AGENT_TRACE_ENABLED or app_id is None: - return False - from extensions.ext_database import db - from models.model import App - - try: - app = db.session.get(App, str(app_id)) - policy = json.loads(app.tracing) if app is not None and app.tracing else {} - return isinstance(policy, dict) and policy.get("enabled") is True - except Exception as error: - _log_failure("queue_policy", error) - return False - - -def capture_terminal(info: BaseTraceInfo) -> bool: - """Transfer ownership before legacy dispatch; transient DB failures retry. - - Unlike request construction, a worker must not fall back to a second export - path when the durable collector is temporarily inaccessible. - """ - if not dify_config.OPS_AGENT_TRACE_ENABLED: - return False - from sqlalchemy.exc import SQLAlchemyError - - from core.ops.exceptions import RetryableTraceDispatchError - - try: - return _service().capture_terminal(info) - except SQLAlchemyError: - raise RetryableTraceDispatchError("Application tracing persistence temporarily unavailable") from None - - -class RuntimeApplicationTracingHooks: - """Stateless app-scoped adapter; database services are created only on invocation.""" - - __slots__ = () - - prepare_agent_app = staticmethod(prepare_agent_app) - prepare_workflow_agent = staticmethod(prepare_workflow_agent) - freeze_workflow_run = staticmethod(freeze_workflow_run) - queue_enabled_at_start = staticmethod(queue_enabled_at_start) - capture_terminal = staticmethod(capture_terminal) - capture_agent_app_terminal = staticmethod(capture_agent_app_terminal) + @override + def freeze_workflow_run( + self, + entity: AdvancedChatAppGenerateEntity | WorkflowAppGenerateEntity, + *, + workflow_id: str, + workflow_run_id: str, + conversation_id: object, + graph_data: object, + ) -> None: + """Freeze policy immediately after the business workflow row is saved.""" + if not workflow_has_agent_v2(graph_data): + return + try: + explicit_session = entity.extras.get("trace_session_id") + session_id = next( + ( + value + for value in (explicit_session, conversation_id, workflow_run_id) + if isinstance(value, str) and value + ), + workflow_run_id, + ) + self._traces.start( + tenant_id=entity.app_config.tenant_id, + app_id=entity.app_config.app_id, + source_kind="workflow", + source_id=workflow_run_id, + session_id=session_id, + workflow_id=workflow_id, + ) + except Exception as error: + _log_failure("workflow_start", error) diff --git a/api/services/tracing/scopes.py b/api/services/tracing/scopes.py new file mode 100644 index 00000000000000..86f3a2edbe9f1f --- /dev/null +++ b/api/services/tracing/scopes.py @@ -0,0 +1,77 @@ +"""Decide which executions are traced and give their producers a scope. + +Deciding costs nothing extra on the request path: generators already hold +the app's tracing policy, so no configuration is decrypted and no provider is +loaded while a user waits. +""" + +import logging +from collections.abc import Callable +from typing import Protocol + +from dify_trace_sdk.emitter import TraceEmitter, TraceScope +from dify_trace_sdk.events import TraceCorrelation + +from services.app_tracing_config_service import AppTraceSettings + +logger = logging.getLogger(__name__) + +UNCORRELATED = TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None +) +"""The correlation of an operation that carries no caller-supplied trace context.""" + + +class TracingPolicyQuery(Protocol): + def policy(self, *, tenant_id: str, app_id: str) -> AppTraceSettings | None: + """The app's stored tracing policy; ``None`` when the app is not in the workspace. + + A stored policy that cannot be read is a disabled policy. + """ + ... + + +class TraceScopes: + """Create trace scopes for executions whose events someone consumes.""" + + def __init__( + self, *, emitter: TraceEmitter, policies: TracingPolicyQuery, enterprise_enabled: Callable[[], bool] + ) -> None: + self._emitter = emitter + self._policies = policies + self._enterprise_enabled = enterprise_enabled + + def for_execution( + self, + *, + tenant_id: str, + app_id: str | None, + policy: AppTraceSettings, + user_id: str | None, + correlation: TraceCorrelation, + ) -> TraceScope | None: + """The scope of an execution whose app policy the caller holds; ``None`` when nobody consumes it.""" + if not ((policy.enabled and policy.tracing_provider is not None) or self._enterprise_enabled()): + return None + return TraceScope( + emitter=self._emitter, + tenant_id=tenant_id, + app_id=app_id, + user_id=user_id, + correlation=correlation, + ) + + def for_app( + self, *, tenant_id: str, app_id: str | None, user_id: str | None, correlation: TraceCorrelation + ) -> TraceScope | None: + """A scope for callers that do not hold the app row; tracing failures stay out of the request.""" + policy = AppTraceSettings() + if app_id is not None: + try: + policy = self._policies.policy(tenant_id=tenant_id, app_id=app_id) or policy + except Exception: + logger.warning("Tracing policy lookup failed: app_id=%s", app_id, exc_info=True) + return None + return self.for_execution( + tenant_id=tenant_id, app_id=app_id, policy=policy, user_id=user_id, correlation=correlation + ) diff --git a/api/services/tracing/sdk_transport.py b/api/services/tracing/sdk_transport.py deleted file mode 100644 index d78ee4bae9c515..00000000000000 --- a/api/services/tracing/sdk_transport.py +++ /dev/null @@ -1,42 +0,0 @@ -"""Per-client SDK transports using Dify's HTTP egress policy, never global state.""" - -from typing import Any, override - -import httpx -import requests - -from services.tracing.provider_common import request as send_http -from services.tracing.provider_dispatch import TraceExportRejectedError - - -class SSRFTransport(httpx.BaseTransport): - def __init__(self, provider: str) -> None: - self.provider = provider - - @override - def handle_request(self, request: httpx.Request) -> httpx.Response: - response = send_http( - request.method, str(request.url), self.provider, headers=dict(request.headers), content=request.read() - ) - return httpx.Response(response.status_code, headers=response.headers, content=response.content, request=request) - - -class SSRFSession(requests.Session): - def __init__(self, provider: str) -> None: - super().__init__() - self.provider = provider - - @override - def send(self, request: requests.PreparedRequest, **kwargs: Any) -> requests.Response: - if not request.method or not request.url or not isinstance(request.body, (str, bytes, type(None))): - raise TraceExportRejectedError("Tracing SDK prepared an unsupported HTTP request") - received = send_http( - request.method, request.url, self.provider, headers=dict(request.headers), content=request.body - ) - response = requests.Response() - response.status_code = received.status_code - response.headers.update(received.headers) - response._content = received.content - response.url = str(received.url) - response.request = request - return response diff --git a/api/services/tracing/service.py b/api/services/tracing/service.py index acce7315b1ddaa..ae3bd00ef3661e 100644 --- a/api/services/tracing/service.py +++ b/api/services/tracing/service.py @@ -1,40 +1,63 @@ -"""Application tracing orchestration; secrets and vendor I/O stay in the API worker.""" +"""Agent trace collection for Agent-managed executions. + +Agent apps and workflows with Agent nodes produce spans that the Dify Agent +runtime reports separately. Their traces freeze a destination and a capture +policy when they start, collect runtime batches in a ledger, and are sent +once complete through the same provider interface as every other trace. +""" import json import logging import threading -from collections.abc import Generator +from collections.abc import Callable, Generator from contextlib import contextmanager -from datetime import UTC, datetime, timedelta -from typing import Any, Literal, cast +from datetime import UTC, datetime +from typing import Any, Literal, override +from uuid import uuid4 from dify_agent.layers.execution_context import ApplicationTracingContext from dify_agent.protocol.tracing import AgentTraceBatch -from sqlalchemy import select -from sqlalchemy.orm import sessionmaker - -from configs import dify_config -from core.ops.entities.trace_entity import BaseTraceInfo, MessageTraceInfo, TraceTaskName, WorkflowTraceInfo -from core.ops.exceptions import RetryableTraceDispatchError -from core.ops.ops_trace_manager import OpsTraceManager, TraceTask -from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace -from core.ops.unified_trace.trace_builder import CanonicalTraceBuilder, RepositoryWorkflowExecutionLoader -from extensions.ext_database import db -from models import Account, App, AppMode, Conversation, Message, Workflow, WorkflowNodeExecutionModel, WorkflowRun -from models.application_trace import ApplicationTraceRun -from models.model import MessageDict -from repositories.application_trace_repository import ApplicationTraceRepository, TraceCollectionError -from services.tracing.normalization import normalize_agent_span, order_spans, remap_outer_trace, sanitize -from services.tracing.provider_dispatch import ( - TraceExportUncertainError, - export_trace, - frozen_project_url, - provider_capabilities, +from dify_trace_sdk.egress import Egress +from dify_trace_sdk.errors import ExportOutcomeUnknownError, ExportRejectedError, ExportRetryableError +from dify_trace_sdk.events import MessageCompleted, TraceCorrelation, TraceEvent +from dify_trace_sdk.provider import TraceProvider +from pydantic import ValidationError + +from domain.tracing.errors import TraceRunConflictError, TraceRunError, TraceRunLimitError, TraceRunNotFoundError +from domain.tracing.sources import SourceNodeExecution +from domain.tracing.trace_run_agg import TraceRunAgg, TraceRunState, TraceSourceKind +from libs.datetime_utils import naive_utc_now +from libs.uuid_utils import uuidv7 +from services.tracing.canonical import CanonicalTraceBuilder +from services.tracing.collection import ( + AgentTraceSettings, + SourceMessage, + TraceApp, + TraceAppMode, + TraceRun, + TraceRunStatus, + TraceSource, + workflow_has_agent_v2, +) +from services.tracing.collection_ports import ScheduleTraceExport, TraceRunQuery, TraceRunRepository, TraceSourceQuery +from services.tracing.delivery import TraceCollector, TraceRecordSource +from services.tracing.errors import ( + ProviderConfigUnreadableError, + TraceCollectionConflictError, + TraceCollectionError, + TraceCollectionLimitError, + TraceCollectionNotFoundError, + TraceProviderUnavailableError, + UnknownTraceProviderError, ) -from services.tracing.runtime_bridge import workflow_has_agent_v2 +from services.tracing.providers import ProviderConfigCodec, TraceProviderCatalog, TraceProviderSupport, support_of +from services.tracing.records import BaseTraceInfo, MessageTraceInfo, TraceMessage, WorkflowTraceInfo +from services.tracing.recovery import recover_terminal logger = logging.getLogger(__name__) +_RECONCILE_BATCH = 100 + def _message_inputs_are_empty(inputs: object) -> bool: if isinstance(inputs, str): @@ -47,23 +70,91 @@ def _message_inputs_are_empty(inputs: object) -> bool: return inputs is None or isinstance(inputs, (str, dict, list)) and not inputs -class ApplicationTracingService: - def __init__(self, repository: ApplicationTraceRepository) -> None: - self.repository = repository +@contextmanager +def _collection_errors() -> Generator[None]: + """Report a run's rejection as the collection error its handlers map.""" + try: + yield + except TraceRunNotFoundError as error: + raise TraceCollectionNotFoundError(error.code) from error + except TraceRunConflictError as error: + raise TraceCollectionConflictError(error.code) from error + except TraceRunLimitError as error: + raise TraceCollectionLimitError(error.code) from error + except TraceRunError as error: + raise TraceCollectionError(error.code) from error + + +def _record(run: TraceRunAgg) -> TraceRun: + return TraceRun( + id=run.id, + tenant_id=run.tenant_id, + app_id=run.app_id, + source_kind=run.source_kind, + source_id=run.source_id, + provider=run.provider, + config_revision=run.config_revision, + provider_config=run.provider_config, + trace_id=run.trace_id, + root_span_id=run.root_span_id, + session_id=run.session_id, + capture_content=run.capture_content, + capture_agent=run.capture_agent, + state=run.state, + outer_trace=run.outer_trace, + span_count=run.span_count, + attempts=run.attempts, + error_code=run.error_code, + trace_url=run.trace_url, + provider_trace_id=run.provider_trace_id, + span_urls=run.span_urls, + accepted_at=run.accepted_at, + ) + + +class ApplicationTracingService(TraceCollector): + def __init__( + self, + *, + runs: TraceRunRepository, + run_query: TraceRunQuery, + sources: TraceSourceQuery, + records: TraceRecordSource, + builder: CanonicalTraceBuilder, + providers: TraceProviderCatalog, + codec: ProviderConfigCodec, + egress: Egress, + schedule_export: ScheduleTraceExport, + settings: AgentTraceSettings, + ) -> None: + self._runs = runs + self._run_query = run_query + self._sources = sources + self._records = records + self._builder = builder + self._providers = providers + self._codec = codec + self._egress = egress + self._schedule_export = schedule_export + self._settings = settings def start( self, *, tenant_id: str, app_id: str, - source_kind: Literal["message", "workflow", "node"], + source_kind: TraceSourceKind, source_id: str, session_id: str, - workflow_id: str | None = None, - ) -> ApplicationTraceRun | None: - if not dify_config.OPS_AGENT_TRACE_ENABLED: + workflow_id: str | None, + ) -> TraceRun | None: + """Freeze collection for a source once; ``None`` when the source is not collected. + + ``workflow_id`` lets a trusted workflow builder reserve a run before its source row is saved. + """ + if not self._settings.enabled: return None - run = self.repository.start( + run = self._freeze( tenant_id=tenant_id, app_id=app_id, source_kind=source_kind, @@ -71,7 +162,84 @@ def start( session_id=session_id, workflow_id=workflow_id, ) - return run if run is not None and run.state != "not_collected" else None + return run if run is not None and run.state is not TraceRunState.NOT_COLLECTED else None + + def _freeze( + self, + *, + tenant_id: str, + app_id: str, + source_kind: TraceSourceKind, + source_id: str, + session_id: str, + workflow_id: str | None, + ) -> TraceRun | None: + """The source's frozen run; ``None`` when the source or its owner does not exist.""" + # Re-entry uses the already validated, frozen source. It must not depend + # on LogStore availability or indexing latency after collection began. + if self._sources.app(tenant_id=tenant_id, app_id=app_id) is None: + return None + frozen = self._run_query.source_run( + tenant_id=tenant_id, app_id=app_id, source_kind=source_kind, source_id=source_id + ) + if frozen is not None: + return frozen + if self._source(tenant_id=tenant_id, app_id=app_id, kind=source_kind, source_id=source_id) is None: + # A single-node preview allocates its execution ID before the node row + # is saved. Only the trusted workflow builder may reserve it, after + # its persisted owner workflow is checked. + if source_kind not in {"node", "workflow"} or workflow_id is None: + return None + if self._sources.app_workflow_graph(tenant_id=tenant_id, app_id=app_id, workflow_id=workflow_id) is None: + return None + tracing = self._sources.app_tracing(tenant_id=tenant_id, app_id=app_id) + if tracing is None: + return None + return self._add( + TraceRunAgg.freeze( + run_id=str(uuidv7()), + trace_id=uuidv7().hex, + tenant_id=tenant_id, + app_id=app_id, + source_kind=source_kind, + source_id=source_id, + session_id=session_id, + policy=tracing.policy, + stored_config=tracing.config, + now=naive_utc_now(), + ) + ) + + def _add(self, run: TraceRunAgg) -> TraceRun | None: + if self._runs.add(run): + return _record(run) + # The first frozen destination wins, so one execution never exports to two providers. + return self._run_query.source_run( + tenant_id=run.tenant_id, app_id=run.app_id, source_kind=run.source_kind, source_id=run.source_id + ) + + def _fail_initialization( + self, *, tenant_id: str, app_id: str, source_kind: TraceSourceKind, source_id: str + ) -> TraceRun | None: + """Close a source whose start snapshot is missing; ``None`` when the source does not exist.""" + if self._source(tenant_id=tenant_id, app_id=app_id, kind=source_kind, source_id=source_id) is None: + return None + frozen = self._run_query.source_run( + tenant_id=tenant_id, app_id=app_id, source_kind=source_kind, source_id=source_id + ) + if frozen is not None: + return frozen + return self._add( + TraceRunAgg.initialization_failed( + run_id=str(uuidv7()), + trace_id=uuidv7().hex, + tenant_id=tenant_id, + app_id=app_id, + source_kind=source_kind, + source_id=source_id, + now=naive_utc_now(), + ) + ) def prepare_invocation( self, @@ -79,13 +247,13 @@ def prepare_invocation( tenant_id: str, app_id: str, agent_id: str, - source_kind: Literal["message", "workflow", "node"], + source_kind: TraceSourceKind, source_id: str, session_id: str, node_execution_id: str | None, attempt: int, call_key: str, - workflow_id: str | None = None, + workflow_id: str | None, ) -> ApplicationTracingContext | None: run = self.start( tenant_id=tenant_id, @@ -97,9 +265,12 @@ def prepare_invocation( ) if run is None or not run.capture_agent: return None - invocation = self.repository.invocation( - run_id=run.id, agent_id=agent_id, node_execution_id=node_execution_id, attempt=attempt, call_key=call_key - ) + with _collection_errors(), self._runs.change(run.id, tenant_id=tenant_id) as locked: + if locked is None: + raise TraceCollectionConflictError("tracing_run_not_collecting") + invocation = locked.run.open_invocation( + agent_id=agent_id, node_execution_id=node_execution_id, attempt=attempt, call_key=call_key + ) return ApplicationTracingContext( run_id=run.id, invocation_id=invocation.id, @@ -110,13 +281,30 @@ def prepare_invocation( ) def ingest(self, batch: AgentTraceBatch) -> int: - accepted, run_id = self.repository.ingest( - batch, - max_spans=dify_config.OPS_AGENT_TRACE_MAX_SPANS, - max_bytes=dify_config.OPS_AGENT_TRACE_MAX_BYTES, - ) + """Store a runtime batch and return how many spans the run accepted. + + Raises: + TraceCollectionNotFoundError: The run, its owner, or the invocation does not exist. + TraceCollectionConflictError: The batch conflicts with the run or the run stopped collecting. + TraceCollectionLimitError: The batch exceeds the run's limits. + """ + run_id = str(batch.run_id) + run = self._run_query.get(run_id) + if run is None: + raise TraceCollectionNotFoundError("tracing_run_not_found") + with _collection_errors(), self._runs.change(run_id, tenant_id=run.tenant_id) as locked: + if locked is None: + raise TraceCollectionNotFoundError("tracing_run_not_found") + if not locked.owner_exists: + raise TraceCollectionNotFoundError("tracing_owner_not_found") + accepted = locked.run.ingest( + batch, + max_spans=self._settings.max_spans, + max_bytes=self._settings.max_bytes, + now=naive_utc_now(), + ) # Persistence commits before scheduling. Beat reconciles a broker outage. - self.schedule(run_id) + self._schedule(run_id) return accepted def capture_agent_app_terminal( @@ -134,212 +322,129 @@ def capture_agent_app_terminal( This handoff is made only after the runner observes a terminal result or finishes handling an explicit stop, and it does not rewrite Message data. """ - source = self.repository._source(tenant_id=tenant_id, app_id=app_id, kind="message", source_id=message_id) - if not isinstance(source, Message): - raise TraceCollectionError("tracing_source_not_found", 404) - conversation_id = source.conversation_id - run = self.repository.source_run( + source = self._sources.message(tenant_id=tenant_id, app_id=app_id, message_id=message_id) + if source is None: + raise TraceCollectionNotFoundError("tracing_source_not_found") + run = self._run_query.source_run( tenant_id=tenant_id, app_id=app_id, source_kind="message", source_id=message_id ) - if run is not None and run.state == "not_collected": + if run is not None and run.state is TraceRunState.NOT_COLLECTED: return - info = TraceTask( - TraceTaskName.MESSAGE_TRACE, - message_id=message_id, - conversation_id=conversation_id, - trace_session_id=run.session_id if run else conversation_id, - ).execute() + info = self._records.assemble( + TraceEvent( + event_id=uuid4().hex, + tenant_id=tenant_id, + app_id=app_id, + user_id=None, + occurred_at=datetime.now(UTC), + correlation=TraceCorrelation( + external_trace_id=None, + session_id=run.session_id if run else source.conversation_id, + parent_workflow_run_id=None, + parent_node_execution_id=None, + ), + payload=MessageCompleted(message_id=message_id, conversation_id=source.conversation_id), + ) + ) if not isinstance(info, MessageTraceInfo): raise TraceCollectionError("tracing_terminal_not_supported") - if isinstance(info.message_data, dict): - info.message_data = Message.from_dict(cast(MessageDict, info.message_data)) - if ( - isinstance(info.message_data, Message) - and _message_inputs_are_empty(info.inputs) - and info.message_data.query - ): + inputs = info.inputs + if _message_inputs_are_empty(inputs) and info.message_data.query: # Early errors may precede formatted-message persistence. Preserve only # the known user query on the root; this is not an inferred LLM prompt. - info.inputs = {"query": info.message_data.query} - if output is not None and isinstance(info.message_data, Message): - # This is a detached trace snapshot, not a write to the business row. - info.message_data.answer = output - info.outputs = output if output is not None else info.outputs - info.end_time = datetime.now(UTC) - info.metadata = {**info.metadata, "message_execution_state": outcome} - if outcome != "succeeded": - info.error = f"Agent message {outcome}" - self.capture_terminal(info) - - @staticmethod - def schedule(run_id: str) -> None: - from tasks.application_trace_task import export_application_trace + inputs = {"query": info.message_data.query} + message = info.message_data if output is None else info.message_data.model_copy(update={"answer": output}) + self.capture_terminal( + tenant_id=tenant_id, + info=info.model_copy( + update={ + "message_data": message, + "inputs": inputs, + "outputs": output if output is not None else info.outputs, + "end_time": datetime.now(UTC).replace(tzinfo=None), + "metadata": {**info.metadata, "message_execution_state": outcome}, + "error": info.error if outcome == "succeeded" else f"Agent message {outcome}", + } + ), + ) + def reconcile(self) -> None: + """Schedule unfinished runs again, behind a broker outage, and purge runs past their retention.""" + for run_id in self._run_query.pending_ids(_RECONCILE_BATCH): + self._schedule(run_id) + self._runs.purge(retention_days=self._settings.retention_days, limit=_RECONCILE_BATCH) + + def _schedule(self, run_id: str) -> None: try: - export_application_trace.delay(run_id) + self._schedule_export(run_id) except Exception: logger.warning("Application trace is persisted; scheduling deferred: run_id=%s", run_id) - def _service_account(self, run: ApplicationTraceRun, app_id: str) -> Account: - with self.repository.sessions() as session: - app = session.scalar( - select(App).where(App.id == app_id, App.id == run.app_id, App.tenant_id == run.tenant_id) - ) - if app is None or not app.created_by: - raise TraceCollectionError("tracing_owner_not_found", 404) - account = session.get(Account, app.created_by) - if account is None: - raise TraceCollectionError("tracing_owner_not_found", 404) - account.set_tenant_id_with_session(run.tenant_id, session=session) - return account + def _change(self, run: TraceRun | TraceRunAgg, change: Callable[[TraceRunAgg], object]) -> None: + """Apply ``change`` to a run that may have been purged since it was read.""" + with self._runs.change(run.id, tenant_id=run.tenant_id) as locked: + if locked is not None: + change(locked.run) def _run_for_source( - self, *, tenant_id: str, app_id: str, source_kind: str, source_id: str - ) -> ApplicationTraceRun | None: - run = self.repository.source_run( + self, *, tenant_id: str, app_id: str, source_kind: TraceSourceKind, source_id: str + ) -> TraceRun | None: + run = self._run_query.source_run( tenant_id=tenant_id, app_id=app_id, source_kind=source_kind, source_id=source_id ) if run is None and source_kind == "message": - with self.repository.sessions() as session: - workflow_run_id = session.scalar( - select(Message.workflow_run_id) - .join(App, App.id == Message.app_id) - .where( - Message.id == source_id, - Message.app_id == app_id, - App.tenant_id == tenant_id, - ) - ) - if workflow_run_id: - run = self.repository.source_run( - tenant_id=tenant_id, app_id=app_id, source_kind="workflow", source_id=workflow_run_id + message = self._sources.message(tenant_id=tenant_id, app_id=app_id, message_id=source_id) + if message is not None and message.workflow_run_id: + run = self._run_query.source_run( + tenant_id=tenant_id, app_id=app_id, source_kind="workflow", source_id=message.workflow_run_id ) return run - def _managed_owner(self, app: App, message: Message | None = None, workflow_run_id: str | None = None) -> bool: - """Honor frozen owners first; ordinary workflows retain legacy tracing.""" + def _managed_owner( + self, app: TraceApp, message: SourceMessage | TraceMessage | None, workflow_run_id: str | None + ) -> bool: + """Agent apps and workflows with Agent nodes are collected here; frozen runs win.""" if message is not None: if self._run_for_source( tenant_id=app.tenant_id, app_id=app.id, source_kind="message", source_id=message.id ): return True workflow_run_id = message.workflow_run_id or workflow_run_id - if app.mode == AppMode.AGENT: + if app.mode is TraceAppMode.AGENT: return True if workflow_run_id: - if self.repository.source_run( + if self._run_query.source_run( tenant_id=app.tenant_id, app_id=app.id, source_kind="workflow", source_id=workflow_run_id ): return True - workflow = self.repository._source( - tenant_id=app.tenant_id, app_id=app.id, kind="workflow", source_id=workflow_run_id + workflow = self._sources.workflow_run( + tenant_id=app.tenant_id, app_id=app.id, workflow_run_id=workflow_run_id ) - if isinstance(workflow, WorkflowRun): + if workflow is not None: return workflow_has_agent_v2(workflow.graph) - if app.mode == AppMode.ADVANCED_CHAT or (app.mode == AppMode.WORKFLOW and workflow_run_id): + if app.mode is TraceAppMode.ADVANCED_CHAT or (app.mode is TraceAppMode.WORKFLOW and workflow_run_id): # Naming may run before Message.workflow_run_id is populated. Use # the owned app workflow; an unknown snapshot fails closed. The same # applies to a Workflow whose run is not yet indexed by LogStore. - with self.repository.sessions() as session: - definition = ( - session.scalar( - select(Workflow).where( - Workflow.id == app.workflow_id, - Workflow.app_id == app.id, - Workflow.tenant_id == app.tenant_id, - ) - ) - if app.workflow_id - else None - ) - return definition is None or workflow_has_agent_v2(definition.graph) - return False - - def _capture_auxiliary(self, info: BaseTraceInfo) -> bool: - """Resolve real business owners even when message_id is a workflow-log ID.""" - actual_message_id = info.message_data.id if isinstance(info.message_data, Message) else None - message_ids = [actual_message_id, info.metadata.get("message_id"), info.message_id] - message: Message | None = None - app: App | None = None - with self.repository.sessions() as session: - for message_id in message_ids: - if not isinstance(message_id, str): - continue - owner = session.execute( - select(Message, App).join(App, App.id == Message.app_id).where(Message.id == message_id) - ).first() - if owner is not None: - message, app = owner - break - conversation_id = info.metadata.get("conversation_id") - if app is None and isinstance(conversation_id, str): - app = session.scalar( - select(App) - .join(Conversation, Conversation.app_id == App.id) - .where(Conversation.id == conversation_id) - ) - if app is not None: - message = session.scalar( - select(Message) - .where(Message.conversation_id == conversation_id, Message.app_id == app.id) - .order_by(Message.created_at.desc()) - .limit(1) - ) - if app is not None: - return self._managed_owner(app, message) - workflow_run_id = info.metadata.get("workflow_run_id") - if isinstance(workflow_run_id, str): - with self.repository.sessions() as session: - # A durable managed owner already passed the trusted start boundary. - # Suppress auxiliary content even before LogStore indexing catches up. - managed_app = session.scalar( - select(App) - .join(ApplicationTraceRun, ApplicationTraceRun.app_id == App.id) - .where( - ApplicationTraceRun.source_kind == "workflow", - ApplicationTraceRun.source_id == workflow_run_id, - ApplicationTraceRun.tenant_id == App.tenant_id, - ) - ) - if managed_app is not None: - return True - app_id = info.metadata.get("app_id") - app = ( - session.scalar(select(App).where(App.id == app_id)) - if isinstance(app_id, str) - else session.scalar( - select(App) - .join(WorkflowRun, WorkflowRun.app_id == App.id) - .where(WorkflowRun.id == workflow_run_id, WorkflowRun.tenant_id == App.tenant_id) - ) - ) - if app is not None: - return self._managed_owner(app, workflow_run_id=workflow_run_id) + graph = ( + self._sources.app_workflow_graph(tenant_id=app.tenant_id, app_id=app.id, workflow_id=app.workflow_id) + if app.workflow_id + else None + ) + return graph is None or workflow_has_agent_v2(graph) return False - def _link_managed_parent(self, info: WorkflowTraceInfo, tenant_id: str) -> None: - """A legacy child cannot await a parent exported by a different pipeline.""" + @override + def link_managed_parent(self, *, tenant_id: str, info: WorkflowTraceInfo) -> WorkflowTraceInfo: + """A provider-exported nested workflow names the Agent-collected parent it cannot attach to.""" parent_run_id, parent_node_id = info.resolved_parent_context if not parent_run_id: - return - with self.repository.sessions() as session: - parent = session.scalar( - select(ApplicationTraceRun) - .join(App, App.id == ApplicationTraceRun.app_id) - .where( - ApplicationTraceRun.source_kind == "workflow", - ApplicationTraceRun.source_id == parent_run_id, - ApplicationTraceRun.tenant_id == tenant_id, - App.tenant_id == tenant_id, - ) - ) + return info + parent = self._run_query.workflow_run_trace(tenant_id=tenant_id, workflow_run_id=parent_run_id) if parent is None: - return - if ( - self.repository._source(tenant_id=tenant_id, app_id=parent.app_id, kind="workflow", source_id=parent_run_id) - is None - ): - return + return info + if self._sources.workflow_run(tenant_id=tenant_id, app_id=parent.app_id, workflow_run_id=parent_run_id) is None: + return info metadata = dict(info.metadata) metadata.pop("parent_trace_context", None) metadata["linked_parent_workflow_run_id"] = parent_run_id @@ -347,27 +452,37 @@ def _link_managed_parent(self, info: WorkflowTraceInfo, tenant_id: str) -> None: metadata["linked_parent_node_execution_id"] = parent_node_id metadata["linked_parent_application_trace_id"] = parent.id metadata["linked_parent_trace_id"] = parent.trace_id - info.metadata = metadata + return info.model_copy(update={"metadata": metadata}) - def capture_terminal(self, info: BaseTraceInfo) -> bool: - """Hand managed traces to durable collection while preserving legacy scope.""" + @override + def capture_terminal(self, *, tenant_id: str, info: BaseTraceInfo) -> bool: + """Take records of Agent-managed executions; ``False`` leaves the record to its provider.""" + source_kind: TraceSourceKind if isinstance(info, WorkflowTraceInfo): app_id = str(info.metadata.get("app_id", "")) source_id, source_kind = info.workflow_run_id, "workflow" - elif isinstance(info, MessageTraceInfo) and isinstance(info.message_data, Message): + elif isinstance(info, MessageTraceInfo): app_id = info.message_data.app_id source_id, source_kind = info.message_data.id, "message" else: - return self._capture_auxiliary(info) - with self.repository.sessions() as session: - app = session.scalar(select(App).where(App.id == app_id)) - if app is None: - return False - message = ( - session.scalar(select(Message).where(Message.id == source_id, Message.app_id == app_id)) - if source_kind == "message" + # Auxiliary operations of a managed execution belong to its Agent trace. Naming may + # report no message, so the record's own app decides then. + message_data = info.message_data + owner_app_id = message_data.app_id if message_data is not None else info.metadata.get("app_id") + owner = ( + self._sources.app(tenant_id=tenant_id, app_id=owner_app_id) + if isinstance(owner_app_id, str) and owner_app_id else None ) + return owner is not None and self._managed_owner(owner, message_data, None) + app = self._sources.app(tenant_id=tenant_id, app_id=app_id) + if app is None: + return False + message = ( + self._sources.message(tenant_id=tenant_id, app_id=app_id, message_id=source_id) + if source_kind == "message" + else None + ) managed = self._managed_owner(app, message, source_id if source_kind == "workflow" else None) run = self._run_for_source(tenant_id=app.tenant_id, app_id=app.id, source_kind=source_kind, source_id=source_id) if run is not None and source_kind == "message" and run.source_kind == "workflow": @@ -375,154 +490,75 @@ def capture_terminal(self, info: BaseTraceInfo) -> bool: return True if run is None: if not managed: - if isinstance(info, WorkflowTraceInfo): - self._link_managed_parent(info, app.tenant_id) return False - run = self.repository.record_initialization_failure( - tenant_id=app.tenant_id, - app_id=app.id, - source_kind=source_kind, - source_id=source_id, + run = self._fail_initialization( + tenant_id=app.tenant_id, app_id=app.id, source_kind=source_kind, source_id=source_id ) + # A missing start snapshot must never fall back to an exporter without its capture policy. logger.warning( - "Application trace snapshot missing; legacy export suppressed: app_id=%s source_id=%s", - app_id, - source_id, + "Application trace snapshot missing; export suppressed: app_id=%s source_id=%s", app_id, source_id ) if run is None: return True - if run.state in {"not_collected", "failed", "uncertain", "accepted", "exporting"}: + if not run.state.accepts_terminal: return True if ( - app.mode == AppMode.AGENT + app.mode is TraceAppMode.AGENT and message is not None and info.metadata.get("message_execution_state") not in ("succeeded", "failed", "cancelled") + and message.agent_stop_requested ): - from services.tracing.recovery import agent_stop_requested - - if agent_stop_requested(message): - # An already queued consumer notification may arrive after a - # stop. Its ordinary Message snapshot cannot own the runtime - # outcome, even if span finalization has already completed. - return True - loader = RepositoryWorkflowExecutionLoader(lambda requested_app: self._service_account(run, requested_app)) - trace = CanonicalTraceBuilder(loader).build(info) + # An already queued consumer notification may arrive after a + # stop. Its ordinary Message snapshot cannot own the runtime + # outcome, even if span finalization has already completed. + return True + trace = self._builder.build(info) if trace is None: raise TraceCollectionError("tracing_terminal_not_supported") - trace = remap_outer_trace( - trace, trace_id=run.trace_id, capture_content=run.capture_content, agent_message=source_kind == "message" - ) - trace = trace.model_copy(update={"session_id": run.session_id}) - self.repository.terminal(run.id, trace.model_dump(mode="json"), trace.root_span_id) - self.schedule(run.id) + self._change(run, lambda current: current.capture_terminal(trace, now=naive_utc_now())) + self._schedule(run.id) return True - def settle_preview(self, run: ApplicationTraceRun) -> None: - if run.source_kind != "node" or run.outer_trace is not None: - return - node = self.repository._source(tenant_id=run.tenant_id, app_id=run.app_id, kind="node", source_id=run.source_id) - if not isinstance(node, WorkflowNodeExecutionModel) or str(node.status) in {"running", "pending", "paused"}: + def _settle_preview(self, run: TraceRunAgg) -> None: + if run.source_kind != "node" or not run.awaits_terminal: return - error = node.error - started_at = node.created_at.replace(tzinfo=UTC) - ended_at = started_at + timedelta(seconds=node.elapsed_time or 0) - span = CanonicalSpan( - id=run.root_span_id, - parent_id=None, - name=node.title or "Agent preview", - kind=CanonicalSpanKind.AGENT, - start_time=started_at, - end_time=ended_at, - inputs=sanitize(node.inputs) if run.capture_content else None, - outputs=sanitize(node.outputs) if run.capture_content else None, - status=CanonicalSpanStatus.ERROR if error else CanonicalSpanStatus.OK, - error=error if run.capture_content else "Node execution failed" if error else None, - metadata={"node_execution_id": node.id, "node_id": node.node_id, "source": "single_step"}, - ) - trace = CanonicalTrace( - trace_id=run.trace_id, session_id=run.session_id, root_span_id=run.root_span_id, spans=(span,) - ) - self.repository.terminal(run.id, trace.model_dump(mode="json"), run.root_span_id) - - def assemble(self, run: ApplicationTraceRun) -> CanonicalTrace: - if run.outer_trace is None: - raise TraceCollectionError("tracing_terminal_missing") - outer = CanonicalTrace.model_validate(run.outer_trace) - invocations, batches = self.repository.collected(run.id) - invocation_by_id = {inv.id: inv for inv in invocations} - spans = list(outer.spans) - internal_by_invocation: dict[str, dict[str, CanonicalSpan]] = {inv.id: {} for inv in invocations} - for batch in batches: - if batch.invocation_id not in invocation_by_id: - raise TraceCollectionError("tracing_invocation_not_found") - for data in batch.spans: - span = normalize_agent_span(data, capture_content=run.capture_content) - current = internal_by_invocation[batch.invocation_id].get(span.id) - if current is not None and current != span: - raise TraceCollectionError("tracing_span_conflict") - internal_by_invocation[batch.invocation_id][span.id] = span - for invocation in invocations: - children = list(internal_by_invocation[invocation.id].values()) - if not children: - raise TraceCollectionError("tracing_invocation_empty") - valid_ids = {child.id for child in children} - if any(child.parent_id not in valid_ids and child.parent_id != invocation.span_id for child in children): - raise TraceCollectionError("tracing_invocation_parent_invalid") - root_children = [child for child in children if child.parent_id == invocation.span_id] - if not root_children: - raise TraceCollectionError("tracing_invocation_root_missing") - errors = [child.error for child in root_children if child.status == CanonicalSpanStatus.ERROR] - spans.append( - CanonicalSpan( - id=invocation.span_id, - parent_id=invocation.parent_span_id, - name=f"Agent attempt {invocation.attempt + 1}", - kind=CanonicalSpanKind.AGENT, - start_time=min(child.start_time for child in children), - end_time=max(child.end_time or child.start_time for child in children), - status=CanonicalSpanStatus.ERROR if errors else CanonicalSpanStatus.OK, - error=errors[0] if errors else None, - metadata={ - "agent_id": invocation.agent_id, - "attempt": invocation.attempt, - "node_execution_id": invocation.node_execution_id, - "invocation_id": invocation.id, - }, - ) - ) - spans.extend(children) - trusted = { - "tenant_id": run.tenant_id, - "app_id": run.app_id, - "session_id": run.session_id, - "dify.tenant_id": run.tenant_id, - "dify.app_id": run.app_id, - "dify.tracing.run_id": run.id, - "dify.source.id": run.source_id, - } - spans = [span.model_copy(update={"metadata": {**span.metadata, **trusted}}) for span in spans] - return CanonicalTrace( - trace_id=run.trace_id, - session_id=run.session_id, - root_span_id=outer.root_span_id, - spans=order_spans(outer.root_span_id, spans), - ) + node = self._sources.node_execution(tenant_id=run.tenant_id, app_id=run.app_id, node_execution_id=run.source_id) + if node is not None: + self._change(run, lambda current: current.settle_preview(node, now=naive_utc_now())) + + def _claim(self, run: TraceRunAgg) -> tuple[TraceRunAgg, str] | None: + token = str(uuidv7()) + with self._runs.change(run.id, tenant_id=run.tenant_id) as locked: + if locked is None or not locked.run.claim_export( + token=token, + owner_exists=locked.owner_exists, + collection_timeout=self._settings.collection_timeout_seconds, + lease_seconds=self._settings.export_lease_seconds, + now=naive_utc_now(), + ): + return None + return locked.run, token + + def _renew_export_lease(self, run: TraceRunAgg, token: str) -> bool: + with self._runs.change(run.id, tenant_id=run.tenant_id) as locked: + return locked is not None and locked.run.renew_export_lease(token, now=naive_utc_now()) @contextmanager - def _export_lease_heartbeat(self, run_id: str, token: str) -> Generator[None]: + def _export_lease_heartbeat(self, run: TraceRunAgg, token: str) -> Generator[None]: """Keep slow exporters leased without holding a transaction across vendor I/O. A lost database connection is not proof that the exporter stopped. The attempt continues; the token-checked completion can reconcile a late ACK after lease expiry unless a manual retry has already replaced its token. """ + run_id = run.id stopped = threading.Event() - interval = dify_config.OPS_AGENT_TRACE_EXPORT_LEASE_SECONDS / 3 + interval = self._settings.export_lease_seconds / 3 def heartbeat() -> None: while not stopped.wait(interval): try: - if not self.repository.renew_export_lease(run_id, token): + if not self._renew_export_lease(run, token): return except Exception as error: logger.warning( @@ -550,50 +586,51 @@ def heartbeat() -> None: logger.warning("Application trace heartbeat shutdown pending: run_id=%s", run_id) def deliver(self, run_id: str) -> None: - record = self.repository.get(run_id) - if record is None: - return - from services.tracing.recovery import recover_terminal - - recover_terminal(self, record) - record = self.repository.get(run_id) - if record is None: + run = self._runs.load(run_id) + if run is None: return - self.settle_preview(record) - claim = self.repository.claim( - run_id, - collection_timeout=dify_config.OPS_AGENT_TRACE_COLLECTION_TIMEOUT, - lease_seconds=dify_config.OPS_AGENT_TRACE_EXPORT_LEASE_SECONDS, + recover_terminal( + run, + runs=self._runs, + sources=self._sources, + records=self._records, + collector=self, + collection_timeout=self._settings.collection_timeout_seconds, + now=naive_utc_now(), ) + run = self._runs.load(run_id) + if run is None: + return + self._settle_preview(run) + claim = self._claim(run) if claim is None: return run, token = claim - with self._export_lease_heartbeat(run_id, token): + with self._export_lease_heartbeat(run, token): try: try: - trace = self.assemble(run) - except ValueError as error: + trace = run.exported_trace(self._run_query.batches(run.id)) + except (ValueError, KeyError, TypeError) as error: # Invalid persisted trees cannot be repaired by retrying a # provider request. ValidationError also inherits ValueError. raise TraceCollectionError("tracing_payload_invalid") from error - config = dict(run.provider_config) - config.update(OpsTraceManager.decrypt_tracing_config(run.tenant_id, run.provider, run.provider_config)) - result = export_trace(run.provider, config, trace) - except TraceExportUncertainError: - self.repository.complete(run_id, token, state="uncertain", error_code="provider_acceptance_unknown") - except RetryableTraceDispatchError: - if run.attempts >= dify_config.OPS_AGENT_TRACE_MAX_ATTEMPTS: - self.repository.complete(run_id, token, state="failed", error_code="provider_retry_exhausted") - else: - self.repository.complete( - run_id, - token, - state="retrying", - error_code="provider_temporarily_unavailable", - retry_delay=min(300, 5 * 2 ** min(run.attempts - 1, 6)), - ) - except TraceCollectionError as error: - self.repository.complete(run_id, token, state="failed", error_code=error.code) + provider = self._frozen_provider(run.provider) + result = provider.send(self._codec.decode(run.tenant_id, type(provider), run.provider_config), trace) + except ExportOutcomeUnknownError: + self._change(run, lambda current: current.export_uncertain(token, now=naive_utc_now())) + except ExportRetryableError: + self._change( + run, + lambda current: current.export_deferred( + token, max_attempts=self._settings.max_export_attempts, now=naive_utc_now() + ), + ) + except (TraceCollectionError, TraceRunError) as error: + self._export_failed(run, token, error.code) + except ProviderConfigUnreadableError: + self._export_failed(run, token, "tracing_config_unreadable") + except ExportRejectedError: + self._export_failed(run, token, "provider_export_rejected") except Exception as error: logger.warning( "Application trace export failed: run_id=%s provider=%s error_type=%s", @@ -601,13 +638,51 @@ def deliver(self, run_id: str) -> None: run.provider, type(error).__name__, ) - self.repository.complete(run_id, token, state="failed", error_code="provider_export_rejected") + self._export_failed(run, token, "provider_export_rejected") else: - self.repository.complete(run_id, token, state="accepted", result=result) + self._change( + run, + lambda current: current.export_accepted( + token, + provider_trace_id=result.trace_id, + trace_url=result.trace_url, + span_urls=result.span_urls, + now=naive_utc_now(), + ), + ) + + def _export_failed(self, run: TraceRunAgg, token: str, error_code: str) -> None: + self._change(run, lambda current: current.export_failed(token, error_code, now=naive_utc_now())) + + def _frozen_provider(self, provider_id: str) -> TraceProvider[Any]: + try: + return self._providers.resolve(provider_id)(self._egress) + except (UnknownTraceProviderError, TraceProviderUnavailableError): + raise TraceCollectionError("tracing_provider_unavailable") from None + + def _frozen_capabilities(self, provider_id: str) -> TraceProviderSupport | None: + try: + return support_of(self._providers.resolve(provider_id)) + except (UnknownTraceProviderError, TraceProviderUnavailableError): + return None + + def _frozen_project_url(self, run: TraceRun) -> str | None: + """Present the run's own destination, never the app's current one.""" + try: + provider = self._providers.resolve(run.provider) + return provider(self._egress).project_url(provider.config_model.model_validate(run.provider_config)) + except (UnknownTraceProviderError, TraceProviderUnavailableError, ValidationError): + return None def status( - self, *, tenant_id: str, app_id: str, source_kind: str, source_id: str, node_execution_id: str | None = None - ) -> dict[str, Any]: + self, + *, + tenant_id: str, + app_id: str, + source_kind: TraceSourceKind, + source_id: str, + node_execution_id: str | None, + ) -> TraceRunStatus: run, node_execution_id = self._resolve_status_source( tenant_id=tenant_id, app_id=app_id, @@ -617,75 +692,102 @@ def status( ) return self._status_for_run(run, source_kind, source_id, node_execution_id) + def _source(self, *, tenant_id: str, app_id: str, kind: TraceSourceKind, source_id: str) -> TraceSource | None: + match kind: + case "message": + return self._sources.message(tenant_id=tenant_id, app_id=app_id, message_id=source_id) + case "workflow": + return self._sources.workflow_run(tenant_id=tenant_id, app_id=app_id, workflow_run_id=source_id) + case "node": + return self._sources.node_execution(tenant_id=tenant_id, app_id=app_id, node_execution_id=source_id) + def _resolve_status_source( - self, *, tenant_id: str, app_id: str, source_kind: str, source_id: str, node_execution_id: str | None - ) -> tuple[ApplicationTraceRun | None, str | None]: + self, + *, + tenant_id: str, + app_id: str, + source_kind: TraceSourceKind, + source_id: str, + node_execution_id: str | None, + ) -> tuple[TraceRun | None, str | None]: """Resolve node log aliases through their owned workflow, outside SQL transactions.""" - source = self.repository._source(tenant_id=tenant_id, app_id=app_id, kind=source_kind, source_id=source_id) + source = self._source(tenant_id=tenant_id, app_id=app_id, kind=source_kind, source_id=source_id) if source is None: - raise TraceCollectionError("tracing_source_not_found", 404) + raise TraceCollectionNotFoundError("tracing_source_not_found") if source_kind == "node" and node_execution_id is not None and node_execution_id != source_id: - raise TraceCollectionError("tracing_node_not_found", 404) + raise TraceCollectionNotFoundError("tracing_node_not_found") run = self._run_for_source(tenant_id=tenant_id, app_id=app_id, source_kind=source_kind, source_id=source_id) # A standalone preview (including a frozen disabled decision) owns its # status even if its source later carries a workflow run reference. - if run is None and isinstance(source, WorkflowNodeExecutionModel) and source.workflow_run_id: - workflow = self.repository._source( - tenant_id=tenant_id, app_id=app_id, kind="workflow", source_id=source.workflow_run_id + if run is None and isinstance(source, SourceNodeExecution) and source.workflow_run_id: + workflow = self._sources.workflow_run( + tenant_id=tenant_id, app_id=app_id, workflow_run_id=source.workflow_run_id ) - if ( - not isinstance(workflow, WorkflowRun) - or not source.workflow_id - or workflow.workflow_id != source.workflow_id - ): - raise TraceCollectionError("tracing_source_not_found", 404) - run = self.repository.source_run( + if workflow is None or not source.workflow_id or workflow.workflow_id != source.workflow_id: + raise TraceCollectionNotFoundError("tracing_source_not_found") + run = self._run_query.source_run( tenant_id=tenant_id, app_id=app_id, source_kind="workflow", source_id=workflow.id ) node_execution_id = source.id return run, node_execution_id def _status_for_run( - self, run: ApplicationTraceRun | None, source_kind: str, source_id: str, node_execution_id: str | None - ) -> dict[str, Any]: - if run is None or run.state == "not_collected": - return {"status": "not_collected", "source_kind": source_kind, "source_id": source_id} - invocations, _ = self.repository.collected(run.id) + self, run: TraceRun | None, source_kind: TraceSourceKind, source_id: str, node_execution_id: str | None + ) -> TraceRunStatus: + if run is None or run.state is TraceRunState.NOT_COLLECTED: + return TraceRunStatus( + status=TraceRunState.NOT_COLLECTED, + source_kind=source_kind, + source_id=source_id, + provider=None, + trace_id=None, + trace_url=None, + span_url=None, + project_url=None, + config_revision=None, + capture_content=None, + capture_agent=None, + span_count=None, + attempts=None, + error_code=None, + accepted_at=None, + capabilities=None, + ) + invocations = self._run_query.invocations(run.id) selected = ( [inv for inv in invocations if inv.node_execution_id == node_execution_id] if node_execution_id else [] ) if node_execution_id and not selected and not (source_kind == "node" and run.source_kind == "workflow"): - raise TraceCollectionError("tracing_node_not_found", 404) - return { - "status": run.state, - "source_kind": source_kind, - "source_id": source_id, - "provider": run.provider or None, - "trace_id": run.provider_trace_id or run.trace_id, - "trace_url": run.trace_url, - "span_url": run.span_urls.get(selected[-1].span_id) if selected else None, - "project_url": frozen_project_url(run.provider, run.provider_config), - "config_revision": run.config_revision or None, - "capture_content": run.capture_content, - "capture_agent": run.capture_agent, - "span_count": run.span_count, - "attempts": run.attempts, - "error_code": run.error_code, - "accepted_at": run.accepted_at, - "queryable": None, - "capabilities": provider_capabilities(run.provider) if run.provider else None, - } + raise TraceCollectionNotFoundError("tracing_node_not_found") + return TraceRunStatus( + status=run.state, + source_kind=source_kind, + source_id=source_id, + provider=run.provider or None, + trace_id=run.provider_trace_id or run.trace_id, + trace_url=run.trace_url, + span_url=run.span_urls.get(selected[-1].span_id) if selected else None, + project_url=self._frozen_project_url(run) if run.provider else None, + config_revision=run.config_revision or None, + capture_content=run.capture_content, + capture_agent=run.capture_agent, + span_count=run.span_count, + attempts=run.attempts, + error_code=run.error_code, + accepted_at=run.accepted_at, + capabilities=self._frozen_capabilities(run.provider) if run.provider else None, + ) def retry( self, *, tenant_id: str, app_id: str, - source_kind: str, + source_kind: TraceSourceKind, source_id: str, acknowledge_possible_duplicates: bool, - node_execution_id: str | None = None, - ) -> dict[str, Any]: + node_execution_id: str | None, + ) -> TraceRunStatus: run, node_execution_id = self._resolve_status_source( tenant_id=tenant_id, app_id=app_id, @@ -694,12 +796,13 @@ def retry( node_execution_id=node_execution_id, ) if run is None: - raise TraceCollectionError("tracing_run_not_found", 404) + raise TraceCollectionNotFoundError("tracing_run_not_found") self._status_for_run(run, source_kind, source_id, node_execution_id) - self.repository.retry(run_id=run.id, acknowledge_possible_duplicates=acknowledge_possible_duplicates) - self.schedule(run.id) - return self._status_for_run(self.repository.get(run.id), source_kind, source_id, node_execution_id) - - -def application_tracing_service() -> ApplicationTracingService: - return ApplicationTracingService(ApplicationTraceRepository(sessionmaker(bind=db.engine, expire_on_commit=False))) + with _collection_errors(), self._runs.change(run.id, tenant_id=tenant_id) as locked: + if locked is None: + raise TraceCollectionNotFoundError("tracing_run_not_found") + locked.run.release_for_retry( + acknowledge_possible_duplicates=acknowledge_possible_duplicates, now=naive_utc_now() + ) + self._schedule(run.id) + return self._status_for_run(self._run_query.get(run.id), source_kind, source_id, node_execution_id) diff --git a/api/services/workflow/node_output_inspector_service.py b/api/services/workflow/node_output_inspector_service.py index 93d6ded71ef6cc..274b338a3a4be9 100644 --- a/api/services/workflow/node_output_inspector_service.py +++ b/api/services/workflow/node_output_inspector_service.py @@ -59,7 +59,7 @@ WorkflowAgentBindingError, WorkflowAgentBindingResolver, ) -from core.workflow.nodes.agent_v2.discriminator import is_dify_agent_node_data +from domain.workflow.agent_node import is_dify_agent_node_data from factories.file_factory.builders import build_from_mapping from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus from graphon.file import helpers as file_helpers diff --git a/api/services/workflow_node_execution_trace_service.py b/api/services/workflow_node_execution_trace_service.py index f016afa80ea1f0..37d2803bb6b42f 100644 --- a/api/services/workflow_node_execution_trace_service.py +++ b/api/services/workflow_node_execution_trace_service.py @@ -10,7 +10,7 @@ from pydantic import BaseModel, ConfigDict, ValidationError from sqlalchemy.orm import Session -from core.app.workflow.retry_history import RETRY_HISTORY_PROCESS_DATA_KEY, WorkflowNodeRetryAttempt +from domain.workflow.retry_history import RETRY_HISTORY_PROCESS_DATA_KEY, WorkflowNodeRetryAttempt from libs.helper import to_timestamp from models.workflow import WorkflowNodeExecutionModel from repositories.api_workflow_node_execution_repository import DifyAPIWorkflowNodeExecutionRepository diff --git a/api/services/workflow_persistence.py b/api/services/workflow_persistence.py index 28bb6750b83d0f..6707c010ebd482 100644 --- a/api/services/workflow_persistence.py +++ b/api/services/workflow_persistence.py @@ -9,12 +9,10 @@ from datetime import datetime from typing import Any, override +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import WorkflowRunCompleted + from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, WorkflowAppGenerateEntity -from core.app.workflow.retry_history import RETRY_HISTORY_PROCESS_DATA_KEY, WorkflowNodeRetryAttempt -from core.helper.trace_id_helper import ParentTraceContext -from core.ops.base_trace_instance import application_tracing_hooks -from core.ops.entities.trace_entity import TraceTaskName -from core.ops.ops_trace_manager import TraceQueueManager, TraceTask from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository from core.tools.workflow_as_tool.repository import WorkflowToolSource from core.workflow.node_execution import WorkflowNodeExecution as DifyWorkflowNodeExecution @@ -22,6 +20,8 @@ from core.workflow.system_variables import SystemVariableKey from core.workflow.variable_prefixes import SYSTEM_VARIABLE_NODE_ID from core.workflow.workflow_run_outputs import project_node_outputs_for_workflow_run +from domain.workflow.retry_history import RETRY_HISTORY_PROCESS_DATA_KEY, WorkflowNodeRetryAttempt +from extensions.tracing import application_tracing_hooks from graphon.engine.layer import Layer from graphon.engine_events import ( EngineEvent, @@ -77,14 +77,14 @@ def __init__( workflow_info: PersistenceWorkflowInfo, workflow_execution_repository: WorkflowExecutionRepository, workflow_node_execution_repository: WorkflowNodeExecutionRepository, - trace_manager: TraceQueueManager | None = None, + trace_scope: TraceScope | None, ) -> None: super().__init__() self._application_generate_entity = application_generate_entity self._workflow_info = workflow_info self._workflow_execution_repository = workflow_execution_repository self._workflow_node_execution_repository = workflow_node_execution_repository - self._trace_manager = trace_manager + self._trace_scope = trace_scope self._workflow_execution: WorkflowExecution | None = None self._node_execution_cache: dict[str, WorkflowNodeExecution] = {} @@ -129,6 +129,8 @@ def create_workflow_tool_event_listener( workflow_node_execution_repository=self._workflow_node_execution_repository.for_workflow_tool( source.app_id ), + # Hidden tool nodes belong to the caller's run, whose trace already covers them. + trace_scope=None, ) tool_layer.set_node_run_indices(self._node_run_indices) tool_layer.set_node_execution_history(node_executions) @@ -242,7 +244,7 @@ def _finish_workflow_execution(self, execution: WorkflowExecution) -> None: if execution.status in (WorkflowExecutionStatus.FAILED, WorkflowExecutionStatus.STOPPED): self._fail_running_node_executions(error_message=execution.error_message or "") self._workflow_execution_repository.save_synchronously(execution) - self._enqueue_trace_task(execution) + self._emit_workflow_trace(execution) _inspector_publish_workflow_completed(workflow_run_id=execution.id_, status=str(execution.status.value)) # ------------------------------------------------------------------ @@ -517,32 +519,17 @@ def _fail_running_node_executions(self, *, error_message: str) -> None: unfinished.append(execution) self._workflow_node_execution_repository.save_many(unfinished) - def _enqueue_trace_task(self, execution: WorkflowExecution) -> None: - if not self._trace_manager: + def _emit_workflow_trace(self, execution: WorkflowExecution) -> None: + if self._trace_scope is None: return - conversation_id = self._system_variables().get(SystemVariableKey.CONVERSATION_ID.value) - external_trace_id = None - trace_session_id = None - parent_trace_context = None - if isinstance(self._application_generate_entity, (WorkflowAppGenerateEntity, AdvancedChatAppGenerateEntity)): - extras = self._application_generate_entity.extras - external_trace_id = extras.get("external_trace_id") - trace_session_id = extras.get("trace_session_id") - parent_trace_context = extras.get("parent_trace_context") - if isinstance(parent_trace_context, ParentTraceContext): - parent_trace_context = parent_trace_context.model_dump(exclude_none=True) - - trace_task = TraceTask( - TraceTaskName.WORKFLOW_TRACE, - workflow_execution=execution, - conversation_id=conversation_id, - user_id=self._trace_manager.user_id, - external_trace_id=external_trace_id, - trace_session_id=trace_session_id, - parent_trace_context=parent_trace_context, + self._trace_scope.emit( + WorkflowRunCompleted( + workflow_run_id=execution.id_, + conversation_id=conversation_id if isinstance(conversation_id, str) and conversation_id else None, + total_tokens=execution.total_tokens, + ) ) - self._trace_manager.add_trace_task(trace_task) def _system_variables(self) -> Mapping[str, Any]: runtime_state = self.runtime_state diff --git a/api/services/workflow_run_agg.py b/api/services/workflow_run_agg.py index d0730d7b1062b0..7365b4d5204330 100644 --- a/api/services/workflow_run_agg.py +++ b/api/services/workflow_run_agg.py @@ -153,7 +153,7 @@ def __init__( ), workflow_execution_repository=workflow_execution_repository, workflow_node_execution_repository=workflow_node_execution_repository, - trace_manager=prepared.generate_entity.trace_manager, + trace_scope=prepared.generate_entity.trace_scope, ) application_layers: list[Layer] = [ self._persistence_layer, diff --git a/api/services/workflow_service.py b/api/services/workflow_service.py index 65225738a758e2..0f540254a82b68 100644 --- a/api/services/workflow_service.py +++ b/api/services/workflow_service.py @@ -38,7 +38,6 @@ from core.workflow.node_runtime import ( apply_dify_debug_email_recipient, ) -from core.workflow.nodes.agent_v2.discriminator import is_dify_agent_node_data from core.workflow.nodes.human_input.callback import ( DifyHITLCallback, render_form_content_before_submission, @@ -56,6 +55,7 @@ ) from core.workflow.variable_pool_initializer import add_node_inputs_to_pool, add_variables_to_pool from core.workflow.workflow_entry import WorkflowEntry +from domain.workflow.agent_node import is_dify_agent_node_data from enterprise.telemetry.draft_trace import enqueue_draft_node_execution_trace from enums import CloudPlan, DeploymentEdition from events.app_event import app_draft_workflow_was_synced, app_published_workflow_was_updated diff --git a/api/tasks/application_trace_task.py b/api/tasks/application_trace_task.py index a7d68d69b8dd1a..e1ea6cbe05e9bd 100644 --- a/api/tasks/application_trace_task.py +++ b/api/tasks/application_trace_task.py @@ -1,22 +1,15 @@ -"""Durable application trace delivery on the existing ops_trace worker queue.""" +"""Durable application trace delivery routed by ``OPS_TRACE_QUEUE``.""" from celery import shared_task -from configs import dify_config +from extensions.ext_application_services import application_services -@shared_task(queue="ops_trace") +@shared_task def export_application_trace(run_id: str) -> None: - from services.tracing.service import application_tracing_service + application_services().tracing.agent_traces.deliver(run_id) - application_tracing_service().deliver(run_id) - -@shared_task(queue="ops_trace") +@shared_task def reconcile_application_traces() -> None: - from services.tracing.service import application_tracing_service - - service = application_tracing_service() - for run_id in service.repository.pending_ids(): - service.schedule(run_id) - service.repository.purge(retention_days=dify_config.OPS_AGENT_TRACE_RETENTION_DAYS) + application_services().tracing.agent_traces.reconcile() diff --git a/api/tasks/ops_trace_task.py b/api/tasks/ops_trace_task.py index 1295b3f97b4ba3..cd78df32fb41bf 100644 --- a/api/tasks/ops_trace_task.py +++ b/api/tasks/ops_trace_task.py @@ -1,157 +1,30 @@ -""" -Celery task for asynchronous ops trace dispatch. - -Trace providers may report retryable dispatch failures through the core -exception contract. Unknown acceptance is terminal to prevent automatic replay -of a possibly accepted prefix. The task preserves the payload file only -when Celery accepts the retry request; successful dispatches and terminal -failures clean up the stored payload. +"""Retired entry point of the previous trace pipeline. -One concrete producer today is Phoenix nested workflow tracing. The outer -workflow tool span publishes a restorable parent span context asynchronously, -while the nested workflow trace may be picked up by Celery first. In that -ordering window, the provider raises a retryable core exception instead of -dropping the trace or emitting it under the wrong parent. The task intentionally -does not know that the provider is Phoenix; it only honors the core retryable -dispatch contract. +Producers from before an upgrade may still enqueue this task. Its payload uses +the previous record format, so the message is dropped and its payload deleted; +only old workers can export it. Keep new workers off the old queue until it has +drained. The new pipeline's sweeper leaves legacy payloads alone. """ -import json import logging +import re from celery import shared_task -from celery.exceptions import Retry -from flask import current_app -from configs import dify_config -from core.ops.base_trace_instance import application_tracing_hooks -from core.ops.entities.config_entity import OPS_FILE_PATH, OPS_TRACE_FAILED_KEY -from core.ops.entities.trace_entity import trace_info_info_map -from core.ops.exceptions import RetryableTraceDispatchError -from core.rag.models.document import Document -from extensions.ext_redis import redis_client from extensions.ext_storage import storage -from models.model import Message -from models.workflow import WorkflowRun logger = logging.getLogger(__name__) -_RETRYABLE_TRACE_DISPATCH_LIMIT = dify_config.OPS_TRACE_RETRYABLE_DISPATCH_MAX_RETRIES -_RETRYABLE_TRACE_DISPATCH_DELAY_SECONDS = dify_config.OPS_TRACE_RETRYABLE_DISPATCH_DELAY_SECONDS - - -@shared_task( - queue="ops_trace", - bind=True, - max_retries=_RETRYABLE_TRACE_DISPATCH_LIMIT, - default_retry_delay=_RETRYABLE_TRACE_DISPATCH_DELAY_SECONDS, -) -def process_trace_tasks(self, file_info): - """ - Async process trace tasks - Usage: process_trace_tasks.delay(tasks_data) - """ - from core.ops.ops_trace_manager import OpsTraceManager - - app_id = file_info.get("app_id") - file_id = file_info.get("file_id") - file_path = f"{OPS_FILE_PATH}{app_id}/{file_id}.json" - file_data = json.loads(storage.load(file_path)) - trace_info = file_data.get("trace_info") - trace_info_type = file_data.get("trace_info_type") - enterprise_trace_dispatched = bool(file_data.get("_enterprise_trace_dispatched")) - - if trace_info.get("message_data"): - trace_info["message_data"] = Message.from_dict(data=trace_info["message_data"]) - if trace_info.get("workflow_data"): - trace_info["workflow_data"] = WorkflowRun.from_dict(data=trace_info["workflow_data"]) - if trace_info.get("documents"): - trace_info["documents"] = [Document.model_validate(doc) for doc in trace_info["documents"]] - - should_delete_file = True - - try: - trace_type = trace_info_info_map.get(trace_info_type) - if trace_type: - trace_info = trace_type(**trace_info) +_PAYLOAD_SEGMENT = re.compile(r"[A-Za-z0-9-]+") - application_trace_captured = False - if dify_config.OPS_AGENT_TRACE_ENABLED: - with current_app.app_context(): - tracing_hooks = application_tracing_hooks() - if tracing_hooks is not None: - application_trace_captured = tracing_hooks.capture_terminal(trace_info) - - from extensions.ext_enterprise_telemetry import is_enabled as is_ee_telemetry_enabled - - if is_ee_telemetry_enabled() and not enterprise_trace_dispatched: - from enterprise.telemetry.enterprise_trace import EnterpriseOtelTrace +@shared_task(queue="ops_trace") +def process_trace_tasks(file_info: dict[str, object]) -> None: + app_id, file_id = file_info.get("app_id"), file_info.get("file_id") + if isinstance(app_id, str) and isinstance(file_id, str): + if _PAYLOAD_SEGMENT.fullmatch(app_id) and _PAYLOAD_SEGMENT.fullmatch(file_id): try: - EnterpriseOtelTrace().trace(trace_info) + storage.delete(f"ops_trace/{app_id}/{file_id}.json") except Exception: - logger.exception("Enterprise trace failed for app_id: %s", app_id) - else: - file_data["_enterprise_trace_dispatched"] = True - enterprise_trace_dispatched = True - - if not application_trace_captured: - trace_instance = OpsTraceManager.get_ops_trace_instance(app_id) - if trace_instance: - with current_app.app_context(): - trace_instance.trace(trace_info) - - logger.info("Processing trace tasks success, app_id: %s", app_id) - except RetryableTraceDispatchError as e: - # Only confirmed retryable failures, including transient parent ordering - # gaps, may replay. Unknown acceptance must terminate without retry. - # Keep the payload only after - # Celery accepts the retry request; otherwise this attempt becomes a - # terminal failure and the stored file is cleaned up in `finally`. - # - # Enterprise telemetry runs before provider dispatch. If it already ran - # and provider dispatch asks for a retry, persist that private flag so - # the next attempt does not emit the same enterprise trace twice. - if e.acceptance_unknown: - # Provider exception details can contain private request content. - logger.exception("Trace acceptance unknown; automatic retry suppressed, app_id: %s", app_id, exc_info=False) - failed_key = f"{OPS_TRACE_FAILED_KEY}_{app_id}" - redis_client.incr(failed_key) - elif self.request.retries >= _RETRYABLE_TRACE_DISPATCH_LIMIT: - logger.exception("Retryable trace dispatch budget exhausted, app_id: %s", app_id) - failed_key = f"{OPS_TRACE_FAILED_KEY}_{app_id}" - redis_client.incr(failed_key) - else: - logger.warning( - "Retryable trace dispatch failure, scheduling retry %s/%s for app_id %s: %s", - self.request.retries + 1, - _RETRYABLE_TRACE_DISPATCH_LIMIT, - app_id, - e, - ) - try: - if enterprise_trace_dispatched: - storage.save(file_path, json.dumps(file_data).encode("utf-8")) - raise self.retry(exc=e, countdown=_RETRYABLE_TRACE_DISPATCH_DELAY_SECONDS) - except Retry: - should_delete_file = False - raise - except Exception: - logger.exception("Failed to schedule trace dispatch retry, app_id: %s", app_id) - failed_key = f"{OPS_TRACE_FAILED_KEY}_{app_id}" - redis_client.incr(failed_key) - except Exception: - logger.exception("Processing trace tasks failed, app_id: %s", app_id) - failed_key = f"{OPS_TRACE_FAILED_KEY}_{app_id}" - redis_client.incr(failed_key) - finally: - if should_delete_file: - try: - storage.delete(file_path) - except Exception as e: - logger.warning( - "Failed to delete trace file %s for app_id %s: %s", - file_path, - app_id, - e, - ) + logger.warning("Legacy trace payload cleanup failed: app_id=%s", app_id, exc_info=True) + logger.info("Dropped a trace task enqueued before the trace pipeline upgrade: app_id=%s", app_id) diff --git a/api/tasks/trace_delivery_task.py b/api/tasks/trace_delivery_task.py new file mode 100644 index 00000000000000..38dc69bf1fd989 --- /dev/null +++ b/api/tasks/trace_delivery_task.py @@ -0,0 +1,112 @@ +"""Celery entry points of the trace pipeline, routed by ``OPS_TRACE_QUEUE``. + +Delivery is at most once per sink: an event is redelivered only while its +records are not persisted yet, or when the Agent collector could not decide +ownership (enterprise telemetry is then skipped, it already ran). +""" + +import logging +import re +from datetime import UTC, date, datetime, timedelta +from typing import Any + +from celery import shared_task +from dify_trace_sdk.events import TraceEvent + +from extensions.ext_application_services import application_services +from extensions.ext_storage import storage +from services.tracing.adapters import PAYLOAD_PREFIX +from services.tracing.delivery import ExportJob +from services.tracing.errors import SourceNotFoundError, SourceNotReadyError, TraceCollectorUnavailableError + +logger = logging.getLogger(__name__) + +REDELIVERY_LIMIT = 10 +PAYLOAD_RETENTION = timedelta(days=2) +_PAYLOAD_DAY = re.compile(rf"{PAYLOAD_PREFIX}/(\d{{4}}-\d{{2}}-\d{{2}})/") + + +def _expired(key: str, cutoff: date) -> bool: + match = _PAYLOAD_DAY.match(key) + if match is None: + return False + try: + return date.fromisoformat(match.group(1)) < cutoff + except ValueError: + return False + + +def _redelivery_delay(retries: int) -> int: + return min(60, 5 * 2**retries) + + +def _delete_quietly(key: str) -> None: + try: + storage.delete(key) + except Exception: + logger.warning("Trace payload cleanup failed: key=%s", key, exc_info=True) + + +@shared_task(bind=True, max_retries=REDELIVERY_LIMIT) +def deliver_trace_event(self, envelope: dict[str, Any]) -> None: + payload_key = envelope.get("payload_key") + redelivering = False + try: + raw = storage.load(payload_key) if isinstance(payload_key, str) else envelope["event"] + event = TraceEvent.model_validate_json(raw) + except Exception: + logger.warning("Trace event unreadable; dropped", exc_info=True) + if isinstance(payload_key, str): + _delete_quietly(payload_key) + return + try: + application_services().tracing.delivery.deliver( + event, enterprise_delivered=envelope.get("enterprise_delivered") is True + ) + except SourceNotFoundError: + logger.info("Trace event source not found; dropped: kind=%s app_id=%s", event.payload.kind, event.app_id) + except SourceNotReadyError: + if self.request.retries < REDELIVERY_LIMIT: + redelivering = True + raise self.retry(args=(envelope,), countdown=_redelivery_delay(self.request.retries)) from None + logger.warning("Trace event source never became ready; dropped: kind=%s", event.payload.kind) + except TraceCollectorUnavailableError: + if self.request.retries < REDELIVERY_LIMIT: + redelivering = True + raise self.retry( + args=({**envelope, "enterprise_delivered": True},), + countdown=_redelivery_delay(self.request.retries), + ) from None + logger.warning("Agent trace collector unavailable; dropped: kind=%s", event.payload.kind) + finally: + if isinstance(payload_key, str) and not redelivering: + _delete_quietly(payload_key) + + +@shared_task +def retry_trace_export(job_key: str) -> None: + try: + job = ExportJob.model_validate_json(storage.load(job_key)) + except Exception: + logger.warning("Trace export job unreadable; dropped: key=%s", job_key, exc_info=True) + _delete_quietly(job_key) + return + try: + application_services().tracing.delivery.retry(job) + finally: + _delete_quietly(job_key) + + +@shared_task +def sweep_trace_payloads() -> None: + """Delete expired new-pipeline payloads; legacy workers still own their payloads during an upgrade.""" + cutoff = (datetime.now(UTC) - PAYLOAD_RETENTION).date() + # Directory listings need the trailing slash on some storage backends. + prefix = f"{PAYLOAD_PREFIX}/" + try: + keys = storage.scan(prefix, files=True) + except (FileNotFoundError, NotImplementedError): + return + for key in keys: + if key.startswith(prefix) and _expired(key, cutoff): + _delete_quietly(key) diff --git a/api/tests/e2e_tests/containers/core/workflow/test_human_input_resume_node_execution.py b/api/tests/e2e_tests/containers/core/workflow/test_human_input_resume_node_execution.py index 926c860df5d481..fad94dd8a46dac 100644 --- a/api/tests/e2e_tests/containers/core/workflow/test_human_input_resume_node_execution.py +++ b/api/tests/e2e_tests/containers/core/workflow/test_human_input_resume_node_execution.py @@ -278,6 +278,7 @@ def _build_run_layers( ), workflow_execution_repository=execution_repo, workflow_node_execution_repository=node_execution_repo, + trace_scope=None, ) persistence.set_node_run_indices(run_index.indices) persistence.set_node_execution_history(history) diff --git a/api/tests/e2e_tests/containers/repositories/test_sqlalchemy_api_workflow_node_execution_repository.py b/api/tests/e2e_tests/containers/repositories/test_sqlalchemy_api_workflow_node_execution_repository.py index d35153beebc616..0d5f0625dbe7ad 100644 --- a/api/tests/e2e_tests/containers/repositories/test_sqlalchemy_api_workflow_node_execution_repository.py +++ b/api/tests/e2e_tests/containers/repositories/test_sqlalchemy_api_workflow_node_execution_repository.py @@ -9,10 +9,13 @@ from sqlalchemy import Engine, delete from sqlalchemy.orm import Session, sessionmaker +from extensions.ext_storage import storage +from extensions.storage.storage_type import StorageType from graphon.enums import WorkflowNodeExecutionStatus from libs.datetime_utils import naive_utc_now -from models.enums import CreatorUserRole -from models.workflow import WorkflowNodeExecutionModel +from models.enums import CreatorUserRole, ExecutionOffLoadType +from models.model import UploadFile +from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionOffload, WorkflowNodeExecutionTriggeredFrom from repositories.sqlalchemy_api_workflow_node_execution_repository import ( DifyAPISQLAlchemyWorkflowNodeExecutionRepository, ) @@ -36,7 +39,7 @@ def _create_node_execution( tenant_id=tenant_id, app_id=app_id, workflow_id=workflow_id, - triggered_from="workflow-run", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, workflow_run_id=workflow_run_id, index=index, predecessor_node_id=None, @@ -76,13 +79,39 @@ def test_load_full_process_data_reads_offload( db_session_with_containers: Session, monkeypatch: pytest.MonkeyPatch, ) -> None: - execution = WorkflowNodeExecutionModel(process_data="{}") - execution.offload_data = [] - monkeypatch.setattr( - WorkflowNodeExecutionModel, - "load_full_process_data", - lambda _self, _session, _storage: {"__dify_retry_history": [{"retry_index": 1}]}, + tenant_id = str(uuid4()) + file = UploadFile( + tenant_id=tenant_id, + storage_type=StorageType.LOCAL, + key=f"offloaded-process-data/{uuid4()}.json", + name="process-data.json", + size=45, + extension="json", + mime_type="application/json", + created_by=str(uuid4()), + created_by_role=CreatorUserRole.ACCOUNT, + created_at=naive_utc_now(), + used=True, ) + file_id, file_key = file.id, file.key + db_session_with_containers.add(file) + db_session_with_containers.commit() + execution = WorkflowNodeExecutionModel(tenant_id=tenant_id, process_data="{}") + execution.offload_data = [ + WorkflowNodeExecutionOffload( + tenant_id=tenant_id, + app_id=str(uuid4()), + node_execution_id=str(uuid4()), + type_=ExecutionOffLoadType.PROCESS_DATA, + file_id=file_id, + ) + ] + + def load_storage(key: str) -> bytes: + assert key == file_key + return b'{"__dify_retry_history": [{"retry_index": 1}]}' + + monkeypatch.setattr(storage, "load", load_storage) engine = db_session_with_containers.get_bind() assert isinstance(engine, Engine) repository = DifyAPISQLAlchemyWorkflowNodeExecutionRepository(sessionmaker(bind=engine, expire_on_commit=False)) diff --git a/api/tests/e2e_tests/containers/services/test_message_suggested_questions_failures.py b/api/tests/e2e_tests/containers/services/test_message_suggested_questions_failures.py index ac9b2203562ef5..e7bb3834d7e7df 100644 --- a/api/tests/e2e_tests/containers/services/test_message_suggested_questions_failures.py +++ b/api/tests/e2e_tests/containers/services/test_message_suggested_questions_failures.py @@ -25,7 +25,6 @@ from core.helper.model_provider_cache import ProviderCredentialsCache, ProviderCredentialsCacheType from core.model_context import get_credit_usage_metadata from core.model_manager import ModelManager -from core.ops import ops_trace_manager from core.plugin.entities.plugin_daemon import PluginModelProviderDeclaration from core.plugin.impl import base as plugin_base from core.plugin.impl.model_runtime_factory import create_plugin_model_runtime @@ -47,6 +46,7 @@ create_console_account_and_tenant, create_console_app, ) +from tests.helpers.unit.services.tracing_ports import untraced_scopes _PROVIDER = "tests/suggested/suggested" _MODEL = "questions-model" @@ -59,7 +59,7 @@ class _Scenario: @pytest.fixture -def scenario(db_session_with_containers: Session, monkeypatch: pytest.MonkeyPatch) -> Iterator[_Scenario]: +def scenario(db_session_with_containers: Session, monkeypatch: pytest.MonkeyPatch) -> _Scenario: monkeypatch.setattr(dify_config, "PLUGIN_MODEL_PROVIDERS_CACHE_ENABLED", True) session = db_session_with_containers account, tenant = create_console_account_and_tenant(session) @@ -139,15 +139,7 @@ def scenario(db_session_with_containers: Session, monkeypatch: pytest.MonkeyPatc # The container fixture reloads config after this module-level URL was bound. monkeypatch.setattr(plugin_base, "plugin_daemon_inner_api_baseurl", URL(str(dify_config.PLUGIN_DAEMON_URL))) - previous_timer = ops_trace_manager.trace_manager_timer - try: - yield _Scenario(account=account, context=context) - finally: - timer = ops_trace_manager.trace_manager_timer - if timer is not None and timer is not previous_timer: - timer.cancel() - timer.join(timeout=1) - ops_trace_manager.trace_manager_timer = previous_timer + return _Scenario(account=account, context=context) @pytest.fixture @@ -287,7 +279,7 @@ def test_unavailable_configured_and_default_models_return_none(scenario: _Scenar .values(provider_name="tests/missing/missing") ) db.session.commit() - generator = SuggestedQuestionsGenerator() + generator = SuggestedQuestionsGenerator(trace_scopes=untraced_scopes()) assert ( generator._prepare_model( @@ -309,7 +301,7 @@ def test_missing_default_model_does_not_create_one_or_commit_caller_changes(scen caller = db.session() try: - with SuggestedQuestionsGenerator().prepare( + with SuggestedQuestionsGenerator(trace_scopes=untraced_scopes()).prepare( context=scenario.context, instruction_prompt=None, model_config=None ) as generate: assert generate is None diff --git a/api/tests/e2e_tests/contracts/configs/test_docker_env_sync.py b/api/tests/e2e_tests/contracts/configs/test_docker_env_sync.py new file mode 100644 index 00000000000000..5988004e8af694 --- /dev/null +++ b/api/tests/e2e_tests/contracts/configs/test_docker_env_sync.py @@ -0,0 +1,56 @@ +"""Upgrades must preserve overrides omitted from the minimal Docker template.""" + +import subprocess +import sys +from pathlib import Path + +import pytest +from dotenv import dotenv_values + +_DOCKER = Path(__file__).resolve().parents[5] / "docker" + + +@pytest.mark.parametrize("script", ["dify-env-sync.py", "dify-env-sync.sh"]) +@pytest.mark.parametrize( + "existing", + [ + None, + "", + "# Local settings\nCOMPOSE_PROFILES=mysql,qdrant\n" + "OPS_TRACE_QUEUE=custom_trace_queue\nREDIS_PASSWORD='space # $ = value'\n" + "MULTILINE='first line\nsecond=line'\nEMPTY=\nLITERAL='$(touch unexpected)'", + ], +) +def test_sync_preserves_local_overrides(tmp_path: Path, script: str, existing: str | None) -> None: + (tmp_path / ".env.example").write_text( + "# Deployment inputs\nCOMPOSE_PROFILES=postgresql,weaviate\nNEW_INPUT=value\n" + ) + env_file = tmp_path / ".env" + if existing is not None: + env_file.write_text(existing) + env_file.chmod(0o600) + + command = [sys.executable if script.endswith(".py") else "bash", str(_DOCKER / script)] + result = subprocess.run(command, cwd=tmp_path, capture_output=True, text=True, check=False) + assert result.returncode == 0, result.stderr + updated = env_file.read_text() + if existing is not None: + assert updated.startswith(existing) + assert env_file.stat().st_mode & 0o777 == 0o600 + assert next((tmp_path / "env-backup").iterdir()).read_text() == existing + values = dotenv_values(env_file, interpolate=False) + assert values["NEW_INPUT"] == "value" + if existing: + assert values["COMPOSE_PROFILES"] == "mysql,qdrant" + assert values["OPS_TRACE_QUEUE"] == "custom_trace_queue" + assert values["REDIS_PASSWORD"] == "space # $ = value" + assert values["MULTILINE"] == "first line\nsecond=line" + assert values["EMPTY"] == "" + assert values["LITERAL"] == "$(touch unexpected)" + else: + assert values["COMPOSE_PROFILES"] == "postgresql,weaviate" + assert not (tmp_path / "unexpected").exists() + + result = subprocess.run(command, cwd=tmp_path, capture_output=True, text=True, check=False) + assert result.returncode == 0, result.stderr + assert env_file.read_text() == updated diff --git a/api/tests/e2e_tests/contracts/configs/test_knowledge_fs_worker_deployment.py b/api/tests/e2e_tests/contracts/configs/test_knowledge_fs_worker_deployment.py index 210805636b7269..e5347f7a36507a 100644 --- a/api/tests/e2e_tests/contracts/configs/test_knowledge_fs_worker_deployment.py +++ b/api/tests/e2e_tests/contracts/configs/test_knowledge_fs_worker_deployment.py @@ -1,26 +1,13 @@ """Check the shipped Compose consumer wiring against the real entrypoint queue selection.""" -import subprocess from pathlib import Path import pytest import yaml -_ROOT = Path(__file__).resolve().parents[5] - +from tests.helpers.unit.worker_entrypoint import worker_configuration -def _worker_configuration(environment: dict[str, str]) -> subprocess.CompletedProcess[str]: - script = (_ROOT / "api/docker/entrypoint.sh").read_text() - # Run the queue/pool selection and its admission guards, without starting a worker. - start = script.index(" # Configure queues") - end = script.index(' echo "Starting Celery worker with queues:') - return subprocess.run( - ["bash", "-c", script[start:end] + '\nprintf "%s\\n" "$DEFAULT_QUEUES"'], - env={"PATH": "/usr/bin:/bin", **environment}, - capture_output=True, - text=True, - check=False, - ) +_ROOT = Path(__file__).resolve().parents[5] @pytest.mark.parametrize("edition", ["COMMUNITY", "ENTERPRISE", "CLOUD"]) @@ -28,7 +15,7 @@ def _worker_configuration(environment: dict[str, str]) -> subprocess.CompletedPr def test_compose_worker_consumes_provisioning_import_and_cleanup_queue(edition: str, filename: str) -> None: compose = yaml.safe_load((_ROOT / "docker" / filename).read_text()) environment = compose["services"]["worker"]["environment"] - selected = _worker_configuration({**environment, "DEPLOYMENT_EDITION": edition}) + selected = worker_configuration({**environment, "DEPLOYMENT_EDITION": edition}) assert selected.returncode == 0, selected.stderr assert "knowledge_fs_lifecycle" in selected.stdout.splitlines()[-1].split(",") @@ -39,15 +26,15 @@ def test_upgrade_worker_enforces_prefork_for_checkpoint_recovery(filename: str) worker = compose["services"]["worker_knowledge_fs_upgrade"] assert worker["profiles"] == ["knowledge-fs"] environment = worker["environment"] - selected = _worker_configuration({key: str(value) for key, value in environment.items()}) + selected = worker_configuration({key: str(value) for key, value in environment.items()}) assert selected.returncode == 0, selected.stderr assert selected.stdout.splitlines()[-1] == "knowledge_fs_upgrade" - rejected = _worker_configuration({"CELERY_WORKER_QUEUES": "knowledge_fs_upgrade", "CELERY_WORKER_POOL": "gevent"}) + rejected = worker_configuration({"CELERY_WORKER_QUEUES": "knowledge_fs_upgrade", "CELERY_WORKER_POOL": "gevent"}) assert rejected.returncode != 0 assert "require CELERY_WORKER_POOL=prefork" in rejected.stderr def test_explicit_operator_queues_remain_authoritative() -> None: - selected = _worker_configuration({"CELERY_WORKER_QUEUES": "dataset,mail"}) + selected = worker_configuration({"CELERY_WORKER_QUEUES": "dataset,mail"}) assert selected.returncode == 0, selected.stderr assert selected.stdout.splitlines()[-1] == "dataset,mail" diff --git a/api/tests/e2e_tests/contracts/controllers/console/app/test_message_suggested_questions.py b/api/tests/e2e_tests/contracts/controllers/console/app/test_message_suggested_questions.py index 87b21be992eecf..39ca74dffcbb53 100644 --- a/api/tests/e2e_tests/contracts/controllers/console/app/test_message_suggested_questions.py +++ b/api/tests/e2e_tests/contracts/controllers/console/app/test_message_suggested_questions.py @@ -8,6 +8,7 @@ from uuid import uuid4 import pytest +from dify_trace_sdk.testing import RecordingEgress from flask import Flask from sqlalchemy import Connection, Engine, delete, event, func, select, text, update from sqlalchemy.orm import Session, SessionTransaction, sessionmaker @@ -47,7 +48,6 @@ from repositories.recommended_app_catalog_repository import DatabaseRecommendedAppCatalogRepository from services.agent.roster_package_exporter import RosterAgentPackageExporter from services.app.agent_app_service import AgentAppAccessService -from services.app_tracing_config_gateway import OpsTraceManagerGateway from services.app_tracing_config_service import AppTracingConfigService from services.entities.message_entities import MessageAccount from services.message_suggested_questions_generator import SuggestedQuestionsGenerator @@ -60,6 +60,7 @@ from tests.helpers.unit.config_override import apply_config_overrides from tests.helpers.unit.model_factories import make_account, make_app, make_conversation, make_message from tests.helpers.unit.rbac_fakes import RBACDomain +from tests.helpers.unit.services.tracing_ports import deployment_provider_configs, untraced_scopes @dataclass @@ -187,7 +188,7 @@ def track_session(session: Session, _transaction: SessionTransaction, _connectio rbac_members=rbac_domain.rbac.members, tracing_configs=AppTracingConfigService( configs=SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory), - provider=OpsTraceManagerGateway(), + provider=deployment_provider_configs(RecordingEgress([], grpc_target=None)), ), oauth=_build_oauth_server_service(database_client=sqlite_session_factory, redis=redis), recommended_packages=RecommendedAppPackageService( @@ -204,7 +205,9 @@ def track_session(session: Session, _transaction: SessionTransaction, _connectio services = _Services( app_services, _Agents(AgentAppAccessService(references=AgentAppRepository(session_factory=sqlite_session_factory))), - MessageSuggestedQuestionsService(queries=queries, generator=SuggestedQuestionsGenerator()), + MessageSuggestedQuestionsService( + queries=queries, generator=SuggestedQuestionsGenerator(trace_scopes=untraced_scopes()) + ), rbac_domain.rbac, ) flask_app.extensions["application_services"] = services diff --git a/api/tests/e2e_tests/contracts/controllers/console/explore/test_trial_suggested_questions.py b/api/tests/e2e_tests/contracts/controllers/console/explore/test_trial_suggested_questions.py index 0c2acd7cd90899..fd2aea04f605d4 100644 --- a/api/tests/e2e_tests/contracts/controllers/console/explore/test_trial_suggested_questions.py +++ b/api/tests/e2e_tests/contracts/controllers/console/explore/test_trial_suggested_questions.py @@ -39,6 +39,7 @@ ) from services.recommended_app_query_service import RecommendedAppQueryService from services.trial_app_access_service import TrialAppAccessService +from tests.helpers.unit.services.tracing_ports import untraced_scopes @dataclass(frozen=True) @@ -185,7 +186,7 @@ def track_read(session: Session, _transaction: SessionTransaction, _connection: trial_apps=_TrialAppServices(access=TrialAppAccessService(apps=trial_apps)), recommended_app_queries=features, message_suggested_questions=MessageSuggestedQuestionsService( - queries=queries, generator=SuggestedQuestionsGenerator() + queries=queries, generator=SuggestedQuestionsGenerator(trace_scopes=untraced_scopes()) ), ) diff --git a/api/tests/e2e_tests/contracts/controllers/service_api/app/test_conversation_rename.py b/api/tests/e2e_tests/contracts/controllers/service_api/app/test_conversation_rename.py index e38470ef8c5aa5..642f523a08dd93 100644 --- a/api/tests/e2e_tests/contracts/controllers/service_api/app/test_conversation_rename.py +++ b/api/tests/e2e_tests/contracts/controllers/service_api/app/test_conversation_rename.py @@ -31,10 +31,11 @@ from services.app_definition_query_service import AppDefinitionQueryService from services.app_scoped_end_user_service import AppScopedEndUserService from services.auth.resource_access_token_contracts import ResourceAccessTokenResource -from services.conversation_name_generator import generate_conversation_name +from services.conversation_name_generator import DefaultModelConversationNameGenerator from services.conversation_name_service import ConversationNameService from services.resource_access_token_service import ResourceAccessTokenService from tests.helpers.unit.model_factories import make_app, make_conversation, make_end_user, make_message +from tests.helpers.unit.services.tracing_ports import untraced_scopes _CREATED_AT = datetime(2026, 1, 2, 3, 4, 5) _UPDATED_AT = datetime(2026, 1, 3, 3, 4, 5) @@ -174,7 +175,7 @@ def track_scoped(session: Session, _transaction: SessionTransaction, _connection ), conversation_names=ConversationNameService( conversations=ConversationRepository(session_factory=query_factory), - generate_name=generate_conversation_name, + generate_name=DefaultModelConversationNameGenerator(trace_scopes=untraced_scopes()), ), resource_access_tokens=build_resource_access_token_service(database_client=admission_factory), ) diff --git a/api/tests/e2e_tests/contracts/controllers/service_api/app/test_message_suggested_questions.py b/api/tests/e2e_tests/contracts/controllers/service_api/app/test_message_suggested_questions.py index 97924a109443b3..01bec883666ade 100644 --- a/api/tests/e2e_tests/contracts/controllers/service_api/app/test_message_suggested_questions.py +++ b/api/tests/e2e_tests/contracts/controllers/service_api/app/test_message_suggested_questions.py @@ -45,6 +45,7 @@ ) from services.resource_access_token_service import ResourceAccessTokenService from tests.helpers.unit.model_factories import make_app, make_conversation, make_end_user, make_message +from tests.helpers.unit.services.tracing_ports import untraced_scopes @dataclass(frozen=True) @@ -186,7 +187,9 @@ def track_scoped(session: Session, _transaction: SessionTransaction, _connection queries = SuggestedQuestionsQuery( session_factory=query_factory, repository=MessageRepository(session_factory=query_factory) ) - service = MessageSuggestedQuestionsService(queries=queries, generator=SuggestedQuestionsGenerator()) + service = MessageSuggestedQuestionsService( + queries=queries, generator=SuggestedQuestionsGenerator(trace_scopes=untraced_scopes()) + ) app.extensions["application_services"] = _Services( app_definitions=AppDefinitionQueryService( definitions=AppDefinitionQueryRepository(session_factory=admission_factory), builtin_icon_url_prefix="" diff --git a/api/tests/e2e_tests/contracts/controllers/web/test_conversation_rename.py b/api/tests/e2e_tests/contracts/controllers/web/test_conversation_rename.py index fd097de70d0bd6..cb53926a9ab22a 100644 --- a/api/tests/e2e_tests/contracts/controllers/web/test_conversation_rename.py +++ b/api/tests/e2e_tests/contracts/controllers/web/test_conversation_rename.py @@ -36,9 +36,10 @@ from models.model import App, AppMode, AppModelConfig, Conversation, EndUser, Message, Site from models.workflow import WorkflowType from repositories.conversation_repository import ConversationRepository -from services.conversation_name_generator import generate_conversation_name +from services.conversation_name_generator import DefaultModelConversationNameGenerator from services.conversation_name_service import ConversationNameService from tests.helpers.unit.model_factories import make_app, make_conversation, make_end_user, make_message, make_workflow +from tests.helpers.unit.services.tracing_ports import untraced_scopes _CREATED_AT = datetime(2026, 1, 2, 3, 4, 5) _UPDATED_AT = datetime(2026, 1, 3, 3, 4, 5) @@ -157,7 +158,7 @@ def track_provider(session: Session, _transaction: SessionTransaction, _connecti app.extensions["application_services"] = _Services( ConversationNameService( conversations=ConversationRepository(session_factory=repository_factory), - generate_name=generate_conversation_name, + generate_name=DefaultModelConversationNameGenerator(trace_scopes=untraced_scopes()), ) ) api = ExternalApi(app) diff --git a/api/tests/e2e_tests/contracts/controllers/web/test_message_suggested_questions.py b/api/tests/e2e_tests/contracts/controllers/web/test_message_suggested_questions.py index a8a4ae88aa11e1..2c2246b91549d3 100644 --- a/api/tests/e2e_tests/contracts/controllers/web/test_message_suggested_questions.py +++ b/api/tests/e2e_tests/contracts/controllers/web/test_message_suggested_questions.py @@ -42,6 +42,7 @@ MessageSuggestedQuestionsService, ) from tests.helpers.unit.model_factories import make_app, make_conversation, make_end_user, make_message +from tests.helpers.unit.services.tracing_ports import untraced_scopes @dataclass @@ -157,7 +158,11 @@ def track_admission(session: Session, transaction: SessionTransaction, connectio queries = SuggestedQuestionsQuery( session_factory=read_factory, repository=MessageRepository(session_factory=read_factory) ) - services = _Services(MessageSuggestedQuestionsService(queries=queries, generator=SuggestedQuestionsGenerator())) + services = _Services( + MessageSuggestedQuestionsService( + queries=queries, generator=SuggestedQuestionsGenerator(trace_scopes=untraced_scopes()) + ) + ) app.extensions["application_services"] = services api = ExternalApi(app) api.add_resource(message_controller.MessageSuggestedQuestionApi, "/messages//suggested-questions") diff --git a/api/tests/e2e_tests/contracts/core/app/apps/completion/test_prepared_completion.py b/api/tests/e2e_tests/contracts/core/app/apps/completion/test_prepared_completion.py index 6ffc67e61bce0d..b721f1b2e760cc 100644 --- a/api/tests/e2e_tests/contracts/core/app/apps/completion/test_prepared_completion.py +++ b/api/tests/e2e_tests/contracts/core/app/apps/completion/test_prepared_completion.py @@ -75,7 +75,7 @@ def test_prepared_records_preserve_input_and_actor_ownership( assert saved.from_end_user_id is None assert verification.scalar(select(func.count()).select_from(Message)) == 1 assert verification.scalar(select(func.count()).select_from(Conversation)) == 1 - assert entity.trace_manager is None + assert entity.trace_scope is None assert entity.file_upload_config is None diff --git a/api/tests/e2e_tests/contracts/core/workflow/nodes/tool/test_tool_node.py b/api/tests/e2e_tests/contracts/core/workflow/nodes/tool/test_tool_node.py index c8ea5a542f5d84..eb55a3654f8311 100644 --- a/api/tests/e2e_tests/contracts/core/workflow/nodes/tool/test_tool_node.py +++ b/api/tests/e2e_tests/contracts/core/workflow/nodes/tool/test_tool_node.py @@ -1,10 +1,8 @@ from __future__ import annotations -import sys -import types from collections.abc import Generator from types import SimpleNamespace -from typing import TYPE_CHECKING, Any +from typing import Any from unittest.mock import MagicMock import pytest @@ -17,14 +15,12 @@ from graphon.model_runtime.entities.llm_entities import LLMUsage from graphon.node_events import StreamChunkEvent, StreamCompletedEvent from graphon.nodes.tool.entities import ToolNodeData +from graphon.nodes.tool.tool_node import ToolNode from graphon.nodes.tool_runtime_entities import ToolRuntimeHandle, ToolRuntimeMessage from graphon.runtime import RuntimeState from graphon.variables.segments import ArrayFileSegment from tests.workflow_test_utils import build_test_graph_init_params, build_test_variable_pool -if TYPE_CHECKING: # pragma: no cover - imported for type checking only - from graphon.nodes.tool.tool_node import ToolNode - class _StubToolRuntime: def get_runtime( @@ -66,16 +62,7 @@ def resolve_provider_icons( @pytest.fixture -def tool_node(monkeypatch) -> ToolNode: - module_name = "core.ops.ops_trace_manager" - if module_name not in sys.modules: - ops_stub = types.ModuleType(module_name) - ops_stub.TraceQueueManager = object # pragma: no cover - stub attribute - ops_stub.TraceTask = object # pragma: no cover - stub attribute - monkeypatch.setitem(sys.modules, module_name, ops_stub) - - from graphon.nodes.tool.tool_node import ToolNode - +def tool_node() -> ToolNode: graph_config: dict[str, Any] = { "nodes": [ { diff --git a/api/tests/helpers/unit/core/ops/trace_fixtures.py b/api/tests/helpers/unit/core/ops/trace_fixtures.py deleted file mode 100644 index d64b0ce80a1f12..00000000000000 --- a/api/tests/helpers/unit/core/ops/trace_fixtures.py +++ /dev/null @@ -1,60 +0,0 @@ -"""Validated trace records with baseline values for provider unit tests.""" - -from datetime import datetime - -from core.ops.entities.trace_entity import MessageTraceInfo, ToolTraceInfo, WorkflowTraceInfo -from graphon.entities import WorkflowNodeExecution -from graphon.enums import BuiltinNodeTypes - - -def workflow_trace_info() -> WorkflowTraceInfo: - return WorkflowTraceInfo( - 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={}, - ) - - -def message_trace_info() -> MessageTraceInfo: - return MessageTraceInfo( - message_id="message-1", - conversation_model="gpt-4", - conversation_mode="chat", - message_tokens=0, - answer_tokens=0, - total_tokens=0, - metadata={}, - ) - - -def tool_trace_info() -> ToolTraceInfo: - return ToolTraceInfo( - tool_name="search", - tool_inputs={}, - tool_outputs="", - tool_config={}, - tool_parameters={}, - time_cost=0, - metadata={}, - ) - - -def workflow_node_execution() -> WorkflowNodeExecution: - return WorkflowNodeExecution( - id="execution-1", - workflow_id="workflow-1", - index=1, - node_id="node-1", - node_type=BuiltinNodeTypes.LLM, - title="LLM", - created_at=datetime(2026, 1, 1), - ) diff --git a/api/tests/helpers/unit/services/conversation_name_fixtures.py b/api/tests/helpers/unit/services/conversation_name_fixtures.py index fe1e2dc646e8f6..be8188e4fe1e44 100644 --- a/api/tests/helpers/unit/services/conversation_name_fixtures.py +++ b/api/tests/helpers/unit/services/conversation_name_fixtures.py @@ -11,7 +11,6 @@ from sqlalchemy import Engine from yarl import URL -from core.ops import ops_trace_manager from core.plugin.impl import base as plugin_http from extensions.ext_database import db @@ -98,23 +97,14 @@ def respond(self, data: object, *, stream: bool, code: int = 0, message: str = " server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) thread = Thread(target=server.serve_forever, daemon=True) original_url = plugin_http.plugin_daemon_inner_api_baseurl - original_trace_interval = ops_trace_manager.trace_manager_interval - original_trace_timer = ops_trace_manager.trace_manager_timer # Change real connection settings, preserving every production client and method. plugin_http.plugin_daemon_inner_api_baseurl = URL(f"http://127.0.0.1:{server.server_port}") config_overrides(PLUGIN_MODEL_PROVIDERS_CACHE_ENABLED=False) - ops_trace_manager.trace_manager_interval = 3600 thread.start() try: yield plugin finally: plugin_http.plugin_daemon_inner_api_baseurl = original_url - ops_trace_manager.trace_manager_interval = original_trace_interval - timer = ops_trace_manager.trace_manager_timer - if timer is not None and timer is not original_trace_timer: - timer.cancel() - timer.join(timeout=1) - ops_trace_manager.trace_manager_timer = original_trace_timer server.shutdown() server.server_close() thread.join(timeout=1) diff --git a/api/tests/helpers/unit/services/tracing_ports.py b/api/tests/helpers/unit/services/tracing_ports.py new file mode 100644 index 00000000000000..8baf46415fc2ba --- /dev/null +++ b/api/tests/helpers/unit/services/tracing_ports.py @@ -0,0 +1,175 @@ +"""In-memory tracing ports for tests; every other collaborator stays real.""" + +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any, Literal, override + +from dify_agent.layers.execution_context import ApplicationTracingContext +from dify_trace_sdk.egress import Egress +from dify_trace_sdk.emitter import TraceEmitter +from dify_trace_sdk.events import TraceEvent +from dify_trace_sdk.provider import TraceProvider + +from core.helper.encrypter import obfuscated_token +from extensions.tracing import ApplicationTracingHooks +from services.app_tracing_config_service import AppTraceSettings +from services.tracing.adapters import StoredPolicyTraceScopes, WorkspaceTracingSecretCipher +from services.tracing.collection import AgentTraceSettings +from services.tracing.errors import TraceProviderUnavailableError, UnknownTraceProviderError +from services.tracing.provider_configs import ProviderConfigGateway +from services.tracing.providers import ( + SUPPORTED_PROVIDERS, + ProviderConfigCodec, + TraceProviderCatalog, + TracingSecretCipher, +) +from services.tracing.registry import EntryPointTraceProviderCatalog +from services.tracing.scopes import TraceScopes, TracingPolicyQuery + +if TYPE_CHECKING: + from core.app.apps.agent_app.runtime_request_builder import AgentAppRuntimeBuildContext + from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, WorkflowAppGenerateEntity + from core.workflow.nodes.agent_v2.runtime_request_builder import WorkflowAgentRuntimeBuildContext + +AGENT_TRACE_SETTINGS = AgentTraceSettings( + enabled=True, + max_spans=4096, + max_bytes=16_777_216, + collection_timeout_seconds=180, + export_lease_seconds=120, + max_export_attempts=8, + retention_days=7, +) +"""The deployment defaults with Agent trace collection enabled.""" + + +class InstalledTraceProviders(TraceProviderCatalog): + """Only the given provider classes are installed; other supported IDs are not.""" + + def __init__(self, *providers: type[TraceProvider[Any]]) -> None: + self._providers = {provider.id: provider for provider in providers} + + @override + def resolve(self, provider_id: str) -> type[TraceProvider[Any]]: + provider = self._providers.get(provider_id) + if provider is not None: + return provider + if provider_id in SUPPORTED_PROVIDERS: + raise TraceProviderUnavailableError(provider_id) + raise UnknownTraceProviderError(provider_id) + + +class WorkspaceBoundCipher(TracingSecretCipher): + """Reversible ``enc::`` tokens that only their own workspace decrypts. + + ``decrypted`` lists every token decrypted, so tests can prove when credentials were not read. + """ + + def __init__(self) -> None: + self.decrypted: list[str] = [] + + @override + def encrypt(self, tenant_id: str, value: str) -> str: + return f"enc:{tenant_id}:{value}" + + @override + def decrypt(self, tenant_id: str, tokens: Sequence[str]) -> list[str]: + prefix = f"enc:{tenant_id}:" + if not all(token.startswith(prefix) for token in tokens): + raise ValueError("token was encrypted for another workspace") + self.decrypted.extend(tokens) + return [token.removeprefix(prefix) for token in tokens] + + @override + def mask(self, value: str) -> str: + return obfuscated_token(value) + + +def deployment_provider_configs(egress: Egress) -> ProviderConfigGateway: + """The deployment's provider configuration gateway, reaching destinations only through ``egress``.""" + return ProviderConfigGateway( + providers=EntryPointTraceProviderCatalog(), + codec=ProviderConfigCodec(WorkspaceTracingSecretCipher()), + egress=egress, + ) + + +class NoTracingPolicies(TracingPolicyQuery): + """Producers that hold their app's policy never look it up.""" + + @override + def policy(self, *, tenant_id: str, app_id: str) -> AppTraceSettings | None: + raise AssertionError(f"unexpected tracing policy lookup: app_id={app_id}") + + +def kernel_trace_scopes( + emitter: TraceEmitter, *, policies: TracingPolicyQuery, enterprise_enabled: bool +) -> StoredPolicyTraceScopes: + """Trace scopes as the composition root binds them for legacy producers.""" + return StoredPolicyTraceScopes( + TraceScopes(emitter=emitter, policies=policies, enterprise_enabled=lambda: enterprise_enabled) + ) + + +class RecordingTracingHooks(ApplicationTracingHooks): + """Agent collection hooks that record every call and answer each Agent run with ``tracing``.""" + + def __init__(self, tracing: ApplicationTracingContext | None) -> None: + self.tracing = tracing + self.agent_apps: list[AgentAppRuntimeBuildContext] = [] + self.workflow_agents: list[tuple[WorkflowAgentRuntimeBuildContext, str | None]] = [] + self.frozen_runs: list[tuple[str, str, object]] = [] + """(workflow_id, workflow_run_id, conversation_id) of each frozen run.""" + self.terminals: list[tuple[str, str, str]] = [] + """(app_id, message_id, outcome) of each captured Agent app terminal.""" + + @override + def prepare_agent_app(self, context: "AgentAppRuntimeBuildContext") -> ApplicationTracingContext | None: + self.agent_apps.append(context) + return self.tracing + + @override + def prepare_workflow_agent( + self, context: "WorkflowAgentRuntimeBuildContext", conversation_id: str | None + ) -> ApplicationTracingContext | None: + self.workflow_agents.append((context, conversation_id)) + return self.tracing + + @override + def freeze_workflow_run( + self, + entity: "AdvancedChatAppGenerateEntity | WorkflowAppGenerateEntity", + *, + workflow_id: str, + workflow_run_id: str, + conversation_id: object, + graph_data: object, + ) -> None: + self.frozen_runs.append((workflow_id, workflow_run_id, conversation_id)) + + @override + def capture_agent_app_terminal( + self, + *, + tenant_id: str, + app_id: str, + message_id: str, + outcome: Literal["succeeded", "failed", "cancelled"], + output: str | None, + ) -> None: + self.terminals.append((app_id, message_id, outcome)) + + +class _NoStoredPolicies(TracingPolicyQuery): + @override + def policy(self, *, tenant_id: str, app_id: str) -> AppTraceSettings | None: + return None + + +class _NoConsumer: + def emit(self, event: TraceEvent) -> None: + raise AssertionError(f"no scope exists, so nothing emits: kind={event.payload.kind}") + + +def untraced_scopes() -> TraceScopes: + """The scopes of a deployment where no app traces and enterprise telemetry is off.""" + return TraceScopes(emitter=_NoConsumer(), policies=_NoStoredPolicies(), enterprise_enabled=lambda: False) diff --git a/api/tests/helpers/unit/worker_entrypoint.py b/api/tests/helpers/unit/worker_entrypoint.py new file mode 100644 index 00000000000000..d33b47e20eb66a --- /dev/null +++ b/api/tests/helpers/unit/worker_entrypoint.py @@ -0,0 +1,20 @@ +"""Run the shipped worker queue selection without starting a worker.""" + +import subprocess +from pathlib import Path + +_ENTRYPOINT = Path(__file__).resolve().parents[3] / "docker" / "entrypoint.sh" + + +def worker_configuration(environment: dict[str, str]) -> subprocess.CompletedProcess[str]: + script = _ENTRYPOINT.read_text() + # Run the queue/pool selection and its admission guards, without starting a worker. + start = script.index(" # Configure queues") + end = script.index(' echo "Starting Celery worker with queues:') + return subprocess.run( + ["bash", "-c", script[start:end] + '\nprintf "%s\\n" "$DEFAULT_QUEUES"'], + env={"PATH": "/usr/bin:/bin", **environment}, + capture_output=True, + text=True, + check=False, + ) diff --git a/api/tests/unit_tests/.ruff.toml b/api/tests/unit_tests/.ruff.toml index f28b9056483a45..952dbe1464ea8d 100644 --- a/api/tests/unit_tests/.ruff.toml +++ b/api/tests/unit_tests/.ruff.toml @@ -102,7 +102,7 @@ extend-select = ["ANN401", "ARG"] "core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline.py" = ["ARG001", "ARG002", "ARG005"] "core/app/task_pipeline/test_message_cycle_manager.py" = ["ARG002"] "core/app/app_config/easy_ui_based_app/model_config/test_manager.py" = ["ARG005"] -"services/test_workflow_persistence.py" = ["ANN401", "ARG001", "ARG005", "TID251"] +"services/test_workflow_persistence.py" = ["ANN401", "ARG005", "TID251"] "core/app/workflow/test_file_runtime.py" = ["ARG001", "ARG005"] "core/app/workflow/layers/test_observability.py" = ["ARG002", "ARG005"] "core/base/tts/test_app_generator_tts_publisher.py" = ["ARG002"] @@ -122,9 +122,6 @@ extend-select = ["ANN401", "ARG"] "core/mcp/session/test_base_session.py" = ["S110"] "core/mcp/test_mcp_client.py" = ["ARG002"] "core/memory/test_token_buffer_memory.py" = ["ARG002"] -"core/ops/test_base_trace_instance.py" = ["ARG001", "ARG004"] -"core/ops/test_ops_trace_manager.py" = ["ARG001", "ARG002", "ARG004", "ARG005"] -"core/ops/entities/test_trace_entity.py" = ["ARG001", "ARG005"] "core/plugin/impl/test_agent.py" = ["ARG001"] "core/plugin/impl/test_datasource.py" = ["ARG001"] "core/plugin/impl/test_oauth.py" = ["ARG001"] @@ -157,7 +154,6 @@ extend-select = ["ANN401", "ARG"] "core/repositories/test_human_input_repository.py" = ["ANN401", "ARG001", "ARG004", "ARG005", "TID251"] "core/repositories/test_sqlalchemy_workflow_node_execution_repository.py" = ["ANN401", "ARG002", "ARG005", "TID251"] "core/schemas/test_resolver.py" = ["ARG005", "T201"] -"core/telemetry/test___init__.py" = ["ARG002", "ARG004"] "core/telemetry/test_gateway.py" = ["ARG002"] "core/test_model_manager.py" = ["ARG001"] "core/trigger/debug/test_event_selectors.py" = ["ARG002"] @@ -277,7 +273,6 @@ extend-select = ["ANN401", "ARG"] "tasks/test_document_indexing_sync_task.py" = ["ARG002"] "tasks/test_human_input_timeout_tasks.py" = ["ARG001", "ARG002", "ARG005", "TID251"] "tasks/test_initialize_created_app_rbac_access_task.py" = ["ARG005"] -"tasks/test_ops_trace_task.py" = ["ARG004"] "tasks/test_process_tenant_plugin_autoupgrade_check_task.py" = ["ARG001"] "tasks/test_remove_app_and_related_data_task.py" = ["ARG002"] "tasks/app_generate/test_workflow_execute_task.py" = ["ARG005"] diff --git a/api/tests/unit_tests/configs/test_app_config.py b/api/tests/unit_tests/configs/test_app_config.py index ed5a9adf4d8c06..1238fb0294a379 100644 --- a/api/tests/unit_tests/configs/test_app_config.py +++ b/api/tests/unit_tests/configs/test_app_config.py @@ -15,32 +15,15 @@ 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_export_attempts_must_be_positive() -> None: + with pytest.raises(ValidationError, match="OPS_TRACE_EXPORT_MAX_ATTEMPTS"): + OpsTraceConfig(OPS_TRACE_EXPORT_MAX_ATTEMPTS=0) -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, - ) +@pytest.mark.parametrize("queue", ["", " ", "ops_trace,ops_trace_v2", "ops trace", "ops_trace\n"]) +def test_ops_trace_queue_requires_one_shell_safe_queue_name(queue: str) -> None: + with pytest.raises(ValidationError, match="OPS_TRACE_QUEUE"): + OpsTraceConfig(OPS_TRACE_QUEUE=queue) @pytest.mark.parametrize( @@ -135,8 +118,9 @@ def test_dify_config() -> None: assert config.TEMPLATE_TRANSFORM_MAX_LENGTH == 400_000 assert config.APP_MAX_EXECUTION_TIME == 3600 assert config.WORKFLOW_MAX_EXECUTION_TIME == 3600 - assert config.OPS_TRACE_RETRYABLE_DISPATCH_MAX_RETRIES == 780 - assert config.OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS == 3900 + assert config.OPS_TRACE_EXPORT_MAX_ATTEMPTS == 6 + assert config.OPS_TRACE_QUEUE == "ops_trace_v2" + assert config.ENABLE_TRACE_PAYLOAD_SWEEP_TASK is True assert config.GRAPH_ENGINE_MAX_WORKERS == 10 # annotated field with custom configured value @@ -149,6 +133,21 @@ def test_dify_config() -> None: assert Version(config.project.version) >= Version("1.0.0") +def test_dify_config_ignores_removed_trace_dispatch_settings() -> None: + removed = { + "OPS_TRACE_UNIFIED_ENABLED": "true", + "OPS_TRACE_RETRYABLE_DISPATCH_MAX_RETRIES": "780", + "OPS_TRACE_RETRYABLE_DISPATCH_DELAY_SECONDS": "5", + # Shorter than the old retry window, which used to fail startup. + "OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS": "19", + } + + config = _make_config(**removed) + + assert not any(hasattr(config, name) for name in removed) + assert config.OPS_TRACE_EXPORT_MAX_ATTEMPTS == 6 + + @pytest.mark.parametrize( ("environment_value", "expected"), [ diff --git a/api/tests/unit_tests/conftest.py b/api/tests/unit_tests/conftest.py index f7689017428771..b5cf36363aefbb 100644 --- a/api/tests/unit_tests/conftest.py +++ b/api/tests/unit_tests/conftest.py @@ -14,6 +14,7 @@ os.environ.setdefault("STORAGE_TYPE", "opendal") import pytest +from dify_trace_sdk.testing import RecordingEgress from flask import Flask from redis import Redis from sqlalchemy import create_engine @@ -26,6 +27,7 @@ from models.base import TypeBase from tests.helpers.unit.config_override import apply_config_overrides from tests.helpers.unit.rbac_fakes import RBACDomain, build_rbac_domain +from tests.helpers.unit.services.tracing_ports import deployment_provider_configs if TYPE_CHECKING: from extensions.ext_application_services import ApplicationServices @@ -131,7 +133,7 @@ def reset_redis_mock(_patch_redis_clients: None) -> None: redis_mock.get.reset_mock() redis_mock.setex.reset_mock() redis_mock.setnx.reset_mock() - redis_mock.lock.reset_mock() + redis_mock.lock.reset_mock(return_value=True, side_effect=True) redis_mock.exists.reset_mock() redis_mock.set.reset_mock() redis_mock.expire.reset_mock() @@ -359,7 +361,6 @@ def app_services(sqlite_session_factory: sessionmaker[Session], rbac_domain: RBA from extensions.ext_redis import redis_client from repositories.app.console_repository import ConsoleAppRepository from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository - from services.app_tracing_config_gateway import OpsTraceManagerGateway from services.app_tracing_config_service import AppTracingConfigService from services.recommended_app_package_service import RecommendedAppPackageService @@ -369,7 +370,8 @@ def app_services(sqlite_session_factory: sessionmaker[Session], rbac_domain: RBA rbac_members=rbac_domain.rbac.members, tracing_configs=AppTracingConfigService( configs=SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory), - provider=OpsTraceManagerGateway(), + # Unit tests never reach a tracing vendor; an unexpected request fails the test. + provider=deployment_provider_configs(RecordingEgress([], grpc_target=None)), ), oauth=_build_oauth_server_service(database_client=sqlite_session_factory, redis=redis_client), recommended_packages=RecommendedAppPackageService(sources=Mock(), exporter=Mock()), diff --git a/api/tests/unit_tests/controllers/console/agent/test_tracing.py b/api/tests/unit_tests/controllers/console/agent/test_tracing.py index cac1ebba8696be..babc94466565f1 100644 --- a/api/tests/unit_tests/controllers/console/agent/test_tracing.py +++ b/api/tests/unit_tests/controllers/console/agent/test_tracing.py @@ -1,42 +1,27 @@ -import importlib as _test_importlib import inspect -import json from collections.abc import Callable -from copy import deepcopy -from dataclasses import dataclass, field -from uuid import UUID, uuid4 +from dataclasses import dataclass, field, replace +from datetime import UTC, datetime +from uuid import UUID import pytest -from dify_trace_arize_phoenix.arize_phoenix_trace import ArizePhoenixDataTrace from flask import Flask -from flask_restx import Api from pydantic import ValidationError -from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker from werkzeug.exceptions import Forbidden, NotFound from controllers.common.rbac import AgentId, RBACCheck -from controllers.console import flask_admission, wraps +from controllers.console import flask_admission from controllers.console.agent import tracing -from controllers.console.agent.tracing import AgentTraceConfigApi from controllers.console.app.app import AppTracePayload -from controllers.console.app.ops_trace import TraceAppConfigApi -from core.helper import ssrf_proxy -from core.ops.ops_trace_manager import OpsTraceManager -from enums import DeploymentEdition from extensions.ext_application_services import ApplicationServices from libs.login import AccountWithTenant from machinery.context import RequestContext from models.account import Account, AccountStatus, TenantAccountRole -from models.agent import Agent, AgentScope, AgentSource, AgentStatus -from models.model import App, AppMode, TraceAppConfig -from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository from services.app.agent_app_contracts import AgentAppNotFoundError from services.app_tracing_config_service import AppTraceSettings, AppTracingConfigRecord +from services.tracing.provider_configs import TraceProviderCapabilitiesQuery +from services.tracing.registry import EntryPointTraceProviderCatalog from tests.helpers.unit.config_override import apply_config_overrides -from tests.helpers.unit.model_factories import make_app - -pytest_plugins = ("tests.unit_tests.core.ops.test_ops_trace_manager",) """Agent tracing Console contracts and authorization stay separate from PlainApp.""" @@ -53,6 +38,15 @@ class ConfigCase: failure: Exception | None = None +def _collect_agent_traces(monkeypatch: pytest.MonkeyPatch, case: ConfigCase, *, enabled: bool) -> None: + """Serve the case's services from a deployment that collects Agent traces or not.""" + capabilities = TraceProviderCapabilitiesQuery( + providers=EntryPointTraceProviderCatalog(), api_collection_enabled=enabled + ) + services = replace(case.services, tracing=replace(case.services.tracing, provider_capabilities=capabilities)) + monkeypatch.setattr(tracing, "application_services", lambda: services) + + class TestTracingCases: @staticmethod @pytest.fixture @@ -110,7 +104,7 @@ def test_enable_requires_provider_even_when_field_omitted() -> None: def test_capabilities_reports_current_api_gate_without_claiming_runtime_delivery( app: Flask, service: ConfigCase, monkeypatch: pytest.MonkeyPatch, enabled: bool ) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=enabled) + _collect_agent_traces(monkeypatch, service, enabled=enabled) with app.test_request_context("/?tracing_provider=langfuse"): result = inspect.unwrap(tracing.AgentTraceCapabilitiesApi.get)( tracing.AgentTraceCapabilitiesApi(), CONTEXT, AGENT_ID @@ -134,7 +128,7 @@ def get_trace(context: RequestContext, app_id: str) -> AppTraceSettings: return AppTraceSettings() monkeypatch.setattr(service.services.apps.console, "get_trace", get_trace) - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=enabled) + _collect_agent_traces(monkeypatch, service, enabled=enabled) with app.test_request_context("/?tracing_provider=mlflow"): result = inspect.unwrap(tracing.AppTraceCapabilitiesApi.get)( tracing.AppTraceCapabilitiesApi(), CONTEXT, AGENT_ID @@ -214,8 +208,6 @@ def denied(**kwargs: object) -> None: @staticmethod def test_read_settings_and_masked_config_keep_agent_identity(app: Flask, service: ConfigCase) -> None: - from datetime import UTC, datetime - now = datetime(2026, 1, 1, tzinfo=UTC) service.record = AppTracingConfigRecord( id="config", @@ -305,193 +297,3 @@ def missing(_context: RequestContext, _app_id: str) -> AppTraceSettings: monkeypatch.setattr(service.services.apps.console, "get_trace", missing) with app.test_request_context("/?tracing_provider=langfuse"), pytest.raises(NotFound): inspect.unwrap(tracing.AppTraceCapabilitiesApi.get)(tracing.AppTraceCapabilitiesApi(), CONTEXT, AGENT_ID) - - -"""Both real config Resources reach the same ACK gate before atomic SQL updates. - -HTTP dispatch, payload validation, owner resolution, DI services, encryption and -SQLite persistence are real. Login admission supplies a local owner account and -the existing RSA fixture supplies private-key storage; no external account is used. -""" - - -_324fde08_TestArizeConfigurationAckCases = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_arize_configuration" -).TestArizeConfigurationAckCases - - -class TestArizeTracingAckRoutesCases: - collector = staticmethod(_324fde08_TestArizeConfigurationAckCases.collector) - - @staticmethod - @pytest.mark.parametrize("surface", ["app", "agent"]) - @pytest.mark.parametrize("method", ["POST", "PATCH"]) - def test_config_write_rejects_bad_ack_without_changing_old_rows_policy_or_cache( - workspace: str, - sqlite_session_factory: sessionmaker[Session], - account_application_services: ApplicationServices, - monkeypatch: pytest.MonkeyPatch, - collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], - surface: str, - method: str, - ) -> None: - endpoint, calls, _ = collector - app_id, agent_id, account_id = (str(uuid4()) for _ in range(3)) - old_policy = json.dumps( - {"enabled": True, "tracing_provider": "langsmith", "capture_content": False, "capture_agent": False} - ) - with sqlite_session_factory.begin() as session: - session.add( - make_app( - app_id=app_id, - tenant_id=workspace, - mode=AppMode.AGENT if surface == "agent" else AppMode.WORKFLOW, - tracing=old_policy, - ) - ) - if surface == "agent": - session.add( - Agent( - id=agent_id, - tenant_id=workspace, - name="ACK test", - scope=AgentScope.ROSTER, - source=AgentSource.AGENT_APP, - app_id=app_id, - status=AgentStatus.ACTIVE, - created_by=account_id, - updated_by=account_id, - ) - ) - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - original = OpsTraceManager.encrypt_tracing_config( - workspace, "langsmith", {"api_key": "other-provider-key", "project": "other"} - ) - assert store.create( - workspace_id=workspace, app_id=app_id, tracing_provider="langsmith", tracing_config=original - ) - if method == "PATCH": - old = OpsTraceManager.encrypt_tracing_config( - workspace, - "arize", - { - "endpoint": endpoint, - "project": "old-project", - "api_key": "synthetic-good", - "space_id": "synthetic-space", - }, - ) - old.update( - { - "project_url": "https://app.arize.com/?redirect_project_name=old-project", - "arize_project_destination": {"legacy": "unchanged"}, - } - ) - assert store.create(workspace_id=workspace, app_id=app_id, tracing_provider="arize", tracing_config=old) - with sqlite_session_factory() as session: - before = deepcopy( - list( - session.execute( - select(TraceAppConfig.__table__) - .where(TraceAppConfig.app_id == app_id) - .order_by(TraceAppConfig.id) - ).mappings() - ) - ) - - account = Account(name="Owner", email="owner@example.test", status=AccountStatus.ACTIVE) - account.id, account.role = account_id, TenantAccountRole.OWNER - identity = AccountWithTenant(account, workspace) - monkeypatch.setattr(flask_admission, "current_account_with_tenant", lambda: identity) - monkeypatch.setattr(wraps, "current_account_with_tenant", lambda: identity) - apply_config_overrides( - monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, LOGIN_DISABLED=True, RBAC_ENABLED=False - ) - - def forbidden(*_args: object, **_kwargs: object) -> None: - pytest.fail("Local custom collector verification must not perform optional Cloud HTTP reads") - - monkeypatch.setattr(ssrf_proxy, "get", forbidden) - monkeypatch.setattr(ArizePhoenixDataTrace, "__init__", forbidden) - app = Flask(__name__) - app.extensions["application_services"] = account_application_services - api = Api(app) - api.add_resource(TraceAppConfigApi, "/console/api/apps//trace-config") - api.add_resource(AgentTraceConfigApi, "/console/api/agent//trace-config") - url = ( - f"/console/api/apps/{app_id}/trace-config" - if surface == "app" - else f"/console/api/agent/{agent_id}/trace-config" - ) - values = { - "endpoint": endpoint, - "project": "not-created-by-verification", - "api_key": "synthetic-invalid", - "space_id": "synthetic-space", - } - with app.test_client() as client: - failed = client.open( - url, method=method, json={"tracing_provider": "arize", "tracing_config": values, "activate": True} - ) - assert failed.status_code == 400 - body = failed.json - assert isinstance(body, dict) - assert body["code"] == "tracing_config_verification_failed" - assert "private-vendor-details" not in failed.get_data(as_text=True) - assert "synthetic-invalid" not in failed.get_data(as_text=True) - with sqlite_session_factory() as session: - after = list( - session.execute( - select(TraceAppConfig.__table__) - .where(TraceAppConfig.app_id == app_id) - .order_by(TraceAppConfig.id) - ).mappings() - ) - assert after == before - assert session.scalar(select(App.tracing).where(App.id == app_id)) == old_policy - values["api_key"] = "synthetic-good" - accepted = client.open( - url, method=method, json={"tracing_provider": "arize", "tracing_config": values, "activate": True} - ) - assert accepted.status_code == (201 if method == "POST" else 200) - saved = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="arize") - assert saved is not None - assert saved.tracing_config is not None - assert ( - OpsTraceManager.decrypt_tracing_config(workspace, "arize", saved.tracing_config)["api_key"] - == "synthetic-good" - ) - assert "project_url" not in saved.tracing_config - assert "arize_project_destination" not in saved.tracing_config - assert len(calls) == 2 - with sqlite_session_factory() as session: - policy = json.loads(session.scalar(select(App.tracing).where(App.id == app_id)) or "{}") - assert policy == { - "enabled": True, - "tracing_provider": "arize", - "capture_content": False, - "capture_agent": False, - } - if method == "PATCH": - with app.test_client() as client: - masked = client.patch( - url, - json={ - "tracing_provider": "arize", - "tracing_config": { - **values, - "project": "next-project", - "api_key": "******", - "space_id": "******", - }, - }, - ) - assert masked.status_code == 200 - updated = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="arize") - assert updated is not None - assert updated.tracing_config is not None - assert updated.tracing_config["api_key"] == saved.tracing_config["api_key"] - assert updated.tracing_config["space_id"] == saved.tracing_config["space_id"] - assert calls[-1][1]["api_key"] == "synthetic-good" - assert calls[-1][1]["space_id"] == "synthetic-space" - assert len(calls) == 3 diff --git a/api/tests/unit_tests/controllers/console/agent/test_tracing_runs.py b/api/tests/unit_tests/controllers/console/agent/test_tracing_runs.py index c817de50d467eb..2c661beeb14082 100644 --- a/api/tests/unit_tests/controllers/console/agent/test_tracing_runs.py +++ b/api/tests/unit_tests/controllers/console/agent/test_tracing_runs.py @@ -2,7 +2,7 @@ import inspect from collections.abc import Callable -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from datetime import UTC, datetime, timedelta, timezone from functools import wraps from uuid import UUID @@ -11,20 +11,20 @@ from flask import Flask from flask_restx import Resource from pydantic import ValidationError -from sqlalchemy.orm import Session, sessionmaker from werkzeug.exceptions import Forbidden, NotFound from controllers.common.rbac import AgentId, PlainApp, RBACCheck from controllers.console import flask_admission from controllers.console.agent import tracing_runs +from domain.tracing.trace_run_agg import TraceRunState, TraceSourceKind from extensions.ext_application_services import ApplicationServices from libs.external_api import ExternalApi from libs.login import AccountWithTenant from machinery.context import RequestContext from models.account import Account, AccountStatus, TenantAccountRole -from repositories.application_trace_repository import ApplicationTraceRepository, TraceCollectionError from services.app.agent_app_contracts import AgentAppNotFoundError -from services.tracing.service import ApplicationTracingService +from services.tracing.collection import TraceRunStatus +from services.tracing.errors import TraceCollectionConflictError, TraceCollectionError, TraceCollectionNotFoundError from tests.helpers.unit.config_override import apply_config_overrides CONTEXT = RequestContext("request", None, "account", "tenant") @@ -32,30 +32,45 @@ APP_ID = UUID("01a22222-2222-7222-8222-222222222222") SOURCE_ID = UUID("01a33333-3333-7333-8333-333333333333") NODE_ID = UUID("01a44444-4444-7444-8444-444444444444") +NOT_COLLECTED = TraceRunStatus( + status=TraceRunState.NOT_COLLECTED, + source_kind="message", + source_id=str(SOURCE_ID), + provider=None, + trace_id=None, + trace_url=None, + span_url=None, + project_url=None, + config_revision=None, + capture_content=None, + capture_agent=None, + span_count=None, + attempts=None, + error_code=None, + accepted_at=None, + capabilities=None, +) @dataclass class RunCase: - owners: ApplicationServices - runs: ApplicationTracingService + services: ApplicationServices + """The composition root the handlers resolve; its Agent trace runs answer with ``response``.""" owner_calls: list[tuple[RequestContext, str]] = field(default_factory=list) status_calls: list[dict[str, object]] = field(default_factory=list) retry_calls: list[dict[str, object]] = field(default_factory=list) - response: dict[str, object] = field( - default_factory=lambda: {"status": "not_collected", "source_kind": "message", "source_id": str(SOURCE_ID)} - ) + response: TraceRunStatus = NOT_COLLECTED owner_error: Exception | None = None retry_error: Exception | None = None @pytest.fixture def services( + app: Flask, monkeypatch: pytest.MonkeyPatch, account_application_services: ApplicationServices, - sqlite_session_factory: sessionmaker[Session], ) -> RunCase: - runs = ApplicationTracingService(ApplicationTraceRepository(sqlite_session_factory)) - case = RunCase(account_application_services, runs) + case = RunCase(account_application_services) def resolve(context: RequestContext, agent_id: str) -> str: case.owner_calls.append((context, agent_id)) @@ -63,21 +78,21 @@ def resolve(context: RequestContext, agent_id: str) -> str: raise case.owner_error return str(APP_ID) - def status(**kwargs: object) -> dict[str, object]: + def status(**kwargs: object) -> TraceRunStatus: case.status_calls.append(kwargs) return case.response - def retry(**kwargs: object) -> dict[str, object]: + def retry(**kwargs: object) -> TraceRunStatus: case.retry_calls.append(kwargs) if case.retry_error is not None: raise case.retry_error - return {**case.response, "status": "retrying"} + return replace(case.response, status=TraceRunState.RETRYING) + runs = account_application_services.tracing.agent_traces monkeypatch.setattr(account_application_services.agent_tracing_configs, "resolve_app_id", resolve) monkeypatch.setattr(runs, "status", status) monkeypatch.setattr(runs, "retry", retry) - monkeypatch.setattr(tracing_runs, "application_services", lambda: account_application_services) - monkeypatch.setattr(tracing_runs, "application_tracing_service", lambda: runs) + monkeypatch.setitem(app.extensions, "application_services", account_application_services) return case @@ -92,18 +107,15 @@ def retry(**kwargs: object) -> dict[str, object]: def test_agent_status_resolves_owner_and_never_exposes_destination_secrets( app: Flask, services: RunCase, accepted_at: datetime ) -> None: - services.response.update( - { - "status": "accepted", - "provider": "langfuse", - "trace_id": "trace-id", - "trace_url": "https://example.com/trace", - "project_url": "https://example.com/project/frozen", - "accepted_at": accepted_at, - "queryable": None, - "provider_config": {"secret_key": "do-not-return"}, - "config_revision": "frozen-revision", - } + services.response = replace( + services.response, + status=TraceRunState.ACCEPTED, + provider="langfuse", + trace_id="trace-id", + trace_url="https://example.com/trace", + project_url="https://example.com/project/frozen", + accepted_at=accepted_at, + config_revision="frozen-revision", ) with app.test_request_context(f"/?node_execution_id={NODE_ID}"): result = inspect.unwrap(tracing_runs.AgentTraceRunApi.get)( @@ -124,12 +136,11 @@ def test_agent_status_resolves_owner_and_never_exposes_destination_secrets( assert result["config_revision"] == "frozen-revision" assert result["project_url"] == "https://example.com/project/frozen" assert "provider_config" not in result - assert "do-not-return" not in str(result) @pytest.mark.parametrize("kind", ["message", "workflow", "node"]) -def test_app_status_source_scope_and_not_collected(app: Flask, services: RunCase, kind: str) -> None: - services.response["source_kind"] = kind +def test_app_status_source_scope_and_not_collected(app: Flask, services: RunCase, kind: TraceSourceKind) -> None: + services.response = replace(services.response, source_kind=kind) with app.test_request_context(): result = inspect.unwrap(tracing_runs.AppTraceRunApi.get)( tracing_runs.AppTraceRunApi(), CONTEXT, APP_ID, kind, SOURCE_ID @@ -204,18 +215,20 @@ def test_agent_owner_rejection_prevents_run_lookup(app: Flask, services: RunCase @pytest.mark.parametrize( - ("code", "status"), + ("error", "code", "status"), [ - ("tracing_source_not_found", 404), - ("tracing_node_not_found", 404), - ("tracing_run_not_found", 404), - ("tracing_already_accepted", 409), - ("tracing_retry_not_allowed", 409), - ("tracing_duplicate_acknowledgement_required", 409), + (TraceCollectionNotFoundError, "tracing_source_not_found", 404), + (TraceCollectionNotFoundError, "tracing_node_not_found", 404), + (TraceCollectionNotFoundError, "tracing_run_not_found", 404), + (TraceCollectionConflictError, "tracing_already_accepted", 409), + (TraceCollectionConflictError, "tracing_retry_not_allowed", 409), + (TraceCollectionConflictError, "tracing_duplicate_acknowledgement_required", 409), ], ) -def test_retry_business_errors_match_served_error_contract(services: RunCase, code: str, status: int) -> None: - services.retry_error = TraceCollectionError(code, status) +def test_retry_business_errors_match_served_error_contract( + services: RunCase, error: type[TraceCollectionError], code: str, status: int +) -> None: + services.retry_error = error(code) target = tracing_runs.AppTraceRunRetryApi.post @wraps(target) @@ -224,6 +237,7 @@ def post(self: Resource) -> dict[str, object]: resource = type("RetryContractResource", (Resource,), {"post": post}) app = Flask("trace-error-contract") + app.extensions["application_services"] = services.services api = ExternalApi(app) api.add_resource(resource, "/retry") with app.test_client() as client: diff --git a/api/tests/unit_tests/controllers/console/app/test_ops_trace.py b/api/tests/unit_tests/controllers/console/app/test_ops_trace.py index 9d0e9847a14d60..3242d6caa671a0 100644 --- a/api/tests/unit_tests/controllers/console/app/test_ops_trace.py +++ b/api/tests/unit_tests/controllers/console/app/test_ops_trace.py @@ -1,14 +1,19 @@ from __future__ import annotations import inspect +import json from collections.abc import Callable +from dataclasses import dataclass from datetime import datetime from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock from uuid import UUID import pytest +from dify_trace_sdk.testing import RecordingEgress from flask import Flask +from sqlalchemy import func, select +from sqlalchemy.orm import Session, sessionmaker from werkzeug.exceptions import Forbidden from controllers.common.rbac import PlainApp @@ -30,7 +35,8 @@ from libs.login import AccountWithTenant from machinery.context import RequestContext from models.account import Account, AccountStatus, TenantAccountRole -from services.app_tracing_config_gateway import OpsTraceManagerGateway +from models.model import TraceAppConfig +from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository from services.app_tracing_config_service import ( AppTracingConfigAlreadyExistsError, AppTracingConfigAppNotFoundError, @@ -45,6 +51,8 @@ AppTracingConfigVerificationFailedError, ) from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_app +from tests.helpers.unit.services.tracing_ports import deployment_provider_configs APP_ID = "11111111-1111-1111-1111-111111111111" WORKSPACE_ID = "22222222-2222-2222-2222-222222222222" @@ -114,21 +122,56 @@ def tracing_configs(monkeypatch: pytest.MonkeyPatch) -> MagicMock: return service -@pytest.mark.parametrize("method_name", ["get", "delete"]) -def test_trace_config_without_sdk_can_be_read_when_absent_or_deleted( - app: Flask, monkeypatch: pytest.MonkeyPatch, method_name: str -) -> None: - configs = MagicMock() - configs.get.return_value = None - configs.delete.return_value = True - service = AppTracingConfigService(configs=configs, provider=OpsTraceManagerGateway()) +@dataclass(frozen=True) +class _ApplicationServices: + app_tracing_configs: AppTracingConfigService + + +@dataclass(frozen=True) +class _StoredConfigs: + """Configurations persisted the way the provider gateway stores them, behind the real service.""" + + sessions: sessionmaker[Session] + egress: RecordingEgress + + def add(self, tracing_provider: str, tracing_config: dict[str, str | None]) -> TraceAppConfig: + with self.sessions.begin() as session: + session.add(TraceAppConfig(app_id=APP_ID, tracing_provider=tracing_provider, tracing_config=tracing_config)) + with self.sessions() as session: + return session.scalars( + select(TraceAppConfig).where( + TraceAppConfig.app_id == APP_ID, TraceAppConfig.tracing_provider == tracing_provider + ) + ).one() + + def count(self) -> int: + with self.sessions() as session: + return session.scalar(select(func.count()).select_from(TraceAppConfig)) or 0 + + +@pytest.fixture +def stored_configs(monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session]) -> _StoredConfigs: + with sqlite_session_factory.begin() as session: + session.add(make_app(app_id=APP_ID, tenant_id=WORKSPACE_ID)) + egress = RecordingEgress([], grpc_target=None) + service = AppTracingConfigService( + configs=SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory), + provider=deployment_provider_configs(egress), + ) monkeypatch.setattr( - ops_trace_module, - "application_services", - lambda: SimpleNamespace(app_tracing_configs=service), + ops_trace_module, "application_services", lambda: _ApplicationServices(app_tracing_configs=service) ) + return _StoredConfigs(sqlite_session_factory, egress) - with patch.object(OpsTraceManagerGateway, "_provider_config") as load_provider, app.test_request_context("/"): + +@pytest.mark.parametrize("method_name", ["get", "delete"]) +def test_trace_config_read_when_absent_or_deleted_never_contacts_the_provider( + app: Flask, stored_configs: _StoredConfigs, method_name: str +) -> None: + if method_name == "delete": + stored_configs.add("weave", {"api_key": "ciphertext-api-key", "entity": "dify", "project": "traces"}) + + with app.test_request_context("/"): result = _original(_CONTROLLER_METHODS[method_name])( ops_trace_module.TraceAppConfigApi(), ops_trace_module.TraceProviderQuery(tracing_provider="weave"), @@ -137,9 +180,87 @@ def test_trace_config_without_sdk_can_be_read_when_absent_or_deleted( ) assert result == ({"has_not_configured": True} if method_name == "get" else ("", 204)) - load_provider.assert_not_called() - if method_name == "delete": - configs.delete.assert_called_once_with(workspace_id=WORKSPACE_ID, app_id=APP_ID, tracing_provider="weave") + assert stored_configs.count() == 0 + assert stored_configs.egress.requests == [] + + +@pytest.mark.parametrize( + ("provider", "stored", "presented", "configured_secrets"), + [ + pytest.param( + "langfuse", + { + "public_key": "ciphertext-public-key", + "secret_key": "ciphertext-secret-key", + "host": "https://langfuse.example.com", + "project_key": "project-1", + }, + {"host": "https://langfuse.example.com", "project_url": "https://langfuse.example.com/project/project-1"}, + ["public_key", "secret_key"], + id="all-secrets-set", + ), + pytest.param( + "databricks", + { + "host": "https://dbc.example.com", + "experiment_id": "7", + "client_id": "client-1", + "client_secret": "", + "personal_access_token": "ciphertext-token", + }, + { + "host": "https://dbc.example.com", + "experiment_id": "7", + "client_id": "client-1", + "project_url": "https://dbc.example.com/ml/experiments/7/traces", + }, + ["personal_access_token"], + id="empty-secret-not-configured", + ), + pytest.param( + "mlflow", + {"tracking_uri": "https://mlflow.example.com", "experiment_id": "7", "username": "admin", "password": None}, + { + "tracking_uri": "https://mlflow.example.com", + "experiment_id": "7", + "username": "admin", + "project_url": "https://mlflow.example.com/#/experiments/7/traces", + }, + [], + id="no-secrets-set", + ), + ], +) +def test_trace_app_config_get_returns_configured_secrets_and_never_secret_values( + app: Flask, + stored_configs: _StoredConfigs, + provider: str, + stored: dict[str, str | None], + presented: dict[str, str], + configured_secrets: list[str], +) -> None: + record = stored_configs.add(provider, stored) + + with app.test_request_context(f"/?tracing_provider={provider}"): + result = _original(ops_trace_module.TraceAppConfigApi.get)( + ops_trace_module.TraceAppConfigApi(), + ops_trace_module.TraceProviderQuery(tracing_provider=provider), + _request_context(), + UUID(APP_ID), + ) + + assert result == { + "id": record.id, + "app_id": APP_ID, + "tracing_provider": provider, + "tracing_config": presented, + "configured_secrets": configured_secrets, + "is_active": True, + "created_at": str(record.created_at), + "updated_at": str(record.updated_at), + } + assert "ciphertext" not in json.dumps(result) + assert stored_configs.egress.requests == [] @pytest.mark.parametrize( @@ -270,26 +391,31 @@ def test_trace_app_config_get_empty_returns_exact_legacy_body( @pytest.mark.parametrize( - ("provider", "config"), + ("provider", "config", "configured_secrets"), [ - ("langfuse", {"public_key": "pk", "secret_key": "******"}), + ( + "langfuse", + {"host": "https://cloud.langfuse.com", "project_url": "https://cloud.langfuse.com/project/p-1"}, + ("public_key", "secret_key"), + ), ( "mlflow", { "tracking_uri": "http://localhost:5000", "experiment_id": "0", "username": None, - "password": "******", - "project_url": "http://localhost:5000", + "project_url": "http://localhost:5000/#/experiments/0/traces", }, + (), ), ], ) -def test_trace_app_config_get_configured_returns_exact_legacy_body( +def test_trace_app_config_get_configured_returns_exact_body( app: Flask, tracing_configs: MagicMock, provider: str, config: dict[str, str | None], + configured_secrets: tuple[str, ...], ) -> None: tracing_configs.get.return_value = AppTracingConfigRecord( id="trace-config-1", @@ -299,6 +425,7 @@ def test_trace_app_config_get_configured_returns_exact_legacy_body( is_active=True, created_at=datetime(2026, 1, 2, 3, 4, 5), updated_at=datetime(2026, 1, 3, 4, 5, 6), + configured_secrets=configured_secrets, ) with app.test_request_context(f"/?tracing_provider={provider}"): @@ -314,6 +441,7 @@ def test_trace_app_config_get_configured_returns_exact_legacy_body( "app_id": APP_ID, "tracing_provider": provider, "tracing_config": config, + "configured_secrets": list(configured_secrets), "is_active": True, "created_at": "2026-01-02 03:04:05", "updated_at": "2026-01-03 04:05:06", diff --git a/api/tests/unit_tests/controllers/inner_api/agent/test_tracing.py b/api/tests/unit_tests/controllers/inner_api/agent/test_tracing.py index 496772c2169ef2..046dd8f845297a 100644 --- a/api/tests/unit_tests/controllers/inner_api/agent/test_tracing.py +++ b/api/tests/unit_tests/controllers/inner_api/agent/test_tracing.py @@ -1,97 +1,188 @@ """HTTP authentication and acknowledgement backed by the real SQLite collector.""" from collections.abc import Iterator +from dataclasses import replace from uuid import uuid4 import pytest +from dify_trace_sdk.testing import RecordingEgress from flask import Flask -from flask.testing import FlaskClient from sqlalchemy import func, select from sqlalchemy.orm import Session, sessionmaker from controllers.inner_api import bp +from extensions.ext_application_services import ApplicationServices from models.application_trace import ApplicationTraceBatch +from repositories.factory import DifyAPIRepositoryFactory +from repositories.trace_record_repository import TraceRecordRepository +from services.tracing.canonical import CanonicalTraceBuilder +from services.tracing.collection import AgentTraceSettings +from services.tracing.collection_ports import ScheduleTraceExport +from services.tracing.providers import ProviderConfigCodec from services.tracing.service import ApplicationTracingService from tests.helpers.unit.config_override import config_overrides_context -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_batch_for as batch_for -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_seed as seed +from tests.helpers.unit.services.tracing_ports import ( + AGENT_TRACE_SETTINGS, + InstalledTraceProviders, + WorkspaceBoundCipher, +) +from tests.unit_tests.repositories.test_application_trace_repository import Collection, batch_for, seed + +BATCHES_URL = "/inner/api/agent/tracing/batches" +INNER_API_HEADERS = {"X-Inner-Api-Key": "inner"} @pytest.fixture -def client() -> Iterator[FlaskClient]: +def inner_app() -> Iterator[Flask]: app = Flask(__name__) app.config["TESTING"] = True app.register_blueprint(bp) - with config_overrides_context( - PLUGIN_DAEMON_KEY="daemon", INNER_API_KEY_FOR_PLUGIN="inner", OPS_AGENT_TRACE_ENABLED=True - ): - yield app.test_client() + with config_overrides_context(PLUGIN_DAEMON_KEY="daemon", INNER_API_KEY_FOR_PLUGIN="inner"): + yield app + + +def bind_collector( + app: Flask, + services: ApplicationServices, + collection: Collection, + *, + settings: AgentTraceSettings, + schedule_export: ScheduleTraceExport, +) -> None: + """Make the app's composition root hand the ingestion handler a collector over the seeded ledger.""" + collector = ApplicationTracingService( + runs=collection.repository, + run_query=collection.repository, + sources=collection.repository, + records=TraceRecordRepository( + collection.sessions, + workflow_runs=DifyAPIRepositoryFactory.create_api_workflow_run_repository(collection.sessions), + node_executions=DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(collection.sessions), + files_url="https://files.test", + enterprise_details=False, + ), + builder=CanonicalTraceBuilder(), + providers=InstalledTraceProviders(), + codec=ProviderConfigCodec(WorkspaceBoundCipher()), + # Ingesting batches never exports; an unexpected request fails the test. + egress=RecordingEgress([], grpc_target=None), + schedule_export=schedule_export, + settings=settings, + ) + app.extensions["application_services"] = replace( + services, tracing=replace(services.tracing, agent_traces=collector) + ) + + +def stored_batches(sessions: sessionmaker[Session]) -> int | None: + with sessions() as session: + return session.scalar(select(func.count()).select_from(ApplicationTraceBatch)) def test_trace_http_ack_follows_commit_and_replay_is_idempotent( - client: FlaskClient, + inner_app: Flask, + account_application_services: ApplicationServices, sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, ) -> None: - collection = seed(sqlite_session_factory, "http", capture_content=False) - invocation = collection.invocation() - payload = batch_for(collection.run, invocation).model_dump(mode="json") - service = ApplicationTracingService(collection.repository) - monkeypatch.setattr("controllers.inner_api.agent.tracing.application_tracing_service", lambda: service) - scheduled = [] + collection = seed(sqlite_session_factory, "http") + payload = batch_for(collection.run, collection.invocation()).model_dump(mode="json") + scheduled: list[str] = [] def schedule(run_id: str) -> None: # A fresh transaction can see the batch before the caller gets an ACK. - with sqlite_session_factory() as session: - assert session.scalar(select(func.count()).select_from(ApplicationTraceBatch)) == 1 + assert stored_batches(sqlite_session_factory) == 1 scheduled.append(run_id) - monkeypatch.setattr(service, "schedule", schedule) + bind_collector( + inner_app, account_application_services, collection, settings=AGENT_TRACE_SETTINGS, schedule_export=schedule + ) + client = inner_app.test_client() for _ in range(2): - response = client.post("/inner/api/agent/tracing/batches", json=payload, headers={"X-Inner-Api-Key": "inner"}) + response = client.post(BATCHES_URL, json=payload, headers=INNER_API_HEADERS) assert response.status_code == 200 assert response.get_json() == {"result": "success", "batch_id": payload["batch_id"], "accepted_spans": 1} assert scheduled == [collection.run.id, collection.run.id] payload["spans"][0]["name"] = "conflicting replay" - response = client.post("/inner/api/agent/tracing/batches", json=payload, headers={"X-Inner-Api-Key": "inner"}) + response = client.post(BATCHES_URL, json=payload, headers=INNER_API_HEADERS) assert response.status_code == 409 assert len(scheduled) == 2 @pytest.mark.parametrize("header", [{}, {"X-Inner-Api-Key": "wrong"}]) -def test_trace_http_requires_internal_auth(client: FlaskClient, header: dict[str, str]) -> None: - response = client.post("/inner/api/agent/tracing/batches", json={}, headers=header) +def test_trace_http_requires_internal_auth(inner_app: Flask, header: dict[str, str]) -> None: + response = inner_app.test_client().post(BATCHES_URL, json={}, headers=header) assert response.status_code == 404 def test_trace_http_rejects_invalid_span_and_unknown_owner( - client: FlaskClient, + inner_app: Flask, + account_application_services: ApplicationServices, sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, ) -> None: collection = seed(sqlite_session_factory, "http-validation") - invocation = collection.invocation() - payload = batch_for(collection.run, invocation).model_dump(mode="json") - service = ApplicationTracingService(collection.repository) - monkeypatch.setattr("controllers.inner_api.agent.tracing.application_tracing_service", lambda: service) + payload = batch_for(collection.run, collection.invocation()).model_dump(mode="json") + scheduled: list[str] = [] + bind_collector( + inner_app, + account_application_services, + collection, + settings=AGENT_TRACE_SETTINGS, + schedule_export=scheduled.append, + ) + client = inner_app.test_client() payload["spans"][0]["end_time_ns"] = 0 - response = client.post("/inner/api/agent/tracing/batches", json=payload, headers={"X-Inner-Api-Key": "inner"}) + response = client.post(BATCHES_URL, json=payload, headers=INNER_API_HEADERS) assert response.status_code == 400 payload["spans"][0]["end_time_ns"] = 2_000_000_000 payload["invocation_id"] = str(uuid4()) - response = client.post("/inner/api/agent/tracing/batches", json=payload, headers={"X-Inner-Api-Key": "inner"}) + response = client.post(BATCHES_URL, json=payload, headers=INNER_API_HEADERS) assert response.status_code == 404 - with sqlite_session_factory() as session: - assert session.scalar(select(func.count()).select_from(ApplicationTraceBatch)) == 0 + assert scheduled == [] + assert stored_batches(sqlite_session_factory) == 0 -def test_trace_http_enforces_body_limit_before_ingestion(client: FlaskClient, monkeypatch: pytest.MonkeyPatch) -> None: - calls = [] - monkeypatch.setattr("controllers.inner_api.agent.tracing.application_tracing_service", lambda: calls.append(True)) +def test_trace_http_rejects_batches_beyond_collection_limits( + inner_app: Flask, + account_application_services: ApplicationServices, + sqlite_session_factory: sessionmaker[Session], +) -> None: + collection = seed(sqlite_session_factory, "http-limits") + payload = batch_for(collection.run, collection.invocation()).model_dump(mode="json") + scheduled: list[str] = [] + bind_collector( + inner_app, + account_application_services, + collection, + settings=replace(AGENT_TRACE_SETTINGS, max_spans=0), + schedule_export=scheduled.append, + ) + response = inner_app.test_client().post(BATCHES_URL, json=payload, headers=INNER_API_HEADERS) + assert response.status_code == 413 + body = response.get_json() + assert body is not None + assert body["message"] == "tracing_payload_limit" + assert scheduled == [] + assert stored_batches(sqlite_session_factory) == 0 + + +def test_trace_http_enforces_body_limit_before_ingestion( + inner_app: Flask, + account_application_services: ApplicationServices, + sqlite_session_factory: sessionmaker[Session], +) -> None: + collection = seed(sqlite_session_factory, "http-body-limit") + payload = batch_for(collection.run, collection.invocation()).model_dump(mode="json") + scheduled: list[str] = [] + bind_collector( + inner_app, + account_application_services, + collection, + settings=AGENT_TRACE_SETTINGS, + schedule_export=scheduled.append, + ) with config_overrides_context(OPS_AGENT_TRACE_BATCH_MAX_BYTES=64): - response = client.post( - "/inner/api/agent/tracing/batches", json={"padding": "x" * 256}, headers={"X-Inner-Api-Key": "inner"} - ) + response = inner_app.test_client().post(BATCHES_URL, json=payload, headers=INNER_API_HEADERS) assert response.status_code == 413 - assert calls == [] + assert scheduled == [] + assert stored_batches(sqlite_session_factory) == 0 diff --git a/api/tests/unit_tests/controllers/service_api/app/test_workflow_events.py b/api/tests/unit_tests/controllers/service_api/app/test_workflow_events.py index e5011096700479..a145fe5d8b0879 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_workflow_events.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_workflow_events.py @@ -523,7 +523,7 @@ def test_advanced_chat_blocking_pipeline_pause_payload_contract(self) -> None: stream=False, invoke_from=InvokeFrom.WEB_APP, extras={}, - trace_manager=None, + trace_scope=None, workflow_run_id="run-id", ) pipeline = AdvancedChatAppGenerateTaskPipeline( @@ -613,7 +613,7 @@ def test_workflow_blocking_pipeline_pause_payload_contract(self) -> None: user_id="user", stream=False, invoke_from=InvokeFrom.WEB_APP, - trace_manager=None, + trace_scope=None, workflow_execution_id="run-id", extras={}, call_depth=0, diff --git a/api/tests/unit_tests/core/agent/test_cot_agent_runner.py b/api/tests/unit_tests/core/agent/test_cot_agent_runner.py index 635ddcfa803eb6..72afdc14757a3d 100644 --- a/api/tests/unit_tests/core/agent/test_cot_agent_runner.py +++ b/api/tests/unit_tests/core/agent/test_cot_agent_runner.py @@ -4,6 +4,8 @@ from unittest.mock import MagicMock import pytest +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import TraceCorrelation, TraceEvent from pytest_mock import MockerFixture from sqlalchemy import event from sqlalchemy.engine import Engine @@ -50,6 +52,26 @@ def _make_message(*, message_id: str = "msg-id", conversation_id: str = "conv1") return message +class _Recorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +def _trace_scope() -> TraceScope: + return TraceScope( + emitter=_Recorder(), + tenant_id="tenant", + app_id="app", + user_id="user", + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + ) + + class DummyRunner(CotAgentRunner): """Concrete implementation for testing abstract methods.""" @@ -78,7 +100,7 @@ def runner(mocker: MockerFixture, sqlite_engine: Engine) -> Iterator[DummyRunner application_generate_entity.model_conf.stop = [] application_generate_entity.model_conf.provider = "openai" application_generate_entity.model_conf.parameters = {} - application_generate_entity.trace_manager = None + application_generate_entity.trace_scope = None application_generate_entity.invoke_from = "test" app_config = MagicMock() @@ -208,7 +230,7 @@ def test_format_with_action_and_observation(self, runner: DummyRunner): class TestHandleInvokeAction: def test_handle_invoke_action_tool_not_present(self, runner: DummyRunner): action = AgentScratchpadUnit.Action(action_name="missing", action_input={}) - response, meta = runner._handle_invoke_action(runner.session, action, {}, []) + response, meta = runner._handle_invoke_action(runner.session, action, {}, [], trace_scope=None) assert "there is not a tool named" in response def test_tool_with_json_string_args(self, runner: DummyRunner, mocker: MockerFixture): @@ -221,7 +243,7 @@ def test_tool_with_json_string_args(self, runner: DummyRunner, mocker: MockerFix return_value=("result", [], MagicMock(to_dict=lambda: {})), ) - response, meta = runner._handle_invoke_action(runner.session, action, tool_instances, []) + response, meta = runner._handle_invoke_action(runner.session, action, tool_instances, [], trace_scope=None) assert response == "result" @@ -430,6 +452,8 @@ def test_run_usage_missing_key_branch(self, runner: DummyRunner, mocker: MockerF def test_run_prompt_tool_update_branch(self, runner: DummyRunner, mocker: MockerFixture): message = _make_message() message.id = "msg-id" + trace_scope = _trace_scope() + runner.application_generate_entity.trace_scope = trace_scope action = AgentScratchpadUnit.Action(action_name="tool", action_input={}) @@ -440,7 +464,7 @@ def test_run_prompt_tool_update_branch(self, runner: DummyRunner, mocker: Mocker side_effect=[[action], []], ) - mocker.patch( + agent_invoke = mocker.patch( "core.agent.cot_agent_runner.ToolEngine.agent_invoke", return_value=("ok", [], MagicMock(to_dict=lambda: {})), ) @@ -458,6 +482,7 @@ def test_run_prompt_tool_update_branch(self, runner: DummyRunner, mocker: Mocker list(runner.run(runner.session, message, "query", {})) runner.update_prompt_message_tool.assert_called_once() + assert agent_invoke.call_args.kwargs["trace_scope"] is trace_scope def test_historic_with_assistant_and_tool_calls(self, runner: DummyRunner): from graphon.model_runtime.entities.message_entities import AssistantPromptMessage, ToolPromptMessage @@ -507,7 +532,9 @@ def test_tool_with_invalid_json_string_args(self, runner: DummyRunner, mocker: M ) message_file_ids = [] - response, meta = runner._handle_invoke_action(runner.session, action, tool_instances, message_file_ids) + response, meta = runner._handle_invoke_action( + runner.session, action, tool_instances, message_file_ids, trace_scope=None + ) assert response == "ok" assert message_file_ids == ["file1"] diff --git a/api/tests/unit_tests/core/agent/test_fc_agent_runner.py b/api/tests/unit_tests/core/agent/test_fc_agent_runner.py index 0b6e9bd429ede0..49e56e36fb9ed3 100644 --- a/api/tests/unit_tests/core/agent/test_fc_agent_runner.py +++ b/api/tests/unit_tests/core/agent/test_fc_agent_runner.py @@ -7,6 +7,8 @@ from unittest.mock import MagicMock import pytest +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import TraceCorrelation, TraceEvent from pytest_mock import MockerFixture from sqlalchemy import event from sqlalchemy.engine import Engine @@ -84,6 +86,26 @@ def _make_message(*, message_id: str = "m1", conversation_id: str = "conv1") -> return message +class _Recorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +def _trace_scope() -> TraceScope: + return TraceScope( + emitter=_Recorder(), + tenant_id="tenant", + app_id="app", + user_id="user", + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + ) + + class DummyMessage: def __init__(self, content: str | None = None, tool_calls: list[Any] | None = None): self.content: str | None = content @@ -178,7 +200,7 @@ def runner(mocker: MockerFixture, sqlite_engine: Engine) -> Iterator[FunctionCal application_generate_entity = MagicMock() application_generate_entity.model_conf = MagicMock(parameters={}, stop=None) - application_generate_entity.trace_manager = MagicMock() + application_generate_entity.trace_scope = None application_generate_entity.invoke_from = "test" application_generate_entity.app_config = MagicMock(app_id="app") application_generate_entity.file_upload_config = None @@ -640,6 +662,8 @@ def test_run_with_missing_tool_instance(self, runner: FunctionCallAgentRunner): def test_run_with_tool_instance_and_files(self, runner: FunctionCallAgentRunner, mocker: MockerFixture): message = _make_message() + trace_scope = _trace_scope() + runner.application_generate_entity.trace_scope = trace_scope tool_call = MagicMock() tool_call.id = "1" @@ -659,13 +683,14 @@ def test_run_with_tool_instance_and_files(self, runner: FunctionCallAgentRunner, tool_invoke_meta = MagicMock() tool_invoke_meta.to_dict.return_value = {"ok": True} - mocker.patch( + agent_invoke = mocker.patch( "core.agent.fc_agent_runner.ToolEngine.agent_invoke", return_value=("ok", ["file1"], tool_invoke_meta), ) outputs = list(runner.run(runner.session, message, "query")) assert len(outputs) >= 1 + assert agent_invoke.call_args.kwargs["trace_scope"] is trace_scope assert any( isinstance(call.args[0], QueueMessageFileEvent) and call.args[0].message_file_id == "file1" diff --git a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py b/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py index 20bccd869f432e..ed8a48ee0c4d8d 100644 --- a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py @@ -2,16 +2,17 @@ import json import logging -import sys import time from collections.abc import Mapping from contextlib import contextmanager from decimal import Decimal -from types import ModuleType, SimpleNamespace +from types import SimpleNamespace from typing import Any from unittest.mock import MagicMock import pytest +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import TraceCorrelation, TraceEvent from pydantic import BaseModel, ValidationError from pytest_mock import MockerFixture from sqlalchemy import Engine as _app_generator_Engine @@ -32,12 +33,16 @@ ) from core.app.apps.exc import GenerateTaskStoppedError from core.app.apps.workflow import app_generator as wf_app_gen_module -from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom +from core.app.entities.app_invoke_entities import ( + AdvancedChatAppGenerateEntity, + InvokeFrom, + UserFrom, + WorkflowAppGenerateEntity, +) from core.app.file_access import DatabaseFileAccessController from core.app.task_pipeline import message_cycle_manager from core.app.task_pipeline.message_cycle_manager import MessageCycleManager from core.app.workflow.file_runtime import DifyWorkflowFileRuntime -from core.ops.ops_trace_manager import TraceQueueManager from core.workflow import node_factory as node_factory_module from core.workflow.node_factory import DifyNodeFactory from core.workflow.system_variables import build_system_variables @@ -79,6 +84,37 @@ from tests.workflow_test_utils import build_test_graph_init_params +class _TraceRecorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +def _trace_scope() -> TraceScope: + return TraceScope( + emitter=_TraceRecorder(), + tenant_id="tenant", + app_id="app", + user_id="user-id", + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + ) + + +def _record_trace_scope_requests(monkeypatch: pytest.MonkeyPatch, trace_scope: TraceScope) -> list[dict[str, object]]: + requests: list[dict[str, object]] = [] + + def create_trace_scope(**kwargs: object) -> TraceScope: + requests.append(kwargs) + return trace_scope + + monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.create_trace_scope", create_trace_scope) + return requests + + def _app_generator_make_app(*, app_id: str = "app", tenant_id: str = "tenant") -> App: return App( id=app_id, @@ -293,17 +329,8 @@ def test_generate_loads_conversation_and_files( SimpleNamespace(engine=sqlite_engine, session=unbound_session), ) monkeypatch.setattr(generator, "_prepare_user_inputs", lambda **kwargs: kwargs["user_inputs"]) - - dummy_trace_queue_manager = type( - "_DummyTraceQueueManager", - (TraceQueueManager,), - { - "__init__": lambda self, app_id=None, user_id=None: ( - setattr(self, "app_id", app_id) or setattr(self, "user_id", user_id) - ) - }, - ) - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.TraceQueueManager", dummy_trace_queue_manager) + trace_scope = _trace_scope() + scope_requests = _record_trace_scope_requests(monkeypatch, trace_scope) def _fake_generate(**kwargs): captured.update(kwargs) @@ -316,8 +343,11 @@ def _fake_generate(**kwargs): user = Account(name="Tester", email="tester@example.com") user.id = "user-id" + app_model = _app_generator_make_app() + app_model.tracing = '{"enabled": true, "tracing_provider": "langfuse"}' + result = generator.generate( - app_model=_app_generator_make_app(), + app_model=app_model, workflow=_app_generator_make_workflow(), user=user, args={ @@ -325,6 +355,7 @@ def _fake_generate(**kwargs): "inputs": {"k": "v"}, "conversation_id": "conversation-id", "files": [{"id": "f"}], + "trace_session_id": "session-1", }, invoke_from=InvokeFrom.WEB_APP, workflow_run_id="run-id", @@ -338,10 +369,23 @@ def _fake_generate(**kwargs): assert captured["session"] is session assert build_files_called["called"] is True assert get_conversation.call_args.kwargs["session"] is session + application_generate_entity = captured["application_generate_entity"] + assert isinstance(application_generate_entity, AdvancedChatAppGenerateEntity) + assert application_generate_entity.trace_scope is trace_scope + assert scope_requests == [ + { + "tenant_id": "tenant", + "app_id": "app", + "app_tracing": '{"enabled": true, "tracing_provider": "langfuse"}', + "user_id": "user-id", + "extras": application_generate_entity.extras, + } + ] + assert application_generate_entity.extras["trace_session_id"] == "session-1" def test_resume_delegates_to_generate(self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session): generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) - existing_trace_manager = SimpleNamespace(app_id="existing-app", user_id="existing-user") + existing_trace_scope = _trace_scope() application_generate_entity = AdvancedChatAppGenerateEntity.model_construct( task_id="task", app_config=self._build_app_config(), @@ -352,7 +396,7 @@ def test_resume_delegates_to_generate(self, monkeypatch: pytest.MonkeyPatch, unb stream=True, invoke_from=InvokeFrom.WEB_APP, extras={}, - trace_manager=existing_trace_manager, + trace_scope=existing_trace_scope, workflow_run_id="run-id", ) @@ -383,7 +427,7 @@ def _fake_generate(**kwargs): assert result.resumed is True assert captured_entity is not None - assert captured_entity.trace_manager is existing_trace_manager + assert captured_entity.trace_scope is existing_trace_scope assert captured_graph_runtime_state is not None def test_single_iteration_generate_builds_debug_task( @@ -529,7 +573,7 @@ def test_generate_internal_flow_initial_conversation_with_pause_layer( stream=False, invoke_from=InvokeFrom.WEB_APP, extras={}, - trace_manager=None, + trace_scope=None, workflow_run_id="run-id", ) @@ -657,7 +701,7 @@ def test_generate_internal_flow_with_existing_records_skips_init( stream=False, invoke_from=InvokeFrom.WEB_APP, extras={}, - trace_manager=None, + trace_scope=None, workflow_run_id="run-id", ) @@ -768,7 +812,7 @@ def test_generate_worker_raises_when_workflow_not_found( stream=False, invoke_from=InvokeFrom.WEB_APP, extras={}, - trace_manager=None, + trace_scope=None, workflow_run_id="run-id", ) @@ -817,7 +861,7 @@ def test_generate_worker_raises_when_app_not_found_for_internal_call( stream=False, invoke_from=InvokeFrom.DEBUGGER, extras={}, - trace_manager=None, + trace_scope=None, workflow_run_id="run-id", ) @@ -868,7 +912,7 @@ def test_generate_worker_handles_stopped_error( stream=False, invoke_from=InvokeFrom.DEBUGGER, extras={}, - trace_manager=None, + trace_scope=None, workflow_run_id="run-id", ) @@ -934,7 +978,7 @@ def test_generate_worker_handles_validation_error( stream=False, invoke_from=InvokeFrom.DEBUGGER, extras={}, - trace_manager=None, + trace_scope=None, workflow_run_id="run-id", ) @@ -1022,7 +1066,7 @@ def prepare(self): stream=False, invoke_from=InvokeFrom.DEBUGGER, extras={}, - trace_manager=None, + trace_scope=None, workflow_run_id="run-id", ) @@ -1078,7 +1122,7 @@ def test_handle_response_closed_file_raises_stopped(self, monkeypatch: pytest.Mo stream=False, invoke_from=InvokeFrom.WEB_APP, extras={}, - trace_manager=None, + trace_scope=None, workflow_run_id="run-id", ) @@ -1128,7 +1172,7 @@ def test_handle_response_re_raises_value_error( stream=False, invoke_from=InvokeFrom.WEB_APP, extras={}, - trace_manager=None, + trace_scope=None, workflow_run_id="run-id", ) @@ -1186,7 +1230,7 @@ def test_generate_worker_handles_invoke_auth_error( stream=False, invoke_from=InvokeFrom.SERVICE_API, extras={}, - trace_manager=None, + trace_scope=None, workflow_run_id="run-id", ) @@ -1256,19 +1300,6 @@ def test_generate_debugger_enables_retrieve_source( "core.app.apps.advanced_chat.app_generator.FileUploadConfigManager.convert", lambda features_dict, is_vision=False: None, ) - dummy_trace_queue_manager = type( - "_DummyTraceQueueManager", - (TraceQueueManager,), - { - "__init__": lambda self, app_id=None, user_id=None: ( - setattr(self, "app_id", app_id) or setattr(self, "user_id", user_id) - ) - }, - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.TraceQueueManager", - dummy_trace_queue_manager, - ) monkeypatch.setattr( "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", lambda **kwargs: SimpleNamespace(), @@ -1331,19 +1362,6 @@ def test_generate_service_api_sets_parent_message_id( "core.app.apps.advanced_chat.app_generator.FileUploadConfigManager.convert", lambda features_dict, is_vision=False: None, ) - dummy_trace_queue_manager = type( - "_DummyTraceQueueManager", - (TraceQueueManager,), - { - "__init__": lambda self, app_id=None, user_id=None: ( - setattr(self, "app_id", app_id) or setattr(self, "user_id", user_id) - ) - }, - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.TraceQueueManager", - dummy_trace_queue_manager, - ) monkeypatch.setattr( "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", lambda **kwargs: SimpleNamespace(), @@ -1395,9 +1413,9 @@ def _build_app_config() -> WorkflowUIBasedAppConfig: workflow_id="workflow-id", ) - def test_resume_restores_trace_manager_when_missing( + def test_resume_recreates_trace_scope_when_missing( self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session - ): + ) -> None: generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) application_generate_entity = AdvancedChatAppGenerateEntity.model_construct( task_id="task", @@ -1411,23 +1429,12 @@ def test_resume_restores_trace_manager_when_missing( user_id="user", stream=False, invoke_from=InvokeFrom.WEB_APP, - extras={}, - trace_manager=None, + extras={"trace_session_id": "session-1"}, + trace_scope=None, workflow_run_id="run-id", ) - dummy_trace_queue_manager = type( - "_DummyTraceQueueManager", - (TraceQueueManager,), - { - "__init__": lambda self, app_id=None, user_id=None: ( - setattr(self, "app_id", app_id) or setattr(self, "user_id", user_id) - ) - }, - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.TraceQueueManager", - dummy_trace_queue_manager, - ) + trace_scope = _trace_scope() + scope_requests = _record_trace_scope_requests(monkeypatch, trace_scope) captured_entity: AdvancedChatAppGenerateEntity | None = None def _fake_generate(**kwargs): @@ -1436,9 +1443,11 @@ def _fake_generate(**kwargs): return SimpleNamespace(ok=True) monkeypatch.setattr(generator, "_generate", _fake_generate) + app_model = _app_generator_make_app(app_id="app-id") + app_model.tracing = '{"enabled": true, "tracing_provider": "langfuse"}' result = generator.resume( - app_model=_app_generator_make_app(app_id="app-id"), + app_model=app_model, workflow=_app_generator_make_workflow(), user=_make_end_user(), conversation=_make_conversation(), @@ -1452,14 +1461,22 @@ def _fake_generate(**kwargs): assert result.ok is True assert captured_entity is not None - trace_manager = captured_entity.trace_manager - assert isinstance(trace_manager, dummy_trace_queue_manager) - assert trace_manager.app_id == "app-id" - assert trace_manager.user_id == "session-id" + assert captured_entity.trace_scope is trace_scope + assert scope_requests == [ + { + "tenant_id": "tenant", + "app_id": "app-id", + "app_tracing": '{"enabled": true, "tracing_provider": "langfuse"}', + "user_id": "session-id", + "extras": {"trace_session_id": "session-1"}, + } + ] - def test_resume_preserves_existing_trace_manager(self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session): + def test_resume_preserves_existing_trace_scope( + self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session + ) -> None: generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) - existing_trace_manager = SimpleNamespace(app_id="existing-app", user_id="existing-user") + existing_trace_scope = _trace_scope() application_generate_entity = AdvancedChatAppGenerateEntity.model_construct( task_id="task", app_config=self._build_app_config(), @@ -1473,9 +1490,10 @@ def test_resume_preserves_existing_trace_manager(self, monkeypatch: pytest.Monke stream=False, invoke_from=InvokeFrom.WEB_APP, extras={}, - trace_manager=existing_trace_manager, + trace_scope=existing_trace_scope, workflow_run_id="run-id", ) + scope_requests = _record_trace_scope_requests(monkeypatch, _trace_scope()) captured_entity: AdvancedChatAppGenerateEntity | None = None def _fake_generate(**kwargs): @@ -1500,7 +1518,8 @@ def _fake_generate(**kwargs): assert result.ok is True assert captured_entity is not None - assert captured_entity.trace_manager is existing_trace_manager + assert captured_entity.trace_scope is existing_trace_scope + assert scope_requests == [] def _make_app_config() -> WorkflowUIBasedAppConfig: @@ -1647,11 +1666,22 @@ def raise_conversation_not_exists(**_kwargs): lambda **_kwargs: app_config, ) monkeypatch.setattr(app_generator_module, "db", SimpleNamespace(engine=sqlite_engine)) - trace_manager = object.__new__(TraceQueueManager) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.TraceQueueManager", - lambda **_kwargs: trace_manager, + trace_scope = TraceScope( + emitter=_TraceRecorder(), + tenant_id="tenant-id", + app_id="app-id", + user_id="user-id", + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), ) + scope_requests: list[dict[str, object]] = [] + + def create_trace_scope(**kwargs: object) -> TraceScope: + scope_requests.append(kwargs) + return trace_scope + + monkeypatch.setattr(app_generator_module, "create_trace_scope", create_trace_scope) monkeypatch.setattr( "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", lambda **_kwargs: SimpleNamespace(), @@ -1687,6 +1717,16 @@ def fake_generate(self, **kwargs): application_generate_entity = captured["application_generate_entity"] assert isinstance(application_generate_entity, AdvancedChatAppGenerateEntity) assert application_generate_entity.conversation_id is None + assert application_generate_entity.trace_scope is trace_scope + assert scope_requests == [ + { + "tenant_id": "tenant-id", + "app_id": "app-id", + "app_tracing": None, + "user_id": "user-id", + "extras": application_generate_entity.extras, + } + ] def test_message_cycle_manager_uses_new_conversation_flag(monkeypatch: pytest.MonkeyPatch): @@ -1721,17 +1761,6 @@ def fake_thread[**P](*args: P.args, **kwargs: P.kwargs): assert entity.is_new_conversation is False -if "core.ops.ops_trace_manager" not in sys.modules: - ops_stub = ModuleType("core.ops.ops_trace_manager") - - class _StubTraceQueueManager: - def __init__(self, *_, **__): - pass - - ops_stub.TraceQueueManager = _StubTraceQueueManager - sys.modules["core.ops.ops_trace_manager"] = ops_stub - - class _StubToolNodeData(BaseNodeData): type: NodeType = BuiltinNodeTypes.TOOL pause_on: bool = False @@ -1931,6 +1960,41 @@ def _pause_resume_make_message() -> Message: ) +def _pause_resume_make_app_config(app_mode: AppMode) -> WorkflowUIBasedAppConfig: + return WorkflowUIBasedAppConfig(tenant_id="tenant", app_id="app", app_mode=app_mode, workflow_id="workflow") + + +def _make_workflow_generate_entity() -> WorkflowAppGenerateEntity: + """A deserialized paused entity: its trace scope is not persisted.""" + return WorkflowAppGenerateEntity( + task_id="task", + app_config=_pause_resume_make_app_config(AppMode.WORKFLOW), + inputs={}, + files=[], + user_id="user", + user_from=UserFrom.END_USER, + stream=False, + invoke_from=InvokeFrom.SERVICE_API, + workflow_execution_id="paused-run", + ) + + +def _make_advanced_chat_generate_entity() -> AdvancedChatAppGenerateEntity: + """A deserialized paused entity: its trace scope is not persisted.""" + return AdvancedChatAppGenerateEntity( + task_id="task", + app_config=_pause_resume_make_app_config(AppMode.ADVANCED_CHAT), + conversation_id="conv", + inputs={}, + query="query", + files=[], + user_id="user", + stream=False, + invoke_from=InvokeFrom.SERVICE_API, + workflow_run_id="adv-paused", + ) + + def _run_with_optional_pause(runtime_state: RuntimeState, *, pause_on: str | None) -> list[EngineEvent]: command_channel = InMemoryChannel() graph = _build_graph(runtime_state, pause_on=pause_on) @@ -1980,11 +2044,7 @@ def _fake_generate(**kwargs): app_model=_pause_resume_make_app(), workflow=_pause_resume_make_workflow(), user=_pause_resume_make_account(), - application_generate_entity=SimpleNamespace( - stream=False, - invoke_from=InvokeFrom.SERVICE_API, - trace_manager=SimpleNamespace(), - ), + application_generate_entity=_make_workflow_generate_entity(), graph_runtime_state=resumed_state, workflow_execution_repository=SimpleNamespace(), workflow_node_execution_repository=SimpleNamespace(), @@ -2027,11 +2087,7 @@ def _fake_generate(**kwargs): conversation=_pause_resume_make_conversation(), message=_pause_resume_make_message(), session=unbound_session, - application_generate_entity=SimpleNamespace( - stream=False, - invoke_from=InvokeFrom.SERVICE_API, - trace_manager=SimpleNamespace(), - ), + application_generate_entity=_make_advanced_chat_generate_entity(), workflow_execution_repository=SimpleNamespace(), workflow_node_execution_repository=SimpleNamespace(), graph_runtime_state=resumed_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 aec5b9f4be5947..db51ec30017348 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 @@ -1,12 +1,16 @@ from __future__ import annotations -from contextlib import nullcontext +from collections.abc import Generator +from contextlib import contextmanager, nullcontext from datetime import datetime from types import SimpleNamespace from unittest import mock import pytest +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import MessageCompleted, TraceCorrelation, TraceEvent from pytest_mock import MockerFixture +from sqlalchemy import event as sqlalchemy_event from sqlalchemy import select from sqlalchemy.orm import Session, sessionmaker @@ -74,6 +78,29 @@ from tests.workflow_test_utils import build_test_variable_pool +class _TraceRecorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +def _trace_scope(recorder: _TraceRecorder) -> TraceScope: + return TraceScope( + emitter=recorder, + tenant_id="tenant-1", + app_id="app-1", + user_id="session-1", + correlation=TraceCorrelation( + external_trace_id=None, + session_id="trace-session-1", + parent_workflow_run_id=None, + parent_node_execution_id=None, + ), + ) + + def _build_generate_entity() -> AdvancedChatAppGenerateEntity: return AdvancedChatAppGenerateEntity( task_id="task-1", @@ -90,6 +117,32 @@ def _build_generate_entity() -> AdvancedChatAppGenerateEntity: ) +def _build_message_pipeline( + application_generate_entity: AdvancedChatAppGenerateEntity, message_status: MessageStatus, answer: str +) -> pipeline_module.AdvancedChatAppGenerateTaskPipeline: + """Build the pipeline for message persistence, as the generator does.""" + user = EndUser() + user.id = "user-1" + user.session_id = "session-1" + return pipeline_module.AdvancedChatAppGenerateTaskPipeline( + application_generate_entity=application_generate_entity, + workflow=pipeline_module.WorkflowSnapshot(id="workflow-1", tenant_id="tenant-1", features_dict={}), + queue_manager=SimpleNamespace(graph_runtime_state=None), + conversation=pipeline_module.ConversationSnapshot(id="conversation-1", mode=AppMode.ADVANCED_CHAT), + message=pipeline_module.MessageSnapshot( + id="message-1", + created_at=datetime(2024, 1, 1), + query="hello", + answer=answer, + status=message_status, + ), + user=user, + stream=True, + dialogue_count=1, + draft_var_saver_factory=lambda **kwargs: None, + ) + + def _build_pipeline() -> pipeline_module.AdvancedChatAppGenerateTaskPipeline: pipeline = pipeline_module.AdvancedChatAppGenerateTaskPipeline.__new__( pipeline_module.AdvancedChatAppGenerateTaskPipeline @@ -103,7 +156,7 @@ def _build_pipeline() -> pipeline_module.AdvancedChatAppGenerateTaskPipeline: 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._application_generate_entity = _build_generate_entity() pipeline._message_cycle_manager = mock.Mock() pipeline._base_task_pipeline = SimpleNamespace(stream=True) pipeline._wrapper_process_stream_response = mock.Mock(return_value=iter(())) @@ -183,7 +236,7 @@ def test_handle_workflow_paused_event_does_not_overwrite_a_concurrent_stop( node_run_steps=0, ), ) - pipeline._save_message = mock.Mock() + pipeline._save_message = lambda **kwargs: None message = _message(status=MessageStatus.NORMAL) pipeline._get_message = mock.Mock(return_value=message) persist_content = mocker.patch.object(pipeline, "_persist_human_input_extra_content") @@ -214,43 +267,7 @@ def test_handle_workflow_paused_event_does_not_overwrite_a_concurrent_stop( def test_resume_appends_chunks_to_paused_answer(unbound_session: Session) -> None: - app_config = SimpleNamespace(app_id="app-1", tenant_id="tenant-1", sensitive_word_avoidance=None) - application_generate_entity = SimpleNamespace( - app_config=app_config, - files=[], - workflow_run_id="run-1", - query="hello", - invoke_from=InvokeFrom.WEB_APP, - inputs={}, - task_id="task-1", - trace_manager=None, - extras={}, - ) - queue_manager = SimpleNamespace(graph_runtime_state=None) - conversation = pipeline_module.ConversationSnapshot(id="conversation-1", mode=AppMode.ADVANCED_CHAT) - message = pipeline_module.MessageSnapshot( - id="message-1", - created_at=datetime(2024, 1, 1), - query="hello", - answer="before", - status=MessageStatus.PAUSED, - ) - user = EndUser() - user.id = "user-1" - user.session_id = "session-1" - workflow = pipeline_module.WorkflowSnapshot(id="workflow-1", tenant_id="tenant-1", features_dict={}) - - pipeline = pipeline_module.AdvancedChatAppGenerateTaskPipeline( - application_generate_entity=application_generate_entity, - workflow=workflow, - queue_manager=queue_manager, - conversation=conversation, - message=message, - user=user, - stream=True, - dialogue_count=1, - draft_var_saver_factory=SimpleNamespace(), - ) + pipeline = _build_message_pipeline(_build_generate_entity(), message_status=MessageStatus.PAUSED, answer="before") stored_message = _message(status=MessageStatus.PAUSED, answer="before") pipeline._get_message = mock.Mock(return_value=stored_message) @@ -263,217 +280,48 @@ def test_resume_appends_chunks_to_paused_answer(unbound_session: Session) -> Non assert stored_message.status == MessageStatus.NORMAL -def test_emit_message_trace_enqueues_message_trace() -> None: - """When trace_manager is provided, _emit_message_trace must enqueue a MESSAGE_TRACE task.""" - from core.ops.entities.trace_entity import TraceTaskName - - mock_trace_manager = mock.MagicMock() - app_config = SimpleNamespace(app_id="app-1", tenant_id="tenant-1", sensitive_word_avoidance=None) - application_generate_entity = SimpleNamespace( - app_config=app_config, - files=[], - workflow_run_id="run-1", - query="hello", - invoke_from=InvokeFrom.WEB_APP, - inputs={}, - task_id="task-1", - trace_manager=mock_trace_manager, - extras={"trace_session_id": "session-trace-1"}, - ) - queue_manager = SimpleNamespace(graph_runtime_state=None) - conversation = pipeline_module.ConversationSnapshot(id="conv-1", mode=AppMode.ADVANCED_CHAT) - message = pipeline_module.MessageSnapshot( - id="msg-1", - created_at=datetime(2024, 1, 1), - query="hello", - answer="", - status=MessageStatus.NORMAL, - ) - user = EndUser() - user.id = "user-1" - user.session_id = "session-1" - workflow = pipeline_module.WorkflowSnapshot(id="workflow-1", tenant_id="tenant-1", features_dict={}) - - pipeline = pipeline_module.AdvancedChatAppGenerateTaskPipeline( - application_generate_entity=application_generate_entity, - workflow=workflow, - queue_manager=queue_manager, - conversation=conversation, - message=message, - user=user, - stream=True, - dialogue_count=1, - draft_var_saver_factory=SimpleNamespace(), - ) - - stored_message = SimpleNamespace( - id="msg-1", +def test_emit_message_trace_reports_the_completed_message() -> None: + recorder = _TraceRecorder() + pipeline = _build_message_pipeline( + _build_generate_entity().model_copy(update={"trace_scope": _trace_scope(recorder)}), + message_status=MessageStatus.NORMAL, answer="", - status=MessageStatus.NORMAL, - updated_at=None, - provider_response_latency=0, - message_tokens=0, - message_unit_price=0, - message_price_unit=0, - answer_tokens=0, - answer_unit_price=0, - answer_price_unit=0, - total_price=0, - currency="USD", - message_metadata=None, - invoke_from=InvokeFrom.WEB_APP, - from_account_id=None, - from_end_user_id="user-1", ) - pipeline._get_message = mock.Mock(return_value=stored_message) - pipeline._recorded_files = [] pipeline._emit_message_trace() - mock_trace_manager.add_trace_task.assert_called_once() - trace_task = mock_trace_manager.add_trace_task.call_args[0][0] - assert trace_task.trace_type == TraceTaskName.MESSAGE_TRACE - assert trace_task.conversation_id == "conv-1" - assert trace_task.message_id == "msg-1" - assert trace_task.kwargs.get("trace_session_id") == "session-trace-1" + assert len(recorder.events) == 1 + event = recorder.events[0] + assert event.payload == MessageCompleted(message_id="message-1", conversation_id="conversation-1") + assert event.correlation.session_id == "trace-session-1" + assert (event.tenant_id, event.app_id, event.user_id) == ("tenant-1", "app-1", "session-1") -def test_save_message_does_not_emit_trace() -> None: - """_save_message must not enqueue a trace; emitting is the terminal caller's responsibility. +def test_save_message_does_not_emit_trace(unbound_session: Session) -> None: + """_save_message must not emit a trace; emitting is the terminal caller's responsibility. This keeps intermediate saves (e.g. on pause) from duplicating message, token, and TTFT metrics when the message is later resumed. """ - mock_trace_manager = mock.MagicMock() - app_config = SimpleNamespace(app_id="app-1", tenant_id="tenant-1", sensitive_word_avoidance=None) - application_generate_entity = SimpleNamespace( - app_config=app_config, - files=[], - workflow_run_id="run-1", - query="hello", - invoke_from=InvokeFrom.WEB_APP, - inputs={}, - task_id="task-1", - trace_manager=mock_trace_manager, - extras={"trace_session_id": "session-trace-1"}, - ) - queue_manager = SimpleNamespace(graph_runtime_state=None) - conversation = pipeline_module.ConversationSnapshot(id="conv-1", mode=AppMode.ADVANCED_CHAT) - message = pipeline_module.MessageSnapshot( - id="msg-1", - created_at=datetime(2024, 1, 1), - query="hello", + recorder = _TraceRecorder() + pipeline = _build_message_pipeline( + _build_generate_entity().model_copy(update={"trace_scope": _trace_scope(recorder)}), + message_status=MessageStatus.NORMAL, answer="", - status=MessageStatus.NORMAL, - ) - user = EndUser() - user.id = "user-1" - user.session_id = "session-1" - workflow = pipeline_module.WorkflowSnapshot(id="workflow-1", tenant_id="tenant-1", features_dict={}) - - pipeline = pipeline_module.AdvancedChatAppGenerateTaskPipeline( - application_generate_entity=application_generate_entity, - workflow=workflow, - queue_manager=queue_manager, - conversation=conversation, - message=message, - user=user, - stream=True, - dialogue_count=1, - draft_var_saver_factory=SimpleNamespace(), - ) - - stored_message = SimpleNamespace( - id="msg-1", - answer="", - status=MessageStatus.NORMAL, - updated_at=None, - provider_response_latency=0, - message_tokens=0, - message_unit_price=0, - message_price_unit=0, - answer_tokens=0, - answer_unit_price=0, - answer_price_unit=0, - total_price=0, - currency="USD", - message_metadata=None, - invoke_from=InvokeFrom.WEB_APP, - from_account_id=None, - from_end_user_id="user-1", ) + stored_message = _message(status=MessageStatus.NORMAL) pipeline._get_message = mock.Mock(return_value=stored_message) pipeline._recorded_files = [] - pipeline._save_message(session=mock.Mock()) - - mock_trace_manager.add_trace_task.assert_not_called() + pipeline._save_message(session=unbound_session) + assert stored_message.status == MessageStatus.NORMAL + assert recorder.events == [] -def test_emit_message_trace_skips_when_no_trace_manager() -> None: - """When trace_manager is None, _emit_message_trace and _save_message must not fail.""" - app_config = SimpleNamespace(app_id="app-1", tenant_id="tenant-1", sensitive_word_avoidance=None) - application_generate_entity = SimpleNamespace( - app_config=app_config, - files=[], - workflow_run_id="run-1", - query="hello", - invoke_from=InvokeFrom.WEB_APP, - inputs={}, - task_id="task-1", - trace_manager=None, - extras={}, - ) - queue_manager = SimpleNamespace(graph_runtime_state=None) - conversation = pipeline_module.ConversationSnapshot(id="conv-1", mode=AppMode.ADVANCED_CHAT) - message = pipeline_module.MessageSnapshot( - id="msg-1", - created_at=datetime(2024, 1, 1), - query="hello", - answer="", - status=MessageStatus.NORMAL, - ) - user = EndUser() - user.id = "user-1" - user.session_id = "session-1" - workflow = pipeline_module.WorkflowSnapshot(id="workflow-1", tenant_id="tenant-1", features_dict={}) - pipeline = pipeline_module.AdvancedChatAppGenerateTaskPipeline( - application_generate_entity=application_generate_entity, - workflow=workflow, - queue_manager=queue_manager, - conversation=conversation, - message=message, - user=user, - stream=True, - dialogue_count=1, - draft_var_saver_factory=SimpleNamespace(), - ) +def test_emit_message_trace_skips_without_trace_scope() -> None: + pipeline = _build_message_pipeline(_build_generate_entity(), message_status=MessageStatus.NORMAL, answer="") - stored_message = SimpleNamespace( - id="msg-1", - answer="", - status=MessageStatus.NORMAL, - updated_at=None, - provider_response_latency=0, - message_tokens=0, - message_unit_price=0, - message_price_unit=0, - answer_tokens=0, - answer_unit_price=0, - answer_price_unit=0, - total_price=0, - currency="USD", - message_metadata=None, - invoke_from=InvokeFrom.WEB_APP, - from_account_id=None, - from_end_user_id="user-1", - ) - pipeline._get_message = mock.Mock(return_value=stored_message) - pipeline._recorded_files = [] - - # Should not raise - pipeline._save_message(session=mock.Mock()) pipeline._emit_message_trace() @@ -542,7 +390,7 @@ def test_process_stream_response_breaks_after_workflow_succeeded() -> None: responses = list(pipeline._process_stream_response()) assert [resp.event for resp in responses] == [StreamEvent.WORKFLOW_FINISHED] - pipeline._handle_workflow_succeeded_event.assert_called_once_with(succeeded_event, trace_manager=None) + pipeline._handle_workflow_succeeded_event.assert_called_once_with(succeeded_event) pipeline._base_task_pipeline.ping_stream_response.assert_not_called() @@ -567,10 +415,46 @@ def test_process_stream_response_breaks_after_workflow_partial_success() -> None responses = list(pipeline._process_stream_response()) assert [resp.event for resp in responses] == [StreamEvent.WORKFLOW_FINISHED] - pipeline._handle_workflow_partial_success_event.assert_called_once_with(partial_event, trace_manager=None) + pipeline._handle_workflow_partial_success_event.assert_called_once_with(partial_event) pipeline._base_task_pipeline.ping_stream_response.assert_not_called() +class _TimelineTraceRecorder: + """Record trace events on a timeline shared with the commits they must follow.""" + + def __init__(self, timeline: list[str]) -> None: + self._timeline = timeline + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + self._timeline.append("trace") + + +def _timeline_trace_scope(recorder: _TimelineTraceRecorder) -> TraceScope: + return TraceScope( + emitter=recorder, + tenant_id="tenant", + app_id="app", + user_id="session", + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + ) + + +@contextmanager +def _recording_commits(timeline: list[str]) -> Generator[None]: + def record_commit(_session: Session) -> None: + timeline.append("commit") + + sqlalchemy_event.listen(Session, "after_commit", record_commit) + try: + yield + finally: + sqlalchemy_event.remove(Session, "after_commit", record_commit) + + def _make_pipeline(): app_config = WorkflowUIBasedAppConfig( tenant_id="tenant", @@ -590,7 +474,7 @@ def _make_pipeline(): stream=False, invoke_from=InvokeFrom.WEB_APP, extras={}, - trace_manager=None, + trace_scope=None, workflow_run_id="run-id", ) @@ -1227,6 +1111,60 @@ def _save_message(**kwargs): assert responses == ["replace", "end"] assert saved == ["saved"] + @pytest.mark.parametrize( + ("saved_on_pause", "expected_timeline", "expected_payloads"), + [ + ( + False, + ["save", "commit", "trace"], + [MessageCompleted(message_id="message-id", conversation_id="conv-id")], + ), + # The resumed run saves and reports the message instead. + (True, [], []), + ], + ) + def test_handle_message_end_event_reports_message_after_commit( + self, saved_on_pause: bool, expected_timeline: list[str], expected_payloads: list[MessageCompleted] + ) -> None: + timeline: list[str] = [] + recorder = _TimelineTraceRecorder(timeline) + pipeline = _make_pipeline() + pipeline._application_generate_entity.trace_scope = _timeline_trace_scope(recorder) + pipeline._message_saved_on_pause = saved_on_pause + pipeline._graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap( + system_variables=build_system_variables(workflow_execution_id="run-id") + ), + start_at=0.0, + ) + pipeline._save_message = lambda **kwargs: timeline.append("save") + + with _recording_commits(timeline): + responses = list(pipeline._handle_advanced_chat_message_end_event(QueueAdvancedChatMessageEndEvent())) + + assert [type(response) for response in responses] == [MessageEndStreamResponse] + assert timeline == expected_timeline + assert [event.payload for event in recorder.events] == expected_payloads + + def test_handle_stop_event_reports_moderated_message_after_commit(self) -> None: + timeline: list[str] = [] + recorder = _TimelineTraceRecorder(timeline) + pipeline = _make_pipeline() + pipeline._application_generate_entity.trace_scope = _timeline_trace_scope(recorder) + pipeline._save_message = lambda **kwargs: timeline.append("save") + + with _recording_commits(timeline): + responses = list( + pipeline._handle_stop_event(QueueStopEvent(stopped_by=QueueStopEvent.StopBy.INPUT_MODERATION)) + ) + + assert [type(response) for response in responses] == [MessageEndStreamResponse] + assert timeline == ["save", "commit", "trace"] + assert [event.payload for event in recorder.events] == [ + MessageCompleted(message_id="message-id", conversation_id="conv-id") + ] + def test_dispatch_event_handles_node_exception(self): pipeline = _make_pipeline() pipeline._workflow_response_converter.workflow_node_finish_to_stream_response = lambda **kwargs: "failed" diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py b/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py index d7bd1f2d32f07f..5c426513035a6b 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py @@ -8,6 +8,8 @@ from unittest.mock import MagicMock import pytest +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import TraceCorrelation, TraceEvent from pytest_mock import MockerFixture from sqlalchemy import event, select from sqlalchemy.orm import Session @@ -15,14 +17,16 @@ import core.app.apps.agent_app.app_generator as module import core.app.features.annotation_reply.annotation_reply as annotation_mod import core.moderation.input_moderation as input_moderation_mod +from core.app.app_config.entities import EasyUIBasedAppModelConfigFrom, ModelConfigEntity, PromptTemplateEntity from core.app.apps.agent_app import app_generator +from core.app.apps.agent_app.app_config_manager import AgentAppConfig from core.app.apps.agent_app.app_generator import ( AgentAppGenerator, AgentAppGeneratorError, AgentAppNotPublishedError, ) from core.app.apps.exc import GenerateTaskStoppedError -from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom +from core.app.entities.app_invoke_entities import AgentAppGenerateEntity, InvokeFrom, UserFrom from core.app.entities.queue_entities import ( QueueAnnotationReplyEvent, QueueLLMChunkEvent, @@ -51,6 +55,7 @@ AgentWorkspaceBindingGenerationMismatchError, AgentWorkspaceService, ) +from tests.helpers.unit.core.model_fixtures import make_model_config from tests.helpers.unit.model_factories import make_account, make_app, make_conversation, make_message """Unit tests for AgentAppGenerator.generate() and its worker thread. @@ -66,6 +71,37 @@ MODULE = "core.app.apps.agent_app.app_generator" +class _TraceRecorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +def _trace_scope() -> TraceScope: + return TraceScope( + emitter=_TraceRecorder(), + tenant_id="tenant", + app_id="app1", + user_id="user", + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + ) + + +def _record_trace_scope_requests(monkeypatch: pytest.MonkeyPatch, trace_scope: TraceScope) -> list[dict[str, object]]: + requests: list[dict[str, object]] = [] + + def create_trace_scope(**kwargs: object) -> TraceScope: + requests.append(kwargs) + return trace_scope + + monkeypatch.setattr(module, "create_trace_scope", create_trace_scope) + return requests + + def _app_generator_account(user_id: str = "user") -> _app_generator_Account: account = _app_generator_Account(name="User", email=f"{user_id}@example.com") account.id = user_id @@ -217,7 +253,11 @@ def test_session_scope_config_version_id_preserves_draft_or_snapshot_id(self): == "snapshot-1" ) - def test_generate_orchestrates_and_starts_worker(self, generator, mocker: MockerFixture): + def test_generate_orchestrates_and_starts_worker( + self, generator: AgentAppGenerator, mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch + ) -> None: + trace_scope = _trace_scope() + scope_requests = _record_trace_scope_requests(monkeypatch, trace_scope) config = AppModelConfig(app_id="app1") config.id = "config-1" session = _session() @@ -240,7 +280,6 @@ def test_generate_orchestrates_and_starts_worker(self, generator, mocker: Mocker mocker.patch(f"{MODULE}.FileUploadConfigManager.convert", return_value=file_upload_config) parsed_file = mocker.MagicMock() build_files = mocker.patch(f"{MODULE}.file_factory.build_from_mappings", return_value=[parsed_file]) - mocker.patch(f"{MODULE}.TraceQueueManager", return_value=mocker.MagicMock()) generate_entity = mocker.patch( f"{MODULE}.AgentAppGenerateEntity", return_value=mocker.MagicMock(task_id="t", user_id="user") ) @@ -284,6 +323,11 @@ def test_generate_orchestrates_and_starts_worker(self, generator, mocker: Mocker assert generate_entity.call_args.kwargs["files"] == [parsed_file] assert generate_entity.call_args.kwargs["file_upload_config"] is file_upload_config assert "agent_runtime_exit_intent" not in generate_entity.call_args.kwargs + extras = generate_entity.call_args.kwargs["extras"] + assert generate_entity.call_args.kwargs["trace_scope"] is trace_scope + assert scope_requests == [ + {"tenant_id": "tenant", "app_id": "app1", "app_tracing": None, "user_id": "user", "extras": extras} + ] def test_generate_loads_existing_conversation(self, generator: AgentAppGenerator, mocker: MockerFixture): app_model = _app_generator_app() @@ -297,7 +341,6 @@ def test_generate_loads_existing_conversation(self, generator: AgentAppGenerator mocker.patch(f"{MODULE}.AgentAppConfigManager.get_app_config", return_value=mocker.MagicMock(variables=[])) mocker.patch(f"{MODULE}.load_annotation_reply_config", return_value={"enabled": False}) mocker.patch(f"{MODULE}.ModelConfigConverter.convert", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.TraceQueueManager", return_value=mocker.MagicMock()) mocker.patch(f"{MODULE}.AgentAppGenerateEntity", return_value=mocker.MagicMock()) mocker.patch(f"{MODULE}.MessageBasedAppQueueManager", return_value=mocker.MagicMock()) mocker.patch(f"{MODULE}.threading.Thread", return_value=mocker.MagicMock()) @@ -324,8 +367,9 @@ def test_generate_loads_existing_conversation(self, generator: AgentAppGenerator assert generator._init_generate_records.call_args.kwargs["session"] is session def test_generate_does_not_include_trace_session_id_in_extras( - self, generator: AgentAppGenerator, mocker: MockerFixture - ): + self, generator: AgentAppGenerator, mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch + ) -> None: + scope_requests = _record_trace_scope_requests(monkeypatch, _trace_scope()) app_model = _app_generator_app() user = _app_generator_account() @@ -339,7 +383,6 @@ def test_generate_does_not_include_trace_session_id_in_extras( return_value=mocker.MagicMock(variables=[], tenant_id="tenant", app_id="app1"), ) mocker.patch(f"{MODULE}.ModelConfigConverter.convert", return_value=mocker.MagicMock(model="gpt-4o-mini")) - mocker.patch(f"{MODULE}.TraceQueueManager", return_value=mocker.MagicMock()) generate_entity = mocker.patch( f"{MODULE}.AgentAppGenerateEntity", return_value=mocker.MagicMock(task_id="t", user_id="user") ) @@ -357,6 +400,7 @@ def test_generate_does_not_include_trace_session_id_in_extras( ) assert generate_entity.call_args.kwargs["extras"] == {"auto_generate_conversation_name": True} + assert [request["extras"] for request in scope_requests] == [{"auto_generate_conversation_name": True}] class TestGenerateWorker: @pytest.fixture(autouse=True) @@ -506,14 +550,17 @@ def publish(event, *_args): assert events == ["commit", "publish"] runner.run.assert_not_called() - def test_generate_task_stopped_is_swallowed(self, generator, mocker: MockerFixture): + def test_generate_task_stopped_is_swallowed( + self, generator: AgentAppGenerator, mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch + ) -> None: self._wire(generator, mocker, run_side_effect=GenerateTaskStoppedError()) + # A locally stopped consumer does not prove runtime cancellation. The # runner must hand off the typed terminal it actually observed. - mocker.patch( - "core.ops.base_trace_instance.application_tracing_hooks", - side_effect=AssertionError("stop intent must not invent a tracing terminal"), - ) + def application_tracing_hooks() -> None: + raise AssertionError("stop intent must not invent a tracing terminal") + + monkeypatch.setattr(module, "application_tracing_hooks", application_tracing_hooks) queue_manager = mocker.MagicMock() self._call(generator, mocker, queue_manager) queue_manager.publish_error.assert_not_called() @@ -542,15 +589,36 @@ def publish(self, event: Any, _from: Any) -> None: self.events.append(event) -def _make_entity(query: str = "hello") -> SimpleNamespace: - return SimpleNamespace( - app_config=SimpleNamespace(app_id="app-1", tenant_id="tenant-1"), - model_conf=SimpleNamespace(model="gpt-4o-mini"), +def _make_entity(query: str) -> AgentAppGenerateEntity: + return AgentAppGenerateEntity( + task_id="task-1", + app_config=AgentAppConfig( + tenant_id="tenant-1", + app_id="app-1", + app_mode=AppMode.AGENT, + app_model_config_from=EasyUIBasedAppModelConfigFrom.APP_LATEST_CONFIG, + app_model_config_dict={}, + model=ModelConfigEntity(provider="openai", model="gpt-4o-mini", mode="chat"), + prompt_template=PromptTemplateEntity(prompt_type=PromptTemplateEntity.PromptType.SIMPLE), + ), + model_conf=make_model_config(provider="openai", model="gpt-4o-mini", mode="chat"), inputs={}, + files=[], query=query, - trace_manager=None, user_id="user-1", - invoke_from=SimpleNamespace(), + stream=True, + invoke_from=InvokeFrom.WEB_APP, + trace_scope=TraceScope( + emitter=_TraceRecorder(), + tenant_id="tenant-1", + app_id="app-1", + user_id="user-1", + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + ), + agent_id="agent-1", + agent_config_snapshot_id="snapshot-1", ) @@ -562,14 +630,23 @@ def _input_guards_message() -> Message: return make_message(message_id="msg-1") -def _patch_moderation(monkeypatch: pytest.MonkeyPatch, *, returns=None, raises: Exception | None = None) -> None: +def _patch_moderation( + monkeypatch: pytest.MonkeyPatch, + *, + returns: tuple[bool, dict[str, Any], str] | None = None, + raises: Exception | None = None, +) -> list[dict[str, Any]]: + calls: list[dict[str, Any]] = [] + class _FakeModeration: def check(self, **kwargs: Any): + calls.append(kwargs) if raises is not None: raise raises return returns monkeypatch.setattr(input_moderation_mod, "InputModeration", _FakeModeration) + return calls def _patch_annotation(monkeypatch: pytest.MonkeyPatch, *, reply=None) -> None: @@ -595,13 +672,14 @@ def _saved_user_query(events: list[Any]) -> str: class TestInputGuardsCases: class TestRunInputGuards: def test_no_guards_passes_through(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): - _patch_moderation(monkeypatch, returns=(False, {}, "hello")) + moderation_calls = _patch_moderation(monkeypatch, returns=(False, {}, "hello")) _patch_annotation(monkeypatch, reply=None) qm = _FakeQueueManager() + entity = _make_entity("hello") handled, query, annotation_reply = AgentAppGenerator()._run_input_guards( session=sqlite_session, - application_generate_entity=_make_entity("hello"), + application_generate_entity=entity, app_model=_input_guards_app(), message=_input_guards_message(), queue_manager=qm, @@ -611,6 +689,9 @@ def test_no_guards_passes_through(self, monkeypatch: pytest.MonkeyPatch, sqlite_ assert query == "hello" assert annotation_reply is None assert qm.events == [] + assert len(moderation_calls) == 1 + assert moderation_calls[0]["message_id"] == "msg-1" + assert moderation_calls[0]["trace_scope"] is entity.trace_scope def test_moderation_override_sanitizes_query(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): _patch_moderation(monkeypatch, returns=(True, {}, "[redacted]")) diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py b/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py index a9b01d9aca1b53..8a914d0002021e 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py @@ -16,6 +16,7 @@ from dify_agent.layers.user_prompt import DifyUserPromptLayerConfig from dify_agent.layers.user_prompt import layer as user_prompt_layer from dify_agent.protocol.snapshot import SessionSnapshot +from dify_trace_sdk.events import TraceEvent from flask import Flask from pydantic_ai.messages import BinaryContent, ImageUrl @@ -34,16 +35,16 @@ AgentAppRuntimeRequestBuildError, ) from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom -from core.ops.base_trace_instance import bind_application_tracing_hooks from core.workflow.file_reference import build_file_reference from core.workflow.nodes.agent_v2.dify_tools_builder import WorkflowAgentDifyToolsBuilder +from extensions.tracing import bind_tracing from graphon.file import File, FileTransferMethod, FileType from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent from graphon.model_runtime.entities.model_entities import ModelFeature from models.agent_config_entities import AgentSoulConfig -from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks from tests.helpers.unit.config_override import apply_config_overrides from tests.helpers.unit.knowledge_fs_core_fakes import RecordingModelFactory, llm_schema +from tests.helpers.unit.services.tracing_ports import NoTracingPolicies, RecordingTracingHooks, kernel_trace_scopes @dataclass(frozen=True) @@ -53,9 +54,12 @@ class _LayerDeps: layers: dict[str, dict[str, object]] -@pytest.fixture(autouse=True) -def _application_tracing_disabled(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) +class _TraceRecorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) @pytest.fixture(autouse=True) @@ -133,7 +137,7 @@ def test_soul_prompt_optional(self): assert [layer.name for layer in request.composition.layers][0] == "agent_app_user_prompt" -def test_builder_carries_application_trace_context_but_redacts_it_from_request_logs(monkeypatch: pytest.MonkeyPatch): +def test_builder_carries_application_trace_context_but_redacts_it_from_request_logs(): from uuid import uuid4 from dify_agent.layers.execution_context import ApplicationTracingContext @@ -146,18 +150,14 @@ def test_builder_carries_application_trace_context_but_redacts_it_from_request_l session_id="private-session", capture_content=False, ) - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) context = _ctx(_soul_with_model()) - calls: list[AgentAppRuntimeBuildContext] = [] - - def prepare(value: AgentAppRuntimeBuildContext) -> ApplicationTracingContext: - calls.append(value) - return tracing - - monkeypatch.setattr(RuntimeApplicationTracingHooks, "prepare_agent_app", staticmethod(prepare)) - hooks = RuntimeApplicationTracingHooks() + hooks = RecordingTracingHooks(tracing) app = Flask(__name__) - bind_application_tracing_hooks(app, hooks) + bind_tracing( + app, + scopes=kernel_trace_scopes(_TraceRecorder(), policies=NoTracingPolicies(), enterprise_enabled=False), + hooks=hooks, + ) with app.app_context(): result = AgentAppRuntimeRequestBuilder(dify_tools_builder=WorkflowAgentDifyToolsBuilder()).build(context) @@ -167,7 +167,7 @@ def prepare(value: AgentAppRuntimeBuildContext) -> ApplicationTracingContext: item for item in result.redacted_request["composition"]["layers"] if item["name"] == "execution_context" ) assert redacted["config"]["tracing"] == "[REDACTED]" - assert calls == [context] + assert hooks.agent_apps == [context] class _NoToolsBuilder: diff --git a/api/tests/unit_tests/core/app/apps/agent_chat/test_app_generator.py b/api/tests/unit_tests/core/app/apps/agent_chat/test_app_generator.py index e8d66310c26b4a..2bb1fdc1caaf7a 100644 --- a/api/tests/unit_tests/core/app/apps/agent_chat/test_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/agent_chat/test_app_generator.py @@ -5,6 +5,8 @@ from dataclasses import dataclass, field import pytest +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import TraceCorrelation, TraceEvent from flask import Flask from pydantic import ValidationError from sqlalchemy.orm import Session @@ -19,7 +21,6 @@ ) from core.app.entities.queue_entities import QueueErrorEvent from core.app.file_access import get_current_file_access_scope -from core.ops.ops_trace_manager import TraceQueueManager from graphon.model_runtime.errors.invoke import InvokeAuthorizationError from models import Account from models.model import App, AppMode, AppModelConfig, Conversation, Message @@ -45,6 +46,26 @@ def __call__(self, *args: object, **kwargs: object) -> T: return self.result +class _TraceRecorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +def _trace_scope() -> TraceScope: + return TraceScope( + emitter=_TraceRecorder(), + tenant_id="tenant", + app_id="app1", + user_id="user", + correlation=TraceCorrelation( + external_trace_id=None, session_id="session-1", parent_workflow_run_id=None, parent_node_execution_id=None + ), + ) + + def _app() -> App: return App( id="app1", @@ -68,13 +89,10 @@ def generator( agent_runtime_app: Flask, agent_queue_manager: MessageBasedAppQueueManager, agent_records: tuple[Conversation, Message], - monkeypatch: pytest.MonkeyPatch, ) -> AgentChatAppGenerator: assert agent_runtime_app.name assert agent_queue_manager is not None assert agent_records[0].id == "conv" - # Keep trace configuration lookup and initialization real, but do not schedule background timers. - monkeypatch.setattr(TraceQueueManager, "start_timer", lambda _manager: None) return AgentChatAppGenerator(execution_driver=WorkflowRunAgg.run) @@ -226,6 +244,9 @@ def observe_runner(_runner: AgentChatAppRunner, **kwargs: object) -> None: "core.app.apps.agent_chat.app_generator.AgentChatAppGenerateResponseConverter.convert", CallRecorder({"result": "ok"}), ) + trace_scope = _trace_scope() + create_trace_scope = CallRecorder(trace_scope) + monkeypatch.setattr("core.app.apps.agent_chat.app_generator.create_trace_scope", create_trace_scope) args = { "query": "hello", @@ -259,10 +280,20 @@ def observe_runner(_runner: AgentChatAppRunner, **kwargs: object) -> None: assert len(runner_calls) == 1 entity = runner_calls[0]["application_generate_entity"] assert isinstance(entity, AgentChatAppGenerateEntity) - assert entity.extras["trace_session_id"] == "session-1" - assert type(entity.trace_manager) is TraceQueueManager - assert entity.trace_manager.trace_instance is None - assert entity.trace_manager.app_id == "app1" + assert entity.extras == {"auto_generate_conversation_name": True, "trace_session_id": "session-1"} + assert entity.trace_scope is trace_scope + assert create_trace_scope.calls == [ + RecordedCall( + args=(), + kwargs={ + "tenant_id": "tenant", + "app_id": "app1", + "app_tracing": None, + "user_id": "user", + "extras": entity.extras, + }, + ) + ] assert type(runner_calls[0]["queue_manager"]) is MessageBasedAppQueueManager conversation = runner_calls[0]["conversation"] message = runner_calls[0]["message"] diff --git a/api/tests/unit_tests/core/app/apps/chat/test_app_generator.py b/api/tests/unit_tests/core/app/apps/chat/test_app_generator.py index cf438c6824c8b2..bfaf28a79b63a9 100644 --- a/api/tests/unit_tests/core/app/apps/chat/test_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/chat/test_app_generator.py @@ -3,6 +3,8 @@ from unittest.mock import ANY, MagicMock, Mock, patch import pytest +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import TraceCorrelation, TraceEvent from sqlalchemy import event from sqlalchemy.orm import Session @@ -24,6 +26,26 @@ def __init__(self, **kwargs): self.__dict__.update(kwargs) +class _TraceRecorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +def _trace_scope() -> TraceScope: + return TraceScope( + emitter=_TraceRecorder(), + tenant_id="tenant-1", + app_id="app-1", + user_id="user-1", + correlation=TraceCorrelation( + external_trace_id=None, session_id="session-1", parent_workflow_run_id=None, parent_node_execution_id=None + ), + ) + + class DummyQueueManager: def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): self.published = [] @@ -143,10 +165,21 @@ def test_generate_rejects_non_string_query(self, unbound_session: Session): streaming=False, ) - def test_generate_debugger_overrides_model_config(self, unbound_session: Session): + def test_generate_debugger_overrides_model_config( + self, unbound_session: Session, monkeypatch: pytest.MonkeyPatch + ) -> None: generator = ChatAppGenerator() app_model = _app() + app_model.tracing = '{"enabled": true, "tracing_provider": "langfuse"}' user = _account() + trace_scope = _trace_scope() + scope_requests: list[dict[str, object]] = [] + + def create_trace_scope(**kwargs: object) -> TraceScope: + scope_requests.append(kwargs) + return trace_scope + + monkeypatch.setattr(generator_module, "create_trace_scope", create_trace_scope) args = { "query": "hi", "inputs": {}, @@ -173,7 +206,6 @@ def test_generate_debugger_overrides_model_config(self, unbound_session: Session "core.app.apps.chat.app_generator.ChatAppGenerateEntity", Mock(side_effect=DummyGenerateEntity), ) as generate_entity, - patch("core.app.apps.chat.app_generator.TraceQueueManager", return_value=SimpleNamespace()), patch("core.app.apps.chat.app_generator.MessageBasedAppQueueManager", DummyQueueManager), patch( "core.app.apps.chat.app_generator.ChatAppGenerateResponseConverter.convert", return_value={"ok": True} @@ -196,7 +228,18 @@ def test_generate_debugger_overrides_model_config(self, unbound_session: Session assert result == {"ok": True} assert get_conversation.call_args.kwargs["session"] is unbound_session - assert generate_entity.call_args.kwargs["extras"]["trace_session_id"] == "session-1" + extras = generate_entity.call_args.kwargs["extras"] + assert extras == {"auto_generate_conversation_name": True, "trace_session_id": "session-1"} + assert generate_entity.call_args.kwargs["trace_scope"] is trace_scope + assert scope_requests == [ + { + "tenant_id": "tenant-1", + "app_id": "app-1", + "app_tracing": '{"enabled": true, "tracing_provider": "langfuse"}', + "user_id": "user-1", + "extras": extras, + } + ] def test_generate_uses_session_for_annotation_reply(self, unbound_session: Session): generator = ChatAppGenerator() diff --git a/api/tests/unit_tests/core/app/apps/completion/test_app_generator.py b/api/tests/unit_tests/core/app/apps/completion/test_app_generator.py index 1a6dacff8e875f..e648c14ab2aece 100644 --- a/api/tests/unit_tests/core/app/apps/completion/test_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/completion/test_app_generator.py @@ -4,6 +4,8 @@ from unittest.mock import MagicMock import pytest +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import TraceCorrelation, TraceEvent from pydantic import ValidationError from pytest_mock import MockerFixture from sqlalchemy.orm import Session @@ -20,6 +22,26 @@ TABLES = (AppModelConfig, Conversation, Message) +class _TraceRecorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +def _trace_scope() -> TraceScope: + return TraceScope( + emitter=_TraceRecorder(), + tenant_id="tenant", + app_id="app1", + user_id="user", + correlation=TraceCorrelation( + external_trace_id=None, session_id="session-1", parent_workflow_run_id=None, parent_node_execution_id=None + ), + ) + + @pytest.fixture def generator(mocker: MockerFixture): gen = CompletionAppGenerator() @@ -34,7 +56,6 @@ def generator(mocker: MockerFixture): mocker.patch.object(module.threading, "Thread", return_value=thread) mocker.patch.object(module, "MessageBasedAppQueueManager", return_value=MagicMock()) - mocker.patch.object(module, "TraceQueueManager", return_value=MagicMock()) generate_entity = mocker.patch.object( module, "CompletionAppGenerateEntity", side_effect=lambda **kwargs: SimpleNamespace(**kwargs) ) @@ -154,7 +175,17 @@ def test_generate_override_not_debugger(self, generator, sqlite_session: Session streaming=False, ) - def test_generate_success_no_file_config(self, generator, mocker: MockerFixture, sqlite_session: Session): + def test_generate_success_no_file_config( + self, generator, mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ) -> None: + trace_scope = _trace_scope() + scope_requests: list[dict[str, object]] = [] + + def create_trace_scope(**kwargs: object) -> TraceScope: + scope_requests.append(kwargs) + return trace_scope + + monkeypatch.setattr(module, "create_trace_scope", create_trace_scope) app_model_config = _build_app_model_config() mocker.patch.object(generator, "_get_app_model_config", return_value=app_model_config) to_dict = mocker.patch.object(AppModelConfig, "to_dict", return_value={"model": {"provider": "x"}}) @@ -213,7 +244,12 @@ def test_generate_success_no_file_config(self, generator, mocker: MockerFixture, ) assert result == "converted" - assert generator.generate_entity.call_args.kwargs["extras"]["trace_session_id"] == "session-1" + extras = generator.generate_entity.call_args.kwargs["extras"] + assert extras == {"trace_session_id": "session-1"} + assert generator.generate_entity.call_args.kwargs["trace_scope"] is trace_scope + assert scope_requests == [ + {"tenant_id": "tenant", "app_id": "app1", "app_tracing": None, "user_id": "user", "extras": extras} + ] module.file_factory.build_from_mappings.assert_not_called() load_annotation_reply_config.assert_called_once_with(sqlite_session, "app1") to_dict.assert_called_once_with(annotation_reply=annotation_reply) diff --git a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py index 4c7e8a1fbad56b..ee8afab8602dff 100644 --- a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py +++ b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py @@ -109,7 +109,7 @@ def _build_app_generate_entity() -> RagPipelineGenerateEntity: pipeline_config=app_config, invoke_from=InvokeFrom.WEB_APP, user_id="user", - trace_manager=None, + trace_scope=None, inputs={"input1": "v1"}, files=[], workflow_execution_id="run", diff --git a/api/tests/unit_tests/core/app/apps/test_base_app_runner.py b/api/tests/unit_tests/core/app/apps/test_base_app_runner.py index 563f8c188b62e4..443caab0770093 100644 --- a/api/tests/unit_tests/core/app/apps/test_base_app_runner.py +++ b/api/tests/unit_tests/core/app/apps/test_base_app_runner.py @@ -6,6 +6,8 @@ from uuid import uuid4 import pytest +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import TraceCorrelation, TraceEvent from sqlalchemy import event, func, select from sqlalchemy.orm import Session, sessionmaker @@ -368,6 +370,14 @@ def test_handle_multimodal_image_content_service_api_mode( assert message_file.created_by_role == CreatorUserRole.END_USER +class _TraceRecorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + @dataclass(eq=False, repr=False) class _DummyParameterRule: name: str @@ -867,11 +877,31 @@ def test_fill_in_inputs_from_external_data_tools(self, monkeypatch: pytest.Monke def test_moderation_for_inputs_returns_result(self, monkeypatch: pytest.MonkeyPatch): runner = AppRunner() - monkeypatch.setattr( - "core.app.apps.base_app_runner.InputModeration.check", - lambda self, app_id, tenant_id, app_config, inputs, query, message_id, trace_manager: (True, {}, ""), + moderation_calls: list[dict[str, object]] = [] + + def check(self: object, **kwargs: object) -> tuple[bool, dict[str, object], str]: + moderation_calls.append(kwargs) + return True, {}, "" + + monkeypatch.setattr("core.app.apps.base_app_runner.InputModeration.check", check) + trace_scope = TraceScope( + emitter=_TraceRecorder(), + tenant_id="tenant", + app_id="app", + user_id="user", + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + ) + app_generate_entity = AppGenerateEntity( + task_id="task", + inputs={}, + files=[], + user_id="user", + stream=False, + invoke_from=InvokeFrom.WEB_APP, + trace_scope=trace_scope, ) - app_generate_entity = AppGenerateEntity.model_construct(app_config=None, trace_manager=None) result = runner.moderation_for_inputs( app_id="app", @@ -883,6 +913,17 @@ def test_moderation_for_inputs_returns_result(self, monkeypatch: pytest.MonkeyPa ) assert result == (True, {}, "") + assert moderation_calls == [ + { + "app_id": "app", + "tenant_id": "tenant", + "app_config": None, + "inputs": {}, + "query": "q", + "message_id": "msg", + "trace_scope": trace_scope, + } + ] @pytest.mark.parametrize("sqlite_session", [()], indirect=True) def test_query_app_annotations_to_reply( diff --git a/api/tests/unit_tests/core/app/apps/test_message_based_app_generator.py b/api/tests/unit_tests/core/app/apps/test_message_based_app_generator.py index 57bccb9696ba21..0dccc7db5175c5 100644 --- a/api/tests/unit_tests/core/app/apps/test_message_based_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/test_message_based_app_generator.py @@ -85,7 +85,7 @@ def _make_chat_generate_entity(app_config: EasyUIBasedAppConfig) -> ChatAppGener invoke_from=InvokeFrom.WEB_APP, extras={}, call_depth=0, - trace_manager=None, + trace_scope=None, ) diff --git a/api/tests/unit_tests/core/app/apps/workflow/test_app_generator.py b/api/tests/unit_tests/core/app/apps/workflow/test_app_generator.py index a6f0085a80fd1a..1b8896b29d27fe 100644 --- a/api/tests/unit_tests/core/app/apps/workflow/test_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/workflow/test_app_generator.py @@ -7,6 +7,8 @@ from unittest.mock import MagicMock, Mock import pytest +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import TraceCorrelation, TraceEvent, WorkflowRunCompleted from flask import Flask from sqlalchemy import inspect from sqlalchemy.orm import Session, sessionmaker @@ -29,7 +31,6 @@ _WorkflowGenerateEntityWrapper, ) from core.app.workflow.file_runtime import DifyWorkflowFileRuntime -from core.ops.ops_trace_manager import OpsTraceManager, TraceQueueManager from core.repositories import SQLAlchemyWorkflowExecutionRepository, SQLAlchemyWorkflowNodeExecutionRepository from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository from core.tools.tool_file_manager import ToolFileManager @@ -52,6 +53,42 @@ """SQLite-backed tests for workflow app generation and worker reload behavior.""" +class _TraceRecorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +class _PersistedRunRecorder: + """Record each workflow trace with the run status persisted when it was emitted.""" + + def __init__(self, session_factory: sessionmaker[Session]) -> None: + self._session_factory = session_factory + self.events: list[TraceEvent] = [] + self.persisted_statuses: list[str | None] = [] + + def emit(self, event: TraceEvent) -> None: + assert isinstance(event.payload, WorkflowRunCompleted) + self.events.append(event) + with self._session_factory() as session: + run = session.get(WorkflowRun, event.payload.workflow_run_id) + self.persisted_statuses.append(run.status if run is not None else None) + + +def _trace_scope(recorder: _TraceRecorder | _PersistedRunRecorder, *, app: App, end_user: EndUser) -> TraceScope: + return TraceScope( + emitter=recorder, + tenant_id=app.tenant_id, + app_id=app.id, + user_id=end_user.session_id, + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + ) + + def _workflow_app_generator_workflow( *, workflow_id: str = "workflow", @@ -122,7 +159,7 @@ def _generate_entity( user_id=end_user.id, stream=stream, invoke_from=invoke_from, - trace_manager=TraceQueueManager(app_id=app.id, user_id=end_user.id), + trace_scope=_trace_scope(_TraceRecorder(), app=app, end_user=end_user), workflow_execution_id="run", ) @@ -155,10 +192,7 @@ def _repositories( class TestWorkflowAppGeneratorCases: @staticmethod @pytest.fixture(autouse=True) - def disable_trace_delivery(monkeypatch: pytest.MonkeyPatch) -> None: - """Construct real trace managers without provider lookup or background delivery.""" - monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", lambda app_id: None) - monkeypatch.setattr(TraceQueueManager, "start_timer", lambda self: None) + def disable_enterprise_telemetry(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr("core.telemetry.gateway.is_enterprise_telemetry_enabled", lambda: False) @staticmethod @@ -183,6 +217,7 @@ def test_worker_executes_and_persists_a_workflow_through_the_injected_driver( } ) sqlite_session.commit() + recorder = _PersistedRunRecorder(sqlite_session_factory) entity = WorkflowAppGenerateEntity( user_from=UserFrom.END_USER, task_id="worker-task", @@ -193,6 +228,7 @@ def test_worker_executes_and_persists_a_workflow_through_the_injected_driver( stream=True, invoke_from=InvokeFrom.SERVICE_API, workflow_execution_id="worker-run", + trace_scope=_trace_scope(recorder, app=app_model, end_user=end_user), ) queue = WorkflowAppQueueManager(entity.task_id, end_user.id, entity.invoke_from, AppMode.WORKFLOW) execution_repository, node_repository = _repositories(sqlite_session_factory, app_model, end_user) @@ -214,6 +250,10 @@ def test_worker_executes_and_persists_a_workflow_through_the_injected_driver( assert run is not None assert run.status == "succeeded" assert run.outputs_dict == {"caller": end_user.session_id} + assert [event.payload for event in recorder.events] == [ + WorkflowRunCompleted(workflow_run_id="worker-run", conversation_id=None, total_tokens=0) + ] + assert recorder.persisted_statuses == ["succeeded"] @staticmethod def test_ensure_snippet_start_node_in_worker_returns_standard_workflow_without_lookup( @@ -317,6 +357,14 @@ def create_workflow_node_execution_repository(**kwargs): "core.app.apps.workflow.app_generator.db", SimpleNamespace(engine=sqlite_session.get_bind()) ) monkeypatch.setattr(generator, "_prepare_user_inputs", lambda *, user_inputs, **kwargs: user_inputs) + trace_scope = _trace_scope(_TraceRecorder(), app=app, end_user=end_user) + scope_requests: list[dict[str, object]] = [] + + def create_trace_scope(**kwargs: object) -> TraceScope: + scope_requests.append(kwargs) + return trace_scope + + monkeypatch.setattr(_workflow_app_generator, "create_trace_scope", create_trace_scope) captured = {} @@ -356,6 +404,16 @@ def fake_generate(**kwargs): "parent_node_execution_id": "outer-node-execution-1", } assert extras["trace_session_id"] == "session-1" + assert application_generate_entity.trace_scope is trace_scope + assert scope_requests == [ + { + "tenant_id": app.tenant_id, + "app_id": app.id, + "app_tracing": None, + "user_id": end_user.session_id, + "extras": extras, + } + ] assert isinstance(captured["workflow_execution_repository"], SQLAlchemyWorkflowExecutionRepository) assert isinstance(captured["workflow_node_execution_repository"], SQLAlchemyWorkflowNodeExecutionRepository) assert repository_tenant_ids == {"workflow": app.tenant_id, "node": app.tenant_id} @@ -654,6 +712,29 @@ def is_alive(self): CREATOR_ID = "00000000-0000-0000-0000-000000000006" +def _app_generator_extra_trace_scope() -> TraceScope: + return TraceScope( + emitter=_TraceRecorder(), + tenant_id=TENANT_ID, + app_id=APP_ID, + user_id="session-id", + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + ) + + +def _record_trace_scope_requests(monkeypatch: pytest.MonkeyPatch, trace_scope: TraceScope) -> list[dict[str, object]]: + requests: list[dict[str, object]] = [] + + def create_trace_scope(**kwargs: object) -> TraceScope: + requests.append(kwargs) + return trace_scope + + monkeypatch.setattr(_app_generator_extra_app_generator_module, "create_trace_scope", create_trace_scope) + return requests + + def _app() -> App: return App( id=APP_ID, @@ -762,10 +843,7 @@ def _persist_snippet( class TestAppGeneratorExtraCases: @staticmethod @pytest.fixture(autouse=True) - def disable_trace_delivery(monkeypatch: pytest.MonkeyPatch) -> None: - """Construct real trace managers without provider lookup or background delivery.""" - monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", lambda app_id: None) - monkeypatch.setattr(TraceQueueManager, "start_timer", lambda self: None) + def disable_enterprise_telemetry(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr("core.telemetry.gateway.is_enterprise_telemetry_enabled", lambda: False) @staticmethod @@ -1234,7 +1312,7 @@ def test_handle_response_closed_file_raises_stopped(self, monkeypatch: pytest.Mo stream=False, invoke_from=InvokeFrom.WEB_APP, extras={}, - trace_manager=None, + trace_scope=None, workflow_execution_id="run-id", call_depth=0, ) @@ -1360,7 +1438,7 @@ def test_generate_prepares_start_inputs(self, generation, root_node_id): execute.assert_not_called() class TestWorkflowAppGeneratorResume: - def test_resume_restores_trace_manager_when_missing(self, monkeypatch: pytest.MonkeyPatch): + def test_resume_recreates_trace_scope_when_missing(self, monkeypatch: pytest.MonkeyPatch) -> None: generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) app_config = WorkflowUIBasedAppConfig( tenant_id="tenant", @@ -1379,24 +1457,13 @@ def test_resume_restores_trace_manager_when_missing(self, monkeypatch: pytest.Mo user_id="user", stream=False, invoke_from=InvokeFrom.WEB_APP, - extras={}, - trace_manager=None, + extras={"external_trace_id": "trace-1"}, + trace_scope=None, workflow_execution_id="run-id", call_depth=0, ) - dummy_trace_queue_manager = type( - "_DummyTraceQueueManager", - (TraceQueueManager,), - { - "__init__": lambda self, app_id=None, user_id=None: ( - setattr(self, "app_id", app_id) or setattr(self, "user_id", user_id) - ) - }, - ) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.TraceQueueManager", - dummy_trace_queue_manager, - ) + trace_scope = _app_generator_extra_trace_scope() + scope_requests = _record_trace_scope_requests(monkeypatch, trace_scope) captured_entity: WorkflowAppGenerateEntity | None = None def _fake_generate(**kwargs): @@ -1418,12 +1485,18 @@ def _fake_generate(**kwargs): assert result.ok is True assert captured_entity is not None - trace_manager = captured_entity.trace_manager - assert isinstance(trace_manager, dummy_trace_queue_manager) - assert trace_manager.app_id == APP_ID - assert trace_manager.user_id == "session-id" + assert captured_entity.trace_scope is trace_scope + assert scope_requests == [ + { + "tenant_id": TENANT_ID, + "app_id": APP_ID, + "app_tracing": None, + "user_id": "session-id", + "extras": {"external_trace_id": "trace-1"}, + } + ] - def test_resume_preserves_existing_trace_manager(self, monkeypatch: pytest.MonkeyPatch): + def test_resume_preserves_existing_trace_scope(self, monkeypatch: pytest.MonkeyPatch) -> None: generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) app_config = WorkflowUIBasedAppConfig( tenant_id="tenant", @@ -1433,7 +1506,7 @@ def test_resume_preserves_existing_trace_manager(self, monkeypatch: pytest.Monke variables=[], workflow_id="workflow-id", ) - existing_trace_manager = SimpleNamespace(app_id="existing-app", user_id="existing-user") + existing_trace_scope = _app_generator_extra_trace_scope() application_generate_entity = WorkflowAppGenerateEntity.model_construct( user_from=UserFrom.END_USER, task_id="task", @@ -1444,10 +1517,11 @@ def test_resume_preserves_existing_trace_manager(self, monkeypatch: pytest.Monke stream=False, invoke_from=InvokeFrom.WEB_APP, extras={}, - trace_manager=existing_trace_manager, + trace_scope=existing_trace_scope, workflow_execution_id="run-id", call_depth=0, ) + scope_requests = _record_trace_scope_requests(monkeypatch, _app_generator_extra_trace_scope()) captured_entity: WorkflowAppGenerateEntity | None = None def _fake_generate(**kwargs): @@ -1469,7 +1543,8 @@ def _fake_generate(**kwargs): assert result.ok is True assert captured_entity is not None - assert captured_entity.trace_manager is existing_trace_manager + assert captured_entity.trace_scope is existing_trace_scope + assert scope_requests == [] class TestWorkflowAppGeneratorWorker: def test_generate_worker_uses_end_user_session_for_external_invocation( @@ -1517,7 +1592,7 @@ def prepare(self): stream=False, invoke_from=InvokeFrom.WEB_APP, extras={}, - trace_manager=None, + trace_scope=None, workflow_execution_id="run-id", call_depth=0, ) diff --git a/api/tests/unit_tests/core/app/apps/workflow/test_generate_task_pipeline.py b/api/tests/unit_tests/core/app/apps/workflow/test_generate_task_pipeline.py index d0568e7cabd1c8..5c116f5f46aa4e 100644 --- a/api/tests/unit_tests/core/app/apps/workflow/test_generate_task_pipeline.py +++ b/api/tests/unit_tests/core/app/apps/workflow/test_generate_task_pipeline.py @@ -184,7 +184,7 @@ def _make_pipeline(): user_id="user", stream=False, invoke_from=InvokeFrom.WEB_APP, - trace_manager=None, + trace_scope=None, workflow_execution_id="run-id", extras={}, call_depth=0, @@ -628,7 +628,7 @@ def test_init_with_end_user_sets_role_and_system_user(self): user_id="end-user-id", stream=False, invoke_from=InvokeFrom.WEB_APP, - trace_manager=None, + trace_scope=None, workflow_execution_id="run-id", extras={}, call_depth=0, diff --git a/api/tests/unit_tests/core/app/entities/test_app_invoke_entities.py b/api/tests/unit_tests/core/app/entities/test_app_invoke_entities.py index 17be35ee40b7a1..8603ee44e8ad02 100644 --- a/api/tests/unit_tests/core/app/entities/test_app_invoke_entities.py +++ b/api/tests/unit_tests/core/app/entities/test_app_invoke_entities.py @@ -3,6 +3,8 @@ from dataclasses import dataclass import pytest +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import TraceCorrelation, TraceEvent from core.app.app_config.entities import WorkflowUIBasedAppConfig from core.app.entities.app_invoke_entities import ( @@ -21,16 +23,27 @@ _WorkflowGenerateEntityWrapper, ) from core.credit_usage import CreditUsageAppType -from core.ops.ops_trace_manager import TraceQueueManager from models.model import AppMode -class TraceQueueManagerStub(TraceQueueManager): - """Minimal TraceQueueManager stub that avoids Flask dependencies.""" +class _TraceRecorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] - def __init__(self): - # Skip parent initialization to avoid starting timers or accessing Flask globals. - pass + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +def _trace_scope() -> TraceScope: + return TraceScope( + emitter=_TraceRecorder(), + tenant_id="tenant-id", + app_id="app-id", + user_id="user-id", + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + ) def _build_workflow_app_config(app_mode: AppMode) -> WorkflowUIBasedAppConfig: @@ -42,7 +55,7 @@ def _build_workflow_app_config(app_mode: AppMode) -> WorkflowUIBasedAppConfig: ) -def _create_workflow_generate_entity(trace_manager: TraceQueueManager | None = None) -> WorkflowAppGenerateEntity: +def _create_workflow_generate_entity(trace_scope: TraceScope) -> WorkflowAppGenerateEntity: return WorkflowAppGenerateEntity( user_from=UserFrom.ACCOUNT, task_id="workflow-task", @@ -53,15 +66,13 @@ def _create_workflow_generate_entity(trace_manager: TraceQueueManager | None = N stream=True, invoke_from=InvokeFrom.DEBUGGER, call_depth=1, - trace_manager=trace_manager, + trace_scope=trace_scope, workflow_execution_id="workflow-exec-id", extras={"external_trace_id": "trace-id"}, ) -def _create_advanced_chat_generate_entity( - trace_manager: TraceQueueManager | None = None, -) -> AdvancedChatAppGenerateEntity: +def _create_advanced_chat_generate_entity(trace_scope: TraceScope) -> AdvancedChatAppGenerateEntity: return AdvancedChatAppGenerateEntity( task_id="advanced-task", app_config=_build_workflow_app_config(AppMode.ADVANCED_CHAT), @@ -73,37 +84,42 @@ def _create_advanced_chat_generate_entity( invoke_from=InvokeFrom.DEBUGGER, query="Explain serialization", extras={"auto_generate_conversation_name": True}, - trace_manager=trace_manager, + trace_scope=trace_scope, workflow_run_id="workflow-run-id", ) -def test_workflow_app_generate_entity_roundtrip_excludes_trace_manager(): - entity = _create_workflow_generate_entity(trace_manager=TraceQueueManagerStub()) +def test_workflow_app_generate_entity_roundtrip_excludes_trace_scope() -> None: + trace_scope = _trace_scope() + entity = _create_workflow_generate_entity(trace_scope=trace_scope) serialized = entity.model_dump_json() payload = json.loads(serialized) - assert "trace_manager" not in payload + assert entity.trace_scope is trace_scope + assert "trace_scope" not in payload + assert "trace_scope" not in repr(entity) restored = WorkflowAppGenerateEntity.model_validate_json(serialized) assert restored.model_dump() == entity.model_dump() - assert restored.trace_manager is None + assert restored.trace_scope is None -def test_advanced_chat_generate_entity_roundtrip_excludes_trace_manager(): - entity = _create_advanced_chat_generate_entity(trace_manager=TraceQueueManagerStub()) +def test_advanced_chat_generate_entity_roundtrip_excludes_trace_scope() -> None: + trace_scope = _trace_scope() + entity = _create_advanced_chat_generate_entity(trace_scope=trace_scope) serialized = entity.model_dump_json() payload = json.loads(serialized) - assert "trace_manager" not in payload + assert entity.trace_scope is trace_scope + assert "trace_scope" not in payload restored = AdvancedChatAppGenerateEntity.model_validate_json(serialized) assert restored.model_dump() == entity.model_dump() - assert restored.trace_manager is None + assert restored.trace_scope is None @pytest.mark.parametrize( @@ -167,7 +183,7 @@ class ResumptionContextCase: def _workflow_resumption_case() -> tuple[WorkflowResumptionContext, type]: - entity = _create_workflow_generate_entity(trace_manager=TraceQueueManagerStub()) + entity = _create_workflow_generate_entity(trace_scope=_trace_scope()) context = WorkflowResumptionContext( serialized_graph_runtime_state=json.dumps({"state": "workflow"}), generate_entity=_WorkflowGenerateEntityWrapper(entity=entity), @@ -176,7 +192,7 @@ def _workflow_resumption_case() -> tuple[WorkflowResumptionContext, type]: def _advanced_chat_resumption_case() -> tuple[WorkflowResumptionContext, type]: - entity = _create_advanced_chat_generate_entity(trace_manager=TraceQueueManagerStub()) + entity = _create_advanced_chat_generate_entity(trace_scope=_trace_scope()) context = WorkflowResumptionContext( serialized_graph_runtime_state=json.dumps({"state": "advanced"}), generate_entity=_AdvancedChatAppGenerateEntityWrapper(entity=entity), @@ -201,7 +217,7 @@ def test_workflow_resumption_context_roundtrip(case: ResumptionContextCase): entity = restored.get_generate_entity() assert isinstance(entity, expected_type) assert entity.model_dump() == context.get_generate_entity().model_dump() - assert entity.trace_manager is None + assert entity.trace_scope is None def test_openapi_variant_present(): diff --git a/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline.py b/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline.py index f6ecac635fe7d3..3f5c8bc1295e25 100644 --- a/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline.py +++ b/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline.py @@ -3,6 +3,7 @@ import json import uuid from collections.abc import Generator, Iterator, Sequence +from contextlib import contextmanager from dataclasses import dataclass from datetime import UTC, datetime from decimal import Decimal @@ -11,6 +12,8 @@ from unittest.mock import MagicMock, Mock, patch import pytest +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import MessageCompleted, TraceCorrelation, TraceEvent from sqlalchemy import Engine as _easy_ui_based_generate_task_pipeline_Engine from sqlalchemy import event from sqlalchemy.engine import Engine as _easy_ui_message_end_files_Engine @@ -60,8 +63,6 @@ ) from core.app.task_pipeline.easy_ui_based_generate_task_pipeline import EasyUIBasedGenerateTaskPipeline from core.base.tts import AppGeneratorTTSPublisher, AudioTrunk -from core.ops.entities.trace_entity import TraceTaskName -from core.ops.ops_trace_manager import TraceQueueManager from extensions.ext_redis import RedisClientWrapper from extensions.storage.storage_type import StorageType from graphon.file import FileTransferMethod, FileType @@ -97,6 +98,19 @@ def publish(self, message: WorkflowQueueMessage | MessageQueueMessage | None, /) self.messages.append(message) +class _CommitCountTraceRecorder: + """Record trace events with the number of commits that preceded each one.""" + + def __init__(self, committed_sessions: list[Session]) -> None: + self._committed_sessions = committed_sessions + self.events: list[TraceEvent] = [] + self.commits_before_emit: list[int] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + self.commits_before_emit.append(len(self._committed_sessions)) + + def queue_message(event: AppQueueEvent) -> MessageQueueMessage: return MessageQueueMessage( task_id="test-task-id", @@ -243,7 +257,7 @@ def test_get_message_event_type_called_once_when_first_llm_chunk_arrives( pipeline.queue_manager.messages = [message] # Execute - list(pipeline._process_stream_response(publisher=None, trace_manager=None)) + list(pipeline._process_stream_response(publisher=None)) # Assert mock_message_cycle_manager.get_message_event_type.assert_called_once_with(message_id="test-message-id") @@ -261,7 +275,7 @@ def test_llm_chunk_event_with_text_content(self, pipeline, mock_message_cycle_ma mock_message_cycle_manager.get_message_event_type.return_value = StreamEvent.MESSAGE # Execute - responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) + responses = list(pipeline._process_stream_response(publisher=None)) # Assert assert len(responses) == 1 @@ -286,7 +300,7 @@ def test_llm_chunk_event_with_list_content(self, pipeline, mock_message_cycle_ma mock_message_cycle_manager.get_message_event_type.return_value = StreamEvent.MESSAGE # Execute - responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) + responses = list(pipeline._process_stream_response(publisher=None)) # Assert assert len(responses) == 1 @@ -309,7 +323,7 @@ def test_agent_message_event(self, pipeline, mock_message_cycle_manager, task_st pipeline._agent_message_to_stream_response = Mock(return_value=Mock()) # Execute - responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) + responses = list(pipeline._process_stream_response(publisher=None)) # Assert assert len(responses) == 1 @@ -332,7 +346,7 @@ def test_message_end_event(self, pipeline, mock_message_cycle_manager, task_stat return_value=MessageEndStreamResponse(task_id="test-task-id", id="test-message-id") ) - responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) + responses = list(pipeline._process_stream_response(publisher=None)) # Assert assert len(responses) == 1 @@ -355,7 +369,7 @@ def test_error_event(self, pipeline, committed_sessions): return_value=ErrorStreamResponse(task_id="test-task-id", err=Exception("Test error")) ) - responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) + responses = list(pipeline._process_stream_response(publisher=None)) # Assert assert len(responses) == 1 @@ -378,7 +392,7 @@ def test_ping_event(self, pipeline): pipeline.ping_stream_response = Mock(return_value=PingStreamResponse(task_id="test-task-id")) # Execute - responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) + responses = list(pipeline._process_stream_response(publisher=None)) # Assert assert len(responses) == 1 @@ -402,7 +416,7 @@ def test_file_event(self, pipeline, mock_message_cycle_manager): mock_message_cycle_manager.message_file_to_stream_response.return_value = file_response # Execute - responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) + responses = list(pipeline._process_stream_response(publisher=None)) # Assert assert len(responses) == 1 @@ -421,37 +435,42 @@ def test_publisher_is_called_with_messages(self, pipeline): pipeline.ping_stream_response = Mock(return_value=PingStreamResponse(task_id="test-task-id")) # Execute - list(pipeline._process_stream_response(publisher=publisher, trace_manager=None)) + list(pipeline._process_stream_response(publisher=publisher)) # Assert assert publisher.messages == [message] - def test_trace_manager_passed_to_save_message(self, pipeline, committed_sessions): - """Test that trace manager is passed to _save_message.""" - # Setup - with ( - patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None), - patch.object(TraceQueueManager, "start_timer"), - ): - trace_manager = TraceQueueManager(app_id="test-app-id") - - message_end_event = QueueMessageEndEvent(llm_result=None) - - message = queue_message(message_end_event) - pipeline.queue_manager.messages = [message] - - pipeline._save_message = Mock() - pipeline._message_end_to_stream_response = Mock( - return_value=MessageEndStreamResponse(task_id="test-task-id", id="test-message-id") + def test_message_end_emits_message_trace_after_commit( + self, + pipeline: EasyUIBasedGenerateTaskPipeline, + queue_manager: _QueueManager, + application_generate_entity: ChatAppGenerateEntity, + committed_sessions: list[Session], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """The message trace is reported only after the saved message is committed.""" + recorder = _CommitCountTraceRecorder(committed_sessions) + application_generate_entity.trace_scope = TraceScope( + emitter=recorder, + tenant_id="test-tenant-id", + app_id="test-app-id", + user_id="test-user-id", + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), ) + queue_manager.messages = [queue_message(QueueMessageEndEvent(llm_result=None))] + saved_sessions: list[Session] = [] + monkeypatch.setattr(pipeline, "_save_message", lambda *, session: saved_sessions.append(session)) - list(pipeline._process_stream_response(publisher=None, trace_manager=trace_manager)) + responses = list(pipeline._process_stream_response(publisher=None)) - # Assert - session = pipeline._save_message.call_args.kwargs["session"] - assert isinstance(session, Session) - assert committed_sessions == [session] - pipeline._save_message.assert_called_once_with(session=session, trace_manager=trace_manager) + assert isinstance(responses[-1], MessageEndStreamResponse) + assert committed_sessions == saved_sessions + assert [event.payload for event in recorder.events] == [ + MessageCompleted(message_id="test-message-id", conversation_id="test-conversation-id") + ] + assert recorder.commits_before_emit == [1] def test_multiple_events_sequence(self, pipeline, mock_message_cycle_manager, task_state): """Test handling multiple events in sequence.""" @@ -477,7 +496,7 @@ def test_multiple_events_sequence(self, pipeline, mock_message_cycle_manager, ta pipeline.ping_stream_response = Mock(return_value=PingStreamResponse(task_id="test-task-id")) # Execute - responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) + responses = list(pipeline._process_stream_response(publisher=None)) # Assert assert len(responses) == 3 @@ -525,9 +544,40 @@ def check_and_get_audio(self) -> AudioTrunk: return self._audio -class _TraceManagerDouble: - def __init__(self) -> None: - self.add_trace_task = Mock() +class _TimelineTraceRecorder: + """Record trace events on a timeline shared with the commits they must follow.""" + + def __init__(self, timeline: list[str]) -> None: + self._timeline = timeline + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + self._timeline.append("trace") + + +def _trace_scope(recorder: _TimelineTraceRecorder) -> TraceScope: + return TraceScope( + emitter=recorder, + tenant_id="tenant", + app_id="app", + user_id="user", + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + ) + + +@contextmanager +def _recording_commits(timeline: list[str]) -> Generator[None]: + def record_commit(_session: Session) -> None: + timeline.append("commit") + + event.listen(Session, "after_commit", record_commit) + try: + yield + finally: + event.remove(Session, "after_commit", record_commit) class _FakeQueueManager(AppQueueManager): @@ -708,7 +758,7 @@ def _make_entity(entity_cls, app_mode: AppMode): invoke_from=InvokeFrom.WEB_APP, extras={}, call_depth=0, - trace_manager=None, + trace_scope=None, ) @@ -936,7 +986,7 @@ def __init__(self) -> None: invoke_from=InvokeFrom.WEB_APP, extras={}, call_depth=0, - trace_manager=None, + trace_scope=None, ) pipeline = EasyUIBasedGenerateTaskPipeline( @@ -1135,51 +1185,32 @@ def commit(self): assert isinstance(responses[-1].err, ValueError) assert pipeline._task_state.llm_result.message.content == "annotated" - def test_process_stream_response_error_event_adds_trace_task(self, monkeypatch: pytest.MonkeyPatch): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() + def test_process_stream_response_error_event_emits_message_trace_after_commit(self) -> None: + timeline: list[str] = [] + recorder = _TimelineTraceRecorder(timeline) application_generate_entity = _make_entity(ChatAppGenerateEntity, AppMode.CHAT) - application_generate_entity.extras = {"trace_session_id": "session-1"} + application_generate_entity.trace_scope = _trace_scope(recorder) pipeline = EasyUIBasedGenerateTaskPipeline( application_generate_entity=application_generate_entity, queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, + conversation=_make_conversation(AppMode.CHAT), + message=_make_message(), stream=True, ) _set_queue_events(pipeline, [_queue_message(QueueErrorEvent(error=ValueError("boom")))]) _set_method(pipeline, "handle_error", lambda **kwargs: ValueError("boom")) _set_method(pipeline, "error_to_stream_response", lambda err: ErrorStreamResponse(task_id="task", err=err)) - trace_manager_double = _TraceManagerDouble() - trace_manager = cast(TraceQueueManager, trace_manager_double) - - class _Session: - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def commit(self): - return None - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) - - responses = list(pipeline._process_stream_response(publisher=None, trace_manager=trace_manager)) + with _recording_commits(timeline): + responses = list(pipeline._process_stream_response(publisher=None)) assert len(responses) == 1 assert isinstance(responses[0], ErrorStreamResponse) - trace_manager_double.add_trace_task.assert_called_once() - trace_task = trace_manager_double.add_trace_task.call_args.args[0] - assert trace_task.trace_type == TraceTaskName.MESSAGE_TRACE - assert trace_task.conversation_id == "conv" - assert trace_task.message_id == "msg" - assert trace_task.kwargs["trace_session_id"] == "session-1" + assert timeline == ["commit", "trace"] + assert [event.payload for event in recorder.events] == [ + MessageCompleted(message_id="msg", conversation_id="conv") + ] def test_agent_thought_to_stream_response_returns_payload(self, monkeypatch: pytest.MonkeyPatch): conversation = _make_conversation(AppMode.CHAT) @@ -1279,7 +1310,7 @@ def test_process_routes_to_stream_and_starts_conversation_name_generation(self): _set_method( pipeline, "_wrapper_process_stream_response", - lambda trace_manager: iter([PingStreamResponse(task_id="task")]), + lambda: iter([PingStreamResponse(task_id="task")]), ) _set_method(pipeline, "_to_stream_response", lambda generator: "streamed") @@ -1304,7 +1335,7 @@ def test_process_routes_to_blocking_for_completion_mode(self): _set_method( pipeline, "_wrapper_process_stream_response", - lambda trace_manager: iter([PingStreamResponse(task_id="task")]), + lambda: iter([PingStreamResponse(task_id="task")]), ) _set_method(pipeline, "_to_blocking_response", lambda generator: "blocking") @@ -1433,7 +1464,7 @@ def test_wrapper_process_stream_response_without_tts_publisher(self): stream=True, ) payload = PingStreamResponse(task_id="task") - _set_method(pipeline, "_process_stream_response", lambda publisher, trace_manager: iter([payload])) + _set_method(pipeline, "_process_stream_response", lambda publisher: iter([payload])) responses = list(pipeline._wrapper_process_stream_response()) @@ -1469,7 +1500,7 @@ def cancel(self): audio_calls = iter([inline_audio, None]) payload = PingStreamResponse(task_id="task") _set_method(pipeline, "_listen_audio_msg", lambda publisher, task_id: next(audio_calls)) - _set_method(pipeline, "_process_stream_response", lambda publisher, trace_manager: iter([payload])) + _set_method(pipeline, "_process_stream_response", lambda publisher: iter([payload])) monkeypatch.setattr( "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.AppGeneratorTTSPublisher", lambda tenant_id, voice, language, app_type: _Publisher(), @@ -1511,7 +1542,7 @@ def publish(self, message): def cancel(self): return None - _set_method(pipeline, "_process_stream_response", lambda publisher, trace_manager: iter([])) + _set_method(pipeline, "_process_stream_response", lambda publisher: iter([])) monkeypatch.setattr( "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.AppGeneratorTTSPublisher", lambda tenant_id, voice, language, app_type: _Publisher(), @@ -1550,7 +1581,7 @@ def publish(self, message): def cancel(self): return None - _set_method(pipeline, "_process_stream_response", lambda publisher, trace_manager: iter([])) + _set_method(pipeline, "_process_stream_response", lambda publisher: iter([])) monkeypatch.setattr( "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.AppGeneratorTTSPublisher", lambda tenant_id, voice, language, app_type: _Publisher(), @@ -1579,7 +1610,7 @@ def test_wrapper_process_stream_response_ends_audio_before_a_main_error(self, mo publisher = Mock() publisher.check_and_get_audio.return_value = None error = ErrorStreamResponse(task_id="task", err=RuntimeError("generation failed")) - _set_method(pipeline, "_process_stream_response", lambda publisher, trace_manager: iter([error])) + _set_method(pipeline, "_process_stream_response", lambda publisher: iter([error])) monkeypatch.setattr( "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.AppGeneratorTTSPublisher", lambda tenant_id, voice, language, app_type: publisher, @@ -1770,13 +1801,15 @@ def test_process_stream_response_reaches_post_loop_branch_with_thread_reference( assert list(pipeline._process_stream_response(publisher=None)) == [] - def test_save_message_persists_fields_and_emits_trace( + def test_save_message_persists_fields_without_emitting_trace( self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session - ): + ) -> None: conversation = _make_conversation(AppMode.CHAT) message = _make_message() + timeline: list[str] = [] + recorder = _TimelineTraceRecorder(timeline) application_generate_entity = _make_entity(ChatAppGenerateEntity, AppMode.CHAT) - application_generate_entity.extras = {"trace_session_id": "session-1"} + application_generate_entity.trace_scope = _trace_scope(recorder) pipeline = EasyUIBasedGenerateTaskPipeline( application_generate_entity=application_generate_entity, queue_manager=_FakeQueueManager(), @@ -1797,8 +1830,6 @@ def test_save_message_persists_fields_and_emits_trace( session = sqlite_session session.add_all([conversation_obj, message_obj]) session.flush() - trace_manager_double = _TraceManagerDouble() - trace_manager = cast(TraceQueueManager, trace_manager_double) sent_payloads: list[tuple[tuple[object, ...], dict[str, object]]] = [] monkeypatch.setattr( @@ -1821,18 +1852,66 @@ def test_save_message_persists_fields_and_emits_trace( lambda *args, **kwargs: sent_payloads.append((args, kwargs)), ) - pipeline._save_message(session=session, trace_manager=trace_manager) + pipeline._save_message(session=session) assert message_obj.message == "serialized-prompt" assert message_obj.answer == "hello" assert message_obj.provider_response_latency == 5.0 - trace_manager_double.add_trace_task.assert_called_once() - trace_task = trace_manager_double.add_trace_task.call_args.args[0] - assert trace_task.trace_type == TraceTaskName.MESSAGE_TRACE - assert trace_task.conversation_id == "conv" - assert trace_task.message_id == "msg" - assert trace_task.kwargs["trace_session_id"] == "session-1" assert len(sent_payloads) == 1 + # The caller reports the message only after it commits this session. + assert recorder.events == [] + + @pytest.mark.parametrize( + ("app_mode", "expected_timeline", "expected_payloads"), + [ + ( + AppMode.CHAT, + ["message_was_created", "commit", "trace"], + [MessageCompleted(message_id="msg", conversation_id="conv")], + ), + # The Agent worker reports a stopped Agent run's terminal itself. + (AppMode.AGENT, ["commit", "message_was_created"], []), + ], + ) + def test_stopped_message_trace_follows_commit_unless_agent_worker_owns_it( + self, + app_mode: AppMode, + expected_timeline: list[str], + expected_payloads: list[MessageCompleted], + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + ) -> None: + timeline: list[str] = [] + recorder = _TimelineTraceRecorder(timeline) + application_generate_entity = _make_entity(ChatAppGenerateEntity, app_mode) + application_generate_entity.trace_scope = _trace_scope(recorder) + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=application_generate_entity, + queue_manager=_FakeQueueManager(), + conversation=_make_conversation(app_mode), + message=_make_message(), + stream=True, + ) + _set_method(pipeline, "_model_config", _ModelConfigMode(mode="chat")) + _set_method(pipeline, "_handle_stop", lambda event: None) + _set_queue_events(pipeline, [_queue_message(QueueStopEvent(stopped_by=QueueStopEvent.StopBy.USER_MANUAL))]) + sqlite_session.add_all([_make_conversation(app_mode), _make_message()]) + sqlite_session.commit() + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.PromptMessageUtil.prompt_messages_to_prompt_for_saving", + lambda _mode, _prompt_messages: "", + ) + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.message_was_created.send", + lambda *_args, **_kwargs: timeline.append("message_was_created"), + ) + + with _recording_commits(timeline): + responses = list(pipeline._process_stream_response(publisher=None)) + + assert isinstance(responses[-1], MessageEndStreamResponse) + assert timeline == expected_timeline + assert [event.payload for event in recorder.events] == expected_payloads def test_save_stopped_message_preserves_backend_reported_usage( self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session diff --git a/api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager.py b/api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager.py index f50355ace268ab..41832d5acae6a3 100644 --- a/api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager.py +++ b/api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager.py @@ -387,7 +387,12 @@ def test_generate_conversation_name_worker_generates_and_caches_name( 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" + "tenant-1", + "hello", + "conv-1", + "app-id", + message_id="message-1", + trace_scope=message_cycle_manager._application_generate_entity.trace_scope, ) mock_redis.setex.assert_called_once() diff --git a/api/tests/unit_tests/core/callback_handler/test_agent_tool_callback_handler.py b/api/tests/unit_tests/core/callback_handler/test_agent_tool_callback_handler.py index 6245b1ace230ee..a398ede0366c59 100644 --- a/api/tests/unit_tests/core/callback_handler/test_agent_tool_callback_handler.py +++ b/api/tests/unit_tests/core/callback_handler/test_agent_tool_callback_handler.py @@ -126,27 +126,9 @@ def test_on_tool_start_debug_disabled( mock_print_text.assert_not_called() - def test_on_tool_end_debug_enabled_and_trace( - self, handler: DifyAgentCallbackHandler, enable_debug, mocker: MockerFixture - ): - mock_print_text = mocker.patch("core.callback_handler.agent_tool_callback_handler.print_text") - mock_trace_manager = MagicMock() - - handler.on_tool_end( - tool_name="tool1", - tool_inputs={"a": 1}, - tool_outputs="output", - message_id="msg1", - timer=123, - trace_manager=mock_trace_manager, - ) - - assert mock_print_text.call_count >= 1 - mock_trace_manager.add_trace_task.assert_called_once() - - def test_on_tool_end_without_trace_manager( - self, handler: DifyAgentCallbackHandler, enable_debug, mocker: MockerFixture - ): + def test_on_tool_end_debug_enabled( + self, handler: DifyAgentCallbackHandler, enable_debug: None, mocker: MockerFixture + ) -> None: mock_print_text = mocker.patch("core.callback_handler.agent_tool_callback_handler.print_text") handler.on_tool_end( diff --git a/api/tests/unit_tests/core/callback_handler/test_workflow_tool_callback_handler.py b/api/tests/unit_tests/core/callback_handler/test_workflow_tool_callback_handler.py index d55c4fc1a69e5f..37326fb3417c33 100644 --- a/api/tests/unit_tests/core/callback_handler/test_workflow_tool_callback_handler.py +++ b/api/tests/unit_tests/core/callback_handler/test_workflow_tool_callback_handler.py @@ -206,25 +206,3 @@ def test_on_tool_execution_model_dump_json_exception( # Ensure first two prints happened before failure assert mock_print_text.call_count >= 2 - - def test_on_tool_execution_none_message_id_and_trace_manager( - self, handler: DifyWorkflowCallbackHandler, mock_print_text, enable_debug - ): - # Arrange - tool_name = "optional_params_tool" - message = DummyToolInvokeMessage('{"data": "ok"}') - - # Act - results = list( - handler.on_tool_execution( - tool_name=tool_name, - tool_inputs={}, - tool_outputs=[message], - message_id=None, - timer=None, - trace_manager=None, - ) - ) - - assert results == [message] - assert mock_print_text.call_count == 4 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 ba36023dc0c7a5..13c5e78e74b727 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 @@ -6,6 +6,8 @@ from uuid import uuid4 import pytest +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import ConversationNamed, TraceCorrelation, TraceEvent from sqlalchemy import event from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.orm import Session, scoped_session, sessionmaker @@ -15,6 +17,7 @@ from core.llm_generator import llm_generator as llm_generator_module from core.llm_generator.entities import RuleCodeGeneratePayload, RuleGeneratePayload, RuleStructuredOutputPayload from core.llm_generator.llm_generator import LLMGenerator, _parse_string_list +from core.llm_generator.prompts import CONVERSATION_TITLE_PROMPT from core.model_manager import ModelManager from core.plugin.impl.model_runtime_factory import create_plugin_model_manager from core.workflow.generator import tool_catalogue as tool_catalogue_module @@ -46,6 +49,14 @@ def _llm_result(content: str) -> LLMResult: ) +class _Recorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + def _persist_app(database: Session, *, tenant_id: str | None = None) -> App: app = App( id=str(uuid4()), @@ -189,19 +200,44 @@ def model_config_entity(self): provider="openai", name="gpt-4", mode=LLMMode.CHAT, completion_params={"temperature": 0.7} ) - def test_generate_conversation_name_success(self, mock_model_instance): + def test_generate_conversation_name_success( + self, mock_model_instance: MagicMock, monkeypatch: pytest.MonkeyPatch + ) -> None: mock_response = MagicMock() mock_response.message.get_text_content.return_value = json.dumps({"Your Output": "Test Conversation Name"}) mock_model_instance.invoke_llm.return_value = mock_response + recorder = _Recorder() + scope = TraceScope( + emitter=recorder, + tenant_id="tenant_id", + app_id="app-1", + user_id=None, + correlation=TraceCorrelation( + external_trace_id="caller-trace", + session_id=None, + parent_workflow_run_id=None, + parent_node_execution_id=None, + ), + ) - with patch("core.llm_generator.llm_generator.TraceQueueManager") as mock_trace: - 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_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" + name = LLMGenerator.generate_conversation_name( + "tenant_id", "test query", "conversation-1", "app-1", message_id="message-1", trace_scope=scope + ) + + assert name == "Test Conversation Name" + [event] = recorder.events + assert event.correlation.external_trace_id == "caller-trace" + payload = event.payload + assert isinstance(payload, ConversationNamed) + assert payload == ConversationNamed( + conversation_id="conversation-1", + message_id="message-1", + inputs=CONVERSATION_TITLE_PROMPT + "test query\n", + name="Test Conversation Name", + started_at=payload.started_at, + ended_at=payload.ended_at, + ) + assert payload.ended_at >= payload.started_at def test_generate_conversation_name_truncated(self, mock_model_instance): long_query = "a" * 2100 @@ -209,16 +245,15 @@ def test_generate_conversation_name_truncated(self, mock_model_instance): mock_response.message.get_text_content.return_value = json.dumps({"Your Output": "Short Name"}) mock_model_instance.invoke_llm.return_value = mock_response - with patch("core.llm_generator.llm_generator.TraceQueueManager"): - name = LLMGenerator.generate_conversation_name("tenant_id", long_query) - assert name == "Short Name" + name = LLMGenerator.generate_conversation_name("tenant_id", long_query, trace_scope=None) + assert name == "Short Name" def test_generate_conversation_name_empty_answer(self, mock_model_instance): mock_response = MagicMock() mock_response.message.get_text_content.return_value = "" mock_model_instance.invoke_llm.return_value = mock_response - name = LLMGenerator.generate_conversation_name("tenant_id", "test query") + name = LLMGenerator.generate_conversation_name("tenant_id", "test query", trace_scope=None) assert name == "test query" def test_generate_conversation_name_json_repair(self, mock_model_instance): @@ -227,25 +262,22 @@ def test_generate_conversation_name_json_repair(self, mock_model_instance): mock_response.message.get_text_content.return_value = "{'Your Output': 'Repaired Name'}" mock_model_instance.invoke_llm.return_value = mock_response - with patch("core.llm_generator.llm_generator.TraceQueueManager"): - name = LLMGenerator.generate_conversation_name("tenant_id", "test query") - assert name == "Repaired Name" + name = LLMGenerator.generate_conversation_name("tenant_id", "test query", trace_scope=None) + assert name == "Repaired Name" def test_generate_conversation_name_not_dict_result(self, mock_model_instance): mock_response = MagicMock() mock_response.message.get_text_content.return_value = '["not a dict"]' mock_model_instance.invoke_llm.return_value = mock_response - with patch("core.llm_generator.llm_generator.TraceQueueManager"): - name = LLMGenerator.generate_conversation_name("tenant_id", "test query") - assert name == "test query" + name = LLMGenerator.generate_conversation_name("tenant_id", "test query", trace_scope=None) + assert name == "test query" def test_generate_conversation_name_no_output_in_dict(self, mock_model_instance): mock_response = MagicMock() mock_response.message.get_text_content.return_value = '{"something": "else"}' mock_model_instance.invoke_llm.return_value = mock_response - with patch("core.llm_generator.llm_generator.TraceQueueManager"): - name = LLMGenerator.generate_conversation_name("tenant_id", "test query") - assert name == "test query" + name = LLMGenerator.generate_conversation_name("tenant_id", "test query", trace_scope=None) + assert name == "test query" def test_generate_conversation_name_long_output(self, mock_model_instance): long_output = "a" * 100 @@ -253,10 +285,9 @@ def test_generate_conversation_name_long_output(self, mock_model_instance): mock_response.message.get_text_content.return_value = json.dumps({"Your Output": long_output}) mock_model_instance.invoke_llm.return_value = mock_response - with patch("core.llm_generator.llm_generator.TraceQueueManager"): - name = LLMGenerator.generate_conversation_name("tenant_id", "test query") - assert len(name) == 78 # 75 + "..." - assert name.endswith("...") + name = LLMGenerator.generate_conversation_name("tenant_id", "test query", trace_scope=None) + assert len(name) == 78 # 75 + "..." + assert name.endswith("...") def test_generate_rule_config_no_variable_success(self, mock_model_instance, model_config_entity): payload = RuleGeneratePayload( @@ -790,6 +821,8 @@ def test_generate_code_emits_on_success(self, mock_emit: MagicMock) -> None: assert event.context.tenant_id == "t-1" assert event.context.app_id == "app-1" assert event.payload["error"] is None + started_at, ended_at = event.payload["started_at"], event.payload["ended_at"] + assert event.payload["latency"] == (ended_at - started_at).total_seconds() @patch("core.llm_generator.llm_generator.telemetry_emit") def test_generate_code_emits_on_failure(self, mock_emit: MagicMock) -> None: diff --git a/api/tests/unit_tests/core/moderation/test_input_moderation.py b/api/tests/unit_tests/core/moderation/test_input_moderation.py index a98fed2a1f37da..ad51d7bbc1489d 100644 --- a/api/tests/unit_tests/core/moderation/test_input_moderation.py +++ b/api/tests/unit_tests/core/moderation/test_input_moderation.py @@ -1,35 +1,42 @@ -import queue +from datetime import timedelta +from decimal import Decimal from unittest.mock import MagicMock, patch import pytest -from flask import Flask +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import ModerationChecked, TraceCorrelation, TraceEvent from core.app.app_config.entities import AppConfig, SensitiveWordAvoidanceEntity from core.moderation.base import ModerationAction, ModerationError, ModerationInputsResult from core.moderation.input_moderation import InputModeration -from core.ops.entities.trace_entity import TraceTaskName -from core.ops.ops_trace_manager import TraceQueueManager from models.model import AppMode +class _Recorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +def _trace_scope(recorder: _Recorder) -> TraceScope: + return TraceScope( + emitter=recorder, + tenant_id="test_tenant_id", + app_id="test_app_id", + user_id="test_user_id", + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + ) + + class TestInputModeration: @pytest.fixture def app_config(self): return AppConfig(tenant_id="test_tenant_id", app_id="test_app_id", app_mode=AppMode.CHAT) - @pytest.fixture - def trace_manager(self): - with ( - Flask(__name__).app_context(), - patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None), - patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True), - patch("core.ops.ops_trace_manager.trace_manager_queue", queue.Queue()), - patch.object(TraceQueueManager, "start_timer"), - ): - manager = TraceQueueManager(app_id="test_app_id") - with patch.object(manager, "add_trace_task", wraps=manager.add_trace_task): - yield manager - @pytest.fixture def input_moderation(self): return InputModeration() @@ -42,7 +49,13 @@ def test_check_no_sensitive_word_avoidance(self, app_config, input_moderation): message_id = "test_message_id" flagged, final_inputs, final_query = input_moderation.check( - app_id=app_id, tenant_id=tenant_id, app_config=app_config, inputs=inputs, query=query, message_id=message_id + app_id=app_id, + tenant_id=tenant_id, + app_config=app_config, + inputs=inputs, + query=query, + message_id=message_id, + trace_scope=None, ) assert flagged is False @@ -67,7 +80,13 @@ def test_check_not_flagged(self, mock_factory_cls, app_config, input_moderation) mock_factory.moderation_for_inputs.return_value = mock_result flagged, final_inputs, final_query = input_moderation.check( - app_id=app_id, tenant_id=tenant_id, app_config=app_config, inputs=inputs, query=query, message_id=message_id + app_id=app_id, + tenant_id=tenant_id, + app_config=app_config, + inputs=inputs, + query=query, + message_id=message_id, + trace_scope=None, ) assert flagged is False @@ -78,46 +97,86 @@ def test_check_not_flagged(self, mock_factory_cls, app_config, input_moderation) ) mock_factory.moderation_for_inputs.assert_called_once_with(dict(inputs), query) - @patch("core.moderation.input_moderation.ModerationFactory") - @patch("core.moderation.input_moderation.TraceTask") - def test_check_with_trace_manager( - self, mock_trace_task, mock_factory_cls, app_config, input_moderation, trace_manager - ): - app_id = "test_app_id" - tenant_id = "test_tenant_id" - inputs = {"input_key": "input_value"} - query = "test query" - message_id = "test_message_id" + def test_check_without_sensitive_word_avoidance_emits_nothing( + self, app_config: AppConfig, input_moderation: InputModeration + ) -> None: + recorder = _Recorder() - # Setup config - sensitive_word_config = SensitiveWordAvoidanceEntity(type="keywords", config={}) - app_config.sensitive_word_avoidance = sensitive_word_config + input_moderation.check( + app_id="test_app_id", + tenant_id="test_tenant_id", + app_config=app_config, + inputs={"input_key": "input_value"}, + query="test query", + message_id="test_message_id", + trace_scope=_trace_scope(recorder), + ) - # Setup factory mock - mock_factory = mock_factory_cls.return_value - mock_result = ModerationInputsResult(flagged=False, action=ModerationAction.DIRECT_OUTPUT) - mock_factory.moderation_for_inputs.return_value = mock_result + assert recorder.events == [] + + @patch("core.moderation.input_moderation.ModerationFactory") + def test_check_with_trace_scope_emits_moderation_checked( + self, mock_factory_cls: MagicMock, app_config: AppConfig, input_moderation: InputModeration + ) -> None: + app_config.sensitive_word_avoidance = SensitiveWordAvoidanceEntity(type="keywords", config={}) + mock_factory_cls.return_value.moderation_for_inputs.return_value = ModerationInputsResult( + flagged=False, action=ModerationAction.DIRECT_OUTPUT + ) + recorder = _Recorder() input_moderation.check( - app_id=app_id, - tenant_id=tenant_id, + app_id="test_app_id", + tenant_id="test_tenant_id", app_config=app_config, - inputs=inputs, - query=query, - message_id=message_id, - trace_manager=trace_manager, + inputs={"input_key": "input_value", "budget": Decimal("1.5")}, + query="test query", + message_id="test_message_id", + trace_scope=_trace_scope(recorder), ) - trace_manager.add_trace_task.assert_called_once_with(mock_trace_task.return_value) - assert trace_manager.collect_tasks() == [mock_trace_task.return_value] - mock_trace_task.assert_called_once() - call_kwargs = mock_trace_task.call_args.kwargs - call_args = mock_trace_task.call_args.args - assert call_args[0] == TraceTaskName.MODERATION_TRACE - assert call_kwargs["message_id"] == message_id - assert call_kwargs["moderation_result"] == mock_result - assert call_kwargs["inputs"] == inputs - assert "timer" in call_kwargs + [event] = recorder.events + payload = event.payload + assert isinstance(payload, ModerationChecked) + assert payload == ModerationChecked( + message_id="test_message_id", + inputs={"input_key": "input_value", "budget": "1.5"}, + query="", + flagged=False, + action="direct_output", + preset_response="", + started_at=payload.started_at, + ended_at=payload.ended_at, + ) + assert payload.started_at.utcoffset() == timedelta(0) + assert payload.ended_at >= payload.started_at + + @patch("core.moderation.input_moderation.ModerationFactory") + def test_check_flagged_direct_output_emits_before_raising( + self, mock_factory_cls: MagicMock, app_config: AppConfig, input_moderation: InputModeration + ) -> None: + app_config.sensitive_word_avoidance = SensitiveWordAvoidanceEntity(type="keywords", config={}) + mock_factory_cls.return_value.moderation_for_inputs.return_value = ModerationInputsResult( + flagged=True, action=ModerationAction.DIRECT_OUTPUT, preset_response="Blocked content" + ) + recorder = _Recorder() + + with pytest.raises(ModerationError, match="Blocked content"): + input_moderation.check( + app_id="test_app_id", + tenant_id="test_tenant_id", + app_config=app_config, + inputs={"input_key": "input_value"}, + query="test query", + message_id="test_message_id", + trace_scope=_trace_scope(recorder), + ) + + [event] = recorder.events + payload = event.payload + assert isinstance(payload, ModerationChecked) + assert payload.flagged is True + assert payload.action == "direct_output" + assert payload.preset_response == "Blocked content" @patch("core.moderation.input_moderation.ModerationFactory") def test_check_flagged_direct_output(self, mock_factory_cls, app_config, input_moderation): @@ -146,6 +205,7 @@ def test_check_flagged_direct_output(self, mock_factory_cls, app_config, input_m inputs=inputs, query=query, message_id=message_id, + trace_scope=None, ) assert str(excinfo.value) == "Blocked content" @@ -173,7 +233,13 @@ def test_check_flagged_overridden(self, mock_factory_cls, app_config, input_mode mock_factory.moderation_for_inputs.return_value = mock_result flagged, final_inputs, final_query = input_moderation.check( - app_id=app_id, tenant_id=tenant_id, app_config=app_config, inputs=inputs, query=query, message_id=message_id + app_id=app_id, + tenant_id=tenant_id, + app_config=app_config, + inputs=inputs, + query=query, + message_id=message_id, + trace_scope=None, ) assert flagged is True @@ -206,6 +272,7 @@ def test_check_flagged_other_action(self, mock_factory_cls, app_config, input_mo inputs=inputs, query=query, message_id=message_id, + trace_scope=None, ) assert flagged is True diff --git a/api/tests/unit_tests/core/ops/__init__.py b/api/tests/unit_tests/core/ops/__init__.py deleted file mode 100644 index bb92ccdec744c0..00000000000000 --- a/api/tests/unit_tests/core/ops/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Unit tests for core ops module diff --git a/api/tests/unit_tests/core/ops/entities/test_config_entity.py b/api/tests/unit_tests/core/ops/entities/test_config_entity.py deleted file mode 100644 index 15ff08c0e6a319..00000000000000 --- a/api/tests/unit_tests/core/ops/entities/test_config_entity.py +++ /dev/null @@ -1,75 +0,0 @@ -from dify_trace_aliyun.config import AliyunConfig -from dify_trace_arize_phoenix.config import ArizeConfig, PhoenixConfig -from dify_trace_langfuse.config import LangfuseConfig -from dify_trace_langsmith.config import LangSmithConfig -from dify_trace_opik.config import OpikConfig -from dify_trace_weave.config import WeaveConfig - -from core.ops.entities.config_entity import TracingProviderEnum - - -class TestTracingProviderEnum: - """Test cases for TracingProviderEnum""" - - def test_enum_values(self): - """Test that all expected enum values are present""" - assert TracingProviderEnum.ARIZE == "arize" - assert TracingProviderEnum.PHOENIX == "phoenix" - assert TracingProviderEnum.LANGFUSE == "langfuse" - assert TracingProviderEnum.LANGSMITH == "langsmith" - assert TracingProviderEnum.OPIK == "opik" - assert TracingProviderEnum.WEAVE == "weave" - assert TracingProviderEnum.ALIYUN == "aliyun" - - -class TestConfigIntegration: - """Cross-provider configuration sanity checks""" - - def test_all_configs_can_be_instantiated(self): - """Test that all config classes can be instantiated with valid data""" - configs = [ - ArizeConfig(api_key="key"), - PhoenixConfig(api_key="key"), - LangfuseConfig(public_key="public", secret_key="secret"), - LangSmithConfig(api_key="key", project="project"), - OpikConfig(api_key="key"), - WeaveConfig(api_key="key", project="project"), - AliyunConfig(license_key="test_license", endpoint="https://tracing-analysis-dc-hz.aliyuncs.com"), - ] - - for config in configs: - assert config is not None - - def test_url_normalization_consistency(self): - """Test that URL normalization works consistently across configs""" - arize_config = ArizeConfig(endpoint="https://arize.com/api/v1/test") - phoenix_with_path_config = PhoenixConfig(endpoint="https://app.phoenix.arize.com/s/dify-integration") - phoenix_without_path_config = PhoenixConfig(endpoint="https://app.phoenix.arize.com") - aliyun_config = AliyunConfig( - license_key="test_license", endpoint="https://tracing-analysis-dc-hz.aliyuncs.com/api/v1/traces" - ) - langsmith_config = LangSmithConfig( - api_key="key", project="project", endpoint="https://langsmith.internal/api/v1" - ) - weave_config = WeaveConfig(api_key="key", project="project", endpoint="https://weave.internal/wandb") - - assert arize_config.endpoint == "https://arize.com" - assert phoenix_with_path_config.endpoint == "https://app.phoenix.arize.com/s/dify-integration" - assert phoenix_without_path_config.endpoint == "https://app.phoenix.arize.com" - assert aliyun_config.endpoint == "https://tracing-analysis-dc-hz.aliyuncs.com/api/v1/traces" - assert langsmith_config.endpoint == "https://langsmith.internal/api/v1" - assert weave_config.endpoint == "https://weave.internal/wandb" - - def test_project_default_values(self): - """Test that project default values are set correctly""" - arize_config = ArizeConfig(project="") - phoenix_config = PhoenixConfig(project="") - opik_config = OpikConfig(project="") - aliyun_config = AliyunConfig( - license_key="test_license", endpoint="https://tracing-analysis-dc-hz.aliyuncs.com", app_name="" - ) - - assert arize_config.project == "default" - assert phoenix_config.project == "default" - assert opik_config.project == "Default Project" - assert aliyun_config.app_name == "dify_app" diff --git a/api/tests/unit_tests/core/ops/entities/test_trace_entity.py b/api/tests/unit_tests/core/ops/entities/test_trace_entity.py deleted file mode 100644 index cd93af24c6fe1c..00000000000000 --- a/api/tests/unit_tests/core/ops/entities/test_trace_entity.py +++ /dev/null @@ -1,168 +0,0 @@ -import json -from datetime import datetime, timedelta -from types import SimpleNamespace -from unittest.mock import MagicMock - -import pytest -from sqlalchemy import Engine -from sqlalchemy.orm import Session, scoped_session, sessionmaker - -from core.ops.entities.trace_entity import TraceTaskName -from core.ops.ops_trace_manager import TraceTask -from graphon.enums import WorkflowExecutionStatus -from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom -from models.model import App, AppMode, Conversation, Message, MessageFile -from models.workflow import WorkflowAppLog, WorkflowNodeExecutionModel, WorkflowRun, WorkflowType - -TABLES = (App, Conversation, Message, MessageFile, WorkflowAppLog, WorkflowNodeExecutionModel) - - -@pytest.fixture(autouse=True) -def _bind_trace_database( - monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, - sqlite_session_factory: sessionmaker[Session], -) -> None: - """Use real SQLite sessions for ORM lookups without changing trace-domain data.""" - monkeypatch.setattr( - "core.ops.ops_trace_manager.db", - SimpleNamespace(engine=sqlite_engine, session=scoped_session(sqlite_session_factory)), - ) - - -def _make_workflow_run() -> WorkflowRun: - workflow_run = WorkflowRun( - workflow_id="wf-1", - tenant_id="tenant-1", - id="run-1", - elapsed_time=1.0, - status=WorkflowExecutionStatus.SUCCEEDED, - inputs="{}", - outputs="{}", - graph="{}", - type=WorkflowType.WORKFLOW, - version="1", - error=None, - total_tokens=0, - total_steps=0, - finished_at=datetime(2026, 1, 1, 0, 0, 1), - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - app_id="app-1", - created_by_role=CreatorUserRole.END_USER, - created_by="user-1", - exceptions_count=0, - ) - workflow_run.created_at = datetime(2026, 1, 1, 0, 0, 0) - return workflow_run - - -def _make_message_data(): - created_at = datetime(2026, 1, 1, 0, 0, 0) - data = { - "id": "message-1", - "app_id": "app-1", - "conversation_id": "conv-1", - "created_at": created_at, - "updated_at": created_at + timedelta(seconds=1), - "message": "hello", - "provider_response_latency": 1, - "message_tokens": 0, - "answer_tokens": 0, - "answer": "world", - "error": "", - "status": "normal", - "model_provider": "provider", - "model_id": "model", - "from_end_user_id": "end-user-1", - "from_account_id": None, - "agent_based": False, - "workflow_run_id": None, - "from_source": "api", - "message_metadata": json.dumps({"usage": {}}), - } - - class _MessageData: - def __init__(self, values): - self.__dict__.update(values) - - def to_dict(self, **_kwargs: object): - return dict(self.__dict__) - - return _MessageData(data) - - -@pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) -def test_workflow_trace_metadata_includes_trace_session_id(monkeypatch, sqlite_session: Session): - repo = MagicMock() - repo.get_workflow_run_by_id_without_tenant.return_value = _make_workflow_run() - monkeypatch.setattr(TraceTask, "_get_workflow_run_repo", classmethod(lambda cls: repo)) - monkeypatch.setattr("core.telemetry.gateway.is_enterprise_telemetry_enabled", lambda: False) - - task = TraceTask( - TraceTaskName.WORKFLOW_TRACE, - workflow_execution=SimpleNamespace(id_="run-1", total_tokens=0), - conversation_id="conv-1", - user_id="user-1", - trace_session_id="session-1", - ) - - trace_info = task.workflow_trace(workflow_run_id="run-1", conversation_id="conv-1", user_id="user-1") - - assert task.kwargs["trace_session_id"] == "session-1" - assert trace_info.metadata["trace_session_id"] == "session-1" - - -@pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) -def test_message_trace_metadata_includes_trace_session_id(monkeypatch, sqlite_session: Session): - app = App( - id="app-1", - tenant_id="tenant-1", - name="Trace App", - description="", - mode=AppMode.CHAT, - icon_type=None, - icon=None, - icon_background=None, - app_model_config_id=None, - workflow_id=None, - enable_site=True, - enable_api=True, - max_active_requests=None, - created_by=None, - ) - conversation = Conversation( - id="conv-1", - app_id=app.id, - app_model_config_id=None, - model_provider=None, - override_model_configs=None, - model_id=None, - mode=AppMode.CHAT, - name="Trace Conversation", - summary=None, - _inputs={}, - introduction="", - system_instruction="", - invoke_from=None, - from_source="api", - from_end_user_id="end-user-1", - from_account_id=None, - read_at=None, - read_account_id=None, - ) - sqlite_session.add_all([app, conversation]) - sqlite_session.commit() - - monkeypatch.setattr("core.ops.ops_trace_manager.get_message_data", lambda message_id, session: _make_message_data()) - monkeypatch.setattr("core.telemetry.gateway.is_enterprise_telemetry_enabled", lambda: False) - - task = TraceTask( - TraceTaskName.MESSAGE_TRACE, - message_id="message-1", - trace_session_id="session-1", - ) - - trace_info = task.message_trace(message_id="message-1", **task.kwargs) - - assert task.kwargs["trace_session_id"] == "session-1" - assert trace_info.metadata["trace_session_id"] == "session-1" diff --git a/api/tests/unit_tests/core/ops/test_base_trace_instance.py b/api/tests/unit_tests/core/ops/test_base_trace_instance.py deleted file mode 100644 index fdf8a4ff54f7eb..00000000000000 --- a/api/tests/unit_tests/core/ops/test_base_trace_instance.py +++ /dev/null @@ -1,264 +0,0 @@ -from contextvars import Context -from types import SimpleNamespace - -import pytest -from flask import Flask, has_app_context -from sqlalchemy import Engine, select -from sqlalchemy.orm import Session - -from core.ops.base_trace_instance import ( - ApplicationTracingHooksNotRegisteredError, - BaseTraceInstance, - application_tracing_hooks, - bind_application_tracing_hooks, -) -from core.ops.entities.config_entity import BaseTracingConfig -from core.ops.entities.trace_entity import BaseTraceInfo -from models import Account, App, Tenant, TenantAccountJoin -from models.account import TenantAccountRole -from models.model import AppMode -from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks -from tests.helpers.unit.config_override import apply_config_overrides - -"""The execution port is gated before Flask lookup and bound per application.""" - - -class TestApplicationTracingHooksCases: - @staticmethod - def test_disabled_port_needs_no_application_context(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) - isolated_context = Context() - assert not isolated_context.run(has_app_context) - assert isolated_context.run(application_tracing_hooks) is None - - @staticmethod - def test_enabled_port_reports_missing_registration_without_context(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - isolated_context = Context() - assert not isolated_context.run(has_app_context) - with pytest.raises(ApplicationTracingHooksNotRegisteredError, match="application_tracing_hooks_not_registered"): - isolated_context.run(application_tracing_hooks) - - @staticmethod - def test_enabled_port_reports_missing_registration_in_app(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - with Flask(__name__).app_context(): - with pytest.raises( - ApplicationTracingHooksNotRegisteredError, match="application_tracing_hooks_not_registered" - ): - application_tracing_hooks() - - @staticmethod - def test_port_does_not_leak_between_flask_applications(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - first, second = Flask("first"), Flask("second") - first_hooks = RuntimeApplicationTracingHooks() - second_hooks = RuntimeApplicationTracingHooks() - bind_application_tracing_hooks(first, first_hooks) - bind_application_tracing_hooks(second, second_hooks) - with first.app_context(): - assert application_tracing_hooks() is first_hooks - with second.app_context(): - assert application_tracing_hooks() is second_hooks - assert application_tracing_hooks() is first_hooks - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) - assert application_tracing_hooks() is None - - -TABLES = (App, Account, Tenant, TenantAccountJoin) - - -class ConcreteTraceInstance(BaseTraceInstance): - def __init__(self, trace_config: BaseTracingConfig): - super().__init__(trace_config) - - def trace(self, trace_info: BaseTraceInfo): - super().trace(trace_info) - - -def _persist_app(session: Session, *, created_by: str | None) -> App: - app = App( - id="some_app_id", - tenant_id="tenant_id", - name="Trace App", - description="", - mode=AppMode.CHAT, - icon_type=None, - icon=None, - icon_background=None, - app_model_config_id=None, - workflow_id=None, - enable_site=True, - enable_api=True, - max_active_requests=None, - created_by=created_by, - ) - session.add(app) - session.commit() - return app - - -def _persist_account(session: Session) -> Account: - account = Account(name="Creator", email="creator@example.com") - account.id = "creator_id" - session.add(account) - session.commit() - return account - - -def _trace_instance() -> ConcreteTraceInstance: - return ConcreteTraceInstance(BaseTracingConfig()) - - -class TestBaseTraceInstanceCases: - @staticmethod - @pytest.fixture(autouse=True) - def _bind_production_sessions(monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine) -> None: - """Bind the service-owned ORM session to the test's SQLite engine.""" - database = SimpleNamespace(engine=sqlite_engine) - monkeypatch.setattr("core.ops.base_trace_instance.db", database) - - @staticmethod - @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) - def test_get_service_account_with_tenant_app_not_found(sqlite_session: Session): - with pytest.raises(ValueError, match="App with id some_app_id not found"): - _trace_instance().get_service_account_with_tenant("some_app_id") - - @staticmethod - @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) - def test_get_service_account_with_tenant_no_creator(sqlite_session: Session): - _persist_app(sqlite_session, created_by=None) - - with pytest.raises(ValueError, match="App with id some_app_id has no creator"): - _trace_instance().get_service_account_with_tenant("some_app_id") - - @staticmethod - @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) - def test_get_service_account_with_tenant_creator_not_found(sqlite_session: Session): - _persist_app(sqlite_session, created_by="creator_id") - - with pytest.raises(ValueError, match="Creator account with id creator_id not found for app some_app_id"): - _trace_instance().get_service_account_with_tenant("some_app_id") - - @staticmethod - @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) - def test_get_service_account_with_tenant_tenant_not_found(sqlite_session: Session): - _persist_app(sqlite_session, created_by="creator_id") - _persist_account(sqlite_session) - - with pytest.raises(ValueError, match="App tenant tenant_id not found for account creator_id"): - _trace_instance().get_service_account_with_tenant("some_app_id") - - @staticmethod - @pytest.mark.parametrize("current", [True, False]) - @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) - def test_get_service_account_with_tenant_success(sqlite_session: Session, current: bool): - _persist_app(sqlite_session, created_by="creator_id") - _persist_account(sqlite_session) - tenant = Tenant(name="Workspace") - tenant.id = "tenant_id" - sqlite_session.add_all( - [ - tenant, - TenantAccountJoin( - tenant_id=tenant.id, - account_id="creator_id", - current=current, - role=TenantAccountRole.OWNER, - ), - ] - ) - sqlite_session.commit() - - result = _trace_instance().get_service_account_with_tenant("some_app_id") - - assert result.id == "creator_id" - assert result.current_tenant_id == "tenant_id" - assert result.current_role == TenantAccountRole.OWNER - - @staticmethod - @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) - def test_get_service_account_with_tenant_uses_app_workspace(sqlite_session: Session): - _persist_app(sqlite_session, created_by="creator_id") - _persist_account(sqlite_session) - app_tenant = Tenant(name="App workspace") - app_tenant.id = "tenant_id" - current_tenant = Tenant(name="Current workspace") - current_tenant.id = "other_tenant_id" - sqlite_session.add_all( - [ - app_tenant, - current_tenant, - TenantAccountJoin( - tenant_id="tenant_id", - account_id="creator_id", - current=False, - role=TenantAccountRole.EDITOR, - ), - TenantAccountJoin( - tenant_id="other_tenant_id", - account_id="creator_id", - current=True, - role=TenantAccountRole.OWNER, - ), - ] - ) - sqlite_session.commit() - - result = _trace_instance().get_service_account_with_tenant("some_app_id") - - assert result.id == "creator_id" - assert result.current_tenant_id == "tenant_id" - assert result.current_role == TenantAccountRole.EDITOR - sqlite_session.expire_all() - memberships = sqlite_session.scalars( - select(TenantAccountJoin).where(TenantAccountJoin.account_id == "creator_id") - ).all() - assert {membership.tenant_id: membership.current for membership in memberships} == { - "tenant_id": False, - "other_tenant_id": True, - } - - @staticmethod - @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) - def test_get_service_account_with_tenant_requires_app_membership(sqlite_session: Session): - _persist_app(sqlite_session, created_by="creator_id") - _persist_account(sqlite_session) - app_tenant = Tenant(name="App workspace") - app_tenant.id = "tenant_id" - current_tenant = Tenant(name="Current workspace") - current_tenant.id = "other_tenant_id" - sqlite_session.add_all( - [ - app_tenant, - current_tenant, - TenantAccountJoin( - tenant_id="other_tenant_id", - account_id="creator_id", - current=True, - role=TenantAccountRole.OWNER, - ), - ] - ) - sqlite_session.commit() - - with pytest.raises(ValueError, match="App tenant tenant_id not found for account creator_id"): - _trace_instance().get_service_account_with_tenant("some_app_id") - - @staticmethod - @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) - def test_get_service_account_with_tenant_requires_existing_app_tenant(sqlite_session: Session): - _persist_app(sqlite_session, created_by="creator_id") - _persist_account(sqlite_session) - sqlite_session.add( - TenantAccountJoin( - tenant_id="tenant_id", - account_id="creator_id", - current=True, - role=TenantAccountRole.OWNER, - ) - ) - sqlite_session.commit() - - with pytest.raises(ValueError, match="App tenant tenant_id not found for account creator_id"): - _trace_instance().get_service_account_with_tenant("some_app_id") 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 deleted file mode 100644 index 34783f46478c21..00000000000000 --- a/api/tests/unit_tests/core/ops/test_ops_trace_manager.py +++ /dev/null @@ -1,1685 +0,0 @@ -from __future__ import annotations - -import builtins -import json -import queue -import sys -import uuid -from collections.abc import Iterator -from contextlib import contextmanager -from datetime import datetime, timedelta -from decimal import Decimal -from types import SimpleNamespace -from unittest.mock import MagicMock, PropertyMock, patch -from uuid import UUID, uuid4 - -import pytest -from Crypto.PublicKey import RSA -from flask import Flask -from pydantic import ValidationError -from sqlalchemy import Engine, event -from sqlalchemy.orm import Session, scoped_session, sessionmaker - -import core.ops.ops_trace_manager as module -import models.engine as model_engine -from core.moderation.base import ModerationAction, ModerationInputsResult -from core.ops.base_trace_instance import bind_application_tracing_hooks -from core.ops.exceptions import TraceProviderNotInstalledError -from core.ops.ops_trace_manager import OpsTraceManager, TraceQueueManager, TraceTask, TraceTaskName -from core.rag.models.document import Document as RetrievalDocument -from core.tools.entities.tool_entities import ApiProviderSchemaType -from extensions.ext_database import db -from extensions.ext_key_provider import key_provider_manager -from graphon.enums import WorkflowExecutionStatus -from graphon.file import FileTransferMethod, FileType -from graphon.model_runtime.entities.model_entities import ModelType -from libs import gmpy2_pkcs10aep_cipher, rsa -from libs.key_providers.rsa_key_provider import RSAKeyProvider -from models.account import Tenant -from models.base import TypeBase -from models.enums import ConversationFromSource, CreatorUserRole, MessageStatus, WorkflowRunTriggeredFrom -from models.model import ( - App, - AppMode, - AppModelConfig, - Conversation, - IconType, - Message, - MessageAgentThought, - MessageFile, - TraceAppConfig, -) -from models.provider import Provider, ProviderCredential, ProviderModel, ProviderModelCredential, ProviderType -from models.tools import ApiToolProvider, BuiltinToolProvider, MCPToolProvider, WorkflowToolProvider -from models.workflow import WorkflowAppLog, WorkflowAppLogCreatedFrom, WorkflowRun, WorkflowType -from repositories.sqlalchemy_api_workflow_run_repository import DifyAPISQLAlchemyWorkflowRunRepository -from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks -from tests.helpers.unit.config_override import apply_config_overrides -from tests.helpers.unit.model_factories import make_app, make_tenant - -"""Unit tests for lookup helper functions in core.ops.ops_trace_manager. - -Covers: -- _lookup_app_and_workspace_names -- _lookup_credential_name -- _lookup_llm_credential_info -- TraceTask._get_user_id_from_metadata -""" - - -def _persist_app(session: Session, *, tenant_id: str, name: str = "MyApp") -> App: - app = App( - id=str(uuid.uuid4()), - tenant_id=tenant_id, - name=name, - mode=AppMode.WORKFLOW, - icon_type=IconType.EMOJI, - icon="workflow", - icon_background="#FFFFFF", - enable_site=True, - enable_api=False, - ) - session.add(app) - session.commit() - return app - - -def _persist_tenant(session: Session, *, name: str = "MyWorkspace") -> Tenant: - tenant = Tenant(name=name) - session.add(tenant) - session.commit() - return tenant - - -def _persist_tool_provider( - session: Session, provider_type: str -) -> BuiltinToolProvider | ApiToolProvider | WorkflowToolProvider | MCPToolProvider: - tenant_id = str(uuid.uuid4()) - user_id = str(uuid.uuid4()) - if provider_type in {"builtin", "plugin"}: - provider = BuiltinToolProvider( - name="CredentialA", - tenant_id=tenant_id, - user_id=user_id, - provider="test/provider", - ) - elif provider_type == "api": - provider = ApiToolProvider( - name="CredentialA", - icon="icon.svg", - schema="{}", - schema_type_str=ApiProviderSchemaType.OPENAPI, - user_id=user_id, - tenant_id=tenant_id, - description="API provider", - tools_str="[]", - credentials_str="{}", - ) - elif provider_type == "workflow": - provider = WorkflowToolProvider( - name="CredentialA", - label="CredentialA", - icon="icon.svg", - app_id=str(uuid.uuid4()), - version="1", - user_id=user_id, - tenant_id=tenant_id, - description="Workflow provider", - ) - elif provider_type == "mcp": - provider = MCPToolProvider( - name="CredentialA", - server_identifier="credential-a", - server_url="https://example.com/mcp", - server_url_hash="credential-a-hash", - icon="icon.svg", - tenant_id=tenant_id, - user_id=user_id, - ) - else: - raise ValueError(f"unsupported provider type: {provider_type}") - - session.add(provider) - session.commit() - return provider - - -def _persist_provider_credential( - session: Session, - *, - tenant_id: str, - credential_name: str = "ProvCredName", -) -> ProviderCredential: - credential = ProviderCredential( - tenant_id=tenant_id, - provider_name="openai", - credential_name=credential_name, - encrypted_config="{}", - ) - session.add(credential) - session.commit() - return credential - - -def _persist_model_credential( - session: Session, - *, - tenant_id: str, - credential_name: str = "ModelCredName", -) -> ProviderModelCredential: - credential = ProviderModelCredential( - tenant_id=tenant_id, - provider_name="openai", - model_name="gpt-4", - model_type=ModelType.LLM, - credential_name=credential_name, - encrypted_config="{}", - ) - session.add(credential) - session.commit() - return credential - - -def _persist_provider( - session: Session, - *, - tenant_id: str, - credential_id: str | None, -) -> Provider: - provider = Provider( - tenant_id=tenant_id, - provider_name="openai", - provider_type=ProviderType.CUSTOM, - credential_id=credential_id, - ) - session.add(provider) - session.commit() - return provider - - -def _persist_provider_model( - session: Session, - *, - tenant_id: str, - credential_id: str | None, -) -> ProviderModel: - model = ProviderModel( - tenant_id=tenant_id, - provider_name="openai", - model_name="gpt-4", - model_type=ModelType.LLM, - credential_id=credential_id, - ) - session.add(model) - session.commit() - return model - - -@contextmanager -def _raise_on_table(engine: Engine, table_name: str) -> Iterator[None]: - """Raise only when SQL targets the named table, leaving other real lookups intact.""" - - def fail_target_query(_conn, _cursor, statement, _parameters, _context, _executemany): - if f"FROM {table_name}" in statement: - raise RuntimeError(f"forced failure for {table_name}") - - event.listen(engine, "before_cursor_execute", fail_target_query) - try: - yield - finally: - event.remove(engine, "before_cursor_execute", fail_target_query) - - -class TestLookupHelpersCases: - # --------------------------------------------------------------------------- - # Helpers - # --------------------------------------------------------------------------- - - @staticmethod - @pytest.fixture - def orm_session(sqlite_engine: Engine) -> Iterator[Session]: - models = ( - App, - Tenant, - Provider, - ProviderCredential, - ProviderModel, - ProviderModelCredential, - BuiltinToolProvider, - ApiToolProvider, - WorkflowToolProvider, - MCPToolProvider, - ) - tables = [model.metadata.tables[model.__tablename__] for model in models] - TypeBase.metadata.create_all(sqlite_engine, tables=tables) - - with patch.object(type(db), "engine", new_callable=PropertyMock, return_value=sqlite_engine): - with Session(sqlite_engine, expire_on_commit=False) as session: - yield session - - # --------------------------------------------------------------------------- - # _lookup_app_and_workspace_names - # --------------------------------------------------------------------------- - - class TestLookupAppAndWorkspaceNames: - """Tests for _lookup_app_and_workspace_names(app_id, tenant_id).""" - - def test_both_found(self, orm_session: Session): - """Returns (app_name, workspace_name) when both records exist.""" - from core.ops.ops_trace_manager import _lookup_app_and_workspace_names - - tenant = _persist_tenant(orm_session) - app = _persist_app(orm_session, tenant_id=tenant.id) - app_name, workspace_name = _lookup_app_and_workspace_names(app.id, tenant.id) - - assert app_name == "MyApp" - assert workspace_name == "MyWorkspace" - - def test_app_only_found(self, orm_session: Session): - """Returns (app_name, '') when tenant record is absent.""" - from core.ops.ops_trace_manager import _lookup_app_and_workspace_names - - app = _persist_app(orm_session, tenant_id=str(uuid.uuid4())) - app_name, workspace_name = _lookup_app_and_workspace_names(app.id, str(uuid.uuid4())) - - assert app_name == "MyApp" - assert workspace_name == "" - - def test_tenant_only_found(self, orm_session: Session): - """Returns ('', workspace_name) when app record is absent.""" - from core.ops.ops_trace_manager import _lookup_app_and_workspace_names - - tenant = _persist_tenant(orm_session) - app_name, workspace_name = _lookup_app_and_workspace_names(str(uuid.uuid4()), tenant.id) - - assert app_name == "" - assert workspace_name == "MyWorkspace" - - def test_neither_found(self, orm_session: Session): - """Returns ('', '') when both DB lookups return None.""" - from core.ops.ops_trace_manager import _lookup_app_and_workspace_names - - app_name, workspace_name = _lookup_app_and_workspace_names(str(uuid.uuid4()), str(uuid.uuid4())) - - assert app_name == "" - assert workspace_name == "" - - def test_none_inputs_skips_db(self): - """Returns ('', '') immediately when both IDs are None — no DB access.""" - from core.ops.ops_trace_manager import _lookup_app_and_workspace_names - - app_name, workspace_name = _lookup_app_and_workspace_names(None, None) - - assert app_name == "" - assert workspace_name == "" - - def test_app_id_none_only_queries_tenant(self, orm_session: Session): - """When app_id is None, only the tenant query is issued.""" - from core.ops.ops_trace_manager import _lookup_app_and_workspace_names - - tenant = _persist_tenant(orm_session, name="OnlyWorkspace") - app_name, workspace_name = _lookup_app_and_workspace_names(None, tenant.id) - - assert app_name == "" - assert workspace_name == "OnlyWorkspace" - - def test_tenant_id_none_only_queries_app(self, orm_session: Session): - """When tenant_id is None, only the app query is issued.""" - from core.ops.ops_trace_manager import _lookup_app_and_workspace_names - - app = _persist_app(orm_session, tenant_id=str(uuid.uuid4()), name="OnlyApp") - app_name, workspace_name = _lookup_app_and_workspace_names(app.id, None) - - assert app_name == "OnlyApp" - assert workspace_name == "" - - # --------------------------------------------------------------------------- - # _lookup_credential_name - # --------------------------------------------------------------------------- - - class TestLookupCredentialName: - """Tests for _lookup_credential_name(credential_id, provider_type).""" - - @pytest.mark.parametrize("provider_type", ["builtin", "plugin", "api", "workflow", "mcp"]) - def test_known_provider_types_return_name(self, provider_type: str, orm_session: Session): - """Each valid provider_type results in a DB query and returns the credential name.""" - from core.ops.ops_trace_manager import _lookup_credential_name - - provider = _persist_tool_provider(orm_session, provider_type) - result = _lookup_credential_name(provider.id, provider_type) - - assert result == "CredentialA" - - def test_credential_not_found_returns_empty_string(self, orm_session: Session): - """Returns '' when DB yields None for the given credential_id.""" - from core.ops.ops_trace_manager import _lookup_credential_name - - result = _lookup_credential_name(str(uuid.uuid4()), "api") - - assert result == "" - - def test_invalid_provider_type_returns_empty_string_without_db(self): - """Returns '' immediately for an unrecognised provider_type — no DB access.""" - from core.ops.ops_trace_manager import _lookup_credential_name - - result = _lookup_credential_name(str(uuid.uuid4()), "unknown_type") - - assert result == "" - - def test_none_credential_id_returns_empty_string_without_db(self): - """Returns '' immediately when credential_id is None — no DB access.""" - from core.ops.ops_trace_manager import _lookup_credential_name - - result = _lookup_credential_name(None, "api") - - assert result == "" - - def test_none_provider_type_returns_empty_string_without_db(self): - """Returns '' immediately when provider_type is None — no DB access.""" - from core.ops.ops_trace_manager import _lookup_credential_name - - result = _lookup_credential_name(str(uuid.uuid4()), None) - - assert result == "" - - def test_builtin_and_plugin_map_to_same_model(self): - """Both 'builtin' and 'plugin' provider_types query BuiltinToolProvider.""" - from core.ops.ops_trace_manager import _PROVIDER_TYPE_TO_MODEL - from models.tools import BuiltinToolProvider - - assert _PROVIDER_TYPE_TO_MODEL["builtin"] is BuiltinToolProvider - assert _PROVIDER_TYPE_TO_MODEL["plugin"] is BuiltinToolProvider - - def test_api_maps_to_api_tool_provider(self): - """'api' maps to ApiToolProvider.""" - from core.ops.ops_trace_manager import _PROVIDER_TYPE_TO_MODEL - from models.tools import ApiToolProvider - - assert _PROVIDER_TYPE_TO_MODEL["api"] is ApiToolProvider - - def test_workflow_maps_to_workflow_tool_provider(self): - """'workflow' maps to WorkflowToolProvider.""" - from core.ops.ops_trace_manager import _PROVIDER_TYPE_TO_MODEL - from models.tools import WorkflowToolProvider - - assert _PROVIDER_TYPE_TO_MODEL["workflow"] is WorkflowToolProvider - - def test_mcp_maps_to_mcp_tool_provider(self): - """'mcp' maps to MCPToolProvider.""" - from core.ops.ops_trace_manager import _PROVIDER_TYPE_TO_MODEL - from models.tools import MCPToolProvider - - assert _PROVIDER_TYPE_TO_MODEL["mcp"] is MCPToolProvider - - # --------------------------------------------------------------------------- - # _lookup_llm_credential_info - # --------------------------------------------------------------------------- - - class TestLookupLlmCredentialInfo: - """Tests for _lookup_llm_credential_info(tenant_id, provider, model, model_type).""" - - def test_model_level_credential_found(self, orm_session: Session): - """Returns model-level credential_id and name when ProviderModel has a credential.""" - from core.ops.ops_trace_manager import _lookup_llm_credential_info - - tenant_id = str(uuid.uuid4()) - model_credential = _persist_model_credential(orm_session, tenant_id=tenant_id) - _persist_provider(orm_session, tenant_id=tenant_id, credential_id=None) - _persist_provider_model(orm_session, tenant_id=tenant_id, credential_id=model_credential.id) - - decoy_tenant_id = str(uuid.uuid4()) - decoy_credential = _persist_model_credential( - orm_session, - tenant_id=decoy_tenant_id, - credential_name="WrongTenantCredential", - ) - _persist_provider(orm_session, tenant_id=decoy_tenant_id, credential_id=None) - _persist_provider_model(orm_session, tenant_id=decoy_tenant_id, credential_id=decoy_credential.id) - - cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4") - - assert cred_id == model_credential.id - assert cred_name == "ModelCredName" - - def test_provider_level_fallback_when_no_model_credential(self, orm_session: Session): - """Falls back to provider-level credential when ProviderModel has no credential_id.""" - from core.ops.ops_trace_manager import _lookup_llm_credential_info - - tenant_id = str(uuid.uuid4()) - provider_credential = _persist_provider_credential(orm_session, tenant_id=tenant_id) - _persist_provider(orm_session, tenant_id=tenant_id, credential_id=provider_credential.id) - _persist_provider_model(orm_session, tenant_id=tenant_id, credential_id=None) - - cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4") - - assert cred_id == provider_credential.id - assert cred_name == "ProvCredName" - - def test_provider_level_fallback_when_no_model_record(self, orm_session: Session): - """Falls back to provider-level credential when no ProviderModel row exists.""" - from core.ops.ops_trace_manager import _lookup_llm_credential_info - - tenant_id = str(uuid.uuid4()) - provider_credential = _persist_provider_credential(orm_session, tenant_id=tenant_id) - _persist_provider(orm_session, tenant_id=tenant_id, credential_id=provider_credential.id) - - cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4") - - assert cred_id == provider_credential.id - assert cred_name == "ProvCredName" - - def test_no_model_arg_uses_provider_level_only(self, orm_session: Session): - """When model is None, skips ProviderModel query and uses provider credential.""" - from core.ops.ops_trace_manager import _lookup_llm_credential_info - - tenant_id = str(uuid.uuid4()) - provider_credential = _persist_provider_credential(orm_session, tenant_id=tenant_id) - _persist_provider(orm_session, tenant_id=tenant_id, credential_id=provider_credential.id) - - cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", None) - - assert cred_id == provider_credential.id - assert cred_name == "ProvCredName" - - def test_provider_not_found_returns_none_and_empty(self, orm_session: Session): - """Returns (None, '') when Provider record does not exist.""" - from core.ops.ops_trace_manager import _lookup_llm_credential_info - - other_tenant_id = str(uuid.uuid4()) - _persist_provider(orm_session, tenant_id=other_tenant_id, credential_id=None) - tenant_id = str(uuid.uuid4()) - - cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4") - - assert cred_id is None - assert cred_name == "" - - def test_none_tenant_id_returns_none_and_empty_without_db(self): - """Returns (None, '') immediately when tenant_id is None — no DB access.""" - from core.ops.ops_trace_manager import _lookup_llm_credential_info - - cred_id, cred_name = _lookup_llm_credential_info(None, "openai", "gpt-4") - - assert cred_id is None - assert cred_name == "" - - def test_none_provider_returns_none_and_empty_without_db(self): - """Returns (None, '') immediately when provider is None — no DB access.""" - from core.ops.ops_trace_manager import _lookup_llm_credential_info - - cred_id, cred_name = _lookup_llm_credential_info(str(uuid.uuid4()), None, "gpt-4") - - assert cred_id is None - assert cred_name == "" - - def test_db_error_on_outer_query_returns_none_and_empty(self, orm_session: Session, sqlite_engine: Engine): - """Returns (None, '') and logs a warning when the outer DB query raises.""" - from core.ops.ops_trace_manager import _lookup_llm_credential_info - - with _raise_on_table(sqlite_engine, "providers"): - cred_id, cred_name = _lookup_llm_credential_info(str(uuid.uuid4()), "openai", "gpt-4") - - assert cred_id is None - assert cred_name == "" - - def test_credential_name_lookup_failure_returns_id_with_empty_name( - self, orm_session: Session, sqlite_engine: Engine - ): - """When credential name sub-query fails, returns cred_id but '' for name.""" - from core.ops.ops_trace_manager import _lookup_llm_credential_info - - tenant_id = str(uuid.uuid4()) - provider_credential = _persist_provider_credential(orm_session, tenant_id=tenant_id) - _persist_provider(orm_session, tenant_id=tenant_id, credential_id=provider_credential.id) - - with _raise_on_table(sqlite_engine, "provider_credentials"): - cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4") - - assert cred_id == provider_credential.id - assert cred_name == "" - - def test_no_credential_on_provider_or_model_returns_none_id(self, orm_session: Session): - """Returns (None, '') when neither provider nor model has a credential_id.""" - from core.ops.ops_trace_manager import _lookup_llm_credential_info - - tenant_id = str(uuid.uuid4()) - _persist_provider(orm_session, tenant_id=tenant_id, credential_id=None) - _persist_provider_model(orm_session, tenant_id=tenant_id, credential_id=None) - - cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4") - - assert cred_id is None - assert cred_name == "" - - # --------------------------------------------------------------------------- - # TraceTask._get_user_id_from_metadata - # --------------------------------------------------------------------------- - - class TestGetUserIdFromMetadata: - """Tests for TraceTask._get_user_id_from_metadata(metadata). - - Pure dict logic — no DB access required. - """ - - @pytest.fixture - def get_user_id(self): - """Return the classmethod under test.""" - from core.ops.ops_trace_manager import TraceTask - - return TraceTask._get_user_id_from_metadata - - def test_from_end_user_id_has_highest_priority(self, get_user_id): - """from_end_user_id takes precedence over all other keys.""" - metadata = { - "from_end_user_id": "eu-abc", - "from_account_id": "acc-xyz", - "user_id": "u-123", - } - assert get_user_id(metadata) == "end_user:eu-abc" - - def test_from_account_id_used_when_no_end_user(self, get_user_id): - """from_account_id is used when from_end_user_id is absent.""" - metadata = { - "from_account_id": "acc-xyz", - "user_id": "u-123", - } - assert get_user_id(metadata) == "account:acc-xyz" - - def test_user_id_used_when_no_end_user_or_account(self, get_user_id): - """user_id is used when both higher-priority keys are absent.""" - metadata = {"user_id": "u-123"} - assert get_user_id(metadata) == "user:u-123" - - def test_returns_anonymous_when_all_keys_absent(self, get_user_id): - """Returns 'anonymous' when metadata has none of the expected keys.""" - assert get_user_id({}) == "anonymous" - - def test_empty_string_end_user_id_is_skipped(self, get_user_id): - """Empty string for from_end_user_id is falsy and falls through to next key.""" - metadata = { - "from_end_user_id": "", - "from_account_id": "acc-xyz", - } - assert get_user_id(metadata) == "account:acc-xyz" - - def test_empty_string_account_id_is_skipped(self, get_user_id): - """Empty string for from_account_id is falsy and falls through to user_id.""" - metadata = { - "from_end_user_id": "", - "from_account_id": "", - "user_id": "u-123", - } - assert get_user_id(metadata) == "user:u-123" - - def test_empty_string_user_id_falls_through_to_anonymous(self, get_user_id): - """Empty string for user_id is falsy, so 'anonymous' is returned.""" - metadata = { - "from_end_user_id": "", - "from_account_id": "", - "user_id": "", - } - assert get_user_id(metadata) == "anonymous" - - def test_only_from_end_user_id_present(self, get_user_id): - """Minimal case: only from_end_user_id present.""" - assert get_user_id({"from_end_user_id": "eu-only"}) == "end_user:eu-only" - - def test_irrelevant_keys_do_not_interfere(self, get_user_id): - """Extra metadata keys have no effect on the result.""" - metadata = {"invoke_from": "web", "app_id": "a1"} - assert get_user_id(metadata) == "anonymous" - - -"""SQLite-backed tests for :mod:`core.ops.ops_trace_manager`.""" - - -class DummyConfig: - def __init__(self, **kwargs): - self._data = kwargs - - def model_dump(self): - return dict(self._data) - - -class DummyTraceInstance: - def __init__(self, config): - self.config = config - - def api_check(self): - return True - - def get_project_key(self): - return "fake-key" - - def get_project_url(self): - 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 - - def __getitem__(self, key): - if key in self._data: - return self._data[key] - raise KeyError(key) - - -PROVIDER_ENTRY = { - "config_class": DummyConfig, - "secret_keys": ["secret_value"], - "other_keys": ["other_value"], - "trace_instance": DummyTraceInstance, -} - -UNIFIED_PROVIDER_ENTRY = { - "config_class": DummyConfig, - "trace_instance": DummyUnifiedTraceInstance, -} - - -class DummyTimer: - def __init__(self, interval, function): - self.interval = interval - self.function = function - self.name = "" - self.daemon = False - self.started = False - - def start(self): - self.started = True - - def is_alive(self): - return False - - -class EncryptTokenRecorder: - def __init__(self) -> None: - self.calls: list[tuple[str, str]] = [] - - def __call__(self, tenant_id: str, value: str) -> str: - self.calls.append((tenant_id, value)) - return f"enc-{value}" - - -class BatchDecryptTokenRecorder: - def __init__(self) -> None: - self.calls: list[tuple[str, list[str]]] = [] - - def __call__(self, tenant_id: str, values: list[str]) -> list[str]: - self.calls.append((tenant_id, values)) - return [f"dec-{value}" for value in values] - - -class ObfuscatedTokenRecorder: - def __init__(self) -> None: - self.calls: list[str] = [] - - def __call__(self, value: str) -> str: - self.calls.append(value) - return f"ob-{value}" - - -class RecordingStorage: - def __init__(self) -> None: - self.writes: list[tuple[str, bytes]] = [] - - def save(self, path: str, data: bytes) -> None: - self.writes.append((path, data)) - - -class RecordingDispatcher: - def __init__(self) -> None: - self.payloads: list[dict[str, str]] = [] - self.options: list[dict[str, object]] = [] - - def apply_async(self, *, args: list[dict[str, str]], **kwargs: object) -> None: - self.payloads.append(args[0]) - self.options.append(kwargs) - - -def _app(session: Session, *, app_id: str = "app-id", tracing: str | None = None) -> App: - app = make_app(app_id=app_id, name="App", description="description", icon_type=None, tracing=tracing) - session.add(app) - session.commit() - return app - - -def _conversation_message( - session: Session, app: App, *, config: AppModelConfig | None = None -) -> tuple[Conversation, Message]: - conversation = Conversation( - id="conversation-1", - app_id=app.id, - app_model_config_id=config.id if config else None, - model_provider="provider", - override_model_configs=None, - model_id="model", - mode=AppMode.CHAT, - name="Conversation", - summary="", - _inputs={}, - introduction="", - system_instruction="", - invoke_from=None, - from_source=ConversationFromSource.CONSOLE, - from_end_user_id="end-user-1", - from_account_id=None, - read_at=None, - read_account_id=None, - ) - message = Message( - id="message-1", - app_id=app.id, - model_provider="provider", - model_id="model", - override_model_configs=None, - conversation_id=conversation.id, - _inputs={}, - query="query", - message={"text": "hello"}, - message_tokens=5, - message_unit_price=Decimal(0), - message_price_unit=Decimal("0.001"), - answer="world", - answer_tokens=7, - answer_unit_price=Decimal(0), - answer_price_unit=Decimal("0.001"), - parent_message_id=None, - provider_response_latency=1, - total_price=Decimal(0), - currency="USD", - status=MessageStatus.NORMAL, - error=None, - message_metadata=None, - invoke_from=None, - from_source=ConversationFromSource.CONSOLE, - from_end_user_id="end-user-1", - from_account_id=None, - agent_based=False, - workflow_run_id="run-1", - app_mode=AppMode.CHAT, - ) - session.add_all([conversation, message]) - session.commit() - return conversation, message - - -def _message_data(**overrides): - created_at = datetime(2025, 2, 20, 12, 0, 0) - data = { - "id": "message-1", - "app_id": "app-id", - "conversation_id": "conversation-1", - "created_at": created_at, - "updated_at": created_at + timedelta(seconds=3), - "message": "hello", - "provider_response_latency": 1, - "message_tokens": 5, - "answer_tokens": 7, - "answer": "world", - "error": "", - "status": "complete", - "model_provider": "provider", - "model_id": "model", - "from_end_user_id": "end-user-1", - "from_account_id": None, - "agent_based": False, - "workflow_run_id": "run-1", - "from_source": "console", - "message_metadata": json.dumps({"usage": {"time_to_first_token": 1, "time_to_generate": 2}}), - "agent_thoughts": [], - "query": "query", - "inputs": "inputs", - } - data.update(overrides) - return SimpleNamespace( - **data, - agent_thoughts_with_session=lambda *, session: data["agent_thoughts"], - to_dict=lambda: data, - ) - - -class TestOpsTraceManagerCases: - @staticmethod - @pytest.fixture - def database(sqlite_engine: Engine, sqlite_session: Session) -> Iterator[Session]: - session_proxy = scoped_session(lambda: sqlite_session) - with ( - patch.object(module.db, "session", session_proxy), - patch.object(type(module.db), "engine", new_callable=PropertyMock, return_value=sqlite_engine), - ): - yield sqlite_session - - @staticmethod - @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({})) - apply_config_overrides(monkeypatch, OPS_TRACE_UNIFIED_ENABLED=False) - OpsTraceManager.ops_trace_instances_cache.clear() - OpsTraceManager.decrypted_configs_cache.clear() - monkeypatch.setattr(module.threading, "Timer", DummyTimer) - monkeypatch.setattr(module, "trace_manager_queue", queue.Queue()) - monkeypatch.setattr(module, "trace_manager_timer", None) - monkeypatch.setattr("core.telemetry.gateway.is_enterprise_telemetry_enabled", lambda: False) - - app = Flask(__name__) - with app.app_context(): - yield - - @staticmethod - @pytest.fixture - def encryption_functions( - monkeypatch: pytest.MonkeyPatch, - ) -> tuple[EncryptTokenRecorder, BatchDecryptTokenRecorder, ObfuscatedTokenRecorder]: - encrypt = EncryptTokenRecorder() - decrypt = BatchDecryptTokenRecorder() - obfuscate = ObfuscatedTokenRecorder() - monkeypatch.setattr(module, "encrypt_token", encrypt) - monkeypatch.setattr(module, "batch_decrypt_token", decrypt) - monkeypatch.setattr(module, "obfuscated_token", obfuscate) - return encrypt, decrypt, obfuscate - - @staticmethod - def test_encrypt_decrypt_obfuscate_and_cache( - trace_environment: None, - encryption_functions: tuple[EncryptTokenRecorder, BatchDecryptTokenRecorder, ObfuscatedTokenRecorder], - ) -> None: - encrypted = OpsTraceManager.encrypt_tracing_config( - "tenant-1", "dummy", {"secret_value": "value", "other_value": "info"} - ) - assert encrypted == {"secret_value": "enc-value", "other_value": "info"} - preserved = OpsTraceManager.encrypt_tracing_config( - "tenant-1", "dummy", {"secret_value": "*"}, current_trace_config={"secret_value": "keep"} - ) - assert preserved["secret_value"] == "keep" - first = OpsTraceManager.decrypt_tracing_config("tenant-1", "dummy", encrypted) - second = OpsTraceManager.decrypt_tracing_config("tenant-1", "dummy", encrypted) - assert first == second - assert len(encryption_functions[1].calls) == 1 - obfuscated = OpsTraceManager.obfuscated_decrypt_token("dummy", first) - assert obfuscated["secret_value"] == "ob-dec-enc-value" - assert encryption_functions[2].calls == ["dec-enc-value"] - - @staticmethod - def test_decrypted_config_reads_real_trace_and_app_rows( - trace_environment: None, - encryption_functions, - database: Session, - ) -> None: - app = _app(database) - trace = TraceAppConfig( - app_id=app.id, - tracing_provider="dummy", - tracing_config={"secret_value": "encrypted", "other_value": "info"}, - ) - database.add(trace) - database.commit() - result = OpsTraceManager.get_decrypted_tracing_config(app.id, "dummy") - assert result == {"secret_value": "dec-encrypted", "other_value": "info"} - assert OpsTraceManager.get_decrypted_tracing_config(app.id, "missing") is None - - null_config_app = _app(database, app_id="app-null-config") - database.add(TraceAppConfig(app_id=null_config_app.id, tracing_provider="dummy", tracing_config=None)) - database.commit() - with pytest.raises(ValueError, match="Tracing config cannot be None"): - OpsTraceManager.get_decrypted_tracing_config(null_config_app.id, "dummy") - - database.delete(app) - database.commit() - with pytest.raises(ValueError, match="App not found"): - OpsTraceManager.get_decrypted_tracing_config("app-id", "dummy") - - @staticmethod - def test_ops_trace_instance_uses_persisted_enabled_state_and_cache( - trace_environment: None, - encryption_functions, - database: Session, - ) -> None: - app = _app(database, tracing=json.dumps({"enabled": False, "tracing_provider": "dummy"})) - assert OpsTraceManager.get_ops_trace_instance(app.id) is None - app.tracing = json.dumps({"enabled": True, "tracing_provider": "dummy"}) - database.add( - TraceAppConfig(app_id=app.id, tracing_provider="dummy", tracing_config={"secret_value": "encrypted"}) - ) - database.commit() - instance = OpsTraceManager.get_ops_trace_instance(app.id) - assert isinstance(instance, DummyTraceInstance) - assert OpsTraceManager.get_ops_trace_instance(app.id) is instance - - app.tracing = json.dumps({"enabled": True, "tracing_provider": "missing"}) - database.commit() - assert OpsTraceManager.get_ops_trace_instance(app.id) is None - - assert OpsTraceManager.get_ops_trace_instance(None) is None - assert OpsTraceManager.get_ops_trace_instance("tenant-storage-id") is None - assert OpsTraceManager.get_ops_trace_instance("missing") is None - - @staticmethod - @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() - apply_config_overrides(monkeypatch, 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 - - @staticmethod - 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() - apply_config_overrides(monkeypatch, 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) - - @staticmethod - 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})) - - apply_config_overrides(monkeypatch, OPS_TRACE_UNIFIED_ENABLED=False) - legacy = OpsTraceManager.get_ops_trace_instance(app.id) - apply_config_overrides(monkeypatch, 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 - - @staticmethod - 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"}') - database.add(config) - database.commit() - conversation, message = _conversation_message(database, app, config=config) - result = OpsTraceManager.get_app_config_through_message_id(message.id) - assert result.id == config.id - - conversation.app_model_config_id = None - conversation.override_model_configs = json.dumps({"provider": "override"}) - database.commit() - override = OpsTraceManager.get_app_config_through_message_id(message.id) - assert json.loads(override) == {"provider": "override"} - - assert OpsTraceManager.get_app_config_through_message_id("missing") is None - - @staticmethod - @pytest.mark.parametrize("imports_cached", [False, True], ids=["cold", "warm"]) - @pytest.mark.parametrize("module_name", ["dify_trace_weave", "weave"]) - def test_provider_loader_identifies_missing_optional_dependency(module_name: str, imports_cached: bool) -> None: - with patch.dict(sys.modules): - if imports_cached: - module.OpsTraceProviderConfigMap()["weave"] - # Cached typed SDK submodules can bypass a missing top-level package. - # Simulate its absence consistently and restore the complete cache after. - for cached_name in list(sys.modules): - if cached_name.partition(".")[0] in {"dify_trace_weave", "weave"}: - del sys.modules[cached_name] - with patch.dict(sys.modules, {module_name: None}): - with pytest.raises(TraceProviderNotInstalledError, match=f"weave.*{module_name}") as caught: - module.OpsTraceProviderConfigMap()["weave"] - - cause = caught.value.__cause__ - assert isinstance(cause, ModuleNotFoundError) - assert cause.name is not None - assert cause.name.partition(".")[0] == module_name - - @staticmethod - @pytest.mark.parametrize( - "failure", - [ - ImportError("cannot import name 'WeaveDataTrace'"), - ModuleNotFoundError("SDK import failed without identifying a missing module"), - ModuleNotFoundError("No module named 'json.missing_module'", name="json.missing_module"), - ], - ) - def test_provider_loader_preserves_import_errors_in_installed_packages(failure: ImportError) -> None: - original_import = builtins.__import__ - - def broken_import(name, *args, **kwargs): - if name.startswith("dify_trace_weave."): - raise failure - return original_import(name, *args, **kwargs) - - with patch("builtins.__import__", side_effect=broken_import), pytest.raises(ImportError) as caught: - module.OpsTraceProviderConfigMap()["weave"] - - assert caught.value is failure - - @staticmethod - def test_message_trace_reads_real_conversation_app_and_message_file( - trace_environment: None, - database: Session, - ) -> None: - app = _app(database) - _, message = _conversation_message(database, app) - file = MessageFile( - message_id=message.id, - type=FileType.DOCUMENT, - transfer_method=FileTransferMethod.REMOTE_URL, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="user-1", - url="path/to/file", - ) - database.add(file) - database.commit() - 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") - assert result.metadata["tenant_id"] == "tenant-1" - - @staticmethod - def test_workflow_log_enriches_moderation_and_suggested_question_traces( - database: Session, - ) -> None: - app = _app(database) - _, message = _conversation_message(database, app) - log = WorkflowAppLog( - tenant_id="tenant-1", - app_id="app-id", - workflow_id="workflow-1", - workflow_run_id="run-1", - created_from=WorkflowAppLogCreatedFrom.WEB_APP, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="user-1", - ) - database.add(log) - database.commit() - task = TraceTask(trace_type=TraceTaskName.MODERATION_TRACE, message_id=message.id) - moderation = ModerationInputsResult( - action=ModerationAction.DIRECT_OUTPUT, - preset_response="no", - query="q", - flagged=True, - ) - result = task.moderation_trace( - message.id, - {"start": 1, "end": 2}, - moderation_result=moderation, - inputs={"source": "payload"}, - ) - assert result.message_id == log.id - assert result.flagged is True - assert result.inputs == {"source": "payload"} - suggested = task.suggested_question_trace(message.id, {"start": 1, "end": 2}, suggested_question=["q1"]) - assert suggested.message_id == log.id - assert suggested.suggested_question == ["q1"] - - @staticmethod - def test_dataset_retrieval_trace_serializes_documents( - trace_environment: None, - database: Session, - ) -> None: - app = _app(database) - _, message = _conversation_message(database, app) - document = RetrievalDocument(page_content="value") - - result = TraceTask(trace_type=TraceTaskName.DATASET_RETRIEVAL_TRACE).dataset_retrieval_trace( - message.id, - {"start": 1, "end": 2}, - documents=[document], - ) - - assert result.documents == [document.model_dump()] - assert result.documents[0]["page_content"] == "value" - assert result.metadata["tenant_id"] == "tenant-1" - - @staticmethod - def test_workflow_trace_reads_real_workflow_log_from_owned_session( - monkeypatch: pytest.MonkeyPatch, - trace_environment: None, - database: Session, - sqlite_session_factory: sessionmaker[Session], - ) -> None: - app = _app(database) - log = WorkflowAppLog( - tenant_id=app.tenant_id, - app_id=app.id, - workflow_id="workflow-1", - workflow_run_id="run-1", - created_from=WorkflowAppLogCreatedFrom.WEB_APP, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="user-1", - ) - workflow_run = WorkflowRun( - id="run-1", - tenant_id=app.tenant_id, - app_id=app.id, - workflow_id="workflow-1", - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="3", - graph="{}", - inputs=json.dumps({"query": "search"}), - status=WorkflowExecutionStatus.SUCCEEDED, - outputs=json.dumps({"out": "value"}), - error=None, - elapsed_time=10, - total_tokens=12, - total_steps=1, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="user-1", - finished_at=datetime(2025, 2, 20, 10, 0, 5), - ) - workflow_run.created_at = datetime(2025, 2, 20, 10, 0, 0) - database.add_all([log, workflow_run]) - database.commit() - repo = DifyAPISQLAlchemyWorkflowRunRepository(sqlite_session_factory) - monkeypatch.setattr(TraceTask, "_get_workflow_run_repo", classmethod(lambda cls: repo)) - monkeypatch.setattr(TraceTask, "_calculate_workflow_token_split", classmethod(lambda cls, *_a, **_k: (5, 7))) - result = TraceTask(trace_type=TraceTaskName.WORKFLOW_TRACE).workflow_trace( - workflow_run_id="run-1", conversation_id=None, user_id="user-1" - ) - assert result.workflow_run_id == "run-1" - assert result.workflow_id == "workflow-1" - assert result.workflow_app_log_id == log.id - assert result.prompt_tokens == 5 - assert result.completion_tokens == 7 - - @staticmethod - def test_tool_trace_reads_real_message_file(database: Session) -> None: - app = _app(database) - _, message = _conversation_message(database, app) - file = MessageFile( - message_id=message.id, - type=FileType.DOCUMENT, - transfer_method=FileTransferMethod.REMOTE_URL, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="user-1", - url="tool/file", - ) - thought = MessageAgentThought( - message_id=message.id, - position=1, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="user-1", - tool="tool-a", - tool_meta_str=json.dumps( - {"tool-a": {"tool_config": {}, "time_cost": 5, "error": "", "tool_parameters": {}}} - ), - ) - database.add_all([file, thought]) - database.commit() - result = TraceTask(trace_type=TraceTaskName.TOOL_TRACE).tool_trace( - message.id, {"start": 1, "end": 2}, tool_name="tool-a", tool_inputs={}, tool_outputs="result" - ) - assert result.tool_name == "tool-a" - assert result.time_cost == 5 - assert result.message_file_data.id == file.id - - @staticmethod - def test_node_execution_trace_resolves_real_message_by_conversation_and_run( - trace_environment: None, database: Session - ) -> None: - app = _app(database) - conversation, message = _conversation_message(database, app) - result = TraceTask(trace_type=TraceTaskName.NODE_EXECUTION_TRACE).node_execution_trace( - node_execution_data={ - "tenant_id": app.tenant_id, - "app_id": app.id, - "conversation_id": conversation.id, - "workflow_execution_id": message.workflow_run_id, - "workflow_id": "workflow-1", - "node_execution_id": "node-execution-1", - "node_id": "node-1", - "node_type": "llm", - "title": "Node", - "status": "succeeded", - } - ) - assert result.message_id == message.id - assert result.metadata["conversation_id"] == conversation.id - - @staticmethod - 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, message_id="message-1") - assert task.conversation_trace(foo="bar") == {"foo": "bar"} - assert task._extract_streaming_metrics(Message(message_metadata="invalid")) == {} - assert task.generate_name_trace("conversation", {"start": 1, "end": 2}, tenant_id=None) == {} - generated = task.generate_name_trace( - "conversation", - {"start": 1, "end": 2}, - tenant_id="tenant-1", - generate_conversation_name="name", - inputs="query", - ) - assert generated.outputs == "name" - assert generated.tenant_id == "tenant-1" - assert generated.message_id == "message-1" - - @staticmethod - def test_trace_queue_preserves_source_app_when_another_manager_drains( - monkeypatch: pytest.MonkeyPatch, trace_environment: None - ) -> None: - monkeypatch.setattr( - OpsTraceManager, - "get_ops_trace_instance", - classmethod(lambda cls, app_id: app_id == "source-app-id"), - ) - draining_manager = TraceQueueManager(app_id="draining-app-id", user_id="user-1") - source_manager = TraceQueueManager(app_id="source-app-id", user_id="user-2") - 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", - ) - source_manager.add_trace_task(task) - assert draining_manager.collect_tasks() == [task] - - recording_storage = RecordingStorage() - dispatcher = RecordingDispatcher() - monkeypatch.setattr(module.storage, "save", recording_storage.save) - 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) - source_manager.add_trace_task(task) - draining_manager.run() - - assert len(recording_storage.writes) == 1 - path, data = recording_storage.writes[0] - assert path.endswith(f"source-app-id/{file_id.hex}.json") - assert json.loads(data)["app_id"] == "source-app-id" - assert dispatcher.payloads == [{"file_id": file_id.hex, "app_id": "source-app-id"}] - - @staticmethod - 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", - ) - task.app_id = "app-id" - 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"]) - - @staticmethod - 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) - task.app_id = "app-id" - - 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") - - @staticmethod - 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"}) - - -"""Unit tests for TraceQueueManager telemetry guard. - -Verifies that TraceQueueManager.add_trace_task() only enqueues tasks when at -least one consumer is active: -- Enterprise telemetry is enabled (_enterprise_telemetry_enabled=True), OR -- A third-party trace instance (Langfuse, etc.) is configured - -When neither is active, tasks are silently dropped to avoid unnecessary work. - -When BOTH are false, tasks are silently dropped (correct behavior). -""" - - -class TestTraceQueueManagerCases: - @staticmethod - @pytest.fixture(autouse=True) - def _application_tracing_disabled(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) - - @staticmethod - @pytest.fixture - def trace_queue_manager_and_task(): - """Exercise the production guard without starting background dispatch.""" - from core.ops.entities.trace_entity import TraceTaskName - from core.ops.ops_trace_manager import TraceQueueManager, TraceTask - - with patch.object(TraceQueueManager, "start_timer"): - yield TraceQueueManager, TraceTask, TraceTaskName - - class TestTraceQueueManagerTelemetryGuard: - """Test TraceQueueManager's telemetry guard in add_trace_task().""" - - def test_task_not_enqueued_when_telemetry_disabled_and_no_trace_instance(self, trace_queue_manager_and_task): - """Verify task is NOT enqueued when telemetry disabled and no trace instance. - - This is the core guard: when _enterprise_telemetry_enabled=False AND - trace_instance=None, the task should be silently dropped. - """ - trace_queue_manager, trace_task_type, trace_task_name = trace_queue_manager_and_task - - trace_queue = queue.Queue() - - trace_task = trace_task_type(trace_type=trace_task_name.WORKFLOW_TRACE) - - with ( - patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False), - patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None), - patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), - ): - manager = trace_queue_manager(app_id="test-app-id") - manager.add_trace_task(trace_task) - - assert trace_queue.empty() - - def test_task_enqueued_when_telemetry_enabled(self, trace_queue_manager_and_task): - """Verify task IS enqueued when enterprise telemetry is enabled. - - When _enterprise_telemetry_enabled=True, the task should be enqueued - regardless of trace_instance state. - """ - trace_queue_manager, trace_task_type, trace_task_name = trace_queue_manager_and_task - - trace_queue = queue.Queue() - - trace_task = trace_task_type(trace_type=trace_task_name.WORKFLOW_TRACE) - - with ( - patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True), - patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None), - patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), - ): - manager = trace_queue_manager(app_id="test-app-id") - manager.add_trace_task(trace_task) - - called_task = trace_queue.get_nowait() - assert called_task is trace_task - assert trace_queue.empty() - assert called_task.app_id == "test-app-id" - - def test_task_enqueued_when_trace_instance_configured(self, trace_queue_manager_and_task): - """Verify task IS enqueued when third-party trace instance is configured. - - When trace_instance is not None (e.g., Langfuse configured), the task - should be enqueued even if enterprise telemetry is disabled. - """ - trace_queue_manager, trace_task_type, trace_task_name = trace_queue_manager_and_task - - trace_queue = queue.Queue() - - mock_trace_instance = MagicMock() - - trace_task = trace_task_type(trace_type=trace_task_name.WORKFLOW_TRACE) - - with ( - patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False), - patch( - "core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", - return_value=mock_trace_instance, - ), - patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), - ): - manager = trace_queue_manager(app_id="test-app-id") - manager.add_trace_task(trace_task) - - called_task = trace_queue.get_nowait() - assert called_task is trace_task - assert trace_queue.empty() - assert called_task.app_id == "test-app-id" - - def test_task_enqueued_when_both_telemetry_and_trace_instance_enabled(self, trace_queue_manager_and_task): - """Verify task IS enqueued when both telemetry and trace instance are enabled. - - When both _enterprise_telemetry_enabled=True AND trace_instance is set, - the task should definitely be enqueued. - """ - trace_queue_manager, trace_task_type, trace_task_name = trace_queue_manager_and_task - - trace_queue = queue.Queue() - - mock_trace_instance = MagicMock() - - trace_task = trace_task_type(trace_type=trace_task_name.WORKFLOW_TRACE) - - with ( - patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True), - patch( - "core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", - return_value=mock_trace_instance, - ), - patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), - ): - manager = trace_queue_manager(app_id="test-app-id") - manager.add_trace_task(trace_task) - - called_task = trace_queue.get_nowait() - assert called_task is trace_task - assert trace_queue.empty() - assert called_task.app_id == "test-app-id" - - def test_app_id_set_before_enqueue(self, trace_queue_manager_and_task): - """Verify app_id is set on the task before enqueuing. - - The guard logic sets trace_task.app_id = self.app_id before calling - trace_manager_queue.put(trace_task). This test verifies that behavior. - """ - trace_queue_manager, trace_task_type, trace_task_name = trace_queue_manager_and_task - - trace_queue = queue.Queue() - - trace_task = trace_task_type(trace_type=trace_task_name.WORKFLOW_TRACE) - - with ( - patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True), - patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None), - patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), - ): - manager = trace_queue_manager(app_id="expected-app-id") - manager.add_trace_task(trace_task) - - called_task = trace_queue.get_nowait() - assert called_task is trace_task - assert trace_queue.empty() - assert called_task.app_id == "expected-app-id" - - @staticmethod - @pytest.mark.parametrize("enabled_at_start", [False, True]) - def test_application_pipeline_snapshots_switch_without_constructing_sdk( - trace_queue_manager_and_task, monkeypatch: pytest.MonkeyPatch, enabled_at_start: bool - ): - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - trace_queue_manager, trace_task_class, trace_task_name = trace_queue_manager_and_task - trace_queue = queue.Queue() - calls: list[object] = [] - current_enabled = enabled_at_start - - def queue_enabled(app_id: object) -> bool: - calls.append(app_id) - return current_enabled - - monkeypatch.setattr(RuntimeApplicationTracingHooks, "queue_enabled_at_start", staticmethod(queue_enabled)) - hooks = RuntimeApplicationTracingHooks() - app = Flask(__name__) - bind_application_tracing_hooks(app, hooks) - with ( - app.app_context(), - patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False), - patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance") as legacy, - patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), - ): - manager = trace_queue_manager(app_id="app-id") - current_enabled = not enabled_at_start - task = trace_task_class(trace_type=trace_task_name.WORKFLOW_TRACE) - manager.add_trace_task(task) - legacy.assert_not_called() - assert calls == ["app-id"] - assert manager.trace_instance is None - assert trace_queue.qsize() == int(enabled_at_start) - - -"""Nullable provider credentials never reach RSA/base64 as non-string tokens.""" - - -@pytest.fixture -def workspace(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> str: - key = RSA.generate(2048) - tenant_id = str(uuid4()) - sqlite_session.add(make_tenant(tenant_id=tenant_id, encrypt_public_key=key.publickey().export_key().decode())) - sqlite_session.commit() - monkeypatch.setattr(model_engine.db, "session", sqlite_session) - monkeypatch.setattr(key_provider_manager, "_provider", RSAKeyProvider()) - monkeypatch.setattr(rsa, "get_decrypt_decoding", lambda _tenant: (key, gmpy2_pkcs10aep_cipher.new(key))) - return tenant_id - - -class TestNullableTracingCredentialsCases: - @staticmethod - def test_databricks_switches_auth_with_real_encryption_and_explicit_null(workspace: str) -> None: - pat = OpsTraceManager.encrypt_tracing_config( - workspace, - "databricks", - { - "host": "https://workspace.test", - "experiment_id": "12", - "personal_access_token": "synthetic-pat", - "client_id": None, - "client_secret": None, - }, - ) - assert pat["personal_access_token"] != "synthetic-pat" - decrypted_pat = OpsTraceManager.decrypt_tracing_config(workspace, "databricks", pat) - assert decrypted_pat["personal_access_token"] == "synthetic-pat" - assert decrypted_pat["client_secret"] is None - - oauth = OpsTraceManager.encrypt_tracing_config( - workspace, - "databricks", - { - "host": "https://workspace.test", - "experiment_id": "12", - "personal_access_token": None, - "client_id": "synthetic-client", - "client_secret": "synthetic-secret", - }, - pat, - ) - assert oauth["personal_access_token"] is None - assert oauth["client_secret"] != "synthetic-secret" - decrypted_oauth = OpsTraceManager.decrypt_tracing_config(workspace, "databricks", oauth) - assert decrypted_oauth["personal_access_token"] is None - assert decrypted_oauth["client_secret"] == "synthetic-secret" - - @staticmethod - def test_other_optional_secret_model_can_be_cleared_without_losing_validation(workspace: str) -> None: - password = OpsTraceManager.encrypt_tracing_config( - workspace, - "mlflow", - { - "tracking_uri": "https://mlflow.test", - "experiment_id": "0", - "username": "user", - "password": "synthetic-password", - }, - ) - assert OpsTraceManager.decrypt_tracing_config(workspace, "mlflow", password)["password"] == "synthetic-password" - cleared = OpsTraceManager.encrypt_tracing_config(workspace, "mlflow", {**password, "password": None}, password) - assert cleared["password"] is None - assert OpsTraceManager.decrypt_tracing_config(workspace, "mlflow", cleared)["password"] is None - - @staticmethod - def test_non_nullable_secret_is_still_rejected_by_provider_model(workspace: str) -> None: - with pytest.raises(ValidationError, match="secret_key"): - OpsTraceManager.encrypt_tracing_config( - workspace, "langfuse", {"public_key": "synthetic-public", "secret_key": None} - ) - encrypted = OpsTraceManager.encrypt_tracing_config( - workspace, "langfuse", {"public_key": "synthetic-public", "secret_key": "synthetic-secret"} - ) - with pytest.raises(ValidationError, match="secret_key"): - OpsTraceManager.decrypt_tracing_config(workspace, "langfuse", {**encrypted, "secret_key": None}) diff --git a/api/tests/unit_tests/core/ops/test_utils.py b/api/tests/unit_tests/core/ops/test_utils.py deleted file mode 100644 index b1dbd9c0974dbe..00000000000000 --- a/api/tests/unit_tests/core/ops/test_utils.py +++ /dev/null @@ -1,337 +0,0 @@ -import re -from datetime import datetime -from decimal import Decimal - -import pytest -from sqlalchemy.orm import Session - -from core.ops.utils import ( - filter_none_values, - generate_dotted_order, - get_message_data, - measure_time, - replace_text_with_content, - validate_integer_id, - validate_project_name, - validate_url, - validate_url_with_path, -) -from models.enums import ConversationFromSource -from models.model import Message -from tests.helpers.unit.model_factories import make_message - - -@pytest.fixture -def message_session(sqlite_session: Session) -> Session: - """Provide the shared SQLite test session for message lookups.""" - return sqlite_session - - -def _message(message_id: str) -> Message: - return make_message( - message_id=message_id, - app_id="app-id", - conversation_id="conversation-id", - inputs={}, - query="question", - message={"role": "user", "content": "question"}, - answer="answer", - message_unit_price=Decimal("0.0001"), - answer_unit_price=Decimal("0.0001"), - currency="USD", - from_source=ConversationFromSource.API, - ) - - -class TestValidateUrl: - """Test cases for validate_url function""" - - def test_valid_https_url(self): - """Test valid HTTPS URL""" - result = validate_url("https://example.com", "https://default.com") - assert result == "https://example.com" - - def test_valid_http_url(self): - """Test valid HTTP URL""" - result = validate_url("http://example.com", "https://default.com") - assert result == "http://example.com" - - def test_url_with_path_removed(self): - """Test that URL path is removed during normalization""" - result = validate_url("https://example.com/api/v1/test", "https://default.com") - assert result == "https://example.com" - - def test_url_with_query_removed(self): - """Test that URL query parameters are removed""" - result = validate_url("https://example.com?param=value", "https://default.com") - assert result == "https://example.com" - - def test_url_with_fragment_removed(self): - """Test that URL fragments are removed""" - result = validate_url("https://example.com#section", "https://default.com") - assert result == "https://example.com" - - def test_empty_url_returns_default(self): - """Test empty URL returns default""" - result = validate_url("", "https://default.com") - assert result == "https://default.com" - - def test_none_url_returns_default(self): - """Test None URL returns default""" - result = validate_url(None, "https://default.com") - assert result == "https://default.com" - - def test_whitespace_url_returns_default(self): - """Test whitespace URL returns default""" - result = validate_url(" ", "https://default.com") - assert result == "https://default.com" - - def test_invalid_scheme_raises_error(self): - """Test invalid scheme raises ValueError""" - with pytest.raises(ValueError, match="URL scheme must be one of"): - validate_url("ftp://example.com", "https://default.com") - - def test_no_scheme_raises_error(self): - """Test URL without scheme raises ValueError""" - with pytest.raises(ValueError, match="URL scheme must be one of"): - validate_url("example.com", "https://default.com") - - def test_custom_allowed_schemes(self): - """Test custom allowed schemes""" - result = validate_url("https://example.com", "https://default.com", allowed_schemes=("https",)) - assert result == "https://example.com" - - with pytest.raises(ValueError, match="URL scheme must be one of"): - validate_url("http://example.com", "https://default.com", allowed_schemes=("https",)) - - -class TestValidateUrlWithPath: - """Test cases for validate_url_with_path function""" - - def test_valid_url_with_path(self): - """Test valid URL with path""" - result = validate_url_with_path("https://example.com/api/v1", "https://default.com") - assert result == "https://example.com/api/v1" - - def test_valid_url_with_required_suffix(self): - """Test valid URL with required suffix""" - result = validate_url_with_path("https://example.com/api/", "https://default.com", required_suffix="/api/") - assert result == "https://example.com/api/" - - def test_url_without_required_suffix_raises_error(self): - """Test URL without required suffix raises error""" - with pytest.raises(ValueError, match="URL should end with /api/"): - validate_url_with_path("https://example.com/api", "https://default.com", required_suffix="/api/") - - def test_empty_url_returns_default(self): - """Test empty URL returns default""" - result = validate_url_with_path("", "https://default.com") - assert result == "https://default.com" - - def test_none_url_returns_default(self): - """Test None URL returns default""" - result = validate_url_with_path(None, "https://default.com") - assert result == "https://default.com" - - def test_invalid_scheme_raises_error(self): - """Test invalid scheme raises ValueError""" - with pytest.raises(ValueError, match="URL must start with https:// or http://"): - validate_url_with_path("ftp://example.com", "https://default.com") - - def test_no_scheme_raises_error(self): - """Test URL without scheme raises ValueError""" - with pytest.raises(ValueError, match="URL must start with https:// or http://"): - validate_url_with_path("example.com", "https://default.com") - - def test_restricted_scheme_accepts_allowed_scheme(self): - """Test https-only validation keeps the path of an https URL""" - result = validate_url_with_path( - "https://langsmith.internal/api", "https://default.com", allowed_schemes=("https",) - ) - assert result == "https://langsmith.internal/api" - - def test_restricted_scheme_rejects_http(self): - """Test https-only validation rejects http and names only https in the error""" - with pytest.raises(ValueError, match="URL must start with https://") as excinfo: - validate_url_with_path("http://langsmith.internal/api", "https://default.com", allowed_schemes=("https",)) - assert str(excinfo.value) == "URL must start with https://" - - def test_default_schemes_keep_original_error_message(self): - """Test the two-scheme default keeps the exact message existing providers assert on""" - with pytest.raises(ValueError, match="URL must start with https:// or http://") as excinfo: - validate_url_with_path("ftp://example.com", "https://default.com") - assert str(excinfo.value) == "URL must start with https:// or http://" - - def test_surrounding_whitespace_is_stripped(self): - """Test surrounding whitespace is removed while the path is preserved""" - result = validate_url_with_path(" https://example.com/api/v1 ", "https://default.com") - assert result == "https://example.com/api/v1" - - -class TestValidateProjectName: - """Test cases for validate_project_name function""" - - def test_valid_project_name(self): - """Test valid project name""" - result = validate_project_name("my-project", "default") - assert result == "my-project" - - def test_empty_project_name_returns_default(self): - """Test empty project name returns default""" - result = validate_project_name("", "default") - assert result == "default" - - def test_none_project_name_returns_default(self): - """Test None project name returns default""" - result = validate_project_name(None, "default") - assert result == "default" - - def test_whitespace_project_name_returns_default(self): - """Test whitespace project name returns default""" - result = validate_project_name(" ", "default") - assert result == "default" - - def test_project_name_with_whitespace_trimmed(self): - """Test project name with whitespace is trimmed""" - result = validate_project_name(" my-project ", "default") - assert result == "my-project" - - def test_custom_default_name(self): - """Test custom default name""" - result = validate_project_name("", "Custom Default") - assert result == "Custom Default" - - -class TestGenerateDottedOrder: - """Test cases for generate_dotted_order function""" - - def test_dotted_order_has_6_digit_microseconds(self): - """Test that timestamp includes full 6-digit microseconds for LangSmith API compatibility. - - LangSmith API expects timestamps in format: YYYYMMDDTHHMMSSffffffZ (6-digit microseconds). - Previously, the code truncated to 3 digits which caused API errors: - 'cannot parse .111 as .000000' - """ - start_time = datetime(2025, 12, 23, 4, 19, 55, 111000) - run_id = "test-run-id" - result = generate_dotted_order(run_id, start_time) - - # Extract timestamp portion (before the run_id) - timestamp_match = re.match(r"^(\d{8}T\d{6})(\d+)Z", result) - assert timestamp_match is not None, "Timestamp format should match YYYYMMDDTHHMMSSffffffZ" - - microseconds = timestamp_match.group(2) - assert len(microseconds) == 6, f"Microseconds should be 6 digits, got {len(microseconds)}: {microseconds}" - - def test_dotted_order_format_matches_langsmith_expected(self): - """Test that dotted_order format matches LangSmith API expected format.""" - start_time = datetime(2025, 1, 15, 10, 30, 45, 123456) - run_id = "abc123" - result = generate_dotted_order(run_id, start_time) - - # LangSmith expects: YYYYMMDDTHHMMSSffffffZ followed by run_id - assert result == "20250115T103045123456Zabc123" - - def test_dotted_order_with_parent(self): - """Test dotted_order generation with parent order uses dot separator.""" - start_time = datetime(2025, 12, 23, 4, 19, 55, 111000) - run_id = "child-run-id" - parent_order = "20251223T041955000000Zparent-run-id" - result = generate_dotted_order(run_id, start_time, parent_order) - - assert result == "20251223T041955000000Zparent-run-id.20251223T041955111000Zchild-run-id" - - def test_dotted_order_without_parent_has_no_dot(self): - """Test dotted_order generation without parent has no dot separator.""" - start_time = datetime(2025, 12, 23, 4, 19, 55, 111000) - run_id = "test-run-id" - result = generate_dotted_order(run_id, start_time, None) - - assert "." not in result - - def test_dotted_order_with_string_start_time(self): - """Test dotted_order generation with string start_time.""" - start_time = "2025-12-23T04:19:55.111000" - run_id = "test-run-id" - result = generate_dotted_order(run_id, start_time) - - assert result == "20251223T041955111000Ztest-run-id" - - -class TestFilterNoneValues: - """Test cases for filter_none_values function""" - - def test_filter_none_values(self): - data = {"a": 1, "b": None, "c": "test", "d": datetime(2025, 1, 1, 12, 0, 0)} - result = filter_none_values(data) - assert result == {"a": 1, "c": "test", "d": "2025-01-01T12:00:00"} - - def test_filter_none_values_empty(self): - assert filter_none_values({}) == {} - - -@pytest.mark.parametrize("sqlite_session", [(Message,)], indirect=True) -class TestGetMessageData: - """Test cases for get_message_data function""" - - def test_get_message_data(self, message_session: Session): - target = _message("message-id") - unrelated = _message("other-message-id") - message_session.add_all((target, unrelated)) - message_session.commit() - - result = get_message_data("message-id", message_session) - - assert result is target - assert result.id == "message-id" - - -class TestMeasureTime: - """Test cases for measure_time function""" - - def test_measure_time(self): - with measure_time() as timing_info: - assert "start" in timing_info - assert isinstance(timing_info["start"], datetime) - assert timing_info["end"] is None - - assert timing_info["end"] is not None - assert isinstance(timing_info["end"], datetime) - assert timing_info["end"] >= timing_info["start"] - - -class TestReplaceTextWithContent: - """Test cases for replace_text_with_content function""" - - def test_replace_text_with_content_dict(self): - data = {"text": "hello", "other": "world"} - assert replace_text_with_content(data) == {"content": "hello", "other": "world"} - - def test_replace_text_with_content_nested(self): - data = {"text": "v1", "nested": {"text": "v2", "list": [{"text": "v3"}]}} - expected = {"content": "v1", "nested": {"content": "v2", "list": [{"content": "v3"}]}} - assert replace_text_with_content(data) == expected - - def test_replace_text_with_content_list(self): - data = [{"text": "v1"}, "v2"] - assert replace_text_with_content(data) == [{"content": "v1"}, "v2"] - - def test_replace_text_with_content_primitive(self): - assert replace_text_with_content(123) == 123 - assert replace_text_with_content("text") == "text" - - -class TestValidateIntegerId: - """Test cases for validate_integer_id function""" - - def test_valid_integer_id(self): - assert validate_integer_id("123") == "123" - assert validate_integer_id(" 456 ") == "456" - - def test_invalid_integer_id_raises_error(self): - with pytest.raises(ValueError, match="ID must be a valid integer"): - validate_integer_id("abc") - - def test_empty_integer_id_raises_error(self): - with pytest.raises(ValueError, match="ID must be a valid integer"): - validate_integer_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 deleted file mode 100644 index e69de29bb2d1d6..00000000000000 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 deleted file mode 100644 index d4fb795f3b5cf5..00000000000000 --- a/api/tests/unit_tests/core/ops/unified_trace/test_entities.py +++ /dev/null @@ -1,108 +0,0 @@ -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 deleted file mode 100644 index e591e13934365d..00000000000000 --- a/api/tests/unit_tests/core/ops/unified_trace/test_hierarchy.py +++ /dev/null @@ -1,107 +0,0 @@ -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 deleted file mode 100644 index e41b56e2dd66db..00000000000000 --- a/api/tests/unit_tests/core/ops/unified_trace/test_parent_context.py +++ /dev/null @@ -1,184 +0,0 @@ -import json -from collections.abc import Callable - -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, - RedisParentContextStore, - destination_scope, - parent_destination_from_config, -) - - -class InMemoryParentContextStore(RedisParentContextStore): - def __init__(self, value: bytes | str | None = None) -> None: - self.value = value - self.get_calls: list[str] = [] - self.setex_calls: list[tuple[str, int, str]] = [] - self.read_error: Exception | None = None - self.write_error: Exception | None = None - - def get(self, name: str) -> bytes | str | None: - self.get_calls.append(name) - if self.read_error is not None: - raise self.read_error - return self.value - - def setex(self, name: str, time: int, value: str) -> object: - self.setex_calls.append((name, time, value)) - if self.write_error is not None: - raise self.write_error - self.value = value - return True - - -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: InMemoryParentContextStore, 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(config_overrides: Callable[..., None]) -> None: - redis = InMemoryParentContextStore() - value = context() - config_overrides(OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS=1_800) - - coordinator(redis, None).publish("outer-tool", value) - - key, ttl, payload = redis.setex_calls[-1] - 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 = InMemoryParentContextStore(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 = InMemoryParentContextStore(context().model_dump_json().encode()) - destination_calls: list[str] = [] - - def resolve_destination(workflow_run_id: str) -> ParentDestination | None: - destination_calls.append(workflow_run_id) - return None - - 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() - assert redis.get_calls == ["trace:unified:parent:message-1"] - assert destination_calls == [] - - -def test_missing_required_message_context_is_retryable() -> None: - redis = InMemoryParentContextStore() - 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 = InMemoryParentContextStore() - 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 = InMemoryParentContextStore() - 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() - - assert redis.get_calls == [] - - -def test_malformed_or_stale_context_is_terminal() -> None: - redis = InMemoryParentContextStore() - 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.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 = InMemoryParentContextStore() - redis.read_error = ConnectionError("down") - redis.write_error = 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 deleted file mode 100644 index c0c4b2b7310beb..00000000000000 --- a/api/tests/unit_tests/core/ops/unified_trace/test_provider.py +++ /dev/null @@ -1,120 +0,0 @@ -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 deleted file mode 100644 index 0dd38b94bfca76..00000000000000 --- a/api/tests/unit_tests/core/ops/unified_trace/test_registry.py +++ /dev/null @@ -1,14 +0,0 @@ -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 deleted file mode 100644 index 7997b87ae186f8..00000000000000 --- a/api/tests/unit_tests/core/ops/unified_trace/test_trace_builder.py +++ /dev/null @@ -1,676 +0,0 @@ -import importlib as _test_importlib -import json -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 -import requests -from sqlalchemy.orm import Session - -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, -) -from extensions.ext_database import db -from graphon.enums import WorkflowNodeExecutionStatus -from models import Account -from models.enums import CreatorUserRole -from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom - -"""External traces retain the caller's Tool contract without source workflow internals.""" - - -exported_spans = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-mlflow.src.dify_trace_mlflow.test_mlflow_trace" -).exported_spans - - -class TestWorkflowToolExportBoundaryCases: - import importlib as _test_importlib - - transport = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-mlflow.src.dify_trace_mlflow.test_mlflow_trace" - ).transport - - transport = staticmethod(transport) - - @staticmethod - @pytest.mark.parametrize("exporter", ["canonical", "mlflow"]) - def test_external_workflow_trace_omits_source_internals( - exporter: str, - sqlite_session: Session, - monkeypatch: pytest.MonkeyPatch, - transport: list[requests.PreparedRequest], - ) -> None: - started_at = datetime(2025, 1, 1) - public_inputs = {"query": "public tool input"} - public_outputs = {"answer": "public tool output"} - for index, (execution_id, app_id, origin) in enumerate( - [ - ("outer-tool", "caller-app", WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN), - ("source-llm", "source-app", WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL), - ("recursive-llm", "caller-app", WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL), - ] - ): - is_outer = execution_id == "outer-tool" - node_record = WorkflowNodeExecutionModel( - id=execution_id, - tenant_id="tenant-1", - app_id=app_id, - workflow_id=f"{app_id}-workflow", - workflow_run_id="run-1", - triggered_from_workflow_id=None if is_outer else "caller-app-workflow", - triggered_from_node_execution_id=None if is_outer else "outer-tool", - triggered_from=origin, - index=index, - node_execution_id=execution_id, - node_id=execution_id, - node_type="tool" if is_outer else "llm", - title=execution_id, - inputs=json.dumps(public_inputs if is_outer else {"secret": "private source input"}), - outputs=json.dumps(public_outputs if is_outer else {"text": "private source output"}), - process_data=json.dumps( - {} - if is_outer - else { - "prompts": [{"role": "system", "text": "private source prompt"}], - "workflow_tool_root_app_id": "caller-app", - } - ), - status=WorkflowNodeExecutionStatus.SUCCEEDED, - elapsed_time=1, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="account-1", - ) - node_record.created_at = started_at - sqlite_session.add(node_record) - sqlite_session.commit() - info = WorkflowTraceInfo( - tenant_id="tenant-1", - workflow_id="caller-workflow", - workflow_run_id="run-1", - workflow_run_elapsed_time=1, - workflow_run_status="succeeded", - workflow_run_inputs={}, - workflow_run_outputs={}, - workflow_run_version="1", - total_tokens=0, - file_list=[], - query="", - metadata={"app_id": "caller-app"}, - start_time=started_at, - end_time=started_at, - ) - monkeypatch.setattr(type(db), "engine", property(lambda _: sqlite_session.get_bind())) - monkeypatch.setattr(db, "session", sqlite_session) - if exporter == "canonical": - account = Account(name="Caller maintainer", email="caller@example.com") - account.id = "account-1" - trace = CanonicalTraceBuilder(RepositoryWorkflowExecutionLoader(lambda _: account)).build(info) - assert trace is not None - assert len(trace.spans) == 2 - payload = trace.model_dump_json() - tool = next(span for span in trace.spans if span.id == "outer-tool") - assert tool.inputs == public_inputs - assert tool.outputs == public_outputs - else: - from dify_trace_mlflow.config import MLflowConfig - from dify_trace_mlflow.mlflow_trace import MLflowDataTrace - from mlflow.entities import Span - from mlflow.tracing.trace_manager import InMemoryTraceManager - - monkeypatch.setattr(InMemoryTraceManager, "_instance", InMemoryTraceManager()) - # OSS owns a private synchronous SDK pipeline. Inspect its real wire spans - # instead of replacing the unrelated process-wide MLflow tracer. - instance = MLflowDataTrace( - MLflowConfig(tracking_uri="https://workflow-boundary.example", experiment_id="0") - ) - try: - instance.trace(info) - spans = [Span.from_otel_proto(span) for span in exported_spans(transport)] - assert len(spans) == 2 - payload = json.dumps([span.to_dict() for span in spans]) - tool = next(span for span in spans if span.name == "outer-tool") - assert tool.inputs == public_inputs - assert tool.outputs == public_outputs - assert transport - assert {request.url for request in transport} <= { - "https://workflow-boundary.example/version", - "https://workflow-boundary.example/api/3.0/mlflow/traces", - "https://workflow-boundary.example/v1/traces", - } - finally: - instance._provider.shutdown() - assert "private source" not in payload - assert "source-llm" not in payload - assert "recursive-llm" not in payload - assert "workflow_tool_root_app_id" not in payload - - -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 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 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 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) - - -class TestTraceBuilderCases: - @staticmethod - 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" - - @staticmethod - 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" - - @staticmethod - 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" - - @staticmethod - 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" - - @staticmethod - 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 - - @staticmethod - 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 - - @staticmethod - 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 - - @staticmethod - 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() - - @staticmethod - 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 - - @staticmethod - 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"]) - - @staticmethod - 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 - - @staticmethod - 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 - - @staticmethod - 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" - - @staticmethod - 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") - - @staticmethod - 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() - - @staticmethod - 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 - - @staticmethod - @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) - - @staticmethod - 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 - - @staticmethod - 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 - - @staticmethod - 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/core/rag/retrieval/test_dataset_retrieval.py b/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval.py index f735a8e2ae51bb..9820ef7a5dbac2 100644 --- a/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval.py +++ b/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval.py @@ -1,14 +1,16 @@ import threading from collections.abc import Generator -from contextlib import contextmanager, nullcontext +from contextlib import nullcontext from dataclasses import dataclass -from datetime import datetime +from datetime import UTC, datetime, timedelta from types import SimpleNamespace from typing import Any from unittest.mock import MagicMock, Mock, patch from uuid import uuid4 import pytest +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import DatasetRetrieved, TraceCorrelation, TraceEvent from flask import Flask, current_app from sqlalchemy import event, inspect, select from sqlalchemy.dialects import postgresql @@ -20,6 +22,9 @@ from core.app.app_config.entities import ( DatasetEntity, DatasetRetrieveConfigEntity, + EasyUIBasedAppModelConfigFrom, + ModelConfigEntity, + PromptTemplateEntity, ) from core.app.app_config.entities import ( MetadataFilteringCondition as AppMetadataFilteringCondition, @@ -28,7 +33,13 @@ ModelConfig as AppModelConfig, ) from core.app.app_config.entities import ModelConfig as WorkflowModelConfig -from core.app.entities.app_invoke_entities import InvokeFrom, ModelConfigWithCredentialsEntity, UserFrom +from core.app.apps.completion.app_config_manager import CompletionAppConfig +from core.app.entities.app_invoke_entities import ( + CompletionAppGenerateEntity, + InvokeFrom, + ModelConfigWithCredentialsEntity, + UserFrom, +) from core.app.file_access import FileAccessScope, bind_file_access_scope, get_current_file_access_scope from core.entities.agent_entities import PlanningStrategy from core.entities.model_entities import ModelStatus @@ -38,6 +49,7 @@ from core.rag.entities import Condition as AppCondition from core.rag.index_processor.constant.doc_type import DocType from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType +from core.rag.models.document import AttachmentDocument, ChildDocument from core.rag.models.document import Document as RetrievedDocument from core.rag.models.document import Document as _dataset_retrieval_Document from core.rag.models.document import Document as _dataset_retrieval_methods_Document @@ -50,6 +62,7 @@ from extensions.storage.storage_type import StorageType from graphon.model_runtime.entities.llm_entities import LLMUsage from graphon.model_runtime.entities.model_entities import ModelFeature +from libs.datetime_utils import Timing from models import UploadFile from models.dataset import ( ChildChunk, @@ -63,12 +76,63 @@ from models.dataset import Document as DatasetDocument from models.dataset import Document as _retrieval_source_query_Document from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, IndexingStatus +from models.model import AppMode from repositories.knowledge.retrieval_source_repository import get_retrieval_source_documents +from tests.helpers.unit.core.model_fixtures import make_model_config from tests.helpers.unit.model_factories import make_document # ==================== Helper Functions ==================== +class _Recorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +def _trace_scope(recorder: _Recorder) -> TraceScope: + return TraceScope( + emitter=recorder, + tenant_id="tenant", + app_id="app", + user_id="user", + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + ) + + +def _completion_generate_entity(trace_scope: TraceScope) -> CompletionAppGenerateEntity: + return CompletionAppGenerateEntity( + task_id="task", + app_config=CompletionAppConfig( + tenant_id="tenant", + app_id="app", + app_mode=AppMode.COMPLETION, + app_model_config_from=EasyUIBasedAppModelConfigFrom.APP_LATEST_CONFIG, + app_model_config_dict={}, + model=ModelConfigEntity(provider="openai", model="gpt-4", mode="completion"), + prompt_template=PromptTemplateEntity( + prompt_type=PromptTemplateEntity.PromptType.SIMPLE, simple_prompt_template="Answer the query." + ), + ), + model_conf=make_model_config(provider="openai", model="gpt-4", mode="completion"), + inputs={}, + files=[], + user_id="user", + stream=False, + invoke_from=InvokeFrom.SERVICE_API, + trace_scope=trace_scope, + ) + + +def _timing() -> Timing: + started_at = datetime(2026, 1, 1, tzinfo=UTC) + return Timing(started_at=started_at, ended_at=started_at + timedelta(seconds=1)) + + def create_mock_document( content: str, doc_id: str, @@ -343,11 +407,6 @@ def is_alive(self) -> bool: return self._alive -@contextmanager -def _timer(): - yield {"cost": 1} - - class TestDatasetRetrievalCases: class TestRetrievalService: """ @@ -4092,18 +4151,60 @@ def test_calculate_keyword_score(self, retrieval: DatasetRetrieval) -> None: assert "keywords" in ranked[0].metadata assert ranked[0].metadata["doc_id"] == "1" - def test_send_trace_task(self, retrieval: DatasetRetrieval) -> None: - trace_manager = Mock() - retrieval.application_generate_entity = SimpleNamespace(trace_manager=trace_manager) - docs = [_dataset_retrieval_Document(page_content="d", metadata={}, provider="dify")] + def test_send_trace_task_emits_documents_without_vectors(self) -> None: + recorder = _Recorder() + retrieval = DatasetRetrieval(_completion_generate_entity(_trace_scope(recorder))) + documents = [ + _dataset_retrieval_Document( + page_content="parent", + vector=[0.1, 0.2], + metadata={"doc_id": "n1", "score": 0.9}, + provider="dify", + children=[ChildDocument(page_content="child", vector=[0.3], metadata={"doc_id": "c1"})], + attachments=[AttachmentDocument(page_content="image", vector=[0.4], metadata={"doc_id": "a1"})], + ), + _dataset_retrieval_Document(page_content="external", vector=[0.5], metadata={}, provider="external"), + ] + timer = _timing() + + retrieval._send_trace_task("m1", documents, timer) - retrieval._send_trace_task("m1", docs, {"cost": 1}) - trace_manager.add_trace_task.assert_called_once() + assert [event.payload for event in recorder.events] == [ + DatasetRetrieved( + message_id="m1", + documents=[ + { + "page_content": "parent", + "metadata": {"doc_id": "n1", "score": 0.9}, + "provider": "dify", + "children": [{"page_content": "child", "metadata": {"doc_id": "c1"}}], + "attachments": [ + {"page_content": "image", "provider": "dify", "metadata": {"doc_id": "a1"}} + ], + }, + { + "page_content": "external", + "metadata": {}, + "provider": "external", + "children": None, + "attachments": None, + }, + ], + started_at=timer.started_at, + ended_at=timer.ended_at, + ) + ] + + def test_send_trace_task_requires_scope_and_timer(self) -> None: + recorder = _Recorder() + documents = [_dataset_retrieval_Document(page_content="d", metadata={}, provider="dify")] + + DatasetRetrieval()._send_trace_task("m1", documents, _timing()) + DatasetRetrieval(_completion_generate_entity(_trace_scope(recorder)))._send_trace_task( + "m1", documents, None + ) - retrieval.application_generate_entity = None - trace_manager.reset_mock() - retrieval._send_trace_task("m1", docs, {"cost": 1}) - trace_manager.add_trace_task.assert_not_called() + assert recorder.events == [] def test_on_query(self, retrieval: DatasetRetrieval, sqlite_engine: Engine, sqlite_session: Session) -> None: dataset_ids = [str(uuid4()), str(uuid4())] @@ -5357,7 +5458,6 @@ def fake_multiple_thread(**kwargs): with app.app_context(): with ( - patch("core.rag.retrieval.dataset_retrieval.measure_time", _timer), patch("core.rag.retrieval.dataset_retrieval.threading.Thread", _ImmediateThread), patch.object(retrieval, "_multiple_retrieve_thread", side_effect=fake_multiple_thread), patch.object(retrieval, "_on_query") as mock_on_query, @@ -5402,7 +5502,6 @@ def failing_thread(**_kwargs): with app.app_context(): with ( - patch("core.rag.retrieval.dataset_retrieval.measure_time", _timer), patch("core.rag.retrieval.dataset_retrieval.threading.Thread", _ImmediateThread), patch.object(retrieval, "_multiple_retrieve_thread", side_effect=failing_thread), ): @@ -5436,7 +5535,7 @@ def test_on_retrieval_end_without_dify_documents(self, retrieval: DatasetRetriev flask_app=app, documents=[_doc(provider="external")], message_id="m1", - timer={"cost": 1}, + timer=_timing(), ) mock_trace.assert_called_once() @@ -5447,7 +5546,7 @@ def test_on_retrieval_end_dify_without_document_ids(self, retrieval: DatasetRetr patch("core.rag.retrieval.dataset_retrieval.db", SimpleNamespace(engine=self.orm_engine)), patch.object(retrieval, "_send_trace_task") as mock_trace, ): - retrieval._on_retrieval_end(flask_app=app, documents=[doc], message_id="m1", timer={"cost": 1}) + retrieval._on_retrieval_end(flask_app=app, documents=[doc], message_id="m1", timer=_timing()) mock_trace.assert_called_once() def test_on_retrieval_end_updates_segments_for_text_and_image(self, retrieval: DatasetRetrieval) -> None: @@ -5569,7 +5668,7 @@ def capture_statement(*args: object) -> None: patch("core.rag.retrieval.dataset_retrieval.db", SimpleNamespace(engine=self.orm_engine)), patch.object(retrieval, "_send_trace_task") as mock_trace, ): - retrieval._on_retrieval_end(flask_app=app, documents=docs, message_id="m1", timer={"cost": 1}) + retrieval._on_retrieval_end(flask_app=app, documents=docs, message_id="m1", timer=_timing()) finally: event.remove(self.orm_engine, "before_execute", capture_statement) @@ -5586,6 +5685,7 @@ def capture_statement(*args: object) -> None: def test_on_retrieval_end_retries_postgres_deadlock( self, retrieval: DatasetRetrieval, sqlite_session_factory: sessionmaker[Session] ) -> None: + """A first-update deadlock must roll back, retry once, and emit one timed trace.""" app = Flask(__name__) deadlock = DBAPIError( "UPDATE document_segments", @@ -5611,7 +5711,7 @@ def deadlock_first_update(_conn, _cursor, statement: str, _parameters, _context, patch("core.rag.retrieval.dataset_retrieval.db", SimpleNamespace(engine=self.orm_engine)), patch.object(retrieval, "_send_trace_task") as mock_trace, ): - retrieval._on_retrieval_end(flask_app=app, documents=[doc], message_id="m1", timer={"cost": 1}) + retrieval._on_retrieval_end(flask_app=app, documents=[doc], message_id="m1", timer=_timing()) finally: event.remove(self.orm_engine, "before_cursor_execute", deadlock_first_update) @@ -5620,7 +5720,7 @@ def deadlock_first_update(_conn, _cursor, statement: str, _parameters, _context, assert ( self.orm_session.scalar(select(DocumentSegment.hit_count).where(DocumentSegment.id == segment.id)) == 6 ) - mock_trace.assert_called_once_with("m1", [doc], {"cost": 1}) + mock_trace.assert_called_once_with("m1", [doc], _timing()) @pytest.mark.parametrize( ("error", "expected"), diff --git a/api/tests/unit_tests/core/telemetry/test___init__.py b/api/tests/unit_tests/core/telemetry/test___init__.py index ac8a02ac65da34..2692242d087127 100644 --- a/api/tests/unit_tests/core/telemetry/test___init__.py +++ b/api/tests/unit_tests/core/telemetry/test___init__.py @@ -2,115 +2,111 @@ from __future__ import annotations -import queue -from collections.abc import Callable, Iterator -from unittest.mock import MagicMock, Mock, patch +from collections.abc import Mapping +from dataclasses import dataclass, field +from datetime import UTC, datetime import pytest +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import DraftNodeExecuted, TraceCorrelation, TraceEvent -from core.ops.entities.trace_entity import TraceTaskName -from core.ops.ops_trace_manager import TraceQueueManager, TraceTask from core.telemetry import emit +from core.telemetry import gateway as gateway_module from core.telemetry.events import DraftNodeExecutionTraceEvent, TelemetryContext +type _ScopeRequest = tuple[str, str | None, str | None, Mapping[str, object]] -@pytest.fixture -def telemetry_test_setup() -> Iterator[tuple[Callable[..., object], queue.Queue[TraceTask]]]: - trace_queue: queue.Queue[TraceTask] = queue.Queue() - with ( - patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None), - patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), - patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False), - patch.object(TraceQueueManager, "start_timer"), - ): - yield emit, trace_queue - - -_TelemetrySetup = tuple[Callable[..., object], "queue.Queue[TraceTask]"] +@dataclass +class _TracePipeline: + """Stands in for the bound trace pipeline: hands out scopes and records what they emit.""" -class TestTelemetryEmit: - @patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True) - def test_emit_enterprise_trace_creates_trace_task( - self, mock_ee: Mock, telemetry_test_setup: _TelemetrySetup - ) -> None: - emit_fn, trace_queue = telemetry_test_setup + scope_requests: list[_ScopeRequest] = field(default_factory=list) + events: list[TraceEvent] = field(default_factory=list) - event = DraftNodeExecutionTraceEvent( - context=TelemetryContext( - tenant_id="test-tenant", - user_id="test-user", - app_id="test-app", + def create_scope( + self, *, tenant_id: str, app_id: str | None, user_id: str | None, extras: Mapping[str, object] + ) -> TraceScope: + self.scope_requests.append((tenant_id, app_id, user_id, extras)) + return TraceScope( + emitter=self, + tenant_id=tenant_id, + app_id=app_id, + user_id=user_id, + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None ), - payload={"node_execution_data": {"key": "value"}}, ) - emit_fn(event) + def emit(self, event: TraceEvent) -> None: + self.events.append(event) - called_task = trace_queue.get_nowait() - assert trace_queue.empty() - assert called_task.app_id == "test-app" - assert called_task.user_id == "test-user" - assert called_task.trace_type == TraceTaskName.DRAFT_NODE_EXECUTION_TRACE - def test_emit_enterprise_only_trace_dropped_when_ee_disabled(self, telemetry_test_setup: _TelemetrySetup) -> None: - emit_fn, trace_queue = telemetry_test_setup +@pytest.fixture +def pipeline(monkeypatch: pytest.MonkeyPatch) -> _TracePipeline: + pipeline = _TracePipeline() + monkeypatch.setattr(gateway_module, "create_trace_scope_for_app", pipeline.create_scope) + return pipeline - event = DraftNodeExecutionTraceEvent( - context=TelemetryContext( - tenant_id="test-tenant", - user_id="test-user", - app_id="test-app", - ), - payload={"node_execution_data": {}}, - ) - emit_fn(event) +def _enable_enterprise_telemetry(monkeypatch: pytest.MonkeyPatch, enabled: bool) -> None: + monkeypatch.setattr(gateway_module, "is_enterprise_telemetry_enabled", lambda: enabled) + - assert trace_queue.empty() +def _draft_node_event(node_execution_data: dict[str, object]) -> DraftNodeExecutionTraceEvent: + return DraftNodeExecutionTraceEvent( + context=TelemetryContext(tenant_id="test-tenant", user_id="test-user", app_id="test-app"), + payload={"node_execution_data": node_execution_data}, + ) - @patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True) - def test_emit_passes_name_directly_to_trace_task( - self, mock_ee: Mock, telemetry_test_setup: _TelemetrySetup + +class TestTelemetryEmit: + def test_emit_enterprise_trace_reports_the_draft_node_execution( + self, monkeypatch: pytest.MonkeyPatch, pipeline: _TracePipeline ) -> None: - emit_fn, trace_queue = telemetry_test_setup + _enable_enterprise_telemetry(monkeypatch, True) - event = DraftNodeExecutionTraceEvent( - context=TelemetryContext( - tenant_id="test-tenant", - user_id="test-user", - app_id="test-app", - ), - payload={"node_execution_data": {"extra": "data"}}, + emit(_draft_node_event({"node_id": "node-1", "status": "succeeded"})) + + assert pipeline.scope_requests == [("test-tenant", "test-app", "test-user", {})] + (event,) = pipeline.events + assert (event.tenant_id, event.app_id, event.user_id) == ("test-tenant", "test-app", "test-user") + assert event.payload == DraftNodeExecuted(node_execution={"node_id": "node-1", "status": "succeeded"}) + + def test_emit_makes_node_execution_values_json_safe( + self, monkeypatch: pytest.MonkeyPatch, pipeline: _TracePipeline + ) -> None: + _enable_enterprise_telemetry(monkeypatch, True) + created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC) + + emit(_draft_node_event({"created_at": created_at, "elapsed_time": float("nan"), "node_inputs": {"q": 1}})) + + (event,) = pipeline.events + assert event.payload == DraftNodeExecuted( + node_execution={"created_at": str(created_at), "elapsed_time": "NaN", "node_inputs": {"q": 1}} ) - emit_fn(event) + def test_emit_enterprise_only_trace_dropped_when_ee_disabled( + self, monkeypatch: pytest.MonkeyPatch, pipeline: _TracePipeline + ) -> None: + _enable_enterprise_telemetry(monkeypatch, False) - called_task = trace_queue.get_nowait() - assert trace_queue.empty() - assert called_task.app_id == "test-app" - assert called_task.user_id == "test-user" - assert called_task.trace_type == TraceTaskName.DRAFT_NODE_EXECUTION_TRACE - assert isinstance(called_task.trace_type, TraceTaskName) + emit(_draft_node_event({})) - @patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True) - def test_emit_with_provided_trace_manager(self, mock_ee: Mock, telemetry_test_setup: _TelemetrySetup) -> None: - emit_fn, trace_queue = telemetry_test_setup + assert pipeline.scope_requests == [] + assert pipeline.events == [] - mock_trace_manager = MagicMock() - mock_trace_manager.add_trace_task = MagicMock() + def test_emit_drops_trace_without_a_workspace( + self, monkeypatch: pytest.MonkeyPatch, pipeline: _TracePipeline + ) -> None: + _enable_enterprise_telemetry(monkeypatch, True) - event = DraftNodeExecutionTraceEvent( - context=TelemetryContext( - tenant_id="test-tenant", - user_id="test-user", - app_id="test-app", - ), - payload={"node_execution_data": {}}, + emit( + DraftNodeExecutionTraceEvent( + context=TelemetryContext(tenant_id=None, user_id="test-user", app_id="test-app"), + payload={"node_execution_data": {}}, + ) ) - emit_fn(event, trace_manager=mock_trace_manager) - - mock_trace_manager.add_trace_task.assert_called_once() - called_task = mock_trace_manager.add_trace_task.call_args[0][0] - assert called_task.trace_type == TraceTaskName.DRAFT_NODE_EXECUTION_TRACE + assert pipeline.scope_requests == [] + assert pipeline.events == [] diff --git a/api/tests/unit_tests/core/telemetry/test_gateway.py b/api/tests/unit_tests/core/telemetry/test_gateway.py index 8569c5388c04af..6df19f851f74a3 100644 --- a/api/tests/unit_tests/core/telemetry/test_gateway.py +++ b/api/tests/unit_tests/core/telemetry/test_gateway.py @@ -1,11 +1,16 @@ from __future__ import annotations import sys +from collections.abc import Mapping +from dataclasses import dataclass, field +from datetime import UTC, datetime from unittest.mock import MagicMock, patch import pytest +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import DraftNodeExecuted, PromptGenerated, TraceCorrelation, TraceEvent -from core.ops.entities.trace_entity import TraceTaskName +from core.telemetry import gateway as gateway_module from core.telemetry.events import ( AppCreatedEvent, AppCreatedPayload, @@ -15,6 +20,7 @@ PromptGenerationPayload, TelemetryContext, TelemetryEvent, + TraceSignalEvent, ) from core.telemetry.gateway import PAYLOAD_SIZE_THRESHOLD_BYTES, emit, is_enterprise_telemetry_enabled from enterprise.telemetry.contracts import SignalType, TelemetryCase, TelemetryEnvelope @@ -32,12 +38,14 @@ def test_is_enterprise_telemetry_enabled_exported(self) -> None: # --------------------------------------------------------------------------- _GATEWAY_INTEGRATION_TRACE_CTX = TelemetryContext(tenant_id="tenant-789", user_id="user-456", app_id="app-123") +_STARTED_AT = datetime(2026, 10, 1, 8, 0, 0, tzinfo=UTC) +_ENDED_AT = datetime(2026, 10, 1, 8, 0, 1, tzinfo=UTC) def _gateway_integration_node_event() -> DraftNodeExecutionTraceEvent: return DraftNodeExecutionTraceEvent( context=_GATEWAY_INTEGRATION_TRACE_CTX, - payload=NodeExecutionPayload(node_execution_data={}), + payload=NodeExecutionPayload(node_execution_data={"node_id": "node-1", "node_type": "llm"}), ) @@ -55,10 +63,40 @@ def _gateway_integration_prompt_event() -> PromptGenerationEvent: completion_tokens=5, total_tokens=15, latency=0.5, + total_price=0.002, + currency="USD", + error=None, + started_at=_STARTED_AT, + ended_at=_ENDED_AT, ), ) +_EXPECTED_PAYLOADS = { + "draft_node_execution": DraftNodeExecuted(node_execution={"node_id": "node-1", "node_type": "llm"}), + "prompt_generation": PromptGenerated( + operation_type="generate", + instruction="test", + generated_output="out", + model_provider="openai", + model_name="gpt-4", + prompt_tokens=10, + completion_tokens=5, + total_tokens=15, + latency=0.5, + total_price=0.002, + currency="USD", + error=None, + started_at=_STARTED_AT, + ended_at=_ENDED_AT, + ), +} +_TRACE_EVENTS = { + "draft_node_execution": _gateway_integration_node_event, + "prompt_generation": _gateway_integration_prompt_event, +} + + def _gateway_integration_app_created_event() -> AppCreatedEvent: return AppCreatedEvent( context=TelemetryContext(tenant_id="tenant-123"), @@ -66,39 +104,71 @@ def _gateway_integration_app_created_event() -> AppCreatedEvent: ) +@dataclass +class _TracePipeline: + """Stands in for the bound trace pipeline: hands out scopes and records what they emit.""" + + scope_requests: list[tuple[str, str | None, str | None, Mapping[str, object]]] = field(default_factory=list) + events: list[TraceEvent] = field(default_factory=list) + + def create_scope( + self, *, tenant_id: str, app_id: str | None, user_id: str | None, extras: Mapping[str, object] + ) -> TraceScope: + self.scope_requests.append((tenant_id, app_id, user_id, extras)) + return TraceScope( + emitter=self, + tenant_id=tenant_id, + app_id=app_id, + user_id=user_id, + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + ) + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +@pytest.fixture +def pipeline(monkeypatch: pytest.MonkeyPatch) -> _TracePipeline: + pipeline = _TracePipeline() + monkeypatch.setattr(gateway_module, "create_trace_scope_for_app", pipeline.create_scope) + return pipeline + + # --------------------------------------------------------------------------- # Trace routing integration # --------------------------------------------------------------------------- class TestGatewayIntegrationTraceRouting: - @pytest.fixture - def mock_trace_manager(self) -> MagicMock: - return MagicMock() - - def test_trace_event_routed_to_trace_manager( - self, - mock_trace_manager: MagicMock, + @pytest.mark.parametrize("kind", ["draft_node_execution", "prompt_generation"]) + def test_trace_event_routed_through_the_app_trace_scope( + self, monkeypatch: pytest.MonkeyPatch, pipeline: _TracePipeline, kind: str ) -> None: - with patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True): - emit(_gateway_integration_prompt_event(), mock_trace_manager) - mock_trace_manager.add_trace_task.assert_called_once() + monkeypatch.setattr(gateway_module, "is_enterprise_telemetry_enabled", lambda: True) - def test_enterprise_only_trace_dropped_when_ee_disabled( - self, - mock_trace_manager: MagicMock, - ) -> None: - with patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False): - emit(_gateway_integration_node_event(), mock_trace_manager) - mock_trace_manager.add_trace_task.assert_not_called() + emit(_TRACE_EVENTS[kind]()) - def test_enterprise_only_trace_routed_when_ee_enabled( - self, - mock_trace_manager: MagicMock, - ) -> None: - with patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True): - emit(_gateway_integration_node_event(), mock_trace_manager) - mock_trace_manager.add_trace_task.assert_called_once() + assert pipeline.scope_requests == [("tenant-789", "app-123", "user-456", {})] + (event,) = pipeline.events + assert (event.tenant_id, event.app_id, event.user_id) == ("tenant-789", "app-123", "user-456") + assert event.payload == _EXPECTED_PAYLOADS[kind] + + def test_trace_event_without_a_consumer_is_dropped(self, monkeypatch: pytest.MonkeyPatch) -> None: + requests: list[tuple[str, str | None, str | None, Mapping[str, object]]] = [] + + def no_consumer( + *, tenant_id: str, app_id: str | None, user_id: str | None, extras: Mapping[str, object] + ) -> None: + requests.append((tenant_id, app_id, user_id, extras)) + + monkeypatch.setattr(gateway_module, "is_enterprise_telemetry_enabled", lambda: True) + monkeypatch.setattr(gateway_module, "create_trace_scope_for_app", no_consumer) + + emit(_gateway_integration_prompt_event()) + + assert requests == [("tenant-789", "app-123", "user-456", {})] # --------------------------------------------------------------------------- @@ -131,25 +201,16 @@ def test_metric_event_routes_to_celery_task( class TestGatewayIntegrationCEEligibility: - @pytest.fixture - def mock_trace_manager(self) -> MagicMock: - return MagicMock() - - def test_draft_node_execution_not_ce_eligible( - self, - mock_trace_manager: MagicMock, + @pytest.mark.parametrize("kind", ["draft_node_execution", "prompt_generation"]) + def test_enterprise_only_trace_dropped_when_ee_disabled( + self, monkeypatch: pytest.MonkeyPatch, pipeline: _TracePipeline, kind: str ) -> None: - with patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False): - emit(_gateway_integration_node_event(), mock_trace_manager) - mock_trace_manager.add_trace_task.assert_not_called() + monkeypatch.setattr(gateway_module, "is_enterprise_telemetry_enabled", lambda: False) - def test_prompt_generation_not_ce_eligible( - self, - mock_trace_manager: MagicMock, - ) -> None: - with patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False): - emit(_gateway_integration_prompt_event(), mock_trace_manager) - mock_trace_manager.add_trace_task.assert_not_called() + emit(_TRACE_EVENTS[kind]()) + + assert pipeline.scope_requests == [] + assert pipeline.events == [] # --------------------------------------------------------------------------- @@ -196,6 +257,11 @@ def _gateway_prompt_event() -> PromptGenerationEvent: completion_tokens=5, total_tokens=15, latency=0.5, + total_price=None, + currency=None, + error=None, + started_at=datetime(2026, 10, 1, 8, 0, 0, tzinfo=UTC), + ended_at=datetime(2026, 10, 1, 8, 0, 1, tzinfo=UTC), ), ) @@ -220,21 +286,23 @@ def test_draft_node_execution_routing(self) -> None: assert ev.case == TelemetryCase.DRAFT_NODE_EXECUTION assert ev.signal_type is SignalType.TRACE assert ev.ce_eligible is False - assert ev.trace_task_name is TraceTaskName.DRAFT_NODE_EXECUTION_TRACE + assert isinstance(ev, TraceSignalEvent) + assert ev.trace_payload().kind == "draft_node_execution" def test_prompt_generation_routing(self) -> None: ev = _gateway_prompt_event() assert ev.case == TelemetryCase.PROMPT_GENERATION assert ev.signal_type is SignalType.TRACE assert ev.ce_eligible is False - assert ev.trace_task_name is TraceTaskName.PROMPT_GENERATION_TRACE + assert isinstance(ev, TraceSignalEvent) + assert ev.trace_payload().kind == "prompt_generation" def test_app_created_routing(self) -> None: ev = _gateway_app_created_event() assert ev.case == TelemetryCase.APP_CREATED assert ev.signal_type is SignalType.METRIC_LOG assert ev.ce_eligible is False - assert ev.trace_task_name is None + assert not isinstance(ev, TraceSignalEvent) def test_all_events_satisfy_protocol(self) -> None: events: list[TelemetryEvent] = [_gateway_node_event(), _gateway_prompt_event(), _gateway_app_created_event()] @@ -248,36 +316,37 @@ def test_all_events_satisfy_protocol(self) -> None: class TestGatewayTraceRouting: - @pytest.fixture - def mock_trace_manager(self) -> MagicMock: - return MagicMock() - - @patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True) - def test_trace_event_routes_to_trace_manager( - self, - mock_ee_enabled: MagicMock, - mock_trace_manager: MagicMock, + def test_trace_event_routes_through_the_app_trace_scope( + self, monkeypatch: pytest.MonkeyPatch, pipeline: _TracePipeline ) -> None: - emit(_gateway_prompt_event(), mock_trace_manager) - mock_trace_manager.add_trace_task.assert_called_once() + monkeypatch.setattr(gateway_module, "is_enterprise_telemetry_enabled", lambda: True) + + emit(_gateway_prompt_event()) + + assert pipeline.scope_requests == [("tenant-789", "app-123", "user-456", {})] + (event,) = pipeline.events + assert (event.tenant_id, event.app_id, event.user_id) == ("tenant-789", "app-123", "user-456") + assert event.payload == _gateway_prompt_event().trace_payload() - @patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False) def test_enterprise_only_trace_dropped_when_ee_disabled( - self, - mock_ee_enabled: MagicMock, - mock_trace_manager: MagicMock, + self, monkeypatch: pytest.MonkeyPatch, pipeline: _TracePipeline ) -> None: - emit(_gateway_node_event(), mock_trace_manager) - mock_trace_manager.add_trace_task.assert_not_called() + monkeypatch.setattr(gateway_module, "is_enterprise_telemetry_enabled", lambda: False) - @patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True) - def test_enterprise_only_trace_enqueued_when_ee_enabled( - self, - mock_ee_enabled: MagicMock, - mock_trace_manager: MagicMock, + emit(_gateway_node_event()) + + assert pipeline.events == [] + + def test_enterprise_only_trace_emitted_when_ee_enabled( + self, monkeypatch: pytest.MonkeyPatch, pipeline: _TracePipeline ) -> None: - emit(_gateway_node_event(), mock_trace_manager) - mock_trace_manager.add_trace_task.assert_called_once() + monkeypatch.setattr(gateway_module, "is_enterprise_telemetry_enabled", lambda: True) + + emit(_gateway_node_event()) + + assert pipeline.scope_requests == [("tenant-789", "app-123", "user-456", {})] + (event,) = pipeline.events + assert event.payload == _gateway_node_event().trace_payload() # --------------------------------------------------------------------------- diff --git a/api/tests/unit_tests/core/tools/test_tool_engine.py b/api/tests/unit_tests/core/tools/test_tool_engine.py index 7082e46e3f47ba..54abd16923187e 100644 --- a/api/tests/unit_tests/core/tools/test_tool_engine.py +++ b/api/tests/unit_tests/core/tools/test_tool_engine.py @@ -2,11 +2,14 @@ from collections.abc import Generator from dataclasses import dataclass +from datetime import timedelta from typing import Any from unittest.mock import patch from uuid import uuid4 import pytest +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import ToolInvoked, TraceCorrelation, TraceEvent from sqlalchemy import select from sqlalchemy.engine import Engine from sqlalchemy.orm import Session @@ -40,6 +43,26 @@ class _DatabaseBinding: engine: Engine +class _Recorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +def _trace_scope(recorder: _Recorder) -> TraceScope: + return TraceScope( + emitter=recorder, + tenant_id="tenant-1", + app_id="app-1", + user_id="u1", + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + ) + + def _message() -> Message: message = Message( app_id=str(uuid4()), @@ -331,6 +354,7 @@ def test_agent_invoke_success(sqlite_session: Session): message=message, invoke_from=InvokeFrom.DEBUGGER, agent_tool_callback=callback, + trace_scope=None, ) assert result_text == "ok" @@ -338,6 +362,53 @@ def test_agent_invoke_success(sqlite_session: Session): assert result_meta.error is None +def test_agent_invoke_emits_tool_invoked_from_meta( + sqlite_engine: Engine, sqlite_session: Session, monkeypatch: pytest.MonkeyPatch +) -> None: + tool = _build_tool(with_llm_parameter=True) + message = _message() + meta = ToolInvokeMeta( + time_cost=1.25, + error="partial failure", + tool_config={"tool_name": "tool-a", "tool_provider": "provider-a", "tool_parameters": {"rt": 1}}, + ) + recorder = _Recorder() + monkeypatch.setattr( + ToolEngine, "_invoke", staticmethod(lambda *_args: iter([tool.create_text_message("ok"), meta])) + ) + monkeypatch.setattr("core.tools.tool_engine.db", _DatabaseBinding(sqlite_engine)) + + ToolEngine.agent_invoke( + session=sqlite_session, + tool=tool, + tool_parameters="hello", + user_id="u1", + tenant_id="tenant-1", + message=message, + invoke_from=InvokeFrom.DEBUGGER, + agent_tool_callback=DifyAgentCallbackHandler(), + trace_scope=_trace_scope(recorder), + ) + + [event] = recorder.events + payload = event.payload + assert isinstance(payload, ToolInvoked) + assert payload == ToolInvoked( + message_id=message.id, + tool_name="tool-a", + tool_inputs={"query": "hello"}, + tool_outputs="ok", + tool_config={"tool_name": "tool-a", "tool_provider": "provider-a", "tool_parameters": {"rt": 1}}, + error="partial failure", + time_cost=1.25, + started_at=payload.started_at, + ended_at=payload.ended_at, + ) + assert payload.started_at.utcoffset() == timedelta(0) + assert payload.ended_at.utcoffset() == timedelta(0) + assert payload.ended_at >= payload.started_at + + @pytest.mark.parametrize("sqlite_session", [()], indirect=True) def test_agent_invoke_param_validation_error(sqlite_session: Session): tool = _build_tool(with_llm_parameter=True) @@ -354,6 +425,7 @@ def test_agent_invoke_param_validation_error(sqlite_session: Session): message=message, invoke_from=InvokeFrom.DEBUGGER, agent_tool_callback=callback, + trace_scope=None, ) assert "tool parameters validation error" in error_text @@ -378,6 +450,7 @@ def test_agent_invoke_engine_meta_error(sqlite_session: Session): message=message, invoke_from=InvokeFrom.DEBUGGER, agent_tool_callback=callback, + trace_scope=None, ) assert "meta failure" in error_text @@ -423,6 +496,7 @@ def test_agent_invoke_tool_invoke_error(sqlite_session: Session): message=message, invoke_from=InvokeFrom.DEBUGGER, agent_tool_callback=callback, + trace_scope=None, ) assert "tool invoke error" in error_text diff --git a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_entities.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_entities.py new file mode 100644 index 00000000000000..84d6aba2261cd4 --- /dev/null +++ b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_entities.py @@ -0,0 +1,9 @@ +import pytest +from pydantic import ValidationError + +from core.workflow.nodes.agent_v2.entities import DifyAgentNodeData + + +def test_dify_agent_node_data_requires_explicit_kind_marker() -> None: + with pytest.raises(ValidationError, match="agent_node_kind"): + DifyAgentNodeData.model_validate({"type": "agent", "version": "2"}) diff --git a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_runtime_request_builder.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_runtime_request_builder.py index a9471bfe52346d..acb0da114f07aa 100644 --- a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_runtime_request_builder.py +++ b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_runtime_request_builder.py @@ -20,7 +20,6 @@ ) from clients.agent_backend.request_builder import DIFY_SHELL_LAYER_ID from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom -from core.ops.base_trace_instance import bind_application_tracing_hooks from core.workflow.file_reference import build_file_reference from core.workflow.nodes.agent_v2.dify_tools_builder import WorkflowAgentDifyToolsBuilder from core.workflow.nodes.agent_v2.runtime_request_builder import ( @@ -30,6 +29,7 @@ build_knowledge_layer_config, build_shell_layer_config, ) +from extensions.tracing import bind_tracing from graphon.file import File, FileTransferMethod, FileType from graphon.variables.segments import ArrayFileSegment, FileSegment, StringSegment from models.agent import ( @@ -50,14 +50,10 @@ DeclaredOutputType, WorkflowNodeJobConfig, ) -from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks +from services.tracing.adapters import CeleryTraceEmitter from tests.helpers.unit.config_override import apply_config_overrides from tests.helpers.unit.knowledge_fs_core_fakes import RecordingModelFactory, llm_schema - - -@pytest.fixture(autouse=True) -def _application_tracing_disabled(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) +from tests.helpers.unit.services.tracing_ports import NoTracingPolicies, RecordingTracingHooks, kernel_trace_scopes @pytest.fixture(autouse=True) @@ -246,9 +242,7 @@ def _context() -> WorkflowAgentRuntimeBuildContext: @pytest.mark.parametrize("preview", [False, True]) -def test_builder_passes_inline_or_preview_context_to_application_tracing( - monkeypatch: pytest.MonkeyPatch, preview: bool -): +def test_builder_passes_inline_or_preview_context_to_application_tracing(preview: bool): from uuid import uuid4 from dify_agent.layers.execution_context import ApplicationTracingContext, DifyExecutionContextLayerConfig @@ -261,24 +255,20 @@ def test_builder_passes_inline_or_preview_context_to_application_tracing( session_id="session", capture_content=False, ) - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) context = replace(_context(), workflow_run_id=None if preview else "run-1", attempt=2) - calls: list[tuple[WorkflowAgentRuntimeBuildContext, str | None]] = [] - - def prepare(value: WorkflowAgentRuntimeBuildContext, conversation_id: str | None) -> ApplicationTracingContext: - calls.append((value, conversation_id)) - return tracing - - monkeypatch.setattr(RuntimeApplicationTracingHooks, "prepare_workflow_agent", staticmethod(prepare)) - hooks = RuntimeApplicationTracingHooks() + hooks = RecordingTracingHooks(tracing) app = Flask(__name__) - bind_application_tracing_hooks(app, hooks) + bind_tracing( + app, + scopes=kernel_trace_scopes(CeleryTraceEmitter(), policies=NoTracingPolicies(), enterprise_enabled=False), + hooks=hooks, + ) with app.app_context(): result = WorkflowAgentRuntimeRequestBuilder(dify_tools_builder=WorkflowAgentDifyToolsBuilder()).build(context) layer = next(item for item in result.request.composition.layers if item.name == DIFY_EXECUTION_CONTEXT_LAYER_ID) assert DifyExecutionContextLayerConfig.model_validate(layer.config).tracing == tracing - assert calls == [(context, None)] + assert hooks.workflow_agents == [(context, None)] def _request_layers(result) -> dict[str, dict[str, object]]: diff --git a/api/tests/unit_tests/core/workflow/test_node_runtime.py b/api/tests/unit_tests/core/workflow/test_node_runtime.py index 0f97bb0339cb21..e541b765b8df20 100644 --- a/api/tests/unit_tests/core/workflow/test_node_runtime.py +++ b/api/tests/unit_tests/core/workflow/test_node_runtime.py @@ -2,9 +2,7 @@ import json import os -import sys import time -import types import uuid from collections.abc import Generator, Iterator, Mapping from datetime import UTC, datetime @@ -1042,14 +1040,7 @@ def _build_tool_node_data() -> ToolNodeData: class TestToolNodeRuntimeCases: @staticmethod @pytest.fixture - def runtime(monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine) -> DifyToolNodeRuntime: - module_name = "core.ops.ops_trace_manager" - if module_name not in sys.modules: - ops_stub = types.ModuleType(module_name) - ops_stub.TraceQueueManager = object # pragma: no cover - stub attribute - ops_stub.TraceTask = object # pragma: no cover - stub attribute - monkeypatch.setitem(sys.modules, module_name, ops_stub) - + def runtime(sqlite_engine: Engine) -> DifyToolNodeRuntime: init_params = build_test_graph_init_params( workflow_id="workflow-id", graph_config={"nodes": [], "edges": []}, diff --git a/api/tests/unit_tests/core/workflow/test_workflow_tool_container_handler.py b/api/tests/unit_tests/core/workflow/test_workflow_tool_container_handler.py index 3232a564dc0b65..f5a4e6fc6959e2 100644 --- a/api/tests/unit_tests/core/workflow/test_workflow_tool_container_handler.py +++ b/api/tests/unit_tests/core/workflow/test_workflow_tool_container_handler.py @@ -1754,6 +1754,7 @@ def fail_caller_save( app_id="outer-app", triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, ), + trace_scope=None, ) index = WorkflowRunIndex(()) index.seed_source(source.app_id, source.workflow_id, ()) diff --git a/api/tests/unit_tests/domain/tracing/test_normalization.py b/api/tests/unit_tests/domain/tracing/test_normalization.py new file mode 100644 index 00000000000000..fa74baab3ad8a4 --- /dev/null +++ b/api/tests/unit_tests/domain/tracing/test_normalization.py @@ -0,0 +1,363 @@ +"""Content policy and model-call accounting across all canonical tree layers.""" + +import json +from datetime import UTC, datetime + +import pytest +from dify_trace_sdk.ids import span_id_for +from dify_trace_sdk.model import Cost, Span, SpanKind, SpanStatus, TokenUsage, Trace +from pydantic import ValidationError + +from domain.tracing.normalization import ( + legacy_outer_trace, + normalize_agent_span, + order_spans, + remap_outer_trace, + safe_attributes, +) + + +def agent_data( + attributes: dict[str, object], *, span_id: str = "2" * 16, parent_id: str = "1" * 16 +) -> dict[str, object]: + return { + "span_id": span_id, + "parent_span_id": parent_id, + "name": "operation", + "start_time_ns": 1_000_000_000, + "end_time_ns": 2_000_000_000, + "status": "error", + "status_message": "sensitive error output", + "attributes": attributes, + } + + +def test_content_off_is_safe_for_all_metadata_namespaces() -> None: + normalized = normalize_agent_span( + agent_data( + { + "gen_ai.operation.name": "chat", + "gen_ai.request.model": "model-one", + "gen_ai.usage.input_tokens": 3, + "gen_ai.usage.output_tokens": 2, + "gen_ai.input.messages": [{"content": "private input"}], + "gen_ai.output.messages": [{"content": "private output"}], + "exception.message": "private exception", + "dify.app_id": "app-one", + "dify.custom_payload": {"prompt": "private metadata"}, + "gen_ai.request.messages": "private request metadata", + "gen_ai.response.output": "private response metadata", + "gen_ai.usage.custom_prompt": "private usage metadata", + } + ), + capture_content=False, + ) + assert normalized.inputs is None + assert normalized.outputs is None + assert "private" not in normalized.model_dump_json() + assert "sensitive error" not in normalized.model_dump_json() + assert normalized.error == "Agent operation failed" + assert normalized.attributes["dify.app_id"] == "app-one" + assert normalized.model_name == "model-one" + assert normalized.usage == TokenUsage(prompt_tokens=3, completion_tokens=2, total_tokens=5) + + +def test_json_encoded_tool_credentials_are_redacted_from_inputs_and_metadata() -> None: + normalized = normalize_agent_span( + agent_data( + { + "gen_ai.operation.name": "execute_tool", + "gen_ai.tool.name": "search", + "gen_ai.tool.call.arguments": json.dumps({"api_key": "tool-secret-123", "query": "public query"}), + "gen_ai.tool.call.result": json.dumps( + {"nested": {"authorization": "bearer-secret-123"}, "result": "okay"} + ), + } + ), + capture_content=True, + ) + serialized = normalized.model_dump_json() + assert "tool-secret-123" not in serialized + assert "bearer-secret-123" not in serialized + assert "public query" in serialized + assert normalized.kind is SpanKind.TOOL + assert normalized.inputs == {"api_key": "[redacted]", "query": "public query"} + assert normalized.outputs == {"nested": {"authorization": "[redacted]"}, "result": "okay"} + + +def test_content_off_drops_error_event_payload_and_keeps_numeric_usage() -> None: + attrs = safe_attributes( + { + "input": "user text", + "authorization": "secret", + "exception.stacktrace": "sensitive", + "gen_ai.usage.input_tokens": 10, + "gen_ai.tool.name": "lookup", + }, + capture_content=False, + ) + assert attrs == {"gen_ai.usage.input_tokens": 10, "gen_ai.tool.name": "lookup"} + + +def test_real_runtime_accounting_attributes_survive_content_off() -> None: + normalized = normalize_agent_span( + agent_data( + { + "gen_ai.operation.name": "chat", + "gen_ai.request.model": "gpt-4o-mini", + "gen_ai.usage.input_tokens": 8, + "gen_ai.usage.output_tokens": 4, + # The runtime reports its decimal price as a string. + "dify.llm.total_price": "0.03", + "dify.llm.currency": "USD", + "gen_ai.response.time_to_first_token": 0.2, + "dify.llm.time_to_first_token_unit": "s", + "dify.llm.time_to_first_token_source": "gateway_usage", + "gen_ai.input.messages": "private input", + } + ), + capture_content=False, + ) + assert normalized.cost == Cost(amount=0.03, currency="USD") + assert normalized.time_to_first_token == 0.2 + assert normalized.usage == TokenUsage(prompt_tokens=8, completion_tokens=4, total_tokens=12) + assert normalized.attributes["dify.llm.time_to_first_token_source"] == "gateway_usage" + assert normalized.inputs is None + + +@pytest.mark.parametrize("price", ["not-a-price", "-0.01", "nan", "inf", True]) +def test_unusable_runtime_price_is_not_reported_as_cost(price: object) -> None: + normalized = normalize_agent_span( + agent_data({"gen_ai.operation.name": "chat", "dify.llm.total_price": price}), capture_content=True + ) + assert normalized.kind is SpanKind.LLM + assert normalized.cost is None + + +def test_agent_aggregate_usage_is_context_not_another_generation() -> None: + normalized = normalize_agent_span( + agent_data( + { + "gen_ai.operation.name": "invoke_agent", + "gen_ai.request.model": "model", + "gen_ai.usage.input_tokens": 30, + "gen_ai.usage.output_tokens": 20, + "dify.llm.total_price": "0.5", + } + ), + capture_content=False, + ) + assert normalized.kind is SpanKind.AGENT + assert normalized.usage is None + assert normalized.cost is None + assert normalized.model_name is None + assert normalized.attributes["aggregate_prompt_tokens"] == 30 + assert normalized.attributes["aggregate_completion_tokens"] == 20 + + +def span(name: str, parent: str | None, kind: SpanKind) -> Span: + return Span( + span_id=span_id_for(name), + parent_span_id=span_id_for(parent) if parent else None, + name=name, + kind=kind, + start_time=datetime.fromtimestamp(0, UTC), + end_time=datetime.fromtimestamp(3, UTC), + status=SpanStatus.OK, + inputs={"content": "private"}, + outputs="private", + usage=TokenUsage(prompt_tokens=6, completion_tokens=4, total_tokens=10) if kind is SpanKind.LLM else None, + cost=Cost(amount=0.03, currency="USD") if kind is SpanKind.LLM else None, + attributes={"prompt": "private", "node_type": kind.value}, + ) + + +def trace_of(*spans: Span) -> Trace: + return Trace( + trace_id=span_id_for("old") * 2, + external_trace_id=None, + session_id="session", + user_id="user", + parent=None, + spans=spans, + ) + + +def test_two_native_model_calls_replace_synthetic_aggregate_generation() -> None: + outer = trace_of(span("message", None, SpanKind.CHAIN), span("message:llm", "message", SpanKind.LLM)) + remapped = remap_outer_trace( + outer, trace_id="1" * 32, capture_content=False, synthetic_span_ids=frozenset({span_id_for("message:llm")}) + ) + assert [item.span_id for item in remapped.spans] == [span_id_for("message")] + assert remapped.trace_id == "1" * 32 + assert (remapped.session_id, remapped.user_id) == ("session", "user") + calls = [ + normalize_agent_span( + agent_data( + { + "gen_ai.operation.name": "chat", + "gen_ai.request.model": "model", + "gen_ai.usage.input_tokens": 3, + "gen_ai.usage.output_tokens": 2, + }, + span_id=identifier, + parent_id=remapped.root.span_id, + ), + capture_content=False, + ) + for identifier in ("2" * 16, "3" * 16) + ] + combined = order_spans(remapped.root.span_id, [calls[1], *remapped.spans, calls[0]]) + assert combined[0].span_id == remapped.root.span_id + generations = [item for item in combined if item.kind is SpanKind.LLM] + assert len(generations) == 2 + assert sum(item.usage.total_tokens for item in generations if item.usage is not None) == 10 + assert combined[0].usage is None + assert "private" not in json.dumps([item.model_dump(mode="json") for item in combined]) + + +def test_non_agent_workflow_retains_its_real_llm_node() -> None: + trace = trace_of(span("workflow", None, SpanKind.CHAIN), span("llm-node", "workflow", SpanKind.LLM)) + remapped = remap_outer_trace(trace, trace_id="1" * 32, capture_content=True, synthetic_span_ids=frozenset()) + assert len(remapped.spans) == 2 + llm = remapped.spans[1] + assert llm.kind is SpanKind.LLM + assert llm.usage == TokenUsage(prompt_tokens=6, completion_tokens=4, total_tokens=10) + assert llm.cost == Cost(amount=0.03, currency="USD") + assert llm.inputs == {"content": "private"} + assert llm.attributes["prompt"] == "private" + + +def test_content_off_outer_trace_keeps_only_identifiers_and_masks_errors() -> None: + failed = span("workflow", None, SpanKind.CHAIN).model_copy( + update={"status": SpanStatus.ERROR, "error": "private failure"} + ) + remapped = remap_outer_trace( + trace_of(failed), trace_id="1" * 32, capture_content=False, synthetic_span_ids=frozenset() + ) + root = remapped.root + assert root.error == "Application execution failed" + assert root.status is SpanStatus.ERROR + assert (root.inputs, root.outputs) == (None, None) + assert root.attributes == {"node_type": "chain"} + + +def test_order_spans_rejects_orphans_cycles_and_conflicting_ids() -> None: + root = span("root", None, SpanKind.CHAIN) + with pytest.raises(ValueError, match="parent is missing"): + order_spans(root.span_id, [root, span("orphan", "missing", SpanKind.TOOL)]) + with pytest.raises(ValueError, match="parent cycle"): + order_spans(root.span_id, [root, span("one", "two", SpanKind.TOOL), span("two", "one", SpanKind.TOOL)]) + with pytest.raises(ValueError, match="Conflicting"): + order_spans(root.span_id, [root, root.model_copy(update={"name": "other"})]) + with pytest.raises(ValueError, match="root is missing"): + order_spans(root.span_id, [span("one", "root", SpanKind.TOOL)]) + with pytest.raises(ValueError, match="Unexpected trace root"): + order_spans(root.span_id, [root, span("second-root", None, SpanKind.CHAIN)]) + + +def test_managed_external_parent_is_preserved_as_explicit_link_when_content_off() -> None: + root = span("workflow", None, SpanKind.CHAIN).model_copy( + update={ + "attributes": { + "parent_workflow_run_id": "outer-workflow", + "parent_node_execution_id": "outer-tool", + "prompt": "private", + } + } + ) + outer = Trace( + trace_id="2" * 32, + external_trace_id="caller-trace", + session_id="session", + user_id=None, + parent=None, + spans=(root,), + ) + remapped = remap_outer_trace(outer, trace_id="1" * 32, capture_content=False, synthetic_span_ids=frozenset()) + assert remapped.parent is None + assert remapped.external_trace_id is None + assert remapped.root.attributes == { + "parent_workflow_run_id": "outer-workflow", + "parent_node_execution_id": "outer-tool", + "external_trace_id": "caller-trace", + } + + +TRACE_ID = "0af7651916cd43dd8448eb211c80319c" + + +def _legacy(**span_overrides: object) -> dict[str, object]: + root, llm = span_id_for("message-1"), span_id_for("message-1:llm") + return { + "trace_id": TRACE_ID, + "session_id": "conversation-1", + "root_span_id": root, + "spans": [ + { + "id": root, + "parent_id": None, + "name": "message", + "kind": "chain", + "start_time": "2026-01-01T08:00:00", + "end_time": "2026-01-01T08:00:02", + "inputs": {"query": "hi"}, + "outputs": "hello", + "status": "ok", + "error": None, + "metadata": {"aggregate_total_tokens": 20}, + "synthetic": False, + "can_parent_workflow": False, + "publishes_parent_context": False, + "links": [], + }, + { + "id": llm, + "parent_id": root, + "name": "llm", + "kind": "llm", + "start_time": "2026-01-01T08:00:00", + "end_time": None, + "inputs": None, + "outputs": None, + "status": "ok", + "error": None, + "metadata": { + "model_name": "gpt-4o", + "model_provider": "openai", + "prompt_tokens": 12, + "completion_tokens": 8, + "total_tokens": 20, + "total_price": "0.0042", + "currency": "USD", + }, + "synthetic": False, + "can_parent_workflow": False, + "publishes_parent_context": False, + "links": [], + **span_overrides, + }, + ], + "external_parent": None, + "required_parent_context_id": None, + } + + +def test_outer_traces_stored_by_the_previous_release_still_export() -> None: + trace = legacy_outer_trace(_legacy()) + + assert (trace.trace_id, trace.session_id, trace.parent) == (TRACE_ID, "conversation-1", None) + root, llm = trace.spans + assert root.attributes == {"aggregate_total_tokens": 20} + assert llm.kind is SpanKind.LLM + assert (llm.model_name, llm.model_provider) == ("gpt-4o", "openai") + assert llm.usage is not None + assert (llm.usage.prompt_tokens, llm.usage.total_tokens) == (12, 20) + assert llm.cost is not None + assert (llm.cost.amount, llm.cost.currency) == (0.0042, "USD") + assert llm.end_time == datetime(2026, 1, 1, 8, 0, tzinfo=UTC) + + +def test_malformed_legacy_traces_still_fail_validation() -> None: + with pytest.raises(ValidationError, match="must follow its parent"): + legacy_outer_trace(_legacy(parent_id="0000000000000abc")) diff --git a/api/tests/unit_tests/domain/tracing/test_terminal_recovery.py b/api/tests/unit_tests/domain/tracing/test_terminal_recovery.py new file mode 100644 index 00000000000000..4f0674496c8c6e --- /dev/null +++ b/api/tests/unit_tests/domain/tracing/test_terminal_recovery.py @@ -0,0 +1,155 @@ +from dataclasses import replace +from datetime import datetime, timedelta +from uuid import NAMESPACE_URL, UUID, uuid4, uuid5 + +import pytest +from dify_agent.protocol.tracing import AgentTraceBatch, AgentTraceSpan + +from domain.tracing.errors import TraceRunError +from domain.tracing.sources import MessageProgress, SourceWorkflowRun +from domain.tracing.terminal_recovery import message_terminal_due, workflow_terminal_due +from domain.tracing.trace_run_agg import TraceRunAgg + +COMMITTED = datetime(2026, 10, 1, 12, 0, 0) +TIMEOUT = 60 +POLICY: dict[str, object] = {"enabled": True, "tracing_provider": "langfuse", "capture_agent": True} + + +def identifier(value: str) -> str: + return str(uuid5(NAMESPACE_URL, value)) + + +def collecting(policy: dict[str, object]) -> TraceRunAgg: + return TraceRunAgg.freeze( + run_id=str(uuid4()), + trace_id=uuid4().hex, + tenant_id=identifier("tenant"), + app_id=identifier("app"), + source_kind="message", + source_id=identifier("message"), + session_id="session", + policy=policy, + stored_config={"host": "https://langfuse.test"}, + now=COMMITTED, + ) + + +def with_invocation(run: TraceRunAgg, *, finalized_at: datetime | None) -> TraceRunAgg: + """The run with one Agent invocation, finalized at ``finalized_at`` unless that is ``None``.""" + opened = run.open_invocation(agent_id=identifier("agent"), node_execution_id=None, attempt=0, call_key="k") + if finalized_at is not None: + payload = AgentTraceBatch( + run_id=UUID(run.id), + invocation_id=UUID(opened.id), + batch_id=uuid4(), + final=True, + spans=[ + AgentTraceSpan( + trace_id=run.trace_id, + span_id="a" * 16, + parent_span_id=opened.span_id, + name="Agent execution", + start_time_ns=1_000_000_000, + end_time_ns=2_000_000_000, + status="ok", + ) + ], + ) + run.ingest(payload, max_spans=100, max_bytes=1_000_000, now=finalized_at) + return run + + +PROGRESS = MessageProgress( + message_id=identifier("message"), + conversation_id=identifier("conversation"), + agent_app=True, + paused=False, + explicit_error=False, + answer_saved=True, + agent_stop_requested=False, + updated_at=COMMITTED, +) + + +def due(run: TraceRunAgg, progress: MessageProgress, after: timedelta) -> bool: + return message_terminal_due(run, progress, now=COMMITTED + after, collection_timeout=TIMEOUT) + + +class TestMessageTerminal: + @staticmethod + def test_saved_answer_is_the_terminal_once_it_settles() -> None: + run = with_invocation(collecting(POLICY), finalized_at=COMMITTED) + + assert due(run, PROGRESS, timedelta(seconds=4)) is False + assert due(run, PROGRESS, timedelta(seconds=5)) is True + assert due(run, replace(PROGRESS, answer_saved=False, explicit_error=True), timedelta(seconds=5)) is True + + @staticmethod + def test_paused_message_waits_unless_an_agent_stop_is_pending() -> None: + run = with_invocation(collecting(POLICY), finalized_at=COMMITTED) + paused = replace(PROGRESS, paused=True) + + assert due(run, paused, timedelta(minutes=5)) is False + with pytest.raises(TraceRunError, match="business_terminal_not_persisted"): + due(run, replace(paused, agent_stop_requested=True), timedelta(seconds=TIMEOUT + 1)) + + @staticmethod + def test_unfinished_invocation_waits_until_a_failed_message_times_out() -> None: + run = with_invocation(collecting(POLICY), finalized_at=None) + + assert due(run, PROGRESS, timedelta(minutes=5)) is False + assert due(run, replace(PROGRESS, explicit_error=True), timedelta(seconds=TIMEOUT)) is False + with pytest.raises(TraceRunError, match="agent_trace_collection_incomplete"): + due(run, replace(PROGRESS, explicit_error=True), timedelta(seconds=TIMEOUT + 1)) + + @staticmethod + def test_pending_agent_stop_waits_for_the_worker_handoff_until_its_deadline() -> None: + run = with_invocation(collecting(POLICY), finalized_at=COMMITTED + timedelta(seconds=10)) + stopping = replace(PROGRESS, agent_stop_requested=True) + + assert due(run, stopping, timedelta(seconds=10 + TIMEOUT)) is False + with pytest.raises(TraceRunError, match="business_terminal_not_persisted"): + due(run, stopping, timedelta(seconds=10 + TIMEOUT + 1)) + + @staticmethod + def test_agent_capture_waits_for_its_first_invocation() -> None: + assert due(collecting(POLICY), PROGRESS, timedelta(minutes=5)) is False + assert due(collecting({**POLICY, "capture_agent": False}), PROGRESS, timedelta(minutes=5)) is True + + @staticmethod + def test_usage_flush_without_the_final_answer_never_becomes_a_terminal() -> None: + run = with_invocation(collecting(POLICY), finalized_at=COMMITTED) + unsaved = replace(PROGRESS, answer_saved=False) + + assert due(run, unsaved, timedelta(seconds=TIMEOUT)) is False + with pytest.raises(TraceRunError, match="business_terminal_not_persisted"): + due(run, unsaved, timedelta(seconds=TIMEOUT + 1)) + + +class TestWorkflowTerminal: + @staticmethod + @pytest.mark.parametrize( + ("status", "finished_offset", "settled"), + [ + ("running", None, False), + ("succeeded", None, False), + ("succeeded", timedelta(seconds=-4), False), + ("succeeded", timedelta(seconds=-5), True), + ("partial-succeeded", timedelta(seconds=-5), True), + ("failed", timedelta(seconds=-5), True), + ("stopped", timedelta(seconds=-5), True), + ], + ) + def test_workflow_run_is_the_terminal_once_it_finished_and_settled( + status: str, finished_offset: timedelta | None, settled: bool + ) -> None: + workflow = SourceWorkflowRun( + id=identifier("workflow-run"), + workflow_id=identifier("workflow"), + graph=None, + status=status, + finished_at=COMMITTED + finished_offset if finished_offset is not None else None, + created_by=identifier("account"), + ) + + assert workflow_terminal_due(workflow, now=COMMITTED) is settled diff --git a/api/tests/unit_tests/domain/tracing/test_trace_run_agg.py b/api/tests/unit_tests/domain/tracing/test_trace_run_agg.py new file mode 100644 index 00000000000000..853502fbb2afe1 --- /dev/null +++ b/api/tests/unit_tests/domain/tracing/test_trace_run_agg.py @@ -0,0 +1,778 @@ +from dataclasses import replace +from datetime import UTC, datetime, timedelta +from uuid import NAMESPACE_URL, UUID, uuid4, uuid5 + +import pytest +from dify_agent.protocol.tracing import AgentTraceBatch, AgentTraceEvent, AgentTraceSpan +from dify_trace_sdk.ids import span_id_for +from dify_trace_sdk.model import Span, SpanKind, SpanStatus, Trace + +from domain.tracing.errors import TraceRunConflictError, TraceRunError, TraceRunLimitError, TraceRunNotFoundError +from domain.tracing.sources import SourceNodeExecution +from domain.tracing.trace_run_agg import TraceBatch, TraceInvocation, TraceRunAgg, TraceRunState, TraceSourceKind + +NOW = datetime(2026, 10, 1, 12, 0, 0) +POLICY: dict[str, object] = { + "enabled": True, + "tracing_provider": "langfuse", + "capture_content": False, + "capture_agent": True, +} +CONFIG: dict[str, object] = {"public_key": "enc:public", "secret_key": "enc:secret", "host": "https://langfuse.test"} + + +def identifier(value: str) -> str: + return str(uuid5(NAMESPACE_URL, value)) + + +def frozen( + *, + source_kind: TraceSourceKind = "message", + policy: dict[str, object] = POLICY, + config: dict[str, object] | None = CONFIG, +) -> TraceRunAgg: + return TraceRunAgg.freeze( + run_id=str(uuid4()), + trace_id=uuid4().hex, + tenant_id=identifier("tenant"), + app_id=identifier("app"), + source_kind=source_kind, + source_id=identifier("source"), + session_id="session", + policy=policy, + stored_config=config, + now=NOW, + ) + + +def invocation(run: TraceRunAgg, call_key: str = "first") -> TraceInvocation: + return run.open_invocation(agent_id=identifier("agent"), node_execution_id=None, attempt=0, call_key=call_key) + + +def batch( + run: TraceRunAgg, + opened: TraceInvocation, + *, + final: bool = True, + batch_id: UUID | None = None, + name: str = "Agent execution", + status: str = "ok", +) -> AgentTraceBatch: + """The validated runtime batch, as the service hands it to the run.""" + return AgentTraceBatch( + run_id=UUID(run.id), + invocation_id=UUID(opened.id), + batch_id=batch_id or uuid4(), + final=final, + spans=[ + AgentTraceSpan( + trace_id=run.trace_id, + span_id="a" * 16, + parent_span_id=opened.span_id, + name=name, + start_time_ns=1_000_000_000, + end_time_ns=2_000_000_000, + status="error" if status == "error" else "ok", + status_message="provider said no" if status == "error" else None, + attributes={"gen_ai.operation.name": "invoke_agent", "gen_ai.input.messages": "private prompt"}, + events=[AgentTraceEvent(name="token", timestamp_ns=1_500_000_000)], + ) + ], + ) + + +def ingest(run: TraceRunAgg, payload: AgentTraceBatch, *, now: datetime = NOW) -> int: + return run.ingest(payload, max_spans=100, max_bytes=1_000_000, now=now) + + +def business_terminal(run: TraceRunAgg) -> Trace: + """A message's canonical trace: its root span and the LLM span the Agent runtime replaces.""" + started = datetime.fromtimestamp(0, UTC) + return Trace( + trace_id="f" * 32, + external_trace_id="caller-trace", + session_id="conversation", + user_id="user", + parent=None, + spans=( + Span( + span_id=span_id_for(run.source_id), + parent_span_id=None, + name="message", + kind=SpanKind.CHAIN, + start_time=started, + end_time=started + timedelta(seconds=3), + status=SpanStatus.ERROR, + error="private failure", + inputs={"query": "private question"}, + outputs="private answer", + attributes={"app_id": run.app_id, "private": "value"}, + ), + Span( + span_id=span_id_for(f"{run.source_id}:llm"), + parent_span_id=span_id_for(run.source_id), + name="llm", + kind=SpanKind.LLM, + start_time=started, + end_time=started + timedelta(seconds=3), + status=SpanStatus.OK, + ), + ), + ) + + +def exportable(run: TraceRunAgg) -> TraceRunAgg: + """Collect one final batch and the business terminal, so the run can be claimed.""" + ingest(run, batch(run, invocation(run))) + run.capture_terminal(business_terminal(run), now=NOW) + return run + + +def claimed(run: TraceRunAgg, *, now: datetime = NOW) -> str: + token = str(uuid4()) + assert run.claim_export(token=token, owner_exists=True, collection_timeout=60, lease_seconds=30, now=now) + return token + + +def not_claimed(run: TraceRunAgg, *, owner_exists: bool = True, now: datetime = NOW) -> bool: + return not run.claim_export( + token=str(uuid4()), owner_exists=owner_exists, collection_timeout=60, lease_seconds=30, now=now + ) + + +class TestFreezing: + @staticmethod + def test_collecting_run_freezes_the_destination_and_capture_policy() -> None: + run = frozen(policy={**POLICY, "capture_content": True}) + + assert run.state is TraceRunState.COLLECTING + assert run.provider == "langfuse" + assert run.provider_config == CONFIG + assert run.capture_content is True + assert run.capture_agent is True + assert run.root_span_id == span_id_for(run.source_id) + assert run.updated_at == NOW + + @staticmethod + @pytest.mark.parametrize( + ("policy", "config"), + [ + ({**POLICY, "enabled": False}, CONFIG), + ({**POLICY, "enabled": "true"}, CONFIG), + ({**POLICY, "tracing_provider": None}, CONFIG), + (POLICY, None), + ({}, None), + ], + ) + def test_run_without_an_enabled_stored_destination_is_not_collected( + policy: dict[str, object], config: dict[str, object] | None + ) -> None: + run = frozen(policy=policy, config=config) + + assert run.state is TraceRunState.NOT_COLLECTED + assert run.provider == "" + assert run.provider_config == {} + assert run.capture_content is False + assert run.capture_agent is False + + @staticmethod + def test_revision_names_the_frozen_policy_and_configuration() -> None: + assert frozen().config_revision == frozen().config_revision + assert frozen().config_revision != frozen(config={**CONFIG, "host": "https://other.test"}).config_revision + assert frozen().config_revision != frozen(policy={**POLICY, "capture_content": True}).config_revision + + @staticmethod + def test_session_is_bounded_and_falls_back_to_the_source() -> None: + long = TraceRunAgg.freeze( + run_id=str(uuid4()), + trace_id=uuid4().hex, + tenant_id=identifier("tenant"), + app_id=identifier("app"), + source_kind="message", + source_id=identifier("source"), + session_id="s" * 300, + policy=POLICY, + stored_config=CONFIG, + now=NOW, + ) + empty = TraceRunAgg.freeze( + run_id=str(uuid4()), + trace_id=uuid4().hex, + tenant_id=identifier("tenant"), + app_id=identifier("app"), + source_kind="message", + source_id=identifier("source"), + session_id="", + policy=POLICY, + stored_config=CONFIG, + now=NOW, + ) + + assert long.session_id == "s" * 255 + assert empty.session_id == identifier("source") + + @staticmethod + def test_initialization_failure_closes_the_source_without_a_destination() -> None: + run = TraceRunAgg.initialization_failed( + run_id=str(uuid4()), + trace_id=uuid4().hex, + tenant_id=identifier("tenant"), + app_id=identifier("app"), + source_kind="workflow", + source_id=identifier("source"), + now=NOW, + ) + + assert run.state is TraceRunState.FAILED + assert run.error_code == "tracing_initialization_failed" + assert run.terminal_at == NOW + assert (run.provider, run.config_revision, run.provider_config) == ("", "", {}) + assert not run.capture_content + assert not run.capture_agent + assert not_claimed(run) + + +class TestInvocations: + @staticmethod + def test_call_key_names_one_invocation_under_its_parent_span() -> None: + run = frozen() + app_invocation = invocation(run) + node_invocation = run.open_invocation( + agent_id=identifier("agent"), node_execution_id=identifier("node"), attempt=1, call_key="node:1" + ) + + assert invocation(run) == app_invocation + assert app_invocation.id == str(uuid5(UUID(run.id), "first")) + assert app_invocation.parent_span_id == run.root_span_id + assert node_invocation.parent_span_id == span_id_for(identifier("node")) + assert set(run.invocations) == {app_invocation.id, node_invocation.id} + + @staticmethod + def test_repeated_call_key_keeps_its_agent_binding() -> None: + run = frozen() + invocation(run) + + with pytest.raises(TraceRunConflictError, match="tracing_invocation_conflict"): + run.open_invocation(agent_id=identifier("other-agent"), node_execution_id=None, attempt=0, call_key="first") + + @staticmethod + def test_invocations_open_only_while_the_run_collects() -> None: + not_collected = frozen(config=None) + exporting = exportable(frozen()) + claimed(exporting) + + for run in (not_collected, exporting): + with pytest.raises(TraceRunConflictError, match="tracing_run_not_collecting"): + invocation(run, "late") + + +class TestIngestion: + @staticmethod + def test_final_batch_is_stored_and_finalizes_its_invocation() -> None: + run = frozen() + opened = invocation(run) + payload = batch(run, opened) + + assert ingest(run, payload) == 1 + + (stored,) = run.new_batches + assert (stored.id, stored.invocation_id, stored.final) == (str(payload.batch_id), opened.id, True) + assert run.receipts[stored.id].payload_hash == stored.payload_hash + assert run.span_count == 1 + assert run.payload_bytes > 0 + assert run.invocations[opened.id].finalized is True + assert run.invocations[opened.id].finalized_at == NOW + + @staticmethod + def test_replayed_batch_is_accepted_again_after_the_run_closed() -> None: + run = frozen() + opened = invocation(run) + payload = batch(run, opened) + ingest(run, payload) + run.capture_terminal(business_terminal(run), now=NOW) + run.export_accepted(claimed(run), provider_trace_id="remote", trace_url=None, span_urls={}, now=NOW) + + assert ingest(run, payload) == 1 + assert len(run.new_batches) == 1 + assert run.span_count == 1 + + @staticmethod + def test_batch_id_replays_only_with_the_same_payload_and_invocation() -> None: + run = frozen() + opened = invocation(run) + other = invocation(run, "second") + payload = batch(run, opened, final=False) + ingest(run, payload) + batch_id = payload.batch_id + + with pytest.raises(TraceRunConflictError, match="tracing_batch_conflict"): + ingest(run, batch(run, opened, final=False, batch_id=batch_id, name="changed")) + with pytest.raises(TraceRunConflictError, match="tracing_batch_conflict"): + ingest(run, batch(run, other, final=False, batch_id=batch_id)) + assert run.span_count == 1 + + @staticmethod + def test_batch_of_an_unknown_invocation_is_not_found() -> None: + run = frozen() + foreign = invocation(frozen(), "first") + + with pytest.raises(TraceRunNotFoundError, match="tracing_invocation_not_found"): + ingest(run, batch(run, foreign)) + + @staticmethod + def test_finalized_invocation_and_closed_run_accept_no_new_batch() -> None: + run = frozen() + opened = invocation(run) + ingest(run, batch(run, opened)) + + with pytest.raises(TraceRunConflictError, match="tracing_collection_closed"): + ingest(run, batch(run, opened)) + + closed = frozen() + late = invocation(closed) + closed.capture_terminal(business_terminal(closed), now=NOW) + after_timeout = NOW + timedelta(seconds=61) + assert not_claimed(closed, now=after_timeout) + assert closed.state is TraceRunState.FAILED + with pytest.raises(TraceRunConflictError, match="tracing_collection_closed"): + ingest(closed, batch(closed, late), now=after_timeout) + + @staticmethod + def test_spans_must_belong_to_the_run_trace() -> None: + run = frozen() + opened = invocation(run) + foreign = AgentTraceBatch( + run_id=UUID(run.id), + invocation_id=UUID(opened.id), + batch_id=uuid4(), + final=True, + spans=[ + AgentTraceSpan( + trace_id="e" * 32, + span_id="a" * 16, + parent_span_id=opened.span_id, + name="Agent execution", + start_time_ns=1_000_000_000, + end_time_ns=2_000_000_000, + status="ok", + ) + ], + ) + + with pytest.raises(TraceRunConflictError, match="tracing_trace_id_mismatch"): + ingest(run, foreign) + assert run.new_batches == [] + + @staticmethod + @pytest.mark.parametrize(("max_spans", "max_bytes"), [(0, 1_000_000), (100, 10)]) + def test_limits_reject_the_whole_batch(max_spans: int, max_bytes: int) -> None: + run = frozen() + opened = invocation(run) + + with pytest.raises(TraceRunLimitError, match="tracing_payload_limit"): + run.ingest(batch(run, opened), max_spans=max_spans, max_bytes=max_bytes, now=NOW) + + assert run.new_batches == [] + assert (run.span_count, run.payload_bytes) == (0, 0) + assert run.invocations[opened.id].finalized is False + + @staticmethod + def test_content_off_stores_only_allowlisted_span_data() -> None: + hidden = frozen() + ingest(hidden, batch(hidden, invocation(hidden), status="error")) + shown = frozen(policy={**POLICY, "capture_content": True}) + ingest(shown, batch(shown, invocation(shown), status="error")) + + (hidden_span,) = hidden.new_batches[0].spans + (shown_span,) = shown.new_batches[0].spans + assert hidden_span["attributes"] == {"gen_ai.operation.name": "invoke_agent"} + assert hidden_span["events"] == [] + assert hidden_span["status_message"] == "Agent operation failed" + assert shown_span["attributes"]["gen_ai.input.messages"] == "private prompt" + assert len(shown_span["events"]) == 1 + assert shown_span["status_message"] == "provider said no" + + +class TestTerminal: + @staticmethod + def test_terminal_adopts_the_run_identity_and_capture_policy() -> None: + run = frozen() + + run.capture_terminal(business_terminal(run), now=NOW) + + assert run.outer_trace is not None + kept = Trace.model_validate(run.outer_trace) + assert kept.trace_id == run.trace_id + assert kept.session_id == "session" + assert [span.name for span in kept.spans] == ["message"] + (root,) = kept.spans + assert (root.inputs, root.outputs) == (None, None) + assert root.error == "Application execution failed" + assert root.attributes == {"app_id": run.app_id, "external_trace_id": "caller-trace"} + assert run.root_span_id == root.span_id + assert run.terminal_at == NOW + + @staticmethod + def test_first_terminal_wins() -> None: + run = frozen() + run.capture_terminal(business_terminal(run), now=NOW) + first = run.outer_trace + + replacement = business_terminal(run).model_copy(update={"user_id": "someone else"}) + run.capture_terminal(replacement, now=NOW + timedelta(seconds=5)) + + assert run.outer_trace == first + assert run.terminal_at == NOW + + @staticmethod + def test_closed_and_uncollected_runs_keep_no_terminal() -> None: + uncollected = frozen(config=None) + uncollected.capture_terminal(business_terminal(uncollected), now=NOW) + failed = TraceRunAgg.initialization_failed( + run_id=str(uuid4()), + trace_id=uuid4().hex, + tenant_id=identifier("tenant"), + app_id=identifier("app"), + source_kind="message", + source_id=identifier("source"), + now=NOW, + ) + failed.capture_terminal(business_terminal(failed), now=NOW) + + assert uncollected.outer_trace is None + assert failed.outer_trace is None + + @staticmethod + @pytest.mark.parametrize("capture_content", [True, False]) + def test_finished_preview_node_becomes_the_terminal(capture_content: bool) -> None: + run = frozen(source_kind="node", policy={**POLICY, "capture_content": capture_content}) + node = SourceNodeExecution( + id=run.source_id, + node_id="agent", + workflow_id=None, + workflow_run_id=None, + title="Agent preview", + status="running", + error="private failure", + inputs='{"query": "private"}', + outputs=None, + created_at=NOW, + elapsed_time=2, + ) + + run.settle_preview(node, now=NOW) + assert run.outer_trace is None + + run.settle_preview( + SourceNodeExecution( + id=node.id, + node_id=node.node_id, + workflow_id=None, + workflow_run_id=None, + title=node.title, + status="failed", + error=node.error, + inputs=node.inputs, + outputs=None, + created_at=NOW, + elapsed_time=2, + ), + now=NOW, + ) + assert run.outer_trace is not None + (span,) = Trace.model_validate(run.outer_trace).spans + assert span.span_id == run.root_span_id + assert span.end_time - span.start_time == timedelta(seconds=2) + assert span.status is SpanStatus.ERROR + assert span.error == ("private failure" if capture_content else "Node execution failed") + assert span.inputs == ({"query": "private"} if capture_content else None) + + +class TestExport: + @staticmethod + def test_claim_waits_for_the_terminal_and_every_final_invocation() -> None: + run = frozen() + opened = invocation(run) + ingest(run, batch(run, opened, final=False)) + assert not_claimed(run) + + run.capture_terminal(business_terminal(run), now=NOW) + assert not_claimed(run) + + ingest(run, batch(run, opened)) + token = claimed(run) + assert (run.state, run.export_token, run.attempts) == (TraceRunState.EXPORTING, token, 1) + assert not_claimed(run) + + @staticmethod + def test_incomplete_collection_fails_after_its_timeout() -> None: + run = frozen() + invocation(run) + run.capture_terminal(business_terminal(run), now=NOW) + + later = NOW + timedelta(seconds=61) + assert not_claimed(run, now=later) + assert (run.state, run.error_code) == (TraceRunState.FAILED, "agent_trace_collection_incomplete") + + @staticmethod + def test_run_whose_app_is_gone_fails_instead_of_exporting() -> None: + run = exportable(frozen()) + + assert not_claimed(run, owner_exists=False) + assert (run.state, run.error_code, run.attempts) == (TraceRunState.FAILED, "tracing_owner_not_found", 0) + + @staticmethod + def test_lease_renews_only_for_the_current_attempt() -> None: + run = exportable(frozen()) + token = claimed(run) + later = NOW + timedelta(seconds=20) + + assert run.renew_export_lease("other", now=later) is False + assert run.renew_export_lease(token, now=later) is True + assert run.export_heartbeat_at == later + + @staticmethod + def test_expired_lease_becomes_uncertain_and_still_takes_the_late_acceptance() -> None: + run = exportable(frozen()) + token = claimed(run) + + assert not_claimed(run, now=NOW + timedelta(seconds=31)) + assert (run.state, run.error_code) == (TraceRunState.UNCERTAIN, "export_lease_expired") + + run.export_failed(token, "provider_export_rejected", now=NOW) + assert run.state is TraceRunState.UNCERTAIN + trace_url = "https://langfuse.test/trace" + run.export_accepted(token, provider_trace_id="remote", trace_url=trace_url, span_urls={"a" * 16: "u"}, now=NOW) + assert (run.state, run.error_code, run.export_token) == (TraceRunState.ACCEPTED, None, None) + assert (run.provider_trace_id, run.trace_url, run.span_urls) == ("remote", trace_url, {"a" * 16: "u"}) + assert run.accepted_at == NOW + + @staticmethod + def test_retryable_failure_backs_off_until_the_attempts_run_out() -> None: + run = exportable(frozen()) + run.export_deferred(claimed(run), max_attempts=2, now=NOW) + + assert (run.state, run.error_code) == (TraceRunState.RETRYING, "provider_temporarily_unavailable") + assert run.next_attempt_at == NOW + timedelta(seconds=5) + assert not_claimed(run) + + token = claimed(run, now=NOW + timedelta(seconds=5)) + run.export_deferred(token, max_attempts=2, now=NOW + timedelta(seconds=5)) + assert (run.state, run.error_code, run.attempts) == (TraceRunState.FAILED, "provider_retry_exhausted", 2) + + @staticmethod + def test_outcomes_of_a_replaced_attempt_change_nothing() -> None: + run = exportable(frozen()) + claimed(run) + + run.export_failed("replaced", "provider_export_rejected", now=NOW) + run.export_uncertain("replaced", now=NOW) + run.export_deferred("replaced", max_attempts=8, now=NOW) + run.export_accepted("replaced", provider_trace_id="remote", trace_url=None, span_urls={}, now=NOW) + + assert run.state is TraceRunState.EXPORTING + assert run.provider_trace_id is None + + +class TestRetry: + @staticmethod + def test_uncertain_export_retries_only_with_the_duplicate_acknowledgement() -> None: + run = exportable(frozen()) + run.export_uncertain(claimed(run), now=NOW) + assert (run.state, run.error_code) == (TraceRunState.UNCERTAIN, "provider_acceptance_unknown") + + with pytest.raises(TraceRunConflictError, match="tracing_duplicate_acknowledgement_required"): + run.release_for_retry(acknowledge_possible_duplicates=False, now=NOW) + run.release_for_retry(acknowledge_possible_duplicates=True, now=NOW) + + assert (run.state, run.error_code, run.export_token, run.next_attempt_at) == ( + TraceRunState.READY, + None, + None, + None, + ) + assert claimed(run) is not None + + @staticmethod + def test_ineligible_runs_are_not_released() -> None: + accepted = exportable(frozen()) + accepted.export_accepted(claimed(accepted), provider_trace_id="r", trace_url=None, span_urls={}, now=NOW) + rejected_by_collection = exportable(frozen()) + rejected_by_collection.export_failed(claimed(rejected_by_collection), "tracing_payload_invalid", now=NOW) + collecting = frozen() + + cases: list[tuple[TraceRunAgg, type[TraceRunError], str]] = [ + (frozen(config=None), TraceRunNotFoundError, "tracing_run_not_found"), + (accepted, TraceRunConflictError, "tracing_already_accepted"), + (collecting, TraceRunConflictError, "tracing_retry_not_allowed"), + (rejected_by_collection, TraceRunConflictError, "tracing_retry_not_allowed"), + ] + for run, error, code in cases: + state = run.state + with pytest.raises(error, match=code): + run.release_for_retry(acknowledge_possible_duplicates=True, now=NOW) + assert run.state is state + + +class TestRecoveryCheck: + @staticmethod + def test_unresolved_recovery_rotates_and_a_final_answer_fails_the_run() -> None: + run = frozen() + later = NOW + timedelta(minutes=1) + + run.recovery_checked(None, now=later) + assert (run.state, run.updated_at) == (TraceRunState.COLLECTING, later) + + run.recovery_checked("business_terminal_not_persisted", now=later) + assert (run.state, run.error_code) == (TraceRunState.FAILED, "business_terminal_not_persisted") + + @staticmethod + def test_captured_terminal_is_not_touched_by_recovery() -> None: + run = frozen() + run.capture_terminal(business_terminal(run), now=NOW) + + run.recovery_checked("business_terminal_not_persisted", now=NOW + timedelta(minutes=1)) + + assert (run.state, run.error_code, run.updated_at) == (TraceRunState.COLLECTING, None, NOW) + + +def agent_span(run: TraceRunAgg, span_id: str, parent_span_id: str, *, name: str, status: str) -> AgentTraceSpan: + return AgentTraceSpan( + trace_id=run.trace_id, + span_id=span_id, + parent_span_id=parent_span_id, + name=name, + start_time_ns=1_000_000_000, + end_time_ns=2_000_000_000, + status="error" if status == "error" else "ok", + status_message="tool failed" if status == "error" else None, + attributes={"gen_ai.operation.name": "invoke_agent"}, + ) + + +def report(run: TraceRunAgg, opened: TraceInvocation, spans: list[AgentTraceSpan]) -> None: + payload = AgentTraceBatch( + run_id=UUID(run.id), invocation_id=UUID(opened.id), batch_id=uuid4(), final=True, spans=spans + ) + ingest(run, payload) + + +class TestExportedTrace: + @staticmethod + def test_runtime_spans_nest_under_their_invocation_with_the_run_identity() -> None: + run = frozen(policy={**POLICY, "capture_content": True}) + run.capture_terminal(business_terminal(run), now=NOW) + opened = run.open_invocation(agent_id=identifier("agent"), node_execution_id=None, attempt=1, call_key="k") + report( + run, + opened, + [ + agent_span(run, "a" * 16, opened.span_id, name="Agent execution", status="error"), + agent_span(run, "b" * 16, "a" * 16, name="tool", status="ok"), + ], + ) + + trace = run.exported_trace(run.new_batches) + + assert trace.trace_id == run.trace_id + assert trace.session_id == "session" + assert trace.user_id == "user" + assert [item.span_id for item in trace.spans] == [run.root_span_id, opened.span_id, "a" * 16, "b" * 16] + attempt = trace.spans[1] + assert attempt.name == "Agent attempt 2" + assert attempt.parent_span_id == run.root_span_id + assert (attempt.status, attempt.error) == (SpanStatus.ERROR, "tool failed") + assert attempt.attributes["invocation_id"] == opened.id + assert all(item.attributes["dify.tracing.run_id"] == run.id for item in trace.spans) + assert all(item.attributes["dify.tenant_id"] == run.tenant_id for item in trace.spans) + + @staticmethod + def test_run_from_before_the_upgrade_exports_its_stored_terminal_shape() -> None: + run = frozen() + run.outer_trace = { + "trace_id": run.trace_id, + "root_span_id": run.root_span_id, + "session_id": "legacy-session", + "spans": [ + { + "id": run.root_span_id, + "parent_id": None, + "name": "message", + "kind": SpanKind.CHAIN.value, + "start_time": "2026-10-01T12:00:00Z", + "end_time": None, + "status": SpanStatus.OK.value, + "metadata": {"workflow_id": "legacy"}, + } + ], + } + + (root,) = run.exported_trace([]).spans + + assert root.span_id == run.root_span_id + assert root.end_time == root.start_time + assert root.attributes["workflow_id"] == "legacy" + + @staticmethod + def test_run_without_a_captured_terminal_exports_nothing() -> None: + with pytest.raises(TraceRunError, match="tracing_terminal_missing"): + frozen().exported_trace([]) + + @staticmethod + @pytest.mark.parametrize( + "code", + [ + "tracing_invocation_empty", + "tracing_invocation_not_found", + "tracing_span_conflict", + "tracing_invocation_parent_invalid", + "tracing_invocation_root_missing", + ], + ) + def test_invalid_runtime_trees_are_rejected(code: str) -> None: + run = frozen() + run.capture_terminal(business_terminal(run), now=NOW) + opened = run.open_invocation(agent_id=identifier("agent"), node_execution_id=None, attempt=0, call_key="k") + root = agent_span(run, "a" * 16, opened.span_id, name="Agent execution", status="ok") + batches: list[TraceBatch] + match code: + case "tracing_invocation_empty": + batches = [] + case "tracing_invocation_not_found": + report(run, opened, [root]) + batches = [replace(run.new_batches[0], invocation_id=identifier("other-invocation"))] + case "tracing_span_conflict": + report(run, opened, [root]) + conflicting = {**run.new_batches[0].spans[0], "name": "renamed"} + batches = [*run.new_batches, replace(run.new_batches[0], id=str(uuid4()), spans=[conflicting])] + case "tracing_invocation_parent_invalid": + report(run, opened, [root, agent_span(run, "b" * 16, "c" * 16, name="orphan", status="ok")]) + batches = run.new_batches + case _: + detached = [ + agent_span(run, "a" * 16, "b" * 16, name="detached", status="ok"), + agent_span(run, "b" * 16, "a" * 16, name="detached", status="ok"), + ] + report(run, opened, detached) + batches = run.new_batches + + with pytest.raises(TraceRunError, match=code): + run.exported_trace(batches) + + +class TestStates: + @staticmethod + def test_only_collecting_or_exported_runs_accept_a_terminal() -> None: + accepting = {state for state in TraceRunState if state.accepts_terminal} + + assert accepting == {TraceRunState.COLLECTING, TraceRunState.READY, TraceRunState.RETRYING} + + @staticmethod + def test_run_awaits_its_terminal_only_while_it_collects_without_one() -> None: + collecting = frozen() + captured = frozen() + captured.capture_terminal(business_terminal(captured), now=NOW) + + assert collecting.awaits_terminal is True + assert captured.awaits_terminal is False + assert frozen(config=None).awaits_terminal is False diff --git a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_discriminator.py b/api/tests/unit_tests/domain/workflow/test_agent_node.py similarity index 69% rename from api/tests/unit_tests/core/workflow/nodes/agent_v2/test_discriminator.py rename to api/tests/unit_tests/domain/workflow/test_agent_node.py index 1476764d37ff4f..356f791aca20e4 100644 --- a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_discriminator.py +++ b/api/tests/unit_tests/domain/workflow/test_agent_node.py @@ -1,8 +1,6 @@ import pytest -from pydantic import ValidationError -from core.workflow.nodes.agent_v2.discriminator import is_dify_agent_node_data -from core.workflow.nodes.agent_v2.entities import DifyAgentNodeData +from domain.workflow.agent_node import is_dify_agent_node_data from graphon.entities.base_node_data import BaseNodeData @@ -24,8 +22,3 @@ def test_is_dify_agent_node_data_supports_base_node_data() -> None: node_data = BaseNodeData.model_validate({"type": "agent", "version": "2", "agent_node_kind": "dify_agent"}) assert is_dify_agent_node_data(node_data) is True - - -def test_dify_agent_node_data_requires_explicit_kind_marker() -> None: - with pytest.raises(ValidationError, match="agent_node_kind"): - DifyAgentNodeData.model_validate({"type": "agent", "version": "2"}) diff --git a/api/tests/unit_tests/enterprise/telemetry/test_contracts.py b/api/tests/unit_tests/enterprise/telemetry/test_contracts.py index 7cdb9fb8cfacf2..2001dd45c43c06 100644 --- a/api/tests/unit_tests/enterprise/telemetry/test_contracts.py +++ b/api/tests/unit_tests/enterprise/telemetry/test_contracts.py @@ -161,7 +161,7 @@ def test_trace_enterprise_only_events(self) -> None: for cls in (DraftNodeExecutionTraceEvent, PromptGenerationEvent): assert cls.signal_type is SignalType.TRACE # type: ignore[attr-defined] assert cls.ce_eligible is False # type: ignore[attr-defined] - assert cls.trace_task_name is not None # type: ignore[attr-defined] + assert hasattr(cls, "trace_payload") def test_metric_log_events(self) -> None: from core.telemetry.events import ( @@ -174,4 +174,4 @@ def test_metric_log_events(self) -> None: for cls in (AppCreatedEvent, AppUpdatedEvent, AppDeletedEvent, FeedbackCreatedEvent): assert cls.signal_type is SignalType.METRIC_LOG # type: ignore[attr-defined] assert cls.ce_eligible is False # type: ignore[attr-defined] - assert cls.trace_task_name is None # type: ignore[attr-defined] + assert not hasattr(cls, "trace_payload") diff --git a/api/tests/unit_tests/enterprise/telemetry/test_draft_trace.py b/api/tests/unit_tests/enterprise/telemetry/test_draft_trace.py index 5068dd1431122b..abbf4fe675fe63 100644 --- a/api/tests/unit_tests/enterprise/telemetry/test_draft_trace.py +++ b/api/tests/unit_tests/enterprise/telemetry/test_draft_trace.py @@ -3,10 +3,20 @@ from __future__ import annotations import json +from collections.abc import Mapping +from dataclasses import dataclass, field from datetime import UTC, datetime from typing import cast from unittest.mock import MagicMock, patch +import pytest +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import TraceCorrelation, TraceEvent +from sqlalchemy.orm import Session, sessionmaker + +from core.telemetry import gateway as telemetry_gateway +from enterprise.telemetry.contracts import TelemetryCase +from enterprise.telemetry.draft_trace import enqueue_draft_node_execution_trace from graphon.enums import WorkflowNodeExecutionMetadataKey from models.enums import CreatorUserRole from models.workflow import ( @@ -14,6 +24,9 @@ WorkflowNodeExecutionStatus, WorkflowNodeExecutionTriggeredFrom, ) +from repositories.factory import DifyAPIRepositoryFactory +from repositories.trace_record_repository import TraceRecordRepository +from services.tracing.records import DraftNodeExecutionTrace # --------------------------------------------------------------------------- # Helpers @@ -235,8 +248,7 @@ def test_node_inputs_and_process_data_included(self) -> None: class TestEnqueueDraftNodeExecutionTrace: @patch("enterprise.telemetry.draft_trace.telemetry_emit") def test_emits_telemetry_event(self, mock_emit: MagicMock) -> None: - from core.telemetry import DraftNodeExecutionTraceEvent, TraceTaskName - from enterprise.telemetry.draft_trace import enqueue_draft_node_execution_trace + from core.telemetry import DraftNodeExecutionTraceEvent execution = _make_execution() enqueue_draft_node_execution_trace( @@ -248,7 +260,7 @@ def test_emits_telemetry_event(self, mock_emit: MagicMock) -> None: mock_emit.assert_called_once() event: DraftNodeExecutionTraceEvent = mock_emit.call_args[0][0] - assert event.trace_task_name == TraceTaskName.DRAFT_NODE_EXECUTION_TRACE + assert event.case == TelemetryCase.DRAFT_NODE_EXECUTION assert event.context.tenant_id == "tenant-1" assert event.context.user_id == "user-1" assert event.context.app_id == "app-1" @@ -256,7 +268,6 @@ def test_emits_telemetry_event(self, mock_emit: MagicMock) -> None: @patch("enterprise.telemetry.draft_trace.telemetry_emit") def test_payload_contains_node_execution_data(self, mock_emit: MagicMock) -> None: from core.telemetry import DraftNodeExecutionTraceEvent - from enterprise.telemetry.draft_trace import enqueue_draft_node_execution_trace execution = _make_execution() enqueue_draft_node_execution_trace( @@ -275,7 +286,6 @@ def test_payload_contains_node_execution_data(self, mock_emit: MagicMock) -> Non @patch("enterprise.telemetry.draft_trace.telemetry_emit") def test_outputs_forwarded_to_build(self, mock_emit: MagicMock) -> None: from core.telemetry import DraftNodeExecutionTraceEvent - from enterprise.telemetry.draft_trace import enqueue_draft_node_execution_trace execution = _make_execution(outputs_dict={"default": True}) enqueue_draft_node_execution_trace( @@ -291,7 +301,6 @@ def test_outputs_forwarded_to_build(self, mock_emit: MagicMock) -> None: @patch("enterprise.telemetry.draft_trace.telemetry_emit") def test_none_outputs_uses_execution_outputs(self, mock_emit: MagicMock) -> None: from core.telemetry import DraftNodeExecutionTraceEvent - from enterprise.telemetry.draft_trace import enqueue_draft_node_execution_trace execution = _make_execution(outputs_dict={"from_model": "yes"}) enqueue_draft_node_execution_trace( @@ -307,10 +316,34 @@ def test_none_outputs_uses_execution_outputs(self, mock_emit: MagicMock) -> None # --------------------------------------------------------------------------- # End-to-end token/model data flow: _build_node_execution_data → -# ops_trace_manager.draft_node_execution_trace → DraftNodeExecutionTrace +# DraftNodeExecuted trace event → TraceRecordRepository → DraftNodeExecutionTrace # --------------------------------------------------------------------------- +@dataclass +class _TracePipeline: + """Records the events the telemetry gateway emits through an app trace scope.""" + + events: list[TraceEvent] = field(default_factory=list) + + def create_scope( + self, *, tenant_id: str, app_id: str | None, user_id: str | None, extras: Mapping[str, object] + ) -> TraceScope: + assert extras == {} + return TraceScope( + emitter=self, + tenant_id=tenant_id, + app_id=app_id, + user_id=user_id, + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + ) + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + def _make_llm_execution() -> WorkflowNodeExecutionModel: """Return a real WorkflowNodeExecutionModel that mimics an LLM node. @@ -372,8 +405,8 @@ def _make_llm_execution() -> WorkflowNodeExecutionModel: class TestDraftTraceTokenDataFlow: """End-to-end test: verify all token and model fields survive from - _build_node_execution_data through ops_trace_manager.draft_node_execution_trace - to the DraftNodeExecutionTrace that enterprise_trace.py consumes. + _build_node_execution_data through the emitted trace event to the + DraftNodeExecutionTrace that enterprise_trace.py consumes. """ def test_all_token_and_model_fields_reach_trace_info(self) -> None: @@ -389,9 +422,8 @@ def test_all_token_and_model_fields_reach_trace_info(self) -> None: workflow_execution_id="run-flow", ) - # Simulate what ops_trace_manager.draft_node_execution_trace does: - # it calls node_execution_trace(node_execution_data=node_data) which - # reads top-level keys from node_data. Verify all expected keys exist. + # The trace assembler reads these top-level keys from the emitted node + # execution data. Verify all expected keys exist. expected_keys = { # Token fields — read by enterprise_trace._emit_node_execution_trace "total_tokens", @@ -477,60 +509,36 @@ def test_none_process_data_and_none_outputs(self) -> None: assert result["prompt_tokens"] is None assert result["completion_tokens"] is None - def test_node_data_feeds_into_draft_node_execution_trace(self) -> None: - """Verify the node_data dict can be consumed by - ops_trace_manager.draft_node_execution_trace without error and - produces a DraftNodeExecutionTrace with correct token/model fields.""" - - from enterprise.telemetry.draft_trace import _build_node_execution_data + def test_node_data_feeds_into_draft_node_execution_trace( + self, monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] + ) -> None: + """Verify the emitted node data is assembled into a DraftNodeExecutionTrace + with correct token/model fields.""" + pipeline = _TracePipeline() + monkeypatch.setattr(telemetry_gateway, "is_enterprise_telemetry_enabled", lambda: True) + monkeypatch.setattr(telemetry_gateway, "create_trace_scope_for_app", pipeline.create_scope) - execution = _make_llm_execution() - node_data = _build_node_execution_data( - execution=execution, + enqueue_draft_node_execution_trace( + execution=_make_llm_execution(), outputs=None, workflow_execution_id="run-e2e", + user_id="user-1", ) - # Directly construct DraftNodeExecutionTrace the way - # ops_trace_manager.node_execution_trace does (lines 1315-1350), - # skipping DB lookups by providing minimal metadata. - from core.ops.entities.trace_entity import DraftNodeExecutionTrace - - trace_info = DraftNodeExecutionTrace( - workflow_id=node_data.get("workflow_id", ""), - workflow_run_id=node_data.get("workflow_execution_id", ""), - tenant_id=node_data.get("tenant_id", ""), - node_execution_id=node_data.get("node_execution_id", ""), - node_id=node_data.get("node_id", ""), - node_type=node_data.get("node_type", ""), - title=node_data.get("title", ""), - status=node_data.get("status", ""), - error=node_data.get("error"), - elapsed_time=node_data.get("elapsed_time", 0.0), - index=node_data.get("index", 0), - predecessor_node_id=node_data.get("predecessor_node_id"), - total_tokens=node_data.get("total_tokens", 0), - total_price=node_data.get("total_price", 0.0), - currency=node_data.get("currency"), - model_provider=node_data.get("model_provider"), - model_name=node_data.get("model_name"), - prompt_tokens=node_data.get("prompt_tokens"), - completion_tokens=node_data.get("completion_tokens"), - tool_name=node_data.get("tool_name"), - iteration_id=node_data.get("iteration_id"), - iteration_index=node_data.get("iteration_index"), - loop_id=node_data.get("loop_id"), - loop_index=node_data.get("loop_index"), - parallel_id=node_data.get("parallel_id"), - node_inputs=node_data.get("node_inputs"), - node_outputs=node_data.get("node_outputs"), - process_data=node_data.get("process_data"), - start_time=node_data.get("created_at"), - end_time=node_data.get("finished_at"), - metadata={}, - ) + (event,) = pipeline.events + assert (event.tenant_id, event.app_id, event.user_id) == ("tenant-flow", "app-flow", "user-1") + trace_info = TraceRecordRepository( + sqlite_session_factory, + workflow_runs=DifyAPIRepositoryFactory.create_api_workflow_run_repository(sqlite_session_factory), + node_executions=DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository( + sqlite_session_factory + ), + files_url="https://files.example.com", + enterprise_details=False, + ).assemble(event) # These are the fields enterprise_trace._emit_node_execution_trace reads + assert isinstance(trace_info, DraftNodeExecutionTrace) assert trace_info.total_tokens == 80 assert trace_info.prompt_tokens == 50 assert trace_info.completion_tokens == 30 @@ -539,3 +547,7 @@ def test_node_data_feeds_into_draft_node_execution_trace(self) -> None: assert trace_info.node_type == "llm" assert trace_info.total_price == 0.0014 assert trace_info.currency == "USD" + assert trace_info.workflow_run_id == "run-e2e" + assert trace_info.node_execution_id == "exec-flow" + assert trace_info.start_time == datetime(2024, 1, 1, tzinfo=UTC) + assert trace_info.end_time == datetime(2024, 1, 1, 0, 0, 5, tzinfo=UTC) diff --git a/api/tests/unit_tests/enterprise/telemetry/test_enterprise_trace.py b/api/tests/unit_tests/enterprise/telemetry/test_enterprise_trace.py index 2181738b181a6c..5aa4c413cefc6e 100644 --- a/api/tests/unit_tests/enterprise/telemetry/test_enterprise_trace.py +++ b/api/tests/unit_tests/enterprise/telemetry/test_enterprise_trace.py @@ -10,7 +10,14 @@ import pytest -from core.ops.entities.trace_entity import ( +from enterprise.telemetry.enterprise_trace import EnterpriseOtelTrace +from enterprise.telemetry.entities import ( + EnterpriseTelemetryCounter, + EnterpriseTelemetryEvent, + EnterpriseTelemetryHistogram, + EnterpriseTelemetrySpan, +) +from services.tracing.records import ( DatasetRetrievalTraceInfo, DraftNodeExecutionTrace, GenerateNameTraceInfo, @@ -19,16 +26,10 @@ PromptGenerationTraceInfo, SuggestedQuestionTraceInfo, ToolTraceInfo, + TraceMessage, WorkflowNodeTraceInfo, WorkflowTraceInfo, ) -from enterprise.telemetry.enterprise_trace import EnterpriseOtelTrace -from enterprise.telemetry.entities import ( - EnterpriseTelemetryCounter, - EnterpriseTelemetryEvent, - EnterpriseTelemetryHistogram, - EnterpriseTelemetrySpan, -) # --------------------------------------------------------------------------- # Fixtures @@ -59,8 +60,31 @@ def trace_handler(mock_exporter): _T1 = datetime(2024, 1, 10, 12, 0, 5, tzinfo=UTC) +def _trace_message() -> TraceMessage: + return TraceMessage( + id="msg-001", + app_id="app-001", + conversation_id="conv-001", + workflow_run_id=None, + query="user input", + answer="assistant output", + error=None, + created_at=_T0, + model_provider="openai", + model_id="gpt-4", + total_price=None, + currency="USD", + ) + + def make_workflow_info(**overrides) -> WorkflowTraceInfo: defaults: dict[str, Any] = { + "node_executions": [], + "conversation_id": None, + "message_started_at": None, + "workflow_app_log_id": None, + "error": None, + "invoked_by": None, "workflow_id": "wf-001", "tenant_id": "tenant-abc", "workflow_run_id": "run-001", @@ -141,6 +165,12 @@ def make_draft_node_info(**overrides) -> DraftNodeExecutionTrace: def make_message_info(**overrides) -> MessageTraceInfo: defaults: dict[str, Any] = { "message_id": "msg-001", + "message_data": _trace_message(), + "error": None, + "file_list": [], + "gen_ai_server_time_to_first_token": None, + "llm_streaming_time_to_generate": None, + "is_streaming_request": False, "conversation_model": "gpt-4", "message_tokens": 40, "answer_tokens": 60, @@ -172,6 +202,8 @@ def make_tool_info(**overrides) -> ToolTraceInfo: "tool_config": {"max_results": 5}, "tool_parameters": {"verbose": True}, "time_cost": 1.5, + "error": None, + "file_url": None, "metadata": {"app_id": "app-001", "tenant_id": "tenant-abc"}, } defaults.update(overrides) @@ -198,6 +230,12 @@ def make_suggested_question_info(**overrides) -> SuggestedQuestionTraceInfo: "suggested_question": ["Question A?", "Question B?"], "level": "info", "status": "succeeded", + "status_message": None, + "error": None, + "from_account_id": None, + "agent_based": False, + "from_source": "api", + "workflow_run_id": None, "model_provider": "openai", "model_id": "gpt-3.5-turbo", "start_time": _T0, @@ -223,6 +261,7 @@ def make_dataset_retrieval_info(**overrides) -> DatasetRetrievalTraceInfo: } ], "inputs": "search query", + "error": None, "metadata": { "app_id": "app-001", "tenant_id": "tenant-abc", @@ -266,6 +305,9 @@ def make_prompt_generation_info(**overrides) -> PromptGenerationTraceInfo: "model_provider": "openai", "model_name": "gpt-4", "latency": 3.2, + "total_price": None, + "currency": None, + "error": None, "metadata": {"app_id": "app-001", "tenant_id": "tenant-abc"}, } defaults.update(overrides) @@ -550,24 +592,6 @@ def test_records_workflow_duration_histogram(self, trace_handler: EnterpriseOtel assert hist_call[0][0] == EnterpriseTelemetryHistogram.WORKFLOW_DURATION assert hist_call[0][1] == pytest.approx(5.0) - def test_duration_falls_back_to_elapsed_time_when_timestamps_missing( - self, trace_handler: EnterpriseOtelTrace, mock_exporter - ): - with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"): - info = make_workflow_info(start_time=None, end_time=None, workflow_run_elapsed_time=7.3) - trace_handler._workflow_trace(info) - - hist_call = mock_exporter.record_histogram.call_args - assert hist_call[0][1] == pytest.approx(7.3) - - def test_duration_defaults_to_zero_when_no_timing(self, trace_handler: EnterpriseOtelTrace, mock_exporter): - with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"): - info = make_workflow_info(start_time=None, end_time=None, workflow_run_elapsed_time=0) - trace_handler._workflow_trace(info) - - hist_call = mock_exporter.record_histogram.call_args - assert hist_call[0][1] == pytest.approx(0.0) - def test_error_path_increments_error_counter(self, trace_handler: EnterpriseOtelTrace, mock_exporter): with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"): info = make_workflow_info(error="Something went wrong", workflow_run_status="failed") diff --git a/api/tests/unit_tests/extensions/application_services/test_tracing.py b/api/tests/unit_tests/extensions/application_services/test_tracing.py new file mode 100644 index 00000000000000..9e079027b2eed3 --- /dev/null +++ b/api/tests/unit_tests/extensions/application_services/test_tracing.py @@ -0,0 +1,85 @@ +from collections.abc import Callable + +import pytest +from sqlalchemy.orm import Session, sessionmaker + +from enums import DeploymentEdition +from extensions.application_services.tracing import build_tracing_services +from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository +from repositories.application_trace_repository import ApplicationTraceRepository +from repositories.trace_record_repository import TraceRecordRepository +from services.tracing.adapters import ( + CeleryExportRetryQueue, + CeleryTraceEmitter, + WorkspaceTracingSecretCipher, + schedule_agent_trace_export, + send_to_enterprise_telemetry, +) +from services.tracing.egress import SSRFEgress +from services.tracing.registry import EntryPointTraceProviderCatalog +from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks + + +def test_delivery_reads_through_the_deployment_adapters(sqlite_session_factory: sessionmaker[Session]) -> None: + tracing = build_tracing_services(database_client=sqlite_session_factory) + delivery = tracing.delivery + + assert isinstance(delivery._records, TraceRecordRepository) + assert delivery._records._sessions is sqlite_session_factory + assert isinstance(delivery._destinations, SQLAlchemyAppTracingConfigRepository) + assert isinstance(delivery._providers, EntryPointTraceProviderCatalog) + assert isinstance(delivery._codec._cipher, WorkspaceTracingSecretCipher) + assert isinstance(delivery._egress, SSRFEgress) + assert isinstance(delivery._retries, CeleryExportRetryQueue) + assert isinstance(tracing.scopes._emitter, CeleryTraceEmitter) + assert isinstance(tracing.scopes._policies, SQLAlchemyAppTracingConfigRepository) + assert tracing.provider_configs._providers is delivery._providers + assert tracing.provider_configs._codec is delivery._codec + + +def test_agent_collection_shares_the_providers_and_schedules_exports_on_the_trace_queue( + sqlite_session_factory: sessionmaker[Session], +) -> None: + tracing = build_tracing_services(database_client=sqlite_session_factory) + agent_traces = tracing.agent_traces + + # One repository persists the run aggregates and answers the run and source queries. + assert isinstance(agent_traces._runs, ApplicationTraceRepository) + assert agent_traces._run_query is agent_traces._runs + assert agent_traces._sources is agent_traces._runs + assert agent_traces._runs._sessions is sqlite_session_factory + assert isinstance(agent_traces._records, TraceRecordRepository) + # Agent traces never carry the enterprise telemetry dimensions. + assert agent_traces._records._enterprise_details is False + assert agent_traces._providers is tracing.delivery._providers + assert agent_traces._schedule_export is schedule_agent_trace_export + + +@pytest.mark.parametrize("enabled", [True, False]) +def test_agent_collection_joins_delivery_only_when_enabled( + sqlite_session_factory: sessionmaker[Session], config_overrides: Callable[..., None], enabled: bool +) -> None: + config_overrides(OPS_AGENT_TRACE_ENABLED=enabled, OPS_AGENT_TRACE_MAX_ATTEMPTS=3) + + tracing = build_tracing_services(database_client=sqlite_session_factory) + + assert (tracing.delivery._collector is tracing.agent_traces) is enabled + assert isinstance(tracing.hooks, RuntimeApplicationTracingHooks) is enabled + assert tracing.agent_traces._settings.enabled is enabled + assert tracing.agent_traces._settings.max_export_attempts == 3 + assert tracing.provider_capabilities._api_collection_enabled is enabled + + +@pytest.mark.parametrize("edition", [DeploymentEdition.ENTERPRISE, DeploymentEdition.COMMUNITY]) +def test_enterprise_telemetry_receives_records_with_its_dimensions_only_when_enabled( + sqlite_session_factory: sessionmaker[Session], config_overrides: Callable[..., None], edition: DeploymentEdition +) -> None: + config_overrides(DEPLOYMENT_EDITION=edition, ENTERPRISE_TELEMETRY_ENABLED=True, OPS_TRACE_EXPORT_MAX_ATTEMPTS=4) + enterprise = edition == DeploymentEdition.ENTERPRISE + + delivery = build_tracing_services(database_client=sqlite_session_factory).delivery + + assert (delivery._enterprise_sink is send_to_enterprise_telemetry) is enterprise + assert isinstance(delivery._records, TraceRecordRepository) + assert delivery._records._enterprise_details is enterprise + assert delivery._max_export_attempts == 4 diff --git a/api/tests/unit_tests/extensions/test_ext_application_services.py b/api/tests/unit_tests/extensions/test_ext_application_services.py index 2f8aea35211ff4..7e07f078a88074 100644 --- a/api/tests/unit_tests/extensions/test_ext_application_services.py +++ b/api/tests/unit_tests/extensions/test_ext_application_services.py @@ -6,7 +6,7 @@ from decimal import Decimal from io import BytesIO from types import SimpleNamespace -from typing import Never, cast +from typing import cast from unittest.mock import MagicMock, call, create_autospec, patch from uuid import uuid4 @@ -20,7 +20,6 @@ from core.credit_usage import CreditUsageAppType, CreditUsageCreatedBy from core.llm_generator.prompts import CONVERSATION_TITLE_PROMPT from core.model_context import get_credit_usage_metadata, use_credit_usage_metadata -from core.ops.base_trace_instance import application_tracing_hooks from core.plugin.impl import base as plugin_http from core.tools.tool_file_manager import ToolFileManager from enums import DeploymentEdition, WebAppAccessMode @@ -28,6 +27,7 @@ from extensions.application_services.app import AppServices from extensions.ext_database import db from extensions.ext_redis import RedisClientWrapper +from extensions.tracing import application_tracing_hooks, create_trace_scope_for_app from graphon.model_runtime.entities.model_entities import ModelType from machinery.context import RequestContext from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole @@ -107,7 +107,6 @@ from services.app_scoped_end_user_query_service import AppScopedEndUserQueryService from services.app_scoped_end_user_service import AppScopedEndUserService from services.app_site_service import AppSiteService -from services.app_tracing_config_gateway import OpsTraceManagerGateway from services.app_tracing_config_service import AppTracingConfigService from services.audio_types import AudioAppRef, AudioOutput, AudioUpload from services.billing_portal_service import BillingPortalService @@ -155,7 +154,8 @@ from services.retention.workflow_run.archive_log_service import WorkflowRunArchiveService from services.tag_application_service import TagApplicationService from services.tool_file_download_service import ToolFileDownloadService -from services.tracing import runtime_bridge +from services.tracing.adapters import CeleryTraceEmitter +from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks from services.upload_file_delivery_service import UploadFileDeliveryService from services.web_authentication_adapters import ( AccountWebAuthenticationSecurityGateway, @@ -170,6 +170,7 @@ from services.workflow_statistic_query_service import WorkflowStatisticQueryService from tasks.delete_conversation_task import delete_conversation_related_data from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_app from tests.helpers.unit.services.conversation_name_fixtures import _MODEL, _PROVIDER, NamingPlugin from tests.unit_tests.services.test_app_task_service import _AppTaskServiceStopredis as _StopRedis @@ -281,9 +282,11 @@ def test_init_app_registers_services_for_the_current_app( assert creators._oauth is services.oauth_server -def test_init_app_registers_distinct_tracing_ports_without_constructing_trace_services( +@pytest.mark.parametrize("enabled", [True, False]) +def test_init_app_binds_each_apps_agent_hooks_only_when_agent_tracing_is_enabled( monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session], + enabled: bool, ) -> None: first_app, second_app = Flask("first"), Flask("second") monkeypatch.setattr(ext_application_services, "get_session_maker", lambda: sqlite_session_factory) @@ -291,28 +294,60 @@ def test_init_app_registers_distinct_tracing_ports_without_constructing_trace_se monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, INIT_PASSWORD="", - OPS_AGENT_TRACE_ENABLED=True, + OPS_AGENT_TRACE_ENABLED=enabled, ) - construction_calls: list[bool] = [] - - def construct_trace_service() -> Never: - construction_calls.append(True) - raise AssertionError("Composition must not construct a trace service") - - monkeypatch.setattr(runtime_bridge, "_service", construct_trace_service) ext_application_services.init_app(first_app) ext_application_services.init_app(second_app) with first_app.app_context(): first_hooks = application_tracing_hooks() - assert isinstance(first_hooks, runtime_bridge.RuntimeApplicationTracingHooks) + first_services = ext_application_services.application_services() with second_app.app_context(): second_hooks = application_tracing_hooks() - assert isinstance(second_hooks, runtime_bridge.RuntimeApplicationTracingHooks) - assert second_hooks is not first_hooks assert application_tracing_hooks() is first_hooks - assert construction_calls == [] + if not enabled: + assert first_hooks is None + assert second_hooks is None + return + assert isinstance(first_hooks, RuntimeApplicationTracingHooks) + assert isinstance(second_hooks, RuntimeApplicationTracingHooks) + assert second_hooks is not first_hooks + assert first_hooks is first_services.tracing.hooks + + +def test_init_app_binds_trace_scopes_to_the_apps_tracing_policy( + monkeypatch: pytest.MonkeyPatch, + sqlite_session_factory: sessionmaker[Session], +) -> None: + app = Flask(__name__) + monkeypatch.setattr(ext_application_services, "get_session_maker", lambda: sqlite_session_factory) + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, INIT_PASSWORD="") + with sqlite_session_factory.begin() as session: + session.add_all( + [ + make_app( + app_id="traced-app", + tenant_id="tenant-1", + tracing=json.dumps({"enabled": True, "tracing_provider": "langfuse"}), + ), + make_app(app_id="untraced-app", tenant_id="tenant-1"), + ] + ) + + ext_application_services.init_app(app) + + with app.app_context(): + scope = create_trace_scope_for_app( + tenant_id="tenant-1", app_id="traced-app", user_id="user-1", extras={"external_trace_id": "external-1"} + ) + assert scope is not None + assert isinstance(scope.emitter, CeleryTraceEmitter) + assert (scope.tenant_id, scope.app_id, scope.user_id) == ("tenant-1", "traced-app", "user-1") + assert scope.correlation.external_trace_id == "external-1" + assert create_trace_scope_for_app(tenant_id="tenant-1", app_id="untraced-app", user_id=None, extras={}) is None + with Flask("unbound").app_context(): + assert create_trace_scope_for_app(tenant_id="tenant-1", app_id="traced-app", user_id=None, extras={}) is None def test_build_application_services_preserves_composed_boundaries( @@ -643,7 +678,7 @@ def test_build_application_services_wires_app_tracing_config_boundary( assert isinstance(services.app_tracing_configs, AppTracingConfigService) assert isinstance(services.app_tracing_configs._configs, SQLAlchemyAppTracingConfigRepository) assert services.app_tracing_configs._configs._session_factory is sqlite_session_factory - assert isinstance(services.app_tracing_configs._provider, OpsTraceManagerGateway) + assert services.app_tracing_configs._provider is services.tracing.provider_configs def test_build_application_services_wires_workflow_app_log_boundary( diff --git a/api/tests/unit_tests/extensions/test_ext_celery.py b/api/tests/unit_tests/extensions/test_ext_celery.py index 9cf4bea1647c0c..e34e06daf59e3d 100644 --- a/api/tests/unit_tests/extensions/test_ext_celery.py +++ b/api/tests/unit_tests/extensions/test_ext_celery.py @@ -2,11 +2,14 @@ from collections.abc import Iterator from datetime import timedelta from importlib import import_module +from queue import Empty from types import SimpleNamespace from unittest.mock import MagicMock, Mock, patch import pytest +from celery import Celery from celery.contrib.testing.app import TestApp, setup_default_app +from celery.exceptions import Retry from dify_app import DifyApp from enums import DeploymentEdition @@ -15,13 +18,95 @@ get_celery_broker_transport_options, init_app, ) +from schedule.collect_agent_sandbox_usage import collect_agent_sandbox_usage from services.knowledge_fs.background_contract import OPERATION_QUEUES, OPERATION_TASK from services.knowledge_fs.lifecycle_readiness import KnowledgeFSLifecycleWorkerReadiness -from tests.helpers.unit.config_override import apply_config_overrides +from tasks.application_trace_task import export_application_trace +from tasks.trace_delivery_task import deliver_trace_event, retry_trace_export +from tests.helpers.unit.config_override import apply_config_overrides, config_overrides_context """Tests for Celery SSL configuration.""" +@pytest.fixture +def tracing_celery(request: pytest.FixtureRequest) -> Iterator[Celery]: + with ( + setup_default_app(TestApp()), + config_overrides_context( + CELERY_BROKER_URL="memory://", + CELERY_BACKEND="disabled", + OPS_TRACE_QUEUE=request.param, + OPS_AGENT_TRACE_ENABLED=True, + ENABLE_APPLICATION_TRACE_RECONCILE_TASK=True, + ENABLE_TRACE_PAYLOAD_SWEEP_TASK=True, + DISABLE_TELEMETRY=True, + ), + ): + app = init_app(DifyApp(__name__)) + try: + yield app + finally: + app.close() + + +@pytest.mark.parametrize("tracing_celery", ["ops_trace_v2", "ops_trace", "tracing-upgrade"], indirect=True) +def test_tracing_dispatch_and_beat_use_configured_queue(tracing_celery: Celery) -> None: + queue = tracing_celery.conf.task_routes["tasks.trace_delivery_task.*"]["queue"] + calls = { + deliver_trace_event.name: ({"event": "{}"},), + retry_trace_export.name: ("trace_payloads/job.json",), + export_application_trace.name: ("run-1",), + tracing_celery.conf.beat_schedule["sweep_trace_payloads"]["task"]: (), + tracing_celery.conf.beat_schedule["reconcile_application_traces"]["task"]: (), + collect_agent_sandbox_usage.name: (), + } + with tracing_celery.connection_for_write() as connection: + if queue != "ops_trace": + tracing_celery.send_task("tasks.ops_trace_task.process_trace_tasks", args=({},), queue="ops_trace") + for name, args in calls.items(): + tracing_celery.tasks[name].apply_async(args=args) + with connection.SimpleQueue(queue) as inbox: + for name, args in calls.items(): + message = inbox.get(block=False) + assert message.headers["task"] == name + assert message.payload[0] == list(args) + message.ack() + with pytest.raises(Empty): + inbox.get(block=False) + if queue != "ops_trace": + with connection.SimpleQueue("ops_trace") as legacy: + message = legacy.get(block=False) + assert message.headers["task"] == "tasks.ops_trace_task.process_trace_tasks" + message.ack() + + +@pytest.mark.parametrize("tracing_celery", ["ops_trace"], indirect=True) +def test_redelivery_stays_on_original_queue_during_switchback(tracing_celery: Celery) -> None: + task = tracing_celery.tasks["tasks.trace_delivery_task.deliver_trace_event"] + task.push_request( + id="trace-redelivery", + args=({"event": "{}"},), + kwargs={}, + retries=0, + called_directly=False, + delivery_info={"exchange": "", "routing_key": "ops_trace_v2"}, + ) + try: + with pytest.raises(Retry): + task.retry(countdown=5) + with tracing_celery.connection_for_read() as connection: + with connection.SimpleQueue("ops_trace_v2") as original: + message = original.get(block=False) + assert message.headers["task"] == task.name + assert message.headers["retries"] == 1 + message.ack() + with connection.SimpleQueue("ops_trace") as restored: + with pytest.raises(Empty): + restored.get(block=False) + finally: + task.pop_request() + + class TestCelerySslCases: @staticmethod @pytest.fixture(autouse=True) diff --git a/api/tests/unit_tests/extensions/test_tracing.py b/api/tests/unit_tests/extensions/test_tracing.py new file mode 100644 index 00000000000000..cec06ffa4e581c --- /dev/null +++ b/api/tests/unit_tests/extensions/test_tracing.py @@ -0,0 +1,186 @@ +from typing import TYPE_CHECKING, Literal, override + +from dify_agent.layers.execution_context import ApplicationTracingContext +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import TraceCorrelation, TraceEvent +from flask import Flask + +from core.helper.trace_id_helper import ParentTraceContext +from extensions.tracing import ( + ApplicationTracingHooks, + ExecutionTraceScopes, + application_tracing_hooks, + bind_tracing, + correlation_from_extras, + create_trace_scope, + create_trace_scope_for_app, +) + +if TYPE_CHECKING: + from core.app.apps.agent_app.runtime_request_builder import AgentAppRuntimeBuildContext + from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, WorkflowAppGenerateEntity + from core.workflow.nodes.agent_v2.runtime_request_builder import WorkflowAgentRuntimeBuildContext + + +class _Recorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +class _Scopes(ExecutionTraceScopes): + def __init__(self, *, fails: bool) -> None: + self.fails = fails + self.policies: list[str | None] = [] + self.correlations: list[TraceCorrelation] = [] + + def _scope( + self, *, tenant_id: str, app_id: str | None, user_id: str | None, correlation: TraceCorrelation + ) -> TraceScope: + if self.fails: + raise RuntimeError("policy lookup failed") + self.correlations.append(correlation) + return TraceScope( + emitter=_Recorder(), tenant_id=tenant_id, app_id=app_id, user_id=user_id, correlation=correlation + ) + + @override + def for_execution( + self, + *, + tenant_id: str, + app_id: str | None, + app_tracing: str | None, + user_id: str | None, + correlation: TraceCorrelation, + ) -> TraceScope | None: + self.policies.append(app_tracing) + return self._scope(tenant_id=tenant_id, app_id=app_id, user_id=user_id, correlation=correlation) + + @override + def for_app( + self, *, tenant_id: str, app_id: str | None, user_id: str | None, correlation: TraceCorrelation + ) -> TraceScope | None: + return self._scope(tenant_id=tenant_id, app_id=app_id, user_id=user_id, correlation=correlation) + + +class _Hooks(ApplicationTracingHooks): + @override + def prepare_agent_app(self, context: "AgentAppRuntimeBuildContext") -> ApplicationTracingContext | None: + return None + + @override + def prepare_workflow_agent( + self, context: "WorkflowAgentRuntimeBuildContext", conversation_id: str | None + ) -> ApplicationTracingContext | None: + return None + + @override + def freeze_workflow_run( + self, + entity: "AdvancedChatAppGenerateEntity | WorkflowAppGenerateEntity", + *, + workflow_id: str, + workflow_run_id: str, + conversation_id: object, + graph_data: object, + ) -> None: + return None + + @override + def capture_agent_app_terminal( + self, + *, + tenant_id: str, + app_id: str, + message_id: str, + outcome: Literal["succeeded", "failed", "cancelled"], + output: str | None, + ) -> None: + return None + + +def _app(scopes: _Scopes, hooks: ApplicationTracingHooks | None) -> Flask: + app = Flask(__name__) + bind_tracing(app, scopes=scopes, hooks=hooks) + return app + + +def test_correlation_reads_the_caller_extras() -> None: + nested = correlation_from_extras( + { + "external_trace_id": "caller-trace", + "trace_session_id": "session-1", + "parent_trace_context": {"parent_workflow_run_id": "outer-run", "parent_node_execution_id": "outer-node"}, + } + ) + typed = correlation_from_extras( + {"parent_trace_context": ParentTraceContext(parent_workflow_run_id="outer-run", parent_node_execution_id=None)} + ) + + assert nested == TraceCorrelation( + external_trace_id="caller-trace", + session_id="session-1", + parent_workflow_run_id="outer-run", + parent_node_execution_id="outer-node", + ) + assert typed.parent_workflow_run_id == "outer-run" + assert correlation_from_extras({"external_trace_id": "", "parent_trace_context": {"x": 1}}) == TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ) + + +def test_scopes_come_from_the_bound_scopes_with_the_raw_policy_and_correlation() -> None: + scopes = _Scopes(fails=False) + with _app(scopes, None).app_context(): + scope = create_trace_scope( + tenant_id="tenant-1", + app_id="app-1", + app_tracing='{"enabled": true}', + user_id="u", + extras={"external_trace_id": "caller-trace"}, + ) + app_scope = create_trace_scope_for_app(tenant_id="tenant-1", app_id=None, user_id=None, extras={}) + + assert scope is not None + assert (scope.tenant_id, scope.app_id, scope.user_id) == ("tenant-1", "app-1", "u") + assert scopes.policies == ['{"enabled": true}'] + assert scope.correlation.external_trace_id == "caller-trace" + assert app_scope is not None + assert app_scope.app_id is None + + +def test_scopes_fail_open() -> None: + assert create_trace_scope(tenant_id="tenant-1", app_id="app-1", app_tracing=None, user_id=None, extras={}) is None + with _app(_Scopes(fails=True), None).app_context(): + assert ( + create_trace_scope(tenant_id="tenant-1", app_id="app-1", app_tracing=None, user_id=None, extras={}) is None + ) + assert create_trace_scope_for_app(tenant_id="tenant-1", app_id="app-1", user_id=None, extras={}) is None + with Flask(__name__).app_context(): + assert create_trace_scope_for_app(tenant_id="tenant-1", app_id="app-1", user_id=None, extras={}) is None + + +def test_agent_hooks_are_the_ones_the_composition_root_bound() -> None: + hooks = _Hooks() + + with _app(_Scopes(fails=False), hooks).app_context(): + assert application_tracing_hooks() is hooks + with _app(_Scopes(fails=False), None).app_context(): + assert application_tracing_hooks() is None + assert application_tracing_hooks() is None + + +def test_invalid_extension_bindings_fail_open() -> None: + app = Flask(__name__) + app.extensions["trace_scopes"] = object() + app.extensions["application_tracing_hooks"] = object() + + with app.app_context(): + assert ( + create_trace_scope(tenant_id="tenant-1", app_id="app-1", app_tracing=None, user_id=None, extras={}) is None + ) + assert create_trace_scope_for_app(tenant_id="tenant-1", app_id="app-1", user_id=None, extras={}) is None + assert application_tracing_hooks() is None diff --git a/api/tests/unit_tests/libs/test_datetime_utils.py b/api/tests/unit_tests/libs/test_datetime_utils.py index a4229d207778c9..b62c3585ce04f0 100644 --- a/api/tests/unit_tests/libs/test_datetime_utils.py +++ b/api/tests/unit_tests/libs/test_datetime_utils.py @@ -4,7 +4,7 @@ import pytest import pytz -from libs.datetime_utils import naive_utc_now, parse_time_range, to_utc_timestamp, utc_now +from libs.datetime_utils import Timing, measure_time, naive_utc_now, parse_time_range, to_utc_timestamp, utc_now def test_utc_now(monkeypatch: pytest.MonkeyPatch): @@ -48,6 +48,27 @@ def test_to_utc_timestamp(value: datetime.datetime): assert to_utc_timestamp(value) == 1704067200 +def test_measure_time_yields_aware_utc_timing() -> None: + with measure_time() as timer: + pass + + assert timer.started_at.tzinfo is datetime.UTC + assert timer.ended_at.tzinfo is datetime.UTC + assert timer.ended_at >= timer.started_at + + +def test_measure_time_records_end_when_block_raises(monkeypatch: pytest.MonkeyPatch) -> None: + started_at = datetime.datetime(2026, 8, 26, 12, tzinfo=datetime.UTC) + ended_at = started_at + datetime.timedelta(seconds=2) + clock = iter([started_at, ended_at]) + monkeypatch.setattr("libs.datetime_utils._now_func", lambda tz: next(clock).astimezone(tz)) + + with pytest.raises(RuntimeError, match="boom"), measure_time() as timer: + raise RuntimeError("boom") + + assert timer == Timing(started_at=started_at, ended_at=ended_at) + + class TestParseTimeRange: """Test cases for parse_time_range function.""" diff --git a/api/tests/unit_tests/migrations/versions/test_2026_10_23_1200-53c10126d551_unique_trace_app_config_provider.py b/api/tests/unit_tests/migrations/versions/test_2026_10_23_1200-53c10126d551_unique_trace_app_config_provider.py new file mode 100644 index 00000000000000..6e50b0210ad5da --- /dev/null +++ b/api/tests/unit_tests/migrations/versions/test_2026_10_23_1200-53c10126d551_unique_trace_app_config_provider.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +import importlib.util +from io import StringIO +from pathlib import Path +from types import ModuleType +from unittest.mock import patch + +import pytest +import sqlalchemy as sa +from alembic.migration import MigrationContext +from alembic.operations import Operations +from sqlalchemy.exc import IntegrityError + +_PATH = ( + Path(__file__).resolve().parents[4] + / "migrations/versions/2026_10_23_1200-53c10126d551_unique_trace_app_config_provider.py" +) + + +def _migration() -> ModuleType: + spec = importlib.util.spec_from_file_location("unique_trace_app_config_provider", _PATH) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_the_most_recently_updated_duplicate_survives() -> None: + engine = sa.create_engine("sqlite:///:memory:") + migration = _migration() + with engine.begin() as connection: + connection.execute( + sa.text( + "CREATE TABLE trace_app_config (" + "id TEXT PRIMARY KEY, app_id TEXT NOT NULL, tracing_provider TEXT, updated_at TEXT NOT NULL)" + ) + ) + connection.execute( + sa.text( + "INSERT INTO trace_app_config VALUES " + "('old', 'app-1', 'langfuse', '2026-01-01'), " + "('new', 'app-1', 'langfuse', '2026-02-01'), " + "('tie-a', 'app-1', 'opik', '2026-01-01'), " + "('tie-b', 'app-1', 'opik', '2026-01-01'), " + "('other-app', 'app-2', 'langfuse', '2026-01-01'), " + "('unset-1', 'app-3', NULL, '2026-01-01'), " + "('unset-2', 'app-3', NULL, '2026-01-01')" + ) + ) + with patch.object(migration, "op", Operations(MigrationContext.configure(connection))): + migration.upgrade() + remaining = connection.execute(sa.text("SELECT id FROM trace_app_config ORDER BY id")).scalars().all() + assert remaining == ["new", "other-app", "tie-b", "unset-1", "unset-2"] + with pytest.raises(IntegrityError): + connection.execute( + sa.text("INSERT INTO trace_app_config VALUES ('again', 'app-1', 'langfuse', '2026-03-01')") + ) + engine.dispose() + + +@pytest.mark.parametrize("dialect", ["postgresql", "mysql"]) +def test_migration_renders_for_supported_databases(dialect: str) -> None: + output = StringIO() + context = MigrationContext.configure(dialect_name=dialect, opts={"as_sql": True, "output_buffer": output}) + migration = _migration() + with patch.object(migration, "op", Operations(context)): + migration.upgrade() + migration.downgrade() + sql = output.getvalue() + assert "DELETE FROM trace_app_config" in sql + assert "CREATE UNIQUE INDEX trace_app_config_app_provider_uq" in sql + assert "DROP INDEX trace_app_config_app_provider_uq" in sql diff --git a/api/tests/unit_tests/models/test_model.py b/api/tests/unit_tests/models/test_model.py index fb7180a41ee3e8..fd2ce88d8d5cd6 100644 --- a/api/tests/unit_tests/models/test_model.py +++ b/api/tests/unit_tests/models/test_model.py @@ -1,23 +1,13 @@ import json -from datetime import UTC, datetime, timedelta +from datetime import UTC, datetime from decimal import Decimal from unittest.mock import patch from uuid import UUID, uuid4 import pytest -from dify_trace_arize_phoenix import arize_phoenix_trace -from dify_trace_arize_phoenix.arize_phoenix_trace import ArizePhoenixDataTrace -from dify_trace_arize_phoenix.config import ArizeConfig -from openinference.semconv.trace import SpanAttributes -from opentelemetry.sdk.trace import Tracer, TracerProvider -from opentelemetry.sdk.trace.export import SimpleSpanProcessor -from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter -from pydantic import TypeAdapter from sqlalchemy import inspect, select from sqlalchemy.orm import Session -from core.ops.entities.trace_entity import MessageTraceInfo -from core.ops.ops_trace_manager import OpsTraceManager, TaskData from core.workflow.file_reference import build_file_reference from graphon.file import FILE_MODEL_IDENTITY, FileTransferMethod from models.account import Account, Tenant @@ -49,16 +39,12 @@ InstalledApp, Message, MessageAnnotation, - MessageDict, MessageFeedback, RecommendedApp, Site, load_annotation_reply_config, ) from models.workflow import Workflow, WorkflowType -from tasks import ops_trace_task -from tests.helpers.unit.config_override import apply_config_overrides -from tests.helpers.unit.model_factories import make_message """Regression coverage for the ``@property``→session-parameter refactor on ``AppAnnotationHitHistory`` accessors. @@ -2085,148 +2071,6 @@ def test_returns_none_when_account_missing(self, sqlite_session: Session) -> Non assert annotation.annotation_create_account_with_session(session=sqlite_session) is None -"""Keep the persisted provider through private trace JSON and the real consumer.""" - - -_MESSAGE = TypeAdapter(MessageDict) - - -class TestMessageTraceSerializationCases: - @staticmethod - @pytest.fixture - def persisted_message(sqlite_session: Session) -> Message: - now = datetime(2026, 10, 20, tzinfo=UTC) - message = make_message( - message_id=str(uuid4()), - app_id=str(uuid4()), - conversation_id=str(uuid4()), - inputs={}, - query="Synthetic input", - message={}, - answer="Synthetic output", - message_unit_price=Decimal(0), - answer_unit_price=Decimal(0), - total_price=Decimal(0), - currency="USD", - from_source=ConversationFromSource.CONSOLE, - created_at=now, - updated_at=now + timedelta(seconds=1), - ) - message.model_provider = "langgenius/openai/openai" - message.model_id = "gpt-4o-mini" - sqlite_session.add(message) - sqlite_session.commit() - sqlite_session.expire_all() - restored = sqlite_session.get(Message, message.id) - assert restored is not None - return restored - - @staticmethod - @pytest.mark.parametrize("provider", ["langgenius/openai/openai", "Custom/Provider:Exact", None]) - def test_message_provider_survives_real_json_round_trip( - sqlite_session: Session, - persisted_message: Message, - provider: str | None, - ) -> None: - persisted_message.model_provider = provider - sqlite_session.commit() - sqlite_session.refresh(persisted_message) - wire = json.dumps(persisted_message.to_dict(session=sqlite_session), default=str) - payload = _MESSAGE.validate_json(wire) - restored = Message.from_dict(payload) - assert "model_provider" in payload - assert restored.model_provider == provider - assert restored.model_id == persisted_message.model_id - assert restored.id == persisted_message.id - - @staticmethod - @pytest.mark.parametrize("shape", ["current", "legacy_missing", "explicit_null"]) - def test_private_trace_task_preserves_provider_into_real_arize_span( - sqlite_session: Session, - persisted_message: Message, - monkeypatch: pytest.MonkeyPatch, - shape: str, - ) -> None: - message_data = persisted_message.to_dict(session=sqlite_session) - if shape == "legacy_missing": - message_data.pop("model_provider", None) - elif shape == "explicit_null": - message_data["model_provider"] = None - start = datetime(2026, 10, 20, tzinfo=UTC) - info = MessageTraceInfo( - message_id=persisted_message.id, - message_data=message_data, - conversation_model="chat", - conversation_mode="chat", - inputs={"query": "Synthetic input"}, - outputs="Synthetic output", - start_time=start, - end_time=start + timedelta(seconds=1), - metadata={}, - message_tokens=1, - answer_tokens=1, - total_tokens=2, - ) - encoded = ( - TaskData(app_id=persisted_message.app_id, trace_info_type="MessageTraceInfo", trace_info=info.model_dump()) - .model_dump_json() - .encode() - ) - exporter = InMemorySpanExporter() - provider = TracerProvider() - processor = SimpleSpanProcessor(exporter) - provider.add_span_processor(processor) - - def setup(_config: ArizeConfig) -> tuple[Tracer, SimpleSpanProcessor]: - tracer = provider.get_tracer("message-provider-serialization") - assert isinstance(tracer, Tracer) - return tracer, processor - - monkeypatch.setattr(arize_phoenix_trace, "setup_tracer", setup) - instance = ArizePhoenixDataTrace(ArizeConfig(api_key="unused-synthetic-key", space_id="synthetic-space")) - paths: list[str] = [] - deleted: list[str] = [] - - def load(path: str) -> bytes: - paths.append(path) - return encoded - - def delete(path: str) -> None: - deleted.append(path) - - def lookup(app_id: str) -> ArizePhoenixDataTrace: - assert app_id == persisted_message.app_id - return instance - - def failed(_key: str) -> None: - pytest.fail("The real trace consumer must not hide a failed dispatch") - - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) - monkeypatch.setattr(ops_trace_task.storage, "load", load) - monkeypatch.setattr(ops_trace_task.storage, "delete", delete) - monkeypatch.setattr(ops_trace_task.redis_client, "incr", failed) - monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", lookup) - monkeypatch.setattr("extensions.ext_enterprise_telemetry.is_enabled", lambda: False) - try: - ops_trace_task.process_trace_tasks.apply( - args=({"app_id": persisted_message.app_id, "file_id": "owned-synthetic"},), throw=True - ).get() - spans = exporter.get_finished_spans() - llm = [span for span in spans if span.name == "llm"] - assert len(llm) == 1 - attributes = llm[0].attributes - assert attributes is not None - assert attributes[SpanAttributes.LLM_MODEL_NAME] == persisted_message.model_id - if shape == "current": - assert attributes[SpanAttributes.LLM_PROVIDER] == persisted_message.model_provider - else: - assert SpanAttributes.LLM_PROVIDER not in attributes - assert paths == [f"ops_trace/{persisted_message.app_id}/owned-synthetic.json"] - assert deleted == paths - finally: - provider.shutdown() - - def _wrap_md(url: str) -> str: """ Wrap a raw URL into the markdown that re_sign_file_url_answer expects: diff --git a/api/tests/unit_tests/providers/trace/trace-aliyun/src/dify_trace_aliyun/test_configuration.py b/api/tests/unit_tests/providers/trace/trace-aliyun/src/dify_trace_aliyun/test_configuration.py deleted file mode 100644 index e679dc996b8c41..00000000000000 --- a/api/tests/unit_tests/providers/trace/trace-aliyun/src/dify_trace_aliyun/test_configuration.py +++ /dev/null @@ -1,276 +0,0 @@ -"""Pure Aliyun destination rules; these tests do not claim account verification.""" - -import traceback -from dataclasses import dataclass, field -from urllib.parse import unquote, urlsplit - -import httpx -import pytest -from dify_trace_aliyun.aliyun_trace import AliyunDataTrace -from dify_trace_aliyun.config import AliyunConfig -from dify_trace_aliyun.configuration import ( - AliyunConfigurationError, - console_url, - ingestion_url, - normalize_configuration, - verify_configuration, -) -from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest -from pydantic import ValidationError - -from core.helper import ssrf_proxy - - -def values(**overrides: object) -> dict[str, object]: - return {"endpoint": "http://tracing-analysis-dc-qd.aliyuncs.com", "license_key": "synthetic-key", **overrides} - - -def test_normalization_preserves_opaque_license_prefix_and_application_name() -> None: - submitted = values( - endpoint=" HTTPS://Project.CN-QINGDAO.LOG.ALIYUNCS.COM:443/apm/trace/opentelemetry/ ", - app_name=" 中文 agent_Name 01 ", - license_key="literal%2F@a+b/c", - ) - original = dict(submitted) - result = normalize_configuration(submitted) - assert result.endpoint == "https://project.cn-qingdao.log.aliyuncs.com:443/apm/trace/opentelemetry" - assert result.app_name == "中文 agent_Name 01" - assert result.license_key == "literal%2F@a+b/c" - assert submitted == original - assert ingestion_url(result.model_dump()) == ( - "https://project.cn-qingdao.log.aliyuncs.com:443/apm/trace/opentelemetry/" - "adapt_literal%252F@a%2Bb%2Fc/api/v1/traces" - ) - - -@pytest.mark.parametrize("key", ["raw%2F", "raw%252F", "a/b?c#d", "left@right_one+two", "密钥"]) -def test_ingestion_path_round_trips_the_raw_license_exactly_once(key: str) -> None: - url = ingestion_url(values(license_key=key)) - path = urlsplit(url).path - assert path.startswith("/adapt_") - assert path.endswith("/api/otlp/traces") - encoded = path.removeprefix("/adapt_").removesuffix("/api/otlp/traces") - assert "/" not in encoded - assert unquote(encoded) == key - - -@pytest.mark.parametrize("name", ["", " ", "dify_app", "应用 & Space"]) -def test_application_names_keep_aliyun_semantics(name: str) -> None: - assert normalize_configuration(values(app_name=name)).app_name == (name.strip() or "dify_app") - - -def test_legacy_default_endpoint_and_app_name_remain_available() -> None: - result = normalize_configuration(values(endpoint="")) - assert result.endpoint == "https://tracing-analysis-dc-hz.aliyuncs.com" - assert result.app_name == "dify_app" - - -@pytest.mark.parametrize( - ("base", "expected"), - [ - ("https://log.aliyuncs.com", "https://log.aliyuncs.com/adapt_synthetic-key/api/v1/traces"), - ( - "https://CN-QINGDAO.LOG.ALIYUNCS.COM.:443/prefix/", - "https://cn-qingdao.log.aliyuncs.com.:443/prefix/adapt_synthetic-key/api/v1/traces", - ), - ( - "https://log.aliyuncs.com.evil.test", - "https://log.aliyuncs.com.evil.test/adapt_synthetic-key/api/otlp/traces", - ), - ("https://evil-log.aliyuncs.com", "https://evil-log.aliyuncs.com/adapt_synthetic-key/api/otlp/traces"), - ( - "https://collector.test/log.aliyuncs.com", - "https://collector.test/log.aliyuncs.com/adapt_synthetic-key/api/otlp/traces", - ), - ("http://[::1]:4318/reverse/", "http://[::1]:4318/reverse/adapt_synthetic-key/api/otlp/traces"), - ], -) -def test_ingestion_family_uses_only_the_hostname(base: str, expected: str) -> None: - assert ingestion_url(values(endpoint=base)) == expected - - -@pytest.mark.parametrize( - "endpoint", - [ - None, - 1, - "ftp://host.test", - "https:///missing", - "https://user:secret@host.test", - "https://host.test?", - "https://host.test#", - "https://host.test/path?key=secret", - "https://host.test:bad", - "https://host.test:0", - "https://host.test:65536", - "https://host.test:", - "https://[invalid", - "\nhttps://host.test", - "https://ho\nst.test", - "https://host.test\\@elsewhere", - "https://host.test/a/../b", - "https://host.test/a/%2e%2e/b", - "https://host.test/a%2fb", - "https://host.test/%252e%252e", - "https://host.test/%zz", - "https://host.test/double//path", - "https://host.test/adapt_DO-NOT-PRINT/api/otlp/traces", - "https://host.test/reverse/adapt_DO-NOT-PRINT/api/v1/traces/", - "https://host.test/%61dapt_DO-NOT-PRINT", - "https://host.test/api/otlp/traces", - "https://host.test/v1/traces", - ], -) -def test_new_endpoints_reject_unsafe_or_already_completed_routes(endpoint: object) -> None: - with pytest.raises(AliyunConfigurationError) as error: - normalize_configuration(values(endpoint=endpoint)) - assert error.value.code == "invalid_endpoint" - assert "DO-NOT-PRINT" not in str(error.value) - assert "host.test" not in str(error.value) - - -@pytest.mark.parametrize("field", ["endpoint", "license_key"]) -def test_required_fields_are_not_invented(field: str) -> None: - submitted = values() - del submitted[field] - with pytest.raises(AliyunConfigurationError): - normalize_configuration(submitted) - - -@pytest.mark.parametrize("key", [None, 1, "", " ", " secret", "secret ", "secret\n", "secret\x00", "secret\ud800"]) -def test_invalid_license_is_rejected_without_echoing_it(key: object) -> None: - with pytest.raises(AliyunConfigurationError) as error: - normalize_configuration(values(license_key=key)) - assert error.value.code == "invalid_license_key" - assert "secret" not in str(error.value) - - -@pytest.mark.parametrize("name", [None, 1, "app\nname", "app\x00name"]) -def test_application_control_characters_are_rejected(name: object) -> None: - with pytest.raises(AliyunConfigurationError) as error: - normalize_configuration(values(app_name=name)) - assert error.value.code == "invalid_app_name" - - -def test_masks_survive_patch_normalization_but_cannot_be_sent() -> None: - result = normalize_configuration(values(license_key="***masked")) - assert result.license_key == "***masked" - with pytest.raises(AliyunConfigurationError): - ingestion_url(result.model_dump()) - - -def test_console_link_and_historical_model_do_not_require_new_write_validation() -> None: - old = AliyunConfig(license_key="synthetic-key", endpoint="https://host.test/adapt_old/api/otlp/traces") - assert old.endpoint == "https://host.test/adapt_old/api/otlp/traces" - assert console_url() == "https://arms.console.aliyun.com/#/llm" - - -def test_historical_model_validation_does_not_format_sensitive_endpoint_inputs() -> None: - secret = "SYNTHETIC_LICENSE_NOT_IN_ERRORS" - with pytest.raises(ValidationError) as caught: - AliyunConfig(license_key=secret, endpoint=f"ftp://host.test/adapt_{secret}/api/otlp/traces") - rendered = str(caught.value) + repr(caught.value) + "".join(traceback.format_exception(caught.value)) - assert secret not in rendered - - -@dataclass -class CollectorBoundary: - status: int = 200 - content: bytes = b'"success"' - failure: Exception | None = None - calls: list[tuple[str, str, dict[str, object]]] = field(default_factory=list) - responses: list[httpx.Response] = field(default_factory=list) - - -@pytest.fixture -def collector(monkeypatch: pytest.MonkeyPatch) -> CollectorBoundary: - boundary = CollectorBoundary() - - def send(method: str, url: str, **options: object) -> httpx.Response: - boundary.calls.append((method, url, dict(options))) - if boundary.failure is not None: - raise boundary.failure - response = httpx.Response(boundary.status, content=boundary.content, request=httpx.Request(method, url)) - boundary.responses.append(response) - return response - - def forbidden(*_args: object, **_kwargs: object) -> None: - raise AssertionError("Resource-only verification must not initialize a legacy exporter") - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - monkeypatch.setattr(AliyunDataTrace, "__init__", forbidden) - monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "https://ambient.test") - monkeypatch.setenv("OTEL_EXPORTER_OTLP_HEADERS", "Authorization=ambient-credential") - return boundary - - -def test_verification_emits_one_resource_and_zero_business_spans(collector: CollectorBoundary) -> None: - config = normalize_configuration(values(app_name="test_app", license_key="raw%2F@a+b/c")) - verify_configuration(config) - assert len(collector.calls) == 1 - method, url, options = collector.calls[0] - assert method == "POST" - assert url == "http://tracing-analysis-dc-qd.aliyuncs.com/adapt_raw%252F@a%2Bb%2Fc/api/otlp/traces" - assert options["headers"] == {"Content-Type": "application/x-protobuf", "Accept": "application/x-protobuf"} - assert options["max_retries"] == 0 - assert options["follow_redirects"] is False - body = options["content"] - assert isinstance(body, bytes) - assert body - payload = ExportTraceServiceRequest.FromString(body) - assert len(payload.resource_spans) == 1 - resource = payload.resource_spans[0] - assert len(resource.scope_spans) == 0 - assert len(resource.resource.attributes) == 1 - assert resource.resource.attributes[0].key == "service.name" - assert resource.resource.attributes[0].value.string_value == "test_app" - assert collector.responses[0].is_closed - - -@pytest.mark.parametrize( - ("status", "body"), - [ - (200, b""), - (200, b"{}"), - (200, b"success"), - (200, b'"success"\n'), - (200, b'{"ret":0}'), - (200, b'"failure"'), - (200, b"private-response"), - (201, b'"success"'), - (204, b""), - (400, b"private-response"), - (401, b"private-response"), - (403, b"private-response"), - (404, b"private-response"), - (405, b"private-response"), - (429, b"private-response"), - (503, b"private-response"), - ], -) -def test_only_the_observed_resource_ack_verifies_access( - collector: CollectorBoundary, status: int, body: bytes, caplog: pytest.LogCaptureFixture -) -> None: - collector.status = status - collector.content = body - config = normalize_configuration(values(license_key="SYNTHETIC_PRIVATE_LICENSE@key")) - with pytest.raises(AliyunConfigurationError) as error: - verify_configuration(config) - assert error.value.code == "verification_failed" - rendered = str(error.value) + repr(error.value) + "".join(traceback.format_exception(error.value)) + caplog.text - assert "SYNTHETIC_PRIVATE_LICENSE" not in rendered - assert "private-response" not in rendered - assert len(collector.calls) == 1 - assert all(response.is_closed for response in collector.responses) - - -def test_transport_failure_does_not_escape_with_private_url(collector: CollectorBoundary) -> None: - collector.failure = httpx.ConnectError("request failed at /adapt_PRIVATE%40KEY/api/otlp/traces") - config = normalize_configuration(values(license_key="PRIVATE@KEY")) - with pytest.raises(AliyunConfigurationError) as error: - verify_configuration(config) - assert error.value.code == "verification_failed" - rendered = "".join(traceback.format_exception(error.value)) - assert "PRIVATE@KEY" not in rendered - assert "PRIVATE%40KEY" not in rendered - assert len(collector.calls) == 1 diff --git a/api/tests/unit_tests/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/test_arize_configuration.py b/api/tests/unit_tests/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/test_arize_configuration.py deleted file mode 100644 index 63776e15227754..00000000000000 --- a/api/tests/unit_tests/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/test_arize_configuration.py +++ /dev/null @@ -1,581 +0,0 @@ -from collections import deque -from collections.abc import Iterator -from concurrent.futures import ThreadPoolExecutor -from threading import Event -from typing import Never - -import grpc # type: ignore[import-untyped] -import httpx -import pytest -from dify_trace_arize_phoenix import arize_configuration as configuration -from dify_trace_arize_phoenix.arize_configuration import ( - DESTINATION_KEY, - ArizeConfigurationError, - bind_destination, - project_url, - resolve_project_destination, - verify_configuration, -) -from dify_trace_arize_phoenix.arize_phoenix_trace import ArizePhoenixDataTrace -from dify_trace_arize_phoenix.config import ArizeConfig -from opentelemetry.context import _SUPPRESS_INSTRUMENTATION_KEY, get_value -from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( - ExportTraceServiceRequest, - ExportTraceServiceResponse, -) - -from core.helper import ssrf_proxy -from core.ops.ops_trace_manager import OpsTraceManager -from services.app_tracing_config_gateway import OpsTraceManagerGateway -from services.app_tracing_config_service import AppTracingConfigVerificationFailedError - -pytest_plugins = ("tests.unit_tests.core.ops.test_ops_trace_manager",) - -"""Arize metadata reads are optional, bounded and pinned to the exact destination.""" - - -SPACE = "U3BhY2U6dGVzdA==" -ORG = "T3JnOnRlc3Q=" -PROJECT = "UHJvamVjdDovKz0=" -VALUES = {"endpoint": "https://otlp.arize.com", "space_id": SPACE, "project": "exact", "api_key": "synthetic-key"} -DESTINATION = { - "endpoint": VALUES["endpoint"], - "space_id": SPACE, - "project": "exact", - "organization_id": ORG, - "project_id": PROJECT, -} -URL = "https://app.arize.com/organizations/T3JnOnRlc3Q%3D/spaces/U3BhY2U6dGVzdA%3D%3D/projects/UHJvamVjdDovKz0%3D" - - -def response(body: object, status: int = 200) -> httpx.Response: - return httpx.Response(status, json=body, request=httpx.Request("GET", "https://api.arize.com/v2/projects")) - - -def space() -> httpx.Response: - return response({"id": SPACE, "organization_id": ORG}) - - -def page(*projects: dict[str, str], more: bool = False, cursor: str | None = None) -> httpx.Response: - return response({"projects": list(projects), "pagination": {"has_more": more, "next_cursor": cursor}}) - - -def project(name: str = "exact", project_id: str = PROJECT) -> dict[str, str]: - return {"id": project_id, "name": name, "space_id": SPACE} - - -class TestArizeConfigurationCases: - @staticmethod - @pytest.fixture - def responses() -> deque[httpx.Response | Exception]: - return deque() - - @staticmethod - @pytest.fixture - def calls( - monkeypatch: pytest.MonkeyPatch, responses: deque[httpx.Response | Exception] - ) -> list[tuple[str, dict[str, object]]]: - calls: list[tuple[str, dict[str, object]]] = [] - - def get(url: str, **options: object) -> httpx.Response: - calls.append((url, options)) - assert responses, "Unplanned metadata request" - item = responses.popleft() - if isinstance(item, Exception): - raise item - return item - - monkeypatch.setattr(ssrf_proxy, "get", get) - return calls - - @staticmethod - def test_exact_scoped_project_and_opaque_ids_use_official_route( - calls: list[tuple[str, dict[str, object]]], responses: deque[httpx.Response | Exception] - ) -> None: - received = [space(), page(project("exact-substring"), more=True, cursor="next"), page(project())] - responses.extend(received) - assert resolve_project_destination(VALUES) == DESTINATION - assert ( - project_url( - { - **VALUES, - DESTINATION_KEY: bind_destination(VALUES, DESTINATION), - "project_url": "https://untrusted.test", - } - ) - == URL - ) - assert len(calls) == 3 - assert calls[0][0] == "https://api.arize.com/v2/spaces/U3BhY2U6dGVzdA%3D%3D" - assert calls[1][1]["params"] == {"space_id": SPACE, "name": "exact", "limit": 100} - assert calls[2][1]["params"] == {"space_id": SPACE, "name": "exact", "limit": 100, "cursor": "next"} - for url, options in calls: - assert url.startswith("https://api.arize.com/v2/") - assert options["headers"] == { - "Authorization": "Bearer synthetic-key", - "Accept": "application/json", - "Accept-Encoding": "identity", - "Cookie": "", - } - assert options["auth"] is None - assert options["follow_redirects"] is False - assert options["max_retries"] == 0 - assert options["stream_response"] is True - timeout = options["timeout"] - assert isinstance(timeout, float) - assert 0 < timeout <= 5 - assert all(item.is_closed for item in received) - - @staticmethod - @pytest.mark.parametrize("status", [301, 401, 403, 404, 429, 500]) - def test_unreadable_metadata_has_no_link_and_closes_response( - calls: list[tuple[str, dict[str, object]]], responses: deque[httpx.Response | Exception], status: int - ) -> None: - denied = response({"error": "private-response"}, status) - responses.append(denied) - assert resolve_project_destination(VALUES) is None - assert len(calls) == 1 - assert denied.is_closed - - @staticmethod - @pytest.mark.parametrize( - "body", [{"id": SPACE}, {"id": "other", "organization_id": ORG}, {"id": SPACE, "organization_id": "../bad"}] - ) - def test_missing_or_mismatched_space_metadata_is_not_a_project( - calls: list[tuple[str, dict[str, object]]], responses: deque[httpx.Response | Exception], body: dict[str, str] - ) -> None: - responses.append(response(body)) - assert resolve_project_destination(VALUES) is None - assert len(calls) == 1 - - @staticmethod - @pytest.mark.parametrize( - "body", - [ - {"projects": [project("EXACT")], "pagination": {"has_more": False}}, - {"projects": [project(), project(project_id="other")], "pagination": {"has_more": False}}, - {"projects": [{**project(), "space_id": "other"}], "pagination": {"has_more": False}}, - {"projects": [{**project(), "id": "../bad"}], "pagination": {"has_more": False}}, - {"projects": [project()] * 101, "pagination": {"has_more": False}}, - {"projects": [project()], "pagination": {"has_more": True}}, - {"projects": [project()], "pagination": {"has_more": "false"}}, - {"projects": [project()]}, - {"data": [project()], "pagination": {"has_more": False}}, - ], - ) - def test_incomplete_ambiguous_or_non_exact_projects_do_not_fabricate_links( - calls: list[tuple[str, dict[str, object]]], responses: deque[httpx.Response | Exception], body: object - ) -> None: - responses.extend([space(), response(body)]) - assert resolve_project_destination(VALUES) is None - assert len(calls) == 2 - - @staticmethod - @pytest.mark.parametrize("repeat", [False, True]) - def test_pagination_is_bounded_and_never_accepts_a_partial_match( - calls: list[tuple[str, dict[str, object]]], responses: deque[httpx.Response | Exception], repeat: bool - ) -> None: - responses.append(space()) - responses.extend(page(project(), more=True, cursor="same" if repeat else str(index)) for index in range(3)) - assert resolve_project_destination(VALUES) is None - assert len(calls) == (3 if repeat else 4) - - @staticmethod - @pytest.mark.parametrize( - "item", - [ - httpx.ReadTimeout("synthetic-timeout"), - httpx.Response(200, content=b"not-json"), - httpx.Response(200, json=[]), - httpx.Response(200, content=b" " * (1024 * 1024 + 1)), - httpx.Response(200, json={}, headers={"content-encoding": "unsupported"}), - ], - ) - def test_transport_and_bounded_decode_failures_remain_optional( - calls: list[tuple[str, dict[str, object]]], - responses: deque[httpx.Response | Exception], - item: httpx.Response | Exception, - ) -> None: - responses.append(item) - assert resolve_project_destination(VALUES) is None - assert len(calls) == 1 - if isinstance(item, httpx.Response): - assert item.is_closed - - @staticmethod - def test_all_pages_share_one_deadline(monkeypatch: pytest.MonkeyPatch) -> None: - clock = [100.0] - calls: list[float] = [] - received = [space(), page(project(), more=True, cursor="next")] - - def get(_url: str, **options: object) -> httpx.Response: - timeout = options["timeout"] - assert isinstance(timeout, float) - calls.append(timeout) - clock[0] += 2.5 - return received[len(calls) - 1] - - monkeypatch.setattr(configuration.time, "monotonic", lambda: clock[0]) - monkeypatch.setattr(ssrf_proxy, "get", get) - assert resolve_project_destination(VALUES) is None - assert calls == [5.0, 2.5] - assert all(item.is_closed for item in received) - - @staticmethod - def test_slow_response_chunks_stop_at_the_shared_deadline(monkeypatch: pytest.MonkeyPatch) -> None: - clock = [100.0] - chunks: list[int] = [] - received = httpx.Response(200) - - def stream() -> Iterator[bytes]: - for index in range(10): - clock[0] += 2.0 - chunks.append(index) - yield b" " - - monkeypatch.setattr(configuration.time, "monotonic", lambda: clock[0]) - monkeypatch.setattr(received, "iter_bytes", stream) - monkeypatch.setattr(ssrf_proxy, "get", lambda _url, **_options: received) - assert resolve_project_destination(VALUES) is None - assert chunks == [0, 1, 2] - assert received.is_closed - - @staticmethod - @pytest.mark.parametrize( - "change", - [ - {"endpoint": "https://custom.test"}, - {"endpoint": "https://otlp.eu-west-1a.arize.com"}, - {"endpoint": "http://otlp.arize.com"}, - {"endpoint": "https://user:secret@otlp.arize.com"}, - {"endpoint": "https://otlp.arize.com:444"}, - {"space_id": "../bad"}, - {"api_key": None}, - {"api_key": ""}, - {"api_key": "******"}, - {"api_key": "key\nheader"}, - ], - ) - def test_unsupported_destination_or_credentials_do_not_send_to_an_assumed_host( - calls: list[tuple[str, dict[str, object]]], change: dict[str, str | None] - ) -> None: - assert resolve_project_destination({**VALUES, **change}) is None - assert calls == [] - - @staticmethod - @pytest.mark.parametrize( - "change", - [ - {"endpoint": "https://otlp.arize.com:443"}, - {"space_id": "other"}, - {"project": "later"}, - {DESTINATION_KEY: {**DESTINATION, "organization_id": "bad?query"}}, - {DESTINATION_KEY: {**DESTINATION, "project_id": "../bad"}}, - {DESTINATION_KEY: None}, - ], - ) - def test_get_formatter_rejects_stale_or_forged_metadata_without_io( - monkeypatch: pytest.MonkeyPatch, change: dict[str, object] - ) -> None: - def forbidden(*_args: object, **_kwargs: object) -> Never: - raise AssertionError("GET must not call the vendor") - - monkeypatch.setattr(ssrf_proxy, "get", forbidden) - assert project_url({**VALUES, DESTINATION_KEY: bind_destination(VALUES, DESTINATION), **change}) is None - - -"""A real gRPC collector must acknowledge Arize credentials before config persistence.""" - - -class TestArizeConfigurationAckCases: - @staticmethod - @pytest.fixture - def collector() -> Iterator[tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]]]: - calls: list[tuple[bytes, dict[str, str]]] = [] - state: dict[str, object] = {"status": grpc.StatusCode.OK} - release = Event() - state["release"] = release - - def export(request: ExportTraceServiceRequest, context: grpc.ServicerContext) -> ExportTraceServiceResponse: - headers = { - key: value.decode() if isinstance(value, bytes) else value - for key, value in context.invocation_metadata() - } - calls.append((request.SerializeToString(), headers)) - if headers.get("api_key") != "synthetic-good" or headers.get("space_id") != "synthetic-space": - context.abort(grpc.StatusCode.PERMISSION_DENIED, "private-vendor-details") - code = state["status"] - if isinstance(code, grpc.StatusCode) and code is not grpc.StatusCode.OK: - context.abort(code, "private-vendor-details") - if state.get("delay"): - release.wait(2) - result = ExportTraceServiceResponse() - if state.get("partial"): - result.partial_success.rejected_spans = 1 - if state.get("message"): - result.partial_success.error_message = "private-vendor-details" - return result - - def dispatch(*args: object, **kwargs: object) -> ExportTraceServiceResponse: - assert kwargs == {} - request, context = args - assert isinstance(request, ExportTraceServiceRequest) - assert isinstance(context, grpc.ServicerContext) - return export(request, context) - - handler = grpc.unary_unary_rpc_method_handler( - dispatch, - request_deserializer=ExportTraceServiceRequest.FromString, - response_serializer=ExportTraceServiceResponse.SerializeToString, - ) - with ThreadPoolExecutor(max_workers=2) as pool: - server = grpc.server(pool) - server.add_generic_rpc_handlers( - ( - grpc.method_handlers_generic_handler( - "opentelemetry.proto.collector.trace.v1.TraceService", {"Export": handler} - ), - ) - ) - port = server.add_insecure_port("127.0.0.1:0") - assert port > 0 - server.start() - try: - yield f"http://127.0.0.1:{port}", calls, state - finally: - release.set() - server.stop(0).wait(2) - - @staticmethod - def test_rejected_real_collector_ack_cannot_be_saved( - workspace: str, - collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], - ) -> None: - endpoint, calls, _ = collector - with pytest.raises(AppTracingConfigVerificationFailedError): - OpsTraceManagerGateway().prepare_new_config( - workspace_id=workspace, - tracing_provider="arize", - tracing_config={ - "endpoint": endpoint, - "project": "synthetic-project", - "space_id": "synthetic-space", - "api_key": "synthetic-invalid", - }, - ) - assert len(calls) == 1 - - @staticmethod - def test_synchronous_ack_is_resource_only_and_has_no_sdk_or_environment_side_effects( - collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], - monkeypatch: pytest.MonkeyPatch, - ) -> None: - endpoint, calls, _ = collector - attempted: list[bool] = [] - - def forbidden(*_args: object, **_kwargs: object) -> Never: - attempted.append(True) - raise AssertionError("Arize verification must not construct the SDK tracer") - - monkeypatch.setattr(ArizePhoenixDataTrace, "__init__", forbidden) - monkeypatch.setenv("OTEL_EXPORTER_OTLP_HEADERS", "api_key=ambient-invalid,space_id=ambient-other") - verify_configuration( - ArizeConfig( - endpoint=endpoint, api_key="synthetic-good", space_id="synthetic-space", project="synthetic-project" - ) - ) - assert attempted == [] - assert len(calls) == 1 - payload = ExportTraceServiceRequest.FromString(calls[0][0]) - assert len(payload.resource_spans) == 1 - assert {a.key: a.value.string_value for a in payload.resource_spans[0].resource.attributes} == { - "model_id": "synthetic-project", - "openinference.project.name": "synthetic-project", - } - assert len(payload.resource_spans[0].scope_spans) == 0 - assert calls[0][1]["api_key"] == "synthetic-good" - assert calls[0][1]["space_id"] == "synthetic-space" - assert calls[0][1]["authorization"] == "Bearer synthetic-good" - assert "ambient" not in str(calls) - - @staticmethod - @pytest.mark.parametrize( - ("code", "expected"), - [ - (grpc.StatusCode.UNAUTHENTICATED, "authentication_failed"), - (grpc.StatusCode.PERMISSION_DENIED, "permission_denied"), - (grpc.StatusCode.UNAVAILABLE, "connection_failed"), - (grpc.StatusCode.RESOURCE_EXHAUSTED, "connection_failed"), - ], - ) - def test_real_rpc_rejection_uses_only_safe_fixed_diagnostics( - collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], - code: grpc.StatusCode, - expected: str, - caplog: pytest.LogCaptureFixture, - ) -> None: - endpoint, calls, state = collector - state["status"] = code - with pytest.raises(ArizeConfigurationError) as caught: - verify_configuration(ArizeConfig(endpoint=endpoint, api_key="synthetic-good", space_id="synthetic-space")) - assert caught.value.code == expected - assert caught.value.__suppress_context__ - assert "synthetic-good" not in str(caught.value) - assert "private-vendor-details" not in str(caught.value) - assert "private-vendor-details" not in caplog.text - assert len(calls) == 1 - - @staticmethod - @pytest.mark.parametrize("partial", ["partial", "message"]) - def test_partial_or_error_ack_is_not_success( - collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], - partial: str, - ) -> None: - endpoint, calls, state = collector - state[partial] = True - with pytest.raises(ArizeConfigurationError, match="unexpected_response"): - verify_configuration(ArizeConfig(endpoint=endpoint, api_key="synthetic-good", space_id="synthetic-space")) - assert len(calls) == 1 - - @staticmethod - def test_deadline_is_finite_and_does_not_retry( - collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], - monkeypatch: pytest.MonkeyPatch, - ) -> None: - endpoint, calls, state = collector - state["delay"] = True - monkeypatch.setattr(configuration, "_VERIFY_TIMEOUT_SECONDS", 0.05) - with pytest.raises(ArizeConfigurationError, match="connection_failed"): - verify_configuration(ArizeConfig(endpoint=endpoint, api_key="synthetic-good", space_id="synthetic-space")) - assert len(calls) == 1 - - @staticmethod - @pytest.mark.parametrize("reject", [False, True]) - def test_tls_channel_is_explicit_suppressed_and_closed_on_all_outcomes( - collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], - monkeypatch: pytest.MonkeyPatch, - reject: bool, - ) -> None: - endpoint, _, state = collector - if reject: - state["status"] = grpc.StatusCode.PERMISSION_DENIED - original = grpc.insecure_channel - options_seen: list[tuple[str, object, object]] = [] - closed: list[bool] = [] - - def secure(target: str, credentials: grpc.ChannelCredentials, options: object) -> grpc.Channel: - assert isinstance(credentials, grpc.ChannelCredentials) - options_seen.append((target, options, get_value(_SUPPRESS_INSTRUMENTATION_KEY))) - channel = original(endpoint.removeprefix("http://")) - close = channel.close - - def close_recording() -> None: - closed.append(True) - close() - - monkeypatch.setattr(channel, "close", close_recording) - return channel - - monkeypatch.setattr(grpc, "secure_channel", secure) - before = get_value(_SUPPRESS_INSTRUMENTATION_KEY) - config = ArizeConfig(api_key="synthetic-good", space_id="synthetic-space") - if reject: - with pytest.raises(ArizeConfigurationError): - verify_configuration(config) - else: - verify_configuration(config) - assert options_seen == [ - ("otlp.arize.com:443", (("grpc.enable_retries", 0), ("grpc.max_receive_message_length", 1024 * 1024)), True) - ] - assert closed == [True] - assert get_value(_SUPPRESS_INSTRUMENTATION_KEY) == before - - @staticmethod - @pytest.mark.parametrize( - ("field", "value"), - [ - ("api_key", None), - ("api_key", ""), - ("api_key", "******"), - ("api_key", "key\nheader"), - ("space_id", None), - ("space_id", " "), - ("space_id", "******"), - ("endpoint", "https://user:secret@otlp.arize.com"), - ("endpoint", "https://otlp.arize.com:0"), - ], - ) - def test_invalid_auth_values_never_construct_a_channel( - monkeypatch: pytest.MonkeyPatch, - field: str, - value: str | None, - ) -> None: - attempted: list[bool] = [] - - def forbidden(*_args: object, **_kwargs: object) -> Never: - attempted.append(True) - raise AssertionError("Invalid configuration must be rejected before transport construction") - - monkeypatch.setattr(grpc, "secure_channel", forbidden) - monkeypatch.setattr(grpc, "insecure_channel", forbidden) - config = ArizeConfig.model_validate({"api_key": "synthetic-good", "space_id": "synthetic-space", field: value}) - with pytest.raises(ArizeConfigurationError, match="invalid_configuration"): - verify_configuration(config) - assert attempted == [] - - @staticmethod - @pytest.mark.parametrize("accepted", [False, True]) - def test_legacy_api_check_also_requires_the_real_ack( - collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], - accepted: bool, - ) -> None: - endpoint, calls, _ = collector - instance = ArizePhoenixDataTrace( - ArizeConfig( - endpoint=endpoint, - api_key="synthetic-good" if accepted else "synthetic-invalid", - space_id="synthetic-space", - ) - ) - try: - if accepted: - assert instance.api_check() - else: - with pytest.raises(ArizeConfigurationError, match="permission_denied"): - instance.api_check() - finally: - instance.processor.shutdown() - assert len(calls) == 1 - assert len(ExportTraceServiceRequest.FromString(calls[0][0]).resource_spans[0].scope_spans) == 0 - - @staticmethod - def test_real_ingest_ack_without_rest_read_permission_can_still_save( - workspace: str, - collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], - monkeypatch: pytest.MonkeyPatch, - ) -> None: - endpoint, rpc_calls, _ = collector - original = grpc.insecure_channel - reads: list[str] = [] - - def secure(_target: str, _credentials: grpc.ChannelCredentials, **_options: object) -> grpc.Channel: - return original(endpoint.removeprefix("http://")) - - def get(url: str, **_options: object) -> httpx.Response: - reads.append(url) - return httpx.Response(403, json={"message": "private-vendor-details"}) - - monkeypatch.setattr(grpc, "secure_channel", secure) - monkeypatch.setattr(ssrf_proxy, "get", get) - saved = OpsTraceManagerGateway().prepare_new_config( - workspace_id=workspace, - tracing_provider="arize", - tracing_config={"api_key": "synthetic-good", "space_id": "synthetic-space", "project": "not-created"}, - ) - assert len(rpc_calls) == len(reads) == 1 - assert reads[0] == "https://api.arize.com/v2/spaces/synthetic-space" - assert "project_url" not in saved - assert "arize_project_destination" not in saved - assert OpsTraceManager.decrypt_tracing_config(workspace, "arize", saved)["api_key"] == "synthetic-good" - payload = ExportTraceServiceRequest.FromString(rpc_calls[0][0]) - assert len(payload.resource_spans[0].scope_spans) == 0 diff --git a/api/tests/unit_tests/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/test_configuration.py b/api/tests/unit_tests/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/test_configuration.py deleted file mode 100644 index 032b511aa25bf4..00000000000000 --- a/api/tests/unit_tests/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/test_configuration.py +++ /dev/null @@ -1,336 +0,0 @@ -"""Phoenix configuration verifies only a bounded project-list read, with isolated auth.""" - -import os -from collections import deque -from collections.abc import Iterator -from concurrent.futures import ThreadPoolExecutor -from threading import Barrier -from typing import Never - -import httpx -import pytest -from dify_trace_arize_phoenix.arize_phoenix_trace import ArizePhoenixDataTrace -from dify_trace_arize_phoenix.config import PhoenixConfig -from dify_trace_arize_phoenix.configuration import ( - PhoenixConfigurationError, - normalize_configuration, - project_url, - verify_configuration, -) - -from core.helper import ssrf_proxy -from core.ops.ops_trace_manager import OpsTraceManager - - -def response(body: object, status: int = 200) -> httpx.Response: - return httpx.Response(status, json=body, request=httpx.Request("GET", "https://phoenix.test/v1/projects")) - - -@pytest.fixture(autouse=True) -def no_legacy_or_environment_auth(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: - names = ("PHOENIX_API_KEY", "PHOENIX_PROJECT_NAME", "PHOENIX_COLLECTOR_ENDPOINT", "OTEL_EXPORTER_OTLP_HEADERS") - for name in names: - monkeypatch.setenv(name, "synthetic-ambient-value") - called: list[bool] = [] - - def forbidden(*_args: object, **_kwargs: object) -> Never: - called.append(True) - raise AssertionError("Phoenix configuration must not initialize a tracer or write test spans") - - monkeypatch.setattr(ArizePhoenixDataTrace, "__init__", forbidden) - monkeypatch.setattr(OpsTraceManager, "check_trace_config_is_effective", forbidden) - monkeypatch.setattr(OpsTraceManager, "get_trace_config_project_url", forbidden) - for method in ("post", "put", "patch", "delete"): - monkeypatch.setattr(ssrf_proxy, method, forbidden) - yield - assert called == [] - assert all(os.environ[name] == "synthetic-ambient-value" for name in names) - - -@pytest.fixture -def api_responses() -> deque[httpx.Response | Exception]: - return deque() - - -@pytest.fixture -def api_calls( - monkeypatch: pytest.MonkeyPatch, api_responses: deque[httpx.Response | Exception] -) -> list[tuple[str, dict[str, object]]]: - calls: list[tuple[str, dict[str, object]]] = [] - - def get(url: str, **options: object) -> httpx.Response: - calls.append((url, options)) - result = api_responses.popleft() if api_responses else response({"data": [], "next_cursor": None}) - if isinstance(result, Exception): - raise result - return result - - monkeypatch.setattr(ssrf_proxy, "get", get) - return calls - - -def test_normalization_preserves_prefix_and_does_not_mutate_input( - api_calls: list[tuple[str, dict[str, object]]], -) -> None: - values = {"endpoint": " HTTP://Phoenix.TEST:6006/proxy/s/space/ ", "project": " project name ", "api_key": None} - untouched = dict(values) - config = normalize_configuration(values) - result = verify_configuration(config) - assert (result.endpoint, result.project) == ("http://phoenix.test:6006/proxy/s/space", "project name") - assert values == untouched - assert api_calls == [ - ( - "http://phoenix.test:6006/proxy/s/space/v1/projects", - { - "params": {"limit": 1}, - "headers": {"Accept": "application/json", "Accept-Encoding": "identity", "Cookie": ""}, - "auth": None, - "timeout": 30.0, - "max_retries": 0, - "follow_redirects": False, - "stream_response": True, - }, - ) - ] - - -@pytest.mark.parametrize("key", [None, "", "synthetic-key"]) -def test_optional_key_does_not_inherit_environment( - api_calls: list[tuple[str, dict[str, object]]], - key: str | None, -) -> None: - result = verify_configuration(PhoenixConfig(endpoint="https://phoenix.test", project="not-created", api_key=key)) - assert result.project == "not-created" - expected = {"Accept": "application/json", "Accept-Encoding": "identity", "Cookie": ""} - if key: - expected["Authorization"] = "Bearer synthetic-key" - assert api_calls[0][1]["headers"] == expected - assert len(api_calls) == 1 - - -@pytest.mark.parametrize("project", [None, "", " "]) -def test_empty_project_uses_default(project: str | None, api_calls: list[tuple[str, dict[str, object]]]) -> None: - config = normalize_configuration({"project": project}) - assert config.project == "default" - assert project_url(config) == "https://app.phoenix.arize.com/redirects/projects/default" - assert api_calls == [] - - -def test_project_redirect_encodes_name_for_graphql_loader_without_io( - api_calls: list[tuple[str, dict[str, object]]], -) -> None: - config = PhoenixConfig(endpoint="https://host.test:8443/s/space/", project="A /? & #中文") - assert ( - project_url(config) - == "https://host.test:8443/s/space/redirects/projects/A%20%2F%3F%20%26%20%23%E4%B8%AD%E6%96%87" - ) - assert api_calls == [] - - -@pytest.mark.parametrize( - "endpoint", - [ - "ftp://host.test", - "https:///missing", - "https://user:secret@host.test", - "https://host.test?", - "https://host.test#", - "https://host.test/path?token=secret", - "https://host.test/path#fragment", - "https://host.test:bad", - "https://host.test:0", - "https://host.test:65536", - "https://host.test:", - "https://[invalid", - "https://host.test\\@elsewhere", - "\nhttps://host.test", - "https://ho\nst.test", - "https://host.test/a/../b", - "https://host.test/a/%2e%2e/b", - "https://host.test/a%2fb", - "https://host.test/a%5cb", - "https://host.test/a%252fb", - "https://host.test//proxy", - "https://host.test/%00", - "https://host.test/%zz", - "https://host.test/v1/traces", - "https://host.test/s/space/v1/traces/", - None, - ], -) -def test_invalid_base_is_rejected_before_io(endpoint: object, api_calls: list[tuple[str, dict[str, object]]]) -> None: - with pytest.raises(PhoenixConfigurationError) as caught: - normalize_configuration({"endpoint": endpoint, "api_key": "private-secret"}) - assert caught.value.code == "invalid_endpoint" - assert "private-secret" not in str(caught.value) - assert api_calls == [] - - -@pytest.mark.parametrize("project", ["bad\nname", "\tname", ".", "..", 123]) -def test_invalid_project_is_rejected_before_io(project: object, api_calls: list[tuple[str, dict[str, object]]]) -> None: - with pytest.raises(PhoenixConfigurationError) as caught: - normalize_configuration({"project": project}) - assert caught.value.code == "invalid_project" - assert api_calls == [] - - -@pytest.mark.parametrize("key", [" ", "line\nbreak", "invalid token", "******", "ab***cd"]) -def test_header_or_mask_key_is_rejected_after_restore(key: str, api_calls: list[tuple[str, dict[str, object]]]) -> None: - config = normalize_configuration({"api_key": key}) - with pytest.raises(PhoenixConfigurationError) as caught: - verify_configuration(config) - assert caught.value.code == "invalid_api_key" - assert api_calls == [] - - -@pytest.mark.parametrize( - ("status", "code"), - [ - (401, "authentication_failed"), - (403, "permission_denied"), - (302, "endpoint_unavailable"), - (404, "endpoint_unavailable"), - (500, "endpoint_unavailable"), - (204, "endpoint_unavailable"), - ], -) -def test_http_failures_are_safe_and_close_response( - status: int, - code: str, - api_responses: deque[httpx.Response | Exception], - api_calls: list[tuple[str, dict[str, object]]], -) -> None: - item = response({"private-vendor-detail": "private-secret"}, status) - api_responses.append(item) - with pytest.raises(PhoenixConfigurationError) as caught: - verify_configuration(PhoenixConfig(api_key="private-secret")) - assert caught.value.code == code - assert "private-secret" not in str(caught.value) - assert item.is_closed - assert len(api_calls) == 1 - - -@pytest.mark.parametrize( - "body", - [ - None, - [], - {}, - {"data": []}, - {"data": {}, "next_cursor": None}, - {"data": [], "next_cursor": 3}, - {"data": [None], "next_cursor": None}, - {"data": [{"id": "1"}], "next_cursor": None}, - {"data": [{"id": 1, "name": "x"}], "next_cursor": None}, - {"data": [{"id": "1", "name": ""}], "next_cursor": None}, - {"data": [{"id": "1", "name": "x", "description": {}}], "next_cursor": None}, - ], -) -def test_project_list_shape_is_validated( - body: object, - api_responses: deque[httpx.Response | Exception], - api_calls: list[tuple[str, dict[str, object]]], -) -> None: - item = response(body) - api_responses.append(item) - with pytest.raises(PhoenixConfigurationError) as caught: - verify_configuration(PhoenixConfig()) - assert caught.value.code == "malformed_response" - assert item.is_closed - assert len(api_calls) == 1 - - -def test_existing_project_page_does_not_require_selected_project_or_follow_cursor( - api_responses: deque[httpx.Response | Exception], - api_calls: list[tuple[str, dict[str, object]]], -) -> None: - api_responses.append( - response({"data": [{"id": "Project:1", "name": "other", "description": None}], "next_cursor": "opaque-cursor"}) - ) - assert verify_configuration(PhoenixConfig(project="new-project")).project == "new-project" - assert len(api_calls) == 1 - - -@pytest.mark.parametrize("kind", ["html", "gzip", "oversized"]) -def test_wrong_path_html_compressed_and_oversized_responses_are_not_success( - kind: str, - api_responses: deque[httpx.Response | Exception], - api_calls: list[tuple[str, dict[str, object]]], -) -> None: - content = b"x" * (1024 * 1024 + 1) if kind == "oversized" else b"private-login-page" - item = httpx.Response( - 200, content=content, request=httpx.Request("GET", "https://host.test/not-a-collector/v1/projects") - ) - if kind == "gzip": - item.headers["Content-Encoding"] = "gzip" - api_responses.append(item) - with pytest.raises(PhoenixConfigurationError) as caught: - verify_configuration(PhoenixConfig(endpoint="https://host.test/not-a-collector")) - assert caught.value.code == "malformed_response" - assert "private-login-page" not in str(caught.value) - assert item.is_closed - assert len(api_calls) == 1 - - -def test_transport_error_does_not_leak_url_or_credentials( - api_responses: deque[httpx.Response | Exception], - api_calls: list[tuple[str, dict[str, object]]], -) -> None: - api_responses.append(httpx.ReadTimeout("private-key in vendor response")) - with pytest.raises(PhoenixConfigurationError) as caught: - verify_configuration(PhoenixConfig(api_key="private-key")) - assert caught.value.code == "connection_failed" - assert "private-key" not in str(caught.value) - assert caught.value.__suppress_context__ - assert len(api_calls) == 1 - - -def test_concurrent_targets_and_credentials_remain_isolated(monkeypatch: pytest.MonkeyPatch) -> None: - barrier = Barrier(2) - requests: list[httpx.Request] = [] - targets = { - "https://first.test/s/one/v1/projects?limit=1": "first-key", - "https://second.test/reverse/v1/projects?limit=1": "second-key", - } - - def handle(request: httpx.Request) -> httpx.Response: - requests.append(request) - assert request.method == "GET" - assert request.headers["Authorization"] == "Bearer " + targets[str(request.url)] - barrier.wait(timeout=5) - return httpx.Response(200, json={"data": [], "next_cursor": None}) - - with httpx.Client(transport=httpx.MockTransport(handle), trust_env=False) as client: - monkeypatch.setattr(ssrf_proxy, "_get_ssrf_client", lambda _ssl_verify_enabled: client) - configs = [ - PhoenixConfig(endpoint="https://first.test/s/one", api_key="first-key", project="one"), - PhoenixConfig(endpoint="https://second.test/reverse", api_key="second-key", project="two"), - ] - with ThreadPoolExecutor(max_workers=2) as pool: - assert [item.project for item in pool.map(verify_configuration, configs)] == ["one", "two"] - assert sorted(str(request.url) for request in requests) == sorted(targets) - - -def test_real_proxy_client_has_no_ambient_auth_or_cookie_fallback(monkeypatch: pytest.MonkeyPatch) -> None: - requests: list[httpx.Request] = [] - - def handle(request: httpx.Request) -> httpx.Response: - requests.append(request) - assert str(request.url) == "https://host.test/s/space/v1/projects?limit=1" - assert request.headers["Cookie"] == "" - return httpx.Response( - 200, headers={"Set-Cookie": "session=vendor; Path=/"}, json={"data": [], "next_cursor": None} - ) - - with httpx.Client( - transport=httpx.MockTransport(handle), - auth=httpx.BasicAuth("ambient", "secret"), - cookies={"session": "ambient"}, - trust_env=False, - ) as client: - monkeypatch.setattr(ssrf_proxy, "_get_ssrf_client", lambda _ssl_verify_enabled: client) - verify_configuration(PhoenixConfig(endpoint="https://host.test/s/space", api_key="explicit-key")) - verify_configuration(PhoenixConfig(endpoint="https://host.test/s/space", api_key=None)) - assert requests[0].headers["Authorization"] == "Bearer explicit-key" - assert "Authorization" not in requests[1].headers - assert len(requests) == 2 diff --git a/api/tests/unit_tests/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/test_unified_trace.py b/api/tests/unit_tests/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/test_unified_trace.py deleted file mode 100644 index 6f45ec277dd003..00000000000000 --- a/api/tests/unit_tests/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/test_unified_trace.py +++ /dev/null @@ -1,258 +0,0 @@ -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/trace-langsmith/src/dify_trace_langsmith/test_unified_trace.py b/api/tests/unit_tests/providers/trace/trace-langsmith/src/dify_trace_langsmith/test_unified_trace.py deleted file mode 100644 index da3532d316c386..00000000000000 --- a/api/tests/unit_tests/providers/trace/trace-langsmith/src/dify_trace_langsmith/test_unified_trace.py +++ /dev/null @@ -1,328 +0,0 @@ -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/providers/trace/trace-mlflow/src/dify_trace_mlflow/test_configuration.py b/api/tests/unit_tests/providers/trace/trace-mlflow/src/dify_trace_mlflow/test_configuration.py deleted file mode 100644 index e16ff2891e905e..00000000000000 --- a/api/tests/unit_tests/providers/trace/trace-mlflow/src/dify_trace_mlflow/test_configuration.py +++ /dev/null @@ -1,230 +0,0 @@ -"""Official experiment schemas at a request-scoped, read-only HTTP boundary.""" - -from collections.abc import Callable -from dataclasses import dataclass, field - -import httpx -import pytest -from dify_trace_mlflow.config import DatabricksConfig -from dify_trace_mlflow.configuration import DatabricksConfigurationError, project_url, verify_configuration -from google.protobuf.json_format import MessageToDict -from mlflow.protos.service_pb2 import Experiment, GetExperiment - -from core.helper import ssrf_proxy - - -def response(url: str, body: object, status: int = 200) -> httpx.Response: - return httpx.Response(status, json=body, request=httpx.Request("GET", url)) - - -@dataclass -class ExperimentAPI: - calls: list[tuple[str, str, dict[str, object]]] = field(default_factory=list) - override: Callable[[str, str, dict[str, object]], httpx.Response] | None = None - - def get(self, url: str, **options: object) -> httpx.Response: - self.calls.append(("GET", url, options)) - if self.override: - return self.override("GET", url, options) - assert url.endswith("/api/2.0/mlflow/experiments/get") - params = options["params"] - assert isinstance(params, dict) - requested = GetExperiment(experiment_id=params["experiment_id"]) - return response( - url, - MessageToDict( - GetExperiment.Response( - experiment=Experiment(experiment_id=requested.experiment_id, lifecycle_stage="active") - ), - preserving_proto_field_name=True, - ), - ) - - def post(self, url: str, **options: object) -> httpx.Response: - self.calls.append(("POST", url, options)) - assert url.endswith("/oidc/v1/token"), "configuration must never create experiments, runs or traces" - if self.override: - return self.override("POST", url, options) - return response(url, {"access_token": "synthetic-oauth-token", "token_type": "Bearer", "expires_in": 3600}) - - -@pytest.fixture -def experiment_api(monkeypatch: pytest.MonkeyPatch) -> ExperimentAPI: - api = ExperimentAPI() - monkeypatch.setattr(ssrf_proxy, "get", api.get) - monkeypatch.setattr(ssrf_proxy, "post", api.post) - return api - - -def test_pat_requests_are_isolated_and_read_only(experiment_api: ExperimentAPI) -> None: - for host, token, experiment_id in [ - ("https://first.test/", "first-pat", "000123"), - ("https://second.test", "second-pat", "456"), - ]: - config = DatabricksConfig(host=host, experiment_id=experiment_id, personal_access_token=token) - result = verify_configuration(config) - assert result.host == host.rstrip("/") - assert result.experiment_id == experiment_id.lstrip("0") - method, url, options = experiment_api.calls[-1] - assert method == "GET" - assert url == host.rstrip("/") + "/api/2.0/mlflow/experiments/get" - assert options["headers"] == {"Authorization": f"Bearer {token}", "Accept": "application/json"} - assert options["max_retries"] == 0 - assert options["timeout"] == 30.0 - assert options["follow_redirects"] is False - assert len(experiment_api.calls) == 2 - - -def test_oauth_precedence_matches_exporter_without_reusing_pat(experiment_api: ExperimentAPI) -> None: - config = DatabricksConfig( - host="https://workspace.test", - experiment_id="12", - personal_access_token="unused-pat", - client_id="synthetic-client", - client_secret="synthetic-secret", - ) - assert verify_configuration(config).experiment_id == "12" - method, url, options = experiment_api.calls[0] - assert method == "POST" - assert url == "https://workspace.test/oidc/v1/token" - assert options["auth"] == ("synthetic-client", "synthetic-secret") - assert options["data"] == {"grant_type": "client_credentials", "scope": "all-apis"} - assert "headers" not in options - assert options["follow_redirects"] is False - assert experiment_api.calls[1][2]["headers"] == { - "Authorization": "Bearer synthetic-oauth-token", - "Accept": "application/json", - } - - -def test_incomplete_oauth_uses_explicit_pat_as_existing_exporter_does(experiment_api: ExperimentAPI) -> None: - verify_configuration( - DatabricksConfig( - host="https://workspace.test", experiment_id="12", client_id="unused-client", personal_access_token="pat" - ) - ) - assert len(experiment_api.calls) == 1 - assert experiment_api.calls[0][0] == "GET" - assert experiment_api.calls[0][2]["headers"] == {"Authorization": "Bearer pat", "Accept": "application/json"} - - -def test_non_ascii_experiment_ids_are_not_sent(experiment_api: ExperimentAPI) -> None: - with pytest.raises(DatabricksConfigurationError) as caught: - verify_configuration( - DatabricksConfig(host="https://workspace.test", experiment_id="１２", personal_access_token="pat") - ) - assert caught.value.code == "invalid_experiment_id" - assert experiment_api.calls == [] - - -@pytest.mark.parametrize( - "credentials", - [ - {}, - {"client_id": "client"}, - {"client_secret": "secret"}, - {"personal_access_token": " "}, - ], -) -def test_missing_credentials_fail_before_vendor_io(experiment_api: ExperimentAPI, credentials: dict[str, str]) -> None: - with pytest.raises(DatabricksConfigurationError) as caught: - verify_configuration(DatabricksConfig(host="https://workspace.test", experiment_id="12", **credentials)) - assert caught.value.code == "authentication_failed" - assert experiment_api.calls == [] - - -@pytest.mark.parametrize( - "host", - [ - "http://workspace.test", - "https://user:secret@workspace.test", - "https://workspace.test/path", - "https://workspace.test?o=1", - "https://workspace.test/#/experiments/12", - "https://workspace.test\\@evil.test", - "https://workspace.test:\n443", - "https://workspace.test:bad", - "https://[invalid", - ], -) -def test_invalid_workspace_origins_do_not_send_credentials(experiment_api: ExperimentAPI, host: str) -> None: - with pytest.raises(DatabricksConfigurationError) as caught: - verify_configuration(DatabricksConfig(host=host, experiment_id="12", personal_access_token="synthetic-key")) - assert caught.value.code == "invalid_workspace_url" - assert experiment_api.calls == [] - - -def test_project_url_is_pure_and_supports_explicit_https_port(experiment_api: ExperimentAPI) -> None: - config = DatabricksConfig(host=" https://Workspace.Example:443/ ", experiment_id="00012") - assert project_url(config) == "https://Workspace.Example:443/ml/experiments/12/traces" - assert experiment_api.calls == [] - - -@pytest.mark.parametrize( - ("status", "code"), - [ - (401, "authentication_failed"), - (403, "permission_denied"), - (404, "experiment_not_found_or_inaccessible"), - (302, "endpoint_unavailable"), - (500, "endpoint_unavailable"), - ], -) -def test_http_failures_are_safe_and_responses_are_closed(experiment_api: ExperimentAPI, status: int, code: str) -> None: - result = response("https://workspace.test", {"message": "private-key-vendor-response"}, status) - experiment_api.override = lambda _method, _url, _options: result - with pytest.raises(DatabricksConfigurationError) as caught: - verify_configuration( - DatabricksConfig(host="https://workspace.test", experiment_id="12", personal_access_token="synthetic-key") - ) - assert caught.value.code == code - assert "private-key" not in str(caught.value) - assert "synthetic-key" not in str(caught.value) - assert result.is_closed - - -@pytest.mark.parametrize( - ("body", "code"), - [ - ({}, "malformed_response"), - ([], "malformed_response"), - ({"experiment": "invalid"}, "malformed_response"), - ({"experiment": {"experiment_id": "other", "lifecycle_stage": "active"}}, "experiment_id_mismatch"), - ({"experiment": {"experiment_id": "12", "lifecycle_stage": "deleted"}}, "experiment_not_active"), - ({"experiment": {"experiment_id": "12"}}, "experiment_not_active"), - ], -) -def test_response_must_identify_selected_active_experiment( - experiment_api: ExperimentAPI, body: object, code: str -) -> None: - experiment_api.override = lambda _method, url, _options: response(url, body) - with pytest.raises(DatabricksConfigurationError) as caught: - verify_configuration( - DatabricksConfig(host="https://workspace.test", experiment_id="12", personal_access_token="synthetic-key") - ) - assert caught.value.code == code - - -@pytest.mark.parametrize("body", [{}, {"access_token": ""}, {"access_token": "token", "token_type": "other"}]) -def test_invalid_oauth_response_never_reaches_experiment_api(experiment_api: ExperimentAPI, body: object) -> None: - experiment_api.override = lambda _method, url, _options: response(url, body) - with pytest.raises(DatabricksConfigurationError, match="authentication failed"): - verify_configuration( - DatabricksConfig( - host="https://workspace.test", experiment_id="12", client_id="client", client_secret="secret" - ) - ) - assert len(experiment_api.calls) == 1 - - -def test_transport_errors_do_not_disclose_request_details(experiment_api: ExperimentAPI) -> None: - def fail(_method: str, _url: str, _options: dict[str, object]) -> httpx.Response: - raise httpx.ConnectError("private-key-vendor-response") - - experiment_api.override = fail - with pytest.raises(DatabricksConfigurationError) as caught: - verify_configuration( - DatabricksConfig(host="https://workspace.test", experiment_id="12", personal_access_token="synthetic-key") - ) - assert caught.value.code == "connection_failed" - assert "private-key" not in str(caught.value) diff --git a/api/tests/unit_tests/providers/trace/trace-mlflow/src/dify_trace_mlflow/test_mlflow_trace.py b/api/tests/unit_tests/providers/trace/trace-mlflow/src/dify_trace_mlflow/test_mlflow_trace.py deleted file mode 100644 index c0e6dc36868a2e..00000000000000 --- a/api/tests/unit_tests/providers/trace/trace-mlflow/src/dify_trace_mlflow/test_mlflow_trace.py +++ /dev/null @@ -1,639 +0,0 @@ -"""Real MLflow SDK export must not select the application's ambient OTLP destination.""" - -import json -import os -from concurrent.futures import ThreadPoolExecutor -from datetime import UTC, datetime, timedelta -from decimal import Decimal -from threading import Barrier -from typing import Never - -import mlflow -import pytest -import requests -from dify_trace_mlflow.config import DatabricksConfig, MLflowConfig -from dify_trace_mlflow.mlflow_trace import MLflowDataTrace -from google.protobuf.json_format import Parse -from mlflow.entities import SpanType, Trace -from mlflow.entities.trace_location import MlflowExperimentLocation -from mlflow.protos.service_pb2 import StartTraceV3 -from mlflow.tracing import provider as sdk_provider -from mlflow.tracing.constant import SpanAttributeKey, TraceMetadataKey, TraceTagKey -from mlflow.tracing.trace_manager import InMemoryTraceManager -from mlflow.tracing.utils.otlp import _decode_otel_proto_anyvalue -from opentelemetry import trace -from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest -from opentelemetry.proto.trace.v1.trace_pb2 import Span as ProtoSpan -from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.trace import INVALID_SPAN_CONTEXT, NonRecordingSpan - -from core.ops.entities.trace_entity import ( - DatasetRetrievalTraceInfo, - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - SuggestedQuestionTraceInfo, - ToolTraceInfo, - WorkflowTraceInfo, -) -from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionStatus -from models.model import Message -from models.workflow import WorkflowNodeExecutionModel - - -@pytest.fixture -def transport(monkeypatch: pytest.MonkeyPatch) -> list[requests.PreparedRequest]: - calls: list[requests.PreparedRequest] = [] - - def send(_session: requests.Session, request: requests.PreparedRequest, **_kwargs: object) -> requests.Response: - calls.append(request) - response = requests.Response() - response.status_code = 200 - response.request = request - response.url = request.url or "" - if response.url.endswith("/version"): - response._content = b"3.17.1" - elif response.url.endswith("/v1/traces"): - response._content = b"{}" - elif response.url.endswith("/api/3.0/mlflow/traces"): - assert isinstance(request.body, (str, bytes)) - body = json.loads(request.body) - info = body["trace"]["trace_info"] - info.setdefault("tags", {})[TraceTagKey.SPANS_LOCATION] = "TRACKING_STORE" - response._content = json.dumps(body).encode() - else: - raise AssertionError(f"Unexpected SDK endpoint: {request.method} {request.path_url}") - return response - - monkeypatch.setattr(requests.Session, "send", send) - monkeypatch.setenv("MLFLOW_DISABLE_TELEMETRY", "true") - monkeypatch.delenv("_MLFLOW_TELEMETRY_SESSION_ID", raising=False) - monkeypatch.setenv("MLFLOW_ENABLE_ASYNC_TRACE_LOGGING", "false") - monkeypatch.setenv("MLFLOW_ENABLE_OTLP_EXPORTER", "true") - monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://dify-internal.example:4318") - monkeypatch.setenv("OTEL_EXPORTER_OTLP_PROTOCOL", "http") - monkeypatch.setenv("OTEL_SERVICE_NAME", "dify-api") - monkeypatch.setenv("OTEL_RESOURCE_ATTRIBUTES", "deployment.environment=unit-test") - return calls - - -def tool_info() -> ToolTraceInfo: - started = datetime(2026, 1, 1, tzinfo=UTC) - return ToolTraceInfo( - tool_name="code", - tool_inputs={"a": 17}, - tool_outputs="323", - tool_config={}, - tool_parameters={}, - time_cost=1, - metadata={}, - inputs={"a": 17}, - outputs="323", - start_time=started, - end_time=started + timedelta(seconds=1), - ) - - -@pytest.mark.parametrize("default_provider", ["true", "false"]) -def test_real_sdk_uses_private_mlflow_exporter_without_touching_shared_state( - monkeypatch: pytest.MonkeyPatch, - transport: list[requests.PreparedRequest], - default_provider: str, -) -> None: - monkeypatch.setenv("MLFLOW_USE_DEFAULT_TRACER_PROVIDER", default_provider) - main_provider = TracerProvider() - monkeypatch.setattr(trace, "_TRACER_PROVIDER", main_provider) - sdk_before = sdk_provider.provider.get() - destination_before = sdk_provider._MLFLOW_TRACE_USER_DESTINATION.get() - unrelated = MlflowExperimentLocation("999") - monkeypatch.setattr(sdk_provider._MLFLOW_TRACE_USER_DESTINATION, "_global_value", unrelated) - destination_before = sdk_provider._MLFLOW_TRACE_USER_DESTINATION.get() - environment_before = dict(os.environ) - - def forbidden(*_args: object, **_kwargs: object) -> Never: - raise AssertionError("Legacy exporter must not initialize or reset the process SDK") - - for name in ("set_tracking_uri", "set_experiment"): - monkeypatch.setattr(mlflow, name, forbidden) - monkeypatch.setattr(sdk_provider, "set_destination", forbidden) - monkeypatch.setattr(sdk_provider, "reset", forbidden) - instance = MLflowDataTrace(MLflowConfig(tracking_uri="https://legacy.example", experiment_id="12")) - assert transport == [] - instance.trace(tool_info()) - assert sdk_provider.provider.get() is sdk_before - assert sdk_provider._MLFLOW_TRACE_USER_DESTINATION.get() is destination_before - assert trace.get_tracer_provider() is main_provider - assert dict(os.environ) == environment_before - otlp_requests = [request for request in transport if (request.url or "").endswith("/v1/traces")] - assert len(otlp_requests) == 1 - assert otlp_requests[0].url == "https://legacy.example/v1/traces" - assert otlp_requests[0].headers["x-mlflow-experiment-id"] == "12" - body = otlp_requests[0].body - assert isinstance(body, bytes) - spans = ExportTraceServiceRequest.FromString(body).resource_spans[0].scope_spans[0].spans - assert len(spans) == 1 - assert spans[0].name == "code" - assert len(spans[0].trace_id) == 16 - assert any(spans[0].trace_id) - assert any((request.url or "").endswith("/api/3.0/mlflow/traces") for request in transport) - main_provider.shutdown() - - -def exported_spans(calls: list[requests.PreparedRequest]) -> list[ProtoSpan]: - spans: list[ProtoSpan] = [] - for request in calls: - if (request.url or "").endswith("/v1/traces"): - assert isinstance(request.body, bytes) - spans.extend(ExportTraceServiceRequest.FromString(request.body).resource_spans[0].scope_spans[0].spans) - return spans - - -def exported_traces(calls: list[requests.PreparedRequest]) -> list[Trace]: - results: list[Trace] = [] - for request in calls: - if (request.url or "").endswith("/api/3.0/mlflow/traces"): - assert isinstance(request.body, (bytes, str)) - results.append(Trace.from_proto(Parse(request.body, StartTraceV3()).trace)) - return results - - -def test_workflow_parent_children_usage_metadata_and_explicit_experiment( - monkeypatch: pytest.MonkeyPatch, - transport: list[requests.PreparedRequest], -) -> None: - instance = MLflowDataTrace(MLflowConfig(tracking_uri="https://workflow.example", experiment_id="42")) - started = datetime(2026, 1, 1, tzinfo=UTC) - nodes = [ - WorkflowNodeExecutionModel( - id=f"node-{index}", - tenant_id="tenant", - app_id="app", - title=name, - node_type=node_type, - status=WorkflowNodeExecutionStatus.SUCCEEDED, - inputs='{"question":"synthetic"}', - outputs='{"text":"answer"}', - elapsed_time=0.1, - process_data=json.dumps( - { - "prompts": [{"role": "user", "text": "synthetic"}], - "model_name": "test-model", - "model_provider": "test", - "usage": {"prompt_tokens": 25, "completion_tokens": 7, "total_tokens": 32}, - } - ) - if node_type == BuiltinNodeTypes.LLM - else None, - ) - for index, (name, node_type) in enumerate( - [ - ("Start", BuiltinNodeTypes.START), - ("LLM", BuiltinNodeTypes.LLM), - ("End", BuiltinNodeTypes.END), - ] - ) - ] - for trace_node_index, trace_node_record in enumerate(nodes): - trace_node_record.created_at = started + timedelta(milliseconds=trace_node_index) - monkeypatch.setattr(instance, "_get_workflow_nodes", lambda _run_id: nodes) - info = WorkflowTraceInfo( - workflow_id="workflow", - tenant_id="tenant", - workflow_run_id="run", - workflow_run_elapsed_time=1, - workflow_run_status=WorkflowNodeExecutionStatus.SUCCEEDED, - workflow_run_inputs={"q": "synthetic", "sys.hidden": "drop"}, - workflow_run_outputs={"answer": "answer"}, - workflow_run_version="1", - total_tokens=32, - file_list=[], - query="", - conversation_id="session", - metadata={"user_id": "user"}, - start_time=started, - end_time=started + timedelta(seconds=1), - ) - ambient_provider = TracerProvider() - with ambient_provider.get_tracer("ambient").start_as_current_span("unrelated"): - instance.trace(info) - spans = exported_spans(transport) - assert len(spans) == 4 - root = next(span for span in spans if span.name == "workflow") - assert not root.parent_span_id - assert all(span.trace_id == root.trace_id for span in spans) - assert all(span.parent_span_id == root.span_id for span in spans if span is not root) - assert all( - request.headers["x-mlflow-experiment-id"] == "42" - for request in transport - if (request.url or "").endswith("/v1/traces") - ) - native = exported_traces(transport) - assert len(native) == 1 - location = native[0].info.trace_location.mlflow_experiment - assert location is not None - assert location.experiment_id == "42" - assert native[0].info.trace_metadata[TraceMetadataKey.TRACE_SESSION] == "session" - assert native[0].info.trace_metadata[TraceMetadataKey.TRACE_USER] == "user" - assert json.loads(native[0].info.trace_metadata[TraceMetadataKey.TOKEN_USAGE]) == { - "input_tokens": 25, - "output_tokens": 7, - "total_tokens": 32, - } - attributes = { - attribute.key: _decode_otel_proto_anyvalue(attribute.value) - for attribute in root.attributes - if attribute.key in (SpanAttributeKey.INPUTS, SpanAttributeKey.OUTPUTS) - } - assert attributes[SpanAttributeKey.INPUTS] == {"q": "synthetic"} - assert attributes[SpanAttributeKey.OUTPUTS] == {"answer": "answer"} - ambient_provider.shutdown() - - -def test_interleaved_instances_keep_children_on_their_own_uri_and_experiment( - transport: list[requests.PreparedRequest], -) -> None: - barrier = Barrier(2) - instances = [ - MLflowDataTrace(MLflowConfig(tracking_uri=f"https://{name}.example", experiment_id=experiment)) - for name, experiment in [("first", "12"), ("second", "34")] - ] - - def emit(instance: MLflowDataTrace) -> tuple[str, str]: - root = instance._start_span_no_context(name="workflow", span_type=SpanType.CHAIN) - barrier.wait(timeout=5) - child = instance._start_span_no_context(name="node", span_type=SpanType.TOOL, parent_span=root) - child.end(outputs="ok") - barrier.wait(timeout=5) - root.end(outputs="ok") - return root.trace_id, root.span_id - - with ThreadPoolExecutor(max_workers=2) as pool: - ids = list(pool.map(emit, instances)) - assert ids[0][0] != ids[1][0] - for host, experiment in [("first", "12"), ("second", "34")]: - requests_for_host = [item for item in transport if (item.url or "").startswith(f"https://{host}.example/")] - spans = exported_spans(requests_for_host) - assert len(spans) == 2 - assert spans[0].trace_id == spans[1].trace_id - root = next(span for span in spans if not span.parent_span_id) - child = next(span for span in spans if span.parent_span_id) - assert child.parent_span_id == root.span_id - assert all( - item.headers["x-mlflow-experiment-id"] == experiment - for item in requests_for_host - if (item.url or "").endswith("/v1/traces") - ) - location = exported_traces(requests_for_host)[0].info.trace_location.mlflow_experiment - assert location is not None - assert location.experiment_id == experiment - - -@pytest.mark.parametrize( - "method", ["message", "moderation", "dataset_retrieval", "suggested_question", "generate_name"] -) -def test_other_legacy_entrypoints_use_the_same_private_sdk( - method: str, - transport: list[requests.PreparedRequest], -) -> None: - started = datetime(2026, 1, 1, tzinfo=UTC) - message = Message( - model_provider="test", - model_id="model", - answer="answer", - total_price=Decimal(0), - ) - message.created_at = started - message.updated_at = started + timedelta(seconds=1) - common = { - "metadata": {}, - "message_data": message, - "inputs": "synthetic", - "outputs": "answer", - "start_time": started, - "end_time": started + timedelta(seconds=1), - } - match method: - case "message": - info = MessageTraceInfo( - **common, - conversation_model="chat", - conversation_mode="chat", - message_tokens=3, - answer_tokens=2, - total_tokens=5, - ) - case "moderation": - info = ModerationTraceInfo(**common, flagged=False, action="allow", preset_response="", query="synthetic") - case "dataset_retrieval": - info = DatasetRetrievalTraceInfo(**common, documents=[{"content": "synthetic"}]) - case "suggested_question": - info = SuggestedQuestionTraceInfo(**common, total_tokens=5, suggested_question=["synthetic"], level="info") - case _: - info = GenerateNameTraceInfo(**common, tenant_id="tenant") - instance = MLflowDataTrace(MLflowConfig(tracking_uri="https://other.example", experiment_id="16")) - instance.trace(info) - assert len(exported_spans(transport)) == 1 - location = exported_traces(transport)[0].info.trace_location.mlflow_experiment - assert location is not None - assert location.experiment_id == "16" - - -@pytest.mark.parametrize("failure", ["exception", "noop", "unregistered"]) -def test_span_creation_failure_is_explicit_and_does_not_export_success( - monkeypatch: pytest.MonkeyPatch, - transport: list[requests.PreparedRequest], - failure: str, -) -> None: - instance = MLflowDataTrace(MLflowConfig(tracking_uri="https://failure.example", experiment_id="10")) - if failure == "exception": - - def fail(*_args: object, **_kwargs: object) -> Never: - raise ValueError("must-not-leak-vendor-details") - - monkeypatch.setattr(instance._tracer, "start_span", fail) - expected = "mlflow_legacy_span_start_failed" - elif failure == "noop": - - def noop(*_args: object, **_kwargs: object) -> NonRecordingSpan: - return NonRecordingSpan(INVALID_SPAN_CONTEXT) - - monkeypatch.setattr(instance._tracer, "start_span", noop) - expected = "mlflow_legacy_span_not_recording" - else: - manager = InMemoryTraceManager.get_instance() - monkeypatch.setattr(manager, "get_span_from_id", lambda *_args: None) - expected = "mlflow_legacy_span_not_registered" - with pytest.raises(RuntimeError, match=f"^{expected}$"): - instance.trace(tool_info()) - assert not exported_spans(transport) - assert not exported_traces(transport) - - -def test_legacy_basic_credentials_remain_compatible_at_transport_boundary( - monkeypatch: pytest.MonkeyPatch, - transport: list[requests.PreparedRequest], -) -> None: - # This preserves legacy process-level credentials; it is not an isolation claim. - for key in ("MLFLOW_TRACKING_USERNAME", "MLFLOW_TRACKING_PASSWORD"): - monkeypatch.delenv(key, raising=False) - instance = MLflowDataTrace( - MLflowConfig( - tracking_uri="https://basic.example", - experiment_id="12", - username="synthetic-user", - password="synthetic-secret", - ) - ) - instance.trace(tool_info()) - import base64 - - expected = "Basic " + base64.b64encode(b"synthetic-user:synthetic-secret").decode() - assert all(request.headers["Authorization"] == expected for request in transport) - - -@pytest.mark.parametrize("noop", [False, True]) -def test_databricks_preserves_legacy_sdk_setup_and_span_selection( - monkeypatch: pytest.MonkeyPatch, - transport: list[requests.PreparedRequest], - noop: bool, -) -> None: - import dify_trace_mlflow.mlflow_trace as legacy - from mlflow.entities import LiveSpan - from mlflow.entities.span import NoOpSpan - - setup: list[tuple[str, str]] = [] - provider = TracerProvider() - raw = provider.get_tracer("databricks-legacy").start_span("code") - span = LiveSpan(raw, "tr-synthetic-databricks") - options: list[dict[str, object]] = [] - - def start(**kwargs: object) -> LiveSpan | NoOpSpan: - options.append(kwargs) - return NoOpSpan() if noop else span - - monkeypatch.setattr(mlflow, "set_tracking_uri", lambda uri: setup.append(("uri", uri))) - monkeypatch.setattr(mlflow, "set_experiment", lambda *, experiment_id: setup.append(("experiment", experiment_id))) - monkeypatch.setattr(legacy, "start_span_no_context", start) - for name in ("DATABRICKS_HOST", "DATABRICKS_TOKEN", "MLFLOW_ENABLE_ASYNC_TRACE_LOGGING"): - monkeypatch.delenv(name, raising=False) - instance = MLflowDataTrace( - DatabricksConfig(host="https://workspace.example", experiment_id="123", personal_access_token="synthetic-pat") - ) - assert setup == [("uri", "databricks"), ("experiment", "123")] - assert instance._tracer is None - assert instance._exporter is None - if noop: - with pytest.raises(RuntimeError, match="^mlflow_legacy_span_not_recording$"): - instance.trace(tool_info()) - raw.end() - else: - instance.trace(tool_info()) - assert span.outputs == "323" - assert options[0]["name"] == "code" - assert options[0]["inputs"] == {"a": 17} - assert transport == [] - provider.shutdown() - - -def test_private_exporter_is_synchronous_and_does_not_retain_evicted_adapters( - monkeypatch: pytest.MonkeyPatch, - transport: list[requests.PreparedRequest], -) -> None: - import atexit - import gc - import weakref - - from mlflow.tracing.export import mlflow_v3 - - monkeypatch.setenv("MLFLOW_ENABLE_ASYNC_TRACE_LOGGING", "true") - - def forbidden(*_args: object, **_kwargs: object) -> Never: - raise AssertionError("A private legacy exporter must not retain an SDK queue or exit callback") - - monkeypatch.setattr(mlflow_v3, "AsyncTraceExportQueue", forbidden) - monkeypatch.setattr(atexit, "register", forbidden) - instance = MLflowDataTrace(MLflowConfig(tracking_uri="https://synchronous.example", experiment_id="77")) - reference = weakref.ref(instance) - provider_reference = weakref.ref(instance._provider) - instance.trace(tool_info()) - instance.trace(tool_info()) - assert len(exported_traces(transport)) == 2 # Complete before return, no global flush call. - assert os.environ["MLFLOW_ENABLE_ASYNC_TRACE_LOGGING"] == "true" - del instance - gc.collect() - assert reference() is None - assert provider_reference() is None - - -@pytest.mark.parametrize("endpoint", ["/v1/traces", "/api/3.0/mlflow/traces"]) -def test_synchronous_http_failure_reaches_task_boundary_and_does_not_poison_next_trace( - monkeypatch: pytest.MonkeyPatch, - transport: list[requests.PreparedRequest], - endpoint: str, -) -> None: - original_send = requests.Session.send - fail = True - - def send(session: requests.Session, request: requests.PreparedRequest, **_kwargs: object) -> requests.Response: - if fail and (request.url or "").endswith(endpoint): - response = requests.Response() - response.status_code = 403 - response._content = b'{"error_code":"PERMISSION_DENIED","message":"synthetic denial"}' - response.request = request - response.url = request.url or "" - return response - return original_send(session, request) - - monkeypatch.setattr(requests.Session, "send", send) - instance = MLflowDataTrace(MLflowConfig(tracking_uri="https://failure-transport.example", experiment_id="70")) - with pytest.raises(RuntimeError, match="^mlflow_legacy_export_failed$"): - instance.trace(tool_info()) - assert instance._exporter is not None - assert instance._exporter.failed.get() is False - fail = False - transport.clear() - instance.trace(tool_info()) - assert len(exported_spans(transport)) == len(exported_traces(transport)) == 1 - - -@pytest.mark.parametrize("artifact_fails", [False, True]) -def test_expected_501_falls_back_to_sdk_artifact_and_only_final_failure_is_reported( - monkeypatch: pytest.MonkeyPatch, - transport: list[requests.PreparedRequest], - artifact_fails: bool, -) -> None: - from mlflow.entities import TraceData, TraceInfo - from mlflow.tracing.client import TracingClient - - assert transport == [] - original_send = requests.Session.send - uploaded: list[TraceData] = [] - - def send(session: requests.Session, request: requests.PreparedRequest, **_kwargs: object) -> requests.Response: - if (request.url or "").endswith("/v1/traces"): - response = requests.Response() - response.status_code = 501 - response._content = b'{"error_code":"NOT_IMPLEMENTED","message":"REST OTLP span logging is not supported"}' - response.request = request - response.url = request.url or "" - return response - response = original_send(session, request) - if (request.url or "").endswith("/api/3.0/mlflow/traces"): - body = response.json() - body["trace"]["trace_info"]["tags"].pop(TraceTagKey.SPANS_LOCATION) - response._content = json.dumps(body).encode() - return response - - def upload(_client: TracingClient, _info: TraceInfo, data: TraceData) -> None: - if artifact_fails: - raise requests.HTTPError("synthetic artifact failure") - uploaded.append(data) - - monkeypatch.setattr(requests.Session, "send", send) - monkeypatch.setattr(TracingClient, "_upload_trace_data", upload) - instance = MLflowDataTrace(MLflowConfig(tracking_uri="https://file-store.example", experiment_id="80")) - if artifact_fails: - with pytest.raises(RuntimeError, match="^mlflow_legacy_export_failed$"): - instance.trace(tool_info()) - else: - instance.trace(tool_info()) - assert len(uploaded) == 1 - assert len(uploaded[0].spans) == 1 - assert uploaded[0].spans[0].outputs == "323" - - -@pytest.mark.parametrize("fallback_fails", [False, True]) -def test_v3_endpoint_not_found_keeps_real_sdk_v2_fallback( - monkeypatch: pytest.MonkeyPatch, - transport: list[requests.PreparedRequest], - fallback_fails: bool, -) -> None: - from mlflow.entities.trace_info_v2 import TraceInfoV2 - from mlflow.exceptions import MlflowException - from mlflow.protos.databricks_pb2 import ENDPOINT_NOT_FOUND - from mlflow.store.tracking.rest_store import RestStore - - assert transport == [] - original_endpoint = RestStore._call_endpoint - actions: list[str] = [] - - def endpoint(store: RestStore, api: object, json_body: str, **kwargs: object) -> object: - if api is StartTraceV3: - actions.append("v3-not-found") - raise MlflowException("old server", error_code=ENDPOINT_NOT_FOUND) - return original_endpoint(store, api, json_body, **kwargs) - - def start_v2(_store: RestStore, **kwargs: object) -> TraceInfoV2: - actions.append("v2-start") - if fallback_fails: - raise requests.HTTPError("synthetic v2 failure") - return TraceInfoV2.from_dict( - { - "request_id": "tr-v2", - "experiment_id": str(kwargs["experiment_id"]), - "timestamp_ms": 1, - "execution_time_ms": 1, - "status": "OK", - "request_metadata": {}, - "tags": {TraceTagKey.SPANS_LOCATION: "TRACKING_STORE"}, - } - ) - - def end_v2(_store: RestStore, **_kwargs: object) -> None: - actions.append("v2-end") - - monkeypatch.setattr(RestStore, "_call_endpoint", endpoint) - monkeypatch.setattr(RestStore, "deprecated_start_trace_v2", start_v2) - monkeypatch.setattr(RestStore, "deprecated_end_trace_v2", end_v2) - instance = MLflowDataTrace(MLflowConfig(tracking_uri="https://old-v2.example", experiment_id="90")) - if fallback_fails: - with pytest.raises(RuntimeError, match="^mlflow_legacy_export_failed$"): - instance.trace(tool_info()) - assert actions == ["v3-not-found", "v2-start"] - else: - instance.trace(tool_info()) - assert actions == ["v3-not-found", "v2-start", "v2-end"] - - -def test_concurrent_dispatches_on_cached_adapter_do_not_share_failure_state( - monkeypatch: pytest.MonkeyPatch, - transport: list[requests.PreparedRequest], -) -> None: - from contextvars import ContextVar - - original_send = requests.Session.send - fail_this_dispatch: ContextVar[bool] = ContextVar("fail_this_dispatch", default=False) - barrier = Barrier(2) - - def send(session: requests.Session, request: requests.PreparedRequest, **_kwargs: object) -> requests.Response: - if (request.url or "").endswith("/v1/traces"): - barrier.wait(timeout=5) - if fail_this_dispatch.get(): - response = requests.Response() - response.status_code = 403 - response._content = b'{"error_code":"PERMISSION_DENIED","message":"synthetic denial"}' - response.request = request - response.url = request.url or "" - return response - return original_send(session, request) - - monkeypatch.setattr(requests.Session, "send", send) - instance = MLflowDataTrace(MLflowConfig(tracking_uri="https://concurrent.example", experiment_id="71")) - - def emit(fail: bool) -> bool: - token = fail_this_dispatch.set(fail) - try: - if fail: - with pytest.raises(RuntimeError, match="^mlflow_legacy_export_failed$"): - instance.trace(tool_info()) - return False - instance.trace(tool_info()) - return True - finally: - fail_this_dispatch.reset(token) - assert instance._exporter is not None - assert not instance._exporter.failed.get() - - with ThreadPoolExecutor(max_workers=2) as pool: - assert list(pool.map(emit, [True, False])) == [False, True] - assert len(exported_spans(transport)) == 1 diff --git a/api/tests/unit_tests/providers/trace/trace-mlflow/src/dify_trace_mlflow/test_oss_configuration.py b/api/tests/unit_tests/providers/trace/trace-mlflow/src/dify_trace_mlflow/test_oss_configuration.py deleted file mode 100644 index 82ffb642a46805..00000000000000 --- a/api/tests/unit_tests/providers/trace/trace-mlflow/src/dify_trace_mlflow/test_oss_configuration.py +++ /dev/null @@ -1,325 +0,0 @@ -"""OSS MLflow verification reads one experiment with isolated Basic credentials.""" - -import base64 -import importlib as _test_importlib -import os -from collections.abc import Callable, Iterator -from concurrent.futures import ThreadPoolExecutor -from threading import Barrier -from typing import Never - -import httpx -import mlflow -import pytest -from dify_trace_mlflow.config import MLflowConfig -from dify_trace_mlflow.mlflow_trace import MLflowDataTrace -from dify_trace_mlflow.oss_configuration import ( - MLflowConfigurationError, - normalize_configuration, - project_url, - verify_configuration, -) -from pydantic import ValidationError - -from core.helper import ssrf_proxy - -ExperimentAPI = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-mlflow.src.dify_trace_mlflow.test_configuration" -).ExperimentAPI - -response = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-mlflow.src.dify_trace_mlflow.test_configuration" -).response - - -@pytest.fixture(autouse=True) -def isolated_sdk_state(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: - names = ( - "MLFLOW_TRACKING_URI", - "MLFLOW_TRACKING_USERNAME", - "MLFLOW_TRACKING_PASSWORD", - "MLFLOW_TRACKING_TOKEN", - "MLFLOW_EXPERIMENT_ID", - "MLFLOW_ENABLE_ASYNC_TRACE_LOGGING", - "DATABRICKS_HOST", - "DATABRICKS_TOKEN", - ) - for name in names: - monkeypatch.setenv(name, "synthetic-ambient-value") - forbidden_calls: list[str] = [] - - def forbidden(name: str) -> Callable[..., Never]: - def block(*_args: object, **_kwargs: object) -> Never: - forbidden_calls.append(name) - raise AssertionError("Configuration must not initialize an exporter or global SDK") - - return block - - monkeypatch.setattr(MLflowDataTrace, "__init__", forbidden("legacy_constructor")) - for name in ("set_tracking_uri", "set_experiment", "search_experiments"): - monkeypatch.setattr(mlflow, name, forbidden(name)) - monkeypatch.setattr(mlflow.config, "enable_async_logging", forbidden("async_logging")) - yield - assert forbidden_calls == [] - assert all(os.environ[name] == "synthetic-ambient-value" for name in names) - - -@pytest.fixture -def experiment_api(monkeypatch: pytest.MonkeyPatch) -> ExperimentAPI: - api = ExperimentAPI() - monkeypatch.setattr(ssrf_proxy, "get", api.get) - - def forbid_post(*_args: object, **_kwargs: object) -> Never: - raise AssertionError("OSS verification must not create an experiment, run, trace, or session") - - monkeypatch.setattr(ssrf_proxy, "post", forbid_post) - return api - - -def test_exact_experiment_and_path_prefix_use_only_explicit_basic_auth(experiment_api: ExperimentAPI) -> None: - values = { - "tracking_uri": " HTTP://FIRST.Example:5000/team/mlflow/ ", - "experiment_id": "00012", - "username": "first-user", - "password": "first-password", - } - config = normalize_configuration(values) - target = verify_configuration(config) - assert target.tracking_uri == "http://first.example:5000/team/mlflow" - assert target.experiment_id == "12" - assert project_url(config) == "http://first.example:5000/team/mlflow/#/experiments/12/traces" - method, url, options = experiment_api.calls[0] - assert method == "GET" - assert url == "http://first.example:5000/team/mlflow/api/2.0/mlflow/experiments/get" - assert options["params"] == {"experiment_id": "12"} - assert options["headers"] == { - "Accept": "application/json", - "Accept-Encoding": "identity", - "Cookie": "", - "Authorization": "Basic " + base64.b64encode(b"first-user:first-password").decode(), - } - assert options["follow_redirects"] is False - assert options["stream_response"] is True - assert options["timeout"] == 30.0 - assert options["max_retries"] == 0 - assert values["tracking_uri"].startswith(" HTTP") - - -def test_anonymous_request_does_not_inherit_previous_or_environment_credentials(experiment_api: ExperimentAPI) -> None: - verify_configuration( - MLflowConfig(tracking_uri="https://same.test", experiment_id="12", username="user", password="password") - ) - verify_configuration(MLflowConfig(tracking_uri="https://same.test", experiment_id="34")) - assert len(experiment_api.calls) == 2 - assert experiment_api.calls[1][2]["headers"] == { - "Accept": "application/json", - "Accept-Encoding": "identity", - "Cookie": "", - } - - -def test_concurrent_destinations_keep_experiment_and_auth_together(experiment_api: ExperimentAPI) -> None: - barrier = Barrier(2) - expected_destinations = { - "https://first.test/prefix/api/2.0/mlflow/experiments/get": ("first", "12"), - "https://second.test/prefix/api/2.0/mlflow/experiments/get": ("second", "34"), - } - - def read(_method: str, url: str, options: dict[str, object]) -> httpx.Response: - params = options["params"] - headers = options["headers"] - assert isinstance(params, dict) - assert isinstance(headers, dict) - label, experiment_id = expected_destinations[url] - assert params["experiment_id"] == experiment_id - assert headers["Authorization"] == "Basic " + base64.b64encode(f"{label}:{label}-secret".encode()).decode() - barrier.wait(timeout=5) - return response(url, {"experiment": {"experiment_id": params["experiment_id"], "lifecycle_stage": "active"}}) - - experiment_api.override = read - with ThreadPoolExecutor(max_workers=2) as pool: - pending = [ - pool.submit( - verify_configuration, - MLflowConfig( - tracking_uri=f"https://{label}.test/prefix", - experiment_id=experiment, - username=label, - password=f"{label}-secret", - ), - ) - for label, experiment in (("first", "12"), ("second", "34")) - ] - assert [item.result().experiment_id for item in pending] == ["12", "34"] - assert sorted(url for _method, url, _options in experiment_api.calls) == sorted(expected_destinations) - - -@pytest.mark.parametrize( - "uri", - [ - "", - "ftp://server.test", - "https:///missing", - "https://user:secret@server.test", - "https://@server.test", - "https://server.test?", - "https://server.test#", - "https://server.test?key=secret", - "https://server.test/#/experiments/1", - "https://server.test:bad", - "https://server.test:0", - "https://server.test:", - "https://[invalid", - "https://server.test\\@other.test", - "\nhttps://server.test", - "https://ser\nver.test", - "https://server.test/./prefix", - "https://server.test/prefix/../other", - "https://server.test/%2e%2e/other", - "https://server.test/prefix%2fother", - "https://server.test/%252fother", - "https://server.test/%00", - "https://server.test/%zz", - "https://server.test//prefix", - ], -) -def test_ambiguous_destinations_fail_before_request(experiment_api: ExperimentAPI, uri: str) -> None: - with pytest.raises((MLflowConfigurationError, ValidationError)): - verify_configuration(normalize_configuration({"tracking_uri": uri, "experiment_id": "12"})) - assert experiment_api.calls == [] - - -@pytest.mark.parametrize("experiment", ["", "-1", "1.5", "１２", "abc"]) -def test_experiment_ids_are_ascii_integer_identifiers(experiment_api: ExperimentAPI, experiment: str) -> None: - with pytest.raises((MLflowConfigurationError, ValidationError)): - normalize_configuration({"tracking_uri": "https://server.test", "experiment_id": experiment}) - assert experiment_api.calls == [] - - -@pytest.mark.parametrize( - ("username", "password"), - [("user", None), (None, "password"), ("user", ""), (" ", "password"), ("user", " "), ("user:name", "password")], -) -def test_partial_or_ambiguous_basic_auth_never_sends_credentials( - experiment_api: ExperimentAPI, username: str | None, password: str | None -) -> None: - with pytest.raises(MLflowConfigurationError) as caught: - verify_configuration(MLflowConfig(tracking_uri="https://server.test", username=username, password=password)) - assert caught.value.code == "incomplete_basic_auth" - assert experiment_api.calls == [] - - -@pytest.mark.parametrize( - ("status", "code"), - [ - (401, "authentication_failed"), - (403, "permission_denied"), - (404, "experiment_not_found_or_inaccessible"), - (302, "endpoint_unavailable"), - (307, "endpoint_unavailable"), - (500, "endpoint_unavailable"), - ], -) -def test_http_failures_do_not_follow_redirect_or_disclose_body( - experiment_api: ExperimentAPI, status: int, code: str -) -> None: - result = response("https://server.test", {"message": "private-vendor-body"}, status) - result.headers["Location"] = "https://other.test/secret" - experiment_api.override = lambda _method, _url, _options: result - with pytest.raises(MLflowConfigurationError) as caught: - verify_configuration( - MLflowConfig(tracking_uri="https://server.test", username="user", password="private-password") - ) - assert caught.value.code == code - assert "private" not in str(caught.value) - assert result.is_closed - assert len(experiment_api.calls) == 1 - - -@pytest.mark.parametrize( - ("body", "code"), - [ - ({}, "malformed_response"), - ([], "malformed_response"), - ({"experiment": "invalid"}, "malformed_response"), - ({"experiment": {"experiment_id": "34", "lifecycle_stage": "active"}}, "experiment_id_mismatch"), - ({"experiment": {"experiment_id": "12", "lifecycle_stage": "deleted"}}, "experiment_not_active"), - ({"experiment": {"experiment_id": "12"}}, "experiment_not_active"), - ], -) -def test_response_must_match_exact_active_experiment(experiment_api: ExperimentAPI, body: object, code: str) -> None: - experiment_api.override = lambda _method, url, _options: response(url, body) - with pytest.raises(MLflowConfigurationError) as caught: - verify_configuration(MLflowConfig(tracking_uri="https://server.test", experiment_id="12")) - assert caught.value.code == code - - -def test_transport_failure_and_oversized_response_are_bounded_and_safe(experiment_api: ExperimentAPI) -> None: - def fail(_method: str, _url: str, _options: dict[str, object]) -> httpx.Response: - raise httpx.ConnectError("private-url-and-password") - - experiment_api.override = fail - with pytest.raises(MLflowConfigurationError, match="MLflow verification request failed"): - verify_configuration(MLflowConfig(tracking_uri="https://server.test")) - oversized = response("https://server.test", {"message": "x" * (1024 * 1024)}) - experiment_api.override = lambda _method, _url, _options: oversized - with pytest.raises(MLflowConfigurationError) as caught: - verify_configuration(MLflowConfig(tracking_uri="https://server.test")) - assert caught.value.code == "malformed_response" - assert oversized.is_closed - - -def test_project_link_requires_no_auth_or_io_and_preserves_valid_prefix(experiment_api: ExperimentAPI) -> None: - assert project_url(MLflowConfig(tracking_uri="https://Server.test/team/mlflow/", experiment_id="000")) == ( - "https://server.test/team/mlflow/#/experiments/0/traces" - ) - assert experiment_api.calls == [] - - -def test_real_http_transport_overrides_ambient_auth_and_cookies(monkeypatch: pytest.MonkeyPatch) -> None: - requests: list[httpx.Request] = [] - - def handle(request: httpx.Request) -> httpx.Response: - requests.append(request) - assert request.method == "GET" - assert request.url.path == "/prefix/api/2.0/mlflow/experiments/get" - assert request.headers.get("cookie") == "" - return httpx.Response( - 200, - headers={"Set-Cookie": "session=synthetic-vendor-session; Path=/"}, - json={"experiment": {"experiment_id": request.url.params["experiment_id"], "lifecycle_stage": "active"}}, - ) - - with httpx.Client( - transport=httpx.MockTransport(handle), - auth=httpx.BasicAuth("ambient-user", "ambient-password"), - cookies={"session": "synthetic-ambient-session"}, - trust_env=False, - ) as client: - - def get_client(_ssl_verify_enabled: bool) -> httpx.Client: - return client - - monkeypatch.setattr(ssrf_proxy, "_get_ssrf_client", get_client) - verify_configuration( - MLflowConfig( - tracking_uri="https://server.test/prefix", experiment_id="12", username="owner", password="secret" - ) - ) - verify_configuration(MLflowConfig(tracking_uri="https://server.test/prefix", experiment_id="34")) - assert len(requests) == 2 - assert requests[0].headers["authorization"] == "Basic " + base64.b64encode(b"owner:secret").decode() - assert "authorization" not in requests[1].headers - - -@pytest.mark.parametrize("compressed", [False, True]) -def test_invalid_encoding_or_json_response_is_closed(experiment_api: ExperimentAPI, compressed: bool) -> None: - result = httpx.Response(200, content=b"private-vendor-text", request=httpx.Request("GET", "https://server.test")) - if compressed: - result.headers["Content-Encoding"] = "gzip" - experiment_api.override = lambda _method, _url, _options: result - with pytest.raises(MLflowConfigurationError) as caught: - verify_configuration(MLflowConfig(tracking_uri="https://server.test")) - assert caught.value.code == "malformed_response" - assert "private-vendor-text" not in str(caught.value) - assert result.is_closed diff --git a/api/tests/unit_tests/providers/trace/trace-opik/src/dify_trace_opik/test_configuration.py b/api/tests/unit_tests/providers/trace/trace-opik/src/dify_trace_opik/test_configuration.py deleted file mode 100644 index 073b5bb8233faa..00000000000000 --- a/api/tests/unit_tests/providers/trace/trace-opik/src/dify_trace_opik/test_configuration.py +++ /dev/null @@ -1,188 +0,0 @@ -"""Opik auth contracts at the read-only HTTP boundary, without SDK streamers.""" - -from collections.abc import Callable -from dataclasses import dataclass, field -from urllib.parse import parse_qs, urlsplit - -import httpx -import pytest -from dify_trace_opik.config import OpikConfig -from dify_trace_opik.configuration import OpikConfigurationError, console_url, project_url, verify_configuration - -from core.helper import ssrf_proxy - - -@dataclass -class AccessAPI: - calls: list[tuple[str, str, dict[str, object]]] = field(default_factory=list) - default_workspace: object = "resolved-team" - override: Callable[[str, str, dict[str, object]], httpx.Response] | None = None - - def request(self, method: str, url: str, **options: object) -> httpx.Response: - self.calls.append((method, url, options)) - if self.override: - return self.override(method, url, options) - request = httpx.Request(method, url) - if method == "POST": - assert url.endswith("/v1/private/auth"), "configuration must never create projects, spans or traces" - assert options["json"] == {} - return httpx.Response(204, request=request) - assert method == "GET" - assert url.endswith("/v1/private/auth/workspace") - return httpx.Response(200, json={"workspace_name": self.default_workspace}, request=request) - - -@pytest.fixture -def access_api(monkeypatch: pytest.MonkeyPatch) -> AccessAPI: - api = AccessAPI() - monkeypatch.setattr(ssrf_proxy, "make_request", api.request) - return api - - -def test_named_workspaces_use_isolated_auth_without_default_lookup(access_api: AccessAPI) -> None: - for host, key, workspace in [ - ("https://first.test/opik/api/", "first-key", "first-team"), - ("https://second.test/api/", "second-key", "second-team"), - ]: - result = verify_configuration(OpikConfig(api_key=key, workspace=workspace, project="not-created", url=host)) - assert (result.url, result.workspace, result.project) == (host, workspace, "not-created") - method, url, options = access_api.calls[-1] - assert method == "POST" - assert url == host + "v1/private/auth" - assert options["headers"] == { - "Authorization": key, - "Comet-Workspace": workspace, - "Accept": "application/json", - } - assert options["json"] == {} - assert options["max_retries"] == 0 - assert options["follow_redirects"] is False - assert options["timeout"] == 30.0 - assert len(access_api.calls) == 2 - - -@pytest.mark.parametrize("workspace", [None, "", "default", " default "]) -def test_default_workspace_is_resolved_only_during_verification(access_api: AccessAPI, workspace: str | None) -> None: - result = verify_configuration(OpikConfig(api_key="synthetic-key", workspace=workspace)) - assert result.workspace == "resolved-team" - assert result.project == "Default Project" - assert [method for method, _, _ in access_api.calls] == ["POST", "GET"] - assert access_api.calls[1][2]["headers"] == access_api.calls[0][2]["headers"] - - -@pytest.mark.parametrize("key", [None, ""]) -def test_keyless_self_host_does_not_fall_back_to_environment( - access_api: AccessAPI, monkeypatch: pytest.MonkeyPatch, key: str | None -) -> None: - monkeypatch.setenv("OPIK_API_KEY", "unrelated-global-key") - monkeypatch.setenv("OPIK_WORKSPACE", "unrelated-global-team") - monkeypatch.setenv("OPIK_URL_OVERRIDE", "https://unrelated.test/api/") - result = verify_configuration(OpikConfig(api_key=key, workspace="local", url="http://selfhost.test:5173/api/")) - assert result.url == "http://selfhost.test:5173/api/" - assert access_api.calls[0][2]["headers"] == {"Comet-Workspace": "local", "Accept": "application/json"} - assert len(access_api.calls) == 1 - - -@pytest.mark.parametrize( - "url", - [ - "https://user:password@private.test/api/", - "https://private.test:invalid/api/", - "http:///api/", - "https://private.test/api/?next=/api/", - "https://private.test/api/#/api/", - "https://private.test/with space/api/", - "https://private.test/back\\slash/api/", - ], -) -def test_invalid_destination_is_rejected_before_sending_credentials(access_api: AccessAPI, url: str) -> None: - with pytest.raises(OpikConfigurationError) as error: - verify_configuration(OpikConfig(api_key="private-key", url=url)) - assert error.value.code == "invalid_endpoint" - assert access_api.calls == [] - assert "private-key" not in str(error.value) - - -@pytest.mark.parametrize( - ("status", "code"), - [ - (401, "authentication_failed"), - (403, "permission_denied"), - (302, "endpoint_unavailable"), - (500, "endpoint_unavailable"), - ], -) -def test_auth_failures_are_sanitized(access_api: AccessAPI, status: int, code: str) -> None: - access_api.override = lambda method, url, _options: httpx.Response( - status, text="private-vendor-error", request=httpx.Request(method, url) - ) - with pytest.raises(OpikConfigurationError) as error: - verify_configuration(OpikConfig(api_key="private-key", workspace="team")) - assert error.value.code == code - assert "private" not in str(error.value) - assert len(access_api.calls) == 1 - - -def test_auth_html_success_is_not_a_verified_workspace(access_api: AccessAPI) -> None: - access_api.override = lambda method, url, _options: httpx.Response( - 200, text="Login", request=httpx.Request(method, url) - ) - with pytest.raises(OpikConfigurationError, match="invalid response"): - verify_configuration(OpikConfig(workspace="team")) - - -@pytest.mark.parametrize("workspace", [None, "", " ", 123, []]) -def test_default_workspace_response_requires_a_name(access_api: AccessAPI, workspace: object) -> None: - access_api.default_workspace = workspace - with pytest.raises(OpikConfigurationError) as error: - verify_configuration(OpikConfig()) - assert error.value.code == "malformed_response" - - -def test_connection_error_never_exposes_request_or_key(access_api: AccessAPI) -> None: - def fail(method: str, url: str, _options: dict[str, object]) -> httpx.Response: - raise httpx.ReadTimeout("private-key in vendor detail", request=httpx.Request(method, url)) - - access_api.override = fail - with pytest.raises(OpikConfigurationError) as error: - verify_configuration(OpikConfig(api_key="private-key")) - assert error.value.code == "connection_failed" - assert "private-key" not in str(error.value) - - -def test_project_url_uses_configured_host_and_encodes_names_without_network(access_api: AccessAPI) -> None: - config = OpikConfig(url="https://private.test:7443/prefix/api/", workspace="team/name", project="A & B/? 中文") - url = project_url(config) - assert url is not None - parsed = urlsplit(url) - assert parsed.netloc == "private.test:7443" - assert parsed.path == "/prefix/team%2Fname/redirect/projects" - assert parse_qs(parsed.query) == {"name": ["A & B/? 中文"]} - assert access_api.calls == [] - - -@pytest.mark.parametrize("workspace", [None, "", "default"]) -def test_historical_default_workspace_get_has_a_pure_console_fallback( - access_api: AccessAPI, workspace: str | None -) -> None: - config = OpikConfig(url="http://selfhost.test:5173/prefix/api/", workspace=workspace) - assert project_url(config) is None - assert console_url(config) == "http://selfhost.test:5173/prefix/" - assert access_api.calls == [] - - -@pytest.mark.parametrize( - ("workspace", "project", "code"), - [ - ("..", "ok", "invalid_workspace"), - ("bad\nworkspace", "ok", "invalid_workspace"), - ("team", "bad\nproject", "invalid_project"), - ], -) -def test_invalid_names_cannot_escape_the_console_path_or_headers( - access_api: AccessAPI, workspace: str, project: str, code: str -) -> None: - with pytest.raises(OpikConfigurationError) as error: - verify_configuration(OpikConfig(workspace=workspace, project=project)) - assert error.value.code == code - assert access_api.calls == [] diff --git a/api/tests/unit_tests/providers/trace/trace-opik/src/dify_trace_opik/test_opik_trace.py b/api/tests/unit_tests/providers/trace/trace-opik/src/dify_trace_opik/test_opik_trace.py deleted file mode 100644 index f11027ac0f2cfd..00000000000000 --- a/api/tests/unit_tests/providers/trace/trace-opik/src/dify_trace_opik/test_opik_trace.py +++ /dev/null @@ -1,214 +0,0 @@ -"""Legacy external IDs retain stable SDK batches and workflow parent relations.""" - -import json -from datetime import datetime, timedelta -from uuid import NAMESPACE_URL, UUID, uuid4, uuid5 - -import pytest -from dify_trace_opik.opik_trace import OpikDataTrace, prepare_opik_uuid -from opik import Opik -from opik.config import OpikConfig as SDKConfig -from opik.id_helpers import uuid4_to_uuid7 -from opik.message_processing.batching.batchers import CreateSpanMessageBatcher, CreateTraceMessageBatcher -from opik.message_processing.messages import ( - BaseMessage, - CreateSpanMessage, - CreateSpansBatchMessage, - CreateTraceBatchMessage, - CreateTraceMessage, -) -from opik.message_processing.streamer import Streamer -from sqlalchemy import Engine -from sqlalchemy.orm import Session, sessionmaker - -from core.ops.entities.trace_entity import WorkflowTraceInfo -from extensions.ext_database import db -from extensions.logstore.repositories.logstore_api_workflow_node_execution_repository import ( - _dict_to_workflow_node_execution_model, -) -from models.account import TenantAccountJoin, TenantAccountRole -from tests.helpers.unit.model_factories import make_account, make_app, make_tenant - -START = datetime(2026, 10, 16, 15, 18, 10, 792586) -WORKFLOW_ID = "11111111-2222-4333-8444-555555555555" -EXTERNAL_IDS = [ - "0123456789abcdef0123456789abcdef", - "019140fb-6972-7d74-a219-6cba9eb543b1", - str(uuid5(NAMESPACE_URL, "external-trace")), - "external-trace_correlation", -] - - -@pytest.mark.parametrize("identifier", [WORKFLOW_ID, WORKFLOW_ID.upper(), UUID(WORKFLOW_ID).hex]) -def test_valid_uuid4_keeps_the_original_sdk_mapping(identifier: str) -> None: - assert prepare_opik_uuid(START, identifier) == uuid4_to_uuid7(START, identifier) - - -@pytest.mark.parametrize("identifier", EXTERNAL_IDS) -def test_external_identifiers_become_stable_uuid7_with_unchanged_time(identifier: str) -> None: - mapped = prepare_opik_uuid(START, identifier) - assert mapped.version == 7 - assert mapped == prepare_opik_uuid(START, identifier) - assert int.from_bytes(mapped.bytes[:6]) == int(START.timestamp() * 1000) - assert mapped != prepare_opik_uuid(START, identifier + "-different") - - -def test_missing_identifier_keeps_random_fallback() -> None: - first, second = prepare_opik_uuid(START, None), prepare_opik_uuid(START, None) - assert first.version == second.version == 7 - assert first != second - assert int.from_bytes(first.bytes[:6]) == int(START.timestamp() * 1000) - assert prepare_opik_uuid(None, WORKFLOW_ID).version == 7 - - -@pytest.mark.parametrize("external_id", [*EXTERNAL_IDS, WORKFLOW_ID, None]) -def test_workflow_external_ids_reach_real_sdk_batches_with_complete_owned_tree( - external_id: str | None, - sqlite_engine: Engine, - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, -) -> None: - tenant_id, app_id, account_id, definition_id = (str(uuid4()) for _ in range(4)) - node_ids = [str(uuid4()) for _ in range(3)] - with sqlite_session_factory.begin() as session: - session.add(make_tenant(tenant_id=tenant_id)) - session.add(make_account(account_id=account_id)) - session.add( - TenantAccountJoin(tenant_id=tenant_id, account_id=account_id, role=TenantAccountRole.OWNER, current=True) - ) - session.add(make_app(app_id=app_id, tenant_id=tenant_id, created_by=account_id)) - for index, kind in enumerate(("start", "llm", "end")): - usage = {"prompt_tokens": 27, "completion_tokens": 5, "total_tokens": 32} - process: dict[str, object] = ( - { - "model_mode": "chat", - "model_name": "test-model", - "model_provider": "provider/test", - "usage": usage, - "prompts": [{"role": "user", "content": "synthetic"}], - } - if kind == "llm" - else {} - ) - session.add( - _dict_to_workflow_node_execution_model( - { - "id": node_ids[index], - "tenant_id": tenant_id, - "app_id": app_id, - "workflow_id": definition_id, - "workflow_run_id": WORKFLOW_ID, - "node_id": kind, - "node_type": kind, - "title": kind, - "created_by": account_id, - "created_at": (START + timedelta(milliseconds=index * 100)).isoformat(), - "status": "succeeded", - "index": index, - "elapsed_time": 0.1, - "inputs": json.dumps({}), - "outputs": json.dumps({"answer": "synthetic"}), - "process_data": json.dumps(process), - "execution_metadata": json.dumps({"total_tokens": 32} if kind == "llm" else {}), - } - ) - ) - # A node from another run must not enter this trace's SDK batch. - session.add( - _dict_to_workflow_node_execution_model( - { - "id": str(uuid4()), - "tenant_id": tenant_id, - "app_id": app_id, - "workflow_id": definition_id, - "workflow_run_id": str(uuid4()), - "node_id": "foreign", - "node_type": "llm", - "title": "foreign", - "created_by": account_id, - "created_at": START.isoformat(), - "status": "succeeded", - "elapsed_time": 0.1, - } - ) - ) - - # Bind the real SQLAlchemy owner to the test engine; no Session/repository double. - monkeypatch.setattr(type(db), "engine", property(lambda _owner: sqlite_engine)) - queued: list[BaseMessage] = [] - streamer = object.__new__(Streamer) - monkeypatch.setattr(streamer, "put", queued.append) - client = object.__new__(Opik) - client._streamer = streamer - client._config = SDKConfig.model_construct(url_override="https://offline.invalid/api", project_name="test-project") - client._project_name = "test-project" - client._project_name_most_recent_trace = None - instance = object.__new__(OpikDataTrace) - instance.opik_client = client - instance.project = "test-project" - info = WorkflowTraceInfo( - workflow_id=definition_id, - tenant_id=tenant_id, - workflow_run_id=WORKFLOW_ID, - workflow_run_elapsed_time=1, - workflow_run_status="succeeded", - workflow_run_inputs={"query": "synthetic"}, - workflow_run_outputs={"answer": "synthetic"}, - workflow_run_version="draft", - total_tokens=32, - file_list=[], - query="synthetic", - metadata={"app_id": app_id}, - start_time=START, - end_time=START + timedelta(seconds=1), - trace_id=external_id, - ) - instance.workflow_trace(info) - trace_batches: list[BaseMessage] = [] - span_batches: list[BaseMessage] = [] - trace_batcher = CreateTraceMessageBatcher(trace_batches.append, 1000, 2) - span_batcher = CreateSpanMessageBatcher(span_batches.append, 1000, 2) - for message in queued: - if isinstance(message, CreateTraceMessage): - trace_batcher.add(message) - elif isinstance(message, CreateSpanMessage): - span_batcher.add(message) - else: - pytest.fail("Unexpected SDK message type") - trace_batcher.flush() - span_batcher.flush() - assert len(trace_batches) == len(span_batches) == 1 - trace_batch, span_batch = trace_batches[0], span_batches[0] - assert isinstance(trace_batch, CreateTraceBatchMessage) - assert isinstance(span_batch, CreateSpansBatchMessage) - assert len(trace_batch.batch) == 1 - assert len(span_batch.batch) == 4 - trace = trace_batch.batch[0] - root = next(span for span in span_batch.batch if span.parent_span_id is None) - assert trace.id == str(prepare_opik_uuid(START, external_id or WORKFLOW_ID)) - assert trace.start_time.timestamp() == START.timestamp() - assert root.id != trace.id - assert {span.trace_id for span in span_batch.batch} == {trace.id} - assert len({span.id for span in span_batch.batch}) == 4 - assert all(UUID(str(span.id)).version == 7 for span in span_batch.batch) - children = [span for span in span_batch.batch if span.parent_span_id is not None] - assert all(span.parent_span_id == root.id for span in children) - assert {span.metadata["node_execution_id"] for span in children if isinstance(span.metadata, dict)} == set(node_ids) - for index, node_id in enumerate(node_ids): - child = next( - span - for span in children - if isinstance(span.metadata, dict) and span.metadata["node_execution_id"] == node_id - ) - assert child.start_time.timestamp() == (START + timedelta(milliseconds=index * 100)).timestamp() - assert child.end_time is not None - assert (child.end_time - child.start_time).total_seconds() == pytest.approx(0.1) - llm = next(span for span in children if span.type == "llm") - assert llm.usage is not None - assert {key: llm.usage[key] for key in ("prompt_tokens", "completion_tokens", "total_tokens")} == { - "prompt_tokens": 27, - "completion_tokens": 5, - "total_tokens": 32, - } - assert sum(span.usage.get("total_tokens", 0) for span in span_batch.batch if span.usage) == 32 - assert trace.output == {"answer": "synthetic"} diff --git a/api/tests/unit_tests/providers/trace/trace-tencent/src/dify_trace_tencent/test_configuration.py b/api/tests/unit_tests/providers/trace/trace-tencent/src/dify_trace_tencent/test_configuration.py deleted file mode 100644 index c3cd8d08cb23fb..00000000000000 --- a/api/tests/unit_tests/providers/trace/trace-tencent/src/dify_trace_tencent/test_configuration.py +++ /dev/null @@ -1,390 +0,0 @@ -"""Tencent configuration checks must reject malformed input without provider I/O.""" - -from collections.abc import Callable, Iterator, Sequence -from dataclasses import dataclass, field -from types import MappingProxyType -from typing import Never - -import grpc # type: ignore[import-untyped] -import pytest -from dify_trace_tencent.configuration import ( - TencentConfigurationError, - console_url, - normalize_configuration, - verify_configuration, -) -from opentelemetry.context import _SUPPRESS_INSTRUMENTATION_KEY, get_value -from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( - ExportTraceServiceRequest, - ExportTraceServiceResponse, -) - - -def configuration_values(**overrides: object) -> dict[str, object]: - return { - "endpoint": "https://collector.example:4320", - "service_name": "dify-agent", - "token": "synthetic-token", - **overrides, - } - - -@pytest.fixture(autouse=True) -def _forbid_provider_io(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: - attempts: list[str] = [] - - def forbidden(*_args: object, **_kwargs: object) -> Never: - attempts.append("provider_io") - raise AssertionError("Local configuration validation must not construct clients or perform network I/O") - - for target in ( - "dify_trace_tencent.tencent_trace.TencentDataTrace.__init__", - "dify_trace_tencent.client.TencentTraceClient.__init__", - "socket.socket", - "socket.create_connection", - "socket.getaddrinfo", - "grpc.secure_channel", - "grpc.insecure_channel", - ): - monkeypatch.setattr(target, forbidden) - yield - assert attempts == [] - - -def test_normalization_accepts_read_only_mapping_and_preserves_opaque_token() -> None: - values = configuration_values( - endpoint=" https://COLLECTOR.Example:4320/ ", - service_name=" dify-agent ", - token=" opaque token/+== ", - ) - original = dict(values) - - result = normalize_configuration(MappingProxyType(values)) - - assert result.model_dump() == { - "endpoint": "https://collector.example:4320", - "service_name": "dify-agent", - "token": " opaque token/+== ", - } - assert values == original - - -@pytest.mark.parametrize("service_name", ["a", "a" * 63, "app-1", "a--b"]) -def test_accepts_documented_service_name_boundaries(service_name: str) -> None: - result = normalize_configuration(configuration_values(service_name=service_name)) - - assert result.service_name == service_name - - -@pytest.mark.parametrize( - "service_name", - [None, 42, "", " ", "dify_app", "Dify", "1app", "-app", "app-", "a" * 64, "应用", "a.b", "a/b", "a b", "a\nb"], -) -def test_rejects_invalid_service_names_before_provider_io(service_name: object) -> None: - with pytest.raises(TencentConfigurationError) as caught: - normalize_configuration(configuration_values(service_name=service_name)) - - assert caught.value.code == "invalid_service_name" - - -@pytest.mark.parametrize( - ("endpoint", "expected"), - [ - ("https://COLLECTOR.Example:4320/", "https://collector.example:4320"), - ("http://collector.example:4319", "http://collector.example:4319"), - ("https://collector.example/", "https://collector.example"), - ("http://collector.example", "http://collector.example"), - ("https://collector.example:443", "https://collector.example:443"), - ("http://127.0.0.1:1/", "http://127.0.0.1:1"), - ("https://[2001:DB8::1]:65535/", "https://[2001:db8::1]:65535"), - ("https://BÜCHER.example:4320/", "https://xn--bcher-kva.example:4320"), - ], -) -def test_normalizes_origins_without_changing_explicit_port(endpoint: str, expected: str) -> None: - result = normalize_configuration(configuration_values(endpoint=endpoint)) - - assert result.endpoint == expected - - -@pytest.mark.parametrize( - "endpoint", - [ - None, - 4320, - "", - " ", - "collector.example:4320", - "ftp://collector.example:4320", - "https://", - "https:///", - "https://user:synthetic-secret@collector.example:4320", - "https://@collector.example:4320", - "https://collector.example:4320/v1/traces", - "https://collector.example:4320//", - "https://collector.example:4320?token=synthetic-secret", - "https://collector.example:4320?", - "https://collector.example:4320#project", - "https://collector.example:4320#", - "https://collector.example\\@other.example:4320", - "https://col lector.example:4320", - "https://collector.example:\n4320", - "https://collector.example:\t4320", - "https://collector\x00.example:4320", - "https://collector.example:0", - "https://collector.example:65536", - "https://collector.example:-1", - "https://collector.example:invalid", - "https://collector.example:", - "https://[2001:db8::1", - "https://[not-an-ipv6-address]:4320", - ], -) -def test_rejects_invalid_endpoints_before_provider_io(endpoint: object) -> None: - with pytest.raises(TencentConfigurationError) as caught: - normalize_configuration(configuration_values(endpoint=endpoint)) - - assert caught.value.code == "invalid_endpoint" - assert "synthetic-secret" not in str(caught.value) - assert "synthetic-token" not in str(caught.value) - - -@pytest.mark.parametrize("token", [None, 42, "", " ", "\t\n"]) -def test_rejects_missing_or_blank_tokens_without_provider_io(token: object) -> None: - with pytest.raises(TencentConfigurationError) as caught: - normalize_configuration(configuration_values(token=token)) - - assert caught.value.code == "invalid_token" - - -def test_masked_token_survives_local_validation_for_patch_restoration() -> None: - result = normalize_configuration(configuration_values(token="***")) - - assert result.token == "***" - - -@pytest.mark.parametrize( - ("field", "error_code"), - [("endpoint", "invalid_endpoint"), ("service_name", "invalid_service_name"), ("token", "invalid_token")], -) -def test_missing_fields_do_not_receive_legacy_defaults(field: str, error_code: str) -> None: - values = configuration_values() - del values[field] - - with pytest.raises(TencentConfigurationError) as caught: - normalize_configuration(values) - - assert caught.value.code == error_code - - -def test_console_url_is_constant_without_configuration_or_provider_io() -> None: - assert console_url() == "https://console.cloud.tencent.com/apm" - - -@dataclass -class ExportCall: - request: ExportTraceServiceRequest - metadata: tuple[tuple[str, str], ...] - timeout: float - wait_for_ready: bool - - -@dataclass -class CollectorChannel: - target: str - secure: bool - options: tuple[tuple[str, int], ...] - response_bytes: bytes - failure: Exception | None = None - fail_stub_construction: bool = False - calls: list[ExportCall] = field(default_factory=list) - close_count: int = 0 - - def unary_unary( - self, - method: str, - request_serializer: Callable[[ExportTraceServiceRequest], bytes], - response_deserializer: Callable[[bytes], ExportTraceServiceResponse], - **_kwargs: object, - ) -> Callable[..., ExportTraceServiceResponse]: - assert method == "/opentelemetry.proto.collector.trace.v1.TraceService/Export" - if self.fail_stub_construction: - raise RuntimeError("private-stub-detail synthetic-token") - - def invoke( - payload: ExportTraceServiceRequest, - metadata: Sequence[tuple[str, str]], - timeout: float, - wait_for_ready: bool, - ) -> ExportTraceServiceResponse: - assert get_value(_SUPPRESS_INSTRUMENTATION_KEY) is True - request = ExportTraceServiceRequest.FromString(request_serializer(payload)) - self.calls.append( - ExportCall(request=request, metadata=tuple(metadata), timeout=timeout, wait_for_ready=wait_for_ready) - ) - if self.failure is not None: - raise self.failure - return response_deserializer(self.response_bytes) - - return invoke - - def close(self) -> None: - self.close_count += 1 - - -@dataclass -class CollectorBoundary: - response_bytes: bytes = b"" - failure: Exception | None = None - fail_stub_construction: bool = False - fail_channel_construction: bool = False - channels: list[CollectorChannel] = field(default_factory=list) - - def open_channel(self, target: str, secure: bool, options: Sequence[tuple[str, int]]) -> CollectorChannel: - if self.fail_channel_construction: - raise RuntimeError("private-channel-detail synthetic-token") - channel = CollectorChannel( - target=target, - secure=secure, - options=tuple(options), - response_bytes=self.response_bytes, - failure=self.failure, - fail_stub_construction=self.fail_stub_construction, - ) - self.channels.append(channel) - return channel - - def secure_channel(self, target: str, credentials: object, options: Sequence[tuple[str, int]]) -> CollectorChannel: - assert credentials is not None - return self.open_channel(target, secure=True, options=options) - - def insecure_channel(self, target: str, options: Sequence[tuple[str, int]]) -> CollectorChannel: - return self.open_channel(target, secure=False, options=options) - - -@pytest.fixture -def collector(monkeypatch: pytest.MonkeyPatch, _forbid_provider_io: None) -> Iterator[CollectorBoundary]: - prior_suppression = get_value(_SUPPRESS_INSTRUMENTATION_KEY) - boundary = CollectorBoundary() - monkeypatch.setattr("grpc.secure_channel", boundary.secure_channel) - monkeypatch.setattr("grpc.insecure_channel", boundary.insecure_channel) - yield boundary - assert get_value(_SUPPRESS_INSTRUMENTATION_KEY) == prior_suppression - - -@pytest.mark.parametrize( - ("endpoint", "target", "secure"), - [ - ("https://collector.example", "collector.example:4317", True), - ("http://collector.example", "collector.example:4317", False), - ("https://collector.example:4320", "collector.example:4320", True), - ("http://collector.example:4319", "collector.example:4319", False), - ("https://[2001:db8::1]:4320", "[2001:db8::1]:4320", True), - ], -) -def test_probe_serializes_one_empty_export_and_closes_channel( - collector: CollectorBoundary, endpoint: str, target: str, secure: bool -) -> None: - config = normalize_configuration(configuration_values(endpoint=endpoint)) - - assert verify_configuration(config) == config - - assert len(collector.channels) == 1 - channel = collector.channels[0] - assert (channel.target, channel.secure) == (target, secure) - assert dict(channel.options) == {"grpc.enable_retries": 0, "grpc.max_receive_message_length": 4 * 1024 * 1024} - assert len(channel.calls) == 1 - call = channel.calls[0] - assert list(call.request.resource_spans) == [] - assert call.request.SerializeToString() == b"" - assert call.metadata == (("authorization", "Bearer synthetic-token"),) - assert call.timeout == 10 - assert call.wait_for_ready is False - assert channel.close_count == 1 - - -def test_probe_credentials_belong_to_each_request(collector: CollectorBoundary) -> None: - for token in ("first-synthetic-token", "second-synthetic-token"): - verify_configuration(normalize_configuration(configuration_values(token=token))) - - assert len(collector.channels) == 2 - assert [channel.calls[0].metadata for channel in collector.channels] == [ - (("authorization", "Bearer first-synthetic-token"),), - (("authorization", "Bearer second-synthetic-token"),), - ] - assert all(channel.close_count == 1 for channel in collector.channels) - - -@pytest.mark.parametrize( - ("rejected_spans", "error_message"), - [(1, ""), (0, "private-collector-detail synthetic-token")], -) -def test_probe_rejects_partial_response_and_closes_channel( - collector: CollectorBoundary, rejected_spans: int, error_message: str -) -> None: - response = ExportTraceServiceResponse() - response.partial_success.rejected_spans = rejected_spans - response.partial_success.error_message = error_message - collector.response_bytes = response.SerializeToString() - - with pytest.raises(TencentConfigurationError) as caught: - verify_configuration(normalize_configuration(configuration_values())) - - assert caught.value.code == "unexpected_response" - assert "private-collector-detail" not in str(caught.value) - assert "synthetic-token" not in str(caught.value) - assert collector.channels[0].close_count == 1 - - -@pytest.mark.parametrize( - "failure", - [ - grpc.RpcError("private-rpc-detail synthetic-token"), - ValueError("private-serialization-detail synthetic-token"), - ], -) -def test_probe_transport_failures_are_safe_and_close_channel(collector: CollectorBoundary, failure: Exception) -> None: - collector.failure = failure - - with pytest.raises(TencentConfigurationError) as caught: - verify_configuration(normalize_configuration(configuration_values())) - - assert caught.value.code == "connection_failed" - assert "private-" not in str(caught.value) - assert "synthetic-token" not in str(caught.value) - assert len(collector.channels[0].calls) == 1 - assert collector.channels[0].close_count == 1 - - -def test_probe_unreadable_response_is_safe_and_closes_channel(collector: CollectorBoundary) -> None: - collector.response_bytes = b"\xff" - - with pytest.raises(TencentConfigurationError) as caught: - verify_configuration(normalize_configuration(configuration_values())) - - assert caught.value.code == "connection_failed" - assert collector.channels[0].close_count == 1 - - -def test_probe_stub_construction_failure_is_safe_and_closes_channel(collector: CollectorBoundary) -> None: - collector.fail_stub_construction = True - - with pytest.raises(TencentConfigurationError) as caught: - verify_configuration(normalize_configuration(configuration_values())) - - assert caught.value.code == "connection_failed" - assert "private-stub-detail" not in str(caught.value) - assert "synthetic-token" not in str(caught.value) - assert collector.channels[0].calls == [] - assert collector.channels[0].close_count == 1 - - -def test_probe_channel_construction_failure_has_fixed_safe_error(collector: CollectorBoundary) -> None: - collector.fail_channel_construction = True - - with pytest.raises(TencentConfigurationError) as caught: - verify_configuration(normalize_configuration(configuration_values())) - - assert caught.value.code == "connection_failed" - assert "private-channel-detail" not in str(caught.value) - assert "synthetic-token" not in str(caught.value) - assert collector.channels == [] diff --git a/api/tests/unit_tests/providers/trace/trace-tencent/src/dify_trace_tencent/test_utils.py b/api/tests/unit_tests/providers/trace/trace-tencent/src/dify_trace_tencent/test_utils.py deleted file mode 100644 index 7ab192b990396a..00000000000000 --- a/api/tests/unit_tests/providers/trace/trace-tencent/src/dify_trace_tencent/test_utils.py +++ /dev/null @@ -1,194 +0,0 @@ -"""Pure legacy Tencent session regressions; IDs here are builder-local IDs.""" - -from datetime import UTC, datetime, timedelta - -import pytest -from dify_trace_tencent.entities.tencent_trace_entity import SpanData -from dify_trace_tencent.span_builder import TencentSpanBuilder -from dify_trace_tencent.utils import TencentTraceUtils -from opentelemetry.trace import StatusCode -from pydantic import ValidationError - -from core.ops.entities.trace_entity import WorkflowTraceInfo -from graphon.entities import WorkflowNodeExecution -from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionStatus - -RUN_ID = "00000000-0000-4000-8000-000000000001" -WORKFLOW_ID = "00000000-0000-4000-8000-000000000002" -TENANT_ID = "00000000-0000-4000-8000-000000000003" -STARTED = datetime(2026, 10, 17, tzinfo=UTC) - - -def workflow_trace(metadata: dict[str, object]) -> WorkflowTraceInfo: - session_id = metadata.get("conversation_id") - return WorkflowTraceInfo( - workflow_id=WORKFLOW_ID, - tenant_id=TENANT_ID, - workflow_run_id=RUN_ID, - workflow_run_elapsed_time=3.0, - workflow_run_status="succeeded", - workflow_run_inputs={"query": "synthetic input"}, - workflow_run_outputs={"answer": "synthetic answer"}, - workflow_run_version="1", - error="", - total_tokens=32, - prompt_tokens=13, - completion_tokens=19, - file_list=[], - query="synthetic input", - metadata=metadata, - conversation_id=session_id if isinstance(session_id, str) else None, - start_time=STARTED, - end_time=STARTED + timedelta(seconds=3), - ) - - -def workflow_node(node_type: str, index: int = 1) -> WorkflowNodeExecution: - return WorkflowNodeExecution( - id=f"00000000-0000-4000-8000-{index + 10:012d}", - workflow_id=WORKFLOW_ID, - workflow_execution_id=RUN_ID, - index=index, - node_id=f"node-{index}", - node_type=node_type, - title={BuiltinNodeTypes.START: "Start", BuiltinNodeTypes.END: "End"}.get(node_type, "Synthetic node"), - inputs={"query": "synthetic input"}, - process_data={"model_name": "synthetic-model", "model_provider": "synthetic-provider"}, - outputs={ - "text": "synthetic answer", - "usage": {"prompt_tokens": 13, "completion_tokens": 19, "total_tokens": 32}, - "finish_reason": "stop", - }, - status=WorkflowNodeExecutionStatus.SUCCEEDED, - created_at=STARTED + timedelta(seconds=index - 1), - finished_at=STARTED + timedelta(seconds=index), - ) - - -def build_node_span(trace: WorkflowTraceInfo, node: WorkflowNodeExecution) -> SpanData: - trace_id = TencentTraceUtils.convert_to_trace_id(trace.workflow_run_id) - workflow_span_id = TencentTraceUtils.convert_to_span_id(trace.workflow_run_id, "workflow") - if node.node_type == BuiltinNodeTypes.LLM: - return TencentSpanBuilder.build_workflow_llm_span(trace_id, workflow_span_id, trace, node) - if node.node_type == BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL: - return TencentSpanBuilder.build_workflow_retrieval_span(trace_id, workflow_span_id, trace, node) - return TencentSpanBuilder.build_workflow_task_span(trace_id, workflow_span_id, trace, node) - - -@pytest.mark.parametrize( - "node_type", [BuiltinNodeTypes.START, BuiltinNodeTypes.LLM, BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL] -) -@pytest.mark.parametrize( - ("metadata", "expected_session"), - [ - pytest.param({}, "", id="missing"), - pytest.param({"conversation_id": None}, "", id="none"), - pytest.param({"conversation_id": ""}, "", id="empty"), - pytest.param({"conversation_id": "owned-conversation"}, "owned-conversation", id="conversation"), - ], -) -def test_workflow_builders_accept_optional_sessions_without_rewriting_metadata( - node_type: str, metadata: dict[str, object], expected_session: str -) -> None: - trace = workflow_trace(metadata) - before = dict(trace.metadata) - - span = build_node_span(trace, workflow_node(node_type)) - - assert span.attributes["gen_ai.session.id"] == expected_session - assert trace.metadata == before - - -def complete_workflow_spans(trace: WorkflowTraceInfo) -> list[SpanData]: - spans = TencentSpanBuilder.build_workflow_spans( - trace, TencentTraceUtils.convert_to_trace_id(trace.workflow_run_id), "synthetic-user" - ) - spans.extend( - build_node_span(trace, workflow_node(node_type, index)) - for index, node_type in enumerate((BuiltinNodeTypes.START, BuiltinNodeTypes.LLM, BuiltinNodeTypes.END), start=1) - ) - return spans - - -def assert_usage_is_not_duplicated(spans: list[SpanData]) -> None: - generations = [span for span in spans if span.attributes.get("gen_ai.span.kind") == "GENERATION"] - assert len(generations) == 1 - generation = generations[0] - assert generation.attributes["gen_ai.usage.input_tokens"] == "13" - assert generation.attributes["gen_ai.usage.output_tokens"] == "19" - assert generation.attributes["gen_ai.usage.total_tokens"] == "32" - assert sum(int(span.attributes.get("gen_ai.usage.total_tokens", "0")) for span in spans) == 32 - - -def test_workflow_without_conversation_keeps_start_llm_end_under_one_root() -> None: - trace = workflow_trace({"conversation_id": None}) - - spans = complete_workflow_spans(trace) - - assert [span.name for span in spans] == ["workflow", "Start", "GENERATION", "End"] - root = spans[0] - assert root.parent_span_id is None - assert all(span.parent_span_id == root.span_id for span in spans[1:]) - assert all(span.trace_id == root.trace_id for span in spans) - assert len({span.span_id for span in spans}) == 4 - assert all(span.attributes["gen_ai.session.id"] == "" for span in spans[1:]) - assert trace.metadata["conversation_id"] is None - assert_usage_is_not_duplicated(spans) - - -def test_chatflow_keeps_conversation_and_message_workflow_parent_layer() -> None: - trace = workflow_trace({"conversation_id": "owned-conversation"}) - - spans = complete_workflow_spans(trace) - - assert [span.name for span in spans] == ["message", "workflow", "Start", "GENERATION", "End"] - message, workflow = spans[:2] - assert message.parent_span_id is None - assert workflow.parent_span_id == message.span_id - assert all(span.parent_span_id == workflow.span_id for span in spans[2:]) - assert all(span.trace_id == message.trace_id for span in spans) - assert len({span.span_id for span in spans}) == 5 - assert message.attributes["gen_ai.session.id"] == "owned-conversation" - assert all(span.attributes["gen_ai.session.id"] == "owned-conversation" for span in spans[2:]) - assert_usage_is_not_duplicated(spans) - - -@pytest.mark.parametrize("status", [WorkflowNodeExecutionStatus.FAILED, WorkflowNodeExecutionStatus.EXCEPTION]) -def test_error_status_and_description_survive_missing_session(status: WorkflowNodeExecutionStatus) -> None: - trace = workflow_trace({"conversation_id": None}) - trace.error = "synthetic workflow failure" - trace.workflow_run_status = "failed" - node = workflow_node(BuiltinNodeTypes.LLM) - node.status = status - node.error = "synthetic node failure" - - root = TencentSpanBuilder.build_workflow_spans( - trace, TencentTraceUtils.convert_to_trace_id(trace.workflow_run_id), "synthetic-user" - )[0] - span = build_node_span(trace, node) - - assert root.status.status_code == StatusCode.ERROR - assert root.status.description == "synthetic workflow failure" - assert span.status.status_code == StatusCode.ERROR - assert span.status.description == "synthetic node failure" - assert span.parent_span_id == root.span_id - assert span.attributes["gen_ai.session.id"] == "" - assert_usage_is_not_duplicated([root, span]) - - -@pytest.mark.parametrize( - "node_type", [BuiltinNodeTypes.START, BuiltinNodeTypes.LLM, BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL] -) -@pytest.mark.parametrize( - "invalid_session", [0, False, [], {}, 123], ids=["zero", "false", "empty-list", "empty-map", "int"] -) -def test_non_string_sessions_remain_validation_errors(node_type: str, invalid_session: object) -> None: - trace = workflow_trace({"conversation_id": invalid_session}) - - with pytest.raises(ValidationError) as caught: - build_node_span(trace, workflow_node(node_type)) - - errors = caught.value.errors(include_input=False, include_url=False) - assert len(errors) == 1 - assert errors[0]["loc"] == ("attributes", "gen_ai.session.id") - assert errors[0]["type"] == "string_type" diff --git a/api/tests/unit_tests/providers/trace/trace-weave/src/dify_trace_weave/test_configuration.py b/api/tests/unit_tests/providers/trace/trace-weave/src/dify_trace_weave/test_configuration.py deleted file mode 100644 index 63d4983119b7e6..00000000000000 --- a/api/tests/unit_tests/providers/trace/trace-weave/src/dify_trace_weave/test_configuration.py +++ /dev/null @@ -1,258 +0,0 @@ -"""Read-only Weave validation at the real config/schema and HTTP boundaries.""" - -from collections.abc import Callable -from dataclasses import dataclass, field -from typing import cast - -import httpx -import pytest -from dify_trace_weave.config import WeaveConfig -from dify_trace_weave.configuration import ( - WeaveConfigurationError, - console_url, - project_url, - verify_configuration, -) -from weave.trace_server.trace_server_interface import CallsQueryStatsReq - -from core.helper import ssrf_proxy - - -def object_dict(value: object) -> dict[str, object]: - assert isinstance(value, dict) - return cast(dict[str, object], value) - - -def response(url: str, body: object, *, status: int = 200) -> httpx.Response: - return httpx.Response(status, json=body, request=httpx.Request("POST", url)) - - -@dataclass -class ReadAPI: - calls: list[tuple[str, dict[str, object]]] = field(default_factory=list) - override: Callable[[str, dict[str, object]], httpx.Response] | None = None - project_fields: dict[str, object] = field(default_factory=dict) - project_exists: bool = True - viewer_exists: bool = True - count: object = 0 - - def post(self, url: str, **options: object) -> httpx.Response: - self.calls.append((url, options)) - body = object_dict(options["json"]) - if self.override is not None: - return self.override(url, body) - if url.endswith("/graphql"): - query = body["query"] - assert isinstance(query, str) - assert "mutation" not in query - if "DifyWeaveViewer" in query: - return response(url, {"data": {"viewer": {"id": "viewer-id", "entity": "default-team"}}}) - variables = object_dict(body["variables"]) - project: dict[str, object] = { - "id": "project-id", - "name": variables["project"], - "entityName": variables["entity"], - "readOnly": False, - } - project.update(self.project_fields) - return response( - url, - { - "data": { - "viewer": {"id": "viewer-id"} if self.viewer_exists else None, - "project": project if self.project_exists else None, - } - }, - ) - assert url.endswith("/calls/query_stats") - query_request = CallsQueryStatsReq.model_validate(body) - assert query_request.limit == 1 - return response(url, {"count": self.count}) - - -@pytest.fixture -def read_api(monkeypatch: pytest.MonkeyPatch) -> ReadAPI: - api = ReadAPI() - monkeypatch.setattr(ssrf_proxy, "post", api.post) - return api - - -@pytest.mark.parametrize( - ("entity", "project", "expected_entity", "requests"), - [("team", "demo", "team", 2), (None, "team/demo", "team", 2), (None, "demo", "default-team", 3)], -) -def test_verify_resolves_project_with_request_scoped_authentication( - read_api: ReadAPI, entity: str | None, project: str, expected_entity: str, requests: int -) -> None: - result = verify_configuration(WeaveConfig(api_key="synthetic-key", entity=entity, project=project)) - assert result.entity == expected_entity - assert result.project == "demo" - assert result.project_url == f"https://wandb.ai/{expected_entity}/demo" - assert len(read_api.calls) == requests - for _url, options in read_api.calls: - assert options["auth"] == ("api", "synthetic-key") - assert options["max_retries"] == 0 - assert options["follow_redirects"] is False - assert options["timeout"] == 30.0 - assert object_dict(read_api.calls[-1][1]["json"])["project_id"] == f"{expected_entity}/demo" - - -@pytest.mark.parametrize( - ("entity", "project"), - [ - ("team", ""), - ("team", ".."), - ("/", "demo"), - (None, "/demo"), - (None, "team/"), - (None, "team/sub/demo"), - ("other", "team/demo"), - ("team", "de\\mo"), - ], -) -def test_invalid_target_is_rejected_before_network(read_api: ReadAPI, entity: str | None, project: str) -> None: - with pytest.raises(WeaveConfigurationError) as caught: - verify_configuration(WeaveConfig(api_key="synthetic-key", entity=entity, project=project)) - assert caught.value.code == "invalid_project" - assert read_api.calls == [] - - -def test_private_endpoints_preserve_paths_and_encode_project_links(read_api: ReadAPI) -> None: - config = WeaveConfig( - api_key="synthetic-key", - entity="team one", - project="demo & notes#", - host="http://wandb.internal/wandb/", - endpoint="https://trace.internal/proxy/", - ) - result = verify_configuration(config) - assert result.project_url == "http://wandb.internal/wandb/team%20one/demo%20%26%20notes%23" - assert [url for url, _options in read_api.calls] == [ - "http://wandb.internal/wandb/graphql", - "https://trace.internal/proxy/calls/query_stats", - ] - - -@pytest.mark.parametrize( - "host", ["https://user:secret@host.test", "https://host.test/?secret=x", "https://host.test/#x"] -) -def test_unsafe_host_is_rejected_without_authentication_request(read_api: ReadAPI, host: str) -> None: - with pytest.raises(WeaveConfigurationError) as caught: - verify_configuration(WeaveConfig(api_key="synthetic-key", entity="team", project="demo", host=host)) - assert caught.value.code == "invalid_endpoint" - assert read_api.calls == [] - - -@pytest.mark.parametrize( - ("status", "code"), - [ - (401, "authentication_failed"), - (403, "permission_denied"), - (302, "endpoint_unavailable"), - (404, "endpoint_unavailable"), - (429, "endpoint_unavailable"), - (503, "endpoint_unavailable"), - ], -) -def test_http_errors_are_sanitized_and_never_retried(read_api: ReadAPI, status: int, code: str) -> None: - read_api.override = lambda url, _body: response(url, {"error": "never-log-secret"}, status=status) - with pytest.raises(WeaveConfigurationError) as caught: - verify_configuration(WeaveConfig(api_key="never-log-secret", entity="team", project="demo")) - assert caught.value.code == code - assert "never-log-secret" not in str(caught.value) - assert len(read_api.calls) == 1 - - -@pytest.mark.parametrize( - ("fields", "code"), - [ - ({"readOnly": True}, "project_read_only"), - ({"readOnly": None}, "project_permission_unknown"), - ({"readOnly": "false"}, "malformed_response"), - ({"id": ""}, "malformed_response"), - ({"name": "other"}, "invalid_project"), - ({"entityName": "other"}, "invalid_project"), - ], -) -def test_project_permission_and_identity_are_checked(read_api: ReadAPI, fields: dict[str, object], code: str) -> None: - read_api.project_fields = fields - with pytest.raises(WeaveConfigurationError) as caught: - verify_configuration(WeaveConfig(api_key="synthetic-key", entity="team", project="demo")) - assert caught.value.code == code - assert len(read_api.calls) == 1 - - -def test_missing_project_is_not_created(read_api: ReadAPI) -> None: - read_api.project_exists = False - with pytest.raises(WeaveConfigurationError) as caught: - verify_configuration(WeaveConfig(api_key="synthetic-key", entity="team", project="demo")) - assert caught.value.code == "project_not_found_or_inaccessible" - assert len(read_api.calls) == 1 - - -def test_anonymous_viewer_cannot_validate_a_public_project(read_api: ReadAPI) -> None: - read_api.viewer_exists = False - with pytest.raises(WeaveConfigurationError) as caught: - verify_configuration(WeaveConfig(api_key="synthetic-key", entity="team", project="demo")) - assert caught.value.code == "authentication_failed" - - -@pytest.mark.parametrize("count", [-1, True, "0", None]) -def test_invalid_trace_endpoint_response_is_rejected(read_api: ReadAPI, count: object) -> None: - read_api.count = count - with pytest.raises(WeaveConfigurationError) as caught: - verify_configuration(WeaveConfig(api_key="synthetic-key", entity="team", project="demo")) - assert caught.value.code == "malformed_response" - assert len(read_api.calls) == 2 - - -@pytest.mark.parametrize( - ("body", "code"), - [ - ({"errors": [{"message": "never-log-secret"}]}, "graphql_error"), - ({"errors": "never-log-secret"}, "malformed_response"), - ({"data": []}, "malformed_response"), - ({"data": {}}, "malformed_response"), - ], -) -def test_graphql_errors_and_invalid_envelopes_are_rejected( - read_api: ReadAPI, body: dict[str, object], code: str -) -> None: - read_api.override = lambda url, _body: response(url, body) - with pytest.raises(WeaveConfigurationError) as caught: - verify_configuration(WeaveConfig(api_key="never-log-secret", entity="team", project="demo")) - assert caught.value.code == code - assert "never-log-secret" not in str(caught.value) - - -def test_connection_failure_has_no_raw_exception_or_secret(read_api: ReadAPI) -> None: - def fail(_url: str, _body: dict[str, object]) -> httpx.Response: - raise httpx.ReadTimeout("never-log-secret") - - read_api.override = fail - with pytest.raises(WeaveConfigurationError) as caught: - verify_configuration(WeaveConfig(api_key="synthetic-key", entity="team", project="demo")) - assert caught.value.code == "connection_failed" - assert "never-log-secret" not in str(caught.value) - assert caught.value.__cause__ is None - - -def test_sequential_tenants_use_only_their_supplied_credentials(read_api: ReadAPI) -> None: - for tenant in ("a", "b"): - verify_configuration(WeaveConfig(api_key=f"synthetic-{tenant}", entity=f"team-{tenant}", project="demo")) - assert [options["auth"] for _url, options in read_api.calls] == [ - ("api", "synthetic-a"), - ("api", "synthetic-a"), - ("api", "synthetic-b"), - ("api", "synthetic-b"), - ] - - -def test_project_url_and_console_fallback_are_pure(read_api: ReadAPI) -> None: - assert project_url(WeaveConfig(api_key="synthetic", project="demo")) is None - assert project_url(WeaveConfig(api_key="synthetic", project="team/demo")) == "https://wandb.ai/team/demo" - assert console_url(WeaveConfig(api_key="synthetic", project="demo")) == "https://wandb.ai" - assert console_url(WeaveConfig(api_key="synthetic", project="demo", host="https://private.test/wandb")) == ( - "https://private.test/wandb" - ) - assert read_api.calls == [] diff --git a/api/tests/unit_tests/pyrefly.toml b/api/tests/unit_tests/pyrefly.toml index bdf577fc59823e..de4694d4e92478 100644 --- a/api/tests/unit_tests/pyrefly.toml +++ b/api/tests/unit_tests/pyrefly.toml @@ -176,11 +176,6 @@ project-excludes = [ "core/memory/test_token_buffer_memory.py", "core/plugin/impl/test_model_runtime_factory.py", "core/moderation/test_input_moderation.py", - "core/ops/test_base_trace_instance.py", - "core/ops/entities/test_config_entity.py", - "core/ops/test_ops_trace_manager.py", - "core/ops/entities/test_trace_entity.py", - "core/ops/test_utils.py", "core/plugin/impl/test_agent.py", "core/plugin/impl/test_base.py", "core/plugin/impl/test_datasource.py", @@ -223,7 +218,6 @@ project-excludes = [ "core/repositories/test_human_input_repository.py", "core/repositories/test_sqlalchemy_workflow_execution_repository.py", "core/repositories/test_sqlalchemy_workflow_node_execution_repository.py", - "core/telemetry/test_gateway.py", "../e2e_tests/contracts/core/test_file.py", "core/test_model_manager.py", "core/test_provider_manager.py", @@ -303,7 +297,6 @@ project-excludes = [ "core/workflow/test_workflow_entry.py", "dev/test_generate_knowledge_fs_contract.py", "enterprise/telemetry/test_contracts.py", - "enterprise/telemetry/test_draft_trace.py", "enterprise/telemetry/test_enterprise_trace.py", "events/test_app_event.py", "events/test___init__.py", diff --git a/api/tests/unit_tests/repositories/app/test_tracing_config_repository.py b/api/tests/unit_tests/repositories/app/test_tracing_config_repository.py index 7cf8c93d530710..24e70b4d7f558c 100644 --- a/api/tests/unit_tests/repositories/app/test_tracing_config_repository.py +++ b/api/tests/unit_tests/repositories/app/test_tracing_config_repository.py @@ -9,13 +9,14 @@ from machinery.context import RequestContext from models.model import App, AppMode, TraceAppConfig from repositories.app.console_repository import ConsoleAppRepository -from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository +from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository, trace_settings_from_column from services.app_tracing_config_service import ( AppTraceSettings, AppTracingConfigAppNotFoundError, AppTracingConfigConflictError, AppTracingConfigRecord, ) +from services.tracing.delivery import StoredTraceDestination from tests.helpers.unit.repositories.app.console_visibility import UNADDRESSABLE_IN_WORKSPACE, MakeUnaddressable _APP_ID = "11111111-1111-1111-1111-111111111111" @@ -205,3 +206,88 @@ def test_capture_policy_round_trip_and_legacy_defaults(sqlite_session_factory: s with pytest.raises(AppTracingConfigAppNotFoundError): writer.set_trace(workspace_id="other-tenant", app_id="app", settings=AppTraceSettings(), expected=None) assert repository.get_trace(context, "app") == policy + + +"""The trace pipeline reads the stored policy leniently and only through the app's workspace.""" + + +@pytest.mark.parametrize( + ("stored", "settings"), + [ + ('{"enabled": true, "tracing_provider": "langfuse"}', AppTraceSettings(True, "langfuse", False, True)), + ( + '{"enabled": true, "tracing_provider": "langfuse", "capture_content": true, "capture_agent": false}', + AppTraceSettings(True, "langfuse", True, False), + ), + ('{"enabled": false, "tracing_provider": "langfuse"}', AppTraceSettings(False, "langfuse", False, True)), + ('{"enabled": "true", "tracing_provider": 7}', AppTraceSettings()), + ('["enabled"]', AppTraceSettings()), + ("not json", AppTraceSettings()), + (None, AppTraceSettings()), + ], + ids=["enabled", "capture-policy", "disabled", "wrong-types", "not-an-object", "unparsable", "unset"], +) +def test_stored_policies_that_cannot_be_read_disable_tracing(stored: str | None, settings: AppTraceSettings) -> None: + assert trace_settings_from_column(stored) == settings + + +def _persist_traced_app(session: Session, tracing: str | None) -> None: + _persist_app(session) + app = session.get(App, _APP_ID) + assert app is not None + app.tracing = tracing + session.commit() + + +def test_policy_is_read_through_the_apps_workspace( + sqlite_session: Session, sqlite_session_factory: sessionmaker[Session] +) -> None: + _persist_traced_app(sqlite_session, json.dumps({"enabled": True, "tracing_provider": _PROVIDER})) + repository = _repository(sqlite_session_factory) + + assert repository.policy(tenant_id=_WORKSPACE_ID, app_id=_APP_ID) == AppTraceSettings(True, _PROVIDER, False, True) + assert repository.policy(tenant_id=_OTHER_WORKSPACE_ID, app_id=_APP_ID) is None + assert repository.policy(tenant_id=_WORKSPACE_ID, app_id="missing-app") is None + + +def test_the_active_destination_is_the_enabled_providers_active_configuration( + sqlite_session: Session, sqlite_session_factory: sessionmaker[Session] +) -> None: + _persist_traced_app(sqlite_session, json.dumps({"enabled": True, "tracing_provider": _PROVIDER})) + sqlite_session.add_all( + [ + TraceAppConfig(app_id=_APP_ID, tracing_provider=_PROVIDER, tracing_config={"public_key": "ciphertext"}), + TraceAppConfig(app_id=_APP_ID, tracing_provider="opik", tracing_config={"api_key": "other"}), + ] + ) + sqlite_session.commit() + repository = _repository(sqlite_session_factory) + + assert repository.active_destination(tenant_id=_WORKSPACE_ID, app_id=_APP_ID) == StoredTraceDestination( + provider_id=_PROVIDER, config={"public_key": "ciphertext"} + ) + assert repository.active_destination(tenant_id=_OTHER_WORKSPACE_ID, app_id=_APP_ID) is None + + +@pytest.mark.parametrize( + ("tracing", "active"), + [ + (json.dumps({"enabled": False, "tracing_provider": _PROVIDER}), True), + (json.dumps({"enabled": True, "tracing_provider": None}), True), + (json.dumps({"enabled": True, "tracing_provider": _PROVIDER}), False), + (None, True), + ], + ids=["disabled", "no-provider", "inactive-config", "unset"], +) +def test_apps_that_do_not_trace_have_no_active_destination( + sqlite_session: Session, sqlite_session_factory: sessionmaker[Session], tracing: str | None, active: bool +) -> None: + _persist_traced_app(sqlite_session, tracing) + sqlite_session.add( + TraceAppConfig( + app_id=_APP_ID, tracing_provider=_PROVIDER, tracing_config={"public_key": "ciphertext"}, is_active=active + ) + ) + sqlite_session.commit() + + assert _repository(sqlite_session_factory).active_destination(tenant_id=_WORKSPACE_ID, app_id=_APP_ID) is None diff --git a/api/tests/unit_tests/repositories/test_application_trace_repository.py b/api/tests/unit_tests/repositories/test_application_trace_repository.py index d9eb6469526faa..86ebed894751c4 100644 --- a/api/tests/unit_tests/repositories/test_application_trace_repository.py +++ b/api/tests/unit_tests/repositories/test_application_trace_repository.py @@ -1,36 +1,349 @@ import json -from collections.abc import Iterator -from dataclasses import dataclass -from datetime import timedelta -from typing import Literal -from uuid import UUID +from collections.abc import Callable, Generator, Iterator +from contextlib import contextmanager +from dataclasses import asdict, dataclass, replace +from datetime import UTC, datetime, timedelta +from decimal import Decimal +from typing import Literal, override +from uuid import NAMESPACE_URL, UUID, uuid4, uuid5 import pytest from aliyun.log.logclient import LogClient +from dify_agent.protocol.tracing import AgentTraceBatch, AgentTraceSpan +from dify_trace_sdk.events import TraceEvent, WorkflowRunCompleted +from dify_trace_sdk.model import Span, SpanKind, SpanStatus, Trace +from dify_trace_sdk.testing import RecordingEgress from sqlalchemy import Connection, Engine, delete, event, func, select +from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session, sessionmaker -from core.ops.entities.trace_entity import ModerationTraceInfo +from domain.tracing.errors import TraceRunConflictError +from domain.tracing.trace_run_agg import TraceBatch, TraceInvocation, TraceRunAgg, TraceRunState from extensions.logstore.aliyun_logstore import AliyunLogStore +from extensions.logstore.repositories.logstore_api_workflow_node_execution_repository import ( + _dict_to_workflow_node_execution_model, +) from extensions.logstore.repositories.logstore_api_workflow_run_repository import _dict_to_workflow_run from libs.datetime_utils import naive_utc_now +from libs.uuid_utils import uuidv7 from models.application_trace import ApplicationTraceBatch, ApplicationTraceInvocation, ApplicationTraceRun -from models.model import App, AppMode +from models.enums import ConversationFromSource +from models.model import App, AppMode, Message, TraceAppConfig from models.workflow import WorkflowNodeExecutionModel, WorkflowRun -from repositories.application_trace_repository import TraceCollectionError -from services.tracing.provider_dispatch import TraceExportResult -from services.tracing.recovery import _workflow_task +from repositories.application_trace_repository import ApplicationTraceRepository +from repositories.factory import DifyAPIRepositoryFactory +from repositories.trace_record_repository import TraceRecordRepository +from services.tracing import service as service_module +from services.tracing.canonical import CanonicalTraceBuilder +from services.tracing.collection import AppTracingSnapshot, TraceRun, TraceRunStatus +from services.tracing.collection_ports import LockedTraceRun +from services.tracing.delivery import TraceRecordSource +from services.tracing.errors import ( + TraceCollectionConflictError, + TraceCollectionError, + TraceCollectionNotFoundError, + TracingStoreUnavailableError, +) +from services.tracing.providers import ProviderConfigCodec +from services.tracing.records import ( + BaseTraceInfo, + MessageTraceInfo, + ModerationTraceInfo, + TraceMessage, + WorkflowTraceInfo, +) +from services.tracing.recovery import recover_terminal from services.tracing.service import ApplicationTracingService from tests.helpers.unit.config_override import apply_config_overrides -from tests.helpers.unit.model_factories import make_workflow -from tests.unit_tests.services.tracing.test_runtime_bridge import AGENT_GRAPH, workflow_info -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_batch_for as batch_for -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_claim as claim -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_identifier as identifier -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_ingest as ingest -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_seed as seed -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_terminal as terminal -from tests.unit_tests.services.tracing.test_service import _ApplicationTraceCollectionCollection as Collection +from tests.helpers.unit.model_factories import make_app, make_conversation, make_message, make_workflow +from tests.helpers.unit.services.tracing_ports import ( + AGENT_TRACE_SETTINGS, + InstalledTraceProviders, + WorkspaceBoundCipher, +) + +"""Seeded Agent trace collections, shared with the ingestion handler and worker task tests.""" + +AGENT_GRAPH = { + "nodes": [{"id": "agent", "data": {"type": "agent", "version": "2", "agent_node_kind": "dify_agent"}}], + "edges": list[object](), +} +LANGFUSE_POLICY = {"enabled": True, "tracing_provider": "langfuse", "capture_content": False, "capture_agent": True} +"""The tracing policy every seeded app stores.""" + + +def identifier(value: str) -> str: + return str(uuid5(NAMESPACE_URL, value)) + + +def trace_records(sessions: sessionmaker[Session]) -> TraceRecordRepository: + return TraceRecordRepository( + sessions, + workflow_runs=DifyAPIRepositoryFactory.create_api_workflow_run_repository(sessions), + node_executions=DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(sessions), + files_url="https://files.test", + enterprise_details=False, + ) + + +def collector( + repository: ApplicationTraceRepository, records: TraceRecordSource, scheduled: list[str] +) -> ApplicationTracingService: + """Agent trace collection over ``repository`` that schedules exports into ``scheduled``. + + No provider is installed, so a delivery stops before it reaches any destination. + """ + return ApplicationTracingService( + runs=repository, + run_query=repository, + sources=repository, + records=records, + builder=CanonicalTraceBuilder(), + providers=InstalledTraceProviders(), + codec=ProviderConfigCodec(WorkspaceBoundCipher()), + egress=RecordingEgress([], grpc_target=None), + schedule_export=scheduled.append, + settings=AGENT_TRACE_SETTINGS, + ) + + +@dataclass +class Collection: + repository: ApplicationTraceRepository + sessions: sessionmaker[Session] + run: TraceRun + scheduled: list[str] + """The runs this collection's service scheduled for export.""" + + def service(self) -> ApplicationTracingService: + return collector(self.repository, trace_records(self.sessions), self.scheduled) + + def invocation(self) -> TraceInvocation: + return open_invocation(self, node_execution_id=None, call_key="first") + + +def seed(sessions: sessionmaker[Session], name: str) -> Collection: + """Start collecting an Agent app message whose run freezes a Langfuse destination. + + The frozen secrets decrypt with ``WorkspaceBoundCipher`` in the seeded workspace. + """ + app_id, tenant_id, message_id = identifier(f"app-{name}"), identifier("tenant"), identifier(f"message-{name}") + cipher = WorkspaceBoundCipher() + with sessions.begin() as session: + session.add( + make_app(app_id=app_id, tenant_id=tenant_id, mode=AppMode.AGENT, tracing=json.dumps(LANGFUSE_POLICY)) + ) + session.add( + make_message( + message_id=message_id, + app_id=app_id, + conversation_id=identifier(f"conversation-{name}"), + inputs={}, + query="hello", + message={}, + answer="done", + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.CONSOLE, + ) + ) + session.add( + TraceAppConfig( + app_id=app_id, + tracing_provider="langfuse", + tracing_config={ + "public_key": cipher.encrypt(tenant_id, "public"), + "secret_key": cipher.encrypt(tenant_id, "original"), + "host": "https://langfuse.test", + "project_key": "original", + }, + ) + ) + repository = ApplicationTraceRepository(sessions) + scheduled: list[str] = [] + run = collector(repository, trace_records(sessions), scheduled).start( + tenant_id=tenant_id, + app_id=app_id, + source_kind="message", + source_id=message_id, + session_id="session-original", + workflow_id=None, + ) + assert run is not None + return Collection(repository, sessions, run, scheduled) + + +def open_invocation(collection: Collection, *, node_execution_id: str | None, call_key: str) -> TraceInvocation: + with collection.repository.change(collection.run.id, tenant_id=collection.run.tenant_id) as locked: + assert locked is not None + return locked.run.open_invocation( + agent_id=identifier("agent"), node_execution_id=node_execution_id, attempt=0, call_key=call_key + ) + + +def batch_for(run: TraceRun, invocation: TraceInvocation) -> AgentTraceBatch: + return AgentTraceBatch( + run_id=UUID(run.id), + invocation_id=UUID(invocation.id), + batch_id=uuid4(), + final=True, + spans=[ + AgentTraceSpan( + trace_id=run.trace_id, + span_id="a" * 16, + parent_span_id=invocation.span_id, + name="Agent execution", + start_time_ns=1_000_000_000, + end_time_ns=2_000_000_000, + status="ok", + attributes={"gen_ai.operation.name": "invoke_agent"}, + ) + ], + ) + + +def ingest(collection: Collection, batch: AgentTraceBatch) -> int: + return collection.service().ingest(batch) + + +def store_batch(collection: Collection, batch: AgentTraceBatch) -> None: + """Ingest ``batch`` into the collection's run inside one change of the run.""" + with collection.repository.change(collection.run.id, tenant_id=collection.run.tenant_id) as locked: + assert locked is not None + locked.run.ingest(batch, max_spans=100, max_bytes=1_000_000, now=naive_utc_now()) + + +def agent_turn(run: TraceRun, *, root_span_id: str) -> Trace: + """A business terminal of one Agent turn under ``root_span_id``.""" + root = Span( + span_id=root_span_id, + parent_span_id=None, + name="Agent turn", + kind=SpanKind.CHAIN, + start_time=datetime.fromtimestamp(0, UTC), + end_time=datetime.fromtimestamp(3, UTC), + status=SpanStatus.OK, + ) + return Trace( + trace_id=run.trace_id, + external_trace_id=None, + session_id=run.session_id, + user_id=None, + parent=None, + spans=(root,), + ) + + +def capture_terminal(collection: Collection, trace: Trace) -> None: + with collection.repository.change(collection.run.id, tenant_id=collection.run.tenant_id) as locked: + assert locked is not None + locked.run.capture_terminal(trace, now=naive_utc_now()) + + +def terminal(collection: Collection) -> None: + capture_terminal(collection, agent_turn(collection.run, root_span_id=collection.run.root_span_id)) + + +def claim(collection: Collection) -> str | None: + """Lease the collection's run for one export attempt; the token fencing it, ``None`` when it cannot export.""" + token = str(uuid4()) + with collection.repository.change(collection.run.id, tenant_id=collection.run.tenant_id) as locked: + assert locked is not None + claimed = locked.run.claim_export( + token=token, owner_exists=locked.owner_exists, collection_timeout=60, lease_seconds=30, now=naive_utc_now() + ) + return token if claimed else None + + +def recover(collection: Collection, service: ApplicationTracingService, records: TraceRecordSource) -> bool: + """One reconciler recovery attempt of the collection's run, collected by ``service``.""" + run = collection.repository.load(collection.run.id) + assert run is not None + return recover_terminal( + run, + runs=collection.repository, + sources=collection.repository, + records=records, + collector=service, + collection_timeout=AGENT_TRACE_SETTINGS.collection_timeout_seconds, + now=naive_utc_now(), + ) + + +def message_info(collection: Collection) -> MessageTraceInfo: + """The completed message record of the collection's message.""" + with collection.sessions() as session: + message = session.get(Message, collection.run.source_id) + assert message is not None + data = TraceMessage( + id=message.id, + app_id=message.app_id, + conversation_id=message.conversation_id, + workflow_run_id=message.workflow_run_id, + query=message.query, + answer=message.answer, + error=message.error, + created_at=message.created_at, + model_provider=message.model_provider, + model_id=message.model_id, + total_price=message.total_price, + currency=message.currency, + ) + return MessageTraceInfo( + message_id=data.id, + message_data=data, + inputs={"query": data.query}, + outputs=data.answer, + start_time=data.created_at, + end_time=data.created_at, + metadata={"app_id": data.app_id}, + conversation_model="chat", + message_tokens=0, + answer_tokens=0, + total_tokens=0, + error=None, + file_list=[], + conversation_mode=AppMode.AGENT, + gen_ai_server_time_to_first_token=None, + llm_streaming_time_to_generate=None, + is_streaming_request=False, + ) + + +def workflow_info(row: WorkflowRun, parent: WorkflowRun | None) -> WorkflowTraceInfo: + assert row.finished_at is not None + metadata: dict[str, object] = {"app_id": row.app_id} + if parent is not None: + metadata["parent_trace_context"] = { + "parent_workflow_run_id": parent.id, + "parent_node_execution_id": identifier("parent-tool-node"), + } + return WorkflowTraceInfo( + node_executions=[], + conversation_id=None, + message_started_at=None, + workflow_app_log_id=None, + workflow_id=row.workflow_id, + tenant_id=row.tenant_id, + workflow_run_id=row.id, + workflow_run_elapsed_time=2, + workflow_run_status="succeeded", + workflow_run_inputs={}, + workflow_run_outputs={}, + workflow_run_version="1", + error=None, + total_tokens=0, + prompt_tokens=None, + completion_tokens=None, + file_list=[], + invoked_by=None, + query="", + metadata=metadata, + start_time=row.created_at, + end_time=row.finished_at, + ) + """Configured LogStore readers work without SQL mirrors and retain source ownership.""" @@ -42,21 +355,22 @@ class StoredSource: row: dict[str, object] requests: list[str] - def start(self) -> ApplicationTraceRun: - run = self.case.repository.start( + def start(self) -> TraceRun: + run = self.service().start( tenant_id=self.case.run.tenant_id, app_id=self.case.run.app_id, source_kind=self.kind, source_id=str(self.row["id"]), session_id="logstore-session", + workflow_id=None, ) assert run is not None return run def service(self) -> ApplicationTracingService: - return ApplicationTracingService(self.case.repository) + return self.case.service() - def status(self, node_execution_id: str | None = None) -> dict[str, object]: + def status(self, *, node_execution_id: str | None) -> TraceRunStatus: return self.service().status( tenant_id=self.case.run.tenant_id, app_id=self.case.run.app_id, @@ -66,6 +380,19 @@ def status(self, node_execution_id: str | None = None) -> dict[str, object]: ) +class AssembledEvents(TraceRecordSource): + """Real trace records that remember every event they were asked to assemble.""" + + def __init__(self, records: TraceRecordSource) -> None: + self._records = records + self.events: list[TraceEvent] = [] + + @override + def assemble(self, event: TraceEvent) -> BaseTraceInfo: + self.events.append(event) + return self._records.assemble(event) + + class TestApplicationTraceLogstoreCases: @staticmethod @pytest.fixture(params=["workflow", "node"]) @@ -159,6 +486,14 @@ def get_logs( def execute_sql( client: AliyunLogStore, *, sql: str, query: str, logstore: str, **_kwargs: object ) -> list[dict[str, object]]: + if kind == "workflow": + assert not transactions, "LogStore I/O must not hold a SQL transaction" + assert logstore == AliyunLogStore.workflow_node_execution_logstore + assert f"tenant_id = '{case.run.tenant_id}'" in sql + assert f"app_id = '{case.run.app_id}'" in sql + assert f"workflow_run_id = '{row['id']}'" in sql + requests.append(query) + return [] # The node reader now uses SQL in both SDK and PG modes. Preserve the # real repository's query/decoding path while isolating this transport. assert kind == "node" @@ -181,9 +516,7 @@ def execute_sql( event.remove(sqlite_engine, "rollback", ended) @staticmethod - def test_status_and_retry_use_logstore_without_sql_source_rows( - stored: StoredSource, monkeypatch: pytest.MonkeyPatch - ) -> None: + def test_status_and_retry_use_logstore_without_sql_source_rows(stored: StoredSource) -> None: run = stored.start() with stored.case.sessions() as session: assert session.scalar(select(WorkflowRun)) is None @@ -191,23 +524,21 @@ def test_status_and_retry_use_logstore_without_sql_source_rows( with stored.case.sessions.begin() as session: current = session.get(ApplicationTraceRun, run.id) assert current is not None - current.state = "failed" + current.state = TraceRunState.FAILED current.error_code = "provider_export_rejected" current.outer_trace = {"persisted": True} - assert stored.status()["status"] == "failed" - scheduled: list[str] = [] - service = stored.service() - monkeypatch.setattr(service, "schedule", scheduled.append) + assert stored.status(node_execution_id=None).status is TraceRunState.FAILED stored.requests.clear() - result = service.retry( + result = stored.service().retry( tenant_id=run.tenant_id, app_id=run.app_id, source_kind=stored.kind, source_id=run.source_id, acknowledge_possible_duplicates=False, + node_execution_id=None, ) - assert result["status"] == "ready" - assert scheduled == [run.id] + assert result.status is TraceRunState.READY + assert stored.case.scheduled == [run.id] assert len(stored.requests) == 1 @staticmethod @@ -215,11 +546,14 @@ def test_status_and_retry_use_logstore_without_sql_source_rows( def test_backend_result_must_match_complete_source_owner(stored: StoredSource, field: str) -> None: run = stored.start() stored.row[field] = identifier("foreign") - with pytest.raises(TraceCollectionError, match="tracing_source_not_found") as caught: + with pytest.raises(TraceCollectionNotFoundError, match="tracing_source_not_found"): stored.service().status( - tenant_id=run.tenant_id, app_id=run.app_id, source_kind=stored.kind, source_id=run.source_id + tenant_id=run.tenant_id, + app_id=run.app_id, + source_kind=stored.kind, + source_id=run.source_id, + node_execution_id=None, ) - assert caught.value.status == 404 assert stored.case.repository.get(run.id) is not None @staticmethod @@ -230,20 +564,23 @@ def test_foreign_app_owner_is_rejected_before_logstore_request(stored: StoredSou app = session.get(App, stored.case.run.app_id) assert app is not None app.tenant_id = identifier("other-tenant") - with pytest.raises(TraceCollectionError, match="tracing_source_not_found"): - stored.status() + with pytest.raises(TraceCollectionNotFoundError, match="tracing_source_not_found"): + stored.status(node_execution_id=None) assert stored.requests == [] @staticmethod + @pytest.mark.parametrize("stored", ["workflow"], indirect=True) def test_initialization_failure_uses_configured_source_reader(stored: StoredSource) -> None: - run = stored.case.repository.record_initialization_failure( + info = workflow_info(_dict_to_workflow_run(stored.row), None) + assert stored.service().capture_terminal(tenant_id=stored.case.run.tenant_id, info=info) is True + run = stored.case.repository.source_run( tenant_id=stored.case.run.tenant_id, app_id=stored.case.run.app_id, source_kind=stored.kind, source_id=str(stored.row["id"]), ) assert run is not None - assert run.state == "failed" + assert run.state is TraceRunState.FAILED assert run.error_code == "tracing_initialization_failed" @staticmethod @@ -251,12 +588,13 @@ def test_existing_snapshot_reentry_does_not_consult_logstore(stored: StoredSourc original = stored.start() stored.requests.clear() stored.row.clear() - current = stored.case.repository.start( + current = stored.service().start( tenant_id=original.tenant_id, app_id=original.app_id, source_kind=stored.kind, source_id=original.source_id, session_id="must-not-replace", + workflow_id=None, ) assert current is not None assert current.id == original.id @@ -267,37 +605,47 @@ def test_existing_snapshot_reentry_does_not_consult_logstore(stored: StoredSourc @pytest.mark.parametrize("stored", ["workflow"], indirect=True) def test_workflow_node_deep_link_is_scoped_to_durable_invocation(stored: StoredSource) -> None: run = stored.start() - case = Collection(stored.case.repository, stored.case.sessions, run) + case = replace(stored.case, run=run) node_id = identifier("logstore-agent-node") - invocation = case.repository.invocation( - run_id=run.id, agent_id=identifier("agent"), node_execution_id=node_id, attempt=0, call_key="node" - ) + invocation = open_invocation(case, node_execution_id=node_id, call_key="node") ingest(case, batch_for(run, invocation)) terminal(case) - lease = claim(case) - assert lease is not None - case.repository.complete( - run.id, - lease[1], - state="accepted", - result=TraceExportResult("provider-id", span_urls={invocation.span_id: "https://trace.test/span"}), - ) - assert stored.status(node_id)["span_url"] == "https://trace.test/span" - with pytest.raises(TraceCollectionError, match="tracing_node_not_found"): - stored.status(identifier("another-run-node")) + token = claim(case) + assert token is not None + with case.repository.change(run.id, tenant_id=run.tenant_id) as locked: + assert locked is not None + locked.run.export_accepted( + token, + provider_trace_id="provider-id", + trace_url=None, + span_urls={invocation.span_id: "https://trace.test/span"}, + now=naive_utc_now(), + ) + assert stored.status(node_execution_id=node_id).span_url == "https://trace.test/span" + with pytest.raises(TraceCollectionNotFoundError, match="tracing_node_not_found"): + stored.status(node_execution_id=identifier("another-run-node")) @staticmethod def test_preview_settlement_and_workflow_recovery_use_configured_source(stored: StoredSource) -> None: run = stored.start() if stored.kind == "workflow": - task = _workflow_task(stored.case.repository, run, naive_utc_now()) - assert task is not None - assert task.workflow_run_id == run.source_id - assert task.kwargs["trace_session_id"] == run.session_id + records = AssembledEvents(trace_records(stored.case.sessions)) + service = collector(stored.case.repository, records, stored.case.scheduled) stored.row["status"] = "running" - assert _workflow_task(stored.case.repository, run, naive_utc_now()) is None + assert recover(replace(stored.case, run=run), service, records) is False + assert records.events == [] + stored.row["status"] = "succeeded" + assert recover(replace(stored.case, run=run), service, records) is True + (recovered,) = records.events + assert recovered.payload == WorkflowRunCompleted( + workflow_run_id=run.source_id, conversation_id=None, total_tokens=None + ) + assert (recovered.tenant_id, recovered.app_id) == (run.tenant_id, run.app_id) + assert recovered.user_id == stored.row["created_by"] + assert recovered.correlation.session_id == run.session_id else: - stored.service().settle_preview(run) + # Delivery settles a finished preview from its node execution before claiming the export. + stored.service().deliver(run.id) current = stored.case.repository.get(run.id) assert current is not None assert current.outer_trace is not None @@ -305,16 +653,16 @@ def test_preview_settlement_and_workflow_recovery_use_configured_source(stored: assert root["status"] == "ok" assert root["inputs"] is None assert root["outputs"] is None - assert root["metadata"]["node_execution_id"] == run.source_id + assert root["attributes"]["node_execution_id"] == run.source_id @staticmethod @pytest.mark.parametrize("stored", ["workflow"], indirect=True) def test_missing_snapshot_classifies_logstore_agent_workflow_and_persists_failure(stored: StoredSource) -> None: - info = workflow_info(_dict_to_workflow_run(stored.row)) - assert stored.service().capture_terminal(info) is True - result = stored.status() - assert result["status"] == "failed" - assert result["error_code"] == "tracing_initialization_failed" + info = workflow_info(_dict_to_workflow_run(stored.row), None) + assert stored.service().capture_terminal(tenant_id=stored.case.run.tenant_id, info=info) is True + result = stored.status(node_execution_id=None) + assert result.status is TraceRunState.FAILED + assert result.error_code == "tracing_initialization_failed" @staticmethod @pytest.mark.parametrize("stored", ["workflow"], indirect=True) @@ -322,7 +670,7 @@ def test_missing_snapshot_classifies_logstore_agent_workflow_and_persists_failur def test_unindexed_workflow_uses_owned_definition_for_safe_legacy_decision( stored: StoredSource, has_agent: bool ) -> None: - info = workflow_info(_dict_to_workflow_run(stored.row)) + info = workflow_info(_dict_to_workflow_run(stored.row), None) with stored.case.sessions.begin() as session: app = session.get(App, stored.case.run.app_id) assert app is not None @@ -335,20 +683,39 @@ def test_unindexed_workflow_uses_owned_definition_for_safe_legacy_decision( app.workflow_id = definition.id session.add(definition) stored.row.clear() - assert stored.service().capture_terminal(info) is has_agent + assert stored.service().capture_terminal(tenant_id=stored.case.run.tenant_id, info=info) is has_agent @staticmethod @pytest.mark.parametrize("stored", ["workflow"], indirect=True) @pytest.mark.parametrize("snapshot", [True, False]) - def test_auxiliary_content_is_suppressed_for_logstore_workflow(stored: StoredSource, snapshot: bool) -> None: + def test_auxiliary_content_is_suppressed_for_logstore_chatflow(stored: StoredSource, snapshot: bool) -> None: + with stored.case.sessions.begin() as session: + app = session.get(App, stored.case.run.app_id) + message = session.get(Message, stored.case.run.source_id) + assert app is not None + assert message is not None + app.mode = AppMode.ADVANCED_CHAT + message.workflow_run_id = str(stored.row["id"]) + session.execute(delete(ApplicationTraceRun).where(ApplicationTraceRun.id == stored.case.run.id)) if snapshot: stored.start() info = ModerationTraceInfo( message_id=identifier("workflow-log"), - metadata={ - "workflow_run_id": stored.row["id"], - **({} if snapshot else {"app_id": stored.case.run.app_id}), - }, + message_data=TraceMessage( + id=message.id, + app_id=message.app_id, + conversation_id=message.conversation_id, + workflow_run_id=message.workflow_run_id, + query=message.query, + answer=message.answer, + error=message.error, + created_at=message.created_at, + model_provider=message.model_provider, + model_id=message.model_id, + total_price=message.total_price, + currency=message.currency, + ), + metadata={"message_id": message.id}, flagged=False, action="direct_output", preset_response="", @@ -356,15 +723,16 @@ def test_auxiliary_content_is_suppressed_for_logstore_workflow(stored: StoredSou inputs="private input", outputs="private output", ) - assert stored.service().capture_terminal(info) is True + tenant_id = stored.case.run.tenant_id + assert stored.service().capture_terminal(tenant_id=tenant_id, info=info) is True if snapshot: # Durable ownership is sufficient to suppress auxiliary payload even when # the remote run has not been indexed yet. stored.row.clear() - assert stored.service().capture_terminal(info) is True + assert stored.service().capture_terminal(tenant_id=tenant_id, info=info) is True else: stored.row["graph"] = json.dumps({"nodes": [{"data": {"type": "llm"}}]}) - assert stored.service().capture_terminal(info) is False + assert stored.service().capture_terminal(tenant_id=tenant_id, info=info) is False @staticmethod @pytest.mark.parametrize("stored", ["workflow"], indirect=True) @@ -374,33 +742,15 @@ def test_linked_parent_requires_logstore_run_and_durable_owner_to_agree( ) -> None: parent = stored.start() info = workflow_info(_dict_to_workflow_run(stored.row), _dict_to_workflow_run(stored.row)) - original = dict(info.metadata) if mismatch: stored.row[mismatch] = identifier("foreign") - stored.service()._link_managed_parent(info, parent.tenant_id) + linked = stored.service().link_managed_parent(tenant_id=stored.case.run.tenant_id, info=info) if mismatch: - assert info.metadata == original + assert linked is info else: - assert info.resolved_parent_context == (None, None) - assert info.metadata["linked_parent_trace_id"] == parent.trace_id - assert info.metadata["linked_parent_workflow_run_id"] == parent.source_id - - @staticmethod - @pytest.mark.parametrize("stored", ["workflow"], indirect=True) - def test_auxiliary_durable_owner_rejects_app_reassigned_to_other_tenant(stored: StoredSource) -> None: - run = stored.start() - with stored.case.sessions.begin() as session: - app = session.get(App, run.app_id) - assert app is not None - app.tenant_id = identifier("foreign") - info = ModerationTraceInfo( - metadata={"workflow_run_id": run.source_id}, - query="", - flagged=False, - action="direct_output", - preset_response="", - ) - assert stored.service().capture_terminal(info) is False + assert linked.resolved_parent_context == (None, None) + assert linked.metadata["linked_parent_trace_id"] == parent.trace_id + assert linked.metadata["linked_parent_workflow_run_id"] == parent.source_id @staticmethod def test_sql_workflow_source_still_uses_scoped_configured_reader( @@ -418,14 +768,14 @@ def test_sql_workflow_source_still_uses_scoped_configured_reader( ) with case.sessions.begin() as session: session.add(row) - source = case.repository._source( - tenant_id=case.run.tenant_id, app_id=case.run.app_id, kind="workflow", source_id=row.id + source = case.repository.workflow_run( + tenant_id=case.run.tenant_id, app_id=case.run.app_id, workflow_run_id=row.id ) - assert isinstance(source, WorkflowRun) - assert source.id == row.id + assert source is not None + assert (source.id, source.workflow_id, source.created_by) == (row.id, row.workflow_id, row.created_by) assert ( - case.repository._source( - tenant_id=identifier("wrong-tenant"), app_id=case.run.app_id, kind="workflow", source_id=row.id + case.repository.workflow_run( + tenant_id=identifier("wrong-tenant"), app_id=case.run.app_id, workflow_run_id=row.id ) is None ) @@ -434,7 +784,108 @@ def test_sql_workflow_source_still_uses_scoped_configured_reader( """Persistence boundaries that must fail closed or remain immutable across retries.""" +def race_before_insert(monkeypatch: pytest.MonkeyPatch, competitor: Callable[[], None]) -> None: + """Run ``competitor`` after the caller's existence checks, when it allocates its new run's identifiers.""" + raced: list[bool] = [] + + def allocate() -> UUID: + if not raced: + raced.append(True) + competitor() + return uuidv7() + + monkeypatch.setattr(service_module, "uuidv7", allocate) + + +def drop_run(case: Collection) -> None: + with case.sessions.begin() as session: + session.execute(delete(ApplicationTraceRun).where(ApplicationTraceRun.id == case.run.id)) + + +def run_count(case: Collection) -> int | None: + with case.sessions() as session: + return session.scalar( + select(func.count()) + .select_from(ApplicationTraceRun) + .where(ApplicationTraceRun.source_id == case.run.source_id) + ) + + +def start_message(case: Collection, session_id: str) -> TraceRun | None: + return case.service().start( + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, + source_kind="message", + source_id=case.run.source_id, + session_id=session_id, + workflow_id=None, + ) + + +class PurgedBeforeLock(ApplicationTraceRepository): + """Retention that keeps runs for zero days purges each run after a use case read it and before it locks it.""" + + @override + @contextmanager + def change(self, run_id: str, *, tenant_id: str) -> Generator[LockedTraceRun | None]: + self.purge(retention_days=0, limit=100) + with super().change(run_id, tenant_id=tenant_id) as locked: + yield locked + + class TestApplicationTraceRepositoryBoundariesCases: + @staticmethod + def test_concurrent_start_returns_the_first_frozen_snapshot( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + ) -> None: + case = seed(sqlite_session_factory, "concurrent-start") + drop_run(case) + winners: list[TraceRun | None] = [] + race_before_insert(monkeypatch, lambda: winners.append(start_message(case, "first-session"))) + current = start_message(case, "second-session") + assert current is not None + assert winners[0] is not None + assert current.id == winners[0].id + assert current.session_id == "first-session" + assert current.state is TraceRunState.COLLECTING + assert run_count(case) == 1 + + @staticmethod + def test_initialization_failure_yields_to_a_concurrently_frozen_start( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + ) -> None: + case = seed(sqlite_session_factory, "concurrent-failure") + drop_run(case) + info = message_info(case) + winners: list[TraceRun | None] = [] + race_before_insert(monkeypatch, lambda: winners.append(start_message(case, "started-session"))) + # The Agent app's terminal finds no start snapshot, so it records an initialization failure. + assert case.service().capture_terminal(tenant_id=case.run.tenant_id, info=info) is True + current = case.repository.source_run( + tenant_id=case.run.tenant_id, app_id=case.run.app_id, source_kind="message", source_id=case.run.source_id + ) + assert current is not None + assert winners[0] is not None + assert current.id == winners[0].id + assert current.state is TraceRunState.COLLECTING + assert current.error_code is None + assert current.provider_config == case.run.provider_config + assert run_count(case) == 1 + + @staticmethod + def test_identifier_collision_is_not_mistaken_for_a_concurrent_start( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + ) -> None: + case = seed(sqlite_session_factory, "colliding-start") + other = seed(sqlite_session_factory, "colliding-owner") + drop_run(case) + monkeypatch.setattr(service_module, "uuidv7", lambda: UUID(other.run.id)) + with pytest.raises(TracingStoreUnavailableError) as caught: + start_message(case, "colliding-session") + assert isinstance(caught.value.__cause__, IntegrityError) + assert run_count(case) == 0 + assert case.repository.get(other.run.id) is not None + @staticmethod @pytest.mark.parametrize("policy", ["invalid-json", "null", "[]"]) def test_malformed_policy_freezes_disabled_decision( @@ -446,15 +897,20 @@ def test_malformed_policy_freezes_disabled_decision( app = session.get(App, case.run.app_id) assert app is not None app.tracing = policy - run = case.repository.start( + started = case.service().start( tenant_id=case.run.tenant_id, app_id=case.run.app_id, source_kind="message", source_id=case.run.source_id, session_id="", + workflow_id=None, + ) + assert started is None + run = case.repository.source_run( + tenant_id=case.run.tenant_id, app_id=case.run.app_id, source_kind="message", source_id=case.run.source_id ) assert run is not None - assert run.state == "not_collected" + assert run.state is TraceRunState.NOT_COLLECTED assert run.provider_config == {} assert run.session_id == case.run.source_id assert not run.capture_content @@ -463,42 +919,48 @@ def test_malformed_policy_freezes_disabled_decision( @staticmethod def test_reserved_preview_requires_persisted_owner_workflow(sqlite_session_factory: sessionmaker[Session]) -> None: case = seed(sqlite_session_factory, "missing-preview-owner") + preview = identifier("preview") assert ( - case.repository.start( + case.service().start( tenant_id=case.run.tenant_id, app_id=case.run.app_id, source_kind="node", - source_id=identifier("preview"), + source_id=preview, session_id="preview", workflow_id=identifier("missing-workflow"), ) is None ) + assert ( + case.repository.source_run( + tenant_id=case.run.tenant_id, app_id=case.run.app_id, source_kind="node", source_id=preview + ) + is None + ) @staticmethod def test_initialization_failure_cannot_replace_snapshot_or_create_foreign_row( sqlite_session_factory: sessionmaker[Session], ) -> None: case = seed(sqlite_session_factory, "failure-snapshot") - same = case.repository.record_initialization_failure( - tenant_id=case.run.tenant_id, - app_id=case.run.app_id, - source_kind="message", - source_id=case.run.source_id, + info = message_info(case) + assert case.service().capture_terminal(tenant_id=case.run.tenant_id, info=info) is True + same = case.repository.source_run( + tenant_id=case.run.tenant_id, app_id=case.run.app_id, source_kind="message", source_id=case.run.source_id ) assert same is not None assert same.id == case.run.id - assert same.state == "collecting" + assert same.state is TraceRunState.COLLECTING assert same.provider_config == case.run.provider_config + foreign = identifier("foreign-tenant") + assert case.service().capture_terminal(tenant_id=foreign, info=info) is False assert ( - case.repository.record_initialization_failure( - tenant_id=identifier("foreign-tenant"), - app_id=case.run.app_id, - source_kind="message", - source_id=case.run.source_id, + case.repository.source_run( + tenant_id=foreign, app_id=case.run.app_id, source_kind="message", source_id=case.run.source_id ) is None ) + assert run_count(case) == 1 @staticmethod @pytest.mark.parametrize("target", ["missing_run", "revoked_owner", "closed_invocation"]) @@ -508,22 +970,23 @@ def test_ingestion_rejects_missing_or_closed_owner_without_writing( case = seed(sqlite_session_factory, "ingestion-owner") invocation = case.invocation() batch = batch_for(case.run, invocation) + error: type[TraceCollectionError] if target == "missing_run": batch = batch.model_copy(update={"run_id": UUID(identifier("missing-run"))}) - error = "tracing_run_not_found" + error, code = TraceCollectionNotFoundError, "tracing_run_not_found" elif target == "revoked_owner": with case.sessions.begin() as session: app = session.get(App, case.run.app_id) assert app is not None app.tenant_id = identifier("another-tenant") - error = "tracing_owner_not_found" + error, code = TraceCollectionNotFoundError, "tracing_owner_not_found" else: with case.sessions.begin() as session: stored = session.get(ApplicationTraceInvocation, invocation.id) assert stored is not None stored.finalized = True - error = "tracing_collection_closed" - with pytest.raises(TraceCollectionError, match=error): + error, code = TraceCollectionConflictError, "tracing_collection_closed" + with pytest.raises(error, match=code): ingest(case, batch) with case.sessions() as session: assert session.scalar(select(func.count()).select_from(ApplicationTraceBatch)) == 0 @@ -539,7 +1002,8 @@ def test_first_terminal_is_immutable_and_revoked_owner_is_never_exported( terminal(case) before = case.repository.get(case.run.id) assert before is not None - case.repository.terminal(case.run.id, {"replacement": True}, "f" * 16) + replacement = agent_turn(case.run, root_span_id="f" * 16) + capture_terminal(case, replacement) after = case.repository.get(case.run.id) assert after is not None assert after.outer_trace == before.outer_trace @@ -552,23 +1016,26 @@ def test_first_terminal_is_immutable_and_revoked_owner_is_never_exported( failed = case.repository.get(case.run.id) assert failed is not None assert failed.error_code == "tracing_owner_not_found" - assert failed.state == "failed" - case.repository.terminal(case.run.id, {"replacement": True}, "f" * 16) + assert failed.state is TraceRunState.FAILED + capture_terminal(case, replacement) frozen = case.repository.get(case.run.id) assert frozen is not None assert frozen.outer_trace == before.outer_trace @staticmethod @pytest.mark.parametrize( - ("state", "error"), + ("state", "error", "code"), [ - ("not_collected", "tracing_run_not_found"), - ("accepted", "tracing_already_accepted"), - ("collecting", "tracing_retry_not_allowed"), + (TraceRunState.NOT_COLLECTED, TraceCollectionNotFoundError, "tracing_run_not_found"), + (TraceRunState.ACCEPTED, TraceCollectionConflictError, "tracing_already_accepted"), + (TraceRunState.COLLECTING, TraceCollectionConflictError, "tracing_retry_not_allowed"), ], ) def test_retry_rejects_ineligible_states( - sqlite_session_factory: sessionmaker[Session], state: str, error: str + sqlite_session_factory: sessionmaker[Session], + state: TraceRunState, + error: type[TraceCollectionError], + code: str, ) -> None: case = seed(sqlite_session_factory, "retry-state") terminal(case) @@ -576,28 +1043,52 @@ def test_retry_rejects_ineligible_states( row = session.get(ApplicationTraceRun, case.run.id) assert row is not None row.state = state - with pytest.raises(TraceCollectionError, match=error): - case.repository.retry(run_id=case.run.id, acknowledge_possible_duplicates=True) - row = case.repository.get(case.run.id) - assert row is not None - assert row.state == state + with pytest.raises(error, match=code): + case.service().retry( + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, + source_kind="message", + source_id=case.run.source_id, + acknowledge_possible_duplicates=True, + node_execution_id=None, + ) + current = case.repository.get(case.run.id) + assert current is not None + assert current.state is state @staticmethod def test_retry_unknown_run_and_invocation_reject_before_creation( sqlite_session_factory: sessionmaker[Session], ) -> None: case = seed(sqlite_session_factory, "retry-unknown") - with pytest.raises(TraceCollectionError, match="tracing_run_not_found"): - case.repository.retry(run_id=identifier("missing"), acknowledge_possible_duplicates=True) - with pytest.raises(TraceCollectionError, match="tracing_run_not_collecting"): - case.repository.invocation( - run_id=identifier("missing"), + purged = replace(case, repository=PurgedBeforeLock(case.sessions)).service() + with pytest.raises(TraceCollectionNotFoundError, match="tracing_run_not_found"): + purged.retry( + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, + source_kind="message", + source_id=case.run.source_id, + acknowledge_possible_duplicates=True, + node_execution_id=None, + ) + with pytest.raises(TraceCollectionConflictError, match="tracing_run_not_collecting"): + purged.prepare_invocation( + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, agent_id=identifier("agent"), + source_kind="message", + source_id=case.run.source_id, + session_id="unknown-session", node_execution_id=None, attempt=0, call_key="x", + workflow_id=None, ) - assert case.repository.claim(identifier("missing"), collection_timeout=1, lease_seconds=1) is None + purged.deliver(case.run.id) + assert run_count(case) == 0 + with case.sessions() as session: + assert session.scalar(select(func.count()).select_from(ApplicationTraceInvocation)) == 0 + assert case.scheduled == [] @staticmethod def test_collection_without_terminal_is_not_exportable(sqlite_session_factory: sessionmaker[Session]) -> None: @@ -607,7 +1098,7 @@ def test_collection_without_terminal_is_not_exportable(sqlite_session_factory: s row = case.repository.get(case.run.id) assert row is not None assert row.attempts == 0 - assert row.state == "collecting" + assert row.state is TraceRunState.COLLECTING @staticmethod def test_retention_removes_all_payload_rows_and_preserves_newer_runs( @@ -621,10 +1112,521 @@ def test_retention_removes_all_payload_rows_and_preserves_newer_runs( row = session.get(ApplicationTraceRun, expired.run.id) assert row is not None row.created_at = naive_utc_now() - timedelta(days=8) - assert expired.repository.purge(retention_days=7) == 1 + assert expired.repository.purge(retention_days=7, limit=100) == 1 assert expired.repository.get(expired.run.id) is None - assert expired.repository.collected(expired.run.id) == ([], []) + assert expired.repository.invocations(expired.run.id) == [] + assert expired.repository.batches(expired.run.id) == [] assert fresh.repository.get(fresh.run.id) is not None - invocations, batches = fresh.repository.collected(fresh.run.id) - assert len(invocations) == len(batches) == 1 - assert expired.repository.purge(retention_days=7) == 0 + assert len(fresh.repository.invocations(fresh.run.id)) == len(fresh.repository.batches(fresh.run.id)) == 1 + assert expired.repository.purge(retention_days=7, limit=100) == 0 + + +"""Whole runs load and save under their row lock; every read is scoped by the run's owner.""" + + +def stored_fields(run: TraceRunAgg) -> dict[str, object]: + """Every field of a run that its repository saves and loads again.""" + return {name: value for name, value in asdict(run).items() if name != "new_batches"} + + +def reloaded(collection: Collection) -> dict[str, object]: + with collection.repository.change(collection.run.id, tenant_id=collection.run.tenant_id) as locked: + assert locked is not None + return stored_fields(locked.run) + + +def record_of(run: TraceRunAgg) -> TraceRun: + """The record the repository reads back for a stored run.""" + return TraceRun( + id=run.id, + tenant_id=run.tenant_id, + app_id=run.app_id, + source_kind=run.source_kind, + source_id=run.source_id, + provider=run.provider, + config_revision=run.config_revision, + provider_config=run.provider_config, + trace_id=run.trace_id, + root_span_id=run.root_span_id, + session_id=run.session_id, + capture_content=run.capture_content, + capture_agent=run.capture_agent, + state=run.state, + outer_trace=run.outer_trace, + span_count=run.span_count, + attempts=run.attempts, + error_code=run.error_code, + trace_url=run.trace_url, + provider_trace_id=run.provider_trace_id, + span_urls=run.span_urls, + accepted_at=run.accepted_at, + ) + + +def frozen_run(case: Collection) -> TraceRunAgg: + """A new run of the collection's message, frozen for its Langfuse destination.""" + return TraceRunAgg.freeze( + run_id=str(uuidv7()), + trace_id=uuidv7().hex, + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, + source_kind="message", + source_id=case.run.source_id, + session_id="frozen-session", + policy=LANGFUSE_POLICY, + stored_config=case.run.provider_config, + now=naive_utc_now(), + ) + + +def failed_run(case: Collection) -> TraceRunAgg: + """A new run of the collection's message, closed because its start snapshot is missing.""" + return TraceRunAgg.initialization_failed( + run_id=str(uuidv7()), + trace_id=uuidv7().hex, + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, + source_kind="message", + source_id=case.run.source_id, + now=naive_utc_now(), + ) + + +def store_policy(case: Collection, policy: str | None) -> None: + with case.sessions.begin() as session: + app = session.get(App, case.run.app_id) + assert app is not None + app.tracing = policy + + +@dataclass +class OwnedSources: + """An Agent app's message in its conversation, a workflow definition, and a run of it with a node execution.""" + + case: Collection + workflow_id: str + workflow_run_id: str + node_execution_id: str + + def read(self, name: str, *, tenant_id: str, app_id: str) -> object: + repository = self.case.repository + message_id = self.case.run.source_id + match name: + case "source_run": + return repository.source_run( + tenant_id=tenant_id, app_id=app_id, source_kind="message", source_id=message_id + ) + case "workflow_run_trace": + return repository.workflow_run_trace(tenant_id=tenant_id, workflow_run_id=self.workflow_run_id) + case "app": + return repository.app(tenant_id=tenant_id, app_id=app_id) + case "app_tracing": + return repository.app_tracing(tenant_id=tenant_id, app_id=app_id) + case "app_workflow_graph": + return repository.app_workflow_graph(tenant_id=tenant_id, app_id=app_id, workflow_id=self.workflow_id) + case "message": + return repository.message(tenant_id=tenant_id, app_id=app_id, message_id=message_id) + case "workflow_run": + return repository.workflow_run(tenant_id=tenant_id, app_id=app_id, workflow_run_id=self.workflow_run_id) + case "node_execution": + return repository.node_execution( + tenant_id=tenant_id, app_id=app_id, node_execution_id=self.node_execution_id + ) + case "workflow_run_conversation_id": + return repository.workflow_run_conversation_id( + tenant_id=tenant_id, app_id=app_id, workflow_run_id=self.workflow_run_id + ) + case "message_progress": + return repository.message_progress(tenant_id=tenant_id, app_id=app_id, message_id=message_id) + case _: + raise AssertionError(f"unknown read: {name}") + + +WORKSPACE_READS = ( + "source_run", + "workflow_run_trace", + "app", + "app_tracing", + "app_workflow_graph", + "message", + "workflow_run", + "node_execution", + "workflow_run_conversation_id", + "message_progress", +) +"""Every read of runs and of the records they are collected for that names a workspace.""" +APP_READS = tuple(name for name in WORKSPACE_READS if name not in {"workflow_run_trace", "app", "app_tracing"}) +"""The reads that also name the app owning the record.""" + + +class TestApplicationTraceRepositoryPersistenceCases: + @staticmethod + @pytest.fixture + def owned(sqlite_session_factory: sessionmaker[Session]) -> OwnedSources: + case = seed(sqlite_session_factory, "owned") + workflow_id, workflow_run_id = identifier("owned-definition"), identifier("owned-run") + node_execution_id, owner = identifier("owned-node"), identifier("owner") + with case.sessions.begin() as session: + message = session.get(Message, case.run.source_id) + assert message is not None + message.workflow_run_id = workflow_run_id + session.add( + make_conversation( + conversation_id=message.conversation_id, + app_id=case.run.app_id, + mode=AppMode.AGENT, + inputs={}, + from_source=ConversationFromSource.CONSOLE, + ) + ) + session.add( + make_workflow( + workflow_id=workflow_id, tenant_id=case.run.tenant_id, app_id=case.run.app_id, graph=AGENT_GRAPH + ) + ) + session.add( + _dict_to_workflow_run( + { + "id": workflow_run_id, + "tenant_id": case.run.tenant_id, + "app_id": case.run.app_id, + "workflow_id": workflow_id, + "created_by": owner, + } + ) + ) + session.add( + _dict_to_workflow_node_execution_model( + { + "id": node_execution_id, + "tenant_id": case.run.tenant_id, + "app_id": case.run.app_id, + "workflow_id": workflow_id, + "workflow_run_id": workflow_run_id, + "node_id": "agent", + "node_type": "agent", + "title": "Agent", + "status": "succeeded", + "created_by": owner, + } + ) + ) + workflow_trace = case.service().start( + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, + source_kind="workflow", + source_id=workflow_run_id, + session_id="owned-session", + workflow_id=None, + ) + assert workflow_trace is not None + return OwnedSources(case, workflow_id, workflow_run_id, node_execution_id) + + @staticmethod + @pytest.mark.parametrize("new_run", [frozen_run, failed_run], ids=["frozen", "initialization_failed"]) + def test_add_stores_every_field_of_a_new_run( + sqlite_session_factory: sessionmaker[Session], new_run: Callable[[Collection], TraceRunAgg] + ) -> None: + case = seed(sqlite_session_factory, "added") + drop_run(case) + run = new_run(case) + assert case.repository.add(run) is True + assert case.repository.get(run.id) == record_of(run) + assert reloaded(replace(case, run=record_of(run))) == stored_fields(run) + + @staticmethod + def test_add_leaves_a_source_to_the_run_that_first_froze_it(sqlite_session_factory: sessionmaker[Session]) -> None: + case = seed(sqlite_session_factory, "first-frozen") + later, failed = frozen_run(case), failed_run(case) + assert case.repository.add(later) is False + assert case.repository.add(failed) is False + assert case.repository.get(later.id) is None + assert case.repository.get(failed.id) is None + assert ( + case.repository.source_run( + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, + source_kind="message", + source_id=case.run.source_id, + ) + == case.run + ) + assert run_count(case) == 1 + + @staticmethod + def test_change_saves_root_fields_invocations_and_new_batches( + sqlite_session_factory: sessionmaker[Session], + ) -> None: + case = seed(sqlite_session_factory, "saved") + now = naive_utc_now() + batches: list[TraceBatch] = [] + with case.repository.change(case.run.id, tenant_id=case.run.tenant_id) as locked: + assert locked is not None + invocation = locked.run.open_invocation( + agent_id=identifier("agent"), node_execution_id=None, attempt=0, call_key="first" + ) + partial = batch_for(case.run, invocation).model_copy(update={"final": False}) + locked.run.ingest(partial, max_spans=100, max_bytes=1_000_000, now=now) + locked.run.capture_terminal(agent_turn(case.run, root_span_id="b" * 16), now=now) + batches.extend(locked.run.new_batches) + saved = stored_fields(locked.run) + assert reloaded(case) == saved + + with case.repository.change(case.run.id, tenant_id=case.run.tenant_id) as locked: + assert locked is not None + # The final batch finalizes the invocation the previous change stored. + final = batch_for(case.run, invocation) + locked.run.ingest(final, max_spans=100, max_bytes=1_000_000, now=now) + token = str(uuid4()) + assert locked.run.claim_export( + token=token, owner_exists=locked.owner_exists, collection_timeout=60, lease_seconds=30, now=now + ) + batches.extend(locked.run.new_batches) + saved = stored_fields(locked.run) + assert reloaded(case) == saved + + with case.repository.change(case.run.id, tenant_id=case.run.tenant_id) as locked: + assert locked is not None + locked.run.export_deferred(token, max_attempts=8, now=now) + saved = stored_fields(locked.run) + assert reloaded(case) == saved + + later = now + timedelta(minutes=1) + with case.repository.change(case.run.id, tenant_id=case.run.tenant_id) as locked: + assert locked is not None + token = str(uuid4()) + assert locked.run.claim_export( + token=token, owner_exists=locked.owner_exists, collection_timeout=60, lease_seconds=30, now=later + ) + locked.run.export_accepted( + token, + provider_trace_id="provider-trace", + trace_url="https://trace.test", + span_urls={invocation.span_id: "https://trace.test/agent"}, + now=later, + ) + saved, record = stored_fields(locked.run), record_of(locked.run) + assert reloaded(case) == saved + assert case.repository.get(case.run.id) == record + assert case.repository.invocations(case.run.id) == [replace(invocation, finalized=True, finalized_at=now)] + assert case.repository.batches(case.run.id) == batches + + @staticmethod + def test_change_rolls_back_the_whole_block_when_it_raises(sqlite_session_factory: sessionmaker[Session]) -> None: + case = seed(sqlite_session_factory, "rolled-back") + before = reloaded(case) + + def change_then_break_a_rule() -> None: + with case.repository.change(case.run.id, tenant_id=case.run.tenant_id) as locked: + assert locked is not None + invocation = locked.run.open_invocation( + agent_id=identifier("agent"), node_execution_id=None, attempt=0, call_key="first" + ) + batch = batch_for(case.run, invocation) + locked.run.ingest(batch, max_spans=100, max_bytes=1_000_000, now=naive_utc_now()) + locked.run.capture_terminal( + agent_turn(case.run, root_span_id=case.run.root_span_id), now=naive_utc_now() + ) + locked.run.open_invocation( + agent_id=identifier("another-agent"), node_execution_id=None, attempt=0, call_key="first" + ) + + with pytest.raises(TraceRunConflictError, match="tracing_invocation_conflict"): + change_then_break_a_rule() + assert reloaded(case) == before + assert case.repository.invocations(case.run.id) == [] + assert case.repository.batches(case.run.id) == [] + + @staticmethod + def test_change_rejects_a_new_batch_whose_id_another_run_stores( + sqlite_session_factory: sessionmaker[Session], + ) -> None: + case = seed(sqlite_session_factory, "batch-owner") + other = seed(sqlite_session_factory, "batch-reuse") + batch = batch_for(case.run, case.invocation()) + ingest(case, batch) + reused = batch_for(other.run, other.invocation()).model_copy(update={"batch_id": batch.batch_id}) + before = reloaded(other) + with pytest.raises(TraceCollectionConflictError, match="tracing_batch_conflict"): + store_batch(other, reused) + assert reloaded(other) == before + assert other.repository.batches(other.run.id) == [] + (kept,) = case.repository.batches(case.run.id) + assert (kept.id, kept.invocation_id) == (str(batch.batch_id), str(batch.invocation_id)) + + @staticmethod + @pytest.mark.parametrize("revocation", ["moved", "deleted"]) + def test_change_tells_whether_the_owning_app_still_exists( + sqlite_session_factory: sessionmaker[Session], revocation: str + ) -> None: + case = seed(sqlite_session_factory, "owner") + + def owner_exists() -> bool: + with case.repository.change(case.run.id, tenant_id=case.run.tenant_id) as locked: + assert locked is not None + return locked.owner_exists + + assert owner_exists() is True + with case.sessions.begin() as session: + app = session.get(App, case.run.app_id) + assert app is not None + if revocation == "moved": + app.tenant_id = identifier("another-tenant") + else: + session.delete(app) + assert owner_exists() is False + + @staticmethod + def test_change_yields_nothing_for_another_tenant(sqlite_session_factory: sessionmaker[Session]) -> None: + case = seed(sqlite_session_factory, "scoped-change") + before = reloaded(case) + + with case.repository.change(case.run.id, tenant_id=identifier("other-tenant")) as locked: + assert locked is None + + assert reloaded(case) == before + + @staticmethod + def test_change_yields_nothing_for_a_missing_run(sqlite_session_factory: sessionmaker[Session]) -> None: + case = seed(sqlite_session_factory, "missing") + with case.repository.change(identifier("missing-run"), tenant_id=case.run.tenant_id) as locked: + assert locked is None + + @staticmethod + @pytest.mark.parametrize("read", WORKSPACE_READS) + def test_reads_are_scoped_to_the_owning_workspace(owned: OwnedSources, read: str) -> None: + owner = owned.case.run + assert owned.read(read, tenant_id=owner.tenant_id, app_id=owner.app_id) is not None + assert owned.read(read, tenant_id=identifier("foreign-tenant"), app_id=owner.app_id) is None + + @staticmethod + @pytest.mark.parametrize("read", APP_READS) + def test_reads_are_scoped_to_the_owning_app(owned: OwnedSources, read: str) -> None: + # Every seeded app belongs to the same workspace. + sibling = seed(owned.case.sessions, "sibling") + assert owned.read(read, tenant_id=sibling.run.tenant_id, app_id=sibling.run.app_id) is None + + @staticmethod + def test_workflow_run_trace_requires_its_app(owned: OwnedSources) -> None: + owner = owned.case.run + with owned.case.sessions.begin() as session: + session.execute(delete(App).where(App.id == owner.app_id)) + assert owned.read("workflow_run_trace", tenant_id=owner.tenant_id, app_id=owner.app_id) is None + + @staticmethod + @pytest.mark.parametrize("policy", [None, "", "invalid-json", "null", "[]", '"langfuse"']) + def test_app_tracing_reads_a_missing_or_malformed_policy_as_empty( + sqlite_session_factory: sessionmaker[Session], policy: str | None + ) -> None: + case = seed(sqlite_session_factory, "unreadable-policy") + store_policy(case, policy) + tracing = case.repository.app_tracing(tenant_id=case.run.tenant_id, app_id=case.run.app_id) + assert tracing == AppTracingSnapshot(policy={}, config=None) + + @staticmethod + @pytest.mark.parametrize( + ("policy", "named"), + [ + ({"enabled": True, "tracing_provider": "langfuse"}, "langfuse"), + ({"enabled": False, "tracing_provider": "phoenix"}, "phoenix"), + ({"enabled": True, "tracing_provider": "opik"}, None), + ({"enabled": True}, None), + ], + ids=["enabled", "disabled", "unconfigured", "unnamed"], + ) + def test_app_tracing_reads_the_stored_configuration_of_the_named_provider( + sqlite_session_factory: sessionmaker[Session], policy: dict[str, object], named: str | None + ) -> None: + case = seed(sqlite_session_factory, "named-provider") + configs = {"langfuse": dict(case.run.provider_config), "phoenix": {"endpoint": "https://phoenix.test"}} + with case.sessions.begin() as session: + session.add( + TraceAppConfig(app_id=case.run.app_id, tracing_provider="phoenix", tracing_config=configs["phoenix"]) + ) + store_policy(case, json.dumps(policy)) + tracing = case.repository.app_tracing(tenant_id=case.run.tenant_id, app_id=case.run.app_id) + assert tracing == AppTracingSnapshot(policy=policy, config=configs[named] if named else None) + + @staticmethod + def test_runs_read_invocations_by_attempt_then_creation_and_batches_by_arrival( + sqlite_session_factory: sessionmaker[Session], + ) -> None: + case = seed(sqlite_session_factory, "collected") + now = naive_utc_now() + agent_id = identifier("agent") + with case.repository.change(case.run.id, tenant_id=case.run.tenant_id) as locked: + assert locked is not None + retried = locked.run.open_invocation(agent_id=agent_id, node_execution_id=None, attempt=1, call_key="r") + first = locked.run.open_invocation(agent_id=agent_id, node_execution_id=None, attempt=0, call_key="a") + second = locked.run.open_invocation(agent_id=agent_id, node_execution_id=None, attempt=0, call_key="b") + for invocation in (first, second, retried): + batch = batch_for(case.run, invocation) + locked.run.ingest(batch, max_spans=100, max_bytes=1_000_000, now=now) + invocations, batches = dict(locked.run.invocations), list(locked.run.new_batches) + # Creation and arrival times that contradict attempt, identifier and storage order. + earlier, later = sorted([first, second], key=lambda invocation: invocation.id, reverse=True) + arrival = [batches[2], batches[0], batches[1]] + with case.sessions.begin() as session: + for offset, invocation in enumerate([retried, earlier, later]): + stored_invocation = session.get(ApplicationTraceInvocation, invocation.id) + assert stored_invocation is not None + stored_invocation.created_at = now + timedelta(seconds=offset) + for offset, batch in enumerate(arrival): + stored_batch = session.get(ApplicationTraceBatch, batch.id) + assert stored_batch is not None + stored_batch.created_at = now + timedelta(seconds=offset) + ordered = [invocations[earlier.id], invocations[later.id], invocations[retried.id]] + assert case.repository.invocations(case.run.id) == ordered + assert case.repository.batches(case.run.id) == arrival + loaded = case.repository.load(case.run.id) + assert loaded is not None + assert list(loaded.invocations.values()) == ordered + + @staticmethod + def test_pending_ids_lists_due_runs_that_may_export_with_terminals_first( + sqlite_session_factory: sessionmaker[Session], + ) -> None: + case = seed(sqlite_session_factory, "pending") + drop_run(case) + now, hour = naive_utc_now(), timedelta(hours=1) + # name: (state, terminal_at, next_attempt_at, updated_at), stored in another order than the expected one + rows: dict[str, tuple[TraceRunState, datetime | None, datetime | None, datetime]] = { + "newer-ready": (TraceRunState.READY, None, None, now - hour), + "older-collecting": (TraceRunState.COLLECTING, None, None, now - 6 * hour), + "terminal-collecting": (TraceRunState.COLLECTING, now - 5 * hour, None, now - 2 * hour), + "due-retry": (TraceRunState.RETRYING, now - 5 * hour, now - timedelta(minutes=1), now - 3 * hour), + "exporting": (TraceRunState.EXPORTING, now - 5 * hour, None, now - 4 * hour), + "backing-off": (TraceRunState.RETRYING, now - 5 * hour, now + hour, now - 7 * hour), + "accepted": (TraceRunState.ACCEPTED, now - 5 * hour, None, now - 7 * hour), + "failed": (TraceRunState.FAILED, now - 5 * hour, None, now - 7 * hour), + "uncertain": (TraceRunState.UNCERTAIN, now - 5 * hour, None, now - 7 * hour), + "not-collected": (TraceRunState.NOT_COLLECTED, None, None, now - 7 * hour), + } + run_ids: dict[str, str] = {} + for name in rows: + run = TraceRunAgg.freeze( + run_id=str(uuidv7()), + trace_id=uuidv7().hex, + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, + source_kind="message", + source_id=identifier(f"pending-{name}"), + session_id=name, + policy=LANGFUSE_POLICY, + stored_config=case.run.provider_config, + now=now, + ) + assert case.repository.add(run) is True + run_ids[name] = run.id + with case.sessions.begin() as session: + for name, (state, terminal_at, next_attempt_at, updated_at) in rows.items(): + row = session.get(ApplicationTraceRun, run_ids[name]) + assert row is not None + row.state, row.terminal_at = state, terminal_at + row.next_attempt_at, row.updated_at = next_attempt_at, updated_at + pending = [ + run_ids[name] + for name in ("exporting", "due-retry", "terminal-collecting", "older-collecting", "newer-ready") + ] + assert case.repository.pending_ids(100) == pending + assert case.repository.pending_ids(2) == pending[:2] diff --git a/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_node_execution_repository.py b/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_node_execution_repository.py index dabdf5e83df275..b44a861c9c5f47 100644 --- a/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_node_execution_repository.py +++ b/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_node_execution_repository.py @@ -1,17 +1,21 @@ import json +from collections.abc import Mapping from datetime import datetime, timedelta from unittest.mock import patch from uuid import uuid4 import pytest from sqlalchemy import event, func, select -from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.orm import ORMExecuteState, Session, sessionmaker from core.workflow.node_execution_process_data import ( WORKFLOW_TOOL_ROOT_APP_ID_KEY, ) +from extensions.ext_storage import storage +from extensions.storage.storage_type import StorageType from graphon.enums import WorkflowNodeExecutionStatus from models.enums import CreatorUserRole +from models.model import UploadFile from models.workflow import ( ExecutionOffLoadType, WorkflowNodeExecutionModel, @@ -26,6 +30,172 @@ from tests.helpers.unit.config_override import apply_config_overrides +@pytest.mark.parametrize("foreign_tenant", [False, True]) +@pytest.mark.parametrize("payload_type", list(ExecutionOffLoadType)) +@pytest.mark.parametrize("bulk", [False, True]) +def test_full_payloads_load_owned_files_after_closing_the_database_session( + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + payload_type: ExecutionOffLoadType, + foreign_tenant: bool, + bulk: bool, +) -> None: + file = UploadFile( + tenant_id="other-tenant" if foreign_tenant else "tenant", + storage_type=StorageType.LOCAL, + key="full-payload.json", + name="payload.json", + size=21, + extension="json", + mime_type="application/json", + created_by_role=CreatorUserRole.ACCOUNT, + created_by="account", + created_at=datetime(2026, 1, 1), + used=True, + ) + file_id = file.id + with sqlite_session_factory.begin() as session: + session.add(file) + execution = WorkflowNodeExecutionModel( + id="execution", + tenant_id="tenant", + inputs='{"truncated":true}', + outputs='{"truncated":true}', + process_data='{"truncated":true}', + ) + execution.offload_data = [] + repository = DifyAPISQLAlchemyWorkflowNodeExecutionRepository(sqlite_session_factory) + load = { + ExecutionOffLoadType.INPUTS: repository.load_full_inputs, + ExecutionOffLoadType.OUTPUTS: repository.load_full_outputs, + ExecutionOffLoadType.PROCESS_DATA: repository.load_full_process_data, + }[payload_type] + + def read_payload() -> Mapping[str, object] | None: + if bulk: + full = repository.load_full_payloads([execution])[execution.id] + return { + ExecutionOffLoadType.INPUTS: full.inputs, + ExecutionOffLoadType.OUTPUTS: full.outputs, + ExecutionOffLoadType.PROCESS_DATA: full.process_data, + }[payload_type] + return load(execution) + + assert read_payload() == {"truncated": True} + execution.offload_data.append( + WorkflowNodeExecutionOffload( + tenant_id="tenant", app_id="app", node_execution_id="execution", type_=payload_type, file_id=file_id + ) + ) + readers: list[Session] = [] + + def track_reader(session: Session, *_args: object) -> None: + readers.append(session) + + def load_storage(key: str) -> bytes: + assert key == "full-payload.json" + assert readers + assert all(not session.in_transaction() for session in readers) + return b'{"complete":"payload"}' + + monkeypatch.setattr(storage, "load", load_storage) + event.listen(sqlite_session_factory, "after_begin", track_reader) + try: + if foreign_tenant: + with pytest.raises(AssertionError, match="execution's tenant"): + read_payload() + else: + assert read_payload() == {"complete": "payload"} + finally: + event.remove(sqlite_session_factory, "after_begin", track_reader) + + +@pytest.mark.parametrize("execution_count", [2, 1001]) +def test_full_payloads_batch_metadata_and_load_shared_files_once( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, execution_count: int +) -> None: + files = [ + UploadFile( + tenant_id="tenant", + storage_type=StorageType.LOCAL, + key=f"payload-{index}.json", + name="payload.json", + size=21, + extension="json", + mime_type="application/json", + created_by_role=CreatorUserRole.ACCOUNT, + created_by="account", + created_at=datetime(2026, 1, 1), + used=True, + ) + for index in range(execution_count + 1) + ] + with sqlite_session_factory.begin() as session: + session.add_all(files) + executions = [ + WorkflowNodeExecutionModel(id=str(index), tenant_id="tenant", outputs=json.dumps({"index": index})) + for index in range(execution_count) + ] + for index, execution in enumerate(executions): + execution.offload_data = [ + WorkflowNodeExecutionOffload( + tenant_id="tenant", + app_id="app", + node_execution_id=execution.id, + type_=type_, + file_id=file.id, + ) + for type_, file in ( + (ExecutionOffLoadType.INPUTS, files[index]), + (ExecutionOffLoadType.PROCESS_DATA, files[-1]), + ) + ] + readers: list[Session] = [] + reads: list[str] = [] + query_count = 0 + + def track_reader(session: Session, *_args: object) -> None: + readers.append(session) + + def count_query(state: ORMExecuteState) -> None: + nonlocal query_count + assert state.is_select + query_count += 1 + + def load_storage(key: str) -> bytes: + assert readers + assert all(not session.in_transaction() for session in readers) + reads.append(key) + return json.dumps({"file": key}).encode() + + monkeypatch.setattr(storage, "load", load_storage) + event.listen(sqlite_session_factory, "after_begin", track_reader) + event.listen(sqlite_session_factory, "do_orm_execute", count_query) + try: + repository = DifyAPISQLAlchemyWorkflowNodeExecutionRepository(sqlite_session_factory) + assert repository.load_full_payloads([]) == {} + inline = WorkflowNodeExecutionModel(id="inline", tenant_id="tenant", inputs='{"inline":true}') + inline.offload_data = [] + full_inline = repository.load_full_payloads([inline])[inline.id] + assert full_inline.inputs == {"inline": True} + assert full_inline.outputs is None + assert full_inline.process_data is None + assert query_count == 0 + + full = repository.load_full_payloads(executions) + finally: + event.remove(sqlite_session_factory, "after_begin", track_reader) + event.remove(sqlite_session_factory, "do_orm_execute", count_query) + + assert query_count == (execution_count + 1000) // 1000 + assert len(reads) == len(set(reads)) == execution_count + 1 + for index, execution in enumerate(executions): + payloads = full[execution.id] + assert payloads.inputs == {"file": files[index].key} + assert payloads.process_data == {"file": files[-1].key} + assert payloads.outputs == {"index": index} + + @pytest.mark.parametrize("logstore", [False, True]) def test_app_deletion_removes_nested_caller_records_and_offloads( sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, logstore: bool diff --git a/api/tests/unit_tests/repositories/test_trace_record_repository.py b/api/tests/unit_tests/repositories/test_trace_record_repository.py new file mode 100644 index 00000000000000..d1e9b6f1cd0b0d --- /dev/null +++ b/api/tests/unit_tests/repositories/test_trace_record_repository.py @@ -0,0 +1,433 @@ +import json +from datetime import UTC, datetime, timedelta +from decimal import Decimal + +import pytest +from dify_trace_sdk.events import ( + MessageCompleted, + ModerationChecked, + ToolInvoked, + TraceCorrelation, + TraceEvent, + TracePayload, + WorkflowRunCompleted, +) +from sqlalchemy import create_engine, text +from sqlalchemy.orm import Session, sessionmaker + +from extensions.logstore.aliyun_logstore import AliyunLogStore +from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus +from models.enums import ( + ConversationFromSource, + ConversationStatus, + CreatorUserRole, + EndUserType, + InvokeFrom, + WorkflowRunTriggeredFrom, +) +from models.model import App, AppMode, Conversation, EndUser, Message +from models.workflow import ( + WorkflowNodeExecutionModel, + WorkflowNodeExecutionTriggeredFrom, + WorkflowRun, + WorkflowType, +) +from repositories.factory import DifyAPIRepositoryFactory +from repositories.trace_record_repository import TraceRecordRepository +from services.tracing.errors import SourceNotFoundError, SourceNotReadyError, TracingStoreUnavailableError +from services.tracing.records import ( + MessageTraceInfo, + ModerationTraceInfo, + ToolTraceInfo, + TraceMessage, + WorkflowTraceInfo, +) +from tests.helpers.unit.config_override import apply_config_overrides + +TENANT_ID = "tenant-1" +APP_ID = "app-1" +MESSAGE_ID = "message-1" +CREATED_AT = datetime(2026, 1, 1, 8, 0, 0) +WORKFLOW_APP_ID = "workflow-app" +OUTER_RUN_ID = "outer-workflow-run" +TOOL_NODE_ID = "outer-tool-node" + + +@pytest.fixture +def records(sqlite_session: Session) -> None: + app = App( + id=APP_ID, + tenant_id=TENANT_ID, + name="Chat app", + description="", + mode=AppMode.CHAT, + enable_site=False, + enable_api=False, + ) + end_user = EndUser( + id="end-user-1", + tenant_id=TENANT_ID, + app_id=APP_ID, + type=EndUserType.SERVICE_API, + session_id="customer-42", + ) + conversation = Conversation( + id="conversation-1", + app_id=APP_ID, + app_model_config_id=None, + model_provider=None, + override_model_configs=None, + model_id=None, + mode=AppMode.CHAT, + name="Conversation", + _inputs={}, + introduction="", + system_instruction="", + system_instruction_tokens=0, + status=ConversationStatus.NORMAL, + invoke_from=InvokeFrom.SERVICE_API, + from_source=ConversationFromSource.API, + from_end_user_id=end_user.id, + from_account_id=None, + ) + message = Message( + id=MESSAGE_ID, + app_id=APP_ID, + conversation_id=conversation.id, + _inputs={}, + query="What is Dify?", + message=[{"role": "user", "text": "What is Dify?"}], + message_tokens=12, + message_unit_price=Decimal(0), + message_price_unit=Decimal(0), + answer="An LLM app platform.", + answer_tokens=8, + answer_unit_price=Decimal(0), + answer_price_unit=Decimal(0), + provider_response_latency=1.5, + total_price=Decimal(0), + currency="USD", + invoke_from=InvokeFrom.SERVICE_API, + from_source=ConversationFromSource.API, + from_end_user_id=end_user.id, + from_account_id=None, + app_mode=AppMode.CHAT, + ) + message.created_at = CREATED_AT + message.updated_at = CREATED_AT + timedelta(seconds=2) + sqlite_session.add_all([app, end_user, conversation, message]) + sqlite_session.commit() + + +def _workflow_run(run_id: str, app_id: str) -> WorkflowRun: + run = WorkflowRun( + id=run_id, + app_id=app_id, + tenant_id=TENANT_ID, + workflow_id="workflow-1", + type=WorkflowType.WORKFLOW, + version="1", + graph="{}", + inputs="{}", + outputs="{}", + status=WorkflowExecutionStatus.RUNNING, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="account-1", + finished_at=None, + elapsed_time=0, + total_tokens=0, + total_steps=0, + ) + run.created_at = CREATED_AT + timedelta(seconds=1) + return run + + +@pytest.fixture +def outer_runs(sqlite_session: Session) -> None: + """A workflow run with an owned Tool that can admit persisted child frames.""" + workflow_app = App( + id=WORKFLOW_APP_ID, + tenant_id=TENANT_ID, + name="Workflow app", + description="", + mode=AppMode.WORKFLOW, + enable_site=False, + enable_api=False, + ) + tool_node = WorkflowNodeExecutionModel( + id=TOOL_NODE_ID, + node_execution_id=TOOL_NODE_ID, + tenant_id=TENANT_ID, + app_id=WORKFLOW_APP_ID, + workflow_id="workflow-1", + workflow_run_id=OUTER_RUN_ID, + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + index=1, + node_id="tool", + node_type="tool", + title="Nested workflow", + inputs="{}", + outputs="{}", + process_data="{}", + status=WorkflowNodeExecutionStatus.RUNNING, + elapsed_time=0, + execution_metadata="{}", + finished_at=None, + created_by="account-1", + created_by_role=CreatorUserRole.ACCOUNT, + ) + tool_node.created_at = CREATED_AT + timedelta(seconds=2) + sqlite_session.add_all([workflow_app, _workflow_run(OUTER_RUN_ID, WORKFLOW_APP_ID), tool_node]) + sqlite_session.commit() + + +def _assembler(sessions: sessionmaker[Session]) -> TraceRecordRepository: + return TraceRecordRepository( + sessions, + workflow_runs=DifyAPIRepositoryFactory.create_api_workflow_run_repository(sessions), + node_executions=DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(sessions), + files_url="https://files.example.com", + enterprise_details=False, + ) + + +def _event(payload: TracePayload, *, tenant_id: str = TENANT_ID) -> TraceEvent: + return TraceEvent( + event_id="event-1", + tenant_id=tenant_id, + app_id=APP_ID, + user_id="customer-42", + occurred_at=datetime.now(UTC), + correlation=TraceCorrelation( + external_trace_id="caller-trace", + session_id="session-1", + parent_workflow_run_id=None, + parent_node_execution_id=None, + ), + payload=payload, + ) + + +@pytest.mark.usefixtures("records") +def test_message_records_carry_identity_correlation_and_usage(sqlite_session_factory: sessionmaker[Session]) -> None: + record = _assembler(sqlite_session_factory).assemble( + _event(MessageCompleted(message_id=MESSAGE_ID, conversation_id="conversation-1")) + ) + + assert isinstance(record, MessageTraceInfo) + assert record.trace_id == "caller-trace" + assert (record.message_tokens, record.answer_tokens, record.total_tokens) == (12, 8, 20) + assert record.inputs == [{"role": "user", "text": "What is Dify?"}] + assert (record.start_time, record.end_time) == (CREATED_AT, CREATED_AT + timedelta(seconds=2)) + assert record.metadata["trace_event_id"] == "event-1" + assert record.metadata["trace_session_id"] == "session-1" + assert record.metadata["end_user_session_id"] == "customer-42" + assert record.metadata["user_id"] == "end-user-1" + assert "app_name" in record.metadata + assert record.metadata["app_name"] == "" + + +@pytest.mark.usefixtures("records") +def test_events_are_read_only_through_their_own_workspace(sqlite_session_factory: sessionmaker[Session]) -> None: + with pytest.raises(SourceNotReadyError): + _assembler(sqlite_session_factory).assemble( + _event(MessageCompleted(message_id=MESSAGE_ID, conversation_id=None), tenant_id="tenant-2") + ) + + +@pytest.mark.usefixtures("records") +def test_messages_that_are_not_committed_yet_are_retried(sqlite_session_factory: sessionmaker[Session]) -> None: + with pytest.raises(SourceNotReadyError): + _assembler(sqlite_session_factory).assemble( + _event(MessageCompleted(message_id="message-not-yet-committed", conversation_id=None)) + ) + + +@pytest.mark.usefixtures("records") +def test_moderation_records_keep_the_verdict(sqlite_session_factory: sessionmaker[Session]) -> None: + payload = ModerationChecked( + message_id=MESSAGE_ID, + inputs={"topic": "weather"}, + query="sanitized", + flagged=True, + action="overridden", + preset_response="", + started_at=CREATED_AT, + ended_at=CREATED_AT + timedelta(milliseconds=30), + ) + + record = _assembler(sqlite_session_factory).assemble(_event(payload)) + + assert isinstance(record, ModerationTraceInfo) + assert (record.flagged, record.action, record.query) == (True, "overridden", "sanitized") + assert record.message_id == MESSAGE_ID + assert record.inputs == {"topic": "weather"} + + +@pytest.mark.usefixtures("records") +def test_tool_records_take_their_facts_from_the_event(sqlite_session_factory: sessionmaker[Session]) -> None: + payload = ToolInvoked( + message_id=MESSAGE_ID, + tool_name="web_search", + tool_inputs={"q": "dify"}, + tool_outputs="results", + tool_config={"provider": "builtin"}, + error=None, + time_cost=0.25, + started_at=CREATED_AT, + ended_at=CREATED_AT + timedelta(milliseconds=250), + ) + + record = _assembler(sqlite_session_factory).assemble(_event(payload)) + + assert isinstance(record, ToolTraceInfo) + assert (record.tool_name, record.tool_outputs, record.time_cost) == ("web_search", "results", 0.25) + assert record.tool_parameters == {"q": "dify"} + assert record.metadata["tool_config"] == {"provider": "builtin"} + assert record.file_url is None + + +def test_events_without_a_message_id_are_not_found(sqlite_session_factory: sessionmaker[Session]) -> None: + payload = ModerationChecked( + message_id=None, + inputs={}, + query="", + flagged=False, + action="direct_output", + preset_response="", + started_at=CREATED_AT, + ended_at=CREATED_AT, + ) + + with pytest.raises(SourceNotFoundError): + _assembler(sqlite_session_factory).assemble(_event(payload)) + + +@pytest.mark.usefixtures("records") +def test_records_copy_the_message_before_the_session_closes(sqlite_session_factory: sessionmaker[Session]) -> None: + record = _assembler(sqlite_session_factory).assemble( + _event(MessageCompleted(message_id=MESSAGE_ID, conversation_id="conversation-1")) + ) + + assert isinstance(record, MessageTraceInfo) + assert record.message_data == TraceMessage( + id=MESSAGE_ID, + app_id=APP_ID, + conversation_id="conversation-1", + workflow_run_id=None, + query="What is Dify?", + answer="An LLM app platform.", + error=None, + created_at=CREATED_AT, + model_provider=None, + model_id=None, + total_price=Decimal(0), + currency="USD", + ) + + +def test_an_unreachable_database_is_reported_as_unavailable() -> None: + sessions = sessionmaker(bind=create_engine("sqlite:///file:missing?mode=ro&uri=true")) + + with pytest.raises(TracingStoreUnavailableError): + _assembler(sessions).assemble(_event(MessageCompleted(message_id=MESSAGE_ID, conversation_id="conversation-1"))) + + +@pytest.mark.usefixtures("records", "outer_runs") +@pytest.mark.parametrize("logstore", [False, True]) +def test_workflow_trace_loads_only_descendants_of_its_owned_tools( + sqlite_session_factory: sessionmaker[Session], + sqlite_session: Session, + monkeypatch: pytest.MonkeyPatch, + logstore: bool, +) -> None: + if logstore: + # Run the LogStore adapter's real queries against the same owner-scoped records. + sqlite_session.execute( + text( + "CREATE VIEW workflow_node_execution AS SELECT *, 1 AS log_version, 1 AS __time__ " + "FROM workflow_node_executions" + ) + ) + monkeypatch.setattr(AliyunLogStore, "_instance", None) + monkeypatch.setattr(AliyunLogStore, "_initialized", False) + monkeypatch.setenv("ALIYUN_SLS_ENDPOINT", "localhost") + monkeypatch.setattr(AliyunLogStore, "_check_endpoint_connectivity", lambda *_args: None) + store = AliyunLogStore() + + def execute_sql(*, sql: str, **_kwargs: object) -> list[dict[str, object]]: + with sqlite_session_factory() as session: + return [dict(row) for row in session.execute(text(sql)).mappings()] + + monkeypatch.setattr(store, "execute_sql", execute_sql) + apply_config_overrides( + monkeypatch, + API_WORKFLOW_NODE_EXECUTION_REPOSITORY=( + "extensions.logstore.repositories.logstore_api_workflow_node_execution_repository." + "LogstoreAPIWorkflowNodeExecutionRepository" + ), + ) + run = sqlite_session.get(WorkflowRun, OUTER_RUN_ID) + tool = sqlite_session.get(WorkflowNodeExecutionModel, TOOL_NODE_ID) + assert run is not None + assert tool is not None + run.status = WorkflowExecutionStatus.SUCCEEDED + run.finished_at = CREATED_AT + timedelta(seconds=10) + run.total_tokens = 12 + tool.node_execution_id = "caller-engine-id" + tool.status = WorkflowNodeExecutionStatus.SUCCEEDED + tool.finished_at = run.finished_at + rows: list[tuple[str, str, str, str, dict[str, str]]] = [ + ("child-tool", "tool", "caller-engine-id", "workflow-1", {}), + ("grandchild", "llm", "child-tool-engine", "child-workflow", {}), + ("foreign-tenant", "llm", "caller-engine-id", "workflow-1", {"tenant_id": "other-tenant"}), + ("foreign-run", "llm", "caller-engine-id", "workflow-1", {"workflow_run_id": "other-run"}), + ("foreign-caller", "llm", "caller-engine-id", "other-workflow", {}), + ("orphan", "llm", "unrelated-tool", "workflow-1", {}), + ("wrong-origin", "llm", "caller-engine-id", "workflow-1", {"triggered_from": "workflow-run"}), + ("paused", "llm", "caller-engine-id", "workflow-1", {"status": "paused"}), + ] + for index, (row_id, node_type, caller, caller_workflow, overrides) in enumerate(rows): + row = WorkflowNodeExecutionModel( + id=row_id, + node_execution_id=f"{row_id}-engine", + tenant_id=overrides.get("tenant_id", TENANT_ID), + app_id="child-app", + workflow_id="child-workflow", + workflow_run_id=overrides.get("workflow_run_id", run.id), + triggered_from=WorkflowNodeExecutionTriggeredFrom(overrides.get("triggered_from", "workflow-tool")), + triggered_from_node_execution_id=caller, + triggered_from_workflow_id=caller_workflow, + index=index, + node_id="same-node", + node_type=node_type, + title=row_id, + status=WorkflowNodeExecutionStatus(overrides.get("status", "succeeded")), + inputs='{"query":"hello"}', + outputs=json.dumps({"usage": {"prompt_tokens": 5, "completion_tokens": 7}}) if node_type == "llm" else "{}", + process_data='{"workflow_tool_invocation_id":"invocation-1"}', + execution_metadata='{"future_metadata_key":"ignored","total_tokens":12}', + finished_at=run.finished_at, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="account-1", + ) + row.created_at = CREATED_AT + timedelta(seconds=3) + sqlite_session.add(row) + sqlite_session.commit() + event = _event(WorkflowRunCompleted(workflow_run_id=run.id, conversation_id=None, total_tokens=None)).model_copy( + update={"app_id": WORKFLOW_APP_ID} + ) + + record = _assembler(sqlite_session_factory).assemble(event) + + assert isinstance(record, WorkflowTraceInfo) + assert {node.id for node in record.node_executions} == {TOOL_NODE_ID, "child-tool", "grandchild"} + assert record.workflow_tool_parents == {"child-tool": TOOL_NODE_ID, "grandchild": "child-tool"} + assert (record.prompt_tokens, record.completion_tokens, record.total_tokens) == (5, 7, 12) + grandchild = next(node for node in record.node_executions if node.id == "grandchild") + assert grandchild.inputs == {"query": "hello"} + assert grandchild.process_data == {"workflow_tool_invocation_id": "invocation-1"} + for owner in ({"app_id": APP_ID}, {"tenant_id": "other-tenant"}): + with pytest.raises(SourceNotReadyError): + _assembler(sqlite_session_factory).assemble(event.model_copy(update=owner)) diff --git a/api/tests/unit_tests/schedule/test_collect_agent_sandbox_usage.py b/api/tests/unit_tests/schedule/test_collect_agent_sandbox_usage.py index e1d154b1fe09b3..c283ab38941a6f 100644 --- a/api/tests/unit_tests/schedule/test_collect_agent_sandbox_usage.py +++ b/api/tests/unit_tests/schedule/test_collect_agent_sandbox_usage.py @@ -42,7 +42,7 @@ def response(payload: object, status: int = 200) -> httpx.Response: ) -def test_task_uses_existing_queue_deadlines_and_single_authenticated_call(monkeypatch: pytest.MonkeyPatch) -> None: +def test_task_uses_configured_routing_deadlines_and_single_authenticated_call(monkeypatch: pytest.MonkeyPatch) -> None: post = MagicMock(return_value=response({"completed": True})) monkeypatch.setattr("schedule.collect_agent_sandbox_usage.ssrf_proxy.post", post) assert collect_agent_sandbox_usage() is True @@ -56,7 +56,7 @@ def test_task_uses_existing_queue_deadlines_and_single_authenticated_call(monkey dispatch = MagicMock() monkeypatch.setattr(collect_agent_sandbox_usage.app, "send_task", dispatch) collect_agent_sandbox_usage.apply_async() - assert dispatch.call_args.kwargs["queue"] == "ops_trace" + assert dispatch.call_args.kwargs.get("queue") is None assert collect_agent_sandbox_usage.soft_time_limit == 270 assert collect_agent_sandbox_usage.time_limit == 300 diff --git a/api/tests/unit_tests/services/knowledge_fs/test_workflow_runtime.py b/api/tests/unit_tests/services/knowledge_fs/test_workflow_runtime.py index f570313f51280d..381733d6ddd521 100644 --- a/api/tests/unit_tests/services/knowledge_fs/test_workflow_runtime.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_workflow_runtime.py @@ -13,6 +13,7 @@ from core.workflow.node_factory import DifyNodeFactory from core.workflow.node_runtime import build_knowledge_fs_node_dependencies from extensions import ext_application_services +from extensions.application_services.tracing import TracingServices, build_tracing_services from graphon.enums import WorkflowNodeExecutionStatus from graphon.runtime import RuntimeState, VariablePool from graphon.variables import StringSegment @@ -76,6 +77,7 @@ class _FakeAccounts: @dataclass class _FakeApplicationServices: + tracing: TracingServices accounts: _FakeAccounts = field(default_factory=_FakeAccounts) @@ -87,8 +89,11 @@ def test_application_registers_lazy_workflow_runtime_dependencies(monkeypatch: p monkeypatch.setattr(workflow_runtime, "get_knowledge_fs_runtime", resolve_runtime) monkeypatch.setattr(workflow_runtime.session_factory, "get_session_maker", lambda: session_maker) monkeypatch.setattr(ext_application_services, "get_session_maker", lambda: session_maker) + tracing = build_tracing_services(database_client=session_maker) monkeypatch.setattr( - ext_application_services, "build_application_services", lambda **_kwargs: _FakeApplicationServices() + ext_application_services, + "build_application_services", + lambda **_kwargs: _FakeApplicationServices(tracing=tracing), ) ext_application_services.init_app(app) diff --git a/api/tests/unit_tests/services/test_agent_tracing_config_service.py b/api/tests/unit_tests/services/test_agent_tracing_config_service.py index f1c4953bc0489e..a8d94ecee7e2ea 100644 --- a/api/tests/unit_tests/services/test_agent_tracing_config_service.py +++ b/api/tests/unit_tests/services/test_agent_tracing_config_service.py @@ -3,6 +3,7 @@ from collections.abc import Callable import pytest +from dify_trace_sdk.testing import RecordingEgress from sqlalchemy.orm import Session, sessionmaker from extensions.application_services.app import AppServices @@ -11,8 +12,8 @@ from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository from services.agent_tracing_config_service import AgentTracingConfigService from services.app.agent_app_contracts import AgentAppNotFoundError -from services.app_tracing_config_gateway import OpsTraceManagerGateway from services.app_tracing_config_service import AppTraceSettings, AppTracingConfigService +from tests.helpers.unit.services.tracing_ports import deployment_provider_configs type TraceConfigArguments = tuple[()] | tuple[AppTraceSettings] | tuple[str] | tuple[str, dict[str, str]] type RecordedCall = tuple[str, tuple[object, ...]] @@ -56,7 +57,7 @@ def test_agent_operations_resolve_tenant_owner_before_app_services( settings = app_services.console configs = AppTracingConfigService( configs=SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory), - provider=OpsTraceManagerGateway(), + provider=deployment_provider_configs(RecordingEgress([], grpc_target=None)), ) calls: list[RecordedCall] = [] diff --git a/api/tests/unit_tests/services/test_app_tracing_config_gateway.py b/api/tests/unit_tests/services/test_app_tracing_config_gateway.py deleted file mode 100644 index 5b85b730f8e7a8..00000000000000 --- a/api/tests/unit_tests/services/test_app_tracing_config_gateway.py +++ /dev/null @@ -1,2810 +0,0 @@ -import importlib as _test_importlib -import inspect -import json -import os -import traceback -from collections import deque -from collections.abc import Callable, Iterator, Mapping, Sequence -from copy import deepcopy -from dataclasses import dataclass, field, replace -from datetime import UTC, datetime, timedelta -from decimal import Decimal -from functools import partial -from typing import Never -from unittest.mock import MagicMock, patch -from urllib.parse import quote -from uuid import UUID, uuid4 - -import grpc # type: ignore[import-untyped] -import httpx -import mlflow -import opik -import pytest -import wandb -import weave -from dify_trace_aliyun.aliyun_trace import AliyunDataTrace -from dify_trace_arize_phoenix import arize_configuration -from dify_trace_arize_phoenix.arize_configuration import DESTINATION_KEY -from dify_trace_arize_phoenix.config import ArizeConfig -from dify_trace_arize_phoenix.configuration import PhoenixConfigurationError -from dify_trace_mlflow.configuration import DatabricksConfigurationError -from dify_trace_mlflow.mlflow_trace import MLflowDataTrace -from dify_trace_mlflow.oss_configuration import MLflowConfigurationError -from dify_trace_opik.configuration import OpikConfigurationError -from dify_trace_opik.opik_trace import OpikDataTrace -from dify_trace_tencent.configuration import TencentConfigurationError -from dify_trace_weave.configuration import WeaveConfigurationError -from dify_trace_weave.weave_trace import WeaveDataTrace -from flask import Flask -from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest -from opik.config import OpikConfig as SDKOpikConfig -from pydantic import ValidationInfo, field_validator -from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker - -from controllers.console.app import ops_trace as controller -from controllers.console.app.error import InvalidTracingConfigError -from core.helper import ssrf_proxy -from core.ops import ops_trace_manager -from core.ops.entities.config_entity import BaseTracingConfig, TracingProviderEnum -from core.ops.exceptions import TraceProviderNotInstalledError -from core.ops.ops_trace_manager import OpsTraceManager, TracingProviderConfigEntry -from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace -from extensions.ext_application_services import ApplicationServices -from machinery.context import RequestContext -from models.application_trace import ApplicationTraceRun -from models.enums import ConversationFromSource -from models.model import App, AppMode -from repositories.app.console_repository import ConsoleAppRepository -from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository -from repositories.application_trace_repository import ApplicationTraceRepository -from services import app_tracing_config_gateway as gateway_module -from services.app_tracing_config_gateway import OpsTraceManagerGateway -from services.app_tracing_config_service import ( - AppTracingConfigInvalidConfigurationError, - AppTracingConfigInvalidProviderError, - AppTracingConfigProcessingError, - AppTracingConfigProviderUnavailableError, - AppTracingConfigRecord, - AppTracingConfigService, - AppTracingConfigVerificationFailedError, -) -from services.tracing.provider_dispatch import export_trace, frozen_project_url -from services.tracing.service import ApplicationTracingService -from tests.helpers.unit.model_factories import make_app, make_message - -pytest_plugins = ("tests.unit_tests.core.ops.test_ops_trace_manager",) - -"""Real RSA/SQL/controller Aliyun configuration through a recorded HTTP boundary.""" - - -@dataclass -class AliyunCase: - gateway: OpsTraceManagerGateway - service: AppTracingConfigService - store: SQLAlchemyAppTracingConfigRepository - reader: ConsoleAppRepository - context: RequestContext - app_id: str - verification_calls: list[tuple[str, dict[str, object]]] = field(default_factory=list) - verification_succeeds: bool = True - - def read(self, provider: str = "aliyun") -> AppTracingConfigRecord | None: - return self.store.get( - workspace_id=self.context.active_workspace_id, app_id=self.app_id, tracing_provider=provider - ) - - -class TestAliyunTracingConfigGatewayCases: - @staticmethod - @pytest.fixture - def case( - workspace: str, sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch - ) -> AliyunCase: - app_id = str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add( - make_app( - app_id=app_id, - tenant_id=workspace, - tracing=json.dumps( - { - "enabled": True, - "tracing_provider": "langsmith", - "capture_content": True, - "capture_agent": False, - } - ), - ) - ) - gateway = OpsTraceManagerGateway() - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - old = OpsTraceManager.encrypt_tracing_config( - workspace, "langsmith", {"api_key": "synthetic-other", "project": "other", "endpoint": "https://other.test"} - ) - assert store.create(workspace_id=workspace, app_id=app_id, tracing_provider="langsmith", tracing_config=old) - result = AliyunCase( - gateway, - AppTracingConfigService(configs=store, provider=gateway), - store, - ConsoleAppRepository(session_factory=sqlite_session_factory), - RequestContext("request", None, "account", workspace), - app_id, - ) - - def verification_boundary(method: str, url: str, **options: object) -> httpx.Response: - assert method == "POST" - content = options["content"] - assert isinstance(content, bytes) - payload = ExportTraceServiceRequest.FromString(content) - assert len(payload.resource_spans) == 1 - assert len(payload.resource_spans[0].scope_spans) == 0 - assert [attribute.key for attribute in payload.resource_spans[0].resource.attributes] == ["service.name"] - result.verification_calls.append((url, dict(options))) - status = 200 if result.verification_succeeds else 403 - return httpx.Response(status, content=b'"success"', request=httpx.Request(method, url)) - - def forbidden(*_args: object, **_kwargs: object) -> Never: - raise AssertionError("Configuration must not initialize a legacy exporter or fetch its console link") - - monkeypatch.setattr(ssrf_proxy, "make_request", verification_boundary) - monkeypatch.setattr(OpsTraceManager, "check_trace_config_is_effective", forbidden) - monkeypatch.setattr(OpsTraceManager, "get_trace_config_project_url", forbidden) - monkeypatch.setattr(AliyunDataTrace, "__init__", forbidden) - return result - - @staticmethod - def test_post_masked_patch_and_get_preserve_real_ciphertext_and_policy(case: AliyunCase) -> None: - other = case.read("langsmith") - submitted = { - "endpoint": " HTTP://TRACING-ANALYSIS-DC-QD.ALIYUNCS.COM/reverse/ ", - "app_name": " 中文 app_Name ", - "license_key": "literal%2F@synthetic", - } - original = dict(submitted) - case.service.create(case.context, case.app_id, "aliyun", submitted, activate=True) - before = case.read() - assert before is not None - assert before.tracing_config is not None - stored = before.tracing_config - assert submitted == original - assert stored["endpoint"] == "http://tracing-analysis-dc-qd.aliyuncs.com/reverse" - assert stored["app_name"] == "中文 app_Name" - assert stored["license_key"] != submitted["license_key"] - assert OpsTraceManager.decrypt_tracing_config(case.context.active_workspace_id, "aliyun", stored)[ - "license_key" - ] == ("literal%2F@synthetic") - shown = case.service.get(case.context, case.app_id, "aliyun") - assert shown is not None - assert shown.tracing_config is not None - masked = shown.tracing_config["license_key"] - assert isinstance(masked, str) - assert "*" in masked - assert masked != submitted["license_key"] - assert shown.tracing_config["project_url"] == "https://arms.console.aliyun.com/#/llm" - assert len(case.verification_calls) == 1 - case.service.update( - case.context, - case.app_id, - "aliyun", - { - "endpoint": "https://project.region.log.aliyuncs.com/prefix/", - "app_name": " second_app ", - "license_key": masked, - }, - ) - after = case.read() - assert after is not None - assert after.tracing_config is not None - assert after.tracing_config["license_key"] == stored["license_key"] - assert after.tracing_config["app_name"] == "second_app" - assert after.tracing_config["endpoint"] == "https://project.region.log.aliyuncs.com/prefix" - assert case.verification_calls[-1][0] == ( - "https://project.region.log.aliyuncs.com/prefix/adapt_literal%252F@synthetic/api/v1/traces" - ) - assert case.service.get(case.context, case.app_id, "aliyun") is not None - assert len(case.verification_calls) == 2 - assert case.read("langsmith") == other - policy = case.reader.get_trace(case.context, case.app_id) - assert policy.tracing_provider == "aliyun" - assert policy.enabled - assert policy.capture_content is True - assert policy.capture_agent is False - - @staticmethod - @pytest.mark.parametrize("update", [False, True]) - @pytest.mark.parametrize("failure", ["invalid-schema", "verification-failed"]) - def test_failed_save_preserves_rows_and_active_policy(case: AliyunCase, update: bool, failure: str) -> None: - original = {"endpoint": "https://host.test", "app_name": "original", "license_key": "synthetic-secret"} - if update: - case.service.create(case.context, case.app_id, "aliyun", original) - before, other, policy = case.read(), case.read("langsmith"), case.reader.get_trace(case.context, case.app_id) - submitted = {**original, "app_name": "changed"} - if failure == "invalid-schema": - submitted["endpoint"] = "https://host.test/adapt_DO-NOT-PERSIST/api/otlp/traces" - expected = AppTracingConfigInvalidConfigurationError - else: - case.verification_succeeds = False - expected = AppTracingConfigVerificationFailedError - save = case.service.update if update else case.service.create - if update: - submitted["license_key"] = "******" - with pytest.raises(expected): - save(case.context, case.app_id, "aliyun", submitted, activate=True) - assert case.read() == before - assert case.read("langsmith") == other - assert case.reader.get_trace(case.context, case.app_id) == policy - - @staticmethod - @pytest.mark.parametrize("cached_url", [None, "https://wrong.test/stale"]) - def test_historical_get_is_pure_even_for_a_now_invalid_ingestion_url( - case: AliyunCase, cached_url: str | None - ) -> None: - stored = OpsTraceManager.encrypt_tracing_config( - case.context.active_workspace_id, - "aliyun", - { - "endpoint": "https://host.test/adapt_old/api/otlp/traces", - "app_name": "historical_name", - "license_key": "synthetic-key", - }, - ) - if cached_url: - stored["project_url"] = cached_url - before = dict(stored) - result = case.gateway.present_config( - workspace_id=case.context.active_workspace_id, tracing_provider="aliyun", tracing_config=stored - ) - assert result["project_url"] == "https://arms.console.aliyun.com/#/llm" - assert isinstance(result["license_key"], str) - assert "*" in result["license_key"] - assert stored == before - assert case.verification_calls == [] - - @staticmethod - @pytest.mark.parametrize("update", [False, True]) - @pytest.mark.parametrize("encoded", [False, True]) - def test_controller_rejects_embedded_license_without_echo_or_state_change( - case: AliyunCase, - app: Flask, - account_application_services: ApplicationServices, - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - update: bool, - encoded: bool, - ) -> None: - old_key = "SYNTHETIC_OLD_LICENSE" - original = {"endpoint": "https://host.test", "app_name": "old-app", "license_key": old_key} - if update: - case.service.create(case.context, case.app_id, "aliyun", original, activate=True) - before, other, policy = case.read(), case.read("langsmith"), case.reader.get_trace(case.context, case.app_id) - before_checks = len(case.verification_calls) - services = replace(account_application_services, app_tracing_configs=case.service) - monkeypatch.setattr(controller, "application_services", lambda: services) - key = "SYNTHETIC_PRIVATE_LICENSE@+/%2F" - url_key = quote(key, safe="") if encoded else key - payload = controller.TraceConfigPayload( - tracing_provider="aliyun", - tracing_config={ - "endpoint": f"https://host.test/adapt_{url_key}/api/otlp/traces", - "app_name": "should-not-save", - "license_key": "******" if update else key, - }, - activate=True, - ) - save = inspect.unwrap(controller.TraceAppConfigApi.patch if update else controller.TraceAppConfigApi.post) - with app.test_request_context(): - with pytest.raises(InvalidTracingConfigError) as caught: - save(controller.TraceAppConfigApi(), payload, case.context, UUID(case.app_id)) - # The real controller exposes this fixed BaseHTTPException body. Neither - # Pydantic inputs nor a credential-bearing exception chain reaches it. - assert caught.value.data == { - "code": "invalid_tracing_config", - "message": "The tracing configuration is invalid.", - "status": 400, - } - rendered = json.dumps(caught.value.data) + str(caught.value) + repr(caught.value) - rendered += "".join(traceback.format_exception(caught.value)) + caplog.text - assert key not in rendered - assert quote(key, safe="") not in rendered - assert old_key not in rendered - assert case.read() == before - assert case.read("langsmith") == other - assert case.reader.get_trace(case.context, case.app_id) == policy - assert len(case.verification_calls) == before_checks - - -class _ProviderConfig(BaseTracingConfig): - endpoint: str = "https://default.example.com" - project: str = "default-project" - - @field_validator("endpoint", "project", mode="before") - @classmethod - def replace_empty_with_default(cls, value: object, info: ValidationInfo) -> object: - if value != "": - return value - if info.field_name == "endpoint": - return "https://default.example.com" - return "default-project" - - -def _provider_entry(*, other_keys: list[str] | None = None) -> TracingProviderConfigEntry: - return { - "config_class": _ProviderConfig, - "secret_keys": ["api_key"], - "other_keys": other_keys or [], - "trace_instance": object, - } - - -class TestAppTracingConfigGatewayCases: - @staticmethod - def test_validate_provider_rejects_unknown_provider(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {}) - - with pytest.raises(AppTracingConfigInvalidProviderError, match="Invalid tracing provider: unknown"): - OpsTraceManagerGateway().validate_provider("unknown") - - @staticmethod - @pytest.mark.parametrize("provider", TracingProviderEnum) - def test_validate_provider_does_not_load_optional_dependencies(provider: TracingProviderEnum) -> None: - with patch.object(OpsTraceManagerGateway, "_provider_config") as load_provider: - OpsTraceManagerGateway().validate_provider(provider.value) - - load_provider.assert_not_called() - - @staticmethod - def test_require_provider_available_loads_provider_dependencies() -> None: - with patch.object(OpsTraceManagerGateway, "_provider_config") as load_provider: - OpsTraceManagerGateway().require_provider_available("weave") - - load_provider.assert_called_once_with("weave") - - @staticmethod - def test_require_provider_available_reports_missing_dependency() -> None: - unavailable = AppTracingConfigProviderUnavailableError() - with ( - patch.object(OpsTraceManagerGateway, "_provider_config", side_effect=unavailable), - pytest.raises(AppTracingConfigProviderUnavailableError) as caught, - ): - OpsTraceManagerGateway().require_provider_available("weave") - - assert caught.value is unavailable - - @staticmethod - @pytest.mark.parametrize("update", [False, True]) - def test_prepare_config_reports_missing_provider_dependencies( - monkeypatch: pytest.MonkeyPatch, update: bool - ) -> None: - missing_dependency = TraceProviderNotInstalledError("weave", "wandb") - providers = MagicMock() - providers.__getitem__.side_effect = missing_dependency - monkeypatch.setattr(gateway_module, "provider_config_map", providers) - gateway = OpsTraceManagerGateway() - prepare = ( - partial(gateway.prepare_updated_config, current_tracing_config={}) if update else gateway.prepare_new_config - ) - - with pytest.raises(AppTracingConfigProviderUnavailableError) as caught: - prepare(workspace_id="workspace-1", tracing_provider="weave", tracing_config={}) - - assert caught.value.__cause__ is missing_dependency - - @staticmethod - def test_present_config_reports_missing_provider_dependencies() -> None: - missing_dependency = TraceProviderNotInstalledError("weave", "wandb") - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.decrypt_tracing_config.side_effect = missing_dependency - - with pytest.raises(AppTracingConfigProviderUnavailableError) as caught: - OpsTraceManagerGateway().present_config( - workspace_id="workspace-1", tracing_provider="weave", tracing_config={"api_key": "encrypted"} - ) - - assert caught.value.__cause__ is missing_dependency - manager.obfuscated_decrypt_token.assert_not_called() - - @staticmethod - def test_prepare_new_config_applies_defaults_validates_and_encrypts(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - gateway_module, - "provider_config_map", - {"langsmith": _provider_entry(other_keys=["endpoint", "project"])}, - ) - submitted = {"api_key": "plain", "endpoint": "", "project": ""} - - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.check_trace_config_is_effective.return_value = True - manager.get_trace_config_project_url.return_value = "https://project.example.com" - manager.encrypt_tracing_config.return_value = {"api_key": "encrypted"} - - result = OpsTraceManagerGateway().prepare_new_config( - workspace_id="workspace-1", - tracing_provider="langsmith", - tracing_config=submitted, - ) - - normalized = { - "api_key": "plain", - "endpoint": "https://default.example.com", - "project": "default-project", - } - assert submitted == {"api_key": "plain", "endpoint": "", "project": ""} - assert result == {"api_key": "encrypted", "project_url": "https://project.example.com"} - manager.check_trace_config_is_effective.assert_called_once_with(normalized, "langsmith") - manager.encrypt_tracing_config.assert_called_once_with("workspace-1", "langsmith", normalized) - - @staticmethod - def test_prepare_new_config_reports_failed_verification_before_encryption( - monkeypatch: pytest.MonkeyPatch, - ) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) - - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.check_trace_config_is_effective.return_value = False - - with pytest.raises(AppTracingConfigVerificationFailedError): - OpsTraceManagerGateway().prepare_new_config( - workspace_id="workspace-1", - tracing_provider="langsmith", - tracing_config={"api_key": "plain"}, - ) - - manager.encrypt_tracing_config.assert_not_called() - - @staticmethod - def test_prepare_new_config_keeps_success_when_project_url_lookup_fails( - monkeypatch: pytest.MonkeyPatch, - ) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) - - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.check_trace_config_is_effective.return_value = True - manager.get_trace_config_project_url.side_effect = RuntimeError("provider unavailable") - manager.encrypt_tracing_config.return_value = {"api_key": "encrypted"} - - result = OpsTraceManagerGateway().prepare_new_config( - workspace_id="workspace-1", - tracing_provider="langsmith", - tracing_config={"api_key": "plain"}, - ) - - assert result == {"api_key": "encrypted"} - - @staticmethod - def test_prepare_new_langfuse_config_builds_project_url(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langfuse": _provider_entry()}) - - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.check_trace_config_is_effective.return_value = True - manager.get_trace_config_project_key.return_value = "project-key" - manager.encrypt_tracing_config.return_value = {"secret_key": "encrypted"} - - result = OpsTraceManagerGateway().prepare_new_config( - workspace_id="workspace-1", - tracing_provider="langfuse", - tracing_config={"host": "https://langfuse.example.com"}, - ) - - assert result["project_url"] == "https://langfuse.example.com/project/project-key" - - @staticmethod - def test_prepare_new_config_reports_provider_check_exception_as_failed_verification( - monkeypatch: pytest.MonkeyPatch, - ) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) - - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.check_trace_config_is_effective.side_effect = ValueError("verification failed") - - with pytest.raises(AppTracingConfigVerificationFailedError) as caught: - OpsTraceManagerGateway().prepare_new_config( - workspace_id="workspace-1", - tracing_provider="langsmith", - tracing_config={}, - ) - - assert isinstance(caught.value.__cause__, ValueError) - - @staticmethod - def test_prepare_new_config_propagates_unexpected_verification_failure( - monkeypatch: pytest.MonkeyPatch, - ) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) - failure = RuntimeError("unexpected provider bug") - - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.check_trace_config_is_effective.side_effect = failure - - with pytest.raises(RuntimeError) as caught: - OpsTraceManagerGateway().prepare_new_config( - workspace_id="workspace-1", - tracing_provider="langsmith", - tracing_config={}, - ) - - assert caught.value is failure - manager.encrypt_tracing_config.assert_not_called() - - @staticmethod - def test_prepare_new_config_rejects_invalid_schema_before_verification( - monkeypatch: pytest.MonkeyPatch, - ) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) - - with patch.object(gateway_module, "OpsTraceManager") as manager: - with pytest.raises(AppTracingConfigInvalidConfigurationError): - OpsTraceManagerGateway().prepare_new_config( - workspace_id="workspace-1", - tracing_provider="langsmith", - tracing_config={"endpoint": {"invalid": "value"}}, - ) - - manager.check_trace_config_is_effective.assert_not_called() - manager.encrypt_tracing_config.assert_not_called() - - @staticmethod - def test_prepare_new_config_reports_encryption_failure_as_processing_error(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) - failure = RuntimeError("key provider unavailable") - - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.check_trace_config_is_effective.return_value = True - manager.get_trace_config_project_url.return_value = None - manager.encrypt_tracing_config.side_effect = failure - - with pytest.raises(AppTracingConfigProcessingError) as caught: - OpsTraceManagerGateway().prepare_new_config( - workspace_id="workspace-1", - tracing_provider="langsmith", - tracing_config={}, - ) - - assert caught.value.__cause__ is failure - - @staticmethod - def test_prepare_updated_config_preserves_masked_secret_from_current_config( - monkeypatch: pytest.MonkeyPatch, - ) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) - submitted = {"api_key": "******", "project": "new-project"} - current = {"api_key": "old-encrypted", "project": "old-project"} - encrypted = {"api_key": "old-encrypted", "project": "new-project"} - decrypted = {"api_key": "old-plain", "project": "new-project"} - - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.encrypt_tracing_config.return_value = encrypted - manager.decrypt_tracing_config.return_value = decrypted - manager.check_trace_config_is_effective.return_value = True - manager.get_trace_config_project_url.return_value = "https://example.com/project/new" - - result = OpsTraceManagerGateway().prepare_updated_config( - workspace_id="workspace-1", - tracing_provider="langsmith", - tracing_config=submitted, - current_tracing_config=current, - ) - - assert result == encrypted - assert result["project_url"] == "https://example.com/project/new" - assert submitted == {"api_key": "******", "project": "new-project"} - manager.encrypt_tracing_config.assert_called_once_with("workspace-1", "langsmith", submitted, current) - manager.decrypt_tracing_config.assert_called_once_with("workspace-1", "langsmith", encrypted) - manager.check_trace_config_is_effective.assert_called_once_with(decrypted, "langsmith") - manager.get_trace_config_project_url.assert_called_once_with(decrypted, "langsmith") - - @staticmethod - def test_updated_destination_drops_old_project_when_lookup_fails(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.encrypt_tracing_config.return_value = {"api_key": "encrypted", "project": "new-project"} - manager.decrypt_tracing_config.return_value = {"api_key": "plain", "project": "new-project"} - manager.check_trace_config_is_effective.return_value = True - manager.get_trace_config_project_url.side_effect = RuntimeError("metadata unavailable") - result = OpsTraceManagerGateway().prepare_updated_config( - workspace_id="workspace-1", - tracing_provider="langsmith", - tracing_config={"project": "new-project"}, - current_tracing_config={"project_url": "https://example.com/project/old", "api_key": "encrypted"}, - ) - assert "project_url" not in result - - @staticmethod - def test_prepare_updated_config_validates_schema_before_encryption(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) - - with patch.object(gateway_module, "OpsTraceManager") as manager: - with pytest.raises(AppTracingConfigInvalidConfigurationError): - OpsTraceManagerGateway().prepare_updated_config( - workspace_id="workspace-1", - tracing_provider="langsmith", - tracing_config={"endpoint": {"invalid": "value"}}, - current_tracing_config={"api_key": "old-encrypted"}, - ) - - manager.encrypt_tracing_config.assert_not_called() - manager.decrypt_tracing_config.assert_not_called() - manager.check_trace_config_is_effective.assert_not_called() - - @staticmethod - def test_prepare_updated_config_reports_decryption_failure_as_processing_error( - monkeypatch: pytest.MonkeyPatch, - ) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) - failure = RuntimeError("stored credential cannot be decrypted") - - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.encrypt_tracing_config.return_value = {"api_key": "encrypted"} - manager.decrypt_tracing_config.side_effect = failure - - with pytest.raises(AppTracingConfigProcessingError) as caught: - OpsTraceManagerGateway().prepare_updated_config( - workspace_id="workspace-1", - tracing_provider="langsmith", - tracing_config={"project": "new-project"}, - current_tracing_config={"api_key": "old-encrypted"}, - ) - - assert caught.value.__cause__ is failure - manager.check_trace_config_is_effective.assert_not_called() - - @staticmethod - @pytest.mark.parametrize( - ("provider", "fallback_url"), - [ - ("arize", None), - ("phoenix", "https://app.phoenix.arize.com/redirects/projects/default"), - ("langsmith", "https://smith.langchain.com/"), - ("opik", "https://www.comet.com/opik/"), - ("weave", "https://wandb.ai/"), - ("aliyun", "https://arms.console.aliyun.com/#/llm"), - ("tencent", "https://console.cloud.tencent.com/apm"), - ("mlflow", "http://localhost:5000/#/experiments/0/traces"), - ("databricks", "https://www.databricks.com/"), - ], - ) - def test_present_config_uses_provider_fallback_when_project_lookup_fails( - provider: str, - fallback_url: str | None, - ) -> None: - with patch.object(gateway_module, "OpsTraceManager") as manager: - decrypted_config: dict[str, object] = {} - presented_config: dict[str, object] = {} - manager.decrypt_tracing_config.return_value = decrypted_config - manager.obfuscated_decrypt_token.return_value = presented_config - manager.get_trace_config_project_url.side_effect = RuntimeError("provider unavailable") - - result = OpsTraceManagerGateway().present_config( - workspace_id="workspace-1", - tracing_provider=provider, - tracing_config={"encrypted": "config"}, - ) - - assert result == {"project_url": fallback_url} - if provider in {"arize", "aliyun", "mlflow", "phoenix"}: - # Provider defaults produce a pure link; the failing legacy lookup is - # deliberately not reached, even when no cached URL is present. - manager.get_trace_config_project_url.assert_not_called() - - @staticmethod - def test_present_langfuse_config_builds_project_url() -> None: - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.decrypt_tracing_config.return_value = {"host": "https://langfuse.example.com"} - manager.obfuscated_decrypt_token.return_value = {"host": "https://langfuse.example.com"} - manager.get_trace_config_project_key.return_value = "project-key" - - result = OpsTraceManagerGateway().present_config( - workspace_id="workspace-1", - tracing_provider="langfuse", - tracing_config={"encrypted": "config"}, - ) - - assert result["project_url"] == "https://langfuse.example.com/project/project-key" - - @staticmethod - def test_present_langfuse_config_falls_back_to_host() -> None: - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.decrypt_tracing_config.return_value = {"host": "https://langfuse.example.com"} - manager.obfuscated_decrypt_token.return_value = {"host": "https://langfuse.example.com"} - manager.get_trace_config_project_key.side_effect = RuntimeError("provider unavailable") - - result = OpsTraceManagerGateway().present_config( - workspace_id="workspace-1", - tracing_provider="langfuse", - tracing_config={"encrypted": "config"}, - ) - - assert result["project_url"] == "https://langfuse.example.com/" - - @staticmethod - def test_present_config_rejects_missing_stored_config() -> None: - with pytest.raises(AppTracingConfigProcessingError, match="processing failed"): - OpsTraceManagerGateway().present_config( - workspace_id="workspace-1", - tracing_provider="langsmith", - tracing_config=None, - ) - - @staticmethod - def test_present_config_reports_decryption_failure_as_processing_error() -> None: - failure = RuntimeError("stored credential cannot be decrypted") - - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.decrypt_tracing_config.side_effect = failure - - with pytest.raises(AppTracingConfigProcessingError) as caught: - OpsTraceManagerGateway().present_config( - workspace_id="workspace-1", - tracing_provider="langsmith", - tracing_config={"encrypted": "config"}, - ) - - assert caught.value.__cause__ is failure - manager.obfuscated_decrypt_token.assert_not_called() - - @staticmethod - @pytest.mark.parametrize("malformed", [False, True]) - def test_validate_saved_config_decrypts_and_validates_without_vendor_io( - monkeypatch: pytest.MonkeyPatch, malformed: bool - ) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.decrypt_tracing_config.return_value = { - "project": ["invalid-project"] if malformed else "saved-project" - } - gateway = OpsTraceManagerGateway() - if malformed: - with pytest.raises(AppTracingConfigInvalidConfigurationError): - gateway.validate_saved_config( - workspace_id="workspace", tracing_provider="langsmith", tracing_config={"api_key": "encrypted"} - ) - else: - gateway.validate_saved_config( - workspace_id="workspace", tracing_provider="langsmith", tracing_config={"api_key": "encrypted"} - ) - manager.decrypt_tracing_config.assert_called_once_with("workspace", "langsmith", {"api_key": "encrypted"}) - manager.check_trace_config_is_effective.assert_not_called() - - @staticmethod - def test_validate_saved_config_rejects_undecryptable_credentials(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.decrypt_tracing_config.side_effect = ValueError("invalid ciphertext") - with pytest.raises(AppTracingConfigProcessingError): - OpsTraceManagerGateway().validate_saved_config( - workspace_id="workspace", tracing_provider="langsmith", tracing_config={"api_key": "broken"} - ) - - -"""Real provider config serialization must retain verified project links on update.""" - - -AccessAPI = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-opik.src.dify_trace_opik.test_configuration" -).AccessAPI - - -@dataclass -class ProviderBoundary: - calls: list[tuple[str, str, dict[str, object]]] = field(default_factory=list) - effective: bool = True - lookup_fails: bool = False - - def verify(self, config: Mapping[str, object], provider: str) -> bool: - self.calls.append(("verify", provider, dict(config))) - return self.effective - - def project_url(self, config: Mapping[str, object], provider: str) -> str: - self.calls.append(("url", provider, dict(config))) - if self.lookup_fails: - raise RuntimeError("Synthetic project lookup failure") - return f"https://tracing.test/{provider}/{config['project']}" - - def project_key(self, config: Mapping[str, object], provider: str) -> str: - self.calls.append(("key", provider, dict(config))) - return "verified-project" - - -class TestAppTracingConfigProjectUrlCases: - @staticmethod - @pytest.fixture - def boundary(monkeypatch: pytest.MonkeyPatch) -> ProviderBoundary: - boundary = ProviderBoundary() - - def encrypt(_tenant: str, value: str) -> str: - return "test-cipher:" + value - - def decrypt(_tenant: str, values: Sequence[str]) -> list[str]: - assert all(value.startswith("test-cipher:") for value in values) - return [value.removeprefix("test-cipher:") for value in values] - - monkeypatch.setattr(ops_trace_manager, "encrypt_token", encrypt) - monkeypatch.setattr(ops_trace_manager, "batch_decrypt_token", decrypt) - monkeypatch.setattr(OpsTraceManager, "check_trace_config_is_effective", boundary.verify) - monkeypatch.setattr(OpsTraceManager, "get_trace_config_project_url", boundary.project_url) - monkeypatch.setattr(OpsTraceManager, "get_trace_config_project_key", boundary.project_key) - return boundary - - @staticmethod - @pytest.mark.parametrize("provider", ["langsmith", "opik"]) - @pytest.mark.parametrize("project", ["original", "changed"]) - def test_update_resolves_project_link_using_verified_unmasked_configuration( - boundary: ProviderBoundary, provider: str, project: str, monkeypatch: pytest.MonkeyPatch - ) -> None: - gateway = OpsTraceManagerGateway() - workspace = str(uuid4()) - api = AccessAPI() - if provider == "opik": - monkeypatch.setattr(ssrf_proxy, "make_request", api.request) - target = {"workspace": "team", "url": "https://tracing.test/opik/api/"} if provider == "opik" else {} - current = gateway.prepare_new_config( - workspace_id=workspace, - tracing_provider=provider, - tracing_config={"api_key": "synthetic-key", "project": "original", **target}, - ) - submitted = {"api_key": "******", "project": project, "project_url": "https://untrusted.test/stale", **target} - updated = gateway.prepare_updated_config( - workspace_id=workspace, tracing_provider=provider, tracing_config=submitted, current_tracing_config=current - ) - - prefix = f"https://tracing.test/{provider}/" + ("team/redirect/projects?name=" if provider == "opik" else "") - assert updated["project_url"] == prefix + project - assert updated["api_key"] == current["api_key"] == "test-cipher:synthetic-key" - assert current["project_url"] == prefix + "original" - assert submitted["api_key"] == "******" - if provider == "opik": - assert boundary.calls == [] - assert len(api.calls) == 2 - assert api.calls[-1][2]["headers"] == { - "Authorization": "synthetic-key", - "Comet-Workspace": "team", - "Accept": "application/json", - } - else: - assert [call[0] for call in boundary.calls] == ["verify", "url", "verify", "url"] - assert boundary.calls[-1][2]["api_key"] == "synthetic-key" - assert boundary.calls[-1][2]["project"] == project - - @staticmethod - def test_langfuse_update_uses_current_host_and_resolved_project_key(boundary: ProviderBoundary) -> None: - gateway = OpsTraceManagerGateway() - workspace = str(uuid4()) - current = gateway.prepare_new_config( - workspace_id=workspace, - tracing_provider="langfuse", - tracing_config={ - "public_key": "synthetic-public", - "secret_key": "synthetic-secret", - "host": "https://old.test", - }, - ) - updated = gateway.prepare_updated_config( - workspace_id=workspace, - tracing_provider="langfuse", - tracing_config={"public_key": "******", "secret_key": "******", "host": "https://new.test"}, - current_tracing_config=current, - ) - assert updated["project_url"] == "https://new.test/project/verified-project" - assert current["project_url"] == "https://old.test/project/verified-project" - assert boundary.calls[-1][0] == "key" - assert boundary.calls[-1][2]["secret_key"] == "synthetic-secret" - - @staticmethod - def test_failed_optional_link_lookup_does_not_reuse_stale_project_url(boundary: ProviderBoundary) -> None: - gateway = OpsTraceManagerGateway() - workspace = str(uuid4()) - current = gateway.prepare_new_config( - workspace_id=workspace, - tracing_provider="langsmith", - tracing_config={"api_key": "synthetic-key", "project": "original"}, - ) - boundary.lookup_fails = True - updated = gateway.prepare_updated_config( - workspace_id=workspace, - tracing_provider="langsmith", - tracing_config={"api_key": "******", "project": "changed"}, - current_tracing_config=current, - ) - assert updated["project"] == "changed" - assert "project_url" not in updated - assert boundary.calls[-1][0] == "url" - - @staticmethod - def test_invalid_updated_configuration_never_resolves_a_project_url(boundary: ProviderBoundary) -> None: - boundary.effective = False - with pytest.raises(AppTracingConfigVerificationFailedError): - OpsTraceManagerGateway().prepare_updated_config( - workspace_id=str(uuid4()), - tracing_provider="langsmith", - tracing_config={"api_key": "synthetic-key", "project": "changed"}, - current_tracing_config=None, - ) - assert [call[0] for call in boundary.calls] == ["verify"] - - @staticmethod - def test_updated_project_url_is_frozen_and_used_by_real_langsmith_serializer( - boundary: ProviderBoundary, - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - ) -> None: - workspace, app_id, message_id = (str(uuid4()) for _ in range(3)) - with sqlite_session_factory.begin() as session: - session.add( - make_app( - app_id=app_id, - tenant_id=workspace, - mode=AppMode.AGENT, - tracing=json.dumps({"enabled": True, "tracing_provider": "langsmith"}), - ) - ) - session.add( - make_message( - message_id=message_id, - app_id=app_id, - inputs={}, - query="test input", - message={}, - answer="test output", - message_unit_price=Decimal(0), - answer_unit_price=Decimal(0), - currency="USD", - from_source=ConversationFromSource.CONSOLE, - ) - ) - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - configs = AppTracingConfigService(configs=store, provider=OpsTraceManagerGateway()) - context = RequestContext("request", None, "account", workspace) - configs.create(context, app_id, "langsmith", {"api_key": "synthetic-key", "project": "original"}) - configs.update(context, app_id, "langsmith", {"api_key": "******", "project": "changed"}) - repository = ApplicationTraceRepository(sqlite_session_factory) - run = repository.start( - tenant_id=workspace, app_id=app_id, source_kind="message", source_id=message_id, session_id="test-session" - ) - assert run is not None - assert run.provider_config["project_url"] == "https://tracing.test/langsmith/changed" - configs.update(context, app_id, "langsmith", {"api_key": "******", "project": "later"}) - frozen = repository.source_run(tenant_id=workspace, app_id=app_id, source_kind="message", source_id=message_id) - assert frozen is not None - assert frozen.provider_config["project_url"] == "https://tracing.test/langsmith/changed" - - requests: list[str] = [] - - def send(method: str, url: str, **_kwargs: object) -> httpx.Response: - requests.append(url) - return httpx.Response(200, request=httpx.Request(method, url)) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - config = dict(frozen.provider_config) - config.update(OpsTraceManager.decrypt_tracing_config(workspace, "langsmith", frozen.provider_config)) - start = datetime(2026, 10, 16, tzinfo=UTC) - trace = CanonicalTrace( - trace_id=frozen.trace_id, - session_id=frozen.session_id, - root_span_id=frozen.root_span_id, - spans=( - CanonicalSpan( - id=frozen.root_span_id, - name="Agent", - kind=CanonicalSpanKind.AGENT, - parent_id=None, - status=CanonicalSpanStatus.OK, - start_time=start, - end_time=start + timedelta(seconds=1), - ), - ), - ) - result = export_trace("langsmith", config, trace) - assert result.status == "accepted" - assert result.trace_url == f"https://tracing.test/langsmith/changed/r/{result.trace_id}" - assert result.span_urls[frozen.root_span_id] == result.trace_url - assert len(requests) == 1 - assert boundary.calls[-1][2]["project"] == "later" - - -"""Arize project links use verified metadata without changing ingestion permissions.""" - - -DESTINATION = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_arize_configuration" -).DESTINATION -SPACE = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_arize_configuration" -).SPACE -URL = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_arize_configuration" -).URL -VALUES = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_arize_configuration" -).VALUES -page = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_arize_configuration" -).page -project = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_arize_configuration" -).project -_arize_tracing_project_url_response = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_arize_configuration" -).response -space = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_arize_configuration" -).space -_D0db4e12TestArizeConfigurationCases = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_arize_configuration" -).TestArizeConfigurationCases -_D0db4e12TestArizeConfigurationCases = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_arize_configuration" -).TestArizeConfigurationCases - - -class TestArizeTracingProjectUrlCases: - calls = staticmethod(_D0db4e12TestArizeConfigurationCases.calls) - - responses = staticmethod(_D0db4e12TestArizeConfigurationCases.responses) - - @staticmethod - @pytest.fixture - def verified(monkeypatch: pytest.MonkeyPatch) -> list[dict[str, object]]: - calls: list[dict[str, object]] = [] - - def verify(config: ArizeConfig) -> None: - calls.append(config.model_dump()) - - monkeypatch.setattr(arize_configuration, "verify_configuration", verify) - return calls - - @staticmethod - def test_legacy_arize_get_hides_unverified_redirect_without_network_or_rewriting( - workspace: str, monkeypatch: pytest.MonkeyPatch - ) -> None: - stored = OpsTraceManager.encrypt_tracing_config( - workspace, - "arize", - {"api_key": "synthetic-ingestion-key", "space_id": "U3BhY2U6dGVzdA==", "project": "original"}, - ) - stored["project_url"] = "https://app.arize.com/?redirect_project_name=original" - untouched = deepcopy(stored) - calls: list[str] = [] - - def forbidden(*_args: object, **_kwargs: object) -> Never: - calls.append("legacy or network") - raise AssertionError("Configuration GET must not initialize a tracer or resolve vendor metadata") - - monkeypatch.setattr(OpsTraceManager, "get_trace_config_project_url", forbidden) - monkeypatch.setattr(ssrf_proxy, "get", forbidden) - shown = OpsTraceManagerGateway().present_config( - workspace_id=workspace, tracing_provider="arize", tracing_config=stored - ) - - assert shown["project_url"] is None - assert calls == [] - assert stored == untouched - - @staticmethod - def test_real_rsa_create_masked_update_get_and_frozen_history( - workspace: str, - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - verified: list[dict[str, object]], - calls: list[tuple[str, dict[str, object]]], - responses: deque[httpx.Response | Exception], - ) -> None: - app_id, message_id = str(uuid4()), str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add( - make_app( - app_id=app_id, - tenant_id=workspace, - mode=AppMode.AGENT, - tracing=json.dumps({"enabled": True, "tracing_provider": "arize"}), - ) - ) - session.add( - make_message( - message_id=message_id, - app_id=app_id, - inputs={}, - query="test", - message={}, - answer="ok", - message_unit_price=Decimal(0), - answer_unit_price=Decimal(0), - currency="USD", - from_source=ConversationFromSource.CONSOLE, - ) - ) - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - gateway = OpsTraceManagerGateway() - service = AppTracingConfigService(configs=store, provider=gateway) - context = RequestContext("request", None, "account", workspace) - legacy = OpsTraceManager.encrypt_tracing_config(workspace, "arize", VALUES) - legacy["project_url"] = "https://app.arize.com/?redirect_project_name=exact" - assert store.create(workspace_id=workspace, app_id=app_id, tracing_provider="arize", tracing_config=legacy) - repository = ApplicationTraceRepository(sqlite_session_factory) - run = repository.start( - tenant_id=workspace, app_id=app_id, source_kind="message", source_id=message_id, session_id="session" - ) - assert run is not None - with sqlite_session_factory.begin() as session: - accepted = session.get(ApplicationTraceRun, run.id) - assert accepted is not None - accepted.state = "accepted" - accepted.attempts = 1 - with sqlite_session_factory() as session: - accepted_row = deepcopy( - dict( - session.execute(select(ApplicationTraceRun.__table__).where(ApplicationTraceRun.id == run.id)) - .mappings() - .one() - ) - ) - frozen_before = deepcopy(run.provider_config) - - # Both legacy secret fields are really encrypted, rather than identity stubs. - assert legacy["api_key"] != VALUES["api_key"] - assert legacy["space_id"] != SPACE - responses.extend([space(), page(project())]) - submitted = { - **VALUES, - "api_key": "******", - "space_id": "******", - DESTINATION_KEY: {"organization_id": "forged", "project_id": "forged"}, - "project_url": "https://untrusted.test/forged", - } - untouched = deepcopy(submitted) - service.update(context, app_id, "arize", submitted) - current = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="arize") - assert current is not None - assert current.tracing_config is not None - saved = current.tracing_config - assert saved["api_key"] == legacy["api_key"] - assert saved["space_id"] == legacy["space_id"] - assert saved[DESTINATION_KEY] == {**DESTINATION, "source_digest": saved[DESTINATION_KEY]["source_digest"]} - assert saved["project_url"] == URL - assert submitted == untouched - assert len(calls) == 2 - assert verified[-1]["api_key"] == VALUES["api_key"] - assert verified[-1]["space_id"] == SPACE - - def forbidden(*_args: object, **_kwargs: object) -> Never: - raise AssertionError("Reading current configuration or frozen history must not use vendor I/O") - - monkeypatch.setattr(ssrf_proxy, "get", forbidden) - monkeypatch.setattr(OpsTraceManager, "get_trace_config_project_url", forbidden) - shown = service.get(context, app_id, "arize") - assert shown is not None - assert shown.tracing_config is not None - assert shown.tracing_config["project_url"] == URL - assert DESTINATION_KEY not in shown.tracing_config - frozen = repository.source_run(tenant_id=workspace, app_id=app_id, source_kind="message", source_id=message_id) - assert frozen is not None - assert (frozen.state, frozen.attempts, frozen.provider_config) == ("accepted", 1, frozen_before) - assert frozen_project_url("arize", frozen.provider_config) is None - assert frozen.provider_config == frozen_before - monkeypatch.setattr(OpsTraceManager, "decrypt_tracing_config", forbidden) - status = ApplicationTracingService(repository).status( - tenant_id=workspace, app_id=app_id, source_kind="message", source_id=message_id - ) - assert status["project_url"] is None - assert status["status"] == "accepted" - with sqlite_session_factory() as session: - after_row = dict( - session.execute(select(ApplicationTraceRun.__table__).where(ApplicationTraceRun.id == run.id)) - .mappings() - .one() - ) - assert after_row == accepted_row - - # A future source freezes the verified destination; it never repairs the old run. - later_id = str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add( - make_message( - message_id=later_id, - app_id=app_id, - inputs={}, - query="later", - message={}, - answer="ok", - message_unit_price=Decimal(0), - answer_unit_price=Decimal(0), - currency="USD", - from_source=ConversationFromSource.CONSOLE, - ) - ) - later = repository.start( - tenant_id=workspace, app_id=app_id, source_kind="message", source_id=later_id, session_id="session" - ) - assert later is not None - assert frozen_project_url("arize", later.provider_config) == URL - - @staticmethod - @pytest.mark.parametrize("status", [401, 403, 404, 500]) - def test_ingest_verified_key_without_metadata_read_permission_still_saves( - workspace: str, - verified: list[dict[str, object]], - calls: list[tuple[str, dict[str, object]]], - responses: deque[httpx.Response | Exception], - status: int, - ) -> None: - responses.append(_arize_tracing_project_url_response({"message": "private-response"}, status)) - saved = OpsTraceManagerGateway().prepare_new_config( - workspace_id=workspace, tracing_provider="arize", tracing_config=VALUES - ) - assert "project_url" not in saved - assert DESTINATION_KEY not in saved - assert OpsTraceManager.decrypt_tracing_config(workspace, "arize", saved)["api_key"] == VALUES["api_key"] - assert len(verified) == len(calls) == 1 - - @staticmethod - def test_changed_destination_cannot_keep_old_or_client_supplied_cache( - workspace: str, - verified: list[dict[str, object]], - calls: list[tuple[str, dict[str, object]]], - responses: deque[httpx.Response | Exception], - ) -> None: - gateway = OpsTraceManagerGateway() - responses.extend([space(), page(project())]) - before = gateway.prepare_new_config(workspace_id=workspace, tracing_provider="arize", tracing_config=VALUES) - assert before["project_url"] == URL - responses.append(_arize_tracing_project_url_response({}, 403)) - after = gateway.prepare_updated_config( - workspace_id=workspace, - tracing_provider="arize", - current_tracing_config=before, - tracing_config={ - **VALUES, - "api_key": "******", - "space_id": "******", - "project": "later", - DESTINATION_KEY: before[DESTINATION_KEY], - "project_url": URL, - }, - ) - assert after["project"] == "later" - assert after["api_key"] == before["api_key"] - assert after["space_id"] == before["space_id"] - assert "project_url" not in after - assert DESTINATION_KEY not in after - assert before["project_url"] == URL - assert len(calls) == 3 - assert len(verified) == 2 - - @staticmethod - def test_failed_original_ingestion_verification_never_resolves_metadata( - workspace: str, - monkeypatch: pytest.MonkeyPatch, - calls: list[tuple[str, dict[str, object]]], - ) -> None: - def rejected(_config: ArizeConfig) -> Never: - raise arize_configuration.ArizeConfigurationError("permission_denied") - - monkeypatch.setattr(arize_configuration, "verify_configuration", rejected) - with pytest.raises(AppTracingConfigVerificationFailedError): - OpsTraceManagerGateway().prepare_new_config( - workspace_id=workspace, tracing_provider="arize", tracing_config=VALUES - ) - assert calls == [] - - -"""Databricks config persistence must not initialize a process-global MLflow client.""" - - -ExperimentAPI = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-mlflow.src.dify_trace_mlflow.test_configuration" -).ExperimentAPI -_databricks_tracing_config_gateway_response = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-mlflow.src.dify_trace_mlflow.test_configuration" -).response - - -@dataclass -class _DatabricksTracingConfigGatewayGatewayenvironment: - gateway: OpsTraceManagerGateway = field(default_factory=OpsTraceManagerGateway) - api: ExperimentAPI = field(default_factory=ExperimentAPI) - forbidden_calls: list[str] = field(default_factory=list) - - -class TestDatabricksTracingConfigGatewayCases: - @staticmethod - @pytest.fixture - def environment(monkeypatch: pytest.MonkeyPatch) -> Iterator[_DatabricksTracingConfigGatewayGatewayenvironment]: - environment = _DatabricksTracingConfigGatewayGatewayenvironment() - - def encrypt(_tenant: str, value: str) -> str: - return "test-cipher:" + value - - def decrypt(_tenant: str, values: Sequence[str]) -> list[str]: - assert all(value.startswith("test-cipher:") for value in values) - return [value.removeprefix("test-cipher:") for value in values] - - def forbidden(name: str) -> Callable[..., Never]: - def block(*_args: object, **_kwargs: object) -> Never: - environment.forbidden_calls.append(name) - raise AssertionError("Configuration must not initialize a stateful exporter") - - return block - - monkeypatch.setattr(ops_trace_manager, "encrypt_token", encrypt) - monkeypatch.setattr(ops_trace_manager, "batch_decrypt_token", decrypt) - monkeypatch.setattr(ops_trace_manager, "obfuscated_token", lambda value: "******" if value else value) - monkeypatch.setattr(ssrf_proxy, "get", environment.api.get) - monkeypatch.setattr(ssrf_proxy, "post", environment.api.post) - monkeypatch.setattr(MLflowDataTrace, "__init__", forbidden("legacy_constructor")) - monkeypatch.setattr(mlflow, "set_tracking_uri", forbidden("tracking_uri")) - monkeypatch.setattr(mlflow, "set_experiment", forbidden("experiment")) - monkeypatch.setattr(mlflow, "search_experiments", forbidden("search")) - monkeypatch.setattr(mlflow.config, "enable_async_logging", forbidden("async_logging")) - yield environment - assert environment.forbidden_calls == [] - - @staticmethod - @pytest.mark.parametrize("oauth", [False, True]) - def test_create_and_patch_persist_verified_target_without_replacing_masked_credentials( - environment: _DatabricksTracingConfigGatewayGatewayenvironment, - sqlite_session_factory: sessionmaker[Session], - oauth: bool, - ) -> None: - workspace, app_id = str(uuid4()), str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add(make_app(app_id=app_id, tenant_id=workspace)) - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - service = AppTracingConfigService(configs=store, provider=environment.gateway) - context = RequestContext("request", None, "account", workspace) - credentials = ( - {"client_id": "synthetic-client", "client_secret": "synthetic-secret"} - if oauth - else {"personal_access_token": "synthetic-pat"} - ) - secret_field = "client_secret" if oauth else "personal_access_token" - service.create( - context, app_id, "databricks", {"host": " https://first.test/ ", "experiment_id": "00012", **credentials} - ) - before = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="databricks") - assert before is not None - assert before.tracing_config is not None - assert before.tracing_config["host"] == "https://first.test" - assert before.tracing_config["experiment_id"] == "12" - service.update( - context, - app_id, - "databricks", - { - "host": " https://second.test/ ", - "experiment_id": "00034", - **credentials, - secret_field: "******", - }, - ) - after = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="databricks") - assert after is not None - assert after.tracing_config is not None - assert after.tracing_config["host"] == "https://second.test" - assert after.tracing_config["experiment_id"] == "34" - assert after.tracing_config[secret_field] == before.tracing_config[secret_field] - assert after.tracing_config["project_url"] == "https://second.test/ml/experiments/34/traces" - assert before.tracing_config["project_url"] == "https://first.test/ml/experiments/12/traces" - before_get = len(environment.api.calls) - presented = service.get(context, app_id, "databricks") - assert presented is not None - assert presented.tracing_config is not None - assert presented.tracing_config[secret_field] == "******" - assert presented.tracing_config["project_url"] == after.tracing_config["project_url"] - environment.gateway.validate_saved_config( - workspace_id=workspace, tracing_provider="databricks", tracing_config=after.tracing_config - ) - assert len(environment.api.calls) == before_get - - @staticmethod - @pytest.mark.parametrize("saved_url", [None, "https://stale.test/old-project"]) - def test_get_is_pure_for_historical_and_current_configs( - environment: _DatabricksTracingConfigGatewayGatewayenvironment, saved_url: str | None - ) -> None: - workspace = str(uuid4()) - encrypted = OpsTraceManager.encrypt_tracing_config( - workspace, - "databricks", - {"host": "https://workspace.test/", "experiment_id": "12", "personal_access_token": "synthetic-pat"}, - ) - if saved_url is not None: - encrypted["project_url"] = saved_url - result = environment.gateway.present_config( - workspace_id=workspace, tracing_provider="databricks", tracing_config=encrypted - ) - assert result["project_url"] == "https://workspace.test/ml/experiments/12/traces" - assert result["personal_access_token"] == "******" - assert environment.api.calls == [] - - @staticmethod - def test_invalid_historical_host_uses_safe_pure_fallback( - environment: _DatabricksTracingConfigGatewayGatewayenvironment, - ) -> None: - workspace = str(uuid4()) - encrypted = OpsTraceManager.encrypt_tracing_config( - workspace, - "databricks", - { - "host": "https://user:password@workspace.test", - "experiment_id": "12", - "personal_access_token": "synthetic-pat", - }, - ) - result = environment.gateway.present_config( - workspace_id=workspace, tracing_provider="databricks", tracing_config=encrypted - ) - assert result["project_url"] == "https://www.databricks.com/" - assert environment.api.calls == [] - - @staticmethod - def test_failed_patch_keeps_original_destination_and_logs_only_safe_code( - environment: _DatabricksTracingConfigGatewayGatewayenvironment, - sqlite_session_factory: sessionmaker[Session], - caplog: pytest.LogCaptureFixture, - ) -> None: - workspace, app_id = str(uuid4()), str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add(make_app(app_id=app_id, tenant_id=workspace)) - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - service = AppTracingConfigService(configs=store, provider=environment.gateway) - context = RequestContext("request", None, "account", workspace) - original = { - "host": "https://workspace.test", - "experiment_id": "12", - "personal_access_token": "synthetic-private-pat", - } - service.create(context, app_id, "databricks", original) - before = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="databricks") - assert before is not None - environment.api.override = lambda _method, url, _options: _databricks_tracing_config_gateway_response( - url, {"message": "private-vendor-detail"}, 403 - ) - with pytest.raises(AppTracingConfigVerificationFailedError) as caught: - service.update( - context, app_id, "databricks", {**original, "experiment_id": "34", "personal_access_token": "******"} - ) - assert isinstance(caught.value.__cause__, DatabricksConfigurationError) - assert caught.value.__cause__.code == "permission_denied" - after = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="databricks") - assert after is not None - assert after.tracing_config == before.tracing_config - assert "permission_denied" in caplog.text - assert "private-vendor-detail" not in caplog.text - assert "synthetic-private-pat" not in caplog.text - - @staticmethod - @pytest.mark.parametrize("start_with_oauth", [False, True]) - def test_switching_auth_clears_unselected_credentials_in_persistence_and_requests( - environment: _DatabricksTracingConfigGatewayGatewayenvironment, - sqlite_session_factory: sessionmaker[Session], - start_with_oauth: bool, - ) -> None: - workspace, app_id = str(uuid4()), str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add(make_app(app_id=app_id, tenant_id=workspace)) - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - service = AppTracingConfigService(configs=store, provider=environment.gateway) - context = RequestContext("request", None, "account", workspace) - pat = {"personal_access_token": "synthetic-pat", "client_id": None, "client_secret": None} - oauth = {"personal_access_token": None, "client_id": "synthetic-client", "client_secret": "synthetic-secret"} - target = {"host": "https://workspace.test", "experiment_id": "12"} - service.create(context, app_id, "databricks", {**target, **(oauth if start_with_oauth else pat)}) - environment.api.calls.clear() - service.update(context, app_id, "databricks", {**target, **(pat if start_with_oauth else oauth)}) - stored = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="databricks") - assert stored is not None - assert stored.tracing_config is not None - assert stored.tracing_config["client_secret" if start_with_oauth else "personal_access_token"] is None - assert [method for method, _url, _options in environment.api.calls] == ( - ["GET"] if start_with_oauth else ["POST", "GET"] - ) - expected_token = "synthetic-pat" if start_with_oauth else "synthetic-oauth-token" - assert environment.api.calls[-1][2]["headers"] == { - "Authorization": f"Bearer {expected_token}", - "Accept": "application/json", - } - presented = service.get(context, app_id, "databricks") - assert presented is not None - assert presented.tracing_config is not None - assert presented.tracing_config["client_secret" if start_with_oauth else "personal_access_token"] is None - - -"""MLflow config persistence uses real RSA/SQLite without constructing an exporter.""" - - -ExperimentAPI = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-mlflow.src.dify_trace_mlflow.test_configuration" -).ExperimentAPI -_mlflow_tracing_config_gateway_response = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-mlflow.src.dify_trace_mlflow.test_configuration" -).response - - -@dataclass -class _MlflowTracingConfigGatewayConfigurationcase: - api: ExperimentAPI - gateway: OpsTraceManagerGateway - service: AppTracingConfigService - store: SQLAlchemyAppTracingConfigRepository - reader: ConsoleAppRepository - context: RequestContext - app_id: str - - def read(self, provider: str = "mlflow") -> AppTracingConfigRecord | None: - return self.store.get( - workspace_id=self.context.active_workspace_id, app_id=self.app_id, tracing_provider=provider - ) - - -class TestMlflowTracingConfigGatewayCases: - import importlib as _test_importlib - - experiment_api = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-mlflow.src.dify_trace_mlflow.test_oss_configuration" - ).experiment_api - - experiment_api = staticmethod(experiment_api) - - import importlib as _test_importlib - - isolated_sdk_state = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-mlflow.src.dify_trace_mlflow.test_oss_configuration" - ).isolated_sdk_state - - isolated_sdk_state = staticmethod(isolated_sdk_state) - - @staticmethod - @pytest.fixture - def case( - workspace: str, sqlite_session_factory: sessionmaker[Session], experiment_api: ExperimentAPI - ) -> _MlflowTracingConfigGatewayConfigurationcase: - app_id = str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add( - make_app( - app_id=app_id, - tenant_id=workspace, - tracing=json.dumps( - { - "enabled": True, - "tracing_provider": "langsmith", - "capture_content": True, - "capture_agent": False, - } - ), - ) - ) - gateway = OpsTraceManagerGateway() - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - previous = OpsTraceManager.encrypt_tracing_config( - workspace, - "langsmith", - {"api_key": "synthetic-other-key", "project": "other", "endpoint": "https://other.test"}, - ) - assert store.create( - workspace_id=workspace, app_id=app_id, tracing_provider="langsmith", tracing_config=previous - ) - return _MlflowTracingConfigGatewayConfigurationcase( - experiment_api, - gateway, - AppTracingConfigService(configs=store, provider=gateway), - store, - ConsoleAppRepository(session_factory=sqlite_session_factory), - RequestContext("request", None, "account", workspace), - app_id, - ) - - @staticmethod - def test_post_masked_patch_and_get_preserve_ciphertext_and_exact_destination( - case: _MlflowTracingConfigGatewayConfigurationcase, - ) -> None: - original = case.read("langsmith") - submitted = { - "tracking_uri": " HTTP://FIRST.Example:5000/team/ ", - "experiment_id": "00012", - "username": "test-user", - "password": "synthetic-password", - } - untouched = dict(submitted) - case.service.create(case.context, case.app_id, "mlflow", submitted, activate=True) - before = case.read() - assert before is not None - assert before.tracing_config is not None - config = before.tracing_config - assert config["tracking_uri"] == "http://first.example:5000/team" - assert config["experiment_id"] == "12" - assert config["password"] != "synthetic-password" - assert OpsTraceManager.decrypt_tracing_config(case.context.active_workspace_id, "mlflow", config)[ - "password" - ] == ("synthetic-password") - assert submitted == untouched - shown = case.service.get(case.context, case.app_id, "mlflow") - assert shown is not None - assert shown.tracing_config is not None - masked = shown.tracing_config["password"] - assert isinstance(masked, str) - assert "*" in masked - assert masked != "synthetic-password" - assert len(case.api.calls) == 1 - case.service.update( - case.context, - case.app_id, - "mlflow", - { - "tracking_uri": "https://SECOND.test/new-prefix/", - "experiment_id": "00034", - "username": "test-user", - "password": masked, - }, - ) - after = case.read() - assert after is not None - assert after.tracing_config is not None - assert after.tracing_config["password"] == config["password"] - assert after.tracing_config["tracking_uri"] == "https://second.test/new-prefix" - assert after.tracing_config["experiment_id"] == "34" - assert after.tracing_config["project_url"] == "https://second.test/new-prefix/#/experiments/34/traces" - assert len(case.api.calls) == 2 - assert case.api.calls[-1][1] == "https://second.test/new-prefix/api/2.0/mlflow/experiments/get" - case.gateway.validate_saved_config( - workspace_id=case.context.active_workspace_id, - tracing_provider="mlflow", - tracing_config=after.tracing_config, - ) - assert case.service.get(case.context, case.app_id, "mlflow") is not None - assert len(case.api.calls) == 2 - assert case.read("langsmith") == original - policy = case.reader.get_trace(case.context, case.app_id) - assert policy.enabled - assert policy.tracing_provider == "mlflow" - assert policy.capture_content is True - assert policy.capture_agent is False - - @staticmethod - @pytest.mark.parametrize("clear_with_null", [False, True]) - def test_basic_to_anonymous_and_back_does_not_reuse_credentials( - case: _MlflowTracingConfigGatewayConfigurationcase, clear_with_null: bool - ) -> None: - destination = {"tracking_uri": "https://same.test/prefix", "experiment_id": "12"} - case.service.create( - case.context, case.app_id, "mlflow", {**destination, "username": "first", "password": "first-secret"} - ) - empty = None if clear_with_null else "" - case.service.update(case.context, case.app_id, "mlflow", {**destination, "username": empty, "password": empty}) - assert case.api.calls[-1][2]["headers"] == { - "Accept": "application/json", - "Accept-Encoding": "identity", - "Cookie": "", - } - stored = case.read() - assert stored is not None - assert stored.tracing_config is not None - assert ( - OpsTraceManager.decrypt_tracing_config(case.context.active_workspace_id, "mlflow", stored.tracing_config)[ - "password" - ] - == empty - ) - case.service.update( - case.context, case.app_id, "mlflow", {**destination, "username": "second", "password": "second-secret"} - ) - assert case.api.calls[-1][2]["headers"] != case.api.calls[0][2]["headers"] - - @staticmethod - def test_failed_patch_keeps_config_policy_and_other_provider( - case: _MlflowTracingConfigGatewayConfigurationcase, caplog: pytest.LogCaptureFixture - ) -> None: - original = { - "tracking_uri": "https://first.test", - "experiment_id": "12", - "username": "user", - "password": "private-password", - } - case.service.create(case.context, case.app_id, "mlflow", original) - before, other, policy = case.read(), case.read("langsmith"), case.reader.get_trace(case.context, case.app_id) - case.api.override = lambda _method, url, _options: _mlflow_tracing_config_gateway_response( - url, {"message": "private-vendor-response"}, 403 - ) - with pytest.raises(AppTracingConfigVerificationFailedError) as caught: - case.service.update( - case.context, - case.app_id, - "mlflow", - {**original, "tracking_uri": "https://second.test", "password": "******"}, - activate=True, - ) - assert isinstance(caught.value.__cause__, MLflowConfigurationError) - assert caught.value.__cause__.code == "permission_denied" - assert case.read() == before - assert case.read("langsmith") == other - assert case.reader.get_trace(case.context, case.app_id) == policy - assert "permission_denied" in caplog.text - assert "private-vendor-response" not in caplog.text - assert "private-password" not in caplog.text - - @staticmethod - def test_failed_auth_pair_patch_does_not_erase_a_password( - case: _MlflowTracingConfigGatewayConfigurationcase, - ) -> None: - original = { - "tracking_uri": "https://first.test", - "experiment_id": "12", - "username": "user", - "password": "password", - } - case.service.create(case.context, case.app_id, "mlflow", original) - before = case.read() - with pytest.raises(AppTracingConfigVerificationFailedError) as caught: - case.service.update(case.context, case.app_id, "mlflow", {**original, "password": None}) - assert isinstance(caught.value.__cause__, MLflowConfigurationError) - assert caught.value.__cause__.code == "incomplete_basic_auth" - assert len(case.api.calls) == 1 - assert case.read() == before - - @staticmethod - @pytest.mark.parametrize("saved_url", [None, "https://stale.test/wrong-experiment"]) - def test_historical_get_is_pure_even_with_partial_auth( - case: _MlflowTracingConfigGatewayConfigurationcase, saved_url: str | None - ) -> None: - stored = OpsTraceManager.encrypt_tracing_config( - case.context.active_workspace_id, - "mlflow", - { - "tracking_uri": "https://server.test/prefix/", - "experiment_id": "00012", - "username": "historical-user", - "password": None, - }, - ) - if saved_url is not None: - stored["project_url"] = saved_url - shown = case.gateway.present_config( - workspace_id=case.context.active_workspace_id, tracing_provider="mlflow", tracing_config=stored - ) - assert shown["project_url"] == "https://server.test/prefix/#/experiments/12/traces" - assert shown["password"] is None - case.gateway.validate_saved_config( - workspace_id=case.context.active_workspace_id, tracing_provider="mlflow", tracing_config=stored - ) - assert case.api.calls == [] - - @staticmethod - def test_invalid_historical_url_has_no_fabricated_link_or_network( - case: _MlflowTracingConfigGatewayConfigurationcase, - ) -> None: - stored = OpsTraceManager.encrypt_tracing_config( - case.context.active_workspace_id, - "mlflow", - { - "tracking_uri": "https://user:secret@server.test?token=private", - "experiment_id": "12", - "username": None, - "password": None, - }, - ) - shown = case.gateway.present_config( - workspace_id=case.context.active_workspace_id, tracing_provider="mlflow", tracing_config=stored - ) - assert shown["project_url"] is None - assert case.api.calls == [] - - @staticmethod - def test_invalid_new_destination_never_reaches_storage_or_network( - case: _MlflowTracingConfigGatewayConfigurationcase, - ) -> None: - with pytest.raises(AppTracingConfigInvalidConfigurationError): - case.service.create( - case.context, - case.app_id, - "mlflow", - { - "tracking_uri": "https://server.test/prefix/../other", - "experiment_id": "12", - "username": "user", - "password": "secret", - }, - activate=True, - ) - assert case.read() is None - assert case.api.calls == [] - assert case.reader.get_trace(case.context, case.app_id).tracing_provider == "langsmith" - - -"""Opik configuration transactions use real RSA and never initialize an SDK streamer.""" - - -AccessAPI = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-opik.src.dify_trace_opik.test_configuration" -).AccessAPI - - -@dataclass -class _OpikTracingConfigGatewayConfigurationcase: - api: AccessAPI - gateway: OpsTraceManagerGateway - service: AppTracingConfigService - store: SQLAlchemyAppTracingConfigRepository - reader: ConsoleAppRepository - context: RequestContext - app_id: str - sessions: sessionmaker[Session] - - -class TestOpikTracingConfigGatewayCases: - @staticmethod - @pytest.fixture - def case( - workspace: str, sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch - ) -> Iterator[_OpikTracingConfigGatewayConfigurationcase]: - app_id = str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add(make_app(app_id=app_id, tenant_id=workspace)) - api = AccessAPI() - gateway = OpsTraceManagerGateway() - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - forbidden_calls: list[str] = [] - - def forbidden(*_args: object, **_kwargs: object) -> Never: - forbidden_calls.append("sdk_constructor") - raise AssertionError("Configuration must not initialize an SDK client or ambient configuration") - - monkeypatch.setattr(ssrf_proxy, "make_request", api.request) - monkeypatch.setattr(OpikDataTrace, "__init__", forbidden) - monkeypatch.setattr(opik.Opik, "__init__", forbidden) - monkeypatch.setattr(SDKOpikConfig, "__init__", forbidden) - yield _OpikTracingConfigGatewayConfigurationcase( - api, - gateway, - AppTracingConfigService(configs=store, provider=gateway), - store, - ConsoleAppRepository(session_factory=sqlite_session_factory), - RequestContext("request", None, "account", workspace), - app_id, - sqlite_session_factory, - ) - assert forbidden_calls == [] - - @staticmethod - def test_post_get_patch_preserve_real_ciphertext_and_verified_destination( - case: _OpikTracingConfigGatewayConfigurationcase, - ) -> None: - secret = "synthetic-private-opik-key" - case.service.create( - case.context, - case.app_id, - "opik", - {"api_key": secret, "workspace": "", "project": " first ", "url": "https://first.test/prefix/api/"}, - activate=True, - ) - before = case.store.get( - workspace_id=case.context.active_workspace_id, app_id=case.app_id, tracing_provider="opik" - ) - assert before is not None - assert before.tracing_config is not None - config = before.tracing_config - assert config["api_key"] != secret - assert config["workspace"] == "resolved-team" - assert config["project"] == "first" - assert config["project_url"] == "https://first.test/prefix/resolved-team/redirect/projects?name=first" - assert ( - OpsTraceManager.decrypt_tracing_config(case.context.active_workspace_id, "opik", config)["api_key"] - == secret - ) - assert case.reader.get_trace(case.context, case.app_id).tracing_provider == "opik" - - before_get = len(case.api.calls) - presented = case.service.get(case.context, case.app_id, "opik") - assert presented is not None - assert presented.tracing_config is not None - masked = presented.tracing_config["api_key"] - assert masked != secret - assert "*" in masked - assert len(case.api.calls) == before_get - case.api.default_workspace = "new-team" - case.service.update( - case.context, - case.app_id, - "opik", - {"api_key": masked, "workspace": "default", "project": "A & B", "url": "https://second.test/api/"}, - ) - after = case.store.get( - workspace_id=case.context.active_workspace_id, app_id=case.app_id, tracing_provider="opik" - ) - assert after is not None - assert after.tracing_config is not None - assert after.tracing_config["api_key"] == config["api_key"] - assert after.tracing_config["workspace"] == "new-team" - assert after.tracing_config["project"] == "A & B" - assert after.tracing_config["url"] == "https://second.test/api/" - assert after.tracing_config["project_url"] == "https://second.test/new-team/redirect/projects?name=A+%26+B" - assert case.api.calls[-1][2]["headers"] == { - "Authorization": secret, - "Comet-Workspace": "default", - "Accept": "application/json", - } - before_get = len(case.api.calls) - presented = case.service.get(case.context, case.app_id, "opik") - assert presented is not None - assert presented.tracing_config is not None - assert presented.tracing_config["project_url"] == after.tracing_config["project_url"] - case.gateway.validate_saved_config( - workspace_id=case.context.active_workspace_id, tracing_provider="opik", tracing_config=after.tracing_config - ) - assert len(case.api.calls) == before_get - - @staticmethod - @pytest.mark.parametrize("saved_url", [None, "https://wrong-global.test/old-project"]) - @pytest.mark.parametrize("workspace_name", ["team", "default", None]) - def test_historical_get_uses_only_configuration_host_without_network( - case: _OpikTracingConfigGatewayConfigurationcase, saved_url: str | None, workspace_name: str | None - ) -> None: - stored = OpsTraceManager.encrypt_tracing_config( - case.context.active_workspace_id, - "opik", - { - "api_key": "synthetic-key", - "workspace": workspace_name, - "project": "A & B", - "url": "https://private.test/opik/api/", - }, - ) - if saved_url: - stored["project_url"] = saved_url - result = case.gateway.present_config( - workspace_id=case.context.active_workspace_id, tracing_provider="opik", tracing_config=stored - ) - expected = ( - "https://private.test/opik/team/redirect/projects?name=A+%26+B" - if workspace_name == "team" - else "https://private.test/opik/" - ) - assert result["project_url"] == expected - assert result["api_key"] != "synthetic-key" - assert case.api.calls == [] - - @staticmethod - @pytest.mark.parametrize("status", [401, 403]) - @pytest.mark.parametrize("existing", [False, True]) - def test_bad_key_or_workspace_does_not_save_or_change_active_provider( - case: _OpikTracingConfigGatewayConfigurationcase, status: int, existing: bool, caplog: pytest.LogCaptureFixture - ) -> None: - original = {"api_key": "synthetic-private-key", "workspace": "team", "project": "original"} - if existing: - case.service.create(case.context, case.app_id, "opik", original) - with case.sessions.begin() as session: - app = session.get(App, case.app_id) - assert app is not None - app.tracing = json.dumps({"enabled": True, "tracing_provider": "langsmith", "capture_content": False}) - before_policy = case.reader.get_trace(case.context, case.app_id) - before = case.store.get( - workspace_id=case.context.active_workspace_id, app_id=case.app_id, tracing_provider="opik" - ) - case.api.override = lambda method, url, _options: httpx.Response( - status, text="private-vendor-details", request=httpx.Request(method, url) - ) - operation = case.service.update if existing else case.service.create - with pytest.raises(AppTracingConfigVerificationFailedError) as error: - operation(case.context, case.app_id, "opik", {**original, "workspace": "inaccessible"}, activate=True) - assert isinstance(error.value.__cause__, OpikConfigurationError) - assert error.value.__cause__.code == ("authentication_failed" if status == 401 else "permission_denied") - after = case.store.get( - workspace_id=case.context.active_workspace_id, app_id=case.app_id, tracing_provider="opik" - ) - assert after == before - assert case.reader.get_trace(case.context, case.app_id) == before_policy - assert "private-vendor-details" not in caplog.text - assert "synthetic-private-key" not in caplog.text - - @staticmethod - def test_clear_optional_api_key_keeps_keyless_self_host_valid( - case: _OpikTracingConfigGatewayConfigurationcase, monkeypatch: pytest.MonkeyPatch - ) -> None: - target = {"url": "http://local.test:5173/api/", "workspace": "team", "project": "sample"} - case.service.create(case.context, case.app_id, "opik", {**target, "api_key": "synthetic-key"}) - monkeypatch.setenv("OPIK_API_KEY", "unrelated-key") - case.api.calls.clear() - case.service.update(case.context, case.app_id, "opik", {**target, "api_key": None}) - stored = case.store.get( - workspace_id=case.context.active_workspace_id, app_id=case.app_id, tracing_provider="opik" - ) - assert stored is not None - assert stored.tracing_config is not None - assert stored.tracing_config["api_key"] is None - presented = case.service.get(case.context, case.app_id, "opik") - assert presented is not None - assert presented.tracing_config is not None - assert presented.tracing_config["api_key"] is None - assert case.api.calls[0][2]["headers"] == {"Comet-Workspace": "team", "Accept": "application/json"} - assert len(case.api.calls) == 1 - - -"""Persist verified Phoenix configuration while preserving owner policy and masked keys.""" - - -_phoenix_tracing_config_gateway_response = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_configuration" -).response - - -@dataclass -class PhoenixCase: - gateway: OpsTraceManagerGateway - service: AppTracingConfigService - store: SQLAlchemyAppTracingConfigRepository - reader: ConsoleAppRepository - context: RequestContext - app_id: str - - def read(self, provider: str = "phoenix") -> AppTracingConfigRecord | None: - return self.store.get( - workspace_id=self.context.active_workspace_id, app_id=self.app_id, tracing_provider=provider - ) - - -class TestPhoenixTracingConfigGatewayCases: - import importlib as _test_importlib - - api_calls = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_configuration" - ).api_calls - - api_calls = staticmethod(api_calls) - - import importlib as _test_importlib - - api_responses = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_configuration" - ).api_responses - - api_responses = staticmethod(api_responses) - - import importlib as _test_importlib - - no_legacy_or_environment_auth = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_configuration" - ).no_legacy_or_environment_auth - - no_legacy_or_environment_auth = staticmethod(no_legacy_or_environment_auth) - - @staticmethod - @pytest.fixture - def case(workspace: str, sqlite_session_factory: sessionmaker[Session]) -> PhoenixCase: - app_id = str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add( - make_app( - app_id=app_id, - tenant_id=workspace, - tracing=json.dumps( - { - "enabled": True, - "tracing_provider": "langsmith", - "capture_content": True, - "capture_agent": False, - } - ), - ) - ) - gateway = OpsTraceManagerGateway() - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - old = OpsTraceManager.encrypt_tracing_config( - workspace, - "langsmith", - {"api_key": "synthetic-other-key", "project": "other", "endpoint": "https://other.test"}, - ) - assert store.create(workspace_id=workspace, app_id=app_id, tracing_provider="langsmith", tracing_config=old) - return PhoenixCase( - gateway, - AppTracingConfigService(configs=store, provider=gateway), - store, - ConsoleAppRepository(session_factory=sqlite_session_factory), - RequestContext("request", None, "account", workspace), - app_id, - ) - - @staticmethod - def test_post_masked_patch_get_keep_normalized_destination_and_ciphertext( - case: PhoenixCase, - api_calls: list[tuple[str, dict[str, object]]], - ) -> None: - other = case.read("langsmith") - submitted = { - "endpoint": " HTTP://FIRST.Example:6006/reverse/ ", - "project": " first project ", - "api_key": "synthetic-secret", - } - untouched = dict(submitted) - case.service.create(case.context, case.app_id, "phoenix", submitted, activate=True) - before = case.read() - assert before is not None - assert before.tracing_config is not None - stored = before.tracing_config - assert stored["endpoint"] == "http://first.example:6006/reverse" - assert stored["project"] == "first project" - assert stored["api_key"] != "synthetic-secret" - assert submitted == untouched - assert OpsTraceManager.decrypt_tracing_config(case.context.active_workspace_id, "phoenix", stored)[ - "api_key" - ] == ("synthetic-secret") - shown = case.service.get(case.context, case.app_id, "phoenix") - assert shown is not None - assert shown.tracing_config is not None - masked = shown.tracing_config["api_key"] - assert isinstance(masked, str) - assert "*" in masked - assert masked != "synthetic-secret" - assert ( - shown.tracing_config["project_url"] - == "http://first.example:6006/reverse/redirects/projects/first%20project" - ) - assert len(api_calls) == 1 - case.service.update( - case.context, - case.app_id, - "phoenix", - { - "endpoint": " https://SECOND.test/s/space/ ", - "project": " second & 项目 ", - "api_key": masked, - }, - ) - after = case.read() - assert after is not None - assert after.tracing_config is not None - assert after.tracing_config["api_key"] == stored["api_key"] - assert after.tracing_config["endpoint"] == "https://second.test/s/space" - assert after.tracing_config["project"] == "second & 项目" - assert after.tracing_config["project_url"] == ( - "https://second.test/s/space/redirects/projects/second%20%26%20%E9%A1%B9%E7%9B%AE" - ) - assert api_calls[-1][0] == "https://second.test/s/space/v1/projects" - headers = api_calls[-1][1]["headers"] - assert isinstance(headers, dict) - assert headers["Authorization"] == "Bearer synthetic-secret" - case.gateway.validate_saved_config( - workspace_id=case.context.active_workspace_id, - tracing_provider="phoenix", - tracing_config=after.tracing_config, - ) - assert case.service.get(case.context, case.app_id, "phoenix") is not None - assert len(api_calls) == 2 - assert case.read("langsmith") == other - policy = case.reader.get_trace(case.context, case.app_id) - assert policy.enabled - assert policy.tracing_provider == "phoenix" - assert policy.capture_content is True - assert policy.capture_agent is False - - @staticmethod - @pytest.mark.parametrize("empty", [None, ""]) - def test_bearer_to_keyless_and_back_is_explicit( - case: PhoenixCase, - api_calls: list[tuple[str, dict[str, object]]], - empty: str | None, - ) -> None: - destination = {"endpoint": "https://host.test", "project": "project"} - case.service.create(case.context, case.app_id, "phoenix", {**destination, "api_key": "first-key"}) - case.service.update(case.context, case.app_id, "phoenix", {**destination, "api_key": empty}) - assert api_calls[-1][1]["headers"] == { - "Accept": "application/json", - "Accept-Encoding": "identity", - "Cookie": "", - } - stored = case.read() - assert stored is not None - assert stored.tracing_config is not None - assert ( - OpsTraceManager.decrypt_tracing_config(case.context.active_workspace_id, "phoenix", stored.tracing_config)[ - "api_key" - ] - == empty - ) - case.service.update(case.context, case.app_id, "phoenix", {**destination, "api_key": "second-key"}) - assert api_calls[-1][1]["headers"] != api_calls[0][1]["headers"] - assert len(api_calls) == 3 - - @staticmethod - @pytest.mark.parametrize("failure", ["denied", "wrong-path-html"]) - def test_failed_patch_preserves_policy_config_and_other_provider( - case: PhoenixCase, - api_calls: list[tuple[str, dict[str, object]]], - api_responses: deque[httpx.Response | Exception], - caplog: pytest.LogCaptureFixture, - failure: str, - ) -> None: - original = {"endpoint": "https://host.test", "project": "project", "api_key": "private-secret"} - case.service.create(case.context, case.app_id, "phoenix", original) - before, other, policy = case.read(), case.read("langsmith"), case.reader.get_trace(case.context, case.app_id) - if failure == "denied": - api_responses.append(_phoenix_tracing_config_gateway_response({"secret": "private-vendor-response"}, 403)) - else: - api_responses.append( - httpx.Response( - 200, - text="private-vendor-response", - request=httpx.Request("GET", "https://host.test/not-a-collector/v1/projects"), - ) - ) - with pytest.raises(AppTracingConfigVerificationFailedError) as caught: - case.service.update( - case.context, - case.app_id, - "phoenix", - { - **original, - "endpoint": "https://host.test/not-a-collector", - "api_key": "******", - }, - activate=True, - ) - assert isinstance(caught.value.__cause__, PhoenixConfigurationError) - assert caught.value.__cause__.code == ("permission_denied" if failure == "denied" else "malformed_response") - assert case.read() == before - assert case.read("langsmith") == other - assert case.reader.get_trace(case.context, case.app_id) == policy - assert "private-vendor-response" not in caplog.text - assert "private-secret" not in caplog.text - assert len(api_calls) == 2 - - @staticmethod - @pytest.mark.parametrize("saved_url", [None, "https://stale.test/wrong-project"]) - def test_historical_get_calculates_url_without_verification_or_sdk( - case: PhoenixCase, - api_calls: list[tuple[str, dict[str, object]]], - saved_url: str | None, - ) -> None: - stored = OpsTraceManager.encrypt_tracing_config( - case.context.active_workspace_id, - "phoenix", - { - "endpoint": "https://host.test/s/space/", - "project": "historical project", - "api_key": None, - }, - ) - if saved_url: - stored["project_url"] = saved_url - shown = case.gateway.present_config( - workspace_id=case.context.active_workspace_id, tracing_provider="phoenix", tracing_config=stored - ) - assert shown["project_url"] == "https://host.test/s/space/redirects/projects/historical%20project" - assert shown["api_key"] is None - case.gateway.validate_saved_config( - workspace_id=case.context.active_workspace_id, tracing_provider="phoenix", tracing_config=stored - ) - assert api_calls == [] - - @staticmethod - @pytest.mark.parametrize( - "endpoint", ["https://user:private@host.test?token=private", "https://host.test/v1/traces"] - ) - def test_unsafe_historical_destination_is_readable_without_link_or_io( - case: PhoenixCase, - api_calls: list[tuple[str, dict[str, object]]], - endpoint: str, - ) -> None: - stored = OpsTraceManager.encrypt_tracing_config( - case.context.active_workspace_id, - "phoenix", - { - "endpoint": endpoint, - "project": "old-project", - "api_key": None, - }, - ) - shown = case.gateway.present_config( - workspace_id=case.context.active_workspace_id, tracing_provider="phoenix", tracing_config=stored - ) - assert shown["endpoint"] == endpoint - assert shown["project_url"] is None - assert api_calls == [] - - @staticmethod - @pytest.mark.parametrize("values", [{"endpoint": "https://host.test/v1/traces"}, {"project": "bad\nproject"}]) - def test_invalid_new_destination_has_no_io_or_storage_effect( - case: PhoenixCase, - api_calls: list[tuple[str, dict[str, object]]], - values: dict[str, str], - ) -> None: - policy = case.reader.get_trace(case.context, case.app_id) - with pytest.raises(AppTracingConfigInvalidConfigurationError): - case.service.create( - case.context, case.app_id, "phoenix", {"api_key": "synthetic-key", **values}, activate=True - ) - assert case.read() is None - assert case.reader.get_trace(case.context, case.app_id) == policy - assert api_calls == [] - - -"""Tencent configuration persistence uses real RSA and bounded protocol checks.""" - - -CollectorBoundary = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-tencent.src.dify_trace_tencent.test_configuration" -).CollectorBoundary - - -@dataclass -class _TencentTracingConfigGatewayConfigurationcase: - collector: CollectorBoundary - gateway: OpsTraceManagerGateway - service: AppTracingConfigService - store: SQLAlchemyAppTracingConfigRepository - reader: ConsoleAppRepository - context: RequestContext - app_id: str - - def read(self, provider: str = "tencent") -> AppTracingConfigRecord | None: - return self.store.get( - workspace_id=self.context.active_workspace_id, app_id=self.app_id, tracing_provider=provider - ) - - -class TestTencentTracingConfigGatewayCases: - import importlib as _test_importlib - - _forbid_provider_io = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-tencent.src.dify_trace_tencent.test_configuration" - )._forbid_provider_io - - _forbid_provider_io = staticmethod(_forbid_provider_io) - - import importlib as _test_importlib - - collector = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-tencent.src.dify_trace_tencent.test_configuration" - ).collector - - collector = staticmethod(collector) - - @staticmethod - @pytest.fixture - def case( - workspace: str, sqlite_session_factory: sessionmaker[Session], collector: CollectorBoundary - ) -> _TencentTracingConfigGatewayConfigurationcase: - app_id = str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add( - make_app( - app_id=app_id, - tenant_id=workspace, - tracing=json.dumps( - { - "enabled": True, - "tracing_provider": "langsmith", - "capture_content": True, - "capture_agent": False, - } - ), - ) - ) - gateway = OpsTraceManagerGateway() - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - old_config = OpsTraceManager.encrypt_tracing_config( - workspace, - "langsmith", - { - "api_key": "synthetic-existing-key", - "project": "existing-project", - "endpoint": "https://existing.example", - }, - ) - assert store.create( - workspace_id=workspace, app_id=app_id, tracing_provider="langsmith", tracing_config=old_config - ) - return _TencentTracingConfigGatewayConfigurationcase( - collector, - gateway, - AppTracingConfigService(configs=store, provider=gateway), - store, - ConsoleAppRepository(session_factory=sqlite_session_factory), - RequestContext("request", None, "account", workspace), - app_id, - ) - - @staticmethod - def test_post_get_masked_patch_persist_normalization_and_preserve_real_ciphertext( - case: _TencentTracingConfigGatewayConfigurationcase, - ) -> None: - secret = "synthetic-private-tencent-token" - old_provider = case.read("langsmith") - submitted = {"token": secret, "endpoint": " HTTPS://FIRST.Example:4320/ ", "service_name": " first-app "} - original_submission = dict(submitted) - - case.service.create(case.context, case.app_id, "tencent", submitted, activate=True) - - before = case.read() - assert before is not None - assert before.tracing_config is not None - config = before.tracing_config - assert config["token"] != secret - assert config["endpoint"] == "https://first.example:4320" - assert config["service_name"] == "first-app" - assert config["project_url"] == "https://console.cloud.tencent.com/apm" - assert ( - OpsTraceManager.decrypt_tracing_config(case.context.active_workspace_id, "tencent", config)["token"] - == secret - ) - assert submitted == original_submission - assert len(case.collector.channels) == 1 - assert case.collector.channels[0].calls[0].metadata == (("authorization", f"Bearer {secret}"),) - policy = case.reader.get_trace(case.context, case.app_id) - assert policy.enabled is True - assert policy.tracing_provider == "tencent" - assert policy.capture_content is True - assert policy.capture_agent is False - assert case.read("langsmith") == old_provider - - presented = case.service.get(case.context, case.app_id, "tencent") - assert presented is not None - assert presented.tracing_config is not None - masked = presented.tracing_config["token"] - assert isinstance(masked, str) - assert masked != secret - assert "*" in masked - assert presented.tracing_config["project_url"] == "https://console.cloud.tencent.com/apm" - assert len(case.collector.channels) == 1 - - replacement = {"token": masked, "endpoint": " http://SECOND.Example:4319/ ", "service_name": " second-app "} - original_replacement = dict(replacement) - case.service.update(case.context, case.app_id, "tencent", replacement, activate=True) - - after = case.read() - assert after is not None - assert after.tracing_config is not None - assert after.tracing_config["token"] == config["token"] - assert after.tracing_config["endpoint"] == "http://second.example:4319" - assert after.tracing_config["service_name"] == "second-app" - assert after.tracing_config["project_url"] == "https://console.cloud.tencent.com/apm" - assert replacement == original_replacement - assert len(case.collector.channels) == 2 - channel = case.collector.channels[-1] - assert channel.target == "second.example:4319" - assert channel.secure is False - assert channel.calls[0].metadata == (("authorization", f"Bearer {secret}"),) - assert all(channel.close_count == 1 for channel in case.collector.channels) - case.gateway.validate_saved_config( - workspace_id=case.context.active_workspace_id, - tracing_provider="tencent", - tracing_config=after.tracing_config, - ) - assert case.service.get(case.context, case.app_id, "tencent") is not None - assert len(case.collector.channels) == 2 - assert case.read("langsmith") == old_provider - - @staticmethod - @pytest.mark.parametrize("saved_url", [None, "https://wrong.example/previous-project"]) - def test_historical_name_remains_readable_and_saved_validation_does_not_connect( - case: _TencentTracingConfigGatewayConfigurationcase, saved_url: str | None - ) -> None: - stored = OpsTraceManager.encrypt_tracing_config( - case.context.active_workspace_id, - "tencent", - { - "token": "synthetic-historical-token", - "endpoint": "https://legacy.example:4320", - "service_name": "dify_app", - }, - ) - if saved_url is not None: - stored["project_url"] = saved_url - assert case.store.create( - workspace_id=case.context.active_workspace_id, - app_id=case.app_id, - tracing_provider="tencent", - tracing_config=stored, - ) - before = case.read() - before_policy = case.reader.get_trace(case.context, case.app_id) - - presented = case.service.get(case.context, case.app_id, "tencent") - case.gateway.validate_saved_config( - workspace_id=case.context.active_workspace_id, tracing_provider="tencent", tracing_config=stored - ) - - assert presented is not None - assert presented.tracing_config is not None - assert presented.tracing_config["service_name"] == "dify_app" - assert presented.tracing_config["project_url"] == "https://console.cloud.tencent.com/apm" - assert presented.tracing_config["token"] != "synthetic-historical-token" - assert "*" in presented.tracing_config["token"] - assert case.collector.channels == [] - assert case.read() == before - assert case.reader.get_trace(case.context, case.app_id) == before_policy - - @staticmethod - @pytest.mark.parametrize("existing", [False, True]) - @pytest.mark.parametrize( - ("field_name", "invalid_value", "error_code"), - [ - ("service_name", "dify_app", "invalid_service_name"), - ("service_name", "", "invalid_service_name"), - ("endpoint", "https://collector.example:4320/v1/traces", "invalid_endpoint"), - ("endpoint", "https://collector.example:4320?token=private-endpoint-detail", "invalid_endpoint"), - ("endpoint", "https://", "invalid_endpoint"), - ], - ) - def test_invalid_create_or_patch_rejects_before_rpc_encryption_or_persistence( - case: _TencentTracingConfigGatewayConfigurationcase, - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - existing: bool, - field_name: str, - invalid_value: str, - error_code: str, - ) -> None: - original = { - "token": "synthetic-private-token", - "endpoint": "https://original.example:4320", - "service_name": "original-app", - } - if existing: - case.service.create(case.context, case.app_id, "tencent", original) - case.collector.channels.clear() - before = case.read() - old_provider = case.read("langsmith") - before_policy = case.reader.get_trace(case.context, case.app_id) - forbidden_attempts: list[str] = [] - - def forbidden(*_args: object, **_kwargs: object) -> Never: - forbidden_attempts.append("encryption_or_write") - raise AssertionError("Invalid new input must fail before encryption or persistence") - - monkeypatch.setattr(OpsTraceManager, "encrypt_tracing_config", forbidden) - monkeypatch.setattr(case.store, "create", forbidden) - monkeypatch.setattr(case.store, "update", forbidden) - operation = case.service.update if existing else case.service.create - with pytest.raises(AppTracingConfigInvalidConfigurationError) as caught: - operation(case.context, case.app_id, "tencent", {**original, field_name: invalid_value}, activate=True) - - assert isinstance(caught.value.__cause__, TencentConfigurationError) - assert caught.value.__cause__.code == error_code - assert forbidden_attempts == [] - assert case.collector.channels == [] - assert case.read() == before - assert case.read("langsmith") == old_provider - assert case.reader.get_trace(case.context, case.app_id) == before_policy - assert "synthetic-private-token" not in caplog.text - assert "private-endpoint-detail" not in caplog.text - - @staticmethod - @pytest.mark.parametrize("existing", [False, True]) - @pytest.mark.parametrize("endpoint", ["https://collector.example:4320", "http://localhost:4317"]) - def test_rpc_failure_never_saves_or_activates_even_for_localhost( - case: _TencentTracingConfigGatewayConfigurationcase, - caplog: pytest.LogCaptureFixture, - existing: bool, - endpoint: str, - ) -> None: - secret = "synthetic-private-token" - original = {"token": secret, "endpoint": "https://original.example:4320", "service_name": "original-app"} - if existing: - case.service.create(case.context, case.app_id, "tencent", original) - case.collector.channels.clear() - before = case.read() - old_provider = case.read("langsmith") - before_policy = case.reader.get_trace(case.context, case.app_id) - case.collector.failure = grpc.RpcError("private-provider-details synthetic-private-token") - operation = case.service.update if existing else case.service.create - - with pytest.raises(AppTracingConfigVerificationFailedError) as caught: - operation( - case.context, - case.app_id, - "tencent", - {"token": "******" if existing else secret, "endpoint": endpoint, "service_name": "changed-app"}, - activate=True, - ) - - assert isinstance(caught.value.__cause__, TencentConfigurationError) - assert caught.value.__cause__.code == "connection_failed" - assert len(case.collector.channels) == 1 - channel = case.collector.channels[0] - assert channel.close_count == 1 - assert len(channel.calls) == 1 - assert channel.calls[0].metadata == (("authorization", f"Bearer {secret}"),) - assert case.read() == before - assert case.read("langsmith") == old_provider - assert case.reader.get_trace(case.context, case.app_id) == before_policy - assert "connection_failed" in caplog.text - assert "private-provider-details" not in caplog.text - assert secret not in caplog.text - assert secret not in str(caught.value) - assert secret not in str(caught.value.__cause__) - - @staticmethod - def test_explicit_configuration_ignores_and_preserves_otel_environment( - case: _TencentTracingConfigGatewayConfigurationcase, monkeypatch: pytest.MonkeyPatch - ) -> None: - environment = { - "OTEL_EXPORTER_OTLP_ENDPOINT": "http://ambient.example:4317", - "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT": "http://ambient-traces.example:4317", - "OTEL_EXPORTER_OTLP_HEADERS": "authorization=Bearer unrelated-token", - "OTEL_EXPORTER_OTLP_TRACES_HEADERS": "authorization=Bearer unrelated-traces-token", - "OTEL_EXPORTER_OTLP_PROTOCOL": "http/protobuf", - "OTEL_SERVICE_NAME": "unrelated-service", - "OTEL_RESOURCE_ATTRIBUTES": "token=unrelated-resource-token,service.name=ambient-app", - } - for key, value in environment.items(): - monkeypatch.setenv(key, value) - - case.service.create( - case.context, - case.app_id, - "tencent", - {"token": "explicit-token", "endpoint": "https://explicit.example:4320", "service_name": "explicit-app"}, - ) - - assert len(case.collector.channels) == 1 - channel = case.collector.channels[0] - assert channel.target == "explicit.example:4320" - assert channel.secure is True - assert channel.calls[0].metadata == (("authorization", "Bearer explicit-token"),) - assert channel.calls[0].request.SerializeToString() == b"" - assert channel.close_count == 1 - presented = case.service.get(case.context, case.app_id, "tencent") - assert presented is not None - assert presented.tracing_config is not None - assert presented.tracing_config["service_name"] == "explicit-app" - assert len(case.collector.channels) == 1 - assert {key: os.environ.get(key) for key in environment} == environment - - -"""Config reads and writes must not construct Weave's stateful legacy exporter.""" - - -ReadAPI = _test_importlib.import_module( - "tests.unit_tests.providers.trace.trace-weave.src.dify_trace_weave.test_configuration" -).ReadAPI - - -@dataclass -class _WeaveTracingConfigGatewayGatewayenvironment: - gateway: OpsTraceManagerGateway = field(default_factory=OpsTraceManagerGateway) - api: ReadAPI = field(default_factory=ReadAPI) - forbidden_calls: list[str] = field(default_factory=list) - - -class TestWeaveTracingConfigGatewayCases: - @staticmethod - @pytest.fixture - def environment(monkeypatch: pytest.MonkeyPatch) -> Iterator[_WeaveTracingConfigGatewayGatewayenvironment]: - environment = _WeaveTracingConfigGatewayGatewayenvironment() - - def encrypt(_tenant: str, value: str) -> str: - return "test-cipher:" + value - - def decrypt(_tenant: str, values: Sequence[str]) -> list[str]: - assert all(value.startswith("test-cipher:") for value in values) - return [value.removeprefix("test-cipher:") for value in values] - - def mask(_value: str) -> str: - return "******" - - def forbidden(name: str) -> Callable[..., Never]: - def block(*_args: object, **_kwargs: object) -> Never: - environment.forbidden_calls.append(name) - raise AssertionError("Configuration must not initialize a stateful exporter") - - return block - - monkeypatch.setattr(ops_trace_manager, "encrypt_token", encrypt) - monkeypatch.setattr(ops_trace_manager, "batch_decrypt_token", decrypt) - monkeypatch.setattr(ops_trace_manager, "obfuscated_token", mask) - monkeypatch.setattr(ssrf_proxy, "post", environment.api.post) - monkeypatch.setattr(WeaveDataTrace, "__init__", forbidden("legacy_constructor")) - monkeypatch.setattr(wandb, "login", forbidden("wandb_login")) - monkeypatch.setattr(weave, "init", forbidden("weave_init")) - yield environment - assert environment.forbidden_calls == [] - - @staticmethod - def test_create_persists_resolved_default_entity_and_canonical_project( - environment: _WeaveTracingConfigGatewayGatewayenvironment, - ) -> None: - submitted = {"api_key": "synthetic-key", "project": " demo "} - stored = environment.gateway.prepare_new_config( - workspace_id=str(uuid4()), tracing_provider="weave", tracing_config=submitted - ) - assert stored["entity"] == "default-team" - assert stored["project"] == "demo" - assert stored["project_url"] == "https://wandb.ai/default-team/demo" - assert stored["api_key"] == "test-cipher:synthetic-key" - assert submitted == {"api_key": "synthetic-key", "project": " demo "} - assert len(environment.api.calls) == 3 - - @staticmethod - def test_patch_persists_verified_destination_with_the_existing_ciphertext( - environment: _WeaveTracingConfigGatewayGatewayenvironment, sqlite_session_factory: sessionmaker[Session] - ) -> None: - workspace, app_id = str(uuid4()), str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add(make_app(app_id=app_id, tenant_id=workspace)) - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - service = AppTracingConfigService(configs=store, provider=environment.gateway) - context = RequestContext("request", None, "account", workspace) - service.create(context, app_id, "weave", {"api_key": "synthetic-key", "entity": "old", "project": "original"}) - before = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="weave") - assert before is not None - assert before.tracing_config is not None - - submitted = {"api_key": "******", "project": "new-team/changed"} - service.update(context, app_id, "weave", submitted) - after = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="weave") - assert after is not None - assert after.tracing_config is not None - assert after.tracing_config["entity"] == "new-team" - assert after.tracing_config["project"] == "changed" - assert after.tracing_config["project_url"] == "https://wandb.ai/new-team/changed" - assert after.tracing_config["api_key"] == before.tracing_config["api_key"] == "test-cipher:synthetic-key" - assert before.tracing_config["project_url"] == "https://wandb.ai/old/original" - assert submitted == {"api_key": "******", "project": "new-team/changed"} - assert len(environment.api.calls) == 4 - - presented = service.get(context, app_id, "weave") - assert presented is not None - assert presented.tracing_config is not None - assert presented.tracing_config["api_key"] == "******" - assert presented.tracing_config["project_url"] == after.tracing_config["project_url"] - environment.gateway.validate_saved_config( - workspace_id=workspace, tracing_provider="weave", tracing_config=after.tracing_config - ) - assert len(environment.api.calls) == 4 - - @staticmethod - @pytest.mark.parametrize( - ("entity", "project", "host", "expected"), - [ - ("team", "demo", None, "https://wandb.ai/team/demo"), - (None, "team/demo", None, "https://wandb.ai/team/demo"), - (None, "demo", None, "https://wandb.ai"), - (None, "demo", "http://private.test/wandb", "http://private.test/wandb"), - ("team", "a & b", "https://private.test/wandb", "https://private.test/wandb/team/a%20%26%20b"), - ], - ) - def test_get_computes_project_or_console_link_without_vendor_requests( - environment: _WeaveTracingConfigGatewayGatewayenvironment, - entity: str | None, - project: str, - host: str | None, - expected: str, - ) -> None: - workspace = str(uuid4()) - encrypted = OpsTraceManager.encrypt_tracing_config( - workspace, "weave", {"api_key": "synthetic-key", "entity": entity, "project": project, "host": host} - ) - encrypted["project_url"] = "https://stale.test/previous-project" - result = environment.gateway.present_config( - workspace_id=workspace, tracing_provider="weave", tracing_config=encrypted - ) - assert result["project_url"] == expected - assert result["api_key"] == "******" - assert environment.api.calls == [] - - @staticmethod - def test_inaccessible_project_is_a_sanitized_domain_error_and_is_not_saved( - environment: _WeaveTracingConfigGatewayGatewayenvironment, - sqlite_session_factory: sessionmaker[Session], - caplog: pytest.LogCaptureFixture, - ) -> None: - workspace, app_id = str(uuid4()), str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add(make_app(app_id=app_id, tenant_id=workspace)) - environment.api.project_exists = False - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - service = AppTracingConfigService(configs=store, provider=environment.gateway) - with pytest.raises(AppTracingConfigVerificationFailedError) as caught: - service.create( - RequestContext("request", None, "account", workspace), - app_id, - "weave", - { - "api_key": "synthetic-private-key", - "entity": "team", - "project": "missing", - }, - ) - assert isinstance(caught.value.__cause__, WeaveConfigurationError) - assert caught.value.__cause__.code == "project_not_found_or_inaccessible" - assert store.get(workspace_id=workspace, app_id=app_id, tracing_provider="weave") is None - assert "synthetic-private-key" not in caplog.text - assert "project_not_found_or_inaccessible" in caplog.text - assert len(environment.api.calls) == 1 diff --git a/api/tests/unit_tests/services/test_app_tracing_config_service.py b/api/tests/unit_tests/services/test_app_tracing_config_service.py index c4ec2569ae58a7..a66b0a25c3193e 100644 --- a/api/tests/unit_tests/services/test_app_tracing_config_service.py +++ b/api/tests/unit_tests/services/test_app_tracing_config_service.py @@ -1,9 +1,10 @@ import json -from dataclasses import dataclass +from dataclasses import dataclass, replace from datetime import UTC, datetime from unittest.mock import MagicMock import pytest +from dify_trace_sdk.testing import RecordingEgress from sqlalchemy import event, select from sqlalchemy.orm import Session, sessionmaker @@ -11,7 +12,6 @@ from models.model import App, AppMode, TraceAppConfig from repositories.app.console_repository import ConsoleAppRepository from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository -from services.app_tracing_config_gateway import OpsTraceManagerGateway from services.app_tracing_config_service import ( AppTraceSettings, AppTracingConfigAlreadyExistsError, @@ -20,7 +20,10 @@ AppTracingConfigRecord, AppTracingConfigService, AppTracingConfigVerificationFailedError, + PresentedTracingConfig, ) +from services.tracing.provider_configs import ProviderConfigGateway +from tests.helpers.unit.services.tracing_ports import deployment_provider_configs """Public tracing operations keep credentials and policy in one transaction.""" @@ -33,7 +36,7 @@ class Case: sessions: sessionmaker[Session] store: SQLAlchemyAppTracingConfigRepository service: AppTracingConfigService - gateway: OpsTraceManagerGateway + gateway: ProviderConfigGateway reader: ConsoleAppRepository def policy(self) -> AppTraceSettings: @@ -61,7 +64,7 @@ def case(sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.Monk ), ) ) - gateway = OpsTraceManagerGateway() + gateway = deployment_provider_configs(RecordingEgress([], grpc_target=None)) def prepare(*, tracing_config: dict[str, object], **_: object) -> dict[str, object]: return dict(tracing_config) @@ -257,12 +260,16 @@ def test_get_scopes_lookup_to_workspace_and_presents_config() -> None: configs = MagicMock() configs.get.return_value = stored provider = MagicMock() - provider.present_config.return_value = {"api_key": "******", "project_url": "https://example.com"} + provider.present_config.return_value = PresentedTracingConfig( + values={"project_url": "https://example.com"}, configured_secrets=("api_key",) + ) service = AppTracingConfigService(configs=configs, provider=provider) result = service.get(_context(), "app-1", "arize") - assert result == _record(tracing_config={"api_key": "******", "project_url": "https://example.com"}) + assert result == replace( + _record(tracing_config={"project_url": "https://example.com"}), configured_secrets=("api_key",) + ) assert stored.tracing_config == {"api_key": "encrypted"} provider.validate_provider.assert_called_once_with("arize") configs.get.assert_called_once_with( diff --git a/api/tests/unit_tests/services/test_conversation_name_generator.py b/api/tests/unit_tests/services/test_conversation_name_generator.py index 94d2994e4001b7..e5c5bbd45f925f 100644 --- a/api/tests/unit_tests/services/test_conversation_name_generator.py +++ b/api/tests/unit_tests/services/test_conversation_name_generator.py @@ -1,9 +1,11 @@ """Real provider resolution keeps naming metadata and the caller's session isolated.""" +import json from uuid import uuid4 import httpx import pytest +from dify_trace_sdk.events import ConversationNamed, TraceEvent from flask import Flask from sqlalchemy import Connection, event, select from sqlalchemy.orm import Session, SessionTransaction, sessionmaker @@ -18,8 +20,11 @@ from graphon.model_runtime.errors.invoke import InvokeServerUnavailableError from models import App, AppMode, Tenant from models.provider import Provider, TenantDefaultModel -from services.conversation_name_generator import generate_conversation_name +from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository +from services.conversation_name_generator import DefaultModelConversationNameGenerator +from services.tracing.scopes import TraceScopes from tests.helpers.unit.services.conversation_name_fixtures import _MODEL, _PROVIDER, NamingPlugin +from tests.helpers.unit.services.tracing_ports import untraced_scopes pytest_plugins = ("tests.helpers.unit.services.conversation_name_fixtures",) @@ -85,7 +90,7 @@ def observe_request(request: httpx.Request) -> None: inherited_metadata = dict(get_credit_usage_metadata() or {}) if fail_invoke: with pytest.raises(InvokeServerUnavailableError, match="Title provider unavailable"): - generate_conversation_name( + DefaultModelConversationNameGenerator(trace_scopes=untraced_scopes())( tenant_id=tenant.id, app_id=target.id, conversation_id=str(uuid4()), @@ -93,7 +98,7 @@ def observe_request(request: httpx.Request) -> None: app_mode=target.mode, ) else: - name = generate_conversation_name( + name = DefaultModelConversationNameGenerator(trace_scopes=untraced_scopes())( tenant_id=tenant.id, app_id=target.id, conversation_id=str(uuid4()), @@ -121,12 +126,8 @@ def observe_request(request: httpx.Request) -> None: assert caller_app.name == "Pending caller edit" assert all(session is not caller_session for session in naming_sessions) assert all(not session.in_transaction() and not session.identity_map for session in naming_sessions) - if fail_invoke: - assert naming_sessions == preparation_sessions - else: - # Trace configuration is read after invocation using a new inner session. - assert len(naming_sessions) > len(preparation_sessions) - assert naming_sessions[-1] not in preparation_sessions + # Tracing reads its policy through its own store, so preparation is the only scoped session. + assert naming_sessions == preparation_sessions with sqlite_session_factory() as session: assert session.scalar(select(App.name).where(App.id == target.id)) == "Original app" @@ -198,7 +199,7 @@ def observe_provider_transaction( with use_credit_usage_metadata({"request_id": "name-request"}): inherited_metadata = dict(get_credit_usage_metadata() or {}) with pytest.raises(ProviderTokenNotInitError, match="^Default model not found for llm$"): - generate_conversation_name( + DefaultModelConversationNameGenerator(trace_scopes=untraced_scopes())( tenant_id=target.tenant_id, app_id=target.id, conversation_id=str(uuid4()), @@ -232,3 +233,57 @@ def observe_provider_transaction( stored_app = session.get(App, target.id) assert stored_app is not None assert stored_app.name == "Original app" + + +class _TraceRecorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +@pytest.mark.usefixtures("naming_plugin") +def test_naming_reports_the_generated_name_to_a_traced_apps_trace( + naming_app: Flask, sqlite_session_factory: sessionmaker[Session] +) -> None: + tenant = Tenant(name="Naming workspace") + target = App( + id=str(uuid4()), + tenant_id=tenant.id, + name="Traced app", + mode=AppMode.CHAT, + enable_site=True, + enable_api=True, + tracing=json.dumps({"enabled": True, "tracing_provider": "langfuse"}), + ) + with sqlite_session_factory.begin() as session: + session.add_all( + [ + tenant, + target, + Provider(tenant_id=tenant.id, provider_name=_PROVIDER, is_valid=True), + TenantDefaultModel( + tenant_id=tenant.id, provider_name=_PROVIDER, model_name=_MODEL, model_type=ModelType.LLM + ), + ] + ) + recorder = _TraceRecorder() + generator = DefaultModelConversationNameGenerator( + trace_scopes=TraceScopes( + emitter=recorder, + policies=SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory), + enterprise_enabled=lambda: False, + ) + ) + conversation_id = str(uuid4()) + + with naming_app.app_context(): + name = generator( + tenant_id=tenant.id, app_id=target.id, conversation_id=conversation_id, query="First query", app_mode="chat" + ) + + (event,) = recorder.events + assert (event.tenant_id, event.app_id) == (tenant.id, target.id) + assert isinstance(event.payload, ConversationNamed) + assert (event.payload.conversation_id, event.payload.name) == (conversation_id, name) diff --git a/api/tests/unit_tests/services/test_message_suggested_questions_generator.py b/api/tests/unit_tests/services/test_message_suggested_questions_generator.py index 3df7c726f8b7e8..65adddf0a52392 100644 --- a/api/tests/unit_tests/services/test_message_suggested_questions_generator.py +++ b/api/tests/unit_tests/services/test_message_suggested_questions_generator.py @@ -9,8 +9,8 @@ import httpx import pytest +from dify_trace_sdk.events import SuggestedQuestionsGenerated, TraceEvent from flask import Flask, current_app -from pydantic import ValidationError from sqlalchemy import Connection, Engine, event, select from sqlalchemy.exc import IntegrityError, PendingRollbackError from sqlalchemy.orm import Session, SessionTransaction, sessionmaker @@ -28,6 +28,7 @@ from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent from graphon.model_runtime.entities.model_entities import ModelType from models.provider import Provider, TenantDefaultModel +from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository from services.entities.message_entities import MessageEndUser from services.message_suggested_questions_generator import ( PreparedSuggestedQuestionsModel, @@ -35,8 +36,10 @@ _default_suggested_questions_model_parameters, ) from services.message_suggested_questions_service import SuggestedQuestionsContext +from services.tracing.scopes import TraceScopes from tests.helpers.unit.core.model_fixtures import make_model_instance from tests.helpers.unit.model_factories import make_app +from tests.helpers.unit.services.tracing_ports import untraced_scopes @pytest.fixture @@ -82,7 +85,9 @@ def test_prepare_preserves_caller_transaction_when_stored_provider_is_invalid( caller = db.session() provider = caller.scalars(select(Provider).where(Provider.tenant_id == invalid_provider_tenant)).one() provider.quota_limit = 27 - with SuggestedQuestionsGenerator().prepare(context=context, instruction_prompt=None, model_config=None) as generate: + with SuggestedQuestionsGenerator(trace_scopes=untraced_scopes()).prepare( + context=context, instruction_prompt=None, model_config=None + ) as generate: assert generate is None assert db.session() is caller assert caller.in_transaction() @@ -163,7 +168,7 @@ def test_token_transport_failure_propagates_before_best_effort_generation( metadata = get_credit_usage_metadata() timeout = _get_plugin_daemon_request_timeout() with pytest.raises(PluginDaemonInnerError) as failure: - SuggestedQuestionsGenerator().generate( + SuggestedQuestionsGenerator(trace_scopes=untraced_scopes()).generate( history, history_model=make_model_instance(provider="langgenius/openai/openai", model="history-model"), context=context, @@ -196,11 +201,28 @@ def test_configured_parameters_keep_positive_limits() -> None: ) -@pytest.mark.parametrize("tracing", ["invalid trace JSON", '{"enabled":{"invalid":"type"}}']) -def test_trace_config_failure_escapes_generation_and_closes_its_sessions( +class _TraceRecorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +@pytest.mark.parametrize( + ("tracing", "traced"), + [ + ("invalid trace JSON", False), + ('{"enabled":{"invalid":"type"}}', False), + ('{"enabled":true,"tracing_provider":"langfuse"}', True), + ], + ids=["unparsable", "invalid-policy", "enabled"], +) +def test_tracing_reads_its_policy_in_its_own_session_and_never_fails_generation( invalid_provider_tenant: str, sqlite_session_factory: sessionmaker[Session], tracing: str, + traced: bool, caplog: pytest.LogCaptureFixture, ) -> None: app_record = make_app(app_id=str(uuid4()), tenant_id=invalid_provider_tenant) @@ -227,13 +249,21 @@ def record_session(session: Session, _transaction: SessionTransaction, _connecti sessions.append(session) factory = db.session.session_factory + recorder = _TraceRecorder() + generator = SuggestedQuestionsGenerator( + trace_scopes=TraceScopes( + emitter=recorder, + policies=SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory), + enterprise_enabled=lambda: False, + ) + ) event.listen(factory, "after_begin", record_session) + event.listen(sqlite_session_factory, "after_begin", record_session) try: # Empty detached history needs no token-counting call. The real failed - # model lookup returns [], then trace construction reads the malformed - # stored config before it could create a timer or dispatch telemetry. - with pytest.raises(ValidationError), current_app.app_context(): - SuggestedQuestionsGenerator().generate( + # model lookup returns [], then the trace scope reads the stored policy. + with current_app.app_context(): + questions = generator.generate( PreparedHistory(prompts=()), history_model=make_model_instance(provider="invalid/provider", model="unused"), context=context, @@ -242,7 +272,17 @@ def record_session(session: Session, _transaction: SessionTransaction, _connecti ) finally: event.remove(factory, "after_begin", record_session) - + event.remove(sqlite_session_factory, "after_begin", record_session) + + assert questions == [] + if traced: + (trace_event,) = recorder.events + assert (trace_event.tenant_id, trace_event.app_id) == (invalid_provider_tenant, app_record.id) + assert isinstance(trace_event.payload, SuggestedQuestionsGenerated) + assert (trace_event.payload.message_id, trace_event.payload.questions) == (context.message_id, []) + assert trace_event.payload.started_at <= trace_event.payload.ended_at + else: + assert recorder.events == [] assert "invalid/provider" in caplog.text assert len(sessions) == 2 assert sessions[0] is not sessions[1] @@ -281,7 +321,9 @@ def test_prepare_isolates_a_caller_transaction_failed_by_a_real_constraint( assert failed_transaction is not None assert not caller.is_active - with SuggestedQuestionsGenerator().prepare(context=context, instruction_prompt=None, model_config=None) as generate: + with SuggestedQuestionsGenerator(trace_scopes=untraced_scopes()).prepare( + context=context, instruction_prompt=None, model_config=None + ) as generate: assert generate is None # Model resolution used an independent usable session. It must neither diff --git a/api/tests/unit_tests/services/test_message_suggested_questions_service.py b/api/tests/unit_tests/services/test_message_suggested_questions_service.py index f1f291a7e89607..cb4c045a4c55aa 100644 --- a/api/tests/unit_tests/services/test_message_suggested_questions_service.py +++ b/api/tests/unit_tests/services/test_message_suggested_questions_service.py @@ -56,6 +56,7 @@ SuggestedQuestionsInvokeFrom, ) from tests.helpers.unit.core.model_fixtures import make_model_instance +from tests.helpers.unit.services.tracing_ports import untraced_scopes @dataclass(frozen=True) @@ -197,7 +198,9 @@ def track_scoped(session: Session, _transaction: SessionTransaction, _connection message, sqlite_session_factory, queries, - MessageSuggestedQuestionsService(queries=queries, generator=SuggestedQuestionsGenerator()), + MessageSuggestedQuestionsService( + queries=queries, generator=SuggestedQuestionsGenerator(trace_scopes=untraced_scopes()) + ), read_sessions, scoped_sessions, ) diff --git a/api/tests/unit_tests/services/test_workflow_node_execution_trace_service.py b/api/tests/unit_tests/services/test_workflow_node_execution_trace_service.py index 1324c13c376323..b26abc34d7cdd1 100644 --- a/api/tests/unit_tests/services/test_workflow_node_execution_trace_service.py +++ b/api/tests/unit_tests/services/test_workflow_node_execution_trace_service.py @@ -9,7 +9,7 @@ import pytest from sqlalchemy.orm import Session -from core.app.workflow.retry_history import RETRY_HISTORY_PROCESS_DATA_KEY, WorkflowNodeRetryAttempt +from domain.workflow.retry_history import RETRY_HISTORY_PROCESS_DATA_KEY, WorkflowNodeRetryAttempt from fields.workflow_run_fields import WorkflowRunNodeExecutionListResponse from graphon.enums import WorkflowNodeExecutionStatus from models.account import Account diff --git a/api/tests/unit_tests/services/test_workflow_persistence.py b/api/tests/unit_tests/services/test_workflow_persistence.py index cf508c3f14e86a..1d3bf637c4a6dd 100644 --- a/api/tests/unit_tests/services/test_workflow_persistence.py +++ b/api/tests/unit_tests/services/test_workflow_persistence.py @@ -3,19 +3,18 @@ from dataclasses import replace from datetime import UTC, datetime, timedelta from types import SimpleNamespace -from typing import Any +from typing import Any, override from unittest.mock import MagicMock import pytest +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.events import TraceCorrelation, TraceEvent, WorkflowRunCompleted from flask import Flask from sqlalchemy import event as sqlalchemy_event from sqlalchemy import select from sqlalchemy.orm import Session, sessionmaker from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, WorkflowAppGenerateEntity -from core.app.workflow.retry_history import RETRY_HISTORY_PROCESS_DATA_KEY -from core.ops.base_trace_instance import bind_application_tracing_hooks -from core.ops.ops_trace_manager import TraceTask, TraceTaskName from core.repositories.sqlalchemy_workflow_execution_repository import SQLAlchemyWorkflowExecutionRepository from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository from core.tools.workflow_as_tool.repository import WorkflowToolSource @@ -24,7 +23,9 @@ WORKFLOW_TOOL_ROOT_APP_ID_KEY, ) from core.workflow.system_variables import SystemVariableKey, build_system_variables +from domain.workflow.retry_history import RETRY_HISTORY_PROCESS_DATA_KEY from enums import WorkflowKind +from extensions.tracing import bind_tracing from graphon.engine.event.node_failure import NodeFailureHandler from graphon.engine.event.processor import NodeEventProcessor from graphon.engine.event.stream import EventStream @@ -63,6 +64,7 @@ from graphon.nodes.start.start_node import StartNode from graphon.runtime import ReadOnlyRuntimeStateWrapper, RuntimeState, VariablePool from graphon.runtime.execution import ROOT_FRAME_ID +from graphon.variables.variables import Variable from models import Account, WorkflowRun from models.enums import WorkflowRunTriggeredFrom from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom @@ -71,9 +73,8 @@ DifyAPISQLAlchemyWorkflowNodeExecutionRepository, ) from services import workflow_persistence as persistence_mod -from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks from services.workflow_persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer -from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.services.tracing_ports import NoTracingPolicies, RecordingTracingHooks, kernel_trace_scopes from tests.workflow_test_utils import build_test_graph_init_params """Verify the workflow persistence layer fans Inspector deltas to redis pub/sub. @@ -138,6 +139,7 @@ def layer() -> WorkflowPersistenceLayer: workflow_info=MagicMock(workflow_id="wf-1"), workflow_execution_repository=MagicMock(), workflow_node_execution_repository=MagicMock(), + trace_scope=None, ) # Use a SimpleNamespace-like spec so Pydantic-validated callsites (e.g. # ``WorkflowNodeExecution.new`` requires real strings) get the right types. @@ -273,6 +275,29 @@ def save_execution_data(self, entity): self.saved_exec_data.append(entity) +class _TraceRecorder: + """Record trace events with the number of workflow saves that preceded each one.""" + + def __init__(self, workflow_repo: _RepoRecorder) -> None: + self._workflow_repo = workflow_repo + self.events: list[TraceEvent] = [] + self.saves_before_emit: list[int] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + self.saves_before_emit.append(len(self._workflow_repo.synchronously_saved)) + + +def _trace_scope(recorder: _TraceRecorder, correlation: TraceCorrelation) -> TraceScope: + return TraceScope(emitter=recorder, tenant_id="tenant", app_id="app", user_id="user", correlation=correlation) + + +def _empty_correlation() -> TraceCorrelation: + return TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ) + + def _naive_utc_now() -> datetime: return datetime.now(UTC).replace(tzinfo=None) @@ -282,7 +307,7 @@ def _make_layer( *, extras: dict | None = None, graph_data: dict | None = None, - trace_manager: object | None = None, + trace_scope: TraceScope | None = None, workflow_execution_repo=None, workflow_node_execution_repo=None, node_run_indices=None, @@ -306,7 +331,7 @@ def _make_layer( user_id="user", stream=False, invoke_from=None, - trace_manager=None, + trace_scope=trace_scope, workflow_execution_id="run-id", extras=extras or {}, call_depth=0, @@ -329,7 +354,7 @@ def _make_layer( workflow_info=workflow_info, workflow_execution_repository=workflow_execution_repo, workflow_node_execution_repository=workflow_node_execution_repo, - trace_manager=trace_manager, + trace_scope=trace_scope, ) layer.initialize(read_only_state, command_channel=None) layer.set_node_run_indices(node_run_indices or {}) @@ -373,11 +398,6 @@ def _tool_source(): class TestPersistenceLayerCases: - @staticmethod - @pytest.fixture(autouse=True) - def _application_tracing_disabled(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) - @staticmethod def test_tool_trace_identity_survives_start_pause_and_resume(): layer, _, node_repo, _ = _make_layer(node_run_indices={"tool-exec": 1}) @@ -874,38 +894,44 @@ def record_commit(_session: Session) -> None: assert len(commits) == 3 class TestWorkflowPersistenceLayer: - def test_start_freezes_tracing_only_after_workflow_save(self, monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) + def test_start_freezes_tracing_only_after_workflow_save(self) -> None: layer, workflow_repo, _, _ = _make_layer(extras={"trace_session_id": "external-session"}) calls: list[tuple[str, str, str, str, object]] = [] - def freeze( - entity: AdvancedChatAppGenerateEntity | WorkflowAppGenerateEntity, - *, - workflow_id: str, - workflow_run_id: str, - conversation_id: object, - graph_data: object, - ) -> None: - assert workflow_repo.synchronously_saved - assert graph_data == {"nodes": [], "edges": []} - saved_execution = workflow_repo.synchronously_saved[-1] - assert isinstance(saved_execution, WorkflowExecution) - assert saved_execution.id_ == workflow_run_id - calls.append( - ( - entity.app_config.app_id, - entity.app_config.tenant_id, - workflow_id, - workflow_run_id, - conversation_id, + class _FreezeAfterSave(RecordingTracingHooks): + @override + def freeze_workflow_run( + self, + entity: AdvancedChatAppGenerateEntity | WorkflowAppGenerateEntity, + *, + workflow_id: str, + workflow_run_id: str, + conversation_id: object, + graph_data: object, + ) -> None: + assert workflow_repo.synchronously_saved + assert graph_data == {"nodes": [], "edges": []} + saved_execution = workflow_repo.synchronously_saved[-1] + assert isinstance(saved_execution, WorkflowExecution) + assert saved_execution.id_ == workflow_run_id + calls.append( + ( + entity.app_config.app_id, + entity.app_config.tenant_id, + workflow_id, + workflow_run_id, + conversation_id, + ) ) - ) - monkeypatch.setattr(RuntimeApplicationTracingHooks, "freeze_workflow_run", staticmethod(freeze)) - hooks = RuntimeApplicationTracingHooks() app = Flask(__name__) - bind_application_tracing_hooks(app, hooks) + bind_tracing( + app, + scopes=kernel_trace_scopes( + _TraceRecorder(workflow_repo), policies=NoTracingPolicies(), enterprise_enabled=False + ), + hooks=_FreezeAfterSave(None), + ) with app.app_context(): layer._handle_graph_run_started() @@ -1033,11 +1059,11 @@ def test_handle_graph_run_partial_succeeded_updates_execution(self): assert saved.exceptions_count == 2 assert saved.total_tokens == 5 - def test_handle_graph_run_failed_marks_nodes_and_enqueues_trace(self): - trace_tasks: list[object] = [] - trace_manager = SimpleNamespace(user_id="user", add_trace_task=lambda task: trace_tasks.append(task)) + def test_handle_graph_run_failed_marks_nodes_and_emits_trace_after_save(self) -> None: + workflow_repo = _RepoRecorder() + recorder = _TraceRecorder(workflow_repo) layer, exec_repo, node_repo, _ = _make_layer( - extras={"external_trace_id": "trace"}, trace_manager=trace_manager + trace_scope=_trace_scope(recorder, _empty_correlation()), workflow_execution_repo=workflow_repo ) layer.on_event(GraphRunStartedEvent()) @@ -1057,64 +1083,62 @@ def test_handle_graph_run_failed_marks_nodes_and_enqueues_trace(self): assert node_repo.saved assert exec_repo.synchronously_saved[-1].status == WorkflowExecutionStatus.FAILED - assert trace_tasks + assert [event.payload for event in recorder.events] == [ + WorkflowRunCompleted(workflow_run_id="run-id", conversation_id="conv-id", total_tokens=0) + ] + # Started and finished saves both precede the event. + assert recorder.saves_before_emit == [2] - def test_handle_graph_run_succeeded_enqueues_parent_trace_context(self, monkeypatch: pytest.MonkeyPatch): - trace_tasks: list[TraceTask] = [] - trace_manager = SimpleNamespace(user_id="user", add_trace_task=lambda task: trace_tasks.append(task)) - layer, _, _, _ = _make_layer( - extras={ - "external_trace_id": "trace", - "trace_session_id": "session-1", - "parent_trace_context": { - "parent_workflow_run_id": "outer-workflow-run-1", - "parent_node_execution_id": "outer-node-execution-1", - }, - }, - trace_manager=trace_manager, + @pytest.mark.parametrize( + ("system_variables", "conversation_id"), + [ + (build_system_variables(workflow_execution_id="run-id", conversation_id="conv-id"), "conv-id"), + (build_system_variables(workflow_execution_id="run-id"), None), + ], + ) + def test_handle_graph_run_succeeded_emits_run_tokens_with_scope_correlation( + self, system_variables: list[Variable], conversation_id: str | None + ) -> None: + workflow_repo = _RepoRecorder() + recorder = _TraceRecorder(workflow_repo) + correlation = TraceCorrelation( + external_trace_id="trace", + session_id="session-1", + parent_workflow_run_id="outer-workflow-run-1", + parent_node_execution_id="outer-node-execution-1", + ) + layer, _, _, runtime_state = _make_layer( + system_variables, + trace_scope=_trace_scope(recorder, correlation), + workflow_execution_repo=workflow_repo, ) layer.on_event(GraphRunStartedEvent()) - - captured: dict[str, object] = {} - - def fake_workflow_trace( - self: TraceTask, - *, - workflow_run_id: str | None, - conversation_id: str | None, - user_id: str | None, - total_tokens_override: int | None = None, - ): - captured["trace_type"] = self.trace_type - captured["external_trace_id"] = self.kwargs.get("external_trace_id") - captured["trace_session_id"] = self.kwargs.get("trace_session_id") - captured["parent_trace_context"] = self.kwargs.get("parent_trace_context") - captured["workflow_run_id"] = workflow_run_id - return {"ok": True} - - monkeypatch.setattr(TraceTask, "workflow_trace", fake_workflow_trace) + usage = LLMUsage.empty_usage() + usage.total_tokens = 7 + runtime_state.add_llm_usage(usage) layer.on_event(GraphRunSucceededEvent(outputs={"ok": True})) - assert trace_tasks - trace_task = trace_tasks[0] - assert trace_task.trace_type == TraceTaskName.WORKFLOW_TRACE - assert trace_task.kwargs["external_trace_id"] == "trace" - assert trace_task.kwargs["trace_session_id"] == "session-1" - assert trace_task.kwargs["parent_trace_context"] == { - "parent_workflow_run_id": "outer-workflow-run-1", - "parent_node_execution_id": "outer-node-execution-1", - } + assert len(recorder.events) == 1 + event = recorder.events[0] + assert event.payload == WorkflowRunCompleted( + workflow_run_id="run-id", conversation_id=conversation_id, total_tokens=7 + ) + assert event.correlation == correlation + assert (event.tenant_id, event.app_id, event.user_id) == ("tenant", "app", "user") + assert recorder.saves_before_emit == [2] + + def test_paused_run_emits_no_trace(self) -> None: + workflow_repo = _RepoRecorder() + recorder = _TraceRecorder(workflow_repo) + layer, _, _, _ = _make_layer( + trace_scope=_trace_scope(recorder, _empty_correlation()), workflow_execution_repo=workflow_repo + ) + layer.on_event(GraphRunStartedEvent()) - trace_task.execute() + layer.on_event(GraphRunPausedEvent(reasons=[SchedulingPause(message="pause")])) - assert captured["trace_type"] == TraceTaskName.WORKFLOW_TRACE - assert captured["external_trace_id"] == "trace" - assert captured["trace_session_id"] == "session-1" - assert captured["parent_trace_context"] == { - "parent_workflow_run_id": "outer-workflow-run-1", - "parent_node_execution_id": "outer-node-execution-1", - } + assert recorder.events == [] def test_handle_graph_run_aborted_sets_status(self): layer, exec_repo, _, _ = _make_layer() diff --git a/api/tests/unit_tests/services/test_workflow_run_index.py b/api/tests/unit_tests/services/test_workflow_run_index.py index dffb8e5fe6eeb6..08347c1c8712ba 100644 --- a/api/tests/unit_tests/services/test_workflow_run_index.py +++ b/api/tests/unit_tests/services/test_workflow_run_index.py @@ -107,6 +107,7 @@ def test_engine_indices_are_available_before_delivery_and_match_persistence( ), workflow_execution_repository=execution_repo, workflow_node_execution_repository=node_repo, + trace_scope=None, ) layer.set_node_run_indices(index.indices) layer.set_node_execution_history(persisted_history) diff --git a/api/tests/unit_tests/services/tracing/test_adapters.py b/api/tests/unit_tests/services/tracing/test_adapters.py new file mode 100644 index 00000000000000..1afaf8a09318e3 --- /dev/null +++ b/api/tests/unit_tests/services/tracing/test_adapters.py @@ -0,0 +1,363 @@ +from collections.abc import Iterator +from datetime import UTC, datetime +from multiprocessing import get_context +from pathlib import Path +from queue import Empty +from threading import Event, Thread +from typing import override +from uuid import uuid4 + +import pytest +from celery import Celery +from celery.contrib.testing.app import TestApp, setup_default_app +from celery.signals import before_task_publish, worker_process_shutdown +from dify_trace_sdk.events import ModerationChecked, TraceCorrelation, TraceEvent, WorkflowRunCompleted +from dify_trace_sdk.model import Span, SpanKind, SpanStatus, Trace +from kombu.message import Message +from kombu.simple import SimpleQueue + +from extensions.storage.opendal_storage import OpenDALStorage +from services.app_tracing_config_service import AppTraceSettings +from services.tracing import adapters +from services.tracing.adapters import ( + AGENT_TRACE_EXPORT_TASK, + DELIVER_TASK, + INLINE_EVENT_BYTES, + PENDING_EVENT_LIMIT, + RETRY_TASK, + WORKFLOW_SETTLE_SECONDS, + CeleryExportRetryQueue, + CeleryTraceEmitter, + StoredPolicyTraceScopes, + payload_key, + schedule_agent_trace_export, +) +from services.tracing.delivery import ExportJob +from services.tracing.scopes import TraceScopes, TracingPolicyQuery + +STARTED_AT = datetime(2026, 3, 4, 5, 6, tzinfo=UTC) +CORRELATION = TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None +) + + +class _Recorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +class _NoPolicies(TracingPolicyQuery): + @override + def policy(self, *, tenant_id: str, app_id: str) -> AppTraceSettings | None: + raise AssertionError("an execution that holds its policy must not read it again") + + +@pytest.fixture +def broker() -> Iterator[Celery]: + app = TestApp(config={"task_default_queue": f"tracing-{uuid4().hex}"}) + with setup_default_app(app): + app.set_current() + app.set_default() + yield app + + +@pytest.fixture +def storage(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> OpenDALStorage: + local = OpenDALStorage(scheme="fs", root=str(tmp_path)) + monkeypatch.setattr(adapters, "storage", local) + return local + + +def published(broker: Celery) -> list[Message]: + messages: list[Message] = [] + with broker.connection_for_read() as connection, SimpleQueue(connection, broker.conf.task_default_queue) as inbox: + while True: + try: + message = inbox.get(block=False) + except Empty: + return messages + message.ack() + messages.append(message) + + +def _event(query: str) -> TraceEvent: + return TraceEvent( + event_id="event-1", + tenant_id="tenant-1", + app_id="app-1", + user_id=None, + occurred_at=STARTED_AT, + correlation=CORRELATION, + payload=ModerationChecked( + message_id="message-1", + inputs={}, + query=query, + flagged=False, + action="direct_output", + preset_response="", + started_at=STARTED_AT, + ended_at=STARTED_AT, + ), + ) + + +def test_small_events_travel_inline(broker: Celery, storage: OpenDALStorage) -> None: + event = _event("hello") + + emitter = CeleryTraceEmitter() + emitter.emit(event) + emitter.shutdown() + + (message,) = published(broker) + assert message.headers["task"] == DELIVER_TASK + (envelope,) = message.payload[0] + assert isinstance(envelope, dict) + assert TraceEvent.model_validate_json(envelope["event"]) == event + assert not storage.exists("trace_payloads") + + +def test_large_events_travel_through_storage(broker: Celery, storage: OpenDALStorage) -> None: + event = _event("x" * INLINE_EVENT_BYTES) + + emitter = CeleryTraceEmitter() + emitter.emit(event) + emitter.shutdown() + + (message,) = published(broker) + (envelope,) = message.payload[0] + assert envelope == {"payload_key": payload_key(event.occurred_at, "events", "event-1")} + assert payload_key(event.occurred_at, "events", "event-1") == "trace_payloads/2026-03-04/events/event-1.json" + assert TraceEvent.model_validate_json(storage.load_once("trace_payloads/2026-03-04/events/event-1.json")) == event + + +@pytest.mark.usefixtures("storage") +def test_inline_size_is_measured_in_bytes(broker: Celery) -> None: + emitter = CeleryTraceEmitter() + emitter.emit(_event("中" * (INLINE_EVENT_BYTES // 2))) + emitter.shutdown() + + (message,) = published(broker) + (envelope,) = message.payload[0] + assert isinstance(envelope, dict) + assert "payload_key" in envelope + + +@pytest.mark.usefixtures("storage") +def test_workflow_runs_wait_for_node_executions_to_land(broker: Celery) -> None: + workflow = _event("hello").model_copy( + update={"payload": WorkflowRunCompleted(workflow_run_id="run-1", conversation_id=None, total_tokens=0)} + ) + + emitter = CeleryTraceEmitter() + emitter.emit(_event("hello")) + before = datetime.now(UTC) + emitter.emit(workflow) + emitter.shutdown() + + immediate, delayed = published(broker) + assert immediate.headers["eta"] is None + eta = datetime.fromisoformat(delayed.headers["eta"]) + assert (eta - before).total_seconds() >= WORKFLOW_SETTLE_SECONDS + + +def test_export_retries_persist_the_trace_and_run_after_their_delay(broker: Celery, storage: OpenDALStorage) -> None: + job = ExportJob( + tenant_id="tenant-1", + app_id="app-1", + provider_id="langfuse", + attempt=2, + trace=Trace( + trace_id="a" * 32, + external_trace_id=None, + session_id=None, + user_id=None, + parent=None, + spans=( + Span( + span_id="b" * 16, + parent_span_id=None, + name="message", + kind=SpanKind.CHAIN, + start_time=STARTED_AT, + end_time=STARTED_AT, + status=SpanStatus.OK, + ), + ), + ), + ) + + CeleryExportRetryQueue().schedule(job, delay_seconds=20) + + (message,) = published(broker) + (key,) = message.payload[0] + assert message.headers["task"] == RETRY_TASK + assert message.headers["eta"] is not None + assert isinstance(key, str) + assert key.startswith("trace_payloads/") + assert "/exports/" in key + assert ExportJob.model_validate_json(storage.load_once(key)) == job + + +def test_agent_trace_exports_run_on_the_trace_queue(broker: Celery) -> None: + schedule_agent_trace_export("run-1") + + (message,) = published(broker) + assert message.headers["task"] == AGENT_TRACE_EXPORT_TASK + assert message.payload[0] == ["run-1"] + + +def test_blocked_publisher_does_not_block_producers_and_bounds_pending_events( + broker: Celery, caplog: pytest.LogCaptureFixture +) -> None: + entered, release, returned = Event(), Event(), Event() + + def hold_publication(**_: object) -> None: + entered.set() + assert release.wait(10) + + emitter = CeleryTraceEmitter() + + def produce() -> None: + broker.set_current() + emitter.emit(_event("first")) + returned.set() + + before_task_publish.connect(hold_publication, sender=DELIVER_TASK, weak=False) + producer = Thread(target=produce) + try: + producer.start() + assert entered.wait(5) + assert returned.wait(1), "emit waited for broker publication" + for _ in range(PENDING_EVENT_LIMIT - 1): + emitter.emit(_event("buffered")) + emitter.emit(_event("overflow")) + assert "buffer full" in caplog.text + assert published(broker) == [] + finally: + release.set() + producer.join(timeout=5) + emitter.shutdown() + before_task_publish.disconnect(hold_publication, sender=DELIVER_TASK) # type: ignore[attr-defined] # Missing from celery-types. + + assert len(published(broker)) == PENDING_EVENT_LIMIT + + +def test_publication_snapshots_payload_and_preserves_callers_celery_app(broker: Celery) -> None: + entered, release = Event(), Event() + + def hold_publication(**_: object) -> None: + entered.set() + assert release.wait(10) + + other = TestApp(config={"task_default_queue": f"other-{uuid4().hex}"}) + other.set_default() + emitter = CeleryTraceEmitter() + event = _event("second") + before_task_publish.connect(hold_publication, sender=DELIVER_TASK, weak=False) + try: + emitter.emit(_event("first")) + assert entered.wait(5) + emitter.emit(event) + assert isinstance(event.payload, ModerationChecked) + event.payload.inputs["changed_after_emit"] = True + finally: + release.set() + emitter.shutdown() + before_task_publish.disconnect(hold_publication, sender=DELIVER_TASK) # type: ignore[attr-defined] # Missing from celery-types. + other.close() + + first, second = published(broker) + assert TraceEvent.model_validate_json(first.payload[0][0]["event"]).payload == _event("first").payload + captured = TraceEvent.model_validate_json(second.payload[0][0]["event"]) + assert isinstance(captured.payload, ModerationChecked) + assert captured.payload.inputs == {} + + +def test_failed_storage_publication_releases_capacity_and_later_events_continue( + broker: Celery, storage: OpenDALStorage, caplog: pytest.LogCaptureFixture +) -> None: + # A file in place of the payload directory makes the real filesystem write fail. + storage.save("trace_payloads", b"not a directory") + emitter = CeleryTraceEmitter() + emitter.emit(_event("x" * INLINE_EVENT_BYTES)) + emitter.emit(_event("next")) + emitter.shutdown() + + assert "Trace publication failed" in caplog.text + (message,) = published(broker) + event = TraceEvent.model_validate_json(message.payload[0][0]["event"]) + assert isinstance(event.payload, ModerationChecked) + assert event.payload.query == "next" + for _ in range(PENDING_EVENT_LIMIT): + assert emitter._slots.acquire(blocking=False) + + +def test_forked_worker_has_its_own_publisher_and_drains_on_shutdown(broker: Celery) -> None: + emitter = CeleryTraceEmitter() + emitter.emit(_event("parent")) + with broker.connection_for_read() as connection, SimpleQueue(connection, broker.conf.task_default_queue) as inbox: + message = inbox.get(block=True, timeout=5) + message.ack() + + context = get_context("fork") + reader, writer = context.Pipe(duplex=False) + + def publish_from_child() -> None: + reader.close() + emitter.emit(_event("child")) + worker_process_shutdown.send(sender=None) + assert emitter._closed, "worker shutdown did not close the publisher" + (message,) = published(broker) + event = TraceEvent.model_validate_json(message.payload[0][0]["event"]) + assert isinstance(event.payload, ModerationChecked) + writer.send(event.payload.query) + writer.close() + + child = context.Process(target=publish_from_child) + try: + child.start() + writer.close() + assert reader.poll(10), "the child inherited an unusable publisher" + assert reader.recv() == "child" + child.join(timeout=5) + assert child.exitcode == 0 + assert published(broker) == [] + emitter.emit(_event("parent after fork")) + finally: + if child.is_alive(): + child.terminate() + child.join(timeout=5) + reader.close() + writer.close() + emitter.shutdown() + + (message,) = published(broker) + event = TraceEvent.model_validate_json(message.payload[0][0]["event"]) + assert isinstance(event.payload, ModerationChecked) + assert event.payload.query == "parent after fork" + + +@pytest.mark.parametrize( + ("stored", "traced"), + [ + ('{"enabled": true, "tracing_provider": "langfuse"}', True), + ('{"enabled": "true", "tracing_provider": "langfuse"}', False), + ("not json", False), + (None, False), + ], + ids=["enabled", "not-a-boolean", "unparsable", "unset"], +) +def test_legacy_producers_hand_over_the_raw_stored_policy(stored: str | None, traced: bool) -> None: + recorder = _Recorder() + scopes = StoredPolicyTraceScopes( + TraceScopes(emitter=recorder, policies=_NoPolicies(), enterprise_enabled=lambda: False) + ) + + scope = scopes.for_execution( + tenant_id="tenant-1", app_id="app-1", app_tracing=stored, user_id=None, correlation=CORRELATION + ) + + assert (scope is not None) is traced diff --git a/api/tests/unit_tests/services/tracing/test_canonical.py b/api/tests/unit_tests/services/tracing/test_canonical.py new file mode 100644 index 00000000000000..bf6b7cfcdcb7ee --- /dev/null +++ b/api/tests/unit_tests/services/tracing/test_canonical.py @@ -0,0 +1,395 @@ +import os +import subprocess +import sys +from datetime import UTC, datetime, timedelta +from decimal import Decimal +from pathlib import Path + +import pytest +from dify_trace_sdk.ids import span_id_for, trace_id_for +from dify_trace_sdk.model import SpanKind, SpanStatus + +from graphon.entities import WorkflowNodeExecution +from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus +from services.tracing.canonical import CanonicalTraceBuilder, message_trace_id +from services.tracing.records import ( + MessageTraceInfo, + PromptGenerationTraceInfo, + ToolTraceInfo, + TraceMessage, + WorkflowTraceInfo, +) + +STARTED_AT = datetime(2026, 1, 1, 8, 0, 0) +MESSAGE_ID = "6f1c0b52-6ad4-4c4e-9d7c-2a3f0d0c7a11" +CONVERSATION_ID = "0b7f86a3-5f8e-4d51-a4c5-6a8b2f1b6e22" +WORKFLOW_RUN_ID = "4b4b8c2e-1a0f-4bd0-93c1-0f2f0e9d6a33" + + +def _message() -> TraceMessage: + return TraceMessage( + id=MESSAGE_ID, + app_id="app-1", + conversation_id=CONVERSATION_ID, + workflow_run_id=None, + query="What is Dify?", + answer="An LLM app platform.", + error=None, + created_at=STARTED_AT, + model_provider="openai", + model_id="gpt-4o", + total_price=Decimal("0.0042"), + currency="USD", + ) + + +def _message_record(**metadata: object) -> MessageTraceInfo: + return MessageTraceInfo( + message_id=MESSAGE_ID, + message_data=_message(), + inputs=[{"role": "user", "text": "What is Dify?"}], + outputs="An LLM app platform.", + start_time=STARTED_AT, + end_time=STARTED_AT + timedelta(seconds=2), + metadata={"user_id": "account-1", **metadata}, + trace_id=None, + conversation_model="chat", + message_tokens=12, + answer_tokens=8, + total_tokens=20, + error=None, + file_list=[], + conversation_mode="chat", + gen_ai_server_time_to_first_token=0.4, + llm_streaming_time_to_generate=1.2, + is_streaming_request=True, + ) + + +def _node(execution_id: str, node_type: str, *, index: int, offset: int, **fields: object) -> WorkflowNodeExecution: + return WorkflowNodeExecution.model_validate( + { + "id": execution_id, + "workflow_id": "workflow-1", + "workflow_execution_id": WORKFLOW_RUN_ID, + "index": index, + "node_id": f"node-{index}", + "node_type": node_type, + "title": node_type.upper(), + "status": WorkflowNodeExecutionStatus.SUCCEEDED, + "created_at": STARTED_AT + timedelta(seconds=offset), + "finished_at": STARTED_AT + timedelta(seconds=offset + 1), + **fields, + } + ) + + +def _workflow_record(*, message_id: str | None, metadata: dict[str, object]) -> WorkflowTraceInfo: + llm = _node( + "exec-llm", + BuiltinNodeTypes.LLM, + index=2, + offset=1, + process_data={ + "prompts": [{"role": "user", "text": "hi"}], + "model_name": "gpt-4o", + "model_provider": "openai", + "usage": {"prompt_tokens": 5, "completion_tokens": 7, "total_tokens": 12, "time_to_first_token": -1}, + }, + outputs={"text": "hello"}, + metadata={WorkflowNodeExecutionMetadataKey.TOTAL_PRICE: 0.01, WorkflowNodeExecutionMetadataKey.CURRENCY: "USD"}, + ) + start = _node("exec-start", BuiltinNodeTypes.START, index=1, offset=0, inputs={"query": "hi"}) + return WorkflowTraceInfo( + message_id=message_id, + node_executions=[llm, start], + conversation_id=CONVERSATION_ID if message_id else None, + message_started_at=STARTED_AT - timedelta(seconds=1) if message_id else None, + workflow_app_log_id=None, + workflow_id="workflow-1", + tenant_id="tenant-1", + workflow_run_id=WORKFLOW_RUN_ID, + workflow_run_elapsed_time=3, + workflow_run_status="succeeded", + workflow_run_inputs={"query": "hi"}, + workflow_run_outputs={"answer": "hello"}, + workflow_run_version="1", + error=None, + total_tokens=12, + prompt_tokens=5, + completion_tokens=7, + file_list=[], + invoked_by=None, + query="hi" if message_id else "", + start_time=STARTED_AT, + end_time=STARTED_AT + timedelta(seconds=3), + metadata=metadata, + ) + + +def _tool_record(**metadata: object) -> ToolTraceInfo: + return ToolTraceInfo( + message_id=MESSAGE_ID, + message_data=_message(), + tool_name="web_search", + tool_inputs={"q": "dify"}, + tool_outputs="results", + error=None, + tool_config={}, + time_cost=0.5, + tool_parameters={"q": "dify"}, + file_url=None, + start_time=STARTED_AT + timedelta(seconds=1), + end_time=STARTED_AT + timedelta(seconds=2), + metadata=dict(metadata), + ) + + +def test_message_record_becomes_a_root_span_with_one_llm_child() -> None: + trace = CanonicalTraceBuilder().build(_message_record(end_user_session_id="customer-42")) + + assert trace is not None + assert trace.trace_id == message_trace_id(None, MESSAGE_ID) == trace_id_for(MESSAGE_ID) + assert trace.session_id == CONVERSATION_ID + assert trace.user_id == "customer-42" + root, llm = trace.spans + assert (root.span_id, root.parent_span_id, root.name) == (span_id_for(MESSAGE_ID), None, "message") + assert llm.parent_span_id == root.span_id + assert llm.kind is SpanKind.LLM + assert (llm.model_name, llm.model_provider) == ("gpt-4o", "openai") + assert llm.usage is not None + assert (llm.usage.prompt_tokens, llm.usage.completion_tokens) == (12, 8) + assert llm.cost is not None + assert llm.cost.currency == "USD" + assert llm.time_to_first_token == 0.4 + assert root.start_time.tzinfo is UTC + + +def test_caller_trace_id_names_the_trace_and_is_kept_verbatim() -> None: + record = _message_record().model_copy(update={"trace_id": "caller-trace"}) + + trace = CanonicalTraceBuilder().build(record) + + assert trace is not None + assert trace.trace_id == trace_id_for("caller-trace") + assert trace.external_trace_id == "caller-trace" + + +def test_chatflow_run_hangs_under_its_message_and_orders_nodes_parent_first() -> None: + record = _workflow_record(message_id=MESSAGE_ID, metadata={"trace_session_id": "session-1"}) + + trace = CanonicalTraceBuilder().build(record) + + assert trace is not None + names = [span.name for span in trace.spans] + assert names == ["message", "workflow", "start_START", "llm_LLM"] + message_span, workflow_span, _, llm = trace.spans + assert trace.trace_id == trace_id_for(MESSAGE_ID) + assert trace.session_id == "session-1" + assert message_span.start_time == (STARTED_AT - timedelta(seconds=1)).replace(tzinfo=UTC) + assert workflow_span.parent_span_id == message_span.span_id + assert llm.parent_span_id == workflow_span.span_id + assert llm.inputs == [{"role": "user", "text": "hi"}] + assert llm.cost is not None + assert llm.cost.amount == 0.01 + assert llm.time_to_first_token is None + + +def test_nested_workflow_is_its_own_trace_linked_by_attributes() -> None: + record = _workflow_record( + message_id=None, + metadata={ + "parent_trace_context": {"parent_workflow_run_id": "outer-run", "parent_node_execution_id": "outer-node"} + }, + ) + + trace = CanonicalTraceBuilder().build(record) + + assert trace is not None + assert trace.trace_id == trace_id_for(WORKFLOW_RUN_ID) + assert trace.parent is None + assert trace.session_id == "outer-run" + root = trace.root + assert root.name == "workflow" + assert root.attributes["parent_workflow_run_id"] == "outer-run" + assert root.attributes["parent_node_execution_id"] == "outer-node" + + +def test_failed_node_marks_its_span_as_error() -> None: + record = _workflow_record(message_id=None, metadata={}) + failed = record.node_executions[0].model_copy( + update={"status": WorkflowNodeExecutionStatus.FAILED, "error": "model timeout"} + ) + record = record.model_copy(update={"node_executions": [failed, record.node_executions[1]]}) + + trace = CanonicalTraceBuilder().build(record) + + assert trace is not None + llm = next(span for span in trace.spans if span.kind is SpanKind.LLM) + assert (llm.status, llm.error) == (SpanStatus.ERROR, "model timeout") + + +def test_workflow_tool_calls_and_retries_resolve_predecessors_within_each_invocation() -> None: + record = _workflow_record(message_id=None, metadata={}) + nodes = [ + _node("root-start", "start", index=0, offset=0, node_id="start"), + _node("tool-a", "tool", index=1, offset=0), + _node("tool-b", "tool", index=2, offset=0), + ] + parents: dict[str, str] = {} + expected: dict[str, str] = {} + for invocation, caller in (("first", "tool-a"), ("retry", "tool-a"), ("second", "tool-b")): + for node_id, predecessor, node_type in (("start", None, "start"), ("llm", "start", "llm")): + execution_id = f"{invocation}-{node_id}" + nodes.append( + _node( + execution_id, + node_type, + index=3, + offset=1, + node_id=node_id, + predecessor_node_id=predecessor, + process_data={"workflow_tool_invocation_id": invocation}, + ) + ) + parents[execution_id] = caller + expected[execution_id] = f"{invocation}-{predecessor}" if predecessor else caller + nodes.extend( + [ + _node( + "inner-tool", + "tool", + index=4, + offset=2, + predecessor_node_id="llm", + process_data={"workflow_tool_invocation_id": "first"}, + ), + _node("deep-child", "code", index=5, offset=2, node_id="start", status="failed", error="nested failure"), + ] + ) + parents.update({"inner-tool": "tool-a", "deep-child": "inner-tool"}) + expected.update({"inner-tool": "first-llm", "deep-child": "inner-tool"}) + record = record.model_copy(update={"node_executions": nodes, "workflow_tool_parents": parents}) + + trace = CanonicalTraceBuilder().build(record) + + assert trace is not None + spans = {span.span_id: span for span in trace.spans} + for child, parent in expected.items(): + assert spans[span_id_for(child)].parent_span_id == span_id_for(parent) + assert spans[span_id_for("deep-child")].status is SpanStatus.ERROR + assert len(trace.spans) == len(nodes) + 1 + assert CanonicalTraceBuilder().build(record) == trace + + +@pytest.mark.parametrize("container_type", ["iteration", "loop"]) +def test_workflow_tool_container_rounds_stay_with_their_own_invocation(container_type: str) -> None: + record = _workflow_record(message_id=None, metadata={}) + nodes = [_node("caller", "tool", index=0, offset=0)] + parents: dict[str, str] = {} + for invocation in ("first", "second"): + nodes.extend( + ( + _node( + f"{invocation}-container", + container_type, + index=1, + offset=0, + node_id="container", + process_data={"workflow_tool_invocation_id": invocation}, + ), + _node( + f"{invocation}-llm", + "llm", + index=2, + offset=1, + node_id="llm", + process_data={"workflow_tool_invocation_id": invocation}, + metadata={f"{container_type}_id": "container", f"{container_type}_index": 0}, + ), + ) + ) + parents[f"{invocation}-container"] = "caller" + parents[f"{invocation}-llm"] = "caller" + + trace = CanonicalTraceBuilder().build( + record.model_copy(update={"node_executions": nodes, "workflow_tool_parents": parents}) + ) + + assert trace is not None + spans = {span.span_id: span for span in trace.spans} + for invocation in ("first", "second"): + container_id = f"{invocation}-container" + wrapper = spans[span_id_for(f"{container_type}:{container_id}:0")] + assert wrapper.parent_span_id == span_id_for(container_id) + assert spans[span_id_for(container_id)].parent_span_id == span_id_for("caller") + assert spans[span_id_for(f"{invocation}-llm")].parent_span_id == wrapper.span_id + + +def test_tool_call_joins_its_message_trace_as_a_fragment() -> None: + trace = CanonicalTraceBuilder().build(_tool_record(trace_event_id="event-1")) + + assert trace is not None + assert trace.trace_id == trace_id_for(MESSAGE_ID) + assert trace.parent is not None + assert trace.parent.span_id == span_id_for(MESSAGE_ID) + (span,) = trace.spans + assert span.span_id == span_id_for("event:event-1") + assert span.parent_span_id == trace.parent.span_id + assert (span.name, span.kind, span.inputs, span.outputs) == ("web_search", SpanKind.TOOL, {"q": "dify"}, "results") + + +def test_auxiliary_record_without_event_id_is_not_exported() -> None: + assert CanonicalTraceBuilder().build(_tool_record()) is None + + +def test_records_without_a_provider_mapping_are_not_exported() -> None: + record = PromptGenerationTraceInfo( + metadata={}, + tenant_id="tenant-1", + user_id="account-1", + app_id=None, + operation_type="rule_generate", + instruction="write a rule", + prompt_tokens=1, + completion_tokens=1, + total_tokens=2, + model_provider="openai", + model_name="gpt-4o", + latency=0.1, + total_price=None, + currency=None, + error=None, + ) + + assert CanonicalTraceBuilder().build(record) is None + + +def test_tracing_application_imports_do_not_load_runtime_dependencies() -> None: + api_root = Path(__file__).resolve().parents[4] + script = """ +import sys +import services.tracing.canonical +import services.tracing.collection +import services.tracing.service + +forbidden = ( + "celery", "configs", "controllers", "core", "enterprise", "extensions", + "flask", "models", "repositories", "sqlalchemy", "tasks", "werkzeug", +) +loaded = sorted( + name for name in sys.modules + if any(name == prefix or name.startswith(prefix + ".") for prefix in forbidden) +) +assert not loaded, loaded +""" + result = subprocess.run( + [sys.executable, "-c", script], + cwd=api_root, + env={**os.environ, "PYTHONPATH": str(api_root)}, + capture_output=True, + text=True, + timeout=30, + check=False, + ) + assert result.returncode == 0, result.stderr diff --git a/api/tests/unit_tests/services/tracing/test_delivery.py b/api/tests/unit_tests/services/tracing/test_delivery.py new file mode 100644 index 00000000000000..3b80818652068d --- /dev/null +++ b/api/tests/unit_tests/services/tracing/test_delivery.py @@ -0,0 +1,367 @@ +from dataclasses import dataclass, field +from datetime import UTC, datetime, timedelta +from typing import override + +import pytest +from dify_trace_sdk.config import ProviderConfig +from dify_trace_sdk.egress import EgressResponse +from dify_trace_sdk.errors import ExportOutcomeUnknownError, ExportRejectedError, ExportRetryableError +from dify_trace_sdk.events import MessageCompleted, TraceCorrelation, TraceEvent +from dify_trace_sdk.model import Trace +from dify_trace_sdk.provider import ProviderCapabilities, SendResult, TraceProvider +from dify_trace_sdk.testing import RecordingEgress, response +from pydantic import ValidationError + +from services.tracing.canonical import CanonicalTraceBuilder +from services.tracing.delivery import ( + ExportJob, + ExportRetryQueue, + StoredTraceDestination, + TraceCollector, + TraceDeliveryService, + TraceDestinationQuery, + TraceRecordSource, + retry_delay_seconds, +) +from services.tracing.errors import SourceNotReadyError, TraceCollectorUnavailableError +from services.tracing.providers import ProviderConfigCodec +from services.tracing.records import BaseTraceInfo, MessageTraceInfo, TraceMessage, WorkflowTraceInfo +from tests.helpers.unit.services.tracing_ports import InstalledTraceProviders, WorkspaceBoundCipher + +STARTED_AT = datetime(2026, 1, 1, 8, 0, 0) +MESSAGE_ID = "6f1c0b52-6ad4-4c4e-9d7c-2a3f0d0c7a11" +TENANT_ID = "tenant-1" +APP_ID = "app-1" + + +class _Config(ProviderConfig): + endpoint: str + + +class _Provider(TraceProvider[_Config]): + """Posts each trace to its endpoint; the egress script decides the outcome.""" + + id = "langfuse" + config_model = _Config + capabilities = ProviderCapabilities(native_cost=False, trace_url=False, span_url=False) + + @override + def verify(self, config: _Config) -> _Config: + return config + + @override + def project_url(self, config: _Config) -> str | None: + return None + + @override + def send(self, config: _Config, trace: Trace) -> SendResult: + self.egress.request( + "POST", + f"{config.endpoint}/traces", + headers={}, + content=trace.model_dump_json().encode(), + timeout=30.0, + sensitive_url=False, + ) + return SendResult(trace_id=trace.trace_id, trace_url=None, span_urls={}) + + +class _Records(TraceRecordSource): + def __init__(self, result: BaseTraceInfo | Exception) -> None: + self.result = result + self.events: list[TraceEvent] = [] + + @override + def assemble(self, event: TraceEvent) -> BaseTraceInfo: + self.events.append(event) + if isinstance(self.result, Exception): + raise self.result + return self.result + + +class _Destinations(TraceDestinationQuery): + """Stored destinations by (tenant, app); records the owner of every lookup.""" + + def __init__(self, destinations: dict[tuple[str, str], StoredTraceDestination]) -> None: + self.destinations = destinations + self.lookups: list[tuple[str, str]] = [] + + @override + def active_destination(self, *, tenant_id: str, app_id: str) -> StoredTraceDestination | None: + self.lookups.append((tenant_id, app_id)) + return self.destinations.get((tenant_id, app_id)) + + +class _Retries(ExportRetryQueue): + def __init__(self) -> None: + self.scheduled: list[tuple[ExportJob, int]] = [] + + @override + def schedule(self, job: ExportJob, *, delay_seconds: int) -> None: + self.scheduled.append((job, delay_seconds)) + + +class _Collector(TraceCollector): + def __init__(self, *, captures: bool, fails: bool) -> None: + self.captures = captures + self.fails = fails + self.seen: list[tuple[str, BaseTraceInfo]] = [] + + @override + def capture_terminal(self, *, tenant_id: str, info: BaseTraceInfo) -> bool: + if self.fails: + raise RuntimeError("ledger unavailable") + self.seen.append((tenant_id, info)) + return self.captures + + @override + def link_managed_parent(self, *, tenant_id: str, info: WorkflowTraceInfo) -> WorkflowTraceInfo: + return info + + +@dataclass +class _Delivery: + records: _Records + egress: RecordingEgress + destinations: _Destinations + retries: _Retries = field(default_factory=_Retries) + enterprise: list[BaseTraceInfo] = field(default_factory=list) + + def service(self, *, enterprise: bool, collector: _Collector | None) -> TraceDeliveryService: + return TraceDeliveryService( + records=self.records, + destinations=self.destinations, + providers=InstalledTraceProviders(_Provider), + codec=ProviderConfigCodec(WorkspaceBoundCipher()), + builder=CanonicalTraceBuilder(), + egress=self.egress, + retries=self.retries, + enterprise_sink=self.enterprise.append if enterprise else None, + collector=collector, + max_export_attempts=3, + ) + + def sent(self) -> list[Trace]: + return [Trace.model_validate_json(request.content or b"") for request in self.egress.requests] + + +def _event() -> TraceEvent: + return TraceEvent( + event_id="event-1", + tenant_id=TENANT_ID, + app_id=APP_ID, + user_id="account-1", + occurred_at=datetime.now(UTC), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + payload=MessageCompleted(message_id=MESSAGE_ID, conversation_id="conversation-1"), + ) + + +def _record(*, workflow_run_id: str | None) -> MessageTraceInfo: + message = TraceMessage( + id=MESSAGE_ID, + app_id=APP_ID, + conversation_id="conversation-1", + workflow_run_id=workflow_run_id, + query="hi", + answer="hello", + error=None, + created_at=STARTED_AT, + model_provider=None, + model_id=None, + total_price=None, + currency="USD", + ) + return MessageTraceInfo( + message_id=MESSAGE_ID, + message_data=message, + inputs="hi", + outputs="hello", + start_time=STARTED_AT, + end_time=STARTED_AT + timedelta(seconds=1), + metadata={}, + conversation_model="chat", + message_tokens=1, + answer_tokens=1, + total_tokens=2, + error=None, + file_list=[], + conversation_mode="chat", + gen_ai_server_time_to_first_token=None, + llm_streaming_time_to_generate=None, + is_streaming_request=False, + ) + + +def _delivery(record: BaseTraceInfo | Exception, outcomes: list[EgressResponse | Exception]) -> _Delivery: + stored = StoredTraceDestination(provider_id="langfuse", config={"endpoint": "https://collector.example"}) + return _Delivery( + records=_Records(record), + egress=RecordingEgress(outcomes, grpc_target=None), + destinations=_Destinations({(TENANT_ID, APP_ID): stored}), + ) + + +def test_every_sink_receives_the_record_and_the_provider_gets_the_trace() -> None: + delivery = _delivery(_record(workflow_run_id=None), [response(200, {})]) + collector = _Collector(captures=False, fails=False) + + delivery.service(enterprise=True, collector=collector).deliver(_event(), enterprise_delivered=False) + + assert len(delivery.enterprise) == 1 + assert [tenant_id for tenant_id, _ in collector.seen] == [TENANT_ID] + (trace,) = delivery.sent() + assert trace.root.name == "message" + assert delivery.destinations.lookups == [(TENANT_ID, APP_ID)] + assert delivery.retries.scheduled == [] + + +def test_enterprise_sink_is_skipped_on_redelivery() -> None: + delivery = _delivery(_record(workflow_run_id=None), [response(200, {})]) + + delivery.service(enterprise=True, collector=None).deliver(_event(), enterprise_delivered=True) + + assert delivery.enterprise == [] + + +def test_records_the_collector_owns_never_reach_the_provider() -> None: + delivery = _delivery(_record(workflow_run_id=None), []) + + delivery.service(enterprise=False, collector=_Collector(captures=True, fails=False)).deliver( + _event(), enterprise_delivered=False + ) + + assert delivery.sent() == [] + assert delivery.destinations.lookups == [] + + +def test_an_unavailable_collector_asks_for_redelivery() -> None: + delivery = _delivery(_record(workflow_run_id=None), []) + + with pytest.raises(TraceCollectorUnavailableError): + delivery.service(enterprise=False, collector=_Collector(captures=False, fails=True)).deliver( + _event(), enterprise_delivered=False + ) + + +def test_chatflow_message_is_covered_by_its_workflow_trace() -> None: + delivery = _delivery(_record(workflow_run_id="run-1"), []) + + delivery.service(enterprise=False, collector=None).deliver(_event(), enterprise_delivered=False) + + assert delivery.sent() == [] + + +def test_source_not_ready_propagates_for_redelivery() -> None: + delivery = _delivery(SourceNotReadyError(), []) + + with pytest.raises(SourceNotReadyError): + delivery.service(enterprise=False, collector=None).deliver(_event(), enterprise_delivered=False) + + +def test_an_app_without_an_active_destination_sends_nothing() -> None: + delivery = _delivery(_record(workflow_run_id=None), []) + delivery.destinations.destinations.clear() + + delivery.service(enterprise=False, collector=None).deliver(_event(), enterprise_delivered=False) + + assert delivery.sent() == [] + assert delivery.destinations.lookups == [(TENANT_ID, APP_ID)] + + +def test_retryable_failure_schedules_the_next_attempt_until_the_limit() -> None: + delivery = _delivery(_record(workflow_run_id=None), [ExportRetryableError("busy"), ExportRetryableError("busy")]) + service = delivery.service(enterprise=False, collector=None) + + service.deliver(_event(), enterprise_delivered=False) + + ((job, delay),) = delivery.retries.scheduled + assert (job.attempt, job.provider_id, job.tenant_id, job.app_id) == (2, "langfuse", TENANT_ID, APP_ID) + assert delay == retry_delay_seconds(1) + destination = service.destination(TENANT_ID, APP_ID) + assert destination is not None + service.export(job.model_copy(update={"attempt": 3}), destination) + assert len(delivery.retries.scheduled) == 1 + + +@pytest.mark.parametrize("failure", [ExportRejectedError("bad key"), ExportOutcomeUnknownError("timeout")]) +def test_rejected_or_unknown_outcomes_are_never_retried(failure: Exception) -> None: + delivery = _delivery(_record(workflow_run_id=None), [failure]) + + delivery.service(enterprise=False, collector=None).deliver(_event(), enterprise_delivered=False) + + assert delivery.retries.scheduled == [] + + +def test_retry_is_dropped_when_the_app_now_traces_elsewhere() -> None: + trace = CanonicalTraceBuilder().build(_record(workflow_run_id=None)) + assert trace is not None + delivery = _delivery(_record(workflow_run_id=None), []) + + delivery.service(enterprise=False, collector=None).retry( + ExportJob(tenant_id=TENANT_ID, app_id=APP_ID, provider_id="opik", attempt=2, trace=trace) + ) + + assert delivery.sent() == [] + + +def test_retry_reads_the_destination_through_the_jobs_workspace() -> None: + trace = CanonicalTraceBuilder().build(_record(workflow_run_id=None)) + assert trace is not None + delivery = _delivery(_record(workflow_run_id=None), [response(200, {})]) + + delivery.service(enterprise=False, collector=None).retry( + ExportJob(tenant_id=TENANT_ID, app_id=APP_ID, provider_id="langfuse", attempt=2, trace=trace) + ) + + assert delivery.sent() == [trace] + assert delivery.destinations.lookups == [(TENANT_ID, APP_ID)] + + +def test_export_job_round_trips_through_storage_format() -> None: + trace = CanonicalTraceBuilder().build(_record(workflow_run_id=None)) + assert trace is not None + job = ExportJob(tenant_id=TENANT_ID, app_id=APP_ID, provider_id="langfuse", attempt=2, trace=trace) + + assert ExportJob.model_validate_json(job.model_dump_json()) == job + + +def test_retry_delay_grows_and_is_capped() -> None: + assert [retry_delay_seconds(attempt) for attempt in (1, 2, 3, 10)] == [10, 20, 40, 600] + + +def test_records_that_fail_validation_are_dropped_instead_of_failing_the_task() -> None: + with pytest.raises(ValidationError) as invalid: + MessageTraceInfo.model_validate({}) + delivery = _delivery(invalid.value, []) + + delivery.service(enterprise=False, collector=None).deliver(_event(), enterprise_delivered=False) + + assert delivery.sent() == [] + + +def test_unclassified_provider_failures_are_contained() -> None: + delivery = _delivery(_record(workflow_run_id=None), [RuntimeError("provider bug with secret-detail")]) + + delivery.service(enterprise=False, collector=None).deliver(_event(), enterprise_delivered=False) + + assert delivery.retries.scheduled == [] + + +@pytest.mark.parametrize( + "stored", + [ + StoredTraceDestination(provider_id="opik", config={"endpoint": "https://collector.example"}), + StoredTraceDestination(provider_id="langfuse", config={"endpoint": 1}), + ], + ids=["not-installed", "unreadable"], +) +def test_destinations_that_cannot_be_used_send_nothing(stored: StoredTraceDestination) -> None: + delivery = _delivery(_record(workflow_run_id=None), []) + delivery.destinations.destinations[(TENANT_ID, APP_ID)] = stored + + delivery.service(enterprise=False, collector=None).deliver(_event(), enterprise_delivered=False) + + assert delivery.sent() == [] diff --git a/api/tests/unit_tests/services/tracing/test_egress.py b/api/tests/unit_tests/services/tracing/test_egress.py new file mode 100644 index 00000000000000..d4f463a7fa53c6 --- /dev/null +++ b/api/tests/unit_tests/services/tracing/test_egress.py @@ -0,0 +1,200 @@ +import socket +from collections.abc import Callable +from dataclasses import dataclass + +import httpx +import pytest +from dify_trace_sdk.errors import EgressBlockedError, EgressError + +from core.helper import ssrf_proxy +from core.tools.errors import ToolSSRFError +from services.tracing.egress import MAX_RESPONSE_BYTES, SSRFEgress + + +@dataclass(frozen=True) +class _ProxyCall: + max_retries: int + stream_response: bool + timeout: float + follow_redirects: bool + headers: dict[str, str] + + +class _ProxyCalls: + """Stands in for the SSRF proxy owner's ``make_request``.""" + + def __init__(self, outcome: httpx.Response | Exception) -> None: + self.outcome = outcome + self.calls: list[_ProxyCall] = [] + + def __call__( + self, + method: str, + url: str, + *, + max_retries: int, + stream_response: bool, + timeout: float, + follow_redirects: bool, + headers: dict[str, str], + content: bytes | None, + ) -> httpx.Response: + del method, url, content + self.calls.append( + _ProxyCall( + max_retries=max_retries, + stream_response=stream_response, + timeout=timeout, + follow_redirects=follow_redirects, + headers=headers, + ) + ) + if isinstance(self.outcome, Exception): + raise self.outcome + return self.outcome + + +def _request(egress: SSRFEgress, headers: dict[str, str]) -> None: + egress.request( + "POST", + "https://collector.example.com/v1/traces", + headers=headers, + content=b"{}", + timeout=5.0, + sensitive_url=False, + ) + + +def test_requests_go_through_the_proxy_without_retries_redirects_or_ambient_cookies( + monkeypatch: pytest.MonkeyPatch, +) -> None: + proxy = _ProxyCalls(httpx.Response(202, headers={"Content-Type": "application/json"}, content=b'{"ok": true}')) + monkeypatch.setattr(ssrf_proxy, "make_request", proxy) + + result = SSRFEgress().request( + "POST", + "https://collector.example.com/v1/traces", + headers={"Authorization": "Bearer token"}, + content=b"{}", + timeout=5.0, + sensitive_url=False, + ) + + assert (result.status_code, result.json(), result.headers["content-type"]) == ( + 202, + {"ok": True}, + "application/json", + ) + (call,) = proxy.calls + assert (call.max_retries, call.follow_redirects, call.stream_response, call.timeout) == (0, False, True, 5.0) + assert call.headers["Accept-Encoding"] == "identity" + assert call.headers["Cookie"] == "" + assert call.headers["Authorization"] == "Bearer token" + + +def test_sdk_encodings_are_replaced_but_provider_cookies_kept(monkeypatch: pytest.MonkeyPatch) -> None: + proxy = _ProxyCalls(httpx.Response(200, content=b"")) + monkeypatch.setattr(ssrf_proxy, "make_request", proxy) + + _request(SSRFEgress(), {"accept-encoding": "gzip, deflate, br", "cookie": "session=mine"}) + + headers = proxy.calls[0].headers + assert "accept-encoding" not in headers + assert headers["Accept-Encoding"] == "identity" + assert "Cookie" not in headers + assert headers["cookie"] == "session=mine" + + +@pytest.mark.parametrize( + ("failure", "blocked"), + [(httpx.ProxyError("403 Forbidden"), True), (httpx.ProxyError("502 Bad Gateway"), False)], +) +def test_proxy_refusals_are_never_sent(monkeypatch: pytest.MonkeyPatch, failure: Exception, blocked: bool) -> None: + monkeypatch.setattr(ssrf_proxy, "make_request", _ProxyCalls(failure)) + + with pytest.raises(EgressError) as raised: + _request(SSRFEgress(), {}) + + assert isinstance(raised.value, EgressBlockedError) is blocked + assert raised.value.request_sent is False + + +@pytest.mark.parametrize( + ("failure", "expected", "request_sent"), + [ + (ToolSSRFError("blocked"), EgressBlockedError, False), + (httpx.ConnectError("refused"), EgressError, False), + (httpx.ConnectTimeout("slow dns"), EgressError, False), + (httpx.PoolTimeout("pool exhausted"), EgressError, False), + (httpx.ReadTimeout("slow body"), EgressError, True), + (httpx.RemoteProtocolError("reset"), EgressError, True), + ], +) +def test_transport_failures_tell_whether_the_request_was_sent( + monkeypatch: pytest.MonkeyPatch, failure: Exception, expected: type[EgressError], request_sent: bool +) -> None: + monkeypatch.setattr(ssrf_proxy, "make_request", _ProxyCalls(failure)) + + with pytest.raises(expected) as raised: + _request(SSRFEgress(), {}) + + assert raised.value.request_sent is request_sent + assert raised.value.__cause__ is None + assert raised.value.__context__ is None + + +def test_oversized_responses_count_as_sent(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + ssrf_proxy, "make_request", _ProxyCalls(httpx.Response(200, content=b"x" * (MAX_RESPONSE_BYTES + 1))) + ) + + with pytest.raises(EgressError) as raised: + _request(SSRFEgress(), {}) + + assert raised.value.request_sent is True + + +def _resolve_to(address: str) -> Callable[..., list[tuple[object, ...]]]: + def resolve(host: str, port: int, type: int = 0) -> list[tuple[object, ...]]: + del host + return [(socket.AF_INET6 if ":" in address else socket.AF_INET, type, 6, "", (address, port))] + + return resolve + + +@pytest.mark.parametrize("address", ["127.0.0.1", "169.254.169.254", "0.0.0.0", "::1", "::ffff:127.0.0.1"]) +def test_grpc_channels_refuse_non_routable_destinations(monkeypatch: pytest.MonkeyPatch, address: str) -> None: + monkeypatch.setattr(socket, "getaddrinfo", _resolve_to(address)) + + with pytest.raises(EgressBlockedError): + SSRFEgress().grpc_channel("https://collector.example.com:4317") + + +def test_grpc_channels_allow_private_collectors(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(socket, "getaddrinfo", _resolve_to("10.0.0.8")) + + channel = SSRFEgress().grpc_channel("http://collector.internal:4317") + + channel.close() + + +@pytest.mark.parametrize( + "endpoint", ["ftp://collector.example.com", "https://user:pass@collector.example.com", "https://"] +) +def test_grpc_channels_reject_malformed_endpoints(endpoint: str) -> None: + with pytest.raises(EgressBlockedError): + SSRFEgress().grpc_channel(endpoint) + + +def test_unresolvable_grpc_hosts_are_transient(monkeypatch: pytest.MonkeyPatch) -> None: + def fail(host: str, port: int, type: int = 0) -> list[tuple[object, ...]]: + del host, port, type + raise socket.gaierror("no such host") + + monkeypatch.setattr(socket, "getaddrinfo", fail) + + with pytest.raises(EgressError) as raised: + SSRFEgress().grpc_channel("https://collector.example.com") + + assert not isinstance(raised.value, EgressBlockedError) + assert raised.value.request_sent is False diff --git a/api/tests/unit_tests/services/tracing/test_normalization.py b/api/tests/unit_tests/services/tracing/test_normalization.py deleted file mode 100644 index 870b812dc2cbc7..00000000000000 --- a/api/tests/unit_tests/services/tracing/test_normalization.py +++ /dev/null @@ -1,218 +0,0 @@ -"""Content policy and model-call accounting across all canonical tree layers.""" - -import json -from datetime import UTC, datetime - -import pytest - -from core.helper.trace_id_helper import ParentTraceContext -from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace -from services.tracing.normalization import normalize_agent_span, order_spans, remap_outer_trace, safe_attributes -from services.tracing.provider_common import token_usage, usd_cost - - -def agent_data( - attributes: dict[str, object], *, span_id: str = "2" * 16, parent_id: str = "1" * 16 -) -> dict[str, object]: - return { - "span_id": span_id, - "parent_span_id": parent_id, - "name": "operation", - "start_time_ns": 1_000_000_000, - "end_time_ns": 2_000_000_000, - "status": "error", - "status_message": "sensitive error output", - "attributes": attributes, - } - - -def test_content_off_is_safe_for_all_metadata_namespaces() -> None: - normalized = normalize_agent_span( - agent_data( - { - "gen_ai.operation.name": "chat", - "gen_ai.request.model": "model-one", - "gen_ai.usage.input_tokens": 3, - "gen_ai.usage.output_tokens": 2, - "gen_ai.input.messages": [{"content": "private input"}], - "gen_ai.output.messages": [{"content": "private output"}], - "exception.message": "private exception", - "dify.app_id": "app-one", - "dify.custom_payload": {"prompt": "private metadata"}, - "gen_ai.request.messages": "private request metadata", - "gen_ai.response.output": "private response metadata", - "gen_ai.usage.custom_prompt": "private usage metadata", - } - ), - capture_content=False, - ) - assert normalized.inputs is None - assert normalized.outputs is None - assert "private" not in normalized.model_dump_json() - assert "sensitive error" not in normalized.model_dump_json() - assert normalized.metadata["dify.app_id"] == "app-one" - assert normalized.metadata["model_name"] == "model-one" - assert normalized.metadata["total_tokens"] == 5 - - -def test_json_encoded_tool_credentials_are_redacted_from_inputs_and_metadata() -> None: - normalized = normalize_agent_span( - agent_data( - { - "gen_ai.operation.name": "execute_tool", - "gen_ai.tool.name": "search", - "gen_ai.tool.call.arguments": json.dumps({"api_key": "tool-secret-123", "query": "public query"}), - "gen_ai.tool.call.result": json.dumps( - {"nested": {"authorization": "bearer-secret-123"}, "result": "okay"} - ), - } - ), - capture_content=True, - ) - serialized = normalized.model_dump_json() - assert "tool-secret-123" not in serialized - assert "bearer-secret-123" not in serialized - assert "public query" in serialized - assert normalized.inputs["api_key"] == "[redacted]" - - -def test_content_off_drops_error_event_payload_and_keeps_numeric_usage() -> None: - attrs = safe_attributes( - { - "input": "user text", - "authorization": "secret", - "exception.stacktrace": "sensitive", - "gen_ai.usage.input_tokens": 10, - "gen_ai.tool.name": "lookup", - }, - capture_content=False, - ) - assert attrs == {"gen_ai.usage.input_tokens": 10, "gen_ai.tool.name": "lookup"} - - -def test_real_runtime_accounting_attributes_survive_content_off() -> None: - normalized = normalize_agent_span( - agent_data( - { - "gen_ai.operation.name": "chat", - "gen_ai.request.model": "gpt-4o-mini", - "gen_ai.usage.input_tokens": 8, - "gen_ai.usage.output_tokens": 4, - "dify.llm.total_price": "0.03", - "dify.llm.currency": "USD", - "gen_ai.response.time_to_first_token": 0.2, - "dify.llm.time_to_first_token_unit": "s", - "dify.llm.time_to_first_token_source": "gateway_usage", - "gen_ai.input.messages": "private input", - } - ), - capture_content=False, - ) - assert usd_cost(normalized) == 0.03 - assert normalized.metadata["currency"] == "USD" - assert normalized.metadata["time_to_first_token"] == 0.2 - assert token_usage(normalized) == {"input": 8, "output": 4, "total": 12} - assert normalized.inputs is None - - -def span(identifier: str, parent: str | None, kind: CanonicalSpanKind, *, synthetic: bool = False) -> CanonicalSpan: - return CanonicalSpan( - id=identifier, - parent_id=parent, - name=identifier, - kind=kind, - start_time=datetime.fromtimestamp(0, UTC), - end_time=datetime.fromtimestamp(3, UTC), - status=CanonicalSpanStatus.OK, - inputs={"content": "private"}, - outputs="private", - metadata={ - "prompt_tokens": 6, - "completion_tokens": 4, - "total_tokens": 10, - "total_price": "0.03", - "prompt": "private", - }, - synthetic=synthetic, - ) - - -def test_two_native_model_calls_replace_synthetic_aggregate_generation() -> None: - outer = CanonicalTrace( - trace_id="old", - session_id="session", - root_span_id="message", - spans=( - span("message", None, CanonicalSpanKind.CHAIN), - span("message:llm", "message", CanonicalSpanKind.LLM, synthetic=True), - ), - ) - remapped = remap_outer_trace(outer, trace_id="1" * 32, capture_content=False, agent_message=True) - calls = [ - normalize_agent_span( - agent_data( - { - "gen_ai.operation.name": "chat", - "gen_ai.request.model": "model", - "gen_ai.usage.input_tokens": 3, - "gen_ai.usage.output_tokens": 2, - }, - span_id=identifier, - parent_id=remapped.root_span_id, - ), - capture_content=False, - ) - for identifier in ("2" * 16, "3" * 16) - ] - combined = order_spans(remapped.root_span_id, [calls[1], *remapped.spans, calls[0]]) - assert combined[0].id == remapped.root_span_id - generations = [item for item in combined if item.kind is CanonicalSpanKind.LLM] - assert len(generations) == 2 - assert sum(token_usage(item)["total"] for item in generations) == 10 - assert token_usage(combined[0]) == {} - assert combined[0].metadata["aggregate_total_tokens"] == 10 - assert "private" not in json.dumps([item.model_dump(mode="json") for item in combined]) - - -def test_non_agent_workflow_retains_its_real_llm_node() -> None: - trace = CanonicalTrace( - trace_id="old", - session_id="session", - root_span_id="workflow", - spans=( - span("workflow", None, CanonicalSpanKind.CHAIN), - span("llm-node", "workflow", CanonicalSpanKind.LLM), - ), - ) - remapped = remap_outer_trace(trace, trace_id="1" * 32, capture_content=True, agent_message=False) - assert len(remapped.spans) == 2 - assert remapped.spans[1].kind is CanonicalSpanKind.LLM - assert token_usage(remapped.spans[1])["total"] == 10 - - -def test_order_spans_rejects_orphans_cycles_and_conflicting_ids() -> None: - root = span("root", None, CanonicalSpanKind.CHAIN) - with pytest.raises(ValueError, match="parent is missing"): - order_spans("root", [root, span("orphan", "missing", CanonicalSpanKind.TOOL)]) - with pytest.raises(ValueError, match="parent cycle"): - order_spans( - "root", [root, span("one", "two", CanonicalSpanKind.TOOL), span("two", "one", CanonicalSpanKind.TOOL)] - ) - with pytest.raises(ValueError, match="Conflicting"): - order_spans("root", [root, root.model_copy(update={"name": "other"})]) - - -def test_managed_external_parent_is_preserved_as_explicit_link_when_content_off() -> None: - trace = CanonicalTrace( - trace_id="old", - session_id="session", - root_span_id="workflow", - spans=(span("workflow", None, CanonicalSpanKind.CHAIN),), - external_parent=ParentTraceContext( - parent_workflow_run_id="outer-workflow", parent_node_execution_id="outer-tool" - ), - ) - remapped = remap_outer_trace(trace, trace_id="1" * 32, capture_content=False, agent_message=False) - assert remapped.external_parent is None - assert remapped.spans[0].metadata["linked_parent_workflow_run_id"] == "outer-workflow" - assert remapped.spans[0].metadata["linked_parent_node_execution_id"] == "outer-tool" diff --git a/api/tests/unit_tests/services/tracing/test_otlp_export.py b/api/tests/unit_tests/services/tracing/test_otlp_export.py deleted file mode 100644 index d1c028d324f51d..00000000000000 --- a/api/tests/unit_tests/services/tracing/test_otlp_export.py +++ /dev/null @@ -1,241 +0,0 @@ -"""Exercise Aliyun's observed HTTP ACK through both real protobuf owners.""" - -import logging -from collections.abc import Mapping -from datetime import UTC, datetime, timedelta -from typing import Literal, TypedDict, Unpack - -import httpx -import pytest -from dify_trace_aliyun.data_exporter.traceclient import TraceClient -from dify_trace_aliyun.entities.aliyun_trace_entity import SpanData -from opentelemetry.context import _SUPPRESS_INSTRUMENTATION_KEY, get_value -from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( - ExportTraceServiceRequest, - ExportTraceServiceResponse, -) - -from core.helper import ssrf_proxy -from core.ops.exceptions import RetryableTraceDispatchError -from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace -from services.tracing.otlp_export import export_http, is_aliyun_http_success_ack -from services.tracing.provider_dispatch import TraceExportRejectedError, TraceExportUncertainError, export_trace - -ENDPOINT = "https://collector.test/adapt_synthetic-license/api/otlp/traces" -START = datetime(2026, 10, 19, tzinfo=UTC) -TRACE_ID = "1" * 32 -SPAN_ID = "2" * 16 - -type Owner = Literal["managed", "legacy"] - - -class RequestOptions(TypedDict, total=False): - content: bytes - headers: Mapping[str, str] - max_retries: int - timeout: float - follow_redirects: bool - - -def boundary( - monkeypatch: pytest.MonkeyPatch, status: int, content: bytes, content_type: str -) -> list[ExportTraceServiceRequest]: - received: list[ExportTraceServiceRequest] = [] - - def send(method: str, url: str, **options: Unpack[RequestOptions]) -> httpx.Response: - assert method == "POST" - assert url == ENDPOINT - assert options["max_retries"] == 0 - assert options["follow_redirects"] is False - assert options["headers"] == {"Content-Type": "application/x-protobuf", "Accept": "application/x-protobuf"} - assert get_value(_SUPPRESS_INSTRUMENTATION_KEY) is True - received.append(ExportTraceServiceRequest.FromString(options["content"])) - headers = {"Content-Type": content_type} if content_type else {} - return httpx.Response(status, content=content, headers=headers) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - return received - - -def dispatch(owner: Owner) -> None: - if owner == "managed": - trace = CanonicalTrace( - trace_id=TRACE_ID, - root_span_id=SPAN_ID, - session_id="synthetic-conversation", - spans=( - CanonicalSpan( - id=SPAN_ID, - parent_id=None, - name="synthetic root", - kind=CanonicalSpanKind.AGENT, - start_time=START, - end_time=START + timedelta(seconds=1), - status=CanonicalSpanStatus.OK, - ), - ), - ) - result = export_trace( - "aliyun", - {"endpoint": "https://collector.test", "license_key": "synthetic-license", "app_name": "synthetic"}, - trace, - ) - assert result.status == "accepted" - else: - client = TraceClient("synthetic", ENDPOINT) - with client.batch(): - client.add_span( - SpanData( - trace_id=int(TRACE_ID, 16), - span_id=int(SPAN_ID, 16), - name="synthetic root", - start_time=int(START.timestamp()) * 1_000_000_000, - end_time=int((START + timedelta(seconds=1)).timestamp()) * 1_000_000_000, - ) - ) - assert client._dispatch_spans.get() is None - - -def assert_one_batch(received: list[ExportTraceServiceRequest]) -> None: - assert len(received) == 1 - spans = received[0].resource_spans[0].scope_spans[0].spans - assert len(spans) == 1 - assert spans[0].trace_id.hex() == TRACE_ID - assert spans[0].span_id.hex() == SPAN_ID - - -@pytest.mark.parametrize("owner", ["managed", "legacy"]) -@pytest.mark.parametrize("content_type", ["", "application/json", "application/x-protobuf"]) -def test_exact_observed_ack_accepts_one_batch_without_changing_its_ids( - monkeypatch: pytest.MonkeyPatch, owner: Owner, content_type: str -) -> None: - received = boundary(monkeypatch, 200, b'"success"', content_type) - dispatch(owner) - assert_one_batch(received) - - -@pytest.mark.parametrize("owner", ["managed", "legacy"]) -@pytest.mark.parametrize( - ("content", "content_type"), - [ - (b"success", ""), - (b'"SUCCESS"', ""), - (b' "success"', "application/json"), - (b'"success"\n', "application/json"), - (b'{"success":true}', "application/json"), - (b'{"message":"success"}', "application/json"), - (b"true", "application/json"), - (b"not protobuf", ""), - ], -) -def test_other_ack_bodies_are_uncertain_and_not_replayed( - monkeypatch: pytest.MonkeyPatch, owner: Owner, content: bytes, content_type: str -) -> None: - received = boundary(monkeypatch, 200, content, content_type) - with pytest.raises(TraceExportUncertainError, match="unreadable ingestion acknowledgment") as error: - dispatch(owner) - assert error.value.acceptance_unknown is True - assert_one_batch(received) - - -@pytest.mark.parametrize("owner", ["managed", "legacy"]) -@pytest.mark.parametrize("status", [201, 202, 204]) -def test_literal_ack_does_not_accept_another_success_status( - monkeypatch: pytest.MonkeyPatch, owner: Owner, status: int -) -> None: - received = boundary(monkeypatch, status, b'"success"', "") - assert is_aliyun_http_success_ack(status, b'"success"') is False - with pytest.raises(TraceExportUncertainError) as error: - dispatch(owner) - assert error.value.acceptance_unknown is True - assert_one_batch(received) - - -@pytest.mark.parametrize("owner", ["managed", "legacy"]) -@pytest.mark.parametrize("status", [403, 429, 503]) -def test_literal_ack_does_not_override_http_rejection_or_known_retry( - monkeypatch: pytest.MonkeyPatch, owner: Owner, status: int -) -> None: - received = boundary(monkeypatch, status, b'"success"', "") - expected = TraceExportRejectedError if status == 403 else RetryableTraceDispatchError - with pytest.raises(expected) as error: - dispatch(owner) - if isinstance(error.value, RetryableTraceDispatchError): - assert error.value.acceptance_unknown is False - assert_one_batch(received) - - -@pytest.mark.parametrize("owner", ["managed", "legacy"]) -@pytest.mark.parametrize("encoding", ["protobuf", "json"]) -@pytest.mark.parametrize("rejected_spans", [0, 1]) -def test_standard_otlp_ack_and_partial_success_remain_validated( - monkeypatch: pytest.MonkeyPatch, owner: Owner, encoding: str, rejected_spans: int -) -> None: - if encoding == "protobuf": - response = ExportTraceServiceResponse() - response.partial_success.rejected_spans = rejected_spans - response.partial_success.error_message = "synthetic warning" - content = response.SerializeToString() - content_type = "application/x-protobuf" - else: - content = b'{"partialSuccess":{"rejectedSpans":"%d","errorMessage":"synthetic warning"}}' % rejected_spans - content_type = "application/json" - received = boundary(monkeypatch, 200, content, content_type) - if rejected_spans: - with pytest.raises(TraceExportUncertainError, match="partially accepted") as error: - dispatch(owner) - assert error.value.acceptance_unknown is True - else: - dispatch(owner) - assert_one_batch(received) - - -@pytest.mark.parametrize("owner", ["managed", "legacy"]) -def test_damaged_compressed_ack_is_uncertain_after_real_http_submission( - monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, owner: Owner -) -> None: - received: list[ExportTraceServiceRequest] = [] - previous_suppression = get_value(_SUPPRESS_INSTRUMENTATION_KEY) - - def send(wire: httpx.Request) -> httpx.Response: - assert wire.method == "POST" - assert str(wire.url) == ENDPOINT - assert get_value(_SUPPRESS_INSTRUMENTATION_KEY) is True - received.append(ExportTraceServiceRequest.FromString(wire.content)) - # Reading a malformed compressed response raises DecodingError, which - # is a RequestError but not a TransportError. Submission already ran. - return httpx.Response( - 200, - headers={"Content-Encoding": "gzip"}, - stream=httpx.ByteStream(b"invalid gzip acknowledgment"), - request=wire, - ) - - with httpx.Client(transport=httpx.MockTransport(send)) as client: - - def get_client(ssl_verify_enabled: bool) -> httpx.Client: - assert isinstance(ssl_verify_enabled, bool) - return client - - monkeypatch.setattr(ssrf_proxy, "_get_ssrf_client", get_client) - with caplog.at_level(logging.INFO, logger="httpx"): - with pytest.raises(TraceExportUncertainError) as error: - dispatch(owner) - assert error.value.acceptance_unknown is True - assert error.value.__cause__ is None - assert error.value.__context__ is None - assert "synthetic-license" not in str(error.value) + caplog.text - assert get_value(_SUPPRESS_INSTRUMENTATION_KEY) == previous_suppression - assert_one_batch(received) - - -@pytest.mark.parametrize("provider", ["langfuse", "phoenix", "tencent"]) -@pytest.mark.parametrize("content_type", ["", "application/json"]) -def test_literal_success_ack_does_not_expand_other_providers( - monkeypatch: pytest.MonkeyPatch, provider: str, content_type: str -) -> None: - received = boundary(monkeypatch, 200, b'"success"', content_type) - with pytest.raises(TraceExportUncertainError, match="unreadable ingestion acknowledgment") as error: - export_http(provider, ENDPOINT, {}, ExportTraceServiceRequest()) - assert error.value.acceptance_unknown is True - assert len(received) == 1 diff --git a/api/tests/unit_tests/services/tracing/test_provider_configs.py b/api/tests/unit_tests/services/tracing/test_provider_configs.py new file mode 100644 index 00000000000000..95eea6f1329bf7 --- /dev/null +++ b/api/tests/unit_tests/services/tracing/test_provider_configs.py @@ -0,0 +1,253 @@ +import pytest +from dify_trace_arize_phoenix.provider import PhoenixProvider +from dify_trace_langfuse.config import LangfuseConfig +from dify_trace_langfuse.provider import LangfuseProvider +from dify_trace_sdk.egress import Egress, EgressResponse +from dify_trace_sdk.testing import RecordingEgress, response + +from core.helper.encrypter import obfuscated_token +from services import app_tracing_config_service as config_service +from services.tracing.provider_configs import ProviderConfigGateway, TraceProviderCapabilitiesQuery +from services.tracing.providers import ProviderConfigCodec, TraceProviderSupport, configured_secrets +from tests.helpers.unit.services.tracing_ports import InstalledTraceProviders, WorkspaceBoundCipher + +WORKSPACE = "tenant-1" +PUBLIC_KEY = "pk-lf-0123456789" +SECRET_KEY = "sk-lf-abcdefghijk" +CIPHER = WorkspaceBoundCipher() + + +def _encrypt(tenant_id: str, token: str) -> str: + return CIPHER.encrypt(tenant_id, token) + + +def _gateway(egress: Egress) -> ProviderConfigGateway: + return ProviderConfigGateway( + providers=InstalledTraceProviders(LangfuseProvider, PhoenixProvider), + codec=ProviderConfigCodec(WorkspaceBoundCipher()), + egress=egress, + ) + + +def _projects(project_id: str) -> EgressResponse: + return response(200, {"data": [{"id": project_id}]}) + + +def _stored(**overrides: object) -> dict[str, object]: + return { + "public_key": _encrypt(WORKSPACE, PUBLIC_KEY), + "secret_key": _encrypt(WORKSPACE, SECRET_KEY), + "host": "https://langfuse.example.com", + "project_key": "project-1", + **overrides, + } + + +def test_presentation_never_returns_secrets_or_derived_fields() -> None: + gateway = _gateway(RecordingEgress([], grpc_target=None)) + + presented = gateway.present_config(workspace_id=WORKSPACE, tracing_provider="langfuse", tracing_config=_stored()) + + assert presented.values == { + "host": "https://langfuse.example.com", + "project_url": "https://langfuse.example.com/project/project-1", + } + assert presented.configured_secrets == ("public_key", "secret_key") + assert configured_secrets(LangfuseProvider, _stored(secret_key="")) == ("public_key",) + + +def test_legacy_config_is_presented_and_decoded_without_resaving_or_verification() -> None: + egress = RecordingEgress([], grpc_target=None) + stored = _stored(project_url="https://langfuse.example.com/project/legacy-project") + stored.pop("project_key") + original = stored.copy() + + presented = _gateway(egress).present_config( + workspace_id=WORKSPACE, tracing_provider="langfuse", tracing_config=stored + ) + decoded = ProviderConfigCodec(CIPHER).decode(WORKSPACE, LangfuseProvider, stored) + + assert isinstance(decoded, LangfuseConfig) + assert presented.values == {"host": stored["host"], "project_url": stored["project_url"]} + assert presented.configured_secrets == ("public_key", "secret_key") + assert decoded.model_dump()["secret_key"] == SECRET_KEY + assert LangfuseProvider(egress).project_url(decoded) == stored["project_url"] + assert stored == original + assert egress.requests == [] + + +def test_new_configs_are_verified_and_store_derived_fields() -> None: + egress = RecordingEgress([_projects("project-9")], grpc_target=None) + gateway = _gateway(egress) + + stored = gateway.prepare_new_config( + workspace_id=WORKSPACE, + tracing_provider="langfuse", + tracing_config={ + "public_key": PUBLIC_KEY, + "secret_key": SECRET_KEY, + "host": "https://langfuse.example.com", + "project_key": "client-supplied", + "project_url": "https://client-supplied.example/project/wrong", + }, + ) + + assert stored["project_key"] == "project-9" + assert stored["project_url"] is None + assert stored["secret_key"] == _encrypt(WORKSPACE, SECRET_KEY) + (request,) = egress.requests + assert request.url == "https://langfuse.example.com/api/public/projects" + + +@pytest.mark.parametrize( + "submitted_secret", + [None, "", obfuscated_token(SECRET_KEY)], + ids=["omitted", "blank", "masked"], +) +def test_updates_keep_a_secret_that_was_not_re_entered(submitted_secret: str | None) -> None: + gateway = _gateway(RecordingEgress([_projects("project-1")], grpc_target=None)) + submitted: dict[str, object] = {"public_key": PUBLIC_KEY, "host": "https://langfuse.example.com"} + if submitted_secret is not None: + submitted["secret_key"] = submitted_secret + + stored = gateway.prepare_updated_config( + workspace_id=WORKSPACE, + tracing_provider="langfuse", + tracing_config=submitted, + current_tracing_config=_stored(), + ) + + assert stored["secret_key"] == _encrypt(WORKSPACE, SECRET_KEY) + + +def test_updates_replace_a_re_entered_secret() -> None: + gateway = _gateway(RecordingEgress([_projects("project-1")], grpc_target=None)) + + stored = gateway.prepare_updated_config( + workspace_id=WORKSPACE, + tracing_provider="langfuse", + tracing_config={ + "public_key": PUBLIC_KEY, + "secret_key": "sk-lf-rotated", + "host": "https://langfuse.example.com", + }, + current_tracing_config=_stored(), + ) + + assert stored["secret_key"] == _encrypt(WORKSPACE, "sk-lf-rotated") + + +def test_failed_verification_is_reported_without_details() -> None: + gateway = _gateway(RecordingEgress([response(401, {"message": "bad key"})], grpc_target=None)) + + with pytest.raises(config_service.AppTracingConfigVerificationFailedError): + gateway.prepare_new_config( + workspace_id=WORKSPACE, + tracing_provider="langfuse", + tracing_config={"public_key": PUBLIC_KEY, "secret_key": SECRET_KEY}, + ) + + +def test_invalid_configs_are_rejected_before_any_request() -> None: + egress = RecordingEgress([], grpc_target=None) + + with pytest.raises(config_service.AppTracingConfigInvalidConfigurationError): + _gateway(egress).prepare_new_config( + workspace_id=WORKSPACE, + tracing_provider="langfuse", + tracing_config={"public_key": PUBLIC_KEY, "secret_key": SECRET_KEY, "host": "ftp://langfuse.example.com"}, + ) + assert egress.requests == [] + + +def test_unknown_providers_are_rejected() -> None: + gateway = _gateway(RecordingEgress([], grpc_target=None)) + + with pytest.raises(config_service.AppTracingConfigInvalidProviderError): + gateway.validate_provider("not-a-provider") + + +def test_capabilities_describe_the_provider_mapping_and_the_collection_gate() -> None: + capabilities = TraceProviderCapabilitiesQuery( + providers=InstalledTraceProviders(LangfuseProvider), api_collection_enabled=True + ).get("langfuse") + + assert capabilities.support == TraceProviderSupport( + provider="langfuse", native_cost=True, trace_url=True, span_url=True + ) + assert capabilities.support.export_ack == "accepted" + assert capabilities.api_collection_enabled is True + + +@pytest.mark.parametrize( + ("provider_id", "error"), + [ + ("datadog", config_service.AppTracingConfigInvalidProviderError), + ("opik", config_service.AppTracingConfigProviderUnavailableError), + ], +) +def test_capabilities_of_unknown_or_uninstalled_providers_are_errors(provider_id: str, error: type[Exception]) -> None: + query = TraceProviderCapabilitiesQuery( + providers=InstalledTraceProviders(LangfuseProvider), api_collection_enabled=False + ) + + with pytest.raises(error): + query.get(provider_id) + + +def test_uninstalled_providers_can_still_be_named_but_not_presented() -> None: + gateway = ProviderConfigGateway( + providers=InstalledTraceProviders(), + codec=ProviderConfigCodec(WorkspaceBoundCipher()), + egress=RecordingEgress([], grpc_target=None), + ) + + gateway.validate_provider("langfuse") + with pytest.raises(config_service.AppTracingConfigProviderUnavailableError): + gateway.present_config(workspace_id=WORKSPACE, tracing_provider="langfuse", tracing_config=_stored()) + + +@pytest.mark.parametrize( + "submitted_secret", [None, "", obfuscated_token(SECRET_KEY)], ids=["omitted", "blank", "masked"] +) +def test_stored_secrets_never_follow_an_edited_destination(submitted_secret: str | None) -> None: + egress = RecordingEgress([], grpc_target=None) + submitted: dict[str, object] = {"public_key": PUBLIC_KEY, "host": "https://collector.attacker.example"} + if submitted_secret is not None: + submitted["secret_key"] = submitted_secret + + with pytest.raises(config_service.AppTracingConfigInvalidConfigurationError): + _gateway(egress).prepare_updated_config( + workspace_id=WORKSPACE, + tracing_provider="langfuse", + tracing_config=submitted, + current_tracing_config=_stored(), + ) + assert egress.requests == [] + + +def test_an_edited_destination_accepts_re_entered_secrets() -> None: + egress = RecordingEgress([_projects("project-2")], grpc_target=None) + + stored = _gateway(egress).prepare_updated_config( + workspace_id=WORKSPACE, + tracing_provider="langfuse", + tracing_config={"public_key": PUBLIC_KEY, "secret_key": "sk-lf-new", "host": "https://eu.langfuse.example"}, + current_tracing_config=_stored(), + ) + + assert stored["secret_key"] == _encrypt(WORKSPACE, "sk-lf-new") + assert egress.requests[0].url == "https://eu.langfuse.example/api/public/projects" + + +def test_an_unchanged_destination_with_a_trailing_slash_keeps_secrets() -> None: + egress = RecordingEgress([_projects("project-1")], grpc_target=None) + + stored = _gateway(egress).prepare_updated_config( + workspace_id=WORKSPACE, + tracing_provider="langfuse", + tracing_config={"public_key": PUBLIC_KEY, "host": "https://langfuse.example.com/"}, + current_tracing_config=_stored(), + ) + + assert stored["secret_key"] == _encrypt(WORKSPACE, SECRET_KEY) diff --git a/api/tests/unit_tests/services/tracing/test_provider_dispatch.py b/api/tests/unit_tests/services/tracing/test_provider_dispatch.py deleted file mode 100644 index 9af18badcf18cd..00000000000000 --- a/api/tests/unit_tests/services/tracing/test_provider_dispatch.py +++ /dev/null @@ -1,2846 +0,0 @@ -import atexit -import base64 -import json -import os -import threading -from collections.abc import Callable, Iterator, Mapping, Sequence -from concurrent.futures import ThreadPoolExecutor -from copy import deepcopy -from dataclasses import dataclass -from datetime import UTC, datetime, timedelta -from decimal import Decimal -from pathlib import Path -from threading import Barrier, Thread -from typing import Never, NotRequired, TypedDict, Unpack, cast -from urllib.parse import parse_qs, urlsplit -from uuid import UUID - -import httpcore -import httpx -import pytest -import requests -import wandb -import weave -from aliyun.log.logclient import LogClient -from aliyun.log.util import Util -from dify_trace_aliyun.aliyun_trace import AliyunDataTrace -from dify_trace_aliyun.config import AliyunConfig -from dify_trace_aliyun.data_exporter.traceclient import TraceClient, convert_to_span_id -from dify_trace_aliyun.entities.aliyun_trace_entity import SpanData -from dify_trace_aliyun.entities.semconv import GEN_AI_USAGE_TOTAL_TOKENS -from dify_trace_arize_phoenix.arize_phoenix_trace import setup_tracer -from dify_trace_arize_phoenix.config import PhoenixConfig -from dify_trace_weave.config import WeaveConfig -from dify_trace_weave.entities.weave_trace_entity import WeaveTraceModel -from dify_trace_weave.weave_trace import WeaveDataTrace -from mlflow.entities import Span as _native_weave_mlflow_Span -from mlflow.entities import TraceData -from mlflow.tracing.constant import SpanAttributeKey -from mlflow.tracing.otel.translation import translate_span_when_storing -from mlflow.types.chat import ChatTool -from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter -from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( - ExportTraceServiceRequest, - ExportTraceServiceResponse, -) -from opentelemetry.proto.trace.v1.trace_pb2 import Span as _provider_dispatch_Span -from sqlalchemy import Connection, Engine, event, select -from sqlalchemy.orm import Session, sessionmaker -from weave.trace_server.trace_server_interface import CallEndReq, CallStartReq, StartedCallSchemaForInsert, TraceStatus - -from core.helper import ssrf_proxy -from core.ops.entities.trace_entity import ( - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - WorkflowTraceInfo, -) -from core.ops.exceptions import RetryableTraceDispatchError -from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace -from extensions.logstore.aliyun_logstore import AliyunLogStore -from extensions.logstore.repositories.logstore_api_workflow_node_execution_repository import ( - _dict_to_workflow_node_execution_model, -) -from extensions.logstore.repositories.logstore_api_workflow_run_repository import _dict_to_workflow_run -from graphon.enums import BuiltinNodeTypes, WorkflowExecutionStatus, WorkflowNodeExecutionStatus -from libs.datetime_utils import naive_utc_now -from models.account import TenantAccountJoin, TenantAccountRole -from models.application_trace import ApplicationTraceRun -from models.enums import ConversationFromSource, CreatorUserRole, MessageStatus -from models.model import App, AppMode -from models.workflow import ( - WorkflowNodeExecutionModel, - WorkflowNodeExecutionTriggeredFrom, - WorkflowRun, - WorkflowRunTriggeredFrom, - WorkflowType, -) -from repositories import application_trace_repository as repository_module -from repositories.application_trace_repository import TraceCollectionError -from services.tracing import service as service_module -from services.tracing.normalization import remap_outer_trace -from services.tracing.provider_dispatch import ( - TraceExportRejectedError, - TraceExportResult, - TraceExportUncertainError, - export_trace, - provider_capabilities, -) -from services.tracing.service import ApplicationTracingService -from tests.helpers.unit.config_override import apply_config_overrides -from tests.helpers.unit.model_factories import make_account, make_app, make_message, make_tenant -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_batch_for as batch_for -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_claim as claim -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_identifier as identifier -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_ingest as ingest -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_seed as seed -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_terminal as terminal -from tests.unit_tests.services.tracing.test_service import _ApplicationTraceCollectionCollection as Collection - -"""Node logs resolve the caller's trace using SQL or configured LogStore ownership.""" - - -@dataclass -class NodeAlias: - case: Collection - node: dict[str, object] - workflow: dict[str, object] - requests: list[str] - logstore: bool - - def service(self) -> ApplicationTracingService: - return ApplicationTracingService(self.case.repository) - - def status(self, node_execution_id: str | None = None) -> dict[str, object]: - return self.service().status( - tenant_id=self.case.run.tenant_id, - app_id=self.case.run.app_id, - source_kind="node", - source_id=str(self.node["id"]), - node_execution_id=node_execution_id, - ) - - def accept(self, *, spans: bool = True) -> tuple[str, str]: - span_urls: dict[str, str] = {} - for attempt in (0, 1): - invocation = self.case.repository.invocation( - run_id=self.case.run.id, - agent_id=identifier("referenced-agent"), - node_execution_id=str(self.node["id"]), - attempt=attempt, - call_key=f"attempt-{attempt}", - ) - ingest(self.case, batch_for(self.case.run, invocation, span_id=str(attempt + 1) * 16)) - span_urls[invocation.span_id] = f"https://provider.test/span/{attempt}" - terminal(self.case) - lease = claim(self.case) - assert lease is not None - self.case.repository.complete( - self.case.run.id, - lease[1], - state="accepted", - result=TraceExportResult( - "provider-trace", "https://provider.test/trace", span_urls=span_urls if spans else {} - ), - ) - return "https://provider.test/trace", "https://provider.test/span/1" - - -class TestApplicationTraceNodeAliasCases: - @staticmethod - @pytest.fixture(params=["sql", "logstore"]) - def alias( - request: pytest.FixtureRequest, - sqlite_engine: Engine, - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - ) -> Iterator[NodeAlias]: - case = seed(sqlite_session_factory, "node-alias") - with case.sessions.begin() as session: - app = session.get(App, case.run.app_id) - assert app is not None - app.mode = AppMode.WORKFLOW - common: dict[str, object] = { - "tenant_id": case.run.tenant_id, - "app_id": case.run.app_id, - "workflow_id": identifier("definition"), - "created_by": identifier("actor"), - "created_at": naive_utc_now().isoformat(), - "status": "succeeded", - } - workflow = {**common, "id": identifier("workflow-source"), "graph": json.dumps({"nodes": []})} - node = { - **common, - "id": identifier("node-source"), - "workflow_run_id": workflow["id"], - "node_id": "reference-agent", - "node_type": "agent", - "title": "Reference Agent", - } - requests: list[str] = [] - transactions: set[Connection] = set() - - def begun(connection: Connection) -> None: - transactions.add(connection) - - def ended(connection: Connection) -> None: - transactions.discard(connection) - - if request.param == "sql": - with case.sessions.begin() as session: - session.add(_dict_to_workflow_run(workflow)) - session.add(_dict_to_workflow_node_execution_model(node)) - else: - apply_config_overrides( - monkeypatch, - LOGSTORE_DUAL_WRITE_ENABLED=False, - API_WORKFLOW_RUN_REPOSITORY=( - "extensions.logstore.repositories.logstore_api_workflow_run_repository.LogstoreAPIWorkflowRunRepository" - ), - API_WORKFLOW_NODE_EXECUTION_REPOSITORY=( - "extensions.logstore.repositories.logstore_api_workflow_node_execution_repository." - "LogstoreAPIWorkflowNodeExecutionRepository" - ), - ) - for key, value in { - "LOGSTORE_DUAL_READ_ENABLED": "false", - "LOGSTORE_PG_MODE_ENABLED": "false", - "ALIYUN_SLS_ACCESS_KEY_ID": "test-access-key", - "ALIYUN_SLS_ACCESS_KEY_SECRET": "test-secret", - "ALIYUN_SLS_ENDPOINT": "https://sls.example.test", - "ALIYUN_SLS_REGION": "test-region", - }.items(): - monkeypatch.setenv(key, value) - monkeypatch.setattr(AliyunLogStore, "_instance", None) - monkeypatch.setattr(AliyunLogStore, "_initialized", False) - monkeypatch.setattr( - AliyunLogStore, "_check_endpoint_connectivity", staticmethod(lambda _endpoint, _timeout: None) - ) - # The SDK constructor probes the endpoint through a socket before any HTTP request. - monkeypatch.setattr(Util, "get_host_ip", staticmethod(lambda _host: "127.0.0.1")) - - def reject_http_transport(*_args: object, **_kwargs: object) -> None: - pytest.fail("Configured LogStore test escaped its isolated transport boundary") - - monkeypatch.setattr(LogClient, "_getHttpResponse", reject_http_transport) - - def get_logs( - _client: AliyunLogStore, *, query: str, logstore: str, **_kwargs: object - ) -> list[dict[str, object]]: - assert not transactions, "Configured source reads must not hold a SQL transaction" - assert f'tenant_id:"{case.run.tenant_id}"' in query.replace(": ", ":") - assert logstore in { - AliyunLogStore.workflow_execution_logstore, - AliyunLogStore.workflow_node_execution_logstore, - } - is_workflow = logstore == AliyunLogStore.workflow_execution_logstore - if is_workflow: - assert f'app_id:"{case.run.app_id}"' in query - requests.append("workflow" if is_workflow else "node") - row = workflow if is_workflow else node - return [dict(row)] if row else [] - - monkeypatch.setattr(AliyunLogStore, "get_logs", get_logs) - - def execute_sql( - client: AliyunLogStore, *, sql: str, query: str, logstore: str, **_kwargs: object - ) -> list[dict[str, object]]: - # Node lookups use the SQL transport in SDK mode too; the real - # service still validates the returned node-to-workflow owner chain. - assert logstore == AliyunLogStore.workflow_node_execution_logstore - assert f'id:"{identifier("node-source")}"' in query.replace(": ", ":") - assert f"tenant_id = '{case.run.tenant_id}'" in sql - assert f"id = '{identifier('node-source')}'" in sql - assert f'FROM "{logstore}"' in sql - assert "ORDER BY log_version DESC LIMIT 1" in sql - return get_logs(client, query=query, logstore=logstore) - - monkeypatch.setattr(AliyunLogStore, "execute_sql", execute_sql) - with case.sessions() as session: - assert session.scalar(select(WorkflowRun)) is None - assert session.scalar(select(WorkflowNodeExecutionModel)) is None - event.listen(sqlite_engine, "begin", begun) - event.listen(sqlite_engine, "commit", ended) - event.listen(sqlite_engine, "rollback", ended) - try: - run = case.repository.start( - tenant_id=case.run.tenant_id, - app_id=case.run.app_id, - source_kind="workflow", - source_id=str(workflow["id"]), - session_id="caller-session", - ) - assert run is not None - requests.clear() - yield NodeAlias( - Collection(case.repository, case.sessions, run), node, workflow, requests, request.param == "logstore" - ) - finally: - event.remove(sqlite_engine, "begin", begun) - event.remove(sqlite_engine, "commit", ended) - event.remove(sqlite_engine, "rollback", ended) - - @staticmethod - @pytest.mark.parametrize("spans", [True, False]) - def test_alias_returns_parent_trace_and_latest_owned_invocation(alias: NodeAlias, spans: bool) -> None: - trace_url, span_url = alias.accept(spans=spans) - status = alias.status() - assert status["status"] == "accepted" - assert status["source_kind"] == "node" - assert status["source_id"] == alias.node["id"] - assert status["trace_id"] == "provider-trace" - assert status["trace_url"] == trace_url - assert status["span_url"] == (span_url if spans else None) - assert status["config_revision"] == alias.case.run.config_revision - if alias.logstore: - assert alias.requests == ["node", "workflow"] - - @staticmethod - @pytest.mark.parametrize("capture_agent", [True, False]) - @pytest.mark.parametrize("accepted", [True, False]) - def test_alias_without_internal_invocation_returns_parent_status( - alias: NodeAlias, capture_agent: bool, accepted: bool - ) -> None: - with alias.case.sessions.begin() as session: - run = session.get(ApplicationTraceRun, alias.case.run.id) - assert run is not None - run.capture_agent = capture_agent - if accepted: - terminal(alias.case) - lease = claim(alias.case) - assert lease is not None - alias.case.repository.complete( - alias.case.run.id, - lease[1], - state="accepted", - result=TraceExportResult("outer-only", "https://provider.test/outer"), - ) - status = alias.status() - assert status["status"] == ("accepted" if accepted else "collecting") - assert status["span_url"] is None - assert status["capture_agent"] is capture_agent - assert status["trace_url"] == ("https://provider.test/outer" if accepted else None) - - @staticmethod - @pytest.mark.parametrize("state", ["collecting", "not_collected"]) - def test_standalone_node_snapshot_wins_over_parent(alias: NodeAlias, state: str) -> None: - alias.accept() - node_run = alias.case.repository.start( - tenant_id=alias.case.run.tenant_id, - app_id=alias.case.run.app_id, - source_kind="node", - source_id=str(alias.node["id"]), - session_id="preview-session", - ) - assert node_run is not None - with alias.case.sessions.begin() as session: - run = session.get(ApplicationTraceRun, node_run.id) - assert run is not None - run.state = state - alias.requests.clear() - status = alias.status() - assert status["status"] == state - assert status.get("trace_url") is None - if alias.logstore: - assert alias.requests == ["node"] - - @staticmethod - def test_explicit_other_node_filter_cannot_select_sibling_span( - alias: NodeAlias, monkeypatch: pytest.MonkeyPatch - ) -> None: - sibling = alias.case.repository.invocation( - run_id=alias.case.run.id, - agent_id=identifier("referenced-agent"), - node_execution_id=identifier("sibling-node"), - attempt=0, - call_key="sibling", - ) - ingest(alias.case, batch_for(alias.case.run, sibling, span_id="f" * 16)) - alias.accept() - with pytest.raises(TraceCollectionError, match="tracing_node_not_found") as error: - alias.status(identifier("sibling-node")) - assert error.value.status == 404 - service = alias.service() - scheduled: list[str] = [] - monkeypatch.setattr(service, "schedule", scheduled.append) - with pytest.raises(TraceCollectionError, match="tracing_node_not_found"): - service.retry( - tenant_id=alias.case.run.tenant_id, - app_id=alias.case.run.app_id, - source_kind="node", - source_id=str(alias.node["id"]), - node_execution_id=identifier("sibling-node"), - acknowledge_possible_duplicates=True, - ) - assert not scheduled - - @staticmethod - @pytest.mark.parametrize("state", ["failed", "uncertain", "accepted"]) - def test_alias_retry_uses_parent_guards_and_resolves_sources_once( - alias: NodeAlias, state: str, monkeypatch: pytest.MonkeyPatch - ) -> None: - alias.accept() - with alias.case.sessions.begin() as session: - run = session.get(ApplicationTraceRun, alias.case.run.id) - assert run is not None - run.state = state - service = alias.service() - scheduled: list[str] = [] - monkeypatch.setattr(service, "schedule", scheduled.append) - - def retry(acknowledge: bool) -> dict[str, object]: - alias.requests.clear() - return service.retry( - tenant_id=alias.case.run.tenant_id, - app_id=alias.case.run.app_id, - source_kind="node", - source_id=str(alias.node["id"]), - acknowledge_possible_duplicates=acknowledge, - ) - - if state == "accepted": - with pytest.raises(TraceCollectionError, match="tracing_already_accepted") as error: - retry(True) - assert error.value.status == 409 - assert not scheduled - else: - if state == "uncertain": - with pytest.raises(TraceCollectionError, match="tracing_duplicate_acknowledgement_required"): - retry(False) - assert not scheduled - assert retry(state == "uncertain")["status"] == "ready" - assert scheduled == [alias.case.run.id] - if alias.logstore: - assert alias.requests == ["node", "workflow"] - - @staticmethod - @pytest.mark.parametrize("source", ["node", "workflow"]) - @pytest.mark.parametrize("field", ["tenant_id", "app_id", "workflow_id"]) - def test_alias_rejects_broken_owner_chain(alias: NodeAlias, source: str, field: str) -> None: - alias.accept() - row = alias.node if source == "node" else alias.workflow - if alias.logstore: - row[field] = identifier("foreign") - else: - with alias.case.sessions.begin() as session: - model = session.get(WorkflowNodeExecutionModel if source == "node" else WorkflowRun, str(row["id"])) - assert model is not None - # Keep field assignments explicit so ownership failures are visible. - if field == "tenant_id": - model.tenant_id = identifier("foreign") - elif field == "app_id": - model.app_id = identifier("foreign") - else: - model.workflow_id = identifier("foreign") - with pytest.raises(TraceCollectionError, match="tracing_source_not_found") as error: - alias.status() - assert error.value.status == 404 - - @staticmethod - def test_alias_missing_parent_business_source_is_not_a_durable_owner_bypass(alias: NodeAlias) -> None: - if alias.logstore: - alias.workflow.clear() - else: - with alias.case.sessions.begin() as session: - parent = session.get(WorkflowRun, str(alias.workflow["id"])) - assert parent is not None - session.delete(parent) - with pytest.raises(TraceCollectionError, match="tracing_source_not_found"): - alias.status() - - @staticmethod - def test_owned_alias_without_parent_trace_is_not_collected(alias: NodeAlias) -> None: - with alias.case.sessions.begin() as session: - run = session.get(ApplicationTraceRun, alias.case.run.id) - assert run is not None - session.delete(run) - assert alias.status() == {"status": "not_collected", "source_kind": "node", "source_id": alias.node["id"]} - - -"""Real SQLite state transitions and background renewal for delivery leases. - -PostgreSQL row-lock scheduling and the deployed Celery pool still need dev -validation; these tests establish token and acknowledgement semantics. -""" - - -def _stored(case: Collection) -> ApplicationTraceRun: - row = case.repository.get(case.run.id) - assert row is not None - return row - - -def _expire(case: Collection) -> str: - leased = claim(case) - assert leased is not None - with case.sessions.begin() as session: - row = session.get(ApplicationTraceRun, case.run.id) - assert row is not None - row.export_started_at = naive_utc_now() - timedelta(seconds=60) - row.export_heartbeat_at = row.export_started_at - assert claim(case) is None - assert _stored(case).state == "uncertain" - return leased[1] - - -class TestExportLeaseCases: - @staticmethod - @pytest.fixture - def case(sqlite_session_factory: sessionmaker[Session]) -> Collection: - result = seed(sqlite_session_factory, "lease-test") - ingest(result, batch_for(result.run, result.invocation())) - terminal(result) - return result - - @staticmethod - def test_renewal_keeps_long_export_alive_then_missing_heartbeat_expires( - case: Collection, monkeypatch: pytest.MonkeyPatch - ) -> None: - clock = [naive_utc_now()] - monkeypatch.setattr(repository_module, "naive_utc_now", lambda: clock[0]) - leased = claim(case) - assert leased is not None - token = leased[1] - clock[0] += timedelta(seconds=20) - assert case.repository.renew_export_lease(case.run.id, token) - clock[0] += timedelta(seconds=20) - assert claim(case) is None - current = _stored(case) - assert current.state == "exporting" - assert current.attempts == 1 - assert current.export_started_at is not None - assert current.export_started_at + timedelta(seconds=30) < clock[0] - clock[0] += timedelta(seconds=11) - assert claim(case) is None - expired = _stored(case) - assert expired.state == "uncertain" - assert expired.error_code == "export_lease_expired" - assert expired.export_token == token - assert not case.repository.renew_export_lease(case.run.id, token) - - @staticmethod - def test_old_rows_without_heartbeat_use_started_at(case: Collection) -> None: - leased = claim(case) - assert leased is not None - with case.sessions.begin() as session: - row = session.get(ApplicationTraceRun, case.run.id) - assert row is not None - row.export_started_at = naive_utc_now() - timedelta(seconds=60) - row.export_heartbeat_at = None - assert claim(case) is None - current = _stored(case) - assert current.state == "uncertain" - assert current.export_token == leased[1] - - @staticmethod - def test_wrong_token_cannot_renew_or_complete(case: Collection) -> None: - leased = claim(case) - assert leased is not None - before = _stored(case) - assert not case.repository.renew_export_lease(case.run.id, "other-token") - case.repository.complete(case.run.id, "other-token", state="accepted", result=TraceExportResult("wrong")) - current = _stored(case) - assert current.export_token == leased[1] - assert current.export_heartbeat_at == before.export_heartbeat_at - assert current.state == "exporting" - assert current.provider_trace_id is None - - @staticmethod - def test_late_success_wins_before_manual_retry_and_is_idempotent(case: Collection) -> None: - token = _expire(case) - result = TraceExportResult("confirmed", "https://example.com/trace") - case.repository.complete(case.run.id, token, state="accepted", result=result) - current = _stored(case) - assert current.state == "accepted" - assert current.error_code is None - assert current.export_token is None - assert current.accepted_at is not None - assert current.trace_url == result.trace_url - case.repository.complete(case.run.id, token, state="accepted", result=TraceExportResult("duplicate-ack")) - assert _stored(case).provider_trace_id == "confirmed" - with pytest.raises(TraceCollectionError, match="tracing_already_accepted"): - case.repository.retry(run_id=case.run.id, acknowledge_possible_duplicates=True) - - @staticmethod - def test_manual_retry_revokes_old_heartbeat_and_ack_before_and_after_new_claim(case: Collection) -> None: - old_token = _expire(case) - case.repository.retry(run_id=case.run.id, acknowledge_possible_duplicates=True) - assert not case.repository.renew_export_lease(case.run.id, old_token) - case.repository.complete(case.run.id, old_token, state="accepted", result=TraceExportResult("old")) - assert _stored(case).state == "ready" - second = claim(case) - assert second is not None - assert second[1] != old_token - case.repository.complete(case.run.id, old_token, state="accepted", result=TraceExportResult("old")) - current = _stored(case) - assert current.state == "exporting" - assert current.export_token == second[1] - assert current.provider_trace_id is None - case.repository.complete(case.run.id, second[1], state="accepted", result=TraceExportResult("new")) - assert _stored(case).provider_trace_id == "new" - - @staticmethod - @pytest.mark.parametrize("late_state", ["failed", "retrying", "uncertain"]) - def test_late_non_success_cannot_overwrite_lease_expiry(case: Collection, late_state: str) -> None: - token = _expire(case) - case.repository.complete(case.run.id, token, state=late_state, error_code="late-failure") - current = _stored(case) - assert current.state == "uncertain" - assert current.error_code == "export_lease_expired" - - @staticmethod - def test_other_uncertain_reason_cannot_be_promoted_by_late_ack(case: Collection) -> None: - token = _expire(case) - with case.sessions.begin() as session: - row = session.get(ApplicationTraceRun, case.run.id) - assert row is not None - row.error_code = "provider_acceptance_unknown" - case.repository.complete(case.run.id, token, state="accepted", result=TraceExportResult("late")) - assert _stored(case).state == "uncertain" - - @staticmethod - @pytest.mark.parametrize("transient_failure", [False, True]) - @pytest.mark.parametrize("outcome", ["accepted", "uncertain"]) - def test_service_heartbeats_during_export_and_stops_on_every_outcome( - case: Collection, - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - transient_failure: bool, - outcome: str, - ) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_EXPORT_LEASE_SECONDS=1) - clock = [naive_utc_now()] - monkeypatch.setattr(repository_module, "naive_utc_now", lambda: clock[0]) - monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", lambda *_args: {}) - renewed = threading.Event() - calls = [] - original = case.repository.renew_export_lease - - def renew(run_id: str, token: str) -> bool: - calls.append(threading.get_ident()) - if transient_failure and len(calls) == 1: - raise RuntimeError("private-provider-value") - result = original(run_id, token) - renewed.set() - return result - - monkeypatch.setattr(case.repository, "renew_export_lease", renew) - - def export(*_args: object) -> TraceExportResult: - # The original start time is now beyond the lease; only a real background - # transaction can prevent a second claimant from marking it uncertain. - clock[0] += timedelta(seconds=3) - assert renewed.wait(3), "background heartbeat did not commit" - assert case.repository.claim(case.run.id, collection_timeout=60, lease_seconds=1) is None - assert _stored(case).state == "exporting" - if outcome == "uncertain": - raise TraceExportUncertainError("lost acknowledgement") - return TraceExportResult("ack") - - monkeypatch.setattr(service_module, "export_trace", export) - ApplicationTracingService(case.repository).deliver(case.run.id) - assert _stored(case).state == outcome - assert calls - assert all(thread_id != threading.get_ident() for thread_id in calls) - assert not any(thread.name == f"application-trace-lease-{case.run.id}" for thread in threading.enumerate()) - assert "private-provider-value" not in caplog.text - - @staticmethod - def test_heartbeat_stops_even_when_completion_database_write_raises( - case: Collection, monkeypatch: pytest.MonkeyPatch - ) -> None: - monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", lambda *_args: {}) - monkeypatch.setattr(service_module, "export_trace", lambda *_args: TraceExportResult("ack")) - - def fail_complete(*_args: object, **_kwargs: object) -> None: - raise RuntimeError("database unavailable") - - monkeypatch.setattr(case.repository, "complete", fail_complete) - with pytest.raises(RuntimeError, match="database unavailable"): - ApplicationTracingService(case.repository).deliver(case.run.id) - assert not any(thread.name == f"application-trace-lease-{case.run.id}" for thread in threading.enumerate()) - - -"""Persisted legacy sources through real OTel protobuf, replacing only HTTP I/O.""" - - -_LEGACY_ALIYUN_TRANSPORT_START = datetime(2026, 10, 19, tzinfo=UTC) -RUN_ID = "a0000000-0000-4000-8000-000000000001" -MESSAGE_ID = "a0000000-0000-4000-8000-000000000002" -NODE_IDS = [f"b0000000-0000-4000-8000-00000000000{i}" for i in range(3)] - - -class _LegacyAliyunTransportRequestoptions(TypedDict, total=False): - content: bytes - headers: Mapping[str, str] - max_retries: int - follow_redirects: bool - timeout: float - - -def _legacy_aliyun_transport_config(**overrides: object) -> AliyunConfig: - return AliyunConfig.model_validate( - { - "app_name": "owned-service", - "license_key": "owned-key", - "endpoint": "https://collector.test/prefix", - **overrides, - } - ) - - -def decoded(options: _LegacyAliyunTransportRequestoptions) -> ExportTraceServiceRequest: - payload = ExportTraceServiceRequest() - payload.ParseFromString(options["content"]) - return payload - - -class TestLegacyAliyunTransportCases: - @staticmethod - def test_constructor_and_presentation_have_no_network_sdk_queue_or_environment_mutation( - monkeypatch: pytest.MonkeyPatch, - ) -> None: - def forbidden(*_args: object, **_kwargs: object) -> None: - pytest.fail("Constructor and presentation must not create I/O, exporters or worker threads") - - monkeypatch.setattr(ssrf_proxy, "make_request", forbidden) - monkeypatch.setattr(OTLPSpanExporter, "__init__", forbidden) - monkeypatch.setattr(Thread, "start", forbidden) - monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "https://ambient.test") - monkeypatch.setenv("OTEL_EXPORTER_OTLP_HEADERS", "Authorization=ambient") - before = dict(os.environ) - original = _legacy_aliyun_transport_config() - adapter = AliyunDataTrace(original) - original.license_key = "mutated-after-construction" - assert adapter.trace_client.endpoint == "https://collector.test/prefix/adapt_owned-key/api/otlp/traces" - assert adapter.get_project_url() == "https://arms.console.aliyun.com/#/llm" - assert not hasattr(adapter.trace_client, "exporter") - assert not hasattr(adapter.trace_client, "queue") - assert not hasattr(adapter.trace_client, "worker_thread") - assert dict(os.environ) == before - - @staticmethod - @pytest.mark.usefixtures("agent_runtime_app") - @pytest.mark.parametrize("failure", [None, 400, 401, 403, 429, 503, "timeout", "partial", "malformed", "redirect"]) - def test_persisted_workflow_uses_one_acknowledged_protobuf_batch_and_preserves_business_state( - sqlite_session: Session, monkeypatch: pytest.MonkeyPatch, failure: int | str | None - ) -> None: - tenant = make_tenant(tenant_id="tenant") - account = make_account(account_id="owner", tenant=tenant) - app = make_app(app_id="app", tenant_id=tenant.id, mode=AppMode.WORKFLOW, created_by=account.id) - membership = TenantAccountJoin(tenant_id=tenant.id, account_id=account.id, role=TenantAccountRole.OWNER) - run = WorkflowRun( - id=RUN_ID, - tenant_id=tenant.id, - app_id=app.id, - workflow_id="workflow", - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, - version="v1", - inputs=json.dumps({"query": "hello"}), - outputs=json.dumps({"answer": "ok"}), - status=WorkflowExecutionStatus.SUCCEEDED, - total_tokens=32, - created_by=account.id, - created_by_role=CreatorUserRole.ACCOUNT, - finished_at=_LEGACY_ALIYUN_TRANSPORT_START + timedelta(seconds=1), - elapsed_time=1, - ) - run.created_at = _LEGACY_ALIYUN_TRANSPORT_START - sqlite_session.add_all([tenant, account, app, membership, run]) - nodes = [ - WorkflowNodeExecutionModel( - id=NODE_IDS[i], - tenant_id=tenant.id, - app_id=app.id, - workflow_id="workflow", - workflow_run_id=RUN_ID, - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - index=i, - node_id=f"node-{i}", - title=title, - node_type=kind, - status=WorkflowNodeExecutionStatus.SUCCEEDED, - created_by=account.id, - created_by_role=CreatorUserRole.ACCOUNT, - inputs=json.dumps({"synthetic": True}), - outputs=json.dumps( - {"text": "ok", "usage": {"prompt_tokens": 21, "completion_tokens": 11, "total_tokens": 32}} - ), - finished_at=_LEGACY_ALIYUN_TRANSPORT_START + timedelta(milliseconds=i + 100), - elapsed_time=0.1, - process_data=json.dumps( - { - "prompts": [{"role": "user", "text": "hello"}], - "model_provider": "test", - "model_name": "test-model", - } - ) - if i == 1 - else None, - ) - for i, (title, kind) in enumerate( - [("Start", BuiltinNodeTypes.START), ("LLM", BuiltinNodeTypes.LLM), ("End", BuiltinNodeTypes.END)] - ) - ] - for node_index, node_record in enumerate(nodes): - node_record.created_at = _LEGACY_ALIYUN_TRANSPORT_START + timedelta(milliseconds=node_index) - sqlite_session.add_all(nodes) - sqlite_session.commit() - info = WorkflowTraceInfo( - query="hello", - workflow_id="workflow", - tenant_id=tenant.id, - workflow_run_id=RUN_ID, - workflow_run_elapsed_time=1.0, - workflow_run_status="succeeded", - workflow_run_inputs={"query": "hello"}, - workflow_run_outputs={"answer": "ok"}, - workflow_run_version="v1", - total_tokens=32, - file_list=[], - metadata={"app_id": app.id, "conversation_id": "owned-conversation"}, - start_time=_LEGACY_ALIYUN_TRANSPORT_START, - end_time=_LEGACY_ALIYUN_TRANSPORT_START + timedelta(seconds=1), - ) - sent: list[tuple[str, _LegacyAliyunTransportRequestoptions, ExportTraceServiceRequest]] = [] - - def send(method: str, url: str, **options: Unpack[_LegacyAliyunTransportRequestoptions]) -> httpx.Response: - assert method == "POST" - assert url == "https://collector.test/prefix/adapt_owned-key/api/otlp/traces" - sent.append((url, options, decoded(options))) - if len(sent) == 1: - if failure == "timeout": - raise httpx.ReadTimeout("synthetic lost ACK", request=httpx.Request("POST", url)) - if isinstance(failure, int): - return httpx.Response(failure, content=b"DO_NOT_EXPOSE_BODY") - if failure == "malformed": - return httpx.Response(200, content=b"not protobuf") - if failure == "partial": - response = ExportTraceServiceResponse() - response.partial_success.rejected_spans = 1 - response.partial_success.error_message = "DO_NOT_EXPOSE_BODY" - return httpx.Response(200, content=response.SerializeToString()) - if failure == "redirect": - return httpx.Response(307, headers={"Location": "https://unrelated.test"}) - return httpx.Response(200, content=ExportTraceServiceResponse().SerializeToString()) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - adapter = AliyunDataTrace(_legacy_aliyun_transport_config()) - if failure is None: - adapter.trace(info) - else: - expected = ( - RetryableTraceDispatchError - if failure in {429, 503} - else TraceExportUncertainError - if failure in {"timeout", "partial", "malformed"} - else TraceExportRejectedError - ) - with pytest.raises(expected) as error: - adapter.trace(info) - assert "DO_NOT_EXPOSE_BODY" not in str(error.value) - if isinstance(error.value, RetryableTraceDispatchError): - assert error.value.acceptance_unknown is (expected is TraceExportUncertainError) - assert len(sent) == 1 # No SDK retry, per-span request, batching thread or silent queue success. - _, options, payload = sent[0] - assert options["max_retries"] == 0 - assert options["follow_redirects"] is False - assert options["headers"] == {"Content-Type": "application/x-protobuf", "Accept": "application/x-protobuf"} - resource = payload.resource_spans[0] - attributes = {item.key: item.value for item in resource.resource.attributes} - assert attributes["service.name"].string_value == "owned-service" - spans = list(resource.scope_spans[0].spans) - assert len(spans) == 4 - assert [s.name for s in spans] == ["workflow", "Start", "LLM", "End"] - root_id = convert_to_span_id(RUN_ID, "workflow").to_bytes(8) - assert [s.span_id for s in spans] == [root_id, *[convert_to_span_id(i, "node").to_bytes(8) for i in NODE_IDS]] - assert all(s.trace_id == UUID(RUN_ID).bytes for s in spans) - assert [s.parent_span_id for s in spans] == [b"", root_id, root_id, root_id] - assert spans[0].start_time_unix_nano == int(_LEGACY_ALIYUN_TRANSPORT_START.timestamp() * 1_000_000_000) - assert spans[0].end_time_unix_nano == int( - (_LEGACY_ALIYUN_TRANSPORT_START + timedelta(seconds=1)).timestamp() * 1_000_000_000 - ) - model_attributes = {item.key: item.value for item in spans[2].attributes} - assert model_attributes[GEN_AI_USAGE_TOTAL_TOKENS].string_value == "32" - assert all(s.status.code == 1 for s in spans) - assert adapter.trace_client._dispatch_spans.get() is None - assert info.workflow_run_status == "succeeded" - sqlite_session.refresh(run) - assert run.status == WorkflowExecutionStatus.SUCCEEDED - assert run.total_tokens == 32 - assert all(node.status == WorkflowNodeExecutionStatus.SUCCEEDED for node in nodes) - # A fresh business operation does not inherit any old batch after success or error. - adapter.trace(ModerationTraceInfo(flagged=False, action="allow", preset_response="", query="", metadata={})) - assert len(sent) == 1 - adapter.trace(info.model_copy(update={"workflow_run_id": "a0000000-0000-4000-8000-000000000003"})) - assert len(sent) == 2 - recovered_spans = list(sent[1][2].resource_spans[0].scope_spans[0].spans) - assert [s.name for s in recovered_spans] == ["workflow"] - - @staticmethod - @pytest.mark.parametrize("failed", [False, True]) - def test_message_mapping_keeps_original_ids_raw_usage_error_and_business_state( - sqlite_session: Session, monkeypatch: pytest.MonkeyPatch, failed: bool - ) -> None: - message = make_message( - message_id=MESSAGE_ID, - app_id="app", - conversation_id="conversation", - query="hello", - answer="ok", - status=MessageStatus.ERROR if failed else MessageStatus.NORMAL, - inputs={}, - message={}, - message_unit_price=Decimal(0), - answer_unit_price=Decimal(0), - currency="USD", - from_source=ConversationFromSource.CONSOLE, - from_account_id="owner", - message_tokens=21, - answer_tokens=11, - ) - sqlite_session.add(message) - sqlite_session.commit() - info = MessageTraceInfo( - conversation_model="chat", - conversation_mode="chat", - message_id=MESSAGE_ID, - message_data=message, - inputs={"prompt": "hello"}, - outputs="ok", - message_tokens=21, - answer_tokens=11, - total_tokens=32, - error="synthetic model failure" if failed else None, - metadata={ - "app_id": "app", - "conversation_id": "conversation", - "ls_model_name": "test-model", - "ls_provider": "test", - }, - start_time=_LEGACY_ALIYUN_TRANSPORT_START, - end_time=_LEGACY_ALIYUN_TRANSPORT_START + timedelta(seconds=1), - ) - sent: list[ExportTraceServiceRequest] = [] - - def send(_method: str, _url: str, **options: Unpack[_LegacyAliyunTransportRequestoptions]) -> httpx.Response: - sent.append(decoded(options)) - return httpx.Response(200, content=b"") - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - adapter = AliyunDataTrace(_legacy_aliyun_transport_config()) - adapter.trace(info) - assert len(sent) == 1 - spans = list(sent[0].resource_spans[0].scope_spans[0].spans) - assert [s.name for s in spans] == ["message", "llm"] - root_id = convert_to_span_id(MESSAGE_ID, "message").to_bytes(8) - assert [s.parent_span_id for s in spans] == [b"", root_id] - assert [s.span_id for s in spans] == [root_id, convert_to_span_id(MESSAGE_ID, "llm").to_bytes(8)] - assert all(s.trace_id == UUID(MESSAGE_ID).bytes and s.status.code == (2 if failed else 1) for s in spans) - assert {a.key: a.value for a in spans[1].attributes}[GEN_AI_USAGE_TOTAL_TOKENS].string_value == "32" - assert message.status == (MessageStatus.ERROR if failed else MessageStatus.NORMAL) - - @staticmethod - def test_one_cached_client_keeps_interleaved_batches_and_failure_cleanup_isolated( - monkeypatch: pytest.MonkeyPatch, - ) -> None: - client = TraceClient(service_name="owned-service", endpoint="https://collector.test/adapt_key/api/otlp/traces") - barrier = Barrier(2) - observed: list[list[str]] = [] - - def send(_method: str, _url: str, **options: Unpack[_LegacyAliyunTransportRequestoptions]) -> httpx.Response: - payload = decoded(options) - observed.append([s.name for rs in payload.resource_spans for scope in rs.scope_spans for s in scope.spans]) - return httpx.Response(200, content=b"") - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - - def run(label: str) -> None: - try: - with client.batch(): - client.add_span(SpanData(name=label, trace_id=1, span_id=2, start_time=1, end_time=2)) - barrier.wait(timeout=5) - if label == "failed": - raise ValueError("mapper failed before submission") - client.add_span(SpanData(name=label + "-child", trace_id=1, span_id=3, start_time=1, end_time=2)) - except ValueError: - assert label == "failed" - assert client._dispatch_spans.get() is None - - with ThreadPoolExecutor(max_workers=2) as pool: - list(pool.map(run, ["accepted", "failed"])) - assert observed == [["accepted", "accepted-child"]] - with client.batch(): - client.add_span(SpanData(name="next", trace_id=1, span_id=4, start_time=1, end_time=2)) - assert observed == [["accepted", "accepted-child"], ["next"]] - - -"""Real Weave SDK serialization with only the outbound HTTP boundary replaced.""" - - -_LEGACY_WEAVE_TRANSPORT_START = datetime(2026, 10, 19, tzinfo=UTC) - - -class _LegacyWeaveTransportRequestoptions(TypedDict, total=False): - content: bytes - json: dict[str, object] - auth: tuple[str, str] - headers: Mapping[str, str] - max_retries: int - follow_redirects: bool - timeout: float - - -def _legacy_weave_transport_config(**overrides: object) -> WeaveConfig: - return WeaveConfig.model_validate({"api_key": "fixture-key", "entity": "team", "project": "project", **overrides}) - - -def call(call_id: str = "call-1", **overrides: object) -> WeaveTraceModel: - return WeaveTraceModel.model_validate( - { - "id": call_id, - "op": "tool", - "inputs": {"synthetic": True}, - "outputs": {"result": "ok"}, - "attributes": { - "trace_id": call_id, - "start_time": _LEGACY_WEAVE_TRANSPORT_START, - "end_time": _LEGACY_WEAVE_TRANSPORT_START + timedelta(seconds=1), - }, - **overrides, - } - ) - - -def response(url: str, body: object, *, status: int = 200) -> httpx.Response: - return httpx.Response(status, request=httpx.Request("POST", url), json=body) - - -class TestLegacyWeaveTransportCases: - @staticmethod - @pytest.fixture - def wire( - monkeypatch: pytest.MonkeyPatch, failure: tuple[str, str, int | str] | None - ) -> list[tuple[str, _LegacyWeaveTransportRequestoptions, dict[str, object]]]: - sent: list[tuple[str, _LegacyWeaveTransportRequestoptions, dict[str, object]]] = [] - fault_pending = True - - def send(method: str, url: str, **options: Unpack[_LegacyWeaveTransportRequestoptions]) -> httpx.Response: - nonlocal fault_pending - assert method == "POST" - assert url in {"https://trace.wandb.ai/call/start", "https://trace.wandb.ai/call/end"} - body = json.loads(options["content"]) - sent.append((url, options, body)) - stage = "start" if "start" in body else "end" - record = body[stage] - if fault_pending and failure is not None and (stage, record["id"]) == failure[:2]: - fault_pending = False - outcome = failure[2] - if outcome == "timeout": - raise httpx.ReadTimeout("synthetic lost acknowledgment", request=httpx.Request("POST", url)) - if outcome == "malformed": - return httpx.Response(200, request=httpx.Request("POST", url), content=b"invalid-json") - assert isinstance(outcome, int) - return response(url, {"error": "DO_NOT_EXPOSE_BODY"}, status=outcome) - if url.endswith("/call/start"): - req = CallStartReq.model_validate(body) - return response(url, {"id": req.start.id, "trace_id": req.start.trace_id}) - CallEndReq.model_validate(body) - return response(url, {}) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - return sent - - @staticmethod - def test_constructor_and_presentation_have_no_login_network_files_or_background_state( - monkeypatch: pytest.MonkeyPatch, - ) -> None: - def forbidden(*_args: object, **_kwargs: object) -> None: - pytest.fail("Constructor must not initialize global credentials, files, I/O or background work") - - monkeypatch.setattr(Path, "touch", forbidden) - monkeypatch.setattr(wandb, "login", forbidden) - monkeypatch.setattr(weave, "init", forbidden) - monkeypatch.setattr("wandb.sdk.lib.wbauth.wbnetrc.write_netrc_auth", forbidden) - monkeypatch.setattr("weave.trace_server_bindings.remote_http_trace_server.AsyncBatchProcessor", forbidden) - monkeypatch.setattr("weave.trace_server_bindings.remote_http_trace_server.CallBatchProcessor", forbidden) - monkeypatch.setattr(ssrf_proxy, "make_request", forbidden) - monkeypatch.setattr(atexit, "register", forbidden) - monkeypatch.setenv("WANDB_API_KEY", "unrelated-key") - monkeypatch.setenv("WANDB_BASE_URL", "https://unrelated.test") - monkeypatch.setenv("WEAVE_TRACE_SERVER", "https://unrelated-traces.test") - before = dict(os.environ) - original = _legacy_weave_transport_config(entity=None) - adapter = WeaveDataTrace(original) - original.api_key = "changed-after-construction" - assert adapter.weave_api_key == "fixture-key" - assert adapter.get_project_url() == "https://wandb.ai" - assert adapter.server.trace_server_url == "https://trace.wandb.ai" - assert adapter.server._auth == ("api", "fixture-key") - assert not adapter.server.should_batch - assert adapter.server.call_processor is None - assert adapter.server.feedback_processor is None - assert dict(os.environ) == before - - @staticmethod - @pytest.mark.usefixtures("agent_runtime_app") - @pytest.mark.parametrize( - "failure", - [None] - + [ - (stage, call_id, status) - for stage, call_id in [ - ("start", "legacy-run"), - ("end", "legacy-run"), - ("start", "node-0"), - ("end", "node-0"), - ] - for status in [429, 503] - ] - + [ - ("start", "legacy-run", "timeout"), - ("start", "legacy-run", "malformed"), - ("end", "node-0", "malformed"), - ("end", "legacy-run", "timeout"), - ] - + [("start", "legacy-run", 401), ("end", "legacy-run", 401)], - ) - def test_workflow_emits_four_real_sdk_calls_with_original_ids_payload_and_status( - sqlite_session: Session, - wire: list[tuple[str, _LegacyWeaveTransportRequestoptions, dict[str, object]]], - failure: tuple[str, str, int | str] | None, - ) -> None: - # The normal Flask/SQLite fixtures bind the actual service-account lookup and - # production repository to persisted rows. Only outbound HTTP is substituted. - tenant = make_tenant(tenant_id="tenant") - account = make_account(account_id="owner", tenant=tenant) - app = make_app(app_id="app", tenant_id=tenant.id, mode=AppMode.WORKFLOW, created_by=account.id) - membership = TenantAccountJoin(tenant_id=tenant.id, account_id=account.id, role=TenantAccountRole.OWNER) - sqlite_session.add_all([tenant, account, app, membership]) - nodes = [ - WorkflowNodeExecutionModel( - id=f"node-{i}", - tenant_id=tenant.id, - app_id=app.id, - workflow_id="workflow", - workflow_run_id="legacy-run", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - index=i, - node_id=f"node-{i}", - title=title, - node_type=kind, - status=WorkflowNodeExecutionStatus.SUCCEEDED, - created_by=account.id, - created_by_role=CreatorUserRole.ACCOUNT, - inputs=json.dumps({"synthetic": True}), - outputs=json.dumps({"answer": "ok"}), - finished_at=_LEGACY_WEAVE_TRANSPORT_START + timedelta(milliseconds=i + 100), - elapsed_time=0.1, - process_data=json.dumps( - { - "prompts": [{"role": "user", "text": "hello"}], - "model_mode": "chat", - "model_provider": "test", - "model_name": "test-model", - } - ) - if i == 1 - else None, - execution_metadata=json.dumps( - {"total_tokens": 32, "total_price": "0.0000098", "currency": "USD"} if i == 1 else {} - ), - ) - for i, (title, kind) in enumerate( - [("Start", BuiltinNodeTypes.START), ("LLM", BuiltinNodeTypes.LLM), ("End", BuiltinNodeTypes.END)] - ) - ] - for node_index, node_record in enumerate(nodes): - node_record.created_at = _LEGACY_WEAVE_TRANSPORT_START + timedelta(milliseconds=node_index) - sqlite_session.add_all(nodes) - sqlite_session.commit() - adapter = WeaveDataTrace(_legacy_weave_transport_config()) - info = WorkflowTraceInfo( - query="hello", - workflow_id="workflow", - tenant_id="tenant", - workflow_run_id="legacy-run", - workflow_run_elapsed_time=1.0, - workflow_run_status="succeeded", - workflow_run_inputs={"query": "hello"}, - workflow_run_outputs={"answer": "ok"}, - workflow_run_version="v1", - total_tokens=32, - file_list=[], - metadata={"app_id": "app", "conversation_id": None}, - trace_id="external-correlation", - start_time=_LEGACY_WEAVE_TRANSPORT_START, - end_time=_LEGACY_WEAVE_TRANSPORT_START + timedelta(seconds=1), - ) - - if failure is not None: - stage, call_id, outcome = failure - first_start = (stage, call_id) == ("start", "legacy-run") - expected = ( - TraceExportRejectedError - if outcome == 401 - else RetryableTraceDispatchError - if first_start and outcome in {429, 503} - else TraceExportUncertainError - ) - with pytest.raises(expected) as error: - adapter.trace(info) - if isinstance(error.value, RetryableTraceDispatchError): - assert error.value.acceptance_unknown is (expected is TraceExportUncertainError) - assert len(wire) == (1 if first_start else 8 if call_id == "legacy-run" else 2 if stage == "start" else 3) - assert "DO_NOT_EXPOSE" not in str(error.value) - assert adapter.calls == {} - assert info.workflow_run_status == "succeeded" - assert all(node.status == WorkflowNodeExecutionStatus.SUCCEEDED for node in nodes) - before_recovery = len(wire) - adapter.trace( - GenerateNameTraceInfo( - tenant_id=tenant.id, - inputs="new business operation", - outputs="new title", - metadata={}, - start_time=_LEGACY_WEAVE_TRANSPORT_START, - end_time=_LEGACY_WEAVE_TRANSPORT_START + timedelta(seconds=1), - ) - ) - assert len(wire) == before_recovery + 2 - assert adapter.calls == {} - return - - adapter.trace(info) - - starts = [CallStartReq.model_validate(body).start for _, _, body in wire if "start" in body] - ends = [CallEndReq.model_validate(body).end for _, _, body in wire if "end" in body] - assert len(starts) == len(ends) == 4 - assert [s.id for s in starts] == ["legacy-run", "node-0", "node-1", "node-2"] - assert [s.parent_id for s in starts] == [None, "legacy-run", "legacy-run", "legacy-run"] - assert all(s.trace_id == "external-correlation" and s.project_id == "team/project" for s in starts) - # Legacy has no native thread/turn contract: keep the existing raw metadata. - assert all(s.thread_id is None and s.turn_id is None for s in starts) - assert starts[0].attributes["conversation_id"] is None - model = starts[2] - assert model.attributes["total_tokens"] == 32 - assert model.attributes["total_price"] == "0.0000098" - assert model.inputs["messages"][0]["content"] == "hello" - model_end = next(e for e in ends if e.id == "node-1") - assert isinstance(model_end.output, dict) - assert model_end.output["usage_metadata"]["total_tokens"] == 32 - assert all(e.exception is None and e.summary["status_counts"][TraceStatus.SUCCESS] == 1 for e in ends) - assert all("usage" not in e.summary for e in ends) # Do not invent native aggregation for the old mapper. - assert all(options["auth"] == ("api", "fixture-key") and options["max_retries"] == 0 for _, options, _ in wire) - assert adapter.calls == {} - - @staticmethod - def test_interleaved_instances_resolve_only_their_own_credentials_and_targets( - monkeypatch: pytest.MonkeyPatch, - ) -> None: - barrier = Barrier(2) - observed: list[tuple[str, tuple[str, str], str]] = [] - monkeypatch.setenv("WANDB_API_KEY", "ambient-key") - monkeypatch.setenv("WANDB_ENTITY", "ambient-entity") - monkeypatch.setenv("WANDB_PROJECT", "ambient-project") - before = dict(os.environ) - - def send(_method: str, url: str, **options: Unpack[_LegacyWeaveTransportRequestoptions]) -> httpx.Response: - label = ( - "a" - if url - in { - "https://control-a.test/prefix/graphql", - "https://trace-a.test/prefix/call/start", - "https://trace-a.test/prefix/call/end", - } - else "b" - ) - assert url in { - f"https://control-{label}.test/prefix/graphql", - f"https://trace-{label}.test/prefix/call/start", - f"https://trace-{label}.test/prefix/call/end", - } - assert options["auth"] == ("api", f"key-{label}") - if url.endswith("/graphql"): - return response(url, {"data": {"viewer": {"entity": f"team-{label}"}}}) - body = json.loads(options["content"]) - record = body.get("start", body.get("end")) - assert record["project_id"] == f"team-{label}/project-{label}" - observed.append((url, options["auth"], record["project_id"])) - if "start" in body: - barrier.wait(timeout=5) - return response(url, {"id": record["id"], "trace_id": record["trace_id"]}) - return response(url, {}) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - - def run(label: str) -> None: - adapter = WeaveDataTrace( - _legacy_weave_transport_config( - api_key=f"key-{label}", - entity=None, - project=f"project-{label}", - host=f"https://control-{label}.test/prefix", - endpoint=f"https://trace-{label}.test/prefix/", - ) - ) - data = call(f"call-{label}") - adapter.start_call(data) - adapter.finish_call(data) - - with ThreadPoolExecutor(max_workers=2) as pool: - list(pool.map(run, ["a", "b"])) - assert len(observed) == 4 - assert dict(os.environ) == before - - @staticmethod - def test_concurrent_dispatches_on_one_cached_instance_keep_prefix_state_isolated( - monkeypatch: pytest.MonkeyPatch, - ) -> None: - barrier = Barrier(2) - - def send(_method: str, url: str, **options: Unpack[_LegacyWeaveTransportRequestoptions]) -> httpx.Response: - body = json.loads(options["content"]) - if "start" in body: - start = CallStartReq.model_validate(body).start - if start.inputs["label"] == "unaccepted": - barrier.wait(timeout=5) - return response(url, {}, status=503) - return response(url, {"id": start.id, "trace_id": start.trace_id}) - CallEndReq.model_validate(body) - barrier.wait(timeout=5) - return response(url, {}, status=503) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - adapter = WeaveDataTrace(_legacy_weave_transport_config()) - - def run(label: str) -> RetryableTraceDispatchError: - info = GenerateNameTraceInfo( - tenant_id="tenant", - inputs={"label": label}, - outputs={"label": label}, - metadata={}, - start_time=_LEGACY_WEAVE_TRANSPORT_START, - end_time=_LEGACY_WEAVE_TRANSPORT_START + timedelta(seconds=1), - ) - with pytest.raises(RetryableTraceDispatchError) as error: - adapter.trace(info) - return error.value - - with ThreadPoolExecutor(max_workers=2) as pool: - accepted, unaccepted = list(pool.map(run, ["accepted", "unaccepted"])) - assert isinstance(accepted, TraceExportUncertainError) - assert accepted.acceptance_unknown - assert not unaccepted.acceptance_unknown - assert adapter.calls == {} - - @staticmethod - @pytest.mark.parametrize( - ("stage", "failure"), - [ - (stage, failure) - for stage in ("start", "end") - for failure in ("rejected", "retryable", "malformed", "mismatched_id") - if not (stage == "end" and failure == "mismatched_id") - ], - ) - def test_synchronous_failures_propagate_and_next_call_recovers( - monkeypatch: pytest.MonkeyPatch, stage: str, failure: str - ) -> None: - broken = True - requests = [] - - def send(_method: str, url: str, **options: Unpack[_LegacyWeaveTransportRequestoptions]) -> httpx.Response: - body = json.loads(options["content"]) - requests.append(url) - if broken and url.endswith("/call/" + stage): - if failure == "rejected": - return response(url, {"error": "DO_NOT_EXPOSE_KEY_OR_BODY"}, status=401) - if failure == "retryable": - return response(url, {}, status=503) - if failure == "malformed": - return httpx.Response(200, request=httpx.Request("POST", url), content=b"invalid-json") - return response(url, {"id": "wrong-id", "trace_id": "wrong-trace"}) - if "start" in body: - return response(url, {"id": body["start"]["id"], "trace_id": body["start"]["trace_id"]}) - return response(url, {}) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - adapter = WeaveDataTrace(_legacy_weave_transport_config()) - data = call() - expected = { - "rejected": TraceExportRejectedError, - "retryable": RetryableTraceDispatchError, - "malformed": TraceExportUncertainError, - "mismatched_id": TraceExportUncertainError, - }[failure] - if stage == "end": - adapter.start_call(data) - with pytest.raises(expected) as error: - adapter.start_call(data) if stage == "start" else adapter.finish_call(data) - assert "DO_NOT_EXPOSE" not in str(error.value) - assert len(requests) == (1 if stage == "start" else 2) - assert adapter.calls == {} - broken = False - recovered = call("next-call", exception="synthetic tool error") - adapter.start_call(recovered) - adapter.finish_call(recovered) - assert adapter.calls == {} - - @staticmethod - def test_legacy_api_check_uses_read_only_project_and_stats_requests(monkeypatch: pytest.MonkeyPatch) -> None: - sent = [] - - def post(url: str, **options: Unpack[_LegacyWeaveTransportRequestoptions]) -> httpx.Response: - assert options["auth"] == ("api", "fixture-key") - assert options["follow_redirects"] is False - sent.append(url) - if url == "https://api.wandb.ai/graphql": - return response( - url, - { - "data": { - "viewer": {"id": "viewer"}, - "project": {"id": "project-id", "name": "project", "entityName": "team", "readOnly": False}, - } - }, - ) - assert url == "https://trace.wandb.ai/calls/query_stats" - assert options["json"]["project_id"] == "team/project" - return response(url, {"count": 0}) - - monkeypatch.setattr(ssrf_proxy, "post", post) - adapter = WeaveDataTrace(_legacy_weave_transport_config()) - assert sent == [] - assert adapter.api_check() - assert sent == ["https://api.wandb.ai/graphql", "https://trace.wandb.ai/calls/query_stats"] - assert adapter.calls == {} - - @staticmethod - @pytest.mark.parametrize( - ("sent_id", "ack_id"), - [ - ("1234567812345678123456781234abcd", "12345678-1234-5678-1234-56781234abcd"), - ("12345678-1234-5678-1234-56781234abcd", "1234567812345678123456781234ABCD"), - ("12345678-1234-5678-1234-56781234ABCD", "1234567812345678123456781234abcd"), - ], - ) - def test_uuid_equivalent_start_ack_preserves_sent_ids( - monkeypatch: pytest.MonkeyPatch, sent_id: str, ack_id: str - ) -> None: - sent: list[dict[str, object]] = [] - - def send(_method: str, url: str, **options: Unpack[_LegacyWeaveTransportRequestoptions]) -> httpx.Response: - body = json.loads(options["content"]) - sent.append(body) - if "start" in body: - start = CallStartReq.model_validate(body).start - assert start.id == sent_id - assert start.trace_id == sent_id - return response(url, {"id": ack_id, "trace_id": ack_id}) - end = CallEndReq.model_validate(body).end - assert end.id == sent_id - return response(url, {}) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - adapter = WeaveDataTrace(_legacy_weave_transport_config()) - data = call(sent_id) - adapter.start_call(data) - assert adapter.calls[sent_id]["trace_id"] == sent_id - adapter.finish_call(data) - assert len(sent) == 2 - assert adapter.calls == {} - - @staticmethod - @pytest.mark.parametrize("field", ["id", "trace_id"]) - @pytest.mark.parametrize( - ("sent_id", "wrong_ack"), - [ - ("1234567812345678123456781234abcd", "12345678-1234-5678-1234-56781234abce"), - ("legacy-correlation", "LEGACY-CORRELATION"), - ("1234567812345678123456781234abcd", "not-a-uuid"), - ], - ) - def test_start_ack_rejects_different_uuid_or_opaque_identifiers( - monkeypatch: pytest.MonkeyPatch, field: str, sent_id: str, wrong_ack: str - ) -> None: - sent = [] - - def send(_method: str, url: str, **options: Unpack[_LegacyWeaveTransportRequestoptions]) -> httpx.Response: - start = CallStartReq.model_validate_json(options["content"]).start - sent.append(start) - ack = {"id": start.id, "trace_id": start.trace_id} - ack[field] = wrong_ack - return response(url, ack) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - adapter = WeaveDataTrace(_legacy_weave_transport_config()) - with pytest.raises(TraceExportUncertainError, match="did not acknowledge the submitted identifiers"): - adapter.start_call(call(sent_id)) - assert len(sent) == 1 - assert adapter.calls == {} - - -"""Real SDK schema/serialization boundaries with HTTP transport substituted. - -These tests do not claim acceptance by a live third-party account. -""" - - -class TransportOptions(TypedDict, total=False): - content: bytes | str - headers: Mapping[str, str] - auth: tuple[str, str] - max_retries: int - - -def _response(url: str, *, body: object = None, content: bytes | None = None, status: int = 200) -> httpx.Response: - request = httpx.Request("POST", url) - if content is not None: - return httpx.Response(status, request=request, content=content) - return httpx.Response(status, request=request, json=body or {}) - - -def _tool_definitions() -> list[dict[str, object]]: - return [ - { - "type": "function", - "name": "calculate", - "description": "TOOL_CONTENT_CANARY", - "parameters": { - "type": "object", - "properties": {"value": {"type": "integer"}}, - "required": ["value"], - "additionalProperties": False, - "$comment": "Preserve schema extension fields", - }, - "strict": True, - }, - {"type": "function", "function": {"name": "already_nested", "description": "Nested definition"}}, - ] - - -class TestNativeWeaveMlflowCases: - @staticmethod - @pytest.fixture - def trace() -> CanonicalTrace: - start = datetime(2026, 1, 1, tzinfo=UTC) - root = CanonicalSpan( - id="0000000000000001", - parent_id=None, - name="Agent", - kind=CanonicalSpanKind.AGENT, - start_time=start, - end_time=start + timedelta(seconds=2), - inputs={"query": "hello"}, - outputs="answer", - status=CanonicalSpanStatus.OK, - metadata={"total_price": 999, "total_tokens": 999}, - ) - child = CanonicalSpan( - id="0000000000000002", - parent_id=root.id, - name="Model", - kind=CanonicalSpanKind.LLM, - start_time=start, - end_time=start + timedelta(seconds=1), - inputs={"prompt": "hello"}, - outputs="answer", - status=CanonicalSpanStatus.OK, - metadata={ - "model_name": "test-model", - "prompt_tokens": 2, - "completion_tokens": 3, - "total_tokens": 5, - "total_price": "0.01", - "currency": "USD", - }, - ) - return CanonicalTrace( - trace_id="1234567890abcdef1234567890abcdef", - session_id="session-1", - root_span_id=root.id, - spans=(root, child), - ) - - @staticmethod - @pytest.fixture - def weave_started_calls(monkeypatch: pytest.MonkeyPatch) -> list[StartedCallSchemaForInsert]: - starts: list[StartedCallSchemaForInsert] = [] - - def send(_method: str, url: str, **kwargs: Unpack[TransportOptions]) -> httpx.Response: - body = json.loads(kwargs["content"]) - if start := body.get("start"): - started = StartedCallSchemaForInsert.model_validate(start) - starts.append(started) - return _response(url, body={"id": started.id, "trace_id": started.trace_id}) - return _response(url) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - return starts - - @staticmethod - @pytest.mark.parametrize("host", ["https://api.wandb.ai", "https://API.WANDB.AI:443"]) - def test_weave_sdk_exports_tree_session_usage_without_global_credentials( - monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace, host: str - ) -> None: - calls = [] - before = dict(os.environ) - - def send(_method: str, url: str, **kwargs: Unpack[TransportOptions]) -> httpx.Response: - body = json.loads(kwargs["content"]) - calls.append((_method, url, kwargs, body)) - start = body.get("start") - return _response(url, body={"id": start["id"], "trace_id": start["trace_id"]} if start else {}) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - result = export_trace( - "weave", {"api_key": "tenant-key", "entity": "team", "project": "demo", "host": host}, trace - ) - assert len(calls) == 4 - root, child = calls[0][3]["start"], calls[1][3]["start"] - assert child["parent_id"] == root["id"] - assert child["thread_id"] == root["thread_id"] == "session-1" - assert child["turn_id"] == root["turn_id"] == root["id"] - assert calls[2][3]["end"]["summary"]["usage"]["test-model"]["total_tokens"] == 5 - assert calls[3][3]["end"]["summary"] == {} - assert result.trace_url == f"https://wandb.ai/team/demo/r/call/{root['id']}" - assert result.span_urls[trace.spans[1].id].endswith(child["id"]) - assert all(call[2]["auth"] == ("api", "tenant-key") and call[2]["max_retries"] == 0 for call in calls) - assert dict(os.environ) == before - - @staticmethod - @pytest.mark.parametrize("root_kind", [CanonicalSpanKind.CHAIN, CanonicalSpanKind.AGENT]) - def test_weave_nested_agent_calls_form_one_outer_turn( - trace: CanonicalTrace, weave_started_calls: list[StartedCallSchemaForInsert], root_kind: CanonicalSpanKind - ) -> None: - root = trace.spans[0].model_copy(update={"kind": root_kind}) - agent = trace.spans[1].model_copy(update={"kind": CanonicalSpanKind.AGENT}) - model = trace.spans[1].model_copy(update={"id": "0000000000000003", "parent_id": agent.id}) - nested = trace.model_copy(update={"spans": (root, agent, model)}) - - export_trace("weave", {"api_key": "tenant-key", "entity": "team", "project": "demo"}, nested) - - outer, inner_agent, inner_model = weave_started_calls - assert inner_agent.parent_id == outer.id - assert inner_model.parent_id == inner_agent.id - assert all(call.thread_id == trace.session_id and call.turn_id == outer.id for call in weave_started_calls) - # Native Threads queries count only call.id == call.turn_id, including for - # Workflow and single-step preview outer roots with nested Agent details. - assert [call.id for call in weave_started_calls if call.id == call.turn_id] == [outer.id] - - @staticmethod - def test_weave_same_session_has_distinct_stable_turns( - trace: CanonicalTrace, weave_started_calls: list[StartedCallSchemaForInsert] - ) -> None: - config = {"api_key": "tenant-key", "entity": "team", "project": "demo"} - first = export_trace("weave", config, trace) - repeated = export_trace("weave", config, trace) - next_trace = trace.model_copy(update={"trace_id": "abcdef1234567890abcdef1234567890"}) - second = export_trace("weave", config, next_trace) - - first_calls, repeated_calls, second_calls = ( - weave_started_calls[:2], - weave_started_calls[2:4], - weave_started_calls[4:], - ) - assert repeated_calls == first_calls - assert repeated == first - assert second.trace_url != first.trace_url - assert all(call.thread_id == trace.session_id for call in weave_started_calls) - assert first_calls[0].turn_id == first_calls[0].id - assert second_calls[0].turn_id == second_calls[0].id - assert second_calls[0].turn_id != first_calls[0].turn_id - assert all(call.turn_id == first_calls[0].id for call in first_calls) - assert all(call.turn_id == second_calls[0].id for call in second_calls) - - @staticmethod - def test_weave_without_session_does_not_fabricate_a_thread_or_turn( - trace: CanonicalTrace, weave_started_calls: list[StartedCallSchemaForInsert] - ) -> None: - trace = trace.model_copy(update={"session_id": ""}) - - export_trace("weave", {"api_key": "tenant-key", "entity": "team", "project": "demo"}, trace) - - assert len(weave_started_calls) == len(trace.spans) - assert all(call.thread_id is None and call.turn_id is None for call in weave_started_calls) - - @staticmethod - def test_weave_resolves_optional_entity_using_only_supplied_credentials( - monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace - ) -> None: - def send(_method: str, url: str, **kwargs: Unpack[TransportOptions]) -> httpx.Response: - assert kwargs["auth"] == ("api", "tenant-key") - if url.endswith("/graphql"): - return _response(url, body={"data": {"viewer": {"entity": "default-team"}}}) - body = json.loads(kwargs["content"]) - if start := body.get("start"): - assert start["project_id"] == "default-team/demo" - return _response(url, body={"id": start["id"], "trace_id": start["trace_id"]}) - return _response(url) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - result = export_trace("weave", {"api_key": "tenant-key", "project": "demo"}, trace) - assert result.trace_url is not None - assert "/default-team/demo/" in result.trace_url - - @staticmethod - def test_weave_partial_tree_is_uncertain_and_not_retried( - monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace - ) -> None: - calls = [] - - def send(_method: str, url: str, **kwargs: Unpack[TransportOptions]) -> httpx.Response: - calls.append(url) - if len(calls) > 1: - return _response(url, status=503) - start = json.loads(kwargs["content"])["start"] - return _response(url, body={"id": start["id"], "trace_id": start["trace_id"]}) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - with pytest.raises(TraceExportUncertainError, match="partially accepted"): - export_trace("weave", {"api_key": "secret", "entity": "team", "project": "demo"}, trace) - assert len(calls) == 2 - - @staticmethod - @pytest.mark.parametrize("tracking_uri", ["https://mlflow.example", "https://mlflow.example/proxy/mlflow"]) - def test_mlflow_sdk_span_roundtrip_and_credential_isolation( - monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace, tracking_uri: str - ) -> None: - calls = [] - before = dict(os.environ) - - def send(_method: str, url: str, **kwargs: Unpack[TransportOptions]) -> httpx.Response: - calls.append((url, kwargs)) - return _response(url, content=b"") - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - for username in ("tenant-a", "tenant-b"): - result = export_trace( - "mlflow", - { - "tracking_uri": tracking_uri, - "experiment_id": "42", - "username": username, - "password": "password", - }, - trace, - ) - assert result.trace_id == "tr-" + trace.trace_id - assert result.trace_url is not None - fragment_path, _, query = urlsplit(result.trace_url).fragment.partition("?") - assert result.trace_url.startswith(tracking_uri + "/#/") - assert fragment_path == "/experiments/42/traces" - assert parse_qs(query) == { - "traceId": [result.trace_id], - "selectedTraceId": [result.trace_id], - "selectedEvaluationId": [result.trace_id], - } - for index, (url, options) in enumerate(calls): - assert url == tracking_uri + "/v1/traces" - assert options["headers"]["x-mlflow-experiment-id"] == "42" - assert ( - options["headers"]["Authorization"] - == "Basic " + base64.b64encode(f"tenant-{'a' if index == 0 else 'b'}:password".encode()).decode() - ) - content = options["content"] - assert isinstance(content, bytes) - payload = ExportTraceServiceRequest.FromString(content) - spans = [ - _native_weave_mlflow_Span.from_otel_proto(span) - for span in payload.resource_spans[0].scope_spans[0].spans - ] - assert spans[1].parent_id == spans[0].span_id - assert spans[1].get_attribute(SpanAttributeKey.CHAT_USAGE) == { - "input_tokens": 2, - "output_tokens": 3, - "total_tokens": 5, - } - assert spans[1].get_attribute(SpanAttributeKey.LLM_COST) == {"total_cost": 0.01} - assert spans[0].get_attribute(SpanAttributeKey.CHAT_USAGE) is None - assert spans[0].get_attribute(SpanAttributeKey.SESSION_ID) == "session-1" - assert dict(os.environ) == before - - @staticmethod - @pytest.mark.parametrize("serialized", [False, True]) - def test_mlflow_tool_definitions_have_native_chat_schema_without_losing_source( - monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace, serialized: bool - ) -> None: - definitions = _tool_definitions() - source = json.dumps(definitions) if serialized else definitions - trace.spans[1].metadata["gen_ai.tool.definitions"] = source - before = deepcopy(trace.model_dump()) - native: list[_native_weave_mlflow_Span] = [] - - def send(_method: str, url: str, **options: Unpack[TransportOptions]) -> httpx.Response: - content = options["content"] - assert isinstance(content, bytes) - payload = ExportTraceServiceRequest.FromString(content) - native.extend( - _native_weave_mlflow_Span.from_otel_proto(item) - for item in payload.resource_spans[0].scope_spans[0].spans - ) - return _response(url, content=b"") - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - export_trace("mlflow", {"tracking_uri": "https://mlflow.example", "experiment_id": "42"}, trace) - tools = native[1].get_attribute(SpanAttributeKey.CHAT_TOOLS) - assert tools == [ - {"type": "function", "function": {key: value for key, value in definitions[0].items() if key != "type"}}, - definitions[1], - ] - assert isinstance(tools, list) - for tool in tools: - ChatTool.model_validate(tool, strict=True) - assert native[1].get_attribute("gen_ai.tool.definitions") == source - stored = translate_span_when_storing(native[1]) - assert json.loads(stored["attributes"][SpanAttributeKey.CHAT_TOOLS]) == tools - assert trace.model_dump() == before - - @staticmethod - def test_mlflow_redacted_trace_does_not_recreate_tool_definition_content( - monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace - ) -> None: - trace.spans[1].metadata["gen_ai.tool.definitions"] = _tool_definitions() - trace.spans[1].metadata[SpanAttributeKey.CHAT_TOOLS] = [ - {"type": "function", "function": {"name": "TOOL_CONTENT_CANARY"}} - ] - redacted = remap_outer_trace(trace, trace_id=trace.trace_id, capture_content=False, agent_message=False) - native: list[_native_weave_mlflow_Span] = [] - - def send(_method: str, url: str, **options: Unpack[TransportOptions]) -> httpx.Response: - content = options["content"] - assert isinstance(content, bytes) - assert b"TOOL_CONTENT_CANARY" not in content - payload = ExportTraceServiceRequest.FromString(content) - native.extend( - _native_weave_mlflow_Span.from_otel_proto(item) - for item in payload.resource_spans[0].scope_spans[0].spans - ) - return _response(url, content=b"") - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - export_trace("mlflow", {"tracking_uri": "https://mlflow.example", "experiment_id": "42"}, redacted) - for span in native: - assert span.get_attribute(SpanAttributeKey.CHAT_TOOLS) is None - assert span.get_attribute("gen_ai.tool.definitions") is None - assert span.inputs is None - assert span.outputs is None - - @staticmethod - @pytest.mark.parametrize("definitions", [[None, 12, {"type": "other"}, {"type": "function"}], "invalid-json"]) - def test_mlflow_unsupported_tool_shapes_remain_in_metadata_without_fake_functions( - trace: CanonicalTrace, definitions: object - ) -> None: - from dify_trace_mlflow.agent_trace import _spans - - trace.spans[1].metadata["gen_ai.tool.definitions"] = definitions - before = deepcopy(trace.model_dump()) - result = _spans(trace, "42", "tr-" + trace.trace_id) - assert result[1].get_attribute("gen_ai.tool.definitions") == definitions - assert result[1].get_attribute(SpanAttributeKey.CHAT_TOOLS) == ([] if isinstance(definitions, list) else None) - assert trace.model_dump() == before - - @staticmethod - def test_mlflow_existing_native_chat_tools_take_precedence(trace: CanonicalTrace) -> None: - from dify_trace_mlflow.agent_trace import _spans - - native_tools = [{"type": "function", "function": {"name": "already_native"}}] - trace.spans[1].metadata[SpanAttributeKey.CHAT_TOOLS] = native_tools - trace.spans[1].metadata["gen_ai.tool.definitions"] = _tool_definitions() - result = _spans(trace, "42", "tr-" + trace.trace_id) - assert result[1].get_attribute(SpanAttributeKey.CHAT_TOOLS) == native_tools - - @staticmethod - def test_mlflow_preserves_json_schema_beyond_the_sdk_chat_tool_model(trace: CanonicalTrace) -> None: - from dify_trace_mlflow.agent_trace import _spans - - parameters = { - "$ref": "#/$defs/Arguments", - "$defs": { - "Arguments": { - "type": "object", - "properties": {"choice": {"enum": [None, {"mode": "auto"}]}}, - "additionalProperties": {"type": "string"}, - } - }, - } - source = [{"type": "function", "name": "schema_extension", "parameters": parameters}] - trace.spans[1].metadata["gen_ai.tool.definitions"] = source - result = _spans(trace, "42", "tr-" + trace.trace_id) - expected = [{"type": "function", "function": {"name": "schema_extension", "parameters": parameters}}] - assert result[1].get_attribute(SpanAttributeKey.CHAT_TOOLS) == expected - assert result[1].get_attribute("gen_ai.tool.definitions") == source - stored = translate_span_when_storing(result[1]) - assert json.loads(stored["attributes"][SpanAttributeKey.CHAT_TOOLS]) == expected - - @staticmethod - def test_mlflow_partial_otlp_ack_is_uncertain(monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace) -> None: - ack = ExportTraceServiceResponse() - ack.partial_success.rejected_spans = 1 - monkeypatch.setattr( - ssrf_proxy, "make_request", lambda _method, url, **_kwargs: _response(url, content=ack.SerializeToString()) - ) - with pytest.raises(TraceExportUncertainError, match="partially accepted"): - export_trace("mlflow", {"tracking_uri": "https://mlflow.example", "experiment_id": "42"}, trace) - - @staticmethod - @pytest.mark.parametrize( - "credential_type", ["AWS_PRESIGNED_URL", "GCP_SIGNED_URL", "AZURE_SAS_URI", "AZURE_ADLS_GEN2_SAS_URI"] - ) - @pytest.mark.parametrize("oauth", [False, True]) - def test_databricks_sdk_v3_metadata_and_native_artifact( - monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace, credential_type: str, oauth: bool - ) -> None: - calls: list[tuple[str, str, TransportOptions]] = [] - uploaded: list[TraceData] = [] - before = dict(os.environ) - server_trace_id = "tr-abcdef0123456789abcdef0123456789" - definitions = _tool_definitions() - trace.spans[1].metadata["gen_ai.tool.definitions"] = definitions - - def send(_method: str, url: str, **kwargs: Unpack[TransportOptions]) -> httpx.Response: - calls.append((_method, url, kwargs)) - if url.endswith("/oidc/v1/token"): - assert kwargs["auth"] == ("client", "client-secret") - return _response(url, body={"access_token": "oauth-token"}) - if url.endswith("/api/3.0/mlflow/traces"): - assert kwargs["headers"]["Authorization"] == ("Bearer oauth-token" if oauth else "Bearer pat-token") - body = json.loads(kwargs["content"]) - assert body["trace"]["trace_info"]["trace_metadata"]["mlflow.trace.session"] == "session-1" - body["trace"]["trace_info"]["trace_id"] = server_trace_id - return _response(url, body=body) - if url.endswith("/credentials-for-data-upload"): - assert f"/traces/{server_trace_id}/" in url - return _response( - url, - body={ - "credential_info": { - "signed_uri": "https://storage.example/trace?sig=fake", - "type": credential_type, - "headers": [], - } - }, - ) - assert "Authorization" not in kwargs["headers"] - if kwargs.get("content"): - uploaded.append(TraceData.from_dict(json.loads(kwargs["content"]))) - return _response(url) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - config = {"host": "https://databricks.example", "experiment_id": "42"} - config.update( - {"client_id": "client", "client_secret": "client-secret"} - if oauth - else {"personal_access_token": "pat-token"} - ) - result = export_trace("databricks", config, trace) - assert result.trace_id == server_trace_id - assert result.trace_url == ( - f"https://databricks.example/ml/experiments/42/traces?selectedEvaluationId={server_trace_id}" - ) - data = uploaded[0] - assert len(data.spans) == 2 - assert all(span.trace_id == server_trace_id for span in data.spans) - assert data.spans[1].parent_id == data.spans[0].span_id - usage = data.spans[1].get_attribute(SpanAttributeKey.CHAT_USAGE) - assert usage is not None - assert usage["total_tokens"] == 5 - assert data.spans[1].get_attribute(SpanAttributeKey.LLM_COST) == {"total_cost": 0.01} - tools = data.spans[1].get_attribute(SpanAttributeKey.CHAT_TOOLS) - assert isinstance(tools, list) - assert len(tools) == 2 - for tool in tools: - ChatTool.model_validate(tool, strict=True) - assert data.spans[1].get_attribute("gen_ai.tool.definitions") == definitions - assert dict(os.environ) == before - assert len(calls) == 3 + int(oauth) + int(credential_type == "AZURE_ADLS_GEN2_SAS_URI") - - @staticmethod - def test_databricks_partial_metadata_delivery_is_uncertain( - monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace - ) -> None: - def send(_method: str, url: str, **kwargs: Unpack[TransportOptions]) -> httpx.Response: - if url.endswith("/api/3.0/mlflow/traces"): - return _response(url, body=json.loads(kwargs["content"])) - return _response(url, status=503) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - with pytest.raises(TraceExportUncertainError, match="metadata"): - export_trace( - "databricks", - {"host": "https://databricks.example", "experiment_id": "42", "personal_access_token": "secret"}, - trace, - ) - - @staticmethod - @pytest.mark.parametrize("provider", ["weave", "databricks"]) - def test_success_status_with_unreadable_ack_is_uncertain( - monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace, provider: str - ) -> None: - monkeypatch.setattr( - ssrf_proxy, "make_request", lambda _method, url, **_kwargs: _response(url, content=b"not-json") - ) - config = { - "api_key": "secret", - "entity": "team", - "project": "demo", - "host": "https://example.com", - "experiment_id": "42", - "personal_access_token": "secret", - } - with pytest.raises(TraceExportUncertainError, match="unreadable"): - export_trace(provider, config, trace) - - -"""Exercise anonymous Phoenix auth at the HTTP/1.1 and legacy SDK boundaries.""" - - -class TestPhoenixExportAuthCases: - @staticmethod - @pytest.fixture - def trace() -> CanonicalTrace: - start = datetime(2026, 10, 20, tzinfo=UTC) - return CanonicalTrace( - trace_id="1" * 32, - session_id="phoenix-auth-session", - root_span_id="2" * 16, - spans=( - CanonicalSpan( - id="2" * 16, - parent_id=None, - name="Phoenix auth regression", - kind=CanonicalSpanKind.AGENT, - start_time=start, - end_time=start + timedelta(seconds=1), - status=CanonicalSpanStatus.OK, - ), - ), - ) - - @staticmethod - @pytest.fixture - def http11_wire(monkeypatch: pytest.MonkeyPatch) -> Iterator[list[bytes]]: - """Keep real HTTPX/httpcore/h11 serialization, replacing only socket I/O.""" - chunks: list[bytes] = [] - - def write(_stream: httpcore.MockStream, buffer: bytes, timeout: float | None = None) -> None: - del timeout - chunks.append(buffer) - - monkeypatch.setattr(httpcore.MockStream, "write", write) - transport = httpx.HTTPTransport() - transport.close() - transport._pool = httpcore.ConnectionPool( - network_backend=httpcore.MockBackend([b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"]) - ) - with httpx.Client(transport=transport, trust_env=False) as client: - monkeypatch.setattr(ssrf_proxy, "_get_ssrf_client", lambda _verify: client) - yield chunks - - @staticmethod - def test_managed_exports_serialize_optional_auth_without_cross_destination_credentials( - monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace, http11_wire: list[bytes] - ) -> None: - monkeypatch.setenv( - "OTEL_EXPORTER_OTLP_HEADERS", "authorization=Bearer%20ambient-general,api_key=ambient-general" - ) - monkeypatch.setenv( - "OTEL_EXPORTER_OTLP_TRACES_HEADERS", "authorization=Bearer%20ambient-traces,api_key=ambient-traces" - ) - - for index, key in enumerate(("first-synthetic-key", "", None, "second-synthetic-key", None, "")): - previous_chunks = len(http11_wire) - result = export_trace( - "phoenix", - {"endpoint": "http://phoenix.test/proxy", "api_key": key, "project": f"project-{index}"}, - trace, - ) - assert result.status == "accepted" - assert result.trace_id == trace.trace_id - head, separator, body = b"".join(http11_wire[previous_chunks:]).partition(b"\r\n\r\n") - assert separator - lines = head.split(b"\r\n") - assert lines[0] == b"POST /proxy/v1/traces HTTP/1.1" - header_values: list[tuple[bytes, bytes]] = [] - for line in lines[1:]: - name, value = line.split(b": ", 1) - header_values.append((name, value)) - headers = httpx.Headers(header_values) - assert headers["Content-Type"] == "application/x-protobuf" - if key: - assert headers["api_key"] == key - assert headers["authorization"] == f"Bearer {key}" - else: - assert "api_key" not in headers - assert "authorization" not in headers - payload = ExportTraceServiceRequest.FromString(body) - resource = payload.resource_spans[0] - assert {item.key: item.value.string_value for item in resource.resource.attributes}[ - "openinference.project.name" - ] == f"project-{index}" - assert resource.scope_spans[0].spans[0].trace_id.hex() == trace.trace_id - - @staticmethod - def test_legacy_sdk_prepares_anonymous_requests_without_ambient_or_previous_auth( - monkeypatch: pytest.MonkeyPatch, - ) -> None: - monkeypatch.setenv( - "OTEL_EXPORTER_OTLP_HEADERS", "authorization=Bearer%20ambient-general,api_key=ambient-general" - ) - monkeypatch.setenv( - "OTEL_EXPORTER_OTLP_TRACES_HEADERS", "authorization=Bearer%20ambient-traces,api_key=ambient-traces" - ) - monkeypatch.delenv("OTEL_EXPORTER_OTLP_COMPRESSION", raising=False) - monkeypatch.delenv("OTEL_EXPORTER_OTLP_TRACES_COMPRESSION", raising=False) - prepared: list[requests.PreparedRequest] = [] - - def send(_session: requests.Session, request: requests.PreparedRequest, **_kwargs: object) -> requests.Response: - prepared.append(request) - response = requests.Response() - response.status_code = 200 - response._content = b"" - response.request = request - return response - - monkeypatch.setattr(requests.Session, "send", send) - for index, key in enumerate(("first-synthetic-key", "", None, "second-synthetic-key", None, "")): - tracer, processor = setup_tracer( - PhoenixConfig(endpoint="http://phoenix.test/proxy", api_key=key, project=f"project-{index}") - ) - try: - tracer.start_span("Legacy Phoenix auth regression").end() - finally: - processor.shutdown() - assert len(prepared) == index + 1 - request = prepared[-1] - assert request.method == "POST" - assert request.url == "http://phoenix.test/proxy/v1/traces" - assert request.headers["Content-Type"] == "application/x-protobuf" - if key: - assert request.headers["api_key"] == key - assert request.headers["authorization"] == f"Bearer {key}" - else: - assert "api_key" not in request.headers - assert "authorization" not in request.headers - assert isinstance(request.body, bytes) - payload = ExportTraceServiceRequest.FromString(request.body) - assert {item.key: item.value.string_value for item in payload.resource_spans[0].resource.attributes}[ - "openinference.project.name" - ] == f"project-{index}" - - -"""Exercise real vendor serializers and OTLP stubs at their network boundary.""" - - -class HTTPArguments(TypedDict): - content: bytes - headers: Mapping[str, str] - max_retries: NotRequired[int] - follow_redirects: NotRequired[bool] - - -type HTTPCall = tuple[str, str, HTTPArguments] - - -def spans_from_call(call: HTTPCall) -> list[_provider_dispatch_Span]: - request = ExportTraceServiceRequest.FromString(call[2]["content"]) - return list(request.resource_spans[0].scope_spans[0].spans) - - -def attributes(span: _provider_dispatch_Span) -> dict[str, object]: - result: dict[str, object] = {} - for item in span.attributes: - kind = item.value.WhichOneof("value") - assert kind is not None - result[item.key] = { - "string_value": item.value.string_value, - "bool_value": item.value.bool_value, - "int_value": item.value.int_value, - "double_value": item.value.double_value, - "array_value": item.value.array_value, - "kvlist_value": item.value.kvlist_value, - "bytes_value": item.value.bytes_value, - }[kind] - return result - - -def _langfuse_queue_job_ack() -> dict[str, object]: - # Synthetic values only. Real acknowledgements contain auth context and must - # never be logged or persisted as acceptance-test evidence. - return { - "name": "otel-ingestion-job", - "data": { - "id": "11111111-2222-4333-8444-555555555555", - "timestamp": "2026-10-15T01:00:00.000Z", - "name": "otel-ingestion-job", - "payload": { - "data": {"fileKey": "otel/test/payload.json", "publicKey": "never-log-public-key"}, - "authCheck": { - "validKey": True, - "scope": {"projectId": "never-log-project", "accessLevel": "project", "orgId": "test-org"}, - }, - "sdkName": "test-sdk", - "ingestionVersion": "4", - }, - }, - } - - -def _replace_ack_value(ack: dict[str, object], path: tuple[str, ...], value: object) -> None: - for field in path[:-1]: - nested = ack[field] - assert isinstance(nested, dict) - ack = cast(dict[str, object], nested) - ack[path[-1]] = value - - -class TestProviderDispatchCases: - @staticmethod - @pytest.fixture - def trace() -> CanonicalTrace: - start = datetime(2026, 10, 15, 1, tzinfo=UTC) - return CanonicalTrace( - trace_id="1" * 32, - session_id="conversation-1", - root_span_id="1" * 16, - spans=( - CanonicalSpan( - id="1" * 16, - parent_id=None, - name="Agent turn", - kind=CanonicalSpanKind.AGENT, - start_time=start, - end_time=start + timedelta(seconds=3), - status=CanonicalSpanStatus.OK, - inputs={"question": "hi"}, - outputs={"answer": "done"}, - metadata={ - "app_id": "app-one", - "tenant_id": "tenant-one", - "total_tokens": 12, - "total_price": "0.03", - }, - ), - CanonicalSpan( - id="2" * 16, - parent_id="1" * 16, - name="model call", - kind=CanonicalSpanKind.LLM, - start_time=start, - end_time=start + timedelta(seconds=1), - status=CanonicalSpanStatus.OK, - inputs=[{"role": "user", "content": "hi"}], - outputs="done", - metadata={ - "model": "gpt-4o-mini", - "model_provider": "openai", - "prompt_tokens": 8, - "completion_tokens": 4, - "total_tokens": 12, - "total_price": "0.03", - "currency": "USD", - "ttft": 0.2, - }, - ), - CanonicalSpan( - id="3" * 16, - parent_id="1" * 16, - name="search", - kind=CanonicalSpanKind.TOOL, - start_time=start + timedelta(seconds=1), - end_time=start + timedelta(seconds=2), - status=CanonicalSpanStatus.ERROR, - inputs={"q": "x"}, - outputs=None, - error="tool failed", - ), - ), - ) - - @staticmethod - @pytest.fixture - def http_calls(monkeypatch: pytest.MonkeyPatch) -> list[HTTPCall]: - from core.helper import ssrf_proxy - - calls: list[HTTPCall] = [] - - def send(method: str, url: str, **kwargs: object) -> httpx.Response: - calls.append((method, url, cast(HTTPArguments, kwargs))) - return httpx.Response(200, content=b"", request=httpx.Request(method, url)) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - return calls - - @staticmethod - def test_langfuse_v4_tree_and_native_usage(trace: CanonicalTrace, http_calls: list[HTTPCall]) -> None: - config = { - "host": "https://us.cloud.langfuse.com", - "public_key": "pk-test", - "secret_key": "sk-test", - "project_url": "https://us.cloud.langfuse.com/project/example", - } - result = export_trace("langfuse", config, trace) - assert result.status == "accepted" - assert result.trace_id == trace.trace_id - assert len(http_calls) == 1 - _, url, args = http_calls[0] - assert url.endswith("/api/public/otel/v1/traces") - assert args["headers"]["x-langfuse-ingestion-version"] == "4" - assert args["max_retries"] == 0 - assert args["follow_redirects"] is False - spans = spans_from_call(http_calls[0]) - assert [span.span_id.hex() for span in spans] == [s.id for s in trace.spans] - assert spans[1].parent_span_id == spans[0].span_id - assert spans[2].status.code == 2 - root, llm = attributes(spans[0]), attributes(spans[1]) - assert root["langfuse.observation.type"] == "agent" - assert "langfuse.observation.usage_details" not in root - usage_details, cost_details = ( - llm["langfuse.observation.usage_details"], - llm["langfuse.observation.cost_details"], - ) - assert isinstance(usage_details, str) - assert isinstance(cost_details, str) - assert json.loads(usage_details) == {"input": 8, "output": 4, "total": 12} - assert json.loads(cost_details) == {"total": 0.03} - assert all(attributes(span)["langfuse.session.id"] == trace.session_id for span in spans) - assert result.span_urls[trace.spans[1].id].endswith(trace.spans[1].id) - - @staticmethod - @pytest.mark.parametrize( - ("provider", "config", "suffix", "kind_key", "kind"), - [ - ( - "phoenix", - {"endpoint": "https://phoenix.test/proxy", "api_key": "key", "project": "one"}, - "/proxy/v1/traces", - "openinference.span.kind", - "AGENT", - ), - ( - "aliyun", - {"endpoint": "https://region.log.aliyuncs.com", "license_key": "license", "app_name": "one"}, - "/adapt_license/api/v1/traces", - "gen_ai.span.kind", - "AGENT", - ), - ], - ) - def test_other_http_providers_have_native_semantics( - trace: CanonicalTrace, - http_calls: list[HTTPCall], - provider: str, - config: dict[str, str], - suffix: str, - kind_key: str, - kind: str, - ) -> None: - result = export_trace(provider, config, trace) - assert result.status == "accepted" - assert http_calls[0][1].endswith(suffix) - spans = spans_from_call(http_calls[0]) - assert attributes(spans[0])[kind_key] == kind - token_key = "llm.token_count.prompt" if provider == "phoenix" else "gen_ai.usage.input_tokens" - assert attributes(spans[1])[token_key] == 8 - assert token_key not in attributes(spans[0]) - - @staticmethod - @pytest.mark.parametrize( - ("endpoint_url", "suffix"), - [ - ("https://log.aliyuncs.com", "api/v1/traces"), - ("https://cn-hangzhou.log.aliyuncs.com", "api/v1/traces"), - ("https://project.cn-heyuan.log.aliyuncs.com/apm/trace/opentelemetry", "api/v1/traces"), - ("https://project.cn-heyuan-intranet.log.aliyuncs.com/apm/trace/opentelemetry", "api/v1/traces"), - ("https://cn-shanghai-share.log.aliyuncs.com", "api/v1/traces"), - ("https://cn-beijing.dualstack.log.aliyuncs.com", "api/v1/traces"), - ("https://cn-hangzhou.log.aliyuncs.com.:443", "api/v1/traces"), - ("https://tracing-analysis-dc-hz.aliyuncs.com", "api/otlp/traces"), - ("https://collector.example.test:4318/custom-prefix", "api/otlp/traces"), - ("https://evil-log.aliyuncs.com.evil.test", "api/otlp/traces"), - ("https://log.aliyuncs.com.evil.test", "api/otlp/traces"), - ("https://evil-log.aliyuncs.com", "api/otlp/traces"), - ("https://collector.example.test/log.aliyuncs.com", "api/otlp/traces"), - ], - ) - def test_aliyun_endpoint_family_uses_dns_labels_and_preserves_custom_destinations( - trace: CanonicalTrace, http_calls: list[HTTPCall], endpoint_url: str, suffix: str - ) -> None: - result = export_trace( - "aliyun", {"endpoint": endpoint_url, "license_key": "synthetic-license", "app_name": "test-app"}, trace - ) - assert result.status == "accepted" - assert len(http_calls) == 1 - assert http_calls[0][1] == f"{endpoint_url}/adapt_synthetic-license/{suffix}" - spans = spans_from_call(http_calls[0]) - assert [span.span_id.hex() for span in spans] == [span.id for span in trace.spans] - assert spans[1].parent_span_id == spans[0].span_id - - @staticmethod - def test_aliyun_raw_license_is_encoded_once_at_the_real_otlp_boundary( - trace: CanonicalTrace, http_calls: list[HTTPCall] - ) -> None: - result = export_trace( - "aliyun", - {"endpoint": "https://COLLECTOR.test:8443/prefix/", "license_key": "raw%2F@a+b/c", "app_name": "test_app"}, - trace, - ) - assert result.status == "accepted" - assert len(http_calls) == 1 - assert http_calls[0][1] == "https://collector.test:8443/prefix/adapt_raw%252F@a%2Bb%2Fc/api/otlp/traces" - assert [span.span_id.hex() for span in spans_from_call(http_calls[0])] == [span.id for span in trace.spans] - - @staticmethod - def test_aliyun_complete_token_url_is_rejected_without_any_ingestion( - trace: CanonicalTrace, http_calls: list[HTTPCall] - ) -> None: - with pytest.raises(TraceExportRejectedError, match="Aliyun tracing configuration is invalid"): - export_trace( - "aliyun", - {"endpoint": "https://collector.test/adapt_private/api/otlp/traces", "license_key": "private"}, - trace, - ) - assert http_calls == [] - - @staticmethod - @pytest.mark.parametrize( - ("provider", "config", "header"), - [ - ( - "arize", - {"endpoint": "https://otlp.arize.com", "api_key": "key", "space_id": "space", "project": "one"}, - ("space_id", "space"), - ), - ( - "tencent", - {"endpoint": "https://apm.test:4317", "token": "token", "service_name": "one"}, - ("authorization", "Bearer token"), - ), - ], - ) - def test_grpc_real_stub_serialization( - trace: CanonicalTrace, - monkeypatch: pytest.MonkeyPatch, - provider: str, - config: dict[str, str], - header: tuple[str, str], - ) -> None: - calls: list[tuple[ExportTraceServiceRequest, Sequence[tuple[str, str]]]] = [] - - class Channel: - def unary_unary( - self, - method: str, - request_serializer: Callable[[ExportTraceServiceRequest], bytes], - response_deserializer: Callable[[bytes], ExportTraceServiceResponse], - **_kwargs: object, - ) -> Callable[..., ExportTraceServiceResponse]: - assert method == "/opentelemetry.proto.collector.trace.v1.TraceService/Export" - - def invoke( - payload: ExportTraceServiceRequest, metadata: Sequence[tuple[str, str]], timeout: float - ) -> ExportTraceServiceResponse: - assert timeout == 30 - calls.append((ExportTraceServiceRequest.FromString(request_serializer(payload)), metadata)) - return response_deserializer(ExportTraceServiceResponse().SerializeToString()) - - return invoke - - def close(self) -> None: - pass - - monkeypatch.setattr("grpc.secure_channel", lambda *_args, **_kwargs: Channel()) - result = export_trace(provider, config, trace) - assert result.status == "accepted" - assert header in calls[0][1] - spans = calls[0][0].resource_spans[0].scope_spans[0].spans - assert spans[1].parent_span_id == spans[0].span_id - - @staticmethod - def test_langsmith_real_sdk_uses_isolated_sync_runs(trace: CanonicalTrace, http_calls: list[HTTPCall]) -> None: - result = export_trace( - "langsmith", {"api_key": "ls-test", "endpoint": "https://langsmith.test", "project": "one"}, trace - ) - writes = [json.loads(call[2]["content"]) for call in http_calls if call[0] == "POST"] - assert len(writes) == 3 - root, llm, tool = writes - assert root["run_type"] == "chain" - assert llm["run_type"] == "llm" - assert llm["parent_run_id"] == root["id"] - assert llm["dotted_order"].startswith(root["dotted_order"]) - assert llm["input_tokens"] == 8 - assert llm["output_tokens"] == 4 - assert llm["total_cost"] == 0.03 - assert "input_tokens" not in root - assert all(write["extra"]["metadata"]["session_id"] == trace.session_id for write in writes) - assert tool["error"] == "tool failed" - assert UUID(result.trace_id) - - @staticmethod - def test_opik_real_sdk_preserves_tree_thread_and_native_usage( - trace: CanonicalTrace, http_calls: list[HTTPCall] - ) -> None: - result = export_trace( - "opik", - {"api_key": "opik-test", "url": "https://opik.test/api", "workspace": "org-one", "project": "one"}, - trace, - ) - assert len(http_calls) == 2 - root = json.loads(http_calls[0][2]["content"]) - batch = json.loads(http_calls[1][2]["content"]) - spans = batch["spans"] - assert UUID(result.trace_id).version == 7 - assert root["thread_id"] == trace.session_id - assert spans[1]["parent_span_id"] == spans[0]["id"] - assert spans[1]["usage"]["prompt_tokens"] == 8 - assert spans[1]["total_estimated_cost"] == 0.03 - assert not spans[0].get("usage") - assert spans[2]["error_info"]["message"] == "tool failed" - assert http_calls[0][2]["headers"]["comet-workspace"] == "org-one" - - @staticmethod - @pytest.mark.parametrize( - ("value", "wire_value"), - [ - (None, None), - ("323", "323"), - ({"result": 323}, {"result": 323}), - ([{"result": 323}], [{"result": 323}]), - ([], []), - (323, "323"), - (0, "0"), - (-5, "-5"), - (3.5, "3.5"), - (True, "true"), - (False, "false"), - ([1, 2], "[1, 2]"), - (["one", "two"], '["one", "two"]'), - ([{"result": 1}, "text"], '[{"result": 1}, "text"]'), - ([[1], [2]], "[[1], [2]]"), - ], - ) - def test_opik_real_sdk_preserves_json_content_across_trace_and_span_boundaries( - trace: CanonicalTrace, http_calls: list[HTTPCall], value: object, wire_value: object - ) -> None: - content_trace = trace.model_copy( - update={"spans": tuple(span.model_copy(update={"inputs": value, "outputs": value}) for span in trace.spans)} - ) - result = export_trace( - "opik", - {"api_key": "synthetic-key", "url": "https://opik.test/api", "workspace": "team", "project": "one"}, - content_trace, - ) - assert result.status == "accepted" - assert len(http_calls) == 2 - root = json.loads(http_calls[0][2]["content"]) - spans = json.loads(http_calls[1][2]["content"])["spans"] - for payload in [root, *spans]: - assert payload.get("input") == wire_value - assert payload.get("output") == wire_value - assert spans[1]["parent_span_id"] == spans[0]["id"] - assert spans[2]["parent_span_id"] == spans[0]["id"] - assert root["thread_id"] == trace.session_id - assert spans[1]["usage"]["total_tokens"] == 12 - assert not spans[0].get("usage") - assert not spans[2].get("usage") - - @staticmethod - @pytest.mark.parametrize("capture_content", [True, False]) - @pytest.mark.parametrize("root_error", [True, False]) - def test_opik_trace_error_matches_sanitized_root_and_not_handled_tool_failure( - trace: CanonicalTrace, http_calls: list[HTTPCall], capture_content: bool, root_error: bool - ) -> None: - root = trace.spans[0].model_copy( - update={ - "status": CanonicalSpanStatus.ERROR if root_error else CanonicalSpanStatus.OK, - "error": "private-model-error" if root_error else None, - "inputs": {"query": "private-error-input"}, - } - ) - source = trace.model_copy(update={"spans": (root, *trace.spans[1:])}) - canonical = remap_outer_trace( - source, trace_id=trace.trace_id, capture_content=capture_content, agent_message=True - ) - export_trace( - "opik", - {"api_key": "synthetic-key", "url": "https://opik.test/api", "workspace": "team", "project": "one"}, - canonical, - ) - root_payload = json.loads(http_calls[0][2]["content"]) - spans = json.loads(http_calls[1][2]["content"])["spans"] - assert spans[2]["error_info"] is not None - if root_error: - expected = { - "exception_type": "DifyExecutionError", - "message": "private-model-error" if capture_content else "Application execution failed", - "traceback": "", - } - assert root_payload["error_info"] == spans[0]["error_info"] == expected - else: - assert root_payload.get("error_info") is None - assert spans[0].get("error_info") is None - if not capture_content: - assert all(b"private-model-error" not in call[2]["content"] for call in http_calls) - assert all(b"private-error-input" not in call[2]["content"] for call in http_calls) - - @staticmethod - @pytest.mark.parametrize( - ("status", "error"), - [(401, TraceExportRejectedError), (429, RetryableTraceDispatchError), (503, RetryableTraceDispatchError)], - ) - def test_explicit_http_rejections_are_classified_without_response_body( - trace: CanonicalTrace, monkeypatch: pytest.MonkeyPatch, status: int, error: type[Exception] - ) -> None: - from core.helper import ssrf_proxy - - monkeypatch.setattr( - ssrf_proxy, - "make_request", - lambda method, url, **_kwargs: httpx.Response( - status, text="secret-token", request=httpx.Request(method, url) - ), - ) - with pytest.raises(error) as caught: - export_trace( - "langfuse", {"host": "https://langfuse.test", "public_key": "pk", "secret_key": "secret-token"}, trace - ) - assert "secret-token" not in str(caught.value) - - @staticmethod - def test_partial_otlp_acceptance_requires_reconciliation( - trace: CanonicalTrace, monkeypatch: pytest.MonkeyPatch - ) -> None: - from core.helper import ssrf_proxy - - response = ExportTraceServiceResponse() - response.partial_success.rejected_spans = 1 - monkeypatch.setattr( - ssrf_proxy, - "make_request", - lambda method, url, **_kwargs: httpx.Response( - 200, content=response.SerializeToString(), request=httpx.Request(method, url) - ), - ) - with pytest.raises(TraceExportUncertainError): - export_trace("langfuse", {"host": "https://langfuse.test", "public_key": "pk", "secret_key": "sk"}, trace) - - @staticmethod - def test_langfuse_v4_queue_ack_confirms_ingestion_acceptance_only( - trace: CanonicalTrace, - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - capsys: pytest.CaptureFixture[str], - ) -> None: - from core.helper import ssrf_proxy - - monkeypatch.setattr( - ssrf_proxy, - "make_request", - lambda method, url, **_kwargs: httpx.Response( - 200, json=_langfuse_queue_job_ack(), request=httpx.Request(method, url) - ), - ) - result = export_trace( - "langfuse", {"host": "https://langfuse.test", "public_key": "pk", "secret_key": "sk"}, trace - ) - assert result.status == "accepted" - assert result.trace_id == trace.trace_id - assert result.trace_url is None - captured = capsys.readouterr() - evidence = caplog.text + captured.out + captured.err + repr(result) - assert "never-log-public-key" not in evidence - assert "never-log-project" not in evidence - assert "authCheck" not in evidence - - @staticmethod - @pytest.mark.parametrize( - "scenario", - [ - "wrong_name", - "wrong_nested_name", - "invalid_id", - "invalid_timestamp", - "invalid_auth", - "missing_scope", - "missing_file", - "partial_success", - "error", - "failed_job", - "unknown", - "null", - "array", - ], - ) - def test_langfuse_ambiguous_json_ack_remains_uncertain( - trace: CanonicalTrace, - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - capsys: pytest.CaptureFixture[str], - scenario: str, - ) -> None: - from core.helper import ssrf_proxy - - ack = _langfuse_queue_job_ack() - body: object = ack - if scenario == "wrong_name": - ack["name"] = "different-job" - elif scenario == "wrong_nested_name": - _replace_ack_value(ack, ("data", "name"), "different-job") - elif scenario == "invalid_id": - _replace_ack_value(ack, ("data", "id"), "not-a-uuid") - elif scenario == "invalid_timestamp": - _replace_ack_value(ack, ("data", "timestamp"), "2026-10-15T01:00:00") - elif scenario == "invalid_auth": - _replace_ack_value(ack, ("data", "payload", "authCheck", "validKey"), False) - elif scenario == "missing_scope": - _replace_ack_value(ack, ("data", "payload", "authCheck", "scope"), dict[str, object]()) - elif scenario == "missing_file": - _replace_ack_value(ack, ("data", "payload", "data", "fileKey"), "") - elif scenario == "partial_success": - ack["partialSuccess"] = {"rejectedSpans": "1"} - elif scenario == "error": - ack["error"] = "never-log-public-key" - elif scenario == "failed_job": - ack["failedReason"] = "never-log-public-key" - elif scenario == "unknown": - body = {"accepted": True, "publicKey": "never-log-public-key"} - elif scenario == "null": - body = None - elif scenario == "array": - body = list[object]() - monkeypatch.setattr( - ssrf_proxy, - "make_request", - lambda method, url, **_kwargs: httpx.Response( - 200, - content=json.dumps(body).encode(), - headers={"Content-Type": "application/json"}, - request=httpx.Request(method, url), - ), - ) - with pytest.raises(TraceExportUncertainError) as caught: - export_trace("langfuse", {"host": "https://langfuse.test", "public_key": "pk", "secret_key": "sk"}, trace) - captured = capsys.readouterr() - evidence = caplog.text + captured.out + captured.err + str(caught.value) - assert "never-log-public-key" not in evidence - assert "never-log-project" not in evidence - assert "authCheck" not in evidence - - @staticmethod - @pytest.mark.parametrize( - ("provider", "version", "status"), [("phoenix", "4", 200), ("langfuse", "3", 200), ("langfuse", "4", 202)] - ) - def test_langfuse_queue_ack_is_not_a_generic_otlp_success( - monkeypatch: pytest.MonkeyPatch, provider: str, version: str, status: int - ) -> None: - from core.helper import ssrf_proxy - from services.tracing.otlp_export import export_http - - monkeypatch.setattr( - ssrf_proxy, - "make_request", - lambda method, url, **_kwargs: httpx.Response( - status, json=_langfuse_queue_job_ack(), request=httpx.Request(method, url) - ), - ) - with pytest.raises(TraceExportUncertainError): - export_http( - provider, - "https://collector.test/v1/traces", - {"x-langfuse-ingestion-version": version}, - ExportTraceServiceRequest(), - ) - - @staticmethod - @pytest.mark.parametrize("body", [{}, {"partialSuccess": {}}, {"partialSuccess": {"rejectedSpans": "1"}}]) - def test_langfuse_standard_otlp_json_ack_remains_supported( - trace: CanonicalTrace, monkeypatch: pytest.MonkeyPatch, body: dict[str, dict[str, str]] - ) -> None: - from core.helper import ssrf_proxy - - monkeypatch.setattr( - ssrf_proxy, - "make_request", - lambda method, url, **_kwargs: httpx.Response(200, json=body, request=httpx.Request(method, url)), - ) - config = {"host": "https://langfuse.test", "public_key": "pk", "secret_key": "sk"} - if body.get("partialSuccess", {}).get("rejectedSpans"): - with pytest.raises(TraceExportUncertainError): - export_trace("langfuse", config, trace) - else: - assert export_trace("langfuse", config, trace).status == "accepted" - - @staticmethod - def test_lost_http_ack_is_uncertain_not_rejected(trace: CanonicalTrace, monkeypatch: pytest.MonkeyPatch) -> None: - from core.helper import ssrf_proxy - - def timeout(*_args: object, **_kwargs: object) -> Never: - raise httpx.ReadTimeout("response lost") - - monkeypatch.setattr(ssrf_proxy, "make_request", timeout) - with pytest.raises(TraceExportUncertainError): - export_trace("langfuse", {"host": "https://langfuse.test", "public_key": "pk", "secret_key": "sk"}, trace) - - @staticmethod - @pytest.mark.parametrize("provider", ["langsmith", "opik"]) - @pytest.mark.parametrize("accepted_prefix", [False, True]) - def test_native_sdk_preserves_first_failure_or_partial_acceptance( - trace: CanonicalTrace, monkeypatch: pytest.MonkeyPatch, provider: str, accepted_prefix: bool - ) -> None: - from core.helper import ssrf_proxy - - calls: list[str] = [] - - def send(method: str, url: str, **_kwargs: object) -> httpx.Response: - calls.append(url) - status = 200 if accepted_prefix and len(calls) == 1 else 429 - return httpx.Response(status, content=b"", request=httpx.Request(method, url)) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - config = { - "api_key": "test-key", - "project": "one", - "endpoint": "https://smith.test", - "url": "https://opik.test/api", - } - expected = TraceExportUncertainError if accepted_prefix else RetryableTraceDispatchError - with pytest.raises(expected) as caught: - export_trace(provider, config, trace) - assert isinstance(caught.value, TraceExportUncertainError) is accepted_prefix - assert len(calls) == (2 if accepted_prefix else 1) - - @staticmethod - def test_http_exports_keep_destination_credentials_local(trace: CanonicalTrace, http_calls: list[HTTPCall]) -> None: - for project in ("a", "b"): - export_trace( - "langfuse", {"host": f"https://{project}.test", "public_key": project, "secret_key": project}, trace - ) - assert http_calls[0][1].startswith("https://a.test/") - assert http_calls[1][1].startswith("https://b.test/") - assert http_calls[0][2]["headers"]["Authorization"] != http_calls[1][2]["headers"]["Authorization"] - assert http_calls[0][2]["content"] == http_calls[1][2]["content"] - - @staticmethod - @pytest.mark.parametrize( - "provider", - ["langfuse", "langsmith", "phoenix", "arize", "opik", "weave", "mlflow", "databricks", "aliyun", "tencent"], - ) - def test_capabilities_state_implemented_contract_not_platform_validation(provider: str) -> None: - result = provider_capabilities(provider) - assert result["agent_spans"] is True - assert result["export_ack"] == "accepted" - assert "platform_validated" not in result - - @staticmethod - def test_rejects_unfinished_and_unknown_provider_before_sending( - trace: CanonicalTrace, http_calls: list[HTTPCall] - ) -> None: - unfinished = trace.model_copy(update={"spans": (trace.spans[0].model_copy(update={"end_time": None}),)}) - with pytest.raises(TraceExportRejectedError): - export_trace("langfuse", {}, unfinished) - with pytest.raises(TraceExportRejectedError): - export_trace("unknown", {}, trace) - assert not http_calls diff --git a/api/tests/unit_tests/services/tracing/test_providers.py b/api/tests/unit_tests/services/tracing/test_providers.py new file mode 100644 index 00000000000000..1dad615470dad9 --- /dev/null +++ b/api/tests/unit_tests/services/tracing/test_providers.py @@ -0,0 +1,90 @@ +from collections.abc import Sequence +from typing import override + +import pytest +from dify_trace_arize_phoenix.config import PhoenixConfig +from dify_trace_arize_phoenix.provider import PhoenixProvider +from dify_trace_langfuse.config import LangfuseConfig +from dify_trace_langfuse.provider import LangfuseProvider + +from services.tracing.errors import ProviderConfigUnreadableError +from services.tracing.providers import ProviderConfigCodec, TracingSecretCipher, configured_secrets +from tests.helpers.unit.services.tracing_ports import WorkspaceBoundCipher + +WORKSPACE = "tenant-1" +PUBLIC_KEY = "pk-lf-0123456789" +SECRET_KEY = "sk-lf-abcdefghijk" + + +class _MissingWorkspaceKeyError(Exception): + pass + + +class _UnreadableCipher(TracingSecretCipher): + """A workspace whose key is gone: nothing it stored can be decrypted.""" + + @override + def encrypt(self, tenant_id: str, value: str) -> str: + raise AssertionError("decoding never encrypts") + + @override + def decrypt(self, tenant_id: str, tokens: Sequence[str]) -> list[str]: + raise _MissingWorkspaceKeyError(f"{tenant_id}: {len(tokens)} tokens") + + @override + def mask(self, value: str) -> str: + raise AssertionError("decoding never masks") + + +def test_stored_configs_encrypt_only_secrets_and_round_trip() -> None: + cipher = WorkspaceBoundCipher() + codec = ProviderConfigCodec(cipher) + config = LangfuseConfig(public_key=PUBLIC_KEY, secret_key=SECRET_KEY, host="https://langfuse.example.com") + + stored = codec.encode(WORKSPACE, config) + + assert stored["public_key"] == cipher.encrypt(WORKSPACE, PUBLIC_KEY) + assert stored["host"] == "https://langfuse.example.com" + assert codec.decode(WORKSPACE, LangfuseProvider, stored) == config + + +def test_a_config_encrypted_for_another_workspace_is_unreadable() -> None: + codec = ProviderConfigCodec(WorkspaceBoundCipher()) + stored = codec.encode(WORKSPACE, LangfuseConfig(public_key=PUBLIC_KEY, secret_key=SECRET_KEY)) + + with pytest.raises(ProviderConfigUnreadableError): + codec.decode("tenant-2", LangfuseProvider, stored) + + +def test_empty_optional_secrets_are_stored_empty_and_not_reported_as_configured() -> None: + codec = ProviderConfigCodec(WorkspaceBoundCipher()) + config = PhoenixConfig(api_key="", project="dify", endpoint="https://app.phoenix.arize.com") + + stored = codec.encode(WORKSPACE, config) + + assert stored["api_key"] == "" + assert configured_secrets(PhoenixProvider, stored) == () + assert codec.decode(WORKSPACE, PhoenixProvider, stored) == config + + +def test_secrets_that_cannot_be_decrypted_make_the_config_unreadable() -> None: + stored = ProviderConfigCodec(WorkspaceBoundCipher()).encode( + WORKSPACE, LangfuseConfig(public_key=PUBLIC_KEY, secret_key=SECRET_KEY) + ) + + with pytest.raises(ProviderConfigUnreadableError): + ProviderConfigCodec(_UnreadableCipher()).decode(WORKSPACE, LangfuseProvider, stored) + + +def test_configs_without_secrets_never_load_the_workspace_key() -> None: + config = ProviderConfigCodec(_UnreadableCipher()).decode( + WORKSPACE, PhoenixProvider, {"api_key": None, "project": "dify"} + ) + + assert isinstance(config, PhoenixConfig) + assert config.api_key is None + + +def test_stored_configs_that_no_longer_validate_are_unreadable() -> None: + with pytest.raises(ProviderConfigUnreadableError): + ProviderConfigCodec(WorkspaceBoundCipher()).decode(WORKSPACE, LangfuseProvider, {"host": "ftp://example"}) diff --git a/api/tests/unit_tests/services/tracing/test_recovery.py b/api/tests/unit_tests/services/tracing/test_recovery.py index 775662e814e13b..7b8b19661e94fc 100644 --- a/api/tests/unit_tests/services/tracing/test_recovery.py +++ b/api/tests/unit_tests/services/tracing/test_recovery.py @@ -1,74 +1,448 @@ -"""Recover lost terminal notifications using real business rows and TraceTask.""" +"""Recover lost terminal notifications from real business rows through the trace record repository.""" import json -from collections.abc import Callable -from dataclasses import dataclass -from datetime import datetime, timedelta +from collections.abc import Callable, Generator, Iterator +from dataclasses import dataclass, replace +from datetime import UTC, datetime, timedelta from decimal import Decimal -from typing import Literal +from typing import Literal, cast, override +from uuid import NAMESPACE_URL, UUID, uuid4, uuid5 import pytest -from sqlalchemy import Engine +from cachetools import TTLCache +from dify_agent.protocol import RunEvent +from dify_agent.protocol.tracing import AgentTraceBatch, AgentTraceSpan +from dify_trace_langfuse.config import LangfuseConfig +from dify_trace_sdk.emitter import TraceScope +from dify_trace_sdk.errors import EgressError +from dify_trace_sdk.events import MessageCompleted, TraceCorrelation, TraceEvent, TracePayload +from dify_trace_sdk.exporting import send_export_request +from dify_trace_sdk.ids import span_id_for +from dify_trace_sdk.model import Trace +from dify_trace_sdk.provider import ProviderCapabilities, SendResult, TraceProvider +from dify_trace_sdk.testing import RecordingEgress, response +from flask import Flask +from sqlalchemy import Engine, delete, inspect, update from sqlalchemy.orm import Session, scoped_session, sessionmaker -from core.app.entities.app_invoke_entities import ChatAppGenerateEntity -from core.ops.ops_trace_manager import TraceTask -from core.ops.unified_trace.entities import CanonicalTrace +from clients.agent_backend import FakeAgentBackendRunClient, FakeAgentBackendScenario +from core.app.app_config.entities import ( + EasyUIBasedAppConfig, + EasyUIBasedAppModelConfigFrom, + ModelConfigEntity, + PromptTemplateEntity, +) +from core.app.apps import base_app_queue_manager +from core.app.apps.agent_app.app_runner import AgentAppRunner +from core.app.apps.agent_app.session_store import AgentAppSessionScope +from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom +from core.app.apps.exc import GenerateTaskStoppedError +from core.app.apps.execution_coordinator import set_app_task_stop_flag +from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager +from core.app.entities.app_invoke_entities import ChatAppGenerateEntity, DifyRunContext, InvokeFrom, UserFrom +from core.app.entities.queue_entities import ( + AppQueueEvent, + MessageQueueMessage, + QueueLLMChunkEvent, + QueueMessageEndEvent, + QueueStopEvent, +) +from core.app.task_pipeline.easy_ui_based_generate_task_pipeline import EasyUIBasedGenerateTaskPipeline +from core.entities.provider_entities import ProviderQuotaType, QuotaConfiguration, QuotaUnit, SystemConfiguration +from domain.tracing.trace_run_agg import TraceInvocation, TraceRunState +from events.event_handlers import update_provider_when_message_created as receiver from extensions.ext_database import db +from extensions.tracing import application_tracing_hooks, bind_tracing from graphon.enums import WorkflowExecutionStatus +from graphon.model_runtime.entities.llm_entities import LLMUsage from libs.datetime_utils import naive_utc_now -from models.application_trace import ApplicationTraceInvocation +from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole +from models.application_trace import ApplicationTraceInvocation, ApplicationTraceRun from models.enums import ConversationFromSource, CreatorUserRole, MessageStatus, WorkflowRunTriggeredFrom -from models.model import App, AppMode, Message +from models.model import App, AppMode, Conversation, Message, TraceAppConfig from models.workflow import WorkflowRun, WorkflowType -from services.tracing import recovery -from services.tracing import service as service_module -from services.tracing.normalization import span_id -from services.tracing.provider_dispatch import TraceExportResult, TraceExportUncertainError +from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository +from repositories.application_trace_repository import ApplicationTraceRepository +from repositories.factory import DifyAPIRepositoryFactory +from repositories.trace_record_repository import TraceRecordRepository +from services.tracing.adapters import StoredPolicyTraceScopes +from services.tracing.canonical import CanonicalTraceBuilder +from services.tracing.collection import AgentTraceSettings, TraceRun +from services.tracing.delivery import TraceRecordSource +from services.tracing.errors import SourceNotReadyError, TraceCollectionNotFoundError +from services.tracing.providers import ProviderConfigCodec +from services.tracing.records import BaseTraceInfo, MessageTraceInfo +from services.tracing.recovery import recover_terminal +from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks +from services.tracing.scopes import TraceScopes from services.tracing.service import ApplicationTracingService -from tests.helpers.unit.model_factories import make_account, make_conversation -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_batch_for as batch_for -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_identifier as identifier -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_ingest as ingest -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_seed as seed -from tests.unit_tests.services.tracing.test_service import _ApplicationTraceCollectionCollection as Collection +from tests.helpers.unit.core.model_fixtures import make_model_config +from tests.helpers.unit.model_factories import make_account, make_app, make_conversation, make_message, make_tenant +from tests.helpers.unit.services.tracing_ports import ( + AGENT_TRACE_SETTINGS, + InstalledTraceProviders, + WorkspaceBoundCipher, +) +from tests.unit_tests.core.app.apps.agent_app.test_app_runner import ( + _CancelAndWaitFailingClient, + _ExplodingSessionStore, + _FakeQueueManager, + _FakeSessionStore, + _runner, + _soul, + _StoppedQueueManager, + _UsageCancellationClient, +) -@dataclass -class RecoveryCase: - collection: Collection - service: ApplicationTracingService - now: datetime - scheduled: list[str] +def identifier(value: str) -> str: + return str(uuid5(NAMESPACE_URL, value)) + + +@dataclass(frozen=True) +class Collection: + repository: ApplicationTraceRepository + sessions: sessionmaker[Session] + run: TraceRun + + def invocation(self, name: str) -> TraceInvocation: + with self.repository.change(self.run.id, tenant_id=self.run.tenant_id) as locked: + assert locked is not None + return locked.run.open_invocation( + agent_id=identifier("agent"), node_execution_id=None, attempt=0, call_key=name + ) + + +class ScriptedProvider(TraceProvider[LangfuseConfig]): + """Sends each trace as one JSON request through the egress, whose script decides the outcome.""" + + id = "langfuse" + config_model = LangfuseConfig + capabilities = ProviderCapabilities(native_cost=True, trace_url=True, span_url=True) + + @override + def verify(self, config: LangfuseConfig) -> LangfuseConfig: + return config + + @override + def project_url(self, config: LangfuseConfig) -> str | None: + return f"{config.host}/project/{config.project_key}" + + @override + def send(self, config: LangfuseConfig, trace: Trace) -> SendResult: + reply = send_export_request( + self.egress, + "POST", + f"{config.host}/traces", + destination="Scripted", + headers={"Authorization": f"{config.public_key}:{config.secret_key}"}, + content=trace.model_dump_json().encode(), + sensitive_url=False, + ) + body = reply.json() + assert isinstance(body, dict) + return SendResult( + trace_id=str(body["trace_id"]), trace_url=body.get("trace_url"), span_urls=body.get("span_urls", {}) + ) + + +def record_source(sessions: sessionmaker[Session]) -> TraceRecordRepository: + return TraceRecordRepository( + sessions, + workflow_runs=DifyAPIRepositoryFactory.create_api_workflow_run_repository(sessions), + node_executions=DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(sessions), + files_url="https://files.test", + enterprise_details=False, + ) + + +def tracing_service( + repository: ApplicationTraceRepository, + records: TraceRecordSource, + destination: RecordingEgress, + scheduled: list[str], + settings: AgentTraceSettings, +) -> ApplicationTracingService: + """Agent trace collection over ``repository``; scheduled exports are recorded instead of enqueued.""" + return ApplicationTracingService( + runs=repository, + run_query=repository, + sources=repository, + records=records, + builder=CanonicalTraceBuilder(), + providers=InstalledTraceProviders(ScriptedProvider), + codec=ProviderConfigCodec(WorkspaceBoundCipher()), + egress=destination, + schedule_export=scheduled.append, + settings=settings, + ) + + +def seed(sessions: sessionmaker[Session], name: str, *, capture_content: bool, capture_agent: bool) -> Collection: + """An Agent app tracing to Langfuse, one of its messages, and the run frozen when the message started.""" + app_id, tenant_id, message_id = identifier(f"app-{name}"), identifier("tenant"), identifier(f"message-{name}") + cipher = WorkspaceBoundCipher() + with sessions.begin() as session: + session.add( + make_app( + app_id=app_id, + tenant_id=tenant_id, + mode=AppMode.AGENT, + tracing=json.dumps( + { + "enabled": True, + "tracing_provider": "langfuse", + "capture_content": capture_content, + "capture_agent": capture_agent, + } + ), + ) + ) + session.add( + make_message( + message_id=message_id, + app_id=app_id, + conversation_id=identifier(f"conversation-{name}"), + inputs={}, + query="hello", + message={}, + answer="done", + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.CONSOLE, + ) + ) + session.add( + TraceAppConfig( + app_id=app_id, + tracing_provider="langfuse", + tracing_config={ + "public_key": cipher.encrypt(tenant_id, "public"), + "secret_key": cipher.encrypt(tenant_id, "original"), + "host": "https://langfuse.test", + "project_key": "original", + }, + ) + ) + repository = ApplicationTraceRepository(sessions) + run = tracing_service( + repository, record_source(sessions), RecordingEgress([], grpc_target=None), [], AGENT_TRACE_SETTINGS + ).start( + tenant_id=tenant_id, + app_id=app_id, + source_kind="message", + source_id=message_id, + session_id="session-original", + workflow_id=None, + ) + assert run is not None + return Collection(repository, sessions, run) -def attach_conversation(case: Collection) -> None: - with case.sessions.begin() as session: - message = session.get(Message, case.run.source_id) +def attach_conversation(collection: Collection, mode: AppMode) -> None: + """Message records are assembled only for a message whose conversation belongs to its app.""" + with collection.sessions.begin() as session: + message = session.get(Message, collection.run.source_id) assert message is not None session.add( make_conversation( conversation_id=message.conversation_id, - app_id=case.run.app_id, - mode=AppMode.AGENT, + app_id=collection.run.app_id, + mode=mode, inputs={}, from_source=ConversationFromSource.CONSOLE, ) ) +def add_app_owner(collection: Collection, name: str) -> str: + """An account of the app's workspace that creates the app and, here, its workflow runs.""" + owner_id = identifier(f"owner-{name}") + with collection.sessions.begin() as session: + if session.get(Tenant, collection.run.tenant_id) is None: + session.add(make_tenant(tenant_id=collection.run.tenant_id)) + session.add(make_account(account_id=owner_id, email=f"{name}@example.com")) + session.add( + TenantAccountJoin(tenant_id=collection.run.tenant_id, account_id=owner_id, role=TenantAccountRole.OWNER) + ) + app = session.get(App, collection.run.app_id) + assert app is not None + app.created_by = owner_id + return owner_id + + +def trace_event(run: TraceRun, payload: TracePayload) -> TraceEvent: + return TraceEvent( + event_id=uuid4().hex, + tenant_id=run.tenant_id, + app_id=run.app_id, + user_id=None, + occurred_at=datetime.now(UTC), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + payload=payload, + ) + + +def batch_for(run: TraceRun, invocation: TraceInvocation, *, final: bool, span_id: str) -> AgentTraceBatch: + return AgentTraceBatch( + run_id=UUID(run.id), + invocation_id=UUID(invocation.id), + batch_id=uuid4(), + final=final, + spans=[ + AgentTraceSpan( + trace_id=run.trace_id, + span_id=span_id, + parent_span_id=invocation.span_id, + name="Agent execution", + start_time_ns=1_000_000_000, + end_time_ns=2_000_000_000, + status="ok", + attributes={"gen_ai.operation.name": "invoke_agent"}, + ) + ], + ) + + +def ingest(collection: Collection, batch: AgentTraceBatch) -> None: + """Store a runtime batch in its run without scheduling the run's export.""" + with collection.repository.change(str(batch.run_id), tenant_id=collection.run.tenant_id) as locked: + assert locked is not None + locked.run.ingest(batch, max_spans=100, max_bytes=1_000_000, now=naive_utc_now()) + + +def claim(collection: Collection, run_id: str) -> str | None: + """Claim the run for an export attempt; ``None`` when it may not export now.""" + token = str(uuid4()) + with collection.repository.change(run_id, tenant_id=collection.run.tenant_id) as locked: + assert locked is not None + claimed = locked.run.claim_export( + token=token, owner_exists=locked.owner_exists, collection_timeout=60, lease_seconds=30, now=naive_utc_now() + ) + return token if claimed else None + + +def recover( + collection: Collection, + run: TraceRun, + *, + records: TraceRecordSource, + collector: ApplicationTracingService, + settings: AgentTraceSettings, + now: datetime, +) -> bool: + """One reconciler recovery attempt of ``run`` at ``now``.""" + current = collection.repository.load(run.id) + assert current is not None + return recover_terminal( + current, + runs=collection.repository, + sources=collection.repository, + records=records, + collector=collector, + collection_timeout=settings.collection_timeout_seconds, + now=now, + ) + + +def sent_traces(egress: RecordingEgress) -> list[Trace]: + return [Trace.model_validate_json(request.content or b"") for request in egress.requests] + + +@dataclass(frozen=True) +class RecoveryCase: + collection: Collection + records: TraceRecordRepository + service: ApplicationTracingService + settings: AgentTraceSettings + now: datetime + scheduled: list[str] + destination: RecordingEgress + + def recover(self, run: TraceRun, *, now: datetime) -> bool: + """One reconciler recovery attempt of ``run`` at ``now``.""" + return recover( + self.collection, run, records=self.records, collector=self.service, settings=self.settings, now=now + ) + + +def recovery_case( + collection: Collection, destination: RecordingEgress, settings: AgentTraceSettings, now: datetime +) -> RecoveryCase: + """The service the runtime hooks reach; scheduled exports are recorded instead of enqueued.""" + records = record_source(collection.sessions) + scheduled: list[str] = [] + service = tracing_service(collection.repository, records, destination, scheduled, settings) + return RecoveryCase(collection, records, service, settings, now, scheduled, destination) + + +class _Recorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +class _QueueManager(AppQueueManager): + """Finite in-memory event source for the consumer's stream dispatch.""" + + def __init__(self, messages: list[MessageQueueMessage]) -> None: + self.messages = messages + + @override + def listen(self) -> Generator[MessageQueueMessage]: + yield from self.messages + + @override + def _publish(self, event: AppQueueEvent, pub_from: PublishFrom) -> None: + raise AssertionError("Stream dispatch must not publish queue events") + + +class _FailingRecords(TraceRecordSource): + """Remembers every event recovery asks it to assemble, then fails with ``failure``.""" + + def __init__(self, failure: Exception) -> None: + self._failure = failure + self.events: list[TraceEvent] = [] + + @override + def assemble(self, event: TraceEvent) -> BaseTraceInfo: + self.events.append(event) + raise self._failure + + +def tracing_application(case: RecoveryCase) -> Flask: + """An application bound like the composition root of a collecting deployment, with events kept in memory.""" + application = Flask(__name__) + bind_tracing( + application, + scopes=StoredPolicyTraceScopes( + TraceScopes( + emitter=_Recorder(), + policies=SQLAlchemyAppTracingConfigRepository(session_factory=case.collection.sessions), + enterprise_enabled=lambda: False, + ) + ), + hooks=RuntimeApplicationTracingHooks(case.service), + ) + return application + + @pytest.fixture def case( sqlite_engine: Engine, sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch ) -> RecoveryCase: - collection = seed(sqlite_session_factory, "recover", capture_content=True) - attach_conversation(collection) - invocation = collection.invocation() - ingest(collection, batch_for(collection.run, invocation)) + collection = seed(sqlite_session_factory, "recover", capture_content=True, capture_agent=True) + attach_conversation(collection, AppMode.AGENT) + invocation = collection.invocation("first") + ingest(collection, batch_for(collection.run, invocation, final=True, span_id="a" * 16)) now = naive_utc_now() with sqlite_session_factory.begin() as session: message = session.get(Message, collection.run.source_id) - assert message is not None stored_invocation = session.get(ApplicationTraceInvocation, invocation.id) assert message is not None assert stored_invocation is not None @@ -81,13 +455,43 @@ def case( stored_invocation.finalized_at = now - timedelta(seconds=10) monkeypatch.setattr(type(db), "engine", property(lambda _db: sqlite_engine)) monkeypatch.setattr(db, "session", scoped_session(sqlite_session_factory)) - monkeypatch.setattr(TraceTask, "_workflow_run_repo", None) - monkeypatch.setattr("core.telemetry.gateway.is_enterprise_telemetry_enabled", lambda: False) - monkeypatch.setattr(recovery, "naive_utc_now", lambda: now) - service = ApplicationTracingService(collection.repository) - scheduled: list[str] = [] - monkeypatch.setattr(service, "schedule", scheduled.append) - return RecoveryCase(collection, service, now, scheduled) + destination = RecordingEgress([response(200, {"trace_id": "provider-id"})], grpc_target=None) + return recovery_case(collection, destination, AGENT_TRACE_SETTINGS, now) + + +def collecting_within(case: RecoveryCase, collection_timeout_seconds: int) -> RecoveryCase: + """The same case on a deployment that waits ``collection_timeout_seconds`` for a business terminal.""" + return recovery_case( + case.collection, + case.destination, + replace(case.settings, collection_timeout_seconds=collection_timeout_seconds), + case.now, + ) + + +def stop_requested(collection: Collection) -> bool: + """The consumer's persisted stop intent, as both ledger readers of the message report it.""" + run = collection.run + message = collection.repository.message(tenant_id=run.tenant_id, app_id=run.app_id, message_id=run.source_id) + progress = collection.repository.message_progress( + tenant_id=run.tenant_id, app_id=run.app_id, message_id=run.source_id + ) + assert message is not None + assert progress is not None + assert message.agent_stop_requested is progress.agent_stop_requested + return progress.agent_stop_requested + + +def message_record(case: RecoveryCase) -> MessageTraceInfo: + """The record the delivery worker assembles when the message reports its end.""" + run = case.collection.run + message = case.collection.repository.message(tenant_id=run.tenant_id, app_id=run.app_id, message_id=run.source_id) + assert message is not None + record = case.records.assemble( + trace_event(run, MessageCompleted(message_id=run.source_id, conversation_id=message.conversation_id)) + ) + assert isinstance(record, MessageTraceInfo) + return record def complete_metadata() -> str: @@ -110,35 +514,8 @@ def save_agent_stop( consumer_session: Session | None = None, process_stream: bool = False, notify: Callable[[Message, ChatAppGenerateEntity], None] | None = None, -) -> list[TraceTask]: - """Run the real consumer save and tracing queue against the owned SQL rows.""" - import queue - - from flask import Flask - - from core.app.app_config.entities import ( - EasyUIBasedAppConfig, - EasyUIBasedAppModelConfigFrom, - ModelConfigEntity, - PromptTemplateEntity, - ) - from core.app.entities.app_invoke_entities import InvokeFrom - from core.app.entities.queue_entities import QueueStopEvent - from core.app.task_pipeline.easy_ui_based_generate_task_pipeline import EasyUIBasedGenerateTaskPipeline - from core.ops import ops_trace_manager - from core.ops.base_trace_instance import bind_application_tracing_hooks - from services.tracing import runtime_bridge - from tests.helpers.unit.config_override import apply_config_overrides - from tests.helpers.unit.core.model_fixtures import make_model_config - from tests.unit_tests.core.app.task_pipeline.test_easy_ui_based_generate_task_pipeline import ( - _QueueManager, - queue_message, - ) - - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - monkeypatch.setattr(runtime_bridge, "_service", lambda: case.service) - monkeypatch.setattr(ops_trace_manager, "trace_manager_queue", queue.Queue()) - monkeypatch.setattr(ops_trace_manager.TraceQueueManager, "start_timer", lambda _manager: None) +) -> list[TraceEvent]: + """Run the real consumer save against the owned SQL rows; returns the trace events it emitted.""" def send(message: Message, *, application_generate_entity: ChatAppGenerateEntity) -> None: if notify is not None: @@ -157,8 +534,9 @@ def send(message: Message, *, application_generate_entity: ChatAppGenerateEntity with case.collection.sessions() as session: message = session.get(Message, case.collection.run.source_id) assert message is not None - conversation = session.get(service_module.Conversation, message.conversation_id) + conversation = session.get(Conversation, message.conversation_id) assert conversation is not None + recorder = _Recorder() entity = ChatAppGenerateEntity( task_id=identifier("stop-race-task"), app_config=EasyUIBasedAppConfig( @@ -176,39 +554,46 @@ def send(message: Message, *, application_generate_entity: ChatAppGenerateEntity user_id=identifier("stop-race-user"), stream=True, invoke_from=InvokeFrom.DEBUGGER, + trace_scope=TraceScope( + emitter=recorder, + tenant_id=case.collection.run.tenant_id, + app_id=case.collection.run.app_id, + user_id=identifier("stop-race-user"), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + ), + ) + stop = MessageQueueMessage( + task_id=entity.task_id, + app_mode=mode, + message_id=message.id, + conversation_id=conversation.id, + event=QueueStopEvent(stopped_by=QueueStopEvent.StopBy.USER_MANUAL), + ) + pipeline = EasyUIBasedGenerateTaskPipeline( + entity, _QueueManager([stop] if process_stream else []), conversation, message, stream=True ) - queue_manager = _QueueManager() - pipeline = EasyUIBasedGenerateTaskPipeline(entity, queue_manager, conversation, message, stream=True) pipeline._task_state.llm_result.message.content = answer - application = Flask(__name__) - bind_application_tracing_hooks(application, runtime_bridge.RuntimeApplicationTracingHooks()) - with application.app_context(): - trace_manager = ops_trace_manager.TraceQueueManager(app_id=case.collection.run.app_id) - if process_stream: - queue_manager.messages = [queue_message(QueueStopEvent(stopped_by=QueueStopEvent.StopBy.USER_MANUAL))] - # Token estimates and moderation are separate external boundaries; - # exercise real stream dispatch, SQL save, commit and notification. - monkeypatch.setattr(pipeline, "_handle_stop", lambda _event: None) - monkeypatch.setattr(pipeline, "handle_output_moderation_when_task_finished", lambda _answer: None) - list(pipeline._process_stream_response(publisher=None, trace_manager=trace_manager)) - elif consumer_session is None: - with case.collection.sessions.begin() as session: - pipeline._save_message(session=session, trace_manager=trace_manager, preserve_existing_usage=True) - else: - pipeline._save_message(session=consumer_session, trace_manager=trace_manager, preserve_existing_usage=True) - consumer_session.commit() - return trace_manager.collect_tasks() + if process_stream: + # Token estimates and moderation are separate external boundaries; + # exercise real stream dispatch, SQL save, commit and notification. + monkeypatch.setattr(pipeline, "_handle_stop", lambda _event: None) + monkeypatch.setattr(pipeline, "handle_output_moderation_when_task_finished", lambda _answer: None) + list(pipeline._process_stream_response(publisher=None)) + elif consumer_session is None: + with case.collection.sessions.begin() as session: + pipeline._save_message(session=session, preserve_existing_usage=True) + else: + pipeline._save_message(session=consumer_session, preserve_existing_usage=True) + consumer_session.commit() + return recorder.events @pytest.mark.parametrize("receiver_fails", [False, True]) def test_stopped_agent_notifies_the_real_provider_receiver_once_after_commit( case: RecoveryCase, monkeypatch: pytest.MonkeyPatch, receiver_fails: bool ) -> None: - from sqlalchemy import inspect - - from core.entities.provider_entities import ProviderQuotaType, QuotaConfiguration, QuotaUnit, SystemConfiguration - from events.event_handlers import update_provider_when_message_created as receiver - calls: list[str] = [] provider_updates: list[str] = [] monkeypatch.setattr( @@ -226,8 +611,8 @@ def notify(message: Message, entity: ChatAppGenerateEntity) -> None: with case.collection.sessions() as session: saved = session.get(Message, message.id) assert saved is not None - assert recovery.agent_stop_requested(saved) is True assert saved.answer == message.answer + assert stop_requested(case.collection) is True system = SystemConfiguration( enabled=True, current_quota_type=ProviderQuotaType.TRIAL, @@ -251,101 +636,84 @@ def notify(message: Message, entity: ChatAppGenerateEntity) -> None: with pytest.raises(RuntimeError, match="receiver failed after commit"): save_agent_stop(case, monkeypatch, "partial answer", process_stream=True, notify=notify) else: + # The Agent worker reports a stopped Agent run's terminal itself. assert save_agent_stop(case, monkeypatch, "partial answer", process_stream=True, notify=notify) == [] assert calls == [case.collection.run.source_id] assert provider_updates == ["basic_last_used_update"] - with case.collection.sessions() as session: - saved = session.get(Message, case.collection.run.source_id) - assert saved is not None - assert recovery.agent_stop_requested(saved) is True + assert stop_requested(case.collection) is True @pytest.mark.parametrize("answer", ["", "partial answer before stop"]) @pytest.mark.parametrize("outcome", ["succeeded", "failed", "cancelled"]) -@pytest.mark.parametrize("delivery", ["accepted", "uncertain"]) +@pytest.mark.parametrize("delivery", [TraceRunState.ACCEPTED, TraceRunState.UNCERTAIN]) def test_stopped_consumer_cannot_freeze_a_terminal_before_the_runtime_worker( case: RecoveryCase, monkeypatch: pytest.MonkeyPatch, answer: str, outcome: Literal["succeeded", "failed", "cancelled"], - delivery: str, + delivery: TraceRunState, ) -> None: - from typing import cast - - from core.ops.entities.trace_entity import MessageTraceInfo, TraceTaskName - from models.model import MessageDict - + destination = RecordingEgress( + [ + response(200, {"trace_id": "provider-id"}) + if delivery is TraceRunState.ACCEPTED + else EgressError("acknowledgement lost", request_sent=True) + ], + grpc_target=None, + ) + stopped = recovery_case(case.collection, destination, case.settings, case.now) + service = stopped.service + tenant_id = case.collection.run.tenant_id # Simulate an old queued notification as well as the current consumer save: # persisted state, not a detached event payload, must decide who owns the end. - notification = TraceTask( - TraceTaskName.MESSAGE_TRACE, - message_id=case.collection.run.source_id, - conversation_id=identifier("conversation-recover"), - ) - info = notification.execute() - assert isinstance(info, MessageTraceInfo) - assert isinstance(info.message_data, dict) - info.message_data = Message.from_dict(cast(MessageDict, info.message_data)) - exported: list[CanonicalTrace] = [] - - def export(_provider: str, _config: dict[str, object], trace: CanonicalTrace) -> TraceExportResult: - exported.append(trace) - if delivery == "uncertain": - raise TraceExportUncertainError("acknowledgment lost") - return TraceExportResult(trace_id=trace.trace_id) - - monkeypatch.setattr(service_module, "export_trace", export) - monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", lambda *_args: {}) + notification = message_record(stopped) queued = save_agent_stop(case, monkeypatch, answer) - assert case.service.capture_terminal(info) is True + assert service.capture_terminal(tenant_id=tenant_id, info=notification) is True stored = case.collection.repository.get(case.collection.run.id) assert stored is not None assert stored.outer_trace is None, "a stopped SSE save is not the runtime terminal" assert queued == [] - assert recovery.recover_terminal(case.service, stored) is False - assert case.collection.repository.claim(stored.id, collection_timeout=60, lease_seconds=30) is None - case.service.deliver(stored.id) - assert exported == [] + assert stopped.recover(stored, now=case.now) is False + assert claim(case.collection, stored.id) is None + service.deliver(stored.id) + assert destination.requests == [] output = "real runtime answer" if outcome == "succeeded" else None - case.service.capture_agent_app_terminal( + service.capture_agent_app_terminal( tenant_id=stored.tenant_id, app_id=stored.app_id, message_id=stored.source_id, outcome=outcome, output=output, ) - terminal = case.collection.repository.get(stored.id) - assert terminal is not None - trace = case.service.assemble(terminal) - root = next(span for span in trace.spans if span.id == trace.root_span_id) - assert root.metadata["message_execution_state"] == outcome - assert root.status.value == ("ok" if outcome == "succeeded" else "error") - assert len(trace.spans) == 3 # A duplicate worker result and late consumer notification cannot rewrite # the frozen end or create a second ingestion attempt after acceptance. - case.service.deliver(stored.id) - case.service.deliver(stored.id) + service.deliver(stored.id) + service.deliver(stored.id) + exported = sent_traces(destination) assert len(exported) == 1 - assert exported[0] == trace + trace = exported[0] + assert trace.root.attributes["message_execution_state"] == outcome + assert trace.root.status.value == ("ok" if outcome == "succeeded" else "error") + assert len(trace.spans) == 3 frozen = case.collection.repository.get(stored.id) assert frozen is not None - assert frozen.state == delivery - case.service.capture_agent_app_terminal( + assert frozen.state is delivery + service.capture_agent_app_terminal( tenant_id=stored.tenant_id, app_id=stored.app_id, message_id=stored.source_id, outcome="succeeded" if outcome != "succeeded" else "cancelled", output="late result must not replace the accepted envelope", ) - assert case.service.capture_terminal(info) is True + assert service.capture_terminal(tenant_id=tenant_id, info=notification) is True unchanged = case.collection.repository.get(stored.id) assert unchanged is not None assert unchanged.outer_trace == frozen.outer_trace assert unchanged.attempts == frozen.attempts == 1 - case.service.deliver(stored.id) - assert len(exported) == 1 + service.deliver(stored.id) + assert len(destination.requests) == 1 with case.collection.sessions() as session: message = session.get(Message, stored.source_id) assert message is not None @@ -359,9 +727,10 @@ def export(_provider: str, _config: dict[str, object], trace: CanonicalTrace) -> def test_classic_stopped_consumer_keeps_its_existing_trace_notification( case: RecoveryCase, monkeypatch: pytest.MonkeyPatch, mode: AppMode ) -> None: - queued = save_agent_stop(case, monkeypatch, "classic partial answer", mode=mode) - assert len(queued) == 1 - assert queued[0].trace_type is service_module.TraceTaskName.MESSAGE_TRACE + queued = save_agent_stop(case, monkeypatch, "classic partial answer", mode=mode, process_stream=True) + assert [event.payload for event in queued] == [ + MessageCompleted(message_id=case.collection.run.source_id, conversation_id=identifier("conversation-recover")) + ] with case.collection.sessions() as session: message = session.get(Message, case.collection.run.source_id) assert message is not None @@ -371,9 +740,6 @@ def test_classic_stopped_consumer_keeps_its_existing_trace_notification( def test_worker_usage_flush_preserves_stop_intent_from_a_newer_consumer_save( case: RecoveryCase, monkeypatch: pytest.MonkeyPatch ) -> None: - from core.app.apps.agent_app.app_runner import AgentAppRunner - from graphon.model_runtime.entities.llm_entities import LLMUsage - # The worker's scoped session may already have loaded this Message before # the consumer commits its independent stop transaction. cached_message = db.session.get(Message, case.collection.run.source_id) @@ -384,10 +750,10 @@ def test_worker_usage_flush_preserves_stop_intent_from_a_newer_consumer_save( message_id=cached_message.id, usage=LLMUsage.from_metadata({"prompt_tokens": 13, "completion_tokens": 8, "total_price": "0.21"}), ) + assert stop_requested(case.collection) is True with case.collection.sessions() as session: message = session.get(Message, cached_message.id) assert message is not None - assert recovery.agent_stop_requested(message) is True assert message.answer == "partial answer" assert message.message_tokens == 13 assert message.answer_tokens == 8 @@ -397,9 +763,6 @@ def test_worker_usage_flush_preserves_stop_intent_from_a_newer_consumer_save( def test_stop_save_preserves_worker_usage_committed_after_its_cached_message_read( case: RecoveryCase, monkeypatch: pytest.MonkeyPatch ) -> None: - from core.app.apps.agent_app.app_runner import AgentAppRunner - from graphon.model_runtime.entities.llm_entities import LLMUsage - with case.collection.sessions.begin() as session: initial = session.get(Message, case.collection.run.source_id) assert initial is not None @@ -414,10 +777,10 @@ def test_stop_save_preserves_worker_usage_committed_after_its_cached_message_rea usage=LLMUsage.from_metadata({"prompt_tokens": 13, "completion_tokens": 8, "total_price": "0.21"}), ) save_agent_stop(case, monkeypatch, "partial answer", consumer_session=consumer) + assert stop_requested(case.collection) is True with case.collection.sessions() as session: message = session.get(Message, case.collection.run.source_id) assert message is not None - assert recovery.agent_stop_requested(message) is True assert message.message_tokens == 13 assert message.answer_tokens == 8 assert message.total_price == Decimal("0.21") @@ -434,28 +797,8 @@ def test_stopped_runner_hands_off_only_an_observed_runtime_terminal( final_batch_arrived: bool, snapshot_save_fails: bool, ) -> None: - from typing import cast - - from flask import Flask - - from clients.agent_backend import FakeAgentBackendRunClient, FakeAgentBackendScenario - from core.app.apps.agent_app.session_store import AgentAppSessionScope - from core.app.apps.base_app_queue_manager import AppQueueManager - from core.app.apps.exc import GenerateTaskStoppedError - from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom - from core.ops.base_trace_instance import bind_application_tracing_hooks - from services.tracing import runtime_bridge - from tests.helpers.unit.config_override import apply_config_overrides - from tests.unit_tests.core.app.apps.agent_app.test_app_runner import ( - _CancelAndWaitFailingClient, - _ExplodingSessionStore, - _FakeSessionStore, - _runner, - _StoppedQueueManager, - ) - save_agent_stop(case, monkeypatch, "partial answer") - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_COLLECTION_TIMEOUT=60) + case = collecting_within(case, 60) pending_invocation = None if not final_batch_arrived: pending_invocation = case.collection.invocation("awaiting-stop-final-batch") @@ -498,9 +841,7 @@ def unexpected_cancel(*_args: object, **_kwargs: object) -> None: agent_config_snapshot_id=identifier("stop-snapshot"), home_snapshot_id=None, ) - application = Flask(__name__) - bind_application_tracing_hooks(application, runtime_bridge.RuntimeApplicationTracingHooks()) - with application.app_context(), pytest.raises(GenerateTaskStoppedError): + with tracing_application(case).app_context(), pytest.raises(GenerateTaskStoppedError): runner._consume_stream( client.run_id, dify_context=context, @@ -516,96 +857,65 @@ def unexpected_cancel(*_args: object, **_kwargs: object) -> None: if runtime_outcome == "cancel_wait_failed": assert frozen.outer_trace is None assert store.saved == [] - monkeypatch.setattr(recovery, "naive_utc_now", lambda: case.now + timedelta(seconds=120)) - assert recovery.recover_terminal(case.service, frozen) is False + assert case.recover(frozen, now=case.now + timedelta(seconds=120)) is False expired = case.collection.repository.get(run.id) assert expired is not None - assert expired.state == "failed" + assert expired.state is TraceRunState.FAILED assert expired.error_code == ( "business_terminal_not_persisted" if final_batch_arrived else "agent_trace_collection_incomplete" ) return assert frozen.outer_trace is not None - trace = case.service.assemble(frozen) - root = next(span for span in trace.spans if span.id == trace.root_span_id) - assert root.metadata["message_execution_state"] == runtime_outcome + root = Trace.model_validate(frozen.outer_trace).root + assert root.attributes["message_execution_state"] == runtime_outcome assert root.status.value == ("ok" if runtime_outcome == "succeeded" else "error") assert len(store.saved) == (0 if snapshot_save_fails else 1) - exported: list[CanonicalTrace] = [] - - def export(_provider: str, _config: dict[str, object], payload: CanonicalTrace) -> TraceExportResult: - exported.append(payload) - return TraceExportResult(trace_id=payload.trace_id) - - monkeypatch.setattr(service_module, "export_trace", export) - monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", lambda *_args: {}) if pending_invocation is not None: - assert case.collection.repository.claim(run.id, collection_timeout=60, lease_seconds=30) is None + assert claim(case.collection, run.id) is None case.service.deliver(run.id) - assert exported == [] + assert case.destination.requests == [] ingest(case.collection, batch_for(run, pending_invocation, final=True, span_id="b" * 16)) case.service.deliver(run.id) case.service.deliver(run.id) + exported = sent_traces(case.destination) assert len(exported) == 1 + assert exported[0].root.attributes["message_execution_state"] == runtime_outcome assert len(exported[0].spans) == (3 if final_batch_arrived else 5) - accepted = case.collection.repository.get(run.id) - assert accepted is not None - assert accepted.state == "accepted" - assert accepted.attempts == 1 + accepted_run = case.collection.repository.get(run.id) + assert accepted_run is not None + assert accepted_run.state is TraceRunState.ACCEPTED + assert accepted_run.attempts == 1 + assert stop_requested(case.collection) is True with case.collection.sessions() as session: message = session.get(Message, run.source_id) assert message is not None assert message.answer == "partial answer" - assert recovery.agent_stop_requested(message) is True @pytest.mark.parametrize("incomplete", [False, True]) def test_missing_runtime_stop_terminal_converges_without_inventing_success( case: RecoveryCase, monkeypatch: pytest.MonkeyPatch, incomplete: bool ) -> None: - from tests.helpers.unit.config_override import apply_config_overrides - - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_COLLECTION_TIMEOUT=60) + case = collecting_within(case, 60) save_agent_stop(case, monkeypatch, "partial answer") if incomplete: invocation = case.collection.invocation("still-collecting") ingest(case.collection, batch_for(case.collection.run, invocation, final=False, span_id="b" * 16)) - monkeypatch.setattr(recovery, "naive_utc_now", lambda: case.now + timedelta(seconds=120)) - assert recovery.recover_terminal(case.service, case.collection.run) is False + assert case.recover(case.collection.run, now=case.now + timedelta(seconds=120)) is False stored = case.collection.repository.get(case.collection.run.id) assert stored is not None - assert stored.state == "failed" + assert stored.state is TraceRunState.FAILED assert stored.error_code == ( "agent_trace_collection_incomplete" if incomplete else "business_terminal_not_persisted" ) assert stored.outer_trace is None - assert case.collection.repository.claim(stored.id, collection_timeout=60, lease_seconds=30) is None + assert claim(case.collection, stored.id) is None @pytest.mark.parametrize("stop_on_end", [False, True]) def test_stop_during_final_queue_publication_preserves_the_already_observed_success( case: RecoveryCase, monkeypatch: pytest.MonkeyPatch, stop_on_end: bool ) -> None: - from cachetools import TTLCache - from flask import Flask - - from clients.agent_backend import FakeAgentBackendRunClient - from core.app.apps import base_app_queue_manager - from core.app.apps.base_app_queue_manager import PublishFrom - from core.app.apps.exc import GenerateTaskStoppedError - from core.app.apps.execution_coordinator import set_app_task_stop_flag - from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager - from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom - from core.app.entities.queue_entities import ( - AppQueueEvent, - MessageQueueMessage, - QueueLLMChunkEvent, - QueueMessageEndEvent, - ) - from core.ops.base_trace_instance import bind_application_tracing_hooks - from services.tracing import runtime_bridge - from tests.unit_tests.core.app.apps.agent_app.test_app_runner import _FakeSessionStore, _runner, _soul - save_agent_stop(case, monkeypatch, "partial answer") # Only Redis storage is in memory. Use the real queue's publish stop check # rather than an already-stopped fake consumer or an exception-only stub. @@ -647,9 +957,7 @@ def publish_after_stop(event: AppQueueEvent, pub_from: PublishFrom) -> None: user_from=UserFrom.ACCOUNT, invoke_from=InvokeFrom.DEBUGGER, ) - application = Flask(__name__) - bind_application_tracing_hooks(application, runtime_bridge.RuntimeApplicationTracingHooks()) - with application.app_context(), pytest.raises(GenerateTaskStoppedError): + with tracing_application(case).app_context(), pytest.raises(GenerateTaskStoppedError): runner.run( dify_context=context, agent_id=identifier("stop-agent"), @@ -672,9 +980,8 @@ def publish_after_stop(event: AppQueueEvent, pub_from: PublishFrom) -> None: frozen = case.collection.repository.get(run.id) assert frozen is not None assert frozen.outer_trace is not None, "success arrived before the final queue publication saw Stop" - trace = CanonicalTrace.model_validate(frozen.outer_trace) - root = next(span for span in trace.spans if span.id == trace.root_span_id) - assert root.metadata["message_execution_state"] == "succeeded" + root = Trace.model_validate(frozen.outer_trace).root + assert root.attributes["message_execution_state"] == "succeeded" assert root.status.value == "ok" assert len(store.saved) == 1 @@ -683,28 +990,6 @@ def publish_after_stop(event: AppQueueEvent, pub_from: PublishFrom) -> None: def test_runtime_cleanup_cancellation_does_not_reclassify_a_business_stream_failure( case: RecoveryCase, monkeypatch: pytest.MonkeyPatch, stop_during_cleanup: bool ) -> None: - from collections.abc import Iterator - from typing import cast - - from dify_agent.protocol import RunEvent - from flask import Flask - - from core.app.apps.agent_app.session_store import AgentAppSessionScope - from core.app.apps.base_app_queue_manager import AppQueueManager - from core.app.apps.exc import GenerateTaskStoppedError - from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom - from core.ops.base_trace_instance import bind_application_tracing_hooks - from services.tracing import runtime_bridge - from tests.helpers.unit.config_override import apply_config_overrides - from tests.unit_tests.core.app.apps.agent_app.test_app_runner import ( - _FakeQueueManager, - _FakeSessionStore, - _runner, - _UsageCancellationClient, - ) - - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - monkeypatch.setattr(runtime_bridge, "_service", lambda: case.service) client = _UsageCancellationClient() queue_manager = _FakeQueueManager() @@ -739,9 +1024,7 @@ def cancel_then_stop(*_args: object, **_kwargs: object) -> RunEvent: agent_config_snapshot_id=identifier("cleanup-snapshot"), home_snapshot_id=None, ) - application = Flask(__name__) - bind_application_tracing_hooks(application, runtime_bridge.RuntimeApplicationTracingHooks()) - with application.app_context(): + with tracing_application(case).app_context(): with pytest.raises(GenerateTaskStoppedError if stop_during_cleanup else RuntimeError): runner._consume_stream( client.run_id, @@ -756,7 +1039,9 @@ def cancel_then_stop(*_args: object, **_kwargs: object) -> RunEvent: if not stop_during_cleanup: # This is the generator's existing generic exception handoff. A # cleanup child-run cancellation must not win the application root. - runtime_bridge.capture_agent_app_terminal( + hooks = application_tracing_hooks() + assert hooks is not None + hooks.capture_agent_app_terminal( tenant_id=run.tenant_id, app_id=run.app_id, message_id=run.source_id, @@ -766,9 +1051,8 @@ def cancel_then_stop(*_args: object, **_kwargs: object) -> RunEvent: frozen = case.collection.repository.get(run.id) assert frozen is not None assert frozen.outer_trace is not None - trace = CanonicalTrace.model_validate(frozen.outer_trace) - root = next(span for span in trace.spans if span.id == trace.root_span_id) - assert root.metadata["message_execution_state"] == ("cancelled" if stop_during_cleanup else "failed") + root = Trace.model_validate(frozen.outer_trace).root + assert root.attributes["message_execution_state"] == ("cancelled" if stop_during_cleanup else "failed") assert root.status.value == "error" assert len(store.saved) == 1 with case.collection.sessions() as session: @@ -777,40 +1061,37 @@ def cancel_then_stop(*_args: object, **_kwargs: object) -> RunEvent: assert message.message_tokens == 13 assert message.answer_tokens == 8 assert message.total_price == Decimal("0.21") - exported: list[CanonicalTrace] = [] - - def export(_provider: str, _config: dict[str, object], payload: CanonicalTrace) -> TraceExportResult: - exported.append(payload) - return TraceExportResult(trace_id=payload.trace_id) - - monkeypatch.setattr(service_module, "export_trace", export) - monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", lambda *_args: {}) case.service.deliver(run.id) case.service.deliver(run.id) - accepted = case.collection.repository.get(run.id) - assert accepted is not None - assert accepted.state == "accepted" - assert accepted.attempts == 1 - assert len(exported) == 1 + delivered = case.collection.repository.get(run.id) + assert delivered is not None + assert delivered.state is TraceRunState.ACCEPTED + assert delivered.attempts == 1 + assert len(case.destination.requests) == 1 @pytest.mark.parametrize("capture_agent", [False, True]) def test_missing_stop_worker_without_runtime_spans_also_reaches_a_delivery_failure( case: RecoveryCase, monkeypatch: pytest.MonkeyPatch, capture_agent: bool ) -> None: - from tests.helpers.unit.config_override import apply_config_overrides - - collection = seed(case.collection.sessions, f"stop-without-spans-{capture_agent}", capture_agent=capture_agent) - attach_conversation(collection) - stopped = RecoveryCase(collection, ApplicationTracingService(collection.repository), case.now, []) - monkeypatch.setattr(stopped.service, "schedule", stopped.scheduled.append) - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_COLLECTION_TIMEOUT=60) + collection = seed( + case.collection.sessions, + f"stop-without-spans-{capture_agent}", + capture_content=False, + capture_agent=capture_agent, + ) + attach_conversation(collection, AppMode.AGENT) + stopped = recovery_case( + collection, + RecordingEgress([], grpc_target=None), + replace(AGENT_TRACE_SETTINGS, collection_timeout_seconds=60), + case.now, + ) save_agent_stop(stopped, monkeypatch, "") - monkeypatch.setattr(recovery, "naive_utc_now", lambda: case.now + timedelta(seconds=120)) - assert recovery.recover_terminal(stopped.service, collection.run) is False + assert stopped.recover(collection.run, now=case.now + timedelta(seconds=120)) is False stored = collection.repository.get(collection.run.id) assert stored is not None - assert stored.state == "failed" + assert stored.state is TraceRunState.FAILED assert stored.error_code == "business_terminal_not_persisted" assert stored.outer_trace is None @@ -825,7 +1106,7 @@ def test_worker_terminal_survives_closed_sse_without_inventing_a_message_write( message.message_metadata = json.dumps({"usage": {"prompt_tokens": 2, "completion_tokens": 1}}) message.answer = "" # A runtime final/usage flush still does not imply that an API message ended. - assert recovery.recover_terminal(case.service, case.collection.run) is False + assert case.recover(case.collection.run, now=case.now) is False output = "actual worker result" if outcome == "succeeded" else None case.service.capture_agent_app_terminal( tenant_id=case.collection.run.tenant_id, @@ -836,9 +1117,8 @@ def test_worker_terminal_survives_closed_sse_without_inventing_a_message_write( ) stored = case.collection.repository.get(case.collection.run.id) assert stored is not None - trace = CanonicalTrace.model_validate(stored.outer_trace) - root = trace.spans[0] - assert root.metadata["message_execution_state"] == outcome + root = Trace.model_validate(stored.outer_trace).root + assert root.attributes["message_execution_state"] == outcome assert root.status.value == ("ok" if outcome == "succeeded" else "error") assert root.outputs == (output or "") assert case.scheduled == [stored.id] @@ -851,9 +1131,7 @@ def test_worker_terminal_survives_closed_sse_without_inventing_a_message_write( def test_worker_terminal_revalidates_message_tenant(case: RecoveryCase) -> None: - from repositories.application_trace_repository import TraceCollectionError - - with pytest.raises(TraceCollectionError, match="tracing_source_not_found"): + with pytest.raises(TraceCollectionNotFoundError, match="tracing_source_not_found"): case.service.capture_agent_app_terminal( tenant_id=identifier("foreign-tenant"), app_id=case.collection.run.app_id, @@ -866,19 +1144,19 @@ def test_worker_terminal_revalidates_message_tenant(case: RecoveryCase) -> None: assert stored.outer_trace is None -def test_recovers_real_message_terminal_once_through_existing_trace_task(case: RecoveryCase) -> None: +def test_recovers_real_message_terminal_once_through_the_record_repository(case: RecoveryCase) -> None: with case.collection.sessions.begin() as session: message = session.get(Message, case.collection.run.source_id) assert message is not None message.message_metadata = complete_metadata() message.updated_at = case.now - timedelta(seconds=11) - assert recovery.recover_terminal(case.service, case.collection.run) is True + assert case.recover(case.collection.run, now=case.now) is True stored = case.collection.repository.get(case.collection.run.id) assert stored is not None assert stored.outer_trace is not None assert stored.outer_trace["session_id"] == "session-original" assert stored.outer_trace["spans"][0]["outputs"] == "persisted final answer" - assert recovery.recover_terminal(case.service, stored) is False + assert case.recover(stored, now=case.now) is False assert case.scheduled == [stored.id] with case.collection.sessions() as session: message = session.get(Message, case.collection.run.source_id) @@ -888,29 +1166,21 @@ def test_recovers_real_message_terminal_once_through_existing_trace_task(case: R assert message.answer == "persisted final answer" -def test_deliver_recovers_and_exports_once_after_terminal_queue_loss( - case: RecoveryCase, monkeypatch: pytest.MonkeyPatch -) -> None: +def test_deliver_recovers_and_exports_once_after_terminal_queue_loss(case: RecoveryCase) -> None: with case.collection.sessions.begin() as session: message = session.get(Message, case.collection.run.source_id) assert message is not None message.message_metadata = complete_metadata() message.updated_at = case.now - timedelta(seconds=11) - exported: list[CanonicalTrace] = [] - - def export(_provider: str, _config: dict[str, object], trace: CanonicalTrace) -> TraceExportResult: - exported.append(trace) - return TraceExportResult(trace_id="provider-id") - - monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", lambda *_args: {}) - monkeypatch.setattr(service_module, "export_trace", export) case.service.deliver(case.collection.run.id) case.service.deliver(case.collection.run.id) + exported = sent_traces(case.destination) assert len(exported) == 1 assert exported[0].spans[0].outputs == "persisted final answer" stored = case.collection.repository.get(case.collection.run.id) assert stored is not None - assert stored.state == "accepted" + assert stored.state is TraceRunState.ACCEPTED + assert stored.provider_trace_id == "provider-id" assert stored.attempts == 1 @@ -921,24 +1191,24 @@ def test_early_usage_commit_is_not_treated_as_a_final_answer(case: RecoveryCase) message.message_metadata = json.dumps({"usage": {"prompt_tokens": 2, "completion_tokens": 1}}) message.answer = "" message.updated_at = case.now - timedelta(seconds=11) - assert recovery.recover_terminal(case.service, case.collection.run) is False + assert case.recover(case.collection.run, now=case.now) is False stored = case.collection.repository.get(case.collection.run.id) assert stored is not None assert stored.outer_trace is None - assert stored.state == "collecting" + assert stored.state is TraceRunState.COLLECTING assert not case.scheduled def test_final_without_business_commit_times_out_as_failed_not_success(case: RecoveryCase) -> None: - invocations, _ = case.collection.repository.collected(case.collection.run.id) + invocations = case.collection.repository.invocations(case.collection.run.id) with case.collection.sessions.begin() as session: invocation = session.get(ApplicationTraceInvocation, invocations[0].id) assert invocation is not None invocation.finalized_at = case.now - timedelta(hours=1) - assert recovery.recover_terminal(case.service, case.collection.run) is False + assert case.recover(case.collection.run, now=case.now) is False stored = case.collection.repository.get(case.collection.run.id) assert stored is not None - assert stored.state == "failed" + assert stored.state is TraceRunState.FAILED assert stored.error_code == "business_terminal_not_persisted" assert stored.outer_trace is None @@ -949,13 +1219,13 @@ def test_final_settle_window_and_all_invocations_must_finish(case: RecoveryCase) assert message is not None message.message_metadata = complete_metadata() message.updated_at = case.now - timedelta(seconds=1) - assert recovery.recover_terminal(case.service, case.collection.run) is False + assert case.recover(case.collection.run, now=case.now) is False second = case.collection.invocation("pending-second") with case.collection.sessions.begin() as session: message = session.get(Message, case.collection.run.source_id) assert message is not None message.updated_at = case.now - timedelta(seconds=10) - assert recovery.recover_terminal(case.service, case.collection.run) is False + assert case.recover(case.collection.run, now=case.now) is False with case.collection.sessions() as session: invocation = session.get(ApplicationTraceInvocation, second.id) assert invocation is not None @@ -973,43 +1243,65 @@ def test_explicit_message_failure_can_recover_without_fabricating_usage(case: Re message.message_tokens = 0 message.answer_tokens = 0 message.updated_at = case.now - timedelta(seconds=11) - assert recovery.recover_terminal(case.service, case.collection.run) is True + assert case.recover(case.collection.run, now=case.now) is True stored = case.collection.repository.get(case.collection.run.id) assert stored is not None assert stored.outer_trace is not None - assert stored.outer_trace["spans"][0]["status"] == "error" - assert stored.outer_trace["spans"][0]["outputs"] == "" - assert stored.outer_trace["spans"][0]["metadata"]["aggregate_total_tokens"] == 0 + root = stored.outer_trace["spans"][0] + assert root["status"] == "error" + assert root["error"] == "model initialization failed" + assert root["outputs"] == "" + # The message's own model call is left to the Agent runtime; the root carries no usage of its own. + assert [(span["usage"], span["cost"]) for span in stored.outer_trace["spans"]] == [(None, None)] -def test_cross_tenant_source_is_rejected_before_trace_task_executes( - case: RecoveryCase, monkeypatch: pytest.MonkeyPatch -) -> None: +def test_cross_tenant_source_is_rejected_before_the_record_repository_runs(case: RecoveryCase) -> None: with case.collection.sessions.begin() as session: app = session.get(App, case.collection.run.app_id) assert app is not None app.tenant_id = identifier("other-tenant") - called: list[bool] = [] + records = _FailingRecords(AssertionError("a foreign source must never be assembled")) + service = tracing_service(case.collection.repository, records, case.destination, case.scheduled, case.settings) + recovered = recover( + case.collection, case.collection.run, records=records, collector=service, settings=case.settings, now=case.now + ) + assert recovered is False + assert records.events == [] + stored = case.collection.repository.get(case.collection.run.id) + assert stored is not None + assert stored.error_code == "tracing_source_not_found" - def forbidden(_task: TraceTask) -> None: - called.append(True) - monkeypatch.setattr(TraceTask, "execute", forbidden) - assert recovery.recover_terminal(case.service, case.collection.run) is False - assert called == [] +def test_source_not_ready_for_assembly_is_retried_later(case: RecoveryCase) -> None: + with case.collection.sessions.begin() as session: + message = session.get(Message, case.collection.run.source_id) + assert message is not None + message.message_metadata = complete_metadata() + message.updated_at = case.now - timedelta(seconds=11) + records = _FailingRecords(SourceNotReadyError()) + service = tracing_service(case.collection.repository, records, case.destination, case.scheduled, case.settings) + recovered = recover( + case.collection, case.collection.run, records=records, collector=service, settings=case.settings, now=case.now + ) + assert recovered is False + assert [event.payload for event in records.events] == [ + MessageCompleted(message_id=case.collection.run.source_id, conversation_id=identifier("conversation-recover")) + ] + assert records.events[0].correlation.session_id == "session-original" stored = case.collection.repository.get(case.collection.run.id) assert stored is not None - assert stored.error_code == "tracing_source_not_found" + assert (stored.state, stored.error_code, stored.outer_trace) == (TraceRunState.COLLECTING, None, None) + assert case.scheduled == [] def test_unresolved_recovery_rotates_behind_unchecked_sources(case: RecoveryCase) -> None: - other = seed(case.collection.sessions, "other-recovery") + other = seed(case.collection.sessions, "other-recovery", capture_content=False, capture_agent=True) assert case.collection.repository.pending_ids(limit=1) == [case.collection.run.id] - assert recovery.recover_terminal(case.service, case.collection.run) is False + assert case.recover(case.collection.run, now=case.now) is False assert case.collection.repository.pending_ids(limit=1) == [other.run.id] stored = case.collection.repository.get(case.collection.run.id) assert stored is not None - assert stored.state == "collecting" + assert stored.state is TraceRunState.COLLECTING @pytest.mark.parametrize("status", [MessageStatus.PAUSED, MessageStatus.NORMAL]) @@ -1019,10 +1311,53 @@ def test_no_final_marker_or_paused_message_is_not_recovered(case: RecoveryCase, assert message is not None message.status = status message.updated_at = case.now - timedelta(seconds=10) - assert recovery.recover_terminal(case.service, case.collection.run) is False + assert case.recover(case.collection.run, now=case.now) is False assert not case.scheduled +def add_workflow_run(case: RecoveryCase, status: WorkflowExecutionStatus, *, finished: bool, created_by: str) -> str: + workflow_run_id = identifier("recover-workflow-run") + workflow_run = WorkflowRun( + id=workflow_run_id, + tenant_id=case.collection.run.tenant_id, + app_id=case.collection.run.app_id, + workflow_id=identifier("workflow"), + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="1", + graph="{}", + inputs="{}", + outputs=json.dumps({"result": "persisted workflow output"}), + status=status, + error="actual workflow failure" + if status in {WorkflowExecutionStatus.FAILED, WorkflowExecutionStatus.STOPPED} + else None, + created_by=created_by, + created_by_role=CreatorUserRole.ACCOUNT, + finished_at=case.now - timedelta(seconds=10) if finished else None, + elapsed_time=20, + total_tokens=0, + total_steps=0, + ) + workflow_run.created_at = case.now - timedelta(seconds=30) + with case.collection.sessions.begin() as session: + session.add(workflow_run) + return workflow_run_id + + +def start_workflow(case: RecoveryCase, workflow_run_id: str) -> TraceRun: + run = case.service.start( + tenant_id=case.collection.run.tenant_id, + app_id=case.collection.run.app_id, + source_kind="workflow", + source_id=workflow_run_id, + session_id="workflow-session", + workflow_id=None, + ) + assert run is not None + return run + + @pytest.mark.parametrize( ("status", "finished"), [ @@ -1038,59 +1373,54 @@ def test_no_final_marker_or_paused_message_is_not_recovered(case: RecoveryCase, def test_workflow_recovery_requires_committed_terminal_state( case: RecoveryCase, status: WorkflowExecutionStatus, finished: bool ) -> None: - owner_id, workflow_run_id = identifier("workflow-owner"), identifier("recover-workflow-run") - with case.collection.sessions.begin() as session: - app = session.get(App, case.collection.run.app_id) - assert app is not None - app.created_by = owner_id - session.add(make_account(account_id=owner_id)) - trace_workflow_record = WorkflowRun( - id=workflow_run_id, - tenant_id=case.collection.run.tenant_id, - app_id=app.id, - workflow_id=identifier("workflow"), - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="1", - graph="{}", - inputs="{}", - outputs=json.dumps({"result": "persisted workflow output"}), - status=status, - error="actual workflow failure" - if status in {WorkflowExecutionStatus.FAILED, WorkflowExecutionStatus.STOPPED} - else None, - created_by=owner_id, - created_by_role=CreatorUserRole.ACCOUNT, - finished_at=case.now - timedelta(seconds=10) if finished else None, - elapsed_time=20, - total_tokens=0, - total_steps=0, - ) - trace_workflow_record.created_at = case.now - timedelta(seconds=30) - session.add(trace_workflow_record) - run = case.collection.repository.start( - tenant_id=case.collection.run.tenant_id, - app_id=case.collection.run.app_id, - source_kind="workflow", - source_id=workflow_run_id, - session_id="workflow-session", - ) - assert run is not None + owner_id = add_app_owner(case.collection, "workflow") + workflow_run_id = add_workflow_run(case, status, finished=finished, created_by=owner_id) + run = start_workflow(case, workflow_run_id) expected = finished and status in { WorkflowExecutionStatus.SUCCEEDED, WorkflowExecutionStatus.PARTIAL_SUCCEEDED, WorkflowExecutionStatus.FAILED, WorkflowExecutionStatus.STOPPED, } - assert recovery.recover_terminal(case.service, run) is expected + assert case.recover(run, now=case.now) is expected stored = case.collection.repository.get(run.id) assert stored is not None if expected: assert stored.outer_trace is not None - assert stored.outer_trace["root_span_id"] == span_id(workflow_run_id) - assert stored.outer_trace["spans"][0]["outputs"] == {"result": "persisted workflow output"} + assert stored.root_span_id == span_id_for(workflow_run_id) + root = stored.outer_trace["spans"][0] + assert root["span_id"] == span_id_for(workflow_run_id) + assert root["outputs"] == {"result": "persisted workflow output"} + succeeded = status in {WorkflowExecutionStatus.SUCCEEDED, WorkflowExecutionStatus.PARTIAL_SUCCEEDED} + assert root["status"] == ("ok" if succeeded else "error") + assert stored.outer_trace["session_id"] == "workflow-session" + assert stored.outer_trace["user_id"] == owner_id else: assert stored.outer_trace is None + assert (stored.state, stored.error_code) == (TraceRunState.COLLECTING, None) + + +@pytest.mark.parametrize("owner_state", ["no_creator", "missing_account", "not_a_member"]) +def test_workflow_recovery_does_not_depend_on_the_app_creator(case: RecoveryCase, owner_state: str) -> None: + owner_id = add_app_owner(case.collection, "workflow") + with case.collection.sessions.begin() as session: + app = session.get(App, case.collection.run.app_id) + assert app is not None + if owner_state == "no_creator": + app.created_by = None + elif owner_state == "missing_account": + account = session.get(Account, owner_id) + assert account is not None + session.delete(account) + else: + session.execute(delete(TenantAccountJoin).where(TenantAccountJoin.account_id == owner_id)) + workflow_run_id = add_workflow_run(case, WorkflowExecutionStatus.SUCCEEDED, finished=True, created_by=owner_id) + run = start_workflow(case, workflow_run_id) + assert case.recover(run, now=case.now) is True + stored = case.collection.repository.get(run.id) + assert stored is not None + assert stored.outer_trace is not None + assert stored.root_span_id == span_id_for(workflow_run_id) @pytest.mark.parametrize("formatted", [None, "", " \n\t", "{}", "[]", "null", " { } ", "[\n]", {}, []]) @@ -1098,10 +1428,6 @@ def test_workflow_recovery_requires_committed_terminal_state( def test_early_error_root_preserves_known_query_only_when_content_is_enabled( case: RecoveryCase, formatted: object, capture_content: bool ) -> None: - from sqlalchemy import update - - from models.application_trace import ApplicationTraceRun - query = "KNOWN_USER_QUERY_BEFORE_FORMATTED_MESSAGE" with case.collection.sessions.begin() as session: # Historical rows can contain JSON strings/lists despite the ORM's dict annotation. @@ -1122,15 +1448,15 @@ def test_early_error_root_preserves_known_query_only_when_content_is_enabled( ) stored = case.collection.repository.get(case.collection.run.id) assert stored is not None - trace = CanonicalTrace.model_validate(stored.outer_trace) - root = trace.spans[0] + root = Trace.model_validate(stored.outer_trace).root assert root.status.value == "error" - assert root.metadata["message_execution_state"] == "failed" + assert root.attributes["message_execution_state"] == "failed" assert root.inputs == ({"query": query} if capture_content else None) assert root.outputs == ("" if capture_content else None) - assembled = case.service.assemble(stored) + case.service.deliver(stored.id) + [assembled] = sent_traces(case.destination) assert (query in assembled.model_dump_json()) is capture_content - assert all(span.inputs is None for span in assembled.spans if span.id != assembled.root_span_id) + assert all(span.inputs is None for span in assembled.spans if span.span_id != assembled.root.span_id) with case.collection.sessions() as session: message = session.get(Message, case.collection.run.source_id) assert message is not None @@ -1166,8 +1492,6 @@ def test_early_error_root_preserves_known_query_only_when_content_is_enabled( def test_worker_handoff_never_replaces_existing_formatted_inputs( case: RecoveryCase, formatted: object, expected: object ) -> None: - from sqlalchemy import update - query = "QUERY_MUST_NOT_REPLACE_FORMATTED_INPUTS" with case.collection.sessions.begin() as session: session.execute( @@ -1184,8 +1508,8 @@ def test_worker_handoff_never_replaces_existing_formatted_inputs( ) stored = case.collection.repository.get(case.collection.run.id) assert stored is not None - trace = CanonicalTrace.model_validate(stored.outer_trace) - assert trace.spans[0].inputs == expected + trace = Trace.model_validate(stored.outer_trace) + assert trace.root.inputs == expected assert query not in trace.model_dump_json() with case.collection.sessions() as session: message = session.get(Message, case.collection.run.source_id) @@ -1209,4 +1533,23 @@ def test_worker_handoff_does_not_invent_input_when_query_is_empty(case: Recovery ) stored = case.collection.repository.get(case.collection.run.id) assert stored is not None - assert CanonicalTrace.model_validate(stored.outer_trace).spans[0].inputs == {} + assert Trace.model_validate(stored.outer_trace).root.inputs == {} + + +def test_worker_handoff_leaves_a_not_collected_run_untouched(case: RecoveryCase) -> None: + with case.collection.sessions.begin() as session: + run = session.get(ApplicationTraceRun, case.collection.run.id) + assert run is not None + run.state = TraceRunState.NOT_COLLECTED + case.service.capture_agent_app_terminal( + tenant_id=case.collection.run.tenant_id, + app_id=case.collection.run.app_id, + message_id=case.collection.run.source_id, + outcome="succeeded", + output="answer", + ) + stored = case.collection.repository.get(case.collection.run.id) + assert stored is not None + assert stored.outer_trace is None + assert case.scheduled == [] + assert claim(case.collection, case.collection.run.id) is None diff --git a/api/tests/unit_tests/services/tracing/test_registry.py b/api/tests/unit_tests/services/tracing/test_registry.py new file mode 100644 index 00000000000000..c02662bbd16db8 --- /dev/null +++ b/api/tests/unit_tests/services/tracing/test_registry.py @@ -0,0 +1,51 @@ +from collections.abc import Iterator +from importlib.metadata import EntryPoint + +import pytest + +from services.tracing import registry +from services.tracing.errors import TraceProviderUnavailableError, UnknownTraceProviderError +from services.tracing.providers import SUPPORTED_PROVIDERS +from services.tracing.registry import ENTRY_POINT_GROUP, EntryPointTraceProviderCatalog + + +@pytest.fixture(autouse=True) +def fresh_registry() -> Iterator[None]: + registry._load.cache_clear() + yield + registry._load.cache_clear() + + +def test_every_supported_provider_is_installed_under_its_own_id() -> None: + catalog = EntryPointTraceProviderCatalog() + + for provider_id in SUPPORTED_PROVIDERS: + assert catalog.resolve(provider_id).id == provider_id + + +def test_unknown_providers_are_client_errors() -> None: + with pytest.raises(UnknownTraceProviderError): + EntryPointTraceProviderCatalog().resolve("datadog") + + +@pytest.mark.parametrize( + "entry_points", + [ + {}, + {"langfuse": EntryPoint(name="langfuse", value="not_a_trace_package:Provider", group=ENTRY_POINT_GROUP)}, + { + "langfuse": EntryPoint( + name="langfuse", value="dify_trace_opik.provider:OpikProvider", group=ENTRY_POINT_GROUP + ) + }, + {"langfuse": EntryPoint(name="langfuse", value="dify_trace_sdk.ids:span_id_for", group=ENTRY_POINT_GROUP)}, + ], + ids=["not-installed", "import-fails", "wrong-provider", "not-a-provider"], +) +def test_supported_providers_that_do_not_load_are_unavailable( + monkeypatch: pytest.MonkeyPatch, entry_points: dict[str, EntryPoint] +) -> None: + monkeypatch.setattr(registry, "_entry_points", lambda: entry_points) + + with pytest.raises(TraceProviderUnavailableError): + EntryPointTraceProviderCatalog().resolve("langfuse") diff --git a/api/tests/unit_tests/services/tracing/test_request_log_protection.py b/api/tests/unit_tests/services/tracing/test_request_log_protection.py index a989a7f140fa64..a29274c8e1ff0e 100644 --- a/api/tests/unit_tests/services/tracing/test_request_log_protection.py +++ b/api/tests/unit_tests/services/tracing/test_request_log_protection.py @@ -1,4 +1,4 @@ -"""Real HTTPX/SSRF diagnostics with synthetic credential-bearing destinations.""" +"""Real HTTPX/SSRF diagnostics with synthetic credential-bearing destinations, sent through the egress.""" import logging import traceback @@ -10,13 +10,13 @@ import httpx import pytest +from dify_trace_sdk.egress import EgressResponse +from dify_trace_sdk.errors import EgressBlockedError, EgressError from opentelemetry.context import _SUPPRESS_INSTRUMENTATION_KEY, get_value from core.helper import ssrf_proxy from core.logging.structured_formatter import StructuredJSONFormatter -from core.ops.exceptions import RetryableTraceDispatchError -from services.tracing.provider_common import request -from services.tracing.provider_dispatch import TraceExportRejectedError, TraceExportUncertainError +from services.tracing.egress import SSRFEgress from services.tracing.request_log_protection import SAFE_HTTP_LOG_MESSAGE, protect_http_request_logs SYNTHETIC_KEY = "SYNTHETIC-NOT-A-REAL-LICENSE/+?=测试" @@ -49,6 +49,10 @@ def get_client(ssl_verify_enabled: bool) -> httpx.Client: yield install +def send(url: str, *, sensitive: bool) -> EgressResponse: + return SSRFEgress().request("POST", url, headers={}, content=b"synthetic", timeout=30.0, sensitive_url=sensitive) + + def evidence(caplog: pytest.LogCaptureFixture) -> str: formatter = StructuredJSONFormatter() return caplog.text + "".join(repr(record.__dict__) + formatter.format(record) for record in caplog.records) @@ -56,7 +60,7 @@ def evidence(caplog: pytest.LogCaptureFixture) -> str: @pytest.mark.parametrize("outcome", [200, 503, 302, "squid", "timeout", "invalid"]) @pytest.mark.parametrize("scheme", ["http", "https"]) -def test_aliyun_keeps_real_wire_url_but_protects_transport_logs_and_exception_chains( +def test_sensitive_urls_reach_the_wire_but_never_logs_or_exception_chains( install_client: InstallClient, caplog: pytest.LogCaptureFixture, outcome: int | str, scheme: str ) -> None: caplog.set_level(logging.INFO) @@ -94,20 +98,15 @@ def handler(wire: httpx.Request) -> httpx.Response: return httpx.Response(outcome, headers={"Location": "https://redirect.test/"}) install_client(handler) - if outcome == 200: - result = request("POST", private_url, "aliyun", content=b"synthetic") - assert result.status_code == 200 + if isinstance(outcome, int): + # Providers classify statuses; the egress only bounds and returns the response. + assert send(private_url, sensitive=True).status_code == outcome rendered_exception = "" else: - expected = ( - TraceExportUncertainError - if outcome == "timeout" - else RetryableTraceDispatchError - if outcome == 503 - else TraceExportRejectedError - ) + expected = EgressBlockedError if outcome == "squid" else EgressError with pytest.raises(expected) as caught: - request("POST", private_url, "aliyun", content=b"synthetic") + send(private_url, sensitive=True) + assert caught.value.request_sent is (outcome == "timeout") assert caught.value.__cause__ is None assert caught.value.__context__ is None rendered_exception = "".join(traceback.format_exception(caught.value)) @@ -152,7 +151,7 @@ def handler(wire: httpx.Request) -> httpx.Response: install_client(handler) def run(label: str) -> int: - result = request("POST", urls[label], "aliyun" if label in keys else "weave") + result = send(urls[label], sensitive=label in keys) if label == "a": first_finished.set() return result.status_code diff --git a/api/tests/unit_tests/services/tracing/test_runtime_bridge.py b/api/tests/unit_tests/services/tracing/test_runtime_bridge.py index 3776be89018bdf..9cf24c29429712 100644 --- a/api/tests/unit_tests/services/tracing/test_runtime_bridge.py +++ b/api/tests/unit_tests/services/tracing/test_runtime_bridge.py @@ -1,39 +1,44 @@ import json from copy import deepcopy -from dataclasses import dataclass, field, replace +from dataclasses import dataclass, replace from datetime import UTC, datetime, timedelta -from typing import cast -from uuid import UUID, uuid4 +from decimal import Decimal +from pathlib import Path +from typing import Literal, override +from uuid import NAMESPACE_URL, UUID, uuid4, uuid5 -import httpx import pytest from dify_agent.layers.execution_context import ApplicationTracingContext from dify_agent.protocol.tracing import AgentTraceBatch, AgentTraceSpan +from dify_trace_langfuse.provider import LangfuseProvider +from dify_trace_sdk.egress import Egress +from dify_trace_sdk.events import ( + ConversationNamed, + MessageCompleted, + ModerationChecked, + TraceCorrelation, + TraceEvent, + TracePayload, + WorkflowRunCompleted, +) +from dify_trace_sdk.ids import span_id_for +from dify_trace_sdk.model import Span, SpanKind, SpanStatus, Trace +from dify_trace_sdk.testing import RecordingEgress, response from flask import Flask -from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest -from sqlalchemy import delete, func, select -from sqlalchemy.exc import OperationalError +from sqlalchemy import URL, Engine, create_engine, delete, event, func, select from sqlalchemy.orm import Session, sessionmaker from core.app.app_config.entities import WorkflowUIBasedAppConfig from core.app.apps.agent_app.runtime_request_builder import AgentAppRuntimeBuildContext from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom, WorkflowAppGenerateEntity -from core.ops.base_trace_instance import bind_application_tracing_hooks -from core.ops.entities.trace_entity import ( - BaseTraceInfo, - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - WorkflowTraceInfo, -) -from core.ops.exceptions import RetryableTraceDispatchError -from core.ops.ops_trace_manager import OpsTraceManager -from core.ops.unified_trace.trace_builder import CanonicalTraceBuilder from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository from core.workflow.nodes.agent_v2.runtime_request_builder import WorkflowAgentRuntimeBuildContext -from extensions.ext_database import db +from domain.tracing.trace_run_agg import TraceRunState, TraceSourceKind +from extensions.tracing import application_tracing_hooks, bind_tracing from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus from graphon.runtime import VariablePool +from libs.datetime_utils import naive_utc_now +from models.account import Tenant, TenantAccountJoin, TenantAccountRole from models.agent import ( Agent, AgentConfigSnapshot, @@ -44,36 +49,45 @@ ) from models.agent_config_entities import AgentSoulConfig, WorkflowNodeJobConfig from models.application_trace import ApplicationTraceRun -from models.enums import ConversationFromSource, CreatorUserRole -from models.enums import WorkflowRunTriggeredFrom as _application_trace_handoff_WorkflowRunTriggeredFrom +from models.enums import ConversationFromSource, CreatorUserRole, WorkflowRunTriggeredFrom from models.model import App, AppMode, Message, TraceAppConfig from models.workflow import ( + WorkflowAppLog, + WorkflowAppLogCreatedFrom, WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom, WorkflowRun, + WorkflowType, ) -from models.workflow import ( - WorkflowRunTriggeredFrom as _workflow_tool_tracing_boundary_WorkflowRunTriggeredFrom, -) -from models.workflow import WorkflowType as _application_trace_handoff_WorkflowType +from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository from repositories.application_trace_repository import ApplicationTraceRepository -from services.tracing import runtime_bridge -from services.tracing.normalization import remap_outer_trace -from services.tracing.provider_dispatch import export_trace +from repositories.factory import DifyAPIRepositoryFactory +from repositories.trace_record_repository import TraceRecordRepository +from services.tracing.adapters import StoredPolicyTraceScopes +from services.tracing.canonical import CanonicalTraceBuilder +from services.tracing.collection import AgentTraceSettings, AppTracingSnapshot, TraceRun, workflow_has_agent_v2 +from services.tracing.delivery import ExportJob, ExportRetryQueue, TraceDeliveryService +from services.tracing.errors import TraceCollectorUnavailableError, TracingStoreUnavailableError +from services.tracing.providers import ProviderConfigCodec +from services.tracing.records import WorkflowTraceInfo +from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks +from services.tracing.scopes import TraceScopes from services.tracing.service import ApplicationTracingService -from tasks import ops_trace_task -from tests.helpers.unit.config_override import apply_config_overrides -from tests.helpers.unit.model_factories import make_account, make_app, make_conversation, make_workflow -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_batch_for as batch_for -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_identifier as identifier -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_ingest as ingest -from tests.unit_tests.services.tracing.test_service import ( - _application_trace_collection_seed as _application_trace_handoff_seed, +from tests.helpers.unit.model_factories import ( + make_account, + make_app, + make_conversation, + make_message, + make_tenant, + make_workflow, +) +from tests.helpers.unit.services.tracing_ports import ( + AGENT_TRACE_SETTINGS, + InstalledTraceProviders, + WorkspaceBoundCipher, ) -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_terminal as terminal -from tests.unit_tests.services.tracing.test_service import _ApplicationTraceCollectionCollection as Collection -"""SQLite ownership regressions at the actual legacy-task handoff boundary.""" +"""SQLite ownership regressions at the handoff from trace delivery to Agent trace collection.""" AGENT_GRAPH = { @@ -82,85 +96,349 @@ } -@dataclass -class Case: - collection: Collection - service: ApplicationTracingService - scheduled: list[str] +STARTED_AT = datetime(2026, 10, 15, tzinfo=UTC) + + +def identifier(value: str) -> str: + return str(uuid5(NAMESPACE_URL, value)) + + +@dataclass(frozen=True) +class Collection: + repository: ApplicationTraceRepository + sessions: sessionmaker[Session] + run: TraceRun + + +class _Recorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +class _SingleAttempt(ExportRetryQueue): + """Delivery here makes one provider attempt; it never schedules a retry.""" + + @override + def schedule(self, job: ExportJob, *, delay_seconds: int) -> None: + raise AssertionError(f"unexpected export retry: provider={job.provider_id}") + + +class _FreezeFails(ApplicationTraceRepository): + """The real ledger, except that reading the tracing snapshot a run freezes fails with ``failure``.""" + + def __init__(self, sessions: sessionmaker[Session], failure: Exception) -> None: + super().__init__(sessions) + self._failure = failure + + @override + def app_tracing(self, *, tenant_id: str, app_id: str) -> AppTracingSnapshot | None: + raise self._failure + + +def record_source(sessions: sessionmaker[Session]) -> TraceRecordRepository: + return TraceRecordRepository( + sessions, + workflow_runs=DifyAPIRepositoryFactory.create_api_workflow_run_repository(sessions), + node_executions=DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(sessions), + files_url="https://files.test", + enterprise_details=False, + ) + + +def tracing_service( + sessions: sessionmaker[Session], + traces: ApplicationTraceRepository, + egress: Egress, + scheduled: list[str], + settings: AgentTraceSettings, +) -> ApplicationTracingService: + """Agent trace collection exporting to Langfuse; scheduled exports are recorded instead of enqueued.""" + return ApplicationTracingService( + runs=traces, + run_query=traces, + sources=traces, + records=record_source(sessions), + builder=CanonicalTraceBuilder(), + providers=InstalledTraceProviders(LangfuseProvider), + codec=ProviderConfigCodec(WorkspaceBoundCipher()), + egress=egress, + schedule_export=scheduled.append, + settings=settings, + ) + + +def trace_scopes(sessions: sessionmaker[Session]) -> StoredPolicyTraceScopes: + """Trace scopes as the composition root binds them, with events kept in memory.""" + return StoredPolicyTraceScopes( + TraceScopes( + emitter=_Recorder(), + policies=SQLAlchemyAppTracingConfigRepository(session_factory=sessions), + enterprise_enabled=lambda: False, + ) + ) + + +def langfuse_config(tenant_id: str, host: str) -> dict[str, str]: + """A stored Langfuse configuration whose secrets only ``tenant_id`` decrypts.""" + cipher = WorkspaceBoundCipher() + return { + "public_key": cipher.encrypt(tenant_id, "public"), + "secret_key": cipher.encrypt(tenant_id, "secret"), + "host": host, + } + +def no_egress() -> RecordingEgress: + """An egress that fails any request the test did not script.""" + return RecordingEgress([], grpc_target=None) -def setup( - sessions: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, name: str, mode: AppMode, *, snapshot: bool -) -> Case: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - collection = _application_trace_handoff_seed(sessions, name) + +def exported() -> RecordingEgress: + """A provider destination that accepts exactly one ingestion request.""" + return RecordingEgress([response(200, {})], grpc_target=None) + + +def ingested(destination: RecordingEgress) -> list[dict[str, object]]: + """The Langfuse ingestion events the provider sent.""" + events: list[dict[str, object]] = [] + for request in destination.requests: + body = request.json() + assert isinstance(body, dict) + batch = body.get("batch") + assert isinstance(batch, list) + events.extend(item for item in batch if isinstance(item, dict)) + return events + + +def observations(destination: RecordingEgress) -> list[dict[str, object]]: + """The Langfuse observations the provider sent: every span except the trace's root.""" + bodies: list[dict[str, object]] = [] + for item in ingested(destination): + body = item["body"] + assert isinstance(body, dict) + if item["type"] != "trace-create": + bodies.append(body) + return bodies + + +def run_count(sessions: sessionmaker[Session]) -> int | None: + with sessions() as session: + return session.scalar(select(func.count()).select_from(ApplicationTraceRun)) + + +def seed(sessions: sessionmaker[Session], name: str) -> Collection: + """An Agent app tracing to Langfuse, one of its messages, and the run frozen when the message started.""" + app_id, tenant_id, message_id = identifier(f"app-{name}"), identifier("tenant"), identifier(f"message-{name}") with sessions.begin() as session: - app = session.get(App, collection.run.app_id) + session.add( + make_app( + app_id=app_id, + tenant_id=tenant_id, + mode=AppMode.AGENT, + tracing=json.dumps( + {"enabled": True, "tracing_provider": "langfuse", "capture_content": False, "capture_agent": True} + ), + ) + ) + session.add( + make_message( + message_id=message_id, + app_id=app_id, + conversation_id=identifier(f"conversation-{name}"), + inputs={}, + query="hello", + message={}, + answer="done", + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.CONSOLE, + ) + ) + session.add( + TraceAppConfig( + app_id=app_id, + tracing_provider="langfuse", + tracing_config=langfuse_config(tenant_id, "https://langfuse.test"), + ) + ) + repository = ApplicationTraceRepository(sessions) + run = tracing_service(sessions, repository, no_egress(), [], AGENT_TRACE_SETTINGS).start( + tenant_id=tenant_id, + app_id=app_id, + source_kind="message", + source_id=message_id, + session_id="session-original", + workflow_id=None, + ) + assert run is not None + return Collection(repository, sessions, run) + + +def attach_conversation(collection: Collection, mode: AppMode) -> None: + """Message records are assembled only for a message whose conversation belongs to its app.""" + with collection.sessions.begin() as session: message = session.get(Message, collection.run.source_id) - assert app is not None assert message is not None - app.mode = mode session.add( make_conversation( conversation_id=message.conversation_id, - app_id=app.id, + app_id=collection.run.app_id, mode=mode, inputs={}, from_source=ConversationFromSource.CONSOLE, ) ) - if not snapshot: - session.execute(delete(ApplicationTraceRun).where(ApplicationTraceRun.id == collection.run.id)) - service = ApplicationTracingService(collection.repository) + + +def add_app_owner(collection: Collection, name: str) -> str: + """An account of the app's workspace that creates the app.""" + owner_id = identifier(f"owner-{name}") + with collection.sessions.begin() as session: + if session.get(Tenant, collection.run.tenant_id) is None: + session.add(make_tenant(tenant_id=collection.run.tenant_id)) + session.add(make_account(account_id=owner_id, email=f"{name}@example.com")) + session.add( + TenantAccountJoin(tenant_id=collection.run.tenant_id, account_id=owner_id, role=TenantAccountRole.OWNER) + ) + app = session.get(App, collection.run.app_id) + assert app is not None + app.created_by = owner_id + return owner_id + + +def trace_event(run: TraceRun, payload: TracePayload) -> TraceEvent: + return TraceEvent( + event_id=uuid4().hex, + tenant_id=run.tenant_id, + app_id=run.app_id, + user_id=None, + occurred_at=datetime.now(UTC), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + payload=payload, + ) + + +def capture(traces: ApplicationTraceRepository, run_id: str, trace: Trace) -> None: + """Keep a business terminal in its run without scheduling the run's export.""" + run = traces.get(run_id) + assert run is not None + with traces.change(run_id, tenant_id=run.tenant_id) as locked: + assert locked is not None + locked.run.capture_terminal(trace, now=naive_utc_now()) + + +def ingest(traces: ApplicationTraceRepository, batch: AgentTraceBatch, *, max_spans: int, max_bytes: int) -> None: + """Store a runtime batch in its run without scheduling the run's export.""" + run = traces.get(str(batch.run_id)) + assert run is not None + with traces.change(run.id, tenant_id=run.tenant_id) as locked: + assert locked is not None + locked.run.ingest(batch, max_spans=max_spans, max_bytes=max_bytes, now=naive_utc_now()) + + +def terminal(collection: Collection) -> None: + """Persist a root-only business terminal for the collection's run.""" + run = collection.run + root = Span( + span_id=run.root_span_id, + parent_span_id=None, + name="Agent turn", + kind=SpanKind.CHAIN, + start_time=datetime.fromtimestamp(0, UTC), + end_time=datetime.fromtimestamp(3, UTC), + status=SpanStatus.OK, + ) + trace = Trace( + trace_id=run.trace_id, + external_trace_id=None, + session_id=run.session_id, + user_id=None, + parent=None, + spans=(root,), + ) + capture(collection.repository, run.id, trace) + + +@dataclass(frozen=True) +class Case: + collection: Collection + service: ApplicationTracingService + scheduled: list[str] + + +def handoff(collection: Collection, traces: ApplicationTraceRepository, egress: Egress) -> Case: + """The Agent collector over ``traces``, exporting through ``egress``.""" scheduled: list[str] = [] - monkeypatch.setattr(service, "schedule", scheduled.append) - monkeypatch.setattr(runtime_bridge, "_service", lambda: service) - return Case(collection, service, scheduled) + return Case( + collection, tracing_service(collection.sessions, traces, egress, scheduled, AGENT_TRACE_SETTINGS), scheduled + ) -def message_info(case: Case) -> MessageTraceInfo: +def setup(sessions: sessionmaker[Session], name: str, mode: AppMode, *, snapshot: bool) -> Case: + collection = seed(sessions, name) + attach_conversation(collection, mode) + with sessions.begin() as session: + app = session.get(App, collection.run.app_id) + assert app is not None + app.mode = mode + if not snapshot: + session.execute(delete(ApplicationTraceRun).where(ApplicationTraceRun.id == collection.run.id)) + return handoff(collection, collection.repository, no_egress()) + + +def deliver(case: Case, event: TraceEvent, destination: RecordingEgress) -> None: + """Deliver one event the way the trace worker does, with this case's service as the Agent collector.""" + sessions = case.collection.sessions + TraceDeliveryService( + records=record_source(sessions), + destinations=SQLAlchemyAppTracingConfigRepository(session_factory=sessions), + providers=InstalledTraceProviders(LangfuseProvider), + codec=ProviderConfigCodec(WorkspaceBoundCipher()), + builder=CanonicalTraceBuilder(), + egress=destination, + retries=_SingleAttempt(), + enterprise_sink=None, + collector=case.service, + max_export_attempts=1, + ).deliver(event, enterprise_delivered=False) + + +def message(case: Case) -> Message: with case.collection.sessions() as session: - message = session.get(Message, case.collection.run.source_id) - assert message is not None - return MessageTraceInfo( - metadata={"app_id": message.app_id}, - message_id=message.id, - message_data=message.to_dict(session=session), - conversation_model="agent", - conversation_mode="agent", - message_tokens=0, - answer_tokens=0, - total_tokens=0, - inputs="private input", - outputs="private output", - ) + row = session.get(Message, case.collection.run.source_id) + assert row is not None + return row -def dispatch(case: Case, info: BaseTraceInfo, monkeypatch: pytest.MonkeyPatch) -> list[BaseTraceInfo]: - calls: list[BaseTraceInfo] = [] - failures: list[str] = [] +def message_event(case: Case) -> TraceEvent: + row = message(case) + return trace_event(case.collection.run, MessageCompleted(message_id=row.id, conversation_id=row.conversation_id)) - class Legacy: - def trace(self, restored: BaseTraceInfo) -> None: - calls.append(restored) - payload = json.dumps({"trace_info_type": type(info).__name__, "trace_info": info.model_dump(mode="json")}) - monkeypatch.setattr(ops_trace_task.storage, "load", lambda _path: payload) - monkeypatch.setattr(ops_trace_task.storage, "delete", lambda _path: None) - monkeypatch.setattr(ops_trace_task.redis_client, "incr", failures.append) - monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", lambda _app: Legacy()) - monkeypatch.setattr("extensions.ext_enterprise_telemetry.is_enabled", lambda: False) - app = Flask(__name__) - bind_application_tracing_hooks(app, runtime_bridge.RuntimeApplicationTracingHooks()) - with app.app_context(): - ops_trace_task.process_trace_tasks.run({"app_id": case.collection.run.app_id, "file_id": "fixture"}) - assert failures == [] - return calls +def workflow_event(row: WorkflowRun, parent: WorkflowRun | None) -> TraceEvent: + return TraceEvent( + event_id=uuid4().hex, + tenant_id=row.tenant_id, + app_id=row.app_id, + user_id=None, + occurred_at=datetime.now(UTC), + correlation=TraceCorrelation( + external_trace_id=None, + session_id=None, + parent_workflow_run_id=parent.id if parent else None, + parent_node_execution_id=identifier("parent-tool-node") if parent else None, + ), + payload=WorkflowRunCompleted(workflow_run_id=row.id, conversation_id=None, total_tokens=None), + ) -def workflow( - case: Case, name: str, *, has_agent: bool, snapshot: bool -) -> tuple[WorkflowRun, ApplicationTraceRun | None]: +def workflow(case: Case, name: str, *, has_agent: bool, snapshot: bool) -> tuple[WorkflowRun, TraceRun | None]: graph = AGENT_GRAPH if has_agent else {"nodes": [{"data": {"type": "llm"}}], "edges": list[object]()} now = datetime(2026, 10, 15) with case.collection.sessions.begin() as session: @@ -179,13 +457,13 @@ def workflow( app_id=app.id, tenant_id=app.tenant_id, workflow_id=definition.id, - type=_application_trace_handoff_WorkflowType.WORKFLOW, + type=WorkflowType.WORKFLOW, version="1", graph=json.dumps(graph), inputs="{}", outputs="{}", status=WorkflowExecutionStatus.SUCCEEDED, - triggered_from=_application_trace_handoff_WorkflowRunTriggeredFrom.APP_RUN, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, created_by_role=CreatorUserRole.ACCOUNT, created_by=identifier("owner"), finished_at=now + timedelta(seconds=2), @@ -196,12 +474,13 @@ def workflow( row.created_at = now session.add(row) trace_run = ( - case.collection.repository.start( + case.service.start( tenant_id=row.tenant_id, app_id=row.app_id, source_kind="workflow", source_id=row.id, session_id="workflow-session", + workflow_id=None, ) if snapshot else None @@ -209,7 +488,8 @@ def workflow( return row, trace_run -def workflow_info(row: WorkflowRun, parent: WorkflowRun | None = None) -> WorkflowTraceInfo: +def workflow_info(row: WorkflowRun, parent: WorkflowRun | None) -> WorkflowTraceInfo: + assert row.finished_at is not None metadata: dict[str, object] = {"app_id": row.app_id} if parent is not None: metadata["parent_trace_context"] = { @@ -217,7 +497,10 @@ def workflow_info(row: WorkflowRun, parent: WorkflowRun | None = None) -> Workfl "parent_node_execution_id": identifier("parent-tool-node"), } return WorkflowTraceInfo( - workflow_data=row, + node_executions=[], + conversation_id=None, + message_started_at=None, + workflow_app_log_id=None, workflow_id=row.workflow_id, tenant_id=row.tenant_id, workflow_run_id=row.id, @@ -226,8 +509,12 @@ def workflow_info(row: WorkflowRun, parent: WorkflowRun | None = None) -> Workfl workflow_run_inputs={}, workflow_run_outputs={}, workflow_run_version="1", + error=None, total_tokens=0, + prompt_tokens=None, + completion_tokens=None, file_list=[], + invoked_by=None, query="", metadata=metadata, start_time=row.created_at, @@ -237,38 +524,35 @@ def workflow_info(row: WorkflowRun, parent: WorkflowRun | None = None) -> Workfl class TestApplicationTraceHandoffCases: @staticmethod - def test_chatflow_message_event_cannot_export_again_via_legacy( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch - ) -> None: - case = setup(sqlite_session_factory, monkeypatch, "chatflow", AppMode.ADVANCED_CHAT, snapshot=False) + def test_chatflow_message_event_is_not_exported_again(sqlite_session_factory: sessionmaker[Session]) -> None: + case = setup(sqlite_session_factory, "chatflow", AppMode.ADVANCED_CHAT, snapshot=False) row, trace_run = workflow(case, "chatflow", has_agent=True, snapshot=True) assert trace_run is not None terminal(Collection(case.collection.repository, case.collection.sessions, trace_run)) before = case.collection.repository.get(trace_run.id) assert before is not None with case.collection.sessions.begin() as session: - message = session.get(Message, case.collection.run.source_id) - assert message is not None - message.workflow_run_id = row.id - assert dispatch(case, message_info(case), monkeypatch) == [] + stored = session.get(Message, case.collection.run.source_id) + assert stored is not None + stored.workflow_run_id = row.id + destination = no_egress() + deliver(case, message_event(case), destination) + assert destination.requests == [] after = case.collection.repository.get(trace_run.id) assert after is not None assert after.outer_trace == before.outer_trace assert case.scheduled == [] - with case.collection.sessions() as session: - assert session.scalar(select(func.count()).select_from(ApplicationTraceRun)) == 1 + assert run_count(case.collection.sessions) == 1 @staticmethod @pytest.mark.parametrize("mode", [AppMode.AGENT, AppMode.ADVANCED_CHAT]) def test_snapshot_failure_is_persisted_and_never_falls_back( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, mode: AppMode + sqlite_session_factory: sessionmaker[Session], mode: AppMode ) -> None: - case = setup(sqlite_session_factory, monkeypatch, "failure", mode, snapshot=False) - - def fail_start(**_kwargs: object) -> None: - raise OperationalError("snapshot", {}, Exception("temporary database unavailable")) - - monkeypatch.setattr(case.collection.repository, "start", fail_start) + case = setup(sqlite_session_factory, "failure", mode, snapshot=False) + failing = handoff( + case.collection, _FreezeFails(sqlite_session_factory, TracingStoreUnavailableError()), no_egress() + ) context = AgentAppRuntimeBuildContext( dify_context=DifyRunContext( tenant_id=case.collection.run.tenant_id, @@ -286,8 +570,10 @@ def fail_start(**_kwargs: object) -> None: binding_id=identifier("binding"), backend_binding_ref="backend", ) - assert runtime_bridge.prepare_agent_app(context) is None - assert dispatch(case, message_info(case), monkeypatch) == [] + assert RuntimeApplicationTracingHooks(failing.service).prepare_agent_app(context) is None + destination = no_egress() + deliver(failing, message_event(failing), destination) + assert destination.requests == [] stored = case.collection.repository.source_run( tenant_id=case.collection.run.tenant_id, app_id=case.collection.run.app_id, @@ -295,7 +581,7 @@ def fail_start(**_kwargs: object) -> None: source_id=case.collection.run.source_id, ) assert stored is not None - assert stored.state == "failed" + assert stored.state is TraceRunState.FAILED assert stored.error_code == "tracing_initialization_failed" assert stored.provider_config == {} @@ -303,101 +589,157 @@ def fail_start(**_kwargs: object) -> None: @pytest.mark.parametrize("mode", [AppMode.AGENT, AppMode.ADVANCED_CHAT, AppMode.CHAT]) @pytest.mark.parametrize("kind", ["title", "moderation", "aliased_moderation"]) def test_auxiliary_events_resolve_owner_without_bypassing_content_policy( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, mode: AppMode, kind: str + sqlite_session_factory: sessionmaker[Session], mode: AppMode, kind: str ) -> None: - case = setup(sqlite_session_factory, monkeypatch, "aux", mode, snapshot=False) - with case.collection.sessions() as session: - message = session.get(Message, case.collection.run.source_id) - assert message is not None - if kind == "title": - info: BaseTraceInfo = GenerateNameTraceInfo( - tenant_id=case.collection.run.tenant_id, - conversation_id=message.conversation_id, - metadata={"conversation_id": message.conversation_id}, - inputs="private conversation", - outputs="private generated title", - ) - else: - info = ModerationTraceInfo( - message_id=identifier("workflow-app-log") if kind == "aliased_moderation" else message.id, - message_data=message.to_dict(session=session), - metadata={"message_id": message.id}, - flagged=False, - action="direct_output", - preset_response="", - query="private prompt", - inputs="private prompt", - outputs="private output", + case = setup(sqlite_session_factory, "aux", mode, snapshot=False) + row = message(case) + if kind == "aliased_moderation": + # Chatflow moderation is keyed by its workflow log, not by the message it guards. + workflow_run_id = identifier("aliased-workflow-run") + with case.collection.sessions.begin() as session: + stored = session.get(Message, row.id) + assert stored is not None + stored.workflow_run_id = workflow_run_id + session.add( + WorkflowAppLog( + tenant_id=case.collection.run.tenant_id, + app_id=case.collection.run.app_id, + workflow_id=identifier("aliased-workflow"), + workflow_run_id=workflow_run_id, + created_from=WorkflowAppLogCreatedFrom.WEB_APP, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=identifier("owner"), + ) ) - legacy = dispatch(case, info, monkeypatch) - assert len(legacy) == (1 if mode is AppMode.CHAT else 0) + payload: TracePayload = ( + ConversationNamed( + conversation_id=row.conversation_id, + message_id=row.id, + inputs="private conversation", + name="private generated title", + started_at=STARTED_AT, + ended_at=STARTED_AT, + ) + if kind == "title" + else ModerationChecked( + message_id=row.id, + inputs={"query": "private prompt"}, + query="private prompt", + flagged=False, + action="direct_output", + preset_response="", + started_at=STARTED_AT, + ended_at=STARTED_AT, + ) + ) + destination = exported() if mode is AppMode.CHAT else no_egress() + deliver(case, trace_event(case.collection.run, payload), destination) + assert len(destination.requests) == (1 if mode is AppMode.CHAT else 0) assert case.scheduled == [] + assert run_count(case.collection.sessions) == 0 @staticmethod - def test_plain_workflow_preserves_legacy_parent_context( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + def test_plain_nested_workflow_keeps_its_parent_context_for_the_provider( + sqlite_session_factory: sessionmaker[Session], ) -> None: - parent_case = setup(sqlite_session_factory, monkeypatch, "plain-parent", AppMode.WORKFLOW, snapshot=False) + parent_case = setup(sqlite_session_factory, "plain-parent", AppMode.WORKFLOW, snapshot=False) parent, _ = workflow(parent_case, "plain-parent", has_agent=False, snapshot=False) - child_case = setup(sqlite_session_factory, monkeypatch, "plain-child", AppMode.WORKFLOW, snapshot=False) + child_case = setup(sqlite_session_factory, "plain-child", AppMode.WORKFLOW, snapshot=False) child, _ = workflow(child_case, "plain-child", has_agent=False, snapshot=False) - info = workflow_info(child, parent) - original = dict(info.metadata) - assert child_case.service.capture_terminal(info) is False - assert info.metadata == original - with child_case.collection.sessions() as session: - assert session.scalar(select(func.count()).select_from(ApplicationTraceRun)) == 0 + add_app_owner(child_case.collection, "plain-child") + event = workflow_event(child, parent) + record = record_source(sqlite_session_factory).assemble(event) + assert isinstance(record, WorkflowTraceInfo) + assert child_case.service.capture_terminal(tenant_id=child.tenant_id, info=record) is False + assert child_case.service.link_managed_parent(tenant_id=child.tenant_id, info=record) is record + destination = exported() + deliver(child_case, event, destination) + trace = next(item for item in ingested(destination) if item["type"] == "trace-create") + body = trace["body"] + assert isinstance(body, dict) + assert body["metadata"]["parent_workflow_run_id"] == parent.id + assert body["metadata"]["parent_node_execution_id"] == identifier("parent-tool-node") + assert "linked_parent_workflow_run_id" not in body["metadata"] + assert run_count(sqlite_session_factory) == 0 @staticmethod - def test_plain_chatflow_auxiliary_events_keep_the_legacy_route( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + def test_plain_chatflow_auxiliary_events_keep_the_provider_route( + sqlite_session_factory: sessionmaker[Session], ) -> None: - case = setup(sqlite_session_factory, monkeypatch, "plain-chatflow", AppMode.ADVANCED_CHAT, snapshot=False) + case = setup(sqlite_session_factory, "plain-chatflow", AppMode.ADVANCED_CHAT, snapshot=False) row, _ = workflow(case, "plain-chatflow", has_agent=False, snapshot=False) with case.collection.sessions.begin() as session: - message = session.get(Message, case.collection.run.source_id) - assert message is not None - message.workflow_run_id = row.id - info = GenerateNameTraceInfo( - tenant_id=row.tenant_id, - conversation_id=message.conversation_id, - metadata={"conversation_id": message.conversation_id}, - inputs="legacy input", - outputs="legacy title", + stored = session.get(Message, case.collection.run.source_id) + assert stored is not None + stored.workflow_run_id = row.id + payload = ConversationNamed( + conversation_id=stored.conversation_id, + message_id=stored.id, + inputs="provider input", + name="provider title", + started_at=STARTED_AT, + ended_at=STARTED_AT, ) - assert len(dispatch(case, info, monkeypatch)) == 1 + destination = exported() + deliver(case, trace_event(case.collection.run, payload), destination) + assert len(destination.requests) == 1 @staticmethod - def test_legacy_child_of_managed_parent_uses_explicit_link_not_unavailable_context( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + def test_provider_exported_child_of_managed_parent_names_the_collected_parent( + sqlite_session_factory: sessionmaker[Session], ) -> None: - parent_case = setup(sqlite_session_factory, monkeypatch, "managed-parent", AppMode.WORKFLOW, snapshot=False) + parent_case = setup(sqlite_session_factory, "managed-parent", AppMode.WORKFLOW, snapshot=False) parent, parent_trace = workflow(parent_case, "managed-parent", has_agent=True, snapshot=True) assert parent_trace is not None - child_case = setup(sqlite_session_factory, monkeypatch, "legacy-child", AppMode.WORKFLOW, snapshot=False) - child, _ = workflow(child_case, "legacy-child", has_agent=False, snapshot=False) + child_case = setup(sqlite_session_factory, "provider-child", AppMode.WORKFLOW, snapshot=False) + child, _ = workflow(child_case, "provider-child", has_agent=False, snapshot=False) info = workflow_info(child, parent) - assert child_case.service.capture_terminal(info) is False - assert info.resolved_parent_context == (None, None) - assert info.metadata["linked_parent_workflow_run_id"] == parent.id - assert info.metadata["linked_parent_node_execution_id"] == identifier("parent-tool-node") - assert info.metadata["linked_parent_application_trace_id"] == parent_trace.id - assert info.metadata["linked_parent_trace_id"] == parent_trace.trace_id + assert child_case.service.capture_terminal(tenant_id=child.tenant_id, info=info) is False + linked = child_case.service.link_managed_parent(tenant_id=child.tenant_id, info=info) + assert linked.resolved_parent_context == (None, None) + assert linked.metadata["linked_parent_workflow_run_id"] == parent.id + assert linked.metadata["linked_parent_node_execution_id"] == identifier("parent-tool-node") + assert linked.metadata["linked_parent_application_trace_id"] == parent_trace.id + assert linked.metadata["linked_parent_trace_id"] == parent_trace.trace_id + # Records are frozen: the delivered record is a copy, the original keeps its unavailable context. + assert info.resolved_parent_context == (parent.id, identifier("parent-tool-node")) + add_app_owner(child_case.collection, "provider-child") + destination = exported() + deliver(child_case, workflow_event(child, parent), destination) + trace = next(item for item in ingested(destination) if item["type"] == "trace-create") + body = trace["body"] + assert isinstance(body, dict) + assert body["metadata"]["linked_parent_trace_id"] == parent_trace.trace_id + assert "parent_workflow_run_id" not in body["metadata"] + + @staticmethod + def test_collector_database_failure_redelivers_without_provider_fallback( + sqlite_session_factory: sessionmaker[Session], tmp_path: Path + ) -> None: + case = setup(sqlite_session_factory, "collector-unavailable", AppMode.AGENT, snapshot=True) + # The ledger's database cannot be opened while the business records stay readable. + unreachable = create_engine(URL.create("sqlite", database=str(tmp_path / "missing" / "ledger.sqlite3"))) + destination = no_egress() + try: + collector = handoff( + case.collection, ApplicationTraceRepository(sessionmaker(bind=unreachable)), no_egress() + ) + with pytest.raises(TraceCollectorUnavailableError): + deliver(collector, message_event(case), destination) + finally: + unreachable.dispose() + assert destination.requests == [] @staticmethod def test_mixed_workflow_exports_normal_llm_and_agent_calls_without_double_usage( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + sqlite_session_factory: sessionmaker[Session], ) -> None: - case = setup(sqlite_session_factory, monkeypatch, "mixed", AppMode.WORKFLOW, snapshot=False) + case = setup(sqlite_session_factory, "mixed", AppMode.WORKFLOW, snapshot=False) row, managed = workflow(case, "mixed", has_agent=True, snapshot=True) assert managed is not None + owner = add_app_owner(case.collection, "mixed") agent_node_id = identifier("agent-execution") - owner = identifier("mixed-owner") with case.collection.sessions.begin() as session: - app = session.get(App, row.app_id) - assert app is not None - app.created_by = owner - session.add(make_account(account_id=owner)) for index, kind in enumerate(("llm", "agent")): node_id = identifier("normal-execution") if kind == "llm" else agent_node_id usage: dict[str, int] = ( @@ -438,13 +780,11 @@ def test_mixed_workflow_exports_normal_llm_and_agent_calls_without_double_usage( ) node_record.created_at = row.created_at session.add(node_record) - invocation = case.collection.repository.invocation( - run_id=managed.id, - agent_id=identifier("agent"), - node_execution_id=agent_node_id, - attempt=0, - call_key="agent:0", - ) + with case.collection.repository.change(managed.id, tenant_id=managed.tenant_id) as locked: + assert locked is not None + invocation = locked.run.open_invocation( + agent_id=identifier("agent"), node_execution_id=agent_node_id, attempt=0, call_key="agent:0" + ) spans = [ AgentTraceSpan( trace_id=managed.trace_id, @@ -477,56 +817,51 @@ def test_mixed_workflow_exports_normal_llm_and_agent_calls_without_double_usage( ) for key in ("b", "c") ) - case.collection.repository.ingest( + ingest( + case.collection.repository, AgentTraceBatch( run_id=UUID(managed.id), invocation_id=UUID(invocation.id), batch_id=uuid4(), final=True, spans=spans ), max_spans=100, max_bytes=100000, ) - monkeypatch.setattr(type(db), "engine", property(lambda _db: sqlite_session_factory.kw["bind"])) - assert case.service.capture_terminal(workflow_info(row)) is True - stored = case.collection.repository.get(managed.id) - assert stored is not None - trace = case.service.assemble(stored) - payloads: list[ExportTraceServiceRequest] = [] - - def send(method: str, url: str, **kwargs: object) -> httpx.Response: - body = kwargs["content"] - assert isinstance(body, bytes) - payloads.append(ExportTraceServiceRequest.FromString(body)) - return httpx.Response(200, content=b"", request=httpx.Request(method, url)) - - monkeypatch.setattr("core.helper.ssrf_proxy.make_request", send) - export_trace("langfuse", {"host": "https://langfuse.test", "public_key": "pk", "secret_key": "sk"}, trace) - observations = payloads[0].resource_spans[0].scope_spans[0].spans - attributes = [{item.key: item.value.string_value for item in span.attributes} for span in observations] - generations = [item for item in attributes if item["langfuse.observation.type"] == "generation"] - assert len(generations) == 3 - assert sum(json.loads(item["langfuse.observation.usage_details"])["total"] for item in generations) == 25 - assert sum( - json.loads(item["langfuse.observation.cost_details"])["total"] for item in generations - ) == pytest.approx(0.04) + record = record_source(sqlite_session_factory).assemble(workflow_event(row, None)) + assert isinstance(record, WorkflowTraceInfo) + assert case.service.capture_terminal(tenant_id=row.tenant_id, info=record) is True + assert case.scheduled == [managed.id] + destination = exported() + handoff(case.collection, case.collection.repository, destination).service.deliver(managed.id) + delivered = case.collection.repository.get(managed.id) + assert delivered is not None + assert delivered.state is TraceRunState.ACCEPTED + assert [request.url for request in destination.requests] == ["https://langfuse.test/api/public/ingestion"] + generations = [item["body"] for item in ingested(destination) if item["type"] == "generation-create"] + usages = [body["usage"] for body in generations if isinstance(body, dict)] + assert len(usages) == 3 + assert sum(usage["total"] for usage in usages) == 25 + assert sum(usage["totalCost"] for usage in usages) == pytest.approx(0.04) """Verify execution hook ownership, policy gating, and fail-open business behavior.""" -@dataclass -class BridgeCase: - target: ApplicationTracingService - calls: list[tuple[str, dict[str, object]]] = field(default_factory=list) - failure: Exception | None = None - context: ApplicationTracingContext = field( - default_factory=lambda: ApplicationTracingContext( - run_id=uuid4(), - invocation_id=uuid4(), - trace_id="1" * 32, - parent_span_id="2" * 16, - session_id="session", - capture_content=False, - ) - ) +@dataclass(frozen=True) +class Bridge: + sessions: sessionmaker[Session] + traces: ApplicationTraceRepository + hooks: RuntimeApplicationTracingHooks + + def run(self, *, app_id: str, source_kind: TraceSourceKind, source_id: str) -> TraceRun: + """The run frozen for a source of the bridge's workspace.""" + run = self.traces.source_run(tenant_id="tenant", app_id=app_id, source_kind=source_kind, source_id=source_id) + assert run is not None + return run + + +def bridge_hooks( + sessions: sessionmaker[Session], traces: ApplicationTraceRepository, settings: AgentTraceSettings +) -> RuntimeApplicationTracingHooks: + return RuntimeApplicationTracingHooks(tracing_service(sessions, traces, no_egress(), [], settings)) def run_context(app_id: str) -> DifyRunContext: @@ -553,7 +888,7 @@ def agent_context() -> AgentAppRuntimeBuildContext: ) -def _runtime_bridge_workflow_context(preview: bool) -> WorkflowAgentRuntimeBuildContext: +def workflow_agent_context(preview: bool) -> WorkflowAgentRuntimeBuildContext: return WorkflowAgentRuntimeBuildContext( dify_context=run_context("workflow-app"), agent=Agent( @@ -588,225 +923,225 @@ def _runtime_bridge_workflow_context(preview: bool) -> WorkflowAgentRuntimeBuild ) -def _runtime_bridge_workflow_entity(extras: dict[str, object] | None = None) -> WorkflowAppGenerateEntity: +def workflow_entity(extras: dict[str, object]) -> WorkflowAppGenerateEntity: return WorkflowAppGenerateEntity.model_construct( app_config=WorkflowUIBasedAppConfig( - app_id="app", tenant_id="tenant", app_mode=AppMode.WORKFLOW, workflow_id="workflow" + app_id="workflow-app", tenant_id="tenant", app_mode=AppMode.WORKFLOW, workflow_id="workflow" ), - extras=extras or {}, + extras=extras, ) class TestRuntimeBridgeCases: @staticmethod @pytest.fixture - def service(monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session]) -> BridgeCase: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - target = ApplicationTracingService(ApplicationTraceRepository(sqlite_session_factory)) - case = BridgeCase(target) - - def prepare(**values: object) -> ApplicationTracingContext: - if case.failure is not None: - raise case.failure - case.calls.append(("prepare", values)) - return case.context - - def start(**values: object) -> None: - if case.failure is not None: - raise case.failure - case.calls.append(("start", values)) - - def terminal(_info: BaseTraceInfo) -> bool: - if case.failure is not None: - raise case.failure - return True - - def agent_terminal(**values: object) -> None: - if case.failure is not None: - raise case.failure - case.calls.append(("terminal", values)) - - monkeypatch.setattr(target, "prepare_invocation", prepare) - monkeypatch.setattr(target, "start", start) - monkeypatch.setattr(target, "capture_terminal", terminal) - monkeypatch.setattr(target, "capture_agent_app_terminal", agent_terminal) - monkeypatch.setattr(runtime_bridge, "_service", lambda: target) - return case + def bridge(sqlite_session_factory: sessionmaker[Session]) -> Bridge: + """An Agent app with one message and a workflow app with an Agent node, both tracing their content.""" + with sqlite_session_factory.begin() as session: + for app_id, mode in (("agent-app", AppMode.AGENT), ("workflow-app", AppMode.WORKFLOW)): + session.add( + make_app( + app_id=app_id, + tenant_id="tenant", + mode=mode, + tracing=json.dumps({"enabled": True, "tracing_provider": "langfuse", "capture_content": True}), + ) + ) + session.add( + TraceAppConfig( + app_id=app_id, + tracing_provider="langfuse", + tracing_config=langfuse_config("tenant", "https://langfuse.test"), + ) + ) + session.add( + make_workflow(workflow_id="workflow", tenant_id="tenant", app_id="workflow-app", graph=AGENT_GRAPH) + ) + session.add( + make_conversation( + conversation_id="conversation", + app_id="agent-app", + mode=AppMode.AGENT, + inputs={}, + from_source=ConversationFromSource.CONSOLE, + ) + ) + session.add( + make_message( + message_id="message-id", + app_id="agent-app", + conversation_id="conversation", + inputs={}, + query="query", + message={}, + answer="", + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.CONSOLE, + ) + ) + traces = ApplicationTraceRepository(sqlite_session_factory) + return Bridge( + sqlite_session_factory, traces, bridge_hooks(sqlite_session_factory, traces, AGENT_TRACE_SETTINGS) + ) @staticmethod - def test_disabled_gate_never_constructs_database_service_or_reads_context(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) - calls = [] - - def forbidden() -> ApplicationTracingService: - calls.append(True) - raise AssertionError("disabled tracing must not construct a service") - - monkeypatch.setattr(runtime_bridge, "_service", forbidden) - assert runtime_bridge.prepare_agent_app(cast(AgentAppRuntimeBuildContext, object())) is None - assert runtime_bridge.prepare_workflow_agent(cast(WorkflowAgentRuntimeBuildContext, object()), None) is None - runtime_bridge.freeze_workflow_run( - cast(WorkflowAppGenerateEntity, object()), + def test_disabled_deployment_binds_no_hooks_and_never_freezes_a_run(bridge: Bridge) -> None: + # The composition root binds no hooks when the deployment does not collect Agent traces. + application = Flask(__name__) + bind_tracing(application, scopes=trace_scopes(bridge.sessions), hooks=None) + with application.app_context(): + assert application_tracing_hooks() is None + disabled = bridge_hooks(bridge.sessions, bridge.traces, replace(AGENT_TRACE_SETTINGS, enabled=False)) + assert disabled.prepare_agent_app(agent_context()) is None + assert disabled.prepare_workflow_agent(workflow_agent_context(False), None) is None + disabled.freeze_workflow_run( + workflow_entity({}), workflow_id="workflow", workflow_run_id="run", conversation_id=None, - graph_data={}, + graph_data=AGENT_GRAPH, ) - assert runtime_bridge.capture_terminal(cast(BaseTraceInfo, object())) is False - assert runtime_bridge.queue_enabled_at_start("app") is False - assert calls == [] + assert run_count(bridge.sessions) == 0 @staticmethod - def test_agent_app_uses_message_identity_and_stable_conversation(service: BridgeCase) -> None: - assert runtime_bridge.prepare_agent_app(agent_context()) is service.context - assert service.calls == [ - ( - "prepare", - { - "tenant_id": "tenant", - "app_id": "agent-app", - "agent_id": "agent", - "source_kind": "message", - "source_id": "message-id", - "session_id": "conversation", - "node_execution_id": None, - "attempt": 0, - "call_key": "message-id", - }, - ) - ] + def test_agent_app_uses_message_identity_and_stable_conversation(bridge: Bridge) -> None: + runtime = bridge.hooks.prepare_agent_app(agent_context()) + run = bridge.run(app_id="agent-app", source_kind="message", source_id="message-id") + assert run.session_id == "conversation" + [invocation] = bridge.traces.invocations(run.id) + assert (invocation.agent_id, invocation.node_execution_id, invocation.attempt) == ("agent", None, 0) + # The message is the call key, so a redelivered request reuses its invocation. + assert invocation.id == str(uuid5(UUID(run.id), "message-id")) + assert runtime == ApplicationTracingContext( + run_id=UUID(run.id), + invocation_id=UUID(invocation.id), + trace_id=run.trace_id, + parent_span_id=invocation.span_id, + session_id="conversation", + capture_content=True, + ) @staticmethod @pytest.mark.parametrize("preview", [False, True]) - def test_workflow_invocation_uses_owner_app_and_execution_attempt(service: BridgeCase, preview: bool) -> None: - assert runtime_bridge.prepare_workflow_agent(_runtime_bridge_workflow_context(preview), None) is service.context - values = service.calls[0][1] - assert values["app_id"] == "workflow-app" - assert values["agent_id"] == "referenced-agent" - assert values["source_kind"] == ("node" if preview else "workflow") - assert values["source_id"] == ("node-execution" if preview else "run") - assert values["session_id"] == values["source_id"] - assert values["call_key"] == "node-execution:2" - assert values["workflow_id"] == "workflow" + def test_workflow_invocation_uses_owner_app_and_execution_attempt(bridge: Bridge, preview: bool) -> None: + # Neither source row exists yet: only the workflow the hook passes lets the ledger reserve the run. + runtime = bridge.hooks.prepare_workflow_agent(workflow_agent_context(preview), None) + source_kind: TraceSourceKind = "node" if preview else "workflow" + source_id = "node-execution" if preview else "run" + run = bridge.run(app_id="workflow-app", source_kind=source_kind, source_id=source_id) + assert run.session_id == source_id + [invocation] = bridge.traces.invocations(run.id) + assert (invocation.agent_id, invocation.node_execution_id, invocation.attempt) == ( + "referenced-agent", + "node-execution", + 2, + ) + assert invocation.id == str(uuid5(UUID(run.id), "node-execution:2")) + assert invocation.parent_span_id == span_id_for("node-execution") + assert runtime is not None + assert (runtime.run_id, runtime.invocation_id) == (UUID(run.id), UUID(invocation.id)) @staticmethod - def test_workflow_start_prefers_explicit_trace_session(service: BridgeCase) -> None: - entity = _runtime_bridge_workflow_entity({"trace_session_id": "external-session"}) - runtime_bridge.freeze_workflow_run( - entity, + def test_workflow_start_prefers_explicit_trace_session(bridge: Bridge) -> None: + bridge.hooks.freeze_workflow_run( + workflow_entity({"trace_session_id": "external-session"}), workflow_id="workflow", workflow_run_id="run", conversation_id="conversation", graph_data={"nodes": [{"data": {"type": "agent", "version": "2", "agent_node_kind": "dify_agent"}}]}, ) - assert service.calls == [ - ( - "start", - { - "tenant_id": "tenant", - "app_id": "app", - "source_kind": "workflow", - "source_id": "run", - "session_id": "external-session", - "workflow_id": "workflow", - }, - ) - ] + run = bridge.run(app_id="workflow-app", source_kind="workflow", source_id="run") + assert run.session_id == "external-session" + assert run.state is TraceRunState.COLLECTING @staticmethod def test_business_hook_failure_logs_type_without_sensitive_exception_message( - service: BridgeCase, caplog: pytest.LogCaptureFixture + bridge: Bridge, caplog: pytest.LogCaptureFixture ) -> None: - service.failure = ValueError("password=private-secret") - assert runtime_bridge.prepare_agent_app(agent_context()) is None + failing = _FreezeFails(bridge.sessions, ValueError("password=private-secret")) + hooks = bridge_hooks(bridge.sessions, failing, AGENT_TRACE_SETTINGS) + assert hooks.prepare_agent_app(agent_context()) is None assert "ValueError" in caplog.text assert "private-secret" not in caplog.text @staticmethod - def test_worker_database_failure_retries_without_legacy_fallback(service: BridgeCase) -> None: - service.failure = OperationalError("statement", {}, Exception("sensitive database error")) - with pytest.raises(RetryableTraceDispatchError, match="persistence temporarily unavailable"): - runtime_bridge.capture_terminal(cast(BaseTraceInfo, object())) - - @staticmethod - def test_ordinary_workflow_does_not_create_managed_snapshot(service: BridgeCase) -> None: - runtime_bridge.freeze_workflow_run( - cast(WorkflowAppGenerateEntity, object()), + def test_ordinary_workflow_does_not_create_managed_snapshot(bridge: Bridge) -> None: + bridge.hooks.freeze_workflow_run( + workflow_entity({}), workflow_id="workflow", workflow_run_id="run", conversation_id=None, graph_data={"nodes": [{"data": {"type": "llm"}}, {"data": {"type": "agent", "version": "1"}}]}, ) - assert service.calls == [] + assert run_count(bridge.sessions) == 0 @staticmethod @pytest.mark.parametrize( "graph", ["not-json", "null", "[]", {}, {"nodes": {}}, {"nodes": [None, {"data": {"type": "llm"}}]}] ) def test_invalid_or_ordinary_graph_cannot_enable_agent_collection(graph: object) -> None: - assert runtime_bridge.workflow_has_agent_v2(graph) is False + assert workflow_has_agent_v2(graph) is False @staticmethod def test_graph_string_with_agent_is_recognized() -> None: assert ( - runtime_bridge.workflow_has_agent_v2( - '{"nodes":[{"data":{"type":"agent","version":"2","agent_node_kind":"dify_agent"}}]}' - ) + workflow_has_agent_v2('{"nodes":[{"data":{"type":"agent","version":"2","agent_node_kind":"dify_agent"}}]}') is True ) @staticmethod @pytest.mark.parametrize("outcome", ["succeeded", "failed", "cancelled"]) - def test_worker_terminal_handoff_preserves_outcome_without_sse(service: BridgeCase, outcome: str) -> None: - from typing import Literal, cast - - runtime_bridge.capture_agent_app_terminal( + def test_worker_terminal_handoff_preserves_outcome_without_sse( + bridge: Bridge, outcome: Literal["succeeded", "failed", "cancelled"] + ) -> None: + started = tracing_service(bridge.sessions, bridge.traces, no_egress(), [], AGENT_TRACE_SETTINGS).start( tenant_id="tenant", app_id="agent-app", - message_id="message", + source_kind="message", + source_id="message-id", + session_id="conversation", + workflow_id=None, + ) + assert started is not None + bridge.hooks.capture_agent_app_terminal( + tenant_id="tenant", + app_id="agent-app", + message_id="message-id", output="answer", - outcome=cast(Literal["succeeded", "failed", "cancelled"], outcome), + outcome=outcome, ) - assert service.calls == [ - ( - "terminal", - { - "tenant_id": "tenant", - "app_id": "agent-app", - "message_id": "message", - "outcome": outcome, - "output": "answer", - }, - ) - ] + stored = bridge.traces.get(started.id) + assert stored is not None + root = Trace.model_validate(stored.outer_trace).root + assert root.attributes["message_execution_state"] == outcome + assert root.status is (SpanStatus.OK if outcome == "succeeded" else SpanStatus.ERROR) + assert root.outputs == "answer" @staticmethod - def test_terminal_handoff_gate_and_fail_open_log_redact_sensitive_details( - service: BridgeCase, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + def test_terminal_handoff_fail_open_log_redacts_sensitive_details( + bridge: Bridge, caplog: pytest.LogCaptureFixture ) -> None: - from repositories.application_trace_repository import TraceCollectionError - - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) - runtime_bridge.capture_agent_app_terminal( - tenant_id="tenant", app_id="app", message_id="message", outcome="failed", output=None - ) - assert service.calls == [] - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - service.failure = TraceCollectionError("safe_error_code") - runtime_bridge.capture_agent_app_terminal( - tenant_id="tenant", app_id="app", message_id="message", outcome="failed", output="private-content" + bridge.hooks.capture_agent_app_terminal( + tenant_id="tenant", + app_id="agent-app", + message_id="missing-message", + outcome="failed", + output="private-content", ) - assert "safe_error_code" in caplog.text + assert "tracing_source_not_found" in caplog.text assert "private-content" not in caplog.text @staticmethod def test_workflow_fail_open_logs_error_kind_without_content( - service: BridgeCase, caplog: pytest.LogCaptureFixture + bridge: Bridge, caplog: pytest.LogCaptureFixture ) -> None: - service.failure = ValueError("private-prompt") - assert runtime_bridge.prepare_workflow_agent(_runtime_bridge_workflow_context(False), "conversation") is None - runtime_bridge.freeze_workflow_run( - _runtime_bridge_workflow_entity(), + failing = _FreezeFails(bridge.sessions, ValueError("private-prompt")) + hooks = bridge_hooks(bridge.sessions, failing, AGENT_TRACE_SETTINGS) + assert hooks.prepare_workflow_agent(workflow_agent_context(False), "conversation") is None + hooks.freeze_workflow_run( + workflow_entity({}), workflow_id="workflow", workflow_run_id="run", conversation_id="conversation", @@ -818,51 +1153,27 @@ def test_workflow_fail_open_logs_error_kind_without_content( assert "private-prompt" not in caplog.text @staticmethod - @pytest.mark.parametrize( - ("policy", "expected"), - [ - (None, False), - ('{"enabled":true}', True), - ('{"enabled":"true"}', False), - ("[]", False), - ("invalid-private", False), - ], - ) - def test_queue_policy_reads_real_app_rows_without_constructing_vendor_sdk( - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - policy: str | None, - expected: bool, - ) -> None: - from extensions.ext_database import db - from tests.helpers.unit.model_factories import make_app - - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - with sqlite_session_factory.begin() as session: - session.add(make_app(app_id="queue-app", tenant_id="tenant", tracing=policy)) - with sqlite_session_factory() as session: - monkeypatch.setattr(db, "session", session) - assert runtime_bridge.queue_enabled_at_start("queue-app") is expected - assert runtime_bridge.queue_enabled_at_start("missing") is False - assert runtime_bridge.queue_enabled_at_start(None) is False - assert "invalid-private" not in caplog.text + def test_hooks_expose_only_the_agent_collection_lifecycle(bridge: Bridge) -> None: + assert {name for name in dir(bridge.hooks) if not name.startswith("_")} == { + "prepare_agent_app", + "prepare_workflow_agent", + "freeze_workflow_run", + "capture_agent_app_terminal", + } + application = Flask(__name__) + bind_tracing(application, scopes=trace_scopes(bridge.sessions), hooks=bridge.hooks) + with application.app_context(): + assert application_tracing_hooks() is bridge.hooks """Hidden Workflow Tool Agents must not invent a trace owned by the source app.""" -_workflow_tool_tracing_boundary_workflow_context = _runtime_bridge_workflow_context -_workflow_tool_tracing_boundary_workflow_entity = _runtime_bridge_workflow_entity - - def _workflow_tool_tracing_boundary_seed( - sessions: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - *, - root_has_agent: bool, - source_enabled: bool = True, -) -> tuple[ApplicationTracingService, WorkflowAgentRuntimeBuildContext, dict[str, str]]: + sessions: sessionmaker[Session], *, root_has_agent: bool, source_enabled: bool +) -> tuple[ + ApplicationTraceRepository, RuntimeApplicationTracingHooks, WorkflowAgentRuntimeBuildContext, dict[str, str] +]: ids = { name: identifier("tool-boundary-" + name) for name in [ @@ -890,6 +1201,14 @@ def _workflow_tool_tracing_boundary_seed( } ) now = datetime(2026, 10, 21, tzinfo=UTC) + configs = { + "root": langfuse_config(ids["tenant"], "https://tracing.test/root"), + "source": { + "api_key": WorkspaceBoundCipher().encrypt(ids["tenant"], "secret"), + "project": "source", + "endpoint": "https://tracing.test/source", + }, + } with sessions.begin() as session: for app, provider, enabled, content in [ ("root", "langfuse", True, False), @@ -911,20 +1230,14 @@ def _workflow_tool_tracing_boundary_seed( graph=root_graph if app == "root" else {"nodes": [], "edges": []}, ) ) - session.add( - TraceAppConfig( - app_id=ids[app], - tracing_provider=provider, - tracing_config={"project_url": "https://tracing.test/" + app}, - ) - ) + session.add(TraceAppConfig(app_id=ids[app], tracing_provider=provider, tracing_config=configs[app])) workflow_record = WorkflowRun( id=ids["run"], tenant_id=ids["tenant"], app_id=ids["root"], workflow_id=ids["root_workflow"], - type=_application_trace_handoff_WorkflowType.WORKFLOW, - triggered_from=_workflow_tool_tracing_boundary_WorkflowRunTriggeredFrom.APP_RUN, + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, version="1", graph=json.dumps(root_graph), inputs="{}", @@ -968,21 +1281,20 @@ def _workflow_tool_tracing_boundary_seed( ) node_record.created_at = now session.add(node_record) - service = ApplicationTracingService(ApplicationTraceRepository(sessions)) - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - monkeypatch.setattr(runtime_bridge, "_service", lambda: service) - entity = _workflow_tool_tracing_boundary_workflow_entity() + traces = ApplicationTraceRepository(sessions) + hooks = bridge_hooks(sessions, traces, AGENT_TRACE_SETTINGS) + entity = workflow_entity({}) entity.app_config = entity.app_config.model_copy( update={"app_id": ids["root"], "tenant_id": ids["tenant"], "workflow_id": ids["root_workflow"]} ) - runtime_bridge.freeze_workflow_run( + hooks.freeze_workflow_run( entity, workflow_id=ids["root_workflow"], workflow_run_id=ids["run"], conversation_id=None, graph_data=root_graph, ) - original = _workflow_tool_tracing_boundary_workflow_context(False) + original = workflow_agent_context(False) context = replace( original, dify_context=original.dify_context.model_copy( @@ -992,7 +1304,7 @@ def _workflow_tool_tracing_boundary_seed( workflow_run_id=ids["run"], node_execution_id=ids["child"], ) - return service, context, ids + return traces, hooks, context, ids class TestWorkflowToolTracingBoundaryCases: @@ -1000,38 +1312,47 @@ class TestWorkflowToolTracingBoundaryCases: @pytest.mark.parametrize("root_has_agent", [False, True]) @pytest.mark.parametrize("source_enabled", [False, True]) def test_hidden_child_never_reserves_source_policy_or_changes_root_freeze( + sqlite_engine: Engine, sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, root_has_agent: bool, source_enabled: bool, ) -> None: - service, context, ids = _workflow_tool_tracing_boundary_seed( - sqlite_session_factory, monkeypatch, root_has_agent=root_has_agent, source_enabled=source_enabled + traces, hooks, context, ids = _workflow_tool_tracing_boundary_seed( + sqlite_session_factory, root_has_agent=root_has_agent, source_enabled=source_enabled ) with sqlite_session_factory() as session: before = deepcopy([dict(row) for row in session.execute(select(ApplicationTraceRun.__table__)).mappings()]) - service_lookups: list[bool] = [] - - def lookup() -> ApplicationTracingService: - service_lookups.append(True) - return service + statements: list[str] = [] + + def executed( + _connection: object, + _cursor: object, + statement: str, + _parameters: object, + _context: object, + _executemany: object, + ) -> None: + statements.append(statement) - monkeypatch.setattr(runtime_bridge, "_service", lookup) # Marked children never consult policy or ownership, even for malformed # empty markers or a foreign tenant. They cannot create an orphan B run. - for tenant in (ids["tenant"], identifier("foreign-tool-tenant")): - for marker in ("tool-call", ""): - child = replace( - context, - dify_context=context.dify_context.model_copy( - update={ - "tenant_id": tenant, - "workflow_tool_invocation_id": marker, - } - ), - ) - assert runtime_bridge.prepare_workflow_agent(child, None) is None - assert service_lookups == [] + event.listen(sqlite_engine, "before_cursor_execute", executed) + try: + for tenant in (ids["tenant"], identifier("foreign-tool-tenant")): + for marker in ("tool-call", ""): + child = replace( + context, + dify_context=context.dify_context.model_copy( + update={ + "tenant_id": tenant, + "workflow_tool_invocation_id": marker, + } + ), + ) + assert hooks.prepare_workflow_agent(child, None) is None + finally: + event.remove(sqlite_engine, "before_cursor_execute", executed) + assert statements == [] with sqlite_session_factory() as session: observed = [ (r.app_id, r.provider, r.capture_content, r.source_id) @@ -1041,7 +1362,7 @@ def lookup() -> ApplicationTracingService: after = [dict(row) for row in session.execute(select(ApplicationTraceRun.__table__)).mappings()] assert after == before assert ( - service.repository.source_run( + traces.source_run( tenant_id=ids["tenant"], app_id=ids["source"], source_kind="workflow", source_id=ids["run"] ) is None @@ -1050,12 +1371,11 @@ def lookup() -> ApplicationTracingService: @staticmethod def test_hidden_child_skip_keeps_normal_root_agent_trace_assemblable( sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, ) -> None: - service, child, ids = _workflow_tool_tracing_boundary_seed( - sqlite_session_factory, monkeypatch, root_has_agent=True + traces, hooks, child, ids = _workflow_tool_tracing_boundary_seed( + sqlite_session_factory, root_has_agent=True, source_enabled=True ) - assert runtime_bridge.prepare_workflow_agent(child, None) is None + assert hooks.prepare_workflow_agent(child, None) is None root_context = replace( child, dify_context=child.dify_context.model_copy( @@ -1064,9 +1384,9 @@ def test_hidden_child_skip_keeps_normal_root_agent_trace_assemblable( workflow_id=ids["root_workflow"], node_execution_id=ids["root_agent"], ) - runtime = runtime_bridge.prepare_workflow_agent(root_context, None) + runtime = hooks.prepare_workflow_agent(root_context, None) assert runtime is not None - run = service.repository.get(str(runtime.run_id)) + run = traces.get(str(runtime.run_id)) assert run is not None assert run.app_id == ids["root"] frozen = deepcopy(run.provider_config) @@ -1076,6 +1396,7 @@ def test_hidden_child_skip_keeps_normal_root_agent_trace_assemblable( owner.tracing = json.dumps({"enabled": True, "tracing_provider": "langsmith", "capture_content": True}) workflow = session.get(WorkflowRun, ids["run"]) assert workflow is not None + assert workflow.finished_at is not None actor = make_account(account_id=ids["actor"]) nodes = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=sqlite_session_factory, @@ -1085,7 +1406,10 @@ def test_hidden_child_skip_keeps_normal_root_agent_trace_assemblable( triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, ) info = WorkflowTraceInfo( - workflow_data=workflow, + node_executions=nodes.get_by_workflow_execution(ids["run"]), + conversation_id=None, + message_started_at=None, + workflow_app_log_id=None, workflow_id=ids["root_workflow"], tenant_id=ids["tenant"], workflow_run_id=ids["run"], @@ -1094,22 +1418,54 @@ def test_hidden_child_skip_keeps_normal_root_agent_trace_assemblable( workflow_run_inputs={}, workflow_run_outputs={}, workflow_run_version="1", + error=None, total_tokens=0, + prompt_tokens=None, + completion_tokens=None, file_list=[], + invoked_by=None, query="", metadata={"app_id": ids["root"]}, + start_time=workflow.created_at, + end_time=workflow.finished_at, ) - outer = CanonicalTraceBuilder(lambda _info: nodes.get_by_workflow_execution(ids["run"])).build(info) + outer = CanonicalTraceBuilder().build(info) assert outer is not None - outer = remap_outer_trace(outer, trace_id=run.trace_id, capture_content=False, agent_message=False) - service.repository.terminal(run.id, outer.model_dump(mode="json"), outer.root_span_id) - invocation = service.repository.collected(run.id)[0][0] - ingest(Collection(service.repository, sqlite_session_factory, run), batch_for(run, invocation)) - current = service.repository.get(run.id) + capture(traces, run.id, outer) + [invocation] = traces.invocations(run.id) + ingest( + traces, + AgentTraceBatch( + run_id=UUID(run.id), + invocation_id=UUID(invocation.id), + batch_id=uuid4(), + final=True, + spans=[ + AgentTraceSpan( + trace_id=run.trace_id, + span_id="a" * 16, + parent_span_id=invocation.span_id, + name="Agent execution", + start_time_ns=1_000_000_000, + end_time_ns=2_000_000_000, + status="ok", + attributes={"gen_ai.operation.name": "invoke_agent"}, + ) + ], + ), + max_spans=100, + max_bytes=1_000_000, + ) + destination = exported() + tracing_service(sqlite_session_factory, traces, destination, [], AGENT_TRACE_SETTINGS).deliver(run.id) + current = traces.get(run.id) assert current is not None - trace = service.assemble(current) - span_ids = {span.id for span in trace.spans} - assert all(span.parent_id is None or span.parent_id in span_ids for span in trace.spans) + # Assembly rejects any span whose parent is missing, so acceptance proves a complete tree. + assert current.state is TraceRunState.ACCEPTED + assert [request.url for request in destination.requests] == ["https://tracing.test/root/api/public/ingestion"] + sent = observations(destination) + emitted = {body["id"] for body in sent} + assert all(body["parentObservationId"] in emitted for body in sent if "parentObservationId" in body) assert current.provider_config == frozen assert current.provider == "langfuse" assert current.capture_content is False diff --git a/api/tests/unit_tests/services/tracing/test_scopes.py b/api/tests/unit_tests/services/tracing/test_scopes.py new file mode 100644 index 00000000000000..f269375ea3c8b0 --- /dev/null +++ b/api/tests/unit_tests/services/tracing/test_scopes.py @@ -0,0 +1,103 @@ +from typing import override + +import pytest +from dify_trace_sdk.events import TraceCorrelation, TraceEvent + +from services.app_tracing_config_service import AppTraceSettings +from services.tracing.scopes import TraceScopes, TracingPolicyQuery + +CORRELATION = TraceCorrelation( + external_trace_id="caller-trace", session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None +) +ENABLED = AppTraceSettings(enabled=True, tracing_provider="langfuse") + + +class _Recorder: + def __init__(self) -> None: + self.events: list[TraceEvent] = [] + + def emit(self, event: TraceEvent) -> None: + self.events.append(event) + + +class _Policies(TracingPolicyQuery): + """Stored policies by (tenant, app); records the owner of every lookup.""" + + def __init__(self, policies: dict[tuple[str, str], AppTraceSettings], *, fails: bool) -> None: + self.policies = policies + self.fails = fails + self.lookups: list[tuple[str, str]] = [] + + @override + def policy(self, *, tenant_id: str, app_id: str) -> AppTraceSettings | None: + self.lookups.append((tenant_id, app_id)) + if self.fails: + raise RuntimeError("database unavailable") + return self.policies.get((tenant_id, app_id)) + + +def _scopes(policies: _Policies, *, enterprise: bool) -> tuple[TraceScopes, _Recorder]: + recorder = _Recorder() + return TraceScopes(emitter=recorder, policies=policies, enterprise_enabled=lambda: enterprise), recorder + + +@pytest.mark.parametrize( + "policy", + [ + AppTraceSettings(), + AppTraceSettings(enabled=False, tracing_provider="langfuse"), + AppTraceSettings(enabled=True, tracing_provider=None), + ], + ids=["unset", "disabled", "no-provider"], +) +def test_no_scope_exists_when_nothing_consumes_the_events(policy: AppTraceSettings) -> None: + scopes, _ = _scopes(_Policies({}, fails=False), enterprise=False) + + assert ( + scopes.for_execution(tenant_id="tenant-1", app_id="app-1", policy=policy, user_id=None, correlation=CORRELATION) + is None + ) + + +@pytest.mark.parametrize(("policy", "enterprise"), [(ENABLED, False), (AppTraceSettings(), True)]) +def test_scopes_stamp_owner_and_correlation(policy: AppTraceSettings, enterprise: bool) -> None: + scopes, recorder = _scopes(_Policies({}, fails=False), enterprise=enterprise) + + scope = scopes.for_execution( + tenant_id="tenant-1", app_id="app-1", policy=policy, user_id="end-user-session", correlation=CORRELATION + ) + + assert scope is not None + assert (scope.tenant_id, scope.app_id, scope.user_id) == ("tenant-1", "app-1", "end-user-session") + assert scope.correlation == CORRELATION + assert scope.emitter is recorder + + +def test_app_scopes_read_the_policy_through_the_callers_workspace() -> None: + policies = _Policies({("tenant-1", "app-1"): ENABLED, ("tenant-2", "app-1"): ENABLED}, fails=False) + scopes, _ = _scopes(policies, enterprise=False) + + owned = scopes.for_app(tenant_id="tenant-1", app_id="app-1", user_id=None, correlation=CORRELATION) + foreign = scopes.for_app(tenant_id="tenant-3", app_id="app-1", user_id=None, correlation=CORRELATION) + + assert owned is not None + assert owned.tenant_id == "tenant-1" + assert foreign is None + assert policies.lookups == [("tenant-1", "app-1"), ("tenant-3", "app-1")] + + +def test_app_scopes_without_an_app_depend_on_enterprise_telemetry() -> None: + policies = _Policies({}, fails=False) + scopes, _ = _scopes(policies, enterprise=True) + + scope = scopes.for_app(tenant_id="tenant-1", app_id=None, user_id=None, correlation=CORRELATION) + + assert scope is not None + assert scope.app_id is None + assert policies.lookups == [] + + +def test_a_failed_policy_lookup_stays_out_of_the_request() -> None: + scopes, _ = _scopes(_Policies({}, fails=True), enterprise=True) + + assert scopes.for_app(tenant_id="tenant-1", app_id="app-1", user_id=None, correlation=CORRELATION) is None diff --git a/api/tests/unit_tests/services/tracing/test_service.py b/api/tests/unit_tests/services/tracing/test_service.py index 9bcc701f97be8b..aa30f07f9b7678 100644 --- a/api/tests/unit_tests/services/tracing/test_service.py +++ b/api/tests/unit_tests/services/tracing/test_service.py @@ -1,29 +1,74 @@ import json -from dataclasses import dataclass +import threading +from collections.abc import Callable, Generator, Iterator, Mapping, Sequence +from contextlib import contextmanager +from dataclasses import dataclass, field, replace from datetime import UTC, datetime, timedelta from decimal import Decimal +from typing import NoReturn, override from uuid import NAMESPACE_URL, UUID, uuid4, uuid5 import pytest +from aliyun.log.logclient import LogClient from dify_agent.protocol.tracing import AgentTraceBatch, AgentTraceEvent, AgentTraceSpan -from sqlalchemy import func, select +from dify_trace_langfuse.config import LangfuseConfig +from dify_trace_sdk.egress import Egress, EgressResponse +from dify_trace_sdk.errors import EgressError, ExportOutcomeUnknownError, ExportRejectedError, ExportRetryableError +from dify_trace_sdk.events import MessageCompleted, TraceCorrelation, TraceEvent, TracePayload +from dify_trace_sdk.exporting import send_export_request +from dify_trace_sdk.model import Span, SpanKind, SpanStatus, Trace +from dify_trace_sdk.provider import ProviderCapabilities, SendResult, TraceProvider +from dify_trace_sdk.testing import RecordingEgress, response +from sqlalchemy import Connection, Engine, event, func, select from sqlalchemy.orm import Session, sessionmaker - -from core.ops.entities.trace_entity import MessageTraceInfo -from core.ops.exceptions import RetryableTraceDispatchError -from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace +from typing_extensions import TypedDict + +from core.helper.encrypter import obfuscated_token +from domain.tracing.errors import TraceRunConflictError +from domain.tracing.trace_run_agg import TraceInvocation, TraceRunAgg, TraceRunState, TraceSourceKind +from extensions.logstore.aliyun_logstore import AliyunLogStore +from extensions.logstore.repositories.logstore_api_workflow_node_execution_repository import ( + _dict_to_workflow_node_execution_model, +) +from extensions.logstore.repositories.logstore_api_workflow_run_repository import _dict_to_workflow_run from graphon.enums import WorkflowNodeExecutionStatus from libs.datetime_utils import naive_utc_now -from models.application_trace import ApplicationTraceBatch, ApplicationTraceInvocation, ApplicationTraceRun +from machinery.context import RequestContext +from models.account import Tenant, TenantAccountJoin, TenantAccountRole +from models.application_trace import ApplicationTraceBatch, ApplicationTraceRun from models.enums import ConversationFromSource, CreatorUserRole from models.model import App, AppMode, Message, TraceAppConfig -from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom -from repositories.application_trace_repository import ApplicationTraceRepository, TraceCollectionError +from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom, WorkflowRun +from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository +from repositories.application_trace_repository import ApplicationTraceRepository +from repositories.factory import DifyAPIRepositoryFactory +from repositories.trace_record_repository import TraceRecordRepository +from services.app_tracing_config_service import AppTracingConfigService from services.tracing import service as service_module -from services.tracing.provider_dispatch import TraceExportResult, TraceExportUncertainError +from services.tracing.canonical import CanonicalTraceBuilder +from services.tracing.collection import AgentTraceSettings, TraceRun, TraceRunStatus +from services.tracing.collection_ports import LockedTraceRun +from services.tracing.errors import ( + TraceCollectionConflictError, + TraceCollectionLimitError, + TraceCollectionNotFoundError, + TracingStoreUnavailableError, +) +from services.tracing.provider_configs import ProviderConfigGateway +from services.tracing.providers import ProviderConfigCodec, TraceProviderCatalog, TracingSecretCipher +from services.tracing.records import GenerateNameTraceInfo, MessageTraceInfo +from services.tracing.registry import EntryPointTraceProviderCatalog from services.tracing.service import ApplicationTracingService from tests.helpers.unit.config_override import apply_config_overrides -from tests.helpers.unit.model_factories import make_app, make_message, make_workflow +from tests.helpers.unit.model_factories import ( + make_account, + make_app, + make_conversation, + make_message, + make_tenant, + make_workflow, +) +from tests.helpers.unit.services.tracing_ports import AGENT_TRACE_SETTINGS, InstalledTraceProviders """Exercise collection transactions and delivery state against real SQLite rows.""" @@ -36,17 +81,177 @@ def _application_trace_collection_identifier(value: str) -> str: class _ApplicationTraceCollectionCollection: repository: ApplicationTraceRepository sessions: sessionmaker[Session] - run: ApplicationTraceRun + run: TraceRun + + @contextmanager + def change(self) -> Generator[TraceRunAgg]: + """The run under its row lock, in its own transaction; it is saved when the block exits without an error.""" + with self.repository.change(self.run.id, tenant_id=self.run.tenant_id) as locked: + assert locked is not None + yield locked.run - def invocation(self, name: str = "first") -> ApplicationTraceInvocation: - return self.repository.invocation( - run_id=self.run.id, + def invocation(self, name: str = "first") -> TraceInvocation: + return self.open_invocation( agent_id=_application_trace_collection_identifier("agent"), node_execution_id=None, attempt=0, call_key=name, ) + def open_invocation( + self, *, agent_id: str, node_execution_id: str | None, attempt: int, call_key: str + ) -> TraceInvocation: + with self.change() as run: + return run.open_invocation( + agent_id=agent_id, node_execution_id=node_execution_id, attempt=attempt, call_key=call_key + ) + + +class ScriptedProvider(TraceProvider[LangfuseConfig]): + """Sends each trace as one JSON request through the egress, whose script decides the outcome.""" + + id = "langfuse" + config_model = LangfuseConfig + capabilities = ProviderCapabilities(native_cost=True, trace_url=True, span_url=True) + + @override + def verify(self, config: LangfuseConfig) -> LangfuseConfig: + return config + + @override + def project_url(self, config: LangfuseConfig) -> str | None: + return f"{config.host}/project/{config.project_key}" + + @override + def send(self, config: LangfuseConfig, trace: Trace) -> SendResult: + reply = send_export_request( + self.egress, + "POST", + f"{config.host}/traces", + destination="Scripted", + headers={"Authorization": f"{config.public_key}:{config.secret_key}"}, + content=trace.model_dump_json().encode(), + sensitive_url=False, + ) + body = reply.json() + assert isinstance(body, dict) + return SendResult( + trace_id=str(body["trace_id"]), trace_url=body.get("trace_url"), span_urls=body.get("span_urls", {}) + ) + + +SCRIPTED_PROVIDERS = InstalledTraceProviders(ScriptedProvider) +"""Agent traces are delivered through ``ScriptedProvider``; no other provider is installed.""" + + +class ForbiddenTraceProviders(TraceProviderCatalog): + """Fails the test when the service resolves any provider.""" + + def __init__(self, reason: str) -> None: + self._reason = reason + + @override + def resolve(self, provider_id: str) -> NoReturn: + pytest.fail(self._reason) + + +class StoredSecretCipher(TracingSecretCipher): + """A stored ``encrypted-x`` secret decrypts to ``plaintext-x``; ``decrypted`` lists the ciphertexts read.""" + + def __init__(self) -> None: + self.decrypted: list[str] = [] + + @override + def encrypt(self, tenant_id: str, value: str) -> str: + return f"encrypted-{value.removeprefix('plaintext-')}" + + @override + def decrypt(self, tenant_id: str, tokens: Sequence[str]) -> list[str]: + self.decrypted.extend(tokens) + return [token.replace("encrypted-", "plaintext-", 1) for token in tokens] + + @override + def mask(self, value: str) -> str: + return obfuscated_token(value) + + +class ForbiddenSecretCipher(StoredSecretCipher): + """Fails the test when the service reads any credential.""" + + def __init__(self, reason: str) -> None: + super().__init__() + self._reason = reason + + @override + def decrypt(self, tenant_id: str, tokens: Sequence[str]) -> NoReturn: + pytest.fail(self._reason) + + +class UnreadableSecretCipher(StoredSecretCipher): + """No stored secret can be decrypted, as when the workspace key is lost.""" + + @override + def decrypt(self, tenant_id: str, tokens: Sequence[str]) -> NoReturn: + raise ValueError("private-ciphertext") + + +def trace_records(sessions: sessionmaker[Session]) -> TraceRecordRepository: + return TraceRecordRepository( + sessions, + workflow_runs=DifyAPIRepositoryFactory.create_api_workflow_run_repository(sessions), + node_executions=DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(sessions), + files_url="https://files.test", + enterprise_details=False, + ) + + +@dataclass +class ServicePorts: + """The service's collaborators besides its repository; a test replaces those it constrains. + + ``scheduled`` lists the runs the service scheduled for export. + """ + + egress: Egress + providers: TraceProviderCatalog = SCRIPTED_PROVIDERS + cipher: TracingSecretCipher = field(default_factory=StoredSecretCipher) + settings: AgentTraceSettings = AGENT_TRACE_SETTINGS + scheduled: list[str] = field(default_factory=list[str]) + + def service(self, case: _ApplicationTraceCollectionCollection) -> ApplicationTracingService: + return self.service_on(case.repository, case.sessions) + + def service_on( + self, repository: ApplicationTraceRepository, sessions: sessionmaker[Session] + ) -> ApplicationTracingService: + """The service whose runs, run queries, and source queries all go through ``repository``.""" + return ApplicationTracingService( + runs=repository, + run_query=repository, + sources=repository, + records=trace_records(sessions), + builder=CanonicalTraceBuilder(), + providers=self.providers, + codec=ProviderConfigCodec(self.cipher), + egress=self.egress, + schedule_export=self.scheduled.append, + settings=self.settings, + ) + + +def _application_trace_collection_accepted(trace_id: str) -> RecordingEgress: + """An egress whose destination acknowledges exactly one trace.""" + return RecordingEgress([response(200, {"trace_id": trace_id})], grpc_target=None) + + +def _application_trace_collection_no_egress() -> RecordingEgress: + """An egress that fails any request the test did not script.""" + return RecordingEgress([], grpc_target=None) + + +def sent_traces(egress: RecordingEgress) -> list[Trace]: + return [Trace.model_validate_json(request.content or b"") for request in egress.requests] + def _application_trace_collection_seed( sessions: sessionmaker[Session], @@ -97,23 +302,88 @@ def _application_trace_collection_seed( app_id=app_id, tracing_provider="langfuse", tracing_config={ - "host": "https://langfuse.test", + "public_key": "encrypted-public", "secret_key": "encrypted-original", - "project_url": "https://langfuse.test/project/original", + "host": "https://langfuse.test", + "project_key": "original", }, ) ) repository = ApplicationTraceRepository(sessions) - run = repository.start( - tenant_id=tenant_id, app_id=app_id, source_kind="message", source_id=message_id, session_id="session-original" + ServicePorts(_application_trace_collection_no_egress()).service_on(repository, sessions).start( + tenant_id=tenant_id, + app_id=app_id, + source_kind="message", + source_id=message_id, + session_id="session-original", + workflow_id=None, ) + # ``start`` returns no run for a source frozen as not collected, so read the frozen run itself. + run = repository.source_run(tenant_id=tenant_id, app_id=app_id, source_kind="message", source_id=message_id) assert run is not None return _ApplicationTraceCollectionCollection(repository, sessions, run) +def attach_conversation(case: _ApplicationTraceCollectionCollection, mode: AppMode) -> None: + """Message records are assembled only for a message whose conversation belongs to its app.""" + with case.sessions.begin() as session: + message = session.get(Message, case.run.source_id) + assert message is not None + session.add( + make_conversation( + conversation_id=message.conversation_id, + app_id=case.run.app_id, + mode=mode, + inputs={}, + from_source=ConversationFromSource.CONSOLE, + ) + ) + + +def add_app_owner(case: _ApplicationTraceCollectionCollection, name: str) -> str: + """Workflow records are read as the app creator, who must be a member of the app's workspace.""" + owner_id = _application_trace_collection_identifier(f"owner-{name}") + with case.sessions.begin() as session: + if session.get(Tenant, case.run.tenant_id) is None: + session.add(make_tenant(tenant_id=case.run.tenant_id)) + session.add(make_account(account_id=owner_id, email=f"{name}@example.com")) + session.add(TenantAccountJoin(tenant_id=case.run.tenant_id, account_id=owner_id, role=TenantAccountRole.OWNER)) + app = session.get(App, case.run.app_id) + assert app is not None + app.created_by = owner_id + return owner_id + + +def trace_event(case: _ApplicationTraceCollectionCollection, payload: TracePayload) -> TraceEvent: + return TraceEvent( + event_id=uuid4().hex, + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, + user_id=None, + occurred_at=datetime.now(UTC), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + payload=payload, + ) + + +def message_record(case: _ApplicationTraceCollectionCollection) -> MessageTraceInfo: + """The record the delivery worker assembles when the message reports its end.""" + with case.sessions() as session: + message = session.get(Message, case.run.source_id) + assert message is not None + conversation_id = message.conversation_id + record = trace_records(case.sessions).assemble( + trace_event(case, MessageCompleted(message_id=case.run.source_id, conversation_id=conversation_id)) + ) + assert isinstance(record, MessageTraceInfo) + return record + + def _application_trace_collection_batch_for( - run: ApplicationTraceRun, - invocation: ApplicationTraceInvocation, + run: TraceRun, + invocation: TraceInvocation, *, final: bool = True, batch_id: UUID | None = None, @@ -141,29 +411,64 @@ def _application_trace_collection_batch_for( def _application_trace_collection_ingest( collection: _ApplicationTraceCollectionCollection, batch: AgentTraceBatch -) -> tuple[int, str]: - return collection.repository.ingest(batch, max_spans=100, max_bytes=1_000_000) +) -> tuple[int, list[str]]: + """Ingest through the service; returns the spans it accepted and the runs it then scheduled for export.""" + ports = ServicePorts(_application_trace_collection_no_egress()) + return ports.service(collection).ingest(batch), ports.scheduled def _application_trace_collection_terminal(collection: _ApplicationTraceCollectionCollection) -> None: run = collection.run - root = CanonicalSpan( - id=run.root_span_id, - parent_id=None, + root = Span( + span_id=run.root_span_id, + parent_span_id=None, name="Agent turn", - kind=CanonicalSpanKind.CHAIN, + kind=SpanKind.CHAIN, start_time=datetime.fromtimestamp(0, UTC), end_time=datetime.fromtimestamp(3, UTC), - status=CanonicalSpanStatus.OK, + status=SpanStatus.OK, + ) + trace = Trace( + trace_id=run.trace_id, + external_trace_id=None, + session_id=run.session_id, + user_id=None, + parent=None, + spans=(root,), ) - trace = CanonicalTrace(trace_id=run.trace_id, session_id=run.session_id, root_span_id=root.id, spans=(root,)) - collection.repository.terminal(run.id, trace.model_dump(mode="json"), root.id) + with collection.change() as current: + current.capture_terminal(trace, now=naive_utc_now()) + + +def _claim_at( + collection: _ApplicationTraceCollectionCollection, *, lease_seconds: int, now: datetime +) -> tuple[TraceRunAgg, str] | None: + """Lease the run for one export attempt; returns the claimed run and the token that fences the attempt.""" + token = str(uuid4()) + with collection.repository.change(collection.run.id, tenant_id=collection.run.tenant_id) as locked: + assert locked is not None + claimed = locked.run.claim_export( + token=token, owner_exists=locked.owner_exists, collection_timeout=60, lease_seconds=lease_seconds, now=now + ) + return (locked.run, token) if claimed else None def _application_trace_collection_claim( collection: _ApplicationTraceCollectionCollection, -) -> tuple[ApplicationTraceRun, str] | None: - return collection.repository.claim(collection.run.id, collection_timeout=60, lease_seconds=30) +) -> tuple[TraceRunAgg, str] | None: + return _claim_at(collection, lease_seconds=30, now=naive_utc_now()) + + +def _accept_export(collection: _ApplicationTraceCollectionCollection, token: str, accepted: SendResult) -> None: + """Record the destination's acceptance of the export attempt ``token``.""" + with collection.change() as run: + run.export_accepted( + token, + provider_trace_id=accepted.trace_id, + trace_url=accepted.trace_url, + span_urls=accepted.span_urls, + now=naive_utc_now(), + ) class TestApplicationTraceCollectionCases: @@ -178,15 +483,13 @@ def test_batch_duplicate_is_idempotent_after_final_and_acceptance( ) -> None: invocation = collection.invocation() batch = _application_trace_collection_batch_for(collection.run, invocation) - assert _application_trace_collection_ingest(collection, batch) == (1, collection.run.id) - assert _application_trace_collection_ingest(collection, batch) == (1, collection.run.id) + assert _application_trace_collection_ingest(collection, batch) == (1, [collection.run.id]) + assert _application_trace_collection_ingest(collection, batch) == (1, [collection.run.id]) _application_trace_collection_terminal(collection) leased = _application_trace_collection_claim(collection) assert leased is not None - collection.repository.complete( - collection.run.id, leased[1], state="accepted", result=TraceExportResult("provider-id") - ) - assert _application_trace_collection_ingest(collection, batch) == (1, collection.run.id) + _accept_export(collection, leased[1], SendResult(trace_id="provider-id", trace_url=None, span_urls={})) + assert _application_trace_collection_ingest(collection, batch) == (1, [collection.run.id]) stored = collection.repository.get(collection.run.id) assert stored is not None assert stored.span_count == 1 @@ -200,18 +503,19 @@ def test_conflicting_batch_id_rolls_back_without_replacing_payload( batch = _application_trace_collection_batch_for(collection.run, collection.invocation(), final=False) _application_trace_collection_ingest(collection, batch) changed = batch.model_copy(update={"final": True}) - with pytest.raises(TraceCollectionError, match="tracing_batch_conflict"): + with pytest.raises(TraceCollectionConflictError, match="tracing_batch_conflict"): _application_trace_collection_ingest(collection, changed) - invocations, batches = collection.repository.collected(collection.run.id) + invocations = collection.repository.invocations(collection.run.id) assert invocations[0].finalized is False - assert batches[0].final is False + with collection.sessions() as session: + assert session.scalar(select(ApplicationTraceBatch.final)) is False @staticmethod def test_run_cannot_ingest_another_runs_invocation(collection: _ApplicationTraceCollectionCollection) -> None: other = _application_trace_collection_seed(collection.sessions, "other") foreign = other.invocation() forged = _application_trace_collection_batch_for(collection.run, foreign) - with pytest.raises(TraceCollectionError, match="tracing_invocation_not_found"): + with pytest.raises(TraceCollectionNotFoundError, match="tracing_invocation_not_found"): _application_trace_collection_ingest(collection, forged) with collection.sessions() as session: assert session.scalar(select(func.count()).select_from(ApplicationTraceBatch)) == 0 @@ -219,11 +523,12 @@ def test_run_cannot_ingest_another_runs_invocation(collection: _ApplicationTrace @staticmethod def test_payload_limits_reject_whole_transaction(collection: _ApplicationTraceCollectionCollection) -> None: batch = _application_trace_collection_batch_for(collection.run, collection.invocation()) - with pytest.raises(TraceCollectionError, match="tracing_payload_limit"): - collection.repository.ingest(batch, max_spans=100, max_bytes=1) - invocations, batches = collection.repository.collected(collection.run.id) - assert not batches - assert invocations[0].finalized is False + limits = replace(AGENT_TRACE_SETTINGS, max_spans=100, max_bytes=1) + limited = ServicePorts(_application_trace_collection_no_egress(), settings=limits).service(collection) + with pytest.raises(TraceCollectionLimitError, match="tracing_payload_limit"): + limited.ingest(batch) + assert not collection.repository.batches(collection.run.id) + assert collection.repository.invocations(collection.run.id)[0].finalized is False stored = collection.repository.get(collection.run.id) assert stored is not None assert stored.span_count == 0 @@ -232,10 +537,9 @@ def test_payload_limits_reject_whole_transaction(collection: _ApplicationTraceCo def test_wrong_trace_id_is_rejected_before_persistence(collection: _ApplicationTraceCollectionCollection) -> None: batch = _application_trace_collection_batch_for(collection.run, collection.invocation()) altered = batch.model_copy(update={"spans": [batch.spans[0].model_copy(update={"trace_id": "f" * 32})]}) - with pytest.raises(TraceCollectionError, match="tracing_trace_id_mismatch"): + with pytest.raises(TraceCollectionConflictError, match="tracing_trace_id_mismatch"): _application_trace_collection_ingest(collection, altered) - _, batches = collection.repository.collected(collection.run.id) - assert not batches + assert not collection.repository.batches(collection.run.id) @staticmethod def test_content_off_cleans_the_durable_batch_not_only_the_export( @@ -263,7 +567,7 @@ def test_content_off_cleans_the_durable_batch_not_only_the_export( ) batch = original.model_copy(update={"spans": [observation]}) _application_trace_collection_ingest(collection, batch) - _, batches = collection.repository.collected(collection.run.id) + batches = collection.repository.batches(collection.run.id) assert "private" not in json.dumps(batches[0].spans) assert batches[0].spans[0]["events"] == [] assert batches[0].spans[0]["status"] == "error" @@ -273,7 +577,7 @@ def test_content_off_cleans_the_durable_batch_not_only_the_export( changed = batch.model_copy( update={"spans": [observation.model_copy(update={"status_message": "different private error"})]} ) - with pytest.raises(TraceCollectionError, match="tracing_batch_conflict"): + with pytest.raises(TraceCollectionConflictError, match="tracing_batch_conflict"): _application_trace_collection_ingest(collection, changed) @staticmethod @@ -283,9 +587,8 @@ def test_invocation_key_replay_must_not_change_agent_binding( original = collection.invocation() repeated = collection.invocation() assert original.id == repeated.id - with pytest.raises(TraceCollectionError, match="tracing_invocation_conflict"): - collection.repository.invocation( - run_id=collection.run.id, + with pytest.raises(TraceRunConflictError, match="tracing_invocation_conflict"): + collection.open_invocation( agent_id=_application_trace_collection_identifier("different-agent"), node_execution_id=_application_trace_collection_identifier("different-node"), attempt=2, @@ -326,14 +629,14 @@ def test_expired_delivery_lease_accepts_confirmed_late_ack_for_same_attempt( assert _application_trace_collection_claim(collection) is None current = collection.repository.get(collection.run.id) assert current is not None - assert current.state == "uncertain" - collection.repository.complete(collection.run.id, leased[1], state="accepted", result=TraceExportResult("late")) - stored = collection.repository.get(collection.run.id) - assert stored is not None - assert stored.state == "accepted" - assert stored.error_code is None - assert stored.accepted_at is not None - assert stored.provider_trace_id == "late" + assert current.state is TraceRunState.UNCERTAIN + _accept_export(collection, leased[1], SendResult(trace_id="late", trace_url=None, span_urls={})) + accepted = collection.repository.get(collection.run.id) + assert accepted is not None + assert accepted.state is TraceRunState.ACCEPTED + assert accepted.error_code is None + assert accepted.accepted_at is not None + assert accepted.provider_trace_id == "late" assert _application_trace_collection_claim(collection) is None @staticmethod @@ -341,13 +644,13 @@ def test_incomplete_final_times_out_without_export(collection: _ApplicationTrace collection.invocation() _application_trace_collection_terminal(collection) with collection.sessions.begin() as session: - stored = session.get(ApplicationTraceRun, collection.run.id) - assert stored is not None - stored.terminal_at = naive_utc_now() - timedelta(seconds=120) + row = session.get(ApplicationTraceRun, collection.run.id) + assert row is not None + row.terminal_at = naive_utc_now() - timedelta(seconds=120) assert _application_trace_collection_claim(collection) is None stored = collection.repository.get(collection.run.id) assert stored is not None - assert stored.state == "failed" + assert stored.state is TraceRunState.FAILED assert stored.error_code == "agent_trace_collection_incomplete" @staticmethod @@ -362,12 +665,14 @@ def test_running_trace_freezes_credentials_content_policy_and_session( config = session.scalar(select(TraceAppConfig).where(TraceAppConfig.app_id == original.app_id)) assert config is not None config.tracing_config = {"secret_key": "encrypted-new"} - repeated = collection.repository.start( + service = ServicePorts(_application_trace_collection_no_egress()).service(collection) + repeated = service.start( tenant_id=original.tenant_id, app_id=original.app_id, source_kind="message", source_id=original.source_id, session_id="changed", + workflow_id=None, ) assert repeated is not None assert repeated.provider == "langfuse" @@ -377,82 +682,96 @@ def test_running_trace_freezes_credentials_content_policy_and_session( assert repeated.session_id == "session-original" @staticmethod - @pytest.mark.parametrize("outcome", ["accepted", "uncertain", "retrying"]) + @pytest.mark.parametrize( + ("outcome", "reply"), + [ + ( + TraceRunState.ACCEPTED, + response(200, {"trace_id": "vendor-trace", "trace_url": "https://trace.test/example"}), + ), + (TraceRunState.UNCERTAIN, EgressError("acknowledgement lost", request_sent=True)), + (TraceRunState.RETRYING, response(503, {})), + ], + ) def test_delivery_checkpoint_controls_replays( - collection: _ApplicationTraceCollectionCollection, monkeypatch: pytest.MonkeyPatch, outcome: str + collection: _ApplicationTraceCollectionCollection, + outcome: TraceRunState, + reply: EgressResponse | Exception, ) -> None: _application_trace_collection_ingest( collection, _application_trace_collection_batch_for(collection.run, collection.invocation()) ) _application_trace_collection_terminal(collection) - exports: list[tuple[str, dict[str, object], CanonicalTrace]] = [] - decryptions: list[dict[str, object]] = [] - - def decrypt(_tenant: str, _provider: str, config: dict[str, object]) -> dict[str, object]: - decryptions.append(config) - return {"secret_key": "plaintext-only-in-worker"} - - def export(provider: str, config: dict[str, object], trace: CanonicalTrace) -> TraceExportResult: - exports.append((provider, config, trace)) - if outcome == "uncertain": - raise TraceExportUncertainError("ack lost") - if outcome == "retrying": - raise RetryableTraceDispatchError("busy") - return TraceExportResult("vendor-trace", "https://trace.test/example") - - monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", decrypt) - monkeypatch.setattr(service_module, "export_trace", export) - service = ApplicationTracingService(collection.repository) + cipher = StoredSecretCipher() + egress = RecordingEgress([reply], grpc_target=None) + service = ServicePorts(egress, cipher=cipher).service(collection) service.deliver(collection.run.id) service.deliver(collection.run.id) - assert len(exports) == 1 - assert decryptions[0]["secret_key"] == "encrypted-original" + assert len(egress.requests) == 1 + request = egress.requests[0] + assert request.url == "https://langfuse.test/traces" + assert request.headers["Authorization"] == "plaintext-public:plaintext-original" + assert cipher.decrypted == ["encrypted-public", "encrypted-original"] stored = collection.repository.get(collection.run.id) assert stored is not None - assert stored.state == outcome + assert stored.state is outcome assert stored.provider_config["secret_key"] == "encrypted-original" - assert "plaintext-only-in-worker" not in json.dumps(stored.outer_trace) - assert exports[0][1]["host"] == "https://langfuse.test" + assert "plaintext" not in json.dumps(stored.outer_trace) @staticmethod def test_source_owner_checks_app_and_tenant(collection: _ApplicationTraceCollectionCollection) -> None: run = collection.run + service = ServicePorts(_application_trace_collection_no_egress()).service(collection) + foreign_tenant = _application_trace_collection_identifier("foreign-tenant") assert ( - collection.repository.start( - tenant_id=_application_trace_collection_identifier("foreign-tenant"), + service.start( + tenant_id=foreign_tenant, app_id=run.app_id, source_kind="message", source_id=run.source_id, session_id="x", + workflow_id=None, + ) + is None + ) + # ``start`` also returns no run for a source frozen as not collected; no run may be frozen at all. + assert ( + collection.repository.source_run( + tenant_id=foreign_tenant, app_id=run.app_id, source_kind="message", source_id=run.source_id ) is None ) other = _application_trace_collection_seed(collection.sessions, "second-app") assert ( - collection.repository.start( + service.start( tenant_id=run.tenant_id, app_id=other.run.app_id, source_kind="message", source_id=run.source_id, session_id="x", + workflow_id=None, + ) + is None + ) + assert ( + collection.repository.source_run( + tenant_id=run.tenant_id, app_id=other.run.app_id, source_kind="message", source_id=run.source_id ) is None ) @staticmethod def test_disabled_source_decision_stays_disabled_after_config_enable( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + sqlite_session_factory: sessionmaker[Session], ) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) case = _application_trace_collection_seed(sqlite_session_factory, "initially-disabled", enabled=False) - service = ApplicationTracingService(case.repository) - assert case.run.state == "not_collected" + attach_conversation(case, AppMode.AGENT) + service = ServicePorts(_application_trace_collection_no_egress()).service(case) + assert case.run.state is TraceRunState.NOT_COLLECTED with case.sessions.begin() as session: app = session.get(App, case.run.app_id) assert app is not None app.tracing = json.dumps({"enabled": True, "tracing_provider": "langfuse", "capture_content": True}) - message = session.get(Message, case.run.source_id) - assert message is not None assert ( service.start( tenant_id=case.run.tenant_id, @@ -460,35 +779,25 @@ def test_disabled_source_decision_stays_disabled_after_config_enable( source_kind="message", source_id=case.run.source_id, session_id="later-session", + workflow_id=None, ) is None ) - info = MessageTraceInfo( - metadata={}, - message_data=message, - message_id=message.id, - conversation_model="agent", - conversation_mode="agent", - message_tokens=1, - answer_tokens=1, - total_tokens=2, - ) - assert service.capture_terminal(info) is True + assert service.capture_terminal(tenant_id=case.run.tenant_id, info=message_record(case)) is True stored = case.repository.get(case.run.id) assert stored is not None - assert stored.state == "not_collected" + assert stored.state is TraceRunState.NOT_COLLECTED assert stored.provider_config == {} assert stored.outer_trace is None @staticmethod def test_capture_agent_disabled_preserves_outer_only_content_policy( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + sqlite_session_factory: sessionmaker[Session], ) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) case = _application_trace_collection_seed(sqlite_session_factory, "outer-only", capture_agent=False) - service = ApplicationTracingService(case.repository) - scheduled: list[str] = [] - monkeypatch.setattr(service, "schedule", scheduled.append) + attach_conversation(case, AppMode.AGENT) + ports = ServicePorts(_application_trace_collection_no_egress()) + service = ports.service(case) assert ( service.prepare_invocation( tenant_id=case.run.tenant_id, @@ -500,53 +809,39 @@ def test_capture_agent_disabled_preserves_outer_only_content_policy( node_execution_id=None, attempt=0, call_key="call", + workflow_id=None, ) is None ) - with case.sessions() as session: - message = session.get(Message, case.run.source_id) - assert message is not None - info = MessageTraceInfo( - metadata={}, - message_data=message, - message_id=message.id, - conversation_model="agent", - conversation_mode="agent", - message_tokens=1, - answer_tokens=1, - total_tokens=2, - inputs="private input", - outputs="private output", - ) - assert service.capture_terminal(info) is True - assert scheduled == [case.run.id] + info = message_record(case).model_copy(update={"inputs": "private input", "outputs": "private output"}) + assert service.capture_terminal(tenant_id=case.run.tenant_id, info=info) is True + assert ports.scheduled == [case.run.id] stored = case.repository.get(case.run.id) assert stored is not None assert stored.outer_trace is not None assert "private" not in json.dumps(stored.outer_trace) + # The message's synthetic model call is left to the Agent runtime, which reports none here. assert len(stored.outer_trace["spans"]) == 1 - invocations, _ = case.repository.collected(case.run.id) - assert invocations == [] + assert case.repository.invocations(case.run.id) == [] @staticmethod def test_deleting_selected_configuration_keeps_historical_destination( collection: _ApplicationTraceCollectionCollection, ) -> None: - from machinery.context import RequestContext - from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository - from services.app_tracing_config_gateway import OpsTraceManagerGateway - from services.app_tracing_config_service import AppTracingConfigService - original = collection.run service = AppTracingConfigService( configs=SQLAlchemyAppTracingConfigRepository(session_factory=collection.sessions), - provider=OpsTraceManagerGateway(), + provider=ProviderConfigGateway( + providers=SCRIPTED_PROVIDERS, + codec=ProviderConfigCodec(StoredSecretCipher()), + egress=_application_trace_collection_no_egress(), + ), ) service.delete(RequestContext("request", None, "account", original.tenant_id), original.app_id, "langfuse") stored = collection.repository.get(original.id) assert stored is not None assert stored.provider_config == original.provider_config - assert stored.provider_config["project_url"] == "https://langfuse.test/project/original" + assert stored.provider_config["project_key"] == "original" assert stored.config_revision == original.config_revision with collection.sessions() as session: assert session.scalar(select(TraceAppConfig).where(TraceAppConfig.app_id == original.app_id)) is None @@ -561,31 +856,74 @@ def test_deleting_selected_configuration_keeps_historical_destination( _ApplicationTraceServiceBoundariesCollection = _ApplicationTraceCollectionCollection +_application_trace_service_boundaries_accepted = _application_trace_collection_accepted _application_trace_service_boundaries_batch_for = _application_trace_collection_batch_for _application_trace_service_boundaries_claim = _application_trace_collection_claim _application_trace_service_boundaries_identifier = _application_trace_collection_identifier _application_trace_service_boundaries_ingest = _application_trace_collection_ingest +_application_trace_service_boundaries_no_egress = _application_trace_collection_no_egress _application_trace_service_boundaries_seed = _application_trace_collection_seed _application_trace_service_boundaries_terminal = _application_trace_collection_terminal -def source(case: _ApplicationTraceServiceBoundariesCollection) -> dict[str, str]: - return { - "tenant_id": case.run.tenant_id, - "app_id": case.run.app_id, - "source_kind": "message", - "source_id": case.run.source_id, - } +FROZEN_LANGFUSE: dict[str, object] = { + "public_key": "encrypted-public", + "secret_key": "encrypted-original", + "host": "https://langfuse.test", + "project_key": "original", +} + + +class RunSource(TypedDict, closed=True): + """The source a caller names to read or retry the status of a run.""" + + tenant_id: str + app_id: str + source_kind: TraceSourceKind + source_id: str + + +def source(run: TraceRun) -> RunSource: + return RunSource(tenant_id=run.tenant_id, app_id=run.app_id, source_kind=run.source_kind, source_id=run.source_id) + + +def not_collected(source_kind: TraceSourceKind, source_id: str) -> TraceRunStatus: + """The status of a source nothing is collected for: it names no destination or delivery detail.""" + return TraceRunStatus( + status=TraceRunState.NOT_COLLECTED, + source_kind=source_kind, + source_id=source_id, + provider=None, + trace_id=None, + trace_url=None, + span_url=None, + project_url=None, + config_revision=None, + capture_content=None, + capture_agent=None, + span_count=None, + attempts=None, + error_code=None, + accepted_at=None, + capabilities=None, + ) + + +def freeze_config(case: _ApplicationTraceServiceBoundariesCollection, provider: str, config: dict[str, object]) -> None: + with case.sessions.begin() as session: + run = session.get(ApplicationTraceRun, case.run.id) + assert run is not None + run.provider = provider + run.provider_config = config class TestApplicationTraceServiceBoundariesCases: @staticmethod def test_prepare_invocation_returns_only_opaque_context_and_respects_capture_gate( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + sqlite_session_factory: sessionmaker[Session], ) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) case = _application_trace_service_boundaries_seed(sqlite_session_factory, "prepare", capture_content=True) - service = ApplicationTracingService(case.repository) + service = ServicePorts(_application_trace_service_boundaries_no_egress()).service(case) context = service.prepare_invocation( tenant_id=case.run.tenant_id, app_id=case.run.app_id, @@ -596,20 +934,24 @@ def test_prepare_invocation_returns_only_opaque_context_and_respects_capture_gat node_execution_id=None, attempt=0, call_key="call", + workflow_id=None, ) assert context is not None assert context.run_id == UUID(case.run.id) assert context.capture_content is True assert context.session_id == "session-original" assert "encrypted-original" not in context.model_dump_json() - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) + disabled = ServicePorts( + _application_trace_service_boundaries_no_egress(), settings=replace(AGENT_TRACE_SETTINGS, enabled=False) + ).service(case) assert ( - service.start( + disabled.start( tenant_id=case.run.tenant_id, app_id=case.run.app_id, source_kind="message", source_id=case.run.source_id, session_id="new", + workflow_id=None, ) is None ) @@ -620,35 +962,30 @@ def test_status_denies_foreign_or_absent_source_without_destination_details( sqlite_session_factory: sessionmaker[Session], missing: str ) -> None: case = _application_trace_service_boundaries_seed(sqlite_session_factory, "scope") - identity = source(case) - identity[{"tenant": "tenant_id", "app": "app_id", "message": "source_id"}[missing]] = ( - _application_trace_service_boundaries_identifier("foreign") - ) - with pytest.raises(TraceCollectionError) as caught: - ApplicationTracingService(case.repository).status(**identity) + foreign = _application_trace_service_boundaries_identifier("foreign") + service = ServicePorts(_application_trace_service_boundaries_no_egress()).service(case) + with pytest.raises(TraceCollectionNotFoundError) as caught: + service.status( + tenant_id=foreign if missing == "tenant" else case.run.tenant_id, + app_id=foreign if missing == "app" else case.run.app_id, + source_kind="message", + source_id=foreign if missing == "message" else case.run.source_id, + node_execution_id=None, + ) assert caught.value.code == "tracing_source_not_found" - assert caught.value.status == 404 assert "encrypted-original" not in str(caught.value) @staticmethod def test_status_uses_frozen_destination_and_selects_latest_node_attempt( sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, ) -> None: case = _application_trace_service_boundaries_seed(sqlite_session_factory, "status") with case.sessions.begin() as session: config = session.scalar(select(TraceAppConfig).where(TraceAppConfig.app_id == case.run.app_id)) assert config is not None - config.tracing_config = {"project_url": "https://new.example/project/new", "secret_key": "encrypted-new"} - - def forbidden(*_args: object, **_kwargs: object) -> None: - pytest.fail("status must not decrypt credentials or resolve current provider metadata") - - monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", forbidden) - monkeypatch.setattr(service_module.OpsTraceManager, "get_trace_config_project_url", forbidden) + config.tracing_config = {**FROZEN_LANGFUSE, "host": "https://new.example", "project_key": "new"} node_id = _application_trace_service_boundaries_identifier("selected-node") - invocation = case.repository.invocation( - run_id=case.run.id, + invocation = case.open_invocation( agent_id=_application_trace_service_boundaries_identifier("agent"), node_execution_id=node_id, attempt=1, @@ -660,30 +997,33 @@ def forbidden(*_args: object, **_kwargs: object) -> None: _application_trace_service_boundaries_terminal(case) lease = _application_trace_service_boundaries_claim(case) assert lease is not None - case.repository.complete( - case.run.id, + _accept_export( + case, lease[1], - state="accepted", - result=TraceExportResult( - "provider-trace", + SendResult( + trace_id="provider-trace", trace_url="https://frozen.example/trace", span_urls={invocation.span_id: "https://frozen.example/span"}, ), ) - service = ApplicationTracingService(case.repository) - result = service.status(**source(case), node_execution_id=node_id) - assert result["status"] == "accepted" - assert result["trace_id"] == "provider-trace" - assert result["span_url"] == "https://frozen.example/span" - assert result["project_url"] == "https://langfuse.test/project/original" - assert result["config_revision"] == case.run.config_revision - assert result["queryable"] is None - assert result["capture_content"] is False - assert "provider_config" not in result + service = ServicePorts( + _application_trace_service_boundaries_no_egress(), + cipher=ForbiddenSecretCipher("status must not decrypt credentials"), + ).service(case) + result = service.status(**source(case.run), node_execution_id=node_id) + assert result.status is TraceRunState.ACCEPTED + assert result.provider == "langfuse" + assert result.trace_id == "provider-trace" + assert result.span_url == "https://frozen.example/span" + assert result.project_url == "https://langfuse.test/project/original" + assert result.config_revision == case.run.config_revision + assert result.capture_content is False + assert result.capabilities is not None + assert result.capabilities.provider == "langfuse" assert "encrypted-original" not in str(result) - with pytest.raises(TraceCollectionError, match="tracing_node_not_found"): + with pytest.raises(TraceCollectionNotFoundError, match="tracing_node_not_found"): service.status( - **source(case), node_execution_id=_application_trace_service_boundaries_identifier("another-node") + **source(case.run), node_execution_id=_application_trace_service_boundaries_identifier("another-node") ) @staticmethod @@ -695,7 +1035,7 @@ def test_new_run_freezes_updated_project_without_changing_previous_run( with case.sessions.begin() as session: config = session.scalar(select(TraceAppConfig).where(TraceAppConfig.app_id == case.run.app_id)) assert config is not None - config.tracing_config = {"project_url": "https://new.example/project/new", "secret_key": "encrypted-new"} + config.tracing_config = {**FROZEN_LANGFUSE, "host": "https://new.example", "project_key": "new"} message = session.get(Message, case.run.source_id) assert message is not None session.add( @@ -713,64 +1053,80 @@ def test_new_run_freezes_updated_project_without_changing_previous_run( from_source=message.from_source, ) ) - new_run = case.repository.start( + service = ServicePorts(_application_trace_service_boundaries_no_egress()).service(case) + new_run = service.start( tenant_id=case.run.tenant_id, app_id=case.run.app_id, source_kind="message", source_id=new_id, session_id="new-session", + workflow_id=None, ) assert new_run is not None - service = ApplicationTracingService(case.repository) - assert service.status(**source(case))["project_url"] == "https://langfuse.test/project/original" + previous = service.status(**source(case.run), node_execution_id=None) + assert previous.project_url == "https://langfuse.test/project/original" assert ( - service.status(**{**source(case), "source_id": new_id})["project_url"] == "https://new.example/project/new" + service.status(**source(new_run), node_execution_id=None).project_url == "https://new.example/project/new" ) @staticmethod @pytest.mark.parametrize( - "project_url", + "frozen", [ - None, - "", - 1, - "javascript:alert(1)", - "//example.com/project", - "https://user:secret@example.com/project", - "https://@example.com", - "https://example.com\\@other.test", - "https://example.com:invalid/project", - "https://[invalid/project", - "https://example.com/\nproject", + {key: value for key, value in FROZEN_LANGFUSE.items() if key != "public_key"}, + {**FROZEN_LANGFUSE, "host": 1}, + {**FROZEN_LANGFUSE, "host": "javascript:alert(1)"}, + {**FROZEN_LANGFUSE, "host": "//example.com/project"}, + {**FROZEN_LANGFUSE, "host": "https://user:secret@example.com/project"}, + {**FROZEN_LANGFUSE, "host": "https://@example.com"}, + {**FROZEN_LANGFUSE, "host": "https://example.com\\@other.test"}, + {**FROZEN_LANGFUSE, "host": "https://[invalid/project"}, ], ) def test_status_hides_unavailable_or_unsafe_frozen_project_urls( - sqlite_session_factory: sessionmaker[Session], project_url: object + sqlite_session_factory: sessionmaker[Session], frozen: dict[str, object] ) -> None: case = _application_trace_service_boundaries_seed(sqlite_session_factory, "invalid-project") - with case.sessions.begin() as session: - run = session.get(ApplicationTraceRun, case.run.id) - assert run is not None - run.provider_config = { - "secret_key": "encrypted-original", - **({"project_url": project_url} if project_url is not None else {}), - } - result = ApplicationTracingService(case.repository).status(**source(case)) - assert result["project_url"] is None + freeze_config(case, "langfuse", frozen) + result = ( + ServicePorts(_application_trace_service_boundaries_no_egress()) + .service(case) + .status(**source(case.run), node_execution_id=None) + ) + assert result.project_url is None assert "encrypted-original" not in str(result) @staticmethod - @pytest.mark.parametrize("provider", ["aliyun", "tencent"]) - def test_status_does_not_present_platform_console_as_project( - sqlite_session_factory: sessionmaker[Session], provider: str + @pytest.mark.parametrize( + ("provider", "frozen"), + [ + ("langfuse", FROZEN_LANGFUSE), + ("aliyun", {"license_key": "encrypted-license", "endpoint": "https://tracing.aliyun.example"}), + ( + "tencent", + {"token": "encrypted-token", "endpoint": "https://apm.tencent.example", "service_name": "dify-app"}, + ), + ], + ) + def test_status_derives_project_link_from_frozen_provider_not_stored_urls( + sqlite_session_factory: sessionmaker[Session], provider: str, frozen: dict[str, object] ) -> None: case = _application_trace_service_boundaries_seed(sqlite_session_factory, "platform-console") - with case.sessions.begin() as session: - run = session.get(ApplicationTraceRun, case.run.id) - assert run is not None - run.provider = provider - run.provider_config = {"project_url": "https://console.example.com"} - assert ApplicationTracingService(case.repository).status(**source(case))["project_url"] is None + freeze_config(case, provider, {**frozen, "project_url": "https://stored.example/project"}) + installed = EntryPointTraceProviderCatalog() + destination = installed.resolve(provider) + expected = destination(_application_trace_service_boundaries_no_egress()).project_url( + destination.config_model.model_validate(frozen) + ) + result = ( + ServicePorts(_application_trace_service_boundaries_no_egress(), providers=installed) + .service(case) + .status(**source(case.run), node_execution_id=None) + ) + assert expected is not None + assert result.project_url == expected + assert result.capabilities is not None + assert result.capabilities.provider == provider @staticmethod @pytest.mark.parametrize("disabled", [True, False]) @@ -783,19 +1139,20 @@ def test_uncollected_status_and_missing_run_retry_do_not_construct_provider( stored = session.get(ApplicationTraceRun, case.run.id) assert stored is not None session.delete(stored) - service = ApplicationTracingService(case.repository) - assert service.status(**source(case)) == { - "status": "not_collected", - "source_kind": "message", - "source_id": case.run.source_id, - } + service = ServicePorts( + _application_trace_service_boundaries_no_egress(), + providers=ForbiddenTraceProviders("an uncollected source has no destination to construct"), + ).service(case) + assert service.status(**source(case.run), node_execution_id=None) == not_collected( + "message", case.run.source_id + ) if not disabled: - with pytest.raises(TraceCollectionError, match="tracing_run_not_found"): - service.retry(**source(case), acknowledge_possible_duplicates=True) + with pytest.raises(TraceCollectionNotFoundError, match="tracing_run_not_found"): + service.retry(**source(case.run), acknowledge_possible_duplicates=True, node_execution_id=None) @staticmethod def test_retry_acknowledgement_guards_schedule_and_keeps_frozen_snapshot( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + sqlite_session_factory: sessionmaker[Session], ) -> None: case = _application_trace_service_boundaries_seed(sqlite_session_factory, "retry") _application_trace_service_boundaries_ingest( @@ -804,57 +1161,117 @@ def test_retry_acknowledgement_guards_schedule_and_keeps_frozen_snapshot( _application_trace_service_boundaries_terminal(case) lease = _application_trace_service_boundaries_claim(case) assert lease is not None - case.repository.complete(case.run.id, lease[1], state="uncertain", error_code="provider_acceptance_unknown") - scheduled: list[str] = [] - service = ApplicationTracingService(case.repository) - monkeypatch.setattr(service, "schedule", scheduled.append) - with pytest.raises(TraceCollectionError, match="tracing_duplicate_acknowledgement_required"): - service.retry(**source(case), acknowledge_possible_duplicates=False) - assert scheduled == [] - result = service.retry(**source(case), acknowledge_possible_duplicates=True) - assert scheduled == [case.run.id] - assert result["status"] == "ready" - assert result["config_revision"] == case.run.config_revision + with case.change() as run: + run.export_uncertain(lease[1], now=naive_utc_now()) + ports = ServicePorts(_application_trace_service_boundaries_no_egress()) + service = ports.service(case) + with pytest.raises(TraceCollectionConflictError, match="tracing_duplicate_acknowledgement_required"): + service.retry(**source(case.run), acknowledge_possible_duplicates=False, node_execution_id=None) + assert ports.scheduled == [] + result = service.retry(**source(case.run), acknowledge_possible_duplicates=True, node_execution_id=None) + assert ports.scheduled == [case.run.id] + assert result.status is TraceRunState.READY + assert result.config_revision == case.run.config_revision current = case.repository.get(case.run.id) assert current is not None assert current.provider_config == case.run.provider_config @staticmethod @pytest.mark.parametrize( - ("failure", "code"), + ("failure", "max_attempts", "state", "code"), [ - (RetryableTraceDispatchError("private-transport-error"), "provider_retry_exhausted"), - (ValueError("private-credential-error"), "provider_export_rejected"), + (ExportRetryableError("private-transport-error"), 1, TraceRunState.FAILED, "provider_retry_exhausted"), + ( + ExportRetryableError("private-transport-error"), + 2, + TraceRunState.RETRYING, + "provider_temporarily_unavailable", + ), + (ExportRejectedError("private-rejection"), 2, TraceRunState.FAILED, "provider_export_rejected"), + ( + ExportOutcomeUnknownError("private-acknowledgement"), + 2, + TraceRunState.UNCERTAIN, + "provider_acceptance_unknown", + ), + (ValueError("private-credential-error"), 2, TraceRunState.FAILED, "provider_export_rejected"), ], ) - def test_delivery_exhaustion_and_rejection_never_log_vendor_details( + def test_delivery_maps_provider_failures_to_stable_codes_without_vendor_details( sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, failure: Exception, + max_attempts: int, + state: TraceRunState, code: str, ) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_MAX_ATTEMPTS=1) case = _application_trace_service_boundaries_seed(sqlite_session_factory, "delivery") _application_trace_service_boundaries_ingest( case, _application_trace_service_boundaries_batch_for(case.run, case.invocation()) ) _application_trace_service_boundaries_terminal(case) - - def export(*_args: object) -> TraceExportResult: - raise failure - - monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", lambda *_args: {}) - monkeypatch.setattr(service_module, "export_trace", export) - service = ApplicationTracingService(case.repository) + destination = RecordingEgress([failure], grpc_target=None) + service = ServicePorts( + destination, settings=replace(AGENT_TRACE_SETTINGS, max_export_attempts=max_attempts) + ).service(case) service.deliver(case.run.id) current = case.repository.get(case.run.id) assert current is not None - assert current.state == "failed" - assert current.error_code == code + assert (current.state, current.error_code) == (state, code) + assert current.attempts == 1 + with case.sessions() as session: + next_attempt_at = session.scalar( + select(ApplicationTraceRun.next_attempt_at).where(ApplicationTraceRun.id == case.run.id) + ) + assert (next_attempt_at is not None) is (state is TraceRunState.RETRYING) + assert len(destination.requests) == 1 assert "private-" not in caplog.text service.deliver(_application_trace_service_boundaries_identifier("missing-run")) + @staticmethod + @pytest.mark.parametrize("unreadable", ["decryption", "validation"]) + def test_unreadable_frozen_configuration_fails_before_provider_io( + sqlite_session_factory: sessionmaker[Session], + caplog: pytest.LogCaptureFixture, + unreadable: str, + ) -> None: + case = _application_trace_service_boundaries_seed(sqlite_session_factory, "unreadable") + _application_trace_service_boundaries_ingest( + case, _application_trace_service_boundaries_batch_for(case.run, case.invocation()) + ) + _application_trace_service_boundaries_terminal(case) + if unreadable == "validation": + freeze_config(case, "langfuse", {**FROZEN_LANGFUSE, "host": "private-host-without-scheme"}) + cipher = UnreadableSecretCipher() if unreadable == "decryption" else StoredSecretCipher() + destination = RecordingEgress([], grpc_target=None) + service = ServicePorts(destination, cipher=cipher).service(case) + service.deliver(case.run.id) + current = case.repository.get(case.run.id) + assert current is not None + assert (current.state, current.error_code) == (TraceRunState.FAILED, "tracing_config_unreadable") + assert destination.requests == [] + assert "private-" not in caplog.text + with pytest.raises(TraceCollectionConflictError, match="tracing_retry_not_allowed"): + service.retry(**source(case.run), acknowledge_possible_duplicates=True, node_execution_id=None) + + @staticmethod + def test_unavailable_frozen_provider_fails_before_decryption( + sqlite_session_factory: sessionmaker[Session], + ) -> None: + case = _application_trace_service_boundaries_seed(sqlite_session_factory, "retired-provider") + _application_trace_service_boundaries_ingest( + case, _application_trace_service_boundaries_batch_for(case.run, case.invocation()) + ) + _application_trace_service_boundaries_terminal(case) + freeze_config(case, "retired-provider", FROZEN_LANGFUSE) + ServicePorts( + _application_trace_service_boundaries_no_egress(), + cipher=ForbiddenSecretCipher("an unavailable provider must not decrypt credentials"), + ).service(case).deliver(case.run.id) + current = case.repository.get(case.run.id) + assert current is not None + assert (current.state, current.error_code) == (TraceRunState.FAILED, "tracing_provider_unavailable") + @staticmethod @pytest.mark.parametrize("capture_content", [False, True]) @pytest.mark.parametrize( @@ -897,7 +1314,9 @@ def test_preview_terminal_uses_persisted_node_state_and_frozen_privacy( ) node_record.created_at = datetime(2026, 1, 1) session.add(node_record) - run = case.repository.start( + destination = _application_trace_service_boundaries_accepted("preview-trace") + service = ServicePorts(destination).service(case) + run = service.start( tenant_id=case.run.tenant_id, app_id=case.run.app_id, source_kind="node", @@ -906,45 +1325,26 @@ def test_preview_terminal_uses_persisted_node_state_and_frozen_privacy( session_id="preview-session", ) assert run is not None - ApplicationTracingService(case.repository).settle_preview(run) + service.deliver(run.id) current = case.repository.get(run.id) assert current is not None if status == WorkflowNodeExecutionStatus.RUNNING: assert current.outer_trace is None + assert current.state is TraceRunState.COLLECTING + assert destination.requests == [] else: assert current.outer_trace is not None - trace = service_module.CanonicalTrace.model_validate(current.outer_trace) - root = trace.spans[0] + assert current.state is TraceRunState.ACCEPTED + trace = Trace.model_validate(current.outer_trace) + root = trace.root + assert root.span_id == run.root_span_id assert root.status.value == ("error" if status == WorkflowNodeExecutionStatus.FAILED else "ok") assert ("private-" in trace.model_dump_json()) is capture_content - assert root.metadata["node_execution_id"] == node_id + assert root.attributes["node_execution_id"] == node_id + assert (root.end_time - root.start_time).total_seconds() == 2 if not capture_content and status == WorkflowNodeExecutionStatus.FAILED: assert root.error == "Node execution failed" - @staticmethod - @pytest.mark.parametrize("owner_state", ["no_creator", "missing_account", "wrong_app"]) - def test_missing_owner_identity_cannot_authorize_trace_expansion( - sqlite_session_factory: sessionmaker[Session], owner_state: str - ) -> None: - from models.model import App - - case = _application_trace_service_boundaries_seed(sqlite_session_factory, "owner") - with case.sessions.begin() as session: - app = session.get(App, case.run.app_id) - assert app is not None - app.created_by = ( - None - if owner_state == "no_creator" - else _application_trace_service_boundaries_identifier("missing-owner") - ) - app_id = ( - _application_trace_service_boundaries_identifier("another-app") - if owner_state == "wrong_app" - else case.run.app_id - ) - with pytest.raises(TraceCollectionError, match="tracing_owner_not_found"): - ApplicationTracingService(case.repository)._service_account(case.run, app_id) - @staticmethod @pytest.mark.parametrize( ("invalid", "code"), @@ -952,11 +1352,12 @@ def test_missing_owner_identity_cannot_authorize_trace_expansion( ("empty", "tracing_invocation_empty"), ("foreign_parent", "tracing_invocation_parent_invalid"), ("no_invocation_root", "tracing_invocation_root_missing"), + ("self_parent", "tracing_payload_invalid"), ("conflicting_span", "tracing_span_conflict"), ], ) def test_invalid_runtime_tree_is_rejected_before_provider_credentials_or_io( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, invalid: str, code: str + sqlite_session_factory: sessionmaker[Session], invalid: str, code: str ) -> None: case = _application_trace_service_boundaries_seed(sqlite_session_factory, "invalid-runtime") batch = _application_trace_service_boundaries_batch_for(case.run, case.invocation()) @@ -965,11 +1366,15 @@ def test_invalid_runtime_tree_is_rejected_before_provider_credentials_or_io( elif invalid == "foreign_parent": batch.spans[0] = batch.spans[0].model_copy(update={"parent_span_id": "f" * 16}) elif invalid == "no_invocation_root": + # Every parent is a collected span, but none of them hangs from the invocation. + batch.spans = [ + batch.spans[0].model_copy(update={"parent_span_id": "b" * 16}), + batch.spans[0].model_copy(update={"span_id": "b" * 16, "parent_span_id": "a" * 16}), + ] + elif invalid == "self_parent": batch.spans[0] = batch.spans[0].model_copy(update={"parent_span_id": batch.spans[0].span_id}) else: _application_trace_service_boundaries_ingest(case, batch.model_copy(update={"final": False})) - from uuid import uuid4 - batch = batch.model_copy( update={ "batch_id": uuid4(), @@ -978,24 +1383,57 @@ def test_invalid_runtime_tree_is_rejected_before_provider_credentials_or_io( ) _application_trace_service_boundaries_ingest(case, batch) _application_trace_service_boundaries_terminal(case) - - def forbidden(*_args: object, **_kwargs: object) -> None: - pytest.fail("invalid trees must fail before credentials or provider requests") - - monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", forbidden) - monkeypatch.setattr(service_module, "export_trace", forbidden) - service = ApplicationTracingService(case.repository) - service.deliver(case.run.id) + reason = "invalid trees must fail before credentials or provider requests" + ServicePorts( + _application_trace_service_boundaries_no_egress(), + providers=ForbiddenTraceProviders(reason), + cipher=ForbiddenSecretCipher(reason), + ).service(case).deliver(case.run.id) current = case.repository.get(case.run.id) assert current is not None - assert current.state == "failed" + assert current.state is TraceRunState.FAILED assert current.error_code == code @staticmethod - def test_assemble_requires_a_persisted_business_terminal(sqlite_session_factory: sessionmaker[Session]) -> None: + def test_delivery_requires_a_persisted_business_terminal(sqlite_session_factory: sessionmaker[Session]) -> None: case = _application_trace_service_boundaries_seed(sqlite_session_factory, "no-terminal") - with pytest.raises(TraceCollectionError, match="tracing_terminal_missing"): - ApplicationTracingService(case.repository).assemble(case.run) + attach_conversation(case, AppMode.AGENT) + destination = _application_trace_service_boundaries_no_egress() + ServicePorts(destination).service(case).deliver(case.run.id) + current = case.repository.get(case.run.id) + assert current is not None + assert current.state is TraceRunState.COLLECTING + assert current.outer_trace is None + assert destination.requests == [] + + @staticmethod + def test_assembled_trace_nests_runtime_spans_under_their_invocation_with_trusted_identity( + sqlite_session_factory: sessionmaker[Session], + ) -> None: + case = _application_trace_service_boundaries_seed(sqlite_session_factory, "assembled") + invocation = case.invocation() + _application_trace_service_boundaries_ingest( + case, _application_trace_service_boundaries_batch_for(case.run, invocation) + ) + _application_trace_service_boundaries_terminal(case) + destination = _application_trace_service_boundaries_accepted("assembled-trace") + ServicePorts(destination).service(case).deliver(case.run.id) + [trace] = sent_traces(destination) + assert trace.trace_id == case.run.trace_id + assert trace.session_id == "session-original" + assert [(span.span_id, span.parent_span_id) for span in trace.spans] == [ + (case.run.root_span_id, None), + (invocation.span_id, case.run.root_span_id), + ("a" * 16, invocation.span_id), + ] + attempt = trace.spans[1] + assert attempt.name == "Agent attempt 1" + assert attempt.attributes["invocation_id"] == invocation.id + assert (attempt.start_time, attempt.end_time) == (trace.spans[2].start_time, trace.spans[2].end_time) + for span in trace.spans: + assert span.attributes["dify.tracing.run_id"] == case.run.id + assert span.attributes["dify.source.id"] == case.run.source_id + assert span.attributes["tenant_id"] == case.run.tenant_id """Malformed durable traces must never be sent or offered for provider retry.""" @@ -1003,6 +1441,7 @@ def test_assemble_requires_a_persisted_business_terminal(sqlite_session_factory: _invalid_trace_delivery_batch_for = _application_trace_collection_batch_for _invalid_trace_delivery_ingest = _application_trace_collection_ingest +_invalid_trace_delivery_no_egress = _application_trace_collection_no_egress _invalid_trace_delivery_seed = _application_trace_collection_seed _invalid_trace_delivery_terminal = _application_trace_collection_terminal @@ -1011,7 +1450,7 @@ class TestInvalidTraceDeliveryCases: @staticmethod @pytest.mark.parametrize("invalid", ["cycle", "invalid_envelope"]) def test_invalid_tree_is_permanent_before_decryption_or_provider_io( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, invalid: str + sqlite_session_factory: sessionmaker[Session], invalid: str ) -> None: case = _invalid_trace_delivery_seed(sqlite_session_factory, str(uuid4())) batch = _invalid_trace_delivery_batch_for(case.run, case.invocation()) @@ -1035,17 +1474,735 @@ def test_invalid_tree_is_permanent_before_decryption_or_provider_io( row = session.get(ApplicationTraceRun, case.run.id) assert row is not None row.outer_trace = {"corrupt": True} - - def forbidden(*_args: object, **_kwargs: object) -> None: - pytest.fail("invalid traces must fail before credentials or provider I/O") - - monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", forbidden) - monkeypatch.setattr(service_module, "export_trace", forbidden) - ApplicationTracingService(case.repository).deliver(case.run.id) + reason = "invalid traces must fail before credentials or provider I/O" + ServicePorts( + _invalid_trace_delivery_no_egress(), + providers=ForbiddenTraceProviders(reason), + cipher=ForbiddenSecretCipher(reason), + ).service(case).deliver(case.run.id) current = case.repository.get(case.run.id) assert current is not None - assert current.state == "failed" + assert current.state is TraceRunState.FAILED assert current.error_code == "tracing_payload_invalid" assert current.attempts == 1 - with pytest.raises(TraceCollectionError, match="tracing_retry_not_allowed"): - case.repository.retry(run_id=case.run.id, acknowledge_possible_duplicates=True) + # A retry reports the run's status, which resolves its frozen provider; the forbidden ports guard only delivery. + with pytest.raises(TraceCollectionConflictError, match="tracing_retry_not_allowed"): + ServicePorts(_invalid_trace_delivery_no_egress()).service(case).retry( + **source(case.run), acknowledge_possible_duplicates=True, node_execution_id=None + ) + + +"""Node logs resolve the caller's trace using SQL or configured LogStore ownership.""" + + +_ApplicationTraceNodeAliasCollection = _ApplicationTraceCollectionCollection +_application_trace_node_alias_batch_for = _application_trace_collection_batch_for +_application_trace_node_alias_claim = _application_trace_collection_claim +_application_trace_node_alias_identifier = _application_trace_collection_identifier +_application_trace_node_alias_ingest = _application_trace_collection_ingest +_application_trace_node_alias_no_egress = _application_trace_collection_no_egress +_application_trace_node_alias_seed = _application_trace_collection_seed +_application_trace_node_alias_terminal = _application_trace_collection_terminal + + +@dataclass +class NodeAlias: + case: _ApplicationTraceNodeAliasCollection + node: dict[str, object] + workflow: dict[str, object] + requests: list[str] + logstore: bool + ports: ServicePorts + + def service(self) -> ApplicationTracingService: + return self.ports.service(self.case) + + def status(self, node_execution_id: str | None = None) -> TraceRunStatus: + return self.service().status( + tenant_id=self.case.run.tenant_id, + app_id=self.case.run.app_id, + source_kind="node", + source_id=str(self.node["id"]), + node_execution_id=node_execution_id, + ) + + def accept(self, *, spans: bool = True) -> tuple[str, str]: + span_urls: dict[str, str] = {} + for attempt in (0, 1): + invocation = self.case.open_invocation( + agent_id=_application_trace_node_alias_identifier("referenced-agent"), + node_execution_id=str(self.node["id"]), + attempt=attempt, + call_key=f"attempt-{attempt}", + ) + _application_trace_node_alias_ingest( + self.case, + _application_trace_node_alias_batch_for(self.case.run, invocation, span_id=str(attempt + 1) * 16), + ) + span_urls[invocation.span_id] = f"https://provider.test/span/{attempt}" + _application_trace_node_alias_terminal(self.case) + lease = _application_trace_node_alias_claim(self.case) + assert lease is not None + _accept_export( + self.case, + lease[1], + SendResult( + trace_id="provider-trace", + trace_url="https://provider.test/trace", + span_urls=span_urls if spans else {}, + ), + ) + return "https://provider.test/trace", "https://provider.test/span/1" + + +class TestApplicationTraceNodeAliasCases: + @staticmethod + @pytest.fixture(params=["sql", "logstore"]) + def alias( + request: pytest.FixtureRequest, + sqlite_engine: Engine, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + ) -> Iterator[NodeAlias]: + case = _application_trace_node_alias_seed(sqlite_session_factory, "node-alias") + with case.sessions.begin() as session: + app = session.get(App, case.run.app_id) + assert app is not None + app.mode = AppMode.WORKFLOW + common: dict[str, object] = { + "tenant_id": case.run.tenant_id, + "app_id": case.run.app_id, + "workflow_id": _application_trace_node_alias_identifier("definition"), + "created_by": _application_trace_node_alias_identifier("actor"), + "created_at": naive_utc_now().isoformat(), + "status": "succeeded", + } + workflow = { + **common, + "id": _application_trace_node_alias_identifier("workflow-source"), + "graph": json.dumps({"nodes": []}), + } + node = { + **common, + "id": _application_trace_node_alias_identifier("node-source"), + "workflow_run_id": workflow["id"], + "node_id": "reference-agent", + "node_type": "agent", + "title": "Reference Agent", + } + requests: list[str] = [] + transactions: set[Connection] = set() + + def begun(connection: Connection) -> None: + transactions.add(connection) + + def ended(connection: Connection) -> None: + transactions.discard(connection) + + if request.param == "sql": + with case.sessions.begin() as session: + session.add(_dict_to_workflow_run(workflow)) + session.add(_dict_to_workflow_node_execution_model(node)) + else: + apply_config_overrides( + monkeypatch, + LOGSTORE_DUAL_WRITE_ENABLED=False, + API_WORKFLOW_RUN_REPOSITORY=( + "extensions.logstore.repositories.logstore_api_workflow_run_repository.LogstoreAPIWorkflowRunRepository" + ), + API_WORKFLOW_NODE_EXECUTION_REPOSITORY=( + "extensions.logstore.repositories.logstore_api_workflow_node_execution_repository." + "LogstoreAPIWorkflowNodeExecutionRepository" + ), + ) + for key, value in { + "LOGSTORE_DUAL_READ_ENABLED": "false", + "LOGSTORE_PG_MODE_ENABLED": "false", + "ALIYUN_SLS_ACCESS_KEY_ID": "test-access-key", + "ALIYUN_SLS_ACCESS_KEY_SECRET": "test-secret", + "ALIYUN_SLS_ENDPOINT": "https://sls.example.test", + "ALIYUN_SLS_REGION": "test-region", + }.items(): + monkeypatch.setenv(key, value) + monkeypatch.setattr(AliyunLogStore, "_instance", None) + monkeypatch.setattr(AliyunLogStore, "_initialized", False) + monkeypatch.setattr( + AliyunLogStore, "_check_endpoint_connectivity", staticmethod(lambda _endpoint, _timeout: None) + ) + + def reject_http_transport(*_args: object, **_kwargs: object) -> None: + pytest.fail("Configured LogStore test escaped its isolated transport boundary") + + monkeypatch.setattr(LogClient, "_getHttpResponse", reject_http_transport) + + def get_logs( + _client: AliyunLogStore, *, query: str, logstore: str, **_kwargs: object + ) -> list[dict[str, object]]: + assert not transactions, "Configured source reads must not hold a SQL transaction" + assert f'tenant_id:"{case.run.tenant_id}"' in query.replace(": ", ":") + assert logstore in { + AliyunLogStore.workflow_execution_logstore, + AliyunLogStore.workflow_node_execution_logstore, + } + is_workflow = logstore == AliyunLogStore.workflow_execution_logstore + if is_workflow: + assert f'app_id:"{case.run.app_id}"' in query + requests.append("workflow" if is_workflow else "node") + row = workflow if is_workflow else node + return [dict(row)] if row else [] + + monkeypatch.setattr(AliyunLogStore, "get_logs", get_logs) + + def execute_sql( + client: AliyunLogStore, *, sql: str, query: str, logstore: str, **_kwargs: object + ) -> list[dict[str, object]]: + # Node lookups use the SQL transport in SDK mode too; the real + # service still validates the returned node-to-workflow owner chain. + assert logstore == AliyunLogStore.workflow_node_execution_logstore + assert f'id:"{_application_trace_node_alias_identifier("node-source")}"' in query.replace(": ", ":") + assert f"tenant_id = '{case.run.tenant_id}'" in sql + assert f"id = '{_application_trace_node_alias_identifier('node-source')}'" in sql + assert f'FROM "{logstore}"' in sql + assert "ORDER BY log_version DESC LIMIT 1" in sql + return get_logs(client, query=query, logstore=logstore) + + monkeypatch.setattr(AliyunLogStore, "execute_sql", execute_sql) + with case.sessions() as session: + assert session.scalar(select(WorkflowRun)) is None + assert session.scalar(select(WorkflowNodeExecutionModel)) is None + event.listen(sqlite_engine, "begin", begun) + event.listen(sqlite_engine, "commit", ended) + event.listen(sqlite_engine, "rollback", ended) + try: + service = ServicePorts(_application_trace_node_alias_no_egress()).service(case) + run = service.start( + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, + source_kind="workflow", + source_id=str(workflow["id"]), + session_id="caller-session", + workflow_id=None, + ) + assert run is not None + requests.clear() + yield NodeAlias( + _ApplicationTraceNodeAliasCollection(case.repository, case.sessions, run), + node, + workflow, + requests, + request.param == "logstore", + ServicePorts(_application_trace_node_alias_no_egress()), + ) + finally: + event.remove(sqlite_engine, "begin", begun) + event.remove(sqlite_engine, "commit", ended) + event.remove(sqlite_engine, "rollback", ended) + + @staticmethod + @pytest.mark.parametrize("spans", [True, False]) + def test_alias_returns_parent_trace_and_latest_owned_invocation(alias: NodeAlias, spans: bool) -> None: + trace_url, span_url = alias.accept(spans=spans) + status = alias.status() + assert status.status is TraceRunState.ACCEPTED + assert status.source_kind == "node" + assert status.source_id == alias.node["id"] + assert status.trace_id == "provider-trace" + assert status.trace_url == trace_url + assert status.span_url == (span_url if spans else None) + assert status.config_revision == alias.case.run.config_revision + if alias.logstore: + assert alias.requests == ["node", "workflow"] + + @staticmethod + @pytest.mark.parametrize("capture_agent", [True, False]) + @pytest.mark.parametrize("accepted", [True, False]) + def test_alias_without_internal_invocation_returns_parent_status( + alias: NodeAlias, capture_agent: bool, accepted: bool + ) -> None: + with alias.case.sessions.begin() as session: + run = session.get(ApplicationTraceRun, alias.case.run.id) + assert run is not None + run.capture_agent = capture_agent + if accepted: + _application_trace_node_alias_terminal(alias.case) + lease = _application_trace_node_alias_claim(alias.case) + assert lease is not None + _accept_export( + alias.case, + lease[1], + SendResult(trace_id="outer-only", trace_url="https://provider.test/outer", span_urls={}), + ) + status = alias.status() + assert status.status is (TraceRunState.ACCEPTED if accepted else TraceRunState.COLLECTING) + assert status.span_url is None + assert status.capture_agent is capture_agent + assert status.trace_url == ("https://provider.test/outer" if accepted else None) + + @staticmethod + @pytest.mark.parametrize("state", [TraceRunState.COLLECTING, TraceRunState.NOT_COLLECTED]) + def test_standalone_node_snapshot_wins_over_parent(alias: NodeAlias, state: TraceRunState) -> None: + alias.accept() + node_run = alias.service().start( + tenant_id=alias.case.run.tenant_id, + app_id=alias.case.run.app_id, + source_kind="node", + source_id=str(alias.node["id"]), + session_id="preview-session", + workflow_id=None, + ) + assert node_run is not None + with alias.case.sessions.begin() as session: + run = session.get(ApplicationTraceRun, node_run.id) + assert run is not None + run.state = state + alias.requests.clear() + status = alias.status() + assert status.status is state + assert status.trace_url is None + if alias.logstore: + assert alias.requests == ["node"] + + @staticmethod + def test_explicit_other_node_filter_cannot_select_sibling_span(alias: NodeAlias) -> None: + sibling = alias.case.open_invocation( + agent_id=_application_trace_node_alias_identifier("referenced-agent"), + node_execution_id=_application_trace_node_alias_identifier("sibling-node"), + attempt=0, + call_key="sibling", + ) + _application_trace_node_alias_ingest( + alias.case, _application_trace_node_alias_batch_for(alias.case.run, sibling, span_id="f" * 16) + ) + alias.accept() + with pytest.raises(TraceCollectionNotFoundError, match="tracing_node_not_found"): + alias.status(_application_trace_node_alias_identifier("sibling-node")) + with pytest.raises(TraceCollectionNotFoundError, match="tracing_node_not_found"): + alias.service().retry( + tenant_id=alias.case.run.tenant_id, + app_id=alias.case.run.app_id, + source_kind="node", + source_id=str(alias.node["id"]), + node_execution_id=_application_trace_node_alias_identifier("sibling-node"), + acknowledge_possible_duplicates=True, + ) + assert not alias.ports.scheduled + + @staticmethod + @pytest.mark.parametrize("state", [TraceRunState.FAILED, TraceRunState.UNCERTAIN, TraceRunState.ACCEPTED]) + def test_alias_retry_uses_parent_guards_and_resolves_sources_once(alias: NodeAlias, state: TraceRunState) -> None: + alias.accept() + with alias.case.sessions.begin() as session: + run = session.get(ApplicationTraceRun, alias.case.run.id) + assert run is not None + run.state = state + service = alias.service() + + def retry(acknowledge: bool) -> TraceRunStatus: + alias.requests.clear() + return service.retry( + tenant_id=alias.case.run.tenant_id, + app_id=alias.case.run.app_id, + source_kind="node", + source_id=str(alias.node["id"]), + acknowledge_possible_duplicates=acknowledge, + node_execution_id=None, + ) + + if state is TraceRunState.ACCEPTED: + with pytest.raises(TraceCollectionConflictError, match="tracing_already_accepted"): + retry(True) + assert not alias.ports.scheduled + else: + if state is TraceRunState.UNCERTAIN: + with pytest.raises(TraceCollectionConflictError, match="tracing_duplicate_acknowledgement_required"): + retry(False) + assert not alias.ports.scheduled + assert retry(state is TraceRunState.UNCERTAIN).status is TraceRunState.READY + assert alias.ports.scheduled == [alias.case.run.id] + if alias.logstore: + assert alias.requests == ["node", "workflow"] + + @staticmethod + @pytest.mark.parametrize("source", ["node", "workflow"]) + @pytest.mark.parametrize("field", ["tenant_id", "app_id", "workflow_id"]) + def test_alias_rejects_broken_owner_chain(alias: NodeAlias, source: str, field: str) -> None: + alias.accept() + row = alias.node if source == "node" else alias.workflow + if alias.logstore: + row[field] = _application_trace_node_alias_identifier("foreign") + else: + with alias.case.sessions.begin() as session: + model = session.get(WorkflowNodeExecutionModel if source == "node" else WorkflowRun, str(row["id"])) + assert model is not None + # Keep field assignments explicit so ownership failures are visible. + if field == "tenant_id": + model.tenant_id = _application_trace_node_alias_identifier("foreign") + elif field == "app_id": + model.app_id = _application_trace_node_alias_identifier("foreign") + else: + model.workflow_id = _application_trace_node_alias_identifier("foreign") + with pytest.raises(TraceCollectionNotFoundError, match="tracing_source_not_found"): + alias.status() + + @staticmethod + def test_alias_missing_parent_business_source_is_not_a_durable_owner_bypass(alias: NodeAlias) -> None: + if alias.logstore: + alias.workflow.clear() + else: + with alias.case.sessions.begin() as session: + parent = session.get(WorkflowRun, str(alias.workflow["id"])) + assert parent is not None + session.delete(parent) + with pytest.raises(TraceCollectionNotFoundError, match="tracing_source_not_found"): + alias.status() + + @staticmethod + def test_owned_alias_without_parent_trace_is_not_collected(alias: NodeAlias) -> None: + with alias.case.sessions.begin() as session: + run = session.get(ApplicationTraceRun, alias.case.run.id) + assert run is not None + session.delete(run) + assert alias.status() == not_collected("node", str(alias.node["id"])) + + +"""Real SQLite state transitions and background renewal for delivery leases. + +PostgreSQL row-lock scheduling and the deployed Celery pool still need dev +validation; these tests establish token and acknowledgement semantics. +""" + + +_ExportLeaseCollection = _ApplicationTraceCollectionCollection +_export_lease_accepted = _application_trace_collection_accepted +_export_lease_batch_for = _application_trace_collection_batch_for +_export_lease_claim = _application_trace_collection_claim +_export_lease_ingest = _application_trace_collection_ingest +_export_lease_seed = _application_trace_collection_seed +_export_lease_terminal = _application_trace_collection_terminal + + +class _SlowDestination(RecordingEgress): + """Runs ``during_request`` while the export request is in flight.""" + + def __init__(self, replies: list[EgressResponse | Exception], during_request: Callable[[], None]) -> None: + super().__init__(replies, grpc_target=None) + self._during_request = during_request + + @override + def request( + self, + method: str, + url: str, + *, + headers: Mapping[str, str], + content: bytes | None, + timeout: float, + sensitive_url: bool, + ) -> EgressResponse: + self._during_request() + return super().request( + method, url, headers=headers, content=content, timeout=timeout, sensitive_url=sensitive_url + ) + + +def _export_lease(locked: LockedTraceRun | None) -> str | None: + """The token of the export attempt that holds the run; ``None`` when no attempt holds it.""" + if locked is None or locked.run.state is not TraceRunState.EXPORTING: + return None + return locked.run.export_token + + +class _RenewalRecordingRepository(ApplicationTraceRepository): + """Records the thread of every lease renewal; the first one can fail like a dropped connection. + + A renewal is a change after which the export attempt that held the run still holds it. + """ + + def __init__(self, sessions: sessionmaker[Session], *, transient_failure: bool) -> None: + super().__init__(sessions) + self._transient_failure = transient_failure + self.renewal_threads: list[int] = [] + self.renewed = threading.Event() + + @override + @contextmanager + def change(self, run_id: str, *, tenant_id: str) -> Generator[LockedTraceRun | None]: + renewal = False + with super().change(run_id, tenant_id=tenant_id) as locked: + lease = _export_lease(locked) + yield locked + renewal = lease is not None and _export_lease(locked) == lease + if renewal: + self.renewal_threads.append(threading.get_ident()) + if self._transient_failure and len(self.renewal_threads) == 1: + raise RuntimeError("private-provider-value") + if renewal: + self.renewed.set() + + +class _UnavailableCompletionRepository(ApplicationTraceRepository): + """Loses the database when the export outcome is recorded.""" + + @override + @contextmanager + def change(self, run_id: str, *, tenant_id: str) -> Generator[LockedTraceRun | None]: + with super().change(run_id, tenant_id=tenant_id) as locked: + lease = _export_lease(locked) + yield locked + if lease is not None and _export_lease(locked) is None: + raise TracingStoreUnavailableError("database unavailable") + + +def result(trace_id: str) -> SendResult: + return SendResult(trace_id=trace_id, trace_url=None, span_urls={}) + + +def _renew_export_lease(case: _ExportLeaseCollection, token: str, now: datetime) -> bool: + with case.change() as run: + return run.renew_export_lease(token, now=now) + + +def _release_for_retry(case: _ExportLeaseCollection) -> None: + """An operator's retry that acknowledges possible duplicates.""" + with case.change() as run: + run.release_for_retry(acknowledge_possible_duplicates=True, now=naive_utc_now()) + + +def _stored(case: _ExportLeaseCollection) -> ApplicationTraceRun: + """The run's row: its record does not carry the export lease.""" + with case.sessions() as session: + row = session.get(ApplicationTraceRun, case.run.id) + assert row is not None + return row + + +def _expire(case: _ExportLeaseCollection) -> str: + leased = _export_lease_claim(case) + assert leased is not None + with case.sessions.begin() as session: + row = session.get(ApplicationTraceRun, case.run.id) + assert row is not None + row.export_started_at = naive_utc_now() - timedelta(seconds=60) + row.export_heartbeat_at = row.export_started_at + assert _export_lease_claim(case) is None + assert _stored(case).state == TraceRunState.UNCERTAIN + return leased[1] + + +class TestExportLeaseCases: + @staticmethod + @pytest.fixture + def case(sqlite_session_factory: sessionmaker[Session]) -> _ExportLeaseCollection: + collection = _export_lease_seed(sqlite_session_factory, "lease-test") + _export_lease_ingest(collection, _export_lease_batch_for(collection.run, collection.invocation())) + _export_lease_terminal(collection) + return collection + + @staticmethod + def test_renewal_keeps_long_export_alive_then_missing_heartbeat_expires(case: _ExportLeaseCollection) -> None: + now = naive_utc_now() + leased = _claim_at(case, lease_seconds=30, now=now) + assert leased is not None + token = leased[1] + now += timedelta(seconds=20) + assert _renew_export_lease(case, token, now) + now += timedelta(seconds=20) + assert _claim_at(case, lease_seconds=30, now=now) is None + current = _stored(case) + assert current.state == TraceRunState.EXPORTING + assert current.attempts == 1 + assert current.export_started_at is not None + assert current.export_started_at + timedelta(seconds=30) < now + now += timedelta(seconds=11) + assert _claim_at(case, lease_seconds=30, now=now) is None + expired = _stored(case) + assert expired.state == TraceRunState.UNCERTAIN + assert expired.error_code == "export_lease_expired" + assert expired.export_token == token + assert not _renew_export_lease(case, token, now) + + @staticmethod + def test_old_rows_without_heartbeat_use_started_at(case: _ExportLeaseCollection) -> None: + leased = _export_lease_claim(case) + assert leased is not None + with case.sessions.begin() as session: + row = session.get(ApplicationTraceRun, case.run.id) + assert row is not None + row.export_started_at = naive_utc_now() - timedelta(seconds=60) + row.export_heartbeat_at = None + assert _export_lease_claim(case) is None + current = _stored(case) + assert current.state == TraceRunState.UNCERTAIN + assert current.export_token == leased[1] + + @staticmethod + def test_wrong_token_cannot_renew_or_complete(case: _ExportLeaseCollection) -> None: + leased = _export_lease_claim(case) + assert leased is not None + before = _stored(case) + assert not _renew_export_lease(case, "other-token", naive_utc_now()) + _accept_export(case, "other-token", result("wrong")) + current = _stored(case) + assert current.export_token == leased[1] + assert current.export_heartbeat_at == before.export_heartbeat_at + assert current.state == TraceRunState.EXPORTING + assert current.provider_trace_id is None + + @staticmethod + def test_late_success_wins_before_manual_retry_and_is_idempotent(case: _ExportLeaseCollection) -> None: + token = _expire(case) + confirmed = SendResult( + trace_id="confirmed", + trace_url="https://example.com/trace", + span_urls={"a" * 16: "https://example.com/span"}, + ) + _accept_export(case, token, confirmed) + current = _stored(case) + assert current.state == TraceRunState.ACCEPTED + assert current.error_code is None + assert current.export_token is None + assert current.accepted_at is not None + assert current.trace_url == confirmed.trace_url + assert current.span_urls == {"a" * 16: "https://example.com/span"} + _accept_export(case, token, result("duplicate-ack")) + assert _stored(case).provider_trace_id == "confirmed" + with pytest.raises(TraceRunConflictError, match="tracing_already_accepted"): + _release_for_retry(case) + + @staticmethod + def test_manual_retry_revokes_old_heartbeat_and_ack_before_and_after_new_claim( + case: _ExportLeaseCollection, + ) -> None: + old_token = _expire(case) + _release_for_retry(case) + assert not _renew_export_lease(case, old_token, naive_utc_now()) + _accept_export(case, old_token, result("old")) + assert _stored(case).state == TraceRunState.READY + second = _export_lease_claim(case) + assert second is not None + assert second[1] != old_token + _accept_export(case, old_token, result("old")) + current = _stored(case) + assert current.state == TraceRunState.EXPORTING + assert current.export_token == second[1] + assert current.provider_trace_id is None + _accept_export(case, second[1], result("new")) + assert _stored(case).provider_trace_id == "new" + + @staticmethod + @pytest.mark.parametrize("late_state", [TraceRunState.FAILED, TraceRunState.RETRYING, TraceRunState.UNCERTAIN]) + def test_late_non_success_cannot_overwrite_lease_expiry( + case: _ExportLeaseCollection, late_state: TraceRunState + ) -> None: + token = _expire(case) + now = naive_utc_now() + with case.change() as run: + if late_state is TraceRunState.FAILED: + run.export_failed(token, "late-failure", now=now) + elif late_state is TraceRunState.RETRYING: + run.export_deferred(token, max_attempts=AGENT_TRACE_SETTINGS.max_export_attempts, now=now) + else: + run.export_uncertain(token, now=now) + current = _stored(case) + assert current.state == TraceRunState.UNCERTAIN + assert current.error_code == "export_lease_expired" + + @staticmethod + def test_other_uncertain_reason_cannot_be_promoted_by_late_ack(case: _ExportLeaseCollection) -> None: + token = _expire(case) + with case.sessions.begin() as session: + row = session.get(ApplicationTraceRun, case.run.id) + assert row is not None + row.error_code = "provider_acceptance_unknown" + _accept_export(case, token, result("late")) + assert _stored(case).state == TraceRunState.UNCERTAIN + + @staticmethod + @pytest.mark.parametrize("transient_failure", [False, True]) + @pytest.mark.parametrize( + ("outcome", "reply"), + [ + (TraceRunState.ACCEPTED, response(200, {"trace_id": "ack"})), + (TraceRunState.UNCERTAIN, EgressError("lost acknowledgement", request_sent=True)), + ], + ) + def test_service_heartbeats_during_export_and_stops_on_every_outcome( + case: _ExportLeaseCollection, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + transient_failure: bool, + outcome: TraceRunState, + reply: EgressResponse | Exception, + ) -> None: + clock = [naive_utc_now()] + # The service dates its claim, renewals, and outcome; the second claimant below passes the same clock. + monkeypatch.setattr(service_module, "naive_utc_now", lambda: clock[0]) + renewals = _RenewalRecordingRepository(case.sessions, transient_failure=transient_failure) + observed: list[tuple[bool, bool, str]] = [] + + def while_exporting() -> None: + # The original start time is now beyond the lease; only a real background + # transaction can prevent a second claimant from marking it uncertain. + clock[0] += timedelta(seconds=3) + heartbeat_committed = renewals.renewed.wait(3) + second_claim = _claim_at(case, lease_seconds=1, now=clock[0]) + observed.append((heartbeat_committed, second_claim is None, _stored(case).state)) + + ports = ServicePorts( + _SlowDestination([reply], while_exporting), settings=replace(AGENT_TRACE_SETTINGS, export_lease_seconds=1) + ) + ports.service(replace(case, repository=renewals)).deliver(case.run.id) + assert observed == [(True, True, TraceRunState.EXPORTING)] + assert _stored(case).state == outcome + assert renewals.renewal_threads + assert all(thread_id != threading.get_ident() for thread_id in renewals.renewal_threads) + assert not any(thread.name == f"application-trace-lease-{case.run.id}" for thread in threading.enumerate()) + assert "private-provider-value" not in caplog.text + + @staticmethod + def test_heartbeat_stops_even_when_completion_database_write_raises(case: _ExportLeaseCollection) -> None: + unavailable = replace(case, repository=_UnavailableCompletionRepository(case.sessions)) + destination = _export_lease_accepted("ack") + with pytest.raises(TracingStoreUnavailableError, match="database unavailable"): + ServicePorts(destination).service(unavailable).deliver(case.run.id) + # The database is lost only once the destination answered, so the export ran under its heartbeat. + assert len(destination.requests) == 1 + assert not any(thread.name == f"application-trace-lease-{case.run.id}" for thread in threading.enumerate()) + + +_application_trace_ownership_no_egress = _application_trace_collection_no_egress +_application_trace_ownership_seed = _application_trace_collection_seed + + +class TestApplicationTraceOwnershipCases: + @staticmethod + def test_naming_without_a_message_stays_with_its_agent_app(sqlite_session_factory: sessionmaker[Session]) -> None: + case = _application_trace_ownership_seed(sqlite_session_factory, "naming") + service = ServicePorts(_application_trace_ownership_no_egress()).service(case) + record = GenerateNameTraceInfo( + message_id=None, + message_data=None, + inputs="hello", + outputs="Greeting", + start_time=datetime(2026, 1, 1), + end_time=datetime(2026, 1, 1), + metadata={"app_id": case.run.app_id, "conversation_id": "conversation-1", "trace_event_id": "event-1"}, + conversation_id="conversation-1", + tenant_id=case.run.tenant_id, + ) + + assert service.capture_terminal(tenant_id=case.run.tenant_id, info=record) is True + + @staticmethod + def test_runs_frozen_for_an_unavailable_provider_report_no_capabilities( + sqlite_session_factory: sessionmaker[Session], + ) -> None: + case = _application_trace_ownership_seed(sqlite_session_factory, "retired-capabilities") + freeze_config(case, "not-a-provider", FROZEN_LANGFUSE) + result = ( + ServicePorts(_application_trace_ownership_no_egress()) + .service(case) + .status(**source(case.run), node_execution_id=None) + ) + assert result.provider == "not-a-provider" + assert result.capabilities is None diff --git a/api/tests/unit_tests/tasks/test_application_trace_task.py b/api/tests/unit_tests/tasks/test_application_trace_task.py index 0768f6732b262c..da70d0865f7961 100644 --- a/api/tests/unit_tests/tasks/test_application_trace_task.py +++ b/api/tests/unit_tests/tasks/test_application_trace_task.py @@ -1,57 +1,156 @@ """Exercise Celery entrypoints against durable application trace state.""" -from datetime import timedelta +from dataclasses import replace +from datetime import UTC, datetime, timedelta +from uuid import UUID, uuid4 +import httpx import pytest +from dify_trace_sdk.events import MessageCompleted, TraceCorrelation, TraceEvent +from flask import Flask from sqlalchemy.orm import Session, sessionmaker -from core.ops.unified_trace.entities import CanonicalTrace +from core.helper import ssrf_proxy +from domain.tracing.trace_run_agg import TraceRunState +from extensions.application_services.tracing import build_tracing_services +from extensions.ext_application_services import ApplicationServices from libs.datetime_utils import naive_utc_now from models.application_trace import ApplicationTraceRun -from services.tracing import service as service_module -from services.tracing.provider_dispatch import TraceExportResult +from models.enums import ConversationFromSource +from models.model import AppMode, Message +from repositories.factory import DifyAPIRepositoryFactory +from repositories.trace_record_repository import TraceRecordRepository +from services.tracing.canonical import CanonicalTraceBuilder +from services.tracing.collection import AgentTraceSettings +from services.tracing.egress import SSRFEgress +from services.tracing.providers import ProviderConfigCodec +from services.tracing.records import MessageTraceInfo +from services.tracing.registry import EntryPointTraceProviderCatalog from services.tracing.service import ApplicationTracingService from tasks.application_trace_task import export_application_trace, reconcile_application_traces -from tests.helpers.unit.config_override import apply_config_overrides -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_batch_for as batch_for -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_ingest as ingest -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_seed as seed -from tests.unit_tests.services.tracing.test_service import _application_trace_collection_terminal as terminal +from tests.helpers.unit.model_factories import make_conversation, make_tenant +from tests.helpers.unit.services.tracing_ports import AGENT_TRACE_SETTINGS, WorkspaceBoundCipher +from tests.unit_tests.repositories.test_application_trace_repository import ( + Collection, + batch_for, + ingest, + seed, + terminal, +) -def test_worker_task_exports_once_and_replayed_task_observes_checkpoint( +def bind_worker( + app: Flask, + services: ApplicationServices, + collection: Collection, + *, + settings: AgentTraceSettings, + scheduled: list[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Hand the tasks a worker over the seeded ledger with the installed providers and the SSRF egress.""" + worker = ApplicationTracingService( + runs=collection.repository, + run_query=collection.repository, + sources=collection.repository, + records=TraceRecordRepository( + collection.sessions, + workflow_runs=DifyAPIRepositoryFactory.create_api_workflow_run_repository(collection.sessions), + node_executions=DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(collection.sessions), + files_url="https://files.test", + enterprise_details=False, + ), + builder=CanonicalTraceBuilder(), + providers=EntryPointTraceProviderCatalog(), + codec=ProviderConfigCodec(WorkspaceBoundCipher()), + egress=SSRFEgress(), + schedule_export=scheduled.append, + settings=settings, + ) + monkeypatch.setitem( + app.extensions, + "application_services", + replace(services, tracing=replace(services.tracing, agent_traces=worker)), + ) + + +def test_worker_task_exports_once_through_the_ssrf_proxy_and_replay_observes_checkpoint( + app: Flask, + account_application_services: ApplicationServices, sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, ) -> None: case = seed(sqlite_session_factory, "worker-task") ingest(case, batch_for(case.run, case.invocation())) terminal(case) - service = ApplicationTracingService(case.repository) - monkeypatch.setattr(service_module, "application_tracing_service", lambda: service) - monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", lambda *_args: {}) - exports: list[CanonicalTrace] = [] + bind_worker( + app, + account_application_services, + case, + settings=AGENT_TRACE_SETTINGS, + scheduled=[], + monkeypatch=monkeypatch, + ) + proxied: list[tuple[str, str]] = [] - def export(_provider: str, _config: dict[str, object], trace: CanonicalTrace) -> TraceExportResult: - exports.append(trace) - return TraceExportResult("vendor-id", "https://example.test/trace/vendor-id") + def send(method: str, url: str, **_kwargs: object) -> httpx.Response: + proxied.append((method, url)) + return httpx.Response(200, json={}, request=httpx.Request(method, url)) - monkeypatch.setattr(service_module, "export_trace", export) + monkeypatch.setattr(ssrf_proxy, "make_request", send) export_application_trace.run(case.run.id) export_application_trace.run(case.run.id) - assert len(exports) == 1 - assert exports[0].trace_id == case.run.trace_id + assert proxied == [("POST", "https://langfuse.test/api/public/ingestion")] result = case.repository.get(case.run.id) assert result is not None - assert result.state == "accepted" + assert result.state is TraceRunState.ACCEPTED assert result.attempts == 1 - assert result.trace_url == "https://example.test/trace/vendor-id" + provider_trace_id = str(UUID(hex=case.run.trace_id)) + assert result.provider_trace_id == provider_trace_id + assert result.trace_url == f"https://langfuse.test/project/original/traces/{provider_trace_id}" + + +def test_worker_records_carry_no_enterprise_names(sqlite_session_factory: sessionmaker[Session]) -> None: + case = seed(sqlite_session_factory, "worker-names") + with case.sessions.begin() as session: + message = session.get(Message, case.run.source_id) + assert message is not None + conversation_id = message.conversation_id + session.add( + make_conversation( + conversation_id=conversation_id, + app_id=case.run.app_id, + mode=AppMode.AGENT, + inputs={}, + from_source=ConversationFromSource.CONSOLE, + ) + ) + session.add(make_tenant(tenant_id=case.run.tenant_id)) + # The tasks reach this Agent trace service through the composition root; its records feed recovery. + worker = build_tracing_services(database_client=sqlite_session_factory).agent_traces + record = worker._records.assemble( + TraceEvent( + event_id=uuid4().hex, + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, + user_id=None, + occurred_at=datetime.now(UTC), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + payload=MessageCompleted(message_id=case.run.source_id, conversation_id=conversation_id), + ) + ) + assert isinstance(record, MessageTraceInfo) + assert (record.metadata["app_name"], record.metadata["workspace_name"]) == ("", "") def test_reconcile_only_schedules_due_pending_rows_and_applies_retention( + app: Flask, + account_application_services: ApplicationServices, sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, ) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_RETENTION_DAYS=3) due = seed(sqlite_session_factory, "due") deferred = seed(sqlite_session_factory, "future") accepted = seed(sqlite_session_factory, "accepted-expired") @@ -60,14 +159,19 @@ def test_reconcile_only_schedules_due_pending_rows_and_applies_retention( expired = session.get(ApplicationTraceRun, accepted.run.id) assert future is not None assert expired is not None - future.state = "retrying" + future.state = TraceRunState.RETRYING future.next_attempt_at = naive_utc_now() + timedelta(hours=1) - expired.state = "accepted" + expired.state = TraceRunState.ACCEPTED expired.created_at = naive_utc_now() - timedelta(days=4) - service = ApplicationTracingService(due.repository) scheduled: list[str] = [] - monkeypatch.setattr(service, "schedule", scheduled.append) - monkeypatch.setattr(service_module, "application_tracing_service", lambda: service) + bind_worker( + app, + account_application_services, + due, + settings=replace(AGENT_TRACE_SETTINGS, retention_days=3), + scheduled=scheduled, + monkeypatch=monkeypatch, + ) reconcile_application_traces.run() assert scheduled == [due.run.id] assert due.repository.get(accepted.run.id) is None 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 46196ca01f077e..a00d910aa65b6b 100644 --- a/api/tests/unit_tests/tasks/test_ops_trace_task.py +++ b/api/tests/unit_tests/tasks/test_ops_trace_task.py @@ -1,487 +1,37 @@ -import json -import sys -from collections.abc import Generator -from contextlib import AbstractContextManager, contextmanager, nullcontext -from types import ModuleType -from unittest.mock import MagicMock, patch - import pytest -from celery.exceptions import Retry -from flask import Flask - -from configs import dify_config -from core.ops.base_trace_instance import ( - ApplicationTracingHooksNotRegisteredError, - bind_application_tracing_hooks, -) -from core.ops.entities.config_entity import OPS_TRACE_FAILED_KEY -from core.ops.entities.trace_entity import BaseTraceInfo -from core.ops.exceptions import PendingTraceParentContextError, RetryableTraceDispatchError -from services.tracing.provider_dispatch import TraceExportUncertainError -from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks -from tasks.ops_trace_task import process_trace_tasks -from tests.helpers.unit.config_override import apply_config_overrides - - -@pytest.fixture(autouse=True) -def _application_tracing_disabled(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) - - -@contextmanager -def fake_app_context() -> Generator[None]: - yield - - -class FakeCurrentApp: - def app_context(self) -> AbstractContextManager[None]: - return fake_app_context() - - -def _install_trace_manager( - trace_instance: MagicMock, - *, - enterprise_enabled: bool = False, - enterprise_trace_cls: MagicMock | None = None, -) -> dict[str, ModuleType]: - ops_trace_manager_module = ModuleType("core.ops.ops_trace_manager") - - class StubOpsTraceManager: - @staticmethod - def get_ops_trace_instance(app_id: str) -> MagicMock: - return trace_instance - - telemetry_module = ModuleType("extensions.ext_enterprise_telemetry") - telemetry_module.is_enabled = lambda: enterprise_enabled # pyrefly: ignore[missing-attribute] - - ops_trace_manager_module.OpsTraceManager = StubOpsTraceManager # pyrefly: ignore[missing-attribute] - modules = { - "core.ops.ops_trace_manager": ops_trace_manager_module, - "extensions.ext_enterprise_telemetry": telemetry_module, - } - if enterprise_trace_cls is not None: - enterprise_module = ModuleType("enterprise") - enterprise_telemetry_module = ModuleType("enterprise.telemetry") - enterprise_trace_module = ModuleType("enterprise.telemetry.enterprise_trace") - enterprise_trace_module.EnterpriseOtelTrace = enterprise_trace_cls # pyrefly: ignore[missing-attribute] - modules.update( - { - "enterprise": enterprise_module, - "enterprise.telemetry": enterprise_telemetry_module, - "enterprise.telemetry.enterprise_trace": enterprise_trace_module, - } - ) - return modules - - -def _make_payload() -> str: - return json.dumps({"trace_info": {}, "trace_info_type": None}) - - -def _decode_saved_payload(payload: bytes | str) -> dict[str, object]: - if isinstance(payload, bytes): - payload = payload.decode("utf-8") - return json.loads(payload) - - -def _retryable_dispatch_error() -> RetryableTraceDispatchError: - return RetryableTraceDispatchError("transient trace dispatch failure") - - -def _run_task(file_info: dict[str, str], retries: int = 0) -> None: - process_trace_tasks.push_request(retries=retries) # pyrefly: ignore[unexpected-keyword, missing-argument] - try: - process_trace_tasks.run(file_info) - finally: - process_trace_tasks.pop_request() - - -@pytest.mark.parametrize( - "pending_error", - [ - _retryable_dispatch_error(), - PendingTraceParentContextError("parent-node"), - TraceExportUncertainError("DO_NOT_EXPOSE_PROVIDER_BODY"), - ], -) -def test_process_trace_tasks_retries_only_confirmed_retryable_dispatch_failures( - pending_error: RetryableTraceDispatchError, - caplog: pytest.LogCaptureFixture, -) -> None: - file_info = {"app_id": "app-id", "file_id": "file-id"} - trace_instance = MagicMock() - trace_instance.trace.side_effect = pending_error - retry_error = Retry() - - with ( - patch.dict(sys.modules, _install_trace_manager(trace_instance)), - patch("tasks.ops_trace_task.current_app", FakeCurrentApp()), - patch("tasks.ops_trace_task.storage.load", return_value=_make_payload()), - patch("tasks.ops_trace_task.storage.delete") as mock_delete, - patch("tasks.ops_trace_task.redis_client.incr") as mock_incr, - patch.object(process_trace_tasks, "retry", side_effect=retry_error) as mock_retry, - pytest.raises(Retry) if not pending_error.acceptance_unknown else nullcontext(), - ): - _run_task(file_info) - - trace_instance.trace.assert_called_once_with({}) - if pending_error.acceptance_unknown: - mock_retry.assert_not_called() - mock_delete.assert_called_once_with("ops_trace/app-id/file-id.json") - mock_incr.assert_called_once_with(f"{OPS_TRACE_FAILED_KEY}_app-id") - assert "DO_NOT_EXPOSE_PROVIDER_BODY" not in caplog.text - else: - mock_retry.assert_called_once_with( - exc=pending_error, - countdown=process_trace_tasks.default_retry_delay, - ) - mock_delete.assert_not_called() - mock_incr.assert_not_called() - - -def test_process_trace_tasks_marks_enterprise_trace_dispatched_before_retryable_dispatch_retry() -> None: - file_info = {"app_id": "app-id", "file_id": "file-id"} - trace_instance = MagicMock() - pending_error = _retryable_dispatch_error() - trace_instance.trace.side_effect = pending_error - retry_error = Retry() - enterprise_tracer = MagicMock() - enterprise_trace_cls = MagicMock(return_value=enterprise_tracer) - - with ( - patch.dict( - sys.modules, - _install_trace_manager( - trace_instance, - enterprise_enabled=True, - enterprise_trace_cls=enterprise_trace_cls, - ), - ), - patch("tasks.ops_trace_task.current_app", FakeCurrentApp()), - patch("tasks.ops_trace_task.storage.load", return_value=_make_payload()), - patch("tasks.ops_trace_task.storage.save") as mock_save, - patch("tasks.ops_trace_task.storage.delete") as mock_delete, - patch("tasks.ops_trace_task.redis_client.incr") as mock_incr, - patch.object(process_trace_tasks, "retry", side_effect=retry_error) as mock_retry, - pytest.raises(Retry), - ): - _run_task(file_info) - - enterprise_tracer.trace.assert_called_once_with({}) - saved_path, saved_payload = mock_save.call_args.args - assert saved_path == "ops_trace/app-id/file-id.json" - assert _decode_saved_payload(saved_payload)["_enterprise_trace_dispatched"] is True - mock_retry.assert_called_once_with( - exc=pending_error, - countdown=process_trace_tasks.default_retry_delay, - ) - mock_delete.assert_not_called() - mock_incr.assert_not_called() - - -def test_process_trace_tasks_does_not_mark_failed_enterprise_trace_as_dispatched_before_retry() -> None: - file_info = {"app_id": "app-id", "file_id": "file-id"} - trace_instance = MagicMock() - pending_error = _retryable_dispatch_error() - trace_instance.trace.side_effect = pending_error - retry_error = Retry() - enterprise_tracer = MagicMock() - enterprise_tracer.trace.side_effect = RuntimeError("enterprise trace failed") - enterprise_trace_cls = MagicMock(return_value=enterprise_tracer) - - with ( - patch.dict( - sys.modules, - _install_trace_manager( - trace_instance, - enterprise_enabled=True, - enterprise_trace_cls=enterprise_trace_cls, - ), - ), - patch("tasks.ops_trace_task.current_app", FakeCurrentApp()), - patch("tasks.ops_trace_task.storage.load", return_value=_make_payload()), - patch("tasks.ops_trace_task.storage.save") as mock_save, - patch("tasks.ops_trace_task.storage.delete") as mock_delete, - patch("tasks.ops_trace_task.redis_client.incr") as mock_incr, - patch.object(process_trace_tasks, "retry", side_effect=retry_error) as mock_retry, - pytest.raises(Retry), - ): - _run_task(file_info) - - enterprise_tracer.trace.assert_called_once_with({}) - mock_save.assert_not_called() - mock_retry.assert_called_once_with( - exc=pending_error, - countdown=process_trace_tasks.default_retry_delay, - ) - mock_delete.assert_not_called() - mock_incr.assert_not_called() - - -def test_process_trace_tasks_skips_enterprise_trace_when_retry_payload_was_already_dispatched() -> None: - file_info = {"app_id": "app-id", "file_id": "file-id"} - trace_instance = MagicMock() - enterprise_trace_cls = MagicMock() - payload = json.dumps({"trace_info": {}, "trace_info_type": None, "_enterprise_trace_dispatched": True}) - - with ( - patch.dict( - sys.modules, - _install_trace_manager( - trace_instance, - enterprise_enabled=True, - enterprise_trace_cls=enterprise_trace_cls, - ), - ), - patch("tasks.ops_trace_task.current_app", FakeCurrentApp()), - patch("tasks.ops_trace_task.storage.load", return_value=payload), - patch("tasks.ops_trace_task.storage.save") as mock_save, - patch("tasks.ops_trace_task.storage.delete") as mock_delete, - patch("tasks.ops_trace_task.redis_client.incr") as mock_incr, - ): - _run_task(file_info) - - enterprise_trace_cls.assert_not_called() - trace_instance.trace.assert_called_once_with({}) - mock_save.assert_not_called() - mock_delete.assert_called_once_with("ops_trace/app-id/file-id.json") - mock_incr.assert_not_called() - - -def test_process_trace_tasks_default_retry_window_covers_workflow_and_export_grace_period() -> None: - max_retries = process_trace_tasks.max_retries - assert max_retries is not None - assert max_retries * process_trace_tasks.default_retry_delay >= dify_config.WORKFLOW_MAX_EXECUTION_TIME + 300 - - -def test_process_trace_tasks_deletes_payload_on_success() -> None: - file_info = {"app_id": "app-id", "file_id": "file-id"} - trace_instance = MagicMock() - - with ( - patch.dict(sys.modules, _install_trace_manager(trace_instance)), - patch("tasks.ops_trace_task.current_app", FakeCurrentApp()), - patch("tasks.ops_trace_task.storage.load", return_value=_make_payload()), - patch("tasks.ops_trace_task.storage.delete") as mock_delete, - patch("tasks.ops_trace_task.redis_client.incr") as mock_incr, - ): - _run_task(file_info) - - trace_instance.trace.assert_called_once_with({}) - mock_delete.assert_called_once_with("ops_trace/app-id/file-id.json") - mock_incr.assert_not_called() - - -def test_process_trace_tasks_deletes_payload_and_counts_terminal_failure() -> None: - file_info = {"app_id": "app-id", "file_id": "file-id"} - trace_instance = MagicMock() - trace_instance.trace.side_effect = RuntimeError("trace failed") - - with ( - patch.dict(sys.modules, _install_trace_manager(trace_instance)), - patch("tasks.ops_trace_task.current_app", FakeCurrentApp()), - patch("tasks.ops_trace_task.storage.load", return_value=_make_payload()), - patch("tasks.ops_trace_task.storage.delete") as mock_delete, - patch("tasks.ops_trace_task.redis_client.incr") as mock_incr, - ): - _run_task(file_info) - - mock_delete.assert_called_once_with("ops_trace/app-id/file-id.json") - mock_incr.assert_called_once_with(f"{OPS_TRACE_FAILED_KEY}_app-id") - - -def test_process_trace_tasks_treats_retry_enqueue_failure_as_terminal_failure() -> None: - file_info = {"app_id": "app-id", "file_id": "file-id"} - trace_instance = MagicMock() - pending_error = _retryable_dispatch_error() - retry_enqueue_error = RuntimeError("retry enqueue failed") - trace_instance.trace.side_effect = pending_error - - with ( - patch.dict(sys.modules, _install_trace_manager(trace_instance)), - patch("tasks.ops_trace_task.current_app", FakeCurrentApp()), - patch("tasks.ops_trace_task.storage.load", return_value=_make_payload()), - patch("tasks.ops_trace_task.storage.delete") as mock_delete, - patch("tasks.ops_trace_task.redis_client.incr") as mock_incr, - patch.object(process_trace_tasks, "retry", side_effect=retry_enqueue_error) as mock_retry, - ): - _run_task(file_info) - - mock_retry.assert_called_once_with( - exc=pending_error, - countdown=process_trace_tasks.default_retry_delay, - ) - mock_delete.assert_called_once_with("ops_trace/app-id/file-id.json") - mock_incr.assert_called_once_with(f"{OPS_TRACE_FAILED_KEY}_app-id") - - -def test_process_trace_tasks_deletes_payload_and_counts_exhausted_retryable_dispatch_failure() -> None: - file_info = {"app_id": "app-id", "file_id": "file-id"} - trace_instance = MagicMock() - pending_error = _retryable_dispatch_error() - trace_instance.trace.side_effect = pending_error - - with ( - patch.dict(sys.modules, _install_trace_manager(trace_instance)), - patch("tasks.ops_trace_task.current_app", FakeCurrentApp()), - patch("tasks.ops_trace_task.storage.load", return_value=_make_payload()), - patch("tasks.ops_trace_task.storage.delete") as mock_delete, - patch("tasks.ops_trace_task.redis_client.incr") as mock_incr, - patch.object(process_trace_tasks, "retry") as mock_retry, - ): - max_retries = process_trace_tasks.max_retries - assert max_retries is not None - _run_task(file_info, retries=max_retries) - - mock_retry.assert_not_called() - mock_delete.assert_called_once_with("ops_trace/app-id/file-id.json") - mock_incr.assert_called_once_with(f"{OPS_TRACE_FAILED_KEY}_app-id") - - -def test_managed_application_capture_receives_restored_info_and_preserves_enterprise( - monkeypatch: pytest.MonkeyPatch, -) -> None: - from core.ops.entities.trace_entity import MessageTraceInfo - - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - calls: list[BaseTraceInfo] = [] - - def capture(info: BaseTraceInfo) -> bool: - calls.append(info) - return True - - monkeypatch.setattr(RuntimeApplicationTracingHooks, "capture_terminal", staticmethod(capture)) - hooks = RuntimeApplicationTracingHooks() - app = Flask(__name__) - bind_application_tracing_hooks(app, hooks) - enterprise_calls: list[BaseTraceInfo] = [] - - class RecordingEnterpriseTrace: - def trace(self, info: BaseTraceInfo) -> None: - enterprise_calls.append(info) - payload = json.dumps( - { - "trace_info_type": "MessageTraceInfo", - "trace_info": { - "metadata": {}, - "message_id": "message-id", - "conversation_model": "agent", - "conversation_mode": "agent", - "message_tokens": 0, - "answer_tokens": 0, - "total_tokens": 0, - }, - } - ) - with ( - app.app_context(), - patch("extensions.ext_enterprise_telemetry.is_enabled", return_value=True), - patch("enterprise.telemetry.enterprise_trace.EnterpriseOtelTrace", new=RecordingEnterpriseTrace), - patch( - "core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", - side_effect=AssertionError("managed traces must not instantiate legacy providers"), - ) as lookup, - patch("tasks.ops_trace_task.storage.load", return_value=payload), - patch("tasks.ops_trace_task.storage.delete") as deleted, - ): - _run_task({"app_id": "app-id", "file_id": "file-id"}) - assert len(calls) == 1 - restored = calls[0] - assert isinstance(restored, MessageTraceInfo) - assert restored.message_id == "message-id" - assert enterprise_calls == [restored] - lookup.assert_not_called() - deleted.assert_called_once() +from tasks import ops_trace_task -@pytest.mark.parametrize("enabled", [False, True]) -def test_nonmanaged_task_uses_legacy_and_disabled_gate_never_calls_collector( - monkeypatch: pytest.MonkeyPatch, enabled: bool -) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=enabled) - legacy_calls: list[object] = [] +class _Storage: + def __init__(self) -> None: + self.files: dict[str, bytes] = {} - class RecordingLegacyTrace: - def trace(self, info: object) -> None: - legacy_calls.append(info) + def save(self, key: str, data: bytes) -> None: + self.files[key] = data - legacy = RecordingLegacyTrace() - calls: list[BaseTraceInfo] = [] + def delete(self, key: str) -> None: + self.files.pop(key, None) - def capture(info: BaseTraceInfo) -> bool: - calls.append(info) - return False - monkeypatch.setattr(RuntimeApplicationTracingHooks, "capture_terminal", staticmethod(capture)) - hooks = RuntimeApplicationTracingHooks() - app = Flask(__name__) - bind_application_tracing_hooks(app, hooks) - with ( - app.app_context(), - patch("extensions.ext_enterprise_telemetry.is_enabled", return_value=False), - patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=legacy) as lookup, - patch("tasks.ops_trace_task.storage.load", return_value=_make_payload()), - patch("tasks.ops_trace_task.storage.delete"), - ): - _run_task({"app_id": "app-id", "file_id": "file-id"}) - assert calls == ([{}] if enabled else []) - assert legacy_calls == [{}] - lookup.assert_called_once_with("app-id") +@pytest.fixture +def storage(monkeypatch: pytest.MonkeyPatch) -> _Storage: + fake = _Storage() + monkeypatch.setattr(ops_trace_task, "storage", fake) + return fake -def test_capture_persistence_failure_preserves_payload_without_legacy_export(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - pending_error = RetryableTraceDispatchError("database busy") - calls: list[BaseTraceInfo] = [] +def test_legacy_trace_tasks_are_drained(storage: _Storage) -> None: + storage.save("ops_trace/app-1/file-1.json", b"{}") - def capture(info: BaseTraceInfo) -> bool: - calls.append(info) - raise pending_error + ops_trace_task.process_trace_tasks.run({"app_id": "app-1", "file_id": "file-1"}) - monkeypatch.setattr(RuntimeApplicationTracingHooks, "capture_terminal", staticmethod(capture)) - hooks = RuntimeApplicationTracingHooks() - app = Flask(__name__) - bind_application_tracing_hooks(app, hooks) - with ( - app.app_context(), - patch( - "core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", - side_effect=AssertionError("capture failures must not instantiate legacy providers"), - ) as lookup, - patch("tasks.ops_trace_task.storage.load", return_value=_make_payload()), - patch("tasks.ops_trace_task.storage.delete") as deleted, - patch.object(process_trace_tasks, "retry", side_effect=Retry()) as retry, - pytest.raises(Retry), - ): - _run_task({"app_id": "app-id", "file_id": "file-id"}) - lookup.assert_not_called() - deleted.assert_not_called() + assert storage.files == {} - assert calls == [{}] - retry.assert_called_once_with(exc=pending_error, countdown=process_trace_tasks.default_retry_delay) +def test_legacy_trace_tasks_never_delete_outside_their_folder(storage: _Storage) -> None: + storage.save("ops_trace/app-1/file-1.json", b"{}") -def test_unregistered_application_hooks_preserve_payload_without_legacy_export(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - app = Flask(__name__) - with ( - app.app_context(), - patch( - "core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", - side_effect=AssertionError("unregistered hooks must not instantiate legacy providers"), - ) as lookup, - patch("tasks.ops_trace_task.storage.load", return_value=_make_payload()), - patch("tasks.ops_trace_task.storage.delete") as deleted, - patch("tasks.ops_trace_task.redis_client.incr") as failed, - patch.object(process_trace_tasks, "retry", side_effect=Retry()) as retry, - pytest.raises(Retry), - ): - _run_task({"app_id": "app-id", "file_id": "file-id"}) + ops_trace_task.process_trace_tasks.run({"app_id": "..", "file_id": "../secrets"}) - deleted.assert_not_called() - failed.assert_not_called() - retry.assert_called_once() - lookup.assert_not_called() - assert isinstance(retry.call_args.kwargs["exc"], ApplicationTracingHooksNotRegisteredError) - assert retry.call_args.kwargs["countdown"] == process_trace_tasks.default_retry_delay + assert list(storage.files) == ["ops_trace/app-1/file-1.json"] diff --git a/api/tests/unit_tests/tasks/test_trace_delivery_task.py b/api/tests/unit_tests/tasks/test_trace_delivery_task.py new file mode 100644 index 00000000000000..16b8e47b8b708a --- /dev/null +++ b/api/tests/unit_tests/tasks/test_trace_delivery_task.py @@ -0,0 +1,276 @@ +from dataclasses import replace +from datetime import UTC, datetime, timedelta +from typing import override + +import pytest +from celery.exceptions import Retry +from dify_trace_sdk.events import MessageCompleted, TraceCorrelation, TraceEvent +from dify_trace_sdk.testing import RecordingEgress, message_trace +from flask import Flask +from sqlalchemy.orm import Session, sessionmaker + +from enums import DeploymentEdition +from extensions.ext_application_services import build_application_services +from extensions.ext_redis import redis_client +from services.tracing.adapters import payload_key +from services.tracing.canonical import CanonicalTraceBuilder +from services.tracing.delivery import ( + ExportJob, + ExportRetryQueue, + StoredTraceDestination, + TraceCollector, + TraceDeliveryService, + TraceDestinationQuery, + TraceRecordSource, +) +from services.tracing.errors import SourceNotFoundError, SourceNotReadyError +from services.tracing.providers import ProviderConfigCodec +from services.tracing.records import BaseTraceInfo, GenerateNameTraceInfo, WorkflowTraceInfo +from tasks import trace_delivery_task +from tasks.trace_delivery_task import deliver_trace_event, retry_trace_export, sweep_trace_payloads +from tests.helpers.unit.services.tracing_ports import InstalledTraceProviders, WorkspaceBoundCipher + + +class _Storage: + def __init__(self, files: dict[str, bytes]) -> None: + self.files = dict(files) + + def load(self, key: str) -> bytes: + return self.files[key] + + def save(self, key: str, data: bytes) -> None: + self.files[key] = data + + def delete(self, key: str) -> None: + self.files.pop(key, None) + + def scan(self, path: str, files: bool = True) -> list[str]: + keys: list[str] = [key for key in self.files if key.startswith(path)] if files else [] + if not keys: + raise FileNotFoundError(path) + return keys + + +class _Records(TraceRecordSource): + """Answers every assembly with the scripted record or error after recording the event.""" + + def __init__(self) -> None: + self.result: BaseTraceInfo | Exception = SourceNotFoundError() + self.events: list[TraceEvent] = [] + + @override + def assemble(self, event: TraceEvent) -> BaseTraceInfo: + self.events.append(event) + if isinstance(self.result, Exception): + raise self.result + return self.result + + +class _NoDestinations(TraceDestinationQuery): + def __init__(self) -> None: + self.lookups: list[tuple[str, str]] = [] + + @override + def active_destination(self, *, tenant_id: str, app_id: str) -> StoredTraceDestination | None: + self.lookups.append((tenant_id, app_id)) + return None + + +class _NoRetries(ExportRetryQueue): + @override + def schedule(self, job: ExportJob, *, delay_seconds: int) -> None: + raise AssertionError("no send is attempted") + + +class _UnavailableCollector(TraceCollector): + @override + def capture_terminal(self, *, tenant_id: str, info: BaseTraceInfo) -> bool: + raise RuntimeError("ledger unavailable") + + @override + def link_managed_parent(self, *, tenant_id: str, info: WorkflowTraceInfo) -> WorkflowTraceInfo: + return info + + +class _Delivery: + def __init__(self) -> None: + self.records = _Records() + self.destinations = _NoDestinations() + + def service(self, collector: TraceCollector | None) -> TraceDeliveryService: + return TraceDeliveryService( + records=self.records, + destinations=self.destinations, + providers=InstalledTraceProviders(), + codec=ProviderConfigCodec(WorkspaceBoundCipher()), + builder=CanonicalTraceBuilder(), + egress=RecordingEgress([], grpc_target=None), + retries=_NoRetries(), + enterprise_sink=None, + collector=collector, + max_export_attempts=3, + ) + + +class _RetryRequests: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + def __call__(self, **kwargs: object) -> Retry: + self.calls.append(kwargs) + raise Retry("redelivery scheduled", None) + + +def _event() -> TraceEvent: + return TraceEvent( + event_id="event-1", + tenant_id="tenant-1", + app_id="app-1", + user_id=None, + occurred_at=datetime.now(UTC), + correlation=TraceCorrelation( + external_trace_id=None, session_id=None, parent_workflow_run_id=None, parent_node_execution_id=None + ), + payload=MessageCompleted(message_id="message-1", conversation_id=None), + ) + + +@pytest.fixture +def storage(monkeypatch: pytest.MonkeyPatch) -> _Storage: + fake = _Storage({}) + monkeypatch.setattr(trace_delivery_task, "storage", fake) + return fake + + +def _install( + monkeypatch: pytest.MonkeyPatch, + app: Flask, + sessions: sessionmaker[Session], + delivery: _Delivery, + collector: TraceCollector | None, +) -> None: + services = build_application_services( + database_client=sessions, + deployment_edition=DeploymentEdition.COMMUNITY, + initialization_password="", + redis=redis_client, + ) + tracing = replace(services.tracing, delivery=delivery.service(collector)) + monkeypatch.setitem(app.extensions, "application_services", replace(services, tracing=tracing)) + + +@pytest.fixture +def delivery(monkeypatch: pytest.MonkeyPatch, app: Flask, sqlite_session_factory: sessionmaker[Session]) -> _Delivery: + delivery = _Delivery() + _install(monkeypatch, app, sqlite_session_factory, delivery, None) + return delivery + + +@pytest.mark.usefixtures("storage") +def test_inline_events_are_delivered(delivery: _Delivery) -> None: + deliver_trace_event.run({"event": _event().model_dump_json()}) + + assert [event.event_id for event in delivery.records.events] == ["event-1"] + + +def test_stored_events_are_loaded_and_their_payload_removed(delivery: _Delivery, storage: _Storage) -> None: + key = payload_key(datetime.now(UTC), "events", "event-1") + storage.save(key, _event().model_dump_json().encode()) + + deliver_trace_event.run({"payload_key": key}) + + assert len(delivery.records.events) == 1 + assert key not in storage.files + + +def test_unreadable_events_are_dropped(delivery: _Delivery, storage: _Storage) -> None: + storage.save("trace_payloads/2026-01-01/events/bad.json", b"not json") + + deliver_trace_event.run({"payload_key": "trace_payloads/2026-01-01/events/bad.json"}) + + assert delivery.records.events == [] + assert storage.files == {} + + +@pytest.mark.usefixtures("storage") +def test_missing_sources_are_dropped(delivery: _Delivery) -> None: + deliver_trace_event.run({"event": _event().model_dump_json()}) + + assert len(delivery.records.events) == 1 + + +def test_sources_not_ready_are_redelivered_and_keep_their_payload( + monkeypatch: pytest.MonkeyPatch, delivery: _Delivery, storage: _Storage +) -> None: + delivery.records.result = SourceNotReadyError() + retries = _RetryRequests() + monkeypatch.setattr(deliver_trace_event, "retry", retries) + key = payload_key(datetime.now(UTC), "events", "event-1") + storage.save(key, _event().model_dump_json().encode()) + + with pytest.raises(Retry): + deliver_trace_event.run({"payload_key": key}) + + assert retries.calls == [{"args": ({"payload_key": key},), "countdown": 5}] + assert key in storage.files + + +@pytest.mark.usefixtures("storage") +def test_collector_outages_redeliver_without_repeating_enterprise_telemetry( + monkeypatch: pytest.MonkeyPatch, app: Flask, sqlite_session_factory: sessionmaker[Session] +) -> None: + delivery = _Delivery() + delivery.records.result = GenerateNameTraceInfo( + message_id=None, + message_data=None, + inputs="hello", + outputs="Greeting", + start_time=datetime(2026, 1, 1), + end_time=datetime(2026, 1, 1), + metadata={"trace_event_id": "event-1"}, + conversation_id=None, + tenant_id="tenant-1", + ) + _install(monkeypatch, app, sqlite_session_factory, delivery, _UnavailableCollector()) + retries = _RetryRequests() + monkeypatch.setattr(deliver_trace_event, "retry", retries) + envelope = {"event": _event().model_dump_json()} + + with pytest.raises(Retry): + deliver_trace_event.run(envelope) + + assert retries.calls[0]["args"] == ({**envelope, "enterprise_delivered": True},) + + +def test_export_retries_run_the_stored_job_once(delivery: _Delivery, storage: _Storage) -> None: + job = ExportJob(tenant_id="tenant-1", app_id="app-1", provider_id="langfuse", attempt=2, trace=message_trace()) + key = payload_key(datetime.now(UTC), "exports", "job-1") + storage.save(key, job.model_dump_json().encode()) + + retry_trace_export.run(key) + + assert delivery.destinations.lookups == [("tenant-1", "app-1")] + assert key not in storage.files + + +def test_unreadable_export_jobs_are_dropped(delivery: _Delivery, storage: _Storage) -> None: + storage.save("trace_payloads/2026-01-01/exports/bad.json", b"{}") + + retry_trace_export.run("trace_payloads/2026-01-01/exports/bad.json") + + assert delivery.destinations.lookups == [] + assert storage.files == {} + + +def test_sweep_preserves_legacy_payloads_while_old_workers_drain(storage: _Storage) -> None: + today = datetime.now(UTC) + fresh = payload_key(today, "events", "fresh") + expired = payload_key(today - timedelta(days=3), "events", "expired") + for key in (fresh, expired, "trace_payloads/not-a-date/x.json", "ops_trace/app-1/legacy.json", "other/keep.json"): + storage.save(key, b"{}") + + sweep_trace_payloads.run() + + assert sorted(storage.files) == sorted( + [fresh, "trace_payloads/not-a-date/x.json", "ops_trace/app-1/legacy.json", "other/keep.json"] + ) diff --git a/api/tests/unit_tests/test_worker_entrypoint.py b/api/tests/unit_tests/test_worker_entrypoint.py index 32d45d3d1750e0..1233d1349fece5 100644 --- a/api/tests/unit_tests/test_worker_entrypoint.py +++ b/api/tests/unit_tests/test_worker_entrypoint.py @@ -1,14 +1,25 @@ -import re -from pathlib import Path +import pytest -_ENTRYPOINT_PATH = Path(__file__).resolve().parents[2] / "docker" / "entrypoint.sh" -_DEFAULT_QUEUE_ASSIGNMENT = re.compile(r'^\s*DEFAULT_QUEUES="(?P[^"$]+)"$', re.MULTILINE) -_IM_CONTACT_SYNC_QUEUE = "human_input_contact_sync" +from tests.helpers.unit.worker_entrypoint import worker_configuration -def test_default_worker_queue_sets_consume_im_contact_sync_tasks() -> None: - entrypoint = _ENTRYPOINT_PATH.read_text() - default_queue_sets = [match.group("queues").split(",") for match in _DEFAULT_QUEUE_ASSIGNMENT.finditer(entrypoint)] +@pytest.mark.parametrize("edition", ["COMMUNITY", "ENTERPRISE", "CLOUD"]) +def test_default_worker_queue_sets_consume_im_contact_sync_tasks(edition: str) -> None: + selected = worker_configuration({"DEPLOYMENT_EDITION": edition}) + assert selected.returncode == 0, selected.stderr + assert "human_input_contact_sync" in selected.stdout.splitlines()[-1].split(",") - assert len(default_queue_sets) == 2 - assert all(_IM_CONTACT_SYNC_QUEUE in queues for queues in default_queue_sets) + +@pytest.mark.parametrize("edition", ["COMMUNITY", "ENTERPRISE", "CLOUD"]) +@pytest.mark.parametrize("queue", [None, "ops_trace", "tracing-upgrade"]) +def test_default_worker_consumes_configured_trace_queue_without_legacy_queue(edition: str, queue: str | None) -> None: + environment = {"DEPLOYMENT_EDITION": edition} + if queue is not None: + environment["OPS_TRACE_QUEUE"] = queue + selected = worker_configuration(environment) + assert selected.returncode == 0, selected.stderr + queues = selected.stdout.splitlines()[-1].split(",") + expected = queue or "ops_trace_v2" + assert expected in queues + if expected != "ops_trace": + assert "ops_trace" not in queues diff --git a/api/uv.lock b/api/uv.lock index 8e2887ee428f35..1349894803c3d8 100644 --- a/api/uv.lock +++ b/api/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 5 requires-python = "==3.13.*" resolution-markers = [ "sys_platform == 'win32'", @@ -16,6 +16,7 @@ members = [ "dify-trace-langsmith", "dify-trace-mlflow", "dify-trace-opik", + "dify-trace-sdk", "dify-trace-tencent", "dify-trace-weave", "dify-vdb-alibabacloud-mysql", @@ -57,18 +58,6 @@ overrides = [ { name = "wrapt", specifier = ">=2.1.1,<3.0.0" }, ] -[[package]] -name = "abnf" -version = "2.2.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/9d/f2/7b5fac50ee42e8b8d4a098d76743a394546f938c94125adbb93414e5ae7d/abnf-2.2.0.tar.gz", hash = "sha256:433380fd32855bbc60bc7b3d35d40616e21383a32ed1c9b8893d16d9f4a6c2f4", size = 197507, upload-time = "2023-03-17T18:26:24.577Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/30/95/f456ae7928a2f3a913f467d4fd9e662e295dd7349fc58b35f77f6c757a23/abnf-2.2.0-py3-none-any.whl", hash = "sha256:5dc2ae31a84ff454f7de46e08a2a21a442a0e21a092468420587a1590b490d1f", size = 39938, upload-time = "2023-03-17T18:26:22.608Z" }, -] - [[package]] name = "aiofiles" version = "24.1.0" @@ -398,25 +387,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9f/64/2e54428beba8d9992aa478bb8f6de9e4ecaa5f8f513bcfd567ed7fb0262d/apscheduler-3.11.2-py3-none-any.whl", hash = "sha256:ce005177f741409db4e4dd40a7431b76feb856b9dd69d57e0da49d6715bfd26d", size = 64439, upload-time = "2025-12-22T00:39:33.303Z" }, ] -[[package]] -name = "arize-phoenix-otel" -version = "0.15.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "openinference-instrumentation" }, - { name = "openinference-semantic-conventions" }, - { name = "opentelemetry-exporter-otlp" }, - { name = "opentelemetry-proto" }, - { name = "opentelemetry-sdk" }, - { name = "opentelemetry-semantic-conventions" }, - { name = "typing-extensions" }, - { name = "wrapt" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/f5/f0/b254118db28a2a202573472be67cf61f09cb37912bfde45b27ddc1c5b71f/arize_phoenix_otel-0.15.0.tar.gz", hash = "sha256:56c7dae09aaaa80df9e9595b7384c1bd4054b69b6032ab18e3a110a59b488388", size = 20254, upload-time = "2026-03-02T20:19:04.112Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e4/4d/70d9c9d7137cc2e2aad819932172ef13ce21b4e60bf258910b9f15e426af/arize_phoenix_otel-0.15.0-py3-none-any.whl", hash = "sha256:5ff4d03b52d2dbd9c2a234417848f6b171cd220dc3c4020cf3568be84b89b88b", size = 17697, upload-time = "2026-03-02T20:19:03.242Z" }, -] - [[package]] name = "asgiref" version = "3.11.0" @@ -749,11 +719,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/6f/36/a9faa75f61d7b4f77b60cc28eaf649d8912c9834c42457c792a55a1d17b4/boto3_stubs-1.43.103-py3-none-any.whl", hash = "sha256:c2cf91503f29f71d25129853ad76124db5996c28dfe5a573b6789f8ffee71014", size = 71851, upload-time = "2026-09-25T20:44:09.597Z" }, ] -[package.optional-dependencies] -bedrock-runtime = [ - { name = "mypy-boto3-bedrock-runtime" }, -] - [[package]] name = "botocore" version = "1.43.108" @@ -913,15 +878,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cb/0e/02ceeec9a7d6ee63bb596121c2c8e9b3a9e150936f4fbef6ca1943e6137c/cffi-2.0.0-cp313-cp313-win_arm64.whl", hash = "sha256:256f80b80ca3853f90c21b23ee78cd008713787b1b1e93eae9f3d6a7134abd91", size = 177780, upload-time = "2025-09-08T23:23:16.761Z" }, ] -[[package]] -name = "chardet" -version = "5.1.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/41/32/cdc91dcf83849c7385bf8e2a5693d87376536ed000807fa07f5eab33430d/chardet-5.1.0.tar.gz", hash = "sha256:0d62712b956bc154f85fb0a266e2a3c5913c2967e00348701b32411d6def31e5", size = 2069617, upload-time = "2022-12-01T22:34:18.086Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/74/8f/8fc49109009e8d2169d94d72e6b1f4cd45c13d147ba7d6170fb41f22b08f/chardet-5.1.0-py3-none-any.whl", hash = "sha256:362777fb014af596ad31334fde1e8c327dfdb076e1960d1694662d46a6917ab9", size = 199124, upload-time = "2022-12-01T22:34:14.609Z" }, -] - [[package]] name = "charset-normalizer" version = "3.4.7" @@ -995,15 +951,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5f/7a/10bf5dc92d13cc03230190fcc5016a0b138d99e5b36b8b89ee0fe1680e10/chromadb-0.5.20-py3-none-any.whl", hash = "sha256:9550ba1b6dce911e35cac2568b301badf4b42f457b99a432bdeec2b6b9dd3680", size = 617884, upload-time = "2024-11-19T05:13:56.29Z" }, ] -[[package]] -name = "cint" -version = "1.0.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/3e/c8/3ae22fa142be0bf9eee856e90c314f4144dfae376cc5e3e55b9a169670fb/cint-1.0.0.tar.gz", hash = "sha256:66f026d28c46ef9ea9635be5cb342506c6a1af80d11cb1c881a8898ca429fc91", size = 4641, upload-time = "2019-03-19T01:07:48.723Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/91/c2/898e59963084e1e2cbd4aad1dee92c5bd7a79d121dcff1e659c2a0c2174e/cint-1.0.0-py3-none-any.whl", hash = "sha256:8aa33028e04015711c0305f918cb278f1dc8c5c9997acdc45efad2c7cb1abf50", size = 5573, upload-time = "2019-03-19T01:07:46.496Z" }, -] - [[package]] name = "click" version = "8.4.2" @@ -1474,6 +1421,7 @@ docs = [ name = "dify-api" version = "1.18.0" source = { virtual = "." } +default-groups = ["kms", "storage", "tools", "trace-all", "vdb-all"] dependencies = [ { name = "alibabacloud-dingtalk" }, { name = "aliyun-log-python-sdk" }, @@ -1485,6 +1433,7 @@ dependencies = [ { name = "celery" }, { name = "croniter" }, { name = "dify-agent" }, + { name = "dify-trace-sdk" }, { name = "email-validator" }, { name = "fastopenapi", extra = ["flask"] }, { name = "flask" }, @@ -1507,6 +1456,8 @@ dependencies = [ { name = "lark-oapi" }, { name = "markdown-it-py" }, { name = "opentelemetry-distro" }, + { name = "opentelemetry-exporter-otlp-proto-grpc" }, + { name = "opentelemetry-exporter-otlp-proto-http" }, { name = "opentelemetry-instrumentation-celery" }, { name = "opentelemetry-instrumentation-flask" }, { name = "opentelemetry-instrumentation-httpx" }, @@ -1515,12 +1466,14 @@ dependencies = [ { name = "opentelemetry-propagator-b3" }, { name = "psycogreen" }, { name = "psycopg2-binary" }, + { name = "pydantic-settings" }, { name = "pyjwt" }, { name = "python-socketio" }, { name = "readabilipy" }, { name = "redis", extra = ["hiredis"] }, { name = "resend" }, { name = "sendgrid" }, + { name = "sentry-sdk" }, { name = "slack-sdk" }, { name = "sseclient-py" }, { name = "wechatpy" }, @@ -1786,6 +1739,7 @@ requires-dist = [ { name = "celery", specifier = ">=5.6.3,<6.0.0" }, { name = "croniter", specifier = ">=6.2.2,<7.0.0" }, { name = "dify-agent", editable = "../dify-agent" }, + { name = "dify-trace-sdk", editable = "providers/trace/trace-sdk" }, { name = "email-validator", specifier = ">=2.3.0,<3.0.0" }, { name = "fastopenapi", extras = ["flask"], specifier = "==0.7.0" }, { name = "flask", specifier = ">=3.1.3,<4.0.0" }, @@ -1808,6 +1762,8 @@ requires-dist = [ { name = "lark-oapi", specifier = "==1.7.2" }, { name = "markdown-it-py", specifier = ">=4.0.0,<5.0.0" }, { name = "opentelemetry-distro", specifier = "==0.65b0" }, + { name = "opentelemetry-exporter-otlp-proto-grpc", specifier = ">=1.44.0,<2.0.0" }, + { name = "opentelemetry-exporter-otlp-proto-http", specifier = ">=1.44.0,<2.0.0" }, { name = "opentelemetry-instrumentation-celery", specifier = "==0.65b0" }, { name = "opentelemetry-instrumentation-flask", specifier = "==0.65b0" }, { name = "opentelemetry-instrumentation-httpx", specifier = "==0.65b0" }, @@ -1816,12 +1772,14 @@ requires-dist = [ { name = "opentelemetry-propagator-b3", specifier = ">=1.45.0,<2.0.0" }, { name = "psycogreen", specifier = ">=1.0.2,<2.0.0" }, { name = "psycopg2-binary", specifier = ">=2.9.12,<3.0.0" }, + { name = "pydantic-settings", specifier = ">=2.14.2,<3.0.0" }, { name = "pyjwt", specifier = ">=2.10.1,<3.0.0" }, { name = "python-socketio", specifier = ">=5.13.0,<6.0.0" }, { name = "readabilipy", specifier = "==0.3.0" }, { name = "redis", extras = ["hiredis"], specifier = ">=7.4.0,<8.0.0" }, { name = "resend", specifier = ">=2.27.0,<3.0.0" }, { name = "sendgrid", specifier = ">=6.12.5,<7.0.0" }, + { name = "sentry-sdk", specifier = ">=2.57.0,<3.0.0" }, { name = "slack-sdk", specifier = ">=3.43.0,<4.0.0" }, { name = "sseclient-py", specifier = ">=1.8.0,<2.0.0" }, { name = "wechatpy", specifier = ">=1.8.18,<2.0.0" }, @@ -2000,104 +1958,119 @@ name = "dify-trace-aliyun" version = "0.0.1" source = { editable = "providers/trace/trace-aliyun" } dependencies = [ - { name = "opentelemetry-api" }, - { name = "opentelemetry-exporter-otlp-proto-grpc" }, - { name = "opentelemetry-sdk" }, - { name = "opentelemetry-semantic-conventions" }, + { name = "dify-trace-sdk" }, ] [package.metadata] -requires-dist = [ - { name = "opentelemetry-api" }, - { name = "opentelemetry-exporter-otlp-proto-grpc" }, - { name = "opentelemetry-sdk" }, - { name = "opentelemetry-semantic-conventions" }, -] +requires-dist = [{ name = "dify-trace-sdk", editable = "providers/trace/trace-sdk" }] [[package]] name = "dify-trace-arize-phoenix" version = "0.0.1" source = { editable = "providers/trace/trace-arize-phoenix" } dependencies = [ - { name = "arize-phoenix-otel" }, + { name = "dify-trace-sdk" }, ] [package.metadata] -requires-dist = [{ name = "arize-phoenix-otel", specifier = "==0.15.0" }] +requires-dist = [{ name = "dify-trace-sdk", editable = "providers/trace/trace-sdk" }] [[package]] name = "dify-trace-langfuse" version = "0.0.1" source = { editable = "providers/trace/trace-langfuse" } dependencies = [ - { name = "langfuse" }, + { name = "dify-trace-sdk" }, ] [package.metadata] -requires-dist = [{ name = "langfuse", specifier = ">=4.2.0,<5.0.0" }] +requires-dist = [{ name = "dify-trace-sdk", editable = "providers/trace/trace-sdk" }] [[package]] name = "dify-trace-langsmith" version = "0.0.1" source = { editable = "providers/trace/trace-langsmith" } dependencies = [ + { name = "dify-trace-sdk" }, { name = "langsmith" }, ] [package.metadata] -requires-dist = [{ name = "langsmith", specifier = "==0.8.18" }] +requires-dist = [ + { name = "dify-trace-sdk", editable = "providers/trace/trace-sdk" }, + { name = "langsmith", specifier = "==0.8.18" }, +] [[package]] name = "dify-trace-mlflow" version = "0.0.1" source = { editable = "providers/trace/trace-mlflow" } dependencies = [ + { name = "dify-trace-sdk" }, { name = "mlflow-skinny" }, ] [package.metadata] -requires-dist = [{ name = "mlflow-skinny", specifier = ">=3.11.1,<4.0.0" }] +requires-dist = [ + { name = "dify-trace-sdk", editable = "providers/trace/trace-sdk" }, + { name = "mlflow-skinny", specifier = ">=3.11.1,<4.0.0" }, +] [[package]] name = "dify-trace-opik" version = "0.0.1" source = { editable = "providers/trace/trace-opik" } dependencies = [ - { name = "opik" }, + { name = "dify-trace-sdk" }, ] [package.metadata] -requires-dist = [{ name = "opik", specifier = "~=1.11.2" }] +requires-dist = [{ name = "dify-trace-sdk", editable = "providers/trace/trace-sdk" }] [[package]] -name = "dify-trace-tencent" +name = "dify-trace-sdk" version = "0.0.1" -source = { editable = "providers/trace/trace-tencent" } +source = { editable = "providers/trace/trace-sdk" } dependencies = [ - { name = "opentelemetry-api" }, - { name = "opentelemetry-exporter-otlp-proto-grpc" }, - { name = "opentelemetry-sdk" }, - { name = "opentelemetry-semantic-conventions" }, + { name = "grpcio" }, + { name = "httpx" }, + { name = "opentelemetry-proto" }, + { name = "protobuf" }, + { name = "pydantic" }, + { name = "requests" }, ] [package.metadata] requires-dist = [ - { name = "opentelemetry-api" }, - { name = "opentelemetry-exporter-otlp-proto-grpc" }, - { name = "opentelemetry-sdk" }, - { name = "opentelemetry-semantic-conventions" }, + { name = "grpcio" }, + { name = "httpx" }, + { name = "opentelemetry-proto" }, + { name = "protobuf" }, + { name = "pydantic", specifier = ">=2" }, + { name = "requests" }, ] +[[package]] +name = "dify-trace-tencent" +version = "0.0.1" +source = { editable = "providers/trace/trace-tencent" } +dependencies = [ + { name = "dify-trace-sdk" }, +] + +[package.metadata] +requires-dist = [{ name = "dify-trace-sdk", editable = "providers/trace/trace-sdk" }] + [[package]] name = "dify-trace-weave" version = "0.0.1" source = { editable = "providers/trace/trace-weave" } dependencies = [ - { name = "weave" }, + { name = "dify-trace-sdk" }, ] [package.metadata] -requires-dist = [{ name = "weave", specifier = "==0.52.36" }] +requires-dist = [{ name = "dify-trace-sdk", editable = "providers/trace/trace-sdk" }] [[package]] name = "dify-vdb-alibabacloud-mysql" @@ -2425,15 +2398,6 @@ dependencies = [ [package.metadata] requires-dist = [{ name = "weaviate-client", specifier = "==4.22.0" }] -[[package]] -name = "diskcache-weave" -version = "5.6.3.post1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/a6/52/634e1f43486489fdaded1a7c9bd3524b7e0ca9bcc43af426afa511c541e2/diskcache_weave-5.6.3.post1.tar.gz", hash = "sha256:1fe7e648d1d85d517c05b296f1692e7c425a71714dc31a4b7a584a8f8f5604f2", size = 68297, upload-time = "2026-03-19T14:57:54.299Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/d9/8d/92887441bc338fb8d0b8ea75eb0392c00e20a85ec0bbe02f273188849568/diskcache_weave-5.6.3.post1-py3-none-any.whl", hash = "sha256:b00e9842b74eeecf314456f9c833a6d4f7792ed12b20297b4d3b9df7859ee66f", size = 45905, upload-time = "2026-03-19T14:57:52.819Z" }, -] - [[package]] name = "distro" version = "1.9.0" @@ -2629,34 +2593,6 @@ flask = [ { name = "flask" }, ] -[[package]] -name = "fastuuid" -version = "0.14.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/c3/7d/d9daedf0f2ebcacd20d599928f8913e9d2aea1d56d2d355a93bfa2b611d7/fastuuid-0.14.0.tar.gz", hash = "sha256:178947fc2f995b38497a74172adee64fdeb8b7ec18f2a5934d037641ba265d26", size = 18232, upload-time = "2025-10-19T22:19:22.402Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a5/83/ae12dd39b9a39b55d7f90abb8971f1a5f3c321fd72d5aa83f90dc67fe9ed/fastuuid-0.14.0-cp313-cp313-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:77a09cb7427e7af74c594e409f7731a0cf887221de2f698e1ca0ebf0f3139021", size = 510720, upload-time = "2025-10-19T22:42:34.633Z" }, - { url = "https://files.pythonhosted.org/packages/53/b0/a4b03ff5d00f563cc7546b933c28cb3f2a07344b2aec5834e874f7d44143/fastuuid-0.14.0-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:9bd57289daf7b153bfa3e8013446aa144ce5e8c825e9e366d455155ede5ea2dc", size = 262024, upload-time = "2025-10-19T22:30:25.482Z" }, - { url = "https://files.pythonhosted.org/packages/9c/6d/64aee0a0f6a58eeabadd582e55d0d7d70258ffdd01d093b30c53d668303b/fastuuid-0.14.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:ac60fc860cdf3c3f327374db87ab8e064c86566ca8c49d2e30df15eda1b0c2d5", size = 251679, upload-time = "2025-10-19T22:36:14.096Z" }, - { url = "https://files.pythonhosted.org/packages/60/f5/a7e9cda8369e4f7919d36552db9b2ae21db7915083bc6336f1b0082c8b2e/fastuuid-0.14.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ab32f74bd56565b186f036e33129da77db8be09178cd2f5206a5d4035fb2a23f", size = 277862, upload-time = "2025-10-19T22:36:23.302Z" }, - { url = "https://files.pythonhosted.org/packages/f0/d3/8ce11827c783affffd5bd4d6378b28eb6cc6d2ddf41474006b8d62e7448e/fastuuid-0.14.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:33e678459cf4addaedd9936bbb038e35b3f6b2061330fd8f2f6a1d80414c0f87", size = 278278, upload-time = "2025-10-19T22:29:43.809Z" }, - { url = "https://files.pythonhosted.org/packages/a2/51/680fb6352d0bbade04036da46264a8001f74b7484e2fd1f4da9e3db1c666/fastuuid-0.14.0-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:1e3cc56742f76cd25ecb98e4b82a25f978ccffba02e4bdce8aba857b6d85d87b", size = 301788, upload-time = "2025-10-19T22:36:06.825Z" }, - { url = "https://files.pythonhosted.org/packages/fa/7c/2014b5785bd8ebdab04ec857635ebd84d5ee4950186a577db9eff0fb8ff6/fastuuid-0.14.0-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:cb9a030f609194b679e1660f7e32733b7a0f332d519c5d5a6a0a580991290022", size = 459819, upload-time = "2025-10-19T22:35:31.623Z" }, - { url = "https://files.pythonhosted.org/packages/01/d2/524d4ceeba9160e7a9bc2ea3e8f4ccf1ad78f3bde34090ca0c51f09a5e91/fastuuid-0.14.0-cp313-cp313-musllinux_1_1_i686.whl", hash = "sha256:09098762aad4f8da3a888eb9ae01c84430c907a297b97166b8abc07b640f2995", size = 478546, upload-time = "2025-10-19T22:26:03.023Z" }, - { url = "https://files.pythonhosted.org/packages/bc/17/354d04951ce114bf4afc78e27a18cfbd6ee319ab1829c2d5fb5e94063ac6/fastuuid-0.14.0-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:1383fff584fa249b16329a059c68ad45d030d5a4b70fb7c73a08d98fd53bcdab", size = 450921, upload-time = "2025-10-19T22:31:02.151Z" }, - { url = "https://files.pythonhosted.org/packages/fb/be/d7be8670151d16d88f15bb121c5b66cdb5ea6a0c2a362d0dcf30276ade53/fastuuid-0.14.0-cp313-cp313-win32.whl", hash = "sha256:a0809f8cc5731c066c909047f9a314d5f536c871a7a22e815cc4967c110ac9ad", size = 154559, upload-time = "2025-10-19T22:36:36.011Z" }, - { url = "https://files.pythonhosted.org/packages/22/1d/5573ef3624ceb7abf4a46073d3554e37191c868abc3aecd5289a72f9810a/fastuuid-0.14.0-cp313-cp313-win_amd64.whl", hash = "sha256:0df14e92e7ad3276327631c9e7cec09e32572ce82089c55cb1bb8df71cf394ed", size = 156539, upload-time = "2025-10-19T22:33:35.898Z" }, -] - -[[package]] -name = "fickling" -version = "0.1.12" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/d7/20/d3c2bdb9235b777763a4afc7cc3673afcd162a707ec0988ae7141a540802/fickling-0.1.12.tar.gz", hash = "sha256:83f6ccc948e21edb9ebd92795069536b47f481ce6add62598eac608b31576821", size = 357026, upload-time = "2026-06-26T23:55:57.829Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ba/c0/c65003b71abc4ffddbae99ae86952a21ff859d81b0f1331f79239fe5a58e/fickling-0.1.12-py3-none-any.whl", hash = "sha256:6232b72857e6ee9d729922811b681132d49dc39abc896581390dad1c0eada814", size = 58960, upload-time = "2026-06-26T23:55:56.401Z" }, -] - [[package]] name = "filelock" version = "3.20.3" @@ -3184,26 +3120,6 @@ grpc = [ { name = "grpcio" }, ] -[[package]] -name = "gql" -version = "4.0.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "anyio" }, - { name = "backoff" }, - { name = "graphql-core" }, - { name = "yarl" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/06/9f/cf224a88ed71eb223b7aa0b9ff0aa10d7ecc9a4acdca2279eb046c26d5dc/gql-4.0.0.tar.gz", hash = "sha256:f22980844eb6a7c0266ffc70f111b9c7e7c7c13da38c3b439afc7eab3d7c9c8e", size = 215644, upload-time = "2025-08-17T14:32:35.397Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ac/94/30bbd09e8d45339fa77a48f5778d74d47e9242c11b3cd1093b3d994770a5/gql-4.0.0-py3-none-any.whl", hash = "sha256:f3beed7c531218eb24d97cb7df031b4a84fdb462f4a2beb86e2633d395937479", size = 89900, upload-time = "2025-08-17T14:32:34.029Z" }, -] - -[package.optional-dependencies] -httpx = [ - { name = "httpx" }, -] - [[package]] name = "graphon" version = "0.8.0" @@ -3230,24 +3146,6 @@ dependencies = [ { name = "webvtt-py" }, ] -[[package]] -name = "graphql-core" -version = "3.2.13" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/4e/5e/aa0d4e701b50db0bab71b125dd19ddb0f98c638d008d6f7e3d8ce9cbc92e/graphql_core-3.2.13.tar.gz", hash = "sha256:bb81dd266d4ab7b591bd976f1b23639d97776cb9ac1a896b4a93c271e11ed618", size = 577471, upload-time = "2026-09-27T12:22:57.136Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1f/b4/b52a324b0b297ddc49b8ea31a51d39fae63132f694a267b9381d94e2e0af/graphql_core-3.2.13-py3-none-any.whl", hash = "sha256:b0eb04f2c31556b2310a77c8fb53c74e8b56570f8fea594d89c6ef7827dbb497", size = 268191, upload-time = "2026-09-27T12:22:55.661Z" }, -] - -[[package]] -name = "graphviz" -version = "0.21" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/f8/b3/3ac91e9be6b761a4b30d66ff165e54439dcd48b83f4e20d644867215f6ca/graphviz-0.21.tar.gz", hash = "sha256:20743e7183be82aaaa8ad6c93f8893c923bd6658a04c32ee115edb3c8a835f78", size = 200434, upload-time = "2025-06-15T09:35:05.824Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/91/4c/e0ce1ef95d4000ebc1c11801f9b944fa5910ecc15b5e351865763d8657f8/graphviz-0.21-py3-none-any.whl", hash = "sha256:54f33de9f4f911d7e84e4191749cac8cc5653f815b06738c54db9a15ab8b1e42", size = 47300, upload-time = "2025-06-15T09:35:04.433Z" }, -] - [[package]] name = "greenlet" version = "3.2.4" @@ -3501,8 +3399,8 @@ name = "httpcore2" version = "2.12.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "h11", marker = "sys_platform != 'emscripten'" }, - { name = "truststore", marker = "sys_platform != 'emscripten'" }, + { name = "h11" }, + { name = "truststore" }, ] sdist = { url = "https://files.pythonhosted.org/packages/be/ad/f4f0e57345f1870f3e8cb624e058d7eca6e5a27d33bcc3311d9b618734cd/httpcore2-2.12.0.tar.gz", hash = "sha256:9293522bba0aa7c4c8e9e3f040c16575bd8868e155a77fa30c7a9085a5eae648", size = 67548, upload-time = "2026-08-18T13:22:08.211Z" } wheels = [ @@ -3763,15 +3661,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/96/b3/97b7da76d63e0b7e119872db6f85ae19371a39fe57f037c9fd08fd4e07af/intersystems_irispython-5.4.0-cp39.cp310.cp311.cp312.cp313.cp314-cp39.cp310.cp311.cp312.cp313.cp314-win_amd64.whl", hash = "sha256:97f3568c351de3f5b6cff088f104990b728b3243faab5947ed11357cd646d671", size = 3704377, upload-time = "2026-07-30T19:07:00.885Z" }, ] -[[package]] -name = "intervaltree" -version = "3.1.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "sortedcontainers" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/50/fb/396d568039d21344639db96d940d40eb62befe704ef849b27949ded5c3bb/intervaltree-3.1.0.tar.gz", hash = "sha256:902b1b88936918f9b2a19e0e5eb7ccb430ae45cde4f39ea4b36932920d33952d", size = 32861, upload-time = "2020-08-03T08:01:11.392Z" } - [[package]] name = "isodate" version = "0.7.2" @@ -3808,32 +3697,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, ] -[[package]] -name = "jiter" -version = "0.12.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/45/9d/e0660989c1370e25848bb4c52d061c71837239738ad937e83edca174c273/jiter-0.12.0.tar.gz", hash = "sha256:64dfcd7d5c168b38d3f9f8bba7fc639edb3418abcc74f22fdbe6b8938293f30b", size = 168294, upload-time = "2025-11-09T20:49:23.302Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/3d/a6/97209693b177716e22576ee1161674d1d58029eb178e01866a0422b69224/jiter-0.12.0-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:6cc49d5130a14b732e0612bc76ae8db3b49898732223ef8b7599aa8d9810683e", size = 313658, upload-time = "2025-11-09T20:47:44.424Z" }, - { url = "https://files.pythonhosted.org/packages/06/4d/125c5c1537c7d8ee73ad3d530a442d6c619714b95027143f1b61c0b4dfe0/jiter-0.12.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:37f27a32ce36364d2fa4f7fdc507279db604d27d239ea2e044c8f148410defe1", size = 318605, upload-time = "2025-11-09T20:47:45.973Z" }, - { url = "https://files.pythonhosted.org/packages/99/bf/a840b89847885064c41a5f52de6e312e91fa84a520848ee56c97e4fa0205/jiter-0.12.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bbc0944aa3d4b4773e348cda635252824a78f4ba44328e042ef1ff3f6080d1cf", size = 349803, upload-time = "2025-11-09T20:47:47.535Z" }, - { url = "https://files.pythonhosted.org/packages/8a/88/e63441c28e0db50e305ae23e19c1d8fae012d78ed55365da392c1f34b09c/jiter-0.12.0-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:da25c62d4ee1ffbacb97fac6dfe4dcd6759ebdc9015991e92a6eae5816287f44", size = 365120, upload-time = "2025-11-09T20:47:49.284Z" }, - { url = "https://files.pythonhosted.org/packages/0a/7c/49b02714af4343970eb8aca63396bc1c82fa01197dbb1e9b0d274b550d4e/jiter-0.12.0-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:048485c654b838140b007390b8182ba9774621103bd4d77c9c3f6f117474ba45", size = 479918, upload-time = "2025-11-09T20:47:50.807Z" }, - { url = "https://files.pythonhosted.org/packages/69/ba/0a809817fdd5a1db80490b9150645f3aae16afad166960bcd562be194f3b/jiter-0.12.0-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:635e737fbb7315bef0037c19b88b799143d2d7d3507e61a76751025226b3ac87", size = 379008, upload-time = "2025-11-09T20:47:52.211Z" }, - { url = "https://files.pythonhosted.org/packages/5f/c3/c9fc0232e736c8877d9e6d83d6eeb0ba4e90c6c073835cc2e8f73fdeef51/jiter-0.12.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4e017c417b1ebda911bd13b1e40612704b1f5420e30695112efdbed8a4b389ed", size = 361785, upload-time = "2025-11-09T20:47:53.512Z" }, - { url = "https://files.pythonhosted.org/packages/96/61/61f69b7e442e97ca6cd53086ddc1cf59fb830549bc72c0a293713a60c525/jiter-0.12.0-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:89b0bfb8b2bf2351fba36bb211ef8bfceba73ef58e7f0c68fb67b5a2795ca2f9", size = 386108, upload-time = "2025-11-09T20:47:54.893Z" }, - { url = "https://files.pythonhosted.org/packages/e9/2e/76bb3332f28550c8f1eba3bf6e5efe211efda0ddbbaf24976bc7078d42a5/jiter-0.12.0-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:f5aa5427a629a824a543672778c9ce0c5e556550d1569bb6ea28a85015287626", size = 519937, upload-time = "2025-11-09T20:47:56.253Z" }, - { url = "https://files.pythonhosted.org/packages/84/d6/fa96efa87dc8bff2094fb947f51f66368fa56d8d4fc9e77b25d7fbb23375/jiter-0.12.0-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:ed53b3d6acbcb0fd0b90f20c7cb3b24c357fe82a3518934d4edfa8c6898e498c", size = 510853, upload-time = "2025-11-09T20:47:58.32Z" }, - { url = "https://files.pythonhosted.org/packages/8a/28/93f67fdb4d5904a708119a6ab58a8f1ec226ff10a94a282e0215402a8462/jiter-0.12.0-cp313-cp313-win32.whl", hash = "sha256:4747de73d6b8c78f2e253a2787930f4fffc68da7fa319739f57437f95963c4de", size = 204699, upload-time = "2025-11-09T20:47:59.686Z" }, - { url = "https://files.pythonhosted.org/packages/c4/1f/30b0eb087045a0abe2a5c9c0c0c8da110875a1d3be83afd4a9a4e548be3c/jiter-0.12.0-cp313-cp313-win_amd64.whl", hash = "sha256:e25012eb0c456fcc13354255d0338cd5397cce26c77b2832b3c4e2e255ea5d9a", size = 204258, upload-time = "2025-11-09T20:48:01.01Z" }, - { url = "https://files.pythonhosted.org/packages/2c/f4/2b4daf99b96bce6fc47971890b14b2a36aef88d7beb9f057fafa032c6141/jiter-0.12.0-cp313-cp313-win_arm64.whl", hash = "sha256:c97b92c54fe6110138c872add030a1f99aea2401ddcdaa21edf74705a646dd60", size = 185503, upload-time = "2025-11-09T20:48:02.35Z" }, - { url = "https://files.pythonhosted.org/packages/39/ca/67bb15a7061d6fe20b9b2a2fd783e296a1e0f93468252c093481a2f00efa/jiter-0.12.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:53839b35a38f56b8be26a7851a48b89bc47e5d88e900929df10ed93b95fea3d6", size = 317965, upload-time = "2025-11-09T20:48:03.783Z" }, - { url = "https://files.pythonhosted.org/packages/18/af/1788031cd22e29c3b14bc6ca80b16a39a0b10e611367ffd480c06a259831/jiter-0.12.0-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:94f669548e55c91ab47fef8bddd9c954dab1938644e715ea49d7e117015110a4", size = 345831, upload-time = "2025-11-09T20:48:05.55Z" }, - { url = "https://files.pythonhosted.org/packages/05/17/710bf8472d1dff0d3caf4ced6031060091c1320f84ee7d5dcbed1f352417/jiter-0.12.0-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:351d54f2b09a41600ffea43d081522d792e81dcfb915f6d2d242744c1cc48beb", size = 361272, upload-time = "2025-11-09T20:48:06.951Z" }, - { url = "https://files.pythonhosted.org/packages/fb/f1/1dcc4618b59761fef92d10bcbb0b038b5160be653b003651566a185f1a5c/jiter-0.12.0-cp313-cp313t-win_amd64.whl", hash = "sha256:2a5e90604620f94bf62264e7c2c038704d38217b7465b863896c6d7c902b06c7", size = 204604, upload-time = "2025-11-09T20:48:08.328Z" }, - { url = "https://files.pythonhosted.org/packages/d9/32/63cb1d9f1c5c6632a783c0052cde9ef7ba82688f7065e2f0d5f10a7e3edb/jiter-0.12.0-cp313-cp313t-win_arm64.whl", hash = "sha256:88ef757017e78d2860f96250f9393b7b577b06a956ad102c29c8237554380db3", size = 185628, upload-time = "2025-11-09T20:48:09.572Z" }, -] - [[package]] name = "jmespath" version = "0.10.0" @@ -3900,15 +3763,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/41/45/1a4ed80516f02155c51f51e8cedb3c1902296743db0bbc66608a0db2814f/jsonschema_specifications-2025.9.1-py3-none-any.whl", hash = "sha256:98802fee3a11ee76ecaca44429fda8a41bff98b00a0f2838151b113f210cc6fe", size = 18437, upload-time = "2025-09-08T01:34:57.871Z" }, ] -[[package]] -name = "kaitaistruct" -version = "0.11" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/27/b8/ca7319556912f68832daa4b81425314857ec08dfccd8dbc8c0f65c992108/kaitaistruct-0.11.tar.gz", hash = "sha256:053ee764288e78b8e53acf748e9733268acbd579b8d82a427b1805453625d74b", size = 11519, upload-time = "2025-09-08T15:46:25.037Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/4a/4a/cf14bf3b1f5ffb13c69cf5f0ea78031247790558ee88984a8bdd22fae60d/kaitaistruct-0.11-py2.py3-none-any.whl", hash = "sha256:5c6ce79177b4e193a577ecd359e26516d1d6d000a0bffd6e1010f2a46a62a561", size = 11372, upload-time = "2025-09-08T15:46:23.635Z" }, -] - [[package]] name = "kombu" version = "5.6.2" @@ -3955,25 +3809,6 @@ dependencies = [ ] sdist = { url = "https://files.pythonhosted.org/packages/0e/72/a3add0e4eec4eb9e2569554f7c70f4a3c27712f40e3284d483e88094cc0e/langdetect-1.0.9.tar.gz", hash = "sha256:cbc1fef89f8d062739774bd51eda3da3274006b3661d199c2655f6b3f6d605a0", size = 981474, upload-time = "2021-05-07T07:54:13.562Z" } -[[package]] -name = "langfuse" -version = "4.2.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "backoff" }, - { name = "httpx" }, - { name = "opentelemetry-api" }, - { name = "opentelemetry-exporter-otlp-proto-http" }, - { name = "opentelemetry-sdk" }, - { name = "packaging" }, - { name = "pydantic" }, - { name = "wrapt" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/45/9c/b912a00ffae92ff9955cdd9b74fb839be58f631d4329ae2a8a0376f697f2/langfuse-4.2.0.tar.gz", hash = "sha256:d0bd26d5065cf6a59d7d1093b08d8910e2458dc3da7ed8ccec160db114c18342", size = 275582, upload-time = "2026-04-10T11:55:25.21Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/be/0a/b84e3e68a690ccfe6d64953c572772c685fcb0915b7f2ee3a87c22e388ab/langfuse-4.2.0-py3-none-any.whl", hash = "sha256:bfd760bf10fd0228f297f6369436620f76d16b589de46393d65706b27e4e4082", size = 475449, upload-time = "2026-04-10T11:55:23.624Z" }, -] - [[package]] name = "langsmith" version = "0.8.18" @@ -4020,37 +3855,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8b/34/151efd46acca34d086a2db86d1820ec2452a72d1a8281bb73d5991f03505/lark_oapi-1.7.2-py3-none-any.whl", hash = "sha256:d73cbcac28d8e7cd75efe64a8c350da50b1a45be046c5e5b8382701ea8bb764c", size = 7629022, upload-time = "2026-08-05T09:22:50.485Z" }, ] -[[package]] -name = "litellm" -version = "1.100.4" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "aiohttp" }, - { name = "boto3" }, - { name = "click" }, - { name = "fastuuid" }, - { name = "httpx" }, - { name = "importlib-metadata" }, - { name = "jinja2" }, - { name = "jsonschema" }, - { name = "openai" }, - { name = "pydantic" }, - { name = "pydantic-settings" }, - { name = "python-dotenv" }, - { name = "tiktoken" }, - { name = "tokenizers" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/71/63/28fb5696e52f40f4bafd096a3d0ed2c8679a24311ed89c200a280ee3b68a/litellm-1.100.4.tar.gz", hash = "sha256:c031cbc63848be16403725594f6060fe002dc2e65f0866878285d72923ac27ab" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/67/5e/407031806577da6b1f3c9217774dabcc1353345337282093915411596aa9/litellm-1.100.4-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c86e2e3370674557547022eb9acd8ec842f802bddb2cfa97992332ad56309a53" }, - { url = "https://files.pythonhosted.org/packages/db/71/a284cce63fa6c37dab677415d146c968d9d802e1eed2778efcdddba49010/litellm-1.100.4-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:737c99748f9a37e74d4927ad9fe311e3f531a5379ea423eb0d1412d3310f1643" }, - { url = "https://files.pythonhosted.org/packages/0c/55/c16164e6cbae93437f50aaf550c5dcacc7a1ca4f1af80ece2c116fbeb01e/litellm-1.100.4-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:7228d4d12eb332e5e7b7dbfeb41a4182784ad2acf86f6b7eb806565e04912ada" }, - { url = "https://files.pythonhosted.org/packages/ff/09/d99612074d3e114c449fc928d7374f09d8e87cc74dc36fdbb466fc1b9f5a/litellm-1.100.4-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:ded14b15466163c76c9a7da0cd091663b4bebbff9aa56724a5905386a9d12532" }, - { url = "https://files.pythonhosted.org/packages/6e/a0/cc57f435963705f6baeaa717a870a616dd83981217daa81c9de2c2456a0d/litellm-1.100.4-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:eb12645456c19dd56f1fa5d7e165a2a76effb1501ace5d0a4df18efc0ac5c39a" }, - { url = "https://files.pythonhosted.org/packages/d5/dd/b7b458120100e47fa17718fa7aa7f1953cc01b34ce0439cc9f681cc746a0/litellm-1.100.4-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:aa8e8507678451b18c4369b69506d8948154b2ab4af211454dd062380597c1ce" }, - { url = "https://files.pythonhosted.org/packages/26/35/51d5be49f6add1bea14a0d47cb4f125fc343150095e2e655aa02cba2e885/litellm-1.100.4-cp310-abi3-win_amd64.whl", hash = "sha256:83342175eb948f146f405735b03c1cb1e5c332aeef6a98cfb31f59233d10426a" }, -] - [[package]] name = "llvmlite" version = "0.47.0" @@ -4383,15 +4187,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ff/89/4ad9d215ef6ade89f27a72dc4e86b98ef1a43534cc3e6a6900a362a0bf0a/murmurhash-1.0.15-cp313-cp313t-win_arm64.whl", hash = "sha256:189a8de4d657b5da9efd66601b0636330b08262b3a55431f2379097c986995d0", size = 25398, upload-time = "2025-11-14T09:50:43.023Z" }, ] -[[package]] -name = "mypy-boto3-bedrock-runtime" -version = "1.43.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/21/f2/61519c0162307b1e4d47f63ed8b25390874640934f3d2d25c5d6c5078dd8/mypy_boto3_bedrock_runtime-1.43.0.tar.gz", hash = "sha256:19fc3167de6e66dd7a0ab293adc55c93e2fd67be35e8ab4fc3a7523a380752ce", size = 29903, upload-time = "2026-04-29T22:57:57.561Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/40/4d/7e4c4d55af23b2b1304d6814db8c406beab7977056963200230417c1a2db/mypy_boto3_bedrock_runtime-1.43.0-py3-none-any.whl", hash = "sha256:a125296f992093d58bdcd95176002680fa81ca8a8b8bdf02afad7e5f2d8966aa", size = 36172, upload-time = "2026-04-29T22:57:54.777Z" }, -] - [[package]] name = "mysql-connector-python" version = "26.7.0" @@ -4406,15 +4201,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ff/01/07efeb3e3fc730f945af40ddf1f4c879c13945b193a94eba8e642c0d4176/mysql_connector_python-26.7.0-py2.py3-none-any.whl", hash = "sha256:43ee453b53556f690e4f4b795d3f17a1c6d4b45f729657bb0d999c8dfc9261d5", size = 480733, upload-time = "2026-07-29T10:18:49.777Z" }, ] -[[package]] -name = "networkx" -version = "3.6" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e8/fc/7b6fd4d22c8c4dc5704430140d8b3f520531d4fe7328b8f8d03f5a7950e8/networkx-3.6.tar.gz", hash = "sha256:285276002ad1f7f7da0f7b42f004bcba70d381e936559166363707fdad3d72ad", size = 2511464, upload-time = "2025-11-24T03:03:47.158Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/07/c7/d64168da60332c17d24c0d2f08bdf3987e8d1ae9d84b5bbd0eec2eb26a55/networkx-3.6-py3-none-any.whl", hash = "sha256:cdb395b105806062473d3be36458d8f1459a4e4b98e236a66c3a48996e07684f", size = 2063713, upload-time = "2025-11-24T03:03:45.21Z" }, -] - [[package]] name = "nh3" version = "0.3.7" @@ -4574,25 +4360,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b6/ca/862b1e7a639460f0ca25fd5b6135fb42cf9deea86d398a92e44dfda2279d/onnxruntime-1.23.2-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e2b9233c4947907fd1818d0e581c049c41ccc39b2856cc942ff6d26317cee145", size = 17394184, upload-time = "2025-10-22T03:47:08.127Z" }, ] -[[package]] -name = "openai" -version = "2.24.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "anyio" }, - { name = "distro" }, - { name = "httpx" }, - { name = "jiter" }, - { name = "pydantic" }, - { name = "sniffio" }, - { name = "tqdm" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/55/13/17e87641b89b74552ed408a92b231283786523edddc95f3545809fab673c/openai-2.24.0.tar.gz", hash = "sha256:1e5769f540dbd01cb33bc4716a23e67b9d695161a734aff9c5f925e2bf99a673", size = 658717, upload-time = "2026-02-24T20:02:07.958Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c9/30/844dc675ee6902579b8eef01ed23917cc9319a1c9c0c14ec6e39340c96d0/openai-2.24.0-py3-none-any.whl", hash = "sha256:fed30480d7d6c884303287bde864980a4b137b60553ffbcf9ab4a233b7a73d94", size = 1120122, upload-time = "2026-02-24T20:02:05.669Z" }, -] - [[package]] name = "opendal" version = "0.46.0" @@ -4617,30 +4384,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/23/6e/64d1330d7e39c2afefcb786491a55aca76e0c6aabc4774f2a7a1300c583a/opendal-0.46.0-cp313-cp313t-win_amd64.whl", hash = "sha256:8331cd209e1dc18a7eab44a4e35aa6b6ebe77adfb295de496556991c6659e65e", size = 14895732, upload-time = "2025-07-17T06:58:50.721Z" }, ] -[[package]] -name = "openinference-instrumentation" -version = "0.1.42" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "openinference-semantic-conventions" }, - { name = "opentelemetry-api" }, - { name = "opentelemetry-sdk" }, - { name = "wrapt" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/00/d0/b19061a21fd6127d2857c77744a36073bba9c1502d1d5e8517b708eb8b7c/openinference_instrumentation-0.1.42.tar.gz", hash = "sha256:2275babc34022e151b5492cfba41d3b12e28377f8e08cb45e5d64fe2d9d7fe37", size = 23954, upload-time = "2025-11-05T01:37:46.869Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c3/71/43ee4616fc95dbd2f560550f199c6652a5eb93f84e8aa0039bc95c19cfe0/openinference_instrumentation-0.1.42-py3-none-any.whl", hash = "sha256:e7521ff90833ef7cc65db526a2f59b76a496180abeaaee30ec6abbbc0b43f8ec", size = 30086, upload-time = "2025-11-05T01:37:43.866Z" }, -] - -[[package]] -name = "openinference-semantic-conventions" -version = "0.1.25" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/0b/68/81c8a0b90334ff11e4f285e4934c57f30bea3ef0c0b9f99b65e7b80fae3b/openinference_semantic_conventions-0.1.25.tar.gz", hash = "sha256:f0a8c2cfbd00195d1f362b4803518341e80867d446c2959bf1743f1894fce31d", size = 12767, upload-time = "2025-11-05T01:37:45.89Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/fd/3d/dd14ee2eb8a3f3054249562e76b253a1545c76adbbfd43a294f71acde5c3/openinference_semantic_conventions-0.1.25-py3-none-any.whl", hash = "sha256:3814240f3bd61f05d9562b761de70ee793d55b03bca1634edf57d7a2735af238", size = 10395, upload-time = "2025-11-05T01:37:43.697Z" }, -] - [[package]] name = "openpyxl" version = "3.1.5" @@ -4708,19 +4451,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/25/d5/603a29962542152331c6e37a121129a26e24917daa4298ca00bea97571bb/opentelemetry_distro-0.65b0-py3-none-any.whl", hash = "sha256:0e15bb1a54c638e6c361bc66bc43e9de9ead442e1ae06f8099a2c1d27b6ef845", size = 2775, upload-time = "2026-07-16T15:24:47.562Z" }, ] -[[package]] -name = "opentelemetry-exporter-otlp" -version = "1.44.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "opentelemetry-exporter-otlp-proto-grpc" }, - { name = "opentelemetry-exporter-otlp-proto-http" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/f2/45/7af37fe54e5d3e66e7dcd7ba8b8aeee73f202bfac909cc94b8c4e428f9ac/opentelemetry_exporter_otlp-1.44.0.tar.gz", hash = "sha256:af1cde7c33ea8ed624bf04ac49a885730fe44c1f1ad698656e592c38f70ce106", size = 6090, upload-time = "2026-07-16T15:25:34.585Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/18/c3/7b466a9463944e70b37b744072a0c1b88a425dade3fff0631adec66c9bcc/opentelemetry_exporter_otlp-1.44.0-py3-none-any.whl", hash = "sha256:4a498fa8d8fd8be9e8e2d175fe5524a3fe581ccffadd8509db86526a5fb97051", size = 6727, upload-time = "2026-07-16T15:25:14.445Z" }, -] - [[package]] name = "opentelemetry-exporter-otlp-proto-common" version = "1.44.0" @@ -4970,33 +4700,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/23/3f/ab8d29df207ce5f470a07fa96ebb48af4e95b7fab7e7635311b9a32f2fab/opentelemetry_util_http-0.65b0-py3-none-any.whl", hash = "sha256:7553b606f963097cb190536dc30556cce85090692e471a422fff30ca29b04348", size = 8245, upload-time = "2026-07-16T15:25:46.482Z" }, ] -[[package]] -name = "opik" -version = "1.11.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "boto3-stubs", extra = ["bedrock-runtime"] }, - { name = "click" }, - { name = "httpx" }, - { name = "jinja2" }, - { name = "litellm" }, - { name = "openai" }, - { name = "pydantic" }, - { name = "pydantic-settings" }, - { name = "pytest" }, - { name = "rapidfuzz" }, - { name = "rich" }, - { name = "sentry-sdk" }, - { name = "tenacity" }, - { name = "tqdm" }, - { name = "uuid6" }, - { name = "watchfiles" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/a3/b9/f6c7e41cb6c02f6e68fde9b6dacf377dcf42079cdbaf891f9fecf4dc958b/opik-1.11.2.tar.gz", hash = "sha256:79e054595b29e1ca8a4fd67d023249f0cf355ea9efbe3e00c28f51628d053d63", size = 871557, upload-time = "2026-04-10T10:48:14.965Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/99/2d/e5536a2a1b6fdd920d995e09315523be53bde5fe01f104894d9ba7421a8c/opik-1.11.2-py3-none-any.whl", hash = "sha256:1016b6db7563d847e50e463a2ae09e595b6921372dd52edeada660b82036e1b2", size = 1451056, upload-time = "2026-04-10T10:48:12.927Z" }, -] - [[package]] name = "optionaldict" version = "0.1.2" @@ -5146,19 +4849,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9a/cb/5ad79e02a556cc23fed5816de0109fa8af660c66cfa5f4af74c3e8d4cd26/pandas_stubs-3.0.5.260914-py3-none-any.whl", hash = "sha256:39a1300c5c5c55fdf609e3476805decce5d5015539a4dcb683449f8feaeee2fb", size = 177344, upload-time = "2026-09-14T16:42:33.771Z" }, ] -[[package]] -name = "pdfminer-six" -version = "20251230" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "charset-normalizer" }, - { name = "cryptography" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/46/9a/d79d8fa6d47a0338846bb558b39b9963b8eb2dfedec61867c138c1b17eeb/pdfminer_six-20251230.tar.gz", hash = "sha256:e8f68a14c57e00c2d7276d26519ea64be1b48f91db1cdc776faa80528ca06c1e", size = 8511285, upload-time = "2025-12-30T15:49:13.104Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/65/d7/b288ea32deb752a09aab73c75e1e7572ab2a2b56c3124a5d1eb24c62ceb3/pdfminer_six-20251230-py3-none-any.whl", hash = "sha256:9ff2e3466a7dfc6de6fd779478850b6b7c2d9e9405aa2a5869376a822771f485", size = 6591909, upload-time = "2025-12-30T15:49:10.76Z" }, -] - [[package]] name = "pgvecto-rs" version = "0.2.2" @@ -5206,15 +4896,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/87/7d/8d02eba9f7ecc87f8e90f731e032cc4773802f2a843763afd0c66a535ff4/pillow-13.0.0-cp313-cp313-win_arm64.whl", hash = "sha256:40a266371ece581e5de13b6635ef78d7487824af80561b6ca27694a88c19bc0f", size = 2698826, upload-time = "2026-10-16T22:16:25.382Z" }, ] -[[package]] -name = "platformdirs" -version = "4.5.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/61/33/9611380c2bdb1225fdef633e2a9610622310fed35ab11dac9620972ee088/platformdirs-4.5.0.tar.gz", hash = "sha256:70ddccdd7c99fc5942e9fc25636a8b34d04c24b335100223152c2803e4063312", size = 21632, upload-time = "2025-10-08T17:44:48.791Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/73/cb/ac7874b3e5d58441674fb70742e6c374b28b0c7cb988d37d991cde47166c/platformdirs-4.5.0-py3-none-any.whl", hash = "sha256:e578a81bb873cbb89a41fcc904c7ef523cc18284b7e3b3ccf06aca1403b7ebd3", size = 18651, upload-time = "2025-10-08T17:44:47.223Z" }, -] - [[package]] name = "pluggy" version = "1.6.0" @@ -5224,31 +4905,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, ] -[[package]] -name = "polyfile-weave" -version = "0.5.9" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "abnf" }, - { name = "chardet" }, - { name = "cint" }, - { name = "fickling" }, - { name = "filelock" }, - { name = "graphviz" }, - { name = "intervaltree" }, - { name = "jinja2" }, - { name = "kaitaistruct" }, - { name = "networkx" }, - { name = "pdfminer-six" }, - { name = "pillow" }, - { name = "pyreadline3", marker = "sys_platform == 'win32'" }, - { name = "pyyaml" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/70/55/e5400762e3884f743d59291e71eaaa9c52dd7e144b75a11911e74ec1bac9/polyfile_weave-0.5.9.tar.gz", hash = "sha256:12341fab03e06ede1bfebbd3627dd24015fde5353ea74ece2da186321b818bdb", size = 6024974, upload-time = "2026-01-22T22:08:48.081Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/52/94/215005530a48c5f7d4ec4a31acdb5828f2bfb985cc6e577b0eaa5882c0e2/polyfile_weave-0.5.9-py3-none-any.whl", hash = "sha256:6ae4b1b5eeac9f5bfc862474484d6d3e33655fab31749d93af0b0a91fddabfc7", size = 1700174, upload-time = "2026-01-22T22:08:46.346Z" }, -] - [[package]] name = "portalocker" version = "2.10.1" @@ -7713,15 +7369,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d9/26/529f4beee17e5248e37e0bc17a2761d34c0fa3b1e5729c88adb2065bae6e/uuid_utils-0.14.1-cp39-abi3-win_arm64.whl", hash = "sha256:b04cb49b42afbc4ff8dbc60cf054930afc479d6f4dd7f1ec3bbe5dbfdde06b7a", size = 188132, upload-time = "2026-02-20T22:50:41.718Z" }, ] -[[package]] -name = "uuid6" -version = "2025.0.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/ca/b7/4c0f736ca824b3a25b15e8213d1bcfc15f8ac2ae48d1b445b310892dc4da/uuid6-2025.0.1.tar.gz", hash = "sha256:cd0af94fa428675a44e32c5319ec5a3485225ba2179eefcf4c3f205ae30a81bd", size = 13932, upload-time = "2025-07-04T18:30:35.186Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/3d/b2/93faaab7962e2aa8d6e174afb6f76be2ca0ce89fde14d3af835acebcaa59/uuid6-2025.0.1-py3-none-any.whl", hash = "sha256:80530ce4d02a93cdf82e7122ca0da3ebbbc269790ec1cb902481fa3e9cc9ff99", size = 6979, upload-time = "2025-07-04T18:30:34.001Z" }, -] - [[package]] name = "uvicorn" version = "0.38.0" @@ -7796,34 +7443,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/37/da/7ccbe82470dc27e1cfd0466dc637248be906eb8447c28a40c1c74cf617ee/volcengine_compat-1.0.156-py3-none-any.whl", hash = "sha256:4abc149a7601ebad8fa2d28fab50c7945145cf74daecb71bca797b0bdc82c5a5", size = 677272, upload-time = "2024-10-13T09:17:19.944Z" }, ] -[[package]] -name = "wandb" -version = "0.28.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "click" }, - { name = "opentelemetry-api" }, - { name = "packaging" }, - { name = "platformdirs" }, - { name = "protobuf" }, - { name = "pydantic" }, - { name = "pyyaml" }, - { name = "requests" }, - { name = "sentry-sdk" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/fd/41/c2eb0e25e0504287442f0829a5dc380669b683f78068c79220acc626383f/wandb-0.28.2.tar.gz", hash = "sha256:9dba560ce076f96a0033a0d2898d97cba651fc95fec7274fd09920bae8aec5cf", size = 41011367, upload-time = "2026-08-12T01:34:20.484Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/47/d9/8fb47c6c0b6e70378fc2b5da6dfba5d1d3178366a06bd7e672c3249964c3/wandb-0.28.2-py3-none-macosx_12_0_arm64.whl", hash = "sha256:f9aa4037d18168dd4e804ff8213100c4a3ebe3e52b5f29463bbc7fadcaadff75", size = 26716653, upload-time = "2026-08-12T01:33:58.338Z" }, - { url = "https://files.pythonhosted.org/packages/2f/a7/6690491966e46409bfe770b73bdb1952fb6a59988b9e8c72c748f4ca080c/wandb-0.28.2-py3-none-macosx_12_0_x86_64.whl", hash = "sha256:b3830781cb93a29ad91a736d3bc5763312e782467432ee432fd731f408708991", size = 28649923, upload-time = "2026-08-12T01:34:01.22Z" }, - { url = "https://files.pythonhosted.org/packages/03/46/e69c813d88f37f77da92cfd9c63e1e9b5ce70834b0f420f876c66590353f/wandb-0.28.2-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:0c475c6b93d57e0455f3d1323e323927696502e6af7ee15e389a552114edacff", size = 27385453, upload-time = "2026-08-12T01:34:03.876Z" }, - { url = "https://files.pythonhosted.org/packages/3c/7f/f3e23a2ff01848d7a5adab9f307992b92ef99c5e6bff2bd89cd46d651e3a/wandb-0.28.2-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:1db698d107871c66b2dcbb0cf4dc2af1ddb159ba94e957e890158ec60ab2de54", size = 29622780, upload-time = "2026-08-12T01:34:06.621Z" }, - { url = "https://files.pythonhosted.org/packages/9e/1d/e087f71898ebfa86f68264e91f44ee5b8106823f09e2644043e3f1d5fe49/wandb-0.28.2-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:86e017d3f38bb99374de4d991a2ea0e6e71722164918585b3dea66190b28d0eb", size = 27425107, upload-time = "2026-08-12T01:34:09.359Z" }, - { url = "https://files.pythonhosted.org/packages/f2/58/3f7a624c880d7b9e60b4a000b0bb751c20244d0bcc002d3ff683cf828651/wandb-0.28.2-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:81468f5c00b5453a2a1d3c7ca7a18056c52a85fa4df60299052cb96829d81b11", size = 29851453, upload-time = "2026-08-12T01:34:12.133Z" }, - { url = "https://files.pythonhosted.org/packages/f6/a3/b680c7518b0dc7eb236e7ef3bcc0994d0fabc67bbef1b73878dee196fc1e/wandb-0.28.2-py3-none-win_amd64.whl", hash = "sha256:ae5c2591214565be9f085c3b7838ee4a89344ae0c8dc06030425def08b6191aa", size = 26800272, upload-time = "2026-08-12T01:34:14.64Z" }, - { url = "https://files.pythonhosted.org/packages/ae/5a/a323653c6989c565b2138ddd7421ef30b47ea2849a62d0181534772b56da/wandb-0.28.2-py3-none-win_arm64.whl", hash = "sha256:23a4b9ac74f211836b71b27ba33bb217977a6b4cdc927e33eb5b12de143dd727", size = 24484327, upload-time = "2026-08-12T01:34:17.46Z" }, -] - [[package]] name = "wasabi" version = "1.1.3" @@ -7899,29 +7518,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a4/74/a148b41572656904a39dfcfed3f84dd1066014eed94e209223ae8e9d088d/weasel-0.4.3-py3-none-any.whl", hash = "sha256:08f65b5d0dbded4879e08a64882de9b9514753d9eaa4c4e2a576e33666ac12cf", size = 50757, upload-time = "2025-11-13T23:52:26.982Z" }, ] -[[package]] -name = "weave" -version = "0.52.36" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cachetools" }, - { name = "click" }, - { name = "diskcache-weave" }, - { name = "gql", extra = ["httpx"] }, - { name = "jsonschema" }, - { name = "packaging" }, - { name = "polyfile-weave" }, - { name = "pydantic" }, - { name = "sentry-sdk" }, - { name = "tenacity" }, - { name = "tzdata", marker = "sys_platform == 'win32'" }, - { name = "wandb" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/69/ee/63064875e0d1d7bf261484da7a8ef1d447d38f739b3e93e4d8673d51c882/weave-0.52.36.tar.gz", hash = "sha256:f58b37786de5444914e408e64026c3131b5c4417e6889d5a61fdcbec12c8e8dd", size = 793026, upload-time = "2026-04-01T17:23:50.671Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/78/34/57a0843b3016e3dc63660cd06bfa52322f3333225ae7e7cfa6db97a59f8c/weave-0.52.36-py3-none-any.whl", hash = "sha256:4ff5b53323d20cc321aec665a4b4da746d6d85d432eda2ccca0e85bc8891649d", size = 983539, upload-time = "2026-04-01T17:23:48.819Z" }, -] - [[package]] name = "weaviate-client" version = "4.22.0" diff --git a/dev/start-worker b/dev/start-worker index 2738afa1019237..3690aead32539b 100755 --- a/dev/start-worker +++ b/dev/start-worker @@ -29,7 +29,7 @@ show_help() { echo " schedule_poller - Schedule polling tasks" echo " schedule_executor - Schedule execution tasks" echo " mail - Email notifications" - echo " ops_trace - Operations tracing" + echo " ops_trace_v2 - Operations tracing (override with OPS_TRACE_QUEUE)" echo " app_deletion - Application cleanup" echo " plugin - Plugin operations" echo " workflow_storage - Workflow storage tasks" @@ -107,10 +107,10 @@ if [[ -z "${QUEUES}" ]]; then # Configure queues based on product edition if [[ "${DEPLOYMENT_EDITION:-COMMUNITY}" == "CLOUD" ]]; then # Cloud edition: separate queues for dataset and trigger tasks - QUEUES="dataset,dataset_summary,priority_dataset,priority_pipeline,pipeline,mail,ops_trace,app_deletion,plugin,workflow_storage,conversation,workflow_professional,workflow_team,workflow_sandbox,schedule_poller,schedule_executor,triggered_workflow_dispatcher,trigger_refresh_executor,retention,workflow_based_app_execution,knowledge_fs_lifecycle" + QUEUES="dataset,dataset_summary,priority_dataset,priority_pipeline,pipeline,mail,${OPS_TRACE_QUEUE:-ops_trace_v2},app_deletion,plugin,workflow_storage,conversation,workflow_professional,workflow_team,workflow_sandbox,schedule_poller,schedule_executor,triggered_workflow_dispatcher,trigger_refresh_executor,retention,workflow_based_app_execution,knowledge_fs_lifecycle" else # Self-hosted editions: dataset and workflow have separate queues - QUEUES="dataset,dataset_summary,priority_dataset,priority_pipeline,pipeline,mail,ops_trace,app_deletion,plugin,workflow_storage,conversation,workflow,schedule_poller,schedule_executor,triggered_workflow_dispatcher,trigger_refresh_executor,retention,workflow_based_app_execution,knowledge_fs_lifecycle" + QUEUES="dataset,dataset_summary,priority_dataset,priority_pipeline,pipeline,mail,${OPS_TRACE_QUEUE:-ops_trace_v2},app_deletion,plugin,workflow_storage,conversation,workflow,schedule_poller,schedule_executor,triggered_workflow_dispatcher,trigger_refresh_executor,retention,workflow_based_app_execution,knowledge_fs_lifecycle" fi echo "No queues specified, using edition-based defaults: ${QUEUES}" diff --git a/docker/.env.example b/docker/.env.example index 01e7be49aabfd9..a0e1f159ca0739 100644 --- a/docker/.env.example +++ b/docker/.env.example @@ -1,304 +1,10 @@ -# ------------------------------------------------------------------ -# Essential defaults for Docker Compose deployments. -# Only include variables required for services to start. -# Do not add optional variables to this file. -# -# For a default deployment, copy this file to .env and run: -# docker compose up -d -# -# Optional and provider-specific variables live under docker/envs/. -# Copy an optional *.env.example file beside itself without the -# .example suffix when you need those advanced settings. -# Values in docker/.env take precedence over docker/envs/*.env files. -# ------------------------------------------------------------------ - -# Core service URLs -CONSOLE_API_URL= -SERVER_CONSOLE_API_URL=http://api:5001 -CONSOLE_WEB_URL= -SERVICE_API_URL= -TRIGGER_URL=http://localhost -HUMAN_INPUT_IM_EVENT_TRANSPORT_MODE=WEBHOOK -APP_API_URL= -APP_WEB_URL= -FILES_URL= -INTERNAL_FILES_URL= -ENDPOINT_URL_TEMPLATE=http://localhost/e/{hook_id} -NEXT_PUBLIC_SOCKET_URL=ws://localhost - -# Runtime and security -LANG=C.UTF-8 -LC_ALL=C.UTF-8 -PYTHONIOENCODING=utf-8 -# Leave empty to auto-generate a persistent key in the storage directory. -SECRET_KEY= -INIT_PASSWORD= -DEPLOY_ENV=PRODUCTION -CHECK_UPDATE_URL=https://updates.dify.ai -OPENAI_API_BASE=https://api.openai.com/v1 -MIGRATION_ENABLED=true -FILES_ACCESS_TIMEOUT=300 -# Remove `collaboration` from COMPOSE_PROFILES to stop the dedicated websocket service. -ENABLE_COLLABORATION_MODE=true - -# Logging and server workers -LOG_LEVEL=INFO -LOG_OUTPUT_FORMAT=text -LOG_FILE=/app/logs/server.log -LOG_FILE_MAX_SIZE=20 -LOG_FILE_BACKUP_COUNT=5 -LOG_DATEFORMAT=%Y-%m-%d %H:%M:%S -LOG_TZ=UTC -DEBUG=false -FLASK_DEBUG=false -ENABLE_REQUEST_LOGGING=False -DIFY_BIND_ADDRESS=0.0.0.0 -DIFY_PORT=5001 -SERVER_WORKER_AMOUNT=1 -SERVER_WORKER_CLASS=gevent -SERVER_WORKER_CONNECTIONS=10 -API_WEBSOCKET_WORKER_AMOUNT=1 -API_WEBSOCKET_WORKER_CLASS=geventwebsocket.gunicorn.workers.GeventWebSocketWorker -API_WEBSOCKET_WORKER_CONNECTIONS=1000 -API_WEBSOCKET_GUNICORN_TIMEOUT=360 -GUNICORN_TIMEOUT=360 -# Must stay above the idle connection reuse window of the HTTP clients that -# pool connections to the API (Node's undici reuses a socket for 4s). 0 -# disables keep-alive. -GUNICORN_KEEP_ALIVE=10 -CELERY_WORKER_CLASS= -CELERY_WORKER_AMOUNT=4 -CELERY_AUTO_SCALE=false -CELERY_MAX_WORKERS= -CELERY_MIN_WORKERS= -COMPOSE_WORKER_HEALTHCHECK_DISABLED=true -COMPOSE_WORKER_HEALTHCHECK_INTERVAL=30s -COMPOSE_WORKER_HEALTHCHECK_TIMEOUT=30s - -# Database -DB_TYPE=postgresql -DB_USERNAME=postgres -DB_PASSWORD=difyai123456 -DB_HOST=db_postgres -DB_PORT=5432 -DB_DATABASE=dify -SQLALCHEMY_POOL_SIZE=30 -SQLALCHEMY_MAX_OVERFLOW=10 -SQLALCHEMY_POOL_RECYCLE=3600 -SQLALCHEMY_ECHO=false -SQLALCHEMY_POOL_PRE_PING=false -SQLALCHEMY_POOL_USE_LIFO=false -SQLALCHEMY_POOL_TIMEOUT=30 -SQLALCHEMY_POOL_RESET_ON_RETURN=rollback -PGDATA=/var/lib/postgresql/data/pgdata -POSTGRES_MAX_CONNECTIONS=200 -POSTGRES_SHARED_BUFFERS=128MB -POSTGRES_WORK_MEM=4MB -POSTGRES_MAINTENANCE_WORK_MEM=64MB -POSTGRES_EFFECTIVE_CACHE_SIZE=4096MB -POSTGRES_STATEMENT_TIMEOUT=0 -POSTGRES_IDLE_IN_TRANSACTION_SESSION_TIMEOUT=0 - -# Redis and Celery -REDIS_HOST=redis -REDIS_PORT=6379 -REDIS_USERNAME= -REDIS_PASSWORD=difyai123456 -REDIS_USE_SSL=false -REDIS_SSL_CERT_REQS=CERT_NONE -REDIS_SSL_CA_CERTS= -REDIS_SSL_CERTFILE= -REDIS_SSL_KEYFILE= -REDIS_DB=0 -REDIS_KEY_PREFIX= -REDIS_MAX_CONNECTIONS= -REDIS_RETRY_RETRIES=3 -REDIS_RETRY_BACKOFF_BASE=1.0 -REDIS_RETRY_BACKOFF_CAP=10.0 -REDIS_SOCKET_TIMEOUT=5.0 -REDIS_SOCKET_CONNECT_TIMEOUT=5.0 -REDIS_HEALTH_CHECK_INTERVAL=30 -REDIS_KEEPALIVE=true -REDIS_KEEPALIVE_IDLE=30 -REDIS_KEEPALIVE_INTERVAL=10 -REDIS_KEEPALIVE_COUNT=10 -CELERY_BROKER_URL=redis://:difyai123456@redis:6379/1 -CELERY_BACKEND=redis -BROKER_USE_SSL=false -CELERY_TASK_ANNOTATIONS=null -EVENT_BUS_REDIS_URL= -EVENT_BUS_REDIS_CHANNEL_TYPE=pubsub -EVENT_BUS_REDIS_USE_CLUSTERS=false - -# Web and app limits -WEB_API_CORS_ALLOW_ORIGINS=* -CONSOLE_CORS_ALLOW_ORIGINS=* -COOKIE_DOMAIN= -NEXT_PUBLIC_COOKIE_DOMAIN= -NEXT_PUBLIC_BATCH_CONCURRENCY=5 -API_SENTRY_DSN= -API_SENTRY_TRACES_SAMPLE_RATE=1.0 -API_SENTRY_PROFILES_SAMPLE_RATE=1.0 -WEB_SENTRY_DSN= -AMPLITUDE_API_KEY= -TURNSTILE_SITE_KEY= -TEXT_GENERATION_TIMEOUT_MS=60000 -CSP_WHITELIST= -ALLOW_EMBED=false -ALLOW_UNSAFE_DATA_SCHEME=false -TOP_K_MAX_VALUE=10 -INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH=4000 -LOOP_NODE_MAX_COUNT=100 -MAX_TOOLS_NUM=10 -MAX_PARALLEL_LIMIT=10 -MAX_ITERATIONS_NUM=99 -MAX_TREE_DEPTH=50 -ENABLE_WEBSITE_JINAREADER=true -ENABLE_WEBSITE_FIRECRAWL=true -ENABLE_WEBSITE_WATERCRAWL=true -NEXT_PUBLIC_ENABLE_SINGLE_DOLLAR_LATEX=false -NEXT_PUBLIC_ENABLE_AGENT_V2=true -EXPERIMENTAL_ENABLE_VINEXT=false - -# Storage and default vector store -STORAGE_TYPE=opendal -OPENDAL_SCHEME=fs -OPENDAL_FS_ROOT=storage -VECTOR_STORE=weaviate -VECTOR_INDEX_NAME_PREFIX=Vector_index -WEAVIATE_ENDPOINT=http://weaviate:8080 -WEAVIATE_API_KEY=WVF5YThaHlkYwhGUSmCRgsX3tD5ngdN8pkih -WEAVIATE_GRPC_ENDPOINT=grpc://weaviate:50051 -WEAVIATE_TOKENIZATION=word -WEAVIATE_PERSISTENCE_DATA_PATH=/var/lib/weaviate -WEAVIATE_QUERY_DEFAULTS_LIMIT=25 -WEAVIATE_AUTHENTICATION_ANONYMOUS_ACCESS_ENABLED=true -WEAVIATE_DEFAULT_VECTORIZER_MODULE=none -WEAVIATE_CLUSTER_HOSTNAME=node1 -WEAVIATE_AUTHENTICATION_APIKEY_ENABLED=true -WEAVIATE_AUTHENTICATION_APIKEY_ALLOWED_KEYS=WVF5YThaHlkYwhGUSmCRgsX3tD5ngdN8pkih -WEAVIATE_AUTHENTICATION_APIKEY_USERS=hello@dify.ai -WEAVIATE_AUTHORIZATION_ADMINLIST_ENABLED=true -WEAVIATE_AUTHORIZATION_ADMINLIST_USERS=hello@dify.ai -WEAVIATE_DISABLE_TELEMETRY=false -WEAVIATE_ENABLE_TOKENIZER_GSE=false -WEAVIATE_ENABLE_TOKENIZER_KAGOME_JA=false -WEAVIATE_ENABLE_TOKENIZER_KAGOME_KR=false - -# Sandbox and SSRF proxy -CODE_EXECUTION_ENDPOINT=http://sandbox:8194 -CODE_EXECUTION_API_KEY=dify-sandbox -CODE_EXECUTION_SSL_VERIFY=True -CODE_EXECUTION_POOL_MAX_CONNECTIONS=100 -CODE_EXECUTION_POOL_MAX_KEEPALIVE_CONNECTIONS=20 -CODE_EXECUTION_POOL_KEEPALIVE_EXPIRY=5.0 -CODE_EXECUTION_CONNECT_TIMEOUT=10 -CODE_EXECUTION_READ_TIMEOUT=60 -CODE_EXECUTION_WRITE_TIMEOUT=10 -SANDBOX_API_KEY=dify-sandbox -SANDBOX_GIN_MODE=release -SANDBOX_WORKER_TIMEOUT=15 -SANDBOX_ENABLE_NETWORK=true -SANDBOX_HTTP_PROXY=http://ssrf_proxy:3128 -SANDBOX_HTTPS_PROXY=http://ssrf_proxy:3128 -SANDBOX_PORT=8194 -PIP_MIRROR_URL= -SSRF_PROXY_HTTP_URL=http://ssrf_proxy:3128 -SSRF_PROXY_HTTPS_URL=http://ssrf_proxy:3128 -SSRF_HTTP_PORT=3128 -SSRF_COREDUMP_DIR=/var/spool/squid -SSRF_DEFAULT_TIME_OUT=5 -SSRF_DEFAULT_CONNECT_TIME_OUT=5 -SSRF_DEFAULT_READ_TIME_OUT=5 -SSRF_DEFAULT_WRITE_TIME_OUT=5 -SSRF_POOL_MAX_CONNECTIONS=100 -SSRF_POOL_MAX_KEEPALIVE_CONNECTIONS=20 -SSRF_POOL_KEEPALIVE_EXPIRY=5.0 -# Comma-separated CIDR ranges that the SSRF proxy should allow even when they -# resolve to private, loopback, link-local, or otherwise non-public addresses. -# Leave empty (the default) to keep the deny-by-default policy. Required when -# Dify needs to reach internal HTTP endpoints (e.g. an internal API on -# http://172.21.x.x) from an HTTP Request node, a tool, or a plugin. -# Example: SSRF_PROXY_ALLOW_PRIVATE_IPS=172.21.0.0/16,10.0.0.0/8 -SSRF_PROXY_ALLOW_PRIVATE_IPS= - -# Plugin daemon -DB_PLUGIN_DATABASE=dify_plugin -EXPOSE_PLUGIN_DAEMON_PORT=5002 -PLUGIN_DAEMON_PORT=5002 -PLUGIN_DAEMON_KEY=lYkiYYT6owG+71oLerGzA7GXCgOT++6ovaezWAjpCjf+Sjc3ZtU+qUEi -PLUGIN_DAEMON_URL=http://plugin_daemon:5002 -PLUGIN_MAX_PACKAGE_SIZE=52428800 -PLUGIN_MAX_FILE_SIZE=52428800 -PLUGIN_MODEL_SCHEMA_CACHE_TTL=3600 -PLUGIN_PPROF_ENABLED=false -PLUGIN_DEBUGGING_HOST=0.0.0.0 -PLUGIN_DEBUGGING_PORT=5003 -EXPOSE_PLUGIN_DEBUGGING_HOST=localhost -EXPOSE_PLUGIN_DEBUGGING_PORT=5003 -PLUGIN_DIFY_INNER_API_KEY=QaHbTe77CtuXmsfyhR7+vRjI/+XbV1AaFy691iy+kGDv2Jvy0/eAh8Y1 -PLUGIN_DIFY_INNER_API_URL=http://api:5001 -FORCE_VERIFYING_SIGNATURE=true -PLUGIN_STDIO_BUFFER_SIZE=1024 -PLUGIN_STDIO_MAX_BUFFER_SIZE=5242880 -PLUGIN_PYTHON_ENV_INIT_TIMEOUT=120 -PLUGIN_MAX_EXECUTION_TIMEOUT=600 -PLUGIN_STORAGE_TYPE=local -PLUGIN_STORAGE_LOCAL_ROOT=/app/storage -PLUGIN_WORKING_PATH=/app/storage/cwd -PLUGIN_INSTALLED_PATH=plugin -PLUGIN_PACKAGE_CACHE_PATH=plugin_packages -PLUGIN_MEDIA_CACHE_PATH=assets -PLUGIN_STORAGE_OSS_BUCKET= -PLUGIN_SENTRY_ENABLED=false -PLUGIN_SENTRY_DSN= -MARKETPLACE_ENABLED=true -MARKETPLACE_API_URL=https://marketplace.dify.ai -MARKETPLACE_URL= - -# Dify Agent backend -AGENT_BACKEND_BASE_URL=http://agent_backend:5050 -# Bearer token for the Agent backend control-plane API. -# Replace this development default in production. -# Generate one with: python -c 'import secrets; print(secrets.token_urlsafe(32))' -DIFY_AGENT_API_TOKEN=dify-agent-run-token-for-dev-only -# Leave empty to derive from REDIS_PASSWORD. -DIFY_AGENT_REDIS_URL= -# Leave empty to derive from PLUGIN_DAEMON_URL and PLUGIN_DAEMON_KEY. -DIFY_AGENT_PLUGIN_DAEMON_URL= -DIFY_AGENT_PLUGIN_DAEMON_API_KEY= -# Leave empty to derive from PLUGIN_DIFY_INNER_API_URL and PLUGIN_DIFY_INNER_API_KEY. -# DIFY_AGENT_INNER_API_KEY must match API/worker INNER_API_KEY_FOR_PLUGIN, not INNER_API_KEY. -DIFY_AGENT_INNER_API_URL= -DIFY_AGENT_INNER_API_KEY= -# Select exactly one coherent Home Snapshot + Sandbox backend. -DIFY_AGENT_RUNTIME_BACKEND=local -DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT=http://local_sandbox:5004 -DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN= -# Sandbox-reachable Dify API base for dify-agent CLI /files/* transfers. -# Remote Sandboxes should use the public Dify ingress; local Compose uses api via agent_ssrf_proxy. -DIFY_AGENT_SANDBOX_FILES_BASE_URL=http://api:5001 -DIFY_AGENT_STUB_API_BASE_URL=http://agent_backend:5050/agent-stub -# This is security-sensitive: it derives the JWE encryption key for Agent Stub bearer tokens. -# Replace this development default in production. -# Generate one with: python -c 'import secrets; print(secrets.token_urlsafe(32))' -DIFY_AGENT_SERVER_SECRET_KEY=MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY - -# Nginx and Docker Compose -NGINX_SERVER_NAME=_ -NGINX_HTTPS_ENABLED=false -NGINX_PORT=80 -NGINX_SSL_PORT=443 -NGINX_SSL_CERT_FILENAME=dify.crt -NGINX_SSL_CERT_KEY_FILENAME=dify.key -NGINX_SSL_PROTOCOLS=TLSv1.2 TLSv1.3 -NGINX_WORKER_PROCESSES=auto -NGINX_CLIENT_MAX_BODY_SIZE=100M -NGINX_KEEPALIVE_TIMEOUT=65 -NGINX_PROXY_READ_TIMEOUT=3600s -NGINX_PROXY_SEND_TIMEOUT=3600s -NGINX_ENABLE_CERTBOT_CHALLENGE=false -NGINX_SOCKET_IO_UPSTREAM=api_websocket:5001 -EXPOSE_NGINX_PORT=80 -EXPOSE_NGINX_SSL_PORT=443 -COMPOSE_PROFILES=${VECTOR_STORE:-weaviate},${DB_TYPE:-postgresql},collaboration +# Only settings requiring manual configuration belong in this file. +# Application and Docker defaults are loaded automatically; do not repeat them here. +# Copy this file to .env, then select the services for your deployment. +# Optional overrides and provider settings are documented in envs/**/*.env.example. +# Values in .env take precedence over the matching envs/*.env files. + +# Compose requires an explicit selection of the bundled database, vector store, +# and websocket service. Match these profiles to your DB_TYPE and VECTOR_STORE +# overrides when using different providers; omit services hosted externally. +COMPOSE_PROFILES=weaviate,postgresql,collaboration diff --git a/docker/README.md b/docker/README.md index e4c9af0d2b072c..ac4096269b00a8 100644 --- a/docker/README.md +++ b/docker/README.md @@ -8,13 +8,13 @@ Run Dify with Docker Compose. Configuration templates live in `.env.example` and 2. **Environment Setup**: - Navigate to the `docker` directory. - Copy `.env.example` to `.env`. - - Customize `.env` when you need to change essential startup defaults. Copy optional files from `envs/` without the `.example` suffix when you need advanced settings. + - Select the bundled services through `COMPOSE_PROFILES`. Add overrides to `.env` only when needed. Copy optional files from `envs/` without the `.example` suffix when you need advanced settings. - **Optional (for advanced deployments)**: If you maintain a full `.env` file copied from `.env.example`, you may use the environment synchronization tool to keep it aligned with the latest `.env.example` updates while preserving your custom settings. See the [Environment Variables Synchronization](#environment-variables-synchronization) section below. 3. **Running the Services**: - Execute `docker compose up -d` from the `docker` directory to start the services. - - To specify a vector database, set the `VECTOR_STORE` variable in your `.env` file to your desired vector database service, such as `milvus`, `weaviate`, or `opensearch`. See `envs/vectorstores/` for the full list of supported options. + - To specify a vector database, set the `VECTOR_STORE` variable in your `.env` file to your desired vector database service, such as `milvus`, `weaviate`, or `opensearch`, and select the matching service in `COMPOSE_PROFILES`. See `envs/vectorstores/` for the full list of supported options. ```bash cp .env.example .env docker compose up -d @@ -99,6 +99,18 @@ this is intentional fail-closed behavior. Do not add a proxy route, set `KNOWLEDGE_FS_ENABLED=true`, or send product traffic until readiness returns `200` and the migration and per-Workspace cutover gates are complete. +### Ops tracing rolling upgrades + +New tracing tasks use `OPS_TRACE_QUEUE=ops_trace_v2` by default. Set it consistently on the API, +workers, and Beat in `docker/envs/core-services/shared.env` (or `docker/.env`). Default worker +subscriptions follow this value; explicit `CELERY_QUEUES` or `CELERY_WORKER_QUEUES` lists must +include it and exclude legacy `ops_trace` while old workers are draining. + +Keep an old tracing worker pool on `ops_trace` until old producers stop and their queued, active, reserved, scheduled, and +retry tasks finish. To return to `ops_trace`, first remove old workers, have new workers consume +both queues, then change `OPS_TRACE_QUEUE` on all new processes and drain `ops_trace_v2` before +removing its consumers. See the [complete rollout sequence](../api/providers/trace/README.md#rolling-upgrade-from-the-previous-pipeline). + ### KnowledgeFS lifecycle and upgrade workers The standard Dify worker consumes `knowledge_fs_lifecycle` for provisioning, deletion, initial @@ -263,14 +275,19 @@ use multiple real addresses. ## Overview of `.env`, `.env.example`, and `envs/` -- `.env.example` contains the essential default configuration for Docker Compose deployments. -- `.env` contains local startup values copied from `.env.example` and any local changes. +- `.env.example` contains only settings requiring manual configuration, such as the Compose service selection. +- `.env` contains your deployment inputs and any explicit overrides. - `envs/*.env.example` files contain optional advanced configuration grouped by theme. -Keep the root `.env.example` limited to variables required to start the default Docker Compose deployment. -Do not add optional, advanced, provider-specific, or service-specific variables there; place them in the appropriate `envs/*.env.example` file instead. +Keep the root `.env.example` limited to settings that require manual user configuration. +Settings with usable application or Docker defaults must not appear there. Document optional, +advanced, provider-specific, and service-specific overrides in the appropriate `envs/*.env.example` file. -Docker Compose reads `envs/*.env` files when present, then reads `.env` last so values in `.env` take precedence. +Docker Compose loads Docker-specific defaults from `envs/core-services/shared.defaults`, +then optional `envs/*.env` files, then `.env` so your overrides take precedence. Other defaults +come from the application and each service in `docker-compose.yaml`; do not copy them into `.env`. +The bundled credentials are intended for local deployment; override them for production using +the matching security and service templates. For the available variables and defaults, read the matching `envs/**/*.env.example` template and its service in `docker-compose.yaml`. @@ -288,10 +305,10 @@ provider's interaction configuration after changing the public origin. ## Environment Variables Synchronization -When upgrading Dify or pulling the latest changes, new environment variables may be introduced in `.env.example` only when they are required for startup, +When upgrading Dify or pulling the latest changes, new environment variables may be introduced in `.env.example` only when they require manual configuration, or in the optional files under `envs/` for advanced, provider-specific, and service-specific settings. -If you use the default workflow, review `.env.example` and keep your `.env` aligned with essential startup values. +If you use the default workflow, review `.env.example` and supply any new deployment inputs that require manual configuration. If you maintain a customized `.env` file copied from `.env.example`, an optional environment variables synchronization tool is provided. @@ -306,8 +323,8 @@ This script compares your current `.env` file with the latest `.env.example` tem - Creates a backup of the current `.env` file before making any changes - Synchronizes newly added environment variables from `.env.example` -- Preserves all existing custom values in `.env` -- Displays differences and variables removed from `.env.example` for review +- Preserves the existing `.env`, including custom settings absent from `.env.example` +- Displays differences and additional local overrides for review **Backup behavior** diff --git a/docker/dify-env-sync.py b/docker/dify-env-sync.py index afa39d84517c9b..e41837651dd61e 100755 --- a/docker/dify-env-sync.py +++ b/docker/dify-env-sync.py @@ -7,7 +7,7 @@ # - Synchronize latest settings from .env.example to .env # - Preserve custom settings in existing .env # - Add new environment variables -# - Detect removed environment variables +# - Report additional local overrides # - Create backup files # ================================================================ @@ -275,69 +275,44 @@ def detect_removed_variables(env_vars: dict[str, str], example_vars: dict[str, s Returns: Sorted list of variable names that no longer appear in .env.example. """ - log_info("Detecting removed environment variables...") + log_info("Checking for additional local overrides...") removed = sorted(set(env_vars) - set(example_vars)) if removed: - log_warning("The following environment variables have been removed from .env.example:") + log_warning("The following local overrides are not listed in .env.example:") for var in removed: log_warning(f" - {var}") - log_warning("Consider manually removing these variables from .env") + log_warning("These overrides will be preserved in .env") else: - log_success("No removed environment variables found") + log_success("No additional local overrides found") return removed -def sync_env_file(work_dir: Path, env_vars: dict[str, str], diffs: dict[str, tuple[str, str]]) -> None: - """Rewrite .env based on .env.example while preserving custom values. - - The output file follows the exact line structure of .env.example - (preserving comments, blank lines, and ordering). For every variable - that exists in .env with a different value from the example, the - current .env value is kept. Variables that are new in .env.example - (not present in .env at all) are added with the example's default. - - Args: - work_dir: Directory containing .env and .env.example. - env_vars: Parsed key/value pairs from the original .env. - diffs: Keys whose .env values differ from .env.example (to preserve). - """ - log_info("Starting partial synchronization of .env file...") +def sync_env_file(work_dir: Path, env_vars: dict[str, str]) -> None: + """Append missing template inputs without rewriting existing overrides or comments.""" + log_info("Synchronizing new deployment inputs into .env...") example_file = work_dir / ".env.example" new_env_file = work_dir / ".env.new" - - # Keys whose current .env value should override the example default - preserved_keys: set[str] = set(diffs.keys()) - - preserved_count = 0 - updated_count = 0 - - env_var_pattern = re.compile(r"^([A-Za-z_][A-Za-z0-9_]*)\s*=") + shutil.copy2(work_dir / ".env", new_env_file) + existing_keys = set(env_vars) + added_count = 0 + env_var_pattern = re.compile(r"^\s*([A-Za-z_][A-Za-z0-9_]*)\s*=") with ( example_file.open(encoding="utf-8") as src, - new_env_file.open("w", encoding="utf-8") as dst, + new_env_file.open("a", encoding="utf-8") as dst, ): for line in src: - raw_line = line.rstrip("\n") - match = env_var_pattern.match(raw_line) - if match: - key = match.group(1) - if key in preserved_keys: - # Write the preserved value from .env - dst.write(f"{key}={env_vars[key]}\n") - log_info(f" Preserved: {key} (.env value)") - preserved_count += 1 - else: - # Use the example value (covers new vars and unchanged ones) - dst.write(line if line.endswith("\n") else raw_line + "\n") - updated_count += 1 - else: - # Blank line, comment, or non-variable line — keep as-is - dst.write(line if line.endswith("\n") else raw_line + "\n") + match = env_var_pattern.match(line) + if match and match.group(1) not in existing_keys: + if added_count == 0: + dst.write("\n# New deployment inputs from .env.example\n") + dst.write(line if line.endswith("\n") else line + "\n") + existing_keys.add(match.group(1)) + added_count += 1 # Atomically replace the original .env try: @@ -347,10 +322,8 @@ def sync_env_file(work_dir: Path, env_vars: dict[str, str], diffs: dict[str, tup new_env_file.unlink(missing_ok=True) sys.exit(1) - log_success("Successfully created new .env file") - log_success("Partial synchronization of .env file completed") - log_info(f" Preserved .env values: {preserved_count}") - log_info(f" Updated to .env.example values: {updated_count}") + log_success("Synchronization of .env file completed") + log_info(f" Preserved existing .env settings; added {added_count} new inputs") def show_statistics(work_dir: Path) -> None: @@ -381,7 +354,7 @@ def build_arg_parser() -> argparse.ArgumentParser: prog="dify-env-sync", description=( "Synchronize .env with .env.example: add new variables, " - "preserve custom values, and report removed variables." + "preserve custom values, and report additional local overrides." ), formatter_class=argparse.RawDescriptionHelpFormatter, epilog=( @@ -430,13 +403,13 @@ def main() -> None: example_vars = parse_env_file(work_dir / ".env.example") # 4. Report differences (values that changed in the example) - diffs = detect_differences(env_vars, example_vars) + detect_differences(env_vars, example_vars) - # 5. Report variables removed from the example + # 5. Report additional local overrides detect_removed_variables(env_vars, example_vars) - # 6. Rewrite .env - sync_env_file(work_dir, env_vars, diffs) + # 6. Append missing inputs to .env + sync_env_file(work_dir, env_vars) # 7. Print summary statistics show_statistics(work_dir) diff --git a/docker/dify-env-sync.sh b/docker/dify-env-sync.sh index edeeae3abc8d99..7d64c00011a989 100755 --- a/docker/dify-env-sync.sh +++ b/docker/dify-env-sync.sh @@ -7,7 +7,7 @@ # - Synchronize latest settings from .env.example to .env # - Preserve custom settings in existing .env # - Add new environment variables -# - Detect removed environment variables +# - Report additional local overrides # - Create backup files # ================================================================ @@ -82,7 +82,8 @@ check_files() { # Create timestamped backup of .env file # Creates env-backup directory if needed and backs up current .env file create_backup() { - local timestamp=$(date +"%Y%m%d_%H%M%S") + local timestamp + timestamp=$(date +"%Y%m%d_%H%M%S") local backup_dir="env-backup" # Create backup directory if it doesn't exist @@ -103,18 +104,20 @@ detect_differences() { log_info "Detecting differences between .env and .env.example..." # Create secure temporary directory - local temp_dir=$(mktemp -d) + local temp_dir + temp_dir=$(mktemp -d) local temp_diff="$temp_dir/env_diff" # Store diff file path as global variable - declare -g DIFF_FILE="$temp_diff" - declare -g TEMP_DIR="$temp_dir" + DIFF_FILE="$temp_diff" + TEMP_DIR="$temp_dir" # Initialize difference file - > "$temp_diff" + : > "$temp_diff" # Use awk for efficient comparison (much faster for large files) - local diff_count=$(awk -F= ' + local diff_count + diff_count=$(awk -F= ' BEGIN { OFS="\x01" } FNR==NR { if (!/^[[:space:]]*#/ && !/^[[:space:]]*$/ && /=/) { @@ -261,123 +264,50 @@ analyze_value_change() { fi } -# Synchronize .env file with .env.example while preserving custom values -# Creates a new .env file based on .env.example structure, preserving existing custom values -# Global variables used: DIFF_FILE, TEMP_DIR +# Append missing template inputs without rewriting existing overrides or comments. sync_env_file() { - log_info "Starting partial synchronization of .env file..." + log_info "Synchronizing new deployment inputs into .env..." local new_env_file=".env.new" - local preserved_count=0 - local updated_count=0 - - # Pre-process diff file for efficient lookup - local lookup_file="" - if [[ -f "$DIFF_FILE" && -s "$DIFF_FILE" ]]; then - lookup_file="${DIFF_FILE}.lookup" - # Create sorted lookup file for fast search - sort "$DIFF_FILE" > "$lookup_file" - log_info "Created lookup file for $(wc -l < "$DIFF_FILE") preserved values" - fi - - # Use AWK for efficient processing (much faster than bash loop for large files) - log_info "Processing $(wc -l < .env.example) lines with AWK..." - - local preserved_keys_file="${TEMP_DIR}/preserved_keys" - local awk_preserved_count_file="${TEMP_DIR}/awk_preserved_count" - local awk_updated_count_file="${TEMP_DIR}/awk_updated_count" - - awk -F'=' -v lookup_file="$lookup_file" -v preserved_file="$preserved_keys_file" \ - -v preserved_count_file="$awk_preserved_count_file" -v updated_count_file="$awk_updated_count_file" ' - BEGIN { - preserved_count = 0 - updated_count = 0 - - # Load preserved values if lookup file exists - if (lookup_file != "") { - while ((getline line < lookup_file) > 0) { - split(line, parts, "\x01") - key = parts[1] - value = parts[2] - preserved_values[key] = value - } - close(lookup_file) - } - } - - # Process each line - { - # Check if this is an environment variable line + cp .env "$new_env_file" + awk -F= ' + FILENAME == ARGV[1] { if (/^[[:space:]]*[A-Za-z_][A-Za-z0-9_]*[[:space:]]*=/) { - # Extract key key = $1 gsub(/^[[:space:]]+|[[:space:]]+$/, "", key) - - # Check if key should be preserved - if (key in preserved_values) { - print key "=" preserved_values[key] - print key > preserved_file - preserved_count++ - } else { - print $0 - updated_count++ - } - } else { - # Not an env var line, preserve as-is - print $0 + existing[key] = 1 } + next } - - END { - print preserved_count > preserved_count_file - print updated_count > updated_count_file + /^[[:space:]]*[A-Za-z_][A-Za-z0-9_]*[[:space:]]*=/ { + key = $1 + gsub(/^[[:space:]]+|[[:space:]]+$/, "", key) + if (!(key in existing)) { + if (!added++) print "\n# New deployment inputs from .env.example" + print $0 + existing[key] = 1 + } } - ' .env.example > "$new_env_file" - - # Read counters and preserved keys - if [[ -f "$awk_preserved_count_file" ]]; then - preserved_count=$(cat "$awk_preserved_count_file") - fi - if [[ -f "$awk_updated_count_file" ]]; then - updated_count=$(cat "$awk_updated_count_file") - fi + ' .env .env.example >> "$new_env_file" - # Show what was preserved - if [[ -f "$preserved_keys_file" ]]; then - while read -r key; do - [[ -n "$key" ]] && log_info " Preserved: $key (.env value)" - done < "$preserved_keys_file" - fi - - # Clean up lookup file - [[ -n "$lookup_file" ]] && rm -f "$lookup_file" - - # Replace the original .env file - if mv "$new_env_file" ".env"; then - log_success "Successfully created new .env file" - else + if ! mv "$new_env_file" .env; then log_error "Failed to replace .env file" rm -f "$new_env_file" return 1 fi - # Clean up difference file and temporary directory if [[ -n "${TEMP_DIR:-}" ]]; then rm -rf "${TEMP_DIR}" unset TEMP_DIR fi - if [[ -n "${DIFF_FILE:-}" ]]; then - unset DIFF_FILE - fi + unset DIFF_FILE - log_success "Partial synchronization of .env file completed" - log_info " Preserved .env values: $preserved_count" - log_info " Updated to .env.example values: $updated_count" + log_success "Synchronization completed; existing .env settings preserved" } -# Detect removed environment variables +# Report additional local overrides detect_removed_variables() { - log_info "Detecting removed environment variables..." + log_info "Checking for additional local overrides..." if [[ ! -f ".env" ]]; then return @@ -410,13 +340,13 @@ detect_removed_variables() { fi if [[ ${#removed_vars[@]} -gt 0 ]]; then - log_warning "The following environment variables have been removed from .env.example:" + log_warning "The following local overrides are not listed in .env.example:" for var in "${removed_vars[@]}"; do log_warning " - $var" done - log_warning "Consider manually removing these variables from .env" + log_warning "These overrides will be preserved in .env" else - log_success "No removed environment variables found" + log_success "No additional local overrides found" fi } @@ -424,8 +354,9 @@ detect_removed_variables() { show_statistics() { log_info "Synchronization statistics:" - local total_example=$(grep -c "^[^#]*=" .env.example 2>/dev/null || echo "0") - local total_env=$(grep -c "^[^#]*=" .env 2>/dev/null || echo "0") + local total_example total_env + total_example=$(grep -c "^[^#]*=" .env.example 2>/dev/null || echo "0") + total_env=$(grep -c "^[^#]*=" .env 2>/dev/null || echo "0") log_info " .env.example environment variables: $total_example" log_info " .env environment variables: $total_env" @@ -446,7 +377,7 @@ main() { # Detect differences detect_differences - # Detect removed variables (before sync) + # Report additional local overrides (before sync) detect_removed_variables # Synchronize environment file @@ -462,4 +393,4 @@ main() { # Execute main function only when script is run directly if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then main "$@" -fi \ No newline at end of file +fi diff --git a/docker/docker-compose-template.yaml b/docker/docker-compose-template.yaml index 8c97edc9e2dac1..928835b9e9be94 100644 --- a/docker/docker-compose-template.yaml +++ b/docker/docker-compose-template.yaml @@ -1,6 +1,7 @@ # Shared configuration using YAML anchors and env_file x-shared-api-worker-config: &shared-api-worker-config env_file: + - ./envs/core-services/shared.defaults - path: ./envs/core-services/shared.env required: false - path: ./envs/core-services/api.env @@ -67,6 +68,7 @@ x-shared-api-worker-config: &shared-api-worker-config x-shared-worker-config: &shared-worker-config env_file: + - ./envs/core-services/shared.defaults - path: ./envs/core-services/shared.env required: false - path: ./envs/core-services/worker.env @@ -141,6 +143,7 @@ x-knowledge-fs-celery-worker: &knowledge-fs-celery-worker # Match the two-hour default task deadline plus shutdown margin for long document jobs. stop_grace_period: 125m env_file: + - ./envs/core-services/shared.defaults - path: ./envs/core-services/shared.env required: false - path: ./envs/core-services/worker.env @@ -178,6 +181,7 @@ x-knowledge-fs-celery-worker: &knowledge-fs-celery-worker x-shared-worker-beat-config: &shared-worker-beat-config env_file: + - ./envs/core-services/shared.defaults - path: ./envs/core-services/shared.env required: false - path: ./envs/core-services/worker-beat.env @@ -463,6 +467,7 @@ services: image: langgenius/dify-web:1.18.0 restart: always env_file: + - ./envs/core-services/shared.defaults - path: ./envs/core-services/web.env required: false - path: ./envs/security.env @@ -508,7 +513,7 @@ services: POSTGRES_DB: ${DB_DATABASE:-dify} PGDATA: ${PGDATA:-/var/lib/postgresql/data/pgdata} command: > - postgres -c 'max_connections=${POSTGRES_MAX_CONNECTIONS:-100}' + postgres -c 'max_connections=${POSTGRES_MAX_CONNECTIONS:-200}' -c 'shared_buffers=${POSTGRES_SHARED_BUFFERS:-128MB}' -c 'work_mem=${POSTGRES_WORK_MEM:-4MB}' -c 'maintenance_work_mem=${POSTGRES_MAINTENANCE_WORK_MEM:-64MB}' diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index bfc96bcd0ed840..c4a45549a4c1b9 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -7,6 +7,7 @@ # Shared configuration using YAML anchors and env_file x-shared-api-worker-config: &shared-api-worker-config env_file: + - ./envs/core-services/shared.defaults - path: ./envs/core-services/shared.env required: false - path: ./envs/core-services/api.env @@ -73,6 +74,7 @@ x-shared-api-worker-config: &shared-api-worker-config x-shared-worker-config: &shared-worker-config env_file: + - ./envs/core-services/shared.defaults - path: ./envs/core-services/shared.env required: false - path: ./envs/core-services/worker.env @@ -147,6 +149,7 @@ x-knowledge-fs-celery-worker: &knowledge-fs-celery-worker # Match the two-hour default task deadline plus shutdown margin for long document jobs. stop_grace_period: 125m env_file: + - ./envs/core-services/shared.defaults - path: ./envs/core-services/shared.env required: false - path: ./envs/core-services/worker.env @@ -184,6 +187,7 @@ x-knowledge-fs-celery-worker: &knowledge-fs-celery-worker x-shared-worker-beat-config: &shared-worker-beat-config env_file: + - ./envs/core-services/shared.defaults - path: ./envs/core-services/shared.env required: false - path: ./envs/core-services/worker-beat.env @@ -469,6 +473,7 @@ services: image: langgenius/dify-web:1.18.0 restart: always env_file: + - ./envs/core-services/shared.defaults - path: ./envs/core-services/web.env required: false - path: ./envs/security.env @@ -514,7 +519,7 @@ services: POSTGRES_DB: ${DB_DATABASE:-dify} PGDATA: ${PGDATA:-/var/lib/postgresql/data/pgdata} command: > - postgres -c 'max_connections=${POSTGRES_MAX_CONNECTIONS:-100}' + postgres -c 'max_connections=${POSTGRES_MAX_CONNECTIONS:-200}' -c 'shared_buffers=${POSTGRES_SHARED_BUFFERS:-128MB}' -c 'work_mem=${POSTGRES_WORK_MEM:-4MB}' -c 'maintenance_work_mem=${POSTGRES_MAINTENANCE_WORK_MEM:-64MB}' diff --git a/docker/envs/core-services/dify-agent.env.example b/docker/envs/core-services/dify-agent.env.example index 1c202a4f266e49..f80ed6db4205f3 100644 --- a/docker/envs/core-services/dify-agent.env.example +++ b/docker/envs/core-services/dify-agent.env.example @@ -3,6 +3,10 @@ # ------------------------------ AGENT_BACKEND_BASE_URL=http://agent_backend:5050 +# Set this in docker/.env to share the same control-plane token with API and workers. +# Replace the development default in production. +# Generate one with: python -c 'import secrets; print(secrets.token_urlsafe(32))' +# DIFY_AGENT_API_TOKEN=dify-agent-run-token-for-dev-only # Leave empty to derive from REDIS_PASSWORD in Docker Compose. DIFY_AGENT_REDIS_URL= diff --git a/docker/envs/core-services/shared.defaults b/docker/envs/core-services/shared.defaults new file mode 100644 index 00000000000000..1e63c8e9abae3c --- /dev/null +++ b/docker/envs/core-services/shared.defaults @@ -0,0 +1,28 @@ +# Docker-specific defaults, loaded before optional env files and docker/.env. +# Keep runtime defaults in the application or Compose; only Docker overrides belong here. +# Override these values in docker/.env or the matching envs/*.env file, not here. +CONSOLE_CORS_ALLOW_ORIGINS=* +TRIGGER_URL=http://localhost +ENDPOINT_URL_TEMPLATE=http://localhost/e/{hook_id} +MIGRATION_ENABLED=true +LOG_FILE=/app/logs/server.log +LOG_DATEFORMAT=%Y-%m-%d %H:%M:%S +CELERY_WORKER_AMOUNT=4 +DB_HOST=db_postgres +DB_PASSWORD=difyai123456 +REDIS_HOST=redis +REDIS_PASSWORD=difyai123456 +REDIS_KEEPALIVE=true +CELERY_BROKER_URL=redis://:${REDIS_PASSWORD:-difyai123456}@${REDIS_HOST:-redis}:${REDIS_PORT:-6379}/1 +OPENDAL_FS_ROOT=storage +VECTOR_STORE=weaviate +WEAVIATE_ENDPOINT=http://weaviate:8080 +WEAVIATE_API_KEY=WVF5YThaHlkYwhGUSmCRgsX3tD5ngdN8pkih +WEAVIATE_GRPC_ENDPOINT=grpc://weaviate:50051 +SSRF_PROXY_HTTP_URL=http://ssrf_proxy:3128 +SSRF_PROXY_HTTPS_URL=http://ssrf_proxy:3128 +PLUGIN_DAEMON_URL=http://plugin_daemon:5002 +PLUGIN_DAEMON_KEY=lYkiYYT6owG+71oLerGzA7GXCgOT++6ovaezWAjpCjf+Sjc3ZtU+qUEi +PLUGIN_MAX_PACKAGE_SIZE=52428800 +AGENT_BACKEND_BASE_URL=http://agent_backend:5050 +NEXT_PUBLIC_ENABLE_AGENT_V2=true diff --git a/docker/envs/core-services/shared.env.example b/docker/envs/core-services/shared.env.example index 1a6280dab03d06..7fdb12c8ab2d33 100644 --- a/docker/envs/core-services/shared.env.example +++ b/docker/envs/core-services/shared.env.example @@ -7,8 +7,8 @@ SERVICE_API_URL= TRIGGER_URL=http://localhost HUMAN_INPUT_IM_EVENT_TRANSPORT_MODE=WEBHOOK APP_WEB_URL= -# FILES_URL= -# INTERNAL_FILES_URL= +FILES_URL= +INTERNAL_FILES_URL= LANG=C.UTF-8 LC_ALL=C.UTF-8 PYTHONIOENCODING=utf-8 @@ -91,6 +91,7 @@ SSRF_PROXY_HTTP_URL=http://ssrf_proxy:3128 SSRF_PROXY_HTTPS_URL=http://ssrf_proxy:3128 PGDATA=/var/lib/postgresql/data/pgdata PLUGIN_MAX_PACKAGE_SIZE=52428800 +PLUGIN_MAX_FILE_SIZE=52428800 PLUGIN_MODEL_SCHEMA_CACHE_TTL=3600 PLUGIN_MODEL_PROVIDERS_CACHE_ENABLED=true PLUGIN_MODEL_PROVIDERS_CACHE_TTL=86400 @@ -111,15 +112,15 @@ LOG_TZ=UTC DEBUG=false FLASK_DEBUG=false ENABLE_REQUEST_LOGGING=False -OPS_TRACE_UNIFIED_ENABLED=false # Application-level Agent tracing; also enable the collector in dify-agent.env. OPS_AGENT_TRACE_ENABLED=false ENABLE_APPLICATION_TRACE_RECONCILE_TASK=true -OPS_TRACE_RETRYABLE_DISPATCH_MAX_RETRIES=780 -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=3900 +# Shared by API, workers, and Beat. Keep separate from legacy ops_trace during upgrades. +OPS_TRACE_QUEUE=ops_trace_v2 +# Attempts to send one trace while its provider is temporarily unavailable. +OPS_TRACE_EXPORT_MAX_ATTEMPTS=6 +# Daily cleanup of trace payloads an interrupted delivery left in storage. +ENABLE_TRACE_PAYLOAD_SWEEP_TASK=true WORKFLOW_LOG_CLEANUP_ENABLED=false WORKFLOW_LOG_RETENTION_DAYS=30 WORKFLOW_LOG_CLEANUP_BATCH_SIZE=100 @@ -139,7 +140,7 @@ AGENT_BACKEND_STREAM_MAX_RECONNECTS=3 # Must stay above DIFY_AGENT_ENTERPRISE_SANDBOX_SNAPSHOT_TIMEOUT on the Agent backend. AGENT_BACKEND_HOME_SNAPSHOT_TIMEOUT_SECONDS=45 # Independent E2B accounting, collected only by API Celery Beat/workers. -# Uses the existing ops_trace queue; business binding operations never call metering. +# Uses OPS_TRACE_QUEUE; business binding operations never call metering. # Enable after migrations. START_AT is an immutable, whole-second UTC T0. AGENT_SANDBOX_METERING_ENABLED=false AGENT_SANDBOX_METERING_PROJECT_ID= diff --git a/docker/envs/core-services/web.env.example b/docker/envs/core-services/web.env.example index 0b826189f59ba0..e2a507ce6e5289 100644 --- a/docker/envs/core-services/web.env.example +++ b/docker/envs/core-services/web.env.example @@ -3,6 +3,7 @@ # ------------------------------ CONSOLE_API_URL= +SERVER_CONSOLE_API_URL=http://api:5001 APP_API_URL= # Optional Cloud WebApp document admission. Enable only after the edge proxy # overwrites both X-Dify-Webapp-Client-Ip and X-Dify-Webapp-Proxy-Secret, Web is @@ -29,6 +30,7 @@ MARKDOWN_FORM_FIELD_NAME_MAX_LENGTH=128 ALLOW_UNSAFE_DATA_SCHEME=false MAX_TREE_DEPTH=50 MARKETPLACE_API_URL=https://marketplace.dify.ai +MARKETPLACE_URL=https://marketplace.dify.ai INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH=4000 ALLOW_EMBED=false AMPLITUDE_API_KEY= diff --git a/oxlint-suppressions.json b/oxlint-suppressions.json index 2744ce70b1107b..93f790fd5de41c 100644 --- a/oxlint-suppressions.json +++ b/oxlint-suppressions.json @@ -820,16 +820,6 @@ "count": 2 } }, - "web/app/components/app/tracing/configuration.tsx": { - "better-tailwindcss/enforce-logical-properties": { - "count": 3 - } - }, - "web/app/components/app/tracing/provider-panel.tsx": { - "better-tailwindcss/enforce-logical-properties": { - "count": 1 - } - }, "web/app/components/app/tracing/summary-disclosure.tsx": { "better-tailwindcss/enforce-logical-properties": { "count": 1 diff --git a/packages/contracts/generated/api/console/agent/errors.gen.ts b/packages/contracts/generated/api/console/agent/errors.gen.ts index 2600e7ed265171..9eb12086b9e99a 100644 --- a/packages/contracts/generated/api/console/agent/errors.gen.ts +++ b/packages/contracts/generated/api/console/agent/errors.gen.ts @@ -611,6 +611,7 @@ export const post19Errors = { } export const get30Errors = { + tracing_provider_unavailable: { status: 400, data: zTracingProviderUnavailableError }, unsupported_tracing_provider: { status: 400, data: zUnsupportedTracingProviderError }, } diff --git a/packages/contracts/generated/api/console/agent/orpc.gen.ts b/packages/contracts/generated/api/console/agent/orpc.gen.ts index 83273430e32aca..0d969fd54b6342 100644 --- a/packages/contracts/generated/api/console/agent/orpc.gen.ts +++ b/packages/contracts/generated/api/console/agent/orpc.gen.ts @@ -1501,7 +1501,10 @@ export const get30 = oc }), ) .output(zGetAgentByAgentIdTraceCapabilitiesResponse) - .errors({ unsupported_tracing_provider: { status: 400, data: zUnsupportedTracingProviderError } }) + .errors({ + tracing_provider_unavailable: { status: 400, data: zTracingProviderUnavailableError }, + unsupported_tracing_provider: { status: 400, data: zUnsupportedTracingProviderError }, + }) export const traceCapabilities = { get: get30, diff --git a/packages/contracts/generated/api/console/agent/types.gen.ts b/packages/contracts/generated/api/console/agent/types.gen.ts index 69522a9ce4f4b3..e6392b23c7b752 100644 --- a/packages/contracts/generated/api/console/agent/types.gen.ts +++ b/packages/contracts/generated/api/console/agent/types.gen.ts @@ -651,6 +651,7 @@ export type AppNotFoundError = ErrorResponse & { export type TraceAppConfigResponse = { app_id?: string | null + configured_secrets?: Array | null created_at?: string | null error?: string | null has_not_configured?: boolean | null @@ -4300,7 +4301,7 @@ export type GetAgentByAgentIdTraceCapabilitiesData = { } export type GetAgentByAgentIdTraceCapabilitiesErrors = { - 400: UnsupportedTracingProviderError | ErrorResponse + 400: TracingProviderUnavailableError | UnsupportedTracingProviderError | ErrorResponse default: ErrorResponse } diff --git a/packages/contracts/generated/api/console/agent/zod.gen.ts b/packages/contracts/generated/api/console/agent/zod.gen.ts index 04b3c04e8c1971..f877891c08777d 100644 --- a/packages/contracts/generated/api/console/agent/zod.gen.ts +++ b/packages/contracts/generated/api/console/agent/zod.gen.ts @@ -713,6 +713,7 @@ export const zAppNotFoundError = zErrorResponse.and( */ export const zTraceAppConfigResponse = z.object({ app_id: z.string().nullish(), + configured_secrets: z.array(z.string()).nullish(), created_at: z.string().nullish(), error: z.string().nullish(), has_not_configured: z.boolean().nullish(), diff --git a/packages/contracts/generated/api/console/apps/errors.gen.ts b/packages/contracts/generated/api/console/apps/errors.gen.ts index fe263698edf6ed..af23d6b030ad5a 100644 --- a/packages/contracts/generated/api/console/apps/errors.gen.ts +++ b/packages/contracts/generated/api/console/apps/errors.gen.ts @@ -907,6 +907,7 @@ export const post37Errors = { } export const get46Errors = { + tracing_provider_unavailable: { status: 400, data: zTracingProviderUnavailableError }, unsupported_tracing_provider: { status: 400, data: zUnsupportedTracingProviderError }, } diff --git a/packages/contracts/generated/api/console/apps/orpc.gen.ts b/packages/contracts/generated/api/console/apps/orpc.gen.ts index f11cbfde52cde4..7b1dbe33560d73 100644 --- a/packages/contracts/generated/api/console/apps/orpc.gen.ts +++ b/packages/contracts/generated/api/console/apps/orpc.gen.ts @@ -2879,7 +2879,10 @@ export const get46 = oc }), ) .output(zGetAppsByAppIdTraceCapabilitiesResponse) - .errors({ unsupported_tracing_provider: { status: 400, data: zUnsupportedTracingProviderError } }) + .errors({ + tracing_provider_unavailable: { status: 400, data: zTracingProviderUnavailableError }, + unsupported_tracing_provider: { status: 400, data: zUnsupportedTracingProviderError }, + }) export const traceCapabilities = { get: get46, diff --git a/packages/contracts/generated/api/console/apps/types.gen.ts b/packages/contracts/generated/api/console/apps/types.gen.ts index c1fcd054d5878e..a83d555d18eaec 100644 --- a/packages/contracts/generated/api/console/apps/types.gen.ts +++ b/packages/contracts/generated/api/console/apps/types.gen.ts @@ -999,6 +999,7 @@ export type TraceProviderCapabilitiesResponse = { export type TraceAppConfigResponse = { app_id?: string | null + configured_secrets?: Array | null created_at?: string | null error?: string | null has_not_configured?: boolean | null @@ -7451,7 +7452,7 @@ export type GetAppsByAppIdTraceCapabilitiesData = { } export type GetAppsByAppIdTraceCapabilitiesErrors = { - 400: UnsupportedTracingProviderError | ErrorResponse + 400: TracingProviderUnavailableError | UnsupportedTracingProviderError | ErrorResponse default: ErrorResponse } diff --git a/packages/contracts/generated/api/console/apps/zod.gen.ts b/packages/contracts/generated/api/console/apps/zod.gen.ts index 88a054a43e8f04..1c76c5dfcfa22d 100644 --- a/packages/contracts/generated/api/console/apps/zod.gen.ts +++ b/packages/contracts/generated/api/console/apps/zod.gen.ts @@ -1030,6 +1030,7 @@ export const zTraceProviderCapabilitiesResponse = z.object({ */ export const zTraceAppConfigResponse = z.object({ app_id: z.string().nullish(), + configured_secrets: z.array(z.string()).nullish(), created_at: z.string().nullish(), error: z.string().nullish(), has_not_configured: z.boolean().nullish(), diff --git a/web/app/components/app/tracing/__tests__/panel.spec.tsx b/web/app/components/app/tracing/__tests__/panel.spec.tsx index cc69c04c9f64e6..317bb1c61be6ef 100644 --- a/web/app/components/app/tracing/__tests__/panel.spec.tsx +++ b/web/app/components/app/tracing/__tests__/panel.spec.tsx @@ -1,5 +1,7 @@ import type { AppMode, + TraceAppConfigResponse, + TraceConfigPayload, TraceProviderCapabilitiesResponse, } from '@dify/contracts/api/console/apps/types.gen' import { act, render, screen, waitFor, within } from '@testing-library/react' @@ -21,9 +23,20 @@ vi.mock('@/service/base', () => ({ request })) const langfuseConfig = { public_key: 'public-key', - secret_key: '[__MASKED__]', + secret_key: 'secret-key', host: 'https://cloud.langfuse.com', } +const langfuseSecrets = ['public_key', 'secret_key'] + +/** The trace-config response for a stored Langfuse config: secrets are reported, never returned. */ +function presentLangfuse(stored: TraceConfigPayload['tracing_config']): TraceAppConfigResponse { + return { + tracing_config: Object.fromEntries( + Object.entries(stored).filter(([name]) => !langfuseSecrets.includes(name)), + ), + configured_secrets: langfuseSecrets.filter((name) => !!stored[name]), + } +} const capabilities: TraceProviderCapabilitiesResponse = { provider: 'langfuse', @@ -44,7 +57,7 @@ function mockServer( initialStatus = { enabled: configured, tracing_provider: configured ? 'langfuse' : null }, ) { let status = initialStatus - let config: Record | null = configured ? langfuseConfig : null + let config: TraceConfigPayload['tracing_config'] | null = configured ? langfuseConfig : null const writes: Array<{ method: string; body: unknown }> = [] request.mockImplementation(async (_url, _init, { request: sent }) => { const url = new URL(sent.url) @@ -63,7 +76,9 @@ function mockServer( } if (sent.method === 'GET') { const providerConfig = url.searchParams.get('tracing_provider') === 'langfuse' ? config : null - return Response.json({ tracing_config: providerConfig, has_not_configured: !providerConfig }) + return Response.json( + providerConfig ? presentLangfuse(providerConfig) : { has_not_configured: true }, + ) } if (sent.method === 'DELETE') { config = null @@ -73,7 +88,8 @@ function mockServer( } const body = await sent.json() writes.push({ method: sent.method, body }) - config = { ...body.tracing_config, secret_key: '[__MASKED__]' } + // Like the API, a secret left out of the save keeps its stored value. + config = { ...config, ...body.tracing_config } if (body.activate) status = { ...status, enabled: true, tracing_provider: body.tracing_provider } return Response.json({ result: 'success' }) @@ -522,7 +538,10 @@ describe('Tracing overview panel', () => { request.mockImplementation(async (...args) => { const url = new URL(args[2].request.url) if (url.searchParams.get('tracing_provider') === 'langsmith') - return Response.json({ tracing_config: { api_key: 'key', project: 'project' } }) + return Response.json({ + tracing_config: { project: 'project' }, + configured_secrets: ['api_key'], + }) if (url.searchParams.get('tracing_provider') === 'phoenix') { if (unavailable) return Response.json( @@ -714,10 +733,10 @@ describe('Tracing overview panel', () => { Response.json({ has_not_configured: false, tracing_config: { - api_key: 'key', project: 'project', endpoint: 'https://api.smith.langchain.com', }, + configured_secrets: ['api_key'], }), ) return original(...args) @@ -792,7 +811,7 @@ describe('Tracing overview panel', () => { method: 'POST', body: { tracing_provider: 'langfuse', - tracing_config: { ...langfuseConfig, secret_key: 'secret-key' }, + tracing_config: langfuseConfig, activate: true, }, }, @@ -814,13 +833,13 @@ describe('Tracing overview panel', () => { ) }) - it('reopens an edited provider from the refreshed masked config', async () => { + it('saves a replaced secret and reopens the provider without prefilling secrets', async () => { const writes = mockServer() const { client } = renderPanel() const user = await openConfiguredPanel() await user.click(screen.getAllByText('app.tracing.config')[0]!) const secret = screen.getByRole('textbox', { name: 'app.tracing.configProvider.secretKey' }) - await user.clear(secret) + expect(secret).toHaveValue('') await user.type(secret, 'new-secret') await user.click(screen.getByRole('button', { name: 'common.operation.save' })) await waitFor(() => @@ -830,15 +849,23 @@ describe('Tracing overview panel', () => { ) await waitFor(() => expect(client.isFetching()).toBe(0)) await user.click(screen.getAllByText('app.tracing.config')[0]!) - expect( - screen.getByRole('textbox', { name: 'app.tracing.configProvider.secretKey' }), - ).toHaveValue('[__MASKED__]') + for (const label of [ + 'app.tracing.configProvider.secretKey', + 'app.tracing.configProvider.publicKey', + ]) { + const reopened = screen.getByRole('textbox', { name: label }) + expect(reopened).toHaveValue('') + expect(reopened).toHaveAttribute( + 'placeholder', + 'app.tracing.configProvider.configuredSecretPlaceholder', + ) + } expect(writes).toEqual([ { method: 'PATCH', body: { tracing_provider: 'langfuse', - tracing_config: { ...langfuseConfig, secret_key: 'new-secret' }, + tracing_config: { host: langfuseConfig.host, secret_key: 'new-secret' }, }, }, ]) @@ -1024,10 +1051,11 @@ describe('Tracing overview panel', () => { url.pathname.endsWith('/trace-config') && url.searchParams.get('tracing_provider') === 'langsmith' ) { - return Response.json({ - has_not_configured: removed, - tracing_config: removed ? null : { api_key: '[__MASKED__]', project: 'project' }, - }) + return Response.json( + removed + ? { has_not_configured: true } + : { tracing_config: { project: 'project' }, configured_secrets: ['api_key'] }, + ) } return server(...args) }) @@ -1075,7 +1103,7 @@ describe('Tracing overview panel', () => { consoleQuery.apps.byAppId.traceConfig.get.queryKey({ input: { params: { app_id: 'app-1' }, query: { tracing_provider: 'langfuse' } }, }), - { tracing_config: { ...langfuseConfig, host: 'https://server.langfuse.com' } }, + presentLangfuse({ ...langfuseConfig, host: 'https://server.langfuse.com' }), ), ) expect(host).toHaveValue('https://draft.langfuse.com') diff --git a/web/app/components/app/tracing/__tests__/provider-config-modal.spec.tsx b/web/app/components/app/tracing/__tests__/provider-config-modal.spec.tsx index df7773dd727da4..fdfa7329e060de 100644 --- a/web/app/components/app/tracing/__tests__/provider-config-modal.spec.tsx +++ b/web/app/components/app/tracing/__tests__/provider-config-modal.spec.tsx @@ -1,4 +1,7 @@ -import type { TraceConfigPayload } from '@dify/contracts/api/console/apps/types.gen' +import type { + TraceAppConfigResponse, + TraceConfigPayload, +} from '@dify/contracts/api/console/apps/types.gen' import type { ConfigurableTracingProvider } from '../providers' import type { TracingResource } from '../resource' import { Dialog, DialogContent } from '@langgenius/dify-ui/dialog' @@ -90,121 +93,165 @@ const validConfigs = { }, } satisfies Record +// The fields each provider's backend config declares with `secret()`, by input label. +const secretLabels = { + arize: { api_key: 'API Key', space_id: 'Space ID' }, + phoenix: { api_key: 'API Key' }, + langsmith: { api_key: 'API Key' }, + langfuse: { + secret_key: 'app.tracing.configProvider.secretKey', + public_key: 'app.tracing.configProvider.publicKey', + }, + opik: { api_key: 'API Key' }, + weave: { api_key: 'API Key' }, + aliyun: { license_key: 'License Key' }, + mlflow: { password: 'app.tracing.configProvider.password' }, + databricks: { + client_secret: 'app.tracing.configProvider.clientSecret', + personal_access_token: 'app.tracing.configProvider.personalAccessToken', + }, + tencent: { token: 'Token' }, +} satisfies Record> + +const configuredSecretPlaceholder = 'app.tracing.configProvider.configuredSecretPlaceholder' + +const isSecret = (provider: ConfigurableTracingProvider, name: string) => + Object.hasOwn(secretLabels[provider], name) + +/** The submitted `tracing_config` when every secret input is left blank. */ +const withoutSecrets = ( + provider: ConfigurableTracingProvider, + values: TraceConfigPayload['tracing_config'], +) => Object.fromEntries(Object.entries(values).filter(([name]) => !isSecret(provider, name))) + +/** The trace-config response for stored `values`: secrets are reported, never returned. */ +const savedConfig = ( + provider: ConfigurableTracingProvider, + values: TraceConfigPayload['tracing_config'], +): TraceAppConfigResponse => ({ + tracing_provider: provider, + tracing_config: { + ...withoutSecrets(provider, values), + project_url: 'https://tracing.example/project', + }, + configured_secrets: Object.keys(secretLabels[provider]).filter((name) => !!values[name]), +}) + const invalidConfigCases: Array<{ provider: ConfigurableTracingProvider - payload: TraceConfigPayload['tracing_config'] + values: TraceConfigPayload['tracing_config'] missingField: string }> = [ { provider: 'arize', - payload: { ...validConfigs.arize, api_key: '' }, + values: { ...validConfigs.arize, api_key: '' }, missingField: 'API Key', }, { provider: 'arize', - payload: { ...validConfigs.arize, space_id: '' }, + values: { ...validConfigs.arize, space_id: '' }, missingField: 'Space ID', }, { provider: 'arize', - payload: { ...validConfigs.arize, project: '' }, + values: { ...validConfigs.arize, project: '' }, missingField: 'app.tracing.configProvider.project', }, { provider: 'phoenix', - payload: { ...validConfigs.phoenix, project: '' }, + values: { ...validConfigs.phoenix, project: '' }, missingField: 'app.tracing.configProvider.project', }, { provider: 'langsmith', - payload: { ...validConfigs.langsmith, api_key: '' }, + values: { ...validConfigs.langsmith, api_key: '' }, missingField: 'API Key', }, { provider: 'langsmith', - payload: { ...validConfigs.langsmith, project: '' }, + values: { ...validConfigs.langsmith, project: '' }, missingField: 'app.tracing.configProvider.project', }, { provider: 'langfuse', - payload: { ...validConfigs.langfuse, secret_key: '' }, + values: { ...validConfigs.langfuse, secret_key: '' }, missingField: 'app.tracing.configProvider.secretKey', }, { provider: 'langfuse', - payload: { ...validConfigs.langfuse, public_key: '' }, + values: { ...validConfigs.langfuse, public_key: '' }, missingField: 'app.tracing.configProvider.publicKey', }, { provider: 'langfuse', - payload: { ...validConfigs.langfuse, host: '' }, + values: { ...validConfigs.langfuse, host: '' }, missingField: 'Host', }, { provider: 'weave', - payload: { ...validConfigs.weave, api_key: '' }, + values: { ...validConfigs.weave, api_key: '' }, missingField: 'API Key', }, { provider: 'weave', - payload: { ...validConfigs.weave, project: '' }, + values: { ...validConfigs.weave, project: '' }, missingField: 'app.tracing.configProvider.project', }, { provider: 'aliyun', - payload: { ...validConfigs.aliyun, app_name: '' }, + values: { ...validConfigs.aliyun, app_name: '' }, missingField: 'App Name', }, { provider: 'aliyun', - payload: { ...validConfigs.aliyun, license_key: '' }, + values: { ...validConfigs.aliyun, license_key: '' }, missingField: 'License Key', }, { provider: 'aliyun', - payload: { ...validConfigs.aliyun, endpoint: '' }, + values: { ...validConfigs.aliyun, endpoint: '' }, missingField: 'Endpoint', }, { provider: 'mlflow', - payload: { ...validConfigs.mlflow, tracking_uri: '' }, + values: { ...validConfigs.mlflow, tracking_uri: '' }, missingField: 'app.tracing.configProvider.trackingUri', }, { provider: 'databricks', - payload: { ...validConfigs.databricks, experiment_id: '' }, + values: { ...validConfigs.databricks, experiment_id: '' }, missingField: 'app.tracing.configProvider.experimentId', }, { provider: 'databricks', - payload: { ...validConfigs.databricks, host: '' }, + values: { ...validConfigs.databricks, host: '' }, missingField: 'app.tracing.configProvider.databricksHost', }, { provider: 'tencent', - payload: { ...validConfigs.tencent, token: '' }, + values: { ...validConfigs.tencent, token: '' }, missingField: 'Token', }, { provider: 'tencent', - payload: { ...validConfigs.tencent, endpoint: '' }, + values: { ...validConfigs.tencent, endpoint: '' }, missingField: 'Endpoint', }, { provider: 'tencent', - payload: { ...validConfigs.tencent, service_name: '' }, + values: { ...validConfigs.tencent, service_name: '' }, missingField: 'Service Name', }, ] const renderProviderConfigModal = ({ type = 'langfuse', - payload, + config, inUse = false, resource = { kind: 'app', appId: 'app-id', mode: 'chat' }, }: { type?: ConfigurableTracingProvider - payload?: TraceConfigPayload['tracing_config'] | null + config?: TraceAppConfigResponse inUse?: boolean resource?: TracingResource } = {}) => { @@ -227,7 +274,7 @@ const renderProviderConfigModal = ({ { it('creates an anonymous Phoenix destination without a key', async () => { const user = userEvent.setup() const callbacks = renderProviderConfigModal({ type: 'phoenix' }) + expect(screen.getByRole('textbox', { name: 'API Key' })).toHaveAttribute( + 'placeholder', + 'app.tracing.configProvider.placeholder:{"key":"API Key"}', + ) await user.type( screen.getByRole('textbox', { name: 'app.tracing.configProvider.project' }), 'local-project', @@ -265,7 +316,6 @@ describe('ProviderConfigModal', () => { tracing_provider: 'phoenix', activate: true, tracing_config: { - api_key: '', project: 'local-project', endpoint: 'http://localhost:6006', }, @@ -296,39 +346,28 @@ describe('ProviderConfigModal', () => { }) it.each(configurableProviders)( - 'should update valid %s config in edit mode', + 'updates a saved %s config without prefilling or resending its secrets', async (provider) => { const user = userEvent.setup() const callbacks = renderProviderConfigModal({ type: provider, - payload: validConfigs[provider], + config: savedConfig(provider, validConfigs[provider]), }) - const credentialLabels = { - arize: ['API Key'], - phoenix: ['API Key'], - langsmith: ['API Key'], - langfuse: ['app.tracing.configProvider.secretKey'], - opik: ['API Key'], - weave: ['API Key'], - aliyun: ['License Key'], - mlflow: ['app.tracing.configProvider.password'], - databricks: [ - 'app.tracing.configProvider.clientSecret', - 'app.tracing.configProvider.personalAccessToken', - ], - tencent: ['Token'], - } satisfies Record - for (const label of credentialLabels[provider]) { - const credential = screen.getByRole('textbox', { name: label }) - expect(credential).toHaveAttribute('autocomplete', 'off') + for (const label of Object.values(secretLabels[provider])) { + const secret = screen.getByRole('textbox', { name: label }) + expect(secret).toHaveValue('') + expect(secret).toHaveAttribute('placeholder', configuredSecretPlaceholder) + expect(secret).not.toBeRequired() + expect(secret).toHaveAttribute('autocomplete', 'off') } const edits = { - arize: ['Space ID', 'space_id', 'new-space'], + arize: ['app.tracing.configProvider.project', 'project', 'new-project'], phoenix: ['app.tracing.configProvider.project', 'project', 'new-project'], langsmith: ['app.tracing.configProvider.project', 'project', 'new-project'], - langfuse: ['Host', 'host', 'https://new.langfuse.com'], + // Every other Langfuse field is a secret or a destination. + langfuse: ['Host', 'host', 'https://cloud.langfuse.com/'], opik: ['Workspace', 'workspace', 'new-workspace'], weave: ['Entity', 'entity', 'new-entity'], aliyun: ['App Name', 'app_name', 'new-app'], @@ -345,31 +384,142 @@ describe('ProviderConfigModal', () => { await waitFor(() => expect(callbacks.onSaved).toHaveBeenCalledOnce()) await expectRequest('PATCH', { tracing_provider: provider, - tracing_config: { ...validConfigs[provider], [key]: value }, + tracing_config: { ...withoutSecrets(provider, validConfigs[provider]), [key]: value }, }) - expect(request).toHaveBeenCalledOnce() }, ) + it('sends a newly entered secret and keeps the other saved secret', async () => { + const user = userEvent.setup() + const callbacks = renderProviderConfigModal({ + type: 'langfuse', + config: savedConfig('langfuse', validConfigs.langfuse), + }) + + await user.type( + screen.getByRole('textbox', { name: 'app.tracing.configProvider.secretKey' }), + 'rotated-secret-key', + ) + await user.click(screen.getByRole('button', { name: 'common.operation.save' })) + + await waitFor(() => expect(callbacks.onSaved).toHaveBeenCalledOnce()) + await expectRequest('PATCH', { + tracing_provider: 'langfuse', + tracing_config: { host: validConfigs.langfuse.host, secret_key: 'rotated-secret-key' }, + }) + }) + + it('requires the saved secrets again once the destination changes', async () => { + const user = userEvent.setup() + const callbacks = renderProviderConfigModal({ + type: 'langfuse', + config: savedConfig('langfuse', validConfigs.langfuse), + }) + const secretKey = screen.getByRole('textbox', { + name: 'app.tracing.configProvider.secretKey', + }) + const publicKey = screen.getByRole('textbox', { + name: 'app.tracing.configProvider.publicKey', + }) + expect(secretKey).not.toBeRequired() + + const host = screen.getByRole('textbox', { name: 'Host' }) + await user.clear(host) + await user.type(host, 'https://us.cloud.langfuse.com') + + for (const [secret, label] of [ + [secretKey, 'app.tracing.configProvider.secretKey'], + [publicKey, 'app.tracing.configProvider.publicKey'], + ] as const) { + expect(secret).toBeRequired() + expect(secret.labels?.[0]).toHaveTextContent(`${label}*`) + expect(secret).toHaveAttribute( + 'placeholder', + `app.tracing.configProvider.placeholder:{"key":"${label}"}`, + ) + } + await user.click(screen.getByRole('button', { name: 'common.operation.save' })) + + expect(request).not.toHaveBeenCalled() + expect(secretKey).toBeInvalid() + expect(secretKey).toHaveFocus() + + await user.type(secretKey, 'us-secret-key') + await user.type(publicKey, 'us-public-key') + await user.click(screen.getByRole('button', { name: 'common.operation.save' })) + + await waitFor(() => expect(callbacks.onSaved).toHaveBeenCalledOnce()) + await expectRequest('PATCH', { + tracing_provider: 'langfuse', + tracing_config: { + host: 'https://us.cloud.langfuse.com', + secret_key: 'us-secret-key', + public_key: 'us-public-key', + }, + }) + }) + + it('keeps blank secrets while the destination is unchanged, even with an added trailing slash', async () => { + const user = userEvent.setup() + const callbacks = renderProviderConfigModal({ + type: 'langfuse', + config: savedConfig('langfuse', validConfigs.langfuse), + }) + const secrets = Object.values(secretLabels.langfuse).map((label) => + screen.getByRole('textbox', { name: label }), + ) + const expectSavedSecretsKept = () => { + for (const secret of secrets) { + expect(secret).not.toBeRequired() + expect(secret).toHaveAttribute('placeholder', configuredSecretPlaceholder) + } + } + expectSavedSecretsKept() + + const host = screen.getByRole('textbox', { name: 'Host' }) + await user.clear(host) + await user.type(host, 'https://us.cloud.langfuse.com') + for (const secret of secrets) expect(secret).toBeRequired() + await user.clear(host) + await user.type(host, `${validConfigs.langfuse.host}/`) + expectSavedSecretsKept() + await user.click(screen.getByRole('button', { name: 'common.operation.save' })) + + await waitFor(() => expect(callbacks.onSaved).toHaveBeenCalledOnce()) + await expectRequest('PATCH', { + tracing_provider: 'langfuse', + tracing_config: { host: `${validConfigs.langfuse.host}/` }, + }) + }) + it.each([ ['langfuse', 'host', 'HTTPS://cloud.langfuse.com'], ['langsmith', 'endpoint', 'HTTPS://api.smith.langchain.com'], ['opik', 'url', 'HTTPS://www.comet.com/opik/api/'], ] as const)('accepts case-insensitive URL schemes for %s', async (provider, key, value) => { const user = userEvent.setup() - const payload = { ...validConfigs[provider], [key]: value } - const callbacks = renderProviderConfigModal({ type: provider, payload }) + const values = { ...validConfigs[provider], [key]: value } + const callbacks = renderProviderConfigModal({ + type: provider, + config: savedConfig(provider, values), + }) await user.click(screen.getByRole('button', { name: 'common.operation.save' })) await waitFor(() => expect(callbacks.onSaved).toHaveBeenCalledOnce()) - await expectRequest('PATCH', { tracing_provider: provider, tracing_config: payload }) + await expectRequest('PATCH', { + tracing_provider: provider, + tracing_config: withoutSecrets(provider, values), + }) }) it('displays nullable provider fields as empty without changing their submitted values', async () => { const user = userEvent.setup() - const payload = { ...validConfigs.mlflow, username: null } - const callbacks = renderProviderConfigModal({ type: 'mlflow', payload }) + const values = { ...validConfigs.mlflow, username: null } + const callbacks = renderProviderConfigModal({ + type: 'mlflow', + config: savedConfig('mlflow', values), + }) expect( screen.getByRole('textbox', { name: 'app.tracing.configProvider.username' }), @@ -379,7 +529,7 @@ describe('ProviderConfigModal', () => { await waitFor(() => expect(callbacks.onSaved).toHaveBeenCalledOnce()) await expectRequest('PATCH', { tracing_provider: 'mlflow', - tracing_config: payload, + tracing_config: withoutSecrets('mlflow', values), }) }) @@ -396,11 +546,14 @@ describe('ProviderConfigModal', () => { ) const logError = vi.spyOn(console, 'error').mockImplementation(() => {}) const callbacks = renderProviderConfigModal({ - payload: validConfigs.langfuse, + config: savedConfig('langfuse', validConfigs.langfuse), }) const host = screen.getByRole('textbox', { name: 'Host' }) await user.clear(host) await user.type(host, 'https://draft.langfuse.com') + // A new destination needs its secrets entered again. + for (const label of Object.values(secretLabels.langfuse)) + await user.type(screen.getByRole('textbox', { name: label }), 'draft-key') const save = screen.getByRole('button', { name: 'common.operation.save' }) try { @@ -431,23 +584,103 @@ describe('ProviderConfigModal', () => { }) it.each([ - ['langfuse', 'host', 'Host', 'ftp://cloud.langfuse.com'], - ['langsmith', 'endpoint', 'Endpoint', 'http://api.smith.langchain.com'], - ['weave', 'endpoint', 'Endpoint', 'http://trace.wandb.ai'], - ['opik', 'url', 'Url', 'https://www.comet.com/opik/'], - ['mlflow', 'experiment_id', 'app.tracing.configProvider.experimentId', '1.5'], - ['databricks', 'experiment_id', 'app.tracing.configProvider.experimentId', '-1'], - ['aliyun', 'license_key', 'License Key', ' '], - ['tencent', 'token', 'Token', ' '], + ['langfuse', 'host', 'Host', 'ftp://cloud.langfuse.com', 'https://us.cloud.langfuse.com'], + [ + 'langsmith', + 'endpoint', + 'Endpoint', + 'http://api.smith.langchain.com', + 'https://eu.api.smith.langchain.com', + ], + ['weave', 'endpoint', 'Endpoint', 'http://trace.wandb.ai', 'https://weave.example.com'], + ['opik', 'url', 'Url', 'https://www.comet.com/opik/', 'https://opik.example.com/api/'], + ] as const)( + 'validates a new %s %s inline before sending the re-entered secrets', + async (provider, key, label, invalid, valid) => { + const user = userEvent.setup() + const callbacks = renderProviderConfigModal({ + type: provider, + config: savedConfig(provider, validConfigs[provider]), + }) + const field = screen.getByRole('textbox', { name: label }) + await user.clear(field) + await user.type(field, invalid) + const reentered: TraceConfigPayload['tracing_config'] = {} + for (const [name, secretLabel] of Object.entries(secretLabels[provider])) { + reentered[name] = `new-${name}` + await user.type(screen.getByRole('textbox', { name: secretLabel }), `new-${name}`) + } + await user.click(screen.getByRole('button', { name: 'common.operation.save' })) + + expect(request).not.toHaveBeenCalled() + expect(field).toBeInvalid() + expect(field).toHaveFocus() + expect(field).toHaveAccessibleDescription() + + await user.clear(field) + await user.type(field, valid) + await user.keyboard('{Enter}') + + await waitFor(() => expect(callbacks.onSaved).toHaveBeenCalledOnce()) + expect(field).toBeValid() + await expectRequest('PATCH', { + tracing_provider: provider, + tracing_config: { + ...withoutSecrets(provider, validConfigs[provider]), + [key]: valid, + ...reentered, + }, + }) + }, + ) + + it.each([ + ['mlflow', '1.5'], + ['databricks', '-1'], + ] as const)('validates the %s experiment ID inline before sending', async (provider, value) => { + const user = userEvent.setup() + const callbacks = renderProviderConfigModal({ + type: provider, + config: savedConfig(provider, validConfigs[provider]), + }) + const field = screen.getByRole('textbox', { name: 'app.tracing.configProvider.experimentId' }) + await user.clear(field) + await user.type(field, value) + await user.click(screen.getByRole('button', { name: 'common.operation.save' })) + + expect(request).not.toHaveBeenCalled() + expect(field).toBeInvalid() + expect(field).toHaveFocus() + expect(field).toHaveAccessibleDescription() + + await user.clear(field) + await user.type(field, '456') + await user.keyboard('{Enter}') + + await waitFor(() => expect(callbacks.onSaved).toHaveBeenCalledOnce()) + expect(field).toBeValid() + await expectRequest('PATCH', { + tracing_provider: provider, + tracing_config: { + ...withoutSecrets(provider, validConfigs[provider]), + experiment_id: '456', + }, + }) + }) + + it.each([ + ['aliyun', 'license_key', 'License Key'], + ['tencent', 'token', 'Token'], ] as const)( - 'validates %s %s inline before sending credentials', - async (provider, key, label, value) => { + 'rejects a whitespace-only replacement for the saved %s %s', + async (provider, key, label) => { const user = userEvent.setup() const callbacks = renderProviderConfigModal({ type: provider, - payload: { ...validConfigs[provider], [key]: value }, + config: savedConfig(provider, validConfigs[provider]), }) const field = screen.getByRole('textbox', { name: label }) + await user.type(field, ' ') await user.click(screen.getByRole('button', { name: 'common.operation.save' })) expect(request).not.toHaveBeenCalled() @@ -456,26 +689,28 @@ describe('ProviderConfigModal', () => { expect(field).toHaveAccessibleDescription() await user.clear(field) - const validConfig: TraceConfigPayload['tracing_config'] = validConfigs[provider] - await user.type(field, validConfig[key] ?? '') + await user.type(field, 'replacement-secret') await user.keyboard('{Enter}') await waitFor(() => expect(callbacks.onSaved).toHaveBeenCalledOnce()) expect(field).toBeValid() await expectRequest('PATCH', { tracing_provider: provider, - tracing_config: validConfigs[provider], + tracing_config: { + ...withoutSecrets(provider, validConfigs[provider]), + [key]: 'replacement-secret', + }, }) }, ) it.each(invalidConfigCases)( 'should reject $provider config when $missingField is missing', - async ({ provider, payload, missingField }) => { + async ({ provider, values, missingField }) => { const user = userEvent.setup() renderProviderConfigModal({ type: provider, - payload, + config: savedConfig(provider, values), }) await user.click(screen.getByRole('button', { name: 'common.operation.save' })) @@ -503,7 +738,7 @@ describe('ProviderConfigModal', () => { const user = userEvent.setup() const callbacks = renderProviderConfigModal({ type: 'langfuse', - payload: validConfigs.langfuse, + config: savedConfig('langfuse', validConfigs.langfuse), }) await user.click(screen.getByRole('button', { name: 'common.operation.remove' })) @@ -528,7 +763,11 @@ describe('ProviderConfigModal', () => { [false, 'removeConfirmContentNotInUse'], ])('describes the removal consequence when in use is %s', async (inUse, key) => { const user = userEvent.setup() - renderProviderConfigModal({ type: 'langfuse', payload: validConfigs.langfuse, inUse }) + renderProviderConfigModal({ + type: 'langfuse', + config: savedConfig('langfuse', validConfigs.langfuse), + inUse, + }) await user.click(screen.getByRole('button', { name: 'common.operation.remove' })) @@ -544,7 +783,10 @@ describe('ProviderConfigModal', () => { await httpErrorFromResponse(Response.json({ message: 'Removal failed' }, { status: 400 })), ) const user = userEvent.setup() - renderProviderConfigModal({ payload: validConfigs.langfuse }) + renderProviderConfigModal({ + type: 'langfuse', + config: savedConfig('langfuse', validConfigs.langfuse), + }) const host = screen.getByRole('textbox', { name: 'Host' }) await user.clear(host) await user.type(host, 'https://draft.langfuse.com') @@ -568,7 +810,7 @@ describe('ProviderConfigModal', () => { const user = userEvent.setup() renderProviderConfigModal({ type: 'langfuse', - payload: validConfigs.langfuse, + config: savedConfig('langfuse', validConfigs.langfuse), }) const host = screen.getByRole('textbox', { name: 'Host' }) diff --git a/web/app/components/app/tracing/__tests__/selection.browser.spec.tsx b/web/app/components/app/tracing/__tests__/selection.browser.spec.tsx index 357b5746267ba4..cebc8da687766d 100644 --- a/web/app/components/app/tracing/__tests__/selection.browser.spec.tsx +++ b/web/app/components/app/tracing/__tests__/selection.browser.spec.tsx @@ -36,11 +36,17 @@ it('preserves provider focus during an async save and allows keyboard selection const provider = url.searchParams.get('tracing_provider') const config = provider === 'langfuse' - ? { public_key: 'public', secret_key: 'secret', host: 'https://cloud.langfuse.com' } + ? { + tracing_config: { host: 'https://cloud.langfuse.com' }, + configured_secrets: ['public_key', 'secret_key'], + } : provider === 'langsmith' - ? { api_key: 'key', project: 'project', endpoint: 'https://api.smith.langchain.com' } - : null - return Response.json({ has_not_configured: !config, tracing_config: config }) + ? { + tracing_config: { project: 'project', endpoint: 'https://api.smith.langchain.com' }, + configured_secrets: ['api_key'], + } + : { has_not_configured: true } + return Response.json(config) }) const QueryProvider = createQueryClientWrapper(createTestQueryClient()) const { wrapper: NuqsProvider } = createNuqsTestWrapper({ hasMemory: true }) diff --git a/web/app/components/app/tracing/config-popup.tsx b/web/app/components/app/tracing/config-popup.tsx index 25390c369e9b62..713a7c25b09b48 100644 --- a/web/app/components/app/tracing/config-popup.tsx +++ b/web/app/components/app/tracing/config-popup.tsx @@ -37,7 +37,7 @@ type ConfigPopupProps = { policyError?: string policy?: AppTraceResponse onPolicyChange: (changes: Partial) => Promise - configs: Partial> + configs: Partial> } export function ConfigPopup({ @@ -240,7 +240,7 @@ export function ConfigPopup({ {inUseProviderIconClassName && ( - + )} @@ -204,7 +204,7 @@ export function TracingConfiguration({ resource, readOnly }: TracingConfiguratio {t(($) => $[`${I18N_PREFIX}.${enabled ? 'enabled' : 'disabled'}`], { ns: 'app' })} - + diff --git a/web/app/components/app/tracing/provider-config-modal.tsx b/web/app/components/app/tracing/provider-config-modal.tsx index 01434f988fd78f..c5d08c8fc94e5f 100644 --- a/web/app/components/app/tracing/provider-config-modal.tsx +++ b/web/app/components/app/tracing/provider-config-modal.tsx @@ -4,7 +4,7 @@ import type { TraceConfigPayload, } from '@dify/contracts/api/console/apps/types.gen' import type { RefObject } from 'react' -import type { ConfigurableTracingProvider } from './providers' +import type { ConfigurableTracingProvider, TracingProviderField } from './providers' import type { TracingResource } from './resource' import { AlertDialog, @@ -31,6 +31,7 @@ import { Form } from '@langgenius/dify-ui/form' import { Input } from '@langgenius/dify-ui/input' import { Separator } from '@langgenius/dify-ui/separator' import { useMutation } from '@tanstack/react-query' +import { useState } from 'react' import { useTranslation } from 'react-i18next' import { EncryptedBottom } from '@/app/components/base/encrypted-bottom' import { consoleQuery } from '@/service/console' @@ -39,7 +40,8 @@ import { tracingProviders } from './providers' type ProviderConfigModalProps = { resource: TracingResource type: ConfigurableTracingProvider - payload?: TraceAppConfigResponse['tracing_config'] + /** The saved configuration; absent while the provider is not configured yet. */ + config: TraceAppConfigResponse | null | undefined /** Whether traces currently go to this provider, so removing it turns tracing off. */ inUse: boolean onRemoved: () => void @@ -49,10 +51,15 @@ type ProviderConfigModalProps = { const I18N_PREFIX = 'tracing.configProvider' +/** A destination as the API compares it: trimmed, without trailing slashes. */ +function comparableDestination(value: string | null) { + return value === null ? null : value.trim().replace(/\/+$/, '') +} + export function ProviderConfigModal({ resource, type, - payload, + config, inUse, onRemoved, firstInputRef, @@ -80,7 +87,27 @@ export function ProviderConfigModal({ const createConfig = resource.kind === 'agent' ? createAgentConfig : createAppConfig const updateConfig = resource.kind === 'agent' ? updateAgentConfig : updateAppConfig const removeConfig = resource.kind === 'agent' ? removeAgentConfig : removeAppConfig - const isEdit = !!payload + const fields: readonly TracingProviderField[] = tracingProviders[type].fields + const savedValues = config?.tracing_config + const configuredSecrets = config?.configured_secrets ?? [] + // Mirrors the destination inputs, which start from the values loaded when the form opened. + const [destinationDrafts, setDestinationDrafts] = useState(() => + Object.fromEntries( + fields + .filter((field) => field.destination) + .map((field) => [field.name, savedValues?.[field.name] ?? '']), + ), + ) + const submittedValue = (name: string, value: string) => + value === '' && savedValues?.[name] === null ? null : value + // The API keeps saved secrets only for the destination they were entered for. + const destinationChanged = fields.some( + (field) => + field.destination && + comparableDestination(submittedValue(field.name, destinationDrafts[field.name] ?? '')) !== + comparableDestination(savedValues?.[field.name] ?? null), + ) + const isEdit = !!config const saveError = createConfig.error ?? updateConfig.error const pending = createConfig.isPending || updateConfig.isPending || removeConfig.isPending const handleRemove = () => { @@ -101,10 +128,17 @@ export function ProviderConfigModal({ if (pending) return if (isEdit) createConfig.reset() else updateConfig.reset() - const config: TraceConfigPayload['tracing_config'] = {} - for (const [key, value] of Object.entries(values)) - config[key] = value === '' && payload?.[key] === null ? null : value - const body: TraceConfigPayload = { tracing_provider: type, tracing_config: config } + const tracingConfig: TraceConfigPayload['tracing_config'] = {} + for (const field of fields) { + const value = values[field.name] ?? '' + if (field.secret) { + // A secret left out keeps its saved value, so only a newly entered one is sent. + if (value !== '') tracingConfig[field.name] = value + continue + } + tracingConfig[field.name] = submittedValue(field.name, value) + } + const body: TraceConfigPayload = { tracing_provider: type, tracing_config: tracingConfig } if (!isEdit) body.activate = true if (resource.kind === 'agent') (isEdit ? updateAgentConfig : createAgentConfig).mutate( @@ -138,594 +172,52 @@ export function ProviderConfigModal({ >
- {type === 'arize' && ( - <> - - - API Key - - * - - - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: 'API Key', - })} - /> - - - - - Space ID - - * - - - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: 'Space ID', - })} - /> - - - - - {t(($) => $[`${I18N_PREFIX}.project`], { ns: 'app' })} - - * - - - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: t(($) => $[`${I18N_PREFIX}.project`], { ns: 'app' }), - })} - /> - - - - Endpoint - - - - - )} - {type === 'phoenix' && ( - <> - - API Key - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: 'API Key', - })} - /> - - - - - {t(($) => $[`${I18N_PREFIX}.project`], { ns: 'app' })} - - * - - - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: t(($) => $[`${I18N_PREFIX}.project`], { ns: 'app' }), - })} - /> - - - - Endpoint - - - - - )} - {type === 'aliyun' && ( - <> - - - License Key - - * - - - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: 'License Key', - })} - /> - - - - - Endpoint - - * - - - - - - - - App Name - - * - - - - - - - )} - {type === 'tencent' && ( - <> - - - Token - - * - - - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: 'Token', - })} - /> - - - - - Endpoint - - * - - - - - {t(($) => $['tracing.tencent.endpointHelp'], { ns: 'app' })} - - - - - - Service Name - - * - - - - - {t(($) => $['tracing.tencent.serviceNameHelp'], { ns: 'app' })} - - - - - )} - {type === 'weave' && ( - <> - - - API Key - - * - - - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: 'API Key', - })} - /> - - - + {fields.map((field, index) => { + const { label, description } = field + const labelText = + typeof label === 'string' ? label : t(($) => $[label.key], { ns: 'app' }) + const keepsSavedSecret = + field.secret && !destinationChanged && configuredSecrets.includes(field.name) + const required = field.required && !keepsSavedSecret + return ( + - {t(($) => $[`${I18N_PREFIX}.project`], { ns: 'app' })} - - * - + {labelText} + {required && ( + + * + + )} $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: t(($) => $[`${I18N_PREFIX}.project`], { ns: 'app' }), - })} - /> - - - - Entity - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: 'Entity', - })} - /> - - - - Endpoint - - - - - Host - - - - - )} - {type === 'langsmith' && ( - <> - - - API Key - - * - - - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: 'API Key', - })} - /> - - - - - {t(($) => $[`${I18N_PREFIX}.project`], { ns: 'app' })} - - * - - - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: t(($) => $[`${I18N_PREFIX}.project`], { ns: 'app' }), - })} - /> - - - - Endpoint - - - - - )} - {type === 'langfuse' && ( - <> - - - {t(($) => $[`${I18N_PREFIX}.secretKey`], { ns: 'app' })} - - * - - - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: t(($) => $[`${I18N_PREFIX}.secretKey`], { ns: 'app' }), - })} - /> - - - - - {t(($) => $[`${I18N_PREFIX}.publicKey`], { ns: 'app' })} - - * - - - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: t(($) => $[`${I18N_PREFIX}.publicKey`], { ns: 'app' }), - })} - /> - - - - - Host - - * - - - - - - - )} - {type === 'opik' && ( - <> - - API Key - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: 'API Key', - })} - /> - - - - {t(($) => $[`${I18N_PREFIX}.project`], { ns: 'app' })} - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: t(($) => $[`${I18N_PREFIX}.project`], { ns: 'app' }), - })} - /> - - - - Workspace - - - - - Url - - - - - )} - {type === 'mlflow' && ( - <> - - - {t(($) => $[`${I18N_PREFIX}.trackingUri`], { ns: 'app' })} - - * - - - - - - - - {t(($) => $[`${I18N_PREFIX}.experimentId`], { ns: 'app' })} - - * - - - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: t(($) => $[`${I18N_PREFIX}.experimentId`], { ns: 'app' }), - })} - /> - - - - {t(($) => $[`${I18N_PREFIX}.username`], { ns: 'app' })} - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: t(($) => $[`${I18N_PREFIX}.username`], { ns: 'app' }), - })} - /> - - - - {t(($) => $[`${I18N_PREFIX}.password`], { ns: 'app' })} - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: t(($) => $[`${I18N_PREFIX}.password`], { ns: 'app' }), - })} - /> - - - - )} - {type === 'databricks' && ( - <> - - - {t(($) => $[`${I18N_PREFIX}.experimentId`], { ns: 'app' })} - - * - - - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: t(($) => $[`${I18N_PREFIX}.experimentId`], { ns: 'app' }), - })} - required - /> - - - - - {t(($) => $[`${I18N_PREFIX}.databricksHost`], { ns: 'app' })} - - * - - - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: t(($) => $[`${I18N_PREFIX}.databricksHost`], { ns: 'app' }), - })} - required - /> - - - - {t(($) => $[`${I18N_PREFIX}.clientId`], { ns: 'app' })} - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: t(($) => $[`${I18N_PREFIX}.clientId`], { ns: 'app' }), - })} - /> - - - - - {t(($) => $[`${I18N_PREFIX}.clientSecret`], { ns: 'app' })} - - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: t(($) => $[`${I18N_PREFIX}.clientSecret`], { ns: 'app' }), - })} - /> - - - - - {t(($) => $[`${I18N_PREFIX}.personalAccessToken`], { ns: 'app' })} - - $[`${I18N_PREFIX}.placeholder`], { - ns: 'app', - key: t(($) => $[`${I18N_PREFIX}.personalAccessToken`], { ns: 'app' }), - })} - /> + ref={index === 0 ? firstInputRef : undefined} + autoComplete={field.secret ? 'off' : undefined} + type={field.type} + inputMode={field.inputMode} + pattern={field.pattern} + maxLength={field.maxLength} + required={required} + defaultValue={field.secret ? '' : (savedValues?.[field.name] ?? '')} + onValueChange={ + field.destination + ? (value) => + setDestinationDrafts((drafts) => ({ ...drafts, [field.name]: value })) + : undefined + } + placeholder={ + keepsSavedSecret + ? t(($) => $[`${I18N_PREFIX}.configuredSecretPlaceholder`], { ns: 'app' }) + : (field.placeholder ?? + t(($) => $[`${I18N_PREFIX}.placeholder`], { ns: 'app', key: labelText })) + } + /> + {description && ( + {t(($) => $[description], { ns: 'app' })} + )} - - )} + ) + })}
{saveError && (

diff --git a/web/app/components/app/tracing/provider-panel.tsx b/web/app/components/app/tracing/provider-panel.tsx index 2c8df2a8ffd8c3..36da6e29348ab7 100644 --- a/web/app/components/app/tracing/provider-panel.tsx +++ b/web/app/components/app/tracing/provider-panel.tsx @@ -26,7 +26,7 @@ type ProviderPanelProps = { readOnly: boolean pending: boolean isChosen: boolean - config: TraceAppConfigResponse['tracing_config'] + config: TraceAppConfigResponse | null | undefined configError?: ProviderConfigError } @@ -46,6 +46,7 @@ export function ProviderPanel({ }) const iconClassName = tracingProviders[type].panelIcon const hasConfigured = !!config + const projectUrl = config?.tracing_config?.project_url return (

{isChosen && ( -
+
{t(($) => $[`${I18N_PREFIX}.inUse`], { ns: 'app' })}
)} @@ -86,11 +87,11 @@ export function ProviderPanel({ {hasConfigured && ( {t(($) => $[`${I18N_PREFIX}.view`], { ns: 'app' })} diff --git a/web/app/components/app/tracing/providers.ts b/web/app/components/app/tracing/providers.ts index 747f22cd4c2cd8..b5221d8d60ea11 100644 --- a/web/app/components/app/tracing/providers.ts +++ b/web/app/components/app/tracing/providers.ts @@ -1,60 +1,383 @@ import type { TraceConfigPayload } from '@dify/contracts/api/console/apps/types.gen' +type FieldLabelKey = `tracing.configProvider.${ + | 'clientId' + | 'clientSecret' + | 'databricksHost' + | 'experimentId' + | 'password' + | 'personalAccessToken' + | 'project' + | 'publicKey' + | 'secretKey' + | 'trackingUri' + | 'username'}` + +/** One input of a provider form; it submits the `tracing_config` entry named `name`. */ +export type TracingProviderField = { + name: string + /** Product terms such as "API Key" are shown as written; other labels are translated. */ + label: string | { key: FieldLabelKey } + required: boolean + /** + * Declared with `secret()` in the provider's backend config: the API never returns it, and a + * secret left out of a save keeps its saved value while no destination changes. + */ + secret: boolean + /** + * Declared with `destination()` in the provider's backend config: it names where the secrets + * are sent, so changing it means the secrets must be entered again. + */ + destination: boolean + /** An example value; without one, the placeholder asks for the label. */ + placeholder?: string + description?: 'tracing.tencent.endpointHelp' | 'tracing.tencent.serviceNameHelp' + type?: 'url' + inputMode?: 'numeric' + pattern?: string + maxLength?: number +} + +const httpUrlPattern = '[Hh][Tt][Tt][Pp][Ss]?://.+' +const httpsUrlPattern = '[Hh][Tt][Tt][Pp][Ss]://.+' +const nonBlankPattern = String.raw`.*\S.*` +const integerPattern = String.raw`\s*[0-9]+\s*` + export const tracingProviders = { langfuse: { docUrl: 'https://docs.langfuse.com', icon: 'i-custom-public-tracing-langfuse-icon w-18.5', panelIcon: 'i-custom-public-tracing-langfuse-icon-big w-27.75', + fields: [ + { + name: 'secret_key', + label: { key: 'tracing.configProvider.secretKey' }, + required: true, + secret: true, + destination: false, + }, + { + name: 'public_key', + label: { key: 'tracing.configProvider.publicKey' }, + required: true, + secret: true, + destination: false, + }, + { + name: 'host', + label: 'Host', + required: true, + secret: false, + destination: true, + type: 'url', + pattern: httpUrlPattern, + placeholder: 'https://cloud.langfuse.com', + }, + ], }, langsmith: { docUrl: 'https://docs.smith.langchain.com/', icon: 'i-custom-public-tracing-langsmith-icon w-21', panelIcon: 'i-custom-public-tracing-langsmith-icon-big w-31', + fields: [ + { name: 'api_key', label: 'API Key', required: true, secret: true, destination: false }, + { + name: 'project', + label: { key: 'tracing.configProvider.project' }, + required: true, + secret: false, + destination: false, + }, + { + name: 'endpoint', + label: 'Endpoint', + required: false, + secret: false, + destination: true, + type: 'url', + pattern: httpsUrlPattern, + placeholder: 'https://api.smith.langchain.com', + }, + ], }, opik: { docUrl: 'https://www.comet.com/docs/opik/integrations/dify', icon: 'i-custom-public-tracing-opik-icon w-[47.133904px]', panelIcon: 'i-custom-public-tracing-opik-icon-big w-[70.700851px]', + fields: [ + { name: 'api_key', label: 'API Key', required: false, secret: true, destination: false }, + { + name: 'project', + label: { key: 'tracing.configProvider.project' }, + required: false, + secret: false, + destination: false, + }, + { + name: 'workspace', + label: 'Workspace', + required: false, + secret: false, + destination: false, + placeholder: 'default', + }, + { + name: 'url', + label: 'Url', + required: false, + secret: false, + destination: true, + type: 'url', + pattern: '[Hh][Tt][Tt][Pp][Ss]?://.*/api/', + placeholder: 'https://www.comet.com/opik/api/', + }, + ], }, mlflow: { docUrl: 'https://mlflow.org/docs/latest/genai/', icon: 'i-custom-public-tracing-mlflow-icon w-11', panelIcon: 'i-custom-public-tracing-mlflow-icon-big w-16.25', + fields: [ + { + name: 'tracking_uri', + label: { key: 'tracing.configProvider.trackingUri' }, + required: true, + secret: false, + destination: true, + type: 'url', + pattern: httpUrlPattern, + placeholder: 'http://localhost:5000', + }, + { + name: 'experiment_id', + label: { key: 'tracing.configProvider.experimentId' }, + required: true, + secret: false, + destination: false, + inputMode: 'numeric', + pattern: integerPattern, + }, + { + name: 'username', + label: { key: 'tracing.configProvider.username' }, + required: false, + secret: false, + destination: false, + }, + { + name: 'password', + label: { key: 'tracing.configProvider.password' }, + required: false, + secret: true, + destination: false, + }, + ], }, databricks: { docUrl: 'https://docs.databricks.com/aws/en/mlflow3/genai/tracing/', icon: 'i-custom-public-tracing-databricks-icon w-25', panelIcon: 'i-custom-public-tracing-databricks-icon-big w-37.5', + fields: [ + { + name: 'experiment_id', + label: { key: 'tracing.configProvider.experimentId' }, + required: true, + secret: false, + destination: false, + inputMode: 'numeric', + pattern: integerPattern, + }, + { + name: 'host', + label: { key: 'tracing.configProvider.databricksHost' }, + required: true, + secret: false, + destination: true, + type: 'url', + pattern: httpUrlPattern, + }, + { + name: 'client_id', + label: { key: 'tracing.configProvider.clientId' }, + required: false, + secret: false, + destination: false, + }, + { + name: 'client_secret', + label: { key: 'tracing.configProvider.clientSecret' }, + required: false, + secret: true, + destination: false, + }, + { + name: 'personal_access_token', + label: { key: 'tracing.configProvider.personalAccessToken' }, + required: false, + secret: true, + destination: false, + }, + ], }, weave: { docUrl: 'https://weave-docs.wandb.ai/', icon: 'i-custom-public-tracing-weave-icon w-30', panelIcon: 'i-custom-public-tracing-weave-icon-big w-31', + fields: [ + { name: 'api_key', label: 'API Key', required: true, secret: true, destination: false }, + { + name: 'project', + label: { key: 'tracing.configProvider.project' }, + required: true, + secret: false, + destination: false, + }, + { name: 'entity', label: 'Entity', required: false, secret: false, destination: false }, + { + name: 'endpoint', + label: 'Endpoint', + required: false, + secret: false, + destination: true, + type: 'url', + pattern: httpsUrlPattern, + placeholder: 'https://trace.wandb.ai/', + }, + { + name: 'host', + label: 'Host', + required: false, + secret: false, + destination: true, + type: 'url', + pattern: httpUrlPattern, + placeholder: 'https://api.wandb.ai', + }, + ], }, arize: { docUrl: 'https://docs.arize.com/arize', icon: 'i-custom-public-tracing-arize-icon w-18.5', panelIcon: 'i-custom-public-tracing-arize-icon-big w-27.75', + fields: [ + { name: 'api_key', label: 'API Key', required: true, secret: true, destination: false }, + { name: 'space_id', label: 'Space ID', required: true, secret: true, destination: false }, + { + name: 'project', + label: { key: 'tracing.configProvider.project' }, + required: true, + secret: false, + destination: false, + }, + { + name: 'endpoint', + label: 'Endpoint', + required: false, + secret: false, + destination: true, + type: 'url', + pattern: httpUrlPattern, + placeholder: 'https://otlp.arize.com', + }, + ], }, phoenix: { docUrl: 'https://docs.arize.com/phoenix', icon: 'i-custom-public-tracing-phoenix-icon w-18.5', panelIcon: 'i-custom-public-tracing-phoenix-icon-big w-27.75', + fields: [ + { name: 'api_key', label: 'API Key', required: false, secret: true, destination: false }, + { + name: 'project', + label: { key: 'tracing.configProvider.project' }, + required: true, + secret: false, + destination: false, + }, + { + name: 'endpoint', + label: 'Endpoint', + required: false, + secret: false, + destination: true, + type: 'url', + pattern: httpUrlPattern, + placeholder: 'https://app.phoenix.arize.com', + }, + ], }, aliyun: { docUrl: 'https://help.aliyun.com/zh/arms/tracing-analysis/untitled-document-1750672984680', icon: 'i-custom-public-tracing-aliyun-icon w-16.25', panelIcon: 'i-custom-public-tracing-aliyun-icon-big w-24', + fields: [ + { + name: 'license_key', + label: 'License Key', + required: true, + secret: true, + destination: false, + pattern: nonBlankPattern, + }, + { + name: 'endpoint', + label: 'Endpoint', + required: true, + secret: false, + destination: true, + type: 'url', + pattern: httpUrlPattern, + placeholder: 'https://tracing.arms.aliyuncs.com', + }, + { name: 'app_name', label: 'App Name', required: true, secret: false, destination: false }, + ], }, tencent: { docUrl: 'https://cloud.tencent.com/document/product/248/116531', icon: 'i-custom-public-tracing-tencent-icon w-20', panelIcon: 'i-custom-public-tracing-tencent-icon-big w-30', + fields: [ + { + name: 'token', + label: 'Token', + required: true, + secret: true, + destination: false, + pattern: nonBlankPattern, + }, + { + name: 'endpoint', + label: 'Endpoint', + required: true, + secret: false, + destination: true, + type: 'url', + pattern: httpUrlPattern, + placeholder: 'https://ap-guangzhou.apm.tencentcs.com:4320', + description: 'tracing.tencent.endpointHelp', + }, + { + name: 'service_name', + label: 'Service Name', + required: true, + secret: false, + destination: false, + maxLength: 63, + pattern: String.raw`[a-z]([a-z0-9\-]*[a-z0-9])?`, + placeholder: 'dify-app', + description: 'tracing.tencent.serviceNameHelp', + }, + ], }, } as const satisfies Partial< Record< TraceConfigPayload['tracing_provider'], - { docUrl: string; icon: string; panelIcon: string } + { + docUrl: string + icon: string + panelIcon: string + fields: readonly TracingProviderField[] + } > > diff --git a/web/features/agent-v2/agent-detail/monitoring/__tests__/tracing-configuration.spec.tsx b/web/features/agent-v2/agent-detail/monitoring/__tests__/tracing-configuration.spec.tsx index cc0d34c970857b..cce3c7b248daf8 100644 --- a/web/features/agent-v2/agent-detail/monitoring/__tests__/tracing-configuration.spec.tsx +++ b/web/features/agent-v2/agent-detail/monitoring/__tests__/tracing-configuration.spec.tsx @@ -1,5 +1,7 @@ import type { AppTraceResponse, + TraceAppConfigResponse, + TraceConfigPayload, TraceProviderCapabilitiesResponse, } from '@dify/contracts/api/console/agent/types.gen' import { act, render, screen, waitFor, within } from '@testing-library/react' @@ -22,9 +24,20 @@ vi.mock('@/service/base', async (importOriginal) => ({ const langfuse = { public_key: 'public', - secret_key: '[__MASKED__]', + secret_key: 'secret', host: 'https://cloud.langfuse.com', } +const langfuseSecrets = ['public_key', 'secret_key'] + +/** The trace-config response for a stored Langfuse config: secrets are reported, never returned. */ +function presentLangfuse(stored: TraceConfigPayload['tracing_config']): TraceAppConfigResponse { + return { + tracing_config: Object.fromEntries( + Object.entries(stored).filter(([name]) => !langfuseSecrets.includes(name)), + ), + configured_secrets: langfuseSecrets.filter((name) => !!stored[name]), + } +} const capabilities: TraceProviderCapabilitiesResponse = { provider: 'langfuse', agent_spans: true, @@ -97,7 +110,7 @@ function setup({ if (url.pathname.endsWith('/trace-config')) { if (req.method === 'GET') { const config = url.searchParams.get('tracing_provider') === 'langfuse' ? state.config : null - return Response.json({ tracing_config: config, has_not_configured: !config }) + return Response.json(config ? presentLangfuse(config) : { has_not_configured: true }) } if (req.method === 'DELETE') { if (state.failDelete) @@ -118,7 +131,8 @@ function setup({ const body = await req.json() if (body.activate) state.policy = { ...state.policy, enabled: true, tracing_provider: body.tracing_provider } - state.config = { ...body.tracing_config, secret_key: '[__MASKED__]' } + // Like the API, a secret left out of the save keeps its stored value. + state.config = { ...state.config, ...body.tracing_config } return Response.json({ result: 'success' }, { status: req.method === 'POST' ? 201 : 200 }) } throw new Error(`Unexpected request ${req.method} ${url.pathname}`) @@ -380,7 +394,7 @@ describe('Agent tracing configuration', () => { ).toEqual(['POST trace-config', 'POST trace-config']) expect(await writes(sent)[1]!.json()).toEqual({ tracing_provider: 'langfuse', - tracing_config: { ...langfuse, secret_key: 'secret', host: 'https://updated.langfuse.com' }, + tracing_config: { ...langfuse, host: 'https://updated.langfuse.com' }, activate: true, }) expect(screen.getByRole('switch', { name: 'app.tracing.tracing' })).toBeChecked() @@ -390,7 +404,7 @@ describe('Agent tracing configuration', () => { capture_content: false, capture_agent: true, }) - expect(state.config?.secret_key).toBe('[__MASKED__]') + expect(state.config).toEqual({ ...langfuse, host: 'https://updated.langfuse.com' }) }) it('preserves the draft after concurrent deletion and clears the conflict when retrying creation', async () => { @@ -400,6 +414,11 @@ describe('Agent tracing configuration', () => { const host = screen.getByRole('textbox', { name: 'Host' }) await user.clear(host) await user.type(host, 'https://draft.langfuse.com') + // A new destination needs its secrets entered again. + const publicKey = screen.getByRole('textbox', { name: 'app.tracing.configProvider.publicKey' }) + const secretKey = screen.getByRole('textbox', { name: 'app.tracing.configProvider.secretKey' }) + await user.type(publicKey, 'public') + await user.type(secretKey, 'secret') const server = request.getMockImplementation()! request.mockImplementation(async (...args) => { if (args[2].request.method === 'PATCH') { @@ -420,6 +439,8 @@ describe('Agent tracing configuration', () => { await user.click(screen.getByRole('button', { name: 'common.operation.save' })) expect(await screen.findByRole('alert')).toHaveTextContent('Tracing configuration changed') expect(host).toHaveValue('https://draft.langfuse.com') + expect(publicKey).toHaveValue('public') + expect(secretKey).toHaveValue('secret') expect(screen.getByRole('button', { name: 'common.operation.saveAndEnable' })).toBeEnabled() expect( screen.getByRole('switch', { name: 'app.tracing.tracing', hidden: true }), @@ -433,10 +454,6 @@ describe('Agent tracing configuration', () => { if (args[2].request.method === 'POST') await retry return server(...args) }) - await user.type( - screen.getByRole('textbox', { name: 'app.tracing.configProvider.secretKey' }), - 'secret', - ) await user.click(screen.getByRole('button', { name: 'common.operation.saveAndEnable' })) expect(screen.queryByRole('alert')).not.toBeInTheDocument() expect(host).toHaveValue('https://draft.langfuse.com') @@ -475,7 +492,7 @@ describe('Agent tracing configuration', () => { }) }) - it('refreshes capabilities after updating masked credentials', async () => { + it('refreshes capabilities after saving with the saved credentials kept', async () => { const { sent } = setup() const user = await openPanel() await openProvider(user) @@ -484,7 +501,10 @@ describe('Agent tracing configuration', () => { expect(screen.queryByRole('textbox', { name: 'Host' })).not.toBeInTheDocument(), ) const patch = writes(sent).find((req) => req.method === 'PATCH')! - expect(await patch.json()).toEqual({ tracing_provider: 'langfuse', tracing_config: langfuse }) + expect(await patch.json()).toEqual({ + tracing_provider: 'langfuse', + tracing_config: { host: langfuse.host }, + }) expect( sent.filter((req) => new URL(req.url).pathname.endsWith('/trace-capabilities')).length, ).toBeGreaterThan(1) diff --git a/web/i18n/locales/ar-TN/app.json b/web/i18n/locales/ar-TN/app.json index cf3a10d786bd5f..05c9807b6edfa1 100644 --- a/web/i18n/locales/ar-TN/app.json +++ b/web/i18n/locales/ar-TN/app.json @@ -268,6 +268,7 @@ "tracing.config": "تكوين", "tracing.configProvider.clientId": "معرف العميل (Client ID)", "tracing.configProvider.clientSecret": "سر العميل (Client Secret)", + "tracing.configProvider.configuredSecretPlaceholder": "تم التكوين. اتركه فارغًا للاحتفاظ بالقيمة المحفوظة.", "tracing.configProvider.databricksHost": "عنوان URL لمساحة عمل Databricks", "tracing.configProvider.experimentId": "معرف التجربة", "tracing.configProvider.password": "كلمة المرور", diff --git a/web/i18n/locales/az-AZ/app.json b/web/i18n/locales/az-AZ/app.json index 4d059cf6bdd6ae..a6b00c0669152c 100644 --- a/web/i18n/locales/az-AZ/app.json +++ b/web/i18n/locales/az-AZ/app.json @@ -268,6 +268,7 @@ "tracing.config": "Konfiqurasiya", "tracing.configProvider.clientId": "OAuth müştəri ID-si", "tracing.configProvider.clientSecret": "OAuth müştəri sirri", + "tracing.configProvider.configuredSecretPlaceholder": "Konfiqurasiya edilib. Saxlanılmış dəyəri dəyişməmək üçün boş saxlayın.", "tracing.configProvider.databricksHost": "Databricks iş sahəsinin URL-i", "tracing.configProvider.experimentId": "Təcrübə ID-si", "tracing.configProvider.password": "Parol", diff --git a/web/i18n/locales/de-DE/app.json b/web/i18n/locales/de-DE/app.json index 70ce0a528fcb4b..9362865aed55a8 100644 --- a/web/i18n/locales/de-DE/app.json +++ b/web/i18n/locales/de-DE/app.json @@ -268,6 +268,7 @@ "tracing.config": "Konfigurieren", "tracing.configProvider.clientId": "OAuth-Client-ID", "tracing.configProvider.clientSecret": "OAuth-Client-Geheimnis", + "tracing.configProvider.configuredSecretPlaceholder": "Konfiguriert. Leer lassen, um den gespeicherten Wert beizubehalten.", "tracing.configProvider.databricksHost": "Databricks-Workspace-URL", "tracing.configProvider.experimentId": "Experiment-ID", "tracing.configProvider.password": "Passwort", diff --git a/web/i18n/locales/en-US/app.json b/web/i18n/locales/en-US/app.json index 49662311ebdc32..8af1a990cdfa24 100644 --- a/web/i18n/locales/en-US/app.json +++ b/web/i18n/locales/en-US/app.json @@ -268,6 +268,7 @@ "tracing.config": "Config", "tracing.configProvider.clientId": "OAuth Client ID", "tracing.configProvider.clientSecret": "OAuth Client Secret", + "tracing.configProvider.configuredSecretPlaceholder": "Configured. Leave blank to keep the saved value.", "tracing.configProvider.databricksHost": "Databricks Workspace URL", "tracing.configProvider.experimentId": "Experiment ID", "tracing.configProvider.password": "Password", diff --git a/web/i18n/locales/es-ES/app.json b/web/i18n/locales/es-ES/app.json index 13114688ca0812..989f3db839544c 100644 --- a/web/i18n/locales/es-ES/app.json +++ b/web/i18n/locales/es-ES/app.json @@ -268,6 +268,7 @@ "tracing.config": "Configurar", "tracing.configProvider.clientId": "ID de cliente OAuth", "tracing.configProvider.clientSecret": "Secreto del cliente OAuth", + "tracing.configProvider.configuredSecretPlaceholder": "Configurado. Déjalo en blanco para conservar el valor guardado.", "tracing.configProvider.databricksHost": "URL del espacio de trabajo de Databricks", "tracing.configProvider.experimentId": "ID del experimento", "tracing.configProvider.password": "Contraseña", diff --git a/web/i18n/locales/fa-IR/app.json b/web/i18n/locales/fa-IR/app.json index fef24733a68e8f..a22c23fe979dc2 100644 --- a/web/i18n/locales/fa-IR/app.json +++ b/web/i18n/locales/fa-IR/app.json @@ -268,6 +268,7 @@ "tracing.config": "پیکربندی", "tracing.configProvider.clientId": "شناسه مشتری OAuth", "tracing.configProvider.clientSecret": "رمز مخفی مشتری OAuth", + "tracing.configProvider.configuredSecretPlaceholder": "پیکربندی شده. برای حفظ مقدار ذخیره‌شده، آن را خالی بگذارید.", "tracing.configProvider.databricksHost": "نشانی اینترنتی محیط کاری دیتابریکس", "tracing.configProvider.experimentId": "شناسه آزمایش", "tracing.configProvider.password": "رمز عبور", diff --git a/web/i18n/locales/fr-FR/app.json b/web/i18n/locales/fr-FR/app.json index 1340a1b3536fcd..f71990d0852612 100644 --- a/web/i18n/locales/fr-FR/app.json +++ b/web/i18n/locales/fr-FR/app.json @@ -268,6 +268,7 @@ "tracing.config": "Configurer", "tracing.configProvider.clientId": "ID client OAuth", "tracing.configProvider.clientSecret": "Secret client OAuth", + "tracing.configProvider.configuredSecretPlaceholder": "Configuré. Laissez vide pour conserver la valeur enregistrée.", "tracing.configProvider.databricksHost": "URL de l'espace de travail Databricks", "tracing.configProvider.experimentId": "ID de l'expérience", "tracing.configProvider.password": "Mot de passe", diff --git a/web/i18n/locales/hi-IN/app.json b/web/i18n/locales/hi-IN/app.json index 3b60274e09736b..260efcf692f9d6 100644 --- a/web/i18n/locales/hi-IN/app.json +++ b/web/i18n/locales/hi-IN/app.json @@ -268,6 +268,7 @@ "tracing.config": "कॉन्फ़िगर करें", "tracing.configProvider.clientId": "OAuth क्लाइंट आईडी", "tracing.configProvider.clientSecret": "OAuth क्लाइंट सीक्रेट", + "tracing.configProvider.configuredSecretPlaceholder": "कॉन्फ़िगर किया गया। सहेजा गया मान बनाए रखने के लिए इसे खाली छोड़ें।", "tracing.configProvider.databricksHost": "डेटाब्रिक्स वर्कस्पेस यूआरएल", "tracing.configProvider.experimentId": "प्रयोग आईडी", "tracing.configProvider.password": "पासवर्ड", diff --git a/web/i18n/locales/id-ID/app.json b/web/i18n/locales/id-ID/app.json index a05bb69016a18c..b12d9ecece0760 100644 --- a/web/i18n/locales/id-ID/app.json +++ b/web/i18n/locales/id-ID/app.json @@ -268,6 +268,7 @@ "tracing.config": "Konfigurasi", "tracing.configProvider.clientId": "ID Klien OAuth", "tracing.configProvider.clientSecret": "Rahasia Klien OAuth", + "tracing.configProvider.configuredSecretPlaceholder": "Dikonfigurasi. Kosongkan untuk mempertahankan nilai yang tersimpan.", "tracing.configProvider.databricksHost": "URL Workspace Databricks", "tracing.configProvider.experimentId": "ID Eksperimen", "tracing.configProvider.password": "Kata sandi", diff --git a/web/i18n/locales/it-IT/app.json b/web/i18n/locales/it-IT/app.json index 21c873cd2b6f59..dbe44f0249aa5b 100644 --- a/web/i18n/locales/it-IT/app.json +++ b/web/i18n/locales/it-IT/app.json @@ -268,6 +268,7 @@ "tracing.config": "Config", "tracing.configProvider.clientId": "ID client OAuth", "tracing.configProvider.clientSecret": "Segreto del client OAuth", + "tracing.configProvider.configuredSecretPlaceholder": "Configurato. Lascia vuoto per mantenere il valore salvato.", "tracing.configProvider.databricksHost": "URL dell'area di lavoro Databricks", "tracing.configProvider.experimentId": "ID Esperimento", "tracing.configProvider.password": "Password", diff --git a/web/i18n/locales/ja-JP/app.json b/web/i18n/locales/ja-JP/app.json index 9c3ca4aa2ea0bd..2900984162c4d0 100644 --- a/web/i18n/locales/ja-JP/app.json +++ b/web/i18n/locales/ja-JP/app.json @@ -268,6 +268,7 @@ "tracing.config": "設定", "tracing.configProvider.clientId": "OAuthクライアントID", "tracing.configProvider.clientSecret": "OAuthクライアントシークレット", + "tracing.configProvider.configuredSecretPlaceholder": "設定済みです。空欄のままにすると保存済みの値を維持します。", "tracing.configProvider.databricksHost": "DatabricksワークスペースのURL", "tracing.configProvider.experimentId": "実験 ID", "tracing.configProvider.password": "パスワード", diff --git a/web/i18n/locales/ko-KR/app.json b/web/i18n/locales/ko-KR/app.json index 37ff1982d9411a..cf7e6d9a95b70a 100644 --- a/web/i18n/locales/ko-KR/app.json +++ b/web/i18n/locales/ko-KR/app.json @@ -268,6 +268,7 @@ "tracing.config": "구성", "tracing.configProvider.clientId": "OAuth 클라이언트 ID", "tracing.configProvider.clientSecret": "OAuth 클라이언트 비밀", + "tracing.configProvider.configuredSecretPlaceholder": "구성됨. 비워 두면 저장된 값이 유지됩니다.", "tracing.configProvider.databricksHost": "Databricks 작업 영역 URL", "tracing.configProvider.experimentId": "실험 ID", "tracing.configProvider.password": "비밀번호", diff --git a/web/i18n/locales/lo-LA/app.json b/web/i18n/locales/lo-LA/app.json index de91addbefbb81..b269c55c5fe06e 100644 --- a/web/i18n/locales/lo-LA/app.json +++ b/web/i18n/locales/lo-LA/app.json @@ -268,6 +268,7 @@ "tracing.config": "ຕັ້ງຄ່າ", "tracing.configProvider.clientId": "OAuth Client ID", "tracing.configProvider.clientSecret": "OAuth Client Secret", + "tracing.configProvider.configuredSecretPlaceholder": "ຕັ້ງຄ່າແລ້ວ. ປະຫວ່າງໄວ້ເພື່ອຮັກສາຄ່າທີ່ບັນທຶກໄວ້.", "tracing.configProvider.databricksHost": "Databricks Workspace URL", "tracing.configProvider.experimentId": "Experiment ID", "tracing.configProvider.password": "ລະຫັດຜ່ານ", diff --git a/web/i18n/locales/nl-NL/app.json b/web/i18n/locales/nl-NL/app.json index cfe46e13b691e0..7cbeb058849109 100644 --- a/web/i18n/locales/nl-NL/app.json +++ b/web/i18n/locales/nl-NL/app.json @@ -268,6 +268,7 @@ "tracing.config": "Config", "tracing.configProvider.clientId": "OAuth Client ID", "tracing.configProvider.clientSecret": "OAuth Client Secret", + "tracing.configProvider.configuredSecretPlaceholder": "Geconfigureerd. Laat leeg om de opgeslagen waarde te behouden.", "tracing.configProvider.databricksHost": "Databricks Workspace URL", "tracing.configProvider.experimentId": "Experiment ID", "tracing.configProvider.password": "Password", diff --git a/web/i18n/locales/pl-PL/app.json b/web/i18n/locales/pl-PL/app.json index ad190de415feea..f288242b732d86 100644 --- a/web/i18n/locales/pl-PL/app.json +++ b/web/i18n/locales/pl-PL/app.json @@ -268,6 +268,7 @@ "tracing.config": "Konfiguruj", "tracing.configProvider.clientId": "ID klienta OAuth", "tracing.configProvider.clientSecret": "Sekretny klucz klienta OAuth", + "tracing.configProvider.configuredSecretPlaceholder": "Skonfigurowano. Pozostaw puste, aby zachować zapisaną wartość.", "tracing.configProvider.databricksHost": "Adres URL obszaru roboczego Databricks", "tracing.configProvider.experimentId": "ID eksperymentu", "tracing.configProvider.password": "Hasło", diff --git a/web/i18n/locales/pt-BR/app.json b/web/i18n/locales/pt-BR/app.json index 745de1d646d2d4..8bbd186fe71d67 100644 --- a/web/i18n/locales/pt-BR/app.json +++ b/web/i18n/locales/pt-BR/app.json @@ -268,6 +268,7 @@ "tracing.config": "Configurar", "tracing.configProvider.clientId": "ID do Cliente OAuth", "tracing.configProvider.clientSecret": "Segredo do Cliente OAuth", + "tracing.configProvider.configuredSecretPlaceholder": "Configurado. Deixe em branco para manter o valor salvo.", "tracing.configProvider.databricksHost": "URL do Workspace do Databricks", "tracing.configProvider.experimentId": "ID do Experimento", "tracing.configProvider.password": "Senha", diff --git a/web/i18n/locales/ro-RO/app.json b/web/i18n/locales/ro-RO/app.json index 93af3227a472d8..7e459b48d12f9b 100644 --- a/web/i18n/locales/ro-RO/app.json +++ b/web/i18n/locales/ro-RO/app.json @@ -268,6 +268,7 @@ "tracing.config": "Configurare", "tracing.configProvider.clientId": "ID client OAuth", "tracing.configProvider.clientSecret": "Secret client OAuth", + "tracing.configProvider.configuredSecretPlaceholder": "Configurat. Lăsați gol pentru a păstra valoarea salvată.", "tracing.configProvider.databricksHost": "URL-ul spațiului de lucru Databricks", "tracing.configProvider.experimentId": "ID-ul experimentului", "tracing.configProvider.password": "Parolă", diff --git a/web/i18n/locales/ru-RU/app.json b/web/i18n/locales/ru-RU/app.json index 690ab5230e7184..2b85429845c100 100644 --- a/web/i18n/locales/ru-RU/app.json +++ b/web/i18n/locales/ru-RU/app.json @@ -268,6 +268,7 @@ "tracing.config": "Настройка", "tracing.configProvider.clientId": "Идентификатор клиента OAuth", "tracing.configProvider.clientSecret": "Секрет клиента OAuth", + "tracing.configProvider.configuredSecretPlaceholder": "Настроено. Оставьте пустым, чтобы не менять сохранённое значение.", "tracing.configProvider.databricksHost": "URL рабочего пространства Databricks", "tracing.configProvider.experimentId": "ID эксперимента", "tracing.configProvider.password": "Пароль", diff --git a/web/i18n/locales/sl-SI/app.json b/web/i18n/locales/sl-SI/app.json index 07764df1331a59..9671bb4ce82dcc 100644 --- a/web/i18n/locales/sl-SI/app.json +++ b/web/i18n/locales/sl-SI/app.json @@ -268,6 +268,7 @@ "tracing.config": "Konfiguracija", "tracing.configProvider.clientId": "ID odjemalca OAuth", "tracing.configProvider.clientSecret": "OAuth skrivnost odjemalca", + "tracing.configProvider.configuredSecretPlaceholder": "Konfigurirano. Pustite prazno, da ohranite shranjeno vrednost.", "tracing.configProvider.databricksHost": "URL delovnega prostora Databricks", "tracing.configProvider.experimentId": "ID eksperimenta", "tracing.configProvider.password": "Geslo", diff --git a/web/i18n/locales/th-TH/app.json b/web/i18n/locales/th-TH/app.json index 05c8bce96a53b1..2a10b9be921405 100644 --- a/web/i18n/locales/th-TH/app.json +++ b/web/i18n/locales/th-TH/app.json @@ -268,6 +268,7 @@ "tracing.config": "กําหนดค่า", "tracing.configProvider.clientId": "รหัสลูกค้า OAuth", "tracing.configProvider.clientSecret": "รหัสลับของลูกค้า OAuth", + "tracing.configProvider.configuredSecretPlaceholder": "กำหนดค่าแล้ว เว้นว่างไว้เพื่อคงค่าที่บันทึกไว้", "tracing.configProvider.databricksHost": "URL ของ Workspace ใน Databricks", "tracing.configProvider.experimentId": "รหัสการทดลอง", "tracing.configProvider.password": "รหัสผ่าน", diff --git a/web/i18n/locales/tr-TR/app.json b/web/i18n/locales/tr-TR/app.json index ae0df7b4e37337..5abb807453b1a8 100644 --- a/web/i18n/locales/tr-TR/app.json +++ b/web/i18n/locales/tr-TR/app.json @@ -268,6 +268,7 @@ "tracing.config": "Yapılandırma", "tracing.configProvider.clientId": "OAuth İstemci Kimliği", "tracing.configProvider.clientSecret": "OAuth İstemci Sırrı", + "tracing.configProvider.configuredSecretPlaceholder": "Yapılandırıldı. Kayıtlı değeri korumak için boş bırakın.", "tracing.configProvider.databricksHost": "Databricks Çalışma Alanı URL'si", "tracing.configProvider.experimentId": "Deney Kimliği", "tracing.configProvider.password": "Parola", diff --git a/web/i18n/locales/uk-UA/app.json b/web/i18n/locales/uk-UA/app.json index 2c47a01fbb3dd1..7bde74f4bb1215 100644 --- a/web/i18n/locales/uk-UA/app.json +++ b/web/i18n/locales/uk-UA/app.json @@ -268,6 +268,7 @@ "tracing.config": "Налаштувати", "tracing.configProvider.clientId": "Ідентифікатор клієнта OAuth", "tracing.configProvider.clientSecret": "Секретний ключ клієнта OAuth", + "tracing.configProvider.configuredSecretPlaceholder": "Налаштовано. Залиште порожнім, щоб не змінювати збережене значення.", "tracing.configProvider.databricksHost": "URL робочого простору Databricks", "tracing.configProvider.experimentId": "Ідентифікатор експерименту", "tracing.configProvider.password": "Пароль", diff --git a/web/i18n/locales/vi-VN/app.json b/web/i18n/locales/vi-VN/app.json index b9f9deafe28795..172f56f35e4112 100644 --- a/web/i18n/locales/vi-VN/app.json +++ b/web/i18n/locales/vi-VN/app.json @@ -268,6 +268,7 @@ "tracing.config": "Cấu hình", "tracing.configProvider.clientId": "ID Khách Hàng OAuth", "tracing.configProvider.clientSecret": "Bí mật Khách hàng OAuth", + "tracing.configProvider.configuredSecretPlaceholder": "Đã cấu hình. Để trống để giữ giá trị đã lưu.", "tracing.configProvider.databricksHost": "URL Workspace của Databricks", "tracing.configProvider.experimentId": "Mã thí nghiệm", "tracing.configProvider.password": "Mật khẩu", diff --git a/web/i18n/locales/zh-Hans/app.json b/web/i18n/locales/zh-Hans/app.json index 227dd73300ffe6..07e567b9fd365c 100644 --- a/web/i18n/locales/zh-Hans/app.json +++ b/web/i18n/locales/zh-Hans/app.json @@ -268,6 +268,7 @@ "tracing.config": "配置", "tracing.configProvider.clientId": "OAuth 客户端 ID", "tracing.configProvider.clientSecret": "OAuth 客户端密钥", + "tracing.configProvider.configuredSecretPlaceholder": "已配置。留空则保留已保存的值。", "tracing.configProvider.databricksHost": "Databricks 工作区 URL", "tracing.configProvider.experimentId": "实验编号", "tracing.configProvider.password": "密码", diff --git a/web/i18n/locales/zh-Hant/app.json b/web/i18n/locales/zh-Hant/app.json index 13556f07b66f4d..28097cb7f64314 100644 --- a/web/i18n/locales/zh-Hant/app.json +++ b/web/i18n/locales/zh-Hant/app.json @@ -268,6 +268,7 @@ "tracing.config": "配置", "tracing.configProvider.clientId": "OAuth 用戶端 ID", "tracing.configProvider.clientSecret": "OAuth 用戶端密鑰", + "tracing.configProvider.configuredSecretPlaceholder": "已配置。留空則保留已儲存的值。", "tracing.configProvider.databricksHost": "Databricks 工作區網址", "tracing.configProvider.experimentId": "實驗編號", "tracing.configProvider.password": "密碼", diff --git a/web/i18n/resources.generated.d.ts b/web/i18n/resources.generated.d.ts index 5f7032d3078158..998b0354e16f5d 100644 --- a/web/i18n/resources.generated.d.ts +++ b/web/i18n/resources.generated.d.ts @@ -807,6 +807,7 @@ export default interface Resources { 'tracing.config': 'Config' 'tracing.configProvider.clientId': 'OAuth Client ID' 'tracing.configProvider.clientSecret': 'OAuth Client Secret' + 'tracing.configProvider.configuredSecretPlaceholder': 'Configured. Leave blank to keep the saved value.' 'tracing.configProvider.databricksHost': 'Databricks Workspace URL' 'tracing.configProvider.experimentId': 'Experiment ID' 'tracing.configProvider.password': 'Password'